From d6155b290294ec3f51f2159319b4889a3ab40853 Mon Sep 17 00:00:00 2001 From: Jackson57279 Date: Thu, 7 May 2026 11:57:26 -0500 Subject: [PATCH 1/8] build: add Cargo workspace and update build system for Rust port --- .devops/nix/package.nix | 2 +- .editorconfig | 2 +- .github/labeler.yml | 22 +- .../workflows/build-and-test-snapdragon.yml | 4 +- CMakeLists.txt | 3 +- CODEOWNERS | 4 +- CONTRIBUTING.md | 2 +- Cargo.lock | 700 + Cargo.toml | 116 + README.md | 4 +- build-xcframework.sh | 34 +- common/CMakeLists.txt | 252 +- common/arg.cpp | 4098 - common/arg.h | 133 - common/base64.hpp | 392 - common/build-info-rust/Cargo.toml | 10 + common/build-info-rust/build.rs | 15 + common/build-info-rust/src/lib.rs | 66 + common/build-info.cpp.in | 35 - common/build-info.h | 11 - common/chat-auto-parser-generator.cpp | 478 - common/chat-auto-parser-helpers.cpp | 364 - common/chat-auto-parser-helpers.h | 74 - common/chat-auto-parser.h | 438 - common/chat-diff-analyzer.cpp | 1395 - common/chat-peg-parser.cpp | 1007 - common/chat-peg-parser.h | 195 - common/chat.cpp | 2384 - common/chat.h | 284 - common/common.cpp | 1962 - common/common.cpp.inc | 23896 ++++++ common/common.h | 1028 - common/common.h.inc | 5003 ++ common/console.cpp | 1166 - common/console.h | 46 - common/debug.cpp | 190 - common/debug.h | 31 - common/download.cpp | 958 - common/download.h | 107 - common/fit.cpp | 951 - common/fit.h | 32 - common/hf-cache.cpp | 772 - common/hf-cache.h | 36 - common/http.h | 99 - common/jinja/README.md | 6 +- common/jinja/caps.cpp | 479 - common/jinja/caps.h | 32 - common/jinja/lexer.cpp | 341 - common/jinja/lexer.h | 157 - common/jinja/parser.cpp | 602 - common/jinja/parser.h | 21 - common/jinja/runtime.cpp | 906 - common/jinja/runtime.h | 652 - common/jinja/string-rust/Cargo.toml | 10 + common/jinja/string-rust/src/lib.rs | 156 + common/jinja/string.cpp | 213 - common/jinja/string.h | 61 - common/jinja/utils.h | 149 - common/jinja/value.cpp | 1484 - common/jinja/value.h | 759 - common/json-partial.cpp | 324 - common/json-partial.h | 39 - common/json-schema-to-grammar.cpp | 1189 - common/json-schema-to-grammar.h | 43 - common/llguidance.cpp | 258 - common/log.cpp | 453 - common/log.h | 123 - common/ngram-cache-rust/Cargo.toml | 7 + common/ngram-cache-rust/src/lib.rs | 355 + common/ngram-cache.cpp | 285 - common/ngram-cache.h | 101 - common/ngram-map-rust/Cargo.toml | 7 + common/ngram-map-rust/src/lib.rs | 453 + common/ngram-map.cpp | 530 - common/ngram-map.h | 115 - common/ngram-mod-rust/Cargo.toml | 10 + common/ngram-mod-rust/src/lib.rs | 138 + common/ngram-mod.cpp | 60 - common/ngram-mod.h | 38 - common/peg-parser.cpp | 2128 - common/peg-parser.h | 523 - common/preset.cpp | 483 - common/preset.h | 83 - common/reasoning-budget-rust/Cargo.toml | 10 + common/reasoning-budget-rust/src/lib.rs | 420 + common/reasoning-budget.cpp | 250 - common/reasoning-budget.h | 42 - common/regex-partial-rust/Cargo.toml | 8 + common/regex-partial-rust/src/lib.rs | 249 + common/regex-partial.cpp | 204 - common/regex-partial.h | 56 - common/sampling.cpp | 843 - common/sampling.h | 119 - common/speculative.cpp | 1248 - common/speculative.h | 46 - common/unicode-rust/Cargo.toml | 10 + common/unicode-rust/src/lib.rs | 5233 ++ common/unicode.cpp | 124 - common/unicode.h | 30 - docs/autoparser.md | 17 +- docs/build-s390x.md | 2 +- docs/build.md | 2 +- docs/development/HOWTO-add-model.md | 14 +- docs/development/parsing.md | 2 +- docs/development/rust-port.md | 21 + docs/function-calling.md | 2 +- docs/preset.md | 2 +- examples/batched-rust/Cargo.toml | 12 + examples/batched-rust/build.rs | 45 + examples/batched-rust/src/main.rs | 595 + examples/batched/CMakeLists.txt | 28 +- examples/batched/batched.cpp | 264 - .../convert-llama2c-to-ggml-rust/Cargo.toml | 9 + .../convert-llama2c-to-ggml-rust/src/main.rs | 279 + .../convert-llama2c-to-ggml/CMakeLists.txt | 20 +- .../convert-llama2c-to-ggml.cpp | 945 - examples/debug-rust/Cargo.toml | 12 + examples/debug-rust/build.rs | 43 + examples/debug-rust/src/lib.rs | 148 + examples/debug-rust/src/main.rs | 316 + examples/debug/CMakeLists.txt | 29 +- examples/debug/debug.cpp | 261 - examples/deprecation-warning/Cargo.toml | 7 + .../deprecation-warning.cpp | 38 - examples/deprecation-warning/src/lib.rs | 53 + examples/deprecation-warning/src/main.rs | 7 + examples/diffusion-rust/Cargo.toml | 12 + examples/diffusion-rust/build.rs | 45 + examples/diffusion-rust/src/main.rs | 1114 + examples/diffusion/CMakeLists.txt | 28 +- examples/diffusion/diffusion-cli.cpp | 698 - examples/embedding-rust/Cargo.toml | 12 + examples/embedding-rust/build.rs | 45 + examples/embedding-rust/src/main.rs | 966 + examples/embedding/CMakeLists.txt | 28 +- examples/embedding/embedding.cpp | 414 - examples/eval-callback-rust/Cargo.toml | 12 + examples/eval-callback-rust/build.rs | 43 + examples/eval-callback-rust/src/lib.rs | 109 + examples/eval-callback-rust/src/main.rs | 177 + examples/eval-callback/CMakeLists.txt | 29 +- examples/eval-callback/eval-callback.cpp | 88 - examples/gen-docs-rust/Cargo.toml | 9 + examples/gen-docs-rust/src/main.rs | 176 + examples/gen-docs/CMakeLists.txt | 22 +- examples/gen-docs/gen-docs.cpp | 145 - examples/gguf-hash-rust/Cargo.toml | 14 + examples/gguf-hash-rust/build.rs | 49 + examples/gguf-hash-rust/src/lib.rs | 300 + examples/gguf-hash-rust/src/main.rs | 405 + examples/gguf-hash/CMakeLists.txt | 37 +- .../gguf-hash/deps/rotate-bits/rotate-bits.h | 46 - examples/gguf-hash/deps/sha1/sha1.c | 295 - examples/gguf-hash/deps/sha1/sha1.h | 52 - examples/gguf-hash/deps/sha256/sha256.c | 221 - examples/gguf-hash/deps/sha256/sha256.h | 24 - examples/gguf-hash/deps/xxhash/xxhash.c | 42 - examples/gguf-hash/deps/xxhash/xxhash.h | 7093 -- examples/gguf-hash/gguf-hash.cpp | 697 - examples/gguf-rust/Cargo.toml | 9 + examples/gguf-rust/build.rs | 47 + examples/gguf-rust/src/lib.rs | 57 + examples/gguf-rust/src/main.rs | 435 + examples/gguf/CMakeLists.txt | 26 +- examples/gguf/gguf.cpp | 273 - examples/idle-rust/Cargo.toml | 12 + examples/idle-rust/src/lib.rs | 99 + examples/idle-rust/src/main.rs | 148 + examples/idle/CMakeLists.txt | 28 +- examples/idle/idle.cpp | 110 - .../lib/src/main/cpp/CMakeLists.txt | 2 +- .../main/cpp/{ai_chat.cpp => ai_chat.cpp.inc} | 61 +- .../llama.android/lib/src/main/cpp/logging.h | 61 - .../aichat/internal/InferenceEngineImpl.kt | 4 +- examples/lookahead-rust/Cargo.toml | 12 + examples/lookahead-rust/build.rs | 45 + examples/lookahead-rust/src/main.rs | 761 + examples/lookahead/CMakeLists.txt | 28 +- examples/lookahead/lookahead.cpp | 483 - examples/lookup-create-rust/Cargo.toml | 13 + examples/lookup-create-rust/build.rs | 43 + examples/lookup-create-rust/src/lib.rs | 154 + examples/lookup-create-rust/src/main.rs | 113 + examples/lookup-merge-rust/Cargo.toml | 7 + examples/lookup-merge-rust/src/lib.rs | 470 + examples/lookup-merge-rust/src/main.rs | 68 + examples/lookup-rust/Cargo.toml | 13 + examples/lookup-rust/build.rs | 45 + examples/lookup-rust/src/main.rs | 696 + examples/lookup-stats-rust/Cargo.toml | 13 + examples/lookup-stats-rust/build.rs | 43 + examples/lookup-stats-rust/src/lib.rs | 143 + examples/lookup-stats-rust/src/main.rs | 269 + examples/lookup/CMakeLists.txt | 108 +- examples/lookup/lookup-create.cpp | 45 - examples/lookup/lookup-merge.cpp | 50 - examples/lookup/lookup-stats.cpp | 160 - examples/lookup/lookup.cpp | 245 - examples/parallel-rust/Cargo.toml | 12 + examples/parallel-rust/build.rs | 45 + examples/parallel-rust/src/main.rs | 917 + examples/parallel/CMakeLists.txt | 28 +- examples/parallel/parallel.cpp | 520 - examples/passkey-rust/Cargo.toml | 12 + examples/passkey-rust/build.rs | 45 + examples/passkey-rust/src/main.rs | 642 + examples/passkey/CMakeLists.txt | 28 +- examples/passkey/passkey.cpp | 277 - examples/retrieval-rust/Cargo.toml | 12 + examples/retrieval-rust/build.rs | 45 + examples/retrieval-rust/src/main.rs | 687 + examples/retrieval/CMakeLists.txt | 28 +- examples/retrieval/retrieval.cpp | 307 - examples/save-load-state-rust/Cargo.toml | 12 + examples/save-load-state-rust/build.rs | 45 + examples/save-load-state-rust/src/main.rs | 630 + examples/save-load-state/CMakeLists.txt | 28 +- examples/save-load-state/save-load-state.cpp | 238 - examples/simple-chat-rust/Cargo.toml | 12 + examples/simple-chat-rust/build.rs | 3 + examples/simple-chat-rust/src/lib.rs | 109 + examples/simple-chat-rust/src/main.rs | 325 + examples/simple-chat/CMakeLists.txt | 28 +- examples/simple-chat/simple-chat.cpp | 210 - examples/simple-cmake-pkg/CMakeLists.txt | 26 +- examples/simple-rust/Cargo.toml | 13 + examples/simple-rust/README.md | 14 + examples/simple-rust/build.rs | 46 + examples/simple-rust/src/ffi.rs | 555 + examples/simple-rust/src/lib.rs | 144 + examples/simple-rust/src/main.rs | 240 + examples/simple/CMakeLists.txt | 27 +- examples/simple/simple.cpp | 223 - examples/speculative-rust/Cargo.toml | 12 + examples/speculative-rust/build.rs | 45 + examples/speculative-rust/src/main.rs | 1 + examples/speculative-simple-rust/Cargo.toml | 12 + examples/speculative-simple-rust/build.rs | 45 + examples/speculative-simple-rust/src/main.rs | 663 + examples/speculative-simple/CMakeLists.txt | 28 +- .../speculative-simple/speculative-simple.cpp | 348 - examples/speculative/CMakeLists.txt | 29 +- examples/speculative/speculative.cpp | 660 - examples/sycl/CMakeLists.txt | 26 +- examples/sycl/ls-sycl-device-rust/Cargo.toml | 11 + examples/sycl/ls-sycl-device-rust/build.rs | 42 + examples/sycl/ls-sycl-device-rust/src/lib.rs | 13 + examples/sycl/ls-sycl-device-rust/src/main.rs | 11 + examples/sycl/ls-sycl-device.cpp | 15 - examples/training/CMakeLists.txt | 28 +- examples/training/finetune-rust/Cargo.toml | 12 + examples/training/finetune-rust/build.rs | 45 + examples/training/finetune-rust/src/main.rs | 642 + examples/training/finetune.cpp | 101 - ggml/CMakeLists.txt | 35 +- .../{ggml-alloc.h => ggml-alloc.h.inc} | 2 +- .../{ggml-backend.h => ggml-backend.h.inc} | 4 +- ggml/include/{ggml-blas.h => ggml-blas.h.inc} | 4 +- ggml/include/{ggml-cann.h => ggml-cann.h.inc} | 4 +- ggml/include/ggml-cpp.h | 39 - ggml/include/{ggml-cpu.h => ggml-cpu.h.inc} | 4 +- ggml/include/{ggml-cuda.h => ggml-cuda.h.inc} | 4 +- .../{ggml-hexagon.h => ggml-hexagon.h.inc} | 4 +- .../{ggml-metal.h => ggml-metal.h.inc} | 4 +- .../{ggml-opencl.h => ggml-opencl.h.inc} | 4 +- .../{ggml-openvino.h => ggml-openvino.h.inc} | 2 +- ggml/include/{ggml-opt.h => ggml-opt.h.inc} | 4 +- ggml/include/{ggml-rpc.h => ggml-rpc.h.inc} | 2 +- ggml/include/{ggml-sycl.h => ggml-sycl.h.inc} | 4 +- .../{ggml-virtgpu.h => ggml-virtgpu.h.inc} | 4 +- .../{ggml-vulkan.h => ggml-vulkan.h.inc} | 4 +- .../{ggml-webgpu.h => ggml-webgpu.h.inc} | 4 +- ggml/include/{ggml-zdnn.h => ggml-zdnn.h.inc} | 4 +- .../{ggml-zendnn.h => ggml-zendnn.h.inc} | 4 +- ggml/include/{ggml.h => ggml.h.inc} | 4 +- ggml/include/{gguf.h => gguf.h.inc} | 2 +- ggml/src/CMakeLists.txt | 70 +- ggml/src/ggml-alloc.c | 1248 - ggml/src/ggml-backend-dl-rust/Cargo.toml | 7 + ggml/src/ggml-backend-dl-rust/src/lib.rs | 226 + ggml/src/ggml-backend-dl.cpp | 48 - ggml/src/ggml-backend-dl.h | 45 - ...kend-impl.h => ggml-backend-impl-defs.inc} | 2 +- ggml/src/ggml-backend-meta.cpp | 2143 - ...ckend-reg.cpp => ggml-backend-reg.cpp.inc} | 203 +- ggml/src/ggml-backend.cpp | 2371 - ggml/src/ggml-backend.cpp.inc | 7185 ++ ggml/src/ggml-blas/CMakeLists.txt | 3 +- .../{ggml-blas.cpp => ggml-blas.cpp.inc} | 6 +- ggml/src/ggml-cann/CMakeLists.txt | 3 +- ggml/src/ggml-cann/acl_tensor.cpp | 195 - ggml/src/ggml-cann/acl_tensor.h | 349 - ggml/src/ggml-cann/aclnn_ops.cpp | 4436 -- ggml/src/ggml-cann/aclnn_ops.h | 1190 - ggml/src/ggml-cann/common.h | 651 - .../{ggml-cann.cpp => ggml-cann.cpp.inc} | 2192 +- .../{ggml-common.h => ggml-common-defs.inc} | 0 ggml/src/ggml-cpu/CMakeLists.txt | 250 +- ggml/src/ggml-cpu/amx/amx.cpp | 249 - ggml/src/ggml-cpu/amx/amx.h | 8 - ggml/src/ggml-cpu/amx/common.h | 115 - ggml/src/ggml-cpu/amx/mmq.cpp | 2512 - ggml/src/ggml-cpu/amx/mmq.h | 10 - ggml/src/ggml-cpu/arch-fallback.h | 349 - .../arch/arm/cpu-feats-rust/Cargo.toml | 7 + .../arch/arm/cpu-feats-rust/src/lib.rs | 176 + ggml/src/ggml-cpu/arch/arm/cpu-feats.cpp | 98 - .../arch/arm/{quants.c => quants.c.inc} | 12 +- .../arch/arm/{repack.cpp => repack.cpp.inc} | 143 +- .../arch/loongarch/{quants.c => quants.c.inc} | 12 +- .../arch/powerpc/cpu-feats-rust/Cargo.toml | 7 + .../arch/powerpc/cpu-feats-rust/src/lib.rs | 148 + ggml/src/ggml-cpu/arch/powerpc/cpu-feats.cpp | 82 - .../arch/powerpc/{quants.c => quants.c.inc} | 12 +- .../arch/riscv/cpu-feats-rust/Cargo.toml | 7 + .../arch/riscv/cpu-feats-rust/src/lib.rs | 84 + ggml/src/ggml-cpu/arch/riscv/cpu-feats.cpp | 38 - .../arch/riscv/{quants.c => quants.c.inc} | 12 +- .../arch/riscv/{repack.cpp => repack.cpp.inc} | 143 +- .../arch/s390/cpu-feats-rust/Cargo.toml | 7 + .../arch/s390/cpu-feats-rust/src/lib.rs | 84 + ggml/src/ggml-cpu/arch/s390/cpu-feats.cpp | 50 - .../arch/s390/{quants.c => quants.c.inc} | 12 +- .../arch/wasm/{quants.c => quants.c.inc} | 12 +- .../arch/x86/cpu-feats-rust/Cargo.toml | 7 + .../arch/x86/cpu-feats-rust/src/lib.rs | 232 + ggml/src/ggml-cpu/arch/x86/cpu-feats.cpp | 327 - ggml/src/ggml-cpu/binary-ops.cpp | 154 - ggml/src/ggml-cpu/binary-ops.h | 16 - ggml/src/ggml-cpu/common.h | 95 - ...ggml-cpu-impl.h => ggml-cpu-impl-defs.inc} | 4 +- .../ggml-cpu/{ggml-cpu.c => ggml-cpu.c.inc} | 1986 +- ggml/src/ggml-cpu/ggml-cpu.cpp | 703 - ggml/src/ggml-cpu/ggml-cpu.cpp.inc | 8199 ++ ggml/src/ggml-cpu/hbm.cpp | 55 - ggml/src/ggml-cpu/hbm.h | 8 - .../kleidiai/{kernels.cpp => kernels.cpp.inc} | 96 +- ggml/src/ggml-cpu/kleidiai/kernels.h | 90 - .../{kleidiai.cpp => kleidiai.cpp.inc} | 236 +- ggml/src/ggml-cpu/kleidiai/kleidiai.h | 17 - ggml/src/ggml-cpu/llamafile/sgemm.cpp | 4051 - ggml/src/ggml-cpu/llamafile/sgemm.h | 25 - ggml/src/ggml-cpu/{ops.cpp => ops.cpp.inc} | 19090 +++-- ggml/src/ggml-cpu/ops.h | 117 - ggml/src/ggml-cpu/quants.c | 1288 - .../{arch/x86/quants.c => quants.c.inc} | 1732 +- ggml/src/ggml-cpu/quants.h | 103 - ggml/src/ggml-cpu/repack.cpp | 4836 -- .../{arch/x86/repack.cpp => repack.cpp.inc} | 5577 +- ggml/src/ggml-cpu/repack.h | 245 - ggml/src/ggml-cpu/simd-gemm.h | 226 - ...simd-mappings.h => simd-mappings-defs.inc} | 2 +- .../spacemit/{ime.cpp => ime.cpp.inc} | 159 +- ggml/src/ggml-cpu/spacemit/ime.h | 13 - ...{ime1_kernels.cpp => ime1_kernels.cpp.inc} | 24 +- ggml/src/ggml-cpu/spacemit/ime_kernels.h | 26 - ggml/src/ggml-cpu/traits.cpp | 36 - ggml/src/ggml-cpu/traits.h | 38 - ggml/src/ggml-cpu/unary-ops.cpp | 337 - ggml/src/ggml-cpu/unary-ops.h | 35 - ggml/src/ggml-cpu/vec.cpp | 629 - ggml/src/ggml-cpu/vec.h | 1588 - ggml/src/ggml-cuda/CMakeLists.txt | 26 +- ggml/src/ggml-cuda/{acc.cu => acc.cu.inc} | 2 +- ggml/src/ggml-cuda/{acc.cuh => acc.cuh.inc} | 2 +- .../ggml-cuda/{add-id.cu => add-id.cu.inc} | 2 +- .../ggml-cuda/{add-id.cuh => add-id.cuh.inc} | 2 +- .../ggml-cuda/{arange.cu => arange.cu.inc} | 2 +- .../ggml-cuda/{arange.cuh => arange.cuh.inc} | 2 +- .../ggml-cuda/{argmax.cu => argmax.cu.inc} | 6 +- .../ggml-cuda/{argmax.cuh => argmax.cuh.inc} | 2 +- .../ggml-cuda/{argsort.cu => argsort.cu.inc} | 2 +- .../{argsort.cuh => argsort.cuh.inc} | 2 +- .../{binbcast.cu => binbcast.cu.inc} | 2 +- .../{binbcast.cuh => binbcast.cuh.inc} | 2 +- ggml/src/ggml-cuda/{clamp.cu => clamp.cu.inc} | 2 +- .../ggml-cuda/{clamp.cuh => clamp.cuh.inc} | 2 +- .../ggml-cuda/{common.cuh => common.cuh.inc} | 16 +- .../ggml-cuda/{concat.cu => concat.cu.inc} | 2 +- .../ggml-cuda/{concat.cuh => concat.cuh.inc} | 2 +- ...anspose-1d.cu => conv-transpose-1d.cu.inc} | 2 +- ...spose-1d.cuh => conv-transpose-1d.cuh.inc} | 2 +- .../{conv2d-dw.cu => conv2d-dw.cu.inc} | 2 +- .../{conv2d-dw.cuh => conv2d-dw.cuh.inc} | 2 +- ...d-transpose.cu => conv2d-transpose.cu.inc} | 4 +- ...transpose.cuh => conv2d-transpose.cuh.inc} | 2 +- .../ggml-cuda/{conv2d.cu => conv2d.cu.inc} | 4 +- .../ggml-cuda/{conv2d.cuh => conv2d.cuh.inc} | 2 +- .../ggml-cuda/{convert.cu => convert.cu.inc} | 4 +- .../{convert.cuh => convert.cuh.inc} | 2 +- .../{count-equal.cu => count-equal.cu.inc} | 4 +- .../{count-equal.cuh => count-equal.cuh.inc} | 2 +- .../{cp-async.cuh => cp-async.cuh.inc} | 2 +- .../{cpy-utils.cuh => cpy-utils.cuh.inc} | 6 +- ggml/src/ggml-cuda/{cpy.cu => cpy.cu.inc} | 8 +- ggml/src/ggml-cuda/{cpy.cuh => cpy.cuh.inc} | 2 +- ...ropy-loss.cu => cross-entropy-loss.cu.inc} | 6 +- ...py-loss.cuh => cross-entropy-loss.cuh.inc} | 2 +- .../ggml-cuda/{cumsum.cu => cumsum.cu.inc} | 8 +- .../ggml-cuda/{cumsum.cuh => cumsum.cuh.inc} | 2 +- .../{dequantize.cuh => dequantize.cuh.inc} | 2 +- ggml/src/ggml-cuda/{diag.cu => diag.cu.inc} | 6 +- ggml/src/ggml-cuda/{diag.cuh => diag.cuh.inc} | 2 +- .../{diagmask.cu => diagmask.cu.inc} | 2 +- .../{diagmask.cuh => diagmask.cuh.inc} | 2 +- ...{fattn-common.cuh => fattn-common.cuh.inc} | 6 +- ...attn-mma-f16.cuh => fattn-mma-f16.cuh.inc} | 10 +- .../{fattn-tile.cu => fattn-tile.cu.inc} | 6 +- .../{fattn-tile.cuh => fattn-tile.cuh.inc} | 6 +- .../{fattn-vec.cuh => fattn-vec.cuh.inc} | 4 +- ...attn-wmma-f16.cu => fattn-wmma-f16.cu.inc} | 6 +- ...tn-wmma-f16.cuh => fattn-wmma-f16.cuh.inc} | 2 +- ggml/src/ggml-cuda/{fattn.cu => fattn.cu.inc} | 14 +- .../ggml-cuda/{fattn.cuh => fattn.cuh.inc} | 2 +- ggml/src/ggml-cuda/{fill.cu => fill.cu.inc} | 4 +- ggml/src/ggml-cuda/{fill.cuh => fill.cuh.inc} | 2 +- ...ed_delta_net.cu => gated_delta_net.cu.inc} | 2 +- ..._delta_net.cuh => gated_delta_net.cuh.inc} | 4 +- .../ggml-cuda/{getrows.cu => getrows.cu.inc} | 6 +- .../{getrows.cuh => getrows.cuh.inc} | 2 +- .../{ggml-cuda.cu => ggml-cuda.cu.inc} | 128 +- ggml/src/ggml-cuda/{gla.cu => gla.cu.inc} | 4 +- ggml/src/ggml-cuda/{gla.cuh => gla.cuh.inc} | 2 +- .../ggml-cuda/{im2col.cu => im2col.cu.inc} | 2 +- .../ggml-cuda/{im2col.cuh => im2col.cuh.inc} | 2 +- ggml/src/ggml-cuda/{mean.cu => mean.cu.inc} | 4 +- ggml/src/ggml-cuda/{mean.cuh => mean.cuh.inc} | 2 +- ggml/src/ggml-cuda/{mma.cuh => mma.cuh.inc} | 2 +- ggml/src/ggml-cuda/{mmf.cu => mmf.cu.inc} | 6 +- ggml/src/ggml-cuda/{mmf.cuh => mmf.cuh.inc} | 6 +- ggml/src/ggml-cuda/{mmid.cu => mmid.cu.inc} | 4 +- ggml/src/ggml-cuda/{mmid.cuh => mmid.cuh.inc} | 0 ggml/src/ggml-cuda/{mmq.cu => mmq.cu.inc} | 8 +- ggml/src/ggml-cuda/{mmq.cuh => mmq.cuh.inc} | 6 +- ggml/src/ggml-cuda/{mmvf.cu => mmvf.cu.inc} | 10 +- ggml/src/ggml-cuda/{mmvf.cuh => mmvf.cuh.inc} | 2 +- ggml/src/ggml-cuda/{mmvq.cu => mmvq.cu.inc} | 8 +- ggml/src/ggml-cuda/{mmvq.cuh => mmvq.cuh.inc} | 2 +- ggml/src/ggml-cuda/{norm.cu => norm.cu.inc} | 2 +- ggml/src/ggml-cuda/{norm.cuh => norm.cuh.inc} | 2 +- ...pt-step-adamw.cu => opt-step-adamw.cu.inc} | 4 +- ...-step-adamw.cuh => opt-step-adamw.cuh.inc} | 2 +- .../{opt-step-sgd.cu => opt-step-sgd.cu.inc} | 4 +- ...{opt-step-sgd.cuh => opt-step-sgd.cuh.inc} | 2 +- .../{out-prod.cu => out-prod.cu.inc} | 2 +- .../{out-prod.cuh => out-prod.cuh.inc} | 2 +- ggml/src/ggml-cuda/{pad.cu => pad.cu.inc} | 2 +- ggml/src/ggml-cuda/{pad.cuh => pad.cuh.inc} | 2 +- ...ad_reflect_1d.cu => pad_reflect_1d.cu.inc} | 4 +- ..._reflect_1d.cuh => pad_reflect_1d.cuh.inc} | 2 +- .../ggml-cuda/{pool2d.cu => pool2d.cu.inc} | 2 +- .../ggml-cuda/{pool2d.cuh => pool2d.cuh.inc} | 2 +- .../{quantize.cu => quantize.cu.inc} | 2 +- .../{quantize.cuh => quantize.cuh.inc} | 4 +- .../{reduce_rows.cuh => reduce_rows.cuh.inc} | 2 +- ggml/src/ggml-cuda/{roll.cu => roll.cu.inc} | 4 +- ggml/src/ggml-cuda/{roll.cuh => roll.cuh.inc} | 2 +- ggml/src/ggml-cuda/{rope.cu => rope.cu.inc} | 8 +- ggml/src/ggml-cuda/{rope.cuh => rope.cuh.inc} | 2 +- ggml/src/ggml-cuda/{scale.cu => scale.cu.inc} | 2 +- .../ggml-cuda/{scale.cuh => scale.cuh.inc} | 2 +- .../{set-rows.cu => set-rows.cu.inc} | 4 +- .../{set-rows.cuh => set-rows.cuh.inc} | 2 +- ggml/src/ggml-cuda/{set.cu => set.cu.inc} | 4 +- ggml/src/ggml-cuda/{set.cuh => set.cuh.inc} | 2 +- .../ggml-cuda/{softcap.cu => softcap.cu.inc} | 2 +- .../{softcap.cuh => softcap.cuh.inc} | 2 +- .../ggml-cuda/{softmax.cu => softmax.cu.inc} | 6 +- .../{softmax.cuh => softmax.cuh.inc} | 2 +- .../{solve_tri.cu => solve_tri.cu.inc} | 8 +- .../{solve_tri.cuh => solve_tri.cuh.inc} | 2 +- .../{ssm-conv.cu => ssm-conv.cu.inc} | 4 +- .../{ssm-conv.cuh => ssm-conv.cuh.inc} | 2 +- .../{ssm-scan.cu => ssm-scan.cu.inc} | 2 +- .../{ssm-scan.cuh => ssm-scan.cuh.inc} | 2 +- ggml/src/ggml-cuda/{sum.cu => sum.cu.inc} | 4 +- ggml/src/ggml-cuda/{sum.cuh => sum.cuh.inc} | 2 +- .../ggml-cuda/{sumrows.cu => sumrows.cu.inc} | 4 +- .../{sumrows.cuh => sumrows.cuh.inc} | 2 +- ...ma-f16-instance-ncols1_1-ncols2_16.cu.inc} | 2 +- ...ma-f16-instance-ncols1_1-ncols2_32.cu.inc} | 2 +- ...mma-f16-instance-ncols1_1-ncols2_8.cu.inc} | 2 +- ...ma-f16-instance-ncols1_16-ncols2_1.cu.inc} | 2 +- ...ma-f16-instance-ncols1_16-ncols2_2.cu.inc} | 2 +- ...ma-f16-instance-ncols1_16-ncols2_4.cu.inc} | 2 +- ...ma-f16-instance-ncols1_2-ncols2_16.cu.inc} | 2 +- ...ma-f16-instance-ncols1_2-ncols2_32.cu.inc} | 2 +- ...mma-f16-instance-ncols1_2-ncols2_4.cu.inc} | 2 +- ...mma-f16-instance-ncols1_2-ncols2_8.cu.inc} | 2 +- ...ma-f16-instance-ncols1_32-ncols2_1.cu.inc} | 2 +- ...ma-f16-instance-ncols1_32-ncols2_2.cu.inc} | 2 +- ...ma-f16-instance-ncols1_4-ncols2_16.cu.inc} | 2 +- ...mma-f16-instance-ncols1_4-ncols2_2.cu.inc} | 2 +- ...mma-f16-instance-ncols1_4-ncols2_4.cu.inc} | 2 +- ...mma-f16-instance-ncols1_4-ncols2_8.cu.inc} | 2 +- ...ma-f16-instance-ncols1_64-ncols2_1.cu.inc} | 2 +- ...mma-f16-instance-ncols1_8-ncols2_1.cu.inc} | 2 +- ...mma-f16-instance-ncols1_8-ncols2_2.cu.inc} | 2 +- ...mma-f16-instance-ncols1_8-ncols2_4.cu.inc} | 2 +- ...mma-f16-instance-ncols1_8-ncols2_8.cu.inc} | 2 +- .../fattn-tile-instance-dkq112-dv112.cu | 5 - .../fattn-tile-instance-dkq112-dv112.cu.inc} | 2 +- ...> fattn-tile-instance-dkq128-dv128.cu.inc} | 2 +- .../fattn-tile-instance-dkq256-dv256.cu | 5 - .../fattn-tile-instance-dkq256-dv256.cu.inc} | 2 +- ...> fattn-tile-instance-dkq320-dv256.cu.inc} | 2 +- ... => fattn-tile-instance-dkq40-dv40.cu.inc} | 2 +- ...> fattn-tile-instance-dkq512-dv512.cu.inc} | 2 +- ...> fattn-tile-instance-dkq576-dv512.cu.inc} | 2 +- ... => fattn-tile-instance-dkq64-dv64.cu.inc} | 2 +- ... => fattn-tile-instance-dkq72-dv72.cu.inc} | 2 +- ... => fattn-tile-instance-dkq80-dv80.cu.inc} | 2 +- ... => fattn-tile-instance-dkq96-dv96.cu.inc} | 2 +- ...cu => fattn-vec-instance-bf16-bf16.cu.inc} | 2 +- ....cu => fattn-vec-instance-bf16-f16.cu.inc} | 2 +- ...cu => fattn-vec-instance-bf16-q4_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-bf16-q4_1.cu.inc} | 2 +- ...cu => fattn-vec-instance-bf16-q5_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-bf16-q5_1.cu.inc} | 2 +- ...cu => fattn-vec-instance-bf16-q8_0.cu.inc} | 2 +- ....cu => fattn-vec-instance-f16-bf16.cu.inc} | 2 +- ...6.cu => fattn-vec-instance-f16-f16.cu.inc} | 2 +- ....cu => fattn-vec-instance-f16-q4_0.cu.inc} | 2 +- ....cu => fattn-vec-instance-f16-q4_1.cu.inc} | 2 +- ....cu => fattn-vec-instance-f16-q5_0.cu.inc} | 2 +- ....cu => fattn-vec-instance-f16-q5_1.cu.inc} | 2 +- ....cu => fattn-vec-instance-f16-q8_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q4_0-bf16.cu.inc} | 2 +- ....cu => fattn-vec-instance-q4_0-f16.cu.inc} | 2 +- ...cu => fattn-vec-instance-q4_0-q4_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q4_0-q4_1.cu.inc} | 2 +- ...cu => fattn-vec-instance-q4_0-q5_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q4_0-q5_1.cu.inc} | 2 +- ...cu => fattn-vec-instance-q4_0-q8_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q4_1-bf16.cu.inc} | 2 +- ....cu => fattn-vec-instance-q4_1-f16.cu.inc} | 2 +- ...cu => fattn-vec-instance-q4_1-q4_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q4_1-q4_1.cu.inc} | 2 +- ...cu => fattn-vec-instance-q4_1-q5_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q4_1-q5_1.cu.inc} | 2 +- ...cu => fattn-vec-instance-q4_1-q8_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q5_0-bf16.cu.inc} | 2 +- ....cu => fattn-vec-instance-q5_0-f16.cu.inc} | 2 +- ...cu => fattn-vec-instance-q5_0-q4_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q5_0-q4_1.cu.inc} | 2 +- ...cu => fattn-vec-instance-q5_0-q5_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q5_0-q5_1.cu.inc} | 2 +- ...cu => fattn-vec-instance-q5_0-q8_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q5_1-bf16.cu.inc} | 2 +- ....cu => fattn-vec-instance-q5_1-f16.cu.inc} | 2 +- ...cu => fattn-vec-instance-q5_1-q4_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q5_1-q4_1.cu.inc} | 2 +- ...cu => fattn-vec-instance-q5_1-q5_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q5_1-q5_1.cu.inc} | 2 +- ...cu => fattn-vec-instance-q5_1-q8_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q8_0-bf16.cu.inc} | 2 +- ....cu => fattn-vec-instance-q8_0-f16.cu.inc} | 2 +- ...cu => fattn-vec-instance-q8_0-q4_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q8_0-q4_1.cu.inc} | 2 +- ...cu => fattn-vec-instance-q8_0-q5_0.cu.inc} | 2 +- ...cu => fattn-vec-instance-q8_0-q5_1.cu.inc} | 2 +- ...cu => fattn-vec-instance-q8_0-q8_0.cu.inc} | 2 +- ...ncols_1.cu => mmf-instance-ncols_1.cu.inc} | 2 +- ...ols_10.cu => mmf-instance-ncols_10.cu.inc} | 2 +- ...ols_11.cu => mmf-instance-ncols_11.cu.inc} | 2 +- ...ols_12.cu => mmf-instance-ncols_12.cu.inc} | 2 +- ...ols_13.cu => mmf-instance-ncols_13.cu.inc} | 2 +- ...ols_14.cu => mmf-instance-ncols_14.cu.inc} | 2 +- ...ols_15.cu => mmf-instance-ncols_15.cu.inc} | 2 +- ...ols_16.cu => mmf-instance-ncols_16.cu.inc} | 2 +- ...ncols_2.cu => mmf-instance-ncols_2.cu.inc} | 2 +- ...ncols_3.cu => mmf-instance-ncols_3.cu.inc} | 2 +- ...ncols_4.cu => mmf-instance-ncols_4.cu.inc} | 2 +- ...ncols_5.cu => mmf-instance-ncols_5.cu.inc} | 2 +- ...ncols_6.cu => mmf-instance-ncols_6.cu.inc} | 2 +- ...ncols_7.cu => mmf-instance-ncols_7.cu.inc} | 2 +- ...ncols_8.cu => mmf-instance-ncols_8.cu.inc} | 2 +- ...ncols_9.cu => mmf-instance-ncols_9.cu.inc} | 2 +- ...nce-iq1_s.cu => mmq-instance-iq1_s.cu.inc} | 2 +- ...nce-iq2_s.cu => mmq-instance-iq2_s.cu.inc} | 2 +- ...e-iq2_xs.cu => mmq-instance-iq2_xs.cu.inc} | 2 +- ...iq2_xxs.cu => mmq-instance-iq2_xxs.cu.inc} | 2 +- ...nce-iq3_s.cu => mmq-instance-iq3_s.cu.inc} | 2 +- ...iq3_xxs.cu => mmq-instance-iq3_xxs.cu.inc} | 2 +- ...e-iq4_nl.cu => mmq-instance-iq4_nl.cu.inc} | 2 +- ...e-iq4_xs.cu => mmq-instance-iq4_xs.cu.inc} | 2 +- ...nce-mxfp4.cu => mmq-instance-mxfp4.cu.inc} | 2 +- ...nce-nvfp4.cu => mmq-instance-nvfp4.cu.inc} | 2 +- ...tance-q1_0.cu => mmq-instance-q1_0.cu.inc} | 2 +- ...tance-q2_k.cu => mmq-instance-q2_k.cu.inc} | 2 +- ...tance-q3_k.cu => mmq-instance-q3_k.cu.inc} | 2 +- ...tance-q4_0.cu => mmq-instance-q4_0.cu.inc} | 2 +- ...tance-q4_1.cu => mmq-instance-q4_1.cu.inc} | 2 +- ...tance-q4_k.cu => mmq-instance-q4_k.cu.inc} | 2 +- ...tance-q5_0.cu => mmq-instance-q5_0.cu.inc} | 2 +- ...tance-q5_1.cu => mmq-instance-q5_1.cu.inc} | 2 +- ...tance-q5_k.cu => mmq-instance-q5_k.cu.inc} | 2 +- ...tance-q6_k.cu => mmq-instance-q6_k.cu.inc} | 2 +- ...tance-q8_0.cu => mmq-instance-q8_0.cu.inc} | 2 +- ggml/src/ggml-cuda/{top-k.cu => top-k.cu.inc} | 4 +- .../ggml-cuda/{top-k.cuh => top-k.cuh.inc} | 2 +- .../{topk-moe.cu => topk-moe.cu.inc} | 6 +- .../{topk-moe.cuh => topk-moe.cuh.inc} | 4 +- ggml/src/ggml-cuda/{tri.cu => tri.cu.inc} | 8 +- ggml/src/ggml-cuda/{tri.cuh => tri.cuh.inc} | 2 +- .../ggml-cuda/{tsembd.cu => tsembd.cu.inc} | 2 +- .../ggml-cuda/{tsembd.cuh => tsembd.cuh.inc} | 2 +- ggml/src/ggml-cuda/{unary.cu => unary.cu.inc} | 4 +- .../ggml-cuda/{unary.cuh => unary.cuh.inc} | 2 +- .../ggml-cuda/{upscale.cu => upscale.cu.inc} | 2 +- .../{upscale.cuh => upscale.cuh.inc} | 2 +- .../{vecdotq.cuh => vecdotq.cuh.inc} | 2 +- .../ggml-cuda/vendors/{cuda.h => cuda.h.inc} | 2 +- .../ggml-cuda/vendors/{hip.h => hip.h.inc} | 0 .../ggml-cuda/vendors/{musa.h => musa.h.inc} | 2 +- ggml/src/ggml-cuda/{wkv.cu => wkv.cu.inc} | 4 +- ggml/src/ggml-cuda/{wkv.cuh => wkv.cuh.inc} | 2 +- ggml/src/ggml-hexagon/CMakeLists.txt | 9 +- ...{ggml-hexagon.cpp => ggml-hexagon.cpp.inc} | 284 +- .../{htp-drv.cpp => htp-drv.cpp.inc} | 202 +- ggml/src/ggml-hexagon/htp-drv.h | 121 - ggml/src/ggml-hexagon/htp/CMakeLists.txt | 60 +- .../htp/{act-ops.c => act-ops.inc} | 12 +- .../htp/{argsort-ops.c => argsort-ops.inc} | 14 +- .../htp/{binary-ops.c => binary-ops.inc} | 12 +- .../htp/{cpy-ops.c => cpy-ops.inc} | 10 +- .../htp/{cumsum-ops.c => cumsum-ops.inc} | 12 +- ggml/src/ggml-hexagon/htp/diag-ops.c | 216 - .../htp/{hvx-copy.h => diag-ops.inc} | 219 +- .../htp/{fill-ops.c => fill-ops.inc} | 9 +- .../{flash-attn-ops.c => flash-attn-ops.inc} | 211 +- .../htp/{get-rows-ops.c => get-rows-ops.inc} | 10 +- .../htp/{hex-dma.h => hex-dma-defs.inc} | 0 .../htp/{hex-dma.c => hex-dma.inc} | 2 +- ggml/src/ggml-hexagon/htp/hex-dump.h | 86 - ggml/src/ggml-hexagon/htp/hex-fastdiv.h | 37 - .../htp/{hex-utils.h => hex-utils-defs.inc} | 119 +- ...lash-attn-ops.c => hmx-flash-attn-ops.inc} | 679 +- .../{hmx-matmul-ops.c => hmx-matmul-ops.inc} | 459 +- ggml/src/ggml-hexagon/htp/hmx-ops.h | 71 - ggml/src/ggml-hexagon/htp/hmx-profile.h | 34 - .../htp/{hmx-queue.h => hmx-queue-defs.inc} | 2 +- .../htp/{hmx-queue.c => hmx-queue.inc} | 2 +- ggml/src/ggml-hexagon/htp/hmx-utils.h | 202 - .../htp/{htp-ctx.h => htp-ctx-defs.inc} | 8 +- .../htp/{htp-ops.h => htp-ops-defs.inc} | 6 +- ggml/src/ggml-hexagon/htp/hvx-arith.h | 443 - .../htp/{hvx-base.h => hvx-base-defs.inc} | 4 +- ggml/src/ggml-hexagon/htp/hvx-div.h | 291 - ggml/src/ggml-hexagon/htp/hvx-dump.h | 129 - ggml/src/ggml-hexagon/htp/hvx-exp.h | 216 - ggml/src/ggml-hexagon/htp/hvx-floor.h | 100 - ggml/src/ggml-hexagon/htp/hvx-inverse.h | 210 - ggml/src/ggml-hexagon/htp/hvx-reduce.h | 296 - ggml/src/ggml-hexagon/htp/hvx-scale.h | 133 - ggml/src/ggml-hexagon/htp/hvx-sigmoid.h | 142 - ggml/src/ggml-hexagon/htp/hvx-sqrt.h | 126 - .../htp/{hvx-types.h => hvx-types-defs.inc} | 0 ggml/src/ggml-hexagon/htp/hvx-utils-defs.inc | 2188 + ggml/src/ggml-hexagon/htp/hvx-utils.h | 18 - .../src/ggml-hexagon/htp/{main.c => main.inc} | 23 +- .../htp/{matmul-ops.c => matmul-ops.inc} | 211 +- .../htp/{repeat-ops.c => repeat-ops.inc} | 10 +- .../htp/{rope-ops.c => rope-ops.inc} | 49 +- .../htp/{set-rows-ops.c => set-rows-ops.inc} | 12 +- .../htp/{softmax-ops.c => softmax-ops.inc} | 47 +- .../{solve-tri-ops.c => solve-tri-ops.inc} | 10 +- .../htp/{ssm-conv.c => ssm-conv.inc} | 12 +- .../htp/{sum-rows-ops.c => sum-rows-ops.inc} | 12 +- ggml/src/ggml-hexagon/htp/unary-ops.c | 556 - ggml/src/ggml-hexagon/htp/unary-ops.inc | 1205 + ggml/src/ggml-hexagon/htp/vtcm-utils.h | 16 - .../{worker-pool.h => worker-pool-defs.inc} | 0 .../htp/{worker-pool.c => worker-pool.inc} | 2 +- ggml/src/ggml-hexagon/libdl.h | 79 - ggml/src/ggml-hexagon/op-desc.h | 153 - ggml/src/ggml-hip/CMakeLists.txt | 24 +- ggml/src/{ggml-impl.h => ggml-impl-defs.inc} | 33 +- ggml/src/ggml-metal/CMakeLists.txt | 31 +- ...al-common.h => ggml-metal-common-defs.inc} | 0 ...metal-common.cpp => ggml-metal-common.inc} | 6 +- ...-context.h => ggml-metal-context-defs.inc} | 2 +- ggml/src/ggml-metal/ggml-metal-context.m | 12 +- ...al-device.h => ggml-metal-device-defs.inc} | 2 +- ...metal-device.cpp => ggml-metal-device.inc} | 6 +- ggml/src/ggml-metal/ggml-metal-device.m | 4 +- ...-metal-impl.h => ggml-metal-impl-defs.inc} | 0 ...ml-metal-ops.h => ggml-metal-ops-defs.inc} | 2 +- ...{ggml-metal-ops.cpp => ggml-metal-ops.inc} | 14 +- .../{ggml-metal.cpp => ggml-metal.inc} | 12 +- ggml/src/ggml-metal/ggml-metal.metal | 10 +- ggml/src/ggml-musa/CMakeLists.txt | 26 +- ggml/src/ggml-musa/{mudnn.cu => mudnn.cu.inc} | 2 +- .../ggml-musa/{mudnn.cuh => mudnn.cuh.inc} | 4 +- ggml/src/ggml-opencl/CMakeLists.txt | 5 +- .../{ggml-opencl.cpp => ggml-opencl.cpp.inc} | 10 +- ggml/src/ggml-openvino/CMakeLists.txt | 5 +- ggml/src/ggml-openvino/ggml-decoder.cpp | 985 - ggml/src/ggml-openvino/ggml-decoder.h | 294 - .../src/ggml-openvino/ggml-openvino-extra.cpp | 380 - ggml/src/ggml-openvino/ggml-openvino-extra.h | 182 - ggml/src/ggml-openvino/ggml-openvino.cpp | 1132 - ggml/src/ggml-openvino/ggml-openvino.cpp.inc | 12186 +++ ggml/src/ggml-openvino/ggml-quants.cpp | 956 - ggml/src/ggml-openvino/ggml-quants.h | 153 - ggml/src/ggml-openvino/openvino/decoder.h | 74 - ggml/src/ggml-openvino/openvino/frontend.cpp | 27 - ggml/src/ggml-openvino/openvino/frontend.h | 23 - .../ggml-openvino/openvino/input_model.cpp | 17 - ggml/src/ggml-openvino/openvino/input_model.h | 29 - .../src/ggml-openvino/openvino/node_context.h | 112 - ggml/src/ggml-openvino/openvino/op/cont.cpp | 48 - ggml/src/ggml-openvino/openvino/op/cpy.cpp | 21 - .../openvino/op/flash_attn_ext.cpp | 90 - .../ggml-openvino/openvino/op/get_rows.cpp | 69 - .../ggml-openvino/openvino/op/glu_geglu.cpp | 61 - .../ggml-openvino/openvino/op/glu_swiglu.cpp | 62 - ggml/src/ggml-openvino/openvino/op/mulmat.cpp | 90 - .../src/ggml-openvino/openvino/op/permute.cpp | 102 - .../src/ggml-openvino/openvino/op/reshape.cpp | 83 - .../ggml-openvino/openvino/op/rms_norm.cpp | 46 - ggml/src/ggml-openvino/openvino/op/rope.cpp | 149 - ggml/src/ggml-openvino/openvino/op/scale.cpp | 41 - .../ggml-openvino/openvino/op/set_rows.cpp | 76 - .../src/ggml-openvino/openvino/op/softmax.cpp | 89 - .../ggml-openvino/openvino/op/transpose.cpp | 23 - .../ggml-openvino/openvino/op/unary_gelu.cpp | 25 - .../ggml-openvino/openvino/op/unary_silu.cpp | 27 - ggml/src/ggml-openvino/openvino/op/view.cpp | 53 - ggml/src/ggml-openvino/openvino/op_table.cpp | 47 - ggml/src/ggml-openvino/openvino/op_table.h | 40 - .../openvino/pass/fuse_to_sdpa.cpp | 60 - .../openvino/pass/fuse_to_sdpa.h | 17 - ...k_decompression_convert_constant_folding.h | 29 - .../openvino/pass/squeeze_matmul.cpp | 58 - .../openvino/pass/squeeze_matmul.h | 17 - .../rt_info/weightless_caching_attributes.hpp | 41 - .../openvino/translate_session.cpp | 317 - .../openvino/translate_session.h | 28 - ggml/src/ggml-openvino/openvino/utils.cpp | 257 - ggml/src/ggml-openvino/openvino/utils.h | 86 - ggml/src/ggml-openvino/utils.cpp | 880 - ggml/src/ggml-openvino/utils.h | 143 - ggml/src/ggml-opt.cpp | 1094 - .../{ggml-quants.h => ggml-quants-defs.inc} | 4 +- ggml/src/ggml-quants.c | 5491 -- ggml/src/ggml-rpc/CMakeLists.txt | 4 +- .../{ggml-rpc.cpp => ggml-rpc.cpp.inc} | 727 +- ggml/src/ggml-rpc/transport.cpp | 683 - ggml/src/ggml-rpc/transport.h | 34 - ggml/src/ggml-sycl/CMakeLists.txt | 55 +- ggml/src/ggml-sycl/add-id.hpp | 8 - ggml/src/ggml-sycl/{add-id.cpp => add-id.inc} | 3 +- ggml/src/ggml-sycl/backend-defs.inc | 48 + ggml/src/ggml-sycl/backend.hpp | 48 - ggml/src/ggml-sycl/binbcast.hpp | 39 - .../ggml-sycl/{binbcast.cpp => binbcast.inc} | 26 +- .../ggml-sycl/{common.hpp => common-defs.inc} | 37 +- ggml/src/ggml-sycl/{common.cpp => common.inc} | 6 +- ggml/src/ggml-sycl/concat.hpp | 20 - ggml/src/ggml-sycl/{concat.cpp => concat.inc} | 2 +- ggml/src/ggml-sycl/conv.hpp | 20 - ggml/src/ggml-sycl/{conv.cpp => conv.inc} | 3 +- .../{convert.hpp => convert-defs.inc} | 2 +- .../ggml-sycl/{convert.cpp => convert.inc} | 6 +- ggml/src/ggml-sycl/count-equal.hpp | 9 - .../{count-equal.cpp => count-equal.inc} | 4 +- ggml/src/ggml-sycl/{cpy.hpp => cpy-defs.inc} | 2 +- ggml/src/ggml-sycl/{cpy.cpp => cpy.inc} | 10 +- .../{dequantize.hpp => dequantize-defs.inc} | 4 +- ggml/src/ggml-sycl/dmmv.hpp | 27 - ggml/src/ggml-sycl/{dmmv.cpp => dmmv.inc} | 8 +- .../dpct/{helper.hpp => helper-defs.inc} | 2 +- ...element_wise.hpp => element_wise-defs.inc} | 4 +- .../{element_wise.cpp => element_wise.inc} | 8 +- ...fattn-common.hpp => fattn-common-defs.inc} | 10 +- .../{fattn-tile.hpp => fattn-tile-defs.inc} | 6 +- .../{fattn-tile.cpp => fattn-tile.inc} | 8 +- .../{fattn-vec.hpp => fattn-vec-defs.inc} | 8 +- ggml/src/ggml-sycl/fattn.hpp | 22 - ggml/src/ggml-sycl/{fattn.cpp => fattn.inc} | 12 +- ggml/src/ggml-sycl/gated_delta_net.hpp | 8 - ...ated_delta_net.cpp => gated_delta_net.inc} | 7 +- .../src/ggml-sycl/{gemm.hpp => gemm-defs.inc} | 2 +- ggml/src/ggml-sycl/getrows.hpp | 20 - .../ggml-sycl/{getrows.cpp => getrows.inc} | 8 +- .../{ggml-sycl.cpp => ggml-sycl.inc} | 30 +- ggml/src/ggml-sycl/gla.hpp | 8 - ggml/src/ggml-sycl/{gla.cpp => gla.inc} | 2 +- ggml/src/ggml-sycl/im2col.hpp | 21 - ggml/src/ggml-sycl/{im2col.cpp => im2col.inc} | 4 +- ggml/src/ggml-sycl/mmq.hpp | 33 - ggml/src/ggml-sycl/{mmq.cpp => mmq.inc} | 4 +- ggml/src/ggml-sycl/mmvq.hpp | 43 - ggml/src/ggml-sycl/{mmvq.cpp => mmvq.inc} | 10 +- ggml/src/ggml-sycl/norm.hpp | 28 - ggml/src/ggml-sycl/{norm.cpp => norm.inc} | 6 +- ggml/src/ggml-sycl/outprod.hpp | 10 - .../ggml-sycl/{outprod.cpp => outprod.inc} | 2 +- ggml/src/ggml-sycl/pad.hpp | 24 - ggml/src/ggml-sycl/{pad.cpp => pad.inc} | 4 +- ggml/src/ggml-sycl/pad_reflect_1d.hpp | 10 - ...{pad_reflect_1d.cpp => pad_reflect_1d.inc} | 4 +- .../{presets.hpp => presets-defs.inc} | 0 .../{quantize.hpp => quantize-defs.inc} | 2 +- .../ggml-sycl/{quants.hpp => quants-defs.inc} | 6 +- ggml/src/ggml-sycl/repeat_back.hpp | 8 - .../{repeat_back.cpp => repeat_back.inc} | 4 +- ggml/src/ggml-sycl/roll.hpp | 20 - ggml/src/ggml-sycl/{roll.cpp => roll.inc} | 4 +- ggml/src/ggml-sycl/rope.hpp | 26 - ggml/src/ggml-sycl/{rope.cpp => rope.inc} | 8 +- ggml/src/ggml-sycl/set.hpp | 5 - ggml/src/ggml-sycl/{set.cpp => set.inc} | 7 +- ggml/src/ggml-sycl/set_rows.hpp | 8 - .../ggml-sycl/{set_rows.cpp => set_rows.inc} | 4 +- ggml/src/ggml-sycl/softmax.hpp | 24 - .../ggml-sycl/{softmax.cpp => softmax.inc} | 4 +- ggml/src/ggml-sycl/ssm_conv.hpp | 5 - .../ggml-sycl/{ssm_conv.cpp => ssm_conv.inc} | 3 +- ggml/src/ggml-sycl/sycl_hw.hpp | 38 - .../ggml-sycl/{sycl_hw.cpp => sycl_hw.inc} | 2 +- .../fattn-tile-instance-dkq128-dv128.cpp | 5 - .../fattn-tile-instance-dkq40-dv40.cpp | 5 - .../fattn-tile-instance-dkq512-dv512.cpp | 6 - .../fattn-tile-instance-dkq576-dv512.cpp | 5 - .../fattn-tile-instance-dkq64-dv64.cpp | 5 - .../fattn-tile-instance-dkq72-dv72.cpp | 5 - .../fattn-tile-instance-dkq80-dv80.cpp | 5 - .../fattn-tile-instance-dkq96-dv96.cpp | 5 - .../fattn-tile-instances.inc | 14 + .../fattn-vec-instance-f16-f16.cpp | 8 - .../fattn-vec-instance-f16-q4_0.cpp | 8 - .../fattn-vec-instance-f16-q4_1.cpp | 8 - .../fattn-vec-instance-f16-q5_0.cpp | 8 - .../fattn-vec-instance-f16-q5_1.cpp | 8 - .../fattn-vec-instance-f16-q8_0.cpp | 8 - .../fattn-vec-instance-q4_0-f16.cpp | 8 - .../fattn-vec-instance-q4_0-q4_0.cpp | 8 - .../fattn-vec-instance-q4_0-q4_1.cpp | 8 - .../fattn-vec-instance-q4_0-q5_0.cpp | 8 - .../fattn-vec-instance-q4_0-q5_1.cpp | 8 - .../fattn-vec-instance-q4_0-q8_0.cpp | 8 - .../fattn-vec-instance-q4_1-f16.cpp | 8 - .../fattn-vec-instance-q4_1-q4_0.cpp | 8 - .../fattn-vec-instance-q4_1-q4_1.cpp | 8 - .../fattn-vec-instance-q4_1-q5_0.cpp | 8 - .../fattn-vec-instance-q4_1-q5_1.cpp | 8 - .../fattn-vec-instance-q4_1-q8_0.cpp | 8 - .../fattn-vec-instance-q5_0-f16.cpp | 8 - .../fattn-vec-instance-q5_0-q4_0.cpp | 8 - .../fattn-vec-instance-q5_0-q4_1.cpp | 8 - .../fattn-vec-instance-q5_0-q5_0.cpp | 8 - .../fattn-vec-instance-q5_0-q5_1.cpp | 8 - .../fattn-vec-instance-q5_0-q8_0.cpp | 8 - .../fattn-vec-instance-q5_1-f16.cpp | 8 - .../fattn-vec-instance-q5_1-q4_0.cpp | 8 - .../fattn-vec-instance-q5_1-q4_1.cpp | 8 - .../fattn-vec-instance-q5_1-q5_0.cpp | 8 - .../fattn-vec-instance-q5_1-q5_1.cpp | 8 - .../fattn-vec-instance-q5_1-q8_0.cpp | 8 - .../fattn-vec-instance-q8_0-f16.cpp | 8 - .../fattn-vec-instance-q8_0-q4_0.cpp | 8 - .../fattn-vec-instance-q8_0-q4_1.cpp | 8 - .../fattn-vec-instance-q8_0-q5_0.cpp | 8 - .../fattn-vec-instance-q8_0-q5_1.cpp | 8 - .../fattn-vec-instance-q8_0-q8_0.cpp | 8 - .../fattn-vec-instances.inc | 53 + ggml/src/ggml-sycl/tsembd.hpp | 20 - ggml/src/ggml-sycl/{tsembd.cpp => tsembd.inc} | 2 +- .../src/ggml-sycl/{type.hpp => type-defs.inc} | 0 ggml/src/ggml-sycl/upscale.hpp | 9 - .../ggml-sycl/{upscale.cpp => upscale.inc} | 2 +- .../{vecdotq.hpp => vecdotq-defs.inc} | 8 +- ggml/src/ggml-sycl/wkv.hpp | 10 - ggml/src/ggml-sycl/{wkv.cpp => wkv.inc} | 2 +- ggml/src/ggml-threading-rust/Cargo.toml | 7 + ggml/src/ggml-threading-rust/src/lib.rs | 150 + ggml/src/ggml-threading.cpp | 12 - ggml/src/ggml-threading.h | 14 - ggml/src/ggml-virtgpu/CMakeLists.txt | 20 +- .../ggml-virtgpu/apir_cs_ggml-rpc-front.cpp | 87 - ggml/src/ggml-virtgpu/backend/CMakeLists.txt | 32 +- .../backend/apir_cs_ggml-rpc-back.cpp | 115 - .../ggml-virtgpu/backend/backend-convert.h | 13 - .../backend/backend-dispatched-backend.cpp | 102 - .../backend-dispatched-buffer-type.cpp | 105 - .../backend/backend-dispatched-buffer.cpp | 179 - .../backend/backend-dispatched-device.cpp | 148 - .../backend-dispatched-rust/Cargo.toml | 7 + .../backend-dispatched-rust/src/lib.rs | 173 + .../backend/backend-dispatched.cpp | 51 - .../backend/backend-dispatched.gen.h | 73 - .../ggml-virtgpu/backend/backend-dispatched.h | 27 - .../ggml-virtgpu/backend/backend-virgl-apir.h | 32 - ggml/src/ggml-virtgpu/backend/backend.cpp | 144 - ggml/src/ggml-virtgpu/backend/backend.cpp.inc | 1831 + .../backend/shared/api_remoting.h | 95 - .../backend/shared/apir_backend.gen.h | 94 - .../backend/shared/apir_backend.h | 50 - .../src/ggml-virtgpu/backend/shared/apir_cs.h | 378 - .../backend/shared/apir_cs_ggml.h | 232 - .../ggml-virtgpu/backend/shared/apir_cs_rpc.h | 58 - .../ggml-virtgpu/ggml-backend-buffer-type.cpp | 81 - ggml/src/ggml-virtgpu/ggml-backend-buffer.cpp | 123 - ggml/src/ggml-virtgpu/ggml-backend-device.cpp | 158 - ggml/src/ggml-virtgpu/ggml-backend-reg.cpp | 213 - ggml/src/ggml-virtgpu/ggml-backend.cpp | 71 - ggml/src/ggml-virtgpu/ggml-remoting.h | 71 - ggml/src/ggml-virtgpu/include/apir_hw.h | 9 - ggml/src/ggml-virtgpu/regenerate_remoting.py | 42 +- ggml/src/ggml-virtgpu/virtgpu-apir.h | 15 - .../ggml-virtgpu/virtgpu-forward-backend.cpp | 58 - .../virtgpu-forward-buffer-type.cpp | 110 - .../ggml-virtgpu/virtgpu-forward-buffer.cpp | 173 - .../ggml-virtgpu/virtgpu-forward-device.cpp | 192 - ggml/src/ggml-virtgpu/virtgpu-forward-impl.h | 36 - ggml/src/ggml-virtgpu/virtgpu-forward.gen.h | 53 - ggml/src/ggml-virtgpu/virtgpu-shm.cpp | 99 - ggml/src/ggml-virtgpu/virtgpu-shm.h | 23 - ggml/src/ggml-virtgpu/virtgpu-utils.cpp | 179 - ggml/src/ggml-virtgpu/virtgpu-utils.h | 86 - ggml/src/ggml-virtgpu/virtgpu.cpp | 545 - ggml/src/ggml-virtgpu/virtgpu.cpp.inc | 3376 + ggml/src/ggml-virtgpu/virtgpu.h | 115 - ggml/src/ggml-vulkan/CMakeLists.txt | 5 +- .../{ggml-vulkan.cpp => ggml-vulkan.cpp.inc} | 10 +- .../ggml-vulkan/vulkan-shaders/CMakeLists.txt | 3 +- .../ggml-vulkan/vulkan-shaders/upscale.comp | 2 +- ...ers-gen.cpp => vulkan-shaders-gen.cpp.inc} | 0 ggml/src/ggml-webgpu/CMakeLists.txt | 5 +- .../ggml-webgpu/ggml-webgpu-shader-lib.hpp | 3068 - .../{ggml-webgpu.cpp => ggml-webgpu.cpp.inc} | 3848 +- ggml/src/ggml-webgpu/pre_wgsl.hpp | 778 - ggml/src/ggml-zdnn/CMakeLists.txt | 5 +- ggml/src/ggml-zdnn/common.hpp | 59 - .../{ggml-zdnn.cpp => ggml-zdnn.cpp.inc} | 276 +- ggml/src/ggml-zdnn/mmf.cpp | 80 - ggml/src/ggml-zdnn/mmf.hpp | 12 - ggml/src/ggml-zdnn/utils.cpp | 79 - ggml/src/ggml-zdnn/utils.hpp | 19 - ggml/src/ggml-zendnn/CMakeLists.txt | 3 +- .../{ggml-zendnn.cpp => ggml-zendnn.cpp.inc} | 6 +- ggml/src/ggml.c | 7765 -- ggml/src/ggml.c.inc | 14500 ++++ ggml/src/ggml.cpp | 26 - ggml/src/gguf.cpp | 1556 - gguf-py/gguf/constants.py | 6 +- grammars/README.md | 2 +- include/llama-cpp.h | 30 - include/{llama.h => llama.h.inc} | 10 +- pocs/vdot/CMakeLists.txt | 47 +- pocs/vdot/q8dot-rust/Cargo.toml | 11 + pocs/vdot/q8dot-rust/build.rs | 46 + pocs/vdot/q8dot-rust/src/lib.rs | 174 + pocs/vdot/q8dot-rust/src/main.rs | 140 + pocs/vdot/q8dot.cpp | 173 - pocs/vdot/vdot-rust/Cargo.toml | 9 + pocs/vdot/vdot-rust/build.rs | 41 + pocs/vdot/vdot-rust/src/lib.rs | 266 + pocs/vdot/vdot-rust/src/main.rs | 203 + pocs/vdot/vdot.cpp | 311 - scripts/rust-port-status.sh | 34 + scripts/sync_vendor.py | 18 +- src/CMakeLists.txt | 96 +- src/llama-adapter.cpp | 497 - src/llama-adapter.h | 91 - src/llama-arch.cpp | 908 - src/llama-arch.h | 638 - src/llama-batch.cpp | 919 - src/llama-batch.h | 173 - src/llama-chat.cpp | 939 - src/llama-chat.h | 74 - src/llama-context.cpp | 3533 - src/llama-context.h | 349 - src/llama-cparams-rust/Cargo.toml | 10 + src/llama-cparams-rust/src/lib.rs | 24 + src/llama-cparams.cpp | 5 - src/llama-cparams.h | 47 - src/llama-ext.h | 90 - src/llama-grammar.cpp | 1510 - src/llama-grammar.h | 194 - src/llama-graph.cpp | 2924 - src/llama-graph.h | 1064 - src/llama-hparams.cpp | 258 - src/llama-hparams.h | 357 - src/llama-impl.cpp | 171 - src/llama-impl.h | 75 - src/llama-impl.h.inc | 7603 ++ src/llama-io-rust/Cargo.toml | 8 + src/llama-io-rust/src/lib.rs | 90 + src/llama-io.cpp | 15 - src/llama-io.h | 35 - src/llama-kv-cache-iswa.cpp | 330 - src/llama-kv-cache-iswa.h | 137 - src/llama-kv-cache.cpp | 2504 - src/llama-kv-cache.h | 420 - src/llama-kv-cells.h | 533 - src/llama-memory-hybrid-iswa.cpp | 275 - src/llama-memory-hybrid-iswa.h | 140 - src/llama-memory-hybrid.cpp | 268 - src/llama-memory-hybrid.h | 139 - src/llama-memory-recurrent.cpp | 1163 - src/llama-memory-recurrent.h | 182 - src/llama-memory-rust/Cargo.toml | 8 + src/llama-memory-rust/src/lib.rs | 109 + src/llama-memory.cpp | 59 - src/llama-memory.h | 122 - src/llama-mmap.cpp | 779 - src/llama-mmap.h | 74 - src/llama-model-loader.cpp | 1695 - src/llama-model-loader.h | 207 - src/llama-model-saver.cpp | 414 - src/llama-model-saver.h | 44 - src/llama-model.cpp | 9489 --- src/llama-model.h | 640 - src/llama-quant.cpp | 1402 - src/llama-quant.h | 1 - src/llama-sampler.cpp | 3885 - src/llama-sampler.h | 42 - src/llama-vocab.cpp | 4102 - src/llama-vocab.h | 188 - src/llama.cpp | 558 - src/llama.cpp.inc | 61760 ++++++++++++++++ src/models/afmoe.cpp | 179 - src/models/apertus.cpp | 113 - src/models/arcee.cpp | 112 - src/models/arctic.cpp | 127 - src/models/arwkv7.cpp | 86 - src/models/baichuan.cpp | 112 - src/models/bailingmoe.cpp | 121 - src/models/bailingmoe2.cpp | 125 - src/models/bert.cpp | 156 - src/models/bitnet.cpp | 123 - src/models/bloom.cpp | 89 - src/models/chameleon.cpp | 156 - src/models/chatglm.cpp | 107 - src/models/codeshell.cpp | 103 - src/models/cogvlm.cpp | 108 - src/models/cohere2-iswa.cpp | 112 - src/models/command-r.cpp | 103 - src/models/dbrx.cpp | 110 - src/models/deci.cpp | 114 - src/models/deepseek.cpp | 123 - src/models/deepseek2.cpp | 296 - src/models/delta-net-base.cpp | 445 - src/models/dots1.cpp | 122 - src/models/dream.cpp | 89 - src/models/ernie4-5-moe.cpp | 129 - src/models/ernie4-5.cpp | 91 - src/models/eurobert.cpp | 89 - src/models/exaone-moe.cpp | 135 - src/models/exaone.cpp | 93 - src/models/exaone4.cpp | 112 - src/models/falcon-h1.cpp | 100 - src/models/falcon.cpp | 114 - src/models/gemma-embedding.cpp | 106 - src/models/gemma.cpp | 101 - src/models/gemma2-iswa.cpp | 118 - src/models/gemma3.cpp | 145 - src/models/gemma3n-iswa.cpp | 378 - src/models/gemma4-iswa.cpp | 322 - src/models/glm4-moe.cpp | 151 - src/models/glm4.cpp | 122 - src/models/gpt2.cpp | 93 - src/models/gptneox.cpp | 135 - src/models/granite-hybrid.cpp | 171 - src/models/granite.cpp | 187 - src/models/grok.cpp | 140 - src/models/grovemoe.cpp | 129 - src/models/hunyuan-dense.cpp | 132 - src/models/hunyuan-moe.cpp | 134 - src/models/internlm2.cpp | 101 - src/models/jais.cpp | 80 - src/models/jais2.cpp | 106 - src/models/jamba.cpp | 93 - src/models/kimi-linear.cpp | 382 - src/models/lfm2.cpp | 191 - src/models/llada-moe.cpp | 112 - src/models/llada.cpp | 90 - src/models/llama.cpp | 153 - src/models/llama4.cpp | 171 - src/models/maincoder.cpp | 107 - src/models/mamba-base.cpp | 289 - src/models/mamba.cpp | 54 - src/models/mimo2-iswa.cpp | 129 - src/models/minicpm3.cpp | 200 - src/models/minimax-m2.cpp | 123 - src/models/mistral3.cpp | 141 - src/models/models.h | 720 - src/models/modern-bert.cpp | 102 - src/models/mpt.cpp | 106 - src/models/nemotron-h.cpp | 137 - src/models/nemotron.cpp | 103 - src/models/neo-bert.cpp | 94 - src/models/olmo.cpp | 102 - src/models/olmo2.cpp | 150 - src/models/olmoe.cpp | 124 - src/models/openai-moe-iswa.cpp | 108 - src/models/openelm.cpp | 124 - src/models/orion.cpp | 101 - src/models/paddleocr.cpp | 103 - src/models/pangu-embedded.cpp | 107 - src/models/phi2.cpp | 98 - src/models/phi3.cpp | 132 - src/models/plamo.cpp | 100 - src/models/plamo2.cpp | 321 - src/models/plamo3.cpp | 130 - src/models/plm.cpp | 169 - src/models/qwen.cpp | 100 - src/models/qwen2.cpp | 104 - src/models/qwen2moe.cpp | 132 - src/models/qwen2vl.cpp | 104 - src/models/qwen3.cpp | 107 - src/models/qwen35.cpp | 385 - src/models/qwen35moe.cpp | 426 - src/models/qwen3moe.cpp | 118 - src/models/qwen3next.cpp | 528 - src/models/qwen3vl-moe.cpp | 130 - src/models/qwen3vl.cpp | 122 - src/models/refact.cpp | 84 - src/models/rnd1.cpp | 116 - src/models/rwkv6-base.cpp | 164 - src/models/rwkv6.cpp | 94 - src/models/rwkv6qwen2.cpp | 86 - src/models/rwkv7-base.cpp | 137 - src/models/rwkv7.cpp | 90 - src/models/seed-oss.cpp | 105 - src/models/smallthinker.cpp | 117 - src/models/smollm3.cpp | 109 - src/models/stablelm.cpp | 124 - src/models/starcoder.cpp | 88 - src/models/starcoder2.cpp | 102 - src/models/step35-iswa.cpp | 167 - src/models/t5.cpp | 252 - src/models/t5encoder.cpp | 3 - src/models/wavtokenizer-dec.cpp | 149 - src/models/xverse.cpp | 98 - src/unicode-data.cpp | 7034 -- src/unicode-data.h | 20 - src/unicode.cpp | 1275 - src/unicode.h | 111 - target/.rustc_info.json | 1 + target/CACHEDIR.TAG | 3 + target/debug/.cargo-lock | 0 .../dep-lib-ggml_backend_dl_rust | Bin 0 -> 30 bytes .../invoked.timestamp | 1 + .../lib-ggml_backend_dl_rust | 1 + .../lib-ggml_backend_dl_rust.json | 1 + .../dep-test-lib-ggml_backend_dl_rust | Bin 0 -> 30 bytes .../invoked.timestamp | 1 + .../test-lib-ggml_backend_dl_rust | 1 + .../test-lib-ggml_backend_dl_rust.json | 1 + .../dep-lib-ggml_threading_rust | Bin 0 -> 30 bytes .../invoked.timestamp | 1 + .../lib-ggml_threading_rust | 1 + .../lib-ggml_threading_rust.json | 1 + .../dep-test-lib-ggml_threading_rust | Bin 0 -> 30 bytes .../invoked.timestamp | 1 + .../test-lib-ggml_threading_rust | 1 + .../test-lib-ggml_threading_rust.json | 1 + .../dep-lib-llama_common_unicode_rust | Bin 0 -> 30 bytes .../invoked.timestamp | 1 + .../lib-llama_common_unicode_rust | 1 + .../lib-llama_common_unicode_rust.json | 1 + .../dep-test-lib-llama_common_unicode_rust | Bin 0 -> 30 bytes .../invoked.timestamp | 1 + .../test-lib-llama_common_unicode_rust | 1 + .../test-lib-llama_common_unicode_rust.json | 1 + .../dep-lib-llama_server_base64_rust | Bin 0 -> 30 bytes .../invoked.timestamp | 1 + .../lib-llama_server_base64_rust | 1 + .../lib-llama_server_base64_rust.json | 1 + .../dep-test-lib-llama_server_base64_rust | Bin 0 -> 30 bytes .../invoked.timestamp | 1 + .../test-lib-llama_server_base64_rust | 1 + .../test-lib-llama_server_base64_rust.json | 1 + .../ggml_backend_dl_rust-b4756e65d1ce8962 | Bin 0 -> 7078536 bytes .../deps/ggml_threading_rust-eefbb14ef3d83e68 | Bin 0 -> 6976304 bytes ...ggml_backend_dl_rust-2d121e9e1f8e6973.rlib | Bin 0 -> 437840 bytes ...bggml_threading_rust-0a510f30ab2a7a70.rlib | Bin 0 -> 104632 bytes ..._common_unicode_rust-1502b8c604829416.rlib | Bin 0 -> 2058340 bytes ...a_server_base64_rust-49d340cccbd5cf1e.rlib | Bin 0 -> 457180 bytes ...llama_common_unicode_rust-9be35eb4ef816de7 | Bin 0 -> 8245120 bytes .../llama_server_base64_rust-ec05054f222a1afa | Bin 0 -> 7106208 bytes .../metadata.rmeta | Bin 0 -> 8527 bytes .../s-hiarnw10ke-071vupl.lock | 0 .../metadata.rmeta | Bin 0 -> 8892 bytes .../s-hiaroj9fec-1fpkguj.lock | 0 .../s-hiarnw10kc-0pzxdd3.lock | 0 .../s-hiaroj9fe7-1sw5mx8.lock | 0 .../metadata.rmeta | Bin 0 -> 4892 bytes .../s-hiawyw92a5-0ad3ptp.lock | 0 .../metadata.rmeta | Bin 0 -> 5340 bytes .../s-hiax0mj3dr-19df9zm.lock | 0 .../s-hiawyw92a3-199hpju.lock | 0 .../s-hiax0mj3dp-0arubnp.lock | 0 .../metadata.rmeta | Bin 0 -> 64294 bytes .../s-hiavvq3mw9-1h8uhp8.lock | 0 .../metadata.rmeta | Bin 0 -> 64928 bytes .../s-hiavznzjim-05mxdnk.lock | 0 .../s-hiavvq3n1n-0rfum30.lock | 0 .../s-hiavznzjmv-07vn1ra.lock | 0 .../s-hiawfg0ftc-1j30nia.lock | 0 .../s-hiawky1p23-1de387b.lock | 0 .../metadata.rmeta | Bin 0 -> 8152 bytes .../s-hiawfg0fu5-16hqef7.lock | 0 .../metadata.rmeta | Bin 0 -> 8512 bytes .../s-hiawky1p25-0g5qg3t.lock | 0 tests/.gitignore | 2 +- tests/CMakeLists.txt | 906 +- tests/export-graph-ops.cpp | 226 - tests/get-model.cpp | 21 - tests/get-model.h | 2 - tests/gguf-model-data.cpp | 740 - tests/gguf-model-data.h | 50 - tests/no-c-cpp-files.cmake | 30 + tests/peg-parser/simple-tokenize.cpp | 37 - tests/peg-parser/simple-tokenize.h | 6 - tests/peg-parser/test-basic.cpp | 471 - tests/peg-parser/test-gbnf-generation.cpp | 370 - tests/peg-parser/test-json-parser.cpp | 109 - tests/peg-parser/test-json-serialization.cpp | 28 - tests/peg-parser/test-python-dict-parser.cpp | 318 - tests/peg-parser/test-unicode.cpp | 446 - tests/peg-parser/tests.h | 25 - tests/test-alloc.cpp | 608 - tests/test-arg-parser.cpp | 212 - tests/test-autorelease.cpp | 24 - tests/test-backend-ops.cpp | 9541 --- tests/test-backend-sampler.cpp | 1166 - tests/test-barrier.cpp | 236 - tests/test-c.c | 3 - tests/test-chat-auto-parser.cpp | 1969 - tests/test-chat-peg-parser.cpp | 983 - tests/test-chat-template.cpp | 712 - tests/test-chat.cpp | 4519 -- tests/test-double-float.cpp | 57 - tests/test-gbnf-validator.cpp | 109 - tests/test-gguf-model-data.cpp | 154 - tests/test-gguf.cpp | 1365 - tests/test-grammar-integration.cpp | 1493 - tests/test-grammar-llguidance.cpp | 1204 - tests/test-grammar-parser.cpp | 537 - tests/test-jinja.cpp | 2514 - tests/test-json-partial.cpp | 287 - tests/test-json-schema-to-grammar.cpp | 1589 - tests/test-llama-archs.cpp | 663 - tests/test-llama-grammar.cpp | 406 - tests/test-log.cpp | 43 - tests/test-model-load-cancel.cpp | 27 - tests/test-mtmd-c-api.c | 65 - tests/test-opt.cpp | 1003 - tests/test-peg-parser.cpp | 26 - tests/test-quant-type-selection.cpp | 520 - tests/test-quantize-fns.cpp | 196 - tests/test-quantize-perf.cpp | 356 - tests/test-quantize-stats.cpp | 427 - tests/test-reasoning-budget.cpp | 260 - tests/test-regex-partial.cpp | 288 - tests/test-rope.cpp | 263 - tests/test-sampling.cpp | 400 - tests/test-state-restore-fragmented.cpp | 122 - tests/test-thread-safety.cpp | 164 - tests/test-tokenizer-0.cpp | 312 - tests/test-tokenizer-1-bpe.cpp | 155 - tests/test-tokenizer-1-spm.cpp | 125 - tests/test-tokenizer-random.py | 2 +- tests/testing.h | 243 - tools/batched-bench-rust/Cargo.toml | 12 + tools/batched-bench-rust/build.rs | 45 + tools/batched-bench-rust/src/main.rs | 588 + tools/batched-bench/CMakeLists.txt | 28 +- tools/batched-bench/batched-bench.cpp | 259 - tools/cli-rust/Cargo.toml | 8 + tools/cli-rust/src/main.rs | 468 + tools/cli/CMakeLists.txt | 23 +- tools/cli/cli.cpp | 652 - tools/completion-rust/Cargo.toml | 8 + tools/completion-rust/src/main.rs | 409 + tools/completion/CMakeLists.txt | 23 +- tools/completion/completion.cpp | 1003 - tools/cvector-generator-rust/Cargo.toml | 9 + tools/cvector-generator-rust/src/main.rs | 311 + tools/cvector-generator/CMakeLists.txt | 21 +- tools/cvector-generator/cvector-generator.cpp | 515 - tools/cvector-generator/mean-rust/Cargo.toml | 10 + tools/cvector-generator/mean-rust/src/lib.rs | 86 + tools/cvector-generator/mean.hpp | 48 - tools/cvector-generator/pca.hpp | 315 - tools/export-lora-rust/Cargo.toml | 9 + tools/export-lora-rust/src/main.rs | 242 + tools/export-lora/CMakeLists.txt | 20 +- tools/export-lora/export-lora.cpp | 439 - tools/fit-params-rust/Cargo.toml | 9 + tools/fit-params-rust/src/main.rs | 224 + tools/fit-params/CMakeLists.txt | 22 +- tools/fit-params/fit-params.cpp | 74 - tools/gguf-split-rust/Cargo.toml | 12 + tools/gguf-split-rust/build.rs | 45 + tools/gguf-split-rust/src/main.rs | 719 + tools/gguf-split/CMakeLists.txt | 28 +- tools/gguf-split/gguf-split.cpp | 589 - tools/imatrix-rust/Cargo.toml | 9 + tools/imatrix-rust/src/main.rs | 420 + tools/imatrix/CMakeLists.txt | 26 +- tools/imatrix/imatrix.cpp | 1318 - tools/llama-bench-rust/Cargo.toml | 8 + tools/llama-bench-rust/src/main.rs | 635 + tools/llama-bench/CMakeLists.txt | 23 +- tools/llama-bench/llama-bench.cpp | 2431 - tools/mtmd/CMakeLists.txt | 124 +- tools/mtmd/README.md | 2 +- tools/mtmd/cli-rust/Cargo.toml | 9 + tools/mtmd/cli-rust/src/main.rs | 322 + tools/mtmd/clip-graph.h | 122 - tools/mtmd/clip-impl.h | 674 - tools/mtmd/clip-model.h | 531 - tools/mtmd/clip.cpp | 3941 - tools/mtmd/clip.h | 118 - tools/mtmd/debug-rust/Cargo.toml | 9 + tools/mtmd/debug-rust/src/main.rs | 373 + tools/mtmd/debug/mtmd-debug.cpp | 230 - tools/mtmd/debug/mtmd-debug.h | 17 - .../mtmd/deprecation-warning-rust/Cargo.toml | 7 + .../mtmd/deprecation-warning-rust/src/lib.rs | 37 + .../mtmd/deprecation-warning-rust/src/main.rs | 7 + tools/mtmd/deprecation-warning.cpp | 25 - tools/mtmd/models/cogvlm.cpp | 98 - tools/mtmd/models/conformer.cpp | 216 - tools/mtmd/models/deepseekocr.cpp | 324 - tools/mtmd/models/dotsocr.cpp | 49 - tools/mtmd/models/gemma4a.cpp | 288 - tools/mtmd/models/gemma4v.cpp | 151 - tools/mtmd/models/glm4v.cpp | 122 - tools/mtmd/models/hunyuanocr.cpp | 73 - tools/mtmd/models/internvl.cpp | 69 - tools/mtmd/models/kimik25.cpp | 101 - tools/mtmd/models/kimivl.cpp | 63 - tools/mtmd/models/llama4.cpp | 96 - tools/mtmd/models/llava.cpp | 374 - tools/mtmd/models/minicpmv.cpp | 114 - tools/mtmd/models/mobilenetv5.cpp | 451 - tools/mtmd/models/models.h | 173 - tools/mtmd/models/nemotron-v2-vl.cpp | 35 - tools/mtmd/models/paddleocr.cpp | 52 - tools/mtmd/models/pixtral.cpp | 86 - tools/mtmd/models/qwen2vl.cpp | 183 - tools/mtmd/models/qwen3a.cpp | 68 - tools/mtmd/models/qwen3vl.cpp | 193 - tools/mtmd/models/siglip.cpp | 94 - tools/mtmd/models/step3vl.cpp | 81 - tools/mtmd/models/whisper-enc.cpp | 137 - tools/mtmd/models/yasa2.cpp | 191 - tools/mtmd/models/youtuvl.cpp | 179 - tools/mtmd/mtmd-audio.cpp | 836 - tools/mtmd/mtmd-audio.h | 123 - tools/mtmd/mtmd-cli.cpp | 441 - tools/mtmd/mtmd-helper.cpp | 537 - tools/mtmd/mtmd-helper.h | 100 - tools/mtmd/mtmd-image.cpp | 1427 - tools/mtmd/mtmd-image.h | 179 - tools/mtmd/mtmd.cpp | 1509 - tools/mtmd/mtmd.cpp.inc | 12489 ++++ tools/mtmd/mtmd.h | 332 - tools/mtmd/mtmd.h.inc | 1708 + tools/parser-rust/Cargo.toml | 16 + tools/parser-rust/build.rs | 45 + .../src/bin/debug_template_parser.rs | 90 + .../parser-rust/src/bin/template_analysis.rs | 76 + tools/parser-rust/src/lib.rs | 389 + tools/parser/CMakeLists.txt | 58 +- tools/parser/debug-template-parser.cpp | 462 - tools/parser/template-analysis.cpp | 611 - tools/perplexity-rust/Cargo.toml | 8 + tools/perplexity-rust/src/main.rs | 455 + tools/perplexity/CMakeLists.txt | 23 +- tools/perplexity/perplexity.cpp | 2096 - tools/quantize-rust/Cargo.toml | 9 + tools/quantize-rust/src/main.rs | 529 + tools/quantize/CMakeLists.txt | 22 +- tools/quantize/quantize.cpp | 756 - tools/results-rust/Cargo.toml | 12 + tools/results-rust/build.rs | 43 + tools/results-rust/src/lib.rs | 127 + tools/results-rust/src/main.rs | 360 + tools/results/CMakeLists.txt | 29 +- tools/results/results.cpp | 183 - tools/rpc-rust/Cargo.toml | 11 + tools/rpc-rust/build.rs | 44 + tools/rpc-rust/src/main.rs | 406 + tools/rpc/CMakeLists.txt | 25 +- tools/rpc/rpc-server.cpp | 342 - tools/server/CMakeLists.txt | 40 +- tools/server/README.md | 4 +- tools/server/base64-rust/Cargo.toml | 8 + tools/server/base64-rust/src/lib.rs | 266 + tools/server/server-chat.cpp | 630 - tools/server/server-chat.h | 26 - tools/server/server-common.cpp | 1586 - tools/server/server-common.cpp.inc | 11567 +++ tools/server/server-common.h | 373 - tools/server/server-context.cpp | 4347 -- tools/server/server-context.h | 146 - tools/server/server-cors-proxy.h | 67 - tools/server/server-http.cpp | 483 - tools/server/server-http.h | 89 - tools/server/server-models.cpp | 1365 - tools/server/server-models.h | 223 - tools/server/server-queue.cpp | 450 - tools/server/server-queue.h | 205 - tools/server/server-task.cpp | 2127 - tools/server/server-task.h | 645 - tools/server/server-tools.cpp | 768 - tools/server/server-tools.h | 26 - tools/server/server.cpp | 353 - tools/server/server.cpp.inc | 3866 + tools/server/tests/README.md | 2 +- tools/server/tests/unit/test_infill.py | 12 +- tools/server/webui/src/lib/enums/server.ts | 2 +- .../webui/tests/stories/fixtures/blog-post.ts | 2 +- tools/tokenize-rust/Cargo.toml | 9 + tools/tokenize-rust/build.rs | 42 + tools/tokenize-rust/src/lib.rs | 219 + tools/tokenize-rust/src/main.rs | 354 + tools/tokenize/CMakeLists.txt | 26 +- tools/tokenize/tokenize.cpp | 419 - tools/tts-rust/Cargo.toml | 9 + tools/tts-rust/src/main.rs | 466 + tools/tts/CMakeLists.txt | 21 +- tools/tts/tts.cpp | 1096 - vendor/cpp-httplib/CMakeLists.txt | 6 +- .../{httplib.cpp => httplib.cpp.inc} | 2 +- .../cpp-httplib/{httplib.h => httplib.h.inc} | 2 +- .../{miniaudio.h => miniaudio.h.inc} | 2 +- vendor/nlohmann/{json.hpp => json.hpp.inc} | 14 +- .../{json_fwd.hpp => json_fwd.hpp.inc} | 0 .../{subprocess.h => subprocess.h.inc} | 2 +- vendor/stb/{stb_image.h => stb_image.h.inc} | 2 +- 1439 files changed, 247357 insertions(+), 279648 deletions(-) create mode 100644 Cargo.lock create mode 100644 Cargo.toml delete mode 100644 common/arg.cpp delete mode 100644 common/arg.h delete mode 100644 common/base64.hpp create mode 100644 common/build-info-rust/Cargo.toml create mode 100644 common/build-info-rust/build.rs create mode 100644 common/build-info-rust/src/lib.rs delete mode 100644 common/build-info.cpp.in delete mode 100644 common/build-info.h delete mode 100644 common/chat-auto-parser-generator.cpp delete mode 100644 common/chat-auto-parser-helpers.cpp delete mode 100644 common/chat-auto-parser-helpers.h delete mode 100644 common/chat-auto-parser.h delete mode 100644 common/chat-diff-analyzer.cpp delete mode 100644 common/chat-peg-parser.cpp delete mode 100644 common/chat-peg-parser.h delete mode 100644 common/chat.cpp delete mode 100644 common/chat.h delete mode 100644 common/common.cpp create mode 100644 common/common.cpp.inc delete mode 100644 common/common.h create mode 100644 common/common.h.inc delete mode 100644 common/console.cpp delete mode 100644 common/console.h delete mode 100644 common/debug.cpp delete mode 100644 common/debug.h delete mode 100644 common/download.cpp delete mode 100644 common/download.h delete mode 100644 common/fit.cpp delete mode 100644 common/fit.h delete mode 100644 common/hf-cache.cpp delete mode 100644 common/hf-cache.h delete mode 100644 common/http.h delete mode 100644 common/jinja/caps.cpp delete mode 100644 common/jinja/caps.h delete mode 100644 common/jinja/lexer.cpp delete mode 100644 common/jinja/lexer.h delete mode 100644 common/jinja/parser.cpp delete mode 100644 common/jinja/parser.h delete mode 100644 common/jinja/runtime.cpp delete mode 100644 common/jinja/runtime.h create mode 100644 common/jinja/string-rust/Cargo.toml create mode 100644 common/jinja/string-rust/src/lib.rs delete mode 100644 common/jinja/string.cpp delete mode 100644 common/jinja/string.h delete mode 100644 common/jinja/utils.h delete mode 100644 common/jinja/value.cpp delete mode 100644 common/jinja/value.h delete mode 100644 common/json-partial.cpp delete mode 100644 common/json-partial.h delete mode 100644 common/json-schema-to-grammar.cpp delete mode 100644 common/json-schema-to-grammar.h delete mode 100644 common/llguidance.cpp delete mode 100644 common/log.cpp delete mode 100644 common/log.h create mode 100644 common/ngram-cache-rust/Cargo.toml create mode 100644 common/ngram-cache-rust/src/lib.rs delete mode 100644 common/ngram-cache.cpp delete mode 100644 common/ngram-cache.h create mode 100644 common/ngram-map-rust/Cargo.toml create mode 100644 common/ngram-map-rust/src/lib.rs delete mode 100644 common/ngram-map.cpp delete mode 100644 common/ngram-map.h create mode 100644 common/ngram-mod-rust/Cargo.toml create mode 100644 common/ngram-mod-rust/src/lib.rs delete mode 100644 common/ngram-mod.cpp delete mode 100644 common/ngram-mod.h delete mode 100644 common/peg-parser.cpp delete mode 100644 common/peg-parser.h delete mode 100644 common/preset.cpp delete mode 100644 common/preset.h create mode 100644 common/reasoning-budget-rust/Cargo.toml create mode 100644 common/reasoning-budget-rust/src/lib.rs delete mode 100644 common/reasoning-budget.cpp delete mode 100644 common/reasoning-budget.h create mode 100644 common/regex-partial-rust/Cargo.toml create mode 100644 common/regex-partial-rust/src/lib.rs delete mode 100644 common/regex-partial.cpp delete mode 100644 common/regex-partial.h delete mode 100644 common/sampling.cpp delete mode 100644 common/sampling.h delete mode 100644 common/speculative.cpp delete mode 100644 common/speculative.h create mode 100644 common/unicode-rust/Cargo.toml create mode 100644 common/unicode-rust/src/lib.rs delete mode 100644 common/unicode.cpp delete mode 100644 common/unicode.h create mode 100644 docs/development/rust-port.md create mode 100644 examples/batched-rust/Cargo.toml create mode 100644 examples/batched-rust/build.rs create mode 100644 examples/batched-rust/src/main.rs delete mode 100644 examples/batched/batched.cpp create mode 100644 examples/convert-llama2c-to-ggml-rust/Cargo.toml create mode 100644 examples/convert-llama2c-to-ggml-rust/src/main.rs delete mode 100644 examples/convert-llama2c-to-ggml/convert-llama2c-to-ggml.cpp create mode 100644 examples/debug-rust/Cargo.toml create mode 100644 examples/debug-rust/build.rs create mode 100644 examples/debug-rust/src/lib.rs create mode 100644 examples/debug-rust/src/main.rs delete mode 100644 examples/debug/debug.cpp create mode 100644 examples/deprecation-warning/Cargo.toml delete mode 100644 examples/deprecation-warning/deprecation-warning.cpp create mode 100644 examples/deprecation-warning/src/lib.rs create mode 100644 examples/deprecation-warning/src/main.rs create mode 100644 examples/diffusion-rust/Cargo.toml create mode 100644 examples/diffusion-rust/build.rs create mode 100644 examples/diffusion-rust/src/main.rs delete mode 100644 examples/diffusion/diffusion-cli.cpp create mode 100644 examples/embedding-rust/Cargo.toml create mode 100644 examples/embedding-rust/build.rs create mode 100644 examples/embedding-rust/src/main.rs delete mode 100644 examples/embedding/embedding.cpp create mode 100644 examples/eval-callback-rust/Cargo.toml create mode 100644 examples/eval-callback-rust/build.rs create mode 100644 examples/eval-callback-rust/src/lib.rs create mode 100644 examples/eval-callback-rust/src/main.rs delete mode 100644 examples/eval-callback/eval-callback.cpp create mode 100644 examples/gen-docs-rust/Cargo.toml create mode 100644 examples/gen-docs-rust/src/main.rs delete mode 100644 examples/gen-docs/gen-docs.cpp create mode 100644 examples/gguf-hash-rust/Cargo.toml create mode 100644 examples/gguf-hash-rust/build.rs create mode 100644 examples/gguf-hash-rust/src/lib.rs create mode 100644 examples/gguf-hash-rust/src/main.rs delete mode 100644 examples/gguf-hash/deps/rotate-bits/rotate-bits.h delete mode 100644 examples/gguf-hash/deps/sha1/sha1.c delete mode 100644 examples/gguf-hash/deps/sha1/sha1.h delete mode 100644 examples/gguf-hash/deps/sha256/sha256.c delete mode 100644 examples/gguf-hash/deps/sha256/sha256.h delete mode 100644 examples/gguf-hash/deps/xxhash/xxhash.c delete mode 100644 examples/gguf-hash/deps/xxhash/xxhash.h delete mode 100644 examples/gguf-hash/gguf-hash.cpp create mode 100644 examples/gguf-rust/Cargo.toml create mode 100644 examples/gguf-rust/build.rs create mode 100644 examples/gguf-rust/src/lib.rs create mode 100644 examples/gguf-rust/src/main.rs delete mode 100644 examples/gguf/gguf.cpp create mode 100644 examples/idle-rust/Cargo.toml create mode 100644 examples/idle-rust/src/lib.rs create mode 100644 examples/idle-rust/src/main.rs delete mode 100644 examples/idle/idle.cpp rename examples/llama.android/lib/src/main/cpp/{ai_chat.cpp => ai_chat.cpp.inc} (90%) delete mode 100644 examples/llama.android/lib/src/main/cpp/logging.h create mode 100644 examples/lookahead-rust/Cargo.toml create mode 100644 examples/lookahead-rust/build.rs create mode 100644 examples/lookahead-rust/src/main.rs delete mode 100644 examples/lookahead/lookahead.cpp create mode 100644 examples/lookup-create-rust/Cargo.toml create mode 100644 examples/lookup-create-rust/build.rs create mode 100644 examples/lookup-create-rust/src/lib.rs create mode 100644 examples/lookup-create-rust/src/main.rs create mode 100644 examples/lookup-merge-rust/Cargo.toml create mode 100644 examples/lookup-merge-rust/src/lib.rs create mode 100644 examples/lookup-merge-rust/src/main.rs create mode 100644 examples/lookup-rust/Cargo.toml create mode 100644 examples/lookup-rust/build.rs create mode 100644 examples/lookup-rust/src/main.rs create mode 100644 examples/lookup-stats-rust/Cargo.toml create mode 100644 examples/lookup-stats-rust/build.rs create mode 100644 examples/lookup-stats-rust/src/lib.rs create mode 100644 examples/lookup-stats-rust/src/main.rs delete mode 100644 examples/lookup/lookup-create.cpp delete mode 100644 examples/lookup/lookup-merge.cpp delete mode 100644 examples/lookup/lookup-stats.cpp delete mode 100644 examples/lookup/lookup.cpp create mode 100644 examples/parallel-rust/Cargo.toml create mode 100644 examples/parallel-rust/build.rs create mode 100644 examples/parallel-rust/src/main.rs delete mode 100644 examples/parallel/parallel.cpp create mode 100644 examples/passkey-rust/Cargo.toml create mode 100644 examples/passkey-rust/build.rs create mode 100644 examples/passkey-rust/src/main.rs delete mode 100644 examples/passkey/passkey.cpp create mode 100644 examples/retrieval-rust/Cargo.toml create mode 100644 examples/retrieval-rust/build.rs create mode 100644 examples/retrieval-rust/src/main.rs delete mode 100644 examples/retrieval/retrieval.cpp create mode 100644 examples/save-load-state-rust/Cargo.toml create mode 100644 examples/save-load-state-rust/build.rs create mode 100644 examples/save-load-state-rust/src/main.rs delete mode 100644 examples/save-load-state/save-load-state.cpp create mode 100644 examples/simple-chat-rust/Cargo.toml create mode 100644 examples/simple-chat-rust/build.rs create mode 100644 examples/simple-chat-rust/src/lib.rs create mode 100644 examples/simple-chat-rust/src/main.rs delete mode 100644 examples/simple-chat/simple-chat.cpp create mode 100644 examples/simple-rust/Cargo.toml create mode 100644 examples/simple-rust/README.md create mode 100644 examples/simple-rust/build.rs create mode 100644 examples/simple-rust/src/ffi.rs create mode 100644 examples/simple-rust/src/lib.rs create mode 100644 examples/simple-rust/src/main.rs delete mode 100644 examples/simple/simple.cpp create mode 100644 examples/speculative-rust/Cargo.toml create mode 100644 examples/speculative-rust/build.rs create mode 100644 examples/speculative-rust/src/main.rs create mode 100644 examples/speculative-simple-rust/Cargo.toml create mode 100644 examples/speculative-simple-rust/build.rs create mode 100644 examples/speculative-simple-rust/src/main.rs delete mode 100644 examples/speculative-simple/speculative-simple.cpp delete mode 100644 examples/speculative/speculative.cpp create mode 100644 examples/sycl/ls-sycl-device-rust/Cargo.toml create mode 100644 examples/sycl/ls-sycl-device-rust/build.rs create mode 100644 examples/sycl/ls-sycl-device-rust/src/lib.rs create mode 100644 examples/sycl/ls-sycl-device-rust/src/main.rs delete mode 100644 examples/sycl/ls-sycl-device.cpp create mode 100644 examples/training/finetune-rust/Cargo.toml create mode 100644 examples/training/finetune-rust/build.rs create mode 100644 examples/training/finetune-rust/src/main.rs delete mode 100644 examples/training/finetune.cpp rename ggml/include/{ggml-alloc.h => ggml-alloc.h.inc} (99%) rename ggml/include/{ggml-backend.h => ggml-backend.h.inc} (99%) rename ggml/include/{ggml-blas.h => ggml-blas.h.inc} (91%) rename ggml/include/{ggml-cann.h => ggml-cann.h.inc} (98%) delete mode 100644 ggml/include/ggml-cpp.h rename ggml/include/{ggml-cpu.h => ggml-cpu.h.inc} (99%) rename ggml/include/{ggml-cuda.h => ggml-cuda.h.inc} (97%) rename ggml/include/{ggml-hexagon.h => ggml-hexagon.h.inc} (85%) rename ggml/include/{ggml-metal.h => ggml-metal.h.inc} (97%) rename ggml/include/{ggml-opencl.h => ggml-opencl.h.inc} (91%) rename ggml/include/{ggml-openvino.h => ggml-openvino.h.inc} (97%) rename ggml/include/{ggml-opt.h => ggml-opt.h.inc} (99%) rename ggml/include/{ggml-rpc.h => ggml-rpc.h.inc} (97%) rename ggml/include/{ggml-sycl.h => ggml-sycl.h.inc} (97%) rename ggml/include/{ggml-virtgpu.h => ggml-virtgpu.h.inc} (74%) rename ggml/include/{ggml-vulkan.h => ggml-vulkan.h.inc} (94%) rename ggml/include/{ggml-webgpu.h => ggml-webgpu.h.inc} (84%) rename ggml/include/{ggml-zdnn.h => ggml-zdnn.h.inc} (82%) rename ggml/include/{ggml-zendnn.h => ggml-zendnn.h.inc} (89%) rename ggml/include/{ggml.h => ggml.h.inc} (99%) rename ggml/include/{gguf.h => gguf.h.inc} (99%) delete mode 100644 ggml/src/ggml-alloc.c create mode 100644 ggml/src/ggml-backend-dl-rust/Cargo.toml create mode 100644 ggml/src/ggml-backend-dl-rust/src/lib.rs delete mode 100644 ggml/src/ggml-backend-dl.cpp delete mode 100644 ggml/src/ggml-backend-dl.h rename ggml/src/{ggml-backend-impl.h => ggml-backend-impl-defs.inc} (99%) delete mode 100644 ggml/src/ggml-backend-meta.cpp rename ggml/src/{ggml-backend-reg.cpp => ggml-backend-reg.cpp.inc} (84%) delete mode 100644 ggml/src/ggml-backend.cpp create mode 100644 ggml/src/ggml-backend.cpp.inc rename ggml/src/ggml-blas/{ggml-blas.cpp => ggml-blas.cpp.inc} (99%) delete mode 100644 ggml/src/ggml-cann/acl_tensor.cpp delete mode 100644 ggml/src/ggml-cann/acl_tensor.h delete mode 100644 ggml/src/ggml-cann/aclnn_ops.cpp delete mode 100644 ggml/src/ggml-cann/aclnn_ops.h delete mode 100644 ggml/src/ggml-cann/common.h rename ggml/src/ggml-cann/{ggml-cann.cpp => ggml-cann.cpp.inc} (54%) rename ggml/src/{ggml-common.h => ggml-common-defs.inc} (100%) delete mode 100644 ggml/src/ggml-cpu/amx/amx.cpp delete mode 100644 ggml/src/ggml-cpu/amx/amx.h delete mode 100644 ggml/src/ggml-cpu/amx/common.h delete mode 100644 ggml/src/ggml-cpu/amx/mmq.cpp delete mode 100644 ggml/src/ggml-cpu/amx/mmq.h delete mode 100644 ggml/src/ggml-cpu/arch-fallback.h create mode 100644 ggml/src/ggml-cpu/arch/arm/cpu-feats-rust/Cargo.toml create mode 100644 ggml/src/ggml-cpu/arch/arm/cpu-feats-rust/src/lib.rs delete mode 100644 ggml/src/ggml-cpu/arch/arm/cpu-feats.cpp rename ggml/src/ggml-cpu/arch/arm/{quants.c => quants.c.inc} (99%) rename ggml/src/ggml-cpu/arch/arm/{repack.cpp => repack.cpp.inc} (98%) rename ggml/src/ggml-cpu/arch/loongarch/{quants.c => quants.c.inc} (99%) create mode 100644 ggml/src/ggml-cpu/arch/powerpc/cpu-feats-rust/Cargo.toml create mode 100644 ggml/src/ggml-cpu/arch/powerpc/cpu-feats-rust/src/lib.rs delete mode 100644 ggml/src/ggml-cpu/arch/powerpc/cpu-feats.cpp rename ggml/src/ggml-cpu/arch/powerpc/{quants.c => quants.c.inc} (99%) create mode 100644 ggml/src/ggml-cpu/arch/riscv/cpu-feats-rust/Cargo.toml create mode 100644 ggml/src/ggml-cpu/arch/riscv/cpu-feats-rust/src/lib.rs delete mode 100644 ggml/src/ggml-cpu/arch/riscv/cpu-feats.cpp rename ggml/src/ggml-cpu/arch/riscv/{quants.c => quants.c.inc} (99%) rename ggml/src/ggml-cpu/arch/riscv/{repack.cpp => repack.cpp.inc} (96%) create mode 100644 ggml/src/ggml-cpu/arch/s390/cpu-feats-rust/Cargo.toml create mode 100644 ggml/src/ggml-cpu/arch/s390/cpu-feats-rust/src/lib.rs delete mode 100644 ggml/src/ggml-cpu/arch/s390/cpu-feats.cpp rename ggml/src/ggml-cpu/arch/s390/{quants.c => quants.c.inc} (99%) rename ggml/src/ggml-cpu/arch/wasm/{quants.c => quants.c.inc} (99%) create mode 100644 ggml/src/ggml-cpu/arch/x86/cpu-feats-rust/Cargo.toml create mode 100644 ggml/src/ggml-cpu/arch/x86/cpu-feats-rust/src/lib.rs delete mode 100644 ggml/src/ggml-cpu/arch/x86/cpu-feats.cpp delete mode 100644 ggml/src/ggml-cpu/binary-ops.cpp delete mode 100644 ggml/src/ggml-cpu/binary-ops.h delete mode 100644 ggml/src/ggml-cpu/common.h rename ggml/src/ggml-cpu/{ggml-cpu-impl.h => ggml-cpu-impl-defs.inc} (99%) rename ggml/src/ggml-cpu/{ggml-cpu.c => ggml-cpu.c.inc} (57%) delete mode 100644 ggml/src/ggml-cpu/ggml-cpu.cpp create mode 100644 ggml/src/ggml-cpu/ggml-cpu.cpp.inc delete mode 100644 ggml/src/ggml-cpu/hbm.cpp delete mode 100644 ggml/src/ggml-cpu/hbm.h rename ggml/src/ggml-cpu/kleidiai/{kernels.cpp => kernels.cpp.inc} (95%) delete mode 100644 ggml/src/ggml-cpu/kleidiai/kernels.h rename ggml/src/ggml-cpu/kleidiai/{kleidiai.cpp => kleidiai.cpp.inc} (89%) delete mode 100644 ggml/src/ggml-cpu/kleidiai/kleidiai.h delete mode 100644 ggml/src/ggml-cpu/llamafile/sgemm.cpp delete mode 100644 ggml/src/ggml-cpu/llamafile/sgemm.h rename ggml/src/ggml-cpu/{ops.cpp => ops.cpp.inc} (73%) delete mode 100644 ggml/src/ggml-cpu/ops.h delete mode 100644 ggml/src/ggml-cpu/quants.c rename ggml/src/ggml-cpu/{arch/x86/quants.c => quants.c.inc} (71%) delete mode 100644 ggml/src/ggml-cpu/quants.h delete mode 100644 ggml/src/ggml-cpu/repack.cpp rename ggml/src/ggml-cpu/{arch/x86/repack.cpp => repack.cpp.inc} (73%) delete mode 100644 ggml/src/ggml-cpu/repack.h delete mode 100644 ggml/src/ggml-cpu/simd-gemm.h rename ggml/src/ggml-cpu/{simd-mappings.h => simd-mappings-defs.inc} (99%) rename ggml/src/ggml-cpu/spacemit/{ime.cpp => ime.cpp.inc} (90%) delete mode 100644 ggml/src/ggml-cpu/spacemit/ime.h rename ggml/src/ggml-cpu/spacemit/{ime1_kernels.cpp => ime1_kernels.cpp.inc} (99%) delete mode 100644 ggml/src/ggml-cpu/spacemit/ime_kernels.h delete mode 100644 ggml/src/ggml-cpu/traits.cpp delete mode 100644 ggml/src/ggml-cpu/traits.h delete mode 100644 ggml/src/ggml-cpu/unary-ops.cpp delete mode 100644 ggml/src/ggml-cpu/unary-ops.h delete mode 100644 ggml/src/ggml-cpu/vec.cpp delete mode 100644 ggml/src/ggml-cpu/vec.h rename ggml/src/ggml-cuda/{acc.cu => acc.cu.inc} (99%) rename ggml/src/ggml-cuda/{acc.cuh => acc.cuh.inc} (80%) rename ggml/src/ggml-cuda/{add-id.cu => add-id.cu.inc} (98%) rename ggml/src/ggml-cuda/{add-id.cuh => add-id.cuh.inc} (75%) rename ggml/src/ggml-cuda/{arange.cu => arange.cu.inc} (97%) rename ggml/src/ggml-cuda/{arange.cuh => arange.cuh.inc} (81%) rename ggml/src/ggml-cuda/{argmax.cu => argmax.cu.inc} (97%) rename ggml/src/ggml-cuda/{argmax.cuh => argmax.cuh.inc} (74%) rename ggml/src/ggml-cuda/{argsort.cu => argsort.cu.inc} (99%) rename ggml/src/ggml-cuda/{argsort.cuh => argsort.cuh.inc} (97%) rename ggml/src/ggml-cuda/{binbcast.cu => binbcast.cu.inc} (99%) rename ggml/src/ggml-cuda/{binbcast.cuh => binbcast.cuh.inc} (96%) rename ggml/src/ggml-cuda/{clamp.cu => clamp.cu.inc} (98%) rename ggml/src/ggml-cuda/{clamp.cuh => clamp.cuh.inc} (81%) rename ggml/src/ggml-cuda/{common.cuh => common.cuh.inc} (99%) rename ggml/src/ggml-cuda/{concat.cu => concat.cu.inc} (99%) rename ggml/src/ggml-cuda/{concat.cuh => concat.cuh.inc} (81%) rename ggml/src/ggml-cuda/{conv-transpose-1d.cu => conv-transpose-1d.cu.inc} (98%) rename ggml/src/ggml-cuda/{conv-transpose-1d.cuh => conv-transpose-1d.cuh.inc} (83%) rename ggml/src/ggml-cuda/{conv2d-dw.cu => conv2d-dw.cu.inc} (99%) rename ggml/src/ggml-cuda/{conv2d-dw.cuh => conv2d-dw.cuh.inc} (83%) rename ggml/src/ggml-cuda/{conv2d-transpose.cu => conv2d-transpose.cu.inc} (98%) rename ggml/src/ggml-cuda/{conv2d-transpose.cuh => conv2d-transpose.cuh.inc} (83%) rename ggml/src/ggml-cuda/{conv2d.cu => conv2d.cu.inc} (99%) rename ggml/src/ggml-cuda/{conv2d.cuh => conv2d.cuh.inc} (83%) rename ggml/src/ggml-cuda/{convert.cu => convert.cu.inc} (99%) rename ggml/src/ggml-cuda/{convert.cuh => convert.cuh.inc} (98%) rename ggml/src/ggml-cuda/{count-equal.cu => count-equal.cu.inc} (97%) rename ggml/src/ggml-cuda/{count-equal.cuh => count-equal.cuh.inc} (82%) rename ggml/src/ggml-cuda/{cp-async.cuh => cp-async.cuh.inc} (98%) rename ggml/src/ggml-cuda/{cpy-utils.cuh => cpy-utils.cuh.inc} (98%) rename ggml/src/ggml-cuda/{cpy.cu => cpy.cu.inc} (99%) rename ggml/src/ggml-cuda/{cpy.cuh => cpy.cuh.inc} (88%) rename ggml/src/ggml-cuda/{cross-entropy-loss.cu => cross-entropy-loss.cu.inc} (98%) rename ggml/src/ggml-cuda/{cross-entropy-loss.cuh => cross-entropy-loss.cuh.inc} (89%) rename ggml/src/ggml-cuda/{cumsum.cu => cumsum.cu.inc} (98%) rename ggml/src/ggml-cuda/{cumsum.cuh => cumsum.cuh.inc} (81%) rename ggml/src/ggml-cuda/{dequantize.cuh => dequantize.cuh.inc} (99%) rename ggml/src/ggml-cuda/{diag.cu => diag.cu.inc} (97%) rename ggml/src/ggml-cuda/{diag.cuh => diag.cuh.inc} (81%) rename ggml/src/ggml-cuda/{diagmask.cu => diagmask.cu.inc} (98%) rename ggml/src/ggml-cuda/{diagmask.cuh => diagmask.cuh.inc} (83%) rename ggml/src/ggml-cuda/{fattn-common.cuh => fattn-common.cuh.inc} (99%) rename ggml/src/ggml-cuda/{fattn-mma-f16.cuh => fattn-mma-f16.cuh.inc} (99%) rename ggml/src/ggml-cuda/{fattn-tile.cu => fattn-tile.cu.inc} (95%) rename ggml/src/ggml-cuda/{fattn-tile.cuh => fattn-tile.cuh.inc} (99%) rename ggml/src/ggml-cuda/{fattn-vec.cuh => fattn-vec.cuh.inc} (99%) rename ggml/src/ggml-cuda/{fattn-wmma-f16.cu => fattn-wmma-f16.cu.inc} (99%) rename ggml/src/ggml-cuda/{fattn-wmma-f16.cuh => fattn-wmma-f16.cuh.inc} (98%) rename ggml/src/ggml-cuda/{fattn.cu => fattn.cu.inc} (99%) rename ggml/src/ggml-cuda/{fattn.cuh => fattn.cuh.inc} (86%) rename ggml/src/ggml-cuda/{fill.cu => fill.cu.inc} (95%) rename ggml/src/ggml-cuda/{fill.cuh => fill.cuh.inc} (74%) rename ggml/src/ggml-cuda/{gated_delta_net.cu => gated_delta_net.cu.inc} (99%) rename ggml/src/ggml-cuda/{gated_delta_net.cuh => gated_delta_net.cuh.inc} (64%) rename ggml/src/ggml-cuda/{getrows.cu => getrows.cu.inc} (99%) rename ggml/src/ggml-cuda/{getrows.cuh => getrows.cuh.inc} (95%) rename ggml/src/ggml-cuda/{ggml-cuda.cu => ggml-cuda.cu.inc} (98%) rename ggml/src/ggml-cuda/{gla.cu => gla.cu.inc} (98%) rename ggml/src/ggml-cuda/{gla.cuh => gla.cuh.inc} (77%) rename ggml/src/ggml-cuda/{im2col.cu => im2col.cu.inc} (99%) rename ggml/src/ggml-cuda/{im2col.cuh => im2col.cuh.inc} (88%) rename ggml/src/ggml-cuda/{mean.cu => mean.cu.inc} (98%) rename ggml/src/ggml-cuda/{mean.cuh => mean.cuh.inc} (74%) rename ggml/src/ggml-cuda/{mma.cuh => mma.cuh.inc} (99%) rename ggml/src/ggml-cuda/{mmf.cu => mmf.cu.inc} (99%) rename ggml/src/ggml-cuda/{mmf.cuh => mmf.cuh.inc} (99%) rename ggml/src/ggml-cuda/{mmid.cu => mmid.cu.inc} (99%) rename ggml/src/ggml-cuda/{mmid.cuh => mmid.cuh.inc} (100%) rename ggml/src/ggml-cuda/{mmq.cu => mmq.cu.inc} (99%) rename ggml/src/ggml-cuda/{mmq.cuh => mmq.cuh.inc} (99%) rename ggml/src/ggml-cuda/{mmvf.cu => mmvf.cu.inc} (99%) rename ggml/src/ggml-cuda/{mmvf.cuh => mmvf.cuh.inc} (96%) rename ggml/src/ggml-cuda/{mmvq.cu => mmvq.cu.inc} (99%) rename ggml/src/ggml-cuda/{mmvq.cuh => mmvq.cuh.inc} (97%) rename ggml/src/ggml-cuda/{norm.cu => norm.cu.inc} (99%) rename ggml/src/ggml-cuda/{norm.cuh => norm.cuh.inc} (96%) rename ggml/src/ggml-cuda/{opt-step-adamw.cu => opt-step-adamw.cu.inc} (97%) rename ggml/src/ggml-cuda/{opt-step-adamw.cuh => opt-step-adamw.cuh.inc} (83%) rename ggml/src/ggml-cuda/{opt-step-sgd.cu => opt-step-sgd.cu.inc} (96%) rename ggml/src/ggml-cuda/{opt-step-sgd.cuh => opt-step-sgd.cuh.inc} (82%) rename ggml/src/ggml-cuda/{out-prod.cu => out-prod.cu.inc} (98%) rename ggml/src/ggml-cuda/{out-prod.cuh => out-prod.cuh.inc} (75%) rename ggml/src/ggml-cuda/{pad.cu => pad.cu.inc} (99%) rename ggml/src/ggml-cuda/{pad.cuh => pad.cuh.inc} (80%) rename ggml/src/ggml-cuda/{pad_reflect_1d.cu => pad_reflect_1d.cu.inc} (98%) rename ggml/src/ggml-cuda/{pad_reflect_1d.cuh => pad_reflect_1d.cuh.inc} (83%) rename ggml/src/ggml-cuda/{pool2d.cu => pool2d.cu.inc} (99%) rename ggml/src/ggml-cuda/{pool2d.cuh => pool2d.cuh.inc} (81%) rename ggml/src/ggml-cuda/{quantize.cu => quantize.cu.inc} (99%) rename ggml/src/ggml-cuda/{quantize.cuh => quantize.cuh.inc} (97%) rename ggml/src/ggml-cuda/{reduce_rows.cuh => reduce_rows.cuh.inc} (97%) rename ggml/src/ggml-cuda/{roll.cu => roll.cu.inc} (97%) rename ggml/src/ggml-cuda/{roll.cuh => roll.cuh.inc} (81%) rename ggml/src/ggml-cuda/{rope.cu => rope.cu.inc} (99%) rename ggml/src/ggml-cuda/{rope.cuh => rope.cuh.inc} (92%) rename ggml/src/ggml-cuda/{scale.cu => scale.cu.inc} (98%) rename ggml/src/ggml-cuda/{scale.cuh => scale.cuh.inc} (81%) rename ggml/src/ggml-cuda/{set-rows.cu => set-rows.cu.inc} (99%) rename ggml/src/ggml-cuda/{set-rows.cuh => set-rows.cuh.inc} (83%) rename ggml/src/ggml-cuda/{set.cu => set.cu.inc} (96%) rename ggml/src/ggml-cuda/{set.cuh => set.cuh.inc} (82%) rename ggml/src/ggml-cuda/{softcap.cu => softcap.cu.inc} (97%) rename ggml/src/ggml-cuda/{softcap.cuh => softcap.cuh.inc} (83%) rename ggml/src/ggml-cuda/{softmax.cu => softmax.cu.inc} (99%) rename ggml/src/ggml-cuda/{softmax.cuh => softmax.cuh.inc} (88%) rename ggml/src/ggml-cuda/{solve_tri.cu => solve_tri.cu.inc} (98%) rename ggml/src/ggml-cuda/{solve_tri.cuh => solve_tri.cuh.inc} (75%) rename ggml/src/ggml-cuda/{ssm-conv.cu => ssm-conv.cu.inc} (99%) rename ggml/src/ggml-cuda/{ssm-conv.cuh => ssm-conv.cuh.inc} (85%) rename ggml/src/ggml-cuda/{ssm-scan.cu => ssm-scan.cu.inc} (99%) rename ggml/src/ggml-cuda/{ssm-scan.cuh => ssm-scan.cuh.inc} (75%) rename ggml/src/ggml-cuda/{sum.cu => sum.cu.inc} (96%) rename ggml/src/ggml-cuda/{sum.cuh => sum.cuh.inc} (87%) rename ggml/src/ggml-cuda/{sumrows.cu => sumrows.cu.inc} (96%) rename ggml/src/ggml-cuda/{sumrows.cuh => sumrows.cuh.inc} (87%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_1-ncols2_16.cu => fattn-mma-f16-instance-ncols1_1-ncols2_16.cu.inc} (77%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_1-ncols2_32.cu => fattn-mma-f16-instance-ncols1_1-ncols2_32.cu.inc} (82%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_1-ncols2_8.cu => fattn-mma-f16-instance-ncols1_1-ncols2_8.cu.inc} (91%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_16-ncols2_1.cu => fattn-mma-f16-instance-ncols1_16-ncols2_1.cu.inc} (90%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_16-ncols2_2.cu => fattn-mma-f16-instance-ncols1_16-ncols2_2.cu.inc} (90%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_16-ncols2_4.cu => fattn-mma-f16-instance-ncols1_16-ncols2_4.cu.inc} (92%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_2-ncols2_16.cu => fattn-mma-f16-instance-ncols1_2-ncols2_16.cu.inc} (77%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_2-ncols2_32.cu => fattn-mma-f16-instance-ncols1_2-ncols2_32.cu.inc} (82%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_2-ncols2_4.cu => fattn-mma-f16-instance-ncols1_2-ncols2_4.cu.inc} (91%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_2-ncols2_8.cu => fattn-mma-f16-instance-ncols1_2-ncols2_8.cu.inc} (91%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_32-ncols2_1.cu => fattn-mma-f16-instance-ncols1_32-ncols2_1.cu.inc} (90%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_32-ncols2_2.cu => fattn-mma-f16-instance-ncols1_32-ncols2_2.cu.inc} (90%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_4-ncols2_16.cu => fattn-mma-f16-instance-ncols1_4-ncols2_16.cu.inc} (77%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_4-ncols2_2.cu => fattn-mma-f16-instance-ncols1_4-ncols2_2.cu.inc} (90%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_4-ncols2_4.cu => fattn-mma-f16-instance-ncols1_4-ncols2_4.cu.inc} (91%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_4-ncols2_8.cu => fattn-mma-f16-instance-ncols1_4-ncols2_8.cu.inc} (91%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_64-ncols2_1.cu => fattn-mma-f16-instance-ncols1_64-ncols2_1.cu.inc} (90%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_8-ncols2_1.cu => fattn-mma-f16-instance-ncols1_8-ncols2_1.cu.inc} (90%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_8-ncols2_2.cu => fattn-mma-f16-instance-ncols1_8-ncols2_2.cu.inc} (90%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_8-ncols2_4.cu => fattn-mma-f16-instance-ncols1_8-ncols2_4.cu.inc} (91%) rename ggml/src/ggml-cuda/template-instances/{fattn-mma-f16-instance-ncols1_8-ncols2_8.cu => fattn-mma-f16-instance-ncols1_8-ncols2_8.cu.inc} (91%) delete mode 100644 ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq112-dv112.cu rename ggml/src/{ggml-sycl/template-instances/fattn-tile-instance-dkq112-dv112.cpp => ggml-cuda/template-instances/fattn-tile-instance-dkq112-dv112.cu.inc} (78%) rename ggml/src/ggml-cuda/template-instances/{fattn-tile-instance-dkq128-dv128.cu => fattn-tile-instance-dkq128-dv128.cu.inc} (78%) delete mode 100644 ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq256-dv256.cu rename ggml/src/{ggml-sycl/template-instances/fattn-tile-instance-dkq256-dv256.cpp => ggml-cuda/template-instances/fattn-tile-instance-dkq256-dv256.cu.inc} (78%) rename ggml/src/ggml-cuda/template-instances/{fattn-tile-instance-dkq320-dv256.cu => fattn-tile-instance-dkq320-dv256.cu.inc} (78%) rename ggml/src/ggml-cuda/template-instances/{fattn-tile-instance-dkq40-dv40.cu => fattn-tile-instance-dkq40-dv40.cu.inc} (77%) rename ggml/src/ggml-cuda/template-instances/{fattn-tile-instance-dkq512-dv512.cu => fattn-tile-instance-dkq512-dv512.cu.inc} (78%) rename ggml/src/ggml-cuda/template-instances/{fattn-tile-instance-dkq576-dv512.cu => fattn-tile-instance-dkq576-dv512.cu.inc} (78%) rename ggml/src/ggml-cuda/template-instances/{fattn-tile-instance-dkq64-dv64.cu => fattn-tile-instance-dkq64-dv64.cu.inc} (77%) rename ggml/src/ggml-cuda/template-instances/{fattn-tile-instance-dkq72-dv72.cu => fattn-tile-instance-dkq72-dv72.cu.inc} (77%) rename ggml/src/ggml-cuda/template-instances/{fattn-tile-instance-dkq80-dv80.cu => fattn-tile-instance-dkq80-dv80.cu.inc} (77%) rename ggml/src/ggml-cuda/template-instances/{fattn-tile-instance-dkq96-dv96.cu => fattn-tile-instance-dkq96-dv96.cu.inc} (77%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-bf16-bf16.cu => fattn-vec-instance-bf16-bf16.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-bf16-f16.cu => fattn-vec-instance-bf16-f16.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-bf16-q4_0.cu => fattn-vec-instance-bf16-q4_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-bf16-q4_1.cu => fattn-vec-instance-bf16-q4_1.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-bf16-q5_0.cu => fattn-vec-instance-bf16-q5_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-bf16-q5_1.cu => fattn-vec-instance-bf16-q5_1.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-bf16-q8_0.cu => fattn-vec-instance-bf16-q8_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-f16-bf16.cu => fattn-vec-instance-f16-bf16.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-f16-f16.cu => fattn-vec-instance-f16-f16.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-f16-q4_0.cu => fattn-vec-instance-f16-q4_0.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-f16-q4_1.cu => fattn-vec-instance-f16-q4_1.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-f16-q5_0.cu => fattn-vec-instance-f16-q5_0.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-f16-q5_1.cu => fattn-vec-instance-f16-q5_1.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-f16-q8_0.cu => fattn-vec-instance-f16-q8_0.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_0-bf16.cu => fattn-vec-instance-q4_0-bf16.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_0-f16.cu => fattn-vec-instance-q4_0-f16.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_0-q4_0.cu => fattn-vec-instance-q4_0-q4_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_0-q4_1.cu => fattn-vec-instance-q4_0-q4_1.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_0-q5_0.cu => fattn-vec-instance-q4_0-q5_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_0-q5_1.cu => fattn-vec-instance-q4_0-q5_1.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_0-q8_0.cu => fattn-vec-instance-q4_0-q8_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_1-bf16.cu => fattn-vec-instance-q4_1-bf16.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_1-f16.cu => fattn-vec-instance-q4_1-f16.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_1-q4_0.cu => fattn-vec-instance-q4_1-q4_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_1-q4_1.cu => fattn-vec-instance-q4_1-q4_1.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_1-q5_0.cu => fattn-vec-instance-q4_1-q5_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_1-q5_1.cu => fattn-vec-instance-q4_1-q5_1.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q4_1-q8_0.cu => fattn-vec-instance-q4_1-q8_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_0-bf16.cu => fattn-vec-instance-q5_0-bf16.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_0-f16.cu => fattn-vec-instance-q5_0-f16.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_0-q4_0.cu => fattn-vec-instance-q5_0-q4_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_0-q4_1.cu => fattn-vec-instance-q5_0-q4_1.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_0-q5_0.cu => fattn-vec-instance-q5_0-q5_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_0-q5_1.cu => fattn-vec-instance-q5_0-q5_1.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_0-q8_0.cu => fattn-vec-instance-q5_0-q8_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_1-bf16.cu => fattn-vec-instance-q5_1-bf16.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_1-f16.cu => fattn-vec-instance-q5_1-f16.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_1-q4_0.cu => fattn-vec-instance-q5_1-q4_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_1-q4_1.cu => fattn-vec-instance-q5_1-q4_1.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_1-q5_0.cu => fattn-vec-instance-q5_1-q5_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_1-q5_1.cu => fattn-vec-instance-q5_1-q5_1.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q5_1-q8_0.cu => fattn-vec-instance-q5_1-q8_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q8_0-bf16.cu => fattn-vec-instance-q8_0-bf16.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q8_0-f16.cu => fattn-vec-instance-q8_0-f16.cu.inc} (88%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q8_0-q4_0.cu => fattn-vec-instance-q8_0-q4_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q8_0-q4_1.cu => fattn-vec-instance-q8_0-q4_1.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q8_0-q5_0.cu => fattn-vec-instance-q8_0-q5_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q8_0-q5_1.cu => fattn-vec-instance-q8_0-q5_1.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{fattn-vec-instance-q8_0-q8_0.cu => fattn-vec-instance-q8_0-q8_0.cu.inc} (89%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_1.cu => mmf-instance-ncols_1.cu.inc} (79%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_10.cu => mmf-instance-ncols_10.cu.inc} (80%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_11.cu => mmf-instance-ncols_11.cu.inc} (80%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_12.cu => mmf-instance-ncols_12.cu.inc} (80%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_13.cu => mmf-instance-ncols_13.cu.inc} (80%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_14.cu => mmf-instance-ncols_14.cu.inc} (80%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_15.cu => mmf-instance-ncols_15.cu.inc} (80%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_16.cu => mmf-instance-ncols_16.cu.inc} (80%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_2.cu => mmf-instance-ncols_2.cu.inc} (79%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_3.cu => mmf-instance-ncols_3.cu.inc} (79%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_4.cu => mmf-instance-ncols_4.cu.inc} (79%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_5.cu => mmf-instance-ncols_5.cu.inc} (79%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_6.cu => mmf-instance-ncols_6.cu.inc} (79%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_7.cu => mmf-instance-ncols_7.cu.inc} (79%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_8.cu => mmf-instance-ncols_8.cu.inc} (79%) rename ggml/src/ggml-cuda/template-instances/{mmf-instance-ncols_9.cu => mmf-instance-ncols_9.cu.inc} (79%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-iq1_s.cu => mmq-instance-iq1_s.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-iq2_s.cu => mmq-instance-iq2_s.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-iq2_xs.cu => mmq-instance-iq2_xs.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-iq2_xxs.cu => mmq-instance-iq2_xxs.cu.inc} (82%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-iq3_s.cu => mmq-instance-iq3_s.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-iq3_xxs.cu => mmq-instance-iq3_xxs.cu.inc} (82%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-iq4_nl.cu => mmq-instance-iq4_nl.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-iq4_xs.cu => mmq-instance-iq4_xs.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-mxfp4.cu => mmq-instance-mxfp4.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-nvfp4.cu => mmq-instance-nvfp4.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-q1_0.cu => mmq-instance-q1_0.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-q2_k.cu => mmq-instance-q2_k.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-q3_k.cu => mmq-instance-q3_k.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-q4_0.cu => mmq-instance-q4_0.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-q4_1.cu => mmq-instance-q4_1.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-q4_k.cu => mmq-instance-q4_k.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-q5_0.cu => mmq-instance-q5_0.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-q5_1.cu => mmq-instance-q5_1.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-q5_k.cu => mmq-instance-q5_k.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-q6_k.cu => mmq-instance-q6_k.cu.inc} (81%) rename ggml/src/ggml-cuda/template-instances/{mmq-instance-q8_0.cu => mmq-instance-q8_0.cu.inc} (81%) rename ggml/src/ggml-cuda/{top-k.cu => top-k.cu.inc} (98%) rename ggml/src/ggml-cuda/{top-k.cuh => top-k.cuh.inc} (75%) rename ggml/src/ggml-cuda/{topk-moe.cu => topk-moe.cu.inc} (99%) rename ggml/src/ggml-cuda/{topk-moe.cuh => topk-moe.cuh.inc} (95%) rename ggml/src/ggml-cuda/{tri.cu => tri.cu.inc} (98%) rename ggml/src/ggml-cuda/{tri.cuh => tri.cuh.inc} (80%) rename ggml/src/ggml-cuda/{tsembd.cu => tsembd.cu.inc} (98%) rename ggml/src/ggml-cuda/{tsembd.cuh => tsembd.cuh.inc} (84%) rename ggml/src/ggml-cuda/{unary.cu => unary.cu.inc} (99%) rename ggml/src/ggml-cuda/{unary.cuh => unary.cuh.inc} (99%) rename ggml/src/ggml-cuda/{upscale.cu => upscale.cu.inc} (99%) rename ggml/src/ggml-cuda/{upscale.cuh => upscale.cuh.inc} (81%) rename ggml/src/ggml-cuda/{vecdotq.cuh => vecdotq.cuh.inc} (99%) rename ggml/src/ggml-cuda/vendors/{cuda.h => cuda.h.inc} (97%) rename ggml/src/ggml-cuda/vendors/{hip.h => hip.h.inc} (100%) rename ggml/src/ggml-cuda/vendors/{musa.h => musa.h.inc} (99%) rename ggml/src/ggml-cuda/{wkv.cu => wkv.cu.inc} (99%) rename ggml/src/ggml-cuda/{wkv.cuh => wkv.cuh.inc} (88%) rename ggml/src/ggml-hexagon/{ggml-hexagon.cpp => ggml-hexagon.cpp.inc} (94%) rename ggml/src/ggml-hexagon/{htp-drv.cpp => htp-drv.cpp.inc} (78%) delete mode 100644 ggml/src/ggml-hexagon/htp-drv.h rename ggml/src/ggml-hexagon/htp/{act-ops.c => act-ops.inc} (99%) rename ggml/src/ggml-hexagon/htp/{argsort-ops.c => argsort-ops.inc} (97%) rename ggml/src/ggml-hexagon/htp/{binary-ops.c => binary-ops.inc} (99%) rename ggml/src/ggml-hexagon/htp/{cpy-ops.c => cpy-ops.inc} (98%) rename ggml/src/ggml-hexagon/htp/{cumsum-ops.c => cumsum-ops.inc} (98%) delete mode 100644 ggml/src/ggml-hexagon/htp/diag-ops.c rename ggml/src/ggml-hexagon/htp/{hvx-copy.h => diag-ops.inc} (63%) rename ggml/src/ggml-hexagon/htp/{fill-ops.c => fill-ops.inc} (97%) rename ggml/src/ggml-hexagon/htp/{flash-attn-ops.c => flash-attn-ops.inc} (83%) rename ggml/src/ggml-hexagon/htp/{get-rows-ops.c => get-rows-ops.inc} (96%) rename ggml/src/ggml-hexagon/htp/{hex-dma.h => hex-dma-defs.inc} (100%) rename ggml/src/ggml-hexagon/htp/{hex-dma.c => hex-dma.inc} (98%) delete mode 100644 ggml/src/ggml-hexagon/htp/hex-dump.h delete mode 100644 ggml/src/ggml-hexagon/htp/hex-fastdiv.h rename ggml/src/ggml-hexagon/htp/{hex-utils.h => hex-utils-defs.inc} (51%) rename ggml/src/ggml-hexagon/htp/{hmx-flash-attn-ops.c => hmx-flash-attn-ops.inc} (77%) rename ggml/src/ggml-hexagon/htp/{hmx-matmul-ops.c => hmx-matmul-ops.inc} (83%) delete mode 100644 ggml/src/ggml-hexagon/htp/hmx-ops.h delete mode 100644 ggml/src/ggml-hexagon/htp/hmx-profile.h rename ggml/src/ggml-hexagon/htp/{hmx-queue.h => hmx-queue-defs.inc} (99%) rename ggml/src/ggml-hexagon/htp/{hmx-queue.c => hmx-queue.inc} (99%) delete mode 100644 ggml/src/ggml-hexagon/htp/hmx-utils.h rename ggml/src/ggml-hexagon/htp/{htp-ctx.h => htp-ctx-defs.inc} (96%) rename ggml/src/ggml-hexagon/htp/{htp-ops.h => htp-ops-defs.inc} (97%) delete mode 100644 ggml/src/ggml-hexagon/htp/hvx-arith.h rename ggml/src/ggml-hexagon/htp/{hvx-base.h => hvx-base-defs.inc} (99%) delete mode 100644 ggml/src/ggml-hexagon/htp/hvx-div.h delete mode 100644 ggml/src/ggml-hexagon/htp/hvx-dump.h delete mode 100644 ggml/src/ggml-hexagon/htp/hvx-exp.h delete mode 100644 ggml/src/ggml-hexagon/htp/hvx-floor.h delete mode 100644 ggml/src/ggml-hexagon/htp/hvx-inverse.h delete mode 100644 ggml/src/ggml-hexagon/htp/hvx-reduce.h delete mode 100644 ggml/src/ggml-hexagon/htp/hvx-scale.h delete mode 100644 ggml/src/ggml-hexagon/htp/hvx-sigmoid.h delete mode 100644 ggml/src/ggml-hexagon/htp/hvx-sqrt.h rename ggml/src/ggml-hexagon/htp/{hvx-types.h => hvx-types-defs.inc} (100%) create mode 100644 ggml/src/ggml-hexagon/htp/hvx-utils-defs.inc delete mode 100644 ggml/src/ggml-hexagon/htp/hvx-utils.h rename ggml/src/ggml-hexagon/htp/{main.c => main.inc} (98%) rename ggml/src/ggml-hexagon/htp/{matmul-ops.c => matmul-ops.inc} (96%) rename ggml/src/ggml-hexagon/htp/{repeat-ops.c => repeat-ops.inc} (97%) rename ggml/src/ggml-hexagon/htp/{rope-ops.c => rope-ops.inc} (94%) rename ggml/src/ggml-hexagon/htp/{set-rows-ops.c => set-rows-ops.inc} (97%) rename ggml/src/ggml-hexagon/htp/{softmax-ops.c => softmax-ops.inc} (92%) rename ggml/src/ggml-hexagon/htp/{solve-tri-ops.c => solve-tri-ops.inc} (98%) rename ggml/src/ggml-hexagon/htp/{ssm-conv.c => ssm-conv.inc} (98%) rename ggml/src/ggml-hexagon/htp/{sum-rows-ops.c => sum-rows-ops.inc} (95%) delete mode 100644 ggml/src/ggml-hexagon/htp/unary-ops.c create mode 100644 ggml/src/ggml-hexagon/htp/unary-ops.inc delete mode 100644 ggml/src/ggml-hexagon/htp/vtcm-utils.h rename ggml/src/ggml-hexagon/htp/{worker-pool.h => worker-pool-defs.inc} (100%) rename ggml/src/ggml-hexagon/htp/{worker-pool.c => worker-pool.inc} (99%) delete mode 100644 ggml/src/ggml-hexagon/libdl.h delete mode 100644 ggml/src/ggml-hexagon/op-desc.h rename ggml/src/{ggml-impl.h => ggml-impl-defs.inc} (97%) rename ggml/src/ggml-metal/{ggml-metal-common.h => ggml-metal-common-defs.inc} (100%) rename ggml/src/ggml-metal/{ggml-metal-common.cpp => ggml-metal-common.inc} (99%) rename ggml/src/ggml-metal/{ggml-metal-context.h => ggml-metal-context-defs.inc} (97%) rename ggml/src/ggml-metal/{ggml-metal-device.h => ggml-metal-device-defs.inc} (99%) rename ggml/src/ggml-metal/{ggml-metal-device.cpp => ggml-metal-device.inc} (99%) rename ggml/src/ggml-metal/{ggml-metal-impl.h => ggml-metal-impl-defs.inc} (100%) rename ggml/src/ggml-metal/{ggml-metal-ops.h => ggml-metal-ops-defs.inc} (99%) rename ggml/src/ggml-metal/{ggml-metal-ops.cpp => ggml-metal-ops.inc} (99%) rename ggml/src/ggml-metal/{ggml-metal.cpp => ggml-metal.inc} (99%) rename ggml/src/ggml-musa/{mudnn.cu => mudnn.cu.inc} (99%) rename ggml/src/ggml-musa/{mudnn.cuh => mudnn.cuh.inc} (82%) rename ggml/src/ggml-opencl/{ggml-opencl.cpp => ggml-opencl.cpp.inc} (99%) delete mode 100644 ggml/src/ggml-openvino/ggml-decoder.cpp delete mode 100644 ggml/src/ggml-openvino/ggml-decoder.h delete mode 100644 ggml/src/ggml-openvino/ggml-openvino-extra.cpp delete mode 100644 ggml/src/ggml-openvino/ggml-openvino-extra.h delete mode 100644 ggml/src/ggml-openvino/ggml-openvino.cpp create mode 100644 ggml/src/ggml-openvino/ggml-openvino.cpp.inc delete mode 100644 ggml/src/ggml-openvino/ggml-quants.cpp delete mode 100644 ggml/src/ggml-openvino/ggml-quants.h delete mode 100644 ggml/src/ggml-openvino/openvino/decoder.h delete mode 100644 ggml/src/ggml-openvino/openvino/frontend.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/frontend.h delete mode 100644 ggml/src/ggml-openvino/openvino/input_model.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/input_model.h delete mode 100644 ggml/src/ggml-openvino/openvino/node_context.h delete mode 100644 ggml/src/ggml-openvino/openvino/op/cont.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/cpy.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/get_rows.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/glu_geglu.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/glu_swiglu.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/mulmat.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/permute.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/reshape.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/rms_norm.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/rope.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/scale.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/set_rows.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/softmax.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/transpose.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/unary_gelu.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/unary_silu.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op/view.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op_table.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/op_table.h delete mode 100644 ggml/src/ggml-openvino/openvino/pass/fuse_to_sdpa.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/pass/fuse_to_sdpa.h delete mode 100644 ggml/src/ggml-openvino/openvino/pass/mark_decompression_convert_constant_folding.h delete mode 100644 ggml/src/ggml-openvino/openvino/pass/squeeze_matmul.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/pass/squeeze_matmul.h delete mode 100644 ggml/src/ggml-openvino/openvino/rt_info/weightless_caching_attributes.hpp delete mode 100644 ggml/src/ggml-openvino/openvino/translate_session.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/translate_session.h delete mode 100644 ggml/src/ggml-openvino/openvino/utils.cpp delete mode 100644 ggml/src/ggml-openvino/openvino/utils.h delete mode 100644 ggml/src/ggml-openvino/utils.cpp delete mode 100644 ggml/src/ggml-openvino/utils.h delete mode 100644 ggml/src/ggml-opt.cpp rename ggml/src/{ggml-quants.h => ggml-quants-defs.inc} (99%) delete mode 100644 ggml/src/ggml-quants.c rename ggml/src/ggml-rpc/{ggml-rpc.cpp => ggml-rpc.cpp.inc} (77%) delete mode 100644 ggml/src/ggml-rpc/transport.cpp delete mode 100644 ggml/src/ggml-rpc/transport.h delete mode 100644 ggml/src/ggml-sycl/add-id.hpp rename ggml/src/ggml-sycl/{add-id.cpp => add-id.inc} (98%) create mode 100644 ggml/src/ggml-sycl/backend-defs.inc delete mode 100644 ggml/src/ggml-sycl/backend.hpp delete mode 100644 ggml/src/ggml-sycl/binbcast.hpp rename ggml/src/ggml-sycl/{binbcast.cpp => binbcast.inc} (96%) rename ggml/src/ggml-sycl/{common.hpp => common-defs.inc} (97%) rename ggml/src/ggml-sycl/{common.cpp => common.inc} (96%) delete mode 100644 ggml/src/ggml-sycl/concat.hpp rename ggml/src/ggml-sycl/{concat.cpp => concat.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/conv.hpp rename ggml/src/ggml-sycl/{conv.cpp => conv.inc} (99%) rename ggml/src/ggml-sycl/{convert.hpp => convert-defs.inc} (98%) rename ggml/src/ggml-sycl/{convert.cpp => convert.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/count-equal.hpp rename ggml/src/ggml-sycl/{count-equal.cpp => count-equal.inc} (97%) rename ggml/src/ggml-sycl/{cpy.hpp => cpy-defs.inc} (99%) rename ggml/src/ggml-sycl/{cpy.cpp => cpy.inc} (99%) rename ggml/src/ggml-sycl/{dequantize.hpp => dequantize-defs.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/dmmv.hpp rename ggml/src/ggml-sycl/{dmmv.cpp => dmmv.inc} (99%) rename ggml/src/ggml-sycl/dpct/{helper.hpp => helper-defs.inc} (99%) rename ggml/src/ggml-sycl/{element_wise.hpp => element_wise-defs.inc} (98%) rename ggml/src/ggml-sycl/{element_wise.cpp => element_wise.inc} (99%) rename ggml/src/ggml-sycl/{fattn-common.hpp => fattn-common-defs.inc} (99%) rename ggml/src/ggml-sycl/{fattn-tile.hpp => fattn-tile-defs.inc} (99%) rename ggml/src/ggml-sycl/{fattn-tile.cpp => fattn-tile.inc} (94%) rename ggml/src/ggml-sycl/{fattn-vec.hpp => fattn-vec-defs.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/fattn.hpp rename ggml/src/ggml-sycl/{fattn.cpp => fattn.inc} (97%) delete mode 100644 ggml/src/ggml-sycl/gated_delta_net.hpp rename ggml/src/ggml-sycl/{gated_delta_net.cpp => gated_delta_net.inc} (99%) rename ggml/src/ggml-sycl/{gemm.hpp => gemm-defs.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/getrows.hpp rename ggml/src/ggml-sycl/{getrows.cpp => getrows.inc} (98%) rename ggml/src/ggml-sycl/{ggml-sycl.cpp => ggml-sycl.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/gla.hpp rename ggml/src/ggml-sycl/{gla.cpp => gla.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/im2col.hpp rename ggml/src/ggml-sycl/{im2col.cpp => im2col.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/mmq.hpp rename ggml/src/ggml-sycl/{mmq.cpp => mmq.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/mmvq.hpp rename ggml/src/ggml-sycl/{mmvq.cpp => mmvq.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/norm.hpp rename ggml/src/ggml-sycl/{norm.cpp => norm.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/outprod.hpp rename ggml/src/ggml-sycl/{outprod.cpp => outprod.inc} (98%) delete mode 100644 ggml/src/ggml-sycl/pad.hpp rename ggml/src/ggml-sycl/{pad.cpp => pad.inc} (98%) delete mode 100644 ggml/src/ggml-sycl/pad_reflect_1d.hpp rename ggml/src/ggml-sycl/{pad_reflect_1d.cpp => pad_reflect_1d.inc} (98%) rename ggml/src/ggml-sycl/{presets.hpp => presets-defs.inc} (100%) rename ggml/src/ggml-sycl/{quantize.hpp => quantize-defs.inc} (99%) rename ggml/src/ggml-sycl/{quants.hpp => quants-defs.inc} (97%) delete mode 100644 ggml/src/ggml-sycl/repeat_back.hpp rename ggml/src/ggml-sycl/{repeat_back.cpp => repeat_back.inc} (98%) delete mode 100644 ggml/src/ggml-sycl/roll.hpp rename ggml/src/ggml-sycl/{roll.cpp => roll.inc} (98%) delete mode 100644 ggml/src/ggml-sycl/rope.hpp rename ggml/src/ggml-sycl/{rope.cpp => rope.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/set.hpp rename ggml/src/ggml-sycl/{set.cpp => set.inc} (97%) delete mode 100644 ggml/src/ggml-sycl/set_rows.hpp rename ggml/src/ggml-sycl/{set_rows.cpp => set_rows.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/softmax.hpp rename ggml/src/ggml-sycl/{softmax.cpp => softmax.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/ssm_conv.hpp rename ggml/src/ggml-sycl/{ssm_conv.cpp => ssm_conv.inc} (98%) delete mode 100644 ggml/src/ggml-sycl/sycl_hw.hpp rename ggml/src/ggml-sycl/{sycl_hw.cpp => sycl_hw.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq128-dv128.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq40-dv40.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq512-dv512.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq576-dv512.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq64-dv64.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq72-dv72.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq80-dv80.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq96-dv96.cpp create mode 100644 ggml/src/ggml-sycl/template-instances/fattn-tile-instances.inc delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-f16.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q4_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q4_1.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q5_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q5_1.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q8_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-f16.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q4_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q4_1.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q5_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q5_1.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q8_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-f16.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q4_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q4_1.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q5_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q5_1.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q8_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-f16.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q4_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q4_1.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q5_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q5_1.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q8_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-f16.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q4_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q4_1.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q5_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q5_1.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q8_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-f16.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q4_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q4_1.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q5_0.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q5_1.cpp delete mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q8_0.cpp create mode 100644 ggml/src/ggml-sycl/template-instances/fattn-vec-instances.inc delete mode 100644 ggml/src/ggml-sycl/tsembd.hpp rename ggml/src/ggml-sycl/{tsembd.cpp => tsembd.inc} (98%) rename ggml/src/ggml-sycl/{type.hpp => type-defs.inc} (100%) delete mode 100644 ggml/src/ggml-sycl/upscale.hpp rename ggml/src/ggml-sycl/{upscale.cpp => upscale.inc} (99%) rename ggml/src/ggml-sycl/{vecdotq.hpp => vecdotq-defs.inc} (99%) delete mode 100644 ggml/src/ggml-sycl/wkv.hpp rename ggml/src/ggml-sycl/{wkv.cpp => wkv.inc} (99%) create mode 100644 ggml/src/ggml-threading-rust/Cargo.toml create mode 100644 ggml/src/ggml-threading-rust/src/lib.rs delete mode 100644 ggml/src/ggml-threading.cpp delete mode 100644 ggml/src/ggml-threading.h delete mode 100644 ggml/src/ggml-virtgpu/apir_cs_ggml-rpc-front.cpp delete mode 100644 ggml/src/ggml-virtgpu/backend/apir_cs_ggml-rpc-back.cpp delete mode 100644 ggml/src/ggml-virtgpu/backend/backend-convert.h delete mode 100644 ggml/src/ggml-virtgpu/backend/backend-dispatched-backend.cpp delete mode 100644 ggml/src/ggml-virtgpu/backend/backend-dispatched-buffer-type.cpp delete mode 100644 ggml/src/ggml-virtgpu/backend/backend-dispatched-buffer.cpp delete mode 100644 ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp create mode 100644 ggml/src/ggml-virtgpu/backend/backend-dispatched-rust/Cargo.toml create mode 100644 ggml/src/ggml-virtgpu/backend/backend-dispatched-rust/src/lib.rs delete mode 100644 ggml/src/ggml-virtgpu/backend/backend-dispatched.cpp delete mode 100644 ggml/src/ggml-virtgpu/backend/backend-dispatched.gen.h delete mode 100644 ggml/src/ggml-virtgpu/backend/backend-dispatched.h delete mode 100644 ggml/src/ggml-virtgpu/backend/backend-virgl-apir.h delete mode 100644 ggml/src/ggml-virtgpu/backend/backend.cpp create mode 100644 ggml/src/ggml-virtgpu/backend/backend.cpp.inc delete mode 100644 ggml/src/ggml-virtgpu/backend/shared/api_remoting.h delete mode 100644 ggml/src/ggml-virtgpu/backend/shared/apir_backend.gen.h delete mode 100644 ggml/src/ggml-virtgpu/backend/shared/apir_backend.h delete mode 100644 ggml/src/ggml-virtgpu/backend/shared/apir_cs.h delete mode 100644 ggml/src/ggml-virtgpu/backend/shared/apir_cs_ggml.h delete mode 100644 ggml/src/ggml-virtgpu/backend/shared/apir_cs_rpc.h delete mode 100644 ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp delete mode 100644 ggml/src/ggml-virtgpu/ggml-backend-buffer.cpp delete mode 100644 ggml/src/ggml-virtgpu/ggml-backend-device.cpp delete mode 100644 ggml/src/ggml-virtgpu/ggml-backend-reg.cpp delete mode 100644 ggml/src/ggml-virtgpu/ggml-backend.cpp delete mode 100644 ggml/src/ggml-virtgpu/ggml-remoting.h delete mode 100644 ggml/src/ggml-virtgpu/include/apir_hw.h delete mode 100644 ggml/src/ggml-virtgpu/virtgpu-apir.h delete mode 100644 ggml/src/ggml-virtgpu/virtgpu-forward-backend.cpp delete mode 100644 ggml/src/ggml-virtgpu/virtgpu-forward-buffer-type.cpp delete mode 100644 ggml/src/ggml-virtgpu/virtgpu-forward-buffer.cpp delete mode 100644 ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp delete mode 100644 ggml/src/ggml-virtgpu/virtgpu-forward-impl.h delete mode 100644 ggml/src/ggml-virtgpu/virtgpu-forward.gen.h delete mode 100644 ggml/src/ggml-virtgpu/virtgpu-shm.cpp delete mode 100644 ggml/src/ggml-virtgpu/virtgpu-shm.h delete mode 100644 ggml/src/ggml-virtgpu/virtgpu-utils.cpp delete mode 100644 ggml/src/ggml-virtgpu/virtgpu-utils.h delete mode 100644 ggml/src/ggml-virtgpu/virtgpu.cpp create mode 100644 ggml/src/ggml-virtgpu/virtgpu.cpp.inc delete mode 100644 ggml/src/ggml-virtgpu/virtgpu.h rename ggml/src/ggml-vulkan/{ggml-vulkan.cpp => ggml-vulkan.cpp.inc} (99%) rename ggml/src/ggml-vulkan/vulkan-shaders/{vulkan-shaders-gen.cpp => vulkan-shaders-gen.cpp.inc} (100%) delete mode 100644 ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp rename ggml/src/ggml-webgpu/{ggml-webgpu.cpp => ggml-webgpu.cpp.inc} (58%) delete mode 100644 ggml/src/ggml-webgpu/pre_wgsl.hpp delete mode 100644 ggml/src/ggml-zdnn/common.hpp rename ggml/src/ggml-zdnn/{ggml-zdnn.cpp => ggml-zdnn.cpp.inc} (69%) delete mode 100644 ggml/src/ggml-zdnn/mmf.cpp delete mode 100644 ggml/src/ggml-zdnn/mmf.hpp delete mode 100644 ggml/src/ggml-zdnn/utils.cpp delete mode 100644 ggml/src/ggml-zdnn/utils.hpp rename ggml/src/ggml-zendnn/{ggml-zendnn.cpp => ggml-zendnn.cpp.inc} (99%) delete mode 100644 ggml/src/ggml.c create mode 100644 ggml/src/ggml.c.inc delete mode 100644 ggml/src/ggml.cpp delete mode 100644 ggml/src/gguf.cpp delete mode 100644 include/llama-cpp.h rename include/{llama.h => llama.h.inc} (99%) create mode 100644 pocs/vdot/q8dot-rust/Cargo.toml create mode 100644 pocs/vdot/q8dot-rust/build.rs create mode 100644 pocs/vdot/q8dot-rust/src/lib.rs create mode 100644 pocs/vdot/q8dot-rust/src/main.rs delete mode 100644 pocs/vdot/q8dot.cpp create mode 100644 pocs/vdot/vdot-rust/Cargo.toml create mode 100644 pocs/vdot/vdot-rust/build.rs create mode 100644 pocs/vdot/vdot-rust/src/lib.rs create mode 100644 pocs/vdot/vdot-rust/src/main.rs delete mode 100644 pocs/vdot/vdot.cpp create mode 100755 scripts/rust-port-status.sh delete mode 100644 src/llama-adapter.cpp delete mode 100644 src/llama-adapter.h delete mode 100644 src/llama-arch.cpp delete mode 100644 src/llama-arch.h delete mode 100644 src/llama-batch.cpp delete mode 100644 src/llama-batch.h delete mode 100644 src/llama-chat.cpp delete mode 100644 src/llama-chat.h delete mode 100644 src/llama-context.cpp delete mode 100644 src/llama-context.h create mode 100644 src/llama-cparams-rust/Cargo.toml create mode 100644 src/llama-cparams-rust/src/lib.rs delete mode 100644 src/llama-cparams.cpp delete mode 100644 src/llama-cparams.h delete mode 100644 src/llama-ext.h delete mode 100644 src/llama-grammar.cpp delete mode 100644 src/llama-grammar.h delete mode 100644 src/llama-graph.cpp delete mode 100644 src/llama-graph.h delete mode 100644 src/llama-hparams.cpp delete mode 100644 src/llama-hparams.h delete mode 100644 src/llama-impl.cpp delete mode 100644 src/llama-impl.h create mode 100644 src/llama-impl.h.inc create mode 100644 src/llama-io-rust/Cargo.toml create mode 100644 src/llama-io-rust/src/lib.rs delete mode 100644 src/llama-io.cpp delete mode 100644 src/llama-io.h delete mode 100644 src/llama-kv-cache-iswa.cpp delete mode 100644 src/llama-kv-cache-iswa.h delete mode 100644 src/llama-kv-cache.cpp delete mode 100644 src/llama-kv-cache.h delete mode 100644 src/llama-kv-cells.h delete mode 100644 src/llama-memory-hybrid-iswa.cpp delete mode 100644 src/llama-memory-hybrid-iswa.h delete mode 100644 src/llama-memory-hybrid.cpp delete mode 100644 src/llama-memory-hybrid.h delete mode 100644 src/llama-memory-recurrent.cpp delete mode 100644 src/llama-memory-recurrent.h create mode 100644 src/llama-memory-rust/Cargo.toml create mode 100644 src/llama-memory-rust/src/lib.rs delete mode 100644 src/llama-memory.cpp delete mode 100644 src/llama-memory.h delete mode 100644 src/llama-mmap.cpp delete mode 100644 src/llama-mmap.h delete mode 100644 src/llama-model-loader.cpp delete mode 100644 src/llama-model-loader.h delete mode 100644 src/llama-model-saver.cpp delete mode 100644 src/llama-model-saver.h delete mode 100644 src/llama-model.cpp delete mode 100644 src/llama-model.h delete mode 100644 src/llama-quant.cpp delete mode 100644 src/llama-quant.h delete mode 100644 src/llama-sampler.cpp delete mode 100644 src/llama-sampler.h delete mode 100644 src/llama-vocab.cpp delete mode 100644 src/llama-vocab.h delete mode 100644 src/llama.cpp create mode 100644 src/llama.cpp.inc delete mode 100644 src/models/afmoe.cpp delete mode 100644 src/models/apertus.cpp delete mode 100644 src/models/arcee.cpp delete mode 100644 src/models/arctic.cpp delete mode 100644 src/models/arwkv7.cpp delete mode 100644 src/models/baichuan.cpp delete mode 100644 src/models/bailingmoe.cpp delete mode 100644 src/models/bailingmoe2.cpp delete mode 100644 src/models/bert.cpp delete mode 100644 src/models/bitnet.cpp delete mode 100644 src/models/bloom.cpp delete mode 100644 src/models/chameleon.cpp delete mode 100644 src/models/chatglm.cpp delete mode 100644 src/models/codeshell.cpp delete mode 100644 src/models/cogvlm.cpp delete mode 100644 src/models/cohere2-iswa.cpp delete mode 100644 src/models/command-r.cpp delete mode 100644 src/models/dbrx.cpp delete mode 100644 src/models/deci.cpp delete mode 100644 src/models/deepseek.cpp delete mode 100644 src/models/deepseek2.cpp delete mode 100644 src/models/delta-net-base.cpp delete mode 100644 src/models/dots1.cpp delete mode 100644 src/models/dream.cpp delete mode 100644 src/models/ernie4-5-moe.cpp delete mode 100644 src/models/ernie4-5.cpp delete mode 100644 src/models/eurobert.cpp delete mode 100644 src/models/exaone-moe.cpp delete mode 100644 src/models/exaone.cpp delete mode 100644 src/models/exaone4.cpp delete mode 100644 src/models/falcon-h1.cpp delete mode 100644 src/models/falcon.cpp delete mode 100644 src/models/gemma-embedding.cpp delete mode 100644 src/models/gemma.cpp delete mode 100644 src/models/gemma2-iswa.cpp delete mode 100644 src/models/gemma3.cpp delete mode 100644 src/models/gemma3n-iswa.cpp delete mode 100644 src/models/gemma4-iswa.cpp delete mode 100644 src/models/glm4-moe.cpp delete mode 100644 src/models/glm4.cpp delete mode 100644 src/models/gpt2.cpp delete mode 100644 src/models/gptneox.cpp delete mode 100644 src/models/granite-hybrid.cpp delete mode 100644 src/models/granite.cpp delete mode 100644 src/models/grok.cpp delete mode 100644 src/models/grovemoe.cpp delete mode 100644 src/models/hunyuan-dense.cpp delete mode 100644 src/models/hunyuan-moe.cpp delete mode 100644 src/models/internlm2.cpp delete mode 100644 src/models/jais.cpp delete mode 100644 src/models/jais2.cpp delete mode 100644 src/models/jamba.cpp delete mode 100644 src/models/kimi-linear.cpp delete mode 100644 src/models/lfm2.cpp delete mode 100644 src/models/llada-moe.cpp delete mode 100644 src/models/llada.cpp delete mode 100644 src/models/llama.cpp delete mode 100644 src/models/llama4.cpp delete mode 100644 src/models/maincoder.cpp delete mode 100644 src/models/mamba-base.cpp delete mode 100644 src/models/mamba.cpp delete mode 100644 src/models/mimo2-iswa.cpp delete mode 100644 src/models/minicpm3.cpp delete mode 100644 src/models/minimax-m2.cpp delete mode 100644 src/models/mistral3.cpp delete mode 100644 src/models/models.h delete mode 100644 src/models/modern-bert.cpp delete mode 100644 src/models/mpt.cpp delete mode 100644 src/models/nemotron-h.cpp delete mode 100644 src/models/nemotron.cpp delete mode 100644 src/models/neo-bert.cpp delete mode 100644 src/models/olmo.cpp delete mode 100644 src/models/olmo2.cpp delete mode 100644 src/models/olmoe.cpp delete mode 100644 src/models/openai-moe-iswa.cpp delete mode 100644 src/models/openelm.cpp delete mode 100644 src/models/orion.cpp delete mode 100644 src/models/paddleocr.cpp delete mode 100644 src/models/pangu-embedded.cpp delete mode 100644 src/models/phi2.cpp delete mode 100644 src/models/phi3.cpp delete mode 100644 src/models/plamo.cpp delete mode 100644 src/models/plamo2.cpp delete mode 100644 src/models/plamo3.cpp delete mode 100644 src/models/plm.cpp delete mode 100644 src/models/qwen.cpp delete mode 100644 src/models/qwen2.cpp delete mode 100644 src/models/qwen2moe.cpp delete mode 100644 src/models/qwen2vl.cpp delete mode 100644 src/models/qwen3.cpp delete mode 100644 src/models/qwen35.cpp delete mode 100644 src/models/qwen35moe.cpp delete mode 100644 src/models/qwen3moe.cpp delete mode 100644 src/models/qwen3next.cpp delete mode 100644 src/models/qwen3vl-moe.cpp delete mode 100644 src/models/qwen3vl.cpp delete mode 100644 src/models/refact.cpp delete mode 100644 src/models/rnd1.cpp delete mode 100644 src/models/rwkv6-base.cpp delete mode 100644 src/models/rwkv6.cpp delete mode 100644 src/models/rwkv6qwen2.cpp delete mode 100644 src/models/rwkv7-base.cpp delete mode 100644 src/models/rwkv7.cpp delete mode 100644 src/models/seed-oss.cpp delete mode 100644 src/models/smallthinker.cpp delete mode 100644 src/models/smollm3.cpp delete mode 100644 src/models/stablelm.cpp delete mode 100644 src/models/starcoder.cpp delete mode 100644 src/models/starcoder2.cpp delete mode 100644 src/models/step35-iswa.cpp delete mode 100644 src/models/t5.cpp delete mode 100644 src/models/t5encoder.cpp delete mode 100644 src/models/wavtokenizer-dec.cpp delete mode 100644 src/models/xverse.cpp delete mode 100644 src/unicode-data.cpp delete mode 100644 src/unicode-data.h delete mode 100644 src/unicode.cpp delete mode 100644 src/unicode.h create mode 100644 target/.rustc_info.json create mode 100644 target/CACHEDIR.TAG create mode 100644 target/debug/.cargo-lock create mode 100644 target/debug/.fingerprint/ggml-backend-dl-rust-2d121e9e1f8e6973/dep-lib-ggml_backend_dl_rust create mode 100644 target/debug/.fingerprint/ggml-backend-dl-rust-2d121e9e1f8e6973/invoked.timestamp create mode 100644 target/debug/.fingerprint/ggml-backend-dl-rust-2d121e9e1f8e6973/lib-ggml_backend_dl_rust create mode 100644 target/debug/.fingerprint/ggml-backend-dl-rust-2d121e9e1f8e6973/lib-ggml_backend_dl_rust.json create mode 100644 target/debug/.fingerprint/ggml-backend-dl-rust-b4756e65d1ce8962/dep-test-lib-ggml_backend_dl_rust create mode 100644 target/debug/.fingerprint/ggml-backend-dl-rust-b4756e65d1ce8962/invoked.timestamp create mode 100644 target/debug/.fingerprint/ggml-backend-dl-rust-b4756e65d1ce8962/test-lib-ggml_backend_dl_rust create mode 100644 target/debug/.fingerprint/ggml-backend-dl-rust-b4756e65d1ce8962/test-lib-ggml_backend_dl_rust.json create mode 100644 target/debug/.fingerprint/ggml-threading-rust-0a510f30ab2a7a70/dep-lib-ggml_threading_rust create mode 100644 target/debug/.fingerprint/ggml-threading-rust-0a510f30ab2a7a70/invoked.timestamp create mode 100644 target/debug/.fingerprint/ggml-threading-rust-0a510f30ab2a7a70/lib-ggml_threading_rust create mode 100644 target/debug/.fingerprint/ggml-threading-rust-0a510f30ab2a7a70/lib-ggml_threading_rust.json create mode 100644 target/debug/.fingerprint/ggml-threading-rust-eefbb14ef3d83e68/dep-test-lib-ggml_threading_rust create mode 100644 target/debug/.fingerprint/ggml-threading-rust-eefbb14ef3d83e68/invoked.timestamp create mode 100644 target/debug/.fingerprint/ggml-threading-rust-eefbb14ef3d83e68/test-lib-ggml_threading_rust create mode 100644 target/debug/.fingerprint/ggml-threading-rust-eefbb14ef3d83e68/test-lib-ggml_threading_rust.json create mode 100644 target/debug/.fingerprint/llama-common-unicode-rust-1502b8c604829416/dep-lib-llama_common_unicode_rust create mode 100644 target/debug/.fingerprint/llama-common-unicode-rust-1502b8c604829416/invoked.timestamp create mode 100644 target/debug/.fingerprint/llama-common-unicode-rust-1502b8c604829416/lib-llama_common_unicode_rust create mode 100644 target/debug/.fingerprint/llama-common-unicode-rust-1502b8c604829416/lib-llama_common_unicode_rust.json create mode 100644 target/debug/.fingerprint/llama-common-unicode-rust-9be35eb4ef816de7/dep-test-lib-llama_common_unicode_rust create mode 100644 target/debug/.fingerprint/llama-common-unicode-rust-9be35eb4ef816de7/invoked.timestamp create mode 100644 target/debug/.fingerprint/llama-common-unicode-rust-9be35eb4ef816de7/test-lib-llama_common_unicode_rust create mode 100644 target/debug/.fingerprint/llama-common-unicode-rust-9be35eb4ef816de7/test-lib-llama_common_unicode_rust.json create mode 100644 target/debug/.fingerprint/llama-server-base64-rust-49d340cccbd5cf1e/dep-lib-llama_server_base64_rust create mode 100644 target/debug/.fingerprint/llama-server-base64-rust-49d340cccbd5cf1e/invoked.timestamp create mode 100644 target/debug/.fingerprint/llama-server-base64-rust-49d340cccbd5cf1e/lib-llama_server_base64_rust create mode 100644 target/debug/.fingerprint/llama-server-base64-rust-49d340cccbd5cf1e/lib-llama_server_base64_rust.json create mode 100644 target/debug/.fingerprint/llama-server-base64-rust-ec05054f222a1afa/dep-test-lib-llama_server_base64_rust create mode 100644 target/debug/.fingerprint/llama-server-base64-rust-ec05054f222a1afa/invoked.timestamp create mode 100644 target/debug/.fingerprint/llama-server-base64-rust-ec05054f222a1afa/test-lib-llama_server_base64_rust create mode 100644 target/debug/.fingerprint/llama-server-base64-rust-ec05054f222a1afa/test-lib-llama_server_base64_rust.json create mode 100755 target/debug/deps/ggml_backend_dl_rust-b4756e65d1ce8962 create mode 100755 target/debug/deps/ggml_threading_rust-eefbb14ef3d83e68 create mode 100644 target/debug/deps/libggml_backend_dl_rust-2d121e9e1f8e6973.rlib create mode 100644 target/debug/deps/libggml_threading_rust-0a510f30ab2a7a70.rlib create mode 100644 target/debug/deps/libllama_common_unicode_rust-1502b8c604829416.rlib create mode 100644 target/debug/deps/libllama_server_base64_rust-49d340cccbd5cf1e.rlib create mode 100755 target/debug/deps/llama_common_unicode_rust-9be35eb4ef816de7 create mode 100755 target/debug/deps/llama_server_base64_rust-ec05054f222a1afa create mode 100644 target/debug/incremental/ggml_backend_dl_rust-25gkt6c6gtxct/s-hiarnw10ke-071vupl-9ocx88kytr1ingr70sewscvcc/metadata.rmeta create mode 100644 target/debug/incremental/ggml_backend_dl_rust-25gkt6c6gtxct/s-hiarnw10ke-071vupl.lock create mode 100644 target/debug/incremental/ggml_backend_dl_rust-25gkt6c6gtxct/s-hiaroj9fec-1fpkguj-88qz973i4yqwfbkrvgyobfv2h/metadata.rmeta create mode 100644 target/debug/incremental/ggml_backend_dl_rust-25gkt6c6gtxct/s-hiaroj9fec-1fpkguj.lock create mode 100644 target/debug/incremental/ggml_backend_dl_rust-2fgv7cjo4hdtv/s-hiarnw10kc-0pzxdd3.lock create mode 100644 target/debug/incremental/ggml_backend_dl_rust-2fgv7cjo4hdtv/s-hiaroj9fe7-1sw5mx8.lock create mode 100644 target/debug/incremental/ggml_threading_rust-07vwfy03wqgv0/s-hiawyw92a5-0ad3ptp-b7zautb3dvjivoatos66mlfjc/metadata.rmeta create mode 100644 target/debug/incremental/ggml_threading_rust-07vwfy03wqgv0/s-hiawyw92a5-0ad3ptp.lock create mode 100644 target/debug/incremental/ggml_threading_rust-07vwfy03wqgv0/s-hiax0mj3dr-19df9zm-3rkm4gbxshgfzliochyhl76nw/metadata.rmeta create mode 100644 target/debug/incremental/ggml_threading_rust-07vwfy03wqgv0/s-hiax0mj3dr-19df9zm.lock create mode 100644 target/debug/incremental/ggml_threading_rust-0m3afago8nzbn/s-hiawyw92a3-199hpju.lock create mode 100644 target/debug/incremental/ggml_threading_rust-0m3afago8nzbn/s-hiax0mj3dp-0arubnp.lock create mode 100644 target/debug/incremental/llama_common_unicode_rust-1f3la0ibzk2gf/s-hiavvq3mw9-1h8uhp8-e7eijgub6ygf75zldccfwsejz/metadata.rmeta create mode 100644 target/debug/incremental/llama_common_unicode_rust-1f3la0ibzk2gf/s-hiavvq3mw9-1h8uhp8.lock create mode 100644 target/debug/incremental/llama_common_unicode_rust-1f3la0ibzk2gf/s-hiavznzjim-05mxdnk-epjml62evlpq49uk6i2ndj71u/metadata.rmeta create mode 100644 target/debug/incremental/llama_common_unicode_rust-1f3la0ibzk2gf/s-hiavznzjim-05mxdnk.lock create mode 100644 target/debug/incremental/llama_common_unicode_rust-1m12d0csu7kub/s-hiavvq3n1n-0rfum30.lock create mode 100644 target/debug/incremental/llama_common_unicode_rust-1m12d0csu7kub/s-hiavznzjmv-07vn1ra.lock create mode 100644 target/debug/incremental/llama_server_base64_rust-0mgv6r8sr514q/s-hiawfg0ftc-1j30nia.lock create mode 100644 target/debug/incremental/llama_server_base64_rust-0mgv6r8sr514q/s-hiawky1p23-1de387b.lock create mode 100644 target/debug/incremental/llama_server_base64_rust-3kfobjtf7tal7/s-hiawfg0fu5-16hqef7-7e6ik477fc9rnf2mn0juszri2/metadata.rmeta create mode 100644 target/debug/incremental/llama_server_base64_rust-3kfobjtf7tal7/s-hiawfg0fu5-16hqef7.lock create mode 100644 target/debug/incremental/llama_server_base64_rust-3kfobjtf7tal7/s-hiawky1p25-0g5qg3t-1bjcieqfmdza1bu26bypwkyum/metadata.rmeta create mode 100644 target/debug/incremental/llama_server_base64_rust-3kfobjtf7tal7/s-hiawky1p25-0g5qg3t.lock delete mode 100644 tests/export-graph-ops.cpp delete mode 100644 tests/get-model.cpp delete mode 100644 tests/get-model.h delete mode 100644 tests/gguf-model-data.cpp delete mode 100644 tests/gguf-model-data.h create mode 100644 tests/no-c-cpp-files.cmake delete mode 100644 tests/peg-parser/simple-tokenize.cpp delete mode 100644 tests/peg-parser/simple-tokenize.h delete mode 100644 tests/peg-parser/test-basic.cpp delete mode 100644 tests/peg-parser/test-gbnf-generation.cpp delete mode 100644 tests/peg-parser/test-json-parser.cpp delete mode 100644 tests/peg-parser/test-json-serialization.cpp delete mode 100644 tests/peg-parser/test-python-dict-parser.cpp delete mode 100644 tests/peg-parser/test-unicode.cpp delete mode 100644 tests/peg-parser/tests.h delete mode 100644 tests/test-alloc.cpp delete mode 100644 tests/test-arg-parser.cpp delete mode 100644 tests/test-autorelease.cpp delete mode 100644 tests/test-backend-ops.cpp delete mode 100644 tests/test-backend-sampler.cpp delete mode 100644 tests/test-barrier.cpp delete mode 100644 tests/test-c.c delete mode 100644 tests/test-chat-auto-parser.cpp delete mode 100644 tests/test-chat-peg-parser.cpp delete mode 100644 tests/test-chat-template.cpp delete mode 100644 tests/test-chat.cpp delete mode 100644 tests/test-double-float.cpp delete mode 100644 tests/test-gbnf-validator.cpp delete mode 100644 tests/test-gguf-model-data.cpp delete mode 100644 tests/test-gguf.cpp delete mode 100644 tests/test-grammar-integration.cpp delete mode 100644 tests/test-grammar-llguidance.cpp delete mode 100644 tests/test-grammar-parser.cpp delete mode 100644 tests/test-jinja.cpp delete mode 100644 tests/test-json-partial.cpp delete mode 100755 tests/test-json-schema-to-grammar.cpp delete mode 100644 tests/test-llama-archs.cpp delete mode 100644 tests/test-llama-grammar.cpp delete mode 100644 tests/test-log.cpp delete mode 100644 tests/test-model-load-cancel.cpp delete mode 100644 tests/test-mtmd-c-api.c delete mode 100644 tests/test-opt.cpp delete mode 100644 tests/test-peg-parser.cpp delete mode 100644 tests/test-quant-type-selection.cpp delete mode 100644 tests/test-quantize-fns.cpp delete mode 100644 tests/test-quantize-perf.cpp delete mode 100644 tests/test-quantize-stats.cpp delete mode 100644 tests/test-reasoning-budget.cpp delete mode 100644 tests/test-regex-partial.cpp delete mode 100644 tests/test-rope.cpp delete mode 100644 tests/test-sampling.cpp delete mode 100644 tests/test-state-restore-fragmented.cpp delete mode 100644 tests/test-thread-safety.cpp delete mode 100644 tests/test-tokenizer-0.cpp delete mode 100644 tests/test-tokenizer-1-bpe.cpp delete mode 100644 tests/test-tokenizer-1-spm.cpp delete mode 100644 tests/testing.h create mode 100644 tools/batched-bench-rust/Cargo.toml create mode 100644 tools/batched-bench-rust/build.rs create mode 100644 tools/batched-bench-rust/src/main.rs delete mode 100644 tools/batched-bench/batched-bench.cpp create mode 100644 tools/cli-rust/Cargo.toml create mode 100644 tools/cli-rust/src/main.rs delete mode 100644 tools/cli/cli.cpp create mode 100644 tools/completion-rust/Cargo.toml create mode 100644 tools/completion-rust/src/main.rs delete mode 100644 tools/completion/completion.cpp create mode 100644 tools/cvector-generator-rust/Cargo.toml create mode 100644 tools/cvector-generator-rust/src/main.rs delete mode 100644 tools/cvector-generator/cvector-generator.cpp create mode 100644 tools/cvector-generator/mean-rust/Cargo.toml create mode 100644 tools/cvector-generator/mean-rust/src/lib.rs delete mode 100644 tools/cvector-generator/mean.hpp delete mode 100644 tools/cvector-generator/pca.hpp create mode 100644 tools/export-lora-rust/Cargo.toml create mode 100644 tools/export-lora-rust/src/main.rs delete mode 100644 tools/export-lora/export-lora.cpp create mode 100644 tools/fit-params-rust/Cargo.toml create mode 100644 tools/fit-params-rust/src/main.rs delete mode 100644 tools/fit-params/fit-params.cpp create mode 100644 tools/gguf-split-rust/Cargo.toml create mode 100644 tools/gguf-split-rust/build.rs create mode 100644 tools/gguf-split-rust/src/main.rs delete mode 100644 tools/gguf-split/gguf-split.cpp create mode 100644 tools/imatrix-rust/Cargo.toml create mode 100644 tools/imatrix-rust/src/main.rs delete mode 100644 tools/imatrix/imatrix.cpp create mode 100644 tools/llama-bench-rust/Cargo.toml create mode 100644 tools/llama-bench-rust/src/main.rs delete mode 100644 tools/llama-bench/llama-bench.cpp create mode 100644 tools/mtmd/cli-rust/Cargo.toml create mode 100644 tools/mtmd/cli-rust/src/main.rs delete mode 100644 tools/mtmd/clip-graph.h delete mode 100644 tools/mtmd/clip-impl.h delete mode 100644 tools/mtmd/clip-model.h delete mode 100644 tools/mtmd/clip.cpp delete mode 100644 tools/mtmd/clip.h create mode 100644 tools/mtmd/debug-rust/Cargo.toml create mode 100644 tools/mtmd/debug-rust/src/main.rs delete mode 100644 tools/mtmd/debug/mtmd-debug.cpp delete mode 100644 tools/mtmd/debug/mtmd-debug.h create mode 100644 tools/mtmd/deprecation-warning-rust/Cargo.toml create mode 100644 tools/mtmd/deprecation-warning-rust/src/lib.rs create mode 100644 tools/mtmd/deprecation-warning-rust/src/main.rs delete mode 100644 tools/mtmd/deprecation-warning.cpp delete mode 100644 tools/mtmd/models/cogvlm.cpp delete mode 100644 tools/mtmd/models/conformer.cpp delete mode 100644 tools/mtmd/models/deepseekocr.cpp delete mode 100644 tools/mtmd/models/dotsocr.cpp delete mode 100644 tools/mtmd/models/gemma4a.cpp delete mode 100644 tools/mtmd/models/gemma4v.cpp delete mode 100644 tools/mtmd/models/glm4v.cpp delete mode 100644 tools/mtmd/models/hunyuanocr.cpp delete mode 100644 tools/mtmd/models/internvl.cpp delete mode 100644 tools/mtmd/models/kimik25.cpp delete mode 100644 tools/mtmd/models/kimivl.cpp delete mode 100644 tools/mtmd/models/llama4.cpp delete mode 100644 tools/mtmd/models/llava.cpp delete mode 100644 tools/mtmd/models/minicpmv.cpp delete mode 100644 tools/mtmd/models/mobilenetv5.cpp delete mode 100644 tools/mtmd/models/models.h delete mode 100644 tools/mtmd/models/nemotron-v2-vl.cpp delete mode 100644 tools/mtmd/models/paddleocr.cpp delete mode 100644 tools/mtmd/models/pixtral.cpp delete mode 100644 tools/mtmd/models/qwen2vl.cpp delete mode 100644 tools/mtmd/models/qwen3a.cpp delete mode 100644 tools/mtmd/models/qwen3vl.cpp delete mode 100644 tools/mtmd/models/siglip.cpp delete mode 100644 tools/mtmd/models/step3vl.cpp delete mode 100644 tools/mtmd/models/whisper-enc.cpp delete mode 100644 tools/mtmd/models/yasa2.cpp delete mode 100644 tools/mtmd/models/youtuvl.cpp delete mode 100644 tools/mtmd/mtmd-audio.cpp delete mode 100644 tools/mtmd/mtmd-audio.h delete mode 100644 tools/mtmd/mtmd-cli.cpp delete mode 100644 tools/mtmd/mtmd-helper.cpp delete mode 100644 tools/mtmd/mtmd-helper.h delete mode 100644 tools/mtmd/mtmd-image.cpp delete mode 100644 tools/mtmd/mtmd-image.h delete mode 100644 tools/mtmd/mtmd.cpp create mode 100644 tools/mtmd/mtmd.cpp.inc delete mode 100644 tools/mtmd/mtmd.h create mode 100644 tools/mtmd/mtmd.h.inc create mode 100644 tools/parser-rust/Cargo.toml create mode 100644 tools/parser-rust/build.rs create mode 100644 tools/parser-rust/src/bin/debug_template_parser.rs create mode 100644 tools/parser-rust/src/bin/template_analysis.rs create mode 100644 tools/parser-rust/src/lib.rs delete mode 100644 tools/parser/debug-template-parser.cpp delete mode 100644 tools/parser/template-analysis.cpp create mode 100644 tools/perplexity-rust/Cargo.toml create mode 100644 tools/perplexity-rust/src/main.rs delete mode 100644 tools/perplexity/perplexity.cpp create mode 100644 tools/quantize-rust/Cargo.toml create mode 100644 tools/quantize-rust/src/main.rs delete mode 100644 tools/quantize/quantize.cpp create mode 100644 tools/results-rust/Cargo.toml create mode 100644 tools/results-rust/build.rs create mode 100644 tools/results-rust/src/lib.rs create mode 100644 tools/results-rust/src/main.rs delete mode 100644 tools/results/results.cpp create mode 100644 tools/rpc-rust/Cargo.toml create mode 100644 tools/rpc-rust/build.rs create mode 100644 tools/rpc-rust/src/main.rs delete mode 100644 tools/rpc/rpc-server.cpp create mode 100644 tools/server/base64-rust/Cargo.toml create mode 100644 tools/server/base64-rust/src/lib.rs delete mode 100644 tools/server/server-chat.cpp delete mode 100644 tools/server/server-chat.h delete mode 100644 tools/server/server-common.cpp create mode 100644 tools/server/server-common.cpp.inc delete mode 100644 tools/server/server-common.h delete mode 100644 tools/server/server-context.cpp delete mode 100644 tools/server/server-context.h delete mode 100644 tools/server/server-cors-proxy.h delete mode 100644 tools/server/server-http.cpp delete mode 100644 tools/server/server-http.h delete mode 100644 tools/server/server-models.cpp delete mode 100644 tools/server/server-models.h delete mode 100644 tools/server/server-queue.cpp delete mode 100644 tools/server/server-queue.h delete mode 100644 tools/server/server-task.cpp delete mode 100644 tools/server/server-task.h delete mode 100644 tools/server/server-tools.cpp delete mode 100644 tools/server/server-tools.h delete mode 100644 tools/server/server.cpp create mode 100644 tools/server/server.cpp.inc create mode 100644 tools/tokenize-rust/Cargo.toml create mode 100644 tools/tokenize-rust/build.rs create mode 100644 tools/tokenize-rust/src/lib.rs create mode 100644 tools/tokenize-rust/src/main.rs delete mode 100644 tools/tokenize/tokenize.cpp create mode 100644 tools/tts-rust/Cargo.toml create mode 100644 tools/tts-rust/src/main.rs delete mode 100644 tools/tts/tts.cpp rename vendor/cpp-httplib/{httplib.cpp => httplib.cpp.inc} (99%) rename vendor/cpp-httplib/{httplib.h => httplib.h.inc} (99%) rename vendor/miniaudio/{miniaudio.h => miniaudio.h.inc} (99%) rename vendor/nlohmann/{json.hpp => json.hpp.inc} (99%) rename vendor/nlohmann/{json_fwd.hpp => json_fwd.hpp.inc} (100%) rename vendor/sheredom/{subprocess.h => subprocess.h.inc} (99%) rename vendor/stb/{stb_image.h => stb_image.h.inc} (99%) diff --git a/.devops/nix/package.nix b/.devops/nix/package.nix index 4e5fd00a5552..a0ace5ec34b1 100644 --- a/.devops/nix/package.nix +++ b/.devops/nix/package.nix @@ -207,7 +207,7 @@ effectiveStdenv.mkDerivation (finalAttrs: { # if they haven't been added yet. postInstall = '' mkdir -p $out/include - cp $src/include/llama.h $out/include/ + cp $src/include/llama.h.inc $out/include/ ''; meta = { diff --git a/.editorconfig b/.editorconfig index b995dbff7ddb..4dda9dcb3894 100644 --- a/.editorconfig +++ b/.editorconfig @@ -41,7 +41,7 @@ charset = unset trim_trailing_whitespace = unset insert_final_newline = unset -[vendor/miniaudio/miniaudio.h] +[vendor/miniaudio/miniaudio.h.inc] trim_trailing_whitespace = unset insert_final_newline = unset diff --git a/.github/labeler.yml b/.github/labeler.yml index 70384994dcb4..1836e19031ae 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -2,35 +2,35 @@ Apple Metal: - changed-files: - any-glob-to-any-file: - - ggml/include/ggml-metal.h + - ggml/include/ggml-metal.h.inc - ggml/src/ggml-metal/** - README-metal.md SYCL: - changed-files: - any-glob-to-any-file: - - ggml/include/ggml-sycl.h + - ggml/include/ggml-sycl.h.inc - ggml/src/ggml-sycl/** - docs/backend/SYCL.md - examples/sycl/** Nvidia GPU: - changed-files: - any-glob-to-any-file: - - ggml/include/ggml-cuda.h + - ggml/include/ggml-cuda.h.inc - ggml/src/ggml-cuda/** Vulkan: - changed-files: - any-glob-to-any-file: - - ggml/include/ggml-vulkan.h + - ggml/include/ggml-vulkan.h.inc - ggml/src/ggml-vulkan/** IBM zDNN: - changed-files: - any-glob-to-any-file: - - ggml/include/ggml-zdnn.h + - ggml/include/ggml-zdnn.h.inc - ggml/src/ggml-zdnn/** AMD ZenDNN: - changed-files: - any-glob-to-any-file: - - ggml/include/ggml-zendnn.h + - ggml/include/ggml-zendnn.h.inc - ggml/src/ggml-zendnn/** documentation: - changed-files: @@ -109,28 +109,28 @@ jinja parser: Ascend NPU: - changed-files: - any-glob-to-any-file: - - ggml/include/ggml-cann.h + - ggml/include/ggml-cann.h.inc - ggml/src/ggml-cann/** - docs/backend/CANN.md OpenCL: - changed-files: - any-glob-to-any-file: - - ggml/include/ggml-opencl.h + - ggml/include/ggml-opencl.h.inc - ggml/src/ggml-opencl/** - docs/backend/OPENCL.md Hexagon: - changed-files: - any-glob-to-any-file: - - ggml/include/ggml-hexagon.h + - ggml/include/ggml-hexagon.h.inc - ggml/src/ggml-hexagon/** WebGPU: - changed-files: - any-glob-to-any-file: - - ggml/include/ggml-webgpu.h + - ggml/include/ggml-webgpu.h.inc - ggml/src/ggml-webgpu/** OpenVINO: - changed-files: - any-glob-to-any-file: - - ggml/include/ggml-openvino.h + - ggml/include/ggml-openvino.h.inc - ggml/src/ggml-openvino/** - docs/backend/OPENVINO.md diff --git a/.github/workflows/build-and-test-snapdragon.yml b/.github/workflows/build-and-test-snapdragon.yml index deed8e808b78..6cd44752153c 100644 --- a/.github/workflows/build-and-test-snapdragon.yml +++ b/.github/workflows/build-and-test-snapdragon.yml @@ -7,7 +7,7 @@ on: - master paths: - '.github/workflows/build-and-test-snapdragon.yml' - - 'ggml/include/ggml-hexagon.h' + - 'ggml/include/ggml-hexagon.h.inc' - 'ggml/src/ggml-hexagon/**' - 'docs/backend/snapdragon/**' - 'scripts/snapdragon/**' @@ -17,7 +17,7 @@ on: types: [opened, synchronize, reopened] paths: - '.github/workflows/build-and-test-snapdragon.yml' - - 'ggml/include/ggml-hexagon.h' + - 'ggml/include/ggml-hexagon.h.inc' - 'ggml/src/ggml-hexagon/**' - 'docs/backend/snapdragon/**' - 'scripts/snapdragon/**' diff --git a/CMakeLists.txt b/CMakeLists.txt index 310a3dcfd241..ce33e517054a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -240,8 +240,7 @@ set(LLAMA_LIB_INSTALL_DIR ${CMAKE_INSTALL_LIBDIR} CACHE PATH "Location o set(LLAMA_BIN_INSTALL_DIR ${CMAKE_INSTALL_BINDIR} CACHE PATH "Location of binary files") set(LLAMA_PUBLIC_HEADERS - ${CMAKE_CURRENT_SOURCE_DIR}/include/llama.h - ${CMAKE_CURRENT_SOURCE_DIR}/include/llama-cpp.h) + ${CMAKE_CURRENT_SOURCE_DIR}/include/llama.h.inc) set_target_properties(llama PROPERTIES diff --git a/CODEOWNERS b/CODEOWNERS index eed43d762c78..55d5d13f23fc 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -55,11 +55,11 @@ /ggml/include/ @ggerganov /ggml/src/ggml-backend-meta.cpp @JohannesGaessler /ggml/src/ggml-cann/ @ggml-org/ggml-cann -/ggml/src/ggml-common.h @ggerganov +/ggml/src/ggml-common.h.inc @ggerganov /ggml/src/ggml-cpu/ @ggerganov /ggml/src/ggml-cpu/spacemit/ @alex-spacemit /ggml/src/ggml-cuda/ @ggml-org/ggml-cuda -/ggml/src/ggml-cuda/vendors/hip.h @IMbackK +/ggml/src/ggml-cuda/vendors/hip.h.inc @IMbackK /ggml/src/ggml-cuda/fattn-wmma* @IMbackK /ggml/src/ggml-hexagon/ @ggml-org/ggml-hexagon /ggml/src/ggml-hip/ @IMbackK diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8000b4718676..1ad46c51e959 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -93,7 +93,7 @@ Maintainers reserve the right to decline review or close pull requests for any r - Try to follow the existing patterns in the code (indentation, spaces, etc.). In case of doubt use `clang-format` (from clang-tools v15+) to format the added code - For anything not covered in the current guidelines, refer to the [C++ Core Guidelines](https://isocpp.github.io/CppCoreGuidelines/CppCoreGuidelines) - Tensors store data in row-major order. We refer to dimension 0 as columns, 1 as rows, 2 as matrices -- Matrix multiplication is unconventional: [`C = ggml_mul_mat(ctx, A, B)`](https://github.com/ggml-org/llama.cpp/blob/880e352277fc017df4d5794f0c21c44e1eae2b84/ggml.h#L1058-L1064) means $C^T = A B^T \Leftrightarrow C = B A^T.$ +- Matrix multiplication is unconventional: [`C = ggml_mul_mat(ctx, A, B)`](https://github.com/ggml-org/llama.cpp/blob/880e352277fc017df4d5794f0c21c44e1eae2b84/ggml.h.inc#L1058-L1064) means $C^T = A B^T \Leftrightarrow C = B A^T.$ ![matmul](media/matmul.png) diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 000000000000..bb9e2cb47c3a --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,700 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "ggml-backend-dl-rust" +version = "0.1.0" + +[[package]] +name = "ggml-cpu-arm-feats-rust" +version = "0.1.0" + +[[package]] +name = "ggml-cpu-powerpc-feats-rust" +version = "0.1.0" + +[[package]] +name = "ggml-cpu-riscv-feats-rust" +version = "0.1.0" + +[[package]] +name = "ggml-cpu-s390-feats-rust" +version = "0.1.0" + +[[package]] +name = "ggml-cpu-x86-feats-rust" +version = "0.1.0" + +[[package]] +name = "ggml-threading-rust" +version = "0.1.0" + +[[package]] +name = "ggml-virtgpu-backend-dispatched-rust" +version = "0.1.0" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "llama-api-rust" +version = "0.1.0" + +[[package]] +name = "llama-batched-bench-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-batched-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-bench-rust" +version = "0.1.0" + +[[package]] +name = "llama-cli-rust" +version = "0.1.0" + +[[package]] +name = "llama-common-build-info-rust" +version = "0.1.0" + +[[package]] +name = "llama-common-jinja-string-rust" +version = "0.1.0" + +[[package]] +name = "llama-common-ngram-cache-rust" +version = "0.1.0" + +[[package]] +name = "llama-common-ngram-map-rust" +version = "0.1.0" + +[[package]] +name = "llama-common-ngram-mod-rust" +version = "0.1.0" + +[[package]] +name = "llama-common-reasoning-budget-rust" +version = "0.1.0" + +[[package]] +name = "llama-common-regex-partial-rust" +version = "0.1.0" + +[[package]] +name = "llama-common-unicode-rust" +version = "0.1.0" + +[[package]] +name = "llama-completion-rust" +version = "0.1.0" + +[[package]] +name = "llama-convert-llama2c-to-ggml-rust" +version = "0.1.0" + +[[package]] +name = "llama-cparams-rust" +version = "0.1.0" + +[[package]] +name = "llama-cvector-generator-rust" +version = "0.1.0" + +[[package]] +name = "llama-cvector-mean-rust" +version = "0.1.0" + +[[package]] +name = "llama-debug-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-deprecation-warning" +version = "0.1.0" + +[[package]] +name = "llama-diffusion-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-embedding-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-eval-callback-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-export-graph-ops-rust" +version = "0.1.0" + +[[package]] +name = "llama-export-lora-rust" +version = "0.1.0" + +[[package]] +name = "llama-finetune-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-fit-params-rust" +version = "0.1.0" + +[[package]] +name = "llama-gen-docs-rust" +version = "0.1.0" + +[[package]] +name = "llama-gguf-hash-rust" +version = "0.1.0" +dependencies = [ + "sha1", + "sha2", + "xxhash-rust", +] + +[[package]] +name = "llama-gguf-rust" +version = "0.1.0" + +[[package]] +name = "llama-gguf-split-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-idle-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-imatrix-rust" +version = "0.1.0" + +[[package]] +name = "llama-io-rust" +version = "0.1.0" + +[[package]] +name = "llama-lookahead-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-lookup-create-rust" +version = "0.1.0" +dependencies = [ + "llama-lookup-merge", + "llama-simple-rust", +] + +[[package]] +name = "llama-lookup-merge" +version = "0.1.0" + +[[package]] +name = "llama-lookup-rust" +version = "0.1.0" +dependencies = [ + "llama-lookup-merge", + "llama-simple-rust", +] + +[[package]] +name = "llama-lookup-stats-rust" +version = "0.1.0" +dependencies = [ + "llama-lookup-merge", + "llama-simple-rust", +] + +[[package]] +name = "llama-ls-sycl-device-rust" +version = "0.1.0" + +[[package]] +name = "llama-memory-rust" +version = "0.1.0" + +[[package]] +name = "llama-mtmd-cli-rust" +version = "0.1.0" + +[[package]] +name = "llama-mtmd-debug-rust" +version = "0.1.0" + +[[package]] +name = "llama-parallel-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-parser-tools-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-passkey-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-perplexity-rust" +version = "0.1.0" + +[[package]] +name = "llama-q8dot-rust" +version = "0.1.0" + +[[package]] +name = "llama-quantize-rust" +version = "0.1.0" + +[[package]] +name = "llama-results-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-retrieval-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-rpc-server-rust" +version = "0.1.0" + +[[package]] +name = "llama-save-load-state-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-server-base64-rust" +version = "0.1.0" + +[[package]] +name = "llama-simple-chat-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-simple-rust" +version = "0.1.0" + +[[package]] +name = "llama-simple-tokenize-rust" +version = "0.1.0" + +[[package]] +name = "llama-speculative-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-speculative-simple-rust" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-test-alloc" +version = "0.1.0" + +[[package]] +name = "llama-test-arg-parser-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-autorelease" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-test-backend-ops-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-backend-sampler-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-barrier" +version = "0.1.0" + +[[package]] +name = "llama-test-chat-auto-parser-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-chat-peg-parser-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-chat-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-chat-template-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-double-float" +version = "0.1.0" + +[[package]] +name = "llama-test-gbnf-validator" +version = "0.1.0" + +[[package]] +name = "llama-test-get-model-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-gguf-model-data" +version = "0.1.0" + +[[package]] +name = "llama-test-gguf-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-grammar-integration-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-grammar-llguidance-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-grammar-parser-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-jinja-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-json-partial-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-json-schema-to-grammar-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-llama-archs-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-llama-grammar-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-log" +version = "0.1.0" + +[[package]] +name = "llama-test-model-load-cancel" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-test-mtmd-c-api" +version = "0.1.0" + +[[package]] +name = "llama-test-opt-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-peg-parser-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-quant-type-selection-rust" +version = "0.1.0" + +[[package]] +name = "llama-test-quantize-fns" +version = "0.1.0" + +[[package]] +name = "llama-test-quantize-perf" +version = "0.1.0" +dependencies = [ + "llama-test-quantize-fns", +] + +[[package]] +name = "llama-test-quantize-stats-rust" +version = "0.1.0" +dependencies = [ + "regex", +] + +[[package]] +name = "llama-test-reasoning-budget" +version = "0.1.0" + +[[package]] +name = "llama-test-regex-partial" +version = "0.1.0" +dependencies = [ + "llama-common-regex-partial-rust", + "regex", +] + +[[package]] +name = "llama-test-rope" +version = "0.1.0" + +[[package]] +name = "llama-test-sampling" +version = "0.1.0" + +[[package]] +name = "llama-test-state-restore-fragmented" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-test-thread-safety" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-test-tokenizer-0" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-test-tokenizer-1-bpe" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-test-tokenizer-1-spm" +version = "0.1.0" +dependencies = [ + "llama-simple-rust", +] + +[[package]] +name = "llama-tokenize-rust" +version = "0.1.0" + +[[package]] +name = "llama-tts-rust" +version = "0.1.0" + +[[package]] +name = "llama-vdot-rust" +version = "0.1.0" + +[[package]] +name = "memchr" +version = "2.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" + +[[package]] +name = "mtmd-deprecation-warning" +version = "0.1.0" + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" + +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "typenum" +version = "1.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40ce102ab67701b8526c123c1bab5cbe42d7040ccfd0f64af1a385808d2f43de" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "xxhash-rust" +version = "0.8.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 000000000000..a2ae0033a4b6 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,116 @@ +[workspace] +resolver = "2" +members = [ + "common/build-info-rust", + "common/jinja/string-rust", + "common/ngram-cache-rust", + "common/ngram-map-rust", + "common/ngram-mod-rust", + "common/reasoning-budget-rust", + "common/regex-partial-rust", + "common/unicode-rust", + "examples/deprecation-warning", + "examples/debug-rust", + "examples/batched-rust", + "examples/convert-llama2c-to-ggml-rust", + "examples/diffusion-rust", + "examples/embedding-rust", + "examples/eval-callback-rust", + "examples/gen-docs-rust", + "examples/gguf-rust", + "examples/gguf-hash-rust", + "examples/idle-rust", + "examples/lookahead-rust", + "examples/lookup-create-rust", + "examples/lookup-rust", + "examples/lookup-merge-rust", + "examples/lookup-stats-rust", + "examples/passkey-rust", + "examples/parallel-rust", + "examples/retrieval-rust", + "examples/save-load-state-rust", + "examples/simple-chat-rust", + "examples/simple-rust", + "examples/speculative-simple-rust", + "examples/speculative-rust", + "examples/sycl/ls-sycl-device-rust", + "ggml/src/ggml-backend-dl-rust", + "ggml/src/ggml-cpu/arch/arm/cpu-feats-rust", + "ggml/src/ggml-cpu/arch/powerpc/cpu-feats-rust", + "ggml/src/ggml-cpu/arch/riscv/cpu-feats-rust", + "ggml/src/ggml-cpu/arch/s390/cpu-feats-rust", + "ggml/src/ggml-cpu/arch/x86/cpu-feats-rust", + "ggml/src/ggml-threading-rust", + "ggml/src/ggml-virtgpu/backend/backend-dispatched-rust", + "pocs/vdot/vdot-rust", + "pocs/vdot/q8dot-rust", + "src/llama-cparams-rust", + "src/llama-io-rust", + "src/llama-memory-rust", + "tests/get-model-rust", + "tests/llama-api-rust", + "tests/peg-parser/simple-tokenize-rust", + "tests/test-autorelease-rust", + "tests/test-barrier-rust", + "tests/test-backend-sampler-rust", + "tests/test-backend-ops-rust", + "tests/test-chat-peg-parser-rust", + "tests/test-chat-auto-parser-rust", + "tests/test-chat-rust", + "tests/test-chat-template-rust", + "tests/test-arg-parser-rust", + "tests/test-double-float-rust", + "tests/export-graph-ops-rust", + "tests/test-gbnf-validator-rust", + "tests/test-gguf-model-data-rust", + "tests/test-gguf-rust", + "tests/test-grammar-integration-rust", + "tests/test-grammar-llguidance-rust", + "tests/test-grammar-parser-rust", + "tests/test-llama-grammar-rust", + "tests/test-json-partial-rust", + "tests/test-json-schema-to-grammar-rust", + "tests/test-jinja-rust", + "tests/test-peg-parser-rust", + "tests/test-alloc-rust", + "tests/test-quantize-fns-rust", + "tests/test-quantize-perf-rust", + "tests/test-quantize-stats-rust", + "tests/test-quant-type-selection-rust", + "tests/test-reasoning-budget-rust", + "tests/test-regex-partial-rust", + "tests/test-rope-rust", + "tests/test-sampling-rust", + "tests/test-state-restore-fragmented-rust", + "tests/test-thread-safety-rust", + "tests/test-tokenizer-0-rust", + "tests/test-tokenizer-1-bpe-rust", + "tests/test-tokenizer-1-spm-rust", + "tests/test-log-rust", + "tests/test-model-load-cancel-rust", + "tests/test-mtmd-c-api-rust", + "tests/test-opt-rust", + "tests/test-llama-archs-rust", + "examples/training/finetune-rust", + "tools/batched-bench-rust", + "tools/cli-rust", + "tools/completion-rust", + "tools/cvector-generator/mean-rust", + "tools/cvector-generator-rust", + "tools/export-lora-rust", + "tools/fit-params-rust", + "tools/gguf-split-rust", + "tools/imatrix-rust", + "tools/llama-bench-rust", + "tools/parser-rust", + "tools/perplexity-rust", + "tools/quantize-rust", + "tools/mtmd/cli-rust", + "tools/mtmd/debug-rust", + "tools/mtmd/deprecation-warning-rust", + "tools/results-rust", + "tools/rpc-rust", + "tools/server/base64-rust", + "tools/tts-rust", + "tools/tokenize-rust", +] diff --git a/README.md b/README.md index be23abcea67f..ec04d202d9e2 100644 --- a/README.md +++ b/README.md @@ -592,5 +592,5 @@ $ echo "source ~/.llama-completion.bash" >> ~/.bashrc - [yhirose/cpp-httplib](https://github.com/yhirose/cpp-httplib) - Single-header HTTP server, used by `llama-server` - MIT license - [stb-image](https://github.com/nothings/stb) - Single-header image format decoder, used by multimodal subsystem - Public domain - [nlohmann/json](https://github.com/nlohmann/json) - Single-header JSON library, used by various tools/examples - MIT License -- [miniaudio.h](https://github.com/mackron/miniaudio) - Single-header audio format decoder, used by multimodal subsystem - Public domain -- [subprocess.h](https://github.com/sheredom/subprocess.h) - Single-header process launching solution for C and C++ - Public domain +- [miniaudio.h.inc](https://github.com/mackron/miniaudio) - Single-header audio format decoder, used by multimodal subsystem - Public domain +- [subprocess.h.inc](https://github.com/sheredom/subprocess.h.inc) - Single-header process launching solution for C and C++ - Public domain diff --git a/build-xcframework.sh b/build-xcframework.sh index c25a1ef28c18..7603f6ff2977 100755 --- a/build-xcframework.sh +++ b/build-xcframework.sh @@ -113,27 +113,27 @@ setup_framework_structure() { fi # Copy all required headers (common for all platforms) - cp include/llama.h ${header_path} - cp ggml/include/ggml.h ${header_path} - cp ggml/include/ggml-opt.h ${header_path} - cp ggml/include/ggml-alloc.h ${header_path} - cp ggml/include/ggml-backend.h ${header_path} - cp ggml/include/ggml-metal.h ${header_path} - cp ggml/include/ggml-cpu.h ${header_path} - cp ggml/include/ggml-blas.h ${header_path} - cp ggml/include/gguf.h ${header_path} + cp include/llama.h.inc ${header_path} + cp ggml/include/ggml.h.inc ${header_path} + cp ggml/include/ggml-opt.h.inc ${header_path} + cp ggml/include/ggml-alloc.h.inc ${header_path} + cp ggml/include/ggml-backend.h.inc ${header_path} + cp ggml/include/ggml-metal.h.inc ${header_path} + cp ggml/include/ggml-cpu.h.inc ${header_path} + cp ggml/include/ggml-blas.h.inc ${header_path} + cp ggml/include/gguf.h.inc ${header_path} # Create module map (common for all platforms) cat > ${module_path}module.modulemap << EOF framework module llama { - header "llama.h" - header "ggml.h" - header "ggml-alloc.h" - header "ggml-backend.h" - header "ggml-metal.h" - header "ggml-cpu.h" - header "ggml-blas.h" - header "gguf.h" + header "llama.h.inc" + header "ggml.h.inc" + header "ggml-alloc.h.inc" + header "ggml-backend.h.inc" + header "ggml-metal.h.inc" + header "ggml-cpu.h.inc" + header "ggml-blas.h.inc" + header "gguf.h.inc" link "c++" link framework "Accelerate" diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index 1a56c25857fa..16acf14d316b 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -1,6 +1,7 @@ find_package(Threads REQUIRED) llama_add_compile_flags() +find_program(CARGO_EXECUTABLE cargo REQUIRED) # # llama-common-base @@ -33,19 +34,51 @@ else() message(WARNING "Git repository not found; to enable automatic generation of build info, make sure Git is installed and the project is a Git repository.") endif() -set(TEMPLATE_FILE "${CMAKE_CURRENT_SOURCE_DIR}/build-info.cpp.in") -set(OUTPUT_FILE "${CMAKE_CURRENT_BINARY_DIR}/build-info.cpp") - -configure_file(${TEMPLATE_FILE} ${OUTPUT_FILE}) +set(COMMON_BUILD_INFO_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/build-info-rust/Cargo.toml) +set(COMMON_BUILD_INFO_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/common-build-info) +set(COMMON_BUILD_INFO_RUST_LIB ${COMMON_BUILD_INFO_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}llama_common_build_info_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) +set(COMMON_JINJA_STRING_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/jinja/string-rust/Cargo.toml) +set(COMMON_JINJA_STRING_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/common-jinja-string) +set(COMMON_JINJA_STRING_RUST_LIB ${COMMON_JINJA_STRING_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}llama_common_jinja_string_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) + +add_custom_command( + OUTPUT ${COMMON_BUILD_INFO_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env + "RUSTFLAGS=-Crelocation-model=pic -Cpanic=abort" + "LLAMA_BUILD_NUMBER=${LLAMA_BUILD_NUMBER}" + "LLAMA_BUILD_COMMIT=${LLAMA_BUILD_COMMIT}" + "BUILD_COMPILER=${BUILD_COMPILER}" + "BUILD_TARGET=${BUILD_TARGET}" + ${CARGO_EXECUTABLE} build --manifest-path ${COMMON_BUILD_INFO_RUST_MANIFEST} --target-dir ${COMMON_BUILD_INFO_RUST_TARGET_DIR} + DEPENDS + ${COMMON_BUILD_INFO_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/build-info-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/build-info-rust/src/lib.rs +) +add_custom_target(common-build-info-rust-build DEPENDS ${COMMON_BUILD_INFO_RUST_LIB}) +add_library(common-build-info-rust STATIC IMPORTED GLOBAL) +set_target_properties(common-build-info-rust PROPERTIES IMPORTED_LOCATION ${COMMON_BUILD_INFO_RUST_LIB}) +add_dependencies(common-build-info-rust common-build-info-rust-build) + +add_custom_command( + OUTPUT ${COMMON_JINJA_STRING_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=-Crelocation-model=pic -Cpanic=abort" + ${CARGO_EXECUTABLE} build --manifest-path ${COMMON_JINJA_STRING_RUST_MANIFEST} --target-dir ${COMMON_JINJA_STRING_RUST_TARGET_DIR} + DEPENDS + ${COMMON_JINJA_STRING_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/jinja/string-rust/src/lib.rs +) +add_custom_target(common-jinja-string-rust-build DEPENDS ${COMMON_JINJA_STRING_RUST_LIB}) +add_library(common-jinja-string-rust STATIC IMPORTED GLOBAL) +set_target_properties(common-jinja-string-rust PROPERTIES IMPORTED_LOCATION ${COMMON_JINJA_STRING_RUST_LIB}) +add_dependencies(common-jinja-string-rust common-jinja-string-rust-build) set(TARGET llama-common-base) -add_library(${TARGET} STATIC ${OUTPUT_FILE}) +add_library(${TARGET} INTERFACE) -target_include_directories(${TARGET} PUBLIC .) - -if (BUILD_SHARED_LIBS) - set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON) -endif() +target_include_directories(${TARGET} INTERFACE .) +target_link_libraries(${TARGET} INTERFACE common-build-info-rust) +target_link_libraries(${TARGET} INTERFACE common-jinja-string-rust) # # llama-common @@ -53,70 +86,109 @@ endif() set(TARGET llama-common) +set(COMMON_NGRAM_MOD_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/ngram-mod-rust/Cargo.toml) +set(COMMON_NGRAM_MOD_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/common-ngram-mod) +set(COMMON_NGRAM_MOD_RUST_LIB ${COMMON_NGRAM_MOD_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}llama_common_ngram_mod_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) +set(COMMON_NGRAM_CACHE_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/ngram-cache-rust/Cargo.toml) +set(COMMON_NGRAM_CACHE_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/common-ngram-cache) +set(COMMON_NGRAM_CACHE_RUST_LIB ${COMMON_NGRAM_CACHE_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}llama_common_ngram_cache_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) +set(COMMON_NGRAM_MAP_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/ngram-map-rust/Cargo.toml) +set(COMMON_NGRAM_MAP_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/common-ngram-map) +set(COMMON_NGRAM_MAP_RUST_LIB ${COMMON_NGRAM_MAP_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}llama_common_ngram_map_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) +set(COMMON_REASONING_BUDGET_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/reasoning-budget-rust/Cargo.toml) +set(COMMON_REASONING_BUDGET_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/common-reasoning-budget) +set(COMMON_REASONING_BUDGET_RUST_LIB ${COMMON_REASONING_BUDGET_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}llama_common_reasoning_budget_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) +set(COMMON_REGEX_PARTIAL_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/regex-partial-rust/Cargo.toml) +set(COMMON_REGEX_PARTIAL_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/common-regex-partial) +set(COMMON_REGEX_PARTIAL_RUST_LIB ${COMMON_REGEX_PARTIAL_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}llama_common_regex_partial_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) +set(COMMON_UNICODE_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/unicode-rust/Cargo.toml) +set(COMMON_UNICODE_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/common-unicode) +set(COMMON_UNICODE_RUST_LIB ${COMMON_UNICODE_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}llama_common_unicode_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) + +add_custom_command( + OUTPUT ${COMMON_NGRAM_MOD_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env RUSTFLAGS=-Crelocation-model=pic + ${CARGO_EXECUTABLE} build --manifest-path ${COMMON_NGRAM_MOD_RUST_MANIFEST} --target-dir ${COMMON_NGRAM_MOD_RUST_TARGET_DIR} + DEPENDS + ${COMMON_NGRAM_MOD_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/ngram-mod-rust/src/lib.rs +) +add_custom_target(common-ngram-mod-rust-build DEPENDS ${COMMON_NGRAM_MOD_RUST_LIB}) +add_library(common-ngram-mod-rust STATIC IMPORTED GLOBAL) +set_target_properties(common-ngram-mod-rust PROPERTIES IMPORTED_LOCATION ${COMMON_NGRAM_MOD_RUST_LIB}) +add_dependencies(common-ngram-mod-rust common-ngram-mod-rust-build) + +add_custom_command( + OUTPUT ${COMMON_NGRAM_CACHE_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=-Crelocation-model=pic -Cpanic=abort" + ${CARGO_EXECUTABLE} build --manifest-path ${COMMON_NGRAM_CACHE_RUST_MANIFEST} --target-dir ${COMMON_NGRAM_CACHE_RUST_TARGET_DIR} + DEPENDS + ${COMMON_NGRAM_CACHE_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/ngram-cache-rust/src/lib.rs +) +add_custom_target(common-ngram-cache-rust-build DEPENDS ${COMMON_NGRAM_CACHE_RUST_LIB}) +add_library(common-ngram-cache-rust STATIC IMPORTED GLOBAL) +set_target_properties(common-ngram-cache-rust PROPERTIES IMPORTED_LOCATION ${COMMON_NGRAM_CACHE_RUST_LIB}) +add_dependencies(common-ngram-cache-rust common-ngram-cache-rust-build) + +add_custom_command( + OUTPUT ${COMMON_NGRAM_MAP_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=-Crelocation-model=pic -Cpanic=abort" + ${CARGO_EXECUTABLE} build --manifest-path ${COMMON_NGRAM_MAP_RUST_MANIFEST} --target-dir ${COMMON_NGRAM_MAP_RUST_TARGET_DIR} + DEPENDS + ${COMMON_NGRAM_MAP_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/ngram-map-rust/src/lib.rs +) +add_custom_target(common-ngram-map-rust-build DEPENDS ${COMMON_NGRAM_MAP_RUST_LIB}) +add_library(common-ngram-map-rust STATIC IMPORTED GLOBAL) +set_target_properties(common-ngram-map-rust PROPERTIES IMPORTED_LOCATION ${COMMON_NGRAM_MAP_RUST_LIB}) +add_dependencies(common-ngram-map-rust common-ngram-map-rust-build) + +add_custom_command( + OUTPUT ${COMMON_REASONING_BUDGET_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=-Crelocation-model=pic -Cpanic=abort" + ${CARGO_EXECUTABLE} build --manifest-path ${COMMON_REASONING_BUDGET_RUST_MANIFEST} --target-dir ${COMMON_REASONING_BUDGET_RUST_TARGET_DIR} + DEPENDS + ${COMMON_REASONING_BUDGET_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/reasoning-budget-rust/src/lib.rs +) +add_custom_target(common-reasoning-budget-rust-build DEPENDS ${COMMON_REASONING_BUDGET_RUST_LIB}) +add_library(common-reasoning-budget-rust STATIC IMPORTED GLOBAL) +set_target_properties(common-reasoning-budget-rust PROPERTIES IMPORTED_LOCATION ${COMMON_REASONING_BUDGET_RUST_LIB}) +add_dependencies(common-reasoning-budget-rust common-reasoning-budget-rust-build) + +add_custom_command( + OUTPUT ${COMMON_REGEX_PARTIAL_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=-Crelocation-model=pic -Cpanic=abort" + ${CARGO_EXECUTABLE} build --manifest-path ${COMMON_REGEX_PARTIAL_RUST_MANIFEST} --target-dir ${COMMON_REGEX_PARTIAL_RUST_TARGET_DIR} + DEPENDS + ${COMMON_REGEX_PARTIAL_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/regex-partial-rust/src/lib.rs +) +add_custom_target(common-regex-partial-rust-build DEPENDS ${COMMON_REGEX_PARTIAL_RUST_LIB}) +add_library(common-regex-partial-rust STATIC IMPORTED GLOBAL) +set_target_properties(common-regex-partial-rust PROPERTIES IMPORTED_LOCATION ${COMMON_REGEX_PARTIAL_RUST_LIB}) +add_dependencies(common-regex-partial-rust common-regex-partial-rust-build) + +add_custom_command( + OUTPUT ${COMMON_UNICODE_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=-Crelocation-model=pic -Cpanic=abort" + ${CARGO_EXECUTABLE} build --manifest-path ${COMMON_UNICODE_RUST_MANIFEST} --target-dir ${COMMON_UNICODE_RUST_TARGET_DIR} + DEPENDS + ${COMMON_UNICODE_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/unicode-rust/src/lib.rs +) +add_custom_target(common-unicode-rust-build DEPENDS ${COMMON_UNICODE_RUST_LIB}) +add_library(common-unicode-rust STATIC IMPORTED GLOBAL) +set_target_properties(common-unicode-rust PROPERTIES IMPORTED_LOCATION ${COMMON_UNICODE_RUST_LIB}) +add_dependencies(common-unicode-rust common-unicode-rust-build) + add_library(${TARGET} - arg.cpp - arg.h - base64.hpp - chat-auto-parser-generator.cpp - chat-auto-parser-helpers.cpp - chat-auto-parser.h - chat-diff-analyzer.cpp - chat-peg-parser.cpp - chat-peg-parser.h - chat.cpp - chat.h - common.cpp - common.h - console.cpp - console.h - debug.cpp - debug.h - download.cpp - download.h - fit.cpp - fit.h - hf-cache.cpp - hf-cache.h - http.h - json-partial.cpp - json-partial.h - json-schema-to-grammar.cpp - llguidance.cpp - log.cpp - log.h - ngram-cache.cpp - ngram-cache.h - ngram-map.cpp - ngram-map.h - ngram-mod.cpp - ngram-mod.h - peg-parser.cpp - peg-parser.h - preset.cpp - preset.h - regex-partial.cpp - reasoning-budget.cpp - reasoning-budget.h - regex-partial.h - sampling.cpp - sampling.h - speculative.cpp - speculative.h - unicode.cpp - unicode.h - jinja/lexer.cpp - jinja/lexer.h - jinja/parser.cpp - jinja/parser.h - jinja/runtime.cpp - jinja/runtime.h - jinja/value.cpp - jinja/value.h - jinja/string.cpp - jinja/string.h - jinja/caps.cpp - jinja/caps.h + common.cpp.inc + common.h.inc ) +set_source_files_properties(common.cpp.inc PROPERTIES LANGUAGE CXX) +set_source_files_properties(common.h.inc PROPERTIES HEADER_FILE_ONLY TRUE) set_target_properties(${TARGET} PROPERTIES VERSION ${LLAMA_INSTALL_VERSION} @@ -135,38 +207,16 @@ if (BUILD_SHARED_LIBS) endif() target_link_libraries(${TARGET} PUBLIC llama-common-base) +target_link_libraries(${TARGET} PRIVATE common-ngram-mod-rust) +target_link_libraries(${TARGET} PRIVATE common-ngram-cache-rust) +target_link_libraries(${TARGET} PRIVATE common-ngram-map-rust) +target_link_libraries(${TARGET} PUBLIC common-reasoning-budget-rust) +target_link_libraries(${TARGET} PUBLIC common-regex-partial-rust) +target_link_libraries(${TARGET} PUBLIC common-unicode-rust) target_link_libraries(${TARGET} PRIVATE cpp-httplib) if (LLAMA_LLGUIDANCE) - include(ExternalProject) - set(LLGUIDANCE_SRC ${CMAKE_BINARY_DIR}/llguidance/source) - set(LLGUIDANCE_PATH ${LLGUIDANCE_SRC}/target/release) - set(LLGUIDANCE_LIB_NAME "${CMAKE_STATIC_LIBRARY_PREFIX}llguidance${CMAKE_STATIC_LIBRARY_SUFFIX}") - - ExternalProject_Add(llguidance_ext - GIT_REPOSITORY https://github.com/guidance-ai/llguidance - # v1.0.1: - GIT_TAG d795912fedc7d393de740177ea9ea761e7905774 - PREFIX ${CMAKE_BINARY_DIR}/llguidance - SOURCE_DIR ${LLGUIDANCE_SRC} - BUILD_IN_SOURCE TRUE - CONFIGURE_COMMAND "" - BUILD_COMMAND cargo build --release --package llguidance - INSTALL_COMMAND "" - BUILD_BYPRODUCTS ${LLGUIDANCE_PATH}/${LLGUIDANCE_LIB_NAME} ${LLGUIDANCE_PATH}/llguidance.h - UPDATE_COMMAND "" - ) - target_compile_definitions(${TARGET} PUBLIC LLAMA_USE_LLGUIDANCE) - - add_library(llguidance STATIC IMPORTED) - set_target_properties(llguidance PROPERTIES IMPORTED_LOCATION ${LLGUIDANCE_PATH}/${LLGUIDANCE_LIB_NAME}) - add_dependencies(llguidance llguidance_ext) - - target_include_directories(${TARGET} PRIVATE ${LLGUIDANCE_PATH}) - target_link_libraries(${TARGET} PRIVATE llguidance) - if (WIN32) - target_link_libraries(${TARGET} PRIVATE ws2_32 userenv ntdll bcrypt) - endif() + message(FATAL_ERROR "LLAMA_LLGUIDANCE is not available in this Rust-port build") endif() target_link_libraries(${TARGET} PUBLIC llama Threads::Threads) diff --git a/common/arg.cpp b/common/arg.cpp deleted file mode 100644 index c21598e7687f..000000000000 --- a/common/arg.cpp +++ /dev/null @@ -1,4098 +0,0 @@ -#include "arg.h" - -#include "build-info.h" -#include "chat.h" -#include "common.h" -#include "download.h" -#include "hf-cache.h" -#include "json-schema-to-grammar.h" -#include "log.h" -#include "sampling.h" -#include "speculative.h" -#include "preset.h" - -// fix problem with std::min and std::max -#if defined(_WIN32) -#define WIN32_LEAN_AND_MEAN -#ifndef NOMINMAX -# define NOMINMAX -#endif -#include -#endif - -#define JSON_ASSERT GGML_ASSERT -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include // for hardware_concurrency -#include - -#ifndef __EMSCRIPTEN__ -#ifdef __linux__ -#include -#elif defined(_WIN32) -# if !defined(PATH_MAX) -# define PATH_MAX MAX_PATH -# endif -#elif defined(_AIX) -#include -#else -#include -#endif -#endif - -#define LLAMA_MAX_URL_LENGTH 2084 // Maximum URL Length in Chrome: 2083 - -extern const char * LICENSES[]; - -using json = nlohmann::ordered_json; -using namespace common_arg_utils; - -static std::initializer_list mmproj_examples = { - LLAMA_EXAMPLE_MTMD, - LLAMA_EXAMPLE_SERVER, - LLAMA_EXAMPLE_CLI, -}; - -static std::string read_file(const std::string & fname) { - std::ifstream file(fname); - if (!file) { - throw std::runtime_error(string_format("error: failed to open file '%s'\n", fname.c_str())); - } - std::string content((std::istreambuf_iterator(file)), std::istreambuf_iterator()); - file.close(); - return content; -} - -static const std::vector & get_common_arg_defs() { - static const std::vector options = [] { - common_params params; - auto ctx = common_params_parser_init(params, LLAMA_EXAMPLE_SERVER, nullptr); - return ctx.options; - }(); - return options; -} - -common_arg & common_arg::set_examples(std::initializer_list examples) { - this->examples = examples; - return *this; -} - -common_arg & common_arg::set_excludes(std::initializer_list excludes) { - this->excludes = excludes; - return *this; -} - -common_arg & common_arg::set_env(const char * env) { - help = help + "\n(env: " + env + ")"; - this->env = env; - return *this; -} - -common_arg & common_arg::set_sampling() { - is_sampling = true; - return *this; -} - -common_arg & common_arg::set_spec() { - is_spec = true; - return *this; -} - -common_arg & common_arg::set_preset_only() { - is_preset_only = true; - return *this; -} - -bool common_arg::in_example(enum llama_example ex) { - return examples.find(ex) != examples.end(); -} - -bool common_arg::is_exclude(enum llama_example ex) { - return excludes.find(ex) != excludes.end(); -} - -bool common_arg::get_value_from_env(std::string & output) const { - if (env == nullptr) return false; - if (!args_neg.empty()) { - // for compatibility, we need to check LLAMA_ARG_NO_ env as well - std::string neg_env = env; - string_replace_all(neg_env, "LLAMA_ARG_", "LLAMA_ARG_NO_"); - char * neg_value = std::getenv(neg_env.c_str()); - if (neg_value) { - output = "0"; // falsey - return true; - } - } - char * value = std::getenv(env); - if (value) { - output = value; - return true; - } - return false; -} - -bool common_arg::has_value_from_env() const { - if (env != nullptr && !args_neg.empty()) { - // for compatibility, we need to check LLAMA_ARG_NO_ env as well - std::string neg_env = env; - string_replace_all(neg_env, "LLAMA_ARG_", "LLAMA_ARG_NO_"); - if (std::getenv(neg_env.c_str())) { - return true; - } - } - return env != nullptr && std::getenv(env); -} - -static std::vector break_str_into_lines(std::string input, size_t max_char_per_line) { - std::vector result; - std::istringstream iss(input); - std::string line; - auto add_line = [&](const std::string& l) { - if (l.length() <= max_char_per_line) { - result.push_back(l); - } else { - std::istringstream line_stream(l); - std::string word, current_line; - while (line_stream >> word) { - if (current_line.length() + !current_line.empty() + word.length() > max_char_per_line) { - if (!current_line.empty()) result.push_back(current_line); - current_line = word; - } else { - current_line += (!current_line.empty() ? " " : "") + word; - } - } - if (!current_line.empty()) result.push_back(current_line); - } - }; - while (std::getline(iss, line)) { - add_line(line); - } - return result; -} - -std::string common_arg::to_string() const { - // params for printing to console - const static int n_leading_spaces = 40; - const static int n_char_per_line_help = 70; // TODO: detect this based on current console - std::string leading_spaces(n_leading_spaces, ' '); - - std::ostringstream ss; - auto all_args = get_args(); // also contains args_neg - for (const auto & arg : all_args) { - if (arg == all_args.front()) { - if (all_args.size() == 1) { - ss << arg; - } else { - // first arg is usually abbreviation, we need padding to make it more beautiful - auto tmp = std::string(arg) + ", "; - auto spaces = std::string(std::max(0, 7 - (int)tmp.size()), ' '); - ss << tmp << spaces; - } - } else { - ss << arg << (arg != all_args.back() ? ", " : ""); - } - } - if (value_hint) ss << " " << value_hint; - if (value_hint_2) ss << " " << value_hint_2; - if (ss.tellp() > n_leading_spaces - 3) { - // current line is too long, add new line - ss << "\n" << leading_spaces; - } else { - // padding between arg and help, same line - ss << std::string(leading_spaces.size() - ss.tellp(), ' '); - } - const auto help_lines = break_str_into_lines(help, n_char_per_line_help); - for (const auto & line : help_lines) { - ss << (&line == &help_lines.front() ? "" : leading_spaces) << line << "\n"; - } - return ss.str(); -} - -std::vector common_arg::get_args() const { - std::vector result; - for (const auto & arg : args) { - result.push_back(std::string(arg)); - } - for (const auto & arg : args_neg) { - result.push_back(std::string(arg)); - } - return result; -} - -std::vector common_arg::get_env() const { - std::vector result; - if (env) { - result.push_back(std::string(env)); - } - if (!args_neg.empty() && env) { - // for compatibility, we need to add LLAMA_ARG_NO_ variant - std::string neg_env = env; - string_replace_all(neg_env, "LLAMA_ARG_", "LLAMA_ARG_NO_"); - result.push_back(neg_env); - } - return result; -} - -// -// utils -// - -// Helper function to parse tensor buffer override strings -static void parse_tensor_buffer_overrides(const std::string & value, std::vector & overrides) { - std::map buft_list; - for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { - auto * dev = ggml_backend_dev_get(i); - auto * buft = ggml_backend_dev_buffer_type(dev); - if (buft) { - buft_list[ggml_backend_buft_name(buft)] = buft; - } - } - - for (const auto & override : string_split(value, ',')) { - std::string::size_type pos = override.find('='); - if (pos == std::string::npos) { - throw std::invalid_argument("invalid value"); - } - std::string tensor_name = override.substr(0, pos); - std::string buffer_type = override.substr(pos + 1); - - if (buft_list.find(buffer_type) == buft_list.end()) { - printf("Available buffer types:\n"); - for (const auto & it : buft_list) { - printf(" %s\n", ggml_backend_buft_name(it.second)); - } - throw std::invalid_argument("unknown buffer type"); - } - // keep strings alive and avoid leaking memory by storing them in a static vector - static std::list buft_overrides; - buft_overrides.push_back(tensor_name); - overrides.push_back({buft_overrides.back().c_str(), buft_list.at(buffer_type)}); - } -} - -static std::string clean_file_name(const std::string & fname) { - std::string clean_fname = fname; - string_replace_all(clean_fname, "\\", "_"); - string_replace_all(clean_fname, "/", "_"); - return clean_fname; -} - -static bool common_params_handle_remote_preset(common_params & params, llama_example ex) { - GGML_ASSERT(!params.model.hf_repo.empty()); - - // the returned hf_repo is without tag - auto [hf_repo, hf_tag] = common_download_split_repo_tag(params.model.hf_repo); - - // "latest" tag (default if not specified) is translated to "default" preset - if (hf_tag == "latest") { - hf_tag = "default"; - } - - std::string model_endpoint = common_get_model_endpoint(); - auto preset_url = model_endpoint + hf_repo + "/resolve/main/preset.ini"; - - // prepare local path for caching - auto preset_fname = clean_file_name(hf_repo + "_preset.ini"); - auto preset_path = fs_get_cache_file(preset_fname); - common_download_opts opts; - opts.bearer_token = params.hf_token; - opts.offline = params.offline; - const int status = common_download_file_single(preset_url, preset_path, opts); - const bool has_preset = status >= 200 && status < 400; - - // remote preset is optional, so we don't error out if not found - if (has_preset) { - LOG_INF("applying remote preset from %s\n", preset_url.c_str()); - common_preset_context ctx(ex, /* only_remote_allowed */ true); - common_preset global; - auto remote_presets = ctx.load_from_ini(preset_path, global); - remote_presets = ctx.cascade(global, remote_presets); - if (remote_presets.find(hf_tag) != remote_presets.end()) { - common_preset preset = remote_presets.at(hf_tag); - LOG_INF("\n%s", preset.to_ini().c_str()); // to_ini already added trailing newline - preset.apply_to_params(params); - } else { - throw std::runtime_error("Remote preset.ini does not contain [" + std::string(hf_tag) + "] section"); - } - } else { - LOG_INF("%s", "no remote preset found, skipping\n"); - } - - return has_preset; -} - -struct handle_model_result { - bool found_mmproj = false; - common_params_model mmproj; -}; - -static handle_model_result common_params_handle_model(struct common_params_model & model, - const std::string & bearer_token, - bool offline) { - handle_model_result result; - - if (!model.docker_repo.empty()) { - model.path = common_docker_resolve_model(model.docker_repo); - model.name = model.docker_repo; - } else if (!model.hf_repo.empty()) { - // If -m was used with -hf, treat the model "path" as the hf_file to download - if (model.hf_file.empty() && !model.path.empty()) { - model.hf_file = model.path; - model.path = ""; - } - common_download_opts opts; - opts.bearer_token = bearer_token; - opts.offline = offline; - auto download_result = common_download_model(model, opts, true); - - if (download_result.model_path.empty()) { - LOG_ERR("error: failed to download model from Hugging Face\n"); - exit(1); - } - - model.name = model.hf_repo; - model.path = download_result.model_path; - - if (!download_result.mmproj_path.empty()) { - result.found_mmproj = true; - result.mmproj.path = download_result.mmproj_path; - } - } else if (!model.url.empty()) { - if (model.path.empty()) { - auto f = string_split(model.url, '#').front(); - f = string_split(f, '?').front(); - model.path = fs_get_cache_file(string_split(f, '/').back()); - } - - common_download_opts opts; - opts.bearer_token = bearer_token; - opts.offline = offline; - auto download_result = common_download_model(model, opts); - if (download_result.model_path.empty()) { - LOG_ERR("error: failed to download model from %s\n", model.url.c_str()); - exit(1); - } - } - - return result; -} - -const std::vector kv_cache_types = { - GGML_TYPE_F32, - GGML_TYPE_F16, - GGML_TYPE_BF16, - GGML_TYPE_Q8_0, - GGML_TYPE_Q4_0, - GGML_TYPE_Q4_1, - GGML_TYPE_IQ4_NL, - GGML_TYPE_Q5_0, - GGML_TYPE_Q5_1, -}; - -static ggml_type kv_cache_type_from_str(const std::string & s) { - for (const auto & type : kv_cache_types) { - if (ggml_type_name(type) == s) { - return type; - } - } - throw std::runtime_error("Unsupported cache type: " + s); -} - -static std::string get_all_kv_cache_types() { - std::ostringstream msg; - for (const auto & type : kv_cache_types) { - msg << ggml_type_name(type) << (&type == &kv_cache_types.back() ? "" : ", "); - } - return msg.str(); -} - -static bool parse_bool_value(const std::string & value) { - if (is_truthy(value)) { - return true; - } else if (is_falsey(value)) { - return false; - } else { - throw std::invalid_argument("invalid boolean value"); - } -} - -// -// CLI argument parsing functions -// - -static bool common_params_parse_ex(int argc, char ** argv, common_params_context & ctx_arg) { - common_params & params = ctx_arg.params; - - // setup log directly from params.verbosity: see tools/cli/cli.cpp - common_log_set_verbosity_thold(params.verbosity); - - std::unordered_map> arg_to_options; - for (auto & opt : ctx_arg.options) { - for (const auto & arg : opt.args) { - arg_to_options[arg] = {&opt, /* is_positive */ true}; - } - for (const auto & arg : opt.args_neg) { - arg_to_options[arg] = {&opt, /* is_positive */ false}; - } - } - - // handle environment variables - for (auto & opt : ctx_arg.options) { - std::string value; - if (opt.get_value_from_env(value)) { - try { - if (opt.handler_void && is_truthy(value)) { - opt.handler_void(params); - } - if (opt.handler_int) { - opt.handler_int(params, std::stoi(value)); - } - if (opt.handler_bool) { - opt.handler_bool(params, parse_bool_value(value)); - } - if (opt.handler_string) { - opt.handler_string(params, value); - continue; - } - } catch (std::exception & e) { - throw std::invalid_argument(string_format( - "error while handling environment variable \"%s\": %s\n\n", opt.env, e.what())); - } - } - } - - // handle command line arguments - auto check_arg = [&](int i) { - if (i+1 >= argc) { - throw std::invalid_argument("expected value for argument"); - } - }; - - auto parse_cli_args = [&]() { - std::set seen_args; - - for (int i = 1; i < argc; i++) { - const std::string arg_prefix = "--"; - - std::string arg = argv[i]; - if (arg.compare(0, arg_prefix.size(), arg_prefix) == 0) { - std::replace(arg.begin(), arg.end(), '_', '-'); - } - if (arg_to_options.find(arg) == arg_to_options.end()) { - throw std::invalid_argument(string_format("error: invalid argument: %s", arg.c_str())); - } - if (!seen_args.insert(arg).second) { - LOG_WRN("DEPRECATED: argument '%s' specified multiple times, use comma-separated values instead (only last value will be used)\n", arg.c_str()); - } - auto & tmp = arg_to_options[arg]; - auto opt = *tmp.first; - bool is_positive = tmp.second; - if (opt.has_value_from_env()) { - fprintf(stderr, "warn: %s environment variable is set, but will be overwritten by command line argument %s\n", opt.env, arg.c_str()); - } - try { - if (opt.handler_void) { - opt.handler_void(params); - continue; - } - if (opt.handler_bool) { - opt.handler_bool(params, is_positive); - continue; - } - - // arg with single value - check_arg(i); - std::string val = argv[++i]; - if (opt.handler_int) { - opt.handler_int(params, std::stoi(val)); - continue; - } - if (opt.handler_string) { - opt.handler_string(params, val); - continue; - } - - // arg with 2 values - check_arg(i); - std::string val2 = argv[++i]; - if (opt.handler_str_str) { - opt.handler_str_str(params, val, val2); - continue; - } - } catch (std::exception & e) { - throw std::invalid_argument(string_format( - "error while handling argument \"%s\": %s\n\n" - "usage:\n%s\n\nto show complete usage, run with -h", - arg.c_str(), e.what(), opt.to_string().c_str())); - } - } - }; - - // parse the first time to get -hf option (used for remote preset) - parse_cli_args(); - - // TODO: Remove later - try { - hf_cache::migrate_old_cache_to_hf_cache(params.hf_token, params.offline); - } catch (const std::exception & e) { - LOG_WRN("HF cache migration failed: %s\n", e.what()); - } - // export_graph_ops loads only metadata - const bool skip_model_download = ctx_arg.ex == LLAMA_EXAMPLE_EXPORT_GRAPH_OPS; - - // maybe handle remote preset - if (!params.model.hf_repo.empty() && !skip_model_download) { - std::string cli_hf_repo = params.model.hf_repo; - bool has_preset = common_params_handle_remote_preset(params, ctx_arg.ex); - - // special case: if hf_repo explicitly set by preset, we need to preserve it (ignore CLI value) - // this is useful when we have one HF repo pointing to other HF repos (one model - multiple GGUFs) - std::string preset_hf_repo = params.model.hf_repo; - bool preset_has_hf_repo = preset_hf_repo != cli_hf_repo; - - if (has_preset) { - // re-parse CLI args to override preset values - parse_cli_args(); - } - - // preserve hf_repo from preset if needed - if (preset_has_hf_repo) { - params.model.hf_repo = preset_hf_repo; - } - } - - postprocess_cpu_params(params.cpuparams, nullptr); - postprocess_cpu_params(params.cpuparams_batch, ¶ms.cpuparams); - - postprocess_cpu_params(params.speculative.draft.cpuparams, ¶ms.cpuparams); - postprocess_cpu_params(params.speculative.draft.cpuparams_batch, ¶ms.cpuparams_batch); - - if (params.prompt_cache_all && (params.interactive || params.interactive_first)) { - throw std::invalid_argument("error: --prompt-cache-all not supported in interactive mode yet\n"); - } - - // handle model and download - if (!skip_model_download) { - auto res = common_params_handle_model(params.model, params.hf_token, params.offline); - if (params.no_mmproj) { - params.mmproj = {}; - } else if (res.found_mmproj && params.mmproj.path.empty() && params.mmproj.url.empty()) { - // optionally, handle mmproj model when -hf is specified - params.mmproj = res.mmproj; - } - // only download mmproj if the current example is using it - for (const auto & ex : mmproj_examples) { - if (ctx_arg.ex == ex) { - common_params_handle_model(params.mmproj, params.hf_token, params.offline); - break; - } - } - common_params_handle_model(params.speculative.draft.mparams, params.hf_token, params.offline); - common_params_handle_model(params.vocoder.model, params.hf_token, params.offline); - } - - // model is required (except for server) - // TODO @ngxson : maybe show a list of available models in CLI in this case - if (params.model.path.empty() && ctx_arg.ex != LLAMA_EXAMPLE_SERVER && !skip_model_download && !params.usage && !params.completion) { - throw std::invalid_argument("error: --model is required\n"); - } - - if (params.escape) { - string_process_escapes(params.prompt); - string_process_escapes(params.input_prefix); - string_process_escapes(params.input_suffix); - for (auto & antiprompt : params.antiprompt) { - string_process_escapes(antiprompt); - } - for (auto & seq_breaker : params.sampling.dry_sequence_breakers) { - string_process_escapes(seq_breaker); - } - for (auto & pair : params.speculative.draft.replacements) { - string_process_escapes(pair.first); - string_process_escapes(pair.second); - } - } - - if (!params.kv_overrides.empty()) { - params.kv_overrides.emplace_back(); - params.kv_overrides.back().key[0] = 0; - } - - // pad tensor_buft_overrides for llama_params_fit: - const size_t ntbo = llama_max_tensor_buft_overrides(); - while (params.tensor_buft_overrides.size() < ntbo) { - params.tensor_buft_overrides.push_back({nullptr, nullptr}); - } - - if (!params.speculative.draft.tensor_buft_overrides.empty()) { - params.speculative.draft.tensor_buft_overrides.push_back({nullptr, nullptr}); - } - - if (!params.chat_template.empty() && !common_chat_verify_template(params.chat_template, params.use_jinja)) { - throw std::runtime_error(string_format( - "error: the supplied chat template is not supported: %s%s\n", - params.chat_template.c_str(), - params.use_jinja ? "" : "\nnote: llama.cpp was started without --jinja, we only support commonly used templates" - )); - } - - return true; -} - -static void common_params_print_usage(common_params_context & ctx_arg) { - auto print_options = [](std::vector & options) { - for (common_arg * opt : options) { - printf("%s", opt->to_string().c_str()); - } - }; - - std::vector common_options; - std::vector sampling_options; - std::vector spec_options; - std::vector specific_options; - for (auto & opt : ctx_arg.options) { - // in case multiple LLAMA_EXAMPLE_* are set, we prioritize the LLAMA_EXAMPLE_* matching current example - if (opt.is_sampling) { - sampling_options.push_back(&opt); - } else if (opt.is_spec) { - spec_options.push_back(&opt); - } else if (opt.in_example(ctx_arg.ex)) { - specific_options.push_back(&opt); - } else { - common_options.push_back(&opt); - } - } - printf("----- common params -----\n\n"); - print_options(common_options); - printf("\n\n----- sampling params -----\n\n"); - print_options(sampling_options); - printf("\n\n----- speculative params -----\n\n"); - print_options(spec_options); - // TODO: maybe convert enum llama_example to string - printf("\n\n----- example-specific params -----\n\n"); - print_options(specific_options); -} - -static void common_params_print_completion(common_params_context & ctx_arg) { - std::vector common_options; - std::vector sampling_options; - std::vector spec_options; - std::vector specific_options; - - for (auto & opt : ctx_arg.options) { - if (opt.is_sampling) { - sampling_options.push_back(&opt); - } else if (opt.is_spec) { - spec_options.push_back(&opt); - } else if (opt.in_example(ctx_arg.ex)) { - specific_options.push_back(&opt); - } else { - common_options.push_back(&opt); - } - } - - printf("_llama_completions() {\n"); - printf(" local cur prev opts\n"); - printf(" COMPREPLY=()\n"); - printf(" cur=\"${COMP_WORDS[COMP_CWORD]}\"\n"); - printf(" prev=\"${COMP_WORDS[COMP_CWORD-1]}\"\n\n"); - - printf(" opts=\""); - auto print_options = [](const std::vector & options) { - for (const common_arg * opt : options) { - for (const char * arg : opt->args) { - printf("%s ", arg); - } - } - }; - - print_options(common_options); - print_options(sampling_options); - print_options(spec_options); - print_options(specific_options); - printf("\"\n\n"); - - printf(" case \"$prev\" in\n"); - printf(" --model|-m)\n"); - printf(" COMPREPLY=( $(compgen -f -X '!*.gguf' -- \"$cur\") $(compgen -d -- \"$cur\") )\n"); - printf(" return 0\n"); - printf(" ;;\n"); - printf(" --grammar-file)\n"); - printf(" COMPREPLY=( $(compgen -f -X '!*.gbnf' -- \"$cur\") $(compgen -d -- \"$cur\") )\n"); - printf(" return 0\n"); - printf(" ;;\n"); - printf(" --chat-template-file)\n"); - printf(" COMPREPLY=( $(compgen -f -X '!*.jinja' -- \"$cur\") $(compgen -d -- \"$cur\") )\n"); - printf(" return 0\n"); - printf(" ;;\n"); - printf(" *)\n"); - printf(" COMPREPLY=( $(compgen -W \"${opts}\" -- \"$cur\") )\n"); - printf(" return 0\n"); - printf(" ;;\n"); - printf(" esac\n"); - printf("}\n\n"); - - std::set executables = { - "llama-batched", - "llama-batched-bench", - "llama-bench", - "llama-cli", - "llama-completion", - "llama-convert-llama2c-to-ggml", - "llama-cvector-generator", - "llama-debug", - "llama-diffusion-cli", - "llama-embedding", - "llama-eval-callback", - "llama-export-lora", - "llama-finetune", - "llama-fit-params", - "llama-gemma3-cli", - "llama-gen-docs", - "llama-gguf", - "llama-gguf-hash", - "llama-gguf-split", - "llama-idle", - "llama-imatrix", - "llama-llava-cli", - "llama-lookahead", - "llama-lookup", - "llama-lookup-create", - "llama-lookup-merge", - "llama-lookup-stats", - "llama-minicpmv-cli", - "llama-mtmd-cli", - "llama-parallel", - "llama-passkey", - "llama-perplexity", - "llama-q8dot", - "llama-quantize", - "llama-qwen2vl-cli", - "llama-retrieval", - "llama-save-load-state", - "llama-server", - "llama-simple", - "llama-simple-chat", - "llama-speculative", - "llama-speculative-simple", - "llama-tokenize", - "llama-tts", - "llama-vdot" - }; - - for (const auto& exe : executables) { - printf("complete -F _llama_completions %s\n", exe.c_str()); - } -} - -static std::vector parse_device_list(const std::string & value) { - std::vector devices; - auto dev_names = string_split(value, ','); - if (dev_names.empty()) { - throw std::invalid_argument("no devices specified"); - } - if (dev_names.size() == 1 && dev_names[0] == "none") { - devices.push_back(nullptr); - } else { - for (const auto & device : dev_names) { - auto * dev = ggml_backend_dev_by_name(device.c_str()); - if (!dev || ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) { - throw std::invalid_argument(string_format("invalid device: %s", device.c_str())); - } - devices.push_back(dev); - } - devices.push_back(nullptr); - } - return devices; -} - -static void add_rpc_devices(const std::string & servers) { - auto rpc_servers = string_split(servers, ','); - if (rpc_servers.empty()) { - throw std::invalid_argument("no RPC servers specified"); - } - ggml_backend_reg_t rpc_reg = ggml_backend_reg_by_name("RPC"); - if (!rpc_reg) { - throw std::invalid_argument("failed to find RPC backend"); - } - typedef ggml_backend_reg_t (*ggml_backend_rpc_add_server_t)(const char * endpoint); - ggml_backend_rpc_add_server_t ggml_backend_rpc_add_server_fn = (ggml_backend_rpc_add_server_t) ggml_backend_reg_get_proc_address(rpc_reg, "ggml_backend_rpc_add_server"); - if (!ggml_backend_rpc_add_server_fn) { - throw std::invalid_argument("failed to find RPC add server function"); - } - for (const auto & server : rpc_servers) { - auto reg = ggml_backend_rpc_add_server_fn(server.c_str()); - ggml_backend_register(reg); - } -} - -bool common_params_to_map(int argc, char ** argv, llama_example ex, std::map & out_map) { - common_params dummy_params; - common_params_context ctx_arg = common_params_parser_init(dummy_params, ex, nullptr); - - std::unordered_map arg_to_options; - for (auto & opt : ctx_arg.options) { - for (const auto & arg : opt.args) { - arg_to_options[arg] = &opt; - } - for (const auto & arg : opt.args_neg) { - arg_to_options[arg] = &opt; - } - } - - // TODO @ngxson : find a way to deduplicate this code - - // handle command line arguments - auto check_arg = [&](int i) { - if (i+1 >= argc) { - throw std::invalid_argument("expected value for argument"); - } - }; - - std::set seen_args; - - for (int i = 1; i < argc; i++) { - const std::string arg_prefix = "--"; - - std::string arg = argv[i]; - if (arg.compare(0, arg_prefix.size(), arg_prefix) == 0) { - std::replace(arg.begin(), arg.end(), '_', '-'); - } - if (arg_to_options.find(arg) == arg_to_options.end()) { - throw std::invalid_argument(string_format("error: invalid argument: %s", arg.c_str())); - } - if (!seen_args.insert(arg).second) { - LOG_WRN("DEPRECATED: argument '%s' specified multiple times, use comma-separated values instead (only last value will be used)\n", arg.c_str()); - } - auto opt = *arg_to_options[arg]; - std::string val; - if (opt.value_hint == nullptr && opt.value_hint_2 == nullptr) { - // bool arg (need to reverse the meaning for negative args) - bool is_neg = std::find(opt.args_neg.begin(), opt.args_neg.end(), arg) != opt.args_neg.end(); - val = is_neg ? "0" : "1"; - } - if (opt.value_hint != nullptr) { - // arg with single value - check_arg(i); - val = argv[++i]; - } - if (opt.value_hint_2 != nullptr) { - // TODO: support arg with 2 values - throw std::invalid_argument("error: argument with 2 values is not yet supported\n"); - } - out_map[opt] = val; - } - - return true; -} - -bool common_params_parse(int argc, char ** argv, common_params & params, llama_example ex, void(*print_usage)(int, char **)) { - auto ctx_arg = common_params_parser_init(params, ex, print_usage); - const common_params params_org = ctx_arg.params; // the example can modify the default params - - try { - if (!common_params_parse_ex(argc, argv, ctx_arg)) { - ctx_arg.params = params_org; - return false; - } - if (ctx_arg.params.usage) { - common_params_print_usage(ctx_arg); - if (ctx_arg.print_usage) { - ctx_arg.print_usage(argc, argv); - } - exit(0); - } - if (ctx_arg.params.completion) { - common_params_print_completion(ctx_arg); - exit(0); - } - params.lr.init(); - } catch (const std::invalid_argument & ex) { - fprintf(stderr, "%s\n", ex.what()); - ctx_arg.params = params_org; - return false; - } catch (std::exception & ex) { - fprintf(stderr, "%s\n", ex.what()); - exit(1); // for other exceptions, we exit with status code 1 - } - - return true; -} - -static std::string list_builtin_chat_templates() { - std::vector supported_tmpl; - int32_t res = llama_chat_builtin_templates(nullptr, 0); - supported_tmpl.resize(res); - res = llama_chat_builtin_templates(supported_tmpl.data(), supported_tmpl.size()); - std::ostringstream msg; - for (auto & tmpl : supported_tmpl) { - msg << tmpl << (&tmpl == &supported_tmpl.back() ? "" : ", "); - } - return msg.str(); -} - -bool common_arg_utils::is_truthy(const std::string & value) { - return value == "on" || value == "enabled" || value == "true" || value == "1"; -} - -bool common_arg_utils::is_falsey(const std::string & value) { - return value == "off" || value == "disabled" || value == "false" || value == "0"; -} - -bool common_arg_utils::is_autoy(const std::string & value) { - return value == "auto" || value == "-1"; -} - -// Simple CSV parser that handles quoted fields and escaped quotes -// example: -// input: value1,"value, with, commas","value with ""escaped"" quotes",value4 -// output: [value1] [value, with, commas] [value with "escaped" quotes] [value4] -static std::vector parse_csv_row(const std::string& input) { - std::vector fields; - std::string field; - bool in_quotes = false; - - for (size_t i = 0; i < input.length(); ++i) { - char ch = input[i]; - - if (ch == '"') { - if (!in_quotes) { - // start of quoted field (only valid if at beginning of field) - if (!field.empty()) { - // quote appeared in middle of unquoted field, treat as literal - field += '"'; - } else { - in_quotes = true; // start - } - } else { - if (i + 1 < input.length() && input[i + 1] == '"') { - // escaped quote: "" - field += '"'; - ++i; // skip the next quote - } else { - in_quotes = false; // end - } - } - } else if (ch == ',') { - if (in_quotes) { - field += ','; - } else { - fields.push_back(std::move(field)); - field.clear(); - } - } else { - field += ch; - } - } - - // Add the last field - fields.push_back(std::move(field)); - - return fields; -} - -common_params_context common_params_parser_init(common_params & params, llama_example ex, void(*print_usage)(int, char **)) { - // per-example default params - // we define here to make sure it's included in llama-gen-docs - if (ex == LLAMA_EXAMPLE_COMPLETION) { - params.use_jinja = false; // disable jinja by default - - } else if (ex == LLAMA_EXAMPLE_MTMD) { - params.use_jinja = false; // disable jinja by default - params.sampling.temp = 0.2; // lower temp by default for better quality - - } else if (ex == LLAMA_EXAMPLE_SERVER) { - params.n_parallel = -1; // auto by default - } - - params.use_color = tty_can_use_colors(); - - // load dynamic backends - ggml_backend_load_all(); - - common_params_context ctx_arg(params); - ctx_arg.print_usage = print_usage; - ctx_arg.ex = ex; - - std::string sampler_type_chars; - std::string sampler_type_names; - for (const auto & sampler : params.sampling.samplers) { - sampler_type_chars += common_sampler_type_to_chr(sampler); - sampler_type_names += common_sampler_type_to_str(sampler) + ";"; - } - if (!sampler_type_names.empty()) { - sampler_type_names.pop_back(); // remove last semicolon - } - - - /** - * filter options by example - * rules: - * - all examples inherit options from LLAMA_EXAMPLE_COMMON - * - if LLAMA_EXAMPLE_* is set (other than COMMON), we only show the option in the corresponding example - * - if both {LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_*,} are set, we will prioritize the LLAMA_EXAMPLE_* matching current example - */ - auto add_opt = [&](common_arg arg) { - if ((arg.in_example(ex) || arg.in_example(LLAMA_EXAMPLE_COMMON)) && !arg.is_exclude(ex)) { - ctx_arg.options.push_back(std::move(arg)); - } - }; - - - add_opt(common_arg( - {"-h", "--help", "--usage"}, - "print usage and exit", - [](common_params & params) { - params.usage = true; - } - )); - add_opt(common_arg( - {"--version"}, - "show version and build info", - [](common_params &) { - fprintf(stderr, "version: %d (%s)\n", llama_build_number(), llama_commit()); - fprintf(stderr, "built with %s for %s\n", llama_compiler(), llama_build_target()); - exit(0); - } - )); - add_opt(common_arg( - {"--license"}, - "show source code license and dependencies", - [](common_params &) { - for (int i = 0; LICENSES[i]; ++i) { - printf("%s\n", LICENSES[i]); - } - exit(0); - } - )); - add_opt(common_arg( - {"-cl", "--cache-list"}, - "show list of models in cache", - [](common_params &) { - auto models = common_list_cached_models(); - printf("number of models in cache: %zu\n", models.size()); - for (size_t i = 0; i < models.size(); i++) { - printf("%4zu. %s\n", i + 1, models[i].to_string().c_str()); - } - exit(0); - } - )); - add_opt(common_arg( - {"--completion-bash"}, - "print source-able bash completion script for llama.cpp", - [](common_params & params) { - params.completion = true; - } - )); - add_opt(common_arg( - {"--verbose-prompt"}, - string_format("print a verbose prompt before generation (default: %s)", params.verbose_prompt ? "true" : "false"), - [](common_params & params) { - params.verbose_prompt = true; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_EMBEDDING, LLAMA_EXAMPLE_RETRIEVAL})); - add_opt(common_arg( - {"--display-prompt"}, - {"--no-display-prompt"}, - string_format("whether to print prompt at generation (default: %s)", params.display_prompt ? "true" : "false"), - [](common_params & params, bool value) { - params.display_prompt = value; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"-co", "--color"}, "[on|off|auto]", - "Colorize output to distinguish prompt and user input from generations ('on', 'off', or 'auto', default: 'auto')\n" - "'auto' enables colors when output is to a terminal", - [](common_params & params, const std::string & value) { - if (is_truthy(value)) { - params.use_color = true; - } else if (is_falsey(value)) { - params.use_color = false; - } else if (is_autoy(value)) { - params.use_color = tty_can_use_colors(); - } else { - throw std::invalid_argument( - string_format("error: unknown value for --color: '%s'\n", value.c_str())); - } - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP})); - add_opt(common_arg( - {"-t", "--threads"}, "N", - string_format("number of CPU threads to use during generation (default: %d)", params.cpuparams.n_threads), - [](common_params & params, int value) { - params.cpuparams.n_threads = value; - if (params.cpuparams.n_threads <= 0) { - params.cpuparams.n_threads = std::thread::hardware_concurrency(); - } - } - ).set_env("LLAMA_ARG_THREADS")); - add_opt(common_arg( - {"-tb", "--threads-batch"}, "N", - "number of threads to use during batch and prompt processing (default: same as --threads)", - [](common_params & params, int value) { - params.cpuparams_batch.n_threads = value; - if (params.cpuparams_batch.n_threads <= 0) { - params.cpuparams_batch.n_threads = std::thread::hardware_concurrency(); - } - } - )); - add_opt(common_arg( - {"-C", "--cpu-mask"}, "M", - "CPU affinity mask: arbitrarily long hex. Complements cpu-range (default: \"\")", - [](common_params & params, const std::string & mask) { - params.cpuparams.mask_valid = true; - if (!parse_cpu_mask(mask, params.cpuparams.cpumask)) { - throw std::invalid_argument("invalid cpumask"); - } - } - )); - add_opt(common_arg( - {"-Cr", "--cpu-range"}, "lo-hi", - "range of CPUs for affinity. Complements --cpu-mask", - [](common_params & params, const std::string & range) { - params.cpuparams.mask_valid = true; - if (!parse_cpu_range(range, params.cpuparams.cpumask)) { - throw std::invalid_argument("invalid range"); - } - } - )); - add_opt(common_arg( - {"--cpu-strict"}, "<0|1>", - string_format("use strict CPU placement (default: %u)\n", (unsigned) params.cpuparams.strict_cpu), - [](common_params & params, const std::string & value) { - params.cpuparams.strict_cpu = std::stoul(value); - } - )); - add_opt(common_arg( - {"--prio"}, "N", - string_format("set process/thread priority : low(-1), normal(0), medium(1), high(2), realtime(3) (default: %d)\n", params.cpuparams.priority), - [](common_params & params, int prio) { - if (prio < GGML_SCHED_PRIO_LOW || prio > GGML_SCHED_PRIO_REALTIME) { - throw std::invalid_argument("invalid value"); - } - params.cpuparams.priority = (enum ggml_sched_priority) prio; - } - )); - add_opt(common_arg( - {"--poll"}, "<0...100>", - string_format("use polling level to wait for work (0 - no polling, default: %u)\n", (unsigned) params.cpuparams.poll), - [](common_params & params, const std::string & value) { - params.cpuparams.poll = std::stoul(value); - } - )); - add_opt(common_arg( - {"-Cb", "--cpu-mask-batch"}, "M", - "CPU affinity mask: arbitrarily long hex. Complements cpu-range-batch (default: same as --cpu-mask)", - [](common_params & params, const std::string & mask) { - params.cpuparams_batch.mask_valid = true; - if (!parse_cpu_mask(mask, params.cpuparams_batch.cpumask)) { - throw std::invalid_argument("invalid cpumask"); - } - } - )); - add_opt(common_arg( - {"-Crb", "--cpu-range-batch"}, "lo-hi", - "ranges of CPUs for affinity. Complements --cpu-mask-batch", - [](common_params & params, const std::string & range) { - params.cpuparams_batch.mask_valid = true; - if (!parse_cpu_range(range, params.cpuparams_batch.cpumask)) { - throw std::invalid_argument("invalid range"); - } - } - )); - add_opt(common_arg( - {"--cpu-strict-batch"}, "<0|1>", - "use strict CPU placement (default: same as --cpu-strict)", - [](common_params & params, int value) { - params.cpuparams_batch.strict_cpu = value; - } - )); - add_opt(common_arg( - {"--prio-batch"}, "N", - string_format("set process/thread priority : 0-normal, 1-medium, 2-high, 3-realtime (default: %d)\n", params.cpuparams_batch.priority), - [](common_params & params, int prio) { - if (prio < 0 || prio > 3) { - throw std::invalid_argument("invalid value"); - } - params.cpuparams_batch.priority = (enum ggml_sched_priority) prio; - } - )); - add_opt(common_arg( - {"--poll-batch"}, "<0|1>", - "use polling to wait for work (default: same as --poll)", - [](common_params & params, int value) { - params.cpuparams_batch.poll = value; - } - )); - add_opt(common_arg( - {"-lcs", "--lookup-cache-static"}, "FNAME", - "path to static lookup cache to use for lookup decoding (not updated by generation)", - [](common_params & params, const std::string & value) { - params.speculative.ngram_cache.lookup_cache_static = value; - } - ).set_examples({LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"-lcd", "--lookup-cache-dynamic"}, "FNAME", - "path to dynamic lookup cache to use for lookup decoding (updated by generation)", - [](common_params & params, const std::string & value) { - params.speculative.ngram_cache.lookup_cache_dynamic = value; - } - ).set_examples({LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"-c", "--ctx-size"}, "N", - string_format("size of the prompt context (default: %d, 0 = loaded from model)", params.n_ctx), - [](common_params & params, int value) { - params.n_ctx = value; - if (value == 0) { - // disable context reduction in llama_params_fit if the user explicitly requests the full context size: - params.fit_params_min_ctx = UINT32_MAX; - } - } - ).set_env("LLAMA_ARG_CTX_SIZE")); - add_opt(common_arg( - {"-n", "--predict", "--n-predict"}, "N", - string_format( - ex == LLAMA_EXAMPLE_COMPLETION - ? "number of tokens to predict (default: %d, -1 = infinity, -2 = until context filled)" - : "number of tokens to predict (default: %d, -1 = infinity)", - params.n_predict), - [](common_params & params, int value) { - params.n_predict = value; - } - ).set_env("LLAMA_ARG_N_PREDICT")); - add_opt(common_arg( - {"-b", "--batch-size"}, "N", - string_format("logical maximum batch size (default: %d)", params.n_batch), - [](common_params & params, int value) { - params.n_batch = value; - } - ).set_env("LLAMA_ARG_BATCH")); - add_opt(common_arg( - {"-ub", "--ubatch-size"}, "N", - string_format("physical maximum batch size (default: %d)", params.n_ubatch), - [](common_params & params, int value) { - params.n_ubatch = value; - } - ).set_env("LLAMA_ARG_UBATCH")); - add_opt(common_arg( - {"--keep"}, "N", - string_format("number of tokens to keep from the initial prompt (default: %d, -1 = all)", params.n_keep), - [](common_params & params, int value) { - params.n_keep = value; - } - )); - add_opt(common_arg( - {"--swa-full"}, - string_format("use full-size SWA cache (default: %s)\n" - "[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)", params.swa_full ? "true" : "false"), - [](common_params & params) { - params.swa_full = true; - } - ).set_env("LLAMA_ARG_SWA_FULL")); - add_opt(common_arg( - {"-ctxcp", "--ctx-checkpoints", "--swa-checkpoints"}, "N", - string_format("max number of context checkpoints to create per slot (default: %d)" - "[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)", params.n_ctx_checkpoints), - [](common_params & params, int value) { - params.n_ctx_checkpoints = value; - } - ).set_env("LLAMA_ARG_CTX_CHECKPOINTS").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"-cpent", "--checkpoint-every-n-tokens"}, "N", - string_format("create a checkpoint every n tokens during prefill (processing), -1 to disable (default: %d)", params.checkpoint_every_nt), - [](common_params & params, int value) { - params.checkpoint_every_nt = value; - } - ).set_env("LLAMA_ARG_CHECKPOINT_EVERY_NT").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"-cram", "--cache-ram"}, "N", - string_format("set the maximum cache size in MiB (default: %d, -1 - no limit, 0 - disable)" - "[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)", params.cache_ram_mib), - [](common_params & params, int value) { - params.cache_ram_mib = value; - } - ).set_env("LLAMA_ARG_CACHE_RAM").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"-kvu", "--kv-unified"}, - {"-no-kvu", "--no-kv-unified"}, - "use single unified KV buffer shared across all sequences (default: enabled if number of slots is auto)", - [](common_params & params, bool value) { - params.kv_unified = value; - } - ).set_env("LLAMA_ARG_KV_UNIFIED").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_BATCHED, LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL})); - add_opt(common_arg( - {"--cache-idle-slots"}, - {"--no-cache-idle-slots"}, - "save and clear idle slots on new task (default: enabled, requires unified KV and cache-ram)", - [](common_params & params, bool value) { - params.cache_idle_slots = value; - } - ).set_env("LLAMA_ARG_CACHE_IDLE_SLOTS").set_examples({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"--context-shift"}, - {"--no-context-shift"}, - string_format("whether to use context shift on infinite text generation (default: %s)", params.ctx_shift ? "enabled" : "disabled"), - [](common_params & params, bool value) { - params.ctx_shift = value; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_PERPLEXITY}).set_env("LLAMA_ARG_CONTEXT_SHIFT")); - add_opt(common_arg( - {"--chunks"}, "N", - string_format("max number of chunks to process (default: %d, -1 = all)", params.n_chunks), - [](common_params & params, int value) { - params.n_chunks = value; - } - ).set_examples({LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_RETRIEVAL})); - add_opt(common_arg({ "-fa", "--flash-attn" }, "[on|off|auto]", - string_format("set Flash Attention use ('on', 'off', or 'auto', default: '%s')", - llama_flash_attn_type_name(params.flash_attn_type)), - [](common_params & params, const std::string & value) { - if (is_truthy(value)) { - params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; - } else if (is_falsey(value)) { - params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; - } else if (is_autoy(value)) { - params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO; - } else { - throw std::runtime_error( - string_format("error: unknown value for --flash-attn: '%s'\n", value.c_str())); - } - }).set_env("LLAMA_ARG_FLASH_ATTN")); - add_opt(common_arg( - {"-p", "--prompt"}, "PROMPT", - "prompt to start generation with; for system message, use -sys", - [](common_params & params, const std::string & value) { - params.prompt = value; - } - ).set_excludes({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"-sys", "--system-prompt"}, "PROMPT", - "system prompt to use with model (if applicable, depending on chat template)", - [](common_params & params, const std::string & value) { - params.system_prompt = value; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_DIFFUSION, LLAMA_EXAMPLE_MTMD})); - add_opt(common_arg( - {"--perf"}, - {"--no-perf"}, - string_format("whether to enable internal libllama performance timings (default: %s)", params.no_perf ? "true" : "false"), - [](common_params & params, bool value) { - params.no_perf = !value; - params.sampling.no_perf = !value; - } - ).set_env("LLAMA_ARG_PERF")); - add_opt(common_arg( - {"--show-timings"}, - {"--no-show-timings"}, - string_format("whether to show timing information after each response (default: %s)", params.show_timings ? "true" : "false"), - [](common_params & params, bool value) { - params.show_timings = value; - } - ).set_examples({LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SHOW_TIMINGS")); - add_opt(common_arg( - {"-f", "--file"}, "FNAME", - "a file containing the prompt (default: none)", - [](common_params & params, const std::string & value) { - params.prompt = read_file(value); - // store the external file name in params - params.prompt_file = value; - if (!params.prompt.empty() && params.prompt.back() == '\n') { - params.prompt.pop_back(); - } - } - ).set_excludes({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"-sysf", "--system-prompt-file"}, "FNAME", - "a file containing the system prompt (default: none)", - [](common_params & params, const std::string & value) { - params.system_prompt = read_file(value); - if (!params.system_prompt.empty() && params.system_prompt.back() == '\n') { - params.system_prompt.pop_back(); - } - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_DIFFUSION})); - add_opt(common_arg( - {"--in-file"}, "FNAME", - "an input file (use comma-separated values to specify multiple files)", - [](common_params & params, const std::string & value) { - for (const auto & item : parse_csv_row(value)) { - std::ifstream file(item); - if (!file) { - throw std::runtime_error(string_format("error: failed to open file '%s'\n", item.c_str())); - } - params.in_files.push_back(item); - } - } - ).set_examples({LLAMA_EXAMPLE_IMATRIX})); - add_opt(common_arg( - {"-bf", "--binary-file"}, "FNAME", - "binary file containing the prompt (default: none)", - [](common_params & params, const std::string & value) { - std::ifstream file(value, std::ios::binary); - if (!file) { - throw std::runtime_error(string_format("error: failed to open file '%s'\n", value.c_str())); - } - // store the external file name in params - params.prompt_file = value; - std::ostringstream ss; - ss << file.rdbuf(); - params.prompt = ss.str(); - fprintf(stderr, "Read %zu bytes from binary file %s\n", params.prompt.size(), value.c_str()); - } - ).set_excludes({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"-e", "--escape"}, - {"--no-escape"}, - string_format("whether to process escapes sequences (\\n, \\r, \\t, \\', \\\", \\\\) (default: %s)", params.escape ? "true" : "false"), - [](common_params & params, bool value) { - params.escape = value; - } - )); - add_opt(common_arg( - {"-ptc", "--print-token-count"}, "N", - string_format("print token count every N tokens (default: %d)", params.n_print), - [](common_params & params, int value) { - params.n_print = value; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION})); - add_opt(common_arg( - {"--prompt-cache"}, "FNAME", - "file to cache prompt state for faster startup (default: none)", - [](common_params & params, const std::string & value) { - params.path_prompt_cache = value; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION})); - add_opt(common_arg( - {"--prompt-cache-all"}, - "if specified, saves user input and generations to cache as well\n", - [](common_params & params) { - params.prompt_cache_all = true; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION})); - add_opt(common_arg( - {"--prompt-cache-ro"}, - "if specified, uses the prompt cache but does not update it", - [](common_params & params) { - params.prompt_cache_ro = true; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION})); - add_opt(common_arg( - {"-r", "--reverse-prompt"}, "PROMPT", - "halt generation at PROMPT, return control in interactive mode\n", - [](common_params & params, const std::string & value) { - params.antiprompt.emplace_back(value); - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"-sp", "--special"}, - string_format("special tokens output enabled (default: %s)", params.special ? "true" : "false"), - [](common_params & params) { - params.special = true; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"-cnv", "--conversation"}, - {"-no-cnv", "--no-conversation"}, - "whether to run in conversation mode:\n" - "- does not print special tokens and suffix/prefix\n" - "- interactive mode is also enabled\n" - "(default: auto enabled if chat template is available)", - [](common_params & params, bool value) { - params.conversation_mode = value ? COMMON_CONVERSATION_MODE_ENABLED : COMMON_CONVERSATION_MODE_DISABLED; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"-st", "--single-turn"}, - "run conversation for a single turn only, then exit when done\n" - "will not be interactive if first turn is predefined with --prompt\n" - "(default: false)", - [](common_params & params) { - params.single_turn = true; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"-i", "--interactive"}, - string_format("run in interactive mode (default: %s)", params.interactive ? "true" : "false"), - [](common_params & params) { - params.interactive = true; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION})); - add_opt(common_arg( - {"-if", "--interactive-first"}, - string_format("run in interactive mode and wait for input right away (default: %s)", params.interactive_first ? "true" : "false"), - [](common_params & params) { - params.interactive_first = true; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION})); - add_opt(common_arg( - {"-mli", "--multiline-input"}, - "allows you to write or paste multiple lines without ending each in '\\'", - [](common_params & params) { - params.multiline_input = true; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--in-prefix-bos"}, - "prefix BOS to user inputs, preceding the `--in-prefix` string", - [](common_params & params) { - params.input_prefix_bos = true; - params.enable_chat_template = false; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION})); - add_opt(common_arg( - {"--in-prefix"}, "STRING", - "string to prefix user inputs with (default: empty)", - [](common_params & params, const std::string & value) { - params.input_prefix = value; - params.enable_chat_template = false; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION})); - add_opt(common_arg( - {"--in-suffix"}, "STRING", - "string to suffix after user inputs with (default: empty)", - [](common_params & params, const std::string & value) { - params.input_suffix = value; - params.enable_chat_template = false; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION})); - add_opt(common_arg( - {"--warmup"}, - {"--no-warmup"}, - string_format("whether to perform warmup with an empty run (default: %s)", params.warmup ? "enabled" : "disabled"), - [](common_params & params, bool value) { - params.warmup = value; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_MTMD, LLAMA_EXAMPLE_EMBEDDING, LLAMA_EXAMPLE_RETRIEVAL, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_DEBUG})); - add_opt(common_arg( - {"--spm-infill"}, - string_format( - "use Suffix/Prefix/Middle pattern for infill (instead of Prefix/Suffix/Middle) as some models prefer this. (default: %s)", - params.spm_infill ? "enabled" : "disabled" - ), - [](common_params & params) { - params.spm_infill = true; - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"--samplers"}, "SAMPLERS", - string_format("samplers that will be used for generation in the order, separated by \';\'\n(default: %s)", sampler_type_names.c_str()), - [](common_params & params, const std::string & value) { - const auto sampler_names = string_split(value, ';'); - params.sampling.samplers = common_sampler_types_from_names(sampler_names, true); - params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_SAMPLERS; - } - ).set_sampling()); - add_opt(common_arg( - {"-s", "--seed"}, "SEED", - string_format("RNG seed (default: %d, use random seed for %d)", params.sampling.seed, LLAMA_DEFAULT_SEED), - [](common_params & params, const std::string & value) { - params.sampling.seed = std::stoul(value); - } - ).set_sampling()); - add_opt(common_arg( - {"--sampler-seq", "--sampling-seq"}, "SEQUENCE", - string_format("simplified sequence for samplers that will be used (default: %s)", sampler_type_chars.c_str()), - [](common_params & params, const std::string & value) { - params.sampling.samplers = common_sampler_types_from_chars(value); - } - ).set_sampling()); - add_opt(common_arg( - {"--ignore-eos"}, - "ignore end of stream token and continue generating (implies --logit-bias EOS-inf)", - [](common_params & params) { - params.sampling.ignore_eos = true; - } - ).set_sampling()); - add_opt(common_arg( - {"--temp", "--temperature"}, "N", - string_format("temperature (default: %.2f)", (double)params.sampling.temp), - [](common_params & params, const std::string & value) { - params.sampling.temp = std::stof(value); - params.sampling.temp = std::max(params.sampling.temp, 0.0f); - params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TEMP; - } - ).set_sampling()); - add_opt(common_arg( - {"--top-k"}, "N", - string_format("top-k sampling (default: %d, 0 = disabled)", params.sampling.top_k), - [](common_params & params, int value) { - params.sampling.top_k = value; - params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TOP_K; - } - ).set_sampling().set_env("LLAMA_ARG_TOP_K")); - add_opt(common_arg( - {"--top-p"}, "N", - string_format("top-p sampling (default: %.2f, 1.0 = disabled)", (double)params.sampling.top_p), - [](common_params & params, const std::string & value) { - params.sampling.top_p = std::stof(value); - params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TOP_P; - } - ).set_sampling()); - add_opt(common_arg( - {"--min-p"}, "N", - string_format("min-p sampling (default: %.2f, 0.0 = disabled)", (double)params.sampling.min_p), - [](common_params & params, const std::string & value) { - params.sampling.min_p = std::stof(value); - params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIN_P; - } - ).set_sampling()); - add_opt(common_arg( - {"--top-nsigma", "--top-n-sigma"}, "N", - string_format("top-n-sigma sampling (default: %.2f, -1.0 = disabled)", params.sampling.top_n_sigma), - [](common_params & params, const std::string & value) { - params.sampling.top_n_sigma = std::stof(value); - } - ).set_sampling()); - add_opt(common_arg( - {"--xtc-probability"}, "N", - string_format("xtc probability (default: %.2f, 0.0 = disabled)", (double)params.sampling.xtc_probability), - [](common_params & params, const std::string & value) { - params.sampling.xtc_probability = std::stof(value); - params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_XTC_PROBABILITY; - } - ).set_sampling()); - add_opt(common_arg( - {"--xtc-threshold"}, "N", - string_format("xtc threshold (default: %.2f, 1.0 = disabled)", (double)params.sampling.xtc_threshold), - [](common_params & params, const std::string & value) { - params.sampling.xtc_threshold = std::stof(value); - params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_XTC_THRESHOLD; - } - ).set_sampling()); - add_opt(common_arg( - {"--typical", "--typical-p"}, "N", - string_format("locally typical sampling, parameter p (default: %.2f, 1.0 = disabled)", (double)params.sampling.typ_p), - [](common_params & params, const std::string & value) { - params.sampling.typ_p = std::stof(value); - } - ).set_sampling()); - add_opt(common_arg( - {"--repeat-last-n"}, "N", - string_format("last n tokens to consider for penalize (default: %d, 0 = disabled, -1 = ctx_size)", params.sampling.penalty_last_n), - [](common_params & params, int value) { - if (value < -1) { - throw std::runtime_error(string_format("error: invalid repeat-last-n = %d\n", value)); - } - params.sampling.penalty_last_n = value; - params.sampling.n_prev = std::max(params.sampling.n_prev, params.sampling.penalty_last_n); - params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_LAST_N; - } - ).set_sampling()); - add_opt(common_arg( - {"--repeat-penalty"}, "N", - string_format("penalize repeat sequence of tokens (default: %.2f, 1.0 = disabled)", (double)params.sampling.penalty_repeat), - [](common_params & params, const std::string & value) { - params.sampling.penalty_repeat = std::stof(value); - params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_REPEAT; - } - ).set_sampling()); - add_opt(common_arg( - {"--presence-penalty"}, "N", - string_format("repeat alpha presence penalty (default: %.2f, 0.0 = disabled)", (double)params.sampling.penalty_present), - [](common_params & params, const std::string & value) { - params.sampling.penalty_present = std::stof(value); - } - ).set_sampling()); - add_opt(common_arg( - {"--frequency-penalty"}, "N", - string_format("repeat alpha frequency penalty (default: %.2f, 0.0 = disabled)", (double)params.sampling.penalty_freq), - [](common_params & params, const std::string & value) { - params.sampling.penalty_freq = std::stof(value); - } - ).set_sampling()); - add_opt(common_arg( - {"--dry-multiplier"}, "N", - string_format("set DRY sampling multiplier (default: %.2f, 0.0 = disabled)", (double)params.sampling.dry_multiplier), - [](common_params & params, const std::string & value) { - params.sampling.dry_multiplier = std::stof(value); - } - ).set_sampling()); - add_opt(common_arg( - {"--dry-base"}, "N", - string_format("set DRY sampling base value (default: %.2f)", (double)params.sampling.dry_base), - [](common_params & params, const std::string & value) { - float potential_base = std::stof(value); - if (potential_base >= 1.0f) - { - params.sampling.dry_base = potential_base; - } - } - ).set_sampling()); - add_opt(common_arg( - {"--dry-allowed-length"}, "N", - string_format("set allowed length for DRY sampling (default: %d)", params.sampling.dry_allowed_length), - [](common_params & params, int value) { - params.sampling.dry_allowed_length = value; - } - ).set_sampling()); - add_opt(common_arg( - {"--dry-penalty-last-n"}, "N", - string_format("set DRY penalty for the last n tokens (default: %d, 0 = disable, -1 = context size)", params.sampling.dry_penalty_last_n), - [](common_params & params, int value) { - if (value < -1) { - throw std::runtime_error(string_format("error: invalid dry-penalty-last-n = %d\n", value)); - } - params.sampling.dry_penalty_last_n = value; - } - ).set_sampling()); - add_opt(common_arg( - {"--dry-sequence-breaker"}, "STRING", - string_format("add sequence breaker for DRY sampling, clearing out default breakers (%s) in the process; use \"none\" to not use any sequence breakers\n", - params.sampling.dry_sequence_breakers.empty() ? "none" : - std::accumulate(std::next(params.sampling.dry_sequence_breakers.begin()), - params.sampling.dry_sequence_breakers.end(), - std::string("'") + (params.sampling.dry_sequence_breakers[0] == "\n" ? "\\n" : params.sampling.dry_sequence_breakers[0]) + "'", - [](const std::string& a, const std::string& b) { - std::string formatted_b = (b == "\n") ? "\\n" : b; - return a + ", '" + formatted_b + "'"; - }).c_str()), - [](common_params & params, const std::string & value) { - static bool defaults_cleared = false; - - if (!defaults_cleared) { - params.sampling.dry_sequence_breakers.clear(); - defaults_cleared = true; - } - - if (value == "none") { - params.sampling.dry_sequence_breakers.clear(); - } else { - params.sampling.dry_sequence_breakers.emplace_back(value); - } - } - ).set_sampling()); - add_opt(common_arg( - {"--adaptive-target"}, "N", - string_format("adaptive-p: select tokens near this probability (valid range 0.0 " - "to 1.0; negative = disabled) (default: %.2f)\n" - "[(more info)](https://github.com/ggml-org/llama.cpp/pull/17927)", - (double)params.sampling.adaptive_target), - [](common_params & params, const std::string & value) { - params.sampling.adaptive_target = std::stof(value); - } - ).set_sampling()); - add_opt(common_arg( - {"--adaptive-decay"}, "N", - string_format("adaptive-p: decay rate for target adaptation over time. lower values " - "are more reactive, higher values are more stable.\n" - "(valid range 0.0 to 0.99) (default: %.2f)", - (double)params.sampling.adaptive_decay), - [](common_params & params, const std::string & value) { - params.sampling.adaptive_decay = std::stof(value); - } - ).set_sampling()); - add_opt(common_arg( - {"--dynatemp-range"}, "N", - string_format("dynamic temperature range (default: %.2f, 0.0 = disabled)", (double)params.sampling.dynatemp_range), - [](common_params & params, const std::string & value) { - params.sampling.dynatemp_range = std::stof(value); - } - ).set_sampling()); - add_opt(common_arg( - {"--dynatemp-exp"}, "N", - string_format("dynamic temperature exponent (default: %.2f)", (double)params.sampling.dynatemp_exponent), - [](common_params & params, const std::string & value) { - params.sampling.dynatemp_exponent = std::stof(value); - } - ).set_sampling()); - add_opt(common_arg( - {"--mirostat"}, "N", - string_format("use Mirostat sampling.\nTop K, Nucleus and Locally Typical samplers are ignored if used.\n" - "(default: %d, 0 = disabled, 1 = Mirostat, 2 = Mirostat 2.0)", params.sampling.mirostat), - [](common_params & params, int value) { - params.sampling.mirostat = value; - params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT; - } - ).set_sampling()); - add_opt(common_arg( - {"--mirostat-lr"}, "N", - string_format("Mirostat learning rate, parameter eta (default: %.2f)", (double)params.sampling.mirostat_eta), - [](common_params & params, const std::string & value) { - params.sampling.mirostat_eta = std::stof(value); - params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_ETA; - } - ).set_sampling()); - add_opt(common_arg( - {"--mirostat-ent"}, "N", - string_format("Mirostat target entropy, parameter tau (default: %.2f)", (double)params.sampling.mirostat_tau), - [](common_params & params, const std::string & value) { - params.sampling.mirostat_tau = std::stof(value); - params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_TAU; - } - ).set_sampling()); - add_opt(common_arg( - {"-l", "--logit-bias"}, "TOKEN_ID(+/-)BIAS", - "modifies the likelihood of token appearing in the completion,\n" - "i.e. `--logit-bias 15043+1` to increase likelihood of token ' Hello',\n" - "or `--logit-bias 15043-1` to decrease likelihood of token ' Hello'", - [](common_params & params, const std::string & value) { - std::stringstream ss(value); - llama_token key; - char sign; - std::string value_str; - try { - if (ss >> key && ss >> sign && std::getline(ss, value_str) && (sign == '+' || sign == '-')) { - const float bias = std::stof(value_str) * ((sign == '-') ? -1.0f : 1.0f); - params.sampling.logit_bias.push_back({key, bias}); - } else { - throw std::invalid_argument("invalid input format"); - } - } catch (const std::exception&) { - throw std::invalid_argument("invalid input format"); - } - } - ).set_sampling()); - add_opt(common_arg( - {"--grammar"}, "GRAMMAR", - "BNF-like grammar to constrain generations (see samples in grammars/ dir)", - [](common_params & params, const std::string & value) { - params.sampling.grammar = {COMMON_GRAMMAR_TYPE_USER, value}; - } - ).set_sampling()); - add_opt(common_arg( - {"--grammar-file"}, "FNAME", - "file to read grammar from", - [](common_params & params, const std::string & value) { - params.sampling.grammar = {COMMON_GRAMMAR_TYPE_USER, read_file(value)}; - } - ).set_sampling()); - add_opt(common_arg( - {"-j", "--json-schema"}, "SCHEMA", - "JSON schema to constrain generations (https://json-schema.org/), e.g. `{}` for any JSON object\nFor schemas w/ external $refs, use --grammar + example/json_schema_to_grammar.py instead", - [](common_params & params, const std::string & value) { - params.sampling.grammar = {COMMON_GRAMMAR_TYPE_OUTPUT_FORMAT, json_schema_to_grammar(json::parse(value))}; - } - ).set_sampling()); - add_opt(common_arg( - {"-jf", "--json-schema-file"}, "FILE", - "File containing a JSON schema to constrain generations (https://json-schema.org/), e.g. `{}` for any JSON object\nFor schemas w/ external $refs, use --grammar + example/json_schema_to_grammar.py instead", - [](common_params & params, const std::string & value) { - std::ifstream file(value); - if (!file) { - throw std::runtime_error(string_format("error: failed to open file '%s'\n", value.c_str())); - } - std::string schema; - std::copy( - std::istreambuf_iterator(file), - std::istreambuf_iterator(), - std::back_inserter(schema) - ); - params.sampling.grammar = {COMMON_GRAMMAR_TYPE_OUTPUT_FORMAT, json_schema_to_grammar(json::parse(schema))}; - } - ).set_sampling()); - add_opt(common_arg( - {"-bs", "--backend-sampling"}, - "enable backend sampling (experimental) (default: disabled)", - [](common_params & params) { - params.sampling.backend_sampling = true; - } - ).set_sampling().set_env("LLAMA_ARG_BACKEND_SAMPLING")); - add_opt(common_arg( - {"--pooling"}, "{none,mean,cls,last,rank}", - "pooling type for embeddings, use model default if unspecified", - [](common_params & params, const std::string & value) { - /**/ if (value == "none") { params.pooling_type = LLAMA_POOLING_TYPE_NONE; } - else if (value == "mean") { params.pooling_type = LLAMA_POOLING_TYPE_MEAN; } - else if (value == "cls") { params.pooling_type = LLAMA_POOLING_TYPE_CLS; } - else if (value == "last") { params.pooling_type = LLAMA_POOLING_TYPE_LAST; } - else if (value == "rank") { params.pooling_type = LLAMA_POOLING_TYPE_RANK; } - else { throw std::invalid_argument("invalid value"); } - } - ).set_examples({LLAMA_EXAMPLE_EMBEDDING, LLAMA_EXAMPLE_RETRIEVAL, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_DEBUG}).set_env("LLAMA_ARG_POOLING")); - add_opt(common_arg( - {"--attention"}, "{causal,non-causal}", - "attention type for embeddings, use model default if unspecified", - [](common_params & params, const std::string & value) { - /**/ if (value == "causal") { params.attention_type = LLAMA_ATTENTION_TYPE_CAUSAL; } - else if (value == "non-causal") { params.attention_type = LLAMA_ATTENTION_TYPE_NON_CAUSAL; } - else { throw std::invalid_argument("invalid value"); } - } - ).set_examples({LLAMA_EXAMPLE_EMBEDDING})); - add_opt(common_arg( - {"--rope-scaling"}, "{none,linear,yarn}", - "RoPE frequency scaling method, defaults to linear unless specified by the model", - [](common_params & params, const std::string & value) { - /**/ if (value == "none") { params.rope_scaling_type = LLAMA_ROPE_SCALING_TYPE_NONE; } - else if (value == "linear") { params.rope_scaling_type = LLAMA_ROPE_SCALING_TYPE_LINEAR; } - else if (value == "yarn") { params.rope_scaling_type = LLAMA_ROPE_SCALING_TYPE_YARN; } - else { throw std::invalid_argument("invalid value"); } - } - ).set_env("LLAMA_ARG_ROPE_SCALING_TYPE")); - add_opt(common_arg( - {"--rope-scale"}, "N", - "RoPE context scaling factor, expands context by a factor of N", - [](common_params & params, const std::string & value) { - params.rope_freq_scale = 1.0f / std::stof(value); - } - ).set_env("LLAMA_ARG_ROPE_SCALE")); - add_opt(common_arg( - {"--rope-freq-base"}, "N", - "RoPE base frequency, used by NTK-aware scaling (default: loaded from model)", - [](common_params & params, const std::string & value) { - params.rope_freq_base = std::stof(value); - } - ).set_env("LLAMA_ARG_ROPE_FREQ_BASE")); - add_opt(common_arg( - {"--rope-freq-scale"}, "N", - "RoPE frequency scaling factor, expands context by a factor of 1/N", - [](common_params & params, const std::string & value) { - params.rope_freq_scale = std::stof(value); - } - ).set_env("LLAMA_ARG_ROPE_FREQ_SCALE")); - add_opt(common_arg( - {"--yarn-orig-ctx"}, "N", - string_format("YaRN: original context size of model (default: %d = model training context size)", params.yarn_orig_ctx), - [](common_params & params, int value) { - params.yarn_orig_ctx = value; - } - ).set_env("LLAMA_ARG_YARN_ORIG_CTX")); - add_opt(common_arg( - {"--yarn-ext-factor"}, "N", - string_format("YaRN: extrapolation mix factor (default: %.2f, 0.0 = full interpolation)", (double)params.yarn_ext_factor), - [](common_params & params, const std::string & value) { - params.yarn_ext_factor = std::stof(value); - } - ).set_env("LLAMA_ARG_YARN_EXT_FACTOR")); - add_opt(common_arg( - {"--yarn-attn-factor"}, "N", - string_format("YaRN: scale sqrt(t) or attention magnitude (default: %.2f)", (double)params.yarn_attn_factor), - [](common_params & params, const std::string & value) { - params.yarn_attn_factor = std::stof(value); - } - ).set_env("LLAMA_ARG_YARN_ATTN_FACTOR")); - add_opt(common_arg( - {"--yarn-beta-slow"}, "N", - string_format("YaRN: high correction dim or alpha (default: %.2f)", (double)params.yarn_beta_slow), - [](common_params & params, const std::string & value) { - params.yarn_beta_slow = std::stof(value); - } - ).set_env("LLAMA_ARG_YARN_BETA_SLOW")); - add_opt(common_arg( - {"--yarn-beta-fast"}, "N", - string_format("YaRN: low correction dim or beta (default: %.2f)", (double)params.yarn_beta_fast), - [](common_params & params, const std::string & value) { - params.yarn_beta_fast = std::stof(value); - } - ).set_env("LLAMA_ARG_YARN_BETA_FAST")); - add_opt(common_arg( - {"-gan", "--grp-attn-n"}, "N", - string_format("group-attention factor (default: %d)", params.grp_attn_n), - [](common_params & params, int value) { - params.grp_attn_n = value; - } - ).set_env("LLAMA_ARG_GRP_ATTN_N").set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_PASSKEY})); - add_opt(common_arg( - {"-gaw", "--grp-attn-w"}, "N", - string_format("group-attention width (default: %d)", params.grp_attn_w), - [](common_params & params, int value) { - params.grp_attn_w = value; - } - ).set_env("LLAMA_ARG_GRP_ATTN_W").set_examples({LLAMA_EXAMPLE_COMPLETION})); - add_opt(common_arg( - {"-kvo", "--kv-offload"}, - {"-nkvo", "--no-kv-offload"}, - string_format("whether to enable KV cache offloading (default: %s)", params.no_kv_offload ? "disabled" : "enabled"), - [](common_params & params, bool value) { - params.no_kv_offload = !value; - } - ).set_env("LLAMA_ARG_KV_OFFLOAD")); - add_opt(common_arg( - {"--repack"}, - {"-nr", "--no-repack"}, - string_format("whether to enable weight repacking (default: %s)", params.no_extra_bufts ? "disabled" : "enabled"), - [](common_params & params, bool value) { - params.no_extra_bufts = !value; - } - ).set_env("LLAMA_ARG_REPACK")); - add_opt(common_arg( - {"--no-host"}, - "bypass host buffer allowing extra buffers to be used", - [](common_params & params) { - params.no_host = true; - } - ).set_env("LLAMA_ARG_NO_HOST")); - add_opt(common_arg( - {"-ctk", "--cache-type-k"}, "TYPE", - string_format( - "KV cache data type for K\n" - "allowed values: %s\n" - "(default: %s)", - get_all_kv_cache_types().c_str(), - ggml_type_name(params.cache_type_k) - ), - [](common_params & params, const std::string & value) { - params.cache_type_k = kv_cache_type_from_str(value); - } - ).set_env("LLAMA_ARG_CACHE_TYPE_K")); - add_opt(common_arg( - {"-ctv", "--cache-type-v"}, "TYPE", - string_format( - "KV cache data type for V\n" - "allowed values: %s\n" - "(default: %s)", - get_all_kv_cache_types().c_str(), - ggml_type_name(params.cache_type_v) - ), - [](common_params & params, const std::string & value) { - params.cache_type_v = kv_cache_type_from_str(value); - } - ).set_env("LLAMA_ARG_CACHE_TYPE_V")); - add_opt(common_arg( - {"--hellaswag"}, - "compute HellaSwag score over random tasks from datafile supplied with -f", - [](common_params & params) { - params.hellaswag = true; - } - ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); - add_opt(common_arg( - {"--hellaswag-tasks"}, "N", - string_format("number of tasks to use when computing the HellaSwag score (default: %zu)", params.hellaswag_tasks), - [](common_params & params, int value) { - params.hellaswag_tasks = value; - } - ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); - add_opt(common_arg( - {"--winogrande"}, - "compute Winogrande score over random tasks from datafile supplied with -f", - [](common_params & params) { - params.winogrande = true; - } - ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); - add_opt(common_arg( - {"--winogrande-tasks"}, "N", - string_format("number of tasks to use when computing the Winogrande score (default: %zu)", params.winogrande_tasks), - [](common_params & params, int value) { - params.winogrande_tasks = value; - } - ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); - add_opt(common_arg( - {"--multiple-choice"}, - "compute multiple choice score over random tasks from datafile supplied with -f", - [](common_params & params) { - params.multiple_choice = true; - } - ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); - add_opt(common_arg( - {"--multiple-choice-tasks"}, "N", - string_format("number of tasks to use when computing the multiple choice score (default: %zu)", params.multiple_choice_tasks), - [](common_params & params, int value) { - params.multiple_choice_tasks = value; - } - ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); - add_opt(common_arg( - {"--kl-divergence"}, - "computes KL-divergence to logits provided via --kl-divergence-base", - [](common_params & params) { - params.kl_divergence = true; - } - ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); - add_opt(common_arg( - {"--save-all-logits", "--kl-divergence-base"}, "FNAME", - "set logits file", - [](common_params & params, const std::string & value) { - params.logits_file = value; - } - ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); - add_opt(common_arg( - {"--ppl-stride"}, "N", - string_format("stride for perplexity calculation (default: %d)", params.ppl_stride), - [](common_params & params, int value) { - params.ppl_stride = value; - } - ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); - add_opt(common_arg( - {"--ppl-output-type"}, "<0|1>", - string_format("output type for perplexity calculation (default: %d)", params.ppl_output_type), - [](common_params & params, int value) { - params.ppl_output_type = value; - } - ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); - add_opt(common_arg( - {"-dt", "--defrag-thold"}, "N", - string_format("KV cache defragmentation threshold (DEPRECATED)"), - [](common_params & params, const std::string & value) { - GGML_UNUSED(params); - GGML_UNUSED(value); - LOG_WRN("DEPRECATED: --defrag-thold is deprecated and no longer necessary to specify\n"); - } - ).set_env("LLAMA_ARG_DEFRAG_THOLD")); - if (ex == LLAMA_EXAMPLE_SERVER) { - // this is to make sure this option appears in the server-specific section of the help message - add_opt(common_arg( - {"-np", "--parallel"}, "N", - string_format("number of server slots (default: %d, -1 = auto)", params.n_parallel), - [](common_params & params, int value) { - if (value == 0) { - throw std::invalid_argument("error: invalid value for n_parallel\n"); - } - params.n_parallel = value; - } - ).set_env("LLAMA_ARG_N_PARALLEL").set_examples({LLAMA_EXAMPLE_SERVER})); - } else { - add_opt(common_arg( - {"-np", "--parallel"}, "N", - string_format("number of parallel sequences to decode (default: %d)", params.n_parallel), - [](common_params & params, int value) { - params.n_parallel = value; - } - ).set_env("LLAMA_ARG_N_PARALLEL")); - } - add_opt(common_arg( - {"-ns", "--sequences"}, "N", - string_format("number of sequences to decode (default: %d)", params.n_sequences), - [](common_params & params, int value) { - params.n_sequences = value; - } - ).set_examples({LLAMA_EXAMPLE_PARALLEL})); - add_opt(common_arg( - {"-cb", "--cont-batching"}, - {"-nocb", "--no-cont-batching"}, - string_format("whether to enable continuous batching (a.k.a dynamic batching) (default: %s)", params.cont_batching ? "enabled" : "disabled"), - [](common_params & params, bool value) { - params.cont_batching = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CONT_BATCHING")); - add_opt(common_arg( - {"-mm", "--mmproj"}, "FILE", - "path to a multimodal projector file. see tools/mtmd/README.md\n" - "note: if -hf is used, this argument can be omitted", - [](common_params & params, const std::string & value) { - params.mmproj.path = value; - } - ).set_examples(mmproj_examples).set_env("LLAMA_ARG_MMPROJ")); - add_opt(common_arg( - {"-mmu", "--mmproj-url"}, "URL", - "URL to a multimodal projector file. see tools/mtmd/README.md", - [](common_params & params, const std::string & value) { - params.mmproj.url = value; - } - ).set_examples(mmproj_examples).set_env("LLAMA_ARG_MMPROJ_URL")); - add_opt(common_arg( - {"--mmproj-auto"}, - {"--no-mmproj", "--no-mmproj-auto"}, - string_format("whether to use multimodal projector file (if available), useful when using -hf (default: %s)", params.no_mmproj ? "disabled" : "enabled"), - [](common_params & params, bool value) { - params.no_mmproj = !value; - } - ).set_examples(mmproj_examples).set_env("LLAMA_ARG_MMPROJ_AUTO")); - add_opt(common_arg( - {"--mmproj-offload"}, - {"--no-mmproj-offload"}, - string_format("whether to enable GPU offloading for multimodal projector (default: %s)", params.mmproj_use_gpu ? "enabled" : "disabled"), - [](common_params & params, bool value) { - params.mmproj_use_gpu = value; - } - ).set_examples(mmproj_examples).set_env("LLAMA_ARG_MMPROJ_OFFLOAD")); - add_opt(common_arg( - {"--image", "--audio"}, "FILE", - "path to an image or audio file. use with multimodal models, use comma-separated values for multiple files\n", - [](common_params & params, const std::string & value) { - for (const auto & item : parse_csv_row(value)) { - params.image.emplace_back(item); - } - } - ).set_examples({LLAMA_EXAMPLE_MTMD, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--image-min-tokens"}, "N", - "minimum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)", - [](common_params & params, int value) { - params.image_min_tokens = value; - } - ).set_examples(mmproj_examples).set_env("LLAMA_ARG_IMAGE_MIN_TOKENS")); - add_opt(common_arg( - {"--image-max-tokens"}, "N", - "maximum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)", - [](common_params & params, int value) { - params.image_max_tokens = value; - } - ).set_examples(mmproj_examples).set_env("LLAMA_ARG_IMAGE_MAX_TOKENS")); - if (llama_supports_rpc()) { - add_opt(common_arg( - {"--rpc"}, "SERVERS", - "comma separated list of RPC servers (host:port)", - [](common_params & params, const std::string & value) { - add_rpc_devices(value); - GGML_UNUSED(params); - } - ).set_env("LLAMA_ARG_RPC")); - } - add_opt(common_arg( - {"--mlock"}, - "force system to keep model in RAM rather than swapping or compressing", - [](common_params & params) { - params.use_mlock = true; - } - ).set_env("LLAMA_ARG_MLOCK")); - add_opt(common_arg( - {"--mmap"}, - {"--no-mmap"}, - string_format("whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock) (default: %s)", params.use_mmap ? "enabled" : "disabled"), - [](common_params & params, bool value) { - params.use_mmap = value; - } - ).set_env("LLAMA_ARG_MMAP")); - add_opt(common_arg( - {"-dio", "--direct-io"}, - {"-ndio", "--no-direct-io"}, - string_format("use DirectIO if available. (default: %s)", params.use_direct_io ? "enabled" : "disabled"), - [](common_params & params, bool value) { - params.use_direct_io = value; - } - ).set_env("LLAMA_ARG_DIO")); - add_opt(common_arg( - {"--numa"}, "TYPE", - "attempt optimizations that help on some NUMA systems\n" - "- distribute: spread execution evenly over all nodes\n" - "- isolate: only spawn threads on CPUs on the node that execution started on\n" - "- numactl: use the CPU map provided by numactl\n" - "if run without this previously, it is recommended to drop the system page cache before using this\n" - "see https://github.com/ggml-org/llama.cpp/issues/1437", - [](common_params & params, const std::string & value) { - /**/ if (value == "distribute" || value == "") { params.numa = GGML_NUMA_STRATEGY_DISTRIBUTE; } - else if (value == "isolate") { params.numa = GGML_NUMA_STRATEGY_ISOLATE; } - else if (value == "numactl") { params.numa = GGML_NUMA_STRATEGY_NUMACTL; } - else { throw std::invalid_argument("invalid value"); } - } - ).set_env("LLAMA_ARG_NUMA")); - add_opt(common_arg( - {"-dev", "--device"}, "", - "comma-separated list of devices to use for offloading (none = don't offload)\n" - "use --list-devices to see a list of available devices", - [](common_params & params, const std::string & value) { - params.devices = parse_device_list(value); - } - ).set_env("LLAMA_ARG_DEVICE")); - add_opt(common_arg( - {"--list-devices"}, - "print list of available devices and exit", - [](common_params &) { - std::vector devices; - for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { - auto * dev = ggml_backend_dev_get(i); - if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) { - devices.push_back(dev); - } - } - printf("Available devices:\n"); - for (auto * dev : devices) { - size_t free, total; - ggml_backend_dev_memory(dev, &free, &total); - printf(" %s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024); - } - exit(0); - } - )); - add_opt(common_arg( - {"-ot", "--override-tensor"}, "=,...", - "override tensor buffer type", [](common_params & params, const std::string & value) { - parse_tensor_buffer_overrides(value, params.tensor_buft_overrides); - } - ).set_env("LLAMA_ARG_OVERRIDE_TENSOR")); - add_opt(common_arg( - {"-cmoe", "--cpu-moe"}, - "keep all Mixture of Experts (MoE) weights in the CPU", - [](common_params & params) { - params.tensor_buft_overrides.push_back(llm_ffn_exps_cpu_override()); - } - ).set_env("LLAMA_ARG_CPU_MOE")); - add_opt(common_arg( - {"-ncmoe", "--n-cpu-moe"}, "N", - "keep the Mixture of Experts (MoE) weights of the first N layers in the CPU", - [](common_params & params, int value) { - if (value < 0) { - throw std::invalid_argument("invalid value"); - } - for (int i = 0; i < value; ++i) { - // keep strings alive and avoid leaking memory by storing them in a static vector - static std::list buft_overrides; - buft_overrides.push_back(llm_ffn_exps_block_regex(i)); - params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()}); - } - } - ).set_env("LLAMA_ARG_N_CPU_MOE")); - GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0 - add_opt(common_arg( - {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N", - string_format("max. number of layers to store in VRAM, either an exact number, 'auto', or 'all' (default: %s)", params.n_gpu_layers == -1 ? "auto" : "all"), - [](common_params & params, const std::string & value) { - if (value == "auto") { - params.n_gpu_layers = -1; - } else if (value == "all") { - params.n_gpu_layers = -2; - } else { - params.n_gpu_layers = std::stoi(value); - } - if (!llama_supports_gpu_offload()) { - fprintf(stderr, "warning: no usable GPU found, --gpu-layers option will be ignored\n"); - fprintf(stderr, "warning: one possible reason is that llama.cpp was compiled without GPU support\n"); - fprintf(stderr, "warning: consult docs/build.md for compilation instructions\n"); - } - } - ).set_env("LLAMA_ARG_N_GPU_LAYERS")); - add_opt(common_arg( - {"-sm", "--split-mode"}, "{none,layer,row,tensor}", - "how to split the model across multiple GPUs, one of:\n" - "- none: use one GPU only\n" - "- layer (default): split layers and KV across GPUs (pipelined)\n" - "- row: split weight across GPUs by rows (parallelized)\n" - "- tensor: split weights and KV across GPUs (parallelized, EXPERIMENTAL)", - [](common_params & params, const std::string & value) { - if (value == "none") { - params.split_mode = LLAMA_SPLIT_MODE_NONE; - } else if (value == "layer") { - params.split_mode = LLAMA_SPLIT_MODE_LAYER; - } else if (value == "row") { - params.split_mode = LLAMA_SPLIT_MODE_ROW; - } else if (value == "tensor") { - params.split_mode = LLAMA_SPLIT_MODE_TENSOR; - } else { - throw std::invalid_argument("invalid value"); - } - if (!llama_supports_gpu_offload()) { - fprintf(stderr, "warning: llama.cpp was compiled without support for GPU offload. Setting the split mode has no effect.\n"); - } - } - ).set_env("LLAMA_ARG_SPLIT_MODE")); - add_opt(common_arg( - {"-ts", "--tensor-split"}, "N0,N1,N2,...", - "fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1", - [](common_params & params, const std::string & value) { - std::string arg_next = value; - - // split string by , and / - const std::regex regex{ R"([,/]+)" }; - std::sregex_token_iterator it{ arg_next.begin(), arg_next.end(), regex, -1 }; - std::vector split_arg{ it, {} }; - if (split_arg.size() >= llama_max_devices()) { - throw std::invalid_argument( - string_format("got %zu input configs, but system only has %zu devices", split_arg.size(), llama_max_devices()) - ); - } - for (size_t i = 0; i < llama_max_devices(); ++i) { - if (i < split_arg.size()) { - params.tensor_split[i] = std::stof(split_arg[i]); - } else { - params.tensor_split[i] = 0.0f; - } - } - if (!llama_supports_gpu_offload()) { - fprintf(stderr, "warning: llama.cpp was compiled without support for GPU offload. Setting a tensor split has no effect.\n"); - } - } - ).set_env("LLAMA_ARG_TENSOR_SPLIT")); - add_opt(common_arg( - {"-mg", "--main-gpu"}, "INDEX", - string_format("the GPU to use for the model (with split-mode = none), or for intermediate results and KV (with split-mode = row) (default: %d)", params.main_gpu), - [](common_params & params, int value) { - params.main_gpu = value; - if (!llama_supports_gpu_offload()) { - fprintf(stderr, "warning: llama.cpp was compiled without support for GPU offload. Setting the main GPU has no effect.\n"); - } - } - ).set_env("LLAMA_ARG_MAIN_GPU")); - add_opt(common_arg( - { "-fit", "--fit" }, "[on|off]", - string_format("whether to adjust unset arguments to fit in device memory ('on' or 'off', default: '%s')", params.fit_params ? "on" : "off"), - [](common_params & params, const std::string & value) { - if (is_truthy(value)) { - params.fit_params = true; - } else if (is_falsey(value)) { - params.fit_params = false; - } else { - throw std::runtime_error( - string_format("error: unknown value for --fit: '%s'\n", value.c_str())); - } - } - ).set_env("LLAMA_ARG_FIT")); - add_opt(common_arg( - { "-fitp", "--fit-print" }, "[on|off]", - string_format("print the estimated required memory ('on' or 'off', default: '%s')", params.fit_params_print ? "on" : "off"), - [](common_params & params, const std::string & value) { - if (is_truthy(value)) { - params.fit_params_print = true; - } else if (is_falsey(value)) { - params.fit_params_print = false; - } else { - throw std::runtime_error( - string_format("error: unknown value for --fit-print: '%s'\n", value.c_str())); - } - } - ).set_examples({LLAMA_EXAMPLE_FIT_PARAMS}).set_env("LLAMA_ARG_FIT_ESTIMATE")); - add_opt(common_arg( - { "-fitt", "--fit-target" }, "MiB0,MiB1,MiB2,...", - string_format("target margin per device for --fit, comma-separated list of values, " - "single value is broadcast across all devices, default: %zu", params.fit_params_target[0]/(1024*1024)), - [](common_params & params, const std::string & value) { - std::string arg_next = value; - - // split string by , and / - const std::regex regex{ R"([,/]+)" }; - std::sregex_token_iterator it{ arg_next.begin(), arg_next.end(), regex, -1 }; - std::vector split_arg{ it, {} }; - if (split_arg.size() >= llama_max_devices()) { - throw std::invalid_argument( - string_format("got %zu input configs, but system only has %zu devices", split_arg.size(), llama_max_devices()) - ); - } - if (split_arg.size() == 1) { - std::fill(params.fit_params_target.begin(), params.fit_params_target.end(), std::stoull(split_arg[0]) * 1024*1024); - return; - } - for (size_t i = 0; i < split_arg.size(); i++) { - params.fit_params_target[i] = std::stoull(split_arg[i]) * 1024*1024; - } - } - ).set_env("LLAMA_ARG_FIT_TARGET")); - add_opt(common_arg( - { "-fitc", "--fit-ctx" }, "N", - string_format("minimum ctx size that can be set by --fit option, default: %" PRIu32, params.fit_params_min_ctx), - [](common_params & params, int value) { - params.fit_params_min_ctx = value; - } - ).set_env("LLAMA_ARG_FIT_CTX")); - add_opt(common_arg( - {"--check-tensors"}, - string_format("check model tensor data for invalid values (default: %s)", params.check_tensors ? "true" : "false"), - [](common_params & params) { - params.check_tensors = true; - } - )); - add_opt(common_arg( - {"--override-kv"}, "KEY=TYPE:VALUE,...", - "advanced option to override model metadata by key. to specify multiple overrides, either use comma-separated values.\n" - "types: int, float, bool, str. example: --override-kv tokenizer.ggml.add_bos_token=bool:false,tokenizer.ggml.add_eos_token=bool:false", - [](common_params & params, const std::string & value) { - for (const auto & item : parse_csv_row(value)) { - if (!string_parse_kv_override(item.c_str(), params.kv_overrides)) { - throw std::runtime_error(string_format("error: Invalid type for KV override: %s\n", item.c_str())); - } - } - } - )); - add_opt(common_arg( - {"--op-offload"}, - {"--no-op-offload"}, - string_format("whether to offload host tensor operations to device (default: %s)", params.no_op_offload ? "false" : "true"), - [](common_params & params, bool value) { - params.no_op_offload = !value; - } - )); - add_opt(common_arg( - {"--lora"}, "FNAME", - "path to LoRA adapter (use comma-separated values to load multiple adapters)", - [](common_params & params, const std::string & value) { - for (const auto & item : parse_csv_row(value)) { - params.lora_adapters.push_back({ item, 1.0, "", "", nullptr }); - } - } - // we define this arg on both COMMON and EXPORT_LORA, so when showing help message of export-lora, it will be categorized as "example-specific" arg - ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_EXPORT_LORA})); - add_opt(common_arg( - {"--lora-scaled"}, "FNAME:SCALE,...", - "path to LoRA adapter with user defined scaling (format: FNAME:SCALE,...)\n" - "note: use comma-separated values", - [](common_params & params, const std::string & value) { - for (const auto & item : parse_csv_row(value)) { - auto parts = string_split(item, ':'); - if (parts.size() != 2) { - throw std::invalid_argument("lora-scaled format: FNAME:SCALE"); - } - params.lora_adapters.push_back({ parts[0], std::stof(parts[1]), "", "", nullptr }); - } - } - // we define this arg on both COMMON and EXPORT_LORA, so when showing help message of export-lora, it will be categorized as "example-specific" arg - ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_EXPORT_LORA})); - add_opt(common_arg( - {"--control-vector"}, "FNAME", - "add a control vector\nnote: use comma-separated values to add multiple control vectors", - [](common_params & params, const std::string & value) { - for (const auto & item : parse_csv_row(value)) { - params.control_vectors.push_back({ 1.0f, item, }); - } - } - )); - add_opt(common_arg( - {"--control-vector-scaled"}, "FNAME:SCALE,...", - "add a control vector with user defined scaling SCALE\n" - "note: use comma-separated values (format: FNAME:SCALE,...)", - [](common_params & params, const std::string & value) { - for (const auto & item : parse_csv_row(value)) { - auto parts = string_split(item, ':'); - if (parts.size() != 2) { - throw std::invalid_argument("control-vector-scaled format: FNAME:SCALE"); - } - params.control_vectors.push_back({ std::stof(parts[1]), parts[0] }); - } - } - )); - add_opt(common_arg( - {"--control-vector-layer-range"}, "START", "END", - "layer range to apply the control vector(s) to, start and end inclusive", - [](common_params & params, const std::string & start, const std::string & end) { - params.control_vector_layer_start = std::stoi(start); - params.control_vector_layer_end = std::stoi(end); - } - )); - add_opt(common_arg( - {"-a", "--alias"}, "STRING", - "set model name aliases, comma-separated (to be used by API)", - [](common_params & params, const std::string & value) { - for (auto & alias : string_split(value, ',')) { - alias = string_strip(alias); - if (!alias.empty()) { - params.model_alias.insert(alias); - } - } - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_ALIAS")); - add_opt(common_arg( - {"--tags"}, "STRING", - "set model tags, comma-separated (informational, not used for routing)", - [](common_params & params, const std::string & value) { - for (auto & tag : string_split(value, ',')) { - tag = string_strip(tag); - if (!tag.empty()) { - params.model_tags.insert(tag); - } - } - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TAGS")); - add_opt(common_arg( - {"-m", "--model"}, "FNAME", - ex == LLAMA_EXAMPLE_EXPORT_LORA - ? "model path from which to load base model" - : "model path to load", - [](common_params & params, const std::string & value) { - params.model.path = value; - } - ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_EXPORT_LORA}).set_env("LLAMA_ARG_MODEL")); - add_opt(common_arg( - {"-mu", "--model-url"}, "MODEL_URL", - "model download url (default: unused)", - [](common_params & params, const std::string & value) { - params.model.url = value; - } - ).set_env("LLAMA_ARG_MODEL_URL")); - add_opt(common_arg( - { "-dr", "--docker-repo" }, "[/][:quant]", - "Docker Hub model repository. repo is optional, default to ai/. quant is optional, default to :latest.\n" - "example: gemma3\n" - "(default: unused)", - [](common_params & params, const std::string & value) { - params.model.docker_repo = value; - } - ).set_env("LLAMA_ARG_DOCKER_REPO")); - add_opt(common_arg( - {"-hf", "-hfr", "--hf-repo"}, "/[:quant]", - "Hugging Face model repository; quant is optional, case-insensitive, default to Q4_K_M, or falls back to the first file in the repo if Q4_K_M doesn't exist.\n" - "mmproj is also downloaded automatically if available. to disable, add --no-mmproj\n" - "example: ggml-org/GLM-4.7-Flash-GGUF:Q4_K_M\n" - "(default: unused)", - [](common_params & params, const std::string & value) { - params.model.hf_repo = value; - } - ).set_env("LLAMA_ARG_HF_REPO")); - add_opt(common_arg( - {"-hff", "--hf-file"}, "FILE", - "Hugging Face model file. If specified, it will override the quant in --hf-repo (default: unused)", - [](common_params & params, const std::string & value) { - params.model.hf_file = value; - } - ).set_env("LLAMA_ARG_HF_FILE")); - add_opt(common_arg( - {"-hfv", "-hfrv", "--hf-repo-v"}, "/[:quant]", - "Hugging Face model repository for the vocoder model (default: unused)", - [](common_params & params, const std::string & value) { - params.vocoder.model.hf_repo = value; - } - ).set_env("LLAMA_ARG_HF_REPO_V")); - add_opt(common_arg( - {"-hffv", "--hf-file-v"}, "FILE", - "Hugging Face model file for the vocoder model (default: unused)", - [](common_params & params, const std::string & value) { - params.vocoder.model.hf_file = value; - } - ).set_env("LLAMA_ARG_HF_FILE_V")); - add_opt(common_arg( - {"-hft", "--hf-token"}, "TOKEN", - "Hugging Face access token (default: value from HF_TOKEN environment variable)", - [](common_params & params, const std::string & value) { - params.hf_token = value; - } - ).set_env("HF_TOKEN")); - add_opt(common_arg( - {"--context-file"}, "FNAME", - "file to load context from (use comma-separated values to specify multiple files)", - [](common_params & params, const std::string & value) { - for (const auto & item : parse_csv_row(value)) { - std::ifstream file(item, std::ios::binary); - if (!file) { - throw std::runtime_error(string_format("error: failed to open file '%s'\n", item.c_str())); - } - params.context_files.push_back(item); - } - } - ).set_examples({LLAMA_EXAMPLE_RETRIEVAL})); - add_opt(common_arg( - {"--chunk-size"}, "N", - string_format("minimum length of embedded text chunks (default: %d)", params.chunk_size), - [](common_params & params, int value) { - params.chunk_size = value; - } - ).set_examples({LLAMA_EXAMPLE_RETRIEVAL})); - add_opt(common_arg( - {"--chunk-separator"}, "STRING", - string_format("separator between chunks (default: '%s')", params.chunk_separator.c_str()), - [](common_params & params, const std::string & value) { - params.chunk_separator = value; - } - ).set_examples({LLAMA_EXAMPLE_RETRIEVAL})); - add_opt(common_arg( - {"--junk"}, "N", - string_format("number of times to repeat the junk text (default: %d)", params.n_junk), - [](common_params & params, int value) { - params.n_junk = value; - } - ).set_examples({LLAMA_EXAMPLE_PASSKEY, LLAMA_EXAMPLE_PARALLEL})); - add_opt(common_arg( - {"--pos"}, "N", - string_format("position of the passkey in the junk text (default: %d)", params.i_pos), - [](common_params & params, int value) { - params.i_pos = value; - } - ).set_examples({LLAMA_EXAMPLE_PASSKEY})); - add_opt(common_arg( - {"-o", "--output", "--output-file"}, "FNAME", - string_format("output file (default: '%s')", params.out_file.c_str()), - [](common_params & params, const std::string & value) { - params.out_file = value; - } - ).set_examples({LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_CVECTOR_GENERATOR, LLAMA_EXAMPLE_EXPORT_LORA, LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_FINETUNE, - LLAMA_EXAMPLE_RESULTS, LLAMA_EXAMPLE_EXPORT_GRAPH_OPS})); - add_opt(common_arg( - {"-ofreq", "--output-frequency"}, "N", - string_format("output the imatrix every N iterations (default: %d)", params.n_out_freq), - [](common_params & params, int value) { - params.n_out_freq = value; - } - ).set_examples({LLAMA_EXAMPLE_IMATRIX})); - add_opt(common_arg( - {"--output-format"}, "{gguf,dat}", - string_format("output format for imatrix file (default: %s)", params.imat_dat > 0 ? "dat" : "gguf"), - [](common_params & params, const std::string & value) { - /**/ if (value == "gguf") { params.imat_dat = -1; } - else if (value == "dat") { params.imat_dat = 1; } - else { throw std::invalid_argument("invalid output format"); } - } - ).set_examples({LLAMA_EXAMPLE_IMATRIX})); - add_opt(common_arg( - {"--save-frequency"}, "N", - string_format("save an imatrix copy every N iterations (default: %d)", params.n_save_freq), - [](common_params & params, int value) { - params.n_save_freq = value; - } - ).set_examples({LLAMA_EXAMPLE_IMATRIX})); - add_opt(common_arg( - {"--process-output"}, - string_format("collect data for the output tensor (default: %s)", params.process_output ? "true" : "false"), - [](common_params & params) { - params.process_output = true; - } - ).set_examples({LLAMA_EXAMPLE_IMATRIX})); - add_opt(common_arg( - {"--ppl"}, - {"--no-ppl"}, - string_format("whether to compute perplexity (default: %s)", params.compute_ppl ? "true" : "false"), - [](common_params & params, bool value) { - params.compute_ppl = value; - } - ).set_examples({LLAMA_EXAMPLE_IMATRIX})); - add_opt(common_arg( - {"--chunk", "--from-chunk"}, "N", - string_format("start processing the input from chunk N (default: %d)", params.i_chunk), - [](common_params & params, int value) { - params.i_chunk = value; - } - ).set_examples({LLAMA_EXAMPLE_IMATRIX})); - add_opt(common_arg( - {"--show-statistics"}, - string_format("show imatrix statistics and then exit (default: %s)", params.show_statistics ? "true" : "false"), - [](common_params & params) { - params.show_statistics = true; - } - ).set_examples({LLAMA_EXAMPLE_IMATRIX})); - add_opt(common_arg( - {"--parse-special"}, - string_format("parse special tokens (chat, tool, etc) (default: %s)", params.parse_special ? "true" : "false"), - [](common_params & params) { - params.parse_special = true; - } - ).set_examples({LLAMA_EXAMPLE_IMATRIX})); - add_opt(common_arg( - {"-pps"}, - string_format("is the prompt shared across parallel sequences (default: %s)", params.is_pp_shared ? "true" : "false"), - [](common_params & params) { - params.is_pp_shared = true; - } - ).set_examples({LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL})); - add_opt(common_arg( - {"-tgs"}, - string_format("is the text generation separated across the different sequences (default: %s)", params.is_tg_separate ? "true" : "false"), - [](common_params & params) { - params.is_tg_separate = true; - } - ).set_examples({LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL})); - add_opt(common_arg( - {"-npp"}, "n0,n1,...", - "number of prompt tokens", - [](common_params & params, const std::string & value) { - auto p = string_split(value, ','); - params.n_pp.insert(params.n_pp.end(), p.begin(), p.end()); - } - ).set_examples({LLAMA_EXAMPLE_BENCH})); - add_opt(common_arg( - {"-ntg"}, "n0,n1,...", - "number of text generation tokens", - [](common_params & params, const std::string & value) { - auto p = string_split(value, ','); - params.n_tg.insert(params.n_tg.end(), p.begin(), p.end()); - } - ).set_examples({LLAMA_EXAMPLE_BENCH})); - add_opt(common_arg( - {"-npl"}, "n0,n1,...", - "number of parallel prompts", - [](common_params & params, const std::string & value) { - auto p = string_split(value, ','); - params.n_pl.insert(params.n_pl.end(), p.begin(), p.end()); - } - ).set_examples({LLAMA_EXAMPLE_BENCH})); - add_opt(common_arg( - {"--embd-normalize"}, "N", - string_format("normalisation for embeddings (default: %d) (-1=none, 0=max absolute int16, 1=taxicab, 2=euclidean, >2=p-norm)", params.embd_normalize), - [](common_params & params, int value) { - params.embd_normalize = value; - } - ).set_examples({LLAMA_EXAMPLE_EMBEDDING, LLAMA_EXAMPLE_DEBUG})); - add_opt(common_arg( - {"--embd-output-format"}, "FORMAT", - "empty = default, \"array\" = [[],[]...], \"json\" = openai style, \"json+\" = same \"json\" + cosine similarity matrix, \"raw\" = plain whitespace-delimited output (one embedding per line)", - [](common_params & params, const std::string & value) { - params.embd_out = value; - } - ).set_examples({LLAMA_EXAMPLE_EMBEDDING})); - add_opt(common_arg( - {"--embd-separator"}, "STRING", - "separator of embeddings (default \\n) for example \"<#sep#>\"", - [](common_params & params, const std::string & value) { - params.embd_sep = value; - } - ).set_examples({LLAMA_EXAMPLE_EMBEDDING})); - add_opt(common_arg( - {"--cls-separator"}, "STRING", - "separator of classification sequences (default \\t) for example \"<#seq#>\"", - [](common_params & params, const std::string & value) { - params.cls_sep = value; - } - ).set_examples({LLAMA_EXAMPLE_EMBEDDING})); - add_opt(common_arg( - {"--host"}, "HOST", - string_format("ip address to listen, or bind to an UNIX socket if the address ends with .sock (default: %s)", params.hostname.c_str()), - [](common_params & params, const std::string & value) { - params.hostname = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_HOST")); - add_opt(common_arg( - {"--port"}, "PORT", - string_format("port to listen (default: %d)", params.port), - [](common_params & params, int value) { - params.port = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_PORT")); - add_opt(common_arg( - {"--reuse-port"}, - string_format("allow multiple sockets to bind to the same port (default: %s)", params.reuse_port ? "enabled" : "disabled"), - [](common_params & params) { - params.reuse_port = true; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_REUSE_PORT")); - add_opt(common_arg( - {"--path"}, "PATH", - string_format("path to serve static files from (default: %s)", params.public_path.c_str()), - [](common_params & params, const std::string & value) { - params.public_path = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_STATIC_PATH")); - add_opt(common_arg( - {"--api-prefix"}, "PREFIX", - string_format("prefix path the server serves from, without the trailing slash (default: %s)", params.api_prefix.c_str()), - [](common_params & params, const std::string & value) { - params.api_prefix = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_API_PREFIX")); - add_opt(common_arg( - {"--webui-config"}, "JSON", - "JSON that provides default WebUI settings (overrides WebUI defaults)", - [](common_params & params, const std::string & value) { - params.webui_config_json = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_WEBUI_CONFIG")); - add_opt(common_arg( - {"--webui-config-file"}, "PATH", - "JSON file that provides default WebUI settings (overrides WebUI defaults)", - [](common_params & params, const std::string & value) { - params.webui_config_json = read_file(value); - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_WEBUI_CONFIG_FILE")); - add_opt(common_arg( - {"--webui-mcp-proxy"}, - {"--no-webui-mcp-proxy"}, - string_format("experimental: whether to enable MCP CORS proxy - do not enable in untrusted environments (default: %s)", params.webui_mcp_proxy ? "enabled" : "disabled"), - [](common_params & params, bool value) { - params.webui_mcp_proxy = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_WEBUI_MCP_PROXY")); - add_opt(common_arg( - {"--tools"}, "TOOL1,TOOL2,...", - "experimental: whether to enable built-in tools for AI agents - do not enable in untrusted environments (default: no tools)\n" - "specify \"all\" to enable all tools\n" - "available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, apply_diff", - [](common_params & params, const std::string & value) { - params.server_tools = parse_csv_row(value); - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TOOLS")); - add_opt(common_arg( - {"--webui"}, - {"--no-webui"}, - string_format("whether to enable the Web UI (default: %s)", params.webui ? "enabled" : "disabled"), - [](common_params & params, bool value) { - params.webui = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_WEBUI")); - add_opt(common_arg( - {"--embedding", "--embeddings"}, - string_format("restrict to only support embedding use case; use only with dedicated embedding models (default: %s)", params.embedding ? "enabled" : "disabled"), - [](common_params & params) { - params.embedding = true; - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_DEBUG}).set_env("LLAMA_ARG_EMBEDDINGS")); - add_opt(common_arg( - {"--rerank", "--reranking"}, - string_format("enable reranking endpoint on server (default: %s)", "disabled"), - [](common_params & params) { - params.embedding = true; - params.pooling_type = LLAMA_POOLING_TYPE_RANK; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_RERANKING")); - add_opt(common_arg( - {"--api-key"}, "KEY", - "API key to use for authentication, multiple keys can be provided as a comma-separated list (default: none)", - [](common_params & params, const std::string & value) { - for (const auto & key : parse_csv_row(value)) { - if (!key.empty()) { - params.api_keys.push_back(key); - } - } - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_API_KEY")); - add_opt(common_arg( - {"--api-key-file"}, "FNAME", - "path to file containing API keys (default: none)", - [](common_params & params, const std::string & value) { - std::ifstream key_file(value); - if (!key_file) { - throw std::runtime_error(string_format("error: failed to open file '%s'\n", value.c_str())); - } - std::string key; - while (std::getline(key_file, key)) { - if (!key.empty()) { - params.api_keys.push_back(key); - } - } - key_file.close(); - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"--ssl-key-file"}, "FNAME", - "path to file a PEM-encoded SSL private key", - [](common_params & params, const std::string & value) { - params.ssl_file_key = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SSL_KEY_FILE")); - add_opt(common_arg( - {"--ssl-cert-file"}, "FNAME", - "path to file a PEM-encoded SSL certificate", - [](common_params & params, const std::string & value) { - params.ssl_file_cert = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SSL_CERT_FILE")); - add_opt(common_arg( - {"--chat-template-kwargs"}, "STRING", - "sets additional params for the json template parser, must be a valid json object string, e.g. '{\"key1\":\"value1\",\"key2\":\"value2\"}'", - [](common_params & params, const std::string & value) { - auto parsed = json::parse(value); - for (const auto & item : parsed.items()) { - if (item.key() == "enable_thinking") { - LOG_WRN("Setting 'enable_thinking' via --chat-template-kwargs is deprecated. " - "Use --reasoning on / --reasoning off instead.\n"); - } - params.default_template_kwargs[item.key()] = item.value().dump(); - } - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_CHAT_TEMPLATE_KWARGS")); - add_opt(common_arg( - {"-to", "--timeout"}, "N", - string_format("server read/write timeout in seconds (default: %d)", params.timeout_read), - [](common_params & params, int value) { - params.timeout_read = value; - params.timeout_write = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TIMEOUT")); - add_opt(common_arg( - {"--threads-http"}, "N", - string_format("number of threads used to process HTTP requests (default: %d)", params.n_threads_http), - [](common_params & params, int value) { - params.n_threads_http = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_THREADS_HTTP")); - add_opt(common_arg( - {"--cache-prompt"}, - {"--no-cache-prompt"}, - string_format("whether to enable prompt caching (default: %s)", params.cache_prompt ? "enabled" : "disabled"), - [](common_params & params, bool value) { - params.cache_prompt = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CACHE_PROMPT")); - add_opt(common_arg( - {"--cache-reuse"}, "N", - string_format( - "min chunk size to attempt reusing from the cache via KV shifting, requires prompt caching to be enabled (default: %d)\n" - "[(card)](https://ggml.ai/f0.png)", params.n_cache_reuse - ), - [](common_params & params, int value) { - params.n_cache_reuse = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CACHE_REUSE")); - add_opt(common_arg( - {"--metrics"}, - string_format("enable prometheus compatible metrics endpoint (default: %s)", params.endpoint_metrics ? "enabled" : "disabled"), - [](common_params & params) { - params.endpoint_metrics = true; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_ENDPOINT_METRICS")); - add_opt(common_arg( - {"--props"}, - string_format("enable changing global properties via POST /props (default: %s)", params.endpoint_props ? "enabled" : "disabled"), - [](common_params & params) { - params.endpoint_props = true; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_ENDPOINT_PROPS")); - add_opt(common_arg( - {"--slots"}, - {"--no-slots"}, - string_format("expose slots monitoring endpoint (default: %s)", params.endpoint_slots ? "enabled" : "disabled"), - [](common_params & params, bool value) { - params.endpoint_slots = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_ENDPOINT_SLOTS")); - add_opt(common_arg( - {"--slot-save-path"}, "PATH", - "path to save slot kv cache (default: disabled)", - [](common_params & params, const std::string & value) { - params.slot_save_path = value; - if (!fs_is_directory(params.slot_save_path)) { - throw std::invalid_argument("not a directory: " + value); - } - // if doesn't end with DIRECTORY_SEPARATOR, add it - if (!params.slot_save_path.empty() && params.slot_save_path[params.slot_save_path.size() - 1] != DIRECTORY_SEPARATOR) { - params.slot_save_path += DIRECTORY_SEPARATOR; - } - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"--media-path"}, "PATH", - "directory for loading local media files; files can be accessed via file:// URLs using relative paths (default: disabled)", - [](common_params & params, const std::string & value) { - params.media_path = value; - if (!fs_is_directory(params.media_path)) { - throw std::invalid_argument("not a directory: " + value); - } - // if doesn't end with DIRECTORY_SEPARATOR, add it - if (!params.media_path.empty() && params.media_path[params.media_path.size() - 1] != DIRECTORY_SEPARATOR) { - params.media_path += DIRECTORY_SEPARATOR; - } - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"--models-dir"}, "PATH", - "directory containing models for the router server (default: disabled)", - [](common_params & params, const std::string & value) { - params.models_dir = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODELS_DIR")); - add_opt(common_arg( - {"--models-preset"}, "PATH", - "path to INI file containing model presets for the router server (default: disabled)", - [](common_params & params, const std::string & value) { - params.models_preset = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODELS_PRESET")); - add_opt(common_arg( - {"--models-max"}, "N", - string_format("for router server, maximum number of models to load simultaneously (default: %d, 0 = unlimited)", params.models_max), - [](common_params & params, int value) { - params.models_max = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODELS_MAX")); - add_opt(common_arg( - {"--models-autoload"}, - {"--no-models-autoload"}, - string_format("for router server, whether to automatically load models (default: %s)", params.models_autoload ? "enabled" : "disabled"), - [](common_params & params, bool value) { - params.models_autoload = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODELS_AUTOLOAD")); - add_opt(common_arg( - {"--jinja"}, - {"--no-jinja"}, - string_format("whether to use jinja template engine for chat (default: %s)", params.use_jinja ? "enabled" : "disabled"), - [](common_params & params, bool value) { - params.use_jinja = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_MTMD}).set_env("LLAMA_ARG_JINJA")); - add_opt(common_arg( - {"--reasoning-format"}, "FORMAT", - "controls whether thought tags are allowed and/or extracted from the response, and in which format they're returned; one of:\n" - "- none: leaves thoughts unparsed in `message.content`\n" - "- deepseek: puts thoughts in `message.reasoning_content`\n" - "- deepseek-legacy: keeps `` tags in `message.content` while also populating `message.reasoning_content`\n" - "(default: auto)", - [](common_params & params, const std::string & value) { - params.reasoning_format = common_reasoning_format_from_name(value); - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK")); - add_opt(common_arg( - {"-rea", "--reasoning"}, "[on|off|auto]", - "Use reasoning/thinking in the chat ('on', 'off', or 'auto', default: 'auto' (detect from template))", - [](common_params & params, const std::string & value) { - if (is_truthy(value)) { - params.enable_reasoning = 1; - params.default_template_kwargs["enable_thinking"] = "true"; - } else if (is_falsey(value)) { - params.enable_reasoning = 0; - params.default_template_kwargs["enable_thinking"] = "false"; - } else if (is_autoy(value)) { - params.enable_reasoning = -1; - } else { - throw std::invalid_argument( - string_format("error: unknown value for --reasoning: '%s'\n", value.c_str())); - } - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_REASONING")); - add_opt(common_arg( - {"--reasoning-budget"}, "N", - "token budget for thinking: -1 for unrestricted, 0 for immediate end, N>0 for token budget (default: -1)", - [](common_params & params, int value) { - if (value < -1) { throw std::invalid_argument("invalid value"); } - params.sampling.reasoning_budget_tokens = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET")); - add_opt(common_arg( - {"--reasoning-budget-message"}, "MESSAGE", - "message injected before the end-of-thinking tag when reasoning budget is exhausted (default: none)", - [](common_params & params, const std::string & value) { - params.sampling.reasoning_budget_message = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_MESSAGE")); - add_opt(common_arg( - {"--chat-template"}, "JINJA_TEMPLATE", - string_format( - "set custom jinja chat template (default: template taken from model's metadata)\n" - "if suffix/prefix are specified, template will be disabled\n" - "only commonly used templates are accepted (unless --jinja is set before this flag):\n" - "list of built-in templates:\n%s", list_builtin_chat_templates().c_str() - ), - [](common_params & params, const std::string & value) { - params.chat_template = value; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_MTMD}).set_env("LLAMA_ARG_CHAT_TEMPLATE")); - add_opt(common_arg( - {"--chat-template-file"}, "JINJA_TEMPLATE_FILE", - string_format( - "set custom jinja chat template file (default: template taken from model's metadata)\n" - "if suffix/prefix are specified, template will be disabled\n" - "only commonly used templates are accepted (unless --jinja is set before this flag):\n" - "list of built-in templates:\n%s", list_builtin_chat_templates().c_str() - ), - [](common_params & params, const std::string & value) { - params.chat_template = read_file(value); - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CHAT_TEMPLATE_FILE")); - add_opt(common_arg( - {"--skip-chat-parsing"}, - {"--no-skip-chat-parsing"}, - string_format( - "force a pure content parser, even if a Jinja template is specified; model will output everything " - "in the content section, including any reasoning and/or tool calls (default: disabled)" - ), - [](common_params & params, bool value) { - params.force_pure_content_parser = value; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SKIP_CHAT_PARSING")); - add_opt(common_arg( - {"--prefill-assistant"}, - {"--no-prefill-assistant"}, - string_format( - "whether to prefill the assistant's response if the last message is an assistant message (default: prefill enabled)\n" - "when this flag is set, if the last message is an assistant message then it will be treated as a full message and not prefilled\n" - ), - [](common_params & params, bool value) { - params.prefill_assistant = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_PREFILL_ASSISTANT")); - add_opt(common_arg( - {"-sps", "--slot-prompt-similarity"}, "SIMILARITY", - string_format("how much the prompt of a request must match the prompt of a slot in order to use that slot (default: %.2f, 0.0 = disabled)\n", params.slot_prompt_similarity), - [](common_params & params, const std::string & value) { - params.slot_prompt_similarity = std::stof(value); - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"--lora-init-without-apply"}, - string_format("load LoRA adapters without applying them (apply later via POST /lora-adapters) (default: %s)", params.lora_init_without_apply ? "enabled" : "disabled"), - [](common_params & params) { - params.lora_init_without_apply = true; - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"--sleep-idle-seconds"}, "SECONDS", - string_format("number of seconds of idleness after which the server will sleep (default: %d; -1 = disabled)", params.sleep_idle_seconds), - [](common_params & params, int value) { - if (value == 0 || value < -1) { - throw std::invalid_argument("invalid value: cannot be 0 or less than -1"); - } - params.sleep_idle_seconds = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"--simple-io"}, - "use basic IO for better compatibility in subprocesses and limited consoles", - [](common_params & params) { - params.simple_io = true; - } - ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--positive-file"}, "FNAME", - string_format("positive prompts file, one prompt per line (default: '%s')", params.cvector_positive_file.c_str()), - [](common_params & params, const std::string & value) { - params.cvector_positive_file = value; - } - ).set_examples({LLAMA_EXAMPLE_CVECTOR_GENERATOR})); - add_opt(common_arg( - {"--negative-file"}, "FNAME", - string_format("negative prompts file, one prompt per line (default: '%s')", params.cvector_negative_file.c_str()), - [](common_params & params, const std::string & value) { - params.cvector_negative_file = value; - } - ).set_examples({LLAMA_EXAMPLE_CVECTOR_GENERATOR})); - add_opt(common_arg( - {"--pca-batch"}, "N", - string_format("batch size used for PCA. Larger batch runs faster, but uses more memory (default: %d)", params.n_pca_batch), - [](common_params & params, int value) { - params.n_pca_batch = value; - } - ).set_examples({LLAMA_EXAMPLE_CVECTOR_GENERATOR})); - add_opt(common_arg( - {"--pca-iter"}, "N", - string_format("number of iterations used for PCA (default: %d)", params.n_pca_iterations), - [](common_params & params, int value) { - params.n_pca_iterations = value; - } - ).set_examples({LLAMA_EXAMPLE_CVECTOR_GENERATOR})); - add_opt(common_arg( - {"--method"}, "{pca, mean}", - "dimensionality reduction method to be used (default: pca)", - [](common_params & params, const std::string & value) { - /**/ if (value == "pca") { params.cvector_dimre_method = DIMRE_METHOD_PCA; } - else if (value == "mean") { params.cvector_dimre_method = DIMRE_METHOD_MEAN; } - else { throw std::invalid_argument("invalid value"); } - } - ).set_examples({LLAMA_EXAMPLE_CVECTOR_GENERATOR})); - add_opt(common_arg( - {"--output-format"}, "{md,jsonl}", - "output format for batched-bench results (default: md)", - [](common_params & params, const std::string & value) { - /**/ if (value == "jsonl") { params.batched_bench_output_jsonl = true; } - else if (value == "md") { params.batched_bench_output_jsonl = false; } - else { throw std::invalid_argument("invalid value"); } - } - ).set_examples({LLAMA_EXAMPLE_BENCH})); - add_opt(common_arg( - {"--log-disable"}, - "Log disable", - [](common_params &) { - common_log_pause(common_log_main()); - } - )); - add_opt(common_arg( - {"--log-file"}, "FNAME", - "Log to file", - [](common_params &, const std::string & value) { - common_log_set_file(common_log_main(), value.c_str()); - } - ).set_env("LLAMA_LOG_FILE")); - add_opt(common_arg( - {"--log-colors"}, "[on|off|auto]", - "Set colored logging ('on', 'off', or 'auto', default: 'auto')\n" - "'auto' enables colors when output is to a terminal", - [](common_params &, const std::string & value) { - if (is_truthy(value)) { - common_log_set_colors(common_log_main(), LOG_COLORS_ENABLED); - } else if (is_falsey(value)) { - common_log_set_colors(common_log_main(), LOG_COLORS_DISABLED); - } else if (is_autoy(value)) { - common_log_set_colors(common_log_main(), LOG_COLORS_AUTO); - } else { - throw std::invalid_argument( - string_format("error: unknown value for --log-colors: '%s'\n", value.c_str())); - } - } - ).set_env("LLAMA_LOG_COLORS")); - add_opt(common_arg( - {"-v", "--verbose", "--log-verbose"}, - "Set verbosity level to infinity (i.e. log all messages, useful for debugging)", - [](common_params & params) { - params.verbosity = INT_MAX; - common_log_set_verbosity_thold(INT_MAX); - } - )); - add_opt(common_arg( - {"--offline"}, - "Offline mode: forces use of cache, prevents network access", - [](common_params & params) { - params.offline = true; - } - ).set_env("LLAMA_OFFLINE")); - add_opt(common_arg( - {"-lv", "--verbosity", "--log-verbosity"}, "N", - string_format("Set the verbosity threshold. Messages with a higher verbosity will be ignored. Values:\n" - " - 0: generic output\n" - " - 1: error\n" - " - 2: warning\n" - " - 3: info\n" - " - 4: debug\n" - "(default: %d)\n", params.verbosity), - [](common_params & params, int value) { - params.verbosity = value; - common_log_set_verbosity_thold(value); - } - ).set_env("LLAMA_LOG_VERBOSITY")); - add_opt(common_arg( - {"--log-prefix"}, - "Enable prefix in log messages", - [](common_params &) { - common_log_set_prefix(common_log_main(), true); - } - ).set_env("LLAMA_LOG_PREFIX")); - add_opt(common_arg( - {"--log-timestamps"}, - "Enable timestamps in log messages", - [](common_params &) { - common_log_set_timestamps(common_log_main(), true); - } - ).set_env("LLAMA_LOG_TIMESTAMPS")); - - // - // speculative parameters - // - - add_opt(common_arg( - {"--spec-draft-hf", "-hfd", "-hfrd", "--hf-repo-draft"}, "/[:quant]", - "Same as --hf-repo, but for the draft model (default: unused)", - [](common_params & params, const std::string & value) { - params.speculative.draft.mparams.hf_repo = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_HF_REPO")); - add_opt(common_arg( - {"--spec-draft-threads", "-td", "--threads-draft"}, "N", - "number of threads to use during generation (default: same as --threads)", - [](common_params & params, int value) { - params.speculative.draft.cpuparams.n_threads = value; - if (params.speculative.draft.cpuparams.n_threads <= 0) { - params.speculative.draft.cpuparams.n_threads = std::thread::hardware_concurrency(); - } - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-draft-threads-batch", "-tbd", "--threads-batch-draft"}, "N", - "number of threads to use during batch and prompt processing (default: same as --threads-draft)", - [](common_params & params, int value) { - params.speculative.draft.cpuparams_batch.n_threads = value; - if (params.speculative.draft.cpuparams_batch.n_threads <= 0) { - params.speculative.draft.cpuparams_batch.n_threads = std::thread::hardware_concurrency(); - } - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-draft-cpu-mask", "-Cd", "--cpu-mask-draft"}, "M", - "Draft model CPU affinity mask. Complements cpu-range-draft (default: same as --cpu-mask)", - [](common_params & params, const std::string & mask) { - params.speculative.draft.cpuparams.mask_valid = true; - if (!parse_cpu_mask(mask, params.speculative.draft.cpuparams.cpumask)) { - throw std::invalid_argument("invalid cpumask"); - } - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-draft-cpu-range", "-Crd", "--cpu-range-draft"}, "lo-hi", - "Ranges of CPUs for affinity. Complements --cpu-mask-draft", - [](common_params & params, const std::string & range) { - params.speculative.draft.cpuparams.mask_valid = true; - if (!parse_cpu_range(range, params.speculative.draft.cpuparams.cpumask)) { - throw std::invalid_argument("invalid range"); - } - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-draft-cpu-strict", "--cpu-strict-draft"}, "<0|1>", - "Use strict CPU placement for draft model (default: same as --cpu-strict)", - [](common_params & params, int value) { - params.speculative.draft.cpuparams.strict_cpu = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-draft-prio", "--prio-draft"}, "N", - string_format("set draft process/thread priority : 0-normal, 1-medium, 2-high, 3-realtime (default: %d)\n", params.speculative.draft.cpuparams.priority), - [](common_params & params, int prio) { - if (prio < 0 || prio > 3) { - throw std::invalid_argument("invalid value"); - } - params.speculative.draft.cpuparams.priority = (enum ggml_sched_priority) prio; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-draft-poll", "--poll-draft"}, "<0|1>", - "Use polling to wait for draft model work (default: same as --poll])", - [](common_params & params, int value) { - params.speculative.draft.cpuparams.poll = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-draft-cpu-mask-batch", "-Cbd", "--cpu-mask-batch-draft"}, "M", - "Draft model CPU affinity mask. Complements cpu-range-draft (default: same as --cpu-mask)", - [](common_params & params, const std::string & mask) { - params.speculative.draft.cpuparams_batch.mask_valid = true; - if (!parse_cpu_mask(mask, params.speculative.draft.cpuparams_batch.cpumask)) { - throw std::invalid_argument("invalid cpumask"); - } - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-draft-cpu-range-batch", "-Crbd", "--cpu-range-batch-draft"}, "lo-hi", - "Ranges of CPUs for affinity. Complements --cpu-mask-draft-batch)", - [](common_params & params, const std::string & range) { - params.speculative.draft.cpuparams_batch.mask_valid = true; - if (!parse_cpu_range(range, params.speculative.draft.cpuparams_batch.cpumask)) { - throw std::invalid_argument("invalid cpumask"); - } - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE})); - add_opt(common_arg( - {"--spec-draft-cpu-strict-batch", "--cpu-strict-batch-draft"}, "<0|1>", - "Use strict CPU placement for draft model (default: --cpu-strict-draft)", - [](common_params & params, int value) { - params.speculative.draft.cpuparams_batch.strict_cpu = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-draft-prio-batch", "--prio-batch-draft"}, "N", - string_format("set draft process/thread priority : 0-normal, 1-medium, 2-high, 3-realtime (default: %d)\n", params.speculative.draft.cpuparams_batch.priority), - [](common_params & params, int prio) { - if (prio < 0 || prio > 3) { - throw std::invalid_argument("invalid value"); - } - params.speculative.draft.cpuparams_batch.priority = (enum ggml_sched_priority) prio; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-draft-poll-batch", "--poll-batch-draft"}, "<0|1>", - "Use polling to wait for draft model work (default: --poll-draft)", - [](common_params & params, int value) { - params.speculative.draft.cpuparams_batch.poll = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-draft-type-k", "-ctkd", "--cache-type-k-draft"}, "TYPE", - string_format( - "KV cache data type for K for the draft model\n" - "allowed values: %s\n" - "(default: %s)", - get_all_kv_cache_types().c_str(), - ggml_type_name(params.speculative.draft.cache_type_k) - ), - [](common_params & params, const std::string & value) { - params.speculative.draft.cache_type_k = kv_cache_type_from_str(value); - } - ).set_env("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K")); - add_opt(common_arg( - {"--spec-draft-type-v", "-ctvd", "--cache-type-v-draft"}, "TYPE", - string_format( - "KV cache data type for V for the draft model\n" - "allowed values: %s\n" - "(default: %s)", - get_all_kv_cache_types().c_str(), - ggml_type_name(params.speculative.draft.cache_type_v) - ), - [](common_params & params, const std::string & value) { - params.speculative.draft.cache_type_v = kv_cache_type_from_str(value); - } - ).set_env("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V")); - add_opt(common_arg( - {"--spec-draft-override-tensor", "-otd", "--override-tensor-draft"}, "=,...", - "override tensor buffer type for draft model", [](common_params & params, const std::string & value) { - parse_tensor_buffer_overrides(value, params.speculative.draft.tensor_buft_overrides); - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-draft-cpu-moe", "-cmoed", "--cpu-moe-draft"}, - "keep all Mixture of Experts (MoE) weights in the CPU for the draft model", - [](common_params & params) { - params.speculative.draft.tensor_buft_overrides.push_back(llm_ffn_exps_cpu_override()); - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_CPU_MOE")); - add_opt(common_arg( - {"--spec-draft-n-cpu-moe", "--spec-draft-ncmoe", "-ncmoed", "--n-cpu-moe-draft"}, "N", - "keep the Mixture of Experts (MoE) weights of the first N layers in the CPU for the draft model", - [](common_params & params, int value) { - if (value < 0) { - throw std::invalid_argument("invalid value"); - } - for (int i = 0; i < value; ++i) { - static std::list buft_overrides_draft; - buft_overrides_draft.push_back(llm_ffn_exps_block_regex(i)); - params.speculative.draft.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type()}); - } - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE")); - - add_opt(common_arg( - {"--spec-draft-n-max"}, "N", - string_format("number of tokens to draft for speculative decoding (default: %d)", params.speculative.draft.n_max), - [](common_params & params, int value) { - params.speculative.draft.n_max = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MAX")); - add_opt(common_arg( - {"--spec-draft-n-min"}, "N", - string_format("minimum number of draft tokens to use for speculative decoding (default: %d)", params.speculative.draft.n_min), - [](common_params & params, int value) { - params.speculative.draft.n_min = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MIN")); - - add_opt(common_arg( - {"--spec-draft-p-split", "--draft-p-split"}, "P", - string_format("speculative decoding split probability (default: %.2f)", (double)params.speculative.draft.p_split), - [](common_params & params, const std::string & value) { - params.speculative.draft.p_split = std::stof(value); - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_P_SPLIT")); - add_opt(common_arg( - {"--spec-draft-p-min", "--draft-p-min"}, "P", - string_format("minimum speculative decoding probability (greedy) (default: %.2f)", (double)params.speculative.draft.p_min), - [](common_params & params, const std::string & value) { - params.speculative.draft.p_min = std::stof(value); - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_P_MIN")); - add_opt(common_arg( - {"--spec-draft-ctx-size", "-cd", "--ctx-size-draft"}, "N", - string_format("size of the prompt context for the draft model (default: %d, 0 = loaded from model)", params.speculative.draft.n_ctx), - [](common_params & params, int value) { - params.speculative.draft.n_ctx = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_CTX_SIZE")); - add_opt(common_arg( - {"--spec-draft-device", "-devd", "--device-draft"}, "", - "comma-separated list of devices to use for offloading the draft model (none = don't offload)\n" - "use --list-devices to see a list of available devices", - [](common_params & params, const std::string & value) { - params.speculative.draft.devices = parse_device_list(value); - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - GGML_ASSERT(params.speculative.draft.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0 - add_opt(common_arg( - {"--spec-draft-ngl", "-ngld", "--gpu-layers-draft", "--n-gpu-layers-draft"}, "N", - string_format("max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: %s)", - params.speculative.draft.n_gpu_layers == -1 ? "auto" : "all"), - [](common_params & params, const std::string & value) { - if (value == "auto") { - params.speculative.draft.n_gpu_layers = -1; - } else if (value == "all") { - params.speculative.draft.n_gpu_layers = -2; - } else { - params.speculative.draft.n_gpu_layers = std::stoi(value); - } - if (!llama_supports_gpu_offload()) { - fprintf(stderr, "warning: no usable GPU found, --gpu-layers-draft option will be ignored\n"); - fprintf(stderr, "warning: one possible reason is that llama.cpp was compiled without GPU support\n"); - fprintf(stderr, "warning: consult docs/build.md for compilation instructions\n"); - } - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_N_GPU_LAYERS_DRAFT")); - add_opt(common_arg( - {"--spec-draft-model", "-md", "--model-draft"}, "FNAME", - "draft model for speculative decoding (default: unused)", - [](common_params & params, const std::string & value) { - params.speculative.draft.mparams.path = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_MODEL")); - add_opt(common_arg( - {"--spec-draft-replace", "--spec-replace"}, "TARGET", "DRAFT", - "translate the string in TARGET into DRAFT if the draft model and main model are not compatible", - [](common_params & params, const std::string & tgt, const std::string & dft) { - params.speculative.draft.replacements.push_back({ tgt, dft }); - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-type"}, "[none|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod]", - string_format("type of speculative decoding to use when no draft model is provided (default: %s)\n", - common_speculative_type_to_str(params.speculative.type).c_str()), - [](common_params & params, const std::string & value) { - if (value == "none") { - params.speculative.type = COMMON_SPECULATIVE_TYPE_NONE; - } else if (value == "ngram-cache") { - params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_CACHE; - } else if (value == "ngram-simple") { - params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE; - } else if (value == "ngram-map-k") { - params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K; - } else if (value == "ngram-map-k4v") { - params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V; - } else if (value == "ngram-mod") { - params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_MOD; - } else { - throw std::invalid_argument("unknown speculative decoding type without draft model"); - } - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_TYPE")); - add_opt(common_arg( - {"--spec-ngram-mod-n-min"}, "N", - string_format("minimum number of ngram tokens to use for ngram-based speculative decoding (default: %d)", params.speculative.ngram_mod.n_min), - [](common_params & params, int value) { - if (value < 0 || value > 1024) { - throw std::invalid_argument("ngram n-min must be between 0 and 1024 inclusive"); - } - params.speculative.ngram_mod.n_min = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-ngram-mod-n-max"}, "N", - string_format("maximum number of ngram tokens to use for ngram-based speculative decoding (default: %d)", params.speculative.ngram_mod.n_max), - [](common_params & params, int value) { - if (value < 0 || value > 1024) { - throw std::invalid_argument("ngram n-max must be between 0 and 1024 inclusive"); - } - params.speculative.ngram_mod.n_max = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-ngram-mod-n-match"}, "N", - string_format("ngram-mod lookup length (default: %d)", params.speculative.ngram_mod.n_match), - [](common_params & params, int value) { - if (value < 1 || value > 1024) { - throw std::invalid_argument("ngram size N must be between 1 and 1024 inclusive"); - } - params.speculative.ngram_mod.n_match = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - - add_opt(common_arg( - {"--spec-ngram-simple-size-n"}, "N", - string_format("ngram size N for ngram-simple speculative decoding, length of lookup n-gram (default: %d)", params.speculative.ngram_simple.size_n), - [](common_params & params, int value) { - if (value < 1 || value > 1024) { - throw std::invalid_argument("ngram size N must be between 1 and 1024 inclusive"); - } - params.speculative.ngram_simple.size_n = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-ngram-simple-size-m"}, "N", - string_format("ngram size M for ngram-simple speculative decoding, length of draft m-gram (default: %d)", params.speculative.ngram_simple.size_m), - [](common_params & params, int value) { - if (value < 1 || value > 1024) { - throw std::invalid_argument("ngram size M must be between 1 and 1024 inclusive"); - } - params.speculative.ngram_simple.size_m = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-ngram-simple-min-hits"}, "N", - string_format("minimum hits for ngram-simple speculative decoding (default: %d)", params.speculative.ngram_simple.min_hits), - [](common_params & params, int value) { - if (value < 1) { - throw std::invalid_argument("ngram min hits must be at least 1"); - } - params.speculative.ngram_simple.min_hits = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - - add_opt(common_arg( - {"--spec-ngram-map-k-size-n"}, "N", - string_format("ngram size N for ngram-map-k speculative decoding, length of lookup n-gram (default: %d)", params.speculative.ngram_map_k.size_n), - [](common_params & params, int value) { - if (value < 1 || value > 1024) { - throw std::invalid_argument("ngram size N must be between 1 and 1024 inclusive"); - } - params.speculative.ngram_map_k.size_n = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-ngram-map-k-size-m"}, "N", - string_format("ngram size M for ngram-map-k speculative decoding, length of draft m-gram (default: %d)", params.speculative.ngram_map_k.size_m), - [](common_params & params, int value) { - if (value < 1 || value > 1024) { - throw std::invalid_argument("ngram size M must be between 1 and 1024 inclusive"); - } - params.speculative.ngram_map_k.size_m = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-ngram-map-k-min-hits"}, "N", - string_format("minimum hits for ngram-map-k speculative decoding (default: %d)", params.speculative.ngram_map_k.min_hits), - [](common_params & params, int value) { - if (value < 1) { - throw std::invalid_argument("ngram min hits must be at least 1"); - } - params.speculative.ngram_map_k.min_hits = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - - add_opt(common_arg( - {"--spec-ngram-map-k4v-size-n"}, "N", - string_format("ngram size N for ngram-map-k4v speculative decoding, length of lookup n-gram (default: %d)", params.speculative.ngram_map_k4v.size_n), - [](common_params & params, int value) { - if (value < 1 || value > 1024) { - throw std::invalid_argument("ngram size N must be between 1 and 1024 inclusive"); - } - params.speculative.ngram_map_k4v.size_n = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-ngram-map-k4v-size-m"}, "N", - string_format("ngram size M for ngram-map-k4v speculative decoding, length of draft m-gram (default: %d)", params.speculative.ngram_map_k4v.size_m), - [](common_params & params, int value) { - if (value < 1 || value > 1024) { - throw std::invalid_argument("ngram size M must be between 1 and 1024 inclusive"); - } - params.speculative.ngram_map_k4v.size_m = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - add_opt(common_arg( - {"--spec-ngram-map-k4v-min-hits"}, "N", - string_format("minimum hits for ngram-map-k4v speculative decoding (default: %d)", params.speculative.ngram_map_k4v.min_hits), - [](common_params & params, int value) { - if (value < 1) { - throw std::invalid_argument("ngram min hits must be at least 1"); - } - params.speculative.ngram_map_k4v.min_hits = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - - // - // removed params - // - - add_opt(common_arg( - {"--draft", "--draft-n", "--draft-max"}, "N", - "the argument has been removed. use --spec-draft-n-max or --spec-ngram-mod-n-max", - [](common_params & /*params*/, int /*value*/) { - throw std::invalid_argument("the argument has been removed. use --spec-draft-n-max or --spec-ngram-mod-n-max"); - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_DRAFT_MAX")); - add_opt(common_arg( - {"--draft-min", "--draft-n-min"}, "N", - "the argument has been removed. use --spec-draft-n-min or --spec-ngram-mod-n-min", - [](common_params & /*params*/, int /*value*/) { - throw std::invalid_argument("the argument has been removed. use --spec-draft-n-min or --spec-ngram-mod-n-min"); - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_DRAFT_MIN")); - add_opt(common_arg( - {"--spec-ngram-size-n"}, "N", - "the argument has been removed. use the respective --spec-ngram-*-size-n or --spec-ngram-mod-n-match", - [](common_params & /*params*/, int /*value*/) { - throw std::invalid_argument("the argument has been removed. use the respective --spec-ngram-*-size-n"); - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"--spec-ngram-size-m"}, "N", - "the argument has been removed. use the respective --spec-ngram-*-size-m", - [](common_params & /*params*/, int /*value*/) { - throw std::invalid_argument("the argument has been removed. use the respective --spec-ngram-*-size-m"); - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"--spec-ngram-min-hits"}, "N", - "the argument has been removed. use the respective --spec-ngram-*-min-hits", - [](common_params & /*params*/, int /*value*/) { - throw std::invalid_argument("the argument has been removed. use the respective --spec-ngram-*-min-hits"); - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER})); - - // - // TTS params - // - - add_opt(common_arg( - {"-mv", "--model-vocoder"}, "FNAME", - "vocoder model for audio generation (default: unused)", - [](common_params & params, const std::string & value) { - params.vocoder.model.path = value; - } - ).set_examples({LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"--tts-use-guide-tokens"}, - "Use guide tokens to improve TTS word recall", - [](common_params & params) { - params.vocoder.use_guide_tokens = true; - } - ).set_examples({LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_SERVER})); - add_opt(common_arg( - {"--tts-speaker-file"}, "FNAME", - "speaker file path for audio generation", - [](common_params & params, const std::string & value) { - params.vocoder.speaker_file = value; - } - ).set_examples({LLAMA_EXAMPLE_TTS})); - - // - // diffusion params - // - - add_opt(common_arg( - {"--diffusion-steps"}, "N", - string_format("number of diffusion steps (default: %d)", params.diffusion.steps), - [](common_params & params, int value) { params.diffusion.steps = value; } - ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); - add_opt(common_arg( - {"--diffusion-visual"}, - string_format("enable visual diffusion mode (show progressive generation) (default: %s)", params.diffusion.visual_mode ? "true" : "false"), - [](common_params & params) { params.diffusion.visual_mode = true; } - ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); - add_opt(common_arg( - {"--diffusion-eps"}, "F", - string_format("epsilon for timesteps (default: %.6f)", (double) params.diffusion.eps), - [](common_params & params, const std::string & value) { params.diffusion.eps = std::stof(value); } - ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); - add_opt(common_arg( - {"--diffusion-algorithm"}, "N", - string_format("diffusion algorithm: 0=ORIGIN, 1=ENTROPY_BASED, 2=MARGIN_BASED, 3=RANDOM, 4=LOW_CONFIDENCE (default: %d)", params.diffusion.algorithm), - [](common_params & params, int value) { params.diffusion.algorithm = value; } - ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); - add_opt(common_arg( - {"--diffusion-alg-temp"}, "F", - string_format("dream algorithm temperature (default: %.3f)", (double) params.diffusion.alg_temp), - [](common_params & params, const std::string & value) { params.diffusion.alg_temp = std::stof(value); } - ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); - add_opt(common_arg( - {"--diffusion-block-length"}, "N", - string_format("llada block length for generation (default: %d)", params.diffusion.block_length), - [](common_params & params, int value) { params.diffusion.block_length = value; } - ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); - add_opt(common_arg( - {"--diffusion-cfg-scale"}, "F", - string_format("llada classifier-free guidance scale (default: %.3f)", (double) params.diffusion.cfg_scale), - [](common_params & params, const std::string & value) { params.diffusion.cfg_scale = std::stof(value); } - ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); - add_opt(common_arg( - {"--diffusion-add-gumbel-noise"}, "F", - string_format("add gumbel noise to the logits if temp > 0.0 (default: %s)", params.diffusion.add_gumbel_noise ? "true" : "false"), - [](common_params & params, const std::string & value) { params.diffusion.add_gumbel_noise = std::stof(value); } - ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); - add_opt(common_arg( - { "-lr", "--learning-rate" }, "ALPHA", - string_format("adamw or sgd optimizer alpha (default: %.2g); note: sgd alpha recommended ~10x (no momentum)", (double) params.lr.lr0), - [](common_params & params, const std::string & value) { params.lr.lr0 = std::stof(value); } - ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); - add_opt(common_arg({ "-lr-min", "--learning-rate-min" }, "ALPHA", - string_format("(if >0) final learning rate after decay (if -decay-epochs is set, default=%.2g)", - (double) params.lr.lr_min), - [](common_params & params, const std::string & value) { params.lr.lr_min = std::stof(value); } - ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); - add_opt(common_arg( - {"-decay-epochs", "--learning-rate-decay-epochs"}, "ALPHA", - string_format("(if >0) decay learning rate to -lr-min after this many epochs (exponential decay, default=%.2g)", (double) params.lr.decay_epochs), - [](common_params & params, const std::string & value) { params.lr.decay_epochs = std::stof(value); } - ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); - add_opt(common_arg( - {"-wd", "--weight-decay"}, "WD", - string_format("adamw or sgd optimizer weight decay (0 is off; recommend very small e.g. 1e-9) (default: %.2g).", (double) params.lr.wd), - [](common_params & params, const std::string & value) { params.lr.wd = std::stof(value); } - ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); - add_opt(common_arg( - {"-val-split", "--val-split"}, "FRACTION", - string_format("fraction of data to use as validation set for training (default: %.2g).", (double) params.val_split), - [](common_params & params, const std::string & value) { params.val_split = std::stof(value); } - ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); - add_opt(common_arg( - {"-epochs", "--epochs"}, "N", - string_format("optimizer max # of epochs (default: %d)", params.lr.epochs), - [](common_params & params, int epochs) { params.lr.epochs = epochs; } - ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); - add_opt(common_arg( - {"-opt", "--optimizer"}, "sgd|adamw", "adamw or sgd", - [](common_params & params, const std::string & name) { - params.optimizer = common_opt_get_optimizer(name.c_str()); - if (params.optimizer == GGML_OPT_OPTIMIZER_TYPE_COUNT) { - throw std::invalid_argument("invalid --optimizer, valid options: adamw, sgd"); - } - } - ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); - add_opt(common_arg( - {"--check"}, - string_format("check rather than generate results (default: %s)", params.check ? "true" : "false"), - [](common_params & params) { - params.check = true; - } - ).set_examples({LLAMA_EXAMPLE_RESULTS})); - add_opt(common_arg( - {"--save-logits"}, - string_format("save final logits to files for verification (default: %s)", params.save_logits ? "true" : "false"), - [](common_params & params) { - params.save_logits = true; - } - ).set_examples({LLAMA_EXAMPLE_DEBUG})); - add_opt(common_arg( - {"--logits-output-dir"}, "PATH", - string_format("directory for saving logits output files (default: %s)", params.logits_output_dir.c_str()), - [](common_params & params, const std::string & value) { - params.logits_output_dir = value; - } - ).set_examples({LLAMA_EXAMPLE_DEBUG})); - add_opt(common_arg( - {"--tensor-filter"}, "REGEX", - "filter tensor names for debug output (regex pattern, can be specified multiple times)", - [](common_params & params, const std::string & value) { - params.tensor_filter.push_back(value); - } - ).set_examples({LLAMA_EXAMPLE_DEBUG})); - - // presets - add_opt(common_arg( - {"--tts-oute-default"}, - string_format("use default OuteTTS models (note: can download weights from the internet)"), - [](common_params & params) { - params.model.hf_repo = "OuteAI/OuteTTS-0.2-500M-GGUF"; - params.model.hf_file = "OuteTTS-0.2-500M-Q8_0.gguf"; - params.vocoder.model.hf_repo = "ggml-org/WavTokenizer"; - params.vocoder.model.hf_file = "WavTokenizer-Large-75-F16.gguf"; - } - ).set_examples({LLAMA_EXAMPLE_TTS})); - - add_opt(common_arg( - {"--embd-gemma-default"}, - string_format("use default EmbeddingGemma model (note: can download weights from the internet)"), - [](common_params & params) { - params.model.hf_repo = "ggml-org/embeddinggemma-300M-qat-q4_0-GGUF"; - params.model.hf_file = "embeddinggemma-300M-qat-Q4_0.gguf"; - params.port = 8011; - params.n_ubatch = 2048; - params.n_batch = 2048; - params.n_parallel = 32; - params.n_ctx = 2048*params.n_parallel; - params.verbose_prompt = true; - params.embedding = true; - } - ).set_examples({LLAMA_EXAMPLE_EMBEDDING, LLAMA_EXAMPLE_SERVER})); - - add_opt(common_arg( - {"--fim-qwen-1.5b-default"}, - string_format("use default Qwen 2.5 Coder 1.5B (note: can download weights from the internet)"), - [](common_params & params) { - params.model.hf_repo = "ggml-org/Qwen2.5-Coder-1.5B-Q8_0-GGUF"; - params.model.hf_file = "qwen2.5-coder-1.5b-q8_0.gguf"; - params.port = 8012; - params.n_ubatch = 1024; - params.n_batch = 1024; - params.n_ctx = 0; - params.n_cache_reuse = 256; - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - - add_opt(common_arg( - {"--fim-qwen-3b-default"}, - string_format("use default Qwen 2.5 Coder 3B (note: can download weights from the internet)"), - [](common_params & params) { - params.model.hf_repo = "ggml-org/Qwen2.5-Coder-3B-Q8_0-GGUF"; - params.model.hf_file = "qwen2.5-coder-3b-q8_0.gguf"; - params.port = 8012; - params.n_ubatch = 1024; - params.n_batch = 1024; - params.n_ctx = 0; - params.n_cache_reuse = 256; - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - - add_opt(common_arg( - {"--fim-qwen-7b-default"}, - string_format("use default Qwen 2.5 Coder 7B (note: can download weights from the internet)"), - [](common_params & params) { - params.model.hf_repo = "ggml-org/Qwen2.5-Coder-7B-Q8_0-GGUF"; - params.model.hf_file = "qwen2.5-coder-7b-q8_0.gguf"; - params.port = 8012; - params.n_ubatch = 1024; - params.n_batch = 1024; - params.n_ctx = 0; - params.n_cache_reuse = 256; - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - - add_opt(common_arg( - {"--fim-qwen-7b-spec"}, - string_format("use Qwen 2.5 Coder 7B + 0.5B draft for speculative decoding (note: can download weights from the internet)"), - [](common_params & params) { - params.model.hf_repo = "ggml-org/Qwen2.5-Coder-7B-Q8_0-GGUF"; - params.model.hf_file = "qwen2.5-coder-7b-q8_0.gguf"; - params.speculative.draft.mparams.hf_repo = "ggml-org/Qwen2.5-Coder-0.5B-Q8_0-GGUF"; - params.speculative.draft.mparams.hf_file = "qwen2.5-coder-0.5b-q8_0.gguf"; - params.port = 8012; - params.n_ubatch = 1024; - params.n_batch = 1024; - params.n_ctx = 0; - params.n_cache_reuse = 256; - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - - add_opt(common_arg( - {"--fim-qwen-14b-spec"}, - string_format("use Qwen 2.5 Coder 14B + 0.5B draft for speculative decoding (note: can download weights from the internet)"), - [](common_params & params) { - params.model.hf_repo = "ggml-org/Qwen2.5-Coder-14B-Q8_0-GGUF"; - params.model.hf_file = "qwen2.5-coder-14b-q8_0.gguf"; - params.speculative.draft.mparams.hf_repo = "ggml-org/Qwen2.5-Coder-0.5B-Q8_0-GGUF"; - params.speculative.draft.mparams.hf_file = "qwen2.5-coder-0.5b-q8_0.gguf"; - params.port = 8012; - params.n_ubatch = 1024; - params.n_batch = 1024; - params.n_ctx = 0; - params.n_cache_reuse = 256; - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - - add_opt(common_arg( - {"--fim-qwen-30b-default"}, - string_format("use default Qwen 3 Coder 30B A3B Instruct (note: can download weights from the internet)"), - [](common_params & params) { - params.model.hf_repo = "ggml-org/Qwen3-Coder-30B-A3B-Instruct-Q8_0-GGUF"; - params.model.hf_file = "qwen3-coder-30b-a3b-instruct-q8_0.gguf"; - params.port = 8012; - params.n_ubatch = 1024; - params.n_batch = 1024; - params.n_ctx = 0; - params.n_cache_reuse = 256; - } - ).set_examples({LLAMA_EXAMPLE_SERVER})); - - add_opt(common_arg( - {"--gpt-oss-20b-default"}, - string_format("use gpt-oss-20b (note: can download weights from the internet)"), - [](common_params & params) { - params.model.hf_repo = "ggml-org/gpt-oss-20b-GGUF"; - params.model.hf_file = "gpt-oss-20b-mxfp4.gguf"; - params.port = 8013; - params.n_ubatch = 2048; - params.n_batch = 32768; - params.n_parallel = 2; - params.n_ctx = 131072*params.n_parallel; - params.sampling.temp = 1.0f; - params.sampling.top_p = 1.0f; - params.sampling.top_k = 0; - params.sampling.min_p = 0.01f; - params.use_jinja = true; - //params.default_template_kwargs["reasoning_effort"] = "\"high\""; - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - - add_opt(common_arg( - {"--gpt-oss-120b-default"}, - string_format("use gpt-oss-120b (note: can download weights from the internet)"), - [](common_params & params) { - params.model.hf_repo = "ggml-org/gpt-oss-120b-GGUF"; - params.port = 8013; - params.n_ubatch = 2048; - params.n_batch = 32768; - params.n_parallel = 2; - params.n_ctx = 131072*params.n_parallel; - params.sampling.temp = 1.0f; - params.sampling.top_p = 1.0f; - params.sampling.top_k = 0; - params.sampling.min_p = 0.01f; - params.use_jinja = true; - //params.default_template_kwargs["reasoning_effort"] = "\"high\""; - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - - add_opt(common_arg( - {"--vision-gemma-4b-default"}, - string_format("use Gemma 3 4B QAT (note: can download weights from the internet)"), - [](common_params & params) { - params.model.hf_repo = "ggml-org/gemma-3-4b-it-qat-GGUF"; - params.port = 8014; - params.n_ctx = 0; - params.use_jinja = true; - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - - add_opt(common_arg( - {"--vision-gemma-12b-default"}, - string_format("use Gemma 3 12B QAT (note: can download weights from the internet)"), - [](common_params & params) { - params.model.hf_repo = "ggml-org/gemma-3-12b-it-qat-GGUF"; - params.port = 8014; - params.n_ctx = 0; - params.use_jinja = true; - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - - add_opt(common_arg( - {"--spec-default"}, - string_format("enable default speculative decoding config"), - [](common_params & params) { - params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_MOD; - params.speculative.ngram_mod.n_match = 24; - params.speculative.ngram_mod.n_min = 48; - params.speculative.ngram_mod.n_max = 64; - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); - - return ctx_arg; -} - -void common_params_add_preset_options(std::vector & args) { - // arguments below won't be treated as CLI args, only preset options - args.push_back(common_arg( - {"load-on-startup"}, "NAME", - "in server router mode, autoload this model on startup", - [](common_params &, const std::string &) { /* unused */ } - ).set_env(COMMON_ARG_PRESET_LOAD_ON_STARTUP).set_preset_only()); - - args.push_back(common_arg( - {"stop-timeout"}, "SECONDS", - "in server router mode, force-kill model instance after this many seconds of graceful shutdown", - [](common_params &, int) { /* unused */ } - ).set_env(COMMON_ARG_PRESET_STOP_TIMEOUT).set_preset_only()); - - // args.push_back(common_arg( - // {"pin"}, - // "in server router mode, do not unload this model if models_max is exceeded", - // [](common_params &) { /* unused */ } - // ).set_preset_only()); -} diff --git a/common/arg.h b/common/arg.h deleted file mode 100644 index 2c2a4e38a2bb..000000000000 --- a/common/arg.h +++ /dev/null @@ -1,133 +0,0 @@ -#pragma once - -#include "common.h" - -#include -#include -#include -#include -#include - -// pseudo-env variable to identify preset-only arguments -#define COMMON_ARG_PRESET_LOAD_ON_STARTUP "__PRESET_LOAD_ON_STARTUP" -#define COMMON_ARG_PRESET_STOP_TIMEOUT "__PRESET_STOP_TIMEOUT" - -// -// CLI argument parsing -// - -struct common_arg { - std::set examples = {LLAMA_EXAMPLE_COMMON}; - std::set excludes = {}; - std::vector args; - std::vector args_neg; // for negated args like --no-xxx - const char * value_hint = nullptr; // help text or example for arg value - const char * value_hint_2 = nullptr; // for second arg value - const char * env = nullptr; - std::string help; - bool is_sampling = false; // is current arg a sampling param? - bool is_spec = false; // is current arg a speculative decoding param? - bool is_preset_only = false; // is current arg preset-only (not treated as CLI arg) - void (*handler_void) (common_params & params) = nullptr; - void (*handler_string) (common_params & params, const std::string &) = nullptr; - void (*handler_str_str)(common_params & params, const std::string &, const std::string &) = nullptr; - void (*handler_int) (common_params & params, int) = nullptr; - void (*handler_bool) (common_params & params, bool) = nullptr; - - common_arg() = default; - - common_arg( - const std::initializer_list & args, - const char * value_hint, - const std::string & help, - void (*handler)(common_params & params, const std::string &) - ) : args(args), value_hint(value_hint), help(help), handler_string(handler) {} - - common_arg( - const std::initializer_list & args, - const char * value_hint, - const std::string & help, - void (*handler)(common_params & params, int) - ) : args(args), value_hint(value_hint), help(help), handler_int(handler) {} - - common_arg( - const std::initializer_list & args, - const std::string & help, - void (*handler)(common_params & params) - ) : args(args), help(help), handler_void(handler) {} - - common_arg( - const std::initializer_list & args, - const std::initializer_list & args_neg, - const std::string & help, - void (*handler)(common_params & params, bool) - ) : args(args), args_neg(args_neg), help(help), handler_bool(handler) {} - - // support 2 values for arg - common_arg( - const std::initializer_list & args, - const char * value_hint, - const char * value_hint_2, - const std::string & help, - void (*handler)(common_params & params, const std::string &, const std::string &) - ) : args(args), value_hint(value_hint), value_hint_2(value_hint_2), help(help), handler_str_str(handler) {} - - common_arg & set_examples(std::initializer_list examples); - common_arg & set_excludes(std::initializer_list excludes); - common_arg & set_env(const char * env); - common_arg & set_sampling(); - common_arg & set_spec(); - common_arg & set_preset_only(); - bool in_example(enum llama_example ex); - bool is_exclude(enum llama_example ex); - bool get_value_from_env(std::string & output) const; - bool has_value_from_env() const; - std::string to_string() const; - - // for using as key in std::map - bool operator<(const common_arg& other) const { - if (args.empty() || other.args.empty()) { - return false; - } - return strcmp(args[0], other.args[0]) < 0; - } - bool operator==(const common_arg& other) const { - if (args.empty() || other.args.empty()) { - return false; - } - return strcmp(args[0], other.args[0]) == 0; - } - - // get all args and env vars (including negated args/env) - std::vector get_args() const; - std::vector get_env() const; -}; - -namespace common_arg_utils { - bool is_truthy(const std::string & value); - bool is_falsey(const std::string & value); - bool is_autoy(const std::string & value); -} - -struct common_params_context { - enum llama_example ex = LLAMA_EXAMPLE_COMMON; - common_params & params; - std::vector options; - void(*print_usage)(int, char **) = nullptr; - common_params_context(common_params & params) : params(params) {} -}; - -// parse input arguments from CLI -// if one argument has invalid value, it will automatically display usage of the specific argument (and not the full usage message) -bool common_params_parse(int argc, char ** argv, common_params & params, llama_example ex, void(*print_usage)(int, char **) = nullptr); - -// parse input arguments from CLI into a map -bool common_params_to_map(int argc, char ** argv, llama_example ex, std::map & out_map); - -// populate preset-only arguments -// these arguments are not treated as command line arguments -// see: https://github.com/ggml-org/llama.cpp/issues/18163 -void common_params_add_preset_options(std::vector & args); - -// initialize argument parser context - used by test-arg-parser and preset -common_params_context common_params_parser_init(common_params & params, llama_example ex, void(*print_usage)(int, char **) = nullptr); diff --git a/common/base64.hpp b/common/base64.hpp deleted file mode 100644 index 563247a6e5f7..000000000000 --- a/common/base64.hpp +++ /dev/null @@ -1,392 +0,0 @@ -/* -This is free and unencumbered software released into the public domain. - -Anyone is free to copy, modify, publish, use, compile, sell, or -distribute this software, either in source code form or as a compiled -binary, for any purpose, commercial or non-commercial, and by any -means. - -In jurisdictions that recognize copyright laws, the author or authors -of this software dedicate any and all copyright interest in the -software to the public domain. We make this dedication for the benefit -of the public at large and to the detriment of our heirs and -successors. We intend this dedication to be an overt act of -relinquishment in perpetuity of all present and future rights to this -software under copyright law. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, -EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF -MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. -IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR -OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, -ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR -OTHER DEALINGS IN THE SOFTWARE. - -For more information, please refer to -*/ - -#ifndef PUBLIC_DOMAIN_BASE64_HPP_ -#define PUBLIC_DOMAIN_BASE64_HPP_ - -#include -#include -#include -#include - -class base64_error : public std::runtime_error -{ -public: - using std::runtime_error::runtime_error; -}; - -class base64 -{ -public: - enum class alphabet - { - /** the alphabet is detected automatically */ - auto_, - /** the standard base64 alphabet is used */ - standard, - /** like `standard` except that the characters `+` and `/` are replaced by `-` and `_` respectively*/ - url_filename_safe - }; - - enum class decoding_behavior - { - /** if the input is not padded, the remaining bits are ignored */ - moderate, - /** if a padding character is encounter decoding is finished */ - loose - }; - - /** - Encodes all the elements from `in_begin` to `in_end` to `out`. - - @warning The source and destination cannot overlap. The destination must be able to hold at least - `required_encode_size(std::distance(in_begin, in_end))`, otherwise the behavior depends on the output iterator. - - @tparam Input_iterator the source; the returned elements are cast to `std::uint8_t` and should not be greater than - 8 bits - @tparam Output_iterator the destination; the elements written to it are from the type `char` - @param in_begin the beginning of the source - @param in_end the ending of the source - @param out the destination iterator - @param alphabet which alphabet should be used - @returns the iterator to the next element past the last element copied - @throws see `Input_iterator` and `Output_iterator` - */ - template - static Output_iterator encode(Input_iterator in_begin, Input_iterator in_end, Output_iterator out, - alphabet alphabet = alphabet::standard) - { - constexpr auto pad = '='; - const char* alpha = alphabet == alphabet::url_filename_safe - ? "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_" - : "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; - - while (in_begin != in_end) { - std::uint8_t i0 = 0, i1 = 0, i2 = 0; - - // first character - i0 = static_cast(*in_begin); - ++in_begin; - - *out = alpha[i0 >> 2 & 0x3f]; - ++out; - - // part of first character and second - if (in_begin != in_end) { - i1 = static_cast(*in_begin); - ++in_begin; - - *out = alpha[((i0 & 0x3) << 4) | (i1 >> 4 & 0x0f)]; - ++out; - } else { - *out = alpha[(i0 & 0x3) << 4]; - ++out; - - // last padding - *out = pad; - ++out; - - // last padding - *out = pad; - ++out; - - break; - } - - // part of second character and third - if (in_begin != in_end) { - i2 = static_cast(*in_begin); - ++in_begin; - - *out = alpha[((i1 & 0xf) << 2) | (i2 >> 6 & 0x03)]; - ++out; - } else { - *out = alpha[(i1 & 0xf) << 2]; - ++out; - - // last padding - *out = pad; - ++out; - - break; - } - - // rest of third - *out = alpha[i2 & 0x3f]; - ++out; - } - - return out; - } - /** - Encodes a string. - - @param str the string that should be encoded - @param alphabet which alphabet should be used - @returns the encoded base64 string - @throws see base64::encode() - */ - static std::string encode(const std::string& str, alphabet alphabet = alphabet::standard) - { - std::string result; - - result.reserve(required_encode_size(str.length()) + 1); - - encode(str.begin(), str.end(), std::back_inserter(result), alphabet); - - return result; - } - /** - Encodes a char array. - - @param buffer the char array - @param size the size of the array - @param alphabet which alphabet should be used - @returns the encoded string - */ - static std::string encode(const char* buffer, std::size_t size, alphabet alphabet = alphabet::standard) - { - std::string result; - - result.reserve(required_encode_size(size) + 1); - - encode(buffer, buffer + size, std::back_inserter(result), alphabet); - - return result; - } - /** - Decodes all the elements from `in_begin` to `in_end` to `out`. `in_begin` may point to the same location as `out`, - in other words: inplace decoding is possible. - - @warning The destination must be able to hold at least `required_decode_size(std::distance(in_begin, in_end))`, - otherwise the behavior depends on the output iterator. - - @tparam Input_iterator the source; the returned elements are cast to `char` - @tparam Output_iterator the destination; the elements written to it are from the type `std::uint8_t` - @param in_begin the beginning of the source - @param in_end the ending of the source - @param out the destination iterator - @param alphabet which alphabet should be used - @param behavior the behavior when an error was detected - @returns the iterator to the next element past the last element copied - @throws base64_error depending on the set behavior - @throws see `Input_iterator` and `Output_iterator` - */ - template - static Output_iterator decode(Input_iterator in_begin, Input_iterator in_end, Output_iterator out, - alphabet alphabet = alphabet::auto_, - decoding_behavior behavior = decoding_behavior::moderate) - { - //constexpr auto pad = '='; - std::uint8_t last = 0; - auto bits = 0; - - while (in_begin != in_end) { - auto c = *in_begin; - ++in_begin; - - if (c == '=') { - break; - } - - auto part = _base64_value(alphabet, c); - - // enough bits for one byte - if (bits + 6 >= 8) { - *out = (last << (8 - bits)) | (part >> (bits - 2)); - ++out; - - bits -= 2; - } else { - bits += 6; - } - - last = part; - } - - // check padding - if (behavior != decoding_behavior::loose) { - while (in_begin != in_end) { - auto c = *in_begin; - ++in_begin; - - if (c != '=') { - throw base64_error("invalid base64 character."); - } - } - } - - return out; - } - /** - Decodes a string. - - @param str the base64 encoded string - @param alphabet which alphabet should be used - @param behavior the behavior when an error was detected - @returns the decoded string - @throws see base64::decode() - */ - static std::string decode(const std::string& str, alphabet alphabet = alphabet::auto_, - decoding_behavior behavior = decoding_behavior::moderate) - { - std::string result; - - result.reserve(max_decode_size(str.length())); - - decode(str.begin(), str.end(), std::back_inserter(result), alphabet, behavior); - - return result; - } - /** - Decodes a string. - - @param buffer the base64 encoded buffer - @param size the size of the buffer - @param alphabet which alphabet should be used - @param behavior the behavior when an error was detected - @returns the decoded string - @throws see base64::decode() - */ - static std::string decode(const char* buffer, std::size_t size, alphabet alphabet = alphabet::auto_, - decoding_behavior behavior = decoding_behavior::moderate) - { - std::string result; - - result.reserve(max_decode_size(size)); - - decode(buffer, buffer + size, std::back_inserter(result), alphabet, behavior); - - return result; - } - /** - Decodes a string inplace. - - @param[in,out] str the base64 encoded string - @param alphabet which alphabet should be used - @param behavior the behavior when an error was detected - @throws base64::decode_inplace() - */ - static void decode_inplace(std::string& str, alphabet alphabet = alphabet::auto_, - decoding_behavior behavior = decoding_behavior::moderate) - { - str.resize(decode(str.begin(), str.end(), str.begin(), alphabet, behavior) - str.begin()); - } - /** - Decodes a char array inplace. - - @param[in,out] str the string array - @param size the length of the array - @param alphabet which alphabet should be used - @param behavior the behavior when an error was detected - @returns the pointer to the next element past the last element decoded - @throws base64::decode_inplace() - */ - static char* decode_inplace(char* str, std::size_t size, alphabet alphabet = alphabet::auto_, - decoding_behavior behavior = decoding_behavior::moderate) - { - return decode(str, str + size, str, alphabet, behavior); - } - /** - Returns the required decoding size for a given size. The value is calculated with the following formula: - - $$ - \lceil \frac{size}{4} \rceil \cdot 3 - $$ - - @param size the size of the encoded input - @returns the size of the resulting decoded buffer; this the absolute maximum - */ - static std::size_t max_decode_size(std::size_t size) noexcept - { - return (size / 4 + (size % 4 ? 1 : 0)) * 3; - } - /** - Returns the required encoding size for a given size. The value is calculated with the following formula: - - $$ - \lceil \frac{size}{3} \rceil \cdot 4 - $$ - - @param size the size of the decoded input - @returns the size of the resulting encoded buffer - */ - static std::size_t required_encode_size(std::size_t size) noexcept - { - return (size / 3 + (size % 3 ? 1 : 0)) * 4; - } - -private: - static std::uint8_t _base64_value(alphabet& alphabet, char c) - { - if (c >= 'A' && c <= 'Z') { - return c - 'A'; - } else if (c >= 'a' && c <= 'z') { - return c - 'a' + 26; - } else if (c >= '0' && c <= '9') { - return c - '0' + 52; - } - - // comes down to alphabet - if (alphabet == alphabet::standard) { - if (c == '+') { - return 62; - } else if (c == '/') { - return 63; - } - } else if (alphabet == alphabet::url_filename_safe) { - if (c == '-') { - return 62; - } else if (c == '_') { - return 63; - } - } // auto detect - else { - if (c == '+') { - alphabet = alphabet::standard; - - return 62; - } else if (c == '/') { - alphabet = alphabet::standard; - - return 63; - } else if (c == '-') { - alphabet = alphabet::url_filename_safe; - - return 62; - } else if (c == '_') { - alphabet = alphabet::url_filename_safe; - - return 63; - } - } - - throw base64_error("invalid base64 character."); - } -}; - -#endif // !PUBLIC_DOMAIN_BASE64_HPP_ diff --git a/common/build-info-rust/Cargo.toml b/common/build-info-rust/Cargo.toml new file mode 100644 index 000000000000..f10919c18e2a --- /dev/null +++ b/common/build-info-rust/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "llama-common-build-info-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[lib] +crate-type = ["staticlib", "rlib"] + +[dependencies] diff --git a/common/build-info-rust/build.rs b/common/build-info-rust/build.rs new file mode 100644 index 000000000000..d38173a36c54 --- /dev/null +++ b/common/build-info-rust/build.rs @@ -0,0 +1,15 @@ +use std::env; + +fn main() { + let build_number = env::var("LLAMA_BUILD_NUMBER").unwrap_or_else(|_| "0".to_string()); + let commit = env::var("LLAMA_BUILD_COMMIT").unwrap_or_else(|_| "unknown".to_string()); + let compiler = env::var("BUILD_COMPILER").unwrap_or_else(|_| "unknown".to_string()); + let target = env::var("BUILD_TARGET").unwrap_or_else(|_| "unknown".to_string()); + let build_info = format!("b{build_number}-{commit}"); + + println!("cargo:rustc-env=LLAMA_BUILD_NUMBER={build_number}"); + println!("cargo:rustc-env=LLAMA_BUILD_COMMIT={commit}"); + println!("cargo:rustc-env=BUILD_COMPILER={compiler}"); + println!("cargo:rustc-env=BUILD_TARGET={target}"); + println!("cargo:rustc-env=LLAMA_BUILD_INFO={build_info}"); +} diff --git a/common/build-info-rust/src/lib.rs b/common/build-info-rust/src/lib.rs new file mode 100644 index 000000000000..d5a1dba73e07 --- /dev/null +++ b/common/build-info-rust/src/lib.rs @@ -0,0 +1,66 @@ +use std::os::raw::{c_char, c_int}; + +#[no_mangle] +pub extern "C" fn llama_build_number() -> c_int { + env!("LLAMA_BUILD_NUMBER").parse::().unwrap_or(0) +} + +#[no_mangle] +pub extern "C" fn llama_commit() -> *const c_char { + concat!(env!("LLAMA_BUILD_COMMIT"), "\0").as_ptr() as *const c_char +} + +#[no_mangle] +pub extern "C" fn llama_compiler() -> *const c_char { + concat!(env!("BUILD_COMPILER"), "\0").as_ptr() as *const c_char +} + +#[no_mangle] +pub extern "C" fn llama_build_target() -> *const c_char { + concat!(env!("BUILD_TARGET"), "\0").as_ptr() as *const c_char +} + +#[no_mangle] +pub extern "C" fn llama_build_info() -> *const c_char { + concat!(env!("LLAMA_BUILD_INFO"), "\0").as_ptr() as *const c_char +} + +#[no_mangle] +pub extern "C" fn llama_print_build_info() { + eprintln!( + "llama_print_build_info: build = {} ({})", + llama_build_number(), + env!("LLAMA_BUILD_COMMIT") + ); + eprintln!( + "llama_print_build_info: built with {} for {}", + env!("BUILD_COMPILER"), + env!("BUILD_TARGET") + ); +} + +#[cfg(test)] +mod tests { + use super::*; + use std::ffi::CStr; + + unsafe fn c_string(ptr: *const c_char) -> String { + CStr::from_ptr(ptr).to_string_lossy().into_owned() + } + + #[test] + fn exposes_configured_strings() { + assert!(!unsafe { c_string(llama_commit()) }.is_empty()); + assert!(!unsafe { c_string(llama_compiler()) }.is_empty()); + assert!(!unsafe { c_string(llama_build_target()) }.is_empty()); + assert!(unsafe { c_string(llama_build_info()) }.starts_with('b')); + } + + #[test] + fn build_number_is_parseable() { + assert_eq!( + llama_build_number(), + env!("LLAMA_BUILD_NUMBER").parse::().unwrap_or(0) + ); + } +} diff --git a/common/build-info.cpp.in b/common/build-info.cpp.in deleted file mode 100644 index f888fd079fa5..000000000000 --- a/common/build-info.cpp.in +++ /dev/null @@ -1,35 +0,0 @@ -#include "build-info.h" - -#include -#include - -int LLAMA_BUILD_NUMBER = @LLAMA_BUILD_NUMBER@; -char const * LLAMA_COMMIT = "@LLAMA_BUILD_COMMIT@"; -char const * LLAMA_COMPILER = "@BUILD_COMPILER@"; -char const * LLAMA_BUILD_TARGET = "@BUILD_TARGET@"; - -int llama_build_number(void) { - return LLAMA_BUILD_NUMBER; -} - -const char * llama_commit(void) { - return LLAMA_COMMIT; -} - -const char * llama_compiler(void) { - return LLAMA_COMPILER; -} - -const char * llama_build_target(void) { - return LLAMA_BUILD_TARGET; -} - -const char * llama_build_info(void) { - static std::string s = "b" + std::to_string(LLAMA_BUILD_NUMBER) + "-" + LLAMA_COMMIT; - return s.c_str(); -} - -void llama_print_build_info(void) { - fprintf(stderr, "%s: build = %d (%s)\n", __func__, llama_build_number(), llama_commit()); - fprintf(stderr, "%s: built with %s for %s\n", __func__, llama_compiler(), llama_build_target()); -} diff --git a/common/build-info.h b/common/build-info.h deleted file mode 100644 index 382cfa78500a..000000000000 --- a/common/build-info.h +++ /dev/null @@ -1,11 +0,0 @@ -#pragma once - -int llama_build_number(void); - -const char * llama_commit(void); -const char * llama_compiler(void); - -const char * llama_build_target(void); -const char * llama_build_info(void); - -void llama_print_build_info(void); diff --git a/common/chat-auto-parser-generator.cpp b/common/chat-auto-parser-generator.cpp deleted file mode 100644 index 453559a4b04e..000000000000 --- a/common/chat-auto-parser-generator.cpp +++ /dev/null @@ -1,478 +0,0 @@ -#include "chat-auto-parser-helpers.h" -#include "chat-auto-parser.h" -#include "chat-peg-parser.h" -#include "chat.h" -#include "common.h" -#include "json-schema-to-grammar.h" -#include "log.h" -#include "nlohmann/json.hpp" -#include "peg-parser.h" - -#include -#include - -using json = nlohmann::ordered_json; - -// Helper to iterate over tools/functions -static void foreach_function(const json & tools, const std::function & fn) { - for (const auto & tool : tools) { - if (!tool.contains("type") || tool.at("type") != "function" || !tool.contains("function")) { - continue; - } - fn(tool); - } -} - -namespace autoparser { - -parser_build_context::parser_build_context(common_chat_peg_builder & p, const generation_params & inputs) : - p(p), - inputs(inputs), - reasoning_parser(p.eps()) {} - -common_chat_params peg_generator::generate_parser(const common_chat_template & tmpl, - const struct generation_params & inputs) { - // Run differential analysis to extract template structure - struct autoparser autoparser; - autoparser.analyze_template(tmpl); - return generate_parser(tmpl, inputs, autoparser); -} - -common_chat_params peg_generator::generate_parser(const common_chat_template & tmpl, - const struct generation_params & inputs, - const autoparser & autoparser) { - // Create the result structure - common_chat_params data; - data.prompt = common_chat_template_direct_apply(tmpl, inputs); - data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; - data.preserved_tokens = autoparser.preserved_tokens; - - auto parser = autoparser.build_parser(inputs); - data.parser = parser.save(); - - // Build grammar if tools are present - bool has_tools = - autoparser.tools.format.mode != tool_format::NONE && inputs.tools.is_array() && !inputs.tools.empty(); - std::string trigger_marker = !autoparser.tools.format.section_start.empty() ? autoparser.tools.format.section_start : - autoparser.tools.format.per_call_start; - - bool has_response_format = !inputs.json_schema.empty() && inputs.json_schema.is_object(); - bool include_grammar = has_response_format || (has_tools && - ((inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO && !trigger_marker.empty()) || - inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); - - if (include_grammar) { - data.grammar_lazy = !has_response_format && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; - data.grammar = build_grammar([&](const common_grammar_builder & builder) { - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - auto schema = function.contains("parameters") ? function.at("parameters") : json::object(); - builder.resolve_refs(schema); - }); - if (has_response_format) { - auto schema = inputs.json_schema; - builder.resolve_refs(schema); - } - parser.build_grammar(builder, data.grammar_lazy); - }); - - // Set grammar triggers based on tool section markers (fall back to per-call markers) - if (data.grammar_lazy) { - data.grammar_triggers = { - { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, trigger_marker } - }; - } - } - - return data; -} - -common_peg_arena autoparser::build_parser(const generation_params & inputs) const { - if (!analysis_complete) { - throw std::invalid_argument("Cannot call build_parser on autoparser without performing analysis first, call analyze_template(...)"); - } - return build_chat_peg_parser([&](common_chat_peg_builder & p) { - parser_build_context ctx(p, inputs); - bool extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; - - ctx.extracting_reasoning = extract_reasoning && reasoning.mode != reasoning_mode::NONE; - ctx.content = &content; - ctx.reasoning = &reasoning; - - // Build reasoning parser - ctx.reasoning_parser = reasoning.build_parser(ctx); - - auto parser = p.eps(); - - bool has_tools = inputs.tools.is_array() && !inputs.tools.empty(); - bool has_response_format = inputs.json_schema.is_object() && !inputs.json_schema.empty(); - bool pure_content = reasoning.mode == reasoning_mode::NONE; - - if (has_response_format) { - auto response_format = p.rule("response-format", p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema))); - parser = ctx.reasoning_parser + p.space() + p.choice({ - p.literal("```json") + p.space() + response_format + p.space() + p.literal("```"), - response_format - }) + p.end(); - pure_content = false; - } else if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE && jinja_caps.supports_tool_calls) { - parser = tools.build_parser(ctx); - pure_content = false; - } else { - parser = content.build_parser(ctx); - } - return pure_content ? p.prefix(inputs.generation_prompt, reasoning.start) + parser : p.prefix(inputs.generation_prompt, reasoning.start) << parser; - }); -} - -common_peg_parser analyze_reasoning::build_parser(parser_build_context & ctx) const { - auto & p = ctx.p; - - if (!ctx.extracting_reasoning) { - return p.eps(); - } - - if (mode == reasoning_mode::TAG_BASED || mode == reasoning_mode::TOOLS_ONLY) { - if (!end.empty()) { - if (!start.empty()) { - // Standard tag-based: optional(reasoning) - return p.optional(start + p.reasoning(p.until(end)) + end + p.space()); - } - // Delimiter-style (empty start) - return p.optional(p.reasoning(p.until(end)) + end + p.space()); - } - } - - return p.eps(); -} - -common_peg_parser analyze_content::build_parser(parser_build_context & ctx) const { - auto & p = ctx.p; - - if (is_always_wrapped()) { - if (ctx.extracting_reasoning) { - return ctx.reasoning_parser + start + p.content(p.until(end)) + end + p.end(); - } - return p.content(p.until(start)) + start + p.content(p.until(end)) + end + p.end(); - } - return ctx.reasoning_parser + p.content(p.rest()) + p.end(); -} - -common_peg_parser analyze_content::build_optional_wrapped(parser_build_context & ctx) const { - auto & p = ctx.p; - - if (is_always_wrapped()) { - return p.optional(start + p.content(p.until(end)) + end); - } - return p.eps(); -} - -common_peg_parser analyze_tools::build_parser(parser_build_context & ctx) const { - switch (format.mode) { - case tool_format::JSON_NATIVE: - return build_tool_parser_json_native(ctx); - case tool_format::TAG_WITH_JSON: - return build_tool_parser_tag_json(ctx); - case tool_format::TAG_WITH_TAGGED: - return build_tool_parser_tag_tagged(ctx); - default: - LOG_ERR("[ERROR] Template seems to support tool calls, but failed to determine tool format. Tool calling will not work properly. " - "Check for a fixed template for your model in the models/templates directory of your llama.cpp installation or " - "report an issue at https://github.com/ggml-org/llama.cpp/issues\n"); - return ctx.p.eps(); - } -} - -common_peg_parser analyze_tools::build_tool_parser_json_native(parser_build_context & ctx) const { - auto & p = ctx.p; - const auto & inputs = ctx.inputs; - bool force_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; - - // Build effective field names with dot notation if function_field is set - std::string name_field = format.name_field; - std::string args_field = format.args_field; - - if (!format.function_field.empty() && format.function_field != "function" && - name_field.find('.') == std::string::npos) { - name_field = format.function_field + "." + name_field; - args_field = format.function_field + "." + args_field; - } - - auto tools_parser = p.eps(); - if (format.section_start.empty() && !format.per_call_start.empty()) { - auto single_tool_parser = p.standard_json_tools( - format.per_call_start, format.per_call_end, inputs.tools, inputs.parallel_tool_calls, - inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED, name_field, args_field, format.tools_array_wrapped, - format.fun_name_is_key, format.id_field, format.gen_id_field, format.parameter_order); - tools_parser = p.trigger_rule("tool-calls", p.one_or_more(single_tool_parser + p.space())); - } else { - tools_parser = p.standard_json_tools( - format.section_start, format.section_end, inputs.tools, inputs.parallel_tool_calls, - inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED, name_field, args_field, format.tools_array_wrapped, - format.fun_name_is_key, format.id_field, format.gen_id_field, format.parameter_order); - } - - // Handle content wrappers if present - if (ctx.content && ctx.content->is_always_wrapped()) { - auto wrapped_content = ctx.content->build_optional_wrapped(ctx); - return ctx.reasoning_parser + wrapped_content + tools_parser + p.end(); - } - - std::string tool_start = "{"; - if (!format.section_start.empty()) { - tool_start = format.section_start; - } else if (!format.per_call_start.empty()) { - tool_start = format.per_call_start; - } - - return ctx.reasoning_parser + (force_tools ? p.eps() : p.optional(p.content(p.until(tool_start)))) + tools_parser + - p.end(); -} - -common_peg_parser analyze_tools::build_func_parser(common_chat_peg_builder & p, const std::string & name, - const common_peg_parser & call_id_section, bool have_call_id, - const common_peg_parser & args, - std::optional atomic_peek) const { - auto open = p.tool_open(function.name_prefix + p.tool_name(p.literal(name)) + function.name_suffix); - bool matched_atomic = false; - common_peg_parser func_parser = p.eps(); - - if (!function.name_suffix.empty()) { - func_parser = open + call_id_section + p.space() + args; - matched_atomic = true; - } else if (have_call_id) { - func_parser = p.atomic(open + call_id_section) + p.space() + args; - matched_atomic = true; - } else if (atomic_peek.has_value()) { - func_parser = p.atomic(open + call_id_section + p.space() + *atomic_peek) + args; - matched_atomic = true; - } else { - func_parser = open + call_id_section + p.space() + args; - } - - if (!function.close.empty()) { - func_parser = func_parser + p.space() + p.tool_close(p.literal(function.close)); - } else if (!format.per_call_end.empty()) { - // When there's no func_close but there is a per_call_end marker, use peek() to ensure - // we only emit tool_close when we can actually see the closing marker. This prevents - // premature closing during partial parsing when we've seen e.g. "" (end) or "" prefix that failed to match. - func_parser = func_parser + p.tool_close(p.peek(p.literal(format.per_call_end))); - } else { - func_parser = func_parser + p.tool_close(p.space()); // force this to process tool closing callbacks in mapper - } - if (!matched_atomic) { - func_parser = p.atomic(func_parser); - } - return func_parser; -} - -common_peg_parser analyze_tools::build_tool_parser_tag_json(parser_build_context & ctx) const { - auto & p = ctx.p; - const auto & inputs = ctx.inputs; - bool force_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; - - common_peg_parser tool_choice = p.choice(); - - foreach_function(inputs.tools, [&](const json & tool) { - const auto & func = tool.at("function"); - std::string name = func.at("name"); - const auto & schema = func.contains("parameters") ? func.at("parameters") : json::object(); - - // Build call_id parser based on position (if supported) - bool have_call_id = false; - common_peg_parser call_id_section = p.eps(); - if (call_id.pos == call_id_position::BETWEEN_FUNC_AND_ARGS && !call_id.prefix.empty() && - (!call_id.suffix.empty() || !arguments.start.empty())) { - if (!call_id.suffix.empty()) { - call_id_section = p.optional(call_id.prefix + p.tool_id(p.until(call_id.suffix))) + call_id.suffix; - } else { - call_id_section = p.optional(call_id.prefix + p.tool_id(p.until(arguments.start))); - } - have_call_id = true; - } - auto args_parser = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)); - if (!arguments.start.empty()) { - args_parser = p.literal(arguments.start) + args_parser; - } - if (!arguments.end.empty()) { - args_parser = args_parser + p.literal(arguments.end); - } - - auto atomic_peek = !arguments.start.empty() ? std::optional(p.peek(p.literal(arguments.start))) : std::nullopt; - auto func_parser = build_func_parser(p, name, call_id_section, have_call_id, args_parser, atomic_peek); - tool_choice |= p.rule("tool-" + name, func_parser); - }); - - auto require_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; - - common_peg_parser tool_calls = p.eps(); - - if (!format.per_call_start.empty()) { - auto wrapped_call = format.per_call_start + tool_choice + format.per_call_end; - if (inputs.parallel_tool_calls) { - tool_calls = p.trigger_rule("tool-call", wrapped_call + p.zero_or_more(p.space() + wrapped_call)); - } else { - tool_calls = p.trigger_rule("tool-call", wrapped_call); - } - if (!format.section_start.empty()) { - tool_calls = p.trigger_rule("tool-calls", - p.literal(format.section_start) + p.space() + tool_calls + p.space() + - (format.section_end.empty() ? p.end() : p.literal(format.section_end))); - } - } else { - std::string separator = ", "; // Default - if (inputs.parallel_tool_calls) { - tool_calls = p.trigger_rule("tool-call", format.section_start + tool_choice + - p.zero_or_more(separator + tool_choice) + format.section_end); - } else { - tool_calls = p.trigger_rule("tool-call", format.section_start + tool_choice + format.section_end); - } - } - - if (!require_calls) { - tool_calls = p.optional(tool_calls); - } - - std::string trigger_marker = !format.section_start.empty() ? format.section_start : format.per_call_start; - auto content_before_tools = trigger_marker.empty() ? p.eps() : p.until(trigger_marker); - return ctx.reasoning_parser + (force_tools ? p.eps() : p.optional(p.content(content_before_tools))) + tool_calls + - p.end(); -} - -common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_context & ctx) const { - auto & p = ctx.p; - const auto & inputs = ctx.inputs; - bool force_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; - - auto until_suffix = p.rule("until-suffix", p.until(arguments.value_suffix)); - - common_peg_parser tool_choice = p.choice(); - - foreach_function(inputs.tools, [&](const json & tool) { - const auto & func = tool.at("function"); - std::string name = func.at("name"); - auto params = func.contains("parameters") ? func.at("parameters") : json::object(); - const auto & properties = params.contains("properties") ? params.at("properties") : json::object(); - - std::set required; - if (params.contains("required")) { - params.at("required").get_to(required); - } - - auto schema_info = common_schema_info(); - schema_info.resolve_refs(params); - - // Build parser for each argument, separating required and optional - std::vector required_parsers; - std::vector optional_parsers; - for (const auto & [param_name, param_schema] : properties.items()) { - bool is_required = required.find(param_name) != required.end(); - - auto arg = - p.tool_arg(p.tool_arg_open(arguments.name_prefix + p.tool_arg_name(p.literal(param_name)) + - arguments.name_suffix) + - arguments.value_prefix + - (schema_info.resolves_to_string(param_schema) ? - p.tool_arg_string_value(p.schema(until_suffix, - "tool-" + name + "-arg-" + param_name + "-schema", - param_schema, true)) : - p.tool_arg_json_value(p.schema( - p.json(), "tool-" + name + "-arg-" + param_name + "-schema", param_schema, false)) + - p.space()) + - p.tool_arg_close(p.literal(arguments.value_suffix))); - - auto named_arg = p.rule("tool-" + name + "-arg-" + param_name, arg); - if (is_required) { - required_parsers.push_back(named_arg); - } else { - optional_parsers.push_back(named_arg); - } - } - - // Build required arg sequence in definition order - common_peg_parser args_seq = p.eps(); - for (size_t i = 0; i < required_parsers.size(); i++) { - if (i > 0) { - args_seq = args_seq + p.space(); - } - args_seq = args_seq + required_parsers[i]; - } - - // Build optional args with flexible ordering - if (!optional_parsers.empty()) { - common_peg_parser any_opt = p.choice(); - for (const auto & opt : optional_parsers) { - any_opt |= opt; - } - args_seq = args_seq + p.repeat(p.space() + any_opt, 0, -1); - } - - if (!arguments.start.empty()) { - args_seq = p.literal(arguments.start) + args_seq; - } - if (!arguments.end.empty()) { - args_seq = args_seq + p.literal(arguments.end); - } - - // Build call_id parser based on position (if supported) - common_peg_parser call_id_section = p.eps(); - bool have_call_id = false; - if (call_id.pos == call_id_position::BETWEEN_FUNC_AND_ARGS && !call_id.prefix.empty() && - (!call_id.suffix.empty() || !arguments.start.empty())) { - have_call_id = true; - if (!call_id.suffix.empty()) { - call_id_section = p.optional(call_id.prefix + p.tool_id(p.until(call_id.suffix)) + call_id.suffix); - } else { - call_id_section = p.optional(call_id.prefix + p.tool_id(p.until(arguments.start))); - } - } - - // Only peek for an arg tag when there are required args that must follow. - // When all args are optional, the model may emit no arg tags at all (#20650). - auto atomic_peek = (!arguments.name_prefix.empty() && !required_parsers.empty()) ? - std::optional(p.peek(p.literal(arguments.name_prefix))) : std::nullopt; - auto func_parser = build_func_parser(p, name, call_id_section, have_call_id, args_seq, atomic_peek); - tool_choice |= p.rule("tool-" + name, func_parser); - }); - - auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; - - common_peg_parser tool_calls = p.eps(); - - if (!format.per_call_start.empty()) { - auto wrapped_call = format.per_call_start + p.space() + tool_choice + p.space() + format.per_call_end; - if (inputs.parallel_tool_calls) { - tool_calls = p.trigger_rule("tool-call", wrapped_call + p.zero_or_more(p.space() + wrapped_call) + p.space()); - } else { - tool_calls = p.trigger_rule("tool-call", wrapped_call + p.space()); - } - if (!format.section_start.empty()) { - tool_calls = p.trigger_rule("tool-calls", - p.literal(format.section_start) + p.space() + tool_calls + p.space() + - (format.section_end.empty() ? p.end() : p.literal(format.section_end) + p.space())); - } - } else { - std::string separator = ", "; // Default - - if (inputs.parallel_tool_calls) { - tool_calls = p.trigger_rule("tool-call", format.section_start + p.space() + tool_choice + - p.zero_or_more(separator + tool_choice) + p.space() + - format.section_end); - } else { - tool_calls = p.trigger_rule( - "tool-call", format.section_start + p.space() + tool_choice + p.space() + format.section_end); - } - } - - if (!require_tools) { - tool_calls = p.optional(tool_calls); - } - - std::string trigger_marker = !format.section_start.empty() ? format.section_start : format.per_call_start; - auto content_before_tools = trigger_marker.empty() ? p.eps() : p.until(trigger_marker); - return ctx.reasoning_parser + (force_tools ? p.eps() : p.optional(p.content(content_before_tools))) + tool_calls + - p.end(); -} - -} // namespace autoparser diff --git a/common/chat-auto-parser-helpers.cpp b/common/chat-auto-parser-helpers.cpp deleted file mode 100644 index 2499464cd82f..000000000000 --- a/common/chat-auto-parser-helpers.cpp +++ /dev/null @@ -1,364 +0,0 @@ -#include "chat-auto-parser-helpers.h" - -#include "chat-auto-parser.h" -#include "chat-peg-parser.h" -#include "chat.h" -#include "log.h" -#include "nlohmann/json.hpp" -#include "peg-parser.h" - -#include -#include - -using json = nlohmann::ordered_json; - -std::string trim_whitespace(const std::string & str) { - size_t start = 0; - while (start < str.length() && std::isspace(static_cast(str[start]))) { - start++; - } - - if (start == str.length()) { - return ""; - } - - size_t end = str.length() - 1; - while (end > start && std::isspace(static_cast(str[end]))) { - end--; - } - - return str.substr(start, end - start + 1); -} - -std::string trim_leading_whitespace(const std::string & str) { - size_t start = 0; - while (start < str.length() && std::isspace(static_cast(str[start]))) { - start++; - } - - return str.substr(start); -} - -std::string trim_trailing_whitespace(const std::string & str) { - if (str.empty()) { - return ""; - } - - size_t end = str.length() - 1; - while (end > 0 && std::isspace(static_cast(str[end]))) { - end--; - } - - // If first char is also whitespace, return empty string - if (end == 0 && std::isspace(static_cast(str[0]))) { - return ""; - } - - return str.substr(0, end + 1); -} - -std::string trim_trailing_newlines(const std::string & str) { - size_t end = str.length(); - while (end > 0 && str[end - 1] == '\n') { - end--; - } - - return str.substr(0, end); -} - -static size_t common_prefix_len(const std::string & left, const std::string & right) { - size_t prefix_len = 0; - size_t min_len = std::min(left.length(), right.length()); - while (prefix_len < min_len && left[prefix_len] == right[prefix_len]) { - prefix_len++; - } - return prefix_len; -} - -static size_t common_suffix_len(const std::string & left, const std::string & right) { - size_t suffix_len = 0; - size_t min_len = std::min(left.length(), right.length()); - while (suffix_len < min_len && left[left.length() - 1 - suffix_len] == right[right.length() - 1 - suffix_len]) { - suffix_len++; - } - return suffix_len; -} - -diff_split calculate_diff_split(const std::string & left, const std::string & right) { - diff_split result; - - auto left_seg = segmentize_markers(left); - auto right_seg = segmentize_markers(right); - - if (left_seg.empty()) { - result.right = right; - return result; - } - if (right_seg.empty()) { - result.left = left; - return result; - } - - auto left_start = left_seg.begin(); - auto left_end = --left_seg.end(); - auto right_start = right_seg.begin(); - auto right_end = --right_seg.end(); - - auto test = [&] () { - return left_start != left_end && right_start != right_end; - }; - - bool left_fully_consumed = false; - bool right_fully_consumed = false; - - while (test()) { - bool advanced = false; - if (*left_start == *right_start) { - result.prefix.append(left_start->value); - left_start++; - right_start++; - advanced = true; - } - if (*left_end == *right_end) { - result.suffix = left_end->value + result.suffix; - if (left_start != left_end) { - left_end--; - } else { - left_fully_consumed = true; - } - if (right_start != right_end) { - right_end--; - } else { - right_fully_consumed = true; - } - advanced = true; - } - if (!advanced) { - break; - } - } - - if (left_start == left_end && right_start != right_end) { - if (*left_start == *right_end) { - result.suffix = right_end->value + result.suffix; - right_end--; - left_fully_consumed = true; - } else if (*left_start == *right_start) { - result.prefix.append(right_start->value); - right_start++; - left_fully_consumed = true; - } - } else if (right_start == right_end && left_start != left_end) { - if (*left_end == *right_start) { - result.suffix = left_end->value + result.suffix; - left_end--; - right_fully_consumed = true; - } else if (*left_start == *right_start) { - result.prefix.append(left_start->value); - left_start++; - right_fully_consumed = true; - } - } else if (left_start == left_end && right_start == right_end && *left_start == *right_start && left_start->type == segment_type::MARKER) { - result.prefix.append(right_start->value); - left_fully_consumed = true; - right_fully_consumed = true; - } - - auto eat_segment = [](std::string str, const segment & seg) -> std::string { return std::move(str) + seg.value; }; - - bool can_have_text_suffix = left_end->type == segment_type::TEXT && right_end->type == segment_type::TEXT; - bool can_have_text_prefix = right_start->type == segment_type::TEXT && left_start->type == segment_type::TEXT; - - std::string remainder_left = std::accumulate(left_start, left_fully_consumed ? left_end : ++left_end, std::string(), eat_segment); - std::string remainder_right = std::accumulate(right_start, right_fully_consumed ? right_end : ++right_end, std::string(), eat_segment); - - size_t suffix_len = can_have_text_suffix ? common_suffix_len(remainder_left, remainder_right) : 0; - // avoid overlaps between prefix and suffix - size_t prefix_len = can_have_text_prefix ? common_prefix_len(remainder_left.substr(0, remainder_left.size() - suffix_len), - remainder_right.substr(0, remainder_right.size() - suffix_len)) : 0; - - result.prefix.append(remainder_left.substr(0, prefix_len)); - result.suffix = remainder_left.substr(remainder_left.length() - suffix_len, suffix_len) + result.suffix; - result.left = remainder_left.substr(prefix_len, remainder_left.length() - prefix_len - suffix_len); - result.right = remainder_right.substr(prefix_len, remainder_right.length() - prefix_len - suffix_len); - - if (result.left == "" && result.right == "") { - // degenerate case, no diff - result.prefix = left; - result.suffix = ""; - // pick prefix = all as representation - } - - // When left has no unique content (result.left is empty), left is entirely - // shared with right. The simultaneous prefix/suffix segment matching can - // incorrectly consume trailing segments of left as suffix when those same - // segments also appear at the end of right (e.g. "\n" at the end of both - // the shared content and the generation prompt). This rotates the diff. - // Fix: if left is a prefix of right, enforce that directly. - if (result.left.empty() && !result.right.empty() && - left.size() <= right.size() && - right.substr(0, left.size()) == left) { - result.prefix = left; - result.suffix = ""; - result.right = right.substr(left.size()); - } - - return result; -} - -// Returns the prefix of `full` up until the first occurrence of the common prefix of `left` and `right` -std::string until_common_prefix(const std::string & full, const std::string & left, const std::string & right) { - // Find the common prefix of left and right - size_t common_prefix_len = 0; - size_t min_len = std::min(left.length(), right.length()); - while (common_prefix_len < min_len && left[common_prefix_len] == right[common_prefix_len]) { - common_prefix_len++; - } - - // If there's no common prefix, return empty string - if (common_prefix_len == 0) { - return ""; - } - - // Find the common prefix in the full string - std::string common_prefix = left.substr(0, common_prefix_len); - size_t pos = full.find(common_prefix); - - // If not found, return empty string - if (pos == std::string::npos) { - return ""; - } - - // Return everything before the common prefix - return full.substr(0, pos); -} - -// Returns the suffix of `full` after the last occurrence of the common suffix of `left` and `right` -std::string after_common_suffix(const std::string & full, const std::string & left, const std::string & right) { - // Find the common suffix of left and right (compare from the end) - size_t common_suffix_len = 0; - size_t min_len = std::min(left.length(), right.length()); - while (common_suffix_len < min_len && - left[left.length() - 1 - common_suffix_len] == right[right.length() - 1 - common_suffix_len]) { - common_suffix_len++; - } - - // If there's no common suffix, return empty string - if (common_suffix_len == 0) { - return ""; - } - - // Extract the common suffix - std::string common_suffix = left.substr(left.length() - common_suffix_len); - - // Find the last occurrence of the common suffix in the full string - size_t pos = full.rfind(common_suffix); - - // If not found, return empty string - if (pos == std::string::npos) { - return ""; - } - - // Return everything after the common suffix - return full.substr(pos + common_suffix_len); -} - -// TODO: segmentize will treat a JSON array inside tags as a tag: [{ "fun": { ... } }] will be three markers -// not too worried about that because it hasn't turned out as a problem anywhere, but noting here in case it will -// Might have to put some restrictions on tag contents as well (like "no { }") -std::vector segmentize_markers(const std::string & text) { - std::vector retval; - bool in_marker = false; - char marker_opener = '\0'; - - auto is_marker_opener = [](char c) -> bool { return c == '<' || c == '['; }; - auto is_marker_closer = [](char op, char c) -> bool { return (op == '<' && c == '>') || (op == '[' && c == ']'); }; - - size_t last_border = 0; - - for (size_t cur_pos = 0; cur_pos < text.length(); cur_pos++) { - if (!in_marker && is_marker_opener(text[cur_pos])) { - if (last_border < cur_pos) { - retval.push_back(segment(segment_type::TEXT, text.substr(last_border, cur_pos - last_border))); - } - last_border = cur_pos; - in_marker = true; - marker_opener = text[cur_pos]; - } else if (in_marker && is_marker_closer(marker_opener, text[cur_pos])) { - // no need to check because last_border will always be smaller - retval.push_back(segment(segment_type::MARKER, text.substr(last_border, cur_pos - last_border + 1))); - last_border = cur_pos + 1; - in_marker = false; - marker_opener = '\0'; - } - } - if (last_border < text.length()) { - retval.push_back(segment(segment_type::TEXT, text.substr(last_border))); - } - return retval; -} - -std::vector prune_whitespace_segments(const std::vector & segments) { - std::vector result; - for (const auto & seg : segments) { - if (!trim_whitespace(seg.value).empty()) { - result.push_back(seg); - } - } - return result; -} - -namespace autoparser { - -std::string apply_template(const common_chat_template & tmpl, const template_params & params) { - generation_params tmpl_params; - tmpl_params.messages = params.messages; - tmpl_params.tools = params.tools; - tmpl_params.add_generation_prompt = params.add_generation_prompt; - tmpl_params.enable_thinking = params.enable_thinking; - - if (params.extra_context) { - tmpl_params.extra_context = *params.extra_context; - } - tmpl_params.extra_context["enable_thinking"] = params.enable_thinking; - - try { - return common_chat_template_direct_apply(tmpl, tmpl_params); - } catch (const std::exception & e) { - LOG_DBG("Template application failed: %s\n", e.what()); - return ""; - } -} - -std::optional compare_variants( - const common_chat_template & tmpl, - const template_params & params_A, - const std::function & params_modifier) { - // Create variant B by copying A - template_params params_B = params_A; - - // Apply modifier to create variant B - if (params_modifier) { - params_modifier(params_B); - } - - // Apply template to both variants - std::string output_A = apply_template(tmpl, params_A); - std::string output_B = apply_template(tmpl, params_B); - - // Check for template application failures - if (output_A.empty() || output_B.empty()) { - return std::nullopt; - } - - // Calculate diff and return result with both outputs - compare_variants_result result; - result.diff = calculate_diff_split(output_A, output_B); - result.output_A = output_A; - result.output_B = output_B; - - return result; -} - -} // namespace autoparser - diff --git a/common/chat-auto-parser-helpers.h b/common/chat-auto-parser-helpers.h deleted file mode 100644 index b8804ac1912d..000000000000 --- a/common/chat-auto-parser-helpers.h +++ /dev/null @@ -1,74 +0,0 @@ -#pragma once - -#include "chat-auto-parser.h" - -#include -#include -#include - -std::string trim_whitespace(const std::string & str); -std::string trim_leading_whitespace(const std::string & str); -std::string trim_trailing_whitespace(const std::string & str); -std::string trim_trailing_newlines(const std::string & str); - -// calculate a diff split (longest common prefix, longest common suffix excluding prefix, -// mismatched part on the left, mismatched part on the right) between two strings -// account for markers - align prefix and suffix endings so that they end on markers -// * eg.: -// calculate_diff_split("
", "

Something

") -> -// { "prefix": "" (not: "<"), "suffix": "", "left": "
", "right": "

Something

" } -// calculate_diff_split("Something", "") -> -// { "prefix": "", "suffix": "", "left": "Something", "right": "" } -diff_split calculate_diff_split(const std::string & left, const std::string & right); - -// Returns the prefix of `full` up until the first occurrence of the common prefix of `left` and `right` -// Returns empty string if there's no common prefix -// * eg.: -// until_common_prefix("really want a FUNCTION call", "FUNCTION alpha", "FUNCTION beta") -> "really want a " -// until_common_prefix("", "", "") -> "" -// until_common_prefix("some text", "1234", "abcd") -> "" -// until_common_prefix("one arg two args three args four", "argument alpha", "argument beta") -> "one "" -std::string until_common_prefix(const std::string & full, const std::string & left, const std::string & right); - -// Returns the suffix of `full` after the last occurrence of the common suffix of `left` and `right` -// Returns empty string if there's no common suffix -// Mirror function of `until_common_prefix` -// * eg.: -// after_common_suffix("really want a FUNCTION call", "first FUNCTION", "second FUNCTION") -> " call" -// after_common_suffix("one arg two-args three args four", "alpha-args", "beta-args") -> " three args four" -std::string after_common_suffix(const std::string & full, const std::string & left, const std::string & right); - -// Segmentize text into markers and non-marker fragments -// * eg.: -// segmentize_markers("The site title
Here's some content
" -> -// [ (MARKER, ""), (MARKER, ""), (MARKER, ""), (TEXT, "The site title"), (MARKER, ""), -// (MARKER, ""), (MARKER, "
"), (TEXT, "Here's some "), (MARKER, ""), (TEXT, "content"), (MARKER, ""), -// (MARKER, "
"), (MARKER, ""), (MARKER, "") -// ] -// segmentize_markers("<|tool_call|>[args]{ are here }[/args]<|tool_call_end|>") -> -// [ (MARKER, "<|tool_call|>"), (MARKER, "[args]"), (TEXT, "{ are here }"), (MARKER, "[/args]"), (MARKER, "<|tool_call_end|>") ] -std::vector segmentize_markers(const std::string & text); - -// Prune whitespace-only segments from a vector of segments -// * eg.: -// segmentize_markers("\n\n\n \n\n\n") -> -// X = [ (MARKER, ""), (TEXT, "\n"), (MARKER, ""), (TEXT, "\n"), (MARKER, ""), (TEXT, "\n \n"), -// (MARKER, ""), (TEXT, "\n"), (MARKER, ""), (TEXT, "\n"), (MARKER, "") ] -// prune_whitespace_segments(X) -> [ (MARKER, ""), (MARKER, ""), (MARKER, ""), (MARKER, ""), -// (MARKER, ""), (MARKER, "") ] -std::vector prune_whitespace_segments(const std::vector & segments); - -namespace autoparser { - -// Apply a template with the given parameters, returning the rendered string (empty on failure) -std::string apply_template(const common_chat_template & tmpl, const template_params & params); - -// Factorized differential comparison function -// Takes base params and a single modifier lambda to create variant B -// Returns compare_variants_result containing diff and both outputs, or std::nullopt on failure -std::optional compare_variants( - const common_chat_template & tmpl, - const template_params & params_A, - const std::function & params_modifier); - -} // namespace autoparser diff --git a/common/chat-auto-parser.h b/common/chat-auto-parser.h deleted file mode 100644 index 6c547409760d..000000000000 --- a/common/chat-auto-parser.h +++ /dev/null @@ -1,438 +0,0 @@ -#pragma once - -#include "chat.h" -#include "common.h" -#include "jinja/caps.h" -#include "peg-parser.h" -#include "nlohmann/json.hpp" - -#include -#include -#include -#include -#include - -using json = nlohmann::ordered_json; - -class common_chat_peg_builder; - -// ============================================================================ -// Parameters for template application (low-level, used by diff analysis) -// ============================================================================ -struct template_params { - json messages; - json tools; - bool add_generation_prompt = false; - bool enable_thinking = true; - std::optional extra_context = std::nullopt; -}; - -struct diff_split { - std::string prefix; - std::string suffix; - std::string left; - std::string right; - - bool operator==(struct diff_split & other) const { - return prefix == other.prefix && suffix == other.suffix && left == other.left && right == other.right; - } -}; - -// Result of compare_variants containing diff and original outputs -struct compare_variants_result { - diff_split diff; - std::string output_A; - std::string output_B; -}; - -namespace autoparser { - -// ============================================================================ -// High-level params for parser generation -// ============================================================================ - -struct generation_params { - json messages; - json tools; - common_chat_tool_choice tool_choice = COMMON_CHAT_TOOL_CHOICE_AUTO; - json json_schema; - bool parallel_tool_calls = true; - common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_AUTO; - bool stream = true; - std::string grammar; - bool add_generation_prompt = false; - bool enable_thinking = true; - std::chrono::system_clock::time_point now = std::chrono::system_clock::now(); - std::string generation_prompt; - json extra_context; - bool add_bos = false; - bool add_eos = false; - bool is_inference = true; - bool add_inference = false; - bool mark_input = true; // whether to mark input strings in the jinja context -}; - -// ============================================================================ -// Analysis Result Enums -// ============================================================================ - -// Reasoning handling mode (derived from R1-R3 comparisons) -enum class reasoning_mode { - NONE, // No reasoning markers detected - TAG_BASED, // Tag-based: ... (start can be empty for delimiter-style) - TOOLS_ONLY // Only reason on tool calls, not on normal content -}; - -inline std::ostream & operator<<(std::ostream & os, const reasoning_mode & mode) { - switch (mode) { - case reasoning_mode::NONE: - return os << "NONE"; - case reasoning_mode::TAG_BASED: - return os << "TAG_BASED"; - case reasoning_mode::TOOLS_ONLY: - return os << "TOOLS_ONLY"; - default: - return os << "UNKNOWN"; - } -} - -// Content wrapping mode (derived from C1 comparison) -enum class content_mode { - PLAIN, // No content markers - ALWAYS_WRAPPED, // Content always wrapped with markers - WRAPPED_WITH_REASONING, // Content wrapped only when reasoning present -}; - -inline std::ostream & operator<<(std::ostream & os, const content_mode & mode) { - switch (mode) { - case content_mode::PLAIN: - return os << "PLAIN"; - case content_mode::ALWAYS_WRAPPED: - return os << "ALWAYS_WRAPPED"; - case content_mode::WRAPPED_WITH_REASONING: - return os << "WRAPPED_WITH_REASONING"; - default: - return os << "UNKNOWN"; - } -} - -// Call ID position in tool calls (for non-JSON formats) -enum class call_id_position { - NONE, // No call ID support detected - PRE_FUNC_NAME, // Call ID before function name: [CALL_ID]id[FUNC]name{args} - BETWEEN_FUNC_AND_ARGS, // Call ID between function and args: [FUNC]name[CALL_ID]id{args} - POST_ARGS, // Call ID after arguments: [FUNC]name{args}[CALL_ID]id -}; - -inline std::ostream & operator<<(std::ostream & os, const call_id_position & pos) { - switch (pos) { - case call_id_position::NONE: - return os << "NONE"; - case call_id_position::PRE_FUNC_NAME: - return os << "PRE_FUNC_NAME"; - case call_id_position::BETWEEN_FUNC_AND_ARGS: - return os << "BETWEEN_FUNC_AND_ARGS"; - case call_id_position::POST_ARGS: - return os << "POST_ARGS"; - default: - return os << "UNKNOWN"; - } -} - -// Tool call format classification (derived from T1-T5, A1-A3 comparisons) -enum class tool_format { - NONE, // No tool support detected - JSON_NATIVE, // Pure JSON: {"name": "X", "arguments": {...}} - TAG_WITH_JSON, // Tag-based with JSON args: {...} - TAG_WITH_TAGGED, // Tag-based with tagged args: value -}; - -inline std::ostream & operator<<(std::ostream & os, const tool_format & format) { - switch (format) { - case tool_format::NONE: - return os << "NONE"; - case tool_format::JSON_NATIVE: - return os << "JSON_NATIVE"; - case tool_format::TAG_WITH_JSON: - return os << "TAG_WITH_JSON"; - case tool_format::TAG_WITH_TAGGED: - return os << "TAG_WITH_TAGGED"; - default: - return os << "UNKNOWN"; - } -} - -// ============================================================================ -// Sub-structs for tool analysis -// ============================================================================ - -struct tool_format_analysis { - tool_format mode = tool_format::NONE; - - std::string section_start; // e.g., "", "[TOOL_CALLS]", "" - std::string section_end; // e.g., "", "" - std::string per_call_start; // e.g., "<|tool_call_begin|>", "" (for multi-call templates) - std::string per_call_end; // e.g., "<|tool_call_end|>", "" - - bool fun_name_is_key = false; // In JSON format function name is JSON key, i.e. { "": { ... arguments ... } } - bool tools_array_wrapped = false; // Tool calls wrapped in JSON array [...] - - std::string function_field = "function"; - std::string name_field = "name"; - std::string args_field = "arguments"; - std::string id_field; - std::string gen_id_field; - std::vector parameter_order; -}; - -struct tool_function_analysis { - std::string name_prefix; // e.g., "", "\"", ":0" - std::string close; // e.g., "", "" (for tag-based) -}; - -struct tool_arguments_analysis { - std::string start; // e.g., "<|tool_call_argument_begin|>", "" - std::string end; // e.g., "<|tool_call_argument_end|>", "" - std::string name_prefix; // e.g., "", "\"" - std::string name_suffix; // e.g., ">", "", "\":" - std::string value_prefix; // e.g., "", "", "" - std::string value_suffix; // e.g., "", "", "" - std::string separator; // e.g., "", "\n", "," -}; - -struct tool_id_analysis { - call_id_position pos = call_id_position::NONE; - - std::string prefix; // e.g., "[CALL_ID]" (marker before call ID value) - std::string suffix; // e.g., "" (marker after call ID value, before next section) -}; - -// ============================================================================ -// Parser build context (shared interface for build_parser methods) -// ============================================================================ - -struct analyze_content; -struct analyze_reasoning; - -struct parser_build_context { - common_chat_peg_builder & p; - const generation_params & inputs; - common_peg_parser reasoning_parser; - bool extracting_reasoning = false; - const analyze_reasoning * reasoning = nullptr; - const analyze_content * content = nullptr; - - parser_build_context(common_chat_peg_builder & p, const generation_params & inputs); -}; - -// ============================================================================ -// Base class for analyzers with parser building -// ============================================================================ - -struct analyze_base { - virtual ~analyze_base() = default; - virtual common_peg_parser build_parser(parser_build_context & ctx) const = 0; - - protected: - const common_chat_template * tmpl = nullptr; - - analyze_base() = default; - explicit analyze_base(const common_chat_template & tmpl) : tmpl(&tmpl) {} -}; - -// ============================================================================ -// Reasoning analyzer -// ============================================================================ - -struct analyze_reasoning : analyze_base { - reasoning_mode mode = reasoning_mode::NONE; - - std::string start; // e.g., "", "[THINK]", "<|START_THINKING|>", "" - std::string end; // e.g., "", "[BEGIN FINAL RESPONSE]", "<|END_THINKING|>" - - analyze_reasoning() = default; - analyze_reasoning(const common_chat_template & tmpl, bool supports_tools); - analyze_reasoning(std::string start_, std::string end_) : start(std::move(start_)), end(std::move(end_)) {} - - common_peg_parser build_parser(parser_build_context & ctx) const override; - - private: - // Look for reasoning markers in rendered content - void compare_reasoning_presence(); - - // Compare generation prompt with enable_thinking=true vs false - void compare_thinking_enabled(); - - // Check if reasoning is always possible or only in tool calls - void compare_reasoning_scope(); -}; - -// ============================================================================ -// Content analyzer -// ============================================================================ - -struct analyze_content : analyze_base { - content_mode mode = content_mode::PLAIN; - - std::string start; // e.g., "", ">>>all\n", "" - std::string end; // e.g., "", "" - - bool requires_nonnull_content = false; - - analyze_content() = default; - analyze_content(const common_chat_template & tmpl, const analyze_reasoning & reasoning); - - common_peg_parser build_parser(parser_build_context & ctx) const override; - - bool is_always_wrapped() const; - common_peg_parser build_optional_wrapped(parser_build_context & ctx) const; -}; - -// ============================================================================ -// Tool analyzer -// ============================================================================ - -struct analyze_tools : analyze_base { - tool_format_analysis format; - tool_function_analysis function; - tool_arguments_analysis arguments; - tool_id_analysis call_id; - - analyze_tools() = default; - analyze_tools(const common_chat_template & tmpl, - const jinja::caps & caps, - const analyze_reasoning & reasoning); - - common_peg_parser build_parser(parser_build_context & ctx) const override; - - private: - // Extract tool calling 'haystack' for further analysis and delegate further analysis based on format - void analyze_tool_calls(const analyze_reasoning & reasoning, bool supports_parallel_tool_calls); - - // Analyze format based on position of function and argument name in needle - void analyze_tool_call_format(const std::string & haystack, - const std::string & fun_name_needle, - const std::string & arg_name_needle, - const analyze_reasoning & reasoning, - bool supports_parallel_tool_calls); - - // Analyze specifics of JSON native format (entire tool call is a JSON object) - void analyze_tool_call_format_json_native(const std::string & clean_haystack, - const std::string & fun_name_needle, - const std::string & arg_name_needle); - - // Check if parallel calls in JSON native format array wrapped or tag wrapped - void analyze_json_native_parallel_calls(); - - // Analyze specifics of non-JSON native format (tags for function name or for function name and arguments) - void analyze_tool_call_format_non_json(const std::string & clean_haystack, - const std::string & fun_name_needle); - - // Check for and extract specific per-call markers for non-native-JSON templates with parallel call support - void check_per_call_markers(); - - // Extract function name markers - void extract_function_markers(); - - // Delegates to separate functions for: separator analysis, argument name analysis, argument value analysis - void analyze_arguments(); - - // Extract argument name markers - void extract_argument_name_markers(); - - // Extract argument value markers - void extract_argument_value_markers(); - - // Extract argument separator, if specified (eg. ......) - void extract_argument_separator(); - - // Extract argument wrapper markers, if present (eg. '......') - void extract_args_markers(); - - // Extract call ID markers, if present - void extract_call_id_markers(); - - // Per-format tool parser builders - common_peg_parser build_tool_parser_json_native(parser_build_context & ctx) const; - common_peg_parser build_tool_parser_tag_json(parser_build_context & ctx) const; - common_peg_parser build_tool_parser_tag_tagged(parser_build_context & ctx) const; - - // Shared helper: builds func_parser from open+call_id+args, handling atomic wrapping and close. - // atomic_peek: if present, used as the peek expression in the third atomicity branch. - common_peg_parser build_func_parser(common_chat_peg_builder & p, const std::string & name, - const common_peg_parser & call_id_section, bool have_call_id, - const common_peg_parser & args, - std::optional atomic_peek) const; -}; - -// ============================================================================ -// Main autoparser class -// ============================================================================ - -struct autoparser { - jinja::caps jinja_caps; - analyze_reasoning reasoning; - analyze_content content; - analyze_tools tools; - bool analysis_complete = false; - - // Preserved tokens for tokenizer (union of all non-empty markers) - std::vector preserved_tokens; - - autoparser() = default; - - // Run full differential analysis on a template - void analyze_template(const common_chat_template & tmpl); - - // Build the PEG parser for this template - common_peg_arena build_parser(const generation_params & inputs) const; - - private: - // Collect tokens from entire analysis to preserve - void collect_preserved_tokens(); -}; - -// ============================================================================ -// Parser generator -// ============================================================================ - -class peg_generator { - public: - static common_chat_params generate_parser(const common_chat_template & tmpl, - const struct generation_params & inputs); - - static common_chat_params generate_parser(const common_chat_template & tmpl, - const struct generation_params & inputs, - const autoparser & autoparser); -}; - -} // namespace autoparser - -enum segment_type { TEXT, MARKER }; - -inline std::ostream & operator<<(std::ostream & os, const segment_type & type) { - switch (type) { - case segment_type::TEXT: - return os << "TEXT"; - case segment_type::MARKER: - return os << "MARKER"; - default: - return os << "UNKNOWN"; - } -} - -struct segment { - segment_type type; - std::string value; - - segment(segment_type type, std::string value) : type(type), value(std::move(value)) {} - - bool operator==(const segment & other) const { - return type == other.type && value == other.value; - } - - bool operator!=(const segment & other) const { - return !(*this == other); - } -}; diff --git a/common/chat-diff-analyzer.cpp b/common/chat-diff-analyzer.cpp deleted file mode 100644 index 264ace4627c9..000000000000 --- a/common/chat-diff-analyzer.cpp +++ /dev/null @@ -1,1395 +0,0 @@ -#include "chat-auto-parser.h" -#include "chat-auto-parser-helpers.h" -#include "chat-peg-parser.h" -#include "chat.h" -#include "common.h" -#include "log.h" -#include "nlohmann/json.hpp" -#include "peg-parser.h" - -#include - -#define ANSI_RESET "\033[0m" -#define ANSI_PURPLE "\033[1m\x1b[38;5;126m" -#define ANSI_ORANGE "\033[1m\x1b[38;5;214m" -#define ANSI_RED "\033[1m\x1b[38;5;196m" - -using json = nlohmann::ordered_json; - -namespace autoparser { - -static const std::string FUN_FIRST = "FFF_FIRST_FUN_F"; -static const std::string FUN_SECOND = "SSS_SECOND_FUN_S"; -static const std::string ARG_FIRST = "AA_ARG_FST_AA"; -static const std::string ARG_SECOND = "BB_ARG_SND_BB"; -static const std::string USER_MSG = "U_USER_MSG Hello END_U"; -static const std::string ASSISTANT_MSG = "A_ASST_MSG I can help END_A"; -static const std::string THINKING_CONTENT = "REASON_PART I am thinking END_R"; -static const std::string CALL_ID_001 = "call00001"; -static const std::string CALL_ID_002 = "call00002"; -static const std::string CALL_ID_999 = "call99999"; - -static std::vector> workarounds( - { // Old reasoning Qwen templates - they don't really display reasoning content, but we still want to - // support reasoning on them - [](const common_chat_template & tmpl, autoparser & analysis) -> void { - if (tmpl.src.find("content.split('')") != std::string::npos && - tmpl.src.find("reasoning_content") == std::string::npos && - tmpl.src.find("") == std::string::npos && - analysis.reasoning.mode == reasoning_mode::NONE) { - analysis.reasoning.mode = reasoning_mode::TAG_BASED; - analysis.reasoning.start = ""; - analysis.reasoning.end = ""; - analysis.preserved_tokens.push_back(""); - analysis.preserved_tokens.push_back(""); - LOG_DBG(ANSI_ORANGE "[Patch: old Qwen/Deepseek thinking template]\n" ANSI_RESET); - } - }, - // Granite 3.3, with separate reasoning and content markers - [](const common_chat_template & tmpl, autoparser & analysis) -> void { - if (tmpl.src.find("Write your thoughts between and write your response between " - "") != std::string::npos) { - analysis.reasoning.mode = reasoning_mode::TAG_BASED; - analysis.reasoning.start = ""; - analysis.reasoning.end = ""; - analysis.preserved_tokens.push_back(""); - analysis.preserved_tokens.push_back(""); - analysis.content.mode = content_mode::WRAPPED_WITH_REASONING; - analysis.content.start = ""; - analysis.content.end = ""; - analysis.preserved_tokens.push_back(""); - analysis.preserved_tokens.push_back(""); - LOG_DBG(ANSI_ORANGE "[Patch: Granite 3.3]\n" ANSI_RESET); - } - }, - // Cohere Command R+ - content wrapped in <|CHATBOT_TOKEN|>...<|END_OF_TURN_TOKEN|> - [](const common_chat_template & tmpl, autoparser & analysis) -> void { - if (tmpl.src.find("<|CHATBOT_TOKEN|>") != std::string::npos && - tmpl.src.find("<|END_OF_TURN_TOKEN|>") != std::string::npos && analysis.content.start.empty()) { - analysis.content.mode = content_mode::ALWAYS_WRAPPED; - analysis.content.start = "<|CHATBOT_TOKEN|>"; - analysis.content.end = "<|END_OF_TURN_TOKEN|>"; - analysis.preserved_tokens.push_back("<|CHATBOT_TOKEN|>"); - analysis.preserved_tokens.push_back("<|END_OF_TURN_TOKEN|>"); - LOG_DBG(ANSI_ORANGE "[Patch: Cohere Command R+]\n" ANSI_RESET); - } - }, - // Functionary - no tool call section delimiter - [](const common_chat_template & tmpl, autoparser & analysis) -> void { - if (tmpl.src.find("set has_code_interpreter = tools | selectattr(\"type\", \"equalto\", " - "\"code_interpreter\") | list | length > 0") != std::string::npos) { - analysis.content.mode = content_mode::PLAIN; - analysis.content.end = ""; - analysis.tools.function.name_prefix = ""; - analysis.tools.format.section_start = ""; - analysis.tools.format.section_end = ""; - analysis.tools.format.per_call_start = ""); - analysis.preserved_tokens.push_back("<|eom_id|>"); - analysis.preserved_tokens.push_back(""); - analysis.preserved_tokens.push_back(""); - LOG_DBG(ANSI_ORANGE "[Patch: Functionary 3.1]\n" ANSI_RESET); - } - }, - // DeepSeek-R1-Distill-Qwen - [](const common_chat_template & tmpl, autoparser & analysis) -> void { - if (tmpl.src.find( - "{{'<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>'") != - std::string::npos) { - analysis.tools.format.section_start = "<｜tool▁calls▁begin｜>"; - analysis.tools.format.section_end = "<｜tool▁calls▁end｜>"; - analysis.tools.format.per_call_start = "<｜tool▁call▁begin｜>function"; - analysis.tools.function.name_prefix = "<｜tool▁sep｜>"; - analysis.tools.format.per_call_end = "<｜tool▁call▁end｜>"; - analysis.tools.function.close = "```"; - LOG_DBG(ANSI_ORANGE "[Patch: DeepSeek-R1-Distill-Qwen]\n" ANSI_RESET); - } - } - }); - -// Common JSON structures -static json params_schema = { - { "type", "object" }, - { "properties", - { { ARG_FIRST, { { "type", "string" }, { "description", "First argument" } } }, - { ARG_SECOND, { { "type", "string" }, { "description", "Second argument" } } } } }, - { "required", json::array({}) } -}; - -static json tools = json::array({ - { { "type", "function" }, - { "function", - json{ { "name", FUN_FIRST }, { "description", "Test function foo" }, { "parameters", params_schema } } } }, - { { "type", "function" }, - { "function", - json{ { "name", FUN_SECOND }, { "description", "Test function bar" }, { "parameters", params_schema } } } } -}); - -static json user_msg = json{ - { "role", "user" }, - { "content", USER_MSG } -}; - -static json build_tool_call(const std::string & name, const json & args, const std::string & id = CALL_ID_001) { - return json{ - { "id", id }, - { "type", "function" }, - { "function", json{ { "name", name }, { "arguments", args } } } - }; -} - -static json first_tool_call_zero_args = build_tool_call(FUN_FIRST, json::object(), CALL_ID_001); -static json first_tool_call_one_arg = build_tool_call(FUN_FIRST, {{ ARG_FIRST, "XXXX" }}, CALL_ID_001); -static json first_tool_call_one_arg_other_val = build_tool_call(FUN_FIRST, {{ ARG_FIRST, "YYYY" }}, CALL_ID_001); -static json first_tool_call_other_arg = build_tool_call(FUN_FIRST, {{ ARG_SECOND, "YYYY" }}, CALL_ID_001); - -static json first_tool_call = - build_tool_call(FUN_FIRST, json{{ ARG_FIRST, "XXXX" }, { ARG_SECOND, "YYYY" }}, CALL_ID_001); -static json second_tool_call = - build_tool_call(FUN_SECOND, json{ { ARG_FIRST, "XXXX" }, { ARG_SECOND, "YYYY" }}, CALL_ID_002); -static json first_tool_call_alt_id = - build_tool_call(FUN_FIRST, json{{ ARG_FIRST, "XXXX" }, { ARG_SECOND, "YYYY" }}, CALL_ID_999); - -template -static std::string mode_to_str(T mode) { - std::ostringstream os; - os << mode; - return os.str(); -} - -void autoparser::analyze_template(const common_chat_template & tmpl) { - jinja_caps = tmpl.original_caps(); - reasoning = analyze_reasoning(tmpl, jinja_caps.supports_tool_calls); - content = analyze_content(tmpl, reasoning); - tools = analyze_tools(jinja_caps.supports_tool_calls ? analyze_tools(tmpl, jinja_caps, reasoning) : analyze_tools()); - collect_preserved_tokens(); - - for (auto & workaround : workarounds) { - workaround(tmpl, *this); - } - - LOG_DBG("\n--- Reasoning & Content Structure ---\n"); - LOG_DBG("reasoning_mode: %s\n", mode_to_str(reasoning.mode).c_str()); - LOG_DBG("reasoning_start: '%s'\n", reasoning.start.c_str()); - LOG_DBG("reasoning_end: '%s'\n", reasoning.end.c_str()); - LOG_DBG("content_mode: %s\n", mode_to_str(content.mode).c_str()); - LOG_DBG("content_start: '%s'\n", content.start.c_str()); - LOG_DBG("content_end: '%s'\n", content.end.c_str()); - - LOG_DBG("\n--- Tool Call Structure ---\n"); - LOG_DBG("tool_mode: %s\n", mode_to_str(tools.format.mode).c_str()); - LOG_DBG("supports_tools: %s\n", jinja_caps.supports_tools ? "true" : "false"); - LOG_DBG("supports_parallel_calls: %s\n", jinja_caps.supports_parallel_tool_calls ? "true" : "false"); - LOG_DBG("tool_section_start: '%s'\n", tools.format.section_start.c_str()); - LOG_DBG("tool_section_end: '%s'\n", tools.format.section_end.c_str()); - LOG_DBG("per_call_start: '%s'\n", tools.format.per_call_start.c_str()); - LOG_DBG("per_call_end: '%s'\n", tools.format.per_call_end.c_str()); - LOG_DBG("func_name_prefix: '%s'\n", tools.function.name_prefix.c_str()); - LOG_DBG("func_name_suffix: '%s'\n", tools.function.name_suffix.c_str()); - LOG_DBG("func_close: '%s'\n", tools.function.close.c_str()); - LOG_DBG("call_id_prefix: '%s'\n", tools.call_id.prefix.c_str()); - LOG_DBG("call_id_suffix: '%s'\n", tools.call_id.suffix.c_str()); - LOG_DBG("call_id_pos: '%s'\n", mode_to_str(tools.call_id.pos).c_str()); - LOG_DBG("args_start: '%s'\n", tools.arguments.start.c_str()); - LOG_DBG("args_end: '%s'\n", tools.arguments.end.c_str()); - LOG_DBG("arg_name_prefix: '%s'\n", tools.arguments.name_prefix.c_str()); - LOG_DBG("arg_name_suffix: '%s'\n", tools.arguments.name_suffix.c_str()); - LOG_DBG("arg_value_prefix: '%s'\n", tools.arguments.value_prefix.c_str()); - LOG_DBG("arg_value_suffix: '%s'\n", tools.arguments.value_suffix.c_str()); - LOG_DBG("name_field: '%s'\n", tools.format.name_field.c_str()); - LOG_DBG("args_field: '%s'\n", tools.format.args_field.c_str()); - LOG_DBG("id_field: '%s'\n", tools.format.id_field.c_str()); - LOG_DBG("gen_id_field: '%s'\n", tools.format.gen_id_field.c_str()); - LOG_DBG("parameter_order: '%s'\n", std::accumulate(tools.format.parameter_order.begin(), tools.format.parameter_order.end(), - std::string(""), [] (const std::string & a, const std::string & b) { return a.empty() ? b : a + ", " + b; } - ).c_str()); - - LOG_DBG(ANSI_PURPLE "=== Differential analysis complete ===\n" ANSI_RESET); - analysis_complete = true; -} - -void autoparser::collect_preserved_tokens() { - auto add_token = [this](const std::string & org_token) { - std::string token = trim_whitespace(org_token); - if (!token.empty()) { - // Avoid duplicates - if (std::find(preserved_tokens.begin(), preserved_tokens.end(), token) == preserved_tokens.end()) { - preserved_tokens.push_back(token); - } - } - }; - - add_token(reasoning.start); - add_token(reasoning.end); - add_token(content.start); - add_token(content.end); - add_token(tools.format.section_start); - add_token(tools.format.section_end); - add_token(tools.format.per_call_start); - add_token(tools.format.per_call_end); - add_token(tools.function.name_prefix); - add_token(tools.function.name_suffix); - add_token(tools.function.close); - add_token(tools.arguments.start); - add_token(tools.arguments.end); - add_token(tools.arguments.name_prefix); - add_token(tools.arguments.name_suffix); - add_token(tools.arguments.separator); - add_token(tools.arguments.value_prefix); - add_token(tools.arguments.value_suffix); - add_token(tools.call_id.prefix); - add_token(tools.call_id.suffix); -} - -analyze_reasoning::analyze_reasoning(const common_chat_template & tmpl, bool supports_tools) - : analyze_base(tmpl) { - LOG_DBG(ANSI_PURPLE "=== Starting differential analysis ===\n" ANSI_RESET); - LOG_DBG(ANSI_ORANGE "Phase 1: Reasoning analysis\n" ANSI_RESET); - - compare_reasoning_presence(); - compare_thinking_enabled(); - if (supports_tools) { - compare_reasoning_scope(); - } -} - -void analyze_reasoning::compare_reasoning_presence() { - json user_msg = json{ - { "role", "user" }, - { "content", USER_MSG } - }; - - json assistant_no_reasoning = json{ - { "role", "assistant" }, - { "content", ASSISTANT_MSG } - }; - - json assistant_with_reasoning = json{ - { "role", "assistant" }, - { "content", ASSISTANT_MSG }, - { "reasoning_content", THINKING_CONTENT } - }; - - template_params params; - params.messages = json::array({ user_msg, assistant_no_reasoning }); - params.add_generation_prompt = false; - params.enable_thinking = true; - - auto comparison = compare_variants( - *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_with_reasoning }); }); - - if (!comparison) { - LOG_DBG(ANSI_ORANGE "%s: Template application failed, skipping reasoning detection\n" ANSI_RESET, __func__); - return; - } - - const auto & diff = comparison->diff; - - const std::string reasoning_content = THINKING_CONTENT; - - if (!diff.right.empty() && diff.right.find(reasoning_content) != std::string::npos) { - auto parser_delimiter = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.literal(reasoning_content) + p.space() + p.optional(p.tag("post", (p.marker() + p.space())) + p.rest()); - }); - auto parser_wrapped = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.tag("pre", p.marker() + p.space()) + p.literal(reasoning_content) + p.tag("post", (p.space() + p.marker() + p.space())) + p.rest(); - }); - // try the more aggressive parse first, if it fails, fall back to the delimiter one - auto result = parser_wrapped.parse_anywhere_and_extract(comparison->output_B); - if (!result.result.success()) { - result = parser_delimiter.parse_anywhere_and_extract(comparison->output_B); - } - if (result.result.success()) { - if (!result.tags["pre"].empty() && !result.tags["post"].empty()) { - mode = reasoning_mode::TAG_BASED; - start = result.tags["pre"]; - end = result.tags["post"]; - } else if (!result.tags["post"].empty()) { - mode = reasoning_mode::TAG_BASED; - end = result.tags["post"]; - } - } - } -} - -void analyze_reasoning::compare_thinking_enabled() { - json user_msg = json{ - { "role", "user" }, - { "content", USER_MSG } - }; - - template_params params; - params.messages = json::array({ user_msg }); - params.add_generation_prompt = true; - params.enable_thinking = false; - - auto comparison = compare_variants(*tmpl, params, [&](template_params & p) { p.enable_thinking = true; }); - - if (!comparison) { - LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET , __func__); - return; - } - - const auto & diff = comparison->diff; - - std::string left_trimmed = trim_whitespace(diff.left); - std::string right_trimmed = trim_whitespace(diff.right); - - if (left_trimmed.empty() && !diff.right.empty()) { - if (!right_trimmed.empty() && string_ends_with(comparison->output_B, right_trimmed)) { - if (start.empty()) { - start = trim_leading_whitespace(diff.right); - mode = reasoning_mode::TAG_BASED; - } - } - } else if (right_trimmed.empty() && !diff.left.empty()) { - if (!left_trimmed.empty() && string_ends_with(comparison->output_A, left_trimmed)) { - if (end.empty()) { - auto seg = prune_whitespace_segments(segmentize_markers(comparison->output_A)); - if (seg.size() >= 2 && seg[seg.size() - 1].value == left_trimmed && seg[seg.size() - 2].type == segment_type::MARKER) { - start = seg[seg.size() - 2].value; - } - end = trim_trailing_whitespace(diff.left); - mode = reasoning_mode::TAG_BASED; - } - } - } else if (!left_trimmed.empty() && !right_trimmed.empty()) { - // Full-output diff is noisy (e.g., SmolLM3 changes the system message when enable_thinking flips). - // Try to find reasoning markers by tail-anchoring: - // one output's generation prompt tail may appear in the other with extra reasoning markers appended. - const auto & output_A = comparison->output_A; - const auto & output_B = comparison->output_B; - const size_t anchor_len = 64; - - for (int dir = 0; dir < 2; dir++) { - const auto & base = dir == 0 ? output_B : output_A; - const auto & extended = dir == 0 ? output_A : output_B; - - size_t len = std::min(base.size(), anchor_len); - std::string anchor = base.substr(base.size() - len); - auto pos = extended.rfind(anchor); - if (pos == std::string::npos || pos + len >= extended.size()) { - continue; - } - - std::string extra = trim_whitespace(extended.substr(pos + len)); - if (extra.empty()) { - continue; - } - - auto seg = prune_whitespace_segments(segmentize_markers(extra)); - if (seg.size() == 2 && seg[0].type == segment_type::MARKER && seg[1].type == segment_type::MARKER) { - if (start.empty()) { - start = seg[0].value; - } - if (end.empty()) { - end = seg[1].value; - } - mode = reasoning_mode::TAG_BASED; - break; - } - } - } - - if (mode == reasoning_mode::NONE && start.empty() && !end.empty()) { - mode = reasoning_mode::TAG_BASED; - } -} - -void analyze_reasoning::compare_reasoning_scope() { - json assistant_reasoning_content = json{ - { "role", "assistant" }, - { "content", ASSISTANT_MSG }, - { "reasoning_content", THINKING_CONTENT } - }; - - json assistant_reasoning_tools = json{ - { "role", "assistant" }, - { "content", nullptr }, - { "reasoning_content", THINKING_CONTENT }, - { "tool_calls", - json::array({ build_tool_call(FUN_FIRST, json{ { ARG_FIRST, "VVVV" }, { ARG_SECOND, "XXXX" } }) }) } - }; - - template_params params; - params.messages = json::array({ user_msg, assistant_reasoning_content }); - params.tools = tools; - params.add_generation_prompt = false; - params.enable_thinking = true; - - auto comparison = compare_variants( - *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_reasoning_tools }); }); - - if (!comparison) { - LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); - return; - } - - std::string reasoning_content = THINKING_CONTENT; - - // Check if reasoning only appears in variant B (with tools) - bool reasoning_in_A = comparison->output_A.find(reasoning_content) != std::string::npos; - bool reasoning_in_B = comparison->output_B.find(reasoning_content) != std::string::npos; - - if (!reasoning_in_A && reasoning_in_B) { - mode = reasoning_mode::TOOLS_ONLY; - LOG_DBG(ANSI_ORANGE "%s: Detected TOOLS_ONLY reasoning mode\n" ANSI_RESET, __func__); - - auto parser_wrapped = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.tag("pre", p.marker() + p.space()) + p.literal(reasoning_content) + p.space() + p.tag("post", (p.marker() + p.space())); - }); - auto result = parser_wrapped.parse_anywhere_and_extract(comparison->output_B); - if (result.result.success()) { - start = result.tags["pre"]; - end = trim_trailing_whitespace(result.tags["post"]); - } else { - auto parser_delimiter = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.literal(reasoning_content) + p.space() + p.optional(p.tag("post", (p.marker() + p.space()))); - }); - result = parser_delimiter.parse_anywhere_and_extract(comparison->output_B); - if (result.result.success()) { - end = trim_trailing_whitespace(result.tags["post"]); - } else { - LOG_DBG(ANSI_ORANGE "%s: Unable to extract reasoning markers, falling back to reasoning = NONE\n" ANSI_RESET, __func__); - mode = reasoning_mode::NONE; - } - } - } -} - -analyze_content::analyze_content(const common_chat_template & tmpl, const analyze_reasoning & reasoning) - : analyze_base(tmpl) { - LOG_DBG(ANSI_ORANGE "Phase 2: Content analysis\n" ANSI_RESET); - - json assistant_content_only = json{ - { "role", "assistant" }, - { "content", ASSISTANT_MSG } - }; - - json assistant_with_tools = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ build_tool_call("test_func", json{ { "arg1", "value1" } }) }) } - }; - - json assistant_with_reasoning = json{ - { "role", "assistant" }, - { "content", "" }, - { "reasoning_content", THINKING_CONTENT } - }; - - template_params params_content_only; - params_content_only.messages = json::array({ user_msg, assistant_content_only }); - params_content_only.add_generation_prompt = false; - params_content_only.enable_thinking = true; - params_content_only.tools = tools; - - auto comparison_with_tools = compare_variants(tmpl, params_content_only, [&](template_params & p) { - p.messages = json::array({ user_msg, assistant_with_tools }); - }); - - auto comparison_with_reasoning = compare_variants(tmpl, params_content_only, [&](template_params & p) { - p.messages = json::array({ user_msg, assistant_with_reasoning }); - }); - - if (!comparison_with_tools || !comparison_with_reasoning) { - LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); - return; - } - - const auto & diff_tools = comparison_with_tools->diff; - const auto & diff_reasoning = comparison_with_reasoning->diff; - - std::string response = ASSISTANT_MSG; - - bool found_plain_content = false; - if (trim_whitespace(diff_tools.left) == response) { - auto parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { - return p.space() + diff_reasoning.left + p.space() + p.optional(p.marker()) + p.space() + p.end(); - }); - if (parser.parse_and_extract(diff_reasoning.left).result.success()) { - // We only have the content text in the diff (possibly with a stray EOG marker), so no markers - mode = content_mode::PLAIN; - found_plain_content = true; - } else if (reasoning.mode != reasoning_mode::NONE && !reasoning.end.empty()) { - auto post_reasoning_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { - return p.literal(reasoning.end) + p.space() + p.literal(response); - }); - if (post_reasoning_parser.parse_anywhere_and_extract(diff_reasoning.left).result.success()) { - mode = content_mode::PLAIN; - found_plain_content = true; - } - } - } - if (!found_plain_content) { - std::string rdiff = diff_reasoning.left; - if (!reasoning.end.empty() && rdiff.find(reasoning.end) != std::string::npos) { - rdiff = rdiff.substr(rdiff.find(reasoning.end) + reasoning.end.length()); - } - // Take the more promising diff - std::string pure_content = rdiff.length() > diff_tools.left.length() ? rdiff : diff_tools.left; - auto parser_wrapped = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.tag("pre", p.marker() + p.space()) + p.literal(response) + p.space() + p.tag("post", (p.marker() + p.space())) + p.rest(); - }); - auto result = parser_wrapped.parse_anywhere_and_extract(pure_content); - start = result.tags["pre"]; - end = result.tags["post"]; - // TODO: WRAPPED_WITH_REASONING - } - - // Determine content mode - if (!start.empty() || !end.empty()) { - mode = content_mode::ALWAYS_WRAPPED; - // TODO: END_DELIMITED content mode - delimited at end but not at start? - } -} - -bool analyze_content::is_always_wrapped() const { - return mode == content_mode::ALWAYS_WRAPPED && !start.empty() && !end.empty(); -} - -analyze_tools::analyze_tools(const common_chat_template & tmpl, - const jinja::caps & caps, - const analyze_reasoning & reasoning) - : analyze_base(tmpl) { - LOG_DBG(ANSI_ORANGE "Phase 3: Tool call analysis\n" ANSI_RESET); - - analyze_tool_calls(reasoning, caps.supports_parallel_tool_calls); - - if (format.mode != tool_format::NONE && format.mode != tool_format::JSON_NATIVE) { - if (caps.supports_parallel_tool_calls) { - check_per_call_markers(); - } - LOG_DBG(ANSI_ORANGE "Phase 3a: Function call analysis\n" ANSI_RESET); - extract_function_markers(); - LOG_DBG(ANSI_ORANGE "Phase 3b: Argument analysis\n" ANSI_RESET); - if (format.mode == tool_format::TAG_WITH_TAGGED) { - analyze_arguments(); - } - extract_argument_separator(); - extract_args_markers(); - LOG_DBG(ANSI_ORANGE "Phase 3c: Call id analysis\n" ANSI_RESET); - extract_call_id_markers(); - } -} - -void analyze_tools::analyze_tool_calls(const analyze_reasoning & reasoning, bool supports_parallel_tool_calls) { - json assistant_no_tools = json{ - { "role", "assistant" }, - { "content", ASSISTANT_MSG } - }; - - json assistant_with_tools = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call }) } - }; - - template_params params; - params.messages = json::array({ user_msg, assistant_no_tools }); - params.tools = tools; - params.add_generation_prompt = false; - params.enable_thinking = true; - - auto comparison = compare_variants( - *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_with_tools }); }); - - if (!comparison) { - LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); - return; - } - - const auto & diff = comparison->diff; - - std::string tool_section = diff.right; - - if (tool_section.empty()) { - return; - } - - analyze_tool_call_format(tool_section, FUN_FIRST, ARG_FIRST, reasoning, supports_parallel_tool_calls); -} - -void analyze_tools::analyze_tool_call_format(const std::string & haystack, - const std::string & fun_name_needle, - const std::string & arg_name_needle, - const analyze_reasoning & reasoning, - bool supports_parallel_tool_calls) { - if (fun_name_needle.empty() || arg_name_needle.empty() || haystack.empty()) { - return; - } - - auto in_json_haystack = [&haystack](const std::string & needle) -> bool { - auto parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.choice({ p.literal("{"), p.literal(":") }) << p.choice({ - p.tag("dq", p.literal("\"") + p.literal(needle) + p.literal("\"")) }); - }); - auto result = parser.parse_anywhere_and_extract(haystack); - return result.result.success(); - }; - - auto fun_quote = in_json_haystack(fun_name_needle); - auto arg_quote = in_json_haystack(arg_name_needle); - - if (fun_quote) { - // no need to check further, we're in JSON land - format.mode = tool_format::JSON_NATIVE; - } else if (arg_quote) { - format.mode = tool_format::TAG_WITH_JSON; - } else { - format.mode = tool_format::TAG_WITH_TAGGED; - } - - // first, remove any reasoning markers - std::string clean_haystack = haystack; - if (!reasoning.start.empty()) { - auto pos = haystack.find(reasoning.start); - if (pos != std::string::npos) { - clean_haystack = haystack.substr(0, pos) + haystack.substr(pos + reasoning.start.length()); - } - } - if (!reasoning.end.empty()) { - auto pos = clean_haystack.find(reasoning.end); - if (pos != std::string::npos) { - clean_haystack = clean_haystack.substr(0, pos) + clean_haystack.substr(pos + reasoning.end.length()); - } - } - - if (format.mode == tool_format::JSON_NATIVE) { - analyze_tool_call_format_json_native(clean_haystack, fun_name_needle, arg_name_needle); - if (supports_parallel_tool_calls) { - analyze_json_native_parallel_calls(); - } - } else { - analyze_tool_call_format_non_json(clean_haystack, fun_name_needle); - } - // always relax whitespace requirements on ending markers since they don't influence content - format.section_end = trim_whitespace(format.section_end); - format.per_call_end = trim_whitespace(format.per_call_end); -} - -void analyze_tools::analyze_json_native_parallel_calls() { - json assistant_one_tool = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call }) } - }; - - json assistant_two_tools = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call, second_tool_call }) } - }; - - template_params params; - params.messages = json::array({ user_msg, assistant_one_tool }); - params.tools = tools; - params.add_generation_prompt = false; - params.enable_thinking = true; - - auto comparison = compare_variants( - *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_two_tools }); }); - - if (!comparison) { - LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); - return; - } - - std::string & second_call = comparison->diff.right; - if (!format.section_start.empty() && second_call.find(format.section_start) != std::string::npos) { - format.per_call_start = format.section_start; - format.per_call_end = format.section_end; - format.section_start.clear(); - format.section_end.clear(); - } -} - -void analyze_tools::analyze_tool_call_format_json_native(const std::string & clean_haystack, - const std::string & fun_name_needle, - const std::string & arg_name_needle) { - // we might not have the typical OpenAI tool calling structure - int json_start = clean_haystack.find_first_of('{'); - int json_end = clean_haystack.find_last_of('}'); - std::string cut = clean_haystack.substr(json_start, json_end - json_start + 1); - json call_struct = json::parse(cut); - auto register_field = [&](const std::string & prefix, const nlohmann::detail::iteration_proxy_value & subel) { - if (subel.value().is_string() && std::string(subel.value()).find("call0000") != std::string::npos) { - format.id_field = !prefix.empty() ? prefix + "." + subel.key() : subel.key(); - } else if (subel.value().is_string() && std::string(subel.value()) == fun_name_needle) { - format.name_field = !prefix.empty() ? prefix + "." + subel.key() : subel.key(); - } else if (subel.value().dump().find(arg_name_needle) != - std::string::npos) { // handle both string and JSON obj variants - format.args_field = !prefix.empty() ? prefix + "." + subel.key() : subel.key(); - } else if (subel.key().find("id") != std::string::npos) { - // heuristics for generated id field - format.gen_id_field = !prefix.empty() ? prefix + "." + subel.key() : subel.key(); - } - }; - for (const auto & el : call_struct.items()) { - if (el.key() == fun_name_needle) { - format.fun_name_is_key = true; - // When function name is the key, there's no name field and args are direct - format.name_field.clear(); - format.args_field.clear(); - // Don't register this element - the function name IS the key, not a field - } else { - if (el.value().is_object() && - el.value().dump().find(arg_name_needle) == std::string::npos) { // not the args object - format.function_field = el.key(); - for (const auto & subel : el.value().items()) { - register_field(el.key(), subel); - } - } - // Register this element as a potential field - register_field("", el); - } - } - auto array_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.tag("pre", p.literal("[") + p.space()) + p.literal(cut) + p.tag("post", p.space() + p.literal("]")); - }); - - auto ar_parse_res = array_parser.parse_anywhere_and_extract(clean_haystack); - if (ar_parse_res.result.success()) { - format.tools_array_wrapped = true; - json_start -= ar_parse_res.tags["pre"].length(); - json_end += ar_parse_res.tags["post"].length(); - } - json_end++; // we want to move past the closing char for end marker extraction - - std::vector> located_params; - if (!format.name_field.empty()) { - located_params.push_back({ clean_haystack.find(format.name_field), format.name_field }); - } - if (!format.args_field.empty()) { - located_params.push_back({ clean_haystack.find(format.args_field), format.args_field }); - } - if (!format.id_field.empty()) { - located_params.push_back({ clean_haystack.find(format.id_field), format.id_field }); - } - if (!format.gen_id_field.empty()) { - located_params.push_back({ clean_haystack.find(format.gen_id_field), format.gen_id_field }); - } - std::sort(located_params.begin(), located_params.end()); - for (auto & pair : located_params) { - format.parameter_order.push_back(pair.second); - } - // we can immediately extract tool calling markers too - format.section_start = trim_leading_whitespace(clean_haystack.substr(0, json_start)); - format.section_end = trim_whitespace(clean_haystack.substr(json_end)); - // When tools_array_wrapped is true, the closing bracket is part of the array structure, - // not a separate section end marker. Clear tool_section_end to avoid duplicate brackets. - if (format.tools_array_wrapped && format.section_end == "]") { - format.section_end.clear(); - } -} - -void analyze_tools::analyze_tool_call_format_non_json(const std::string & clean_haystack, - const std::string & fun_name_needle) { - // first, let's find out if the function is inside a tag or standalone - auto fun_marker_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.tag("fun_marker", p.choice({ - p.tag("fun_pre", p.literal("<") + p.until_one_of({ ">", fun_name_needle })) + p.literal(fun_name_needle) + - p.tag("fun_post", p.negate(p.space() + p.literal("<")) + p.until(">") + p.literal(">")) + p.space(), - p.tag("fun_pre", p.literal("[") + p.until_one_of({ "]", fun_name_needle })) + p.literal(fun_name_needle) + - p.tag("fun_post", p.negate(p.space() + p.literal("[") + p.until("]") + p.literal("]")) + p.space()) })); - }); - auto fun_res = fun_marker_parser.parse_anywhere_and_extract(clean_haystack); - std::string fun_marker = fun_name_needle; - if (fun_res.result.success()) { - fun_marker = fun_res.tags["fun_marker"]; - } - // now, consume up to two markers, then treat everything up to the function marker as function name prefix - auto per_tool_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.tag("sec_start", p.marker() + p.space()) + p.tag("call_start", p.marker() + p.space()) + - p.tag("fun_pre", p.until(fun_marker)) + fun_marker + p.tag("rest", p.rest()); - }); - auto section_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.tag("sec_start", p.marker() + p.space()) + fun_marker + p.tag("rest", p.rest()); - }); - auto result = per_tool_parser.parse_anywhere_and_extract(clean_haystack); - tagged_parse_result result_end; - if (result.result.success()) { - auto double_closer_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.tag("call_end", p.marker() + p.space()) + p.tag("sec_end", p.marker() + p.space()) + p.end(); - }); - result_end = double_closer_parser.parse_anywhere_and_extract(result.tags["rest"]); - function.name_prefix = fun_res.tags["fun_pre"] + function.name_prefix; - } else { - result = section_parser.parse_anywhere_and_extract(clean_haystack); - auto single_closer_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.tag("sec_end", p.marker() + p.space()) + p.end(); - }); - result_end = single_closer_parser.parse_anywhere_and_extract(result.tags["rest"]); - } - format.per_call_start = result.tags["call_start"]; - format.per_call_end = result_end.tags["call_end"]; - format.section_start = result.tags["sec_start"]; - format.section_end = result_end.tags["sec_end"]; -} - -void analyze_tools::check_per_call_markers() { - json assistant_one_tool = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call }) } - }; - - json assistant_two_tools = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call, second_tool_call }) } - }; - - template_params params; - params.messages = json::array({ user_msg, assistant_one_tool }); - params.tools = tools; - params.add_generation_prompt = false; - params.enable_thinking = true; - - auto one_vs_two = compare_variants( - *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_two_tools }); }); - - if (!one_vs_two) { - LOG_DBG(ANSI_ORANGE "%s: Generating double tool call comparison failed\n" ANSI_RESET, __func__); - return; - } - - diff_split filter_common_call_part = calculate_diff_split(one_vs_two->diff.suffix, one_vs_two->diff.right); - - std::string second_tool_content = trim_leading_whitespace(filter_common_call_part.right); - if (!format.section_start.empty() && - second_tool_content.find(format.section_start) == 0) { - format.per_call_start = format.section_start; - format.per_call_end = format.section_end; - format.section_start.clear(); - format.section_end.clear(); - } -} - -void analyze_tools::extract_function_markers() { - json assistant_nocall = json{ - { "role", "assistant" }, - { "content", ASSISTANT_MSG }, - }; - - json assistant_foofoo = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call }) } - }; - - json assistant_barbar = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ second_tool_call }) } - }; - - template_params params; - params.messages = json::array({ user_msg, assistant_foofoo }); - params.tools = tools; - params.add_generation_prompt = false; - params.enable_thinking = true; - - auto comparison = compare_variants( - *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_barbar }); }); - - if (!comparison) { - LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); - return; - } - - const auto & diff = comparison->diff; - - if (diff.left.find(FUN_FIRST) != std::string::npos && diff.right.find(FUN_SECOND) != std::string::npos) { - std::string prefix_marker; - if (!format.per_call_start.empty()) { - prefix_marker = format.per_call_start; - } else { - prefix_marker = format.section_start; - } - if (!prefix_marker.empty() && diff.prefix.rfind(prefix_marker) != std::string::npos) { - function.name_prefix = - diff.prefix.substr(diff.prefix.rfind(prefix_marker) + prefix_marker.size()); - } - - // Extract name prefix/suffix from diff.left (stop at the next marker boundary) - auto name_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.tag("pre", p.until(FUN_FIRST)) + p.literal(FUN_FIRST) + - p.tag("post", p.zero_or_more(p.negate(p.marker()) + p.any())); - }); - auto name_result = name_parser.parse_and_extract(diff.left); - if (name_result.result.success()) { - function.name_prefix += name_result.tags["pre"]; - function.name_suffix = name_result.tags["post"]; - } - - // Extend name_suffix with content from diff.suffix before args begin - if (format.mode == tool_format::TAG_WITH_JSON) { - // For JSON: name_suffix extends to the first non-marker { or [, including any - // markers along the way. Only applies if there's at least one marker after - // the JSON content (matching the original "stop < seg_suf.size() - 1" guard). - auto suffix_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - auto non_json = p.marker() | (p.negate(p.literal("{")) + p.negate(p.literal("[")) + p.any()); - auto after_json = p.zero_or_more(p.negate(p.marker()) + p.any()) + p.marker(); - return p.tag("ext", p.zero_or_more(non_json)) + after_json; - }); - auto suf_result = suffix_parser.parse_and_extract(diff.suffix); - if (suf_result.result.success()) { - function.name_suffix += suf_result.tags["ext"]; - } - } else { - // For tagged: name_suffix extends to the first marker (arg marker) - auto suffix_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.tag("ext", p.zero_or_more(p.negate(p.marker()) + p.any())); - }); - auto suf_result = suffix_parser.parse_and_extract(diff.suffix); - if (suf_result.result.success()) { - function.name_suffix += suf_result.tags["ext"]; - } - } - - // Extract the closer (between last arg and call/section end marker) - std::string suffix_marker; - if (!format.per_call_end.empty()) { - suffix_marker = format.per_call_end; - } else { - suffix_marker = format.section_end; - } - std::string closer_suffix; - if (suffix_marker.empty()) { - // we'll have to rely on an extra diff with no-calls version - auto notool_comp = compare_variants( - *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_nocall }); }); - if (notool_comp) { - auto nt_diff = notool_comp->diff; - closer_suffix = nt_diff.left.substr(nt_diff.left.find("YYYY") + 4); - } - } else { - closer_suffix = diff.suffix.substr(0, diff.suffix.find(suffix_marker)); - } - if (!closer_suffix.empty()) { - if (format.mode == tool_format::TAG_WITH_TAGGED) { - // After last arg value, skip the closing arg marker, rest is closer - auto closer_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.until("YYYY") + p.literal("YYYY") + p.space() + - p.marker() + p.space() + - p.tag("close", p.rest()); - }); - auto close_result = closer_parser.parse_and_extract(closer_suffix); - if (close_result.result.success()) { - function.close = close_result.tags["close"]; - } - } else if (format.mode == tool_format::TAG_WITH_JSON) { - // After last arg value, find end of JSON args, rest is closer - auto closer_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { - return p.until("YYYY") + p.literal("YYYY") + p.tag("post_val", p.rest()); - }); - auto close_result = closer_parser.parse_and_extract(closer_suffix); - if (close_result.result.success()) { - const auto & post = close_result.tags["post_val"]; - size_t pos = post.find_last_of("}]"); - if (pos != std::string::npos && pos < post.size() - 1) { - function.close = trim_leading_whitespace(post.substr(pos + 1)); - } - } - } - } - function.close = trim_leading_whitespace(function.close); - } -} - -void analyze_tools::analyze_arguments() { - extract_argument_name_markers(); - extract_argument_value_markers(); -} - -void analyze_tools::extract_argument_name_markers() { - json assistant_first_arg = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call_one_arg }) } - }; - - json assistant_second_arg = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call_other_arg }) } - }; - - template_params params; - params.messages = json::array({ user_msg, assistant_first_arg }); - params.tools = tools; - params.add_generation_prompt = false; - params.enable_thinking = true; - - auto comparison = compare_variants( - *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_second_arg }); }); - - if (!comparison) { - LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); - return; - } - - const auto & diff = comparison->diff; - - if (!diff.left.empty() && !diff.right.empty()) { - // Parse both sides to find ARG_FIRST/ARG_SECOND and extract the surrounding structure - auto left_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { - return p.tag("pre", p.until(ARG_FIRST)) + p.literal(ARG_FIRST) + - p.tag("suffix", p.until_one_of({"\"", "X"})); - }); - auto right_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { - return p.tag("pre", p.until(ARG_SECOND)) + p.literal(ARG_SECOND) + - p.tag("suffix", p.until_one_of({"\"", "Y"})); - }); - auto left_result = left_parser.parse_anywhere_and_extract(diff.left); - auto right_result = right_parser.parse_anywhere_and_extract(diff.right); - - if (left_result.result.success() && right_result.result.success() && - !left_result.tags["pre"].empty() && - left_result.tags["pre"] == right_result.tags["pre"] && - left_result.tags["suffix"] == right_result.tags["suffix"]) { - // Name is inside a structure (e.g., JSON key): prefix is the shared wrapper - arguments.name_prefix = trim_whitespace(left_result.tags["pre"]); - arguments.name_suffix = trim_leading_whitespace(left_result.tags["suffix"]); - } else if (diff.left.substr(0, ARG_FIRST.length()) == ARG_FIRST && diff.right.substr(0, ARG_SECOND.length()) == ARG_SECOND) { - // Name is directly in the diff: prefix comes from last marker in diff.prefix - auto pre_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { - auto last_marker = p.marker() + p.zero_or_more(p.negate(p.marker()) + p.any()) + p.end(); - return p.zero_or_more(p.negate(last_marker) + p.any()) + p.tag("name_prefix", last_marker); - }); - auto pre_result = pre_parser.parse_and_extract(diff.prefix); - arguments.name_prefix = pre_result.result.success() - ? pre_result.tags["name_prefix"] : diff.prefix; - - // Suffix extends from after ARG_FIRST to the first marker (+ optional whitespace). - // The marker could be in diff.left itself or in diff.suffix, so we concatenate. - std::string after_first = diff.left.substr(ARG_FIRST.length()) + diff.suffix; - auto suffix_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { - return p.tag("suffix", p.zero_or_more(p.negate(p.marker()) + p.any()) + - p.marker() + p.space()); - }); - auto suf_result = suffix_parser.parse_anywhere_and_extract(after_first); - if (suf_result.result.success()) { - arguments.name_suffix = suf_result.tags["suffix"]; - } - } - } -} - -void analyze_tools::extract_argument_value_markers() { - json assistant_val_X = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call_one_arg }) } - }; - - json assistant_val_Y = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call_one_arg_other_val }) } - }; - - template_params params; - params.messages = json::array({ user_msg, assistant_val_X }); - params.tools = tools; - params.add_generation_prompt = false; - params.enable_thinking = true; - - auto comparison = compare_variants( - *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_val_Y }); }); - - if (!comparison) { - LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); - return; - } - - const auto & diff = comparison->diff; - - if (diff.left == "XXXX" && diff.right == "YYYY") { - std::string arg_name_ending = ARG_FIRST + arguments.name_suffix; - std::string prefix = diff.prefix; - if (prefix.rfind(arg_name_ending) != std::string::npos) { - prefix = prefix.substr(prefix.rfind(arg_name_ending) + arg_name_ending.size()); - } - if (!prefix.empty()) { - // Find the last marker + any trailing non-marker text to end - auto prefix_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { - auto last_marker = p.marker() + p.zero_or_more(p.negate(p.marker()) + p.any()) + p.end(); - return p.zero_or_more(p.negate(last_marker) + p.any()) + p.tag("val_prefix", last_marker); - }); - auto pre_result = prefix_parser.parse_and_extract(prefix); - arguments.value_prefix = pre_result.result.success() ? pre_result.tags["val_prefix"] : prefix; - } - - std::string value_suffix = diff.suffix; - if (!function.close.empty()) { - size_t func_close_pos = value_suffix.find(function.close); - if (func_close_pos != std::string::npos) { - value_suffix = value_suffix.substr(0, func_close_pos); - } - } else if (!format.per_call_end.empty() || !format.section_end.empty()) { - std::string end_marker = - !format.per_call_end.empty() ? format.per_call_end : format.section_end; - size_t end_marker_pos = value_suffix.find(end_marker); - if (end_marker_pos != std::string::npos) { - value_suffix = value_suffix.substr(0, end_marker_pos); - } - } - value_suffix = trim_leading_whitespace(value_suffix); - if (!value_suffix.empty()) { - arguments.value_suffix = value_suffix; - } - } -} - -void analyze_tools::extract_argument_separator() { - json assistant_one_arg = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call_one_arg }) } - }; - - json assistant_two_args = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call }) } - }; - - template_params params; - params.messages = json::array({ user_msg, assistant_one_arg }); - params.tools = tools; - params.add_generation_prompt = false; - params.enable_thinking = true; - - auto comparison = compare_variants( - *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_two_args }); }); - - if (!comparison) { - LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); - return; - } - - const auto & diff = comparison->diff; - - if (!diff.right.empty()) { - std::string separator = until_common_prefix(diff.right, ARG_FIRST, ARG_SECOND); - arguments.separator = separator; - } -} - -void analyze_tools::extract_args_markers() { - json assistant_no_args = json{ - { "role", "assistant"}, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call_zero_args }) } - }; - - json assistant_with_args = json{ - { "role", "assistant"}, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call_one_arg }) } - }; - - template_params params; - params.messages = json::array({ user_msg, assistant_no_args }); - params.tools = tools; - params.add_generation_prompt = false; - params.enable_thinking = true; - - auto comparison = compare_variants( - *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_with_args }); }); - - if (!comparison) { - LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); - return; - } - - const auto & diff = comparison->diff; - - if (format.mode == tool_format::JSON_NATIVE) { - std::string prefix_marker = !format.section_start.empty() ? format.section_start : format.per_call_start; - std::string suffix_marker = !format.section_end.empty() ? format.section_end : format.per_call_end; - // these might happen earlier in the tools section as an example or somewhere else, so we need to find the closest ones - size_t prefix_pos = prefix_marker.empty() ? 0 : diff.prefix.rfind(prefix_marker); - size_t suffix_pos = suffix_marker.empty() ? diff.suffix.size() : diff.suffix.find(suffix_marker); - if (prefix_pos == std::string::npos) { - prefix_pos = 0; - } - if (suffix_pos == std::string::npos) { - suffix_pos = diff.suffix.size(); - } - std::string prefix_cut = diff.prefix.substr(prefix_pos + prefix_marker.size()); - std::string suffix_cut = diff.suffix.substr(0, suffix_pos); - std::string args_start = until_common_prefix(prefix_cut, "{}", "{\"first\":"); - std::string args_end = after_common_suffix(suffix_cut, "{}", "\"XXXX\"}"); - - if (!args_start.empty() || !args_end.empty()) { - size_t find_fun = args_start.find(FUN_FIRST); - if (find_fun != std::string::npos) { - args_start = args_start.substr(find_fun + FUN_FIRST.size(), args_start.size() - find_fun - FUN_FIRST.size()); - } - size_t find_call_id = args_start.find(CALL_ID_001); - if (find_call_id != std::string::npos) { - args_start = args_start.substr(find_call_id + CALL_ID_001.size(), args_start.size() - find_call_id - CALL_ID_001.size()); - } - arguments.start = args_start; - arguments.end = args_end; - } - } -} - -void analyze_tools::extract_call_id_markers() { - json assistant_id1 = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call }) } - }; - - json assistant_id2 = json{ - { "role", "assistant" }, - { "content", "" }, - { "tool_calls", json::array({ first_tool_call_alt_id }) } - }; - - template_params params; - params.messages = json::array({ user_msg, assistant_id1 }); - params.tools = tools; - params.add_generation_prompt = false; - params.enable_thinking = true; - - auto comparison = compare_variants( - *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_id2 }); }); - - if (!comparison) { - LOG_DBG(ANSI_ORANGE "%s: Template application failed for call_id detection\n" ANSI_RESET, __func__); - return; - } - - const auto & diff = comparison->diff; - - if (diff.left.empty() && diff.right.empty()) { - return; - } - - std::string id_value_1 = CALL_ID_001; - std::string id_value_2 = CALL_ID_999; - - size_t common_id_prefix_len = 0; - for (size_t i = 0; i < std::min(id_value_1.length(), id_value_2.length()); i++) { - if (id_value_1[i] == id_value_2[i]) { - common_id_prefix_len++; - } else { - break; - } - } - std::string common_id_part = id_value_1.substr(0, common_id_prefix_len); - - // Check if the function name is in the prefix (normal case: BETWEEN_FUNC_AND_ARGS or POST_ARGS) - // or in the suffix (call_id is PRE_FUNC_NAME) - std::string func_name = FUN_FIRST; - size_t func_name_in_prefix = diff.prefix.rfind(func_name); - size_t func_name_in_suffix = diff.suffix.find(func_name); - - // Helper: find the last marker in a string (returns just the marker, not trailing text) - auto find_last_marker = [](const std::string & str) -> std::string { - auto parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { - auto last = p.marker() + p.zero_or_more(p.negate(p.marker()) + p.any()) + p.end(); - return p.zero_or_more(p.negate(last) + p.any()) + p.tag("m", p.marker()); - }); - auto res = parser.parse_anywhere_and_extract(str); - return res.result.success() ? res.tags["m"] : ""; - }; - - // Helper: find the first marker in a string - auto find_first_marker = [](const std::string & str) -> std::string { - auto parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { - return p.tag("m", p.marker()); - }); - auto res = parser.parse_anywhere_and_extract(str); - return res.result.success() ? res.tags["m"] : ""; - }; - - if (func_name_in_prefix != std::string::npos && func_name_in_suffix == std::string::npos) { - // Function name is only in prefix - call_id is BETWEEN_FUNC_AND_ARGS or POST_ARGS - // Check if args indicator "{" is in prefix or suffix - size_t args_in_prefix = diff.prefix.find('{', func_name_in_prefix); - size_t args_in_suffix = diff.suffix.find('{'); - - if (args_in_suffix != std::string::npos && - (args_in_prefix == std::string::npos || args_in_prefix > diff.prefix.length())) { - // Args are in suffix, so call_id is BETWEEN_FUNC_AND_ARGS - call_id.pos = call_id_position::BETWEEN_FUNC_AND_ARGS; - - // Find call_id_prefix: marker immediately preceding common_id_part (no intervening markers) - std::string after_func = diff.prefix.substr(func_name_in_prefix + func_name.length()); - auto id_prefix_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { - return p.tag("prefix", p.marker()) + - p.zero_or_more(p.negate(p.marker()) + p.negate(p.literal(common_id_part)) + p.any()) + - p.literal(common_id_part); - }); - auto id_res = id_prefix_parser.parse_anywhere_and_extract(after_func); - if (id_res.result.success()) { - call_id.prefix = id_res.tags["prefix"]; - } else { - // Fallback: use the last marker in after_func - call_id.prefix = find_last_marker(after_func); - } - - // Extract call_id_suffix: the first marker in the suffix before args "{" - auto suffix_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { - return p.zero_or_more(p.negate(p.marker()) + p.negate(p.literal("{")) + p.any()) + - p.tag("suffix", p.marker()); - }); - auto suf_res = suffix_parser.parse_anywhere_and_extract(diff.suffix); - if (suf_res.result.success()) { - call_id.suffix = suf_res.tags["suffix"]; - } - } else if (args_in_prefix != std::string::npos) { - // Args are in prefix, so call_id is POST_ARGS - call_id.pos = call_id_position::POST_ARGS; - - // Extract last marker between args closing brace and the ID - std::string after_args = diff.prefix.substr(args_in_prefix); - size_t closing_brace = after_args.rfind('}'); - if (closing_brace != std::string::npos) { - std::string between_args_and_id = after_args.substr(closing_brace + 1); - call_id.prefix = find_last_marker(between_args_and_id); - } - - // call_id_suffix: first marker in diff.suffix - call_id.suffix = find_first_marker(diff.suffix); - } - } else if (func_name_in_suffix != std::string::npos && func_name_in_prefix == std::string::npos) { - // Function name is only in suffix - call_id is PRE_FUNC_NAME - call_id.pos = call_id_position::PRE_FUNC_NAME; - - // call_id_prefix: last marker in diff.prefix - call_id.prefix = find_last_marker(diff.prefix); - - // call_id_suffix: first marker in the portion of diff.suffix before func_name - std::string before_func = diff.suffix.substr(0, func_name_in_suffix); - call_id.suffix = find_first_marker(before_func); - } - - if (call_id.prefix == arguments.end) { - call_id.prefix = ""; - } - - if (call_id.suffix == arguments.start) { - call_id.suffix = ""; - } - - // When call_id is detected, per_call_end may have been incorrectly set to include - // the call_id_suffix and sample args. Clear it if it starts with call_id_suffix. - if (call_id.pos != call_id_position::NONE && !call_id.suffix.empty() && - format.per_call_end.find(call_id.suffix) == 0) { - format.per_call_end.clear(); - } -} - -} // namespace autoparser diff --git a/common/chat-peg-parser.cpp b/common/chat-peg-parser.cpp deleted file mode 100644 index 56eb567df0a4..000000000000 --- a/common/chat-peg-parser.cpp +++ /dev/null @@ -1,1007 +0,0 @@ -#include "chat-peg-parser.h" - -#include "chat-auto-parser.h" -#include "ggml.h" -#include "peg-parser.h" - -#include - -using ordered_json = nlohmann::ordered_json; - -static std::string_view trim_trailing_space(std::string_view sv, int max = -1) { - int count = 0; - while (!sv.empty() && std::isspace(static_cast(sv.back()))) { - if (max != -1 && count >= max) { - break; - } - sv.remove_suffix(1); - count++; - } - return sv; -} - -static std::string_view trim_leading_space(std::string_view sv, int max = -1) { - int count = 0; - while (!sv.empty() && std::isspace(static_cast(sv.front()))) { - if (max != -1 && count >= max) { - break; - } - sv.remove_prefix(1); - count++; - } - return sv; -} - -static std::string_view trim(std::string_view sv) { - return trim_trailing_space(trim_leading_space(sv, 1)); -} - -// Count the number of unclosed '{' braces in a JSON-like string, -// properly skipping braces inside quoted strings. -static int json_brace_depth(const std::string & s) { - int depth = 0; - bool in_string = false; - bool escaped = false; - for (char c : s) { - if (escaped) { - escaped = false; - continue; - } - if (c == '\\' && in_string) { - escaped = true; - continue; - } - if (c == '"') { - in_string = !in_string; - continue; - } - if (!in_string) { - if (c == '{') { - depth++; - } else if (c == '}') { - depth--; - } - } - } - return depth; -} - -// JSON-escape a string and return the inner content (without surrounding quotes). -static std::string escape_json_string_inner(const std::string & s) { - std::string escaped = ordered_json(s).dump(); - if (escaped.size() >= 2 && escaped.front() == '"' && escaped.back() == '"') { - return escaped.substr(1, escaped.size() - 2); - } - return escaped; -} - -// Convert Python-style single-quoted strings to JSON double-quoted strings -// Only converts outer string delimiters, properly handling escape sequences: -// - {'key': 'value'} -> {"key": "value"} -// - {'code': 'print(\'hello\')'} -> {"code": "print('hello')"} -// - {'msg': 'He said "hi"'} -> {"msg": "He said \"hi\""} -static std::string normalize_quotes_to_json(const std::string & input) { - std::string result; - result.reserve(input.size() + 16); // May need extra space for escaping - - bool in_single_quoted = false; - bool in_double_quoted = false; - - for (size_t i = 0; i < input.size(); ++i) { - char c = input[i]; - - // Handle escape sequences - if (c == '\\' && i + 1 < input.size()) { - char next = input[i + 1]; - - if (in_single_quoted) { - // Inside a single-quoted string being converted to double quotes - if (next == '\'') { - // \' -> ' (escaped single quote becomes unescaped in double-quoted string) - result += '\''; - ++i; - continue; - } - if (next == '"') { - // \" stays as \" (already escaped, works in double-quoted string) - result += "\\\""; - ++i; - continue; - } - // Other escapes (\n, \\, etc.): pass through both characters - result += c; - result += next; - ++i; - continue; - } - - if (in_double_quoted) { - // Inside a double-quoted string - pass through escape sequences as-is - result += c; - result += next; - ++i; - continue; - } - - // Outside any string - just pass through the backslash - result += c; - continue; - } - - // Handle quote characters - if (c == '"') { - if (in_single_quoted) { - // Unescaped double quote inside single-quoted string -> must escape for JSON - result += "\\\""; - } else { - // Double quote as string delimiter or outside strings - in_double_quoted = !in_double_quoted; - result += c; - } - } else if (c == '\'') { - if (in_double_quoted) { - // Single quote inside double-quoted string -> pass through - result += c; - } else if (in_single_quoted) { - // Closing single quote -> convert to double quote - in_single_quoted = false; - result += '"'; - } else { - // Opening single quote -> convert to double quote - in_single_quoted = true; - result += '"'; - } - } else { - result += c; - } - } - - return result; -} - -void tag_based_peg_mapper::from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result) { - arena.visit(result, [this](const common_peg_ast_node & node) { - if (!node.tag.empty()) { - tags[node.tag] = std::string(node.text); - } - }); -} - -tagged_parse_result tagged_peg_parser::parse_and_extract(const std::string & input, common_peg_parse_flags extra_flags) const { - common_peg_parse_context ctx(input, flags | extra_flags); - auto parse_result = arena.parse(ctx); - - tag_based_peg_mapper mapper; - mapper.from_ast(ctx.ast, parse_result); - - return { std::move(parse_result), std::move(mapper.tags) }; -} - -tagged_parse_result tagged_peg_parser::parse_anywhere_and_extract(const std::string & input) const { - if (input.empty()) { - return parse_and_extract(input); - } - for (size_t i = 0; i < input.size(); i++) { - common_peg_parse_context ctx(input, flags); - auto parse_result = arena.parse(ctx, i); - if (parse_result.success() || i == input.size() - 1) { - tag_based_peg_mapper mapper; - mapper.from_ast(ctx.ast, parse_result); - return { std::move(parse_result), std::move(mapper.tags) }; - } - } - GGML_ABORT("Should not happen"); -} - -tagged_peg_parser build_tagged_peg_parser( - const std::function & fn) { - common_peg_parser_builder builder; - builder.set_root(fn(builder)); - return { builder.build() }; -} - -common_peg_parser common_chat_peg_builder::tag_with_safe_content(const std::string & tag_name, - const std::string & marker, - const common_peg_parser & p) { - if (marker.empty()) { - return zero_or_more(choice({ p, rule(tag_name, content(any())) })); - } - auto content_chunk = rule(tag_name, content(negate(literal(marker)) + any() + until(marker))); - return zero_or_more(choice({ p, content_chunk })); -} - -std::string & common_chat_peg_mapper::args_target() { - return (current_tool && !current_tool->name.empty()) ? current_tool->arguments : args_buffer; -} - -std::string common_chat_peg_mapper::normalize_container_value(const std::string & input) { - return normalize_quotes_to_json(input); -} - -void common_chat_peg_mapper::from_ast(const common_peg_ast_arena & arena, - const common_peg_parse_result & parse_result_arg) { - arena.visit(parse_result_arg, [this](const common_peg_ast_node & node) { map(node); }); - // Flush any pending tool call that was started but never got a name - // This happens during partial parsing when the tool call is incomplete - if (pending_tool_call.has_value() && !pending_tool_call->name.empty()) { - if (!args_buffer.empty()) { - pending_tool_call->arguments = args_buffer; - } - if (closing_quote_pending && !pending_tool_call->arguments.empty()) { - pending_tool_call->arguments += "\""; - } - result.tool_calls.push_back(pending_tool_call.value()); - pending_tool_call.reset(); - } - - // Discard whitespace-only reasoning content (e.g. from prefill) - if (!result.reasoning_content.empty()) { - bool all_whitespace = true; - for (char c : result.reasoning_content) { - if (c != ' ' && c != '\n' && c != '\r' && c != '\t') { - all_whitespace = false; - break; - } - } - if (all_whitespace) { - result.reasoning_content.clear(); - } - } -} - -void common_chat_peg_mapper::map(const common_peg_ast_node & node) { - // Handle reasoning/content tags - bool is_reasoning = node.tag == common_chat_peg_builder::REASONING; - bool is_content = node.tag == common_chat_peg_builder::CONTENT; - - if (is_reasoning) { // GPT OSS can have more than 1 reasoning block, so concatenate here - result.reasoning_content += std::string(node.text); - } - - if (is_content) { - // Concatenate content from multiple content nodes (e.g., when reasoning markers - // are preserved before content markers in reasoning_format=NONE mode) - result.content += std::string(node.text); - } - - // Handle tool-related tags (supporting both JSON and tagged formats) - bool is_tool_open = node.tag == common_chat_peg_builder::TOOL_OPEN; - bool is_tool_close = node.tag == common_chat_peg_builder::TOOL_CLOSE; - bool is_tool_name = node.tag == common_chat_peg_builder::TOOL_NAME; - bool is_tool_id = node.tag == common_chat_peg_builder::TOOL_ID; - bool is_tool_args = node.tag == common_chat_peg_builder::TOOL_ARGS; - bool is_arg_open = node.tag == common_chat_peg_builder::TOOL_ARG_OPEN; - bool is_arg_close = node.tag == common_chat_peg_builder::TOOL_ARG_CLOSE; - bool is_arg_name = node.tag == common_chat_peg_builder::TOOL_ARG_NAME; - bool is_arg_value = node.tag == common_chat_peg_builder::TOOL_ARG_VALUE; - bool is_arg_string_value = node.tag == common_chat_peg_builder::TOOL_ARG_STRING_VALUE; - - if (is_tool_open) { - pending_tool_call = common_chat_tool_call(); - current_tool = &pending_tool_call.value(); - arg_count = 0; - args_buffer.clear(); - closing_quote_pending = false; - } - - if (is_tool_id && current_tool) { - auto text = trim_trailing_space(node.text); - if (text.size() >= 2 && text.front() == '"' && text.back() == '"') { - text = text.substr(1, text.size() - 2); - } - current_tool->id = std::string(text); - } - - if (is_tool_name && current_tool) { - current_tool->name = std::string(trim_trailing_space(node.text)); - // Now that we have the name, populate the arguments from the buffer - if (!args_buffer.empty()) { - current_tool->arguments = args_buffer; - args_buffer.clear(); - } else if (current_tool->arguments.empty()) { - current_tool->arguments = "{"; - } - // Add the tool call to results so streaming can see it - if (pending_tool_call.has_value()) { - result.tool_calls.push_back(pending_tool_call.value()); - pending_tool_call.reset(); - current_tool = &result.tool_calls.back(); - } - } - - if (is_tool_args && current_tool) { - // For JSON format: arguments come as a complete JSON object - // For tagged format: built up from individual arg_name/arg_value nodes - auto text = trim_trailing_space(node.text); - if (!text.empty() && text.front() == '{') { - args_target() = std::string(text); - } - } - - if (is_arg_open) { - closing_quote_pending = false; - } - - if (is_arg_name && current_tool) { - std::string arg_entry; - if (arg_count > 0) { - arg_entry = ","; - } - arg_entry += ordered_json(trim(node.text)).dump() + ":"; - ++arg_count; - - auto & target = args_target(); - if (target.empty()) { - target = "{"; - } - target += arg_entry; - } - - if ((is_arg_value || is_arg_string_value) && current_tool) { - std::string value_content = std::string(trim_trailing_space(trim_leading_space(node.text, 1), 1)); - - std::string value_to_add; - if (value_content.empty() && is_arg_string_value) { - // Empty string value - arg_close will add the closing quote - value_to_add = "\""; - closing_quote_pending = true; - } else if (!value_content.empty() && is_arg_string_value) { - // Schema declares this as string type - always treat as literal string value - if (!closing_quote_pending) { - value_to_add = "\""; - closing_quote_pending = true; - } - value_to_add += escape_json_string_inner(value_content); - } else if (!value_content.empty()) { - // For potential containers, normalize Python-style single quotes to JSON double quotes - bool is_potential_container = value_content[0] == '[' || value_content[0] == '{'; - if (is_potential_container) { - value_content = normalize_container_value(value_content); - } - - // Try to parse as JSON value (number, bool, null, object, array) - try { - ordered_json parsed = ordered_json::parse(value_content); - if (parsed.is_string()) { - // Don't add closing quote yet (added by arg_close) for monotonic streaming - std::string escaped = parsed.dump(); - if (!escaped.empty() && escaped.back() == '"') { - escaped.pop_back(); - } - value_to_add = escaped; - closing_quote_pending = true; - } else { - // Non-string values: use raw content to preserve whitespace for monotonicity - value_to_add = value_content; - } - } catch (...) { - if (node.is_partial && is_potential_container) { - // Partial container: pass through the already-normalized content - value_to_add = value_content; - } else { - // Not valid JSON - treat as string value - if (!closing_quote_pending) { - value_to_add = "\""; - closing_quote_pending = true; - } - value_to_add += escape_json_string_inner(value_content); - } - } - } - - args_target() += value_to_add; - } - - if (is_arg_close && current_tool) { - if (closing_quote_pending) { - args_target() += "\""; - closing_quote_pending = false; - } - } - - if (is_tool_close && current_tool) { - // Flush buffer to arguments if tool name was never seen - if (current_tool->name.empty() && !args_buffer.empty()) { - current_tool->arguments = args_buffer; - args_buffer.clear(); - } - // Close any pending string quote - if (closing_quote_pending) { - current_tool->arguments += "\""; - closing_quote_pending = false; - } - // Close any unclosed braces (accounts for nested objects) - for (int d = json_brace_depth(current_tool->arguments); d > 0; d--) { - current_tool->arguments += "}"; - } - // Add tool call to results if named; otherwise discard - if (pending_tool_call.has_value()) { - if (!current_tool->name.empty()) { - result.tool_calls.push_back(pending_tool_call.value()); - } - pending_tool_call.reset(); - } - } -} - -common_peg_parser common_chat_peg_builder::standard_constructed_tools( - const std::map & markers, - const ordered_json & tools, - bool parallel_tool_calls, - bool force_tool_calls) { - if (!tools.is_array() || tools.empty()) { - return eps(); - } - - // Extract markers with defaults - auto get_marker = [&markers](const std::string & key, const std::string & default_val = "") -> std::string { - auto it = markers.find(key); - return it != markers.end() ? it->second : default_val; - }; - - std::string section_start = get_marker("tool_call_start_marker", ""); - std::string section_end = get_marker("tool_call_end_marker", ""); - std::string func_opener = get_marker("function_opener", ""); - std::string func_closer = get_marker("function_closer", ""); - std::string param_key_prefix = get_marker("parameter_key_prefix", ""); - std::string param_closer = get_marker("parameter_closer", ""); - - // Build tool choices for tagged format - auto tool_choices = choice(); - - for (const auto & tool_def : tools) { - if (!tool_def.contains("function")) { - continue; - } - const auto & function = tool_def.at("function"); - std::string name = function.at("name"); - ordered_json params = function.contains("parameters") ? function.at("parameters") : ordered_json::object(); - - // Build argument parsers - auto args = eps(); - if (params.contains("properties") && !params["properties"].empty()) { - auto arg_choice = choice(); - for (const auto & el : params["properties"].items()) { - const std::string & prop_name = el.key(); - - auto arg_name_parser = - choice({ literal(prop_name), literal("\"" + prop_name + "\""), literal("'" + prop_name + "'") }); - - auto arg_rule = tool_arg(tool_arg_open(literal(param_key_prefix)) + tool_arg_name(arg_name_parser) + - literal(param_key_suffix) + tool_arg_value(until(param_closer)) + - tool_arg_close(literal(param_closer))); - arg_choice |= arg_rule; - } - args = zero_or_more(arg_choice + space()); - } - - // Build function parser: args - auto tool_parser = tool(tool_open(literal(func_opener) + tool_name(literal(name)) + literal(func_name_suffix)) + - space() + tool_args(args) + space() + tool_close(literal(func_closer))); - - tool_choices |= rule("tool-" + name, tool_parser); - } - - // Build the section with markers - auto section = - parallel_tool_calls ? - trigger_rule("tool-call", literal(section_start) + space() + one_or_more(tool_choices + space()) + - literal(section_end)) : - trigger_rule("tool-call", literal(section_start) + space() + tool_choices + space() + literal(section_end)); - - return force_tool_calls ? section : optional(section); -} - -// Python-style tool calls: name(arg1="value1", arg2=123) -// Used only by LFM2 for now, so we don't merge it into autoparser -common_peg_parser common_chat_peg_builder::python_style_tool_calls( - const ordered_json & tools, - bool parallel_tool_calls) { - if (!tools.is_array() || tools.empty()) { - return eps(); - } - - auto tool_choices = choice(); - - for (const auto & tool_def : tools) { - if (!tool_def.contains("function")) { - continue; - } - const auto & function = tool_def.at("function"); - std::string name = function.at("name"); - ordered_json params = function.contains("parameters") ? function.at("parameters") : ordered_json::object(); - - auto args = eps(); - if (params.contains("properties") && !params["properties"].empty()) { - auto arg_choice = choice(); - for (const auto & el : params["properties"].items()) { - const std::string & prop_name = el.key(); - const auto & prop_def = el.value(); - bool is_string_type = (prop_def.contains("type") && prop_def["type"] == "string"); - - auto arg_name_parser = literal(prop_name); - - common_peg_parser arg_value_parser = eps(); - auto string_value_parser = choice({ - literal("\"") + tool_arg_string_value(string_content('"')) + literal("\""), - literal("'") + tool_arg_string_value(string_content('\'')) + literal("'") - }); - - if (is_string_type) { - arg_value_parser = string_value_parser; - } else { - arg_value_parser = tool_arg_value(python_value()); - } - - // Full argument: name="value" or name=value - auto arg_rule = tool_arg( - tool_arg_open(eps()) + - tool_arg_name(arg_name_parser) + - literal("=") + - arg_value_parser + - tool_arg_close(eps()) - ); - arg_choice |= arg_rule; - } - - args = arg_choice + zero_or_more("," + space() + arg_choice); - } - - auto tool_parser = tool(tool_open(tool_name(literal(name)) + literal("(")) + - space() + tool_args(args) + space() + tool_close(literal(")")) - ); - - tool_choices |= rule("tool-" + name, tool_parser); - } - - if (parallel_tool_calls) { - return "[" + space() + tool_choices + zero_or_more("," + space() + tool_choices) + space() + "]"; - } - return "[" + space() + tool_choices + space() + "]"; -} - -// Helper: Parse dot notation key into prefix and field name -static std::pair parse_key_spec(const std::string & key) { - auto dot_pos = key.find('.'); - if (dot_pos == std::string::npos) { - return {"", key}; // Top-level field - } - return {key.substr(0, dot_pos), key.substr(dot_pos + 1)}; -} - -// Mode 1: function_is_key — parse {"function_name": {...}} -common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key( - const ordered_json & tools, - const std::string & args_key, - const std::string & effective_args_key, - const std::string & call_id_key, - const std::string & gen_call_id_key) { - - auto tool_choices = choice(); - - for (const auto & tool_def : tools) { - if (!tool_def.contains("function")) { - continue; - } - const auto & function = tool_def.at("function"); - std::string name = function.at("name"); - ordered_json params = function.contains("parameters") ? function.at("parameters") : ordered_json::object(); - - // Build inner object fields - std::vector inner_fields; - - if (!call_id_key.empty()) { - auto id_parser = atomic( - literal("\"" + call_id_key + "\"") + space() + literal(":") + space() + - literal("\"") + tool_id(string_content('"')) + literal("\"") - ); - inner_fields.push_back(optional(id_parser + space() + optional(literal(",") + space()))); - } - - if (!gen_call_id_key.empty()) { - auto gen_id_parser = atomic( - literal("\"" + gen_call_id_key + "\"") + space() + literal(":") + space() + - choice({ - literal("\"") + tool_id(string_content('"')) + literal("\""), - tool_id(json_number()) - }) - ); - inner_fields.push_back(optional(gen_id_parser + space() + optional(literal(",") + space()))); - } - - // Arguments — either wrapped in args_key or parsed directly - common_peg_parser args_parser = eps(); - if (args_key.empty()) { - args_parser = tool_args(schema(json(), "tool-" + name + "-schema", params)); - } else { - args_parser = literal("\"" + effective_args_key + "\"") + space() + literal(":") + space() + - tool_args(schema(json(), "tool-" + name + "-schema", params)); - } - inner_fields.push_back(args_parser); - - // Build inner object parser - common_peg_parser inner_object = eps(); - if (args_key.empty() && inner_fields.size() == 1) { - inner_object = inner_fields[0]; - } else { - inner_object = literal("{") + space(); - for (size_t i = 0; i < inner_fields.size(); i++) { - inner_object = inner_object + inner_fields[i]; - if (i < inner_fields.size() - 1) { - inner_object = inner_object + space(); - } - } - inner_object = inner_object + space() + literal("}"); - } - - auto tool_parser = tool( - tool_open(literal("{")) + space() + - literal("\"") + tool_name(literal(name)) + literal("\"") + - space() + literal(":") + space() + - inner_object + - space() + tool_close(literal("}")) - ); - - tool_choices |= rule("tool-" + name, tool_parser); - } - - return tool_choices; -} - -// Mode 2: Nested keys (dot notation like "function.name") -common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys( - const ordered_json & tools, - const std::string & effective_name_key, - const std::string & effective_args_key, - const std::string & call_id_key, - const std::string & gen_call_id_key) { - - auto tool_choices = choice(); - - auto name_spec = parse_key_spec(effective_name_key); - auto args_spec = parse_key_spec(effective_args_key); - - std::string nested_prefix = !name_spec.first.empty() ? name_spec.first : args_spec.first; - std::string nested_name_field = !name_spec.first.empty() ? name_spec.second : effective_name_key; - std::string nested_args_field = !args_spec.first.empty() ? args_spec.second : effective_args_key; - - for (const auto & tool_def : tools) { - if (!tool_def.contains("function")) { - continue; - } - const auto & function = tool_def.at("function"); - std::string name = function.at("name"); - ordered_json params = function.contains("parameters") ? function.at("parameters") : ordered_json::object(); - - auto nested_name = literal("\"" + nested_name_field + "\"") + space() + literal(":") + space() + - atomic(literal("\"") + tool_name(literal(name)) + literal("\"")); - auto nested_args = literal("\"" + nested_args_field + "\"") + space() + literal(":") + space() + - tool_args(schema(json(), "tool-" + name + "-schema", params)); - - auto nested_object = literal("{") + space() + - nested_name + space() + literal(",") + space() + - nested_args + - space() + literal("}"); - - // Format: { id?, "function": {...} } - auto tool_parser_body = tool_open(literal("{")) + space(); - - if (!call_id_key.empty()) { - auto id_spec = parse_key_spec(call_id_key); - if (id_spec.first.empty()) { - auto id_parser = atomic( - literal("\"" + call_id_key + "\"") + space() + literal(":") + space() + - literal("\"") + tool_id(string_content('"')) + literal("\"") - ); - tool_parser_body = tool_parser_body + optional(id_parser + space() + literal(",") + space()); - } - } - - if (!gen_call_id_key.empty()) { - auto gen_id_spec = parse_key_spec(gen_call_id_key); - if (gen_id_spec.first.empty()) { - auto gen_id_parser = atomic( - literal("\"" + gen_call_id_key + "\"") + space() + literal(":") + space() + - choice({ - literal("\"") + tool_id(string_content('"')) + literal("\""), - tool_id(json_number()) - }) - ); - tool_parser_body = tool_parser_body + optional(gen_id_parser + space() + literal(",") + space()); - } - } - - auto nested_field = literal("\"" + nested_prefix + "\"") + space() + literal(":") + space() + nested_object; - tool_parser_body = tool_parser_body + nested_field + space() + tool_close(literal("}")); - - tool_choices |= rule("tool-" + name, tool(tool_parser_body)); - } - - return tool_choices; -} - -// Mode 3: Flat keys with optional ID fields and parameter ordering -common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys( - const ordered_json & tools, - const std::string & effective_name_key, - const std::string & effective_args_key, - const std::string & call_id_key, - const std::string & gen_call_id_key, - const std::vector & parameters_order) { - - auto tool_choices = choice(); - auto name_key_parser = literal("\"" + effective_name_key + "\""); - auto args_key_parser = literal("\"" + effective_args_key + "\""); - - for (const auto & tool_def : tools) { - if (!tool_def.contains("function")) { - continue; - } - const auto & function = tool_def.at("function"); - std::string name = function.at("name"); - ordered_json params = function.contains("parameters") ? function.at("parameters") : ordered_json::object(); - - auto tool_name_ = name_key_parser + space() + literal(":") + space() + - atomic(literal("\"") + tool_name(literal(name)) + literal("\"")); - auto tool_args_ = args_key_parser + space() + literal(":") + space() + - tool_args(schema(json(), "tool-" + name + "-schema", params)); - - // Build ID parsers if keys are provided - common_peg_parser id_parser = eps(); - if (!call_id_key.empty()) { - id_parser = atomic( - literal("\"" + call_id_key + "\"") + space() + literal(":") + space() + - choice({ - literal("\"") + tool_id(string_content('"')) + literal("\""), - tool_id(json_number()) - }) - ); - } - - common_peg_parser gen_id_parser = eps(); - if (!gen_call_id_key.empty()) { - gen_id_parser = atomic( - literal("\"" + gen_call_id_key + "\"") + space() + literal(":") + space() + - choice({ - literal("\"") + tool_id(string_content('"')) + literal("\""), - tool_id(json_number()) - }) - ); - } - - // Create (parser, key) pairs for all fields, then sort by parameters_order - std::vector> parser_pairs; - parser_pairs.emplace_back(tool_name_, effective_name_key); - parser_pairs.emplace_back(tool_args_, effective_args_key); - if (!call_id_key.empty()) { - parser_pairs.emplace_back(optional(id_parser), call_id_key); - } - if (!gen_call_id_key.empty()) { - parser_pairs.emplace_back(optional(gen_id_parser), gen_call_id_key); - } - - std::sort(parser_pairs.begin(), parser_pairs.end(), - [¶meters_order](const auto & a, const auto & b) { - auto pos_a = std::find(parameters_order.begin(), parameters_order.end(), a.second); - auto pos_b = std::find(parameters_order.begin(), parameters_order.end(), b.second); - size_t idx_a = (pos_a == parameters_order.end()) ? parameters_order.size() : std::distance(parameters_order.begin(), pos_a); - size_t idx_b = (pos_b == parameters_order.end()) ? parameters_order.size() : std::distance(parameters_order.begin(), pos_b); - return idx_a < idx_b; - }); - - auto ordered_body = tool_open(literal("{")) + space(); - for (size_t i = 0; i < parser_pairs.size(); i++) { - ordered_body = ordered_body + parser_pairs[i].first; - if (i < parser_pairs.size() - 1) { - ordered_body = ordered_body + space() + literal(",") + space(); - } - } - ordered_body = ordered_body + space() + tool_close(literal("}")); - - tool_choices |= rule("tool-" + name, tool(ordered_body)); - } - - return tool_choices; -} - -common_peg_parser common_chat_peg_builder::prefix(const std::string & s, const std::string & delimiter) { - if (s.empty()) { - return eps(); - } - if (delimiter.empty()) { - return literal(s); - } - return literal(s.substr(0, s.rfind(delimiter))); -} - -common_peg_parser common_chat_peg_builder::standard_json_tools( - const std::string & section_start, - const std::string & section_end, - const ordered_json & tools, - bool parallel_tool_calls, - bool force_tool_calls, - const std::string & name_key, - const std::string & args_key, - bool array_wrapped, - bool function_is_key, - const std::string & call_id_key, - const std::string & gen_call_id_key, - const std::vector & parameters_order) { - if (!tools.is_array() || tools.empty()) { - return eps(); - } - - std::string effective_name_key = name_key.empty() ? "name" : name_key; - std::string effective_args_key = args_key.empty() ? "arguments" : args_key; - - // Dispatch to the appropriate builder based on the JSON layout mode - common_peg_parser tool_choices = eps(); - if (function_is_key) { - tool_choices = build_json_tools_function_is_key(tools, args_key, effective_args_key, call_id_key, gen_call_id_key); - } else { - auto name_spec = parse_key_spec(effective_name_key); - auto args_spec = parse_key_spec(effective_args_key); - if (!name_spec.first.empty() || !args_spec.first.empty()) { - tool_choices = build_json_tools_nested_keys(tools, effective_name_key, effective_args_key, call_id_key, gen_call_id_key); - } else { - tool_choices = build_json_tools_flat_keys(tools, effective_name_key, effective_args_key, call_id_key, gen_call_id_key, parameters_order); - } - } - - // Build the section with markers - auto tool_calls = tool_choices; - if (parallel_tool_calls) { - tool_calls = tool_calls + zero_or_more(space() + literal(",") + space() + tool_choices); - } - - if (array_wrapped) { - tool_calls = literal("[") + space() + tool_calls + space() + literal("]"); - } - - auto section = - trigger_rule("tool-call", literal(section_start) + space() + tool_calls + space() + literal(section_end)); - - return force_tool_calls ? section : optional(section); -} - -void common_chat_peg_gemma4_mapper::from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result) { - for (const auto & node : result.nodes) { - visit(arena, node); - } -} - -static std::string gemma4_to_json(const common_peg_ast_arena & arena, common_peg_ast_id id) { - const auto & node = arena.get(id); - - if (node.text.empty()) { - return ""; - } - - if (node.rule == "gemma4-number" || node.rule == "gemma4-bool" || node.rule == "gemma4-null") { - return std::string(node.text); - } - - if (node.rule == "gemma4-string-content") { - return escape_json_string_inner(std::string(node.text)); - } - - if (node.rule == "gemma4-string") { - std::string result = "\""; - if (!node.children.empty()) { - result += gemma4_to_json(arena, node.children[0]); - if (!node.is_partial) { - result += "\""; - } - } - return result; - } - - if (node.rule == "gemma4-array") { - std::string result = "["; - - bool add_comma = false; - for (auto child_id : node.children) { - if (add_comma) { - result += ','; - } - add_comma = true; - result += gemma4_to_json(arena, child_id); - } - - if (!node.is_partial) { - result += ']'; - } - return result; - } - - if (node.rule == "gemma4-dict-key-name") { - return std::string(node.text); - } - - if (node.rule == "gemma4-dict-key") { - std::string result = "\""; - if (!node.children.empty()) { - result += escape_json_string_inner(gemma4_to_json(arena, node.children[0])); - } - if (!node.is_partial) { - result += "\":"; - } - return result; - } - - if (node.rule == "gemma4-dict-kv") { - std::string result; - for (auto child_id : node.children) { - result += gemma4_to_json(arena, child_id); - } - return result; - } - - if (node.rule == "gemma4-dict") { - std::string result = "{"; - - bool add_comma = false; - for (auto child_id : node.children) { - if (add_comma) { - result += ','; - } - add_comma = true; - result += gemma4_to_json(arena, child_id); - } - - if (!node.is_partial) { - result += '}'; - } - return result; - } - - if (node.rule == "gemma4-value") { - if (!node.children.empty()) { - return gemma4_to_json(arena, node.children[0]); - } - return ""; - } - - return ""; -} - -void common_chat_peg_gemma4_mapper::visit(const common_peg_ast_arena & arena, common_peg_ast_id id) { - const auto & node = arena.get(id); - - if (node.tag == "reasoning") { - result.reasoning_content += std::string(node.text); - return; - } - - if (node.tag == "content") { - result.content += std::string(node.text); - return; - } - - if (node.tag == "tool") { - auto name_id = arena.find_by_tag(node, "tool-name"); - auto args_id = arena.find_by_tag(node, "tool-args"); - - if (name_id != COMMON_PEG_INVALID_AST_ID && args_id != COMMON_PEG_INVALID_AST_ID) { - const auto & name_node = arena.get(name_id); - const auto & args_node = arena.get(args_id); - - if (!name_node.is_partial) { - common_chat_tool_call call; - call.name = std::string(name_node.text); - if (!args_node.children.empty()) { - call.arguments = gemma4_to_json(arena, args_node.children[0]); - } - result.tool_calls.push_back(call); - } - } - - return; - } - - for (auto child_id : node.children) { - visit(arena, child_id); - } -} diff --git a/common/chat-peg-parser.h b/common/chat-peg-parser.h deleted file mode 100644 index 1ea3eb7eb862..000000000000 --- a/common/chat-peg-parser.h +++ /dev/null @@ -1,195 +0,0 @@ -#pragma once - -#include "chat.h" -#include "peg-parser.h" - -#include -#include -#include - -class common_chat_peg_mapper { - public: - common_chat_msg & result; - - common_chat_peg_mapper(common_chat_msg & msg) : result(msg) {} - - virtual ~common_chat_peg_mapper() = default; - - virtual void from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result); - virtual void map(const common_peg_ast_node & node); - protected: - virtual std::string normalize_container_value(const std::string & input); - private: - // Tool call handling state - std::optional pending_tool_call; // Tool call waiting for name - common_chat_tool_call * current_tool = nullptr; - int arg_count = 0; - bool closing_quote_pending = false; - std::string args_buffer; // Buffer to delay arguments until tool name is known - - // Returns a reference to the active argument destination string. - // Before tool_name is known, writes go to args_buffer; after, to current_tool->arguments. - std::string & args_target(); -}; - -class common_chat_peg_gemma4_mapper : public common_chat_peg_mapper { - public: - common_chat_peg_gemma4_mapper(common_chat_msg & msg) : common_chat_peg_mapper(msg) {} - virtual void from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result); - private: - void visit(const common_peg_ast_arena & arena, common_peg_ast_id id); -}; - -struct content_structure; -struct tool_call_structure; - -class common_chat_peg_builder : public common_peg_parser_builder { - public: - // Tag constants (from former common_chat_peg_base_builder) - static constexpr const char * REASONING_BLOCK = "reasoning-block"; - static constexpr const char * REASONING = "reasoning"; - static constexpr const char * CONTENT = "content"; - - // Tag constants - static constexpr const char * TOOL = "tool"; - static constexpr const char * TOOL_OPEN = "tool-open"; - static constexpr const char * TOOL_CLOSE = "tool-close"; - static constexpr const char * TOOL_ID = "tool-id"; - static constexpr const char * TOOL_NAME = "tool-name"; - static constexpr const char * TOOL_ARGS = "tool-args"; - static constexpr const char * TOOL_ARG = "tool-arg"; - static constexpr const char * TOOL_ARG_OPEN = "tool-arg-open"; - static constexpr const char * TOOL_ARG_CLOSE = "tool-arg-close"; - static constexpr const char * TOOL_ARG_NAME = "tool-arg-name"; - static constexpr const char * TOOL_ARG_VALUE = "tool-arg-value"; - static constexpr const char * TOOL_ARG_STRING_VALUE = "tool-arg-string-value"; // For schema-declared string types - - // Low-level tag methods (from former common_chat_peg_base_builder) - common_peg_parser reasoning_block(const common_peg_parser & p) { return tag(REASONING_BLOCK, p); } - - common_peg_parser reasoning(const common_peg_parser & p) { return tag(REASONING, p); } - - common_peg_parser content(const common_peg_parser & p) { return tag(CONTENT, p); } - - common_peg_parser tag_with_safe_content(const std::string & tag_name, - const std::string & marker, - const common_peg_parser & p); - - // Low-level tag methods - common_peg_parser tool(const common_peg_parser & p) { return tag(TOOL, p); } - common_peg_parser tool_open(const common_peg_parser & p) { return atomic(tag(TOOL_OPEN, p)); } - common_peg_parser tool_close(const common_peg_parser & p) { return atomic(tag(TOOL_CLOSE, p)); } - common_peg_parser tool_id(const common_peg_parser & p) { return atomic(tag(TOOL_ID, p)); } - common_peg_parser tool_name(const common_peg_parser & p) { return atomic(tag(TOOL_NAME, p)); } - common_peg_parser tool_args(const common_peg_parser & p) { return tag(TOOL_ARGS, p); } - common_peg_parser tool_arg(const common_peg_parser & p) { return tag(TOOL_ARG, p); } - common_peg_parser tool_arg_open(const common_peg_parser & p) { return atomic(tag(TOOL_ARG_OPEN, p)); } - common_peg_parser tool_arg_close(const common_peg_parser & p) { return atomic(tag(TOOL_ARG_CLOSE, p)); } - common_peg_parser tool_arg_name(const common_peg_parser & p) { return atomic(tag(TOOL_ARG_NAME, p)); } - common_peg_parser tool_arg_value(const common_peg_parser & p) { return tag(TOOL_ARG_VALUE, p); } - - // Use for schema-declared string types - won't be treated as potential JSON container - common_peg_parser tool_arg_string_value(const common_peg_parser & p) { return tag(TOOL_ARG_STRING_VALUE, p); } - common_peg_parser tool_arg_json_value(const common_peg_parser & p) { return atomic(tag(TOOL_ARG_VALUE, p)); } - - - // Return a parser that parses the prefix of a string, up to a given delimiter. - common_peg_parser prefix(const std::string & s, const std::string & delimiter = {}); - - // Legacy-compatible helper for building standard JSON tool calls - // Used by tests and manual parsers - // name_key/args_key: JSON key names for function name and arguments - // Empty or "name"/"arguments" will accept both common variations - // Supports dot notation for nested objects (e.g., "function.name") - // array_wrapped: if true, tool calls are wrapped in JSON array [...] - // function_is_key: if true, function name is the JSON key (e.g., {"func_name": {...}}) - // call_id_key: JSON key for string call ID (e.g., "id") - // gen_call_id_key: JSON key for generated integer call ID (e.g., "tool_call_id") - // parameters_order: order in which JSON fields should be parsed - common_peg_parser standard_json_tools(const std::string & section_start, - const std::string & section_end, - const nlohmann::ordered_json & tools, - bool parallel_tool_calls, - bool force_tool_calls, - const std::string & name_key = "", - const std::string & args_key = "", - bool array_wrapped = false, - bool function_is_key = false, - const std::string & call_id_key = "", - const std::string & gen_call_id_key = "", - const std::vector & parameters_order = {}); - - // Legacy-compatible helper for building XML/tagged style tool calls - // Used by tests and manual parsers - common_peg_parser standard_constructed_tools(const std::map & markers, - const nlohmann::ordered_json & tools, - bool parallel_tool_calls, - bool force_tool_calls); - - // Helper for Python-style function call format: name(arg1="value1", arg2=123) - // Used by LFM2 and similar templates - common_peg_parser python_style_tool_calls(const nlohmann::ordered_json & tools, - bool parallel_tool_calls); - - private: - // Implementation helpers for standard_json_tools — one per JSON tool call layout mode - common_peg_parser build_json_tools_function_is_key(const nlohmann::ordered_json & tools, - const std::string & args_key, - const std::string & effective_args_key, - const std::string & call_id_key, - const std::string & gen_call_id_key); - - common_peg_parser build_json_tools_nested_keys(const nlohmann::ordered_json & tools, - const std::string & effective_name_key, - const std::string & effective_args_key, - const std::string & call_id_key, - const std::string & gen_call_id_key); - - common_peg_parser build_json_tools_flat_keys(const nlohmann::ordered_json & tools, - const std::string & effective_name_key, - const std::string & effective_args_key, - const std::string & call_id_key, - const std::string & gen_call_id_key, - const std::vector & parameters_order); -}; - -inline common_peg_arena build_chat_peg_parser( - const std::function & fn) { - common_chat_peg_builder builder; - builder.set_root(fn(builder)); - return builder.build(); -} - -class tag_based_peg_mapper { - public: - std::map tags; - - void from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result); -}; - -struct tagged_parse_result { - common_peg_parse_result result; - std::map tags; -}; - -struct tagged_peg_parser { - common_peg_arena arena; - common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE; - - tagged_peg_parser & withDebug() { - flags |= COMMON_PEG_PARSE_FLAG_DEBUG; - return *this; - } - - tagged_peg_parser & withoutDebug() { - flags = flags & ~COMMON_PEG_PARSE_FLAG_DEBUG; - return *this; - } - - tagged_parse_result parse_and_extract(const std::string & input, common_peg_parse_flags extra_flags = COMMON_PEG_PARSE_FLAG_NONE) const; - tagged_parse_result parse_anywhere_and_extract(const std::string & input) const; -}; - -tagged_peg_parser build_tagged_peg_parser( - const std::function & fn); - diff --git a/common/chat.cpp b/common/chat.cpp deleted file mode 100644 index 159d625de99f..000000000000 --- a/common/chat.cpp +++ /dev/null @@ -1,2384 +0,0 @@ -#include "chat.h" - -#include "chat-auto-parser-helpers.h" -#include "chat-auto-parser.h" -#include "chat-peg-parser.h" -#include "common.h" -#include "ggml.h" -#include "json-schema-to-grammar.h" -#include "log.h" - -#include "jinja/value.h" -#include "jinja/runtime.h" -#include "jinja/caps.h" -#include "peg-parser.h" - -#include "nlohmann/json.hpp" - -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include - -using json = nlohmann::ordered_json; - -static std::string format_time(const std::chrono::system_clock::time_point & now, const std::string & format) { - auto time = std::chrono::system_clock::to_time_t(now); - auto local_time = *std::localtime(&time); - std::ostringstream ss; - ss << std::put_time(&local_time, format.c_str()); - auto res = ss.str(); - return res; -} - -static json safe_args_parse(const std::string & to_parse) { - std::string stripped = to_parse; - if (to_parse.at(0) == '"' && to_parse.at(to_parse.length() - 1) == '"') { - stripped = to_parse.substr(1, to_parse.length() - 1); - } - try { - return json::parse(stripped); - } catch (json::exception & e) { - return stripped; - } -} - -static std::string string_diff(const std::string & last, const std::string & current) { - if (last.empty()) { - return current; - } - if (!string_starts_with(current, last)) { - if (string_starts_with(last, current)) { - // This happens if the last generation ended on a partial stop word (not erased), - // and the current ended on a stop word (erased). - return ""; - } - throw std::runtime_error("Invalid diff: '" + last + "' not found at start of '" + current + "'"); - } - return current.substr(last.size()); -} - -static bool has_content_or_tool_calls(const common_chat_msg & msg) { - return !msg.content.empty() || !msg.tool_calls.empty(); -} - -json common_chat_msg::to_json_oaicompat(bool concat_typed_text) const { - if (!content.empty() && !content_parts.empty()) { - throw std::runtime_error("Cannot specify both content and content_parts"); - } - json jmsg { - {"role", role}, - }; - if (!content.empty()) { - jmsg["content"] = content; - } else if (!content_parts.empty()) { - if (concat_typed_text) { - std::string text; - bool last_was_media_marker = false; - // join parts with newline, do not add newline before or after media markers - for (const auto & part : content_parts) { - bool add_new_line = true; - if (part.type == "text") { - add_new_line = !last_was_media_marker && !text.empty(); - last_was_media_marker = false; - } else if (part.type == "media_marker") { - add_new_line = false; - last_was_media_marker = true; - } else { - LOG_WRN("Ignoring content part type: %s\n", part.type.c_str()); - continue; - } - - if (add_new_line) { - text += '\n'; - } - - text += part.text; - } - jmsg["content"] = text; - } else { - auto & parts = jmsg["content"] = json::array(); - for (const auto & part : content_parts) { - parts.push_back({ - {"type", part.type}, - {"text", part.text}, - }); - } - } - } else { - jmsg["content"] = ""; - } - if (!reasoning_content.empty()) { - jmsg["reasoning_content"] = reasoning_content; - } - if (!tool_name.empty()) { - jmsg["name"] = tool_name; - } - if (!tool_call_id.empty()) { - jmsg["tool_call_id"] = tool_call_id; - } - if (!tool_calls.empty()) { - jmsg["tool_calls"] = json::array(); - auto & jtool_calls = jmsg["tool_calls"]; - for (const auto & tool_call : tool_calls) { - json tc { - {"type", "function"}, - {"function", { - {"name", tool_call.name}, - {"arguments", json(tool_call.arguments)}, - }}, - }; - if (!tool_call.id.empty()) { - tc["id"] = tool_call.id; - } - // Some templates generate and require an id (sometimes in a very specific format, e.g. Mistral Nemo). - // We only generate a random id for the ones that don't generate one by themselves - // (they also won't get to see it as their template likely doesn't use it, so it's all for the client) - // {"id", tc.id.empty() ? gen_tool_call_id() : tc.id}, - jtool_calls.push_back(tc); - } - } - - return jmsg; -} - -std::vector common_chat_msg_diff::compute_diffs(const common_chat_msg & msg_prv, - const common_chat_msg & msg_new) { - std::vector diffs; - if (msg_new.tool_calls.size() > msg_prv.tool_calls.size()) { - diffs.reserve(msg_new.tool_calls.size() - msg_prv.tool_calls.size() + 3); - } else { - diffs.reserve(3); - } - - // TODO: these can become expensive for long messages - how to optimize? - if (msg_prv.reasoning_content != msg_new.reasoning_content) { - auto & diff = diffs.emplace_back(); - diff.reasoning_content_delta = string_diff(msg_prv.reasoning_content, msg_new.reasoning_content); - } - if (msg_prv.content != msg_new.content) { - auto & diff = diffs.emplace_back(); - diff.content_delta = string_diff(msg_prv.content, msg_new.content); - } - - if (msg_new.tool_calls.size() < msg_prv.tool_calls.size()) { - std::string err = "Invalid diff: now finding less tool calls!\n"; - err += " Previous (" + std::to_string(msg_prv.tool_calls.size()) + "):\n"; - for (const auto & tc : msg_prv.tool_calls) { - err += " - name: '" + tc.name + "', args: '" + tc.arguments + "'\n"; - } - err += " Current (" + std::to_string(msg_new.tool_calls.size()) + "):\n"; - for (const auto & tc : msg_new.tool_calls) { - err += " - name: '" + tc.name + "', args: '" + tc.arguments + "'\n"; - } - err += " Current msg text content:\n" + msg_new.content + "\n"; - throw std::runtime_error(err); - } - - if (!msg_prv.tool_calls.empty()) { - const auto idx = msg_prv.tool_calls.size() - 1; - const auto & pref = msg_prv.tool_calls[idx]; - const auto & newf = msg_new.tool_calls[idx]; - // Allow tool name to change during incremental parsing: - // - empty -> non-empty (initial discovery) - // - prefix -> longer string (name grows as more input is parsed) - if (pref.name != newf.name && !pref.name.empty() && !newf.name.empty()) { - // Check if one is a prefix of the other (for incremental parsing where names grow or shrink) - bool is_prefix = (newf.name.rfind(pref.name, 0) == 0); - if (!is_prefix) { - LOG_ERR("Tool call mismatch: prev='%s' new='%s'\n", pref.name.c_str(), newf.name.c_str()); - throw std::runtime_error("Invalid diff: tool call mismatch!"); - } - } - const auto args_diff = string_diff(pref.arguments, newf.arguments); - if (!args_diff.empty() || pref.id != newf.id || pref.name != newf.name) { - auto & diff = diffs.emplace_back(); - diff.tool_call_index = idx; - if (pref.id != newf.id || pref.name != newf.name) { - diff.tool_call_delta.id = newf.id; - diff.tool_call_delta.name = newf.name; - } - diff.tool_call_delta.arguments = args_diff; - } - } - for (size_t idx = msg_prv.tool_calls.size(); idx < msg_new.tool_calls.size(); ++idx) { - auto & diff = diffs.emplace_back(); - diff.tool_call_index = idx; - diff.tool_call_delta = msg_new.tool_calls[idx]; - } - - return diffs; -} - -using chat_template_caps = jinja::caps; - -struct common_chat_templates { - bool add_bos; - bool add_eos; - bool has_explicit_template; // Model had builtin template or template overridden was specified. - std::unique_ptr template_default; // always set (defaults to chatml) - std::unique_ptr template_tool_use; -}; - -common_chat_tool_choice common_chat_tool_choice_parse_oaicompat(const std::string & tool_choice) { - if (tool_choice == "auto") { - return COMMON_CHAT_TOOL_CHOICE_AUTO; - } - if (tool_choice == "none") { - return COMMON_CHAT_TOOL_CHOICE_NONE; - } - if (tool_choice == "required") { - return COMMON_CHAT_TOOL_CHOICE_REQUIRED; - } - throw std::invalid_argument("Invalid tool_choice: " + tool_choice); -} - -bool common_chat_templates_support_enable_thinking(const common_chat_templates * chat_templates) { - common_chat_templates_inputs inputs; - inputs.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK; - common_chat_msg msg; - msg.role = "user"; - msg.content = "test"; - inputs.messages = { msg }; - inputs.enable_thinking = true; - inputs.add_generation_prompt = true; - inputs.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK; - - auto params = common_chat_templates_apply(chat_templates, inputs); - return params.supports_thinking; -} - -std::vector common_chat_msgs_parse_oaicompat(const json & messages) { - std::vector msgs; - - try { - if (!messages.is_array()) { - throw std::invalid_argument("Expected 'messages' to be an array, got " + messages.dump()); - } - - for (const auto & message : messages) { - if (!message.is_object()) { - throw std::invalid_argument("Expected 'message' to be an object, got " + message.dump()); - } - - common_chat_msg msg; - if (!message.contains("role")) { - throw std::invalid_argument("Missing 'role' in message: " + message.dump()); - } - msg.role = message.at("role"); - - auto has_content = message.contains("content"); - auto has_tool_calls = message.contains("tool_calls"); - if (has_content) { - const auto & content = message.at("content"); - if (content.is_string()) { - msg.content = content; - } else if (content.is_array()) { - for (const auto & part : content) { - if (!part.contains("type")) { - throw std::invalid_argument("Missing content part type: " + part.dump()); - } - const auto & type = part.at("type"); - if (type != "text" && type != "media_marker") { - throw std::invalid_argument("Unsupported content part type: " + type.dump()); - } - common_chat_msg_content_part msg_part; - msg_part.type = type; - msg_part.text = part.at("text"); - msg.content_parts.push_back(msg_part); - } - } else if (!content.is_null()) { - throw std::invalid_argument("Invalid 'content' type: expected string or array, got " + - content.dump() + - " (ref: https://github.com/ggml-org/llama.cpp/issues/8367)"); - } - } - if (has_tool_calls) { - for (const auto & tool_call : message.at("tool_calls")) { - common_chat_tool_call tc; - if (!tool_call.contains("type")) { - throw std::invalid_argument("Missing tool call type: " + tool_call.dump()); - } - const auto & type = tool_call.at("type"); - if (type != "function") { - throw std::invalid_argument("Unsupported tool call type: " + tool_call.dump()); - } - if (!tool_call.contains("function")) { - throw std::invalid_argument("Missing tool call function: " + tool_call.dump()); - } - const auto & fc = tool_call.at("function"); - if (!fc.contains("name")) { - throw std::invalid_argument("Missing tool call name: " + tool_call.dump()); - } - tc.name = fc.at("name"); - const auto & args = fc.at("arguments"); - if (args.is_string()) { - tc.arguments = args; - } else { - tc.arguments = args.dump(); - } - if (tool_call.contains("id")) { - tc.id = tool_call.at("id"); - } - msg.tool_calls.push_back(tc); - } - } - if (!has_content && !has_tool_calls) { - throw std::invalid_argument( - "Expected 'content' or 'tool_calls' (ref: https://github.com/ggml-org/llama.cpp/issues/8367 & " - "https://github.com/ggml-org/llama.cpp/issues/12279)"); - } - if (message.contains("reasoning_content")) { - msg.reasoning_content = message.at("reasoning_content"); - } - if (message.contains("name")) { - msg.tool_name = message.at("name"); - } - if (message.contains("tool_call_id")) { - msg.tool_call_id = message.at("tool_call_id"); - } - - msgs.push_back(msg); - } - } catch (const std::exception & e) { - // @ngxson : disable otherwise it's bloating the API response - // printf("%s\n", std::string("; messages = ") + messages.dump(2)); - throw std::runtime_error("Failed to parse messages: " + std::string(e.what())); - } - - return msgs; -} - -static json render_message_to_json(const std::vector & msgs, const jinja::caps & c) { - if (!c.supports_string_content && !c.supports_typed_content) { - LOG_WRN("%s: Neither string content nor typed content is supported by the template. This is unexpected and may lead to issues.\n", __func__); - } - - bool only_string_accepted = c.supports_string_content && !c.supports_typed_content; - bool only_typed_accepted = !c.supports_string_content && c.supports_typed_content; - - json messages = json::array(); - for (const auto & msg : msgs) { - if (only_string_accepted) { - json jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ true); - messages.push_back(jmsg); - } else if (only_typed_accepted) { - json jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ false); - if (jmsg.at("content").is_string()) { - jmsg["content"] = json::array({ - json{ - {"type", "text"}, - {"text", jmsg.at("content").get()}, - } - }); - } - messages.push_back(jmsg); - } else { - json jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ false); - messages.push_back(jmsg); - } - } - return messages; -} - -// DEPRECATED: only used in tests -json common_chat_msgs_to_json_oaicompat(const std::vector & msgs, bool concat_typed_text) { - jinja::caps c; - c.supports_string_content = true; - c.supports_typed_content = !concat_typed_text; - return render_message_to_json(msgs, c); -} - -json common_chat_tools_to_json_oaicompat(const std::vector & tools) { - if (tools.empty()) { - return json(); - } - - auto result = json::array(); - for (const auto & tool : tools) { - result.push_back({ - { "type", "function" }, - { "function", { - { "name", tool.name }, - { "description", tool.description }, - { "parameters", json::parse(tool.parameters) }, - }}, - }); - } - return result; -} - -std::vector common_chat_tools_parse_oaicompat(const json & tools) { - std::vector result; - - try { - if (!tools.is_null()) { - if (!tools.is_array()) { - throw std::invalid_argument("Expected 'tools' to be an array, got " + tools.dump()); - } - for (const auto & tool : tools) { - if (!tool.contains("type")) { - throw std::invalid_argument("Missing tool type: " + tool.dump()); - } - const auto & type = tool.at("type"); - if (!type.is_string() || type != "function") { - throw std::invalid_argument("Unsupported tool type: " + tool.dump()); - } - if (!tool.contains("function")) { - throw std::invalid_argument("Missing tool function: " + tool.dump()); - } - - const auto & function = tool.at("function"); - result.push_back({ - /* .name = */ function.at("name"), - /* .description = */ function.value("description", ""), - /* .parameters = */ function.value("parameters", json::object()).dump(), - }); - } - } - } catch (const std::exception & e) { - throw std::runtime_error("Failed to parse tools: " + std::string(e.what()) + "; tools = " + tools.dump(2)); - } - - return result; -} - -bool common_chat_verify_template(const std::string & tmpl, bool use_jinja) { - if (use_jinja) { - try { - common_chat_msg msg; - msg.role = "user"; - msg.content = "test"; - - auto tmpls = common_chat_templates_init(/* model= */ nullptr, tmpl); - - common_chat_templates_inputs inputs; - inputs.messages = { msg }; - - common_chat_templates_apply(tmpls.get(), inputs); - return true; - } catch (const std::exception & e) { - LOG_ERR("%s: failed to apply template: %s\n", __func__, e.what()); - return false; - } - } - llama_chat_message chat[] = { - { "user", "test" } - }; - const int res = llama_chat_apply_template(tmpl.c_str(), chat, 1, true, nullptr, 0); - return res >= 0; -} - -std::string common_chat_format_single(const struct common_chat_templates * tmpls, - const std::vector & past_msg, - const common_chat_msg & new_msg, - bool add_ass, - bool use_jinja) { - common_chat_templates_inputs inputs; - inputs.use_jinja = use_jinja; - inputs.add_bos = tmpls->add_bos; - inputs.add_eos = tmpls->add_eos; - - std::string fmt_past_msg; - if (!past_msg.empty()) { - inputs.messages = past_msg; - inputs.add_generation_prompt = false; - fmt_past_msg = common_chat_templates_apply(tmpls, inputs).prompt; - } - std::ostringstream ss; - // if the past_msg ends with a newline, we must preserve it in the formatted version - if (add_ass && !fmt_past_msg.empty() && fmt_past_msg.back() == '\n') { - ss << "\n"; - }; - // format chat with new_msg - inputs.messages.push_back(new_msg); - inputs.add_generation_prompt = add_ass; - auto fmt_new_msg = common_chat_templates_apply(tmpls, inputs).prompt; - // get the diff part - ss << fmt_new_msg.substr(fmt_past_msg.size(), fmt_new_msg.size() - fmt_past_msg.size()); - return ss.str(); -} - -std::string common_chat_format_example(const struct common_chat_templates * tmpls, - bool use_jinja, - const std::map & chat_template_kwargs) { - common_chat_templates_inputs inputs; - inputs.use_jinja = use_jinja; - inputs.add_bos = tmpls->add_bos; - inputs.add_eos = tmpls->add_eos; - inputs.chat_template_kwargs = chat_template_kwargs; - auto add_simple_msg = [&](auto role, auto content) { - common_chat_msg msg; - msg.role = role; - msg.content = content; - inputs.messages.push_back(msg); - }; - add_simple_msg("system", "You are a helpful assistant"); - add_simple_msg("user", "Hello"); - add_simple_msg("assistant", "Hi there"); - add_simple_msg("user", "How are you?"); - return common_chat_templates_apply(tmpls, inputs).prompt; -} - -#define CHATML_TEMPLATE_SRC \ - "{%- for message in messages -%}\n" \ - " {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' -}}\n" \ - "{%- endfor -%}\n" \ - "{%- if add_generation_prompt -%}\n" \ - " {{- '<|im_start|>assistant\n' -}}\n" \ - "{%- endif -%}" - -void common_chat_templates_free(struct common_chat_templates * tmpls) { - delete tmpls; -} - -bool common_chat_templates_was_explicit(const struct common_chat_templates * tmpls) { - return tmpls->has_explicit_template; -} - -// LFM2 format detection: template uses <|tool_list_start|>[...]<|tool_list_end|> around the tool list -// and <|tool_call_start|>[...]<|tool_call_end|> around each tool call -static bool is_lfm2_template(const std::string & src) { - return src.find("<|tool_list_start|>") != std::string::npos && - src.find("<|tool_list_end|>") != std::string::npos; -} - -common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates) { - common_chat_prompt_preset asr_preset; - asr_preset.system = ""; - asr_preset.user = "Transcribe audio to text"; - - if (chat_templates && chat_templates->template_default && is_lfm2_template(chat_templates->template_default->source())) { - asr_preset.system = "Perform ASR."; - asr_preset.user = ""; - } - - return asr_preset; -} - -std::string common_chat_templates_source(const struct common_chat_templates * tmpls, const std::string & variant) { - if (!variant.empty()) { - if (variant == "tool_use") { - if (tmpls->template_tool_use) { - return tmpls->template_tool_use->source(); - } - return ""; - } - LOG_DBG("%s: unknown template variant: %s\n", __func__, variant.c_str()); - } - return tmpls->template_default->source(); -} - -common_chat_templates_ptr common_chat_templates_init(const struct llama_model * model, - const std::string & chat_template_override, - const std::string & bos_token_override, - const std::string & eos_token_override) { - std::string default_template_src; - std::string template_tool_use_src; - - bool has_explicit_template = !chat_template_override.empty(); - if (chat_template_override.empty()) { - GGML_ASSERT(model != nullptr); - const auto * str = llama_model_chat_template(model, /* name */ nullptr); - if (str) { - default_template_src = str; - has_explicit_template = true; - } - str = llama_model_chat_template(model, /* name */ "tool_use"); - if (str) { - template_tool_use_src = str; - has_explicit_template = true; - } - } else { - default_template_src = chat_template_override; - } - if (default_template_src.empty() || default_template_src == "chatml") { - if (!template_tool_use_src.empty()) { - default_template_src = template_tool_use_src; - } else { - default_template_src = CHATML_TEMPLATE_SRC; - } - } - - // TODO @ngxson : this is a temporary hack to prevent chat template from throwing an error - // Ref: https://github.com/ggml-org/llama.cpp/pull/15230#issuecomment-3173959633 - if (default_template_src.find("<|channel|>") != std::string::npos - // search for the error message and patch it - && default_template_src.find("in message.content or") != std::string::npos) { - string_replace_all(default_template_src, - "{%- if \"<|channel|>analysis<|message|>\" in message.content or " - "\"<|channel|>final<|message|>\" in message.content %}", - "{%- if false %}"); - } - - // TODO @aldehir : this is a temporary fix, pending Minja changes - // Ref: https://github.com/ggml-org/llama.cpp/pull/17713#issuecomment-3631342664 - if (default_template_src.find("[TOOL_CALLS]") != std::string::npos - // search for the error message and patch it - && default_template_src.find("if (message['content'] is none or") != std::string::npos) { - string_replace_all(default_template_src, - "{%- if (message['content'] is none or message['content'] == '' or " - "message['content']|length == 0) and (message['tool_calls'] is not defined or " - "message['tool_calls'] is none or message['tool_calls']|length == 0) %}", - "{%- if false %}"); - } - - std::string token_bos = bos_token_override; - std::string token_eos = eos_token_override; - bool add_bos = false; - bool add_eos = false; - if (model) { - const auto * vocab = llama_model_get_vocab(model); - const auto get_token = [&](llama_token token, const char * name, const char * jinja_variable_name) { - if (token == LLAMA_TOKEN_NULL) { - if (default_template_src.find(jinja_variable_name) != std::string::npos || - template_tool_use_src.find(jinja_variable_name) != std::string::npos) { - LOG_WRN( - "common_chat_templates_init: warning: vocab does not have a %s token, jinja template won't " - "work as intended.\n", - name); - } - return std::string(); - } - return common_token_to_piece(vocab, token, true); - }; - token_bos = get_token(llama_vocab_bos(vocab), "BOS", "bos_token"); - token_eos = get_token(llama_vocab_eos(vocab), "EOS", "eos_token"); - add_bos = llama_vocab_get_add_bos(vocab); - add_eos = llama_vocab_get_add_eos(vocab); - } - common_chat_templates_ptr tmpls(new common_chat_templates()); - tmpls->has_explicit_template = has_explicit_template; - tmpls->add_bos = add_bos; - tmpls->add_eos = add_eos; - try { - tmpls->template_default = std::make_unique(default_template_src, token_bos, token_eos); - } catch (const std::exception & e) { - LOG_ERR("%s: error: %s\n", __func__, e.what()); - LOG_ERR("%s: failed to initialize chat template\n", __func__); - LOG_ERR("%s: please consider disabling jinja via --no-jinja, or using another chat template\n", __func__); - throw e; - } - if (!template_tool_use_src.empty()) { - try { - tmpls->template_tool_use = std::make_unique(template_tool_use_src, token_bos, token_eos); - } catch (const std::exception & e) { - LOG_ERR("%s: failed to parse tool use chat template (ignoring it): %s\n", __func__, e.what()); - } - } - return tmpls; -} - -const char * common_chat_format_name(common_chat_format format) { - switch (format) { - case COMMON_CHAT_FORMAT_CONTENT_ONLY: - return "Content-only"; - case COMMON_CHAT_FORMAT_PEG_SIMPLE: - return "peg-simple"; - case COMMON_CHAT_FORMAT_PEG_NATIVE: - return "peg-native"; - case COMMON_CHAT_FORMAT_PEG_GEMMA4: - return "peg-gemma4"; - default: - throw std::runtime_error("Unknown chat format"); - } -} - -const char * common_reasoning_format_name(common_reasoning_format format) { - switch (format) { - case COMMON_REASONING_FORMAT_NONE: - return "none"; - case COMMON_REASONING_FORMAT_AUTO: - return "auto"; - case COMMON_REASONING_FORMAT_DEEPSEEK: - return "deepseek"; - case COMMON_REASONING_FORMAT_DEEPSEEK_LEGACY: - return "deepseek-legacy"; - default: - throw std::runtime_error("Unknown reasoning format"); - } -} - -common_reasoning_format common_reasoning_format_from_name(const std::string & format) { - if (format == "none") { - return COMMON_REASONING_FORMAT_NONE; - } - if (format == "auto") { - return COMMON_REASONING_FORMAT_AUTO; - } - if (format == "deepseek") { - return COMMON_REASONING_FORMAT_DEEPSEEK; - } - if (format == "deepseek-legacy") { - return COMMON_REASONING_FORMAT_DEEPSEEK_LEGACY; - } - throw std::runtime_error("Unknown reasoning format: " + format); -} - -static void foreach_function(const json & tools, const std::function & fn) { - for (const auto & tool : tools) { - if (!tool.contains("type") || tool.at("type") != "function" || !tool.contains("function")) { - LOG_INF("Skipping tool without function: %s", tool.dump(2).c_str()); - continue; - } - fn(tool); - } -} - -static void foreach_parameter(const json & function, - const std::function & fn) { - if (!function.contains("parameters") || !function.at("parameters").is_object()) { - return; - } - const auto & params = function.at("parameters"); - if (!params.contains("properties") || !params.at("properties").is_object()) { - return; - } - const auto & props = params.at("properties"); - std::set required; - if (params.contains("required") && params.at("required").is_array()) { - params.at("required").get_to(required); - } - for (const auto & [name, prop] : props.items()) { - bool is_required = (required.find(name) != required.end()); - fn(name, prop, is_required); - } -} - -static std::string common_chat_template_direct_apply_impl( - const common_chat_template & tmpl, - const autoparser::generation_params & inputs, - const std::optional & messages_override = std::nullopt, - const std::optional & tools_override = std::nullopt, - const std::optional & additional_context = std::nullopt) { - jinja::context ctx(tmpl.source()); - - nlohmann::ordered_json inp = nlohmann::ordered_json{ - {"messages", messages_override.has_value() ? *messages_override : inputs.messages}, - {"bos_token", tmpl.bos_token()}, - {"eos_token", tmpl.eos_token()}, - {"enable_thinking", inputs.enable_thinking}, - }; - if (tools_override.has_value() || !inputs.tools.empty()) { - inp["tools"] = tools_override.has_value() ? *tools_override : inputs.tools; - } - if (inputs.extra_context.is_object()) { - // TODO: do we need to merge, or replacing is fine? - for (const auto & [k, v] : inputs.extra_context.items()) { - inp[k] = v; - } - } - if (additional_context.has_value()) { - // TODO: merge properly instead of overwriting (matching old behavior) - for (const auto & [k, v] : additional_context->items()) { - inp[k] = v; - } - } - if (inputs.add_generation_prompt) { - inp["add_generation_prompt"] = true; - } - - jinja::global_from_json(ctx, inp, inputs.mark_input); - - // render - jinja::runtime runtime(ctx); - const jinja::value results = runtime.execute(tmpl.prog); - auto parts = jinja::runtime::gather_string_parts(results); - - std::string result = parts->as_string().str(); - - // TODO: improve this later - if (inputs.add_bos && string_starts_with(result, tmpl.bos_token())) { - result = result.substr(tmpl.bos_token().size()); - } - if (inputs.add_eos && string_ends_with(result, tmpl.eos_token())) { - result = result.substr(0, result.size() - tmpl.eos_token().size()); - } - return result; -} - -std::string common_chat_template_direct_apply( - const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { - return common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt); -} - -static common_chat_params common_chat_params_init_ministral_3(const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { - common_chat_params data; - - // Build up messages to follow the format: https://huggingface.co/mistralai/Ministral-3-14B-Reasoning-2512/blob/main/chat_template.jinja - auto adjusted_messages = json::array(); - for (const auto & msg : inputs.messages) { - auto role = msg.value("role", ""); - if (role != "system" && role != "assistant") { - // Only adjust system and assistant messages. Interestingly, the system message may contain thinking. - adjusted_messages.push_back(msg); - continue; - } - - auto content = json::array(); - - // If message contains `reasoning_content`, add it as a block of type `thinking` - if (msg.contains("reasoning_content") && msg.at("reasoning_content").is_string()) { - content.push_back({ - { "type", "thinking" }, - { "thinking", msg.at("reasoning_content").get() }, - }); - } - - // If message contains `content`, add it as a block of type `text` - if (msg.contains("content")) { - if (msg.at("content").is_string()) { - content.push_back({ - { "type", "text" }, - { "text", msg.at("content").get() }, - }); - } else if (msg.at("content").is_array()) { - auto blocks = msg.at("content"); - content.insert(content.end(), blocks.begin(), blocks.end()); - } - } - - auto adjusted = msg; - adjusted["content"] = content; - adjusted.erase("reasoning_content"); - adjusted_messages.push_back(adjusted); - } - - auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); - auto has_response_format = inputs.json_schema.is_object() && !inputs.json_schema.empty(); - auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; - auto include_grammar = true; - - data.supports_thinking = true; - data.thinking_start_tag = "[THINK]"; - data.thinking_end_tag = "[/THINK]"; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages); - data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; - data.preserved_tokens = { - "[THINK]", - "[/THINK]", - "[TOOL_CALLS]", - "[ARGS]", - }; - - auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { - auto generation_prompt = p.prefix(inputs.generation_prompt, "[THINK]"); - auto reasoning = - extract_reasoning ? p.optional("[THINK]" + p.reasoning(p.until("[/THINK]")) + "[/THINK]") : p.eps(); - - // Response format parser - if (has_response_format) { - // Ministral wants to emit json surrounded by code fences - return generation_prompt + (reasoning << "```json" << p.content(p.schema(p.json(), "response-format", inputs.json_schema)) << "```"); - } - - // Tool call parser - if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { - auto tool_choice = p.choice(); - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - std::string name = function.at("name"); - const auto & schema = function.at("parameters"); - - tool_choice |= - p.rule("tool-" + name, p.tool_open(p.tool_name(p.literal(name)) + "[ARGS]") + - p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema))); - }); - - auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0; - auto max_calls = inputs.parallel_tool_calls ? -1 : 1; - auto tool_calls = p.trigger_rule("tool-call", p.repeat("[TOOL_CALLS]" + tool_choice, min_calls, max_calls)); - - return generation_prompt + (reasoning << p.content(p.until("[TOOL_CALLS]")) << tool_calls); - } - - // Content only parser - include_grammar = false; - return generation_prompt + (reasoning << p.content(p.rest())); - }); - - data.parser = parser.save(); - - if (include_grammar) { - data.grammar_lazy = has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; - - data.grammar = build_grammar([&](const common_grammar_builder & builder) { - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - auto schema = function.at("parameters"); - builder.resolve_refs(schema); - }); - if (has_response_format) { - auto schema = inputs.json_schema; - builder.resolve_refs(schema); - } - parser.build_grammar(builder, data.grammar_lazy); - }); - - data.grammar_triggers = { - { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "[TOOL_CALLS]" } - }; - } - - return data; -} - -static common_chat_params common_chat_params_init_gpt_oss(const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { - common_chat_params data; - - // Copy reasoning to the "thinking" field as expected by the gpt-oss template - auto adjusted_messages = json::array(); - for (auto msg : inputs.messages) { - if (msg.contains("reasoning_content") && msg.at("reasoning_content").is_string()) { - msg["thinking"] = msg.at("reasoning_content"); - if (msg.contains("tool_calls") && msg.at("tool_calls").is_array() && !msg.at("tool_calls").empty()) { - msg.erase("content"); - } - } - adjusted_messages.push_back(msg); - } - - auto prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override= */ adjusted_messages); - - // Check if we need to replace the return token with end token during - // inference and without generation prompt. For more details see: - // https://github.com/ggml-org/llama.cpp/issues/15417 - if (inputs.is_inference && !inputs.add_generation_prompt) { - static constexpr std::string_view return_token = "<|return|>"; - static constexpr std::string_view end_token = "<|end|>"; - if (size_t pos = prompt.rfind(return_token); pos != std::string::npos) { - prompt.replace(pos, return_token.length(), end_token); - } - } - - data.prompt = prompt; - data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; - data.supports_thinking = true; - - // These special tokens are required to parse properly, so we include them - // even if parse_tool_calls is false. - data.preserved_tokens = { - "<|channel|>", "<|constrain|>", "<|message|>", "<|start|>", "<|end|>", - }; - - auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); - auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); - auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); - auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; - - auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { - auto start = p.rule("start", p.literal("<|start|>assistant")); - auto end = p.rule("end", p.literal("<|end|>")); - auto content = p.rule("message-content", p.until("<|end|>")); - auto channel = p.literal("<|channel|>") + (p.literal("commentary") | p.literal("analysis")); - auto constrain_type = p.chars("[A-Za-z0-9_-]", 1, -1); - - // Occasionally, gpt-oss-20b will prefix channels with this commentary - auto stray_commentary = p.optional(p.literal("<|channel|>commentary") + p.optional(p.literal(" to=assistant"))); - auto start_analysis = stray_commentary + p.literal("<|channel|>analysis<|message|>"); - - if (extract_reasoning) { - p.rule("analysis", start_analysis + p.reasoning(content) + end); - } else { - p.rule("analysis", p.content(start_analysis + content + end)); - } - - auto analysis = p.ref("analysis"); - auto preamble = p.rule("preamble", p.literal("<|channel|>commentary<|message|>") + p.content(content) + end); - auto final_msg = p.rule("final", stray_commentary + p.literal("<|channel|>final<|message|>") + p.content(content)); - - // Consume any unsolicited tool calls, e.g. builtin functions - auto unsolicited = p.rule("unsolicited", p.atomic(p.optional(channel) + p.literal(" to=") + content + end)); - - auto any = p.rule("any", preamble | analysis); - - if (has_response_format) { - auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type); - auto response_format = p.rule("response-format", - p.literal("<|channel|>final") + constraint + p.literal("<|message|>") + - p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema))); - - return p.zero_or_more(start + analysis) + start + response_format; - } - - if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { - auto tool_choice = p.choice(); - - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - std::string name = function.at("name"); - const auto & params = function.at("parameters"); - - auto func_name = p.literal(" to=functions.") + p.tool_name(p.literal(name)); - auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type); - auto args = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", params)); - - // recipient in role header - // <|start|>assistant to=functions.NAME<|channel|>(commentary|analysis)[constraint]<|message|>ARGS - auto tool_in_role = p.tool(p.tool_open(func_name + channel + constraint + p.literal("<|message|>")) + args); - - // recipient in channel header - // <|channel|>(commentary|analysis) to=functions.NAME[constraint]<|message|>ARGS - auto tool_in_channel = p.tool(p.tool_open(channel + func_name + constraint + p.literal("<|message|>")) + args); - - tool_choice |= p.rule("tool-" + name, tool_in_role | tool_in_channel); - }); - - auto tool_call = p.trigger_rule("tool-call", tool_choice); - - if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) { - return p.zero_or_more(start + any) + start + tool_call; - } - - return p.zero_or_more(start + any) + start + (tool_call | final_msg); - } - - return p.zero_or_more(start + any) + start + (final_msg | unsolicited); - }); - - data.parser = parser.save(); - - if (include_grammar) { - data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); - data.grammar = build_grammar([&](const common_grammar_builder & builder) { - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - auto schema = function.at("parameters"); - builder.resolve_refs(schema); - }); - if (has_response_format) { - auto schema = inputs.json_schema; - builder.resolve_refs(schema); - } - parser.build_grammar(builder, data.grammar_lazy); - }); - - data.grammar_triggers = { - { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "^\\s+to$" }, - { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "^<\\|channel\\|>(?:commentary|analysis)\\s+to=functions$" }, - { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "<\\|start\\|>assistant(\\s+to)" }, - { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "<\\|start\\|>assistant(<\\|channel\\|>(?:commentary|analysis)\\s+to)" } - }; - } - - return data; -} - -static common_chat_params common_chat_params_init_gemma4(const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { - common_chat_params data; - - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); - - if (inputs.add_generation_prompt && string_ends_with(data.prompt, "\n")) { - // This may happen if the model generates content + tool_call, the - // template does not add the model's next turn and confuses the model - // from emitting its proper reasoning token sequence. - data.prompt += "<|turn>model\n"; - } - - data.format = COMMON_CHAT_FORMAT_PEG_GEMMA4; - data.supports_thinking = true; - data.thinking_start_tag = "<|channel>thought"; - data.thinking_end_tag = ""; - - data.preserved_tokens = { - "<|channel>", - "", - "<|tool_call>", - "", - "<|turn>", - }; - - auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); - auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); - auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); - auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; - - auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { - auto start = p.rule("start", p.prefix(inputs.generation_prompt, "<|channel>")); - - if (extract_reasoning) { - p.rule("thought", p.literal("<|channel>thought") + p.space() + p.reasoning(p.until("")) + p.literal("")); - } else { - p.rule("thought", p.content(p.literal("<|channel>thought") + p.space() + p.until("") + p.literal(""))); - } - - auto consume_empty_channels = p.gbnf(p.zero_or_more(p.literal("<|channel>") + p.negate(p.literal("thought"))), ""); - auto thought = (p.peek(p.literal("<|channel>")) + consume_empty_channels + p.ref("thought")) | p.negate(p.literal("<|channel>")); - - if (has_response_format) { - auto response_format = p.literal("```json") << - p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) << - p.literal("```"); - return start + p.optional(thought) + response_format; - } - - if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { - // Gemma4 tool calling syntax - // Rules should match traversal logic in gemma4_to_json() - p.rule("gemma4-string-content", p.until("<|\"|>")); - p.rule("gemma4-string", p.literal("<|\"|>") + p.ref("gemma4-string-content") + p.literal("<|\"|>")); - p.rule("gemma4-bool", p.json_bool()); - p.rule("gemma4-null", p.json_null()); - p.rule("gemma4-number", p.json_number()); - p.rule("gemma4-dict-key", p.rule("gemma4-dict-key-name", p.chars("[^:}]", 1, -1)) + p.literal(":")); - p.rule("gemma4-dict-kv", p.ref("gemma4-dict-key") + p.space() + p.ref("gemma4-value")); - p.rule("gemma4-dict", [&]() { - auto ws = p.space(); - auto member = p.ref("gemma4-dict-kv"); - auto members = p.sequence({member, p.zero_or_more(p.sequence({p.literal(","), ws, member}))}); - return p.sequence({ - p.literal("{"), ws, - p.choice({p.literal("}"), p.sequence({members, ws, p.literal("}")})}) - }); - }); - p.rule("gemma4-array", [&]() { - auto ws = p.space(); - auto value = p.ref("gemma4-value"); - auto elements = p.sequence({value, p.zero_or_more(p.sequence({p.literal(","), ws, value}))}); - return p.sequence({ - p.literal("["), ws, - p.choice({p.literal("]"), p.sequence({elements, ws, p.literal("]")})}) - }); - }); - p.rule("gemma4-value", [&]() { - return p.choice({ - p.ref("gemma4-string"), p.ref("gemma4-dict"), p.ref("gemma4-array"), - p.ref("gemma4-number"), p.ref("gemma4-bool"), p.ref("gemma4-null") - }); - }); - - auto tool_choice = p.choice(); - - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - std::string name = function.at("name"); - // TODO @aldehir : need to extend json-schema-to-grammar to produce more than JSON rules - // const auto & params = function.at("parameters"); - - tool_choice |= p.rule("tool-" + name, p.tool(p.sequence({ - p.tool_open(p.tool_name(p.literal(name)) + p.peek(p.literal("{"))), - p.tool_args(p.ref("gemma4-dict")), - }))); - }); - - auto tool_call = p.trigger_rule("tool-call", p.repeat( - "<|tool_call>call:" + tool_choice + "", - /* min = */ inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0, - /* max = */ inputs.parallel_tool_calls ? -1 : 1 - )); - - auto scan_to_toolcall = p.rule("scan-to-toolcall", p.until("<|tool_call>")); - auto content = p.rule("content", p.content(p.until_one_of({"<|channel>", "", "<|tool_call>"}))); - auto message = p.rule("message", thought + content); - return start + p.zero_or_more(message) + scan_to_toolcall + tool_call; - } - - // Gemma 4 may emit an extra <|channel>thought\n at the end of the content. It may - // also emit a single trailing token. Consume all complete reasoning blocks and - // then stop at the first unmatched token. - auto content = p.rule("content", p.content(p.until_one_of({"<|channel>", ""}))); - auto message = p.rule("message", thought + content); - return start + p.one_or_more(message); - }); - - data.parser = parser.save(); - - if (include_grammar) { - data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); - data.grammar = build_grammar([&](const common_grammar_builder & builder) { - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - auto schema = function.at("parameters"); - builder.resolve_refs(schema); - }); - if (has_response_format) { - auto schema = inputs.json_schema; - builder.resolve_refs(schema); - } - parser.build_grammar(builder, data.grammar_lazy); - }); - - data.grammar_triggers = { - { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "<|tool_call>" }, - }; - } - - return data; -} - -// Functionary v3.2 - uses recipient-based format: >>>recipient\n{content} -static common_chat_params common_chat_params_init_functionary_v3_2(const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { - common_chat_params data; - - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); - data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; - data.preserved_tokens = { - ">>>all", - }; - - auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); - auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; - - auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { - // Functionary v3.2 format: - // - Normal content: >>>all\n{content} - // - Tool calls: >>>function_name\n{json_args} - // Generation prompt ends with ">>>" so model outputs recipient immediately - - // Build content parser for >>>all\n{content} - // When tools are present, content stops before the next ">>>" (tool call) - // When no tools, content goes until end - auto content_until_tool = p.literal("all\n") + p.content(p.until(">>>")); - auto content_until_end = p.literal("all\n") + p.content(p.rest()); - auto generation_prompt = p.literal(inputs.generation_prompt); - - // If no tools or tool_choice is NONE, just parse content - if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { - // When no tools, just match the prefix and capture everything after - return generation_prompt + content_until_end + p.end(); - } - - // Build tool call parsers for each available function - auto tool_choice = p.choice(); - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - std::string name = function.at("name"); - const auto & schema = function.at("parameters"); - - // Tool format: >>>function_name\n{json_args} - auto tool_parser = p.tool( - p.tool_open(p.tool_name(p.literal(name)) + p.literal("\n")) + - p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)) - ); - - tool_choice |= p.rule("tool-" + name, tool_parser); - }); - - auto content_only = content_until_end; - auto tools_only = p.trigger_rule("tools", p.one_or_more(tool_choice)); - auto content_and_tools = content_until_tool + tools_only; - - auto ret = p.eps(); - if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) { - if (inputs.parallel_tool_calls) { - ret = p.choice({ content_and_tools, tools_only }) + p.end(); - } else { - ret = p.choice({ content_until_tool + tool_choice, tools_only }) + p.end(); - } - } else if (inputs.parallel_tool_calls) { - ret = p.choice({ content_and_tools, content_only, tools_only }) + p.end(); - } else { - auto content_and_tool = content_until_tool + tool_choice; - ret = p.choice({ content_and_tool, content_only, tool_choice }) + p.end(); - } - return generation_prompt + ret; - }); - - data.parser = parser.save(); - - if (include_grammar) { - data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; - - data.grammar = build_grammar([&](const common_grammar_builder & builder) { - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - auto schema = function.at("parameters"); - builder.resolve_refs(schema); - }); - parser.build_grammar(builder, data.grammar_lazy); - }); - - // Grammar trigger for when the model starts outputting a tool call - // (after the initial ">>>" in the generation prompt but recipient other than "all") - data.grammar_triggers = { - { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, ">>>(?!all)" } - }; - } - - return data; -} - -// Kimi K2 Thinking - uses unique tool call ID format: functions.: -// The ID contains both the function name and an incrementing counter -static common_chat_params common_chat_params_init_kimi_k2(const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { - common_chat_params data; - - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); - data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; - data.supports_thinking = true; - data.preserved_tokens = { - "<|tool_calls_section_begin|>", - "<|tool_calls_section_end|>", - "<|tool_call_begin|>", - "<|tool_call_argument_begin|>", - "<|tool_call_end|>", - "", - "", - }; - - auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); - auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; - auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; - - const std::string SECTION_BEGIN = "<|tool_calls_section_begin|>"; - const std::string SECTION_END = "<|tool_calls_section_end|>"; - const std::string CALL_BEGIN = "<|tool_call_begin|>"; - const std::string ARGS_BEGIN = "<|tool_call_argument_begin|>"; - const std::string CALL_END = "<|tool_call_end|>"; - - const std::string THINK_START = ""; - const std::string THINK_END = ""; - - data.thinking_start_tag = THINK_START; - data.thinking_end_tag = THINK_END; - - auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { - // Kimi K2 Thinking format: - // - Reasoning: {reasoning} - // - Content: text after reasoning - // - Tool calls section: - // <|tool_calls_section_begin|> - // <|tool_call_begin|>functions.:<|tool_call_argument_begin|>{json_args}<|tool_call_end|> - // ... - // <|tool_calls_section_end|> - // The ID format is: functions.: where counter is 0, 1, 2, ... - - // Tool call markers - auto end = p.end(); - - // Note: this model is CRAZY. It can diverge from its supposed tool calling pattern in so many ways it's not funny. - // For example, it can call tools at the end of reasoning without closing reasoning... - auto reasoning = extract_reasoning ? p.optional(THINK_START + p.reasoning( - p.until_one_of({ THINK_END, "<|tool_calls_section_begin|>", "<|tool_call_begin|>" })) + - p.optional(p.literal(THINK_END))) : p.eps(); - auto generation_prompt = p.prefix(inputs.generation_prompt, THINK_START); - - - // Content only parser (no tools) - if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { - return generation_prompt + reasoning + p.content(p.rest()) + end; - } - - // Build tool call parsers for each available function - // The ID format is: functions.: - // We need to match: functions.: - auto tool_choice = p.choice(); - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - std::string name = function.at("name"); - const auto & schema = function.at("parameters"); - - // Match: functions.: - // Capture the full call id (functions.:) using tool_id tag - auto tool_id = p.tool_id(p.literal("functions.") + p.tool_name(p.literal(name)) + p.literal(":") + p.chars("[0-9]", 1, -1)); - auto tool_parser = p.tool( - p.tool_open(tool_id + p.literal(ARGS_BEGIN)) + - p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)) + - p.tool_close(p.optional((p.literal(CALL_END)))) - ); - - tool_choice |= p.rule("tool-" + name, tool_parser); - }); - - // Tool calls section: <|tool_calls_section_begin|> tool_calls <|tool_calls_section_end|> - auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0; - auto max_calls = inputs.parallel_tool_calls ? -1 : 1; - // Use trigger_rule so grammar generator knows where to start generating rules - auto tool_calls = p.rule("tool-calls", - p.optional(p.literal(SECTION_BEGIN)) + - p.trigger_rule("tool-call", p.repeat(CALL_BEGIN + tool_choice, min_calls, max_calls) + - p.optional(p.literal(SECTION_END))) - ); - - auto content_before_tools = p.content(p.until_one_of({ SECTION_BEGIN, CALL_BEGIN })); - - return generation_prompt + reasoning + content_before_tools + tool_calls + end; - }); - - data.parser = parser.save(); - - if (include_grammar) { - data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; - data.grammar = build_grammar([&](const common_grammar_builder & builder) { - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - auto schema = function.at("parameters"); - builder.resolve_refs(schema); - }); - parser.build_grammar(builder, data.grammar_lazy); - }); - - data.grammar_triggers = { - { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "<|tool_call_begin|>" } - }; - } - - return data; -} - -// LFM2 format: uses <|tool_list_start|>[...]<|tool_list_end|> in system prompt -// and <|tool_call_start|>[name(arg="val")]<|tool_call_end|> for tool calls. -// - Reasoning: {reasoning} (optional) -// - Content: text before a tool call (optional) -// - Tool calls: Python-style, e.g. [function_name(arg1="value1", arg2="value2")] -// Tool calls can appear multiple times (parallel tool calls supported) -static common_chat_params common_chat_params_init_lfm2(const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { - common_chat_params data; - - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); - data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; - data.supports_thinking = true; - data.preserved_tokens = { - "<|tool_list_start|>", - "<|tool_list_end|>", - "<|tool_call_start|>", - "<|tool_call_end|>", - "", - "", - }; - - auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); - auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; - auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; - - const std::string TOOL_CALL_START = "<|tool_call_start|>"; - const std::string TOOL_CALL_END = "<|tool_call_end|>"; - const std::string THINK_START = ""; - const std::string THINK_END = ""; - - data.thinking_start_tag = THINK_START; - data.thinking_end_tag = THINK_END; - - auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { - auto generation_prompt = p.prefix(inputs.generation_prompt, THINK_START); - auto end = p.end(); - - auto reasoning = p.eps(); - if (extract_reasoning && inputs.enable_thinking) { - reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END); - } - - if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { - return generation_prompt + reasoning + p.content(p.rest()) + end; - } - auto tool_calls = p.rule("tool-calls", - p.trigger_rule("tool-call", - p.literal(TOOL_CALL_START) + - p.python_style_tool_calls(inputs.tools, inputs.parallel_tool_calls) + - p.literal(TOOL_CALL_END) - ) - ); - - auto content = p.content(p.until(TOOL_CALL_START)); - - return generation_prompt + reasoning + content + tool_calls + end; - }); - - data.parser = parser.save(); - - if (include_grammar) { - data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; - data.grammar = build_grammar([&](const common_grammar_builder & builder) { - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - auto schema = function.at("parameters"); - builder.resolve_refs(schema); - }); - parser.build_grammar(builder, data.grammar_lazy); - }); - - data.grammar_triggers = { - { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, TOOL_CALL_START } - }; - } - return data; -} - -// LFM2.5 format: uses plain "List of tools: [...]" in system prompt, no wrapper tokens. -// Tool calls are bare [name(arg="val")], though model may optionally emit <|tool_call_start|>. -// - Reasoning: {reasoning} (optional) -// - Content: text before a tool call (optional) -// - Tool calls: Python-style, e.g. [function_name(arg1="value1", arg2="value2")] -// Tool calls can appear multiple times (parallel tool calls supported) -static common_chat_params common_chat_params_init_lfm2_5(const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { - common_chat_params data; - - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); - data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; - data.supports_thinking = true; - data.preserved_tokens = { - "<|tool_call_start|>", - "<|tool_call_end|>", - "", - "", - }; - - auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); - auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; - auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; - - const std::string THINK_START = ""; - const std::string THINK_END = ""; - - data.thinking_start_tag = THINK_START; - data.thinking_end_tag = THINK_END; - - auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { - auto generation_prompt = p.prefix(inputs.generation_prompt, THINK_START); - auto end = p.end(); - - auto reasoning = p.eps(); - if (extract_reasoning && inputs.enable_thinking) { - reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END); - } - - if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { - return generation_prompt + reasoning + p.content(p.rest()) + end; - } - - auto tool_calls = p.rule("tool-calls", - p.trigger_rule("tool-call", - p.python_style_tool_calls(inputs.tools, inputs.parallel_tool_calls) - ) - ); - - auto content = p.content(p.until_one_of({"<|tool_call_start|>", "["})); - auto maybe_start = p.optional(p.literal("<|tool_call_start|>")); - return generation_prompt + reasoning + content + maybe_start + tool_calls + end; - }); - - data.parser = parser.save(); - - if (include_grammar) { - data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; - data.grammar = build_grammar([&](const common_grammar_builder & builder) { - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - auto schema = function.at("parameters"); - builder.resolve_refs(schema); - }); - parser.build_grammar(builder, data.grammar_lazy); - }); - foreach_function(inputs.tools, [&](const json & tool) { - const std::string name = tool.at("function").at("name"); - data.grammar_triggers.push_back({ COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "[" + name + "(" }); - }); - } - - return data; -} - -static common_chat_params common_chat_params_init_gigachat_v3( - const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { - - common_chat_params data; - - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); - data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; - data.supports_thinking = false; - data.preserved_tokens = { - "<|message_sep|>\n\n", - "<|role_sep|>\n", - }; - - auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); - auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; - const auto *tool_call_start_prefix = "<|message_sep|>\n\nfunction call<|role_sep|>\n"; - - auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { - auto ret = p.eps(); - if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { - // Build a choice of all available tools - auto tool_choice = p.choice(); - for (const auto & tool : inputs.tools) { - const auto & function = tool.at("function"); - std::string name = function.at("name"); - const auto & schema = function.at("parameters"); - - auto tool_name = p.json_member("name", "\"" + p.tool_name(p.literal(name)) + "\""); - auto tool_args = p.json_member("arguments", p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema))); - - auto tool_open = p.tool_open(p.literal("{") << tool_name); - - tool_choice |= p.rule("tool-" + name, tool_open << "," << tool_args << "}"); - } - - // Define the tool call structure - auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0; - auto max_calls = 1; // parallel toolcalls are not supported - auto tool_call = p.rule("tool-call", p.literal(tool_call_start_prefix) + tool_choice); - auto tool_calls = p.trigger_rule("tool-call-root", p.repeat(tool_call, /* min = */ min_calls, /* max = */ max_calls)); - - ret = p.content(p.until("<|message_sep|>\n\n")) << tool_calls; - } else { - // Content only parser - include_grammar = false; - ret = p.content(p.rest()); - } - - return p.literal(inputs.generation_prompt) + ret; - }); - - data.parser = parser.save(); - - if (include_grammar) { - data.grammar_lazy = has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; - - data.grammar = build_grammar([&](const common_grammar_builder & builder) { - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - auto schema = function.at("parameters"); - builder.resolve_refs(schema); - }); - parser.build_grammar(builder, data.grammar_lazy); - }); - - data.grammar_triggers = { - {COMMON_GRAMMAR_TRIGGER_TYPE_WORD, tool_call_start_prefix} - }; - } - return data; -} - -static common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { - common_chat_params data; - - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); - data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; - data.supports_thinking = true; - data.thinking_start_tag = ""; - data.thinking_end_tag = ""; - data.preserved_tokens = { - "｜DSML｜", - "", - "", - }; - - auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); - auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); - auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; - auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); - - const std::string DSML = "｜DSML｜"; - const std::string THINK_START = ""; - const std::string THINK_END = ""; - const std::string FC_START = "<" + DSML + "function_calls>"; - const std::string FC_END = ""; - const std::string INVOKE_START = "<" + DSML + "invoke"; - const std::string INVOKE_END = ""; - const std::string PARAM_START = "<" + DSML + "parameter"; - const std::string PARAM_END = ""; - - auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { - auto generation_prompt = p.prefix(inputs.generation_prompt, THINK_START); - auto end = p.end(); - - auto reasoning = p.eps(); - if (extract_reasoning && inputs.enable_thinking) { - reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END); - } else if (extract_reasoning) { - // Thinking disabled but reasoning extraction requested: the generation prompt - // contains an empty pair that must still be consumed. - reasoning = p.optional(p.literal(THINK_START) + p.until(THINK_END) + p.literal(THINK_END)); - } - - if (has_response_format) { - auto response_format = p.rule("response-format", - p.literal("```json") + p.space() + - p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) + - p.space() + p.literal("```")); - return generation_prompt + reasoning + response_format + end; - } - - if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { - return generation_prompt + reasoning + p.content(p.rest()) + end; - } - - auto tool_choice = p.choice(); - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - std::string name = function.at("name"); - auto params = function.contains("parameters") ? function.at("parameters") : json::object(); - const auto & props = params.contains("properties") ? params.at("properties") : json::object(); - - std::set required; - if (params.contains("required")) { - params.at("required").get_to(required); - } - - auto schema_info = common_schema_info(); - schema_info.resolve_refs(params); - - std::vector required_parsers; - std::vector optional_parsers; - for (const auto & [param_name, param_schema] : props.items()) { - bool is_required = required.find(param_name) != required.end(); - bool is_string = schema_info.resolves_to_string(param_schema); - - auto arg = p.tool_arg( - p.tool_arg_open( - p.literal(PARAM_START + " name=\"") + - p.tool_arg_name(p.literal(param_name)) + - p.literal("\" string=\"" + std::string(is_string ? "true" : "false") + "\">")) + - (is_string - ? p.tool_arg_string_value(p.until(PARAM_END)) - : p.tool_arg_json_value(p.schema(p.json(), - "tool-" + name + "-arg-" + param_name + "-schema", - param_schema, false))) + - p.tool_arg_close(p.literal(PARAM_END))); - - auto named_arg = p.rule("tool-" + name + "-arg-" + param_name, arg); - if (is_required) { - required_parsers.push_back(named_arg); - } else { - optional_parsers.push_back(named_arg); - } - } - - common_peg_parser args_seq = p.eps(); - for (size_t i = 0; i < required_parsers.size(); i++) { - if (i > 0) { - args_seq = args_seq + p.space(); - } - args_seq = args_seq + required_parsers[i]; - } - - if (!optional_parsers.empty()) { - common_peg_parser any_opt = p.choice(); - for (const auto & opt : optional_parsers) { - any_opt |= opt; - } - args_seq = args_seq + p.repeat(p.space() + any_opt, 0, -1); - } - - common_peg_parser invoke_body = args_seq; - auto func_parser = p.tool( - p.tool_open(p.literal(INVOKE_START + " name=\"") + - p.tool_name(p.literal(name)) + p.literal("\">\n")) + - invoke_body + p.space() + - p.tool_close(p.literal(INVOKE_END))); - - tool_choice |= p.rule("tool-" + name, func_parser); - }); - - auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; - - common_peg_parser tool_calls = p.eps(); - if (inputs.parallel_tool_calls) { - tool_calls = p.trigger_rule("tool-call", - p.literal(FC_START) + p.space() + tool_choice + - p.zero_or_more(p.space() + tool_choice) + p.space() + p.literal(FC_END)); - } else { - tool_calls = p.trigger_rule("tool-call", - p.literal(FC_START) + p.space() + tool_choice + p.space() + p.literal(FC_END)); - } - - if (!require_tools) { - tool_calls = p.optional(tool_calls); - } - - auto content_before_tools = p.content(p.until(FC_START)); - return generation_prompt + reasoning + content_before_tools + tool_calls + end; - }); - - data.parser = parser.save(); - - if (include_grammar) { - data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); - data.grammar = build_grammar([&](const common_grammar_builder & builder) { - foreach_function(inputs.tools, [&](const json & tool) { - const auto & function = tool.at("function"); - auto schema = function.contains("parameters") ? function.at("parameters") : json::object(); - builder.resolve_refs(schema); - }); - if (has_response_format) { - auto schema = inputs.json_schema; - builder.resolve_refs(schema); - } - parser.build_grammar(builder, data.grammar_lazy); - }); - - data.grammar_triggers = { - { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, FC_START }, - }; - } - - return data; -} - -namespace workaround { - -static void map_developer_role_to_system(json & messages) { - for (auto & message : messages) { - if (message.contains("role")) { - if (message["role"] == "developer") { - message["role"] = "system"; - } - } - } -} - - -// if first message is system and template does not support it, merge it with next message -static void system_message_not_supported(json & messages) { - if (!messages.empty() && messages.front().at("role") == "system") { - if (messages.size() > 1) { - LOG_DBG("Merging system prompt into next message\n"); - auto & first_msg = messages.front(); - auto & second_msg = messages[1]; - second_msg["content"] = first_msg.at("content").get() - + "\n" + second_msg.at("content").get(); - messages.erase(messages.begin()); - } else { - LOG_WRN("Removing system prompt due to template not supporting system role\n"); - messages.erase(messages.begin()); - } - } -} - -static void requires_non_null_content(json & messages) { - GGML_ASSERT(messages.is_array()); - for (auto & message : messages) { - if (message.contains("tool_calls") && !message.contains("content")) { - message["content"] = ""; - } - } -} - -// Gemma4 uses a custom tool_responses field instead of role:tool messages. -// -// This will transform a sequence of messages: -// assistant(tool_call+) -> tool+ -> assistant(content) -// -// Into a single assistant message containing a tool_responses field: -// assistant(content + tool_call + tool_responses) -// -// This is necessary for the Gemma4 chat template to properly format the prompt. -// See https://ai.google.dev/gemma/docs/core/prompt-formatting-gemma4 -struct gemma4_model_turn_builder { - json & messages; - size_t pos; - json tool_calls = json::array(); - json tool_responses = json::array(); - json content; - json reasoning_content; - - gemma4_model_turn_builder(json & msgs, size_t pos) : messages(msgs), pos(pos) {} - - void collect() { - // Collect the first assistant message - auto & msg = messages[pos]; - if (msg.contains("reasoning_content") && msg.at("reasoning_content").is_string()) { - // According to the prompt formatting guide, we need to preserve reasoning_content - // between function calls. The current chat templates do not support this, but we will do it anyway. - reasoning_content = msg.at("reasoning_content"); - } - for (auto & tc : msg.at("tool_calls")) { - tool_calls.push_back(tc); - } - pos++; - - // Collect tool call results - while (pos < messages.size() && messages[pos].value("role", "") == "tool") { - collect_result(messages[pos]); - pos++; - } - - // Check if the next assistant message is the final message - if (pos < messages.size() && messages[pos].value("role", "") == "assistant") { - auto & next = messages[pos]; - if (!has_tool_calls(next) && has_content(next)) { - content = next.at("content"); - pos++; - } - } - } - - void collect_result(const json & curr) { - json response; - if (curr.contains("content")) { - const auto & content = curr.at("content"); - if (content.is_string()) { - // Try to parse the content as JSON; fall back to raw string - try { - response = json::parse(content.get()); - } catch (...) { - response = content; - } - } else { - response = content; - } - } - - std::string name; - - // Match name with corresponding tool call - size_t idx = tool_responses.size(); - if (idx < tool_calls.size()) { - auto & tc = tool_calls[idx]; - if (tc.contains("function")) { - name = tc.at("function").value("name", ""); - } - } - - // Fallback to the tool call id - if (name.empty()) { - name = curr.value("tool_call_id", ""); - } - - tool_responses.push_back({{"name", name}, {"response", response}}); - } - - json build() { - collect(); - - json msg = { - {"role", "assistant"}, - {"tool_calls", tool_calls}, - }; - if (!tool_responses.empty()) { - msg["tool_responses"] = tool_responses; - } - if (!content.is_null()) { - msg["content"] = content; - } - if (!reasoning_content.is_null()) { - msg["reasoning_content"] = reasoning_content; - } - return msg; - } - - static bool has_content(const json & msg) { - if (!msg.contains("content") || msg.at("content").is_null()) { - return false; - } - const auto & content = msg.at("content"); - if (content.is_string() && !content.get().empty()) { - return true; - } - if (content.is_array() && !content.empty()) { - return true; - } - return false; - } - - static bool has_tool_calls(const json & msg) { - return msg.contains("tool_calls") && msg.at("tool_calls").is_array() && !msg.at("tool_calls").empty(); - } -}; - -static void convert_tool_responses_gemma4(json & messages) { - json result = json::array(); - size_t i = 0; - - while (i < messages.size()) { - auto & msg = messages[i]; - - if (msg.value("role", "") != "assistant" || !msg.contains("tool_calls") || - !msg.at("tool_calls").is_array() || msg.at("tool_calls").empty()) { - result.push_back(msg); - i++; - continue; - } - - gemma4_model_turn_builder builder(messages, i); - result.push_back(builder.build()); - i = builder.pos; - } - - messages = result; -} - -static void func_args_not_string(json & messages) { - GGML_ASSERT(messages.is_array()); - for (auto & message : messages) { - if (message.contains("tool_calls")) { - for (auto & tool_call : message["tool_calls"]) { - if (tool_call.contains("function") && tool_call["function"].contains("arguments")) { - auto & args = tool_call["function"]["arguments"]; - if (args.is_string()) { - try { - args = json::parse(args.get()); - } catch (const std::exception & e) { - throw std::runtime_error("Failed to parse tool call arguments as JSON: " + std::string(e.what())); - } - } - } - } - } - } -} - -} - -static json common_chat_extra_context() { - json ctx = json::object(); - std::chrono::system_clock::time_point now = std::chrono::system_clock::now(); - std::string datetime_str = format_time(now, "%b %d %Y"); - std::string date_str = format_time(now, "%d %b %Y"); - ctx["datetime"] = datetime_str; - ctx["date_string"] = date_str; - return ctx; -} - -std::optional common_chat_try_specialized_template( - const common_chat_template & tmpl, - const std::string & src, - autoparser::generation_params & params) { - // Ministral/Mistral Large 3 - uses special reasoning structure fixes, can't use autoparser - // Note: Mistral Small 3.2 uses [CALL_ID] which Ministral doesn't have, so we can distinguish them - if (src.find("[SYSTEM_PROMPT]") != std::string::npos && src.find("[TOOL_CALLS]") != std::string::npos && - src.find("[ARGS]") != std::string::npos && src.find("[CALL_ID]") == std::string::npos) { - LOG_DBG("Using specialized template: Ministral/Magistral Large 3\n"); - return common_chat_params_init_ministral_3(tmpl, params); - } - - // GPT-OSS - has unique channel-based structure that needs dedicated handler - if (src.find("<|channel|>") != std::string::npos) { - LOG_DBG("Using specialized template: GPT-OSS\n"); - return common_chat_params_init_gpt_oss(tmpl, params); - } - - // Functionary v3.2 - uses recipient-based format with >>>recipient\n{content} - // Detection: template has ">>>all" for content and ">>>" prefix for tool calls - if (src.find(">>>all") != std::string::npos && src.find(">>>${recipient}") != std::string::npos) { - LOG_DBG("Using specialized template: Functionary v3.2\n"); - return common_chat_params_init_functionary_v3_2(tmpl, params); - } - - // Kimi K2 Thinking - uses unique tool call ID format: functions.: - // Detection: template has "<|tool_calls_section_begin|>" and "functions." prefix in tool call IDs - if (src.find("<|tool_calls_section_begin|>") != std::string::npos && - src.find("<|tool_call_begin|>") != std::string::npos) { - LOG_DBG("Using specialized template: Kimi K2 Thinking\n"); - return common_chat_params_init_kimi_k2(tmpl, params); - } - - if (is_lfm2_template(src)) { - LOG_DBG("Using specialized template: LFM2\n"); - return common_chat_params_init_lfm2(tmpl, params); - } - - // LFM2.5 format detection: template uses plain "List of tools: [...]" with no special tokens - if (src.find("List of tools: [") != std::string::npos && - src.find("<|tool_list_start|>") == std::string::npos) { - LOG_DBG("Using specialized template: LFM2.5\n"); - return common_chat_params_init_lfm2_5(tmpl, params); - } - - // GigaChatV3 format detection - if (src.find("<|role_sep|>") != std::string::npos && - src.find("<|message_sep|>") != std::string::npos && - src.find("<|function_call|>") == std::string::npos) { - LOG_DBG("Using specialized template: GigaChatV3\n"); - return common_chat_params_init_gigachat_v3(tmpl, params); - } - - // DeepSeek V3.2 format detection: template defines dsml_token and uses it for tool calls. - // The template source contains the token as a variable assignment, not as a literal in markup. - if (src.find("dsml_token") != std::string::npos && - src.find("function_calls") != std::string::npos && - src.find("DSML") != std::string::npos) { - LOG_DBG("Using specialized template: DeepSeek V3.2\n"); - return common_chat_params_init_deepseek_v3_2(tmpl, params); - } - - // Gemma4 format detection - if (src.find("'<|tool_call>call:'") != std::string::npos) { - if (src.find("{#- OpenAI Chat Completions:") == std::string::npos) { - // apply workarounds if using the older gemma4 templates - LOG_WRN("%s: detected an outdated gemma4 chat template, applying compatibility workarounds. " - "Consider updating to the official template.\n", __func__); - workaround::convert_tool_responses_gemma4(params.messages); - } - return common_chat_params_init_gemma4(tmpl, params); - } - - return std::nullopt; -} - -static common_chat_params common_chat_templates_apply_jinja(const struct common_chat_templates * tmpls, - const struct common_chat_templates_inputs & inputs) { - autoparser::generation_params params; - params.tools = common_chat_tools_to_json_oaicompat(inputs.tools); - const auto & tmpl = - params.tools.is_array() && tmpls->template_tool_use ? *tmpls->template_tool_use : *tmpls->template_default; - const auto & src = tmpl.source(); - const auto & caps = tmpl.original_caps(); - params.messages = render_message_to_json(inputs.messages, tmpl.original_caps()); - params.tool_choice = inputs.tool_choice; - params.reasoning_format = inputs.reasoning_format; - params.enable_thinking = inputs.enable_thinking; - params.grammar = inputs.grammar; - params.now = inputs.now; - params.add_bos = tmpls->add_bos; - params.add_eos = tmpls->add_eos; - - if (src.find("<|channel|>") == std::string::npos) { - // map developer to system for all models except for GPT-OSS - workaround::map_developer_role_to_system(params.messages); - } - - if (!tmpl.original_caps().supports_system_role) { - workaround::system_message_not_supported(params.messages); - } - - if (tmpl.original_caps().supports_tool_calls) { - // some templates will require the content field in tool call messages - // to still be non-null, this puts an empty string everywhere where the - // content field is null - workaround::requires_non_null_content(params.messages); - } - - if (tmpl.original_caps().supports_object_arguments) { - workaround::func_args_not_string(params.messages); - } - - params.add_generation_prompt = false; - std::string no_gen_prompt = common_chat_template_direct_apply_impl(tmpl, params); - params.add_generation_prompt = true; - std::string gen_prompt = common_chat_template_direct_apply_impl(tmpl, params); - auto diff = calculate_diff_split(no_gen_prompt, gen_prompt); - params.generation_prompt = diff.right + diff.suffix; - - params.add_generation_prompt = inputs.add_generation_prompt; - - params.extra_context = common_chat_extra_context(); - for (auto el : inputs.chat_template_kwargs) { - params.extra_context[el.first] = json::parse(el.second); - } - - if (!inputs.json_schema.empty()) { - params.json_schema = json::parse(inputs.json_schema); - } - - params.parallel_tool_calls = inputs.parallel_tool_calls; - - if (params.tools.is_array()) { - if (params.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE && !params.grammar.empty()) { - throw std::runtime_error("Cannot specify grammar with tools"); - } - if (caps.supports_tool_calls && !caps.supports_tools) { - LOG_WRN( - "Template supports tool calls but does not natively describe tools. The fallback behaviour used may " - "produce bad results, inspect prompt w/ --verbose & consider overriding the template.\n"); - } - } - - if (inputs.force_pure_content) { - LOG_WRN("Forcing pure content template, will not render reasoning or tools separately."); - // Create the result structure - common_chat_params data; - auto params_copy = params; - params_copy.reasoning_format = COMMON_REASONING_FORMAT_NONE; - data.prompt = common_chat_template_direct_apply_impl(tmpl, params_copy); - data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; - data.generation_prompt = params.generation_prompt; - auto parser = build_chat_peg_parser([¶ms](common_chat_peg_builder &p) { - return p.prefix(params.generation_prompt) << p.content(p.rest()); - }); - data.parser = parser.save(); - return data; - } - - if (auto result = common_chat_try_specialized_template(tmpl, src, params)) { - result->generation_prompt = params.generation_prompt; - return *result; - } - - try { - LOG_DBG("%s: using differential autoparser\n", __func__); - struct autoparser::autoparser autoparser; - autoparser.analyze_template(tmpl); - auto auto_params = autoparser::peg_generator::generate_parser(tmpl, params, autoparser); - auto_params.supports_thinking = autoparser.reasoning.mode != autoparser::reasoning_mode::NONE; - if (auto_params.supports_thinking) { - auto_params.thinking_start_tag = autoparser.reasoning.start; - auto_params.thinking_end_tag = autoparser.reasoning.end; - } - auto_params.generation_prompt = params.generation_prompt; - common_peg_arena arena; - arena.load(auto_params.parser); - LOG_DBG("%s: generated parser:\n%s\n\nparser generation prompt: %s\n", __func__, arena.dump(arena.root()).c_str(), auto_params.generation_prompt.c_str()); - return auto_params; - } catch (const std::exception & e) { - throw std::invalid_argument(std::string("Unable to generate parser for this template. Automatic parser generation failed: ") + e.what()); - } -} - -// Legacy template route (adhoc C++ implementation of known templates), forward to llama_chat_apply_template. -static common_chat_params common_chat_templates_apply_legacy(const struct common_chat_templates * tmpls, - const struct common_chat_templates_inputs & inputs) { - size_t alloc_size = 0; - std::vector chat; - std::vector contents; - - for (const auto & msg : inputs.messages) { - auto content = msg.content; - for (const auto & part : msg.content_parts) { - if (part.type != "text" && part.type != "media_marker") { - LOG_WRN("Ignoring non-text content part: %s\n", part.type.c_str()); - continue; - } - if (!content.empty()) { - content += "\n"; - ; - } - content += part.text; - } - contents.emplace_back(std::move(content)); - } - for (size_t i = 0; i < contents.size(); ++i) { - const auto & msg = inputs.messages[i]; - const auto & content = contents[i]; - chat.push_back({ msg.role.c_str(), content.c_str() }); - size_t msg_size = msg.role.size() + content.size(); - alloc_size += msg_size + (msg_size / 4); // == msg_size * 1.25 but avoiding float ops - } - - std::vector buf(alloc_size); - - // run the first time to get the total output length - const auto & src = tmpls->template_default->source(); - int32_t res = llama_chat_apply_template(src.c_str(), chat.data(), chat.size(), inputs.add_generation_prompt, - buf.data(), buf.size()); - - // error: chat template is not supported - if (res < 0) { - // if the custom "tmpl" is not supported, we throw an error - // this is a bit redundant (for good), since we're not sure if user validated the custom template with llama_chat_verify_template() - throw std::runtime_error("this custom template is not supported, try using --jinja"); - } - - // if it turns out that our buffer is too small, we resize it - if ((size_t) res > buf.size()) { - buf.resize(res); - res = llama_chat_apply_template(src.c_str(), chat.data(), chat.size(), inputs.add_generation_prompt, buf.data(), - buf.size()); - } - - // for safety, we check the result again - if (res < 0 || (size_t) res > buf.size()) { - throw std::runtime_error("failed to apply chat template, try using --jinja"); - } - - common_chat_params params; - params.prompt = std::string(buf.data(), res); - if (!inputs.json_schema.empty()) { - params.grammar = json_schema_to_grammar(json::parse(inputs.json_schema)); - } else { - params.grammar = inputs.grammar; - } - return params; -} - -common_chat_params common_chat_templates_apply(const struct common_chat_templates * tmpls, - const struct common_chat_templates_inputs & inputs) { - GGML_ASSERT(tmpls != nullptr); - return inputs.use_jinja ? common_chat_templates_apply_jinja(tmpls, inputs) : - common_chat_templates_apply_legacy(tmpls, inputs); -} - -common_chat_msg common_chat_parse(const std::string & input, - bool is_partial, - const common_chat_parser_params & params) { - return common_chat_peg_parse(params.parser, input, is_partial, params); -} - -common_chat_msg common_chat_peg_parse(const common_peg_arena & src_parser, - const std::string & input, - bool is_partial, - const common_chat_parser_params & params) { - const common_peg_arena & parser = src_parser.empty() ? - build_chat_peg_parser([](common_chat_peg_builder & p) { return p.content(p.rest()) + p.end(); }) : - src_parser; - - if (src_parser.empty()) { - LOG_DBG("No parser definition detected, assuming pure content parser."); - } - - const std::string effective_input = params.generation_prompt.empty() - ? input - : params.generation_prompt + input; - - //LOG_DBG("Parsing PEG input with format %s: %s\n", common_chat_format_name(params.format), effective_input.c_str()); - - common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_LENIENT; - if (params.debug) { - flags |= COMMON_PEG_PARSE_FLAG_DEBUG; - } - - common_peg_parse_context ctx(effective_input, flags); - auto result = parser.parse(ctx); - - if (result.fail()) { - // During partial parsing, return partial results if any AST nodes were captured - // This allows streaming to work correctly for formats like FUNC_MARKDOWN_CODE_BLOCK - if (is_partial && result.end > 0) { - // Try to extract any partial results from what was successfully parsed - common_chat_msg msg; - msg.role = "assistant"; - std::unique_ptr mapper; - if (params.format == COMMON_CHAT_FORMAT_PEG_GEMMA4) { - mapper = std::make_unique(msg); - } else { - mapper = std::make_unique(msg); - } - mapper->from_ast(ctx.ast, result); - - if (ctx.is_debug()) { - fprintf(stderr, "\nAST for partial parse (fail):\n%s\n", ctx.ast.dump().c_str()); - fflush(stderr); - } - return msg; - } - throw std::runtime_error(std::string("Failed to parse input at pos ") + std::to_string(result.end) + ": " + - effective_input.substr(result.end)); - } - - common_chat_msg msg; - msg.role = "assistant"; - - std::unique_ptr mapper; - if (params.format == COMMON_CHAT_FORMAT_PEG_GEMMA4) { - mapper = std::make_unique(msg); - } else { - mapper = std::make_unique(msg); - } - mapper->from_ast(ctx.ast, result); - - if (ctx.is_debug()) { - fprintf(stderr, "\nAST for %s parse:\n%s\n", is_partial ? "partial" : "full", ctx.ast.dump().c_str()); - fflush(stderr); - } - - if (!is_partial) { - LOG_DBG("Parsed message: %s\n", common_chat_msgs_to_json_oaicompat({ msg }).at(0).dump().c_str()); - } - return msg; -} - -std::map common_chat_templates_get_caps(const common_chat_templates * chat_templates) { - GGML_ASSERT(chat_templates != nullptr); - GGML_ASSERT(chat_templates->template_default != nullptr); - return chat_templates->template_default->caps.to_map(); -} diff --git a/common/chat.h b/common/chat.h deleted file mode 100644 index 01a47b383bf3..000000000000 --- a/common/chat.h +++ /dev/null @@ -1,284 +0,0 @@ -// Chat support (incl. tool call grammar constraining & output parsing) w/ generic & custom template handlers. - -#pragma once - -#include "common.h" -#include "peg-parser.h" -#include "jinja/parser.h" -#include "jinja/runtime.h" -#include "jinja/caps.h" - -#include "nlohmann/json_fwd.hpp" - -#include -#include -#include -#include -#include - -using chat_template_caps = jinja::caps; -using json = nlohmann::ordered_json; - -struct common_chat_templates; - -namespace autoparser { -struct generation_params; -} // namespace autoparser - -struct common_chat_tool_call { - std::string name; - std::string arguments; - std::string id; - - bool operator==(const common_chat_tool_call & other) const { - return name == other.name && arguments == other.arguments && id == other.id; - } -}; - -struct common_chat_msg_content_part { - std::string type; - std::string text; - - // TODO @ngxson : no known chat templates support reasoning_content in content parts yet - // this can be useful for models with interleaved thinking (like Kimi-K2) - // if you see any templates explicitly support this, please ping me - // std::string reasoning_content; - - bool operator==(const common_chat_msg_content_part & other) const { - return type == other.type && text == other.text; - } -}; - -struct common_chat_template { - jinja::program prog; - std::string bos_tok; - std::string eos_tok; - std::string src; - chat_template_caps caps; - - common_chat_template(const std::string & src, const std::string & bos_token, const std::string & eos_token) { - jinja::lexer lexer; - auto lexer_res = lexer.tokenize(src); - this->prog = jinja::parse_from_tokens(lexer_res); - - this->src = lexer_res.source; - this->bos_tok = bos_token; - this->eos_tok = eos_token; - - this->caps = jinja::caps_get(prog); - // LOG_INF("%s: caps:\n%s\n", __func__, this->caps.to_string().c_str()); - } - - const std::string & source() const { return src; } - const std::string & bos_token() const { return bos_tok; } - const std::string & eos_token() const { return eos_tok; } - - chat_template_caps original_caps() const { - return caps; - } -}; - -struct common_chat_msg { - std::string role; - std::string content; - std::vector content_parts; - std::vector tool_calls; - std::string reasoning_content; - std::string tool_name; - std::string tool_call_id; - - nlohmann::ordered_json to_json_oaicompat(bool concat_typed_text = false) const; - - bool empty() const { - return content.empty() && content_parts.empty() && tool_calls.empty() && reasoning_content.empty() && - tool_name.empty() && tool_call_id.empty(); - } - - void set_tool_call_ids(std::vector & ids_cache, - const std::function & gen_tool_call_id) { - for (auto i = 0u; i < tool_calls.size(); i++) { - if (ids_cache.size() <= i) { - auto id = tool_calls[i].id; - if (id.empty()) { - id = gen_tool_call_id(); - } - ids_cache.push_back(id); - } - tool_calls[i].id = ids_cache[i]; - } - } - - bool operator==(const common_chat_msg & other) const { - return role == other.role && content == other.content && content_parts == other.content_parts && - tool_calls == other.tool_calls && reasoning_content == other.reasoning_content && - tool_name == other.tool_name && tool_call_id == other.tool_call_id; - } - - bool operator!=(const common_chat_msg & other) const { return !(*this == other); } -}; - -struct common_chat_msg_diff { - std::string reasoning_content_delta; - std::string content_delta; - size_t tool_call_index = std::string::npos; - common_chat_tool_call tool_call_delta; - - static std::vector compute_diffs(const common_chat_msg & msg_prv, - const common_chat_msg & msg_new); - - bool operator==(const common_chat_msg_diff & other) const { - return content_delta == other.content_delta && tool_call_index == other.tool_call_index && - tool_call_delta == other.tool_call_delta; - } -}; - -struct common_chat_tool { - std::string name; - std::string description; - std::string parameters; -}; - -enum common_chat_tool_choice { - COMMON_CHAT_TOOL_CHOICE_AUTO, - COMMON_CHAT_TOOL_CHOICE_REQUIRED, - COMMON_CHAT_TOOL_CHOICE_NONE, -}; - -enum common_chat_format { - COMMON_CHAT_FORMAT_CONTENT_ONLY, - - // These are intended to be parsed by the PEG parser - COMMON_CHAT_FORMAT_PEG_SIMPLE, - COMMON_CHAT_FORMAT_PEG_NATIVE, - COMMON_CHAT_FORMAT_PEG_GEMMA4, - - COMMON_CHAT_FORMAT_COUNT, // Not a format, just the # formats -}; - -struct common_chat_templates_inputs { - std::vector messages; - std::string grammar; - std::string json_schema; - bool add_generation_prompt = true; - bool use_jinja = true; - // Parameters below only supported when use_jinja is true - std::vector tools; - common_chat_tool_choice tool_choice = COMMON_CHAT_TOOL_CHOICE_AUTO; - bool parallel_tool_calls = false; - common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_NONE; // TODO: refactor this to "bool enable_thinking" - bool enable_thinking = true; - std::chrono::system_clock::time_point now = std::chrono::system_clock::now(); - std::map chat_template_kwargs; - bool add_bos = false; - bool add_eos = false; - bool force_pure_content = false; -}; - -struct common_chat_params { - common_chat_format format = COMMON_CHAT_FORMAT_CONTENT_ONLY; - std::string prompt; - std::string grammar; - bool grammar_lazy = false; - std::string generation_prompt; - bool supports_thinking = false; - std::string thinking_start_tag; // e.g., "" - std::string thinking_end_tag; // e.g., "" - std::vector grammar_triggers; - std::vector preserved_tokens; - std::vector additional_stops; - std::string parser; -}; - -// per-message parsing syntax -// should be derived from common_chat_params -struct common_chat_parser_params { - common_chat_format format = COMMON_CHAT_FORMAT_CONTENT_ONLY; - common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_NONE; // TODO: refactor this to "bool parse_reasoning" - // Whether reasoning_content should be inlined in the content (e.g. for reasoning_format=deepseek in stream mode) - bool reasoning_in_content = false; - std::string generation_prompt; - bool parse_tool_calls = true; - bool debug = false; // Enable debug output for PEG parser - common_peg_arena parser = {}; - common_chat_parser_params() = default; - common_chat_parser_params(const common_chat_params & chat_params) { - format = chat_params.format; - generation_prompt = chat_params.generation_prompt; - } -}; - -// Check if the template supplied via "--chat-template" is supported or not. Returns true if it's valid -bool common_chat_verify_template(const std::string & tmpl, bool use_jinja); - -void common_chat_templates_free(struct common_chat_templates * tmpls); - -struct common_chat_templates_deleter { - void operator()(common_chat_templates * tmpls) { common_chat_templates_free(tmpls); } -}; - -typedef std::unique_ptr common_chat_templates_ptr; - -common_chat_templates_ptr common_chat_templates_init(const struct llama_model * model, - const std::string & chat_template_override, - const std::string & bos_token_override = "", - const std::string & eos_token_override = ""); - -bool common_chat_templates_was_explicit(const struct common_chat_templates * tmpls); -std::string common_chat_templates_source(const struct common_chat_templates * tmpls, const std::string & variant = ""); - -struct common_chat_params common_chat_templates_apply(const struct common_chat_templates * tmpls, - const struct common_chat_templates_inputs & inputs); - -// Format single message, while taking into account the position of that message in chat history -std::string common_chat_format_single(const struct common_chat_templates * tmpls, - const std::vector & past_msg, - const common_chat_msg & new_msg, - bool add_ass, - bool use_jinja); - -// Returns an example of formatted chat -std::string common_chat_format_example(const struct common_chat_templates * tmpls, - bool use_jinja, - const std::map & chat_template_kwargs); - -const char * common_chat_format_name(common_chat_format format); -common_chat_msg common_chat_parse(const std::string & input, bool is_partial, const common_chat_parser_params & params); -common_chat_msg common_chat_peg_parse(const common_peg_arena & src_parser, const std::string & input, bool is_partial, const common_chat_parser_params & params); - -// used by arg and server -const char * common_reasoning_format_name(common_reasoning_format format); -common_reasoning_format common_reasoning_format_from_name(const std::string & format); - -common_chat_tool_choice common_chat_tool_choice_parse_oaicompat(const std::string & tool_choice); - -bool common_chat_templates_support_enable_thinking(const common_chat_templates * chat_templates); - -// Parses a JSON array of messages in OpenAI's chat completion API format. -std::vector common_chat_msgs_parse_oaicompat(const nlohmann::ordered_json & messages); - -std::vector common_chat_tools_parse_oaicompat(const nlohmann::ordered_json & tools); - -// DEPRECATED: only used in tests -nlohmann::ordered_json common_chat_msgs_to_json_oaicompat(const std::vector & msgs, bool concat_typed_text = false); - -nlohmann::ordered_json common_chat_tools_to_json_oaicompat(const std::vector & tools); - -// get template caps, useful for reporting to server /props endpoint -std::map common_chat_templates_get_caps(const common_chat_templates * chat_templates); - -std::string common_chat_template_direct_apply( - const common_chat_template & tmpl, - const autoparser::generation_params & inputs); - -std::optional common_chat_try_specialized_template( - const common_chat_template & tmpl, - const std::string & src, - autoparser::generation_params & params); - -// specialized per-task preset -struct common_chat_prompt_preset { - std::string system; - std::string user; -}; - -common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates); diff --git a/common/common.cpp b/common/common.cpp deleted file mode 100644 index 793b8fee7b84..000000000000 --- a/common/common.cpp +++ /dev/null @@ -1,1962 +0,0 @@ -#include "ggml.h" -#include "gguf.h" - -#include "build-info.h" -#include "common.h" -#include "fit.h" -#include "log.h" -#include "llama.h" -#include "sampling.h" -#include "unicode.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#if defined(__APPLE__) && defined(__MACH__) -#include -#include -#endif - -#if defined(_WIN32) -#define WIN32_LEAN_AND_MEAN -#ifndef NOMINMAX -# define NOMINMAX -#endif -#include -#include -#include -#include -#include -#else -#include -#include -#include -#endif - -#if defined(__linux__) -#include -#include -#endif - -#if defined(_MSC_VER) -#pragma warning(disable: 4244 4267) // possible loss of data -#endif - -common_time_meas::common_time_meas(int64_t & t_acc, bool disable) : t_start_us(disable ? -1 : ggml_time_us()), t_acc(t_acc) {} - -common_time_meas::~common_time_meas() { - if (t_start_us >= 0) { - t_acc += ggml_time_us() - t_start_us; - } -} - -// -// CPU utils -// - -int32_t common_cpu_get_num_physical_cores() { -#ifdef __linux__ - // enumerate the set of thread siblings, num entries is num cores - std::unordered_set siblings; - for (uint32_t cpu=0; cpu < UINT32_MAX; ++cpu) { - std::ifstream thread_siblings("/sys/devices/system/cpu/cpu" - + std::to_string(cpu) + "/topology/thread_siblings"); - if (!thread_siblings.is_open()) { - break; // no more cpus - } - std::string line; - if (std::getline(thread_siblings, line)) { - siblings.insert(line); - } - } - if (!siblings.empty()) { - return static_cast(siblings.size()); - } -#elif defined(__APPLE__) && defined(__MACH__) - int32_t num_physical_cores; - size_t len = sizeof(num_physical_cores); - int result = sysctlbyname("hw.perflevel0.physicalcpu", &num_physical_cores, &len, NULL, 0); - if (result == 0) { - return num_physical_cores; - } - result = sysctlbyname("hw.physicalcpu", &num_physical_cores, &len, NULL, 0); - if (result == 0) { - return num_physical_cores; - } -#elif defined(_WIN32) && (_WIN32_WINNT >= 0x0601) && !defined(__MINGW64__) // windows 7 and later - // TODO: windows + arm64 + mingw64 - unsigned int n_threads_win = std::thread::hardware_concurrency(); - unsigned int default_threads = n_threads_win > 0 ? (n_threads_win <= 4 ? n_threads_win : n_threads_win / 2) : 4; - - DWORD buffer_size = 0; - if (!GetLogicalProcessorInformationEx(RelationProcessorCore, nullptr, &buffer_size)) { - if (GetLastError() != ERROR_INSUFFICIENT_BUFFER) { - return default_threads; - } - } - - std::vector buffer(buffer_size); - if (!GetLogicalProcessorInformationEx(RelationProcessorCore, reinterpret_cast(buffer.data()), &buffer_size)) { - return default_threads; - } - - int32_t num_physical_cores = 0; - PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX info = reinterpret_cast(buffer.data()); - while (buffer_size > 0) { - if (info->Relationship == RelationProcessorCore) { - num_physical_cores += info->Processor.GroupCount; - } - buffer_size -= info->Size; - info = reinterpret_cast(reinterpret_cast(info) + info->Size); - } - - return num_physical_cores > 0 ? num_physical_cores : default_threads; -#endif - unsigned int n_threads = std::thread::hardware_concurrency(); - return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4; -} - -#if defined(__x86_64__) && defined(__linux__) && !defined(__ANDROID__) -#include - -static void cpuid(unsigned leaf, unsigned subleaf, - unsigned *eax, unsigned *ebx, unsigned *ecx, unsigned *edx) { - __asm__("movq\t%%rbx,%%rsi\n\t" - "cpuid\n\t" - "xchgq\t%%rbx,%%rsi" - : "=a"(*eax), "=S"(*ebx), "=c"(*ecx), "=d"(*edx) - : "0"(leaf), "2"(subleaf)); -} - -static int pin_cpu(int cpu) { - cpu_set_t mask; - CPU_ZERO(&mask); - CPU_SET(cpu, &mask); - return pthread_setaffinity_np(pthread_self(), sizeof(mask), &mask); -} - -static bool is_hybrid_cpu(void) { - unsigned eax, ebx, ecx, edx; - cpuid(7, 0, &eax, &ebx, &ecx, &edx); - return !!(edx & (1u << 15)); -} - -static bool is_running_on_efficiency_core(void) { - unsigned eax, ebx, ecx, edx; - cpuid(0x1a, 0, &eax, &ebx, &ecx, &edx); - int intel_atom = 0x20; - int core_type = (eax & 0xff000000u) >> 24; - return core_type == intel_atom; -} - -static int cpu_count_math_cpus(int n_cpu) { - int result = 0; - for (int cpu = 0; cpu < n_cpu; ++cpu) { - if (pin_cpu(cpu)) { - return -1; - } - if (is_running_on_efficiency_core()) { - continue; // efficiency cores harm lockstep threading - } - ++cpu; // hyperthreading isn't useful for linear algebra - ++result; - } - return result; -} - -#endif // __x86_64__ && __linux__ - -/** - * Returns number of CPUs on system that are useful for math. - */ -int32_t common_cpu_get_num_math() { -#if defined(__x86_64__) && defined(__linux__) && !defined(__ANDROID__) - int n_cpu = sysconf(_SC_NPROCESSORS_ONLN); - if (n_cpu < 1) { - return common_cpu_get_num_physical_cores(); - } - if (is_hybrid_cpu()) { - cpu_set_t affinity; - if (!pthread_getaffinity_np(pthread_self(), sizeof(affinity), &affinity)) { - int result = cpu_count_math_cpus(n_cpu); - pthread_setaffinity_np(pthread_self(), sizeof(affinity), &affinity); - if (result > 0) { - return result; - } - } - } -#endif - return common_cpu_get_num_physical_cores(); -} - -// Helper for setting process priority - -#if defined(_WIN32) - -bool set_process_priority(enum ggml_sched_priority prio) { - if (prio == GGML_SCHED_PRIO_NORMAL) { - return true; - } - - DWORD p = NORMAL_PRIORITY_CLASS; - switch (prio) { - case GGML_SCHED_PRIO_LOW: p = BELOW_NORMAL_PRIORITY_CLASS; break; - case GGML_SCHED_PRIO_NORMAL: p = NORMAL_PRIORITY_CLASS; break; - case GGML_SCHED_PRIO_MEDIUM: p = ABOVE_NORMAL_PRIORITY_CLASS; break; - case GGML_SCHED_PRIO_HIGH: p = HIGH_PRIORITY_CLASS; break; - case GGML_SCHED_PRIO_REALTIME: p = REALTIME_PRIORITY_CLASS; break; - } - - if (!SetPriorityClass(GetCurrentProcess(), p)) { - LOG_WRN("failed to set process priority class %d : (%d)\n", prio, (int) GetLastError()); - return false; - } - - return true; -} - -#else // MacOS and POSIX -#include -#include - -bool set_process_priority(enum ggml_sched_priority prio) { - if (prio == GGML_SCHED_PRIO_NORMAL) { - return true; - } - - int p = 0; - switch (prio) { - case GGML_SCHED_PRIO_LOW: p = 5; break; - case GGML_SCHED_PRIO_NORMAL: p = 0; break; - case GGML_SCHED_PRIO_MEDIUM: p = -5; break; - case GGML_SCHED_PRIO_HIGH: p = -10; break; - case GGML_SCHED_PRIO_REALTIME: p = -20; break; - } - - if (setpriority(PRIO_PROCESS, 0, p) != 0) { - LOG_WRN("failed to set process priority %d : %s (%d)\n", prio, strerror(errno), errno); - return false; - } - return true; -} - -#endif - -// -// CLI argument parsing -// - - -void postprocess_cpu_params(common_cpu_params & cpuparams, const common_cpu_params * role_model) { - int32_t n_set = 0; - - if (cpuparams.n_threads < 0) { - // Assuming everything about cpuparams is invalid - if (role_model != nullptr) { - cpuparams = *role_model; - } else { - cpuparams.n_threads = common_cpu_get_num_math(); - } - } - - for (int32_t i = 0; i < GGML_MAX_N_THREADS; i++) { - if (cpuparams.cpumask[i]) { - n_set++; - } - } - - if (n_set && n_set < cpuparams.n_threads) { - // Not enough set bits, may experience performance issues. - LOG_WRN("Not enough set bits in CPU mask (%d) to satisfy requested thread count: %d\n", n_set, cpuparams.n_threads); - } -} - -bool parse_cpu_range(const std::string & range, bool (&boolmask)[GGML_MAX_N_THREADS]) { - size_t dash_loc = range.find('-'); - if (dash_loc == std::string::npos) { - LOG_ERR("Format of CPU range is invalid! Expected []-[].\n"); - return false; - } - - size_t start_i; - size_t end_i; - - if (dash_loc == 0) { - start_i = 0; - } else { - start_i = std::stoull(range.substr(0, dash_loc)); - if (start_i >= GGML_MAX_N_THREADS) { - LOG_ERR("Start index out of bounds!\n"); - return false; - } - } - - if (dash_loc == range.length() - 1) { - end_i = GGML_MAX_N_THREADS - 1; - } else { - end_i = std::stoull(range.substr(dash_loc + 1)); - if (end_i >= GGML_MAX_N_THREADS) { - LOG_ERR("End index out of bounds!\n"); - return false; - } - } - - for (size_t i = start_i; i <= end_i; i++) { - boolmask[i] = true; - } - - return true; -} - -bool parse_cpu_mask(const std::string & mask, bool (&boolmask)[GGML_MAX_N_THREADS]) { - // Discard potential 0x prefix - size_t start_i = 0; - if (mask.length() >= 2 && mask.substr(0, 2) == "0x") { - start_i = 2; - } - - size_t num_digits = mask.length() - start_i; - if (num_digits > 128) num_digits = 128; - - size_t end_i = num_digits + start_i; - - for (size_t i = start_i, n = (num_digits*4 - 1); i < end_i; i++, n-=4) { - char c = mask.at(i); - int8_t id = c; - - if ((c >= '0' && c <= '9')) { - id -= '0'; - } else if (c >= 'a' && c <= 'f') { - id -= 'a' - 10; - } else if (c >= 'A' && c <= 'F') { - id -= 'A' - 10; - } else { - LOG_ERR("Invalid hex character '%c' at position %d\n", c, int32_t(i)); - return false; - } - - boolmask[ n ] = boolmask[ n ] || ((id & 8) != 0); - boolmask[n - 1] = boolmask[n - 1] || ((id & 4) != 0); - boolmask[n - 2] = boolmask[n - 2] || ((id & 2) != 0); - boolmask[n - 3] = boolmask[n - 3] || ((id & 1) != 0); - } - - return true; -} - -void common_init() { -#if defined(_WIN32) - SetConsoleOutputCP(CP_UTF8); - SetConsoleCP(CP_UTF8); -#endif - - llama_log_set(common_log_default_callback, NULL); - -#ifdef NDEBUG - const char * build_type = ""; -#else - const char * build_type = " (debug)"; -#endif - - LOG_DBG("build: %d (%s) with %s for %s%s\n", llama_build_number(), llama_commit(), llama_compiler(), llama_build_target(), build_type); -} - -std::string common_params_get_system_info(const common_params & params) { - std::ostringstream os; - - os << "system_info: n_threads = " << params.cpuparams.n_threads; - if (params.cpuparams_batch.n_threads != -1) { - os << " (n_threads_batch = " << params.cpuparams_batch.n_threads << ")"; - } -#if defined(_WIN32) && (_WIN32_WINNT >= 0x0601) && !defined(__MINGW64__) // windows 7 and later - // TODO: windows + arm64 + mingw64 - DWORD logicalProcessorCount = GetActiveProcessorCount(ALL_PROCESSOR_GROUPS); - os << " / " << logicalProcessorCount << " | " << llama_print_system_info(); -#else - os << " / " << std::thread::hardware_concurrency() << " | " << llama_print_system_info(); -#endif - - return os.str(); -} - -// -// String utils -// - -std::string string_format(const char * fmt, ...) { - va_list ap; - va_list ap2; - va_start(ap, fmt); - va_copy(ap2, ap); - int size = vsnprintf(NULL, 0, fmt, ap); - GGML_ASSERT(size >= 0 && size < INT_MAX); // NOLINT - std::vector buf(size + 1); - int size2 = vsnprintf(buf.data(), size + 1, fmt, ap2); - GGML_ASSERT(size2 == size); - va_end(ap2); - va_end(ap); - return std::string(buf.data(), size); -} - -std::string string_strip(const std::string & str) { - size_t start = 0; - size_t end = str.size(); - while (start < end && std::isspace(str[start])) { - start++; - } - while (end > start && std::isspace(str[end - 1])) { - end--; - } - return str.substr(start, end - start); -} - -std::string string_get_sortable_timestamp() { - using clock = std::chrono::system_clock; - - const clock::time_point current_time = clock::now(); - const time_t as_time_t = clock::to_time_t(current_time); - char timestamp_no_ns[100]; - std::strftime(timestamp_no_ns, 100, "%Y_%m_%d-%H_%M_%S", std::localtime(&as_time_t)); - - const int64_t ns = std::chrono::duration_cast( - current_time.time_since_epoch() % 1000000000).count(); - char timestamp_ns[11]; - snprintf(timestamp_ns, 11, "%09" PRId64, ns); - - return std::string(timestamp_no_ns) + "." + std::string(timestamp_ns); -} - -void string_replace_all(std::string & s, const std::string & search, const std::string & replace) { - if (search.empty()) { - return; - } - std::string builder; - builder.reserve(s.length()); - size_t pos = 0; - size_t last_pos = 0; - while ((pos = s.find(search, last_pos)) != std::string::npos) { - builder.append(s, last_pos, pos - last_pos); - builder.append(replace); - last_pos = pos + search.length(); - } - builder.append(s, last_pos, std::string::npos); - s = std::move(builder); -} - -std::string regex_escape(const std::string & s) { - static const std::regex special_chars("[.^$|()*+?\\[\\]{}\\\\]"); - return std::regex_replace(s, special_chars, "\\$&"); -} - -std::string string_join(const std::vector & values, const std::string & separator) { - std::ostringstream result; - for (size_t i = 0; i < values.size(); ++i) { - if (i > 0) { - result << separator; - } - result << values[i]; - } - return result.str(); -} - -std::vector string_split(const std::string & str, const std::string & delimiter) { - std::vector parts; - size_t start = 0; - size_t end = str.find(delimiter); - - while (end != std::string::npos) { - parts.push_back(str.substr(start, end - start)); - start = end + delimiter.length(); - end = str.find(delimiter, start); - } - - parts.push_back(str.substr(start)); - - return parts; -} - -std::string string_repeat(const std::string & str, size_t n) { - if (n == 0) { - return ""; - } - - std::string result; - result.reserve(str.length() * n); - - for (size_t i = 0; i < n; ++i) { - result += str; - } - - return result; -} - -std::string string_from(bool value) { - return value ? "true" : "false"; -} - -std::string string_from(const std::vector & values) { - std::stringstream buf; - - buf << "[ "; - bool first = true; - for (auto e : values) { - if (first) { - first = false; - } else { - buf << ", "; - } - buf << std::to_string(e); - } - buf << " ]"; - - return buf.str(); -} - -std::string string_from(const struct llama_context * ctx, const std::vector & tokens) { - std::stringstream buf; - - buf << "[ "; - - bool first = true; - for (const auto & token : tokens) { - if (!first) { - buf << ", "; - } else { - first = false; - } - - auto detokenized = common_token_to_piece(ctx, token); - - buf << "'" << detokenized << "'" - << ":" << std::to_string(token); - } - - buf << " ]"; - - return buf.str(); -} - -std::string string_from(const struct llama_context * ctx, const struct llama_batch & batch) { - std::stringstream buf; - - buf << "[ "; - - bool first = true; - for (int i = 0; i < batch.n_tokens; ++i) { - if (!first) { - buf << ", "; - } else { - first = false; - } - - auto detokenized = common_token_to_piece(ctx, batch.token[i]); - - buf << "\n" << std::to_string(i) - << ", token '" << detokenized << "'" - << ", pos " << std::to_string(batch.pos[i]) - << ", n_seq_id " << std::to_string(batch.n_seq_id[i]) - << ", seq_id " << std::to_string(batch.seq_id[i][0]) - << ", logits " << std::to_string(batch.logits[i]); - } - - buf << " ]"; - - return buf.str(); -} - -void string_process_escapes(std::string & input) { - std::size_t input_len = input.length(); - std::size_t output_idx = 0; - - for (std::size_t input_idx = 0; input_idx < input_len; ++input_idx) { - if (input[input_idx] == '\\' && input_idx + 1 < input_len) { - switch (input[++input_idx]) { - case 'n': input[output_idx++] = '\n'; break; - case 'r': input[output_idx++] = '\r'; break; - case 't': input[output_idx++] = '\t'; break; - case '\'': input[output_idx++] = '\''; break; - case '\"': input[output_idx++] = '\"'; break; - case '\\': input[output_idx++] = '\\'; break; - case 'x': - // Handle \x12, etc - if (input_idx + 2 < input_len) { - const char x[3] = { input[input_idx + 1], input[input_idx + 2], 0 }; - char *err_p = nullptr; - const long val = std::strtol(x, &err_p, 16); - if (err_p == x + 2) { - input_idx += 2; - input[output_idx++] = char(val); - break; - } - } - // fall through - default: input[output_idx++] = '\\'; - input[output_idx++] = input[input_idx]; break; - } - } else { - input[output_idx++] = input[input_idx]; - } - } - - input.resize(output_idx); -} - -bool string_parse_kv_override(const char * data, std::vector & overrides) { - const char * sep = strchr(data, '='); - if (sep == nullptr || sep - data >= 128) { - LOG_ERR("%s: malformed KV override '%s'\n", __func__, data); - return false; - } - llama_model_kv_override kvo; - std::strncpy(kvo.key, data, sep - data); - kvo.key[sep - data] = 0; - sep++; - if (strncmp(sep, "int:", 4) == 0) { - sep += 4; - kvo.tag = LLAMA_KV_OVERRIDE_TYPE_INT; - kvo.val_i64 = std::atol(sep); - } else if (strncmp(sep, "float:", 6) == 0) { - sep += 6; - kvo.tag = LLAMA_KV_OVERRIDE_TYPE_FLOAT; - kvo.val_f64 = std::atof(sep); - } else if (strncmp(sep, "bool:", 5) == 0) { - sep += 5; - kvo.tag = LLAMA_KV_OVERRIDE_TYPE_BOOL; - if (std::strcmp(sep, "true") == 0) { - kvo.val_bool = true; - } else if (std::strcmp(sep, "false") == 0) { - kvo.val_bool = false; - } else { - LOG_ERR("%s: invalid boolean value for KV override '%s'\n", __func__, data); - return false; - } - } else if (strncmp(sep, "str:", 4) == 0) { - sep += 4; - kvo.tag = LLAMA_KV_OVERRIDE_TYPE_STR; - if (strlen(sep) > 127) { - LOG_ERR("%s: malformed KV override '%s', value cannot exceed 127 chars\n", __func__, data); - return false; - } - strncpy(kvo.val_str, sep, 127); - kvo.val_str[127] = '\0'; - } else { - LOG_ERR("%s: invalid type for KV override '%s'\n", __func__, data); - return false; - } - overrides.emplace_back(std::move(kvo)); - return true; -} - -static inline bool glob_class_match(const char c, const char * pattern, const char * class_end) { - const char * class_start = pattern; - bool negated = false; - - if (*class_start == '!') { - negated = true; - class_start++; - } - - // If first character after negation is ']' or '-', treat it as literal - if (*class_start == ']' || *class_start == '-') { - if (class_start < class_end && *class_start == c) { - return !negated; - } - class_start++; - } - - bool matched = false; - - while (class_start < class_end) { - if (class_start + 2 < class_end && class_start[1] == '-' && class_start[2] != ']') { - char start_char = *class_start; - char end_char = class_start[2]; - if (c >= start_char && c <= end_char) { - matched = true; - break; - } - class_start += 3; - } else { - if (*class_start == c) { - matched = true; - break; - } - class_start++; - } - } - - return negated ? !matched : matched; -} - -// simple glob: * matches non-/ chars, ** matches anything including /, [] matches character class -static inline bool glob_match(const char * pattern, const char * str) { - if (*pattern == '\0') { - return *str == '\0'; - } - if (pattern[0] == '*' && pattern[1] == '*') { - const char * p = pattern + 2; - if (glob_match(p, str)) return true; - if (*str != '\0') return glob_match(pattern, str + 1); - return false; - } - if (*pattern == '*') { - const char * p = pattern + 1; - for (; *str != '\0' && *str != '/'; str++) { - if (glob_match(p, str)) return true; - } - return glob_match(p, str); - } - if (*pattern == '?' && *str != '\0' && *str != '/') { - return glob_match(pattern + 1, str + 1); - } - if (*pattern == '[') { - const char * class_end = pattern + 1; - // If first character after '[' is ']' or '-', treat it as literal - if (*class_end == ']' || *class_end == '-') { - class_end++; - } - while (*class_end != '\0' && *class_end != ']') { - class_end++; - } - if (*class_end == ']') { - if (*str == '\0') return false; - bool matched = glob_class_match(*str, pattern + 1, class_end); - return matched && glob_match(class_end + 1, str + 1); - } else { - if (*str == '[') { - return glob_match(pattern + 1, str + 1); - } - return false; - } - } - if (*pattern == *str) { - return glob_match(pattern + 1, str + 1); - } - return false; -} - -bool glob_match(const std::string & pattern, const std::string & str) { - return glob_match(pattern.c_str(), str.c_str()); -} - -// -// Filesystem utils -// - -// Validate if a filename is safe to use -// To validate a full path, split the path by the OS-specific path separator, and validate each part with this function -bool fs_validate_filename(const std::string & filename, bool allow_subdirs) { - if (!filename.length()) { - // Empty filename invalid - return false; - } - if (filename.length() > 255) { - // Limit at common largest possible filename on Linux filesystems - // to avoid unnecessary further validation - // (On systems with smaller limits it will be caught by the OS) - return false; - } - - size_t offset = 0; - while (offset < filename.size()) { - utf8_parse_result result = common_parse_utf8_codepoint(filename, offset); - - if (result.status != utf8_parse_result::SUCCESS) { - return false; - } - uint32_t c = result.codepoint; - - if ((result.bytes_consumed == 2 && c < 0x80) || - (result.bytes_consumed == 3 && c < 0x800) || - (result.bytes_consumed == 4 && c < 0x10000)) { - return false; - } - - // Check for forbidden codepoints: - // - Control characters - // - Unicode equivalents of illegal characters - // - UTF-16 surrogate pairs - // - UTF-8 replacement character - // - Byte order mark (BOM) - // - Illegal characters: / \ : * ? " < > | - if (c <= 0x1F // Control characters (C0) - || c == 0x7F // Control characters (DEL) - || (c >= 0x80 && c <= 0x9F) // Control characters (C1) - || c == 0xFF0E // Fullwidth Full Stop (period equivalent) - || c == 0x2215 // Division Slash (forward slash equivalent) - || c == 0x2216 // Set Minus (backslash equivalent) - || (c >= 0xD800 && c <= 0xDFFF) // UTF-16 surrogate pairs - || c > 0x10FFFF // Max Unicode limit - || c == 0xFFFD // Replacement Character (UTF-8) - || c == 0xFEFF // Byte Order Mark (BOM) - || c == ':' || c == '*' // Illegal characters - || c == '?' || c == '"' || c == '<' || c == '>' || c == '|') { - return false; - } - if (!allow_subdirs && (c == '/' || c == '\\')) { - // Subdirectories not allowed, reject path separators - return false; - } - offset += result.bytes_consumed; - } - - // Reject any leading or trailing ' ', or any trailing '.', these are stripped on Windows and will cause a different filename - // Unicode and other whitespace is not affected, only 0x20 space - if (filename.front() == ' ' || filename.back() == ' ' || filename.back() == '.') { - return false; - } - - // Reject any ".." (currently stricter than necessary, it should be fine to just check for == ".." instead) - if (filename.find("..") != std::string::npos) { - return false; - } - - // Reject "." - if (filename == ".") { - return false; - } - - return true; -} - -#include - - -#ifdef _WIN32 -static std::wstring utf8_to_wstring(const std::string & str) { - if (str.empty()) { - return std::wstring(); - } - - int size = MultiByteToWideChar(CP_UTF8, 0, str.c_str(), (int)str.size(), NULL, 0); - - if (size <= 0) { - return std::wstring(); - } - - std::wstring wstr(size, 0); - MultiByteToWideChar(CP_UTF8, 0, str.c_str(), (int)str.size(), &wstr[0], size); - - return wstr; -} -#endif - -// returns true if successful, false otherwise -bool fs_create_directory_with_parents(const std::string & path) { -#ifdef _WIN32 - std::wstring wpath = utf8_to_wstring(path); - - // if the path already exists, check whether it's a directory - const DWORD attributes = GetFileAttributesW(wpath.c_str()); - if ((attributes != INVALID_FILE_ATTRIBUTES) && (attributes & FILE_ATTRIBUTE_DIRECTORY)) { - return true; - } - - size_t pos_slash = 0; - - // process path from front to back, procedurally creating directories - while ((pos_slash = path.find('\\', pos_slash)) != std::string::npos) { - const std::wstring subpath = wpath.substr(0, pos_slash); - - pos_slash += 1; - - // skip the drive letter, in some systems it can return an access denied error - if (subpath.length() == 2 && subpath[1] == ':') { - continue; - } - - const bool success = CreateDirectoryW(subpath.c_str(), NULL); - - if (!success) { - const DWORD error = GetLastError(); - - // if the path already exists, ensure that it's a directory - if (error == ERROR_ALREADY_EXISTS) { - const DWORD attributes = GetFileAttributesW(subpath.c_str()); - if (attributes == INVALID_FILE_ATTRIBUTES || !(attributes & FILE_ATTRIBUTE_DIRECTORY)) { - return false; - } - } else { - return false; - } - } - } - - return true; -#else - // if the path already exists, check whether it's a directory - struct stat info; - if (stat(path.c_str(), &info) == 0) { - return S_ISDIR(info.st_mode); - } - - size_t pos_slash = 1; // skip leading slashes for directory creation - - // process path from front to back, procedurally creating directories - while ((pos_slash = path.find('/', pos_slash)) != std::string::npos) { - const std::string subpath = path.substr(0, pos_slash); - struct stat info; - - // if the path already exists, ensure that it's a directory - if (stat(subpath.c_str(), &info) == 0) { - if (!S_ISDIR(info.st_mode)) { - return false; - } - } else { - // create parent directories - const int ret = mkdir(subpath.c_str(), 0755); - if (ret != 0) { - return false; - } - } - - pos_slash += 1; - } - - return true; -#endif // _WIN32 -} - -bool fs_is_directory(const std::string & path) { - std::filesystem::path dir(path); - return std::filesystem::exists(dir) && std::filesystem::is_directory(dir); -} - -std::string fs_get_cache_directory() { - std::string cache_directory = ""; - auto ensure_trailing_slash = [](std::string p) { - // Make sure to add trailing slash - if (p.back() != DIRECTORY_SEPARATOR) { - p += DIRECTORY_SEPARATOR; - } - return p; - }; - if (getenv("LLAMA_CACHE")) { - cache_directory = std::getenv("LLAMA_CACHE"); - } else { -#if defined(__linux__) || defined(__FreeBSD__) || defined(_AIX) || \ - defined(__OpenBSD__) || defined(__NetBSD__) - if (std::getenv("XDG_CACHE_HOME")) { - cache_directory = std::getenv("XDG_CACHE_HOME"); - } else if (std::getenv("HOME")) { - cache_directory = std::getenv("HOME") + std::string("/.cache/"); - } else { -#if defined(__linux__) - /* no $HOME is defined, fallback to getpwuid */ - struct passwd *pw = getpwuid(getuid()); - if ((!pw) || (!pw->pw_dir)) { - throw std::runtime_error("Failed to find $HOME directory"); - } - - cache_directory = std::string(pw->pw_dir) + std::string("/.cache/"); -#else /* defined(__linux__) */ - throw std::runtime_error("Failed to find $HOME directory"); -#endif /* defined(__linux__) */ - } -#elif defined(__APPLE__) - cache_directory = std::getenv("HOME") + std::string("/Library/Caches/"); -#elif defined(_WIN32) - cache_directory = std::getenv("LOCALAPPDATA"); -#elif defined(__EMSCRIPTEN__) - GGML_ABORT("not implemented on this platform"); -#else -# error Unknown architecture -#endif - cache_directory = ensure_trailing_slash(cache_directory); - cache_directory += "llama.cpp"; - } - return ensure_trailing_slash(cache_directory); -} - -std::string fs_get_cache_file(const std::string & filename) { - GGML_ASSERT(filename.find(DIRECTORY_SEPARATOR) == std::string::npos); - std::string cache_directory = fs_get_cache_directory(); - const bool success = fs_create_directory_with_parents(cache_directory); - if (!success) { - throw std::runtime_error("failed to create cache directory: " + cache_directory); - } - return cache_directory + filename; -} - -std::vector fs_list(const std::string & path, bool include_directories) { - std::vector files; - if (path.empty()) return files; - - std::filesystem::path dir(path); - if (!std::filesystem::exists(dir) || !std::filesystem::is_directory(dir)) { - return files; - } - - for (const auto & entry : std::filesystem::directory_iterator(dir)) { - try { - // Only include regular files (skip directories) - const auto & p = entry.path(); - if (std::filesystem::is_regular_file(p)) { - common_file_info info; - info.path = p.string(); - info.name = p.filename().string(); - info.is_dir = false; - try { - info.size = static_cast(std::filesystem::file_size(p)); - } catch (const std::filesystem::filesystem_error &) { - info.size = 0; - } - files.push_back(std::move(info)); - } else if (include_directories && std::filesystem::is_directory(p)) { - common_file_info info; - info.path = p.string(); - info.name = p.filename().string(); - info.size = 0; // Directories have no size - info.is_dir = true; - files.push_back(std::move(info)); - } - } catch (const std::filesystem::filesystem_error &) { - // skip entries we cannot inspect - continue; - } - } - - return files; -} - -// -// TTY utils -// - -bool tty_can_use_colors() { - // Check NO_COLOR environment variable (https://no-color.org/) - if (const char * no_color = std::getenv("NO_COLOR")) { - if (no_color[0] != '\0') { - return false; - } - } - - // Check TERM environment variable - if (const char * term = std::getenv("TERM")) { - if (std::strcmp(term, "dumb") == 0) { - return false; - } - } - - // Check if stdout and stderr are connected to a terminal - // We check both because log messages can go to either - bool stdout_is_tty = isatty(fileno(stdout)); - bool stderr_is_tty = isatty(fileno(stderr)); - - return stdout_is_tty || stderr_is_tty; -} - -// -// Model utils -// - -// TODO: move to common/sampling -static void common_init_sampler_from_model( - const llama_model * model, - common_params_sampling & sparams) { - - const uint64_t config = sparams.user_sampling_config; - - auto get_int32 = [&](const char * key, int32_t & dst, uint64_t user_config) { - if (config & user_config) { - return; - } - - char buf[64] = {0}; - if (llama_model_meta_val_str(model, key, buf, sizeof(buf)) > 0) { - char * end = nullptr; - int32_t v = strtol(buf, &end, 10); - if (end && end != buf) { - dst = v; - } - } - }; - - auto get_float = [&](const char * key, float & dst, uint64_t user_config) { - if (config & user_config) { - return; - } - - char buf[128] = {0}; - if (llama_model_meta_val_str(model, key, buf, sizeof(buf)) > 0) { - char * end = nullptr; - float v = strtof(buf, &end); - if (end && end != buf) { - dst = v; - } - } - }; - - // Sampling sequence - if (!(config & common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_SAMPLERS)) { - char buf[512] = {0}; - if (llama_model_meta_val_str(model, llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_SEQUENCE), buf, sizeof(buf)) > 0) { - const std::vector sampler_names = string_split(std::string(buf), ';'); - if (!sampler_names.empty()) { - sparams.samplers = common_sampler_types_from_names(sampler_names, true); - } - } - } - - get_int32(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_TOP_K), sparams.top_k, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TOP_K); - get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_TOP_P), sparams.top_p, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TOP_P); - get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_MIN_P), sparams.min_p, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIN_P); - get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_XTC_PROBABILITY), sparams.xtc_probability, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_XTC_PROBABILITY); - get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_XTC_THRESHOLD), sparams.xtc_threshold, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_XTC_THRESHOLD); - get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_TEMP), sparams.temp, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TEMP); - get_int32(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_PENALTY_LAST_N), sparams.penalty_last_n, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_LAST_N); - get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_PENALTY_REPEAT), sparams.penalty_repeat, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_REPEAT); - get_int32(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT), sparams.mirostat, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT); - get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_TAU), sparams.mirostat_tau, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_TAU); - get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_ETA), sparams.mirostat_eta, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_ETA); -} - -struct common_init_result::impl { - impl() = default; - ~impl() = default; - - // note: the order in which model, context, etc. are declared matters because their destructors will be called bottom-to-top - - llama_model_ptr model; - llama_context_ptr context; - - std::vector lora; - - std::vector samplers; - std::vector samplers_seq_config; -}; - -common_init_result::common_init_result(common_params & params) : - pimpl(new impl{}) { - auto mparams = common_model_params_to_llama(params); - auto cparams = common_context_params_to_llama(params); - - if (params.fit_params) { - LOG_INF("%s: fitting params to device memory, for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on\n", __func__); - common_fit_params(params.model.path.c_str(), &mparams, &cparams, - params.tensor_split, - params.tensor_buft_overrides.data(), - params.fit_params_target.data(), - params.fit_params_min_ctx, - params.verbosity >= 4 ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); - } - - llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams); - if (model == NULL) { - return; - } - - pimpl->model.reset(model); - - const llama_vocab * vocab = llama_model_get_vocab(model); - - // load and optionally apply lora adapters - for (auto & la : params.lora_adapters) { - llama_adapter_lora_ptr lora; - lora.reset(llama_adapter_lora_init(model, la.path.c_str())); - if (lora == nullptr) { - LOG_ERR("%s: failed to load lora adapter '%s'\n", __func__, la.path.c_str()); - pimpl->model.reset(model); - return; - } - - char buf[1024]; - la.ptr = lora.get(); - llama_adapter_meta_val_str(la.ptr, "adapter.lora.task_name", buf, sizeof(buf)); - la.task_name = buf; - llama_adapter_meta_val_str(la.ptr, "adapter.lora.prompt_prefix", buf, sizeof(buf)); - la.prompt_prefix = buf; - pimpl->lora.emplace_back(std::move(lora)); // copy to list of loaded adapters - } - - // updates params.sampling - // TODO: fix naming - common_init_sampler_from_model(model, params.sampling); - - if (params.sampling.ignore_eos && llama_vocab_eos(vocab) == LLAMA_TOKEN_NULL) { - LOG_WRN("%s: warning: vocab does not have an EOS token, ignoring --ignore-eos\n", __func__); - params.sampling.ignore_eos = false; - } - - // initialize once - for (llama_token i = 0; i < llama_vocab_n_tokens(vocab); i++) { - if (llama_vocab_is_eog(vocab, i)) { - LOG_INF("%s: added %s logit bias = %f\n", __func__, common_token_to_piece(vocab, i).c_str(), -INFINITY); - params.sampling.logit_bias_eog.push_back({i, -INFINITY}); - } - } - - if (params.sampling.ignore_eos) { - // add EOG biases to the active set of logit biases - params.sampling.logit_bias.insert( - params.sampling.logit_bias.end(), - params.sampling.logit_bias_eog.begin(), params.sampling.logit_bias_eog.end()); - } - - //if (params.sampling.penalty_last_n == -1) { - // LOG_INF("%s: setting penalty_last_n to ctx_size = %d\n", __func__, llama_n_ctx(lctx)); - // params.sampling.penalty_last_n = llama_n_ctx(lctx); - //} - - //if (params.sampling.dry_penalty_last_n == -1) { - // LOG_INF("%s: setting dry_penalty_last_n to ctx_size = %d\n", __func__, llama_n_ctx(lctx)); - // params.sampling.dry_penalty_last_n = llama_n_ctx(lctx); - //} - - // init the backend samplers as part of the context creation - pimpl->samplers.resize(cparams.n_seq_max); - pimpl->samplers_seq_config.resize(cparams.n_seq_max); - - for (int i = 0; i < (int) cparams.n_seq_max; ++i) { - pimpl->samplers[i].reset(common_sampler_init(model, params.sampling)); - pimpl->samplers_seq_config[i] = { i, common_sampler_get(pimpl->samplers[i].get()) }; - } - - if (params.sampling.backend_sampling) { - cparams.samplers = pimpl->samplers_seq_config.data(); - cparams.n_samplers = pimpl->samplers_seq_config.size(); - } - - llama_context * lctx = llama_init_from_model(model, cparams); - if (lctx == NULL) { - LOG_ERR("%s: failed to create context with model '%s'\n", __func__, params.model.path.c_str()); - return; - } - - pimpl->context.reset(lctx); -} - -llama_model * common_init_result::model() { - return pimpl->model.get(); -} - -llama_context * common_init_result::context() { - return pimpl->context.get(); -} - -common_sampler * common_init_result::sampler(llama_seq_id seq_id) { - if (seq_id < 0 || seq_id >= (int) pimpl->samplers.size()) { - return nullptr; - } - return pimpl->samplers[seq_id].get(); -} - -void common_init_result::reset_samplers() { - for (int i = 0; i < (int) pimpl->samplers.size(); ++i) { - llama_sampler_reset(common_sampler_get(pimpl->samplers[i].get())); - } -} - -std::vector & common_init_result::lora() { - return pimpl->lora; -} - -common_init_result_ptr common_init_from_params(common_params & params) { - common_init_result_ptr res(new common_init_result(params)); - - llama_model * model = res->model(); - if (model == NULL) { - LOG_ERR("%s: failed to load model '%s'\n", __func__, params.model.path.c_str()); - return res; - } - - llama_context * lctx = res->context(); - if (lctx == NULL) { - LOG_ERR("%s: failed to create context with model '%s'\n", __func__, params.model.path.c_str()); - return res; - } - - const llama_vocab * vocab = llama_model_get_vocab(model); - - if (params.ctx_shift && !llama_memory_can_shift(llama_get_memory(lctx))) { - LOG_WRN("%s: KV cache shifting is not supported for this context, disabling KV cache shifting\n", __func__); - params.ctx_shift = false; - } - - if (!params.control_vectors.empty()) { - if (params.control_vector_layer_start <= 0) params.control_vector_layer_start = 1; - if (params.control_vector_layer_end <= 0) params.control_vector_layer_end = llama_model_n_layer(model); - - const auto cvec = common_control_vector_load(params.control_vectors); - if (cvec.n_embd == -1) { - return res; - } - - int err = llama_set_adapter_cvec( - lctx, - cvec.data.data(), - cvec.data.size(), - cvec.n_embd, - params.control_vector_layer_start, - params.control_vector_layer_end); - if (err) { - return res; - } - } - - if (llama_pooling_type(lctx) == LLAMA_POOLING_TYPE_RANK) { - bool ok = true; - - if (llama_vocab_bos(vocab) == LLAMA_TOKEN_NULL) { - LOG_WRN("%s: warning: vocab does not have a BOS token, reranking will not work\n", __func__); - ok = false; - } - - bool has_eos = llama_vocab_eos(vocab) != LLAMA_TOKEN_NULL; - bool has_sep = llama_vocab_sep(vocab) != LLAMA_TOKEN_NULL; - bool has_rerank_prompt = llama_model_chat_template(model, "rerank") != NULL; - - if (!has_eos && !has_sep && !has_rerank_prompt) { - LOG_WRN("%s: warning: vocab does not have an EOS token, SEP token, or rerank prompt. Reranking will not work\n", __func__); - ok = false; - } else if (!has_eos) { - LOG_WRN("%s: warning: vocab does not have an EOS token, using SEP token as fallback\n", __func__); - } - - if (!ok) { - return res; - } - } - - if (!params.lora_init_without_apply) { - common_set_adapter_lora(lctx, params.lora_adapters); - } - - if (params.warmup) { - LOG_WRN("%s: warming up the model with an empty run - please wait ... (--no-warmup to disable)\n", __func__); - - llama_set_warmup(lctx, true); - - std::vector tmp; - llama_token bos = llama_vocab_bos(vocab); - llama_token eos = llama_vocab_eos(vocab); - - // some models (e.g. T5) don't have a BOS token - if (bos != LLAMA_TOKEN_NULL) { - tmp.push_back(bos); - } - if (eos != LLAMA_TOKEN_NULL) { - tmp.push_back(eos); - } - if (tmp.empty()) { - tmp.push_back(0); - } - - if (llama_model_has_encoder(model)) { - llama_encode(lctx, llama_batch_get_one(tmp.data(), tmp.size())); - llama_token decoder_start_token_id = llama_model_decoder_start_token(model); - if (decoder_start_token_id == LLAMA_TOKEN_NULL) { - decoder_start_token_id = bos; - } - tmp.clear(); - tmp.push_back(decoder_start_token_id); - } - if (llama_model_has_decoder(model)) { - llama_decode(lctx, llama_batch_get_one(tmp.data(), std::min(tmp.size(), (size_t) params.n_batch))); - } - llama_memory_clear(llama_get_memory(lctx), true); - llama_synchronize(lctx); - llama_perf_context_reset(lctx); - llama_set_warmup(lctx, false); - - // reset samplers to reset RNG state after warmup to the seeded state - res->reset_samplers(); - } - - return res; -} - -common_init_result::~common_init_result() = default; - -std::string common_get_model_endpoint() { - const char * model_endpoint_env = getenv("MODEL_ENDPOINT"); - // We still respect the use of environment-variable "HF_ENDPOINT" for backward-compatibility. - const char * hf_endpoint_env = getenv("HF_ENDPOINT"); - const char * endpoint_env = model_endpoint_env ? model_endpoint_env : hf_endpoint_env; - std::string model_endpoint = "https://huggingface.co/"; - if (endpoint_env) { - model_endpoint = endpoint_env; - if (model_endpoint.back() != '/') { - model_endpoint += '/'; - } - } - return model_endpoint; -} - -common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx) { - auto * mem = llama_get_memory(ctx); - if (mem == nullptr) { - return COMMON_CONTEXT_SEQ_RM_TYPE_NO; - } - - common_context_seq_rm_type res = COMMON_CONTEXT_SEQ_RM_TYPE_PART; - - llama_memory_clear(mem, true); - - // eval 2 tokens to check if the context is compatible - std::vector tmp; - tmp.push_back(0); - tmp.push_back(0); - - int ret = llama_decode(ctx, llama_batch_get_one(tmp.data(), tmp.size())); - if (ret != 0) { - LOG_ERR("%s: llama_decode() failed: %d\n", __func__, ret); - res = COMMON_CONTEXT_SEQ_RM_TYPE_NO; - goto done; - } - - // try to remove the last tokens - if (!llama_memory_seq_rm(mem, 0, 1, -1)) { - LOG_WRN("%s: the target context does not support partial sequence removal\n", __func__); - res = COMMON_CONTEXT_SEQ_RM_TYPE_FULL; - goto done; - } - -done: - llama_memory_clear(mem, true); - llama_synchronize(ctx); - - return res; -} - -void common_set_adapter_lora(struct llama_context * ctx, std::vector & lora) { - std::vector loras; - std::vector scales; - - for (auto & la: lora) { - loras.push_back(la.ptr); - scales.push_back(la.scale); - } - - llama_set_adapters_lora(ctx, loras.data(), loras.size(), scales.data()); -} - -struct llama_model_params common_model_params_to_llama(common_params & params) { - auto mparams = llama_model_default_params(); - - if (!params.devices.empty()) { - mparams.devices = params.devices.data(); - } - - mparams.n_gpu_layers = params.n_gpu_layers; - mparams.main_gpu = params.main_gpu; - mparams.split_mode = params.split_mode; - mparams.tensor_split = params.tensor_split; - mparams.use_mmap = params.use_mmap; - mparams.use_direct_io = params.use_direct_io; - mparams.use_mlock = params.use_mlock; - mparams.check_tensors = params.check_tensors; - mparams.use_extra_bufts = !params.no_extra_bufts; - mparams.no_host = params.no_host; - - if (params.kv_overrides.empty()) { - mparams.kv_overrides = NULL; - } else { - GGML_ASSERT(params.kv_overrides.back().key[0] == 0 && "KV overrides not terminated with empty key"); - mparams.kv_overrides = params.kv_overrides.data(); - } - - if (params.tensor_buft_overrides.empty()) { - mparams.tensor_buft_overrides = NULL; - } else { - GGML_ASSERT(params.tensor_buft_overrides.back().pattern == nullptr && "Tensor buffer overrides not terminated with empty pattern"); - mparams.tensor_buft_overrides = params.tensor_buft_overrides.data(); - } - - mparams.progress_callback = params.load_progress_callback; - mparams.progress_callback_user_data = params.load_progress_callback_user_data; - mparams.no_alloc = params.no_alloc; - - return mparams; -} - -struct llama_context_params common_context_params_to_llama(const common_params & params) { - auto cparams = llama_context_default_params(); - - cparams.n_ctx = params.n_ctx; - cparams.n_seq_max = params.n_parallel; - cparams.n_batch = params.n_batch; - cparams.n_ubatch = params.n_ubatch; - cparams.n_threads = params.cpuparams.n_threads; - cparams.n_threads_batch = params.cpuparams_batch.n_threads == -1 ? - params.cpuparams.n_threads : params.cpuparams_batch.n_threads; - cparams.embeddings = params.embedding; - cparams.rope_scaling_type = params.rope_scaling_type; - cparams.rope_freq_base = params.rope_freq_base; - cparams.rope_freq_scale = params.rope_freq_scale; - cparams.yarn_ext_factor = params.yarn_ext_factor; - cparams.yarn_attn_factor = params.yarn_attn_factor; - cparams.yarn_beta_fast = params.yarn_beta_fast; - cparams.yarn_beta_slow = params.yarn_beta_slow; - cparams.yarn_orig_ctx = params.yarn_orig_ctx; - cparams.pooling_type = params.pooling_type; - cparams.attention_type = params.attention_type; - cparams.flash_attn_type = params.flash_attn_type; - cparams.cb_eval = params.cb_eval; - cparams.cb_eval_user_data = params.cb_eval_user_data; - cparams.offload_kqv = !params.no_kv_offload; - cparams.no_perf = params.no_perf; - cparams.op_offload = !params.no_op_offload; - cparams.swa_full = params.swa_full; - cparams.kv_unified = params.kv_unified; - - cparams.type_k = params.cache_type_k; - cparams.type_v = params.cache_type_v; - - return cparams; -} - -struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const common_cpu_params & params) { - struct ggml_threadpool_params tpp; - - ggml_threadpool_params_init(&tpp, params.n_threads); // setup the defaults - - if (params.mask_valid) { - std::memcpy(&tpp.cpumask, ¶ms.cpumask, GGML_MAX_N_THREADS); - } - - tpp.prio = params.priority; - tpp.poll = params.poll; - tpp.strict_cpu = params.strict_cpu; - - return tpp; -} - -// -// Batch utils -// - -void common_batch_clear(struct llama_batch & batch) { - batch.n_tokens = 0; -} - -void common_batch_add( - struct llama_batch & batch, - llama_token id, - llama_pos pos, - const std::vector & seq_ids, - bool logits) { - GGML_ASSERT(batch.seq_id[batch.n_tokens] && "llama_batch size exceeded"); - - batch.token [batch.n_tokens] = id; - batch.pos [batch.n_tokens] = pos; - batch.n_seq_id[batch.n_tokens] = seq_ids.size(); - for (size_t i = 0; i < seq_ids.size(); ++i) { - batch.seq_id[batch.n_tokens][i] = seq_ids[i]; - } - batch.logits [batch.n_tokens] = logits; - - batch.n_tokens++; -} - -// -// Vocab utils -// - -std::vector common_tokenize( - const struct llama_context * ctx, - const std::string & text, - bool add_special, - bool parse_special) { - const llama_model * model = llama_get_model(ctx); - const llama_vocab * vocab = llama_model_get_vocab(model); - return common_tokenize(vocab, text, add_special, parse_special); -} - -std::vector common_tokenize( - const struct llama_vocab * vocab, - const std::string & text, - bool add_special, - bool parse_special) { - // upper limit for the number of tokens - int n_tokens = text.length() + 2 * add_special; - std::vector result(n_tokens); - n_tokens = llama_tokenize(vocab, text.data(), text.length(), result.data(), result.size(), add_special, parse_special); - if (n_tokens == std::numeric_limits::min()) { - throw std::runtime_error("Tokenization failed: input text too large, tokenization result exceeds int32_t limit"); - } - if (n_tokens < 0) { - result.resize(-n_tokens); - int check = llama_tokenize(vocab, text.data(), text.length(), result.data(), result.size(), add_special, parse_special); - GGML_ASSERT(check == -n_tokens); - } else { - result.resize(n_tokens); - } - return result; -} - -std::string common_token_to_piece(const struct llama_context * ctx, llama_token token, bool special) { - const llama_model * model = llama_get_model(ctx); - const llama_vocab * vocab = llama_model_get_vocab(model); - return common_token_to_piece(vocab, token, special); -} - -std::string common_token_to_piece(const struct llama_vocab * vocab, llama_token token, bool special) { - std::string piece; - piece.resize(piece.capacity()); // using string internal cache, 15 bytes + '\n' - const int n_chars = llama_token_to_piece(vocab, token, &piece[0], piece.size(), 0, special); - if (n_chars < 0) { - piece.resize(-n_chars); - int check = llama_token_to_piece(vocab, token, &piece[0], piece.size(), 0, special); - GGML_ASSERT(check == -n_chars); - } - else { - piece.resize(n_chars); - } - - return piece; -} - -std::string common_detokenize(const struct llama_context * ctx, const std::vector & tokens, bool special) { - const llama_model * model = llama_get_model(ctx); - const llama_vocab * vocab = llama_model_get_vocab(model); - return common_detokenize(vocab, tokens, special); -} - -std::string common_detokenize(const struct llama_vocab * vocab, const std::vector & tokens, bool special) { - std::string text; - text.resize(std::max(text.capacity(), tokens.size())); - int32_t n_chars = llama_detokenize(vocab, tokens.data(), (int32_t)tokens.size(), &text[0], (int32_t)text.size(), false, special); - if (n_chars < 0) { - text.resize(-n_chars); - n_chars = llama_detokenize(vocab, tokens.data(), (int32_t)tokens.size(), &text[0], (int32_t)text.size(), false, special); - GGML_ASSERT(n_chars <= (int32_t)text.size()); // whitespace trimming is performed after per-token detokenization - } - - text.resize(n_chars); - - // NOTE: the original tokenizer decodes bytes after collecting the pieces. - return text; -} - -// -// Embedding utils -// - -void common_embd_normalize(const float * inp, float * out, int n, int embd_norm) { - double sum = 0.0; - - switch (embd_norm) { - case -1: // no normalisation - sum = 1.0; - break; - case 0: // max absolute - for (int i = 0; i < n; i++) { - if (sum < std::abs(inp[i])) { - sum = std::abs(inp[i]); - } - } - sum /= 32760.0; // make an int16 range - break; - case 2: // euclidean - for (int i = 0; i < n; i++) { - sum += inp[i] * inp[i]; - } - sum = std::sqrt(sum); - break; - default: // p-norm (euclidean is p-norm p=2) - for (int i = 0; i < n; i++) { - sum += std::pow(std::abs(inp[i]), embd_norm); - } - sum = std::pow(sum, 1.0 / embd_norm); - break; - } - - const float norm = sum > 0.0 ? 1.0 / sum : 0.0f; - - for (int i = 0; i < n; i++) { - out[i] = inp[i] * norm; - } -} - -float common_embd_similarity_cos(const float * embd1, const float * embd2, int n){ - double sum = 0.0; - double sum1 = 0.0; - double sum2 = 0.0; - - for (int i = 0; i < n; i++) { - sum += embd1[i] * embd2[i]; - sum1 += embd1[i] * embd1[i]; - sum2 += embd2[i] * embd2[i]; - } - - // Handle the case where one or both vectors are zero vectors - if (sum1 == 0.0 || sum2 == 0.0) { - if (sum1 == 0.0 && sum2 == 0.0) { - return 1.0f; // two zero vectors are similar - } - return 0.0f; - } - - return sum / (sqrt(sum1) * sqrt(sum2)); -} - -// -// Control vector utils -// - -static common_control_vector_data common_control_vector_load_one(const common_control_vector_load_info & load_info) { - common_control_vector_data result = { -1, {} }; - - ggml_context * ctx = nullptr; - struct gguf_init_params meta_gguf_params = { - /* .no_alloc = */ false, - /* .ctx = */ &ctx, - }; - struct gguf_context * ctx_gguf = gguf_init_from_file(load_info.fname.c_str(), meta_gguf_params); - if (!ctx_gguf) { - LOG_ERR("%s: failed to load control vector file from %s\n", __func__, load_info.fname.c_str()); - return result; - } - - int32_t n_tensors = gguf_get_n_tensors(ctx_gguf); - if (n_tensors == 0) { - LOG_WRN("%s: no direction tensors found in %s\n", __func__, load_info.fname.c_str()); - } - - for (int i = 0; i < n_tensors; i++) { - std::string name = gguf_get_tensor_name(ctx_gguf, i); - - int layer_idx = -1; - - // split on '.' - size_t dotpos = name.find('.'); - if (dotpos != std::string::npos && name.substr(0, dotpos) == "direction") { - try { - layer_idx = std::stoi(name.substr(dotpos + 1)); - } catch (...) { - layer_idx = -1; - } - } - if (layer_idx < 0) { - LOG_ERR("%s: invalid/unparsable direction tensor layer index in %s\n", __func__, load_info.fname.c_str()); - result.n_embd = -1; - break; - } else if (layer_idx == 0) { - LOG_ERR("%s: invalid (zero) direction tensor layer index in %s\n", __func__, load_info.fname.c_str()); - result.n_embd = -1; - break; - } - - struct ggml_tensor * tensor = ggml_get_tensor(ctx, name.c_str()); - if (tensor->type != GGML_TYPE_F32) { - LOG_ERR("%s: invalid (non-F32) direction tensor type in %s\n", __func__, load_info.fname.c_str()); - result.n_embd = -1; - break; - } - if (ggml_n_dims(tensor) != 1) { - LOG_ERR("%s: invalid (non-1D) direction tensor shape in %s\n", __func__, load_info.fname.c_str()); - result.n_embd = -1; - break; - } - - if (result.n_embd == -1) { - result.n_embd = ggml_nelements(tensor); - } else if (ggml_nelements(tensor) != result.n_embd) { - LOG_ERR("%s: direction tensor in %s does not match previous dimensions\n", __func__, load_info.fname.c_str()); - result.n_embd = -1; - break; - } - - // extend if necessary - do not store data for layer 0 (it's not used) - result.data.resize(std::max(result.data.size(), static_cast(result.n_embd * layer_idx)), 0.0f); - - const float * src = (const float *) tensor->data; - float * dst = result.data.data() + result.n_embd * (layer_idx - 1); // layer 1 at [0] - for (int j = 0; j < result.n_embd; j++) { - dst[j] += src[j] * load_info.strength; // allows multiple directions for same layer in same file - } - - } - - if (result.n_embd == -1) { - LOG_WRN("%s: skipping %s due to invalid direction tensors\n", __func__, load_info.fname.c_str()); - result.data.clear(); - } - - gguf_free(ctx_gguf); - ggml_free(ctx); - - return result; -} - -common_control_vector_data common_control_vector_load(const std::vector & load_infos) { - common_control_vector_data result = { -1, {} }; - - for (const auto & info : load_infos) { - auto cur = common_control_vector_load_one(info); - - if (cur.n_embd == -1) { - result.n_embd = -1; - break; - } - if (result.n_embd != -1 && result.n_embd != cur.n_embd) { - LOG_ERR("%s: control vectors in %s does not match previous dimensions\n", __func__, info.fname.c_str()); - result.n_embd = -1; - break; - } - - if (result.n_embd == -1) { - result = std::move(cur); - } else { - result.data.resize(std::max(result.data.size(), cur.data.size()), 0.0f); // extend if necessary - for (size_t i = 0; i < cur.data.size(); i++) { - result.data[i] += cur.data[i]; - } - } - } - - if (result.n_embd == -1) { - LOG_ERR("%s: no valid control vector files passed\n", __func__); - result.data.clear(); - } - - return result; -} - -ggml_opt_dataset_t common_opt_dataset_init(struct llama_context * ctx, const std::vector & tokens, int64_t stride) { - const int64_t ne_datapoint = llama_n_ctx(ctx); - const int64_t ndata = (tokens.size() - ne_datapoint - 1) / stride; - ggml_opt_dataset_t result = ggml_opt_dataset_init( - GGML_TYPE_I32, GGML_TYPE_I32, ne_datapoint, ne_datapoint, ndata, /*ndata_shard =*/ 1); - - llama_token * data = (llama_token *) ggml_opt_dataset_data(result)->data; - llama_token * labels = (llama_token *) ggml_opt_dataset_labels(result)->data; - - for (int64_t idata = 0; idata < ndata; ++idata) { - memcpy(data + idata*ne_datapoint, tokens.data() + idata*stride + 0, ne_datapoint*sizeof(llama_token)); - memcpy(labels + idata*ne_datapoint, tokens.data() + idata*stride + 1, ne_datapoint*sizeof(llama_token)); - } - - return result; -} - -ggml_opt_optimizer_params common_opt_lr_pars(void * userdata) { - ggml_opt_optimizer_params result = ggml_opt_get_default_optimizer_params(nullptr); - const lr_opt & d = *(lr_opt *) userdata; - result.adamw.alpha = result.sgd.alpha = d.get_lr(d.epoch); - result.sgd.wd = result.adamw.wd = d.wd; - return result; -} - -// TODO make all command line args case-insensitive -static inline bool eq_case_insensitive(char const* a, char const* b) { - return ! -#if defined(_MSC_VER) - _stricmp -#else - strcasecmp -#endif // defined(_MSC_VER) - (a, b); -} - -enum ggml_opt_optimizer_type common_opt_get_optimizer(const char * n) { - if (eq_case_insensitive("adamw", n)) { - return GGML_OPT_OPTIMIZER_TYPE_ADAMW; - } - if (eq_case_insensitive("sgd", n)) { - return GGML_OPT_OPTIMIZER_TYPE_SGD; - } - return GGML_OPT_OPTIMIZER_TYPE_COUNT; -} - -// TODO simplify to use just log and exp -static float const k_log_2 = std::log(2.f); - -void lr_opt::init() { - if (lr_min > 0 && lr_min < lr0) { - float nhalf = std::log(lr0 / lr_min) / k_log_2; - float e = epochs; - if (decay_epochs > 0 && decay_epochs < e) { - e = decay_epochs; - } else { - decay_epochs = e; - } - scale_epoch = nhalf / e; - } -} - -float lr_opt::get_lr(float epoch) const { - float r = lr_min <= 0 ? lr0 : - epoch >= decay_epochs ? lr_min : - lr0 * std::pow(0.5f, epoch * scale_epoch); - LOG_INF("epoch %.2g lr=%.2g\n", epoch, r); - return r; -} - -bool common_replay_last_token(struct llama_context * ctx, llama_token last_token, int32_t pos) { - llama_batch batch = llama_batch_get_one(&last_token, 1); - batch.pos = &pos; - if (llama_decode(ctx, batch)) { - LOG_ERR("%s: failed to replay last token\n", __func__); - return false; - } - return true; -} - -bool common_prompt_batch_decode( - struct llama_context * ctx, - const std::vector & tokens, - int & n_past, - int n_batch, - std::string_view state_path, - bool save_state) { - const int n_eval = tokens.size(); - if (n_eval == 0) { - return true; - } - - if (save_state && n_eval > 1) { - const int n_tokens_before_last = n_eval - 1; - - GGML_ASSERT(n_eval <= n_batch); - - // Decode all but the last token so we can save the memory state before decoding the last token. - // This is done so we can restore the session state later and replay the last token. - // Memory implementations in recurrent/hybrid models don't support removing tokens from their - // memory, so we can't just remove the last token from the memory and replay the last token which - // is the reason for this logic. - if (llama_decode(ctx, llama_batch_get_one(const_cast(tokens.data()), n_tokens_before_last))) { - LOG_ERR("%s : failed to eval\n", __func__); - return false; - } - n_past += n_tokens_before_last; - - llama_state_save_file(ctx, state_path.data(), tokens.data(), n_tokens_before_last); - LOG_INF("saved session before last token to %s, n_tokens = %d\n", state_path.data(), n_tokens_before_last); - - llama_token last_token = tokens.back(); - llama_batch batch = llama_batch_get_one(&last_token, 1); - int32_t pos = n_past; - batch.pos = &pos; - - if (llama_decode(ctx, batch)) { - LOG_ERR("%s : failed to eval last token\n", __func__); - return false; - } - n_past++; - } else { - if (llama_decode(ctx, llama_batch_get_one(const_cast(tokens.data()), n_eval))) { - LOG_ERR("%s : failed to eval\n", __func__); - return false; - } - n_past += n_eval; - } - - return true; -} diff --git a/common/common.cpp.inc b/common/common.cpp.inc new file mode 100644 index 000000000000..5fb94cc63931 --- /dev/null +++ b/common/common.cpp.inc @@ -0,0 +1,23896 @@ +#include "ggml.h.inc" +#include "gguf.h.inc" + +#include "common.h.inc" +#include "llama.h.inc" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct common_rust_string_view { + const uint8_t * data; + size_t len; +}; + +struct common_rust_string_list { + llama_common_unicode_string * data; + size_t len; +}; + +struct common_rust_gguf_split_info { + llama_common_unicode_string prefix; + llama_common_unicode_string tag; + int index; + int count; +}; + +extern "C" { +int llama_build_number(void); +const char * llama_commit(void); +const char * llama_compiler(void); +const char * llama_build_target(void); +llama_common_unicode_string llama_common_string_repeat_rust(const uint8_t * data, size_t len, size_t n); +llama_common_unicode_string llama_common_int_vec_to_string_rust(const int32_t * data, size_t len); +llama_common_unicode_string llama_common_string_trim_rust(const uint8_t * data, size_t len, int mode); +size_t llama_common_prefix_len_rust(const uint8_t * left, size_t left_len, const uint8_t * right, size_t right_len); +size_t llama_common_suffix_len_rust(const uint8_t * left, size_t left_len, const uint8_t * right, size_t right_len); +llama_common_unicode_string llama_common_string_replace_all_rust( + const uint8_t * data, size_t len, + const uint8_t * search, size_t search_len, + const uint8_t * replace, size_t replace_len); +llama_common_unicode_string llama_common_regex_escape_rust(const uint8_t * data, size_t len); +llama_common_unicode_string llama_common_string_join_rust( + const common_rust_string_view * values, size_t values_len, + const uint8_t * separator, size_t separator_len); +common_rust_string_list llama_common_string_split_rust( + const uint8_t * data, size_t len, + const uint8_t * delimiter, size_t delimiter_len); +common_rust_string_list llama_common_parse_csv_row_rust(const uint8_t * data, size_t len); +common_rust_string_view llama_common_trim_trailing_space_view_rust(const uint8_t * data, size_t len, int max); +common_rust_string_view llama_common_trim_leading_space_view_rust(const uint8_t * data, size_t len, int max); +llama_common_unicode_string llama_common_string_process_escapes_rust(const uint8_t * data, size_t len); +llama_common_unicode_string llama_common_clean_file_name_rust(const uint8_t * data, size_t len); +llama_common_unicode_string llama_common_string_diff_rust( + const uint8_t * last, size_t last_len, + const uint8_t * current, size_t current_len, + int * status); +bool llama_common_string_starts_with_rust( + const uint8_t * data, size_t len, + const uint8_t * prefix, size_t prefix_len); +bool llama_common_string_ends_with_rust( + const uint8_t * data, size_t len, + const uint8_t * suffix, size_t suffix_len); +bool llama_common_glob_match_rust( + const uint8_t * pattern, size_t pattern_len, + const uint8_t * str, size_t str_len); +bool llama_common_fs_validate_filename_rust(const uint8_t * filename, size_t filename_len, bool allow_subdirs); +llama_common_unicode_string llama_common_string_lstrip_chars_rust( + const uint8_t * data, size_t len, + const char * chars); +llama_common_unicode_string llama_common_string_rstrip_chars_rust( + const uint8_t * data, size_t len, + const char * chars); +common_rust_string_list llama_common_break_str_into_lines_rust( + const uint8_t * data, size_t len, + size_t max_char_per_line); +llama_common_unicode_string llama_common_rm_leading_dashes_rust(const uint8_t * data, size_t len); +llama_common_unicode_string llama_common_until_common_prefix_rust( + const uint8_t * full, size_t full_len, + const uint8_t * left, size_t left_len, + const uint8_t * right, size_t right_len); +llama_common_unicode_string llama_common_after_common_suffix_rust( + const uint8_t * full, size_t full_len, + const uint8_t * left, size_t left_len, + const uint8_t * right, size_t right_len); +llama_common_unicode_string llama_common_json_ensure_ascii_preserving_format_rust( + const uint8_t * data, size_t len); +llama_common_unicode_string llama_common_normalize_quotes_to_json_rust(const uint8_t * data, size_t len); +void llama_common_string_list_free(common_rust_string_list value); +int llama_common_opt_get_optimizer_rust(const char * name); +int llama_common_arg_is_truthy_rust(const char * value); +int llama_common_arg_is_falsey_rust(const char * value); +int llama_common_arg_is_autoy_rust(const char * value); +int llama_common_parse_bool_value_rust(const char * value); +int llama_common_parse_cpu_mask_rust(const uint8_t * mask, size_t mask_len, bool * boolmask, size_t boolmask_len); +int llama_common_chat_tool_choice_parse_oaicompat_rust(const char * value); +int llama_common_kv_cache_type_from_str_rust(const char * value); +llama_common_unicode_string llama_common_get_line_col_rust(const uint8_t * data, size_t len, size_t pos); +void llama_common_embd_normalize_rust(const float * inp, float * out, int n, int embd_norm); +float llama_common_embd_similarity_cos_rust(const float * embd1, const float * embd2, int n); +const char * llama_common_get_all_kv_cache_types_rust(); +const char * llama_common_bool_to_string_rust(bool value); +void llama_common_lr_opt_init_rust( + float lr0, float lr_min, float decay_epochs, unsigned epochs, + float * out_decay_epochs, float * out_scale_epoch); +float llama_common_lr_opt_get_lr_rust(float lr0, float lr_min, float decay_epochs, float scale_epoch, float epoch); +int llama_common_log_level_to_verbosity_rust(int level); +int64_t llama_common_time_us_rust(); +bool llama_common_is_http_status_ok_rust(int status); +bool llama_common_gguf_filename_is_model_rust(const uint8_t * data, size_t len); +bool llama_common_is_hex_digit_rust(uint8_t c); +size_t llama_common_parse_hex_escape_rust( + const uint8_t * data, size_t len, + size_t pos, int hex_count, + uint32_t * out_value); +size_t llama_common_parse_char_class_char_rust( + const uint8_t * data, size_t len, + size_t pos, + uint32_t * out_value); +const char * llama_common_chat_format_name_rust(int format); +const char * llama_common_reasoning_format_name_rust(int format); +int llama_common_reasoning_format_from_name_rust(const char * format); +char llama_common_sampler_type_to_chr_rust(int type); +const char * llama_common_sampler_type_to_str_rust(int type); +int llama_common_sampler_type_from_name_rust(const char * name, bool allow_alt_names); +int llama_common_sampler_type_from_chr_rust(char name); +const char * llama_common_speculative_type_to_str_rust(int type); +int llama_common_speculative_type_from_name_rust(const char * name); +const char * llama_common_peg_parse_result_type_name_rust(int type); +llama_common_unicode_string llama_common_peg_rule_name_rust(const uint8_t * data, size_t len); +llama_common_unicode_string llama_common_gbnf_format_literal_rust(const uint8_t * data, size_t len); +llama_common_unicode_string llama_common_gbnf_build_repetition_rust( + const uint8_t * item_rule, size_t item_rule_len, + int min_items, int max_items, + const uint8_t * separator_rule, size_t separator_rule_len); +llama_common_unicode_string llama_common_gbnf_escape_char_class_rust(uint32_t c); +bool llama_common_gbnf_is_reserved_name_rust(const uint8_t * data, size_t len); +llama_common_unicode_string llama_common_gbnf_ref_name_rust(const uint8_t * data, size_t len); +llama_common_unicode_string llama_common_hf_folder_name_to_repo_rust(const uint8_t * data, size_t len); +llama_common_unicode_string llama_common_hf_repo_to_folder_name_rust(const uint8_t * data, size_t len); +llama_common_unicode_string llama_common_hf_make_old_cache_filename_rust( + const uint8_t * owner, size_t owner_len, + const uint8_t * repo, size_t repo_len, + const uint8_t * filename, size_t filename_len); +int llama_common_json_brace_depth_rust(const uint8_t * data, size_t len); +bool llama_common_hf_is_valid_repo_id_rust(const uint8_t * data, size_t len); +bool llama_common_hf_is_valid_token_rust(const uint8_t * data, size_t len); +bool llama_common_hf_is_valid_commit_rust(const uint8_t * data, size_t len); +bool llama_common_hf_is_valid_oid_rust(const uint8_t * data, size_t len); +common_rust_string_list llama_common_hf_parse_manifest_name_rust(const uint8_t * data, size_t len); +common_rust_gguf_split_info llama_common_gguf_split_info_rust(const uint8_t * data, size_t len, bool extract_tag); +int llama_common_gguf_extract_quant_bits_rust(const uint8_t * data, size_t len); +} + +static std::string common_rust_string_to_std(llama_common_unicode_string value) { + if (!value.data) { + return ""; + } + std::string result((const char *) value.data, value.len); + llama_common_unicode_string_free(value); + return result; +} + +#if defined(__APPLE__) && defined(__MACH__) +#include +#include +#endif + +#if defined(_WIN32) +#define WIN32_LEAN_AND_MEAN +#ifndef NOMINMAX +# define NOMINMAX +#endif +#include +#include +#include +#include +#include +#else +#include +#include +#include +#endif + +#if defined(__linux__) +#include +#include +#endif + +#if defined(_MSC_VER) +#pragma warning(disable: 4244 4267) // possible loss of data +#endif + +common_time_meas::common_time_meas(int64_t & t_acc, bool disable) : t_start_us(disable ? -1 : ggml_time_us()), t_acc(t_acc) {} + +common_time_meas::~common_time_meas() { + if (t_start_us >= 0) { + t_acc += ggml_time_us() - t_start_us; + } +} + +// +// CPU utils +// + +int32_t common_cpu_get_num_physical_cores() { +#ifdef __linux__ + // enumerate the set of thread siblings, num entries is num cores + std::unordered_set siblings; + for (uint32_t cpu=0; cpu < UINT32_MAX; ++cpu) { + std::ifstream thread_siblings("/sys/devices/system/cpu/cpu" + + std::to_string(cpu) + "/topology/thread_siblings"); + if (!thread_siblings.is_open()) { + break; // no more cpus + } + std::string line; + if (std::getline(thread_siblings, line)) { + siblings.insert(line); + } + } + if (!siblings.empty()) { + return static_cast(siblings.size()); + } +#elif defined(__APPLE__) && defined(__MACH__) + int32_t num_physical_cores; + size_t len = sizeof(num_physical_cores); + int result = sysctlbyname("hw.perflevel0.physicalcpu", &num_physical_cores, &len, NULL, 0); + if (result == 0) { + return num_physical_cores; + } + result = sysctlbyname("hw.physicalcpu", &num_physical_cores, &len, NULL, 0); + if (result == 0) { + return num_physical_cores; + } +#elif defined(_WIN32) && (_WIN32_WINNT >= 0x0601) && !defined(__MINGW64__) // windows 7 and later + // TODO: windows + arm64 + mingw64 + unsigned int n_threads_win = std::thread::hardware_concurrency(); + unsigned int default_threads = n_threads_win > 0 ? (n_threads_win <= 4 ? n_threads_win : n_threads_win / 2) : 4; + + DWORD buffer_size = 0; + if (!GetLogicalProcessorInformationEx(RelationProcessorCore, nullptr, &buffer_size)) { + if (GetLastError() != ERROR_INSUFFICIENT_BUFFER) { + return default_threads; + } + } + + std::vector buffer(buffer_size); + if (!GetLogicalProcessorInformationEx(RelationProcessorCore, reinterpret_cast(buffer.data()), &buffer_size)) { + return default_threads; + } + + int32_t num_physical_cores = 0; + PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX info = reinterpret_cast(buffer.data()); + while (buffer_size > 0) { + if (info->Relationship == RelationProcessorCore) { + num_physical_cores += info->Processor.GroupCount; + } + buffer_size -= info->Size; + info = reinterpret_cast(reinterpret_cast(info) + info->Size); + } + + return num_physical_cores > 0 ? num_physical_cores : default_threads; +#endif + unsigned int n_threads = std::thread::hardware_concurrency(); + return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4; +} + +#if defined(__x86_64__) && defined(__linux__) && !defined(__ANDROID__) +#include + +static void cpuid(unsigned leaf, unsigned subleaf, + unsigned *eax, unsigned *ebx, unsigned *ecx, unsigned *edx) { + __asm__("movq\t%%rbx,%%rsi\n\t" + "cpuid\n\t" + "xchgq\t%%rbx,%%rsi" + : "=a"(*eax), "=S"(*ebx), "=c"(*ecx), "=d"(*edx) + : "0"(leaf), "2"(subleaf)); +} + +static int pin_cpu(int cpu) { + cpu_set_t mask; + CPU_ZERO(&mask); + CPU_SET(cpu, &mask); + return pthread_setaffinity_np(pthread_self(), sizeof(mask), &mask); +} + +static bool is_hybrid_cpu(void) { + unsigned eax, ebx, ecx, edx; + cpuid(7, 0, &eax, &ebx, &ecx, &edx); + return !!(edx & (1u << 15)); +} + +static bool is_running_on_efficiency_core(void) { + unsigned eax, ebx, ecx, edx; + cpuid(0x1a, 0, &eax, &ebx, &ecx, &edx); + int intel_atom = 0x20; + int core_type = (eax & 0xff000000u) >> 24; + return core_type == intel_atom; +} + +static int cpu_count_math_cpus(int n_cpu) { + int result = 0; + for (int cpu = 0; cpu < n_cpu; ++cpu) { + if (pin_cpu(cpu)) { + return -1; + } + if (is_running_on_efficiency_core()) { + continue; // efficiency cores harm lockstep threading + } + ++cpu; // hyperthreading isn't useful for linear algebra + ++result; + } + return result; +} + +#endif // __x86_64__ && __linux__ + +/** + * Returns number of CPUs on system that are useful for math. + */ +int32_t common_cpu_get_num_math() { +#if defined(__x86_64__) && defined(__linux__) && !defined(__ANDROID__) + int n_cpu = sysconf(_SC_NPROCESSORS_ONLN); + if (n_cpu < 1) { + return common_cpu_get_num_physical_cores(); + } + if (is_hybrid_cpu()) { + cpu_set_t affinity; + if (!pthread_getaffinity_np(pthread_self(), sizeof(affinity), &affinity)) { + int result = cpu_count_math_cpus(n_cpu); + pthread_setaffinity_np(pthread_self(), sizeof(affinity), &affinity); + if (result > 0) { + return result; + } + } + } +#endif + return common_cpu_get_num_physical_cores(); +} + +// Helper for setting process priority + +#if defined(_WIN32) + +bool set_process_priority(enum ggml_sched_priority prio) { + if (prio == GGML_SCHED_PRIO_NORMAL) { + return true; + } + + DWORD p = NORMAL_PRIORITY_CLASS; + switch (prio) { + case GGML_SCHED_PRIO_LOW: p = BELOW_NORMAL_PRIORITY_CLASS; break; + case GGML_SCHED_PRIO_NORMAL: p = NORMAL_PRIORITY_CLASS; break; + case GGML_SCHED_PRIO_MEDIUM: p = ABOVE_NORMAL_PRIORITY_CLASS; break; + case GGML_SCHED_PRIO_HIGH: p = HIGH_PRIORITY_CLASS; break; + case GGML_SCHED_PRIO_REALTIME: p = REALTIME_PRIORITY_CLASS; break; + } + + if (!SetPriorityClass(GetCurrentProcess(), p)) { + LOG_WRN("failed to set process priority class %d : (%d)\n", prio, (int) GetLastError()); + return false; + } + + return true; +} + +#else // MacOS and POSIX +#include +#include + +bool set_process_priority(enum ggml_sched_priority prio) { + if (prio == GGML_SCHED_PRIO_NORMAL) { + return true; + } + + int p = 0; + switch (prio) { + case GGML_SCHED_PRIO_LOW: p = 5; break; + case GGML_SCHED_PRIO_NORMAL: p = 0; break; + case GGML_SCHED_PRIO_MEDIUM: p = -5; break; + case GGML_SCHED_PRIO_HIGH: p = -10; break; + case GGML_SCHED_PRIO_REALTIME: p = -20; break; + } + + if (setpriority(PRIO_PROCESS, 0, p) != 0) { + LOG_WRN("failed to set process priority %d : %s (%d)\n", prio, strerror(errno), errno); + return false; + } + return true; +} + +#endif + +// +// CLI argument parsing +// + + +void postprocess_cpu_params(common_cpu_params & cpuparams, const common_cpu_params * role_model) { + int32_t n_set = 0; + + if (cpuparams.n_threads < 0) { + // Assuming everything about cpuparams is invalid + if (role_model != nullptr) { + cpuparams = *role_model; + } else { + cpuparams.n_threads = common_cpu_get_num_math(); + } + } + + for (int32_t i = 0; i < GGML_MAX_N_THREADS; i++) { + if (cpuparams.cpumask[i]) { + n_set++; + } + } + + if (n_set && n_set < cpuparams.n_threads) { + // Not enough set bits, may experience performance issues. + LOG_WRN("Not enough set bits in CPU mask (%d) to satisfy requested thread count: %d\n", n_set, cpuparams.n_threads); + } +} + +bool parse_cpu_range(const std::string & range, bool (&boolmask)[GGML_MAX_N_THREADS]) { + size_t dash_loc = range.find('-'); + if (dash_loc == std::string::npos) { + LOG_ERR("Format of CPU range is invalid! Expected []-[].\n"); + return false; + } + + size_t start_i; + size_t end_i; + + if (dash_loc == 0) { + start_i = 0; + } else { + start_i = std::stoull(range.substr(0, dash_loc)); + if (start_i >= GGML_MAX_N_THREADS) { + LOG_ERR("Start index out of bounds!\n"); + return false; + } + } + + if (dash_loc == range.length() - 1) { + end_i = GGML_MAX_N_THREADS - 1; + } else { + end_i = std::stoull(range.substr(dash_loc + 1)); + if (end_i >= GGML_MAX_N_THREADS) { + LOG_ERR("End index out of bounds!\n"); + return false; + } + } + + for (size_t i = start_i; i <= end_i; i++) { + boolmask[i] = true; + } + + return true; +} + +bool parse_cpu_mask(const std::string & mask, bool (&boolmask)[GGML_MAX_N_THREADS]) { + const int status = llama_common_parse_cpu_mask_rust( + (const uint8_t *) mask.data(), mask.size(), boolmask, GGML_MAX_N_THREADS); + if (status >= 0) { + LOG_ERR("Invalid hex character '%c' at position %d\n", mask.at(status), status); + return false; + } + return true; +} + +void common_init() { +#if defined(_WIN32) + SetConsoleOutputCP(CP_UTF8); + SetConsoleCP(CP_UTF8); +#endif + + llama_log_set(common_log_default_callback, NULL); + +#ifdef NDEBUG + const char * build_type = ""; +#else + const char * build_type = " (debug)"; +#endif + + LOG_DBG("build: %d (%s) with %s for %s%s\n", llama_build_number(), llama_commit(), llama_compiler(), llama_build_target(), build_type); +} + +std::string common_params_get_system_info(const common_params & params) { + std::ostringstream os; + + os << "system_info: n_threads = " << params.cpuparams.n_threads; + if (params.cpuparams_batch.n_threads != -1) { + os << " (n_threads_batch = " << params.cpuparams_batch.n_threads << ")"; + } +#if defined(_WIN32) && (_WIN32_WINNT >= 0x0601) && !defined(__MINGW64__) // windows 7 and later + // TODO: windows + arm64 + mingw64 + DWORD logicalProcessorCount = GetActiveProcessorCount(ALL_PROCESSOR_GROUPS); + os << " / " << logicalProcessorCount << " | " << llama_print_system_info(); +#else + os << " / " << std::thread::hardware_concurrency() << " | " << llama_print_system_info(); +#endif + + return os.str(); +} + +// +// String utils +// + +std::string string_format(const char * fmt, ...) { + va_list ap; + va_list ap2; + va_start(ap, fmt); + va_copy(ap2, ap); + int size = vsnprintf(NULL, 0, fmt, ap); + GGML_ASSERT(size >= 0 && size < INT_MAX); // NOLINT + std::vector buf(size + 1); + int size2 = vsnprintf(buf.data(), size + 1, fmt, ap2); + GGML_ASSERT(size2 == size); + va_end(ap2); + va_end(ap); + return std::string(buf.data(), size); +} + +std::string string_strip(const std::string & str) { + llama_common_unicode_string trimmed = llama_common_string_trim_rust((const uint8_t *) str.data(), str.size(), 0); + if (!trimmed.data) { + return ""; + } + std::string result((const char *) trimmed.data, trimmed.len); + llama_common_unicode_string_free(trimmed); + return result; +} + +std::string string_get_sortable_timestamp() { + using clock = std::chrono::system_clock; + + const clock::time_point current_time = clock::now(); + const time_t as_time_t = clock::to_time_t(current_time); + char timestamp_no_ns[100]; + std::strftime(timestamp_no_ns, 100, "%Y_%m_%d-%H_%M_%S", std::localtime(&as_time_t)); + + const int64_t ns = std::chrono::duration_cast( + current_time.time_since_epoch() % 1000000000).count(); + char timestamp_ns[11]; + snprintf(timestamp_ns, 11, "%09" PRId64, ns); + + return std::string(timestamp_no_ns) + "." + std::string(timestamp_ns); +} + +void string_replace_all(std::string & s, const std::string & search, const std::string & replace) { + llama_common_unicode_string replaced = llama_common_string_replace_all_rust( + (const uint8_t *) s.data(), s.size(), + (const uint8_t *) search.data(), search.size(), + (const uint8_t *) replace.data(), replace.size()); + if (!replaced.data) { + return; + } + s.assign((const char *) replaced.data, replaced.len); + llama_common_unicode_string_free(replaced); +} + +std::string regex_escape(const std::string & s) { + llama_common_unicode_string escaped = llama_common_regex_escape_rust((const uint8_t *) s.data(), s.size()); + if (!escaped.data) { + return ""; + } + std::string result((const char *) escaped.data, escaped.len); + llama_common_unicode_string_free(escaped); + return result; +} + +std::string string_join(const std::vector & values, const std::string & separator) { + std::vector views; + views.reserve(values.size()); + for (const std::string & value : values) { + views.push_back({ (const uint8_t *) value.data(), value.size() }); + } + llama_common_unicode_string joined = llama_common_string_join_rust( + views.data(), views.size(), + (const uint8_t *) separator.data(), separator.size()); + if (!joined.data) { + return ""; + } + std::string result((const char *) joined.data, joined.len); + llama_common_unicode_string_free(joined); + return result; +} + +std::vector string_split(const std::string & str, const std::string & delimiter) { + common_rust_string_list split = llama_common_string_split_rust( + (const uint8_t *) str.data(), str.size(), + (const uint8_t *) delimiter.data(), delimiter.size()); + std::vector parts; + if (!split.data) { + return parts; + } + parts.reserve(split.len); + for (size_t i = 0; i < split.len; ++i) { + const llama_common_unicode_string & item = split.data[i]; + parts.emplace_back((const char *) item.data, item.len); + } + llama_common_string_list_free(split); + return parts; +} + +std::string string_repeat(const std::string & str, size_t n) { + llama_common_unicode_string repeated = llama_common_string_repeat_rust((const uint8_t *) str.data(), str.size(), n); + if (!repeated.data) { + return ""; + } + std::string result((const char *) repeated.data, repeated.len); + llama_common_unicode_string_free(repeated); + return result; +} + +std::string string_from(bool value) { + return llama_common_bool_to_string_rust(value); +} + +std::string string_from(const std::vector & values) { + llama_common_unicode_string formatted = llama_common_int_vec_to_string_rust(values.data(), values.size()); + if (!formatted.data) { + return ""; + } + std::string result((const char *) formatted.data, formatted.len); + llama_common_unicode_string_free(formatted); + return result; +} + +std::string string_from(const struct llama_context * ctx, const std::vector & tokens) { + std::stringstream buf; + + buf << "[ "; + + bool first = true; + for (const auto & token : tokens) { + if (!first) { + buf << ", "; + } else { + first = false; + } + + auto detokenized = common_token_to_piece(ctx, token); + + buf << "'" << detokenized << "'" + << ":" << std::to_string(token); + } + + buf << " ]"; + + return buf.str(); +} + +std::string string_from(const struct llama_context * ctx, const struct llama_batch & batch) { + std::stringstream buf; + + buf << "[ "; + + bool first = true; + for (int i = 0; i < batch.n_tokens; ++i) { + if (!first) { + buf << ", "; + } else { + first = false; + } + + auto detokenized = common_token_to_piece(ctx, batch.token[i]); + + buf << "\n" << std::to_string(i) + << ", token '" << detokenized << "'" + << ", pos " << std::to_string(batch.pos[i]) + << ", n_seq_id " << std::to_string(batch.n_seq_id[i]) + << ", seq_id " << std::to_string(batch.seq_id[i][0]) + << ", logits " << std::to_string(batch.logits[i]); + } + + buf << " ]"; + + return buf.str(); +} + +void string_process_escapes(std::string & input) { + llama_common_unicode_string processed = llama_common_string_process_escapes_rust((const uint8_t *) input.data(), input.size()); + if (!processed.data) { + return; + } + input.assign((const char *) processed.data, processed.len); + llama_common_unicode_string_free(processed); +} + +bool string_parse_kv_override(const char * data, std::vector & overrides) { + const char * sep = strchr(data, '='); + if (sep == nullptr || sep - data >= 128) { + LOG_ERR("%s: malformed KV override '%s'\n", __func__, data); + return false; + } + llama_model_kv_override kvo; + std::strncpy(kvo.key, data, sep - data); + kvo.key[sep - data] = 0; + sep++; + if (strncmp(sep, "int:", 4) == 0) { + sep += 4; + kvo.tag = LLAMA_KV_OVERRIDE_TYPE_INT; + kvo.val_i64 = std::atol(sep); + } else if (strncmp(sep, "float:", 6) == 0) { + sep += 6; + kvo.tag = LLAMA_KV_OVERRIDE_TYPE_FLOAT; + kvo.val_f64 = std::atof(sep); + } else if (strncmp(sep, "bool:", 5) == 0) { + sep += 5; + kvo.tag = LLAMA_KV_OVERRIDE_TYPE_BOOL; + if (std::strcmp(sep, "true") == 0) { + kvo.val_bool = true; + } else if (std::strcmp(sep, "false") == 0) { + kvo.val_bool = false; + } else { + LOG_ERR("%s: invalid boolean value for KV override '%s'\n", __func__, data); + return false; + } + } else if (strncmp(sep, "str:", 4) == 0) { + sep += 4; + kvo.tag = LLAMA_KV_OVERRIDE_TYPE_STR; + if (strlen(sep) > 127) { + LOG_ERR("%s: malformed KV override '%s', value cannot exceed 127 chars\n", __func__, data); + return false; + } + strncpy(kvo.val_str, sep, 127); + kvo.val_str[127] = '\0'; + } else { + LOG_ERR("%s: invalid type for KV override '%s'\n", __func__, data); + return false; + } + overrides.emplace_back(std::move(kvo)); + return true; +} + +bool glob_match(const std::string & pattern, const std::string & str) { + return llama_common_glob_match_rust( + (const uint8_t *) pattern.data(), pattern.size(), + (const uint8_t *) str.data(), str.size()); +} + +// +// Filesystem utils +// + +// Validate if a filename is safe to use +// To validate a full path, split the path by the OS-specific path separator, and validate each part with this function +bool fs_validate_filename(const std::string & filename, bool allow_subdirs) { + return llama_common_fs_validate_filename_rust((const uint8_t *) filename.data(), filename.size(), allow_subdirs); +} + +#include + + +#ifdef _WIN32 +static std::wstring utf8_to_wstring(const std::string & str) { + if (str.empty()) { + return std::wstring(); + } + + int size = MultiByteToWideChar(CP_UTF8, 0, str.c_str(), (int)str.size(), NULL, 0); + + if (size <= 0) { + return std::wstring(); + } + + std::wstring wstr(size, 0); + MultiByteToWideChar(CP_UTF8, 0, str.c_str(), (int)str.size(), &wstr[0], size); + + return wstr; +} +#endif + +// returns true if successful, false otherwise +bool fs_create_directory_with_parents(const std::string & path) { +#ifdef _WIN32 + std::wstring wpath = utf8_to_wstring(path); + + // if the path already exists, check whether it's a directory + const DWORD attributes = GetFileAttributesW(wpath.c_str()); + if ((attributes != INVALID_FILE_ATTRIBUTES) && (attributes & FILE_ATTRIBUTE_DIRECTORY)) { + return true; + } + + size_t pos_slash = 0; + + // process path from front to back, procedurally creating directories + while ((pos_slash = path.find('\\', pos_slash)) != std::string::npos) { + const std::wstring subpath = wpath.substr(0, pos_slash); + + pos_slash += 1; + + // skip the drive letter, in some systems it can return an access denied error + if (subpath.length() == 2 && subpath[1] == ':') { + continue; + } + + const bool success = CreateDirectoryW(subpath.c_str(), NULL); + + if (!success) { + const DWORD error = GetLastError(); + + // if the path already exists, ensure that it's a directory + if (error == ERROR_ALREADY_EXISTS) { + const DWORD attributes = GetFileAttributesW(subpath.c_str()); + if (attributes == INVALID_FILE_ATTRIBUTES || !(attributes & FILE_ATTRIBUTE_DIRECTORY)) { + return false; + } + } else { + return false; + } + } + } + + return true; +#else + // if the path already exists, check whether it's a directory + struct stat info; + if (stat(path.c_str(), &info) == 0) { + return S_ISDIR(info.st_mode); + } + + size_t pos_slash = 1; // skip leading slashes for directory creation + + // process path from front to back, procedurally creating directories + while ((pos_slash = path.find('/', pos_slash)) != std::string::npos) { + const std::string subpath = path.substr(0, pos_slash); + struct stat info; + + // if the path already exists, ensure that it's a directory + if (stat(subpath.c_str(), &info) == 0) { + if (!S_ISDIR(info.st_mode)) { + return false; + } + } else { + // create parent directories + const int ret = mkdir(subpath.c_str(), 0755); + if (ret != 0) { + return false; + } + } + + pos_slash += 1; + } + + return true; +#endif // _WIN32 +} + +bool fs_is_directory(const std::string & path) { + std::filesystem::path dir(path); + return std::filesystem::exists(dir) && std::filesystem::is_directory(dir); +} + +std::string fs_get_cache_directory() { + std::string cache_directory = ""; + auto ensure_trailing_slash = [](std::string p) { + // Make sure to add trailing slash + if (p.back() != DIRECTORY_SEPARATOR) { + p += DIRECTORY_SEPARATOR; + } + return p; + }; + if (getenv("LLAMA_CACHE")) { + cache_directory = std::getenv("LLAMA_CACHE"); + } else { +#if defined(__linux__) || defined(__FreeBSD__) || defined(_AIX) || \ + defined(__OpenBSD__) || defined(__NetBSD__) + if (std::getenv("XDG_CACHE_HOME")) { + cache_directory = std::getenv("XDG_CACHE_HOME"); + } else if (std::getenv("HOME")) { + cache_directory = std::getenv("HOME") + std::string("/.cache/"); + } else { +#if defined(__linux__) + /* no $HOME is defined, fallback to getpwuid */ + struct passwd *pw = getpwuid(getuid()); + if ((!pw) || (!pw->pw_dir)) { + throw std::runtime_error("Failed to find $HOME directory"); + } + + cache_directory = std::string(pw->pw_dir) + std::string("/.cache/"); +#else /* defined(__linux__) */ + throw std::runtime_error("Failed to find $HOME directory"); +#endif /* defined(__linux__) */ + } +#elif defined(__APPLE__) + cache_directory = std::getenv("HOME") + std::string("/Library/Caches/"); +#elif defined(_WIN32) + cache_directory = std::getenv("LOCALAPPDATA"); +#elif defined(__EMSCRIPTEN__) + GGML_ABORT("not implemented on this platform"); +#else +# error Unknown architecture +#endif + cache_directory = ensure_trailing_slash(cache_directory); + cache_directory += "llama.cpp"; + } + return ensure_trailing_slash(cache_directory); +} + +std::string fs_get_cache_file(const std::string & filename) { + GGML_ASSERT(filename.find(DIRECTORY_SEPARATOR) == std::string::npos); + std::string cache_directory = fs_get_cache_directory(); + const bool success = fs_create_directory_with_parents(cache_directory); + if (!success) { + throw std::runtime_error("failed to create cache directory: " + cache_directory); + } + return cache_directory + filename; +} + +std::vector fs_list(const std::string & path, bool include_directories) { + std::vector files; + if (path.empty()) return files; + + std::filesystem::path dir(path); + if (!std::filesystem::exists(dir) || !std::filesystem::is_directory(dir)) { + return files; + } + + for (const auto & entry : std::filesystem::directory_iterator(dir)) { + try { + // Only include regular files (skip directories) + const auto & p = entry.path(); + if (std::filesystem::is_regular_file(p)) { + common_file_info info; + info.path = p.string(); + info.name = p.filename().string(); + info.is_dir = false; + try { + info.size = static_cast(std::filesystem::file_size(p)); + } catch (const std::filesystem::filesystem_error &) { + info.size = 0; + } + files.push_back(std::move(info)); + } else if (include_directories && std::filesystem::is_directory(p)) { + common_file_info info; + info.path = p.string(); + info.name = p.filename().string(); + info.size = 0; // Directories have no size + info.is_dir = true; + files.push_back(std::move(info)); + } + } catch (const std::filesystem::filesystem_error &) { + // skip entries we cannot inspect + continue; + } + } + + return files; +} + +// +// TTY utils +// + +bool tty_can_use_colors() { + // Check NO_COLOR environment variable (https://no-color.org/) + if (const char * no_color = std::getenv("NO_COLOR")) { + if (no_color[0] != '\0') { + return false; + } + } + + // Check TERM environment variable + if (const char * term = std::getenv("TERM")) { + if (std::strcmp(term, "dumb") == 0) { + return false; + } + } + + // Check if stdout and stderr are connected to a terminal + // We check both because log messages can go to either + bool stdout_is_tty = isatty(fileno(stdout)); + bool stderr_is_tty = isatty(fileno(stderr)); + + return stdout_is_tty || stderr_is_tty; +} + +// +// Model utils +// + +// TODO: move to common/sampling +static void common_init_sampler_from_model( + const llama_model * model, + common_params_sampling & sparams) { + + const uint64_t config = sparams.user_sampling_config; + + auto get_int32 = [&](const char * key, int32_t & dst, uint64_t user_config) { + if (config & user_config) { + return; + } + + char buf[64] = {0}; + if (llama_model_meta_val_str(model, key, buf, sizeof(buf)) > 0) { + char * end = nullptr; + int32_t v = strtol(buf, &end, 10); + if (end && end != buf) { + dst = v; + } + } + }; + + auto get_float = [&](const char * key, float & dst, uint64_t user_config) { + if (config & user_config) { + return; + } + + char buf[128] = {0}; + if (llama_model_meta_val_str(model, key, buf, sizeof(buf)) > 0) { + char * end = nullptr; + float v = strtof(buf, &end); + if (end && end != buf) { + dst = v; + } + } + }; + + // Sampling sequence + if (!(config & common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_SAMPLERS)) { + char buf[512] = {0}; + if (llama_model_meta_val_str(model, llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_SEQUENCE), buf, sizeof(buf)) > 0) { + const std::vector sampler_names = string_split(std::string(buf), ';'); + if (!sampler_names.empty()) { + sparams.samplers = common_sampler_types_from_names(sampler_names, true); + } + } + } + + get_int32(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_TOP_K), sparams.top_k, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TOP_K); + get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_TOP_P), sparams.top_p, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TOP_P); + get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_MIN_P), sparams.min_p, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIN_P); + get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_XTC_PROBABILITY), sparams.xtc_probability, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_XTC_PROBABILITY); + get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_XTC_THRESHOLD), sparams.xtc_threshold, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_XTC_THRESHOLD); + get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_TEMP), sparams.temp, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TEMP); + get_int32(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_PENALTY_LAST_N), sparams.penalty_last_n, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_LAST_N); + get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_PENALTY_REPEAT), sparams.penalty_repeat, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_REPEAT); + get_int32(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT), sparams.mirostat, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT); + get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_TAU), sparams.mirostat_tau, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_TAU); + get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_ETA), sparams.mirostat_eta, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_ETA); +} + +struct common_init_result::impl { + impl() = default; + ~impl() = default; + + // note: the order in which model, context, etc. are declared matters because their destructors will be called bottom-to-top + + llama_model_ptr model; + llama_context_ptr context; + + std::vector lora; + + std::vector samplers; + std::vector samplers_seq_config; +}; + +common_init_result::common_init_result(common_params & params) : + pimpl(new impl{}) { + auto mparams = common_model_params_to_llama(params); + auto cparams = common_context_params_to_llama(params); + + if (params.fit_params) { + LOG_INF("%s: fitting params to device memory, for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on\n", __func__); + common_fit_params(params.model.path.c_str(), &mparams, &cparams, + params.tensor_split, + params.tensor_buft_overrides.data(), + params.fit_params_target.data(), + params.fit_params_min_ctx, + params.verbosity >= 4 ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); + } + + llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams); + if (model == NULL) { + return; + } + + pimpl->model.reset(model); + + const llama_vocab * vocab = llama_model_get_vocab(model); + + // load and optionally apply lora adapters + for (auto & la : params.lora_adapters) { + llama_adapter_lora_ptr lora; + lora.reset(llama_adapter_lora_init(model, la.path.c_str())); + if (lora == nullptr) { + LOG_ERR("%s: failed to load lora adapter '%s'\n", __func__, la.path.c_str()); + pimpl->model.reset(model); + return; + } + + char buf[1024]; + la.ptr = lora.get(); + llama_adapter_meta_val_str(la.ptr, "adapter.lora.task_name", buf, sizeof(buf)); + la.task_name = buf; + llama_adapter_meta_val_str(la.ptr, "adapter.lora.prompt_prefix", buf, sizeof(buf)); + la.prompt_prefix = buf; + pimpl->lora.emplace_back(std::move(lora)); // copy to list of loaded adapters + } + + // updates params.sampling + // TODO: fix naming + common_init_sampler_from_model(model, params.sampling); + + if (params.sampling.ignore_eos && llama_vocab_eos(vocab) == LLAMA_TOKEN_NULL) { + LOG_WRN("%s: warning: vocab does not have an EOS token, ignoring --ignore-eos\n", __func__); + params.sampling.ignore_eos = false; + } + + // initialize once + for (llama_token i = 0; i < llama_vocab_n_tokens(vocab); i++) { + if (llama_vocab_is_eog(vocab, i)) { + LOG_INF("%s: added %s logit bias = %f\n", __func__, common_token_to_piece(vocab, i).c_str(), -INFINITY); + params.sampling.logit_bias_eog.push_back({i, -INFINITY}); + } + } + + if (params.sampling.ignore_eos) { + // add EOG biases to the active set of logit biases + params.sampling.logit_bias.insert( + params.sampling.logit_bias.end(), + params.sampling.logit_bias_eog.begin(), params.sampling.logit_bias_eog.end()); + } + + //if (params.sampling.penalty_last_n == -1) { + // LOG_INF("%s: setting penalty_last_n to ctx_size = %d\n", __func__, llama_n_ctx(lctx)); + // params.sampling.penalty_last_n = llama_n_ctx(lctx); + //} + + //if (params.sampling.dry_penalty_last_n == -1) { + // LOG_INF("%s: setting dry_penalty_last_n to ctx_size = %d\n", __func__, llama_n_ctx(lctx)); + // params.sampling.dry_penalty_last_n = llama_n_ctx(lctx); + //} + + // init the backend samplers as part of the context creation + pimpl->samplers.resize(cparams.n_seq_max); + pimpl->samplers_seq_config.resize(cparams.n_seq_max); + + for (int i = 0; i < (int) cparams.n_seq_max; ++i) { + pimpl->samplers[i].reset(common_sampler_init(model, params.sampling)); + pimpl->samplers_seq_config[i] = { i, common_sampler_get(pimpl->samplers[i].get()) }; + } + + if (params.sampling.backend_sampling) { + cparams.samplers = pimpl->samplers_seq_config.data(); + cparams.n_samplers = pimpl->samplers_seq_config.size(); + } + + llama_context * lctx = llama_init_from_model(model, cparams); + if (lctx == NULL) { + LOG_ERR("%s: failed to create context with model '%s'\n", __func__, params.model.path.c_str()); + return; + } + + pimpl->context.reset(lctx); +} + +llama_model * common_init_result::model() { + return pimpl->model.get(); +} + +llama_context * common_init_result::context() { + return pimpl->context.get(); +} + +common_sampler * common_init_result::sampler(llama_seq_id seq_id) { + if (seq_id < 0 || seq_id >= (int) pimpl->samplers.size()) { + return nullptr; + } + return pimpl->samplers[seq_id].get(); +} + +void common_init_result::reset_samplers() { + for (int i = 0; i < (int) pimpl->samplers.size(); ++i) { + llama_sampler_reset(common_sampler_get(pimpl->samplers[i].get())); + } +} + +std::vector & common_init_result::lora() { + return pimpl->lora; +} + +common_init_result_ptr common_init_from_params(common_params & params) { + common_init_result_ptr res(new common_init_result(params)); + + llama_model * model = res->model(); + if (model == NULL) { + LOG_ERR("%s: failed to load model '%s'\n", __func__, params.model.path.c_str()); + return res; + } + + llama_context * lctx = res->context(); + if (lctx == NULL) { + LOG_ERR("%s: failed to create context with model '%s'\n", __func__, params.model.path.c_str()); + return res; + } + + const llama_vocab * vocab = llama_model_get_vocab(model); + + if (params.ctx_shift && !llama_memory_can_shift(llama_get_memory(lctx))) { + LOG_WRN("%s: KV cache shifting is not supported for this context, disabling KV cache shifting\n", __func__); + params.ctx_shift = false; + } + + if (!params.control_vectors.empty()) { + if (params.control_vector_layer_start <= 0) params.control_vector_layer_start = 1; + if (params.control_vector_layer_end <= 0) params.control_vector_layer_end = llama_model_n_layer(model); + + const auto cvec = common_control_vector_load(params.control_vectors); + if (cvec.n_embd == -1) { + return res; + } + + int err = llama_set_adapter_cvec( + lctx, + cvec.data.data(), + cvec.data.size(), + cvec.n_embd, + params.control_vector_layer_start, + params.control_vector_layer_end); + if (err) { + return res; + } + } + + if (llama_pooling_type(lctx) == LLAMA_POOLING_TYPE_RANK) { + bool ok = true; + + if (llama_vocab_bos(vocab) == LLAMA_TOKEN_NULL) { + LOG_WRN("%s: warning: vocab does not have a BOS token, reranking will not work\n", __func__); + ok = false; + } + + bool has_eos = llama_vocab_eos(vocab) != LLAMA_TOKEN_NULL; + bool has_sep = llama_vocab_sep(vocab) != LLAMA_TOKEN_NULL; + bool has_rerank_prompt = llama_model_chat_template(model, "rerank") != NULL; + + if (!has_eos && !has_sep && !has_rerank_prompt) { + LOG_WRN("%s: warning: vocab does not have an EOS token, SEP token, or rerank prompt. Reranking will not work\n", __func__); + ok = false; + } else if (!has_eos) { + LOG_WRN("%s: warning: vocab does not have an EOS token, using SEP token as fallback\n", __func__); + } + + if (!ok) { + return res; + } + } + + if (!params.lora_init_without_apply) { + common_set_adapter_lora(lctx, params.lora_adapters); + } + + if (params.warmup) { + LOG_WRN("%s: warming up the model with an empty run - please wait ... (--no-warmup to disable)\n", __func__); + + llama_set_warmup(lctx, true); + + std::vector tmp; + llama_token bos = llama_vocab_bos(vocab); + llama_token eos = llama_vocab_eos(vocab); + + // some models (e.g. T5) don't have a BOS token + if (bos != LLAMA_TOKEN_NULL) { + tmp.push_back(bos); + } + if (eos != LLAMA_TOKEN_NULL) { + tmp.push_back(eos); + } + if (tmp.empty()) { + tmp.push_back(0); + } + + if (llama_model_has_encoder(model)) { + llama_encode(lctx, llama_batch_get_one(tmp.data(), tmp.size())); + llama_token decoder_start_token_id = llama_model_decoder_start_token(model); + if (decoder_start_token_id == LLAMA_TOKEN_NULL) { + decoder_start_token_id = bos; + } + tmp.clear(); + tmp.push_back(decoder_start_token_id); + } + if (llama_model_has_decoder(model)) { + llama_decode(lctx, llama_batch_get_one(tmp.data(), std::min(tmp.size(), (size_t) params.n_batch))); + } + llama_memory_clear(llama_get_memory(lctx), true); + llama_synchronize(lctx); + llama_perf_context_reset(lctx); + llama_set_warmup(lctx, false); + + // reset samplers to reset RNG state after warmup to the seeded state + res->reset_samplers(); + } + + return res; +} + +common_init_result::~common_init_result() = default; + +std::string common_get_model_endpoint() { + const char * model_endpoint_env = getenv("MODEL_ENDPOINT"); + // We still respect the use of environment-variable "HF_ENDPOINT" for backward-compatibility. + const char * hf_endpoint_env = getenv("HF_ENDPOINT"); + const char * endpoint_env = model_endpoint_env ? model_endpoint_env : hf_endpoint_env; + std::string model_endpoint = "https://huggingface.co/"; + if (endpoint_env) { + model_endpoint = endpoint_env; + if (model_endpoint.back() != '/') { + model_endpoint += '/'; + } + } + return model_endpoint; +} + +common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx) { + auto * mem = llama_get_memory(ctx); + if (mem == nullptr) { + return COMMON_CONTEXT_SEQ_RM_TYPE_NO; + } + + common_context_seq_rm_type res = COMMON_CONTEXT_SEQ_RM_TYPE_PART; + + llama_memory_clear(mem, true); + + // eval 2 tokens to check if the context is compatible + std::vector tmp; + tmp.push_back(0); + tmp.push_back(0); + + int ret = llama_decode(ctx, llama_batch_get_one(tmp.data(), tmp.size())); + if (ret != 0) { + LOG_ERR("%s: llama_decode() failed: %d\n", __func__, ret); + res = COMMON_CONTEXT_SEQ_RM_TYPE_NO; + goto done; + } + + // try to remove the last tokens + if (!llama_memory_seq_rm(mem, 0, 1, -1)) { + LOG_WRN("%s: the target context does not support partial sequence removal\n", __func__); + res = COMMON_CONTEXT_SEQ_RM_TYPE_FULL; + goto done; + } + +done: + llama_memory_clear(mem, true); + llama_synchronize(ctx); + + return res; +} + +void common_set_adapter_lora(struct llama_context * ctx, std::vector & lora) { + std::vector loras; + std::vector scales; + + for (auto & la: lora) { + loras.push_back(la.ptr); + scales.push_back(la.scale); + } + + llama_set_adapters_lora(ctx, loras.data(), loras.size(), scales.data()); +} + +struct llama_model_params common_model_params_to_llama(common_params & params) { + auto mparams = llama_model_default_params(); + + if (!params.devices.empty()) { + mparams.devices = params.devices.data(); + } + + mparams.n_gpu_layers = params.n_gpu_layers; + mparams.main_gpu = params.main_gpu; + mparams.split_mode = params.split_mode; + mparams.tensor_split = params.tensor_split; + mparams.use_mmap = params.use_mmap; + mparams.use_direct_io = params.use_direct_io; + mparams.use_mlock = params.use_mlock; + mparams.check_tensors = params.check_tensors; + mparams.use_extra_bufts = !params.no_extra_bufts; + mparams.no_host = params.no_host; + + if (params.kv_overrides.empty()) { + mparams.kv_overrides = NULL; + } else { + GGML_ASSERT(params.kv_overrides.back().key[0] == 0 && "KV overrides not terminated with empty key"); + mparams.kv_overrides = params.kv_overrides.data(); + } + + if (params.tensor_buft_overrides.empty()) { + mparams.tensor_buft_overrides = NULL; + } else { + GGML_ASSERT(params.tensor_buft_overrides.back().pattern == nullptr && "Tensor buffer overrides not terminated with empty pattern"); + mparams.tensor_buft_overrides = params.tensor_buft_overrides.data(); + } + + mparams.progress_callback = params.load_progress_callback; + mparams.progress_callback_user_data = params.load_progress_callback_user_data; + mparams.no_alloc = params.no_alloc; + + return mparams; +} + +struct llama_context_params common_context_params_to_llama(const common_params & params) { + auto cparams = llama_context_default_params(); + + cparams.n_ctx = params.n_ctx; + cparams.n_seq_max = params.n_parallel; + cparams.n_batch = params.n_batch; + cparams.n_ubatch = params.n_ubatch; + cparams.n_threads = params.cpuparams.n_threads; + cparams.n_threads_batch = params.cpuparams_batch.n_threads == -1 ? + params.cpuparams.n_threads : params.cpuparams_batch.n_threads; + cparams.embeddings = params.embedding; + cparams.rope_scaling_type = params.rope_scaling_type; + cparams.rope_freq_base = params.rope_freq_base; + cparams.rope_freq_scale = params.rope_freq_scale; + cparams.yarn_ext_factor = params.yarn_ext_factor; + cparams.yarn_attn_factor = params.yarn_attn_factor; + cparams.yarn_beta_fast = params.yarn_beta_fast; + cparams.yarn_beta_slow = params.yarn_beta_slow; + cparams.yarn_orig_ctx = params.yarn_orig_ctx; + cparams.pooling_type = params.pooling_type; + cparams.attention_type = params.attention_type; + cparams.flash_attn_type = params.flash_attn_type; + cparams.cb_eval = params.cb_eval; + cparams.cb_eval_user_data = params.cb_eval_user_data; + cparams.offload_kqv = !params.no_kv_offload; + cparams.no_perf = params.no_perf; + cparams.op_offload = !params.no_op_offload; + cparams.swa_full = params.swa_full; + cparams.kv_unified = params.kv_unified; + + cparams.type_k = params.cache_type_k; + cparams.type_v = params.cache_type_v; + + return cparams; +} + +struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const common_cpu_params & params) { + struct ggml_threadpool_params tpp; + + ggml_threadpool_params_init(&tpp, params.n_threads); // setup the defaults + + if (params.mask_valid) { + std::memcpy(&tpp.cpumask, ¶ms.cpumask, GGML_MAX_N_THREADS); + } + + tpp.prio = params.priority; + tpp.poll = params.poll; + tpp.strict_cpu = params.strict_cpu; + + return tpp; +} + +// +// Batch utils +// + +void common_batch_clear(struct llama_batch & batch) { + batch.n_tokens = 0; +} + +void common_batch_add( + struct llama_batch & batch, + llama_token id, + llama_pos pos, + const std::vector & seq_ids, + bool logits) { + GGML_ASSERT(batch.seq_id[batch.n_tokens] && "llama_batch size exceeded"); + + batch.token [batch.n_tokens] = id; + batch.pos [batch.n_tokens] = pos; + batch.n_seq_id[batch.n_tokens] = seq_ids.size(); + for (size_t i = 0; i < seq_ids.size(); ++i) { + batch.seq_id[batch.n_tokens][i] = seq_ids[i]; + } + batch.logits [batch.n_tokens] = logits; + + batch.n_tokens++; +} + +// +// Vocab utils +// + +std::vector common_tokenize( + const struct llama_context * ctx, + const std::string & text, + bool add_special, + bool parse_special) { + const llama_model * model = llama_get_model(ctx); + const llama_vocab * vocab = llama_model_get_vocab(model); + return common_tokenize(vocab, text, add_special, parse_special); +} + +std::vector common_tokenize( + const struct llama_vocab * vocab, + const std::string & text, + bool add_special, + bool parse_special) { + // upper limit for the number of tokens + int n_tokens = text.length() + 2 * add_special; + std::vector result(n_tokens); + n_tokens = llama_tokenize(vocab, text.data(), text.length(), result.data(), result.size(), add_special, parse_special); + if (n_tokens == std::numeric_limits::min()) { + throw std::runtime_error("Tokenization failed: input text too large, tokenization result exceeds int32_t limit"); + } + if (n_tokens < 0) { + result.resize(-n_tokens); + int check = llama_tokenize(vocab, text.data(), text.length(), result.data(), result.size(), add_special, parse_special); + GGML_ASSERT(check == -n_tokens); + } else { + result.resize(n_tokens); + } + return result; +} + +std::string common_token_to_piece(const struct llama_context * ctx, llama_token token, bool special) { + const llama_model * model = llama_get_model(ctx); + const llama_vocab * vocab = llama_model_get_vocab(model); + return common_token_to_piece(vocab, token, special); +} + +std::string common_token_to_piece(const struct llama_vocab * vocab, llama_token token, bool special) { + std::string piece; + piece.resize(piece.capacity()); // using string internal cache, 15 bytes + '\n' + const int n_chars = llama_token_to_piece(vocab, token, &piece[0], piece.size(), 0, special); + if (n_chars < 0) { + piece.resize(-n_chars); + int check = llama_token_to_piece(vocab, token, &piece[0], piece.size(), 0, special); + GGML_ASSERT(check == -n_chars); + } + else { + piece.resize(n_chars); + } + + return piece; +} + +std::string common_detokenize(const struct llama_context * ctx, const std::vector & tokens, bool special) { + const llama_model * model = llama_get_model(ctx); + const llama_vocab * vocab = llama_model_get_vocab(model); + return common_detokenize(vocab, tokens, special); +} + +std::string common_detokenize(const struct llama_vocab * vocab, const std::vector & tokens, bool special) { + std::string text; + text.resize(std::max(text.capacity(), tokens.size())); + int32_t n_chars = llama_detokenize(vocab, tokens.data(), (int32_t)tokens.size(), &text[0], (int32_t)text.size(), false, special); + if (n_chars < 0) { + text.resize(-n_chars); + n_chars = llama_detokenize(vocab, tokens.data(), (int32_t)tokens.size(), &text[0], (int32_t)text.size(), false, special); + GGML_ASSERT(n_chars <= (int32_t)text.size()); // whitespace trimming is performed after per-token detokenization + } + + text.resize(n_chars); + + // NOTE: the original tokenizer decodes bytes after collecting the pieces. + return text; +} + +// +// Embedding utils +// + +void common_embd_normalize(const float * inp, float * out, int n, int embd_norm) { + llama_common_embd_normalize_rust(inp, out, n, embd_norm); +} + +float common_embd_similarity_cos(const float * embd1, const float * embd2, int n){ + return llama_common_embd_similarity_cos_rust(embd1, embd2, n); +} + +// +// Control vector utils +// + +static common_control_vector_data common_control_vector_load_one(const common_control_vector_load_info & load_info) { + common_control_vector_data result = { -1, {} }; + + ggml_context * ctx = nullptr; + struct gguf_init_params meta_gguf_params = { + /* .no_alloc = */ false, + /* .ctx = */ &ctx, + }; + struct gguf_context * ctx_gguf = gguf_init_from_file(load_info.fname.c_str(), meta_gguf_params); + if (!ctx_gguf) { + LOG_ERR("%s: failed to load control vector file from %s\n", __func__, load_info.fname.c_str()); + return result; + } + + int32_t n_tensors = gguf_get_n_tensors(ctx_gguf); + if (n_tensors == 0) { + LOG_WRN("%s: no direction tensors found in %s\n", __func__, load_info.fname.c_str()); + } + + for (int i = 0; i < n_tensors; i++) { + std::string name = gguf_get_tensor_name(ctx_gguf, i); + + int layer_idx = -1; + + // split on '.' + size_t dotpos = name.find('.'); + if (dotpos != std::string::npos && name.substr(0, dotpos) == "direction") { + try { + layer_idx = std::stoi(name.substr(dotpos + 1)); + } catch (...) { + layer_idx = -1; + } + } + if (layer_idx < 0) { + LOG_ERR("%s: invalid/unparsable direction tensor layer index in %s\n", __func__, load_info.fname.c_str()); + result.n_embd = -1; + break; + } else if (layer_idx == 0) { + LOG_ERR("%s: invalid (zero) direction tensor layer index in %s\n", __func__, load_info.fname.c_str()); + result.n_embd = -1; + break; + } + + struct ggml_tensor * tensor = ggml_get_tensor(ctx, name.c_str()); + if (tensor->type != GGML_TYPE_F32) { + LOG_ERR("%s: invalid (non-F32) direction tensor type in %s\n", __func__, load_info.fname.c_str()); + result.n_embd = -1; + break; + } + if (ggml_n_dims(tensor) != 1) { + LOG_ERR("%s: invalid (non-1D) direction tensor shape in %s\n", __func__, load_info.fname.c_str()); + result.n_embd = -1; + break; + } + + if (result.n_embd == -1) { + result.n_embd = ggml_nelements(tensor); + } else if (ggml_nelements(tensor) != result.n_embd) { + LOG_ERR("%s: direction tensor in %s does not match previous dimensions\n", __func__, load_info.fname.c_str()); + result.n_embd = -1; + break; + } + + // extend if necessary - do not store data for layer 0 (it's not used) + result.data.resize(std::max(result.data.size(), static_cast(result.n_embd * layer_idx)), 0.0f); + + const float * src = (const float *) tensor->data; + float * dst = result.data.data() + result.n_embd * (layer_idx - 1); // layer 1 at [0] + for (int j = 0; j < result.n_embd; j++) { + dst[j] += src[j] * load_info.strength; // allows multiple directions for same layer in same file + } + + } + + if (result.n_embd == -1) { + LOG_WRN("%s: skipping %s due to invalid direction tensors\n", __func__, load_info.fname.c_str()); + result.data.clear(); + } + + gguf_free(ctx_gguf); + ggml_free(ctx); + + return result; +} + +common_control_vector_data common_control_vector_load(const std::vector & load_infos) { + common_control_vector_data result = { -1, {} }; + + for (const auto & info : load_infos) { + auto cur = common_control_vector_load_one(info); + + if (cur.n_embd == -1) { + result.n_embd = -1; + break; + } + if (result.n_embd != -1 && result.n_embd != cur.n_embd) { + LOG_ERR("%s: control vectors in %s does not match previous dimensions\n", __func__, info.fname.c_str()); + result.n_embd = -1; + break; + } + + if (result.n_embd == -1) { + result = std::move(cur); + } else { + result.data.resize(std::max(result.data.size(), cur.data.size()), 0.0f); // extend if necessary + for (size_t i = 0; i < cur.data.size(); i++) { + result.data[i] += cur.data[i]; + } + } + } + + if (result.n_embd == -1) { + LOG_ERR("%s: no valid control vector files passed\n", __func__); + result.data.clear(); + } + + return result; +} + +ggml_opt_dataset_t common_opt_dataset_init(struct llama_context * ctx, const std::vector & tokens, int64_t stride) { + const int64_t ne_datapoint = llama_n_ctx(ctx); + const int64_t ndata = (tokens.size() - ne_datapoint - 1) / stride; + ggml_opt_dataset_t result = ggml_opt_dataset_init( + GGML_TYPE_I32, GGML_TYPE_I32, ne_datapoint, ne_datapoint, ndata, /*ndata_shard =*/ 1); + + llama_token * data = (llama_token *) ggml_opt_dataset_data(result)->data; + llama_token * labels = (llama_token *) ggml_opt_dataset_labels(result)->data; + + for (int64_t idata = 0; idata < ndata; ++idata) { + memcpy(data + idata*ne_datapoint, tokens.data() + idata*stride + 0, ne_datapoint*sizeof(llama_token)); + memcpy(labels + idata*ne_datapoint, tokens.data() + idata*stride + 1, ne_datapoint*sizeof(llama_token)); + } + + return result; +} + +ggml_opt_optimizer_params common_opt_lr_pars(void * userdata) { + ggml_opt_optimizer_params result = ggml_opt_get_default_optimizer_params(nullptr); + const lr_opt & d = *(lr_opt *) userdata; + result.adamw.alpha = result.sgd.alpha = d.get_lr(d.epoch); + result.sgd.wd = result.adamw.wd = d.wd; + return result; +} + +enum ggml_opt_optimizer_type common_opt_get_optimizer(const char * n) { + return (enum ggml_opt_optimizer_type) llama_common_opt_get_optimizer_rust(n); +} + +void lr_opt::init() { + llama_common_lr_opt_init_rust(lr0, lr_min, decay_epochs, epochs, &decay_epochs, &scale_epoch); +} + +float lr_opt::get_lr(float epoch) const { + float r = llama_common_lr_opt_get_lr_rust(lr0, lr_min, decay_epochs, scale_epoch, epoch); + LOG_INF("epoch %.2g lr=%.2g\n", epoch, r); + return r; +} + +bool common_replay_last_token(struct llama_context * ctx, llama_token last_token, int32_t pos) { + llama_batch batch = llama_batch_get_one(&last_token, 1); + batch.pos = &pos; + if (llama_decode(ctx, batch)) { + LOG_ERR("%s: failed to replay last token\n", __func__); + return false; + } + return true; +} + +bool common_prompt_batch_decode( + struct llama_context * ctx, + const std::vector & tokens, + int & n_past, + int n_batch, + std::string_view state_path, + bool save_state) { + const int n_eval = tokens.size(); + if (n_eval == 0) { + return true; + } + + if (save_state && n_eval > 1) { + const int n_tokens_before_last = n_eval - 1; + + GGML_ASSERT(n_eval <= n_batch); + + // Decode all but the last token so we can save the memory state before decoding the last token. + // This is done so we can restore the session state later and replay the last token. + // Memory implementations in recurrent/hybrid models don't support removing tokens from their + // memory, so we can't just remove the last token from the memory and replay the last token which + // is the reason for this logic. + if (llama_decode(ctx, llama_batch_get_one(const_cast(tokens.data()), n_tokens_before_last))) { + LOG_ERR("%s : failed to eval\n", __func__); + return false; + } + n_past += n_tokens_before_last; + + llama_state_save_file(ctx, state_path.data(), tokens.data(), n_tokens_before_last); + LOG_INF("saved session before last token to %s, n_tokens = %d\n", state_path.data(), n_tokens_before_last); + + llama_token last_token = tokens.back(); + llama_batch batch = llama_batch_get_one(&last_token, 1); + int32_t pos = n_past; + batch.pos = &pos; + + if (llama_decode(ctx, batch)) { + LOG_ERR("%s : failed to eval last token\n", __func__); + return false; + } + n_past++; + } else { + if (llama_decode(ctx, llama_batch_get_one(const_cast(tokens.data()), n_eval))) { + LOG_ERR("%s : failed to eval\n", __func__); + return false; + } + n_past += n_eval; + } + + return true; +} + +#include "common.h.inc" +#include "common.h.inc" +// fix problem with std::min and std::max +#if defined(_WIN32) +#define WIN32_LEAN_AND_MEAN +#ifndef NOMINMAX +# define NOMINMAX +#endif +#include +#endif + +#define JSON_ASSERT GGML_ASSERT +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include // for hardware_concurrency +#include + +extern "C" { +int llama_build_number(void); +const char * llama_commit(void); +const char * llama_compiler(void); +const char * llama_build_target(void); +} + +#ifndef __EMSCRIPTEN__ +#ifdef __linux__ +#include +#elif defined(_WIN32) +# if !defined(PATH_MAX) +# define PATH_MAX MAX_PATH +# endif +#elif defined(_AIX) +#include +#else +#include +#endif +#endif + +#define LLAMA_MAX_URL_LENGTH 2084 // Maximum URL Length in Chrome: 2083 + +extern const char * LICENSES[]; + +using json = nlohmann::ordered_json; +using namespace common_arg_utils; + +static std::initializer_list mmproj_examples = { + LLAMA_EXAMPLE_MTMD, + LLAMA_EXAMPLE_SERVER, + LLAMA_EXAMPLE_CLI, +}; + +static std::string read_file(const std::string & fname) { + std::ifstream file(fname); + if (!file) { + throw std::runtime_error(string_format("error: failed to open file '%s'\n", fname.c_str())); + } + std::string content((std::istreambuf_iterator(file)), std::istreambuf_iterator()); + file.close(); + return content; +} + +static const std::vector & get_common_arg_defs() { + static const std::vector options = [] { + common_params params; + auto ctx = common_params_parser_init(params, LLAMA_EXAMPLE_SERVER, nullptr); + return ctx.options; + }(); + return options; +} + +common_arg & common_arg::set_examples(std::initializer_list examples) { + this->examples = examples; + return *this; +} + +common_arg & common_arg::set_excludes(std::initializer_list excludes) { + this->excludes = excludes; + return *this; +} + +common_arg & common_arg::set_env(const char * env) { + help = help + "\n(env: " + env + ")"; + this->env = env; + return *this; +} + +common_arg & common_arg::set_sampling() { + is_sampling = true; + return *this; +} + +common_arg & common_arg::set_spec() { + is_spec = true; + return *this; +} + +common_arg & common_arg::set_preset_only() { + is_preset_only = true; + return *this; +} + +bool common_arg::in_example(enum llama_example ex) { + return examples.find(ex) != examples.end(); +} + +bool common_arg::is_exclude(enum llama_example ex) { + return excludes.find(ex) != excludes.end(); +} + +bool common_arg::get_value_from_env(std::string & output) const { + if (env == nullptr) return false; + if (!args_neg.empty()) { + // for compatibility, we need to check LLAMA_ARG_NO_ env as well + std::string neg_env = env; + string_replace_all(neg_env, "LLAMA_ARG_", "LLAMA_ARG_NO_"); + char * neg_value = std::getenv(neg_env.c_str()); + if (neg_value) { + output = "0"; // falsey + return true; + } + } + char * value = std::getenv(env); + if (value) { + output = value; + return true; + } + return false; +} + +bool common_arg::has_value_from_env() const { + if (env != nullptr && !args_neg.empty()) { + // for compatibility, we need to check LLAMA_ARG_NO_ env as well + std::string neg_env = env; + string_replace_all(neg_env, "LLAMA_ARG_", "LLAMA_ARG_NO_"); + if (std::getenv(neg_env.c_str())) { + return true; + } + } + return env != nullptr && std::getenv(env); +} + +static std::vector break_str_into_lines(std::string input, size_t max_char_per_line) { + common_rust_string_list split = + llama_common_break_str_into_lines_rust((const uint8_t *) input.data(), input.size(), max_char_per_line); + std::vector result; + if (!split.data) { + return result; + } + result.reserve(split.len); + for (size_t i = 0; i < split.len; ++i) { + const llama_common_unicode_string & item = split.data[i]; + result.emplace_back((const char *) item.data, item.len); + } + llama_common_string_list_free(split); + return result; +} + +std::string common_arg::to_string() const { + // params for printing to console + const static int n_leading_spaces = 40; + const static int n_char_per_line_help = 70; // TODO: detect this based on current console + std::string leading_spaces(n_leading_spaces, ' '); + + std::ostringstream ss; + auto all_args = get_args(); // also contains args_neg + for (const auto & arg : all_args) { + if (arg == all_args.front()) { + if (all_args.size() == 1) { + ss << arg; + } else { + // first arg is usually abbreviation, we need padding to make it more beautiful + auto tmp = std::string(arg) + ", "; + auto spaces = std::string(std::max(0, 7 - (int)tmp.size()), ' '); + ss << tmp << spaces; + } + } else { + ss << arg << (arg != all_args.back() ? ", " : ""); + } + } + if (value_hint) ss << " " << value_hint; + if (value_hint_2) ss << " " << value_hint_2; + if (ss.tellp() > n_leading_spaces - 3) { + // current line is too long, add new line + ss << "\n" << leading_spaces; + } else { + // padding between arg and help, same line + ss << std::string(leading_spaces.size() - ss.tellp(), ' '); + } + const auto help_lines = break_str_into_lines(help, n_char_per_line_help); + for (const auto & line : help_lines) { + ss << (&line == &help_lines.front() ? "" : leading_spaces) << line << "\n"; + } + return ss.str(); +} + +std::vector common_arg::get_args() const { + std::vector result; + for (const auto & arg : args) { + result.push_back(std::string(arg)); + } + for (const auto & arg : args_neg) { + result.push_back(std::string(arg)); + } + return result; +} + +std::vector common_arg::get_env() const { + std::vector result; + if (env) { + result.push_back(std::string(env)); + } + if (!args_neg.empty() && env) { + // for compatibility, we need to add LLAMA_ARG_NO_ variant + std::string neg_env = env; + string_replace_all(neg_env, "LLAMA_ARG_", "LLAMA_ARG_NO_"); + result.push_back(neg_env); + } + return result; +} + +// +// utils +// + +// Helper function to parse tensor buffer override strings +static void parse_tensor_buffer_overrides(const std::string & value, std::vector & overrides) { + std::map buft_list; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + auto * dev = ggml_backend_dev_get(i); + auto * buft = ggml_backend_dev_buffer_type(dev); + if (buft) { + buft_list[ggml_backend_buft_name(buft)] = buft; + } + } + + for (const auto & override : string_split(value, ',')) { + std::string::size_type pos = override.find('='); + if (pos == std::string::npos) { + throw std::invalid_argument("invalid value"); + } + std::string tensor_name = override.substr(0, pos); + std::string buffer_type = override.substr(pos + 1); + + if (buft_list.find(buffer_type) == buft_list.end()) { + printf("Available buffer types:\n"); + for (const auto & it : buft_list) { + printf(" %s\n", ggml_backend_buft_name(it.second)); + } + throw std::invalid_argument("unknown buffer type"); + } + // keep strings alive and avoid leaking memory by storing them in a static vector + static std::list buft_overrides; + buft_overrides.push_back(tensor_name); + overrides.push_back({buft_overrides.back().c_str(), buft_list.at(buffer_type)}); + } +} + +static bool common_params_handle_remote_preset(common_params & params, llama_example ex) { + GGML_ASSERT(!params.model.hf_repo.empty()); + + // the returned hf_repo is without tag + auto [hf_repo, hf_tag] = common_download_split_repo_tag(params.model.hf_repo); + + // "latest" tag (default if not specified) is translated to "default" preset + if (hf_tag == "latest") { + hf_tag = "default"; + } + + std::string model_endpoint = common_get_model_endpoint(); + auto preset_url = model_endpoint + hf_repo + "/resolve/main/preset.ini"; + + // prepare local path for caching + const std::string preset_repo = hf_repo + "_preset.ini"; + auto preset_fname = common_rust_string_to_std( + llama_common_clean_file_name_rust((const uint8_t *) preset_repo.data(), preset_repo.size())); + auto preset_path = fs_get_cache_file(preset_fname); + common_download_opts opts; + opts.bearer_token = params.hf_token; + opts.offline = params.offline; + const int status = common_download_file_single(preset_url, preset_path, opts); + const bool has_preset = status >= 200 && status < 400; + + // remote preset is optional, so we don't error out if not found + if (has_preset) { + LOG_INF("applying remote preset from %s\n", preset_url.c_str()); + common_preset_context ctx(ex, /* only_remote_allowed */ true); + common_preset global; + auto remote_presets = ctx.load_from_ini(preset_path, global); + remote_presets = ctx.cascade(global, remote_presets); + if (remote_presets.find(hf_tag) != remote_presets.end()) { + common_preset preset = remote_presets.at(hf_tag); + LOG_INF("\n%s", preset.to_ini().c_str()); // to_ini already added trailing newline + preset.apply_to_params(params); + } else { + throw std::runtime_error("Remote preset.ini does not contain [" + std::string(hf_tag) + "] section"); + } + } else { + LOG_INF("%s", "no remote preset found, skipping\n"); + } + + return has_preset; +} + +struct handle_model_result { + bool found_mmproj = false; + common_params_model mmproj; +}; + +static handle_model_result common_params_handle_model(struct common_params_model & model, + const std::string & bearer_token, + bool offline) { + handle_model_result result; + + if (!model.docker_repo.empty()) { + model.path = common_docker_resolve_model(model.docker_repo); + model.name = model.docker_repo; + } else if (!model.hf_repo.empty()) { + // If -m was used with -hf, treat the model "path" as the hf_file to download + if (model.hf_file.empty() && !model.path.empty()) { + model.hf_file = model.path; + model.path = ""; + } + common_download_opts opts; + opts.bearer_token = bearer_token; + opts.offline = offline; + auto download_result = common_download_model(model, opts, true); + + if (download_result.model_path.empty()) { + LOG_ERR("error: failed to download model from Hugging Face\n"); + exit(1); + } + + model.name = model.hf_repo; + model.path = download_result.model_path; + + if (!download_result.mmproj_path.empty()) { + result.found_mmproj = true; + result.mmproj.path = download_result.mmproj_path; + } + } else if (!model.url.empty()) { + if (model.path.empty()) { + auto f = string_split(model.url, '#').front(); + f = string_split(f, '?').front(); + model.path = fs_get_cache_file(string_split(f, '/').back()); + } + + common_download_opts opts; + opts.bearer_token = bearer_token; + opts.offline = offline; + auto download_result = common_download_model(model, opts); + if (download_result.model_path.empty()) { + LOG_ERR("error: failed to download model from %s\n", model.url.c_str()); + exit(1); + } + } + + return result; +} + +const std::vector kv_cache_types = { + GGML_TYPE_F32, + GGML_TYPE_F16, + GGML_TYPE_BF16, + GGML_TYPE_Q8_0, + GGML_TYPE_Q4_0, + GGML_TYPE_Q4_1, + GGML_TYPE_IQ4_NL, + GGML_TYPE_Q5_0, + GGML_TYPE_Q5_1, +}; + +static ggml_type kv_cache_type_from_str(const std::string & s) { + const int parsed = llama_common_kv_cache_type_from_str_rust(s.c_str()); + if (parsed >= 0) { + return (ggml_type) parsed; + } + throw std::runtime_error("Unsupported cache type: " + s); +} + +// +// CLI argument parsing functions +// + +static bool common_params_parse_ex(int argc, char ** argv, common_params_context & ctx_arg) { + common_params & params = ctx_arg.params; + + // setup log directly from params.verbosity: see tools/cli/cli.cpp + common_log_set_verbosity_thold(params.verbosity); + + std::unordered_map> arg_to_options; + for (auto & opt : ctx_arg.options) { + for (const auto & arg : opt.args) { + arg_to_options[arg] = {&opt, /* is_positive */ true}; + } + for (const auto & arg : opt.args_neg) { + arg_to_options[arg] = {&opt, /* is_positive */ false}; + } + } + + // handle environment variables + for (auto & opt : ctx_arg.options) { + std::string value; + if (opt.get_value_from_env(value)) { + try { + if (opt.handler_void && is_truthy(value)) { + opt.handler_void(params); + } + if (opt.handler_int) { + opt.handler_int(params, std::stoi(value)); + } + if (opt.handler_bool) { + const int parsed = llama_common_parse_bool_value_rust(value.c_str()); + if (parsed < 0) { + throw std::invalid_argument("invalid boolean value"); + } + opt.handler_bool(params, parsed != 0); + } + if (opt.handler_string) { + opt.handler_string(params, value); + continue; + } + } catch (std::exception & e) { + throw std::invalid_argument(string_format( + "error while handling environment variable \"%s\": %s\n\n", opt.env, e.what())); + } + } + } + + // handle command line arguments + auto check_arg = [&](int i) { + if (i+1 >= argc) { + throw std::invalid_argument("expected value for argument"); + } + }; + + auto parse_cli_args = [&]() { + std::set seen_args; + + for (int i = 1; i < argc; i++) { + const std::string arg_prefix = "--"; + + std::string arg = argv[i]; + if (arg.compare(0, arg_prefix.size(), arg_prefix) == 0) { + std::replace(arg.begin(), arg.end(), '_', '-'); + } + if (arg_to_options.find(arg) == arg_to_options.end()) { + throw std::invalid_argument(string_format("error: invalid argument: %s", arg.c_str())); + } + if (!seen_args.insert(arg).second) { + LOG_WRN("DEPRECATED: argument '%s' specified multiple times, use comma-separated values instead (only last value will be used)\n", arg.c_str()); + } + auto & tmp = arg_to_options[arg]; + auto opt = *tmp.first; + bool is_positive = tmp.second; + if (opt.has_value_from_env()) { + fprintf(stderr, "warn: %s environment variable is set, but will be overwritten by command line argument %s\n", opt.env, arg.c_str()); + } + try { + if (opt.handler_void) { + opt.handler_void(params); + continue; + } + if (opt.handler_bool) { + opt.handler_bool(params, is_positive); + continue; + } + + // arg with single value + check_arg(i); + std::string val = argv[++i]; + if (opt.handler_int) { + opt.handler_int(params, std::stoi(val)); + continue; + } + if (opt.handler_string) { + opt.handler_string(params, val); + continue; + } + + // arg with 2 values + check_arg(i); + std::string val2 = argv[++i]; + if (opt.handler_str_str) { + opt.handler_str_str(params, val, val2); + continue; + } + } catch (std::exception & e) { + throw std::invalid_argument(string_format( + "error while handling argument \"%s\": %s\n\n" + "usage:\n%s\n\nto show complete usage, run with -h", + arg.c_str(), e.what(), opt.to_string().c_str())); + } + } + }; + + // parse the first time to get -hf option (used for remote preset) + parse_cli_args(); + + // TODO: Remove later + try { + hf_cache::migrate_old_cache_to_hf_cache(params.hf_token, params.offline); + } catch (const std::exception & e) { + LOG_WRN("HF cache migration failed: %s\n", e.what()); + } + // export_graph_ops loads only metadata + const bool skip_model_download = ctx_arg.ex == LLAMA_EXAMPLE_EXPORT_GRAPH_OPS; + + // maybe handle remote preset + if (!params.model.hf_repo.empty() && !skip_model_download) { + std::string cli_hf_repo = params.model.hf_repo; + bool has_preset = common_params_handle_remote_preset(params, ctx_arg.ex); + + // special case: if hf_repo explicitly set by preset, we need to preserve it (ignore CLI value) + // this is useful when we have one HF repo pointing to other HF repos (one model - multiple GGUFs) + std::string preset_hf_repo = params.model.hf_repo; + bool preset_has_hf_repo = preset_hf_repo != cli_hf_repo; + + if (has_preset) { + // re-parse CLI args to override preset values + parse_cli_args(); + } + + // preserve hf_repo from preset if needed + if (preset_has_hf_repo) { + params.model.hf_repo = preset_hf_repo; + } + } + + postprocess_cpu_params(params.cpuparams, nullptr); + postprocess_cpu_params(params.cpuparams_batch, ¶ms.cpuparams); + + postprocess_cpu_params(params.speculative.draft.cpuparams, ¶ms.cpuparams); + postprocess_cpu_params(params.speculative.draft.cpuparams_batch, ¶ms.cpuparams_batch); + + if (params.prompt_cache_all && (params.interactive || params.interactive_first)) { + throw std::invalid_argument("error: --prompt-cache-all not supported in interactive mode yet\n"); + } + + // handle model and download + if (!skip_model_download) { + auto res = common_params_handle_model(params.model, params.hf_token, params.offline); + if (params.no_mmproj) { + params.mmproj = {}; + } else if (res.found_mmproj && params.mmproj.path.empty() && params.mmproj.url.empty()) { + // optionally, handle mmproj model when -hf is specified + params.mmproj = res.mmproj; + } + // only download mmproj if the current example is using it + for (const auto & ex : mmproj_examples) { + if (ctx_arg.ex == ex) { + common_params_handle_model(params.mmproj, params.hf_token, params.offline); + break; + } + } + common_params_handle_model(params.speculative.draft.mparams, params.hf_token, params.offline); + common_params_handle_model(params.vocoder.model, params.hf_token, params.offline); + } + + // model is required (except for server) + // TODO @ngxson : maybe show a list of available models in CLI in this case + if (params.model.path.empty() && ctx_arg.ex != LLAMA_EXAMPLE_SERVER && !skip_model_download && !params.usage && !params.completion) { + throw std::invalid_argument("error: --model is required\n"); + } + + if (params.escape) { + string_process_escapes(params.prompt); + string_process_escapes(params.input_prefix); + string_process_escapes(params.input_suffix); + for (auto & antiprompt : params.antiprompt) { + string_process_escapes(antiprompt); + } + for (auto & seq_breaker : params.sampling.dry_sequence_breakers) { + string_process_escapes(seq_breaker); + } + for (auto & pair : params.speculative.draft.replacements) { + string_process_escapes(pair.first); + string_process_escapes(pair.second); + } + } + + if (!params.kv_overrides.empty()) { + params.kv_overrides.emplace_back(); + params.kv_overrides.back().key[0] = 0; + } + + // pad tensor_buft_overrides for llama_params_fit: + const size_t ntbo = llama_max_tensor_buft_overrides(); + while (params.tensor_buft_overrides.size() < ntbo) { + params.tensor_buft_overrides.push_back({nullptr, nullptr}); + } + + if (!params.speculative.draft.tensor_buft_overrides.empty()) { + params.speculative.draft.tensor_buft_overrides.push_back({nullptr, nullptr}); + } + + if (!params.chat_template.empty() && !common_chat_verify_template(params.chat_template, params.use_jinja)) { + throw std::runtime_error(string_format( + "error: the supplied chat template is not supported: %s%s\n", + params.chat_template.c_str(), + params.use_jinja ? "" : "\nnote: llama.cpp was started without --jinja, we only support commonly used templates" + )); + } + + return true; +} + +static void common_params_print_usage(common_params_context & ctx_arg) { + auto print_options = [](std::vector & options) { + for (common_arg * opt : options) { + printf("%s", opt->to_string().c_str()); + } + }; + + std::vector common_options; + std::vector sampling_options; + std::vector spec_options; + std::vector specific_options; + for (auto & opt : ctx_arg.options) { + // in case multiple LLAMA_EXAMPLE_* are set, we prioritize the LLAMA_EXAMPLE_* matching current example + if (opt.is_sampling) { + sampling_options.push_back(&opt); + } else if (opt.is_spec) { + spec_options.push_back(&opt); + } else if (opt.in_example(ctx_arg.ex)) { + specific_options.push_back(&opt); + } else { + common_options.push_back(&opt); + } + } + printf("----- common params -----\n\n"); + print_options(common_options); + printf("\n\n----- sampling params -----\n\n"); + print_options(sampling_options); + printf("\n\n----- speculative params -----\n\n"); + print_options(spec_options); + // TODO: maybe convert enum llama_example to string + printf("\n\n----- example-specific params -----\n\n"); + print_options(specific_options); +} + +static void common_params_print_completion(common_params_context & ctx_arg) { + std::vector common_options; + std::vector sampling_options; + std::vector spec_options; + std::vector specific_options; + + for (auto & opt : ctx_arg.options) { + if (opt.is_sampling) { + sampling_options.push_back(&opt); + } else if (opt.is_spec) { + spec_options.push_back(&opt); + } else if (opt.in_example(ctx_arg.ex)) { + specific_options.push_back(&opt); + } else { + common_options.push_back(&opt); + } + } + + printf("_llama_completions() {\n"); + printf(" local cur prev opts\n"); + printf(" COMPREPLY=()\n"); + printf(" cur=\"${COMP_WORDS[COMP_CWORD]}\"\n"); + printf(" prev=\"${COMP_WORDS[COMP_CWORD-1]}\"\n\n"); + + printf(" opts=\""); + auto print_options = [](const std::vector & options) { + for (const common_arg * opt : options) { + for (const char * arg : opt->args) { + printf("%s ", arg); + } + } + }; + + print_options(common_options); + print_options(sampling_options); + print_options(spec_options); + print_options(specific_options); + printf("\"\n\n"); + + printf(" case \"$prev\" in\n"); + printf(" --model|-m)\n"); + printf(" COMPREPLY=( $(compgen -f -X '!*.gguf' -- \"$cur\") $(compgen -d -- \"$cur\") )\n"); + printf(" return 0\n"); + printf(" ;;\n"); + printf(" --grammar-file)\n"); + printf(" COMPREPLY=( $(compgen -f -X '!*.gbnf' -- \"$cur\") $(compgen -d -- \"$cur\") )\n"); + printf(" return 0\n"); + printf(" ;;\n"); + printf(" --chat-template-file)\n"); + printf(" COMPREPLY=( $(compgen -f -X '!*.jinja' -- \"$cur\") $(compgen -d -- \"$cur\") )\n"); + printf(" return 0\n"); + printf(" ;;\n"); + printf(" *)\n"); + printf(" COMPREPLY=( $(compgen -W \"${opts}\" -- \"$cur\") )\n"); + printf(" return 0\n"); + printf(" ;;\n"); + printf(" esac\n"); + printf("}\n\n"); + + std::set executables = { + "llama-batched", + "llama-batched-bench", + "llama-bench", + "llama-cli", + "llama-completion", + "llama-convert-llama2c-to-ggml", + "llama-cvector-generator", + "llama-debug", + "llama-diffusion-cli", + "llama-embedding", + "llama-eval-callback", + "llama-export-lora", + "llama-finetune", + "llama-fit-params", + "llama-gemma3-cli", + "llama-gen-docs", + "llama-gguf", + "llama-gguf-hash", + "llama-gguf-split", + "llama-idle", + "llama-imatrix", + "llama-llava-cli", + "llama-lookahead", + "llama-lookup", + "llama-lookup-create", + "llama-lookup-merge", + "llama-lookup-stats", + "llama-minicpmv-cli", + "llama-mtmd-cli", + "llama-parallel", + "llama-passkey", + "llama-perplexity", + "llama-q8dot", + "llama-quantize", + "llama-qwen2vl-cli", + "llama-retrieval", + "llama-save-load-state", + "llama-server", + "llama-simple", + "llama-simple-chat", + "llama-speculative", + "llama-speculative-simple", + "llama-tokenize", + "llama-tts", + "llama-vdot" + }; + + for (const auto& exe : executables) { + printf("complete -F _llama_completions %s\n", exe.c_str()); + } +} + +static std::vector parse_device_list(const std::string & value) { + std::vector devices; + auto dev_names = string_split(value, ','); + if (dev_names.empty()) { + throw std::invalid_argument("no devices specified"); + } + if (dev_names.size() == 1 && dev_names[0] == "none") { + devices.push_back(nullptr); + } else { + for (const auto & device : dev_names) { + auto * dev = ggml_backend_dev_by_name(device.c_str()); + if (!dev || ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) { + throw std::invalid_argument(string_format("invalid device: %s", device.c_str())); + } + devices.push_back(dev); + } + devices.push_back(nullptr); + } + return devices; +} + +static void add_rpc_devices(const std::string & servers) { + auto rpc_servers = string_split(servers, ','); + if (rpc_servers.empty()) { + throw std::invalid_argument("no RPC servers specified"); + } + ggml_backend_reg_t rpc_reg = ggml_backend_reg_by_name("RPC"); + if (!rpc_reg) { + throw std::invalid_argument("failed to find RPC backend"); + } + typedef ggml_backend_reg_t (*ggml_backend_rpc_add_server_t)(const char * endpoint); + ggml_backend_rpc_add_server_t ggml_backend_rpc_add_server_fn = (ggml_backend_rpc_add_server_t) ggml_backend_reg_get_proc_address(rpc_reg, "ggml_backend_rpc_add_server"); + if (!ggml_backend_rpc_add_server_fn) { + throw std::invalid_argument("failed to find RPC add server function"); + } + for (const auto & server : rpc_servers) { + auto reg = ggml_backend_rpc_add_server_fn(server.c_str()); + ggml_backend_register(reg); + } +} + +bool common_params_to_map(int argc, char ** argv, llama_example ex, std::map & out_map) { + common_params dummy_params; + common_params_context ctx_arg = common_params_parser_init(dummy_params, ex, nullptr); + + std::unordered_map arg_to_options; + for (auto & opt : ctx_arg.options) { + for (const auto & arg : opt.args) { + arg_to_options[arg] = &opt; + } + for (const auto & arg : opt.args_neg) { + arg_to_options[arg] = &opt; + } + } + + // TODO @ngxson : find a way to deduplicate this code + + // handle command line arguments + auto check_arg = [&](int i) { + if (i+1 >= argc) { + throw std::invalid_argument("expected value for argument"); + } + }; + + std::set seen_args; + + for (int i = 1; i < argc; i++) { + const std::string arg_prefix = "--"; + + std::string arg = argv[i]; + if (arg.compare(0, arg_prefix.size(), arg_prefix) == 0) { + std::replace(arg.begin(), arg.end(), '_', '-'); + } + if (arg_to_options.find(arg) == arg_to_options.end()) { + throw std::invalid_argument(string_format("error: invalid argument: %s", arg.c_str())); + } + if (!seen_args.insert(arg).second) { + LOG_WRN("DEPRECATED: argument '%s' specified multiple times, use comma-separated values instead (only last value will be used)\n", arg.c_str()); + } + auto opt = *arg_to_options[arg]; + std::string val; + if (opt.value_hint == nullptr && opt.value_hint_2 == nullptr) { + // bool arg (need to reverse the meaning for negative args) + bool is_neg = std::find(opt.args_neg.begin(), opt.args_neg.end(), arg) != opt.args_neg.end(); + val = is_neg ? "0" : "1"; + } + if (opt.value_hint != nullptr) { + // arg with single value + check_arg(i); + val = argv[++i]; + } + if (opt.value_hint_2 != nullptr) { + // TODO: support arg with 2 values + throw std::invalid_argument("error: argument with 2 values is not yet supported\n"); + } + out_map[opt] = val; + } + + return true; +} + +bool common_params_parse(int argc, char ** argv, common_params & params, llama_example ex, void(*print_usage)(int, char **)) { + auto ctx_arg = common_params_parser_init(params, ex, print_usage); + const common_params params_org = ctx_arg.params; // the example can modify the default params + + try { + if (!common_params_parse_ex(argc, argv, ctx_arg)) { + ctx_arg.params = params_org; + return false; + } + if (ctx_arg.params.usage) { + common_params_print_usage(ctx_arg); + if (ctx_arg.print_usage) { + ctx_arg.print_usage(argc, argv); + } + exit(0); + } + if (ctx_arg.params.completion) { + common_params_print_completion(ctx_arg); + exit(0); + } + params.lr.init(); + } catch (const std::invalid_argument & ex) { + fprintf(stderr, "%s\n", ex.what()); + ctx_arg.params = params_org; + return false; + } catch (std::exception & ex) { + fprintf(stderr, "%s\n", ex.what()); + exit(1); // for other exceptions, we exit with status code 1 + } + + return true; +} + +static std::string list_builtin_chat_templates() { + std::vector supported_tmpl; + int32_t res = llama_chat_builtin_templates(nullptr, 0); + supported_tmpl.resize(res); + res = llama_chat_builtin_templates(supported_tmpl.data(), supported_tmpl.size()); + std::ostringstream msg; + for (auto & tmpl : supported_tmpl) { + msg << tmpl << (&tmpl == &supported_tmpl.back() ? "" : ", "); + } + return msg.str(); +} + +bool common_arg_utils::is_truthy(const std::string & value) { + return llama_common_arg_is_truthy_rust(value.c_str()) != 0; +} + +bool common_arg_utils::is_falsey(const std::string & value) { + return llama_common_arg_is_falsey_rust(value.c_str()) != 0; +} + +bool common_arg_utils::is_autoy(const std::string & value) { + return llama_common_arg_is_autoy_rust(value.c_str()) != 0; +} + +// Simple CSV parser that handles quoted fields and escaped quotes +// example: +// input: value1,"value, with, commas","value with ""escaped"" quotes",value4 +// output: [value1] [value, with, commas] [value with "escaped" quotes] [value4] +static std::vector parse_csv_row(const std::string& input) { + common_rust_string_list parsed = llama_common_parse_csv_row_rust((const uint8_t *) input.data(), input.size()); + std::vector fields; + if (!parsed.data) { + return fields; + } + fields.reserve(parsed.len); + for (size_t i = 0; i < parsed.len; ++i) { + const llama_common_unicode_string & item = parsed.data[i]; + fields.emplace_back((const char *) item.data, item.len); + } + llama_common_string_list_free(parsed); + return fields; +} + +common_params_context common_params_parser_init(common_params & params, llama_example ex, void(*print_usage)(int, char **)) { + // per-example default params + // we define here to make sure it's included in llama-gen-docs + if (ex == LLAMA_EXAMPLE_COMPLETION) { + params.use_jinja = false; // disable jinja by default + + } else if (ex == LLAMA_EXAMPLE_MTMD) { + params.use_jinja = false; // disable jinja by default + params.sampling.temp = 0.2; // lower temp by default for better quality + + } else if (ex == LLAMA_EXAMPLE_SERVER) { + params.n_parallel = -1; // auto by default + } + + params.use_color = tty_can_use_colors(); + + // load dynamic backends + ggml_backend_load_all(); + + common_params_context ctx_arg(params); + ctx_arg.print_usage = print_usage; + ctx_arg.ex = ex; + + std::string sampler_type_chars; + std::string sampler_type_names; + for (const auto & sampler : params.sampling.samplers) { + sampler_type_chars += common_sampler_type_to_chr(sampler); + sampler_type_names += common_sampler_type_to_str(sampler) + ";"; + } + if (!sampler_type_names.empty()) { + sampler_type_names.pop_back(); // remove last semicolon + } + + + /** + * filter options by example + * rules: + * - all examples inherit options from LLAMA_EXAMPLE_COMMON + * - if LLAMA_EXAMPLE_* is set (other than COMMON), we only show the option in the corresponding example + * - if both {LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_*,} are set, we will prioritize the LLAMA_EXAMPLE_* matching current example + */ + auto add_opt = [&](common_arg arg) { + if ((arg.in_example(ex) || arg.in_example(LLAMA_EXAMPLE_COMMON)) && !arg.is_exclude(ex)) { + ctx_arg.options.push_back(std::move(arg)); + } + }; + + + add_opt(common_arg( + {"-h", "--help", "--usage"}, + "print usage and exit", + [](common_params & params) { + params.usage = true; + } + )); + add_opt(common_arg( + {"--version"}, + "show version and build info", + [](common_params &) { + fprintf(stderr, "version: %d (%s)\n", llama_build_number(), llama_commit()); + fprintf(stderr, "built with %s for %s\n", llama_compiler(), llama_build_target()); + exit(0); + } + )); + add_opt(common_arg( + {"--license"}, + "show source code license and dependencies", + [](common_params &) { + for (int i = 0; LICENSES[i]; ++i) { + printf("%s\n", LICENSES[i]); + } + exit(0); + } + )); + add_opt(common_arg( + {"-cl", "--cache-list"}, + "show list of models in cache", + [](common_params &) { + auto models = common_list_cached_models(); + printf("number of models in cache: %zu\n", models.size()); + for (size_t i = 0; i < models.size(); i++) { + printf("%4zu. %s\n", i + 1, models[i].to_string().c_str()); + } + exit(0); + } + )); + add_opt(common_arg( + {"--completion-bash"}, + "print source-able bash completion script for llama.cpp", + [](common_params & params) { + params.completion = true; + } + )); + add_opt(common_arg( + {"--verbose-prompt"}, + string_format("print a verbose prompt before generation (default: %s)", params.verbose_prompt ? "true" : "false"), + [](common_params & params) { + params.verbose_prompt = true; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_EMBEDDING, LLAMA_EXAMPLE_RETRIEVAL})); + add_opt(common_arg( + {"--display-prompt"}, + {"--no-display-prompt"}, + string_format("whether to print prompt at generation (default: %s)", params.display_prompt ? "true" : "false"), + [](common_params & params, bool value) { + params.display_prompt = value; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"-co", "--color"}, "[on|off|auto]", + "Colorize output to distinguish prompt and user input from generations ('on', 'off', or 'auto', default: 'auto')\n" + "'auto' enables colors when output is to a terminal", + [](common_params & params, const std::string & value) { + if (is_truthy(value)) { + params.use_color = true; + } else if (is_falsey(value)) { + params.use_color = false; + } else if (is_autoy(value)) { + params.use_color = tty_can_use_colors(); + } else { + throw std::invalid_argument( + string_format("error: unknown value for --color: '%s'\n", value.c_str())); + } + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP})); + add_opt(common_arg( + {"-t", "--threads"}, "N", + string_format("number of CPU threads to use during generation (default: %d)", params.cpuparams.n_threads), + [](common_params & params, int value) { + params.cpuparams.n_threads = value; + if (params.cpuparams.n_threads <= 0) { + params.cpuparams.n_threads = std::thread::hardware_concurrency(); + } + } + ).set_env("LLAMA_ARG_THREADS")); + add_opt(common_arg( + {"-tb", "--threads-batch"}, "N", + "number of threads to use during batch and prompt processing (default: same as --threads)", + [](common_params & params, int value) { + params.cpuparams_batch.n_threads = value; + if (params.cpuparams_batch.n_threads <= 0) { + params.cpuparams_batch.n_threads = std::thread::hardware_concurrency(); + } + } + )); + add_opt(common_arg( + {"-C", "--cpu-mask"}, "M", + "CPU affinity mask: arbitrarily long hex. Complements cpu-range (default: \"\")", + [](common_params & params, const std::string & mask) { + params.cpuparams.mask_valid = true; + if (!parse_cpu_mask(mask, params.cpuparams.cpumask)) { + throw std::invalid_argument("invalid cpumask"); + } + } + )); + add_opt(common_arg( + {"-Cr", "--cpu-range"}, "lo-hi", + "range of CPUs for affinity. Complements --cpu-mask", + [](common_params & params, const std::string & range) { + params.cpuparams.mask_valid = true; + if (!parse_cpu_range(range, params.cpuparams.cpumask)) { + throw std::invalid_argument("invalid range"); + } + } + )); + add_opt(common_arg( + {"--cpu-strict"}, "<0|1>", + string_format("use strict CPU placement (default: %u)\n", (unsigned) params.cpuparams.strict_cpu), + [](common_params & params, const std::string & value) { + params.cpuparams.strict_cpu = std::stoul(value); + } + )); + add_opt(common_arg( + {"--prio"}, "N", + string_format("set process/thread priority : low(-1), normal(0), medium(1), high(2), realtime(3) (default: %d)\n", params.cpuparams.priority), + [](common_params & params, int prio) { + if (prio < GGML_SCHED_PRIO_LOW || prio > GGML_SCHED_PRIO_REALTIME) { + throw std::invalid_argument("invalid value"); + } + params.cpuparams.priority = (enum ggml_sched_priority) prio; + } + )); + add_opt(common_arg( + {"--poll"}, "<0...100>", + string_format("use polling level to wait for work (0 - no polling, default: %u)\n", (unsigned) params.cpuparams.poll), + [](common_params & params, const std::string & value) { + params.cpuparams.poll = std::stoul(value); + } + )); + add_opt(common_arg( + {"-Cb", "--cpu-mask-batch"}, "M", + "CPU affinity mask: arbitrarily long hex. Complements cpu-range-batch (default: same as --cpu-mask)", + [](common_params & params, const std::string & mask) { + params.cpuparams_batch.mask_valid = true; + if (!parse_cpu_mask(mask, params.cpuparams_batch.cpumask)) { + throw std::invalid_argument("invalid cpumask"); + } + } + )); + add_opt(common_arg( + {"-Crb", "--cpu-range-batch"}, "lo-hi", + "ranges of CPUs for affinity. Complements --cpu-mask-batch", + [](common_params & params, const std::string & range) { + params.cpuparams_batch.mask_valid = true; + if (!parse_cpu_range(range, params.cpuparams_batch.cpumask)) { + throw std::invalid_argument("invalid range"); + } + } + )); + add_opt(common_arg( + {"--cpu-strict-batch"}, "<0|1>", + "use strict CPU placement (default: same as --cpu-strict)", + [](common_params & params, int value) { + params.cpuparams_batch.strict_cpu = value; + } + )); + add_opt(common_arg( + {"--prio-batch"}, "N", + string_format("set process/thread priority : 0-normal, 1-medium, 2-high, 3-realtime (default: %d)\n", params.cpuparams_batch.priority), + [](common_params & params, int prio) { + if (prio < 0 || prio > 3) { + throw std::invalid_argument("invalid value"); + } + params.cpuparams_batch.priority = (enum ggml_sched_priority) prio; + } + )); + add_opt(common_arg( + {"--poll-batch"}, "<0|1>", + "use polling to wait for work (default: same as --poll)", + [](common_params & params, int value) { + params.cpuparams_batch.poll = value; + } + )); + add_opt(common_arg( + {"-lcs", "--lookup-cache-static"}, "FNAME", + "path to static lookup cache to use for lookup decoding (not updated by generation)", + [](common_params & params, const std::string & value) { + params.speculative.ngram_cache.lookup_cache_static = value; + } + ).set_examples({LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"-lcd", "--lookup-cache-dynamic"}, "FNAME", + "path to dynamic lookup cache to use for lookup decoding (updated by generation)", + [](common_params & params, const std::string & value) { + params.speculative.ngram_cache.lookup_cache_dynamic = value; + } + ).set_examples({LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"-c", "--ctx-size"}, "N", + string_format("size of the prompt context (default: %d, 0 = loaded from model)", params.n_ctx), + [](common_params & params, int value) { + params.n_ctx = value; + if (value == 0) { + // disable context reduction in llama_params_fit if the user explicitly requests the full context size: + params.fit_params_min_ctx = UINT32_MAX; + } + } + ).set_env("LLAMA_ARG_CTX_SIZE")); + add_opt(common_arg( + {"-n", "--predict", "--n-predict"}, "N", + string_format( + ex == LLAMA_EXAMPLE_COMPLETION + ? "number of tokens to predict (default: %d, -1 = infinity, -2 = until context filled)" + : "number of tokens to predict (default: %d, -1 = infinity)", + params.n_predict), + [](common_params & params, int value) { + params.n_predict = value; + } + ).set_env("LLAMA_ARG_N_PREDICT")); + add_opt(common_arg( + {"-b", "--batch-size"}, "N", + string_format("logical maximum batch size (default: %d)", params.n_batch), + [](common_params & params, int value) { + params.n_batch = value; + } + ).set_env("LLAMA_ARG_BATCH")); + add_opt(common_arg( + {"-ub", "--ubatch-size"}, "N", + string_format("physical maximum batch size (default: %d)", params.n_ubatch), + [](common_params & params, int value) { + params.n_ubatch = value; + } + ).set_env("LLAMA_ARG_UBATCH")); + add_opt(common_arg( + {"--keep"}, "N", + string_format("number of tokens to keep from the initial prompt (default: %d, -1 = all)", params.n_keep), + [](common_params & params, int value) { + params.n_keep = value; + } + )); + add_opt(common_arg( + {"--swa-full"}, + string_format("use full-size SWA cache (default: %s)\n" + "[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)", params.swa_full ? "true" : "false"), + [](common_params & params) { + params.swa_full = true; + } + ).set_env("LLAMA_ARG_SWA_FULL")); + add_opt(common_arg( + {"-ctxcp", "--ctx-checkpoints", "--swa-checkpoints"}, "N", + string_format("max number of context checkpoints to create per slot (default: %d)" + "[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)", params.n_ctx_checkpoints), + [](common_params & params, int value) { + params.n_ctx_checkpoints = value; + } + ).set_env("LLAMA_ARG_CTX_CHECKPOINTS").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"-cpent", "--checkpoint-every-n-tokens"}, "N", + string_format("create a checkpoint every n tokens during prefill (processing), -1 to disable (default: %d)", params.checkpoint_every_nt), + [](common_params & params, int value) { + params.checkpoint_every_nt = value; + } + ).set_env("LLAMA_ARG_CHECKPOINT_EVERY_NT").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"-cram", "--cache-ram"}, "N", + string_format("set the maximum cache size in MiB (default: %d, -1 - no limit, 0 - disable)" + "[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)", params.cache_ram_mib), + [](common_params & params, int value) { + params.cache_ram_mib = value; + } + ).set_env("LLAMA_ARG_CACHE_RAM").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"-kvu", "--kv-unified"}, + {"-no-kvu", "--no-kv-unified"}, + "use single unified KV buffer shared across all sequences (default: enabled if number of slots is auto)", + [](common_params & params, bool value) { + params.kv_unified = value; + } + ).set_env("LLAMA_ARG_KV_UNIFIED").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_BATCHED, LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL})); + add_opt(common_arg( + {"--cache-idle-slots"}, + {"--no-cache-idle-slots"}, + "save and clear idle slots on new task (default: enabled, requires unified KV and cache-ram)", + [](common_params & params, bool value) { + params.cache_idle_slots = value; + } + ).set_env("LLAMA_ARG_CACHE_IDLE_SLOTS").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--context-shift"}, + {"--no-context-shift"}, + string_format("whether to use context shift on infinite text generation (default: %s)", params.ctx_shift ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.ctx_shift = value; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_PERPLEXITY}).set_env("LLAMA_ARG_CONTEXT_SHIFT")); + add_opt(common_arg( + {"--chunks"}, "N", + string_format("max number of chunks to process (default: %d, -1 = all)", params.n_chunks), + [](common_params & params, int value) { + params.n_chunks = value; + } + ).set_examples({LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_RETRIEVAL})); + add_opt(common_arg({ "-fa", "--flash-attn" }, "[on|off|auto]", + string_format("set Flash Attention use ('on', 'off', or 'auto', default: '%s')", + llama_flash_attn_type_name(params.flash_attn_type)), + [](common_params & params, const std::string & value) { + if (is_truthy(value)) { + params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; + } else if (is_falsey(value)) { + params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + } else if (is_autoy(value)) { + params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO; + } else { + throw std::runtime_error( + string_format("error: unknown value for --flash-attn: '%s'\n", value.c_str())); + } + }).set_env("LLAMA_ARG_FLASH_ATTN")); + add_opt(common_arg( + {"-p", "--prompt"}, "PROMPT", + "prompt to start generation with; for system message, use -sys", + [](common_params & params, const std::string & value) { + params.prompt = value; + } + ).set_excludes({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"-sys", "--system-prompt"}, "PROMPT", + "system prompt to use with model (if applicable, depending on chat template)", + [](common_params & params, const std::string & value) { + params.system_prompt = value; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_DIFFUSION, LLAMA_EXAMPLE_MTMD})); + add_opt(common_arg( + {"--perf"}, + {"--no-perf"}, + string_format("whether to enable internal libllama performance timings (default: %s)", params.no_perf ? "true" : "false"), + [](common_params & params, bool value) { + params.no_perf = !value; + params.sampling.no_perf = !value; + } + ).set_env("LLAMA_ARG_PERF")); + add_opt(common_arg( + {"--show-timings"}, + {"--no-show-timings"}, + string_format("whether to show timing information after each response (default: %s)", params.show_timings ? "true" : "false"), + [](common_params & params, bool value) { + params.show_timings = value; + } + ).set_examples({LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SHOW_TIMINGS")); + add_opt(common_arg( + {"-f", "--file"}, "FNAME", + "a file containing the prompt (default: none)", + [](common_params & params, const std::string & value) { + params.prompt = read_file(value); + // store the external file name in params + params.prompt_file = value; + if (!params.prompt.empty() && params.prompt.back() == '\n') { + params.prompt.pop_back(); + } + } + ).set_excludes({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"-sysf", "--system-prompt-file"}, "FNAME", + "a file containing the system prompt (default: none)", + [](common_params & params, const std::string & value) { + params.system_prompt = read_file(value); + if (!params.system_prompt.empty() && params.system_prompt.back() == '\n') { + params.system_prompt.pop_back(); + } + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_DIFFUSION})); + add_opt(common_arg( + {"--in-file"}, "FNAME", + "an input file (use comma-separated values to specify multiple files)", + [](common_params & params, const std::string & value) { + for (const auto & item : parse_csv_row(value)) { + std::ifstream file(item); + if (!file) { + throw std::runtime_error(string_format("error: failed to open file '%s'\n", item.c_str())); + } + params.in_files.push_back(item); + } + } + ).set_examples({LLAMA_EXAMPLE_IMATRIX})); + add_opt(common_arg( + {"-bf", "--binary-file"}, "FNAME", + "binary file containing the prompt (default: none)", + [](common_params & params, const std::string & value) { + std::ifstream file(value, std::ios::binary); + if (!file) { + throw std::runtime_error(string_format("error: failed to open file '%s'\n", value.c_str())); + } + // store the external file name in params + params.prompt_file = value; + std::ostringstream ss; + ss << file.rdbuf(); + params.prompt = ss.str(); + fprintf(stderr, "Read %zu bytes from binary file %s\n", params.prompt.size(), value.c_str()); + } + ).set_excludes({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"-e", "--escape"}, + {"--no-escape"}, + string_format("whether to process escapes sequences (\\n, \\r, \\t, \\', \\\", \\\\) (default: %s)", params.escape ? "true" : "false"), + [](common_params & params, bool value) { + params.escape = value; + } + )); + add_opt(common_arg( + {"-ptc", "--print-token-count"}, "N", + string_format("print token count every N tokens (default: %d)", params.n_print), + [](common_params & params, int value) { + params.n_print = value; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION})); + add_opt(common_arg( + {"--prompt-cache"}, "FNAME", + "file to cache prompt state for faster startup (default: none)", + [](common_params & params, const std::string & value) { + params.path_prompt_cache = value; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION})); + add_opt(common_arg( + {"--prompt-cache-all"}, + "if specified, saves user input and generations to cache as well\n", + [](common_params & params) { + params.prompt_cache_all = true; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION})); + add_opt(common_arg( + {"--prompt-cache-ro"}, + "if specified, uses the prompt cache but does not update it", + [](common_params & params) { + params.prompt_cache_ro = true; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION})); + add_opt(common_arg( + {"-r", "--reverse-prompt"}, "PROMPT", + "halt generation at PROMPT, return control in interactive mode\n", + [](common_params & params, const std::string & value) { + params.antiprompt.emplace_back(value); + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"-sp", "--special"}, + string_format("special tokens output enabled (default: %s)", params.special ? "true" : "false"), + [](common_params & params) { + params.special = true; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"-cnv", "--conversation"}, + {"-no-cnv", "--no-conversation"}, + "whether to run in conversation mode:\n" + "- does not print special tokens and suffix/prefix\n" + "- interactive mode is also enabled\n" + "(default: auto enabled if chat template is available)", + [](common_params & params, bool value) { + params.conversation_mode = value ? COMMON_CONVERSATION_MODE_ENABLED : COMMON_CONVERSATION_MODE_DISABLED; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"-st", "--single-turn"}, + "run conversation for a single turn only, then exit when done\n" + "will not be interactive if first turn is predefined with --prompt\n" + "(default: false)", + [](common_params & params) { + params.single_turn = true; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"-i", "--interactive"}, + string_format("run in interactive mode (default: %s)", params.interactive ? "true" : "false"), + [](common_params & params) { + params.interactive = true; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION})); + add_opt(common_arg( + {"-if", "--interactive-first"}, + string_format("run in interactive mode and wait for input right away (default: %s)", params.interactive_first ? "true" : "false"), + [](common_params & params) { + params.interactive_first = true; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION})); + add_opt(common_arg( + {"-mli", "--multiline-input"}, + "allows you to write or paste multiple lines without ending each in '\\'", + [](common_params & params) { + params.multiline_input = true; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--in-prefix-bos"}, + "prefix BOS to user inputs, preceding the `--in-prefix` string", + [](common_params & params) { + params.input_prefix_bos = true; + params.enable_chat_template = false; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION})); + add_opt(common_arg( + {"--in-prefix"}, "STRING", + "string to prefix user inputs with (default: empty)", + [](common_params & params, const std::string & value) { + params.input_prefix = value; + params.enable_chat_template = false; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION})); + add_opt(common_arg( + {"--in-suffix"}, "STRING", + "string to suffix after user inputs with (default: empty)", + [](common_params & params, const std::string & value) { + params.input_suffix = value; + params.enable_chat_template = false; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION})); + add_opt(common_arg( + {"--warmup"}, + {"--no-warmup"}, + string_format("whether to perform warmup with an empty run (default: %s)", params.warmup ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.warmup = value; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_MTMD, LLAMA_EXAMPLE_EMBEDDING, LLAMA_EXAMPLE_RETRIEVAL, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_DEBUG})); + add_opt(common_arg( + {"--spm-infill"}, + string_format( + "use Suffix/Prefix/Middle pattern for infill (instead of Prefix/Suffix/Middle) as some models prefer this. (default: %s)", + params.spm_infill ? "enabled" : "disabled" + ), + [](common_params & params) { + params.spm_infill = true; + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--samplers"}, "SAMPLERS", + string_format("samplers that will be used for generation in the order, separated by \';\'\n(default: %s)", sampler_type_names.c_str()), + [](common_params & params, const std::string & value) { + const auto sampler_names = string_split(value, ';'); + params.sampling.samplers = common_sampler_types_from_names(sampler_names, true); + params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_SAMPLERS; + } + ).set_sampling()); + add_opt(common_arg( + {"-s", "--seed"}, "SEED", + string_format("RNG seed (default: %d, use random seed for %d)", params.sampling.seed, LLAMA_DEFAULT_SEED), + [](common_params & params, const std::string & value) { + params.sampling.seed = std::stoul(value); + } + ).set_sampling()); + add_opt(common_arg( + {"--sampler-seq", "--sampling-seq"}, "SEQUENCE", + string_format("simplified sequence for samplers that will be used (default: %s)", sampler_type_chars.c_str()), + [](common_params & params, const std::string & value) { + params.sampling.samplers = common_sampler_types_from_chars(value); + } + ).set_sampling()); + add_opt(common_arg( + {"--ignore-eos"}, + "ignore end of stream token and continue generating (implies --logit-bias EOS-inf)", + [](common_params & params) { + params.sampling.ignore_eos = true; + } + ).set_sampling()); + add_opt(common_arg( + {"--temp", "--temperature"}, "N", + string_format("temperature (default: %.2f)", (double)params.sampling.temp), + [](common_params & params, const std::string & value) { + params.sampling.temp = std::stof(value); + params.sampling.temp = std::max(params.sampling.temp, 0.0f); + params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TEMP; + } + ).set_sampling()); + add_opt(common_arg( + {"--top-k"}, "N", + string_format("top-k sampling (default: %d, 0 = disabled)", params.sampling.top_k), + [](common_params & params, int value) { + params.sampling.top_k = value; + params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TOP_K; + } + ).set_sampling().set_env("LLAMA_ARG_TOP_K")); + add_opt(common_arg( + {"--top-p"}, "N", + string_format("top-p sampling (default: %.2f, 1.0 = disabled)", (double)params.sampling.top_p), + [](common_params & params, const std::string & value) { + params.sampling.top_p = std::stof(value); + params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_TOP_P; + } + ).set_sampling()); + add_opt(common_arg( + {"--min-p"}, "N", + string_format("min-p sampling (default: %.2f, 0.0 = disabled)", (double)params.sampling.min_p), + [](common_params & params, const std::string & value) { + params.sampling.min_p = std::stof(value); + params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIN_P; + } + ).set_sampling()); + add_opt(common_arg( + {"--top-nsigma", "--top-n-sigma"}, "N", + string_format("top-n-sigma sampling (default: %.2f, -1.0 = disabled)", params.sampling.top_n_sigma), + [](common_params & params, const std::string & value) { + params.sampling.top_n_sigma = std::stof(value); + } + ).set_sampling()); + add_opt(common_arg( + {"--xtc-probability"}, "N", + string_format("xtc probability (default: %.2f, 0.0 = disabled)", (double)params.sampling.xtc_probability), + [](common_params & params, const std::string & value) { + params.sampling.xtc_probability = std::stof(value); + params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_XTC_PROBABILITY; + } + ).set_sampling()); + add_opt(common_arg( + {"--xtc-threshold"}, "N", + string_format("xtc threshold (default: %.2f, 1.0 = disabled)", (double)params.sampling.xtc_threshold), + [](common_params & params, const std::string & value) { + params.sampling.xtc_threshold = std::stof(value); + params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_XTC_THRESHOLD; + } + ).set_sampling()); + add_opt(common_arg( + {"--typical", "--typical-p"}, "N", + string_format("locally typical sampling, parameter p (default: %.2f, 1.0 = disabled)", (double)params.sampling.typ_p), + [](common_params & params, const std::string & value) { + params.sampling.typ_p = std::stof(value); + } + ).set_sampling()); + add_opt(common_arg( + {"--repeat-last-n"}, "N", + string_format("last n tokens to consider for penalize (default: %d, 0 = disabled, -1 = ctx_size)", params.sampling.penalty_last_n), + [](common_params & params, int value) { + if (value < -1) { + throw std::runtime_error(string_format("error: invalid repeat-last-n = %d\n", value)); + } + params.sampling.penalty_last_n = value; + params.sampling.n_prev = std::max(params.sampling.n_prev, params.sampling.penalty_last_n); + params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_LAST_N; + } + ).set_sampling()); + add_opt(common_arg( + {"--repeat-penalty"}, "N", + string_format("penalize repeat sequence of tokens (default: %.2f, 1.0 = disabled)", (double)params.sampling.penalty_repeat), + [](common_params & params, const std::string & value) { + params.sampling.penalty_repeat = std::stof(value); + params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_REPEAT; + } + ).set_sampling()); + add_opt(common_arg( + {"--presence-penalty"}, "N", + string_format("repeat alpha presence penalty (default: %.2f, 0.0 = disabled)", (double)params.sampling.penalty_present), + [](common_params & params, const std::string & value) { + params.sampling.penalty_present = std::stof(value); + } + ).set_sampling()); + add_opt(common_arg( + {"--frequency-penalty"}, "N", + string_format("repeat alpha frequency penalty (default: %.2f, 0.0 = disabled)", (double)params.sampling.penalty_freq), + [](common_params & params, const std::string & value) { + params.sampling.penalty_freq = std::stof(value); + } + ).set_sampling()); + add_opt(common_arg( + {"--dry-multiplier"}, "N", + string_format("set DRY sampling multiplier (default: %.2f, 0.0 = disabled)", (double)params.sampling.dry_multiplier), + [](common_params & params, const std::string & value) { + params.sampling.dry_multiplier = std::stof(value); + } + ).set_sampling()); + add_opt(common_arg( + {"--dry-base"}, "N", + string_format("set DRY sampling base value (default: %.2f)", (double)params.sampling.dry_base), + [](common_params & params, const std::string & value) { + float potential_base = std::stof(value); + if (potential_base >= 1.0f) + { + params.sampling.dry_base = potential_base; + } + } + ).set_sampling()); + add_opt(common_arg( + {"--dry-allowed-length"}, "N", + string_format("set allowed length for DRY sampling (default: %d)", params.sampling.dry_allowed_length), + [](common_params & params, int value) { + params.sampling.dry_allowed_length = value; + } + ).set_sampling()); + add_opt(common_arg( + {"--dry-penalty-last-n"}, "N", + string_format("set DRY penalty for the last n tokens (default: %d, 0 = disable, -1 = context size)", params.sampling.dry_penalty_last_n), + [](common_params & params, int value) { + if (value < -1) { + throw std::runtime_error(string_format("error: invalid dry-penalty-last-n = %d\n", value)); + } + params.sampling.dry_penalty_last_n = value; + } + ).set_sampling()); + add_opt(common_arg( + {"--dry-sequence-breaker"}, "STRING", + string_format("add sequence breaker for DRY sampling, clearing out default breakers (%s) in the process; use \"none\" to not use any sequence breakers\n", + params.sampling.dry_sequence_breakers.empty() ? "none" : + std::accumulate(std::next(params.sampling.dry_sequence_breakers.begin()), + params.sampling.dry_sequence_breakers.end(), + std::string("'") + (params.sampling.dry_sequence_breakers[0] == "\n" ? "\\n" : params.sampling.dry_sequence_breakers[0]) + "'", + [](const std::string& a, const std::string& b) { + std::string formatted_b = (b == "\n") ? "\\n" : b; + return a + ", '" + formatted_b + "'"; + }).c_str()), + [](common_params & params, const std::string & value) { + static bool defaults_cleared = false; + + if (!defaults_cleared) { + params.sampling.dry_sequence_breakers.clear(); + defaults_cleared = true; + } + + if (value == "none") { + params.sampling.dry_sequence_breakers.clear(); + } else { + params.sampling.dry_sequence_breakers.emplace_back(value); + } + } + ).set_sampling()); + add_opt(common_arg( + {"--adaptive-target"}, "N", + string_format("adaptive-p: select tokens near this probability (valid range 0.0 " + "to 1.0; negative = disabled) (default: %.2f)\n" + "[(more info)](https://github.com/ggml-org/llama.cpp/pull/17927)", + (double)params.sampling.adaptive_target), + [](common_params & params, const std::string & value) { + params.sampling.adaptive_target = std::stof(value); + } + ).set_sampling()); + add_opt(common_arg( + {"--adaptive-decay"}, "N", + string_format("adaptive-p: decay rate for target adaptation over time. lower values " + "are more reactive, higher values are more stable.\n" + "(valid range 0.0 to 0.99) (default: %.2f)", + (double)params.sampling.adaptive_decay), + [](common_params & params, const std::string & value) { + params.sampling.adaptive_decay = std::stof(value); + } + ).set_sampling()); + add_opt(common_arg( + {"--dynatemp-range"}, "N", + string_format("dynamic temperature range (default: %.2f, 0.0 = disabled)", (double)params.sampling.dynatemp_range), + [](common_params & params, const std::string & value) { + params.sampling.dynatemp_range = std::stof(value); + } + ).set_sampling()); + add_opt(common_arg( + {"--dynatemp-exp"}, "N", + string_format("dynamic temperature exponent (default: %.2f)", (double)params.sampling.dynatemp_exponent), + [](common_params & params, const std::string & value) { + params.sampling.dynatemp_exponent = std::stof(value); + } + ).set_sampling()); + add_opt(common_arg( + {"--mirostat"}, "N", + string_format("use Mirostat sampling.\nTop K, Nucleus and Locally Typical samplers are ignored if used.\n" + "(default: %d, 0 = disabled, 1 = Mirostat, 2 = Mirostat 2.0)", params.sampling.mirostat), + [](common_params & params, int value) { + params.sampling.mirostat = value; + params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT; + } + ).set_sampling()); + add_opt(common_arg( + {"--mirostat-lr"}, "N", + string_format("Mirostat learning rate, parameter eta (default: %.2f)", (double)params.sampling.mirostat_eta), + [](common_params & params, const std::string & value) { + params.sampling.mirostat_eta = std::stof(value); + params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_ETA; + } + ).set_sampling()); + add_opt(common_arg( + {"--mirostat-ent"}, "N", + string_format("Mirostat target entropy, parameter tau (default: %.2f)", (double)params.sampling.mirostat_tau), + [](common_params & params, const std::string & value) { + params.sampling.mirostat_tau = std::stof(value); + params.sampling.user_sampling_config |= common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_TAU; + } + ).set_sampling()); + add_opt(common_arg( + {"-l", "--logit-bias"}, "TOKEN_ID(+/-)BIAS", + "modifies the likelihood of token appearing in the completion,\n" + "i.e. `--logit-bias 15043+1` to increase likelihood of token ' Hello',\n" + "or `--logit-bias 15043-1` to decrease likelihood of token ' Hello'", + [](common_params & params, const std::string & value) { + std::stringstream ss(value); + llama_token key; + char sign; + std::string value_str; + try { + if (ss >> key && ss >> sign && std::getline(ss, value_str) && (sign == '+' || sign == '-')) { + const float bias = std::stof(value_str) * ((sign == '-') ? -1.0f : 1.0f); + params.sampling.logit_bias.push_back({key, bias}); + } else { + throw std::invalid_argument("invalid input format"); + } + } catch (const std::exception&) { + throw std::invalid_argument("invalid input format"); + } + } + ).set_sampling()); + add_opt(common_arg( + {"--grammar"}, "GRAMMAR", + "BNF-like grammar to constrain generations (see samples in grammars/ dir)", + [](common_params & params, const std::string & value) { + params.sampling.grammar = {COMMON_GRAMMAR_TYPE_USER, value}; + } + ).set_sampling()); + add_opt(common_arg( + {"--grammar-file"}, "FNAME", + "file to read grammar from", + [](common_params & params, const std::string & value) { + params.sampling.grammar = {COMMON_GRAMMAR_TYPE_USER, read_file(value)}; + } + ).set_sampling()); + add_opt(common_arg( + {"-j", "--json-schema"}, "SCHEMA", + "JSON schema to constrain generations (https://json-schema.org/), e.g. `{}` for any JSON object\nFor schemas w/ external $refs, use --grammar + example/json_schema_to_grammar.py instead", + [](common_params & params, const std::string & value) { + params.sampling.grammar = {COMMON_GRAMMAR_TYPE_OUTPUT_FORMAT, json_schema_to_grammar(json::parse(value))}; + } + ).set_sampling()); + add_opt(common_arg( + {"-jf", "--json-schema-file"}, "FILE", + "File containing a JSON schema to constrain generations (https://json-schema.org/), e.g. `{}` for any JSON object\nFor schemas w/ external $refs, use --grammar + example/json_schema_to_grammar.py instead", + [](common_params & params, const std::string & value) { + std::ifstream file(value); + if (!file) { + throw std::runtime_error(string_format("error: failed to open file '%s'\n", value.c_str())); + } + std::string schema; + std::copy( + std::istreambuf_iterator(file), + std::istreambuf_iterator(), + std::back_inserter(schema) + ); + params.sampling.grammar = {COMMON_GRAMMAR_TYPE_OUTPUT_FORMAT, json_schema_to_grammar(json::parse(schema))}; + } + ).set_sampling()); + add_opt(common_arg( + {"-bs", "--backend-sampling"}, + "enable backend sampling (experimental) (default: disabled)", + [](common_params & params) { + params.sampling.backend_sampling = true; + } + ).set_sampling().set_env("LLAMA_ARG_BACKEND_SAMPLING")); + add_opt(common_arg( + {"--pooling"}, "{none,mean,cls,last,rank}", + "pooling type for embeddings, use model default if unspecified", + [](common_params & params, const std::string & value) { + /**/ if (value == "none") { params.pooling_type = LLAMA_POOLING_TYPE_NONE; } + else if (value == "mean") { params.pooling_type = LLAMA_POOLING_TYPE_MEAN; } + else if (value == "cls") { params.pooling_type = LLAMA_POOLING_TYPE_CLS; } + else if (value == "last") { params.pooling_type = LLAMA_POOLING_TYPE_LAST; } + else if (value == "rank") { params.pooling_type = LLAMA_POOLING_TYPE_RANK; } + else { throw std::invalid_argument("invalid value"); } + } + ).set_examples({LLAMA_EXAMPLE_EMBEDDING, LLAMA_EXAMPLE_RETRIEVAL, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_DEBUG}).set_env("LLAMA_ARG_POOLING")); + add_opt(common_arg( + {"--attention"}, "{causal,non-causal}", + "attention type for embeddings, use model default if unspecified", + [](common_params & params, const std::string & value) { + /**/ if (value == "causal") { params.attention_type = LLAMA_ATTENTION_TYPE_CAUSAL; } + else if (value == "non-causal") { params.attention_type = LLAMA_ATTENTION_TYPE_NON_CAUSAL; } + else { throw std::invalid_argument("invalid value"); } + } + ).set_examples({LLAMA_EXAMPLE_EMBEDDING})); + add_opt(common_arg( + {"--rope-scaling"}, "{none,linear,yarn}", + "RoPE frequency scaling method, defaults to linear unless specified by the model", + [](common_params & params, const std::string & value) { + /**/ if (value == "none") { params.rope_scaling_type = LLAMA_ROPE_SCALING_TYPE_NONE; } + else if (value == "linear") { params.rope_scaling_type = LLAMA_ROPE_SCALING_TYPE_LINEAR; } + else if (value == "yarn") { params.rope_scaling_type = LLAMA_ROPE_SCALING_TYPE_YARN; } + else { throw std::invalid_argument("invalid value"); } + } + ).set_env("LLAMA_ARG_ROPE_SCALING_TYPE")); + add_opt(common_arg( + {"--rope-scale"}, "N", + "RoPE context scaling factor, expands context by a factor of N", + [](common_params & params, const std::string & value) { + params.rope_freq_scale = 1.0f / std::stof(value); + } + ).set_env("LLAMA_ARG_ROPE_SCALE")); + add_opt(common_arg( + {"--rope-freq-base"}, "N", + "RoPE base frequency, used by NTK-aware scaling (default: loaded from model)", + [](common_params & params, const std::string & value) { + params.rope_freq_base = std::stof(value); + } + ).set_env("LLAMA_ARG_ROPE_FREQ_BASE")); + add_opt(common_arg( + {"--rope-freq-scale"}, "N", + "RoPE frequency scaling factor, expands context by a factor of 1/N", + [](common_params & params, const std::string & value) { + params.rope_freq_scale = std::stof(value); + } + ).set_env("LLAMA_ARG_ROPE_FREQ_SCALE")); + add_opt(common_arg( + {"--yarn-orig-ctx"}, "N", + string_format("YaRN: original context size of model (default: %d = model training context size)", params.yarn_orig_ctx), + [](common_params & params, int value) { + params.yarn_orig_ctx = value; + } + ).set_env("LLAMA_ARG_YARN_ORIG_CTX")); + add_opt(common_arg( + {"--yarn-ext-factor"}, "N", + string_format("YaRN: extrapolation mix factor (default: %.2f, 0.0 = full interpolation)", (double)params.yarn_ext_factor), + [](common_params & params, const std::string & value) { + params.yarn_ext_factor = std::stof(value); + } + ).set_env("LLAMA_ARG_YARN_EXT_FACTOR")); + add_opt(common_arg( + {"--yarn-attn-factor"}, "N", + string_format("YaRN: scale sqrt(t) or attention magnitude (default: %.2f)", (double)params.yarn_attn_factor), + [](common_params & params, const std::string & value) { + params.yarn_attn_factor = std::stof(value); + } + ).set_env("LLAMA_ARG_YARN_ATTN_FACTOR")); + add_opt(common_arg( + {"--yarn-beta-slow"}, "N", + string_format("YaRN: high correction dim or alpha (default: %.2f)", (double)params.yarn_beta_slow), + [](common_params & params, const std::string & value) { + params.yarn_beta_slow = std::stof(value); + } + ).set_env("LLAMA_ARG_YARN_BETA_SLOW")); + add_opt(common_arg( + {"--yarn-beta-fast"}, "N", + string_format("YaRN: low correction dim or beta (default: %.2f)", (double)params.yarn_beta_fast), + [](common_params & params, const std::string & value) { + params.yarn_beta_fast = std::stof(value); + } + ).set_env("LLAMA_ARG_YARN_BETA_FAST")); + add_opt(common_arg( + {"-gan", "--grp-attn-n"}, "N", + string_format("group-attention factor (default: %d)", params.grp_attn_n), + [](common_params & params, int value) { + params.grp_attn_n = value; + } + ).set_env("LLAMA_ARG_GRP_ATTN_N").set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_PASSKEY})); + add_opt(common_arg( + {"-gaw", "--grp-attn-w"}, "N", + string_format("group-attention width (default: %d)", params.grp_attn_w), + [](common_params & params, int value) { + params.grp_attn_w = value; + } + ).set_env("LLAMA_ARG_GRP_ATTN_W").set_examples({LLAMA_EXAMPLE_COMPLETION})); + add_opt(common_arg( + {"-kvo", "--kv-offload"}, + {"-nkvo", "--no-kv-offload"}, + string_format("whether to enable KV cache offloading (default: %s)", params.no_kv_offload ? "disabled" : "enabled"), + [](common_params & params, bool value) { + params.no_kv_offload = !value; + } + ).set_env("LLAMA_ARG_KV_OFFLOAD")); + add_opt(common_arg( + {"--repack"}, + {"-nr", "--no-repack"}, + string_format("whether to enable weight repacking (default: %s)", params.no_extra_bufts ? "disabled" : "enabled"), + [](common_params & params, bool value) { + params.no_extra_bufts = !value; + } + ).set_env("LLAMA_ARG_REPACK")); + add_opt(common_arg( + {"--no-host"}, + "bypass host buffer allowing extra buffers to be used", + [](common_params & params) { + params.no_host = true; + } + ).set_env("LLAMA_ARG_NO_HOST")); + add_opt(common_arg( + {"-ctk", "--cache-type-k"}, "TYPE", + string_format( + "KV cache data type for K\n" + "allowed values: %s\n" + "(default: %s)", + llama_common_get_all_kv_cache_types_rust(), + ggml_type_name(params.cache_type_k) + ), + [](common_params & params, const std::string & value) { + params.cache_type_k = kv_cache_type_from_str(value); + } + ).set_env("LLAMA_ARG_CACHE_TYPE_K")); + add_opt(common_arg( + {"-ctv", "--cache-type-v"}, "TYPE", + string_format( + "KV cache data type for V\n" + "allowed values: %s\n" + "(default: %s)", + llama_common_get_all_kv_cache_types_rust(), + ggml_type_name(params.cache_type_v) + ), + [](common_params & params, const std::string & value) { + params.cache_type_v = kv_cache_type_from_str(value); + } + ).set_env("LLAMA_ARG_CACHE_TYPE_V")); + add_opt(common_arg( + {"--hellaswag"}, + "compute HellaSwag score over random tasks from datafile supplied with -f", + [](common_params & params) { + params.hellaswag = true; + } + ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); + add_opt(common_arg( + {"--hellaswag-tasks"}, "N", + string_format("number of tasks to use when computing the HellaSwag score (default: %zu)", params.hellaswag_tasks), + [](common_params & params, int value) { + params.hellaswag_tasks = value; + } + ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); + add_opt(common_arg( + {"--winogrande"}, + "compute Winogrande score over random tasks from datafile supplied with -f", + [](common_params & params) { + params.winogrande = true; + } + ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); + add_opt(common_arg( + {"--winogrande-tasks"}, "N", + string_format("number of tasks to use when computing the Winogrande score (default: %zu)", params.winogrande_tasks), + [](common_params & params, int value) { + params.winogrande_tasks = value; + } + ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); + add_opt(common_arg( + {"--multiple-choice"}, + "compute multiple choice score over random tasks from datafile supplied with -f", + [](common_params & params) { + params.multiple_choice = true; + } + ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); + add_opt(common_arg( + {"--multiple-choice-tasks"}, "N", + string_format("number of tasks to use when computing the multiple choice score (default: %zu)", params.multiple_choice_tasks), + [](common_params & params, int value) { + params.multiple_choice_tasks = value; + } + ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); + add_opt(common_arg( + {"--kl-divergence"}, + "computes KL-divergence to logits provided via --kl-divergence-base", + [](common_params & params) { + params.kl_divergence = true; + } + ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); + add_opt(common_arg( + {"--save-all-logits", "--kl-divergence-base"}, "FNAME", + "set logits file", + [](common_params & params, const std::string & value) { + params.logits_file = value; + } + ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); + add_opt(common_arg( + {"--ppl-stride"}, "N", + string_format("stride for perplexity calculation (default: %d)", params.ppl_stride), + [](common_params & params, int value) { + params.ppl_stride = value; + } + ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); + add_opt(common_arg( + {"--ppl-output-type"}, "<0|1>", + string_format("output type for perplexity calculation (default: %d)", params.ppl_output_type), + [](common_params & params, int value) { + params.ppl_output_type = value; + } + ).set_examples({LLAMA_EXAMPLE_PERPLEXITY})); + add_opt(common_arg( + {"-dt", "--defrag-thold"}, "N", + string_format("KV cache defragmentation threshold (DEPRECATED)"), + [](common_params & params, const std::string & value) { + GGML_UNUSED(params); + GGML_UNUSED(value); + LOG_WRN("DEPRECATED: --defrag-thold is deprecated and no longer necessary to specify\n"); + } + ).set_env("LLAMA_ARG_DEFRAG_THOLD")); + if (ex == LLAMA_EXAMPLE_SERVER) { + // this is to make sure this option appears in the server-specific section of the help message + add_opt(common_arg( + {"-np", "--parallel"}, "N", + string_format("number of server slots (default: %d, -1 = auto)", params.n_parallel), + [](common_params & params, int value) { + if (value == 0) { + throw std::invalid_argument("error: invalid value for n_parallel\n"); + } + params.n_parallel = value; + } + ).set_env("LLAMA_ARG_N_PARALLEL").set_examples({LLAMA_EXAMPLE_SERVER})); + } else { + add_opt(common_arg( + {"-np", "--parallel"}, "N", + string_format("number of parallel sequences to decode (default: %d)", params.n_parallel), + [](common_params & params, int value) { + params.n_parallel = value; + } + ).set_env("LLAMA_ARG_N_PARALLEL")); + } + add_opt(common_arg( + {"-ns", "--sequences"}, "N", + string_format("number of sequences to decode (default: %d)", params.n_sequences), + [](common_params & params, int value) { + params.n_sequences = value; + } + ).set_examples({LLAMA_EXAMPLE_PARALLEL})); + add_opt(common_arg( + {"-cb", "--cont-batching"}, + {"-nocb", "--no-cont-batching"}, + string_format("whether to enable continuous batching (a.k.a dynamic batching) (default: %s)", params.cont_batching ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.cont_batching = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CONT_BATCHING")); + add_opt(common_arg( + {"-mm", "--mmproj"}, "FILE", + "path to a multimodal projector file. see tools/mtmd/README.md\n" + "note: if -hf is used, this argument can be omitted", + [](common_params & params, const std::string & value) { + params.mmproj.path = value; + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_MMPROJ")); + add_opt(common_arg( + {"-mmu", "--mmproj-url"}, "URL", + "URL to a multimodal projector file. see tools/mtmd/README.md", + [](common_params & params, const std::string & value) { + params.mmproj.url = value; + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_MMPROJ_URL")); + add_opt(common_arg( + {"--mmproj-auto"}, + {"--no-mmproj", "--no-mmproj-auto"}, + string_format("whether to use multimodal projector file (if available), useful when using -hf (default: %s)", params.no_mmproj ? "disabled" : "enabled"), + [](common_params & params, bool value) { + params.no_mmproj = !value; + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_MMPROJ_AUTO")); + add_opt(common_arg( + {"--mmproj-offload"}, + {"--no-mmproj-offload"}, + string_format("whether to enable GPU offloading for multimodal projector (default: %s)", params.mmproj_use_gpu ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.mmproj_use_gpu = value; + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_MMPROJ_OFFLOAD")); + add_opt(common_arg( + {"--image", "--audio"}, "FILE", + "path to an image or audio file. use with multimodal models, use comma-separated values for multiple files\n", + [](common_params & params, const std::string & value) { + for (const auto & item : parse_csv_row(value)) { + params.image.emplace_back(item); + } + } + ).set_examples({LLAMA_EXAMPLE_MTMD, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--image-min-tokens"}, "N", + "minimum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)", + [](common_params & params, int value) { + params.image_min_tokens = value; + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_IMAGE_MIN_TOKENS")); + add_opt(common_arg( + {"--image-max-tokens"}, "N", + "maximum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)", + [](common_params & params, int value) { + params.image_max_tokens = value; + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_IMAGE_MAX_TOKENS")); + if (llama_supports_rpc()) { + add_opt(common_arg( + {"--rpc"}, "SERVERS", + "comma separated list of RPC servers (host:port)", + [](common_params & params, const std::string & value) { + add_rpc_devices(value); + GGML_UNUSED(params); + } + ).set_env("LLAMA_ARG_RPC")); + } + add_opt(common_arg( + {"--mlock"}, + "force system to keep model in RAM rather than swapping or compressing", + [](common_params & params) { + params.use_mlock = true; + } + ).set_env("LLAMA_ARG_MLOCK")); + add_opt(common_arg( + {"--mmap"}, + {"--no-mmap"}, + string_format("whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock) (default: %s)", params.use_mmap ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.use_mmap = value; + } + ).set_env("LLAMA_ARG_MMAP")); + add_opt(common_arg( + {"-dio", "--direct-io"}, + {"-ndio", "--no-direct-io"}, + string_format("use DirectIO if available. (default: %s)", params.use_direct_io ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.use_direct_io = value; + } + ).set_env("LLAMA_ARG_DIO")); + add_opt(common_arg( + {"--numa"}, "TYPE", + "attempt optimizations that help on some NUMA systems\n" + "- distribute: spread execution evenly over all nodes\n" + "- isolate: only spawn threads on CPUs on the node that execution started on\n" + "- numactl: use the CPU map provided by numactl\n" + "if run without this previously, it is recommended to drop the system page cache before using this\n" + "see https://github.com/ggml-org/llama.cpp/issues/1437", + [](common_params & params, const std::string & value) { + /**/ if (value == "distribute" || value == "") { params.numa = GGML_NUMA_STRATEGY_DISTRIBUTE; } + else if (value == "isolate") { params.numa = GGML_NUMA_STRATEGY_ISOLATE; } + else if (value == "numactl") { params.numa = GGML_NUMA_STRATEGY_NUMACTL; } + else { throw std::invalid_argument("invalid value"); } + } + ).set_env("LLAMA_ARG_NUMA")); + add_opt(common_arg( + {"-dev", "--device"}, "", + "comma-separated list of devices to use for offloading (none = don't offload)\n" + "use --list-devices to see a list of available devices", + [](common_params & params, const std::string & value) { + params.devices = parse_device_list(value); + } + ).set_env("LLAMA_ARG_DEVICE")); + add_opt(common_arg( + {"--list-devices"}, + "print list of available devices and exit", + [](common_params &) { + std::vector devices; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + auto * dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) { + devices.push_back(dev); + } + } + printf("Available devices:\n"); + for (auto * dev : devices) { + size_t free, total; + ggml_backend_dev_memory(dev, &free, &total); + printf(" %s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024); + } + exit(0); + } + )); + add_opt(common_arg( + {"-ot", "--override-tensor"}, "=,...", + "override tensor buffer type", [](common_params & params, const std::string & value) { + parse_tensor_buffer_overrides(value, params.tensor_buft_overrides); + } + ).set_env("LLAMA_ARG_OVERRIDE_TENSOR")); + add_opt(common_arg( + {"-cmoe", "--cpu-moe"}, + "keep all Mixture of Experts (MoE) weights in the CPU", + [](common_params & params) { + params.tensor_buft_overrides.push_back(llm_ffn_exps_cpu_override()); + } + ).set_env("LLAMA_ARG_CPU_MOE")); + add_opt(common_arg( + {"-ncmoe", "--n-cpu-moe"}, "N", + "keep the Mixture of Experts (MoE) weights of the first N layers in the CPU", + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + for (int i = 0; i < value; ++i) { + // keep strings alive and avoid leaking memory by storing them in a static vector + static std::list buft_overrides; + buft_overrides.push_back(llm_ffn_exps_block_regex(i)); + params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()}); + } + } + ).set_env("LLAMA_ARG_N_CPU_MOE")); + GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0 + add_opt(common_arg( + {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N", + string_format("max. number of layers to store in VRAM, either an exact number, 'auto', or 'all' (default: %s)", params.n_gpu_layers == -1 ? "auto" : "all"), + [](common_params & params, const std::string & value) { + if (value == "auto") { + params.n_gpu_layers = -1; + } else if (value == "all") { + params.n_gpu_layers = -2; + } else { + params.n_gpu_layers = std::stoi(value); + } + if (!llama_supports_gpu_offload()) { + fprintf(stderr, "warning: no usable GPU found, --gpu-layers option will be ignored\n"); + fprintf(stderr, "warning: one possible reason is that llama.cpp was compiled without GPU support\n"); + fprintf(stderr, "warning: consult docs/build.md for compilation instructions\n"); + } + } + ).set_env("LLAMA_ARG_N_GPU_LAYERS")); + add_opt(common_arg( + {"-sm", "--split-mode"}, "{none,layer,row,tensor}", + "how to split the model across multiple GPUs, one of:\n" + "- none: use one GPU only\n" + "- layer (default): split layers and KV across GPUs (pipelined)\n" + "- row: split weight across GPUs by rows (parallelized)\n" + "- tensor: split weights and KV across GPUs (parallelized, EXPERIMENTAL)", + [](common_params & params, const std::string & value) { + if (value == "none") { + params.split_mode = LLAMA_SPLIT_MODE_NONE; + } else if (value == "layer") { + params.split_mode = LLAMA_SPLIT_MODE_LAYER; + } else if (value == "row") { + params.split_mode = LLAMA_SPLIT_MODE_ROW; + } else if (value == "tensor") { + params.split_mode = LLAMA_SPLIT_MODE_TENSOR; + } else { + throw std::invalid_argument("invalid value"); + } + if (!llama_supports_gpu_offload()) { + fprintf(stderr, "warning: llama.cpp was compiled without support for GPU offload. Setting the split mode has no effect.\n"); + } + } + ).set_env("LLAMA_ARG_SPLIT_MODE")); + add_opt(common_arg( + {"-ts", "--tensor-split"}, "N0,N1,N2,...", + "fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1", + [](common_params & params, const std::string & value) { + std::string arg_next = value; + + // split string by , and / + const std::regex regex{ R"([,/]+)" }; + std::sregex_token_iterator it{ arg_next.begin(), arg_next.end(), regex, -1 }; + std::vector split_arg{ it, {} }; + if (split_arg.size() >= llama_max_devices()) { + throw std::invalid_argument( + string_format("got %zu input configs, but system only has %zu devices", split_arg.size(), llama_max_devices()) + ); + } + for (size_t i = 0; i < llama_max_devices(); ++i) { + if (i < split_arg.size()) { + params.tensor_split[i] = std::stof(split_arg[i]); + } else { + params.tensor_split[i] = 0.0f; + } + } + if (!llama_supports_gpu_offload()) { + fprintf(stderr, "warning: llama.cpp was compiled without support for GPU offload. Setting a tensor split has no effect.\n"); + } + } + ).set_env("LLAMA_ARG_TENSOR_SPLIT")); + add_opt(common_arg( + {"-mg", "--main-gpu"}, "INDEX", + string_format("the GPU to use for the model (with split-mode = none), or for intermediate results and KV (with split-mode = row) (default: %d)", params.main_gpu), + [](common_params & params, int value) { + params.main_gpu = value; + if (!llama_supports_gpu_offload()) { + fprintf(stderr, "warning: llama.cpp was compiled without support for GPU offload. Setting the main GPU has no effect.\n"); + } + } + ).set_env("LLAMA_ARG_MAIN_GPU")); + add_opt(common_arg( + { "-fit", "--fit" }, "[on|off]", + string_format("whether to adjust unset arguments to fit in device memory ('on' or 'off', default: '%s')", params.fit_params ? "on" : "off"), + [](common_params & params, const std::string & value) { + if (is_truthy(value)) { + params.fit_params = true; + } else if (is_falsey(value)) { + params.fit_params = false; + } else { + throw std::runtime_error( + string_format("error: unknown value for --fit: '%s'\n", value.c_str())); + } + } + ).set_env("LLAMA_ARG_FIT")); + add_opt(common_arg( + { "-fitp", "--fit-print" }, "[on|off]", + string_format("print the estimated required memory ('on' or 'off', default: '%s')", params.fit_params_print ? "on" : "off"), + [](common_params & params, const std::string & value) { + if (is_truthy(value)) { + params.fit_params_print = true; + } else if (is_falsey(value)) { + params.fit_params_print = false; + } else { + throw std::runtime_error( + string_format("error: unknown value for --fit-print: '%s'\n", value.c_str())); + } + } + ).set_examples({LLAMA_EXAMPLE_FIT_PARAMS}).set_env("LLAMA_ARG_FIT_ESTIMATE")); + add_opt(common_arg( + { "-fitt", "--fit-target" }, "MiB0,MiB1,MiB2,...", + string_format("target margin per device for --fit, comma-separated list of values, " + "single value is broadcast across all devices, default: %zu", params.fit_params_target[0]/(1024*1024)), + [](common_params & params, const std::string & value) { + std::string arg_next = value; + + // split string by , and / + const std::regex regex{ R"([,/]+)" }; + std::sregex_token_iterator it{ arg_next.begin(), arg_next.end(), regex, -1 }; + std::vector split_arg{ it, {} }; + if (split_arg.size() >= llama_max_devices()) { + throw std::invalid_argument( + string_format("got %zu input configs, but system only has %zu devices", split_arg.size(), llama_max_devices()) + ); + } + if (split_arg.size() == 1) { + std::fill(params.fit_params_target.begin(), params.fit_params_target.end(), std::stoull(split_arg[0]) * 1024*1024); + return; + } + for (size_t i = 0; i < split_arg.size(); i++) { + params.fit_params_target[i] = std::stoull(split_arg[i]) * 1024*1024; + } + } + ).set_env("LLAMA_ARG_FIT_TARGET")); + add_opt(common_arg( + { "-fitc", "--fit-ctx" }, "N", + string_format("minimum ctx size that can be set by --fit option, default: %" PRIu32, params.fit_params_min_ctx), + [](common_params & params, int value) { + params.fit_params_min_ctx = value; + } + ).set_env("LLAMA_ARG_FIT_CTX")); + add_opt(common_arg( + {"--check-tensors"}, + string_format("check model tensor data for invalid values (default: %s)", params.check_tensors ? "true" : "false"), + [](common_params & params) { + params.check_tensors = true; + } + )); + add_opt(common_arg( + {"--override-kv"}, "KEY=TYPE:VALUE,...", + "advanced option to override model metadata by key. to specify multiple overrides, either use comma-separated values.\n" + "types: int, float, bool, str. example: --override-kv tokenizer.ggml.add_bos_token=bool:false,tokenizer.ggml.add_eos_token=bool:false", + [](common_params & params, const std::string & value) { + for (const auto & item : parse_csv_row(value)) { + if (!string_parse_kv_override(item.c_str(), params.kv_overrides)) { + throw std::runtime_error(string_format("error: Invalid type for KV override: %s\n", item.c_str())); + } + } + } + )); + add_opt(common_arg( + {"--op-offload"}, + {"--no-op-offload"}, + string_format("whether to offload host tensor operations to device (default: %s)", params.no_op_offload ? "false" : "true"), + [](common_params & params, bool value) { + params.no_op_offload = !value; + } + )); + add_opt(common_arg( + {"--lora"}, "FNAME", + "path to LoRA adapter (use comma-separated values to load multiple adapters)", + [](common_params & params, const std::string & value) { + for (const auto & item : parse_csv_row(value)) { + params.lora_adapters.push_back({ item, 1.0, "", "", nullptr }); + } + } + // we define this arg on both COMMON and EXPORT_LORA, so when showing help message of export-lora, it will be categorized as "example-specific" arg + ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_EXPORT_LORA})); + add_opt(common_arg( + {"--lora-scaled"}, "FNAME:SCALE,...", + "path to LoRA adapter with user defined scaling (format: FNAME:SCALE,...)\n" + "note: use comma-separated values", + [](common_params & params, const std::string & value) { + for (const auto & item : parse_csv_row(value)) { + auto parts = string_split(item, ':'); + if (parts.size() != 2) { + throw std::invalid_argument("lora-scaled format: FNAME:SCALE"); + } + params.lora_adapters.push_back({ parts[0], std::stof(parts[1]), "", "", nullptr }); + } + } + // we define this arg on both COMMON and EXPORT_LORA, so when showing help message of export-lora, it will be categorized as "example-specific" arg + ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_EXPORT_LORA})); + add_opt(common_arg( + {"--control-vector"}, "FNAME", + "add a control vector\nnote: use comma-separated values to add multiple control vectors", + [](common_params & params, const std::string & value) { + for (const auto & item : parse_csv_row(value)) { + params.control_vectors.push_back({ 1.0f, item, }); + } + } + )); + add_opt(common_arg( + {"--control-vector-scaled"}, "FNAME:SCALE,...", + "add a control vector with user defined scaling SCALE\n" + "note: use comma-separated values (format: FNAME:SCALE,...)", + [](common_params & params, const std::string & value) { + for (const auto & item : parse_csv_row(value)) { + auto parts = string_split(item, ':'); + if (parts.size() != 2) { + throw std::invalid_argument("control-vector-scaled format: FNAME:SCALE"); + } + params.control_vectors.push_back({ std::stof(parts[1]), parts[0] }); + } + } + )); + add_opt(common_arg( + {"--control-vector-layer-range"}, "START", "END", + "layer range to apply the control vector(s) to, start and end inclusive", + [](common_params & params, const std::string & start, const std::string & end) { + params.control_vector_layer_start = std::stoi(start); + params.control_vector_layer_end = std::stoi(end); + } + )); + add_opt(common_arg( + {"-a", "--alias"}, "STRING", + "set model name aliases, comma-separated (to be used by API)", + [](common_params & params, const std::string & value) { + for (auto & alias : string_split(value, ',')) { + alias = string_strip(alias); + if (!alias.empty()) { + params.model_alias.insert(alias); + } + } + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_ALIAS")); + add_opt(common_arg( + {"--tags"}, "STRING", + "set model tags, comma-separated (informational, not used for routing)", + [](common_params & params, const std::string & value) { + for (auto & tag : string_split(value, ',')) { + tag = string_strip(tag); + if (!tag.empty()) { + params.model_tags.insert(tag); + } + } + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TAGS")); + add_opt(common_arg( + {"-m", "--model"}, "FNAME", + ex == LLAMA_EXAMPLE_EXPORT_LORA + ? "model path from which to load base model" + : "model path to load", + [](common_params & params, const std::string & value) { + params.model.path = value; + } + ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_EXPORT_LORA}).set_env("LLAMA_ARG_MODEL")); + add_opt(common_arg( + {"-mu", "--model-url"}, "MODEL_URL", + "model download url (default: unused)", + [](common_params & params, const std::string & value) { + params.model.url = value; + } + ).set_env("LLAMA_ARG_MODEL_URL")); + add_opt(common_arg( + { "-dr", "--docker-repo" }, "[/][:quant]", + "Docker Hub model repository. repo is optional, default to ai/. quant is optional, default to :latest.\n" + "example: gemma3\n" + "(default: unused)", + [](common_params & params, const std::string & value) { + params.model.docker_repo = value; + } + ).set_env("LLAMA_ARG_DOCKER_REPO")); + add_opt(common_arg( + {"-hf", "-hfr", "--hf-repo"}, "/[:quant]", + "Hugging Face model repository; quant is optional, case-insensitive, default to Q4_K_M, or falls back to the first file in the repo if Q4_K_M doesn't exist.\n" + "mmproj is also downloaded automatically if available. to disable, add --no-mmproj\n" + "example: ggml-org/GLM-4.7-Flash-GGUF:Q4_K_M\n" + "(default: unused)", + [](common_params & params, const std::string & value) { + params.model.hf_repo = value; + } + ).set_env("LLAMA_ARG_HF_REPO")); + add_opt(common_arg( + {"-hff", "--hf-file"}, "FILE", + "Hugging Face model file. If specified, it will override the quant in --hf-repo (default: unused)", + [](common_params & params, const std::string & value) { + params.model.hf_file = value; + } + ).set_env("LLAMA_ARG_HF_FILE")); + add_opt(common_arg( + {"-hfv", "-hfrv", "--hf-repo-v"}, "/[:quant]", + "Hugging Face model repository for the vocoder model (default: unused)", + [](common_params & params, const std::string & value) { + params.vocoder.model.hf_repo = value; + } + ).set_env("LLAMA_ARG_HF_REPO_V")); + add_opt(common_arg( + {"-hffv", "--hf-file-v"}, "FILE", + "Hugging Face model file for the vocoder model (default: unused)", + [](common_params & params, const std::string & value) { + params.vocoder.model.hf_file = value; + } + ).set_env("LLAMA_ARG_HF_FILE_V")); + add_opt(common_arg( + {"-hft", "--hf-token"}, "TOKEN", + "Hugging Face access token (default: value from HF_TOKEN environment variable)", + [](common_params & params, const std::string & value) { + params.hf_token = value; + } + ).set_env("HF_TOKEN")); + add_opt(common_arg( + {"--context-file"}, "FNAME", + "file to load context from (use comma-separated values to specify multiple files)", + [](common_params & params, const std::string & value) { + for (const auto & item : parse_csv_row(value)) { + std::ifstream file(item, std::ios::binary); + if (!file) { + throw std::runtime_error(string_format("error: failed to open file '%s'\n", item.c_str())); + } + params.context_files.push_back(item); + } + } + ).set_examples({LLAMA_EXAMPLE_RETRIEVAL})); + add_opt(common_arg( + {"--chunk-size"}, "N", + string_format("minimum length of embedded text chunks (default: %d)", params.chunk_size), + [](common_params & params, int value) { + params.chunk_size = value; + } + ).set_examples({LLAMA_EXAMPLE_RETRIEVAL})); + add_opt(common_arg( + {"--chunk-separator"}, "STRING", + string_format("separator between chunks (default: '%s')", params.chunk_separator.c_str()), + [](common_params & params, const std::string & value) { + params.chunk_separator = value; + } + ).set_examples({LLAMA_EXAMPLE_RETRIEVAL})); + add_opt(common_arg( + {"--junk"}, "N", + string_format("number of times to repeat the junk text (default: %d)", params.n_junk), + [](common_params & params, int value) { + params.n_junk = value; + } + ).set_examples({LLAMA_EXAMPLE_PASSKEY, LLAMA_EXAMPLE_PARALLEL})); + add_opt(common_arg( + {"--pos"}, "N", + string_format("position of the passkey in the junk text (default: %d)", params.i_pos), + [](common_params & params, int value) { + params.i_pos = value; + } + ).set_examples({LLAMA_EXAMPLE_PASSKEY})); + add_opt(common_arg( + {"-o", "--output", "--output-file"}, "FNAME", + string_format("output file (default: '%s')", params.out_file.c_str()), + [](common_params & params, const std::string & value) { + params.out_file = value; + } + ).set_examples({LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_CVECTOR_GENERATOR, LLAMA_EXAMPLE_EXPORT_LORA, LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_FINETUNE, + LLAMA_EXAMPLE_RESULTS, LLAMA_EXAMPLE_EXPORT_GRAPH_OPS})); + add_opt(common_arg( + {"-ofreq", "--output-frequency"}, "N", + string_format("output the imatrix every N iterations (default: %d)", params.n_out_freq), + [](common_params & params, int value) { + params.n_out_freq = value; + } + ).set_examples({LLAMA_EXAMPLE_IMATRIX})); + add_opt(common_arg( + {"--output-format"}, "{gguf,dat}", + string_format("output format for imatrix file (default: %s)", params.imat_dat > 0 ? "dat" : "gguf"), + [](common_params & params, const std::string & value) { + /**/ if (value == "gguf") { params.imat_dat = -1; } + else if (value == "dat") { params.imat_dat = 1; } + else { throw std::invalid_argument("invalid output format"); } + } + ).set_examples({LLAMA_EXAMPLE_IMATRIX})); + add_opt(common_arg( + {"--save-frequency"}, "N", + string_format("save an imatrix copy every N iterations (default: %d)", params.n_save_freq), + [](common_params & params, int value) { + params.n_save_freq = value; + } + ).set_examples({LLAMA_EXAMPLE_IMATRIX})); + add_opt(common_arg( + {"--process-output"}, + string_format("collect data for the output tensor (default: %s)", params.process_output ? "true" : "false"), + [](common_params & params) { + params.process_output = true; + } + ).set_examples({LLAMA_EXAMPLE_IMATRIX})); + add_opt(common_arg( + {"--ppl"}, + {"--no-ppl"}, + string_format("whether to compute perplexity (default: %s)", params.compute_ppl ? "true" : "false"), + [](common_params & params, bool value) { + params.compute_ppl = value; + } + ).set_examples({LLAMA_EXAMPLE_IMATRIX})); + add_opt(common_arg( + {"--chunk", "--from-chunk"}, "N", + string_format("start processing the input from chunk N (default: %d)", params.i_chunk), + [](common_params & params, int value) { + params.i_chunk = value; + } + ).set_examples({LLAMA_EXAMPLE_IMATRIX})); + add_opt(common_arg( + {"--show-statistics"}, + string_format("show imatrix statistics and then exit (default: %s)", params.show_statistics ? "true" : "false"), + [](common_params & params) { + params.show_statistics = true; + } + ).set_examples({LLAMA_EXAMPLE_IMATRIX})); + add_opt(common_arg( + {"--parse-special"}, + string_format("parse special tokens (chat, tool, etc) (default: %s)", params.parse_special ? "true" : "false"), + [](common_params & params) { + params.parse_special = true; + } + ).set_examples({LLAMA_EXAMPLE_IMATRIX})); + add_opt(common_arg( + {"-pps"}, + string_format("is the prompt shared across parallel sequences (default: %s)", params.is_pp_shared ? "true" : "false"), + [](common_params & params) { + params.is_pp_shared = true; + } + ).set_examples({LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL})); + add_opt(common_arg( + {"-tgs"}, + string_format("is the text generation separated across the different sequences (default: %s)", params.is_tg_separate ? "true" : "false"), + [](common_params & params) { + params.is_tg_separate = true; + } + ).set_examples({LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL})); + add_opt(common_arg( + {"-npp"}, "n0,n1,...", + "number of prompt tokens", + [](common_params & params, const std::string & value) { + auto p = string_split(value, ','); + params.n_pp.insert(params.n_pp.end(), p.begin(), p.end()); + } + ).set_examples({LLAMA_EXAMPLE_BENCH})); + add_opt(common_arg( + {"-ntg"}, "n0,n1,...", + "number of text generation tokens", + [](common_params & params, const std::string & value) { + auto p = string_split(value, ','); + params.n_tg.insert(params.n_tg.end(), p.begin(), p.end()); + } + ).set_examples({LLAMA_EXAMPLE_BENCH})); + add_opt(common_arg( + {"-npl"}, "n0,n1,...", + "number of parallel prompts", + [](common_params & params, const std::string & value) { + auto p = string_split(value, ','); + params.n_pl.insert(params.n_pl.end(), p.begin(), p.end()); + } + ).set_examples({LLAMA_EXAMPLE_BENCH})); + add_opt(common_arg( + {"--embd-normalize"}, "N", + string_format("normalisation for embeddings (default: %d) (-1=none, 0=max absolute int16, 1=taxicab, 2=euclidean, >2=p-norm)", params.embd_normalize), + [](common_params & params, int value) { + params.embd_normalize = value; + } + ).set_examples({LLAMA_EXAMPLE_EMBEDDING, LLAMA_EXAMPLE_DEBUG})); + add_opt(common_arg( + {"--embd-output-format"}, "FORMAT", + "empty = default, \"array\" = [[],[]...], \"json\" = openai style, \"json+\" = same \"json\" + cosine similarity matrix, \"raw\" = plain whitespace-delimited output (one embedding per line)", + [](common_params & params, const std::string & value) { + params.embd_out = value; + } + ).set_examples({LLAMA_EXAMPLE_EMBEDDING})); + add_opt(common_arg( + {"--embd-separator"}, "STRING", + "separator of embeddings (default \\n) for example \"<#sep#>\"", + [](common_params & params, const std::string & value) { + params.embd_sep = value; + } + ).set_examples({LLAMA_EXAMPLE_EMBEDDING})); + add_opt(common_arg( + {"--cls-separator"}, "STRING", + "separator of classification sequences (default \\t) for example \"<#seq#>\"", + [](common_params & params, const std::string & value) { + params.cls_sep = value; + } + ).set_examples({LLAMA_EXAMPLE_EMBEDDING})); + add_opt(common_arg( + {"--host"}, "HOST", + string_format("ip address to listen, or bind to an UNIX socket if the address ends with .sock (default: %s)", params.hostname.c_str()), + [](common_params & params, const std::string & value) { + params.hostname = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_HOST")); + add_opt(common_arg( + {"--port"}, "PORT", + string_format("port to listen (default: %d)", params.port), + [](common_params & params, int value) { + params.port = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_PORT")); + add_opt(common_arg( + {"--reuse-port"}, + string_format("allow multiple sockets to bind to the same port (default: %s)", params.reuse_port ? "enabled" : "disabled"), + [](common_params & params) { + params.reuse_port = true; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_REUSE_PORT")); + add_opt(common_arg( + {"--path"}, "PATH", + string_format("path to serve static files from (default: %s)", params.public_path.c_str()), + [](common_params & params, const std::string & value) { + params.public_path = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_STATIC_PATH")); + add_opt(common_arg( + {"--api-prefix"}, "PREFIX", + string_format("prefix path the server serves from, without the trailing slash (default: %s)", params.api_prefix.c_str()), + [](common_params & params, const std::string & value) { + params.api_prefix = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_API_PREFIX")); + add_opt(common_arg( + {"--webui-config"}, "JSON", + "JSON that provides default WebUI settings (overrides WebUI defaults)", + [](common_params & params, const std::string & value) { + params.webui_config_json = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_WEBUI_CONFIG")); + add_opt(common_arg( + {"--webui-config-file"}, "PATH", + "JSON file that provides default WebUI settings (overrides WebUI defaults)", + [](common_params & params, const std::string & value) { + params.webui_config_json = read_file(value); + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_WEBUI_CONFIG_FILE")); + add_opt(common_arg( + {"--webui-mcp-proxy"}, + {"--no-webui-mcp-proxy"}, + string_format("experimental: whether to enable MCP CORS proxy - do not enable in untrusted environments (default: %s)", params.webui_mcp_proxy ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.webui_mcp_proxy = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_WEBUI_MCP_PROXY")); + add_opt(common_arg( + {"--tools"}, "TOOL1,TOOL2,...", + "experimental: whether to enable built-in tools for AI agents - do not enable in untrusted environments (default: no tools)\n" + "specify \"all\" to enable all tools\n" + "available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, apply_diff", + [](common_params & params, const std::string & value) { + params.server_tools = parse_csv_row(value); + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TOOLS")); + add_opt(common_arg( + {"--webui"}, + {"--no-webui"}, + string_format("whether to enable the Web UI (default: %s)", params.webui ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.webui = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_WEBUI")); + add_opt(common_arg( + {"--embedding", "--embeddings"}, + string_format("restrict to only support embedding use case; use only with dedicated embedding models (default: %s)", params.embedding ? "enabled" : "disabled"), + [](common_params & params) { + params.embedding = true; + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_DEBUG}).set_env("LLAMA_ARG_EMBEDDINGS")); + add_opt(common_arg( + {"--rerank", "--reranking"}, + string_format("enable reranking endpoint on server (default: %s)", "disabled"), + [](common_params & params) { + params.embedding = true; + params.pooling_type = LLAMA_POOLING_TYPE_RANK; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_RERANKING")); + add_opt(common_arg( + {"--api-key"}, "KEY", + "API key to use for authentication, multiple keys can be provided as a comma-separated list (default: none)", + [](common_params & params, const std::string & value) { + for (const auto & key : parse_csv_row(value)) { + if (!key.empty()) { + params.api_keys.push_back(key); + } + } + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_API_KEY")); + add_opt(common_arg( + {"--api-key-file"}, "FNAME", + "path to file containing API keys (default: none)", + [](common_params & params, const std::string & value) { + std::ifstream key_file(value); + if (!key_file) { + throw std::runtime_error(string_format("error: failed to open file '%s'\n", value.c_str())); + } + std::string key; + while (std::getline(key_file, key)) { + if (!key.empty()) { + params.api_keys.push_back(key); + } + } + key_file.close(); + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--ssl-key-file"}, "FNAME", + "path to file a PEM-encoded SSL private key", + [](common_params & params, const std::string & value) { + params.ssl_file_key = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SSL_KEY_FILE")); + add_opt(common_arg( + {"--ssl-cert-file"}, "FNAME", + "path to file a PEM-encoded SSL certificate", + [](common_params & params, const std::string & value) { + params.ssl_file_cert = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SSL_CERT_FILE")); + add_opt(common_arg( + {"--chat-template-kwargs"}, "STRING", + "sets additional params for the json template parser, must be a valid json object string, e.g. '{\"key1\":\"value1\",\"key2\":\"value2\"}'", + [](common_params & params, const std::string & value) { + auto parsed = json::parse(value); + for (const auto & item : parsed.items()) { + if (item.key() == "enable_thinking") { + LOG_WRN("Setting 'enable_thinking' via --chat-template-kwargs is deprecated. " + "Use --reasoning on / --reasoning off instead.\n"); + } + params.default_template_kwargs[item.key()] = item.value().dump(); + } + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_CHAT_TEMPLATE_KWARGS")); + add_opt(common_arg( + {"-to", "--timeout"}, "N", + string_format("server read/write timeout in seconds (default: %d)", params.timeout_read), + [](common_params & params, int value) { + params.timeout_read = value; + params.timeout_write = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TIMEOUT")); + add_opt(common_arg( + {"--threads-http"}, "N", + string_format("number of threads used to process HTTP requests (default: %d)", params.n_threads_http), + [](common_params & params, int value) { + params.n_threads_http = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_THREADS_HTTP")); + add_opt(common_arg( + {"--cache-prompt"}, + {"--no-cache-prompt"}, + string_format("whether to enable prompt caching (default: %s)", params.cache_prompt ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.cache_prompt = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CACHE_PROMPT")); + add_opt(common_arg( + {"--cache-reuse"}, "N", + string_format( + "min chunk size to attempt reusing from the cache via KV shifting, requires prompt caching to be enabled (default: %d)\n" + "[(card)](https://ggml.ai/f0.png)", params.n_cache_reuse + ), + [](common_params & params, int value) { + params.n_cache_reuse = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CACHE_REUSE")); + add_opt(common_arg( + {"--metrics"}, + string_format("enable prometheus compatible metrics endpoint (default: %s)", params.endpoint_metrics ? "enabled" : "disabled"), + [](common_params & params) { + params.endpoint_metrics = true; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_ENDPOINT_METRICS")); + add_opt(common_arg( + {"--props"}, + string_format("enable changing global properties via POST /props (default: %s)", params.endpoint_props ? "enabled" : "disabled"), + [](common_params & params) { + params.endpoint_props = true; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_ENDPOINT_PROPS")); + add_opt(common_arg( + {"--slots"}, + {"--no-slots"}, + string_format("expose slots monitoring endpoint (default: %s)", params.endpoint_slots ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.endpoint_slots = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_ENDPOINT_SLOTS")); + add_opt(common_arg( + {"--slot-save-path"}, "PATH", + "path to save slot kv cache (default: disabled)", + [](common_params & params, const std::string & value) { + params.slot_save_path = value; + if (!fs_is_directory(params.slot_save_path)) { + throw std::invalid_argument("not a directory: " + value); + } + // if doesn't end with DIRECTORY_SEPARATOR, add it + if (!params.slot_save_path.empty() && params.slot_save_path[params.slot_save_path.size() - 1] != DIRECTORY_SEPARATOR) { + params.slot_save_path += DIRECTORY_SEPARATOR; + } + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--media-path"}, "PATH", + "directory for loading local media files; files can be accessed via file:// URLs using relative paths (default: disabled)", + [](common_params & params, const std::string & value) { + params.media_path = value; + if (!fs_is_directory(params.media_path)) { + throw std::invalid_argument("not a directory: " + value); + } + // if doesn't end with DIRECTORY_SEPARATOR, add it + if (!params.media_path.empty() && params.media_path[params.media_path.size() - 1] != DIRECTORY_SEPARATOR) { + params.media_path += DIRECTORY_SEPARATOR; + } + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--models-dir"}, "PATH", + "directory containing models for the router server (default: disabled)", + [](common_params & params, const std::string & value) { + params.models_dir = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODELS_DIR")); + add_opt(common_arg( + {"--models-preset"}, "PATH", + "path to INI file containing model presets for the router server (default: disabled)", + [](common_params & params, const std::string & value) { + params.models_preset = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODELS_PRESET")); + add_opt(common_arg( + {"--models-max"}, "N", + string_format("for router server, maximum number of models to load simultaneously (default: %d, 0 = unlimited)", params.models_max), + [](common_params & params, int value) { + params.models_max = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODELS_MAX")); + add_opt(common_arg( + {"--models-autoload"}, + {"--no-models-autoload"}, + string_format("for router server, whether to automatically load models (default: %s)", params.models_autoload ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.models_autoload = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODELS_AUTOLOAD")); + add_opt(common_arg( + {"--jinja"}, + {"--no-jinja"}, + string_format("whether to use jinja template engine for chat (default: %s)", params.use_jinja ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.use_jinja = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_MTMD}).set_env("LLAMA_ARG_JINJA")); + add_opt(common_arg( + {"--reasoning-format"}, "FORMAT", + "controls whether thought tags are allowed and/or extracted from the response, and in which format they're returned; one of:\n" + "- none: leaves thoughts unparsed in `message.content`\n" + "- deepseek: puts thoughts in `message.reasoning_content`\n" + "- deepseek-legacy: keeps `` tags in `message.content` while also populating `message.reasoning_content`\n" + "(default: auto)", + [](common_params & params, const std::string & value) { + params.reasoning_format = common_reasoning_format_from_name(value); + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK")); + add_opt(common_arg( + {"-rea", "--reasoning"}, "[on|off|auto]", + "Use reasoning/thinking in the chat ('on', 'off', or 'auto', default: 'auto' (detect from template))", + [](common_params & params, const std::string & value) { + if (is_truthy(value)) { + params.enable_reasoning = 1; + params.default_template_kwargs["enable_thinking"] = "true"; + } else if (is_falsey(value)) { + params.enable_reasoning = 0; + params.default_template_kwargs["enable_thinking"] = "false"; + } else if (is_autoy(value)) { + params.enable_reasoning = -1; + } else { + throw std::invalid_argument( + string_format("error: unknown value for --reasoning: '%s'\n", value.c_str())); + } + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_REASONING")); + add_opt(common_arg( + {"--reasoning-budget"}, "N", + "token budget for thinking: -1 for unrestricted, 0 for immediate end, N>0 for token budget (default: -1)", + [](common_params & params, int value) { + if (value < -1) { throw std::invalid_argument("invalid value"); } + params.sampling.reasoning_budget_tokens = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET")); + add_opt(common_arg( + {"--reasoning-budget-message"}, "MESSAGE", + "message injected before the end-of-thinking tag when reasoning budget is exhausted (default: none)", + [](common_params & params, const std::string & value) { + params.sampling.reasoning_budget_message = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_MESSAGE")); + add_opt(common_arg( + {"--chat-template"}, "JINJA_TEMPLATE", + string_format( + "set custom jinja chat template (default: template taken from model's metadata)\n" + "if suffix/prefix are specified, template will be disabled\n" + "only commonly used templates are accepted (unless --jinja is set before this flag):\n" + "list of built-in templates:\n%s", list_builtin_chat_templates().c_str() + ), + [](common_params & params, const std::string & value) { + params.chat_template = value; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_MTMD}).set_env("LLAMA_ARG_CHAT_TEMPLATE")); + add_opt(common_arg( + {"--chat-template-file"}, "JINJA_TEMPLATE_FILE", + string_format( + "set custom jinja chat template file (default: template taken from model's metadata)\n" + "if suffix/prefix are specified, template will be disabled\n" + "only commonly used templates are accepted (unless --jinja is set before this flag):\n" + "list of built-in templates:\n%s", list_builtin_chat_templates().c_str() + ), + [](common_params & params, const std::string & value) { + params.chat_template = read_file(value); + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CHAT_TEMPLATE_FILE")); + add_opt(common_arg( + {"--skip-chat-parsing"}, + {"--no-skip-chat-parsing"}, + string_format( + "force a pure content parser, even if a Jinja template is specified; model will output everything " + "in the content section, including any reasoning and/or tool calls (default: disabled)" + ), + [](common_params & params, bool value) { + params.force_pure_content_parser = value; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SKIP_CHAT_PARSING")); + add_opt(common_arg( + {"--prefill-assistant"}, + {"--no-prefill-assistant"}, + string_format( + "whether to prefill the assistant's response if the last message is an assistant message (default: prefill enabled)\n" + "when this flag is set, if the last message is an assistant message then it will be treated as a full message and not prefilled\n" + ), + [](common_params & params, bool value) { + params.prefill_assistant = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_PREFILL_ASSISTANT")); + add_opt(common_arg( + {"-sps", "--slot-prompt-similarity"}, "SIMILARITY", + string_format("how much the prompt of a request must match the prompt of a slot in order to use that slot (default: %.2f, 0.0 = disabled)\n", params.slot_prompt_similarity), + [](common_params & params, const std::string & value) { + params.slot_prompt_similarity = std::stof(value); + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--lora-init-without-apply"}, + string_format("load LoRA adapters without applying them (apply later via POST /lora-adapters) (default: %s)", params.lora_init_without_apply ? "enabled" : "disabled"), + [](common_params & params) { + params.lora_init_without_apply = true; + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--sleep-idle-seconds"}, "SECONDS", + string_format("number of seconds of idleness after which the server will sleep (default: %d; -1 = disabled)", params.sleep_idle_seconds), + [](common_params & params, int value) { + if (value == 0 || value < -1) { + throw std::invalid_argument("invalid value: cannot be 0 or less than -1"); + } + params.sleep_idle_seconds = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--simple-io"}, + "use basic IO for better compatibility in subprocesses and limited consoles", + [](common_params & params) { + params.simple_io = true; + } + ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--positive-file"}, "FNAME", + string_format("positive prompts file, one prompt per line (default: '%s')", params.cvector_positive_file.c_str()), + [](common_params & params, const std::string & value) { + params.cvector_positive_file = value; + } + ).set_examples({LLAMA_EXAMPLE_CVECTOR_GENERATOR})); + add_opt(common_arg( + {"--negative-file"}, "FNAME", + string_format("negative prompts file, one prompt per line (default: '%s')", params.cvector_negative_file.c_str()), + [](common_params & params, const std::string & value) { + params.cvector_negative_file = value; + } + ).set_examples({LLAMA_EXAMPLE_CVECTOR_GENERATOR})); + add_opt(common_arg( + {"--pca-batch"}, "N", + string_format("batch size used for PCA. Larger batch runs faster, but uses more memory (default: %d)", params.n_pca_batch), + [](common_params & params, int value) { + params.n_pca_batch = value; + } + ).set_examples({LLAMA_EXAMPLE_CVECTOR_GENERATOR})); + add_opt(common_arg( + {"--pca-iter"}, "N", + string_format("number of iterations used for PCA (default: %d)", params.n_pca_iterations), + [](common_params & params, int value) { + params.n_pca_iterations = value; + } + ).set_examples({LLAMA_EXAMPLE_CVECTOR_GENERATOR})); + add_opt(common_arg( + {"--method"}, "{pca, mean}", + "dimensionality reduction method to be used (default: pca)", + [](common_params & params, const std::string & value) { + /**/ if (value == "pca") { params.cvector_dimre_method = DIMRE_METHOD_PCA; } + else if (value == "mean") { params.cvector_dimre_method = DIMRE_METHOD_MEAN; } + else { throw std::invalid_argument("invalid value"); } + } + ).set_examples({LLAMA_EXAMPLE_CVECTOR_GENERATOR})); + add_opt(common_arg( + {"--output-format"}, "{md,jsonl}", + "output format for batched-bench results (default: md)", + [](common_params & params, const std::string & value) { + /**/ if (value == "jsonl") { params.batched_bench_output_jsonl = true; } + else if (value == "md") { params.batched_bench_output_jsonl = false; } + else { throw std::invalid_argument("invalid value"); } + } + ).set_examples({LLAMA_EXAMPLE_BENCH})); + add_opt(common_arg( + {"--log-disable"}, + "Log disable", + [](common_params &) { + common_log_pause(common_log_main()); + } + )); + add_opt(common_arg( + {"--log-file"}, "FNAME", + "Log to file", + [](common_params &, const std::string & value) { + common_log_set_file(common_log_main(), value.c_str()); + } + ).set_env("LLAMA_LOG_FILE")); + add_opt(common_arg( + {"--log-colors"}, "[on|off|auto]", + "Set colored logging ('on', 'off', or 'auto', default: 'auto')\n" + "'auto' enables colors when output is to a terminal", + [](common_params &, const std::string & value) { + if (is_truthy(value)) { + common_log_set_colors(common_log_main(), LOG_COLORS_ENABLED); + } else if (is_falsey(value)) { + common_log_set_colors(common_log_main(), LOG_COLORS_DISABLED); + } else if (is_autoy(value)) { + common_log_set_colors(common_log_main(), LOG_COLORS_AUTO); + } else { + throw std::invalid_argument( + string_format("error: unknown value for --log-colors: '%s'\n", value.c_str())); + } + } + ).set_env("LLAMA_LOG_COLORS")); + add_opt(common_arg( + {"-v", "--verbose", "--log-verbose"}, + "Set verbosity level to infinity (i.e. log all messages, useful for debugging)", + [](common_params & params) { + params.verbosity = INT_MAX; + common_log_set_verbosity_thold(INT_MAX); + } + )); + add_opt(common_arg( + {"--offline"}, + "Offline mode: forces use of cache, prevents network access", + [](common_params & params) { + params.offline = true; + } + ).set_env("LLAMA_OFFLINE")); + add_opt(common_arg( + {"-lv", "--verbosity", "--log-verbosity"}, "N", + string_format("Set the verbosity threshold. Messages with a higher verbosity will be ignored. Values:\n" + " - 0: generic output\n" + " - 1: error\n" + " - 2: warning\n" + " - 3: info\n" + " - 4: debug\n" + "(default: %d)\n", params.verbosity), + [](common_params & params, int value) { + params.verbosity = value; + common_log_set_verbosity_thold(value); + } + ).set_env("LLAMA_LOG_VERBOSITY")); + add_opt(common_arg( + {"--log-prefix"}, + "Enable prefix in log messages", + [](common_params &) { + common_log_set_prefix(common_log_main(), true); + } + ).set_env("LLAMA_LOG_PREFIX")); + add_opt(common_arg( + {"--log-timestamps"}, + "Enable timestamps in log messages", + [](common_params &) { + common_log_set_timestamps(common_log_main(), true); + } + ).set_env("LLAMA_LOG_TIMESTAMPS")); + + // + // speculative parameters + // + + add_opt(common_arg( + {"--spec-draft-hf", "-hfd", "-hfrd", "--hf-repo-draft"}, "/[:quant]", + "Same as --hf-repo, but for the draft model (default: unused)", + [](common_params & params, const std::string & value) { + params.speculative.draft.mparams.hf_repo = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_HF_REPO")); + add_opt(common_arg( + {"--spec-draft-threads", "-td", "--threads-draft"}, "N", + "number of threads to use during generation (default: same as --threads)", + [](common_params & params, int value) { + params.speculative.draft.cpuparams.n_threads = value; + if (params.speculative.draft.cpuparams.n_threads <= 0) { + params.speculative.draft.cpuparams.n_threads = std::thread::hardware_concurrency(); + } + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-draft-threads-batch", "-tbd", "--threads-batch-draft"}, "N", + "number of threads to use during batch and prompt processing (default: same as --threads-draft)", + [](common_params & params, int value) { + params.speculative.draft.cpuparams_batch.n_threads = value; + if (params.speculative.draft.cpuparams_batch.n_threads <= 0) { + params.speculative.draft.cpuparams_batch.n_threads = std::thread::hardware_concurrency(); + } + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-draft-cpu-mask", "-Cd", "--cpu-mask-draft"}, "M", + "Draft model CPU affinity mask. Complements cpu-range-draft (default: same as --cpu-mask)", + [](common_params & params, const std::string & mask) { + params.speculative.draft.cpuparams.mask_valid = true; + if (!parse_cpu_mask(mask, params.speculative.draft.cpuparams.cpumask)) { + throw std::invalid_argument("invalid cpumask"); + } + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-draft-cpu-range", "-Crd", "--cpu-range-draft"}, "lo-hi", + "Ranges of CPUs for affinity. Complements --cpu-mask-draft", + [](common_params & params, const std::string & range) { + params.speculative.draft.cpuparams.mask_valid = true; + if (!parse_cpu_range(range, params.speculative.draft.cpuparams.cpumask)) { + throw std::invalid_argument("invalid range"); + } + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-draft-cpu-strict", "--cpu-strict-draft"}, "<0|1>", + "Use strict CPU placement for draft model (default: same as --cpu-strict)", + [](common_params & params, int value) { + params.speculative.draft.cpuparams.strict_cpu = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-draft-prio", "--prio-draft"}, "N", + string_format("set draft process/thread priority : 0-normal, 1-medium, 2-high, 3-realtime (default: %d)\n", params.speculative.draft.cpuparams.priority), + [](common_params & params, int prio) { + if (prio < 0 || prio > 3) { + throw std::invalid_argument("invalid value"); + } + params.speculative.draft.cpuparams.priority = (enum ggml_sched_priority) prio; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-draft-poll", "--poll-draft"}, "<0|1>", + "Use polling to wait for draft model work (default: same as --poll])", + [](common_params & params, int value) { + params.speculative.draft.cpuparams.poll = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-draft-cpu-mask-batch", "-Cbd", "--cpu-mask-batch-draft"}, "M", + "Draft model CPU affinity mask. Complements cpu-range-draft (default: same as --cpu-mask)", + [](common_params & params, const std::string & mask) { + params.speculative.draft.cpuparams_batch.mask_valid = true; + if (!parse_cpu_mask(mask, params.speculative.draft.cpuparams_batch.cpumask)) { + throw std::invalid_argument("invalid cpumask"); + } + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-draft-cpu-range-batch", "-Crbd", "--cpu-range-batch-draft"}, "lo-hi", + "Ranges of CPUs for affinity. Complements --cpu-mask-draft-batch)", + [](common_params & params, const std::string & range) { + params.speculative.draft.cpuparams_batch.mask_valid = true; + if (!parse_cpu_range(range, params.speculative.draft.cpuparams_batch.cpumask)) { + throw std::invalid_argument("invalid cpumask"); + } + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE})); + add_opt(common_arg( + {"--spec-draft-cpu-strict-batch", "--cpu-strict-batch-draft"}, "<0|1>", + "Use strict CPU placement for draft model (default: --cpu-strict-draft)", + [](common_params & params, int value) { + params.speculative.draft.cpuparams_batch.strict_cpu = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-draft-prio-batch", "--prio-batch-draft"}, "N", + string_format("set draft process/thread priority : 0-normal, 1-medium, 2-high, 3-realtime (default: %d)\n", params.speculative.draft.cpuparams_batch.priority), + [](common_params & params, int prio) { + if (prio < 0 || prio > 3) { + throw std::invalid_argument("invalid value"); + } + params.speculative.draft.cpuparams_batch.priority = (enum ggml_sched_priority) prio; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-draft-poll-batch", "--poll-batch-draft"}, "<0|1>", + "Use polling to wait for draft model work (default: --poll-draft)", + [](common_params & params, int value) { + params.speculative.draft.cpuparams_batch.poll = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-draft-type-k", "-ctkd", "--cache-type-k-draft"}, "TYPE", + string_format( + "KV cache data type for K for the draft model\n" + "allowed values: %s\n" + "(default: %s)", + llama_common_get_all_kv_cache_types_rust(), + ggml_type_name(params.speculative.draft.cache_type_k) + ), + [](common_params & params, const std::string & value) { + params.speculative.draft.cache_type_k = kv_cache_type_from_str(value); + } + ).set_env("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K")); + add_opt(common_arg( + {"--spec-draft-type-v", "-ctvd", "--cache-type-v-draft"}, "TYPE", + string_format( + "KV cache data type for V for the draft model\n" + "allowed values: %s\n" + "(default: %s)", + llama_common_get_all_kv_cache_types_rust(), + ggml_type_name(params.speculative.draft.cache_type_v) + ), + [](common_params & params, const std::string & value) { + params.speculative.draft.cache_type_v = kv_cache_type_from_str(value); + } + ).set_env("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V")); + add_opt(common_arg( + {"--spec-draft-override-tensor", "-otd", "--override-tensor-draft"}, "=,...", + "override tensor buffer type for draft model", [](common_params & params, const std::string & value) { + parse_tensor_buffer_overrides(value, params.speculative.draft.tensor_buft_overrides); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-draft-cpu-moe", "-cmoed", "--cpu-moe-draft"}, + "keep all Mixture of Experts (MoE) weights in the CPU for the draft model", + [](common_params & params) { + params.speculative.draft.tensor_buft_overrides.push_back(llm_ffn_exps_cpu_override()); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_CPU_MOE")); + add_opt(common_arg( + {"--spec-draft-n-cpu-moe", "--spec-draft-ncmoe", "-ncmoed", "--n-cpu-moe-draft"}, "N", + "keep the Mixture of Experts (MoE) weights of the first N layers in the CPU for the draft model", + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + for (int i = 0; i < value; ++i) { + static std::list buft_overrides_draft; + buft_overrides_draft.push_back(llm_ffn_exps_block_regex(i)); + params.speculative.draft.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type()}); + } + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE")); + + add_opt(common_arg( + {"--spec-draft-n-max"}, "N", + string_format("number of tokens to draft for speculative decoding (default: %d)", params.speculative.draft.n_max), + [](common_params & params, int value) { + params.speculative.draft.n_max = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MAX")); + add_opt(common_arg( + {"--spec-draft-n-min"}, "N", + string_format("minimum number of draft tokens to use for speculative decoding (default: %d)", params.speculative.draft.n_min), + [](common_params & params, int value) { + params.speculative.draft.n_min = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MIN")); + + add_opt(common_arg( + {"--spec-draft-p-split", "--draft-p-split"}, "P", + string_format("speculative decoding split probability (default: %.2f)", (double)params.speculative.draft.p_split), + [](common_params & params, const std::string & value) { + params.speculative.draft.p_split = std::stof(value); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_P_SPLIT")); + add_opt(common_arg( + {"--spec-draft-p-min", "--draft-p-min"}, "P", + string_format("minimum speculative decoding probability (greedy) (default: %.2f)", (double)params.speculative.draft.p_min), + [](common_params & params, const std::string & value) { + params.speculative.draft.p_min = std::stof(value); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_P_MIN")); + add_opt(common_arg( + {"--spec-draft-ctx-size", "-cd", "--ctx-size-draft"}, "N", + string_format("size of the prompt context for the draft model (default: %d, 0 = loaded from model)", params.speculative.draft.n_ctx), + [](common_params & params, int value) { + params.speculative.draft.n_ctx = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_CTX_SIZE")); + add_opt(common_arg( + {"--spec-draft-device", "-devd", "--device-draft"}, "", + "comma-separated list of devices to use for offloading the draft model (none = don't offload)\n" + "use --list-devices to see a list of available devices", + [](common_params & params, const std::string & value) { + params.speculative.draft.devices = parse_device_list(value); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + GGML_ASSERT(params.speculative.draft.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0 + add_opt(common_arg( + {"--spec-draft-ngl", "-ngld", "--gpu-layers-draft", "--n-gpu-layers-draft"}, "N", + string_format("max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: %s)", + params.speculative.draft.n_gpu_layers == -1 ? "auto" : "all"), + [](common_params & params, const std::string & value) { + if (value == "auto") { + params.speculative.draft.n_gpu_layers = -1; + } else if (value == "all") { + params.speculative.draft.n_gpu_layers = -2; + } else { + params.speculative.draft.n_gpu_layers = std::stoi(value); + } + if (!llama_supports_gpu_offload()) { + fprintf(stderr, "warning: no usable GPU found, --gpu-layers-draft option will be ignored\n"); + fprintf(stderr, "warning: one possible reason is that llama.cpp was compiled without GPU support\n"); + fprintf(stderr, "warning: consult docs/build.md for compilation instructions\n"); + } + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_N_GPU_LAYERS_DRAFT")); + add_opt(common_arg( + {"--spec-draft-model", "-md", "--model-draft"}, "FNAME", + "draft model for speculative decoding (default: unused)", + [](common_params & params, const std::string & value) { + params.speculative.draft.mparams.path = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_MODEL")); + add_opt(common_arg( + {"--spec-draft-replace", "--spec-replace"}, "TARGET", "DRAFT", + "translate the string in TARGET into DRAFT if the draft model and main model are not compatible", + [](common_params & params, const std::string & tgt, const std::string & dft) { + params.speculative.draft.replacements.push_back({ tgt, dft }); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-type"}, "[none|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod]", + string_format("type of speculative decoding to use when no draft model is provided (default: %s)\n", + common_speculative_type_to_str(params.speculative.type).c_str()), + [](common_params & params, const std::string & value) { + if (value == "none") { + params.speculative.type = COMMON_SPECULATIVE_TYPE_NONE; + } else if (value == "ngram-cache") { + params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_CACHE; + } else if (value == "ngram-simple") { + params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE; + } else if (value == "ngram-map-k") { + params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K; + } else if (value == "ngram-map-k4v") { + params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V; + } else if (value == "ngram-mod") { + params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_MOD; + } else { + throw std::invalid_argument("unknown speculative decoding type without draft model"); + } + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_TYPE")); + add_opt(common_arg( + {"--spec-ngram-mod-n-min"}, "N", + string_format("minimum number of ngram tokens to use for ngram-based speculative decoding (default: %d)", params.speculative.ngram_mod.n_min), + [](common_params & params, int value) { + if (value < 0 || value > 1024) { + throw std::invalid_argument("ngram n-min must be between 0 and 1024 inclusive"); + } + params.speculative.ngram_mod.n_min = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-ngram-mod-n-max"}, "N", + string_format("maximum number of ngram tokens to use for ngram-based speculative decoding (default: %d)", params.speculative.ngram_mod.n_max), + [](common_params & params, int value) { + if (value < 0 || value > 1024) { + throw std::invalid_argument("ngram n-max must be between 0 and 1024 inclusive"); + } + params.speculative.ngram_mod.n_max = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-ngram-mod-n-match"}, "N", + string_format("ngram-mod lookup length (default: %d)", params.speculative.ngram_mod.n_match), + [](common_params & params, int value) { + if (value < 1 || value > 1024) { + throw std::invalid_argument("ngram size N must be between 1 and 1024 inclusive"); + } + params.speculative.ngram_mod.n_match = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + + add_opt(common_arg( + {"--spec-ngram-simple-size-n"}, "N", + string_format("ngram size N for ngram-simple speculative decoding, length of lookup n-gram (default: %d)", params.speculative.ngram_simple.size_n), + [](common_params & params, int value) { + if (value < 1 || value > 1024) { + throw std::invalid_argument("ngram size N must be between 1 and 1024 inclusive"); + } + params.speculative.ngram_simple.size_n = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-ngram-simple-size-m"}, "N", + string_format("ngram size M for ngram-simple speculative decoding, length of draft m-gram (default: %d)", params.speculative.ngram_simple.size_m), + [](common_params & params, int value) { + if (value < 1 || value > 1024) { + throw std::invalid_argument("ngram size M must be between 1 and 1024 inclusive"); + } + params.speculative.ngram_simple.size_m = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-ngram-simple-min-hits"}, "N", + string_format("minimum hits for ngram-simple speculative decoding (default: %d)", params.speculative.ngram_simple.min_hits), + [](common_params & params, int value) { + if (value < 1) { + throw std::invalid_argument("ngram min hits must be at least 1"); + } + params.speculative.ngram_simple.min_hits = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + + add_opt(common_arg( + {"--spec-ngram-map-k-size-n"}, "N", + string_format("ngram size N for ngram-map-k speculative decoding, length of lookup n-gram (default: %d)", params.speculative.ngram_map_k.size_n), + [](common_params & params, int value) { + if (value < 1 || value > 1024) { + throw std::invalid_argument("ngram size N must be between 1 and 1024 inclusive"); + } + params.speculative.ngram_map_k.size_n = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-ngram-map-k-size-m"}, "N", + string_format("ngram size M for ngram-map-k speculative decoding, length of draft m-gram (default: %d)", params.speculative.ngram_map_k.size_m), + [](common_params & params, int value) { + if (value < 1 || value > 1024) { + throw std::invalid_argument("ngram size M must be between 1 and 1024 inclusive"); + } + params.speculative.ngram_map_k.size_m = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-ngram-map-k-min-hits"}, "N", + string_format("minimum hits for ngram-map-k speculative decoding (default: %d)", params.speculative.ngram_map_k.min_hits), + [](common_params & params, int value) { + if (value < 1) { + throw std::invalid_argument("ngram min hits must be at least 1"); + } + params.speculative.ngram_map_k.min_hits = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + + add_opt(common_arg( + {"--spec-ngram-map-k4v-size-n"}, "N", + string_format("ngram size N for ngram-map-k4v speculative decoding, length of lookup n-gram (default: %d)", params.speculative.ngram_map_k4v.size_n), + [](common_params & params, int value) { + if (value < 1 || value > 1024) { + throw std::invalid_argument("ngram size N must be between 1 and 1024 inclusive"); + } + params.speculative.ngram_map_k4v.size_n = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-ngram-map-k4v-size-m"}, "N", + string_format("ngram size M for ngram-map-k4v speculative decoding, length of draft m-gram (default: %d)", params.speculative.ngram_map_k4v.size_m), + [](common_params & params, int value) { + if (value < 1 || value > 1024) { + throw std::invalid_argument("ngram size M must be between 1 and 1024 inclusive"); + } + params.speculative.ngram_map_k4v.size_m = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-ngram-map-k4v-min-hits"}, "N", + string_format("minimum hits for ngram-map-k4v speculative decoding (default: %d)", params.speculative.ngram_map_k4v.min_hits), + [](common_params & params, int value) { + if (value < 1) { + throw std::invalid_argument("ngram min hits must be at least 1"); + } + params.speculative.ngram_map_k4v.min_hits = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + + // + // removed params + // + + add_opt(common_arg( + {"--draft", "--draft-n", "--draft-max"}, "N", + "the argument has been removed. use --spec-draft-n-max or --spec-ngram-mod-n-max", + [](common_params & /*params*/, int /*value*/) { + throw std::invalid_argument("the argument has been removed. use --spec-draft-n-max or --spec-ngram-mod-n-max"); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_DRAFT_MAX")); + add_opt(common_arg( + {"--draft-min", "--draft-n-min"}, "N", + "the argument has been removed. use --spec-draft-n-min or --spec-ngram-mod-n-min", + [](common_params & /*params*/, int /*value*/) { + throw std::invalid_argument("the argument has been removed. use --spec-draft-n-min or --spec-ngram-mod-n-min"); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_DRAFT_MIN")); + add_opt(common_arg( + {"--spec-ngram-size-n"}, "N", + "the argument has been removed. use the respective --spec-ngram-*-size-n or --spec-ngram-mod-n-match", + [](common_params & /*params*/, int /*value*/) { + throw std::invalid_argument("the argument has been removed. use the respective --spec-ngram-*-size-n"); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--spec-ngram-size-m"}, "N", + "the argument has been removed. use the respective --spec-ngram-*-size-m", + [](common_params & /*params*/, int /*value*/) { + throw std::invalid_argument("the argument has been removed. use the respective --spec-ngram-*-size-m"); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--spec-ngram-min-hits"}, "N", + "the argument has been removed. use the respective --spec-ngram-*-min-hits", + [](common_params & /*params*/, int /*value*/) { + throw std::invalid_argument("the argument has been removed. use the respective --spec-ngram-*-min-hits"); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER})); + + // + // TTS params + // + + add_opt(common_arg( + {"-mv", "--model-vocoder"}, "FNAME", + "vocoder model for audio generation (default: unused)", + [](common_params & params, const std::string & value) { + params.vocoder.model.path = value; + } + ).set_examples({LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--tts-use-guide-tokens"}, + "Use guide tokens to improve TTS word recall", + [](common_params & params) { + params.vocoder.use_guide_tokens = true; + } + ).set_examples({LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--tts-speaker-file"}, "FNAME", + "speaker file path for audio generation", + [](common_params & params, const std::string & value) { + params.vocoder.speaker_file = value; + } + ).set_examples({LLAMA_EXAMPLE_TTS})); + + // + // diffusion params + // + + add_opt(common_arg( + {"--diffusion-steps"}, "N", + string_format("number of diffusion steps (default: %d)", params.diffusion.steps), + [](common_params & params, int value) { params.diffusion.steps = value; } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-visual"}, + string_format("enable visual diffusion mode (show progressive generation) (default: %s)", params.diffusion.visual_mode ? "true" : "false"), + [](common_params & params) { params.diffusion.visual_mode = true; } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-eps"}, "F", + string_format("epsilon for timesteps (default: %.6f)", (double) params.diffusion.eps), + [](common_params & params, const std::string & value) { params.diffusion.eps = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-algorithm"}, "N", + string_format("diffusion algorithm: 0=ORIGIN, 1=ENTROPY_BASED, 2=MARGIN_BASED, 3=RANDOM, 4=LOW_CONFIDENCE (default: %d)", params.diffusion.algorithm), + [](common_params & params, int value) { params.diffusion.algorithm = value; } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-alg-temp"}, "F", + string_format("dream algorithm temperature (default: %.3f)", (double) params.diffusion.alg_temp), + [](common_params & params, const std::string & value) { params.diffusion.alg_temp = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-block-length"}, "N", + string_format("llada block length for generation (default: %d)", params.diffusion.block_length), + [](common_params & params, int value) { params.diffusion.block_length = value; } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-cfg-scale"}, "F", + string_format("llada classifier-free guidance scale (default: %.3f)", (double) params.diffusion.cfg_scale), + [](common_params & params, const std::string & value) { params.diffusion.cfg_scale = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + {"--diffusion-add-gumbel-noise"}, "F", + string_format("add gumbel noise to the logits if temp > 0.0 (default: %s)", params.diffusion.add_gumbel_noise ? "true" : "false"), + [](common_params & params, const std::string & value) { params.diffusion.add_gumbel_noise = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_DIFFUSION })); + add_opt(common_arg( + { "-lr", "--learning-rate" }, "ALPHA", + string_format("adamw or sgd optimizer alpha (default: %.2g); note: sgd alpha recommended ~10x (no momentum)", (double) params.lr.lr0), + [](common_params & params, const std::string & value) { params.lr.lr0 = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); + add_opt(common_arg({ "-lr-min", "--learning-rate-min" }, "ALPHA", + string_format("(if >0) final learning rate after decay (if -decay-epochs is set, default=%.2g)", + (double) params.lr.lr_min), + [](common_params & params, const std::string & value) { params.lr.lr_min = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); + add_opt(common_arg( + {"-decay-epochs", "--learning-rate-decay-epochs"}, "ALPHA", + string_format("(if >0) decay learning rate to -lr-min after this many epochs (exponential decay, default=%.2g)", (double) params.lr.decay_epochs), + [](common_params & params, const std::string & value) { params.lr.decay_epochs = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); + add_opt(common_arg( + {"-wd", "--weight-decay"}, "WD", + string_format("adamw or sgd optimizer weight decay (0 is off; recommend very small e.g. 1e-9) (default: %.2g).", (double) params.lr.wd), + [](common_params & params, const std::string & value) { params.lr.wd = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); + add_opt(common_arg( + {"-val-split", "--val-split"}, "FRACTION", + string_format("fraction of data to use as validation set for training (default: %.2g).", (double) params.val_split), + [](common_params & params, const std::string & value) { params.val_split = std::stof(value); } + ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); + add_opt(common_arg( + {"-epochs", "--epochs"}, "N", + string_format("optimizer max # of epochs (default: %d)", params.lr.epochs), + [](common_params & params, int epochs) { params.lr.epochs = epochs; } + ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); + add_opt(common_arg( + {"-opt", "--optimizer"}, "sgd|adamw", "adamw or sgd", + [](common_params & params, const std::string & name) { + params.optimizer = common_opt_get_optimizer(name.c_str()); + if (params.optimizer == GGML_OPT_OPTIMIZER_TYPE_COUNT) { + throw std::invalid_argument("invalid --optimizer, valid options: adamw, sgd"); + } + } + ).set_examples({ LLAMA_EXAMPLE_FINETUNE })); + add_opt(common_arg( + {"--check"}, + string_format("check rather than generate results (default: %s)", params.check ? "true" : "false"), + [](common_params & params) { + params.check = true; + } + ).set_examples({LLAMA_EXAMPLE_RESULTS})); + add_opt(common_arg( + {"--save-logits"}, + string_format("save final logits to files for verification (default: %s)", params.save_logits ? "true" : "false"), + [](common_params & params) { + params.save_logits = true; + } + ).set_examples({LLAMA_EXAMPLE_DEBUG})); + add_opt(common_arg( + {"--logits-output-dir"}, "PATH", + string_format("directory for saving logits output files (default: %s)", params.logits_output_dir.c_str()), + [](common_params & params, const std::string & value) { + params.logits_output_dir = value; + } + ).set_examples({LLAMA_EXAMPLE_DEBUG})); + add_opt(common_arg( + {"--tensor-filter"}, "REGEX", + "filter tensor names for debug output (regex pattern, can be specified multiple times)", + [](common_params & params, const std::string & value) { + params.tensor_filter.push_back(value); + } + ).set_examples({LLAMA_EXAMPLE_DEBUG})); + + // presets + add_opt(common_arg( + {"--tts-oute-default"}, + string_format("use default OuteTTS models (note: can download weights from the internet)"), + [](common_params & params) { + params.model.hf_repo = "OuteAI/OuteTTS-0.2-500M-GGUF"; + params.model.hf_file = "OuteTTS-0.2-500M-Q8_0.gguf"; + params.vocoder.model.hf_repo = "ggml-org/WavTokenizer"; + params.vocoder.model.hf_file = "WavTokenizer-Large-75-F16.gguf"; + } + ).set_examples({LLAMA_EXAMPLE_TTS})); + + add_opt(common_arg( + {"--embd-gemma-default"}, + string_format("use default EmbeddingGemma model (note: can download weights from the internet)"), + [](common_params & params) { + params.model.hf_repo = "ggml-org/embeddinggemma-300M-qat-q4_0-GGUF"; + params.model.hf_file = "embeddinggemma-300M-qat-Q4_0.gguf"; + params.port = 8011; + params.n_ubatch = 2048; + params.n_batch = 2048; + params.n_parallel = 32; + params.n_ctx = 2048*params.n_parallel; + params.verbose_prompt = true; + params.embedding = true; + } + ).set_examples({LLAMA_EXAMPLE_EMBEDDING, LLAMA_EXAMPLE_SERVER})); + + add_opt(common_arg( + {"--fim-qwen-1.5b-default"}, + string_format("use default Qwen 2.5 Coder 1.5B (note: can download weights from the internet)"), + [](common_params & params) { + params.model.hf_repo = "ggml-org/Qwen2.5-Coder-1.5B-Q8_0-GGUF"; + params.model.hf_file = "qwen2.5-coder-1.5b-q8_0.gguf"; + params.port = 8012; + params.n_ubatch = 1024; + params.n_batch = 1024; + params.n_ctx = 0; + params.n_cache_reuse = 256; + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + + add_opt(common_arg( + {"--fim-qwen-3b-default"}, + string_format("use default Qwen 2.5 Coder 3B (note: can download weights from the internet)"), + [](common_params & params) { + params.model.hf_repo = "ggml-org/Qwen2.5-Coder-3B-Q8_0-GGUF"; + params.model.hf_file = "qwen2.5-coder-3b-q8_0.gguf"; + params.port = 8012; + params.n_ubatch = 1024; + params.n_batch = 1024; + params.n_ctx = 0; + params.n_cache_reuse = 256; + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + + add_opt(common_arg( + {"--fim-qwen-7b-default"}, + string_format("use default Qwen 2.5 Coder 7B (note: can download weights from the internet)"), + [](common_params & params) { + params.model.hf_repo = "ggml-org/Qwen2.5-Coder-7B-Q8_0-GGUF"; + params.model.hf_file = "qwen2.5-coder-7b-q8_0.gguf"; + params.port = 8012; + params.n_ubatch = 1024; + params.n_batch = 1024; + params.n_ctx = 0; + params.n_cache_reuse = 256; + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + + add_opt(common_arg( + {"--fim-qwen-7b-spec"}, + string_format("use Qwen 2.5 Coder 7B + 0.5B draft for speculative decoding (note: can download weights from the internet)"), + [](common_params & params) { + params.model.hf_repo = "ggml-org/Qwen2.5-Coder-7B-Q8_0-GGUF"; + params.model.hf_file = "qwen2.5-coder-7b-q8_0.gguf"; + params.speculative.draft.mparams.hf_repo = "ggml-org/Qwen2.5-Coder-0.5B-Q8_0-GGUF"; + params.speculative.draft.mparams.hf_file = "qwen2.5-coder-0.5b-q8_0.gguf"; + params.port = 8012; + params.n_ubatch = 1024; + params.n_batch = 1024; + params.n_ctx = 0; + params.n_cache_reuse = 256; + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + + add_opt(common_arg( + {"--fim-qwen-14b-spec"}, + string_format("use Qwen 2.5 Coder 14B + 0.5B draft for speculative decoding (note: can download weights from the internet)"), + [](common_params & params) { + params.model.hf_repo = "ggml-org/Qwen2.5-Coder-14B-Q8_0-GGUF"; + params.model.hf_file = "qwen2.5-coder-14b-q8_0.gguf"; + params.speculative.draft.mparams.hf_repo = "ggml-org/Qwen2.5-Coder-0.5B-Q8_0-GGUF"; + params.speculative.draft.mparams.hf_file = "qwen2.5-coder-0.5b-q8_0.gguf"; + params.port = 8012; + params.n_ubatch = 1024; + params.n_batch = 1024; + params.n_ctx = 0; + params.n_cache_reuse = 256; + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + + add_opt(common_arg( + {"--fim-qwen-30b-default"}, + string_format("use default Qwen 3 Coder 30B A3B Instruct (note: can download weights from the internet)"), + [](common_params & params) { + params.model.hf_repo = "ggml-org/Qwen3-Coder-30B-A3B-Instruct-Q8_0-GGUF"; + params.model.hf_file = "qwen3-coder-30b-a3b-instruct-q8_0.gguf"; + params.port = 8012; + params.n_ubatch = 1024; + params.n_batch = 1024; + params.n_ctx = 0; + params.n_cache_reuse = 256; + } + ).set_examples({LLAMA_EXAMPLE_SERVER})); + + add_opt(common_arg( + {"--gpt-oss-20b-default"}, + string_format("use gpt-oss-20b (note: can download weights from the internet)"), + [](common_params & params) { + params.model.hf_repo = "ggml-org/gpt-oss-20b-GGUF"; + params.model.hf_file = "gpt-oss-20b-mxfp4.gguf"; + params.port = 8013; + params.n_ubatch = 2048; + params.n_batch = 32768; + params.n_parallel = 2; + params.n_ctx = 131072*params.n_parallel; + params.sampling.temp = 1.0f; + params.sampling.top_p = 1.0f; + params.sampling.top_k = 0; + params.sampling.min_p = 0.01f; + params.use_jinja = true; + //params.default_template_kwargs["reasoning_effort"] = "\"high\""; + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + + add_opt(common_arg( + {"--gpt-oss-120b-default"}, + string_format("use gpt-oss-120b (note: can download weights from the internet)"), + [](common_params & params) { + params.model.hf_repo = "ggml-org/gpt-oss-120b-GGUF"; + params.port = 8013; + params.n_ubatch = 2048; + params.n_batch = 32768; + params.n_parallel = 2; + params.n_ctx = 131072*params.n_parallel; + params.sampling.temp = 1.0f; + params.sampling.top_p = 1.0f; + params.sampling.top_k = 0; + params.sampling.min_p = 0.01f; + params.use_jinja = true; + //params.default_template_kwargs["reasoning_effort"] = "\"high\""; + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + + add_opt(common_arg( + {"--vision-gemma-4b-default"}, + string_format("use Gemma 3 4B QAT (note: can download weights from the internet)"), + [](common_params & params) { + params.model.hf_repo = "ggml-org/gemma-3-4b-it-qat-GGUF"; + params.port = 8014; + params.n_ctx = 0; + params.use_jinja = true; + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + + add_opt(common_arg( + {"--vision-gemma-12b-default"}, + string_format("use Gemma 3 12B QAT (note: can download weights from the internet)"), + [](common_params & params) { + params.model.hf_repo = "ggml-org/gemma-3-12b-it-qat-GGUF"; + params.port = 8014; + params.n_ctx = 0; + params.use_jinja = true; + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + + add_opt(common_arg( + {"--spec-default"}, + string_format("enable default speculative decoding config"), + [](common_params & params) { + params.speculative.type = COMMON_SPECULATIVE_TYPE_NGRAM_MOD; + params.speculative.ngram_mod.n_match = 24; + params.speculative.ngram_mod.n_min = 48; + params.speculative.ngram_mod.n_max = 64; + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + + return ctx_arg; +} + +void common_params_add_preset_options(std::vector & args) { + // arguments below won't be treated as CLI args, only preset options + args.push_back(common_arg( + {"load-on-startup"}, "NAME", + "in server router mode, autoload this model on startup", + [](common_params &, const std::string &) { /* unused */ } + ).set_env(COMMON_ARG_PRESET_LOAD_ON_STARTUP).set_preset_only()); + + args.push_back(common_arg( + {"stop-timeout"}, "SECONDS", + "in server router mode, force-kill model instance after this many seconds of graceful shutdown", + [](common_params &, int) { /* unused */ } + ).set_env(COMMON_ARG_PRESET_STOP_TIMEOUT).set_preset_only()); + + // args.push_back(common_arg( + // {"pin"}, + // "in server router mode, do not unload this model if models_max is exceeded", + // [](common_params &) { /* unused */ } + // ).set_preset_only()); +} + +#include "common.h.inc" +#include "common.h.inc" +#include "nlohmann/json.hpp.inc" + +#include +#include + +using json = nlohmann::ordered_json; + +// Helper to iterate over tools/functions +static void autoparser_foreach_function(const json & tools, const std::function & fn) { + for (const auto & tool : tools) { + if (!tool.contains("type") || tool.at("type") != "function" || !tool.contains("function")) { + continue; + } + fn(tool); + } +} + +namespace autoparser { + +parser_build_context::parser_build_context(common_chat_peg_builder & p, const generation_params & inputs) : + p(p), + inputs(inputs), + reasoning_parser(p.eps()) {} + +common_chat_params peg_generator::generate_parser(const common_chat_template & tmpl, + const struct generation_params & inputs) { + // Run differential analysis to extract template structure + struct autoparser autoparser; + autoparser.analyze_template(tmpl); + return generate_parser(tmpl, inputs, autoparser); +} + +common_chat_params peg_generator::generate_parser(const common_chat_template & tmpl, + const struct generation_params & inputs, + const autoparser & autoparser) { + // Create the result structure + common_chat_params data; + data.prompt = common_chat_template_direct_apply(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.preserved_tokens = autoparser.preserved_tokens; + + auto parser = autoparser.build_parser(inputs); + data.parser = parser.save(); + + // Build grammar if tools are present + bool has_tools = + autoparser.tools.format.mode != tool_format::NONE && inputs.tools.is_array() && !inputs.tools.empty(); + std::string trigger_marker = !autoparser.tools.format.section_start.empty() ? autoparser.tools.format.section_start : + autoparser.tools.format.per_call_start; + + bool has_response_format = !inputs.json_schema.empty() && inputs.json_schema.is_object(); + bool include_grammar = has_response_format || (has_tools && + ((inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO && !trigger_marker.empty()) || + inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); + + if (include_grammar) { + data.grammar_lazy = !has_response_format && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + autoparser_foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.contains("parameters") ? function.at("parameters") : json::object(); + builder.resolve_refs(schema); + }); + if (has_response_format) { + auto schema = inputs.json_schema; + builder.resolve_refs(schema); + } + parser.build_grammar(builder, data.grammar_lazy); + }); + + // Set grammar triggers based on tool section markers (fall back to per-call markers) + if (data.grammar_lazy) { + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, trigger_marker } + }; + } + } + + return data; +} + +common_peg_arena autoparser::build_parser(const generation_params & inputs) const { + if (!analysis_complete) { + throw std::invalid_argument("Cannot call build_parser on autoparser without performing analysis first, call analyze_template(...)"); + } + return build_chat_peg_parser([&](common_chat_peg_builder & p) { + parser_build_context ctx(p, inputs); + bool extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + + ctx.extracting_reasoning = extract_reasoning && reasoning.mode != reasoning_mode::NONE; + ctx.content = &content; + ctx.reasoning = &reasoning; + + // Build reasoning parser + ctx.reasoning_parser = reasoning.build_parser(ctx); + + auto parser = p.eps(); + + bool has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + bool has_response_format = inputs.json_schema.is_object() && !inputs.json_schema.empty(); + bool pure_content = reasoning.mode == reasoning_mode::NONE; + + if (has_response_format) { + auto response_format = p.rule("response-format", p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema))); + parser = ctx.reasoning_parser + p.space() + p.choice({ + p.literal("```json") + p.space() + response_format + p.space() + p.literal("```"), + response_format + }) + p.end(); + pure_content = false; + } else if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE && jinja_caps.supports_tool_calls) { + parser = tools.build_parser(ctx); + pure_content = false; + } else { + parser = content.build_parser(ctx); + } + return pure_content ? p.prefix(inputs.generation_prompt, reasoning.start) + parser : p.prefix(inputs.generation_prompt, reasoning.start) << parser; + }); +} + +common_peg_parser analyze_reasoning::build_parser(parser_build_context & ctx) const { + auto & p = ctx.p; + + if (!ctx.extracting_reasoning) { + return p.eps(); + } + + if (mode == reasoning_mode::TAG_BASED || mode == reasoning_mode::TOOLS_ONLY) { + if (!end.empty()) { + if (!start.empty()) { + // Standard tag-based: optional(reasoning) + return p.optional(start + p.reasoning(p.until(end)) + end + p.space()); + } + // Delimiter-style (empty start) + return p.optional(p.reasoning(p.until(end)) + end + p.space()); + } + } + + return p.eps(); +} + +common_peg_parser analyze_content::build_parser(parser_build_context & ctx) const { + auto & p = ctx.p; + + if (is_always_wrapped()) { + if (ctx.extracting_reasoning) { + return ctx.reasoning_parser + start + p.content(p.until(end)) + end + p.end(); + } + return p.content(p.until(start)) + start + p.content(p.until(end)) + end + p.end(); + } + return ctx.reasoning_parser + p.content(p.rest()) + p.end(); +} + +common_peg_parser analyze_content::build_optional_wrapped(parser_build_context & ctx) const { + auto & p = ctx.p; + + if (is_always_wrapped()) { + return p.optional(start + p.content(p.until(end)) + end); + } + return p.eps(); +} + +common_peg_parser analyze_tools::build_parser(parser_build_context & ctx) const { + switch (format.mode) { + case tool_format::JSON_NATIVE: + return build_tool_parser_json_native(ctx); + case tool_format::TAG_WITH_JSON: + return build_tool_parser_tag_json(ctx); + case tool_format::TAG_WITH_TAGGED: + return build_tool_parser_tag_tagged(ctx); + default: + LOG_ERR("[ERROR] Template seems to support tool calls, but failed to determine tool format. Tool calling will not work properly. " + "Check for a fixed template for your model in the models/templates directory of your llama.cpp installation or " + "report an issue at https://github.com/ggml-org/llama.cpp/issues\n"); + return ctx.p.eps(); + } +} + +common_peg_parser analyze_tools::build_tool_parser_json_native(parser_build_context & ctx) const { + auto & p = ctx.p; + const auto & inputs = ctx.inputs; + bool force_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; + + // Build effective field names with dot notation if function_field is set + std::string name_field = format.name_field; + std::string args_field = format.args_field; + + if (!format.function_field.empty() && format.function_field != "function" && + name_field.find('.') == std::string::npos) { + name_field = format.function_field + "." + name_field; + args_field = format.function_field + "." + args_field; + } + + auto tools_parser = p.eps(); + if (format.section_start.empty() && !format.per_call_start.empty()) { + auto single_tool_parser = p.standard_json_tools( + format.per_call_start, format.per_call_end, inputs.tools, inputs.parallel_tool_calls, + inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED, name_field, args_field, format.tools_array_wrapped, + format.fun_name_is_key, format.id_field, format.gen_id_field, format.parameter_order); + tools_parser = p.trigger_rule("tool-calls", p.one_or_more(single_tool_parser + p.space())); + } else { + tools_parser = p.standard_json_tools( + format.section_start, format.section_end, inputs.tools, inputs.parallel_tool_calls, + inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED, name_field, args_field, format.tools_array_wrapped, + format.fun_name_is_key, format.id_field, format.gen_id_field, format.parameter_order); + } + + // Handle content wrappers if present + if (ctx.content && ctx.content->is_always_wrapped()) { + auto wrapped_content = ctx.content->build_optional_wrapped(ctx); + return ctx.reasoning_parser + wrapped_content + tools_parser + p.end(); + } + + std::string tool_start = "{"; + if (!format.section_start.empty()) { + tool_start = format.section_start; + } else if (!format.per_call_start.empty()) { + tool_start = format.per_call_start; + } + + return ctx.reasoning_parser + (force_tools ? p.eps() : p.optional(p.content(p.until(tool_start)))) + tools_parser + + p.end(); +} + +common_peg_parser analyze_tools::build_func_parser(common_chat_peg_builder & p, const std::string & name, + const common_peg_parser & call_id_section, bool have_call_id, + const common_peg_parser & args, + std::optional atomic_peek) const { + auto open = p.tool_open(function.name_prefix + p.tool_name(p.literal(name)) + function.name_suffix); + bool matched_atomic = false; + common_peg_parser func_parser = p.eps(); + + if (!function.name_suffix.empty()) { + func_parser = open + call_id_section + p.space() + args; + matched_atomic = true; + } else if (have_call_id) { + func_parser = p.atomic(open + call_id_section) + p.space() + args; + matched_atomic = true; + } else if (atomic_peek.has_value()) { + func_parser = p.atomic(open + call_id_section + p.space() + *atomic_peek) + args; + matched_atomic = true; + } else { + func_parser = open + call_id_section + p.space() + args; + } + + if (!function.close.empty()) { + func_parser = func_parser + p.space() + p.tool_close(p.literal(function.close)); + } else if (!format.per_call_end.empty()) { + // When there's no func_close but there is a per_call_end marker, use peek() to ensure + // we only emit tool_close when we can actually see the closing marker. This prevents + // premature closing during partial parsing when we've seen e.g. "" (end) or "" prefix that failed to match. + func_parser = func_parser + p.tool_close(p.peek(p.literal(format.per_call_end))); + } else { + func_parser = func_parser + p.tool_close(p.space()); // force this to process tool closing callbacks in mapper + } + if (!matched_atomic) { + func_parser = p.atomic(func_parser); + } + return func_parser; +} + +common_peg_parser analyze_tools::build_tool_parser_tag_json(parser_build_context & ctx) const { + auto & p = ctx.p; + const auto & inputs = ctx.inputs; + bool force_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; + + common_peg_parser tool_choice = p.choice(); + + autoparser_foreach_function(inputs.tools, [&](const json & tool) { + const auto & func = tool.at("function"); + std::string name = func.at("name"); + const auto & schema = func.contains("parameters") ? func.at("parameters") : json::object(); + + // Build call_id parser based on position (if supported) + bool have_call_id = false; + common_peg_parser call_id_section = p.eps(); + if (call_id.pos == call_id_position::BETWEEN_FUNC_AND_ARGS && !call_id.prefix.empty() && + (!call_id.suffix.empty() || !arguments.start.empty())) { + if (!call_id.suffix.empty()) { + call_id_section = p.optional(call_id.prefix + p.tool_id(p.until(call_id.suffix))) + call_id.suffix; + } else { + call_id_section = p.optional(call_id.prefix + p.tool_id(p.until(arguments.start))); + } + have_call_id = true; + } + auto args_parser = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)); + if (!arguments.start.empty()) { + args_parser = p.literal(arguments.start) + args_parser; + } + if (!arguments.end.empty()) { + args_parser = args_parser + p.literal(arguments.end); + } + + auto atomic_peek = !arguments.start.empty() ? std::optional(p.peek(p.literal(arguments.start))) : std::nullopt; + auto func_parser = build_func_parser(p, name, call_id_section, have_call_id, args_parser, atomic_peek); + tool_choice |= p.rule("tool-" + name, func_parser); + }); + + auto require_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; + + common_peg_parser tool_calls = p.eps(); + + if (!format.per_call_start.empty()) { + auto wrapped_call = format.per_call_start + tool_choice + format.per_call_end; + if (inputs.parallel_tool_calls) { + tool_calls = p.trigger_rule("tool-call", wrapped_call + p.zero_or_more(p.space() + wrapped_call)); + } else { + tool_calls = p.trigger_rule("tool-call", wrapped_call); + } + if (!format.section_start.empty()) { + tool_calls = p.trigger_rule("tool-calls", + p.literal(format.section_start) + p.space() + tool_calls + p.space() + + (format.section_end.empty() ? p.end() : p.literal(format.section_end))); + } + } else { + std::string separator = ", "; // Default + if (inputs.parallel_tool_calls) { + tool_calls = p.trigger_rule("tool-call", format.section_start + tool_choice + + p.zero_or_more(separator + tool_choice) + format.section_end); + } else { + tool_calls = p.trigger_rule("tool-call", format.section_start + tool_choice + format.section_end); + } + } + + if (!require_calls) { + tool_calls = p.optional(tool_calls); + } + + std::string trigger_marker = !format.section_start.empty() ? format.section_start : format.per_call_start; + auto content_before_tools = trigger_marker.empty() ? p.eps() : p.until(trigger_marker); + return ctx.reasoning_parser + (force_tools ? p.eps() : p.optional(p.content(content_before_tools))) + tool_calls + + p.end(); +} + +common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_context & ctx) const { + auto & p = ctx.p; + const auto & inputs = ctx.inputs; + bool force_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; + + auto until_suffix = p.rule("until-suffix", p.until(arguments.value_suffix)); + + common_peg_parser tool_choice = p.choice(); + + autoparser_foreach_function(inputs.tools, [&](const json & tool) { + const auto & func = tool.at("function"); + std::string name = func.at("name"); + auto params = func.contains("parameters") ? func.at("parameters") : json::object(); + const auto & properties = params.contains("properties") ? params.at("properties") : json::object(); + + std::set required; + if (params.contains("required")) { + params.at("required").get_to(required); + } + + auto schema_info = common_schema_info(); + schema_info.resolve_refs(params); + + // Build parser for each argument, separating required and optional + std::vector required_parsers; + std::vector optional_parsers; + for (const auto & [param_name, param_schema] : properties.items()) { + bool is_required = required.find(param_name) != required.end(); + + auto arg = + p.tool_arg(p.tool_arg_open(arguments.name_prefix + p.tool_arg_name(p.literal(param_name)) + + arguments.name_suffix) + + arguments.value_prefix + + (schema_info.resolves_to_string(param_schema) ? + p.tool_arg_string_value(p.schema(until_suffix, + "tool-" + name + "-arg-" + param_name + "-schema", + param_schema, true)) : + p.tool_arg_json_value(p.schema( + p.json(), "tool-" + name + "-arg-" + param_name + "-schema", param_schema, false)) + + p.space()) + + p.tool_arg_close(p.literal(arguments.value_suffix))); + + auto named_arg = p.rule("tool-" + name + "-arg-" + param_name, arg); + if (is_required) { + required_parsers.push_back(named_arg); + } else { + optional_parsers.push_back(named_arg); + } + } + + // Build required arg sequence in definition order + common_peg_parser args_seq = p.eps(); + for (size_t i = 0; i < required_parsers.size(); i++) { + if (i > 0) { + args_seq = args_seq + p.space(); + } + args_seq = args_seq + required_parsers[i]; + } + + // Build optional args with flexible ordering + if (!optional_parsers.empty()) { + common_peg_parser any_opt = p.choice(); + for (const auto & opt : optional_parsers) { + any_opt |= opt; + } + args_seq = args_seq + p.repeat(p.space() + any_opt, 0, -1); + } + + if (!arguments.start.empty()) { + args_seq = p.literal(arguments.start) + args_seq; + } + if (!arguments.end.empty()) { + args_seq = args_seq + p.literal(arguments.end); + } + + // Build call_id parser based on position (if supported) + common_peg_parser call_id_section = p.eps(); + bool have_call_id = false; + if (call_id.pos == call_id_position::BETWEEN_FUNC_AND_ARGS && !call_id.prefix.empty() && + (!call_id.suffix.empty() || !arguments.start.empty())) { + have_call_id = true; + if (!call_id.suffix.empty()) { + call_id_section = p.optional(call_id.prefix + p.tool_id(p.until(call_id.suffix)) + call_id.suffix); + } else { + call_id_section = p.optional(call_id.prefix + p.tool_id(p.until(arguments.start))); + } + } + + // Only peek for an arg tag when there are required args that must follow. + // When all args are optional, the model may emit no arg tags at all (#20650). + auto atomic_peek = (!arguments.name_prefix.empty() && !required_parsers.empty()) ? + std::optional(p.peek(p.literal(arguments.name_prefix))) : std::nullopt; + auto func_parser = build_func_parser(p, name, call_id_section, have_call_id, args_seq, atomic_peek); + tool_choice |= p.rule("tool-" + name, func_parser); + }); + + auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; + + common_peg_parser tool_calls = p.eps(); + + if (!format.per_call_start.empty()) { + auto wrapped_call = format.per_call_start + p.space() + tool_choice + p.space() + format.per_call_end; + if (inputs.parallel_tool_calls) { + tool_calls = p.trigger_rule("tool-call", wrapped_call + p.zero_or_more(p.space() + wrapped_call) + p.space()); + } else { + tool_calls = p.trigger_rule("tool-call", wrapped_call + p.space()); + } + if (!format.section_start.empty()) { + tool_calls = p.trigger_rule("tool-calls", + p.literal(format.section_start) + p.space() + tool_calls + p.space() + + (format.section_end.empty() ? p.end() : p.literal(format.section_end) + p.space())); + } + } else { + std::string separator = ", "; // Default + + if (inputs.parallel_tool_calls) { + tool_calls = p.trigger_rule("tool-call", format.section_start + p.space() + tool_choice + + p.zero_or_more(separator + tool_choice) + p.space() + + format.section_end); + } else { + tool_calls = p.trigger_rule( + "tool-call", format.section_start + p.space() + tool_choice + p.space() + format.section_end); + } + } + + if (!require_tools) { + tool_calls = p.optional(tool_calls); + } + + std::string trigger_marker = !format.section_start.empty() ? format.section_start : format.per_call_start; + auto content_before_tools = trigger_marker.empty() ? p.eps() : p.until(trigger_marker); + return ctx.reasoning_parser + (force_tools ? p.eps() : p.optional(p.content(content_before_tools))) + tool_calls + + p.end(); +} + +} // namespace autoparser + +#include "common.h.inc" +#include "nlohmann/json.hpp.inc" + +#include +#include + +using json = nlohmann::ordered_json; + +std::string trim_whitespace(const std::string & str) { + llama_common_unicode_string trimmed = llama_common_string_trim_rust((const uint8_t *) str.data(), str.size(), 0); + if (!trimmed.data) { + return ""; + } + std::string result((const char *) trimmed.data, trimmed.len); + llama_common_unicode_string_free(trimmed); + return result; +} + +std::string trim_leading_whitespace(const std::string & str) { + llama_common_unicode_string trimmed = llama_common_string_trim_rust((const uint8_t *) str.data(), str.size(), 1); + if (!trimmed.data) { + return ""; + } + std::string result((const char *) trimmed.data, trimmed.len); + llama_common_unicode_string_free(trimmed); + return result; +} + +std::string trim_trailing_whitespace(const std::string & str) { + llama_common_unicode_string trimmed = llama_common_string_trim_rust((const uint8_t *) str.data(), str.size(), 2); + if (!trimmed.data) { + return ""; + } + std::string result((const char *) trimmed.data, trimmed.len); + llama_common_unicode_string_free(trimmed); + return result; +} + +std::string trim_trailing_newlines(const std::string & str) { + llama_common_unicode_string trimmed = llama_common_string_trim_rust((const uint8_t *) str.data(), str.size(), 3); + if (!trimmed.data) { + return ""; + } + std::string result((const char *) trimmed.data, trimmed.len); + llama_common_unicode_string_free(trimmed); + return result; +} + +diff_split calculate_diff_split(const std::string & left, const std::string & right) { + diff_split result; + + auto left_seg = segmentize_markers(left); + auto right_seg = segmentize_markers(right); + + if (left_seg.empty()) { + result.right = right; + return result; + } + if (right_seg.empty()) { + result.left = left; + return result; + } + + auto left_start = left_seg.begin(); + auto left_end = --left_seg.end(); + auto right_start = right_seg.begin(); + auto right_end = --right_seg.end(); + + auto test = [&] () { + return left_start != left_end && right_start != right_end; + }; + + bool left_fully_consumed = false; + bool right_fully_consumed = false; + + while (test()) { + bool advanced = false; + if (*left_start == *right_start) { + result.prefix.append(left_start->value); + left_start++; + right_start++; + advanced = true; + } + if (*left_end == *right_end) { + result.suffix = left_end->value + result.suffix; + if (left_start != left_end) { + left_end--; + } else { + left_fully_consumed = true; + } + if (right_start != right_end) { + right_end--; + } else { + right_fully_consumed = true; + } + advanced = true; + } + if (!advanced) { + break; + } + } + + if (left_start == left_end && right_start != right_end) { + if (*left_start == *right_end) { + result.suffix = right_end->value + result.suffix; + right_end--; + left_fully_consumed = true; + } else if (*left_start == *right_start) { + result.prefix.append(right_start->value); + right_start++; + left_fully_consumed = true; + } + } else if (right_start == right_end && left_start != left_end) { + if (*left_end == *right_start) { + result.suffix = left_end->value + result.suffix; + left_end--; + right_fully_consumed = true; + } else if (*left_start == *right_start) { + result.prefix.append(left_start->value); + left_start++; + right_fully_consumed = true; + } + } else if (left_start == left_end && right_start == right_end && *left_start == *right_start && left_start->type == segment_type::MARKER) { + result.prefix.append(right_start->value); + left_fully_consumed = true; + right_fully_consumed = true; + } + + auto eat_segment = [](std::string str, const segment & seg) -> std::string { return std::move(str) + seg.value; }; + + bool can_have_text_suffix = left_end->type == segment_type::TEXT && right_end->type == segment_type::TEXT; + bool can_have_text_prefix = right_start->type == segment_type::TEXT && left_start->type == segment_type::TEXT; + + std::string remainder_left = std::accumulate(left_start, left_fully_consumed ? left_end : ++left_end, std::string(), eat_segment); + std::string remainder_right = std::accumulate(right_start, right_fully_consumed ? right_end : ++right_end, std::string(), eat_segment); + + size_t suffix_len = can_have_text_suffix ? llama_common_suffix_len_rust( + (const uint8_t *) remainder_left.data(), remainder_left.size(), + (const uint8_t *) remainder_right.data(), remainder_right.size()) : 0; + // avoid overlaps between prefix and suffix + std::string prefix_left = remainder_left.substr(0, remainder_left.size() - suffix_len); + std::string prefix_right = remainder_right.substr(0, remainder_right.size() - suffix_len); + size_t prefix_len = can_have_text_prefix ? llama_common_prefix_len_rust( + (const uint8_t *) prefix_left.data(), prefix_left.size(), + (const uint8_t *) prefix_right.data(), prefix_right.size()) : 0; + + result.prefix.append(remainder_left.substr(0, prefix_len)); + result.suffix = remainder_left.substr(remainder_left.length() - suffix_len, suffix_len) + result.suffix; + result.left = remainder_left.substr(prefix_len, remainder_left.length() - prefix_len - suffix_len); + result.right = remainder_right.substr(prefix_len, remainder_right.length() - prefix_len - suffix_len); + + if (result.left == "" && result.right == "") { + // degenerate case, no diff + result.prefix = left; + result.suffix = ""; + // pick prefix = all as representation + } + + // When left has no unique content (result.left is empty), left is entirely + // shared with right. The simultaneous prefix/suffix segment matching can + // incorrectly consume trailing segments of left as suffix when those same + // segments also appear at the end of right (e.g. "\n" at the end of both + // the shared content and the generation prompt). This rotates the diff. + // Fix: if left is a prefix of right, enforce that directly. + if (result.left.empty() && !result.right.empty() && + left.size() <= right.size() && + right.substr(0, left.size()) == left) { + result.prefix = left; + result.suffix = ""; + result.right = right.substr(left.size()); + } + + return result; +} + +// Returns the prefix of `full` up until the first occurrence of the common prefix of `left` and `right` +std::string until_common_prefix(const std::string & full, const std::string & left, const std::string & right) { + llama_common_unicode_string result = llama_common_until_common_prefix_rust( + (const uint8_t *) full.data(), full.size(), + (const uint8_t *) left.data(), left.size(), + (const uint8_t *) right.data(), right.size()); + return common_rust_string_to_std(result); +} + +// Returns the suffix of `full` after the last occurrence of the common suffix of `left` and `right` +std::string after_common_suffix(const std::string & full, const std::string & left, const std::string & right) { + llama_common_unicode_string result = llama_common_after_common_suffix_rust( + (const uint8_t *) full.data(), full.size(), + (const uint8_t *) left.data(), left.size(), + (const uint8_t *) right.data(), right.size()); + return common_rust_string_to_std(result); +} + +// TODO: segmentize will treat a JSON array inside tags as a tag: [{ "fun": { ... } }] will be three markers +// not too worried about that because it hasn't turned out as a problem anywhere, but noting here in case it will +// Might have to put some restrictions on tag contents as well (like "no { }") +std::vector segmentize_markers(const std::string & text) { + std::vector retval; + bool in_marker = false; + char marker_opener = '\0'; + + auto is_marker_opener = [](char c) -> bool { return c == '<' || c == '['; }; + auto is_marker_closer = [](char op, char c) -> bool { return (op == '<' && c == '>') || (op == '[' && c == ']'); }; + + size_t last_border = 0; + + for (size_t cur_pos = 0; cur_pos < text.length(); cur_pos++) { + if (!in_marker && is_marker_opener(text[cur_pos])) { + if (last_border < cur_pos) { + retval.push_back(segment(segment_type::TEXT, text.substr(last_border, cur_pos - last_border))); + } + last_border = cur_pos; + in_marker = true; + marker_opener = text[cur_pos]; + } else if (in_marker && is_marker_closer(marker_opener, text[cur_pos])) { + // no need to check because last_border will always be smaller + retval.push_back(segment(segment_type::MARKER, text.substr(last_border, cur_pos - last_border + 1))); + last_border = cur_pos + 1; + in_marker = false; + marker_opener = '\0'; + } + } + if (last_border < text.length()) { + retval.push_back(segment(segment_type::TEXT, text.substr(last_border))); + } + return retval; +} + +std::vector prune_whitespace_segments(const std::vector & segments) { + std::vector result; + for (const auto & seg : segments) { + if (!trim_whitespace(seg.value).empty()) { + result.push_back(seg); + } + } + return result; +} + +namespace autoparser { + +std::string apply_template(const common_chat_template & tmpl, const template_params & params) { + generation_params tmpl_params; + tmpl_params.messages = params.messages; + tmpl_params.tools = params.tools; + tmpl_params.add_generation_prompt = params.add_generation_prompt; + tmpl_params.enable_thinking = params.enable_thinking; + + if (params.extra_context) { + tmpl_params.extra_context = *params.extra_context; + } + tmpl_params.extra_context["enable_thinking"] = params.enable_thinking; + + try { + return common_chat_template_direct_apply(tmpl, tmpl_params); + } catch (const std::exception & e) { + LOG_DBG("Template application failed: %s\n", e.what()); + return ""; + } +} + +std::optional compare_variants( + const common_chat_template & tmpl, + const template_params & params_A, + const std::function & params_modifier) { + // Create variant B by copying A + template_params params_B = params_A; + + // Apply modifier to create variant B + if (params_modifier) { + params_modifier(params_B); + } + + // Apply template to both variants + std::string output_A = apply_template(tmpl, params_A); + std::string output_B = apply_template(tmpl, params_B); + + // Check for template application failures + if (output_A.empty() || output_B.empty()) { + return std::nullopt; + } + + // Calculate diff and return result with both outputs + compare_variants_result result; + result.diff = calculate_diff_split(output_A, output_B); + result.output_A = output_A; + result.output_B = output_B; + + return result; +} + +} // namespace autoparser + +#include "common.h.inc" +#include "common.h.inc" +#include "nlohmann/json.hpp.inc" + +#include + +#define ANSI_RESET "\033[0m" +#define ANSI_PURPLE "\033[1m\x1b[38;5;126m" +#define ANSI_ORANGE "\033[1m\x1b[38;5;214m" +#define ANSI_RED "\033[1m\x1b[38;5;196m" + +using json = nlohmann::ordered_json; + +namespace autoparser { + +static const std::string FUN_FIRST = "FFF_FIRST_FUN_F"; +static const std::string FUN_SECOND = "SSS_SECOND_FUN_S"; +static const std::string ARG_FIRST = "AA_ARG_FST_AA"; +static const std::string ARG_SECOND = "BB_ARG_SND_BB"; +static const std::string USER_MSG = "U_USER_MSG Hello END_U"; +static const std::string ASSISTANT_MSG = "A_ASST_MSG I can help END_A"; +static const std::string THINKING_CONTENT = "REASON_PART I am thinking END_R"; +static const std::string CALL_ID_001 = "call00001"; +static const std::string CALL_ID_002 = "call00002"; +static const std::string CALL_ID_999 = "call99999"; + +static std::vector> workarounds( + { // Old reasoning Qwen templates - they don't really display reasoning content, but we still want to + // support reasoning on them + [](const common_chat_template & tmpl, autoparser & analysis) -> void { + if (tmpl.src.find("content.split('')") != std::string::npos && + tmpl.src.find("reasoning_content") == std::string::npos && + tmpl.src.find("") == std::string::npos && + analysis.reasoning.mode == reasoning_mode::NONE) { + analysis.reasoning.mode = reasoning_mode::TAG_BASED; + analysis.reasoning.start = ""; + analysis.reasoning.end = ""; + analysis.preserved_tokens.push_back(""); + analysis.preserved_tokens.push_back(""); + LOG_DBG(ANSI_ORANGE "[Patch: old Qwen/Deepseek thinking template]\n" ANSI_RESET); + } + }, + // Granite 3.3, with separate reasoning and content markers + [](const common_chat_template & tmpl, autoparser & analysis) -> void { + if (tmpl.src.find("Write your thoughts between and write your response between " + "") != std::string::npos) { + analysis.reasoning.mode = reasoning_mode::TAG_BASED; + analysis.reasoning.start = ""; + analysis.reasoning.end = ""; + analysis.preserved_tokens.push_back(""); + analysis.preserved_tokens.push_back(""); + analysis.content.mode = content_mode::WRAPPED_WITH_REASONING; + analysis.content.start = ""; + analysis.content.end = ""; + analysis.preserved_tokens.push_back(""); + analysis.preserved_tokens.push_back(""); + LOG_DBG(ANSI_ORANGE "[Patch: Granite 3.3]\n" ANSI_RESET); + } + }, + // Cohere Command R+ - content wrapped in <|CHATBOT_TOKEN|>...<|END_OF_TURN_TOKEN|> + [](const common_chat_template & tmpl, autoparser & analysis) -> void { + if (tmpl.src.find("<|CHATBOT_TOKEN|>") != std::string::npos && + tmpl.src.find("<|END_OF_TURN_TOKEN|>") != std::string::npos && analysis.content.start.empty()) { + analysis.content.mode = content_mode::ALWAYS_WRAPPED; + analysis.content.start = "<|CHATBOT_TOKEN|>"; + analysis.content.end = "<|END_OF_TURN_TOKEN|>"; + analysis.preserved_tokens.push_back("<|CHATBOT_TOKEN|>"); + analysis.preserved_tokens.push_back("<|END_OF_TURN_TOKEN|>"); + LOG_DBG(ANSI_ORANGE "[Patch: Cohere Command R+]\n" ANSI_RESET); + } + }, + // Functionary - no tool call section delimiter + [](const common_chat_template & tmpl, autoparser & analysis) -> void { + if (tmpl.src.find("set has_code_interpreter = tools | selectattr(\"type\", \"equalto\", " + "\"code_interpreter\") | list | length > 0") != std::string::npos) { + analysis.content.mode = content_mode::PLAIN; + analysis.content.end = ""; + analysis.tools.function.name_prefix = ""; + analysis.tools.format.section_start = ""; + analysis.tools.format.section_end = ""; + analysis.tools.format.per_call_start = ""); + analysis.preserved_tokens.push_back("<|eom_id|>"); + analysis.preserved_tokens.push_back(""); + analysis.preserved_tokens.push_back(""); + LOG_DBG(ANSI_ORANGE "[Patch: Functionary 3.1]\n" ANSI_RESET); + } + }, + // DeepSeek-R1-Distill-Qwen + [](const common_chat_template & tmpl, autoparser & analysis) -> void { + if (tmpl.src.find( + "{{'<｜Assistant｜><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tool['type'] + '<｜tool▁sep｜>'") != + std::string::npos) { + analysis.tools.format.section_start = "<｜tool▁calls▁begin｜>"; + analysis.tools.format.section_end = "<｜tool▁calls▁end｜>"; + analysis.tools.format.per_call_start = "<｜tool▁call▁begin｜>function"; + analysis.tools.function.name_prefix = "<｜tool▁sep｜>"; + analysis.tools.format.per_call_end = "<｜tool▁call▁end｜>"; + analysis.tools.function.close = "```"; + LOG_DBG(ANSI_ORANGE "[Patch: DeepSeek-R1-Distill-Qwen]\n" ANSI_RESET); + } + } + }); + +// Common JSON structures +static json params_schema = { + { "type", "object" }, + { "properties", + { { ARG_FIRST, { { "type", "string" }, { "description", "First argument" } } }, + { ARG_SECOND, { { "type", "string" }, { "description", "Second argument" } } } } }, + { "required", json::array({}) } +}; + +static json tools = json::array({ + { { "type", "function" }, + { "function", + json{ { "name", FUN_FIRST }, { "description", "Test function foo" }, { "parameters", params_schema } } } }, + { { "type", "function" }, + { "function", + json{ { "name", FUN_SECOND }, { "description", "Test function bar" }, { "parameters", params_schema } } } } +}); + +static json user_msg = json{ + { "role", "user" }, + { "content", USER_MSG } +}; + +static json build_tool_call(const std::string & name, const json & args, const std::string & id = CALL_ID_001) { + return json{ + { "id", id }, + { "type", "function" }, + { "function", json{ { "name", name }, { "arguments", args } } } + }; +} + +static json first_tool_call_zero_args = build_tool_call(FUN_FIRST, json::object(), CALL_ID_001); +static json first_tool_call_one_arg = build_tool_call(FUN_FIRST, {{ ARG_FIRST, "XXXX" }}, CALL_ID_001); +static json first_tool_call_one_arg_other_val = build_tool_call(FUN_FIRST, {{ ARG_FIRST, "YYYY" }}, CALL_ID_001); +static json first_tool_call_other_arg = build_tool_call(FUN_FIRST, {{ ARG_SECOND, "YYYY" }}, CALL_ID_001); + +static json first_tool_call = + build_tool_call(FUN_FIRST, json{{ ARG_FIRST, "XXXX" }, { ARG_SECOND, "YYYY" }}, CALL_ID_001); +static json second_tool_call = + build_tool_call(FUN_SECOND, json{ { ARG_FIRST, "XXXX" }, { ARG_SECOND, "YYYY" }}, CALL_ID_002); +static json first_tool_call_alt_id = + build_tool_call(FUN_FIRST, json{{ ARG_FIRST, "XXXX" }, { ARG_SECOND, "YYYY" }}, CALL_ID_999); + +template +static std::string mode_to_str(T mode) { + std::ostringstream os; + os << mode; + return os.str(); +} + +void autoparser::analyze_template(const common_chat_template & tmpl) { + jinja_caps = tmpl.original_caps(); + reasoning = analyze_reasoning(tmpl, jinja_caps.supports_tool_calls); + content = analyze_content(tmpl, reasoning); + tools = analyze_tools(jinja_caps.supports_tool_calls ? analyze_tools(tmpl, jinja_caps, reasoning) : analyze_tools()); + collect_preserved_tokens(); + + for (auto & workaround : workarounds) { + workaround(tmpl, *this); + } + + LOG_DBG("\n--- Reasoning & Content Structure ---\n"); + LOG_DBG("reasoning_mode: %s\n", mode_to_str(reasoning.mode).c_str()); + LOG_DBG("reasoning_start: '%s'\n", reasoning.start.c_str()); + LOG_DBG("reasoning_end: '%s'\n", reasoning.end.c_str()); + LOG_DBG("content_mode: %s\n", mode_to_str(content.mode).c_str()); + LOG_DBG("content_start: '%s'\n", content.start.c_str()); + LOG_DBG("content_end: '%s'\n", content.end.c_str()); + + LOG_DBG("\n--- Tool Call Structure ---\n"); + LOG_DBG("tool_mode: %s\n", mode_to_str(tools.format.mode).c_str()); + LOG_DBG("supports_tools: %s\n", jinja_caps.supports_tools ? "true" : "false"); + LOG_DBG("supports_parallel_calls: %s\n", jinja_caps.supports_parallel_tool_calls ? "true" : "false"); + LOG_DBG("tool_section_start: '%s'\n", tools.format.section_start.c_str()); + LOG_DBG("tool_section_end: '%s'\n", tools.format.section_end.c_str()); + LOG_DBG("per_call_start: '%s'\n", tools.format.per_call_start.c_str()); + LOG_DBG("per_call_end: '%s'\n", tools.format.per_call_end.c_str()); + LOG_DBG("func_name_prefix: '%s'\n", tools.function.name_prefix.c_str()); + LOG_DBG("func_name_suffix: '%s'\n", tools.function.name_suffix.c_str()); + LOG_DBG("func_close: '%s'\n", tools.function.close.c_str()); + LOG_DBG("call_id_prefix: '%s'\n", tools.call_id.prefix.c_str()); + LOG_DBG("call_id_suffix: '%s'\n", tools.call_id.suffix.c_str()); + LOG_DBG("call_id_pos: '%s'\n", mode_to_str(tools.call_id.pos).c_str()); + LOG_DBG("args_start: '%s'\n", tools.arguments.start.c_str()); + LOG_DBG("args_end: '%s'\n", tools.arguments.end.c_str()); + LOG_DBG("arg_name_prefix: '%s'\n", tools.arguments.name_prefix.c_str()); + LOG_DBG("arg_name_suffix: '%s'\n", tools.arguments.name_suffix.c_str()); + LOG_DBG("arg_value_prefix: '%s'\n", tools.arguments.value_prefix.c_str()); + LOG_DBG("arg_value_suffix: '%s'\n", tools.arguments.value_suffix.c_str()); + LOG_DBG("name_field: '%s'\n", tools.format.name_field.c_str()); + LOG_DBG("args_field: '%s'\n", tools.format.args_field.c_str()); + LOG_DBG("id_field: '%s'\n", tools.format.id_field.c_str()); + LOG_DBG("gen_id_field: '%s'\n", tools.format.gen_id_field.c_str()); + LOG_DBG("parameter_order: '%s'\n", std::accumulate(tools.format.parameter_order.begin(), tools.format.parameter_order.end(), + std::string(""), [] (const std::string & a, const std::string & b) { return a.empty() ? b : a + ", " + b; } + ).c_str()); + + LOG_DBG(ANSI_PURPLE "=== Differential analysis complete ===\n" ANSI_RESET); + analysis_complete = true; +} + +void autoparser::collect_preserved_tokens() { + auto add_token = [this](const std::string & org_token) { + std::string token = trim_whitespace(org_token); + if (!token.empty()) { + // Avoid duplicates + if (std::find(preserved_tokens.begin(), preserved_tokens.end(), token) == preserved_tokens.end()) { + preserved_tokens.push_back(token); + } + } + }; + + add_token(reasoning.start); + add_token(reasoning.end); + add_token(content.start); + add_token(content.end); + add_token(tools.format.section_start); + add_token(tools.format.section_end); + add_token(tools.format.per_call_start); + add_token(tools.format.per_call_end); + add_token(tools.function.name_prefix); + add_token(tools.function.name_suffix); + add_token(tools.function.close); + add_token(tools.arguments.start); + add_token(tools.arguments.end); + add_token(tools.arguments.name_prefix); + add_token(tools.arguments.name_suffix); + add_token(tools.arguments.separator); + add_token(tools.arguments.value_prefix); + add_token(tools.arguments.value_suffix); + add_token(tools.call_id.prefix); + add_token(tools.call_id.suffix); +} + +analyze_reasoning::analyze_reasoning(const common_chat_template & tmpl, bool supports_tools) + : analyze_base(tmpl) { + LOG_DBG(ANSI_PURPLE "=== Starting differential analysis ===\n" ANSI_RESET); + LOG_DBG(ANSI_ORANGE "Phase 1: Reasoning analysis\n" ANSI_RESET); + + compare_reasoning_presence(); + compare_thinking_enabled(); + if (supports_tools) { + compare_reasoning_scope(); + } +} + +void analyze_reasoning::compare_reasoning_presence() { + json user_msg = json{ + { "role", "user" }, + { "content", USER_MSG } + }; + + json assistant_no_reasoning = json{ + { "role", "assistant" }, + { "content", ASSISTANT_MSG } + }; + + json assistant_with_reasoning = json{ + { "role", "assistant" }, + { "content", ASSISTANT_MSG }, + { "reasoning_content", THINKING_CONTENT } + }; + + template_params params; + params.messages = json::array({ user_msg, assistant_no_reasoning }); + params.add_generation_prompt = false; + params.enable_thinking = true; + + auto comparison = compare_variants( + *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_with_reasoning }); }); + + if (!comparison) { + LOG_DBG(ANSI_ORANGE "%s: Template application failed, skipping reasoning detection\n" ANSI_RESET, __func__); + return; + } + + const auto & diff = comparison->diff; + + const std::string reasoning_content = THINKING_CONTENT; + + if (!diff.right.empty() && diff.right.find(reasoning_content) != std::string::npos) { + auto parser_delimiter = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.literal(reasoning_content) + p.space() + p.optional(p.tag("post", (p.marker() + p.space())) + p.rest()); + }); + auto parser_wrapped = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.tag("pre", p.marker() + p.space()) + p.literal(reasoning_content) + p.tag("post", (p.space() + p.marker() + p.space())) + p.rest(); + }); + // try the more aggressive parse first, if it fails, fall back to the delimiter one + auto result = parser_wrapped.parse_anywhere_and_extract(comparison->output_B); + if (!result.result.success()) { + result = parser_delimiter.parse_anywhere_and_extract(comparison->output_B); + } + if (result.result.success()) { + if (!result.tags["pre"].empty() && !result.tags["post"].empty()) { + mode = reasoning_mode::TAG_BASED; + start = result.tags["pre"]; + end = result.tags["post"]; + } else if (!result.tags["post"].empty()) { + mode = reasoning_mode::TAG_BASED; + end = result.tags["post"]; + } + } + } +} + +void analyze_reasoning::compare_thinking_enabled() { + json user_msg = json{ + { "role", "user" }, + { "content", USER_MSG } + }; + + template_params params; + params.messages = json::array({ user_msg }); + params.add_generation_prompt = true; + params.enable_thinking = false; + + auto comparison = compare_variants(*tmpl, params, [&](template_params & p) { p.enable_thinking = true; }); + + if (!comparison) { + LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET , __func__); + return; + } + + const auto & diff = comparison->diff; + + std::string left_trimmed = trim_whitespace(diff.left); + std::string right_trimmed = trim_whitespace(diff.right); + + if (left_trimmed.empty() && !diff.right.empty()) { + if (!right_trimmed.empty() && string_ends_with(comparison->output_B, right_trimmed)) { + if (start.empty()) { + start = trim_leading_whitespace(diff.right); + mode = reasoning_mode::TAG_BASED; + } + } + } else if (right_trimmed.empty() && !diff.left.empty()) { + if (!left_trimmed.empty() && string_ends_with(comparison->output_A, left_trimmed)) { + if (end.empty()) { + auto seg = prune_whitespace_segments(segmentize_markers(comparison->output_A)); + if (seg.size() >= 2 && seg[seg.size() - 1].value == left_trimmed && seg[seg.size() - 2].type == segment_type::MARKER) { + start = seg[seg.size() - 2].value; + } + end = trim_trailing_whitespace(diff.left); + mode = reasoning_mode::TAG_BASED; + } + } + } else if (!left_trimmed.empty() && !right_trimmed.empty()) { + // Full-output diff is noisy (e.g., SmolLM3 changes the system message when enable_thinking flips). + // Try to find reasoning markers by tail-anchoring: + // one output's generation prompt tail may appear in the other with extra reasoning markers appended. + const auto & output_A = comparison->output_A; + const auto & output_B = comparison->output_B; + const size_t anchor_len = 64; + + for (int dir = 0; dir < 2; dir++) { + const auto & base = dir == 0 ? output_B : output_A; + const auto & extended = dir == 0 ? output_A : output_B; + + size_t len = std::min(base.size(), anchor_len); + std::string anchor = base.substr(base.size() - len); + auto pos = extended.rfind(anchor); + if (pos == std::string::npos || pos + len >= extended.size()) { + continue; + } + + std::string extra = trim_whitespace(extended.substr(pos + len)); + if (extra.empty()) { + continue; + } + + auto seg = prune_whitespace_segments(segmentize_markers(extra)); + if (seg.size() == 2 && seg[0].type == segment_type::MARKER && seg[1].type == segment_type::MARKER) { + if (start.empty()) { + start = seg[0].value; + } + if (end.empty()) { + end = seg[1].value; + } + mode = reasoning_mode::TAG_BASED; + break; + } + } + } + + if (mode == reasoning_mode::NONE && start.empty() && !end.empty()) { + mode = reasoning_mode::TAG_BASED; + } +} + +void analyze_reasoning::compare_reasoning_scope() { + json assistant_reasoning_content = json{ + { "role", "assistant" }, + { "content", ASSISTANT_MSG }, + { "reasoning_content", THINKING_CONTENT } + }; + + json assistant_reasoning_tools = json{ + { "role", "assistant" }, + { "content", nullptr }, + { "reasoning_content", THINKING_CONTENT }, + { "tool_calls", + json::array({ build_tool_call(FUN_FIRST, json{ { ARG_FIRST, "VVVV" }, { ARG_SECOND, "XXXX" } }) }) } + }; + + template_params params; + params.messages = json::array({ user_msg, assistant_reasoning_content }); + params.tools = tools; + params.add_generation_prompt = false; + params.enable_thinking = true; + + auto comparison = compare_variants( + *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_reasoning_tools }); }); + + if (!comparison) { + LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); + return; + } + + std::string reasoning_content = THINKING_CONTENT; + + // Check if reasoning only appears in variant B (with tools) + bool reasoning_in_A = comparison->output_A.find(reasoning_content) != std::string::npos; + bool reasoning_in_B = comparison->output_B.find(reasoning_content) != std::string::npos; + + if (!reasoning_in_A && reasoning_in_B) { + mode = reasoning_mode::TOOLS_ONLY; + LOG_DBG(ANSI_ORANGE "%s: Detected TOOLS_ONLY reasoning mode\n" ANSI_RESET, __func__); + + auto parser_wrapped = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.tag("pre", p.marker() + p.space()) + p.literal(reasoning_content) + p.space() + p.tag("post", (p.marker() + p.space())); + }); + auto result = parser_wrapped.parse_anywhere_and_extract(comparison->output_B); + if (result.result.success()) { + start = result.tags["pre"]; + end = trim_trailing_whitespace(result.tags["post"]); + } else { + auto parser_delimiter = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.literal(reasoning_content) + p.space() + p.optional(p.tag("post", (p.marker() + p.space()))); + }); + result = parser_delimiter.parse_anywhere_and_extract(comparison->output_B); + if (result.result.success()) { + end = trim_trailing_whitespace(result.tags["post"]); + } else { + LOG_DBG(ANSI_ORANGE "%s: Unable to extract reasoning markers, falling back to reasoning = NONE\n" ANSI_RESET, __func__); + mode = reasoning_mode::NONE; + } + } + } +} + +analyze_content::analyze_content(const common_chat_template & tmpl, const analyze_reasoning & reasoning) + : analyze_base(tmpl) { + LOG_DBG(ANSI_ORANGE "Phase 2: Content analysis\n" ANSI_RESET); + + json assistant_content_only = json{ + { "role", "assistant" }, + { "content", ASSISTANT_MSG } + }; + + json assistant_with_tools = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ build_tool_call("test_func", json{ { "arg1", "value1" } }) }) } + }; + + json assistant_with_reasoning = json{ + { "role", "assistant" }, + { "content", "" }, + { "reasoning_content", THINKING_CONTENT } + }; + + template_params params_content_only; + params_content_only.messages = json::array({ user_msg, assistant_content_only }); + params_content_only.add_generation_prompt = false; + params_content_only.enable_thinking = true; + params_content_only.tools = tools; + + auto comparison_with_tools = compare_variants(tmpl, params_content_only, [&](template_params & p) { + p.messages = json::array({ user_msg, assistant_with_tools }); + }); + + auto comparison_with_reasoning = compare_variants(tmpl, params_content_only, [&](template_params & p) { + p.messages = json::array({ user_msg, assistant_with_reasoning }); + }); + + if (!comparison_with_tools || !comparison_with_reasoning) { + LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); + return; + } + + const auto & diff_tools = comparison_with_tools->diff; + const auto & diff_reasoning = comparison_with_reasoning->diff; + + std::string response = ASSISTANT_MSG; + + bool found_plain_content = false; + if (trim_whitespace(diff_tools.left) == response) { + auto parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { + return p.space() + diff_reasoning.left + p.space() + p.optional(p.marker()) + p.space() + p.end(); + }); + if (parser.parse_and_extract(diff_reasoning.left).result.success()) { + // We only have the content text in the diff (possibly with a stray EOG marker), so no markers + mode = content_mode::PLAIN; + found_plain_content = true; + } else if (reasoning.mode != reasoning_mode::NONE && !reasoning.end.empty()) { + auto post_reasoning_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { + return p.literal(reasoning.end) + p.space() + p.literal(response); + }); + if (post_reasoning_parser.parse_anywhere_and_extract(diff_reasoning.left).result.success()) { + mode = content_mode::PLAIN; + found_plain_content = true; + } + } + } + if (!found_plain_content) { + std::string rdiff = diff_reasoning.left; + if (!reasoning.end.empty() && rdiff.find(reasoning.end) != std::string::npos) { + rdiff = rdiff.substr(rdiff.find(reasoning.end) + reasoning.end.length()); + } + // Take the more promising diff + std::string pure_content = rdiff.length() > diff_tools.left.length() ? rdiff : diff_tools.left; + auto parser_wrapped = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.tag("pre", p.marker() + p.space()) + p.literal(response) + p.space() + p.tag("post", (p.marker() + p.space())) + p.rest(); + }); + auto result = parser_wrapped.parse_anywhere_and_extract(pure_content); + start = result.tags["pre"]; + end = result.tags["post"]; + // TODO: WRAPPED_WITH_REASONING + } + + // Determine content mode + if (!start.empty() || !end.empty()) { + mode = content_mode::ALWAYS_WRAPPED; + // TODO: END_DELIMITED content mode - delimited at end but not at start? + } +} + +bool analyze_content::is_always_wrapped() const { + return mode == content_mode::ALWAYS_WRAPPED && !start.empty() && !end.empty(); +} + +analyze_tools::analyze_tools(const common_chat_template & tmpl, + const jinja::caps & caps, + const analyze_reasoning & reasoning) + : analyze_base(tmpl) { + LOG_DBG(ANSI_ORANGE "Phase 3: Tool call analysis\n" ANSI_RESET); + + analyze_tool_calls(reasoning, caps.supports_parallel_tool_calls); + + if (format.mode != tool_format::NONE && format.mode != tool_format::JSON_NATIVE) { + if (caps.supports_parallel_tool_calls) { + check_per_call_markers(); + } + LOG_DBG(ANSI_ORANGE "Phase 3a: Function call analysis\n" ANSI_RESET); + extract_function_markers(); + LOG_DBG(ANSI_ORANGE "Phase 3b: Argument analysis\n" ANSI_RESET); + if (format.mode == tool_format::TAG_WITH_TAGGED) { + analyze_arguments(); + } + extract_argument_separator(); + extract_args_markers(); + LOG_DBG(ANSI_ORANGE "Phase 3c: Call id analysis\n" ANSI_RESET); + extract_call_id_markers(); + } +} + +void analyze_tools::analyze_tool_calls(const analyze_reasoning & reasoning, bool supports_parallel_tool_calls) { + json assistant_no_tools = json{ + { "role", "assistant" }, + { "content", ASSISTANT_MSG } + }; + + json assistant_with_tools = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call }) } + }; + + template_params params; + params.messages = json::array({ user_msg, assistant_no_tools }); + params.tools = tools; + params.add_generation_prompt = false; + params.enable_thinking = true; + + auto comparison = compare_variants( + *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_with_tools }); }); + + if (!comparison) { + LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); + return; + } + + const auto & diff = comparison->diff; + + std::string tool_section = diff.right; + + if (tool_section.empty()) { + return; + } + + analyze_tool_call_format(tool_section, FUN_FIRST, ARG_FIRST, reasoning, supports_parallel_tool_calls); +} + +void analyze_tools::analyze_tool_call_format(const std::string & haystack, + const std::string & fun_name_needle, + const std::string & arg_name_needle, + const analyze_reasoning & reasoning, + bool supports_parallel_tool_calls) { + if (fun_name_needle.empty() || arg_name_needle.empty() || haystack.empty()) { + return; + } + + auto in_json_haystack = [&haystack](const std::string & needle) -> bool { + auto parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.choice({ p.literal("{"), p.literal(":") }) << p.choice({ + p.tag("dq", p.literal("\"") + p.literal(needle) + p.literal("\"")) }); + }); + auto result = parser.parse_anywhere_and_extract(haystack); + return result.result.success(); + }; + + auto fun_quote = in_json_haystack(fun_name_needle); + auto arg_quote = in_json_haystack(arg_name_needle); + + if (fun_quote) { + // no need to check further, we're in JSON land + format.mode = tool_format::JSON_NATIVE; + } else if (arg_quote) { + format.mode = tool_format::TAG_WITH_JSON; + } else { + format.mode = tool_format::TAG_WITH_TAGGED; + } + + // first, remove any reasoning markers + std::string clean_haystack = haystack; + if (!reasoning.start.empty()) { + auto pos = haystack.find(reasoning.start); + if (pos != std::string::npos) { + clean_haystack = haystack.substr(0, pos) + haystack.substr(pos + reasoning.start.length()); + } + } + if (!reasoning.end.empty()) { + auto pos = clean_haystack.find(reasoning.end); + if (pos != std::string::npos) { + clean_haystack = clean_haystack.substr(0, pos) + clean_haystack.substr(pos + reasoning.end.length()); + } + } + + if (format.mode == tool_format::JSON_NATIVE) { + analyze_tool_call_format_json_native(clean_haystack, fun_name_needle, arg_name_needle); + if (supports_parallel_tool_calls) { + analyze_json_native_parallel_calls(); + } + } else { + analyze_tool_call_format_non_json(clean_haystack, fun_name_needle); + } + // always relax whitespace requirements on ending markers since they don't influence content + format.section_end = trim_whitespace(format.section_end); + format.per_call_end = trim_whitespace(format.per_call_end); +} + +void analyze_tools::analyze_json_native_parallel_calls() { + json assistant_one_tool = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call }) } + }; + + json assistant_two_tools = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call, second_tool_call }) } + }; + + template_params params; + params.messages = json::array({ user_msg, assistant_one_tool }); + params.tools = tools; + params.add_generation_prompt = false; + params.enable_thinking = true; + + auto comparison = compare_variants( + *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_two_tools }); }); + + if (!comparison) { + LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); + return; + } + + std::string & second_call = comparison->diff.right; + if (!format.section_start.empty() && second_call.find(format.section_start) != std::string::npos) { + format.per_call_start = format.section_start; + format.per_call_end = format.section_end; + format.section_start.clear(); + format.section_end.clear(); + } +} + +void analyze_tools::analyze_tool_call_format_json_native(const std::string & clean_haystack, + const std::string & fun_name_needle, + const std::string & arg_name_needle) { + // we might not have the typical OpenAI tool calling structure + int json_start = clean_haystack.find_first_of('{'); + int json_end = clean_haystack.find_last_of('}'); + std::string cut = clean_haystack.substr(json_start, json_end - json_start + 1); + json call_struct = json::parse(cut); + auto register_field = [&](const std::string & prefix, const nlohmann::detail::iteration_proxy_value & subel) { + if (subel.value().is_string() && std::string(subel.value()).find("call0000") != std::string::npos) { + format.id_field = !prefix.empty() ? prefix + "." + subel.key() : subel.key(); + } else if (subel.value().is_string() && std::string(subel.value()) == fun_name_needle) { + format.name_field = !prefix.empty() ? prefix + "." + subel.key() : subel.key(); + } else if (subel.value().dump().find(arg_name_needle) != + std::string::npos) { // handle both string and JSON obj variants + format.args_field = !prefix.empty() ? prefix + "." + subel.key() : subel.key(); + } else if (subel.key().find("id") != std::string::npos) { + // heuristics for generated id field + format.gen_id_field = !prefix.empty() ? prefix + "." + subel.key() : subel.key(); + } + }; + for (const auto & el : call_struct.items()) { + if (el.key() == fun_name_needle) { + format.fun_name_is_key = true; + // When function name is the key, there's no name field and args are direct + format.name_field.clear(); + format.args_field.clear(); + // Don't register this element - the function name IS the key, not a field + } else { + if (el.value().is_object() && + el.value().dump().find(arg_name_needle) == std::string::npos) { // not the args object + format.function_field = el.key(); + for (const auto & subel : el.value().items()) { + register_field(el.key(), subel); + } + } + // Register this element as a potential field + register_field("", el); + } + } + auto array_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.tag("pre", p.literal("[") + p.space()) + p.literal(cut) + p.tag("post", p.space() + p.literal("]")); + }); + + auto ar_parse_res = array_parser.parse_anywhere_and_extract(clean_haystack); + if (ar_parse_res.result.success()) { + format.tools_array_wrapped = true; + json_start -= ar_parse_res.tags["pre"].length(); + json_end += ar_parse_res.tags["post"].length(); + } + json_end++; // we want to move past the closing char for end marker extraction + + std::vector> located_params; + if (!format.name_field.empty()) { + located_params.push_back({ clean_haystack.find(format.name_field), format.name_field }); + } + if (!format.args_field.empty()) { + located_params.push_back({ clean_haystack.find(format.args_field), format.args_field }); + } + if (!format.id_field.empty()) { + located_params.push_back({ clean_haystack.find(format.id_field), format.id_field }); + } + if (!format.gen_id_field.empty()) { + located_params.push_back({ clean_haystack.find(format.gen_id_field), format.gen_id_field }); + } + std::sort(located_params.begin(), located_params.end()); + for (auto & pair : located_params) { + format.parameter_order.push_back(pair.second); + } + // we can immediately extract tool calling markers too + format.section_start = trim_leading_whitespace(clean_haystack.substr(0, json_start)); + format.section_end = trim_whitespace(clean_haystack.substr(json_end)); + // When tools_array_wrapped is true, the closing bracket is part of the array structure, + // not a separate section end marker. Clear tool_section_end to avoid duplicate brackets. + if (format.tools_array_wrapped && format.section_end == "]") { + format.section_end.clear(); + } +} + +void analyze_tools::analyze_tool_call_format_non_json(const std::string & clean_haystack, + const std::string & fun_name_needle) { + // first, let's find out if the function is inside a tag or standalone + auto fun_marker_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.tag("fun_marker", p.choice({ + p.tag("fun_pre", p.literal("<") + p.until_one_of({ ">", fun_name_needle })) + p.literal(fun_name_needle) + + p.tag("fun_post", p.negate(p.space() + p.literal("<")) + p.until(">") + p.literal(">")) + p.space(), + p.tag("fun_pre", p.literal("[") + p.until_one_of({ "]", fun_name_needle })) + p.literal(fun_name_needle) + + p.tag("fun_post", p.negate(p.space() + p.literal("[") + p.until("]") + p.literal("]")) + p.space()) })); + }); + auto fun_res = fun_marker_parser.parse_anywhere_and_extract(clean_haystack); + std::string fun_marker = fun_name_needle; + if (fun_res.result.success()) { + fun_marker = fun_res.tags["fun_marker"]; + } + // now, consume up to two markers, then treat everything up to the function marker as function name prefix + auto per_tool_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.tag("sec_start", p.marker() + p.space()) + p.tag("call_start", p.marker() + p.space()) + + p.tag("fun_pre", p.until(fun_marker)) + fun_marker + p.tag("rest", p.rest()); + }); + auto section_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.tag("sec_start", p.marker() + p.space()) + fun_marker + p.tag("rest", p.rest()); + }); + auto result = per_tool_parser.parse_anywhere_and_extract(clean_haystack); + tagged_parse_result result_end; + if (result.result.success()) { + auto double_closer_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.tag("call_end", p.marker() + p.space()) + p.tag("sec_end", p.marker() + p.space()) + p.end(); + }); + result_end = double_closer_parser.parse_anywhere_and_extract(result.tags["rest"]); + function.name_prefix = fun_res.tags["fun_pre"] + function.name_prefix; + } else { + result = section_parser.parse_anywhere_and_extract(clean_haystack); + auto single_closer_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.tag("sec_end", p.marker() + p.space()) + p.end(); + }); + result_end = single_closer_parser.parse_anywhere_and_extract(result.tags["rest"]); + } + format.per_call_start = result.tags["call_start"]; + format.per_call_end = result_end.tags["call_end"]; + format.section_start = result.tags["sec_start"]; + format.section_end = result_end.tags["sec_end"]; +} + +void analyze_tools::check_per_call_markers() { + json assistant_one_tool = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call }) } + }; + + json assistant_two_tools = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call, second_tool_call }) } + }; + + template_params params; + params.messages = json::array({ user_msg, assistant_one_tool }); + params.tools = tools; + params.add_generation_prompt = false; + params.enable_thinking = true; + + auto one_vs_two = compare_variants( + *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_two_tools }); }); + + if (!one_vs_two) { + LOG_DBG(ANSI_ORANGE "%s: Generating double tool call comparison failed\n" ANSI_RESET, __func__); + return; + } + + diff_split filter_common_call_part = calculate_diff_split(one_vs_two->diff.suffix, one_vs_two->diff.right); + + std::string second_tool_content = trim_leading_whitespace(filter_common_call_part.right); + if (!format.section_start.empty() && + second_tool_content.find(format.section_start) == 0) { + format.per_call_start = format.section_start; + format.per_call_end = format.section_end; + format.section_start.clear(); + format.section_end.clear(); + } +} + +void analyze_tools::extract_function_markers() { + json assistant_nocall = json{ + { "role", "assistant" }, + { "content", ASSISTANT_MSG }, + }; + + json assistant_foofoo = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call }) } + }; + + json assistant_barbar = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ second_tool_call }) } + }; + + template_params params; + params.messages = json::array({ user_msg, assistant_foofoo }); + params.tools = tools; + params.add_generation_prompt = false; + params.enable_thinking = true; + + auto comparison = compare_variants( + *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_barbar }); }); + + if (!comparison) { + LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); + return; + } + + const auto & diff = comparison->diff; + + if (diff.left.find(FUN_FIRST) != std::string::npos && diff.right.find(FUN_SECOND) != std::string::npos) { + std::string prefix_marker; + if (!format.per_call_start.empty()) { + prefix_marker = format.per_call_start; + } else { + prefix_marker = format.section_start; + } + if (!prefix_marker.empty() && diff.prefix.rfind(prefix_marker) != std::string::npos) { + function.name_prefix = + diff.prefix.substr(diff.prefix.rfind(prefix_marker) + prefix_marker.size()); + } + + // Extract name prefix/suffix from diff.left (stop at the next marker boundary) + auto name_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.tag("pre", p.until(FUN_FIRST)) + p.literal(FUN_FIRST) + + p.tag("post", p.zero_or_more(p.negate(p.marker()) + p.any())); + }); + auto name_result = name_parser.parse_and_extract(diff.left); + if (name_result.result.success()) { + function.name_prefix += name_result.tags["pre"]; + function.name_suffix = name_result.tags["post"]; + } + + // Extend name_suffix with content from diff.suffix before args begin + if (format.mode == tool_format::TAG_WITH_JSON) { + // For JSON: name_suffix extends to the first non-marker { or [, including any + // markers along the way. Only applies if there's at least one marker after + // the JSON content (matching the original "stop < seg_suf.size() - 1" guard). + auto suffix_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + auto non_json = p.marker() | (p.negate(p.literal("{")) + p.negate(p.literal("[")) + p.any()); + auto after_json = p.zero_or_more(p.negate(p.marker()) + p.any()) + p.marker(); + return p.tag("ext", p.zero_or_more(non_json)) + after_json; + }); + auto suf_result = suffix_parser.parse_and_extract(diff.suffix); + if (suf_result.result.success()) { + function.name_suffix += suf_result.tags["ext"]; + } + } else { + // For tagged: name_suffix extends to the first marker (arg marker) + auto suffix_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.tag("ext", p.zero_or_more(p.negate(p.marker()) + p.any())); + }); + auto suf_result = suffix_parser.parse_and_extract(diff.suffix); + if (suf_result.result.success()) { + function.name_suffix += suf_result.tags["ext"]; + } + } + + // Extract the closer (between last arg and call/section end marker) + std::string suffix_marker; + if (!format.per_call_end.empty()) { + suffix_marker = format.per_call_end; + } else { + suffix_marker = format.section_end; + } + std::string closer_suffix; + if (suffix_marker.empty()) { + // we'll have to rely on an extra diff with no-calls version + auto notool_comp = compare_variants( + *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_nocall }); }); + if (notool_comp) { + auto nt_diff = notool_comp->diff; + closer_suffix = nt_diff.left.substr(nt_diff.left.find("YYYY") + 4); + } + } else { + closer_suffix = diff.suffix.substr(0, diff.suffix.find(suffix_marker)); + } + if (!closer_suffix.empty()) { + if (format.mode == tool_format::TAG_WITH_TAGGED) { + // After last arg value, skip the closing arg marker, rest is closer + auto closer_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.until("YYYY") + p.literal("YYYY") + p.space() + + p.marker() + p.space() + + p.tag("close", p.rest()); + }); + auto close_result = closer_parser.parse_and_extract(closer_suffix); + if (close_result.result.success()) { + function.close = close_result.tags["close"]; + } + } else if (format.mode == tool_format::TAG_WITH_JSON) { + // After last arg value, find end of JSON args, rest is closer + auto closer_parser = build_tagged_peg_parser([&](common_peg_parser_builder &p) { + return p.until("YYYY") + p.literal("YYYY") + p.tag("post_val", p.rest()); + }); + auto close_result = closer_parser.parse_and_extract(closer_suffix); + if (close_result.result.success()) { + const auto & post = close_result.tags["post_val"]; + size_t pos = post.find_last_of("}]"); + if (pos != std::string::npos && pos < post.size() - 1) { + function.close = trim_leading_whitespace(post.substr(pos + 1)); + } + } + } + } + function.close = trim_leading_whitespace(function.close); + } +} + +void analyze_tools::analyze_arguments() { + extract_argument_name_markers(); + extract_argument_value_markers(); +} + +void analyze_tools::extract_argument_name_markers() { + json assistant_first_arg = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call_one_arg }) } + }; + + json assistant_second_arg = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call_other_arg }) } + }; + + template_params params; + params.messages = json::array({ user_msg, assistant_first_arg }); + params.tools = tools; + params.add_generation_prompt = false; + params.enable_thinking = true; + + auto comparison = compare_variants( + *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_second_arg }); }); + + if (!comparison) { + LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); + return; + } + + const auto & diff = comparison->diff; + + if (!diff.left.empty() && !diff.right.empty()) { + // Parse both sides to find ARG_FIRST/ARG_SECOND and extract the surrounding structure + auto left_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { + return p.tag("pre", p.until(ARG_FIRST)) + p.literal(ARG_FIRST) + + p.tag("suffix", p.until_one_of({"\"", "X"})); + }); + auto right_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { + return p.tag("pre", p.until(ARG_SECOND)) + p.literal(ARG_SECOND) + + p.tag("suffix", p.until_one_of({"\"", "Y"})); + }); + auto left_result = left_parser.parse_anywhere_and_extract(diff.left); + auto right_result = right_parser.parse_anywhere_and_extract(diff.right); + + if (left_result.result.success() && right_result.result.success() && + !left_result.tags["pre"].empty() && + left_result.tags["pre"] == right_result.tags["pre"] && + left_result.tags["suffix"] == right_result.tags["suffix"]) { + // Name is inside a structure (e.g., JSON key): prefix is the shared wrapper + arguments.name_prefix = trim_whitespace(left_result.tags["pre"]); + arguments.name_suffix = trim_leading_whitespace(left_result.tags["suffix"]); + } else if (diff.left.substr(0, ARG_FIRST.length()) == ARG_FIRST && diff.right.substr(0, ARG_SECOND.length()) == ARG_SECOND) { + // Name is directly in the diff: prefix comes from last marker in diff.prefix + auto pre_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { + auto last_marker = p.marker() + p.zero_or_more(p.negate(p.marker()) + p.any()) + p.end(); + return p.zero_or_more(p.negate(last_marker) + p.any()) + p.tag("name_prefix", last_marker); + }); + auto pre_result = pre_parser.parse_and_extract(diff.prefix); + arguments.name_prefix = pre_result.result.success() + ? pre_result.tags["name_prefix"] : diff.prefix; + + // Suffix extends from after ARG_FIRST to the first marker (+ optional whitespace). + // The marker could be in diff.left itself or in diff.suffix, so we concatenate. + std::string after_first = diff.left.substr(ARG_FIRST.length()) + diff.suffix; + auto suffix_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { + return p.tag("suffix", p.zero_or_more(p.negate(p.marker()) + p.any()) + + p.marker() + p.space()); + }); + auto suf_result = suffix_parser.parse_anywhere_and_extract(after_first); + if (suf_result.result.success()) { + arguments.name_suffix = suf_result.tags["suffix"]; + } + } + } +} + +void analyze_tools::extract_argument_value_markers() { + json assistant_val_X = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call_one_arg }) } + }; + + json assistant_val_Y = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call_one_arg_other_val }) } + }; + + template_params params; + params.messages = json::array({ user_msg, assistant_val_X }); + params.tools = tools; + params.add_generation_prompt = false; + params.enable_thinking = true; + + auto comparison = compare_variants( + *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_val_Y }); }); + + if (!comparison) { + LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); + return; + } + + const auto & diff = comparison->diff; + + if (diff.left == "XXXX" && diff.right == "YYYY") { + std::string arg_name_ending = ARG_FIRST + arguments.name_suffix; + std::string prefix = diff.prefix; + if (prefix.rfind(arg_name_ending) != std::string::npos) { + prefix = prefix.substr(prefix.rfind(arg_name_ending) + arg_name_ending.size()); + } + if (!prefix.empty()) { + // Find the last marker + any trailing non-marker text to end + auto prefix_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { + auto last_marker = p.marker() + p.zero_or_more(p.negate(p.marker()) + p.any()) + p.end(); + return p.zero_or_more(p.negate(last_marker) + p.any()) + p.tag("val_prefix", last_marker); + }); + auto pre_result = prefix_parser.parse_and_extract(prefix); + arguments.value_prefix = pre_result.result.success() ? pre_result.tags["val_prefix"] : prefix; + } + + std::string value_suffix = diff.suffix; + if (!function.close.empty()) { + size_t func_close_pos = value_suffix.find(function.close); + if (func_close_pos != std::string::npos) { + value_suffix = value_suffix.substr(0, func_close_pos); + } + } else if (!format.per_call_end.empty() || !format.section_end.empty()) { + std::string end_marker = + !format.per_call_end.empty() ? format.per_call_end : format.section_end; + size_t end_marker_pos = value_suffix.find(end_marker); + if (end_marker_pos != std::string::npos) { + value_suffix = value_suffix.substr(0, end_marker_pos); + } + } + value_suffix = trim_leading_whitespace(value_suffix); + if (!value_suffix.empty()) { + arguments.value_suffix = value_suffix; + } + } +} + +void analyze_tools::extract_argument_separator() { + json assistant_one_arg = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call_one_arg }) } + }; + + json assistant_two_args = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call }) } + }; + + template_params params; + params.messages = json::array({ user_msg, assistant_one_arg }); + params.tools = tools; + params.add_generation_prompt = false; + params.enable_thinking = true; + + auto comparison = compare_variants( + *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_two_args }); }); + + if (!comparison) { + LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); + return; + } + + const auto & diff = comparison->diff; + + if (!diff.right.empty()) { + std::string separator = until_common_prefix(diff.right, ARG_FIRST, ARG_SECOND); + arguments.separator = separator; + } +} + +void analyze_tools::extract_args_markers() { + json assistant_no_args = json{ + { "role", "assistant"}, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call_zero_args }) } + }; + + json assistant_with_args = json{ + { "role", "assistant"}, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call_one_arg }) } + }; + + template_params params; + params.messages = json::array({ user_msg, assistant_no_args }); + params.tools = tools; + params.add_generation_prompt = false; + params.enable_thinking = true; + + auto comparison = compare_variants( + *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_with_args }); }); + + if (!comparison) { + LOG_DBG(ANSI_ORANGE "%s: Template application failed\n" ANSI_RESET, __func__); + return; + } + + const auto & diff = comparison->diff; + + if (format.mode == tool_format::JSON_NATIVE) { + std::string prefix_marker = !format.section_start.empty() ? format.section_start : format.per_call_start; + std::string suffix_marker = !format.section_end.empty() ? format.section_end : format.per_call_end; + // these might happen earlier in the tools section as an example or somewhere else, so we need to find the closest ones + size_t prefix_pos = prefix_marker.empty() ? 0 : diff.prefix.rfind(prefix_marker); + size_t suffix_pos = suffix_marker.empty() ? diff.suffix.size() : diff.suffix.find(suffix_marker); + if (prefix_pos == std::string::npos) { + prefix_pos = 0; + } + if (suffix_pos == std::string::npos) { + suffix_pos = diff.suffix.size(); + } + std::string prefix_cut = diff.prefix.substr(prefix_pos + prefix_marker.size()); + std::string suffix_cut = diff.suffix.substr(0, suffix_pos); + std::string args_start = until_common_prefix(prefix_cut, "{}", "{\"first\":"); + std::string args_end = after_common_suffix(suffix_cut, "{}", "\"XXXX\"}"); + + if (!args_start.empty() || !args_end.empty()) { + size_t find_fun = args_start.find(FUN_FIRST); + if (find_fun != std::string::npos) { + args_start = args_start.substr(find_fun + FUN_FIRST.size(), args_start.size() - find_fun - FUN_FIRST.size()); + } + size_t find_call_id = args_start.find(CALL_ID_001); + if (find_call_id != std::string::npos) { + args_start = args_start.substr(find_call_id + CALL_ID_001.size(), args_start.size() - find_call_id - CALL_ID_001.size()); + } + arguments.start = args_start; + arguments.end = args_end; + } + } +} + +void analyze_tools::extract_call_id_markers() { + json assistant_id1 = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call }) } + }; + + json assistant_id2 = json{ + { "role", "assistant" }, + { "content", "" }, + { "tool_calls", json::array({ first_tool_call_alt_id }) } + }; + + template_params params; + params.messages = json::array({ user_msg, assistant_id1 }); + params.tools = tools; + params.add_generation_prompt = false; + params.enable_thinking = true; + + auto comparison = compare_variants( + *tmpl, params, [&](template_params & p) { p.messages = json::array({ user_msg, assistant_id2 }); }); + + if (!comparison) { + LOG_DBG(ANSI_ORANGE "%s: Template application failed for call_id detection\n" ANSI_RESET, __func__); + return; + } + + const auto & diff = comparison->diff; + + if (diff.left.empty() && diff.right.empty()) { + return; + } + + std::string id_value_1 = CALL_ID_001; + std::string id_value_2 = CALL_ID_999; + + size_t common_id_prefix_len = 0; + for (size_t i = 0; i < std::min(id_value_1.length(), id_value_2.length()); i++) { + if (id_value_1[i] == id_value_2[i]) { + common_id_prefix_len++; + } else { + break; + } + } + std::string common_id_part = id_value_1.substr(0, common_id_prefix_len); + + // Check if the function name is in the prefix (normal case: BETWEEN_FUNC_AND_ARGS or POST_ARGS) + // or in the suffix (call_id is PRE_FUNC_NAME) + std::string func_name = FUN_FIRST; + size_t func_name_in_prefix = diff.prefix.rfind(func_name); + size_t func_name_in_suffix = diff.suffix.find(func_name); + + // Helper: find the last marker in a string (returns just the marker, not trailing text) + auto find_last_marker = [](const std::string & str) -> std::string { + auto parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { + auto last = p.marker() + p.zero_or_more(p.negate(p.marker()) + p.any()) + p.end(); + return p.zero_or_more(p.negate(last) + p.any()) + p.tag("m", p.marker()); + }); + auto res = parser.parse_anywhere_and_extract(str); + return res.result.success() ? res.tags["m"] : ""; + }; + + // Helper: find the first marker in a string + auto find_first_marker = [](const std::string & str) -> std::string { + auto parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { + return p.tag("m", p.marker()); + }); + auto res = parser.parse_anywhere_and_extract(str); + return res.result.success() ? res.tags["m"] : ""; + }; + + if (func_name_in_prefix != std::string::npos && func_name_in_suffix == std::string::npos) { + // Function name is only in prefix - call_id is BETWEEN_FUNC_AND_ARGS or POST_ARGS + // Check if args indicator "{" is in prefix or suffix + size_t args_in_prefix = diff.prefix.find('{', func_name_in_prefix); + size_t args_in_suffix = diff.suffix.find('{'); + + if (args_in_suffix != std::string::npos && + (args_in_prefix == std::string::npos || args_in_prefix > diff.prefix.length())) { + // Args are in suffix, so call_id is BETWEEN_FUNC_AND_ARGS + call_id.pos = call_id_position::BETWEEN_FUNC_AND_ARGS; + + // Find call_id_prefix: marker immediately preceding common_id_part (no intervening markers) + std::string after_func = diff.prefix.substr(func_name_in_prefix + func_name.length()); + auto id_prefix_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { + return p.tag("prefix", p.marker()) + + p.zero_or_more(p.negate(p.marker()) + p.negate(p.literal(common_id_part)) + p.any()) + + p.literal(common_id_part); + }); + auto id_res = id_prefix_parser.parse_anywhere_and_extract(after_func); + if (id_res.result.success()) { + call_id.prefix = id_res.tags["prefix"]; + } else { + // Fallback: use the last marker in after_func + call_id.prefix = find_last_marker(after_func); + } + + // Extract call_id_suffix: the first marker in the suffix before args "{" + auto suffix_parser = build_tagged_peg_parser([&](common_peg_parser_builder & p) { + return p.zero_or_more(p.negate(p.marker()) + p.negate(p.literal("{")) + p.any()) + + p.tag("suffix", p.marker()); + }); + auto suf_res = suffix_parser.parse_anywhere_and_extract(diff.suffix); + if (suf_res.result.success()) { + call_id.suffix = suf_res.tags["suffix"]; + } + } else if (args_in_prefix != std::string::npos) { + // Args are in prefix, so call_id is POST_ARGS + call_id.pos = call_id_position::POST_ARGS; + + // Extract last marker between args closing brace and the ID + std::string after_args = diff.prefix.substr(args_in_prefix); + size_t closing_brace = after_args.rfind('}'); + if (closing_brace != std::string::npos) { + std::string between_args_and_id = after_args.substr(closing_brace + 1); + call_id.prefix = find_last_marker(between_args_and_id); + } + + // call_id_suffix: first marker in diff.suffix + call_id.suffix = find_first_marker(diff.suffix); + } + } else if (func_name_in_suffix != std::string::npos && func_name_in_prefix == std::string::npos) { + // Function name is only in suffix - call_id is PRE_FUNC_NAME + call_id.pos = call_id_position::PRE_FUNC_NAME; + + // call_id_prefix: last marker in diff.prefix + call_id.prefix = find_last_marker(diff.prefix); + + // call_id_suffix: first marker in the portion of diff.suffix before func_name + std::string before_func = diff.suffix.substr(0, func_name_in_suffix); + call_id.suffix = find_first_marker(before_func); + } + + if (call_id.prefix == arguments.end) { + call_id.prefix = ""; + } + + if (call_id.suffix == arguments.start) { + call_id.suffix = ""; + } + + // When call_id is detected, per_call_end may have been incorrectly set to include + // the call_id_suffix and sample args. Clear it if it starts with call_id_suffix. + if (call_id.pos != call_id_position::NONE && !call_id.suffix.empty() && + format.per_call_end.find(call_id.suffix) == 0) { + format.per_call_end.clear(); + } +} + +} // namespace autoparser + +#include "common.h.inc" + +#include "common.h.inc" +#include "ggml.h.inc" + + +#include "nlohmann/json.hpp.inc" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +using json = nlohmann::ordered_json; + +static std::string format_time(const std::chrono::system_clock::time_point & now, const std::string & format) { + auto time = std::chrono::system_clock::to_time_t(now); + auto local_time = *std::localtime(&time); + std::ostringstream ss; + ss << std::put_time(&local_time, format.c_str()); + auto res = ss.str(); + return res; +} + +static json safe_args_parse(const std::string & to_parse) { + std::string stripped = to_parse; + if (to_parse.at(0) == '"' && to_parse.at(to_parse.length() - 1) == '"') { + stripped = to_parse.substr(1, to_parse.length() - 1); + } + try { + return json::parse(stripped); + } catch (json::exception & e) { + return stripped; + } +} + +static std::string string_diff(const std::string & last, const std::string & current) { + int status = -1; + llama_common_unicode_string diff = llama_common_string_diff_rust( + (const uint8_t *) last.data(), last.size(), + (const uint8_t *) current.data(), current.size(), + &status); + if (status < 0) { + throw std::runtime_error("Invalid diff: '" + last + "' not found at start of '" + current + "'"); + } + if (!diff.data) { + return ""; + } + std::string result((const char *) diff.data, diff.len); + llama_common_unicode_string_free(diff); + return result; +} + +json common_chat_msg::to_json_oaicompat(bool concat_typed_text) const { + if (!content.empty() && !content_parts.empty()) { + throw std::runtime_error("Cannot specify both content and content_parts"); + } + json jmsg { + {"role", role}, + }; + if (!content.empty()) { + jmsg["content"] = content; + } else if (!content_parts.empty()) { + if (concat_typed_text) { + std::string text; + bool last_was_media_marker = false; + // join parts with newline, do not add newline before or after media markers + for (const auto & part : content_parts) { + bool add_new_line = true; + if (part.type == "text") { + add_new_line = !last_was_media_marker && !text.empty(); + last_was_media_marker = false; + } else if (part.type == "media_marker") { + add_new_line = false; + last_was_media_marker = true; + } else { + LOG_WRN("Ignoring content part type: %s\n", part.type.c_str()); + continue; + } + + if (add_new_line) { + text += '\n'; + } + + text += part.text; + } + jmsg["content"] = text; + } else { + auto & parts = jmsg["content"] = json::array(); + for (const auto & part : content_parts) { + parts.push_back({ + {"type", part.type}, + {"text", part.text}, + }); + } + } + } else { + jmsg["content"] = ""; + } + if (!reasoning_content.empty()) { + jmsg["reasoning_content"] = reasoning_content; + } + if (!tool_name.empty()) { + jmsg["name"] = tool_name; + } + if (!tool_call_id.empty()) { + jmsg["tool_call_id"] = tool_call_id; + } + if (!tool_calls.empty()) { + jmsg["tool_calls"] = json::array(); + auto & jtool_calls = jmsg["tool_calls"]; + for (const auto & tool_call : tool_calls) { + json tc { + {"type", "function"}, + {"function", { + {"name", tool_call.name}, + {"arguments", json(tool_call.arguments)}, + }}, + }; + if (!tool_call.id.empty()) { + tc["id"] = tool_call.id; + } + // Some templates generate and require an id (sometimes in a very specific format, e.g. Mistral Nemo). + // We only generate a random id for the ones that don't generate one by themselves + // (they also won't get to see it as their template likely doesn't use it, so it's all for the client) + // {"id", tc.id.empty() ? gen_tool_call_id() : tc.id}, + jtool_calls.push_back(tc); + } + } + + return jmsg; +} + +std::vector common_chat_msg_diff::compute_diffs(const common_chat_msg & msg_prv, + const common_chat_msg & msg_new) { + std::vector diffs; + if (msg_new.tool_calls.size() > msg_prv.tool_calls.size()) { + diffs.reserve(msg_new.tool_calls.size() - msg_prv.tool_calls.size() + 3); + } else { + diffs.reserve(3); + } + + // TODO: these can become expensive for long messages - how to optimize? + if (msg_prv.reasoning_content != msg_new.reasoning_content) { + auto & diff = diffs.emplace_back(); + diff.reasoning_content_delta = string_diff(msg_prv.reasoning_content, msg_new.reasoning_content); + } + if (msg_prv.content != msg_new.content) { + auto & diff = diffs.emplace_back(); + diff.content_delta = string_diff(msg_prv.content, msg_new.content); + } + + if (msg_new.tool_calls.size() < msg_prv.tool_calls.size()) { + std::string err = "Invalid diff: now finding less tool calls!\n"; + err += " Previous (" + std::to_string(msg_prv.tool_calls.size()) + "):\n"; + for (const auto & tc : msg_prv.tool_calls) { + err += " - name: '" + tc.name + "', args: '" + tc.arguments + "'\n"; + } + err += " Current (" + std::to_string(msg_new.tool_calls.size()) + "):\n"; + for (const auto & tc : msg_new.tool_calls) { + err += " - name: '" + tc.name + "', args: '" + tc.arguments + "'\n"; + } + err += " Current msg text content:\n" + msg_new.content + "\n"; + throw std::runtime_error(err); + } + + if (!msg_prv.tool_calls.empty()) { + const auto idx = msg_prv.tool_calls.size() - 1; + const auto & pref = msg_prv.tool_calls[idx]; + const auto & newf = msg_new.tool_calls[idx]; + // Allow tool name to change during incremental parsing: + // - empty -> non-empty (initial discovery) + // - prefix -> longer string (name grows as more input is parsed) + if (pref.name != newf.name && !pref.name.empty() && !newf.name.empty()) { + // Check if one is a prefix of the other (for incremental parsing where names grow or shrink) + bool is_prefix = (newf.name.rfind(pref.name, 0) == 0); + if (!is_prefix) { + LOG_ERR("Tool call mismatch: prev='%s' new='%s'\n", pref.name.c_str(), newf.name.c_str()); + throw std::runtime_error("Invalid diff: tool call mismatch!"); + } + } + const auto args_diff = string_diff(pref.arguments, newf.arguments); + if (!args_diff.empty() || pref.id != newf.id || pref.name != newf.name) { + auto & diff = diffs.emplace_back(); + diff.tool_call_index = idx; + if (pref.id != newf.id || pref.name != newf.name) { + diff.tool_call_delta.id = newf.id; + diff.tool_call_delta.name = newf.name; + } + diff.tool_call_delta.arguments = args_diff; + } + } + for (size_t idx = msg_prv.tool_calls.size(); idx < msg_new.tool_calls.size(); ++idx) { + auto & diff = diffs.emplace_back(); + diff.tool_call_index = idx; + diff.tool_call_delta = msg_new.tool_calls[idx]; + } + + return diffs; +} + +using chat_template_caps = jinja::caps; + +struct common_chat_templates { + bool add_bos; + bool add_eos; + bool has_explicit_template; // Model had builtin template or template overridden was specified. + std::unique_ptr template_default; // always set (defaults to chatml) + std::unique_ptr template_tool_use; +}; + +common_chat_tool_choice common_chat_tool_choice_parse_oaicompat(const std::string & tool_choice) { + const int parsed = llama_common_chat_tool_choice_parse_oaicompat_rust(tool_choice.c_str()); + if (parsed >= 0) { + return (common_chat_tool_choice) parsed; + } + throw std::invalid_argument("Invalid tool_choice: " + tool_choice); +} + +bool common_chat_templates_support_enable_thinking(const common_chat_templates * chat_templates) { + common_chat_templates_inputs inputs; + inputs.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK; + common_chat_msg msg; + msg.role = "user"; + msg.content = "test"; + inputs.messages = { msg }; + inputs.enable_thinking = true; + inputs.add_generation_prompt = true; + inputs.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK; + + auto params = common_chat_templates_apply(chat_templates, inputs); + return params.supports_thinking; +} + +std::vector common_chat_msgs_parse_oaicompat(const json & messages) { + std::vector msgs; + + try { + if (!messages.is_array()) { + throw std::invalid_argument("Expected 'messages' to be an array, got " + messages.dump()); + } + + for (const auto & message : messages) { + if (!message.is_object()) { + throw std::invalid_argument("Expected 'message' to be an object, got " + message.dump()); + } + + common_chat_msg msg; + if (!message.contains("role")) { + throw std::invalid_argument("Missing 'role' in message: " + message.dump()); + } + msg.role = message.at("role"); + + auto has_content = message.contains("content"); + auto has_tool_calls = message.contains("tool_calls"); + if (has_content) { + const auto & content = message.at("content"); + if (content.is_string()) { + msg.content = content; + } else if (content.is_array()) { + for (const auto & part : content) { + if (!part.contains("type")) { + throw std::invalid_argument("Missing content part type: " + part.dump()); + } + const auto & type = part.at("type"); + if (type != "text" && type != "media_marker") { + throw std::invalid_argument("Unsupported content part type: " + type.dump()); + } + common_chat_msg_content_part msg_part; + msg_part.type = type; + msg_part.text = part.at("text"); + msg.content_parts.push_back(msg_part); + } + } else if (!content.is_null()) { + throw std::invalid_argument("Invalid 'content' type: expected string or array, got " + + content.dump() + + " (ref: https://github.com/ggml-org/llama.cpp/issues/8367)"); + } + } + if (has_tool_calls) { + for (const auto & tool_call : message.at("tool_calls")) { + common_chat_tool_call tc; + if (!tool_call.contains("type")) { + throw std::invalid_argument("Missing tool call type: " + tool_call.dump()); + } + const auto & type = tool_call.at("type"); + if (type != "function") { + throw std::invalid_argument("Unsupported tool call type: " + tool_call.dump()); + } + if (!tool_call.contains("function")) { + throw std::invalid_argument("Missing tool call function: " + tool_call.dump()); + } + const auto & fc = tool_call.at("function"); + if (!fc.contains("name")) { + throw std::invalid_argument("Missing tool call name: " + tool_call.dump()); + } + tc.name = fc.at("name"); + const auto & args = fc.at("arguments"); + if (args.is_string()) { + tc.arguments = args; + } else { + tc.arguments = args.dump(); + } + if (tool_call.contains("id")) { + tc.id = tool_call.at("id"); + } + msg.tool_calls.push_back(tc); + } + } + if (!has_content && !has_tool_calls) { + throw std::invalid_argument( + "Expected 'content' or 'tool_calls' (ref: https://github.com/ggml-org/llama.cpp/issues/8367 & " + "https://github.com/ggml-org/llama.cpp/issues/12279)"); + } + if (message.contains("reasoning_content")) { + msg.reasoning_content = message.at("reasoning_content"); + } + if (message.contains("name")) { + msg.tool_name = message.at("name"); + } + if (message.contains("tool_call_id")) { + msg.tool_call_id = message.at("tool_call_id"); + } + + msgs.push_back(msg); + } + } catch (const std::exception & e) { + // @ngxson : disable otherwise it's bloating the API response + // printf("%s\n", std::string("; messages = ") + messages.dump(2)); + throw std::runtime_error("Failed to parse messages: " + std::string(e.what())); + } + + return msgs; +} + +static json render_message_to_json(const std::vector & msgs, const jinja::caps & c) { + if (!c.supports_string_content && !c.supports_typed_content) { + LOG_WRN("%s: Neither string content nor typed content is supported by the template. This is unexpected and may lead to issues.\n", __func__); + } + + bool only_string_accepted = c.supports_string_content && !c.supports_typed_content; + bool only_typed_accepted = !c.supports_string_content && c.supports_typed_content; + + json messages = json::array(); + for (const auto & msg : msgs) { + if (only_string_accepted) { + json jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ true); + messages.push_back(jmsg); + } else if (only_typed_accepted) { + json jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ false); + if (jmsg.at("content").is_string()) { + jmsg["content"] = json::array({ + json{ + {"type", "text"}, + {"text", jmsg.at("content").get()}, + } + }); + } + messages.push_back(jmsg); + } else { + json jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ false); + messages.push_back(jmsg); + } + } + return messages; +} + +// DEPRECATED: only used in tests +json common_chat_msgs_to_json_oaicompat(const std::vector & msgs, bool concat_typed_text) { + jinja::caps c; + c.supports_string_content = true; + c.supports_typed_content = !concat_typed_text; + return render_message_to_json(msgs, c); +} + +json common_chat_tools_to_json_oaicompat(const std::vector & tools) { + if (tools.empty()) { + return json(); + } + + auto result = json::array(); + for (const auto & tool : tools) { + result.push_back({ + { "type", "function" }, + { "function", { + { "name", tool.name }, + { "description", tool.description }, + { "parameters", json::parse(tool.parameters) }, + }}, + }); + } + return result; +} + +std::vector common_chat_tools_parse_oaicompat(const json & tools) { + std::vector result; + + try { + if (!tools.is_null()) { + if (!tools.is_array()) { + throw std::invalid_argument("Expected 'tools' to be an array, got " + tools.dump()); + } + for (const auto & tool : tools) { + if (!tool.contains("type")) { + throw std::invalid_argument("Missing tool type: " + tool.dump()); + } + const auto & type = tool.at("type"); + if (!type.is_string() || type != "function") { + throw std::invalid_argument("Unsupported tool type: " + tool.dump()); + } + if (!tool.contains("function")) { + throw std::invalid_argument("Missing tool function: " + tool.dump()); + } + + const auto & function = tool.at("function"); + result.push_back({ + /* .name = */ function.at("name"), + /* .description = */ function.value("description", ""), + /* .parameters = */ function.value("parameters", json::object()).dump(), + }); + } + } + } catch (const std::exception & e) { + throw std::runtime_error("Failed to parse tools: " + std::string(e.what()) + "; tools = " + tools.dump(2)); + } + + return result; +} + +bool common_chat_verify_template(const std::string & tmpl, bool use_jinja) { + if (use_jinja) { + try { + common_chat_msg msg; + msg.role = "user"; + msg.content = "test"; + + auto tmpls = common_chat_templates_init(/* model= */ nullptr, tmpl); + + common_chat_templates_inputs inputs; + inputs.messages = { msg }; + + common_chat_templates_apply(tmpls.get(), inputs); + return true; + } catch (const std::exception & e) { + LOG_ERR("%s: failed to apply template: %s\n", __func__, e.what()); + return false; + } + } + llama_chat_message chat[] = { + { "user", "test" } + }; + const int res = llama_chat_apply_template(tmpl.c_str(), chat, 1, true, nullptr, 0); + return res >= 0; +} + +std::string common_chat_format_single(const struct common_chat_templates * tmpls, + const std::vector & past_msg, + const common_chat_msg & new_msg, + bool add_ass, + bool use_jinja) { + common_chat_templates_inputs inputs; + inputs.use_jinja = use_jinja; + inputs.add_bos = tmpls->add_bos; + inputs.add_eos = tmpls->add_eos; + + std::string fmt_past_msg; + if (!past_msg.empty()) { + inputs.messages = past_msg; + inputs.add_generation_prompt = false; + fmt_past_msg = common_chat_templates_apply(tmpls, inputs).prompt; + } + std::ostringstream ss; + // if the past_msg ends with a newline, we must preserve it in the formatted version + if (add_ass && !fmt_past_msg.empty() && fmt_past_msg.back() == '\n') { + ss << "\n"; + }; + // format chat with new_msg + inputs.messages.push_back(new_msg); + inputs.add_generation_prompt = add_ass; + auto fmt_new_msg = common_chat_templates_apply(tmpls, inputs).prompt; + // get the diff part + ss << fmt_new_msg.substr(fmt_past_msg.size(), fmt_new_msg.size() - fmt_past_msg.size()); + return ss.str(); +} + +std::string common_chat_format_example(const struct common_chat_templates * tmpls, + bool use_jinja, + const std::map & chat_template_kwargs) { + common_chat_templates_inputs inputs; + inputs.use_jinja = use_jinja; + inputs.add_bos = tmpls->add_bos; + inputs.add_eos = tmpls->add_eos; + inputs.chat_template_kwargs = chat_template_kwargs; + auto add_simple_msg = [&](auto role, auto content) { + common_chat_msg msg; + msg.role = role; + msg.content = content; + inputs.messages.push_back(msg); + }; + add_simple_msg("system", "You are a helpful assistant"); + add_simple_msg("user", "Hello"); + add_simple_msg("assistant", "Hi there"); + add_simple_msg("user", "How are you?"); + return common_chat_templates_apply(tmpls, inputs).prompt; +} + +#undef FILENAME +#define FILENAME "jinja-lexer" + +namespace jinja { + +static void string_lstrip(std::string & s, const char * chars) { + llama_common_unicode_string stripped = + llama_common_string_lstrip_chars_rust((const uint8_t *) s.data(), s.size(), chars); + s = common_rust_string_to_std(stripped); +} + +static void string_rstrip(std::string & s, const char * chars) { + llama_common_unicode_string stripped = + llama_common_string_rstrip_chars_rust((const uint8_t *) s.data(), s.size(), chars); + s = common_rust_string_to_std(stripped); +} + +lexer_result lexer::tokenize(const std::string & source) { + std::vector tokens; + + // NOTE: do NOT transform the source string (i.e. preprocessing), as we need to keep + // the original character positions for error reporting etc. + std::string src = source; + + if (source.empty()) { + return {tokens, src}; + } + + // Normalize \r\n or \r to \n + for (std::string::size_type pos = 0; (pos = src.find("\r\n", pos)) != std::string::npos; ) { + src.erase(pos, 1); + ++pos; + } + for (std::string::size_type pos = 0; (pos = src.find("\r", pos)) != std::string::npos; ) { + src.replace(pos, 1, 1, '\n'); + ++pos; + } + + // In the default configuration: + // - a single trailing newline is stripped if present + // - other whitespace (spaces, tabs, newlines etc.) is returned unchanged + if (source.back() == '\n') { + src.pop_back(); + } + + size_t pos = 0; + size_t start_pos = 0; + size_t curly_bracket_depth = 0; + + using pred = std::function; + auto consume_while = [&](const pred & predicate) -> std::string { + std::string str; + while (predicate(src[pos])) { + // check for escape char + if (src[pos] == '\\') { + // consume backslash + ++pos; + // check for end of input + if (pos >= src.size()) { + throw lexer_exception("unexpected end of input after escape character", source, pos); + } + // add escaped char + char escaped_char = src[pos++]; + if (escape_chars.find(escaped_char) == escape_chars.end()) { + throw lexer_exception(std::string("unknown escape character \\") + escaped_char, source, pos); + } + char unescaped_char = escape_chars.at(escaped_char); + str += unescaped_char; + continue; + } + + str += src[pos++]; + if (pos > src.size()) { + throw lexer_exception("unexpected end of input during consume_while", source, pos); + } + } + return str; + }; + + auto consume_numeric = [&]() -> std::string { + std::string num = consume_while([](char c) { + return llama_common_jinja_is_integer_rust((uint8_t) c); + }); + if (pos < src.size() && src[pos] == '.' && pos + 1 < src.size() && + llama_common_jinja_is_integer_rust((uint8_t) src[pos + 1])) { + ++pos; // Consume '.' + std::string frac = consume_while([](char c) { + return llama_common_jinja_is_integer_rust((uint8_t) c); + }); + num += "." + frac; + } + return num; + }; + + auto next_pos_is = [&](std::initializer_list chars, size_t n = 1) -> bool { + if (pos + n >= src.size()) return false; + for (char c : chars) { + if (src[pos + n] == c) return true; + } + return false; + }; + + // note: default config for chat template: lstrip_blocks = true, trim_blocks = true + + // text\n[space]{block} --> text\n{block} + bool opt_lstrip_blocks = true; + + // {block}\n[space]text --> {block}[space]text + bool opt_trim_blocks = true; + + // options set dynamically based on current/last block + bool is_lstrip_block = false; // example: {%- + bool is_rstrip_block = false; // example: -%} + + while (pos < src.size()) { + start_pos = pos; + // JJ_DEBUG("lexer main loop at pos %zu: '%s...'", pos, src.substr(pos, 10).c_str()); + + // First, consume all text that is outside of a Jinja statement or expression + token::type last_token_type = tokens.empty() + ? token::close_statement // initial state + : tokens.back().t; + if (last_token_type == token::close_statement || + last_token_type == token::close_expression || + last_token_type == token::comment) { + + bool last_block_can_rm_newline = false; + is_rstrip_block = false; + if (pos > 3) { + char c0 = src[pos - 3]; + char c1 = src[pos - 2]; + char c2 = src[pos - 1]; + // strip if: -[%}#]}text + is_rstrip_block = c0 == '-' + && (c1 == '%' || c1 == '}' || c1 == '#') + && c2 == '}'; + // match behavior of hf.js: exclude {{ and }} cases, regex: ([#%-]}) + last_block_can_rm_newline = (c1 == '#' || c1 == '%' || c1 == '-') && c2 == '}'; + } + + size_t start = pos; + size_t end = start; + while (pos < src.size() && + // Keep going until we hit the next Jinja statement or expression + !( + src[pos] == '{' && + next_pos_is( {'%', '{', '#'} ) + )) { + end = ++pos; + } + + // equivalent to hf.js code: template.replace(/^[ \t]*({[#%-])/gm, "$1"); + if (opt_lstrip_blocks && src[pos] == '{' && next_pos_is({'%', '#', '-'})) { + size_t current = end; + while (current > start) { + char c = src[current - 1]; + if (current == 1) { + end = 0; // Trim from the start of the string + break; + } + if (c == '\n') { + end = current; // Trim from the start of the line + break; + } + if (!std::isspace(static_cast(c))) { + break; // Found non-whitespace before newline, keep + } + --current; + } + } + + std::string text = src.substr(start, end - start); + + // equivalent to hf.js code: template.replace(/([#%-]})\n/g, "$1"); + if (opt_trim_blocks && last_block_can_rm_newline) { + if (!text.empty() && text.front() == '\n') { + text.erase(text.begin()); + } + } + + if (is_rstrip_block) { + // example: {last_block}[space]text + // doing lstrip on text, effectively rstrip the LAST block + // JJ_DEBUG("RSTRIP block detected, current text: '%s'", text.c_str()); + string_lstrip(text, " \t\r\n"); + } + + is_lstrip_block = src[pos] == '{' && next_pos_is({'{', '%', '#'}) && next_pos_is({'-'}, 2); + if (is_lstrip_block) { + // example: text[space]{current_block} + // doing rstrip on text, effectively lstrip the CURRENT block + // JJ_DEBUG("LSTRIP block detected, current text: '%s'", text.c_str()); + string_rstrip(text, " \t\r\n"); + } + + if (!text.empty()) { + // JJ_DEBUG("consumed text: '%s'", text.c_str()); + tokens.push_back({token::text, text, start_pos}); + continue; + } + } + + // Possibly consume a comment + // TODO: handle lstrip/rstrip for comments? (not important for now) + if (src[pos] == '{' && next_pos_is( {'#'} )) { + start_pos = pos; + pos += 2; // Skip the opening {# + std::string comment; + while (!(src[pos] == '#' && next_pos_is( {'}'} ))) { + if (pos + 2 >= src.size()) { + throw lexer_exception("missing end of comment tag", source, pos); + } + comment += src[pos++]; + } + JJ_DEBUG("consumed comment: '%s'", comment.c_str()); + tokens.push_back({token::comment, comment, start_pos}); + pos += 2; // Skip the closing #} + continue; + } + + if (src[pos] == '-' && ( + last_token_type == token::open_expression || + last_token_type == token::open_statement) + ) { + JJ_DEBUG("lexer main loop at pos %zu: '%s...'", pos, src.substr(pos, 10).c_str()); + pos++; // consume '-' in {%- or {{- + if (pos >= src.size()) break; + } + + // Consume (and ignore) all whitespace inside Jinja statements or expressions + consume_while([](char c) { return std::isspace(static_cast(c)); }); + + if (pos >= src.size()) break; + + char ch = src[pos]; + + bool is_closing_block = ch == '-' && next_pos_is( {'%', '}'} ); + + // Check for unary operators + if (!is_closing_block && (ch == '-' || ch == '+')) { + start_pos = pos; + token::type last_token_type = tokens.empty() ? token::eof : tokens.back().t; + if (last_token_type == token::text || last_token_type == token::eof) { + throw lexer_exception(std::string("unexpected character: ") + ch, source, pos); + } + switch (last_token_type) { + case token::identifier: + case token::numeric_literal: + case token::string_literal: + case token::close_paren: + case token::close_square_bracket: + // Part of a binary operator + // a - 1, 1 - 1, true - 1, "apple" - 1, (1) - 1, a[1] - 1 + // Continue parsing normally + break; + default: { + // Is part of a unary operator + // (-1), [-1], (1 + -1), not -1, -apple + ++pos; // Consume the operator + + // Check for numbers following the unary operator + std::string num = consume_numeric(); + std::string value = std::string(1, ch) + num; + token::type t = num.empty() ? token::unary_operator : token::numeric_literal; + // JJ_DEBUG("consumed unary operator or numeric literal: '%s'", value.c_str()); + tokens.push_back({t, value, start_pos}); + continue; + } + } + } + + // Try to match one of the tokens in the mapping table + bool matched = false; + for (const auto & [seq, typ] : ordered_mapping_table) { + start_pos = pos; + // Inside an object literal, don't treat "}}" as expression-end + if (seq == "}}" && curly_bracket_depth > 0) { + continue; + } + if (pos + seq.size() <= src.size() && src.substr(pos, seq.size()) == seq) { + tokens.push_back({typ, seq, start_pos}); + if (typ == token::open_expression) { + curly_bracket_depth = 0; + } else if (typ == token::open_curly_bracket) { + ++curly_bracket_depth; + } else if (typ == token::close_curly_bracket) { + --curly_bracket_depth; + } + + pos += seq.size(); + matched = true; + break; // continue main loop + } + } + if (matched) continue; // continue main loop + + // Strings + if (ch == '\'' || ch == '"') { + start_pos = pos; + ++pos; // Skip opening quote + std::string str = consume_while([ch](char c) { return c != ch; }); + // JJ_DEBUG("consumed string literal: '%s'", str.c_str()); + tokens.push_back({token::string_literal, str, start_pos}); + ++pos; // Skip closing quote + continue; + } + + // Numbers + if (llama_common_jinja_is_integer_rust((uint8_t) ch)) { + start_pos = pos; + std::string num = consume_numeric(); + // JJ_DEBUG("consumed numeric literal: '%s'", num.c_str()); + tokens.push_back({token::numeric_literal, num, start_pos}); + continue; + } + + // Identifiers + if (llama_common_jinja_is_word_rust((uint8_t) ch)) { + start_pos = pos; + std::string word = consume_while([](char c) { + return llama_common_jinja_is_word_rust((uint8_t) c); + }); + // JJ_DEBUG("consumed identifier: '%s'", word.c_str()); + tokens.push_back({token::identifier, word, start_pos}); + continue; + } + + throw lexer_exception(std::string("unexpected character: ") + ch, source, pos); + } + + return {std::move(tokens), src}; +} + +} // namespace jinja + +#include "common.h.inc" + +#include +#include +#include +#include +#include + +#undef FILENAME +#define FILENAME "jinja-parser" + +namespace jinja { + +// parse from a list of tokens into an AST (program) +program parse_from_tokens(const lexer_result & lexer_res); + +struct parser_exception : public std::runtime_error { + parser_exception(const std::string & msg, const std::string & source, size_t pos) + : std::runtime_error([&]() { + const char * tag = "parser"; + llama_common_unicode_string result = llama_common_fmt_error_with_source_rust( + reinterpret_cast(tag), std::strlen(tag), + reinterpret_cast(msg.data()), msg.size(), + reinterpret_cast(source.data()), source.size(), pos); + std::string value(reinterpret_cast(result.data), result.len); + llama_common_unicode_string_free(result); + return value; + }()) {} +}; + +// Helper to check type without asserting (useful for logic) +template +static bool is_type(const statement_ptr & ptr) { + return dynamic_cast(ptr.get()) != nullptr; +} + +class parser { + const std::vector & tokens; + size_t current = 0; + + std::string source; // for error reporting + +public: + parser(const std::vector & t, const std::string & src) : tokens(t), source(src) {} + + program parse() { + statements body; + while (current < tokens.size()) { + body.push_back(parse_any()); + } + return program(std::move(body)); + } + + // NOTE: start_pos is the token index, used for error reporting + template + std::unique_ptr mk_stmt(size_t start_pos, Args&&... args) { + auto ptr = std::make_unique(std::forward(args)...); + assert(start_pos < tokens.size()); + ptr->pos = tokens[start_pos].pos; + return ptr; + } + +private: + const token & peek(size_t offset = 0) const { + if (current + offset >= tokens.size()) { + static const token end_token{token::eof, "", 0}; + return end_token; + } + return tokens[current + offset]; + } + + const token & next() { + if (current >= tokens.size()) { + throw parser_exception("Parser Error: Unexpected EOF", source, tokens.empty() ? 0 : tokens.back().pos); + } + return tokens[current++]; + } + + token expect(token::type type, const std::string& error) { + const auto & t = peek(); + if (t.t != type) { + throw parser_exception("Parser Error: " + error + " (Got " + t.value + ")", source, t.pos); + } + current++; + return t; + } + + void expect_identifier(const std::string & name) { + const auto & t = peek(); + if (t.t != token::identifier || t.value != name) { + throw parser_exception("Expected identifier: " + name, source, t.pos); + } + current++; + } + + bool is(token::type type) const { + return peek().t == type; + } + + bool is_identifier(const std::string & name) const { + return peek().t == token::identifier && peek().value == name; + } + + bool is_statement(const std::vector & names) const { + if (peek(0).t != token::open_statement || peek(1).t != token::identifier) { + return false; + } + std::string val = peek(1).value; + return std::find(names.begin(), names.end(), val) != names.end(); + } + + statement_ptr parse_any() { + size_t start_pos = current; + switch (peek().t) { + case token::comment: + return mk_stmt(start_pos, next().value); + case token::text: + return mk_stmt(start_pos, next().value); + case token::open_statement: + return parse_jinja_statement(); + case token::open_expression: + return parse_jinja_expression(); + default: + throw std::runtime_error("Unexpected token type"); + } + } + + statement_ptr parse_jinja_expression() { + // Consume {{ }} tokens + expect(token::open_expression, "Expected {{"); + auto result = parse_expression(); + expect(token::close_expression, "Expected }}"); + return result; + } + + statement_ptr parse_jinja_statement() { + // Consume {% token + expect(token::open_statement, "Expected {%"); + + if (peek().t != token::identifier) { + throw std::runtime_error("Unknown statement"); + } + + size_t start_pos = current; + std::string name = next().value; + + statement_ptr result; + if (name == "set") { + result = parse_set_statement(start_pos); + + } else if (name == "if") { + result = parse_if_statement(start_pos); + // expect {% endif %} + expect(token::open_statement, "Expected {%"); + expect_identifier("endif"); + expect(token::close_statement, "Expected %}"); + + } else if (name == "macro") { + result = parse_macro_statement(start_pos); + // expect {% endmacro %} + expect(token::open_statement, "Expected {%"); + expect_identifier("endmacro"); + expect(token::close_statement, "Expected %}"); + + } else if (name == "for") { + result = parse_for_statement(start_pos); + // expect {% endfor %} + expect(token::open_statement, "Expected {%"); + expect_identifier("endfor"); + expect(token::close_statement, "Expected %}"); + + } else if (name == "break") { + expect(token::close_statement, "Expected %}"); + result = mk_stmt(start_pos); + + } else if (name == "continue") { + expect(token::close_statement, "Expected %}"); + result = mk_stmt(start_pos); + + } else if (name == "call") { + statements caller_args; + // bool has_caller_args = false; + if (is(token::open_paren)) { + // Optional caller arguments, e.g. {% call(user) dump_users(...) %} + caller_args = parse_args(); + // has_caller_args = true; + } + auto callee = parse_primary_expression(); + if (!is_type(callee)) throw std::runtime_error("Expected identifier"); + + auto call_args = parse_args(); + expect(token::close_statement, "Expected %}"); + + statements body; + while (!is_statement({"endcall"})) { + body.push_back(parse_any()); + } + + expect(token::open_statement, "Expected {%"); + expect_identifier("endcall"); + expect(token::close_statement, "Expected %}"); + + auto call_expr = mk_stmt(start_pos, std::move(callee), std::move(call_args)); + result = mk_stmt(start_pos, std::move(call_expr), std::move(caller_args), std::move(body)); + + } else if (name == "filter") { + auto filter_node = parse_primary_expression(); + if (is_type(filter_node) && is(token::open_paren)) { + filter_node = parse_call_expression(std::move(filter_node)); + } + expect(token::close_statement, "Expected %}"); + + statements body; + while (!is_statement({"endfilter"})) { + body.push_back(parse_any()); + } + + expect(token::open_statement, "Expected {%"); + expect_identifier("endfilter"); + expect(token::close_statement, "Expected %}"); + result = mk_stmt(start_pos, std::move(filter_node), std::move(body)); + + } else if (name == "generation" || name == "endgeneration") { + // Ignore generation blocks (transformers-specific) + // See https://github.com/huggingface/transformers/pull/30650 for more information. + result = mk_stmt(start_pos); + ++current; + + } else { + throw std::runtime_error("Unknown statement: " + name); + } + return result; + } + + statement_ptr parse_set_statement(size_t start_pos) { + // NOTE: `set` acts as both declaration statement and assignment expression + auto left = parse_expression_sequence(); + statement_ptr value = nullptr; + statements body; + + if (is(token::equals)) { + ++current; + value = parse_expression_sequence(); + } else { + // parsing multiline set here + expect(token::close_statement, "Expected %}"); + while (!is_statement({"endset"})) { + body.push_back(parse_any()); + } + expect(token::open_statement, "Expected {%"); + expect_identifier("endset"); + } + expect(token::close_statement, "Expected %}"); + return mk_stmt(start_pos, std::move(left), std::move(value), std::move(body)); + } + + statement_ptr parse_if_statement(size_t start_pos) { + auto test = parse_expression(); + expect(token::close_statement, "Expected %}"); + + statements body; + statements alternate; + + // Keep parsing 'if' body until we reach the first {% elif %} or {% else %} or {% endif %} + while (!is_statement({"elif", "else", "endif"})) { + body.push_back(parse_any()); + } + + if (is_statement({"elif"})) { + size_t pos0 = current; + ++current; // consume {% + ++current; // consume 'elif' + alternate.push_back(parse_if_statement(pos0)); // nested If + } else if (is_statement({"else"})) { + ++current; // consume {% + ++current; // consume 'else' + expect(token::close_statement, "Expected %}"); + + // keep going until we hit {% endif %} + while (!is_statement({"endif"})) { + alternate.push_back(parse_any()); + } + } + return mk_stmt(start_pos, std::move(test), std::move(body), std::move(alternate)); + } + + statement_ptr parse_macro_statement(size_t start_pos) { + auto name = parse_primary_expression(); + auto args = parse_args(); + expect(token::close_statement, "Expected %}"); + statements body; + // Keep going until we hit {% endmacro + while (!is_statement({"endmacro"})) { + body.push_back(parse_any()); + } + return mk_stmt(start_pos, std::move(name), std::move(args), std::move(body)); + } + + statement_ptr parse_expression_sequence(bool primary = false) { + size_t start_pos = current; + statements exprs; + exprs.push_back(primary ? parse_primary_expression() : parse_expression()); + bool is_tuple = is(token::comma); + while (is(token::comma)) { + ++current; // consume comma + exprs.push_back(primary ? parse_primary_expression() : parse_expression()); + } + return is_tuple ? mk_stmt(start_pos, std::move(exprs)) : std::move(exprs[0]); + } + + statement_ptr parse_for_statement(size_t start_pos) { + // e.g., `message` in `for message in messages` + auto loop_var = parse_expression_sequence(true); // should be an identifier/tuple + if (!is_identifier("in")) throw std::runtime_error("Expected 'in'"); + ++current; // consume 'in' + + // `messages` in `for message in messages` + auto iterable = parse_expression(); + expect(token::close_statement, "Expected %}"); + + statements body; + statements alternate; + + // Keep going until we hit {% endfor or {% else + while (!is_statement({"endfor", "else"})) { + body.push_back(parse_any()); + } + + if (is_statement({"else"})) { + ++current; // consume {% + ++current; // consume 'else' + expect(token::close_statement, "Expected %}"); + while (!is_statement({"endfor"})) { + alternate.push_back(parse_any()); + } + } + return mk_stmt( + start_pos, + std::move(loop_var), std::move(iterable), + std::move(body), std::move(alternate)); + } + + statement_ptr parse_expression() { + // Choose parse function with lowest precedence + return parse_if_expression(); + } + + statement_ptr parse_if_expression() { + auto a = parse_logical_or_expression(); + if (is_identifier("if")) { + // Ternary expression + size_t start_pos = current; + ++current; // consume 'if' + auto test = parse_logical_or_expression(); + if (is_identifier("else")) { + // Ternary expression with else + size_t pos0 = current; + ++current; // consume 'else' + auto false_expr = parse_if_expression(); // recurse to support chained ternaries + return mk_stmt(pos0, std::move(test), std::move(a), std::move(false_expr)); + } else { + // Select expression on iterable + return mk_stmt(start_pos, std::move(a), std::move(test)); + } + } + return a; + } + + statement_ptr parse_logical_or_expression() { + auto left = parse_logical_and_expression(); + while (is_identifier("or")) { + size_t start_pos = current; + token op = next(); + left = mk_stmt(start_pos, op, std::move(left), parse_logical_and_expression()); + } + return left; + } + + statement_ptr parse_logical_and_expression() { + auto left = parse_logical_negation_expression(); + while (is_identifier("and")) { + size_t start_pos = current; + auto op = next(); + left = mk_stmt(start_pos, op, std::move(left), parse_logical_negation_expression()); + } + return left; + } + + statement_ptr parse_logical_negation_expression() { + // Try parse unary operators + if (is_identifier("not")) { + size_t start_pos = current; + auto op = next(); + return mk_stmt(start_pos, op, parse_logical_negation_expression()); + } + return parse_comparison_expression(); + } + + statement_ptr parse_comparison_expression() { + // NOTE: membership has same precedence as comparison + // e.g., ('a' in 'apple' == 'b' in 'banana') evaluates as ('a' in ('apple' == ('b' in 'banana'))) + auto left = parse_additive_expression(); + while (true) { + token op; + size_t start_pos = current; + if (is_identifier("not") && peek(1).t == token::identifier && peek(1).value == "in") { + op = {token::identifier, "not in", tokens[current].pos}; + ++current; // consume 'not' + ++current; // consume 'in' + } else if (is_identifier("in")) { + op = next(); + } else if (is(token::comparison_binary_operator)) { + op = next(); + } else break; + left = mk_stmt(start_pos, op, std::move(left), parse_additive_expression()); + } + return left; + } + + statement_ptr parse_additive_expression() { + auto left = parse_multiplicative_expression(); + while (is(token::additive_binary_operator)) { + size_t start_pos = current; + auto op = next(); + left = mk_stmt(start_pos, op, std::move(left), parse_multiplicative_expression()); + } + return left; + } + + statement_ptr parse_multiplicative_expression() { + auto left = parse_test_expression(); + while (is(token::multiplicative_binary_operator)) { + size_t start_pos = current; + auto op = next(); + left = mk_stmt(start_pos, op, std::move(left), parse_test_expression()); + } + return left; + } + + statement_ptr parse_test_expression() { + auto operand = parse_filter_expression(); + while (is_identifier("is")) { + size_t start_pos = current; + ++current; // consume 'is' + bool negate = false; + if (is_identifier("not")) { ++current; negate = true; } + auto test_id = parse_primary_expression(); + // FIXME: tests can also be expressed like this: if x is eq 3 + if (is(token::open_paren)) test_id = parse_call_expression(std::move(test_id)); + operand = mk_stmt(start_pos, std::move(operand), negate, std::move(test_id)); + } + return operand; + } + + statement_ptr parse_filter_expression() { + auto operand = parse_call_member_expression(); + while (is(token::pipe)) { + size_t start_pos = current; + ++current; // consume pipe + auto filter = parse_primary_expression(); + if (is(token::open_paren)) filter = parse_call_expression(std::move(filter)); + operand = mk_stmt(start_pos, std::move(operand), std::move(filter)); + } + return operand; + } + + statement_ptr parse_call_member_expression() { + // Handle member expressions recursively + auto member = parse_member_expression(parse_primary_expression()); + return is(token::open_paren) + ? parse_call_expression(std::move(member)) // foo.x() + : std::move(member); + } + + statement_ptr parse_call_expression(statement_ptr callee) { + size_t start_pos = current; + auto expr = mk_stmt(start_pos, std::move(callee), parse_args()); + auto member = parse_member_expression(std::move(expr)); // foo.x().y + return is(token::open_paren) + ? parse_call_expression(std::move(member)) // foo.x()() + : std::move(member); + } + + statements parse_args() { + // comma-separated arguments list + expect(token::open_paren, "Expected ("); + statements args; + while (!is(token::close_paren)) { + statement_ptr arg; + // unpacking: *expr + if (peek().t == token::multiplicative_binary_operator && peek().value == "*") { + size_t start_pos = current; + ++current; // consume * + arg = mk_stmt(start_pos, parse_expression()); + } else { + arg = parse_expression(); + if (is(token::equals)) { + // keyword argument + // e.g., func(x = 5, y = a or b) + size_t start_pos = current; + ++current; // consume equals + arg = mk_stmt(start_pos, std::move(arg), parse_expression()); + } + } + args.push_back(std::move(arg)); + if (is(token::comma)) { + ++current; // consume comma + } + } + expect(token::close_paren, "Expected )"); + return args; + } + + statement_ptr parse_member_expression(statement_ptr object) { + size_t start_pos = current; + while (is(token::dot) || is(token::open_square_bracket)) { + auto op = next(); + bool computed = op.t == token::open_square_bracket; + statement_ptr prop; + if (computed) { + prop = parse_member_expression_arguments(); + expect(token::close_square_bracket, "Expected ]"); + } else { + prop = parse_primary_expression(); + } + object = mk_stmt(start_pos, std::move(object), std::move(prop), computed); + } + return object; + } + + statement_ptr parse_member_expression_arguments() { + // NOTE: This also handles slice expressions colon-separated arguments list + // e.g., ['test'], [0], [:2], [1:], [1:2], [1:2:3] + statements slices; + bool is_slice = false; + size_t start_pos = current; + while (!is(token::close_square_bracket)) { + if (is(token::colon)) { + // A case where a default is used + // e.g., [:2] will be parsed as [undefined, 2] + slices.push_back(nullptr); + ++current; // consume colon + is_slice = true; + } else { + slices.push_back(parse_expression()); + if (is(token::colon)) { + ++current; // consume colon after expression, if it exists + is_slice = true; + } + } + } + if (is_slice) { + statement_ptr start = slices.size() > 0 ? std::move(slices[0]) : nullptr; + statement_ptr stop = slices.size() > 1 ? std::move(slices[1]) : nullptr; + statement_ptr step = slices.size() > 2 ? std::move(slices[2]) : nullptr; + return mk_stmt(start_pos, std::move(start), std::move(stop), std::move(step)); + } + if (slices.empty()) { + return mk_stmt(start_pos); + } + return std::move(slices[0]); + } + + statement_ptr parse_primary_expression() { + size_t start_pos = current; + auto t = next(); + switch (t.t) { + case token::numeric_literal: + if (t.value.find('.') != std::string::npos) { + return mk_stmt(start_pos, std::stod(t.value)); + } else { + return mk_stmt(start_pos, std::stoll(t.value)); + } + case token::string_literal: { + std::string val = t.value; + while (is(token::string_literal)) { + val += next().value; + } + return mk_stmt(start_pos, val); + } + case token::identifier: + return mk_stmt(start_pos, t.value); + case token::open_paren: { + auto expr = parse_expression_sequence(); + expect(token::close_paren, "Expected )"); + return expr; + } + case token::open_square_bracket: { + statements vals; + while (!is(token::close_square_bracket)) { + vals.push_back(parse_expression()); + if (is(token::comma)) ++current; + } + ++current; + return mk_stmt(start_pos, std::move(vals)); + } + case token::open_curly_bracket: { + std::vector> pairs; + while (!is(token::close_curly_bracket)) { + auto key = parse_expression(); + expect(token::colon, "Expected :"); + pairs.push_back({std::move(key), parse_expression()}); + if (is(token::comma)) ++current; + } + ++current; + return mk_stmt(start_pos, std::move(pairs)); + } + default: + throw std::runtime_error("Unexpected token: " + t.value + " of type " + std::to_string(t.t)); + } + } +}; + +program parse_from_tokens(const lexer_result & lexer_res) { + return parser(lexer_res.tokens, lexer_res.source).parse(); +} + +} // namespace jinja + +#include "common.h.inc" + +#include +#include +#include +#include + +#undef FILENAME +#define FILENAME "jinja-runtime" + +bool g_jinja_debug = false; + +namespace jinja { + +void enable_debug(bool enable) { + g_jinja_debug = enable; +} + +static value_string exec_statements(const statements & stmts, context & ctx) { + auto result = mk_val(); + for (const auto & stmt : stmts) { + JJ_DEBUG("Executing statement of type %s", stmt->type().c_str()); + result->push_back(stmt->execute(ctx)); + } + // convert to string parts + value_string str = mk_val(); + gather_string_parts_recursive(result, str); + return str; +} + +static void ensure_key_type_allowed(const value & val) { + if (!val->is_hashable()) { + throw std::runtime_error("Type: " + val->type() + " is not allowed as object key"); + } +} + +// execute with error handling +value statement::execute(context & ctx) { + try { + return execute_impl(ctx); + } catch (const continue_statement::signal & /* ex */) { + throw; + } catch (const break_statement::signal & /* ex */) { + throw; + } catch (const rethrown_exception & /* ex */) { + throw; + } catch (const not_implemented_exception & /* ex */) { + throw; + } catch (const std::exception & e) { + const std::string & source = *ctx.src; + if (source.empty()) { + std::ostringstream oss; + oss << "\nError executing " << type() << " at position " << pos << ": " << e.what(); + throw rethrown_exception(oss.str()); + } else { + std::ostringstream oss; + oss << "\n------------\n"; + oss << "While executing " << type() << " at " + << common_rust_string_to_std(llama_common_get_line_col_rust( + (const uint8_t *) source.data(), source.size(), pos)) + << " in source:\n"; + llama_common_unicode_string peak = llama_common_peak_source_rust( + reinterpret_cast(source.data()), source.size(), pos, 40); + oss << std::string(reinterpret_cast(peak.data), peak.len) << "\n"; + llama_common_unicode_string_free(peak); + oss << "Error: " << e.what(); + // throw as another exception to avoid repeated formatting + throw rethrown_exception(oss.str()); + } + } +} + +value identifier::execute_impl(context & ctx) { + auto it = ctx.get_val(val); + auto builtins = global_builtins(); + if (!it->is_undefined()) { + if (ctx.is_get_stats) { + value_t::stats_t::mark_used(it); + } + JJ_DEBUG("Identifier '%s' found, type = %s", val.c_str(), it->type().c_str()); + return it; + } else if (builtins.find(val) != builtins.end()) { + JJ_DEBUG("Identifier '%s' found in builtins", val.c_str()); + return mk_val(val, builtins.at(val)); + } else { + JJ_DEBUG("Identifier '%s' not found, returning undefined", val.c_str()); + return mk_val(val); + } +} + +value object_literal::execute_impl(context & ctx) { + auto obj = mk_val(); + for (const auto & pair : val) { + value key = pair.first->execute(ctx); + value val = pair.second->execute(ctx); + JJ_DEBUG("Object literal: setting key '%s' with value type %s", key->as_string().str().c_str(), val->type().c_str()); + obj->insert(key, val); + } + return obj; +} + +value binary_expression::execute_impl(context & ctx) { + value left_val = left->execute(ctx); + + // Logical operators + if (op.value == "and") { + JJ_DEBUG("Executing logical test: %s AND %s", left->type().c_str(), right->type().c_str()); + return left_val->as_bool() ? right->execute(ctx) : std::move(left_val); + } else if (op.value == "or") { + JJ_DEBUG("Executing logical test: %s OR %s", left->type().c_str(), right->type().c_str()); + return left_val->as_bool() ? std::move(left_val) : right->execute(ctx); + } + + // Equality operators + value right_val = right->execute(ctx); + JJ_DEBUG("Executing binary expression %s '%s' %s", left_val->type().c_str(), op.value.c_str(), right_val->type().c_str()); + if (op.value == "==") { + return mk_val(*left_val == *right_val); + } else if (op.value == "!=") { + return mk_val(!(*left_val == *right_val)); + } + + auto workaround_concat_null_with_str = [&](value & res) -> bool { + bool is_left_null = left_val->is_none() || left_val->is_undefined(); + bool is_right_null = right_val->is_none() || right_val->is_undefined(); + bool is_left_str = is_val(left_val); + bool is_right_str = is_val(right_val); + if ((is_left_null && is_right_str) || (is_right_null && is_left_str)) { + JJ_DEBUG("%s", "Workaround: treating null/undefined as empty string for string concatenation"); + string left_str = is_left_null ? string() : left_val->as_string(); + string right_str = is_right_null ? string() : right_val->as_string(); + auto output = left_str.append(right_str); + res = mk_val(std::move(output)); + return true; + } + return false; + }; + + auto test_is_in = [&]() -> bool { + func_args args(ctx); + args.push_back(left_val); + args.push_back(right_val); + return global_builtins().at("test_is_in")(args)->as_bool(); + }; + + // Handle undefined and null values + if (is_val(left_val) || is_val(right_val)) { + if (is_val(right_val) && (op.value == "in" || op.value == "not in")) { + // Special case: `anything in undefined` is `false` and `anything not in undefined` is `true` + return mk_val(op.value == "not in"); + } + if (op.value == "+" || op.value == "~") { + value res = mk_val(); + if (workaround_concat_null_with_str(res)) { + return res; + } + } + throw std::runtime_error("Cannot perform operation " + op.value + " on undefined values"); + } else if (is_val(left_val) || is_val(right_val)) { + if (op.value == "+" || op.value == "~") { + value res = mk_val(); + if (workaround_concat_null_with_str(res)) { + return res; + } + } + throw std::runtime_error("Cannot perform operation on null values"); + } + + // Float operations + if ((is_val(left_val) || is_val(left_val)) && + (is_val(right_val) || is_val(right_val))) { + double a = left_val->as_float(); + double b = right_val->as_float(); + if (op.value == "+" || op.value == "-" || op.value == "*") { + double res = (op.value == "+") ? a + b : (op.value == "-") ? a - b : a * b; + JJ_DEBUG("Arithmetic operation: %f %s %f = %f", a, op.value.c_str(), b, res); + bool is_float = is_val(left_val) || is_val(right_val); + if (is_float) { + return mk_val(res); + } else { + return mk_val(static_cast(res)); + } + } else if (op.value == "/") { + JJ_DEBUG("Division operation: %f / %f", a, b); + return mk_val(a / b); + } else if (op.value == "%") { + double rem = std::fmod(a, b); + JJ_DEBUG("Modulo operation: %f %% %f = %f", a, b, rem); + bool is_float = is_val(left_val) || is_val(right_val); + if (is_float) { + return mk_val(rem); + } else { + return mk_val(static_cast(rem)); + } + } else if (op.value == "<") { + JJ_DEBUG("Comparison operation: %f < %f is %d", a, b, a < b); + return mk_val(a < b); + } else if (op.value == ">") { + JJ_DEBUG("Comparison operation: %f > %f is %d", a, b, a > b); + return mk_val(a > b); + } else if (op.value == ">=") { + JJ_DEBUG("Comparison operation: %f >= %f is %d", a, b, a >= b); + return mk_val(a >= b); + } else if (op.value == "<=") { + JJ_DEBUG("Comparison operation: %f <= %f is %d", a, b, a <= b); + return mk_val(a <= b); + } + } + + // Array operations + if (is_val(left_val) && is_val(right_val)) { + if (op.value == "+") { + auto & left_arr = left_val->as_array(); + auto & right_arr = right_val->as_array(); + auto result = mk_val(); + for (const auto & item : left_arr) { + result->push_back(item); + } + for (const auto & item : right_arr) { + result->push_back(item); + } + return result; + } + } else if (is_val(right_val)) { + // case: 1 in [0, 1, 2] + bool member = test_is_in(); + if (op.value == "in") { + return mk_val(member); + } else if (op.value == "not in") { + return mk_val(!member); + } + } + + // String concatenation with ~ and + + if ((is_val(left_val) || is_val(right_val)) && + (op.value == "~" || op.value == "+")) { + JJ_DEBUG("String concatenation with %s operator", op.value.c_str()); + auto output = left_val->as_string().append(right_val->as_string()); + auto res = mk_val(); + res->val_str = std::move(output); + return res; + } + + // Python-style string repetition + // TODO: support array/tuple repetition (e.g., [1, 2] * 3 → [1, 2, 1, 2, 1, 2]) + if (op.value == "*" && + ((is_val(left_val) && is_val(right_val)) || + (is_val(left_val) && is_val(right_val)))) { + const auto & str = is_val(left_val) ? left_val->as_string() : right_val->as_string(); + const int64_t repeat = is_val(right_val) ? right_val->as_int() : left_val->as_int(); + auto res = mk_val(); + if (repeat <= 0) { + return res; + } + for (int64_t i = 0; i < repeat; ++i) { + res->val_str = res->val_str.append(str); + } + return res; + } + + // String membership + if (is_val(left_val) && is_val(right_val)) { + // case: "a" in "abc" + bool member = test_is_in(); + if (op.value == "in") { + return mk_val(member); + } else if (op.value == "not in") { + return mk_val(!member); + } + } + + // Value key in object + if (is_val(right_val)) { + // case: key in {key: value} + bool member = test_is_in(); + if (op.value == "in") { + return mk_val(member); + } else if (op.value == "not in") { + return mk_val(!member); + } + } + + throw std::runtime_error("Unknown operator \"" + op.value + "\" between " + left_val->type() + " and " + right_val->type()); +} + +static value try_builtin_func(context & ctx, const std::string & name, value & input, bool undef_on_missing = false) { + JJ_DEBUG("Trying built-in function '%s' for type %s", name.c_str(), input->type().c_str()); + if (ctx.is_get_stats) { + value_t::stats_t::mark_used(input); + input->stats.ops.insert(name); + } + auto builtins = input->get_builtins(); + auto it = builtins.find(name); + if (it != builtins.end()) { + JJ_DEBUG("Binding built-in '%s'", name.c_str()); + return mk_val(name, it->second, input); + } + if (undef_on_missing) { + return mk_val(name); + } + throw std::runtime_error("Unknown (built-in) filter '" + name + "' for type " + input->type()); +} + +value filter_expression::execute_impl(context & ctx) { + value input = operand ? operand->execute(ctx) : val; + + JJ_DEBUG("Applying filter to %s", input->type().c_str()); + + if (is_stmt(filter)) { + auto filter_id = cast_stmt(filter)->val; + + if (filter_id == "trim") { + filter_id = "strip"; // alias + } + JJ_DEBUG("Applying filter '%s' to %s", filter_id.c_str(), input->type().c_str()); + // TODO: Refactor filters so this coercion can be done automatically + if (!input->is_undefined() && !is_val(input) && ( + filter_id == "capitalize" || + filter_id == "lower" || + filter_id == "replace" || + filter_id == "strip" || + filter_id == "title" || + filter_id == "upper" || + filter_id == "wordcount" + )) { + JJ_DEBUG("Coercing %s to String for '%s' filter", input->type().c_str(), filter_id.c_str()); + input = mk_val(input->as_string()); + } + return try_builtin_func(ctx, filter_id, input)->invoke(func_args(ctx)); + + } else if (is_stmt(filter)) { + auto call = cast_stmt(filter); + if (!is_stmt(call->callee)) { + throw std::runtime_error("Filter callee must be an identifier"); + } + auto filter_id = cast_stmt(call->callee)->val; + + if (filter_id == "trim") { + filter_id = "strip"; // alias + } + JJ_DEBUG("Applying filter '%s' with arguments to %s", filter_id.c_str(), input->type().c_str()); + func_args args(ctx); + for (const auto & arg_expr : call->args) { + args.push_back(arg_expr->execute(ctx)); + } + + return try_builtin_func(ctx, filter_id, input)->invoke(args); + + } else { + throw std::runtime_error("Invalid filter expression"); + } +} + +value filter_statement::execute_impl(context & ctx) { + // eval body as string, then apply filter + auto body_val = exec_statements(body, ctx); + value_string parts = mk_val(); + gather_string_parts_recursive(body_val, parts); + + JJ_DEBUG("FilterStatement: applying filter to body string of length %zu", parts->val_str.length()); + filter_expression filter_expr(std::move(parts), std::move(filter)); + value out = filter_expr.execute(ctx); + + // this node can be reused later, make sure filter is preserved + this->filter = std::move(filter_expr.filter); + return out; +} + +value test_expression::execute_impl(context & ctx) { + // NOTE: "value is something" translates to function call "test_is_something(value)" + const auto & builtins = global_builtins(); + + std::string test_id; + value input = operand->execute(ctx); + + func_args args(ctx); + args.push_back(input); + + if (is_stmt(test)) { + test_id = cast_stmt(test)->val; + } else if (is_stmt(test)) { + auto call = cast_stmt(test); + if (!is_stmt(call->callee)) { + throw std::runtime_error("Test callee must be an identifier"); + } + test_id = cast_stmt(call->callee)->val; + + JJ_DEBUG("Applying test '%s' with arguments to %s", test_id.c_str(), input->type().c_str()); + for (const auto & arg_expr : call->args) { + args.push_back(arg_expr->execute(ctx)); + } + + } else { + throw std::runtime_error("Invalid test expression"); + } + + auto it = builtins.find("test_is_" + test_id); + JJ_DEBUG("Test expression %s '%s' %s (using function 'test_is_%s')", operand->type().c_str(), test_id.c_str(), negate ? "(negate)" : "", test_id.c_str()); + if (it == builtins.end()) { + throw std::runtime_error("Unknown test '" + test_id + "'"); + } + + auto res = it->second(args); + + if (negate) { + return mk_val(!res->as_bool()); + } else { + return res; + } +} + +value unary_expression::execute_impl(context & ctx) { + value operand_val = argument->execute(ctx); + JJ_DEBUG("Executing unary expression with operator '%s'", op.value.c_str()); + + if (op.value == "not") { + return mk_val(!operand_val->as_bool()); + } else if (op.value == "-") { + if (is_val(operand_val)) { + return mk_val(-operand_val->as_int()); + } else if (is_val(operand_val)) { + return mk_val(-operand_val->as_float()); + } else { + throw std::runtime_error("Unary - operator requires numeric operand"); + } + } + + throw std::runtime_error("Unknown unary operator '" + op.value + "'"); +} + +value if_statement::execute_impl(context & ctx) { + value test_val = test->execute(ctx); + + auto out = mk_val(); + if (test_val->as_bool()) { + for (auto & stmt : body) { + JJ_DEBUG("IF --> Executing THEN body, current block: %s", stmt->type().c_str()); + out->push_back(stmt->execute(ctx)); + } + } else { + for (auto & stmt : alternate) { + JJ_DEBUG("IF --> Executing ELSE body, current block: %s", stmt->type().c_str()); + out->push_back(stmt->execute(ctx)); + } + } + // convert to string parts + value_string str = mk_val(); + gather_string_parts_recursive(out, str); + return str; +} + +value for_statement::execute_impl(context & ctx) { + context scope(ctx); // new scope for loop variables + + jinja::select_expression * select_expr = cast_stmt(iterable); + statement_ptr test_expr_nullptr; + + statement_ptr & iter_expr = [&]() -> statement_ptr & { + auto tmp = cast_stmt(iterable); + return tmp ? tmp->lhs : iterable; + }(); + statement_ptr & test_expr = [&]() -> statement_ptr & { + auto tmp = cast_stmt(iterable); + return tmp ? tmp->test : test_expr_nullptr; + }(); + + JJ_DEBUG("Executing for statement, iterable type: %s", iter_expr->type().c_str()); + + value iterable_val = iter_expr->execute(scope); + + // mark the variable being iterated as used for stats + if (ctx.is_get_stats) { + value_t::stats_t::mark_used(iterable_val); + iterable_val->stats.ops.insert("array_access"); + } + + if (iterable_val->is_undefined()) { + JJ_DEBUG("%s", "For loop iterable is undefined, skipping loop"); + iterable_val = mk_val(); + } + + if (!is_val(iterable_val) && !is_val(iterable_val)) { + throw std::runtime_error("Expected iterable or object type in for loop: got " + iterable_val->type()); + } + + std::vector items; + if (is_val(iterable_val)) { + JJ_DEBUG("%s", "For loop over object keys"); + auto & obj = iterable_val->as_ordered_object(); + for (auto & p : obj) { + auto tuple = mk_val(p); + items.push_back(std::move(tuple)); + } + if (ctx.is_get_stats) { + value_t::stats_t::mark_used(iterable_val); + iterable_val->stats.ops.insert("object_access"); + } + } else { + JJ_DEBUG("%s", "For loop over array items"); + auto & arr = iterable_val->as_array(); + for (const auto & item : arr) { + items.push_back(item); + } + if (ctx.is_get_stats) { + value_t::stats_t::mark_used(iterable_val); + iterable_val->stats.ops.insert("array_access"); + } + } + + std::vector> scope_update_fns; + + std::vector filtered_items; + for (size_t i = 0; i < items.size(); ++i) { + context loop_scope(scope); + + value current = items[i]; + + std::function scope_update_fn = [](context &) { /* no-op */}; + if (is_stmt(loopvar)) { + auto id = cast_stmt(loopvar)->val; + + if (is_val(iterable_val)) { + // case example: {% for key in dict %} + current = items[i]->as_array()[0]; + scope_update_fn = [id, &items, i](context & ctx) { + ctx.set_val(id, items[i]->as_array()[0]); + }; + } else { + // case example: {% for item in list %} + scope_update_fn = [id, &items, i](context & ctx) { + ctx.set_val(id, items[i]); + }; + } + + } else if (is_stmt(loopvar)) { + // case example: {% for key, value in dict %} + auto tuple = cast_stmt(loopvar); + if (!is_val(current)) { + throw std::runtime_error("Cannot unpack non-iterable type: " + current->type()); + } + auto & c_arr = current->as_array(); + if (tuple->val.size() != c_arr.size()) { + throw std::runtime_error(std::string("Too ") + (tuple->val.size() > c_arr.size() ? "few" : "many") + " items to unpack"); + } + scope_update_fn = [tuple, &items, i](context & ctx) { + auto & c_arr = items[i]->as_array(); + for (size_t j = 0; j < tuple->val.size(); ++j) { + if (!is_stmt(tuple->val[j])) { + throw std::runtime_error("Cannot unpack non-identifier type: " + tuple->val[j]->type()); + } + auto id = cast_stmt(tuple->val[j])->val; + ctx.set_val(id, c_arr[j]); + } + }; + + } else { + throw std::runtime_error("Invalid loop variable(s): " + loopvar->type()); + } + + if (select_expr && test_expr) { + scope_update_fn(loop_scope); + value test_val = test_expr->execute(loop_scope); + if (!test_val->as_bool()) { + continue; + } + } + JJ_DEBUG("For loop: adding item type %s at index %zu", current->type().c_str(), i); + filtered_items.push_back(current); + scope_update_fns.push_back(scope_update_fn); + } + JJ_DEBUG("For loop: %zu items after filtering", filtered_items.size()); + + auto result = mk_val(); + + bool noIteration = true; + for (size_t i = 0; i < filtered_items.size(); i++) { + JJ_DEBUG("For loop iteration %zu/%zu", i + 1, filtered_items.size()); + value_object loop_obj = mk_val(); + loop_obj->has_builtins = false; // loop object has no builtins + loop_obj->insert("index", mk_val(i + 1)); + loop_obj->insert("index0", mk_val(i)); + loop_obj->insert("revindex", mk_val(filtered_items.size() - i)); + loop_obj->insert("revindex0", mk_val(filtered_items.size() - i - 1)); + loop_obj->insert("first", mk_val(i == 0)); + loop_obj->insert("last", mk_val(i == filtered_items.size() - 1)); + loop_obj->insert("length", mk_val(filtered_items.size())); + loop_obj->insert("previtem", i > 0 ? filtered_items[i - 1] : mk_val("previtem")); + loop_obj->insert("nextitem", i < filtered_items.size() - 1 ? filtered_items[i + 1] : mk_val("nextitem")); + scope.set_val("loop", loop_obj); + scope_update_fns[i](scope); + try { + for (auto & stmt : body) { + value val = stmt->execute(scope); + result->push_back(val); + } + } catch (const continue_statement::signal &) { + continue; + } catch (const break_statement::signal &) { + break; + } + noIteration = false; + } + + JJ_DEBUG("For loop complete, total iterations: %zu", filtered_items.size()); + if (noIteration) { + for (auto & stmt : default_block) { + value val = stmt->execute(ctx); + result->push_back(val); + } + } + + // convert to string parts + value_string str = mk_val(); + gather_string_parts_recursive(result, str); + return str; +} + +value set_statement::execute_impl(context & ctx) { + auto rhs = val ? val->execute(ctx) : exec_statements(body, ctx); + + if (is_stmt(assignee)) { + // case: {% set my_var = value %} + auto var_name = cast_stmt(assignee)->val; + JJ_DEBUG("Setting global variable '%s' with value type %s", var_name.c_str(), rhs->type().c_str()); + ctx.set_val(var_name, rhs); + + } else if (is_stmt(assignee)) { + // case: {% set a, b = value %} + auto tuple = cast_stmt(assignee); + if (!is_val(rhs)) { + throw std::runtime_error("Cannot unpack non-iterable type in set: " + rhs->type()); + } + auto & arr = rhs->as_array(); + if (arr.size() != tuple->val.size()) { + throw std::runtime_error(std::string("Too ") + (tuple->val.size() > arr.size() ? "few" : "many") + " items to unpack in set"); + } + for (size_t i = 0; i < tuple->val.size(); ++i) { + auto & elem = tuple->val[i]; + if (!is_stmt(elem)) { + throw std::runtime_error("Cannot unpack to non-identifier in set: " + elem->type()); + } + auto var_name = cast_stmt(elem)->val; + ctx.set_val(var_name, arr[i]); + } + + } else if (is_stmt(assignee)) { + // case: {% set ns.my_var = value %} + auto member = cast_stmt(assignee); + if (member->computed) { + throw std::runtime_error("Cannot assign to computed member"); + } + if (!is_stmt(member->property)) { + throw std::runtime_error("Cannot assign to member with non-identifier property"); + } + auto prop_name = cast_stmt(member->property)->val; + + value object = member->object->execute(ctx); + if (!is_val(object)) { + throw std::runtime_error("Cannot assign to member of non-object"); + } + auto obj_ptr = cast_val(object); + JJ_DEBUG("Setting object property '%s' with value type %s", prop_name.c_str(), rhs->type().c_str()); + obj_ptr->insert(prop_name, rhs); + + } else { + throw std::runtime_error("Invalid LHS inside assignment expression: " + assignee->type()); + } + return mk_val(); +} + +value macro_statement::execute_impl(context & ctx) { + if (!is_stmt(this->name)) { + throw std::runtime_error("Macro name must be an identifier"); + } + std::string name = cast_stmt(this->name)->val; + + const func_handler func = [this, name, &ctx](const func_args & args) -> value { + size_t expected_count = this->args.size(); + size_t input_count = args.count(); + + JJ_DEBUG("Invoking macro '%s' with %zu input arguments (expected %zu)", name.c_str(), input_count, expected_count); + context macro_ctx(ctx); // new scope for macro execution + + // bind parameters + for (size_t i = 0; i < expected_count; ++i) { + if (i < input_count) { + if (is_stmt(this->args[i])) { + // normal parameter + std::string param_name = cast_stmt(this->args[i])->val; + value param_value = args.get_kwarg_or_pos(param_name, i); + JJ_DEBUG(" Binding parameter '%s' to argument of type %s", param_name.c_str(), param_value->type().c_str()); + macro_ctx.set_val(param_name, param_value); + } else if (is_stmt(this->args[i])) { + // default argument used as normal parameter + auto kwarg = cast_stmt(this->args[i]); + if (!is_stmt(kwarg->key)) { + throw std::runtime_error("Keyword argument key must be an identifier in macro '" + name + "'"); + } + std::string param_name = cast_stmt(kwarg->key)->val; + value param_value = args.get_kwarg_or_pos(param_name, i); + JJ_DEBUG(" Binding parameter '%s' to argument of type %s", param_name.c_str(), param_value->type().c_str()); + macro_ctx.set_val(param_name, param_value); + } else { + throw std::runtime_error("Invalid parameter type in macro '" + name + "'"); + } + } else { + auto & default_arg = this->args[i]; + if (is_stmt(default_arg)) { + auto kwarg = cast_stmt(default_arg); + if (!is_stmt(kwarg->key)) { + throw std::runtime_error("Keyword argument key must be an identifier in macro '" + name + "'"); + } + std::string param_name = cast_stmt(kwarg->key)->val; + JJ_DEBUG(" Binding parameter '%s' to default argument of type %s", param_name.c_str(), kwarg->val->type().c_str()); + macro_ctx.set_val(param_name, kwarg->val->execute(ctx)); + } else { + throw std::runtime_error("Not enough arguments provided to macro '" + name + "'"); + } + //std::string param_name = cast_stmt(default_args[i])->val; + //JJ_DEBUG(" Binding parameter '%s' to default", param_name.c_str()); + //macro_ctx.var[param_name] = default_args[i]->execute(ctx); + } + } + + // execute macro body + JJ_DEBUG("Executing macro '%s' body with %zu statements", name.c_str(), this->body.size()); + auto res = exec_statements(this->body, macro_ctx); + JJ_DEBUG("Macro '%s' execution complete, result: %s", name.c_str(), res->val_str.str().c_str()); + return res; + }; + + JJ_DEBUG("Defining macro '%s' with %zu parameters", name.c_str(), args.size()); + ctx.set_val(name, mk_val(name, func)); + return mk_val(); +} + +value member_expression::execute_impl(context & ctx) { + value object = this->object->execute(ctx); + + value property; + if (this->computed) { + // syntax: obj[expr] + JJ_DEBUG("Member expression, computing property type %s", this->property->type().c_str()); + + int64_t arr_size = 0; + if (is_val(object)) { + arr_size = object->as_array().size(); + } else if (is_val(object)) { + arr_size = object->as_string().length(); + } + + if (is_stmt(this->property)) { + auto s = cast_stmt(this->property); + value start_val = s->start_expr ? s->start_expr->execute(ctx) : mk_val(0); + value stop_val = s->stop_expr ? s->stop_expr->execute(ctx) : mk_val(arr_size); + value step_val = s->step_expr ? s->step_expr->execute(ctx) : mk_val(1); + + // translate to function call: obj.slice(start, stop, step) + JJ_DEBUG("Member expression is a slice: start %s, stop %s, step %s", + start_val->as_repr().c_str(), + stop_val->as_repr().c_str(), + step_val->as_repr().c_str()); + auto slice_func = try_builtin_func(ctx, "slice", object); + func_args args(ctx); + args.push_back(start_val); + args.push_back(stop_val); + args.push_back(step_val); + return slice_func->invoke(args); + } else { + property = this->property->execute(ctx); + } + } else { + // syntax: obj.prop + if (!is_stmt(this->property)) { + throw std::runtime_error("Static member property must be an identifier"); + } + property = mk_val(cast_stmt(this->property)->val); + std::string prop = property->as_string().str(); + JJ_DEBUG("Member expression, object type %s, static property '%s'", object->type().c_str(), prop.c_str()); + + // behavior of jinja2: obj having prop as a built-in function AND 'prop', as an object key, + // then obj.prop returns the built-in function, not the property value. + // while obj['prop'] returns the property value. + // example: {"obj": {"items": 123}} -> obj.items is the built-in function, obj['items'] is 123 + + value val = try_builtin_func(ctx, prop, object, true); + if (!is_val(val)) { + return val; + } + // else, fallthrough to normal property access below + } + + JJ_DEBUG("Member expression on object type %s, property type %s", object->type().c_str(), property->type().c_str()); + value val = mk_val("object_property"); + + if (property->is_undefined()) { + JJ_DEBUG("%s", "Member expression property is undefined, returning undefined"); + return val; + } + + ensure_key_type_allowed(property); + + if (is_val(object)) { + JJ_DEBUG("%s", "Accessing property on undefined object, returning undefined"); + return val; + + } else if (is_val(object)) { + auto key = property->as_string().str(); + val = object->at(property, val); + if (is_val(val)) { + val = try_builtin_func(ctx, key, object, true); + } + JJ_DEBUG("Accessed property '%s' value, got type: %s", key.c_str(), val->type().c_str()); + + } else if (is_val(object) || is_val(object)) { + if (is_val(property)) { + int64_t index = property->as_int(); + JJ_DEBUG("Accessing %s index %d", object->type().c_str(), (int)index); + if (is_val(object)) { + auto & arr = object->as_array(); + if (index < 0) { + index += static_cast(arr.size()); + } + if (index >= 0 && index < static_cast(arr.size())) { + val = arr[index]; + } + } else { // value_string + auto str = object->as_string().str(); + if (index >= 0 && index < static_cast(str.size())) { + val = mk_val(std::string(1, str[index])); + } + } + + } else if (is_val(property)) { + auto key = property->as_string().str(); + JJ_DEBUG("Accessing %s built-in '%s'", is_val(object) ? "array" : "string", key.c_str()); + val = try_builtin_func(ctx, key, object, true); + + } else { + throw std::runtime_error("Cannot access property with non-string/non-number: got " + property->type()); + } + } else { + if (!is_val(property)) { + throw std::runtime_error("Cannot access property with non-string: got " + property->type()); + } + auto key = property->as_string().str(); + val = try_builtin_func(ctx, key, object, true); + } + + if (ctx.is_get_stats && val && object && property) { + value_t::stats_t::mark_used(val); + value_t::stats_t::mark_used(object); + value_t::stats_t::mark_used(property); + if (is_val(property)) { + object->stats.ops.insert("array_access"); + } else if (is_val(property)) { + object->stats.ops.insert("object_access"); + } + } + + return val; +} + +value call_expression::execute_impl(context & ctx) { + // gather arguments + func_args args(ctx); + for (auto & arg_stmt : this->args) { + auto arg_val = arg_stmt->execute(ctx); + JJ_DEBUG(" Argument type: %s", arg_val->type().c_str()); + args.push_back(arg_val); + } + // execute callee + value callee_val = callee->execute(ctx); + if (!is_val(callee_val)) { + throw std::runtime_error("Callee is not a function: got " + callee_val->type()); + } + auto * callee_func = cast_val(callee_val); + JJ_DEBUG("Calling function '%s' with %zu arguments", callee_func->name.c_str(), args.count()); + return callee_func->invoke(args); +} + +value keyword_argument_expression::execute_impl(context & ctx) { + if (!is_stmt(key)) { + throw std::runtime_error("Keyword argument key must be identifiers"); + } + + std::string k = cast_stmt(key)->val; + JJ_DEBUG("Keyword argument expression key: %s, value: %s", k.c_str(), val->type().c_str()); + + value v = val->execute(ctx); + JJ_DEBUG("Keyword argument value executed, type: %s", v->type().c_str()); + + return mk_val(k, v); +} + +} // namespace jinja + +#include "common.h.inc" +#include "common.h.inc" + +// for converting from JSON to jinja values +#include + +#include +#include +#include +#include +#include +#include + +#undef FILENAME +#define FILENAME "jinja-value" + +namespace jinja { + +// func_args method implementations + +value func_args::get_kwarg(const std::string & key, value default_val) const { + for (const auto & arg : args) { + if (is_val(arg)) { + auto * kwarg = cast_val(arg); + if (kwarg->key == key) { + return kwarg->val; + } + } + } + return default_val; +} + +value func_args::get_kwarg_or_pos(const std::string & key, size_t pos) const { + value val = get_kwarg(key, mk_val()); + + if (val->is_undefined() && pos < count() && !is_val(args[pos])) { + return args[pos]; + } + + return val; +} + +value func_args::get_pos(size_t pos) const { + if (count() > pos) { + return args[pos]; + } + throw raised_exception("Function '" + func_name + "' expected at least " + std::to_string(pos + 1) + " arguments, got " + std::to_string(count())); +} + +value func_args::get_pos(size_t pos, value default_val) const { + if (count() > pos) { + return args[pos]; + } + return default_val; +} + +void func_args::push_back(const value & val) { + args.push_back(val); +} + +void func_args::push_front(const value & val) { + args.insert(args.begin(), val); +} + +const std::vector & func_args::get_args() const { + return args; +} + +/** + * Function that mimics Python's array slicing. + */ +template +static T slice(const T & array, int64_t start, int64_t stop, int64_t step = 1) { + int64_t len = static_cast(array.size()); + int64_t direction = (step > 0) ? 1 : ((step < 0) ? -1 : 0); + int64_t start_val = 0; + int64_t stop_val = 0; + if (direction >= 0) { + start_val = start; + if (start_val < 0) { + start_val = std::max(len + start_val, (int64_t)0); + } else { + start_val = std::min(start_val, len); + } + + stop_val = stop; + if (stop_val < 0) { + stop_val = std::max(len + stop_val, (int64_t)0); + } else { + stop_val = std::min(stop_val, len); + } + } else { + start_val = len - 1; + if (start_val < 0) { + start_val = std::max(len + start_val, (int64_t)-1); + } else { + start_val = std::min(start_val, len - 1); + } + + stop_val = -1; + if (stop_val < -1) { + stop_val = std::max(len + stop_val, (int64_t)-1); + } else { + stop_val = std::min(stop_val, len - 1); + } + } + T result; + if (direction == 0) { + return result; + } + for (int64_t i = start_val; direction * i < direction * stop_val; i += step) { + if (i >= 0 && i < len) { + result.push_back(array[static_cast(i)]); + } + } + return result; +} + +template +static value empty_value_fn(const func_args &) { + if constexpr (std::is_same_v) { + return mk_val(0); + } else if constexpr (std::is_same_v) { + return mk_val(0.0); + } else if constexpr (std::is_same_v) { + return mk_val(false); + } else { + return mk_val(); + } +} +template +static value test_type_fn(const func_args & args) { + args.ensure_count(1); + bool is_type = is_val(args.get_pos(0)); + JJ_DEBUG("test_type_fn: type=%s result=%d", typeid(T).name(), is_type ? 1 : 0); + return mk_val(is_type); +} +template +static value test_type_fn(const func_args & args) { + args.ensure_count(1); + bool is_type = is_val(args.get_pos(0)) || is_val(args.get_pos(0)); + JJ_DEBUG("test_type_fn: type=%s or %s result=%d", typeid(T).name(), typeid(U).name(), is_type ? 1 : 0); + return mk_val(is_type); +} +template +static value test_type_fn(const func_args & args) { + args.ensure_count(1); + bool is_type = is_val(args.get_pos(0)) || is_val(args.get_pos(0)) || is_val(args.get_pos(0)); + JJ_DEBUG("test_type_fn: type=%s, %s or %s result=%d", typeid(T).name(), typeid(U).name(), typeid(V).name(), is_type ? 1 : 0); + return mk_val(is_type); +} +template +static value test_compare_fn(const func_args & args) { + args.ensure_count(2, 2); + return mk_val(value_compare(args.get_pos(0), args.get_pos(1), op)); +} + +static std::string json_ensure_ascii_preserving_format(const std::string & json_str) { + llama_common_unicode_string result = + llama_common_json_ensure_ascii_preserving_format_rust((const uint8_t *) json_str.data(), json_str.size()); + return common_rust_string_to_std(result); +} + +static value tojson(const func_args & args) { + args.ensure_count(1, 5); + value val_ascii = args.get_kwarg_or_pos("ensure_ascii", 1); + value val_indent = args.get_kwarg_or_pos("indent", 2); + value val_separators = args.get_kwarg_or_pos("separators", 3); + value val_sort = args.get_kwarg_or_pos("sort_keys", 4); + int indent = -1; + if (args.ctx.is_get_stats) { + // mark as used (recursively) for stats + auto val_input = args.get_pos(0); + value_t::stats_t::mark_used(const_cast(val_input), true); + } + if (is_val(val_indent)) { + indent = static_cast(val_indent->as_int()); + } + if (val_sort->as_bool()) { // undefined == false + throw not_implemented_exception("tojson sort_keys=true not implemented"); + } + const bool ensure_ascii = val_ascii->as_bool(); // undefined == false + auto separators = (is_val(val_separators) ? val_separators : mk_val())->as_array(); + std::string item_sep = separators.size() > 0 ? separators[0]->as_string().str() : (indent < 0 ? ", " : ","); + std::string key_sep = separators.size() > 1 ? separators[1]->as_string().str() : ": "; + std::string json_str = value_to_json(args.get_pos(0), indent, item_sep, key_sep); + if (ensure_ascii) { + json_str = json_ensure_ascii_preserving_format(json_str); + } + return mk_val(json_str); +} + +template +static value selectattr(const func_args & args) { + args.ensure_count(2, 4); + args.ensure_vals(true, true, false, false); + + auto arr = args.get_pos(0)->as_array(); + auto attribute = args.get_pos(1); + auto out = mk_val(); + value val_default = mk_val(); + + if (args.count() == 2) { + // example: array | selectattr("active") + for (const auto & item : arr) { + if (!is_val(item)) { + throw raised_exception("selectattr: item is not an object"); + } + value attr_val = item->at(attribute, val_default); + bool is_selected = attr_val->as_bool(); + if constexpr (is_reject) is_selected = !is_selected; + if (is_selected) out->push_back(item); + } + return out; + + } else if (args.count() == 3) { + // example: array | selectattr("equalto", "text") + // translated to: test_is_equalto(item, "text") + std::string test_name = args.get_pos(1)->as_string().str(); + value test_val = args.get_pos(2); + auto & builtins = global_builtins(); + auto it = builtins.find("test_is_" + test_name); + if (it == builtins.end()) { + throw raised_exception("selectattr: unknown test '" + test_name + "'"); + } + auto test_fn = it->second; + for (const auto & item : arr) { + func_args test_args(args.ctx); + test_args.push_back(item); // current object + test_args.push_back(test_val); // extra argument + value test_result = test_fn(test_args); + bool is_selected = test_result->as_bool(); + if constexpr (is_reject) is_selected = !is_selected; + if (is_selected) out->push_back(item); + } + return out; + + } else if (args.count() == 4) { + // example: array | selectattr("status", "equalto", "active") + // translated to: test_is_equalto(item.status, "active") + std::string test_name = args.get_pos(2)->as_string().str(); + auto extra_arg = args.get_pos(3); + auto & builtins = global_builtins(); + auto it = builtins.find("test_is_" + test_name); + if (it == builtins.end()) { + throw raised_exception("selectattr: unknown test '" + test_name + "'"); + } + auto test_fn = it->second; + for (const auto & item : arr) { + if (!is_val(item)) { + throw raised_exception("selectattr: item is not an object"); + } + value attr_val = item->at(attribute, val_default); + func_args test_args(args.ctx); + test_args.push_back(attr_val); // attribute value + test_args.push_back(extra_arg); // extra argument + value test_result = test_fn(test_args); + bool is_selected = test_result->as_bool(); + if constexpr (is_reject) is_selected = !is_selected; + if (is_selected) out->push_back(item); + } + return out; + } else { + throw raised_exception("selectattr: invalid number of arguments"); + } + + return out; +} + +static value default_value(const func_args & args) { + args.ensure_count(2, 3); + value val_check = args.get_kwarg_or_pos("boolean", 2); + bool check_bool = val_check->as_bool(); // undefined == false + bool no_value = check_bool + ? (!args.get_pos(0)->as_bool()) + : (args.get_pos(0)->is_undefined() || args.get_pos(0)->is_none()); + return no_value ? args.get_pos(1) : args.get_pos(0); +} + +const func_builtins & global_builtins() { + static const func_builtins builtins = { + {"raise_exception", [](const func_args & args) -> value { + args.ensure_vals(); + std::string msg = args.get_pos(0)->as_string().str(); + throw raised_exception("Jinja Exception: " + msg); + }}, + {"namespace", [](const func_args & args) -> value { + auto out = mk_val(); + for (const auto & arg : args.get_args()) { + if (!is_val(arg)) { + throw raised_exception("namespace() arguments must be kwargs"); + } + auto kwarg = cast_val(arg); + JJ_DEBUG("namespace: adding key '%s'", kwarg->key.c_str()); + out->insert(kwarg->key, kwarg->val); + } + return out; + }}, + {"strftime_now", [](const func_args & args) -> value { + args.ensure_vals(); + std::string format = args.get_pos(0)->as_string().str(); + // get current time + // TODO: make sure this is the same behavior as Python's strftime + char buf[100]; + if (std::strftime(buf, sizeof(buf), format.c_str(), std::localtime(&args.ctx.current_time))) { + return mk_val(std::string(buf)); + } else { + throw raised_exception("strftime_now: failed to format time"); + } + }}, + {"range", [](const func_args & args) -> value { + args.ensure_count(1, 3); + args.ensure_vals(true, false, false); + + auto arg0 = args.get_pos(0); + auto arg1 = args.get_pos(1, mk_val()); + auto arg2 = args.get_pos(2, mk_val()); + + int64_t start, stop, step; + if (args.count() == 1) { + start = 0; + stop = arg0->as_int(); + step = 1; + } else if (args.count() == 2) { + start = arg0->as_int(); + stop = arg1->as_int(); + step = 1; + } else { + start = arg0->as_int(); + stop = arg1->as_int(); + step = arg2->as_int(); + } + + auto out = mk_val(); + if (step == 0) { + throw raised_exception("range() step argument must not be zero"); + } + if (step > 0) { + for (int64_t i = start; i < stop; i += step) { + out->push_back(mk_val(i)); + } + } else { + for (int64_t i = start; i > stop; i += step) { + out->push_back(mk_val(i)); + } + } + return out; + }}, + {"tojson", tojson}, + + // tests + {"test_is_boolean", test_type_fn}, + {"test_is_callable", test_type_fn}, + {"test_is_odd", [](const func_args & args) -> value { + args.ensure_vals(); + int64_t val = args.get_pos(0)->as_int(); + return mk_val(val % 2 != 0); + }}, + {"test_is_even", [](const func_args & args) -> value { + args.ensure_vals(); + int64_t val = args.get_pos(0)->as_int(); + return mk_val(val % 2 == 0); + }}, + {"test_is_false", [](const func_args & args) -> value { + args.ensure_count(1); + bool val = is_val(args.get_pos(0)) && !args.get_pos(0)->as_bool(); + return mk_val(val); + }}, + {"test_is_true", [](const func_args & args) -> value { + args.ensure_count(1); + bool val = is_val(args.get_pos(0)) && args.get_pos(0)->as_bool(); + return mk_val(val); + }}, + {"test_is_divisibleby", [](const func_args & args) -> value { + args.ensure_vals(); + bool res = args.get_pos(0)->val_int % args.get_pos(1)->val_int == 0; + return mk_val(res); + }}, + {"test_is_string", test_type_fn}, + {"test_is_integer", test_type_fn}, + {"test_is_float", test_type_fn}, + {"test_is_number", test_type_fn}, + {"test_is_iterable", test_type_fn}, + {"test_is_sequence", test_type_fn}, + {"test_is_mapping", test_type_fn}, + {"test_is_lower", [](const func_args & args) -> value { + args.ensure_vals(); + return mk_val(args.get_pos(0)->val_str.is_lowercase()); + }}, + {"test_is_upper", [](const func_args & args) -> value { + args.ensure_vals(); + return mk_val(args.get_pos(0)->val_str.is_uppercase()); + }}, + {"test_is_none", test_type_fn}, + {"test_is_defined", [](const func_args & args) -> value { + args.ensure_count(1); + bool res = !args.get_pos(0)->is_undefined(); + JJ_DEBUG("test_is_defined: result=%d", res ? 1 : 0); + return mk_val(res); + }}, + {"test_is_undefined", test_type_fn}, + {"test_is_eq", test_compare_fn}, + {"test_is_equalto", test_compare_fn}, + {"test_is_ge", test_compare_fn}, + {"test_is_gt", test_compare_fn}, + {"test_is_greaterthan", test_compare_fn}, + {"test_is_lt", test_compare_fn}, + {"test_is_lessthan", test_compare_fn}, + {"test_is_ne", test_compare_fn}, + {"test_is_in", [](const func_args & args) -> value { + args.ensure_count(2); + auto needle = args.get_pos(0); + auto haystack = args.get_pos(1); + if (is_val(haystack)) { + return mk_val(false); + } + if (is_val(haystack)) { + for (const auto & item : haystack->as_array()) { + if (*needle == *item) { + return mk_val(true); + } + } + return mk_val(false); + } + if (is_val(haystack)) { + if (!is_val(needle)) { + throw raised_exception("'in' test expects args[1] as string when args[0] is string, got args[1] as " + needle->type()); + } + return mk_val( + haystack->as_string().str().find(needle->as_string().str()) != std::string::npos); + } + if (is_val(haystack)) { + return mk_val(haystack->has_key(needle)); + } + throw raised_exception("'in' test expects iterable as first argument, got " + haystack->type()); + }}, + {"test_is_test", [](const func_args & args) -> value { + args.ensure_vals(); + auto & builtins = global_builtins(); + std::string test_name = args.get_pos(0)->val_str.str(); + auto it = builtins.find("test_is_" + test_name); + bool res = it != builtins.end(); + return mk_val(res); + }}, + {"test_is_sameas", [](const func_args & args) -> value { + // Check if an object points to the same memory address as another object + (void)args; + throw not_implemented_exception("sameas test not implemented"); + }}, + {"test_is_escaped", [](const func_args & args) -> value { + (void)args; + throw not_implemented_exception("escaped test not implemented"); + }}, + {"test_is_filter", [](const func_args & args) -> value { + (void)args; + throw not_implemented_exception("filter test not implemented"); + }}, + }; + return builtins; +} + + +const func_builtins & value_int_t::get_builtins() const { + static const func_builtins builtins = { + {"default", default_value}, + {"abs", [](const func_args & args) -> value { + args.ensure_vals(); + int64_t val = args.get_pos(0)->as_int(); + return mk_val(val < 0 ? -val : val); + }}, + {"int", [](const func_args & args) -> value { + args.ensure_vals(); + return mk_val(args.get_pos(0)->as_int()); + }}, + {"float", [](const func_args & args) -> value { + args.ensure_vals(); + double val = static_cast(args.get_pos(0)->as_int()); + return mk_val(val); + }}, + {"safe", tojson}, + {"string", tojson}, + {"tojson", tojson}, + }; + return builtins; +} + + +const func_builtins & value_float_t::get_builtins() const { + static const func_builtins builtins = { + {"default", default_value}, + {"abs", [](const func_args & args) -> value { + args.ensure_vals(); + double val = args.get_pos(0)->as_float(); + return mk_val(val < 0.0 ? -val : val); + }}, + {"int", [](const func_args & args) -> value { + args.ensure_vals(); + int64_t val = static_cast(args.get_pos(0)->as_float()); + return mk_val(val); + }}, + {"float", [](const func_args & args) -> value { + args.ensure_vals(); + return mk_val(args.get_pos(0)->as_float()); + }}, + {"safe", tojson}, + {"string", tojson}, + {"tojson", tojson}, + }; + return builtins; +} + +[[noreturn]] static value string_join_not_implemented(const func_args &) { + throw not_implemented_exception("String join builtin not implemented"); +} + +const func_builtins & value_string_t::get_builtins() const { + static const func_builtins builtins = { + {"default", default_value}, + {"upper", [](const func_args & args) -> value { + args.ensure_vals(); + jinja::string str = args.get_pos(0)->as_string().uppercase(); + return mk_val(str); + }}, + {"lower", [](const func_args & args) -> value { + args.ensure_vals(); + jinja::string str = args.get_pos(0)->as_string().lowercase(); + return mk_val(str); + }}, + {"strip", [](const func_args & args) -> value { + value val_input = args.get_pos(0); + if (!is_val(val_input)) { + throw raised_exception("strip() first argument must be a string"); + } + value val_chars = args.get_kwarg_or_pos("chars", 1); + if (val_chars->is_undefined()) { + return mk_val(args.get_pos(0)->as_string().strip(true, true)); + } else { + return mk_val(args.get_pos(0)->as_string().strip(true, true, val_chars->as_string().str())); + } + }}, + {"rstrip", [](const func_args & args) -> value { + args.ensure_vals(); + value val_chars = args.get_kwarg_or_pos("chars", 1); + if (val_chars->is_undefined()) { + return mk_val(args.get_pos(0)->as_string().strip(false, true)); + } else { + return mk_val(args.get_pos(0)->as_string().strip(false, true, val_chars->as_string().str())); + } + }}, + {"lstrip", [](const func_args & args) -> value { + args.ensure_vals(); + value val_chars = args.get_kwarg_or_pos("chars", 1); + if (val_chars->is_undefined()) { + return mk_val(args.get_pos(0)->as_string().strip(true, false)); + } else { + return mk_val(args.get_pos(0)->as_string().strip(true, false, val_chars->as_string().str())); + } + }}, + {"title", [](const func_args & args) -> value { + args.ensure_vals(); + jinja::string str = args.get_pos(0)->as_string().titlecase(); + return mk_val(str); + }}, + {"capitalize", [](const func_args & args) -> value { + args.ensure_vals(); + jinja::string str = args.get_pos(0)->as_string().capitalize(); + return mk_val(str); + }}, + {"length", [](const func_args & args) -> value { + args.ensure_vals(); + jinja::string str = args.get_pos(0)->as_string(); + return mk_val(str.length()); + }}, + {"startswith", [](const func_args & args) -> value { + args.ensure_vals(); + std::string str = args.get_pos(0)->as_string().str(); + std::string prefix = args.get_pos(1)->as_string().str(); + return mk_val(llama_common_string_starts_with_rust( + (const uint8_t *) str.data(), str.size(), + (const uint8_t *) prefix.data(), prefix.size())); + }}, + {"endswith", [](const func_args & args) -> value { + args.ensure_vals(); + std::string str = args.get_pos(0)->as_string().str(); + std::string suffix = args.get_pos(1)->as_string().str(); + return mk_val(llama_common_string_ends_with_rust( + (const uint8_t *) str.data(), str.size(), + (const uint8_t *) suffix.data(), suffix.size())); + }}, + {"split", [](const func_args & args) -> value { + args.ensure_count(1, 3); + value val_input = args.get_pos(0); + if (!is_val(val_input)) { + throw raised_exception("split() first argument must be a string"); + } + std::string str = val_input->as_string().str(); + // FIXME: Support non-specified delimiter (split on consecutive (no leading or trailing) whitespace) + std::string delim = (args.count() > 1) ? args.get_pos(1)->as_string().str() : " "; + int64_t maxsplit = (args.count() > 2) ? args.get_pos(2)->as_int() : -1; + auto result = mk_val(); + size_t pos = 0; + std::string token; + while ((pos = str.find(delim)) != std::string::npos && maxsplit != 0) { + token = str.substr(0, pos); + result->push_back(mk_val(token)); + str.erase(0, pos + delim.length()); + --maxsplit; + } + auto res = mk_val(str); + res->val_str.mark_input_based_on(args.get_pos(0)->val_str); + result->push_back(std::move(res)); + return result; + }}, + {"rsplit", [](const func_args & args) -> value { + args.ensure_count(1, 3); + value val_input = args.get_pos(0); + if (!is_val(val_input)) { + throw raised_exception("rsplit() first argument must be a string"); + } + std::string str = val_input->as_string().str(); + // FIXME: Support non-specified delimiter (split on consecutive (no leading or trailing) whitespace) + std::string delim = (args.count() > 1) ? args.get_pos(1)->as_string().str() : " "; + int64_t maxsplit = (args.count() > 2) ? args.get_pos(2)->as_int() : -1; + auto result = mk_val(); + size_t pos = 0; + std::string token; + while ((pos = str.rfind(delim)) != std::string::npos && maxsplit != 0) { + token = str.substr(pos + delim.length()); + result->push_back(mk_val(token)); + str.erase(pos); + --maxsplit; + } + auto res = mk_val(str); + res->val_str.mark_input_based_on(args.get_pos(0)->val_str); + result->push_back(std::move(res)); + result->reverse(); + return result; + }}, + {"replace", [](const func_args & args) -> value { + args.ensure_vals(true, true, true, false); + std::string str = args.get_pos(0)->as_string().str(); + std::string old_str = args.get_pos(1)->as_string().str(); + std::string new_str = args.get_pos(2)->as_string().str(); + int64_t count = args.count() > 3 ? args.get_pos(3)->as_int() : -1; + if (count > 0) { + throw not_implemented_exception("String replace with count argument not implemented"); + } + size_t pos = 0; + while ((pos = str.find(old_str, pos)) != std::string::npos) { + str.replace(pos, old_str.length(), new_str); + pos += new_str.length(); + } + auto res = mk_val(str); + res->val_str.mark_input_based_on(args.get_pos(0)->val_str); + return res; + }}, + {"int", [](const func_args & args) -> value { + value val_input = args.get_pos(0); + value val_default = args.get_kwarg_or_pos("default", 1); + value val_base = args.get_kwarg_or_pos("base", 2); + const int base = val_base->is_undefined() ? 10 : val_base->as_int(); + if (is_val(val_input) == false) { + throw raised_exception("int() first argument must be a string"); + } + std::string str = val_input->as_string().str(); + try { + return mk_val(std::stoi(str, nullptr, base)); + } catch (...) { + return mk_val(val_default->is_undefined() ? 0 : val_default->as_int()); + } + }}, + {"float", [](const func_args & args) -> value { + args.ensure_vals(); + value val_default = args.get_kwarg_or_pos("default", 1); + std::string str = args.get_pos(0)->as_string().str(); + try { + return mk_val(std::stod(str)); + } catch (...) { + return mk_val(val_default->is_undefined() ? 0.0 : val_default->as_float()); + } + }}, + {"string", [](const func_args & args) -> value { + // no-op + args.ensure_vals(); + return mk_val(args.get_pos(0)->as_string()); + }}, + {"default", [](const func_args & args) -> value { + value input = args.get_pos(0); + if (!is_val(input)) { + throw raised_exception("default() first argument must be a string"); + } + value default_val = mk_val(""); + if (args.count() > 1 && !args.get_pos(1)->is_undefined()) { + default_val = args.get_pos(1); + } + value boolean_val = args.get_kwarg_or_pos("boolean", 2); // undefined == false + if (input->is_undefined() || (boolean_val->as_bool() && !input->as_bool())) { + return default_val; + } else { + return input; + } + }}, + {"slice", [](const func_args & args) -> value { + args.ensure_count(1, 4); + args.ensure_vals(true, true, false, false); + + auto arg0 = args.get_pos(1); + auto arg1 = args.get_pos(2, mk_val()); + auto arg2 = args.get_pos(3, mk_val()); + + int64_t start, stop, step; + if (args.count() == 1) { + start = 0; + stop = arg0->as_int(); + step = 1; + } else if (args.count() == 2) { + start = arg0->as_int(); + stop = arg1->as_int(); + step = 1; + } else { + start = arg0->as_int(); + stop = arg1->as_int(); + step = arg2->as_int(); + } + if (step == 0) { + throw raised_exception("slice step cannot be zero"); + } + auto input = args.get_pos(0); + auto sliced = slice(input->as_string().str(), start, stop, step); + auto res = mk_val(sliced); + res->val_str.mark_input_based_on(input->as_string()); + return res; + }}, + {"safe", [](const func_args & args) -> value { + // no-op for now + args.ensure_vals(); + return args.get_pos(0); + }}, + {"tojson", tojson}, + {"indent", [](const func_args &args) -> value { + args.ensure_count(1, 4); + value val_input = args.get_pos(0); + value val_width = args.get_kwarg_or_pos("width", 1); + const bool first = args.get_kwarg_or_pos("first", 2)->as_bool(); // undefined == false + const bool blank = args.get_kwarg_or_pos("blank", 3)->as_bool(); // undefined == false + if (!is_val(val_input)) { + throw raised_exception("indent() first argument must be a string"); + } + std::string indent; + if (is_val(val_width)) { + indent.assign(val_width->as_int(), ' '); + } else if (is_val(val_width)) { + indent = val_width->as_string().str(); + } else { + indent = " "; + } + std::string indented; + std::string input = val_input->as_string().str(); + std::istringstream iss = std::istringstream(input); + std::string line; + while (std::getline(iss, line)) { + if (!indented.empty()) { + indented.push_back('\n'); + } + if ((indented.empty() ? first : (!line.empty() || blank))) { + indented += indent; + } + indented += line; + } + if (!input.empty() && input.back() == '\n') { + indented.push_back('\n'); + if (blank) { + indented += indent; + } + } + + auto res = mk_val(indented); + res->val_str.mark_input_based_on(val_input->as_string()); + return res; + }}, + {"join", string_join_not_implemented}, + }; + return builtins; +} + + +const func_builtins & value_bool_t::get_builtins() const { + static const func_handler tostring = [](const func_args & args) -> value { + args.ensure_vals(); + bool val = args.get_pos(0)->as_bool(); + return mk_val(val ? "True" : "False"); + }; + static const func_builtins builtins = { + {"default", default_value}, + {"int", [](const func_args & args) -> value { + args.ensure_vals(); + bool val = args.get_pos(0)->as_bool(); + return mk_val(val ? 1 : 0); + }}, + {"float", [](const func_args & args) -> value { + args.ensure_vals(); + bool val = args.get_pos(0)->as_bool(); + return mk_val(val ? 1.0 : 0.0); + }}, + {"safe", tostring}, + {"string", tostring}, + {"tojson", tojson}, + }; + return builtins; +} + +[[noreturn]] static value array_unique_not_implemented(const func_args &) { + throw not_implemented_exception("Array unique builtin not implemented"); +} + +const func_builtins & value_array_t::get_builtins() const { + static const func_builtins builtins = { + {"default", default_value}, + {"list", [](const func_args & args) -> value { + args.ensure_vals(); + const auto & arr = args.get_pos(0)->as_array(); + auto result = mk_val(); + for (const auto& v : arr) { + result->push_back(v); + } + return result; + }}, + {"first", [](const func_args & args) -> value { + args.ensure_vals(); + const auto & arr = args.get_pos(0)->as_array(); + if (arr.empty()) { + return mk_val(); + } + return arr[0]; + }}, + {"last", [](const func_args & args) -> value { + args.ensure_vals(); + const auto & arr = args.get_pos(0)->as_array(); + if (arr.empty()) { + return mk_val(); + } + return arr[arr.size() - 1]; + }}, + {"length", [](const func_args & args) -> value { + args.ensure_vals(); + const auto & arr = args.get_pos(0)->as_array(); + return mk_val(static_cast(arr.size())); + }}, + {"slice", [](const func_args & args) -> value { + args.ensure_count(1, 4); + args.ensure_vals(true, true, false, false); + + auto val = args.get_pos(0); + auto arg0 = args.get_pos(1); + auto arg1 = args.get_pos(2, mk_val()); + auto arg2 = args.get_pos(3, mk_val()); + + int64_t start, stop, step; + if (args.count() == 1) { + start = 0; + stop = arg0->as_int(); + step = 1; + } else if (args.count() == 2) { + start = arg0->as_int(); + stop = arg1->as_int(); + step = 1; + } else { + start = arg0->as_int(); + stop = arg1->as_int(); + step = arg2->as_int(); + } + if (step == 0) { + throw raised_exception("slice step cannot be zero"); + } + auto arr = slice(val->as_array(), start, stop, step); + return is_val(val) ? mk_val(std::move(arr)) : mk_val(std::move(arr)); + }}, + {"selectattr", selectattr}, + {"select", selectattr}, + {"rejectattr", selectattr}, + {"reject", selectattr}, + {"join", [](const func_args & args) -> value { + args.ensure_count(1, 3); + if (!is_val(args.get_pos(0))) { + throw raised_exception("join() first argument must be an array"); + } + value val_delim = args.get_kwarg_or_pos("d", 1); + value attribute = args.get_kwarg_or_pos("attribute", 2); + const auto & arr = args.get_pos(0)->as_array(); + const bool attr_is_int = is_val(attribute); + if (!attribute->is_undefined() && !is_val(attribute) && !attr_is_int) { + throw raised_exception("join() attribute must be string or integer"); + } + const int64_t attr_int = attr_is_int ? attribute->as_int() : 0; + const std::string delim = val_delim->is_undefined() ? "" : val_delim->as_string().str(); + std::string result; + for (size_t i = 0; i < arr.size(); ++i) { + value val_arr = arr[i]; + if (!attribute->is_undefined()) { + if (attr_is_int && is_val(val_arr)) { + val_arr = val_arr->at(attr_int); + } else if (!attr_is_int && is_val(val_arr)) { + val_arr = val_arr->at(attribute); + } + } + if (!is_val(val_arr) && !is_val(val_arr) && !is_val(val_arr)) { + throw raised_exception("join() can only join arrays of strings or numerics"); + } + result += val_arr->as_string().str(); + if (i < arr.size() - 1) { + result += delim; + } + } + return mk_val(result); + }}, + {"string", [](const func_args & args) -> value { + args.ensure_vals(); + if (args.ctx.is_get_stats) { + // mark as used (recursively) for stats + auto val_input = args.get_pos(0); + value_t::stats_t::mark_used(const_cast(val_input), true); + } + return mk_val(args.get_pos(0)->as_string()); + }}, + {"tojson", tojson}, + {"map", [](const func_args & args) -> value { + args.ensure_count(2); + if (!is_val(args.get_pos(0))) { + throw raised_exception("map: first argument must be an array"); + } + if (!is_val(args.get_args().at(1))) { + throw not_implemented_exception("map: filter-mapping not implemented"); + } + value val = args.get_pos(0); + value attribute = args.get_kwarg_or_pos("attribute", 1); + const bool attr_is_int = is_val(attribute); + if (!is_val(attribute) && !attr_is_int) { + throw raised_exception("map: attribute must be string or integer"); + } + const int64_t attr_int = attr_is_int ? attribute->as_int() : 0; + value default_val = args.get_kwarg("default", mk_val()); + auto out = mk_val(); + auto arr = val->as_array(); + for (const auto & item : arr) { + value attr_val; + if (attr_is_int) { + attr_val = is_val(item) ? item->at(attr_int, default_val) : default_val; + } else { + attr_val = is_val(item) ? item->at(attribute, default_val) : default_val; + } + out->push_back(attr_val); + } + return is_val(val) ? mk_val(std::move(out->as_array())) : out; + }}, + {"append", [](const func_args & args) -> value { + args.ensure_count(2); + if (!is_val(args.get_pos(0))) { + throw raised_exception("append: first argument must be an array"); + } + const value_array_t * arr = cast_val(args.get_pos(0)); + // need to use const_cast here to modify the array + value_array_t * arr_editable = const_cast(arr); + arr_editable->push_back(args.get_pos(1)); + return args.get_pos(0); + }}, + {"pop", [](const func_args & args) -> value { + args.ensure_count(1, 2); + args.ensure_vals(true, false); + int64_t index = args.count() == 2 ? args.get_pos(1)->as_int() : -1; + const value_array_t * arr = cast_val(args.get_pos(0)); + // need to use const_cast here to modify the array + value_array_t * arr_editable = const_cast(arr); + return arr_editable->pop_at(index); + }}, + {"sort", [](const func_args & args) -> value { + args.ensure_count(1, 4); + if (!is_val(args.get_pos(0))) { + throw raised_exception("sort: first argument must be an array"); + } + value val = args.get_pos(0); + value val_reverse = args.get_kwarg_or_pos("reverse", 1); + value val_case = args.get_kwarg_or_pos("case_sensitive", 2); + value attribute = args.get_kwarg_or_pos("attribute", 3); + // FIXME: sorting is currently always case sensitive + //const bool case_sensitive = val_case->as_bool(); // undefined == false + const bool reverse = val_reverse->as_bool(); // undefined == false + const bool attr_is_int = is_val(attribute); + const int64_t attr_int = attr_is_int ? attribute->as_int() : 0; + std::vector arr = val->as_array(); // copy + std::sort(arr.begin(), arr.end(),[&](const value & a, const value & b) { + value val_a = a; + value val_b = b; + if (!attribute->is_undefined()) { + if (attr_is_int && is_val(a) && is_val(b)) { + val_a = a->at(attr_int); + val_b = b->at(attr_int); + } else if (!attr_is_int && is_val(a) && is_val(b)) { + val_a = a->at(attribute); + val_b = b->at(attribute); + } else { + throw raised_exception("sort: unsupported object attribute comparison between " + a->type() + " and " + b->type()); + } + } + return value_compare(val_a, val_b, reverse ? value_compare_op::gt : value_compare_op::lt); + }); + return is_val(val) ? mk_val(std::move(arr)) : mk_val(std::move(arr)); + }}, + {"reverse", [](const func_args & args) -> value { + args.ensure_vals(); + value val = args.get_pos(0); + std::vector arr = val->as_array(); // copy + std::reverse(arr.begin(), arr.end()); + return is_val(val) ? mk_val(std::move(arr)) : mk_val(std::move(arr)); + }}, + {"unique", array_unique_not_implemented}, + }; + return builtins; +} + +[[noreturn]] static value object_join_not_implemented(const func_args &) { + throw not_implemented_exception("object join not implemented"); +} + +const func_builtins & value_object_t::get_builtins() const { + if (!has_builtins) { + static const func_builtins no_builtins = {}; + return no_builtins; + } + + static const func_builtins builtins = { + // {"default", default_value}, // cause issue with gpt-oss + {"get", [](const func_args & args) -> value { + args.ensure_count(2, 3); + if (!is_val(args.get_pos(0))) { + throw raised_exception("get: first argument must be an object"); + } + if (!is_val(args.get_pos(1))) { + throw raised_exception("get: second argument must be a string (key)"); + } + value default_val = mk_val(); + if (args.count() == 3) { + default_val = args.get_pos(2); + } + const value obj = args.get_pos(0); + const value key = args.get_pos(1); + return obj->at(key, default_val); + }}, + {"keys", [](const func_args & args) -> value { + args.ensure_vals(); + const auto & obj = args.get_pos(0)->as_ordered_object(); + auto result = mk_val(); + for (const auto & pair : obj) { + result->push_back(pair.first); + } + return result; + }}, + {"values", [](const func_args & args) -> value { + args.ensure_vals(); + const auto & obj = args.get_pos(0)->as_ordered_object(); + auto result = mk_val(); + for (const auto & pair : obj) { + result->push_back(pair.second); + } + return result; + }}, + {"items", [](const func_args & args) -> value { + args.ensure_vals(); + const auto & obj = args.get_pos(0)->as_ordered_object(); + auto result = mk_val(); + for (const auto & pair : obj) { + auto item = mk_val(pair); + result->push_back(std::move(item)); + } + return result; + }}, + {"tojson", tojson}, + {"string", [](const func_args & args) -> value { + args.ensure_vals(); + if (args.ctx.is_get_stats) { + // mark as used (recursively) for stats + auto val_input = args.get_pos(0); + value_t::stats_t::mark_used(const_cast(val_input), true); + } + return mk_val(args.get_pos(0)->as_string()); + }}, + {"length", [](const func_args & args) -> value { + args.ensure_vals(); + const auto & obj = args.get_pos(0)->as_ordered_object(); + return mk_val(static_cast(obj.size())); + }}, + {"tojson", [](const func_args & args) -> value { + args.ensure_vals(); + // use global to_json + return global_builtins().at("tojson")(args); + }}, + {"dictsort", [](const func_args & args) -> value { + value val_input = args.get_pos(0); + value val_case = args.get_kwarg_or_pos("case_sensitive", 1); + value val_by = args.get_kwarg_or_pos("by", 2); + value val_reverse = args.get_kwarg_or_pos("reverse", 3); + // FIXME: sorting is currently always case sensitive + //const bool case_sensitive = val_case->as_bool(); // undefined == false + const bool reverse = val_reverse->as_bool(); // undefined == false + const bool by_value = is_val(val_by) && val_by->as_string().str() == "value" ? true : false; + auto result = mk_val(val_input); // copy + std::sort(result->val_obj.begin(), result->val_obj.end(), [&](const auto & a, const auto & b) { + if (by_value) { + return value_compare(a.second, b.second, reverse ? value_compare_op::gt : value_compare_op::lt); + } else { + return value_compare(a.first, b.first, reverse ? value_compare_op::gt : value_compare_op::lt); + } + }); + return result; + }}, + {"join", object_join_not_implemented}, + }; + return builtins; +} + +const func_builtins & value_none_t::get_builtins() const { + static const func_handler tostring = [](const func_args &) -> value { + return mk_val("None"); + }; + static const func_builtins builtins = { + {"default", default_value}, + {"tojson", tojson}, + {"string", tostring}, + {"safe", tostring}, + {"items", empty_value_fn}, + {"map", empty_value_fn}, + {"reject", empty_value_fn}, + {"rejectattr", empty_value_fn}, + {"select", empty_value_fn}, + {"selectattr", empty_value_fn}, + {"unique", empty_value_fn}, + }; + return builtins; +} + + +const func_builtins & value_undefined_t::get_builtins() const { + static const func_builtins builtins = { + {"default", default_value}, + {"capitalize", empty_value_fn}, + {"first", empty_value_fn}, + {"items", empty_value_fn}, + {"join", empty_value_fn}, + {"last", empty_value_fn}, + {"length", empty_value_fn}, + {"list", empty_value_fn}, + {"lower", empty_value_fn}, + {"map", empty_value_fn}, + {"max", empty_value_fn}, + {"min", empty_value_fn}, + {"reject", empty_value_fn}, + {"rejectattr", empty_value_fn}, + {"replace", empty_value_fn}, + {"reverse", empty_value_fn}, + {"safe", empty_value_fn}, + {"select", empty_value_fn}, + {"selectattr", empty_value_fn}, + {"sort", empty_value_fn}, + {"string", empty_value_fn}, + {"strip", empty_value_fn}, + {"sum", empty_value_fn}, + {"title", empty_value_fn}, + {"truncate", empty_value_fn}, + {"unique", empty_value_fn}, + {"upper", empty_value_fn}, + {"wordcount", empty_value_fn}, + }; + return builtins; +} + + +////////////////////////////////// + + +static value from_json(const nlohmann::ordered_json & j, bool mark_input) { + if (j.is_null()) { + return mk_val(); + } else if (j.is_boolean()) { + return mk_val(j.get()); + } else if (j.is_number_integer()) { + return mk_val(j.get()); + } else if (j.is_number_float()) { + return mk_val(j.get()); + } else if (j.is_string()) { + auto str = mk_val(j.get()); + if (mark_input) { + str->mark_input(); + } + return str; + } else if (j.is_array()) { + auto arr = mk_val(); + for (const auto & item : j) { + arr->push_back(from_json(item, mark_input)); + } + return arr; + } else if (j.is_object()) { + auto obj = mk_val(); + for (auto it = j.begin(); it != j.end(); ++it) { + obj->insert(it.key(), from_json(it.value(), mark_input)); + } + return obj; + } else { + throw std::runtime_error("Unsupported JSON value type"); + } +} + +// compare operator for value_t +bool value_compare(const value & a, const value & b, value_compare_op op) { + auto cmp = [&]() { + // compare numeric types + if ((is_val(a) || is_val(a)) && + (is_val(b) || is_val(b))){ + try { + if (op == value_compare_op::eq) { + return a->as_float() == b->as_float(); + } else if (op == value_compare_op::ge) { + return a->as_float() >= b->as_float(); + } else if (op == value_compare_op::gt) { + return a->as_float() > b->as_float(); + } else if (op == value_compare_op::lt) { + return a->as_float() < b->as_float(); + } else if (op == value_compare_op::ne) { + return a->as_float() != b->as_float(); + } else { + throw std::runtime_error("Unsupported comparison operator for numeric types"); + } + } catch (...) {} + } + // compare string and number + // TODO: not sure if this is the right behavior + if ((is_val(b) && (is_val(a) || is_val(a))) || + (is_val(a) && (is_val(b) || is_val(b))) || + (is_val(a) && is_val(b))) { + try { + if (op == value_compare_op::eq) { + return a->as_string().str() == b->as_string().str(); + } else if (op == value_compare_op::ge) { + return a->as_string().str() >= b->as_string().str(); + } else if (op == value_compare_op::gt) { + return a->as_string().str() > b->as_string().str(); + } else if (op == value_compare_op::lt) { + return a->as_string().str() < b->as_string().str(); + } else if (op == value_compare_op::ne) { + return a->as_string().str() != b->as_string().str(); + } else { + throw std::runtime_error("Unsupported comparison operator for string/number types"); + } + } catch (...) {} + } + // compare boolean simple + if (is_val(a) && is_val(b)) { + if (op == value_compare_op::eq) { + return a->as_bool() == b->as_bool(); + } else if (op == value_compare_op::ne) { + return a->as_bool() != b->as_bool(); + } else { + throw std::runtime_error("Unsupported comparison operator for bool type"); + } + } + // compare by type + if (a->type() != b->type()) { + return false; + } + return false; + }; + auto result = cmp(); + JJ_DEBUG("Comparing types: %s and %s result=%d", a->type().c_str(), b->type().c_str(), result); + return result; +} + +template<> +void global_from_json(context & ctx, const nlohmann::ordered_json & json_obj, bool mark_input) { + // printf("global_from_json: %s\n" , json_obj.dump(2).c_str()); + if (json_obj.is_null() || !json_obj.is_object()) { + throw std::runtime_error("global_from_json: input JSON value must be an object"); + } + for (auto it = json_obj.begin(); it != json_obj.end(); ++it) { + JJ_DEBUG("global_from_json: setting key '%s'", it.key().c_str()); + ctx.set_val(it.key(), from_json(it.value(), mark_input)); + } +} + +// recursively convert value to JSON string +// TODO: avoid circular references +static void value_to_json_internal(std::ostringstream & oss, const value & val, int curr_lvl, int indent, const std::string_view item_sep, const std::string_view key_sep) { + auto indent_str = [indent, curr_lvl]() -> std::string { + return (indent > 0) ? std::string(curr_lvl * indent, ' ') : ""; + }; + auto newline = [indent]() -> std::string { + return (indent >= 0) ? "\n" : ""; + }; + + if (is_val(val) || val->is_undefined()) { + oss << "null"; + } else if (is_val(val)) { + oss << (val->as_bool() ? "true" : "false"); + } else if (is_val(val)) { + oss << val->as_int(); + } else if (is_val(val)) { + oss << val->as_float(); + } else if (is_val(val)) { + oss << "\""; + for (char c : val->as_string().str()) { + switch (c) { + case '"': oss << "\\\""; break; + case '\\': oss << "\\\\"; break; + case '\b': oss << "\\b"; break; + case '\f': oss << "\\f"; break; + case '\n': oss << "\\n"; break; + case '\r': oss << "\\r"; break; + case '\t': oss << "\\t"; break; + default: + if (static_cast(c) < 0x20) { + char buf[7]; + snprintf(buf, sizeof(buf), "\\u%04x", static_cast(c)); + oss << buf; + } else { + oss << c; + } + } + } + oss << "\""; + } else if (is_val(val)) { + const auto & arr = val->as_array(); + oss << "["; + if (!arr.empty()) { + oss << newline(); + for (size_t i = 0; i < arr.size(); ++i) { + oss << indent_str() << (indent > 0 ? std::string(indent, ' ') : ""); + value_to_json_internal(oss, arr[i], curr_lvl + 1, indent, item_sep, key_sep); + if (i < arr.size() - 1) { + oss << item_sep; + } + oss << newline(); + } + oss << indent_str(); + } + oss << "]"; + } else if (is_val(val)) { + const auto & obj = val->as_ordered_object(); // IMPORTANT: need to keep exact order + oss << "{"; + if (!obj.empty()) { + oss << newline(); + size_t i = 0; + for (const auto & pair : obj) { + oss << indent_str() << (indent > 0 ? std::string(indent, ' ') : ""); + value_to_json_internal(oss, mk_val(pair.first->as_string().str()), curr_lvl + 1, indent, item_sep, key_sep); + oss << key_sep; + value_to_json_internal(oss, pair.second, curr_lvl + 1, indent, item_sep, key_sep); + if (i < obj.size() - 1) { + oss << item_sep; + } + oss << newline(); + ++i; + } + oss << indent_str(); + } + oss << "}"; + } else { + oss << "null"; + } +} + +std::string value_to_json(const value & val, int indent, const std::string_view item_sep, const std::string_view key_sep) { + std::ostringstream oss; + value_to_json_internal(oss, val, 0, indent, item_sep, key_sep); + JJ_DEBUG("value_to_json: result=%s", oss.str().c_str()); + return oss.str(); +} + +// TODO: avoid circular references +std::string value_to_string_repr(const value & val) { + if (is_val(val)) { + const std::string val_str = val->as_string().str(); + + if (val_str.find('\'') != std::string::npos) { + return value_to_json(val); + } else { + return "'" + val_str + "'"; + } + } else { + return val->as_repr(); + } +} + +// stats utility +void value_t::stats_t::mark_used(value & val, bool deep) { + val->stats.used = true; + if (deep) { + if (is_val(val)) { + for (auto & item : val->val_arr) { + mark_used(item, deep); + } + } else if (is_val(val)) { + for (auto & pair : val->val_obj) { + mark_used(pair.first, deep); + mark_used(pair.second, deep); + } + } + } +} + +} // namespace jinja + +#include "common.h.inc" + +// note: the json dependency is only for defining input in a convenient way +// we can remove it in the future when we figure out a better way to define inputs using jinja::value +#include + +#include +#include + +#undef FILENAME +#define FILENAME "jinja-caps" + +using json = nlohmann::ordered_json; + +namespace jinja { + +using caps_json_fn = std::function; +using caps_analyze_fn = std::function; + +static void caps_try_execute(jinja::program & prog, + const caps_json_fn & messages_fn, + const caps_json_fn & tools_fn, + const caps_analyze_fn & analyze_fn) { + context ctx; + ctx.is_get_stats = true; + jinja::global_from_json(ctx, json{ + {"messages", messages_fn()}, + {"tools", tools_fn()}, + {"bos_token", ""}, + {"eos_token", ""}, + {"add_generation_prompt", true} + }, true); + + auto messages = ctx.get_val("messages"); + auto tools = ctx.get_val("tools"); + + bool success = false; + std::string result; + try { + jinja::runtime runtime(ctx); + auto results = runtime.execute(prog); + auto parts = jinja::runtime::gather_string_parts(results); + result = parts->as_string().str(); + success = true; + } catch (const std::exception & e) { + JJ_DEBUG("Exception during execution: %s", e.what()); + result = ""; + // ignore exceptions during capability analysis + } + + analyze_fn(success, messages, tools); +} + +// for debugging only +static void caps_print_stats(value & v, const std::string & path) { + std::string ops; + for (const auto & name : v->stats.ops) { + ops += name + " "; + } + JJ_DEBUG("Value %s, type: %s %s, ops: %s", + path.c_str(), + v->type().c_str(), + v->stats.used ? "(used)" : "", + ops.c_str()); +} + +std::map caps::to_map() const { + return { + {"supports_string_content", supports_string_content}, + {"supports_typed_content", supports_typed_content}, + {"supports_tools", supports_tools}, + {"supports_tool_calls", supports_tool_calls}, + {"supports_parallel_tool_calls", supports_parallel_tool_calls}, + {"supports_system_role", supports_system_role}, + {"supports_preserve_reasoning", supports_preserve_reasoning}, + {"supports_object_arguments", supports_object_arguments}, + }; +} + +std::string caps::to_string() const { + std::ostringstream ss; + ss << "Caps(\n"; + for (const auto & [key, value] : to_map()) { + ss << " " << key << "=" << (value ? "true" : "false") << "\n"; + } + ss << ")"; + return ss.str(); +} + +caps caps_get(jinja::program & prog) { + caps result; + + static const auto has_op = [](value & v, const std::string & op_name) { + return v->stats.ops.find(op_name) != v->stats.ops.end(); + }; + + JJ_DEBUG("%s\n", ">>> Running capability check: typed content"); + + // case: typed content support + caps_try_execute( + prog, + [&]() { + // messages + return json::array({ + { + {"role", "user"}, + {"content", "content"} + } + }); + }, + [&]() { + // tools + return json{nullptr}; + }, + [&](bool success, value & messages, value &) { + auto & content = messages->at(0)->at("content"); + caps_print_stats(content, "messages[0].content"); + if (has_op(content, "selectattr") || has_op(content, "array_access")) { + // accessed as an array + result.supports_typed_content = true; + } + if (!success) { + // failed to execute with content as string + result.supports_string_content = false; + } + } + ); + + JJ_DEBUG("%s\n", ">>> Running capability check: system prompt"); + + // case: system prompt support + caps_try_execute( + prog, + [&]() { + // messages + return json::array({ + { + {"role", "system"}, + {"content", "System message"} + }, + { + {"role", "user"}, + {"content", "User message"} + }, + }); + }, + [&]() { + // tools + return json::array(); + }, + [&](bool, value & messages, value &) { + auto & content = messages->at(0)->at("content"); + caps_print_stats(content, "messages[0].content"); + if (!content->stats.used) { + result.supports_system_role = false; + } + } + ); + + JJ_DEBUG("%s\n", ">>> Running capability check: single tool with object arguments support"); + + // case: tools support: single call with object arguments + caps_try_execute( + prog, + [&]() { + // messages + return json::array({ + { + {"role", "user"}, + {"content", "User message"}, + }, + { + {"role", "assistant"}, + {"content", ""}, // Some templates expect content to be empty with tool calls + {"tool_calls", json::array({ + { + {"id", "call00001"}, + {"type", "function"}, + {"function", { + {"name", "tool1"}, + {"arguments", { + {"arg", "value"} + }} + }} + } + })} + }, + { + {"role", "tool"}, + {"content", "Tool response"}, + {"tool_call_id", "call00001"} + }, + { + {"role", "assistant"}, + {"content", "The tool response was 'tool response'"} + }, + { + {"role", "user"}, + {"content", "User message"}, + }, + }); + }, + [&]() { + // tools + return json::array({ + { + {"name", "tool"}, + {"type", "function"}, + {"function", { + {"name", "tool1"}, + {"description", "Tool description"}, + {"parameters", { + {"type", "object"}, + {"properties", { + {"arg", { + {"type", "string"}, + {"description", "Arg description"}, + }}, + }}, + {"required", json::array({ "arg" })}, + }}, + }}, + }, + }); + }, + [&](bool success, value & messages, value & tools) { + if (!success) { + return; // Nothing can be inferred + } + + auto & tool_name = tools->at(0)->at("function")->at("name"); + caps_print_stats(tool_name, "tools[0].function.name"); + caps_print_stats(tools, "tools"); + if (!tool_name->stats.used) { + result.supports_tools = false; + } + + auto & tool_calls = messages->at(1)->at("tool_calls");; + caps_print_stats(tool_calls, "messages[1].tool_calls"); + if (!tool_calls->stats.used) { + result.supports_tool_calls = false; + return; + } + + auto & tool_arg = tool_calls->at(0)->at("function")->at("arguments")->at("arg"); + caps_print_stats(tool_arg, "messages[1].tool_calls[0].function.arguments.arg"); + if (tool_arg->stats.used) { + result.supports_object_arguments = true; + } + } + ); + + if (!result.supports_object_arguments) { + JJ_DEBUG("%s\n", ">>> Running capability check: single tool with string arguments support"); + + // case: tools support: single call with string arguments + caps_try_execute( + prog, + [&]() { + // messages + return json::array({ + { + {"role", "user"}, + {"content", "User message"}, + }, + { + {"role", "assistant"}, + {"content", ""}, // Some templates expect content to be empty with tool calls + {"tool_calls", json::array({ + { + {"id", "call00001"}, + {"type", "function"}, + {"function", { + {"name", "tool1"}, + {"arguments", R"({"arg": "value"})"} + }} + } + })} + }, + { + {"role", "tool"}, + {"content", "Tool response"}, + {"tool_call_id", "call00001"} + }, + { + {"role", "assistant"}, + {"content", "The tool response was 'tool response'"} + }, + { + {"role", "user"}, + {"content", "User message"}, + }, + }); + }, + [&]() { + // tools + return json::array({ + { + {"name", "tool"}, + {"type", "function"}, + {"function", { + {"name", "tool1"}, + {"description", "Tool description"}, + {"parameters", { + {"type", "object"}, + {"properties", { + {"arg", { + {"type", "string"}, + {"description", "Arg description"}, + }}, + }}, + {"required", json::array({ "arg" })}, + }}, + }}, + }, + }); + }, + [&](bool success, value & messages, value & tools) { + if (!success) { + result.supports_tool_calls = false; + result.supports_tools = false; + return; + } + + auto & tool_name = tools->at(0)->at("function")->at("name"); + caps_print_stats(tool_name, "tools[0].function.name"); + caps_print_stats(tools, "tools"); + if (!tool_name->stats.used) { + result.supports_tools = false; + } + + auto & tool_calls = messages->at(1)->at("tool_calls"); + caps_print_stats(tool_calls, "messages[1].tool_calls"); + if (!tool_calls->stats.used) { + result.supports_tool_calls = false; + return; + } + } + ); + } + + JJ_DEBUG("%s\n", ">>> Running capability check: parallel tool support"); + + // case: tools support: parallel calls + caps_try_execute( + prog, + [&]() { + json args = json(R"({"arg": "value"})"); + if (result.supports_object_arguments) { + args = json{{"arg", "value"}}; + } + + // messages + return json::array({ + { + {"role", "user"}, + {"content", "User message"}, + }, + { + {"role", "assistant"}, + {"content", ""}, // Some templates expect content to be empty with tool calls + {"tool_calls", json::array({ + { + {"id", "call00001"}, + {"type", "function"}, + {"function", { + {"name", "tool1"}, + {"arguments", args} + }} + }, + { + {"id", "call00002"}, + {"type", "function"}, + {"function", { + {"name", "tool1"}, + {"arguments", args} + }} + } + })} + }, + { + {"role", "tool"}, + {"content", "Tool response"}, + {"tool_call_id", "call00001"} + }, + { + {"role", "assistant"}, + {"content", "The tool response was 'tool response'"} + }, + { + {"role", "user"}, + {"content", "User message"}, + }, + }); + }, + [&]() { + // tools + return json::array({ + { + {"name", "tool"}, + {"type", "function"}, + {"function", { + {"name", "tool1"}, + {"description", "Tool description"}, + {"parameters", { + {"type", "object"}, + {"properties", { + {"arg", { + {"type", "string"}, + {"description", "Arg description"}, + }}, + }}, + {"required", json::array({ "arg" })}, + }}, + }}, + }, + }); + }, + [&](bool success, value & messages, value & /*tools*/) { + if (!success) { + result.supports_parallel_tool_calls = false; + return; + } + + auto & tool_calls = messages->at(1)->at("tool_calls"); + caps_print_stats(tool_calls, "messages[1].tool_calls"); + + // check for second tool call usage + auto & tool_call_1 = tool_calls->at(1)->at("function"); + caps_print_stats(tool_call_1, "messages[1].tool_calls[1].function"); + if (!tool_call_1->stats.used) { + result.supports_parallel_tool_calls = false; + } + } + ); + + JJ_DEBUG("%s\n", ">>> Running capability check: preserve reasoning"); + + // case: preserve reasoning content in chat history + caps_try_execute( + prog, + [&]() { + // messages + return json::array({ + { + {"role", "user"}, + {"content", "User message"} + }, + { + {"role", "assistant"}, + {"content", "Assistant message"}, + {"reasoning_content", "Reasoning content"} + }, + { + {"role", "user"}, + {"content", "User message"} + }, + }); + }, + [&]() { + // tools + return json::array(); + }, + [&](bool, value & messages, value &) { + auto & content = messages->at(1)->at("reasoning_content"); + caps_print_stats(content, "messages[1].reasoning_content"); + if (content->stats.used) { + result.supports_preserve_reasoning = true; + } + } + ); + + JJ_DEBUG("%s\n", result.to_string().c_str()); + + return result; +} + +} // namespace jinja + +#include "common.h.inc" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +# include +# include +# define isatty _isatty +# define fileno _fileno +#else +# include +#endif // defined(_WIN32) + +int common_log_verbosity_thold = LOG_DEFAULT_LLAMA; + +int common_log_get_verbosity_thold(void) { + return common_log_verbosity_thold; +} + +void common_log_set_verbosity_thold(int verbosity) { + common_log_verbosity_thold = verbosity; +} + +static int64_t t_us() { + return llama_common_time_us_rust(); +} + +// colors +enum common_log_col : int { + COMMON_LOG_COL_DEFAULT = 0, + COMMON_LOG_COL_BOLD, + COMMON_LOG_COL_RED, + COMMON_LOG_COL_GREEN, + COMMON_LOG_COL_YELLOW, + COMMON_LOG_COL_BLUE, + COMMON_LOG_COL_MAGENTA, + COMMON_LOG_COL_CYAN, + COMMON_LOG_COL_WHITE, +}; + +// disable colors by default +static const char* g_col[] = { + "", + "", + "", + "", + "", + "", + "", + "", + "", +}; + +struct common_log_entry { + enum ggml_log_level level; + + bool prefix; + + int64_t timestamp; + + std::vector msg; + + // signals the worker thread to stop + bool is_end; + + void print(FILE * file = nullptr) const { + FILE * fcur = file; + if (!fcur) { + // stderr displays DBG messages only when their verbosity level is not higher than the threshold + // these messages will still be logged to a file + if (level == GGML_LOG_LEVEL_DEBUG && common_log_verbosity_thold < LOG_DEFAULT_DEBUG) { + return; + } + + fcur = stdout; + + if (level != GGML_LOG_LEVEL_NONE) { + fcur = stderr; + } + } + + if (level != GGML_LOG_LEVEL_NONE && level != GGML_LOG_LEVEL_CONT && prefix) { + if (timestamp) { + // [M.s.ms.us] + fprintf(fcur, "%s%d.%02d.%03d.%03d%s ", + g_col[COMMON_LOG_COL_BLUE], + (int) (timestamp / 1000000 / 60), + (int) (timestamp / 1000000 % 60), + (int) (timestamp / 1000 % 1000), + (int) (timestamp % 1000), + g_col[COMMON_LOG_COL_DEFAULT]); + } + + switch (level) { + case GGML_LOG_LEVEL_INFO: fprintf(fcur, "%sI %s", g_col[COMMON_LOG_COL_GREEN], g_col[COMMON_LOG_COL_DEFAULT]); break; + case GGML_LOG_LEVEL_WARN: fprintf(fcur, "%sW %s", g_col[COMMON_LOG_COL_MAGENTA], "" ); break; + case GGML_LOG_LEVEL_ERROR: fprintf(fcur, "%sE %s", g_col[COMMON_LOG_COL_RED], "" ); break; + case GGML_LOG_LEVEL_DEBUG: fprintf(fcur, "%sD %s", g_col[COMMON_LOG_COL_YELLOW], "" ); break; + default: + break; + } + } + + fprintf(fcur, "%s", msg.data()); + + if (level == GGML_LOG_LEVEL_WARN || level == GGML_LOG_LEVEL_ERROR || level == GGML_LOG_LEVEL_DEBUG) { + fprintf(fcur, "%s", g_col[COMMON_LOG_COL_DEFAULT]); + } + + fflush(fcur); + } +}; + +struct common_log { + // default capacity - will be expanded if needed + common_log() : common_log(256) {} + + common_log(size_t capacity) { + file = nullptr; + prefix = false; + timestamps = false; + running = false; + t_start = t_us(); + + // initial message size - will be expanded if longer messages arrive + entries.resize(capacity); + for (auto & entry : entries) { + entry.msg.resize(256); + } + + head = 0; + tail = 0; + + resume(); + } + + ~common_log() { + pause(); + if (file) { + fclose(file); + } + } + +private: + std::mutex mtx; + std::thread thrd; + std::condition_variable cv; + + FILE * file; + + bool prefix; + bool timestamps; + bool running; + + int64_t t_start; + + // ring buffer of entries + std::vector entries; + size_t head; + size_t tail; + + // worker thread copies into this + common_log_entry cur; + +public: + void add(enum ggml_log_level level, const char * fmt, va_list args) { + std::lock_guard lock(mtx); + + if (!running) { + // discard messages while the worker thread is paused + return; + } + + auto & entry = entries[tail]; + + { + // cannot use args twice, so make a copy in case we need to expand the buffer + va_list args_copy; + va_copy(args_copy, args); + +#if 1 + const size_t n = vsnprintf(entry.msg.data(), entry.msg.size(), fmt, args); + if (n >= entry.msg.size()) { + entry.msg.resize(n + 1); + vsnprintf(entry.msg.data(), entry.msg.size(), fmt, args_copy); + } +#else + // hack for bolding arguments + + std::stringstream ss; + for (int i = 0; fmt[i] != 0; i++) { + if (fmt[i] == '%') { + ss << LOG_COL_BOLD; + while (fmt[i] != ' ' && fmt[i] != ')' && fmt[i] != ']' && fmt[i] != 0) ss << fmt[i++]; + ss << LOG_COL_DEFAULT; + if (fmt[i] == 0) break; + } + ss << fmt[i]; + } + const size_t n = vsnprintf(entry.msg.data(), entry.msg.size(), ss.str().c_str(), args); + if (n >= entry.msg.size()) { + entry.msg.resize(n + 1); + vsnprintf(entry.msg.data(), entry.msg.size(), ss.str().c_str(), args_copy); + } +#endif + va_end(args_copy); + } + + entry.level = level; + entry.prefix = prefix; + entry.timestamp = 0; + if (timestamps) { + entry.timestamp = t_us() - t_start; + } + entry.is_end = false; + + tail = (tail + 1) % entries.size(); + if (tail == head) { + // expand the buffer + std::vector new_entries(2*entries.size()); + + size_t new_tail = 0; + + do { + new_entries[new_tail] = std::move(entries[head]); + + head = (head + 1) % entries.size(); + new_tail = (new_tail + 1); + } while (head != tail); + + head = 0; + tail = new_tail; + + for (size_t i = tail; i < new_entries.size(); i++) { + new_entries[i].msg.resize(256); + } + + entries = std::move(new_entries); + } + cv.notify_one(); + } + + void resume() { + std::lock_guard lock(mtx); + + if (running) { + return; + } + + running = true; + + thrd = std::thread([this]() { + while (true) { + { + std::unique_lock lock(mtx); + cv.wait(lock, [this]() { return head != tail; }); + cur = entries[head]; + + head = (head + 1) % entries.size(); + } + + if (cur.is_end) { + break; + } + + cur.print(); // stdout and stderr + + if (file) { + cur.print(file); + } + } + }); + } + + void pause() { + { + std::lock_guard lock(mtx); + + if (!running) { + return; + } + + running = false; + + // push an entry to signal the worker thread to stop + { + auto & entry = entries[tail]; + entry.is_end = true; + + tail = (tail + 1) % entries.size(); + } + cv.notify_one(); + } + + thrd.join(); + } + + void set_file(const char * path) { + pause(); + + if (file) { + fclose(file); + } + + if (path) { + file = fopen(path, "w"); + } else { + file = nullptr; + } + + resume(); + } + + void set_colors(bool colors) { + pause(); + + if (colors) { + g_col[COMMON_LOG_COL_DEFAULT] = LOG_COL_DEFAULT; + g_col[COMMON_LOG_COL_BOLD] = LOG_COL_BOLD; + g_col[COMMON_LOG_COL_RED] = LOG_COL_RED; + g_col[COMMON_LOG_COL_GREEN] = LOG_COL_GREEN; + g_col[COMMON_LOG_COL_YELLOW] = LOG_COL_YELLOW; + g_col[COMMON_LOG_COL_BLUE] = LOG_COL_BLUE; + g_col[COMMON_LOG_COL_MAGENTA] = LOG_COL_MAGENTA; + g_col[COMMON_LOG_COL_CYAN] = LOG_COL_CYAN; + g_col[COMMON_LOG_COL_WHITE] = LOG_COL_WHITE; + } else { + for (size_t i = 0; i < std::size(g_col); i++) { + g_col[i] = ""; + } + } + + resume(); + } + + void set_prefix(bool prefix) { + std::lock_guard lock(mtx); + + this->prefix = prefix; + } + + void set_timestamps(bool timestamps) { + std::lock_guard lock(mtx); + + this->timestamps = timestamps; + } +}; + +// +// public API +// + +struct common_log * common_log_init() { + return new common_log; +} + +struct common_log * common_log_main() { + // We intentionally leak (i.e. do not delete) the logger singleton because + // common_log destructor called at DLL teardown phase will cause hanging on Windows. + // OS will release resources anyway so it should not be a significant issue, + // though this design may cause logs to be lost if not flushed before the program exits. + // Refer to https://github.com/ggml-org/llama.cpp/issues/22142 for details. + static struct common_log * log; + static std::once_flag init_flag; + std::call_once(init_flag, [&]() { + log = new common_log; + // Set default to auto-detect colors + log->set_colors(tty_can_use_colors()); + }); + + return log; +} + +void common_log_pause(struct common_log * log) { + log->pause(); +} + +void common_log_resume(struct common_log * log) { + log->resume(); +} + +void common_log_free(struct common_log * log) { + delete log; +} + +void common_log_add(struct common_log * log, enum ggml_log_level level, const char * fmt, ...) { + va_list args; + va_start(args, fmt); + log->add(level, fmt, args); + va_end(args); +} + +void common_log_set_file(struct common_log * log, const char * file) { + log->set_file(file); +} + +void common_log_set_colors(struct common_log * log, log_colors colors) { + if (colors == LOG_COLORS_AUTO) { + log->set_colors(tty_can_use_colors()); + return; + } + + if (colors == LOG_COLORS_DISABLED) { + log->set_colors(false); + return; + } + + GGML_ASSERT(colors == LOG_COLORS_ENABLED); + log->set_colors(true); +} + +void common_log_set_prefix(struct common_log * log, bool prefix) { + log->set_prefix(prefix); +} + +void common_log_set_timestamps(struct common_log * log, bool timestamps) { + log->set_timestamps(timestamps); +} + +void common_log_flush(struct common_log * log) { + log->pause(); + log->resume(); +} + +void common_log_default_callback(enum ggml_log_level level, const char * text, void * /*user_data*/) { + auto verbosity = llama_common_log_level_to_verbosity_rust((int) level); + if (verbosity <= common_log_verbosity_thold) { + common_log_add(common_log_main(), level, "%s", text); + } +} + +#include "common.h.inc" + +#include +#include +#include + +// only allow a subset of args for remote presets for security reasons +// do not add more args unless absolutely necessary +// args that output to files are strictly prohibited +static std::set get_remote_preset_whitelist(const std::map & key_to_opt) { + static const std::set allowed_options = { + "model-url", + "hf-repo", + "hf-repo-draft", + "hf-repo-v", // vocoder + "hf-file-v", // vocoder + "mmproj-url", + "pooling", + "jinja", + "batch-size", + "ubatch-size", + "cache-reuse", + "chat-template-kwargs", + "mmap", + // note: sampling params are automatically allowed by default + // negated args will be added automatically if the positive arg is specified above + }; + + std::set allowed_keys; + + for (const auto & it : key_to_opt) { + const std::string & key = it.first; + const common_arg & opt = it.second; + if (allowed_options.find(key) != allowed_options.end() || opt.is_sampling) { + allowed_keys.insert(key); + // also add variant keys (args without leading dashes and env vars) + for (const auto & arg : opt.get_args()) { + allowed_keys.insert(common_rust_string_to_std( + llama_common_rm_leading_dashes_rust((const uint8_t *) arg.data(), arg.size()))); + } + for (const auto & env : opt.get_env()) { + allowed_keys.insert(env); + } + } + } + + return allowed_keys; +} + +std::vector common_preset::to_args(const std::string & bin_path) const { + std::vector args; + + if (!bin_path.empty()) { + args.push_back(bin_path); + } + + for (const auto & [opt, value] : options) { + if (opt.is_preset_only) { + continue; // skip preset-only options (they are not CLI args) + } + + // use the last arg as the main arg (i.e. --long-form) + args.push_back(opt.args.back()); + + // handle value(s) + if (opt.value_hint == nullptr && opt.value_hint_2 == nullptr) { + // flag option, no value + if (common_arg_utils::is_falsey(value)) { + // use negative arg if available + if (!opt.args_neg.empty()) { + args.back() = opt.args_neg.back(); + } else { + // otherwise, skip the flag + // TODO: maybe throw an error instead? + args.pop_back(); + } + } + } + if (opt.value_hint != nullptr) { + // single value + args.push_back(value); + } + if (opt.value_hint != nullptr && opt.value_hint_2 != nullptr) { + throw std::runtime_error(string_format( + "common_preset::to_args(): option '%s' has two values, which is not supported yet", + opt.args.back() + )); + } + } + + return args; +} + +std::string common_preset::to_ini() const { + std::ostringstream ss; + + ss << "[" << name << "]\n"; + for (const auto & [opt, value] : options) { + auto espaced_value = value; + string_replace_all(espaced_value, "\n", "\\\n"); + const std::string & arg = opt.args.back(); + ss << common_rust_string_to_std( + llama_common_rm_leading_dashes_rust((const uint8_t *) arg.data(), arg.size())) << " = "; + ss << espaced_value << "\n"; + } + ss << "\n"; + + return ss.str(); +} + +void common_preset::set_option(const common_preset_context & ctx, const std::string & env, const std::string & value) { + // try if option exists, update it + for (auto & [opt, val] : options) { + if (opt.env && env == opt.env) { + val = value; + return; + } + } + // if option does not exist, we need to add it + if (ctx.key_to_opt.find(env) == ctx.key_to_opt.end()) { + throw std::runtime_error(string_format( + "%s: option with env '%s' not found in ctx_params", + __func__, env.c_str() + )); + } + options[ctx.key_to_opt.at(env)] = value; +} + +void common_preset::unset_option(const std::string & env) { + for (auto it = options.begin(); it != options.end(); ) { + const common_arg & opt = it->first; + if (opt.env && env == opt.env) { + it = options.erase(it); + return; + } else { + ++it; + } + } +} + +bool common_preset::get_option(const std::string & env, std::string & value) const { + for (const auto & [opt, val] : options) { + if (opt.env && env == opt.env) { + value = val; + return true; + } + } + return false; +} + +void common_preset::merge(const common_preset & other) { + for (const auto & [opt, val] : other.options) { + options[opt] = val; // overwrite existing options + } +} + +void common_preset::apply_to_params(common_params & params) const { + for (const auto & [opt, val] : options) { + // apply each option to params + if (opt.handler_string) { + opt.handler_string(params, val); + } else if (opt.handler_int) { + opt.handler_int(params, std::stoi(val)); + } else if (opt.handler_bool) { + opt.handler_bool(params, common_arg_utils::is_truthy(val)); + } else if (opt.handler_str_str) { + // not supported yet + throw std::runtime_error(string_format( + "%s: option with two values is not supported yet", + __func__ + )); + } else if (opt.handler_void) { + opt.handler_void(params); + } else { + GGML_ABORT("unknown handler type"); + } + } +} + +static std::map> parse_ini_from_file(const std::string & path) { + std::map> parsed; + + if (!std::filesystem::exists(path)) { + throw std::runtime_error("preset file does not exist: " + path); + } + + std::ifstream file(path); + if (!file.good()) { + throw std::runtime_error("failed to open server preset file: " + path); + } + + std::string contents((std::istreambuf_iterator(file)), std::istreambuf_iterator()); + + static const auto parser = build_peg_parser([](auto & p) { + // newline ::= "\r\n" / "\n" / "\r" + auto newline = p.rule("newline", p.literal("\r\n") | p.literal("\n") | p.literal("\r")); + + // ws ::= [ \t]* + auto ws = p.rule("ws", p.chars("[ \t]", 0, -1)); + + // comment ::= [;#] (!newline .)* + auto comment = p.rule("comment", p.chars("[;#]", 1, 1) + p.zero_or_more(p.negate(newline) + p.any())); + + // eol ::= ws comment? (newline / EOF) + auto eol = p.rule("eol", ws + p.optional(comment) + (newline | p.end())); + + // ident ::= [a-zA-Z_] [a-zA-Z0-9_.-]* + auto ident = p.rule("ident", p.chars("[a-zA-Z_]", 1, 1) + p.chars("[a-zA-Z0-9_.-]", 0, -1)); + + // value ::= (!eol-start .)* + auto eol_start = p.rule("eol-start", ws + (p.chars("[;#]", 1, 1) | newline | p.end())); + auto value = p.rule("value", p.zero_or_more(p.negate(eol_start) + p.any())); + + // header-line ::= "[" ws ident ws "]" eol + auto header_line = p.rule("header-line", "[" + ws + p.tag("section-name", p.chars("[^]]")) + ws + "]" + eol); + + // kv-line ::= ident ws "=" ws value eol + auto kv_line = p.rule("kv-line", p.tag("key", ident) + ws + "=" + ws + p.tag("value", value) + eol); + + // comment-line ::= ws comment (newline / EOF) + auto comment_line = p.rule("comment-line", ws + comment + (newline | p.end())); + + // blank-line ::= ws (newline / EOF) + auto blank_line = p.rule("blank-line", ws + (newline | p.end())); + + // line ::= header-line / kv-line / comment-line / blank-line + auto line = p.rule("line", header_line | kv_line | comment_line | blank_line); + + // ini ::= line* EOF + auto ini = p.rule("ini", p.zero_or_more(line) + p.end()); + + return ini; + }); + + common_peg_parse_context ctx(contents); + const auto result = parser.parse(ctx); + if (!result.success()) { + throw std::runtime_error("failed to parse server config file: " + path); + } + + std::string current_section = COMMON_PRESET_DEFAULT_NAME; + std::string current_key; + + ctx.ast.visit(result, [&](const auto & node) { + if (node.tag == "section-name") { + const std::string section = std::string(node.text); + current_section = section; + parsed[current_section] = {}; + } else if (node.tag == "key") { + const std::string key = std::string(node.text); + current_key = key; + } else if (node.tag == "value" && !current_key.empty() && !current_section.empty()) { + parsed[current_section][current_key] = std::string(node.text); + current_key.clear(); + } + }); + + return parsed; +} + +static std::map get_map_key_opt(common_params_context & ctx_params) { + std::map mapping; + for (const auto & opt : ctx_params.options) { + for (const auto & env : opt.get_env()) { + mapping[env] = opt; + } + for (const auto & arg : opt.get_args()) { + mapping[common_rust_string_to_std( + llama_common_rm_leading_dashes_rust((const uint8_t *) arg.data(), arg.size()))] = opt; + } + } + return mapping; +} + +static std::string parse_bool_arg(const common_arg & arg, const std::string & key, const std::string & value) { + // if this is a negated arg, we need to reverse the value + for (const auto & neg_arg : arg.args_neg) { + const size_t neg_arg_len = std::strlen(neg_arg); + if (common_rust_string_to_std( + llama_common_rm_leading_dashes_rust((const uint8_t *) neg_arg, neg_arg_len)) == key) { + return common_arg_utils::is_truthy(value) ? "false" : "true"; + } + } + // otherwise, not negated + return value; +} + +common_preset_context::common_preset_context(llama_example ex, bool only_remote_allowed) + : ctx_params(common_params_parser_init(default_params, ex)) { + common_params_add_preset_options(ctx_params.options); + key_to_opt = get_map_key_opt(ctx_params); + + // setup allowed keys if only_remote_allowed is true + if (only_remote_allowed) { + filter_allowed_keys = true; + allowed_keys = get_remote_preset_whitelist(key_to_opt); + } +} + +common_presets common_preset_context::load_from_ini(const std::string & path, common_preset & global) const { + common_presets out; + auto ini_data = parse_ini_from_file(path); + + for (auto section : ini_data) { + common_preset preset; + if (section.first.empty()) { + preset.name = COMMON_PRESET_DEFAULT_NAME; + } else { + preset.name = section.first; + } + LOG_DBG("loading preset: %s\n", preset.name.c_str()); + for (const auto & [key, value] : section.second) { + if (key == "version") { + // skip version key (reserved for future use) + continue; + } + + LOG_DBG("option: %s = %s\n", key.c_str(), value.c_str()); + if (filter_allowed_keys && allowed_keys.find(key) == allowed_keys.end()) { + throw std::runtime_error(string_format( + "option '%s' is not allowed in remote presets", + key.c_str() + )); + } + if (key_to_opt.find(key) != key_to_opt.end()) { + const auto & opt = key_to_opt.at(key); + if (!opt.args_neg.empty()) { + preset.options[opt] = parse_bool_arg(opt, key, value); + } else { + preset.options[opt] = value; + } + LOG_DBG("accepted option: %s = %s\n", key.c_str(), preset.options[opt].c_str()); + } else { + throw std::runtime_error(string_format( + "option '%s' not recognized in preset '%s'", + key.c_str(), preset.name.c_str() + )); + } + } + + if (preset.name == "*") { + // handle global preset + global = preset; + } else { + out[preset.name] = preset; + } + } + + return out; +} + +common_presets common_preset_context::load_from_cache() const { + common_presets out; + + auto cached_models = common_list_cached_models(); + for (const auto & model : cached_models) { + common_preset preset; + preset.name = model.to_string(); + preset.set_option(*this, "LLAMA_ARG_HF_REPO", model.to_string()); + out[preset.name] = preset; + } + + return out; +} + +struct local_model { + std::string name; + std::string path; + std::string path_mmproj; +}; + +common_presets common_preset_context::load_from_models_dir(const std::string & models_dir) const { + if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) { + throw std::runtime_error(string_format("error: '%s' does not exist or is not a directory\n", models_dir.c_str())); + } + + std::vector models; + auto scan_subdir = [&models](const std::string & subdir_path, const std::string & name) { + auto files = fs_list(subdir_path, false); + common_file_info model_file; + common_file_info first_shard_file; + common_file_info mmproj_file; + for (const auto & file : files) { + if (string_ends_with(file.name, ".gguf")) { + if (file.name.find("mmproj") != std::string::npos) { + mmproj_file = file; + } else if (file.name.find("-00001-of-") != std::string::npos) { + first_shard_file = file; + } else { + model_file = file; + } + } + } + // single file model + local_model model{ + /* name */ name, + /* path */ first_shard_file.path.empty() ? model_file.path : first_shard_file.path, + /* path_mmproj */ mmproj_file.path // can be empty + }; + if (!model.path.empty()) { + models.push_back(model); + } + }; + + auto files = fs_list(models_dir, true); + for (const auto & file : files) { + if (file.is_dir) { + scan_subdir(file.path, file.name); + } else if (string_ends_with(file.name, ".gguf")) { + // single file model + std::string name = file.name; + string_replace_all(name, ".gguf", ""); + local_model model{ + /* name */ name, + /* path */ file.path, + /* path_mmproj */ "" + }; + models.push_back(model); + } + } + + // convert local models to presets + common_presets out; + for (const auto & model : models) { + common_preset preset; + preset.name = model.name; + preset.set_option(*this, "LLAMA_ARG_MODEL", model.path); + if (!model.path_mmproj.empty()) { + preset.set_option(*this, "LLAMA_ARG_MMPROJ", model.path_mmproj); + } + out[preset.name] = preset; + } + + return out; +} + +common_preset common_preset_context::load_from_args(int argc, char ** argv) const { + common_preset preset; + preset.name = COMMON_PRESET_DEFAULT_NAME; + + bool ok = common_params_to_map(argc, argv, ctx_params.ex, preset.options); + if (!ok) { + throw std::runtime_error("failed to parse CLI arguments into preset"); + } + + return preset; +} + +common_presets common_preset_context::cascade(const common_presets & base, const common_presets & added) const { + common_presets out = base; // copy + for (const auto & [name, preset_added] : added) { + if (out.find(name) != out.end()) { + // if exists, merge + common_preset & target = out[name]; + target.merge(preset_added); + } else { + // otherwise, add directly + out[name] = preset_added; + } + } + return out; +} + +common_presets common_preset_context::cascade(const common_preset & base, const common_presets & presets) const { + common_presets out; + for (const auto & [name, preset] : presets) { + common_preset tmp = base; // copy + tmp.name = name; + tmp.merge(preset); + out[name] = std::move(tmp); + } + return out; +} + +#include "common.h.inc" + +#include "../src/llama-impl.h.inc" + +#include +#include +#include +#include +#include +#include +#include + +// this enum is only used in llama_params_fit_impl but needs to be defined outside of it to fix a Windows compilation issue +// enum to identify part of a layer for distributing its tensors: +enum common_layer_fraction_t { + LAYER_FRACTION_NONE = 0, // nothing + LAYER_FRACTION_ATTN = 1, // attention + LAYER_FRACTION_UP = 2, // attention + up + LAYER_FRACTION_GATE = 3, // attention + up + gate + LAYER_FRACTION_MOE = 4, // everything but sparse MoE weights +}; + +class common_params_fit_exception : public std::runtime_error { + using std::runtime_error::runtime_error; +}; + +static std::vector common_get_device_memory_data( + const char * path_model, + const llama_model_params * mparams, + const llama_context_params * cparams, + std::vector & devs, + uint32_t & hp_ngl, + uint32_t & hp_n_ctx_train, + uint32_t & hp_n_expert, + ggml_log_level log_level) { + struct user_data_t { + struct { + ggml_log_callback callback; + void * user_data; + } original_logger; + ggml_log_level min_level; // prints below this log level go to debug log + }; + user_data_t ud; + llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data); + ud.min_level = log_level; + + llama_log_set([](ggml_log_level level, const char * text, void * user_data) { + const user_data_t * ud = (const user_data_t *) user_data; + const ggml_log_level level_eff = level >= ud->min_level ? level : GGML_LOG_LEVEL_DEBUG; + ud->original_logger.callback(level_eff, text, ud->original_logger.user_data); + }, &ud); + + llama_model_params mparams_copy = *mparams; + mparams_copy.no_alloc = true; + mparams_copy.use_mmap = false; + mparams_copy.use_mlock = false; + + llama_model * model = llama_model_load_from_file(path_model, mparams_copy); + if (model == nullptr) { + llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); + throw std::runtime_error("failed to load model"); + } + + llama_context * ctx = llama_init_from_model(model, *cparams); + if (ctx == nullptr) { + llama_model_free(model); + llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); + throw std::runtime_error("failed to create llama_context from model"); + } + + const size_t nd = llama_model_n_devices(model); + std::vector ret(nd + 1); + + llama_memory_breakdown memory_breakdown = llama_get_memory_breakdown(ctx); + + for (const auto & [buft, mb] : memory_breakdown) { + if (ggml_backend_buft_is_host(buft)) { + ret.back().mb.model += mb.model; + ret.back().mb.context += mb.context; + ret.back().mb.compute += mb.compute; + continue; + } + + ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft); + if (!dev) { + continue; + } + for (size_t i = 0; i < nd; i++) { + if (dev == llama_model_get_device(model, i)) { + ret[i].mb.model += mb.model; + ret[i].mb.context += mb.context; + ret[i].mb.compute += mb.compute; + break; + } + } + } + + { + ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); + if (cpu_dev == nullptr) { + throw std::runtime_error("no CPU backend found"); + } + size_t free; + size_t total; + ggml_backend_dev_memory(cpu_dev, &free, &total); + ret.back().free = free; + ret.back().total = total; + } + for (size_t i = 0; i < nd; i++) { + size_t free; + size_t total; + ggml_backend_dev_memory(llama_model_get_device(model, i), &free, &total); + + // devices can return 0 bytes for free and total memory if they do not + // have any to report. in this case, we will use the host memory as a fallback + // fixes: https://github.com/ggml-org/llama.cpp/issues/18577 + if (free == 0 && total == 0) { + free = ret.back().free; + total = ret.back().total; + } + ret[i].free = free; + ret[i].total = total; + } + + devs.clear(); + for (int i = 0; i < llama_model_n_devices(model); i++) { + devs.push_back(llama_model_get_device(model, i)); + } + + hp_ngl = llama_model_n_layer(model); + hp_n_ctx_train = llama_model_n_ctx_train(model); + hp_n_expert = llama_model_n_expert(model); + + common_memory_breakdown_print(ctx); + + llama_free(ctx); + llama_model_free(model); + llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); + + return ret; +} + +static void common_params_fit_impl( + const char * path_model, struct llama_model_params * mparams, struct llama_context_params * cparams, + float * tensor_split, struct llama_model_tensor_buft_override * tensor_buft_overrides, + size_t * margins_s, uint32_t n_ctx_min, enum ggml_log_level log_level) { + if (mparams->split_mode == LLAMA_SPLIT_MODE_TENSOR) { + throw common_params_fit_exception("llama_params_fit is not implemented for SPLIT_MODE_TENSOR, abort"); + } + constexpr int64_t MiB = 1024*1024; + typedef std::vector dmds_t; + const llama_model_params default_mparams = llama_model_default_params(); + + std::vector devs; + uint32_t hp_ngl = 0; // hparams.n_gpu_layers + uint32_t hp_nct = 0; // hparams.n_ctx_train + uint32_t hp_nex = 0; // hparams.n_expert + + // step 1: get data for default parameters and check whether any changes are necessary in the first place + + LOG_INF("%s: getting device memory data for initial parameters:\n", __func__); + const dmds_t dmds_full = common_get_device_memory_data(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + const size_t nd = devs.size(); // number of devices + + std::vector margins; // this function uses int64_t rather than size_t for memory sizes to more conveniently handle deficits + margins.reserve(nd); + if (nd == 0) { + margins.push_back(margins_s[0]); + } else { + for (size_t id = 0; id < nd; id++) { + margins.push_back(margins_s[id]); + } + } + + std::vector dev_names; + { + dev_names.reserve(nd); + size_t max_length = 0; + for (const auto & dev : devs) { + std::string name = ggml_backend_dev_name(dev); + name += " ("; + name += ggml_backend_dev_description(dev); + name += ")"; + dev_names.push_back(name); + max_length = std::max(max_length, name.length()); + } + for (std::string & dn : dev_names) { + dn.insert(dn.end(), max_length - dn.length(), ' '); + } + } + + int64_t sum_free = 0; + int64_t sum_projected_free = 0; + int64_t sum_projected_used = 0; + int64_t sum_projected_model = 0; + std::vector projected_free_per_device; + projected_free_per_device.reserve(nd); + + if (nd == 0) { + sum_projected_used = dmds_full.back().mb.total(); + sum_free = dmds_full.back().total; + sum_projected_free = sum_free - sum_projected_used; + LOG_INF("%s: projected to use %" PRId64 " MiB of host memory vs. %" PRId64 " MiB of total host memory\n", + __func__, sum_projected_used/MiB, sum_free/MiB); + if (sum_projected_free >= margins[0]) { + LOG_INF("%s: will leave %" PRId64 " >= %" PRId64 " MiB of system memory, no changes needed\n", + __func__, sum_projected_free/MiB, margins[0]/MiB); + return; + } + } else { + if (nd > 1) { + LOG_INF("%s: projected memory use with initial parameters [MiB]:\n", __func__); + } + for (size_t id = 0; id < nd; id++) { + const llama_device_memory_data & dmd = dmds_full[id]; + + const int64_t projected_used = dmd.mb.total(); + const int64_t projected_free = dmd.free - projected_used; + projected_free_per_device.push_back(projected_free); + + sum_free += dmd.free; + sum_projected_used += projected_used; + sum_projected_free += projected_free; + sum_projected_model += dmd.mb.model; + + if (nd > 1) { + LOG_INF("%s: - %s: %6" PRId64 " total, %6" PRId64 " used, %6" PRId64 " free vs. target of %6" PRId64 "\n", + __func__, dev_names[id].c_str(), dmd.total/MiB, projected_used/MiB, projected_free/MiB, margins[id]/MiB); + } + } + assert(sum_free >= 0 && sum_projected_used >= 0); + LOG_INF("%s: projected to use %" PRId64 " MiB of device memory vs. %" PRId64 " MiB of free device memory\n", + __func__, sum_projected_used/MiB, sum_free/MiB); + if (nd == 1) { + if (projected_free_per_device[0] >= margins[0]) { + LOG_INF("%s: will leave %" PRId64 " >= %" PRId64 " MiB of free device memory, no changes needed\n", + __func__, projected_free_per_device[0]/MiB, margins[0]/MiB); + return; + } + } else { + bool changes_needed = false; + for (size_t id = 0; id < nd; id++) { + if (projected_free_per_device[id] < margins[id]) { + changes_needed = true; + break; + } + } + if (!changes_needed) { + LOG_INF("%s: targets for free memory can be met on all devices, no changes needed\n", __func__); + return; + } + } + } + + // step 2: try reducing memory use by reducing the context size + + { + int64_t global_surplus = sum_projected_free; + if (nd == 0) { + global_surplus -= margins[0]; + } else { + for (size_t id = 0; id < nd; id++) { + global_surplus -= margins[id]; + } + } + if (global_surplus < 0) { + if (nd <= 1) { + LOG_INF("%s: cannot meet free memory target of %" PRId64 " MiB, need to reduce device memory by %" PRId64 " MiB\n", + __func__, margins[0]/MiB, -global_surplus/MiB); + } else { + LOG_INF( + "%s: cannot meet free memory targets on all devices, need to use %" PRId64 " MiB less in total\n", + __func__, -global_surplus/MiB); + } + if (cparams->n_ctx == 0) { + if (hp_nct > n_ctx_min) { + int64_t sum_used_target = sum_free; + if (nd == 0) { + sum_used_target -= margins[0]; + } else { + for (size_t id = 0; id < nd; id++) { + sum_used_target -= margins[id]; + } + } + if (nd > 1) { + // for multiple devices we need to be more conservative in terms of how much context we think can fit: + // - for dense models only whole layers can be assigned to devices + // - for MoE models only whole tensors can be assigned to devices, which we estimate to be <= 1/3 of a layer + // - on average we expect a waste of 0.5 layers/tensors per device + // - use slightly more than the expected average for nd devices to be safe + const int64_t model_per_layer = sum_projected_model / std::min(uint32_t(mparams->n_gpu_layers), hp_ngl); + sum_used_target -= (nd + 1) * model_per_layer / (hp_nex == 0 ? 2 : 6); + } + + int64_t sum_projected_used_min_ctx = 0; + cparams->n_ctx = n_ctx_min; + const dmds_t dmds_min_ctx = common_get_device_memory_data(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + if (nd == 0) { + sum_projected_used_min_ctx = dmds_min_ctx.back().mb.total(); + } else { + for (size_t id = 0; id < nd; id++) { + sum_projected_used_min_ctx += dmds_min_ctx[id].mb.total(); + } + } + if (sum_used_target > sum_projected_used_min_ctx) { + // linear interpolation between minimum and maximum context size: + cparams->n_ctx += (hp_nct - n_ctx_min) * (sum_used_target - sum_projected_used_min_ctx) + / (sum_projected_used - sum_projected_used_min_ctx); + cparams->n_ctx = std::max(cparams->n_ctx - cparams->n_ctx % 256, n_ctx_min); // round down context for CUDA backend + + const int64_t bytes_per_ctx = (sum_projected_used - sum_projected_used_min_ctx) / (hp_nct - n_ctx_min); + const int64_t memory_reduction = (hp_nct - cparams->n_ctx) * bytes_per_ctx; + LOG_INF("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n", + __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB); + if (nd <= 1) { + LOG_INF("%s: entire model can be fit by reducing context\n", __func__); + return; + } + LOG_INF("%s: entire model should be fit across devices by reducing context\n", __func__); + } else { + const int64_t memory_reduction = sum_projected_used - sum_projected_used_min_ctx; + LOG_INF("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n", + __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB); + } + } else { + if (n_ctx_min == UINT32_MAX) { + LOG_INF("%s: user has requested full context size of %" PRIu32 " -> no change\n", __func__, hp_nct); + } else { + LOG_INF("%s: default model context size is %" PRIu32 " which is <= the min. context size of %" PRIu32 " -> no change\n", + __func__, hp_nct, n_ctx_min); + } + } + } else { + LOG_INF("%s: context size set by user to %" PRIu32 " -> no change\n", __func__, cparams->n_ctx); + } + } + } + if (nd == 0) { + throw common_params_fit_exception("was unable to fit model into system memory by reducing context, abort"); + } + + if (mparams->n_gpu_layers != default_mparams.n_gpu_layers) { + throw common_params_fit_exception("n_gpu_layers already set by user to " + std::to_string(mparams->n_gpu_layers) + ", abort"); + } + if (nd > 1) { + if (!tensor_split) { + throw common_params_fit_exception("did not provide a buffer to write the tensor_split to, abort"); + } + if (mparams->tensor_split) { + for (size_t id = 0; id < nd; id++) { + if (mparams->tensor_split[id] != 0.0f) { + throw common_params_fit_exception("model_params::tensor_split already set by user, abort"); + } + } + } + if (mparams->split_mode == LLAMA_SPLIT_MODE_ROW) { + throw common_params_fit_exception("changing weight allocation for LLAMA_SPLIT_MODE_ROW not implemented, abort"); + } + } + if (!tensor_buft_overrides) { + throw common_params_fit_exception("did not provide buffer to set tensor_buft_overrides, abort"); + } + if (mparams->tensor_buft_overrides && (mparams->tensor_buft_overrides->pattern || mparams->tensor_buft_overrides->buft)) { + throw common_params_fit_exception("model_params::tensor_buft_overrides already set by user, abort"); + } + + // step 3: iteratively fill the back to front with "dense" layers + // - for a dense model simply fill full layers, giving each device a contiguous slice of the model + // - for a MoE model, same as dense model but with all MoE tensors in system memory + + // utility function that returns a static C string matching the tensors for a specific layer index and layer fraction: + auto get_overflow_pattern = [&](const size_t il, const common_layer_fraction_t lf) -> const char * { + constexpr size_t n_strings = 1000; + if (il >= n_strings) { + throw std::runtime_error("at most " + std::to_string(n_strings) + " model layers are supported"); + } + switch (lf) { + case LAYER_FRACTION_ATTN: { + static std::array patterns; + if (patterns[il].empty()) { + patterns[il] = "blk\\." + std::to_string(il) + "\\.ffn_(gate|up|gate_up|down).*"; + } + return patterns[il].c_str(); + } + case LAYER_FRACTION_UP: { + static std::array patterns; + if (patterns[il].empty()) { + patterns[il] = "blk\\." + std::to_string(il) + "\\.ffn_(gate|gate_up|down).*"; + } + return patterns[il].c_str(); + } + case LAYER_FRACTION_GATE: { + static std::array patterns; + if (patterns[il].empty()) { + patterns[il] = "blk\\." + std::to_string(il) + "\\.ffn_down.*"; + } + return patterns[il].c_str(); + } + case LAYER_FRACTION_MOE: { + static std::array patterns; + if (patterns[il].empty()) { + patterns[il] = "blk\\." + std::to_string(il) + "\\.ffn_(up|down|gate_up|gate)_(ch|)exps"; + } + return patterns[il].c_str(); + } + default: + GGML_ABORT("fatal error"); + } + }; + + struct ngl_t { + uint32_t n_layer = 0; // number of total layers + uint32_t n_part = 0; // number of partial layers, <= n_layer + + // for the first partial layer varying parts can overflow, all further layers use LAYER_FRACTION_MOE: + common_layer_fraction_t overflow_type = LAYER_FRACTION_MOE; + + uint32_t n_full() const { + assert(n_layer >= n_part); + return n_layer - n_part; + } + }; + + const size_t ntbo = llama_max_tensor_buft_overrides(); + + // utility function to set n_gpu_layers and tensor_split + auto set_ngl_tensor_split_tbo = [&]( + const std::vector & ngl_per_device, + const std::vector & overflow_bufts, + llama_model_params & mparams) { + mparams.n_gpu_layers = 0; + for (size_t id = 0; id < nd; id++) { + mparams.n_gpu_layers += ngl_per_device[id].n_layer; + if (nd > 1) { + tensor_split[id] = ngl_per_device[id].n_layer; + } + } + assert(uint32_t(mparams.n_gpu_layers) <= hp_ngl + 1); + uint32_t il0 = hp_ngl + 1 - mparams.n_gpu_layers; // start index for tensor buft overrides + + mparams.tensor_split = tensor_split; + + size_t itbo = 0; + for (size_t id = 0; id < nd; id++) { + il0 += ngl_per_device[id].n_full(); + for (uint32_t il = il0; il < il0 + ngl_per_device[id].n_part; il++) { + if (itbo + 1 >= ntbo) { + tensor_buft_overrides[itbo].pattern = nullptr; + tensor_buft_overrides[itbo].buft = nullptr; + itbo++; + mparams.tensor_buft_overrides = tensor_buft_overrides; + throw common_params_fit_exception("llama_max_tensor_buft_overrides() == " + + std::to_string(ntbo) + " is insufficient for model"); + } + tensor_buft_overrides[itbo].pattern = get_overflow_pattern(il, il == il0 ? ngl_per_device[id].overflow_type : LAYER_FRACTION_MOE); + tensor_buft_overrides[itbo].buft = il == il0 ? overflow_bufts[id] : ggml_backend_cpu_buffer_type(); + itbo++; + } + il0 += ngl_per_device[id].n_part; + } + tensor_buft_overrides[itbo].pattern = nullptr; + tensor_buft_overrides[itbo].buft = nullptr; + itbo++; + mparams.tensor_buft_overrides = tensor_buft_overrides; + }; + + // utility function that returns the memory use per device for given numbers of layers per device + auto get_memory_for_layers = [&]( + const char * func_name, + const std::vector & ngl_per_device, + const std::vector & overflow_bufts) -> std::vector { + llama_model_params mparams_copy = *mparams; + set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, mparams_copy); + + const dmds_t dmd_nl = common_get_device_memory_data( + path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + + LOG_INF("%s: memory for test allocation by device:\n", func_name); + for (size_t id = 0; id < nd; id++) { + const ngl_t & n = ngl_per_device[id]; + LOG_INF( + "%s: id=%zu, n_layer=%2" PRIu32 ", n_part=%2" PRIu32 ", overflow_type=%d, mem=%6" PRId64 " MiB\n", + func_name, id, n.n_layer, n.n_part, int(n.overflow_type), dmd_nl[id].mb.total()/MiB); + } + + std::vector ret; + ret.reserve(nd); + for (size_t id = 0; id < nd; id++) { + ret.push_back(dmd_nl[id].mb.total()); + } + return ret; + }; + + int64_t global_surplus_cpu_moe = 0; + if (hp_nex > 0) { + const static std::string pattern_moe_all = "blk\\.\\d+\\.ffn_(up|down|gate_up|gate)_(ch|)exps"; // matches all MoE tensors + ggml_backend_buffer_type_t cpu_buft = ggml_backend_cpu_buffer_type(); + tensor_buft_overrides[0] = {pattern_moe_all.c_str(), cpu_buft}; + tensor_buft_overrides[1] = {nullptr, nullptr}; + mparams->tensor_buft_overrides = tensor_buft_overrides; + + LOG_INF("%s: getting device memory data with all MoE tensors moved to system memory:\n", __func__); + const dmds_t dmds_cpu_moe = common_get_device_memory_data( + path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); + + for (size_t id = 0; id < nd; id++) { + global_surplus_cpu_moe += dmds_cpu_moe[id].free; + global_surplus_cpu_moe -= int64_t(dmds_cpu_moe[id].mb.total()) + margins[id]; + } + + if (global_surplus_cpu_moe > 0) { + LOG_INF("%s: with only dense weights in device memory there is a total surplus of %" PRId64 " MiB\n", + __func__, global_surplus_cpu_moe/MiB); + } else { + LOG_INF("%s: with only dense weights in device memory there is still a total deficit of %" PRId64 " MiB\n", + __func__, -global_surplus_cpu_moe/MiB); + } + + // reset + tensor_buft_overrides[0] = {nullptr, nullptr}; + mparams->tensor_buft_overrides = tensor_buft_overrides; + } + + std::vector targets; // maximum acceptable memory use per device + targets.reserve(nd); + for (size_t id = 0; id < nd; id++) { + targets.push_back(dmds_full[id].free - margins[id]); + LOG_INF("%s: id=%zu, target=%" PRId64 " MiB\n", __func__, id, targets[id]/MiB); + } + + std::vector overflow_bufts; // which bufts the first partial layer of a device overflows to: + overflow_bufts.reserve(nd); + for (size_t id = 0; id < nd; id++) { + overflow_bufts.push_back(ggml_backend_cpu_buffer_type()); + } + + std::vector ngl_per_device(nd); + std::vector mem = get_memory_for_layers(__func__, ngl_per_device, overflow_bufts); + + // optimize the number of layers per device using the method of false position: + // - ngl_per_device has 0 layers for each device, lower bound + // - try a "high" configuration where a device is given all unassigned layers + // - interpolate the memory use / layer between low and high linearly to get a guess where it meets our target + // - check memory use of our guess, replace either the low or high bound + // - once we only have a difference of a single layer, stop and return the lower bound that just barely still fits + // - the last device has the output layer, which cannot be a partial layer + if (hp_nex == 0) { + LOG_INF("%s: filling dense layers back-to-front:\n", __func__); + } else { + LOG_INF("%s: filling dense-only layers back-to-front:\n", __func__); + } + for (int id = nd - 1; id >= 0; id--) { + uint32_t n_unassigned = hp_ngl + 1; + for (size_t jd = id + 1; jd < nd; ++jd) { + assert(n_unassigned >= ngl_per_device[jd].n_layer); + n_unassigned -= ngl_per_device[jd].n_layer; + } + + std::vector ngl_per_device_high = ngl_per_device; + ngl_per_device_high[id].n_layer = n_unassigned; + if (hp_nex > 0) { + ngl_per_device_high[id].n_part = size_t(id) < nd - 1 ? ngl_per_device_high[id].n_layer : ngl_per_device_high[id].n_layer - 1; + } + if (ngl_per_device_high[id].n_layer > 0) { + std::vector mem_high = get_memory_for_layers(__func__, ngl_per_device_high, overflow_bufts); + if (mem_high[id] > targets[id]) { + assert(ngl_per_device_high[id].n_layer > ngl_per_device[id].n_layer); + uint32_t delta = ngl_per_device_high[id].n_layer - ngl_per_device[id].n_layer; + LOG_INF("%s: start filling device %" PRIu32 ", delta=%" PRIu32 "\n", __func__, id, delta); + while (delta > 1) { + uint32_t step_size = int64_t(delta) * (targets[id] - mem[id]) / (mem_high[id] - mem[id]); + step_size = std::max(step_size, uint32_t(1)); + step_size = std::min(step_size, delta - 1); + + std::vector ngl_per_device_test = ngl_per_device; + ngl_per_device_test[id].n_layer += step_size; + if (hp_nex) { + ngl_per_device_test[id].n_part += size_t(id) == nd - 1 && ngl_per_device_test[id].n_part == 0 ? + step_size - 1 : step_size; // the first layer is the output layer which must always be full + } + const std::vector mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts); + + if (mem_test[id] <= targets[id]) { + ngl_per_device = ngl_per_device_test; + mem = mem_test; + LOG_INF("%s: set ngl_per_device[%d].n_layer=%" PRIu32 "\n", __func__, id, ngl_per_device[id].n_layer); + } else { + ngl_per_device_high = ngl_per_device_test; + mem_high = mem_test; + LOG_INF("%s: set ngl_per_device_high[%d].n_layer=%" PRIu32 "\n", __func__, id, ngl_per_device_high[id].n_layer); + } + delta = ngl_per_device_high[id].n_layer - ngl_per_device[id].n_layer; + } + } else { + assert(ngl_per_device_high[id].n_layer == n_unassigned); + ngl_per_device = ngl_per_device_high; + mem = mem_high; + LOG_INF("%s: set ngl_per_device[%d].n_layer=%" PRIu32 "\n", __func__, id, ngl_per_device[id].n_layer); + } + } + + const int64_t projected_margin = dmds_full[id].free - mem[id]; + LOG_INF( + "%s: - %s: %2" PRIu32 " layers, %6" PRId64 " MiB used, %6" PRId64 " MiB free\n", + __func__, dev_names[id].c_str(), ngl_per_device[id].n_layer, mem[id]/MiB, projected_margin/MiB); + } + if (hp_nex == 0 || global_surplus_cpu_moe <= 0) { + set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, *mparams); + return; + } + + // step 4: for a MoE model where all dense tensors fit, + // convert the dense-only layers in the back to full layers in the front until all devices are full + // essentially the same procedure as for the dense-only layers except front-to-back + // also, try fitting at least part of one more layer to reduce waste for "small" GPUs with e.g. 24 GiB VRAM + + size_t id_dense_start = nd; + for (int id = nd - 1; id >= 0; id--) { + if (ngl_per_device[id].n_layer > 0) { + id_dense_start = id; + continue; + } + break; + } + assert(id_dense_start < nd); + + LOG_INF("%s: converting dense-only layers to full layers and filling them front-to-back with overflow to next device/system memory:\n", __func__); + for (size_t id = 0; id <= id_dense_start && id_dense_start < nd; id++) { + std::vector ngl_per_device_high = ngl_per_device; + for (size_t jd = id_dense_start; jd < nd; jd++) { + const uint32_t n_layer_move = jd < nd - 1 ? ngl_per_device_high[jd].n_layer : ngl_per_device_high[jd].n_layer - 1; + ngl_per_device_high[id].n_layer += n_layer_move; + ngl_per_device_high[jd].n_layer -= n_layer_move; + ngl_per_device_high[jd].n_part = 0; + } + size_t id_dense_start_high = nd - 1; + std::vector mem_high = get_memory_for_layers(__func__, ngl_per_device_high, overflow_bufts); + + if (mem_high[id] > targets[id]) { + assert(ngl_per_device_high[id].n_full() >= ngl_per_device[id].n_full()); + uint32_t delta = ngl_per_device_high[id].n_full() - ngl_per_device[id].n_full(); + while (delta > 1) { + uint32_t step_size = int64_t(delta) * (targets[id] - mem[id]) / (mem_high[id] - mem[id]); + step_size = std::max(step_size, uint32_t(1)); + step_size = std::min(step_size, delta - 1); + + std::vector ngl_per_device_test = ngl_per_device; + size_t id_dense_start_test = id_dense_start; + uint32_t n_converted_test = 0; + for (;id_dense_start_test < nd; id_dense_start_test++) { + const uint32_t n_convert_jd = std::min(step_size - n_converted_test, ngl_per_device_test[id_dense_start_test].n_part); + ngl_per_device_test[id_dense_start_test].n_layer -= n_convert_jd; + ngl_per_device_test[id_dense_start_test].n_part -= n_convert_jd; + ngl_per_device_test[id].n_layer += n_convert_jd; + n_converted_test += n_convert_jd; + + if (ngl_per_device_test[id_dense_start_test].n_part > 0) { + break; + } + } + const std::vector mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts); + + if (mem_test[id] <= targets[id]) { + ngl_per_device = ngl_per_device_test; + mem = mem_test; + id_dense_start = id_dense_start_test; + LOG_INF("%s: set ngl_per_device[%zu].(n_layer, n_part)=(%" PRIu32 ", %" PRIu32 "), id_dense_start=%zu\n", + __func__, id, ngl_per_device[id].n_layer, ngl_per_device[id].n_part, id_dense_start); + } else { + ngl_per_device_high = ngl_per_device_test; + mem_high = mem_test; + id_dense_start_high = id_dense_start_test; + LOG_INF("%s: set ngl_per_device_high[%zu].(n_layer, n_part)=(%" PRIu32 ", %" PRIu32 "), id_dense_start_high=%zu\n", + __func__, id, ngl_per_device_high[id].n_layer, ngl_per_device_high[id].n_part, id_dense_start_high); + } + assert(ngl_per_device_high[id].n_full() >= ngl_per_device[id].n_full()); + delta = ngl_per_device_high[id].n_full() - ngl_per_device[id].n_full(); + } + } else { + ngl_per_device = ngl_per_device_high; + mem = mem_high; + id_dense_start = id_dense_start_high; + LOG_INF("%s: set ngl_per_device[%zu].(n_layer, n_part)=(%" PRIu32 ", %" PRIu32 "), id_dense_start=%zu\n", + __func__, id, ngl_per_device[id].n_layer, ngl_per_device[id].n_part, id_dense_start); + } + + // try to fit at least part of one more layer + if (ngl_per_device[id_dense_start].n_layer > (id < nd - 1 ? 0 : 1)) { + std::vector ngl_per_device_test = ngl_per_device; + size_t id_dense_start_test = id_dense_start; + ngl_per_device_test[id_dense_start_test].n_layer--; + ngl_per_device_test[id_dense_start_test].n_part--; + ngl_per_device_test[id].n_layer++; + ngl_per_device_test[id].n_part++; + if (ngl_per_device_test[id_dense_start_test].n_part == 0) { + id_dense_start_test++; + } + ngl_per_device_test[id].overflow_type = LAYER_FRACTION_UP; + std::vector overflow_bufts_test = overflow_bufts; + if (id < nd - 1) { + overflow_bufts_test[id] = ggml_backend_dev_buffer_type(devs[id + 1]); + } + LOG_INF("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_UP\n", __func__); + std::vector mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts_test); + if (mem_test[id] < targets[id] && (id + 1 == nd || mem_test[id + 1] < targets[id + 1])) { + ngl_per_device = ngl_per_device_test; + overflow_bufts = overflow_bufts_test; + mem = mem_test; + id_dense_start = id_dense_start_test; + LOG_INF("%s: set ngl_per_device[%zu].(n_layer, n_part, overflow_type)=(%" PRIu32 ", %" PRIu32 ", UP), id_dense_start=%zu\n", + __func__, id, ngl_per_device[id].n_layer, ngl_per_device[id].n_part, id_dense_start); + + ngl_per_device_test[id].overflow_type = LAYER_FRACTION_GATE; + LOG_INF("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_GATE\n", __func__); + mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts_test); + if (mem_test[id] < targets[id] && (id + 1 == nd || mem_test[id + 1] < targets[id + 1])) { + ngl_per_device = ngl_per_device_test; + overflow_bufts = overflow_bufts_test; + mem = mem_test; + id_dense_start = id_dense_start_test; + LOG_INF("%s: set ngl_per_device[%zu].(n_layer, n_part, overflow_type)=(%" PRIu32 ", %" PRIu32 ", GATE), id_dense_start=%zu\n", + __func__, id, ngl_per_device[id].n_layer, ngl_per_device[id].n_part, id_dense_start); + } + } else { + ngl_per_device_test[id].overflow_type = LAYER_FRACTION_ATTN; + LOG_INF("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_ATTN\n", __func__); + mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts_test); + if (mem_test[id] < targets[id] && (id + 1 == nd || mem_test[id + 1] < targets[id + 1])) { + ngl_per_device = ngl_per_device_test; + overflow_bufts = overflow_bufts_test; + mem = mem_test; + id_dense_start = id_dense_start_test; + LOG_INF("%s: set ngl_per_device[%zu].(n_layer, n_part, overflow_type)=(%" PRIu32 ", %" PRIu32 ", ATTN), id_dense_start=%zu\n", + __func__, id, ngl_per_device[id].n_layer, ngl_per_device[id].n_part, id_dense_start); + } + } + } + + const int64_t projected_margin = dmds_full[id].free - mem[id]; + LOG_INF( + "%s: - %s: %2" PRIu32 " layers (%2" PRIu32 " overflowing), %6" PRId64 " MiB used, %6" PRId64 " MiB free\n", + __func__, dev_names[id].c_str(), ngl_per_device[id].n_layer, ngl_per_device[id].n_part, mem[id]/MiB, projected_margin/MiB); + } + + // print info for devices that were not changed during the conversion from dense only to full layers: + for (size_t id = id_dense_start + 1; id < nd; id++) { + const int64_t projected_margin = dmds_full[id].free - mem[id]; + LOG_INF( + "%s: - %s: %2" PRIu32 " layers (%2" PRIu32 " overflowing), %6" PRId64 " MiB used, %6" PRId64 " MiB free\n", + __func__, dev_names[id].c_str(), ngl_per_device[id].n_layer, ngl_per_device[id].n_part, mem[id]/MiB, projected_margin/MiB); + } + + set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, *mparams); +} + +enum common_params_fit_status common_fit_params( + const char * path_model, + llama_model_params * mparams, + llama_context_params * cparams, + float * tensor_split, + llama_model_tensor_buft_override * tensor_buft_overrides, + size_t * margins, + uint32_t n_ctx_min, + ggml_log_level log_level) { + const int64_t t0_us = llama_time_us(); + common_params_fit_status status = COMMON_PARAMS_FIT_STATUS_SUCCESS; + try { + common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, log_level); + LOG_INF("%s: successfully fit params to free device memory\n", __func__); + } catch (const common_params_fit_exception & e) { + LOG_WRN("%s: failed to fit params to free device memory: %s\n", __func__, e.what()); + status = COMMON_PARAMS_FIT_STATUS_FAILURE; + } catch (const std::runtime_error & e) { + LOG_ERR("%s: encountered an error while trying to fit params to free device memory: %s\n", __func__, e.what()); + status = COMMON_PARAMS_FIT_STATUS_ERROR; + } + const int64_t t1_us = llama_time_us(); + LOG_INF("%s: fitting params to free memory took %.2f seconds\n", __func__, (t1_us - t0_us) * 1e-6); + return status; +} + +void common_memory_breakdown_print(const struct llama_context * ctx) { + //const auto & devices = ctx->get_model().devices; + const auto * model = llama_get_model(ctx); + + std::vector devices; + for (int i = 0; i < llama_model_n_devices(model); i++) { + devices.push_back(llama_model_get_device(model, i)); + } + + llama_memory_breakdown memory_breakdown = llama_get_memory_breakdown(ctx); + + std::vector> table_data; + table_data.reserve(devices.size()); + const std::string template_header = "%s: | %s | %s %s %s %s %s %s %s |\n"; + const std::string template_gpu = "%s: | %s | %s = %s + (%s = %s + %s + %s) + %s |\n"; + const std::string template_other = "%s: | %s | %s %s %s = %s + %s + %s %s |\n"; + + table_data.push_back({template_header, "memory breakdown [MiB]", "total", "free", "self", "model", "context", "compute", "unaccounted"}); + + constexpr size_t MiB = 1024 * 1024; + const std::vector desc_prefixes_strip = {"NVIDIA ", "GeForce ", "Tesla ", "AMD ", "Radeon ", "Instinct "}; + + // track seen buffer types to avoid double counting: + std::set seen_buffer_types; + + // accumulative memory breakdown for each device and for host: + std::vector mb_dev(devices.size()); + llama_memory_breakdown_data mb_host; + + for (const auto & buft_mb : memory_breakdown) { + ggml_backend_buffer_type_t buft = buft_mb.first; + const llama_memory_breakdown_data & mb = buft_mb.second; + if (ggml_backend_buft_is_host(buft)) { + mb_host.model += mb.model; + mb_host.context += mb.context; + mb_host.compute += mb.compute; + seen_buffer_types.insert(buft); + continue; + } + ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft); + if (dev) { + int i_dev = -1; + for (size_t i = 0; i < devices.size(); i++) { + if (devices[i] == dev) { + i_dev = i; + break; + } + } + if (i_dev != -1) { + mb_dev[i_dev].model += mb.model; + mb_dev[i_dev].context += mb.context; + mb_dev[i_dev].compute += mb.compute; + seen_buffer_types.insert(buft); + continue; + } + } + } + + // print memory breakdown for each device: + for (size_t i = 0; i < devices.size(); i++) { + ggml_backend_dev_t dev = devices[i]; + llama_memory_breakdown_data mb = mb_dev[i]; + + const std::string name = ggml_backend_dev_name(dev); + std::string desc = ggml_backend_dev_description(dev); + for (const std::string & prefix : desc_prefixes_strip) { + if (desc.length() >= prefix.length() && desc.substr(0, prefix.length()) == prefix) { + desc = desc.substr(prefix.length()); + } + } + + size_t free, total; + ggml_backend_dev_memory(dev, &free, &total); + + const size_t self = mb.model + mb.context + mb.compute; + const int64_t unaccounted = static_cast(total) - static_cast(free) - static_cast(self); + + table_data.push_back({ + template_gpu, + " - " + name + " (" + desc + ")", + std::to_string(total / MiB), + std::to_string(free / MiB), + std::to_string(self / MiB), + std::to_string(mb.model / MiB), + std::to_string(mb.context / MiB), + std::to_string(mb.compute / MiB), + std::to_string(unaccounted / static_cast(MiB))}); + } + + // print memory breakdown for host: + { + const size_t self = mb_host.model + mb_host.context + mb_host.compute; + table_data.push_back({ + template_other, + " - Host", + "", // total + "", // free + std::to_string(self / MiB), + std::to_string(mb_host.model / MiB), + std::to_string(mb_host.context / MiB), + std::to_string(mb_host.compute / MiB), + ""}); // unaccounted + } + + // print memory breakdown for all remaining buffer types: + for (const auto & buft_mb : memory_breakdown) { + ggml_backend_buffer_type_t buft = buft_mb.first; + const llama_memory_breakdown_data & mb = buft_mb.second; + if (seen_buffer_types.count(buft) == 1) { + continue; + } + const std::string name = ggml_backend_buft_name(buft); + const size_t self = mb.model + mb.context + mb.compute; + table_data.push_back({ + template_other, + " - " + name, + "", // total + "", // free + std::to_string(self / MiB), + std::to_string(mb.model / MiB), + std::to_string(mb.context / MiB), + std::to_string(mb.compute / MiB), + ""}); // unaccounted + seen_buffer_types.insert(buft); + } + + for (size_t j = 1; j < table_data[0].size(); j++) { + size_t max_len = 0; + for (const auto & td : table_data) { + max_len = std::max(max_len, td[j].length()); + } + for (auto & td : table_data) { + td[j].insert(j == 1 ? td[j].length() : 0, max_len - td[j].length(), ' '); + } + } + for (const auto & td : table_data) { + LOG_INF(td[0].c_str(), + __func__, td[1].c_str(), td[2].c_str(), td[3].c_str(), td[4].c_str(), td[5].c_str(), + td[6].c_str(), td[7].c_str(), td[8].c_str()); + } +} + +void common_fit_print( + const char * path_model, + llama_model_params * mparams, + llama_context_params * cparams) { + std::vector devs; + uint32_t hp_ngl = 0; // hparams.n_gpu_layers + uint32_t hp_nct = 0; // hparams.n_ctx_train + uint32_t hp_nex = 0; // hparams.n_expert + + auto dmd = common_get_device_memory_data(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, GGML_LOG_LEVEL_ERROR); + GGML_ASSERT(dmd.size() == devs.size() + 1); + + for (size_t id = 0; id < devs.size(); id++) { + printf("%s ", ggml_backend_dev_name(devs[id])); + printf("%zu ", dmd[id].mb.model/1024/1024); + printf("%zu ", dmd[id].mb.context/1024/1024); + printf("%zu ", dmd[id].mb.compute/1024/1024); + printf("\n"); + } + + printf("Host "); + printf("%zu ", dmd.back().mb.model/1024/1024); + printf("%zu ", dmd.back().mb.context/1024/1024); + printf("%zu ", dmd.back().mb.compute/1024/1024); + printf("\n"); +} + +#include "common.h.inc" + +#include "ggml.h.inc" + +#include +#include +#include +#include +#include +#include +#include + +enum common_reasoning_budget_state { + REASONING_BUDGET_IDLE, + REASONING_BUDGET_COUNTING, + REASONING_BUDGET_FORCING, + REASONING_BUDGET_WAITING_UTF8, + REASONING_BUDGET_DONE, +}; + +extern "C" { +struct llama_sampler * llama_common_reasoning_budget_init_rust( + const struct llama_vocab * vocab, + const llama_token * start_tokens, + size_t start_len, + const llama_token * end_tokens, + size_t end_len, + const llama_token * forced_tokens, + size_t forced_len, + int32_t budget, + common_reasoning_budget_state initial_state); + +common_reasoning_budget_state llama_common_reasoning_budget_get_state_rust(const struct llama_sampler * smpl); +} + +static struct llama_sampler * common_reasoning_budget_init( + const struct llama_vocab * vocab, + const std::vector & start_tokens, + const std::vector & end_tokens, + const std::vector & forced_tokens, + int32_t budget, + common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE) { + return llama_common_reasoning_budget_init_rust( + vocab, + start_tokens.data(), start_tokens.size(), + end_tokens.data(), end_tokens.size(), + forced_tokens.data(), forced_tokens.size(), + budget, + initial_state); +} + +static common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) { + return llama_common_reasoning_budget_get_state_rust(smpl); +} + +// the ring buffer works similarly to std::deque, but with a fixed capacity +// TODO: deduplicate with llama-impl.h.inc +template +struct ring_buffer { + ring_buffer(size_t cap) : capacity(cap), data(cap) {} + + T & front() { + if (sz == 0) { + throw std::runtime_error("ring buffer is empty"); + } + return data[first]; + } + + const T & front() const { + if (sz == 0) { + throw std::runtime_error("ring buffer is empty"); + } + return data[first]; + } + + T & back() { + if (sz == 0) { + throw std::runtime_error("ring buffer is empty"); + } + return data[pos]; + } + + const T & back() const { + if (sz == 0) { + throw std::runtime_error("ring buffer is empty"); + } + return data[pos]; + } + + void push_back(const T & value) { + if (sz == capacity) { + // advance the start when buffer is full + first = (first + 1) % capacity; + } else { + sz++; + } + data[pos] = value; + pos = (pos + 1) % capacity; + } + + T pop_front() { + if (sz == 0) { + throw std::runtime_error("ring buffer is empty"); + } + T value = data[first]; + first = (first + 1) % capacity; + sz--; + return value; + } + + const T & rat(size_t i) const { + if (i >= sz) { + throw std::runtime_error("ring buffer: index out of bounds"); + } + return data[(first + sz - i - 1) % capacity]; + } + + std::vector to_vector() const { + std::vector result; + result.reserve(sz); + for (size_t i = 0; i < sz; i++) { + result.push_back(data[(first + i) % capacity]); + } + return result; + } + + void clear() { + // here only reset the status of the buffer + sz = 0; + first = 0; + pos = 0; + } + + bool empty() const { + return sz == 0; + } + + size_t size() const { + return sz; + } + + size_t capacity = 0; + size_t sz = 0; + size_t first = 0; + size_t pos = 0; + std::vector data; +}; + +struct common_sampler { + common_params_sampling params; + + struct llama_sampler * grmr; + struct llama_sampler * rbudget; + struct llama_sampler * chain; + + ring_buffer prev; + + std::vector cur; + + llama_token_data_array cur_p; + + void reset() { + prev.clear(); + + llama_sampler_reset(chain); + } + + void set_logits(struct llama_context * ctx, int idx) { + const float * sampled_probs = llama_get_sampled_probs_ith (ctx, idx); + const float * sampled_logits = llama_get_sampled_logits_ith (ctx, idx); + const llama_token * sampled_ids = llama_get_sampled_candidates_ith(ctx, idx); + + const llama_model * model = llama_get_model(ctx); + const llama_vocab * vocab = llama_model_get_vocab(model); + + const int n_vocab = llama_vocab_n_tokens(vocab); + + if (sampled_probs) { + const uint32_t sampled_probs_count = llama_get_sampled_probs_count_ith(ctx, idx); + cur.resize(sampled_probs_count); + for (uint32_t i = 0; i < sampled_probs_count; ++i) { + cur[i] = llama_token_data{sampled_ids[i], sampled_logits[i], sampled_probs[i]}; + } + } else if (sampled_logits) { + const uint32_t sampled_logits_count = llama_get_sampled_logits_count_ith(ctx, idx); + cur.resize(sampled_logits_count); + for (uint32_t i = 0; i < sampled_logits_count; i++) { + cur[i] = llama_token_data{sampled_ids[i], sampled_logits[i], 0.0f}; + } + } else { + const auto * logits = llama_get_logits_ith(ctx, idx); + GGML_ASSERT(logits != nullptr); + cur.resize(n_vocab); + for (llama_token token_id = 0; token_id < n_vocab; token_id++) { + cur[token_id] = llama_token_data{token_id, logits[token_id], 0.0f}; + } + } + + cur_p = { cur.data(), cur.size(), -1, false }; + } + + common_time_meas tm() { + return common_time_meas(t_total_us, params.no_perf); + } + + mutable int64_t t_total_us = 0; +}; + +std::string common_params_sampling::print() const { + char result[1024]; + + snprintf(result, sizeof(result), + "\trepeat_last_n = %d, repeat_penalty = %.3f, frequency_penalty = %.3f, presence_penalty = %.3f\n" + "\tdry_multiplier = %.3f, dry_base = %.3f, dry_allowed_length = %d, dry_penalty_last_n = %d\n" + "\ttop_k = %d, top_p = %.3f, min_p = %.3f, xtc_probability = %.3f, xtc_threshold = %.3f, typical_p = %.3f, top_n_sigma = %.3f, temp = %.3f\n" + "\tmirostat = %d, mirostat_lr = %.3f, mirostat_ent = %.3f, adaptive_target = %.3f, adaptive_decay = %.3f", + penalty_last_n, penalty_repeat, penalty_freq, penalty_present, + dry_multiplier, dry_base, dry_allowed_length, dry_penalty_last_n, + top_k, top_p, min_p, xtc_probability, xtc_threshold, typ_p, top_n_sigma, temp, + mirostat, mirostat_eta, mirostat_tau, adaptive_target, adaptive_decay); + + return std::string(result); +} + +struct common_sampler * common_sampler_init(const struct llama_model * model, struct common_params_sampling & params) { + const llama_vocab * vocab = llama_model_get_vocab(model); + + llama_sampler_chain_params lparams = llama_sampler_chain_default_params(); + + lparams.no_perf = params.no_perf; + + llama_sampler * grmr = nullptr; + llama_sampler * rbudget = nullptr; + llama_sampler * chain = llama_sampler_chain_init(lparams); + + std::vector samplers; + + const std::string & grammar_str = common_grammar_value(params.grammar); + if (grammar_str.compare(0, 11, "%llguidance") == 0) { +#ifdef LLAMA_USE_LLGUIDANCE + grmr = llama_sampler_init_llg(vocab, "lark", grammar_str.c_str()); +#else + GGML_ABORT("llguidance (cmake -DLLAMA_LLGUIDANCE=ON) is not enabled"); +#endif // LLAMA_USE_LLGUIDANCE + } else { + std::vector trigger_patterns; + std::vector trigger_tokens; + for (const auto & trigger : params.grammar_triggers) { + switch (trigger.type) { + case COMMON_GRAMMAR_TRIGGER_TYPE_WORD: + { + const auto & word = trigger.value; + trigger_patterns.push_back(regex_escape(word)); + break; + } + case COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN: + { + trigger_patterns.push_back(trigger.value); + break; + } + case COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN_FULL: + { + const auto & pattern = trigger.value; + std::string anchored = "^$"; + if (!pattern.empty()) { + anchored = (pattern.front() != '^' ? "^" : "") + + pattern + + (pattern.back() != '$' ? "$" : ""); + } + trigger_patterns.push_back(anchored); + break; + } + case COMMON_GRAMMAR_TRIGGER_TYPE_TOKEN: + { + const auto token = trigger.token; + trigger_tokens.push_back(token); + break; + } + default: + GGML_ASSERT(false && "unknown trigger type"); + } + } + + std::vector trigger_patterns_c; + trigger_patterns_c.reserve(trigger_patterns.size()); + for (const auto & regex : trigger_patterns) { + trigger_patterns_c.push_back(regex.c_str()); + } + + if (!grammar_str.empty()) { + if (params.grammar_lazy) { + grmr = llama_sampler_init_grammar_lazy_patterns(vocab, grammar_str.c_str(), "root", + trigger_patterns_c.data(), trigger_patterns_c.size(), + trigger_tokens.data(), trigger_tokens.size()); + } else { + grmr = llama_sampler_init_grammar(vocab, grammar_str.c_str(), "root"); + } + } + } + + // Compute prefill tokens from the generation prompt + std::vector prefill_tokens; + if (!params.generation_prompt.empty()) { + GGML_ASSERT(vocab != nullptr); + auto tokens = common_tokenize(vocab, params.generation_prompt, false, true); + for (size_t i = 0; i < tokens.size(); i++) { + std::string piece = common_token_to_piece(vocab, tokens[i], true); + if (i == 0 && std::isspace(piece[0]) && !std::isspace(params.generation_prompt[0])) { + // Some tokenizers will add a space before the first special token, need to exclude + continue; + } + LOG_DBG("%s: prefill token: %d = %s\n", __func__, tokens[i], piece.c_str()); + prefill_tokens.push_back(tokens[i]); + } + } + + // Feed generation prompt tokens to the grammar sampler so it advances past + // tokens the template already placed in the prompt. + // Only applies to output-format and tool-call grammars; user-supplied grammars must not be prefilled. + if (grmr && !params.grammar_lazy && common_grammar_needs_prefill(params.grammar)) { + try { + for (const auto & token : prefill_tokens) { + llama_sampler_accept(grmr, token); + LOG_DBG("%s: grammar accepted prefill token (%d)\n", __func__, token); + } + } catch (std::exception &e) { + LOG_ERR("%s: error initializing grammar sampler for grammar:\n%s\n\nGeneration prompt:\n'%s'\n", __func__, + common_grammar_value(params.grammar).c_str(), params.generation_prompt.c_str()); + throw e; + } + } + + // reasoning budget sampler (skip when budget is unlimited unless a lazy grammar is active, which needs rbudget for thinking-block suppression) + if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && (params.grammar_lazy || params.reasoning_budget_tokens >= 0)) { + rbudget = common_reasoning_budget_init( + vocab, + params.reasoning_budget_start, + params.reasoning_budget_end, + params.reasoning_budget_forced, + params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens); + + for (const auto & token : prefill_tokens) { + llama_sampler_accept(rbudget, token); + LOG_DBG("%s: reasoning-budget accepted prefill token (%d)\n", __func__, token); + } + } + + if (params.has_logit_bias()) { + samplers.push_back(llama_sampler_init_logit_bias(llama_vocab_n_tokens(vocab), params.logit_bias.size(), params.logit_bias.data())); + } + + if (params.mirostat == 0) { + + bool use_adaptive_p = false; // see below + + for (const auto & cnstr : params.samplers) { + switch (cnstr) { + case COMMON_SAMPLER_TYPE_DRY: + { + std::vector c_breakers; + c_breakers.reserve(params.dry_sequence_breakers.size()); + for (const auto & str : params.dry_sequence_breakers) { + c_breakers.push_back(str.c_str()); + } + samplers.push_back(llama_sampler_init_dry(vocab, llama_model_n_ctx_train(model), params.dry_multiplier, params.dry_base, params.dry_allowed_length, params.dry_penalty_last_n, c_breakers.data(), c_breakers.size())); + } + break; + case COMMON_SAMPLER_TYPE_TOP_K: + samplers.push_back(llama_sampler_init_top_k(params.top_k)); + break; + case COMMON_SAMPLER_TYPE_TOP_P: + samplers.push_back(llama_sampler_init_top_p(params.top_p, params.min_keep)); + break; + case COMMON_SAMPLER_TYPE_TOP_N_SIGMA: + samplers.push_back(llama_sampler_init_top_n_sigma(params.top_n_sigma)); + break; + case COMMON_SAMPLER_TYPE_MIN_P: + samplers.push_back(llama_sampler_init_min_p(params.min_p, params.min_keep)); + break; + case COMMON_SAMPLER_TYPE_XTC: + samplers.push_back(llama_sampler_init_xtc(params.xtc_probability, params.xtc_threshold, params.min_keep, params.seed)); + break; + case COMMON_SAMPLER_TYPE_TYPICAL_P: + samplers.push_back(llama_sampler_init_typical(params.typ_p, params.min_keep)); + break; + case COMMON_SAMPLER_TYPE_TEMPERATURE: + samplers.push_back(llama_sampler_init_temp_ext(params.temp, params.dynatemp_range, params.dynatemp_exponent)); + break; + case COMMON_SAMPLER_TYPE_INFILL: + samplers.push_back(llama_sampler_init_infill(vocab)); + break; + case COMMON_SAMPLER_TYPE_PENALTIES: + samplers.push_back(llama_sampler_init_penalties(params.penalty_last_n, params.penalty_repeat, params.penalty_freq, params.penalty_present)); + break; + case COMMON_SAMPLER_TYPE_ADAPTIVE_P: + // the `adaptive-p` sampler is like `dist` and `mirostat` in that it selects + // a single token, so we will add `dist` at the end of the chain by default, + // unless the user specifically included `adaptive-p`. we set this flag here + // so we know to add the sampler at the very end. + use_adaptive_p = true; + break; + default: + GGML_ASSERT(false && "unknown sampler type"); + } + } + if (use_adaptive_p) { + // only if user explicitly included adaptive-p sampler + samplers.push_back(llama_sampler_init_adaptive_p(params.adaptive_target, params.adaptive_decay, params.seed)); + } else { + // default: sample from distribution + samplers.push_back(llama_sampler_init_dist(params.seed)); + } + } else if (params.mirostat == 1) { + samplers.push_back(llama_sampler_init_temp(params.temp)); + samplers.push_back(llama_sampler_init_mirostat(llama_vocab_n_tokens(vocab), params.seed, params.mirostat_tau, params.mirostat_eta, 100)); + } else if (params.mirostat == 2) { + samplers.push_back(llama_sampler_init_temp(params.temp)); + samplers.push_back(llama_sampler_init_mirostat_v2(params.seed, params.mirostat_tau, params.mirostat_eta)); + } else { + GGML_ASSERT(false && "unknown mirostat version"); + } + + for (auto * smpl : samplers) { + llama_sampler_chain_add(chain, smpl); + } + + if (grmr && params.backend_sampling) { + LOG_WRN("%s: backend sampling is not compatible with grammar, disabling\n", __func__); + + params.backend_sampling = false; + } + + if (rbudget && params.backend_sampling) { + LOG_WRN("%s: backend sampling is not compatible with reasoning budget, disabling\n", __func__); + + params.backend_sampling = false; + } + + auto * result = new common_sampler { + /* .params = */ params, + /* .grmr = */ grmr, + /* .rbudget = */ rbudget, + /* .chain = */ chain, + /* .prev = */ ring_buffer(std::max(32, params.n_prev)), + /* .cur = */ {}, + /* .cur_p = */ {}, + }; + + return result; +} + +void common_sampler_free(struct common_sampler * gsmpl) { + if (!gsmpl) { + return; + } + + llama_sampler_free(gsmpl->grmr); + llama_sampler_free(gsmpl->rbudget); + llama_sampler_free(gsmpl->chain); + + delete gsmpl; +} + +static bool grammar_should_apply(struct common_sampler * gsmpl) { + if (!gsmpl->grmr) { + return false; + } + if (!gsmpl->rbudget) { + return true; + } + if (gsmpl->params.grammar_lazy) { + // if grammar is lazy, only apply when reasoning budget is not active + const auto state = common_reasoning_budget_get_state(gsmpl->rbudget); + return state == REASONING_BUDGET_IDLE || state == REASONING_BUDGET_DONE; + } + return true; +} + +void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, bool is_generated) { + if (!gsmpl) { + return; + } + + const auto tm = gsmpl->tm(); + + // grammar_should_apply() checks the reasoning budget state, so calculate this before we accept + const auto accept_grammar = is_generated && grammar_should_apply(gsmpl); + + if (gsmpl->rbudget && is_generated) { + llama_sampler_accept(gsmpl->rbudget, token); + } + + if (gsmpl->grmr && accept_grammar) { + llama_sampler_accept(gsmpl->grmr, token); + } + + llama_sampler_accept(gsmpl->chain, token); + + gsmpl->prev.push_back(token); +} + +void common_sampler_reset(struct common_sampler * gsmpl) { + if (!gsmpl) { + return; + } + + gsmpl->reset(); +} + +struct common_sampler * common_sampler_clone(common_sampler * gsmpl) { + return new common_sampler { + /* .params = */ gsmpl->params, + /* .grmr = */ llama_sampler_clone(gsmpl->grmr), + /* .rbudget = */ llama_sampler_clone(gsmpl->rbudget), + /* .chain = */ llama_sampler_clone(gsmpl->chain), + /* .prev = */ gsmpl->prev, + /* .cur = */ gsmpl->cur, + /* .cur_p = */ gsmpl->cur_p, + }; +} + +void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl) { + // TODO: measure grammar performance + + const double t_sampling_ms = gsmpl ? 1e-3*gsmpl->t_total_us : 0; + + llama_perf_sampler_data data_smpl; + llama_perf_context_data data_ctx; + + memset(&data_smpl, 0, sizeof(data_smpl)); + memset(&data_ctx, 0, sizeof(data_ctx)); + + if (gsmpl) { + auto & data = data_smpl; + + data = llama_perf_sampler(gsmpl->chain); + + // note: the sampling time includes the samplers time + extra time spent in common/sampling + LOG_INF("%s: sampling time = %10.2f ms\n", __func__, t_sampling_ms); + LOG_INF("%s: samplers time = %10.2f ms / %5d tokens\n", __func__, data.t_sample_ms, data.n_sample); + } + + if (ctx) { + auto & data = data_ctx; + + data = llama_perf_context(ctx); + + const double t_end_ms = 1e-3 * ggml_time_us(); + + const double t_total_ms = t_end_ms - data.t_start_ms; + const double t_unacc_ms = t_total_ms - (t_sampling_ms + data.t_p_eval_ms + data.t_eval_ms); + const double t_unacc_pc = 100.0 * t_unacc_ms / t_total_ms; + + LOG_INF("%s: load time = %10.2f ms\n", __func__, data.t_load_ms); + LOG_INF("%s: prompt eval time = %10.2f ms / %5d tokens (%8.2f ms per token, %8.2f tokens per second)\n", + __func__, data.t_p_eval_ms, data.n_p_eval, data.t_p_eval_ms / data.n_p_eval, 1e3 / data.t_p_eval_ms * data.n_p_eval); + LOG_INF("%s: eval time = %10.2f ms / %5d runs (%8.2f ms per token, %8.2f tokens per second)\n", + __func__, data.t_eval_ms, data.n_eval, data.t_eval_ms / data.n_eval, 1e3 / data.t_eval_ms * data.n_eval); + LOG_INF("%s: total time = %10.2f ms / %5d tokens\n", __func__, (t_end_ms - data.t_start_ms), (data.n_p_eval + data.n_eval)); + LOG_INF("%s: unaccounted time = %10.2f ms / %5.1f %% (total - sampling - prompt eval - eval) / (total)\n", __func__, t_unacc_ms, t_unacc_pc); + LOG_INF("%s: graphs reused = %10d\n", __func__, data.n_reused); + + common_memory_breakdown_print(ctx); + } +} + +struct llama_sampler * common_sampler_get(const struct common_sampler * gsmpl) { + if (!gsmpl) { + return nullptr; + } + + return gsmpl->chain; +} + +llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_context * ctx, int idx, bool grammar_first) { + llama_synchronize(ctx); + + // start measuring sampling time after the llama_context synchronization in order to not measure any ongoing async operations + const auto tm = gsmpl->tm(); + + llama_token id = LLAMA_TOKEN_NULL; + + auto & grmr = gsmpl->grmr; + auto & rbudget = gsmpl->rbudget; + auto & chain = gsmpl->chain; + auto & cur_p = gsmpl->cur_p; // initialized by set_logits + + // Check if a backend sampler has already sampled a token in which case we + // return that token id directly. + { + id = llama_get_sampled_token_ith(ctx, idx); + + if (id != LLAMA_TOKEN_NULL) { + LOG_DBG("%s: Backend sampler selected token: '%d'. Will not run any CPU samplers\n", __func__, id); + + GGML_ASSERT(!gsmpl->grmr && "using grammar in combination with backend sampling is not supported"); + GGML_ASSERT(!gsmpl->rbudget && "using reasoning budget in combination with backend sampling is not supported"); + + // TODO: simplify + gsmpl->cur.resize(1); + gsmpl->cur[0] = { id, 0.0f, 1.0f }; + cur_p = { gsmpl->cur.data(), gsmpl->cur.size(), 0, true }; + + return id; + } + } + + gsmpl->set_logits(ctx, idx); + + // apply reasoning budget first + llama_sampler_apply(rbudget, &cur_p); + + if (grammar_first && grammar_should_apply(gsmpl)) { + llama_sampler_apply(grmr, &cur_p); + } + + llama_sampler_apply(chain, &cur_p); + + id = cur_p.data[cur_p.selected].id; + + if (grammar_first || !grammar_should_apply(gsmpl)) { + return id; + } + + // check if it the sampled token fits the grammar (grammar-based rejection sampling) + { + llama_token_data single_token_data = { id, 1.0f, 0.0f }; + llama_token_data_array single_token_data_array = { &single_token_data, 1, -1, false }; + + llama_sampler_apply(grmr, &single_token_data_array); + + const bool is_valid = single_token_data_array.data[0].logit != -INFINITY; + if (is_valid) { + return id; + } + } + + // resampling: + // if the token is not valid, sample again, but first apply the grammar sampler and then the sampling chain + gsmpl->set_logits(ctx, idx); + + llama_sampler_apply(rbudget, &cur_p); + + if (grammar_should_apply(gsmpl)) { + llama_sampler_apply(grmr, &cur_p); + } + + llama_sampler_apply(chain, &cur_p); + + GGML_ASSERT(cur_p.selected != -1 && "no selected token during sampling - check your sampling configuration"); + + id = cur_p.data[cur_p.selected].id; + + return id; +} + +std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector & idxs, const llama_tokens & draft, bool grammar_first) { + GGML_ASSERT(idxs.size() == draft.size() + 1 && "idxs.size() must be draft.size() + 1"); + + std::vector result; + result.reserve(idxs.size()); + + size_t i = 0; + for (; i < draft.size(); i++) { + const llama_token id = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first); + + common_sampler_accept(gsmpl, id, true); + + result.push_back(id); + + if (draft[i] != id) { + break; + } + } + + if (i == draft.size()) { + const llama_token id = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first); + + common_sampler_accept(gsmpl, id, true); + + result.push_back(id); + } + + return result; +} + +std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first) { + std::vector idxs(draft.size() + 1); + for (size_t i = 0; i < idxs.size(); ++i) { + idxs[i] = i; + } + + return common_sampler_sample_and_accept_n(gsmpl, ctx, idxs, draft, grammar_first); +} + +uint32_t common_sampler_get_seed(const struct common_sampler * gsmpl) { + return llama_sampler_get_seed(gsmpl->chain); +} + +// helpers + +llama_token_data_array * common_sampler_get_candidates(struct common_sampler * gsmpl, bool do_sort) { + const auto tm = gsmpl->tm(); + + auto * res = &gsmpl->cur_p; + + if (do_sort && !res->sorted) { + // remember the selected token before sorting + const llama_token id = res->data[res->selected].id; + + std::sort(res->data, res->data + res->size, [](const llama_token_data & a, const llama_token_data & b) { + return a.p > b.p; + }); + + // restore the selected token after sorting + for (size_t i = 0; i < res->size; ++i) { + if (res->data[i].id == id) { + res->selected = i; + break; + } + } + + res->sorted = true; + } + + return res; +} + +llama_token common_sampler_last(const struct common_sampler * gsmpl) { + return gsmpl->prev.rat(0); +} + +std::string common_sampler_print(const struct common_sampler * gsmpl) { + std::string result = "logits "; + + for (int i = 0; i < llama_sampler_chain_n(gsmpl->chain); i++) { + const auto * smpl = llama_sampler_chain_get(gsmpl->chain, i); + result += std::string("-> "); + result += std::string(llama_sampler_name(smpl)) + " "; + } + + return result; +} + +std::string common_sampler_prev_str(common_sampler * gsmpl, llama_context * ctx_main, int n) { + n = std::min(n, (int) gsmpl->prev.size()); + + if (n <= 0) { + return ""; + } + + std::string result; + result.reserve(8*n); // 8 is the average length of a token [citation needed], TODO: compute this from the vocab + + for (int i = n - 1; i >= 0; i--) { + const llama_token id = gsmpl->prev.rat(i); + + GGML_ASSERT(id != LLAMA_TOKEN_NULL && "null token in the sampling history - should not happen"); + + result += common_token_to_piece(ctx_main, id); + } + + return result; +} + +char common_sampler_type_to_chr(enum common_sampler_type cnstr) { + return llama_common_sampler_type_to_chr_rust((int) cnstr); +} + +std::string common_sampler_type_to_str(enum common_sampler_type cnstr) { + const char * name = llama_common_sampler_type_to_str_rust((int) cnstr); + return name ? name : ""; +} + +std::vector common_sampler_types_from_names(const std::vector & names, bool allow_alt_names) { + std::vector samplers; + samplers.reserve(names.size()); + + for (const auto & name : names) { + const int sampler = llama_common_sampler_type_from_name_rust(name.c_str(), allow_alt_names); + if (sampler >= 0) { + samplers.push_back((common_sampler_type) sampler); + } else { + LOG_WRN("%s: unable to match sampler by name '%s'\n", __func__, name.c_str()); + } + } + + return samplers; +} + +std::vector common_sampler_types_from_chars(const std::string & chars) { + std::vector samplers; + samplers.reserve(chars.size()); + + for (const auto & c : chars) { + const int sampler = llama_common_sampler_type_from_chr_rust(c); + if (sampler >= 0) { + samplers.push_back((common_sampler_type) sampler); + } else { + LOG_WRN("%s: unable to match sampler by char '%c'\n", __func__, c); + } + } + + return samplers; +} + +#include "common.h.inc" +#include "ggml.h.inc" +#include "llama.h.inc" +#include +#include +#include +#include +#include +#include + + +#define CHATML_TEMPLATE_SRC \ + "{%- for message in messages -%}\n" \ + " {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' -}}\n" \ + "{%- endfor -%}\n" \ + "{%- if add_generation_prompt -%}\n" \ + " {{- '<|im_start|>assistant\n' -}}\n" \ + "{%- endif -%}" + +void common_chat_templates_free(struct common_chat_templates * tmpls) { + delete tmpls; +} + +bool common_chat_templates_was_explicit(const struct common_chat_templates * tmpls) { + return tmpls->has_explicit_template; +} + +// LFM2 format detection: template uses <|tool_list_start|>[...]<|tool_list_end|> around the tool list +// and <|tool_call_start|>[...]<|tool_call_end|> around each tool call +static bool is_lfm2_template(const std::string & src) { + return src.find("<|tool_list_start|>") != std::string::npos && + src.find("<|tool_list_end|>") != std::string::npos; +} + +common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates) { + common_chat_prompt_preset asr_preset; + asr_preset.system = ""; + asr_preset.user = "Transcribe audio to text"; + + if (chat_templates && chat_templates->template_default && is_lfm2_template(chat_templates->template_default->source())) { + asr_preset.system = "Perform ASR."; + asr_preset.user = ""; + } + + return asr_preset; +} + +std::string common_chat_templates_source(const struct common_chat_templates * tmpls, const std::string & variant) { + if (!variant.empty()) { + if (variant == "tool_use") { + if (tmpls->template_tool_use) { + return tmpls->template_tool_use->source(); + } + return ""; + } + LOG_DBG("%s: unknown template variant: %s\n", __func__, variant.c_str()); + } + return tmpls->template_default->source(); +} + +common_chat_templates_ptr common_chat_templates_init(const struct llama_model * model, + const std::string & chat_template_override, + const std::string & bos_token_override, + const std::string & eos_token_override) { + std::string default_template_src; + std::string template_tool_use_src; + + bool has_explicit_template = !chat_template_override.empty(); + if (chat_template_override.empty()) { + GGML_ASSERT(model != nullptr); + const auto * str = llama_model_chat_template(model, /* name */ nullptr); + if (str) { + default_template_src = str; + has_explicit_template = true; + } + str = llama_model_chat_template(model, /* name */ "tool_use"); + if (str) { + template_tool_use_src = str; + has_explicit_template = true; + } + } else { + default_template_src = chat_template_override; + } + if (default_template_src.empty() || default_template_src == "chatml") { + if (!template_tool_use_src.empty()) { + default_template_src = template_tool_use_src; + } else { + default_template_src = CHATML_TEMPLATE_SRC; + } + } + + // TODO @ngxson : this is a temporary hack to prevent chat template from throwing an error + // Ref: https://github.com/ggml-org/llama.cpp/pull/15230#issuecomment-3173959633 + if (default_template_src.find("<|channel|>") != std::string::npos + // search for the error message and patch it + && default_template_src.find("in message.content or") != std::string::npos) { + string_replace_all(default_template_src, + "{%- if \"<|channel|>analysis<|message|>\" in message.content or " + "\"<|channel|>final<|message|>\" in message.content %}", + "{%- if false %}"); + } + + // TODO @aldehir : this is a temporary fix, pending Minja changes + // Ref: https://github.com/ggml-org/llama.cpp/pull/17713#issuecomment-3631342664 + if (default_template_src.find("[TOOL_CALLS]") != std::string::npos + // search for the error message and patch it + && default_template_src.find("if (message['content'] is none or") != std::string::npos) { + string_replace_all(default_template_src, + "{%- if (message['content'] is none or message['content'] == '' or " + "message['content']|length == 0) and (message['tool_calls'] is not defined or " + "message['tool_calls'] is none or message['tool_calls']|length == 0) %}", + "{%- if false %}"); + } + + std::string token_bos = bos_token_override; + std::string token_eos = eos_token_override; + bool add_bos = false; + bool add_eos = false; + if (model) { + const auto * vocab = llama_model_get_vocab(model); + const auto get_token = [&](llama_token token, const char * name, const char * jinja_variable_name) { + if (token == LLAMA_TOKEN_NULL) { + if (default_template_src.find(jinja_variable_name) != std::string::npos || + template_tool_use_src.find(jinja_variable_name) != std::string::npos) { + LOG_WRN( + "common_chat_templates_init: warning: vocab does not have a %s token, jinja template won't " + "work as intended.\n", + name); + } + return std::string(); + } + return common_token_to_piece(vocab, token, true); + }; + token_bos = get_token(llama_vocab_bos(vocab), "BOS", "bos_token"); + token_eos = get_token(llama_vocab_eos(vocab), "EOS", "eos_token"); + add_bos = llama_vocab_get_add_bos(vocab); + add_eos = llama_vocab_get_add_eos(vocab); + } + common_chat_templates_ptr tmpls(new common_chat_templates()); + tmpls->has_explicit_template = has_explicit_template; + tmpls->add_bos = add_bos; + tmpls->add_eos = add_eos; + try { + tmpls->template_default = std::make_unique(default_template_src, token_bos, token_eos); + } catch (const std::exception & e) { + LOG_ERR("%s: error: %s\n", __func__, e.what()); + LOG_ERR("%s: failed to initialize chat template\n", __func__); + LOG_ERR("%s: please consider disabling jinja via --no-jinja, or using another chat template\n", __func__); + throw e; + } + if (!template_tool_use_src.empty()) { + try { + tmpls->template_tool_use = std::make_unique(template_tool_use_src, token_bos, token_eos); + } catch (const std::exception & e) { + LOG_ERR("%s: failed to parse tool use chat template (ignoring it): %s\n", __func__, e.what()); + } + } + return tmpls; +} + +const char * common_chat_format_name(common_chat_format format) { + const char * name = llama_common_chat_format_name_rust((int) format); + if (name) { + return name; + } + throw std::runtime_error("Unknown chat format"); +} + +const char * common_reasoning_format_name(common_reasoning_format format) { + const char * name = llama_common_reasoning_format_name_rust((int) format); + if (name) { + return name; + } + throw std::runtime_error("Unknown reasoning format"); +} + +common_reasoning_format common_reasoning_format_from_name(const std::string & format) { + const int parsed = llama_common_reasoning_format_from_name_rust(format.c_str()); + if (parsed >= 0) { + return (common_reasoning_format) parsed; + } + throw std::runtime_error("Unknown reasoning format: " + format); +} + +static void foreach_function(const json & tools, const std::function & fn) { + for (const auto & tool : tools) { + if (!tool.contains("type") || tool.at("type") != "function" || !tool.contains("function")) { + LOG_INF("Skipping tool without function: %s", tool.dump(2).c_str()); + continue; + } + fn(tool); + } +} + +static void foreach_parameter(const json & function, + const std::function & fn) { + if (!function.contains("parameters") || !function.at("parameters").is_object()) { + return; + } + const auto & params = function.at("parameters"); + if (!params.contains("properties") || !params.at("properties").is_object()) { + return; + } + const auto & props = params.at("properties"); + std::set required; + if (params.contains("required") && params.at("required").is_array()) { + params.at("required").get_to(required); + } + for (const auto & [name, prop] : props.items()) { + bool is_required = (required.find(name) != required.end()); + fn(name, prop, is_required); + } +} + +static std::string common_chat_template_direct_apply_impl( + const common_chat_template & tmpl, + const autoparser::generation_params & inputs, + const std::optional & messages_override = std::nullopt, + const std::optional & tools_override = std::nullopt, + const std::optional & additional_context = std::nullopt) { + jinja::context ctx(tmpl.source()); + + nlohmann::ordered_json inp = nlohmann::ordered_json{ + {"messages", messages_override.has_value() ? *messages_override : inputs.messages}, + {"bos_token", tmpl.bos_token()}, + {"eos_token", tmpl.eos_token()}, + {"enable_thinking", inputs.enable_thinking}, + }; + if (tools_override.has_value() || !inputs.tools.empty()) { + inp["tools"] = tools_override.has_value() ? *tools_override : inputs.tools; + } + if (inputs.extra_context.is_object()) { + // TODO: do we need to merge, or replacing is fine? + for (const auto & [k, v] : inputs.extra_context.items()) { + inp[k] = v; + } + } + if (additional_context.has_value()) { + // TODO: merge properly instead of overwriting (matching old behavior) + for (const auto & [k, v] : additional_context->items()) { + inp[k] = v; + } + } + if (inputs.add_generation_prompt) { + inp["add_generation_prompt"] = true; + } + + jinja::global_from_json(ctx, inp, inputs.mark_input); + + // render + jinja::runtime runtime(ctx); + const jinja::value results = runtime.execute(tmpl.prog); + auto parts = jinja::runtime::gather_string_parts(results); + + std::string result = parts->as_string().str(); + + // TODO: improve this later + if (inputs.add_bos && string_starts_with(result, tmpl.bos_token())) { + result = result.substr(tmpl.bos_token().size()); + } + if (inputs.add_eos && string_ends_with(result, tmpl.eos_token())) { + result = result.substr(0, result.size() - tmpl.eos_token().size()); + } + return result; +} + +std::string common_chat_template_direct_apply( + const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + return common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt); +} + +static common_chat_params common_chat_params_init_ministral_3(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + // Build up messages to follow the format: https://huggingface.co/mistralai/Ministral-3-14B-Reasoning-2512/blob/main/chat_template.jinja + auto adjusted_messages = json::array(); + for (const auto & msg : inputs.messages) { + auto role = msg.value("role", ""); + if (role != "system" && role != "assistant") { + // Only adjust system and assistant messages. Interestingly, the system message may contain thinking. + adjusted_messages.push_back(msg); + continue; + } + + auto content = json::array(); + + // If message contains `reasoning_content`, add it as a block of type `thinking` + if (msg.contains("reasoning_content") && msg.at("reasoning_content").is_string()) { + content.push_back({ + { "type", "thinking" }, + { "thinking", msg.at("reasoning_content").get() }, + }); + } + + // If message contains `content`, add it as a block of type `text` + if (msg.contains("content")) { + if (msg.at("content").is_string()) { + content.push_back({ + { "type", "text" }, + { "text", msg.at("content").get() }, + }); + } else if (msg.at("content").is_array()) { + auto blocks = msg.at("content"); + content.insert(content.end(), blocks.begin(), blocks.end()); + } + } + + auto adjusted = msg; + adjusted["content"] = content; + adjusted.erase("reasoning_content"); + adjusted_messages.push_back(adjusted); + } + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto has_response_format = inputs.json_schema.is_object() && !inputs.json_schema.empty(); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + auto include_grammar = true; + + data.supports_thinking = true; + data.thinking_start_tag = "[THINK]"; + data.thinking_end_tag = "[/THINK]"; + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.preserved_tokens = { + "[THINK]", + "[/THINK]", + "[TOOL_CALLS]", + "[ARGS]", + }; + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto generation_prompt = p.prefix(inputs.generation_prompt, "[THINK]"); + auto reasoning = + extract_reasoning ? p.optional("[THINK]" + p.reasoning(p.until("[/THINK]")) + "[/THINK]") : p.eps(); + + // Response format parser + if (has_response_format) { + // Ministral wants to emit json surrounded by code fences + return generation_prompt + (reasoning << "```json" << p.content(p.schema(p.json(), "response-format", inputs.json_schema)) << "```"); + } + + // Tool call parser + if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + const auto & schema = function.at("parameters"); + + tool_choice |= + p.rule("tool-" + name, p.tool_open(p.tool_name(p.literal(name)) + "[ARGS]") + + p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema))); + }); + + auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0; + auto max_calls = inputs.parallel_tool_calls ? -1 : 1; + auto tool_calls = p.trigger_rule("tool-call", p.repeat("[TOOL_CALLS]" + tool_choice, min_calls, max_calls)); + + return generation_prompt + (reasoning << p.content(p.until("[TOOL_CALLS]")) << tool_calls); + } + + // Content only parser + include_grammar = false; + return generation_prompt + (reasoning << p.content(p.rest())); + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; + + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.at("parameters"); + builder.resolve_refs(schema); + }); + if (has_response_format) { + auto schema = inputs.json_schema; + builder.resolve_refs(schema); + } + parser.build_grammar(builder, data.grammar_lazy); + }); + + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "[TOOL_CALLS]" } + }; + } + + return data; +} + +static common_chat_params common_chat_params_init_gpt_oss(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + // Copy reasoning to the "thinking" field as expected by the gpt-oss template + auto adjusted_messages = json::array(); + for (auto msg : inputs.messages) { + if (msg.contains("reasoning_content") && msg.at("reasoning_content").is_string()) { + msg["thinking"] = msg.at("reasoning_content"); + if (msg.contains("tool_calls") && msg.at("tool_calls").is_array() && !msg.at("tool_calls").empty()) { + msg.erase("content"); + } + } + adjusted_messages.push_back(msg); + } + + auto prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override= */ adjusted_messages); + + // Check if we need to replace the return token with end token during + // inference and without generation prompt. For more details see: + // https://github.com/ggml-org/llama.cpp/issues/15417 + if (inputs.is_inference && !inputs.add_generation_prompt) { + static constexpr std::string_view return_token = "<|return|>"; + static constexpr std::string_view end_token = "<|end|>"; + if (size_t pos = prompt.rfind(return_token); pos != std::string::npos) { + prompt.replace(pos, return_token.length(), end_token); + } + } + + data.prompt = prompt; + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + + // These special tokens are required to parse properly, so we include them + // even if parse_tool_calls is false. + data.preserved_tokens = { + "<|channel|>", "<|constrain|>", "<|message|>", "<|start|>", "<|end|>", + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); + auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto start = p.rule("start", p.literal("<|start|>assistant")); + auto end = p.rule("end", p.literal("<|end|>")); + auto content = p.rule("message-content", p.until("<|end|>")); + auto channel = p.literal("<|channel|>") + (p.literal("commentary") | p.literal("analysis")); + auto constrain_type = p.chars("[A-Za-z0-9_-]", 1, -1); + + // Occasionally, gpt-oss-20b will prefix channels with this commentary + auto stray_commentary = p.optional(p.literal("<|channel|>commentary") + p.optional(p.literal(" to=assistant"))); + auto start_analysis = stray_commentary + p.literal("<|channel|>analysis<|message|>"); + + if (extract_reasoning) { + p.rule("analysis", start_analysis + p.reasoning(content) + end); + } else { + p.rule("analysis", p.content(start_analysis + content + end)); + } + + auto analysis = p.ref("analysis"); + auto preamble = p.rule("preamble", p.literal("<|channel|>commentary<|message|>") + p.content(content) + end); + auto final_msg = p.rule("final", stray_commentary + p.literal("<|channel|>final<|message|>") + p.content(content)); + + // Consume any unsolicited tool calls, e.g. builtin functions + auto unsolicited = p.rule("unsolicited", p.atomic(p.optional(channel) + p.literal(" to=") + content + end)); + + auto any = p.rule("any", preamble | analysis); + + if (has_response_format) { + auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type); + auto response_format = p.rule("response-format", + p.literal("<|channel|>final") + constraint + p.literal("<|message|>") + + p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema))); + + return p.zero_or_more(start + analysis) + start + response_format; + } + + if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { + auto tool_choice = p.choice(); + + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + const auto & params = function.at("parameters"); + + auto func_name = p.literal(" to=functions.") + p.tool_name(p.literal(name)); + auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type); + auto args = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", params)); + + // recipient in role header + // <|start|>assistant to=functions.NAME<|channel|>(commentary|analysis)[constraint]<|message|>ARGS + auto tool_in_role = p.tool(p.tool_open(func_name + channel + constraint + p.literal("<|message|>")) + args); + + // recipient in channel header + // <|channel|>(commentary|analysis) to=functions.NAME[constraint]<|message|>ARGS + auto tool_in_channel = p.tool(p.tool_open(channel + func_name + constraint + p.literal("<|message|>")) + args); + + tool_choice |= p.rule("tool-" + name, tool_in_role | tool_in_channel); + }); + + auto tool_call = p.trigger_rule("tool-call", tool_choice); + + if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) { + return p.zero_or_more(start + any) + start + tool_call; + } + + return p.zero_or_more(start + any) + start + (tool_call | final_msg); + } + + return p.zero_or_more(start + any) + start + (final_msg | unsolicited); + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.at("parameters"); + builder.resolve_refs(schema); + }); + if (has_response_format) { + auto schema = inputs.json_schema; + builder.resolve_refs(schema); + } + parser.build_grammar(builder, data.grammar_lazy); + }); + + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "^\\s+to$" }, + { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "^<\\|channel\\|>(?:commentary|analysis)\\s+to=functions$" }, + { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "<\\|start\\|>assistant(\\s+to)" }, + { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "<\\|start\\|>assistant(<\\|channel\\|>(?:commentary|analysis)\\s+to)" } + }; + } + + return data; +} + +static common_chat_params common_chat_params_init_gemma4(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + + if (inputs.add_generation_prompt && string_ends_with(data.prompt, "\n")) { + // This may happen if the model generates content + tool_call, the + // template does not add the model's next turn and confuses the model + // from emitting its proper reasoning token sequence. + data.prompt += "<|turn>model\n"; + } + + data.format = COMMON_CHAT_FORMAT_PEG_GEMMA4; + data.supports_thinking = true; + data.thinking_start_tag = "<|channel>thought"; + data.thinking_end_tag = ""; + + data.preserved_tokens = { + "<|channel>", + "", + "<|tool_call>", + "", + "<|turn>", + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); + auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto start = p.rule("start", p.prefix(inputs.generation_prompt, "<|channel>")); + + if (extract_reasoning) { + p.rule("thought", p.literal("<|channel>thought") + p.space() + p.reasoning(p.until("")) + p.literal("")); + } else { + p.rule("thought", p.content(p.literal("<|channel>thought") + p.space() + p.until("") + p.literal(""))); + } + + auto consume_empty_channels = p.gbnf(p.zero_or_more(p.literal("<|channel>") + p.negate(p.literal("thought"))), ""); + auto thought = (p.peek(p.literal("<|channel>")) + consume_empty_channels + p.ref("thought")) | p.negate(p.literal("<|channel>")); + + if (has_response_format) { + auto response_format = p.literal("```json") << + p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) << + p.literal("```"); + return start + p.optional(thought) + response_format; + } + + if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { + // Gemma4 tool calling syntax + // Rules should match traversal logic in gemma4_to_json() + p.rule("gemma4-string-content", p.until("<|\"|>")); + p.rule("gemma4-string", p.literal("<|\"|>") + p.ref("gemma4-string-content") + p.literal("<|\"|>")); + p.rule("gemma4-bool", p.json_bool()); + p.rule("gemma4-null", p.json_null()); + p.rule("gemma4-number", p.json_number()); + p.rule("gemma4-dict-key", p.rule("gemma4-dict-key-name", p.chars("[^:}]", 1, -1)) + p.literal(":")); + p.rule("gemma4-dict-kv", p.ref("gemma4-dict-key") + p.space() + p.ref("gemma4-value")); + p.rule("gemma4-dict", [&]() { + auto ws = p.space(); + auto member = p.ref("gemma4-dict-kv"); + auto members = p.sequence({member, p.zero_or_more(p.sequence({p.literal(","), ws, member}))}); + return p.sequence({ + p.literal("{"), ws, + p.choice({p.literal("}"), p.sequence({members, ws, p.literal("}")})}) + }); + }); + p.rule("gemma4-array", [&]() { + auto ws = p.space(); + auto value = p.ref("gemma4-value"); + auto elements = p.sequence({value, p.zero_or_more(p.sequence({p.literal(","), ws, value}))}); + return p.sequence({ + p.literal("["), ws, + p.choice({p.literal("]"), p.sequence({elements, ws, p.literal("]")})}) + }); + }); + p.rule("gemma4-value", [&]() { + return p.choice({ + p.ref("gemma4-string"), p.ref("gemma4-dict"), p.ref("gemma4-array"), + p.ref("gemma4-number"), p.ref("gemma4-bool"), p.ref("gemma4-null") + }); + }); + + auto tool_choice = p.choice(); + + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + // TODO @aldehir : need to extend json-schema-to-grammar to produce more than JSON rules + // const auto & params = function.at("parameters"); + + tool_choice |= p.rule("tool-" + name, p.tool(p.sequence({ + p.tool_open(p.tool_name(p.literal(name)) + p.peek(p.literal("{"))), + p.tool_args(p.ref("gemma4-dict")), + }))); + }); + + auto tool_call = p.trigger_rule("tool-call", p.repeat( + "<|tool_call>call:" + tool_choice + "", + /* min = */ inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0, + /* max = */ inputs.parallel_tool_calls ? -1 : 1 + )); + + auto scan_to_toolcall = p.rule("scan-to-toolcall", p.until("<|tool_call>")); + auto content = p.rule("content", p.content(p.until_one_of({"<|channel>", "", "<|tool_call>"}))); + auto message = p.rule("message", thought + content); + return start + p.zero_or_more(message) + scan_to_toolcall + tool_call; + } + + // Gemma 4 may emit an extra <|channel>thought\n at the end of the content. It may + // also emit a single trailing token. Consume all complete reasoning blocks and + // then stop at the first unmatched token. + auto content = p.rule("content", p.content(p.until_one_of({"<|channel>", ""}))); + auto message = p.rule("message", thought + content); + return start + p.one_or_more(message); + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.at("parameters"); + builder.resolve_refs(schema); + }); + if (has_response_format) { + auto schema = inputs.json_schema; + builder.resolve_refs(schema); + } + parser.build_grammar(builder, data.grammar_lazy); + }); + + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "<|tool_call>" }, + }; + } + + return data; +} + +// Functionary v3.2 - uses recipient-based format: >>>recipient\n{content} +static common_chat_params common_chat_params_init_functionary_v3_2(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.preserved_tokens = { + ">>>all", + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + // Functionary v3.2 format: + // - Normal content: >>>all\n{content} + // - Tool calls: >>>function_name\n{json_args} + // Generation prompt ends with ">>>" so model outputs recipient immediately + + // Build content parser for >>>all\n{content} + // When tools are present, content stops before the next ">>>" (tool call) + // When no tools, content goes until end + auto content_until_tool = p.literal("all\n") + p.content(p.until(">>>")); + auto content_until_end = p.literal("all\n") + p.content(p.rest()); + auto generation_prompt = p.literal(inputs.generation_prompt); + + // If no tools or tool_choice is NONE, just parse content + if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { + // When no tools, just match the prefix and capture everything after + return generation_prompt + content_until_end + p.end(); + } + + // Build tool call parsers for each available function + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + const auto & schema = function.at("parameters"); + + // Tool format: >>>function_name\n{json_args} + auto tool_parser = p.tool( + p.tool_open(p.tool_name(p.literal(name)) + p.literal("\n")) + + p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)) + ); + + tool_choice |= p.rule("tool-" + name, tool_parser); + }); + + auto content_only = content_until_end; + auto tools_only = p.trigger_rule("tools", p.one_or_more(tool_choice)); + auto content_and_tools = content_until_tool + tools_only; + + auto ret = p.eps(); + if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) { + if (inputs.parallel_tool_calls) { + ret = p.choice({ content_and_tools, tools_only }) + p.end(); + } else { + ret = p.choice({ content_until_tool + tool_choice, tools_only }) + p.end(); + } + } else if (inputs.parallel_tool_calls) { + ret = p.choice({ content_and_tools, content_only, tools_only }) + p.end(); + } else { + auto content_and_tool = content_until_tool + tool_choice; + ret = p.choice({ content_and_tool, content_only, tool_choice }) + p.end(); + } + return generation_prompt + ret; + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; + + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.at("parameters"); + builder.resolve_refs(schema); + }); + parser.build_grammar(builder, data.grammar_lazy); + }); + + // Grammar trigger for when the model starts outputting a tool call + // (after the initial ">>>" in the generation prompt but recipient other than "all") + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, ">>>(?!all)" } + }; + } + + return data; +} + +// Kimi K2 Thinking - uses unique tool call ID format: functions.: +// The ID contains both the function name and an incrementing counter +static common_chat_params common_chat_params_init_kimi_k2(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + data.preserved_tokens = { + "<|tool_calls_section_begin|>", + "<|tool_calls_section_end|>", + "<|tool_call_begin|>", + "<|tool_call_argument_begin|>", + "<|tool_call_end|>", + "", + "", + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; + + const std::string SECTION_BEGIN = "<|tool_calls_section_begin|>"; + const std::string SECTION_END = "<|tool_calls_section_end|>"; + const std::string CALL_BEGIN = "<|tool_call_begin|>"; + const std::string ARGS_BEGIN = "<|tool_call_argument_begin|>"; + const std::string CALL_END = "<|tool_call_end|>"; + + const std::string THINK_START = ""; + const std::string THINK_END = ""; + + data.thinking_start_tag = THINK_START; + data.thinking_end_tag = THINK_END; + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + // Kimi K2 Thinking format: + // - Reasoning: {reasoning} + // - Content: text after reasoning + // - Tool calls section: + // <|tool_calls_section_begin|> + // <|tool_call_begin|>functions.:<|tool_call_argument_begin|>{json_args}<|tool_call_end|> + // ... + // <|tool_calls_section_end|> + // The ID format is: functions.: where counter is 0, 1, 2, ... + + // Tool call markers + auto end = p.end(); + + // Note: this model is CRAZY. It can diverge from its supposed tool calling pattern in so many ways it's not funny. + // For example, it can call tools at the end of reasoning without closing reasoning... + auto reasoning = extract_reasoning ? p.optional(THINK_START + p.reasoning( + p.until_one_of({ THINK_END, "<|tool_calls_section_begin|>", "<|tool_call_begin|>" })) + + p.optional(p.literal(THINK_END))) : p.eps(); + auto generation_prompt = p.prefix(inputs.generation_prompt, THINK_START); + + + // Content only parser (no tools) + if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { + return generation_prompt + reasoning + p.content(p.rest()) + end; + } + + // Build tool call parsers for each available function + // The ID format is: functions.: + // We need to match: functions.: + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + const auto & schema = function.at("parameters"); + + // Match: functions.: + // Capture the full call id (functions.:) using tool_id tag + auto tool_id = p.tool_id(p.literal("functions.") + p.tool_name(p.literal(name)) + p.literal(":") + p.chars("[0-9]", 1, -1)); + auto tool_parser = p.tool( + p.tool_open(tool_id + p.literal(ARGS_BEGIN)) + + p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)) + + p.tool_close(p.optional((p.literal(CALL_END)))) + ); + + tool_choice |= p.rule("tool-" + name, tool_parser); + }); + + // Tool calls section: <|tool_calls_section_begin|> tool_calls <|tool_calls_section_end|> + auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0; + auto max_calls = inputs.parallel_tool_calls ? -1 : 1; + // Use trigger_rule so grammar generator knows where to start generating rules + auto tool_calls = p.rule("tool-calls", + p.optional(p.literal(SECTION_BEGIN)) + + p.trigger_rule("tool-call", p.repeat(CALL_BEGIN + tool_choice, min_calls, max_calls) + + p.optional(p.literal(SECTION_END))) + ); + + auto content_before_tools = p.content(p.until_one_of({ SECTION_BEGIN, CALL_BEGIN })); + + return generation_prompt + reasoning + content_before_tools + tool_calls + end; + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.at("parameters"); + builder.resolve_refs(schema); + }); + parser.build_grammar(builder, data.grammar_lazy); + }); + + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "<|tool_call_begin|>" } + }; + } + + return data; +} + +// LFM2 format: uses <|tool_list_start|>[...]<|tool_list_end|> in system prompt +// and <|tool_call_start|>[name(arg="val")]<|tool_call_end|> for tool calls. +// - Reasoning: {reasoning} (optional) +// - Content: text before a tool call (optional) +// - Tool calls: Python-style, e.g. [function_name(arg1="value1", arg2="value2")] +// Tool calls can appear multiple times (parallel tool calls supported) +static common_chat_params common_chat_params_init_lfm2(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + data.preserved_tokens = { + "<|tool_list_start|>", + "<|tool_list_end|>", + "<|tool_call_start|>", + "<|tool_call_end|>", + "", + "", + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; + + const std::string TOOL_CALL_START = "<|tool_call_start|>"; + const std::string TOOL_CALL_END = "<|tool_call_end|>"; + const std::string THINK_START = ""; + const std::string THINK_END = ""; + + data.thinking_start_tag = THINK_START; + data.thinking_end_tag = THINK_END; + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto generation_prompt = p.prefix(inputs.generation_prompt, THINK_START); + auto end = p.end(); + + auto reasoning = p.eps(); + if (extract_reasoning && inputs.enable_thinking) { + reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END); + } + + if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { + return generation_prompt + reasoning + p.content(p.rest()) + end; + } + auto tool_calls = p.rule("tool-calls", + p.trigger_rule("tool-call", + p.literal(TOOL_CALL_START) + + p.python_style_tool_calls(inputs.tools, inputs.parallel_tool_calls) + + p.literal(TOOL_CALL_END) + ) + ); + + auto content = p.content(p.until(TOOL_CALL_START)); + + return generation_prompt + reasoning + content + tool_calls + end; + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.at("parameters"); + builder.resolve_refs(schema); + }); + parser.build_grammar(builder, data.grammar_lazy); + }); + + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, TOOL_CALL_START } + }; + } + return data; +} + +// LFM2.5 format: uses plain "List of tools: [...]" in system prompt, no wrapper tokens. +// Tool calls are bare [name(arg="val")], though model may optionally emit <|tool_call_start|>. +// - Reasoning: {reasoning} (optional) +// - Content: text before a tool call (optional) +// - Tool calls: Python-style, e.g. [function_name(arg1="value1", arg2="value2")] +// Tool calls can appear multiple times (parallel tool calls supported) +static common_chat_params common_chat_params_init_lfm2_5(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + data.preserved_tokens = { + "<|tool_call_start|>", + "<|tool_call_end|>", + "", + "", + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; + + const std::string THINK_START = ""; + const std::string THINK_END = ""; + + data.thinking_start_tag = THINK_START; + data.thinking_end_tag = THINK_END; + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto generation_prompt = p.prefix(inputs.generation_prompt, THINK_START); + auto end = p.end(); + + auto reasoning = p.eps(); + if (extract_reasoning && inputs.enable_thinking) { + reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END); + } + + if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { + return generation_prompt + reasoning + p.content(p.rest()) + end; + } + + auto tool_calls = p.rule("tool-calls", + p.trigger_rule("tool-call", + p.python_style_tool_calls(inputs.tools, inputs.parallel_tool_calls) + ) + ); + + auto content = p.content(p.until_one_of({"<|tool_call_start|>", "["})); + auto maybe_start = p.optional(p.literal("<|tool_call_start|>")); + return generation_prompt + reasoning + content + maybe_start + tool_calls + end; + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.at("parameters"); + builder.resolve_refs(schema); + }); + parser.build_grammar(builder, data.grammar_lazy); + }); + foreach_function(inputs.tools, [&](const json & tool) { + const std::string name = tool.at("function").at("name"); + data.grammar_triggers.push_back({ COMMON_GRAMMAR_TRIGGER_TYPE_WORD, "[" + name + "(" }); + }); + } + + return data; +} + +static common_chat_params common_chat_params_init_gigachat_v3( + const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = false; + data.preserved_tokens = { + "<|message_sep|>\n\n", + "<|role_sep|>\n", + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; + const auto *tool_call_start_prefix = "<|message_sep|>\n\nfunction call<|role_sep|>\n"; + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto ret = p.eps(); + if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { + // Build a choice of all available tools + auto tool_choice = p.choice(); + for (const auto & tool : inputs.tools) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + const auto & schema = function.at("parameters"); + + auto tool_name = p.json_member("name", "\"" + p.tool_name(p.literal(name)) + "\""); + auto tool_args = p.json_member("arguments", p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema))); + + auto tool_open = p.tool_open(p.literal("{") << tool_name); + + tool_choice |= p.rule("tool-" + name, tool_open << "," << tool_args << "}"); + } + + // Define the tool call structure + auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0; + auto max_calls = 1; // parallel toolcalls are not supported + auto tool_call = p.rule("tool-call", p.literal(tool_call_start_prefix) + tool_choice); + auto tool_calls = p.trigger_rule("tool-call-root", p.repeat(tool_call, /* min = */ min_calls, /* max = */ max_calls)); + + ret = p.content(p.until("<|message_sep|>\n\n")) << tool_calls; + } else { + // Content only parser + include_grammar = false; + ret = p.content(p.rest()); + } + + return p.literal(inputs.generation_prompt) + ret; + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO; + + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.at("parameters"); + builder.resolve_refs(schema); + }); + parser.build_grammar(builder, data.grammar_lazy); + }); + + data.grammar_triggers = { + {COMMON_GRAMMAR_TRIGGER_TYPE_WORD, tool_call_start_prefix} + }; + } + return data; +} + +static common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + data.thinking_start_tag = ""; + data.thinking_end_tag = ""; + data.preserved_tokens = { + "｜DSML｜", + "", + "", + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); + + const std::string DSML = "｜DSML｜"; + const std::string THINK_START = ""; + const std::string THINK_END = ""; + const std::string FC_START = "<" + DSML + "function_calls>"; + const std::string FC_END = ""; + const std::string INVOKE_START = "<" + DSML + "invoke"; + const std::string INVOKE_END = ""; + const std::string PARAM_START = "<" + DSML + "parameter"; + const std::string PARAM_END = ""; + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto generation_prompt = p.prefix(inputs.generation_prompt, THINK_START); + auto end = p.end(); + + auto reasoning = p.eps(); + if (extract_reasoning && inputs.enable_thinking) { + reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END); + } else if (extract_reasoning) { + // Thinking disabled but reasoning extraction requested: the generation prompt + // contains an empty pair that must still be consumed. + reasoning = p.optional(p.literal(THINK_START) + p.until(THINK_END) + p.literal(THINK_END)); + } + + if (has_response_format) { + auto response_format = p.rule("response-format", + p.literal("```json") + p.space() + + p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) + + p.space() + p.literal("```")); + return generation_prompt + reasoning + response_format + end; + } + + if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { + return generation_prompt + reasoning + p.content(p.rest()) + end; + } + + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + auto params = function.contains("parameters") ? function.at("parameters") : json::object(); + const auto & props = params.contains("properties") ? params.at("properties") : json::object(); + + std::set required; + if (params.contains("required")) { + params.at("required").get_to(required); + } + + auto schema_info = common_schema_info(); + schema_info.resolve_refs(params); + + std::vector required_parsers; + std::vector optional_parsers; + for (const auto & [param_name, param_schema] : props.items()) { + bool is_required = required.find(param_name) != required.end(); + bool is_string = schema_info.resolves_to_string(param_schema); + + auto arg = p.tool_arg( + p.tool_arg_open( + p.literal(PARAM_START + " name=\"") + + p.tool_arg_name(p.literal(param_name)) + + p.literal("\" string=\"" + std::string(is_string ? "true" : "false") + "\">")) + + (is_string + ? p.tool_arg_string_value(p.until(PARAM_END)) + : p.tool_arg_json_value(p.schema(p.json(), + "tool-" + name + "-arg-" + param_name + "-schema", + param_schema, false))) + + p.tool_arg_close(p.literal(PARAM_END))); + + auto named_arg = p.rule("tool-" + name + "-arg-" + param_name, arg); + if (is_required) { + required_parsers.push_back(named_arg); + } else { + optional_parsers.push_back(named_arg); + } + } + + common_peg_parser args_seq = p.eps(); + for (size_t i = 0; i < required_parsers.size(); i++) { + if (i > 0) { + args_seq = args_seq + p.space(); + } + args_seq = args_seq + required_parsers[i]; + } + + if (!optional_parsers.empty()) { + common_peg_parser any_opt = p.choice(); + for (const auto & opt : optional_parsers) { + any_opt |= opt; + } + args_seq = args_seq + p.repeat(p.space() + any_opt, 0, -1); + } + + common_peg_parser invoke_body = args_seq; + auto func_parser = p.tool( + p.tool_open(p.literal(INVOKE_START + " name=\"") + + p.tool_name(p.literal(name)) + p.literal("\">\n")) + + invoke_body + p.space() + + p.tool_close(p.literal(INVOKE_END))); + + tool_choice |= p.rule("tool-" + name, func_parser); + }); + + auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; + + common_peg_parser tool_calls = p.eps(); + if (inputs.parallel_tool_calls) { + tool_calls = p.trigger_rule("tool-call", + p.literal(FC_START) + p.space() + tool_choice + + p.zero_or_more(p.space() + tool_choice) + p.space() + p.literal(FC_END)); + } else { + tool_calls = p.trigger_rule("tool-call", + p.literal(FC_START) + p.space() + tool_choice + p.space() + p.literal(FC_END)); + } + + if (!require_tools) { + tool_calls = p.optional(tool_calls); + } + + auto content_before_tools = p.content(p.until(FC_START)); + return generation_prompt + reasoning + content_before_tools + tool_calls + end; + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.contains("parameters") ? function.at("parameters") : json::object(); + builder.resolve_refs(schema); + }); + if (has_response_format) { + auto schema = inputs.json_schema; + builder.resolve_refs(schema); + } + parser.build_grammar(builder, data.grammar_lazy); + }); + + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, FC_START }, + }; + } + + return data; +} + +namespace workaround { + +static void map_developer_role_to_system(json & messages) { + for (auto & message : messages) { + if (message.contains("role")) { + if (message["role"] == "developer") { + message["role"] = "system"; + } + } + } +} + + +// if first message is system and template does not support it, merge it with next message +static void system_message_not_supported(json & messages) { + if (!messages.empty() && messages.front().at("role") == "system") { + if (messages.size() > 1) { + LOG_DBG("Merging system prompt into next message\n"); + auto & first_msg = messages.front(); + auto & second_msg = messages[1]; + second_msg["content"] = first_msg.at("content").get() + + "\n" + second_msg.at("content").get(); + messages.erase(messages.begin()); + } else { + LOG_WRN("Removing system prompt due to template not supporting system role\n"); + messages.erase(messages.begin()); + } + } +} + +static void requires_non_null_content(json & messages) { + GGML_ASSERT(messages.is_array()); + for (auto & message : messages) { + if (message.contains("tool_calls") && !message.contains("content")) { + message["content"] = ""; + } + } +} + +// Gemma4 uses a custom tool_responses field instead of role:tool messages. +// +// This will transform a sequence of messages: +// assistant(tool_call+) -> tool+ -> assistant(content) +// +// Into a single assistant message containing a tool_responses field: +// assistant(content + tool_call + tool_responses) +// +// This is necessary for the Gemma4 chat template to properly format the prompt. +// See https://ai.google.dev/gemma/docs/core/prompt-formatting-gemma4 +struct gemma4_model_turn_builder { + json & messages; + size_t pos; + json tool_calls = json::array(); + json tool_responses = json::array(); + json content; + json reasoning_content; + + gemma4_model_turn_builder(json & msgs, size_t pos) : messages(msgs), pos(pos) {} + + void collect() { + // Collect the first assistant message + auto & msg = messages[pos]; + if (msg.contains("reasoning_content") && msg.at("reasoning_content").is_string()) { + // According to the prompt formatting guide, we need to preserve reasoning_content + // between function calls. The current chat templates do not support this, but we will do it anyway. + reasoning_content = msg.at("reasoning_content"); + } + for (auto & tc : msg.at("tool_calls")) { + tool_calls.push_back(tc); + } + pos++; + + // Collect tool call results + while (pos < messages.size() && messages[pos].value("role", "") == "tool") { + collect_result(messages[pos]); + pos++; + } + + // Check if the next assistant message is the final message + if (pos < messages.size() && messages[pos].value("role", "") == "assistant") { + auto & next = messages[pos]; + if (!has_tool_calls(next) && has_content(next)) { + content = next.at("content"); + pos++; + } + } + } + + void collect_result(const json & curr) { + json response; + if (curr.contains("content")) { + const auto & content = curr.at("content"); + if (content.is_string()) { + // Try to parse the content as JSON; fall back to raw string + try { + response = json::parse(content.get()); + } catch (...) { + response = content; + } + } else { + response = content; + } + } + + std::string name; + + // Match name with corresponding tool call + size_t idx = tool_responses.size(); + if (idx < tool_calls.size()) { + auto & tc = tool_calls[idx]; + if (tc.contains("function")) { + name = tc.at("function").value("name", ""); + } + } + + // Fallback to the tool call id + if (name.empty()) { + name = curr.value("tool_call_id", ""); + } + + tool_responses.push_back({{"name", name}, {"response", response}}); + } + + json build() { + collect(); + + json msg = { + {"role", "assistant"}, + {"tool_calls", tool_calls}, + }; + if (!tool_responses.empty()) { + msg["tool_responses"] = tool_responses; + } + if (!content.is_null()) { + msg["content"] = content; + } + if (!reasoning_content.is_null()) { + msg["reasoning_content"] = reasoning_content; + } + return msg; + } + + static bool has_content(const json & msg) { + if (!msg.contains("content") || msg.at("content").is_null()) { + return false; + } + const auto & content = msg.at("content"); + if (content.is_string() && !content.get().empty()) { + return true; + } + if (content.is_array() && !content.empty()) { + return true; + } + return false; + } + + static bool has_tool_calls(const json & msg) { + return msg.contains("tool_calls") && msg.at("tool_calls").is_array() && !msg.at("tool_calls").empty(); + } +}; + +static void convert_tool_responses_gemma4(json & messages) { + json result = json::array(); + size_t i = 0; + + while (i < messages.size()) { + auto & msg = messages[i]; + + if (msg.value("role", "") != "assistant" || !msg.contains("tool_calls") || + !msg.at("tool_calls").is_array() || msg.at("tool_calls").empty()) { + result.push_back(msg); + i++; + continue; + } + + gemma4_model_turn_builder builder(messages, i); + result.push_back(builder.build()); + i = builder.pos; + } + + messages = result; +} + +static void func_args_not_string(json & messages) { + GGML_ASSERT(messages.is_array()); + for (auto & message : messages) { + if (message.contains("tool_calls")) { + for (auto & tool_call : message["tool_calls"]) { + if (tool_call.contains("function") && tool_call["function"].contains("arguments")) { + auto & args = tool_call["function"]["arguments"]; + if (args.is_string()) { + try { + args = json::parse(args.get()); + } catch (const std::exception & e) { + throw std::runtime_error("Failed to parse tool call arguments as JSON: " + std::string(e.what())); + } + } + } + } + } + } +} + +} + +static json common_chat_extra_context() { + json ctx = json::object(); + std::chrono::system_clock::time_point now = std::chrono::system_clock::now(); + std::string datetime_str = format_time(now, "%b %d %Y"); + std::string date_str = format_time(now, "%d %b %Y"); + ctx["datetime"] = datetime_str; + ctx["date_string"] = date_str; + return ctx; +} + +std::optional common_chat_try_specialized_template( + const common_chat_template & tmpl, + const std::string & src, + autoparser::generation_params & params) { + // Ministral/Mistral Large 3 - uses special reasoning structure fixes, can't use autoparser + // Note: Mistral Small 3.2 uses [CALL_ID] which Ministral doesn't have, so we can distinguish them + if (src.find("[SYSTEM_PROMPT]") != std::string::npos && src.find("[TOOL_CALLS]") != std::string::npos && + src.find("[ARGS]") != std::string::npos && src.find("[CALL_ID]") == std::string::npos) { + LOG_DBG("Using specialized template: Ministral/Magistral Large 3\n"); + return common_chat_params_init_ministral_3(tmpl, params); + } + + // GPT-OSS - has unique channel-based structure that needs dedicated handler + if (src.find("<|channel|>") != std::string::npos) { + LOG_DBG("Using specialized template: GPT-OSS\n"); + return common_chat_params_init_gpt_oss(tmpl, params); + } + + // Functionary v3.2 - uses recipient-based format with >>>recipient\n{content} + // Detection: template has ">>>all" for content and ">>>" prefix for tool calls + if (src.find(">>>all") != std::string::npos && src.find(">>>${recipient}") != std::string::npos) { + LOG_DBG("Using specialized template: Functionary v3.2\n"); + return common_chat_params_init_functionary_v3_2(tmpl, params); + } + + // Kimi K2 Thinking - uses unique tool call ID format: functions.: + // Detection: template has "<|tool_calls_section_begin|>" and "functions." prefix in tool call IDs + if (src.find("<|tool_calls_section_begin|>") != std::string::npos && + src.find("<|tool_call_begin|>") != std::string::npos) { + LOG_DBG("Using specialized template: Kimi K2 Thinking\n"); + return common_chat_params_init_kimi_k2(tmpl, params); + } + + if (is_lfm2_template(src)) { + LOG_DBG("Using specialized template: LFM2\n"); + return common_chat_params_init_lfm2(tmpl, params); + } + + // LFM2.5 format detection: template uses plain "List of tools: [...]" with no special tokens + if (src.find("List of tools: [") != std::string::npos && + src.find("<|tool_list_start|>") == std::string::npos) { + LOG_DBG("Using specialized template: LFM2.5\n"); + return common_chat_params_init_lfm2_5(tmpl, params); + } + + // GigaChatV3 format detection + if (src.find("<|role_sep|>") != std::string::npos && + src.find("<|message_sep|>") != std::string::npos && + src.find("<|function_call|>") == std::string::npos) { + LOG_DBG("Using specialized template: GigaChatV3\n"); + return common_chat_params_init_gigachat_v3(tmpl, params); + } + + // DeepSeek V3.2 format detection: template defines dsml_token and uses it for tool calls. + // The template source contains the token as a variable assignment, not as a literal in markup. + if (src.find("dsml_token") != std::string::npos && + src.find("function_calls") != std::string::npos && + src.find("DSML") != std::string::npos) { + LOG_DBG("Using specialized template: DeepSeek V3.2\n"); + return common_chat_params_init_deepseek_v3_2(tmpl, params); + } + + // Gemma4 format detection + if (src.find("'<|tool_call>call:'") != std::string::npos) { + if (src.find("{#- OpenAI Chat Completions:") == std::string::npos) { + // apply workarounds if using the older gemma4 templates + LOG_WRN("%s: detected an outdated gemma4 chat template, applying compatibility workarounds. " + "Consider updating to the official template.\n", __func__); + workaround::convert_tool_responses_gemma4(params.messages); + } + return common_chat_params_init_gemma4(tmpl, params); + } + + return std::nullopt; +} + +static common_chat_params common_chat_templates_apply_jinja(const struct common_chat_templates * tmpls, + const struct common_chat_templates_inputs & inputs) { + autoparser::generation_params params; + params.tools = common_chat_tools_to_json_oaicompat(inputs.tools); + const auto & tmpl = + params.tools.is_array() && tmpls->template_tool_use ? *tmpls->template_tool_use : *tmpls->template_default; + const auto & src = tmpl.source(); + const auto & caps = tmpl.original_caps(); + params.messages = render_message_to_json(inputs.messages, tmpl.original_caps()); + params.tool_choice = inputs.tool_choice; + params.reasoning_format = inputs.reasoning_format; + params.enable_thinking = inputs.enable_thinking; + params.grammar = inputs.grammar; + params.now = inputs.now; + params.add_bos = tmpls->add_bos; + params.add_eos = tmpls->add_eos; + + if (src.find("<|channel|>") == std::string::npos) { + // map developer to system for all models except for GPT-OSS + workaround::map_developer_role_to_system(params.messages); + } + + if (!tmpl.original_caps().supports_system_role) { + workaround::system_message_not_supported(params.messages); + } + + if (tmpl.original_caps().supports_tool_calls) { + // some templates will require the content field in tool call messages + // to still be non-null, this puts an empty string everywhere where the + // content field is null + workaround::requires_non_null_content(params.messages); + } + + if (tmpl.original_caps().supports_object_arguments) { + workaround::func_args_not_string(params.messages); + } + + params.add_generation_prompt = false; + std::string no_gen_prompt = common_chat_template_direct_apply_impl(tmpl, params); + params.add_generation_prompt = true; + std::string gen_prompt = common_chat_template_direct_apply_impl(tmpl, params); + auto diff = calculate_diff_split(no_gen_prompt, gen_prompt); + params.generation_prompt = diff.right + diff.suffix; + + params.add_generation_prompt = inputs.add_generation_prompt; + + params.extra_context = common_chat_extra_context(); + for (auto el : inputs.chat_template_kwargs) { + params.extra_context[el.first] = json::parse(el.second); + } + + if (!inputs.json_schema.empty()) { + params.json_schema = json::parse(inputs.json_schema); + } + + params.parallel_tool_calls = inputs.parallel_tool_calls; + + if (params.tools.is_array()) { + if (params.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE && !params.grammar.empty()) { + throw std::runtime_error("Cannot specify grammar with tools"); + } + if (caps.supports_tool_calls && !caps.supports_tools) { + LOG_WRN( + "Template supports tool calls but does not natively describe tools. The fallback behaviour used may " + "produce bad results, inspect prompt w/ --verbose & consider overriding the template.\n"); + } + } + + if (inputs.force_pure_content) { + LOG_WRN("Forcing pure content template, will not render reasoning or tools separately."); + // Create the result structure + common_chat_params data; + auto params_copy = params; + params_copy.reasoning_format = COMMON_REASONING_FORMAT_NONE; + data.prompt = common_chat_template_direct_apply_impl(tmpl, params_copy); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.generation_prompt = params.generation_prompt; + auto parser = build_chat_peg_parser([¶ms](common_chat_peg_builder &p) { + return p.prefix(params.generation_prompt) << p.content(p.rest()); + }); + data.parser = parser.save(); + return data; + } + + if (auto result = common_chat_try_specialized_template(tmpl, src, params)) { + result->generation_prompt = params.generation_prompt; + return *result; + } + + try { + LOG_DBG("%s: using differential autoparser\n", __func__); + struct autoparser::autoparser autoparser; + autoparser.analyze_template(tmpl); + auto auto_params = autoparser::peg_generator::generate_parser(tmpl, params, autoparser); + auto_params.supports_thinking = autoparser.reasoning.mode != autoparser::reasoning_mode::NONE; + if (auto_params.supports_thinking) { + auto_params.thinking_start_tag = autoparser.reasoning.start; + auto_params.thinking_end_tag = autoparser.reasoning.end; + } + auto_params.generation_prompt = params.generation_prompt; + common_peg_arena arena; + arena.load(auto_params.parser); + LOG_DBG("%s: generated parser:\n%s\n\nparser generation prompt: %s\n", __func__, arena.dump(arena.root()).c_str(), auto_params.generation_prompt.c_str()); + return auto_params; + } catch (const std::exception & e) { + throw std::invalid_argument(std::string("Unable to generate parser for this template. Automatic parser generation failed: ") + e.what()); + } +} + +// Legacy template route (adhoc C++ implementation of known templates), forward to llama_chat_apply_template. +static common_chat_params common_chat_templates_apply_legacy(const struct common_chat_templates * tmpls, + const struct common_chat_templates_inputs & inputs) { + size_t alloc_size = 0; + std::vector chat; + std::vector contents; + + for (const auto & msg : inputs.messages) { + auto content = msg.content; + for (const auto & part : msg.content_parts) { + if (part.type != "text" && part.type != "media_marker") { + LOG_WRN("Ignoring non-text content part: %s\n", part.type.c_str()); + continue; + } + if (!content.empty()) { + content += "\n"; + ; + } + content += part.text; + } + contents.emplace_back(std::move(content)); + } + for (size_t i = 0; i < contents.size(); ++i) { + const auto & msg = inputs.messages[i]; + const auto & content = contents[i]; + chat.push_back({ msg.role.c_str(), content.c_str() }); + size_t msg_size = msg.role.size() + content.size(); + alloc_size += msg_size + (msg_size / 4); // == msg_size * 1.25 but avoiding float ops + } + + std::vector buf(alloc_size); + + // run the first time to get the total output length + const auto & src = tmpls->template_default->source(); + int32_t res = llama_chat_apply_template(src.c_str(), chat.data(), chat.size(), inputs.add_generation_prompt, + buf.data(), buf.size()); + + // error: chat template is not supported + if (res < 0) { + // if the custom "tmpl" is not supported, we throw an error + // this is a bit redundant (for good), since we're not sure if user validated the custom template with llama_chat_verify_template() + throw std::runtime_error("this custom template is not supported, try using --jinja"); + } + + // if it turns out that our buffer is too small, we resize it + if ((size_t) res > buf.size()) { + buf.resize(res); + res = llama_chat_apply_template(src.c_str(), chat.data(), chat.size(), inputs.add_generation_prompt, buf.data(), + buf.size()); + } + + // for safety, we check the result again + if (res < 0 || (size_t) res > buf.size()) { + throw std::runtime_error("failed to apply chat template, try using --jinja"); + } + + common_chat_params params; + params.prompt = std::string(buf.data(), res); + if (!inputs.json_schema.empty()) { + params.grammar = json_schema_to_grammar(json::parse(inputs.json_schema)); + } else { + params.grammar = inputs.grammar; + } + return params; +} + +common_chat_params common_chat_templates_apply(const struct common_chat_templates * tmpls, + const struct common_chat_templates_inputs & inputs) { + GGML_ASSERT(tmpls != nullptr); + return inputs.use_jinja ? common_chat_templates_apply_jinja(tmpls, inputs) : + common_chat_templates_apply_legacy(tmpls, inputs); +} + +common_chat_msg common_chat_parse(const std::string & input, + bool is_partial, + const common_chat_parser_params & params) { + return common_chat_peg_parse(params.parser, input, is_partial, params); +} + +common_chat_msg common_chat_peg_parse(const common_peg_arena & src_parser, + const std::string & input, + bool is_partial, + const common_chat_parser_params & params) { + const common_peg_arena & parser = src_parser.empty() ? + build_chat_peg_parser([](common_chat_peg_builder & p) { return p.content(p.rest()) + p.end(); }) : + src_parser; + + if (src_parser.empty()) { + LOG_DBG("No parser definition detected, assuming pure content parser."); + } + + const std::string effective_input = params.generation_prompt.empty() + ? input + : params.generation_prompt + input; + + //LOG_DBG("Parsing PEG input with format %s: %s\n", common_chat_format_name(params.format), effective_input.c_str()); + + common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_LENIENT; + if (params.debug) { + flags |= COMMON_PEG_PARSE_FLAG_DEBUG; + } + + common_peg_parse_context ctx(effective_input, flags); + auto result = parser.parse(ctx); + + if (result.fail()) { + // During partial parsing, return partial results if any AST nodes were captured + // This allows streaming to work correctly for formats like FUNC_MARKDOWN_CODE_BLOCK + if (is_partial && result.end > 0) { + // Try to extract any partial results from what was successfully parsed + common_chat_msg msg; + msg.role = "assistant"; + std::unique_ptr mapper; + if (params.format == COMMON_CHAT_FORMAT_PEG_GEMMA4) { + mapper = std::make_unique(msg); + } else { + mapper = std::make_unique(msg); + } + mapper->from_ast(ctx.ast, result); + + if (ctx.is_debug()) { + fprintf(stderr, "\nAST for partial parse (fail):\n%s\n", ctx.ast.dump().c_str()); + fflush(stderr); + } + return msg; + } + throw std::runtime_error(std::string("Failed to parse input at pos ") + std::to_string(result.end) + ": " + + effective_input.substr(result.end)); + } + + common_chat_msg msg; + msg.role = "assistant"; + + std::unique_ptr mapper; + if (params.format == COMMON_CHAT_FORMAT_PEG_GEMMA4) { + mapper = std::make_unique(msg); + } else { + mapper = std::make_unique(msg); + } + mapper->from_ast(ctx.ast, result); + + if (ctx.is_debug()) { + fprintf(stderr, "\nAST for %s parse:\n%s\n", is_partial ? "partial" : "full", ctx.ast.dump().c_str()); + fflush(stderr); + } + + if (!is_partial) { + LOG_DBG("Parsed message: %s\n", common_chat_msgs_to_json_oaicompat({ msg }).at(0).dump().c_str()); + } + return msg; +} + +std::map common_chat_templates_get_caps(const common_chat_templates * chat_templates) { + GGML_ASSERT(chat_templates != nullptr); + GGML_ASSERT(chat_templates->template_default != nullptr); + return chat_templates->template_default->caps.to_map(); +} + +#include "common.h.inc" + +#define JSON_ASSERT GGML_ASSERT +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +extern "C" const char * llama_build_info(void); + +#ifndef __EMSCRIPTEN__ +#ifdef __linux__ +#include +#elif defined(_WIN32) +# if !defined(PATH_MAX) +# define PATH_MAX MAX_PATH +# endif +#elif defined(_AIX) +#include +#else +#include +#endif +#endif + +// isatty +#if defined(_WIN32) +#include +#else +#include +#endif + +using json = nlohmann::ordered_json; + +// +// downloader +// + +// validate repo name format: owner/repo +static void write_file(const std::string & fname, const std::string & content) { + const std::string fname_tmp = fname + ".tmp"; + std::ofstream file(fname_tmp); + if (!file) { + throw std::runtime_error(string_format("error: failed to open file '%s'\n", fname.c_str())); + } + + try { + file << content; + file.close(); + + // Makes write atomic + if (rename(fname_tmp.c_str(), fname.c_str()) != 0) { + LOG_ERR("%s: unable to rename file: %s to %s\n", __func__, fname_tmp.c_str(), fname.c_str()); + // If rename fails, try to delete the temporary file + if (remove(fname_tmp.c_str()) != 0) { + LOG_ERR("%s: unable to delete temporary file: %s\n", __func__, fname_tmp.c_str()); + } + } + } catch (...) { + // If anything fails, try to delete the temporary file + if (remove(fname_tmp.c_str()) != 0) { + LOG_ERR("%s: unable to delete temporary file: %s\n", __func__, fname_tmp.c_str()); + } + + throw std::runtime_error(string_format("error: failed to write file '%s'\n", fname.c_str())); + } +} + +static void write_etag(const std::string & path, const std::string & etag) { + const std::string etag_path = path + ".etag"; + write_file(etag_path, etag); + LOG_DBG("%s: file etag saved: %s\n", __func__, etag_path.c_str()); +} + +static std::string read_etag(const std::string & path) { + const std::string etag_path = path + ".etag"; + if (!std::filesystem::exists(etag_path)) { + return {}; + } + std::ifstream etag_in(etag_path); + if (!etag_in) { + LOG_ERR("%s: could not open .etag file for reading: %s\n", __func__, etag_path.c_str()); + return {}; + } + std::string etag; + std::getline(etag_in, etag); + return etag; +} + +static bool is_http_status_ok(int status) { + return llama_common_is_http_status_ok_rust(status); +} + +static std::pair common_http_client(const std::string & url) { + common_http_url parts = common_http_parse_url(url); + + if (parts.host.empty()) { + throw std::runtime_error("error: invalid URL format"); + } + +#ifndef CPPHTTPLIB_OPENSSL_SUPPORT + if (parts.scheme == "https") { + throw std::runtime_error( + "HTTPS is not supported. Please rebuild with one of:\n" + " -DLLAMA_BUILD_BORINGSSL=ON\n" + " -DLLAMA_BUILD_LIBRESSL=ON\n" + " -DLLAMA_OPENSSL=ON (default, requires OpenSSL dev files installed)" + ); + } +#endif + + httplib::Client cli(parts.scheme + "://" + parts.host + ":" + std::to_string(parts.port)); + + if (!parts.user.empty()) { + cli.set_basic_auth(parts.user, parts.password); + } + + cli.set_follow_location(true); + + return { std::move(cli), std::move(parts) }; +} + +std::pair common_download_split_repo_tag(const std::string & hf_repo_with_tag) { + auto parts = string_split(hf_repo_with_tag, ':'); + std::string tag = parts.size() > 1 ? parts.back() : ""; + std::string hf_repo = parts[0]; + if (string_split(hf_repo, '/').size() != 2) { + throw std::invalid_argument("error: invalid HF repo format, expected /[:quant]\n"); + } + return {hf_repo, tag}; +} + +class ProgressBar : public common_download_callback { + static inline std::mutex mutex; + static inline std::map lines; + static inline int max_line = 0; + + std::string filename; + size_t len = 0; + + static void cleanup(const ProgressBar * line) { + lines.erase(line); + if (lines.empty()) { + max_line = 0; + } + } + + static bool is_output_a_tty() { +#if defined(_WIN32) + return _isatty(_fileno(stdout)); +#else + return isatty(1); +#endif + } + +public: + ProgressBar() = default; + + void on_start(const common_download_progress & p) override { + filename = p.url; + + if (auto pos = filename.rfind('/'); pos != std::string::npos) { + filename = filename.substr(pos + 1); + } + if (auto pos = filename.find('?'); pos != std::string::npos) { + filename = filename.substr(0, pos); + } + for (size_t i = 0; i < filename.size(); ++i) { + if ((filename[i] & 0xC0) != 0x80) { + if (len++ == 39) { + filename.resize(i); + filename += "…"; + break; + } + } + } + } + + void on_done(const common_download_progress &, bool) override { + std::lock_guard lock(mutex); + cleanup(this); + } + + void on_update(const common_download_progress & p) override { + if (!p.total || !is_output_a_tty()) { + return; + } + + std::lock_guard lock(mutex); + + if (lines.find(this) == lines.end()) { + lines[this] = max_line++; + std::cout << "\n"; + } + int lines_up = max_line - lines[this]; + + size_t bar = (55 - len) * 2; + size_t pct = (100 * p.downloaded) / p.total; + size_t pos = (bar * p.downloaded) / p.total; + + if (lines_up > 0) { + std::cout << "\033[" << lines_up << "A"; + } + std::cout << '\r' << "Downloading " << filename << " "; + + for (size_t i = 0; i < bar; i += 2) { + std::cout << (i + 1 < pos ? "─" : (i < pos ? "╴" : " ")); + } + std::cout << std::setw(4) << pct << "%\033[K"; + + if (lines_up > 0) { + std::cout << "\033[" << lines_up << "B"; + } + std::cout << '\r' << std::flush; + + if (p.downloaded == p.total) { + cleanup(this); + } + } + + ProgressBar(const ProgressBar &) = delete; + ProgressBar & operator=(const ProgressBar &) = delete; +}; + +static bool common_pull_file(httplib::Client & cli, + const std::string & resolve_path, + const std::string & path_tmp, + bool supports_ranges, + common_download_progress & p, + common_download_callback * callback) { + std::ofstream ofs(path_tmp, std::ios::binary | std::ios::app); + if (!ofs.is_open()) { + LOG_ERR("%s: error opening local file for writing: %s\n", __func__, path_tmp.c_str()); + return false; + } + + httplib::Headers headers; + if (supports_ranges && p.downloaded > 0) { + headers.emplace("Range", "bytes=" + std::to_string(p.downloaded) + "-"); + } + + const char * func = __func__; // avoid __func__ inside a lambda + size_t progress_step = 0; + + auto res = cli.Get(resolve_path, headers, + [&](const httplib::Response &response) { + if (p.downloaded > 0 && response.status != 206) { + LOG_WRN("%s: server did not respond with 206 Partial Content for a resume request. Status: %d\n", func, response.status); + return false; + } + if (p.downloaded == 0 && response.status != 200) { + LOG_WRN("%s: download received non-successful status code: %d\n", func, response.status); + return false; + } + if (p.total == 0 && response.has_header("Content-Length")) { + try { + size_t content_length = std::stoull(response.get_header_value("Content-Length")); + p.total = p.downloaded + content_length; + } catch (const std::exception &e) { + LOG_WRN("%s: invalid Content-Length header: %s\n", func, e.what()); + } + } + return true; + }, + [&](const char *data, size_t len) { + ofs.write(data, len); + if (!ofs) { + LOG_ERR("%s: error writing to file: %s\n", func, path_tmp.c_str()); + return false; + } + p.downloaded += len; + progress_step += len; + + if (progress_step >= p.total / 1000 || p.downloaded == p.total) { + if (callback) { + callback->on_update(p); + if (callback->is_cancelled()) { + return false; + } + } + progress_step = 0; + } + return true; + }, + nullptr + ); + + if (!res) { + LOG_ERR("%s: download failed: %s (status: %d)\n", + __func__, + httplib::to_string(res.error()).c_str(), + res ? res->status : -1); + return false; + } + + return true; +} + +// download one single file from remote URL to local path +// returns status code or -1 on error +static int common_download_file_single_online(const std::string & url, + const std::string & path, + const common_download_opts & opts, + bool skip_etag) { + static const int max_attempts = 3; + static const int retry_delay_seconds = 2; + + const bool file_exists = std::filesystem::exists(path); + + if (file_exists && skip_etag) { + LOG_DBG("%s: using cached file: %s\n", __func__, path.c_str()); + return 304; // 304 Not Modified - fake cached response + } + + auto [cli, parts] = common_http_client(url); + + httplib::Headers headers; + for (const auto & h : opts.headers) { + headers.emplace(h.first, h.second); + } + if (headers.find("User-Agent") == headers.end()) { + headers.emplace("User-Agent", "llama-cpp/" + std::string(llama_build_info())); + } + if (!opts.bearer_token.empty()) { + headers.emplace("Authorization", "Bearer " + opts.bearer_token); + } + cli.set_default_headers(headers); + + std::string last_etag; + if (file_exists) { + last_etag = read_etag(path); + } else { + LOG_DBG("%s: no previous model file found %s\n", __func__, path.c_str()); + } + + auto head = cli.Head(parts.path); + if (!head || head->status < 200 || head->status >= 300) { + LOG_WRN("%s: HEAD failed, status: %d\n", __func__, head ? head->status : -1); + if (file_exists) { + LOG_INF("%s: using cached file (HEAD failed): %s\n", __func__, path.c_str()); + return 304; // 304 Not Modified - fake cached response + } + return head ? head->status : -1; + } + + std::string etag; + if (head->has_header("ETag")) { + etag = head->get_header_value("ETag"); + } + + common_download_progress p; + p.url = url; + if (head->has_header("Content-Length")) { + try { + p.total = std::stoull(head->get_header_value("Content-Length")); + } catch (const std::exception& e) { + LOG_WRN("%s: invalid Content-Length in HEAD response: %s\n", __func__, e.what()); + } + } + + bool supports_ranges = false; + if (head->has_header("Accept-Ranges")) { + supports_ranges = head->get_header_value("Accept-Ranges") != "none"; + } + + if (file_exists) { + if (etag.empty()) { + LOG_DBG("%s: using cached file (no server etag): %s\n", __func__, path.c_str()); + return 304; // 304 Not Modified - fake cached response + } + if (!last_etag.empty() && last_etag == etag) { + LOG_DBG("%s: using cached file (same etag): %s\n", __func__, path.c_str()); + return 304; // 304 Not Modified - fake cached response + } + if (remove(path.c_str()) != 0) { + LOG_ERR("%s: unable to delete file: %s\n", __func__, path.c_str()); + return -1; + } + } + + { // silent + std::error_code ec; + std::filesystem::create_directories(std::filesystem::path(path).parent_path(), ec); + } + + bool success = false; + const std::string path_temporary = path + ".downloadInProgress"; + int delay = retry_delay_seconds; + + if (opts.callback) { + opts.callback->on_start(p); + } + + for (int i = 0; i < max_attempts; ++i) { + if (opts.callback && opts.callback->is_cancelled()) { + break; + } + if (i) { + LOG_WRN("%s: retrying after %d seconds...\n", __func__, delay); + std::this_thread::sleep_for(std::chrono::seconds(delay)); + delay *= retry_delay_seconds; + } + + size_t existing_size = 0; + + if (std::filesystem::exists(path_temporary)) { + if (supports_ranges) { + existing_size = std::filesystem::file_size(path_temporary); + } else if (remove(path_temporary.c_str()) != 0) { + LOG_ERR("%s: unable to delete file: %s\n", __func__, path_temporary.c_str()); + break; + } + } + + p.downloaded = existing_size; + + LOG_DBG("%s: downloading from %s to %s (etag:%s)...\n", + __func__, common_http_show_masked_url(parts).c_str(), + path_temporary.c_str(), etag.c_str()); + + if (common_pull_file(cli, parts.path, path_temporary, supports_ranges, p, opts.callback)) { + if (std::rename(path_temporary.c_str(), path.c_str()) != 0) { + LOG_ERR("%s: unable to rename file: %s to %s\n", __func__, path_temporary.c_str(), path.c_str()); + break; + } + if (!etag.empty() && !skip_etag) { + write_etag(path, etag); + } + success = true; + break; + } + } + + if (opts.callback) { + opts.callback->on_done(p, success); + } + if (opts.callback && opts.callback->is_cancelled() && + std::filesystem::exists(path_temporary)) { + if (remove(path_temporary.c_str()) != 0) { + LOG_ERR("%s: unable to delete temporary file: %s\n", __func__, path_temporary.c_str()); + } + } + if (!success) { + LOG_ERR("%s: download failed after %d attempts\n", __func__, max_attempts); + return -1; // max attempts reached + } + + return head->status; +} + +std::pair> common_remote_get_content(const std::string & url, + const common_remote_params & params) { + auto [cli, parts] = common_http_client(url); + + httplib::Headers headers; + for (const auto & h : params.headers) { + headers.emplace(h.first, h.second); + } + if (headers.find("User-Agent") == headers.end()) { + headers.emplace("User-Agent", "llama-cpp/" + std::string(llama_build_info())); + } + + if (params.timeout > 0) { + cli.set_read_timeout(params.timeout, 0); + cli.set_write_timeout(params.timeout, 0); + } + + std::vector buf; + auto res = cli.Get(parts.path, headers, + [&](const char *data, size_t len) { + buf.insert(buf.end(), data, data + len); + return params.max_size == 0 || + buf.size() <= static_cast(params.max_size); + }, + nullptr + ); + + if (!res) { + throw std::runtime_error("error: cannot make GET request"); + } + + return { res->status, std::move(buf) }; +} + +int common_download_file_single(const std::string & url, + const std::string & path, + const common_download_opts & opts, + bool skip_etag) { + if (!opts.offline) { + ProgressBar tty_cb; + common_download_opts online_opts = opts; + if (!online_opts.callback) { + online_opts.callback = &tty_cb; + } + return common_download_file_single_online(url, path, online_opts, skip_etag); + } + + if (!std::filesystem::exists(path)) { + LOG_ERR("%s: required file is not available in cache (offline mode): %s\n", __func__, path.c_str()); + return -1; + } + + LOG_DBG("%s: using cached file (offline mode): %s\n", __func__, path.c_str()); + + // notify the callback that the file was cached + if (opts.callback) { + common_download_progress p; + p.url = url; + p.cached = true; + opts.callback->on_start(p); + opts.callback->on_done(p, true); + } + + return 304; // Not Modified - fake cached response +} + +struct gguf_split_info { + std::string prefix; // tag included + std::string tag; + int index; + int count; +}; + +static gguf_split_info get_gguf_split_info(const std::string & path) { + common_rust_gguf_split_info split = + llama_common_gguf_split_info_rust((const uint8_t *) path.data(), path.size(), true); + gguf_split_info result { + common_rust_string_to_std(split.prefix), + common_rust_string_to_std(split.tag), + split.index, + split.count, + }; + return result; +} + +static hf_cache::hf_files get_split_files(const hf_cache::hf_files & files, + const hf_cache::hf_file & file) { + auto split = get_gguf_split_info(file.path); + + if (split.count <= 1) { + return {file}; + } + hf_cache::hf_files result; + + for (const auto & f : files) { + auto split_f = get_gguf_split_info(f.path); + if (split_f.count == split.count && split_f.prefix == split.prefix) { + result.push_back(f); + } + } + return result; +} + +static hf_cache::hf_file find_best_mmproj(const hf_cache::hf_files & files, + const std::string & model) { + hf_cache::hf_file best; + size_t best_depth = 0; + int best_diff = 0; + bool found = false; + + auto model_bits = llama_common_gguf_extract_quant_bits_rust((const uint8_t *) model.data(), model.size()); + auto model_parts = string_split(model, '/'); + auto model_dir = model_parts.end() - 1; + + for (const auto & f : files) { + if (!string_ends_with(f.path, ".gguf") || + f.path.find("mmproj") == std::string::npos) { + continue; + } + + auto mmproj_parts = string_split(f.path, '/'); + auto mmproj_dir = mmproj_parts.end() - 1; + + auto [_, dir] = std::mismatch(model_parts.begin(), model_dir, + mmproj_parts.begin(), mmproj_dir); + if (dir != mmproj_dir) { + continue; + } + + size_t depth = dir - mmproj_parts.begin(); + auto bits = llama_common_gguf_extract_quant_bits_rust((const uint8_t *) f.path.data(), f.path.size()); + auto diff = std::abs(bits - model_bits); + + if (!found || depth > best_depth || (depth == best_depth && diff < best_diff)) { + best = f; + best_depth = depth; + best_diff = diff; + found = true; + } + } + return best; +} + +static hf_cache::hf_file find_best_model(const hf_cache::hf_files & files, + const std::string & tag) { + std::vector tags; + + if (!tag.empty()) { + tags.push_back(tag); + } else { + tags = {"Q4_K_M", "Q8_0"}; + } + + for (const auto & t : tags) { + std::regex pattern(t + "[.-]", std::regex::icase); + for (const auto & f : files) { + if (llama_common_gguf_filename_is_model_rust((const uint8_t *) f.path.data(), f.path.size()) && + std::regex_search(f.path, pattern)) { + auto split = get_gguf_split_info(f.path); + if (split.count > 1 && split.index != 1) { + continue; + } + return f; + } + } + } + + // fallback to first available model only if tag is empty + if (tag.empty()) { + for (const auto & f : files) { + if (llama_common_gguf_filename_is_model_rust((const uint8_t *) f.path.data(), f.path.size())) { + auto split = get_gguf_split_info(f.path); + if (split.count > 1 && split.index != 1) { + continue; + } + return f; + } + } + } + + return {}; +} + +static void list_available_gguf_files(const hf_cache::hf_files & files) { + LOG_INF("Available GGUF files:\n"); + for (const auto & f : files) { + if (string_ends_with(f.path, ".gguf")) { + LOG_INF(" - %s\n", f.path.c_str()); + } + } +} + +struct hf_plan { + hf_cache::hf_file primary; + hf_cache::hf_files model_files; + hf_cache::hf_file mmproj; +}; + +static hf_plan get_hf_plan(const common_params_model & model, + const common_download_opts & opts, + bool download_mmproj) { + hf_plan plan; + hf_cache::hf_files all; + + auto [repo, tag] = common_download_split_repo_tag(model.hf_repo); + + if (!opts.offline) { + all = hf_cache::get_repo_files(repo, opts.bearer_token); + } + if (all.empty()) { + all = hf_cache::get_cached_files(repo); + } + if (all.empty()) { + return plan; + } + + hf_cache::hf_file primary; + + if (!model.hf_file.empty()) { + for (const auto & f : all) { + if (f.path == model.hf_file) { + primary = f; + break; + } + } + if (primary.path.empty()) { + LOG_ERR("%s: file '%s' not found in repository\n", __func__, model.hf_file.c_str()); + list_available_gguf_files(all); + return plan; + } + } else { + primary = find_best_model(all, tag); + if (primary.path.empty()) { + LOG_ERR("%s: no GGUF files found in repository %s\n", __func__, repo.c_str()); + list_available_gguf_files(all); + return plan; + } + } + + plan.primary = primary; + plan.model_files = get_split_files(all, primary); + + if (download_mmproj) { + plan.mmproj = find_best_mmproj(all, primary.path); + } + + return plan; +} + +struct download_task { + std::string url; + std::string path; +}; + +static std::vector get_url_tasks(const common_params_model & model) { + auto split = get_gguf_split_info(model.url); + + if (split.count <= 1) { + return {{model.url, model.path}}; + } + + auto filename = split.prefix; + if (auto pos = split.prefix.rfind('/'); pos != std::string::npos) { + filename = split.prefix.substr(pos + 1); + } + + auto parent_path = std::filesystem::path(model.path).parent_path(); + auto prefix_path = (parent_path / filename).string(); + + std::vector tasks; + for (int i = 1; i <= split.count; i++) { + auto suffix = string_format("-%05d-of-%05d.gguf", i, split.count); + tasks.push_back({split.prefix + suffix, prefix_path + suffix}); + } + return tasks; +} + +common_download_model_result common_download_model(const common_params_model & model, + const common_download_opts & opts, + bool download_mmproj) { + common_download_model_result result; + std::vector tasks; + hf_plan hf; + + bool is_hf = !model.hf_repo.empty(); + + if (is_hf) { + hf = get_hf_plan(model, opts, download_mmproj); + for (const auto & f : hf.model_files) { + tasks.push_back({f.url, f.local_path}); + } + if (!hf.mmproj.path.empty()) { + tasks.push_back({hf.mmproj.url, hf.mmproj.local_path}); + } + } else if (!model.url.empty()) { + tasks = get_url_tasks(model); + } else { + result.model_path = model.path; + return result; + } + + if (tasks.empty()) { + return result; + } + + std::vector> futures; + for (const auto & task : tasks) { + futures.push_back(std::async(std::launch::async, + [&task, &opts, is_hf]() { + int status = common_download_file_single(task.url, task.path, opts, is_hf); + return is_http_status_ok(status); + } + )); + } + + for (auto & f : futures) { + if (!f.get()) { + return {}; + } + } + + if (is_hf) { + for (const auto & f : hf.model_files) { + hf_cache::finalize_file(f); + } + result.model_path = hf.primary.final_path; + + if (!hf.mmproj.path.empty()) { + result.mmproj_path = hf_cache::finalize_file(hf.mmproj); + } + } else { + result.model_path = model.path; + } + + return result; +} + +// +// Docker registry functions +// + +static std::string common_docker_get_token(const std::string & repo) { + std::string url = "https://auth.docker.io/token?service=registry.docker.io&scope=repository:" + repo + ":pull"; + + common_remote_params params; + auto res = common_remote_get_content(url, params); + + if (res.first != 200) { + throw std::runtime_error("Failed to get Docker registry token, HTTP code: " + std::to_string(res.first)); + } + + std::string response_str(res.second.begin(), res.second.end()); + nlohmann::ordered_json response = nlohmann::ordered_json::parse(response_str); + + if (!response.contains("token")) { + throw std::runtime_error("Docker registry token response missing 'token' field"); + } + + return response["token"].get(); +} + +std::string common_docker_resolve_model(const std::string & docker) { + // Parse ai/smollm2:135M-Q4_0 + size_t colon_pos = docker.find(':'); + std::string repo, tag; + if (colon_pos != std::string::npos) { + repo = docker.substr(0, colon_pos); + tag = docker.substr(colon_pos + 1); + } else { + repo = docker; + tag = "latest"; + } + + // ai/ is the default + size_t slash_pos = docker.find('/'); + if (slash_pos == std::string::npos) { + repo.insert(0, "ai/"); + } + + LOG_INF("%s: Downloading Docker Model: %s:%s\n", __func__, repo.c_str(), tag.c_str()); + try { + // --- helper: digest validation --- + auto validate_oci_digest = [](const std::string & digest) -> std::string { + // Expected: algo:hex ; start with sha256 (64 hex chars) + // You can extend this map if supporting other algorithms in future. + static const std::regex re("^sha256:([a-fA-F0-9]{64})$"); + std::smatch m; + if (!std::regex_match(digest, m, re)) { + throw std::runtime_error("Invalid OCI digest format received in manifest: " + digest); + } + // normalize hex to lowercase + std::string normalized = digest; + std::transform(normalized.begin()+7, normalized.end(), normalized.begin()+7, [](unsigned char c){ + return std::tolower(c); + }); + return normalized; + }; + + std::string token = common_docker_get_token(repo); // Get authentication token + + // Get manifest + // TODO: cache the manifest response so that it appears in the model list + const std::string url_prefix = "https://registry-1.docker.io/v2/" + repo; + std::string manifest_url = url_prefix + "/manifests/" + tag; + common_remote_params manifest_params; + manifest_params.headers.push_back({"Authorization", "Bearer " + token}); + manifest_params.headers.push_back({"Accept", + "application/vnd.docker.distribution.manifest.v2+json,application/vnd.oci.image.manifest.v1+json" + }); + auto manifest_res = common_remote_get_content(manifest_url, manifest_params); + if (manifest_res.first != 200) { + throw std::runtime_error("Failed to get Docker manifest, HTTP code: " + std::to_string(manifest_res.first)); + } + + std::string manifest_str(manifest_res.second.begin(), manifest_res.second.end()); + nlohmann::ordered_json manifest = nlohmann::ordered_json::parse(manifest_str); + std::string gguf_digest; // Find the GGUF layer + if (manifest.contains("layers")) { + for (const auto & layer : manifest["layers"]) { + if (layer.contains("mediaType")) { + std::string media_type = layer["mediaType"].get(); + if (media_type == "application/vnd.docker.ai.gguf.v3" || + media_type.find("gguf") != std::string::npos) { + gguf_digest = layer["digest"].get(); + break; + } + } + } + } + + if (gguf_digest.empty()) { + throw std::runtime_error("No GGUF layer found in Docker manifest"); + } + + // Validate & normalize digest + gguf_digest = validate_oci_digest(gguf_digest); + LOG_DBG("%s: Using validated digest: %s\n", __func__, gguf_digest.c_str()); + + // Prepare local filename + std::string model_filename = repo; + std::replace(model_filename.begin(), model_filename.end(), '/', '_'); + model_filename += "_" + tag + ".gguf"; + std::string local_path = fs_get_cache_file(model_filename); + + const std::string blob_url = url_prefix + "/blobs/" + gguf_digest; + common_download_opts opts; + opts.bearer_token = token; + const int http_status = common_download_file_single(blob_url, local_path, opts); + if (!is_http_status_ok(http_status)) { + throw std::runtime_error("Failed to download Docker Model"); + } + + LOG_INF("%s: Downloaded Docker Model to: %s\n", __func__, local_path.c_str()); + return local_path; + } catch (const std::exception & e) { + LOG_ERR("%s: Docker Model download failed: %s\n", __func__, e.what()); + throw; + } +} + +std::vector common_list_cached_models() { + std::unordered_set seen; + std::vector result; + + auto files = hf_cache::get_cached_files(); + + for (const auto & f : files) { + auto split = get_gguf_split_info(f.path); + if (split.index != 1 || split.tag.empty() || + split.prefix.find("mmproj") != std::string::npos) { + continue; + } + if (seen.insert(f.repo_id + ":" + split.tag).second) { + result.push_back({f.repo_id, split.tag}); + } + } + + return result; +} + + +#include "common.h.inc" + +#include +#include +#include +#include +#include +#include +#include +#include + +// Trick to catch missing branches +template +inline constexpr bool is_always_false_v = false; + +const char * common_peg_parse_result_type_name(common_peg_parse_result_type type) { + return llama_common_peg_parse_result_type_name_rust((int) type); +} + +// Trie for matching multiple literals. +// This is used in common_peg_until_parser and to build a GBNF exclusion grammar +struct trie { + struct node { + std::map children; // Use uint32_t to store Unicode codepoints + bool is_word; + }; + + std::vector nodes; + + trie(const std::vector & words) { + create_node(); // root node + for (const auto & w : words) { + insert(w); + } + } + + enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH }; + + // Check if a delimiter starts at the given position + match_result check_at(std::string_view sv, size_t start_pos) const { + size_t current = 0; // Start at root + size_t pos = start_pos; + + // LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str()); + + while (pos < sv.size()) { + auto result = common_parse_utf8_codepoint(sv, pos); + if (result.status != utf8_parse_result::SUCCESS) { + break; + } + + auto it = nodes[current].children.find(result.codepoint); + if (it == nodes[current].children.end()) { + // Can't continue matching + return match_result{match_result::NO_MATCH}; + } + + current = it->second; + pos += result.bytes_consumed; + + // Check if we've matched a complete word + if (nodes[current].is_word) { + return match_result{match_result::COMPLETE_MATCH}; + } + } + + // Reached end of input while still in the trie (not at root) + if (current != 0) { + // We're in the middle of a potential match + return match_result{match_result::PARTIAL_MATCH}; + } + + // Reached end at root (no match) + return match_result{match_result::NO_MATCH}; + } + + struct prefix_and_next { + std::vector prefix; + std::vector next_chars; + }; + + std::vector collect_prefix_and_next() { + std::vector prefix; + std::vector result; + collect_prefix_and_next(0, prefix, result); + return result; + } + + private: + void collect_prefix_and_next(size_t index, std::vector & prefix, std::vector & out) { + if (!nodes[index].is_word) { + if (!nodes[index].children.empty()) { + std::vector chars; + chars.reserve(nodes[index].children.size()); + for (const auto & p : nodes[index].children) { + chars.push_back(p.first); + } + out.emplace_back(prefix_and_next{prefix, chars}); + } + } + + for (const auto & p : nodes[index].children) { + uint32_t ch = p.first; + auto child = p.second; + prefix.push_back(ch); + collect_prefix_and_next(child, prefix, out); + prefix.pop_back(); + } + } + + size_t create_node() { + size_t index = nodes.size(); + nodes.emplace_back(); + return index; + } + + void insert(const std::string & word) { + size_t current = 0; + size_t pos = 0; + while (pos < word.length()) { + auto result = common_parse_utf8_codepoint(word, pos); + if (result.status != utf8_parse_result::SUCCESS) { + break; + } + + uint32_t ch = result.codepoint; + pos += result.bytes_consumed; + + auto it = nodes[current].children.find(ch); + if (it == nodes[current].children.end()) { + size_t child = create_node(); + nodes[current].children[ch] = child; + current = child; + } else { + current = it->second; + } + } + nodes[current].is_word = true; + } +}; + +static std::pair parse_hex_escape(const std::string & str, size_t pos, int hex_count) { + uint32_t value = 0; + const size_t consumed = llama_common_parse_hex_escape_rust( + (const uint8_t *) str.data(), str.size(), pos, hex_count, &value); + if (consumed == 0) { + return {0, 0}; + } + return {value, consumed}; +} + +static std::pair parse_char_class_char(const std::string & content, size_t pos) { + uint32_t value = 0; + const size_t consumed = llama_common_parse_char_class_char_rust( + (const uint8_t *) content.data(), content.size(), pos, &value); + return {value, consumed}; +} + +static std::pair, bool> parse_char_classes(const std::string & classes) { + std::vector ranges; + bool negated = false; + + std::string content = classes; + if (content.front() == '[') { + content = content.substr(1); + } + + if (content.back() == ']') { + content.pop_back(); + } + + // Check for negation + if (!content.empty() && content.front() == '^') { + negated = true; + content = content.substr(1); + } + + size_t i = 0; + while (i < content.length()) { + auto [start, start_len] = parse_char_class_char(content, i); + i += start_len; + + if (i + 1 < content.length() && content[i] == '-') { + // Range detected + auto [end, end_len] = parse_char_class_char(content, i + 1); + ranges.push_back(common_peg_chars_parser::char_range{start, end}); + i += 1 + end_len; + } else { + ranges.push_back(common_peg_chars_parser::char_range{start, start}); + } + } + + return {ranges, negated}; +} + +common_peg_ast_id common_peg_ast_arena::find_by_tag(const common_peg_ast_node & parent, const std::string & tag, int max_depth) const { + for (auto child_id : parent.children) { + const auto & child = get(child_id); + if (child.tag == tag) { + return child_id; + } + if (max_depth > 1) { + auto result = find_by_tag(child, tag, max_depth - 1); + if (result != COMMON_PEG_INVALID_AST_ID) { + return result; + } + } + } + return COMMON_PEG_INVALID_AST_ID; +} + +common_peg_ast_id common_peg_ast_arena::find_by_rule(const common_peg_ast_node & parent, const std::string & rule, int max_depth) const { + for (auto child_id : parent.children) { + const auto & child = get(child_id); + if (child.rule == rule) { + return child_id; + } + if (max_depth > 1) { + auto result = find_by_rule(child, rule, max_depth - 1); + if (result != COMMON_PEG_INVALID_AST_ID) { + return result; + } + } + } + return COMMON_PEG_INVALID_AST_ID; +} + +void common_peg_ast_arena::visit(common_peg_ast_id id, const common_peg_ast_visitor & visitor) const { + if (id == COMMON_PEG_INVALID_AST_ID) { + return; + } + const auto & node = get(id); + visitor(node); + for (const auto & child : node.children) { + visit(child, visitor); + } +} + +void common_peg_ast_arena::visit(const common_peg_parse_result & result, const common_peg_ast_visitor & visitor) const { + for (const auto & node : result.nodes) { + visit(node, visitor); + } +} + +struct parser_executor; + +common_peg_parser_id common_peg_arena::add_parser(common_peg_parser_variant parser) { + common_peg_parser_id id = parsers_.size(); + parsers_.push_back(std::move(parser)); + return id; +} + +void common_peg_arena::add_rule(const std::string & name, common_peg_parser_id id) { + rules_[name] = id; +} + +common_peg_parser_id common_peg_arena::get_rule(const std::string & name) const { + auto it = rules_.find(name); + if (it == rules_.end()) { + throw std::runtime_error("Rule not found: " + name); + } + return it->second; +} + +struct parser_executor { + const common_peg_arena & arena; + common_peg_parse_context & ctx; + size_t start_pos; + + parser_executor(const common_peg_arena & arena, common_peg_parse_context & ctx, size_t start) + : arena(arena), ctx(ctx), start_pos(start) {} + + std::string debug_indent() const { return std::string(ctx.parse_depth * 2, ' '); } + + std::string debug_input_snippet(size_t pos, size_t len = 60) const { + if (pos >= ctx.input.size()) { + return ""; + } + auto snippet = ctx.input.substr(pos, len); + // Escape newlines for display + std::string result; + for (char c : snippet) { + if (c == '\n') { + result += "\\n"; + } else if (c == '\r') { + result += "\\r"; + } else if (c == '\t') { + result += "\\t"; + } else { + result += c; + } + } + if (pos + len < ctx.input.size()) { + result += "..."; + } + return result; + } + + common_peg_parse_result operator()(const common_peg_epsilon_parser & /* p */) const { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos); + } + + common_peg_parse_result operator()(const common_peg_start_parser & /* p */) const { + return common_peg_parse_result( + start_pos == 0 ? COMMON_PEG_PARSE_RESULT_SUCCESS : COMMON_PEG_PARSE_RESULT_FAIL, + start_pos + ); + } + + common_peg_parse_result operator()(const common_peg_end_parser & /* p */) const { + return common_peg_parse_result( + start_pos >= ctx.input.size() ? COMMON_PEG_PARSE_RESULT_SUCCESS : COMMON_PEG_PARSE_RESULT_FAIL, + start_pos + ); + } + + common_peg_parse_result operator()(const common_peg_literal_parser & p) { + auto pos = start_pos; + for (auto i = 0u; i < p.literal.size(); ++i) { + if (pos >= ctx.input.size()) { + if (!ctx.is_lenient()) { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos); + } + if (ctx.input[pos] != p.literal[i]) { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); + } + ++pos; + } + + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); + } + + common_peg_parse_result operator()(const common_peg_sequence_parser & p) { + if (ctx.is_debug()) { + LOG_DBG("%sSEQ start at %zu '%s' (%zu children)\n", debug_indent().c_str(), start_pos, + debug_input_snippet(start_pos).c_str(), p.children.size()); + } + ctx.parse_depth++; + + auto pos = start_pos; + std::vector nodes; + + for (size_t i = 0; i < p.children.size(); i++) { + const auto & child_id = p.children[i]; + if (ctx.is_debug()) { + fprintf(stderr, "%sSEQ child %zu: %s\n", debug_indent().c_str(), i, arena.dump(child_id).c_str()); + } + auto result = arena.parse(child_id, ctx, pos); + + if (ctx.is_debug()) { + fprintf(stderr, "%sSEQ child %zu: %s at %zu->%zu\n", debug_indent().c_str(), i, + common_peg_parse_result_type_name(result.type), result.start, result.end); + } + + if (result.fail()) { + ctx.parse_depth--; + if (ctx.is_debug()) { + fprintf(stderr, "%sSEQ -> FAIL\n", debug_indent().c_str()); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos, result.end); + } + + if (!result.nodes.empty()) { + nodes.insert(nodes.end(), result.nodes.begin(), result.nodes.end()); + } + + if (result.need_more_input()) { + ctx.parse_depth--; + if (ctx.is_debug()) { + fprintf(stderr, "%sSEQ -> NEED_MORE\n", debug_indent().c_str()); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, result.end, std::move(nodes)); + } + + pos = result.end; + } + + ctx.parse_depth--; + if (ctx.is_debug()) { + fprintf(stderr, "%sSEQ -> SUCCESS at %zu->%zu\n", debug_indent().c_str(), start_pos, pos); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos, std::move(nodes)); + } + + common_peg_parse_result operator()(const common_peg_choice_parser & p) { + if (ctx.is_debug()) { + fprintf(stderr, "%sCHOICE start at %zu '%s' (%zu options)\n", debug_indent().c_str(), start_pos, + debug_input_snippet(start_pos).c_str(), p.children.size()); + } + ctx.parse_depth++; + + auto pos = start_pos; + for (size_t i = 0; i < p.children.size(); i++) { + const auto & child_id = p.children[i]; + if (ctx.is_debug()) { + fprintf(stderr, "%sCHOICE option %zu: %s\n", debug_indent().c_str(), i, arena.dump(child_id).c_str()); + } + auto result = arena.parse(child_id, ctx, pos); + if (ctx.is_debug()) { + fprintf(stderr, "%sCHOICE option %zu: %s\n", debug_indent().c_str(), i, + common_peg_parse_result_type_name(result.type)); + } + if (!result.fail()) { + ctx.parse_depth--; + if (ctx.is_debug()) { + fprintf(stderr, "%sCHOICE -> %s (option %zu)\n", debug_indent().c_str(), + common_peg_parse_result_type_name(result.type), i); + } + return result; + } + } + + ctx.parse_depth--; + if (ctx.is_debug()) { + fprintf(stderr, "%sCHOICE -> FAIL (no options matched)\n", debug_indent().c_str()); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); + } + + common_peg_parse_result operator()(const common_peg_repetition_parser & p) { + if (ctx.is_debug()) { + fprintf(stderr, "%sREPEAT start at %zu '%s' (min=%d, max=%d)\n", debug_indent().c_str(), start_pos, + debug_input_snippet(start_pos).c_str(), p.min_count, p.max_count); + } + ctx.parse_depth++; + + auto pos = start_pos; + int match_count = 0; + std::vector nodes; + + // Try to match up to max_count times (or unlimited if max_count is -1) + while (p.max_count == -1 || match_count < p.max_count) { + if (pos >= ctx.input.size()) { + if (ctx.is_debug()) { + fprintf(stderr, "%sREPEAT: at end of input, count=%d\n", debug_indent().c_str(), match_count); + } + break; + } + + auto result = arena.parse(p.child, ctx, pos); + + if (ctx.is_debug()) { + fprintf(stderr, "%sREPEAT iter %d: %s at %zu->%zu, nodes=%zu\n", debug_indent().c_str(), match_count, + common_peg_parse_result_type_name(result.type), result.start, result.end, result.nodes.size()); + fprintf(stderr, "%sREPEAT CHILD: %s\n", debug_indent().c_str(), arena.dump(p.child).c_str()); + } + + if (result.success()) { + // Prevent infinite loop on empty matches + if (result.end == pos) { + if (ctx.is_debug()) { + fprintf(stderr, "%s REPEAT: empty match, stopping\n", debug_indent().c_str()); + } + break; + } + + if (!result.nodes.empty()) { + nodes.insert(nodes.end(), result.nodes.begin(), result.nodes.end()); + } + + pos = result.end; + match_count++; + continue; + } + + if (result.need_more_input()) { + if (!result.nodes.empty()) { + nodes.insert(nodes.end(), result.nodes.begin(), result.nodes.end()); + } + + ctx.parse_depth--; + if (ctx.is_debug()) { + fprintf(stderr, "%sREPEAT -> NEED_MORE (count=%d, nodes=%zu)\n", debug_indent().c_str(), + match_count, nodes.size()); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, result.end, std::move(nodes)); + } + + // Child failed - stop trying + if (ctx.is_debug()) { + fprintf(stderr, "%sREPEAT: child failed, stopping\n", debug_indent().c_str()); + } + break; + } + + // Check if we got enough matches + if (p.min_count > 0 && match_count < p.min_count) { + ctx.parse_depth--; + if (pos >= ctx.input.size() && ctx.is_lenient()) { + if (ctx.is_debug()) { + fprintf(stderr, "%sREPEAT -> NEED_MORE (not enough matches: %d < %d)\n", debug_indent().c_str(), + match_count, p.min_count); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos, std::move(nodes)); + } + if (ctx.is_debug()) { + fprintf(stderr, "%sREPEAT -> FAIL (not enough matches: %d < %d)\n", debug_indent().c_str(), match_count, + p.min_count); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos, pos); + } + + ctx.parse_depth--; + if (ctx.is_debug()) { + fprintf(stderr, "%sREPEAT -> SUCCESS (count=%d, nodes=%zu)\n", debug_indent().c_str(), match_count, + nodes.size()); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos, std::move(nodes)); + } + + common_peg_parse_result operator()(const common_peg_and_parser & p) { + auto result = arena.parse(p.child, ctx, start_pos); + // Pass result but don't consume input + return common_peg_parse_result(result.type, start_pos); + } + + common_peg_parse_result operator()(const common_peg_not_parser & p) { + auto result = arena.parse(p.child, ctx, start_pos); + + if (result.success()) { + // Fail if the underlying parser matches + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); + } + + if (result.need_more_input()) { + // Propagate - need to know what child would match before negating + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos); + } + + // Child failed, so negation succeeds + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos); + } + + common_peg_parse_result operator()(const common_peg_any_parser & /* p */) const { + // Parse a single UTF-8 codepoint (not just a single byte) + auto result = common_parse_utf8_codepoint(ctx.input, start_pos); + + if (result.status == utf8_parse_result::INCOMPLETE) { + if (!ctx.is_lenient()) { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos); + } + if (result.status == utf8_parse_result::INVALID) { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, start_pos + result.bytes_consumed); + } + + common_peg_parse_result operator()(const common_peg_space_parser & /* p */) { + auto pos = start_pos; + while (pos < ctx.input.size()) { + auto c = static_cast(ctx.input[pos]); + if (std::isspace(c)) { + ++pos; + } else { + break; + } + } + + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); + } + + common_peg_parse_result operator()(const common_peg_chars_parser & p) const { + auto pos = start_pos; + int match_count = 0; + + // Try to match up to max_count times (or unlimited if max_count is -1) + while (p.max_count == -1 || match_count < p.max_count) { + auto result = common_parse_utf8_codepoint(ctx.input, pos); + + if (result.status == utf8_parse_result::INCOMPLETE) { + if (match_count >= p.min_count) { + // We have enough matches, succeed with what we have + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); + } + // Not enough matches yet + if (!ctx.is_lenient()) { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos); + } + + if (result.status == utf8_parse_result::INVALID) { + // Malformed UTF-8 in input + if (match_count >= p.min_count) { + // We have enough matches, succeed up to here + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); + } + // Not enough matches, fail + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); + } + + // Check if this codepoint matches our character class + bool matches = false; + for (const auto & range : p.ranges) { + if (range.contains(result.codepoint)) { + matches = true; + break; + } + } + + // If negated, invert the match result + if (p.negated) { + matches = !matches; + } + + if (matches) { + pos += result.bytes_consumed; + ++match_count; + } else { + // Character doesn't match, stop matching + break; + } + } + + // Check if we got enough matches + if (match_count < p.min_count) { + if (pos >= ctx.input.size() && ctx.is_lenient()) { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos, pos); + } + + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); + } + + static common_peg_parse_result handle_escape_sequence(common_peg_parse_context & ctx, size_t start, size_t & pos, const char delimiter) { + ++pos; // consume '\' + if (pos >= ctx.input.size()) { + if (!ctx.is_lenient()) { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start, pos); + } + + char c = ctx.input[pos]; + if (c == delimiter || c == '\\' || c == '/' || c == 'b' || c == 'f' || c == 'n' || c == 'r' || c == 't') { + ++pos; + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start, pos); + } else if (c == 'u') { + return handle_unicode_escape(ctx, start, pos); + } else { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); + } + } + + static common_peg_parse_result handle_unicode_escape(common_peg_parse_context & ctx, size_t start, size_t & pos) { + ++pos; // consume 'u' + for (int i = 0; i < 4; ++i) { + if (pos >= ctx.input.size()) { + if (!ctx.is_lenient()) { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start, pos); + } + if (!llama_common_is_hex_digit_rust((uint8_t) ctx.input[pos])) { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); + } + ++pos; + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start, pos); + } + + common_peg_parse_result operator()(const common_peg_string_parser & p) { + auto pos = start_pos; + + // Parse string content (without quotes) + while (pos < ctx.input.size()) { + char c = ctx.input[pos]; + + if (c == p.delimiter) { + // Found closing delimiter - success (don't consume it) + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); + } + + if (c == '\\') { + auto result = handle_escape_sequence(ctx, start_pos, pos, p.delimiter); + if (!result.success()) { + return result; + } + } else { + auto utf8_result = common_parse_utf8_codepoint(ctx.input, pos); + + if (utf8_result.status == utf8_parse_result::INCOMPLETE) { + if (!ctx.is_lenient()) { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos); + } + + if (utf8_result.status == utf8_parse_result::INVALID) { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); + } + + pos += utf8_result.bytes_consumed; + } + } + + // Reached end without finding closing quote + if (!ctx.is_lenient()) { + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos, pos); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos); + } + + common_peg_parse_result operator()(const common_peg_until_parser & p) const { + trie matcher(p.delimiters); + + // Scan input and check for delimiters + size_t pos = start_pos; + size_t last_valid_pos = start_pos; + + while (pos < ctx.input.size()) { + auto utf8_result = common_parse_utf8_codepoint(ctx.input, pos); + + if (utf8_result.status == utf8_parse_result::INCOMPLETE) { + // Incomplete UTF-8 sequence + if (!ctx.is_lenient()) { + // Input is complete but UTF-8 is incomplete = malformed + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); + } + // Return what we have so far (before incomplete sequence) + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, last_valid_pos); + } + + if (utf8_result.status == utf8_parse_result::INVALID) { + // Malformed UTF-8 + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); + } + + // Check if a delimiter starts at this position + auto match = matcher.check_at(ctx.input, pos); + + if (match == trie::COMPLETE_MATCH) { + // Found a complete delimiter, return everything before it + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); + } + + if (match == trie::PARTIAL_MATCH) { + // Found a partial match extending to end of input, return everything before it + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); + } + + pos += utf8_result.bytes_consumed; + last_valid_pos = pos; + } + + if (last_valid_pos == ctx.input.size() && ctx.is_lenient()) { + // Reached the end of a partial stream, there might still be more input that we need to consume. + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, last_valid_pos); + } + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, last_valid_pos); + } + + common_peg_parse_result operator()(const common_peg_schema_parser & p) { + return arena.parse(p.child, ctx, start_pos); + } + + common_peg_parse_result operator()(const common_peg_rule_parser & p) { + // Parse the child + auto result = arena.parse(p.child, ctx, start_pos); + + if (!result.fail()) { + std::string_view text; + if (result.start < ctx.input.size()) { + text = std::string_view(ctx.input).substr(result.start, result.end - result.start); + } + + auto node_id = ctx.ast.add_node( + p.name, + "", + result.start, + result.end, + text, + std::move(result.nodes), + result.need_more_input() + ); + + return common_peg_parse_result(result.type, result.start, result.end, { node_id }); + } + + return result; + } + + common_peg_parse_result operator()(const common_peg_tag_parser & p) { + // Parse the child + if (ctx.is_debug()) { + fprintf(stderr, "%sTAG: %s\n", debug_indent().c_str(), p.tag.c_str()); + } + auto result = arena.parse(p.child, ctx, start_pos); + + if (!result.fail()) { + std::string_view text; + if (result.start < ctx.input.size()) { + text = std::string_view(ctx.input).substr(result.start, result.end - result.start); + } + + auto node_id = ctx.ast.add_node( + "", + p.tag, + result.start, + result.end, + text, + std::move(result.nodes), + result.need_more_input() + ); + + return common_peg_parse_result(result.type, result.start, result.end, { node_id }); + } + + return result; + } + + common_peg_parse_result operator()(const common_peg_ref_parser & p) { + auto rule_id = arena.get_rule(p.name); + return arena.parse(rule_id, ctx, start_pos); + } + + common_peg_parse_result operator()(const common_peg_atomic_parser & p) { + auto result = arena.parse(p.child, ctx, start_pos); + if (result.need_more_input()) { + // Clear nodes so they don't propagate up. + result.nodes.clear(); + } + return result; + } + + common_peg_parse_result operator()(const common_peg_gbnf_parser & p) { + return arena.parse(p.child, ctx, start_pos); + } +}; + +common_peg_parse_result common_peg_arena::parse(common_peg_parse_context & ctx, size_t start) const { + if (root_ == COMMON_PEG_INVALID_PARSER_ID) { + throw std::runtime_error("No root parser set"); + } + return parse(root_, ctx, start); +} + +common_peg_parse_result common_peg_arena::parse(common_peg_parser_id id, common_peg_parse_context & ctx, size_t start) const { + // Execute parser + const auto & parser = parsers_.at(id); + parser_executor exec(*this, ctx, start); + return std::visit(exec, parser); +} + +common_peg_parser_id common_peg_arena::resolve_ref(common_peg_parser_id id) { + const auto & parser = parsers_.at(id); + if (auto ref = std::get_if(&parser)) { + return get_rule(ref->name); + } + return id; +} + +static void bfs_node(common_peg_ast_arena &arena, std::ostringstream & oss, const common_peg_ast_node & node, int indent) { + for (int i = 0; i < indent; i++) { + oss << " "; + } + oss << "NODE " << node.id; + if (!node.rule.empty()) { + oss << " (rule " << node.rule << ")"; + } + if (!node.tag.empty()) { + oss << " (tag " << node.tag << ")"; + } + oss << " ['" << node.text << "']\n"; + for (const auto child : node.children) { + bfs_node(arena, oss, arena.get(child), indent + 1); + } +} + +std::string common_peg_ast_arena::dump() { + std::ostringstream oss; + for (auto & node : nodes_) { + bfs_node(*this, oss, node, 0); + } + return oss.str(); +} + +void common_peg_arena::resolve_refs() { + // Walk through all parsers and replace refs with their corresponding rule IDs + for (auto & parser : parsers_) { + std::visit([this](auto & p) { + using T = std::decay_t; + + if constexpr (std::is_same_v) { + for (auto & child : p.children) { + child = resolve_ref(child); + } + } else if constexpr (std::is_same_v) { + for (auto & child : p.children) { + child = resolve_ref(child); + } + } else 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) { + p.child = resolve_ref(p.child); + } else if constexpr (std::is_same_v) { + p.child = resolve_ref(p.child); + } else if constexpr (std::is_same_v) { + p.child = resolve_ref(p.child); + } else 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 || + std::is_same_v || + std::is_same_v) { + // These rules do not have children + } else { + static_assert(is_always_false_v); + } + }, parser); + } + + // Also flatten root if it's a ref + if (root_ != COMMON_PEG_INVALID_PARSER_ID) { + root_ = resolve_ref(root_); + } +} + +std::string common_peg_arena::dump(common_peg_parser_id id) const { + std::unordered_set visited; + return dump_impl(id, visited); +} + +std::string common_peg_arena::dump_impl(common_peg_parser_id id, + std::unordered_set & visited) const { + // Check for cycles + if (visited.count(id)) { + return "[cycle]"; + } + visited.insert(id); + + const auto & parser = parsers_.at(id); + + return std::visit([this, &visited](const auto & p) -> std::string { + using T = std::decay_t; + + if constexpr (std::is_same_v) { + return "Epsilon"; + } else if constexpr (std::is_same_v) { + return "Start"; + } else if constexpr (std::is_same_v) { + return "End"; + } else if constexpr (std::is_same_v) { + return "Literal(" + p.literal + ")"; + } else if constexpr (std::is_same_v) { + std::vector parts; + for (const auto & child : p.children) { + parts.push_back(dump_impl(child, visited)); + } + return "Sequence(" + string_join(parts, ", ") + ")"; + } else if constexpr (std::is_same_v) { + std::vector parts; + for (const auto & child : p.children) { + parts.push_back(dump_impl(child, visited)); + } + return "Choice(" + string_join(parts, ", ") + ")"; + } else if constexpr (std::is_same_v) { + if (p.max_count == -1) { + return "Repetition(" + dump_impl(p.child, visited) + ", " + std::to_string(p.min_count) + + ", unbounded)"; + } + return "Repetition(" + dump_impl(p.child, visited) + ", " + std::to_string(p.min_count) + ", " + std::to_string(p.max_count) + ")"; + } else if constexpr (std::is_same_v) { + return "And(" + dump_impl(p.child, visited) + ")"; + } else if constexpr (std::is_same_v) { + return "Not(" + dump_impl(p.child, visited) + ")"; + } else if constexpr (std::is_same_v) { + return "Atomic(" + dump_impl(p.child, visited) + ")"; + } else if constexpr (std::is_same_v) { + return "Gbnf(" + p.grammar + ", " + dump_impl(p.child, visited) + ")"; + } else if constexpr (std::is_same_v) { + return "Any"; + } else if constexpr (std::is_same_v) { + return "Space"; + } else if constexpr (std::is_same_v) { + if (p.max_count == -1) { + return "CharRepeat(" + p.pattern + ", " + std::to_string(p.min_count) + ", unbounded)"; + } + return "CharRepeat(" + p.pattern + ", " + std::to_string(p.min_count) + ", " + std::to_string(p.max_count) + ")"; + } else if constexpr (std::is_same_v) { + return "String(" + std::string(1, p.delimiter) + ")"; + } else if constexpr (std::is_same_v) { + return "Until(" + string_join(p.delimiters, " | ") + ")"; + } else if constexpr (std::is_same_v) { + return "Schema(" + dump_impl(p.child, visited) + ", " + (p.schema ? p.schema->dump() : "null") + ")"; + } else if constexpr (std::is_same_v) { + return "Rule(" + p.name + ", " + dump_impl(p.child, visited) + ")"; + } else if constexpr (std::is_same_v) { + return "Ref(" + p.name + ")"; + } else if constexpr (std::is_same_v) { + return "Tag(" + p.tag + ", " + dump(p.child) + ")"; + } else if constexpr (std::is_same_v) { + return "Atomic(" + dump(p.child) + ")"; + } else { + return "Unknown"; + } + }, parser); +} + +common_peg_parser & common_peg_parser::operator=(const common_peg_parser & other) { + id_ = other.id_; + return *this; +} + +common_peg_parser & common_peg_parser::operator+=(const common_peg_parser & other) { + id_ = builder_.sequence({id_, other.id_}); + return *this; +} + +common_peg_parser & common_peg_parser::operator|=(const common_peg_parser & other) { + id_ = builder_.choice({id_, other.id_}); + return *this; +} + +common_peg_parser common_peg_parser::operator+(const common_peg_parser & other) const { + return builder_.sequence({id_, other.id_}); +} + +common_peg_parser common_peg_parser::operator|(const common_peg_parser & other) const { + return builder_.choice({id_, other.id_}); +} + +common_peg_parser common_peg_parser::operator<<(const common_peg_parser & other) const { + return builder_.sequence({id_, builder_.space(), other.id_}); +} + +common_peg_parser common_peg_parser::operator+(const char * str) const { + return *this + builder_.literal(str); +} + +common_peg_parser common_peg_parser::operator+(const std::string & str) const { + return *this + builder_.literal(str); +} + +common_peg_parser common_peg_parser::operator<<(const char * str) const { + return *this << builder_.literal(str); +} + +common_peg_parser common_peg_parser::operator<<(const std::string & str) const { + return *this << builder_.literal(str); +} + +common_peg_parser common_peg_parser::operator|(const char * str) const { + return *this | builder_.literal(str); +} + +common_peg_parser common_peg_parser::operator|(const std::string & str) const { + return *this | builder_.literal(str); +} + +common_peg_parser operator+(const char * str, const common_peg_parser & p) { + return p.builder().literal(str) + p; +} + +common_peg_parser operator+(const std::string & str, const common_peg_parser & p) { + return operator+(str.c_str(), p); +} + +common_peg_parser operator<<(const char * str, const common_peg_parser & p) { + return p.builder().literal(str) << p; +} + +common_peg_parser operator<<(const std::string & str, const common_peg_parser & p) { + return operator<<(str.c_str(), p); +} + +common_peg_parser operator|(const char * str, const common_peg_parser & p) { + return p.builder().literal(str) | p; +} + +common_peg_parser operator|(const std::string & str, const common_peg_parser & p) { + return operator|(str.c_str(), p); +} + +common_peg_parser_builder::common_peg_parser_builder() {} + +common_peg_parser common_peg_parser_builder::sequence(const std::vector & parsers) { + // Flatten nested sequences + std::vector flattened; + for (const auto & p : parsers) { + const auto & parser = arena_.get(p); + if (auto seq = std::get_if(&parser)) { + flattened.insert(flattened.end(), seq->children.begin(), seq->children.end()); + } else { + flattened.push_back(p); + } + } + return wrap(arena_.add_parser(common_peg_sequence_parser{flattened})); +} + +common_peg_parser common_peg_parser_builder::sequence(const std::vector & parsers) { + std::vector ids; + ids.reserve(parsers.size()); + for (const auto & p : parsers) { + ids.push_back(p.id()); + } + return sequence(ids); +} + +common_peg_parser common_peg_parser_builder::sequence(std::initializer_list parsers) { + std::vector ids; + ids.reserve(parsers.size()); + for (const auto & p : parsers) { + ids.push_back(p.id()); + } + return sequence(ids); +} + +common_peg_parser common_peg_parser_builder::choice(const std::vector & parsers) { + // Flatten nested choices + std::vector flattened; + for (const auto & p : parsers) { + const auto & parser = arena_.get(p); + if (auto choice = std::get_if(&parser)) { + flattened.insert(flattened.end(), choice->children.begin(), choice->children.end()); + } else { + flattened.push_back(p); + } + } + return wrap(arena_.add_parser(common_peg_choice_parser{flattened})); +} + +common_peg_parser common_peg_parser_builder::choice(const std::vector & parsers) { + std::vector ids; + ids.reserve(parsers.size()); + for (const auto & p : parsers) { + ids.push_back(p.id()); + } + return choice(ids); +} + +common_peg_parser common_peg_parser_builder::choice(std::initializer_list parsers) { + std::vector ids; + ids.reserve(parsers.size()); + for (const auto & p : parsers) { + ids.push_back(p.id()); + } + return choice(ids); +} + +common_peg_parser common_peg_parser_builder::chars(const std::string & classes, int min, int max) { + auto [ranges, negated] = parse_char_classes(classes); + return wrap(arena_.add_parser(common_peg_chars_parser{classes, ranges, negated, min, max})); +} + +common_peg_parser common_peg_parser_builder::schema(const common_peg_parser & p, const std::string & name, const nlohmann::ordered_json & schema, bool raw) { + return wrap(arena_.add_parser(common_peg_schema_parser{p.id(), name, std::make_shared(schema), raw})); +} + +common_peg_parser common_peg_parser_builder::rule(const std::string & name, const common_peg_parser & p, bool trigger) { + auto clean_name = common_rust_string_to_std( + llama_common_peg_rule_name_rust((const uint8_t *) name.data(), name.size())); + auto rule_id = arena_.add_parser(common_peg_rule_parser{clean_name, p.id(), trigger}); + arena_.add_rule(clean_name, rule_id); + return ref(clean_name); +} + +common_peg_parser common_peg_parser_builder::rule(const std::string & name, const std::function & builder_fn, bool trigger) { + auto clean_name = common_rust_string_to_std( + llama_common_peg_rule_name_rust((const uint8_t *) name.data(), name.size())); + if (arena_.has_rule(clean_name)) { + return ref(clean_name); + } + + // Create placeholder rule to allow recursive references + auto placeholder = any(); // Temporary placeholder + auto placeholder_rule_id = arena_.add_parser(common_peg_rule_parser{clean_name, placeholder.id(), trigger}); + arena_.add_rule(clean_name, placeholder_rule_id); + + // Build the actual parser + auto parser = builder_fn(); + + // Replace placeholder with actual rule + auto rule_id = arena_.add_parser(common_peg_rule_parser{clean_name, parser.id(), trigger}); + arena_.rules_[clean_name] = rule_id; + + return ref(clean_name); +} + +void common_peg_parser_builder::set_root(const common_peg_parser & p) { + arena_.set_root(p.id()); +} + +common_peg_arena common_peg_parser_builder::build() { + arena_.resolve_refs(); + return std::move(arena_); +} + +// String primitives + +common_peg_parser common_peg_parser_builder::string_content(char delimiter) { + return wrap(arena_.add_parser(common_peg_string_parser{delimiter})); +} + +common_peg_parser common_peg_parser_builder::double_quoted_string() { + return rule("double-quoted-string", [this]() { + return sequence({literal("\""), string_content('"'), literal("\""), space()}); + }); +} + +common_peg_parser common_peg_parser_builder::single_quoted_string() { + return rule("single-quoted-string", [this]() { + return sequence({literal("'"), string_content('\''), literal("'"), space()}); + }); +} + +common_peg_parser common_peg_parser_builder::quoted_string() { + return rule("quoted-string", [this]() { + return choice({double_quoted_string(), single_quoted_string()}); + }); +} + +// JSON parsers + +common_peg_parser common_peg_parser_builder::json_number() { + return rule("json-number", [this]() { + auto digit1_9 = chars("[1-9]", 1, 1); + auto digits = chars("[0-9]"); + auto int_part = choice({literal("0"), sequence({digit1_9, chars("[0-9]", 0, -1)})}); + auto frac = sequence({literal("."), digits}); + auto exp = sequence({choice({literal("e"), literal("E")}), optional(chars("[+-]", 1, 1)), digits}); + // Negative lookahead: only commit the number when the next character can't extend it. + // At EOF in partial mode, chars returns NEED_MORE → negate propagates NEED_MORE → number not committed. + // This prevents premature commits of partial numbers (e.g. "3" when "3.14" is incoming). + auto not_number_continuation = negate(chars("[0-9.eE+-]", 1, 1)); + return sequence({ optional(literal("-")), int_part, optional(frac), optional(exp), not_number_continuation, space() }); + }); +} + +common_peg_parser common_peg_parser_builder::json_string() { + return rule("json-string", [this]() { + return sequence({literal("\""), string_content('"'), literal("\""), space()}); + }); +} + +common_peg_parser common_peg_parser_builder::json_bool() { + return rule("json-bool", [this]() { + return sequence({choice({literal("true"), literal("false")}), space()}); + }); +} + +common_peg_parser common_peg_parser_builder::json_null() { + return rule("json-null", [this]() { + return sequence({literal("null"), space()}); + }); +} + +common_peg_parser common_peg_parser_builder::json_object() { + return rule("json-object", [this]() { + auto ws = space(); + auto member = sequence({json_string(), ws, literal(":"), ws, json()}); + auto members = sequence({member, zero_or_more(sequence({ws, literal(","), ws, member}))}); + return sequence({ + literal("{"), + ws, + choice({ + literal("}"), + sequence({members, ws, literal("}")}) + }), + ws + }); + }); +} + +common_peg_parser common_peg_parser_builder::json_array() { + return rule("json-array", [this]() { + auto ws = space(); + auto elements = sequence({json(), zero_or_more(sequence({literal(","), ws, json()}))}); + return sequence({ + literal("["), + ws, + choice({ + literal("]"), + sequence({elements, ws, literal("]")}) + }), + ws + }); + }); +} + +common_peg_parser common_peg_parser_builder::json() { + return rule("json-value", [this]() { + return choice({ + json_object(), + json_array(), + json_string(), + json_number(), + json_bool(), + json_null() + }); + }); +} + +common_peg_parser common_peg_parser_builder::python_string() { + return rule("python-string", [this]() { + return choice({double_quoted_string(), single_quoted_string()}); + }); +} + +common_peg_parser common_peg_parser_builder::python_number() { + return json_number(); +} + +common_peg_parser common_peg_parser_builder::python_bool() { + return rule("python-bool", [this]() { + return sequence({ + choice({literal("True"), literal("False")}), + space() + }); + }); +} + +common_peg_parser common_peg_parser_builder::python_null() { + return rule("python-none", [this]() { + return sequence({literal("None"), space()}); + }); +} + +common_peg_parser common_peg_parser_builder::python_dict() { + return rule("python-dict", [this]() { + auto ws = space(); + auto member = sequence({python_string(), ws, literal(":"), ws, python_value()}); + auto members = sequence({member, zero_or_more(sequence({ws, literal(","), ws, member}))}); + return sequence({ + literal("{"), + ws, + choice({ + literal("}"), + sequence({members, ws, literal("}")}) + }), + ws + }); + }); +} + +common_peg_parser common_peg_parser_builder::python_array() { + return rule("python-array", [this]() { + auto ws = space(); + auto elements = sequence({python_value(), zero_or_more(sequence({literal(","), ws, python_value()}))}); + return sequence({ + literal("["), + ws, + choice({ + literal("]"), + sequence({elements, ws, literal("]")}) + }), + ws + }); + }); +} + +common_peg_parser common_peg_parser_builder::python_value() { + return rule("python-value", [this]() { + return choice({ + python_dict(), + python_array(), + python_string(), + python_number(), + python_bool(), + python_null() + }); + }); +} + +common_peg_parser common_peg_parser_builder::marker() { + auto sharp_bracket_parser = literal("<") + until(">") + literal(">"); + auto square_bracket_parser = literal("[") + until("]") + literal("]"); + return choice({ sharp_bracket_parser, square_bracket_parser }); +} + +common_peg_parser common_peg_parser_builder::json_member(const std::string & key, const common_peg_parser & p) { + auto ws = space(); + return sequence({ + literal("\"" + key + "\""), + ws, + literal(":"), + ws, + p, + }); +} + +static std::string gbnf_excluding_pattern(const std::vector & strings) { + trie matcher(strings); + auto pieces = matcher.collect_prefix_and_next(); + + std::string pattern; + for (size_t i = 0; i < pieces.size(); ++i) { + if (i > 0) { + pattern += " | "; + } + + const auto & pre = pieces[i].prefix; + const auto & chars = pieces[i].next_chars; + + std::string cls; + cls.reserve(chars.size()); + for (uint32_t ch : chars) { + cls += common_rust_string_to_std(llama_common_gbnf_escape_char_class_rust(ch)); + } + + if (!pre.empty()) { + pattern += gbnf_format_literal(common_unicode_cpts_to_utf8(pre)) + " [^" + cls + "]"; + } else { + pattern += "[^" + cls + "]"; + } + } + + return "(" + pattern + ")*"; +} + +static std::unordered_set collect_reachable_rules( + const common_peg_arena & arena, + const common_peg_parser_id & rule +) { + std::unordered_set reachable; + std::unordered_set visited; + + std::function visit = [&](common_peg_parser_id id) { + const auto & parser = arena.get(id); + + std::visit([&](const auto & p) { + using T = std::decay_t; + + 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 || + std::is_same_v) { + // These parsers do not have any children + } else if constexpr (std::is_same_v) { + for (auto child : p.children) { + visit(child); + } + } else if constexpr (std::is_same_v) { + for (auto child : p.children) { + visit(child); + } + } else 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) { + visit(p.child); + } else if constexpr (std::is_same_v) { + if (visited.find(p.name) == visited.end()) { + visited.insert(p.name); + reachable.insert(p.name); + visit(p.child); + } + } else if constexpr (std::is_same_v) { + // Traverse rules so we pick up everything + auto referenced_rule = arena.get_rule(p.name); + visit(referenced_rule); + } else { + static_assert(is_always_false_v); + } + }, parser); + }; + + visit(rule); + return reachable; +} + +// GBNF generation implementation +void common_peg_arena::build_grammar(const common_grammar_builder & builder, bool lazy) const { + auto schema_delegates = [](const common_peg_schema_parser & s) -> bool { + if (!s.schema) { + return true; + } + if (s.raw && s.schema->contains("type")) { + const auto & type_val = s.schema->at("type"); + if (type_val.is_string() && type_val == "string") { + return true; + } + // Handle nullable types like ["string", "null"] - delegate when the + // non-null type is string, since the tagged format uses raw text + if (type_val.is_array()) { + for (const auto & t : type_val) { + if (t.is_string() && t.get() != "null") { + return t.get() == "string"; + } + } + } + } + // Delegate for enum schemas in raw mode - enum values are literal strings + if (s.raw && !s.schema->contains("type") && s.schema->contains("enum")) { + return true; + } + return false; + }; + + // Unwrap the parser so we can properly check if it's a sequence or choice + auto effective_parser = [&](common_peg_parser_id id) -> const common_peg_parser_variant & { + while (true) { + const auto & p = parsers_.at(id); + if (const auto * tag = std::get_if(&p)) { + id = tag->child; + } else if (const auto * atomic = std::get_if(&p)) { + id = atomic->child; + } else if (const auto * schema = std::get_if(&p)) { + if (schema_delegates(*schema)) { + id = schema->child; + } else { + return p; + } + } else { + return p; + } + } + }; + + // Generate GBNF for a parser + std::function to_gbnf = [&](common_peg_parser_id id) -> std::string { + const auto & parser = parsers_.at(id); + + return std::visit([&](const auto & p) -> std::string { + using T = std::decay_t; + + if constexpr (std::is_same_v || + std::is_same_v || + std::is_same_v) { + return ""; + } else if constexpr (std::is_same_v) { + return gbnf_format_literal(p.literal); + } else if constexpr (std::is_same_v) { + std::string s; + for (const auto & child : p.children) { + auto child_gbnf = to_gbnf(child); + if (child_gbnf.empty()) { + continue; + } + if (!s.empty()) { + s += " "; + } + const auto & child_parser = effective_parser(child); + if (std::holds_alternative(child_parser) || + std::holds_alternative(child_parser)) { + s += "(" + child_gbnf + ")"; + } else { + s += child_gbnf; + } + } + return s; + } else if constexpr (std::is_same_v) { + std::string s; + for (const auto & child : p.children) { + if (!s.empty()) { + s += " | "; + } + auto child_gbnf = to_gbnf(child); + const auto & child_parser = effective_parser(child); + if (std::holds_alternative(child_parser)) { + s += "(" + child_gbnf + ")"; + } else { + s += child_gbnf; + } + } + return s; + } else if constexpr (std::is_same_v) { + auto child_gbnf = to_gbnf(p.child); + const auto & child_parser = effective_parser(p.child); + if (std::holds_alternative(child_parser) || + std::holds_alternative(child_parser)) { + child_gbnf = "(" + child_gbnf + ")"; + } + if (p.min_count == 0 && p.max_count == 1) { + return child_gbnf + "?"; + } + if (p.min_count == 0 && p.max_count == -1) { + return child_gbnf + "*"; + } + if (p.min_count == 1 && p.max_count == -1) { + return child_gbnf + "+"; + } + if (p.max_count == -1) { + return child_gbnf + "{" + std::to_string(p.min_count) + ",}"; + } + if (p.min_count == p.max_count) { + if (p.min_count == 1) { + return child_gbnf; + } + return child_gbnf + "{" + std::to_string(p.min_count) + "}"; + } + return child_gbnf + "{" + std::to_string(p.min_count) + "," + std::to_string(p.max_count) + "}"; + } else if constexpr (std::is_same_v || std::is_same_v) { + return ""; // Lookahead not supported in GBNF + } else if constexpr (std::is_same_v) { + return "."; + } else if constexpr (std::is_same_v) { + return "space"; + } else if constexpr (std::is_same_v) { + std::string result = p.pattern; + if (p.min_count == 0 && p.max_count == 1) { + return result + "?"; + } + if (p.min_count == 0 && p.max_count == -1) { + return result + "*"; + } + if (p.min_count == 1 && p.max_count == -1) { + return result + "+"; + } + if (p.max_count == -1) { + return result + "{" + std::to_string(p.min_count) + ",}"; + } + if (p.min_count == p.max_count) { + if (p.min_count == 1) { + return result; + } + return result + "{" + std::to_string(p.min_count) + "}"; + } + return result + "{" + std::to_string(p.min_count) + "," + std::to_string(p.max_count) + "}"; + } else if constexpr (std::is_same_v) { + const std::string delim(1, p.delimiter); + return R"(( [^)" + delim + R"(\\] | "\\" ( [)" + delim + R"(\\/ bfnrt] | "u" [0-9a-fA-F]{4} ) )*)"; + } else if constexpr (std::is_same_v) { + if (p.delimiters.empty()) { + return ".*"; + } + return gbnf_excluding_pattern(p.delimiters); + } else if constexpr (std::is_same_v) { + if (schema_delegates(p)) { + return to_gbnf(p.child); + } + return builder.add_schema(p.name, *p.schema); + } else if constexpr (std::is_same_v) { + return p.name; + } else if constexpr (std::is_same_v) { + // Refs should not exist after flattening, but kept just in case + return p.name; + } else if constexpr (std::is_same_v) { + return to_gbnf(p.child); + } else if constexpr (std::is_same_v) { + return to_gbnf(p.child); + } else if constexpr (std::is_same_v) { + return p.grammar; + } else { + static_assert(is_always_false_v); + } + }, parser); + }; + + // Collect reachable rules + std::unordered_set reachable_rules; + + if (lazy) { + // Collect rules reachable from trigger rules + for (const auto & [name, id] : rules_) { + const auto & parser = parsers_.at(id); + if (auto rule = std::get_if(&parser)) { + if (rule->trigger) { + // Mark trigger as reachable and visit it + reachable_rules.insert(name); + auto add_rules = collect_reachable_rules(*this, id); + reachable_rules.insert(add_rules.begin(), add_rules.end()); + } + } + } + } else { + // Collect rules reachable from root + reachable_rules = collect_reachable_rules(*this, root_); + } + + // Create GBNF rules for all reachable rules + for (const auto & [name, rule_id] : rules_) { + if (reachable_rules.find(name) == reachable_rules.end()) { + continue; + } + + const auto & parser = parsers_.at(rule_id); + if (auto rule = std::get_if(&parser)) { + builder.add_rule(rule->name, to_gbnf(rule->child)); + } + } + + if (lazy) { + // Generate root rule from trigger rules only + std::vector trigger_names; + for (const auto & [name, rule_id] : rules_) { + const auto & parser = parsers_.at(rule_id); + if (auto rule = std::get_if(&parser)) { + if (rule->trigger) { + trigger_names.push_back(rule->name); + } + } + } + + // Sort for predictable order + std::sort(trigger_names.begin(), trigger_names.end()); + builder.add_rule("root", string_join(trigger_names, " | ")); + } else if (root_ != COMMON_PEG_INVALID_PARSER_ID) { + builder.add_rule("root", to_gbnf(root_)); + } +} + +static nlohmann::json serialize_parser_variant(const common_peg_parser_variant & variant) { + using json = nlohmann::json; + + return std::visit([](const auto & p) -> json { + using T = std::decay_t; + + if constexpr (std::is_same_v) { + return json{{"type", "epsilon"}}; + } else if constexpr (std::is_same_v) { + return json{{"type", "start"}}; + } else if constexpr (std::is_same_v) { + return json{{"type", "end"}}; + } else if constexpr (std::is_same_v) { + return json{{"type", "literal"}, {"literal", p.literal}}; + } else if constexpr (std::is_same_v) { + return json{{"type", "sequence"}, {"children", p.children}}; + } else if constexpr (std::is_same_v) { + return json{{"type", "choice"}, {"children", p.children}}; + } else if constexpr (std::is_same_v) { + return json{ + {"type", "repetition"}, + {"child", p.child}, + {"min_count", p.min_count}, + {"max_count", p.max_count} + }; + } else if constexpr (std::is_same_v) { + return json{{"type", "and"}, {"child", p.child}}; + } else if constexpr (std::is_same_v) { + return json{{"type", "not"}, {"child", p.child}}; + } else if constexpr (std::is_same_v) { + return json{{"type", "any"}}; + } else if constexpr (std::is_same_v) { + return json{{"type", "space"}}; + } else if constexpr (std::is_same_v) { + json ranges = json::array(); + for (const auto & range : p.ranges) { + ranges.push_back({{"start", range.start}, {"end", range.end}}); + } + return json{ + {"type", "chars"}, + {"pattern", p.pattern}, + {"ranges", ranges}, + {"negated", p.negated}, + {"min_count", p.min_count}, + {"max_count", p.max_count} + }; + } else if constexpr (std::is_same_v) { + return json{{"type", "string"}, {"delimiter", std::string(1, p.delimiter)}}; + } else if constexpr (std::is_same_v) { + return json{{"type", "until"}, {"delimiters", p.delimiters}}; + } else if constexpr (std::is_same_v) { + return json{ + {"type", "schema"}, + {"child", p.child}, + {"name", p.name}, + {"schema", p.schema ? *p.schema : nullptr}, + {"raw", p.raw} + }; + } else if constexpr (std::is_same_v) { + return json{ + {"type", "rule"}, + {"name", p.name}, + {"child", p.child}, + {"trigger", p.trigger} + }; + } else if constexpr (std::is_same_v) { + return json{{"type", "ref"}, {"name", p.name}}; + } else if constexpr (std::is_same_v) { + return json{{"type", "atomic"}, {"child", p.child}}; + } else if constexpr (std::is_same_v) { + return json{ + {"type", "tag"}, + {"child", p.child}, + {"tag", p.tag} + }; + } else if constexpr (std::is_same_v) { + return json{{"type", "gbnf"}, {"child", p.child}, {"grammar", p.grammar}}; + } + }, variant); +} + +nlohmann::json common_peg_arena::to_json() const { + auto parsers = nlohmann::json::array(); + for (const auto & parser : parsers_) { + parsers.push_back(serialize_parser_variant(parser)); + } + return nlohmann::json{ + {"parsers", parsers}, + {"rules", rules_}, + {"root", root_} + }; +} + +static common_peg_parser_variant deserialize_parser_variant(const nlohmann::json & j) { + if (!j.contains("type") || !j["type"].is_string()) { + throw std::runtime_error("Parser variant JSON missing or invalid 'type' field"); + } + + std::string type = j["type"]; + + if (type == "epsilon") { + return common_peg_epsilon_parser{}; + } + if (type == "start") { + return common_peg_start_parser{}; + } + if (type == "end") { + return common_peg_end_parser{}; + } + if (type == "literal") { + if (!j.contains("literal") || !j["literal"].is_string()) { + throw std::runtime_error("literal parser missing or invalid 'literal' field"); + } + return common_peg_literal_parser{j["literal"]}; + } + if (type == "sequence") { + if (!j.contains("children") || !j["children"].is_array()) { + throw std::runtime_error("sequence parser missing or invalid 'children' field"); + } + return common_peg_sequence_parser{j["children"].get>()}; + } + if (type == "choice") { + if (!j.contains("children") || !j["children"].is_array()) { + throw std::runtime_error("choice parser missing or invalid 'children' field"); + } + return common_peg_choice_parser{j["children"].get>()}; + } + if (type == "repetition") { + if (!j.contains("child") || !j.contains("min_count") || !j.contains("max_count")) { + throw std::runtime_error("repetition parser missing required fields"); + } + return common_peg_repetition_parser{ + j["child"].get(), + j["min_count"].get(), + j["max_count"].get() + }; + } + if (type == "and") { + if (!j.contains("child")) { + throw std::runtime_error("and parser missing 'child' field"); + } + return common_peg_and_parser{j["child"].get()}; + } + if (type == "not") { + if (!j.contains("child")) { + throw std::runtime_error("not parser missing 'child' field"); + } + return common_peg_not_parser{j["child"].get()}; + } + if (type == "any") { + return common_peg_any_parser{}; + } + if (type == "space") { + return common_peg_space_parser{}; + } + if (type == "chars") { + if (!j.contains("pattern") || !j.contains("ranges") || !j.contains("negated") || + !j.contains("min_count") || !j.contains("max_count")) { + throw std::runtime_error("chars parser missing required fields"); + } + common_peg_chars_parser parser; + parser.pattern = j["pattern"]; + parser.negated = j["negated"]; + parser.min_count = j["min_count"]; + parser.max_count = j["max_count"]; + for (const auto & range_json : j["ranges"]) { + if (!range_json.contains("start") || !range_json.contains("end")) { + throw std::runtime_error("char_range missing 'start' or 'end' field"); + } + parser.ranges.push_back({ + range_json["start"].get(), + range_json["end"].get() + }); + } + return parser; + } + if (type == "string") { + if (!j.contains("delimiter")) { + throw std::runtime_error("string parser missing delimiter field."); + } + std::string delimiter = j["delimiter"]; + if (delimiter.empty()) { + throw std::runtime_error("string parser delimiter is empty."); + } + return common_peg_string_parser{delimiter[0]}; + } + if (type == "until") { + if (!j.contains("delimiters") || !j["delimiters"].is_array()) { + throw std::runtime_error("until parser missing or invalid 'delimiters' field"); + } + return common_peg_until_parser{j["delimiters"].get>()}; + } + if (type == "schema") { + if (!j.contains("child") || !j.contains("name") || !j.contains("schema") || !j.contains("raw")) { + throw std::runtime_error("schema parser missing required fields"); + } + common_peg_schema_parser parser; + parser.child = j["child"].get(); + parser.name = j["name"]; + if (!j["schema"].is_null()) { + parser.schema = std::make_shared(j["schema"]); + } + parser.raw = j["raw"].get(); + return parser; + } + if (type == "rule") { + if (!j.contains("name") || !j.contains("child") || !j.contains("trigger")) { + throw std::runtime_error("rule parser missing required fields"); + } + return common_peg_rule_parser{ + j["name"].get(), + j["child"].get(), + j["trigger"].get() + }; + } + if (type == "ref") { + if (!j.contains("name") || !j["name"].is_string()) { + throw std::runtime_error("ref parser missing or invalid 'name' field"); + } + return common_peg_ref_parser{j["name"]}; + } + if (type == "atomic") { + if (!j.contains("child")) { + throw std::runtime_error("tag parser missing required fields"); + } + return common_peg_atomic_parser{ + j["child"].get(), + }; + } + if (type == "tag") { + if (!j.contains("child") || !j.contains("tag")) { + throw std::runtime_error("tag parser missing required fields"); + } + return common_peg_tag_parser{ + j["child"].get(), + j["tag"].get(), + }; + } + + if (type == "gbnf") { + if (!j.contains("child") || !j.contains("grammar")) { + throw std::runtime_error("gbnf parser missing required fields"); + } + return common_peg_gbnf_parser{ + j["child"].get(), + j["grammar"].get(), + }; + } + + throw std::runtime_error("Unknown parser type: " + type); +} + +common_peg_arena common_peg_arena::from_json(const nlohmann::json & j) { + if (!j.contains("parsers") || !j["parsers"].is_array()) { + throw std::runtime_error("JSON missing or invalid 'parsers' array"); + } + if (!j.contains("rules") || !j["rules"].is_object()) { + throw std::runtime_error("JSON missing or invalid 'rules' object"); + } + if (!j.contains("root")) { + throw std::runtime_error("JSON missing 'root' field"); + } + + common_peg_arena arena; + + const auto & parsers_json = j["parsers"]; + arena.parsers_.reserve(parsers_json.size()); + for (const auto & parser_json : parsers_json) { + arena.parsers_.push_back(deserialize_parser_variant(parser_json)); + } + + arena.rules_ = j["rules"].get>(); + + for (const auto & [name, id] : arena.rules_) { + if (id >= arena.parsers_.size()) { + throw std::runtime_error("Rule '" + name + "' references invalid parser ID: " + std::to_string(id)); + } + } + + arena.root_ = j["root"].get(); + if (arena.root_ != COMMON_PEG_INVALID_PARSER_ID && arena.root_ >= arena.parsers_.size()) { + throw std::runtime_error("Root references invalid parser ID: " + std::to_string(arena.root_)); + } + + return arena; +} + +std::string common_peg_arena::save() const { + return to_json().dump(); +} + +void common_peg_arena::load(const std::string & data) { + *this = from_json(nlohmann::json::parse(data)); +} + +common_peg_arena build_peg_parser(const std::function & fn) { + common_peg_parser_builder builder; + builder.set_root(fn(builder)); + return builder.build(); +} + +#include "common.h.inc" +#include "ggml.h.inc" + +#include + +using ordered_json = nlohmann::ordered_json; + +static std::string_view trim_trailing_space(std::string_view sv, int max = -1) { + common_rust_string_view view = + llama_common_trim_trailing_space_view_rust((const uint8_t *) sv.data(), sv.size(), max); + return std::string_view((const char *) view.data, view.len); +} + +static std::string_view trim_leading_space(std::string_view sv, int max = -1) { + common_rust_string_view view = + llama_common_trim_leading_space_view_rust((const uint8_t *) sv.data(), sv.size(), max); + return std::string_view((const char *) view.data, view.len); +} + +static std::string_view trim(std::string_view sv) { + return trim_trailing_space(trim_leading_space(sv, 1)); +} + +// JSON-escape a string and return the inner content (without surrounding quotes). +static std::string escape_json_string_inner(const std::string & s) { + std::string escaped = ordered_json(s).dump(); + if (escaped.size() >= 2 && escaped.front() == '"' && escaped.back() == '"') { + return escaped.substr(1, escaped.size() - 2); + } + return escaped; +} + +void tag_based_peg_mapper::from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result) { + arena.visit(result, [this](const common_peg_ast_node & node) { + if (!node.tag.empty()) { + tags[node.tag] = std::string(node.text); + } + }); +} + +tagged_parse_result tagged_peg_parser::parse_and_extract(const std::string & input, common_peg_parse_flags extra_flags) const { + common_peg_parse_context ctx(input, flags | extra_flags); + auto parse_result = arena.parse(ctx); + + tag_based_peg_mapper mapper; + mapper.from_ast(ctx.ast, parse_result); + + return { std::move(parse_result), std::move(mapper.tags) }; +} + +tagged_parse_result tagged_peg_parser::parse_anywhere_and_extract(const std::string & input) const { + if (input.empty()) { + return parse_and_extract(input); + } + for (size_t i = 0; i < input.size(); i++) { + common_peg_parse_context ctx(input, flags); + auto parse_result = arena.parse(ctx, i); + if (parse_result.success() || i == input.size() - 1) { + tag_based_peg_mapper mapper; + mapper.from_ast(ctx.ast, parse_result); + return { std::move(parse_result), std::move(mapper.tags) }; + } + } + GGML_ABORT("Should not happen"); +} + +tagged_peg_parser build_tagged_peg_parser( + const std::function & fn) { + common_peg_parser_builder builder; + builder.set_root(fn(builder)); + return { builder.build() }; +} + +common_peg_parser common_chat_peg_builder::tag_with_safe_content(const std::string & tag_name, + const std::string & marker, + const common_peg_parser & p) { + if (marker.empty()) { + return zero_or_more(choice({ p, rule(tag_name, content(any())) })); + } + auto content_chunk = rule(tag_name, content(negate(literal(marker)) + any() + until(marker))); + return zero_or_more(choice({ p, content_chunk })); +} + +std::string & common_chat_peg_mapper::args_target() { + return (current_tool && !current_tool->name.empty()) ? current_tool->arguments : args_buffer; +} + +std::string common_chat_peg_mapper::normalize_container_value(const std::string & input) { + return common_rust_string_to_std( + llama_common_normalize_quotes_to_json_rust((const uint8_t *) input.data(), input.size())); +} + +void common_chat_peg_mapper::from_ast(const common_peg_ast_arena & arena, + const common_peg_parse_result & parse_result_arg) { + arena.visit(parse_result_arg, [this](const common_peg_ast_node & node) { map(node); }); + // Flush any pending tool call that was started but never got a name + // This happens during partial parsing when the tool call is incomplete + if (pending_tool_call.has_value() && !pending_tool_call->name.empty()) { + if (!args_buffer.empty()) { + pending_tool_call->arguments = args_buffer; + } + if (closing_quote_pending && !pending_tool_call->arguments.empty()) { + pending_tool_call->arguments += "\""; + } + result.tool_calls.push_back(pending_tool_call.value()); + pending_tool_call.reset(); + } + + // Discard whitespace-only reasoning content (e.g. from prefill) + if (!result.reasoning_content.empty()) { + bool all_whitespace = true; + for (char c : result.reasoning_content) { + if (c != ' ' && c != '\n' && c != '\r' && c != '\t') { + all_whitespace = false; + break; + } + } + if (all_whitespace) { + result.reasoning_content.clear(); + } + } +} + +void common_chat_peg_mapper::map(const common_peg_ast_node & node) { + // Handle reasoning/content tags + bool is_reasoning = node.tag == common_chat_peg_builder::REASONING; + bool is_content = node.tag == common_chat_peg_builder::CONTENT; + + if (is_reasoning) { // GPT OSS can have more than 1 reasoning block, so concatenate here + result.reasoning_content += std::string(node.text); + } + + if (is_content) { + // Concatenate content from multiple content nodes (e.g., when reasoning markers + // are preserved before content markers in reasoning_format=NONE mode) + result.content += std::string(node.text); + } + + // Handle tool-related tags (supporting both JSON and tagged formats) + bool is_tool_open = node.tag == common_chat_peg_builder::TOOL_OPEN; + bool is_tool_close = node.tag == common_chat_peg_builder::TOOL_CLOSE; + bool is_tool_name = node.tag == common_chat_peg_builder::TOOL_NAME; + bool is_tool_id = node.tag == common_chat_peg_builder::TOOL_ID; + bool is_tool_args = node.tag == common_chat_peg_builder::TOOL_ARGS; + bool is_arg_open = node.tag == common_chat_peg_builder::TOOL_ARG_OPEN; + bool is_arg_close = node.tag == common_chat_peg_builder::TOOL_ARG_CLOSE; + bool is_arg_name = node.tag == common_chat_peg_builder::TOOL_ARG_NAME; + bool is_arg_value = node.tag == common_chat_peg_builder::TOOL_ARG_VALUE; + bool is_arg_string_value = node.tag == common_chat_peg_builder::TOOL_ARG_STRING_VALUE; + + if (is_tool_open) { + pending_tool_call = common_chat_tool_call(); + current_tool = &pending_tool_call.value(); + arg_count = 0; + args_buffer.clear(); + closing_quote_pending = false; + } + + if (is_tool_id && current_tool) { + auto text = trim_trailing_space(node.text); + if (text.size() >= 2 && text.front() == '"' && text.back() == '"') { + text = text.substr(1, text.size() - 2); + } + current_tool->id = std::string(text); + } + + if (is_tool_name && current_tool) { + current_tool->name = std::string(trim_trailing_space(node.text)); + // Now that we have the name, populate the arguments from the buffer + if (!args_buffer.empty()) { + current_tool->arguments = args_buffer; + args_buffer.clear(); + } else if (current_tool->arguments.empty()) { + current_tool->arguments = "{"; + } + // Add the tool call to results so streaming can see it + if (pending_tool_call.has_value()) { + result.tool_calls.push_back(pending_tool_call.value()); + pending_tool_call.reset(); + current_tool = &result.tool_calls.back(); + } + } + + if (is_tool_args && current_tool) { + // For JSON format: arguments come as a complete JSON object + // For tagged format: built up from individual arg_name/arg_value nodes + auto text = trim_trailing_space(node.text); + if (!text.empty() && text.front() == '{') { + args_target() = std::string(text); + } + } + + if (is_arg_open) { + closing_quote_pending = false; + } + + if (is_arg_name && current_tool) { + std::string arg_entry; + if (arg_count > 0) { + arg_entry = ","; + } + arg_entry += ordered_json(trim(node.text)).dump() + ":"; + ++arg_count; + + auto & target = args_target(); + if (target.empty()) { + target = "{"; + } + target += arg_entry; + } + + if ((is_arg_value || is_arg_string_value) && current_tool) { + std::string value_content = std::string(trim_trailing_space(trim_leading_space(node.text, 1), 1)); + + std::string value_to_add; + if (value_content.empty() && is_arg_string_value) { + // Empty string value - arg_close will add the closing quote + value_to_add = "\""; + closing_quote_pending = true; + } else if (!value_content.empty() && is_arg_string_value) { + // Schema declares this as string type - always treat as literal string value + if (!closing_quote_pending) { + value_to_add = "\""; + closing_quote_pending = true; + } + value_to_add += escape_json_string_inner(value_content); + } else if (!value_content.empty()) { + // For potential containers, normalize Python-style single quotes to JSON double quotes + bool is_potential_container = value_content[0] == '[' || value_content[0] == '{'; + if (is_potential_container) { + value_content = normalize_container_value(value_content); + } + + // Try to parse as JSON value (number, bool, null, object, array) + try { + ordered_json parsed = ordered_json::parse(value_content); + if (parsed.is_string()) { + // Don't add closing quote yet (added by arg_close) for monotonic streaming + std::string escaped = parsed.dump(); + if (!escaped.empty() && escaped.back() == '"') { + escaped.pop_back(); + } + value_to_add = escaped; + closing_quote_pending = true; + } else { + // Non-string values: use raw content to preserve whitespace for monotonicity + value_to_add = value_content; + } + } catch (...) { + if (node.is_partial && is_potential_container) { + // Partial container: pass through the already-normalized content + value_to_add = value_content; + } else { + // Not valid JSON - treat as string value + if (!closing_quote_pending) { + value_to_add = "\""; + closing_quote_pending = true; + } + value_to_add += escape_json_string_inner(value_content); + } + } + } + + args_target() += value_to_add; + } + + if (is_arg_close && current_tool) { + if (closing_quote_pending) { + args_target() += "\""; + closing_quote_pending = false; + } + } + + if (is_tool_close && current_tool) { + // Flush buffer to arguments if tool name was never seen + if (current_tool->name.empty() && !args_buffer.empty()) { + current_tool->arguments = args_buffer; + args_buffer.clear(); + } + // Close any pending string quote + if (closing_quote_pending) { + current_tool->arguments += "\""; + closing_quote_pending = false; + } + // Close any unclosed braces (accounts for nested objects) + for (int d = llama_common_json_brace_depth_rust( + (const uint8_t *) current_tool->arguments.data(), current_tool->arguments.size()); + d > 0; d--) { + current_tool->arguments += "}"; + } + // Add tool call to results if named; otherwise discard + if (pending_tool_call.has_value()) { + if (!current_tool->name.empty()) { + result.tool_calls.push_back(pending_tool_call.value()); + } + pending_tool_call.reset(); + } + } +} + +common_peg_parser common_chat_peg_builder::standard_constructed_tools( + const std::map & markers, + const ordered_json & tools, + bool parallel_tool_calls, + bool force_tool_calls) { + if (!tools.is_array() || tools.empty()) { + return eps(); + } + + // Extract markers with defaults + auto get_marker = [&markers](const std::string & key, const std::string & default_val = "") -> std::string { + auto it = markers.find(key); + return it != markers.end() ? it->second : default_val; + }; + + std::string section_start = get_marker("tool_call_start_marker", ""); + std::string section_end = get_marker("tool_call_end_marker", ""); + std::string func_opener = get_marker("function_opener", ""); + std::string func_closer = get_marker("function_closer", ""); + std::string param_key_prefix = get_marker("parameter_key_prefix", ""); + std::string param_closer = get_marker("parameter_closer", ""); + + // Build tool choices for tagged format + auto tool_choices = choice(); + + for (const auto & tool_def : tools) { + if (!tool_def.contains("function")) { + continue; + } + const auto & function = tool_def.at("function"); + std::string name = function.at("name"); + ordered_json params = function.contains("parameters") ? function.at("parameters") : ordered_json::object(); + + // Build argument parsers + auto args = eps(); + if (params.contains("properties") && !params["properties"].empty()) { + auto arg_choice = choice(); + for (const auto & el : params["properties"].items()) { + const std::string & prop_name = el.key(); + + auto arg_name_parser = + choice({ literal(prop_name), literal("\"" + prop_name + "\""), literal("'" + prop_name + "'") }); + + auto arg_rule = tool_arg(tool_arg_open(literal(param_key_prefix)) + tool_arg_name(arg_name_parser) + + literal(param_key_suffix) + tool_arg_value(until(param_closer)) + + tool_arg_close(literal(param_closer))); + arg_choice |= arg_rule; + } + args = zero_or_more(arg_choice + space()); + } + + // Build function parser: args + auto tool_parser = tool(tool_open(literal(func_opener) + tool_name(literal(name)) + literal(func_name_suffix)) + + space() + tool_args(args) + space() + tool_close(literal(func_closer))); + + tool_choices |= rule("tool-" + name, tool_parser); + } + + // Build the section with markers + auto section = + parallel_tool_calls ? + trigger_rule("tool-call", literal(section_start) + space() + one_or_more(tool_choices + space()) + + literal(section_end)) : + trigger_rule("tool-call", literal(section_start) + space() + tool_choices + space() + literal(section_end)); + + return force_tool_calls ? section : optional(section); +} + +// Python-style tool calls: name(arg1="value1", arg2=123) +// Used only by LFM2 for now, so we don't merge it into autoparser +common_peg_parser common_chat_peg_builder::python_style_tool_calls( + const ordered_json & tools, + bool parallel_tool_calls) { + if (!tools.is_array() || tools.empty()) { + return eps(); + } + + auto tool_choices = choice(); + + for (const auto & tool_def : tools) { + if (!tool_def.contains("function")) { + continue; + } + const auto & function = tool_def.at("function"); + std::string name = function.at("name"); + ordered_json params = function.contains("parameters") ? function.at("parameters") : ordered_json::object(); + + auto args = eps(); + if (params.contains("properties") && !params["properties"].empty()) { + auto arg_choice = choice(); + for (const auto & el : params["properties"].items()) { + const std::string & prop_name = el.key(); + const auto & prop_def = el.value(); + bool is_string_type = (prop_def.contains("type") && prop_def["type"] == "string"); + + auto arg_name_parser = literal(prop_name); + + common_peg_parser arg_value_parser = eps(); + auto string_value_parser = choice({ + literal("\"") + tool_arg_string_value(string_content('"')) + literal("\""), + literal("'") + tool_arg_string_value(string_content('\'')) + literal("'") + }); + + if (is_string_type) { + arg_value_parser = string_value_parser; + } else { + arg_value_parser = tool_arg_value(python_value()); + } + + // Full argument: name="value" or name=value + auto arg_rule = tool_arg( + tool_arg_open(eps()) + + tool_arg_name(arg_name_parser) + + literal("=") + + arg_value_parser + + tool_arg_close(eps()) + ); + arg_choice |= arg_rule; + } + + args = arg_choice + zero_or_more("," + space() + arg_choice); + } + + auto tool_parser = tool(tool_open(tool_name(literal(name)) + literal("(")) + + space() + tool_args(args) + space() + tool_close(literal(")")) + ); + + tool_choices |= rule("tool-" + name, tool_parser); + } + + if (parallel_tool_calls) { + return "[" + space() + tool_choices + zero_or_more("," + space() + tool_choices) + space() + "]"; + } + return "[" + space() + tool_choices + space() + "]"; +} + +// Helper: Parse dot notation key into prefix and field name +static std::pair parse_key_spec(const std::string & key) { + auto dot_pos = key.find('.'); + if (dot_pos == std::string::npos) { + return {"", key}; // Top-level field + } + return {key.substr(0, dot_pos), key.substr(dot_pos + 1)}; +} + +// Mode 1: function_is_key — parse {"function_name": {...}} +common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key( + const ordered_json & tools, + const std::string & args_key, + const std::string & effective_args_key, + const std::string & call_id_key, + const std::string & gen_call_id_key) { + + auto tool_choices = choice(); + + for (const auto & tool_def : tools) { + if (!tool_def.contains("function")) { + continue; + } + const auto & function = tool_def.at("function"); + std::string name = function.at("name"); + ordered_json params = function.contains("parameters") ? function.at("parameters") : ordered_json::object(); + + // Build inner object fields + std::vector inner_fields; + + if (!call_id_key.empty()) { + auto id_parser = atomic( + literal("\"" + call_id_key + "\"") + space() + literal(":") + space() + + literal("\"") + tool_id(string_content('"')) + literal("\"") + ); + inner_fields.push_back(optional(id_parser + space() + optional(literal(",") + space()))); + } + + if (!gen_call_id_key.empty()) { + auto gen_id_parser = atomic( + literal("\"" + gen_call_id_key + "\"") + space() + literal(":") + space() + + choice({ + literal("\"") + tool_id(string_content('"')) + literal("\""), + tool_id(json_number()) + }) + ); + inner_fields.push_back(optional(gen_id_parser + space() + optional(literal(",") + space()))); + } + + // Arguments — either wrapped in args_key or parsed directly + common_peg_parser args_parser = eps(); + if (args_key.empty()) { + args_parser = tool_args(schema(json(), "tool-" + name + "-schema", params)); + } else { + args_parser = literal("\"" + effective_args_key + "\"") + space() + literal(":") + space() + + tool_args(schema(json(), "tool-" + name + "-schema", params)); + } + inner_fields.push_back(args_parser); + + // Build inner object parser + common_peg_parser inner_object = eps(); + if (args_key.empty() && inner_fields.size() == 1) { + inner_object = inner_fields[0]; + } else { + inner_object = literal("{") + space(); + for (size_t i = 0; i < inner_fields.size(); i++) { + inner_object = inner_object + inner_fields[i]; + if (i < inner_fields.size() - 1) { + inner_object = inner_object + space(); + } + } + inner_object = inner_object + space() + literal("}"); + } + + auto tool_parser = tool( + tool_open(literal("{")) + space() + + literal("\"") + tool_name(literal(name)) + literal("\"") + + space() + literal(":") + space() + + inner_object + + space() + tool_close(literal("}")) + ); + + tool_choices |= rule("tool-" + name, tool_parser); + } + + return tool_choices; +} + +// Mode 2: Nested keys (dot notation like "function.name") +common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys( + const ordered_json & tools, + const std::string & effective_name_key, + const std::string & effective_args_key, + const std::string & call_id_key, + const std::string & gen_call_id_key) { + + auto tool_choices = choice(); + + auto name_spec = parse_key_spec(effective_name_key); + auto args_spec = parse_key_spec(effective_args_key); + + std::string nested_prefix = !name_spec.first.empty() ? name_spec.first : args_spec.first; + std::string nested_name_field = !name_spec.first.empty() ? name_spec.second : effective_name_key; + std::string nested_args_field = !args_spec.first.empty() ? args_spec.second : effective_args_key; + + for (const auto & tool_def : tools) { + if (!tool_def.contains("function")) { + continue; + } + const auto & function = tool_def.at("function"); + std::string name = function.at("name"); + ordered_json params = function.contains("parameters") ? function.at("parameters") : ordered_json::object(); + + auto nested_name = literal("\"" + nested_name_field + "\"") + space() + literal(":") + space() + + atomic(literal("\"") + tool_name(literal(name)) + literal("\"")); + auto nested_args = literal("\"" + nested_args_field + "\"") + space() + literal(":") + space() + + tool_args(schema(json(), "tool-" + name + "-schema", params)); + + auto nested_object = literal("{") + space() + + nested_name + space() + literal(",") + space() + + nested_args + + space() + literal("}"); + + // Format: { id?, "function": {...} } + auto tool_parser_body = tool_open(literal("{")) + space(); + + if (!call_id_key.empty()) { + auto id_spec = parse_key_spec(call_id_key); + if (id_spec.first.empty()) { + auto id_parser = atomic( + literal("\"" + call_id_key + "\"") + space() + literal(":") + space() + + literal("\"") + tool_id(string_content('"')) + literal("\"") + ); + tool_parser_body = tool_parser_body + optional(id_parser + space() + literal(",") + space()); + } + } + + if (!gen_call_id_key.empty()) { + auto gen_id_spec = parse_key_spec(gen_call_id_key); + if (gen_id_spec.first.empty()) { + auto gen_id_parser = atomic( + literal("\"" + gen_call_id_key + "\"") + space() + literal(":") + space() + + choice({ + literal("\"") + tool_id(string_content('"')) + literal("\""), + tool_id(json_number()) + }) + ); + tool_parser_body = tool_parser_body + optional(gen_id_parser + space() + literal(",") + space()); + } + } + + auto nested_field = literal("\"" + nested_prefix + "\"") + space() + literal(":") + space() + nested_object; + tool_parser_body = tool_parser_body + nested_field + space() + tool_close(literal("}")); + + tool_choices |= rule("tool-" + name, tool(tool_parser_body)); + } + + return tool_choices; +} + +// Mode 3: Flat keys with optional ID fields and parameter ordering +common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys( + const ordered_json & tools, + const std::string & effective_name_key, + const std::string & effective_args_key, + const std::string & call_id_key, + const std::string & gen_call_id_key, + const std::vector & parameters_order) { + + auto tool_choices = choice(); + auto name_key_parser = literal("\"" + effective_name_key + "\""); + auto args_key_parser = literal("\"" + effective_args_key + "\""); + + for (const auto & tool_def : tools) { + if (!tool_def.contains("function")) { + continue; + } + const auto & function = tool_def.at("function"); + std::string name = function.at("name"); + ordered_json params = function.contains("parameters") ? function.at("parameters") : ordered_json::object(); + + auto tool_name_ = name_key_parser + space() + literal(":") + space() + + atomic(literal("\"") + tool_name(literal(name)) + literal("\"")); + auto tool_args_ = args_key_parser + space() + literal(":") + space() + + tool_args(schema(json(), "tool-" + name + "-schema", params)); + + // Build ID parsers if keys are provided + common_peg_parser id_parser = eps(); + if (!call_id_key.empty()) { + id_parser = atomic( + literal("\"" + call_id_key + "\"") + space() + literal(":") + space() + + choice({ + literal("\"") + tool_id(string_content('"')) + literal("\""), + tool_id(json_number()) + }) + ); + } + + common_peg_parser gen_id_parser = eps(); + if (!gen_call_id_key.empty()) { + gen_id_parser = atomic( + literal("\"" + gen_call_id_key + "\"") + space() + literal(":") + space() + + choice({ + literal("\"") + tool_id(string_content('"')) + literal("\""), + tool_id(json_number()) + }) + ); + } + + // Create (parser, key) pairs for all fields, then sort by parameters_order + std::vector> parser_pairs; + parser_pairs.emplace_back(tool_name_, effective_name_key); + parser_pairs.emplace_back(tool_args_, effective_args_key); + if (!call_id_key.empty()) { + parser_pairs.emplace_back(optional(id_parser), call_id_key); + } + if (!gen_call_id_key.empty()) { + parser_pairs.emplace_back(optional(gen_id_parser), gen_call_id_key); + } + + std::sort(parser_pairs.begin(), parser_pairs.end(), + [¶meters_order](const auto & a, const auto & b) { + auto pos_a = std::find(parameters_order.begin(), parameters_order.end(), a.second); + auto pos_b = std::find(parameters_order.begin(), parameters_order.end(), b.second); + size_t idx_a = (pos_a == parameters_order.end()) ? parameters_order.size() : std::distance(parameters_order.begin(), pos_a); + size_t idx_b = (pos_b == parameters_order.end()) ? parameters_order.size() : std::distance(parameters_order.begin(), pos_b); + return idx_a < idx_b; + }); + + auto ordered_body = tool_open(literal("{")) + space(); + for (size_t i = 0; i < parser_pairs.size(); i++) { + ordered_body = ordered_body + parser_pairs[i].first; + if (i < parser_pairs.size() - 1) { + ordered_body = ordered_body + space() + literal(",") + space(); + } + } + ordered_body = ordered_body + space() + tool_close(literal("}")); + + tool_choices |= rule("tool-" + name, tool(ordered_body)); + } + + return tool_choices; +} + +common_peg_parser common_chat_peg_builder::prefix(const std::string & s, const std::string & delimiter) { + if (s.empty()) { + return eps(); + } + if (delimiter.empty()) { + return literal(s); + } + return literal(s.substr(0, s.rfind(delimiter))); +} + +common_peg_parser common_chat_peg_builder::standard_json_tools( + const std::string & section_start, + const std::string & section_end, + const ordered_json & tools, + bool parallel_tool_calls, + bool force_tool_calls, + const std::string & name_key, + const std::string & args_key, + bool array_wrapped, + bool function_is_key, + const std::string & call_id_key, + const std::string & gen_call_id_key, + const std::vector & parameters_order) { + if (!tools.is_array() || tools.empty()) { + return eps(); + } + + std::string effective_name_key = name_key.empty() ? "name" : name_key; + std::string effective_args_key = args_key.empty() ? "arguments" : args_key; + + // Dispatch to the appropriate builder based on the JSON layout mode + common_peg_parser tool_choices = eps(); + if (function_is_key) { + tool_choices = build_json_tools_function_is_key(tools, args_key, effective_args_key, call_id_key, gen_call_id_key); + } else { + auto name_spec = parse_key_spec(effective_name_key); + auto args_spec = parse_key_spec(effective_args_key); + if (!name_spec.first.empty() || !args_spec.first.empty()) { + tool_choices = build_json_tools_nested_keys(tools, effective_name_key, effective_args_key, call_id_key, gen_call_id_key); + } else { + tool_choices = build_json_tools_flat_keys(tools, effective_name_key, effective_args_key, call_id_key, gen_call_id_key, parameters_order); + } + } + + // Build the section with markers + auto tool_calls = tool_choices; + if (parallel_tool_calls) { + tool_calls = tool_calls + zero_or_more(space() + literal(",") + space() + tool_choices); + } + + if (array_wrapped) { + tool_calls = literal("[") + space() + tool_calls + space() + literal("]"); + } + + auto section = + trigger_rule("tool-call", literal(section_start) + space() + tool_calls + space() + literal(section_end)); + + return force_tool_calls ? section : optional(section); +} + +void common_chat_peg_gemma4_mapper::from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result) { + for (const auto & node : result.nodes) { + visit(arena, node); + } +} + +static std::string gemma4_to_json(const common_peg_ast_arena & arena, common_peg_ast_id id) { + const auto & node = arena.get(id); + + if (node.text.empty()) { + return ""; + } + + if (node.rule == "gemma4-number" || node.rule == "gemma4-bool" || node.rule == "gemma4-null") { + return std::string(node.text); + } + + if (node.rule == "gemma4-string-content") { + return escape_json_string_inner(std::string(node.text)); + } + + if (node.rule == "gemma4-string") { + std::string result = "\""; + if (!node.children.empty()) { + result += gemma4_to_json(arena, node.children[0]); + if (!node.is_partial) { + result += "\""; + } + } + return result; + } + + if (node.rule == "gemma4-array") { + std::string result = "["; + + bool add_comma = false; + for (auto child_id : node.children) { + if (add_comma) { + result += ','; + } + add_comma = true; + result += gemma4_to_json(arena, child_id); + } + + if (!node.is_partial) { + result += ']'; + } + return result; + } + + if (node.rule == "gemma4-dict-key-name") { + return std::string(node.text); + } + + if (node.rule == "gemma4-dict-key") { + std::string result = "\""; + if (!node.children.empty()) { + result += escape_json_string_inner(gemma4_to_json(arena, node.children[0])); + } + if (!node.is_partial) { + result += "\":"; + } + return result; + } + + if (node.rule == "gemma4-dict-kv") { + std::string result; + for (auto child_id : node.children) { + result += gemma4_to_json(arena, child_id); + } + return result; + } + + if (node.rule == "gemma4-dict") { + std::string result = "{"; + + bool add_comma = false; + for (auto child_id : node.children) { + if (add_comma) { + result += ','; + } + add_comma = true; + result += gemma4_to_json(arena, child_id); + } + + if (!node.is_partial) { + result += '}'; + } + return result; + } + + if (node.rule == "gemma4-value") { + if (!node.children.empty()) { + return gemma4_to_json(arena, node.children[0]); + } + return ""; + } + + return ""; +} + +void common_chat_peg_gemma4_mapper::visit(const common_peg_ast_arena & arena, common_peg_ast_id id) { + const auto & node = arena.get(id); + + if (node.tag == "reasoning") { + result.reasoning_content += std::string(node.text); + return; + } + + if (node.tag == "content") { + result.content += std::string(node.text); + return; + } + + if (node.tag == "tool") { + auto name_id = arena.find_by_tag(node, "tool-name"); + auto args_id = arena.find_by_tag(node, "tool-args"); + + if (name_id != COMMON_PEG_INVALID_AST_ID && args_id != COMMON_PEG_INVALID_AST_ID) { + const auto & name_node = arena.get(name_id); + const auto & args_node = arena.get(args_id); + + if (!name_node.is_partial) { + common_chat_tool_call call; + call.name = std::string(name_node.text); + if (!args_node.children.empty()) { + call.arguments = gemma4_to_json(arena, args_node.children[0]); + } + result.tool_calls.push_back(call); + } + } + + return; + } + + for (auto child_id : node.children) { + visit(arena, child_id); + } +} + +#include "common.h.inc" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using json = nlohmann::ordered_json; + +static void build_min_max_int(int64_t min_value, int64_t max_value, std::stringstream & out, int decimals_left = 16, bool top_level = true) { + auto has_min = min_value != std::numeric_limits::min(); + auto has_max = max_value != std::numeric_limits::max(); + + auto digit_range = [&](char from, char to) { + out << "["; + if (from == to) { + out << from; + } else { + out << from << "-" << to; + } + out << "]"; + }; + auto more_digits = [&](int min_digits, int max_digits) { + out << "[0-9]"; + if (min_digits == max_digits && min_digits == 1) { + return; + } + out << "{"; + out << min_digits; + if (max_digits != min_digits) { + out << ","; + if (max_digits != std::numeric_limits::max()) { + out << max_digits; + } + } + out << "}"; + }; + std::function uniform_range = + [&](const std::string_view & from, const std::string_view & to) { + size_t i = 0; + while (i < from.length() && i < to.length() && from[i] == to[i]) { + i++; + } + if (i > 0) { + out << "\"" << from.substr(0, i) << "\""; + } + if (i < from.length() && i < to.length()) { + if (i > 0) { + out << " "; + } + auto sub_len = from.length() - i - 1; + if (sub_len > 0) { + auto from_sub = from.substr(i + 1); + auto to_sub = to.substr(i + 1); + auto sub_zeros = string_repeat("0", sub_len); + auto sub_nines = string_repeat("9", sub_len); + + auto to_reached = false; + out << "("; + if (from_sub == sub_zeros) { + digit_range(from[i], to[i] - 1); + out << " "; + more_digits(sub_len, sub_len); + } else { + out << "[" << from[i] << "] "; + out << "("; + uniform_range(from_sub, sub_nines); + out << ")"; + if (from[i] < to[i] - 1) { + out << " | "; + if (to_sub == sub_nines) { + digit_range(from[i] + 1, to[i]); + to_reached = true; + } else { + digit_range(from[i] + 1, to[i] - 1); + } + out << " "; + more_digits(sub_len, sub_len); + } + } + if (!to_reached) { + out << " | "; + digit_range(to[i], to[i]); + out << " "; + uniform_range(sub_zeros, to_sub); + } + out << ")"; + } else { + out << "[" << from[i] << "-" << to[i] << "]"; + } + } + }; + + if (has_min && has_max) { + if (min_value < 0 && max_value < 0) { + out << "\"-\" ("; + build_min_max_int(-max_value, -min_value, out, decimals_left, /* top_level= */ true); + out << ")"; + return; + } + + if (min_value < 0) { + out << "\"-\" ("; + build_min_max_int(0, -min_value, out, decimals_left, /* top_level= */ true); + out << ") | "; + min_value = 0; + } + + auto min_s = std::to_string(min_value); + auto max_s = std::to_string(max_value); + auto min_digits = min_s.length(); + auto max_digits = max_s.length(); + + for (auto digits = min_digits; digits < max_digits; digits++) { + uniform_range(min_s, string_repeat("9", digits)); + min_s = "1" + string_repeat("0", digits); + out << " | "; + } + uniform_range(min_s, max_s); + return; + } + + auto less_decimals = std::max(decimals_left - 1, 1); + + if (has_min) { + if (min_value < 0) { + out << "\"-\" ("; + build_min_max_int(std::numeric_limits::min(), -min_value, out, decimals_left, /* top_level= */ false); + out << ") | [0] | [1-9] "; + more_digits(0, decimals_left - 1); + } else if (min_value == 0) { + if (top_level) { + out << "[0] | [1-9] "; + more_digits(0, less_decimals); + } else { + more_digits(1, decimals_left); + } + } else if (min_value <= 9) { + char c = '0' + min_value; + auto range_start = top_level ? '1' : '0'; + if (c > range_start) { + digit_range(range_start, c - 1); + out << " "; + more_digits(1, less_decimals); + out << " | "; + } + digit_range(c, '9'); + out << " "; + more_digits(0, less_decimals); + } else { + auto min_s = std::to_string(min_value); + auto len = min_s.length(); + auto c = min_s[0]; + + if (c > '1') { + digit_range(top_level ? '1' : '0', c - 1); + out << " "; + more_digits(len, less_decimals); + out << " | "; + } + digit_range(c, c); + out << " ("; + build_min_max_int(std::stoll(min_s.substr(1)), std::numeric_limits::max(), out, less_decimals, /* top_level= */ false); + out << ")"; + if (c < '9') { + out << " | "; + digit_range(c + 1, '9'); + out << " "; + more_digits(len - 1, less_decimals); + } + } + return; + } + + if (has_max) { + if (max_value >= 0) { + if (top_level) { + out << "\"-\" [1-9] "; + more_digits(0, less_decimals); + out << " | "; + } + build_min_max_int(0, max_value, out, decimals_left, /* top_level= */ true); + } else { + out << "\"-\" ("; + build_min_max_int(-max_value, std::numeric_limits::max(), out, decimals_left, /* top_level= */ false); + out << ")"; + } + return; + } + + throw std::runtime_error("At least one of min_value or max_value must be set"); +} + +const std::string SPACE_RULE = "| \" \" | \"\\n\"{1,2} [ \\t]{0,20}"; + +struct BuiltinRule { + std::string content; + std::vector deps; +}; + +static std::unordered_map PRIMITIVE_RULES = { + {"boolean", {"(\"true\" | \"false\") space", {}}}, + {"decimal-part", {"[0-9]{1,16}", {}}}, + {"integral-part", {"[0] | [1-9] [0-9]{0,15}", {}}}, + {"number", {"(\"-\"? integral-part) (\".\" decimal-part)? ([eE] [-+]? integral-part)? space", {"integral-part", "decimal-part"}}}, + {"integer", {"(\"-\"? integral-part) space", {"integral-part"}}}, + {"value", {"object | array | string | number | boolean | null", {"object", "array", "string", "number", "boolean", "null"}}}, + {"object", {"\"{\" space ( string \":\" space value (\",\" space string \":\" space value)* )? \"}\" space", {"string", "value"}}}, + {"array", {"\"[\" space ( value (\",\" space value)* )? \"]\" space", {"value"}}}, + {"uuid", {"\"\\\"\" [0-9a-fA-F]{8} \"-\" [0-9a-fA-F]{4} \"-\" [0-9a-fA-F]{4} \"-\" [0-9a-fA-F]{4} \"-\" [0-9a-fA-F]{12} \"\\\"\" space", {}}}, + {"char", {"[^\"\\\\\\x7F\\x00-\\x1F] | [\\\\] ([\"\\\\bfnrt] | \"u\" [0-9a-fA-F]{4})", {}}}, + {"string", {"\"\\\"\" char* \"\\\"\" space", {"char"}}}, + {"null", {"\"null\" space", {}}}, +}; + +static std::unordered_map STRING_FORMAT_RULES = { + {"date", {"[0-9]{4} \"-\" ( \"0\" [1-9] | \"1\" [0-2] ) \"-\" ( \"0\" [1-9] | [1-2] [0-9] | \"3\" [0-1] )", {}}}, + {"time", {"([01] [0-9] | \"2\" [0-3]) \":\" [0-5] [0-9] \":\" [0-5] [0-9] ( \".\" [0-9]{3} )? ( \"Z\" | ( \"+\" | \"-\" ) ( [01] [0-9] | \"2\" [0-3] ) \":\" [0-5] [0-9] )", {}}}, + {"date-time", {"date \"T\" time", {"date", "time"}}}, + {"date-string", {"\"\\\"\" date \"\\\"\" space", {"date"}}}, + {"time-string", {"\"\\\"\" time \"\\\"\" space", {"time"}}}, + {"date-time-string", {"\"\\\"\" date-time \"\\\"\" space", {"date-time"}}} +}; + +static std::regex INVALID_RULE_CHARS_RE("[^a-zA-Z0-9-]+"); + +static std::unordered_set NON_LITERAL_SET = {'|', '.', '(', ')', '[', ']', '{', '}', '*', '+', '?'}; +static std::unordered_set ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS = {'^', '$', '.', '[', ']', '(', ')', '|', '{', '}', '*', '+', '?'}; + +std::string gbnf_format_literal(const std::string & literal) { + return common_rust_string_to_std(llama_common_gbnf_format_literal_rust( + (const uint8_t *) literal.data(), literal.size())); +} + +class common_schema_converter { +private: + friend class common_schema_info; + friend std::string build_grammar(const std::function & cb, const common_grammar_options & options); + std::function _fetch_json; + bool _dotall; + std::map _rules; + std::unordered_map _refs; + std::unordered_set _refs_being_resolved; + std::vector _errors; + std::vector _warnings; + + std::string _add_rule(const std::string & name, const std::string & rule) { + std::string esc_name = regex_replace(name, INVALID_RULE_CHARS_RE, "-"); + if (_rules.find(esc_name) == _rules.end() || _rules[esc_name] == rule) { + _rules[esc_name] = rule; + return esc_name; + } + int i = 0; + while (_rules.find(esc_name + std::to_string(i)) != _rules.end() && _rules[esc_name + std::to_string(i)] != rule) { + i++; + } + std::string key = esc_name + std::to_string(i); + _rules[key] = rule; + return key; + } + + std::string _generate_union_rule(const std::string & name, const std::vector & alt_schemas) { + std::vector rules; + rules.reserve(alt_schemas.size()); + for (size_t i = 0; i < alt_schemas.size(); i++) { + rules.push_back(visit(alt_schemas[i], name + (name.empty() ? "alternative-" : "-") + std::to_string(i))); + } + return string_join(rules, " | "); + } + + std::string _visit_pattern(const std::string & pattern, const std::string & name) { + if (!(pattern.front() == '^' && pattern.back() == '$')) { + _errors.push_back("Pattern must start with '^' and end with '$'"); + return ""; + } + std::string sub_pattern = pattern.substr(1, pattern.length() - 2); + std::unordered_map sub_rule_ids; + + size_t i = 0; + size_t length = sub_pattern.length(); + + using literal_or_rule = std::pair; + auto to_rule = [&](const literal_or_rule & ls) { + auto is_literal = ls.second; + auto s = ls.first; + return is_literal ? "\"" + s + "\"" : s; + }; + std::function transform = [&]() -> literal_or_rule { + size_t start = i; + std::vector seq; + + auto get_dot = [&]() { + std::string rule; + if (_dotall) { + rule = "[\\U00000000-\\U0010FFFF]"; + } else { + rule = "[^\\x0A\\x0D]"; + } + return _add_rule("dot", rule); + }; + + // Joins the sequence, merging consecutive literals together. + auto join_seq = [&]() { + std::vector ret; + + std::string literal; + auto flush_literal = [&]() { + if (literal.empty()) { + return false; + } + ret.emplace_back(literal, true); + literal.clear(); + return true; + }; + + for (const auto & item : seq) { + auto is_literal = item.second; + if (is_literal) { + literal += item.first; + } else { + flush_literal(); + ret.push_back(item); + } + } + flush_literal(); + + std::vector results; + results.reserve(ret.size()); + for (const auto & item : ret) { + results.push_back(to_rule(item)); + } + return std::make_pair(string_join(results, " "), false); + }; + + while (i < length) { + char c = sub_pattern[i]; + if (c == '.') { + seq.emplace_back(get_dot(), false); + i++; + } else if (c == '(') { + i++; + if (i < length && sub_pattern[i] == '?') { + if (i + 1 < length && sub_pattern[i + 1] == ':') { + i += 2; // skip "?:" for non-capturing group, treat as regular group + } else { + // lookahead/lookbehind (?=, ?!, ?<=, ? 0) { + if (sub_pattern[i] == '\\' && i + 1 < length) { + i += 2; // skip escaped character + } else { + if (sub_pattern[i] == '(') depth++; + else if (sub_pattern[i] == ')') depth--; + i++; + } + } + continue; + } + } + seq.emplace_back("(" + to_rule(transform()) + ")", false); + } else if (c == ')') { + i++; + if (start > 0 && sub_pattern[start - 1] != '(' && (start < 2 || sub_pattern[start - 2] != '?' || sub_pattern[start - 1] != ':')) { + _errors.push_back("Unbalanced parentheses"); + } + return join_seq(); + } else if (c == '[') { + std::string square_brackets = std::string(1, c); + i++; + while (i < length && sub_pattern[i] != ']') { + if (sub_pattern[i] == '\\') { + square_brackets += sub_pattern.substr(i, 2); + i += 2; + } else { + square_brackets += sub_pattern[i]; + i++; + } + } + if (i >= length) { + _errors.push_back("Unbalanced square brackets"); + } + square_brackets += ']'; + i++; + seq.emplace_back(square_brackets, false); + } else if (c == '|') { + seq.emplace_back("|", false); + i++; + } else if (c == '*' || c == '+' || c == '?') { + seq.back() = std::make_pair(to_rule(seq.back()) + c, false); + i++; + } else if (c == '{') { + std::string curly_brackets = std::string(1, c); + i++; + while (i < length && sub_pattern[i] != '}') { + curly_brackets += sub_pattern[i]; + i++; + } + if (i >= length) { + _errors.push_back("Unbalanced curly brackets"); + } + curly_brackets += '}'; + i++; + auto nums = string_split(curly_brackets.substr(1, curly_brackets.length() - 2), ","); + int min_times = 0; + int max_times = std::numeric_limits::max(); + try { + if (nums.size() == 1) { + min_times = max_times = std::stoi(nums[0]); + } else if (nums.size() != 2) { + _errors.push_back("Wrong number of values in curly brackets"); + } else { + if (!nums[0].empty()) { + min_times = std::stoi(nums[0]); + } + if (!nums[1].empty()) { + max_times = std::stoi(nums[1]); + } + } + } catch (const std::invalid_argument & e) { + _errors.push_back("Invalid number in curly brackets"); + return std::make_pair("", false); + } + auto &last = seq.back(); + auto &sub = last.first; + auto sub_is_literal = last.second; + + if (!sub_is_literal) { + std::string & sub_id = sub_rule_ids[sub]; + if (sub_id.empty()) { + sub_id = _add_rule(name + "-" + std::to_string(sub_rule_ids.size()), sub); + } + sub = sub_id; + } + const std::string repeated = sub_is_literal ? "\"" + sub + "\"" : sub; + seq.back().first = common_rust_string_to_std(llama_common_gbnf_build_repetition_rust( + (const uint8_t *) repeated.data(), repeated.size(), + min_times, max_times, + nullptr, 0)); + seq.back().second = false; + } else { + std::string literal; + auto is_non_literal = [&](char c) { + return NON_LITERAL_SET.find(c) != NON_LITERAL_SET.end(); + }; + while (i < length) { + if (sub_pattern[i] == '\\' && i < length - 1) { + char next = sub_pattern[i + 1]; + if (ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS.find(next) != ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS.end()) { + i++; + literal += sub_pattern[i]; + i++; + } else { + literal += sub_pattern.substr(i, 2); + i += 2; + } + } else if (sub_pattern[i] == '"') { + literal += "\\\""; + i++; + } else if (!is_non_literal(sub_pattern[i]) && + (i == length - 1 || literal.empty() || sub_pattern[i + 1] == '.' || !is_non_literal(sub_pattern[i + 1]))) { + literal += sub_pattern[i]; + i++; + } else { + break; + } + } + if (!literal.empty()) { + seq.emplace_back(literal, true); + } + } + } + return join_seq(); + }; + return _add_rule(name, "\"\\\"\" (" + to_rule(transform()) + ") \"\\\"\" space"); + } + + /* + Returns a rule that matches a JSON string that is none of the provided strings + + not_strings({"a"}) + -> ["] ( [a] char+ | [^"a] char* )? ["] space + not_strings({"and", "also"}) + -> ["] ( [a] ([l] ([s] ([o] char+ | [^"o] char*) | [^"s] char*) | [n] ([d] char+ | [^"d] char*) | [^"ln] char*) | [^"a] char* )? ["] space + */ + std::string _not_strings(const std::vector & strings) { + + struct TrieNode { + std::map children; + bool is_end_of_string; + + TrieNode() : is_end_of_string(false) {} + + void insert(const std::string & string) { + auto *node = this; + for (char c : string) { + node = &node->children[c]; + } + node->is_end_of_string = true; + } + }; + + TrieNode trie; + for (const auto & s : strings) { + trie.insert(s); + } + + std::string char_rule = _add_primitive("char", PRIMITIVE_RULES.at("char")); + std::ostringstream out; + out << "[\"] ( "; + std::function visit = [&](const TrieNode & node) { + std::ostringstream rejects; + auto first = true; + for (const auto & kv : node.children) { + rejects << kv.first; + if (first) { + first = false; + } else { + out << " | "; + } + out << "[" << kv.first << "]"; + if (!kv.second.children.empty()) { + out << " ("; + visit(kv.second); + out << ")"; + } else if (kv.second.is_end_of_string) { + out << " " << char_rule << "+"; + } + } + if (!node.children.empty()) { + if (!first) { + out << " | "; + } + out << "[^\"" << rejects.str() << "] " << char_rule << "*"; + } + }; + visit(trie); + + out << " )"; + if (!trie.is_end_of_string) { + out << "?"; + } + out << " [\"] space"; + return out.str(); + } + + std::string _resolve_ref(const std::string & ref) { + std::string ref_name = common_rust_string_to_std( + llama_common_gbnf_ref_name_rust((const uint8_t *) ref.data(), ref.size())); + if (_rules.find(ref_name) == _rules.end() && _refs_being_resolved.find(ref) == _refs_being_resolved.end()) { + _refs_being_resolved.insert(ref); + json resolved = _refs[ref]; + ref_name = visit(resolved, ref_name); + _refs_being_resolved.erase(ref); + } + return ref_name; + } + + std::string _build_object_rule( + const std::vector> & properties, + const std::unordered_set & required, + const std::string & name, + const json & additional_properties) + { + std::vector required_props; + std::vector optional_props; + std::unordered_map prop_kv_rule_names; + std::vector prop_names; + for (const auto & kv : properties) { + const auto &prop_name = kv.first; + const auto &prop_schema = kv.second; + + std::string prop_rule_name = visit(prop_schema, name + (name.empty() ? "" : "-") + prop_name); + prop_kv_rule_names[prop_name] = _add_rule( + name + (name.empty() ? "" : "-") + prop_name + "-kv", + gbnf_format_literal(json(prop_name).dump()) + " space \":\" space " + prop_rule_name + ); + if (required.find(prop_name) != required.end()) { + required_props.push_back(prop_name); + } else { + optional_props.push_back(prop_name); + } + prop_names.push_back(prop_name); + } + if ((additional_properties.is_boolean() && additional_properties.get()) || additional_properties.is_object()) { + std::string sub_name = name + (name.empty() ? "" : "-") + "additional"; + std::string value_rule = + additional_properties.is_object() ? visit(additional_properties, sub_name + "-value") + : _add_primitive("value", PRIMITIVE_RULES.at("value")); + + auto key_rule = + prop_names.empty() ? _add_primitive("string", PRIMITIVE_RULES.at("string")) + : _add_rule(sub_name + "-k", _not_strings(prop_names)); + std::string kv_rule = _add_rule(sub_name + "-kv", key_rule + " \":\" space " + value_rule); + prop_kv_rule_names["*"] = kv_rule; + optional_props.push_back("*"); + } + + std::string rule = "\"{\" space "; + for (size_t i = 0; i < required_props.size(); i++) { + if (i > 0) { + rule += " \",\" space "; + } + rule += prop_kv_rule_names[required_props[i]]; + } + + if (!optional_props.empty()) { + rule += " ("; + if (!required_props.empty()) { + rule += " \",\" space ( "; + } + + std::function &, bool)> get_recursive_refs = [&](const std::vector & ks, bool first_is_optional) { + std::string res; + if (ks.empty()) { + return res; + } + const std::string& k = ks[0]; + std::string kv_rule_name = prop_kv_rule_names[k]; + std::string comma_ref = "( \",\" space " + kv_rule_name + " )"; + if (first_is_optional) { + res = comma_ref + (k == "*" ? "*" : "?"); + } else { + res = kv_rule_name + (k == "*" ? " " + comma_ref + "*" : ""); + } + if (ks.size() > 1) { + res += " " + _add_rule( + name + (name.empty() ? "" : "-") + k + "-rest", + get_recursive_refs(std::vector(ks.begin() + 1, ks.end()), true) + ); + } + return res; + }; + + for (size_t i = 0; i < optional_props.size(); i++) { + if (i > 0) { + rule += " | "; + } + rule += get_recursive_refs(std::vector(optional_props.begin() + i, optional_props.end()), false); + } + if (!required_props.empty()) { + rule += " )"; + } + rule += " )?"; + } + + rule += " \"}\" space"; + + return rule; + } + + std::string _add_primitive(const std::string & name, const BuiltinRule & rule) { + auto n = _add_rule(name, rule.content); + for (const auto & dep : rule.deps) { + BuiltinRule dep_rule; + auto it = PRIMITIVE_RULES.find(dep); + if (it == PRIMITIVE_RULES.end()) { + it = STRING_FORMAT_RULES.find(dep); + if (it == STRING_FORMAT_RULES.end()) { + _errors.push_back("Rule " + dep + " not known"); + continue; + } + } + if (_rules.find(dep) == _rules.end()) { + _add_primitive(dep, it->second); + } + } + return n; + } + +public: + common_schema_converter( + const std::function & fetch_json, + bool dotall) + : _fetch_json(fetch_json), _dotall(dotall) + { + _rules["space"] = SPACE_RULE; + } + + void resolve_refs(json & schema, const std::string & url) { + /* + * Resolves all $ref fields in the given schema, fetching any remote schemas, + * replacing each $ref with absolute reference URL and populates _refs with the + * respective referenced (sub)schema dictionaries. + */ + std::function visit_refs = [&](json & n) { + if (n.is_array()) { + for (auto & x : n) { + visit_refs(x); + } + } else if (n.is_object()) { + if (n.contains("$ref")) { + std::string ref = n["$ref"]; + if (_refs.find(ref) == _refs.end()) { + json target; + if (ref.find("https://") == 0) { + std::string base_url = ref.substr(0, ref.find('#')); + auto it = _refs.find(base_url); + if (it != _refs.end()) { + target = it->second; + } else { + // Fetch the referenced schema and resolve its refs + auto referenced = _fetch_json(ref); + resolve_refs(referenced, base_url); + _refs[base_url] = referenced; + } + if (ref.find('#') == std::string::npos || ref.substr(ref.find('#') + 1).empty()) { + return; + } + } else if (ref.find("#/") == 0) { + target = schema; + n["$ref"] = url + ref; + ref = url + ref; + } else { + _errors.push_back("Unsupported ref: " + ref); + return; + } + std::string pointer = ref.substr(ref.find('#') + 1); + std::vector tokens = string_split(pointer, "/"); + for (size_t i = 1; i < tokens.size(); ++i) { + const std::string& sel = tokens[i]; + if (target.is_object() && target.contains(sel)) { + target = target[sel]; + } else if (target.is_array()) { + size_t sel_index; + try { + sel_index = std::stoull(sel); + } catch (const std::invalid_argument & e) { + sel_index = target.size(); + } + if (sel_index >= target.size()) { + _errors.push_back("Error resolving ref " + ref + ": " + sel + " not in " + target.dump()); + return; + } + target = target[sel_index]; + } else { + _errors.push_back("Error resolving ref " + ref + ": " + sel + " not in " + target.dump()); + return; + } + } + _refs[ref] = target; + } + } else { + for (const auto & kv : n.items()) { + visit_refs(kv.value()); + } + } + } + }; + + visit_refs(schema); + } + + static std::string _generate_constant_rule(const json & value) { + return gbnf_format_literal(value.dump()); + } + + std::string visit(const json & schema, const std::string & name) { + json schema_type = schema.contains("type") ? schema["type"] : json(); + std::string schema_format = schema.contains("format") ? schema["format"].get() : ""; + std::string rule_name = + llama_common_gbnf_is_reserved_name_rust((const uint8_t *) name.data(), name.size()) ? + name + "-" : name.empty() ? "root" : name; + + if (schema.contains("$ref")) { + return _add_rule(rule_name, _resolve_ref(schema["$ref"])); + } + if (schema.contains("oneOf") || schema.contains("anyOf")) { + std::vector alt_schemas = schema.contains("oneOf") ? schema["oneOf"].get>() : schema["anyOf"].get>(); + return _add_rule(rule_name, _generate_union_rule(name, alt_schemas)); + } + if (schema_type.is_array()) { + std::vector schema_types; + for (const auto & t : schema_type) { + json schema_copy(schema); + schema_copy["type"] = t; + schema_types.push_back(schema_copy); + } + return _add_rule(rule_name, _generate_union_rule(name, schema_types)); + } + if (schema.contains("const")) { + return _add_rule(rule_name, _generate_constant_rule(schema["const"]) + " space"); + } + if (schema.contains("enum")) { + std::vector enum_values; + for (const auto & v : schema["enum"]) { + enum_values.push_back(_generate_constant_rule(v)); + } + return _add_rule(rule_name, "(" + string_join(enum_values, " | ") + ") space"); + } + if ((schema_type.is_null() || schema_type == "object") + && (schema.contains("properties") || + (schema.contains("additionalProperties") && schema["additionalProperties"] != true))) { + std::unordered_set required; + if (schema.contains("required") && schema["required"].is_array()) { + for (const auto & item : schema["required"]) { + if (item.is_string()) { + required.insert(item.get()); + } + } + } + std::vector> properties; + if (schema.contains("properties")) { + for (const auto & prop : schema["properties"].items()) { + properties.emplace_back(prop.key(), prop.value()); + } + } + return _add_rule(rule_name, + _build_object_rule( + properties, required, name, + schema.contains("additionalProperties") ? schema["additionalProperties"] : json())); + } + if ((schema_type.is_null() || schema_type == "object" || schema_type == "string") && schema.contains("allOf")) { + std::unordered_set required; + std::vector> properties; + std::map enum_values; + const std::string& hybrid_name = name; + std::function add_component = [&](const json & comp_schema, bool is_required) { + if (comp_schema.contains("$ref")) { + add_component(_refs[comp_schema["$ref"]], is_required); + } else if (comp_schema.contains("properties")) { + for (const auto & prop : comp_schema["properties"].items()) { + properties.emplace_back(prop.key(), prop.value()); + if (is_required) { + required.insert(prop.key()); + } + } + } else if (comp_schema.contains("enum")) { + for (const auto & v : comp_schema["enum"]) { + const auto rule = _generate_constant_rule(v); + if (enum_values.find(rule) == enum_values.end()) { + enum_values[rule] = 0; + } + enum_values[rule] += 1; + } + } else { + // todo warning + } + }; + for (const auto & t : schema["allOf"]) { + if (t.contains("anyOf")) { + for (const auto & tt : t["anyOf"]) { + add_component(tt, false); + } + } else { + add_component(t, true); + } + } + if (!enum_values.empty()) { + std::vector enum_intersection; + for (const auto & p : enum_values) { + if (p.second == schema["allOf"].size()) { + enum_intersection.push_back(p.first); + } + } + if (!enum_intersection.empty()) { + return _add_rule(rule_name, "(" + string_join(enum_intersection, " | ") + ") space"); + } + } + return _add_rule(rule_name, _build_object_rule(properties, required, hybrid_name, json())); + } + if ((schema_type.is_null() || schema_type == "array") && (schema.contains("items") || schema.contains("prefixItems"))) { + json items = schema.contains("items") ? schema["items"] : schema["prefixItems"]; + if (items.is_array()) { + std::string rule = "\"[\" space "; + for (size_t i = 0; i < items.size(); i++) { + if (i > 0) { + rule += " \",\" space "; + } + rule += visit(items[i], name + (name.empty() ? "" : "-") + "tuple-" + std::to_string(i)); + } + rule += " \"]\" space"; + return _add_rule(rule_name, rule); + } + std::string item_rule_name = visit(items, name + (name.empty() ? "" : "-") + "item"); + int min_items = schema.contains("minItems") ? schema["minItems"].get() : 0; + json max_items_json = schema.contains("maxItems") ? schema["maxItems"] : json(); + int max_items = max_items_json.is_number_integer() ? max_items_json.get() : std::numeric_limits::max(); + + const std::string separator_rule = "\",\" space"; + return _add_rule(rule_name, "\"[\" space " + common_rust_string_to_std( + llama_common_gbnf_build_repetition_rust( + (const uint8_t *) item_rule_name.data(), item_rule_name.size(), + min_items, max_items, + (const uint8_t *) separator_rule.data(), separator_rule.size())) + " \"]\" space"); + } + if ((schema_type.is_null() || schema_type == "string") && schema.contains("pattern")) { + return _visit_pattern(schema["pattern"], rule_name); + } + if ((schema_type.is_null() || schema_type == "string") && std::regex_match(schema_format, std::regex("^uuid[1-5]?$"))) { + return _add_primitive(rule_name == "root" ? "root" : schema_format, PRIMITIVE_RULES.at("uuid")); + } + if ((schema_type.is_null() || schema_type == "string") && STRING_FORMAT_RULES.find(schema_format + "-string") != STRING_FORMAT_RULES.end()) { + auto prim_name = schema_format + "-string"; + return _add_rule(rule_name, _add_primitive(prim_name, STRING_FORMAT_RULES.at(prim_name))); + } + if (schema_type == "string" && (schema.contains("minLength") || schema.contains("maxLength"))) { + std::string char_rule = _add_primitive("char", PRIMITIVE_RULES.at("char")); + int min_len = schema.contains("minLength") ? schema["minLength"].get() : 0; + int max_len = schema.contains("maxLength") ? schema["maxLength"].get() : std::numeric_limits::max(); + return _add_rule(rule_name, "\"\\\"\" " + common_rust_string_to_std( + llama_common_gbnf_build_repetition_rust( + (const uint8_t *) char_rule.data(), char_rule.size(), + min_len, max_len, + nullptr, 0)) + " \"\\\"\" space"); + } + if (schema_type == "integer" && (schema.contains("minimum") || schema.contains("exclusiveMinimum") || schema.contains("maximum") || schema.contains("exclusiveMaximum"))) { + int64_t min_value = std::numeric_limits::min(); + int64_t max_value = std::numeric_limits::max(); + if (schema.contains("minimum")) { + min_value = schema["minimum"].get(); + } else if (schema.contains("exclusiveMinimum")) { + min_value = schema["exclusiveMinimum"].get() + 1; + } + if (schema.contains("maximum")) { + max_value = schema["maximum"].get(); + } else if (schema.contains("exclusiveMaximum")) { + max_value = schema["exclusiveMaximum"].get() - 1; + } + std::stringstream out; + out << "("; + build_min_max_int(min_value, max_value, out); + out << ") space"; + return _add_rule(rule_name, out.str()); + } + if (schema.empty() || schema_type == "object") { + return _add_rule(rule_name, _add_primitive("object", PRIMITIVE_RULES.at("object"))); + } + if (schema_type.is_null() && schema.is_object()) { + // No type constraint and no recognized structural keywords (e.g. {"description": "..."}). + // Per JSON Schema semantics this is equivalent to {} and accepts any value. + return _add_rule(rule_name, _add_primitive("value", PRIMITIVE_RULES.at("value"))); + } + if (!schema_type.is_string() || PRIMITIVE_RULES.find(schema_type.get()) == PRIMITIVE_RULES.end()) { + _errors.push_back("Unrecognized schema: " + schema.dump()); + return ""; + } + // TODO: support minimum, maximum, exclusiveMinimum, exclusiveMaximum at least for zero + return _add_primitive(rule_name == "root" ? "root" : schema_type.get(), PRIMITIVE_RULES.at(schema_type.get())); + } + + void check_errors() { + if (!_errors.empty()) { + throw std::invalid_argument("JSON schema conversion failed:\n" + string_join(_errors, "\n")); + } + if (!_warnings.empty()) { + fprintf(stderr, "WARNING: JSON schema conversion was incomplete: %s\n", string_join(_warnings, "; ").c_str()); + } + } + + std::string format_grammar() { + std::stringstream ss; + for (const auto & kv : _rules) { + ss << kv.first << " ::= " << kv.second << '\n'; + } + return ss.str(); + } +}; + +// common_schema_info implementation (pimpl) + +common_schema_info::common_schema_info() + : impl_(std::make_unique( + [](const std::string &) { return json(); }, + false)) {} + +common_schema_info::~common_schema_info() = default; + +common_schema_info::common_schema_info(common_schema_info &&) noexcept = default; +common_schema_info & common_schema_info::operator=(common_schema_info &&) noexcept = default; + +void common_schema_info::resolve_refs(nlohmann::ordered_json & schema) { + impl_->resolve_refs(schema, ""); +} + +// Determines if a JSON schema can resolve to a string type through any path. +// Some models emit raw string values rather than JSON-encoded strings for string parameters. +// If any branch of the schema (via oneOf, anyOf, $ref, etc.) permits a string, this returns +// true, allowing callers to handle the value as a raw string for simplicity. +bool common_schema_info::resolves_to_string(const nlohmann::ordered_json & schema) { + std::unordered_set visited_refs; + + std::function check = [&](const json & s) -> bool { + if (!s.is_object()) { + return false; + } + + // Handle $ref + if (s.contains("$ref")) { + const std::string & ref = s["$ref"]; + if (visited_refs.find(ref) != visited_refs.end()) { + // Circular reference, assume not a string to be safe + return false; + } + visited_refs.insert(ref); + auto it = impl_->_refs.find(ref); + if (it != impl_->_refs.end()) { + return check(it->second); + } + return false; + } + + // Check type field + if (s.contains("type")) { + const json & schema_type = s["type"]; + if (schema_type.is_string()) { + if (schema_type == "string") { + return true; + } + } else if (schema_type.is_array()) { + // Type can be an array like ["string", "null"] + for (const auto & t : schema_type) { + if (t == "string") { + return true; + } + } + } + } + + // Check oneOf/anyOf - if any alternative can be a string + if (s.contains("oneOf")) { + for (const auto & alt : s["oneOf"]) { + if (check(alt)) { + return true; + } + } + } + if (s.contains("anyOf")) { + for (const auto & alt : s["anyOf"]) { + if (check(alt)) { + return true; + } + } + } + + // Check allOf - all components must be compatible with string type + if (s.contains("allOf")) { + bool all_string = true; + for (const auto & component : s["allOf"]) { + if (!check(component)) { + all_string = false; + break; + } + } + if (all_string) { + return true; + } + } + + // Check const - if the constant value is a string + if (s.contains("const")) { + if (s["const"].is_string()) { + return true; + } + } + + // Check enum - if any enum value is a string + if (s.contains("enum")) { + for (const auto & val : s["enum"]) { + if (val.is_string()) { + return true; + } + } + } + + // String-specific keywords imply string type + if (s.contains("pattern") || s.contains("minLength") || s.contains("maxLength")) { + return true; + } + + // Check format - many formats imply string + if (s.contains("format")) { + const std::string & fmt = s["format"]; + if (fmt == "date" || fmt == "time" || fmt == "date-time" || + fmt == "uri" || fmt == "email" || fmt == "hostname" || + fmt == "ipv4" || fmt == "ipv6" || fmt == "uuid" || + fmt.find("uuid") == 0) { + return true; + } + } + + return false; + }; + + return check(schema); +} + +std::string json_schema_to_grammar(const json & schema, bool force_gbnf) { +#ifdef LLAMA_USE_LLGUIDANCE + if (!force_gbnf) { + return "%llguidance {}\nstart: %json " + schema.dump(); + } +#else + (void)force_gbnf; +#endif // LLAMA_USE_LLGUIDANCE + return build_grammar([&](const common_grammar_builder & callbacks) { + auto copy = schema; + callbacks.resolve_refs(copy); + callbacks.add_schema("", copy); + }); +} + +std::string build_grammar(const std::function & cb, const common_grammar_options & options) { + common_schema_converter converter([&](const std::string &) { return json(); }, options.dotall); + common_grammar_builder builder { + /* .add_rule = */ [&](const std::string & name, const std::string & rule) { + return converter._add_rule(name, rule); + }, + /* .add_schema = */ [&](const std::string & name, const nlohmann::ordered_json & schema) { + return converter.visit(schema, name == "root" ? "" : name); + }, + /* .resolve_refs = */ [&](nlohmann::ordered_json & schema) { + converter.resolve_refs(schema, ""); + } + }; + cb(builder); + converter.check_errors(); + return converter.format_grammar(); +} + +#include "common.h.inc" + +#define JSON_ASSERT GGML_ASSERT +#include +#include + +#include +#include +#include +#include // migration only +#include +#include +#include + +namespace nl = nlohmann; + +extern "C" const char * llama_build_info(void); + +#if defined(_WIN32) +#define WIN32_LEAN_AND_MEAN +#ifndef NOMINMAX +#define NOMINMAX +#endif +#define HOME_DIR "USERPROFILE" +#include +#else +#define HOME_DIR "HOME" +#include +#include +#endif + +namespace hf_cache { + +static std::pair hf_cache_http_client(const std::string & url) { + common_http_url parts = common_http_parse_url(url); + + if (parts.host.empty()) { + throw std::runtime_error("error: invalid URL format"); + } + +#ifndef CPPHTTPLIB_OPENSSL_SUPPORT + if (parts.scheme == "https") { + throw std::runtime_error( + "HTTPS is not supported. Please rebuild with one of:\n" + " -DLLAMA_BUILD_BORINGSSL=ON\n" + " -DLLAMA_BUILD_LIBRESSL=ON\n" + " -DLLAMA_OPENSSL=ON (default, requires OpenSSL dev files installed)" + ); + } +#endif + + httplib::Client cli(parts.scheme + "://" + parts.host + ":" + std::to_string(parts.port)); + + if (!parts.user.empty()) { + cli.set_basic_auth(parts.user, parts.password); + } + + cli.set_follow_location(true); + + return { std::move(cli), std::move(parts) }; +} + +namespace fs = std::filesystem; + +static fs::path get_cache_directory() { + static const fs::path cache = []() { + struct { + const char * var; + fs::path path; + } entries[] = { + {"LLAMA_CACHE", fs::path()}, + {"HF_HUB_CACHE", fs::path()}, + {"HUGGINGFACE_HUB_CACHE", fs::path()}, + {"HF_HOME", fs::path("hub")}, + {"XDG_CACHE_HOME", fs::path("huggingface") / "hub"}, + {HOME_DIR, fs::path(".cache") / "huggingface" / "hub"} + }; + for (const auto & entry : entries) { + if (auto * p = std::getenv(entry.var); p && *p) { + fs::path base(p); + return entry.path.empty() ? base : base / entry.path; + } + } +#ifndef _WIN32 + const struct passwd * pw = getpwuid(getuid()); + + if (pw && pw->pw_dir && *pw->pw_dir) { + return fs::path(pw->pw_dir) / ".cache" / "huggingface" / "hub"; + } +#endif + throw std::runtime_error("Failed to determine HF cache directory"); + }(); + + return cache; +} + +static fs::path get_repo_path(const std::string & repo_id) { + return get_cache_directory() / common_rust_string_to_std( + llama_common_hf_repo_to_folder_name_rust((const uint8_t *) repo_id.data(), repo_id.size())); +} + +static bool is_valid_subpath(const fs::path & path, const fs::path & subpath) { + if (subpath.is_absolute()) { + return false; // never do a / b with b absolute + } + auto b = fs::absolute(path).lexically_normal(); + auto t = (b / subpath).lexically_normal(); + auto [b_end, _] = std::mismatch(b.begin(), b.end(), t.begin(), t.end()); + + return b_end == b.end(); +} + +static void safe_write_file(const fs::path & path, const std::string & data) { + fs::path path_tmp = path.string() + ".tmp"; + + if (path.has_parent_path()) { + fs::create_directories(path.parent_path()); + } + + std::ofstream file(path_tmp); + file << data; + file.close(); + + std::error_code ec; + + if (!file.fail()) { + fs::rename(path_tmp, path, ec); + } + if (file.fail() || ec) { + fs::remove(path_tmp, ec); + throw std::runtime_error("failed to write file: " + path.string()); + } +} + +static nl::json api_get(const std::string & url, + const std::string & token) { + auto [cli, parts] = hf_cache_http_client(url); + + httplib::Headers headers = { + {"User-Agent", "llama-cpp/" + std::string(llama_build_info())}, + {"Accept", "application/json"} + }; + + if (llama_common_hf_is_valid_token_rust((const uint8_t *) token.data(), token.size())) { + headers.emplace("Authorization", "Bearer " + token); + } else if (!token.empty()) { + LOG_WRN("%s: invalid token, authentication disabled\n", __func__); + } + + if (auto res = cli.Get(parts.path, headers)) { + auto body = res->body; + + if (res->status == 200) { + return nl::json::parse(res->body); + } + try { + body = nl::json::parse(res->body)["error"].get(); + } catch (...) { } + + throw std::runtime_error("GET failed (" + std::to_string(res->status) + "): " + body); + } else { + throw std::runtime_error("HTTPLIB failed: " + httplib::to_string(res.error())); + } +} + +static std::string get_repo_commit(const std::string & repo_id, + const std::string & token) { + try { + auto endpoint = common_get_model_endpoint(); + auto json = api_get(endpoint + "api/models/" + repo_id + "/refs", token); + + if (!json.is_object() || + !json.contains("branches") || !json["branches"].is_array()) { + LOG_WRN("%s: missing 'branches' for '%s'\n", __func__, repo_id.c_str()); + return {}; + } + + fs::path refs_path = get_repo_path(repo_id) / "refs"; + std::string name; + std::string commit; + + for (const auto & branch : json["branches"]) { + if (!branch.is_object() || + !branch.contains("name") || !branch["name"].is_string() || + !branch.contains("targetCommit") || !branch["targetCommit"].is_string()) { + continue; + } + std::string _name = branch["name"].get(); + std::string _commit = branch["targetCommit"].get(); + + if (!is_valid_subpath(refs_path, _name)) { + LOG_WRN("%s: skip invalid branch: %s\n", __func__, _name.c_str()); + continue; + } + if (!llama_common_hf_is_valid_commit_rust((const uint8_t *) _commit.data(), _commit.size())) { + LOG_WRN("%s: skip invalid commit: %s\n", __func__, _commit.c_str()); + continue; + } + + if (_name == "main") { + name = _name; + commit = _commit; + break; + } + + if (name.empty() || commit.empty()) { + name = _name; + commit = _commit; + } + } + + if (name.empty() || commit.empty()) { + LOG_WRN("%s: no valid branch for '%s'\n", __func__, repo_id.c_str()); + return {}; + } + + safe_write_file(refs_path / name, commit); + return commit; + + } catch (const nl::json::exception & e) { + LOG_ERR("%s: JSON error: %s\n", __func__, e.what()); + } catch (const std::exception & e) { + LOG_ERR("%s: error: %s\n", __func__, e.what()); + } + return {}; +} + +hf_files get_repo_files(const std::string & repo_id, + const std::string & token) { + if (!llama_common_hf_is_valid_repo_id_rust((const uint8_t *) repo_id.data(), repo_id.size())) { + LOG_WRN("%s: invalid repository: %s\n", __func__, repo_id.c_str()); + return {}; + } + + std::string commit = get_repo_commit(repo_id, token); + if (commit.empty()) { + LOG_WRN("%s: failed to resolve commit for %s\n", __func__, repo_id.c_str()); + return {}; + } + + fs::path blobs_path = get_repo_path(repo_id) / "blobs"; + fs::path commit_path = get_repo_path(repo_id) / "snapshots" / commit; + + hf_files files; + + try { + auto endpoint = common_get_model_endpoint(); + auto json = api_get(endpoint + "api/models/" + repo_id + "/tree/" + commit + "?recursive=true", token); + + if (!json.is_array()) { + LOG_WRN("%s: response is not an array for '%s'\n", __func__, repo_id.c_str()); + return {}; + } + + for (const auto & item : json) { + if (!item.is_object() || + !item.contains("type") || !item["type"].is_string() || item["type"] != "file" || + !item.contains("path") || !item["path"].is_string()) { + continue; + } + + hf_file file; + file.repo_id = repo_id; + file.path = item["path"].get(); + + if (!is_valid_subpath(commit_path, file.path)) { + LOG_WRN("%s: skip invalid path: %s\n", __func__, file.path.c_str()); + continue; + } + + if (item.contains("lfs") && item["lfs"].is_object()) { + if (item["lfs"].contains("oid") && item["lfs"]["oid"].is_string()) { + file.oid = item["lfs"]["oid"].get(); + } + if (item["lfs"].contains("size") && item["lfs"]["size"].is_number()) { + file.size = item["lfs"]["size"].get(); + } + } else if (item.contains("oid") && item["oid"].is_string()) { + file.oid = item["oid"].get(); + } + if (file.size == 0 && item.contains("size") && item["size"].is_number()) { + file.size = item["size"].get(); + } + + if (!file.oid.empty() && + !llama_common_hf_is_valid_oid_rust((const uint8_t *) file.oid.data(), file.oid.size())) { + LOG_WRN("%s: skip invalid oid: %s\n", __func__, file.oid.c_str()); + continue; + } + + file.url = endpoint + repo_id + "/resolve/" + commit + "/" + file.path; + + fs::path final_path = commit_path / file.path; + file.final_path = final_path.string(); + + if (!file.oid.empty() && !fs::exists(final_path)) { + fs::path local_path = blobs_path / file.oid; + file.local_path = local_path.string(); + } else { + file.local_path = file.final_path; + } + + files.push_back(file); + } + } catch (const nl::json::exception & e) { + LOG_ERR("%s: JSON error: %s\n", __func__, e.what()); + } catch (const std::exception & e) { + LOG_ERR("%s: error: %s\n", __func__, e.what()); + } + return files; +} + +static std::string get_cached_ref(const fs::path & repo_path) { + fs::path refs_path = repo_path / "refs"; + if (!fs::is_directory(refs_path)) { + return {}; + } + std::string fallback; + + for (const auto & entry : fs::directory_iterator(refs_path)) { + if (!entry.is_regular_file()) { + continue; + } + std::ifstream f(entry.path()); + std::string commit; + if (!f || !std::getline(f, commit) || commit.empty()) { + continue; + } + if (!llama_common_hf_is_valid_commit_rust((const uint8_t *) commit.data(), commit.size())) { + LOG_WRN("%s: skip invalid commit: %s\n", __func__, commit.c_str()); + continue; + } + if (entry.path().filename() == "main") { + return commit; + } + if (fallback.empty()) { + fallback = commit; + } + } + return fallback; +} + +hf_files get_cached_files(const std::string & repo_id) { + fs::path cache_dir = get_cache_directory(); + if (!fs::exists(cache_dir)) { + return {}; + } + + if (!repo_id.empty() && + !llama_common_hf_is_valid_repo_id_rust((const uint8_t *) repo_id.data(), repo_id.size())) { + LOG_WRN("%s: invalid repository: %s\n", __func__, repo_id.c_str()); + return {}; + } + + hf_files files; + + for (const auto & repo : fs::directory_iterator(cache_dir)) { + if (!repo.is_directory()) { + continue; + } + fs::path snapshots_path = repo.path() / "snapshots"; + + if (!fs::exists(snapshots_path)) { + continue; + } + const std::string folder = repo.path().filename().string(); + std::string _repo_id = common_rust_string_to_std( + llama_common_hf_folder_name_to_repo_rust((const uint8_t *) folder.data(), folder.size())); + + if (!llama_common_hf_is_valid_repo_id_rust((const uint8_t *) _repo_id.data(), _repo_id.size())) { + continue; + } + if (!repo_id.empty() && _repo_id != repo_id) { + continue; + } + std::string commit = get_cached_ref(repo.path()); + fs::path commit_path = snapshots_path / commit; + + if (commit.empty() || !fs::is_directory(commit_path)) { + continue; + } + for (const auto & entry : fs::recursive_directory_iterator(commit_path)) { + if (!entry.is_regular_file() && !entry.is_symlink()) { + continue; + } + fs::path path = entry.path().lexically_relative(commit_path); + + if (!path.empty()) { + hf_file file; + file.repo_id = _repo_id; + file.path = path.generic_string(); + file.local_path = entry.path().string(); + file.final_path = file.local_path; + files.push_back(std::move(file)); + } + } + } + + return files; +} + +std::string finalize_file(const hf_file & file) { + static std::atomic symlinks_disabled{false}; + + std::error_code ec; + fs::path local_path(file.local_path); + fs::path final_path(file.final_path); + + if (local_path == final_path || fs::exists(final_path, ec)) { + return file.final_path; + } + + if (!fs::exists(local_path, ec)) { + return file.final_path; + } + + fs::create_directories(final_path.parent_path(), ec); + + if (!symlinks_disabled) { + fs::path target = fs::relative(local_path, final_path.parent_path(), ec); + if (!ec) { + fs::create_symlink(target, final_path, ec); + } + if (!ec) { + return file.final_path; + } + } + + if (!symlinks_disabled.exchange(true)) { + LOG_WRN("%s: failed to create symlink: %s\n", __func__, ec.message().c_str()); + LOG_WRN("%s: switching to degraded mode\n", __func__); + } + + fs::rename(local_path, final_path, ec); + if (ec) { + LOG_WRN("%s: failed to move file to snapshots: %s\n", __func__, ec.message().c_str()); + fs::copy(local_path, final_path, ec); + if (ec) { + LOG_ERR("%s: failed to copy file to snapshots: %s\n", __func__, ec.message().c_str()); + } + } + return file.final_path; +} + +// delete everything after this line, one day + +// copied from the download helpers without the tag part +struct gguf_split_info { + std::string prefix; // tag included + int index; + int count; +}; + +static gguf_split_info hf_cache_get_gguf_split_info(const std::string & path) { + common_rust_gguf_split_info split = + llama_common_gguf_split_info_rust((const uint8_t *) path.data(), path.size(), false); + gguf_split_info result { + common_rust_string_to_std(split.prefix), + split.index, + split.count, + }; + llama_common_unicode_string_free(split.tag); + return result; +} + +static std::pair parse_manifest_name(std::string & filename) { + common_rust_string_list parsed = + llama_common_hf_parse_manifest_name_rust((const uint8_t *) filename.data(), filename.size()); + if (!parsed.data || parsed.len != 2) { + llama_common_string_list_free(parsed); + return {}; + } + std::pair result = { + std::string((const char *) parsed.data[0].data, parsed.data[0].len), + std::string((const char *) parsed.data[1].data, parsed.data[1].len), + }; + llama_common_string_list_free(parsed); + return result; +} + +struct migrate_file { + std::string path; + std::string sha256; + size_t size; + fs::path old_path; + fs::path etag_path; + const hf_file * file; +}; + +using migrate_files = std::vector; + +static bool collect_file(const fs::path & old_cache, + const std::string & owner, + const std::string & repo, + const std::string & path, + const std::string & sha256, + const hf_files & files, + migrate_files & to_migrate) { + + const hf_file * file = nullptr; + + for (const auto & f : files) { + if (f.path == path) { + file = &f; + break; + } + } + + std::string old_filename = common_rust_string_to_std(llama_common_hf_make_old_cache_filename_rust( + (const uint8_t *) owner.data(), owner.size(), + (const uint8_t *) repo.data(), repo.size(), + (const uint8_t *) path.data(), path.size())); + fs::path old_path = old_cache / old_filename; + fs::path etag_path = old_path.string() + ".etag"; + + if (!fs::exists(old_path)) { + if (file && fs::exists(file->final_path)) { + return true; + } + LOG_WRN("%s: %s not found in old cache or HF cache\n", __func__, old_filename.c_str()); + return false; + } + + if (!file) { + LOG_WRN("%s: %s not found in current repo\n", __func__, old_filename.c_str()); + return false; + } + + if (!sha256.empty() && !file->oid.empty() && sha256 != file->oid) { + LOG_WRN("%s: %s is not up to date (sha256 mismatch)\n", __func__, old_filename.c_str()); + return false; + } + + if (file->size > 0) { + size_t size = fs::file_size(old_path); + if (size != file->size) { + LOG_WRN("%s: %s has wrong size %zu (expected %zu)\n", __func__, old_filename.c_str(), size, file->size); + return false; + } + } + + to_migrate.push_back({path, sha256, file->size, old_path, etag_path, file}); + return true; +} + +static bool collect_files(const fs::path & old_cache, + const std::string & owner, + const std::string & repo, + const nl::json & node, + const hf_files & files, + migrate_files & to_migrate) { + + if (!node.contains("rfilename") || + !node.contains("lfs") || + !node["lfs"].contains("sha256")) { + return true; + } + + std::string path = node["rfilename"]; + std::string sha256 = node["lfs"]["sha256"]; + + auto split = hf_cache_get_gguf_split_info(path); + + if (split.count <= 1) { + return collect_file(old_cache, owner, repo, path, sha256, files, to_migrate); + } + + std::vector> splits; + + for (const auto & f : files) { + auto split_f = hf_cache_get_gguf_split_info(f.path); + if (split_f.count == split.count && split_f.prefix == split.prefix) { + // sadly the manifest only provides the sha256 of the first file (index == 1) + // the rest will be verified using the size... + std::string f_sha256 = (split_f.index == 1) ? sha256 : ""; + splits.emplace_back(f.path, f_sha256); + } + } + + if ((int)splits.size() != split.count) { + LOG_WRN("%s: expected %d split files but found %d in repo\n", __func__, split.count, (int)splits.size()); + return false; + } + + for (const auto & [f_path, f_sha256] : splits) { + if (!collect_file(old_cache, owner, repo, f_path, f_sha256, files, to_migrate)) { + return false; + } + } + + return true; +} + +static bool migrate_file(const migrate_file & file) { + std::error_code ec; + + fs::path new_path(file.file->local_path); + fs::create_directories(new_path.parent_path(), ec); + + if (!fs::exists(new_path, ec)) { + fs::rename(file.old_path, new_path, ec); + if (ec) { + fs::copy_file(file.old_path, new_path, ec); + if (ec) { + LOG_ERR("%s: failed to move/copy %s: %s\n", __func__, file.old_path.string().c_str(), ec.message().c_str()); + return false; + } + } + fs::remove(file.old_path, ec); + } + fs::remove(file.etag_path, ec); + + std::string filename = finalize_file(*file.file); + LOG_INF("%s: migrated %s -> %s\n", __func__, file.old_path.filename().string().c_str(), filename.c_str()); + return true; +} + +void migrate_old_cache_to_hf_cache(const std::string & token, bool offline) { + fs::path old_cache = fs_get_cache_directory(); + if (!fs::exists(old_cache)) { + return; + } + + if (offline) { + LOG_WRN("%s: skipping migration in offline mode (will run when online)\n", __func__); + return; // -hf is not going to work + } + + bool warned = false; + + for (const auto & entry : fs::directory_iterator(old_cache)) { + if (!entry.is_regular_file()) { + continue; + } + auto filename = entry.path().filename().string(); + auto [owner, repo] = parse_manifest_name(filename); + + if (owner.empty() || repo.empty()) { + continue; + } + + if (!warned) { + warned = true; + LOG_WRN("================================================================================\n" + "WARNING: Migrating cache to HuggingFace cache directory\n" + " Old cache: %s\n" + " New cache: %s\n" + "This one-time migration moves models previously downloaded with -hf\n" + "from the legacy llama.cpp cache to the standard HuggingFace cache.\n" + "Models downloaded with --model-url are not affected.\n" + "================================================================================\n", + old_cache.string().c_str(), get_cache_directory().string().c_str()); + } + + auto repo_id = owner + "/" + repo; + auto files = get_repo_files(repo_id, token); + + if (files.empty()) { + LOG_WRN("%s: could not get repo files for %s, skipping\n", __func__, repo_id.c_str()); + continue; + } + + migrate_files to_migrate; + bool ok = true; + + try { + std::ifstream manifest(entry.path()); + auto json = nl::json::parse(manifest); + for (const char * key : {"ggufFile", "mmprojFile"}) { + if (json.contains(key)) { + if (!collect_files(old_cache, owner, repo, json[key], files, to_migrate)) { + ok = false; + break; + } + } + } + } catch (const std::exception & e) { + LOG_WRN("%s: failed to parse manifest %s: %s\n", __func__, filename.c_str(), e.what()); + continue; + } + + if (!ok) { + LOG_WRN("%s: migration skipped: one or more files failed validation\n", __func__); + continue; + } + + for (const auto & file : to_migrate) { + if (!migrate_file(file)) { + ok = false; + break; + } + } + + if (!ok) { + LOG_WRN("%s: migration failed: could not migrate all files\n", __func__); + continue; + } + + LOG_INF("%s: migration complete, deleting manifest: %s\n", __func__, entry.path().string().c_str()); + fs::remove(entry.path()); + } +} + +} // namespace hf_cache + +#include "common.h.inc" + +#include +#include +#include +#include +#include + + +#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 +#define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 +#define LLAMA_NGRAM_MIN 1 +#define LLAMA_NGRAM_MAX 4 + +extern "C" { +void * common_ngram_cache_rust_new(); +void * common_ngram_cache_rust_load(const char * path); +void common_ngram_cache_rust_free(void * ptr); +void common_ngram_cache_rust_update( + void * ptr, + size_t ngram_min, + size_t ngram_max, + const int32_t * input, + size_t input_len, + size_t nnew); +size_t common_ngram_cache_rust_draft( + const void * context, + const void * dynamic, + const void * static_cache, + const int32_t * input, + size_t input_len, + size_t n_draft, + size_t ngram_min, + size_t ngram_max, + int32_t * out, + size_t out_len); +void * common_ngram_map_rust_new(uint16_t size_key, uint16_t size_value, bool key_only, uint16_t min_hits); +void common_ngram_map_rust_free(void * ptr); +uint16_t common_ngram_map_rust_size_value(const void * ptr); +void common_ngram_map_rust_begin(void * ptr, const int32_t * tokens, size_t tokens_len); +size_t common_ngram_map_rust_draft( + void * ptr, + const int32_t * input, + size_t input_len, + int32_t sampled, + int32_t * out, + size_t out_len); +void common_ngram_map_rust_accept(void * ptr, uint16_t n_accepted); +size_t common_ngram_simple_rust_draft( + uint16_t size_ngram, + uint16_t size_mgram, + const int32_t * tokens, + size_t tokens_len, + int32_t sampled, + int32_t * out, + size_t out_len); +} + +struct common_ngram_simple_config { + uint16_t size_ngram; + uint16_t size_mgram; +}; + +static llama_tokens common_ngram_simple_draft( + const common_ngram_simple_config & config, + const llama_tokens & tokens, + llama_token sampled) { + static_assert(sizeof(llama_token) == sizeof(int32_t)); + llama_tokens out(config.size_mgram); + const size_t drafted = common_ngram_simple_rust_draft( + config.size_ngram, + config.size_mgram, + reinterpret_cast(tokens.data()), + tokens.size(), + sampled, + reinterpret_cast(out.data()), + out.size()); + out.resize(drafted); + return out; +} + +struct common_ngram_map { + common_ngram_map(uint16_t size_key, uint16_t size_value, bool key_only, uint16_t min_hits) + : impl(common_ngram_map_rust_new(size_key, size_value, key_only, min_hits)) {} + + ~common_ngram_map() { common_ngram_map_rust_free(impl); } + + common_ngram_map(const common_ngram_map &) = delete; + common_ngram_map & operator=(const common_ngram_map &) = delete; + + common_ngram_map(common_ngram_map && other) noexcept : impl(other.impl) { + other.impl = nullptr; + } + + common_ngram_map & operator=(common_ngram_map && other) noexcept { + if (this != &other) { + common_ngram_map_rust_free(impl); + impl = other.impl; + other.impl = nullptr; + } + return *this; + } + + uint16_t size_value() const { + return common_ngram_map_rust_size_value(impl); + } + + void begin(const llama_tokens & tokens) { + static_assert(sizeof(llama_token) == sizeof(int32_t)); + common_ngram_map_rust_begin( + impl, + reinterpret_cast(tokens.data()), + tokens.size()); + } + + void draft(const llama_tokens & input, llama_token sampled, llama_tokens & out) { + static_assert(sizeof(llama_token) == sizeof(int32_t)); + std::vector buf(size_value()); + const size_t drafted = common_ngram_map_rust_draft( + impl, + reinterpret_cast(input.data()), + input.size(), + sampled, + reinterpret_cast(buf.data()), + buf.size()); + out.assign(buf.begin(), buf.begin() + drafted); + } + + void accept(uint16_t n_accepted) { + common_ngram_map_rust_accept(impl, n_accepted); + } + +private: + void * impl = nullptr; +}; + +struct common_ngram_cache { + common_ngram_cache() : impl(common_ngram_cache_rust_new()) {} + explicit common_ngram_cache(void * ptr) : impl(ptr) {} + ~common_ngram_cache() { common_ngram_cache_rust_free(impl); } + + common_ngram_cache(const common_ngram_cache &) = delete; + common_ngram_cache & operator=(const common_ngram_cache &) = delete; + + common_ngram_cache(common_ngram_cache && other) noexcept : impl(other.impl) { + other.impl = nullptr; + } + + common_ngram_cache & operator=(common_ngram_cache && other) noexcept { + if (this != &other) { + common_ngram_cache_rust_free(impl); + impl = other.impl; + other.impl = nullptr; + } + return *this; + } + + void * get() const { return impl; } + +private: + void * impl = nullptr; +}; + +static common_ngram_cache common_ngram_cache_load(const std::string & filename) { + void * ptr = common_ngram_cache_rust_load(filename.c_str()); + if (ptr == nullptr) { + throw std::runtime_error("Unable to open file " + filename); + } + return common_ngram_cache(ptr); +} + +static void common_ngram_cache_update( + common_ngram_cache & ngram_cache, + int ngram_min, + int ngram_max, + std::vector & inp_data, + int nnew, + bool) { + static_assert(sizeof(llama_token) == sizeof(int32_t)); + common_ngram_cache_rust_update( + ngram_cache.get(), + ngram_min, + ngram_max, + reinterpret_cast(inp_data.data()), + inp_data.size(), + nnew); +} + +static void common_ngram_cache_draft( + std::vector & inp, + std::vector & draft, + int n_draft, + int ngram_min, + int ngram_max, + common_ngram_cache & nc_context, + common_ngram_cache & nc_dynamic, + common_ngram_cache & nc_static) { + GGML_ASSERT(draft.size() == 1); + static_assert(sizeof(llama_token) == sizeof(int32_t)); + + std::vector out(n_draft); + const size_t drafted = common_ngram_cache_rust_draft( + nc_context.get(), + nc_dynamic.get(), + nc_static.get(), + reinterpret_cast(inp.data()), + inp.size(), + n_draft, + ngram_min, + ngram_max, + reinterpret_cast(out.data()), + out.size()); + + for (size_t i = 0; i < drafted; ++i) { + LOG_DBG(" - draft candidate: token=%d\n", out[i]); + draft.push_back(out[i]); + } +} + +extern "C" { +void * common_ngram_mod_rust_new(uint16_t n, size_t size); +void common_ngram_mod_rust_free(void * ptr); +size_t common_ngram_mod_rust_idx(const void * ptr, const int32_t * tokens); +void common_ngram_mod_rust_add(void * ptr, const int32_t * tokens); +int32_t common_ngram_mod_rust_get(const void * ptr, const int32_t * tokens); +void common_ngram_mod_rust_reset(void * ptr); +size_t common_ngram_mod_rust_get_n(const void * ptr); +size_t common_ngram_mod_rust_get_used(const void * ptr); +size_t common_ngram_mod_rust_size(const void * ptr); +size_t common_ngram_mod_rust_size_bytes(const void * ptr); +} + +struct common_ngram_mod { + using entry_t = int32_t; + + static constexpr entry_t EMPTY = -1; + + common_ngram_mod(uint16_t n, size_t size) : impl(common_ngram_mod_rust_new(n, size)) {} + ~common_ngram_mod() { common_ngram_mod_rust_free(impl); } + + size_t idx(const entry_t * tokens) const { return common_ngram_mod_rust_idx(impl, tokens); } + void add(const entry_t * tokens) { common_ngram_mod_rust_add(impl, tokens); } + entry_t get(const entry_t * tokens) const { return common_ngram_mod_rust_get(impl, tokens); } + + void reset() { common_ngram_mod_rust_reset(impl); } + + size_t get_n() const { return common_ngram_mod_rust_get_n(impl); } + size_t get_used() const { return common_ngram_mod_rust_get_used(impl); } + + size_t size() const { return common_ngram_mod_rust_size(impl); } + size_t size_bytes() const { return common_ngram_mod_rust_size_bytes(impl); } + +private: + void * impl; +}; + +const std::vector common_speculative_types = { + COMMON_SPECULATIVE_TYPE_NONE, + COMMON_SPECULATIVE_TYPE_DRAFT, + COMMON_SPECULATIVE_TYPE_EAGLE3, + COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, + COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, + COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, + COMMON_SPECULATIVE_TYPE_NGRAM_MOD, + COMMON_SPECULATIVE_TYPE_NGRAM_CACHE +}; + +struct common_speculative_config { + common_speculative_type type; + common_params_speculative params; + + common_speculative_config(common_speculative_type t, + const common_params_speculative & p = common_params_speculative{}) : type(t), params(p) {} +}; + +static bool common_speculative_are_compatible( + const llama_model * model_tgt, + const llama_model * model_dft) { + const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt); + const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft); + + const bool vocab_type_tgt = llama_vocab_type(vocab_tgt); + LOG_DBG("%s: vocab_type tgt: %d\n", __func__, vocab_type_tgt); + + const bool vocab_type_dft = llama_vocab_type(vocab_dft); + LOG_DBG("%s: vocab_type dft: %d\n", __func__, vocab_type_dft); + + if (vocab_type_tgt != vocab_type_dft) { + LOG_WRN("%s: draft model vocab type must match target model to use speculation but " + "vocab_type_dft = %d while vocab_type_tgt = %d\n", __func__, vocab_type_dft, vocab_type_tgt); + return false; + } + + if (llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) || + (llama_vocab_get_add_bos(vocab_tgt) && llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft))) { + LOG_WRN("%s: draft model bos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n", + __func__, + llama_vocab_get_add_bos(vocab_tgt), llama_vocab_get_add_bos(vocab_dft), + llama_vocab_bos(vocab_tgt), llama_vocab_bos(vocab_dft)); + return false; + } + + if (llama_vocab_get_add_eos(vocab_tgt) != llama_vocab_get_add_eos(vocab_dft) || + (llama_vocab_get_add_eos(vocab_tgt) && llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft))) { + LOG_WRN("%s: draft model eos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n", + __func__, + llama_vocab_get_add_eos(vocab_tgt), llama_vocab_get_add_eos(vocab_dft), + llama_vocab_eos(vocab_tgt), llama_vocab_eos(vocab_dft)); + return false; + } + + { + const int n_vocab_tgt = llama_vocab_n_tokens(vocab_tgt); + const int n_vocab_dft = llama_vocab_n_tokens(vocab_dft); + const int vocab_diff = n_vocab_tgt > n_vocab_dft + ? n_vocab_tgt - n_vocab_dft + : n_vocab_dft - n_vocab_tgt; + + if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) { + LOG_DBG("%s: draft model vocab must closely match target model to use speculation but ", __func__); + LOG_DBG("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n", + n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE); + return false; + } + + for (int i = SPEC_VOCAB_CHECK_START_TOKEN_ID; i < std::min(n_vocab_tgt, n_vocab_dft); ++i) { + const char * token_text_tgt = llama_vocab_get_text(vocab_tgt, i); + const char * token_text_dft = llama_vocab_get_text(vocab_dft, i); + + if (std::strcmp(token_text_tgt, token_text_dft) != 0) { + LOG_DBG("%s: draft model vocab must match target model to use speculation but ", __func__); + LOG_DBG("token %d content differs - target '%s', draft '%s'\n", i, + common_token_to_piece(vocab_tgt, i).c_str(), + common_token_to_piece(vocab_dft, i).c_str()); + return false; + } + } + } + + return true; +} + +// state of an implementation of speculative decoding +// +// each implementation has a unique type and a state that is implementation-specific +// in a subclass of common_speculative_state +struct common_speculative_state { + const enum common_speculative_type type; + + size_t n_call_begin = 0; // number of times this implementation was called for refresh. + size_t n_call_draft = 0; // number of times this implementation was called for generation. + size_t n_call_accept = 0; // number of times this implementation was called for accumulation. + + size_t n_gen_drafts = 0; // number of times a draft or part was generated by this implementation. + size_t n_acc_drafts = 0; // number of times a draft or part was accepted by the target model. + size_t n_gen_tokens = 0; // number of tokens generated by this implementation. + size_t n_acc_tokens = 0; // number of tokens accepted by the target model. + + // TODO: track performance of most recent calls + const bool gen_perf = true; // whether to generate performance stats. + + int64_t t_begin_us = 0; // total time spent in refresh of this implementation in microseconds. + int64_t t_draft_us = 0; // total time spent in generating drafts in this implementation in microseconds. + int64_t t_accept_us = 0; // total time spent in accumulation of this implementation in microseconds. + + common_speculative_state(enum common_speculative_type type) : type(type) {} + + virtual ~common_speculative_state() = default; + + virtual void begin(const llama_tokens & prompt) = 0; + + virtual void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) = 0; + + virtual void accept(uint16_t n_accepted) = 0; + + virtual int32_t n_max(const common_params_speculative & params) const = 0; + virtual int32_t n_min(const common_params_speculative & params) const = 0; +}; + +struct common_speculative_checkpoint { + llama_pos pos_min = 0; + llama_pos pos_max = 0; + + int64_t n_tokens = 0; + + std::vector data; + + size_t size() const { + return data.size(); + } +}; + +struct common_speculative_state_draft : public common_speculative_state { + llama_context * ctx_tgt; // only used for retokenizing from ctx_dft + llama_context * ctx_dft; + + bool use_ckpt = false; + common_speculative_checkpoint ckpt; + + common_sampler * smpl; + + llama_batch batch; + llama_tokens prompt_dft; + + bool vocab_cmpt = true; // whether retokenization is needed + std::unordered_map vocab_map; + + common_speculative_state_draft( + enum common_speculative_type type, + llama_context * ctx_tgt, + llama_context * ctx_dft, + const std::vector> & replacements, + bool use_ckpt) + : common_speculative_state(type) + , ctx_tgt(ctx_tgt) + , ctx_dft(ctx_dft) + , use_ckpt(use_ckpt) + { + batch = llama_batch_init(llama_n_batch(ctx_dft), 0, 1); + smpl = nullptr; + + // TODO: optimize or pass from outside? + // { + // common_params_sampling params; + // params.no_perf = false; + // + // params.top_k = 40; + // params.top_p = 0.9; + // + // params.samplers = { + // COMMON_SAMPLER_TYPE_TOP_K, + // COMMON_SAMPLER_TYPE_TOP_P, + // COMMON_SAMPLER_TYPE_INFILL, + // }; + // + // result->smpl = common_sampler_init(llama_get_model(ctx_dft), params); + // } + { + common_params_sampling params; + params.no_perf = false; + params.top_k = 10; + params.samplers = { + COMMON_SAMPLER_TYPE_TOP_K, + }; + + smpl = common_sampler_init(llama_get_model(ctx_dft), params); + } + + vocab_cmpt = common_speculative_are_compatible(llama_get_model(ctx_tgt), llama_get_model(ctx_dft)); + LOG_DBG("vocab_cmpt = %d\n", vocab_cmpt); + + if (!vocab_cmpt) { + LOG_WRN("the target and draft vocabs are not compatible - tokens will be translated between the two\n"); + + for (const auto & pair : replacements) { + vocab_map[pair.first] = pair.second; + } + } + } + + ~common_speculative_state_draft() override { + llama_perf_context_print(ctx_dft); + + llama_free(ctx_dft); + + common_sampler_free(smpl); + + llama_batch_free(batch); + } + + void begin(const llama_tokens & /*prompt*/) override { + } + + size_t create_checkpoint(int n_tokens_prompt) { + int slot_id = 0; + const size_t checkpoint_size = llama_state_seq_get_size_ext(ctx_dft, slot_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + + ckpt.pos_min = llama_memory_seq_pos_min(llama_get_memory(ctx_dft), slot_id); + ckpt.pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), slot_id); + ckpt.n_tokens = n_tokens_prompt; + ckpt.data.resize(checkpoint_size); + + const size_t n = llama_state_seq_get_data_ext(ctx_dft, ckpt.data.data(), checkpoint_size, slot_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + if (n != checkpoint_size) { + GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", checkpoint_size, n); + } + + LOG_DBG("%s: pos_min = %d, pos_max = %d, size = %.3f MiB\n", __func__, + ckpt.pos_min, ckpt.pos_max, (float) ckpt.data.size() / 1024 / 1024); + return n; + } + + size_t restore_checkpoint() { + int slot_id = 0; + LOG_DBG("%s: pos_min = %d, pos_max = %d\n", __func__, ckpt.pos_min, ckpt.pos_max); + const size_t n = llama_state_seq_set_data_ext(ctx_dft, ckpt.data.data(), ckpt.size(), slot_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + if (n != ckpt.size()) { + GGML_ABORT("%s: failed to restore context checkpoint (pos_min=%d, pos_max=%d, size=%zu", + __func__, ckpt.pos_min, ckpt.pos_max, ckpt.size()); + } + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot_id, ckpt.pos_max + 1, -1); + + return n; + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) override { + const auto & sparams = params.draft; + + auto * spec = this; + + auto & batch = spec->batch; + auto & ctx_tgt = spec->ctx_tgt; + auto & ctx_dft = spec->ctx_dft; + auto & smpl = spec->smpl; + auto & prompt_dft = spec->prompt_dft; + + auto * mem_dft = llama_get_memory(ctx_dft); + + int reuse_i = 0; // index of part to be reused in prompt_dft + int reuse_n = 0; // length of part to be reused in prompt_dft + + const int n_ctx = llama_n_ctx(ctx_dft) - sparams.n_max; + + llama_tokens prompt_cnv; + if (!spec->vocab_cmpt) { + std::string text; + + text = common_detokenize(ctx_tgt, prompt_tgt, true); + text = replace_to_dft(text); + + LOG_DBG("%s: main->draft detokenized string: '%s'\n", __func__, text.c_str()); + + prompt_cnv = common_tokenize(ctx_dft, text, false, true); + + // convert id_last to draft vocab. llama_detokenize is called directly to avoid an allocation + const auto * model_tgt = llama_get_model(ctx_tgt); + const auto * vocab_tgt = llama_model_get_vocab(model_tgt); + + int32_t n_chars = llama_detokenize(vocab_tgt, &id_last, 1, nullptr, 0, false, false); + GGML_ASSERT(n_chars < 0 && "failed to detokenize id_last"); + + text.resize(-n_chars); + llama_detokenize(vocab_tgt, &id_last, 1, text.data(), text.size(), false, false); + text = replace_to_dft(text); + + LOG_DBG("main->draft detokenized id_last(%d): '%s'\n", id_last, text.c_str()); + id_last = common_tokenize(ctx_dft, text, false, true)[0]; + } + + const llama_tokens & prompt_cur = spec->vocab_cmpt ? prompt_tgt : prompt_cnv; + + const int i_start = std::max(0, (int) prompt_cur.size() - n_ctx); + + if (use_ckpt && i_start > 0) { + LOG_WRN("%s: context shift is not supported with checkpoint-based contexts - skipping\n", __func__); + return; + } + + // reuse as much as possible from the old draft context + // ideally, the draft context should be as big as the target context and we will always reuse the entire prompt + for (int i = 0; i < (int) prompt_dft.size(); ++i) { + int cur = 0; + while (i_start + cur < (int) prompt_cur.size() && + i + cur < (int) prompt_dft.size() && + prompt_cur[i_start + cur] == prompt_dft[i + cur]) { + cur++; + } + + if ((cur >= 256 || n_ctx >= (int) prompt_cur.size()) && cur > reuse_n) { + reuse_i = i; + reuse_n = cur; + } + + if (use_ckpt) { + break; + } + } + + LOG_DBG("%s: reuse_i = %d, reuse_n = %d, #prompt_dft = %zu, #prompt_cur = %zu\n", + __func__, reuse_i, reuse_n, prompt_dft.size(), prompt_cur.size()); + if (use_ckpt && ckpt.n_tokens > reuse_n) { + LOG_DBG("%s: checkpoint (n_tokens = %d) is outdated -> delete it\n", __func__, (int) ckpt.n_tokens); + + reuse_i = 0; + reuse_n = 0; + + ckpt = {}; + } + + result.clear(); + result.reserve(sparams.n_max); + + if (reuse_n == 0 || (use_ckpt && reuse_i > 0)) { + llama_memory_clear(mem_dft, false); + prompt_dft.clear(); + } else { + // this happens when a previous draft has been discarded (for example, due to being too small), but the + // target model agreed with it. in this case, we simply pass back the previous results to save compute + if (reuse_i + reuse_n < (int64_t) prompt_dft.size() && prompt_dft[reuse_i + reuse_n] == id_last) { + for (int i = reuse_i + reuse_n + 1; i < (int) prompt_dft.size(); ++i) { + result.push_back(prompt_dft[i]); + + if (sparams.n_max <= (int) result.size()) { + break; + } + } + + return; + } + + if (reuse_i > 0) { + GGML_ASSERT(!use_ckpt); + + bool is_removed = llama_memory_seq_rm (mem_dft, 0, 0, reuse_i); + if (!is_removed) { + LOG_ERR("%s: llama_memory_seq_rm failed, reuse_i=%d\n", __func__, reuse_i); + return; + } + llama_memory_seq_add(mem_dft, 0, reuse_i, -1, -reuse_i); + + prompt_dft.erase(prompt_dft.begin(), prompt_dft.begin() + reuse_i); + } + + if (reuse_n < (int) prompt_dft.size()) { + if (use_ckpt) { + if (ckpt.n_tokens > 0) { + LOG_DBG("%s: restoring checkpoint, reuse_n=%d, prompt_dft.size=%zu\n", __func__, reuse_n, prompt_dft.size()); + restore_checkpoint(); + reuse_n = ckpt.n_tokens; + prompt_dft.resize(reuse_n); + } + } else { + const bool is_removed = llama_memory_seq_rm(mem_dft, 0, reuse_n, -1); + if (!is_removed) { + LOG_ERR("%s: llama_memory_seq_rm failed, reuse_n=%d, prompt_dft.size=%zu\n", __func__, reuse_n, prompt_dft.size()); + return; + } + prompt_dft.erase(prompt_dft.begin() + reuse_n, prompt_dft.end()); + } + } + } + + // prepare a batch to evaluate any new tokens in the prompt + common_batch_clear(batch); + + for (size_t i = i_start + reuse_n; i < prompt_cur.size(); ++i) { + //LOG_DBG("i = %d, i_start = %d, reuse_n = %d, i - i_start = %d, id = %6d\n", i, i_start, reuse_n, i - i_start, prompt_cur[i]); + common_batch_add(batch, prompt_cur[i], i - i_start, { 0 }, false); + + prompt_dft.push_back(prompt_cur[i]); + } + + // we should rarely end-up here during normal decoding + if (batch.n_tokens > 0) { + //LOG_DBG("%s: draft prompt batch: %s\n", __func__, string_from(ctx, batch).c_str()); + LOG_DBG("%s: draft prompt batch: %d tokens\n", __func__, batch.n_tokens); + + int ret = llama_decode(ctx_dft, batch); + if (ret != 0 && ret != 1) { + LOG_WRN("%s: llama_decode returned %d, prompt_cur.size=%zu\n", + __func__, ret, prompt_cur.size()); + } + + if (use_ckpt) { + create_checkpoint(prompt_dft.size()); + } + } + + const llama_pos n_past = prompt_dft.size(); + + LOG_DBG("%s: n_past = %d\n", __func__, n_past); + + common_batch_clear(batch); + common_batch_add (batch, id_last, n_past, { 0 }, true); + + prompt_dft.push_back(id_last); + + //LOG_DBG("%s: draft prompt: %s\n", __func__, string_from(ctx_dft, prompt_dft).c_str()); + + int ret = llama_decode(ctx_dft, batch); + if (ret != 0 && ret != 1) { + LOG_WRN("%s: llama_decode returned %d, prompt_cur.size=%zu, prompt_dft.size=%zu\n", + __func__, ret, prompt_cur.size(), prompt_dft.size()); + } + + common_sampler_reset(smpl); + + // sample n_draft tokens from the draft model + for (int i = 0; i < sparams.n_max; ++i) { + common_batch_clear(batch); + + common_sampler_sample(smpl, ctx_dft, 0, true); + + const auto * cur_p = common_sampler_get_candidates(smpl, true); + + for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { + LOG_DBG(" - draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", + k, i, cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); + } + + // add drafted token for each sequence + const llama_token id = cur_p->data[0].id; + + common_sampler_accept(smpl, id, true); + + // only collect very high-confidence draft tokens + if (cur_p->data[0].p < sparams.p_min) { + break; + } + + result.push_back(id); + + if (sparams.n_max <= (int) result.size()) { + break; + } + + common_batch_add(batch, id, n_past + i + 1, { 0 }, true); + + // evaluate the drafted tokens on the draft model + ret = llama_decode(ctx_dft, batch); + if (ret != 0) { + LOG_WRN("%s: llama_decode[%d] returned %d, prompt_cur.size=%zu, prompt_dft.size=%zu\n", + __func__, i, ret, prompt_cur.size(), prompt_dft.size()); + } + + prompt_dft.push_back(id); + } + + if (!spec->vocab_cmpt) { + std::string detokenized = common_detokenize(ctx_dft, result, true); + detokenized = replace_to_tgt(detokenized); + LOG_DBG("draft->main detokenized string: '%s'\n", detokenized.c_str()); + result = common_tokenize(ctx_tgt, detokenized, false, true); + if (result.size() > (size_t) sparams.n_max) { + result.resize(sparams.n_max); + } + } + + if (result.size() < (size_t) sparams.n_min) { + result.clear(); + } + } + + void accept(uint16_t n_accepted) override { + // noop + GGML_UNUSED(n_accepted); + } + + int32_t n_max(const common_params_speculative & params) const override { + return params.draft.n_max; + } + + int32_t n_min(const common_params_speculative & params) const override { + return params.draft.n_min; + } + + std::string replace_to_dft(const std::string & input) const { + std::string result = input; + + for (const auto & pair : this->vocab_map) { + size_t pos = result.find(pair.first); + while (pos != std::string::npos) { + result.replace(pos, pair.first.length(), pair.second); + pos = result.find(pair.first, pos + pair.second.length()); + } + } + + return result; + } + + std::string replace_to_tgt(const std::string & input) const { + std::string result = input; + + for (const auto & pair : this->vocab_map) { + size_t pos = result.find(pair.second); + while (pos != std::string::npos) { + result.replace(pos, pair.second.length(), pair.first); + pos = result.find(pair.second, pos + pair.first.length()); + } + } + + return result; + } +}; + +struct common_speculative_state_eagle3 : public common_speculative_state { + common_speculative_state_eagle3(enum common_speculative_type type) : common_speculative_state(type) {} + + void begin(const llama_tokens & prompt) override { + GGML_UNUSED(prompt); + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & draft_tokens) override { + // TODO: implement + GGML_UNUSED(params); + GGML_UNUSED(prompt_tgt); + GGML_UNUSED(id_last); + GGML_UNUSED(draft_tokens); + } + + void accept(uint16_t n_accepted) override { + // noop + GGML_UNUSED(n_accepted); + } + + int32_t n_max(const common_params_speculative & params) const override { + return params.draft.n_max; + } + + int32_t n_min(const common_params_speculative & params) const override { + return params.draft.n_min; + } +}; + +// state of self-speculation (simple implementation, not ngram-map) +struct common_speculative_state_ngram_simple : public common_speculative_state { + common_ngram_simple_config config; + + common_speculative_state_ngram_simple( + enum common_speculative_type type, + common_ngram_simple_config config) + : common_speculative_state(type), config(config) {} + + void begin(const llama_tokens & prompt) override { + GGML_UNUSED(prompt); + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) override { + + result = common_ngram_simple_draft(config, prompt_tgt, id_last); + GGML_UNUSED(params); + } + + void accept(uint16_t n_accepted) override { + // noop + GGML_UNUSED(n_accepted); + } + + int32_t n_max(const common_params_speculative & /*params*/) const override { + return config.size_mgram; + } + + int32_t n_min(const common_params_speculative & /*params*/) const override { + return config.size_mgram; + } +}; + +struct common_speculative_state_ngram_map_k : public common_speculative_state { + // draft ngram map for speculative decoding without draft model + common_ngram_map config; + + common_speculative_state_ngram_map_k( + enum common_speculative_type type, + common_ngram_map config) + : common_speculative_state(type), config(std::move(config)) {} + + void begin(const llama_tokens & prompt) override { + config.begin(prompt); + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) override { + config.draft(prompt_tgt, id_last, result); + GGML_UNUSED(params); + } + + void accept(uint16_t n_accepted) override { + config.accept(n_accepted); + } + + int32_t n_max(const common_params_speculative & /*params*/) const override { + return config.size_value(); + } + + int32_t n_min(const common_params_speculative & /*params*/) const override { + return config.size_value(); + } +}; + +struct common_speculative_state_ngram_mod : public common_speculative_state { + common_ngram_mod & mod; + + // the last position in the prompt that was added to the ngram container + size_t i_last = 0; + + // length of the last drafted n‑gram (number of tokens returned by draft) + size_t n_draft_last = 0; + + // consecutive accept rounds with low acceptance fraction (< 0.5) + int n_low = 0; + + // enable trace logging if LLAMA_TRACE is set + const bool verbose; + + common_speculative_state_ngram_mod(enum common_speculative_type type, common_ngram_mod & mod) + : common_speculative_state(type), mod(mod), verbose(std::getenv("LLAMA_TRACE") != nullptr) { + static_assert(sizeof(llama_token) == sizeof(common_ngram_mod::entry_t)); + } + + void begin(const llama_tokens & prompt) override { + i_last = 0; + + n_draft_last = 0; + + const size_t n = mod.get_n(); + + if (prompt.size() < n) { + return; + } + + for (size_t i = 0; i < prompt.size() - n; ++i) { + mod.add(prompt.data() + i); + } + + i_last = prompt.size() - n; + + const double f = (double)mod.get_used() / (double)mod.size(); + LOG_INF("%s: ngram_mod occupancy = %zu/%zu (%.2f)\n", __func__, mod.get_used(), mod.size(), f); + + constexpr double f_thold = 0.25; + if (f > f_thold) { + LOG_WRN("%s: ngram_mod occupancy %.2f exceeds threshold (%.2f) - resetting\n", __func__, f, f_thold); + + mod.reset(); + } + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) override { + const auto & sparams = params.ngram_mod; + + n_draft_last = 0; + + const size_t cur_len = prompt_tgt.size(); + if (cur_len < mod.get_n()) { + return; + } + + const size_t n = mod.get_n(); + + // add new ngrams in chunks + if (i_last + 32 < cur_len) { + for (size_t i = i_last; i < cur_len - n; ++i) { + mod.add(prompt_tgt.data() + i); + } + + i_last = cur_len - n; + } + + result.resize(n + sparams.n_max); + for (size_t i = 0; i < n - 1; ++i) { + result[i] = prompt_tgt[cur_len - n + 1 + i]; + } + result[n - 1] = id_last; + + for (int i = 0; i < sparams.n_max; ++i) { + const llama_token token = mod.get(result.data() + i); + if (token == common_ngram_mod::EMPTY) { + if (i < sparams.n_min) { + result.clear(); + return; + } + + result.resize(n + i); + break; + } + result[n + i] = token; + } + + // only return the m tokens that were drafted + for (size_t i = 0; n + i < result.size(); ++i) { + result[i] = result[n + i]; + } + result.resize(result.size() - n); + + // store length of drafted n‑gram for later acceptance analysis + n_draft_last = result.size(); + } + + void accept(uint16_t n_accepted) override { + // compute acceptance fraction if we have a recorded draft length + if (n_draft_last > 0) { + const double f_acc = (double)n_accepted / (double)n_draft_last; + if (f_acc < 0.5) { + n_low++; + if (n_low >= 3) { + if (verbose) { + LOG_WRN("%s: low acceptance streak (%d) – resetting ngram_mod\n", __func__, n_low); + } + + mod.reset(); + n_low = 0; + i_last = 0; + } + } else { + n_low = 0; + } + } + } + + int32_t n_max(const common_params_speculative & params) const override { + return params.ngram_mod.n_max; + } + + int32_t n_min(const common_params_speculative & params) const override { + return params.ngram_mod.n_min; + } +}; + +struct common_speculative_state_ngram_cache : public common_speculative_state { + uint16_t n_draft; + bool save_dynamic; + bool save_static; + + common_ngram_cache ngram_cache_context; + common_ngram_cache ngram_cache_dynamic; + common_ngram_cache ngram_cache_static; + + size_t cache_size = 0; // number of tokens in n-gram cache + + common_speculative_state_ngram_cache( + const enum common_speculative_type type, + const std::string & path_static, + const std::string & path_dynamic, + uint16_t n_draft, + bool save_dynamic, + bool save_static) + : common_speculative_state(type) + , n_draft(n_draft) + , save_dynamic(save_dynamic) + , save_static(save_static) + { + if (!path_static.empty()) { + try { + ngram_cache_static = common_ngram_cache_load(path_static); + } catch (...) { + LOG_ERR("failed to open static lookup cache: %s", path_static.c_str()); + GGML_ABORT("Couldn't read static lookup cache"); + } + } + + if (!path_dynamic.empty()) { + try { + ngram_cache_dynamic = common_ngram_cache_load(path_dynamic); + } catch (...) { + LOG_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str()); + GGML_ABORT("Couldn't read dynamic lookup cache"); + } + } + } + + void begin(const llama_tokens & prompt) override { + GGML_UNUSED(prompt); + } + + void draft( + const common_params_speculative & params, + const llama_tokens & prompt_tgt, + llama_token id_last, + llama_tokens & result) override { + GGML_UNUSED(params); + + if (cache_size < prompt_tgt.size() + 1) { + llama_tokens tokens_new; + tokens_new.reserve(prompt_tgt.size() + 1 - cache_size); + for (size_t j = cache_size; j < prompt_tgt.size(); ++j) { + tokens_new.push_back(prompt_tgt[j]); + } + tokens_new.push_back(id_last); // add the last token + + // Update context ngram cache with new prompt_tgt: + common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, + tokens_new, tokens_new.size(), false); + cache_size = prompt_tgt.size() + 1; + } + + llama_tokens inp; + inp.reserve(prompt_tgt.size() + 1); + for (size_t j = 0; j < prompt_tgt.size(); ++j) { + inp.push_back(prompt_tgt[j]); + } + inp.push_back(id_last); + + result.push_back(id_last); + + common_ngram_cache_draft(inp, result, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, + ngram_cache_context, + ngram_cache_dynamic, + ngram_cache_static); + + if (result.size() > 0) { + // delete first token in result (which is the id_last token) + result.erase(result.begin()); + } + } + + void accept(uint16_t n_accepted) override { + // TODO: noop + GGML_UNUSED(n_accepted); + } + + int32_t n_max(const common_params_speculative & /*params*/) const override { + return n_draft; + } + + int32_t n_min(const common_params_speculative & /*params*/) const override { + return 0; + } +}; + +struct common_speculative { + std::vector> impls; // list of implementations to use and their states + + common_speculative_state * curr_impl = nullptr; // current implementation in use (for stats) +}; + +static common_ngram_map make_common_ngram_map( + common_speculative_type type, + const common_params_speculative_ngram_map & config) { + uint16_t size_key = config.size_n; + uint16_t size_value = config.size_m; + bool key_only = type == COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K; + uint16_t min_hits = config.min_hits; + + return common_ngram_map(size_key, size_value, key_only, min_hits); +} + +static common_speculative_state_ngram_cache create_state_ngram_cache( + const std::string & path_static, const std::string & path_dynamic, + const common_speculative_config & config) { + uint16_t n_draft = 8; // TODO get from config? + + // TODO bool param in common/common.h.inc to set save_static/save_dynamic? + bool save_static = false; + bool save_dynamic = false; + + common_speculative_state_ngram_cache state(config.type, path_static, path_dynamic, n_draft, save_static, save_dynamic); + + return state; +} + +std::string common_speculative_type_name_str() { + std::string result; + for (size_t i = 0; i < common_speculative_types.size(); i++) { + if (i > 0) { + result += ", "; + } + result += common_speculative_type_to_str(common_speculative_types[i]); + } + return result; +} + +std::string common_speculative_type_to_str(enum common_speculative_type type) { + const char * name = llama_common_speculative_type_to_str_rust((int) type); + return name ? name : "unknown"; +} + +enum common_speculative_type common_speculative_type_from_name(const std::string & name) { + const int parsed = llama_common_speculative_type_from_name_rust(name.c_str()); + if (parsed < 0) { + return COMMON_SPECULATIVE_TYPE_COUNT; + } + return (common_speculative_type) parsed; +} + +// initialization of the speculative decoding system +// +common_speculative * common_speculative_init( + common_params_speculative & params, + llama_context * ctx_tgt) { + llama_context * ctx_dft = nullptr; + if (params.draft.model) { + ctx_dft = llama_init_from_model(params.draft.model, params.draft.cparams); + if (ctx_dft == nullptr) { + LOG_ERR("%s", "failed to create draft context\n"); + return nullptr; + } + } + + // Compute the implementations to use based on the config and their order of preference + std::vector configs = {}; // list of speculative configs to try + { + bool has_draft = !params.draft.mparams.path.empty(); + bool has_draft_eagle3 = false; // TODO PR-18039: if params.speculative.eagle3 + + bool has_ngram_cache = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_CACHE); + bool has_ngram_simple = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE); + bool has_ngram_map_k = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K); + bool has_ngram_map_k4v = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V); + bool has_ngram_mod = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_MOD); + + // In a more complex implementation we could use the same implementation but with different parameters. + // This was initially used in PR-18471 but removed to simplify the code. + if (has_ngram_simple) { + // This implementation can guess a lot of tokens without any draft model. + configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, params)); + } + if (has_ngram_map_k) { + configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, params)); + } + if (has_ngram_map_k4v) { + // This implementation can guess tokens with high acceptance rate but is more expensive. + configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, params)); + } + if (has_ngram_mod) { + auto & sparams = params.ngram_mod; + + if (!sparams.obj) { + sparams.obj = std::make_shared(sparams.n_match, 4*1024*1024); + + LOG_INF("%s: initialized ngram_mod with n_match=%d, size=%zu (%.3f MB)\n", __func__, + sparams.n_match, sparams.obj->size(), (float)(sparams.obj->size_bytes())/1024/1024); + + if (sparams.n_match < 16) { + LOG_WRN("%s: ngram_mod n_match=%d is too small - poor quality is possible, " + "see: https://github.com/ggml-org/llama.cpp/pull/19164\n", __func__, sparams.n_match); + } + } + + configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, params)); + } + if (has_ngram_cache) { + configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, params)); + } + if (has_draft) { + configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT, params)); + } + if (has_draft_eagle3) { + configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_EAGLE3, params)); + } + } + + std::vector> impls = {}; + + for (const common_speculative_config & config : configs) { + LOG_DBG("%s: adding implementation %s\n", __func__, common_speculative_type_to_str(config.type).c_str()); + switch (config.type) { + case COMMON_SPECULATIVE_TYPE_NONE: + break; + case COMMON_SPECULATIVE_TYPE_DRAFT: { + const bool use_ckpt = common_context_can_seq_rm(ctx_dft) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; + + impls.push_back(std::make_unique(config.type, + /* .ctx_tgt = */ ctx_tgt, + /* .ctx_dft = */ ctx_dft, + /* .replacements = */ params.draft.replacements, + /* .use_ckpt = */ use_ckpt + )); + break; + } + case COMMON_SPECULATIVE_TYPE_EAGLE3: { + impls.push_back(std::make_unique(config.type)); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { + auto config_simple = common_ngram_simple_config { + /* .size_ngram = */ config.params.ngram_simple.size_n, + /* .size_mgram = */ config.params.ngram_simple.size_m + }; + auto state = std::make_unique( + /* .type = */ config.type, + /* .state = */ config_simple + ); + impls.push_back(std::move(state)); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: { + impls.push_back(std::make_unique( + (config.type), + make_common_ngram_map(config.type, config.params.ngram_map_k) + )); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: { + GGML_ASSERT(config.params.ngram_mod.obj); + impls.push_back(std::make_unique(config.type, *config.params.ngram_mod.obj)); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: { + auto state = create_state_ngram_cache(params.ngram_cache.lookup_cache_static, params.ngram_cache.lookup_cache_dynamic, config); + impls.push_back(std::make_unique(std::move(state))); + break; + } + default: + break; + } + } + + if (impls.empty()) { + LOG_WRN("%s", "no implementations specified for speculative decoding\n"); + return nullptr; + } + + auto * result = new common_speculative { + /* .impls = */ std::move(impls), + /* .curr_impl = */ nullptr, + }; + + return result; +} + +void common_speculative_free(common_speculative * spec) { + if (spec == nullptr) { + return; + } + + delete spec; +} + +void common_speculative_begin(common_speculative * spec, const llama_tokens & prompt) { + if (spec == nullptr) { + return; + } + + for (auto & impl : spec->impls) { + common_time_meas tm(impl->t_begin_us, !impl->gen_perf); + impl->begin(prompt); + impl->n_call_begin++; + } +} + +llama_tokens common_speculative_draft( + common_speculative * spec, + const common_params_speculative & params, + const llama_tokens & prompt_tgt, // specified in target model vocab + llama_token id_last) { + llama_tokens result; + + spec->curr_impl = nullptr; // reset current implementation + + for (auto & impl : spec->impls) { + { + common_time_meas tm(impl->t_draft_us, !impl->gen_perf); + impl->draft(params, prompt_tgt, id_last, result); + impl->n_call_draft++; + } + + { + const int n_min = impl->n_min(params); + + if (!result.empty() && (int) result.size() < n_min) { + LOG_DBG("%s: ignoring small draft: %d < %d\n", __func__, (int) result.size(), n_min); + result.clear(); + } + } + + if (!result.empty()) { + LOG_DBG("%s: called impl %s, hist size = %zu, call_count = %zu, gen = %zu\n", __func__, + common_speculative_type_to_str(impl.get()->type).c_str(), prompt_tgt.size(), + impl.get()->n_call_draft, result.size()); + + spec->curr_impl = impl.get(); // set current implementation for stats + impl->n_gen_drafts++; + impl->n_gen_tokens += result.size(); + + break; // we have a draft, so break out of the loop and return it. + } + } + + return result; +} + +void common_speculative_accept(common_speculative * spec, uint16_t n_accepted) { + if (n_accepted == 0) { + return; + } + + common_speculative_state * impl = spec->curr_impl; + + GGML_ASSERT(impl); + + { + common_time_meas tm(impl->t_accept_us, !impl->gen_perf); + if (n_accepted > 0) { + impl->n_acc_drafts++; + impl->n_acc_tokens += n_accepted; + } + + impl->accept(n_accepted); + impl->n_call_accept++; + } +} + +int32_t common_speculative_n_max(const common_speculative * spec, const common_params_speculative & params) { + if (spec == nullptr) { + return 0; + } + + int32_t n_max = 0; + for (const auto & impl : spec->impls) { + n_max = std::max(n_max, impl->n_max(params)); + } + + return n_max; +} + +int32_t common_speculative_n_min(const common_speculative * spec, const common_params_speculative & params) { + if (spec == nullptr) { + return 0; + } + + int32_t n_min = 0; + for (const auto & impl : spec->impls) { + n_min = std::max(n_min, impl->n_min(params)); + } + + return n_min; +} + +void common_speculative_print_stats(const common_speculative * spec) { + if (spec == nullptr) { + return; + } + + for (const auto & impl : spec->impls) { + std::string str_perf; + if (impl->gen_perf) { + std::ostringstream oss; + oss << std::fixed << std::setprecision(3) << impl->t_begin_us / 1000.0 << ", "; + oss << std::fixed << std::setprecision(3) << impl->t_draft_us / 1000.0 << ", "; + oss << std::fixed << std::setprecision(3) << impl->t_accept_us / 1000.0; + str_perf = ", dur(b,g,a) = " + oss.str() + " ms"; + } else { + str_perf = ""; + } + + LOG_INF("statistics %s: #calls(b,g,a) = %zu %zu %zu, #gen drafts = %zu, #acc drafts = %zu, #gen tokens = %zu, #acc tokens = %zu%s\n", + common_speculative_type_to_str(impl->type).c_str(), + impl->n_call_begin, impl->n_call_draft, impl->n_call_accept, + impl->n_gen_drafts, + impl->n_acc_drafts, + impl->n_gen_tokens, + impl->n_acc_tokens, + str_perf.c_str()); + } +} diff --git a/common/common.h b/common/common.h deleted file mode 100644 index a564b3b8c2b4..000000000000 --- a/common/common.h +++ /dev/null @@ -1,1028 +0,0 @@ -// Various helper functions and utilities - -#pragma once - -#include "llama-cpp.h" - -#include "ggml-opt.h" -#include "ggml.h" - -#include -#include -#include -#include -#include -#include - -#if defined(_WIN32) && !defined(_WIN32_WINNT) -#define _WIN32_WINNT 0x0A00 -#endif - -#ifdef _WIN32 -#define DIRECTORY_SEPARATOR '\\' -#else -#define DIRECTORY_SEPARATOR '/' -#endif // _WIN32 - -#define die(msg) do { fputs("error: " msg "\n", stderr); exit(1); } while (0) -#define die_fmt(fmt, ...) do { fprintf(stderr, "error: " fmt "\n", __VA_ARGS__); exit(1); } while (0) - -struct common_time_meas { - common_time_meas(int64_t & t_acc, bool disable = false); - ~common_time_meas(); - - const int64_t t_start_us; - - int64_t & t_acc; -}; - -struct common_adapter_lora_info { - std::string path; - float scale; - - std::string task_name; - std::string prompt_prefix; - - struct llama_adapter_lora * ptr; -}; - -using llama_tokens = std::vector; - -struct common_control_vector_load_info; - -// -// CPU utils -// - -struct common_cpu_params { - int n_threads = -1; - bool cpumask[GGML_MAX_N_THREADS] = {false}; // CPU affinity mask. - bool mask_valid = false; // Default: any CPU - enum ggml_sched_priority priority = GGML_SCHED_PRIO_NORMAL; // Scheduling prio : (0 - normal, 1 - medium, 2 - high, 3 - realtime) - bool strict_cpu = false; // Use strict CPU placement - uint32_t poll = 50; // Polling (busywait) level (0 - no polling, 100 - mostly polling) -}; - -int32_t common_cpu_get_num_physical_cores(); -int32_t common_cpu_get_num_math(); - -// -// Common params -// - -enum llama_example { - LLAMA_EXAMPLE_BATCHED, - LLAMA_EXAMPLE_DEBUG, - LLAMA_EXAMPLE_COMMON, - LLAMA_EXAMPLE_SPECULATIVE, - LLAMA_EXAMPLE_COMPLETION, - LLAMA_EXAMPLE_CLI, - LLAMA_EXAMPLE_EMBEDDING, - LLAMA_EXAMPLE_PERPLEXITY, - LLAMA_EXAMPLE_RETRIEVAL, - LLAMA_EXAMPLE_PASSKEY, - LLAMA_EXAMPLE_IMATRIX, - LLAMA_EXAMPLE_BENCH, - LLAMA_EXAMPLE_SERVER, - LLAMA_EXAMPLE_CVECTOR_GENERATOR, - LLAMA_EXAMPLE_EXPORT_LORA, - LLAMA_EXAMPLE_MTMD, - LLAMA_EXAMPLE_LOOKUP, - LLAMA_EXAMPLE_PARALLEL, - LLAMA_EXAMPLE_TTS, - LLAMA_EXAMPLE_DIFFUSION, - LLAMA_EXAMPLE_FINETUNE, - LLAMA_EXAMPLE_FIT_PARAMS, - LLAMA_EXAMPLE_RESULTS, - LLAMA_EXAMPLE_EXPORT_GRAPH_OPS, - - LLAMA_EXAMPLE_COUNT, -}; - -enum common_sampler_type { - COMMON_SAMPLER_TYPE_NONE = 0, - COMMON_SAMPLER_TYPE_DRY = 1, - COMMON_SAMPLER_TYPE_TOP_K = 2, - COMMON_SAMPLER_TYPE_TOP_P = 3, - COMMON_SAMPLER_TYPE_MIN_P = 4, - //COMMON_SAMPLER_TYPE_TFS_Z = 5, - COMMON_SAMPLER_TYPE_TYPICAL_P = 6, - COMMON_SAMPLER_TYPE_TEMPERATURE = 7, - COMMON_SAMPLER_TYPE_XTC = 8, - COMMON_SAMPLER_TYPE_INFILL = 9, - COMMON_SAMPLER_TYPE_PENALTIES = 10, - COMMON_SAMPLER_TYPE_TOP_N_SIGMA = 11, - COMMON_SAMPLER_TYPE_ADAPTIVE_P = 12, -}; - -// dimensionality reduction methods, used by cvector-generator -enum dimre_method { - DIMRE_METHOD_PCA, - DIMRE_METHOD_MEAN, -}; - -enum common_conversation_mode { - COMMON_CONVERSATION_MODE_DISABLED = 0, - COMMON_CONVERSATION_MODE_ENABLED = 1, - COMMON_CONVERSATION_MODE_AUTO = 2, -}; - -enum common_grammar_trigger_type { - COMMON_GRAMMAR_TRIGGER_TYPE_TOKEN, - COMMON_GRAMMAR_TRIGGER_TYPE_WORD, - COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, - COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN_FULL, -}; - -struct common_grammar_trigger { - common_grammar_trigger_type type; - std::string value; - llama_token token = LLAMA_TOKEN_NULL; -}; - -enum common_params_sampling_config : uint64_t { - COMMON_PARAMS_SAMPLING_CONFIG_SAMPLERS = 1 << 0, - COMMON_PARAMS_SAMPLING_CONFIG_TOP_K = 1 << 1, - COMMON_PARAMS_SAMPLING_CONFIG_TOP_P = 1 << 2, - COMMON_PARAMS_SAMPLING_CONFIG_MIN_P = 1 << 3, - COMMON_PARAMS_SAMPLING_CONFIG_XTC_PROBABILITY = 1 << 4, - COMMON_PARAMS_SAMPLING_CONFIG_XTC_THRESHOLD = 1 << 5, - COMMON_PARAMS_SAMPLING_CONFIG_TEMP = 1 << 6, - COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_LAST_N = 1 << 7, - COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_REPEAT = 1 << 8, - COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT = 1 << 9, - COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_TAU = 1 << 10, - COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_ETA = 1 << 11, -}; - -enum common_speculative_type { - COMMON_SPECULATIVE_TYPE_NONE, // no speculative decoding - COMMON_SPECULATIVE_TYPE_DRAFT, // draft model - COMMON_SPECULATIVE_TYPE_EAGLE3, // eagle draft model - COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, // simple self-speculative decoding - COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, // self-speculative decoding with n-gram keys only - COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values - COMMON_SPECULATIVE_TYPE_NGRAM_MOD, - COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, // self-speculative decoding with 3-level n-gram cache - COMMON_SPECULATIVE_TYPE_COUNT // number of types, unknown type -}; - -// Grammar type enumeration -enum common_grammar_type { - COMMON_GRAMMAR_TYPE_NONE, // no grammar set - COMMON_GRAMMAR_TYPE_USER, // user-provided GBNF (--grammar / "grammar" API field) - COMMON_GRAMMAR_TYPE_OUTPUT_FORMAT, // auto-generated from JSON schema (--json-schema / "json_schema" API field) - COMMON_GRAMMAR_TYPE_TOOL_CALLS, // auto-generated by chat template parser for function calling -}; - -// Grammar variant struct with type and grammar string -struct common_grammar { - common_grammar_type type = COMMON_GRAMMAR_TYPE_NONE; - std::string grammar; - - // Default constructor - no grammar - common_grammar() = default; - - // Constructor with type and grammar string - common_grammar(common_grammar_type t, std::string g) : type(t), grammar(std::move(g)) { - GGML_ASSERT(type != COMMON_GRAMMAR_TYPE_NONE || !grammar.empty()); - } - - // Check if a grammar is set - bool empty() const { return type == COMMON_GRAMMAR_TYPE_NONE || grammar.empty(); } -}; - -// Returns the raw grammar string, or empty string if no grammar is set. -inline const std::string & common_grammar_value(const common_grammar & g) { - return g.grammar; -} - -// Returns true when the generation_prompt should be prefilled into the grammar sampler. -// Only output-format and tool-call grammars need prefill; user-supplied grammars must not be prefilled. -inline bool common_grammar_needs_prefill(const common_grammar & g) { - return g.type == COMMON_GRAMMAR_TYPE_OUTPUT_FORMAT - || g.type == COMMON_GRAMMAR_TYPE_TOOL_CALLS; -} - -// sampling parameters -struct common_params_sampling { - uint32_t seed = LLAMA_DEFAULT_SEED; // the seed used to initialize llama_sampler - - int32_t n_prev = 64; // number of previous tokens to remember - int32_t n_probs = 0; // if greater than 0, output the probabilities of top n_probs tokens. - int32_t min_keep = 0; // 0 = disabled, otherwise samplers should return at least min_keep tokens - int32_t top_k = 40; // <= 0 to use vocab size - float top_p = 0.95f; // 1.0 = disabled - float min_p = 0.05f; // 0.0 = disabled - float xtc_probability = 0.00f; // 0.0 = disabled - float xtc_threshold = 0.10f; // > 0.5 disables XTC - float typ_p = 1.00f; // typical_p, 1.0 = disabled - float temp = 0.80f; // <= 0.0 to sample greedily, 0.0 to not output probabilities - float dynatemp_range = 0.00f; // 0.0 = disabled - float dynatemp_exponent = 1.00f; // controls how entropy maps to temperature in dynamic temperature sampler - int32_t penalty_last_n = 64; // last n tokens to penalize (0 = disable penalty, -1 = context size) - float penalty_repeat = 1.00f; // 1.0 = disabled - float penalty_freq = 0.00f; // 0.0 = disabled - float penalty_present = 0.00f; // 0.0 = disabled - float dry_multiplier = 0.0f; // 0.0 = disabled; DRY repetition penalty for tokens extending repetition: - float dry_base = 1.75f; // 0.0 = disabled; multiplier * base ^ (length of sequence before token - allowed length) - int32_t dry_allowed_length = 2; // tokens extending repetitions beyond this receive penalty - int32_t dry_penalty_last_n = -1; // how many tokens to scan for repetitions (0 = disable penalty, -1 = context size) - float adaptive_target = -1.0f; // select tokens near this probability (valid range 0.0 to 1.0; negative = disabled) - float adaptive_decay = 0.90f; // EMA decay for adaptation; history ≈ 1/(1-decay) tokens (0.0 - 0.99) - int32_t mirostat = 0; // 0 = disabled, 1 = mirostat, 2 = mirostat 2.0 - float top_n_sigma = -1.00f; // -1.0 = disabled - float mirostat_tau = 5.00f; // target entropy - float mirostat_eta = 0.10f; // learning rate - bool ignore_eos = false; - bool no_perf = false; // disable performance metrics - bool timing_per_token = false; - - uint64_t user_sampling_config = 0; // bitfield to track user-specified samplers - - std::vector dry_sequence_breakers = {"\n", ":", "\"", "*"}; // default sequence breakers for DRY - - std::vector samplers = { - COMMON_SAMPLER_TYPE_PENALTIES, - COMMON_SAMPLER_TYPE_DRY, - COMMON_SAMPLER_TYPE_TOP_N_SIGMA, - COMMON_SAMPLER_TYPE_TOP_K, - COMMON_SAMPLER_TYPE_TYPICAL_P, - COMMON_SAMPLER_TYPE_TOP_P, - COMMON_SAMPLER_TYPE_MIN_P, - COMMON_SAMPLER_TYPE_XTC, - COMMON_SAMPLER_TYPE_TEMPERATURE, - }; - - common_grammar grammar; // optional grammar constraint (user / output-format / tool-calls) - bool grammar_lazy = false; - std::vector grammar_triggers; // optional triggers (for lazy grammars) - std::set preserved_tokens; - - std::vector logit_bias; // logit biases to apply - std::vector logit_bias_eog; // pre-calculated logit biases for EOG tokens - - // The assistant generation prompt already prefilled into the prompt. - // Fed to the grammar sampler (to advance past pre-existing tokens) and used - // to determine the reasoning budget sampler's initial state. - // Only applied when the grammar is of output-format or tool-calls type. - std::string generation_prompt; - - // reasoning budget sampler parameters - // these are populated by the server/CLI based on chat template params - int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget - std::vector reasoning_budget_start; // start tag token sequence - std::vector reasoning_budget_end; // end tag token sequence - std::vector reasoning_budget_forced; // forced sequence (message + end tag) - std::string reasoning_budget_message; // message injected before end tag when budget exhausted - - bool backend_sampling = false; - - bool has_logit_bias() const { - return !logit_bias.empty(); - } - - // print the parameters into a string - std::string print() const; -}; - -struct common_params_model { - std::string path = ""; // model local path // NOLINT - std::string url = ""; // model url to download // NOLINT - std::string hf_repo = ""; // HF repo // NOLINT - std::string hf_file = ""; // HF file // NOLINT - std::string docker_repo = ""; // Docker repo // NOLINT - std::string name = ""; // in format /[:] (tag is optional) // NOLINT -}; - -struct common_ngram_mod; - -// draft-model-based speculative decoding parameters -struct common_params_speculative_draft { - int32_t n_max = 16; // maximum number of tokens to draft during speculative decoding - int32_t n_min = 0; // minimum number of draft tokens to use for speculative decoding - - float p_split = 0.1f; // speculative decoding split probability - float p_min = 0.75f; // minimum speculative decoding probability (greedy) - - common_params_model mparams; - - llama_model * model = nullptr; // a llama_model that can be shared by multiple speculative contexts - - llama_context_params cparams; // these are the parameters for the draft llama_context - - int32_t n_ctx = 0; // draft context size - int32_t n_gpu_layers = -1; // number of layers to store in VRAM for the draft model (-1 - use default) - - ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K - ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V - - common_cpu_params cpuparams; - common_cpu_params cpuparams_batch; - - std::vector devices; // devices to use for offloading - - std::vector> replacements; // main to speculative model replacements - std::vector tensor_buft_overrides; -}; - -struct common_params_speculative_ngram_mod { - int32_t n_match = 24; - - int32_t n_max = 64; - int32_t n_min = 48; - - // shared instance of the ngram container for all speculative decoding contexts - std::shared_ptr obj; -}; - -struct common_params_speculative_ngram_map { - uint16_t size_n = 12; // ngram size for lookup - uint16_t size_m = 48; // mgram size for speculative tokens - uint16_t min_hits = 1; // minimum hits at ngram/mgram lookup for mgram to be proposed -}; - -struct common_params_speculative_ngram_cache { - std::string lookup_cache_static; // path of static ngram cache file for lookup decoding - std::string lookup_cache_dynamic; // path of dynamic ngram cache file for lookup decoding -}; - -struct common_params_speculative { - // TODO: become a vector in order to support "chains of speculators" - common_speculative_type type = COMMON_SPECULATIVE_TYPE_NONE; - - common_params_speculative_draft draft; - - common_params_speculative_ngram_mod ngram_mod; - common_params_speculative_ngram_map ngram_simple; - common_params_speculative_ngram_map ngram_map_k; - common_params_speculative_ngram_map ngram_map_k4v; - - common_params_speculative_ngram_cache ngram_cache; - - bool has_dft() const { - return !draft.mparams.path.empty() || !draft.mparams.hf_repo.empty(); - } -}; - -struct common_params_vocoder { - struct common_params_model model; - - std::string speaker_file; // speaker file path - - bool use_guide_tokens = false; // enable guide tokens to improve TTS accuracy -}; - -struct common_params_diffusion { - int32_t steps = 128; - bool visual_mode = false; - - float eps = 0; // epsilon for timesteps - int32_t block_length = 0; // block length for generation - - int32_t algorithm = 4; // default algorithm: low-confidence - float alg_temp = 0.0f; // algorithm temperature - - float cfg_scale = 0; // classifier-free guidance scale - bool add_gumbel_noise = false; // add gumbel noise to the logits if temp > 0.0 -}; - -// reasoning API response format (not to be confused as chat template's reasoning format) -// only used by server -enum common_reasoning_format { - COMMON_REASONING_FORMAT_NONE, - COMMON_REASONING_FORMAT_AUTO, // Same as deepseek, using `message.reasoning_content` - COMMON_REASONING_FORMAT_DEEPSEEK_LEGACY, // Extract thinking tag contents and return as `message.reasoning_content`, or leave inline in tags in stream mode - COMMON_REASONING_FORMAT_DEEPSEEK, // Extract thinking tag contents and return as `message.reasoning_content`, including in streaming deltas. - // do not extend this enum unless you absolutely have to - // in most cases, use COMMON_REASONING_FORMAT_AUTO - // see: https://github.com/ggml-org/llama.cpp/pull/15408 -}; - - -struct lr_opt { - float lr0 = 1e-5; // learning rate at first epoch - float lr_min = -1; - float decay_epochs = -1; // if >0, the learning rate starts at lr0 and decays to lr_min after this many epochs - float scale_epoch = 0; - float wd = 0; - unsigned epochs = 2; - - unsigned epoch; // set by optimizer outer (epochs) loop - // learning rate decay - constant LR per epoch only for now - float get_lr(float e) const; - float get_lr() const { return get_lr(epoch); } - // must call after arg parse, before get_lr - void init(); -}; - -struct ggml_opt_optimizer_params common_opt_lr_pars(void * userdata); - -struct common_params { - int32_t n_predict = -1; // max. number of new tokens to predict, -1 == no limit - int32_t n_ctx = 0; // context size, 0 == context the model was trained with - 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) - int32_t n_keep = 0; // number of tokens to keep from initial prompt - int32_t n_chunks = -1; // max number of chunks to process (-1 = unlimited) - int32_t n_parallel = 1; // number of parallel sequences to decode - int32_t n_sequences = 1; // number of sequences to decode - int32_t grp_attn_n = 1; // group-attention factor - int32_t grp_attn_w = 512; // group-attention width - int32_t n_print = -1; // print token count every n tokens (-1 = disabled) - float rope_freq_base = 0.0f; // RoPE base frequency - float rope_freq_scale = 0.0f; // RoPE frequency scaling factor - float yarn_ext_factor = -1.0f; // YaRN extrapolation mix factor - float yarn_attn_factor = -1.0f; // YaRN magnitude scaling factor - float yarn_beta_fast = -1.0f; // YaRN low correction dim - float yarn_beta_slow = -1.0f; // YaRN high correction dim - int32_t yarn_orig_ctx = 0; // YaRN original context length - - // offload params - std::vector devices; // devices to use for offloading - - int32_t n_gpu_layers = -1; // number of layers to store in VRAM, -1 is auto, <= -2 is all - int32_t main_gpu = 0; // the GPU that is used for scratch and small tensors - float tensor_split[128] = {0}; // how split tensors should be distributed across GPUs - bool fit_params = true; // whether to fit unset model/context parameters to free device memory - bool fit_params_print = false; // print the estimated required memory to run the model - int32_t fit_params_min_ctx = 4096; // minimum context size to set when trying to reduce memory use - - // margin per device in bytes for fitting parameters to free memory: - std::vector fit_params_target = std::vector(llama_max_devices(), 1024 * 1024*1024); - - enum llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER; // how to split the model across GPUs - - common_cpu_params cpuparams; - common_cpu_params cpuparams_batch; - - ggml_backend_sched_eval_callback cb_eval = nullptr; - void * cb_eval_user_data = nullptr; - - ggml_numa_strategy numa = GGML_NUMA_STRATEGY_DISABLED; - - enum llama_rope_scaling_type rope_scaling_type = LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED; - enum llama_pooling_type pooling_type = LLAMA_POOLING_TYPE_UNSPECIFIED; // pooling type for embeddings - enum llama_attention_type attention_type = LLAMA_ATTENTION_TYPE_UNSPECIFIED; // attention type for embeddings - enum llama_flash_attn_type flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO; // whether to use Flash Attention - - struct common_params_sampling sampling; - struct common_params_speculative speculative; - struct common_params_vocoder vocoder; - struct common_params_diffusion diffusion; - - struct common_params_model model; - - std::set model_alias; // model aliases // NOLINT - std::set model_tags; // model tags (informational, not used for routing) // NOLINT - std::string hf_token = ""; // HF token // NOLINT - std::string prompt = ""; // NOLINT - std::string system_prompt = ""; // NOLINT - std::string prompt_file = ""; // store the external prompt file name // NOLINT - std::string path_prompt_cache = ""; // path to file for saving/loading prompt eval state // NOLINT - std::string input_prefix = ""; // string to prefix user inputs with // NOLINT - std::string input_suffix = ""; // string to suffix user inputs with // NOLINT - std::string logits_file = ""; // file for saving *all* logits // NOLINT - - // llama-debug specific options - std::string logits_output_dir = "data"; // directory for saving logits output files // NOLINT - bool save_logits = false; // whether to save logits to files // NOLINT - std::vector tensor_filter; // filter tensor names for debug output (regex) // NOLINT - - std::vector in_files; // all input files - std::vector antiprompt; // strings upon which more user input is prompted (a.k.a. reverse prompts) - std::vector kv_overrides; - std::vector tensor_buft_overrides; - - bool lora_init_without_apply = false; // only load lora to memory, but do not apply it to ctx (user can manually apply lora later using llama_adapter_lora_apply) - std::vector lora_adapters; // lora adapter path with user defined scale - - std::vector control_vectors; // control vector with user defined scale - - int32_t verbosity = 3; // LOG_LEVEL_INFO - int32_t control_vector_layer_start = -1; // layer range for control vector - int32_t control_vector_layer_end = -1; // layer range for control vector - bool offline = false; - - int32_t ppl_stride = 0; // stride for perplexity calculations. If left at 0, the pre-existing approach will be used. - int32_t ppl_output_type = 0; // = 0 -> ppl output is as usual, = 1 -> ppl output is num_tokens, ppl, one per line - // (which is more convenient to use for plotting) - // - bool hellaswag = false; // compute HellaSwag score over random tasks from datafile supplied in prompt - size_t hellaswag_tasks = 400; // number of tasks to use when computing the HellaSwag score - - bool winogrande = false; // compute Winogrande score over random tasks from datafile supplied in prompt - size_t winogrande_tasks = 0; // number of tasks to use when computing the Winogrande score. If 0, all tasks will be computed - - bool multiple_choice = false; // compute TruthfulQA score over random tasks from datafile supplied in prompt - size_t multiple_choice_tasks = 0; // number of tasks to use when computing the TruthfulQA score. If 0, all tasks will be computed - - bool kl_divergence = false; // compute KL divergence - - bool check = false; // check rather than generate results for llama-results - - bool usage = false; // print usage - bool completion = false; // print source-able completion script - bool use_color = false; // use color to distinguish generations and inputs - bool special = false; // enable special token output - bool interactive = false; // interactive mode - bool interactive_first = false; // wait for user input immediately - bool prompt_cache_all = false; // save user input and generations to prompt cache - bool prompt_cache_ro = false; // open the prompt cache read-only and do not update it - - bool escape = true; // escape "\n", "\r", "\t", "\'", "\"", and "\\" - bool multiline_input = false; // reverse the usage of `\` - bool simple_io = false; // improves compatibility with subprocesses and limited consoles - bool cont_batching = true; // insert new sequences for decoding on-the-fly - bool no_perf = false; // disable performance metrics - bool show_timings = true; // show timing information on CLI - bool ctx_shift = false; // context shift on infinite text generation - bool swa_full = false; // use full-size SWA cache (https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055) - bool kv_unified = false; // enable unified KV cache - - bool input_prefix_bos = false; // prefix BOS to user inputs, preceding input_prefix - bool use_mmap = true; // enable mmap to use filesystem cache - bool use_direct_io = false; // read from disk without buffering - bool use_mlock = false; // use mlock to keep model in memory - bool verbose_prompt = false; // print prompt tokens before generation - bool display_prompt = true; // print prompt before generation - bool no_kv_offload = false; // disable KV offloading - bool warmup = true; // warmup run - bool check_tensors = false; // validate tensor data - bool no_op_offload = false; // globally disable offload host tensor operations to device - bool no_extra_bufts = false; // disable extra buffer types (used for weight repacking) - bool no_host = false; // bypass host buffer allowing extra buffers to be used - - bool single_turn = false; // single turn chat conversation - - ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K - ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V - - common_conversation_mode conversation_mode = COMMON_CONVERSATION_MODE_AUTO; - - // multimodal models (see tools/mtmd) - struct common_params_model mmproj; - bool mmproj_use_gpu = true; // use GPU for multimodal model - bool no_mmproj = false; // explicitly disable multimodal model - std::vector image; // path to image file(s) - int image_min_tokens = -1; - int image_max_tokens = -1; - - // finetune - struct lr_opt lr; - enum ggml_opt_optimizer_type optimizer = GGML_OPT_OPTIMIZER_TYPE_ADAMW; - float val_split = 0.05f; // fraction of the data used for the validation set - - // embedding - bool embedding = false; // get only sentence embedding - int32_t embd_normalize = 2; // normalisation for embeddings (-1=none, 0=max absolute int16, 1=taxicab, 2=euclidean, >2=p-norm) - std::string embd_out = ""; // empty = default, "array" = [[],[]...], "json" = openai style, "json+" = same "json" + cosine similarity matrix - std::string embd_sep = "\n"; // separator of embeddings - std::string cls_sep = "\t"; // separator of classification sequences - - // server params - int32_t port = 8080; // server listens on this network port - bool reuse_port = false; // allow multiple sockets to bind to the same port - int32_t timeout_read = 600; // http read timeout in seconds - int32_t timeout_write = timeout_read; // http write timeout in seconds - int32_t n_threads_http = -1; // number of threads to process HTTP requests (TODO: support threadpool) - int32_t n_cache_reuse = 0; // min chunk size to reuse from the cache via KV shifting - bool cache_prompt = true; // whether to enable prompt caching - bool cache_idle_slots = true; // save and clear idle slots upon starting a new task - int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot - int32_t checkpoint_every_nt = 8192; // make a checkpoint every n tokens during prefill - int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. - - std::string hostname = "127.0.0.1"; - std::string public_path = ""; // NOLINT - std::string api_prefix = ""; // NOLINT - std::string chat_template = ""; // NOLINT - bool use_jinja = true; // NOLINT - bool enable_chat_template = true; - bool force_pure_content_parser = false; - common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK; - int enable_reasoning = -1; // -1 = auto, 0 = disable, 1 = enable - bool prefill_assistant = true; // if true, any trailing assistant message will be prefilled into the response - int sleep_idle_seconds = -1; // if >0, server will sleep after this many seconds of idle time - - std::vector api_keys; - - std::string ssl_file_key = ""; // NOLINT - std::string ssl_file_cert = ""; // NOLINT - - std::map default_template_kwargs; - - // webui configs - bool webui = true; - bool webui_mcp_proxy = false; - std::string webui_config_json; - - // "advanced" endpoints are disabled by default for better security - bool endpoint_slots = true; - bool endpoint_props = false; // only control POST requests, not GET - bool endpoint_metrics = false; - - // enable built-in tools - std::vector server_tools; - - // router server configs - std::string models_dir = ""; // directory containing models for the router server - std::string models_preset = ""; // directory containing model presets for the router server - int models_max = 4; // maximum number of models to load simultaneously - bool models_autoload = true; // automatically load models when requested via the router server - - bool log_json = false; - - std::string slot_save_path; - std::string media_path; // path to directory for loading media files - - float slot_prompt_similarity = 0.1f; - - // batched-bench params - bool is_pp_shared = false; - bool is_tg_separate = false; - - std::vector n_pp; - std::vector n_tg; - std::vector n_pl; - - // retrieval params - std::vector context_files; // context files to embed - - int32_t chunk_size = 64; // chunk size for context embedding - - std::string chunk_separator = "\n"; // chunk separator for context embedding - - // passkey params - int32_t n_junk = 250; // number of times to repeat the junk text - int32_t i_pos = -1; // position of the passkey in the junk text - - // imatrix params - int32_t n_out_freq = 10; // output the imatrix every n_out_freq iterations - int32_t n_save_freq = 0; // save the imatrix every n_save_freq iterations - int32_t i_chunk = 0; // start processing from this chunk - int8_t imat_dat = 0; // whether the legacy imatrix.dat format should be output (gguf <= 0 < dat) - - bool process_output = false; // collect data for the output tensor - bool compute_ppl = true; // whether to compute perplexity - bool show_statistics = false; // show imatrix statistics per tensor - bool parse_special = false; // whether to parse special tokens during imatrix tokenization - - // cvector-generator params - int n_pca_batch = 100; - int n_pca_iterations = 1000; - dimre_method cvector_dimre_method = DIMRE_METHOD_PCA; - std::string cvector_positive_file = "tools/cvector-generator/positive.txt"; - std::string cvector_negative_file = "tools/cvector-generator/negative.txt"; - - bool spm_infill = false; // suffix/prefix/middle pattern for infill - - // batched-bench params - bool batched_bench_output_jsonl = false; - - // common params - std::string out_file; // output filename for all example programs - // optional callback for model loading progress and cancellation: - // called with a progress value between 0.0 and 1.0. - // return false from callback to abort model loading or true to continue - llama_progress_callback load_progress_callback = NULL; - void * load_progress_callback_user_data = NULL; - bool no_alloc = false; // Don't allocate model buffers -}; - -// call once at the start of a program if it uses libcommon -// initializes the logging system and prints info about the build -void common_init(); - -std::string common_params_get_system_info(const common_params & params); - -bool parse_cpu_range(const std::string & range, bool(&boolmask)[GGML_MAX_N_THREADS]); -bool parse_cpu_mask(const std::string & mask, bool(&boolmask)[GGML_MAX_N_THREADS]); -void postprocess_cpu_params(common_cpu_params & cpuparams, const common_cpu_params * role_model = nullptr); -bool set_process_priority(enum ggml_sched_priority prio); - -// -// String utils -// - -#ifdef __GNUC__ -# if defined(__MINGW32__) && !defined(__clang__) -# define LLAMA_COMMON_ATTRIBUTE_FORMAT(...) __attribute__((format(gnu_printf, __VA_ARGS__))) -# else -# define LLAMA_COMMON_ATTRIBUTE_FORMAT(...) __attribute__((format(printf, __VA_ARGS__))) -# endif -#else -# define LLAMA_COMMON_ATTRIBUTE_FORMAT(...) -#endif - -LLAMA_COMMON_ATTRIBUTE_FORMAT(1, 2) -std::string string_format(const char * fmt, ...); - -std::string string_strip(const std::string & str); -std::string string_get_sortable_timestamp(); - -std::string string_join(const std::vector & values, const std::string & separator); -std::vector string_split(const std::string & str, const std::string & delimiter); -std::string string_repeat(const std::string & str, size_t n); - -void string_replace_all(std::string & s, const std::string & search, const std::string & replace); - -std::string regex_escape(const std::string & s); - -template -static std::vector string_split(const std::string & str, char delim) { - static_assert(!std::is_same::value, "Please use the specialized version for std::string"); - std::vector values; - std::istringstream str_stream(str); - std::string token; - while (std::getline(str_stream, token, delim)) { - T value; - std::istringstream token_stream(token); - token_stream >> value; - values.push_back(value); - } - return values; -} - -template<> -inline std::vector string_split(const std::string & str, char delim) -{ - std::vector parts; - size_t begin_pos = 0; - size_t delim_pos = str.find(delim); - while (delim_pos != std::string::npos) { - std::string part = str.substr(begin_pos, delim_pos - begin_pos); - parts.emplace_back(part); - begin_pos = delim_pos + 1; - delim_pos = str.find(delim, begin_pos); - } - parts.emplace_back(str.substr(begin_pos)); - return parts; -} - -// remove when moving to c++20 -inline bool string_starts_with(std::string_view str, std::string_view prefix) { - return str.size() >= prefix.size() && - str.compare(0, prefix.size(), prefix) == 0; -} - -// remove when moving to c++20 -inline bool string_starts_with(std::string_view str, char prefix) { - return !str.empty() && str.front() == prefix; -} - -// remove when moving to c++20 -inline bool string_ends_with(std::string_view str, std::string_view suffix) { - return str.size() >= suffix.size() && - str.compare(str.size() - suffix.size(), suffix.size(), suffix) == 0; -} - -inline bool string_remove_suffix(std::string & str, std::string_view suffix) { - if (string_ends_with(str, suffix)) { - str.resize(str.size() - suffix.size()); - return true; - } - return false; -} - -inline size_t string_find_partial_stop(std::string_view str, std::string_view stop) { - if (!str.empty() && !stop.empty()) { - const size_t max_len = std::min(str.size(), stop.size()); - const char last_char = str.back(); - for (size_t len = max_len; len > 0; --len) { - if (stop[len - 1] == last_char) { - if (string_ends_with(str, stop.substr(0, len))) { - return str.size() - len; - } - } - } - } - return std::string::npos; -} - -bool string_parse_kv_override(const char * data, std::vector & overrides); -void string_process_escapes(std::string & input); - -std::string string_from(bool value); -std::string string_from(const std::vector & values); -std::string string_from(const struct llama_context * ctx, const std::vector & tokens); -std::string string_from(const struct llama_context * ctx, const struct llama_batch & batch); - -bool glob_match(const std::string & pattern, const std::string & str); - -// -// Filesystem utils -// - -bool fs_validate_filename(const std::string & filename, bool allow_subdirs = false); -bool fs_create_directory_with_parents(const std::string & path); -bool fs_is_directory(const std::string & path); - -std::string fs_get_cache_directory(); -std::string fs_get_cache_file(const std::string & filename); - -struct common_file_info { - std::string path; - std::string name; - size_t size = 0; // in bytes - bool is_dir = false; -}; -std::vector fs_list(const std::string & path, bool include_directories); - -// -// TTY utils -// - -// Auto-detect if colors can be enabled based on terminal and environment -bool tty_can_use_colors(); - -// -// Model utils -// - -struct common_sampler; - -// note: defines the model, context, samplers, ets. lifetimes -struct common_init_result { - common_init_result(common_params & params); - ~common_init_result(); - - llama_model * model(); - llama_context * context(); - - common_sampler * sampler(llama_seq_id seq_id); - void reset_samplers(); - - std::vector & lora(); - -private: - struct impl; - std::unique_ptr pimpl; -}; - -using common_init_result_ptr = std::unique_ptr; - -common_init_result_ptr common_init_from_params(common_params & params); - -struct llama_model_params common_model_params_to_llama ( common_params & params); -struct llama_context_params common_context_params_to_llama(const common_params & params); -struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const common_cpu_params & params); - -// clear LoRA adapters from context, then apply new list of adapters -void common_set_adapter_lora(struct llama_context * ctx, std::vector & lora); - -// model endpoint from env -std::string common_get_model_endpoint(); - -// -// Context utils -// - -enum common_context_seq_rm_type { - COMMON_CONTEXT_SEQ_RM_TYPE_NO = 0, // seq_rm not supported (e.g. no memory module) - COMMON_CONTEXT_SEQ_RM_TYPE_PART = 1, // can seq_rm partial sequences - COMMON_CONTEXT_SEQ_RM_TYPE_FULL = 2, // can seq_rm full sequences only -}; - -// check if the llama_context can remove sequences -// note: clears the memory of the context -common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx); - - -// -// Batch utils -// - -void common_batch_clear(struct llama_batch & batch); - -void common_batch_add( - struct llama_batch & batch, - llama_token id, - llama_pos pos, - const std::vector & seq_ids, - bool logits); - -// decodes a single batch of tokens for a prompt and manages session tokens -// -// Note: We save state before the last token so that we can replay it to ensure -// compatibility with all memory types. Recurrent/hybrid models cannot remove -// tokens from memory, so this approach works across all model architectures. -bool common_prompt_batch_decode( - struct llama_context * ctx, - const std::vector & embd, - int & n_past, - int n_batch, - std::string_view state_path, - bool save_state); - -// replays the last token after loading state to regenerate logits -// used after loading session state to ensure the sampling context has valid logits -bool common_replay_last_token(struct llama_context * ctx, llama_token last_token, int32_t pos); - -// -// Vocab utils -// - -// tokenizes a string into a vector of tokens -// should work similar to Python's `tokenizer.encode` -std::vector common_tokenize( - const struct llama_context * ctx, - const std::string & text, - bool add_special, - bool parse_special = false); - -std::vector common_tokenize( - const struct llama_vocab * vocab, - const std::string & text, - bool add_special, - bool parse_special = false); - -// tokenizes a token into a piece, optionally renders special/control tokens -// should work similar to Python's `tokenizer.id_to_piece` -std::string common_token_to_piece( - const struct llama_context * ctx, - llama_token token, - bool special = true); - -std::string common_token_to_piece( - const struct llama_vocab * vocab, - llama_token token, - bool special = true); - -// detokenizes a vector of tokens into a string -// should work similar to Python's `tokenizer.decode` -// optionally renders special/control tokens -std::string common_detokenize( - const struct llama_context * ctx, - const std::vector & tokens, - bool special = true); - -std::string common_detokenize( - const struct llama_vocab * vocab, - const std::vector & tokens, - bool special = true); - -// -// Embedding utils -// - -// TODO: replace embd_norm with an enum -void common_embd_normalize(const float * inp, float * out, int n, int embd_norm); - -float common_embd_similarity_cos(const float * embd1, const float * embd2, int n); - -// -// Control vector utils -// - -struct common_control_vector_data { - int n_embd; - - // stores data for layers [1, n_layer] where n_layer = data.size() / n_embd - std::vector data; -}; - -struct common_control_vector_load_info { - float strength; - - std::string fname; -}; - -// Load control vectors, scale each by strength, and add them together. -// On error, returns {-1, empty} -common_control_vector_data common_control_vector_load(const std::vector & load_infos); - -// -// Split utils -// - -namespace { - -const char * const LLM_KV_SPLIT_NO = "split.no"; -const char * const LLM_KV_SPLIT_COUNT = "split.count"; -const char * const LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count"; - -} - -// -// MoE utils -// - -const char * const LLM_FFN_EXPS_REGEX = "\\.ffn_(up|down|gate|gate_up)_(ch|)exps"; - -inline std::string llm_ffn_exps_block_regex(int idx) { - return string_format("blk\\.%d%s", idx, LLM_FFN_EXPS_REGEX); -} - -inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() { - return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() }; -} - -// -// training utils -// - -ggml_opt_dataset_t common_opt_dataset_init(struct llama_context * ctx, const std::vector & tokens, int64_t stride); - -// "adamw" or "sgd" (case insensitive) -enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *); diff --git a/common/common.h.inc b/common/common.h.inc new file mode 100644 index 000000000000..c6fc28c9c7b1 --- /dev/null +++ b/common/common.h.inc @@ -0,0 +1,5003 @@ +// Various helper functions and utilities + +#pragma once + +#include "ggml-opt.h.inc" +#include "ggml.h.inc" +#include "llama.h.inc" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) && !defined(_WIN32_WINNT) +#define _WIN32_WINNT 0x0A00 +#endif + +#ifdef _WIN32 +#define DIRECTORY_SEPARATOR '\\' +#else +#define DIRECTORY_SEPARATOR '/' +#endif // _WIN32 + +#define die(msg) do { fputs("error: " msg "\n", stderr); exit(1); } while (0) +#define die_fmt(fmt, ...) do { fprintf(stderr, "error: " fmt "\n", __VA_ARGS__); exit(1); } while (0) + +struct utf8_parse_result { + uint32_t codepoint; // Decoded codepoint (only valid if status == SUCCESS) + size_t bytes_consumed; // How many bytes this codepoint uses (1-4) + enum status { SUCCESS, INCOMPLETE, INVALID } status; + + utf8_parse_result(enum status s, uint32_t cp = 0, size_t bytes = 0) + : codepoint(cp), bytes_consumed(bytes), status(s) {} +}; + +extern "C" { +struct llama_common_utf8_parse_result { + uint32_t codepoint; + size_t bytes_consumed; + int status; +}; + +struct llama_common_unicode_string { + uint8_t * data; + size_t len; +}; + +size_t llama_common_utf8_sequence_length(uint8_t first_byte); +bool llama_common_utf8_is_complete(const uint8_t * data, size_t len); +llama_common_utf8_parse_result llama_common_parse_utf8_codepoint(const uint8_t * data, size_t len, size_t offset); +llama_common_unicode_string llama_common_unicode_cpt_to_utf8(uint32_t cpt); +llama_common_unicode_string llama_common_unicode_cpts_to_utf8_rust(const uint32_t * data, size_t len); +void llama_common_unicode_string_free(llama_common_unicode_string value); +const char * llama_common_jinja_token_type_to_string_rust(int type); +bool llama_common_jinja_is_word_rust(uint8_t c); +bool llama_common_jinja_is_integer_rust(uint8_t c); +llama_common_unicode_string llama_common_http_show_masked_url_rust( + const uint8_t * scheme, size_t scheme_len, + bool has_user, + const uint8_t * host, size_t host_len, + const uint8_t * path, size_t path_len); +llama_common_unicode_string llama_common_peak_source_rust( + const uint8_t * source, size_t source_len, + size_t pos, + size_t max_peak_chars); +llama_common_unicode_string llama_common_fmt_error_with_source_rust( + const uint8_t * tag, size_t tag_len, + const uint8_t * msg, size_t msg_len, + const uint8_t * source, size_t source_len, + size_t pos); +} + +// Determine the expected length of a UTF-8 sequence from its first byte. +// Returns 0 for invalid first bytes. +static inline size_t common_utf8_sequence_length(unsigned char first_byte) { + return llama_common_utf8_sequence_length(first_byte); +} + +// Check if a string ends with a complete UTF-8 sequence. +static inline bool common_utf8_is_complete(const std::string & s) { + return llama_common_utf8_is_complete(reinterpret_cast(s.data()), s.size()); +} + +// Parse a single UTF-8 codepoint from input. +static inline utf8_parse_result common_parse_utf8_codepoint(std::string_view input, size_t offset) { + llama_common_utf8_parse_result result = llama_common_parse_utf8_codepoint( + reinterpret_cast(input.data()), input.size(), offset); + return utf8_parse_result(static_cast(result.status), result.codepoint, result.bytes_consumed); +} + +static inline std::string common_unicode_cpt_to_utf8(uint32_t cpt) { + llama_common_unicode_string result = llama_common_unicode_cpt_to_utf8(cpt); + if (result.data == nullptr) { + throw std::invalid_argument("invalid codepoint"); + } + std::string value(reinterpret_cast(result.data), result.len); + llama_common_unicode_string_free(result); + return value; +} + +static inline std::string common_unicode_cpts_to_utf8(const std::vector & cps) { + llama_common_unicode_string result = llama_common_unicode_cpts_to_utf8_rust(cps.data(), cps.size()); + if (result.data == nullptr) { + throw std::invalid_argument("invalid codepoint"); + } + std::string value(reinterpret_cast(result.data), result.len); + llama_common_unicode_string_free(result); + return value; +} + +struct common_params_model; + +using common_header = std::pair; +using common_header_list = std::vector; + +namespace hf_cache { + +struct hf_file { + std::string path; + std::string url; + std::string local_path; + std::string final_path; + std::string oid; + std::string repo_id; + size_t size = 0; +}; + +using hf_files = std::vector; + +hf_files get_repo_files( + const std::string & repo_id, + const std::string & token +); + +hf_files get_cached_files(const std::string & repo_id = {}); +std::string finalize_file(const hf_file & file); +void migrate_old_cache_to_hf_cache(const std::string & token, bool offline = false); + +} // namespace hf_cache + +struct common_download_progress { + std::string url; + size_t downloaded = 0; + size_t total = 0; + bool cached = false; +}; + +class common_download_callback { +public: + virtual ~common_download_callback() = default; + virtual void on_start(const common_download_progress & p) = 0; + virtual void on_update(const common_download_progress & p) = 0; + virtual void on_done(const common_download_progress & p, bool ok) = 0; + virtual bool is_cancelled() const { return false; } +}; + +struct common_remote_params { + common_header_list headers; + long timeout = 0; + long max_size = 0; +}; + +std::pair> common_remote_get_content(const std::string & url, const common_remote_params & params); +std::pair common_download_split_repo_tag(const std::string & hf_repo_with_tag); + +struct common_cached_model_info { + std::string repo; + std::string tag; + std::string to_string() const { + return repo + ":" + tag; + } +}; + +struct common_download_opts { + std::string bearer_token; + common_header_list headers; + bool offline = false; + common_download_callback * callback = nullptr; +}; + +struct common_download_model_result { + std::string model_path; + std::string mmproj_path; +}; + +common_download_model_result common_download_model( + const common_params_model & model, + const common_download_opts & opts = {}, + bool download_mmproj = false +); + +std::vector common_list_cached_models(); +int common_download_file_single(const std::string & url, + const std::string & path, + const common_download_opts & opts = {}, + bool skip_etag = false); +std::string common_docker_resolve_model(const std::string & docker); + +#define LOG_CLR_TO_EOL "\033[K\r" +#define LOG_COL_DEFAULT "\033[0m" +#define LOG_COL_BOLD "\033[1m" +#define LOG_COL_RED "\033[31m" +#define LOG_COL_GREEN "\033[32m" +#define LOG_COL_YELLOW "\033[33m" +#define LOG_COL_BLUE "\033[34m" +#define LOG_COL_MAGENTA "\033[35m" +#define LOG_COL_CYAN "\033[36m" +#define LOG_COL_WHITE "\033[37m" + +#ifndef __GNUC__ +# define LOG_ATTRIBUTE_FORMAT(...) +#elif defined(__MINGW32__) && !defined(__clang__) +# define LOG_ATTRIBUTE_FORMAT(...) __attribute__((format(gnu_printf, __VA_ARGS__))) +#else +# define LOG_ATTRIBUTE_FORMAT(...) __attribute__((format(printf, __VA_ARGS__))) +#endif + +#define LOG_LEVEL_DEBUG 4 +#define LOG_LEVEL_INFO 3 +#define LOG_LEVEL_WARN 2 +#define LOG_LEVEL_ERROR 1 +#define LOG_LEVEL_OUTPUT 0 + +#define LOG_DEFAULT_DEBUG LOG_LEVEL_DEBUG +#define LOG_DEFAULT_LLAMA LOG_LEVEL_INFO + +enum log_colors { + LOG_COLORS_AUTO = -1, + LOG_COLORS_DISABLED = 0, + LOG_COLORS_ENABLED = 1, +}; + +int common_log_get_verbosity_thold(void); +void common_log_set_verbosity_thold(int verbosity); + +void common_log_default_callback(enum ggml_log_level level, const char * text, void * user_data); + +struct common_log; + +struct common_log * common_log_init(); +struct common_log * common_log_main(); + +void common_log_pause (struct common_log * log); +void common_log_resume(struct common_log * log); +void common_log_free (struct common_log * log); + +LOG_ATTRIBUTE_FORMAT(3, 4) +void common_log_add(struct common_log * log, enum ggml_log_level level, const char * fmt, ...); + +void common_log_set_file (struct common_log * log, const char * file); +void common_log_set_colors (struct common_log * log, log_colors colors); +void common_log_set_prefix (struct common_log * log, bool prefix); +void common_log_set_timestamps(struct common_log * log, bool timestamps); +void common_log_flush (struct common_log * log); + +#define LOG_TMPL(level, verbosity, ...) \ + do { \ + if ((verbosity) <= common_log_get_verbosity_thold()) { \ + common_log_add(common_log_main(), (level), __VA_ARGS__); \ + } \ + } while (0) + +#define LOG(...) LOG_TMPL(GGML_LOG_LEVEL_NONE, LOG_LEVEL_OUTPUT, __VA_ARGS__) +#define LOGV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_NONE, verbosity, __VA_ARGS__) + +#define LOG_DBG(...) LOG_TMPL(GGML_LOG_LEVEL_DEBUG, LOG_LEVEL_DEBUG, __VA_ARGS__) +#define LOG_INF(...) LOG_TMPL(GGML_LOG_LEVEL_INFO, LOG_LEVEL_INFO, __VA_ARGS__) +#define LOG_WRN(...) LOG_TMPL(GGML_LOG_LEVEL_WARN, LOG_LEVEL_WARN, __VA_ARGS__) +#define LOG_ERR(...) LOG_TMPL(GGML_LOG_LEVEL_ERROR, LOG_LEVEL_ERROR, __VA_ARGS__) +#define LOG_CNT(...) LOG_TMPL(GGML_LOG_LEVEL_CONT, LOG_LEVEL_INFO, __VA_ARGS__) + +#define LOG_INFV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_INFO, verbosity, __VA_ARGS__) +#define LOG_WRNV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_WARN, verbosity, __VA_ARGS__) +#define LOG_ERRV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_ERROR, verbosity, __VA_ARGS__) +#define LOG_DBGV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_DEBUG, verbosity, __VA_ARGS__) +#define LOG_CNTV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_CONT, verbosity, __VA_ARGS__) + +enum common_params_fit_status { + COMMON_PARAMS_FIT_STATUS_SUCCESS = 0, + COMMON_PARAMS_FIT_STATUS_FAILURE = 1, + COMMON_PARAMS_FIT_STATUS_ERROR = 2, +}; + +enum common_params_fit_status common_fit_params( + const char * path_model, + struct llama_model_params * mparams, + struct llama_context_params * cparams, + float * tensor_split, + struct llama_model_tensor_buft_override * tensor_buft_overrides, + size_t * margins, + uint32_t n_ctx_min, + enum ggml_log_level log_level); + +void common_fit_print( + const char * path_model, + struct llama_model_params * mparams, + struct llama_context_params * cparams); + +void common_memory_breakdown_print(const struct llama_context * ctx); + +std::string json_schema_to_grammar(const nlohmann::ordered_json & schema, + bool force_gbnf = false); + +class common_schema_converter; + +class common_schema_info { + std::unique_ptr impl_; + + public: + common_schema_info(); + ~common_schema_info(); + + common_schema_info(const common_schema_info &) = delete; + common_schema_info & operator=(const common_schema_info &) = delete; + common_schema_info(common_schema_info &&) noexcept; + common_schema_info & operator=(common_schema_info &&) noexcept; + + void resolve_refs(nlohmann::ordered_json & schema); + bool resolves_to_string(const nlohmann::ordered_json & schema); +}; + +struct common_grammar_builder { + std::function add_rule; + std::function add_schema; + std::function resolve_refs; +}; + +struct common_grammar_options { + bool dotall = false; +}; + +std::string gbnf_format_literal(const std::string & literal); + +std::string build_grammar(const std::function & cb, const common_grammar_options & options = {}); + +struct llama_model_deleter { + void operator()(llama_model * model) { llama_model_free(model); } +}; + +struct llama_context_deleter { + void operator()(llama_context * context) { llama_free(context); } +}; + +struct llama_sampler_deleter { + void operator()(llama_sampler * sampler) { llama_sampler_free(sampler); } +}; + +struct common_http_url { + std::string scheme; + std::string user; + std::string password; + std::string host; + int port; + std::string path; +}; + +static inline common_http_url common_http_parse_url(const std::string & url) { + common_http_url parts; + auto scheme_end = url.find("://"); + + if (scheme_end == std::string::npos) { + throw std::runtime_error("invalid URL: no scheme"); + } + parts.scheme = url.substr(0, scheme_end); + + if (parts.scheme != "http" && parts.scheme != "https") { + throw std::runtime_error("unsupported URL scheme: " + parts.scheme); + } + + auto rest = url.substr(scheme_end + 3); + auto at_pos = rest.find('@'); + + if (at_pos != std::string::npos) { + auto auth = rest.substr(0, at_pos); + auto colon_pos = auth.find(':'); + if (colon_pos != std::string::npos) { + parts.user = auth.substr(0, colon_pos); + parts.password = auth.substr(colon_pos + 1); + } else { + parts.user = auth; + } + rest = rest.substr(at_pos + 1); + } + + auto slash_pos = rest.find('/'); + + if (slash_pos != std::string::npos) { + parts.host = rest.substr(0, slash_pos); + parts.path = rest.substr(slash_pos); + } else { + parts.host = rest; + parts.path = "/"; + } + + auto colon_pos = parts.host.find(':'); + + if (colon_pos != std::string::npos) { + parts.port = std::stoi(parts.host.substr(colon_pos + 1)); + parts.host = parts.host.substr(0, colon_pos); + } else if (parts.scheme == "http") { + parts.port = 80; + } else if (parts.scheme == "https") { + parts.port = 443; + } else { + throw std::runtime_error("unsupported URL scheme: " + parts.scheme); + } + + return parts; +} + +static inline std::string common_http_show_masked_url(const common_http_url & parts) { + llama_common_unicode_string value = llama_common_http_show_masked_url_rust( + reinterpret_cast(parts.scheme.data()), parts.scheme.size(), + !parts.user.empty(), + reinterpret_cast(parts.host.data()), parts.host.size(), + reinterpret_cast(parts.path.data()), parts.path.size()); + std::string result(reinterpret_cast(value.data), value.len); + llama_common_unicode_string_free(value); + return result; +} + +struct llama_adapter_lora_deleter { + void operator()(llama_adapter_lora * adapter) { llama_adapter_lora_free(adapter); } +}; + +using llama_model_ptr = std::unique_ptr; +using llama_context_ptr = std::unique_ptr; +using llama_sampler_ptr = std::unique_ptr; +using llama_adapter_lora_ptr = std::unique_ptr; + +struct common_time_meas { + common_time_meas(int64_t & t_acc, bool disable = false); + ~common_time_meas(); + + const int64_t t_start_us; + + int64_t & t_acc; +}; + +struct common_adapter_lora_info { + std::string path; + float scale; + + std::string task_name; + std::string prompt_prefix; + + struct llama_adapter_lora * ptr; +}; + +using llama_tokens = std::vector; + +struct common_control_vector_load_info; + +// +// CPU utils +// + +struct common_cpu_params { + int n_threads = -1; + bool cpumask[GGML_MAX_N_THREADS] = {false}; // CPU affinity mask. + bool mask_valid = false; // Default: any CPU + enum ggml_sched_priority priority = GGML_SCHED_PRIO_NORMAL; // Scheduling prio : (0 - normal, 1 - medium, 2 - high, 3 - realtime) + bool strict_cpu = false; // Use strict CPU placement + uint32_t poll = 50; // Polling (busywait) level (0 - no polling, 100 - mostly polling) +}; + +int32_t common_cpu_get_num_physical_cores(); +int32_t common_cpu_get_num_math(); + +// +// Common params +// + +enum llama_example { + LLAMA_EXAMPLE_BATCHED, + LLAMA_EXAMPLE_DEBUG, + LLAMA_EXAMPLE_COMMON, + LLAMA_EXAMPLE_SPECULATIVE, + LLAMA_EXAMPLE_COMPLETION, + LLAMA_EXAMPLE_CLI, + LLAMA_EXAMPLE_EMBEDDING, + LLAMA_EXAMPLE_PERPLEXITY, + LLAMA_EXAMPLE_RETRIEVAL, + LLAMA_EXAMPLE_PASSKEY, + LLAMA_EXAMPLE_IMATRIX, + LLAMA_EXAMPLE_BENCH, + LLAMA_EXAMPLE_SERVER, + LLAMA_EXAMPLE_CVECTOR_GENERATOR, + LLAMA_EXAMPLE_EXPORT_LORA, + LLAMA_EXAMPLE_MTMD, + LLAMA_EXAMPLE_LOOKUP, + LLAMA_EXAMPLE_PARALLEL, + LLAMA_EXAMPLE_TTS, + LLAMA_EXAMPLE_DIFFUSION, + LLAMA_EXAMPLE_FINETUNE, + LLAMA_EXAMPLE_FIT_PARAMS, + LLAMA_EXAMPLE_RESULTS, + LLAMA_EXAMPLE_EXPORT_GRAPH_OPS, + + LLAMA_EXAMPLE_COUNT, +}; + +enum common_sampler_type { + COMMON_SAMPLER_TYPE_NONE = 0, + COMMON_SAMPLER_TYPE_DRY = 1, + COMMON_SAMPLER_TYPE_TOP_K = 2, + COMMON_SAMPLER_TYPE_TOP_P = 3, + COMMON_SAMPLER_TYPE_MIN_P = 4, + //COMMON_SAMPLER_TYPE_TFS_Z = 5, + COMMON_SAMPLER_TYPE_TYPICAL_P = 6, + COMMON_SAMPLER_TYPE_TEMPERATURE = 7, + COMMON_SAMPLER_TYPE_XTC = 8, + COMMON_SAMPLER_TYPE_INFILL = 9, + COMMON_SAMPLER_TYPE_PENALTIES = 10, + COMMON_SAMPLER_TYPE_TOP_N_SIGMA = 11, + COMMON_SAMPLER_TYPE_ADAPTIVE_P = 12, +}; + +// dimensionality reduction methods, used by cvector-generator +enum dimre_method { + DIMRE_METHOD_PCA, + DIMRE_METHOD_MEAN, +}; + +enum common_conversation_mode { + COMMON_CONVERSATION_MODE_DISABLED = 0, + COMMON_CONVERSATION_MODE_ENABLED = 1, + COMMON_CONVERSATION_MODE_AUTO = 2, +}; + +enum common_grammar_trigger_type { + COMMON_GRAMMAR_TRIGGER_TYPE_TOKEN, + COMMON_GRAMMAR_TRIGGER_TYPE_WORD, + COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, + COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN_FULL, +}; + +struct common_grammar_trigger { + common_grammar_trigger_type type; + std::string value; + llama_token token = LLAMA_TOKEN_NULL; +}; + +enum common_params_sampling_config : uint64_t { + COMMON_PARAMS_SAMPLING_CONFIG_SAMPLERS = 1 << 0, + COMMON_PARAMS_SAMPLING_CONFIG_TOP_K = 1 << 1, + COMMON_PARAMS_SAMPLING_CONFIG_TOP_P = 1 << 2, + COMMON_PARAMS_SAMPLING_CONFIG_MIN_P = 1 << 3, + COMMON_PARAMS_SAMPLING_CONFIG_XTC_PROBABILITY = 1 << 4, + COMMON_PARAMS_SAMPLING_CONFIG_XTC_THRESHOLD = 1 << 5, + COMMON_PARAMS_SAMPLING_CONFIG_TEMP = 1 << 6, + COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_LAST_N = 1 << 7, + COMMON_PARAMS_SAMPLING_CONFIG_PENALTY_REPEAT = 1 << 8, + COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT = 1 << 9, + COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_TAU = 1 << 10, + COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_ETA = 1 << 11, +}; + +enum common_speculative_type { + COMMON_SPECULATIVE_TYPE_NONE, // no speculative decoding + COMMON_SPECULATIVE_TYPE_DRAFT, // draft model + COMMON_SPECULATIVE_TYPE_EAGLE3, // eagle draft model + COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, // simple self-speculative decoding + COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, // self-speculative decoding with n-gram keys only + COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values + COMMON_SPECULATIVE_TYPE_NGRAM_MOD, + COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, // self-speculative decoding with 3-level n-gram cache + COMMON_SPECULATIVE_TYPE_COUNT // number of types, unknown type +}; + +// Grammar type enumeration +enum common_grammar_type { + COMMON_GRAMMAR_TYPE_NONE, // no grammar set + COMMON_GRAMMAR_TYPE_USER, // user-provided GBNF (--grammar / "grammar" API field) + COMMON_GRAMMAR_TYPE_OUTPUT_FORMAT, // auto-generated from JSON schema (--json-schema / "json_schema" API field) + COMMON_GRAMMAR_TYPE_TOOL_CALLS, // auto-generated by chat template parser for function calling +}; + +// Grammar variant struct with type and grammar string +struct common_grammar { + common_grammar_type type = COMMON_GRAMMAR_TYPE_NONE; + std::string grammar; + + // Default constructor - no grammar + common_grammar() = default; + + // Constructor with type and grammar string + common_grammar(common_grammar_type t, std::string g) : type(t), grammar(std::move(g)) { + GGML_ASSERT(type != COMMON_GRAMMAR_TYPE_NONE || !grammar.empty()); + } + + // Check if a grammar is set + bool empty() const { return type == COMMON_GRAMMAR_TYPE_NONE || grammar.empty(); } +}; + +// Returns the raw grammar string, or empty string if no grammar is set. +inline const std::string & common_grammar_value(const common_grammar & g) { + return g.grammar; +} + +// Returns true when the generation_prompt should be prefilled into the grammar sampler. +// Only output-format and tool-call grammars need prefill; user-supplied grammars must not be prefilled. +inline bool common_grammar_needs_prefill(const common_grammar & g) { + return g.type == COMMON_GRAMMAR_TYPE_OUTPUT_FORMAT + || g.type == COMMON_GRAMMAR_TYPE_TOOL_CALLS; +} + +// sampling parameters +struct common_params_sampling { + uint32_t seed = LLAMA_DEFAULT_SEED; // the seed used to initialize llama_sampler + + int32_t n_prev = 64; // number of previous tokens to remember + int32_t n_probs = 0; // if greater than 0, output the probabilities of top n_probs tokens. + int32_t min_keep = 0; // 0 = disabled, otherwise samplers should return at least min_keep tokens + int32_t top_k = 40; // <= 0 to use vocab size + float top_p = 0.95f; // 1.0 = disabled + float min_p = 0.05f; // 0.0 = disabled + float xtc_probability = 0.00f; // 0.0 = disabled + float xtc_threshold = 0.10f; // > 0.5 disables XTC + float typ_p = 1.00f; // typical_p, 1.0 = disabled + float temp = 0.80f; // <= 0.0 to sample greedily, 0.0 to not output probabilities + float dynatemp_range = 0.00f; // 0.0 = disabled + float dynatemp_exponent = 1.00f; // controls how entropy maps to temperature in dynamic temperature sampler + int32_t penalty_last_n = 64; // last n tokens to penalize (0 = disable penalty, -1 = context size) + float penalty_repeat = 1.00f; // 1.0 = disabled + float penalty_freq = 0.00f; // 0.0 = disabled + float penalty_present = 0.00f; // 0.0 = disabled + float dry_multiplier = 0.0f; // 0.0 = disabled; DRY repetition penalty for tokens extending repetition: + float dry_base = 1.75f; // 0.0 = disabled; multiplier * base ^ (length of sequence before token - allowed length) + int32_t dry_allowed_length = 2; // tokens extending repetitions beyond this receive penalty + int32_t dry_penalty_last_n = -1; // how many tokens to scan for repetitions (0 = disable penalty, -1 = context size) + float adaptive_target = -1.0f; // select tokens near this probability (valid range 0.0 to 1.0; negative = disabled) + float adaptive_decay = 0.90f; // EMA decay for adaptation; history ≈ 1/(1-decay) tokens (0.0 - 0.99) + int32_t mirostat = 0; // 0 = disabled, 1 = mirostat, 2 = mirostat 2.0 + float top_n_sigma = -1.00f; // -1.0 = disabled + float mirostat_tau = 5.00f; // target entropy + float mirostat_eta = 0.10f; // learning rate + bool ignore_eos = false; + bool no_perf = false; // disable performance metrics + bool timing_per_token = false; + + uint64_t user_sampling_config = 0; // bitfield to track user-specified samplers + + std::vector dry_sequence_breakers = {"\n", ":", "\"", "*"}; // default sequence breakers for DRY + + std::vector samplers = { + COMMON_SAMPLER_TYPE_PENALTIES, + COMMON_SAMPLER_TYPE_DRY, + COMMON_SAMPLER_TYPE_TOP_N_SIGMA, + COMMON_SAMPLER_TYPE_TOP_K, + COMMON_SAMPLER_TYPE_TYPICAL_P, + COMMON_SAMPLER_TYPE_TOP_P, + COMMON_SAMPLER_TYPE_MIN_P, + COMMON_SAMPLER_TYPE_XTC, + COMMON_SAMPLER_TYPE_TEMPERATURE, + }; + + common_grammar grammar; // optional grammar constraint (user / output-format / tool-calls) + bool grammar_lazy = false; + std::vector grammar_triggers; // optional triggers (for lazy grammars) + std::set preserved_tokens; + + std::vector logit_bias; // logit biases to apply + std::vector logit_bias_eog; // pre-calculated logit biases for EOG tokens + + // The assistant generation prompt already prefilled into the prompt. + // Fed to the grammar sampler (to advance past pre-existing tokens) and used + // to determine the reasoning budget sampler's initial state. + // Only applied when the grammar is of output-format or tool-calls type. + std::string generation_prompt; + + // reasoning budget sampler parameters + // these are populated by the server/CLI based on chat template params + int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget + std::vector reasoning_budget_start; // start tag token sequence + std::vector reasoning_budget_end; // end tag token sequence + std::vector reasoning_budget_forced; // forced sequence (message + end tag) + std::string reasoning_budget_message; // message injected before end tag when budget exhausted + + bool backend_sampling = false; + + bool has_logit_bias() const { + return !logit_bias.empty(); + } + + // print the parameters into a string + std::string print() const; +}; + +struct common_params_model { + std::string path = ""; // model local path // NOLINT + std::string url = ""; // model url to download // NOLINT + std::string hf_repo = ""; // HF repo // NOLINT + std::string hf_file = ""; // HF file // NOLINT + std::string docker_repo = ""; // Docker repo // NOLINT + std::string name = ""; // in format /[:] (tag is optional) // NOLINT +}; + +struct common_ngram_mod; + +// draft-model-based speculative decoding parameters +struct common_params_speculative_draft { + int32_t n_max = 16; // maximum number of tokens to draft during speculative decoding + int32_t n_min = 0; // minimum number of draft tokens to use for speculative decoding + + float p_split = 0.1f; // speculative decoding split probability + float p_min = 0.75f; // minimum speculative decoding probability (greedy) + + common_params_model mparams; + + llama_model * model = nullptr; // a llama_model that can be shared by multiple speculative contexts + + llama_context_params cparams; // these are the parameters for the draft llama_context + + int32_t n_ctx = 0; // draft context size + int32_t n_gpu_layers = -1; // number of layers to store in VRAM for the draft model (-1 - use default) + + ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K + ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V + + common_cpu_params cpuparams; + common_cpu_params cpuparams_batch; + + std::vector devices; // devices to use for offloading + + std::vector> replacements; // main to speculative model replacements + std::vector tensor_buft_overrides; +}; + +struct common_params_speculative_ngram_mod { + int32_t n_match = 24; + + int32_t n_max = 64; + int32_t n_min = 48; + + // shared instance of the ngram container for all speculative decoding contexts + std::shared_ptr obj; +}; + +struct common_params_speculative_ngram_map { + uint16_t size_n = 12; // ngram size for lookup + uint16_t size_m = 48; // mgram size for speculative tokens + uint16_t min_hits = 1; // minimum hits at ngram/mgram lookup for mgram to be proposed +}; + +struct common_params_speculative_ngram_cache { + std::string lookup_cache_static; // path of static ngram cache file for lookup decoding + std::string lookup_cache_dynamic; // path of dynamic ngram cache file for lookup decoding +}; + +struct common_params_speculative { + // TODO: become a vector in order to support "chains of speculators" + common_speculative_type type = COMMON_SPECULATIVE_TYPE_NONE; + + common_params_speculative_draft draft; + + common_params_speculative_ngram_mod ngram_mod; + common_params_speculative_ngram_map ngram_simple; + common_params_speculative_ngram_map ngram_map_k; + common_params_speculative_ngram_map ngram_map_k4v; + + common_params_speculative_ngram_cache ngram_cache; + + bool has_dft() const { + return !draft.mparams.path.empty() || !draft.mparams.hf_repo.empty(); + } +}; + +struct common_speculative; + +std::string common_speculative_type_name_str(); +enum common_speculative_type common_speculative_type_from_name(const std::string & name); +std::string common_speculative_type_to_str(enum common_speculative_type type); + +common_speculative * common_speculative_init( + common_params_speculative & params, + llama_context * ctx_tgt); + +void common_speculative_free(common_speculative * spec); +void common_speculative_begin(common_speculative * spec, const llama_tokens & prompt); + +llama_tokens common_speculative_draft( + common_speculative * spec, + const common_params_speculative & params, + const llama_tokens & prompt, + llama_token id_last); + +void common_speculative_accept(common_speculative * spec, uint16_t n_accepted); + +int32_t common_speculative_n_max(const common_speculative * spec, const common_params_speculative & params); +int32_t common_speculative_n_min(const common_speculative * spec, const common_params_speculative & params); + +void common_speculative_print_stats(const common_speculative * spec); + +struct common_speculative_deleter { + void operator()(common_speculative * s) { common_speculative_free(s); } +}; + +typedef std::unique_ptr common_speculative_ptr; + +struct common_params_vocoder { + struct common_params_model model; + + std::string speaker_file; // speaker file path + + bool use_guide_tokens = false; // enable guide tokens to improve TTS accuracy +}; + +struct common_params_diffusion { + int32_t steps = 128; + bool visual_mode = false; + + float eps = 0; // epsilon for timesteps + int32_t block_length = 0; // block length for generation + + int32_t algorithm = 4; // default algorithm: low-confidence + float alg_temp = 0.0f; // algorithm temperature + + float cfg_scale = 0; // classifier-free guidance scale + bool add_gumbel_noise = false; // add gumbel noise to the logits if temp > 0.0 +}; + +// reasoning API response format (not to be confused as chat template's reasoning format) +// only used by server +enum common_reasoning_format { + COMMON_REASONING_FORMAT_NONE, + COMMON_REASONING_FORMAT_AUTO, // Same as deepseek, using `message.reasoning_content` + COMMON_REASONING_FORMAT_DEEPSEEK_LEGACY, // Extract thinking tag contents and return as `message.reasoning_content`, or leave inline in tags in stream mode + COMMON_REASONING_FORMAT_DEEPSEEK, // Extract thinking tag contents and return as `message.reasoning_content`, including in streaming deltas. + // do not extend this enum unless you absolutely have to + // in most cases, use COMMON_REASONING_FORMAT_AUTO + // see: https://github.com/ggml-org/llama.cpp/pull/15408 +}; + + +struct lr_opt { + float lr0 = 1e-5; // learning rate at first epoch + float lr_min = -1; + float decay_epochs = -1; // if >0, the learning rate starts at lr0 and decays to lr_min after this many epochs + float scale_epoch = 0; + float wd = 0; + unsigned epochs = 2; + + unsigned epoch; // set by optimizer outer (epochs) loop + // learning rate decay - constant LR per epoch only for now + float get_lr(float e) const; + float get_lr() const { return get_lr(epoch); } + // must call after arg parse, before get_lr + void init(); +}; + +struct ggml_opt_optimizer_params common_opt_lr_pars(void * userdata); + +struct common_params { + int32_t n_predict = -1; // max. number of new tokens to predict, -1 == no limit + int32_t n_ctx = 0; // context size, 0 == context the model was trained with + 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) + int32_t n_keep = 0; // number of tokens to keep from initial prompt + int32_t n_chunks = -1; // max number of chunks to process (-1 = unlimited) + int32_t n_parallel = 1; // number of parallel sequences to decode + int32_t n_sequences = 1; // number of sequences to decode + int32_t grp_attn_n = 1; // group-attention factor + int32_t grp_attn_w = 512; // group-attention width + int32_t n_print = -1; // print token count every n tokens (-1 = disabled) + float rope_freq_base = 0.0f; // RoPE base frequency + float rope_freq_scale = 0.0f; // RoPE frequency scaling factor + float yarn_ext_factor = -1.0f; // YaRN extrapolation mix factor + float yarn_attn_factor = -1.0f; // YaRN magnitude scaling factor + float yarn_beta_fast = -1.0f; // YaRN low correction dim + float yarn_beta_slow = -1.0f; // YaRN high correction dim + int32_t yarn_orig_ctx = 0; // YaRN original context length + + // offload params + std::vector devices; // devices to use for offloading + + int32_t n_gpu_layers = -1; // number of layers to store in VRAM, -1 is auto, <= -2 is all + int32_t main_gpu = 0; // the GPU that is used for scratch and small tensors + float tensor_split[128] = {0}; // how split tensors should be distributed across GPUs + bool fit_params = true; // whether to fit unset model/context parameters to free device memory + bool fit_params_print = false; // print the estimated required memory to run the model + int32_t fit_params_min_ctx = 4096; // minimum context size to set when trying to reduce memory use + + // margin per device in bytes for fitting parameters to free memory: + std::vector fit_params_target = std::vector(llama_max_devices(), 1024 * 1024*1024); + + enum llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER; // how to split the model across GPUs + + common_cpu_params cpuparams; + common_cpu_params cpuparams_batch; + + ggml_backend_sched_eval_callback cb_eval = nullptr; + void * cb_eval_user_data = nullptr; + + ggml_numa_strategy numa = GGML_NUMA_STRATEGY_DISABLED; + + enum llama_rope_scaling_type rope_scaling_type = LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED; + enum llama_pooling_type pooling_type = LLAMA_POOLING_TYPE_UNSPECIFIED; // pooling type for embeddings + enum llama_attention_type attention_type = LLAMA_ATTENTION_TYPE_UNSPECIFIED; // attention type for embeddings + enum llama_flash_attn_type flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO; // whether to use Flash Attention + + struct common_params_sampling sampling; + struct common_params_speculative speculative; + struct common_params_vocoder vocoder; + struct common_params_diffusion diffusion; + + struct common_params_model model; + + std::set model_alias; // model aliases // NOLINT + std::set model_tags; // model tags (informational, not used for routing) // NOLINT + std::string hf_token = ""; // HF token // NOLINT + std::string prompt = ""; // NOLINT + std::string system_prompt = ""; // NOLINT + std::string prompt_file = ""; // store the external prompt file name // NOLINT + std::string path_prompt_cache = ""; // path to file for saving/loading prompt eval state // NOLINT + std::string input_prefix = ""; // string to prefix user inputs with // NOLINT + std::string input_suffix = ""; // string to suffix user inputs with // NOLINT + std::string logits_file = ""; // file for saving *all* logits // NOLINT + + // llama-debug specific options + std::string logits_output_dir = "data"; // directory for saving logits output files // NOLINT + bool save_logits = false; // whether to save logits to files // NOLINT + std::vector tensor_filter; // filter tensor names for debug output (regex) // NOLINT + + std::vector in_files; // all input files + std::vector antiprompt; // strings upon which more user input is prompted (a.k.a. reverse prompts) + std::vector kv_overrides; + std::vector tensor_buft_overrides; + + bool lora_init_without_apply = false; // only load lora to memory, but do not apply it to ctx (user can manually apply lora later using llama_adapter_lora_apply) + std::vector lora_adapters; // lora adapter path with user defined scale + + std::vector control_vectors; // control vector with user defined scale + + int32_t verbosity = 3; // LOG_LEVEL_INFO + int32_t control_vector_layer_start = -1; // layer range for control vector + int32_t control_vector_layer_end = -1; // layer range for control vector + bool offline = false; + + int32_t ppl_stride = 0; // stride for perplexity calculations. If left at 0, the pre-existing approach will be used. + int32_t ppl_output_type = 0; // = 0 -> ppl output is as usual, = 1 -> ppl output is num_tokens, ppl, one per line + // (which is more convenient to use for plotting) + // + bool hellaswag = false; // compute HellaSwag score over random tasks from datafile supplied in prompt + size_t hellaswag_tasks = 400; // number of tasks to use when computing the HellaSwag score + + bool winogrande = false; // compute Winogrande score over random tasks from datafile supplied in prompt + size_t winogrande_tasks = 0; // number of tasks to use when computing the Winogrande score. If 0, all tasks will be computed + + bool multiple_choice = false; // compute TruthfulQA score over random tasks from datafile supplied in prompt + size_t multiple_choice_tasks = 0; // number of tasks to use when computing the TruthfulQA score. If 0, all tasks will be computed + + bool kl_divergence = false; // compute KL divergence + + bool check = false; // check rather than generate results for llama-results + + bool usage = false; // print usage + bool completion = false; // print source-able completion script + bool use_color = false; // use color to distinguish generations and inputs + bool special = false; // enable special token output + bool interactive = false; // interactive mode + bool interactive_first = false; // wait for user input immediately + bool prompt_cache_all = false; // save user input and generations to prompt cache + bool prompt_cache_ro = false; // open the prompt cache read-only and do not update it + + bool escape = true; // escape "\n", "\r", "\t", "\'", "\"", and "\\" + bool multiline_input = false; // reverse the usage of `\` + bool simple_io = false; // improves compatibility with subprocesses and limited consoles + bool cont_batching = true; // insert new sequences for decoding on-the-fly + bool no_perf = false; // disable performance metrics + bool show_timings = true; // show timing information on CLI + bool ctx_shift = false; // context shift on infinite text generation + bool swa_full = false; // use full-size SWA cache (https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055) + bool kv_unified = false; // enable unified KV cache + + bool input_prefix_bos = false; // prefix BOS to user inputs, preceding input_prefix + bool use_mmap = true; // enable mmap to use filesystem cache + bool use_direct_io = false; // read from disk without buffering + bool use_mlock = false; // use mlock to keep model in memory + bool verbose_prompt = false; // print prompt tokens before generation + bool display_prompt = true; // print prompt before generation + bool no_kv_offload = false; // disable KV offloading + bool warmup = true; // warmup run + bool check_tensors = false; // validate tensor data + bool no_op_offload = false; // globally disable offload host tensor operations to device + bool no_extra_bufts = false; // disable extra buffer types (used for weight repacking) + bool no_host = false; // bypass host buffer allowing extra buffers to be used + + bool single_turn = false; // single turn chat conversation + + ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K + ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V + + common_conversation_mode conversation_mode = COMMON_CONVERSATION_MODE_AUTO; + + // multimodal models (see tools/mtmd) + struct common_params_model mmproj; + bool mmproj_use_gpu = true; // use GPU for multimodal model + bool no_mmproj = false; // explicitly disable multimodal model + std::vector image; // path to image file(s) + int image_min_tokens = -1; + int image_max_tokens = -1; + + // finetune + struct lr_opt lr; + enum ggml_opt_optimizer_type optimizer = GGML_OPT_OPTIMIZER_TYPE_ADAMW; + float val_split = 0.05f; // fraction of the data used for the validation set + + // embedding + bool embedding = false; // get only sentence embedding + int32_t embd_normalize = 2; // normalisation for embeddings (-1=none, 0=max absolute int16, 1=taxicab, 2=euclidean, >2=p-norm) + std::string embd_out = ""; // empty = default, "array" = [[],[]...], "json" = openai style, "json+" = same "json" + cosine similarity matrix + std::string embd_sep = "\n"; // separator of embeddings + std::string cls_sep = "\t"; // separator of classification sequences + + // server params + int32_t port = 8080; // server listens on this network port + bool reuse_port = false; // allow multiple sockets to bind to the same port + int32_t timeout_read = 600; // http read timeout in seconds + int32_t timeout_write = timeout_read; // http write timeout in seconds + int32_t n_threads_http = -1; // number of threads to process HTTP requests (TODO: support threadpool) + int32_t n_cache_reuse = 0; // min chunk size to reuse from the cache via KV shifting + bool cache_prompt = true; // whether to enable prompt caching + bool cache_idle_slots = true; // save and clear idle slots upon starting a new task + int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot + int32_t checkpoint_every_nt = 8192; // make a checkpoint every n tokens during prefill + int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. + + std::string hostname = "127.0.0.1"; + std::string public_path = ""; // NOLINT + std::string api_prefix = ""; // NOLINT + std::string chat_template = ""; // NOLINT + bool use_jinja = true; // NOLINT + bool enable_chat_template = true; + bool force_pure_content_parser = false; + common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK; + int enable_reasoning = -1; // -1 = auto, 0 = disable, 1 = enable + bool prefill_assistant = true; // if true, any trailing assistant message will be prefilled into the response + int sleep_idle_seconds = -1; // if >0, server will sleep after this many seconds of idle time + + std::vector api_keys; + + std::string ssl_file_key = ""; // NOLINT + std::string ssl_file_cert = ""; // NOLINT + + std::map default_template_kwargs; + + // webui configs + bool webui = true; + bool webui_mcp_proxy = false; + std::string webui_config_json; + + // "advanced" endpoints are disabled by default for better security + bool endpoint_slots = true; + bool endpoint_props = false; // only control POST requests, not GET + bool endpoint_metrics = false; + + // enable built-in tools + std::vector server_tools; + + // router server configs + std::string models_dir = ""; // directory containing models for the router server + std::string models_preset = ""; // directory containing model presets for the router server + int models_max = 4; // maximum number of models to load simultaneously + bool models_autoload = true; // automatically load models when requested via the router server + + bool log_json = false; + + std::string slot_save_path; + std::string media_path; // path to directory for loading media files + + float slot_prompt_similarity = 0.1f; + + // batched-bench params + bool is_pp_shared = false; + bool is_tg_separate = false; + + std::vector n_pp; + std::vector n_tg; + std::vector n_pl; + + // retrieval params + std::vector context_files; // context files to embed + + int32_t chunk_size = 64; // chunk size for context embedding + + std::string chunk_separator = "\n"; // chunk separator for context embedding + + // passkey params + int32_t n_junk = 250; // number of times to repeat the junk text + int32_t i_pos = -1; // position of the passkey in the junk text + + // imatrix params + int32_t n_out_freq = 10; // output the imatrix every n_out_freq iterations + int32_t n_save_freq = 0; // save the imatrix every n_save_freq iterations + int32_t i_chunk = 0; // start processing from this chunk + int8_t imat_dat = 0; // whether the legacy imatrix.dat format should be output (gguf <= 0 < dat) + + bool process_output = false; // collect data for the output tensor + bool compute_ppl = true; // whether to compute perplexity + bool show_statistics = false; // show imatrix statistics per tensor + bool parse_special = false; // whether to parse special tokens during imatrix tokenization + + // cvector-generator params + int n_pca_batch = 100; + int n_pca_iterations = 1000; + dimre_method cvector_dimre_method = DIMRE_METHOD_PCA; + std::string cvector_positive_file = "tools/cvector-generator/positive.txt"; + std::string cvector_negative_file = "tools/cvector-generator/negative.txt"; + + bool spm_infill = false; // suffix/prefix/middle pattern for infill + + // batched-bench params + bool batched_bench_output_jsonl = false; + + // common params + std::string out_file; // output filename for all example programs + // optional callback for model loading progress and cancellation: + // called with a progress value between 0.0 and 1.0. + // return false from callback to abort model loading or true to continue + llama_progress_callback load_progress_callback = NULL; + void * load_progress_callback_user_data = NULL; + bool no_alloc = false; // Don't allocate model buffers +}; + +// call once at the start of a program if it uses libcommon +// initializes the logging system and prints info about the build +void common_init(); + +std::string common_params_get_system_info(const common_params & params); + +bool parse_cpu_range(const std::string & range, bool(&boolmask)[GGML_MAX_N_THREADS]); +bool parse_cpu_mask(const std::string & mask, bool(&boolmask)[GGML_MAX_N_THREADS]); +void postprocess_cpu_params(common_cpu_params & cpuparams, const common_cpu_params * role_model = nullptr); +bool set_process_priority(enum ggml_sched_priority prio); + +// +// String utils +// + +#ifdef __GNUC__ +# if defined(__MINGW32__) && !defined(__clang__) +# define LLAMA_COMMON_ATTRIBUTE_FORMAT(...) __attribute__((format(gnu_printf, __VA_ARGS__))) +# else +# define LLAMA_COMMON_ATTRIBUTE_FORMAT(...) __attribute__((format(printf, __VA_ARGS__))) +# endif +#else +# define LLAMA_COMMON_ATTRIBUTE_FORMAT(...) +#endif + +LLAMA_COMMON_ATTRIBUTE_FORMAT(1, 2) +std::string string_format(const char * fmt, ...); + +std::string string_strip(const std::string & str); +std::string string_get_sortable_timestamp(); + +std::string string_join(const std::vector & values, const std::string & separator); +std::vector string_split(const std::string & str, const std::string & delimiter); +std::string string_repeat(const std::string & str, size_t n); + +void string_replace_all(std::string & s, const std::string & search, const std::string & replace); + +std::string regex_escape(const std::string & s); + +template +static std::vector string_split(const std::string & str, char delim) { + static_assert(!std::is_same::value, "Please use the specialized version for std::string"); + std::vector values; + std::istringstream str_stream(str); + std::string token; + while (std::getline(str_stream, token, delim)) { + T value; + std::istringstream token_stream(token); + token_stream >> value; + values.push_back(value); + } + return values; +} + +template<> +inline std::vector string_split(const std::string & str, char delim) +{ + std::vector parts; + size_t begin_pos = 0; + size_t delim_pos = str.find(delim); + while (delim_pos != std::string::npos) { + std::string part = str.substr(begin_pos, delim_pos - begin_pos); + parts.emplace_back(part); + begin_pos = delim_pos + 1; + delim_pos = str.find(delim, begin_pos); + } + parts.emplace_back(str.substr(begin_pos)); + return parts; +} + +// remove when moving to c++20 +inline bool string_starts_with(std::string_view str, std::string_view prefix) { + return str.size() >= prefix.size() && + str.compare(0, prefix.size(), prefix) == 0; +} + +// remove when moving to c++20 +inline bool string_starts_with(std::string_view str, char prefix) { + return !str.empty() && str.front() == prefix; +} + +// remove when moving to c++20 +inline bool string_ends_with(std::string_view str, std::string_view suffix) { + return str.size() >= suffix.size() && + str.compare(str.size() - suffix.size(), suffix.size(), suffix) == 0; +} + +inline bool string_remove_suffix(std::string & str, std::string_view suffix) { + if (string_ends_with(str, suffix)) { + str.resize(str.size() - suffix.size()); + return true; + } + return false; +} + +inline size_t string_find_partial_stop(std::string_view str, std::string_view stop) { + if (!str.empty() && !stop.empty()) { + const size_t max_len = std::min(str.size(), stop.size()); + const char last_char = str.back(); + for (size_t len = max_len; len > 0; --len) { + if (stop[len - 1] == last_char) { + if (string_ends_with(str, stop.substr(0, len))) { + return str.size() - len; + } + } + } + } + return std::string::npos; +} + +bool string_parse_kv_override(const char * data, std::vector & overrides); +void string_process_escapes(std::string & input); + +std::string string_from(bool value); +std::string string_from(const std::vector & values); +std::string string_from(const struct llama_context * ctx, const std::vector & tokens); +std::string string_from(const struct llama_context * ctx, const struct llama_batch & batch); + +bool glob_match(const std::string & pattern, const std::string & str); + +// +// Filesystem utils +// + +bool fs_validate_filename(const std::string & filename, bool allow_subdirs = false); +bool fs_create_directory_with_parents(const std::string & path); +bool fs_is_directory(const std::string & path); + +std::string fs_get_cache_directory(); +std::string fs_get_cache_file(const std::string & filename); + +struct common_file_info { + std::string path; + std::string name; + size_t size = 0; // in bytes + bool is_dir = false; +}; +std::vector fs_list(const std::string & path, bool include_directories); + +// +// TTY utils +// + +// Auto-detect if colors can be enabled based on terminal and environment +bool tty_can_use_colors(); + +// +// Model utils +// + +struct common_sampler; + +// note: defines the model, context, samplers, ets. lifetimes +struct common_init_result { + common_init_result(common_params & params); + ~common_init_result(); + + llama_model * model(); + llama_context * context(); + + common_sampler * sampler(llama_seq_id seq_id); + void reset_samplers(); + + std::vector & lora(); + +private: + struct impl; + std::unique_ptr pimpl; +}; + +using common_init_result_ptr = std::unique_ptr; + +common_init_result_ptr common_init_from_params(common_params & params); + +struct llama_model_params common_model_params_to_llama ( common_params & params); +struct llama_context_params common_context_params_to_llama(const common_params & params); +struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const common_cpu_params & params); + +// clear LoRA adapters from context, then apply new list of adapters +void common_set_adapter_lora(struct llama_context * ctx, std::vector & lora); + +// model endpoint from env +std::string common_get_model_endpoint(); + +// +// Context utils +// + +enum common_context_seq_rm_type { + COMMON_CONTEXT_SEQ_RM_TYPE_NO = 0, // seq_rm not supported (e.g. no memory module) + COMMON_CONTEXT_SEQ_RM_TYPE_PART = 1, // can seq_rm partial sequences + COMMON_CONTEXT_SEQ_RM_TYPE_FULL = 2, // can seq_rm full sequences only +}; + +// check if the llama_context can remove sequences +// note: clears the memory of the context +common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx); + + +// +// Batch utils +// + +void common_batch_clear(struct llama_batch & batch); + +void common_batch_add( + struct llama_batch & batch, + llama_token id, + llama_pos pos, + const std::vector & seq_ids, + bool logits); + +// decodes a single batch of tokens for a prompt and manages session tokens +// +// Note: We save state before the last token so that we can replay it to ensure +// compatibility with all memory types. Recurrent/hybrid models cannot remove +// tokens from memory, so this approach works across all model architectures. +bool common_prompt_batch_decode( + struct llama_context * ctx, + const std::vector & embd, + int & n_past, + int n_batch, + std::string_view state_path, + bool save_state); + +// replays the last token after loading state to regenerate logits +// used after loading session state to ensure the sampling context has valid logits +bool common_replay_last_token(struct llama_context * ctx, llama_token last_token, int32_t pos); + +// +// Vocab utils +// + +// tokenizes a string into a vector of tokens +// should work similar to Python's `tokenizer.encode` +std::vector common_tokenize( + const struct llama_context * ctx, + const std::string & text, + bool add_special, + bool parse_special = false); + +std::vector common_tokenize( + const struct llama_vocab * vocab, + const std::string & text, + bool add_special, + bool parse_special = false); + +// tokenizes a token into a piece, optionally renders special/control tokens +// should work similar to Python's `tokenizer.id_to_piece` +std::string common_token_to_piece( + const struct llama_context * ctx, + llama_token token, + bool special = true); + +std::string common_token_to_piece( + const struct llama_vocab * vocab, + llama_token token, + bool special = true); + +// detokenizes a vector of tokens into a string +// should work similar to Python's `tokenizer.decode` +// optionally renders special/control tokens +std::string common_detokenize( + const struct llama_context * ctx, + const std::vector & tokens, + bool special = true); + +std::string common_detokenize( + const struct llama_vocab * vocab, + const std::vector & tokens, + bool special = true); + +// +// Embedding utils +// + +// TODO: replace embd_norm with an enum +void common_embd_normalize(const float * inp, float * out, int n, int embd_norm); + +float common_embd_similarity_cos(const float * embd1, const float * embd2, int n); + +// +// Control vector utils +// + +struct common_control_vector_data { + int n_embd; + + // stores data for layers [1, n_layer] where n_layer = data.size() / n_embd + std::vector data; +}; + +struct common_control_vector_load_info { + float strength; + + std::string fname; +}; + +// Load control vectors, scale each by strength, and add them together. +// On error, returns {-1, empty} +common_control_vector_data common_control_vector_load(const std::vector & load_infos); + +// +// Split utils +// + +namespace { + +const char * const LLM_KV_SPLIT_NO = "split.no"; +const char * const LLM_KV_SPLIT_COUNT = "split.count"; +const char * const LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count"; + +} + +// +// MoE utils +// + +const char * const LLM_FFN_EXPS_REGEX = "\\.ffn_(up|down|gate|gate_up)_(ch|)exps"; + +inline std::string llm_ffn_exps_block_regex(int idx) { + return string_format("blk\\.%d%s", idx, LLM_FFN_EXPS_REGEX); +} + +inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() { + return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() }; +} + +// +// training utils +// + +ggml_opt_dataset_t common_opt_dataset_init(struct llama_context * ctx, const std::vector & tokens, int64_t stride); + +// "adamw" or "sgd" (case insensitive) +enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *); + +// common_sampler extends llama_sampler with additional functionality: +// +// - grammar support +// - custom sampler logic based on the parameters +// - history of the last accepted tokens +// - performance metrics +// +// This goal is to have a common implementation of the sampling logic shared across the examples. +// For example, depending on the temperature, the sampling chain can be very simple (greedy) or more +// complex (top-k, top-p, etc). +// +// Another example is related to the grammar. In general, the grammar constraints applied on the full +// vocabulary can be very taxing. To improve performance, the grammar can be applied only to the sampled +// token in order to verify if it fits the grammar. And only if the token doesn't fit the grammar, the +// grammar constraints are applied to the full vocabulary and the token is resampled. +// +// The common_sampler also maintains a container with the last accepted tokens. In the future, this can +// be moved into the core llama library. +// +// For convenience, the common_sampler also maintains a container with the current candidate tokens. +// This can be used to access the probabilities of the rest of the non-sampled tokens. +// +// TODO: measure grammar performance +// + +struct common_sampler; + +// llama_sampler API overloads + +// note: can mutate params in some cases +struct common_sampler * common_sampler_init(const struct llama_model * model, struct common_params_sampling & params); + +void common_sampler_free(struct common_sampler * gsmpl); + +// if is_generated is true, the token is accepted by the sampling chain, the reasoning budget sampler, and the grammar sampler +void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, bool is_generated); +void common_sampler_reset (struct common_sampler * gsmpl); +struct common_sampler * common_sampler_clone (struct common_sampler * gsmpl); + +// arguments can be nullptr to skip printing +void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl); + +// get the underlying llama_sampler_chain +struct llama_sampler * common_sampler_get(const struct common_sampler * gsmpl); + +// extended sampling implementation: +// +// - set logits +// - apply the configured sampler chain +// - check if the token fits the grammar (if any) +// - if not: resample by first applying the grammar constraints and then sampling again (slower path) +// +// if grammar_first is true, the grammar is applied before the samplers (slower) +// useful in cases where all the resulting candidates (not just the sampled one) must fit the grammar +// +llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_context * ctx, int idx, bool grammar_first = false); + +// generalized version of common_sampler_sample +// +// will cross-reference the sampled tokens with a batch of draft tokens and accept those that match +// if the sampler disagrees at some point, we stop and return the accepted tokens up to now +// +// common_sampler_sample_n(gsmpl, ctx, { idx }, {}); +// +// is equivalent to +// +// common_sampler_sample(gsmpl, ctx, idx); +// common_sampler_accept(gsmpl, token, true); +// +// requires: idxs.size() == draft.size() + 1 +// +// returns at least 1 token, up to idxs.size() +// +std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector & idxs, const llama_tokens & draft, bool grammar_first = false); + +// assume idxs == [ 0, 1, 2, ..., draft.size() ] +std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first = false); + +uint32_t common_sampler_get_seed(const struct common_sampler * gsmpl); + +// helpers + +// access the internal list of current candidate tokens +// if do_sort == true, the candidates are guaranteed to be sorted afterwards (in descending order of probability) +// the .sorted flag of the result indicates whether the returned candidates are sorted +llama_token_data_array * common_sampler_get_candidates(struct common_sampler * gsmpl, bool do_sort); + +// get the last accepted token +llama_token common_sampler_last(const struct common_sampler * gsmpl); + +// print the sampler chain into a string +std::string common_sampler_print(const struct common_sampler * gsmpl); + +// get a string representation of the last accepted tokens +std::string common_sampler_prev_str(common_sampler * gsmpl, llama_context * ctx, int n); + +char common_sampler_type_to_chr(enum common_sampler_type cnstr); +std::string common_sampler_type_to_str(enum common_sampler_type cnstr); + +std::vector common_sampler_types_from_names(const std::vector & names, bool allow_alt_names); +std::vector common_sampler_types_from_chars(const std::string & chars); + +llama_sampler * llama_sampler_init_llg(const llama_vocab * vocab, + const char * grammar_kind, const char * grammar_data); + +struct common_sampler_deleter { + void operator()(common_sampler * s) { common_sampler_free(s); } +}; + +typedef std::unique_ptr common_sampler_ptr; + +// pseudo-env variable to identify preset-only arguments +#define COMMON_ARG_PRESET_LOAD_ON_STARTUP "__PRESET_LOAD_ON_STARTUP" +#define COMMON_ARG_PRESET_STOP_TIMEOUT "__PRESET_STOP_TIMEOUT" + +// +// CLI argument parsing +// + +struct common_arg { + std::set examples = {LLAMA_EXAMPLE_COMMON}; + std::set excludes = {}; + std::vector args; + std::vector args_neg; // for negated args like --no-xxx + const char * value_hint = nullptr; // help text or example for arg value + const char * value_hint_2 = nullptr; // for second arg value + const char * env = nullptr; + std::string help; + bool is_sampling = false; // is current arg a sampling param? + bool is_spec = false; // is current arg a speculative decoding param? + bool is_preset_only = false; // is current arg preset-only (not treated as CLI arg) + void (*handler_void) (common_params & params) = nullptr; + void (*handler_string) (common_params & params, const std::string &) = nullptr; + void (*handler_str_str)(common_params & params, const std::string &, const std::string &) = nullptr; + void (*handler_int) (common_params & params, int) = nullptr; + void (*handler_bool) (common_params & params, bool) = nullptr; + + common_arg() = default; + + common_arg( + const std::initializer_list & args, + const char * value_hint, + const std::string & help, + void (*handler)(common_params & params, const std::string &) + ) : args(args), value_hint(value_hint), help(help), handler_string(handler) {} + + common_arg( + const std::initializer_list & args, + const char * value_hint, + const std::string & help, + void (*handler)(common_params & params, int) + ) : args(args), value_hint(value_hint), help(help), handler_int(handler) {} + + common_arg( + const std::initializer_list & args, + const std::string & help, + void (*handler)(common_params & params) + ) : args(args), help(help), handler_void(handler) {} + + common_arg( + const std::initializer_list & args, + const std::initializer_list & args_neg, + const std::string & help, + void (*handler)(common_params & params, bool) + ) : args(args), args_neg(args_neg), help(help), handler_bool(handler) {} + + // support 2 values for arg + common_arg( + const std::initializer_list & args, + const char * value_hint, + const char * value_hint_2, + const std::string & help, + void (*handler)(common_params & params, const std::string &, const std::string &) + ) : args(args), value_hint(value_hint), value_hint_2(value_hint_2), help(help), handler_str_str(handler) {} + + common_arg & set_examples(std::initializer_list examples); + common_arg & set_excludes(std::initializer_list excludes); + common_arg & set_env(const char * env); + common_arg & set_sampling(); + common_arg & set_spec(); + common_arg & set_preset_only(); + bool in_example(enum llama_example ex); + bool is_exclude(enum llama_example ex); + bool get_value_from_env(std::string & output) const; + bool has_value_from_env() const; + std::string to_string() const; + + // for using as key in std::map + bool operator<(const common_arg& other) const { + if (args.empty() || other.args.empty()) { + return false; + } + return strcmp(args[0], other.args[0]) < 0; + } + bool operator==(const common_arg& other) const { + if (args.empty() || other.args.empty()) { + return false; + } + return strcmp(args[0], other.args[0]) == 0; + } + + // get all args and env vars (including negated args/env) + std::vector get_args() const; + std::vector get_env() const; +}; + +namespace common_arg_utils { + bool is_truthy(const std::string & value); + bool is_falsey(const std::string & value); + bool is_autoy(const std::string & value); +} + +struct common_params_context { + enum llama_example ex = LLAMA_EXAMPLE_COMMON; + common_params & params; + std::vector options; + void(*print_usage)(int, char **) = nullptr; + common_params_context(common_params & params) : params(params) {} +}; + +constexpr const char * COMMON_PRESET_DEFAULT_NAME = "default"; + +struct common_preset_context; + +struct common_preset { + std::string name; + std::map options; + + std::vector to_args(const std::string & bin_path = "") const; + std::string to_ini() const; + + void set_option(const common_preset_context & ctx, const std::string & env, const std::string & value); + void unset_option(const std::string & env); + bool get_option(const std::string & env, std::string & value) const; + void merge(const common_preset & other); + void apply_to_params(common_params & params) const; +}; + +using common_presets = std::map; + +struct common_preset_context { + common_params default_params; + common_params_context ctx_params; + std::map key_to_opt; + + bool filter_allowed_keys = false; + std::set allowed_keys; + + common_preset_context(llama_example ex, bool only_remote_allowed = false); + + common_presets load_from_ini(const std::string & path, common_preset & global) const; + common_presets load_from_cache() const; + common_presets load_from_models_dir(const std::string & models_dir) const; + common_preset load_from_args(int argc, char ** argv) const; + common_presets cascade(const common_presets & base, const common_presets & added) const; + common_presets cascade(const common_preset & base, const common_presets & presets) const; +}; + +// parse input arguments from CLI +// if one argument has invalid value, it will automatically display usage of the specific argument (and not the full usage message) +bool common_params_parse(int argc, char ** argv, common_params & params, llama_example ex, void(*print_usage)(int, char **) = nullptr); + +// parse input arguments from CLI into a map +bool common_params_to_map(int argc, char ** argv, llama_example ex, std::map & out_map); + +// populate preset-only arguments +// these arguments are not treated as command line arguments +// see: https://github.com/ggml-org/llama.cpp/issues/18163 +void common_params_add_preset_options(std::vector & args); + +// initialize argument parser context - used by test-arg-parser and preset +common_params_context common_params_parser_init(common_params & params, llama_example ex, void(*print_usage)(int, char **) = nullptr); + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace jinja { + +static void string_replace_all(std::string & s, const std::string & search, const std::string & replace) { + if (search.empty()) { + return; + } + std::string builder; + builder.reserve(s.length()); + size_t pos = 0; + size_t last_pos = 0; + while ((pos = s.find(search, last_pos)) != std::string::npos) { + builder.append(s, last_pos, pos - last_pos); + builder.append(replace); + last_pos = pos + search.length(); + } + builder.append(s, last_pos, std::string::npos); + s = std::move(builder); +} + +// Note: this is a simple hasher, not cryptographically secure, just for hash table usage +struct hasher { + static constexpr auto size_t_digits = sizeof(size_t) * 8; + static constexpr size_t prime = size_t_digits == 64 ? 0x100000001b3 : 0x01000193; + static constexpr size_t seed = size_t_digits == 64 ? 0xcbf29ce484222325 : 0x811c9dc5; + static constexpr auto block_size = sizeof(size_t); // in bytes; allowing the compiler to vectorize the computation + + static_assert(size_t_digits == 64 || size_t_digits == 32); + static_assert(block_size == 8 || block_size == 4); + + uint8_t buffer[block_size]; + size_t idx = 0; // current index in buffer + size_t state = seed; + + hasher() = default; + hasher(const std::type_info & type_inf) noexcept { + const auto type_hash = type_inf.hash_code(); + update(&type_hash, sizeof(type_hash)); + } + + // Properties: + // - update is not associative: update(a).update(b) != update(b).update(a) + // - update(a ~ b) == update(a).update(b) with ~ as concatenation operator --> useful for streaming + // - update("", 0) --> state unchanged with empty input + hasher& update(void const * bytes, size_t len) noexcept { + const uint8_t * c = static_cast(bytes); + if (len == 0) { + return *this; + } + size_t processed = 0; + + // first, fill the existing buffer if it's partial + if (idx > 0) { + size_t to_fill = block_size - idx; + if (to_fill > len) { + to_fill = len; + } + std::memcpy(buffer + idx, c, to_fill); + idx += to_fill; + processed += to_fill; + if (idx == block_size) { + update_block(buffer); + idx = 0; + } + } + + // process full blocks from the remaining input + for (; processed + block_size <= len; processed += block_size) { + update_block(c + processed); + } + + // buffer any remaining bytes + size_t remaining = len - processed; + if (remaining > 0) { + std::memcpy(buffer, c + processed, remaining); + idx = remaining; + } + return *this; + } + + // convenience function for testing only + hasher& update(const std::string & s) noexcept { + return update(s.data(), s.size()); + } + + // finalize and get the hash value + // note: after calling digest, the hasher state is modified, do not call update() again + size_t digest() noexcept { + // if there are remaining bytes in buffer, fill the rest with zeros and process + if (idx > 0) { + for (size_t i = idx; i < block_size; ++i) { + buffer[i] = 0; + } + update_block(buffer); + idx = 0; + } + + return state; + } + +private: + // IMPORTANT: block must have at least block_size bytes + void update_block(const uint8_t * block) noexcept { + size_t blk = static_cast(block[0]) + | (static_cast(block[1]) << 8) + | (static_cast(block[2]) << 16) + | (static_cast(block[3]) << 24); + if constexpr (block_size == 8) { + blk = blk | (static_cast(block[4]) << 32) + | (static_cast(block[5]) << 40) + | (static_cast(block[6]) << 48) + | (static_cast(block[7]) << 56); + } + state ^= blk; + state *= prime; + } +}; + +struct token { + enum type { + eof, // end of source + text, // The text between Jinja statements or expressions + + numeric_literal, // e.g., 123, 1.0 + string_literal, // 'string' + identifier, // Variables, functions, statements, booleans, etc. + equals, // = + open_paren, // ( + close_paren, // ) + open_statement, // {% + close_statement, // %} + open_expression, // {{ + close_expression, // }} + open_square_bracket, // [ + close_square_bracket, // ] + open_curly_bracket, // { + close_curly_bracket, // } + comma, // , + dot, // . + colon, // : + pipe, // | + + call_operator, // () + additive_binary_operator, // + - ~ + multiplicative_binary_operator, // * / % + comparison_binary_operator, // < > <= >= == != + unary_operator, // ! - + + comment, // {# ... #} + }; + type t; + std::string value; + size_t pos; +}; + +struct lexer_result { + std::vector tokens; + std::string source; +}; + +struct lexer { + const std::map escape_chars = { + {'n', '\n'}, + {'t', '\t'}, + {'r', '\r'}, + {'b', '\b'}, + {'f', '\f'}, + {'v', '\v'}, + {'\\', '\\'}, + {'\'', '\''}, + {'\"', '\"'}, + }; + + const std::vector> ordered_mapping_table = { + // Trimmed control sequences + {"{%-", token::open_statement}, + {"-%}", token::close_statement}, + {"{{-", token::open_expression}, + {"-}}", token::close_expression}, + // Control sequences + {"{%", token::open_statement}, + {"%}", token::close_statement}, + {"{{", token::open_expression}, + {"}}", token::close_expression}, + // Single character tokens + {"(", token::open_paren}, + {")", token::close_paren}, + {"{", token::open_curly_bracket}, + {"}", token::close_curly_bracket}, + {"[", token::open_square_bracket}, + {"]", token::close_square_bracket}, + {",", token::comma}, + {".", token::dot}, + {":", token::colon}, + {"|", token::pipe}, + // Comparison operators + {"<=", token::comparison_binary_operator}, + {">=", token::comparison_binary_operator}, + {"==", token::comparison_binary_operator}, + {"!=", token::comparison_binary_operator}, + {"<", token::comparison_binary_operator}, + {">", token::comparison_binary_operator}, + // Arithmetic operators + {"+", token::additive_binary_operator}, + {"-", token::additive_binary_operator}, + {"~", token::additive_binary_operator}, + {"*", token::multiplicative_binary_operator}, + {"/", token::multiplicative_binary_operator}, + {"%", token::multiplicative_binary_operator}, + // Assignment operator + {"=", token::equals}, + }; + + // tokenize the source string into a list of tokens + // may throw lexer_exception on error + lexer_result tokenize(const std::string & source); +}; + +struct lexer_exception : public std::runtime_error { + lexer_exception(const std::string & msg, const std::string & source, size_t pos) + : std::runtime_error([&]() { + const char * tag = "lexer"; + llama_common_unicode_string result = llama_common_fmt_error_with_source_rust( + reinterpret_cast(tag), std::strlen(tag), + reinterpret_cast(msg.data()), msg.size(), + reinterpret_cast(source.data()), source.size(), pos); + std::string value(reinterpret_cast(result.data), result.len); + llama_common_unicode_string_free(result); + return value; + }()) {} +}; + +} // namespace jinja + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace jinja { + +extern "C" { +struct llama_jinja_string { + uint8_t * data; + size_t len; +}; + +void llama_jinja_string_free(llama_jinja_string value); +bool llama_jinja_string_is_uppercase(const uint8_t * data, size_t len); +bool llama_jinja_string_is_lowercase(const uint8_t * data, size_t len); +llama_jinja_string llama_jinja_string_uppercase(const uint8_t * data, size_t len); +llama_jinja_string llama_jinja_string_lowercase(const uint8_t * data, size_t len); +llama_jinja_string llama_jinja_string_capitalize(const uint8_t * data, size_t len); +llama_jinja_string llama_jinja_string_titlecase(const uint8_t * data, size_t len); +llama_jinja_string llama_jinja_string_strip( + const uint8_t * data, size_t len, + bool left, bool right, + const uint8_t * chars, size_t chars_len, bool has_chars); +} + +static inline std::string jinja_string_from_rust(llama_jinja_string value) { + std::string result(reinterpret_cast(value.data), value.len); + llama_jinja_string_free(value); + return result; +} + +// allow differentiate between user input strings and template strings +// transformations should handle this information as follows: +// - one-to-one (e.g., uppercase, lowercase): preserve is_input flag +// - one-to-many (e.g., strip): if input string is marked as is_input, all resulting parts should be marked as is_input +// - many-to-one (e.g., concat): if ALL input parts are marked as is_input, resulting part should be marked as is_input +struct string_part { + bool is_input = false; // may skip parsing special tokens if true + std::string val; + + bool is_uppercase() const { + return llama_jinja_string_is_uppercase(reinterpret_cast(val.data()), val.size()); + } + bool is_lowercase() const { + return llama_jinja_string_is_lowercase(reinterpret_cast(val.data()), val.size()); + } +}; + +struct string { + std::vector parts; + string() = default; + string(const std::string & v, bool user_input = false) { + parts.push_back({user_input, v}); + } + string(int v) { + parts.push_back({false, std::to_string(v)}); + } + string(double v) { + parts.push_back({false, std::to_string(v)}); + } + + // mark all parts as user input + void mark_input() { + for (auto & part : parts) { + part.is_input = true; + } + } + + std::string str() const { + if (parts.size() == 1) { + return parts[0].val; + } + std::ostringstream oss; + for (const auto & part : parts) { + oss << part.val; + } + return oss.str(); + } + size_t length() const { + size_t len = 0; + for (const auto & part : parts) { + len += part.val.length(); + } + return len; + } + void hash_update(hasher & hash) const noexcept { + for (const auto & part : parts) { + hash.update(part.val.data(), part.val.length()); + } + } + bool all_parts_are_input() const { + for (const auto & part : parts) { + if (!part.is_input) { + return false; + } + } + return true; + } + bool is_uppercase() const { + for (const auto & part : parts) { + if (!part.is_uppercase()) { + return false; + } + } + return true; + } + bool is_lowercase() const { + for (const auto & part : parts) { + if (!part.is_lowercase()) { + return false; + } + } + return true; + } + + // mark this string as input if other has ALL parts as input + void mark_input_based_on(const string & other) { + if (other.all_parts_are_input()) { + for (auto & part : parts) { + part.is_input = true; + } + } + } + + string append(const string & other) { + for (const auto & part : other.parts) { + parts.push_back(part); + } + return *this; + } + + // in-place transformations + + string uppercase() { + for (auto & part : parts) { + part.val = jinja_string_from_rust(llama_jinja_string_uppercase(reinterpret_cast(part.val.data()), part.val.size())); + } + return *this; + } + string lowercase() { + for (auto & part : parts) { + part.val = jinja_string_from_rust(llama_jinja_string_lowercase(reinterpret_cast(part.val.data()), part.val.size())); + } + return *this; + } + string capitalize() { + for (auto & part : parts) { + part.val = jinja_string_from_rust(llama_jinja_string_capitalize(reinterpret_cast(part.val.data()), part.val.size())); + } + return *this; + } + string titlecase() { + for (auto & part : parts) { + part.val = jinja_string_from_rust(llama_jinja_string_titlecase(reinterpret_cast(part.val.data()), part.val.size())); + } + return *this; + } + string strip(bool left, bool right, std::optional chars = std::nullopt) { + if (parts.empty()) { + return *this; + } + auto strip_part = [&](std::string & value, bool left_part, bool right_part) { + value = jinja_string_from_rust(llama_jinja_string_strip( + reinterpret_cast(value.data()), value.size(), + left_part, right_part, + chars ? reinterpret_cast((*chars).data()) : nullptr, + chars ? (*chars).size() : 0, + chars.has_value())); + }; + if (left) { + for (size_t i = 0; i < parts.size(); ++i) { + strip_part(parts[i].val, true, false); + if (parts[i].val.empty()) { + parts.erase(parts.begin() + i); + --i; + continue; + } + break; + } + } + if (right) { + for (size_t i = parts.size(); i-- > 0;) { + strip_part(parts[i].val, false, true); + if (parts[i].val.empty()) { + parts.erase(parts.begin() + i); + continue; + } + break; + } + } + return *this; + } +}; + +struct value_t; +using value = std::shared_ptr; + + +// Helper to check the type of a value +template +struct extract_pointee { + using type = T; +}; +template +struct extract_pointee> { + using type = U; +}; +template +bool is_val(const value & ptr) { + using PointeeType = typename extract_pointee::type; + return dynamic_cast(ptr.get()) != nullptr; +} +template +bool is_val(const value_t * ptr) { + using PointeeType = typename extract_pointee::type; + return dynamic_cast(ptr) != nullptr; +} +template +std::shared_ptr::type> mk_val(Args&&... args) { + using PointeeType = typename extract_pointee::type; + return std::make_shared(std::forward(args)...); +} +template +const typename extract_pointee::type * cast_val(const value & ptr) { + using PointeeType = typename extract_pointee::type; + return dynamic_cast(ptr.get()); +} +template +typename extract_pointee::type * cast_val(value & ptr) { + using PointeeType = typename extract_pointee::type; + return dynamic_cast(ptr.get()); +} +// End Helper + + +struct context; // forward declaration + + +// for converting from JSON to jinja values +// example input JSON: +// { +// "messages": [ +// {"role": "user", "content": "Hello!"}, +// {"role": "assistant", "content": "Hi there!"} +// ], +// "bos_token": "", +// "eos_token": "", +// } +// +// to mark strings as user input, wrap them in a special object: +// { +// "messages": [ +// { +// "role": "user", +// "content": {"__input__": "Hello!"} // this string is user input +// }, +// ... +// ], +// } +// +// marking input can be useful for tracking data provenance +// and preventing template injection attacks +// +// Note: T_JSON can be nlohmann::ordered_json +template +void global_from_json(context & ctx, const T_JSON & json_obj, bool mark_input); + +// +// base value type +// + +struct func_args; // function argument values + +using func_hptr = value(const func_args &); +using func_handler = std::function; +using func_builtins = std::map; + +enum value_compare_op { eq, ge, gt, lt, ne }; +bool value_compare(const value & a, const value & b, value_compare_op op); + +struct value_t { + int64_t val_int; + double val_flt; + string val_str; + + std::vector val_arr; + std::vector> val_obj; + + func_handler val_func; + + // only used if ctx.is_get_stats = true + struct stats_t { + bool used = false; + // ops can be builtin calls or operators: "array_access", "object_access" + std::set ops; + // utility to recursively mark value and its children as used + static void mark_used(value & val, bool deep = false); + } stats; + + value_t() = default; + value_t(const value_t &) = default; + virtual ~value_t() = default; + + // Note: only for debugging and error reporting purposes + virtual std::string type() const { return ""; } + + virtual int64_t as_int() const { throw_type_error("is not an int value"); } + virtual double as_float() const { throw_type_error("is not a float value"); } + virtual string as_string() const { throw_type_error("is not a string value"); } + virtual bool as_bool() const { throw_type_error("is not a bool value"); } + virtual const std::vector & as_array() const { throw_type_error("is not an array value"); } + virtual const std::vector> & as_ordered_object() const { throw_type_error("is not an object value"); } + virtual value invoke(const func_args &) const { throw_type_error("is not a function value"); } + virtual bool is_none() const { return false; } + virtual bool is_undefined() const { return false; } + virtual const func_builtins & get_builtins() const { throw_type_error("has no builtins"); } + + virtual bool has_key(const value &) { throw_type_error("is not an object value"); } + virtual void insert(const value & /* key */, const value & /* val */) { throw_type_error("is not an object value"); } + virtual value & at(const value & /* key */, value & /* default_val */) { throw_type_error("is not an object value"); } + virtual value & at(const value & /* key */) { throw_type_error("is not an object value"); } + virtual value & at(const std::string & /* key */, value & /* default_val */) { throw_type_error("is not an object value"); } + virtual value & at(const std::string & /* key */) { throw_type_error("is not an object value"); } + virtual value & at(int64_t /* idx */, value & /* default_val */) { throw_type_error("is not an array value"); } + virtual value & at(int64_t /* idx */) { throw_type_error("is not an array value"); } + + virtual bool is_numeric() const { return false; } + virtual bool is_hashable() const { return false; } + virtual bool is_immutable() const { return true; } + virtual hasher unique_hash() const noexcept = 0; + // TODO: C++20 <=> operator + // NOTE: We are treating == as equivalent (for normal comparisons) and != as strict nonequal (for strict (is) comparisons) + virtual bool operator==(const value_t & other) const { return equivalent(other); } + virtual bool operator!=(const value_t & other) const { return nonequal(other); } + + // Note: only for debugging purposes + virtual std::string as_repr() const { return as_string().str(); } + +private: + [[noreturn]] void throw_type_error(const char* expected) const { + throw std::runtime_error(type() + " " + expected); + } + +protected: + virtual bool equivalent(const value_t &) const = 0; + virtual bool nonequal(const value_t & other) const { return !equivalent(other); } +}; + +// +// utils +// + +const func_builtins & global_builtins(); + +std::string value_to_json(const value & val, int indent = -1, const std::string_view item_sep = ", ", const std::string_view key_sep = ": "); + +// Note: only used for debugging purposes +std::string value_to_string_repr(const value & val); + +struct not_implemented_exception : public std::runtime_error { + not_implemented_exception(const std::string & msg) : std::runtime_error("NotImplemented: " + msg) {} +}; + +struct value_hasher { + size_t operator()(const value & val) const noexcept { + return val->unique_hash().digest(); + } +}; + +struct value_equivalence { + bool operator()(const value & lhs, const value & rhs) const { + return *lhs == *rhs; + } + bool operator()(const std::pair & lhs, const std::pair & rhs) const { + return *(lhs.first) == *(rhs.first) && *(lhs.second) == *(rhs.second); + } +}; + +struct value_equality { + bool operator()(const value & lhs, const value & rhs) const { + return !(*lhs != *rhs); + } +}; + +// +// primitive value types +// + +struct value_int_t : public value_t { + value_int_t(int64_t v) { + val_int = v; + val_flt = static_cast(v); + if (static_cast(val_flt) != v) { + val_flt = v < 0 ? -INFINITY : INFINITY; + } + } + virtual std::string type() const override { return "Integer"; } + virtual int64_t as_int() const override { return val_int; } + virtual double as_float() const override { return val_flt; } + virtual string as_string() const override { return std::to_string(val_int); } + virtual bool as_bool() const override { + return val_int != 0; + } + virtual const func_builtins & get_builtins() const override; + virtual bool is_numeric() const override { return true; } + virtual bool is_hashable() const override { return true; } + virtual hasher unique_hash() const noexcept override { + return hasher(typeid(*this)) + .update(&val_int, sizeof(val_int)) + .update(&val_flt, sizeof(val_flt)); + } +protected: + virtual bool equivalent(const value_t & other) const override { + return other.is_numeric() && val_int == other.val_int && val_flt == other.val_flt; + } + virtual bool nonequal(const value_t & other) const override { + return !(typeid(*this) == typeid(other) && val_int == other.val_int); + } +}; +using value_int = std::shared_ptr; + + +struct value_float_t : public value_t { + value val; + value_float_t(double v) { + val_flt = v; + val_int = std::isfinite(v) ? static_cast(v) : 0; + val = mk_val(val_int); + } + virtual std::string type() const override { return "Float"; } + virtual double as_float() const override { return val_flt; } + virtual int64_t as_int() const override { return val_int; } + virtual string as_string() const override { + std::string out = std::to_string(val_flt); + out.erase(out.find_last_not_of('0') + 1, std::string::npos); // remove trailing zeros + if (out.back() == '.') out.push_back('0'); // leave one zero if no decimals + return out; + } + virtual bool as_bool() const override { + return val_flt != 0.0; + } + virtual const func_builtins & get_builtins() const override; + virtual bool is_numeric() const override { return true; } + virtual bool is_hashable() const override { return true; } + virtual hasher unique_hash() const noexcept override { + if (static_cast(val_int) == val_flt) { + return val->unique_hash(); + } else { + return hasher(typeid(*this)) + .update(&val_int, sizeof(val_int)) + .update(&val_flt, sizeof(val_flt)); + } + } +protected: + virtual bool equivalent(const value_t & other) const override { + return other.is_numeric() && val_int == other.val_int && val_flt == other.val_flt; + } + virtual bool nonequal(const value_t & other) const override { + return !(typeid(*this) == typeid(other) && val_flt == other.val_flt); + } +}; +using value_float = std::shared_ptr; + + +struct value_string_t : public value_t { + value_string_t() { val_str = string(); } + value_string_t(const std::string & v) { val_str = string(v); } + value_string_t(const string & v) { val_str = v; } + virtual std::string type() const override { return "String"; } + virtual string as_string() const override { return val_str; } + virtual std::string as_repr() const override { + std::ostringstream ss; + for (const auto & part : val_str.parts) { + ss << (part.is_input ? "INPUT: " : "TMPL: ") << part.val << "\n"; + } + return ss.str(); + } + virtual bool as_bool() const override { + return val_str.length() > 0; + } + virtual const func_builtins & get_builtins() const override; + virtual bool is_hashable() const override { return true; } + virtual hasher unique_hash() const noexcept override { + const auto type_hash = typeid(*this).hash_code(); + auto hash = hasher(); + hash.update(&type_hash, sizeof(type_hash)); + val_str.hash_update(hash); + return hash; + } + void mark_input() { + val_str.mark_input(); + } +protected: + virtual bool equivalent(const value_t & other) const override { + return typeid(*this) == typeid(other) && val_str.str() == other.val_str.str(); + } +}; +using value_string = std::shared_ptr; + + +struct value_bool_t : public value_t { + value val; + value_bool_t(bool v) { + val_int = static_cast(v); + val_flt = static_cast(v); + val = mk_val(val_int); + } + virtual std::string type() const override { return "Boolean"; } + virtual int64_t as_int() const override { return val_int; } + virtual bool as_bool() const override { return val_int; } + virtual string as_string() const override { return std::string(val_int ? "True" : "False"); } + virtual const func_builtins & get_builtins() const override; + virtual bool is_numeric() const override { return true; } + virtual bool is_hashable() const override { return true; } + virtual hasher unique_hash() const noexcept override { + return val->unique_hash(); + } +protected: + virtual bool equivalent(const value_t & other) const override { + return other.is_numeric() && val_int == other.val_int && val_flt == other.val_flt; + } + virtual bool nonequal(const value_t & other) const override { + return !(typeid(*this) == typeid(other) && val_int == other.val_int); + } +}; +using value_bool = std::shared_ptr; + + +struct value_array_t : public value_t { + value_array_t() = default; + value_array_t(value & v) { + val_arr = v->val_arr; + } + value_array_t(std::vector && arr) { + val_arr = arr; + } + value_array_t(const std::vector & arr) { + val_arr = arr; + } + void reverse() { + if (is_immutable()) { + throw std::runtime_error("Attempting to modify immutable type"); + } + std::reverse(val_arr.begin(), val_arr.end()); + } + void push_back(const value & val) { + if (is_immutable()) { + throw std::runtime_error("Attempting to modify immutable type"); + } + val_arr.push_back(val); + } + void push_back(value && val) { + if (is_immutable()) { + throw std::runtime_error("Attempting to modify immutable type"); + } + val_arr.push_back(std::move(val)); + } + value pop_at(int64_t index) { + if (is_immutable()) { + throw std::runtime_error("Attempting to modify immutable type"); + } + if (index < 0) { + index = static_cast(val_arr.size()) + index; + } + if (index < 0 || index >= static_cast(val_arr.size())) { + throw std::runtime_error("Index " + std::to_string(index) + " out of bounds for array of size " + std::to_string(val_arr.size())); + } + value val = val_arr.at(static_cast(index)); + val_arr.erase(val_arr.begin() + index); + return val; + } + virtual std::string type() const override { return "Array"; } + virtual bool is_immutable() const override { return false; } + virtual const std::vector & as_array() const override { return val_arr; } + virtual string as_string() const override { + const bool immutable = is_immutable(); + std::ostringstream ss; + ss << (immutable ? "(" : "["); + for (size_t i = 0; i < val_arr.size(); i++) { + if (i > 0) ss << ", "; + value val = val_arr.at(i); + ss << value_to_string_repr(val); + } + if (immutable && val_arr.size() == 1) { + ss << ","; + } + ss << (immutable ? ")" : "]"); + return ss.str(); + } + virtual bool as_bool() const override { + return !val_arr.empty(); + } + virtual value & at(int64_t index, value & default_val) override { + if (index < 0) { + index += val_arr.size(); + } + if (index < 0 || static_cast(index) >= val_arr.size()) { + return default_val; + } + return val_arr[index]; + } + virtual value & at(int64_t index) override { + if (index < 0) { + index += val_arr.size(); + } + if (index < 0 || static_cast(index) >= val_arr.size()) { + throw std::runtime_error("Index " + std::to_string(index) + " out of bounds for array of size " + std::to_string(val_arr.size())); + } + return val_arr[index]; + } + virtual const func_builtins & get_builtins() const override; + virtual bool is_hashable() const override { + if (std::all_of(val_arr.begin(), val_arr.end(), [&](auto & val) -> bool { + return val->is_immutable() && val->is_hashable(); + })) { + return true; + } + return false; + } + virtual hasher unique_hash() const noexcept override { + auto hash = hasher(typeid(*this)); + for (const auto & val : val_arr) { + // must use digest to prevent problems from "concatenation" property of hasher + // for ex. hash of [ "ab", "c" ] should be different from [ "a", "bc" ] + const size_t val_hash = val->unique_hash().digest(); + hash.update(&val_hash, sizeof(size_t)); + } + return hash; + } +protected: + virtual bool equivalent(const value_t & other) const override { + return typeid(*this) == typeid(other) && is_hashable() && other.is_hashable() && std::equal(val_arr.begin(), val_arr.end(), other.val_arr.begin(), other.val_arr.end(), value_equivalence()); + } +}; +using value_array = std::shared_ptr; + + +struct value_tuple_t : public value_array_t { + value_tuple_t(value & v) { + val_arr = v->val_arr; + } + value_tuple_t(std::vector && arr) { + val_arr = arr; + } + value_tuple_t(const std::vector & arr) { + val_arr = arr; + } + value_tuple_t(const std::pair & pair) { + val_arr.push_back(pair.first); + val_arr.push_back(pair.second); + } + virtual std::string type() const override { return "Tuple"; } + virtual bool is_immutable() const override { return true; } +}; +using value_tuple = std::shared_ptr; + + +struct value_object_t : public value_t { + std::unordered_map unordered; + bool has_builtins = true; // context and loop objects do not have builtins + value_object_t() = default; + value_object_t(value & v) { + val_obj = v->val_obj; + for (const auto & pair : val_obj) { + unordered[pair.first] = pair.second; + } + } + value_object_t(const std::map & obj) { + for (const auto & pair : obj) { + insert(pair.first, pair.second); + } + } + value_object_t(const std::vector> & obj) { + for (const auto & pair : obj) { + insert(pair.first, pair.second); + } + } + void insert(const std::string & key, const value & val) { + insert(mk_val(key), val); + } + virtual std::string type() const override { return "Object"; } + virtual bool is_immutable() const override { return false; } + virtual const std::vector> & as_ordered_object() const override { return val_obj; } + virtual string as_string() const override { + std::ostringstream ss; + ss << "{"; + for (size_t i = 0; i < val_obj.size(); i++) { + if (i > 0) ss << ", "; + auto & [key, val] = val_obj.at(i); + ss << value_to_string_repr(key) << ": " << value_to_string_repr(val); + } + ss << "}"; + return ss.str(); + } + virtual bool as_bool() const override { + return !unordered.empty(); + } + virtual bool has_key(const value & key) override { + if (!key->is_immutable() || !key->is_hashable()) { + throw std::runtime_error("Object key of unhashable type: " + key->type()); + } + return unordered.find(key) != unordered.end(); + } + virtual void insert(const value & key, const value & val) override { + bool replaced = false; + if (is_immutable()) { + throw std::runtime_error("Attempting to modify immutable type"); + } + if (has_key(key)) { + // if key exists, replace value in ordered list instead of appending + for (auto & pair : val_obj) { + if (*(pair.first) == *key) { + pair.second = val; + replaced = true; + break; + } + } + } + unordered[key] = val; + if (!replaced) { + val_obj.push_back({key, val}); + } + } + virtual value & at(const value & key, value & default_val) override { + if (!has_key(key)) { + return default_val; + } + return unordered.at(key); + } + virtual value & at(const value & key) override { + if (!has_key(key)) { + throw std::runtime_error("Key '" + key->as_string().str() + "' not found in value of type " + type()); + } + return unordered.at(key); + } + virtual value & at(const std::string & key, value & default_val) override { + value key_val = mk_val(key); + return at(key_val, default_val); + } + virtual value & at(const std::string & key) override { + value key_val = mk_val(key); + return at(key_val); + } + virtual const func_builtins & get_builtins() const override; + virtual bool is_hashable() const override { + if (std::all_of(val_obj.begin(), val_obj.end(), [&](auto & pair) -> bool { + const auto & val = pair.second; + return val->is_immutable() && val->is_hashable(); + })) { + return true; + } + return false; + } + virtual hasher unique_hash() const noexcept override { + auto hash = hasher(typeid(*this)); + for (const auto & [key, val] : val_obj) { + // must use digest to prevent problems from "concatenation" property of hasher + // for ex. hash of key="ab", value="c" should be different from key="a", value="bc" + const size_t key_hash = key->unique_hash().digest(); + const size_t val_hash = val->unique_hash().digest(); + hash.update(&key_hash, sizeof(key_hash)); + hash.update(&val_hash, sizeof(val_hash)); + } + return hash; + } +protected: + virtual bool equivalent(const value_t & other) const override { + return typeid(*this) == typeid(other) && is_hashable() && other.is_hashable() && std::equal(val_obj.begin(), val_obj.end(), other.val_obj.begin(), other.val_obj.end(), value_equivalence()); + } +}; +using value_object = std::shared_ptr; + +// +// none and undefined types +// + +struct value_none_t : public value_t { + virtual std::string type() const override { return "None"; } + virtual bool is_none() const override { return true; } + virtual bool as_bool() const override { return false; } + virtual string as_string() const override { return string(type()); } + virtual std::string as_repr() const override { return type(); } + virtual const func_builtins & get_builtins() const override; + virtual bool is_hashable() const override { return true; } + virtual hasher unique_hash() const noexcept override { + return hasher(typeid(*this)); + } +protected: + virtual bool equivalent(const value_t & other) const override { + return typeid(*this) == typeid(other); + } +}; +using value_none = std::shared_ptr; + +struct value_undefined_t : public value_t { + std::string hint; // for debugging, to indicate where undefined came from + value_undefined_t(const std::string & h = "") : hint(h) {} + virtual std::string type() const override { return hint.empty() ? "Undefined" : "Undefined (hint: '" + hint + "')"; } + virtual bool is_undefined() const override { return true; } + virtual bool as_bool() const override { return false; } + virtual std::string as_repr() const override { return type(); } + virtual const func_builtins & get_builtins() const override; + virtual hasher unique_hash() const noexcept override { + return hasher(typeid(*this)); + } +protected: + virtual bool equivalent(const value_t & other) const override { + return is_undefined() == other.is_undefined(); + } +}; +using value_undefined = std::shared_ptr; + +// +// function type +// + +struct func_args { +public: + std::string func_name; // for error messages + context & ctx; + func_args(context & ctx) : ctx(ctx) {} + value get_kwarg(const std::string & key, value default_val) const; + value get_kwarg_or_pos(const std::string & key, size_t pos) const; + value get_pos(size_t pos) const; + value get_pos(size_t pos, value default_val) const; + const std::vector & get_args() const; + size_t count() const { return args.size(); } + void push_back(const value & val); + void push_front(const value & val); + void ensure_count(size_t min, size_t max = 999) const { + size_t n = args.size(); + if (n < min || n > max) { + throw std::runtime_error("Function '" + func_name + "' expected between " + std::to_string(min) + " and " + std::to_string(max) + " arguments, got " + std::to_string(n)); + } + } + template void ensure_val(const value & ptr) const { + if (!is_val(ptr)) { + throw std::runtime_error("Function '" + func_name + "' expected value of type " + std::string(typeid(T).name()) + ", got " + ptr->type()); + } + } + void ensure_count(bool require0, bool require1, bool require2, bool require3) const { + static auto bool_to_int = [](bool b) { return b ? 1 : 0; }; + size_t required = bool_to_int(require0) + bool_to_int(require1) + bool_to_int(require2) + bool_to_int(require3); + ensure_count(required); + } + template void ensure_vals(bool required0 = true) const { + ensure_count(required0, false, false, false); + if (required0 && args.size() > 0) ensure_val(args[0]); + } + template void ensure_vals(bool required0 = true, bool required1 = true) const { + ensure_count(required0, required1, false, false); + if (required0 && args.size() > 0) ensure_val(args[0]); + if (required1 && args.size() > 1) ensure_val(args[1]); + } + template void ensure_vals(bool required0 = true, bool required1 = true, bool required2 = true) const { + ensure_count(required0, required1, required2, false); + if (required0 && args.size() > 0) ensure_val(args[0]); + if (required1 && args.size() > 1) ensure_val(args[1]); + if (required2 && args.size() > 2) ensure_val(args[2]); + } + template void ensure_vals(bool required0 = true, bool required1 = true, bool required2 = true, bool required3 = true) const { + ensure_count(required0, required1, required2, required3); + if (required0 && args.size() > 0) ensure_val(args[0]); + if (required1 && args.size() > 1) ensure_val(args[1]); + if (required2 && args.size() > 2) ensure_val(args[2]); + if (required3 && args.size() > 3) ensure_val(args[3]); + } +private: + std::vector args; +}; + +struct value_func_t : public value_t { + std::string name; + value arg0; // bound "this" argument, if any + value_func_t(const std::string & name, const func_handler & func) : name(name) { + val_func = func; + } + value_func_t(const std::string & name, const func_handler & func, const value & arg_this) : name(name), arg0(arg_this) { + val_func = func; + } + virtual value invoke(const func_args & args) const override { + func_args new_args(args); // copy + new_args.func_name = name; + if (arg0) { + new_args.push_front(arg0); + } + return val_func(new_args); + } + virtual std::string type() const override { return "Function"; } + virtual std::string as_repr() const override { return type() + "<" + name + ">(" + (arg0 ? arg0->as_repr() : "") + ")"; } + virtual bool is_hashable() const override { return false; } + virtual hasher unique_hash() const noexcept override { + // Note: this is unused for now, we don't support function as object keys + // use function pointer as unique identifier + const auto target = val_func.target(); + return hasher(typeid(*this)).update(&target, sizeof(target)); + } +protected: + virtual bool equivalent(const value_t & other) const override { + // Note: this is unused for now, we don't support function as object keys + // compare function pointers + // (val_func == other.val_func does not work as std::function::operator== is only used for nullptr check) + const auto target_this = this->val_func.target(); + const auto target_other = other.val_func.target(); + return typeid(*this) == typeid(other) && target_this == target_other; + } +}; +using value_func = std::shared_ptr; + +// special value for kwarg +struct value_kwarg_t : public value_t { + std::string key; + value val; + value_kwarg_t(const std::string & k, const value & v) : key(k), val(v) {} + virtual std::string type() const override { return "KwArg"; } + virtual std::string as_repr() const override { return type(); } + virtual bool is_hashable() const override { return true; } + virtual hasher unique_hash() const noexcept override { + const auto type_hash = typeid(*this).hash_code(); + auto hash = val->unique_hash(); + hash.update(&type_hash, sizeof(type_hash)) + .update(key.data(), key.size()); + return hash; + } +protected: + virtual bool equivalent(const value_t & other) const override { + const value_kwarg_t & other_val = static_cast(other); + return typeid(*this) == typeid(other) && key == other_val.key && val == other_val.val; + } +}; +using value_kwarg = std::shared_ptr; + + +} // namespace jinja + +#include +#include +#include +#include +#include +#include + +#define JJ_DEBUG(msg, ...) do { if (g_jinja_debug) printf("%s:%-3d : " msg "\n", FILENAME, __LINE__, __VA_ARGS__); } while (0) + +extern bool g_jinja_debug; + +namespace jinja { + +struct statement; +using statement_ptr = std::unique_ptr; +using statements = std::vector; + +// Helpers for dynamic casting and type checking +template +struct extract_pointee_unique { + using type = T; +}; +template +struct extract_pointee_unique> { + using type = U; +}; +template +bool is_stmt(const statement_ptr & ptr) { + return dynamic_cast(ptr.get()) != nullptr; +} +template +T * cast_stmt(statement_ptr & ptr) { + return dynamic_cast(ptr.get()); +} +template +const T * cast_stmt(const statement_ptr & ptr) { + return dynamic_cast(ptr.get()); +} +// End Helpers + + +// not thread-safe +void enable_debug(bool enable); + +struct context { + std::shared_ptr src; // for debugging; use shared_ptr to avoid copying on scope creation + std::time_t current_time; // for functions that need current time + + bool is_get_stats = false; // whether to collect stats + + // src is optional, used for error reporting + context(std::string src = "") : src(std::make_shared(std::move(src))) { + env = mk_val(); + env->has_builtins = false; // context object has no builtins + env->insert("true", mk_val(true)); + env->insert("True", mk_val(true)); + env->insert("false", mk_val(false)); + env->insert("False", mk_val(false)); + env->insert("none", mk_val()); + env->insert("None", mk_val()); + current_time = std::time(nullptr); + } + ~context() = default; + + context(const context & parent) : context() { + // inherit variables (for example, when entering a new scope) + auto & pvar = parent.env->as_ordered_object(); + for (const auto & pair : pvar) { + set_val(pair.first, pair.second); + } + current_time = parent.current_time; + is_get_stats = parent.is_get_stats; + src = parent.src; + } + + value get_val(const std::string & name) { + value default_val = mk_val(name); + return env->at(name, default_val); + } + + void set_val(const std::string & name, const value & val) { + env->insert(name, val); + } + + void set_val(const value & name, const value & val) { + env->insert(name, val); + } + + void print_vars() const { + printf("Context Variables:\n%s\n", value_to_json(env, 2).c_str()); + } + +private: + value_object env; +}; + +/** + * Base class for all nodes in the AST. + */ +struct statement { + size_t pos; // position in source, for debugging + virtual ~statement() = default; + virtual std::string type() const { return "Statement"; } + + // execute_impl must be overridden by derived classes + virtual value execute_impl(context &) { throw_exec_error(); } + // execute is the public method to execute a statement with error handling + value execute(context &); + +private: + [[noreturn]] void throw_exec_error() const { + throw std::runtime_error("cannot exec " + type()); + } +}; + +// Type Checking Utilities + +template +static void chk_type(const statement_ptr & ptr) { + if (!ptr) return; // Allow null for optional fields + assert(dynamic_cast(ptr.get()) != nullptr); +} + +template +static void chk_type(const statement_ptr & ptr) { + if (!ptr) return; + assert(dynamic_cast(ptr.get()) != nullptr || dynamic_cast(ptr.get()) != nullptr); +} + +// Base Types + +/** + * Expressions will result in a value at runtime (unlike statements). + */ +struct expression : public statement { + std::string type() const override { return "Expression"; } +}; + +// Statements + +struct program : public statement { + statements body; + + program() = default; + explicit program(statements && body) : body(std::move(body)) {} + std::string type() const override { return "Program"; } + [[noreturn]] value execute_impl(context &) override { + throw std::runtime_error("Cannot execute program directly, use jinja::runtime instead"); + } +}; + +struct if_statement : public statement { + statement_ptr test; + statements body; + statements alternate; + + if_statement(statement_ptr && test, statements && body, statements && alternate) + : test(std::move(test)), body(std::move(body)), alternate(std::move(alternate)) { + chk_type(this->test); + } + + std::string type() const override { return "If"; } + value execute_impl(context & ctx) override; +}; + +struct identifier; +struct tuple_literal; + +/** + * Loop over each item in a sequence + * https://jinja.palletsprojects.com/en/3.0.x/templates/#for + */ +struct for_statement : public statement { + statement_ptr loopvar; // Identifier | TupleLiteral + statement_ptr iterable; + statements body; + statements default_block; // if no iteration took place + + for_statement(statement_ptr && loopvar, statement_ptr && iterable, statements && body, statements && default_block) + : loopvar(std::move(loopvar)), iterable(std::move(iterable)), + body(std::move(body)), default_block(std::move(default_block)) { + chk_type(this->loopvar); + chk_type(this->iterable); + } + + std::string type() const override { return "For"; } + value execute_impl(context & ctx) override; +}; + +struct break_statement : public statement { + std::string type() const override { return "Break"; } + + struct signal : public std::exception { + const char* what() const noexcept override { + return "Break statement executed"; + } + }; + + [[noreturn]] value execute_impl(context &) override { + throw break_statement::signal(); + } +}; + +struct continue_statement : public statement { + std::string type() const override { return "Continue"; } + + struct signal : public std::exception { + const char* what() const noexcept override { + return "Continue statement executed"; + } + }; + + [[noreturn]] value execute_impl(context &) override { + throw continue_statement::signal(); + } +}; + +// do nothing +struct noop_statement : public statement { + std::string type() const override { return "Noop"; } + value execute_impl(context &) override { + return mk_val(); + } +}; + +struct set_statement : public statement { + statement_ptr assignee; + statement_ptr val; + statements body; + + set_statement(statement_ptr && assignee, statement_ptr && value, statements && body) + : assignee(std::move(assignee)), val(std::move(value)), body(std::move(body)) { + chk_type(this->assignee); + chk_type(this->val); + } + + std::string type() const override { return "Set"; } + value execute_impl(context & ctx) override; +}; + +struct macro_statement : public statement { + statement_ptr name; + statements args; + statements body; + + macro_statement(statement_ptr && name, statements && args, statements && body) + : name(std::move(name)), args(std::move(args)), body(std::move(body)) { + chk_type(this->name); + for (const auto& arg : this->args) chk_type(arg); + } + + std::string type() const override { return "Macro"; } + value execute_impl(context & ctx) override; +}; + +struct comment_statement : public statement { + std::string val; + explicit comment_statement(const std::string & v) : val(v) {} + std::string type() const override { return "Comment"; } + value execute_impl(context &) override { + return mk_val(); + } +}; + +// Expressions + +// Represents an omitted expression in a computed member, e.g. `a[]`. +struct blank_expression : public expression { + std::string type() const override { return "BlankExpression"; } + value execute_impl(context &) override { + return mk_val(); + } +}; + +struct member_expression : public expression { + statement_ptr object; + statement_ptr property; + bool computed; // true if obj[expr] and false if obj.prop + + member_expression(statement_ptr && object, statement_ptr && property, bool computed) + : object(std::move(object)), property(std::move(property)), computed(computed) { + chk_type(this->object); + chk_type(this->property); + } + std::string type() const override { return "MemberExpression"; } + value execute_impl(context & ctx) override; +}; + +struct call_expression : public expression { + statement_ptr callee; + statements args; + + call_expression(statement_ptr && callee, statements && args) + : callee(std::move(callee)), args(std::move(args)) { + chk_type(this->callee); + for (const auto& arg : this->args) chk_type(arg); + } + std::string type() const override { return "CallExpression"; } + value execute_impl(context & ctx) override; +}; + +/** + * Represents a user-defined variable or symbol in the template. + */ +struct identifier : public expression { + std::string val; + explicit identifier(const std::string & val) : val(val) {} + std::string type() const override { return "Identifier"; } + value execute_impl(context & ctx) override; +}; + +// Literals + +struct integer_literal : public expression { + int64_t val; + explicit integer_literal(int64_t val) : val(val) {} + std::string type() const override { return "IntegerLiteral"; } + value execute_impl(context &) override { + return mk_val(val); + } +}; + +struct float_literal : public expression { + double val; + explicit float_literal(double val) : val(val) {} + std::string type() const override { return "FloatLiteral"; } + value execute_impl(context &) override { + return mk_val(val); + } +}; + +struct string_literal : public expression { + std::string val; + explicit string_literal(const std::string & val) : val(val) {} + std::string type() const override { return "StringLiteral"; } + value execute_impl(context &) override { + return mk_val(val); + } +}; + +struct array_literal : public expression { + statements val; + explicit array_literal(statements && val) : val(std::move(val)) { + for (const auto& item : this->val) chk_type(item); + } + std::string type() const override { return "ArrayLiteral"; } + value execute_impl(context & ctx) override { + auto arr = mk_val(); + for (const auto & item_stmt : val) { + arr->push_back(item_stmt->execute(ctx)); + } + return arr; + } +}; + +struct tuple_literal : public expression { + statements val; + explicit tuple_literal(statements && val) : val(std::move(val)) { + for (const auto& item : this->val) chk_type(item); + } + std::string type() const override { return "TupleLiteral"; } + value execute_impl(context & ctx) override { + auto arr = mk_val(); + for (const auto & item_stmt : val) { + arr->push_back(item_stmt->execute(ctx)); + } + return mk_val(std::move(arr->as_array())); + } +}; + +struct object_literal : public expression { + std::vector> val; + explicit object_literal(std::vector> && val) + : val(std::move(val)) { + for (const auto & pair : this->val) { + chk_type(pair.first); + chk_type(pair.second); + } + } + std::string type() const override { return "ObjectLiteral"; } + value execute_impl(context & ctx) override; +}; + +// Complex Expressions + +/** + * An operation with two sides, separated by an operator. + * Note: Either side can be a Complex Expression, with order + * of operations being determined by the operator. + */ +struct binary_expression : public expression { + token op; + statement_ptr left; + statement_ptr right; + + binary_expression(token op, statement_ptr && left, statement_ptr && right) + : op(std::move(op)), left(std::move(left)), right(std::move(right)) { + chk_type(this->left); + chk_type(this->right); + } + std::string type() const override { return "BinaryExpression"; } + value execute_impl(context & ctx) override; +}; + +/** + * An operation with two sides, separated by the | operator. + * Operator precedence: https://github.com/pallets/jinja/issues/379#issuecomment-168076202 + */ +struct filter_expression : public expression { + // either an expression or a value is allowed + statement_ptr operand; + value_string val; // will be set by filter_statement + + statement_ptr filter; + + filter_expression(statement_ptr && operand, statement_ptr && filter) + : operand(std::move(operand)), filter(std::move(filter)) { + chk_type(this->operand); + chk_type(this->filter); + } + + filter_expression(value_string && val, statement_ptr && filter) + : val(std::move(val)), filter(std::move(filter)) { + chk_type(this->filter); + } + + std::string type() const override { return "FilterExpression"; } + value execute_impl(context & ctx) override; +}; + +struct filter_statement : public statement { + statement_ptr filter; + statements body; + + filter_statement(statement_ptr && filter, statements && body) + : filter(std::move(filter)), body(std::move(body)) { + chk_type(this->filter); + } + std::string type() const override { return "FilterStatement"; } + value execute_impl(context & ctx) override; +}; + +/** + * An operation which filters a sequence of objects by applying a test to each object, + * and only selecting the objects with the test succeeding. + * + * It may also be used as a shortcut for a ternary operator. + */ +struct select_expression : public expression { + statement_ptr lhs; + statement_ptr test; + + select_expression(statement_ptr && lhs, statement_ptr && test) + : lhs(std::move(lhs)), test(std::move(test)) { + chk_type(this->lhs); + chk_type(this->test); + } + std::string type() const override { return "SelectExpression"; } + value execute_impl(context & ctx) override { + auto predicate = test->execute_impl(ctx); + if (!predicate->as_bool()) { + return mk_val(); + } + return lhs->execute_impl(ctx); + } +}; + +/** + * An operation with two sides, separated by the "is" operator. + * NOTE: "value is something" translates to function call "test_is_something(value)" + */ +struct test_expression : public expression { + statement_ptr operand; + bool negate; + statement_ptr test; + + test_expression(statement_ptr && operand, bool negate, statement_ptr && test) + : operand(std::move(operand)), negate(negate), test(std::move(test)) { + chk_type(this->operand); + chk_type(this->test); + } + std::string type() const override { return "TestExpression"; } + value execute_impl(context & ctx) override; +}; + +/** + * An operation with one side (operator on the left). + */ +struct unary_expression : public expression { + token op; + statement_ptr argument; + + unary_expression(token op, statement_ptr && argument) + : op(std::move(op)), argument(std::move(argument)) { + chk_type(this->argument); + } + std::string type() const override { return "UnaryExpression"; } + value execute_impl(context & ctx) override; +}; + +struct slice_expression : public expression { + statement_ptr start_expr; + statement_ptr stop_expr; + statement_ptr step_expr; + + slice_expression(statement_ptr && start_expr, statement_ptr && stop_expr, statement_ptr && step_expr) + : start_expr(std::move(start_expr)), stop_expr(std::move(stop_expr)), step_expr(std::move(step_expr)) { + chk_type(this->start_expr); + chk_type(this->stop_expr); + chk_type(this->step_expr); + } + std::string type() const override { return "SliceExpression"; } + [[noreturn]] value execute_impl(context &) override { + throw std::runtime_error("must be handled by MemberExpression"); + } +}; + +struct keyword_argument_expression : public expression { + statement_ptr key; + statement_ptr val; + + keyword_argument_expression(statement_ptr && key, statement_ptr && val) + : key(std::move(key)), val(std::move(val)) { + chk_type(this->key); + chk_type(this->val); + } + std::string type() const override { return "KeywordArgumentExpression"; } + value execute_impl(context & ctx) override; +}; + +struct spread_expression : public expression { + statement_ptr argument; + explicit spread_expression(statement_ptr && argument) : argument(std::move(argument)) { + chk_type(this->argument); + } + std::string type() const override { return "SpreadExpression"; } +}; + +struct call_statement : public statement { + statement_ptr call; + statements caller_args; + statements body; + + call_statement(statement_ptr && call, statements && caller_args, statements && body) + : call(std::move(call)), caller_args(std::move(caller_args)), body(std::move(body)) { + chk_type(this->call); + for (const auto & arg : this->caller_args) chk_type(arg); + } + std::string type() const override { return "CallStatement"; } +}; + +struct ternary_expression : public expression { + statement_ptr condition; + statement_ptr true_expr; + statement_ptr false_expr; + + ternary_expression(statement_ptr && condition, statement_ptr && true_expr, statement_ptr && false_expr) + : condition(std::move(condition)), true_expr(std::move(true_expr)), false_expr(std::move(false_expr)) { + chk_type(this->condition); + chk_type(this->true_expr); + chk_type(this->false_expr); + } + std::string type() const override { return "Ternary"; } + value execute_impl(context & ctx) override { + value cond_val = condition->execute(ctx); + if (cond_val->as_bool()) { + return true_expr->execute(ctx); + } else { + return false_expr->execute(ctx); + } + } +}; + +struct raised_exception : public std::exception { + std::string message; + raised_exception(const std::string & msg) : message(msg) {} + const char* what() const noexcept override { + return message.c_str(); + } +}; + +// Used to rethrow exceptions with modified messages +struct rethrown_exception : public std::exception { + std::string message; + rethrown_exception(const std::string & msg) : message(msg) {} + const char* what() const noexcept override { + return message.c_str(); + } +}; + +////////////////////// + +static void gather_string_parts_recursive(const value & val, value_string & parts) { + // TODO: probably allow print value_none as "None" string? currently this breaks some templates + if (is_val(val)) { + const auto & str_val = cast_val(val)->val_str; + parts->val_str.append(str_val); + } else if (is_val(val) || is_val(val) || is_val(val)) { + std::string str_val = val->as_string().str(); + parts->val_str.append(str_val); + } else if (is_val(val)) { + auto items = cast_val(val)->as_array(); + for (const auto & item : items) { + gather_string_parts_recursive(item, parts); + } + } +} + +static std::string render_string_parts(const value_string & parts) { + std::ostringstream oss; + for (const auto & part : parts->val_str.parts) { + oss << part.val; + } + return oss.str(); +} + +struct runtime { + context & ctx; + explicit runtime(context & ctx) : ctx(ctx) {} + + value_array execute(const program & prog) { + value_array results = mk_val(); + for (const auto & stmt : prog.body) { + value res = stmt->execute(ctx); + results->push_back(std::move(res)); + } + return results; + } + + static value_string gather_string_parts(const value & val) { + value_string parts = mk_val(); + gather_string_parts_recursive(val, parts); + // join consecutive parts with the same type + auto & p = parts->val_str.parts; + for (size_t i = 1; i < p.size(); ) { + if (p[i].is_input == p[i - 1].is_input) { + p[i - 1].val += p[i].val; + p.erase(p.begin() + i); + } else { + i++; + } + } + return parts; + } +}; + +} // namespace jinja + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +struct common_grammar_builder; + +class common_peg_parser_builder; + +using common_peg_parser_id = size_t; +constexpr common_peg_parser_id COMMON_PEG_INVALID_PARSER_ID = static_cast(-1); + +using common_peg_ast_id = size_t; +constexpr common_peg_ast_id COMMON_PEG_INVALID_AST_ID = static_cast(-1); + +// Lightweight wrapper around common_peg_parser_id for convenience +class common_peg_parser { + common_peg_parser_id id_; + common_peg_parser_builder & builder_; + + public: + common_peg_parser(const common_peg_parser & other) : id_(other.id_), builder_(other.builder_) {} + common_peg_parser(common_peg_parser_id id, common_peg_parser_builder & builder) : id_(id), builder_(builder) {} + + common_peg_parser & operator=(const common_peg_parser & other); + common_peg_parser & operator+=(const common_peg_parser & other); + common_peg_parser & operator|=(const common_peg_parser & other); + + operator common_peg_parser_id() const { return id_; } + common_peg_parser_id id() const { return id_; } + + common_peg_parser_builder & builder() const { return builder_; } + + // Creates a sequence + common_peg_parser operator+(const common_peg_parser & other) const; + + // Creates a sequence separated by spaces. + common_peg_parser operator<<(const common_peg_parser & other) const; + + // Creates a choice + common_peg_parser operator|(const common_peg_parser & other) const; + + common_peg_parser operator+(const char * str) const; + common_peg_parser operator+(const std::string & str) const; + common_peg_parser operator<<(const char * str) const; + common_peg_parser operator<<(const std::string & str) const; + common_peg_parser operator|(const char * str) const; + common_peg_parser operator|(const std::string & str) const; +}; + +common_peg_parser operator+(const char * str, const common_peg_parser & p); +common_peg_parser operator+(const std::string & str, const common_peg_parser & p); +common_peg_parser operator<<(const char * str, const common_peg_parser & p); +common_peg_parser operator<<(const std::string & str, const common_peg_parser & p); +common_peg_parser operator|(const char * str, const common_peg_parser & p); +common_peg_parser operator|(const std::string & str, const common_peg_parser & p); + +enum common_peg_parse_result_type { + COMMON_PEG_PARSE_RESULT_FAIL = 0, + COMMON_PEG_PARSE_RESULT_SUCCESS = 1, + COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT = 2, +}; + +const char * common_peg_parse_result_type_name(common_peg_parse_result_type type); + +struct common_peg_ast_node { + common_peg_ast_id id; + std::string rule; + std::string tag; + size_t start; + size_t end; + std::string_view text; + std::vector children; + + bool is_partial = false; +}; + +struct common_peg_parse_result; + +using common_peg_ast_visitor = std::function; + +class common_peg_ast_arena { + std::vector nodes_; + public: + common_peg_ast_id add_node( + const std::string & rule, + const std::string & tag, + size_t start, + size_t end, + std::string_view text, + std::vector children, + bool is_partial = false + ) { + common_peg_ast_id id = nodes_.size(); + nodes_.push_back({id, rule, tag, start, end, text, std::move(children), is_partial}); + return id; + } + + const common_peg_ast_node & get(common_peg_ast_id id) const { return nodes_.at(id); } + + common_peg_ast_id find_by_tag(const common_peg_ast_node & parent, const std::string & tag, int max_depth = 3) const; + common_peg_ast_id find_by_rule(const common_peg_ast_node & parent, const std::string & tag, int max_depth = 3) const; + + size_t size() const { return nodes_.size(); } + + void clear() { nodes_.clear(); } + + void visit(common_peg_ast_id id, const common_peg_ast_visitor & visitor) const; + void visit(const common_peg_parse_result & result, const common_peg_ast_visitor & visitor) const; + + std::string dump(); +}; + +struct common_peg_parse_result { + common_peg_parse_result_type type = COMMON_PEG_PARSE_RESULT_FAIL; + size_t start = 0; + size_t end = 0; + + std::vector nodes; + + common_peg_parse_result() = default; + + common_peg_parse_result(common_peg_parse_result_type type, size_t start) + : type(type), start(start), end(start) {} + + common_peg_parse_result(common_peg_parse_result_type type, size_t start, size_t end) + : type(type), start(start), end(end) {} + + common_peg_parse_result(common_peg_parse_result_type type, size_t start, size_t end, std::vector nodes) + : type(type), start(start), end(end), nodes(std::move(nodes)) {} + + bool fail() const { return type == COMMON_PEG_PARSE_RESULT_FAIL; } + bool need_more_input() const { return type == COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT; } + bool success() const { return type == COMMON_PEG_PARSE_RESULT_SUCCESS; } +}; + +enum common_peg_parse_flags { + COMMON_PEG_PARSE_FLAG_NONE = 0, + COMMON_PEG_PARSE_FLAG_LENIENT = 1 << 0, + COMMON_PEG_PARSE_FLAG_DEBUG = 1 << 1, +}; + +inline common_peg_parse_flags operator|(common_peg_parse_flags a, common_peg_parse_flags b) { + return static_cast(int(a) | int(b)); +} + +inline common_peg_parse_flags & operator|=(common_peg_parse_flags & a, common_peg_parse_flags b) { + return a = a | b; +} + +inline common_peg_parse_flags operator&(common_peg_parse_flags a, common_peg_parse_flags b) { + return static_cast(int(a) & int(b)); +} + +inline common_peg_parse_flags operator~(common_peg_parse_flags a) { + return static_cast(~int(a)); +} + +struct common_peg_parse_context { + std::string input; + common_peg_parse_flags flags; + common_peg_ast_arena ast; + + int parse_depth; + + common_peg_parse_context(common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE) + : flags(flags), parse_depth(0) {} + + common_peg_parse_context(const std::string & input, common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE) + : input(input), flags(flags), parse_depth(0) {} + + bool is_lenient() const { return flags & COMMON_PEG_PARSE_FLAG_LENIENT; } + bool is_debug() const { return flags & COMMON_PEG_PARSE_FLAG_DEBUG; } +}; + +class common_peg_arena; + +// Parser variants +struct common_peg_epsilon_parser {}; + +struct common_peg_start_parser {}; + +struct common_peg_end_parser {}; + +struct common_peg_literal_parser { + std::string literal; +}; + +struct common_peg_sequence_parser { + std::vector children; +}; + +struct common_peg_choice_parser { + std::vector children; +}; + +struct common_peg_repetition_parser { + common_peg_parser_id child; + int min_count; + int max_count; // -1 for unbounded +}; + +struct common_peg_and_parser { + common_peg_parser_id child; +}; + +struct common_peg_not_parser { + common_peg_parser_id child; +}; + +struct common_peg_any_parser {}; + +struct common_peg_space_parser {}; + +struct common_peg_chars_parser { + struct char_range { + uint32_t start; + uint32_t end; + bool contains(uint32_t codepoint) const { return codepoint >= start && codepoint <= end; } + }; + + std::string pattern; + std::vector ranges; + bool negated; + int min_count; + int max_count; // -1 for unbounded +}; + +struct common_peg_string_parser { + char delimiter; +}; + +struct common_peg_until_parser { + std::vector delimiters; +}; + +struct common_peg_schema_parser { + common_peg_parser_id child; + std::string name; + std::shared_ptr schema; + + // Indicates if the GBNF should accept a raw string that matches the schema. + bool raw; +}; + +struct common_peg_rule_parser { + std::string name; + common_peg_parser_id child; + bool trigger; +}; + +struct common_peg_ref_parser { + std::string name; +}; + +struct common_peg_atomic_parser { + common_peg_parser_id child; +}; + +struct common_peg_tag_parser { + common_peg_parser_id child; + std::string tag; +}; + +struct common_peg_gbnf_parser { + common_peg_parser_id child; + std::string grammar; +}; + +// Variant holding all parser types +using common_peg_parser_variant = std::variant< + common_peg_epsilon_parser, + common_peg_start_parser, + common_peg_end_parser, + common_peg_literal_parser, + common_peg_sequence_parser, + common_peg_choice_parser, + common_peg_repetition_parser, + common_peg_and_parser, + common_peg_not_parser, + common_peg_any_parser, + common_peg_space_parser, + common_peg_chars_parser, + common_peg_string_parser, + common_peg_until_parser, + common_peg_schema_parser, + common_peg_rule_parser, + common_peg_ref_parser, + common_peg_atomic_parser, + common_peg_tag_parser, + common_peg_gbnf_parser +>; + +class common_peg_arena { + std::vector parsers_; + std::unordered_map rules_; + common_peg_parser_id root_ = COMMON_PEG_INVALID_PARSER_ID; + + public: + const common_peg_parser_variant & get(common_peg_parser_id id) const { return parsers_.at(id); } + common_peg_parser_variant & get(common_peg_parser_id id) { return parsers_.at(id); } + + size_t size() const { return parsers_.size(); } + bool empty() const { return parsers_.empty(); } + + common_peg_parser_id get_rule(const std::string & name) const; + bool has_rule(const std::string & name) const { return rules_.find(name) != rules_.end(); } + + common_peg_parser_id root() const { return root_; } + void set_root(common_peg_parser_id id) { root_ = id; } + + common_peg_parse_result parse(common_peg_parse_context & ctx, size_t start = 0) const; + common_peg_parse_result parse(common_peg_parser_id id, common_peg_parse_context & ctx, size_t start) const; + + void resolve_refs(); + + void build_grammar(const common_grammar_builder & builder, bool lazy = false) const; + + std::string dump(common_peg_parser_id id) const; + + nlohmann::json to_json() const; + static common_peg_arena from_json(const nlohmann::json & j); + + std::string save() const; + void load(const std::string & data); + + friend class common_peg_parser_builder; + + private: + std::string dump_impl(common_peg_parser_id id, std::unordered_set & visited) const; + + common_peg_parser_id add_parser(common_peg_parser_variant parser); + void add_rule(const std::string & name, common_peg_parser_id id); + + common_peg_parser_id resolve_ref(common_peg_parser_id id); +}; + +class common_peg_parser_builder { + common_peg_arena arena_; + + common_peg_parser wrap(common_peg_parser_id id) { return common_peg_parser(id, *this); } + common_peg_parser add(const common_peg_parser_variant & p) { return wrap(arena_.add_parser(p)); } + + public: + common_peg_parser_builder(); + + // Match nothing, always succeed. + // S -> ε + common_peg_parser eps() { return add(common_peg_epsilon_parser{}); } + + // Matches the start of the input. + // S -> ^ + common_peg_parser start() { return add(common_peg_start_parser{}); } + + // Matches the end of the input. + // S -> $ + common_peg_parser end() { return add(common_peg_end_parser{}); } + + // Matches an exact literal string. + // S -> "hello" + common_peg_parser literal(const std::string & literal) { return add(common_peg_literal_parser{literal}); } + + // Matches a sequence of parsers in order, all must succeed. + // S -> A B C + common_peg_parser sequence() { return add(common_peg_sequence_parser{}); } + common_peg_parser sequence(const std::vector & parsers); + common_peg_parser sequence(const std::vector & parsers); + common_peg_parser sequence(std::initializer_list parsers); + + // Matches the first parser that succeeds from a list of alternatives. + // S -> A | B | C + common_peg_parser choice() { return add(common_peg_choice_parser{}); } + common_peg_parser choice(const std::vector & parsers); + common_peg_parser choice(const std::vector & parsers); + common_peg_parser choice(std::initializer_list parsers); + + // Matches one or more repetitions of a parser. + // S -> A+ + common_peg_parser one_or_more(const common_peg_parser & p) { return repeat(p, 1, -1); } + + // Matches zero or more repetitions of a parser, always succeeds. + // S -> A* + common_peg_parser zero_or_more(const common_peg_parser & p) { return repeat(p, 0, -1); } + + // Matches zero or one occurrence of a parser, always succeeds. + // S -> A? + common_peg_parser optional(const common_peg_parser & p) { return repeat(p, 0, 1); } + + // Positive lookahead: succeeds if child parser succeeds, consumes no input. + // S -> &A + common_peg_parser peek(const common_peg_parser & p) { return add(common_peg_and_parser{p}); } + + // Negative lookahead: succeeds if child parser fails, consumes no input. + // S -> !A + common_peg_parser negate(const common_peg_parser & p) { return add(common_peg_not_parser{p}); } + + // Matches any single character. + // S -> . + common_peg_parser any() { return add(common_peg_any_parser{}); } + + // Matches between min and max repetitions of characters from a character class. + // S -> [a-z]{m,n} + // + // Use -1 for max to represent unbounded repetition (equivalent to {m,}) + common_peg_parser chars(const std::string & classes, int min = 1, int max = -1); + + // Creates a lightweight reference to a named rule (resolved during build()). + // Use this for forward references in recursive grammars. + // expr_ref -> expr + common_peg_parser ref(const std::string & name) { return add(common_peg_ref_parser{name}); } + + // Matches zero or more whitespace characters (space, tab, newline). + // S -> [ \t\n]* + common_peg_parser space() { return add(common_peg_space_parser{}); } + + // Matches all characters until a delimiter is found (delimiter not consumed). + // S -> (!delim .)* + common_peg_parser until(const std::string & delimiter) { return add(common_peg_until_parser{{delimiter}}); } + + // Matches all characters until one of the delimiters in the list is found (delimiter not consumed). + // S -> (!delim .)* + common_peg_parser until_one_of(const std::vector & delimiters) { return add(common_peg_until_parser{delimiters}); } + + // Matches everything + // S -> .* + common_peg_parser rest() { return until_one_of({}); } + + // Matches between min and max repetitions of a parser (inclusive). + // S -> A{m,n} + // Use -1 for max to represent unbounded repetition (equivalent to {m,}) + common_peg_parser repeat(const common_peg_parser & p, int min, int max) { return add(common_peg_repetition_parser{p, min,max}); } + + // Matches exactly n repetitions of a parser. + // S -> A{n} + common_peg_parser repeat(const common_peg_parser & p, int n) { return repeat(p, n, n); } + + // Matches a double-quoted string: '"' content '"' space + common_peg_parser double_quoted_string(); + + // Matches a single-quoted string: "'" content "'" space + common_peg_parser single_quoted_string(); + + // Matches a string that accepts both double-quoted and single-quoted styles. + common_peg_parser quoted_string(); + + // Matches string content without the surrounding delimiter. + common_peg_parser string_content(char delimiter); + + // Creates a complete JSON parser supporting objects, arrays, strings, numbers, booleans, and null. + // value -> object | array | string | number | true | false | null + common_peg_parser json(); + common_peg_parser json_object(); + common_peg_parser json_string(); + common_peg_parser json_array(); + common_peg_parser json_number(); + common_peg_parser json_bool(); + common_peg_parser json_null(); + + // Matches a JSON object member with a key and associated parser as the + // value. + common_peg_parser json_member(const std::string & key, const common_peg_parser & p); + + // Creates a complete Python format parser supporting dicts, arrays, strings, numbers, booleans, and None. + // Differs from JSON: uses True/False/None, accepts both single and double-quoted strings. + // value -> dict | array | string | number | True | False | None + common_peg_parser python_value(); + common_peg_parser python_dict(); + common_peg_parser python_string(); + common_peg_parser python_array(); + common_peg_parser python_number(); + common_peg_parser python_bool(); + common_peg_parser python_null(); + + // A marker, i.e. text delimited by a pair of <> or [] + common_peg_parser marker(); + + // Wraps a parser with JSON schema metadata for grammar generation. + // Used internally to convert JSON schemas to GBNF grammar rules. + common_peg_parser schema(const common_peg_parser & p, const std::string & name, const nlohmann::ordered_json & schema, bool raw = false); + + // Creates a named rule, stores it in the grammar, and returns a ref. + // If trigger=true, marks this rule as an entry point for lazy grammar generation. + // auto json = p.rule("json", json_obj | json_arr | ...) + common_peg_parser rule(const std::string & name, const common_peg_parser & p, bool trigger = false); + + // Creates a named rule using a builder function, and returns a ref. + // If trigger=true, marks this rule as an entry point for lazy grammar generation. + // auto json = p.rule("json", [&]() { return json_object() | json_array() | ... }) + common_peg_parser rule(const std::string & name, const std::function & builder, bool trigger = false); + + // Creates a trigger rule. When generating a lazy grammar from the parser, + // only trigger rules and descendents are emitted. + common_peg_parser trigger_rule(const std::string & name, const common_peg_parser & p) { return rule(name, p, true); } + common_peg_parser trigger_rule(const std::string & name, const std::function & builder) { return rule(name, builder, true); } + + // Creates an atomic parser. Atomic parsers do not create an AST node if + // the child results in a partial parse, i.e. NEEDS_MORE_INPUT. This is + // intended for situations where partial output is undesirable. + common_peg_parser atomic(const common_peg_parser & p) { return add(common_peg_atomic_parser{p}); } + + // Tags create nodes in the generated AST for semantic purposes. + // Unlike rules, you can tag multiple nodes with the same tag. + common_peg_parser tag(const std::string & tag, const common_peg_parser & p) { return add(common_peg_tag_parser{p.id(), tag}); } + + // Wraps a child parser but emits a custom GBNF grammar string instead of + // the child's grammar. Parsing delegates entirely to the child. + common_peg_parser gbnf(const common_peg_parser & p, const std::string & grammar) { return add(common_peg_gbnf_parser{p, grammar}); } + + void set_root(const common_peg_parser & p); + + common_peg_arena build(); +}; + +// Helper function for building parsers +common_peg_arena build_peg_parser(const std::function & fn); + +// Chat support (incl. tool call grammar constraining & output parsing) w/ generic & custom template handlers. + +using json = nlohmann::ordered_json; + +namespace jinja { +program parse_from_tokens(const lexer_result & lexer_res); + +struct caps { + bool supports_tools = true; + bool supports_tool_calls = true; + bool supports_system_role = true; + bool supports_parallel_tool_calls = true; + bool supports_preserve_reasoning = false; // support assistant message with reasoning_content + + // one of the 2 content capabilities must be true + bool supports_string_content = true; + bool supports_typed_content = false; + + bool supports_object_arguments = false; + + // for reporting on server + std::map to_map() const; + + // for debugging + std::string to_string() const; +}; + +caps caps_get(program & prog); +} + +using chat_template_caps = jinja::caps; + +struct common_chat_templates; + +namespace autoparser { +struct generation_params; +} // namespace autoparser + +struct common_chat_tool_call { + std::string name; + std::string arguments; + std::string id; + + bool operator==(const common_chat_tool_call & other) const { + return name == other.name && arguments == other.arguments && id == other.id; + } +}; + +struct common_chat_msg_content_part { + std::string type; + std::string text; + + // TODO @ngxson : no known chat templates support reasoning_content in content parts yet + // this can be useful for models with interleaved thinking (like Kimi-K2) + // if you see any templates explicitly support this, please ping me + // std::string reasoning_content; + + bool operator==(const common_chat_msg_content_part & other) const { + return type == other.type && text == other.text; + } +}; + +struct common_chat_template { + jinja::program prog; + std::string bos_tok; + std::string eos_tok; + std::string src; + chat_template_caps caps; + + common_chat_template(const std::string & src, const std::string & bos_token, const std::string & eos_token) { + jinja::lexer lexer; + auto lexer_res = lexer.tokenize(src); + this->prog = jinja::parse_from_tokens(lexer_res); + + this->src = lexer_res.source; + this->bos_tok = bos_token; + this->eos_tok = eos_token; + + this->caps = jinja::caps_get(prog); + // LOG_INF("%s: caps:\n%s\n", __func__, this->caps.to_string().c_str()); + } + + const std::string & source() const { return src; } + const std::string & bos_token() const { return bos_tok; } + const std::string & eos_token() const { return eos_tok; } + + chat_template_caps original_caps() const { + return caps; + } +}; + +struct common_chat_msg { + std::string role; + std::string content; + std::vector content_parts; + std::vector tool_calls; + std::string reasoning_content; + std::string tool_name; + std::string tool_call_id; + + nlohmann::ordered_json to_json_oaicompat(bool concat_typed_text = false) const; + + bool empty() const { + return content.empty() && content_parts.empty() && tool_calls.empty() && reasoning_content.empty() && + tool_name.empty() && tool_call_id.empty(); + } + + void set_tool_call_ids(std::vector & ids_cache, + const std::function & gen_tool_call_id) { + for (auto i = 0u; i < tool_calls.size(); i++) { + if (ids_cache.size() <= i) { + auto id = tool_calls[i].id; + if (id.empty()) { + id = gen_tool_call_id(); + } + ids_cache.push_back(id); + } + tool_calls[i].id = ids_cache[i]; + } + } + + bool operator==(const common_chat_msg & other) const { + return role == other.role && content == other.content && content_parts == other.content_parts && + tool_calls == other.tool_calls && reasoning_content == other.reasoning_content && + tool_name == other.tool_name && tool_call_id == other.tool_call_id; + } + + bool operator!=(const common_chat_msg & other) const { return !(*this == other); } +}; + +struct common_chat_msg_diff { + std::string reasoning_content_delta; + std::string content_delta; + size_t tool_call_index = std::string::npos; + common_chat_tool_call tool_call_delta; + + static std::vector compute_diffs(const common_chat_msg & msg_prv, + const common_chat_msg & msg_new); + + bool operator==(const common_chat_msg_diff & other) const { + return content_delta == other.content_delta && tool_call_index == other.tool_call_index && + tool_call_delta == other.tool_call_delta; + } +}; + +struct common_chat_tool { + std::string name; + std::string description; + std::string parameters; +}; + +enum common_chat_tool_choice { + COMMON_CHAT_TOOL_CHOICE_AUTO, + COMMON_CHAT_TOOL_CHOICE_REQUIRED, + COMMON_CHAT_TOOL_CHOICE_NONE, +}; + +enum common_chat_format { + COMMON_CHAT_FORMAT_CONTENT_ONLY, + + // These are intended to be parsed by the PEG parser + COMMON_CHAT_FORMAT_PEG_SIMPLE, + COMMON_CHAT_FORMAT_PEG_NATIVE, + COMMON_CHAT_FORMAT_PEG_GEMMA4, + + COMMON_CHAT_FORMAT_COUNT, // Not a format, just the # formats +}; + +struct common_chat_templates_inputs { + std::vector messages; + std::string grammar; + std::string json_schema; + bool add_generation_prompt = true; + bool use_jinja = true; + // Parameters below only supported when use_jinja is true + std::vector tools; + common_chat_tool_choice tool_choice = COMMON_CHAT_TOOL_CHOICE_AUTO; + bool parallel_tool_calls = false; + common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_NONE; // TODO: refactor this to "bool enable_thinking" + bool enable_thinking = true; + std::chrono::system_clock::time_point now = std::chrono::system_clock::now(); + std::map chat_template_kwargs; + bool add_bos = false; + bool add_eos = false; + bool force_pure_content = false; +}; + +struct common_chat_params { + common_chat_format format = COMMON_CHAT_FORMAT_CONTENT_ONLY; + std::string prompt; + std::string grammar; + bool grammar_lazy = false; + std::string generation_prompt; + bool supports_thinking = false; + std::string thinking_start_tag; // e.g., "" + std::string thinking_end_tag; // e.g., "" + std::vector grammar_triggers; + std::vector preserved_tokens; + std::vector additional_stops; + std::string parser; +}; + +// per-message parsing syntax +// should be derived from common_chat_params +struct common_chat_parser_params { + common_chat_format format = COMMON_CHAT_FORMAT_CONTENT_ONLY; + common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_NONE; // TODO: refactor this to "bool parse_reasoning" + // Whether reasoning_content should be inlined in the content (e.g. for reasoning_format=deepseek in stream mode) + bool reasoning_in_content = false; + std::string generation_prompt; + bool parse_tool_calls = true; + bool debug = false; // Enable debug output for PEG parser + common_peg_arena parser = {}; + common_chat_parser_params() = default; + common_chat_parser_params(const common_chat_params & chat_params) { + format = chat_params.format; + generation_prompt = chat_params.generation_prompt; + } +}; + +// Check if the template supplied via "--chat-template" is supported or not. Returns true if it's valid +bool common_chat_verify_template(const std::string & tmpl, bool use_jinja); + +void common_chat_templates_free(struct common_chat_templates * tmpls); + +struct common_chat_templates_deleter { + void operator()(common_chat_templates * tmpls) { common_chat_templates_free(tmpls); } +}; + +typedef std::unique_ptr common_chat_templates_ptr; + +common_chat_templates_ptr common_chat_templates_init(const struct llama_model * model, + const std::string & chat_template_override, + const std::string & bos_token_override = "", + const std::string & eos_token_override = ""); + +bool common_chat_templates_was_explicit(const struct common_chat_templates * tmpls); +std::string common_chat_templates_source(const struct common_chat_templates * tmpls, const std::string & variant = ""); + +struct common_chat_params common_chat_templates_apply(const struct common_chat_templates * tmpls, + const struct common_chat_templates_inputs & inputs); + +// Format single message, while taking into account the position of that message in chat history +std::string common_chat_format_single(const struct common_chat_templates * tmpls, + const std::vector & past_msg, + const common_chat_msg & new_msg, + bool add_ass, + bool use_jinja); + +// Returns an example of formatted chat +std::string common_chat_format_example(const struct common_chat_templates * tmpls, + bool use_jinja, + const std::map & chat_template_kwargs); + +const char * common_chat_format_name(common_chat_format format); +common_chat_msg common_chat_parse(const std::string & input, bool is_partial, const common_chat_parser_params & params); +common_chat_msg common_chat_peg_parse(const common_peg_arena & src_parser, const std::string & input, bool is_partial, const common_chat_parser_params & params); + +// used by arg and server +const char * common_reasoning_format_name(common_reasoning_format format); +common_reasoning_format common_reasoning_format_from_name(const std::string & format); + +common_chat_tool_choice common_chat_tool_choice_parse_oaicompat(const std::string & tool_choice); + +bool common_chat_templates_support_enable_thinking(const common_chat_templates * chat_templates); + +// Parses a JSON array of messages in OpenAI's chat completion API format. +std::vector common_chat_msgs_parse_oaicompat(const nlohmann::ordered_json & messages); + +std::vector common_chat_tools_parse_oaicompat(const nlohmann::ordered_json & tools); + +// DEPRECATED: only used in tests +nlohmann::ordered_json common_chat_msgs_to_json_oaicompat(const std::vector & msgs, bool concat_typed_text = false); + +nlohmann::ordered_json common_chat_tools_to_json_oaicompat(const std::vector & tools); + +// get template caps, useful for reporting to server /props endpoint +std::map common_chat_templates_get_caps(const common_chat_templates * chat_templates); + +std::string common_chat_template_direct_apply( + const common_chat_template & tmpl, + const autoparser::generation_params & inputs); + +std::optional common_chat_try_specialized_template( + const common_chat_template & tmpl, + const std::string & src, + autoparser::generation_params & params); + +// specialized per-task preset +struct common_chat_prompt_preset { + std::string system; + std::string user; +}; + +common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates); + +class common_chat_peg_mapper { + public: + common_chat_msg & result; + + common_chat_peg_mapper(common_chat_msg & msg) : result(msg) {} + + virtual ~common_chat_peg_mapper() = default; + + virtual void from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result); + virtual void map(const common_peg_ast_node & node); + + protected: + virtual std::string normalize_container_value(const std::string & input); + + private: + std::optional pending_tool_call; + common_chat_tool_call * current_tool = nullptr; + int arg_count = 0; + bool closing_quote_pending = false; + std::string args_buffer; + + std::string & args_target(); +}; + +class common_chat_peg_gemma4_mapper : public common_chat_peg_mapper { + public: + common_chat_peg_gemma4_mapper(common_chat_msg & msg) : common_chat_peg_mapper(msg) {} + virtual void from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result); + + private: + void visit(const common_peg_ast_arena & arena, common_peg_ast_id id); +}; + +struct content_structure; +struct tool_call_structure; + +class common_chat_peg_builder : public common_peg_parser_builder { + public: + static constexpr const char * REASONING_BLOCK = "reasoning-block"; + static constexpr const char * REASONING = "reasoning"; + static constexpr const char * CONTENT = "content"; + + static constexpr const char * TOOL = "tool"; + static constexpr const char * TOOL_OPEN = "tool-open"; + static constexpr const char * TOOL_CLOSE = "tool-close"; + static constexpr const char * TOOL_ID = "tool-id"; + static constexpr const char * TOOL_NAME = "tool-name"; + static constexpr const char * TOOL_ARGS = "tool-args"; + static constexpr const char * TOOL_ARG = "tool-arg"; + static constexpr const char * TOOL_ARG_OPEN = "tool-arg-open"; + static constexpr const char * TOOL_ARG_CLOSE = "tool-arg-close"; + static constexpr const char * TOOL_ARG_NAME = "tool-arg-name"; + static constexpr const char * TOOL_ARG_VALUE = "tool-arg-value"; + static constexpr const char * TOOL_ARG_STRING_VALUE = "tool-arg-string-value"; + + common_peg_parser reasoning_block(const common_peg_parser & p) { return tag(REASONING_BLOCK, p); } + common_peg_parser reasoning(const common_peg_parser & p) { return tag(REASONING, p); } + common_peg_parser content(const common_peg_parser & p) { return tag(CONTENT, p); } + + common_peg_parser tag_with_safe_content(const std::string & tag_name, + const std::string & marker, + const common_peg_parser & p); + + common_peg_parser tool(const common_peg_parser & p) { return tag(TOOL, p); } + common_peg_parser tool_open(const common_peg_parser & p) { return atomic(tag(TOOL_OPEN, p)); } + common_peg_parser tool_close(const common_peg_parser & p) { return atomic(tag(TOOL_CLOSE, p)); } + common_peg_parser tool_id(const common_peg_parser & p) { return atomic(tag(TOOL_ID, p)); } + common_peg_parser tool_name(const common_peg_parser & p) { return atomic(tag(TOOL_NAME, p)); } + common_peg_parser tool_args(const common_peg_parser & p) { return tag(TOOL_ARGS, p); } + common_peg_parser tool_arg(const common_peg_parser & p) { return tag(TOOL_ARG, p); } + common_peg_parser tool_arg_open(const common_peg_parser & p) { return atomic(tag(TOOL_ARG_OPEN, p)); } + common_peg_parser tool_arg_close(const common_peg_parser & p) { return atomic(tag(TOOL_ARG_CLOSE, p)); } + common_peg_parser tool_arg_name(const common_peg_parser & p) { return atomic(tag(TOOL_ARG_NAME, p)); } + common_peg_parser tool_arg_value(const common_peg_parser & p) { return tag(TOOL_ARG_VALUE, p); } + common_peg_parser tool_arg_string_value(const common_peg_parser & p) { return tag(TOOL_ARG_STRING_VALUE, p); } + common_peg_parser tool_arg_json_value(const common_peg_parser & p) { return atomic(tag(TOOL_ARG_VALUE, p)); } + + common_peg_parser prefix(const std::string & s, const std::string & delimiter = {}); + + common_peg_parser standard_json_tools(const std::string & section_start, + const std::string & section_end, + const nlohmann::ordered_json & tools, + bool parallel_tool_calls, + bool force_tool_calls, + const std::string & name_key = "", + const std::string & args_key = "", + bool array_wrapped = false, + bool function_is_key = false, + const std::string & call_id_key = "", + const std::string & gen_call_id_key = "", + const std::vector & parameters_order = {}); + + common_peg_parser standard_constructed_tools(const std::map & markers, + const nlohmann::ordered_json & tools, + bool parallel_tool_calls, + bool force_tool_calls); + + common_peg_parser python_style_tool_calls(const nlohmann::ordered_json & tools, + bool parallel_tool_calls); + + private: + common_peg_parser build_json_tools_function_is_key(const nlohmann::ordered_json & tools, + const std::string & args_key, + const std::string & effective_args_key, + const std::string & call_id_key, + const std::string & gen_call_id_key); + + common_peg_parser build_json_tools_nested_keys(const nlohmann::ordered_json & tools, + const std::string & effective_name_key, + const std::string & effective_args_key, + const std::string & call_id_key, + const std::string & gen_call_id_key); + + common_peg_parser build_json_tools_flat_keys(const nlohmann::ordered_json & tools, + const std::string & effective_name_key, + const std::string & effective_args_key, + const std::string & call_id_key, + const std::string & gen_call_id_key, + const std::vector & parameters_order); +}; + +inline common_peg_arena build_chat_peg_parser( + const std::function & fn) { + common_chat_peg_builder builder; + builder.set_root(fn(builder)); + return builder.build(); +} + +class tag_based_peg_mapper { + public: + std::map tags; + + void from_ast(const common_peg_ast_arena & arena, const common_peg_parse_result & result); +}; + +struct tagged_parse_result { + common_peg_parse_result result; + std::map tags; +}; + +struct tagged_peg_parser { + common_peg_arena arena; + common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE; + + tagged_peg_parser & withDebug() { + flags |= COMMON_PEG_PARSE_FLAG_DEBUG; + return *this; + } + + tagged_peg_parser & withoutDebug() { + flags = flags & ~COMMON_PEG_PARSE_FLAG_DEBUG; + return *this; + } + + tagged_parse_result parse_and_extract(const std::string & input, common_peg_parse_flags extra_flags = COMMON_PEG_PARSE_FLAG_NONE) const; + tagged_parse_result parse_anywhere_and_extract(const std::string & input) const; +}; + +tagged_peg_parser build_tagged_peg_parser( + const std::function & fn); + +// ============================================================================ +// Parameters for template application (low-level, used by diff analysis) +// ============================================================================ +struct template_params { + json messages; + json tools; + bool add_generation_prompt = false; + bool enable_thinking = true; + std::optional extra_context = std::nullopt; +}; + +struct diff_split { + std::string prefix; + std::string suffix; + std::string left; + std::string right; + + bool operator==(struct diff_split & other) const { + return prefix == other.prefix && suffix == other.suffix && left == other.left && right == other.right; + } +}; + +// Result of compare_variants containing diff and original outputs +struct compare_variants_result { + diff_split diff; + std::string output_A; + std::string output_B; +}; + +std::string trim_whitespace(const std::string & str); +std::string trim_leading_whitespace(const std::string & str); +std::string trim_trailing_whitespace(const std::string & str); +std::string trim_trailing_newlines(const std::string & str); + +diff_split calculate_diff_split(const std::string & left, const std::string & right); +std::string until_common_prefix(const std::string & full, const std::string & left, const std::string & right); +std::string after_common_suffix(const std::string & full, const std::string & left, const std::string & right); + +namespace autoparser { + +std::string apply_template(const common_chat_template & tmpl, const template_params & params); + +std::optional compare_variants( + const common_chat_template & tmpl, + const template_params & params_A, + const std::function & params_modifier); + +// ============================================================================ +// High-level params for parser generation +// ============================================================================ + +struct generation_params { + json messages; + json tools; + common_chat_tool_choice tool_choice = COMMON_CHAT_TOOL_CHOICE_AUTO; + json json_schema; + bool parallel_tool_calls = true; + common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_AUTO; + bool stream = true; + std::string grammar; + bool add_generation_prompt = false; + bool enable_thinking = true; + std::chrono::system_clock::time_point now = std::chrono::system_clock::now(); + std::string generation_prompt; + json extra_context; + bool add_bos = false; + bool add_eos = false; + bool is_inference = true; + bool add_inference = false; + bool mark_input = true; // whether to mark input strings in the jinja context +}; + +// ============================================================================ +// Analysis Result Enums +// ============================================================================ + +// Reasoning handling mode (derived from R1-R3 comparisons) +enum class reasoning_mode { + NONE, // No reasoning markers detected + TAG_BASED, // Tag-based: ... (start can be empty for delimiter-style) + TOOLS_ONLY // Only reason on tool calls, not on normal content +}; + +inline std::ostream & operator<<(std::ostream & os, const reasoning_mode & mode) { + switch (mode) { + case reasoning_mode::NONE: + return os << "NONE"; + case reasoning_mode::TAG_BASED: + return os << "TAG_BASED"; + case reasoning_mode::TOOLS_ONLY: + return os << "TOOLS_ONLY"; + default: + return os << "UNKNOWN"; + } +} + +// Content wrapping mode (derived from C1 comparison) +enum class content_mode { + PLAIN, // No content markers + ALWAYS_WRAPPED, // Content always wrapped with markers + WRAPPED_WITH_REASONING, // Content wrapped only when reasoning present +}; + +inline std::ostream & operator<<(std::ostream & os, const content_mode & mode) { + switch (mode) { + case content_mode::PLAIN: + return os << "PLAIN"; + case content_mode::ALWAYS_WRAPPED: + return os << "ALWAYS_WRAPPED"; + case content_mode::WRAPPED_WITH_REASONING: + return os << "WRAPPED_WITH_REASONING"; + default: + return os << "UNKNOWN"; + } +} + +// Call ID position in tool calls (for non-JSON formats) +enum class call_id_position { + NONE, // No call ID support detected + PRE_FUNC_NAME, // Call ID before function name: [CALL_ID]id[FUNC]name{args} + BETWEEN_FUNC_AND_ARGS, // Call ID between function and args: [FUNC]name[CALL_ID]id{args} + POST_ARGS, // Call ID after arguments: [FUNC]name{args}[CALL_ID]id +}; + +inline std::ostream & operator<<(std::ostream & os, const call_id_position & pos) { + switch (pos) { + case call_id_position::NONE: + return os << "NONE"; + case call_id_position::PRE_FUNC_NAME: + return os << "PRE_FUNC_NAME"; + case call_id_position::BETWEEN_FUNC_AND_ARGS: + return os << "BETWEEN_FUNC_AND_ARGS"; + case call_id_position::POST_ARGS: + return os << "POST_ARGS"; + default: + return os << "UNKNOWN"; + } +} + +// Tool call format classification (derived from T1-T5, A1-A3 comparisons) +enum class tool_format { + NONE, // No tool support detected + JSON_NATIVE, // Pure JSON: {"name": "X", "arguments": {...}} + TAG_WITH_JSON, // Tag-based with JSON args: {...} + TAG_WITH_TAGGED, // Tag-based with tagged args: value +}; + +inline std::ostream & operator<<(std::ostream & os, const tool_format & format) { + switch (format) { + case tool_format::NONE: + return os << "NONE"; + case tool_format::JSON_NATIVE: + return os << "JSON_NATIVE"; + case tool_format::TAG_WITH_JSON: + return os << "TAG_WITH_JSON"; + case tool_format::TAG_WITH_TAGGED: + return os << "TAG_WITH_TAGGED"; + default: + return os << "UNKNOWN"; + } +} + +// ============================================================================ +// Sub-structs for tool analysis +// ============================================================================ + +struct tool_format_analysis { + tool_format mode = tool_format::NONE; + + std::string section_start; // e.g., "", "[TOOL_CALLS]", "" + std::string section_end; // e.g., "", "" + std::string per_call_start; // e.g., "<|tool_call_begin|>", "" (for multi-call templates) + std::string per_call_end; // e.g., "<|tool_call_end|>", "" + + bool fun_name_is_key = false; // In JSON format function name is JSON key, i.e. { "": { ... arguments ... } } + bool tools_array_wrapped = false; // Tool calls wrapped in JSON array [...] + + std::string function_field = "function"; + std::string name_field = "name"; + std::string args_field = "arguments"; + std::string id_field; + std::string gen_id_field; + std::vector parameter_order; +}; + +struct tool_function_analysis { + std::string name_prefix; // e.g., "", "\"", ":0" + std::string close; // e.g., "", "" (for tag-based) +}; + +struct tool_arguments_analysis { + std::string start; // e.g., "<|tool_call_argument_begin|>", "" + std::string end; // e.g., "<|tool_call_argument_end|>", "" + std::string name_prefix; // e.g., "", "\"" + std::string name_suffix; // e.g., ">", "", "\":" + std::string value_prefix; // e.g., "", "", "" + std::string value_suffix; // e.g., "", "", "" + std::string separator; // e.g., "", "\n", "," +}; + +struct tool_id_analysis { + call_id_position pos = call_id_position::NONE; + + std::string prefix; // e.g., "[CALL_ID]" (marker before call ID value) + std::string suffix; // e.g., "" (marker after call ID value, before next section) +}; + +// ============================================================================ +// Parser build context (shared interface for build_parser methods) +// ============================================================================ + +struct analyze_content; +struct analyze_reasoning; + +struct parser_build_context { + common_chat_peg_builder & p; + const generation_params & inputs; + common_peg_parser reasoning_parser; + bool extracting_reasoning = false; + const analyze_reasoning * reasoning = nullptr; + const analyze_content * content = nullptr; + + parser_build_context(common_chat_peg_builder & p, const generation_params & inputs); +}; + +// ============================================================================ +// Base class for analyzers with parser building +// ============================================================================ + +struct analyze_base { + virtual ~analyze_base() = default; + virtual common_peg_parser build_parser(parser_build_context & ctx) const = 0; + + protected: + const common_chat_template * tmpl = nullptr; + + analyze_base() = default; + explicit analyze_base(const common_chat_template & tmpl) : tmpl(&tmpl) {} +}; + +// ============================================================================ +// Reasoning analyzer +// ============================================================================ + +struct analyze_reasoning : analyze_base { + reasoning_mode mode = reasoning_mode::NONE; + + std::string start; // e.g., "", "[THINK]", "<|START_THINKING|>", "" + std::string end; // e.g., "", "[BEGIN FINAL RESPONSE]", "<|END_THINKING|>" + + analyze_reasoning() = default; + analyze_reasoning(const common_chat_template & tmpl, bool supports_tools); + analyze_reasoning(std::string start_, std::string end_) : start(std::move(start_)), end(std::move(end_)) {} + + common_peg_parser build_parser(parser_build_context & ctx) const override; + + private: + // Look for reasoning markers in rendered content + void compare_reasoning_presence(); + + // Compare generation prompt with enable_thinking=true vs false + void compare_thinking_enabled(); + + // Check if reasoning is always possible or only in tool calls + void compare_reasoning_scope(); +}; + +// ============================================================================ +// Content analyzer +// ============================================================================ + +struct analyze_content : analyze_base { + content_mode mode = content_mode::PLAIN; + + std::string start; // e.g., "", ">>>all\n", "" + std::string end; // e.g., "", "" + + bool requires_nonnull_content = false; + + analyze_content() = default; + analyze_content(const common_chat_template & tmpl, const analyze_reasoning & reasoning); + + common_peg_parser build_parser(parser_build_context & ctx) const override; + + bool is_always_wrapped() const; + common_peg_parser build_optional_wrapped(parser_build_context & ctx) const; +}; + +// ============================================================================ +// Tool analyzer +// ============================================================================ + +struct analyze_tools : analyze_base { + tool_format_analysis format; + tool_function_analysis function; + tool_arguments_analysis arguments; + tool_id_analysis call_id; + + analyze_tools() = default; + analyze_tools(const common_chat_template & tmpl, + const jinja::caps & caps, + const analyze_reasoning & reasoning); + + common_peg_parser build_parser(parser_build_context & ctx) const override; + + private: + // Extract tool calling 'haystack' for further analysis and delegate further analysis based on format + void analyze_tool_calls(const analyze_reasoning & reasoning, bool supports_parallel_tool_calls); + + // Analyze format based on position of function and argument name in needle + void analyze_tool_call_format(const std::string & haystack, + const std::string & fun_name_needle, + const std::string & arg_name_needle, + const analyze_reasoning & reasoning, + bool supports_parallel_tool_calls); + + // Analyze specifics of JSON native format (entire tool call is a JSON object) + void analyze_tool_call_format_json_native(const std::string & clean_haystack, + const std::string & fun_name_needle, + const std::string & arg_name_needle); + + // Check if parallel calls in JSON native format array wrapped or tag wrapped + void analyze_json_native_parallel_calls(); + + // Analyze specifics of non-JSON native format (tags for function name or for function name and arguments) + void analyze_tool_call_format_non_json(const std::string & clean_haystack, + const std::string & fun_name_needle); + + // Check for and extract specific per-call markers for non-native-JSON templates with parallel call support + void check_per_call_markers(); + + // Extract function name markers + void extract_function_markers(); + + // Delegates to separate functions for: separator analysis, argument name analysis, argument value analysis + void analyze_arguments(); + + // Extract argument name markers + void extract_argument_name_markers(); + + // Extract argument value markers + void extract_argument_value_markers(); + + // Extract argument separator, if specified (eg. ......) + void extract_argument_separator(); + + // Extract argument wrapper markers, if present (eg. '......') + void extract_args_markers(); + + // Extract call ID markers, if present + void extract_call_id_markers(); + + // Per-format tool parser builders + common_peg_parser build_tool_parser_json_native(parser_build_context & ctx) const; + common_peg_parser build_tool_parser_tag_json(parser_build_context & ctx) const; + common_peg_parser build_tool_parser_tag_tagged(parser_build_context & ctx) const; + + // Shared helper: builds func_parser from open+call_id+args, handling atomic wrapping and close. + // atomic_peek: if present, used as the peek expression in the third atomicity branch. + common_peg_parser build_func_parser(common_chat_peg_builder & p, const std::string & name, + const common_peg_parser & call_id_section, bool have_call_id, + const common_peg_parser & args, + std::optional atomic_peek) const; +}; + +// ============================================================================ +// Main autoparser class +// ============================================================================ + +struct autoparser { + jinja::caps jinja_caps; + analyze_reasoning reasoning; + analyze_content content; + analyze_tools tools; + bool analysis_complete = false; + + // Preserved tokens for tokenizer (union of all non-empty markers) + std::vector preserved_tokens; + + autoparser() = default; + + // Run full differential analysis on a template + void analyze_template(const common_chat_template & tmpl); + + // Build the PEG parser for this template + common_peg_arena build_parser(const generation_params & inputs) const; + + private: + // Collect tokens from entire analysis to preserve + void collect_preserved_tokens(); +}; + +// ============================================================================ +// Parser generator +// ============================================================================ + +class peg_generator { + public: + static common_chat_params generate_parser(const common_chat_template & tmpl, + const struct generation_params & inputs); + + static common_chat_params generate_parser(const common_chat_template & tmpl, + const struct generation_params & inputs, + const autoparser & autoparser); +}; + +} // namespace autoparser + +enum segment_type { TEXT, MARKER }; + +inline std::ostream & operator<<(std::ostream & os, const segment_type & type) { + switch (type) { + case segment_type::TEXT: + return os << "TEXT"; + case segment_type::MARKER: + return os << "MARKER"; + default: + return os << "UNKNOWN"; + } +} + +struct segment { + segment_type type; + std::string value; + + segment(segment_type type, std::string value) : type(type), value(std::move(value)) {} + + bool operator==(const segment & other) const { + return type == other.type && value == other.value; + } + + bool operator!=(const segment & other) const { + return !(*this == other); + } +}; + +std::vector segmentize_markers(const std::string & text); +std::vector prune_whitespace_segments(const std::vector & segments); diff --git a/common/console.cpp b/common/console.cpp deleted file mode 100644 index 36f645f3329e..000000000000 --- a/common/console.cpp +++ /dev/null @@ -1,1166 +0,0 @@ -#include "console.h" -#include "log.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#if defined(_WIN32) -#define WIN32_LEAN_AND_MEAN -#ifndef NOMINMAX -#define NOMINMAX -#endif -#include -#include -#include -#ifndef ENABLE_VIRTUAL_TERMINAL_PROCESSING -#define ENABLE_VIRTUAL_TERMINAL_PROCESSING 0x0004 -#endif -#else -#include -#include -#include -#include -#include -#include -#include -#include -#endif - -#define ANSI_COLOR_RED "\x1b[31m" -#define ANSI_COLOR_GREEN "\x1b[32m" -#define ANSI_COLOR_YELLOW "\x1b[33m" -#define ANSI_COLOR_BLUE "\x1b[34m" -#define ANSI_COLOR_MAGENTA "\x1b[35m" -#define ANSI_COLOR_CYAN "\x1b[36m" -#define ANSI_COLOR_GRAY "\x1b[90m" -#define ANSI_COLOR_RESET "\x1b[0m" -#define ANSI_BOLD "\x1b[1m" - -namespace console { - -#if defined (_WIN32) - namespace { - // Use private-use unicode values to represent special keys that are not reported - // as characters (e.g. arrows on Windows). These values should never clash with - // real input and let the rest of the code handle navigation uniformly. - static constexpr char32_t KEY_ARROW_LEFT = 0xE000; - static constexpr char32_t KEY_ARROW_RIGHT = 0xE001; - static constexpr char32_t KEY_ARROW_UP = 0xE002; - static constexpr char32_t KEY_ARROW_DOWN = 0xE003; - static constexpr char32_t KEY_HOME = 0xE004; - static constexpr char32_t KEY_END = 0xE005; - static constexpr char32_t KEY_CTRL_ARROW_LEFT = 0xE006; - static constexpr char32_t KEY_CTRL_ARROW_RIGHT = 0xE007; - static constexpr char32_t KEY_DELETE = 0xE008; - } - - // - // Console state - // -#endif - - static bool advanced_display = false; - static bool simple_io = true; - static display_type current_display = DISPLAY_TYPE_RESET; - - static FILE* out = stdout; - -#if defined (_WIN32) - static void* hConsole; -#else - static FILE* tty = nullptr; - static termios initial_state; -#endif - - static completion_callback completion_cb = nullptr; - - // - // Init and cleanup - // - - void init(bool use_simple_io, bool use_advanced_display) { - advanced_display = use_advanced_display; - simple_io = use_simple_io; -#if defined(_WIN32) - // Windows-specific console initialization - DWORD dwMode = 0; - hConsole = GetStdHandle(STD_OUTPUT_HANDLE); - if (hConsole == INVALID_HANDLE_VALUE || !GetConsoleMode(hConsole, &dwMode)) { - hConsole = GetStdHandle(STD_ERROR_HANDLE); - if (hConsole != INVALID_HANDLE_VALUE && (!GetConsoleMode(hConsole, &dwMode))) { - hConsole = nullptr; - simple_io = true; - } - } - if (hConsole) { - // Check conditions combined to reduce nesting - if (advanced_display && !(dwMode & ENABLE_VIRTUAL_TERMINAL_PROCESSING) && - !SetConsoleMode(hConsole, dwMode | ENABLE_VIRTUAL_TERMINAL_PROCESSING)) { - advanced_display = false; - } - // Set console output codepage to UTF8 - SetConsoleOutputCP(CP_UTF8); - } - HANDLE hConIn = GetStdHandle(STD_INPUT_HANDLE); - if (hConIn != INVALID_HANDLE_VALUE && GetConsoleMode(hConIn, &dwMode)) { - // Set console input codepage to UTF16 - _setmode(_fileno(stdin), _O_WTEXT); - - // Set ICANON (ENABLE_LINE_INPUT) and ECHO (ENABLE_ECHO_INPUT) - if (simple_io) { - dwMode |= ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT; - } else { - dwMode &= ~(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT); - } - if (!SetConsoleMode(hConIn, dwMode)) { - simple_io = true; - } - } - if (simple_io) { - _setmode(_fileno(stdin), _O_U8TEXT); - } -#else - // POSIX-specific console initialization - if (!simple_io) { - struct termios new_termios; - tcgetattr(STDIN_FILENO, &initial_state); - new_termios = initial_state; - new_termios.c_lflag &= ~(ICANON | ECHO); - new_termios.c_cc[VMIN] = 1; - new_termios.c_cc[VTIME] = 0; - tcsetattr(STDIN_FILENO, TCSANOW, &new_termios); - - tty = fopen("/dev/tty", "w+"); - if (tty != nullptr) { - out = tty; - } - } - - setlocale(LC_ALL, ""); -#endif - } - - void cleanup() { - // Reset console display - set_display(DISPLAY_TYPE_RESET); - -#if !defined(_WIN32) - // Restore settings on POSIX systems - if (!simple_io) { - if (tty != nullptr) { - out = stdout; - fclose(tty); - tty = nullptr; - } - tcsetattr(STDIN_FILENO, TCSANOW, &initial_state); - } -#endif - } - - // - // Display and IO - // - - // Keep track of current display and only emit ANSI code if it changes - void set_display(display_type display) { - if (advanced_display && current_display != display) { - common_log_flush(common_log_main()); - switch(display) { - case DISPLAY_TYPE_RESET: - fprintf(out, ANSI_COLOR_RESET); - break; - case DISPLAY_TYPE_INFO: - fprintf(out, ANSI_COLOR_MAGENTA); - break; - case DISPLAY_TYPE_PROMPT: - fprintf(out, ANSI_COLOR_YELLOW); - break; - case DISPLAY_TYPE_REASONING: - fprintf(out, ANSI_COLOR_GRAY); - break; - case DISPLAY_TYPE_USER_INPUT: - fprintf(out, ANSI_BOLD ANSI_COLOR_GREEN); - break; - case DISPLAY_TYPE_ERROR: - fprintf(out, ANSI_BOLD ANSI_COLOR_RED); - } - current_display = display; - fflush(out); - } - } - - static char32_t getchar32() { -#if defined(_WIN32) - HANDLE hConsole = GetStdHandle(STD_INPUT_HANDLE); - wchar_t high_surrogate = 0; - - while (true) { - INPUT_RECORD record; - DWORD count; - if (!ReadConsoleInputW(hConsole, &record, 1, &count) || count == 0) { - return WEOF; - } - - if (record.EventType == KEY_EVENT && record.Event.KeyEvent.bKeyDown) { - wchar_t wc = record.Event.KeyEvent.uChar.UnicodeChar; - if (wc == 0) { - const DWORD ctrl_mask = LEFT_CTRL_PRESSED | RIGHT_CTRL_PRESSED; - const bool ctrl_pressed = (record.Event.KeyEvent.dwControlKeyState & ctrl_mask) != 0; - switch (record.Event.KeyEvent.wVirtualKeyCode) { - case VK_LEFT: return ctrl_pressed ? KEY_CTRL_ARROW_LEFT : KEY_ARROW_LEFT; - case VK_RIGHT: return ctrl_pressed ? KEY_CTRL_ARROW_RIGHT : KEY_ARROW_RIGHT; - case VK_UP: return KEY_ARROW_UP; - case VK_DOWN: return KEY_ARROW_DOWN; - case VK_HOME: return KEY_HOME; - case VK_END: return KEY_END; - case VK_DELETE: return KEY_DELETE; - default: continue; - } - } - - if ((wc >= 0xD800) && (wc <= 0xDBFF)) { // Check if wc is a high surrogate - high_surrogate = wc; - continue; - } - if ((wc >= 0xDC00) && (wc <= 0xDFFF)) { // Check if wc is a low surrogate - if (high_surrogate != 0) { // Check if we have a high surrogate - return ((high_surrogate - 0xD800) << 10) + (wc - 0xDC00) + 0x10000; - } - } - - high_surrogate = 0; // Reset the high surrogate - return static_cast(wc); - } - } -#else - wchar_t wc = getwchar(); - if (static_cast(wc) == WEOF) { - return WEOF; - } - -#if WCHAR_MAX == 0xFFFF - if ((wc >= 0xD800) && (wc <= 0xDBFF)) { // Check if wc is a high surrogate - wchar_t low_surrogate = getwchar(); - if ((low_surrogate >= 0xDC00) && (low_surrogate <= 0xDFFF)) { // Check if the next wchar is a low surrogate - return (static_cast(wc & 0x03FF) << 10) + (low_surrogate & 0x03FF) + 0x10000; - } - } - if ((wc >= 0xD800) && (wc <= 0xDFFF)) { // Invalid surrogate pair - return 0xFFFD; // Return the replacement character U+FFFD - } -#endif - - return static_cast(wc); -#endif - } - - static void pop_cursor() { -#if defined(_WIN32) - if (hConsole != NULL) { - CONSOLE_SCREEN_BUFFER_INFO bufferInfo; - GetConsoleScreenBufferInfo(hConsole, &bufferInfo); - - COORD newCursorPosition = bufferInfo.dwCursorPosition; - if (newCursorPosition.X == 0) { - newCursorPosition.X = bufferInfo.dwSize.X - 1; - newCursorPosition.Y -= 1; - } else { - newCursorPosition.X -= 1; - } - - SetConsoleCursorPosition(hConsole, newCursorPosition); - return; - } -#endif - putc('\b', out); - } - - static int estimateWidth(char32_t codepoint) { -#if defined(_WIN32) - (void)codepoint; - return 1; -#else - return wcwidth(codepoint); -#endif - } - - static int put_codepoint(const char* utf8_codepoint, size_t length, int expectedWidth) { -#if defined(_WIN32) - CONSOLE_SCREEN_BUFFER_INFO bufferInfo; - if (!GetConsoleScreenBufferInfo(hConsole, &bufferInfo)) { - // go with the default - return expectedWidth; - } - COORD initialPosition = bufferInfo.dwCursorPosition; - DWORD nNumberOfChars = length; - WriteConsole(hConsole, utf8_codepoint, nNumberOfChars, &nNumberOfChars, NULL); - - CONSOLE_SCREEN_BUFFER_INFO newBufferInfo; - GetConsoleScreenBufferInfo(hConsole, &newBufferInfo); - - // Figure out our real position if we're in the last column - if (utf8_codepoint[0] != 0x09 && initialPosition.X == newBufferInfo.dwSize.X - 1) { - DWORD nNumberOfChars; - WriteConsole(hConsole, &" \b", 2, &nNumberOfChars, NULL); - GetConsoleScreenBufferInfo(hConsole, &newBufferInfo); - } - - int width = newBufferInfo.dwCursorPosition.X - initialPosition.X; - if (width < 0) { - width += newBufferInfo.dwSize.X; - } - return width; -#else - // We can trust expectedWidth if we've got one - if (expectedWidth >= 0 || tty == nullptr) { - fwrite(utf8_codepoint, length, 1, out); - return expectedWidth; - } - - fputs("\033[6n", tty); // Query cursor position - int x1; - int y1; - int x2; - int y2; - int results = 0; - results = fscanf(tty, "\033[%d;%dR", &y1, &x1); - - fwrite(utf8_codepoint, length, 1, tty); - - fputs("\033[6n", tty); // Query cursor position - results += fscanf(tty, "\033[%d;%dR", &y2, &x2); - - if (results != 4) { - return expectedWidth; - } - - int width = x2 - x1; - if (width < 0) { - // Calculate the width considering text wrapping - struct winsize w; - ioctl(STDOUT_FILENO, TIOCGWINSZ, &w); - width += w.ws_col; - } - return width; -#endif - } - - static void replace_last(char ch) { -#if defined(_WIN32) - pop_cursor(); - put_codepoint(&ch, 1, 1); -#else - fprintf(out, "\b%c", ch); -#endif - } - - static char32_t decode_utf8(const std::string & input, size_t pos, size_t & advance) { - unsigned char c = static_cast(input[pos]); - if ((c & 0x80u) == 0u) { - advance = 1; - return c; - } - if ((c & 0xE0u) == 0xC0u && pos + 1 < input.size()) { - unsigned char c1 = static_cast(input[pos + 1]); - if ((c1 & 0xC0u) != 0x80u) { - advance = 1; - return 0xFFFD; - } - advance = 2; - return ((c & 0x1Fu) << 6) | (static_cast(input[pos + 1]) & 0x3Fu); - } - if ((c & 0xF0u) == 0xE0u && pos + 2 < input.size()) { - unsigned char c1 = static_cast(input[pos + 1]); - unsigned char c2 = static_cast(input[pos + 2]); - if ((c1 & 0xC0u) != 0x80u || (c2 & 0xC0u) != 0x80u) { - advance = 1; - return 0xFFFD; - } - advance = 3; - return ((c & 0x0Fu) << 12) | - ((static_cast(input[pos + 1]) & 0x3Fu) << 6) | - (static_cast(input[pos + 2]) & 0x3Fu); - } - if ((c & 0xF8u) == 0xF0u && pos + 3 < input.size()) { - unsigned char c1 = static_cast(input[pos + 1]); - unsigned char c2 = static_cast(input[pos + 2]); - unsigned char c3 = static_cast(input[pos + 3]); - if ((c1 & 0xC0u) != 0x80u || (c2 & 0xC0u) != 0x80u || (c3 & 0xC0u) != 0x80u) { - advance = 1; - return 0xFFFD; - } - advance = 4; - return ((c & 0x07u) << 18) | - ((static_cast(input[pos + 1]) & 0x3Fu) << 12) | - ((static_cast(input[pos + 2]) & 0x3Fu) << 6) | - (static_cast(input[pos + 3]) & 0x3Fu); - } - - advance = 1; - return 0xFFFD; // replacement character for invalid input - } - - static void append_utf8(char32_t ch, std::string & out) { - if (ch <= 0x7F) { - out.push_back(static_cast(ch)); - } else if (ch <= 0x7FF) { - out.push_back(static_cast(0xC0 | ((ch >> 6) & 0x1F))); - out.push_back(static_cast(0x80 | (ch & 0x3F))); - } else if (ch <= 0xFFFF) { - out.push_back(static_cast(0xE0 | ((ch >> 12) & 0x0F))); - out.push_back(static_cast(0x80 | ((ch >> 6) & 0x3F))); - out.push_back(static_cast(0x80 | (ch & 0x3F))); - } else if (ch <= 0x10FFFF) { - out.push_back(static_cast(0xF0 | ((ch >> 18) & 0x07))); - out.push_back(static_cast(0x80 | ((ch >> 12) & 0x3F))); - out.push_back(static_cast(0x80 | ((ch >> 6) & 0x3F))); - out.push_back(static_cast(0x80 | (ch & 0x3F))); - } else { - // Invalid Unicode code point - } - } - - // Helper function to remove the last UTF-8 character from a string - static size_t prev_utf8_char_pos(const std::string & line, size_t pos) { - if (pos == 0) return 0; - pos--; - while (pos > 0 && (line[pos] & 0xC0) == 0x80) { - pos--; - } - return pos; - } - - static size_t next_utf8_char_pos(const std::string & line, size_t pos) { - if (pos >= line.length()) return line.length(); - pos++; - while (pos < line.length() && (line[pos] & 0xC0) == 0x80) { - pos++; - } - return pos; - } - - static void move_cursor(int delta); - static void move_word_left(size_t & char_pos, size_t & byte_pos, const std::vector & widths, const std::string & line); - static void move_word_right(size_t & char_pos, size_t & byte_pos, const std::vector & widths, const std::string & line); - static void move_to_line_start(size_t & char_pos, size_t & byte_pos, const std::vector & widths); - static void move_to_line_end(size_t & char_pos, size_t & byte_pos, const std::vector & widths, const std::string & line); - - static void delete_at_cursor(std::string & line, std::vector & widths, size_t & char_pos, size_t & byte_pos) { - if (char_pos >= widths.size()) { - return; - } - - size_t next_pos = next_utf8_char_pos(line, byte_pos); - int w = widths[char_pos]; - size_t char_len = next_pos - byte_pos; - - line.erase(byte_pos, char_len); - widths.erase(widths.begin() + char_pos); - - size_t p = byte_pos; - int tail_width = 0; - for (size_t i = char_pos; i < widths.size(); ++i) { - size_t following = next_utf8_char_pos(line, p); - put_codepoint(line.c_str() + p, following - p, widths[i]); - tail_width += widths[i]; - p = following; - } - - for (int i = 0; i < w; ++i) { - fputc(' ', out); - } - - move_cursor(-(tail_width + w)); - } - - static void clear_current_line(const std::vector & widths) { - int total_width = 0; - for (int w : widths) { - total_width += (w > 0 ? w : 1); - } - - if (total_width > 0) { - std::string spaces(total_width, ' '); - fwrite(spaces.c_str(), 1, total_width, out); - move_cursor(-total_width); - } - } - - static void set_line_contents(std::string new_line, std::string & line, std::vector & widths, size_t & char_pos, - size_t & byte_pos, int cursor_byte_pos = -1) { - move_to_line_start(char_pos, byte_pos, widths); - clear_current_line(widths); - - line = std::move(new_line); - widths.clear(); - byte_pos = 0; - char_pos = 0; - - size_t idx = 0; - int back_width = 0; - while (idx < line.size()) { - size_t advance = 0; - char32_t cp = decode_utf8(line, idx, advance); - int expected_width = estimateWidth(cp); - int real_width = put_codepoint(line.c_str() + idx, advance, expected_width); - if (real_width < 0) real_width = 0; - widths.push_back(real_width); - idx += advance; - if (cursor_byte_pos >= 0 && static_cast(cursor_byte_pos) < idx) { - back_width += real_width; - } else { - ++char_pos; - byte_pos = idx; - } - } - if (cursor_byte_pos >= 0) { - move_cursor(-back_width); - } - } - - static void move_to_line_start(size_t & char_pos, size_t & byte_pos, const std::vector & widths) { - int back_width = 0; - for (size_t i = 0; i < char_pos; ++i) { - back_width += widths[i]; - } - move_cursor(-back_width); - char_pos = 0; - byte_pos = 0; - } - - static void move_to_line_end(size_t & char_pos, size_t & byte_pos, const std::vector & widths, const std::string & line) { - int forward_width = 0; - for (size_t i = char_pos; i < widths.size(); ++i) { - forward_width += widths[i]; - } - move_cursor(forward_width); - char_pos = widths.size(); - byte_pos = line.length(); - } - - static bool has_ctrl_modifier(const std::string & params) { - size_t start = 0; - while (start < params.size()) { - size_t end = params.find(';', start); - size_t len = (end == std::string::npos) ? params.size() - start : end - start; - if (len > 0) { - int value = 0; - for (size_t i = 0; i < len; ++i) { - char ch = params[start + i]; - if (!std::isdigit(static_cast(ch))) { - value = -1; - break; - } - value = value * 10 + (ch - '0'); - } - if (value == 5) { - return true; - } - } - - if (end == std::string::npos) { - break; - } - start = end + 1; - } - return false; - } - - static bool is_space_codepoint(char32_t cp) { - return std::iswspace(static_cast(cp)) != 0; - } - - static void move_word_left(size_t & char_pos, size_t & byte_pos, const std::vector & widths, const std::string & line) { - if (char_pos == 0) { - return; - } - - size_t new_char_pos = char_pos; - size_t new_byte_pos = byte_pos; - int move_width = 0; - - while (new_char_pos > 0) { - size_t prev_byte = prev_utf8_char_pos(line, new_byte_pos); - size_t advance = 0; - char32_t cp = decode_utf8(line, prev_byte, advance); - if (!is_space_codepoint(cp)) { - break; - } - move_width += widths[new_char_pos - 1]; - new_char_pos--; - new_byte_pos = prev_byte; - } - - while (new_char_pos > 0) { - size_t prev_byte = prev_utf8_char_pos(line, new_byte_pos); - size_t advance = 0; - char32_t cp = decode_utf8(line, prev_byte, advance); - if (is_space_codepoint(cp)) { - break; - } - move_width += widths[new_char_pos - 1]; - new_char_pos--; - new_byte_pos = prev_byte; - } - - move_cursor(-move_width); - char_pos = new_char_pos; - byte_pos = new_byte_pos; - } - - static void move_word_right(size_t & char_pos, size_t & byte_pos, const std::vector & widths, const std::string & line) { - if (char_pos >= widths.size()) { - return; - } - - size_t new_char_pos = char_pos; - size_t new_byte_pos = byte_pos; - int move_width = 0; - - while (new_char_pos < widths.size()) { - size_t advance = 0; - char32_t cp = decode_utf8(line, new_byte_pos, advance); - if (!is_space_codepoint(cp)) { - break; - } - move_width += widths[new_char_pos]; - new_char_pos++; - new_byte_pos += advance; - } - - while (new_char_pos < widths.size()) { - size_t advance = 0; - char32_t cp = decode_utf8(line, new_byte_pos, advance); - if (is_space_codepoint(cp)) { - break; - } - move_width += widths[new_char_pos]; - new_char_pos++; - new_byte_pos += advance; - } - - while (new_char_pos < widths.size()) { - size_t advance = 0; - char32_t cp = decode_utf8(line, new_byte_pos, advance); - if (!is_space_codepoint(cp)) { - break; - } - move_width += widths[new_char_pos]; - new_char_pos++; - new_byte_pos += advance; - } - - move_cursor(move_width); - char_pos = new_char_pos; - byte_pos = new_byte_pos; - } - - static void move_cursor(int delta) { - if (delta == 0) return; -#if defined(_WIN32) - if (hConsole != NULL) { - CONSOLE_SCREEN_BUFFER_INFO bufferInfo; - GetConsoleScreenBufferInfo(hConsole, &bufferInfo); - COORD newCursorPosition = bufferInfo.dwCursorPosition; - int width = bufferInfo.dwSize.X; - int newX = newCursorPosition.X + delta; - int newY = newCursorPosition.Y; - - while (newX >= width) { - newX -= width; - newY++; - } - while (newX < 0) { - newX += width; - newY--; - } - - newCursorPosition.X = newX; - newCursorPosition.Y = newY; - SetConsoleCursorPosition(hConsole, newCursorPosition); - } -#else - if (delta < 0) { - for (int i = 0; i < -delta; i++) fprintf(out, "\b"); - } else { - for (int i = 0; i < delta; i++) fprintf(out, "\033[C"); - } -#endif - } - - struct history_t { - std::vector entries; - size_t viewing_idx = SIZE_MAX; - std::string backup_line; // current line before viewing history - void add(std::string_view line) { - if (line.empty()) { - return; - } - // avoid duplicates with the last entry - if (entries.empty() || entries.back() != line) { - entries.emplace_back(line); - } - // also clear viewing state - end_viewing(); - } - bool prev(std::string & cur_line) { - if (entries.empty()) { - return false; - } - if (viewing_idx == SIZE_MAX) { - return false; - } - if (viewing_idx > 0) { - viewing_idx--; - } - cur_line = entries[viewing_idx]; - return true; - } - bool next(std::string & cur_line) { - if (entries.empty() || viewing_idx == SIZE_MAX) { - return false; - } - viewing_idx++; - if (viewing_idx >= entries.size()) { - cur_line = backup_line; - end_viewing(); - } else { - cur_line = entries[viewing_idx]; - } - return true; - } - void begin_viewing(const std::string & line) { - backup_line = line; - viewing_idx = entries.size(); - } - void end_viewing() { - viewing_idx = SIZE_MAX; - backup_line.clear(); - } - bool is_viewing() const { - return viewing_idx != SIZE_MAX; - } - } history; - - static bool readline_advanced(std::string & line, bool multiline_input) { - if (out != stdout) { - fflush(stdout); - } - - line.clear(); - std::vector widths; - bool is_special_char = false; - bool end_of_stream = false; - - size_t byte_pos = 0; // current byte index - size_t char_pos = 0; // current character index (one char can be multiple bytes) - - char32_t input_char; - while (true) { - assert(char_pos <= byte_pos); - assert(char_pos <= widths.size()); - auto history_prev = [&]() { - if (!history.is_viewing()) { - history.begin_viewing(line); - } - std::string new_line; - if (!history.prev(new_line)) { - return; - } - set_line_contents(new_line, line, widths, char_pos, byte_pos); - }; - auto history_next = [&]() { - if (history.is_viewing()) { - std::string new_line; - if (!history.next(new_line)) { - return; - } - set_line_contents(new_line, line, widths, char_pos, byte_pos); - } - }; - - fflush(out); // Ensure all output is displayed before waiting for input - input_char = getchar32(); - - if (input_char == '\r' || input_char == '\n') { - break; - } - - if (completion_cb && input_char == '\t') { - auto candidates = completion_cb(line, byte_pos); - - if (!candidates.empty()) { - if (candidates.size() > 1 || candidates[0].first != line) { - // TODO?: Display all candidates - set_line_contents(candidates[0].first, line, widths, char_pos, byte_pos, candidates[0].second); - } else { - // TODO: Move cursor to new byte_pos - } - continue; - } - } - - if (input_char == (char32_t) WEOF || input_char == 0x04 /* Ctrl+D */) { - end_of_stream = true; - break; - } - - if (is_special_char) { - replace_last(line.back()); - is_special_char = false; - } - - if (input_char == '\033') { // Escape sequence - char32_t code = getchar32(); - if (code == '[') { - std::string params; - while (true) { - code = getchar32(); - if ((code >= 'A' && code <= 'Z') || (code >= 'a' && code <= 'z') || code == '~' || code == (char32_t) WEOF) { - break; - } - params.push_back(static_cast(code)); - } - - const bool ctrl_modifier = has_ctrl_modifier(params); - - if (code == 'D') { // left - if (ctrl_modifier) { - move_word_left(char_pos, byte_pos, widths, line); - } else if (char_pos > 0) { - int w = widths[char_pos - 1]; - move_cursor(-w); - char_pos--; - byte_pos = prev_utf8_char_pos(line, byte_pos); - } - } else if (code == 'C') { // right - if (ctrl_modifier) { - move_word_right(char_pos, byte_pos, widths, line); - } else if (char_pos < widths.size()) { - int w = widths[char_pos]; - move_cursor(w); - char_pos++; - byte_pos = next_utf8_char_pos(line, byte_pos); - } - } else if (code == 'H') { // home - move_to_line_start(char_pos, byte_pos, widths); - } else if (code == 'F') { // end - move_to_line_end(char_pos, byte_pos, widths, line); - } else if (code == 'A' || code == 'B') { - // up/down - if (code == 'A') { - history_prev(); - is_special_char = false; - } else if (code == 'B') { - history_next(); - is_special_char = false; - } - } else if ((code == '~' || (code >= 'A' && code <= 'Z') || (code >= 'a' && code <= 'z')) && !params.empty()) { - std::string digits; - for (char ch : params) { - if (ch == ';') { - break; - } - if (std::isdigit(static_cast(ch))) { - digits.push_back(ch); - } - } - - if (code == '~') { - if (digits == "1" || digits == "7") { // home - move_to_line_start(char_pos, byte_pos, widths); - } else if (digits == "4" || digits == "8") { // end - move_to_line_end(char_pos, byte_pos, widths, line); - } else if (digits == "3") { // delete - delete_at_cursor(line, widths, char_pos, byte_pos); - } - } - } - } else if (code == 0x1B) { - // Discard the rest of the escape sequence - while ((code = getchar32()) != (char32_t) WEOF) { - if ((code >= 'A' && code <= 'Z') || (code >= 'a' && code <= 'z') || code == '~') { - break; - } - } - } -#if defined(_WIN32) - } else if (input_char == KEY_ARROW_LEFT) { - if (char_pos > 0) { - int w = widths[char_pos - 1]; - move_cursor(-w); - char_pos--; - byte_pos = prev_utf8_char_pos(line, byte_pos); - } - } else if (input_char == KEY_ARROW_RIGHT) { - if (char_pos < widths.size()) { - int w = widths[char_pos]; - move_cursor(w); - char_pos++; - byte_pos = next_utf8_char_pos(line, byte_pos); - } - } else if (input_char == KEY_CTRL_ARROW_LEFT) { - move_word_left(char_pos, byte_pos, widths, line); - } else if (input_char == KEY_CTRL_ARROW_RIGHT) { - move_word_right(char_pos, byte_pos, widths, line); - } else if (input_char == KEY_HOME) { - move_to_line_start(char_pos, byte_pos, widths); - } else if (input_char == KEY_END) { - move_to_line_end(char_pos, byte_pos, widths, line); - } else if (input_char == KEY_DELETE) { - delete_at_cursor(line, widths, char_pos, byte_pos); - } else if (input_char == KEY_ARROW_UP || input_char == KEY_ARROW_DOWN) { - if (input_char == KEY_ARROW_UP) { - history_prev(); - is_special_char = false; - } else if (input_char == KEY_ARROW_DOWN) { - history_next(); - is_special_char = false; - } -#endif - } else if (input_char == 0x08 || input_char == 0x7F) { // Backspace - if (char_pos > 0) { - int w = widths[char_pos - 1]; - move_cursor(-w); - char_pos--; - size_t prev_pos = prev_utf8_char_pos(line, byte_pos); - size_t char_len = byte_pos - prev_pos; - byte_pos = prev_pos; - - // remove the character - line.erase(byte_pos, char_len); - widths.erase(widths.begin() + char_pos); - - // redraw tail - size_t p = byte_pos; - int tail_width = 0; - for (size_t i = char_pos; i < widths.size(); ++i) { - size_t next_p = next_utf8_char_pos(line, p); - put_codepoint(line.c_str() + p, next_p - p, widths[i]); - tail_width += widths[i]; - p = next_p; - } - - // clear display - for (int i = 0; i < w; ++i) { - fputc(' ', out); - } - move_cursor(-(tail_width + w)); - } - } else { - // insert character - std::string new_char_str; - append_utf8(input_char, new_char_str); - int w = estimateWidth(input_char); - - if (char_pos == widths.size()) { - // insert at the end - line += new_char_str; - int real_w = put_codepoint(new_char_str.c_str(), new_char_str.length(), w); - if (real_w < 0) real_w = 0; - widths.push_back(real_w); - byte_pos += new_char_str.length(); - char_pos++; - } else { - // insert in middle - line.insert(byte_pos, new_char_str); - - int real_w = put_codepoint(new_char_str.c_str(), new_char_str.length(), w); - if (real_w < 0) real_w = 0; - - widths.insert(widths.begin() + char_pos, real_w); - - // print the tail - size_t p = byte_pos + new_char_str.length(); - int tail_width = 0; - for (size_t i = char_pos + 1; i < widths.size(); ++i) { - size_t next_p = next_utf8_char_pos(line, p); - put_codepoint(line.c_str() + p, next_p - p, widths[i]); - tail_width += widths[i]; - p = next_p; - } - - move_cursor(-tail_width); - - byte_pos += new_char_str.length(); - char_pos++; - } - } - - if (!line.empty() && (line.back() == '\\' || line.back() == '/')) { - replace_last(line.back()); - is_special_char = true; - } - } - - bool has_more = multiline_input; - if (is_special_char) { - replace_last(' '); - pop_cursor(); - - char last = line.back(); - line.pop_back(); - if (last == '\\') { - line += '\n'; - fputc('\n', out); - has_more = !has_more; - } else { - // llama will just eat the single space, it won't act as a space - if (line.length() == 1 && line.back() == ' ') { - line.clear(); - pop_cursor(); - } - has_more = false; - } - } else { - if (end_of_stream) { - has_more = false; - } else { - line += '\n'; - fputc('\n', out); - } - } - - if (!end_of_stream && !line.empty()) { - // remove the trailing newline for history storage - std::string_view hline = line; - if (!line.empty() && line.back() == '\n') { - hline.remove_suffix(1); - } - // TODO: maybe support multiline history entries? - history.add(hline); - } - - fflush(out); - return has_more; - } - - static bool readline_simple(std::string & line, bool multiline_input) { -#if defined(_WIN32) - std::wstring wline; - if (!std::getline(std::wcin, wline)) { - // Input stream is bad or EOF received - line.clear(); - GenerateConsoleCtrlEvent(CTRL_C_EVENT, 0); - return false; - } - - int size_needed = WideCharToMultiByte(CP_UTF8, 0, &wline[0], (int)wline.size(), NULL, 0, NULL, NULL); - line.resize(size_needed); - WideCharToMultiByte(CP_UTF8, 0, &wline[0], (int)wline.size(), &line[0], size_needed, NULL, NULL); -#else - if (!std::getline(std::cin, line)) { - // Input stream is bad or EOF received - line.clear(); - return false; - } -#endif - if (!line.empty()) { - char last = line.back(); - if (last == '/') { // Always return control on '/' symbol - line.pop_back(); - return false; - } - if (last == '\\') { // '\\' changes the default action - line.pop_back(); - multiline_input = !multiline_input; - } - } - line += '\n'; - - // By default, continue input if multiline_input is set - return multiline_input; - } - - bool readline(std::string & line, bool multiline_input) { - if (simple_io) { - return readline_simple(line, multiline_input); - } - return readline_advanced(line, multiline_input); - } - - void set_completion_callback(completion_callback cb) { - completion_cb = cb; - } - - namespace spinner { - static const char LOADING_CHARS[] = {'|', '/', '-', '\\'}; - static std::condition_variable cv_stop; - static std::thread th; - static size_t frame = 0; // only modified by one thread - static bool running = false; - static std::mutex mtx; - static auto wait_time = std::chrono::milliseconds(100); - static void draw_next_frame() { - // don't need lock because only one thread modifies running - frame = (frame + 1) % sizeof(LOADING_CHARS); - replace_last(LOADING_CHARS[frame]); - fflush(out); - } - void start() { - std::unique_lock lock(mtx); - if (simple_io || running) { - return; - } - common_log_flush(common_log_main()); - fprintf(out, "%c", LOADING_CHARS[0]); - fflush(out); - frame = 1; - running = true; - th = std::thread([]() { - std::unique_lock lock(mtx); - while (true) { - if (cv_stop.wait_for(lock, wait_time, []{ return !running; })) { - break; - } - draw_next_frame(); - } - }); - } - void stop() { - { - std::unique_lock lock(mtx); - if (simple_io || !running) { - return; - } - running = false; - cv_stop.notify_all(); - } - if (th.joinable()) { - th.join(); - } - replace_last(' '); - pop_cursor(); - fflush(out); - } - } - - void log(const char * fmt, ...) { - va_list args; - va_start(args, fmt); - vfprintf(out, fmt, args); - va_end(args); - } - - void error(const char * fmt, ...) { - va_list args; - va_start(args, fmt); - display_type cur = current_display; - set_display(DISPLAY_TYPE_ERROR); - vfprintf(out, fmt, args); - set_display(cur); // restore previous color - va_end(args); - } - - void flush() { - fflush(out); - } -} diff --git a/common/console.h b/common/console.h deleted file mode 100644 index 72781bea6f60..000000000000 --- a/common/console.h +++ /dev/null @@ -1,46 +0,0 @@ -// Console functions - -#pragma once - -#include "common.h" - -#include -#include -#include - -enum display_type { - DISPLAY_TYPE_RESET = 0, - DISPLAY_TYPE_INFO, - DISPLAY_TYPE_PROMPT, - DISPLAY_TYPE_REASONING, - DISPLAY_TYPE_USER_INPUT, - DISPLAY_TYPE_ERROR -}; - -namespace console { - void init(bool use_simple_io, bool use_advanced_display); - void cleanup(); - void set_display(display_type display); - bool readline(std::string & line, bool multiline_input); - - using completion_callback = std::function>(std::string_view, size_t)>; - void set_completion_callback(completion_callback cb); - - namespace spinner { - void start(); - void stop(); - } - - // note: the logging API below output directly to stdout - // it can negatively impact performance if used on inference thread - // only use in in a dedicated CLI thread - // for logging in inference thread, use log.h instead - - LLAMA_COMMON_ATTRIBUTE_FORMAT(1, 2) - void log(const char * fmt, ...); - - LLAMA_COMMON_ATTRIBUTE_FORMAT(1, 2) - void error(const char * fmt, ...); - - void flush(); -} diff --git a/common/debug.cpp b/common/debug.cpp deleted file mode 100644 index 102c6924dc9e..000000000000 --- a/common/debug.cpp +++ /dev/null @@ -1,190 +0,0 @@ -#include "debug.h" - -#include "common.h" -#include "log.h" - -#include -#include -#include -#include - -struct common_debug_cb_user_data::impl { - std::vector data; - std::vector tensor_filters; - bool abort_on_nan{false}; -}; - -common_debug_cb_user_data::common_debug_cb_user_data() : pimpl(std::make_unique()) {} -common_debug_cb_user_data::~common_debug_cb_user_data() = default; - -common_debug_cb_user_data::common_debug_cb_user_data(common_params & params, const std::vector & filter_patterns, bool abort_on_nan) - : pimpl(std::make_unique()) -{ - for (const auto & pattern : filter_patterns) { - try { - std::string anchored_pattern = "^" + pattern; - pimpl->tensor_filters.emplace_back(anchored_pattern, std::regex::optimize); - } catch (const std::regex_error & e) { - throw std::runtime_error("Invalid regex pattern '" + pattern + "': " + e.what()); - } - } - pimpl->abort_on_nan = abort_on_nan; - - params.cb_eval = common_debug_cb_eval; - params.cb_eval_user_data = this; -} - -static std::string common_ggml_ne_string(const ggml_tensor * t) { - std::string str; - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - str += std::to_string(t->ne[i]); - if (i + 1 < GGML_MAX_DIMS) { - str += ", "; - } - } - return str; -} - -static float common_ggml_get_float_value(const uint8_t * data, - ggml_type type, - const size_t * nb, - size_t i0, - size_t i1, - size_t i2, - size_t i3) { - size_t i = i3 * nb[3] + i2 * nb[2] + i1 * nb[1] + i0 * nb[0]; - float v; - if (type == GGML_TYPE_F16) { - v = ggml_fp16_to_fp32(*(const ggml_fp16_t *) &data[i]); - } else if (type == GGML_TYPE_F32) { - v = *(const float *) &data[i]; - } else if (type == GGML_TYPE_I64) { - v = (float) *(const int64_t *) &data[i]; - } else if (type == GGML_TYPE_I32) { - v = (float) *(const int32_t *) &data[i]; - } else if (type == GGML_TYPE_I16) { - v = (float) *(const int16_t *) &data[i]; - } else if (type == GGML_TYPE_I8) { - v = (float) *(const int8_t *) &data[i]; - } else if (type == GGML_TYPE_BF16) { - v = ggml_bf16_to_fp32(*(const ggml_bf16_t *) &data[i]); - } else { - GGML_ABORT("fatal error"); - } - return v; -} - -#define INDENT " " - -static void common_debug_print_tensor(uint8_t * data, ggml_type type, const int64_t * ne, const size_t * nb, int64_t n, bool abort_on_nan) { - GGML_ASSERT(n > 0); - float sum = 0; - for (int64_t i3 = 0; i3 < ne[3]; i3++) { - for (int64_t i2 = 0; i2 < ne[2]; i2++) { - for (int64_t i1 = 0; i1 < ne[1]; i1++) { - for (int64_t i0 = 0; i0 < ne[0]; i0++) { - const float v = common_ggml_get_float_value(data, type, nb, i0, i1, i2, i3); - sum += v; - } - } - } - } - for (int64_t i3 = 0; i3 < ne[3]; i3++) { - LOG(INDENT "[\n"); - for (int64_t i2 = 0; i2 < ne[2]; i2++) { - if (i2 == n && ne[2] > 2 * n) { - LOG(INDENT INDENT "..., \n"); - i2 = ne[2] - n; - } - LOG(INDENT INDENT "[\n"); - for (int64_t i1 = 0; i1 < ne[1]; i1++) { - if (i1 == n && ne[1] > 2 * n) { - LOG(INDENT INDENT INDENT "..., \n"); - i1 = ne[1] - n; - } - LOG(INDENT INDENT INDENT "["); - for (int64_t i0 = 0; i0 < ne[0]; i0++) { - if (i0 == n && ne[0] > 2 * n) { - LOG(" ..., "); - i0 = ne[0] - n; - } - const float v = common_ggml_get_float_value(data, type, nb, i0, i1, i2, i3); - LOG("%12.4f", v); - if (i0 < ne[0] - 1) { - LOG(", "); - } - } - LOG(" ],\n"); - } - LOG(INDENT INDENT "],\n"); - } - LOG(INDENT "]\n"); - LOG(INDENT "sum = %f\n", sum); - } - - if (abort_on_nan) { - if (std::isnan(sum)) { - LOG("encountered NaN - aborting\n"); - exit(0); - } - } -} - -/** - * GGML operations callback during the graph execution. - * - * @param t current tensor - * @param ask when ask is true, the scheduler wants to know if we are interested in data from this tensor - * if we return true, a follow-up call will be made with ask=false in which we can do the actual collection. - * see ggml_backend_sched_eval_callback - * @param user_data user data to pass at each call back - * @return true to receive data or continue the graph, false otherwise - */ -bool common_debug_cb_eval(struct ggml_tensor * t, bool ask, void * user_data) { - auto * cb_data = (common_debug_cb_user_data *) user_data; - auto * pimpl = cb_data->pimpl.get(); - - const struct ggml_tensor * src0 = t->src[0]; - const struct ggml_tensor * src1 = t->src[1]; - - if (ask) { - return true; // Always retrieve data - } - - bool matches_filter = pimpl->tensor_filters.empty(); - - if (!matches_filter) { - for (const auto & filter : pimpl->tensor_filters) { - if (std::regex_search(t->name, filter)) { - matches_filter = true; - break; - } - } - } - - char src1_str[128] = { 0 }; - if (src1) { - snprintf(src1_str, sizeof(src1_str), "%s{%s}", src1->name, common_ggml_ne_string(src1).c_str()); - } - - if (matches_filter) { - LOG("%s: %24s = (%s) %10s(%s{%s}, %s}) = {%s}\n", __func__, t->name, ggml_type_name(t->type), - ggml_op_desc(t), src0->name, common_ggml_ne_string(src0).c_str(), src1 ? src1_str : "", - common_ggml_ne_string(t).c_str()); - } - - const bool is_host = ggml_backend_buffer_is_host(t->buffer); - - if (!is_host) { - auto n_bytes = ggml_nbytes(t); - pimpl->data.resize(n_bytes); - ggml_backend_tensor_get(t, pimpl->data.data(), 0, n_bytes); - } - - if (!ggml_is_quantized(t->type) && matches_filter) { - uint8_t * data = is_host ? (uint8_t *) t->data : pimpl->data.data(); - common_debug_print_tensor(data, t->type, t->ne, t->nb, 3, pimpl->abort_on_nan); - } - - return true; -} diff --git a/common/debug.h b/common/debug.h deleted file mode 100644 index 8b8f8c7aa9fe..000000000000 --- a/common/debug.h +++ /dev/null @@ -1,31 +0,0 @@ -#pragma once - -#include -#include -#include - -// common debug functions and structs - -struct common_params; - -// Intended to use as callback for ggml_backend_sched_eval_callback -// prints tensors that are processed in the computation graph -// by default prints all tensors, but can be configured by creating a `common_debug_cb_user_data` instance with -// non-empty filter_patterns. See examples/debug.cpp for possible usage patterns -// `common_debug_cb_user_data` contains `abort_on_nan` flag that determines whether an error should be thrown whenever a NaN is encountered -// in a tensor (useful for stopping debug sessions on first erroneous tensor) -// The callback data will be passed as the third parameter (user_data) -bool common_debug_cb_eval(struct ggml_tensor * t, bool ask, void * user_data); - -struct common_debug_cb_user_data { - struct impl; - std::unique_ptr pimpl; - - common_debug_cb_user_data(); - ~common_debug_cb_user_data(); - - common_debug_cb_user_data(const common_debug_cb_user_data &) = delete; - common_debug_cb_user_data & operator=(const common_debug_cb_user_data &) = delete; - - common_debug_cb_user_data(common_params & params, const std::vector & filter_patterns, bool abort_on_nan = false); -}; diff --git a/common/download.cpp b/common/download.cpp deleted file mode 100644 index d6d47b2d2fcf..000000000000 --- a/common/download.cpp +++ /dev/null @@ -1,958 +0,0 @@ -#include "arg.h" - -#include "build-info.h" -#include "common.h" -#include "log.h" -#include "download.h" -#include "hf-cache.h" - -#define JSON_ASSERT GGML_ASSERT -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "http.h" - -#ifndef __EMSCRIPTEN__ -#ifdef __linux__ -#include -#elif defined(_WIN32) -# if !defined(PATH_MAX) -# define PATH_MAX MAX_PATH -# endif -#elif defined(_AIX) -#include -#else -#include -#endif -#endif - -// isatty -#if defined(_WIN32) -#include -#else -#include -#endif - -using json = nlohmann::ordered_json; - -// -// downloader -// - -// validate repo name format: owner/repo -static void write_file(const std::string & fname, const std::string & content) { - const std::string fname_tmp = fname + ".tmp"; - std::ofstream file(fname_tmp); - if (!file) { - throw std::runtime_error(string_format("error: failed to open file '%s'\n", fname.c_str())); - } - - try { - file << content; - file.close(); - - // Makes write atomic - if (rename(fname_tmp.c_str(), fname.c_str()) != 0) { - LOG_ERR("%s: unable to rename file: %s to %s\n", __func__, fname_tmp.c_str(), fname.c_str()); - // If rename fails, try to delete the temporary file - if (remove(fname_tmp.c_str()) != 0) { - LOG_ERR("%s: unable to delete temporary file: %s\n", __func__, fname_tmp.c_str()); - } - } - } catch (...) { - // If anything fails, try to delete the temporary file - if (remove(fname_tmp.c_str()) != 0) { - LOG_ERR("%s: unable to delete temporary file: %s\n", __func__, fname_tmp.c_str()); - } - - throw std::runtime_error(string_format("error: failed to write file '%s'\n", fname.c_str())); - } -} - -static void write_etag(const std::string & path, const std::string & etag) { - const std::string etag_path = path + ".etag"; - write_file(etag_path, etag); - LOG_DBG("%s: file etag saved: %s\n", __func__, etag_path.c_str()); -} - -static std::string read_etag(const std::string & path) { - const std::string etag_path = path + ".etag"; - if (!std::filesystem::exists(etag_path)) { - return {}; - } - std::ifstream etag_in(etag_path); - if (!etag_in) { - LOG_ERR("%s: could not open .etag file for reading: %s\n", __func__, etag_path.c_str()); - return {}; - } - std::string etag; - std::getline(etag_in, etag); - return etag; -} - -static bool is_http_status_ok(int status) { - return status >= 200 && status < 400; -} - -std::pair common_download_split_repo_tag(const std::string & hf_repo_with_tag) { - auto parts = string_split(hf_repo_with_tag, ':'); - std::string tag = parts.size() > 1 ? parts.back() : ""; - std::string hf_repo = parts[0]; - if (string_split(hf_repo, '/').size() != 2) { - throw std::invalid_argument("error: invalid HF repo format, expected /[:quant]\n"); - } - return {hf_repo, tag}; -} - -class ProgressBar : public common_download_callback { - static inline std::mutex mutex; - static inline std::map lines; - static inline int max_line = 0; - - std::string filename; - size_t len = 0; - - static void cleanup(const ProgressBar * line) { - lines.erase(line); - if (lines.empty()) { - max_line = 0; - } - } - - static bool is_output_a_tty() { -#if defined(_WIN32) - return _isatty(_fileno(stdout)); -#else - return isatty(1); -#endif - } - -public: - ProgressBar() = default; - - void on_start(const common_download_progress & p) override { - filename = p.url; - - if (auto pos = filename.rfind('/'); pos != std::string::npos) { - filename = filename.substr(pos + 1); - } - if (auto pos = filename.find('?'); pos != std::string::npos) { - filename = filename.substr(0, pos); - } - for (size_t i = 0; i < filename.size(); ++i) { - if ((filename[i] & 0xC0) != 0x80) { - if (len++ == 39) { - filename.resize(i); - filename += "…"; - break; - } - } - } - } - - void on_done(const common_download_progress &, bool) override { - std::lock_guard lock(mutex); - cleanup(this); - } - - void on_update(const common_download_progress & p) override { - if (!p.total || !is_output_a_tty()) { - return; - } - - std::lock_guard lock(mutex); - - if (lines.find(this) == lines.end()) { - lines[this] = max_line++; - std::cout << "\n"; - } - int lines_up = max_line - lines[this]; - - size_t bar = (55 - len) * 2; - size_t pct = (100 * p.downloaded) / p.total; - size_t pos = (bar * p.downloaded) / p.total; - - if (lines_up > 0) { - std::cout << "\033[" << lines_up << "A"; - } - std::cout << '\r' << "Downloading " << filename << " "; - - for (size_t i = 0; i < bar; i += 2) { - std::cout << (i + 1 < pos ? "─" : (i < pos ? "╴" : " ")); - } - std::cout << std::setw(4) << pct << "%\033[K"; - - if (lines_up > 0) { - std::cout << "\033[" << lines_up << "B"; - } - std::cout << '\r' << std::flush; - - if (p.downloaded == p.total) { - cleanup(this); - } - } - - ProgressBar(const ProgressBar &) = delete; - ProgressBar & operator=(const ProgressBar &) = delete; -}; - -static bool common_pull_file(httplib::Client & cli, - const std::string & resolve_path, - const std::string & path_tmp, - bool supports_ranges, - common_download_progress & p, - common_download_callback * callback) { - std::ofstream ofs(path_tmp, std::ios::binary | std::ios::app); - if (!ofs.is_open()) { - LOG_ERR("%s: error opening local file for writing: %s\n", __func__, path_tmp.c_str()); - return false; - } - - httplib::Headers headers; - if (supports_ranges && p.downloaded > 0) { - headers.emplace("Range", "bytes=" + std::to_string(p.downloaded) + "-"); - } - - const char * func = __func__; // avoid __func__ inside a lambda - size_t progress_step = 0; - - auto res = cli.Get(resolve_path, headers, - [&](const httplib::Response &response) { - if (p.downloaded > 0 && response.status != 206) { - LOG_WRN("%s: server did not respond with 206 Partial Content for a resume request. Status: %d\n", func, response.status); - return false; - } - if (p.downloaded == 0 && response.status != 200) { - LOG_WRN("%s: download received non-successful status code: %d\n", func, response.status); - return false; - } - if (p.total == 0 && response.has_header("Content-Length")) { - try { - size_t content_length = std::stoull(response.get_header_value("Content-Length")); - p.total = p.downloaded + content_length; - } catch (const std::exception &e) { - LOG_WRN("%s: invalid Content-Length header: %s\n", func, e.what()); - } - } - return true; - }, - [&](const char *data, size_t len) { - ofs.write(data, len); - if (!ofs) { - LOG_ERR("%s: error writing to file: %s\n", func, path_tmp.c_str()); - return false; - } - p.downloaded += len; - progress_step += len; - - if (progress_step >= p.total / 1000 || p.downloaded == p.total) { - if (callback) { - callback->on_update(p); - if (callback->is_cancelled()) { - return false; - } - } - progress_step = 0; - } - return true; - }, - nullptr - ); - - if (!res) { - LOG_ERR("%s: download failed: %s (status: %d)\n", - __func__, - httplib::to_string(res.error()).c_str(), - res ? res->status : -1); - return false; - } - - return true; -} - -// download one single file from remote URL to local path -// returns status code or -1 on error -static int common_download_file_single_online(const std::string & url, - const std::string & path, - const common_download_opts & opts, - bool skip_etag) { - static const int max_attempts = 3; - static const int retry_delay_seconds = 2; - - const bool file_exists = std::filesystem::exists(path); - - if (file_exists && skip_etag) { - LOG_DBG("%s: using cached file: %s\n", __func__, path.c_str()); - return 304; // 304 Not Modified - fake cached response - } - - auto [cli, parts] = common_http_client(url); - - httplib::Headers headers; - for (const auto & h : opts.headers) { - headers.emplace(h.first, h.second); - } - if (headers.find("User-Agent") == headers.end()) { - headers.emplace("User-Agent", "llama-cpp/" + std::string(llama_build_info())); - } - if (!opts.bearer_token.empty()) { - headers.emplace("Authorization", "Bearer " + opts.bearer_token); - } - cli.set_default_headers(headers); - - std::string last_etag; - if (file_exists) { - last_etag = read_etag(path); - } else { - LOG_DBG("%s: no previous model file found %s\n", __func__, path.c_str()); - } - - auto head = cli.Head(parts.path); - if (!head || head->status < 200 || head->status >= 300) { - LOG_WRN("%s: HEAD failed, status: %d\n", __func__, head ? head->status : -1); - if (file_exists) { - LOG_INF("%s: using cached file (HEAD failed): %s\n", __func__, path.c_str()); - return 304; // 304 Not Modified - fake cached response - } - return head ? head->status : -1; - } - - std::string etag; - if (head->has_header("ETag")) { - etag = head->get_header_value("ETag"); - } - - common_download_progress p; - p.url = url; - if (head->has_header("Content-Length")) { - try { - p.total = std::stoull(head->get_header_value("Content-Length")); - } catch (const std::exception& e) { - LOG_WRN("%s: invalid Content-Length in HEAD response: %s\n", __func__, e.what()); - } - } - - bool supports_ranges = false; - if (head->has_header("Accept-Ranges")) { - supports_ranges = head->get_header_value("Accept-Ranges") != "none"; - } - - if (file_exists) { - if (etag.empty()) { - LOG_DBG("%s: using cached file (no server etag): %s\n", __func__, path.c_str()); - return 304; // 304 Not Modified - fake cached response - } - if (!last_etag.empty() && last_etag == etag) { - LOG_DBG("%s: using cached file (same etag): %s\n", __func__, path.c_str()); - return 304; // 304 Not Modified - fake cached response - } - if (remove(path.c_str()) != 0) { - LOG_ERR("%s: unable to delete file: %s\n", __func__, path.c_str()); - return -1; - } - } - - { // silent - std::error_code ec; - std::filesystem::create_directories(std::filesystem::path(path).parent_path(), ec); - } - - bool success = false; - const std::string path_temporary = path + ".downloadInProgress"; - int delay = retry_delay_seconds; - - if (opts.callback) { - opts.callback->on_start(p); - } - - for (int i = 0; i < max_attempts; ++i) { - if (opts.callback && opts.callback->is_cancelled()) { - break; - } - if (i) { - LOG_WRN("%s: retrying after %d seconds...\n", __func__, delay); - std::this_thread::sleep_for(std::chrono::seconds(delay)); - delay *= retry_delay_seconds; - } - - size_t existing_size = 0; - - if (std::filesystem::exists(path_temporary)) { - if (supports_ranges) { - existing_size = std::filesystem::file_size(path_temporary); - } else if (remove(path_temporary.c_str()) != 0) { - LOG_ERR("%s: unable to delete file: %s\n", __func__, path_temporary.c_str()); - break; - } - } - - p.downloaded = existing_size; - - LOG_DBG("%s: downloading from %s to %s (etag:%s)...\n", - __func__, common_http_show_masked_url(parts).c_str(), - path_temporary.c_str(), etag.c_str()); - - if (common_pull_file(cli, parts.path, path_temporary, supports_ranges, p, opts.callback)) { - if (std::rename(path_temporary.c_str(), path.c_str()) != 0) { - LOG_ERR("%s: unable to rename file: %s to %s\n", __func__, path_temporary.c_str(), path.c_str()); - break; - } - if (!etag.empty() && !skip_etag) { - write_etag(path, etag); - } - success = true; - break; - } - } - - if (opts.callback) { - opts.callback->on_done(p, success); - } - if (opts.callback && opts.callback->is_cancelled() && - std::filesystem::exists(path_temporary)) { - if (remove(path_temporary.c_str()) != 0) { - LOG_ERR("%s: unable to delete temporary file: %s\n", __func__, path_temporary.c_str()); - } - } - if (!success) { - LOG_ERR("%s: download failed after %d attempts\n", __func__, max_attempts); - return -1; // max attempts reached - } - - return head->status; -} - -std::pair> common_remote_get_content(const std::string & url, - const common_remote_params & params) { - auto [cli, parts] = common_http_client(url); - - httplib::Headers headers; - for (const auto & h : params.headers) { - headers.emplace(h.first, h.second); - } - if (headers.find("User-Agent") == headers.end()) { - headers.emplace("User-Agent", "llama-cpp/" + std::string(llama_build_info())); - } - - if (params.timeout > 0) { - cli.set_read_timeout(params.timeout, 0); - cli.set_write_timeout(params.timeout, 0); - } - - std::vector buf; - auto res = cli.Get(parts.path, headers, - [&](const char *data, size_t len) { - buf.insert(buf.end(), data, data + len); - return params.max_size == 0 || - buf.size() <= static_cast(params.max_size); - }, - nullptr - ); - - if (!res) { - throw std::runtime_error("error: cannot make GET request"); - } - - return { res->status, std::move(buf) }; -} - -int common_download_file_single(const std::string & url, - const std::string & path, - const common_download_opts & opts, - bool skip_etag) { - if (!opts.offline) { - ProgressBar tty_cb; - common_download_opts online_opts = opts; - if (!online_opts.callback) { - online_opts.callback = &tty_cb; - } - return common_download_file_single_online(url, path, online_opts, skip_etag); - } - - if (!std::filesystem::exists(path)) { - LOG_ERR("%s: required file is not available in cache (offline mode): %s\n", __func__, path.c_str()); - return -1; - } - - LOG_DBG("%s: using cached file (offline mode): %s\n", __func__, path.c_str()); - - // notify the callback that the file was cached - if (opts.callback) { - common_download_progress p; - p.url = url; - p.cached = true; - opts.callback->on_start(p); - opts.callback->on_done(p, true); - } - - return 304; // Not Modified - fake cached response -} - -struct gguf_split_info { - std::string prefix; // tag included - std::string tag; - int index; - int count; -}; - -static gguf_split_info get_gguf_split_info(const std::string & path) { - static const std::regex re_split("^(.+)-([0-9]{5})-of-([0-9]{5})$", std::regex::icase); - static const std::regex re_tag("[-.]([A-Z0-9_]+)$", std::regex::icase); - std::smatch m; - - std::string prefix = path; - if (!string_remove_suffix(prefix, ".gguf")) { - return {}; - } - - int index = 1; - int count = 1; - - if (std::regex_match(prefix, m, re_split)) { - index = std::stoi(m[2].str()); - count = std::stoi(m[3].str()); - prefix = m[1].str(); - } - - std::string tag; - if (std::regex_search(prefix, m, re_tag)) { - tag = m[1].str(); - for (char & c : tag) { - c = std::toupper((unsigned char)c); - } - } - - return {std::move(prefix), std::move(tag), index, count}; -} - -// Q4_0 -> 4, F16 -> 16, NVFP4 -> 4, Q8_K_M -> 8, etc -static int extract_quant_bits(const std::string & filename) { - auto split = get_gguf_split_info(filename); - - auto pos = split.tag.find_first_of("0123456789"); - if (pos == std::string::npos) { - return 0; - } - - return std::stoi(split.tag.substr(pos)); -} - -static hf_cache::hf_files get_split_files(const hf_cache::hf_files & files, - const hf_cache::hf_file & file) { - auto split = get_gguf_split_info(file.path); - - if (split.count <= 1) { - return {file}; - } - hf_cache::hf_files result; - - for (const auto & f : files) { - auto split_f = get_gguf_split_info(f.path); - if (split_f.count == split.count && split_f.prefix == split.prefix) { - result.push_back(f); - } - } - return result; -} - -static hf_cache::hf_file find_best_mmproj(const hf_cache::hf_files & files, - const std::string & model) { - hf_cache::hf_file best; - size_t best_depth = 0; - int best_diff = 0; - bool found = false; - - auto model_bits = extract_quant_bits(model); - auto model_parts = string_split(model, '/'); - auto model_dir = model_parts.end() - 1; - - for (const auto & f : files) { - if (!string_ends_with(f.path, ".gguf") || - f.path.find("mmproj") == std::string::npos) { - continue; - } - - auto mmproj_parts = string_split(f.path, '/'); - auto mmproj_dir = mmproj_parts.end() - 1; - - auto [_, dir] = std::mismatch(model_parts.begin(), model_dir, - mmproj_parts.begin(), mmproj_dir); - if (dir != mmproj_dir) { - continue; - } - - size_t depth = dir - mmproj_parts.begin(); - auto bits = extract_quant_bits(f.path); - auto diff = std::abs(bits - model_bits); - - if (!found || depth > best_depth || (depth == best_depth && diff < best_diff)) { - best = f; - best_depth = depth; - best_diff = diff; - found = true; - } - } - return best; -} - -static bool gguf_filename_is_model(const std::string & filepath) { - if (!string_ends_with(filepath, ".gguf")) { - return false; - } - - std::string filename = filepath; - if (auto pos = filename.rfind('/'); pos != std::string::npos) { - filename = filename.substr(pos + 1); - } - - return filename.find("mmproj") == std::string::npos && - filename.find("imatrix") == std::string::npos; -} - -static hf_cache::hf_file find_best_model(const hf_cache::hf_files & files, - const std::string & tag) { - std::vector tags; - - if (!tag.empty()) { - tags.push_back(tag); - } else { - tags = {"Q4_K_M", "Q8_0"}; - } - - for (const auto & t : tags) { - std::regex pattern(t + "[.-]", std::regex::icase); - for (const auto & f : files) { - if (gguf_filename_is_model(f.path) && - std::regex_search(f.path, pattern)) { - auto split = get_gguf_split_info(f.path); - if (split.count > 1 && split.index != 1) { - continue; - } - return f; - } - } - } - - // fallback to first available model only if tag is empty - if (tag.empty()) { - for (const auto & f : files) { - if (gguf_filename_is_model(f.path)) { - auto split = get_gguf_split_info(f.path); - if (split.count > 1 && split.index != 1) { - continue; - } - return f; - } - } - } - - return {}; -} - -static void list_available_gguf_files(const hf_cache::hf_files & files) { - LOG_INF("Available GGUF files:\n"); - for (const auto & f : files) { - if (string_ends_with(f.path, ".gguf")) { - LOG_INF(" - %s\n", f.path.c_str()); - } - } -} - -struct hf_plan { - hf_cache::hf_file primary; - hf_cache::hf_files model_files; - hf_cache::hf_file mmproj; -}; - -static hf_plan get_hf_plan(const common_params_model & model, - const common_download_opts & opts, - bool download_mmproj) { - hf_plan plan; - hf_cache::hf_files all; - - auto [repo, tag] = common_download_split_repo_tag(model.hf_repo); - - if (!opts.offline) { - all = hf_cache::get_repo_files(repo, opts.bearer_token); - } - if (all.empty()) { - all = hf_cache::get_cached_files(repo); - } - if (all.empty()) { - return plan; - } - - hf_cache::hf_file primary; - - if (!model.hf_file.empty()) { - for (const auto & f : all) { - if (f.path == model.hf_file) { - primary = f; - break; - } - } - if (primary.path.empty()) { - LOG_ERR("%s: file '%s' not found in repository\n", __func__, model.hf_file.c_str()); - list_available_gguf_files(all); - return plan; - } - } else { - primary = find_best_model(all, tag); - if (primary.path.empty()) { - LOG_ERR("%s: no GGUF files found in repository %s\n", __func__, repo.c_str()); - list_available_gguf_files(all); - return plan; - } - } - - plan.primary = primary; - plan.model_files = get_split_files(all, primary); - - if (download_mmproj) { - plan.mmproj = find_best_mmproj(all, primary.path); - } - - return plan; -} - -struct download_task { - std::string url; - std::string path; -}; - -static std::vector get_url_tasks(const common_params_model & model) { - auto split = get_gguf_split_info(model.url); - - if (split.count <= 1) { - return {{model.url, model.path}}; - } - - auto filename = split.prefix; - if (auto pos = split.prefix.rfind('/'); pos != std::string::npos) { - filename = split.prefix.substr(pos + 1); - } - - auto parent_path = std::filesystem::path(model.path).parent_path(); - auto prefix_path = (parent_path / filename).string(); - - std::vector tasks; - for (int i = 1; i <= split.count; i++) { - auto suffix = string_format("-%05d-of-%05d.gguf", i, split.count); - tasks.push_back({split.prefix + suffix, prefix_path + suffix}); - } - return tasks; -} - -common_download_model_result common_download_model(const common_params_model & model, - const common_download_opts & opts, - bool download_mmproj) { - common_download_model_result result; - std::vector tasks; - hf_plan hf; - - bool is_hf = !model.hf_repo.empty(); - - if (is_hf) { - hf = get_hf_plan(model, opts, download_mmproj); - for (const auto & f : hf.model_files) { - tasks.push_back({f.url, f.local_path}); - } - if (!hf.mmproj.path.empty()) { - tasks.push_back({hf.mmproj.url, hf.mmproj.local_path}); - } - } else if (!model.url.empty()) { - tasks = get_url_tasks(model); - } else { - result.model_path = model.path; - return result; - } - - if (tasks.empty()) { - return result; - } - - std::vector> futures; - for (const auto & task : tasks) { - futures.push_back(std::async(std::launch::async, - [&task, &opts, is_hf]() { - int status = common_download_file_single(task.url, task.path, opts, is_hf); - return is_http_status_ok(status); - } - )); - } - - for (auto & f : futures) { - if (!f.get()) { - return {}; - } - } - - if (is_hf) { - for (const auto & f : hf.model_files) { - hf_cache::finalize_file(f); - } - result.model_path = hf.primary.final_path; - - if (!hf.mmproj.path.empty()) { - result.mmproj_path = hf_cache::finalize_file(hf.mmproj); - } - } else { - result.model_path = model.path; - } - - return result; -} - -// -// Docker registry functions -// - -static std::string common_docker_get_token(const std::string & repo) { - std::string url = "https://auth.docker.io/token?service=registry.docker.io&scope=repository:" + repo + ":pull"; - - common_remote_params params; - auto res = common_remote_get_content(url, params); - - if (res.first != 200) { - throw std::runtime_error("Failed to get Docker registry token, HTTP code: " + std::to_string(res.first)); - } - - std::string response_str(res.second.begin(), res.second.end()); - nlohmann::ordered_json response = nlohmann::ordered_json::parse(response_str); - - if (!response.contains("token")) { - throw std::runtime_error("Docker registry token response missing 'token' field"); - } - - return response["token"].get(); -} - -std::string common_docker_resolve_model(const std::string & docker) { - // Parse ai/smollm2:135M-Q4_0 - size_t colon_pos = docker.find(':'); - std::string repo, tag; - if (colon_pos != std::string::npos) { - repo = docker.substr(0, colon_pos); - tag = docker.substr(colon_pos + 1); - } else { - repo = docker; - tag = "latest"; - } - - // ai/ is the default - size_t slash_pos = docker.find('/'); - if (slash_pos == std::string::npos) { - repo.insert(0, "ai/"); - } - - LOG_INF("%s: Downloading Docker Model: %s:%s\n", __func__, repo.c_str(), tag.c_str()); - try { - // --- helper: digest validation --- - auto validate_oci_digest = [](const std::string & digest) -> std::string { - // Expected: algo:hex ; start with sha256 (64 hex chars) - // You can extend this map if supporting other algorithms in future. - static const std::regex re("^sha256:([a-fA-F0-9]{64})$"); - std::smatch m; - if (!std::regex_match(digest, m, re)) { - throw std::runtime_error("Invalid OCI digest format received in manifest: " + digest); - } - // normalize hex to lowercase - std::string normalized = digest; - std::transform(normalized.begin()+7, normalized.end(), normalized.begin()+7, [](unsigned char c){ - return std::tolower(c); - }); - return normalized; - }; - - std::string token = common_docker_get_token(repo); // Get authentication token - - // Get manifest - // TODO: cache the manifest response so that it appears in the model list - const std::string url_prefix = "https://registry-1.docker.io/v2/" + repo; - std::string manifest_url = url_prefix + "/manifests/" + tag; - common_remote_params manifest_params; - manifest_params.headers.push_back({"Authorization", "Bearer " + token}); - manifest_params.headers.push_back({"Accept", - "application/vnd.docker.distribution.manifest.v2+json,application/vnd.oci.image.manifest.v1+json" - }); - auto manifest_res = common_remote_get_content(manifest_url, manifest_params); - if (manifest_res.first != 200) { - throw std::runtime_error("Failed to get Docker manifest, HTTP code: " + std::to_string(manifest_res.first)); - } - - std::string manifest_str(manifest_res.second.begin(), manifest_res.second.end()); - nlohmann::ordered_json manifest = nlohmann::ordered_json::parse(manifest_str); - std::string gguf_digest; // Find the GGUF layer - if (manifest.contains("layers")) { - for (const auto & layer : manifest["layers"]) { - if (layer.contains("mediaType")) { - std::string media_type = layer["mediaType"].get(); - if (media_type == "application/vnd.docker.ai.gguf.v3" || - media_type.find("gguf") != std::string::npos) { - gguf_digest = layer["digest"].get(); - break; - } - } - } - } - - if (gguf_digest.empty()) { - throw std::runtime_error("No GGUF layer found in Docker manifest"); - } - - // Validate & normalize digest - gguf_digest = validate_oci_digest(gguf_digest); - LOG_DBG("%s: Using validated digest: %s\n", __func__, gguf_digest.c_str()); - - // Prepare local filename - std::string model_filename = repo; - std::replace(model_filename.begin(), model_filename.end(), '/', '_'); - model_filename += "_" + tag + ".gguf"; - std::string local_path = fs_get_cache_file(model_filename); - - const std::string blob_url = url_prefix + "/blobs/" + gguf_digest; - common_download_opts opts; - opts.bearer_token = token; - const int http_status = common_download_file_single(blob_url, local_path, opts); - if (!is_http_status_ok(http_status)) { - throw std::runtime_error("Failed to download Docker Model"); - } - - LOG_INF("%s: Downloaded Docker Model to: %s\n", __func__, local_path.c_str()); - return local_path; - } catch (const std::exception & e) { - LOG_ERR("%s: Docker Model download failed: %s\n", __func__, e.what()); - throw; - } -} - -std::vector common_list_cached_models() { - std::unordered_set seen; - std::vector result; - - auto files = hf_cache::get_cached_files(); - - for (const auto & f : files) { - auto split = get_gguf_split_info(f.path); - if (split.index != 1 || split.tag.empty() || - split.prefix.find("mmproj") != std::string::npos) { - continue; - } - if (seen.insert(f.repo_id + ":" + split.tag).second) { - result.push_back({f.repo_id, split.tag}); - } - } - - return result; -} diff --git a/common/download.h b/common/download.h deleted file mode 100644 index edc3e9f1a71f..000000000000 --- a/common/download.h +++ /dev/null @@ -1,107 +0,0 @@ -#pragma once - -#include -#include - -struct common_params_model; - -using common_header = std::pair; -using common_header_list = std::vector; - -struct common_download_progress { - std::string url; - size_t downloaded = 0; - size_t total = 0; - bool cached = false; -}; - -class common_download_callback { -public: - virtual ~common_download_callback() = default; - virtual void on_start(const common_download_progress & p) = 0; - virtual void on_update(const common_download_progress & p) = 0; - virtual void on_done(const common_download_progress & p, bool ok) = 0; - virtual bool is_cancelled() const { return false; } -}; - -struct common_remote_params { - common_header_list headers; - long timeout = 0; // in seconds, 0 means no timeout - long max_size = 0; // unlimited if 0 -}; - -// get remote file content, returns -std::pair> common_remote_get_content(const std::string & url, const common_remote_params & params); - -// split HF repo with tag into , for example: -// - "ggml-org/models:F16" -> <"ggml-org/models", "F16"> -// tag is optional and can be empty -std::pair common_download_split_repo_tag(const std::string & hf_repo_with_tag); - -// Result of common_list_cached_models -struct common_cached_model_info { - std::string repo; - std::string tag; - std::string to_string() const { - return repo + ":" + tag; - } -}; - -// Options for common_download_model and common_download_file_single -struct common_download_opts { - std::string bearer_token; - common_header_list headers; - bool offline = false; - common_download_callback * callback = nullptr; -}; - -// Result of common_download_model -struct common_download_model_result { - std::string model_path; - std::string mmproj_path; -}; - -// Download model from HuggingFace repo or URL -// -// input (via model struct): -// - model.hf_repo: HF repo with optional tag, see common_download_split_repo_tag -// - model.hf_file: specific file in the repo (requires hf_repo) -// - model.url: simple download (used if hf_repo is empty) -// - model.path: local file path -// -// tag matching (for HF repos without model.hf_file): -// - if tag is specified, searches for GGUF matching that quantization -// - if no tag, searches for Q4_K_M, then Q4_0, then first available GGUF -// -// split GGUF: multi-part files like "model-00001-of-00003.gguf" are automatically -// detected and all parts are downloaded -// -// caching: -// - HF repos: uses HuggingFace cache -// - URLs: uses ETag-based caching -// -// when opts.offline=true, no network requests are made -// when download_mmproj=true, searches for mmproj in same directory as model or any parent directory -// then with the closest quantization bits -// -// returns result with model_path and mmproj_path (empty on failure) -common_download_model_result common_download_model( - const common_params_model & model, - const common_download_opts & opts = {}, - bool download_mmproj = false -); - -// returns list of cached models -std::vector common_list_cached_models(); - -// download single file from url to local path -// returns status code or -1 on error -// skip_etag: if true, don't read/write .etag files (for HF cache where filename is the hash) -int common_download_file_single(const std::string & url, - const std::string & path, - const common_download_opts & opts = {}, - bool skip_etag = false); - -// resolve and download model from Docker registry -// return local path to downloaded model file -std::string common_docker_resolve_model(const std::string & docker); diff --git a/common/fit.cpp b/common/fit.cpp deleted file mode 100644 index aca3f4d40748..000000000000 --- a/common/fit.cpp +++ /dev/null @@ -1,951 +0,0 @@ -#include "fit.h" - -#include "log.h" - -#include "../src/llama-ext.h" - -#include -#include -#include -#include -#include -#include -#include - -// this enum is only used in llama_params_fit_impl but needs to be defined outside of it to fix a Windows compilation issue -// enum to identify part of a layer for distributing its tensors: -enum common_layer_fraction_t { - LAYER_FRACTION_NONE = 0, // nothing - LAYER_FRACTION_ATTN = 1, // attention - LAYER_FRACTION_UP = 2, // attention + up - LAYER_FRACTION_GATE = 3, // attention + up + gate - LAYER_FRACTION_MOE = 4, // everything but sparse MoE weights -}; - -class common_params_fit_exception : public std::runtime_error { - using std::runtime_error::runtime_error; -}; - -static std::vector common_get_device_memory_data( - const char * path_model, - const llama_model_params * mparams, - const llama_context_params * cparams, - std::vector & devs, - uint32_t & hp_ngl, - uint32_t & hp_n_ctx_train, - uint32_t & hp_n_expert, - ggml_log_level log_level) { - struct user_data_t { - struct { - ggml_log_callback callback; - void * user_data; - } original_logger; - ggml_log_level min_level; // prints below this log level go to debug log - }; - user_data_t ud; - llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data); - ud.min_level = log_level; - - llama_log_set([](ggml_log_level level, const char * text, void * user_data) { - const user_data_t * ud = (const user_data_t *) user_data; - const ggml_log_level level_eff = level >= ud->min_level ? level : GGML_LOG_LEVEL_DEBUG; - ud->original_logger.callback(level_eff, text, ud->original_logger.user_data); - }, &ud); - - llama_model_params mparams_copy = *mparams; - mparams_copy.no_alloc = true; - mparams_copy.use_mmap = false; - mparams_copy.use_mlock = false; - - llama_model * model = llama_model_load_from_file(path_model, mparams_copy); - if (model == nullptr) { - llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); - throw std::runtime_error("failed to load model"); - } - - llama_context * ctx = llama_init_from_model(model, *cparams); - if (ctx == nullptr) { - llama_model_free(model); - llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); - throw std::runtime_error("failed to create llama_context from model"); - } - - const size_t nd = llama_model_n_devices(model); - std::vector ret(nd + 1); - - llama_memory_breakdown memory_breakdown = llama_get_memory_breakdown(ctx); - - for (const auto & [buft, mb] : memory_breakdown) { - if (ggml_backend_buft_is_host(buft)) { - ret.back().mb.model += mb.model; - ret.back().mb.context += mb.context; - ret.back().mb.compute += mb.compute; - continue; - } - - ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft); - if (!dev) { - continue; - } - for (size_t i = 0; i < nd; i++) { - if (dev == llama_model_get_device(model, i)) { - ret[i].mb.model += mb.model; - ret[i].mb.context += mb.context; - ret[i].mb.compute += mb.compute; - break; - } - } - } - - { - ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); - if (cpu_dev == nullptr) { - throw std::runtime_error("no CPU backend found"); - } - size_t free; - size_t total; - ggml_backend_dev_memory(cpu_dev, &free, &total); - ret.back().free = free; - ret.back().total = total; - } - for (size_t i = 0; i < nd; i++) { - size_t free; - size_t total; - ggml_backend_dev_memory(llama_model_get_device(model, i), &free, &total); - - // devices can return 0 bytes for free and total memory if they do not - // have any to report. in this case, we will use the host memory as a fallback - // fixes: https://github.com/ggml-org/llama.cpp/issues/18577 - if (free == 0 && total == 0) { - free = ret.back().free; - total = ret.back().total; - } - ret[i].free = free; - ret[i].total = total; - } - - devs.clear(); - for (int i = 0; i < llama_model_n_devices(model); i++) { - devs.push_back(llama_model_get_device(model, i)); - } - - hp_ngl = llama_model_n_layer(model); - hp_n_ctx_train = llama_model_n_ctx_train(model); - hp_n_expert = llama_model_n_expert(model); - - common_memory_breakdown_print(ctx); - - llama_free(ctx); - llama_model_free(model); - llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); - - return ret; -} - -static void common_params_fit_impl( - const char * path_model, struct llama_model_params * mparams, struct llama_context_params * cparams, - float * tensor_split, struct llama_model_tensor_buft_override * tensor_buft_overrides, - size_t * margins_s, uint32_t n_ctx_min, enum ggml_log_level log_level) { - if (mparams->split_mode == LLAMA_SPLIT_MODE_TENSOR) { - throw common_params_fit_exception("llama_params_fit is not implemented for SPLIT_MODE_TENSOR, abort"); - } - constexpr int64_t MiB = 1024*1024; - typedef std::vector dmds_t; - const llama_model_params default_mparams = llama_model_default_params(); - - std::vector devs; - uint32_t hp_ngl = 0; // hparams.n_gpu_layers - uint32_t hp_nct = 0; // hparams.n_ctx_train - uint32_t hp_nex = 0; // hparams.n_expert - - // step 1: get data for default parameters and check whether any changes are necessary in the first place - - LOG_INF("%s: getting device memory data for initial parameters:\n", __func__); - const dmds_t dmds_full = common_get_device_memory_data(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); - const size_t nd = devs.size(); // number of devices - - std::vector margins; // this function uses int64_t rather than size_t for memory sizes to more conveniently handle deficits - margins.reserve(nd); - if (nd == 0) { - margins.push_back(margins_s[0]); - } else { - for (size_t id = 0; id < nd; id++) { - margins.push_back(margins_s[id]); - } - } - - std::vector dev_names; - { - dev_names.reserve(nd); - size_t max_length = 0; - for (const auto & dev : devs) { - std::string name = ggml_backend_dev_name(dev); - name += " ("; - name += ggml_backend_dev_description(dev); - name += ")"; - dev_names.push_back(name); - max_length = std::max(max_length, name.length()); - } - for (std::string & dn : dev_names) { - dn.insert(dn.end(), max_length - dn.length(), ' '); - } - } - - int64_t sum_free = 0; - int64_t sum_projected_free = 0; - int64_t sum_projected_used = 0; - int64_t sum_projected_model = 0; - std::vector projected_free_per_device; - projected_free_per_device.reserve(nd); - - if (nd == 0) { - sum_projected_used = dmds_full.back().mb.total(); - sum_free = dmds_full.back().total; - sum_projected_free = sum_free - sum_projected_used; - LOG_INF("%s: projected to use %" PRId64 " MiB of host memory vs. %" PRId64 " MiB of total host memory\n", - __func__, sum_projected_used/MiB, sum_free/MiB); - if (sum_projected_free >= margins[0]) { - LOG_INF("%s: will leave %" PRId64 " >= %" PRId64 " MiB of system memory, no changes needed\n", - __func__, sum_projected_free/MiB, margins[0]/MiB); - return; - } - } else { - if (nd > 1) { - LOG_INF("%s: projected memory use with initial parameters [MiB]:\n", __func__); - } - for (size_t id = 0; id < nd; id++) { - const llama_device_memory_data & dmd = dmds_full[id]; - - const int64_t projected_used = dmd.mb.total(); - const int64_t projected_free = dmd.free - projected_used; - projected_free_per_device.push_back(projected_free); - - sum_free += dmd.free; - sum_projected_used += projected_used; - sum_projected_free += projected_free; - sum_projected_model += dmd.mb.model; - - if (nd > 1) { - LOG_INF("%s: - %s: %6" PRId64 " total, %6" PRId64 " used, %6" PRId64 " free vs. target of %6" PRId64 "\n", - __func__, dev_names[id].c_str(), dmd.total/MiB, projected_used/MiB, projected_free/MiB, margins[id]/MiB); - } - } - assert(sum_free >= 0 && sum_projected_used >= 0); - LOG_INF("%s: projected to use %" PRId64 " MiB of device memory vs. %" PRId64 " MiB of free device memory\n", - __func__, sum_projected_used/MiB, sum_free/MiB); - if (nd == 1) { - if (projected_free_per_device[0] >= margins[0]) { - LOG_INF("%s: will leave %" PRId64 " >= %" PRId64 " MiB of free device memory, no changes needed\n", - __func__, projected_free_per_device[0]/MiB, margins[0]/MiB); - return; - } - } else { - bool changes_needed = false; - for (size_t id = 0; id < nd; id++) { - if (projected_free_per_device[id] < margins[id]) { - changes_needed = true; - break; - } - } - if (!changes_needed) { - LOG_INF("%s: targets for free memory can be met on all devices, no changes needed\n", __func__); - return; - } - } - } - - // step 2: try reducing memory use by reducing the context size - - { - int64_t global_surplus = sum_projected_free; - if (nd == 0) { - global_surplus -= margins[0]; - } else { - for (size_t id = 0; id < nd; id++) { - global_surplus -= margins[id]; - } - } - if (global_surplus < 0) { - if (nd <= 1) { - LOG_INF("%s: cannot meet free memory target of %" PRId64 " MiB, need to reduce device memory by %" PRId64 " MiB\n", - __func__, margins[0]/MiB, -global_surplus/MiB); - } else { - LOG_INF( - "%s: cannot meet free memory targets on all devices, need to use %" PRId64 " MiB less in total\n", - __func__, -global_surplus/MiB); - } - if (cparams->n_ctx == 0) { - if (hp_nct > n_ctx_min) { - int64_t sum_used_target = sum_free; - if (nd == 0) { - sum_used_target -= margins[0]; - } else { - for (size_t id = 0; id < nd; id++) { - sum_used_target -= margins[id]; - } - } - if (nd > 1) { - // for multiple devices we need to be more conservative in terms of how much context we think can fit: - // - for dense models only whole layers can be assigned to devices - // - for MoE models only whole tensors can be assigned to devices, which we estimate to be <= 1/3 of a layer - // - on average we expect a waste of 0.5 layers/tensors per device - // - use slightly more than the expected average for nd devices to be safe - const int64_t model_per_layer = sum_projected_model / std::min(uint32_t(mparams->n_gpu_layers), hp_ngl); - sum_used_target -= (nd + 1) * model_per_layer / (hp_nex == 0 ? 2 : 6); - } - - int64_t sum_projected_used_min_ctx = 0; - cparams->n_ctx = n_ctx_min; - const dmds_t dmds_min_ctx = common_get_device_memory_data(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); - if (nd == 0) { - sum_projected_used_min_ctx = dmds_min_ctx.back().mb.total(); - } else { - for (size_t id = 0; id < nd; id++) { - sum_projected_used_min_ctx += dmds_min_ctx[id].mb.total(); - } - } - if (sum_used_target > sum_projected_used_min_ctx) { - // linear interpolation between minimum and maximum context size: - cparams->n_ctx += (hp_nct - n_ctx_min) * (sum_used_target - sum_projected_used_min_ctx) - / (sum_projected_used - sum_projected_used_min_ctx); - cparams->n_ctx = std::max(cparams->n_ctx - cparams->n_ctx % 256, n_ctx_min); // round down context for CUDA backend - - const int64_t bytes_per_ctx = (sum_projected_used - sum_projected_used_min_ctx) / (hp_nct - n_ctx_min); - const int64_t memory_reduction = (hp_nct - cparams->n_ctx) * bytes_per_ctx; - LOG_INF("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n", - __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB); - if (nd <= 1) { - LOG_INF("%s: entire model can be fit by reducing context\n", __func__); - return; - } - LOG_INF("%s: entire model should be fit across devices by reducing context\n", __func__); - } else { - const int64_t memory_reduction = sum_projected_used - sum_projected_used_min_ctx; - LOG_INF("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n", - __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB); - } - } else { - if (n_ctx_min == UINT32_MAX) { - LOG_INF("%s: user has requested full context size of %" PRIu32 " -> no change\n", __func__, hp_nct); - } else { - LOG_INF("%s: default model context size is %" PRIu32 " which is <= the min. context size of %" PRIu32 " -> no change\n", - __func__, hp_nct, n_ctx_min); - } - } - } else { - LOG_INF("%s: context size set by user to %" PRIu32 " -> no change\n", __func__, cparams->n_ctx); - } - } - } - if (nd == 0) { - throw common_params_fit_exception("was unable to fit model into system memory by reducing context, abort"); - } - - if (mparams->n_gpu_layers != default_mparams.n_gpu_layers) { - throw common_params_fit_exception("n_gpu_layers already set by user to " + std::to_string(mparams->n_gpu_layers) + ", abort"); - } - if (nd > 1) { - if (!tensor_split) { - throw common_params_fit_exception("did not provide a buffer to write the tensor_split to, abort"); - } - if (mparams->tensor_split) { - for (size_t id = 0; id < nd; id++) { - if (mparams->tensor_split[id] != 0.0f) { - throw common_params_fit_exception("model_params::tensor_split already set by user, abort"); - } - } - } - if (mparams->split_mode == LLAMA_SPLIT_MODE_ROW) { - throw common_params_fit_exception("changing weight allocation for LLAMA_SPLIT_MODE_ROW not implemented, abort"); - } - } - if (!tensor_buft_overrides) { - throw common_params_fit_exception("did not provide buffer to set tensor_buft_overrides, abort"); - } - if (mparams->tensor_buft_overrides && (mparams->tensor_buft_overrides->pattern || mparams->tensor_buft_overrides->buft)) { - throw common_params_fit_exception("model_params::tensor_buft_overrides already set by user, abort"); - } - - // step 3: iteratively fill the back to front with "dense" layers - // - for a dense model simply fill full layers, giving each device a contiguous slice of the model - // - for a MoE model, same as dense model but with all MoE tensors in system memory - - // utility function that returns a static C string matching the tensors for a specific layer index and layer fraction: - auto get_overflow_pattern = [&](const size_t il, const common_layer_fraction_t lf) -> const char * { - constexpr size_t n_strings = 1000; - if (il >= n_strings) { - throw std::runtime_error("at most " + std::to_string(n_strings) + " model layers are supported"); - } - switch (lf) { - case LAYER_FRACTION_ATTN: { - static std::array patterns; - if (patterns[il].empty()) { - patterns[il] = "blk\\." + std::to_string(il) + "\\.ffn_(gate|up|gate_up|down).*"; - } - return patterns[il].c_str(); - } - case LAYER_FRACTION_UP: { - static std::array patterns; - if (patterns[il].empty()) { - patterns[il] = "blk\\." + std::to_string(il) + "\\.ffn_(gate|gate_up|down).*"; - } - return patterns[il].c_str(); - } - case LAYER_FRACTION_GATE: { - static std::array patterns; - if (patterns[il].empty()) { - patterns[il] = "blk\\." + std::to_string(il) + "\\.ffn_down.*"; - } - return patterns[il].c_str(); - } - case LAYER_FRACTION_MOE: { - static std::array patterns; - if (patterns[il].empty()) { - patterns[il] = "blk\\." + std::to_string(il) + "\\.ffn_(up|down|gate_up|gate)_(ch|)exps"; - } - return patterns[il].c_str(); - } - default: - GGML_ABORT("fatal error"); - } - }; - - struct ngl_t { - uint32_t n_layer = 0; // number of total layers - uint32_t n_part = 0; // number of partial layers, <= n_layer - - // for the first partial layer varying parts can overflow, all further layers use LAYER_FRACTION_MOE: - common_layer_fraction_t overflow_type = LAYER_FRACTION_MOE; - - uint32_t n_full() const { - assert(n_layer >= n_part); - return n_layer - n_part; - } - }; - - const size_t ntbo = llama_max_tensor_buft_overrides(); - - // utility function to set n_gpu_layers and tensor_split - auto set_ngl_tensor_split_tbo = [&]( - const std::vector & ngl_per_device, - const std::vector & overflow_bufts, - llama_model_params & mparams) { - mparams.n_gpu_layers = 0; - for (size_t id = 0; id < nd; id++) { - mparams.n_gpu_layers += ngl_per_device[id].n_layer; - if (nd > 1) { - tensor_split[id] = ngl_per_device[id].n_layer; - } - } - assert(uint32_t(mparams.n_gpu_layers) <= hp_ngl + 1); - uint32_t il0 = hp_ngl + 1 - mparams.n_gpu_layers; // start index for tensor buft overrides - - mparams.tensor_split = tensor_split; - - size_t itbo = 0; - for (size_t id = 0; id < nd; id++) { - il0 += ngl_per_device[id].n_full(); - for (uint32_t il = il0; il < il0 + ngl_per_device[id].n_part; il++) { - if (itbo + 1 >= ntbo) { - tensor_buft_overrides[itbo].pattern = nullptr; - tensor_buft_overrides[itbo].buft = nullptr; - itbo++; - mparams.tensor_buft_overrides = tensor_buft_overrides; - throw common_params_fit_exception("llama_max_tensor_buft_overrides() == " - + std::to_string(ntbo) + " is insufficient for model"); - } - tensor_buft_overrides[itbo].pattern = get_overflow_pattern(il, il == il0 ? ngl_per_device[id].overflow_type : LAYER_FRACTION_MOE); - tensor_buft_overrides[itbo].buft = il == il0 ? overflow_bufts[id] : ggml_backend_cpu_buffer_type(); - itbo++; - } - il0 += ngl_per_device[id].n_part; - } - tensor_buft_overrides[itbo].pattern = nullptr; - tensor_buft_overrides[itbo].buft = nullptr; - itbo++; - mparams.tensor_buft_overrides = tensor_buft_overrides; - }; - - // utility function that returns the memory use per device for given numbers of layers per device - auto get_memory_for_layers = [&]( - const char * func_name, - const std::vector & ngl_per_device, - const std::vector & overflow_bufts) -> std::vector { - llama_model_params mparams_copy = *mparams; - set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, mparams_copy); - - const dmds_t dmd_nl = common_get_device_memory_data( - path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); - - LOG_INF("%s: memory for test allocation by device:\n", func_name); - for (size_t id = 0; id < nd; id++) { - const ngl_t & n = ngl_per_device[id]; - LOG_INF( - "%s: id=%zu, n_layer=%2" PRIu32 ", n_part=%2" PRIu32 ", overflow_type=%d, mem=%6" PRId64 " MiB\n", - func_name, id, n.n_layer, n.n_part, int(n.overflow_type), dmd_nl[id].mb.total()/MiB); - } - - std::vector ret; - ret.reserve(nd); - for (size_t id = 0; id < nd; id++) { - ret.push_back(dmd_nl[id].mb.total()); - } - return ret; - }; - - int64_t global_surplus_cpu_moe = 0; - if (hp_nex > 0) { - const static std::string pattern_moe_all = "blk\\.\\d+\\.ffn_(up|down|gate_up|gate)_(ch|)exps"; // matches all MoE tensors - ggml_backend_buffer_type_t cpu_buft = ggml_backend_cpu_buffer_type(); - tensor_buft_overrides[0] = {pattern_moe_all.c_str(), cpu_buft}; - tensor_buft_overrides[1] = {nullptr, nullptr}; - mparams->tensor_buft_overrides = tensor_buft_overrides; - - LOG_INF("%s: getting device memory data with all MoE tensors moved to system memory:\n", __func__); - const dmds_t dmds_cpu_moe = common_get_device_memory_data( - path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level); - - for (size_t id = 0; id < nd; id++) { - global_surplus_cpu_moe += dmds_cpu_moe[id].free; - global_surplus_cpu_moe -= int64_t(dmds_cpu_moe[id].mb.total()) + margins[id]; - } - - if (global_surplus_cpu_moe > 0) { - LOG_INF("%s: with only dense weights in device memory there is a total surplus of %" PRId64 " MiB\n", - __func__, global_surplus_cpu_moe/MiB); - } else { - LOG_INF("%s: with only dense weights in device memory there is still a total deficit of %" PRId64 " MiB\n", - __func__, -global_surplus_cpu_moe/MiB); - } - - // reset - tensor_buft_overrides[0] = {nullptr, nullptr}; - mparams->tensor_buft_overrides = tensor_buft_overrides; - } - - std::vector targets; // maximum acceptable memory use per device - targets.reserve(nd); - for (size_t id = 0; id < nd; id++) { - targets.push_back(dmds_full[id].free - margins[id]); - LOG_INF("%s: id=%zu, target=%" PRId64 " MiB\n", __func__, id, targets[id]/MiB); - } - - std::vector overflow_bufts; // which bufts the first partial layer of a device overflows to: - overflow_bufts.reserve(nd); - for (size_t id = 0; id < nd; id++) { - overflow_bufts.push_back(ggml_backend_cpu_buffer_type()); - } - - std::vector ngl_per_device(nd); - std::vector mem = get_memory_for_layers(__func__, ngl_per_device, overflow_bufts); - - // optimize the number of layers per device using the method of false position: - // - ngl_per_device has 0 layers for each device, lower bound - // - try a "high" configuration where a device is given all unassigned layers - // - interpolate the memory use / layer between low and high linearly to get a guess where it meets our target - // - check memory use of our guess, replace either the low or high bound - // - once we only have a difference of a single layer, stop and return the lower bound that just barely still fits - // - the last device has the output layer, which cannot be a partial layer - if (hp_nex == 0) { - LOG_INF("%s: filling dense layers back-to-front:\n", __func__); - } else { - LOG_INF("%s: filling dense-only layers back-to-front:\n", __func__); - } - for (int id = nd - 1; id >= 0; id--) { - uint32_t n_unassigned = hp_ngl + 1; - for (size_t jd = id + 1; jd < nd; ++jd) { - assert(n_unassigned >= ngl_per_device[jd].n_layer); - n_unassigned -= ngl_per_device[jd].n_layer; - } - - std::vector ngl_per_device_high = ngl_per_device; - ngl_per_device_high[id].n_layer = n_unassigned; - if (hp_nex > 0) { - ngl_per_device_high[id].n_part = size_t(id) < nd - 1 ? ngl_per_device_high[id].n_layer : ngl_per_device_high[id].n_layer - 1; - } - if (ngl_per_device_high[id].n_layer > 0) { - std::vector mem_high = get_memory_for_layers(__func__, ngl_per_device_high, overflow_bufts); - if (mem_high[id] > targets[id]) { - assert(ngl_per_device_high[id].n_layer > ngl_per_device[id].n_layer); - uint32_t delta = ngl_per_device_high[id].n_layer - ngl_per_device[id].n_layer; - LOG_INF("%s: start filling device %" PRIu32 ", delta=%" PRIu32 "\n", __func__, id, delta); - while (delta > 1) { - uint32_t step_size = int64_t(delta) * (targets[id] - mem[id]) / (mem_high[id] - mem[id]); - step_size = std::max(step_size, uint32_t(1)); - step_size = std::min(step_size, delta - 1); - - std::vector ngl_per_device_test = ngl_per_device; - ngl_per_device_test[id].n_layer += step_size; - if (hp_nex) { - ngl_per_device_test[id].n_part += size_t(id) == nd - 1 && ngl_per_device_test[id].n_part == 0 ? - step_size - 1 : step_size; // the first layer is the output layer which must always be full - } - const std::vector mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts); - - if (mem_test[id] <= targets[id]) { - ngl_per_device = ngl_per_device_test; - mem = mem_test; - LOG_INF("%s: set ngl_per_device[%d].n_layer=%" PRIu32 "\n", __func__, id, ngl_per_device[id].n_layer); - } else { - ngl_per_device_high = ngl_per_device_test; - mem_high = mem_test; - LOG_INF("%s: set ngl_per_device_high[%d].n_layer=%" PRIu32 "\n", __func__, id, ngl_per_device_high[id].n_layer); - } - delta = ngl_per_device_high[id].n_layer - ngl_per_device[id].n_layer; - } - } else { - assert(ngl_per_device_high[id].n_layer == n_unassigned); - ngl_per_device = ngl_per_device_high; - mem = mem_high; - LOG_INF("%s: set ngl_per_device[%d].n_layer=%" PRIu32 "\n", __func__, id, ngl_per_device[id].n_layer); - } - } - - const int64_t projected_margin = dmds_full[id].free - mem[id]; - LOG_INF( - "%s: - %s: %2" PRIu32 " layers, %6" PRId64 " MiB used, %6" PRId64 " MiB free\n", - __func__, dev_names[id].c_str(), ngl_per_device[id].n_layer, mem[id]/MiB, projected_margin/MiB); - } - if (hp_nex == 0 || global_surplus_cpu_moe <= 0) { - set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, *mparams); - return; - } - - // step 4: for a MoE model where all dense tensors fit, - // convert the dense-only layers in the back to full layers in the front until all devices are full - // essentially the same procedure as for the dense-only layers except front-to-back - // also, try fitting at least part of one more layer to reduce waste for "small" GPUs with e.g. 24 GiB VRAM - - size_t id_dense_start = nd; - for (int id = nd - 1; id >= 0; id--) { - if (ngl_per_device[id].n_layer > 0) { - id_dense_start = id; - continue; - } - break; - } - assert(id_dense_start < nd); - - LOG_INF("%s: converting dense-only layers to full layers and filling them front-to-back with overflow to next device/system memory:\n", __func__); - for (size_t id = 0; id <= id_dense_start && id_dense_start < nd; id++) { - std::vector ngl_per_device_high = ngl_per_device; - for (size_t jd = id_dense_start; jd < nd; jd++) { - const uint32_t n_layer_move = jd < nd - 1 ? ngl_per_device_high[jd].n_layer : ngl_per_device_high[jd].n_layer - 1; - ngl_per_device_high[id].n_layer += n_layer_move; - ngl_per_device_high[jd].n_layer -= n_layer_move; - ngl_per_device_high[jd].n_part = 0; - } - size_t id_dense_start_high = nd - 1; - std::vector mem_high = get_memory_for_layers(__func__, ngl_per_device_high, overflow_bufts); - - if (mem_high[id] > targets[id]) { - assert(ngl_per_device_high[id].n_full() >= ngl_per_device[id].n_full()); - uint32_t delta = ngl_per_device_high[id].n_full() - ngl_per_device[id].n_full(); - while (delta > 1) { - uint32_t step_size = int64_t(delta) * (targets[id] - mem[id]) / (mem_high[id] - mem[id]); - step_size = std::max(step_size, uint32_t(1)); - step_size = std::min(step_size, delta - 1); - - std::vector ngl_per_device_test = ngl_per_device; - size_t id_dense_start_test = id_dense_start; - uint32_t n_converted_test = 0; - for (;id_dense_start_test < nd; id_dense_start_test++) { - const uint32_t n_convert_jd = std::min(step_size - n_converted_test, ngl_per_device_test[id_dense_start_test].n_part); - ngl_per_device_test[id_dense_start_test].n_layer -= n_convert_jd; - ngl_per_device_test[id_dense_start_test].n_part -= n_convert_jd; - ngl_per_device_test[id].n_layer += n_convert_jd; - n_converted_test += n_convert_jd; - - if (ngl_per_device_test[id_dense_start_test].n_part > 0) { - break; - } - } - const std::vector mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts); - - if (mem_test[id] <= targets[id]) { - ngl_per_device = ngl_per_device_test; - mem = mem_test; - id_dense_start = id_dense_start_test; - LOG_INF("%s: set ngl_per_device[%zu].(n_layer, n_part)=(%" PRIu32 ", %" PRIu32 "), id_dense_start=%zu\n", - __func__, id, ngl_per_device[id].n_layer, ngl_per_device[id].n_part, id_dense_start); - } else { - ngl_per_device_high = ngl_per_device_test; - mem_high = mem_test; - id_dense_start_high = id_dense_start_test; - LOG_INF("%s: set ngl_per_device_high[%zu].(n_layer, n_part)=(%" PRIu32 ", %" PRIu32 "), id_dense_start_high=%zu\n", - __func__, id, ngl_per_device_high[id].n_layer, ngl_per_device_high[id].n_part, id_dense_start_high); - } - assert(ngl_per_device_high[id].n_full() >= ngl_per_device[id].n_full()); - delta = ngl_per_device_high[id].n_full() - ngl_per_device[id].n_full(); - } - } else { - ngl_per_device = ngl_per_device_high; - mem = mem_high; - id_dense_start = id_dense_start_high; - LOG_INF("%s: set ngl_per_device[%zu].(n_layer, n_part)=(%" PRIu32 ", %" PRIu32 "), id_dense_start=%zu\n", - __func__, id, ngl_per_device[id].n_layer, ngl_per_device[id].n_part, id_dense_start); - } - - // try to fit at least part of one more layer - if (ngl_per_device[id_dense_start].n_layer > (id < nd - 1 ? 0 : 1)) { - std::vector ngl_per_device_test = ngl_per_device; - size_t id_dense_start_test = id_dense_start; - ngl_per_device_test[id_dense_start_test].n_layer--; - ngl_per_device_test[id_dense_start_test].n_part--; - ngl_per_device_test[id].n_layer++; - ngl_per_device_test[id].n_part++; - if (ngl_per_device_test[id_dense_start_test].n_part == 0) { - id_dense_start_test++; - } - ngl_per_device_test[id].overflow_type = LAYER_FRACTION_UP; - std::vector overflow_bufts_test = overflow_bufts; - if (id < nd - 1) { - overflow_bufts_test[id] = ggml_backend_dev_buffer_type(devs[id + 1]); - } - LOG_INF("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_UP\n", __func__); - std::vector mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts_test); - if (mem_test[id] < targets[id] && (id + 1 == nd || mem_test[id + 1] < targets[id + 1])) { - ngl_per_device = ngl_per_device_test; - overflow_bufts = overflow_bufts_test; - mem = mem_test; - id_dense_start = id_dense_start_test; - LOG_INF("%s: set ngl_per_device[%zu].(n_layer, n_part, overflow_type)=(%" PRIu32 ", %" PRIu32 ", UP), id_dense_start=%zu\n", - __func__, id, ngl_per_device[id].n_layer, ngl_per_device[id].n_part, id_dense_start); - - ngl_per_device_test[id].overflow_type = LAYER_FRACTION_GATE; - LOG_INF("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_GATE\n", __func__); - mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts_test); - if (mem_test[id] < targets[id] && (id + 1 == nd || mem_test[id + 1] < targets[id + 1])) { - ngl_per_device = ngl_per_device_test; - overflow_bufts = overflow_bufts_test; - mem = mem_test; - id_dense_start = id_dense_start_test; - LOG_INF("%s: set ngl_per_device[%zu].(n_layer, n_part, overflow_type)=(%" PRIu32 ", %" PRIu32 ", GATE), id_dense_start=%zu\n", - __func__, id, ngl_per_device[id].n_layer, ngl_per_device[id].n_part, id_dense_start); - } - } else { - ngl_per_device_test[id].overflow_type = LAYER_FRACTION_ATTN; - LOG_INF("%s: trying to fit one extra layer with overflow_type=LAYER_FRACTION_ATTN\n", __func__); - mem_test = get_memory_for_layers(__func__, ngl_per_device_test, overflow_bufts_test); - if (mem_test[id] < targets[id] && (id + 1 == nd || mem_test[id + 1] < targets[id + 1])) { - ngl_per_device = ngl_per_device_test; - overflow_bufts = overflow_bufts_test; - mem = mem_test; - id_dense_start = id_dense_start_test; - LOG_INF("%s: set ngl_per_device[%zu].(n_layer, n_part, overflow_type)=(%" PRIu32 ", %" PRIu32 ", ATTN), id_dense_start=%zu\n", - __func__, id, ngl_per_device[id].n_layer, ngl_per_device[id].n_part, id_dense_start); - } - } - } - - const int64_t projected_margin = dmds_full[id].free - mem[id]; - LOG_INF( - "%s: - %s: %2" PRIu32 " layers (%2" PRIu32 " overflowing), %6" PRId64 " MiB used, %6" PRId64 " MiB free\n", - __func__, dev_names[id].c_str(), ngl_per_device[id].n_layer, ngl_per_device[id].n_part, mem[id]/MiB, projected_margin/MiB); - } - - // print info for devices that were not changed during the conversion from dense only to full layers: - for (size_t id = id_dense_start + 1; id < nd; id++) { - const int64_t projected_margin = dmds_full[id].free - mem[id]; - LOG_INF( - "%s: - %s: %2" PRIu32 " layers (%2" PRIu32 " overflowing), %6" PRId64 " MiB used, %6" PRId64 " MiB free\n", - __func__, dev_names[id].c_str(), ngl_per_device[id].n_layer, ngl_per_device[id].n_part, mem[id]/MiB, projected_margin/MiB); - } - - set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, *mparams); -} - -enum common_params_fit_status common_fit_params( - const char * path_model, - llama_model_params * mparams, - llama_context_params * cparams, - float * tensor_split, - llama_model_tensor_buft_override * tensor_buft_overrides, - size_t * margins, - uint32_t n_ctx_min, - ggml_log_level log_level) { - const int64_t t0_us = llama_time_us(); - common_params_fit_status status = COMMON_PARAMS_FIT_STATUS_SUCCESS; - try { - common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, log_level); - LOG_INF("%s: successfully fit params to free device memory\n", __func__); - } catch (const common_params_fit_exception & e) { - LOG_WRN("%s: failed to fit params to free device memory: %s\n", __func__, e.what()); - status = COMMON_PARAMS_FIT_STATUS_FAILURE; - } catch (const std::runtime_error & e) { - LOG_ERR("%s: encountered an error while trying to fit params to free device memory: %s\n", __func__, e.what()); - status = COMMON_PARAMS_FIT_STATUS_ERROR; - } - const int64_t t1_us = llama_time_us(); - LOG_INF("%s: fitting params to free memory took %.2f seconds\n", __func__, (t1_us - t0_us) * 1e-6); - return status; -} - -void common_memory_breakdown_print(const struct llama_context * ctx) { - //const auto & devices = ctx->get_model().devices; - const auto * model = llama_get_model(ctx); - - std::vector devices; - for (int i = 0; i < llama_model_n_devices(model); i++) { - devices.push_back(llama_model_get_device(model, i)); - } - - llama_memory_breakdown memory_breakdown = llama_get_memory_breakdown(ctx); - - std::vector> table_data; - table_data.reserve(devices.size()); - const std::string template_header = "%s: | %s | %s %s %s %s %s %s %s |\n"; - const std::string template_gpu = "%s: | %s | %s = %s + (%s = %s + %s + %s) + %s |\n"; - const std::string template_other = "%s: | %s | %s %s %s = %s + %s + %s %s |\n"; - - table_data.push_back({template_header, "memory breakdown [MiB]", "total", "free", "self", "model", "context", "compute", "unaccounted"}); - - constexpr size_t MiB = 1024 * 1024; - const std::vector desc_prefixes_strip = {"NVIDIA ", "GeForce ", "Tesla ", "AMD ", "Radeon ", "Instinct "}; - - // track seen buffer types to avoid double counting: - std::set seen_buffer_types; - - // accumulative memory breakdown for each device and for host: - std::vector mb_dev(devices.size()); - llama_memory_breakdown_data mb_host; - - for (const auto & buft_mb : memory_breakdown) { - ggml_backend_buffer_type_t buft = buft_mb.first; - const llama_memory_breakdown_data & mb = buft_mb.second; - if (ggml_backend_buft_is_host(buft)) { - mb_host.model += mb.model; - mb_host.context += mb.context; - mb_host.compute += mb.compute; - seen_buffer_types.insert(buft); - continue; - } - ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft); - if (dev) { - int i_dev = -1; - for (size_t i = 0; i < devices.size(); i++) { - if (devices[i] == dev) { - i_dev = i; - break; - } - } - if (i_dev != -1) { - mb_dev[i_dev].model += mb.model; - mb_dev[i_dev].context += mb.context; - mb_dev[i_dev].compute += mb.compute; - seen_buffer_types.insert(buft); - continue; - } - } - } - - // print memory breakdown for each device: - for (size_t i = 0; i < devices.size(); i++) { - ggml_backend_dev_t dev = devices[i]; - llama_memory_breakdown_data mb = mb_dev[i]; - - const std::string name = ggml_backend_dev_name(dev); - std::string desc = ggml_backend_dev_description(dev); - for (const std::string & prefix : desc_prefixes_strip) { - if (desc.length() >= prefix.length() && desc.substr(0, prefix.length()) == prefix) { - desc = desc.substr(prefix.length()); - } - } - - size_t free, total; - ggml_backend_dev_memory(dev, &free, &total); - - const size_t self = mb.model + mb.context + mb.compute; - const int64_t unaccounted = static_cast(total) - static_cast(free) - static_cast(self); - - table_data.push_back({ - template_gpu, - " - " + name + " (" + desc + ")", - std::to_string(total / MiB), - std::to_string(free / MiB), - std::to_string(self / MiB), - std::to_string(mb.model / MiB), - std::to_string(mb.context / MiB), - std::to_string(mb.compute / MiB), - std::to_string(unaccounted / static_cast(MiB))}); - } - - // print memory breakdown for host: - { - const size_t self = mb_host.model + mb_host.context + mb_host.compute; - table_data.push_back({ - template_other, - " - Host", - "", // total - "", // free - std::to_string(self / MiB), - std::to_string(mb_host.model / MiB), - std::to_string(mb_host.context / MiB), - std::to_string(mb_host.compute / MiB), - ""}); // unaccounted - } - - // print memory breakdown for all remaining buffer types: - for (const auto & buft_mb : memory_breakdown) { - ggml_backend_buffer_type_t buft = buft_mb.first; - const llama_memory_breakdown_data & mb = buft_mb.second; - if (seen_buffer_types.count(buft) == 1) { - continue; - } - const std::string name = ggml_backend_buft_name(buft); - const size_t self = mb.model + mb.context + mb.compute; - table_data.push_back({ - template_other, - " - " + name, - "", // total - "", // free - std::to_string(self / MiB), - std::to_string(mb.model / MiB), - std::to_string(mb.context / MiB), - std::to_string(mb.compute / MiB), - ""}); // unaccounted - seen_buffer_types.insert(buft); - } - - for (size_t j = 1; j < table_data[0].size(); j++) { - size_t max_len = 0; - for (const auto & td : table_data) { - max_len = std::max(max_len, td[j].length()); - } - for (auto & td : table_data) { - td[j].insert(j == 1 ? td[j].length() : 0, max_len - td[j].length(), ' '); - } - } - for (const auto & td : table_data) { - LOG_INF(td[0].c_str(), - __func__, td[1].c_str(), td[2].c_str(), td[3].c_str(), td[4].c_str(), td[5].c_str(), - td[6].c_str(), td[7].c_str(), td[8].c_str()); - } -} - -void common_fit_print( - const char * path_model, - llama_model_params * mparams, - llama_context_params * cparams) { - std::vector devs; - uint32_t hp_ngl = 0; // hparams.n_gpu_layers - uint32_t hp_nct = 0; // hparams.n_ctx_train - uint32_t hp_nex = 0; // hparams.n_expert - - auto dmd = common_get_device_memory_data(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, GGML_LOG_LEVEL_ERROR); - GGML_ASSERT(dmd.size() == devs.size() + 1); - - for (size_t id = 0; id < devs.size(); id++) { - printf("%s ", ggml_backend_dev_name(devs[id])); - printf("%zu ", dmd[id].mb.model/1024/1024); - printf("%zu ", dmd[id].mb.context/1024/1024); - printf("%zu ", dmd[id].mb.compute/1024/1024); - printf("\n"); - } - - printf("Host "); - printf("%zu ", dmd.back().mb.model/1024/1024); - printf("%zu ", dmd.back().mb.context/1024/1024); - printf("%zu ", dmd.back().mb.compute/1024/1024); - printf("\n"); -} diff --git a/common/fit.h b/common/fit.h deleted file mode 100644 index e066092ec6c9..000000000000 --- a/common/fit.h +++ /dev/null @@ -1,32 +0,0 @@ -#pragma once - -#include "ggml.h" - -enum common_params_fit_status { - COMMON_PARAMS_FIT_STATUS_SUCCESS = 0, // found allocations that are projected to fit - COMMON_PARAMS_FIT_STATUS_FAILURE = 1, // could not find allocations that are projected to fit - COMMON_PARAMS_FIT_STATUS_ERROR = 2, // a hard error occurred, e.g. because no model could be found at the specified path -}; - -// fits mparams and cparams to free device memory (assumes system memory is unlimited) -// - returns true if the parameters could be successfully modified to fit device memory -// - this function is NOT thread safe because it modifies the global llama logger state -// - only parameters that have the same value as in llama_default_model_params are modified -// with the exception of the context size which is modified if and only if equal to 0 -enum common_params_fit_status common_fit_params( - const char * path_model, - struct llama_model_params * mparams, - struct llama_context_params * cparams, - float * tensor_split, // writable buffer for tensor split, needs at least llama_max_devices elements - struct llama_model_tensor_buft_override * tensor_buft_overrides, // writable buffer for overrides, needs at least llama_max_tensor_buft_overrides elements - size_t * margins, // margins of memory to leave per device in bytes - uint32_t n_ctx_min, // minimum context size to set when trying to reduce memory use - enum ggml_log_level log_level); // minimum log level to print during fitting, lower levels go to debug log - -// print estimated memory to stdout -void common_fit_print( - const char * path_model, - struct llama_model_params * mparams, - struct llama_context_params * cparams); - -void common_memory_breakdown_print(const struct llama_context * ctx); diff --git a/common/hf-cache.cpp b/common/hf-cache.cpp deleted file mode 100644 index 20f33e4c7f40..000000000000 --- a/common/hf-cache.cpp +++ /dev/null @@ -1,772 +0,0 @@ -#include "hf-cache.h" - -#include "build-info.h" -#include "common.h" -#include "log.h" -#include "http.h" - -#define JSON_ASSERT GGML_ASSERT -#include - -#include -#include -#include -#include // migration only -#include -#include -#include - -namespace nl = nlohmann; - -#if defined(_WIN32) -#define WIN32_LEAN_AND_MEAN -#ifndef NOMINMAX -#define NOMINMAX -#endif -#define HOME_DIR "USERPROFILE" -#include -#else -#define HOME_DIR "HOME" -#include -#include -#endif - -namespace hf_cache { - -namespace fs = std::filesystem; - -static fs::path get_cache_directory() { - static const fs::path cache = []() { - struct { - const char * var; - fs::path path; - } entries[] = { - {"LLAMA_CACHE", fs::path()}, - {"HF_HUB_CACHE", fs::path()}, - {"HUGGINGFACE_HUB_CACHE", fs::path()}, - {"HF_HOME", fs::path("hub")}, - {"XDG_CACHE_HOME", fs::path("huggingface") / "hub"}, - {HOME_DIR, fs::path(".cache") / "huggingface" / "hub"} - }; - for (const auto & entry : entries) { - if (auto * p = std::getenv(entry.var); p && *p) { - fs::path base(p); - return entry.path.empty() ? base : base / entry.path; - } - } -#ifndef _WIN32 - const struct passwd * pw = getpwuid(getuid()); - - if (pw && pw->pw_dir && *pw->pw_dir) { - return fs::path(pw->pw_dir) / ".cache" / "huggingface" / "hub"; - } -#endif - throw std::runtime_error("Failed to determine HF cache directory"); - }(); - - return cache; -} - -static std::string folder_name_to_repo(const std::string & folder) { - constexpr std::string_view prefix = "models--"; - if (folder.rfind(prefix, 0)) { - return {}; - } - std::string result = folder.substr(prefix.length()); - string_replace_all(result, "--", "/"); - return result; -} - -static std::string repo_to_folder_name(const std::string & repo_id) { - constexpr std::string_view prefix = "models--"; - std::string result = std::string(prefix) + repo_id; - string_replace_all(result, "/", "--"); - return result; -} - -static fs::path get_repo_path(const std::string & repo_id) { - return get_cache_directory() / repo_to_folder_name(repo_id); -} - -static bool is_hex_char(const char c) { - return (c >= 'A' && c <= 'F') || - (c >= 'a' && c <= 'f') || - (c >= '0' && c <= '9'); -} - -static bool is_hex_string(const std::string & s, size_t expected_len) { - if (s.length() != expected_len) { - return false; - } - for (const char c : s) { - if (!is_hex_char(c)) { - return false; - } - } - return true; -} - -static bool is_alphanum(const char c) { - return (c >= 'A' && c <= 'Z') || - (c >= 'a' && c <= 'z') || - (c >= '0' && c <= '9'); -} - -static bool is_special_char(char c) { - return c == '/' || c == '.' || c == '-'; -} - -// base chars [A-Za-z0-9_] are always valid -// special chars [/.-] must be surrounded by base chars -// exactly one '/' required -static bool is_valid_repo_id(const std::string & repo_id) { - if (repo_id.empty() || repo_id.length() > 256) { - return false; - } - int slash = 0; - bool special = true; - - for (const char c : repo_id) { - if (is_alphanum(c) || c == '_') { - special = false; - } else if (is_special_char(c)) { - if (special) { - return false; - } - slash += (c == '/'); - special = true; - } else { - return false; - } - } - return !special && slash == 1; -} - -static bool is_valid_hf_token(const std::string & token) { - if (token.length() < 37 || token.length() > 256 || - !string_starts_with(token, "hf_")) { - return false; - } - for (size_t i = 3; i < token.length(); ++i) { - if (!is_alphanum(token[i])) { - return false; - } - } - return true; -} - -static bool is_valid_commit(const std::string & hash) { - return is_hex_string(hash, 40); -} - -static bool is_valid_oid(const std::string & oid) { - return is_hex_string(oid, 40) || is_hex_string(oid, 64); -} - -static bool is_valid_subpath(const fs::path & path, const fs::path & subpath) { - if (subpath.is_absolute()) { - return false; // never do a / b with b absolute - } - auto b = fs::absolute(path).lexically_normal(); - auto t = (b / subpath).lexically_normal(); - auto [b_end, _] = std::mismatch(b.begin(), b.end(), t.begin(), t.end()); - - return b_end == b.end(); -} - -static void safe_write_file(const fs::path & path, const std::string & data) { - fs::path path_tmp = path.string() + ".tmp"; - - if (path.has_parent_path()) { - fs::create_directories(path.parent_path()); - } - - std::ofstream file(path_tmp); - file << data; - file.close(); - - std::error_code ec; - - if (!file.fail()) { - fs::rename(path_tmp, path, ec); - } - if (file.fail() || ec) { - fs::remove(path_tmp, ec); - throw std::runtime_error("failed to write file: " + path.string()); - } -} - -static nl::json api_get(const std::string & url, - const std::string & token) { - auto [cli, parts] = common_http_client(url); - - httplib::Headers headers = { - {"User-Agent", "llama-cpp/" + std::string(llama_build_info())}, - {"Accept", "application/json"} - }; - - if (is_valid_hf_token(token)) { - headers.emplace("Authorization", "Bearer " + token); - } else if (!token.empty()) { - LOG_WRN("%s: invalid token, authentication disabled\n", __func__); - } - - if (auto res = cli.Get(parts.path, headers)) { - auto body = res->body; - - if (res->status == 200) { - return nl::json::parse(res->body); - } - try { - body = nl::json::parse(res->body)["error"].get(); - } catch (...) { } - - throw std::runtime_error("GET failed (" + std::to_string(res->status) + "): " + body); - } else { - throw std::runtime_error("HTTPLIB failed: " + httplib::to_string(res.error())); - } -} - -static std::string get_repo_commit(const std::string & repo_id, - const std::string & token) { - try { - auto endpoint = common_get_model_endpoint(); - auto json = api_get(endpoint + "api/models/" + repo_id + "/refs", token); - - if (!json.is_object() || - !json.contains("branches") || !json["branches"].is_array()) { - LOG_WRN("%s: missing 'branches' for '%s'\n", __func__, repo_id.c_str()); - return {}; - } - - fs::path refs_path = get_repo_path(repo_id) / "refs"; - std::string name; - std::string commit; - - for (const auto & branch : json["branches"]) { - if (!branch.is_object() || - !branch.contains("name") || !branch["name"].is_string() || - !branch.contains("targetCommit") || !branch["targetCommit"].is_string()) { - continue; - } - std::string _name = branch["name"].get(); - std::string _commit = branch["targetCommit"].get(); - - if (!is_valid_subpath(refs_path, _name)) { - LOG_WRN("%s: skip invalid branch: %s\n", __func__, _name.c_str()); - continue; - } - if (!is_valid_commit(_commit)) { - LOG_WRN("%s: skip invalid commit: %s\n", __func__, _commit.c_str()); - continue; - } - - if (_name == "main") { - name = _name; - commit = _commit; - break; - } - - if (name.empty() || commit.empty()) { - name = _name; - commit = _commit; - } - } - - if (name.empty() || commit.empty()) { - LOG_WRN("%s: no valid branch for '%s'\n", __func__, repo_id.c_str()); - return {}; - } - - safe_write_file(refs_path / name, commit); - return commit; - - } catch (const nl::json::exception & e) { - LOG_ERR("%s: JSON error: %s\n", __func__, e.what()); - } catch (const std::exception & e) { - LOG_ERR("%s: error: %s\n", __func__, e.what()); - } - return {}; -} - -hf_files get_repo_files(const std::string & repo_id, - const std::string & token) { - if (!is_valid_repo_id(repo_id)) { - LOG_WRN("%s: invalid repository: %s\n", __func__, repo_id.c_str()); - return {}; - } - - std::string commit = get_repo_commit(repo_id, token); - if (commit.empty()) { - LOG_WRN("%s: failed to resolve commit for %s\n", __func__, repo_id.c_str()); - return {}; - } - - fs::path blobs_path = get_repo_path(repo_id) / "blobs"; - fs::path commit_path = get_repo_path(repo_id) / "snapshots" / commit; - - hf_files files; - - try { - auto endpoint = common_get_model_endpoint(); - auto json = api_get(endpoint + "api/models/" + repo_id + "/tree/" + commit + "?recursive=true", token); - - if (!json.is_array()) { - LOG_WRN("%s: response is not an array for '%s'\n", __func__, repo_id.c_str()); - return {}; - } - - for (const auto & item : json) { - if (!item.is_object() || - !item.contains("type") || !item["type"].is_string() || item["type"] != "file" || - !item.contains("path") || !item["path"].is_string()) { - continue; - } - - hf_file file; - file.repo_id = repo_id; - file.path = item["path"].get(); - - if (!is_valid_subpath(commit_path, file.path)) { - LOG_WRN("%s: skip invalid path: %s\n", __func__, file.path.c_str()); - continue; - } - - if (item.contains("lfs") && item["lfs"].is_object()) { - if (item["lfs"].contains("oid") && item["lfs"]["oid"].is_string()) { - file.oid = item["lfs"]["oid"].get(); - } - if (item["lfs"].contains("size") && item["lfs"]["size"].is_number()) { - file.size = item["lfs"]["size"].get(); - } - } else if (item.contains("oid") && item["oid"].is_string()) { - file.oid = item["oid"].get(); - } - if (file.size == 0 && item.contains("size") && item["size"].is_number()) { - file.size = item["size"].get(); - } - - if (!file.oid.empty() && !is_valid_oid(file.oid)) { - LOG_WRN("%s: skip invalid oid: %s\n", __func__, file.oid.c_str()); - continue; - } - - file.url = endpoint + repo_id + "/resolve/" + commit + "/" + file.path; - - fs::path final_path = commit_path / file.path; - file.final_path = final_path.string(); - - if (!file.oid.empty() && !fs::exists(final_path)) { - fs::path local_path = blobs_path / file.oid; - file.local_path = local_path.string(); - } else { - file.local_path = file.final_path; - } - - files.push_back(file); - } - } catch (const nl::json::exception & e) { - LOG_ERR("%s: JSON error: %s\n", __func__, e.what()); - } catch (const std::exception & e) { - LOG_ERR("%s: error: %s\n", __func__, e.what()); - } - return files; -} - -static std::string get_cached_ref(const fs::path & repo_path) { - fs::path refs_path = repo_path / "refs"; - if (!fs::is_directory(refs_path)) { - return {}; - } - std::string fallback; - - for (const auto & entry : fs::directory_iterator(refs_path)) { - if (!entry.is_regular_file()) { - continue; - } - std::ifstream f(entry.path()); - std::string commit; - if (!f || !std::getline(f, commit) || commit.empty()) { - continue; - } - if (!is_valid_commit(commit)) { - LOG_WRN("%s: skip invalid commit: %s\n", __func__, commit.c_str()); - continue; - } - if (entry.path().filename() == "main") { - return commit; - } - if (fallback.empty()) { - fallback = commit; - } - } - return fallback; -} - -hf_files get_cached_files(const std::string & repo_id) { - fs::path cache_dir = get_cache_directory(); - if (!fs::exists(cache_dir)) { - return {}; - } - - if (!repo_id.empty() && !is_valid_repo_id(repo_id)) { - LOG_WRN("%s: invalid repository: %s\n", __func__, repo_id.c_str()); - return {}; - } - - hf_files files; - - for (const auto & repo : fs::directory_iterator(cache_dir)) { - if (!repo.is_directory()) { - continue; - } - fs::path snapshots_path = repo.path() / "snapshots"; - - if (!fs::exists(snapshots_path)) { - continue; - } - std::string _repo_id = folder_name_to_repo(repo.path().filename().string()); - - if (!is_valid_repo_id(_repo_id)) { - continue; - } - if (!repo_id.empty() && _repo_id != repo_id) { - continue; - } - std::string commit = get_cached_ref(repo.path()); - fs::path commit_path = snapshots_path / commit; - - if (commit.empty() || !fs::is_directory(commit_path)) { - continue; - } - for (const auto & entry : fs::recursive_directory_iterator(commit_path)) { - if (!entry.is_regular_file() && !entry.is_symlink()) { - continue; - } - fs::path path = entry.path().lexically_relative(commit_path); - - if (!path.empty()) { - hf_file file; - file.repo_id = _repo_id; - file.path = path.generic_string(); - file.local_path = entry.path().string(); - file.final_path = file.local_path; - files.push_back(std::move(file)); - } - } - } - - return files; -} - -std::string finalize_file(const hf_file & file) { - static std::atomic symlinks_disabled{false}; - - std::error_code ec; - fs::path local_path(file.local_path); - fs::path final_path(file.final_path); - - if (local_path == final_path || fs::exists(final_path, ec)) { - return file.final_path; - } - - if (!fs::exists(local_path, ec)) { - return file.final_path; - } - - fs::create_directories(final_path.parent_path(), ec); - - if (!symlinks_disabled) { - fs::path target = fs::relative(local_path, final_path.parent_path(), ec); - if (!ec) { - fs::create_symlink(target, final_path, ec); - } - if (!ec) { - return file.final_path; - } - } - - if (!symlinks_disabled.exchange(true)) { - LOG_WRN("%s: failed to create symlink: %s\n", __func__, ec.message().c_str()); - LOG_WRN("%s: switching to degraded mode\n", __func__); - } - - fs::rename(local_path, final_path, ec); - if (ec) { - LOG_WRN("%s: failed to move file to snapshots: %s\n", __func__, ec.message().c_str()); - fs::copy(local_path, final_path, ec); - if (ec) { - LOG_ERR("%s: failed to copy file to snapshots: %s\n", __func__, ec.message().c_str()); - } - } - return file.final_path; -} - -// delete everything after this line, one day - -// copied from download.cpp without the tag part -struct gguf_split_info { - std::string prefix; // tag included - int index; - int count; -}; - -static gguf_split_info get_gguf_split_info(const std::string & path) { - static const std::regex re_split("^(.+)-([0-9]{5})-of-([0-9]{5})$", std::regex::icase); - std::smatch m; - - std::string prefix = path; - if (!string_remove_suffix(prefix, ".gguf")) { - return {}; - } - - int index = 1; - int count = 1; - - if (std::regex_match(prefix, m, re_split)) { - index = std::stoi(m[2].str()); - count = std::stoi(m[3].str()); - prefix = m[1].str(); - } - - return {std::move(prefix), index, count}; -} - -static std::pair parse_manifest_name(std::string & filename) { - static const std::regex re(R"(^manifest=([^=]+)=([^=]+)=.*\.json$)"); - std::smatch match; - if (std::regex_match(filename, match, re)) { - return {match[1].str(), match[2].str()}; - } - return {}; -} - -static std::string make_old_cache_filename(const std::string & owner, - const std::string & repo, - const std::string & filename) { - auto result = owner + "_" + repo + "_" + filename; - string_replace_all(result, "/", "_"); - return result; -} - -struct migrate_file { - std::string path; - std::string sha256; - size_t size; - fs::path old_path; - fs::path etag_path; - const hf_file * file; -}; - -using migrate_files = std::vector; - -static bool collect_file(const fs::path & old_cache, - const std::string & owner, - const std::string & repo, - const std::string & path, - const std::string & sha256, - const hf_files & files, - migrate_files & to_migrate) { - - const hf_file * file = nullptr; - - for (const auto & f : files) { - if (f.path == path) { - file = &f; - break; - } - } - - std::string old_filename = make_old_cache_filename(owner, repo, path); - fs::path old_path = old_cache / old_filename; - fs::path etag_path = old_path.string() + ".etag"; - - if (!fs::exists(old_path)) { - if (file && fs::exists(file->final_path)) { - return true; - } - LOG_WRN("%s: %s not found in old cache or HF cache\n", __func__, old_filename.c_str()); - return false; - } - - if (!file) { - LOG_WRN("%s: %s not found in current repo\n", __func__, old_filename.c_str()); - return false; - } - - if (!sha256.empty() && !file->oid.empty() && sha256 != file->oid) { - LOG_WRN("%s: %s is not up to date (sha256 mismatch)\n", __func__, old_filename.c_str()); - return false; - } - - if (file->size > 0) { - size_t size = fs::file_size(old_path); - if (size != file->size) { - LOG_WRN("%s: %s has wrong size %zu (expected %zu)\n", __func__, old_filename.c_str(), size, file->size); - return false; - } - } - - to_migrate.push_back({path, sha256, file->size, old_path, etag_path, file}); - return true; -} - -static bool collect_files(const fs::path & old_cache, - const std::string & owner, - const std::string & repo, - const nl::json & node, - const hf_files & files, - migrate_files & to_migrate) { - - if (!node.contains("rfilename") || - !node.contains("lfs") || - !node["lfs"].contains("sha256")) { - return true; - } - - std::string path = node["rfilename"]; - std::string sha256 = node["lfs"]["sha256"]; - - auto split = get_gguf_split_info(path); - - if (split.count <= 1) { - return collect_file(old_cache, owner, repo, path, sha256, files, to_migrate); - } - - std::vector> splits; - - for (const auto & f : files) { - auto split_f = get_gguf_split_info(f.path); - if (split_f.count == split.count && split_f.prefix == split.prefix) { - // sadly the manifest only provides the sha256 of the first file (index == 1) - // the rest will be verified using the size... - std::string f_sha256 = (split_f.index == 1) ? sha256 : ""; - splits.emplace_back(f.path, f_sha256); - } - } - - if ((int)splits.size() != split.count) { - LOG_WRN("%s: expected %d split files but found %d in repo\n", __func__, split.count, (int)splits.size()); - return false; - } - - for (const auto & [f_path, f_sha256] : splits) { - if (!collect_file(old_cache, owner, repo, f_path, f_sha256, files, to_migrate)) { - return false; - } - } - - return true; -} - -static bool migrate_file(const migrate_file & file) { - std::error_code ec; - - fs::path new_path(file.file->local_path); - fs::create_directories(new_path.parent_path(), ec); - - if (!fs::exists(new_path, ec)) { - fs::rename(file.old_path, new_path, ec); - if (ec) { - fs::copy_file(file.old_path, new_path, ec); - if (ec) { - LOG_ERR("%s: failed to move/copy %s: %s\n", __func__, file.old_path.string().c_str(), ec.message().c_str()); - return false; - } - } - fs::remove(file.old_path, ec); - } - fs::remove(file.etag_path, ec); - - std::string filename = finalize_file(*file.file); - LOG_INF("%s: migrated %s -> %s\n", __func__, file.old_path.filename().string().c_str(), filename.c_str()); - return true; -} - -void migrate_old_cache_to_hf_cache(const std::string & token, bool offline) { - fs::path old_cache = fs_get_cache_directory(); - if (!fs::exists(old_cache)) { - return; - } - - if (offline) { - LOG_WRN("%s: skipping migration in offline mode (will run when online)\n", __func__); - return; // -hf is not going to work - } - - bool warned = false; - - for (const auto & entry : fs::directory_iterator(old_cache)) { - if (!entry.is_regular_file()) { - continue; - } - auto filename = entry.path().filename().string(); - auto [owner, repo] = parse_manifest_name(filename); - - if (owner.empty() || repo.empty()) { - continue; - } - - if (!warned) { - warned = true; - LOG_WRN("================================================================================\n" - "WARNING: Migrating cache to HuggingFace cache directory\n" - " Old cache: %s\n" - " New cache: %s\n" - "This one-time migration moves models previously downloaded with -hf\n" - "from the legacy llama.cpp cache to the standard HuggingFace cache.\n" - "Models downloaded with --model-url are not affected.\n" - "================================================================================\n", - old_cache.string().c_str(), get_cache_directory().string().c_str()); - } - - auto repo_id = owner + "/" + repo; - auto files = get_repo_files(repo_id, token); - - if (files.empty()) { - LOG_WRN("%s: could not get repo files for %s, skipping\n", __func__, repo_id.c_str()); - continue; - } - - migrate_files to_migrate; - bool ok = true; - - try { - std::ifstream manifest(entry.path()); - auto json = nl::json::parse(manifest); - for (const char * key : {"ggufFile", "mmprojFile"}) { - if (json.contains(key)) { - if (!collect_files(old_cache, owner, repo, json[key], files, to_migrate)) { - ok = false; - break; - } - } - } - } catch (const std::exception & e) { - LOG_WRN("%s: failed to parse manifest %s: %s\n", __func__, filename.c_str(), e.what()); - continue; - } - - if (!ok) { - LOG_WRN("%s: migration skipped: one or more files failed validation\n", __func__); - continue; - } - - for (const auto & file : to_migrate) { - if (!migrate_file(file)) { - ok = false; - break; - } - } - - if (!ok) { - LOG_WRN("%s: migration failed: could not migrate all files\n", __func__); - continue; - } - - LOG_INF("%s: migration complete, deleting manifest: %s\n", __func__, entry.path().string().c_str()); - fs::remove(entry.path()); - } -} - -} // namespace hf_cache diff --git a/common/hf-cache.h b/common/hf-cache.h deleted file mode 100644 index 9e46f9774371..000000000000 --- a/common/hf-cache.h +++ /dev/null @@ -1,36 +0,0 @@ -#pragma once - -#include -#include - -// Ref: https://huggingface.co/docs/hub/local-cache.md - -namespace hf_cache { - -struct hf_file { - std::string path; - std::string url; - std::string local_path; - std::string final_path; - std::string oid; - std::string repo_id; - size_t size = 0; // only for the migration -}; - -using hf_files = std::vector; - -// Get files from HF API -hf_files get_repo_files( - const std::string & repo_id, - const std::string & token -); - -hf_files get_cached_files(const std::string & repo_id = {}); - -// Create snapshot path (link or move/copy) and return it -std::string finalize_file(const hf_file & file); - -// TODO: Remove later -void migrate_old_cache_to_hf_cache(const std::string & token, bool offline = false); - -} // namespace hf_cache diff --git a/common/http.h b/common/http.h deleted file mode 100644 index d3daccd6bf48..000000000000 --- a/common/http.h +++ /dev/null @@ -1,99 +0,0 @@ -#pragma once - -#include - -struct common_http_url { - std::string scheme; - std::string user; - std::string password; - std::string host; - int port; - std::string path; -}; - -static common_http_url common_http_parse_url(const std::string & url) { - common_http_url parts; - auto scheme_end = url.find("://"); - - if (scheme_end == std::string::npos) { - throw std::runtime_error("invalid URL: no scheme"); - } - parts.scheme = url.substr(0, scheme_end); - - if (parts.scheme != "http" && parts.scheme != "https") { - throw std::runtime_error("unsupported URL scheme: " + parts.scheme); - } - - auto rest = url.substr(scheme_end + 3); - auto at_pos = rest.find('@'); - - if (at_pos != std::string::npos) { - auto auth = rest.substr(0, at_pos); - auto colon_pos = auth.find(':'); - if (colon_pos != std::string::npos) { - parts.user = auth.substr(0, colon_pos); - parts.password = auth.substr(colon_pos + 1); - } else { - parts.user = auth; - } - rest = rest.substr(at_pos + 1); - } - - auto slash_pos = rest.find('/'); - - if (slash_pos != std::string::npos) { - parts.host = rest.substr(0, slash_pos); - parts.path = rest.substr(slash_pos); - } else { - parts.host = rest; - parts.path = "/"; - } - - auto colon_pos = parts.host.find(':'); - - if (colon_pos != std::string::npos) { - parts.port = std::stoi(parts.host.substr(colon_pos + 1)); - parts.host = parts.host.substr(0, colon_pos); - } else if (parts.scheme == "http") { - parts.port = 80; - } else if (parts.scheme == "https") { - parts.port = 443; - } else { - throw std::runtime_error("unsupported URL scheme: " + parts.scheme); - } - - return parts; -} - -static std::pair common_http_client(const std::string & url) { - common_http_url parts = common_http_parse_url(url); - - if (parts.host.empty()) { - throw std::runtime_error("error: invalid URL format"); - } - -#ifndef CPPHTTPLIB_OPENSSL_SUPPORT - if (parts.scheme == "https") { - throw std::runtime_error( - "HTTPS is not supported. Please rebuild with one of:\n" - " -DLLAMA_BUILD_BORINGSSL=ON\n" - " -DLLAMA_BUILD_LIBRESSL=ON\n" - " -DLLAMA_OPENSSL=ON (default, requires OpenSSL dev files installed)" - ); - } -#endif - - httplib::Client cli(parts.scheme + "://" + parts.host + ":" + std::to_string(parts.port)); - - if (!parts.user.empty()) { - cli.set_basic_auth(parts.user, parts.password); - } - - cli.set_follow_location(true); - - return { std::move(cli), std::move(parts) }; -} - -static std::string common_http_show_masked_url(const common_http_url & parts) { - return parts.scheme + "://" + (parts.user.empty() ? "" : "****:****@") + parts.host + parts.path; -} diff --git a/common/jinja/README.md b/common/jinja/README.md index 8291240767e8..1ba0c763a58a 100644 --- a/common/jinja/README.md +++ b/common/jinja/README.md @@ -10,7 +10,7 @@ The implementation can be found in the `common/jinja` directory. - Decoupled from `nlohmann::json`: this dependency is only used for JSON-to-internal type translation and is completely optional - Minimal primitive types: int, float, bool, string, array, object, none, undefined - Detailed logging: allow source tracing on error -- Clean architecture: workarounds are applied to input data before entering the runtime (see `common/chat.cpp`) +- Clean architecture: workarounds are applied to input data before entering the runtime (see `common/common.cpp.inc`) ## Architecture @@ -26,7 +26,7 @@ The implementation can be found in the `common/jinja` directory. **For maintainers and contributors:** - See `tests/test-chat-template.cpp` for usage examples -- To add new built-ins, modify `jinja/value.cpp` and add corresponding tests in `tests/test-jinja.cpp` +- To add new built-ins, modify `common.cpp.inc` and add corresponding tests in `tests/test-jinja.cpp` ## Input Marking @@ -54,7 +54,7 @@ Since template output is a plain string, distinguishing legitimate special token ### Solution -The llama.cpp Jinja engine introduces `jinja::string` (see `jinja/string.h`), which wraps `std::string` and preserves origin metadata. +The llama.cpp Jinja engine introduces `jinja::string` (declared in `common.h.inc`), which wraps `std::string` and preserves origin metadata. **Implementation:** - Strings originating from user input are marked with `is_input = true` diff --git a/common/jinja/caps.cpp b/common/jinja/caps.cpp deleted file mode 100644 index ead864763e1d..000000000000 --- a/common/jinja/caps.cpp +++ /dev/null @@ -1,479 +0,0 @@ -#include "value.h" -#include "runtime.h" -#include "caps.h" - -// note: the json dependency is only for defining input in a convenient way -// we can remove it in the future when we figure out a better way to define inputs using jinja::value -#include - -#include -#include - -#define FILENAME "jinja-caps" - -using json = nlohmann::ordered_json; - -namespace jinja { - -using caps_json_fn = std::function; -using caps_analyze_fn = std::function; - -static void caps_try_execute(jinja::program & prog, - const caps_json_fn & messages_fn, - const caps_json_fn & tools_fn, - const caps_analyze_fn & analyze_fn) { - context ctx; - ctx.is_get_stats = true; - jinja::global_from_json(ctx, json{ - {"messages", messages_fn()}, - {"tools", tools_fn()}, - {"bos_token", ""}, - {"eos_token", ""}, - {"add_generation_prompt", true} - }, true); - - auto messages = ctx.get_val("messages"); - auto tools = ctx.get_val("tools"); - - bool success = false; - std::string result; - try { - jinja::runtime runtime(ctx); - auto results = runtime.execute(prog); - auto parts = jinja::runtime::gather_string_parts(results); - result = parts->as_string().str(); - success = true; - } catch (const std::exception & e) { - JJ_DEBUG("Exception during execution: %s", e.what()); - result = ""; - // ignore exceptions during capability analysis - } - - analyze_fn(success, messages, tools); -} - -// for debugging only -static void caps_print_stats(value & v, const std::string & path) { - std::string ops; - for (const auto & name : v->stats.ops) { - ops += name + " "; - } - JJ_DEBUG("Value %s, type: %s %s, ops: %s", - path.c_str(), - v->type().c_str(), - v->stats.used ? "(used)" : "", - ops.c_str()); -} - -std::map caps::to_map() const { - return { - {"supports_string_content", supports_string_content}, - {"supports_typed_content", supports_typed_content}, - {"supports_tools", supports_tools}, - {"supports_tool_calls", supports_tool_calls}, - {"supports_parallel_tool_calls", supports_parallel_tool_calls}, - {"supports_system_role", supports_system_role}, - {"supports_preserve_reasoning", supports_preserve_reasoning}, - {"supports_object_arguments", supports_object_arguments}, - }; -} - -std::string caps::to_string() const { - std::ostringstream ss; - ss << "Caps(\n"; - for (const auto & [key, value] : to_map()) { - ss << " " << key << "=" << (value ? "true" : "false") << "\n"; - } - ss << ")"; - return ss.str(); -} - -caps caps_get(jinja::program & prog) { - caps result; - - static const auto has_op = [](value & v, const std::string & op_name) { - return v->stats.ops.find(op_name) != v->stats.ops.end(); - }; - - JJ_DEBUG("%s\n", ">>> Running capability check: typed content"); - - // case: typed content support - caps_try_execute( - prog, - [&]() { - // messages - return json::array({ - { - {"role", "user"}, - {"content", "content"} - } - }); - }, - [&]() { - // tools - return json{nullptr}; - }, - [&](bool success, value & messages, value &) { - auto & content = messages->at(0)->at("content"); - caps_print_stats(content, "messages[0].content"); - if (has_op(content, "selectattr") || has_op(content, "array_access")) { - // accessed as an array - result.supports_typed_content = true; - } - if (!success) { - // failed to execute with content as string - result.supports_string_content = false; - } - } - ); - - JJ_DEBUG("%s\n", ">>> Running capability check: system prompt"); - - // case: system prompt support - caps_try_execute( - prog, - [&]() { - // messages - return json::array({ - { - {"role", "system"}, - {"content", "System message"} - }, - { - {"role", "user"}, - {"content", "User message"} - }, - }); - }, - [&]() { - // tools - return json::array(); - }, - [&](bool, value & messages, value &) { - auto & content = messages->at(0)->at("content"); - caps_print_stats(content, "messages[0].content"); - if (!content->stats.used) { - result.supports_system_role = false; - } - } - ); - - JJ_DEBUG("%s\n", ">>> Running capability check: single tool with object arguments support"); - - // case: tools support: single call with object arguments - caps_try_execute( - prog, - [&]() { - // messages - return json::array({ - { - {"role", "user"}, - {"content", "User message"}, - }, - { - {"role", "assistant"}, - {"content", ""}, // Some templates expect content to be empty with tool calls - {"tool_calls", json::array({ - { - {"id", "call00001"}, - {"type", "function"}, - {"function", { - {"name", "tool1"}, - {"arguments", { - {"arg", "value"} - }} - }} - } - })} - }, - { - {"role", "tool"}, - {"content", "Tool response"}, - {"tool_call_id", "call00001"} - }, - { - {"role", "assistant"}, - {"content", "The tool response was 'tool response'"} - }, - { - {"role", "user"}, - {"content", "User message"}, - }, - }); - }, - [&]() { - // tools - return json::array({ - { - {"name", "tool"}, - {"type", "function"}, - {"function", { - {"name", "tool1"}, - {"description", "Tool description"}, - {"parameters", { - {"type", "object"}, - {"properties", { - {"arg", { - {"type", "string"}, - {"description", "Arg description"}, - }}, - }}, - {"required", json::array({ "arg" })}, - }}, - }}, - }, - }); - }, - [&](bool success, value & messages, value & tools) { - if (!success) { - return; // Nothing can be inferred - } - - auto & tool_name = tools->at(0)->at("function")->at("name"); - caps_print_stats(tool_name, "tools[0].function.name"); - caps_print_stats(tools, "tools"); - if (!tool_name->stats.used) { - result.supports_tools = false; - } - - auto & tool_calls = messages->at(1)->at("tool_calls");; - caps_print_stats(tool_calls, "messages[1].tool_calls"); - if (!tool_calls->stats.used) { - result.supports_tool_calls = false; - return; - } - - auto & tool_arg = tool_calls->at(0)->at("function")->at("arguments")->at("arg"); - caps_print_stats(tool_arg, "messages[1].tool_calls[0].function.arguments.arg"); - if (tool_arg->stats.used) { - result.supports_object_arguments = true; - } - } - ); - - if (!result.supports_object_arguments) { - JJ_DEBUG("%s\n", ">>> Running capability check: single tool with string arguments support"); - - // case: tools support: single call with string arguments - caps_try_execute( - prog, - [&]() { - // messages - return json::array({ - { - {"role", "user"}, - {"content", "User message"}, - }, - { - {"role", "assistant"}, - {"content", ""}, // Some templates expect content to be empty with tool calls - {"tool_calls", json::array({ - { - {"id", "call00001"}, - {"type", "function"}, - {"function", { - {"name", "tool1"}, - {"arguments", R"({"arg": "value"})"} - }} - } - })} - }, - { - {"role", "tool"}, - {"content", "Tool response"}, - {"tool_call_id", "call00001"} - }, - { - {"role", "assistant"}, - {"content", "The tool response was 'tool response'"} - }, - { - {"role", "user"}, - {"content", "User message"}, - }, - }); - }, - [&]() { - // tools - return json::array({ - { - {"name", "tool"}, - {"type", "function"}, - {"function", { - {"name", "tool1"}, - {"description", "Tool description"}, - {"parameters", { - {"type", "object"}, - {"properties", { - {"arg", { - {"type", "string"}, - {"description", "Arg description"}, - }}, - }}, - {"required", json::array({ "arg" })}, - }}, - }}, - }, - }); - }, - [&](bool success, value & messages, value & tools) { - if (!success) { - result.supports_tool_calls = false; - result.supports_tools = false; - return; - } - - auto & tool_name = tools->at(0)->at("function")->at("name"); - caps_print_stats(tool_name, "tools[0].function.name"); - caps_print_stats(tools, "tools"); - if (!tool_name->stats.used) { - result.supports_tools = false; - } - - auto & tool_calls = messages->at(1)->at("tool_calls"); - caps_print_stats(tool_calls, "messages[1].tool_calls"); - if (!tool_calls->stats.used) { - result.supports_tool_calls = false; - return; - } - } - ); - } - - JJ_DEBUG("%s\n", ">>> Running capability check: parallel tool support"); - - // case: tools support: parallel calls - caps_try_execute( - prog, - [&]() { - json args = json(R"({"arg": "value"})"); - if (result.supports_object_arguments) { - args = json{{"arg", "value"}}; - } - - // messages - return json::array({ - { - {"role", "user"}, - {"content", "User message"}, - }, - { - {"role", "assistant"}, - {"content", ""}, // Some templates expect content to be empty with tool calls - {"tool_calls", json::array({ - { - {"id", "call00001"}, - {"type", "function"}, - {"function", { - {"name", "tool1"}, - {"arguments", args} - }} - }, - { - {"id", "call00002"}, - {"type", "function"}, - {"function", { - {"name", "tool1"}, - {"arguments", args} - }} - } - })} - }, - { - {"role", "tool"}, - {"content", "Tool response"}, - {"tool_call_id", "call00001"} - }, - { - {"role", "assistant"}, - {"content", "The tool response was 'tool response'"} - }, - { - {"role", "user"}, - {"content", "User message"}, - }, - }); - }, - [&]() { - // tools - return json::array({ - { - {"name", "tool"}, - {"type", "function"}, - {"function", { - {"name", "tool1"}, - {"description", "Tool description"}, - {"parameters", { - {"type", "object"}, - {"properties", { - {"arg", { - {"type", "string"}, - {"description", "Arg description"}, - }}, - }}, - {"required", json::array({ "arg" })}, - }}, - }}, - }, - }); - }, - [&](bool success, value & messages, value & /*tools*/) { - if (!success) { - result.supports_parallel_tool_calls = false; - return; - } - - auto & tool_calls = messages->at(1)->at("tool_calls"); - caps_print_stats(tool_calls, "messages[1].tool_calls"); - - // check for second tool call usage - auto & tool_call_1 = tool_calls->at(1)->at("function"); - caps_print_stats(tool_call_1, "messages[1].tool_calls[1].function"); - if (!tool_call_1->stats.used) { - result.supports_parallel_tool_calls = false; - } - } - ); - - JJ_DEBUG("%s\n", ">>> Running capability check: preserve reasoning"); - - // case: preserve reasoning content in chat history - caps_try_execute( - prog, - [&]() { - // messages - return json::array({ - { - {"role", "user"}, - {"content", "User message"} - }, - { - {"role", "assistant"}, - {"content", "Assistant message"}, - {"reasoning_content", "Reasoning content"} - }, - { - {"role", "user"}, - {"content", "User message"} - }, - }); - }, - [&]() { - // tools - return json::array(); - }, - [&](bool, value & messages, value &) { - auto & content = messages->at(1)->at("reasoning_content"); - caps_print_stats(content, "messages[1].reasoning_content"); - if (content->stats.used) { - result.supports_preserve_reasoning = true; - } - } - ); - - JJ_DEBUG("%s\n", result.to_string().c_str()); - - return result; -} - -} // namespace jinja diff --git a/common/jinja/caps.h b/common/jinja/caps.h deleted file mode 100644 index 93a7fe09260e..000000000000 --- a/common/jinja/caps.h +++ /dev/null @@ -1,32 +0,0 @@ -#pragma once - -#include "runtime.h" - -#include -#include - -namespace jinja { - -struct caps { - bool supports_tools = true; - bool supports_tool_calls = true; - bool supports_system_role = true; - bool supports_parallel_tool_calls = true; - bool supports_preserve_reasoning = false; // support assistant message with reasoning_content - - // one of the 2 content capabilities must be true - bool supports_string_content = true; - bool supports_typed_content = false; - - bool supports_object_arguments = false; - - // for reporting on server - std::map to_map() const; - - // for debugging - std::string to_string() const; -}; - -caps caps_get(jinja::program & prog); - -} // namespace jinja diff --git a/common/jinja/lexer.cpp b/common/jinja/lexer.cpp deleted file mode 100644 index 598982c2fe4a..000000000000 --- a/common/jinja/lexer.cpp +++ /dev/null @@ -1,341 +0,0 @@ -#include "lexer.h" -#include "runtime.h" - -#include -#include -#include -#include -#include - -#define FILENAME "jinja-lexer" - -namespace jinja { - -static void string_lstrip(std::string & s, const char * chars) { - size_t start = s.find_first_not_of(chars); - if (start == std::string::npos) { - s.clear(); - } else { - s.erase(0, start); - } -} - -static void string_rstrip(std::string & s, const char * chars) { - size_t end = s.find_last_not_of(chars); - if (end == std::string::npos) { - s.clear(); - } else { - s.erase(end + 1); - } -} - -lexer_result lexer::tokenize(const std::string & source) { - std::vector tokens; - - // NOTE: do NOT transform the source string (i.e. preprocessing), as we need to keep - // the original character positions for error reporting etc. - std::string src = source; - - if (source.empty()) { - return {tokens, src}; - } - - // Normalize \r\n or \r to \n - for (std::string::size_type pos = 0; (pos = src.find("\r\n", pos)) != std::string::npos; ) { - src.erase(pos, 1); - ++pos; - } - for (std::string::size_type pos = 0; (pos = src.find("\r", pos)) != std::string::npos; ) { - src.replace(pos, 1, 1, '\n'); - ++pos; - } - - // In the default configuration: - // - a single trailing newline is stripped if present - // - other whitespace (spaces, tabs, newlines etc.) is returned unchanged - if (source.back() == '\n') { - src.pop_back(); - } - - size_t pos = 0; - size_t start_pos = 0; - size_t curly_bracket_depth = 0; - - using pred = std::function; - auto consume_while = [&](const pred & predicate) -> std::string { - std::string str; - while (predicate(src[pos])) { - // check for escape char - if (src[pos] == '\\') { - // consume backslash - ++pos; - // check for end of input - if (pos >= src.size()) { - throw lexer_exception("unexpected end of input after escape character", source, pos); - } - // add escaped char - char escaped_char = src[pos++]; - if (escape_chars.find(escaped_char) == escape_chars.end()) { - throw lexer_exception(std::string("unknown escape character \\") + escaped_char, source, pos); - } - char unescaped_char = escape_chars.at(escaped_char); - str += unescaped_char; - continue; - } - - str += src[pos++]; - if (pos > src.size()) { - throw lexer_exception("unexpected end of input during consume_while", source, pos); - } - } - return str; - }; - - auto consume_numeric = [&]() -> std::string { - std::string num = consume_while(is_integer); - if (pos < src.size() && src[pos] == '.' && pos + 1 < src.size() && is_integer(src[pos + 1])) { - ++pos; // Consume '.' - std::string frac = consume_while(is_integer); - num += "." + frac; - } - return num; - }; - - auto next_pos_is = [&](std::initializer_list chars, size_t n = 1) -> bool { - if (pos + n >= src.size()) return false; - for (char c : chars) { - if (src[pos + n] == c) return true; - } - return false; - }; - - // note: default config for chat template: lstrip_blocks = true, trim_blocks = true - - // text\n[space]{block} --> text\n{block} - bool opt_lstrip_blocks = true; - - // {block}\n[space]text --> {block}[space]text - bool opt_trim_blocks = true; - - // options set dynamically based on current/last block - bool is_lstrip_block = false; // example: {%- - bool is_rstrip_block = false; // example: -%} - - while (pos < src.size()) { - start_pos = pos; - // JJ_DEBUG("lexer main loop at pos %zu: '%s...'", pos, src.substr(pos, 10).c_str()); - - // First, consume all text that is outside of a Jinja statement or expression - token::type last_token_type = tokens.empty() - ? token::close_statement // initial state - : tokens.back().t; - if (last_token_type == token::close_statement || - last_token_type == token::close_expression || - last_token_type == token::comment) { - - bool last_block_can_rm_newline = false; - is_rstrip_block = false; - if (pos > 3) { - char c0 = src[pos - 3]; - char c1 = src[pos - 2]; - char c2 = src[pos - 1]; - // strip if: -[%}#]}text - is_rstrip_block = c0 == '-' - && (c1 == '%' || c1 == '}' || c1 == '#') - && c2 == '}'; - // match behavior of hf.js: exclude {{ and }} cases, regex: ([#%-]}) - last_block_can_rm_newline = (c1 == '#' || c1 == '%' || c1 == '-') && c2 == '}'; - } - - size_t start = pos; - size_t end = start; - while (pos < src.size() && - // Keep going until we hit the next Jinja statement or expression - !( - src[pos] == '{' && - next_pos_is( {'%', '{', '#'} ) - )) { - end = ++pos; - } - - // equivalent to hf.js code: template.replace(/^[ \t]*({[#%-])/gm, "$1"); - if (opt_lstrip_blocks && src[pos] == '{' && next_pos_is({'%', '#', '-'})) { - size_t current = end; - while (current > start) { - char c = src[current - 1]; - if (current == 1) { - end = 0; // Trim from the start of the string - break; - } - if (c == '\n') { - end = current; // Trim from the start of the line - break; - } - if (!std::isspace(static_cast(c))) { - break; // Found non-whitespace before newline, keep - } - --current; - } - } - - std::string text = src.substr(start, end - start); - - // equivalent to hf.js code: template.replace(/([#%-]})\n/g, "$1"); - if (opt_trim_blocks && last_block_can_rm_newline) { - if (!text.empty() && text.front() == '\n') { - text.erase(text.begin()); - } - } - - if (is_rstrip_block) { - // example: {last_block}[space]text - // doing lstrip on text, effectively rstrip the LAST block - // JJ_DEBUG("RSTRIP block detected, current text: '%s'", text.c_str()); - string_lstrip(text, " \t\r\n"); - } - - is_lstrip_block = src[pos] == '{' && next_pos_is({'{', '%', '#'}) && next_pos_is({'-'}, 2); - if (is_lstrip_block) { - // example: text[space]{current_block} - // doing rstrip on text, effectively lstrip the CURRENT block - // JJ_DEBUG("LSTRIP block detected, current text: '%s'", text.c_str()); - string_rstrip(text, " \t\r\n"); - } - - if (!text.empty()) { - // JJ_DEBUG("consumed text: '%s'", text.c_str()); - tokens.push_back({token::text, text, start_pos}); - continue; - } - } - - // Possibly consume a comment - // TODO: handle lstrip/rstrip for comments? (not important for now) - if (src[pos] == '{' && next_pos_is( {'#'} )) { - start_pos = pos; - pos += 2; // Skip the opening {# - std::string comment; - while (!(src[pos] == '#' && next_pos_is( {'}'} ))) { - if (pos + 2 >= src.size()) { - throw lexer_exception("missing end of comment tag", source, pos); - } - comment += src[pos++]; - } - JJ_DEBUG("consumed comment: '%s'", comment.c_str()); - tokens.push_back({token::comment, comment, start_pos}); - pos += 2; // Skip the closing #} - continue; - } - - if (src[pos] == '-' && ( - last_token_type == token::open_expression || - last_token_type == token::open_statement) - ) { - JJ_DEBUG("lexer main loop at pos %zu: '%s...'", pos, src.substr(pos, 10).c_str()); - pos++; // consume '-' in {%- or {{- - if (pos >= src.size()) break; - } - - // Consume (and ignore) all whitespace inside Jinja statements or expressions - consume_while([](char c) { return std::isspace(static_cast(c)); }); - - if (pos >= src.size()) break; - - char ch = src[pos]; - - bool is_closing_block = ch == '-' && next_pos_is( {'%', '}'} ); - - // Check for unary operators - if (!is_closing_block && (ch == '-' || ch == '+')) { - start_pos = pos; - token::type last_token_type = tokens.empty() ? token::eof : tokens.back().t; - if (last_token_type == token::text || last_token_type == token::eof) { - throw lexer_exception(std::string("unexpected character: ") + ch, source, pos); - } - switch (last_token_type) { - case token::identifier: - case token::numeric_literal: - case token::string_literal: - case token::close_paren: - case token::close_square_bracket: - // Part of a binary operator - // a - 1, 1 - 1, true - 1, "apple" - 1, (1) - 1, a[1] - 1 - // Continue parsing normally - break; - default: { - // Is part of a unary operator - // (-1), [-1], (1 + -1), not -1, -apple - ++pos; // Consume the operator - - // Check for numbers following the unary operator - std::string num = consume_numeric(); - std::string value = std::string(1, ch) + num; - token::type t = num.empty() ? token::unary_operator : token::numeric_literal; - // JJ_DEBUG("consumed unary operator or numeric literal: '%s'", value.c_str()); - tokens.push_back({t, value, start_pos}); - continue; - } - } - } - - // Try to match one of the tokens in the mapping table - bool matched = false; - for (const auto & [seq, typ] : ordered_mapping_table) { - start_pos = pos; - // Inside an object literal, don't treat "}}" as expression-end - if (seq == "}}" && curly_bracket_depth > 0) { - continue; - } - if (pos + seq.size() <= src.size() && src.substr(pos, seq.size()) == seq) { - tokens.push_back({typ, seq, start_pos}); - if (typ == token::open_expression) { - curly_bracket_depth = 0; - } else if (typ == token::open_curly_bracket) { - ++curly_bracket_depth; - } else if (typ == token::close_curly_bracket) { - --curly_bracket_depth; - } - - pos += seq.size(); - matched = true; - break; // continue main loop - } - } - if (matched) continue; // continue main loop - - // Strings - if (ch == '\'' || ch == '"') { - start_pos = pos; - ++pos; // Skip opening quote - std::string str = consume_while([ch](char c) { return c != ch; }); - // JJ_DEBUG("consumed string literal: '%s'", str.c_str()); - tokens.push_back({token::string_literal, str, start_pos}); - ++pos; // Skip closing quote - continue; - } - - // Numbers - if (is_integer(ch)) { - start_pos = pos; - std::string num = consume_numeric(); - // JJ_DEBUG("consumed numeric literal: '%s'", num.c_str()); - tokens.push_back({token::numeric_literal, num, start_pos}); - continue; - } - - // Identifiers - if (is_word(ch)) { - start_pos = pos; - std::string word = consume_while(is_word); - // JJ_DEBUG("consumed identifier: '%s'", word.c_str()); - tokens.push_back({token::identifier, word, start_pos}); - continue; - } - - throw lexer_exception(std::string("unexpected character: ") + ch, source, pos); - } - - return {std::move(tokens), src}; -} - -} // namespace jinja diff --git a/common/jinja/lexer.h b/common/jinja/lexer.h deleted file mode 100644 index 439c85764c24..000000000000 --- a/common/jinja/lexer.h +++ /dev/null @@ -1,157 +0,0 @@ -#pragma once - -#include "utils.h" - -#include -#include -#include -#include -#include - -namespace jinja { - -struct token { - enum type { - eof, // end of source - text, // The text between Jinja statements or expressions - - numeric_literal, // e.g., 123, 1.0 - string_literal, // 'string' - identifier, // Variables, functions, statements, booleans, etc. - equals, // = - open_paren, // ( - close_paren, // ) - open_statement, // {% - close_statement, // %} - open_expression, // {{ - close_expression, // }} - open_square_bracket, // [ - close_square_bracket, // ] - open_curly_bracket, // { - close_curly_bracket, // } - comma, // , - dot, // . - colon, // : - pipe, // | - - call_operator, // () - additive_binary_operator, // + - ~ - multiplicative_binary_operator, // * / % - comparison_binary_operator, // < > <= >= == != - unary_operator, // ! - + - comment, // {# ... #} - }; - type t; - std::string value; - size_t pos; -}; - -static std::string type_to_string(token::type t) { - switch (t) { - case token::eof: return "eof"; - case token::text: return "text"; - case token::numeric_literal: return "numeric_literal"; - case token::string_literal: return "string_literal"; - case token::identifier: return "identifier"; - case token::equals: return "equals"; - case token::open_paren: return "open_paren"; - case token::close_paren: return "close_paren"; - case token::open_statement: return "open_statement"; - case token::close_statement: return "close_statement"; - case token::open_expression: return "open_expression"; - case token::close_expression: return "close_expression"; - case token::open_square_bracket: return "open_square_bracket"; - case token::close_square_bracket: return "close_square_bracket"; - case token::open_curly_bracket: return "open_curly_bracket"; - case token::close_curly_bracket: return "close_curly_bracket"; - case token::comma: return "comma"; - case token::dot: return "dot"; - case token::colon: return "colon"; - case token::pipe: return "pipe"; - case token::call_operator: return "call_operator"; - case token::additive_binary_operator: return "additive_binary_operator"; - case token::multiplicative_binary_operator: return "multiplicative_binary_operator"; - case token::comparison_binary_operator: return "comparison_binary_operator"; - case token::unary_operator: return "unary_operator"; - case token::comment: return "comment"; - default: return "unknown"; - } -} - -struct lexer_result { - std::vector tokens; - std::string source; -}; - -struct lexer { - const std::map escape_chars = { - {'n', '\n'}, - {'t', '\t'}, - {'r', '\r'}, - {'b', '\b'}, - {'f', '\f'}, - {'v', '\v'}, - {'\\', '\\'}, - {'\'', '\''}, - {'\"', '\"'}, - }; - - static bool is_word(char c) { - return std::isalnum(static_cast(c)) || c == '_'; - } - - static bool is_integer(char c) { - return std::isdigit(static_cast(c)); - } - - const std::vector> ordered_mapping_table = { - // Trimmed control sequences - {"{%-", token::open_statement}, - {"-%}", token::close_statement}, - {"{{-", token::open_expression}, - {"-}}", token::close_expression}, - // Control sequences - {"{%", token::open_statement}, - {"%}", token::close_statement}, - {"{{", token::open_expression}, - {"}}", token::close_expression}, - // Single character tokens - {"(", token::open_paren}, - {")", token::close_paren}, - {"{", token::open_curly_bracket}, - {"}", token::close_curly_bracket}, - {"[", token::open_square_bracket}, - {"]", token::close_square_bracket}, - {",", token::comma}, - {".", token::dot}, - {":", token::colon}, - {"|", token::pipe}, - // Comparison operators - {"<=", token::comparison_binary_operator}, - {">=", token::comparison_binary_operator}, - {"==", token::comparison_binary_operator}, - {"!=", token::comparison_binary_operator}, - {"<", token::comparison_binary_operator}, - {">", token::comparison_binary_operator}, - // Arithmetic operators - {"+", token::additive_binary_operator}, - {"-", token::additive_binary_operator}, - {"~", token::additive_binary_operator}, - {"*", token::multiplicative_binary_operator}, - {"/", token::multiplicative_binary_operator}, - {"%", token::multiplicative_binary_operator}, - // Assignment operator - {"=", token::equals}, - }; - - // tokenize the source string into a list of tokens - // may throw lexer_exception on error - lexer_result tokenize(const std::string & source); -}; - -struct lexer_exception : public std::runtime_error { - lexer_exception(const std::string & msg, const std::string & source, size_t pos) - : std::runtime_error(fmt_error_with_source("lexer", msg, source, pos)) {} -}; - -} // namespace jinja diff --git a/common/jinja/parser.cpp b/common/jinja/parser.cpp deleted file mode 100644 index 2b25654a7a0a..000000000000 --- a/common/jinja/parser.cpp +++ /dev/null @@ -1,602 +0,0 @@ -#include "lexer.h" -#include "runtime.h" -#include "parser.h" - -#include -#include -#include -#include -#include - -#define FILENAME "jinja-parser" - -namespace jinja { - -// Helper to check type without asserting (useful for logic) -template -static bool is_type(const statement_ptr & ptr) { - return dynamic_cast(ptr.get()) != nullptr; -} - -class parser { - const std::vector & tokens; - size_t current = 0; - - std::string source; // for error reporting - -public: - parser(const std::vector & t, const std::string & src) : tokens(t), source(src) {} - - program parse() { - statements body; - while (current < tokens.size()) { - body.push_back(parse_any()); - } - return program(std::move(body)); - } - - // NOTE: start_pos is the token index, used for error reporting - template - std::unique_ptr mk_stmt(size_t start_pos, Args&&... args) { - auto ptr = std::make_unique(std::forward(args)...); - assert(start_pos < tokens.size()); - ptr->pos = tokens[start_pos].pos; - return ptr; - } - -private: - const token & peek(size_t offset = 0) const { - if (current + offset >= tokens.size()) { - static const token end_token{token::eof, "", 0}; - return end_token; - } - return tokens[current + offset]; - } - - const token & next() { - if (current >= tokens.size()) { - throw parser_exception("Parser Error: Unexpected EOF", source, tokens.empty() ? 0 : tokens.back().pos); - } - return tokens[current++]; - } - - token expect(token::type type, const std::string& error) { - const auto & t = peek(); - if (t.t != type) { - throw parser_exception("Parser Error: " + error + " (Got " + t.value + ")", source, t.pos); - } - current++; - return t; - } - - void expect_identifier(const std::string & name) { - const auto & t = peek(); - if (t.t != token::identifier || t.value != name) { - throw parser_exception("Expected identifier: " + name, source, t.pos); - } - current++; - } - - bool is(token::type type) const { - return peek().t == type; - } - - bool is_identifier(const std::string & name) const { - return peek().t == token::identifier && peek().value == name; - } - - bool is_statement(const std::vector & names) const { - if (peek(0).t != token::open_statement || peek(1).t != token::identifier) { - return false; - } - std::string val = peek(1).value; - return std::find(names.begin(), names.end(), val) != names.end(); - } - - statement_ptr parse_any() { - size_t start_pos = current; - switch (peek().t) { - case token::comment: - return mk_stmt(start_pos, next().value); - case token::text: - return mk_stmt(start_pos, next().value); - case token::open_statement: - return parse_jinja_statement(); - case token::open_expression: - return parse_jinja_expression(); - default: - throw std::runtime_error("Unexpected token type"); - } - } - - statement_ptr parse_jinja_expression() { - // Consume {{ }} tokens - expect(token::open_expression, "Expected {{"); - auto result = parse_expression(); - expect(token::close_expression, "Expected }}"); - return result; - } - - statement_ptr parse_jinja_statement() { - // Consume {% token - expect(token::open_statement, "Expected {%"); - - if (peek().t != token::identifier) { - throw std::runtime_error("Unknown statement"); - } - - size_t start_pos = current; - std::string name = next().value; - - statement_ptr result; - if (name == "set") { - result = parse_set_statement(start_pos); - - } else if (name == "if") { - result = parse_if_statement(start_pos); - // expect {% endif %} - expect(token::open_statement, "Expected {%"); - expect_identifier("endif"); - expect(token::close_statement, "Expected %}"); - - } else if (name == "macro") { - result = parse_macro_statement(start_pos); - // expect {% endmacro %} - expect(token::open_statement, "Expected {%"); - expect_identifier("endmacro"); - expect(token::close_statement, "Expected %}"); - - } else if (name == "for") { - result = parse_for_statement(start_pos); - // expect {% endfor %} - expect(token::open_statement, "Expected {%"); - expect_identifier("endfor"); - expect(token::close_statement, "Expected %}"); - - } else if (name == "break") { - expect(token::close_statement, "Expected %}"); - result = mk_stmt(start_pos); - - } else if (name == "continue") { - expect(token::close_statement, "Expected %}"); - result = mk_stmt(start_pos); - - } else if (name == "call") { - statements caller_args; - // bool has_caller_args = false; - if (is(token::open_paren)) { - // Optional caller arguments, e.g. {% call(user) dump_users(...) %} - caller_args = parse_args(); - // has_caller_args = true; - } - auto callee = parse_primary_expression(); - if (!is_type(callee)) throw std::runtime_error("Expected identifier"); - - auto call_args = parse_args(); - expect(token::close_statement, "Expected %}"); - - statements body; - while (!is_statement({"endcall"})) { - body.push_back(parse_any()); - } - - expect(token::open_statement, "Expected {%"); - expect_identifier("endcall"); - expect(token::close_statement, "Expected %}"); - - auto call_expr = mk_stmt(start_pos, std::move(callee), std::move(call_args)); - result = mk_stmt(start_pos, std::move(call_expr), std::move(caller_args), std::move(body)); - - } else if (name == "filter") { - auto filter_node = parse_primary_expression(); - if (is_type(filter_node) && is(token::open_paren)) { - filter_node = parse_call_expression(std::move(filter_node)); - } - expect(token::close_statement, "Expected %}"); - - statements body; - while (!is_statement({"endfilter"})) { - body.push_back(parse_any()); - } - - expect(token::open_statement, "Expected {%"); - expect_identifier("endfilter"); - expect(token::close_statement, "Expected %}"); - result = mk_stmt(start_pos, std::move(filter_node), std::move(body)); - - } else if (name == "generation" || name == "endgeneration") { - // Ignore generation blocks (transformers-specific) - // See https://github.com/huggingface/transformers/pull/30650 for more information. - result = mk_stmt(start_pos); - ++current; - - } else { - throw std::runtime_error("Unknown statement: " + name); - } - return result; - } - - statement_ptr parse_set_statement(size_t start_pos) { - // NOTE: `set` acts as both declaration statement and assignment expression - auto left = parse_expression_sequence(); - statement_ptr value = nullptr; - statements body; - - if (is(token::equals)) { - ++current; - value = parse_expression_sequence(); - } else { - // parsing multiline set here - expect(token::close_statement, "Expected %}"); - while (!is_statement({"endset"})) { - body.push_back(parse_any()); - } - expect(token::open_statement, "Expected {%"); - expect_identifier("endset"); - } - expect(token::close_statement, "Expected %}"); - return mk_stmt(start_pos, std::move(left), std::move(value), std::move(body)); - } - - statement_ptr parse_if_statement(size_t start_pos) { - auto test = parse_expression(); - expect(token::close_statement, "Expected %}"); - - statements body; - statements alternate; - - // Keep parsing 'if' body until we reach the first {% elif %} or {% else %} or {% endif %} - while (!is_statement({"elif", "else", "endif"})) { - body.push_back(parse_any()); - } - - if (is_statement({"elif"})) { - size_t pos0 = current; - ++current; // consume {% - ++current; // consume 'elif' - alternate.push_back(parse_if_statement(pos0)); // nested If - } else if (is_statement({"else"})) { - ++current; // consume {% - ++current; // consume 'else' - expect(token::close_statement, "Expected %}"); - - // keep going until we hit {% endif %} - while (!is_statement({"endif"})) { - alternate.push_back(parse_any()); - } - } - return mk_stmt(start_pos, std::move(test), std::move(body), std::move(alternate)); - } - - statement_ptr parse_macro_statement(size_t start_pos) { - auto name = parse_primary_expression(); - auto args = parse_args(); - expect(token::close_statement, "Expected %}"); - statements body; - // Keep going until we hit {% endmacro - while (!is_statement({"endmacro"})) { - body.push_back(parse_any()); - } - return mk_stmt(start_pos, std::move(name), std::move(args), std::move(body)); - } - - statement_ptr parse_expression_sequence(bool primary = false) { - size_t start_pos = current; - statements exprs; - exprs.push_back(primary ? parse_primary_expression() : parse_expression()); - bool is_tuple = is(token::comma); - while (is(token::comma)) { - ++current; // consume comma - exprs.push_back(primary ? parse_primary_expression() : parse_expression()); - } - return is_tuple ? mk_stmt(start_pos, std::move(exprs)) : std::move(exprs[0]); - } - - statement_ptr parse_for_statement(size_t start_pos) { - // e.g., `message` in `for message in messages` - auto loop_var = parse_expression_sequence(true); // should be an identifier/tuple - if (!is_identifier("in")) throw std::runtime_error("Expected 'in'"); - ++current; // consume 'in' - - // `messages` in `for message in messages` - auto iterable = parse_expression(); - expect(token::close_statement, "Expected %}"); - - statements body; - statements alternate; - - // Keep going until we hit {% endfor or {% else - while (!is_statement({"endfor", "else"})) { - body.push_back(parse_any()); - } - - if (is_statement({"else"})) { - ++current; // consume {% - ++current; // consume 'else' - expect(token::close_statement, "Expected %}"); - while (!is_statement({"endfor"})) { - alternate.push_back(parse_any()); - } - } - return mk_stmt( - start_pos, - std::move(loop_var), std::move(iterable), - std::move(body), std::move(alternate)); - } - - statement_ptr parse_expression() { - // Choose parse function with lowest precedence - return parse_if_expression(); - } - - statement_ptr parse_if_expression() { - auto a = parse_logical_or_expression(); - if (is_identifier("if")) { - // Ternary expression - size_t start_pos = current; - ++current; // consume 'if' - auto test = parse_logical_or_expression(); - if (is_identifier("else")) { - // Ternary expression with else - size_t pos0 = current; - ++current; // consume 'else' - auto false_expr = parse_if_expression(); // recurse to support chained ternaries - return mk_stmt(pos0, std::move(test), std::move(a), std::move(false_expr)); - } else { - // Select expression on iterable - return mk_stmt(start_pos, std::move(a), std::move(test)); - } - } - return a; - } - - statement_ptr parse_logical_or_expression() { - auto left = parse_logical_and_expression(); - while (is_identifier("or")) { - size_t start_pos = current; - token op = next(); - left = mk_stmt(start_pos, op, std::move(left), parse_logical_and_expression()); - } - return left; - } - - statement_ptr parse_logical_and_expression() { - auto left = parse_logical_negation_expression(); - while (is_identifier("and")) { - size_t start_pos = current; - auto op = next(); - left = mk_stmt(start_pos, op, std::move(left), parse_logical_negation_expression()); - } - return left; - } - - statement_ptr parse_logical_negation_expression() { - // Try parse unary operators - if (is_identifier("not")) { - size_t start_pos = current; - auto op = next(); - return mk_stmt(start_pos, op, parse_logical_negation_expression()); - } - return parse_comparison_expression(); - } - - statement_ptr parse_comparison_expression() { - // NOTE: membership has same precedence as comparison - // e.g., ('a' in 'apple' == 'b' in 'banana') evaluates as ('a' in ('apple' == ('b' in 'banana'))) - auto left = parse_additive_expression(); - while (true) { - token op; - size_t start_pos = current; - if (is_identifier("not") && peek(1).t == token::identifier && peek(1).value == "in") { - op = {token::identifier, "not in", tokens[current].pos}; - ++current; // consume 'not' - ++current; // consume 'in' - } else if (is_identifier("in")) { - op = next(); - } else if (is(token::comparison_binary_operator)) { - op = next(); - } else break; - left = mk_stmt(start_pos, op, std::move(left), parse_additive_expression()); - } - return left; - } - - statement_ptr parse_additive_expression() { - auto left = parse_multiplicative_expression(); - while (is(token::additive_binary_operator)) { - size_t start_pos = current; - auto op = next(); - left = mk_stmt(start_pos, op, std::move(left), parse_multiplicative_expression()); - } - return left; - } - - statement_ptr parse_multiplicative_expression() { - auto left = parse_test_expression(); - while (is(token::multiplicative_binary_operator)) { - size_t start_pos = current; - auto op = next(); - left = mk_stmt(start_pos, op, std::move(left), parse_test_expression()); - } - return left; - } - - statement_ptr parse_test_expression() { - auto operand = parse_filter_expression(); - while (is_identifier("is")) { - size_t start_pos = current; - ++current; // consume 'is' - bool negate = false; - if (is_identifier("not")) { ++current; negate = true; } - auto test_id = parse_primary_expression(); - // FIXME: tests can also be expressed like this: if x is eq 3 - if (is(token::open_paren)) test_id = parse_call_expression(std::move(test_id)); - operand = mk_stmt(start_pos, std::move(operand), negate, std::move(test_id)); - } - return operand; - } - - statement_ptr parse_filter_expression() { - auto operand = parse_call_member_expression(); - while (is(token::pipe)) { - size_t start_pos = current; - ++current; // consume pipe - auto filter = parse_primary_expression(); - if (is(token::open_paren)) filter = parse_call_expression(std::move(filter)); - operand = mk_stmt(start_pos, std::move(operand), std::move(filter)); - } - return operand; - } - - statement_ptr parse_call_member_expression() { - // Handle member expressions recursively - auto member = parse_member_expression(parse_primary_expression()); - return is(token::open_paren) - ? parse_call_expression(std::move(member)) // foo.x() - : std::move(member); - } - - statement_ptr parse_call_expression(statement_ptr callee) { - size_t start_pos = current; - auto expr = mk_stmt(start_pos, std::move(callee), parse_args()); - auto member = parse_member_expression(std::move(expr)); // foo.x().y - return is(token::open_paren) - ? parse_call_expression(std::move(member)) // foo.x()() - : std::move(member); - } - - statements parse_args() { - // comma-separated arguments list - expect(token::open_paren, "Expected ("); - statements args; - while (!is(token::close_paren)) { - statement_ptr arg; - // unpacking: *expr - if (peek().t == token::multiplicative_binary_operator && peek().value == "*") { - size_t start_pos = current; - ++current; // consume * - arg = mk_stmt(start_pos, parse_expression()); - } else { - arg = parse_expression(); - if (is(token::equals)) { - // keyword argument - // e.g., func(x = 5, y = a or b) - size_t start_pos = current; - ++current; // consume equals - arg = mk_stmt(start_pos, std::move(arg), parse_expression()); - } - } - args.push_back(std::move(arg)); - if (is(token::comma)) { - ++current; // consume comma - } - } - expect(token::close_paren, "Expected )"); - return args; - } - - statement_ptr parse_member_expression(statement_ptr object) { - size_t start_pos = current; - while (is(token::dot) || is(token::open_square_bracket)) { - auto op = next(); - bool computed = op.t == token::open_square_bracket; - statement_ptr prop; - if (computed) { - prop = parse_member_expression_arguments(); - expect(token::close_square_bracket, "Expected ]"); - } else { - prop = parse_primary_expression(); - } - object = mk_stmt(start_pos, std::move(object), std::move(prop), computed); - } - return object; - } - - statement_ptr parse_member_expression_arguments() { - // NOTE: This also handles slice expressions colon-separated arguments list - // e.g., ['test'], [0], [:2], [1:], [1:2], [1:2:3] - statements slices; - bool is_slice = false; - size_t start_pos = current; - while (!is(token::close_square_bracket)) { - if (is(token::colon)) { - // A case where a default is used - // e.g., [:2] will be parsed as [undefined, 2] - slices.push_back(nullptr); - ++current; // consume colon - is_slice = true; - } else { - slices.push_back(parse_expression()); - if (is(token::colon)) { - ++current; // consume colon after expression, if it exists - is_slice = true; - } - } - } - if (is_slice) { - statement_ptr start = slices.size() > 0 ? std::move(slices[0]) : nullptr; - statement_ptr stop = slices.size() > 1 ? std::move(slices[1]) : nullptr; - statement_ptr step = slices.size() > 2 ? std::move(slices[2]) : nullptr; - return mk_stmt(start_pos, std::move(start), std::move(stop), std::move(step)); - } - if (slices.empty()) { - return mk_stmt(start_pos); - } - return std::move(slices[0]); - } - - statement_ptr parse_primary_expression() { - size_t start_pos = current; - auto t = next(); - switch (t.t) { - case token::numeric_literal: - if (t.value.find('.') != std::string::npos) { - return mk_stmt(start_pos, std::stod(t.value)); - } else { - return mk_stmt(start_pos, std::stoll(t.value)); - } - case token::string_literal: { - std::string val = t.value; - while (is(token::string_literal)) { - val += next().value; - } - return mk_stmt(start_pos, val); - } - case token::identifier: - return mk_stmt(start_pos, t.value); - case token::open_paren: { - auto expr = parse_expression_sequence(); - expect(token::close_paren, "Expected )"); - return expr; - } - case token::open_square_bracket: { - statements vals; - while (!is(token::close_square_bracket)) { - vals.push_back(parse_expression()); - if (is(token::comma)) ++current; - } - ++current; - return mk_stmt(start_pos, std::move(vals)); - } - case token::open_curly_bracket: { - std::vector> pairs; - while (!is(token::close_curly_bracket)) { - auto key = parse_expression(); - expect(token::colon, "Expected :"); - pairs.push_back({std::move(key), parse_expression()}); - if (is(token::comma)) ++current; - } - ++current; - return mk_stmt(start_pos, std::move(pairs)); - } - default: - throw std::runtime_error("Unexpected token: " + t.value + " of type " + std::to_string(t.t)); - } - } -}; - -program parse_from_tokens(const lexer_result & lexer_res) { - return parser(lexer_res.tokens, lexer_res.source).parse(); -} - -} // namespace jinja diff --git a/common/jinja/parser.h b/common/jinja/parser.h deleted file mode 100644 index f1cc0212c6a4..000000000000 --- a/common/jinja/parser.h +++ /dev/null @@ -1,21 +0,0 @@ -#pragma once - -#include "lexer.h" -#include "runtime.h" -#include "utils.h" - -#include -#include - -namespace jinja { - -// parse from a list of tokens into an AST (program) -// may throw parser_exception on error -program parse_from_tokens(const lexer_result & lexer_res); - -struct parser_exception : public std::runtime_error { - parser_exception(const std::string & msg, const std::string & source, size_t pos) - : std::runtime_error(fmt_error_with_source("parser", msg, source, pos)) {} -}; - -} // namespace jinja diff --git a/common/jinja/runtime.cpp b/common/jinja/runtime.cpp deleted file mode 100644 index f81d98d954a1..000000000000 --- a/common/jinja/runtime.cpp +++ /dev/null @@ -1,906 +0,0 @@ -#include "lexer.h" -#include "runtime.h" -#include "value.h" -#include "utils.h" - -#include -#include -#include -#include - -#define FILENAME "jinja-runtime" - -bool g_jinja_debug = false; - -namespace jinja { - -void enable_debug(bool enable) { - g_jinja_debug = enable; -} - -static value_string exec_statements(const statements & stmts, context & ctx) { - auto result = mk_val(); - for (const auto & stmt : stmts) { - JJ_DEBUG("Executing statement of type %s", stmt->type().c_str()); - result->push_back(stmt->execute(ctx)); - } - // convert to string parts - value_string str = mk_val(); - gather_string_parts_recursive(result, str); - return str; -} - -static std::string get_line_col(const std::string & source, size_t pos) { - size_t line = 1; - size_t col = 1; - for (size_t i = 0; i < pos && i < source.size(); i++) { - if (source[i] == '\n') { - line++; - col = 1; - } else { - col++; - } - } - return "line " + std::to_string(line) + ", column " + std::to_string(col); -} - -static void ensure_key_type_allowed(const value & val) { - if (!val->is_hashable()) { - throw std::runtime_error("Type: " + val->type() + " is not allowed as object key"); - } -} - -// execute with error handling -value statement::execute(context & ctx) { - try { - return execute_impl(ctx); - } catch (const continue_statement::signal & /* ex */) { - throw; - } catch (const break_statement::signal & /* ex */) { - throw; - } catch (const rethrown_exception & /* ex */) { - throw; - } catch (const not_implemented_exception & /* ex */) { - throw; - } catch (const std::exception & e) { - const std::string & source = *ctx.src; - if (source.empty()) { - std::ostringstream oss; - oss << "\nError executing " << type() << " at position " << pos << ": " << e.what(); - throw rethrown_exception(oss.str()); - } else { - std::ostringstream oss; - oss << "\n------------\n"; - oss << "While executing " << type() << " at " << get_line_col(source, pos) << " in source:\n"; - oss << peak_source(source, pos) << "\n"; - oss << "Error: " << e.what(); - // throw as another exception to avoid repeated formatting - throw rethrown_exception(oss.str()); - } - } -} - -value identifier::execute_impl(context & ctx) { - auto it = ctx.get_val(val); - auto builtins = global_builtins(); - if (!it->is_undefined()) { - if (ctx.is_get_stats) { - value_t::stats_t::mark_used(it); - } - JJ_DEBUG("Identifier '%s' found, type = %s", val.c_str(), it->type().c_str()); - return it; - } else if (builtins.find(val) != builtins.end()) { - JJ_DEBUG("Identifier '%s' found in builtins", val.c_str()); - return mk_val(val, builtins.at(val)); - } else { - JJ_DEBUG("Identifier '%s' not found, returning undefined", val.c_str()); - return mk_val(val); - } -} - -value object_literal::execute_impl(context & ctx) { - auto obj = mk_val(); - for (const auto & pair : val) { - value key = pair.first->execute(ctx); - value val = pair.second->execute(ctx); - JJ_DEBUG("Object literal: setting key '%s' with value type %s", key->as_string().str().c_str(), val->type().c_str()); - obj->insert(key, val); - } - return obj; -} - -value binary_expression::execute_impl(context & ctx) { - value left_val = left->execute(ctx); - - // Logical operators - if (op.value == "and") { - JJ_DEBUG("Executing logical test: %s AND %s", left->type().c_str(), right->type().c_str()); - return left_val->as_bool() ? right->execute(ctx) : std::move(left_val); - } else if (op.value == "or") { - JJ_DEBUG("Executing logical test: %s OR %s", left->type().c_str(), right->type().c_str()); - return left_val->as_bool() ? std::move(left_val) : right->execute(ctx); - } - - // Equality operators - value right_val = right->execute(ctx); - JJ_DEBUG("Executing binary expression %s '%s' %s", left_val->type().c_str(), op.value.c_str(), right_val->type().c_str()); - if (op.value == "==") { - return mk_val(*left_val == *right_val); - } else if (op.value == "!=") { - return mk_val(!(*left_val == *right_val)); - } - - auto workaround_concat_null_with_str = [&](value & res) -> bool { - bool is_left_null = left_val->is_none() || left_val->is_undefined(); - bool is_right_null = right_val->is_none() || right_val->is_undefined(); - bool is_left_str = is_val(left_val); - bool is_right_str = is_val(right_val); - if ((is_left_null && is_right_str) || (is_right_null && is_left_str)) { - JJ_DEBUG("%s", "Workaround: treating null/undefined as empty string for string concatenation"); - string left_str = is_left_null ? string() : left_val->as_string(); - string right_str = is_right_null ? string() : right_val->as_string(); - auto output = left_str.append(right_str); - res = mk_val(std::move(output)); - return true; - } - return false; - }; - - auto test_is_in = [&]() -> bool { - func_args args(ctx); - args.push_back(left_val); - args.push_back(right_val); - return global_builtins().at("test_is_in")(args)->as_bool(); - }; - - // Handle undefined and null values - if (is_val(left_val) || is_val(right_val)) { - if (is_val(right_val) && (op.value == "in" || op.value == "not in")) { - // Special case: `anything in undefined` is `false` and `anything not in undefined` is `true` - return mk_val(op.value == "not in"); - } - if (op.value == "+" || op.value == "~") { - value res = mk_val(); - if (workaround_concat_null_with_str(res)) { - return res; - } - } - throw std::runtime_error("Cannot perform operation " + op.value + " on undefined values"); - } else if (is_val(left_val) || is_val(right_val)) { - if (op.value == "+" || op.value == "~") { - value res = mk_val(); - if (workaround_concat_null_with_str(res)) { - return res; - } - } - throw std::runtime_error("Cannot perform operation on null values"); - } - - // Float operations - if ((is_val(left_val) || is_val(left_val)) && - (is_val(right_val) || is_val(right_val))) { - double a = left_val->as_float(); - double b = right_val->as_float(); - if (op.value == "+" || op.value == "-" || op.value == "*") { - double res = (op.value == "+") ? a + b : (op.value == "-") ? a - b : a * b; - JJ_DEBUG("Arithmetic operation: %f %s %f = %f", a, op.value.c_str(), b, res); - bool is_float = is_val(left_val) || is_val(right_val); - if (is_float) { - return mk_val(res); - } else { - return mk_val(static_cast(res)); - } - } else if (op.value == "/") { - JJ_DEBUG("Division operation: %f / %f", a, b); - return mk_val(a / b); - } else if (op.value == "%") { - double rem = std::fmod(a, b); - JJ_DEBUG("Modulo operation: %f %% %f = %f", a, b, rem); - bool is_float = is_val(left_val) || is_val(right_val); - if (is_float) { - return mk_val(rem); - } else { - return mk_val(static_cast(rem)); - } - } else if (op.value == "<") { - JJ_DEBUG("Comparison operation: %f < %f is %d", a, b, a < b); - return mk_val(a < b); - } else if (op.value == ">") { - JJ_DEBUG("Comparison operation: %f > %f is %d", a, b, a > b); - return mk_val(a > b); - } else if (op.value == ">=") { - JJ_DEBUG("Comparison operation: %f >= %f is %d", a, b, a >= b); - return mk_val(a >= b); - } else if (op.value == "<=") { - JJ_DEBUG("Comparison operation: %f <= %f is %d", a, b, a <= b); - return mk_val(a <= b); - } - } - - // Array operations - if (is_val(left_val) && is_val(right_val)) { - if (op.value == "+") { - auto & left_arr = left_val->as_array(); - auto & right_arr = right_val->as_array(); - auto result = mk_val(); - for (const auto & item : left_arr) { - result->push_back(item); - } - for (const auto & item : right_arr) { - result->push_back(item); - } - return result; - } - } else if (is_val(right_val)) { - // case: 1 in [0, 1, 2] - bool member = test_is_in(); - if (op.value == "in") { - return mk_val(member); - } else if (op.value == "not in") { - return mk_val(!member); - } - } - - // String concatenation with ~ and + - if ((is_val(left_val) || is_val(right_val)) && - (op.value == "~" || op.value == "+")) { - JJ_DEBUG("String concatenation with %s operator", op.value.c_str()); - auto output = left_val->as_string().append(right_val->as_string()); - auto res = mk_val(); - res->val_str = std::move(output); - return res; - } - - // Python-style string repetition - // TODO: support array/tuple repetition (e.g., [1, 2] * 3 → [1, 2, 1, 2, 1, 2]) - if (op.value == "*" && - ((is_val(left_val) && is_val(right_val)) || - (is_val(left_val) && is_val(right_val)))) { - const auto & str = is_val(left_val) ? left_val->as_string() : right_val->as_string(); - const int64_t repeat = is_val(right_val) ? right_val->as_int() : left_val->as_int(); - auto res = mk_val(); - if (repeat <= 0) { - return res; - } - for (int64_t i = 0; i < repeat; ++i) { - res->val_str = res->val_str.append(str); - } - return res; - } - - // String membership - if (is_val(left_val) && is_val(right_val)) { - // case: "a" in "abc" - bool member = test_is_in(); - if (op.value == "in") { - return mk_val(member); - } else if (op.value == "not in") { - return mk_val(!member); - } - } - - // Value key in object - if (is_val(right_val)) { - // case: key in {key: value} - bool member = test_is_in(); - if (op.value == "in") { - return mk_val(member); - } else if (op.value == "not in") { - return mk_val(!member); - } - } - - throw std::runtime_error("Unknown operator \"" + op.value + "\" between " + left_val->type() + " and " + right_val->type()); -} - -static value try_builtin_func(context & ctx, const std::string & name, value & input, bool undef_on_missing = false) { - JJ_DEBUG("Trying built-in function '%s' for type %s", name.c_str(), input->type().c_str()); - if (ctx.is_get_stats) { - value_t::stats_t::mark_used(input); - input->stats.ops.insert(name); - } - auto builtins = input->get_builtins(); - auto it = builtins.find(name); - if (it != builtins.end()) { - JJ_DEBUG("Binding built-in '%s'", name.c_str()); - return mk_val(name, it->second, input); - } - if (undef_on_missing) { - return mk_val(name); - } - throw std::runtime_error("Unknown (built-in) filter '" + name + "' for type " + input->type()); -} - -value filter_expression::execute_impl(context & ctx) { - value input = operand ? operand->execute(ctx) : val; - - JJ_DEBUG("Applying filter to %s", input->type().c_str()); - - if (is_stmt(filter)) { - auto filter_id = cast_stmt(filter)->val; - - if (filter_id == "trim") { - filter_id = "strip"; // alias - } - JJ_DEBUG("Applying filter '%s' to %s", filter_id.c_str(), input->type().c_str()); - // TODO: Refactor filters so this coercion can be done automatically - if (!input->is_undefined() && !is_val(input) && ( - filter_id == "capitalize" || - filter_id == "lower" || - filter_id == "replace" || - filter_id == "strip" || - filter_id == "title" || - filter_id == "upper" || - filter_id == "wordcount" - )) { - JJ_DEBUG("Coercing %s to String for '%s' filter", input->type().c_str(), filter_id.c_str()); - input = mk_val(input->as_string()); - } - return try_builtin_func(ctx, filter_id, input)->invoke(func_args(ctx)); - - } else if (is_stmt(filter)) { - auto call = cast_stmt(filter); - if (!is_stmt(call->callee)) { - throw std::runtime_error("Filter callee must be an identifier"); - } - auto filter_id = cast_stmt(call->callee)->val; - - if (filter_id == "trim") { - filter_id = "strip"; // alias - } - JJ_DEBUG("Applying filter '%s' with arguments to %s", filter_id.c_str(), input->type().c_str()); - func_args args(ctx); - for (const auto & arg_expr : call->args) { - args.push_back(arg_expr->execute(ctx)); - } - - return try_builtin_func(ctx, filter_id, input)->invoke(args); - - } else { - throw std::runtime_error("Invalid filter expression"); - } -} - -value filter_statement::execute_impl(context & ctx) { - // eval body as string, then apply filter - auto body_val = exec_statements(body, ctx); - value_string parts = mk_val(); - gather_string_parts_recursive(body_val, parts); - - JJ_DEBUG("FilterStatement: applying filter to body string of length %zu", parts->val_str.length()); - filter_expression filter_expr(std::move(parts), std::move(filter)); - value out = filter_expr.execute(ctx); - - // this node can be reused later, make sure filter is preserved - this->filter = std::move(filter_expr.filter); - return out; -} - -value test_expression::execute_impl(context & ctx) { - // NOTE: "value is something" translates to function call "test_is_something(value)" - const auto & builtins = global_builtins(); - - std::string test_id; - value input = operand->execute(ctx); - - func_args args(ctx); - args.push_back(input); - - if (is_stmt(test)) { - test_id = cast_stmt(test)->val; - } else if (is_stmt(test)) { - auto call = cast_stmt(test); - if (!is_stmt(call->callee)) { - throw std::runtime_error("Test callee must be an identifier"); - } - test_id = cast_stmt(call->callee)->val; - - JJ_DEBUG("Applying test '%s' with arguments to %s", test_id.c_str(), input->type().c_str()); - for (const auto & arg_expr : call->args) { - args.push_back(arg_expr->execute(ctx)); - } - - } else { - throw std::runtime_error("Invalid test expression"); - } - - auto it = builtins.find("test_is_" + test_id); - JJ_DEBUG("Test expression %s '%s' %s (using function 'test_is_%s')", operand->type().c_str(), test_id.c_str(), negate ? "(negate)" : "", test_id.c_str()); - if (it == builtins.end()) { - throw std::runtime_error("Unknown test '" + test_id + "'"); - } - - auto res = it->second(args); - - if (negate) { - return mk_val(!res->as_bool()); - } else { - return res; - } -} - -value unary_expression::execute_impl(context & ctx) { - value operand_val = argument->execute(ctx); - JJ_DEBUG("Executing unary expression with operator '%s'", op.value.c_str()); - - if (op.value == "not") { - return mk_val(!operand_val->as_bool()); - } else if (op.value == "-") { - if (is_val(operand_val)) { - return mk_val(-operand_val->as_int()); - } else if (is_val(operand_val)) { - return mk_val(-operand_val->as_float()); - } else { - throw std::runtime_error("Unary - operator requires numeric operand"); - } - } - - throw std::runtime_error("Unknown unary operator '" + op.value + "'"); -} - -value if_statement::execute_impl(context & ctx) { - value test_val = test->execute(ctx); - - auto out = mk_val(); - if (test_val->as_bool()) { - for (auto & stmt : body) { - JJ_DEBUG("IF --> Executing THEN body, current block: %s", stmt->type().c_str()); - out->push_back(stmt->execute(ctx)); - } - } else { - for (auto & stmt : alternate) { - JJ_DEBUG("IF --> Executing ELSE body, current block: %s", stmt->type().c_str()); - out->push_back(stmt->execute(ctx)); - } - } - // convert to string parts - value_string str = mk_val(); - gather_string_parts_recursive(out, str); - return str; -} - -value for_statement::execute_impl(context & ctx) { - context scope(ctx); // new scope for loop variables - - jinja::select_expression * select_expr = cast_stmt(iterable); - statement_ptr test_expr_nullptr; - - statement_ptr & iter_expr = [&]() -> statement_ptr & { - auto tmp = cast_stmt(iterable); - return tmp ? tmp->lhs : iterable; - }(); - statement_ptr & test_expr = [&]() -> statement_ptr & { - auto tmp = cast_stmt(iterable); - return tmp ? tmp->test : test_expr_nullptr; - }(); - - JJ_DEBUG("Executing for statement, iterable type: %s", iter_expr->type().c_str()); - - value iterable_val = iter_expr->execute(scope); - - // mark the variable being iterated as used for stats - if (ctx.is_get_stats) { - value_t::stats_t::mark_used(iterable_val); - iterable_val->stats.ops.insert("array_access"); - } - - if (iterable_val->is_undefined()) { - JJ_DEBUG("%s", "For loop iterable is undefined, skipping loop"); - iterable_val = mk_val(); - } - - if (!is_val(iterable_val) && !is_val(iterable_val)) { - throw std::runtime_error("Expected iterable or object type in for loop: got " + iterable_val->type()); - } - - std::vector items; - if (is_val(iterable_val)) { - JJ_DEBUG("%s", "For loop over object keys"); - auto & obj = iterable_val->as_ordered_object(); - for (auto & p : obj) { - auto tuple = mk_val(p); - items.push_back(std::move(tuple)); - } - if (ctx.is_get_stats) { - value_t::stats_t::mark_used(iterable_val); - iterable_val->stats.ops.insert("object_access"); - } - } else { - JJ_DEBUG("%s", "For loop over array items"); - auto & arr = iterable_val->as_array(); - for (const auto & item : arr) { - items.push_back(item); - } - if (ctx.is_get_stats) { - value_t::stats_t::mark_used(iterable_val); - iterable_val->stats.ops.insert("array_access"); - } - } - - std::vector> scope_update_fns; - - std::vector filtered_items; - for (size_t i = 0; i < items.size(); ++i) { - context loop_scope(scope); - - value current = items[i]; - - std::function scope_update_fn = [](context &) { /* no-op */}; - if (is_stmt(loopvar)) { - auto id = cast_stmt(loopvar)->val; - - if (is_val(iterable_val)) { - // case example: {% for key in dict %} - current = items[i]->as_array()[0]; - scope_update_fn = [id, &items, i](context & ctx) { - ctx.set_val(id, items[i]->as_array()[0]); - }; - } else { - // case example: {% for item in list %} - scope_update_fn = [id, &items, i](context & ctx) { - ctx.set_val(id, items[i]); - }; - } - - } else if (is_stmt(loopvar)) { - // case example: {% for key, value in dict %} - auto tuple = cast_stmt(loopvar); - if (!is_val(current)) { - throw std::runtime_error("Cannot unpack non-iterable type: " + current->type()); - } - auto & c_arr = current->as_array(); - if (tuple->val.size() != c_arr.size()) { - throw std::runtime_error(std::string("Too ") + (tuple->val.size() > c_arr.size() ? "few" : "many") + " items to unpack"); - } - scope_update_fn = [tuple, &items, i](context & ctx) { - auto & c_arr = items[i]->as_array(); - for (size_t j = 0; j < tuple->val.size(); ++j) { - if (!is_stmt(tuple->val[j])) { - throw std::runtime_error("Cannot unpack non-identifier type: " + tuple->val[j]->type()); - } - auto id = cast_stmt(tuple->val[j])->val; - ctx.set_val(id, c_arr[j]); - } - }; - - } else { - throw std::runtime_error("Invalid loop variable(s): " + loopvar->type()); - } - - if (select_expr && test_expr) { - scope_update_fn(loop_scope); - value test_val = test_expr->execute(loop_scope); - if (!test_val->as_bool()) { - continue; - } - } - JJ_DEBUG("For loop: adding item type %s at index %zu", current->type().c_str(), i); - filtered_items.push_back(current); - scope_update_fns.push_back(scope_update_fn); - } - JJ_DEBUG("For loop: %zu items after filtering", filtered_items.size()); - - auto result = mk_val(); - - bool noIteration = true; - for (size_t i = 0; i < filtered_items.size(); i++) { - JJ_DEBUG("For loop iteration %zu/%zu", i + 1, filtered_items.size()); - value_object loop_obj = mk_val(); - loop_obj->has_builtins = false; // loop object has no builtins - loop_obj->insert("index", mk_val(i + 1)); - loop_obj->insert("index0", mk_val(i)); - loop_obj->insert("revindex", mk_val(filtered_items.size() - i)); - loop_obj->insert("revindex0", mk_val(filtered_items.size() - i - 1)); - loop_obj->insert("first", mk_val(i == 0)); - loop_obj->insert("last", mk_val(i == filtered_items.size() - 1)); - loop_obj->insert("length", mk_val(filtered_items.size())); - loop_obj->insert("previtem", i > 0 ? filtered_items[i - 1] : mk_val("previtem")); - loop_obj->insert("nextitem", i < filtered_items.size() - 1 ? filtered_items[i + 1] : mk_val("nextitem")); - scope.set_val("loop", loop_obj); - scope_update_fns[i](scope); - try { - for (auto & stmt : body) { - value val = stmt->execute(scope); - result->push_back(val); - } - } catch (const continue_statement::signal &) { - continue; - } catch (const break_statement::signal &) { - break; - } - noIteration = false; - } - - JJ_DEBUG("For loop complete, total iterations: %zu", filtered_items.size()); - if (noIteration) { - for (auto & stmt : default_block) { - value val = stmt->execute(ctx); - result->push_back(val); - } - } - - // convert to string parts - value_string str = mk_val(); - gather_string_parts_recursive(result, str); - return str; -} - -value set_statement::execute_impl(context & ctx) { - auto rhs = val ? val->execute(ctx) : exec_statements(body, ctx); - - if (is_stmt(assignee)) { - // case: {% set my_var = value %} - auto var_name = cast_stmt(assignee)->val; - JJ_DEBUG("Setting global variable '%s' with value type %s", var_name.c_str(), rhs->type().c_str()); - ctx.set_val(var_name, rhs); - - } else if (is_stmt(assignee)) { - // case: {% set a, b = value %} - auto tuple = cast_stmt(assignee); - if (!is_val(rhs)) { - throw std::runtime_error("Cannot unpack non-iterable type in set: " + rhs->type()); - } - auto & arr = rhs->as_array(); - if (arr.size() != tuple->val.size()) { - throw std::runtime_error(std::string("Too ") + (tuple->val.size() > arr.size() ? "few" : "many") + " items to unpack in set"); - } - for (size_t i = 0; i < tuple->val.size(); ++i) { - auto & elem = tuple->val[i]; - if (!is_stmt(elem)) { - throw std::runtime_error("Cannot unpack to non-identifier in set: " + elem->type()); - } - auto var_name = cast_stmt(elem)->val; - ctx.set_val(var_name, arr[i]); - } - - } else if (is_stmt(assignee)) { - // case: {% set ns.my_var = value %} - auto member = cast_stmt(assignee); - if (member->computed) { - throw std::runtime_error("Cannot assign to computed member"); - } - if (!is_stmt(member->property)) { - throw std::runtime_error("Cannot assign to member with non-identifier property"); - } - auto prop_name = cast_stmt(member->property)->val; - - value object = member->object->execute(ctx); - if (!is_val(object)) { - throw std::runtime_error("Cannot assign to member of non-object"); - } - auto obj_ptr = cast_val(object); - JJ_DEBUG("Setting object property '%s' with value type %s", prop_name.c_str(), rhs->type().c_str()); - obj_ptr->insert(prop_name, rhs); - - } else { - throw std::runtime_error("Invalid LHS inside assignment expression: " + assignee->type()); - } - return mk_val(); -} - -value macro_statement::execute_impl(context & ctx) { - if (!is_stmt(this->name)) { - throw std::runtime_error("Macro name must be an identifier"); - } - std::string name = cast_stmt(this->name)->val; - - const func_handler func = [this, name, &ctx](const func_args & args) -> value { - size_t expected_count = this->args.size(); - size_t input_count = args.count(); - - JJ_DEBUG("Invoking macro '%s' with %zu input arguments (expected %zu)", name.c_str(), input_count, expected_count); - context macro_ctx(ctx); // new scope for macro execution - - // bind parameters - for (size_t i = 0; i < expected_count; ++i) { - if (i < input_count) { - if (is_stmt(this->args[i])) { - // normal parameter - std::string param_name = cast_stmt(this->args[i])->val; - value param_value = args.get_kwarg_or_pos(param_name, i); - JJ_DEBUG(" Binding parameter '%s' to argument of type %s", param_name.c_str(), param_value->type().c_str()); - macro_ctx.set_val(param_name, param_value); - } else if (is_stmt(this->args[i])) { - // default argument used as normal parameter - auto kwarg = cast_stmt(this->args[i]); - if (!is_stmt(kwarg->key)) { - throw std::runtime_error("Keyword argument key must be an identifier in macro '" + name + "'"); - } - std::string param_name = cast_stmt(kwarg->key)->val; - value param_value = args.get_kwarg_or_pos(param_name, i); - JJ_DEBUG(" Binding parameter '%s' to argument of type %s", param_name.c_str(), param_value->type().c_str()); - macro_ctx.set_val(param_name, param_value); - } else { - throw std::runtime_error("Invalid parameter type in macro '" + name + "'"); - } - } else { - auto & default_arg = this->args[i]; - if (is_stmt(default_arg)) { - auto kwarg = cast_stmt(default_arg); - if (!is_stmt(kwarg->key)) { - throw std::runtime_error("Keyword argument key must be an identifier in macro '" + name + "'"); - } - std::string param_name = cast_stmt(kwarg->key)->val; - JJ_DEBUG(" Binding parameter '%s' to default argument of type %s", param_name.c_str(), kwarg->val->type().c_str()); - macro_ctx.set_val(param_name, kwarg->val->execute(ctx)); - } else { - throw std::runtime_error("Not enough arguments provided to macro '" + name + "'"); - } - //std::string param_name = cast_stmt(default_args[i])->val; - //JJ_DEBUG(" Binding parameter '%s' to default", param_name.c_str()); - //macro_ctx.var[param_name] = default_args[i]->execute(ctx); - } - } - - // execute macro body - JJ_DEBUG("Executing macro '%s' body with %zu statements", name.c_str(), this->body.size()); - auto res = exec_statements(this->body, macro_ctx); - JJ_DEBUG("Macro '%s' execution complete, result: %s", name.c_str(), res->val_str.str().c_str()); - return res; - }; - - JJ_DEBUG("Defining macro '%s' with %zu parameters", name.c_str(), args.size()); - ctx.set_val(name, mk_val(name, func)); - return mk_val(); -} - -value member_expression::execute_impl(context & ctx) { - value object = this->object->execute(ctx); - - value property; - if (this->computed) { - // syntax: obj[expr] - JJ_DEBUG("Member expression, computing property type %s", this->property->type().c_str()); - - int64_t arr_size = 0; - if (is_val(object)) { - arr_size = object->as_array().size(); - } else if (is_val(object)) { - arr_size = object->as_string().length(); - } - - if (is_stmt(this->property)) { - auto s = cast_stmt(this->property); - value start_val = s->start_expr ? s->start_expr->execute(ctx) : mk_val(0); - value stop_val = s->stop_expr ? s->stop_expr->execute(ctx) : mk_val(arr_size); - value step_val = s->step_expr ? s->step_expr->execute(ctx) : mk_val(1); - - // translate to function call: obj.slice(start, stop, step) - JJ_DEBUG("Member expression is a slice: start %s, stop %s, step %s", - start_val->as_repr().c_str(), - stop_val->as_repr().c_str(), - step_val->as_repr().c_str()); - auto slice_func = try_builtin_func(ctx, "slice", object); - func_args args(ctx); - args.push_back(start_val); - args.push_back(stop_val); - args.push_back(step_val); - return slice_func->invoke(args); - } else { - property = this->property->execute(ctx); - } - } else { - // syntax: obj.prop - if (!is_stmt(this->property)) { - throw std::runtime_error("Static member property must be an identifier"); - } - property = mk_val(cast_stmt(this->property)->val); - std::string prop = property->as_string().str(); - JJ_DEBUG("Member expression, object type %s, static property '%s'", object->type().c_str(), prop.c_str()); - - // behavior of jinja2: obj having prop as a built-in function AND 'prop', as an object key, - // then obj.prop returns the built-in function, not the property value. - // while obj['prop'] returns the property value. - // example: {"obj": {"items": 123}} -> obj.items is the built-in function, obj['items'] is 123 - - value val = try_builtin_func(ctx, prop, object, true); - if (!is_val(val)) { - return val; - } - // else, fallthrough to normal property access below - } - - JJ_DEBUG("Member expression on object type %s, property type %s", object->type().c_str(), property->type().c_str()); - value val = mk_val("object_property"); - - if (property->is_undefined()) { - JJ_DEBUG("%s", "Member expression property is undefined, returning undefined"); - return val; - } - - ensure_key_type_allowed(property); - - if (is_val(object)) { - JJ_DEBUG("%s", "Accessing property on undefined object, returning undefined"); - return val; - - } else if (is_val(object)) { - auto key = property->as_string().str(); - val = object->at(property, val); - if (is_val(val)) { - val = try_builtin_func(ctx, key, object, true); - } - JJ_DEBUG("Accessed property '%s' value, got type: %s", key.c_str(), val->type().c_str()); - - } else if (is_val(object) || is_val(object)) { - if (is_val(property)) { - int64_t index = property->as_int(); - JJ_DEBUG("Accessing %s index %d", object->type().c_str(), (int)index); - if (is_val(object)) { - auto & arr = object->as_array(); - if (index < 0) { - index += static_cast(arr.size()); - } - if (index >= 0 && index < static_cast(arr.size())) { - val = arr[index]; - } - } else { // value_string - auto str = object->as_string().str(); - if (index >= 0 && index < static_cast(str.size())) { - val = mk_val(std::string(1, str[index])); - } - } - - } else if (is_val(property)) { - auto key = property->as_string().str(); - JJ_DEBUG("Accessing %s built-in '%s'", is_val(object) ? "array" : "string", key.c_str()); - val = try_builtin_func(ctx, key, object, true); - - } else { - throw std::runtime_error("Cannot access property with non-string/non-number: got " + property->type()); - } - } else { - if (!is_val(property)) { - throw std::runtime_error("Cannot access property with non-string: got " + property->type()); - } - auto key = property->as_string().str(); - val = try_builtin_func(ctx, key, object, true); - } - - if (ctx.is_get_stats && val && object && property) { - value_t::stats_t::mark_used(val); - value_t::stats_t::mark_used(object); - value_t::stats_t::mark_used(property); - if (is_val(property)) { - object->stats.ops.insert("array_access"); - } else if (is_val(property)) { - object->stats.ops.insert("object_access"); - } - } - - return val; -} - -value call_expression::execute_impl(context & ctx) { - // gather arguments - func_args args(ctx); - for (auto & arg_stmt : this->args) { - auto arg_val = arg_stmt->execute(ctx); - JJ_DEBUG(" Argument type: %s", arg_val->type().c_str()); - args.push_back(arg_val); - } - // execute callee - value callee_val = callee->execute(ctx); - if (!is_val(callee_val)) { - throw std::runtime_error("Callee is not a function: got " + callee_val->type()); - } - auto * callee_func = cast_val(callee_val); - JJ_DEBUG("Calling function '%s' with %zu arguments", callee_func->name.c_str(), args.count()); - return callee_func->invoke(args); -} - -value keyword_argument_expression::execute_impl(context & ctx) { - if (!is_stmt(key)) { - throw std::runtime_error("Keyword argument key must be identifiers"); - } - - std::string k = cast_stmt(key)->val; - JJ_DEBUG("Keyword argument expression key: %s, value: %s", k.c_str(), val->type().c_str()); - - value v = val->execute(ctx); - JJ_DEBUG("Keyword argument value executed, type: %s", v->type().c_str()); - - return mk_val(k, v); -} - -} // namespace jinja diff --git a/common/jinja/runtime.h b/common/jinja/runtime.h deleted file mode 100644 index b6f4a6ab48e4..000000000000 --- a/common/jinja/runtime.h +++ /dev/null @@ -1,652 +0,0 @@ -#pragma once - -#include "lexer.h" -#include "value.h" - -#include -#include -#include -#include -#include -#include - -#define JJ_DEBUG(msg, ...) do { if (g_jinja_debug) printf("%s:%-3d : " msg "\n", FILENAME, __LINE__, __VA_ARGS__); } while (0) - -extern bool g_jinja_debug; - -namespace jinja { - -struct statement; -using statement_ptr = std::unique_ptr; -using statements = std::vector; - -// Helpers for dynamic casting and type checking -template -struct extract_pointee_unique { - using type = T; -}; -template -struct extract_pointee_unique> { - using type = U; -}; -template -bool is_stmt(const statement_ptr & ptr) { - return dynamic_cast(ptr.get()) != nullptr; -} -template -T * cast_stmt(statement_ptr & ptr) { - return dynamic_cast(ptr.get()); -} -template -const T * cast_stmt(const statement_ptr & ptr) { - return dynamic_cast(ptr.get()); -} -// End Helpers - - -// not thread-safe -void enable_debug(bool enable); - -struct context { - std::shared_ptr src; // for debugging; use shared_ptr to avoid copying on scope creation - std::time_t current_time; // for functions that need current time - - bool is_get_stats = false; // whether to collect stats - - // src is optional, used for error reporting - context(std::string src = "") : src(std::make_shared(std::move(src))) { - env = mk_val(); - env->has_builtins = false; // context object has no builtins - env->insert("true", mk_val(true)); - env->insert("True", mk_val(true)); - env->insert("false", mk_val(false)); - env->insert("False", mk_val(false)); - env->insert("none", mk_val()); - env->insert("None", mk_val()); - current_time = std::time(nullptr); - } - ~context() = default; - - context(const context & parent) : context() { - // inherit variables (for example, when entering a new scope) - auto & pvar = parent.env->as_ordered_object(); - for (const auto & pair : pvar) { - set_val(pair.first, pair.second); - } - current_time = parent.current_time; - is_get_stats = parent.is_get_stats; - src = parent.src; - } - - value get_val(const std::string & name) { - value default_val = mk_val(name); - return env->at(name, default_val); - } - - void set_val(const std::string & name, const value & val) { - env->insert(name, val); - } - - void set_val(const value & name, const value & val) { - env->insert(name, val); - } - - void print_vars() const { - printf("Context Variables:\n%s\n", value_to_json(env, 2).c_str()); - } - -private: - value_object env; -}; - -/** - * Base class for all nodes in the AST. - */ -struct statement { - size_t pos; // position in source, for debugging - virtual ~statement() = default; - virtual std::string type() const { return "Statement"; } - - // execute_impl must be overridden by derived classes - virtual value execute_impl(context &) { throw_exec_error(); } - // execute is the public method to execute a statement with error handling - value execute(context &); - -private: - [[noreturn]] void throw_exec_error() const { - throw std::runtime_error("cannot exec " + type()); - } -}; - -// Type Checking Utilities - -template -static void chk_type(const statement_ptr & ptr) { - if (!ptr) return; // Allow null for optional fields - assert(dynamic_cast(ptr.get()) != nullptr); -} - -template -static void chk_type(const statement_ptr & ptr) { - if (!ptr) return; - assert(dynamic_cast(ptr.get()) != nullptr || dynamic_cast(ptr.get()) != nullptr); -} - -// Base Types - -/** - * Expressions will result in a value at runtime (unlike statements). - */ -struct expression : public statement { - std::string type() const override { return "Expression"; } -}; - -// Statements - -struct program : public statement { - statements body; - - program() = default; - explicit program(statements && body) : body(std::move(body)) {} - std::string type() const override { return "Program"; } - [[noreturn]] value execute_impl(context &) override { - throw std::runtime_error("Cannot execute program directly, use jinja::runtime instead"); - } -}; - -struct if_statement : public statement { - statement_ptr test; - statements body; - statements alternate; - - if_statement(statement_ptr && test, statements && body, statements && alternate) - : test(std::move(test)), body(std::move(body)), alternate(std::move(alternate)) { - chk_type(this->test); - } - - std::string type() const override { return "If"; } - value execute_impl(context & ctx) override; -}; - -struct identifier; -struct tuple_literal; - -/** - * Loop over each item in a sequence - * https://jinja.palletsprojects.com/en/3.0.x/templates/#for - */ -struct for_statement : public statement { - statement_ptr loopvar; // Identifier | TupleLiteral - statement_ptr iterable; - statements body; - statements default_block; // if no iteration took place - - for_statement(statement_ptr && loopvar, statement_ptr && iterable, statements && body, statements && default_block) - : loopvar(std::move(loopvar)), iterable(std::move(iterable)), - body(std::move(body)), default_block(std::move(default_block)) { - chk_type(this->loopvar); - chk_type(this->iterable); - } - - std::string type() const override { return "For"; } - value execute_impl(context & ctx) override; -}; - -struct break_statement : public statement { - std::string type() const override { return "Break"; } - - struct signal : public std::exception { - const char* what() const noexcept override { - return "Break statement executed"; - } - }; - - [[noreturn]] value execute_impl(context &) override { - throw break_statement::signal(); - } -}; - -struct continue_statement : public statement { - std::string type() const override { return "Continue"; } - - struct signal : public std::exception { - const char* what() const noexcept override { - return "Continue statement executed"; - } - }; - - [[noreturn]] value execute_impl(context &) override { - throw continue_statement::signal(); - } -}; - -// do nothing -struct noop_statement : public statement { - std::string type() const override { return "Noop"; } - value execute_impl(context &) override { - return mk_val(); - } -}; - -struct set_statement : public statement { - statement_ptr assignee; - statement_ptr val; - statements body; - - set_statement(statement_ptr && assignee, statement_ptr && value, statements && body) - : assignee(std::move(assignee)), val(std::move(value)), body(std::move(body)) { - chk_type(this->assignee); - chk_type(this->val); - } - - std::string type() const override { return "Set"; } - value execute_impl(context & ctx) override; -}; - -struct macro_statement : public statement { - statement_ptr name; - statements args; - statements body; - - macro_statement(statement_ptr && name, statements && args, statements && body) - : name(std::move(name)), args(std::move(args)), body(std::move(body)) { - chk_type(this->name); - for (const auto& arg : this->args) chk_type(arg); - } - - std::string type() const override { return "Macro"; } - value execute_impl(context & ctx) override; -}; - -struct comment_statement : public statement { - std::string val; - explicit comment_statement(const std::string & v) : val(v) {} - std::string type() const override { return "Comment"; } - value execute_impl(context &) override { - return mk_val(); - } -}; - -// Expressions - -// Represents an omitted expression in a computed member, e.g. `a[]`. -struct blank_expression : public expression { - std::string type() const override { return "BlankExpression"; } - value execute_impl(context &) override { - return mk_val(); - } -}; - -struct member_expression : public expression { - statement_ptr object; - statement_ptr property; - bool computed; // true if obj[expr] and false if obj.prop - - member_expression(statement_ptr && object, statement_ptr && property, bool computed) - : object(std::move(object)), property(std::move(property)), computed(computed) { - chk_type(this->object); - chk_type(this->property); - } - std::string type() const override { return "MemberExpression"; } - value execute_impl(context & ctx) override; -}; - -struct call_expression : public expression { - statement_ptr callee; - statements args; - - call_expression(statement_ptr && callee, statements && args) - : callee(std::move(callee)), args(std::move(args)) { - chk_type(this->callee); - for (const auto& arg : this->args) chk_type(arg); - } - std::string type() const override { return "CallExpression"; } - value execute_impl(context & ctx) override; -}; - -/** - * Represents a user-defined variable or symbol in the template. - */ -struct identifier : public expression { - std::string val; - explicit identifier(const std::string & val) : val(val) {} - std::string type() const override { return "Identifier"; } - value execute_impl(context & ctx) override; -}; - -// Literals - -struct integer_literal : public expression { - int64_t val; - explicit integer_literal(int64_t val) : val(val) {} - std::string type() const override { return "IntegerLiteral"; } - value execute_impl(context &) override { - return mk_val(val); - } -}; - -struct float_literal : public expression { - double val; - explicit float_literal(double val) : val(val) {} - std::string type() const override { return "FloatLiteral"; } - value execute_impl(context &) override { - return mk_val(val); - } -}; - -struct string_literal : public expression { - std::string val; - explicit string_literal(const std::string & val) : val(val) {} - std::string type() const override { return "StringLiteral"; } - value execute_impl(context &) override { - return mk_val(val); - } -}; - -struct array_literal : public expression { - statements val; - explicit array_literal(statements && val) : val(std::move(val)) { - for (const auto& item : this->val) chk_type(item); - } - std::string type() const override { return "ArrayLiteral"; } - value execute_impl(context & ctx) override { - auto arr = mk_val(); - for (const auto & item_stmt : val) { - arr->push_back(item_stmt->execute(ctx)); - } - return arr; - } -}; - -struct tuple_literal : public expression { - statements val; - explicit tuple_literal(statements && val) : val(std::move(val)) { - for (const auto& item : this->val) chk_type(item); - } - std::string type() const override { return "TupleLiteral"; } - value execute_impl(context & ctx) override { - auto arr = mk_val(); - for (const auto & item_stmt : val) { - arr->push_back(item_stmt->execute(ctx)); - } - return mk_val(std::move(arr->as_array())); - } -}; - -struct object_literal : public expression { - std::vector> val; - explicit object_literal(std::vector> && val) - : val(std::move(val)) { - for (const auto & pair : this->val) { - chk_type(pair.first); - chk_type(pair.second); - } - } - std::string type() const override { return "ObjectLiteral"; } - value execute_impl(context & ctx) override; -}; - -// Complex Expressions - -/** - * An operation with two sides, separated by an operator. - * Note: Either side can be a Complex Expression, with order - * of operations being determined by the operator. - */ -struct binary_expression : public expression { - token op; - statement_ptr left; - statement_ptr right; - - binary_expression(token op, statement_ptr && left, statement_ptr && right) - : op(std::move(op)), left(std::move(left)), right(std::move(right)) { - chk_type(this->left); - chk_type(this->right); - } - std::string type() const override { return "BinaryExpression"; } - value execute_impl(context & ctx) override; -}; - -/** - * An operation with two sides, separated by the | operator. - * Operator precedence: https://github.com/pallets/jinja/issues/379#issuecomment-168076202 - */ -struct filter_expression : public expression { - // either an expression or a value is allowed - statement_ptr operand; - value_string val; // will be set by filter_statement - - statement_ptr filter; - - filter_expression(statement_ptr && operand, statement_ptr && filter) - : operand(std::move(operand)), filter(std::move(filter)) { - chk_type(this->operand); - chk_type(this->filter); - } - - filter_expression(value_string && val, statement_ptr && filter) - : val(std::move(val)), filter(std::move(filter)) { - chk_type(this->filter); - } - - std::string type() const override { return "FilterExpression"; } - value execute_impl(context & ctx) override; -}; - -struct filter_statement : public statement { - statement_ptr filter; - statements body; - - filter_statement(statement_ptr && filter, statements && body) - : filter(std::move(filter)), body(std::move(body)) { - chk_type(this->filter); - } - std::string type() const override { return "FilterStatement"; } - value execute_impl(context & ctx) override; -}; - -/** - * An operation which filters a sequence of objects by applying a test to each object, - * and only selecting the objects with the test succeeding. - * - * It may also be used as a shortcut for a ternary operator. - */ -struct select_expression : public expression { - statement_ptr lhs; - statement_ptr test; - - select_expression(statement_ptr && lhs, statement_ptr && test) - : lhs(std::move(lhs)), test(std::move(test)) { - chk_type(this->lhs); - chk_type(this->test); - } - std::string type() const override { return "SelectExpression"; } - value execute_impl(context & ctx) override { - auto predicate = test->execute_impl(ctx); - if (!predicate->as_bool()) { - return mk_val(); - } - return lhs->execute_impl(ctx); - } -}; - -/** - * An operation with two sides, separated by the "is" operator. - * NOTE: "value is something" translates to function call "test_is_something(value)" - */ -struct test_expression : public expression { - statement_ptr operand; - bool negate; - statement_ptr test; - - test_expression(statement_ptr && operand, bool negate, statement_ptr && test) - : operand(std::move(operand)), negate(negate), test(std::move(test)) { - chk_type(this->operand); - chk_type(this->test); - } - std::string type() const override { return "TestExpression"; } - value execute_impl(context & ctx) override; -}; - -/** - * An operation with one side (operator on the left). - */ -struct unary_expression : public expression { - token op; - statement_ptr argument; - - unary_expression(token op, statement_ptr && argument) - : op(std::move(op)), argument(std::move(argument)) { - chk_type(this->argument); - } - std::string type() const override { return "UnaryExpression"; } - value execute_impl(context & ctx) override; -}; - -struct slice_expression : public expression { - statement_ptr start_expr; - statement_ptr stop_expr; - statement_ptr step_expr; - - slice_expression(statement_ptr && start_expr, statement_ptr && stop_expr, statement_ptr && step_expr) - : start_expr(std::move(start_expr)), stop_expr(std::move(stop_expr)), step_expr(std::move(step_expr)) { - chk_type(this->start_expr); - chk_type(this->stop_expr); - chk_type(this->step_expr); - } - std::string type() const override { return "SliceExpression"; } - [[noreturn]] value execute_impl(context &) override { - throw std::runtime_error("must be handled by MemberExpression"); - } -}; - -struct keyword_argument_expression : public expression { - statement_ptr key; - statement_ptr val; - - keyword_argument_expression(statement_ptr && key, statement_ptr && val) - : key(std::move(key)), val(std::move(val)) { - chk_type(this->key); - chk_type(this->val); - } - std::string type() const override { return "KeywordArgumentExpression"; } - value execute_impl(context & ctx) override; -}; - -struct spread_expression : public expression { - statement_ptr argument; - explicit spread_expression(statement_ptr && argument) : argument(std::move(argument)) { - chk_type(this->argument); - } - std::string type() const override { return "SpreadExpression"; } -}; - -struct call_statement : public statement { - statement_ptr call; - statements caller_args; - statements body; - - call_statement(statement_ptr && call, statements && caller_args, statements && body) - : call(std::move(call)), caller_args(std::move(caller_args)), body(std::move(body)) { - chk_type(this->call); - for (const auto & arg : this->caller_args) chk_type(arg); - } - std::string type() const override { return "CallStatement"; } -}; - -struct ternary_expression : public expression { - statement_ptr condition; - statement_ptr true_expr; - statement_ptr false_expr; - - ternary_expression(statement_ptr && condition, statement_ptr && true_expr, statement_ptr && false_expr) - : condition(std::move(condition)), true_expr(std::move(true_expr)), false_expr(std::move(false_expr)) { - chk_type(this->condition); - chk_type(this->true_expr); - chk_type(this->false_expr); - } - std::string type() const override { return "Ternary"; } - value execute_impl(context & ctx) override { - value cond_val = condition->execute(ctx); - if (cond_val->as_bool()) { - return true_expr->execute(ctx); - } else { - return false_expr->execute(ctx); - } - } -}; - -struct raised_exception : public std::exception { - std::string message; - raised_exception(const std::string & msg) : message(msg) {} - const char* what() const noexcept override { - return message.c_str(); - } -}; - -// Used to rethrow exceptions with modified messages -struct rethrown_exception : public std::exception { - std::string message; - rethrown_exception(const std::string & msg) : message(msg) {} - const char* what() const noexcept override { - return message.c_str(); - } -}; - -////////////////////// - -static void gather_string_parts_recursive(const value & val, value_string & parts) { - // TODO: probably allow print value_none as "None" string? currently this breaks some templates - if (is_val(val)) { - const auto & str_val = cast_val(val)->val_str; - parts->val_str.append(str_val); - } else if (is_val(val) || is_val(val) || is_val(val)) { - std::string str_val = val->as_string().str(); - parts->val_str.append(str_val); - } else if (is_val(val)) { - auto items = cast_val(val)->as_array(); - for (const auto & item : items) { - gather_string_parts_recursive(item, parts); - } - } -} - -static std::string render_string_parts(const value_string & parts) { - std::ostringstream oss; - for (const auto & part : parts->val_str.parts) { - oss << part.val; - } - return oss.str(); -} - -struct runtime { - context & ctx; - explicit runtime(context & ctx) : ctx(ctx) {} - - value_array execute(const program & prog) { - value_array results = mk_val(); - for (const auto & stmt : prog.body) { - value res = stmt->execute(ctx); - results->push_back(std::move(res)); - } - return results; - } - - static value_string gather_string_parts(const value & val) { - value_string parts = mk_val(); - gather_string_parts_recursive(val, parts); - // join consecutive parts with the same type - auto & p = parts->val_str.parts; - for (size_t i = 1; i < p.size(); ) { - if (p[i].is_input == p[i - 1].is_input) { - p[i - 1].val += p[i].val; - p.erase(p.begin() + i); - } else { - i++; - } - } - return parts; - } -}; - -} // namespace jinja diff --git a/common/jinja/string-rust/Cargo.toml b/common/jinja/string-rust/Cargo.toml new file mode 100644 index 000000000000..b456340da231 --- /dev/null +++ b/common/jinja/string-rust/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "llama-common-jinja-string-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[lib] +crate-type = ["staticlib", "rlib"] + +[dependencies] diff --git a/common/jinja/string-rust/src/lib.rs b/common/jinja/string-rust/src/lib.rs new file mode 100644 index 000000000000..c1f14056b66e --- /dev/null +++ b/common/jinja/string-rust/src/lib.rs @@ -0,0 +1,156 @@ +use std::slice; + +#[repr(C)] +pub struct JinjaString { + data: *mut u8, + len: usize, +} + +#[no_mangle] +pub unsafe extern "C" fn llama_jinja_string_free(value: JinjaString) { + if !value.data.is_null() { + drop(Vec::from_raw_parts(value.data, value.len, value.len)); + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_jinja_string_is_uppercase(data: *const u8, len: usize) -> bool { + bytes(data, len) + .iter() + .all(|&byte| !byte.is_ascii_lowercase()) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_jinja_string_is_lowercase(data: *const u8, len: usize) -> bool { + bytes(data, len) + .iter() + .all(|&byte| !byte.is_ascii_uppercase()) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_jinja_string_uppercase(data: *const u8, len: usize) -> JinjaString { + into_ffi(bytes(data, len).iter().map(|byte| byte.to_ascii_uppercase()).collect()) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_jinja_string_lowercase(data: *const u8, len: usize) -> JinjaString { + into_ffi(bytes(data, len).iter().map(|byte| byte.to_ascii_lowercase()).collect()) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_jinja_string_capitalize(data: *const u8, len: usize) -> JinjaString { + let input = bytes(data, len); + if input.is_empty() { + return into_ffi(Vec::new()); + } + + let mut result = Vec::with_capacity(input.len()); + result.push(input[0].to_ascii_uppercase()); + result.extend(input[1..].iter().map(|byte| byte.to_ascii_lowercase())); + into_ffi(result) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_jinja_string_titlecase(data: *const u8, len: usize) -> JinjaString { + let mut capitalize_next = true; + let mut result = Vec::with_capacity(len); + for &byte in bytes(data, len) { + if byte.is_ascii_whitespace() { + capitalize_next = true; + result.push(byte); + } else if capitalize_next { + result.push(byte.to_ascii_uppercase()); + capitalize_next = false; + } else { + result.push(byte.to_ascii_lowercase()); + } + } + into_ffi(result) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_jinja_string_strip( + data: *const u8, + len: usize, + left: bool, + right: bool, + chars: *const u8, + chars_len: usize, + has_chars: bool, +) -> JinjaString { + let input = bytes(data, len); + let chars = if has_chars { + Some(bytes(chars, chars_len)) + } else { + None + }; + + let mut start = 0usize; + let mut end = input.len(); + if left { + while start < end && match_char(input[start], chars) { + start += 1; + } + } + if right { + while end > start && match_char(input[end - 1], chars) { + end -= 1; + } + } + + into_ffi(input[start..end].to_vec()) +} + +unsafe fn bytes<'a>(data: *const u8, len: usize) -> &'a [u8] { + if data.is_null() && len != 0 { + &[] + } else { + slice::from_raw_parts(data, len) + } +} + +fn match_char(byte: u8, chars: Option<&[u8]>) -> bool { + if let Some(chars) = chars { + chars.contains(&byte) + } else { + byte.is_ascii_whitespace() + } +} + +fn into_ffi(mut bytes: Vec) -> JinjaString { + let len = bytes.len(); + let data = bytes.as_mut_ptr(); + std::mem::forget(bytes); + JinjaString { data, len } +} + +#[cfg(test)] +mod tests { + use super::*; + + unsafe fn owned(value: JinjaString) -> String { + let data = slice::from_raw_parts(value.data, value.len).to_vec(); + llama_jinja_string_free(value); + String::from_utf8(data).unwrap() + } + + #[test] + fn case_transforms_match_cpp_ascii_behavior() { + assert_eq!(unsafe { owned(llama_jinja_string_uppercase(b"aBc!".as_ptr(), 4)) }, "ABC!"); + assert_eq!(unsafe { owned(llama_jinja_string_lowercase(b"aBc!".as_ptr(), 4)) }, "abc!"); + assert_eq!(unsafe { owned(llama_jinja_string_capitalize(b"hELLO".as_ptr(), 5)) }, "Hello"); + assert_eq!(unsafe { owned(llama_jinja_string_titlecase(b"hi THERE".as_ptr(), 8)) }, "Hi There"); + } + + #[test] + fn strips_default_and_custom_chars() { + assert_eq!( + unsafe { owned(llama_jinja_string_strip(b" hi \n".as_ptr(), 6, true, true, std::ptr::null(), 0, false)) }, + "hi" + ); + assert_eq!( + unsafe { owned(llama_jinja_string_strip(b"..hi..".as_ptr(), 6, true, true, b".".as_ptr(), 1, true)) }, + "hi" + ); + } +} diff --git a/common/jinja/string.cpp b/common/jinja/string.cpp deleted file mode 100644 index 8087e15b3502..000000000000 --- a/common/jinja/string.cpp +++ /dev/null @@ -1,213 +0,0 @@ -#include "jinja/string.h" -#include "jinja/value.h" - -#include -#include -#include -#include -#include -#include - -namespace jinja { - -// -// string_part -// - -bool string_part::is_uppercase() const { - for (char c : val) { - if (std::islower(static_cast(c))) { - return false; - } - } - return true; -} - -bool string_part::is_lowercase() const { - for (char c : val) { - if (std::isupper(static_cast(c))) { - return false; - } - } - return true; -} - -// -// string -// - -void string::mark_input() { - for (auto & part : parts) { - part.is_input = true; - } -} - -std::string string::str() const { - if (parts.size() == 1) { - return parts[0].val; - } - std::ostringstream oss; - for (const auto & part : parts) { - oss << part.val; - } - return oss.str(); -} - -size_t string::length() const { - size_t len = 0; - for (const auto & part : parts) { - len += part.val.length(); - } - return len; -} - -void string::hash_update(hasher & hash) const noexcept { - for (const auto & part : parts) { - hash.update(part.val.data(), part.val.length()); - } -} - -bool string::all_parts_are_input() const { - for (const auto & part : parts) { - if (!part.is_input) { - return false; - } - } - return true; -} - -bool string::is_uppercase() const { - for (const auto & part : parts) { - if (!part.is_uppercase()) { - return false; - } - } - return true; -} - -bool string::is_lowercase() const { - for (const auto & part : parts) { - if (!part.is_lowercase()) { - return false; - } - } - return true; -} - -// mark this string as input if other has ALL parts as input -void string::mark_input_based_on(const string & other) { - if (other.all_parts_are_input()) { - for (auto & part : parts) { - part.is_input = true; - } - } -} - -string string::append(const string & other) { - for (const auto & part : other.parts) { - parts.push_back(part); - } - return *this; -} - -// in-place transformation - -using transform_fn = std::function; -static string apply_transform(string & self, const transform_fn & fn) { - for (auto & part : self.parts) { - part.val = fn(part.val); - } - return self; -} - -string string::uppercase() { - return apply_transform(*this, [](const std::string & s) { - std::string res = s; - std::transform(res.begin(), res.end(), res.begin(), ::toupper); - return res; - }); -} -string string::lowercase() { - return apply_transform(*this, [](const std::string & s) { - std::string res = s; - std::transform(res.begin(), res.end(), res.begin(), ::tolower); - return res; - }); -} -string string::capitalize() { - return apply_transform(*this, [](const std::string & s) { - if (s.empty()) return s; - std::string res = s; - res[0] = ::toupper(static_cast(res[0])); - std::transform(res.begin() + 1, res.end(), res.begin() + 1, ::tolower); - return res; - }); -} -string string::titlecase() { - return apply_transform(*this, [](const std::string & s) { - std::string res = s; - bool capitalize_next = true; - for (char &c : res) { - if (isspace(static_cast(c))) { - capitalize_next = true; - } else if (capitalize_next) { - c = ::toupper(static_cast(c)); - capitalize_next = false; - } else { - c = ::tolower(static_cast(c)); - } - } - return res; - }); -} -string string::strip(bool left, bool right, std::optional chars) { - static auto strip_part = [](const std::string & s, bool left, bool right, std::optional chars) -> std::string { - size_t start = 0; - size_t end = s.length(); - auto match_char = [&chars](unsigned char c) -> bool { - return chars ? (*chars).find(c) != std::string::npos : isspace(c); - }; - if (left) { - while (start < end && match_char(static_cast(s[start]))) { - ++start; - } - } - if (right) { - while (end > start && match_char(static_cast(s[end - 1]))) { - --end; - } - } - return s.substr(start, end - start); - }; - if (parts.empty()) { - return *this; - } - if (left) { - for (size_t i = 0; i < parts.size(); ++i) { - parts[i].val = strip_part(parts[i].val, true, false, chars); - if (parts[i].val.empty()) { - // remove empty part - parts.erase(parts.begin() + i); - --i; - continue; - } else { - break; - } - } - } - if (right) { - for (size_t i = parts.size(); i-- > 0;) { - parts[i].val = strip_part(parts[i].val, false, true, chars); - if (parts[i].val.empty()) { - // remove empty part - parts.erase(parts.begin() + i); - continue; - } else { - break; - } - } - } - return *this; -} - -} // namespace jinja diff --git a/common/jinja/string.h b/common/jinja/string.h deleted file mode 100644 index c4963000adb8..000000000000 --- a/common/jinja/string.h +++ /dev/null @@ -1,61 +0,0 @@ -#pragma once - -#include -#include -#include - -#include "utils.h" - -namespace jinja { - -// allow differentiate between user input strings and template strings -// transformations should handle this information as follows: -// - one-to-one (e.g., uppercase, lowercase): preserve is_input flag -// - one-to-many (e.g., strip): if input string is marked as is_input, all resulting parts should be marked as is_input -// - many-to-one (e.g., concat): if ALL input parts are marked as is_input, resulting part should be marked as is_input -struct string_part { - bool is_input = false; // may skip parsing special tokens if true - std::string val; - - bool is_uppercase() const; - bool is_lowercase() const; -}; - -struct string { - std::vector parts; - string() = default; - string(const std::string & v, bool user_input = false) { - parts.push_back({user_input, v}); - } - string(int v) { - parts.push_back({false, std::to_string(v)}); - } - string(double v) { - parts.push_back({false, std::to_string(v)}); - } - - // mark all parts as user input - void mark_input(); - - std::string str() const; - size_t length() const; - void hash_update(hasher & hash) const noexcept; - bool all_parts_are_input() const; - bool is_uppercase() const; - bool is_lowercase() const; - - // mark this string as input if other has ALL parts as input - void mark_input_based_on(const string & other); - - string append(const string & other); - - // in-place transformations - - string uppercase(); - string lowercase(); - string capitalize(); - string titlecase(); - string strip(bool left, bool right, std::optional chars = std::nullopt); -}; - -} // namespace jinja diff --git a/common/jinja/utils.h b/common/jinja/utils.h deleted file mode 100644 index de6947fc28f3..000000000000 --- a/common/jinja/utils.h +++ /dev/null @@ -1,149 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include - -namespace jinja { - -static void string_replace_all(std::string & s, const std::string & search, const std::string & replace) { - if (search.empty()) { - return; - } - std::string builder; - builder.reserve(s.length()); - size_t pos = 0; - size_t last_pos = 0; - while ((pos = s.find(search, last_pos)) != std::string::npos) { - builder.append(s, last_pos, pos - last_pos); - builder.append(replace); - last_pos = pos + search.length(); - } - builder.append(s, last_pos, std::string::npos); - s = std::move(builder); -} - -// for displaying source code around error position -static std::string peak_source(const std::string & source, size_t pos, size_t max_peak_chars = 40) { - if (source.empty()) { - return "(no source available)"; - } - std::string output; - size_t start = (pos >= max_peak_chars) ? (pos - max_peak_chars) : 0; - size_t end = std::min(pos + max_peak_chars, source.length()); - std::string substr = source.substr(start, end - start); - string_replace_all(substr, "\n", "↵"); - output += "..." + substr + "...\n"; - std::string spaces(pos - start + 3, ' '); - output += spaces + "^"; - return output; -} - -static std::string fmt_error_with_source(const std::string & tag, const std::string & msg, const std::string & source, size_t pos) { - std::ostringstream oss; - oss << tag << ": " << msg << "\n"; - oss << peak_source(source, pos); - return oss.str(); -} - -// Note: this is a simple hasher, not cryptographically secure, just for hash table usage -struct hasher { - static constexpr auto size_t_digits = sizeof(size_t) * 8; - static constexpr size_t prime = size_t_digits == 64 ? 0x100000001b3 : 0x01000193; - static constexpr size_t seed = size_t_digits == 64 ? 0xcbf29ce484222325 : 0x811c9dc5; - static constexpr auto block_size = sizeof(size_t); // in bytes; allowing the compiler to vectorize the computation - - static_assert(size_t_digits == 64 || size_t_digits == 32); - static_assert(block_size == 8 || block_size == 4); - - uint8_t buffer[block_size]; - size_t idx = 0; // current index in buffer - size_t state = seed; - - hasher() = default; - hasher(const std::type_info & type_inf) noexcept { - const auto type_hash = type_inf.hash_code(); - update(&type_hash, sizeof(type_hash)); - } - - // Properties: - // - update is not associative: update(a).update(b) != update(b).update(a) - // - update(a ~ b) == update(a).update(b) with ~ as concatenation operator --> useful for streaming - // - update("", 0) --> state unchanged with empty input - hasher& update(void const * bytes, size_t len) noexcept { - const uint8_t * c = static_cast(bytes); - if (len == 0) { - return *this; - } - size_t processed = 0; - - // first, fill the existing buffer if it's partial - if (idx > 0) { - size_t to_fill = block_size - idx; - if (to_fill > len) { - to_fill = len; - } - std::memcpy(buffer + idx, c, to_fill); - idx += to_fill; - processed += to_fill; - if (idx == block_size) { - update_block(buffer); - idx = 0; - } - } - - // process full blocks from the remaining input - for (; processed + block_size <= len; processed += block_size) { - update_block(c + processed); - } - - // buffer any remaining bytes - size_t remaining = len - processed; - if (remaining > 0) { - std::memcpy(buffer, c + processed, remaining); - idx = remaining; - } - return *this; - } - - // convenience function for testing only - hasher& update(const std::string & s) noexcept { - return update(s.data(), s.size()); - } - - // finalize and get the hash value - // note: after calling digest, the hasher state is modified, do not call update() again - size_t digest() noexcept { - // if there are remaining bytes in buffer, fill the rest with zeros and process - if (idx > 0) { - for (size_t i = idx; i < block_size; ++i) { - buffer[i] = 0; - } - update_block(buffer); - idx = 0; - } - - return state; - } - -private: - // IMPORTANT: block must have at least block_size bytes - void update_block(const uint8_t * block) noexcept { - size_t blk = static_cast(block[0]) - | (static_cast(block[1]) << 8) - | (static_cast(block[2]) << 16) - | (static_cast(block[3]) << 24); - if constexpr (block_size == 8) { - blk = blk | (static_cast(block[4]) << 32) - | (static_cast(block[5]) << 40) - | (static_cast(block[6]) << 48) - | (static_cast(block[7]) << 56); - } - state ^= blk; - state *= prime; - } -}; - -} // namespace jinja diff --git a/common/jinja/value.cpp b/common/jinja/value.cpp deleted file mode 100644 index 0b79098cd1e7..000000000000 --- a/common/jinja/value.cpp +++ /dev/null @@ -1,1484 +0,0 @@ -#include "runtime.h" -#include "unicode.h" -#include "value.h" - -// for converting from JSON to jinja values -#include - -#include -#include -#include -#include -#include -#include - -#define FILENAME "jinja-value" - -namespace jinja { - -// func_args method implementations - -value func_args::get_kwarg(const std::string & key, value default_val) const { - for (const auto & arg : args) { - if (is_val(arg)) { - auto * kwarg = cast_val(arg); - if (kwarg->key == key) { - return kwarg->val; - } - } - } - return default_val; -} - -value func_args::get_kwarg_or_pos(const std::string & key, size_t pos) const { - value val = get_kwarg(key, mk_val()); - - if (val->is_undefined() && pos < count() && !is_val(args[pos])) { - return args[pos]; - } - - return val; -} - -value func_args::get_pos(size_t pos) const { - if (count() > pos) { - return args[pos]; - } - throw raised_exception("Function '" + func_name + "' expected at least " + std::to_string(pos + 1) + " arguments, got " + std::to_string(count())); -} - -value func_args::get_pos(size_t pos, value default_val) const { - if (count() > pos) { - return args[pos]; - } - return default_val; -} - -void func_args::push_back(const value & val) { - args.push_back(val); -} - -void func_args::push_front(const value & val) { - args.insert(args.begin(), val); -} - -const std::vector & func_args::get_args() const { - return args; -} - -/** - * Function that mimics Python's array slicing. - */ -template -static T slice(const T & array, int64_t start, int64_t stop, int64_t step = 1) { - int64_t len = static_cast(array.size()); - int64_t direction = (step > 0) ? 1 : ((step < 0) ? -1 : 0); - int64_t start_val = 0; - int64_t stop_val = 0; - if (direction >= 0) { - start_val = start; - if (start_val < 0) { - start_val = std::max(len + start_val, (int64_t)0); - } else { - start_val = std::min(start_val, len); - } - - stop_val = stop; - if (stop_val < 0) { - stop_val = std::max(len + stop_val, (int64_t)0); - } else { - stop_val = std::min(stop_val, len); - } - } else { - start_val = len - 1; - if (start_val < 0) { - start_val = std::max(len + start_val, (int64_t)-1); - } else { - start_val = std::min(start_val, len - 1); - } - - stop_val = -1; - if (stop_val < -1) { - stop_val = std::max(len + stop_val, (int64_t)-1); - } else { - stop_val = std::min(stop_val, len - 1); - } - } - T result; - if (direction == 0) { - return result; - } - for (int64_t i = start_val; direction * i < direction * stop_val; i += step) { - if (i >= 0 && i < len) { - result.push_back(array[static_cast(i)]); - } - } - return result; -} - -template -static value empty_value_fn(const func_args &) { - if constexpr (std::is_same_v) { - return mk_val(0); - } else if constexpr (std::is_same_v) { - return mk_val(0.0); - } else if constexpr (std::is_same_v) { - return mk_val(false); - } else { - return mk_val(); - } -} -template -static value test_type_fn(const func_args & args) { - args.ensure_count(1); - bool is_type = is_val(args.get_pos(0)); - JJ_DEBUG("test_type_fn: type=%s result=%d", typeid(T).name(), is_type ? 1 : 0); - return mk_val(is_type); -} -template -static value test_type_fn(const func_args & args) { - args.ensure_count(1); - bool is_type = is_val(args.get_pos(0)) || is_val(args.get_pos(0)); - JJ_DEBUG("test_type_fn: type=%s or %s result=%d", typeid(T).name(), typeid(U).name(), is_type ? 1 : 0); - return mk_val(is_type); -} -template -static value test_type_fn(const func_args & args) { - args.ensure_count(1); - bool is_type = is_val(args.get_pos(0)) || is_val(args.get_pos(0)) || is_val(args.get_pos(0)); - JJ_DEBUG("test_type_fn: type=%s, %s or %s result=%d", typeid(T).name(), typeid(U).name(), typeid(V).name(), is_type ? 1 : 0); - return mk_val(is_type); -} -template -static value test_compare_fn(const func_args & args) { - args.ensure_count(2, 2); - return mk_val(value_compare(args.get_pos(0), args.get_pos(1), op)); -} - -static void append_codepoint_as_ascii_json_escape(std::string & out, uint32_t codepoint) { - auto append_u16 = [&out](uint32_t value) { - char buf[8]; - snprintf(buf, sizeof(buf), "\\u%04x", static_cast(value)); - out += buf; - }; - - if (codepoint <= 0xFFFF) { - append_u16(codepoint); - return; - } - - codepoint -= 0x10000; - append_u16(0xD800 + ((codepoint >> 10) & 0x3FF)); - append_u16(0xDC00 + (codepoint & 0x3FF)); -} - -static std::string json_ensure_ascii_preserving_format(const std::string & json_str) { - std::string output; - output.reserve(json_str.size()); - - bool in_string = false; - bool escaped = false; - - for (size_t pos = 0; pos < json_str.size();) { - const char ch = json_str[pos]; - if (!in_string) { - output.push_back(ch); - if (ch == '"') { - in_string = true; - } - ++pos; - continue; - } - - if (escaped) { - output.push_back(ch); - escaped = false; - ++pos; - continue; - } - - if (ch == '\\') { - output.push_back(ch); - escaped = true; - ++pos; - continue; - } - - if (ch == '"') { - output.push_back(ch); - in_string = false; - ++pos; - continue; - } - - const unsigned char uch = static_cast(ch); - if (uch < 0x80) { - output.push_back(ch); - ++pos; - continue; - } - - auto parsed = common_parse_utf8_codepoint(json_str, pos); - if (parsed.status != utf8_parse_result::SUCCESS) { - output += "\\ufffd"; - ++pos; - continue; - } - - append_codepoint_as_ascii_json_escape(output, parsed.codepoint); - pos += parsed.bytes_consumed; - } - - return output; -} - -static value tojson(const func_args & args) { - args.ensure_count(1, 5); - value val_ascii = args.get_kwarg_or_pos("ensure_ascii", 1); - value val_indent = args.get_kwarg_or_pos("indent", 2); - value val_separators = args.get_kwarg_or_pos("separators", 3); - value val_sort = args.get_kwarg_or_pos("sort_keys", 4); - int indent = -1; - if (args.ctx.is_get_stats) { - // mark as used (recursively) for stats - auto val_input = args.get_pos(0); - value_t::stats_t::mark_used(const_cast(val_input), true); - } - if (is_val(val_indent)) { - indent = static_cast(val_indent->as_int()); - } - if (val_sort->as_bool()) { // undefined == false - throw not_implemented_exception("tojson sort_keys=true not implemented"); - } - const bool ensure_ascii = val_ascii->as_bool(); // undefined == false - auto separators = (is_val(val_separators) ? val_separators : mk_val())->as_array(); - std::string item_sep = separators.size() > 0 ? separators[0]->as_string().str() : (indent < 0 ? ", " : ","); - std::string key_sep = separators.size() > 1 ? separators[1]->as_string().str() : ": "; - std::string json_str = value_to_json(args.get_pos(0), indent, item_sep, key_sep); - if (ensure_ascii) { - json_str = json_ensure_ascii_preserving_format(json_str); - } - return mk_val(json_str); -} - -template -static value selectattr(const func_args & args) { - args.ensure_count(2, 4); - args.ensure_vals(true, true, false, false); - - auto arr = args.get_pos(0)->as_array(); - auto attribute = args.get_pos(1); - auto out = mk_val(); - value val_default = mk_val(); - - if (args.count() == 2) { - // example: array | selectattr("active") - for (const auto & item : arr) { - if (!is_val(item)) { - throw raised_exception("selectattr: item is not an object"); - } - value attr_val = item->at(attribute, val_default); - bool is_selected = attr_val->as_bool(); - if constexpr (is_reject) is_selected = !is_selected; - if (is_selected) out->push_back(item); - } - return out; - - } else if (args.count() == 3) { - // example: array | selectattr("equalto", "text") - // translated to: test_is_equalto(item, "text") - std::string test_name = args.get_pos(1)->as_string().str(); - value test_val = args.get_pos(2); - auto & builtins = global_builtins(); - auto it = builtins.find("test_is_" + test_name); - if (it == builtins.end()) { - throw raised_exception("selectattr: unknown test '" + test_name + "'"); - } - auto test_fn = it->second; - for (const auto & item : arr) { - func_args test_args(args.ctx); - test_args.push_back(item); // current object - test_args.push_back(test_val); // extra argument - value test_result = test_fn(test_args); - bool is_selected = test_result->as_bool(); - if constexpr (is_reject) is_selected = !is_selected; - if (is_selected) out->push_back(item); - } - return out; - - } else if (args.count() == 4) { - // example: array | selectattr("status", "equalto", "active") - // translated to: test_is_equalto(item.status, "active") - std::string test_name = args.get_pos(2)->as_string().str(); - auto extra_arg = args.get_pos(3); - auto & builtins = global_builtins(); - auto it = builtins.find("test_is_" + test_name); - if (it == builtins.end()) { - throw raised_exception("selectattr: unknown test '" + test_name + "'"); - } - auto test_fn = it->second; - for (const auto & item : arr) { - if (!is_val(item)) { - throw raised_exception("selectattr: item is not an object"); - } - value attr_val = item->at(attribute, val_default); - func_args test_args(args.ctx); - test_args.push_back(attr_val); // attribute value - test_args.push_back(extra_arg); // extra argument - value test_result = test_fn(test_args); - bool is_selected = test_result->as_bool(); - if constexpr (is_reject) is_selected = !is_selected; - if (is_selected) out->push_back(item); - } - return out; - } else { - throw raised_exception("selectattr: invalid number of arguments"); - } - - return out; -} - -static value default_value(const func_args & args) { - args.ensure_count(2, 3); - value val_check = args.get_kwarg_or_pos("boolean", 2); - bool check_bool = val_check->as_bool(); // undefined == false - bool no_value = check_bool - ? (!args.get_pos(0)->as_bool()) - : (args.get_pos(0)->is_undefined() || args.get_pos(0)->is_none()); - return no_value ? args.get_pos(1) : args.get_pos(0); -} - -const func_builtins & global_builtins() { - static const func_builtins builtins = { - {"raise_exception", [](const func_args & args) -> value { - args.ensure_vals(); - std::string msg = args.get_pos(0)->as_string().str(); - throw raised_exception("Jinja Exception: " + msg); - }}, - {"namespace", [](const func_args & args) -> value { - auto out = mk_val(); - for (const auto & arg : args.get_args()) { - if (!is_val(arg)) { - throw raised_exception("namespace() arguments must be kwargs"); - } - auto kwarg = cast_val(arg); - JJ_DEBUG("namespace: adding key '%s'", kwarg->key.c_str()); - out->insert(kwarg->key, kwarg->val); - } - return out; - }}, - {"strftime_now", [](const func_args & args) -> value { - args.ensure_vals(); - std::string format = args.get_pos(0)->as_string().str(); - // get current time - // TODO: make sure this is the same behavior as Python's strftime - char buf[100]; - if (std::strftime(buf, sizeof(buf), format.c_str(), std::localtime(&args.ctx.current_time))) { - return mk_val(std::string(buf)); - } else { - throw raised_exception("strftime_now: failed to format time"); - } - }}, - {"range", [](const func_args & args) -> value { - args.ensure_count(1, 3); - args.ensure_vals(true, false, false); - - auto arg0 = args.get_pos(0); - auto arg1 = args.get_pos(1, mk_val()); - auto arg2 = args.get_pos(2, mk_val()); - - int64_t start, stop, step; - if (args.count() == 1) { - start = 0; - stop = arg0->as_int(); - step = 1; - } else if (args.count() == 2) { - start = arg0->as_int(); - stop = arg1->as_int(); - step = 1; - } else { - start = arg0->as_int(); - stop = arg1->as_int(); - step = arg2->as_int(); - } - - auto out = mk_val(); - if (step == 0) { - throw raised_exception("range() step argument must not be zero"); - } - if (step > 0) { - for (int64_t i = start; i < stop; i += step) { - out->push_back(mk_val(i)); - } - } else { - for (int64_t i = start; i > stop; i += step) { - out->push_back(mk_val(i)); - } - } - return out; - }}, - {"tojson", tojson}, - - // tests - {"test_is_boolean", test_type_fn}, - {"test_is_callable", test_type_fn}, - {"test_is_odd", [](const func_args & args) -> value { - args.ensure_vals(); - int64_t val = args.get_pos(0)->as_int(); - return mk_val(val % 2 != 0); - }}, - {"test_is_even", [](const func_args & args) -> value { - args.ensure_vals(); - int64_t val = args.get_pos(0)->as_int(); - return mk_val(val % 2 == 0); - }}, - {"test_is_false", [](const func_args & args) -> value { - args.ensure_count(1); - bool val = is_val(args.get_pos(0)) && !args.get_pos(0)->as_bool(); - return mk_val(val); - }}, - {"test_is_true", [](const func_args & args) -> value { - args.ensure_count(1); - bool val = is_val(args.get_pos(0)) && args.get_pos(0)->as_bool(); - return mk_val(val); - }}, - {"test_is_divisibleby", [](const func_args & args) -> value { - args.ensure_vals(); - bool res = args.get_pos(0)->val_int % args.get_pos(1)->val_int == 0; - return mk_val(res); - }}, - {"test_is_string", test_type_fn}, - {"test_is_integer", test_type_fn}, - {"test_is_float", test_type_fn}, - {"test_is_number", test_type_fn}, - {"test_is_iterable", test_type_fn}, - {"test_is_sequence", test_type_fn}, - {"test_is_mapping", test_type_fn}, - {"test_is_lower", [](const func_args & args) -> value { - args.ensure_vals(); - return mk_val(args.get_pos(0)->val_str.is_lowercase()); - }}, - {"test_is_upper", [](const func_args & args) -> value { - args.ensure_vals(); - return mk_val(args.get_pos(0)->val_str.is_uppercase()); - }}, - {"test_is_none", test_type_fn}, - {"test_is_defined", [](const func_args & args) -> value { - args.ensure_count(1); - bool res = !args.get_pos(0)->is_undefined(); - JJ_DEBUG("test_is_defined: result=%d", res ? 1 : 0); - return mk_val(res); - }}, - {"test_is_undefined", test_type_fn}, - {"test_is_eq", test_compare_fn}, - {"test_is_equalto", test_compare_fn}, - {"test_is_ge", test_compare_fn}, - {"test_is_gt", test_compare_fn}, - {"test_is_greaterthan", test_compare_fn}, - {"test_is_lt", test_compare_fn}, - {"test_is_lessthan", test_compare_fn}, - {"test_is_ne", test_compare_fn}, - {"test_is_in", [](const func_args & args) -> value { - args.ensure_count(2); - auto needle = args.get_pos(0); - auto haystack = args.get_pos(1); - if (is_val(haystack)) { - return mk_val(false); - } - if (is_val(haystack)) { - for (const auto & item : haystack->as_array()) { - if (*needle == *item) { - return mk_val(true); - } - } - return mk_val(false); - } - if (is_val(haystack)) { - if (!is_val(needle)) { - throw raised_exception("'in' test expects args[1] as string when args[0] is string, got args[1] as " + needle->type()); - } - return mk_val( - haystack->as_string().str().find(needle->as_string().str()) != std::string::npos); - } - if (is_val(haystack)) { - return mk_val(haystack->has_key(needle)); - } - throw raised_exception("'in' test expects iterable as first argument, got " + haystack->type()); - }}, - {"test_is_test", [](const func_args & args) -> value { - args.ensure_vals(); - auto & builtins = global_builtins(); - std::string test_name = args.get_pos(0)->val_str.str(); - auto it = builtins.find("test_is_" + test_name); - bool res = it != builtins.end(); - return mk_val(res); - }}, - {"test_is_sameas", [](const func_args & args) -> value { - // Check if an object points to the same memory address as another object - (void)args; - throw not_implemented_exception("sameas test not implemented"); - }}, - {"test_is_escaped", [](const func_args & args) -> value { - (void)args; - throw not_implemented_exception("escaped test not implemented"); - }}, - {"test_is_filter", [](const func_args & args) -> value { - (void)args; - throw not_implemented_exception("filter test not implemented"); - }}, - }; - return builtins; -} - - -const func_builtins & value_int_t::get_builtins() const { - static const func_builtins builtins = { - {"default", default_value}, - {"abs", [](const func_args & args) -> value { - args.ensure_vals(); - int64_t val = args.get_pos(0)->as_int(); - return mk_val(val < 0 ? -val : val); - }}, - {"int", [](const func_args & args) -> value { - args.ensure_vals(); - return mk_val(args.get_pos(0)->as_int()); - }}, - {"float", [](const func_args & args) -> value { - args.ensure_vals(); - double val = static_cast(args.get_pos(0)->as_int()); - return mk_val(val); - }}, - {"safe", tojson}, - {"string", tojson}, - {"tojson", tojson}, - }; - return builtins; -} - - -const func_builtins & value_float_t::get_builtins() const { - static const func_builtins builtins = { - {"default", default_value}, - {"abs", [](const func_args & args) -> value { - args.ensure_vals(); - double val = args.get_pos(0)->as_float(); - return mk_val(val < 0.0 ? -val : val); - }}, - {"int", [](const func_args & args) -> value { - args.ensure_vals(); - int64_t val = static_cast(args.get_pos(0)->as_float()); - return mk_val(val); - }}, - {"float", [](const func_args & args) -> value { - args.ensure_vals(); - return mk_val(args.get_pos(0)->as_float()); - }}, - {"safe", tojson}, - {"string", tojson}, - {"tojson", tojson}, - }; - return builtins; -} - -static bool string_startswith(const std::string & str, const std::string & prefix) { - if (str.length() < prefix.length()) return false; - return str.compare(0, prefix.length(), prefix) == 0; -} - -static bool string_endswith(const std::string & str, const std::string & suffix) { - if (str.length() < suffix.length()) return false; - return str.compare(str.length() - suffix.length(), suffix.length(), suffix) == 0; -} - -[[noreturn]] static value string_join_not_implemented(const func_args &) { - throw not_implemented_exception("String join builtin not implemented"); -} - -const func_builtins & value_string_t::get_builtins() const { - static const func_builtins builtins = { - {"default", default_value}, - {"upper", [](const func_args & args) -> value { - args.ensure_vals(); - jinja::string str = args.get_pos(0)->as_string().uppercase(); - return mk_val(str); - }}, - {"lower", [](const func_args & args) -> value { - args.ensure_vals(); - jinja::string str = args.get_pos(0)->as_string().lowercase(); - return mk_val(str); - }}, - {"strip", [](const func_args & args) -> value { - value val_input = args.get_pos(0); - if (!is_val(val_input)) { - throw raised_exception("strip() first argument must be a string"); - } - value val_chars = args.get_kwarg_or_pos("chars", 1); - if (val_chars->is_undefined()) { - return mk_val(args.get_pos(0)->as_string().strip(true, true)); - } else { - return mk_val(args.get_pos(0)->as_string().strip(true, true, val_chars->as_string().str())); - } - }}, - {"rstrip", [](const func_args & args) -> value { - args.ensure_vals(); - value val_chars = args.get_kwarg_or_pos("chars", 1); - if (val_chars->is_undefined()) { - return mk_val(args.get_pos(0)->as_string().strip(false, true)); - } else { - return mk_val(args.get_pos(0)->as_string().strip(false, true, val_chars->as_string().str())); - } - }}, - {"lstrip", [](const func_args & args) -> value { - args.ensure_vals(); - value val_chars = args.get_kwarg_or_pos("chars", 1); - if (val_chars->is_undefined()) { - return mk_val(args.get_pos(0)->as_string().strip(true, false)); - } else { - return mk_val(args.get_pos(0)->as_string().strip(true, false, val_chars->as_string().str())); - } - }}, - {"title", [](const func_args & args) -> value { - args.ensure_vals(); - jinja::string str = args.get_pos(0)->as_string().titlecase(); - return mk_val(str); - }}, - {"capitalize", [](const func_args & args) -> value { - args.ensure_vals(); - jinja::string str = args.get_pos(0)->as_string().capitalize(); - return mk_val(str); - }}, - {"length", [](const func_args & args) -> value { - args.ensure_vals(); - jinja::string str = args.get_pos(0)->as_string(); - return mk_val(str.length()); - }}, - {"startswith", [](const func_args & args) -> value { - args.ensure_vals(); - std::string str = args.get_pos(0)->as_string().str(); - std::string prefix = args.get_pos(1)->as_string().str(); - return mk_val(string_startswith(str, prefix)); - }}, - {"endswith", [](const func_args & args) -> value { - args.ensure_vals(); - std::string str = args.get_pos(0)->as_string().str(); - std::string suffix = args.get_pos(1)->as_string().str(); - return mk_val(string_endswith(str, suffix)); - }}, - {"split", [](const func_args & args) -> value { - args.ensure_count(1, 3); - value val_input = args.get_pos(0); - if (!is_val(val_input)) { - throw raised_exception("split() first argument must be a string"); - } - std::string str = val_input->as_string().str(); - // FIXME: Support non-specified delimiter (split on consecutive (no leading or trailing) whitespace) - std::string delim = (args.count() > 1) ? args.get_pos(1)->as_string().str() : " "; - int64_t maxsplit = (args.count() > 2) ? args.get_pos(2)->as_int() : -1; - auto result = mk_val(); - size_t pos = 0; - std::string token; - while ((pos = str.find(delim)) != std::string::npos && maxsplit != 0) { - token = str.substr(0, pos); - result->push_back(mk_val(token)); - str.erase(0, pos + delim.length()); - --maxsplit; - } - auto res = mk_val(str); - res->val_str.mark_input_based_on(args.get_pos(0)->val_str); - result->push_back(std::move(res)); - return result; - }}, - {"rsplit", [](const func_args & args) -> value { - args.ensure_count(1, 3); - value val_input = args.get_pos(0); - if (!is_val(val_input)) { - throw raised_exception("rsplit() first argument must be a string"); - } - std::string str = val_input->as_string().str(); - // FIXME: Support non-specified delimiter (split on consecutive (no leading or trailing) whitespace) - std::string delim = (args.count() > 1) ? args.get_pos(1)->as_string().str() : " "; - int64_t maxsplit = (args.count() > 2) ? args.get_pos(2)->as_int() : -1; - auto result = mk_val(); - size_t pos = 0; - std::string token; - while ((pos = str.rfind(delim)) != std::string::npos && maxsplit != 0) { - token = str.substr(pos + delim.length()); - result->push_back(mk_val(token)); - str.erase(pos); - --maxsplit; - } - auto res = mk_val(str); - res->val_str.mark_input_based_on(args.get_pos(0)->val_str); - result->push_back(std::move(res)); - result->reverse(); - return result; - }}, - {"replace", [](const func_args & args) -> value { - args.ensure_vals(true, true, true, false); - std::string str = args.get_pos(0)->as_string().str(); - std::string old_str = args.get_pos(1)->as_string().str(); - std::string new_str = args.get_pos(2)->as_string().str(); - int64_t count = args.count() > 3 ? args.get_pos(3)->as_int() : -1; - if (count > 0) { - throw not_implemented_exception("String replace with count argument not implemented"); - } - size_t pos = 0; - while ((pos = str.find(old_str, pos)) != std::string::npos) { - str.replace(pos, old_str.length(), new_str); - pos += new_str.length(); - } - auto res = mk_val(str); - res->val_str.mark_input_based_on(args.get_pos(0)->val_str); - return res; - }}, - {"int", [](const func_args & args) -> value { - value val_input = args.get_pos(0); - value val_default = args.get_kwarg_or_pos("default", 1); - value val_base = args.get_kwarg_or_pos("base", 2); - const int base = val_base->is_undefined() ? 10 : val_base->as_int(); - if (is_val(val_input) == false) { - throw raised_exception("int() first argument must be a string"); - } - std::string str = val_input->as_string().str(); - try { - return mk_val(std::stoi(str, nullptr, base)); - } catch (...) { - return mk_val(val_default->is_undefined() ? 0 : val_default->as_int()); - } - }}, - {"float", [](const func_args & args) -> value { - args.ensure_vals(); - value val_default = args.get_kwarg_or_pos("default", 1); - std::string str = args.get_pos(0)->as_string().str(); - try { - return mk_val(std::stod(str)); - } catch (...) { - return mk_val(val_default->is_undefined() ? 0.0 : val_default->as_float()); - } - }}, - {"string", [](const func_args & args) -> value { - // no-op - args.ensure_vals(); - return mk_val(args.get_pos(0)->as_string()); - }}, - {"default", [](const func_args & args) -> value { - value input = args.get_pos(0); - if (!is_val(input)) { - throw raised_exception("default() first argument must be a string"); - } - value default_val = mk_val(""); - if (args.count() > 1 && !args.get_pos(1)->is_undefined()) { - default_val = args.get_pos(1); - } - value boolean_val = args.get_kwarg_or_pos("boolean", 2); // undefined == false - if (input->is_undefined() || (boolean_val->as_bool() && !input->as_bool())) { - return default_val; - } else { - return input; - } - }}, - {"slice", [](const func_args & args) -> value { - args.ensure_count(1, 4); - args.ensure_vals(true, true, false, false); - - auto arg0 = args.get_pos(1); - auto arg1 = args.get_pos(2, mk_val()); - auto arg2 = args.get_pos(3, mk_val()); - - int64_t start, stop, step; - if (args.count() == 1) { - start = 0; - stop = arg0->as_int(); - step = 1; - } else if (args.count() == 2) { - start = arg0->as_int(); - stop = arg1->as_int(); - step = 1; - } else { - start = arg0->as_int(); - stop = arg1->as_int(); - step = arg2->as_int(); - } - if (step == 0) { - throw raised_exception("slice step cannot be zero"); - } - auto input = args.get_pos(0); - auto sliced = slice(input->as_string().str(), start, stop, step); - auto res = mk_val(sliced); - res->val_str.mark_input_based_on(input->as_string()); - return res; - }}, - {"safe", [](const func_args & args) -> value { - // no-op for now - args.ensure_vals(); - return args.get_pos(0); - }}, - {"tojson", tojson}, - {"indent", [](const func_args &args) -> value { - args.ensure_count(1, 4); - value val_input = args.get_pos(0); - value val_width = args.get_kwarg_or_pos("width", 1); - const bool first = args.get_kwarg_or_pos("first", 2)->as_bool(); // undefined == false - const bool blank = args.get_kwarg_or_pos("blank", 3)->as_bool(); // undefined == false - if (!is_val(val_input)) { - throw raised_exception("indent() first argument must be a string"); - } - std::string indent; - if (is_val(val_width)) { - indent.assign(val_width->as_int(), ' '); - } else if (is_val(val_width)) { - indent = val_width->as_string().str(); - } else { - indent = " "; - } - std::string indented; - std::string input = val_input->as_string().str(); - std::istringstream iss = std::istringstream(input); - std::string line; - while (std::getline(iss, line)) { - if (!indented.empty()) { - indented.push_back('\n'); - } - if ((indented.empty() ? first : (!line.empty() || blank))) { - indented += indent; - } - indented += line; - } - if (!input.empty() && input.back() == '\n') { - indented.push_back('\n'); - if (blank) { - indented += indent; - } - } - - auto res = mk_val(indented); - res->val_str.mark_input_based_on(val_input->as_string()); - return res; - }}, - {"join", string_join_not_implemented}, - }; - return builtins; -} - - -const func_builtins & value_bool_t::get_builtins() const { - static const func_handler tostring = [](const func_args & args) -> value { - args.ensure_vals(); - bool val = args.get_pos(0)->as_bool(); - return mk_val(val ? "True" : "False"); - }; - static const func_builtins builtins = { - {"default", default_value}, - {"int", [](const func_args & args) -> value { - args.ensure_vals(); - bool val = args.get_pos(0)->as_bool(); - return mk_val(val ? 1 : 0); - }}, - {"float", [](const func_args & args) -> value { - args.ensure_vals(); - bool val = args.get_pos(0)->as_bool(); - return mk_val(val ? 1.0 : 0.0); - }}, - {"safe", tostring}, - {"string", tostring}, - {"tojson", tojson}, - }; - return builtins; -} - -[[noreturn]] static value array_unique_not_implemented(const func_args &) { - throw not_implemented_exception("Array unique builtin not implemented"); -} - -const func_builtins & value_array_t::get_builtins() const { - static const func_builtins builtins = { - {"default", default_value}, - {"list", [](const func_args & args) -> value { - args.ensure_vals(); - const auto & arr = args.get_pos(0)->as_array(); - auto result = mk_val(); - for (const auto& v : arr) { - result->push_back(v); - } - return result; - }}, - {"first", [](const func_args & args) -> value { - args.ensure_vals(); - const auto & arr = args.get_pos(0)->as_array(); - if (arr.empty()) { - return mk_val(); - } - return arr[0]; - }}, - {"last", [](const func_args & args) -> value { - args.ensure_vals(); - const auto & arr = args.get_pos(0)->as_array(); - if (arr.empty()) { - return mk_val(); - } - return arr[arr.size() - 1]; - }}, - {"length", [](const func_args & args) -> value { - args.ensure_vals(); - const auto & arr = args.get_pos(0)->as_array(); - return mk_val(static_cast(arr.size())); - }}, - {"slice", [](const func_args & args) -> value { - args.ensure_count(1, 4); - args.ensure_vals(true, true, false, false); - - auto val = args.get_pos(0); - auto arg0 = args.get_pos(1); - auto arg1 = args.get_pos(2, mk_val()); - auto arg2 = args.get_pos(3, mk_val()); - - int64_t start, stop, step; - if (args.count() == 1) { - start = 0; - stop = arg0->as_int(); - step = 1; - } else if (args.count() == 2) { - start = arg0->as_int(); - stop = arg1->as_int(); - step = 1; - } else { - start = arg0->as_int(); - stop = arg1->as_int(); - step = arg2->as_int(); - } - if (step == 0) { - throw raised_exception("slice step cannot be zero"); - } - auto arr = slice(val->as_array(), start, stop, step); - return is_val(val) ? mk_val(std::move(arr)) : mk_val(std::move(arr)); - }}, - {"selectattr", selectattr}, - {"select", selectattr}, - {"rejectattr", selectattr}, - {"reject", selectattr}, - {"join", [](const func_args & args) -> value { - args.ensure_count(1, 3); - if (!is_val(args.get_pos(0))) { - throw raised_exception("join() first argument must be an array"); - } - value val_delim = args.get_kwarg_or_pos("d", 1); - value attribute = args.get_kwarg_or_pos("attribute", 2); - const auto & arr = args.get_pos(0)->as_array(); - const bool attr_is_int = is_val(attribute); - if (!attribute->is_undefined() && !is_val(attribute) && !attr_is_int) { - throw raised_exception("join() attribute must be string or integer"); - } - const int64_t attr_int = attr_is_int ? attribute->as_int() : 0; - const std::string delim = val_delim->is_undefined() ? "" : val_delim->as_string().str(); - std::string result; - for (size_t i = 0; i < arr.size(); ++i) { - value val_arr = arr[i]; - if (!attribute->is_undefined()) { - if (attr_is_int && is_val(val_arr)) { - val_arr = val_arr->at(attr_int); - } else if (!attr_is_int && is_val(val_arr)) { - val_arr = val_arr->at(attribute); - } - } - if (!is_val(val_arr) && !is_val(val_arr) && !is_val(val_arr)) { - throw raised_exception("join() can only join arrays of strings or numerics"); - } - result += val_arr->as_string().str(); - if (i < arr.size() - 1) { - result += delim; - } - } - return mk_val(result); - }}, - {"string", [](const func_args & args) -> value { - args.ensure_vals(); - if (args.ctx.is_get_stats) { - // mark as used (recursively) for stats - auto val_input = args.get_pos(0); - value_t::stats_t::mark_used(const_cast(val_input), true); - } - return mk_val(args.get_pos(0)->as_string()); - }}, - {"tojson", tojson}, - {"map", [](const func_args & args) -> value { - args.ensure_count(2); - if (!is_val(args.get_pos(0))) { - throw raised_exception("map: first argument must be an array"); - } - if (!is_val(args.get_args().at(1))) { - throw not_implemented_exception("map: filter-mapping not implemented"); - } - value val = args.get_pos(0); - value attribute = args.get_kwarg_or_pos("attribute", 1); - const bool attr_is_int = is_val(attribute); - if (!is_val(attribute) && !attr_is_int) { - throw raised_exception("map: attribute must be string or integer"); - } - const int64_t attr_int = attr_is_int ? attribute->as_int() : 0; - value default_val = args.get_kwarg("default", mk_val()); - auto out = mk_val(); - auto arr = val->as_array(); - for (const auto & item : arr) { - value attr_val; - if (attr_is_int) { - attr_val = is_val(item) ? item->at(attr_int, default_val) : default_val; - } else { - attr_val = is_val(item) ? item->at(attribute, default_val) : default_val; - } - out->push_back(attr_val); - } - return is_val(val) ? mk_val(std::move(out->as_array())) : out; - }}, - {"append", [](const func_args & args) -> value { - args.ensure_count(2); - if (!is_val(args.get_pos(0))) { - throw raised_exception("append: first argument must be an array"); - } - const value_array_t * arr = cast_val(args.get_pos(0)); - // need to use const_cast here to modify the array - value_array_t * arr_editable = const_cast(arr); - arr_editable->push_back(args.get_pos(1)); - return args.get_pos(0); - }}, - {"pop", [](const func_args & args) -> value { - args.ensure_count(1, 2); - args.ensure_vals(true, false); - int64_t index = args.count() == 2 ? args.get_pos(1)->as_int() : -1; - const value_array_t * arr = cast_val(args.get_pos(0)); - // need to use const_cast here to modify the array - value_array_t * arr_editable = const_cast(arr); - return arr_editable->pop_at(index); - }}, - {"sort", [](const func_args & args) -> value { - args.ensure_count(1, 4); - if (!is_val(args.get_pos(0))) { - throw raised_exception("sort: first argument must be an array"); - } - value val = args.get_pos(0); - value val_reverse = args.get_kwarg_or_pos("reverse", 1); - value val_case = args.get_kwarg_or_pos("case_sensitive", 2); - value attribute = args.get_kwarg_or_pos("attribute", 3); - // FIXME: sorting is currently always case sensitive - //const bool case_sensitive = val_case->as_bool(); // undefined == false - const bool reverse = val_reverse->as_bool(); // undefined == false - const bool attr_is_int = is_val(attribute); - const int64_t attr_int = attr_is_int ? attribute->as_int() : 0; - std::vector arr = val->as_array(); // copy - std::sort(arr.begin(), arr.end(),[&](const value & a, const value & b) { - value val_a = a; - value val_b = b; - if (!attribute->is_undefined()) { - if (attr_is_int && is_val(a) && is_val(b)) { - val_a = a->at(attr_int); - val_b = b->at(attr_int); - } else if (!attr_is_int && is_val(a) && is_val(b)) { - val_a = a->at(attribute); - val_b = b->at(attribute); - } else { - throw raised_exception("sort: unsupported object attribute comparison between " + a->type() + " and " + b->type()); - } - } - return value_compare(val_a, val_b, reverse ? value_compare_op::gt : value_compare_op::lt); - }); - return is_val(val) ? mk_val(std::move(arr)) : mk_val(std::move(arr)); - }}, - {"reverse", [](const func_args & args) -> value { - args.ensure_vals(); - value val = args.get_pos(0); - std::vector arr = val->as_array(); // copy - std::reverse(arr.begin(), arr.end()); - return is_val(val) ? mk_val(std::move(arr)) : mk_val(std::move(arr)); - }}, - {"unique", array_unique_not_implemented}, - }; - return builtins; -} - -[[noreturn]] static value object_join_not_implemented(const func_args &) { - throw not_implemented_exception("object join not implemented"); -} - -const func_builtins & value_object_t::get_builtins() const { - if (!has_builtins) { - static const func_builtins no_builtins = {}; - return no_builtins; - } - - static const func_builtins builtins = { - // {"default", default_value}, // cause issue with gpt-oss - {"get", [](const func_args & args) -> value { - args.ensure_count(2, 3); - if (!is_val(args.get_pos(0))) { - throw raised_exception("get: first argument must be an object"); - } - if (!is_val(args.get_pos(1))) { - throw raised_exception("get: second argument must be a string (key)"); - } - value default_val = mk_val(); - if (args.count() == 3) { - default_val = args.get_pos(2); - } - const value obj = args.get_pos(0); - const value key = args.get_pos(1); - return obj->at(key, default_val); - }}, - {"keys", [](const func_args & args) -> value { - args.ensure_vals(); - const auto & obj = args.get_pos(0)->as_ordered_object(); - auto result = mk_val(); - for (const auto & pair : obj) { - result->push_back(pair.first); - } - return result; - }}, - {"values", [](const func_args & args) -> value { - args.ensure_vals(); - const auto & obj = args.get_pos(0)->as_ordered_object(); - auto result = mk_val(); - for (const auto & pair : obj) { - result->push_back(pair.second); - } - return result; - }}, - {"items", [](const func_args & args) -> value { - args.ensure_vals(); - const auto & obj = args.get_pos(0)->as_ordered_object(); - auto result = mk_val(); - for (const auto & pair : obj) { - auto item = mk_val(pair); - result->push_back(std::move(item)); - } - return result; - }}, - {"tojson", tojson}, - {"string", [](const func_args & args) -> value { - args.ensure_vals(); - if (args.ctx.is_get_stats) { - // mark as used (recursively) for stats - auto val_input = args.get_pos(0); - value_t::stats_t::mark_used(const_cast(val_input), true); - } - return mk_val(args.get_pos(0)->as_string()); - }}, - {"length", [](const func_args & args) -> value { - args.ensure_vals(); - const auto & obj = args.get_pos(0)->as_ordered_object(); - return mk_val(static_cast(obj.size())); - }}, - {"tojson", [](const func_args & args) -> value { - args.ensure_vals(); - // use global to_json - return global_builtins().at("tojson")(args); - }}, - {"dictsort", [](const func_args & args) -> value { - value val_input = args.get_pos(0); - value val_case = args.get_kwarg_or_pos("case_sensitive", 1); - value val_by = args.get_kwarg_or_pos("by", 2); - value val_reverse = args.get_kwarg_or_pos("reverse", 3); - // FIXME: sorting is currently always case sensitive - //const bool case_sensitive = val_case->as_bool(); // undefined == false - const bool reverse = val_reverse->as_bool(); // undefined == false - const bool by_value = is_val(val_by) && val_by->as_string().str() == "value" ? true : false; - auto result = mk_val(val_input); // copy - std::sort(result->val_obj.begin(), result->val_obj.end(), [&](const auto & a, const auto & b) { - if (by_value) { - return value_compare(a.second, b.second, reverse ? value_compare_op::gt : value_compare_op::lt); - } else { - return value_compare(a.first, b.first, reverse ? value_compare_op::gt : value_compare_op::lt); - } - }); - return result; - }}, - {"join", object_join_not_implemented}, - }; - return builtins; -} - -const func_builtins & value_none_t::get_builtins() const { - static const func_handler tostring = [](const func_args &) -> value { - return mk_val("None"); - }; - static const func_builtins builtins = { - {"default", default_value}, - {"tojson", tojson}, - {"string", tostring}, - {"safe", tostring}, - {"items", empty_value_fn}, - {"map", empty_value_fn}, - {"reject", empty_value_fn}, - {"rejectattr", empty_value_fn}, - {"select", empty_value_fn}, - {"selectattr", empty_value_fn}, - {"unique", empty_value_fn}, - }; - return builtins; -} - - -const func_builtins & value_undefined_t::get_builtins() const { - static const func_builtins builtins = { - {"default", default_value}, - {"capitalize", empty_value_fn}, - {"first", empty_value_fn}, - {"items", empty_value_fn}, - {"join", empty_value_fn}, - {"last", empty_value_fn}, - {"length", empty_value_fn}, - {"list", empty_value_fn}, - {"lower", empty_value_fn}, - {"map", empty_value_fn}, - {"max", empty_value_fn}, - {"min", empty_value_fn}, - {"reject", empty_value_fn}, - {"rejectattr", empty_value_fn}, - {"replace", empty_value_fn}, - {"reverse", empty_value_fn}, - {"safe", empty_value_fn}, - {"select", empty_value_fn}, - {"selectattr", empty_value_fn}, - {"sort", empty_value_fn}, - {"string", empty_value_fn}, - {"strip", empty_value_fn}, - {"sum", empty_value_fn}, - {"title", empty_value_fn}, - {"truncate", empty_value_fn}, - {"unique", empty_value_fn}, - {"upper", empty_value_fn}, - {"wordcount", empty_value_fn}, - }; - return builtins; -} - - -////////////////////////////////// - - -static value from_json(const nlohmann::ordered_json & j, bool mark_input) { - if (j.is_null()) { - return mk_val(); - } else if (j.is_boolean()) { - return mk_val(j.get()); - } else if (j.is_number_integer()) { - return mk_val(j.get()); - } else if (j.is_number_float()) { - return mk_val(j.get()); - } else if (j.is_string()) { - auto str = mk_val(j.get()); - if (mark_input) { - str->mark_input(); - } - return str; - } else if (j.is_array()) { - auto arr = mk_val(); - for (const auto & item : j) { - arr->push_back(from_json(item, mark_input)); - } - return arr; - } else if (j.is_object()) { - auto obj = mk_val(); - for (auto it = j.begin(); it != j.end(); ++it) { - obj->insert(it.key(), from_json(it.value(), mark_input)); - } - return obj; - } else { - throw std::runtime_error("Unsupported JSON value type"); - } -} - -// compare operator for value_t -bool value_compare(const value & a, const value & b, value_compare_op op) { - auto cmp = [&]() { - // compare numeric types - if ((is_val(a) || is_val(a)) && - (is_val(b) || is_val(b))){ - try { - if (op == value_compare_op::eq) { - return a->as_float() == b->as_float(); - } else if (op == value_compare_op::ge) { - return a->as_float() >= b->as_float(); - } else if (op == value_compare_op::gt) { - return a->as_float() > b->as_float(); - } else if (op == value_compare_op::lt) { - return a->as_float() < b->as_float(); - } else if (op == value_compare_op::ne) { - return a->as_float() != b->as_float(); - } else { - throw std::runtime_error("Unsupported comparison operator for numeric types"); - } - } catch (...) {} - } - // compare string and number - // TODO: not sure if this is the right behavior - if ((is_val(b) && (is_val(a) || is_val(a))) || - (is_val(a) && (is_val(b) || is_val(b))) || - (is_val(a) && is_val(b))) { - try { - if (op == value_compare_op::eq) { - return a->as_string().str() == b->as_string().str(); - } else if (op == value_compare_op::ge) { - return a->as_string().str() >= b->as_string().str(); - } else if (op == value_compare_op::gt) { - return a->as_string().str() > b->as_string().str(); - } else if (op == value_compare_op::lt) { - return a->as_string().str() < b->as_string().str(); - } else if (op == value_compare_op::ne) { - return a->as_string().str() != b->as_string().str(); - } else { - throw std::runtime_error("Unsupported comparison operator for string/number types"); - } - } catch (...) {} - } - // compare boolean simple - if (is_val(a) && is_val(b)) { - if (op == value_compare_op::eq) { - return a->as_bool() == b->as_bool(); - } else if (op == value_compare_op::ne) { - return a->as_bool() != b->as_bool(); - } else { - throw std::runtime_error("Unsupported comparison operator for bool type"); - } - } - // compare by type - if (a->type() != b->type()) { - return false; - } - return false; - }; - auto result = cmp(); - JJ_DEBUG("Comparing types: %s and %s result=%d", a->type().c_str(), b->type().c_str(), result); - return result; -} - -template<> -void global_from_json(context & ctx, const nlohmann::ordered_json & json_obj, bool mark_input) { - // printf("global_from_json: %s\n" , json_obj.dump(2).c_str()); - if (json_obj.is_null() || !json_obj.is_object()) { - throw std::runtime_error("global_from_json: input JSON value must be an object"); - } - for (auto it = json_obj.begin(); it != json_obj.end(); ++it) { - JJ_DEBUG("global_from_json: setting key '%s'", it.key().c_str()); - ctx.set_val(it.key(), from_json(it.value(), mark_input)); - } -} - -// recursively convert value to JSON string -// TODO: avoid circular references -static void value_to_json_internal(std::ostringstream & oss, const value & val, int curr_lvl, int indent, const std::string_view item_sep, const std::string_view key_sep) { - auto indent_str = [indent, curr_lvl]() -> std::string { - return (indent > 0) ? std::string(curr_lvl * indent, ' ') : ""; - }; - auto newline = [indent]() -> std::string { - return (indent >= 0) ? "\n" : ""; - }; - - if (is_val(val) || val->is_undefined()) { - oss << "null"; - } else if (is_val(val)) { - oss << (val->as_bool() ? "true" : "false"); - } else if (is_val(val)) { - oss << val->as_int(); - } else if (is_val(val)) { - oss << val->as_float(); - } else if (is_val(val)) { - oss << "\""; - for (char c : val->as_string().str()) { - switch (c) { - case '"': oss << "\\\""; break; - case '\\': oss << "\\\\"; break; - case '\b': oss << "\\b"; break; - case '\f': oss << "\\f"; break; - case '\n': oss << "\\n"; break; - case '\r': oss << "\\r"; break; - case '\t': oss << "\\t"; break; - default: - if (static_cast(c) < 0x20) { - char buf[7]; - snprintf(buf, sizeof(buf), "\\u%04x", static_cast(c)); - oss << buf; - } else { - oss << c; - } - } - } - oss << "\""; - } else if (is_val(val)) { - const auto & arr = val->as_array(); - oss << "["; - if (!arr.empty()) { - oss << newline(); - for (size_t i = 0; i < arr.size(); ++i) { - oss << indent_str() << (indent > 0 ? std::string(indent, ' ') : ""); - value_to_json_internal(oss, arr[i], curr_lvl + 1, indent, item_sep, key_sep); - if (i < arr.size() - 1) { - oss << item_sep; - } - oss << newline(); - } - oss << indent_str(); - } - oss << "]"; - } else if (is_val(val)) { - const auto & obj = val->as_ordered_object(); // IMPORTANT: need to keep exact order - oss << "{"; - if (!obj.empty()) { - oss << newline(); - size_t i = 0; - for (const auto & pair : obj) { - oss << indent_str() << (indent > 0 ? std::string(indent, ' ') : ""); - value_to_json_internal(oss, mk_val(pair.first->as_string().str()), curr_lvl + 1, indent, item_sep, key_sep); - oss << key_sep; - value_to_json_internal(oss, pair.second, curr_lvl + 1, indent, item_sep, key_sep); - if (i < obj.size() - 1) { - oss << item_sep; - } - oss << newline(); - ++i; - } - oss << indent_str(); - } - oss << "}"; - } else { - oss << "null"; - } -} - -std::string value_to_json(const value & val, int indent, const std::string_view item_sep, const std::string_view key_sep) { - std::ostringstream oss; - value_to_json_internal(oss, val, 0, indent, item_sep, key_sep); - JJ_DEBUG("value_to_json: result=%s", oss.str().c_str()); - return oss.str(); -} - -// TODO: avoid circular references -std::string value_to_string_repr(const value & val) { - if (is_val(val)) { - const std::string val_str = val->as_string().str(); - - if (val_str.find('\'') != std::string::npos) { - return value_to_json(val); - } else { - return "'" + val_str + "'"; - } - } else { - return val->as_repr(); - } -} - -// stats utility -void value_t::stats_t::mark_used(value & val, bool deep) { - val->stats.used = true; - if (deep) { - if (is_val(val)) { - for (auto & item : val->val_arr) { - mark_used(item, deep); - } - } else if (is_val(val)) { - for (auto & pair : val->val_obj) { - mark_used(pair.first, deep); - mark_used(pair.second, deep); - } - } - } -} - -} // namespace jinja diff --git a/common/jinja/value.h b/common/jinja/value.h deleted file mode 100644 index 5cf85e4f5443..000000000000 --- a/common/jinja/value.h +++ /dev/null @@ -1,759 +0,0 @@ -#pragma once - -#include "string.h" -#include "utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace jinja { - -struct value_t; -using value = std::shared_ptr; - - -// Helper to check the type of a value -template -struct extract_pointee { - using type = T; -}; -template -struct extract_pointee> { - using type = U; -}; -template -bool is_val(const value & ptr) { - using PointeeType = typename extract_pointee::type; - return dynamic_cast(ptr.get()) != nullptr; -} -template -bool is_val(const value_t * ptr) { - using PointeeType = typename extract_pointee::type; - return dynamic_cast(ptr) != nullptr; -} -template -std::shared_ptr::type> mk_val(Args&&... args) { - using PointeeType = typename extract_pointee::type; - return std::make_shared(std::forward(args)...); -} -template -const typename extract_pointee::type * cast_val(const value & ptr) { - using PointeeType = typename extract_pointee::type; - return dynamic_cast(ptr.get()); -} -template -typename extract_pointee::type * cast_val(value & ptr) { - using PointeeType = typename extract_pointee::type; - return dynamic_cast(ptr.get()); -} -// End Helper - - -struct context; // forward declaration - - -// for converting from JSON to jinja values -// example input JSON: -// { -// "messages": [ -// {"role": "user", "content": "Hello!"}, -// {"role": "assistant", "content": "Hi there!"} -// ], -// "bos_token": "", -// "eos_token": "", -// } -// -// to mark strings as user input, wrap them in a special object: -// { -// "messages": [ -// { -// "role": "user", -// "content": {"__input__": "Hello!"} // this string is user input -// }, -// ... -// ], -// } -// -// marking input can be useful for tracking data provenance -// and preventing template injection attacks -// -// Note: T_JSON can be nlohmann::ordered_json -template -void global_from_json(context & ctx, const T_JSON & json_obj, bool mark_input); - -// -// base value type -// - -struct func_args; // function argument values - -using func_hptr = value(const func_args &); -using func_handler = std::function; -using func_builtins = std::map; - -enum value_compare_op { eq, ge, gt, lt, ne }; -bool value_compare(const value & a, const value & b, value_compare_op op); - -struct value_t { - int64_t val_int; - double val_flt; - string val_str; - - std::vector val_arr; - std::vector> val_obj; - - func_handler val_func; - - // only used if ctx.is_get_stats = true - struct stats_t { - bool used = false; - // ops can be builtin calls or operators: "array_access", "object_access" - std::set ops; - // utility to recursively mark value and its children as used - static void mark_used(value & val, bool deep = false); - } stats; - - value_t() = default; - value_t(const value_t &) = default; - virtual ~value_t() = default; - - // Note: only for debugging and error reporting purposes - virtual std::string type() const { return ""; } - - virtual int64_t as_int() const { throw_type_error("is not an int value"); } - virtual double as_float() const { throw_type_error("is not a float value"); } - virtual string as_string() const { throw_type_error("is not a string value"); } - virtual bool as_bool() const { throw_type_error("is not a bool value"); } - virtual const std::vector & as_array() const { throw_type_error("is not an array value"); } - virtual const std::vector> & as_ordered_object() const { throw_type_error("is not an object value"); } - virtual value invoke(const func_args &) const { throw_type_error("is not a function value"); } - virtual bool is_none() const { return false; } - virtual bool is_undefined() const { return false; } - virtual const func_builtins & get_builtins() const { throw_type_error("has no builtins"); } - - virtual bool has_key(const value &) { throw_type_error("is not an object value"); } - virtual void insert(const value & /* key */, const value & /* val */) { throw_type_error("is not an object value"); } - virtual value & at(const value & /* key */, value & /* default_val */) { throw_type_error("is not an object value"); } - virtual value & at(const value & /* key */) { throw_type_error("is not an object value"); } - virtual value & at(const std::string & /* key */, value & /* default_val */) { throw_type_error("is not an object value"); } - virtual value & at(const std::string & /* key */) { throw_type_error("is not an object value"); } - virtual value & at(int64_t /* idx */, value & /* default_val */) { throw_type_error("is not an array value"); } - virtual value & at(int64_t /* idx */) { throw_type_error("is not an array value"); } - - virtual bool is_numeric() const { return false; } - virtual bool is_hashable() const { return false; } - virtual bool is_immutable() const { return true; } - virtual hasher unique_hash() const noexcept = 0; - // TODO: C++20 <=> operator - // NOTE: We are treating == as equivalent (for normal comparisons) and != as strict nonequal (for strict (is) comparisons) - virtual bool operator==(const value_t & other) const { return equivalent(other); } - virtual bool operator!=(const value_t & other) const { return nonequal(other); } - - // Note: only for debugging purposes - virtual std::string as_repr() const { return as_string().str(); } - -private: - [[noreturn]] void throw_type_error(const char* expected) const { - throw std::runtime_error(type() + " " + expected); - } - -protected: - virtual bool equivalent(const value_t &) const = 0; - virtual bool nonequal(const value_t & other) const { return !equivalent(other); } -}; - -// -// utils -// - -const func_builtins & global_builtins(); - -std::string value_to_json(const value & val, int indent = -1, const std::string_view item_sep = ", ", const std::string_view key_sep = ": "); - -// Note: only used for debugging purposes -std::string value_to_string_repr(const value & val); - -struct not_implemented_exception : public std::runtime_error { - not_implemented_exception(const std::string & msg) : std::runtime_error("NotImplemented: " + msg) {} -}; - -struct value_hasher { - size_t operator()(const value & val) const noexcept { - return val->unique_hash().digest(); - } -}; - -struct value_equivalence { - bool operator()(const value & lhs, const value & rhs) const { - return *lhs == *rhs; - } - bool operator()(const std::pair & lhs, const std::pair & rhs) const { - return *(lhs.first) == *(rhs.first) && *(lhs.second) == *(rhs.second); - } -}; - -struct value_equality { - bool operator()(const value & lhs, const value & rhs) const { - return !(*lhs != *rhs); - } -}; - -// -// primitive value types -// - -struct value_int_t : public value_t { - value_int_t(int64_t v) { - val_int = v; - val_flt = static_cast(v); - if (static_cast(val_flt) != v) { - val_flt = v < 0 ? -INFINITY : INFINITY; - } - } - virtual std::string type() const override { return "Integer"; } - virtual int64_t as_int() const override { return val_int; } - virtual double as_float() const override { return val_flt; } - virtual string as_string() const override { return std::to_string(val_int); } - virtual bool as_bool() const override { - return val_int != 0; - } - virtual const func_builtins & get_builtins() const override; - virtual bool is_numeric() const override { return true; } - virtual bool is_hashable() const override { return true; } - virtual hasher unique_hash() const noexcept override { - return hasher(typeid(*this)) - .update(&val_int, sizeof(val_int)) - .update(&val_flt, sizeof(val_flt)); - } -protected: - virtual bool equivalent(const value_t & other) const override { - return other.is_numeric() && val_int == other.val_int && val_flt == other.val_flt; - } - virtual bool nonequal(const value_t & other) const override { - return !(typeid(*this) == typeid(other) && val_int == other.val_int); - } -}; -using value_int = std::shared_ptr; - - -struct value_float_t : public value_t { - value val; - value_float_t(double v) { - val_flt = v; - val_int = std::isfinite(v) ? static_cast(v) : 0; - val = mk_val(val_int); - } - virtual std::string type() const override { return "Float"; } - virtual double as_float() const override { return val_flt; } - virtual int64_t as_int() const override { return val_int; } - virtual string as_string() const override { - std::string out = std::to_string(val_flt); - out.erase(out.find_last_not_of('0') + 1, std::string::npos); // remove trailing zeros - if (out.back() == '.') out.push_back('0'); // leave one zero if no decimals - return out; - } - virtual bool as_bool() const override { - return val_flt != 0.0; - } - virtual const func_builtins & get_builtins() const override; - virtual bool is_numeric() const override { return true; } - virtual bool is_hashable() const override { return true; } - virtual hasher unique_hash() const noexcept override { - if (static_cast(val_int) == val_flt) { - return val->unique_hash(); - } else { - return hasher(typeid(*this)) - .update(&val_int, sizeof(val_int)) - .update(&val_flt, sizeof(val_flt)); - } - } -protected: - virtual bool equivalent(const value_t & other) const override { - return other.is_numeric() && val_int == other.val_int && val_flt == other.val_flt; - } - virtual bool nonequal(const value_t & other) const override { - return !(typeid(*this) == typeid(other) && val_flt == other.val_flt); - } -}; -using value_float = std::shared_ptr; - - -struct value_string_t : public value_t { - value_string_t() { val_str = string(); } - value_string_t(const std::string & v) { val_str = string(v); } - value_string_t(const string & v) { val_str = v; } - virtual std::string type() const override { return "String"; } - virtual string as_string() const override { return val_str; } - virtual std::string as_repr() const override { - std::ostringstream ss; - for (const auto & part : val_str.parts) { - ss << (part.is_input ? "INPUT: " : "TMPL: ") << part.val << "\n"; - } - return ss.str(); - } - virtual bool as_bool() const override { - return val_str.length() > 0; - } - virtual const func_builtins & get_builtins() const override; - virtual bool is_hashable() const override { return true; } - virtual hasher unique_hash() const noexcept override { - const auto type_hash = typeid(*this).hash_code(); - auto hash = hasher(); - hash.update(&type_hash, sizeof(type_hash)); - val_str.hash_update(hash); - return hash; - } - void mark_input() { - val_str.mark_input(); - } -protected: - virtual bool equivalent(const value_t & other) const override { - return typeid(*this) == typeid(other) && val_str.str() == other.val_str.str(); - } -}; -using value_string = std::shared_ptr; - - -struct value_bool_t : public value_t { - value val; - value_bool_t(bool v) { - val_int = static_cast(v); - val_flt = static_cast(v); - val = mk_val(val_int); - } - virtual std::string type() const override { return "Boolean"; } - virtual int64_t as_int() const override { return val_int; } - virtual bool as_bool() const override { return val_int; } - virtual string as_string() const override { return std::string(val_int ? "True" : "False"); } - virtual const func_builtins & get_builtins() const override; - virtual bool is_numeric() const override { return true; } - virtual bool is_hashable() const override { return true; } - virtual hasher unique_hash() const noexcept override { - return val->unique_hash(); - } -protected: - virtual bool equivalent(const value_t & other) const override { - return other.is_numeric() && val_int == other.val_int && val_flt == other.val_flt; - } - virtual bool nonequal(const value_t & other) const override { - return !(typeid(*this) == typeid(other) && val_int == other.val_int); - } -}; -using value_bool = std::shared_ptr; - - -struct value_array_t : public value_t { - value_array_t() = default; - value_array_t(value & v) { - val_arr = v->val_arr; - } - value_array_t(std::vector && arr) { - val_arr = arr; - } - value_array_t(const std::vector & arr) { - val_arr = arr; - } - void reverse() { - if (is_immutable()) { - throw std::runtime_error("Attempting to modify immutable type"); - } - std::reverse(val_arr.begin(), val_arr.end()); - } - void push_back(const value & val) { - if (is_immutable()) { - throw std::runtime_error("Attempting to modify immutable type"); - } - val_arr.push_back(val); - } - void push_back(value && val) { - if (is_immutable()) { - throw std::runtime_error("Attempting to modify immutable type"); - } - val_arr.push_back(std::move(val)); - } - value pop_at(int64_t index) { - if (is_immutable()) { - throw std::runtime_error("Attempting to modify immutable type"); - } - if (index < 0) { - index = static_cast(val_arr.size()) + index; - } - if (index < 0 || index >= static_cast(val_arr.size())) { - throw std::runtime_error("Index " + std::to_string(index) + " out of bounds for array of size " + std::to_string(val_arr.size())); - } - value val = val_arr.at(static_cast(index)); - val_arr.erase(val_arr.begin() + index); - return val; - } - virtual std::string type() const override { return "Array"; } - virtual bool is_immutable() const override { return false; } - virtual const std::vector & as_array() const override { return val_arr; } - virtual string as_string() const override { - const bool immutable = is_immutable(); - std::ostringstream ss; - ss << (immutable ? "(" : "["); - for (size_t i = 0; i < val_arr.size(); i++) { - if (i > 0) ss << ", "; - value val = val_arr.at(i); - ss << value_to_string_repr(val); - } - if (immutable && val_arr.size() == 1) { - ss << ","; - } - ss << (immutable ? ")" : "]"); - return ss.str(); - } - virtual bool as_bool() const override { - return !val_arr.empty(); - } - virtual value & at(int64_t index, value & default_val) override { - if (index < 0) { - index += val_arr.size(); - } - if (index < 0 || static_cast(index) >= val_arr.size()) { - return default_val; - } - return val_arr[index]; - } - virtual value & at(int64_t index) override { - if (index < 0) { - index += val_arr.size(); - } - if (index < 0 || static_cast(index) >= val_arr.size()) { - throw std::runtime_error("Index " + std::to_string(index) + " out of bounds for array of size " + std::to_string(val_arr.size())); - } - return val_arr[index]; - } - virtual const func_builtins & get_builtins() const override; - virtual bool is_hashable() const override { - if (std::all_of(val_arr.begin(), val_arr.end(), [&](auto & val) -> bool { - return val->is_immutable() && val->is_hashable(); - })) { - return true; - } - return false; - } - virtual hasher unique_hash() const noexcept override { - auto hash = hasher(typeid(*this)); - for (const auto & val : val_arr) { - // must use digest to prevent problems from "concatenation" property of hasher - // for ex. hash of [ "ab", "c" ] should be different from [ "a", "bc" ] - const size_t val_hash = val->unique_hash().digest(); - hash.update(&val_hash, sizeof(size_t)); - } - return hash; - } -protected: - virtual bool equivalent(const value_t & other) const override { - return typeid(*this) == typeid(other) && is_hashable() && other.is_hashable() && std::equal(val_arr.begin(), val_arr.end(), other.val_arr.begin(), other.val_arr.end(), value_equivalence()); - } -}; -using value_array = std::shared_ptr; - - -struct value_tuple_t : public value_array_t { - value_tuple_t(value & v) { - val_arr = v->val_arr; - } - value_tuple_t(std::vector && arr) { - val_arr = arr; - } - value_tuple_t(const std::vector & arr) { - val_arr = arr; - } - value_tuple_t(const std::pair & pair) { - val_arr.push_back(pair.first); - val_arr.push_back(pair.second); - } - virtual std::string type() const override { return "Tuple"; } - virtual bool is_immutable() const override { return true; } -}; -using value_tuple = std::shared_ptr; - - -struct value_object_t : public value_t { - std::unordered_map unordered; - bool has_builtins = true; // context and loop objects do not have builtins - value_object_t() = default; - value_object_t(value & v) { - val_obj = v->val_obj; - for (const auto & pair : val_obj) { - unordered[pair.first] = pair.second; - } - } - value_object_t(const std::map & obj) { - for (const auto & pair : obj) { - insert(pair.first, pair.second); - } - } - value_object_t(const std::vector> & obj) { - for (const auto & pair : obj) { - insert(pair.first, pair.second); - } - } - void insert(const std::string & key, const value & val) { - insert(mk_val(key), val); - } - virtual std::string type() const override { return "Object"; } - virtual bool is_immutable() const override { return false; } - virtual const std::vector> & as_ordered_object() const override { return val_obj; } - virtual string as_string() const override { - std::ostringstream ss; - ss << "{"; - for (size_t i = 0; i < val_obj.size(); i++) { - if (i > 0) ss << ", "; - auto & [key, val] = val_obj.at(i); - ss << value_to_string_repr(key) << ": " << value_to_string_repr(val); - } - ss << "}"; - return ss.str(); - } - virtual bool as_bool() const override { - return !unordered.empty(); - } - virtual bool has_key(const value & key) override { - if (!key->is_immutable() || !key->is_hashable()) { - throw std::runtime_error("Object key of unhashable type: " + key->type()); - } - return unordered.find(key) != unordered.end(); - } - virtual void insert(const value & key, const value & val) override { - bool replaced = false; - if (is_immutable()) { - throw std::runtime_error("Attempting to modify immutable type"); - } - if (has_key(key)) { - // if key exists, replace value in ordered list instead of appending - for (auto & pair : val_obj) { - if (*(pair.first) == *key) { - pair.second = val; - replaced = true; - break; - } - } - } - unordered[key] = val; - if (!replaced) { - val_obj.push_back({key, val}); - } - } - virtual value & at(const value & key, value & default_val) override { - if (!has_key(key)) { - return default_val; - } - return unordered.at(key); - } - virtual value & at(const value & key) override { - if (!has_key(key)) { - throw std::runtime_error("Key '" + key->as_string().str() + "' not found in value of type " + type()); - } - return unordered.at(key); - } - virtual value & at(const std::string & key, value & default_val) override { - value key_val = mk_val(key); - return at(key_val, default_val); - } - virtual value & at(const std::string & key) override { - value key_val = mk_val(key); - return at(key_val); - } - virtual const func_builtins & get_builtins() const override; - virtual bool is_hashable() const override { - if (std::all_of(val_obj.begin(), val_obj.end(), [&](auto & pair) -> bool { - const auto & val = pair.second; - return val->is_immutable() && val->is_hashable(); - })) { - return true; - } - return false; - } - virtual hasher unique_hash() const noexcept override { - auto hash = hasher(typeid(*this)); - for (const auto & [key, val] : val_obj) { - // must use digest to prevent problems from "concatenation" property of hasher - // for ex. hash of key="ab", value="c" should be different from key="a", value="bc" - const size_t key_hash = key->unique_hash().digest(); - const size_t val_hash = val->unique_hash().digest(); - hash.update(&key_hash, sizeof(key_hash)); - hash.update(&val_hash, sizeof(val_hash)); - } - return hash; - } -protected: - virtual bool equivalent(const value_t & other) const override { - return typeid(*this) == typeid(other) && is_hashable() && other.is_hashable() && std::equal(val_obj.begin(), val_obj.end(), other.val_obj.begin(), other.val_obj.end(), value_equivalence()); - } -}; -using value_object = std::shared_ptr; - -// -// none and undefined types -// - -struct value_none_t : public value_t { - virtual std::string type() const override { return "None"; } - virtual bool is_none() const override { return true; } - virtual bool as_bool() const override { return false; } - virtual string as_string() const override { return string(type()); } - virtual std::string as_repr() const override { return type(); } - virtual const func_builtins & get_builtins() const override; - virtual bool is_hashable() const override { return true; } - virtual hasher unique_hash() const noexcept override { - return hasher(typeid(*this)); - } -protected: - virtual bool equivalent(const value_t & other) const override { - return typeid(*this) == typeid(other); - } -}; -using value_none = std::shared_ptr; - -struct value_undefined_t : public value_t { - std::string hint; // for debugging, to indicate where undefined came from - value_undefined_t(const std::string & h = "") : hint(h) {} - virtual std::string type() const override { return hint.empty() ? "Undefined" : "Undefined (hint: '" + hint + "')"; } - virtual bool is_undefined() const override { return true; } - virtual bool as_bool() const override { return false; } - virtual std::string as_repr() const override { return type(); } - virtual const func_builtins & get_builtins() const override; - virtual hasher unique_hash() const noexcept override { - return hasher(typeid(*this)); - } -protected: - virtual bool equivalent(const value_t & other) const override { - return is_undefined() == other.is_undefined(); - } -}; -using value_undefined = std::shared_ptr; - -// -// function type -// - -struct func_args { -public: - std::string func_name; // for error messages - context & ctx; - func_args(context & ctx) : ctx(ctx) {} - value get_kwarg(const std::string & key, value default_val) const; - value get_kwarg_or_pos(const std::string & key, size_t pos) const; - value get_pos(size_t pos) const; - value get_pos(size_t pos, value default_val) const; - const std::vector & get_args() const; - size_t count() const { return args.size(); } - void push_back(const value & val); - void push_front(const value & val); - void ensure_count(size_t min, size_t max = 999) const { - size_t n = args.size(); - if (n < min || n > max) { - throw std::runtime_error("Function '" + func_name + "' expected between " + std::to_string(min) + " and " + std::to_string(max) + " arguments, got " + std::to_string(n)); - } - } - template void ensure_val(const value & ptr) const { - if (!is_val(ptr)) { - throw std::runtime_error("Function '" + func_name + "' expected value of type " + std::string(typeid(T).name()) + ", got " + ptr->type()); - } - } - void ensure_count(bool require0, bool require1, bool require2, bool require3) const { - static auto bool_to_int = [](bool b) { return b ? 1 : 0; }; - size_t required = bool_to_int(require0) + bool_to_int(require1) + bool_to_int(require2) + bool_to_int(require3); - ensure_count(required); - } - template void ensure_vals(bool required0 = true) const { - ensure_count(required0, false, false, false); - if (required0 && args.size() > 0) ensure_val(args[0]); - } - template void ensure_vals(bool required0 = true, bool required1 = true) const { - ensure_count(required0, required1, false, false); - if (required0 && args.size() > 0) ensure_val(args[0]); - if (required1 && args.size() > 1) ensure_val(args[1]); - } - template void ensure_vals(bool required0 = true, bool required1 = true, bool required2 = true) const { - ensure_count(required0, required1, required2, false); - if (required0 && args.size() > 0) ensure_val(args[0]); - if (required1 && args.size() > 1) ensure_val(args[1]); - if (required2 && args.size() > 2) ensure_val(args[2]); - } - template void ensure_vals(bool required0 = true, bool required1 = true, bool required2 = true, bool required3 = true) const { - ensure_count(required0, required1, required2, required3); - if (required0 && args.size() > 0) ensure_val(args[0]); - if (required1 && args.size() > 1) ensure_val(args[1]); - if (required2 && args.size() > 2) ensure_val(args[2]); - if (required3 && args.size() > 3) ensure_val(args[3]); - } -private: - std::vector args; -}; - -struct value_func_t : public value_t { - std::string name; - value arg0; // bound "this" argument, if any - value_func_t(const std::string & name, const func_handler & func) : name(name) { - val_func = func; - } - value_func_t(const std::string & name, const func_handler & func, const value & arg_this) : name(name), arg0(arg_this) { - val_func = func; - } - virtual value invoke(const func_args & args) const override { - func_args new_args(args); // copy - new_args.func_name = name; - if (arg0) { - new_args.push_front(arg0); - } - return val_func(new_args); - } - virtual std::string type() const override { return "Function"; } - virtual std::string as_repr() const override { return type() + "<" + name + ">(" + (arg0 ? arg0->as_repr() : "") + ")"; } - virtual bool is_hashable() const override { return false; } - virtual hasher unique_hash() const noexcept override { - // Note: this is unused for now, we don't support function as object keys - // use function pointer as unique identifier - const auto target = val_func.target(); - return hasher(typeid(*this)).update(&target, sizeof(target)); - } -protected: - virtual bool equivalent(const value_t & other) const override { - // Note: this is unused for now, we don't support function as object keys - // compare function pointers - // (val_func == other.val_func does not work as std::function::operator== is only used for nullptr check) - const auto target_this = this->val_func.target(); - const auto target_other = other.val_func.target(); - return typeid(*this) == typeid(other) && target_this == target_other; - } -}; -using value_func = std::shared_ptr; - -// special value for kwarg -struct value_kwarg_t : public value_t { - std::string key; - value val; - value_kwarg_t(const std::string & k, const value & v) : key(k), val(v) {} - virtual std::string type() const override { return "KwArg"; } - virtual std::string as_repr() const override { return type(); } - virtual bool is_hashable() const override { return true; } - virtual hasher unique_hash() const noexcept override { - const auto type_hash = typeid(*this).hash_code(); - auto hash = val->unique_hash(); - hash.update(&type_hash, sizeof(type_hash)) - .update(key.data(), key.size()); - return hash; - } -protected: - virtual bool equivalent(const value_t & other) const override { - const value_kwarg_t & other_val = static_cast(other); - return typeid(*this) == typeid(other) && key == other_val.key && val == other_val.val; - } -}; -using value_kwarg = std::shared_ptr; - - -} // namespace jinja diff --git a/common/json-partial.cpp b/common/json-partial.cpp deleted file mode 100644 index aaf11310ab8a..000000000000 --- a/common/json-partial.cpp +++ /dev/null @@ -1,324 +0,0 @@ -#include "json-partial.h" - -#include "log.h" - -#include - -#include -#include - -using json = nlohmann::ordered_json; - -enum common_json_stack_element_type { - COMMON_JSON_STACK_ELEMENT_OBJECT, - COMMON_JSON_STACK_ELEMENT_KEY, - COMMON_JSON_STACK_ELEMENT_ARRAY, -}; - -struct common_json_stack_element { - common_json_stack_element_type type; - std::string key; -}; - -bool common_json_parse( - const std::string & input, - const std::string & healing_marker, - common_json & out) -{ - std::string::const_iterator it = input.begin(); - const auto end = input.end(); - return common_json_parse(it, end, healing_marker, out); -} - -bool common_json_parse( - std::string::const_iterator & it, - const std::string::const_iterator & end, - const std::string & healing_marker, - common_json & out) -{ - // // https://json.nlohmann.me/features/parsing/sax_interface/ - struct json_error_locator : public nlohmann::json_sax { - std::size_t position; - bool found_error; - std::string last_token; - std::string exception_message; - std::vector stack; - - json_error_locator() : position(0), found_error(false) {} - - bool parse_error(std::size_t position, const std::string & last_token, const json::exception & ex) override { // NOLINT - this->position = position - 1; - this->found_error = true; - this->last_token = last_token; - this->exception_message = ex.what(); - return false; - } - void close_value() { - if (!stack.empty() && (stack.back().type == COMMON_JSON_STACK_ELEMENT_KEY)) { - stack.pop_back(); - } - } - bool null() override { // NOLINT - close_value(); - return true; - } - bool boolean(bool) override { // NOLINT - close_value(); - return true; - } - bool number_integer(number_integer_t) override { // NOLINT - close_value(); - return true; - } - bool number_unsigned(number_unsigned_t) override { // NOLINT - close_value(); - return true; - } - bool number_float(number_float_t, const string_t &) override { // NOLINT - close_value(); - return true; - } - bool string(string_t &) override { // NOLINT - close_value(); - return true; - } - bool binary(binary_t &) override { // NOLINT - close_value(); - return true; - } - bool start_object(std::size_t) override { // NOLINT - stack.push_back({COMMON_JSON_STACK_ELEMENT_OBJECT, ""}); - return true; - } - bool end_object() override { - GGML_ASSERT(!stack.empty() && stack.back().type == COMMON_JSON_STACK_ELEMENT_OBJECT); - stack.pop_back(); - close_value(); - return true; - } - bool key(string_t & key) override { // NOLINT - stack.push_back({COMMON_JSON_STACK_ELEMENT_KEY, key}); - return true; - } - bool start_array(std::size_t) override { // NOLINT - stack.push_back({COMMON_JSON_STACK_ELEMENT_ARRAY, ""}); - return true; - } - bool end_array() override { - GGML_ASSERT(!stack.empty() && stack.back().type == COMMON_JSON_STACK_ELEMENT_ARRAY); - stack.pop_back(); - close_value(); - return true; - } - }; - json_error_locator err_loc; - auto start = it; - json::sax_parse(it, end, &err_loc); - - if (err_loc.found_error) { - it = start; - auto temptative_end = it + err_loc.position; - // LOG_DBG("Error at position %zu (is_end = %s): %s\n", err_loc.position, temptative_end == end ? "true" : "false", err_loc.exception_message.c_str()); - - auto input = std::string(it, temptative_end); - try { - out.json = json::parse(input); - // out.json = json::parse(it, temptative_end); - it = temptative_end; - return true; - } catch (const std::exception & ex) { - // No, needs healing. - LOG_DBG("Failed to parse up to error: %s: <<<%s>>>\n", ex.what(), std::string(it, temptative_end).c_str()); - } - auto can_parse = [](const std::string & str) { - try { - auto _ = json::parse(str); // NOLINT - return true; - } catch (const std::exception &) { - return false; - } - }; - if (!healing_marker.empty() && !err_loc.stack.empty()) { - std::string str(it, temptative_end); - auto last_non_sp_pos = str.find_last_not_of(" \n\r\t"); - if (last_non_sp_pos == std::string::npos) { - throw std::runtime_error("Cannot heal a truncated JSON that stopped in an unknown location"); - } - auto last_non_sp_char = str[last_non_sp_pos]; - // Used to detect stops on a number, which may not be complete. - auto was_maybe_number = [&]() { - if (!str.empty() && std::isspace(str.back())) { - return false; - } - return std::isdigit(last_non_sp_char) || - last_non_sp_char == '.' || - last_non_sp_char == 'e' || - last_non_sp_char == 'E' || - last_non_sp_char == '-'; - }; - - std::string closing; - for (size_t i = err_loc.stack.size(); i > 0; i--) { - auto & el = err_loc.stack[i - 1]; - if (el.type == COMMON_JSON_STACK_ELEMENT_OBJECT) { - closing += "}"; - } else if (el.type == COMMON_JSON_STACK_ELEMENT_ARRAY) { - closing += "]"; - } else if (el.type != COMMON_JSON_STACK_ELEMENT_KEY) { - throw std::runtime_error("Unexpected stack element type"); - } - } - - // Matches a potentially partial unicode escape sequence, e.g. \u, \uX, \uXX, \uXXX, \uXXXX - static const std::regex partial_unicode_regex(R"(\\u(?:[0-9a-fA-F](?:[0-9a-fA-F](?:[0-9a-fA-F](?:[0-9a-fA-F])?)?)?)?$)"); - - auto is_high_surrogate = [&](const std::string & s) { - // Check if a partial of a high surrogate (U+D800-U+DBFF) - return s.length() >= 4 && - s[0] == '\\' && s[1] == 'u' && - std::tolower(s[2]) == 'd' && - (s[3] == '8' || s[3] == '9' || std::tolower(s[3]) == 'a' || std::tolower(s[3]) == 'b'); - }; - - // Initialize the unicode marker to a low surrogate to handle the edge case - // where a high surrogate (U+D800-U+DBFF) is immediately followed by a - // backslash (\) - std::string unicode_marker_padding = "udc00"; - std::smatch last_unicode_seq; - - if (std::regex_search(str, last_unicode_seq, partial_unicode_regex)) { - std::smatch second_last_seq; - std::string prelude = str.substr(0, last_unicode_seq.position()); - - // Pad the escape sequence with 0s until it forms a complete sequence of 6 characters - unicode_marker_padding = std::string(6 - last_unicode_seq.length(), '0'); - - if (is_high_surrogate(last_unicode_seq.str())) { - // If the sequence is a partial match for a high surrogate, add a low surrogate (U+DC00-U+UDFF) - unicode_marker_padding += "\\udc00"; - } else if (std::regex_search(prelude, second_last_seq, partial_unicode_regex)) { - if (is_high_surrogate(second_last_seq.str())) { - // If this follows a high surrogate, pad it to be a low surrogate - if (last_unicode_seq.length() == 2) { - unicode_marker_padding = "dc00"; - } else if (last_unicode_seq.length() == 3) { - unicode_marker_padding = "c00"; - } else { - // The original unicode_marker_padding is already padded with 0s - } - } - } - } - - const auto & magic_seed = out.healing_marker.marker = healing_marker;//"$llama.cpp.json$"; - - if (err_loc.stack.back().type == COMMON_JSON_STACK_ELEMENT_KEY) { - // We're inside an object value - if (last_non_sp_char == ':' && can_parse(str + "1" + closing)) { - // Was about to create an object value - str += (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\"" + closing; - } else if (can_parse(str + ": 1" + closing)) { - str += (out.healing_marker.json_dump_marker = ":\"" + magic_seed) + "\"" + closing; - } else if (last_non_sp_char == '{' && can_parse(str + closing)) { - // Was about to create an object - str += (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\": 1" + closing; - } else if (can_parse(str + "\"" + closing)) { - // Was inside an object value string - str += (out.healing_marker.json_dump_marker = magic_seed) + "\"" + closing; - } else if (str[str.length() - 1] == '\\' && can_parse(str + "\\\"" + closing)) { - // Was inside an object value string after an escape - str += (out.healing_marker.json_dump_marker = "\\" + magic_seed) + "\"" + closing; - } else if (can_parse(str + unicode_marker_padding + "\"" + closing)) { - // Was inside an object value string after a partial unicode escape - str += (out.healing_marker.json_dump_marker = unicode_marker_padding + magic_seed) + "\"" + closing; - } else { - // find last : - auto last_pos = str.find_last_of(':'); - if (last_pos == std::string::npos) { - throw std::runtime_error("Cannot heal a truncated JSON that stopped in an unknown location"); - } - // Cutting back to opening : for object value - str = str.substr(0, last_pos + 1) + (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\"" + closing; - } - } else if (err_loc.stack.back().type == COMMON_JSON_STACK_ELEMENT_ARRAY) { - if ((last_non_sp_char == ',' || last_non_sp_char == '[') && can_parse(str + "1" + closing)) { - // Was about to create an array value - str += (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\"" + closing; - } else if (can_parse(str + "\"" + closing)) { - // Was inside an array value string - str += (out.healing_marker.json_dump_marker = magic_seed) + "\"" + closing; - } else if (str[str.length() - 1] == '\\' && can_parse(str + "\\\"" + closing)) { - // Was inside an array value string after an escape - str += (out.healing_marker.json_dump_marker = "\\" + magic_seed) + "\"" + closing; - } else if (can_parse(str + unicode_marker_padding + "\"" + closing)) { - // Was inside an array value string after a partial unicode escape - str += (out.healing_marker.json_dump_marker = unicode_marker_padding + magic_seed) + "\"" + closing; - } else if (!was_maybe_number() && can_parse(str + ", 1" + closing)) { - // Had just finished a value - str += (out.healing_marker.json_dump_marker = ",\"" + magic_seed) + "\"" + closing; - } else { - auto last_pos = str.find_last_of("[,"); - if (last_pos == std::string::npos) { - throw std::runtime_error("Cannot heal a truncated JSON array stopped in an unknown location"); - } - // Cutting back to last [ or , for array value - str = str.substr(0, last_pos + 1) + (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\"" + closing; - } - } else if (err_loc.stack.back().type == COMMON_JSON_STACK_ELEMENT_OBJECT) { - if ((last_non_sp_char == '{' && can_parse(str + closing)) || - (last_non_sp_char == ',' && can_parse(str + "\"\": 1" + closing))) { - // Was about to create an object key+value - str += (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\": 1" + closing; - } else if (!was_maybe_number() && can_parse(str + ",\"\": 1" + closing)) { - // Was about to create an object key+value - str += (out.healing_marker.json_dump_marker = ",\"" + magic_seed) + "\": 1" + closing; - } else if (can_parse(str + "\": 1" + closing)) { - // Was inside an object key string - str += (out.healing_marker.json_dump_marker = magic_seed) + "\": 1" + closing; - } else if (str[str.length() - 1] == '\\' && can_parse(str + "\\\": 1" + closing)) { - // Was inside an object key string after an escape - str += (out.healing_marker.json_dump_marker = "\\" + magic_seed) + "\": 1" + closing; - } else if (can_parse(str + unicode_marker_padding + "\": 1" + closing)) { - // Was inside an object key string after a partial unicode escape - str += (out.healing_marker.json_dump_marker = unicode_marker_padding + magic_seed) + "\": 1" + closing; - } else { - auto last_pos = str.find_last_of(':'); - if (last_pos == std::string::npos) { - throw std::runtime_error("Cannot heal a truncated JSON object stopped in an unknown location"); - } - // fprintf(stderr, "Cutting back to last : for object key+value\n"); - str = str.substr(0, last_pos + 1) + (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\"" + closing; - } - } else { - throw std::runtime_error("Cannot heal a truncated JSON object stopped in an unknown location"); - } - // fprintf(stderr, "HEALED:\nSTRING <<<\n%s\n>>>\n\nmagic_cut: <<<\n%s\n>>>\n\n", str.c_str(), out.healing_marker.json_dump_marker.c_str()); - out.json = json::parse(str); - it = temptative_end; - return true; - } - // handle unclosed top-level primitive - if (err_loc.position != 0 && !healing_marker.empty() && err_loc.stack.empty()) { - std::string str(it, temptative_end); - const auto & magic_seed = out.healing_marker.marker = healing_marker; - if (can_parse(str + "\"")) { - // Was inside an string - str += (out.healing_marker.json_dump_marker = magic_seed) + "\""; - } else if (str[str.length() - 1] == '\\' && can_parse(str + "\\\"")) { - // Was inside an string after an escape - str += (out.healing_marker.json_dump_marker = "\\" + magic_seed) + "\""; - } else { - // TODO: handle more unclosed top-level primitive if the stack was empty but we got an error (e.g. "tru", "\"", etc...) - // fprintf(stderr, "Closing: TODO\n"); - return false; - } - out.json = json::parse(str); - it = temptative_end; - return true; - } - return false; - } - out.json = json::parse(it, end); - it = end; - return true; -} diff --git a/common/json-partial.h b/common/json-partial.h deleted file mode 100644 index be51aabfbf41..000000000000 --- a/common/json-partial.h +++ /dev/null @@ -1,39 +0,0 @@ -#pragma once - -// TODO: use json_fwd.hpp when possible -#include - -// Healing marker (empty if the JSON was fully parsed / wasn't healed). -struct common_healing_marker { - // Raw marker. - std::string marker; - - // Cutting the `common_json.json.dump()` string at the (only) occurrence of this marker should yield the original partial JSON string (modulo spaces / if it had the same dump format). - std::string json_dump_marker; -}; - -// Represents a parsed JSON object, with its optional healing marker (a JSON dump fragment that can be used to find the position of healing in the JSON dump string) -struct common_json { - nlohmann::ordered_json json; - - common_healing_marker healing_marker; -}; - -// Parse the JSON string, healing (closing) any partial JSON if `healing_marker` is not empty. -// -// Healing completes partial JSON strings by adding a (possibly modified) healing marker, then whatever is needed to close the JSON. -// This allows to parse the resulting healed JSON string, yet be able to cut it again if needed at the healing marker. -// (this is used when parsing JSON outputs from the models, then crafting partial JSONs for the partial tool calls in OAI format). -// -// For instance, parsing `{` with a healing marker `foo` will produce a healed JSON `{"foo":1}`, w/ json_dump_marker = `"foo"` (which can be used to break the JSON again). -bool common_json_parse( - const std::string & input, - const std::string & healing_marker, - common_json & out); - -// Parse the JSON string (see overload above), but advancing an iterator to the end of the input when the (potentially partial) parsing succeeds. -bool common_json_parse( - std::string::const_iterator & it, - const std::string::const_iterator & end, - const std::string & healing_marker, - common_json & out); diff --git a/common/json-schema-to-grammar.cpp b/common/json-schema-to-grammar.cpp deleted file mode 100644 index e2c4d6ce22e5..000000000000 --- a/common/json-schema-to-grammar.cpp +++ /dev/null @@ -1,1189 +0,0 @@ -#include "json-schema-to-grammar.h" -#include "common.h" - -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -using json = nlohmann::ordered_json; - -static std::string build_repetition(const std::string & item_rule, int min_items, int max_items, const std::string & separator_rule = "") { - auto has_max = max_items != std::numeric_limits::max(); - - if (max_items == 0) { - return ""; - } - if (min_items == 0 && max_items == 1) { - return item_rule + "?"; - } - - if (separator_rule.empty()) { - if (min_items == 1 && !has_max) { - return item_rule + "+"; - } - if (min_items == 0 && !has_max) { - return item_rule + "*"; - } - return item_rule + "{" + std::to_string(min_items) + "," + (has_max ? std::to_string(max_items) : "") + "}"; - } - - auto result = item_rule + " " + build_repetition("(" + separator_rule + " " + item_rule + ")", min_items == 0 ? 0 : min_items - 1, has_max ? max_items - 1 : max_items); - if (min_items == 0) { - result = "(" + result + ")?"; - } - return result; -} - -static void build_min_max_int(int64_t min_value, int64_t max_value, std::stringstream & out, int decimals_left = 16, bool top_level = true) { - auto has_min = min_value != std::numeric_limits::min(); - auto has_max = max_value != std::numeric_limits::max(); - - auto digit_range = [&](char from, char to) { - out << "["; - if (from == to) { - out << from; - } else { - out << from << "-" << to; - } - out << "]"; - }; - auto more_digits = [&](int min_digits, int max_digits) { - out << "[0-9]"; - if (min_digits == max_digits && min_digits == 1) { - return; - } - out << "{"; - out << min_digits; - if (max_digits != min_digits) { - out << ","; - if (max_digits != std::numeric_limits::max()) { - out << max_digits; - } - } - out << "}"; - }; - std::function uniform_range = - [&](const std::string_view & from, const std::string_view & to) { - size_t i = 0; - while (i < from.length() && i < to.length() && from[i] == to[i]) { - i++; - } - if (i > 0) { - out << "\"" << from.substr(0, i) << "\""; - } - if (i < from.length() && i < to.length()) { - if (i > 0) { - out << " "; - } - auto sub_len = from.length() - i - 1; - if (sub_len > 0) { - auto from_sub = from.substr(i + 1); - auto to_sub = to.substr(i + 1); - auto sub_zeros = string_repeat("0", sub_len); - auto sub_nines = string_repeat("9", sub_len); - - auto to_reached = false; - out << "("; - if (from_sub == sub_zeros) { - digit_range(from[i], to[i] - 1); - out << " "; - more_digits(sub_len, sub_len); - } else { - out << "[" << from[i] << "] "; - out << "("; - uniform_range(from_sub, sub_nines); - out << ")"; - if (from[i] < to[i] - 1) { - out << " | "; - if (to_sub == sub_nines) { - digit_range(from[i] + 1, to[i]); - to_reached = true; - } else { - digit_range(from[i] + 1, to[i] - 1); - } - out << " "; - more_digits(sub_len, sub_len); - } - } - if (!to_reached) { - out << " | "; - digit_range(to[i], to[i]); - out << " "; - uniform_range(sub_zeros, to_sub); - } - out << ")"; - } else { - out << "[" << from[i] << "-" << to[i] << "]"; - } - } - }; - - if (has_min && has_max) { - if (min_value < 0 && max_value < 0) { - out << "\"-\" ("; - build_min_max_int(-max_value, -min_value, out, decimals_left, /* top_level= */ true); - out << ")"; - return; - } - - if (min_value < 0) { - out << "\"-\" ("; - build_min_max_int(0, -min_value, out, decimals_left, /* top_level= */ true); - out << ") | "; - min_value = 0; - } - - auto min_s = std::to_string(min_value); - auto max_s = std::to_string(max_value); - auto min_digits = min_s.length(); - auto max_digits = max_s.length(); - - for (auto digits = min_digits; digits < max_digits; digits++) { - uniform_range(min_s, string_repeat("9", digits)); - min_s = "1" + string_repeat("0", digits); - out << " | "; - } - uniform_range(min_s, max_s); - return; - } - - auto less_decimals = std::max(decimals_left - 1, 1); - - if (has_min) { - if (min_value < 0) { - out << "\"-\" ("; - build_min_max_int(std::numeric_limits::min(), -min_value, out, decimals_left, /* top_level= */ false); - out << ") | [0] | [1-9] "; - more_digits(0, decimals_left - 1); - } else if (min_value == 0) { - if (top_level) { - out << "[0] | [1-9] "; - more_digits(0, less_decimals); - } else { - more_digits(1, decimals_left); - } - } else if (min_value <= 9) { - char c = '0' + min_value; - auto range_start = top_level ? '1' : '0'; - if (c > range_start) { - digit_range(range_start, c - 1); - out << " "; - more_digits(1, less_decimals); - out << " | "; - } - digit_range(c, '9'); - out << " "; - more_digits(0, less_decimals); - } else { - auto min_s = std::to_string(min_value); - auto len = min_s.length(); - auto c = min_s[0]; - - if (c > '1') { - digit_range(top_level ? '1' : '0', c - 1); - out << " "; - more_digits(len, less_decimals); - out << " | "; - } - digit_range(c, c); - out << " ("; - build_min_max_int(std::stoll(min_s.substr(1)), std::numeric_limits::max(), out, less_decimals, /* top_level= */ false); - out << ")"; - if (c < '9') { - out << " | "; - digit_range(c + 1, '9'); - out << " "; - more_digits(len - 1, less_decimals); - } - } - return; - } - - if (has_max) { - if (max_value >= 0) { - if (top_level) { - out << "\"-\" [1-9] "; - more_digits(0, less_decimals); - out << " | "; - } - build_min_max_int(0, max_value, out, decimals_left, /* top_level= */ true); - } else { - out << "\"-\" ("; - build_min_max_int(-max_value, std::numeric_limits::max(), out, decimals_left, /* top_level= */ false); - out << ")"; - } - return; - } - - throw std::runtime_error("At least one of min_value or max_value must be set"); -} - -const std::string SPACE_RULE = "| \" \" | \"\\n\"{1,2} [ \\t]{0,20}"; - -struct BuiltinRule { - std::string content; - std::vector deps; -}; - -static std::unordered_map PRIMITIVE_RULES = { - {"boolean", {"(\"true\" | \"false\") space", {}}}, - {"decimal-part", {"[0-9]{1,16}", {}}}, - {"integral-part", {"[0] | [1-9] [0-9]{0,15}", {}}}, - {"number", {"(\"-\"? integral-part) (\".\" decimal-part)? ([eE] [-+]? integral-part)? space", {"integral-part", "decimal-part"}}}, - {"integer", {"(\"-\"? integral-part) space", {"integral-part"}}}, - {"value", {"object | array | string | number | boolean | null", {"object", "array", "string", "number", "boolean", "null"}}}, - {"object", {"\"{\" space ( string \":\" space value (\",\" space string \":\" space value)* )? \"}\" space", {"string", "value"}}}, - {"array", {"\"[\" space ( value (\",\" space value)* )? \"]\" space", {"value"}}}, - {"uuid", {"\"\\\"\" [0-9a-fA-F]{8} \"-\" [0-9a-fA-F]{4} \"-\" [0-9a-fA-F]{4} \"-\" [0-9a-fA-F]{4} \"-\" [0-9a-fA-F]{12} \"\\\"\" space", {}}}, - {"char", {"[^\"\\\\\\x7F\\x00-\\x1F] | [\\\\] ([\"\\\\bfnrt] | \"u\" [0-9a-fA-F]{4})", {}}}, - {"string", {"\"\\\"\" char* \"\\\"\" space", {"char"}}}, - {"null", {"\"null\" space", {}}}, -}; - -static std::unordered_map STRING_FORMAT_RULES = { - {"date", {"[0-9]{4} \"-\" ( \"0\" [1-9] | \"1\" [0-2] ) \"-\" ( \"0\" [1-9] | [1-2] [0-9] | \"3\" [0-1] )", {}}}, - {"time", {"([01] [0-9] | \"2\" [0-3]) \":\" [0-5] [0-9] \":\" [0-5] [0-9] ( \".\" [0-9]{3} )? ( \"Z\" | ( \"+\" | \"-\" ) ( [01] [0-9] | \"2\" [0-3] ) \":\" [0-5] [0-9] )", {}}}, - {"date-time", {"date \"T\" time", {"date", "time"}}}, - {"date-string", {"\"\\\"\" date \"\\\"\" space", {"date"}}}, - {"time-string", {"\"\\\"\" time \"\\\"\" space", {"time"}}}, - {"date-time-string", {"\"\\\"\" date-time \"\\\"\" space", {"date-time"}}} -}; - -static bool is_reserved_name(const std::string & name) { - static const std::unordered_set RESERVED_NAMES = [] { - std::unordered_set s; - s.insert("root"); - for (const auto & p : PRIMITIVE_RULES) { - s.insert(p.first); - } - for (const auto & p : STRING_FORMAT_RULES) { - s.insert(p.first); - } - return s; - }(); - return RESERVED_NAMES.find(name) != RESERVED_NAMES.end(); -} - -static std::regex INVALID_RULE_CHARS_RE("[^a-zA-Z0-9-]+"); -static std::regex GRAMMAR_LITERAL_ESCAPE_RE("[\r\n\"\\\\]"); -static std::regex GRAMMAR_RANGE_LITERAL_ESCAPE_RE("[\r\n\"\\]\\-\\\\]"); -static std::unordered_map GRAMMAR_LITERAL_ESCAPES = { - {'\r', "\\r"}, {'\n', "\\n"}, {'"', "\\\""}, {'-', "\\-"}, {']', "\\]"}, {'\\', "\\\\"} -}; - -static std::unordered_set NON_LITERAL_SET = {'|', '.', '(', ')', '[', ']', '{', '}', '*', '+', '?'}; -static std::unordered_set ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS = {'^', '$', '.', '[', ']', '(', ')', '|', '{', '}', '*', '+', '?'}; - -static std::string replacePattern(const std::string & input, const std::regex & regex, const std::function & replacement) { - std::smatch match; - std::string result; - - std::string::const_iterator searchStart(input.cbegin()); - std::string::const_iterator searchEnd(input.cend()); - - while (std::regex_search(searchStart, searchEnd, match, regex)) { - result.append(searchStart, searchStart + match.position()); - result.append(replacement(match)); - searchStart = match.suffix().first; - } - - result.append(searchStart, searchEnd); - - return result; -} - -static std::string format_literal(const std::string & literal) { - std::string escaped = replacePattern(literal, GRAMMAR_LITERAL_ESCAPE_RE, [&](const std::smatch & match) { - char c = match.str()[0]; - return GRAMMAR_LITERAL_ESCAPES.at(c); - }); - return "\"" + escaped + "\""; -} - -std::string gbnf_format_literal(const std::string & literal) { return format_literal(literal); } - -class common_schema_converter { -private: - friend class common_schema_info; - friend std::string build_grammar(const std::function & cb, const common_grammar_options & options); - std::function _fetch_json; - bool _dotall; - std::map _rules; - std::unordered_map _refs; - std::unordered_set _refs_being_resolved; - std::vector _errors; - std::vector _warnings; - - std::string _add_rule(const std::string & name, const std::string & rule) { - std::string esc_name = regex_replace(name, INVALID_RULE_CHARS_RE, "-"); - if (_rules.find(esc_name) == _rules.end() || _rules[esc_name] == rule) { - _rules[esc_name] = rule; - return esc_name; - } - int i = 0; - while (_rules.find(esc_name + std::to_string(i)) != _rules.end() && _rules[esc_name + std::to_string(i)] != rule) { - i++; - } - std::string key = esc_name + std::to_string(i); - _rules[key] = rule; - return key; - } - - std::string _generate_union_rule(const std::string & name, const std::vector & alt_schemas) { - std::vector rules; - rules.reserve(alt_schemas.size()); - for (size_t i = 0; i < alt_schemas.size(); i++) { - rules.push_back(visit(alt_schemas[i], name + (name.empty() ? "alternative-" : "-") + std::to_string(i))); - } - return string_join(rules, " | "); - } - - std::string _visit_pattern(const std::string & pattern, const std::string & name) { - if (!(pattern.front() == '^' && pattern.back() == '$')) { - _errors.push_back("Pattern must start with '^' and end with '$'"); - return ""; - } - std::string sub_pattern = pattern.substr(1, pattern.length() - 2); - std::unordered_map sub_rule_ids; - - size_t i = 0; - size_t length = sub_pattern.length(); - - using literal_or_rule = std::pair; - auto to_rule = [&](const literal_or_rule & ls) { - auto is_literal = ls.second; - auto s = ls.first; - return is_literal ? "\"" + s + "\"" : s; - }; - std::function transform = [&]() -> literal_or_rule { - size_t start = i; - std::vector seq; - - auto get_dot = [&]() { - std::string rule; - if (_dotall) { - rule = "[\\U00000000-\\U0010FFFF]"; - } else { - rule = "[^\\x0A\\x0D]"; - } - return _add_rule("dot", rule); - }; - - // Joins the sequence, merging consecutive literals together. - auto join_seq = [&]() { - std::vector ret; - - std::string literal; - auto flush_literal = [&]() { - if (literal.empty()) { - return false; - } - ret.emplace_back(literal, true); - literal.clear(); - return true; - }; - - for (const auto & item : seq) { - auto is_literal = item.second; - if (is_literal) { - literal += item.first; - } else { - flush_literal(); - ret.push_back(item); - } - } - flush_literal(); - - std::vector results; - results.reserve(ret.size()); - for (const auto & item : ret) { - results.push_back(to_rule(item)); - } - return std::make_pair(string_join(results, " "), false); - }; - - while (i < length) { - char c = sub_pattern[i]; - if (c == '.') { - seq.emplace_back(get_dot(), false); - i++; - } else if (c == '(') { - i++; - if (i < length && sub_pattern[i] == '?') { - if (i + 1 < length && sub_pattern[i + 1] == ':') { - i += 2; // skip "?:" for non-capturing group, treat as regular group - } else { - // lookahead/lookbehind (?=, ?!, ?<=, ? 0) { - if (sub_pattern[i] == '\\' && i + 1 < length) { - i += 2; // skip escaped character - } else { - if (sub_pattern[i] == '(') depth++; - else if (sub_pattern[i] == ')') depth--; - i++; - } - } - continue; - } - } - seq.emplace_back("(" + to_rule(transform()) + ")", false); - } else if (c == ')') { - i++; - if (start > 0 && sub_pattern[start - 1] != '(' && (start < 2 || sub_pattern[start - 2] != '?' || sub_pattern[start - 1] != ':')) { - _errors.push_back("Unbalanced parentheses"); - } - return join_seq(); - } else if (c == '[') { - std::string square_brackets = std::string(1, c); - i++; - while (i < length && sub_pattern[i] != ']') { - if (sub_pattern[i] == '\\') { - square_brackets += sub_pattern.substr(i, 2); - i += 2; - } else { - square_brackets += sub_pattern[i]; - i++; - } - } - if (i >= length) { - _errors.push_back("Unbalanced square brackets"); - } - square_brackets += ']'; - i++; - seq.emplace_back(square_brackets, false); - } else if (c == '|') { - seq.emplace_back("|", false); - i++; - } else if (c == '*' || c == '+' || c == '?') { - seq.back() = std::make_pair(to_rule(seq.back()) + c, false); - i++; - } else if (c == '{') { - std::string curly_brackets = std::string(1, c); - i++; - while (i < length && sub_pattern[i] != '}') { - curly_brackets += sub_pattern[i]; - i++; - } - if (i >= length) { - _errors.push_back("Unbalanced curly brackets"); - } - curly_brackets += '}'; - i++; - auto nums = string_split(curly_brackets.substr(1, curly_brackets.length() - 2), ","); - int min_times = 0; - int max_times = std::numeric_limits::max(); - try { - if (nums.size() == 1) { - min_times = max_times = std::stoi(nums[0]); - } else if (nums.size() != 2) { - _errors.push_back("Wrong number of values in curly brackets"); - } else { - if (!nums[0].empty()) { - min_times = std::stoi(nums[0]); - } - if (!nums[1].empty()) { - max_times = std::stoi(nums[1]); - } - } - } catch (const std::invalid_argument & e) { - _errors.push_back("Invalid number in curly brackets"); - return std::make_pair("", false); - } - auto &last = seq.back(); - auto &sub = last.first; - auto sub_is_literal = last.second; - - if (!sub_is_literal) { - std::string & sub_id = sub_rule_ids[sub]; - if (sub_id.empty()) { - sub_id = _add_rule(name + "-" + std::to_string(sub_rule_ids.size()), sub); - } - sub = sub_id; - } - seq.back().first = build_repetition( - sub_is_literal ? "\"" + sub + "\"" : sub, - min_times, - max_times, - "" - ); - seq.back().second = false; - } else { - std::string literal; - auto is_non_literal = [&](char c) { - return NON_LITERAL_SET.find(c) != NON_LITERAL_SET.end(); - }; - while (i < length) { - if (sub_pattern[i] == '\\' && i < length - 1) { - char next = sub_pattern[i + 1]; - if (ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS.find(next) != ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS.end()) { - i++; - literal += sub_pattern[i]; - i++; - } else { - literal += sub_pattern.substr(i, 2); - i += 2; - } - } else if (sub_pattern[i] == '"') { - literal += "\\\""; - i++; - } else if (!is_non_literal(sub_pattern[i]) && - (i == length - 1 || literal.empty() || sub_pattern[i + 1] == '.' || !is_non_literal(sub_pattern[i + 1]))) { - literal += sub_pattern[i]; - i++; - } else { - break; - } - } - if (!literal.empty()) { - seq.emplace_back(literal, true); - } - } - } - return join_seq(); - }; - return _add_rule(name, "\"\\\"\" (" + to_rule(transform()) + ") \"\\\"\" space"); - } - - /* - Returns a rule that matches a JSON string that is none of the provided strings - - not_strings({"a"}) - -> ["] ( [a] char+ | [^"a] char* )? ["] space - not_strings({"and", "also"}) - -> ["] ( [a] ([l] ([s] ([o] char+ | [^"o] char*) | [^"s] char*) | [n] ([d] char+ | [^"d] char*) | [^"ln] char*) | [^"a] char* )? ["] space - */ - std::string _not_strings(const std::vector & strings) { - - struct TrieNode { - std::map children; - bool is_end_of_string; - - TrieNode() : is_end_of_string(false) {} - - void insert(const std::string & string) { - auto *node = this; - for (char c : string) { - node = &node->children[c]; - } - node->is_end_of_string = true; - } - }; - - TrieNode trie; - for (const auto & s : strings) { - trie.insert(s); - } - - std::string char_rule = _add_primitive("char", PRIMITIVE_RULES.at("char")); - std::ostringstream out; - out << "[\"] ( "; - std::function visit = [&](const TrieNode & node) { - std::ostringstream rejects; - auto first = true; - for (const auto & kv : node.children) { - rejects << kv.first; - if (first) { - first = false; - } else { - out << " | "; - } - out << "[" << kv.first << "]"; - if (!kv.second.children.empty()) { - out << " ("; - visit(kv.second); - out << ")"; - } else if (kv.second.is_end_of_string) { - out << " " << char_rule << "+"; - } - } - if (!node.children.empty()) { - if (!first) { - out << " | "; - } - out << "[^\"" << rejects.str() << "] " << char_rule << "*"; - } - }; - visit(trie); - - out << " )"; - if (!trie.is_end_of_string) { - out << "?"; - } - out << " [\"] space"; - return out.str(); - } - - std::string _resolve_ref(const std::string & ref) { - auto it = ref.find('#'); - std::string ref_fragment = it != std::string::npos ? ref.substr(it + 1) : ref; - static const std::regex nonalphanumeric_regex(R"([^a-zA-Z0-9-]+)"); - std::string ref_name = "ref" + std::regex_replace(ref_fragment, nonalphanumeric_regex, "-"); - if (_rules.find(ref_name) == _rules.end() && _refs_being_resolved.find(ref) == _refs_being_resolved.end()) { - _refs_being_resolved.insert(ref); - json resolved = _refs[ref]; - ref_name = visit(resolved, ref_name); - _refs_being_resolved.erase(ref); - } - return ref_name; - } - - std::string _build_object_rule( - const std::vector> & properties, - const std::unordered_set & required, - const std::string & name, - const json & additional_properties) - { - std::vector required_props; - std::vector optional_props; - std::unordered_map prop_kv_rule_names; - std::vector prop_names; - for (const auto & kv : properties) { - const auto &prop_name = kv.first; - const auto &prop_schema = kv.second; - - std::string prop_rule_name = visit(prop_schema, name + (name.empty() ? "" : "-") + prop_name); - prop_kv_rule_names[prop_name] = _add_rule( - name + (name.empty() ? "" : "-") + prop_name + "-kv", - format_literal(json(prop_name).dump()) + " space \":\" space " + prop_rule_name - ); - if (required.find(prop_name) != required.end()) { - required_props.push_back(prop_name); - } else { - optional_props.push_back(prop_name); - } - prop_names.push_back(prop_name); - } - if ((additional_properties.is_boolean() && additional_properties.get()) || additional_properties.is_object()) { - std::string sub_name = name + (name.empty() ? "" : "-") + "additional"; - std::string value_rule = - additional_properties.is_object() ? visit(additional_properties, sub_name + "-value") - : _add_primitive("value", PRIMITIVE_RULES.at("value")); - - auto key_rule = - prop_names.empty() ? _add_primitive("string", PRIMITIVE_RULES.at("string")) - : _add_rule(sub_name + "-k", _not_strings(prop_names)); - std::string kv_rule = _add_rule(sub_name + "-kv", key_rule + " \":\" space " + value_rule); - prop_kv_rule_names["*"] = kv_rule; - optional_props.push_back("*"); - } - - std::string rule = "\"{\" space "; - for (size_t i = 0; i < required_props.size(); i++) { - if (i > 0) { - rule += " \",\" space "; - } - rule += prop_kv_rule_names[required_props[i]]; - } - - if (!optional_props.empty()) { - rule += " ("; - if (!required_props.empty()) { - rule += " \",\" space ( "; - } - - std::function &, bool)> get_recursive_refs = [&](const std::vector & ks, bool first_is_optional) { - std::string res; - if (ks.empty()) { - return res; - } - const std::string& k = ks[0]; - std::string kv_rule_name = prop_kv_rule_names[k]; - std::string comma_ref = "( \",\" space " + kv_rule_name + " )"; - if (first_is_optional) { - res = comma_ref + (k == "*" ? "*" : "?"); - } else { - res = kv_rule_name + (k == "*" ? " " + comma_ref + "*" : ""); - } - if (ks.size() > 1) { - res += " " + _add_rule( - name + (name.empty() ? "" : "-") + k + "-rest", - get_recursive_refs(std::vector(ks.begin() + 1, ks.end()), true) - ); - } - return res; - }; - - for (size_t i = 0; i < optional_props.size(); i++) { - if (i > 0) { - rule += " | "; - } - rule += get_recursive_refs(std::vector(optional_props.begin() + i, optional_props.end()), false); - } - if (!required_props.empty()) { - rule += " )"; - } - rule += " )?"; - } - - rule += " \"}\" space"; - - return rule; - } - - std::string _add_primitive(const std::string & name, const BuiltinRule & rule) { - auto n = _add_rule(name, rule.content); - for (const auto & dep : rule.deps) { - BuiltinRule dep_rule; - auto it = PRIMITIVE_RULES.find(dep); - if (it == PRIMITIVE_RULES.end()) { - it = STRING_FORMAT_RULES.find(dep); - if (it == STRING_FORMAT_RULES.end()) { - _errors.push_back("Rule " + dep + " not known"); - continue; - } - } - if (_rules.find(dep) == _rules.end()) { - _add_primitive(dep, it->second); - } - } - return n; - } - -public: - common_schema_converter( - const std::function & fetch_json, - bool dotall) - : _fetch_json(fetch_json), _dotall(dotall) - { - _rules["space"] = SPACE_RULE; - } - - void resolve_refs(json & schema, const std::string & url) { - /* - * Resolves all $ref fields in the given schema, fetching any remote schemas, - * replacing each $ref with absolute reference URL and populates _refs with the - * respective referenced (sub)schema dictionaries. - */ - std::function visit_refs = [&](json & n) { - if (n.is_array()) { - for (auto & x : n) { - visit_refs(x); - } - } else if (n.is_object()) { - if (n.contains("$ref")) { - std::string ref = n["$ref"]; - if (_refs.find(ref) == _refs.end()) { - json target; - if (ref.find("https://") == 0) { - std::string base_url = ref.substr(0, ref.find('#')); - auto it = _refs.find(base_url); - if (it != _refs.end()) { - target = it->second; - } else { - // Fetch the referenced schema and resolve its refs - auto referenced = _fetch_json(ref); - resolve_refs(referenced, base_url); - _refs[base_url] = referenced; - } - if (ref.find('#') == std::string::npos || ref.substr(ref.find('#') + 1).empty()) { - return; - } - } else if (ref.find("#/") == 0) { - target = schema; - n["$ref"] = url + ref; - ref = url + ref; - } else { - _errors.push_back("Unsupported ref: " + ref); - return; - } - std::string pointer = ref.substr(ref.find('#') + 1); - std::vector tokens = string_split(pointer, "/"); - for (size_t i = 1; i < tokens.size(); ++i) { - const std::string& sel = tokens[i]; - if (target.is_object() && target.contains(sel)) { - target = target[sel]; - } else if (target.is_array()) { - size_t sel_index; - try { - sel_index = std::stoull(sel); - } catch (const std::invalid_argument & e) { - sel_index = target.size(); - } - if (sel_index >= target.size()) { - _errors.push_back("Error resolving ref " + ref + ": " + sel + " not in " + target.dump()); - return; - } - target = target[sel_index]; - } else { - _errors.push_back("Error resolving ref " + ref + ": " + sel + " not in " + target.dump()); - return; - } - } - _refs[ref] = target; - } - } else { - for (const auto & kv : n.items()) { - visit_refs(kv.value()); - } - } - } - }; - - visit_refs(schema); - } - - static std::string _generate_constant_rule(const json & value) { - return format_literal(value.dump()); - } - - std::string visit(const json & schema, const std::string & name) { - json schema_type = schema.contains("type") ? schema["type"] : json(); - std::string schema_format = schema.contains("format") ? schema["format"].get() : ""; - std::string rule_name = is_reserved_name(name) ? name + "-" : name.empty() ? "root" : name; - - if (schema.contains("$ref")) { - return _add_rule(rule_name, _resolve_ref(schema["$ref"])); - } - if (schema.contains("oneOf") || schema.contains("anyOf")) { - std::vector alt_schemas = schema.contains("oneOf") ? schema["oneOf"].get>() : schema["anyOf"].get>(); - return _add_rule(rule_name, _generate_union_rule(name, alt_schemas)); - } - if (schema_type.is_array()) { - std::vector schema_types; - for (const auto & t : schema_type) { - json schema_copy(schema); - schema_copy["type"] = t; - schema_types.push_back(schema_copy); - } - return _add_rule(rule_name, _generate_union_rule(name, schema_types)); - } - if (schema.contains("const")) { - return _add_rule(rule_name, _generate_constant_rule(schema["const"]) + " space"); - } - if (schema.contains("enum")) { - std::vector enum_values; - for (const auto & v : schema["enum"]) { - enum_values.push_back(_generate_constant_rule(v)); - } - return _add_rule(rule_name, "(" + string_join(enum_values, " | ") + ") space"); - } - if ((schema_type.is_null() || schema_type == "object") - && (schema.contains("properties") || - (schema.contains("additionalProperties") && schema["additionalProperties"] != true))) { - std::unordered_set required; - if (schema.contains("required") && schema["required"].is_array()) { - for (const auto & item : schema["required"]) { - if (item.is_string()) { - required.insert(item.get()); - } - } - } - std::vector> properties; - if (schema.contains("properties")) { - for (const auto & prop : schema["properties"].items()) { - properties.emplace_back(prop.key(), prop.value()); - } - } - return _add_rule(rule_name, - _build_object_rule( - properties, required, name, - schema.contains("additionalProperties") ? schema["additionalProperties"] : json())); - } - if ((schema_type.is_null() || schema_type == "object" || schema_type == "string") && schema.contains("allOf")) { - std::unordered_set required; - std::vector> properties; - std::map enum_values; - const std::string& hybrid_name = name; - std::function add_component = [&](const json & comp_schema, bool is_required) { - if (comp_schema.contains("$ref")) { - add_component(_refs[comp_schema["$ref"]], is_required); - } else if (comp_schema.contains("properties")) { - for (const auto & prop : comp_schema["properties"].items()) { - properties.emplace_back(prop.key(), prop.value()); - if (is_required) { - required.insert(prop.key()); - } - } - } else if (comp_schema.contains("enum")) { - for (const auto & v : comp_schema["enum"]) { - const auto rule = _generate_constant_rule(v); - if (enum_values.find(rule) == enum_values.end()) { - enum_values[rule] = 0; - } - enum_values[rule] += 1; - } - } else { - // todo warning - } - }; - for (const auto & t : schema["allOf"]) { - if (t.contains("anyOf")) { - for (const auto & tt : t["anyOf"]) { - add_component(tt, false); - } - } else { - add_component(t, true); - } - } - if (!enum_values.empty()) { - std::vector enum_intersection; - for (const auto & p : enum_values) { - if (p.second == schema["allOf"].size()) { - enum_intersection.push_back(p.first); - } - } - if (!enum_intersection.empty()) { - return _add_rule(rule_name, "(" + string_join(enum_intersection, " | ") + ") space"); - } - } - return _add_rule(rule_name, _build_object_rule(properties, required, hybrid_name, json())); - } - if ((schema_type.is_null() || schema_type == "array") && (schema.contains("items") || schema.contains("prefixItems"))) { - json items = schema.contains("items") ? schema["items"] : schema["prefixItems"]; - if (items.is_array()) { - std::string rule = "\"[\" space "; - for (size_t i = 0; i < items.size(); i++) { - if (i > 0) { - rule += " \",\" space "; - } - rule += visit(items[i], name + (name.empty() ? "" : "-") + "tuple-" + std::to_string(i)); - } - rule += " \"]\" space"; - return _add_rule(rule_name, rule); - } - std::string item_rule_name = visit(items, name + (name.empty() ? "" : "-") + "item"); - int min_items = schema.contains("minItems") ? schema["minItems"].get() : 0; - json max_items_json = schema.contains("maxItems") ? schema["maxItems"] : json(); - int max_items = max_items_json.is_number_integer() ? max_items_json.get() : std::numeric_limits::max(); - - return _add_rule(rule_name, "\"[\" space " + build_repetition(item_rule_name, min_items, max_items, "\",\" space") + " \"]\" space"); - } - if ((schema_type.is_null() || schema_type == "string") && schema.contains("pattern")) { - return _visit_pattern(schema["pattern"], rule_name); - } - if ((schema_type.is_null() || schema_type == "string") && std::regex_match(schema_format, std::regex("^uuid[1-5]?$"))) { - return _add_primitive(rule_name == "root" ? "root" : schema_format, PRIMITIVE_RULES.at("uuid")); - } - if ((schema_type.is_null() || schema_type == "string") && STRING_FORMAT_RULES.find(schema_format + "-string") != STRING_FORMAT_RULES.end()) { - auto prim_name = schema_format + "-string"; - return _add_rule(rule_name, _add_primitive(prim_name, STRING_FORMAT_RULES.at(prim_name))); - } - if (schema_type == "string" && (schema.contains("minLength") || schema.contains("maxLength"))) { - std::string char_rule = _add_primitive("char", PRIMITIVE_RULES.at("char")); - int min_len = schema.contains("minLength") ? schema["minLength"].get() : 0; - int max_len = schema.contains("maxLength") ? schema["maxLength"].get() : std::numeric_limits::max(); - return _add_rule(rule_name, "\"\\\"\" " + build_repetition(char_rule, min_len, max_len) + " \"\\\"\" space"); - } - if (schema_type == "integer" && (schema.contains("minimum") || schema.contains("exclusiveMinimum") || schema.contains("maximum") || schema.contains("exclusiveMaximum"))) { - int64_t min_value = std::numeric_limits::min(); - int64_t max_value = std::numeric_limits::max(); - if (schema.contains("minimum")) { - min_value = schema["minimum"].get(); - } else if (schema.contains("exclusiveMinimum")) { - min_value = schema["exclusiveMinimum"].get() + 1; - } - if (schema.contains("maximum")) { - max_value = schema["maximum"].get(); - } else if (schema.contains("exclusiveMaximum")) { - max_value = schema["exclusiveMaximum"].get() - 1; - } - std::stringstream out; - out << "("; - build_min_max_int(min_value, max_value, out); - out << ") space"; - return _add_rule(rule_name, out.str()); - } - if (schema.empty() || schema_type == "object") { - return _add_rule(rule_name, _add_primitive("object", PRIMITIVE_RULES.at("object"))); - } - if (schema_type.is_null() && schema.is_object()) { - // No type constraint and no recognized structural keywords (e.g. {"description": "..."}). - // Per JSON Schema semantics this is equivalent to {} and accepts any value. - return _add_rule(rule_name, _add_primitive("value", PRIMITIVE_RULES.at("value"))); - } - if (!schema_type.is_string() || PRIMITIVE_RULES.find(schema_type.get()) == PRIMITIVE_RULES.end()) { - _errors.push_back("Unrecognized schema: " + schema.dump()); - return ""; - } - // TODO: support minimum, maximum, exclusiveMinimum, exclusiveMaximum at least for zero - return _add_primitive(rule_name == "root" ? "root" : schema_type.get(), PRIMITIVE_RULES.at(schema_type.get())); - } - - void check_errors() { - if (!_errors.empty()) { - throw std::invalid_argument("JSON schema conversion failed:\n" + string_join(_errors, "\n")); - } - if (!_warnings.empty()) { - fprintf(stderr, "WARNING: JSON schema conversion was incomplete: %s\n", string_join(_warnings, "; ").c_str()); - } - } - - std::string format_grammar() { - std::stringstream ss; - for (const auto & kv : _rules) { - ss << kv.first << " ::= " << kv.second << '\n'; - } - return ss.str(); - } -}; - -// common_schema_info implementation (pimpl) - -common_schema_info::common_schema_info() - : impl_(std::make_unique( - [](const std::string &) { return json(); }, - false)) {} - -common_schema_info::~common_schema_info() = default; - -common_schema_info::common_schema_info(common_schema_info &&) noexcept = default; -common_schema_info & common_schema_info::operator=(common_schema_info &&) noexcept = default; - -void common_schema_info::resolve_refs(nlohmann::ordered_json & schema) { - impl_->resolve_refs(schema, ""); -} - -// Determines if a JSON schema can resolve to a string type through any path. -// Some models emit raw string values rather than JSON-encoded strings for string parameters. -// If any branch of the schema (via oneOf, anyOf, $ref, etc.) permits a string, this returns -// true, allowing callers to handle the value as a raw string for simplicity. -bool common_schema_info::resolves_to_string(const nlohmann::ordered_json & schema) { - std::unordered_set visited_refs; - - std::function check = [&](const json & s) -> bool { - if (!s.is_object()) { - return false; - } - - // Handle $ref - if (s.contains("$ref")) { - const std::string & ref = s["$ref"]; - if (visited_refs.find(ref) != visited_refs.end()) { - // Circular reference, assume not a string to be safe - return false; - } - visited_refs.insert(ref); - auto it = impl_->_refs.find(ref); - if (it != impl_->_refs.end()) { - return check(it->second); - } - return false; - } - - // Check type field - if (s.contains("type")) { - const json & schema_type = s["type"]; - if (schema_type.is_string()) { - if (schema_type == "string") { - return true; - } - } else if (schema_type.is_array()) { - // Type can be an array like ["string", "null"] - for (const auto & t : schema_type) { - if (t == "string") { - return true; - } - } - } - } - - // Check oneOf/anyOf - if any alternative can be a string - if (s.contains("oneOf")) { - for (const auto & alt : s["oneOf"]) { - if (check(alt)) { - return true; - } - } - } - if (s.contains("anyOf")) { - for (const auto & alt : s["anyOf"]) { - if (check(alt)) { - return true; - } - } - } - - // Check allOf - all components must be compatible with string type - if (s.contains("allOf")) { - bool all_string = true; - for (const auto & component : s["allOf"]) { - if (!check(component)) { - all_string = false; - break; - } - } - if (all_string) { - return true; - } - } - - // Check const - if the constant value is a string - if (s.contains("const")) { - if (s["const"].is_string()) { - return true; - } - } - - // Check enum - if any enum value is a string - if (s.contains("enum")) { - for (const auto & val : s["enum"]) { - if (val.is_string()) { - return true; - } - } - } - - // String-specific keywords imply string type - if (s.contains("pattern") || s.contains("minLength") || s.contains("maxLength")) { - return true; - } - - // Check format - many formats imply string - if (s.contains("format")) { - const std::string & fmt = s["format"]; - if (fmt == "date" || fmt == "time" || fmt == "date-time" || - fmt == "uri" || fmt == "email" || fmt == "hostname" || - fmt == "ipv4" || fmt == "ipv6" || fmt == "uuid" || - fmt.find("uuid") == 0) { - return true; - } - } - - return false; - }; - - return check(schema); -} - -std::string json_schema_to_grammar(const json & schema, bool force_gbnf) { -#ifdef LLAMA_USE_LLGUIDANCE - if (!force_gbnf) { - return "%llguidance {}\nstart: %json " + schema.dump(); - } -#else - (void)force_gbnf; -#endif // LLAMA_USE_LLGUIDANCE - return build_grammar([&](const common_grammar_builder & callbacks) { - auto copy = schema; - callbacks.resolve_refs(copy); - callbacks.add_schema("", copy); - }); -} - -std::string build_grammar(const std::function & cb, const common_grammar_options & options) { - common_schema_converter converter([&](const std::string &) { return json(); }, options.dotall); - common_grammar_builder builder { - /* .add_rule = */ [&](const std::string & name, const std::string & rule) { - return converter._add_rule(name, rule); - }, - /* .add_schema = */ [&](const std::string & name, const nlohmann::ordered_json & schema) { - return converter.visit(schema, name == "root" ? "" : name); - }, - /* .resolve_refs = */ [&](nlohmann::ordered_json & schema) { - converter.resolve_refs(schema, ""); - } - }; - cb(builder); - converter.check_errors(); - return converter.format_grammar(); -} diff --git a/common/json-schema-to-grammar.h b/common/json-schema-to-grammar.h deleted file mode 100644 index 240d64231154..000000000000 --- a/common/json-schema-to-grammar.h +++ /dev/null @@ -1,43 +0,0 @@ -#pragma once - -#include - -#include -#include -#include - -std::string json_schema_to_grammar(const nlohmann::ordered_json & schema, - bool force_gbnf = false); - -class common_schema_converter; - -// Probes a JSON schema to extract information about its structure and type constraints. -class common_schema_info { - std::unique_ptr impl_; - - public: - common_schema_info(); - ~common_schema_info(); - - common_schema_info(const common_schema_info &) = delete; - common_schema_info & operator=(const common_schema_info &) = delete; - common_schema_info(common_schema_info &&) noexcept; - common_schema_info & operator=(common_schema_info &&) noexcept; - - void resolve_refs(nlohmann::ordered_json & schema); - bool resolves_to_string(const nlohmann::ordered_json & schema); -}; - -struct common_grammar_builder { - std::function add_rule; - std::function add_schema; - std::function resolve_refs; -}; - -struct common_grammar_options { - bool dotall = false; -}; - -std::string gbnf_format_literal(const std::string & literal); - -std::string build_grammar(const std::function & cb, const common_grammar_options & options = {}); diff --git a/common/llguidance.cpp b/common/llguidance.cpp deleted file mode 100644 index d58f147a76ab..000000000000 --- a/common/llguidance.cpp +++ /dev/null @@ -1,258 +0,0 @@ -#include "sampling.h" -#include "log.h" - -#ifdef LLAMA_USE_LLGUIDANCE - -# include "llguidance.h" -# include - -struct llama_sampler_llg { - const llama_vocab * vocab; - std::string grammar_kind; - std::string grammar_data; - LlgTokenizer * tokenizer; - LlgMatcher * grammar; -}; - -static LlgMatcher * llama_sampler_llg_new(LlgTokenizer * tokenizer, const char * grammar_kind, - const char * grammar_data) { - LlgConstraintInit cinit; - llg_constraint_init_set_defaults(&cinit, tokenizer); - const char * log_level = getenv("LLGUIDANCE_LOG_LEVEL"); - if (log_level && *log_level) { - cinit.log_stderr_level = atoi(log_level); - } - auto c = llg_new_matcher(&cinit, grammar_kind, grammar_data); - if (llg_matcher_get_error(c)) { - LOG_ERR("llg error: %s\n", llg_matcher_get_error(c)); - llg_free_matcher(c); - return nullptr; - } - - return c; -} - -static const char * llama_sampler_llg_name(const llama_sampler * /*smpl*/) { - return "llguidance"; -} - -static void llama_sampler_llg_accept_impl(llama_sampler * smpl, llama_token token) { - auto * ctx = (llama_sampler_llg *) smpl->ctx; - if (ctx->grammar) { - llg_matcher_consume_token(ctx->grammar, token); - } -} - -static void llama_sampler_llg_apply(llama_sampler * smpl, llama_token_data_array * cur_p) { - auto * ctx = (llama_sampler_llg *) smpl->ctx; - if (ctx->grammar) { - const uint32_t * mask = llg_matcher_get_mask(ctx->grammar); - if (mask == nullptr) { - if (llg_matcher_compute_mask(ctx->grammar) == 0) { - mask = llg_matcher_get_mask(ctx->grammar); - } else { - LOG_ERR("llg error: %s\n", llg_matcher_get_error(ctx->grammar)); - llg_free_matcher(ctx->grammar); - ctx->grammar = nullptr; - return; - } - } - - for (size_t i = 0; i < cur_p->size; ++i) { - auto token = cur_p->data[i].id; - if ((mask[token / 32] & (1 << (token % 32))) == 0) { - cur_p->data[i].logit = -INFINITY; - } - } - } -} - -static void llama_sampler_llg_reset(llama_sampler * smpl) { - auto * ctx = (llama_sampler_llg *) smpl->ctx; - if (ctx->grammar) { - llg_matcher_reset(ctx->grammar); - } -} - -static llama_sampler * llama_sampler_llg_clone(const llama_sampler * smpl) { - const auto * ctx = (const llama_sampler_llg *) smpl->ctx; - - auto * result = llama_sampler_init_llg(ctx->vocab, nullptr, nullptr); - - // copy the state - { - auto * result_ctx = (llama_sampler_llg *) result->ctx; - - if (ctx->grammar) { - result_ctx->grammar_kind = ctx->grammar_kind; - result_ctx->grammar_data = ctx->grammar_data; - result_ctx->grammar = llg_clone_matcher(ctx->grammar); - result_ctx->tokenizer = llg_clone_tokenizer(ctx->tokenizer); - } - } - - return result; -} - -static void llama_sampler_llg_free(llama_sampler * smpl) { - const auto * ctx = (llama_sampler_llg *) smpl->ctx; - - if (ctx->grammar) { - llg_free_matcher(ctx->grammar); - llg_free_tokenizer(ctx->tokenizer); - } - - delete ctx; -} - -static llama_sampler_i llama_sampler_llg_i = { - /* .name = */ llama_sampler_llg_name, - /* .accept = */ llama_sampler_llg_accept_impl, - /* .apply = */ llama_sampler_llg_apply, - /* .reset = */ llama_sampler_llg_reset, - /* .clone = */ llama_sampler_llg_clone, - /* .free = */ llama_sampler_llg_free, - /* .backend_init = */ NULL, - /* .backend_accept = */ NULL, - /* .backend_apply = */ NULL, - /* .backend_set_input = */ NULL, -}; - -static size_t llama_sampler_llg_tokenize_fn(const void * user_data, const uint8_t * bytes, size_t bytes_len, - uint32_t * output_tokens, size_t output_tokens_len) { - const llama_vocab * vocab = (const llama_vocab *) user_data; - int r = 0; - try { - r = llama_tokenize(vocab, (const char *) bytes, bytes_len, (int32_t *) output_tokens, output_tokens_len, false, - true); - } catch (const std::exception & e) { - GGML_ABORT("llama_tokenize failed: %s\n", e.what()); - } - if (r < 0) { - return -r; - } - return r; -} - -static LlgTokenizer * llama_sampler_llg_new_tokenizer(const llama_vocab * vocab) { - // TODO store the tokenizer in the vocab somehow - static const llama_vocab * vocab_cache; - static LlgTokenizer * tokenizer_cache; - - if (vocab_cache == vocab) { - return llg_clone_tokenizer(tokenizer_cache); - } - - auto tok_eos = llama_vocab_eot(vocab); - if (tok_eos == LLAMA_TOKEN_NULL) { - tok_eos = llama_vocab_eos(vocab); - } - - size_t vocab_size = llama_vocab_n_tokens(vocab); - - auto token_lens = new uint32_t[vocab_size]; - // we typically have ~7 bytes per token; let's go on the safe side here - auto token_bytes_size = vocab_size * 16 + 1024 * 1024; - auto token_bytes = new uint8_t[token_bytes_size]; - - size_t offset = 0; - for (size_t i = 0; i < vocab_size; i++) { - size_t max_token = 1024; - if (token_bytes_size - offset < max_token) { - GGML_ABORT("token_bytes buffer too small\n"); - } - - llama_token token = i; - auto dp = (char *) token_bytes + offset; - auto size = llama_detokenize(vocab, &token, 1, dp, max_token, false, false); - if (size < 0) { - GGML_ABORT("llama_detokenize failed\n"); - } - if (size == 0) { - size = llama_detokenize(vocab, &token, 1, dp + 1, max_token - 1, false, true); - if (size < 0) { - GGML_ABORT("llama_detokenize failed\n"); - } - if (size != 0) { - *dp = '\xff'; // special token prefix marker - size += 1; - } - } - - token_lens[i] = size; - offset += size; - } - - LlgTokenizerInit tinit = { - /* .vocab_size = */ (uint32_t) vocab_size, - /* .tok_eos = */ (uint32_t) tok_eos, - /* .token_lens = */ token_lens, - /* .token_bytes = */ token_bytes, - /* .tokenizer_json = */ nullptr, - /* .tokenize_assumes_string = */ true, - /* .tokenize_fn = */ llama_sampler_llg_tokenize_fn, - /* .use_approximate_greedy_tokenize_fn = */ false, - /* .tokenize_user_data = */ vocab, - /* .slices = */ nullptr, - }; - - char error_buffer[1024]; - LlgTokenizer * tokenizer = llg_new_tokenizer(&tinit, error_buffer, sizeof(error_buffer)); - - delete[] token_bytes; - delete[] token_lens; - - if (tokenizer == nullptr) { - LOG_ERR("llg tokenizer error: %s\n", error_buffer); - return tokenizer; - } - - if (tokenizer_cache) { - llg_free_tokenizer(tokenizer_cache); - } - vocab_cache = vocab; - tokenizer_cache = tokenizer; - - return llg_clone_tokenizer(tokenizer_cache); -} - -llama_sampler * llama_sampler_init_llg(const llama_vocab * vocab, const char * grammar_kind, - const char * grammar_data) { - auto * ctx = new llama_sampler_llg; - - if (grammar_kind != nullptr && grammar_kind[0] != '\0') { - auto tokenizer = llama_sampler_llg_new_tokenizer(vocab); - *ctx = { - /* .vocab = */ vocab, - /* .grammar_kind = */ grammar_kind, - /* .grammar_data = */ grammar_data, - /* .tokenizer = */ tokenizer, - /* .grammar = */ llama_sampler_llg_new(tokenizer, grammar_kind, grammar_data), - }; - if (ctx->grammar) { - GGML_ASSERT(((size_t) llama_vocab_n_tokens(vocab) + 31) / 32 * 4 == - llg_matcher_get_mask_byte_size(ctx->grammar)); - } - } else { - *ctx = { - /* .vocab = */ vocab, - /* .grammar_kind = */ {}, - /* .grammar_data = */ {}, - /* .tokenizer = */ nullptr, - /* .grammar = */ nullptr, - }; - } - - return llama_sampler_init( - /* .iface = */ &llama_sampler_llg_i, - /* .ctx = */ ctx); -} - -#else - -llama_sampler * llama_sampler_init_llg(const llama_vocab *, const char *, const char *) { - LOG_WRN("llguidance (cmake -DLLAMA_LLGUIDANCE=ON) is not enabled"); - return nullptr; -} - -#endif // LLAMA_USE_LLGUIDANCE diff --git a/common/log.cpp b/common/log.cpp deleted file mode 100644 index 4b4159db46fb..000000000000 --- a/common/log.cpp +++ /dev/null @@ -1,453 +0,0 @@ -#include "common.h" -#include "log.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#if defined(_WIN32) -# include -# include -# define isatty _isatty -# define fileno _fileno -#else -# include -#endif // defined(_WIN32) - -int common_log_verbosity_thold = LOG_DEFAULT_LLAMA; - -int common_log_get_verbosity_thold(void) { - return common_log_verbosity_thold; -} - -void common_log_set_verbosity_thold(int verbosity) { - common_log_verbosity_thold = verbosity; -} - -static int64_t t_us() { - return std::chrono::duration_cast(std::chrono::system_clock::now().time_since_epoch()).count(); -} - -// colors -enum common_log_col : int { - COMMON_LOG_COL_DEFAULT = 0, - COMMON_LOG_COL_BOLD, - COMMON_LOG_COL_RED, - COMMON_LOG_COL_GREEN, - COMMON_LOG_COL_YELLOW, - COMMON_LOG_COL_BLUE, - COMMON_LOG_COL_MAGENTA, - COMMON_LOG_COL_CYAN, - COMMON_LOG_COL_WHITE, -}; - -// disable colors by default -static const char* g_col[] = { - "", - "", - "", - "", - "", - "", - "", - "", - "", -}; - -struct common_log_entry { - enum ggml_log_level level; - - bool prefix; - - int64_t timestamp; - - std::vector msg; - - // signals the worker thread to stop - bool is_end; - - void print(FILE * file = nullptr) const { - FILE * fcur = file; - if (!fcur) { - // stderr displays DBG messages only when their verbosity level is not higher than the threshold - // these messages will still be logged to a file - if (level == GGML_LOG_LEVEL_DEBUG && common_log_verbosity_thold < LOG_DEFAULT_DEBUG) { - return; - } - - fcur = stdout; - - if (level != GGML_LOG_LEVEL_NONE) { - fcur = stderr; - } - } - - if (level != GGML_LOG_LEVEL_NONE && level != GGML_LOG_LEVEL_CONT && prefix) { - if (timestamp) { - // [M.s.ms.us] - fprintf(fcur, "%s%d.%02d.%03d.%03d%s ", - g_col[COMMON_LOG_COL_BLUE], - (int) (timestamp / 1000000 / 60), - (int) (timestamp / 1000000 % 60), - (int) (timestamp / 1000 % 1000), - (int) (timestamp % 1000), - g_col[COMMON_LOG_COL_DEFAULT]); - } - - switch (level) { - case GGML_LOG_LEVEL_INFO: fprintf(fcur, "%sI %s", g_col[COMMON_LOG_COL_GREEN], g_col[COMMON_LOG_COL_DEFAULT]); break; - case GGML_LOG_LEVEL_WARN: fprintf(fcur, "%sW %s", g_col[COMMON_LOG_COL_MAGENTA], "" ); break; - case GGML_LOG_LEVEL_ERROR: fprintf(fcur, "%sE %s", g_col[COMMON_LOG_COL_RED], "" ); break; - case GGML_LOG_LEVEL_DEBUG: fprintf(fcur, "%sD %s", g_col[COMMON_LOG_COL_YELLOW], "" ); break; - default: - break; - } - } - - fprintf(fcur, "%s", msg.data()); - - if (level == GGML_LOG_LEVEL_WARN || level == GGML_LOG_LEVEL_ERROR || level == GGML_LOG_LEVEL_DEBUG) { - fprintf(fcur, "%s", g_col[COMMON_LOG_COL_DEFAULT]); - } - - fflush(fcur); - } -}; - -struct common_log { - // default capacity - will be expanded if needed - common_log() : common_log(256) {} - - common_log(size_t capacity) { - file = nullptr; - prefix = false; - timestamps = false; - running = false; - t_start = t_us(); - - // initial message size - will be expanded if longer messages arrive - entries.resize(capacity); - for (auto & entry : entries) { - entry.msg.resize(256); - } - - head = 0; - tail = 0; - - resume(); - } - - ~common_log() { - pause(); - if (file) { - fclose(file); - } - } - -private: - std::mutex mtx; - std::thread thrd; - std::condition_variable cv; - - FILE * file; - - bool prefix; - bool timestamps; - bool running; - - int64_t t_start; - - // ring buffer of entries - std::vector entries; - size_t head; - size_t tail; - - // worker thread copies into this - common_log_entry cur; - -public: - void add(enum ggml_log_level level, const char * fmt, va_list args) { - std::lock_guard lock(mtx); - - if (!running) { - // discard messages while the worker thread is paused - return; - } - - auto & entry = entries[tail]; - - { - // cannot use args twice, so make a copy in case we need to expand the buffer - va_list args_copy; - va_copy(args_copy, args); - -#if 1 - const size_t n = vsnprintf(entry.msg.data(), entry.msg.size(), fmt, args); - if (n >= entry.msg.size()) { - entry.msg.resize(n + 1); - vsnprintf(entry.msg.data(), entry.msg.size(), fmt, args_copy); - } -#else - // hack for bolding arguments - - std::stringstream ss; - for (int i = 0; fmt[i] != 0; i++) { - if (fmt[i] == '%') { - ss << LOG_COL_BOLD; - while (fmt[i] != ' ' && fmt[i] != ')' && fmt[i] != ']' && fmt[i] != 0) ss << fmt[i++]; - ss << LOG_COL_DEFAULT; - if (fmt[i] == 0) break; - } - ss << fmt[i]; - } - const size_t n = vsnprintf(entry.msg.data(), entry.msg.size(), ss.str().c_str(), args); - if (n >= entry.msg.size()) { - entry.msg.resize(n + 1); - vsnprintf(entry.msg.data(), entry.msg.size(), ss.str().c_str(), args_copy); - } -#endif - va_end(args_copy); - } - - entry.level = level; - entry.prefix = prefix; - entry.timestamp = 0; - if (timestamps) { - entry.timestamp = t_us() - t_start; - } - entry.is_end = false; - - tail = (tail + 1) % entries.size(); - if (tail == head) { - // expand the buffer - std::vector new_entries(2*entries.size()); - - size_t new_tail = 0; - - do { - new_entries[new_tail] = std::move(entries[head]); - - head = (head + 1) % entries.size(); - new_tail = (new_tail + 1); - } while (head != tail); - - head = 0; - tail = new_tail; - - for (size_t i = tail; i < new_entries.size(); i++) { - new_entries[i].msg.resize(256); - } - - entries = std::move(new_entries); - } - cv.notify_one(); - } - - void resume() { - std::lock_guard lock(mtx); - - if (running) { - return; - } - - running = true; - - thrd = std::thread([this]() { - while (true) { - { - std::unique_lock lock(mtx); - cv.wait(lock, [this]() { return head != tail; }); - cur = entries[head]; - - head = (head + 1) % entries.size(); - } - - if (cur.is_end) { - break; - } - - cur.print(); // stdout and stderr - - if (file) { - cur.print(file); - } - } - }); - } - - void pause() { - { - std::lock_guard lock(mtx); - - if (!running) { - return; - } - - running = false; - - // push an entry to signal the worker thread to stop - { - auto & entry = entries[tail]; - entry.is_end = true; - - tail = (tail + 1) % entries.size(); - } - cv.notify_one(); - } - - thrd.join(); - } - - void set_file(const char * path) { - pause(); - - if (file) { - fclose(file); - } - - if (path) { - file = fopen(path, "w"); - } else { - file = nullptr; - } - - resume(); - } - - void set_colors(bool colors) { - pause(); - - if (colors) { - g_col[COMMON_LOG_COL_DEFAULT] = LOG_COL_DEFAULT; - g_col[COMMON_LOG_COL_BOLD] = LOG_COL_BOLD; - g_col[COMMON_LOG_COL_RED] = LOG_COL_RED; - g_col[COMMON_LOG_COL_GREEN] = LOG_COL_GREEN; - g_col[COMMON_LOG_COL_YELLOW] = LOG_COL_YELLOW; - g_col[COMMON_LOG_COL_BLUE] = LOG_COL_BLUE; - g_col[COMMON_LOG_COL_MAGENTA] = LOG_COL_MAGENTA; - g_col[COMMON_LOG_COL_CYAN] = LOG_COL_CYAN; - g_col[COMMON_LOG_COL_WHITE] = LOG_COL_WHITE; - } else { - for (size_t i = 0; i < std::size(g_col); i++) { - g_col[i] = ""; - } - } - - resume(); - } - - void set_prefix(bool prefix) { - std::lock_guard lock(mtx); - - this->prefix = prefix; - } - - void set_timestamps(bool timestamps) { - std::lock_guard lock(mtx); - - this->timestamps = timestamps; - } -}; - -// -// public API -// - -struct common_log * common_log_init() { - return new common_log; -} - -struct common_log * common_log_main() { - // We intentionally leak (i.e. do not delete) the logger singleton because - // common_log destructor called at DLL teardown phase will cause hanging on Windows. - // OS will release resources anyway so it should not be a significant issue, - // though this design may cause logs to be lost if not flushed before the program exits. - // Refer to https://github.com/ggml-org/llama.cpp/issues/22142 for details. - static struct common_log * log; - static std::once_flag init_flag; - std::call_once(init_flag, [&]() { - log = new common_log; - // Set default to auto-detect colors - log->set_colors(tty_can_use_colors()); - }); - - return log; -} - -void common_log_pause(struct common_log * log) { - log->pause(); -} - -void common_log_resume(struct common_log * log) { - log->resume(); -} - -void common_log_free(struct common_log * log) { - delete log; -} - -void common_log_add(struct common_log * log, enum ggml_log_level level, const char * fmt, ...) { - va_list args; - va_start(args, fmt); - log->add(level, fmt, args); - va_end(args); -} - -void common_log_set_file(struct common_log * log, const char * file) { - log->set_file(file); -} - -void common_log_set_colors(struct common_log * log, log_colors colors) { - if (colors == LOG_COLORS_AUTO) { - log->set_colors(tty_can_use_colors()); - return; - } - - if (colors == LOG_COLORS_DISABLED) { - log->set_colors(false); - return; - } - - GGML_ASSERT(colors == LOG_COLORS_ENABLED); - log->set_colors(true); -} - -void common_log_set_prefix(struct common_log * log, bool prefix) { - log->set_prefix(prefix); -} - -void common_log_set_timestamps(struct common_log * log, bool timestamps) { - log->set_timestamps(timestamps); -} - -void common_log_flush(struct common_log * log) { - log->pause(); - log->resume(); -} - -static int common_get_verbosity(enum ggml_log_level level) { - switch (level) { - case GGML_LOG_LEVEL_DEBUG: return LOG_LEVEL_DEBUG; - case GGML_LOG_LEVEL_INFO: return LOG_LEVEL_INFO; - case GGML_LOG_LEVEL_WARN: return LOG_LEVEL_WARN; - case GGML_LOG_LEVEL_ERROR: return LOG_LEVEL_ERROR; - case GGML_LOG_LEVEL_CONT: return LOG_LEVEL_INFO; // same as INFO - case GGML_LOG_LEVEL_NONE: - default: - return LOG_LEVEL_OUTPUT; - } -} - -void common_log_default_callback(enum ggml_log_level level, const char * text, void * /*user_data*/) { - auto verbosity = common_get_verbosity(level); - if (verbosity <= common_log_verbosity_thold) { - common_log_add(common_log_main(), level, "%s", text); - } -} diff --git a/common/log.h b/common/log.h deleted file mode 100644 index 8ee3e6b13b91..000000000000 --- a/common/log.h +++ /dev/null @@ -1,123 +0,0 @@ -#pragma once - -#include "ggml.h" // for ggml_log_level - -#define LOG_CLR_TO_EOL "\033[K\r" -#define LOG_COL_DEFAULT "\033[0m" -#define LOG_COL_BOLD "\033[1m" -#define LOG_COL_RED "\033[31m" -#define LOG_COL_GREEN "\033[32m" -#define LOG_COL_YELLOW "\033[33m" -#define LOG_COL_BLUE "\033[34m" -#define LOG_COL_MAGENTA "\033[35m" -#define LOG_COL_CYAN "\033[36m" -#define LOG_COL_WHITE "\033[37m" - -#ifndef __GNUC__ -# define LOG_ATTRIBUTE_FORMAT(...) -#elif defined(__MINGW32__) && !defined(__clang__) -# define LOG_ATTRIBUTE_FORMAT(...) __attribute__((format(gnu_printf, __VA_ARGS__))) -#else -# define LOG_ATTRIBUTE_FORMAT(...) __attribute__((format(printf, __VA_ARGS__))) -#endif - -#define LOG_LEVEL_DEBUG 4 -#define LOG_LEVEL_INFO 3 -#define LOG_LEVEL_WARN 2 -#define LOG_LEVEL_ERROR 1 -#define LOG_LEVEL_OUTPUT 0 // output data from tools - -#define LOG_DEFAULT_DEBUG LOG_LEVEL_DEBUG -#define LOG_DEFAULT_LLAMA LOG_LEVEL_INFO - -enum log_colors { - LOG_COLORS_AUTO = -1, - LOG_COLORS_DISABLED = 0, - LOG_COLORS_ENABLED = 1, -}; - -// needed by the LOG_TMPL macro to avoid computing log arguments if the verbosity lower -// set via common_log_set_verbosity() -int common_log_get_verbosity_thold(void); - -void common_log_set_verbosity_thold(int verbosity); // not thread-safe - -void common_log_default_callback(enum ggml_log_level level, const char * text, void * user_data); - -// the common_log uses an internal worker thread to print/write log messages -// when the worker thread is paused, incoming log messages are discarded -struct common_log; - -struct common_log * common_log_init(); - -// Singleton, intentionally leaked to avoid Windows teardown hangs. -// Call common_log_flush() before exit if you want to ensure all logs are flushed. -struct common_log * common_log_main(); - -void common_log_pause (struct common_log * log); // pause the worker thread, not thread-safe -void common_log_resume(struct common_log * log); // resume the worker thread, not thread-safe -void common_log_free (struct common_log * log); - -LOG_ATTRIBUTE_FORMAT(3, 4) -void common_log_add(struct common_log * log, enum ggml_log_level level, const char * fmt, ...); - -// defaults: file = NULL, colors = false, prefix = false, timestamps = false -// -// regular log output: -// -// ggml_backend_metal_log_allocated_size: allocated buffer, size = 6695.84 MiB, ( 6695.91 / 21845.34) -// llm_load_tensors: ggml ctx size = 0.27 MiB -// llm_load_tensors: offloading 32 repeating layers to GPU -// llm_load_tensors: offloading non-repeating layers to GPU -// -// with prefix = true, timestamps = true, the log output will look like this: -// -// 0.00.035.060 D ggml_backend_metal_log_allocated_size: allocated buffer, size = 6695.84 MiB, ( 6695.91 / 21845.34) -// 0.00.035.064 I llm_load_tensors: ggml ctx size = 0.27 MiB -// 0.00.090.578 I llm_load_tensors: offloading 32 repeating layers to GPU -// 0.00.090.579 I llm_load_tensors: offloading non-repeating layers to GPU -// -// D - debug (stderr, V = LOG_DEFAULT_DEBUG) -// I - info (stdout, V = LOG_DEFAULT_INFO) -// W - warning (stderr, V = LOG_DEFAULT_WARN) -// E - error (stderr, V = LOG_DEFAULT_ERROR) -// O - output (stdout, V = LOG_DEFAULT_OUTPUT) -// - -void common_log_set_file (struct common_log * log, const char * file); // not thread-safe -void common_log_set_colors (struct common_log * log, log_colors colors); // not thread-safe -void common_log_set_prefix (struct common_log * log, bool prefix); // whether to output prefix to each log -void common_log_set_timestamps(struct common_log * log, bool timestamps); // whether to output timestamps in the prefix -void common_log_flush (struct common_log * log); // flush all pending log messages - -// helper macros for logging -// use these to avoid computing log arguments if the verbosity of the log is higher than the threshold -// -// for example: -// -// LOG_DBG("this is a debug message: %d\n", expensive_function()); -// -// this will avoid calling expensive_function() if LOG_DEFAULT_DEBUG > common_log_verbosity_thold -// - -#define LOG_TMPL(level, verbosity, ...) \ - do { \ - if ((verbosity) <= common_log_get_verbosity_thold()) { \ - common_log_add(common_log_main(), (level), __VA_ARGS__); \ - } \ - } while (0) - -#define LOG(...) LOG_TMPL(GGML_LOG_LEVEL_NONE, LOG_LEVEL_OUTPUT, __VA_ARGS__) -#define LOGV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_NONE, verbosity, __VA_ARGS__) - -#define LOG_DBG(...) LOG_TMPL(GGML_LOG_LEVEL_DEBUG, LOG_LEVEL_DEBUG, __VA_ARGS__) -#define LOG_INF(...) LOG_TMPL(GGML_LOG_LEVEL_INFO, LOG_LEVEL_INFO, __VA_ARGS__) -#define LOG_WRN(...) LOG_TMPL(GGML_LOG_LEVEL_WARN, LOG_LEVEL_WARN, __VA_ARGS__) -#define LOG_ERR(...) LOG_TMPL(GGML_LOG_LEVEL_ERROR, LOG_LEVEL_ERROR, __VA_ARGS__) -#define LOG_CNT(...) LOG_TMPL(GGML_LOG_LEVEL_CONT, LOG_LEVEL_INFO, __VA_ARGS__) // same as INFO - -#define LOG_INFV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_INFO, verbosity, __VA_ARGS__) -#define LOG_WRNV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_WARN, verbosity, __VA_ARGS__) -#define LOG_ERRV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_ERROR, verbosity, __VA_ARGS__) -#define LOG_DBGV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_DEBUG, verbosity, __VA_ARGS__) -#define LOG_CNTV(verbosity, ...) LOG_TMPL(GGML_LOG_LEVEL_CONT, verbosity, __VA_ARGS__) diff --git a/common/ngram-cache-rust/Cargo.toml b/common/ngram-cache-rust/Cargo.toml new file mode 100644 index 000000000000..ec4d84ec24a4 --- /dev/null +++ b/common/ngram-cache-rust/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "llama-common-ngram-cache-rust" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["staticlib", "rlib"] diff --git a/common/ngram-cache-rust/src/lib.rs b/common/ngram-cache-rust/src/lib.rs new file mode 100644 index 000000000000..efce23c569e9 --- /dev/null +++ b/common/ngram-cache-rust/src/lib.rs @@ -0,0 +1,355 @@ +use std::collections::HashMap; +use std::ffi::CStr; +use std::fs::File; +use std::io::{self, ErrorKind, Read}; +use std::os::raw::{c_char, c_void}; + +const LLAMA_NGRAM_MAX: usize = 4; +const LLAMA_NGRAM_STATIC: usize = 2; +const LLAMA_TOKEN_NULL: LlamaToken = -1; + +type LlamaToken = i32; +type TokenCount = i32; +type NgramCachePart = HashMap; +type NgramCache = HashMap; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +struct Ngram { + tokens: [LlamaToken; LLAMA_NGRAM_MAX], +} + +#[no_mangle] +pub extern "C" fn common_ngram_cache_rust_new() -> *mut c_void { + Box::into_raw(Box::new(NgramCache::new())) as *mut c_void +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_cache_rust_free(ptr: *mut c_void) { + if !ptr.is_null() { + drop(Box::from_raw(ptr as *mut NgramCache)); + } +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_cache_rust_load(path: *const c_char) -> *mut c_void { + if path.is_null() { + return std::ptr::null_mut(); + } + let Ok(path) = CStr::from_ptr(path).to_str() else { + return std::ptr::null_mut(); + }; + match load_cache(path) { + Ok(cache) => Box::into_raw(Box::new(cache)) as *mut c_void, + Err(_) => std::ptr::null_mut(), + } +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_cache_rust_update( + cache: *mut c_void, + ngram_min: usize, + ngram_max: usize, + input: *const LlamaToken, + input_len: usize, + nnew: usize, +) { + if cache.is_null() || input.is_null() { + return; + } + let input = std::slice::from_raw_parts(input, input_len); + update_cache(&mut *(cache as *mut NgramCache), ngram_min, ngram_max, input, nnew); +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_cache_rust_draft( + context: *const c_void, + dynamic: *const c_void, + static_cache: *const c_void, + input: *const LlamaToken, + input_len: usize, + n_draft: usize, + ngram_min: usize, + ngram_max: usize, + out: *mut LlamaToken, + out_len: usize, +) -> usize { + if context.is_null() + || dynamic.is_null() + || static_cache.is_null() + || input.is_null() + || out.is_null() + || input_len == 0 + { + return 0; + } + + let input = std::slice::from_raw_parts(input, input_len); + let mut draft = vec![input[input_len - 1]]; + draft_tokens( + input, + &mut draft, + n_draft.min(out_len), + ngram_min, + ngram_max, + &*(context as *const NgramCache), + &*(dynamic as *const NgramCache), + &*(static_cache as *const NgramCache), + ); + + let drafted = draft.len().saturating_sub(1).min(out_len); + std::ptr::copy_nonoverlapping(draft[1..].as_ptr(), out, drafted); + drafted +} + +fn update_cache( + cache: &mut NgramCache, + ngram_min: usize, + ngram_max: usize, + input: &[LlamaToken], + nnew: usize, +) { + let input_size = input.len(); + for ngram_size in ngram_min..=ngram_max.min(LLAMA_NGRAM_MAX) { + let i_start = input_size.saturating_sub(nnew).max(ngram_size); + for i in i_start..input_size { + let ngram = ngram_from_slice(&input[i - ngram_size..i]); + let token = input[i]; + *cache.entry(ngram).or_default().entry(token).or_default() += 1; + } + } +} + +fn token_from_combined(input: &[LlamaToken], draft: &[LlamaToken], index: usize) -> LlamaToken { + if index < input.len() { + input[index] + } else { + draft[1 + index - input.len()] + } +} + +const DRAFT_MIN_SAMPLE_SIZE_LAX: [i32; LLAMA_NGRAM_MAX] = [2, 2, 1, 1]; +const DRAFT_MIN_PERCENT_LAX: [i32; LLAMA_NGRAM_MAX] = [66, 50, 50, 50]; +const DRAFT_MIN_SAMPLE_SIZE_STRICT: [i32; LLAMA_NGRAM_MAX] = [4, 3, 2, 2]; +const DRAFT_MIN_PERCENT_STRICT: [i32; LLAMA_NGRAM_MAX] = [75, 66, 66, 66]; + +fn draft_tokens( + input: &[LlamaToken], + draft: &mut Vec, + n_draft: usize, + ngram_min: usize, + ngram_max: usize, + context: &NgramCache, + dynamic: &NgramCache, + static_cache: &NgramCache, +) { + assert_eq!(draft.len(), 1); + if input.len() < LLAMA_NGRAM_STATIC { + return; + } + + while draft.len() - 1 < n_draft { + let ngram_start_static = input.len() - LLAMA_NGRAM_STATIC + draft.len() - 1; + let mut ngram_static = Ngram { + tokens: [LLAMA_TOKEN_NULL; LLAMA_NGRAM_MAX], + }; + for j in ngram_start_static..ngram_start_static + LLAMA_NGRAM_STATIC { + ngram_static.tokens[j - ngram_start_static] = token_from_combined(input, draft, j); + } + let part_static = static_cache.get(&ngram_static).cloned().unwrap_or_default(); + + let mut ngrams_cd = Vec::new(); + for ngram_size in ngram_min..=ngram_max.min(LLAMA_NGRAM_MAX) { + if ngram_size > input.len() + draft.len() - 1 { + continue; + } + let ngram_start_cd = input.len() - ngram_size + draft.len() - 1; + let mut ngram_cd = Ngram { + tokens: [LLAMA_TOKEN_NULL; LLAMA_NGRAM_MAX], + }; + for j in ngram_start_cd..ngram_start_cd + ngram_size { + ngram_cd.tokens[j - ngram_start_cd] = token_from_combined(input, draft, j); + } + ngrams_cd.push(ngram_cd); + } + + let drafted = try_draft_primary( + context, + &ngrams_cd, + &part_static, + &DRAFT_MIN_SAMPLE_SIZE_LAX, + &DRAFT_MIN_PERCENT_LAX, + ) + .or_else(|| { + try_draft_primary( + dynamic, + &ngrams_cd, + &part_static, + &DRAFT_MIN_SAMPLE_SIZE_STRICT, + &DRAFT_MIN_PERCENT_STRICT, + ) + }) + .or_else(|| try_draft_static(static_cache, ngram_static)); + + let Some(token) = drafted else { + break; + }; + draft.push(token); + } +} + +fn try_draft_static(cache: &NgramCache, ngram: Ngram) -> Option { + let part = cache.get(&ngram)?; + let mut max_count = 0; + let mut sum_count = 0; + let mut max_token = LLAMA_TOKEN_NULL; + + for (&token, &count) in part { + if count > max_count { + max_token = token; + max_count = count; + } + sum_count += count; + } + + if sum_count < DRAFT_MIN_SAMPLE_SIZE_LAX[LLAMA_NGRAM_STATIC - 1] { + return None; + } + if 100 * max_count < DRAFT_MIN_PERCENT_LAX[LLAMA_NGRAM_STATIC - 1] * sum_count { + return None; + } + Some(max_token) +} + +fn try_draft_primary( + cache: &NgramCache, + ngrams: &[Ngram], + part_static: &NgramCachePart, + min_sample_size: &[i32; LLAMA_NGRAM_MAX], + min_percent: &[i32; LLAMA_NGRAM_MAX], +) -> Option { + for (i, ngram) in ngrams.iter().enumerate().rev() { + let Some(part_primary) = cache.get(ngram) else { + continue; + }; + + let mut max_count_primary = 0; + let mut max_count_static = 0; + let mut sum_count_primary = 0; + let mut max_token = LLAMA_TOKEN_NULL; + + for (&token, &count_primary) in part_primary { + let count_static = part_static.get(&token).map_or(1, |count| 100 * *count); + if count_primary * count_static > max_count_primary * max_count_static { + max_token = token; + max_count_primary = count_primary; + max_count_static = count_static; + } + sum_count_primary += count_primary; + } + + if sum_count_primary < min_sample_size[i] { + continue; + } + if 100 * max_count_primary < min_percent[i] * sum_count_primary { + continue; + } + return Some(max_token); + } + + None +} + +fn ngram_from_slice(slice: &[LlamaToken]) -> Ngram { + let mut tokens = [LLAMA_TOKEN_NULL; LLAMA_NGRAM_MAX]; + for (dst, src) in tokens.iter_mut().zip(slice.iter().copied()) { + *dst = src; + } + Ngram { tokens } +} + +fn load_cache(path: &str) -> io::Result { + let mut file = File::open(path)?; + let mut cache = NgramCache::new(); + + loop { + let Some(ngram) = read_ngram(&mut file)? else { + break; + }; + + let ntokens = read_i32(&mut file, "missing token count")?; + if ntokens <= 0 { + return Err(invalid_data("token count must be positive")); + } + + let mut part = NgramCachePart::new(); + for _ in 0..ntokens { + let token = read_i32(&mut file, "missing token")?; + let count = read_i32(&mut file, "missing count")?; + if count <= 0 { + return Err(invalid_data("count must be positive")); + } + part.insert(token, count); + } + cache.insert(ngram, part); + } + + Ok(cache) +} + +fn read_ngram(mut reader: impl Read) -> io::Result> { + let mut bytes = [0_u8; LLAMA_NGRAM_MAX * std::mem::size_of::()]; + let mut read = 0; + while read < bytes.len() { + match reader.read(&mut bytes[read..]) { + Ok(0) if read == 0 => return Ok(None), + Ok(0) => return Err(invalid_data("partial ngram record")), + Ok(n) => read += n, + Err(err) => return Err(err), + } + } + + let mut tokens = [0; LLAMA_NGRAM_MAX]; + for (token, chunk) in tokens.iter_mut().zip(bytes.chunks_exact(4)) { + *token = i32::from_ne_bytes(chunk.try_into().expect("chunk size")); + } + Ok(Some(Ngram { tokens })) +} + +fn read_i32(mut reader: impl Read, missing: &'static str) -> io::Result { + let mut bytes = [0_u8; 4]; + reader.read_exact(&mut bytes).map_err(|err| { + if err.kind() == ErrorKind::UnexpectedEof { + invalid_data(missing) + } else { + err + } + })?; + Ok(i32::from_ne_bytes(bytes)) +} + +fn invalid_data(message: &'static str) -> io::Error { + io::Error::new(ErrorKind::InvalidData, message) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn updates_and_drafts_from_context_cache() { + let mut context = NgramCache::new(); + let dynamic = NgramCache::new(); + let mut static_cache = NgramCache::new(); + let input = [1, 2, 3, 1, 2, 3, 1, 2]; + + update_cache(&mut context, 1, 4, &input, input.len()); + static_cache.insert( + Ngram { tokens: [1, 2, -1, -1] }, + HashMap::from([(3, 2)]), + ); + + let mut draft = vec![2]; + draft_tokens(&input, &mut draft, 1, 1, 4, &context, &dynamic, &static_cache); + assert_eq!(draft, vec![2, 3]); + } +} diff --git a/common/ngram-cache.cpp b/common/ngram-cache.cpp deleted file mode 100644 index dce54b364749..000000000000 --- a/common/ngram-cache.cpp +++ /dev/null @@ -1,285 +0,0 @@ -#include "ngram-cache.h" -#include "common.h" -#include "log.h" - -#include -#include -#include -#include -#include -#include - -void common_ngram_cache_update(common_ngram_cache & ngram_cache, int ngram_min, int ngram_max, - std::vector & inp, int nnew, bool print_progress) { - const int64_t t_start_ms = ggml_time_ms(); - const int64_t inp_size = inp.size(); - - const int64_t n_todo = inp_size * (ngram_max - ngram_min + 1); - int64_t n_done = 0; - - for (int64_t ngram_size = ngram_min; ngram_size <= ngram_max; ++ngram_size) { - const int64_t i_start = std::max(inp_size - nnew, ngram_size); - for (int64_t i = i_start; i < inp_size; ++i) { - const int64_t ngram_start = i - ngram_size; - common_ngram ngram(&inp[ngram_start], ngram_size); - const llama_token token = inp[i]; - - common_ngram_cache::iterator part_it = ngram_cache.find(ngram); - if (part_it == ngram_cache.end()) { - common_ngram_cache_part part; - part.emplace(token, 1); - ngram_cache.emplace(ngram, part); - } else { - common_ngram_cache_part::iterator token_count_it = part_it->second.find(token); - if (token_count_it == part_it->second.end()) { - part_it->second.emplace(token, 1); - } else { - token_count_it->second++; - } - } - ++n_done; - - if (print_progress && n_done % 10000000 == 0) { - const int64_t t_now_ms = ggml_time_ms(); - const int64_t eta_ms = (inp_size*(ngram_max-ngram_min+1) - n_done) * (t_now_ms - t_start_ms) / n_done; - const int64_t eta_min = eta_ms / (60*1000); - const int64_t eta_s = (eta_ms - 60*1000*eta_min) / 1000; - - fprintf(stderr, "%s: %" PRId64 "/%" PRId64 " done, ETA: %02" PRId64 ":%02" PRId64 "\n", __func__, n_done, n_todo, eta_min, eta_s); - } - } - } -} - -// Helper function to get a token from the combined, speculative sequence of inp and draft. -static llama_token get_token(const std::vector & inp, const std::vector & draft, const size_t i) { - return i < inp.size() ? inp[i] : draft[1 + i - inp.size()]; -} - -// If sample size or percentage are below these thresholds the draft is aborted early: -constexpr int draft_min_sample_size_lax[LLAMA_NGRAM_MAX] = { 2, 2, 1, 1}; -constexpr int draft_min_percent_lax[LLAMA_NGRAM_MAX] = {66, 50, 50, 50}; -constexpr int draft_min_sample_size_strict[LLAMA_NGRAM_MAX] = { 4, 3, 2, 2}; -constexpr int draft_min_percent_strict[LLAMA_NGRAM_MAX] = {75, 66, 66, 66}; - -// Helper function that tries to draft a token from only the static ngram cache: -static llama_token try_draft(common_ngram_cache & nc_static, const common_ngram ngram_static) { - common_ngram_cache::iterator part_static_it = nc_static.find(ngram_static); - if (part_static_it == nc_static.end()) { - return LLAMA_TOKEN_NULL; - } - const common_ngram_cache_part part_static = part_static_it->second; - - int max_count_static = 0; - int sum_count_static = 0; - llama_token max_token = LLAMA_TOKEN_NULL; - - for (std::pair token_count_static : part_static) { - const llama_token token = token_count_static.first; - const int32_t count_static = token_count_static.second; - - if (count_static > max_count_static) { - max_token = token; - max_count_static = count_static; - } - sum_count_static += count_static; - } - - if (sum_count_static < draft_min_sample_size_lax[LLAMA_NGRAM_STATIC-1]) { - return LLAMA_TOKEN_NULL; - } - if (100*max_count_static < draft_min_percent_lax[LLAMA_NGRAM_STATIC-1]*sum_count_static) { - return LLAMA_TOKEN_NULL; - } - return max_token; -} - -// Try to draft a token from primary cache (context/dynamic), validate with static cache: -static llama_token try_draft( - common_ngram_cache & nc_primary, const std::vector & ngrams_primary, common_ngram_cache_part & part_static, - const int * min_sample_size, const int * min_percent) { - - llama_token drafted_token = LLAMA_TOKEN_NULL; - - for (int i = ngrams_primary.size()-1; i >= 0 && drafted_token == LLAMA_TOKEN_NULL; --i) { - const common_ngram ngram_primary = ngrams_primary[i]; - - common_ngram_cache::iterator part_primary_it = nc_primary.find(ngram_primary); - if (part_primary_it == nc_primary.end()) { - continue; - } - const common_ngram_cache_part part_primary = part_primary_it->second; - - int max_count_primary = 0; - int max_count_static = 0; - int sum_count_primary = 0; - llama_token max_token = LLAMA_TOKEN_NULL; - - for (std::pair token_count_primary : part_primary) { - const llama_token token = token_count_primary.first; - - common_ngram_cache_part::iterator token_count_static_it = part_static.find(token); - - const int32_t count_primary = token_count_primary.second; - const int32_t count_static = token_count_static_it != part_static.end() ? 100*token_count_static_it->second : 1; - - if (count_primary*count_static > max_count_primary*max_count_static) { - max_token = token; - max_count_primary = count_primary; - max_count_static = count_static; - } - sum_count_primary += count_primary; - } - - if (sum_count_primary < min_sample_size[i]) { - continue; - } - if (100*max_count_primary < min_percent[i]*sum_count_primary) { - continue;; - } - drafted_token = max_token; - } - - return drafted_token; -} - -void common_ngram_cache_draft( - std::vector & inp, std::vector & draft, int n_draft, int ngram_min, int ngram_max, - common_ngram_cache & nc_context, common_ngram_cache & nc_dynamic, common_ngram_cache & nc_static -) { - GGML_ASSERT(draft.size() == 1); - const int inp_size = inp.size(); - - if (inp_size < LLAMA_NGRAM_STATIC) { - return; - } - - while ((int) draft.size()-1 < n_draft) { - llama_token drafted_token = LLAMA_TOKEN_NULL; - - const int ngram_start_static = inp_size-LLAMA_NGRAM_STATIC + draft.size()-1; - common_ngram ngram_static; - for (int j = ngram_start_static; j < ngram_start_static + LLAMA_NGRAM_STATIC; ++j) { - ngram_static.tokens[j-ngram_start_static] = get_token(inp, draft, j); - } - common_ngram_cache::iterator part_static_it = nc_static.find(ngram_static); - common_ngram_cache_part part_static; - if (part_static_it != nc_static.end()) { - part_static = part_static_it->second; - } - - // cd = context + dynamic - std::vector ngrams_cd; - for (int ngram_size_cd = ngram_min; ngram_size_cd <= ngram_max; ++ngram_size_cd) { - const int ngram_start_cd = inp_size-ngram_size_cd + draft.size()-1; - common_ngram ngram_cd; - for (int j = ngram_start_cd; j < ngram_start_cd + ngram_size_cd; ++j) { - ngram_cd.tokens[j-ngram_start_cd] = get_token(inp, draft, j); - } - ngrams_cd.push_back(ngram_cd); - } - if (drafted_token == LLAMA_TOKEN_NULL) { - drafted_token = try_draft(nc_context, ngrams_cd, part_static, draft_min_sample_size_lax, draft_min_percent_lax); - } - if (drafted_token == LLAMA_TOKEN_NULL) { - drafted_token = try_draft(nc_dynamic, ngrams_cd, part_static, draft_min_sample_size_strict, draft_min_percent_strict); - } - if (drafted_token == LLAMA_TOKEN_NULL) { - drafted_token = try_draft(nc_static, ngram_static); - } - - if (drafted_token == LLAMA_TOKEN_NULL) { - break; - } - - LOG_DBG(" - draft candidate: token=%d\n", drafted_token); - draft.push_back(drafted_token); - } -} - -void common_ngram_cache_save(common_ngram_cache & ngram_cache, const std::string & filename) { - std::ofstream file_out(filename, std::ios::binary); - for (std::pair item : ngram_cache) { - const common_ngram ngram = item.first; - common_ngram_cache_part token_counts = item.second; - GGML_ASSERT(!token_counts.empty()); - const int32_t ntokens = token_counts.size(); - GGML_ASSERT(ntokens > 0); - - file_out.write(reinterpret_cast(&ngram), sizeof(common_ngram)); - file_out.write(reinterpret_cast(&ntokens), sizeof(int32_t)); - for (std::pair item2 : token_counts) { - const llama_token token = item2.first; - const int32_t count = item2.second; - GGML_ASSERT(count > 0); - - file_out.write(reinterpret_cast(&token), sizeof(llama_token)); - file_out.write(reinterpret_cast(&count), sizeof(int32_t)); - } - } -} - -common_ngram_cache common_ngram_cache_load(const std::string & filename) { - std::ifstream hashmap_file(filename, std::ios::binary); - if (!hashmap_file) { - throw std::ifstream::failure("Unable to open file " + filename); - } - common_ngram_cache ngram_cache; - - common_ngram ngram; - int32_t ntokens; - llama_token token; - int32_t count; - - char * ngramc = reinterpret_cast(&ngram); - char * ntokensc = reinterpret_cast(&ntokens); - char * tokenc = reinterpret_cast(&token); - char * countc = reinterpret_cast(&count); - while(hashmap_file.read(ngramc, sizeof(common_ngram))) { - GGML_ASSERT(!hashmap_file.eof()); - GGML_ASSERT(hashmap_file.read(ntokensc, sizeof(int32_t))); - GGML_ASSERT(ntokens > 0); - common_ngram_cache_part token_counts; - - for (int i = 0; i < ntokens; ++i) { - GGML_ASSERT(!hashmap_file.eof()); - GGML_ASSERT(hashmap_file.read(tokenc, sizeof(llama_token))); - GGML_ASSERT(!hashmap_file.eof()); - GGML_ASSERT(hashmap_file.read(countc, sizeof(int32_t))); - GGML_ASSERT(count > 0); - token_counts.emplace(token, count); - } - - ngram_cache.emplace(ngram, token_counts); - } - GGML_ASSERT(hashmap_file.eof()); - - return ngram_cache; -} - -void common_ngram_cache_merge(common_ngram_cache & ngram_cache_target, common_ngram_cache & ngram_cache_add) { - for (std::pair ngram_part : ngram_cache_add) { - const common_ngram ngram = ngram_part.first; - common_ngram_cache_part part = ngram_part.second; - - common_ngram_cache::iterator part_merged_it = ngram_cache_target.find(ngram); - if (part_merged_it == ngram_cache_target.end()) { - ngram_cache_target.emplace(ngram, part); - continue; - } - - for (std::pair token_count : part) { - const llama_token token = token_count.first; - const int32_t count = token_count.second; - GGML_ASSERT(count > 0); - - common_ngram_cache_part::iterator token_count_merged_it = part_merged_it->second.find(token); - if (token_count_merged_it == part_merged_it->second.end()) { - part_merged_it->second.emplace(token, count); - continue; - } - - token_count_merged_it->second += count; - } - } -} diff --git a/common/ngram-cache.h b/common/ngram-cache.h deleted file mode 100644 index 6e7cfea966df..000000000000 --- a/common/ngram-cache.h +++ /dev/null @@ -1,101 +0,0 @@ -#pragma once - -#include "llama.h" - -#include -#include -#include - -#define LLAMA_NGRAM_MIN 1 -#define LLAMA_NGRAM_MAX 4 -#define LLAMA_NGRAM_STATIC 2 - -// Data structures to map n-grams to empirical token probabilities: - -struct common_ngram { - llama_token tokens[LLAMA_NGRAM_MAX]; - - common_ngram() { - for (int i = 0; i < LLAMA_NGRAM_MAX; ++i) { - tokens[i] = LLAMA_TOKEN_NULL; - } - } - - common_ngram(const llama_token * input, const int ngram_size) { - for (int i = 0; i < LLAMA_NGRAM_MAX; ++i) { - tokens[i] = i < ngram_size ? input[i] : LLAMA_TOKEN_NULL; - } - } - - bool operator==(const common_ngram & other) const { - for (int i = 0; i < LLAMA_NGRAM_MAX; ++i) { - if (tokens[i] != other.tokens[i]) { - return false; - } - } - return true; - } -}; - -struct common_token_hash_function { - size_t operator()(const llama_token token) const { - // see https://probablydance.com/2018/06/16/fibonacci-hashing-the-optimization-that-the-world-forgot-or-a-better-alternative-to-integer-modulo/ - return token * 11400714819323198485llu; - } -}; - -struct common_ngram_hash_function { - size_t operator()(const common_ngram & ngram) const { - size_t hash = common_token_hash_function{}(ngram.tokens[0]); - for (int i = 1; i < LLAMA_NGRAM_MAX; ++i) { - hash ^= common_token_hash_function{}(ngram.tokens[i]); - } - return hash; - } -}; - -// token -> number of times token has been seen -typedef std::unordered_map common_ngram_cache_part; - -// n-gram -> empirical distribution of following tokens -typedef std::unordered_map common_ngram_cache; - - -// Update an ngram cache with tokens. -// ngram_cache: the cache to modify. -// ngram_min/ngram_max: the min/max size of the ngrams to extract from inp_data. -// inp_data: the token sequence with which to update ngram_cache. -// nnew: how many new tokens have been appended to inp_data since the last call to this function. -// print_progress: whether to print progress to stderr. -// -// In order to get correct results inp_data can ONLY BE APPENDED TO. -// Changes in the middle need a complete rebuild. -void common_ngram_cache_update( - common_ngram_cache & ngram_cache, int ngram_min, int ngram_max, std::vector & inp_data, int nnew, bool print_progress); - -// Try to draft tokens from ngram caches. -// inp: the tokens generated so far. -// draft: the token sequence to draft. Expected to initially contain the previously sampled token. -// n_draft: maximum number of tokens to add to draft. -// ngram_min/gram_max: the min/max size of the ngrams in nc_context and nc_dynamic. -// nc_context: ngram cache based on current context. -// nc_dynamic: ngram cache based on previous user generations. -// nc_static: ngram cache generated from a large text corpus, used for validation. -void common_ngram_cache_draft( - std::vector & inp, std::vector & draft, int n_draft, int ngram_min, int ngram_max, - common_ngram_cache & nc_context, common_ngram_cache & nc_dynamic, common_ngram_cache & nc_static); - -// Save an ngram cache to a file. -// ngram_cache: the ngram cache to save. -// filename: the path under which to save the ngram cache. -void common_ngram_cache_save(common_ngram_cache & ngram_cache, const std::string & filename); - -// Load an ngram cache saved with common_ngram_cache_save. -// filename: the path from which to load the ngram cache. -// returns: an ngram cache containing the information saved to filename. -common_ngram_cache common_ngram_cache_load(const std::string & filename); - -// Merge two ngram caches. -// ngram_cache_target: the ngram cache to which to add the information from ngram_cache_add. -// ngram_cache_add: the ngram cache to add to ngram_cache_target. -void common_ngram_cache_merge(common_ngram_cache & ngram_cache_target, common_ngram_cache & ngram_cache_add); diff --git a/common/ngram-map-rust/Cargo.toml b/common/ngram-map-rust/Cargo.toml new file mode 100644 index 000000000000..982cc2ef6145 --- /dev/null +++ b/common/ngram-map-rust/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "llama-common-ngram-map-rust" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["staticlib", "rlib"] diff --git a/common/ngram-map-rust/src/lib.rs b/common/ngram-map-rust/src/lib.rs new file mode 100644 index 000000000000..17feed1b6f16 --- /dev/null +++ b/common/ngram-map-rust/src/lib.rs @@ -0,0 +1,453 @@ +use std::ffi::c_void; +use std::ptr; + +const LCG_FACTOR: u32 = 2_654_435_761; +const MAX_VALUES: usize = 4; +const HASH_MAP_SIZE: usize = 262_144; +const MAX_VALUE_COUNT: u16 = 16_380; + +type LlamaToken = i32; + +#[derive(Clone, Copy)] +struct MapValue { + value_idx: usize, + value_num: u16, + n_accepted: i16, +} + +impl Default for MapValue { + fn default() -> Self { + Self { + value_idx: 0, + value_num: 0, + n_accepted: -1, + } + } +} + +#[derive(Clone)] +struct MapKey { + key_idx: usize, + stat_idx: usize, + key_num: u16, + values: [MapValue; MAX_VALUES], +} + +pub struct NgramMap { + size_key: u16, + size_value: u16, + key_only: bool, + keys: Vec, + min_hits: u16, + size_last_begin: usize, + last_draft_created: bool, + last_draft_key_idx: usize, + last_draft_value_idx: u16, + idx_last_check: usize, + key_map: Vec, + key_map_last_idx: u32, +} + +impl NgramMap { + fn new(size_key: u16, size_value: u16, key_only: bool, min_hits: u16) -> Self { + Self { + size_key, + size_value, + key_only, + keys: Vec::new(), + min_hits, + size_last_begin: 0, + last_draft_created: false, + last_draft_key_idx: 0, + last_draft_value_idx: 0, + idx_last_check: 0, + key_map: vec![0; HASH_MAP_SIZE], + key_map_last_idx: 0, + } + } +} + +#[no_mangle] +pub extern "C" fn common_ngram_map_rust_new( + size_key: u16, + size_value: u16, + key_only: bool, + min_hits: u16, +) -> *mut c_void { + Box::into_raw(Box::new(NgramMap::new( + size_key, size_value, key_only, min_hits, + ))) as *mut c_void +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_map_rust_free(ptr: *mut c_void) { + if !ptr.is_null() { + drop(Box::from_raw(ptr as *mut NgramMap)); + } +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_map_rust_size_value(ptr: *const c_void) -> u16 { + if ptr.is_null() { + return 0; + } + (*(ptr as *const NgramMap)).size_value +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_map_rust_begin( + ptr: *mut c_void, + tokens: *const LlamaToken, + tokens_len: usize, +) { + if ptr.is_null() || tokens.is_null() { + return; + } + let map = &mut *(ptr as *mut NgramMap); + let tokens = std::slice::from_raw_parts(tokens, tokens_len); + map_begin(map, tokens); +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_map_rust_draft( + ptr: *mut c_void, + input: *const LlamaToken, + input_len: usize, + sampled: LlamaToken, + out: *mut LlamaToken, + out_len: usize, +) -> usize { + if ptr.is_null() || input.is_null() || out.is_null() { + return 0; + } + let map = &mut *(ptr as *mut NgramMap); + let input = std::slice::from_raw_parts(input, input_len); + let mut draft = Vec::new(); + map_draft(map, input, sampled, &mut draft); + let copied = draft.len().min(out_len); + ptr::copy_nonoverlapping(draft.as_ptr(), out, copied); + copied +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_map_rust_accept(ptr: *mut c_void, n_accepted: u16) { + if ptr.is_null() { + return; + } + map_accept(&mut *(ptr as *mut NgramMap), n_accepted); +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_simple_rust_draft( + size_ngram: u16, + size_mgram: u16, + tokens: *const LlamaToken, + tokens_len: usize, + sampled: LlamaToken, + out: *mut LlamaToken, + out_len: usize, +) -> usize { + if tokens.is_null() || out.is_null() { + return 0; + } + let tokens = std::slice::from_raw_parts(tokens, tokens_len); + let draft = simple_draft(size_ngram as usize, size_mgram as usize, tokens, sampled); + let copied = draft.len().min(out_len); + ptr::copy_nonoverlapping(draft.as_ptr(), out, copied); + copied +} + +fn hash(tokens: &[LlamaToken], start: usize, len: usize) -> u32 { + let mut hash = 0_u32; + for i in 0..len { + hash = hash + .wrapping_mul(LCG_FACTOR) + .wrapping_add(tokens[start + i] as u32); + } + hash +} + +fn simple_draft( + size_ngram: usize, + size_mgram: usize, + tokens: &[LlamaToken], + sampled: LlamaToken, +) -> Vec { + let cur_len = tokens.len(); + if cur_len <= size_ngram + size_mgram + 1 { + return Vec::new(); + } + + let mut pattern = Vec::with_capacity(size_ngram); + for j in cur_len - size_ngram + 1..cur_len { + pattern.push(tokens[j]); + } + pattern.push(sampled); + + let mut match_pos = 0; + for j in (1..cur_len - size_ngram).rev() { + if pattern.iter().enumerate().all(|(k, token)| tokens[j + k] == *token) { + match_pos = j; + break; + } + } + if match_pos == 0 { + return Vec::new(); + } + + let copy_max = size_mgram.min(cur_len - (match_pos + size_ngram)); + if copy_max < size_ngram { + return Vec::new(); + } + tokens[match_pos + size_ngram..match_pos + size_ngram + copy_max].to_vec() +} + +fn map_begin(map: &mut NgramMap, tokens: &[LlamaToken]) { + let size_begin = tokens.len(); + + if !map.key_map.is_empty() && size_begin < map.idx_last_check { + for value in &mut map.key_map { + if *value as usize >= map.size_last_begin { + *value = 0; + } + } + map.key_map_last_idx = map.size_last_begin.saturating_sub(1) as u32; + } + + if size_begin < map.idx_last_check && !map.keys.is_empty() { + for i in (0..map.keys.len()).rev() { + if map.keys[i].key_idx >= map.size_last_begin { + map.keys.remove(i); + continue; + } + if map.key_only { + continue; + } + for j in (0..MAX_VALUES).rev() { + if map.keys[i].values[j].value_idx >= map.size_last_begin { + for k in j..MAX_VALUES - 1 { + map.keys[i].values[k] = map.keys[i].values[k + 1]; + } + map.keys[i].values[MAX_VALUES - 1] = MapValue::default(); + } + } + if map.keys[i].values[0].value_idx == 0 { + map.keys.remove(i); + } + } + } + + map.idx_last_check = size_begin; + map.size_last_begin = size_begin; +} + +fn map_draft( + map: &mut NgramMap, + input: &[LlamaToken], + sampled: LlamaToken, + draft: &mut Vec, +) { + map.last_draft_created = false; + map.last_draft_key_idx = 0; + map.last_draft_value_idx = 0; + + let cur_len = input.len(); + let n = map.size_key as usize; + let m = map.size_value as usize; + if cur_len < 2 * n + m || cur_len >= u32::MAX as usize || map.idx_last_check > cur_len { + return; + } + map.idx_last_check = cur_len; + + let mut key_tokens = Vec::with_capacity(n); + for j in cur_len - n + 1..cur_len { + key_tokens.push(input[j]); + } + key_tokens.push(sampled); + + let mut match_pos = 0; + if map.size_last_begin > cur_len { + return; + } + if !map.key_map.is_empty() { + let idx_hash = hash(&key_tokens, 0, n) as usize % map.key_map.len(); + let idx_key = map.key_map[idx_hash] as usize; + if idx_key != 0 + && idx_key < cur_len - n - m - 1 + && (0..n).all(|k| input[idx_key + k] == key_tokens[k]) + { + match_pos = idx_key; + } + } + + if match_pos == 0 && map.size_last_begin > n + m + 1 { + let upper = map.size_last_begin - n - m - 1; + let lower = map.key_map_last_idx as usize; + for j in (lower + 1..=upper).rev() { + if (0..n).all(|k| input[j + k] == key_tokens[k]) { + match_pos = j; + break; + } + } + } + if match_pos == 0 { + let upper = cur_len - n - m - 1; + let lower = map.size_last_begin.max(map.key_map_last_idx as usize); + if upper > lower { + for j in (lower + 1..=upper).rev() { + if (0..n).all(|k| input[j + k] == key_tokens[k]) { + match_pos = j; + break; + } + } + } + } + + if !map.key_map.is_empty() { + if map.size_last_begin > n + m + 1 { + let upper = map.size_last_begin - n - m - 1; + let lower = map.key_map_last_idx as usize; + for j in (lower + 1..=upper).rev() { + let idx_hash = hash(input, j, n) as usize % map.key_map.len(); + if map.key_map[idx_hash] == 0 { + map.key_map[idx_hash] = j as u32; + } + } + } + let upper = cur_len - n - m - 1; + let lower = map.size_last_begin.max(map.key_map_last_idx as usize); + if upper > lower { + for j in (lower + 1..=upper).rev() { + let idx_hash = hash(input, j, n) as usize % map.key_map.len(); + if map.key_map[idx_hash] == 0 { + map.key_map[idx_hash] = j as u32; + } + } + } + map.key_map_last_idx = map.key_map_last_idx.max((cur_len - n - m - 1) as u32); + } + + if match_pos == 0 { + return; + } + + let key_offset = map + .keys + .iter() + .position(|key| (0..n).all(|j| input[key.key_idx + j] == key_tokens[j])); + let key_offset = match key_offset { + Some(idx) => idx, + None => { + let key = MapKey { + key_idx: match_pos, + stat_idx: 0, + key_num: 0, + values: [MapValue { + n_accepted: m as i16, + ..MapValue::default() + }; MAX_VALUES], + }; + map.keys.push(key); + map.keys.len() - 1 + } + }; + + let curr_key = &mut map.keys[key_offset]; + curr_key.key_num = curr_key.key_num.saturating_add(1).min(MAX_VALUE_COUNT); + + if map.key_only { + let n_draft_tokens = m.min(curr_key.values[0].n_accepted.max(0) as usize); + draft.extend_from_slice(&input[match_pos + n..match_pos + n + n_draft_tokens]); + map.last_draft_created = true; + map.last_draft_key_idx = key_offset; + map.last_draft_value_idx = 0; + return; + } + + if curr_key.key_num < map.min_hits { + return; + } + + for i in curr_key.stat_idx..=match_pos { + if !(0..n).all(|k| input[i + k] == key_tokens[k]) { + continue; + } + let idx_begin_value_key = i + n; + let mut idx_value = None; + for v in 0..MAX_VALUES { + let idx_begin_value_v = curr_key.values[v].value_idx; + if idx_begin_value_v == 0 { + curr_key.values[v].value_idx = idx_begin_value_key; + curr_key.values[v].value_num = 0; + curr_key.values[v].n_accepted = m as i16; + idx_value = Some(v); + break; + } + if (0..m).all(|j| input[idx_begin_value_key + j] == input[idx_begin_value_v + j]) { + idx_value = Some(v); + break; + } + } + if let Some(v) = idx_value { + curr_key.values[v].value_num = + curr_key.values[v].value_num.saturating_add(1).min(MAX_VALUE_COUNT); + } + } + curr_key.stat_idx = match_pos; + + let mut max_occur = 0; + let mut slot_max = 0; + for v in 0..MAX_VALUES { + if curr_key.values[v].value_num > max_occur { + max_occur = curr_key.values[v].value_num; + slot_max = v; + } + } + let mut sum_occur = 0_u32; + for v in 0..MAX_VALUES { + if v != slot_max { + sum_occur += curr_key.values[v].value_num as u32; + } + } + if sum_occur > 0 && (max_occur as u32) < 2 * sum_occur { + return; + } + + let n_draft_tokens = m.min(curr_key.values[slot_max].n_accepted.max(0) as usize); + draft.extend_from_slice(&input[match_pos + n..match_pos + n + n_draft_tokens]); + map.last_draft_created = true; + map.last_draft_key_idx = key_offset; + map.last_draft_value_idx = slot_max as u16; +} + +fn map_accept(map: &mut NgramMap, n_accepted: u16) { + if !map.last_draft_created { + return; + } + let key_idx = map.last_draft_key_idx; + let val_idx = map.last_draft_value_idx as usize; + if let Some(key) = map.keys.get_mut(key_idx) { + key.values[val_idx].n_accepted = n_accepted as i16; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn simple_draft_repeats_prior_suffix() { + let tokens = [1, 2, 3, 4, 9, 2, 3]; + assert_eq!(simple_draft(2, 2, &tokens, 4), vec![9, 2]); + } + + #[test] + fn key_only_map_drafts_after_repeated_key() { + let mut map = NgramMap::new(2, 2, true, 1); + let tokens = [1, 2, 7, 8, 3, 4, 7]; + map_begin(&mut map, &tokens); + let mut draft = Vec::new(); + map_draft(&mut map, &tokens, 8, &mut draft); + assert_eq!(draft, vec![3, 4]); + } +} diff --git a/common/ngram-map.cpp b/common/ngram-map.cpp deleted file mode 100644 index 8e3978f7ed02..000000000000 --- a/common/ngram-map.cpp +++ /dev/null @@ -1,530 +0,0 @@ -#include "common.h" -#include "log.h" -#include "ngram-map.h" - -#include -#include -#include -#include - -// prime number used for LCG hash function (32 bit), it is near (sqrt(5) - 1)/2 * 2^32. -#define LCG_FACTOR 2654435761UL - -// Compute the LCG hash of a n-gram of size len at offset start. -static uint32_t common_ngram_map_hash(const llama_tokens & tokens, size_t start, size_t len) { - uint32_t hash = 0; - for (size_t i = 0; i < len; ++i) { - hash = hash * LCG_FACTOR + tokens[start + i]; - } - return hash; -} - -// Print the values of a sublist of `llama_tokens & inp` to a string in the form [v0, v1, v2, ...]. -static std::string common_tokens_to_str(const llama_tokens & inp, size_t start, size_t length) { - std::ostringstream oss; - oss << '['; - for (size_t i = 0; i < length; ++i) { - if (i > 0) { - oss << ", "; - } - oss << inp[start + i]; - } - oss << ']'; - return oss.str(); -} - - -// n-gram simple -// - -/** - * Perform speculative generation using the model's own token history. - * Searches for a matching pattern in the token history and returns draft tokens. - * - * @param state Current state of this implementation - * @param tokens Token history to search in - * @param sampled Last sampled token - * @return Vector of draft tokens, empty if no matching pattern is found - */ -llama_tokens common_ngram_simple_draft( - const common_ngram_simple_config & config, - const llama_tokens & tokens, llama_token sampled) { - - // Simple implementation of self-speculative decoding without a draft model. - // - const size_t cur_len = tokens.size(); - - const size_t n_draft_min = config.size_ngram; // size of n-gram to lookup in token history - const size_t n_draft_max = config.size_mgram; // the m-gram following the found n-gram is used for draft - - // vector for tokens we want to verify. - // return empty vector if there is no match. - llama_tokens draft_tokens; - - // We need at least n_draft_min + n_draft_max + 1 tokens. - if (cur_len <= static_cast(n_draft_min + n_draft_max + 1)) { - return draft_tokens; - } - - // pattern search - llama_tokens pattern; - pattern.reserve(n_draft_min); - for (size_t j = cur_len - n_draft_min + 1; j < cur_len; ++j) { - pattern.push_back(tokens[j]); - } - pattern.push_back(sampled); // add the last token to the pattern - - size_t match_pos = 0; // we ignore position 0, position 0 == no match - // search backwards, but skip the current match (we are currently there) - for (size_t j = cur_len - n_draft_min - 1; j > 0; --j) { - bool match = true; - for (size_t k = 0; k < pattern.size(); ++k) { - if (tokens[j + k] != pattern[k]) { - match = false; - break; - } - } - if (match) { - match_pos = j; - break; - } - } - if (match_pos == 0) { - return draft_tokens; - } - - const size_t copy_max = std::min( - n_draft_max, - cur_len - (match_pos + n_draft_min) - ); - if (copy_max < n_draft_min) { - return draft_tokens; - } - LOG_DBG("%s: #tokens = %zu: found matching pattern at pos %zu, length %zu, draft length %zu\n", - __func__, cur_len, - match_pos, pattern.size(), copy_max); - - draft_tokens.reserve(copy_max); - for (size_t j = 0; j < copy_max; ++j) { - draft_tokens.push_back(tokens[match_pos + n_draft_min + j]); - } - return draft_tokens; -} - - -// n-gram map -// - -// maximum number of counted values of a ngram map value. -#define COMMON_NGRAM_MAX_VALUE_COUNT 16380 - -void common_ngram_map_begin( - common_ngram_map & map, const llama_tokens & tokens) { - size_t size_begin = tokens.size(); - - LOG_DBG("%s: begin, idx_last_draft=%zu, new begin=%zu, #keys=%zu\n", __func__, - map.idx_last_check, size_begin, map.keys.size()); - - size_t count_map_entries_upd = 0; - if (!map.key_map.empty() && size_begin < map.idx_last_check) { - if (map.show_key_map_stats) { - // Print statistics of hash map map_key. - size_t count_nonzero = 0; - uint32_t min_idx = UINT32_MAX; - uint32_t max_idx = 0; - for (size_t i = 0; i < map.key_map.size(); ++i) { - uint32_t key_idx = map.key_map[i]; - if (key_idx != 0) { - ++count_nonzero; - if (key_idx < min_idx) min_idx = key_idx; - if (key_idx > max_idx) max_idx = key_idx; - } - } - if (count_nonzero == 0) { - min_idx = 0; - } - LOG_INF("%s: key_map stats: entries=%zu, min_idx=%u, max_idx=%u, key_map_last_idx=%u\n", - __func__, count_nonzero, min_idx, max_idx, map.key_map_last_idx); - } - - // Update the map from hash to key index (clear outdated entries). - for (size_t i = 0; i < map.key_map.size(); ++i) { - uint32_t key_idx = map.key_map[i]; - if (key_idx >= map.size_last_begin) { - map.key_map[i] = 0; - count_map_entries_upd++; - } - } - map.key_map_last_idx = (map.size_last_begin > 0) ? map.size_last_begin - 1 : 0; - } - - if (size_begin < map.idx_last_check && !map.keys.empty()) { - // The next token generation will start at index size_begin. - // The tokens between map.size_last_begin and size_begin are no longer valid. - // - // Refresh map: Remove all entries with index >= map.size_last_begin. - size_t count_keys = map.keys.size(); - size_t count_keys_del = 0; - size_t count_values_del = 0; - for (int32_t i = map.keys.size() - 1; i >= 0; --i) { - common_ngram_map_key & key = map.keys[i]; - if (key.key_idx >= map.size_last_begin) { - // Delete the key. - LOG_DBG("%s: delete key %d at index %zu (>= size_last_begin=%zu)\n", __func__, i, key.key_idx, map.size_last_begin); - map.keys.erase(map.keys.begin() + i); - count_keys_del++; - continue; - } - if (map.key_only) { - continue; - } - - // Check the indices of the values. - for (int16_t j = COMMON_NGRAM_MAX_VALUES - 1; j >= 0; --j) { - common_ngram_map_value & value = key.values[j]; - if (value.value_idx >= map.size_last_begin) { - // Delete the value. - count_values_del++; - - // Move all values after this value to the left. - for (uint16_t k = j; k < COMMON_NGRAM_MAX_VALUES - 1; ++k) { - key.values[k] = key.values[k + 1]; - } - // Clear the last value. - key.values[COMMON_NGRAM_MAX_VALUES - 1].value_idx = 0; - key.values[COMMON_NGRAM_MAX_VALUES - 1].value_num = 0; - } - } - if (key.values[0].value_idx == 0) { - // No values left, delete the key. - LOG_DBG("%s: delete key %d at index %zu (no values left)\n", __func__, i, key.key_idx); - map.keys.erase(map.keys.begin() + i); - count_keys_del++; - } - } - - LOG_INF("%s: refresh map: idx_last_draft=%zu, new begin=%zu, #keys_checked=%zu, #keys_del=%zu, #values_del=%zu, #hashes_upd=%zu\n", __func__, - map.idx_last_check, size_begin, - count_keys, count_keys_del, count_values_del, count_map_entries_upd); - } - - map.idx_last_check = size_begin; - map.size_last_begin = size_begin; -} - -void common_ngram_map_draft(common_ngram_map & map, - const llama_tokens & inp, llama_token sampled, - llama_tokens & draft) { - // reset last key and value. - map.last_draft_created = false; - map.last_draft_key_idx = 0; - map.last_draft_value_idx = 0; - - const size_t cur_len = inp.size(); - const uint16_t n = map.size_key; - const uint16_t m = map.size_value; - if (cur_len < static_cast(2 * n + m)) { - return; - } - if (cur_len >= static_cast(UINT32_MAX)) { - // key_map uses uint32_t instead of size_t. - GGML_ABORT("%s: cur_len exceeds UINT32_MAX: %zu", __func__, cur_len); - } - - if (map.idx_last_check > cur_len) { - // Should not happen because of common_ngram_map_begin(). - GGML_ABORT("%s: map.idx_last_check > cur_len: %zu > %zu", __func__, map.idx_last_check, cur_len); - } - map.idx_last_check = cur_len; - - // search pattern, the key n-gram - std::vector key_tokens; - key_tokens.reserve(n); - for (size_t j = cur_len - n + 1; j < cur_len; ++j) { - key_tokens.push_back(inp[j]); - } - key_tokens.push_back(sampled); - - // search for the key in the map - size_t match_pos = 0; - if (map.size_last_begin > cur_len) { - GGML_ABORT("%s: map.size_last_begin > cur_len: %zu > %zu", __func__, map.size_last_begin, cur_len); - } - if (!map.key_map.empty()) { - // Search for the key in the map key_map from hash of ngrams to index of ngram. - uint32_t idx_hash = (common_ngram_map_hash(key_tokens, 0, n) % map.key_map.size()); - uint32_t idx_key = map.key_map[idx_hash]; - if (idx_key != 0 && idx_key < cur_len - n - m - 1) { - // Check if the key matches the key at idx_key (because of possible collisions). - bool match = true; - for (size_t k = 0; k < n; ++k) { - if (inp[idx_key + k] != key_tokens[k]) { - match = false; - break; - } - } - LOG_DBG("%s: key hash %x -> idx_key %d: match %d\n", __func__, idx_hash, idx_key, match ? 1 : 0); - if (match) { - match_pos = idx_key; - } - } - } - if (match_pos == 0 && map.size_last_begin > (size_t) (n + m + 1)) { - // Search for the key in [1, map.size_last_begin - n - m -1], descending. - for (size_t j = map.size_last_begin - n - m - 1; j > map.key_map_last_idx; --j) { - // Check if the key matches the key. - bool match = true; - for (size_t k = 0; k < n; ++k) { - if (inp[j + k] != key_tokens[k]) { - match = false; - break; - } - } - if (match) { - match_pos = j; - break; - } - } - } - if (match_pos == 0) { - // In case of a reasoning chat, the part after size_last_begin may be deleted/reordered later. - // - // Search in [size_last_begin, cur_len - n - m - 1], descending. - for (size_t j = cur_len - n - m - 1; j > map.size_last_begin && j > map.key_map_last_idx; --j) { - bool match = true; - for (size_t k = 0; k < n; ++k) { - if (inp[j + k] != key_tokens[k]) { - match = false; - break; - } - } - if (match) { - match_pos = j; - break; - } - } - } - if (match_pos > 0) { - LOG_DBG("%s: cur_len = %zu, n = %d, m = %d, sz_tkns = %zu, sampled = %d, match_pos = %zu\n", __func__, - cur_len, n, m, key_tokens.size(), sampled, match_pos); - } - - if (!map.key_map.empty()) { - // Add hashes of new ngrams in key_map. - // - // Use the same order as above. - if (map.size_last_begin > (size_t) (n + m + 1)) { - for (size_t j = map.size_last_begin - n - m - 1; j > map.key_map_last_idx; --j) { - // compute hash and store index of ngram at idx j in the map. - uint32_t idx_hash = (common_ngram_map_hash(inp, j, n) % map.key_map.size()); - if (map.key_map[idx_hash] == 0) { - map.key_map[idx_hash] = j; // collisions may occur - } - } - } - - for (size_t j = cur_len - n - m - 1; j > map.size_last_begin && j > map.key_map_last_idx; --j) { - // compute hash and store index of ngram at idx j in the map. - uint32_t idx_hash = (common_ngram_map_hash(inp, j, n) % map.key_map.size()); - if (map.key_map[idx_hash] == 0) { - map.key_map[idx_hash] = j; - } - } - map.key_map_last_idx = std::max(static_cast(cur_len - n - m - 1), map.key_map_last_idx); - } - - if (match_pos == 0) { - return; - } - - // We have a match, now we look for the statistics of the key. - size_t key_offset = map.keys.size(); // offset in the map - // We iterate through the std::vector map->keys. - for (size_t i = 0; i < map.keys.size(); ++i) { - bool match = true; - for (size_t j = 0; j < n; ++j) { - if (inp[map.keys[i].key_idx + j] != key_tokens[j]) { - match = false; - break; - } - } - if (match) { - key_offset = i; - break; - } - } - if (key_offset == map.keys.size()) { - // We create a new key-entry, it will get offset key_offset. - common_ngram_map_key new_key; - new_key.key_idx = match_pos; - new_key.stat_idx = 0; - new_key.key_num = 0; - for (int i = 0; i < COMMON_NGRAM_MAX_VALUES; ++i) { - new_key.values[i].value_num = 0; - new_key.values[i].n_accepted = m; - } - map.keys.push_back(new_key); - } - - // our key n-gram: - common_ngram_map_key & curr_key = map.keys[key_offset]; - - // update number of key hits - curr_key.key_num = (uint16_t) std::min((int) map.keys[key_offset].key_num + 1, - (int) COMMON_NGRAM_MAX_VALUE_COUNT); - - if (map.key_only) { - // simple mode: - // Fill in the draft with the m tokens following the key. - // We work with value values[0] only. - int n_draft_tokens = std::min((int) m, (int) curr_key.values[0].n_accepted); - - for (int i = 0; i < n_draft_tokens; ++i) { - draft.push_back(inp[match_pos + n + i]); - } - - LOG_DBG("%s: key_idx = %zu, key_offset = %zu, key_num = %d, draft.size = %zu\n", __func__, - curr_key.key_idx, key_offset, curr_key.key_num, draft.size()); - - map.last_draft_created = true; - map.last_draft_key_idx = key_offset; - map.last_draft_value_idx = 0; // value 0 is used for simple mode - return; - } - - if (curr_key.key_num < map.min_hits) { - // not enough hits to consider this a good draft - LOG_DBG("%s: key_offset = %zu, key_num = %d, min_hits = %d, no draft\n", __func__, - key_offset, curr_key.key_num, map.min_hits); - return; - } - - // complex mode: examine the different m-grams after this key n-gram. - // - - // determine all (max COMMON_NGRAM_MAX_VALUES) m-grams after the key n-gram. - for (size_t i = curr_key.stat_idx; i <= match_pos; ++i) { - // begins the key n-gram at index i? - bool match_key = true; - for (size_t k = 0; k < n; ++k) { - if (inp[i + k] != key_tokens[k]) { - match_key = false; - break; - } - } - if (!match_key) { - continue; - } - - // Do we haven a existing value m-gram or a new one after the key at index i? - size_t idx_begin_value_key = i + n; - int idx_value = -1; - for (int v = 0; v < COMMON_NGRAM_MAX_VALUES; ++v) { - size_t idx_begin_value_v = curr_key.values[v].value_idx; - if (idx_begin_value_v == 0) { - // We found an empty value slot => we found a new value m-gram after the key n-gram. - curr_key.values[v].value_idx = idx_begin_value_key; - curr_key.values[v].value_num = 0; - curr_key.values[v].n_accepted = m; - idx_value = v; - break; - } - bool match = true; - for (size_t j = 0; j < m; ++j) { - if (inp[idx_begin_value_key + j] != inp[idx_begin_value_v + j]) { - match = false; - break; - } - } - if (match) { - // We found an existing value m-gram after the key n-gram. - idx_value = v; - break; - } - } - if (idx_value >= 0) { - // We found a value m-gram of the key n-gram. - curr_key.values[idx_value].value_num = (uint16_t) std::min((int) curr_key.values[idx_value].value_num + 1, - (int) COMMON_NGRAM_MAX_VALUE_COUNT); - } - } - // the statistics are updated up to match_pos. - curr_key.stat_idx = match_pos; - - // Do we have a value we could use for the draft? - uint16_t max_occur = 0; - int slot_max = 0; - for (int v = 0; v < COMMON_NGRAM_MAX_VALUES; ++v) { - uint16_t curr_occur = curr_key.values[v].value_num; - if (curr_occur > max_occur) { - max_occur = curr_occur; - slot_max = v; - } - } - // What is sum of the other occurrences? - uint32_t sum_occur = 0; - for (int v = 0; v < COMMON_NGRAM_MAX_VALUES; ++v) { - if (v == slot_max) { - continue; - } - uint16_t curr_occur = curr_key.values[v].value_num; - sum_occur += curr_occur; - } - - LOG_INF("%s: key_offset = %zu, max_occur = %d, sum_occur = %d, slot_max = %d [%zu/%d, %zu/%d, %zu/%d, %zu/%d]\n", __func__, - key_offset, - max_occur, sum_occur, slot_max, - curr_key.values[0].value_idx, curr_key.values[0].value_num, - curr_key.values[1].value_idx, curr_key.values[1].value_num, - curr_key.values[2].value_idx, curr_key.values[2].value_num, - curr_key.values[3].value_idx, curr_key.values[3].value_num - ); - // Print the tokens of the four values (if idx != 0), use LOG_INF - for (int v = 0; v < COMMON_NGRAM_MAX_VALUES; ++v) { - if (curr_key.values[v].value_idx != 0) { - LOG_INF("%s: value[%d] = %s\n", __func__, v, common_tokens_to_str(inp, curr_key.values[v].value_idx, m).c_str()); - } - } - - if (sum_occur > 0 && max_occur < 2 * sum_occur) { - // The most frequent value is not much more frequent than the other values. - // We do not use the draft. - return; - } - - // We use the most frequent value values[slot_max] for the draft. - // Fill in the draft with the m tokens following the key. - int n_draft_tokens = std::min((int) m, (int) curr_key.values[slot_max].n_accepted); - - for (int i = 0; i < n_draft_tokens; ++i) { - draft.push_back(inp[match_pos + n + i]); - } - - LOG_INF("%s: key_offset = %zu, slot_max = %d, key_num = %d, draft.size = %zu\n", __func__, - key_offset, slot_max, - curr_key.key_num, draft.size()); - - map.last_draft_created = true; - map.last_draft_key_idx = key_offset; - map.last_draft_value_idx = slot_max; // value used for draft generation. -} - -void common_ngram_map_accept(common_ngram_map & map, uint16_t n_accepted) { - if (!map.last_draft_created) { - return; - } - - // find the key and its chosen value. - const size_t key_idx = map.last_draft_key_idx; - const size_t val_idx = map.last_draft_value_idx; - - // find key corresponding to key_idx. - common_ngram_map_key & curr_key = map.keys[key_idx]; - // find value corresponding to val_idx. - struct common_ngram_map_value & curr_value = curr_key.values[val_idx]; // value used for draft generation. - - // update the value statistics - LOG_DBG("common_ngram_map_send_accepted: n_accepted = %d, prev value_num = %d\n", - n_accepted, curr_value.n_accepted); - curr_value.n_accepted = n_accepted; -} diff --git a/common/ngram-map.h b/common/ngram-map.h deleted file mode 100644 index 97608ef170f2..000000000000 --- a/common/ngram-map.h +++ /dev/null @@ -1,115 +0,0 @@ -#pragma once -// -// common/ngram-map.h: structures used to manage a map from n-grams to a list of m-grams -// -// These structures are used to do a lookup of n-grams followed by m-grams in token history. -// -// There are two algorithms implemented: -// 1. ngram_simple: lookup of n-grams followed by m-grams in token history. -// 2. ngram_map: lookup of n-grams followed by m-grams in token history using a map. -// The map is a vector of key n-grams, and for each key n-gram there is a list of value m-grams. -// -// ref: https://github.com/ggml-org/llama.cpp/pull/18471 -// - -#include "llama.h" -#include "common.h" - -#include - -// n-gram simple -// - -// config of n-gram simple. -struct common_ngram_simple_config { - uint16_t size_ngram; // size of n-grams to lookup in self-mode - uint16_t size_mgram; // size of m-grams to draft in self-mode -}; - -// Searches for a n-gram in the history and checks whether a draft sequence should be generated. -llama_tokens common_ngram_simple_draft( - const common_ngram_simple_config & config, - const llama_tokens & tokens, llama_token sampled); - - -// n-gram map -// - -// maximum number of m-gram values stored for each key n-gram. -#define COMMON_NGRAM_MAX_VALUES 4 - -// number of entries in the (optional, size 0 to disable) map from ngram-hash to ngram-index. -#define COMMON_NGRAM_HASH_MAP_SIZE 262144 - -// statistics of a m-gram after a known n-gram -struct common_ngram_map_value { - size_t value_idx = 0; // index of value m-gram in token-history (0 if unused) - uint16_t value_num = 0; // number of occurrences of this value m-gram after the key n-gram (0 in an unused values-slot) - int16_t n_accepted = -1; // number of accepted tokens at last draft (-1 if unused) -}; - -// statistics of a n-gram -struct common_ngram_map_key { - size_t key_idx; // index of key n-gram in token-history - size_t stat_idx; // index of last token of statistics computation (key_num, values) - - uint16_t key_num; // number of occurrences of this key n-gram in token-history - common_ngram_map_value values[COMMON_NGRAM_MAX_VALUES]; // some known values after the key -}; - -// map from n-grams to following m-grams in token-history -struct common_ngram_map { - uint16_t size_key; // size of key n-grams - uint16_t size_value; // size of value m-grams - - bool key_only; // true if only key n-grams are used, no values. - - std::vector keys; // key n-grams which occur several times in token-history - uint16_t min_hits; // minimum number of key hits to consider a draft - - bool show_key_map_stats = false; // true, if statistics of the key_map should be printed. - - common_ngram_map(uint16_t sz_key, uint16_t sz_value, bool only_keys, - uint16_t min_hits) - : size_key(sz_key), size_value(sz_value), key_only(only_keys), - min_hits(min_hits) { - key_map.resize(COMMON_NGRAM_HASH_MAP_SIZE); // 2^18 hash entries, 0 entries if key_map shouldn't be used - } - - // In reasoning chats the previous reasoning block will be removed from context history. - // A rebuild of the ngram map is needed after that. - - size_t size_last_begin = 0; // number of tokens at previous start of generation - - bool last_draft_created = false; // true if a draft was created at last call. - size_t last_draft_key_idx = 0; // index of last key used for draft generation (0 = no draft) - uint16_t last_draft_value_idx = 0; // index of last value used for draft generation. - - size_t idx_last_check = 0; // index of last check in context history - - // optional map "hash to ngram-index" for faster lookup of n-grams. map is empty if unused. - // - // uint32_t instead of size_t (size of current histories is << UINT32_MAX) - std::vector key_map; // key_map[hash] = index of ngram in context window - uint32_t key_map_last_idx = 0; // index of the last ngram added to key_map -}; - -// Initialize the n-gram map with the given token history. -// map: the ngram map to initialize. -// tokens: the token history to base the map on. -void common_ngram_map_begin( - common_ngram_map & map, - const llama_tokens & tokens); - -// Searches for the n-gram in the history and checks whether a draft sequence should be generated. -// map: the ngram map to search in. -// inp: the tokens generated so far. -// sampled: the token that was just sampled. -// draft: vector to store the draft tokens, initially empty. -void common_ngram_map_draft( - common_ngram_map & map, - const llama_tokens & inp, llama_token sampled, - llama_tokens & draft); - -// Update the statistics of a value after a draft was processed. -void common_ngram_map_accept(common_ngram_map & map, uint16_t n_accepted); diff --git a/common/ngram-mod-rust/Cargo.toml b/common/ngram-mod-rust/Cargo.toml new file mode 100644 index 000000000000..6ba874809f11 --- /dev/null +++ b/common/ngram-mod-rust/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "llama-common-ngram-mod-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[lib] +crate-type = ["staticlib", "rlib"] + +[dependencies] diff --git a/common/ngram-mod-rust/src/lib.rs b/common/ngram-mod-rust/src/lib.rs new file mode 100644 index 000000000000..220e4dee3286 --- /dev/null +++ b/common/ngram-mod-rust/src/lib.rs @@ -0,0 +1,138 @@ +use std::slice; + +pub const EMPTY: i32 = -1; +const HASH_MULTIPLIER: usize = 6_364_136_223_846_793_005; + +pub struct NgramMod { + n: usize, + used: usize, + entries: Vec, +} + +impl NgramMod { + fn new(n: u16, size: usize) -> Self { + Self { + n: n as usize, + used: 0, + entries: vec![EMPTY; size], + } + } + + fn idx(&self, tokens: &[i32]) -> usize { + let mut result = 0_usize; + for token in &tokens[..self.n] { + result = result + .wrapping_mul(HASH_MULTIPLIER) + .wrapping_add(*token as usize); + } + result % self.entries.len() + } + + fn add(&mut self, tokens: &[i32]) { + let idx = self.idx(tokens); + if self.entries[idx] == EMPTY { + self.used += 1; + } + self.entries[idx] = tokens[self.n]; + } + + fn get(&self, tokens: &[i32]) -> i32 { + self.entries[self.idx(tokens)] + } + + fn reset(&mut self) { + self.entries.fill(EMPTY); + self.used = 0; + } +} + +#[no_mangle] +pub extern "C" fn common_ngram_mod_rust_new(n: u16, size: usize) -> *mut NgramMod { + Box::into_raw(Box::new(NgramMod::new(n, size))) +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_mod_rust_free(ptr: *mut NgramMod) { + if !ptr.is_null() { + drop(Box::from_raw(ptr)); + } +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_mod_rust_idx( + ptr: *const NgramMod, + tokens: *const i32, +) -> usize { + let state = &*ptr; + state.idx(slice::from_raw_parts(tokens, state.n + 1)) +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_mod_rust_add(ptr: *mut NgramMod, tokens: *const i32) { + let state = &mut *ptr; + state.add(slice::from_raw_parts(tokens, state.n + 1)); +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_mod_rust_get( + ptr: *const NgramMod, + tokens: *const i32, +) -> i32 { + let state = &*ptr; + state.get(slice::from_raw_parts(tokens, state.n + 1)) +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_mod_rust_reset(ptr: *mut NgramMod) { + (&mut *ptr).reset(); +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_mod_rust_get_n(ptr: *const NgramMod) -> usize { + (&*ptr).n +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_mod_rust_get_used(ptr: *const NgramMod) -> usize { + (&*ptr).used +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_mod_rust_size(ptr: *const NgramMod) -> usize { + (&*ptr).entries.len() +} + +#[no_mangle] +pub unsafe extern "C" fn common_ngram_mod_rust_size_bytes(ptr: *const NgramMod) -> usize { + (&*ptr).entries.len() * std::mem::size_of::() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn adds_gets_and_resets() { + let mut ngram = NgramMod::new(2, 1024); + let tokens = [11, 22, 33]; + + assert_eq!(ngram.get(&tokens), EMPTY); + ngram.add(&tokens); + assert_eq!(ngram.get(&tokens), 33); + assert_eq!(ngram.used, 1); + + ngram.reset(); + assert_eq!(ngram.get(&tokens), EMPTY); + assert_eq!(ngram.used, 0); + } + + #[test] + fn replacing_same_slot_does_not_increase_used() { + let mut ngram = NgramMod::new(1, 16); + ngram.add(&[7, 8]); + ngram.add(&[7, 9]); + + assert_eq!(ngram.get(&[7, 0]), 9); + assert_eq!(ngram.used, 1); + } +} diff --git a/common/ngram-mod.cpp b/common/ngram-mod.cpp deleted file mode 100644 index 76f7257f6113..000000000000 --- a/common/ngram-mod.cpp +++ /dev/null @@ -1,60 +0,0 @@ -#include "ngram-mod.h" - -// -// common_ngram_mod -// - -common_ngram_mod::common_ngram_mod(uint16_t n, size_t size) : n(n), used(0) { - entries.resize(size); - - reset(); -} - -size_t common_ngram_mod::idx(const entry_t * tokens) const { - size_t res = 0; - - for (size_t i = 0; i < n; ++i) { - res = res*6364136223846793005ULL + tokens[i]; - } - - res = res % entries.size(); - - return res; -} - -void common_ngram_mod::add(const entry_t * tokens) { - const size_t i = idx(tokens); - - if (entries[i] == EMPTY) { - used++; - } - - entries[i] = tokens[n]; -} - -common_ngram_mod::entry_t common_ngram_mod::get(const entry_t * tokens) const { - const size_t i = idx(tokens); - - return entries[i]; -} - -void common_ngram_mod::reset() { - std::fill(entries.begin(), entries.end(), EMPTY); - used = 0; -} - -size_t common_ngram_mod::get_n() const { - return n; -} - -size_t common_ngram_mod::get_used() const { - return used; -} - -size_t common_ngram_mod::size() const { - return entries.size(); -} - -size_t common_ngram_mod::size_bytes() const { - return entries.size() * sizeof(entries[0]); -} diff --git a/common/ngram-mod.h b/common/ngram-mod.h deleted file mode 100644 index 7af92e9dde4a..000000000000 --- a/common/ngram-mod.h +++ /dev/null @@ -1,38 +0,0 @@ -#pragma once - -#include -#include -#include - -// -// common_ngram_mod -// ref: https://github.com/ggml-org/llama.cpp/pull/19164 -// - -// basic n-gram hasher -struct common_ngram_mod { - using entry_t = int32_t; - - static constexpr entry_t EMPTY = -1; - - common_ngram_mod(uint16_t n, size_t size); - - size_t idx(const entry_t * tokens) const; - void add(const entry_t * tokens); - entry_t get(const entry_t * tokens) const; // return -1 if not found - - void reset(); - - size_t get_n() const; - size_t get_used() const; - - size_t size() const; - size_t size_bytes() const; - -private: - size_t n; // ngram size to hash - - size_t used; - - std::vector entries; -}; diff --git a/common/peg-parser.cpp b/common/peg-parser.cpp deleted file mode 100644 index e37c1ce80e1d..000000000000 --- a/common/peg-parser.cpp +++ /dev/null @@ -1,2128 +0,0 @@ -#include "peg-parser.h" - -#include "common.h" -#include "json-schema-to-grammar.h" -#include "log.h" -#include "unicode.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -// Trick to catch missing branches -template -inline constexpr bool is_always_false_v = false; - -const char * common_peg_parse_result_type_name(common_peg_parse_result_type type) { - switch (type) { - case COMMON_PEG_PARSE_RESULT_FAIL: return "fail"; - case COMMON_PEG_PARSE_RESULT_SUCCESS: return "success"; - case COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT: return "need_more_input"; - default: return "unknown"; - } -} - -static bool is_hex_digit(const char c) { - return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F'); -} - -// Trie for matching multiple literals. -// This is used in common_peg_until_parser and to build a GBNF exclusion grammar -struct trie { - struct node { - std::map children; // Use uint32_t to store Unicode codepoints - bool is_word; - }; - - std::vector nodes; - - trie(const std::vector & words) { - create_node(); // root node - for (const auto & w : words) { - insert(w); - } - } - - enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH }; - - // Check if a delimiter starts at the given position - match_result check_at(std::string_view sv, size_t start_pos) const { - size_t current = 0; // Start at root - size_t pos = start_pos; - - // LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str()); - - while (pos < sv.size()) { - auto result = common_parse_utf8_codepoint(sv, pos); - if (result.status != utf8_parse_result::SUCCESS) { - break; - } - - auto it = nodes[current].children.find(result.codepoint); - if (it == nodes[current].children.end()) { - // Can't continue matching - return match_result{match_result::NO_MATCH}; - } - - current = it->second; - pos += result.bytes_consumed; - - // Check if we've matched a complete word - if (nodes[current].is_word) { - return match_result{match_result::COMPLETE_MATCH}; - } - } - - // Reached end of input while still in the trie (not at root) - if (current != 0) { - // We're in the middle of a potential match - return match_result{match_result::PARTIAL_MATCH}; - } - - // Reached end at root (no match) - return match_result{match_result::NO_MATCH}; - } - - struct prefix_and_next { - std::vector prefix; - std::vector next_chars; - }; - - std::vector collect_prefix_and_next() { - std::vector prefix; - std::vector result; - collect_prefix_and_next(0, prefix, result); - return result; - } - - private: - void collect_prefix_and_next(size_t index, std::vector & prefix, std::vector & out) { - if (!nodes[index].is_word) { - if (!nodes[index].children.empty()) { - std::vector chars; - chars.reserve(nodes[index].children.size()); - for (const auto & p : nodes[index].children) { - chars.push_back(p.first); - } - out.emplace_back(prefix_and_next{prefix, chars}); - } - } - - for (const auto & p : nodes[index].children) { - uint32_t ch = p.first; - auto child = p.second; - prefix.push_back(ch); - collect_prefix_and_next(child, prefix, out); - prefix.pop_back(); - } - } - - size_t create_node() { - size_t index = nodes.size(); - nodes.emplace_back(); - return index; - } - - void insert(const std::string & word) { - size_t current = 0; - size_t pos = 0; - while (pos < word.length()) { - auto result = common_parse_utf8_codepoint(word, pos); - if (result.status != utf8_parse_result::SUCCESS) { - break; - } - - uint32_t ch = result.codepoint; - pos += result.bytes_consumed; - - auto it = nodes[current].children.find(ch); - if (it == nodes[current].children.end()) { - size_t child = create_node(); - nodes[current].children[ch] = child; - current = child; - } else { - current = it->second; - } - } - nodes[current].is_word = true; - } -}; - -static std::pair parse_hex_escape(const std::string & str, size_t pos, int hex_count) { - if (pos + hex_count > str.length()) { - return {0, 0}; - } - - uint32_t value = 0; - for (int i = 0; i < hex_count; i++) { - char c = str[pos + i]; - if (!is_hex_digit(c)) { - return {0, 0}; - } - value <<= 4; - if ('a' <= c && c <= 'f') { - value += c - 'a' + 10; - } else if ('A' <= c && c <= 'F') { - value += c - 'A' + 10; - } else if ('0' <= c && c <= '9') { - value += c - '0'; - } else { - break; - } - } - return {value, static_cast(hex_count)}; -} - -static std::pair parse_char_class_char(const std::string & content, size_t pos) { - if (content[pos] == '\\' && pos + 1 < content.length()) { - switch (content[pos + 1]) { - case 'x': { - auto result = parse_hex_escape(content, pos + 2, 2); - if (result.second > 0) { - return {result.first, 2 + result.second}; - } - // Invalid escape, treat as literal 'x' - return {static_cast('x'), 2}; - } - case 'u': { - auto result = parse_hex_escape(content, pos + 2, 4); - if (result.second > 0) { - return {result.first, 2 + result.second}; - } - // Invalid escape, treat as literal 'u' - return {static_cast('u'), 2}; - } - case 'U': { - auto result = parse_hex_escape(content, pos + 2, 8); - if (result.second > 0) { - return {result.first, 2 + result.second}; - } - // Invalid escape, treat as literal 'U' - return {static_cast('U'), 2}; - } - case 'n': return {'\n', 2}; - case 't': return {'\t', 2}; - case 'r': return {'\r', 2}; - case '\\': return {'\\', 2}; - case ']': return {']', 2}; - case '[': return {'[', 2}; - default: return {static_cast(content[pos + 1]), 2}; - } - } - - // Regular character - return as codepoint - return {static_cast(static_cast(content[pos])), 1}; -} - -static std::pair, bool> parse_char_classes(const std::string & classes) { - std::vector ranges; - bool negated = false; - - std::string content = classes; - if (content.front() == '[') { - content = content.substr(1); - } - - if (content.back() == ']') { - content.pop_back(); - } - - // Check for negation - if (!content.empty() && content.front() == '^') { - negated = true; - content = content.substr(1); - } - - size_t i = 0; - while (i < content.length()) { - auto [start, start_len] = parse_char_class_char(content, i); - i += start_len; - - if (i + 1 < content.length() && content[i] == '-') { - // Range detected - auto [end, end_len] = parse_char_class_char(content, i + 1); - ranges.push_back(common_peg_chars_parser::char_range{start, end}); - i += 1 + end_len; - } else { - ranges.push_back(common_peg_chars_parser::char_range{start, start}); - } - } - - return {ranges, negated}; -} - -common_peg_ast_id common_peg_ast_arena::find_by_tag(const common_peg_ast_node & parent, const std::string & tag, int max_depth) const { - for (auto child_id : parent.children) { - const auto & child = get(child_id); - if (child.tag == tag) { - return child_id; - } - if (max_depth > 1) { - auto result = find_by_tag(child, tag, max_depth - 1); - if (result != COMMON_PEG_INVALID_AST_ID) { - return result; - } - } - } - return COMMON_PEG_INVALID_AST_ID; -} - -common_peg_ast_id common_peg_ast_arena::find_by_rule(const common_peg_ast_node & parent, const std::string & rule, int max_depth) const { - for (auto child_id : parent.children) { - const auto & child = get(child_id); - if (child.rule == rule) { - return child_id; - } - if (max_depth > 1) { - auto result = find_by_rule(child, rule, max_depth - 1); - if (result != COMMON_PEG_INVALID_AST_ID) { - return result; - } - } - } - return COMMON_PEG_INVALID_AST_ID; -} - -void common_peg_ast_arena::visit(common_peg_ast_id id, const common_peg_ast_visitor & visitor) const { - if (id == COMMON_PEG_INVALID_AST_ID) { - return; - } - const auto & node = get(id); - visitor(node); - for (const auto & child : node.children) { - visit(child, visitor); - } -} - -void common_peg_ast_arena::visit(const common_peg_parse_result & result, const common_peg_ast_visitor & visitor) const { - for (const auto & node : result.nodes) { - visit(node, visitor); - } -} - -struct parser_executor; - -common_peg_parser_id common_peg_arena::add_parser(common_peg_parser_variant parser) { - common_peg_parser_id id = parsers_.size(); - parsers_.push_back(std::move(parser)); - return id; -} - -void common_peg_arena::add_rule(const std::string & name, common_peg_parser_id id) { - rules_[name] = id; -} - -common_peg_parser_id common_peg_arena::get_rule(const std::string & name) const { - auto it = rules_.find(name); - if (it == rules_.end()) { - throw std::runtime_error("Rule not found: " + name); - } - return it->second; -} - -struct parser_executor { - const common_peg_arena & arena; - common_peg_parse_context & ctx; - size_t start_pos; - - parser_executor(const common_peg_arena & arena, common_peg_parse_context & ctx, size_t start) - : arena(arena), ctx(ctx), start_pos(start) {} - - std::string debug_indent() const { return std::string(ctx.parse_depth * 2, ' '); } - - std::string debug_input_snippet(size_t pos, size_t len = 60) const { - if (pos >= ctx.input.size()) { - return ""; - } - auto snippet = ctx.input.substr(pos, len); - // Escape newlines for display - std::string result; - for (char c : snippet) { - if (c == '\n') { - result += "\\n"; - } else if (c == '\r') { - result += "\\r"; - } else if (c == '\t') { - result += "\\t"; - } else { - result += c; - } - } - if (pos + len < ctx.input.size()) { - result += "..."; - } - return result; - } - - common_peg_parse_result operator()(const common_peg_epsilon_parser & /* p */) const { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos); - } - - common_peg_parse_result operator()(const common_peg_start_parser & /* p */) const { - return common_peg_parse_result( - start_pos == 0 ? COMMON_PEG_PARSE_RESULT_SUCCESS : COMMON_PEG_PARSE_RESULT_FAIL, - start_pos - ); - } - - common_peg_parse_result operator()(const common_peg_end_parser & /* p */) const { - return common_peg_parse_result( - start_pos >= ctx.input.size() ? COMMON_PEG_PARSE_RESULT_SUCCESS : COMMON_PEG_PARSE_RESULT_FAIL, - start_pos - ); - } - - common_peg_parse_result operator()(const common_peg_literal_parser & p) { - auto pos = start_pos; - for (auto i = 0u; i < p.literal.size(); ++i) { - if (pos >= ctx.input.size()) { - if (!ctx.is_lenient()) { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos); - } - if (ctx.input[pos] != p.literal[i]) { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); - } - ++pos; - } - - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); - } - - common_peg_parse_result operator()(const common_peg_sequence_parser & p) { - if (ctx.is_debug()) { - LOG_DBG("%sSEQ start at %zu '%s' (%zu children)\n", debug_indent().c_str(), start_pos, - debug_input_snippet(start_pos).c_str(), p.children.size()); - } - ctx.parse_depth++; - - auto pos = start_pos; - std::vector nodes; - - for (size_t i = 0; i < p.children.size(); i++) { - const auto & child_id = p.children[i]; - if (ctx.is_debug()) { - fprintf(stderr, "%sSEQ child %zu: %s\n", debug_indent().c_str(), i, arena.dump(child_id).c_str()); - } - auto result = arena.parse(child_id, ctx, pos); - - if (ctx.is_debug()) { - fprintf(stderr, "%sSEQ child %zu: %s at %zu->%zu\n", debug_indent().c_str(), i, - common_peg_parse_result_type_name(result.type), result.start, result.end); - } - - if (result.fail()) { - ctx.parse_depth--; - if (ctx.is_debug()) { - fprintf(stderr, "%sSEQ -> FAIL\n", debug_indent().c_str()); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos, result.end); - } - - if (!result.nodes.empty()) { - nodes.insert(nodes.end(), result.nodes.begin(), result.nodes.end()); - } - - if (result.need_more_input()) { - ctx.parse_depth--; - if (ctx.is_debug()) { - fprintf(stderr, "%sSEQ -> NEED_MORE\n", debug_indent().c_str()); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, result.end, std::move(nodes)); - } - - pos = result.end; - } - - ctx.parse_depth--; - if (ctx.is_debug()) { - fprintf(stderr, "%sSEQ -> SUCCESS at %zu->%zu\n", debug_indent().c_str(), start_pos, pos); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos, std::move(nodes)); - } - - common_peg_parse_result operator()(const common_peg_choice_parser & p) { - if (ctx.is_debug()) { - fprintf(stderr, "%sCHOICE start at %zu '%s' (%zu options)\n", debug_indent().c_str(), start_pos, - debug_input_snippet(start_pos).c_str(), p.children.size()); - } - ctx.parse_depth++; - - auto pos = start_pos; - for (size_t i = 0; i < p.children.size(); i++) { - const auto & child_id = p.children[i]; - if (ctx.is_debug()) { - fprintf(stderr, "%sCHOICE option %zu: %s\n", debug_indent().c_str(), i, arena.dump(child_id).c_str()); - } - auto result = arena.parse(child_id, ctx, pos); - if (ctx.is_debug()) { - fprintf(stderr, "%sCHOICE option %zu: %s\n", debug_indent().c_str(), i, - common_peg_parse_result_type_name(result.type)); - } - if (!result.fail()) { - ctx.parse_depth--; - if (ctx.is_debug()) { - fprintf(stderr, "%sCHOICE -> %s (option %zu)\n", debug_indent().c_str(), - common_peg_parse_result_type_name(result.type), i); - } - return result; - } - } - - ctx.parse_depth--; - if (ctx.is_debug()) { - fprintf(stderr, "%sCHOICE -> FAIL (no options matched)\n", debug_indent().c_str()); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); - } - - common_peg_parse_result operator()(const common_peg_repetition_parser & p) { - if (ctx.is_debug()) { - fprintf(stderr, "%sREPEAT start at %zu '%s' (min=%d, max=%d)\n", debug_indent().c_str(), start_pos, - debug_input_snippet(start_pos).c_str(), p.min_count, p.max_count); - } - ctx.parse_depth++; - - auto pos = start_pos; - int match_count = 0; - std::vector nodes; - - // Try to match up to max_count times (or unlimited if max_count is -1) - while (p.max_count == -1 || match_count < p.max_count) { - if (pos >= ctx.input.size()) { - if (ctx.is_debug()) { - fprintf(stderr, "%sREPEAT: at end of input, count=%d\n", debug_indent().c_str(), match_count); - } - break; - } - - auto result = arena.parse(p.child, ctx, pos); - - if (ctx.is_debug()) { - fprintf(stderr, "%sREPEAT iter %d: %s at %zu->%zu, nodes=%zu\n", debug_indent().c_str(), match_count, - common_peg_parse_result_type_name(result.type), result.start, result.end, result.nodes.size()); - fprintf(stderr, "%sREPEAT CHILD: %s\n", debug_indent().c_str(), arena.dump(p.child).c_str()); - } - - if (result.success()) { - // Prevent infinite loop on empty matches - if (result.end == pos) { - if (ctx.is_debug()) { - fprintf(stderr, "%s REPEAT: empty match, stopping\n", debug_indent().c_str()); - } - break; - } - - if (!result.nodes.empty()) { - nodes.insert(nodes.end(), result.nodes.begin(), result.nodes.end()); - } - - pos = result.end; - match_count++; - continue; - } - - if (result.need_more_input()) { - if (!result.nodes.empty()) { - nodes.insert(nodes.end(), result.nodes.begin(), result.nodes.end()); - } - - ctx.parse_depth--; - if (ctx.is_debug()) { - fprintf(stderr, "%sREPEAT -> NEED_MORE (count=%d, nodes=%zu)\n", debug_indent().c_str(), - match_count, nodes.size()); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, result.end, std::move(nodes)); - } - - // Child failed - stop trying - if (ctx.is_debug()) { - fprintf(stderr, "%sREPEAT: child failed, stopping\n", debug_indent().c_str()); - } - break; - } - - // Check if we got enough matches - if (p.min_count > 0 && match_count < p.min_count) { - ctx.parse_depth--; - if (pos >= ctx.input.size() && ctx.is_lenient()) { - if (ctx.is_debug()) { - fprintf(stderr, "%sREPEAT -> NEED_MORE (not enough matches: %d < %d)\n", debug_indent().c_str(), - match_count, p.min_count); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos, std::move(nodes)); - } - if (ctx.is_debug()) { - fprintf(stderr, "%sREPEAT -> FAIL (not enough matches: %d < %d)\n", debug_indent().c_str(), match_count, - p.min_count); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos, pos); - } - - ctx.parse_depth--; - if (ctx.is_debug()) { - fprintf(stderr, "%sREPEAT -> SUCCESS (count=%d, nodes=%zu)\n", debug_indent().c_str(), match_count, - nodes.size()); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos, std::move(nodes)); - } - - common_peg_parse_result operator()(const common_peg_and_parser & p) { - auto result = arena.parse(p.child, ctx, start_pos); - // Pass result but don't consume input - return common_peg_parse_result(result.type, start_pos); - } - - common_peg_parse_result operator()(const common_peg_not_parser & p) { - auto result = arena.parse(p.child, ctx, start_pos); - - if (result.success()) { - // Fail if the underlying parser matches - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); - } - - if (result.need_more_input()) { - // Propagate - need to know what child would match before negating - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos); - } - - // Child failed, so negation succeeds - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos); - } - - common_peg_parse_result operator()(const common_peg_any_parser & /* p */) const { - // Parse a single UTF-8 codepoint (not just a single byte) - auto result = common_parse_utf8_codepoint(ctx.input, start_pos); - - if (result.status == utf8_parse_result::INCOMPLETE) { - if (!ctx.is_lenient()) { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos); - } - if (result.status == utf8_parse_result::INVALID) { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, start_pos + result.bytes_consumed); - } - - common_peg_parse_result operator()(const common_peg_space_parser & /* p */) { - auto pos = start_pos; - while (pos < ctx.input.size()) { - auto c = static_cast(ctx.input[pos]); - if (std::isspace(c)) { - ++pos; - } else { - break; - } - } - - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); - } - - common_peg_parse_result operator()(const common_peg_chars_parser & p) const { - auto pos = start_pos; - int match_count = 0; - - // Try to match up to max_count times (or unlimited if max_count is -1) - while (p.max_count == -1 || match_count < p.max_count) { - auto result = common_parse_utf8_codepoint(ctx.input, pos); - - if (result.status == utf8_parse_result::INCOMPLETE) { - if (match_count >= p.min_count) { - // We have enough matches, succeed with what we have - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); - } - // Not enough matches yet - if (!ctx.is_lenient()) { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos); - } - - if (result.status == utf8_parse_result::INVALID) { - // Malformed UTF-8 in input - if (match_count >= p.min_count) { - // We have enough matches, succeed up to here - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); - } - // Not enough matches, fail - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); - } - - // Check if this codepoint matches our character class - bool matches = false; - for (const auto & range : p.ranges) { - if (range.contains(result.codepoint)) { - matches = true; - break; - } - } - - // If negated, invert the match result - if (p.negated) { - matches = !matches; - } - - if (matches) { - pos += result.bytes_consumed; - ++match_count; - } else { - // Character doesn't match, stop matching - break; - } - } - - // Check if we got enough matches - if (match_count < p.min_count) { - if (pos >= ctx.input.size() && ctx.is_lenient()) { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos, pos); - } - - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); - } - - static common_peg_parse_result handle_escape_sequence(common_peg_parse_context & ctx, size_t start, size_t & pos, const char delimiter) { - ++pos; // consume '\' - if (pos >= ctx.input.size()) { - if (!ctx.is_lenient()) { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start, pos); - } - - char c = ctx.input[pos]; - if (c == delimiter || c == '\\' || c == '/' || c == 'b' || c == 'f' || c == 'n' || c == 'r' || c == 't') { - ++pos; - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start, pos); - } else if (c == 'u') { - return handle_unicode_escape(ctx, start, pos); - } else { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); - } - } - - static common_peg_parse_result handle_unicode_escape(common_peg_parse_context & ctx, size_t start, size_t & pos) { - ++pos; // consume 'u' - for (int i = 0; i < 4; ++i) { - if (pos >= ctx.input.size()) { - if (!ctx.is_lenient()) { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start, pos); - } - if (!is_hex_digit(ctx.input[pos])) { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); - } - ++pos; - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start, pos); - } - - common_peg_parse_result operator()(const common_peg_string_parser & p) { - auto pos = start_pos; - - // Parse string content (without quotes) - while (pos < ctx.input.size()) { - char c = ctx.input[pos]; - - if (c == p.delimiter) { - // Found closing delimiter - success (don't consume it) - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); - } - - if (c == '\\') { - auto result = handle_escape_sequence(ctx, start_pos, pos, p.delimiter); - if (!result.success()) { - return result; - } - } else { - auto utf8_result = common_parse_utf8_codepoint(ctx.input, pos); - - if (utf8_result.status == utf8_parse_result::INCOMPLETE) { - if (!ctx.is_lenient()) { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos); - } - - if (utf8_result.status == utf8_parse_result::INVALID) { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); - } - - pos += utf8_result.bytes_consumed; - } - } - - // Reached end without finding closing quote - if (!ctx.is_lenient()) { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos, pos); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, pos); - } - - common_peg_parse_result operator()(const common_peg_until_parser & p) const { - trie matcher(p.delimiters); - - // Scan input and check for delimiters - size_t pos = start_pos; - size_t last_valid_pos = start_pos; - - while (pos < ctx.input.size()) { - auto utf8_result = common_parse_utf8_codepoint(ctx.input, pos); - - if (utf8_result.status == utf8_parse_result::INCOMPLETE) { - // Incomplete UTF-8 sequence - if (!ctx.is_lenient()) { - // Input is complete but UTF-8 is incomplete = malformed - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); - } - // Return what we have so far (before incomplete sequence) - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, last_valid_pos); - } - - if (utf8_result.status == utf8_parse_result::INVALID) { - // Malformed UTF-8 - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start_pos); - } - - // Check if a delimiter starts at this position - auto match = matcher.check_at(ctx.input, pos); - - if (match == trie::COMPLETE_MATCH) { - // Found a complete delimiter, return everything before it - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); - } - - if (match == trie::PARTIAL_MATCH) { - // Found a partial match extending to end of input, return everything before it - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); - } - - pos += utf8_result.bytes_consumed; - last_valid_pos = pos; - } - - if (last_valid_pos == ctx.input.size() && ctx.is_lenient()) { - // Reached the end of a partial stream, there might still be more input that we need to consume. - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start_pos, last_valid_pos); - } - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, last_valid_pos); - } - - common_peg_parse_result operator()(const common_peg_schema_parser & p) { - return arena.parse(p.child, ctx, start_pos); - } - - common_peg_parse_result operator()(const common_peg_rule_parser & p) { - // Parse the child - auto result = arena.parse(p.child, ctx, start_pos); - - if (!result.fail()) { - std::string_view text; - if (result.start < ctx.input.size()) { - text = std::string_view(ctx.input).substr(result.start, result.end - result.start); - } - - auto node_id = ctx.ast.add_node( - p.name, - "", - result.start, - result.end, - text, - std::move(result.nodes), - result.need_more_input() - ); - - return common_peg_parse_result(result.type, result.start, result.end, { node_id }); - } - - return result; - } - - common_peg_parse_result operator()(const common_peg_tag_parser & p) { - // Parse the child - if (ctx.is_debug()) { - fprintf(stderr, "%sTAG: %s\n", debug_indent().c_str(), p.tag.c_str()); - } - auto result = arena.parse(p.child, ctx, start_pos); - - if (!result.fail()) { - std::string_view text; - if (result.start < ctx.input.size()) { - text = std::string_view(ctx.input).substr(result.start, result.end - result.start); - } - - auto node_id = ctx.ast.add_node( - "", - p.tag, - result.start, - result.end, - text, - std::move(result.nodes), - result.need_more_input() - ); - - return common_peg_parse_result(result.type, result.start, result.end, { node_id }); - } - - return result; - } - - common_peg_parse_result operator()(const common_peg_ref_parser & p) { - auto rule_id = arena.get_rule(p.name); - return arena.parse(rule_id, ctx, start_pos); - } - - common_peg_parse_result operator()(const common_peg_atomic_parser & p) { - auto result = arena.parse(p.child, ctx, start_pos); - if (result.need_more_input()) { - // Clear nodes so they don't propagate up. - result.nodes.clear(); - } - return result; - } - - common_peg_parse_result operator()(const common_peg_gbnf_parser & p) { - return arena.parse(p.child, ctx, start_pos); - } -}; - -common_peg_parse_result common_peg_arena::parse(common_peg_parse_context & ctx, size_t start) const { - if (root_ == COMMON_PEG_INVALID_PARSER_ID) { - throw std::runtime_error("No root parser set"); - } - return parse(root_, ctx, start); -} - -common_peg_parse_result common_peg_arena::parse(common_peg_parser_id id, common_peg_parse_context & ctx, size_t start) const { - // Execute parser - const auto & parser = parsers_.at(id); - parser_executor exec(*this, ctx, start); - return std::visit(exec, parser); -} - -common_peg_parser_id common_peg_arena::resolve_ref(common_peg_parser_id id) { - const auto & parser = parsers_.at(id); - if (auto ref = std::get_if(&parser)) { - return get_rule(ref->name); - } - return id; -} - -static void bfs_node(common_peg_ast_arena &arena, std::ostringstream & oss, const common_peg_ast_node & node, int indent) { - for (int i = 0; i < indent; i++) { - oss << " "; - } - oss << "NODE " << node.id; - if (!node.rule.empty()) { - oss << " (rule " << node.rule << ")"; - } - if (!node.tag.empty()) { - oss << " (tag " << node.tag << ")"; - } - oss << " ['" << node.text << "']\n"; - for (const auto child : node.children) { - bfs_node(arena, oss, arena.get(child), indent + 1); - } -} - -std::string common_peg_ast_arena::dump() { - std::ostringstream oss; - for (auto & node : nodes_) { - bfs_node(*this, oss, node, 0); - } - return oss.str(); -} - -void common_peg_arena::resolve_refs() { - // Walk through all parsers and replace refs with their corresponding rule IDs - for (auto & parser : parsers_) { - std::visit([this](auto & p) { - using T = std::decay_t; - - if constexpr (std::is_same_v) { - for (auto & child : p.children) { - child = resolve_ref(child); - } - } else if constexpr (std::is_same_v) { - for (auto & child : p.children) { - child = resolve_ref(child); - } - } else 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) { - p.child = resolve_ref(p.child); - } else if constexpr (std::is_same_v) { - p.child = resolve_ref(p.child); - } else if constexpr (std::is_same_v) { - p.child = resolve_ref(p.child); - } else 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 || - std::is_same_v || - std::is_same_v) { - // These rules do not have children - } else { - static_assert(is_always_false_v); - } - }, parser); - } - - // Also flatten root if it's a ref - if (root_ != COMMON_PEG_INVALID_PARSER_ID) { - root_ = resolve_ref(root_); - } -} - -std::string common_peg_arena::dump(common_peg_parser_id id) const { - std::unordered_set visited; - return dump_impl(id, visited); -} - -std::string common_peg_arena::dump_impl(common_peg_parser_id id, - std::unordered_set & visited) const { - // Check for cycles - if (visited.count(id)) { - return "[cycle]"; - } - visited.insert(id); - - const auto & parser = parsers_.at(id); - - return std::visit([this, &visited](const auto & p) -> std::string { - using T = std::decay_t; - - if constexpr (std::is_same_v) { - return "Epsilon"; - } else if constexpr (std::is_same_v) { - return "Start"; - } else if constexpr (std::is_same_v) { - return "End"; - } else if constexpr (std::is_same_v) { - return "Literal(" + p.literal + ")"; - } else if constexpr (std::is_same_v) { - std::vector parts; - for (const auto & child : p.children) { - parts.push_back(dump_impl(child, visited)); - } - return "Sequence(" + string_join(parts, ", ") + ")"; - } else if constexpr (std::is_same_v) { - std::vector parts; - for (const auto & child : p.children) { - parts.push_back(dump_impl(child, visited)); - } - return "Choice(" + string_join(parts, ", ") + ")"; - } else if constexpr (std::is_same_v) { - if (p.max_count == -1) { - return "Repetition(" + dump_impl(p.child, visited) + ", " + std::to_string(p.min_count) + - ", unbounded)"; - } - return "Repetition(" + dump_impl(p.child, visited) + ", " + std::to_string(p.min_count) + ", " + std::to_string(p.max_count) + ")"; - } else if constexpr (std::is_same_v) { - return "And(" + dump_impl(p.child, visited) + ")"; - } else if constexpr (std::is_same_v) { - return "Not(" + dump_impl(p.child, visited) + ")"; - } else if constexpr (std::is_same_v) { - return "Atomic(" + dump_impl(p.child, visited) + ")"; - } else if constexpr (std::is_same_v) { - return "Gbnf(" + p.grammar + ", " + dump_impl(p.child, visited) + ")"; - } else if constexpr (std::is_same_v) { - return "Any"; - } else if constexpr (std::is_same_v) { - return "Space"; - } else if constexpr (std::is_same_v) { - if (p.max_count == -1) { - return "CharRepeat(" + p.pattern + ", " + std::to_string(p.min_count) + ", unbounded)"; - } - return "CharRepeat(" + p.pattern + ", " + std::to_string(p.min_count) + ", " + std::to_string(p.max_count) + ")"; - } else if constexpr (std::is_same_v) { - return "String(" + std::string(1, p.delimiter) + ")"; - } else if constexpr (std::is_same_v) { - return "Until(" + string_join(p.delimiters, " | ") + ")"; - } else if constexpr (std::is_same_v) { - return "Schema(" + dump_impl(p.child, visited) + ", " + (p.schema ? p.schema->dump() : "null") + ")"; - } else if constexpr (std::is_same_v) { - return "Rule(" + p.name + ", " + dump_impl(p.child, visited) + ")"; - } else if constexpr (std::is_same_v) { - return "Ref(" + p.name + ")"; - } else if constexpr (std::is_same_v) { - return "Tag(" + p.tag + ", " + dump(p.child) + ")"; - } else if constexpr (std::is_same_v) { - return "Atomic(" + dump(p.child) + ")"; - } else { - return "Unknown"; - } - }, parser); -} - -common_peg_parser & common_peg_parser::operator=(const common_peg_parser & other) { - id_ = other.id_; - return *this; -} - -common_peg_parser & common_peg_parser::operator+=(const common_peg_parser & other) { - id_ = builder_.sequence({id_, other.id_}); - return *this; -} - -common_peg_parser & common_peg_parser::operator|=(const common_peg_parser & other) { - id_ = builder_.choice({id_, other.id_}); - return *this; -} - -common_peg_parser common_peg_parser::operator+(const common_peg_parser & other) const { - return builder_.sequence({id_, other.id_}); -} - -common_peg_parser common_peg_parser::operator|(const common_peg_parser & other) const { - return builder_.choice({id_, other.id_}); -} - -common_peg_parser common_peg_parser::operator<<(const common_peg_parser & other) const { - return builder_.sequence({id_, builder_.space(), other.id_}); -} - -common_peg_parser common_peg_parser::operator+(const char * str) const { - return *this + builder_.literal(str); -} - -common_peg_parser common_peg_parser::operator+(const std::string & str) const { - return *this + builder_.literal(str); -} - -common_peg_parser common_peg_parser::operator<<(const char * str) const { - return *this << builder_.literal(str); -} - -common_peg_parser common_peg_parser::operator<<(const std::string & str) const { - return *this << builder_.literal(str); -} - -common_peg_parser common_peg_parser::operator|(const char * str) const { - return *this | builder_.literal(str); -} - -common_peg_parser common_peg_parser::operator|(const std::string & str) const { - return *this | builder_.literal(str); -} - -common_peg_parser operator+(const char * str, const common_peg_parser & p) { - return p.builder().literal(str) + p; -} - -common_peg_parser operator+(const std::string & str, const common_peg_parser & p) { - return operator+(str.c_str(), p); -} - -common_peg_parser operator<<(const char * str, const common_peg_parser & p) { - return p.builder().literal(str) << p; -} - -common_peg_parser operator<<(const std::string & str, const common_peg_parser & p) { - return operator<<(str.c_str(), p); -} - -common_peg_parser operator|(const char * str, const common_peg_parser & p) { - return p.builder().literal(str) | p; -} - -common_peg_parser operator|(const std::string & str, const common_peg_parser & p) { - return operator|(str.c_str(), p); -} - -static std::string rule_name(const std::string & name) { - static const std::regex invalid_rule_chars_re("[^a-zA-Z0-9-]+"); - return std::regex_replace(name, invalid_rule_chars_re, "-"); -} - -common_peg_parser_builder::common_peg_parser_builder() {} - -common_peg_parser common_peg_parser_builder::sequence(const std::vector & parsers) { - // Flatten nested sequences - std::vector flattened; - for (const auto & p : parsers) { - const auto & parser = arena_.get(p); - if (auto seq = std::get_if(&parser)) { - flattened.insert(flattened.end(), seq->children.begin(), seq->children.end()); - } else { - flattened.push_back(p); - } - } - return wrap(arena_.add_parser(common_peg_sequence_parser{flattened})); -} - -common_peg_parser common_peg_parser_builder::sequence(const std::vector & parsers) { - std::vector ids; - ids.reserve(parsers.size()); - for (const auto & p : parsers) { - ids.push_back(p.id()); - } - return sequence(ids); -} - -common_peg_parser common_peg_parser_builder::sequence(std::initializer_list parsers) { - std::vector ids; - ids.reserve(parsers.size()); - for (const auto & p : parsers) { - ids.push_back(p.id()); - } - return sequence(ids); -} - -common_peg_parser common_peg_parser_builder::choice(const std::vector & parsers) { - // Flatten nested choices - std::vector flattened; - for (const auto & p : parsers) { - const auto & parser = arena_.get(p); - if (auto choice = std::get_if(&parser)) { - flattened.insert(flattened.end(), choice->children.begin(), choice->children.end()); - } else { - flattened.push_back(p); - } - } - return wrap(arena_.add_parser(common_peg_choice_parser{flattened})); -} - -common_peg_parser common_peg_parser_builder::choice(const std::vector & parsers) { - std::vector ids; - ids.reserve(parsers.size()); - for (const auto & p : parsers) { - ids.push_back(p.id()); - } - return choice(ids); -} - -common_peg_parser common_peg_parser_builder::choice(std::initializer_list parsers) { - std::vector ids; - ids.reserve(parsers.size()); - for (const auto & p : parsers) { - ids.push_back(p.id()); - } - return choice(ids); -} - -common_peg_parser common_peg_parser_builder::chars(const std::string & classes, int min, int max) { - auto [ranges, negated] = parse_char_classes(classes); - return wrap(arena_.add_parser(common_peg_chars_parser{classes, ranges, negated, min, max})); -} - -common_peg_parser common_peg_parser_builder::schema(const common_peg_parser & p, const std::string & name, const nlohmann::ordered_json & schema, bool raw) { - return wrap(arena_.add_parser(common_peg_schema_parser{p.id(), name, std::make_shared(schema), raw})); -} - -common_peg_parser common_peg_parser_builder::rule(const std::string & name, const common_peg_parser & p, bool trigger) { - auto clean_name = rule_name(name); - auto rule_id = arena_.add_parser(common_peg_rule_parser{clean_name, p.id(), trigger}); - arena_.add_rule(clean_name, rule_id); - return ref(clean_name); -} - -common_peg_parser common_peg_parser_builder::rule(const std::string & name, const std::function & builder_fn, bool trigger) { - auto clean_name = rule_name(name); - if (arena_.has_rule(clean_name)) { - return ref(clean_name); - } - - // Create placeholder rule to allow recursive references - auto placeholder = any(); // Temporary placeholder - auto placeholder_rule_id = arena_.add_parser(common_peg_rule_parser{clean_name, placeholder.id(), trigger}); - arena_.add_rule(clean_name, placeholder_rule_id); - - // Build the actual parser - auto parser = builder_fn(); - - // Replace placeholder with actual rule - auto rule_id = arena_.add_parser(common_peg_rule_parser{clean_name, parser.id(), trigger}); - arena_.rules_[clean_name] = rule_id; - - return ref(clean_name); -} - -void common_peg_parser_builder::set_root(const common_peg_parser & p) { - arena_.set_root(p.id()); -} - -common_peg_arena common_peg_parser_builder::build() { - arena_.resolve_refs(); - return std::move(arena_); -} - -// String primitives - -common_peg_parser common_peg_parser_builder::string_content(char delimiter) { - return wrap(arena_.add_parser(common_peg_string_parser{delimiter})); -} - -common_peg_parser common_peg_parser_builder::double_quoted_string() { - return rule("double-quoted-string", [this]() { - return sequence({literal("\""), string_content('"'), literal("\""), space()}); - }); -} - -common_peg_parser common_peg_parser_builder::single_quoted_string() { - return rule("single-quoted-string", [this]() { - return sequence({literal("'"), string_content('\''), literal("'"), space()}); - }); -} - -common_peg_parser common_peg_parser_builder::quoted_string() { - return rule("quoted-string", [this]() { - return choice({double_quoted_string(), single_quoted_string()}); - }); -} - -// JSON parsers - -common_peg_parser common_peg_parser_builder::json_number() { - return rule("json-number", [this]() { - auto digit1_9 = chars("[1-9]", 1, 1); - auto digits = chars("[0-9]"); - auto int_part = choice({literal("0"), sequence({digit1_9, chars("[0-9]", 0, -1)})}); - auto frac = sequence({literal("."), digits}); - auto exp = sequence({choice({literal("e"), literal("E")}), optional(chars("[+-]", 1, 1)), digits}); - // Negative lookahead: only commit the number when the next character can't extend it. - // At EOF in partial mode, chars returns NEED_MORE → negate propagates NEED_MORE → number not committed. - // This prevents premature commits of partial numbers (e.g. "3" when "3.14" is incoming). - auto not_number_continuation = negate(chars("[0-9.eE+-]", 1, 1)); - return sequence({ optional(literal("-")), int_part, optional(frac), optional(exp), not_number_continuation, space() }); - }); -} - -common_peg_parser common_peg_parser_builder::json_string() { - return rule("json-string", [this]() { - return sequence({literal("\""), string_content('"'), literal("\""), space()}); - }); -} - -common_peg_parser common_peg_parser_builder::json_bool() { - return rule("json-bool", [this]() { - return sequence({choice({literal("true"), literal("false")}), space()}); - }); -} - -common_peg_parser common_peg_parser_builder::json_null() { - return rule("json-null", [this]() { - return sequence({literal("null"), space()}); - }); -} - -common_peg_parser common_peg_parser_builder::json_object() { - return rule("json-object", [this]() { - auto ws = space(); - auto member = sequence({json_string(), ws, literal(":"), ws, json()}); - auto members = sequence({member, zero_or_more(sequence({ws, literal(","), ws, member}))}); - return sequence({ - literal("{"), - ws, - choice({ - literal("}"), - sequence({members, ws, literal("}")}) - }), - ws - }); - }); -} - -common_peg_parser common_peg_parser_builder::json_array() { - return rule("json-array", [this]() { - auto ws = space(); - auto elements = sequence({json(), zero_or_more(sequence({literal(","), ws, json()}))}); - return sequence({ - literal("["), - ws, - choice({ - literal("]"), - sequence({elements, ws, literal("]")}) - }), - ws - }); - }); -} - -common_peg_parser common_peg_parser_builder::json() { - return rule("json-value", [this]() { - return choice({ - json_object(), - json_array(), - json_string(), - json_number(), - json_bool(), - json_null() - }); - }); -} - -common_peg_parser common_peg_parser_builder::python_string() { - return rule("python-string", [this]() { - return choice({double_quoted_string(), single_quoted_string()}); - }); -} - -common_peg_parser common_peg_parser_builder::python_number() { - return json_number(); -} - -common_peg_parser common_peg_parser_builder::python_bool() { - return rule("python-bool", [this]() { - return sequence({ - choice({literal("True"), literal("False")}), - space() - }); - }); -} - -common_peg_parser common_peg_parser_builder::python_null() { - return rule("python-none", [this]() { - return sequence({literal("None"), space()}); - }); -} - -common_peg_parser common_peg_parser_builder::python_dict() { - return rule("python-dict", [this]() { - auto ws = space(); - auto member = sequence({python_string(), ws, literal(":"), ws, python_value()}); - auto members = sequence({member, zero_or_more(sequence({ws, literal(","), ws, member}))}); - return sequence({ - literal("{"), - ws, - choice({ - literal("}"), - sequence({members, ws, literal("}")}) - }), - ws - }); - }); -} - -common_peg_parser common_peg_parser_builder::python_array() { - return rule("python-array", [this]() { - auto ws = space(); - auto elements = sequence({python_value(), zero_or_more(sequence({literal(","), ws, python_value()}))}); - return sequence({ - literal("["), - ws, - choice({ - literal("]"), - sequence({elements, ws, literal("]")}) - }), - ws - }); - }); -} - -common_peg_parser common_peg_parser_builder::python_value() { - return rule("python-value", [this]() { - return choice({ - python_dict(), - python_array(), - python_string(), - python_number(), - python_bool(), - python_null() - }); - }); -} - -common_peg_parser common_peg_parser_builder::marker() { - auto sharp_bracket_parser = literal("<") + until(">") + literal(">"); - auto square_bracket_parser = literal("[") + until("]") + literal("]"); - return choice({ sharp_bracket_parser, square_bracket_parser }); -} - -common_peg_parser common_peg_parser_builder::json_member(const std::string & key, const common_peg_parser & p) { - auto ws = space(); - return sequence({ - literal("\"" + key + "\""), - ws, - literal(":"), - ws, - p, - }); -} - -static std::string gbnf_escape_char_class(uint32_t c) { - if (c == '-' || c == ']' || c == '[' || c == '\\') { - return "\\" + std::string(1, (char) c); - } - // Escape whitespace control characters - if (c == '\n') { - return "\\n"; - } - if (c == '\t') { - return "\\t"; - } - if (c == '\r') { - return "\\r"; - } - - // Printable ASCII - if (c >= 0x20 && c <= 0x7E) { - return std::string(1, (char) c); - } - - // Hex escape - char buf[16]; - const char * hex = "0123456789ABCDEF"; - - if (c <= 0xFF) { - buf[0] = '\\'; - buf[1] = 'x'; - buf[2] = hex[(c >> 4) & 0xF]; - buf[3] = hex[c & 0xF]; - buf[4] = '\0'; - } else if (c <= 0xFFFF) { - buf[0] = '\\'; - buf[1] = 'u'; - buf[2] = hex[(c >> 12) & 0xF]; - buf[3] = hex[(c >> 8) & 0xF]; - buf[4] = hex[(c >> 4) & 0xF]; - buf[5] = hex[c & 0xF]; - buf[6] = '\0'; - } else { - buf[0] = '\\'; - buf[1] = 'U'; - for (int i = 0; i < 8; i++) { - buf[2 + i] = hex[(c >> ((7 - i) * 4)) & 0xF]; - } - buf[10] = '\0'; - } - - return std::string(buf); -} - -static std::string gbnf_excluding_pattern(const std::vector & strings) { - trie matcher(strings); - auto pieces = matcher.collect_prefix_and_next(); - - std::string pattern; - for (size_t i = 0; i < pieces.size(); ++i) { - if (i > 0) { - pattern += " | "; - } - - const auto & pre = pieces[i].prefix; - const auto & chars = pieces[i].next_chars; - - std::string cls; - cls.reserve(chars.size()); - for (uint32_t ch : chars) { - cls += gbnf_escape_char_class(ch); - } - - if (!pre.empty()) { - pattern += gbnf_format_literal(common_unicode_cpts_to_utf8(pre)) + " [^" + cls + "]"; - } else { - pattern += "[^" + cls + "]"; - } - } - - return "(" + pattern + ")*"; -} - -static std::unordered_set collect_reachable_rules( - const common_peg_arena & arena, - const common_peg_parser_id & rule -) { - std::unordered_set reachable; - std::unordered_set visited; - - std::function visit = [&](common_peg_parser_id id) { - const auto & parser = arena.get(id); - - std::visit([&](const auto & p) { - using T = std::decay_t; - - 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 || - std::is_same_v) { - // These parsers do not have any children - } else if constexpr (std::is_same_v) { - for (auto child : p.children) { - visit(child); - } - } else if constexpr (std::is_same_v) { - for (auto child : p.children) { - visit(child); - } - } else 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) { - visit(p.child); - } else if constexpr (std::is_same_v) { - if (visited.find(p.name) == visited.end()) { - visited.insert(p.name); - reachable.insert(p.name); - visit(p.child); - } - } else if constexpr (std::is_same_v) { - // Traverse rules so we pick up everything - auto referenced_rule = arena.get_rule(p.name); - visit(referenced_rule); - } else { - static_assert(is_always_false_v); - } - }, parser); - }; - - visit(rule); - return reachable; -} - -// GBNF generation implementation -void common_peg_arena::build_grammar(const common_grammar_builder & builder, bool lazy) const { - auto schema_delegates = [](const common_peg_schema_parser & s) -> bool { - if (!s.schema) { - return true; - } - if (s.raw && s.schema->contains("type")) { - const auto & type_val = s.schema->at("type"); - if (type_val.is_string() && type_val == "string") { - return true; - } - // Handle nullable types like ["string", "null"] - delegate when the - // non-null type is string, since the tagged format uses raw text - if (type_val.is_array()) { - for (const auto & t : type_val) { - if (t.is_string() && t.get() != "null") { - return t.get() == "string"; - } - } - } - } - // Delegate for enum schemas in raw mode - enum values are literal strings - if (s.raw && !s.schema->contains("type") && s.schema->contains("enum")) { - return true; - } - return false; - }; - - // Unwrap the parser so we can properly check if it's a sequence or choice - auto effective_parser = [&](common_peg_parser_id id) -> const common_peg_parser_variant & { - while (true) { - const auto & p = parsers_.at(id); - if (const auto * tag = std::get_if(&p)) { - id = tag->child; - } else if (const auto * atomic = std::get_if(&p)) { - id = atomic->child; - } else if (const auto * schema = std::get_if(&p)) { - if (schema_delegates(*schema)) { - id = schema->child; - } else { - return p; - } - } else { - return p; - } - } - }; - - // Generate GBNF for a parser - std::function to_gbnf = [&](common_peg_parser_id id) -> std::string { - const auto & parser = parsers_.at(id); - - return std::visit([&](const auto & p) -> std::string { - using T = std::decay_t; - - if constexpr (std::is_same_v || - std::is_same_v || - std::is_same_v) { - return ""; - } else if constexpr (std::is_same_v) { - return gbnf_format_literal(p.literal); - } else if constexpr (std::is_same_v) { - std::string s; - for (const auto & child : p.children) { - auto child_gbnf = to_gbnf(child); - if (child_gbnf.empty()) { - continue; - } - if (!s.empty()) { - s += " "; - } - const auto & child_parser = effective_parser(child); - if (std::holds_alternative(child_parser) || - std::holds_alternative(child_parser)) { - s += "(" + child_gbnf + ")"; - } else { - s += child_gbnf; - } - } - return s; - } else if constexpr (std::is_same_v) { - std::string s; - for (const auto & child : p.children) { - if (!s.empty()) { - s += " | "; - } - auto child_gbnf = to_gbnf(child); - const auto & child_parser = effective_parser(child); - if (std::holds_alternative(child_parser)) { - s += "(" + child_gbnf + ")"; - } else { - s += child_gbnf; - } - } - return s; - } else if constexpr (std::is_same_v) { - auto child_gbnf = to_gbnf(p.child); - const auto & child_parser = effective_parser(p.child); - if (std::holds_alternative(child_parser) || - std::holds_alternative(child_parser)) { - child_gbnf = "(" + child_gbnf + ")"; - } - if (p.min_count == 0 && p.max_count == 1) { - return child_gbnf + "?"; - } - if (p.min_count == 0 && p.max_count == -1) { - return child_gbnf + "*"; - } - if (p.min_count == 1 && p.max_count == -1) { - return child_gbnf + "+"; - } - if (p.max_count == -1) { - return child_gbnf + "{" + std::to_string(p.min_count) + ",}"; - } - if (p.min_count == p.max_count) { - if (p.min_count == 1) { - return child_gbnf; - } - return child_gbnf + "{" + std::to_string(p.min_count) + "}"; - } - return child_gbnf + "{" + std::to_string(p.min_count) + "," + std::to_string(p.max_count) + "}"; - } else if constexpr (std::is_same_v || std::is_same_v) { - return ""; // Lookahead not supported in GBNF - } else if constexpr (std::is_same_v) { - return "."; - } else if constexpr (std::is_same_v) { - return "space"; - } else if constexpr (std::is_same_v) { - std::string result = p.pattern; - if (p.min_count == 0 && p.max_count == 1) { - return result + "?"; - } - if (p.min_count == 0 && p.max_count == -1) { - return result + "*"; - } - if (p.min_count == 1 && p.max_count == -1) { - return result + "+"; - } - if (p.max_count == -1) { - return result + "{" + std::to_string(p.min_count) + ",}"; - } - if (p.min_count == p.max_count) { - if (p.min_count == 1) { - return result; - } - return result + "{" + std::to_string(p.min_count) + "}"; - } - return result + "{" + std::to_string(p.min_count) + "," + std::to_string(p.max_count) + "}"; - } else if constexpr (std::is_same_v) { - const std::string delim(1, p.delimiter); - return R"(( [^)" + delim + R"(\\] | "\\" ( [)" + delim + R"(\\/ bfnrt] | "u" [0-9a-fA-F]{4} ) )*)"; - } else if constexpr (std::is_same_v) { - if (p.delimiters.empty()) { - return ".*"; - } - return gbnf_excluding_pattern(p.delimiters); - } else if constexpr (std::is_same_v) { - if (schema_delegates(p)) { - return to_gbnf(p.child); - } - return builder.add_schema(p.name, *p.schema); - } else if constexpr (std::is_same_v) { - return p.name; - } else if constexpr (std::is_same_v) { - // Refs should not exist after flattening, but kept just in case - return p.name; - } else if constexpr (std::is_same_v) { - return to_gbnf(p.child); - } else if constexpr (std::is_same_v) { - return to_gbnf(p.child); - } else if constexpr (std::is_same_v) { - return p.grammar; - } else { - static_assert(is_always_false_v); - } - }, parser); - }; - - // Collect reachable rules - std::unordered_set reachable_rules; - - if (lazy) { - // Collect rules reachable from trigger rules - for (const auto & [name, id] : rules_) { - const auto & parser = parsers_.at(id); - if (auto rule = std::get_if(&parser)) { - if (rule->trigger) { - // Mark trigger as reachable and visit it - reachable_rules.insert(name); - auto add_rules = collect_reachable_rules(*this, id); - reachable_rules.insert(add_rules.begin(), add_rules.end()); - } - } - } - } else { - // Collect rules reachable from root - reachable_rules = collect_reachable_rules(*this, root_); - } - - // Create GBNF rules for all reachable rules - for (const auto & [name, rule_id] : rules_) { - if (reachable_rules.find(name) == reachable_rules.end()) { - continue; - } - - const auto & parser = parsers_.at(rule_id); - if (auto rule = std::get_if(&parser)) { - builder.add_rule(rule->name, to_gbnf(rule->child)); - } - } - - if (lazy) { - // Generate root rule from trigger rules only - std::vector trigger_names; - for (const auto & [name, rule_id] : rules_) { - const auto & parser = parsers_.at(rule_id); - if (auto rule = std::get_if(&parser)) { - if (rule->trigger) { - trigger_names.push_back(rule->name); - } - } - } - - // Sort for predictable order - std::sort(trigger_names.begin(), trigger_names.end()); - builder.add_rule("root", string_join(trigger_names, " | ")); - } else if (root_ != COMMON_PEG_INVALID_PARSER_ID) { - builder.add_rule("root", to_gbnf(root_)); - } -} - -static nlohmann::json serialize_parser_variant(const common_peg_parser_variant & variant) { - using json = nlohmann::json; - - return std::visit([](const auto & p) -> json { - using T = std::decay_t; - - if constexpr (std::is_same_v) { - return json{{"type", "epsilon"}}; - } else if constexpr (std::is_same_v) { - return json{{"type", "start"}}; - } else if constexpr (std::is_same_v) { - return json{{"type", "end"}}; - } else if constexpr (std::is_same_v) { - return json{{"type", "literal"}, {"literal", p.literal}}; - } else if constexpr (std::is_same_v) { - return json{{"type", "sequence"}, {"children", p.children}}; - } else if constexpr (std::is_same_v) { - return json{{"type", "choice"}, {"children", p.children}}; - } else if constexpr (std::is_same_v) { - return json{ - {"type", "repetition"}, - {"child", p.child}, - {"min_count", p.min_count}, - {"max_count", p.max_count} - }; - } else if constexpr (std::is_same_v) { - return json{{"type", "and"}, {"child", p.child}}; - } else if constexpr (std::is_same_v) { - return json{{"type", "not"}, {"child", p.child}}; - } else if constexpr (std::is_same_v) { - return json{{"type", "any"}}; - } else if constexpr (std::is_same_v) { - return json{{"type", "space"}}; - } else if constexpr (std::is_same_v) { - json ranges = json::array(); - for (const auto & range : p.ranges) { - ranges.push_back({{"start", range.start}, {"end", range.end}}); - } - return json{ - {"type", "chars"}, - {"pattern", p.pattern}, - {"ranges", ranges}, - {"negated", p.negated}, - {"min_count", p.min_count}, - {"max_count", p.max_count} - }; - } else if constexpr (std::is_same_v) { - return json{{"type", "string"}, {"delimiter", std::string(1, p.delimiter)}}; - } else if constexpr (std::is_same_v) { - return json{{"type", "until"}, {"delimiters", p.delimiters}}; - } else if constexpr (std::is_same_v) { - return json{ - {"type", "schema"}, - {"child", p.child}, - {"name", p.name}, - {"schema", p.schema ? *p.schema : nullptr}, - {"raw", p.raw} - }; - } else if constexpr (std::is_same_v) { - return json{ - {"type", "rule"}, - {"name", p.name}, - {"child", p.child}, - {"trigger", p.trigger} - }; - } else if constexpr (std::is_same_v) { - return json{{"type", "ref"}, {"name", p.name}}; - } else if constexpr (std::is_same_v) { - return json{{"type", "atomic"}, {"child", p.child}}; - } else if constexpr (std::is_same_v) { - return json{ - {"type", "tag"}, - {"child", p.child}, - {"tag", p.tag} - }; - } else if constexpr (std::is_same_v) { - return json{{"type", "gbnf"}, {"child", p.child}, {"grammar", p.grammar}}; - } - }, variant); -} - -nlohmann::json common_peg_arena::to_json() const { - auto parsers = nlohmann::json::array(); - for (const auto & parser : parsers_) { - parsers.push_back(serialize_parser_variant(parser)); - } - return nlohmann::json{ - {"parsers", parsers}, - {"rules", rules_}, - {"root", root_} - }; -} - -static common_peg_parser_variant deserialize_parser_variant(const nlohmann::json & j) { - if (!j.contains("type") || !j["type"].is_string()) { - throw std::runtime_error("Parser variant JSON missing or invalid 'type' field"); - } - - std::string type = j["type"]; - - if (type == "epsilon") { - return common_peg_epsilon_parser{}; - } - if (type == "start") { - return common_peg_start_parser{}; - } - if (type == "end") { - return common_peg_end_parser{}; - } - if (type == "literal") { - if (!j.contains("literal") || !j["literal"].is_string()) { - throw std::runtime_error("literal parser missing or invalid 'literal' field"); - } - return common_peg_literal_parser{j["literal"]}; - } - if (type == "sequence") { - if (!j.contains("children") || !j["children"].is_array()) { - throw std::runtime_error("sequence parser missing or invalid 'children' field"); - } - return common_peg_sequence_parser{j["children"].get>()}; - } - if (type == "choice") { - if (!j.contains("children") || !j["children"].is_array()) { - throw std::runtime_error("choice parser missing or invalid 'children' field"); - } - return common_peg_choice_parser{j["children"].get>()}; - } - if (type == "repetition") { - if (!j.contains("child") || !j.contains("min_count") || !j.contains("max_count")) { - throw std::runtime_error("repetition parser missing required fields"); - } - return common_peg_repetition_parser{ - j["child"].get(), - j["min_count"].get(), - j["max_count"].get() - }; - } - if (type == "and") { - if (!j.contains("child")) { - throw std::runtime_error("and parser missing 'child' field"); - } - return common_peg_and_parser{j["child"].get()}; - } - if (type == "not") { - if (!j.contains("child")) { - throw std::runtime_error("not parser missing 'child' field"); - } - return common_peg_not_parser{j["child"].get()}; - } - if (type == "any") { - return common_peg_any_parser{}; - } - if (type == "space") { - return common_peg_space_parser{}; - } - if (type == "chars") { - if (!j.contains("pattern") || !j.contains("ranges") || !j.contains("negated") || - !j.contains("min_count") || !j.contains("max_count")) { - throw std::runtime_error("chars parser missing required fields"); - } - common_peg_chars_parser parser; - parser.pattern = j["pattern"]; - parser.negated = j["negated"]; - parser.min_count = j["min_count"]; - parser.max_count = j["max_count"]; - for (const auto & range_json : j["ranges"]) { - if (!range_json.contains("start") || !range_json.contains("end")) { - throw std::runtime_error("char_range missing 'start' or 'end' field"); - } - parser.ranges.push_back({ - range_json["start"].get(), - range_json["end"].get() - }); - } - return parser; - } - if (type == "string") { - if (!j.contains("delimiter")) { - throw std::runtime_error("string parser missing delimiter field."); - } - std::string delimiter = j["delimiter"]; - if (delimiter.empty()) { - throw std::runtime_error("string parser delimiter is empty."); - } - return common_peg_string_parser{delimiter[0]}; - } - if (type == "until") { - if (!j.contains("delimiters") || !j["delimiters"].is_array()) { - throw std::runtime_error("until parser missing or invalid 'delimiters' field"); - } - return common_peg_until_parser{j["delimiters"].get>()}; - } - if (type == "schema") { - if (!j.contains("child") || !j.contains("name") || !j.contains("schema") || !j.contains("raw")) { - throw std::runtime_error("schema parser missing required fields"); - } - common_peg_schema_parser parser; - parser.child = j["child"].get(); - parser.name = j["name"]; - if (!j["schema"].is_null()) { - parser.schema = std::make_shared(j["schema"]); - } - parser.raw = j["raw"].get(); - return parser; - } - if (type == "rule") { - if (!j.contains("name") || !j.contains("child") || !j.contains("trigger")) { - throw std::runtime_error("rule parser missing required fields"); - } - return common_peg_rule_parser{ - j["name"].get(), - j["child"].get(), - j["trigger"].get() - }; - } - if (type == "ref") { - if (!j.contains("name") || !j["name"].is_string()) { - throw std::runtime_error("ref parser missing or invalid 'name' field"); - } - return common_peg_ref_parser{j["name"]}; - } - if (type == "atomic") { - if (!j.contains("child")) { - throw std::runtime_error("tag parser missing required fields"); - } - return common_peg_atomic_parser{ - j["child"].get(), - }; - } - if (type == "tag") { - if (!j.contains("child") || !j.contains("tag")) { - throw std::runtime_error("tag parser missing required fields"); - } - return common_peg_tag_parser{ - j["child"].get(), - j["tag"].get(), - }; - } - - if (type == "gbnf") { - if (!j.contains("child") || !j.contains("grammar")) { - throw std::runtime_error("gbnf parser missing required fields"); - } - return common_peg_gbnf_parser{ - j["child"].get(), - j["grammar"].get(), - }; - } - - throw std::runtime_error("Unknown parser type: " + type); -} - -common_peg_arena common_peg_arena::from_json(const nlohmann::json & j) { - if (!j.contains("parsers") || !j["parsers"].is_array()) { - throw std::runtime_error("JSON missing or invalid 'parsers' array"); - } - if (!j.contains("rules") || !j["rules"].is_object()) { - throw std::runtime_error("JSON missing or invalid 'rules' object"); - } - if (!j.contains("root")) { - throw std::runtime_error("JSON missing 'root' field"); - } - - common_peg_arena arena; - - const auto & parsers_json = j["parsers"]; - arena.parsers_.reserve(parsers_json.size()); - for (const auto & parser_json : parsers_json) { - arena.parsers_.push_back(deserialize_parser_variant(parser_json)); - } - - arena.rules_ = j["rules"].get>(); - - for (const auto & [name, id] : arena.rules_) { - if (id >= arena.parsers_.size()) { - throw std::runtime_error("Rule '" + name + "' references invalid parser ID: " + std::to_string(id)); - } - } - - arena.root_ = j["root"].get(); - if (arena.root_ != COMMON_PEG_INVALID_PARSER_ID && arena.root_ >= arena.parsers_.size()) { - throw std::runtime_error("Root references invalid parser ID: " + std::to_string(arena.root_)); - } - - return arena; -} - -std::string common_peg_arena::save() const { - return to_json().dump(); -} - -void common_peg_arena::load(const std::string & data) { - *this = from_json(nlohmann::json::parse(data)); -} - -common_peg_arena build_peg_parser(const std::function & fn) { - common_peg_parser_builder builder; - builder.set_root(fn(builder)); - return builder.build(); -} diff --git a/common/peg-parser.h b/common/peg-parser.h deleted file mode 100644 index b6bb05214bb2..000000000000 --- a/common/peg-parser.h +++ /dev/null @@ -1,523 +0,0 @@ -#pragma once - -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -struct common_grammar_builder; - -class common_peg_parser_builder; - -using common_peg_parser_id = size_t; -constexpr common_peg_parser_id COMMON_PEG_INVALID_PARSER_ID = static_cast(-1); - -using common_peg_ast_id = size_t; -constexpr common_peg_ast_id COMMON_PEG_INVALID_AST_ID = static_cast(-1); - -// Lightweight wrapper around common_peg_parser_id for convenience -class common_peg_parser { - common_peg_parser_id id_; - common_peg_parser_builder & builder_; - - public: - common_peg_parser(const common_peg_parser & other) : id_(other.id_), builder_(other.builder_) {} - common_peg_parser(common_peg_parser_id id, common_peg_parser_builder & builder) : id_(id), builder_(builder) {} - - common_peg_parser & operator=(const common_peg_parser & other); - common_peg_parser & operator+=(const common_peg_parser & other); - common_peg_parser & operator|=(const common_peg_parser & other); - - operator common_peg_parser_id() const { return id_; } - common_peg_parser_id id() const { return id_; } - - common_peg_parser_builder & builder() const { return builder_; } - - // Creates a sequence - common_peg_parser operator+(const common_peg_parser & other) const; - - // Creates a sequence separated by spaces. - common_peg_parser operator<<(const common_peg_parser & other) const; - - // Creates a choice - common_peg_parser operator|(const common_peg_parser & other) const; - - common_peg_parser operator+(const char * str) const; - common_peg_parser operator+(const std::string & str) const; - common_peg_parser operator<<(const char * str) const; - common_peg_parser operator<<(const std::string & str) const; - common_peg_parser operator|(const char * str) const; - common_peg_parser operator|(const std::string & str) const; -}; - -common_peg_parser operator+(const char * str, const common_peg_parser & p); -common_peg_parser operator+(const std::string & str, const common_peg_parser & p); -common_peg_parser operator<<(const char * str, const common_peg_parser & p); -common_peg_parser operator<<(const std::string & str, const common_peg_parser & p); -common_peg_parser operator|(const char * str, const common_peg_parser & p); -common_peg_parser operator|(const std::string & str, const common_peg_parser & p); - -enum common_peg_parse_result_type { - COMMON_PEG_PARSE_RESULT_FAIL = 0, - COMMON_PEG_PARSE_RESULT_SUCCESS = 1, - COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT = 2, -}; - -const char * common_peg_parse_result_type_name(common_peg_parse_result_type type); - -struct common_peg_ast_node { - common_peg_ast_id id; - std::string rule; - std::string tag; - size_t start; - size_t end; - std::string_view text; - std::vector children; - - bool is_partial = false; -}; - -struct common_peg_parse_result; - -using common_peg_ast_visitor = std::function; - -class common_peg_ast_arena { - std::vector nodes_; - public: - common_peg_ast_id add_node( - const std::string & rule, - const std::string & tag, - size_t start, - size_t end, - std::string_view text, - std::vector children, - bool is_partial = false - ) { - common_peg_ast_id id = nodes_.size(); - nodes_.push_back({id, rule, tag, start, end, text, std::move(children), is_partial}); - return id; - } - - const common_peg_ast_node & get(common_peg_ast_id id) const { return nodes_.at(id); } - - common_peg_ast_id find_by_tag(const common_peg_ast_node & parent, const std::string & tag, int max_depth = 3) const; - common_peg_ast_id find_by_rule(const common_peg_ast_node & parent, const std::string & tag, int max_depth = 3) const; - - size_t size() const { return nodes_.size(); } - - void clear() { nodes_.clear(); } - - void visit(common_peg_ast_id id, const common_peg_ast_visitor & visitor) const; - void visit(const common_peg_parse_result & result, const common_peg_ast_visitor & visitor) const; - - std::string dump(); -}; - -struct common_peg_parse_result { - common_peg_parse_result_type type = COMMON_PEG_PARSE_RESULT_FAIL; - size_t start = 0; - size_t end = 0; - - std::vector nodes; - - common_peg_parse_result() = default; - - common_peg_parse_result(common_peg_parse_result_type type, size_t start) - : type(type), start(start), end(start) {} - - common_peg_parse_result(common_peg_parse_result_type type, size_t start, size_t end) - : type(type), start(start), end(end) {} - - common_peg_parse_result(common_peg_parse_result_type type, size_t start, size_t end, std::vector nodes) - : type(type), start(start), end(end), nodes(std::move(nodes)) {} - - bool fail() const { return type == COMMON_PEG_PARSE_RESULT_FAIL; } - bool need_more_input() const { return type == COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT; } - bool success() const { return type == COMMON_PEG_PARSE_RESULT_SUCCESS; } -}; - -enum common_peg_parse_flags { - COMMON_PEG_PARSE_FLAG_NONE = 0, - COMMON_PEG_PARSE_FLAG_LENIENT = 1 << 0, - COMMON_PEG_PARSE_FLAG_DEBUG = 1 << 1, -}; - -inline common_peg_parse_flags operator|(common_peg_parse_flags a, common_peg_parse_flags b) { - return static_cast(int(a) | int(b)); -} - -inline common_peg_parse_flags & operator|=(common_peg_parse_flags & a, common_peg_parse_flags b) { - return a = a | b; -} - -inline common_peg_parse_flags operator&(common_peg_parse_flags a, common_peg_parse_flags b) { - return static_cast(int(a) & int(b)); -} - -inline common_peg_parse_flags operator~(common_peg_parse_flags a) { - return static_cast(~int(a)); -} - -struct common_peg_parse_context { - std::string input; - common_peg_parse_flags flags; - common_peg_ast_arena ast; - - int parse_depth; - - common_peg_parse_context(common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE) - : flags(flags), parse_depth(0) {} - - common_peg_parse_context(const std::string & input, common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE) - : input(input), flags(flags), parse_depth(0) {} - - bool is_lenient() const { return flags & COMMON_PEG_PARSE_FLAG_LENIENT; } - bool is_debug() const { return flags & COMMON_PEG_PARSE_FLAG_DEBUG; } -}; - -class common_peg_arena; - -// Parser variants -struct common_peg_epsilon_parser {}; - -struct common_peg_start_parser {}; - -struct common_peg_end_parser {}; - -struct common_peg_literal_parser { - std::string literal; -}; - -struct common_peg_sequence_parser { - std::vector children; -}; - -struct common_peg_choice_parser { - std::vector children; -}; - -struct common_peg_repetition_parser { - common_peg_parser_id child; - int min_count; - int max_count; // -1 for unbounded -}; - -struct common_peg_and_parser { - common_peg_parser_id child; -}; - -struct common_peg_not_parser { - common_peg_parser_id child; -}; - -struct common_peg_any_parser {}; - -struct common_peg_space_parser {}; - -struct common_peg_chars_parser { - struct char_range { - uint32_t start; - uint32_t end; - bool contains(uint32_t codepoint) const { return codepoint >= start && codepoint <= end; } - }; - - std::string pattern; - std::vector ranges; - bool negated; - int min_count; - int max_count; // -1 for unbounded -}; - -struct common_peg_string_parser { - char delimiter; -}; - -struct common_peg_until_parser { - std::vector delimiters; -}; - -struct common_peg_schema_parser { - common_peg_parser_id child; - std::string name; - std::shared_ptr schema; - - // Indicates if the GBNF should accept a raw string that matches the schema. - bool raw; -}; - -struct common_peg_rule_parser { - std::string name; - common_peg_parser_id child; - bool trigger; -}; - -struct common_peg_ref_parser { - std::string name; -}; - -struct common_peg_atomic_parser { - common_peg_parser_id child; -}; - -struct common_peg_tag_parser { - common_peg_parser_id child; - std::string tag; -}; - -struct common_peg_gbnf_parser { - common_peg_parser_id child; - std::string grammar; -}; - -// Variant holding all parser types -using common_peg_parser_variant = std::variant< - common_peg_epsilon_parser, - common_peg_start_parser, - common_peg_end_parser, - common_peg_literal_parser, - common_peg_sequence_parser, - common_peg_choice_parser, - common_peg_repetition_parser, - common_peg_and_parser, - common_peg_not_parser, - common_peg_any_parser, - common_peg_space_parser, - common_peg_chars_parser, - common_peg_string_parser, - common_peg_until_parser, - common_peg_schema_parser, - common_peg_rule_parser, - common_peg_ref_parser, - common_peg_atomic_parser, - common_peg_tag_parser, - common_peg_gbnf_parser ->; - -class common_peg_arena { - std::vector parsers_; - std::unordered_map rules_; - common_peg_parser_id root_ = COMMON_PEG_INVALID_PARSER_ID; - - public: - const common_peg_parser_variant & get(common_peg_parser_id id) const { return parsers_.at(id); } - common_peg_parser_variant & get(common_peg_parser_id id) { return parsers_.at(id); } - - size_t size() const { return parsers_.size(); } - bool empty() const { return parsers_.empty(); } - - common_peg_parser_id get_rule(const std::string & name) const; - bool has_rule(const std::string & name) const { return rules_.find(name) != rules_.end(); } - - common_peg_parser_id root() const { return root_; } - void set_root(common_peg_parser_id id) { root_ = id; } - - common_peg_parse_result parse(common_peg_parse_context & ctx, size_t start = 0) const; - common_peg_parse_result parse(common_peg_parser_id id, common_peg_parse_context & ctx, size_t start) const; - - void resolve_refs(); - - void build_grammar(const common_grammar_builder & builder, bool lazy = false) const; - - std::string dump(common_peg_parser_id id) const; - - nlohmann::json to_json() const; - static common_peg_arena from_json(const nlohmann::json & j); - - std::string save() const; - void load(const std::string & data); - - friend class common_peg_parser_builder; - - private: - std::string dump_impl(common_peg_parser_id id, std::unordered_set & visited) const; - - common_peg_parser_id add_parser(common_peg_parser_variant parser); - void add_rule(const std::string & name, common_peg_parser_id id); - - common_peg_parser_id resolve_ref(common_peg_parser_id id); -}; - -class common_peg_parser_builder { - common_peg_arena arena_; - - common_peg_parser wrap(common_peg_parser_id id) { return common_peg_parser(id, *this); } - common_peg_parser add(const common_peg_parser_variant & p) { return wrap(arena_.add_parser(p)); } - - public: - common_peg_parser_builder(); - - // Match nothing, always succeed. - // S -> ε - common_peg_parser eps() { return add(common_peg_epsilon_parser{}); } - - // Matches the start of the input. - // S -> ^ - common_peg_parser start() { return add(common_peg_start_parser{}); } - - // Matches the end of the input. - // S -> $ - common_peg_parser end() { return add(common_peg_end_parser{}); } - - // Matches an exact literal string. - // S -> "hello" - common_peg_parser literal(const std::string & literal) { return add(common_peg_literal_parser{literal}); } - - // Matches a sequence of parsers in order, all must succeed. - // S -> A B C - common_peg_parser sequence() { return add(common_peg_sequence_parser{}); } - common_peg_parser sequence(const std::vector & parsers); - common_peg_parser sequence(const std::vector & parsers); - common_peg_parser sequence(std::initializer_list parsers); - - // Matches the first parser that succeeds from a list of alternatives. - // S -> A | B | C - common_peg_parser choice() { return add(common_peg_choice_parser{}); } - common_peg_parser choice(const std::vector & parsers); - common_peg_parser choice(const std::vector & parsers); - common_peg_parser choice(std::initializer_list parsers); - - // Matches one or more repetitions of a parser. - // S -> A+ - common_peg_parser one_or_more(const common_peg_parser & p) { return repeat(p, 1, -1); } - - // Matches zero or more repetitions of a parser, always succeeds. - // S -> A* - common_peg_parser zero_or_more(const common_peg_parser & p) { return repeat(p, 0, -1); } - - // Matches zero or one occurrence of a parser, always succeeds. - // S -> A? - common_peg_parser optional(const common_peg_parser & p) { return repeat(p, 0, 1); } - - // Positive lookahead: succeeds if child parser succeeds, consumes no input. - // S -> &A - common_peg_parser peek(const common_peg_parser & p) { return add(common_peg_and_parser{p}); } - - // Negative lookahead: succeeds if child parser fails, consumes no input. - // S -> !A - common_peg_parser negate(const common_peg_parser & p) { return add(common_peg_not_parser{p}); } - - // Matches any single character. - // S -> . - common_peg_parser any() { return add(common_peg_any_parser{}); } - - // Matches between min and max repetitions of characters from a character class. - // S -> [a-z]{m,n} - // - // Use -1 for max to represent unbounded repetition (equivalent to {m,}) - common_peg_parser chars(const std::string & classes, int min = 1, int max = -1); - - // Creates a lightweight reference to a named rule (resolved during build()). - // Use this for forward references in recursive grammars. - // expr_ref -> expr - common_peg_parser ref(const std::string & name) { return add(common_peg_ref_parser{name}); } - - // Matches zero or more whitespace characters (space, tab, newline). - // S -> [ \t\n]* - common_peg_parser space() { return add(common_peg_space_parser{}); } - - // Matches all characters until a delimiter is found (delimiter not consumed). - // S -> (!delim .)* - common_peg_parser until(const std::string & delimiter) { return add(common_peg_until_parser{{delimiter}}); } - - // Matches all characters until one of the delimiters in the list is found (delimiter not consumed). - // S -> (!delim .)* - common_peg_parser until_one_of(const std::vector & delimiters) { return add(common_peg_until_parser{delimiters}); } - - // Matches everything - // S -> .* - common_peg_parser rest() { return until_one_of({}); } - - // Matches between min and max repetitions of a parser (inclusive). - // S -> A{m,n} - // Use -1 for max to represent unbounded repetition (equivalent to {m,}) - common_peg_parser repeat(const common_peg_parser & p, int min, int max) { return add(common_peg_repetition_parser{p, min,max}); } - - // Matches exactly n repetitions of a parser. - // S -> A{n} - common_peg_parser repeat(const common_peg_parser & p, int n) { return repeat(p, n, n); } - - // Matches a double-quoted string: '"' content '"' space - common_peg_parser double_quoted_string(); - - // Matches a single-quoted string: "'" content "'" space - common_peg_parser single_quoted_string(); - - // Matches a string that accepts both double-quoted and single-quoted styles. - common_peg_parser quoted_string(); - - // Matches string content without the surrounding delimiter. - common_peg_parser string_content(char delimiter); - - // Creates a complete JSON parser supporting objects, arrays, strings, numbers, booleans, and null. - // value -> object | array | string | number | true | false | null - common_peg_parser json(); - common_peg_parser json_object(); - common_peg_parser json_string(); - common_peg_parser json_array(); - common_peg_parser json_number(); - common_peg_parser json_bool(); - common_peg_parser json_null(); - - // Matches a JSON object member with a key and associated parser as the - // value. - common_peg_parser json_member(const std::string & key, const common_peg_parser & p); - - // Creates a complete Python format parser supporting dicts, arrays, strings, numbers, booleans, and None. - // Differs from JSON: uses True/False/None, accepts both single and double-quoted strings. - // value -> dict | array | string | number | True | False | None - common_peg_parser python_value(); - common_peg_parser python_dict(); - common_peg_parser python_string(); - common_peg_parser python_array(); - common_peg_parser python_number(); - common_peg_parser python_bool(); - common_peg_parser python_null(); - - // A marker, i.e. text delimited by a pair of <> or [] - common_peg_parser marker(); - - // Wraps a parser with JSON schema metadata for grammar generation. - // Used internally to convert JSON schemas to GBNF grammar rules. - common_peg_parser schema(const common_peg_parser & p, const std::string & name, const nlohmann::ordered_json & schema, bool raw = false); - - // Creates a named rule, stores it in the grammar, and returns a ref. - // If trigger=true, marks this rule as an entry point for lazy grammar generation. - // auto json = p.rule("json", json_obj | json_arr | ...) - common_peg_parser rule(const std::string & name, const common_peg_parser & p, bool trigger = false); - - // Creates a named rule using a builder function, and returns a ref. - // If trigger=true, marks this rule as an entry point for lazy grammar generation. - // auto json = p.rule("json", [&]() { return json_object() | json_array() | ... }) - common_peg_parser rule(const std::string & name, const std::function & builder, bool trigger = false); - - // Creates a trigger rule. When generating a lazy grammar from the parser, - // only trigger rules and descendents are emitted. - common_peg_parser trigger_rule(const std::string & name, const common_peg_parser & p) { return rule(name, p, true); } - common_peg_parser trigger_rule(const std::string & name, const std::function & builder) { return rule(name, builder, true); } - - // Creates an atomic parser. Atomic parsers do not create an AST node if - // the child results in a partial parse, i.e. NEEDS_MORE_INPUT. This is - // intended for situations where partial output is undesirable. - common_peg_parser atomic(const common_peg_parser & p) { return add(common_peg_atomic_parser{p}); } - - // Tags create nodes in the generated AST for semantic purposes. - // Unlike rules, you can tag multiple nodes with the same tag. - common_peg_parser tag(const std::string & tag, const common_peg_parser & p) { return add(common_peg_tag_parser{p.id(), tag}); } - - // Wraps a child parser but emits a custom GBNF grammar string instead of - // the child's grammar. Parsing delegates entirely to the child. - common_peg_parser gbnf(const common_peg_parser & p, const std::string & grammar) { return add(common_peg_gbnf_parser{p, grammar}); } - - void set_root(const common_peg_parser & p); - - common_peg_arena build(); -}; - -// Helper function for building parsers -common_peg_arena build_peg_parser(const std::function & fn); diff --git a/common/preset.cpp b/common/preset.cpp deleted file mode 100644 index 9187a67f05a0..000000000000 --- a/common/preset.cpp +++ /dev/null @@ -1,483 +0,0 @@ -#include "arg.h" -#include "preset.h" -#include "peg-parser.h" -#include "log.h" -#include "download.h" - -#include -#include -#include - -static std::string rm_leading_dashes(const std::string & str) { - size_t pos = 0; - while (pos < str.size() && str[pos] == '-') { - ++pos; - } - return str.substr(pos); -} - -// only allow a subset of args for remote presets for security reasons -// do not add more args unless absolutely necessary -// args that output to files are strictly prohibited -static std::set get_remote_preset_whitelist(const std::map & key_to_opt) { - static const std::set allowed_options = { - "model-url", - "hf-repo", - "hf-repo-draft", - "hf-repo-v", // vocoder - "hf-file-v", // vocoder - "mmproj-url", - "pooling", - "jinja", - "batch-size", - "ubatch-size", - "cache-reuse", - "chat-template-kwargs", - "mmap", - // note: sampling params are automatically allowed by default - // negated args will be added automatically if the positive arg is specified above - }; - - std::set allowed_keys; - - for (const auto & it : key_to_opt) { - const std::string & key = it.first; - const common_arg & opt = it.second; - if (allowed_options.find(key) != allowed_options.end() || opt.is_sampling) { - allowed_keys.insert(key); - // also add variant keys (args without leading dashes and env vars) - for (const auto & arg : opt.get_args()) { - allowed_keys.insert(rm_leading_dashes(arg)); - } - for (const auto & env : opt.get_env()) { - allowed_keys.insert(env); - } - } - } - - return allowed_keys; -} - -std::vector common_preset::to_args(const std::string & bin_path) const { - std::vector args; - - if (!bin_path.empty()) { - args.push_back(bin_path); - } - - for (const auto & [opt, value] : options) { - if (opt.is_preset_only) { - continue; // skip preset-only options (they are not CLI args) - } - - // use the last arg as the main arg (i.e. --long-form) - args.push_back(opt.args.back()); - - // handle value(s) - if (opt.value_hint == nullptr && opt.value_hint_2 == nullptr) { - // flag option, no value - if (common_arg_utils::is_falsey(value)) { - // use negative arg if available - if (!opt.args_neg.empty()) { - args.back() = opt.args_neg.back(); - } else { - // otherwise, skip the flag - // TODO: maybe throw an error instead? - args.pop_back(); - } - } - } - if (opt.value_hint != nullptr) { - // single value - args.push_back(value); - } - if (opt.value_hint != nullptr && opt.value_hint_2 != nullptr) { - throw std::runtime_error(string_format( - "common_preset::to_args(): option '%s' has two values, which is not supported yet", - opt.args.back() - )); - } - } - - return args; -} - -std::string common_preset::to_ini() const { - std::ostringstream ss; - - ss << "[" << name << "]\n"; - for (const auto & [opt, value] : options) { - auto espaced_value = value; - string_replace_all(espaced_value, "\n", "\\\n"); - ss << rm_leading_dashes(opt.args.back()) << " = "; - ss << espaced_value << "\n"; - } - ss << "\n"; - - return ss.str(); -} - -void common_preset::set_option(const common_preset_context & ctx, const std::string & env, const std::string & value) { - // try if option exists, update it - for (auto & [opt, val] : options) { - if (opt.env && env == opt.env) { - val = value; - return; - } - } - // if option does not exist, we need to add it - if (ctx.key_to_opt.find(env) == ctx.key_to_opt.end()) { - throw std::runtime_error(string_format( - "%s: option with env '%s' not found in ctx_params", - __func__, env.c_str() - )); - } - options[ctx.key_to_opt.at(env)] = value; -} - -void common_preset::unset_option(const std::string & env) { - for (auto it = options.begin(); it != options.end(); ) { - const common_arg & opt = it->first; - if (opt.env && env == opt.env) { - it = options.erase(it); - return; - } else { - ++it; - } - } -} - -bool common_preset::get_option(const std::string & env, std::string & value) const { - for (const auto & [opt, val] : options) { - if (opt.env && env == opt.env) { - value = val; - return true; - } - } - return false; -} - -void common_preset::merge(const common_preset & other) { - for (const auto & [opt, val] : other.options) { - options[opt] = val; // overwrite existing options - } -} - -void common_preset::apply_to_params(common_params & params) const { - for (const auto & [opt, val] : options) { - // apply each option to params - if (opt.handler_string) { - opt.handler_string(params, val); - } else if (opt.handler_int) { - opt.handler_int(params, std::stoi(val)); - } else if (opt.handler_bool) { - opt.handler_bool(params, common_arg_utils::is_truthy(val)); - } else if (opt.handler_str_str) { - // not supported yet - throw std::runtime_error(string_format( - "%s: option with two values is not supported yet", - __func__ - )); - } else if (opt.handler_void) { - opt.handler_void(params); - } else { - GGML_ABORT("unknown handler type"); - } - } -} - -static std::map> parse_ini_from_file(const std::string & path) { - std::map> parsed; - - if (!std::filesystem::exists(path)) { - throw std::runtime_error("preset file does not exist: " + path); - } - - std::ifstream file(path); - if (!file.good()) { - throw std::runtime_error("failed to open server preset file: " + path); - } - - std::string contents((std::istreambuf_iterator(file)), std::istreambuf_iterator()); - - static const auto parser = build_peg_parser([](auto & p) { - // newline ::= "\r\n" / "\n" / "\r" - auto newline = p.rule("newline", p.literal("\r\n") | p.literal("\n") | p.literal("\r")); - - // ws ::= [ \t]* - auto ws = p.rule("ws", p.chars("[ \t]", 0, -1)); - - // comment ::= [;#] (!newline .)* - auto comment = p.rule("comment", p.chars("[;#]", 1, 1) + p.zero_or_more(p.negate(newline) + p.any())); - - // eol ::= ws comment? (newline / EOF) - auto eol = p.rule("eol", ws + p.optional(comment) + (newline | p.end())); - - // ident ::= [a-zA-Z_] [a-zA-Z0-9_.-]* - auto ident = p.rule("ident", p.chars("[a-zA-Z_]", 1, 1) + p.chars("[a-zA-Z0-9_.-]", 0, -1)); - - // value ::= (!eol-start .)* - auto eol_start = p.rule("eol-start", ws + (p.chars("[;#]", 1, 1) | newline | p.end())); - auto value = p.rule("value", p.zero_or_more(p.negate(eol_start) + p.any())); - - // header-line ::= "[" ws ident ws "]" eol - auto header_line = p.rule("header-line", "[" + ws + p.tag("section-name", p.chars("[^]]")) + ws + "]" + eol); - - // kv-line ::= ident ws "=" ws value eol - auto kv_line = p.rule("kv-line", p.tag("key", ident) + ws + "=" + ws + p.tag("value", value) + eol); - - // comment-line ::= ws comment (newline / EOF) - auto comment_line = p.rule("comment-line", ws + comment + (newline | p.end())); - - // blank-line ::= ws (newline / EOF) - auto blank_line = p.rule("blank-line", ws + (newline | p.end())); - - // line ::= header-line / kv-line / comment-line / blank-line - auto line = p.rule("line", header_line | kv_line | comment_line | blank_line); - - // ini ::= line* EOF - auto ini = p.rule("ini", p.zero_or_more(line) + p.end()); - - return ini; - }); - - common_peg_parse_context ctx(contents); - const auto result = parser.parse(ctx); - if (!result.success()) { - throw std::runtime_error("failed to parse server config file: " + path); - } - - std::string current_section = COMMON_PRESET_DEFAULT_NAME; - std::string current_key; - - ctx.ast.visit(result, [&](const auto & node) { - if (node.tag == "section-name") { - const std::string section = std::string(node.text); - current_section = section; - parsed[current_section] = {}; - } else if (node.tag == "key") { - const std::string key = std::string(node.text); - current_key = key; - } else if (node.tag == "value" && !current_key.empty() && !current_section.empty()) { - parsed[current_section][current_key] = std::string(node.text); - current_key.clear(); - } - }); - - return parsed; -} - -static std::map get_map_key_opt(common_params_context & ctx_params) { - std::map mapping; - for (const auto & opt : ctx_params.options) { - for (const auto & env : opt.get_env()) { - mapping[env] = opt; - } - for (const auto & arg : opt.get_args()) { - mapping[rm_leading_dashes(arg)] = opt; - } - } - return mapping; -} - -static bool is_bool_arg(const common_arg & arg) { - return !arg.args_neg.empty(); -} - -static std::string parse_bool_arg(const common_arg & arg, const std::string & key, const std::string & value) { - // if this is a negated arg, we need to reverse the value - for (const auto & neg_arg : arg.args_neg) { - if (rm_leading_dashes(neg_arg) == key) { - return common_arg_utils::is_truthy(value) ? "false" : "true"; - } - } - // otherwise, not negated - return value; -} - -common_preset_context::common_preset_context(llama_example ex, bool only_remote_allowed) - : ctx_params(common_params_parser_init(default_params, ex)) { - common_params_add_preset_options(ctx_params.options); - key_to_opt = get_map_key_opt(ctx_params); - - // setup allowed keys if only_remote_allowed is true - if (only_remote_allowed) { - filter_allowed_keys = true; - allowed_keys = get_remote_preset_whitelist(key_to_opt); - } -} - -common_presets common_preset_context::load_from_ini(const std::string & path, common_preset & global) const { - common_presets out; - auto ini_data = parse_ini_from_file(path); - - for (auto section : ini_data) { - common_preset preset; - if (section.first.empty()) { - preset.name = COMMON_PRESET_DEFAULT_NAME; - } else { - preset.name = section.first; - } - LOG_DBG("loading preset: %s\n", preset.name.c_str()); - for (const auto & [key, value] : section.second) { - if (key == "version") { - // skip version key (reserved for future use) - continue; - } - - LOG_DBG("option: %s = %s\n", key.c_str(), value.c_str()); - if (filter_allowed_keys && allowed_keys.find(key) == allowed_keys.end()) { - throw std::runtime_error(string_format( - "option '%s' is not allowed in remote presets", - key.c_str() - )); - } - if (key_to_opt.find(key) != key_to_opt.end()) { - const auto & opt = key_to_opt.at(key); - if (is_bool_arg(opt)) { - preset.options[opt] = parse_bool_arg(opt, key, value); - } else { - preset.options[opt] = value; - } - LOG_DBG("accepted option: %s = %s\n", key.c_str(), preset.options[opt].c_str()); - } else { - throw std::runtime_error(string_format( - "option '%s' not recognized in preset '%s'", - key.c_str(), preset.name.c_str() - )); - } - } - - if (preset.name == "*") { - // handle global preset - global = preset; - } else { - out[preset.name] = preset; - } - } - - return out; -} - -common_presets common_preset_context::load_from_cache() const { - common_presets out; - - auto cached_models = common_list_cached_models(); - for (const auto & model : cached_models) { - common_preset preset; - preset.name = model.to_string(); - preset.set_option(*this, "LLAMA_ARG_HF_REPO", model.to_string()); - out[preset.name] = preset; - } - - return out; -} - -struct local_model { - std::string name; - std::string path; - std::string path_mmproj; -}; - -common_presets common_preset_context::load_from_models_dir(const std::string & models_dir) const { - if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) { - throw std::runtime_error(string_format("error: '%s' does not exist or is not a directory\n", models_dir.c_str())); - } - - std::vector models; - auto scan_subdir = [&models](const std::string & subdir_path, const std::string & name) { - auto files = fs_list(subdir_path, false); - common_file_info model_file; - common_file_info first_shard_file; - common_file_info mmproj_file; - for (const auto & file : files) { - if (string_ends_with(file.name, ".gguf")) { - if (file.name.find("mmproj") != std::string::npos) { - mmproj_file = file; - } else if (file.name.find("-00001-of-") != std::string::npos) { - first_shard_file = file; - } else { - model_file = file; - } - } - } - // single file model - local_model model{ - /* name */ name, - /* path */ first_shard_file.path.empty() ? model_file.path : first_shard_file.path, - /* path_mmproj */ mmproj_file.path // can be empty - }; - if (!model.path.empty()) { - models.push_back(model); - } - }; - - auto files = fs_list(models_dir, true); - for (const auto & file : files) { - if (file.is_dir) { - scan_subdir(file.path, file.name); - } else if (string_ends_with(file.name, ".gguf")) { - // single file model - std::string name = file.name; - string_replace_all(name, ".gguf", ""); - local_model model{ - /* name */ name, - /* path */ file.path, - /* path_mmproj */ "" - }; - models.push_back(model); - } - } - - // convert local models to presets - common_presets out; - for (const auto & model : models) { - common_preset preset; - preset.name = model.name; - preset.set_option(*this, "LLAMA_ARG_MODEL", model.path); - if (!model.path_mmproj.empty()) { - preset.set_option(*this, "LLAMA_ARG_MMPROJ", model.path_mmproj); - } - out[preset.name] = preset; - } - - return out; -} - -common_preset common_preset_context::load_from_args(int argc, char ** argv) const { - common_preset preset; - preset.name = COMMON_PRESET_DEFAULT_NAME; - - bool ok = common_params_to_map(argc, argv, ctx_params.ex, preset.options); - if (!ok) { - throw std::runtime_error("failed to parse CLI arguments into preset"); - } - - return preset; -} - -common_presets common_preset_context::cascade(const common_presets & base, const common_presets & added) const { - common_presets out = base; // copy - for (const auto & [name, preset_added] : added) { - if (out.find(name) != out.end()) { - // if exists, merge - common_preset & target = out[name]; - target.merge(preset_added); - } else { - // otherwise, add directly - out[name] = preset_added; - } - } - return out; -} - -common_presets common_preset_context::cascade(const common_preset & base, const common_presets & presets) const { - common_presets out; - for (const auto & [name, preset] : presets) { - common_preset tmp = base; // copy - tmp.name = name; - tmp.merge(preset); - out[name] = std::move(tmp); - } - return out; -} diff --git a/common/preset.h b/common/preset.h deleted file mode 100644 index 11ba6ef81240..000000000000 --- a/common/preset.h +++ /dev/null @@ -1,83 +0,0 @@ -#pragma once - -#include "common.h" -#include "arg.h" - -#include -#include -#include -#include - -// -// INI preset parser and writer -// - -constexpr const char * COMMON_PRESET_DEFAULT_NAME = "default"; - -struct common_preset_context; - -struct common_preset { - std::string name; - - // options are stored as common_arg to string mapping, representing CLI arg and its value - std::map options; - - // convert preset to CLI argument list - std::vector to_args(const std::string & bin_path = "") const; - - // convert preset to INI format string - std::string to_ini() const; - - // TODO: maybe implement to_env() if needed - - // modify preset options where argument is identified by its env variable - void set_option(const common_preset_context & ctx, const std::string & env, const std::string & value); - - // unset option by its env variable - void unset_option(const std::string & env); - - // get option value by its env variable, return false if not found - bool get_option(const std::string & env, std::string & value) const; - - // merge another preset into this one, overwriting existing options - void merge(const common_preset & other); - - // apply preset options to common_params - void apply_to_params(common_params & params) const; -}; - -// interface for multiple presets in one file -using common_presets = std::map; - -// context for loading and editing presets -struct common_preset_context { - common_params default_params; // unused for now - common_params_context ctx_params; - std::map key_to_opt; - - bool filter_allowed_keys = false; - std::set allowed_keys; - - // if only_remote_allowed is true, only accept whitelisted keys - common_preset_context(llama_example ex, bool only_remote_allowed = false); - - // load presets from INI file - common_presets load_from_ini(const std::string & path, common_preset & global) const; - - // generate presets from cached models - common_presets load_from_cache() const; - - // generate presets from local models directory - // for the directory structure, see "Using multiple models" in server/README.md - common_presets load_from_models_dir(const std::string & models_dir) const; - - // generate one preset from CLI arguments - common_preset load_from_args(int argc, char ** argv) const; - - // cascade multiple presets if exist on both: base < added - // if preset does not exist in base, it will be added without modification - common_presets cascade(const common_presets & base, const common_presets & added) const; - - // apply presets over a base preset (same idea as CSS cascading) - common_presets cascade(const common_preset & base, const common_presets & presets) const; -}; diff --git a/common/reasoning-budget-rust/Cargo.toml b/common/reasoning-budget-rust/Cargo.toml new file mode 100644 index 000000000000..834da73a5f21 --- /dev/null +++ b/common/reasoning-budget-rust/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "llama-common-reasoning-budget-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[lib] +crate-type = ["staticlib", "rlib"] + +[dependencies] diff --git a/common/reasoning-budget-rust/src/lib.rs b/common/reasoning-budget-rust/src/lib.rs new file mode 100644 index 000000000000..e630c1a8764b --- /dev/null +++ b/common/reasoning-budget-rust/src/lib.rs @@ -0,0 +1,420 @@ +use std::ffi::{c_char, c_void}; +use std::ptr; +use std::slice; + +type LlamaToken = i32; + +#[repr(C)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ReasoningBudgetState { + Idle = 0, + Counting = 1, + Forcing = 2, + WaitingUtf8 = 3, + Done = 4, +} + +#[repr(C)] +pub struct LlamaVocab { + _private: [u8; 0], +} + +#[repr(C)] +pub struct LlamaSampler { + iface: *mut LlamaSamplerI, + ctx: *mut c_void, +} + +#[repr(C)] +pub struct LlamaTokenData { + id: LlamaToken, + logit: f32, + p: f32, +} + +#[repr(C)] +pub struct LlamaTokenDataArray { + data: *mut LlamaTokenData, + size: usize, + selected: i64, + sorted: bool, +} + +#[repr(C)] +pub struct LlamaSamplerData { + logits: *mut c_void, + probs: *mut c_void, + sampled: *mut c_void, + candidates: *mut c_void, +} + +type NameFn = unsafe extern "C" fn(*const LlamaSampler) -> *const c_char; +type AcceptFn = unsafe extern "C" fn(*mut LlamaSampler, LlamaToken); +type ApplyFn = unsafe extern "C" fn(*mut LlamaSampler, *mut LlamaTokenDataArray); +type ResetFn = unsafe extern "C" fn(*mut LlamaSampler); +type CloneFn = unsafe extern "C" fn(*const LlamaSampler) -> *mut LlamaSampler; +type FreeFn = unsafe extern "C" fn(*mut LlamaSampler); +type BackendInitFn = unsafe extern "C" fn(*mut LlamaSampler, *mut c_void) -> bool; +type BackendAcceptFn = unsafe extern "C" fn(*mut LlamaSampler, *mut c_void, *mut c_void, *mut c_void); +type BackendApplyFn = + unsafe extern "C" fn(*mut LlamaSampler, *mut c_void, *mut c_void, *mut LlamaSamplerData); +type BackendSetInputFn = unsafe extern "C" fn(*mut LlamaSampler); + +#[repr(C)] +pub struct LlamaSamplerI { + name: Option, + accept: Option, + apply: Option, + reset: Option, + clone: Option, + free: Option, + backend_init: Option, + backend_accept: Option, + backend_apply: Option, + backend_set_input: Option, +} + +extern "C" { + fn llama_sampler_init(iface: *mut LlamaSamplerI, ctx: *mut c_void) -> *mut LlamaSampler; + fn llama_token_to_piece( + vocab: *const LlamaVocab, + token: LlamaToken, + buf: *mut c_char, + length: i32, + lstrip: i32, + special: bool, + ) -> i32; +} + +#[derive(Clone)] +struct TokenMatcher { + tokens: Vec, + pos: usize, +} + +impl TokenMatcher { + fn advance(&mut self, token: LlamaToken) -> bool { + if self.tokens.is_empty() { + return false; + } + + if token == self.tokens[self.pos] { + self.pos += 1; + if self.pos >= self.tokens.len() { + self.pos = 0; + return true; + } + } else { + self.pos = 0; + if token == self.tokens[0] { + self.pos = 1; + } + } + false + } + + fn reset(&mut self) { + self.pos = 0; + } +} + +struct ReasoningBudgetCtx { + vocab: *const LlamaVocab, + start_matcher: TokenMatcher, + end_matcher: TokenMatcher, + forced_tokens: Vec, + budget: i32, + remaining: i32, + state: ReasoningBudgetState, + force_pos: usize, +} + +const NAME: &[u8] = b"reasoning-budget\0"; + +#[no_mangle] +pub unsafe extern "C" fn llama_common_reasoning_budget_init_rust( + vocab: *const LlamaVocab, + start_tokens: *const LlamaToken, + start_len: usize, + end_tokens: *const LlamaToken, + end_len: usize, + forced_tokens: *const LlamaToken, + forced_len: usize, + budget: i32, + initial_state: ReasoningBudgetState, +) -> *mut LlamaSampler { + let start = copy_tokens(start_tokens, start_len); + let end = copy_tokens(end_tokens, end_len); + let forced = copy_tokens(forced_tokens, forced_len); + init_state(vocab, start, end, forced, budget, initial_state) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_reasoning_budget_get_state_rust( + smpl: *const LlamaSampler, +) -> ReasoningBudgetState { + if smpl.is_null() { + return ReasoningBudgetState::Idle; + } + let ctx = (*smpl).ctx as *const ReasoningBudgetCtx; + if ctx.is_null() { + return ReasoningBudgetState::Idle; + } + (*ctx).state +} + +unsafe fn copy_tokens(ptr: *const LlamaToken, len: usize) -> Vec { + if ptr.is_null() || len == 0 { + Vec::new() + } else { + slice::from_raw_parts(ptr, len).to_vec() + } +} + +unsafe fn init_state( + vocab: *const LlamaVocab, + start_tokens: Vec, + end_tokens: Vec, + forced_tokens: Vec, + budget: i32, + mut initial_state: ReasoningBudgetState, +) -> *mut LlamaSampler { + if initial_state == ReasoningBudgetState::Counting && budget <= 0 { + initial_state = ReasoningBudgetState::Forcing; + } + + let ctx = Box::new(ReasoningBudgetCtx { + vocab, + start_matcher: TokenMatcher { + tokens: start_tokens, + pos: 0, + }, + end_matcher: TokenMatcher { + tokens: end_tokens, + pos: 0, + }, + forced_tokens, + budget, + remaining: budget, + state: initial_state, + force_pos: 0, + }); + + llama_sampler_init( + ptr::addr_of_mut!(LLAMA_REASONING_BUDGET_I), + Box::into_raw(ctx) as *mut c_void, + ) +} + +unsafe extern "C" fn reasoning_budget_name(_smpl: *const LlamaSampler) -> *const c_char { + NAME.as_ptr() as *const c_char +} + +unsafe extern "C" fn reasoning_budget_accept(smpl: *mut LlamaSampler, token: LlamaToken) { + let ctx = &mut *((*smpl).ctx as *mut ReasoningBudgetCtx); + + match ctx.state { + ReasoningBudgetState::Idle => { + if ctx.start_matcher.advance(token) { + ctx.state = ReasoningBudgetState::Counting; + ctx.remaining = ctx.budget; + if ctx.remaining <= 0 { + ctx.state = ReasoningBudgetState::Forcing; + ctx.force_pos = 0; + } + } + } + ReasoningBudgetState::Counting | ReasoningBudgetState::WaitingUtf8 => { + if ctx.end_matcher.advance(token) { + ctx.state = ReasoningBudgetState::Done; + return; + } + + let utf8_complete = token_piece_is_utf8_complete(ctx.vocab, token); + + if ctx.state == ReasoningBudgetState::WaitingUtf8 { + if utf8_complete { + ctx.state = ReasoningBudgetState::Forcing; + ctx.force_pos = 0; + ctx.end_matcher.reset(); + } + } else { + ctx.remaining -= 1; + if ctx.remaining <= 0 { + if utf8_complete { + ctx.state = ReasoningBudgetState::Forcing; + ctx.force_pos = 0; + ctx.end_matcher.reset(); + } else { + ctx.state = ReasoningBudgetState::WaitingUtf8; + ctx.end_matcher.reset(); + } + } + } + } + ReasoningBudgetState::Forcing => { + ctx.force_pos += 1; + if ctx.force_pos >= ctx.forced_tokens.len() { + ctx.state = ReasoningBudgetState::Done; + } + } + ReasoningBudgetState::Done => { + if ctx.start_matcher.advance(token) { + ctx.state = ReasoningBudgetState::Counting; + ctx.remaining = ctx.budget; + ctx.end_matcher.reset(); + + if ctx.remaining <= 0 { + ctx.state = ReasoningBudgetState::Forcing; + ctx.force_pos = 0; + } + } + } + } +} + +unsafe extern "C" fn reasoning_budget_apply( + smpl: *mut LlamaSampler, + cur_p: *mut LlamaTokenDataArray, +) { + let ctx = &mut *((*smpl).ctx as *mut ReasoningBudgetCtx); + if ctx.state != ReasoningBudgetState::Forcing || ctx.force_pos >= ctx.forced_tokens.len() { + return; + } + + let forced = ctx.forced_tokens[ctx.force_pos]; + let cur = &mut *cur_p; + let data = slice::from_raw_parts_mut(cur.data, cur.size); + for item in data { + if item.id != forced { + item.logit = f32::NEG_INFINITY; + } + } +} + +unsafe extern "C" fn reasoning_budget_reset(smpl: *mut LlamaSampler) { + let ctx = &mut *((*smpl).ctx as *mut ReasoningBudgetCtx); + ctx.state = ReasoningBudgetState::Idle; + ctx.remaining = ctx.budget; + ctx.start_matcher.reset(); + ctx.end_matcher.reset(); + ctx.force_pos = 0; +} + +unsafe extern "C" fn reasoning_budget_clone(smpl: *const LlamaSampler) -> *mut LlamaSampler { + let ctx = &*((*smpl).ctx as *const ReasoningBudgetCtx); + init_state( + ctx.vocab, + ctx.start_matcher.tokens.clone(), + ctx.end_matcher.tokens.clone(), + ctx.forced_tokens.clone(), + ctx.budget, + ctx.state, + ) +} + +unsafe extern "C" fn reasoning_budget_free(smpl: *mut LlamaSampler) { + let ctx = (*smpl).ctx as *mut ReasoningBudgetCtx; + if !ctx.is_null() { + drop(Box::from_raw(ctx)); + } +} + +unsafe fn token_piece_is_utf8_complete(vocab: *const LlamaVocab, token: LlamaToken) -> bool { + if vocab.is_null() { + return true; + } + + let mut buf = vec![0u8; 32]; + let n = llama_token_to_piece( + vocab, + token, + buf.as_mut_ptr() as *mut c_char, + buf.len() as i32, + 0, + false, + ); + let len = if n < 0 { + let needed = (-n) as usize; + buf.resize(needed, 0); + let check = llama_token_to_piece( + vocab, + token, + buf.as_mut_ptr() as *mut c_char, + buf.len() as i32, + 0, + false, + ); + if check != -n { + return true; + } + needed + } else { + n as usize + }; + + common_utf8_is_complete(&buf[..len]) +} + +fn common_utf8_is_complete(bytes: &[u8]) -> bool { + if bytes.is_empty() { + return true; + } + for i in 1..=bytes.len().min(4) { + let c = bytes[bytes.len() - i]; + if (c & 0xC0) != 0x80 { + let expected = if c >= 0xF0 { + 4 + } else if c >= 0xE0 { + 3 + } else if c >= 0xC0 { + 2 + } else { + 1 + }; + return i >= expected; + } + } + false +} + +static mut LLAMA_REASONING_BUDGET_I: LlamaSamplerI = LlamaSamplerI { + name: Some(reasoning_budget_name), + accept: Some(reasoning_budget_accept), + apply: Some(reasoning_budget_apply), + reset: Some(reasoning_budget_reset), + clone: Some(reasoning_budget_clone), + free: Some(reasoning_budget_free), + backend_init: None, + backend_accept: None, + backend_apply: None, + backend_set_input: None, +}; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn matcher_handles_overlap() { + let mut matcher = TokenMatcher { + tokens: vec![1, 2, 1], + pos: 0, + }; + assert!(!matcher.advance(1)); + assert!(!matcher.advance(2)); + assert!(matcher.advance(1)); + assert!(!matcher.advance(2)); + assert!(!matcher.advance(1)); + assert_eq!(matcher.pos, 1); + } + + #[test] + fn utf8_completion_matches_common_cases() { + assert!(common_utf8_is_complete(b"hello")); + assert!(common_utf8_is_complete(b"abc\xC3\xA9")); + assert!(!common_utf8_is_complete(&[0xC2])); + assert!(!common_utf8_is_complete(&[0xE2, 0x80])); + assert!(!common_utf8_is_complete(&[0x80])); + } +} diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp deleted file mode 100644 index c6e1f86c91e9..000000000000 --- a/common/reasoning-budget.cpp +++ /dev/null @@ -1,250 +0,0 @@ -#include "reasoning-budget.h" -#include "common.h" -#include "unicode.h" - -#include "log.h" - -#include -#include -#include -#include - -struct token_matcher { - std::vector tokens; - size_t pos = 0; - - bool advance(llama_token token) { - if (tokens.empty()) { - return false; - } - - if (token == tokens[pos]) { - pos++; - if (pos >= tokens.size()) { - pos = 0; - return true; - } - } else { - pos = 0; - if (token == tokens[0]) { - pos = 1; - } - } - return false; - } - - void reset() { pos = 0; } -}; - -struct common_reasoning_budget_ctx { - const llama_vocab * vocab; - - token_matcher start_matcher; - token_matcher end_matcher; - std::vector forced_tokens; - - int32_t budget; // maximum tokens in reasoning block - int32_t remaining; // tokens remaining in budget - - common_reasoning_budget_state state; - - // for forcing - size_t force_pos; // next position in forced_tokens to force -}; - -static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) { - return "reasoning-budget"; -} - -static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_token token) { - auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx; - - switch (ctx->state) { - case REASONING_BUDGET_IDLE: - { - if (ctx->start_matcher.advance(token)) { - ctx->state = REASONING_BUDGET_COUNTING; - ctx->remaining = ctx->budget; - LOG_INF("reasoning-budget: activated, budget=%d tokens\n", ctx->budget); - - if (ctx->remaining <= 0) { - ctx->state = REASONING_BUDGET_FORCING; - ctx->force_pos = 0; - LOG_INF("reasoning-budget: budget=0, forcing immediately\n"); - } - } - break; - } - case REASONING_BUDGET_COUNTING: - case REASONING_BUDGET_WAITING_UTF8: - { - if (ctx->end_matcher.advance(token)) { - ctx->state = REASONING_BUDGET_DONE; - LOG_INF("reasoning-budget: deactivated (natural end)\n"); - break; - } - - bool utf8_complete = true; - if (ctx->vocab != nullptr) { - const std::string piece = common_token_to_piece(ctx->vocab, token, false); - utf8_complete = common_utf8_is_complete(piece); - } - - if (ctx->state == REASONING_BUDGET_WAITING_UTF8) { - if (utf8_complete) { - ctx->state = REASONING_BUDGET_FORCING; - ctx->force_pos = 0; - ctx->end_matcher.reset(); - LOG_INF("reasoning-budget: UTF-8 complete, now forcing end sequence\n"); - } - } else if (ctx->state == REASONING_BUDGET_COUNTING) { - ctx->remaining--; - if (ctx->remaining <= 0) { - if (utf8_complete) { - ctx->state = REASONING_BUDGET_FORCING; - ctx->force_pos = 0; - ctx->end_matcher.reset(); - LOG_INF("reasoning-budget: budget exhausted, forcing end sequence\n"); - } else { - ctx->state = REASONING_BUDGET_WAITING_UTF8; - ctx->end_matcher.reset(); - LOG_INF("reasoning-budget: budget exhausted, waiting for UTF-8 completion\n"); - } - } - } - break; - } - case REASONING_BUDGET_FORCING: - ctx->force_pos++; - if (ctx->force_pos >= ctx->forced_tokens.size()) { - ctx->state = REASONING_BUDGET_DONE; - LOG_INF("reasoning-budget: forced sequence complete, done\n"); - } - break; - case REASONING_BUDGET_DONE: - // Re-arm on a new start tag: some models emit multiple blocks - // per response, and each should get a fresh budget window. - if (ctx->start_matcher.advance(token)) { - ctx->state = REASONING_BUDGET_COUNTING; - ctx->remaining = ctx->budget; - ctx->end_matcher.reset(); - LOG_INF("reasoning-budget: re-activated on new start tag, budget=%d tokens\n", ctx->budget); - - if (ctx->remaining <= 0) { - ctx->state = REASONING_BUDGET_FORCING; - ctx->force_pos = 0; - LOG_INF("reasoning-budget: budget=0, forcing immediately\n"); - } - } - break; - } -} - -static void common_reasoning_budget_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) { - auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx; - - if (ctx->state != REASONING_BUDGET_FORCING) { - // passthrough — don't modify logits - return; - } - - if (ctx->force_pos >= ctx->forced_tokens.size()) { - return; - } - - const llama_token forced = ctx->forced_tokens[ctx->force_pos]; - - // set all logits to -inf except the forced token - for (size_t i = 0; i < cur_p->size; i++) { - if (cur_p->data[i].id != forced) { - cur_p->data[i].logit = -INFINITY; - } - } -} - -static void common_reasoning_budget_reset(struct llama_sampler * smpl) { - auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx; - ctx->state = REASONING_BUDGET_IDLE; - ctx->remaining = ctx->budget; - ctx->start_matcher.reset(); - ctx->end_matcher.reset(); - ctx->force_pos = 0; -} - -// forward declaration for use in clone -static struct llama_sampler * common_reasoning_budget_init_state( - const struct llama_vocab * vocab, const std::vector & start_tokens, - const std::vector & end_tokens, const std::vector & forced_tokens, - int32_t budget, common_reasoning_budget_state initial_state); - -static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl) { - const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx; - return common_reasoning_budget_init_state( - ctx->vocab, - ctx->start_matcher.tokens, - ctx->end_matcher.tokens, - ctx->forced_tokens, - ctx->budget, - ctx->state); -} - -static void common_reasoning_budget_free(struct llama_sampler * smpl) { - delete (common_reasoning_budget_ctx *) smpl->ctx; -} - -static struct llama_sampler_i common_reasoning_budget_i = { - /* .name = */ common_reasoning_budget_name, - /* .accept = */ common_reasoning_budget_accept, - /* .apply = */ common_reasoning_budget_apply, - /* .reset = */ common_reasoning_budget_reset, - /* .clone = */ common_reasoning_budget_clone, - /* .free = */ common_reasoning_budget_free, - /* .backend_init = */ nullptr, - /* .backend_accept = */ nullptr, - /* .backend_apply = */ nullptr, - /* .backend_set_input = */ nullptr, -}; - -static struct llama_sampler * common_reasoning_budget_init_state( - const struct llama_vocab * vocab, - const std::vector & start_tokens, - const std::vector & end_tokens, - const std::vector & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state) { - // promote COUNTING with budget <= 0 to FORCING - if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) { - initial_state = REASONING_BUDGET_FORCING; - } - - return llama_sampler_init( - /* .iface = */ &common_reasoning_budget_i, - /* .ctx = */ new common_reasoning_budget_ctx { - /* .vocab = */ vocab, - /* .start_matcher = */ { start_tokens, 0 }, - /* .end_matcher = */ { end_tokens, 0 }, - /* .forced_tokens = */ forced_tokens, - /* .budget = */ budget, - /* .remaining = */ budget, - /* .state = */ initial_state, - /* .force_pos = */ 0, - } - ); -} - -struct llama_sampler * common_reasoning_budget_init( - const struct llama_vocab * vocab, - const std::vector & start_tokens, - const std::vector & end_tokens, - const std::vector & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state) { - return common_reasoning_budget_init_state(vocab, start_tokens, end_tokens, forced_tokens, budget, initial_state); -} - -common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) { - if (!smpl) { - return REASONING_BUDGET_IDLE; - } - return ((const common_reasoning_budget_ctx *)smpl->ctx)->state; -} diff --git a/common/reasoning-budget.h b/common/reasoning-budget.h deleted file mode 100644 index ef37f46ee4dc..000000000000 --- a/common/reasoning-budget.h +++ /dev/null @@ -1,42 +0,0 @@ -#pragma once - -#include "llama.h" - -#include -#include - -enum common_reasoning_budget_state { - REASONING_BUDGET_IDLE, // waiting for start sequence - REASONING_BUDGET_COUNTING, // counting down tokens - REASONING_BUDGET_FORCING, // forcing budget message + end sequence - REASONING_BUDGET_WAITING_UTF8, // budget exhausted, waiting for UTF-8 completion - REASONING_BUDGET_DONE, // passthrough forever -}; - -// Creates a reasoning budget sampler that limits token generation inside a -// reasoning block (e.g. between and ). -// -// State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE -// IDLE: passthrough, watching for start_tokens sequence -// COUNTING: counting down remaining tokens, watching for natural end_tokens -// WAITING_UTF8: budget exhausted, allowing tokens to complete a UTF-8 sequence -// FORCING: forces forced_tokens token-by-token (all other logits -> -inf) -// DONE: passthrough forever -// -// Parameters: -// vocab - vocabulary (used for UTF-8 boundary detection; can be nullptr) -// start_tokens - token sequence that activates counting -// end_tokens - token sequence for natural deactivation -// forced_tokens - token sequence forced when budget expires -// budget - max tokens allowed in the reasoning block -// initial_state - initial state -// -struct llama_sampler * common_reasoning_budget_init( - const struct llama_vocab * vocab, - const std::vector & start_tokens, - const std::vector & end_tokens, - const std::vector & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE); - -common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl); diff --git a/common/regex-partial-rust/Cargo.toml b/common/regex-partial-rust/Cargo.toml new file mode 100644 index 000000000000..426b88bde2ac --- /dev/null +++ b/common/regex-partial-rust/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "llama-common-regex-partial-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[lib] +crate-type = ["staticlib", "rlib"] diff --git a/common/regex-partial-rust/src/lib.rs b/common/regex-partial-rust/src/lib.rs new file mode 100644 index 000000000000..27b50df78757 --- /dev/null +++ b/common/regex-partial-rust/src/lib.rs @@ -0,0 +1,249 @@ +use std::slice; + +#[repr(C)] +pub struct RegexPartialString { + data: *mut u8, + len: usize, +} + +#[no_mangle] +pub unsafe extern "C" fn llama_regex_to_reversed_partial_regex( + pattern: *const u8, + len: usize, + error: *mut RegexPartialString, +) -> RegexPartialString { + if pattern.is_null() && len != 0 { + return error_result("null pattern pointer", error); + } + + let pattern = unsafe { slice::from_raw_parts(pattern, len) }; + match regex_to_reversed_partial_regex_bytes(pattern) { + Ok(value) => into_ffi(value.into_bytes()), + Err(err) => error_result(&err, error), + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_regex_partial_string_free(value: RegexPartialString) { + if !value.data.is_null() { + unsafe { + drop(Vec::from_raw_parts(value.data, value.len, value.len)); + } + } +} + +fn error_result(message: &str, error: *mut RegexPartialString) -> RegexPartialString { + if !error.is_null() { + unsafe { + *error = into_ffi(message.as_bytes().to_vec()); + } + } + RegexPartialString { + data: std::ptr::null_mut(), + len: 0, + } +} + +fn into_ffi(mut bytes: Vec) -> RegexPartialString { + let len = bytes.len(); + let data = bytes.as_mut_ptr(); + std::mem::forget(bytes); + RegexPartialString { data, len } +} + +pub fn regex_to_reversed_partial_regex_bytes(pattern: &[u8]) -> Result { + let mut parser = Parser { pattern, pos: 0 }; + let res = parser.process()?; + if parser.pos != pattern.len() { + return Err("Unmatched '(' in pattern".to_string()); + } + Ok(format!("^({res})")) +} + +struct Parser<'a> { + pattern: &'a [u8], + pos: usize, +} + +impl Parser<'_> { + fn process(&mut self) -> Result { + let mut alternatives = vec![Vec::::new()]; + + while self.pos < self.pattern.len() { + let byte = self.pattern[self.pos]; + let sequence = alternatives.last_mut().unwrap(); + + match byte { + b'[' => { + let start = self.pos; + self.pos += 1; + while self.pos < self.pattern.len() { + if self.pattern[self.pos] == b'\\' { + self.pos += 1; + if self.pos < self.pattern.len() { + self.pos += 1; + } + } else if self.pattern[self.pos] == b']' { + break; + } else { + self.pos += 1; + } + } + if self.pos == self.pattern.len() { + return Err("Unmatched '[' in pattern".to_string()); + } + self.pos += 1; + sequence.push(bytes_to_string(&self.pattern[start..self.pos])); + } + b'*' | b'?' | b'+' => { + if sequence.is_empty() { + return Err("Quantifier without preceding element".to_string()); + } + sequence.last_mut().unwrap().push(byte as char); + let is_star = byte == b'*'; + self.pos += 1; + if is_star && self.pos < self.pattern.len() && self.pattern[self.pos] == b'?' { + self.pos += 1; + } + } + b'{' => { + if sequence.is_empty() { + return Err("Repetition without preceding element".to_string()); + } + self.pos += 1; + let start = self.pos; + while self.pos < self.pattern.len() && self.pattern[self.pos] != b'}' { + self.pos += 1; + } + if self.pos == self.pattern.len() { + return Err("Unmatched '{' in pattern".to_string()); + } + let parts = bytes_to_string(&self.pattern[start..self.pos]) + .split(',') + .map(|part| part.to_string()) + .collect::>(); + self.pos += 1; + + if parts.len() > 2 { + return Err("Invalid repetition range in pattern".to_string()); + } + let min = parse_opt_int(&parts[0], Some(0))? + .ok_or_else(|| "Invalid repetition range in pattern".to_string())?; + let max = if parts.len() == 1 { + Some(min) + } else { + parse_opt_int(&parts[1], None)? + }; + if let Some(max) = max { + if max < min { + return Err("Invalid repetition range in pattern".to_string()); + } + } + + let part = sequence.pop().unwrap(); + for _ in 0..min { + sequence.push(part.clone()); + } + if let Some(max) = max { + for _ in min..max { + sequence.push(format!("{part}?")); + } + } else { + sequence.push(format!("{part}*")); + } + } + b'(' => { + self.pos += 1; + if self.pos + 1 < self.pattern.len() + && self.pattern[self.pos] == b'?' + && self.pattern[self.pos + 1] == b':' + { + self.pos += 2; + } + let sub = self.process()?; + if self.pos >= self.pattern.len() || self.pattern[self.pos] != b')' { + return Err("Unmatched '(' in pattern".to_string()); + } + self.pos += 1; + sequence.push(format!("(?:{sub})")); + } + b')' => break, + b'|' => { + self.pos += 1; + alternatives.push(Vec::new()); + } + b'\\' => { + self.pos += 1; + if self.pos < self.pattern.len() { + sequence.push(format!("\\{}", self.pattern[self.pos] as char)); + self.pos += 1; + } + } + _ => { + sequence.push((byte as char).to_string()); + self.pos += 1; + } + } + } + + let mut res_alts = Vec::with_capacity(alternatives.len()); + for parts in alternatives { + let mut res = String::new(); + for _ in 0..parts.len().saturating_sub(1) { + res.push_str("(?:"); + } + for (idx, part) in parts.iter().rev().enumerate() { + res.push_str(part); + if idx + 1 != parts.len() { + res.push_str(")?"); + } + } + res_alts.push(res); + } + + Ok(res_alts.join("|")) + } +} + +fn parse_opt_int(value: &str, default: Option) -> Result, String> { + if value.is_empty() { + return Ok(default); + } + value + .parse::() + .map(Some) + .map_err(|_| "Invalid repetition range in pattern".to_string()) +} + +fn bytes_to_string(bytes: &[u8]) -> String { + bytes.iter().map(|&byte| byte as char).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn transform(pattern: &str) -> String { + regex_to_reversed_partial_regex_bytes(pattern.as_bytes()).unwrap() + } + + #[test] + fn reverses_simple_sequences_and_groups() { + assert_eq!(transform("abcd"), "^((?:(?:(?:d)?c)?b)?a)"); + assert_eq!(transform("a|b"), "^(a|b)"); + assert_eq!(transform("a(bc)d"), "^((?:(?:d)?(?:(?:c)?b))?a)"); + } + + #[test] + fn expands_repetition_ranges() { + assert_eq!( + transform("ab{2,4}c"), + "^((?:(?:(?:(?:(?:c)?b?)?b?)?b)?b)?a)" + ); + } + + #[test] + fn rejects_unmatched_character_classes() { + assert!(regex_to_reversed_partial_regex_bytes(b"[abc").is_err()); + } +} diff --git a/common/regex-partial.cpp b/common/regex-partial.cpp deleted file mode 100644 index bd9034e93192..000000000000 --- a/common/regex-partial.cpp +++ /dev/null @@ -1,204 +0,0 @@ -#include "regex-partial.h" -#include "common.h" -#include -#include - -common_regex::common_regex(const std::string & pattern) : - pattern(pattern), - rx(pattern), - rx_reversed_partial(regex_to_reversed_partial_regex(pattern)) {} - -common_regex_match common_regex::search(const std::string & input, size_t pos, bool as_match) const { - std::smatch match; - if (pos > input.size()) { - throw std::runtime_error("Position out of bounds"); - } - auto start = input.begin() + pos; - auto found = as_match - ? std::regex_match(start, input.end(), match, rx) - : std::regex_search(start, input.end(), match, rx); - if (found) { - common_regex_match res; - res.type = COMMON_REGEX_MATCH_TYPE_FULL; - for (size_t i = 0; i < match.size(); ++i) { - auto begin = pos + match.position(i); - res.groups.emplace_back(begin, begin + match.length(i)); - } - return res; - } - std::match_results srmatch; - if (std::regex_search(input.rbegin(), input.rend() - pos, srmatch, rx_reversed_partial, std::regex_constants::match_continuous)) { - auto group = srmatch[1].str(); - if (group.length() != 0) { - auto it = srmatch[1].second.base(); - // auto position = static_cast(std::distance(input.begin(), it)); - if ((!as_match) || it == input.begin()) { - common_regex_match res; - res.type = COMMON_REGEX_MATCH_TYPE_PARTIAL; - const size_t begin = std::distance(input.begin(), it); - const size_t end = input.size(); - if (begin == std::string::npos || end == std::string::npos || begin > end) { - throw std::runtime_error("Invalid range"); - } - res.groups.push_back({begin, end}); - return res; - } - } - } - return {}; -} - -/* - Transforms a regex pattern to a partial match pattern that operates on a reversed input string to find partial final matches of the original pattern. - - Ideally we'd like to use boost::match_partial (https://beta.boost.org/doc/libs/1_59_0/libs/regex/doc/html/boost_regex/partial_matches.html) - to see if a string ends with a partial regex match, but but it's not in std::regex yet. - Instead, we'll the regex into a partial match regex operating as a full match on the reverse iterators of the input. - - - /abcd/ -> ^(dcba|cba|ba|a) -> ^((?:(?:(?:(?:d)?c)?b)?a) - - /a|b/ -> ^(a|b) - - /a*?/ -> error, could match "" - - /a*b/ -> ^((?:b)?a*+) (final repetitions become eager) - - /.*?ab/ -> ^((?:b)?a) (omit .*) - - /a.*?b/ -> ^((?:b)?.*?a) (keep reluctant matches) - - /a(bc)d/ -> ^((?:(?:d)?(?:(?:c)?b))?a) - - /a(bc|de)/ -> ^((?:(?:(?:e)?d)?|(?:(?:c)?b)?)?a) - - /ab{2,4}c/ -> ^cbbb?b?a -> ^((?:(?:(?:(?:(?:c)?b)?b)?b?)?b?)?a) - - The regex will match a reversed string fully, and the end of the first (And only) capturing group will indicate the reversed start of the original partial pattern. - All other groups are turned into non-capturing groups, and reluctant quantifiers are ignored. -*/ -std::string regex_to_reversed_partial_regex(const std::string & pattern) { - auto it = pattern.begin(); - const auto end = pattern.end(); - - std::function process = [&]() { - std::vector> alternatives(1); - std::vector * sequence = &alternatives.back(); - - while (it != end) { - if (*it == '[') { - auto start = it; - ++it; - while (it != end) { - if ((*it == '\\') && (++it != end)) { - ++it; - } else if ((it != end) && (*it == ']')) { - break; - } else { - ++it; - } - } - if (it == end) { - throw std::runtime_error("Unmatched '[' in pattern"); - } - ++it; - sequence->push_back(std::string(start, it)); - } else if (*it == '*' || *it == '?' || *it == '+') { - if (sequence->empty()) { - throw std::runtime_error("Quantifier without preceding element"); - } - sequence->back() += *it; - auto is_star = *it == '*'; - ++it; - if (is_star) { - if (it != end && *it == '?') { - ++it; - } - } - } else if (*it == '{') { - if (sequence->empty()) { - throw std::runtime_error("Repetition without preceding element"); - } - ++it; - auto start = it; - while (it != end && *it != '}') { - ++it; - } - if (it == end) { - throw std::runtime_error("Unmatched '{' in pattern"); - } - auto parts = string_split(std::string(start, it), ","); - ++it; - if (parts.size() > 2) { - throw std::runtime_error("Invalid repetition range in pattern"); - } - - auto parseOptInt = [&](const std::string & s, const std::optional & def = std::nullopt) -> std::optional { - if (s.empty()) { - return def; - } - return std::stoi(s); - }; - auto min = parseOptInt(parts[0], 0); - auto max = parts.size() == 1 ? min : parseOptInt(parts[1]); - if (min && max && *max < *min) { - throw std::runtime_error("Invalid repetition range in pattern"); - } - // Brutal but... let's repeat at least min times, then ? for the delta between min & max (or * for unbounded) - auto part = sequence->back(); - sequence->pop_back(); - for (int i = 0; i < *min; i++) { - sequence->push_back(part); - } - if (max) { - for (int i = *min; i < *max; i++) { - sequence->push_back(part + "?"); - } - } else { - sequence->push_back(part + "*"); - } - } else if (*it == '(') { - ++it; - if (it != end && *it == '?' && (it + 1 != end) && *(it + 1) == ':') { - it += 2; - } - auto sub = process(); - if (*it != ')') { - throw std::runtime_error("Unmatched '(' in pattern"); - } - ++it; - auto & part = sequence->emplace_back("(?:"); - part += sub; - part += ")"; - } else if (*it == ')') { - break; - } else if (*it == '|') { - ++it; - alternatives.emplace_back(); - sequence = &alternatives.back(); - } else if (*it == '\\' && (++it != end)) { - auto str = std::string("\\") + *it; - sequence->push_back(str); - ++it; - } else if (it != end) { - sequence->push_back(std::string(1, *it)); - ++it; - } - } - - // /abcd/ -> ^(dcba|cba|ba|a) -> ^((?:(?:(?:d)?c)?b)?a) - // if n(=4) parts, opening n-1(=3) non-capturing groups after the 1 capturing group - // We'll do the outermost capturing group and final .* in the enclosing function. - std::vector res_alts; - for (const auto & parts : alternatives) { - auto & res = res_alts.emplace_back(); - for (size_t i = 0; i < parts.size() - 1; i++) { - res += "(?:"; - } - for (auto it = parts.rbegin(); it != parts.rend(); ++it) { - res += *it; - if (it != parts.rend() - 1) { - res += ")?"; - } - } - } - return string_join(res_alts, "|"); - }; - auto res = process(); - if (it != end) { - throw std::runtime_error("Unmatched '(' in pattern"); - } - - return "^(" + res + ")"; -} diff --git a/common/regex-partial.h b/common/regex-partial.h deleted file mode 100644 index 634cb4022bd1..000000000000 --- a/common/regex-partial.h +++ /dev/null @@ -1,56 +0,0 @@ -#pragma once - -#include -#include - -enum common_regex_match_type { - COMMON_REGEX_MATCH_TYPE_NONE, - COMMON_REGEX_MATCH_TYPE_PARTIAL, - COMMON_REGEX_MATCH_TYPE_FULL, -}; - -struct common_string_range { - size_t begin; - size_t end; - common_string_range(size_t begin, size_t end) : begin(begin), end(end) { - if (begin > end) { - throw std::runtime_error("Invalid range"); - } - } - // prevent default ctor - common_string_range() = delete; - bool empty() const { - return begin == end; - } - bool operator==(const common_string_range & other) const { - return begin == other.begin && end == other.end; - } -}; - -struct common_regex_match { - common_regex_match_type type = COMMON_REGEX_MATCH_TYPE_NONE; - std::vector groups; - - bool operator==(const common_regex_match & other) const { - return type == other.type && groups == other.groups; - } - bool operator!=(const common_regex_match & other) const { - return !(*this == other); - } -}; - -class common_regex { - std::string pattern; - std::regex rx; - std::regex rx_reversed_partial; - - public: - explicit common_regex(const std::string & pattern); - - common_regex_match search(const std::string & input, size_t pos, bool as_match = false) const; - - const std::string & str() const { return pattern; } -}; - -// For testing only (pretty print of failures). -std::string regex_to_reversed_partial_regex(const std::string & pattern); diff --git a/common/sampling.cpp b/common/sampling.cpp deleted file mode 100644 index d4a2fdcdacc7..000000000000 --- a/common/sampling.cpp +++ /dev/null @@ -1,843 +0,0 @@ -#include "sampling.h" - -#include "common.h" -#include "fit.h" -#include "log.h" -#include "reasoning-budget.h" - -#include "ggml.h" - -#include -#include -#include -#include -#include -#include -#include - -// the ring buffer works similarly to std::deque, but with a fixed capacity -// TODO: deduplicate with llama-impl.h -template -struct ring_buffer { - ring_buffer(size_t cap) : capacity(cap), data(cap) {} - - T & front() { - if (sz == 0) { - throw std::runtime_error("ring buffer is empty"); - } - return data[first]; - } - - const T & front() const { - if (sz == 0) { - throw std::runtime_error("ring buffer is empty"); - } - return data[first]; - } - - T & back() { - if (sz == 0) { - throw std::runtime_error("ring buffer is empty"); - } - return data[pos]; - } - - const T & back() const { - if (sz == 0) { - throw std::runtime_error("ring buffer is empty"); - } - return data[pos]; - } - - void push_back(const T & value) { - if (sz == capacity) { - // advance the start when buffer is full - first = (first + 1) % capacity; - } else { - sz++; - } - data[pos] = value; - pos = (pos + 1) % capacity; - } - - T pop_front() { - if (sz == 0) { - throw std::runtime_error("ring buffer is empty"); - } - T value = data[first]; - first = (first + 1) % capacity; - sz--; - return value; - } - - const T & rat(size_t i) const { - if (i >= sz) { - throw std::runtime_error("ring buffer: index out of bounds"); - } - return data[(first + sz - i - 1) % capacity]; - } - - std::vector to_vector() const { - std::vector result; - result.reserve(sz); - for (size_t i = 0; i < sz; i++) { - result.push_back(data[(first + i) % capacity]); - } - return result; - } - - void clear() { - // here only reset the status of the buffer - sz = 0; - first = 0; - pos = 0; - } - - bool empty() const { - return sz == 0; - } - - size_t size() const { - return sz; - } - - size_t capacity = 0; - size_t sz = 0; - size_t first = 0; - size_t pos = 0; - std::vector data; -}; - -struct common_sampler { - common_params_sampling params; - - struct llama_sampler * grmr; - struct llama_sampler * rbudget; - struct llama_sampler * chain; - - ring_buffer prev; - - std::vector cur; - - llama_token_data_array cur_p; - - void reset() { - prev.clear(); - - llama_sampler_reset(chain); - } - - void set_logits(struct llama_context * ctx, int idx) { - const float * sampled_probs = llama_get_sampled_probs_ith (ctx, idx); - const float * sampled_logits = llama_get_sampled_logits_ith (ctx, idx); - const llama_token * sampled_ids = llama_get_sampled_candidates_ith(ctx, idx); - - const llama_model * model = llama_get_model(ctx); - const llama_vocab * vocab = llama_model_get_vocab(model); - - const int n_vocab = llama_vocab_n_tokens(vocab); - - if (sampled_probs) { - const uint32_t sampled_probs_count = llama_get_sampled_probs_count_ith(ctx, idx); - cur.resize(sampled_probs_count); - for (uint32_t i = 0; i < sampled_probs_count; ++i) { - cur[i] = llama_token_data{sampled_ids[i], sampled_logits[i], sampled_probs[i]}; - } - } else if (sampled_logits) { - const uint32_t sampled_logits_count = llama_get_sampled_logits_count_ith(ctx, idx); - cur.resize(sampled_logits_count); - for (uint32_t i = 0; i < sampled_logits_count; i++) { - cur[i] = llama_token_data{sampled_ids[i], sampled_logits[i], 0.0f}; - } - } else { - const auto * logits = llama_get_logits_ith(ctx, idx); - GGML_ASSERT(logits != nullptr); - cur.resize(n_vocab); - for (llama_token token_id = 0; token_id < n_vocab; token_id++) { - cur[token_id] = llama_token_data{token_id, logits[token_id], 0.0f}; - } - } - - cur_p = { cur.data(), cur.size(), -1, false }; - } - - common_time_meas tm() { - return common_time_meas(t_total_us, params.no_perf); - } - - mutable int64_t t_total_us = 0; -}; - -std::string common_params_sampling::print() const { - char result[1024]; - - snprintf(result, sizeof(result), - "\trepeat_last_n = %d, repeat_penalty = %.3f, frequency_penalty = %.3f, presence_penalty = %.3f\n" - "\tdry_multiplier = %.3f, dry_base = %.3f, dry_allowed_length = %d, dry_penalty_last_n = %d\n" - "\ttop_k = %d, top_p = %.3f, min_p = %.3f, xtc_probability = %.3f, xtc_threshold = %.3f, typical_p = %.3f, top_n_sigma = %.3f, temp = %.3f\n" - "\tmirostat = %d, mirostat_lr = %.3f, mirostat_ent = %.3f, adaptive_target = %.3f, adaptive_decay = %.3f", - penalty_last_n, penalty_repeat, penalty_freq, penalty_present, - dry_multiplier, dry_base, dry_allowed_length, dry_penalty_last_n, - top_k, top_p, min_p, xtc_probability, xtc_threshold, typ_p, top_n_sigma, temp, - mirostat, mirostat_eta, mirostat_tau, adaptive_target, adaptive_decay); - - return std::string(result); -} - -struct common_sampler * common_sampler_init(const struct llama_model * model, struct common_params_sampling & params) { - const llama_vocab * vocab = llama_model_get_vocab(model); - - llama_sampler_chain_params lparams = llama_sampler_chain_default_params(); - - lparams.no_perf = params.no_perf; - - llama_sampler * grmr = nullptr; - llama_sampler * rbudget = nullptr; - llama_sampler * chain = llama_sampler_chain_init(lparams); - - std::vector samplers; - - const std::string & grammar_str = common_grammar_value(params.grammar); - if (grammar_str.compare(0, 11, "%llguidance") == 0) { -#ifdef LLAMA_USE_LLGUIDANCE - grmr = llama_sampler_init_llg(vocab, "lark", grammar_str.c_str()); -#else - GGML_ABORT("llguidance (cmake -DLLAMA_LLGUIDANCE=ON) is not enabled"); -#endif // LLAMA_USE_LLGUIDANCE - } else { - std::vector trigger_patterns; - std::vector trigger_tokens; - for (const auto & trigger : params.grammar_triggers) { - switch (trigger.type) { - case COMMON_GRAMMAR_TRIGGER_TYPE_WORD: - { - const auto & word = trigger.value; - trigger_patterns.push_back(regex_escape(word)); - break; - } - case COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN: - { - trigger_patterns.push_back(trigger.value); - break; - } - case COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN_FULL: - { - const auto & pattern = trigger.value; - std::string anchored = "^$"; - if (!pattern.empty()) { - anchored = (pattern.front() != '^' ? "^" : "") - + pattern - + (pattern.back() != '$' ? "$" : ""); - } - trigger_patterns.push_back(anchored); - break; - } - case COMMON_GRAMMAR_TRIGGER_TYPE_TOKEN: - { - const auto token = trigger.token; - trigger_tokens.push_back(token); - break; - } - default: - GGML_ASSERT(false && "unknown trigger type"); - } - } - - std::vector trigger_patterns_c; - trigger_patterns_c.reserve(trigger_patterns.size()); - for (const auto & regex : trigger_patterns) { - trigger_patterns_c.push_back(regex.c_str()); - } - - if (!grammar_str.empty()) { - if (params.grammar_lazy) { - grmr = llama_sampler_init_grammar_lazy_patterns(vocab, grammar_str.c_str(), "root", - trigger_patterns_c.data(), trigger_patterns_c.size(), - trigger_tokens.data(), trigger_tokens.size()); - } else { - grmr = llama_sampler_init_grammar(vocab, grammar_str.c_str(), "root"); - } - } - } - - // Compute prefill tokens from the generation prompt - std::vector prefill_tokens; - if (!params.generation_prompt.empty()) { - GGML_ASSERT(vocab != nullptr); - auto tokens = common_tokenize(vocab, params.generation_prompt, false, true); - for (size_t i = 0; i < tokens.size(); i++) { - std::string piece = common_token_to_piece(vocab, tokens[i], true); - if (i == 0 && std::isspace(piece[0]) && !std::isspace(params.generation_prompt[0])) { - // Some tokenizers will add a space before the first special token, need to exclude - continue; - } - LOG_DBG("%s: prefill token: %d = %s\n", __func__, tokens[i], piece.c_str()); - prefill_tokens.push_back(tokens[i]); - } - } - - // Feed generation prompt tokens to the grammar sampler so it advances past - // tokens the template already placed in the prompt. - // Only applies to output-format and tool-call grammars; user-supplied grammars must not be prefilled. - if (grmr && !params.grammar_lazy && common_grammar_needs_prefill(params.grammar)) { - try { - for (const auto & token : prefill_tokens) { - llama_sampler_accept(grmr, token); - LOG_DBG("%s: grammar accepted prefill token (%d)\n", __func__, token); - } - } catch (std::exception &e) { - LOG_ERR("%s: error initializing grammar sampler for grammar:\n%s\n\nGeneration prompt:\n'%s'\n", __func__, - common_grammar_value(params.grammar).c_str(), params.generation_prompt.c_str()); - throw e; - } - } - - // reasoning budget sampler (skip when budget is unlimited unless a lazy grammar is active, which needs rbudget for thinking-block suppression) - if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && (params.grammar_lazy || params.reasoning_budget_tokens >= 0)) { - rbudget = common_reasoning_budget_init( - vocab, - params.reasoning_budget_start, - params.reasoning_budget_end, - params.reasoning_budget_forced, - params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens); - - for (const auto & token : prefill_tokens) { - llama_sampler_accept(rbudget, token); - LOG_DBG("%s: reasoning-budget accepted prefill token (%d)\n", __func__, token); - } - } - - if (params.has_logit_bias()) { - samplers.push_back(llama_sampler_init_logit_bias(llama_vocab_n_tokens(vocab), params.logit_bias.size(), params.logit_bias.data())); - } - - if (params.mirostat == 0) { - - bool use_adaptive_p = false; // see below - - for (const auto & cnstr : params.samplers) { - switch (cnstr) { - case COMMON_SAMPLER_TYPE_DRY: - { - std::vector c_breakers; - c_breakers.reserve(params.dry_sequence_breakers.size()); - for (const auto & str : params.dry_sequence_breakers) { - c_breakers.push_back(str.c_str()); - } - samplers.push_back(llama_sampler_init_dry(vocab, llama_model_n_ctx_train(model), params.dry_multiplier, params.dry_base, params.dry_allowed_length, params.dry_penalty_last_n, c_breakers.data(), c_breakers.size())); - } - break; - case COMMON_SAMPLER_TYPE_TOP_K: - samplers.push_back(llama_sampler_init_top_k(params.top_k)); - break; - case COMMON_SAMPLER_TYPE_TOP_P: - samplers.push_back(llama_sampler_init_top_p(params.top_p, params.min_keep)); - break; - case COMMON_SAMPLER_TYPE_TOP_N_SIGMA: - samplers.push_back(llama_sampler_init_top_n_sigma(params.top_n_sigma)); - break; - case COMMON_SAMPLER_TYPE_MIN_P: - samplers.push_back(llama_sampler_init_min_p(params.min_p, params.min_keep)); - break; - case COMMON_SAMPLER_TYPE_XTC: - samplers.push_back(llama_sampler_init_xtc(params.xtc_probability, params.xtc_threshold, params.min_keep, params.seed)); - break; - case COMMON_SAMPLER_TYPE_TYPICAL_P: - samplers.push_back(llama_sampler_init_typical(params.typ_p, params.min_keep)); - break; - case COMMON_SAMPLER_TYPE_TEMPERATURE: - samplers.push_back(llama_sampler_init_temp_ext(params.temp, params.dynatemp_range, params.dynatemp_exponent)); - break; - case COMMON_SAMPLER_TYPE_INFILL: - samplers.push_back(llama_sampler_init_infill(vocab)); - break; - case COMMON_SAMPLER_TYPE_PENALTIES: - samplers.push_back(llama_sampler_init_penalties(params.penalty_last_n, params.penalty_repeat, params.penalty_freq, params.penalty_present)); - break; - case COMMON_SAMPLER_TYPE_ADAPTIVE_P: - // the `adaptive-p` sampler is like `dist` and `mirostat` in that it selects - // a single token, so we will add `dist` at the end of the chain by default, - // unless the user specifically included `adaptive-p`. we set this flag here - // so we know to add the sampler at the very end. - use_adaptive_p = true; - break; - default: - GGML_ASSERT(false && "unknown sampler type"); - } - } - if (use_adaptive_p) { - // only if user explicitly included adaptive-p sampler - samplers.push_back(llama_sampler_init_adaptive_p(params.adaptive_target, params.adaptive_decay, params.seed)); - } else { - // default: sample from distribution - samplers.push_back(llama_sampler_init_dist(params.seed)); - } - } else if (params.mirostat == 1) { - samplers.push_back(llama_sampler_init_temp(params.temp)); - samplers.push_back(llama_sampler_init_mirostat(llama_vocab_n_tokens(vocab), params.seed, params.mirostat_tau, params.mirostat_eta, 100)); - } else if (params.mirostat == 2) { - samplers.push_back(llama_sampler_init_temp(params.temp)); - samplers.push_back(llama_sampler_init_mirostat_v2(params.seed, params.mirostat_tau, params.mirostat_eta)); - } else { - GGML_ASSERT(false && "unknown mirostat version"); - } - - for (auto * smpl : samplers) { - llama_sampler_chain_add(chain, smpl); - } - - if (grmr && params.backend_sampling) { - LOG_WRN("%s: backend sampling is not compatible with grammar, disabling\n", __func__); - - params.backend_sampling = false; - } - - if (rbudget && params.backend_sampling) { - LOG_WRN("%s: backend sampling is not compatible with reasoning budget, disabling\n", __func__); - - params.backend_sampling = false; - } - - auto * result = new common_sampler { - /* .params = */ params, - /* .grmr = */ grmr, - /* .rbudget = */ rbudget, - /* .chain = */ chain, - /* .prev = */ ring_buffer(std::max(32, params.n_prev)), - /* .cur = */ {}, - /* .cur_p = */ {}, - }; - - return result; -} - -void common_sampler_free(struct common_sampler * gsmpl) { - if (!gsmpl) { - return; - } - - llama_sampler_free(gsmpl->grmr); - llama_sampler_free(gsmpl->rbudget); - llama_sampler_free(gsmpl->chain); - - delete gsmpl; -} - -static bool grammar_should_apply(struct common_sampler * gsmpl) { - if (!gsmpl->grmr) { - return false; - } - if (!gsmpl->rbudget) { - return true; - } - if (gsmpl->params.grammar_lazy) { - // if grammar is lazy, only apply when reasoning budget is not active - const auto state = common_reasoning_budget_get_state(gsmpl->rbudget); - return state == REASONING_BUDGET_IDLE || state == REASONING_BUDGET_DONE; - } - return true; -} - -void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, bool is_generated) { - if (!gsmpl) { - return; - } - - const auto tm = gsmpl->tm(); - - // grammar_should_apply() checks the reasoning budget state, so calculate this before we accept - const auto accept_grammar = is_generated && grammar_should_apply(gsmpl); - - if (gsmpl->rbudget && is_generated) { - llama_sampler_accept(gsmpl->rbudget, token); - } - - if (gsmpl->grmr && accept_grammar) { - llama_sampler_accept(gsmpl->grmr, token); - } - - llama_sampler_accept(gsmpl->chain, token); - - gsmpl->prev.push_back(token); -} - -void common_sampler_reset(struct common_sampler * gsmpl) { - if (!gsmpl) { - return; - } - - gsmpl->reset(); -} - -struct common_sampler * common_sampler_clone(common_sampler * gsmpl) { - return new common_sampler { - /* .params = */ gsmpl->params, - /* .grmr = */ llama_sampler_clone(gsmpl->grmr), - /* .rbudget = */ llama_sampler_clone(gsmpl->rbudget), - /* .chain = */ llama_sampler_clone(gsmpl->chain), - /* .prev = */ gsmpl->prev, - /* .cur = */ gsmpl->cur, - /* .cur_p = */ gsmpl->cur_p, - }; -} - -void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl) { - // TODO: measure grammar performance - - const double t_sampling_ms = gsmpl ? 1e-3*gsmpl->t_total_us : 0; - - llama_perf_sampler_data data_smpl; - llama_perf_context_data data_ctx; - - memset(&data_smpl, 0, sizeof(data_smpl)); - memset(&data_ctx, 0, sizeof(data_ctx)); - - if (gsmpl) { - auto & data = data_smpl; - - data = llama_perf_sampler(gsmpl->chain); - - // note: the sampling time includes the samplers time + extra time spent in common/sampling - LOG_INF("%s: sampling time = %10.2f ms\n", __func__, t_sampling_ms); - LOG_INF("%s: samplers time = %10.2f ms / %5d tokens\n", __func__, data.t_sample_ms, data.n_sample); - } - - if (ctx) { - auto & data = data_ctx; - - data = llama_perf_context(ctx); - - const double t_end_ms = 1e-3 * ggml_time_us(); - - const double t_total_ms = t_end_ms - data.t_start_ms; - const double t_unacc_ms = t_total_ms - (t_sampling_ms + data.t_p_eval_ms + data.t_eval_ms); - const double t_unacc_pc = 100.0 * t_unacc_ms / t_total_ms; - - LOG_INF("%s: load time = %10.2f ms\n", __func__, data.t_load_ms); - LOG_INF("%s: prompt eval time = %10.2f ms / %5d tokens (%8.2f ms per token, %8.2f tokens per second)\n", - __func__, data.t_p_eval_ms, data.n_p_eval, data.t_p_eval_ms / data.n_p_eval, 1e3 / data.t_p_eval_ms * data.n_p_eval); - LOG_INF("%s: eval time = %10.2f ms / %5d runs (%8.2f ms per token, %8.2f tokens per second)\n", - __func__, data.t_eval_ms, data.n_eval, data.t_eval_ms / data.n_eval, 1e3 / data.t_eval_ms * data.n_eval); - LOG_INF("%s: total time = %10.2f ms / %5d tokens\n", __func__, (t_end_ms - data.t_start_ms), (data.n_p_eval + data.n_eval)); - LOG_INF("%s: unaccounted time = %10.2f ms / %5.1f %% (total - sampling - prompt eval - eval) / (total)\n", __func__, t_unacc_ms, t_unacc_pc); - LOG_INF("%s: graphs reused = %10d\n", __func__, data.n_reused); - - common_memory_breakdown_print(ctx); - } -} - -struct llama_sampler * common_sampler_get(const struct common_sampler * gsmpl) { - if (!gsmpl) { - return nullptr; - } - - return gsmpl->chain; -} - -llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_context * ctx, int idx, bool grammar_first) { - llama_synchronize(ctx); - - // start measuring sampling time after the llama_context synchronization in order to not measure any ongoing async operations - const auto tm = gsmpl->tm(); - - llama_token id = LLAMA_TOKEN_NULL; - - auto & grmr = gsmpl->grmr; - auto & rbudget = gsmpl->rbudget; - auto & chain = gsmpl->chain; - auto & cur_p = gsmpl->cur_p; // initialized by set_logits - - // Check if a backend sampler has already sampled a token in which case we - // return that token id directly. - { - id = llama_get_sampled_token_ith(ctx, idx); - - if (id != LLAMA_TOKEN_NULL) { - LOG_DBG("%s: Backend sampler selected token: '%d'. Will not run any CPU samplers\n", __func__, id); - - GGML_ASSERT(!gsmpl->grmr && "using grammar in combination with backend sampling is not supported"); - GGML_ASSERT(!gsmpl->rbudget && "using reasoning budget in combination with backend sampling is not supported"); - - // TODO: simplify - gsmpl->cur.resize(1); - gsmpl->cur[0] = { id, 0.0f, 1.0f }; - cur_p = { gsmpl->cur.data(), gsmpl->cur.size(), 0, true }; - - return id; - } - } - - gsmpl->set_logits(ctx, idx); - - // apply reasoning budget first - llama_sampler_apply(rbudget, &cur_p); - - if (grammar_first && grammar_should_apply(gsmpl)) { - llama_sampler_apply(grmr, &cur_p); - } - - llama_sampler_apply(chain, &cur_p); - - id = cur_p.data[cur_p.selected].id; - - if (grammar_first || !grammar_should_apply(gsmpl)) { - return id; - } - - // check if it the sampled token fits the grammar (grammar-based rejection sampling) - { - llama_token_data single_token_data = { id, 1.0f, 0.0f }; - llama_token_data_array single_token_data_array = { &single_token_data, 1, -1, false }; - - llama_sampler_apply(grmr, &single_token_data_array); - - const bool is_valid = single_token_data_array.data[0].logit != -INFINITY; - if (is_valid) { - return id; - } - } - - // resampling: - // if the token is not valid, sample again, but first apply the grammar sampler and then the sampling chain - gsmpl->set_logits(ctx, idx); - - llama_sampler_apply(rbudget, &cur_p); - - if (grammar_should_apply(gsmpl)) { - llama_sampler_apply(grmr, &cur_p); - } - - llama_sampler_apply(chain, &cur_p); - - GGML_ASSERT(cur_p.selected != -1 && "no selected token during sampling - check your sampling configuration"); - - id = cur_p.data[cur_p.selected].id; - - return id; -} - -std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector & idxs, const llama_tokens & draft, bool grammar_first) { - GGML_ASSERT(idxs.size() == draft.size() + 1 && "idxs.size() must be draft.size() + 1"); - - std::vector result; - result.reserve(idxs.size()); - - size_t i = 0; - for (; i < draft.size(); i++) { - const llama_token id = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first); - - common_sampler_accept(gsmpl, id, true); - - result.push_back(id); - - if (draft[i] != id) { - break; - } - } - - if (i == draft.size()) { - const llama_token id = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first); - - common_sampler_accept(gsmpl, id, true); - - result.push_back(id); - } - - return result; -} - -std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first) { - std::vector idxs(draft.size() + 1); - for (size_t i = 0; i < idxs.size(); ++i) { - idxs[i] = i; - } - - return common_sampler_sample_and_accept_n(gsmpl, ctx, idxs, draft, grammar_first); -} - -uint32_t common_sampler_get_seed(const struct common_sampler * gsmpl) { - return llama_sampler_get_seed(gsmpl->chain); -} - -// helpers - -llama_token_data_array * common_sampler_get_candidates(struct common_sampler * gsmpl, bool do_sort) { - const auto tm = gsmpl->tm(); - - auto * res = &gsmpl->cur_p; - - if (do_sort && !res->sorted) { - // remember the selected token before sorting - const llama_token id = res->data[res->selected].id; - - std::sort(res->data, res->data + res->size, [](const llama_token_data & a, const llama_token_data & b) { - return a.p > b.p; - }); - - // restore the selected token after sorting - for (size_t i = 0; i < res->size; ++i) { - if (res->data[i].id == id) { - res->selected = i; - break; - } - } - - res->sorted = true; - } - - return res; -} - -llama_token common_sampler_last(const struct common_sampler * gsmpl) { - return gsmpl->prev.rat(0); -} - -std::string common_sampler_print(const struct common_sampler * gsmpl) { - std::string result = "logits "; - - for (int i = 0; i < llama_sampler_chain_n(gsmpl->chain); i++) { - const auto * smpl = llama_sampler_chain_get(gsmpl->chain, i); - result += std::string("-> "); - result += std::string(llama_sampler_name(smpl)) + " "; - } - - return result; -} - -std::string common_sampler_prev_str(common_sampler * gsmpl, llama_context * ctx_main, int n) { - n = std::min(n, (int) gsmpl->prev.size()); - - if (n <= 0) { - return ""; - } - - std::string result; - result.reserve(8*n); // 8 is the average length of a token [citation needed], TODO: compute this from the vocab - - for (int i = n - 1; i >= 0; i--) { - const llama_token id = gsmpl->prev.rat(i); - - GGML_ASSERT(id != LLAMA_TOKEN_NULL && "null token in the sampling history - should not happen"); - - result += common_token_to_piece(ctx_main, id); - } - - return result; -} - -char common_sampler_type_to_chr(enum common_sampler_type cnstr) { - switch (cnstr) { - case COMMON_SAMPLER_TYPE_DRY: return 'd'; - case COMMON_SAMPLER_TYPE_TOP_K: return 'k'; - case COMMON_SAMPLER_TYPE_TYPICAL_P: return 'y'; - case COMMON_SAMPLER_TYPE_TOP_P: return 'p'; - case COMMON_SAMPLER_TYPE_TOP_N_SIGMA: return 's'; - case COMMON_SAMPLER_TYPE_MIN_P: return 'm'; - case COMMON_SAMPLER_TYPE_TEMPERATURE: return 't'; - case COMMON_SAMPLER_TYPE_XTC: return 'x'; - case COMMON_SAMPLER_TYPE_INFILL: return 'i'; - case COMMON_SAMPLER_TYPE_PENALTIES: return 'e'; - case COMMON_SAMPLER_TYPE_ADAPTIVE_P: return 'a'; - default : return '?'; - } -} - -std::string common_sampler_type_to_str(enum common_sampler_type cnstr) { - switch (cnstr) { - case COMMON_SAMPLER_TYPE_DRY: return "dry"; - case COMMON_SAMPLER_TYPE_TOP_K: return "top_k"; - case COMMON_SAMPLER_TYPE_TYPICAL_P: return "typ_p"; - case COMMON_SAMPLER_TYPE_TOP_P: return "top_p"; - case COMMON_SAMPLER_TYPE_TOP_N_SIGMA: return "top_n_sigma"; - case COMMON_SAMPLER_TYPE_MIN_P: return "min_p"; - case COMMON_SAMPLER_TYPE_TEMPERATURE: return "temperature"; - case COMMON_SAMPLER_TYPE_XTC: return "xtc"; - case COMMON_SAMPLER_TYPE_INFILL: return "infill"; - case COMMON_SAMPLER_TYPE_PENALTIES: return "penalties"; - case COMMON_SAMPLER_TYPE_ADAPTIVE_P: return "adaptive_p"; - default : return ""; - } -} - -std::vector common_sampler_types_from_names(const std::vector & names, bool allow_alt_names) { - std::unordered_map sampler_canonical_name_map { - { "dry", COMMON_SAMPLER_TYPE_DRY }, - { "top_k", COMMON_SAMPLER_TYPE_TOP_K }, - { "top_p", COMMON_SAMPLER_TYPE_TOP_P }, - { "top_n_sigma", COMMON_SAMPLER_TYPE_TOP_N_SIGMA }, - { "typ_p", COMMON_SAMPLER_TYPE_TYPICAL_P }, - { "min_p", COMMON_SAMPLER_TYPE_MIN_P }, - { "temperature", COMMON_SAMPLER_TYPE_TEMPERATURE }, - { "xtc", COMMON_SAMPLER_TYPE_XTC }, - { "infill", COMMON_SAMPLER_TYPE_INFILL }, - { "penalties", COMMON_SAMPLER_TYPE_PENALTIES }, - { "adaptive_p", COMMON_SAMPLER_TYPE_ADAPTIVE_P }, - }; - - // since samplers names are written multiple ways - // make it ready for both system names and input names - std::unordered_map sampler_alt_name_map { - { "top-k", COMMON_SAMPLER_TYPE_TOP_K }, - { "top-p", COMMON_SAMPLER_TYPE_TOP_P }, - { "top-n-sigma", COMMON_SAMPLER_TYPE_TOP_N_SIGMA }, - { "nucleus", COMMON_SAMPLER_TYPE_TOP_P }, - { "typical-p", COMMON_SAMPLER_TYPE_TYPICAL_P }, - { "typical", COMMON_SAMPLER_TYPE_TYPICAL_P }, - { "typ-p", COMMON_SAMPLER_TYPE_TYPICAL_P }, - { "typ", COMMON_SAMPLER_TYPE_TYPICAL_P }, - { "min-p", COMMON_SAMPLER_TYPE_MIN_P }, - { "temp", COMMON_SAMPLER_TYPE_TEMPERATURE }, - { "adaptive-p", COMMON_SAMPLER_TYPE_ADAPTIVE_P }, - }; - - std::vector samplers; - samplers.reserve(names.size()); - - for (const auto & name : names) { - auto sampler = sampler_canonical_name_map.find(name); - if (sampler != sampler_canonical_name_map.end()) { - samplers.push_back(sampler->second); - continue; - } - if (allow_alt_names) { - sampler = sampler_alt_name_map.find(name); - if (sampler != sampler_alt_name_map.end()) { - samplers.push_back(sampler->second); - continue; - } - } - LOG_WRN("%s: unable to match sampler by name '%s'\n", __func__, name.c_str()); - } - - return samplers; -} - -std::vector common_sampler_types_from_chars(const std::string & chars) { - std::unordered_map sampler_name_map = { - { common_sampler_type_to_chr(COMMON_SAMPLER_TYPE_DRY), COMMON_SAMPLER_TYPE_DRY }, - { common_sampler_type_to_chr(COMMON_SAMPLER_TYPE_TOP_K), COMMON_SAMPLER_TYPE_TOP_K }, - { common_sampler_type_to_chr(COMMON_SAMPLER_TYPE_TYPICAL_P), COMMON_SAMPLER_TYPE_TYPICAL_P }, - { common_sampler_type_to_chr(COMMON_SAMPLER_TYPE_TOP_P), COMMON_SAMPLER_TYPE_TOP_P }, - { common_sampler_type_to_chr(COMMON_SAMPLER_TYPE_TOP_N_SIGMA), COMMON_SAMPLER_TYPE_TOP_N_SIGMA }, - { common_sampler_type_to_chr(COMMON_SAMPLER_TYPE_MIN_P), COMMON_SAMPLER_TYPE_MIN_P }, - { common_sampler_type_to_chr(COMMON_SAMPLER_TYPE_TEMPERATURE), COMMON_SAMPLER_TYPE_TEMPERATURE }, - { common_sampler_type_to_chr(COMMON_SAMPLER_TYPE_XTC), COMMON_SAMPLER_TYPE_XTC }, - { common_sampler_type_to_chr(COMMON_SAMPLER_TYPE_INFILL), COMMON_SAMPLER_TYPE_INFILL }, - { common_sampler_type_to_chr(COMMON_SAMPLER_TYPE_PENALTIES), COMMON_SAMPLER_TYPE_PENALTIES }, - { common_sampler_type_to_chr(COMMON_SAMPLER_TYPE_ADAPTIVE_P), COMMON_SAMPLER_TYPE_ADAPTIVE_P }, - }; - - std::vector samplers; - samplers.reserve(chars.size()); - - for (const auto & c : chars) { - const auto sampler = sampler_name_map.find(c); - if (sampler != sampler_name_map.end()) { - samplers.push_back(sampler->second); - } else { - LOG_WRN("%s: unable to match sampler by char '%c'\n", __func__, c); - } - } - - return samplers; -} diff --git a/common/sampling.h b/common/sampling.h deleted file mode 100644 index 49506a00cd8b..000000000000 --- a/common/sampling.h +++ /dev/null @@ -1,119 +0,0 @@ -#pragma once - -#include "llama.h" - -#include "common.h" - -#include -#include - -// common_sampler extends llama_sampler with additional functionality: -// -// - grammar support -// - custom sampler logic based on the parameters -// - history of the last accepted tokens -// - performance metrics -// -// This goal is to have a common implementation of the sampling logic shared across the examples. -// For example, depending on the temperature, the sampling chain can be very simple (greedy) or more -// complex (top-k, top-p, etc). -// -// Another example is related to the grammar. In general, the grammar constraints applied on the full -// vocabulary can be very taxing. To improve performance, the grammar can be applied only to the sampled -// token in order to verify if it fits the grammar. And only if the token doesn't fit the grammar, the -// grammar constraints are applied to the full vocabulary and the token is resampled. -// -// The common_sampler also maintains a container with the last accepted tokens. In the future, this can -// be moved into the core llama library. -// -// For convenience, the common_sampler also maintains a container with the current candidate tokens. -// This can be used to access the probabilities of the rest of the non-sampled tokens. -// -// TODO: measure grammar performance -// - -struct common_sampler; - -// llama_sampler API overloads - -// note: can mutate params in some cases -struct common_sampler * common_sampler_init(const struct llama_model * model, struct common_params_sampling & params); - -void common_sampler_free(struct common_sampler * gsmpl); - -// if is_generated is true, the token is accepted by the sampling chain, the reasoning budget sampler, and the grammar sampler -void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, bool is_generated); -void common_sampler_reset (struct common_sampler * gsmpl); -struct common_sampler * common_sampler_clone (struct common_sampler * gsmpl); - -// arguments can be nullptr to skip printing -void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl); - -// get the underlying llama_sampler_chain -struct llama_sampler * common_sampler_get(const struct common_sampler * gsmpl); - -// extended sampling implementation: -// -// - set logits -// - apply the configured sampler chain -// - check if the token fits the grammar (if any) -// - if not: resample by first applying the grammar constraints and then sampling again (slower path) -// -// if grammar_first is true, the grammar is applied before the samplers (slower) -// useful in cases where all the resulting candidates (not just the sampled one) must fit the grammar -// -llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_context * ctx, int idx, bool grammar_first = false); - -// generalized version of common_sampler_sample -// -// will cross-reference the sampled tokens with a batch of draft tokens and accept those that match -// if the sampler disagrees at some point, we stop and return the accepted tokens up to now -// -// common_sampler_sample_n(gsmpl, ctx, { idx }, {}); -// -// is equivalent to -// -// common_sampler_sample(gsmpl, ctx, idx); -// common_sampler_accept(gsmpl, token, true); -// -// requires: idxs.size() == draft.size() + 1 -// -// returns at least 1 token, up to idxs.size() -// -std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector & idxs, const llama_tokens & draft, bool grammar_first = false); - -// assume idxs == [ 0, 1, 2, ..., draft.size() ] -std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first = false); - -uint32_t common_sampler_get_seed(const struct common_sampler * gsmpl); - -// helpers - -// access the internal list of current candidate tokens -// if do_sort == true, the candidates are guaranteed to be sorted afterwards (in descending order of probability) -// the .sorted flag of the result indicates whether the returned candidates are sorted -llama_token_data_array * common_sampler_get_candidates(struct common_sampler * gsmpl, bool do_sort); - -// get the last accepted token -llama_token common_sampler_last(const struct common_sampler * gsmpl); - -// print the sampler chain into a string -std::string common_sampler_print(const struct common_sampler * gsmpl); - -// get a string representation of the last accepted tokens -std::string common_sampler_prev_str(common_sampler * gsmpl, llama_context * ctx, int n); - -char common_sampler_type_to_chr(enum common_sampler_type cnstr); -std::string common_sampler_type_to_str(enum common_sampler_type cnstr); - -std::vector common_sampler_types_from_names(const std::vector & names, bool allow_alt_names); -std::vector common_sampler_types_from_chars(const std::string & chars); - -llama_sampler * llama_sampler_init_llg(const llama_vocab * vocab, - const char * grammar_kind, const char * grammar_data); - -struct common_sampler_deleter { - void operator()(common_sampler * s) { common_sampler_free(s); } -}; - -typedef std::unique_ptr common_sampler_ptr; diff --git a/common/speculative.cpp b/common/speculative.cpp deleted file mode 100644 index bbf88fa6e715..000000000000 --- a/common/speculative.cpp +++ /dev/null @@ -1,1248 +0,0 @@ -#include "speculative.h" - -#include "common.h" -#include "ggml.h" -#include "llama.h" -#include "log.h" -#include "ngram-cache.h" -#include "ngram-map.h" -#include "ngram-mod.h" -#include "sampling.h" - -#include -#include -#include -#include -#include - -#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 -#define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 - -const std::vector common_speculative_types = { - COMMON_SPECULATIVE_TYPE_NONE, - COMMON_SPECULATIVE_TYPE_DRAFT, - COMMON_SPECULATIVE_TYPE_EAGLE3, - COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, - COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, - COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, - COMMON_SPECULATIVE_TYPE_NGRAM_MOD, - COMMON_SPECULATIVE_TYPE_NGRAM_CACHE -}; - -const std::map common_speculative_type_from_name_map = { - {"none", COMMON_SPECULATIVE_TYPE_NONE}, - {"draft", COMMON_SPECULATIVE_TYPE_DRAFT}, - {"eagle3", COMMON_SPECULATIVE_TYPE_EAGLE3}, - {"ngram_simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE}, - {"ngram_map_k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K}, - {"ngram_map_k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V}, - {"ngram_mod", COMMON_SPECULATIVE_TYPE_NGRAM_MOD}, - {"ngram_cache", COMMON_SPECULATIVE_TYPE_NGRAM_CACHE} -}; - -struct common_speculative_config { - common_speculative_type type; - common_params_speculative params; - - common_speculative_config(common_speculative_type t, - const common_params_speculative & p = common_params_speculative{}) : type(t), params(p) {} -}; - -static bool common_speculative_are_compatible( - const llama_model * model_tgt, - const llama_model * model_dft) { - const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt); - const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft); - - const bool vocab_type_tgt = llama_vocab_type(vocab_tgt); - LOG_DBG("%s: vocab_type tgt: %d\n", __func__, vocab_type_tgt); - - const bool vocab_type_dft = llama_vocab_type(vocab_dft); - LOG_DBG("%s: vocab_type dft: %d\n", __func__, vocab_type_dft); - - if (vocab_type_tgt != vocab_type_dft) { - LOG_WRN("%s: draft model vocab type must match target model to use speculation but " - "vocab_type_dft = %d while vocab_type_tgt = %d\n", __func__, vocab_type_dft, vocab_type_tgt); - return false; - } - - if (llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) || - (llama_vocab_get_add_bos(vocab_tgt) && llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft))) { - LOG_WRN("%s: draft model bos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n", - __func__, - llama_vocab_get_add_bos(vocab_tgt), llama_vocab_get_add_bos(vocab_dft), - llama_vocab_bos(vocab_tgt), llama_vocab_bos(vocab_dft)); - return false; - } - - if (llama_vocab_get_add_eos(vocab_tgt) != llama_vocab_get_add_eos(vocab_dft) || - (llama_vocab_get_add_eos(vocab_tgt) && llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft))) { - LOG_WRN("%s: draft model eos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n", - __func__, - llama_vocab_get_add_eos(vocab_tgt), llama_vocab_get_add_eos(vocab_dft), - llama_vocab_eos(vocab_tgt), llama_vocab_eos(vocab_dft)); - return false; - } - - { - const int n_vocab_tgt = llama_vocab_n_tokens(vocab_tgt); - const int n_vocab_dft = llama_vocab_n_tokens(vocab_dft); - const int vocab_diff = n_vocab_tgt > n_vocab_dft - ? n_vocab_tgt - n_vocab_dft - : n_vocab_dft - n_vocab_tgt; - - if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) { - LOG_DBG("%s: draft model vocab must closely match target model to use speculation but ", __func__); - LOG_DBG("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n", - n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE); - return false; - } - - for (int i = SPEC_VOCAB_CHECK_START_TOKEN_ID; i < std::min(n_vocab_tgt, n_vocab_dft); ++i) { - const char * token_text_tgt = llama_vocab_get_text(vocab_tgt, i); - const char * token_text_dft = llama_vocab_get_text(vocab_dft, i); - - if (std::strcmp(token_text_tgt, token_text_dft) != 0) { - LOG_DBG("%s: draft model vocab must match target model to use speculation but ", __func__); - LOG_DBG("token %d content differs - target '%s', draft '%s'\n", i, - common_token_to_piece(vocab_tgt, i).c_str(), - common_token_to_piece(vocab_dft, i).c_str()); - return false; - } - } - } - - return true; -} - -// state of an implementation of speculative decoding -// -// each implementation has a unique type and a state that is implementation-specific -// in a subclass of common_speculative_state -struct common_speculative_state { - const enum common_speculative_type type; - - size_t n_call_begin = 0; // number of times this implementation was called for refresh. - size_t n_call_draft = 0; // number of times this implementation was called for generation. - size_t n_call_accept = 0; // number of times this implementation was called for accumulation. - - size_t n_gen_drafts = 0; // number of times a draft or part was generated by this implementation. - size_t n_acc_drafts = 0; // number of times a draft or part was accepted by the target model. - size_t n_gen_tokens = 0; // number of tokens generated by this implementation. - size_t n_acc_tokens = 0; // number of tokens accepted by the target model. - - // TODO: track performance of most recent calls - const bool gen_perf = true; // whether to generate performance stats. - - int64_t t_begin_us = 0; // total time spent in refresh of this implementation in microseconds. - int64_t t_draft_us = 0; // total time spent in generating drafts in this implementation in microseconds. - int64_t t_accept_us = 0; // total time spent in accumulation of this implementation in microseconds. - - common_speculative_state(enum common_speculative_type type) : type(type) {} - - virtual ~common_speculative_state() = default; - - virtual void begin(const llama_tokens & prompt) = 0; - - virtual void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) = 0; - - virtual void accept(uint16_t n_accepted) = 0; - - virtual int32_t n_max(const common_params_speculative & params) const = 0; - virtual int32_t n_min(const common_params_speculative & params) const = 0; -}; - -struct common_speculative_checkpoint { - llama_pos pos_min = 0; - llama_pos pos_max = 0; - - int64_t n_tokens = 0; - - std::vector data; - - size_t size() const { - return data.size(); - } -}; - -struct common_speculative_state_draft : public common_speculative_state { - llama_context * ctx_tgt; // only used for retokenizing from ctx_dft - llama_context * ctx_dft; - - bool use_ckpt = false; - common_speculative_checkpoint ckpt; - - common_sampler * smpl; - - llama_batch batch; - llama_tokens prompt_dft; - - bool vocab_cmpt = true; // whether retokenization is needed - std::unordered_map vocab_map; - - common_speculative_state_draft( - enum common_speculative_type type, - llama_context * ctx_tgt, - llama_context * ctx_dft, - const std::vector> & replacements, - bool use_ckpt) - : common_speculative_state(type) - , ctx_tgt(ctx_tgt) - , ctx_dft(ctx_dft) - , use_ckpt(use_ckpt) - { - batch = llama_batch_init(llama_n_batch(ctx_dft), 0, 1); - smpl = nullptr; - - // TODO: optimize or pass from outside? - // { - // common_params_sampling params; - // params.no_perf = false; - // - // params.top_k = 40; - // params.top_p = 0.9; - // - // params.samplers = { - // COMMON_SAMPLER_TYPE_TOP_K, - // COMMON_SAMPLER_TYPE_TOP_P, - // COMMON_SAMPLER_TYPE_INFILL, - // }; - // - // result->smpl = common_sampler_init(llama_get_model(ctx_dft), params); - // } - { - common_params_sampling params; - params.no_perf = false; - params.top_k = 10; - params.samplers = { - COMMON_SAMPLER_TYPE_TOP_K, - }; - - smpl = common_sampler_init(llama_get_model(ctx_dft), params); - } - - vocab_cmpt = common_speculative_are_compatible(llama_get_model(ctx_tgt), llama_get_model(ctx_dft)); - LOG_DBG("vocab_cmpt = %d\n", vocab_cmpt); - - if (!vocab_cmpt) { - LOG_WRN("the target and draft vocabs are not compatible - tokens will be translated between the two\n"); - - for (const auto & pair : replacements) { - vocab_map[pair.first] = pair.second; - } - } - } - - ~common_speculative_state_draft() override { - llama_perf_context_print(ctx_dft); - - llama_free(ctx_dft); - - common_sampler_free(smpl); - - llama_batch_free(batch); - } - - void begin(const llama_tokens & /*prompt*/) override { - } - - size_t create_checkpoint(int n_tokens_prompt) { - int slot_id = 0; - const size_t checkpoint_size = llama_state_seq_get_size_ext(ctx_dft, slot_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - - ckpt.pos_min = llama_memory_seq_pos_min(llama_get_memory(ctx_dft), slot_id); - ckpt.pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), slot_id); - ckpt.n_tokens = n_tokens_prompt; - ckpt.data.resize(checkpoint_size); - - const size_t n = llama_state_seq_get_data_ext(ctx_dft, ckpt.data.data(), checkpoint_size, slot_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - if (n != checkpoint_size) { - GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", checkpoint_size, n); - } - - LOG_DBG("%s: pos_min = %d, pos_max = %d, size = %.3f MiB\n", __func__, - ckpt.pos_min, ckpt.pos_max, (float) ckpt.data.size() / 1024 / 1024); - return n; - } - - size_t restore_checkpoint() { - int slot_id = 0; - LOG_DBG("%s: pos_min = %d, pos_max = %d\n", __func__, ckpt.pos_min, ckpt.pos_max); - const size_t n = llama_state_seq_set_data_ext(ctx_dft, ckpt.data.data(), ckpt.size(), slot_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - if (n != ckpt.size()) { - GGML_ABORT("%s: failed to restore context checkpoint (pos_min=%d, pos_max=%d, size=%zu", - __func__, ckpt.pos_min, ckpt.pos_max, ckpt.size()); - } - llama_memory_seq_rm(llama_get_memory(ctx_dft), slot_id, ckpt.pos_max + 1, -1); - - return n; - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) override { - const auto & sparams = params.draft; - - auto * spec = this; - - auto & batch = spec->batch; - auto & ctx_tgt = spec->ctx_tgt; - auto & ctx_dft = spec->ctx_dft; - auto & smpl = spec->smpl; - auto & prompt_dft = spec->prompt_dft; - - auto * mem_dft = llama_get_memory(ctx_dft); - - int reuse_i = 0; // index of part to be reused in prompt_dft - int reuse_n = 0; // length of part to be reused in prompt_dft - - const int n_ctx = llama_n_ctx(ctx_dft) - sparams.n_max; - - llama_tokens prompt_cnv; - if (!spec->vocab_cmpt) { - std::string text; - - text = common_detokenize(ctx_tgt, prompt_tgt, true); - text = replace_to_dft(text); - - LOG_DBG("%s: main->draft detokenized string: '%s'\n", __func__, text.c_str()); - - prompt_cnv = common_tokenize(ctx_dft, text, false, true); - - // convert id_last to draft vocab. llama_detokenize is called directly to avoid an allocation - const auto * model_tgt = llama_get_model(ctx_tgt); - const auto * vocab_tgt = llama_model_get_vocab(model_tgt); - - int32_t n_chars = llama_detokenize(vocab_tgt, &id_last, 1, nullptr, 0, false, false); - GGML_ASSERT(n_chars < 0 && "failed to detokenize id_last"); - - text.resize(-n_chars); - llama_detokenize(vocab_tgt, &id_last, 1, text.data(), text.size(), false, false); - text = replace_to_dft(text); - - LOG_DBG("main->draft detokenized id_last(%d): '%s'\n", id_last, text.c_str()); - id_last = common_tokenize(ctx_dft, text, false, true)[0]; - } - - const llama_tokens & prompt_cur = spec->vocab_cmpt ? prompt_tgt : prompt_cnv; - - const int i_start = std::max(0, (int) prompt_cur.size() - n_ctx); - - if (use_ckpt && i_start > 0) { - LOG_WRN("%s: context shift is not supported with checkpoint-based contexts - skipping\n", __func__); - return; - } - - // reuse as much as possible from the old draft context - // ideally, the draft context should be as big as the target context and we will always reuse the entire prompt - for (int i = 0; i < (int) prompt_dft.size(); ++i) { - int cur = 0; - while (i_start + cur < (int) prompt_cur.size() && - i + cur < (int) prompt_dft.size() && - prompt_cur[i_start + cur] == prompt_dft[i + cur]) { - cur++; - } - - if ((cur >= 256 || n_ctx >= (int) prompt_cur.size()) && cur > reuse_n) { - reuse_i = i; - reuse_n = cur; - } - - if (use_ckpt) { - break; - } - } - - LOG_DBG("%s: reuse_i = %d, reuse_n = %d, #prompt_dft = %zu, #prompt_cur = %zu\n", - __func__, reuse_i, reuse_n, prompt_dft.size(), prompt_cur.size()); - if (use_ckpt && ckpt.n_tokens > reuse_n) { - LOG_DBG("%s: checkpoint (n_tokens = %d) is outdated -> delete it\n", __func__, (int) ckpt.n_tokens); - - reuse_i = 0; - reuse_n = 0; - - ckpt = {}; - } - - result.clear(); - result.reserve(sparams.n_max); - - if (reuse_n == 0 || (use_ckpt && reuse_i > 0)) { - llama_memory_clear(mem_dft, false); - prompt_dft.clear(); - } else { - // this happens when a previous draft has been discarded (for example, due to being too small), but the - // target model agreed with it. in this case, we simply pass back the previous results to save compute - if (reuse_i + reuse_n < (int64_t) prompt_dft.size() && prompt_dft[reuse_i + reuse_n] == id_last) { - for (int i = reuse_i + reuse_n + 1; i < (int) prompt_dft.size(); ++i) { - result.push_back(prompt_dft[i]); - - if (sparams.n_max <= (int) result.size()) { - break; - } - } - - return; - } - - if (reuse_i > 0) { - GGML_ASSERT(!use_ckpt); - - bool is_removed = llama_memory_seq_rm (mem_dft, 0, 0, reuse_i); - if (!is_removed) { - LOG_ERR("%s: llama_memory_seq_rm failed, reuse_i=%d\n", __func__, reuse_i); - return; - } - llama_memory_seq_add(mem_dft, 0, reuse_i, -1, -reuse_i); - - prompt_dft.erase(prompt_dft.begin(), prompt_dft.begin() + reuse_i); - } - - if (reuse_n < (int) prompt_dft.size()) { - if (use_ckpt) { - if (ckpt.n_tokens > 0) { - LOG_DBG("%s: restoring checkpoint, reuse_n=%d, prompt_dft.size=%zu\n", __func__, reuse_n, prompt_dft.size()); - restore_checkpoint(); - reuse_n = ckpt.n_tokens; - prompt_dft.resize(reuse_n); - } - } else { - const bool is_removed = llama_memory_seq_rm(mem_dft, 0, reuse_n, -1); - if (!is_removed) { - LOG_ERR("%s: llama_memory_seq_rm failed, reuse_n=%d, prompt_dft.size=%zu\n", __func__, reuse_n, prompt_dft.size()); - return; - } - prompt_dft.erase(prompt_dft.begin() + reuse_n, prompt_dft.end()); - } - } - } - - // prepare a batch to evaluate any new tokens in the prompt - common_batch_clear(batch); - - for (size_t i = i_start + reuse_n; i < prompt_cur.size(); ++i) { - //LOG_DBG("i = %d, i_start = %d, reuse_n = %d, i - i_start = %d, id = %6d\n", i, i_start, reuse_n, i - i_start, prompt_cur[i]); - common_batch_add(batch, prompt_cur[i], i - i_start, { 0 }, false); - - prompt_dft.push_back(prompt_cur[i]); - } - - // we should rarely end-up here during normal decoding - if (batch.n_tokens > 0) { - //LOG_DBG("%s: draft prompt batch: %s\n", __func__, string_from(ctx, batch).c_str()); - LOG_DBG("%s: draft prompt batch: %d tokens\n", __func__, batch.n_tokens); - - int ret = llama_decode(ctx_dft, batch); - if (ret != 0 && ret != 1) { - LOG_WRN("%s: llama_decode returned %d, prompt_cur.size=%zu\n", - __func__, ret, prompt_cur.size()); - } - - if (use_ckpt) { - create_checkpoint(prompt_dft.size()); - } - } - - const llama_pos n_past = prompt_dft.size(); - - LOG_DBG("%s: n_past = %d\n", __func__, n_past); - - common_batch_clear(batch); - common_batch_add (batch, id_last, n_past, { 0 }, true); - - prompt_dft.push_back(id_last); - - //LOG_DBG("%s: draft prompt: %s\n", __func__, string_from(ctx_dft, prompt_dft).c_str()); - - int ret = llama_decode(ctx_dft, batch); - if (ret != 0 && ret != 1) { - LOG_WRN("%s: llama_decode returned %d, prompt_cur.size=%zu, prompt_dft.size=%zu\n", - __func__, ret, prompt_cur.size(), prompt_dft.size()); - } - - common_sampler_reset(smpl); - - // sample n_draft tokens from the draft model - for (int i = 0; i < sparams.n_max; ++i) { - common_batch_clear(batch); - - common_sampler_sample(smpl, ctx_dft, 0, true); - - const auto * cur_p = common_sampler_get_candidates(smpl, true); - - for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { - LOG_DBG(" - draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", - k, i, cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); - } - - // add drafted token for each sequence - const llama_token id = cur_p->data[0].id; - - common_sampler_accept(smpl, id, true); - - // only collect very high-confidence draft tokens - if (cur_p->data[0].p < sparams.p_min) { - break; - } - - result.push_back(id); - - if (sparams.n_max <= (int) result.size()) { - break; - } - - common_batch_add(batch, id, n_past + i + 1, { 0 }, true); - - // evaluate the drafted tokens on the draft model - ret = llama_decode(ctx_dft, batch); - if (ret != 0) { - LOG_WRN("%s: llama_decode[%d] returned %d, prompt_cur.size=%zu, prompt_dft.size=%zu\n", - __func__, i, ret, prompt_cur.size(), prompt_dft.size()); - } - - prompt_dft.push_back(id); - } - - if (!spec->vocab_cmpt) { - std::string detokenized = common_detokenize(ctx_dft, result, true); - detokenized = replace_to_tgt(detokenized); - LOG_DBG("draft->main detokenized string: '%s'\n", detokenized.c_str()); - result = common_tokenize(ctx_tgt, detokenized, false, true); - if (result.size() > (size_t) sparams.n_max) { - result.resize(sparams.n_max); - } - } - - if (result.size() < (size_t) sparams.n_min) { - result.clear(); - } - } - - void accept(uint16_t n_accepted) override { - // noop - GGML_UNUSED(n_accepted); - } - - int32_t n_max(const common_params_speculative & params) const override { - return params.draft.n_max; - } - - int32_t n_min(const common_params_speculative & params) const override { - return params.draft.n_min; - } - - std::string replace_to_dft(const std::string & input) const { - std::string result = input; - - for (const auto & pair : this->vocab_map) { - size_t pos = result.find(pair.first); - while (pos != std::string::npos) { - result.replace(pos, pair.first.length(), pair.second); - pos = result.find(pair.first, pos + pair.second.length()); - } - } - - return result; - } - - std::string replace_to_tgt(const std::string & input) const { - std::string result = input; - - for (const auto & pair : this->vocab_map) { - size_t pos = result.find(pair.second); - while (pos != std::string::npos) { - result.replace(pos, pair.second.length(), pair.first); - pos = result.find(pair.second, pos + pair.first.length()); - } - } - - return result; - } -}; - -struct common_speculative_state_eagle3 : public common_speculative_state { - common_speculative_state_eagle3(enum common_speculative_type type) : common_speculative_state(type) {} - - void begin(const llama_tokens & prompt) override { - GGML_UNUSED(prompt); - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & draft_tokens) override { - // TODO: implement - GGML_UNUSED(params); - GGML_UNUSED(prompt_tgt); - GGML_UNUSED(id_last); - GGML_UNUSED(draft_tokens); - } - - void accept(uint16_t n_accepted) override { - // noop - GGML_UNUSED(n_accepted); - } - - int32_t n_max(const common_params_speculative & params) const override { - return params.draft.n_max; - } - - int32_t n_min(const common_params_speculative & params) const override { - return params.draft.n_min; - } -}; - -// state of self-speculation (simple implementation, not ngram-map) -struct common_speculative_state_ngram_simple : public common_speculative_state { - common_ngram_simple_config config; - - common_speculative_state_ngram_simple( - enum common_speculative_type type, - common_ngram_simple_config config) - : common_speculative_state(type), config(config) {} - - void begin(const llama_tokens & prompt) override { - GGML_UNUSED(prompt); - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) override { - - result = common_ngram_simple_draft(config, prompt_tgt, id_last); - GGML_UNUSED(params); - } - - void accept(uint16_t n_accepted) override { - // noop - GGML_UNUSED(n_accepted); - } - - int32_t n_max(const common_params_speculative & /*params*/) const override { - return config.size_mgram; - } - - int32_t n_min(const common_params_speculative & /*params*/) const override { - return config.size_mgram; - } -}; - -struct common_speculative_state_ngram_map_k : public common_speculative_state { - // draft ngram map for speculative decoding without draft model - common_ngram_map config; - - common_speculative_state_ngram_map_k( - enum common_speculative_type type, - common_ngram_map config) - : common_speculative_state(type), config(std::move(config)) {} - - void begin(const llama_tokens & prompt) override { - common_ngram_map_begin(config, prompt); - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) override { - common_ngram_map_draft(config, prompt_tgt, id_last, result); - GGML_UNUSED(params); - } - - void accept(uint16_t n_accepted) override { - common_ngram_map_accept(config, n_accepted); - } - - int32_t n_max(const common_params_speculative & /*params*/) const override { - return config.size_value; - } - - int32_t n_min(const common_params_speculative & /*params*/) const override { - return config.size_value; - } -}; - -struct common_speculative_state_ngram_mod : public common_speculative_state { - common_ngram_mod & mod; - - // the last position in the prompt that was added to the ngram container - size_t i_last = 0; - - // length of the last drafted n‑gram (number of tokens returned by draft) - size_t n_draft_last = 0; - - // consecutive accept rounds with low acceptance fraction (< 0.5) - int n_low = 0; - - // enable trace logging if LLAMA_TRACE is set - const bool verbose; - - common_speculative_state_ngram_mod(enum common_speculative_type type, common_ngram_mod & mod) - : common_speculative_state(type), mod(mod), verbose(std::getenv("LLAMA_TRACE") != nullptr) { - static_assert(sizeof(llama_token) == sizeof(common_ngram_mod::entry_t)); - } - - void begin(const llama_tokens & prompt) override { - i_last = 0; - - n_draft_last = 0; - - const size_t n = mod.get_n(); - - if (prompt.size() < n) { - return; - } - - for (size_t i = 0; i < prompt.size() - n; ++i) { - mod.add(prompt.data() + i); - } - - i_last = prompt.size() - n; - - const double f = (double)mod.get_used() / (double)mod.size(); - LOG_INF("%s: ngram_mod occupancy = %zu/%zu (%.2f)\n", __func__, mod.get_used(), mod.size(), f); - - constexpr double f_thold = 0.25; - if (f > f_thold) { - LOG_WRN("%s: ngram_mod occupancy %.2f exceeds threshold (%.2f) - resetting\n", __func__, f, f_thold); - - mod.reset(); - } - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) override { - const auto & sparams = params.ngram_mod; - - n_draft_last = 0; - - const size_t cur_len = prompt_tgt.size(); - if (cur_len < mod.get_n()) { - return; - } - - const size_t n = mod.get_n(); - - // add new ngrams in chunks - if (i_last + 32 < cur_len) { - for (size_t i = i_last; i < cur_len - n; ++i) { - mod.add(prompt_tgt.data() + i); - } - - i_last = cur_len - n; - } - - result.resize(n + sparams.n_max); - for (size_t i = 0; i < n - 1; ++i) { - result[i] = prompt_tgt[cur_len - n + 1 + i]; - } - result[n - 1] = id_last; - - for (int i = 0; i < sparams.n_max; ++i) { - const llama_token token = mod.get(result.data() + i); - if (token == common_ngram_mod::EMPTY) { - if (i < sparams.n_min) { - result.clear(); - return; - } - - result.resize(n + i); - break; - } - result[n + i] = token; - } - - // only return the m tokens that were drafted - for (size_t i = 0; n + i < result.size(); ++i) { - result[i] = result[n + i]; - } - result.resize(result.size() - n); - - // store length of drafted n‑gram for later acceptance analysis - n_draft_last = result.size(); - } - - void accept(uint16_t n_accepted) override { - // compute acceptance fraction if we have a recorded draft length - if (n_draft_last > 0) { - const double f_acc = (double)n_accepted / (double)n_draft_last; - if (f_acc < 0.5) { - n_low++; - if (n_low >= 3) { - if (verbose) { - LOG_WRN("%s: low acceptance streak (%d) – resetting ngram_mod\n", __func__, n_low); - } - - mod.reset(); - n_low = 0; - i_last = 0; - } - } else { - n_low = 0; - } - } - } - - int32_t n_max(const common_params_speculative & params) const override { - return params.ngram_mod.n_max; - } - - int32_t n_min(const common_params_speculative & params) const override { - return params.ngram_mod.n_min; - } -}; - -struct common_speculative_state_ngram_cache : public common_speculative_state { - uint16_t n_draft; - bool save_dynamic; - bool save_static; - - common_ngram_cache ngram_cache_context; - common_ngram_cache ngram_cache_dynamic; - common_ngram_cache ngram_cache_static; - - size_t cache_size = 0; // number of tokens in n-gram cache - - common_speculative_state_ngram_cache( - const enum common_speculative_type type, - const std::string & path_static, - const std::string & path_dynamic, - uint16_t n_draft, - bool save_dynamic, - bool save_static) - : common_speculative_state(type) - , n_draft(n_draft) - , save_dynamic(save_dynamic) - , save_static(save_static) - { - if (!path_static.empty()) { - try { - ngram_cache_static = common_ngram_cache_load(path_static); - } catch (...) { - LOG_ERR("failed to open static lookup cache: %s", path_static.c_str()); - GGML_ABORT("Couldn't read static lookup cache"); - } - } - - if (!path_dynamic.empty()) { - try { - ngram_cache_dynamic = common_ngram_cache_load(path_dynamic); - } catch (...) { - LOG_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str()); - GGML_ABORT("Couldn't read dynamic lookup cache"); - } - } - } - - void begin(const llama_tokens & prompt) override { - GGML_UNUSED(prompt); - } - - void draft( - const common_params_speculative & params, - const llama_tokens & prompt_tgt, - llama_token id_last, - llama_tokens & result) override { - GGML_UNUSED(params); - - if (cache_size < prompt_tgt.size() + 1) { - llama_tokens tokens_new; - tokens_new.reserve(prompt_tgt.size() + 1 - cache_size); - for (size_t j = cache_size; j < prompt_tgt.size(); ++j) { - tokens_new.push_back(prompt_tgt[j]); - } - tokens_new.push_back(id_last); // add the last token - - // Update context ngram cache with new prompt_tgt: - common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, - tokens_new, tokens_new.size(), false); - cache_size = prompt_tgt.size() + 1; - } - - llama_tokens inp; - inp.reserve(prompt_tgt.size() + 1); - for (size_t j = 0; j < prompt_tgt.size(); ++j) { - inp.push_back(prompt_tgt[j]); - } - inp.push_back(id_last); - - result.push_back(id_last); - - common_ngram_cache_draft(inp, result, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, - ngram_cache_context, - ngram_cache_dynamic, - ngram_cache_static); - - if (result.size() > 0) { - // delete first token in result (which is the id_last token) - result.erase(result.begin()); - } - } - - void accept(uint16_t n_accepted) override { - // TODO: noop - GGML_UNUSED(n_accepted); - } - - int32_t n_max(const common_params_speculative & /*params*/) const override { - return n_draft; - } - - int32_t n_min(const common_params_speculative & /*params*/) const override { - return 0; - } -}; - -struct common_speculative { - std::vector> impls; // list of implementations to use and their states - - common_speculative_state * curr_impl = nullptr; // current implementation in use (for stats) -}; - -static common_ngram_map get_common_ngram_map( - common_speculative_type type, - const common_params_speculative_ngram_map & config) { - uint16_t size_key = config.size_n; - uint16_t size_value = config.size_m; - bool key_only = type == COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K; - uint16_t min_hits = config.min_hits; - - return common_ngram_map(size_key, size_value, key_only, min_hits); -} - -static common_speculative_state_ngram_cache create_state_ngram_cache( - const std::string & path_static, const std::string & path_dynamic, - const common_speculative_config & config) { - uint16_t n_draft = 8; // TODO get from config? - - // TODO bool param in common/common.h to set save_static/save_dynamic? - bool save_static = false; - bool save_dynamic = false; - - common_speculative_state_ngram_cache state(config.type, path_static, path_dynamic, n_draft, save_static, save_dynamic); - - return state; -} - -std::string common_speculative_type_name_str() { - std::string result; - for (size_t i = 0; i < common_speculative_types.size(); i++) { - if (i > 0) { - result += ", "; - } - result += common_speculative_type_to_str(common_speculative_types[i]); - } - return result; -} - -std::string common_speculative_type_to_str(enum common_speculative_type type) { - switch (type) { - case COMMON_SPECULATIVE_TYPE_NONE: return "none"; - case COMMON_SPECULATIVE_TYPE_DRAFT: return "draft"; - case COMMON_SPECULATIVE_TYPE_EAGLE3: return "eagle3"; - case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram_simple"; - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram_map_k"; - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram_map_k4v"; - case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: return "ngram_mod"; - case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: return "ngram_cache"; - default: return "unknown"; - } -} - -enum common_speculative_type common_speculative_type_from_name(const std::string & name) { - const auto it = common_speculative_type_from_name_map.find(name); - if (it == common_speculative_type_from_name_map.end()) { - return COMMON_SPECULATIVE_TYPE_COUNT; - } - return it->second; -} - -// initialization of the speculative decoding system -// -common_speculative * common_speculative_init( - common_params_speculative & params, - llama_context * ctx_tgt) { - llama_context * ctx_dft = nullptr; - if (params.draft.model) { - ctx_dft = llama_init_from_model(params.draft.model, params.draft.cparams); - if (ctx_dft == nullptr) { - LOG_ERR("%s", "failed to create draft context\n"); - return nullptr; - } - } - - // Compute the implementations to use based on the config and their order of preference - std::vector configs = {}; // list of speculative configs to try - { - bool has_draft = !params.draft.mparams.path.empty(); - bool has_draft_eagle3 = false; // TODO PR-18039: if params.speculative.eagle3 - - bool has_ngram_cache = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_CACHE); - bool has_ngram_simple = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE); - bool has_ngram_map_k = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K); - bool has_ngram_map_k4v = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V); - bool has_ngram_mod = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_MOD); - - // In a more complex implementation we could use the same implementation but with different parameters. - // This was initially used in PR-18471 but removed to simplify the code. - if (has_ngram_simple) { - // This implementation can guess a lot of tokens without any draft model. - configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, params)); - } - if (has_ngram_map_k) { - configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, params)); - } - if (has_ngram_map_k4v) { - // This implementation can guess tokens with high acceptance rate but is more expensive. - configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, params)); - } - if (has_ngram_mod) { - auto & sparams = params.ngram_mod; - - if (!sparams.obj) { - sparams.obj = std::make_shared(sparams.n_match, 4*1024*1024); - - LOG_INF("%s: initialized ngram_mod with n_match=%d, size=%zu (%.3f MB)\n", __func__, - sparams.n_match, sparams.obj->size(), (float)(sparams.obj->size_bytes())/1024/1024); - - if (sparams.n_match < 16) { - LOG_WRN("%s: ngram_mod n_match=%d is too small - poor quality is possible, " - "see: https://github.com/ggml-org/llama.cpp/pull/19164\n", __func__, sparams.n_match); - } - } - - configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, params)); - } - if (has_ngram_cache) { - configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, params)); - } - if (has_draft) { - configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT, params)); - } - if (has_draft_eagle3) { - configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_EAGLE3, params)); - } - } - - std::vector> impls = {}; - - for (const common_speculative_config & config : configs) { - LOG_DBG("%s: adding implementation %s\n", __func__, common_speculative_type_to_str(config.type).c_str()); - switch (config.type) { - case COMMON_SPECULATIVE_TYPE_NONE: - break; - case COMMON_SPECULATIVE_TYPE_DRAFT: { - const bool use_ckpt = common_context_can_seq_rm(ctx_dft) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; - - impls.push_back(std::make_unique(config.type, - /* .ctx_tgt = */ ctx_tgt, - /* .ctx_dft = */ ctx_dft, - /* .replacements = */ params.draft.replacements, - /* .use_ckpt = */ use_ckpt - )); - break; - } - case COMMON_SPECULATIVE_TYPE_EAGLE3: { - impls.push_back(std::make_unique(config.type)); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { - common_ngram_map ngram_map = get_common_ngram_map(config.type, config.params.ngram_simple); - - uint16_t ngram_size_key = ngram_map.size_key; - uint16_t mgram_size_value = ngram_map.size_value; - - auto config_simple = common_ngram_simple_config { - /* .size_ngram = */ ngram_size_key, - /* .size_mgram = */ mgram_size_value - }; - auto state = std::make_unique( - /* .type = */ config.type, - /* .state = */ config_simple - ); - impls.push_back(std::move(state)); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: - case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: { - impls.push_back(std::make_unique( - (config.type), - get_common_ngram_map(config.type, config.params.ngram_map_k) - )); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: { - GGML_ASSERT(config.params.ngram_mod.obj); - impls.push_back(std::make_unique(config.type, *config.params.ngram_mod.obj)); - break; - } - case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: { - auto state = create_state_ngram_cache(params.ngram_cache.lookup_cache_static, params.ngram_cache.lookup_cache_dynamic, config); - impls.push_back(std::make_unique(state)); - break; - } - default: - break; - } - } - - if (impls.empty()) { - LOG_WRN("%s", "no implementations specified for speculative decoding\n"); - return nullptr; - } - - auto * result = new common_speculative { - /* .impls = */ std::move(impls), - /* .curr_impl = */ nullptr, - }; - - return result; -} - -void common_speculative_free(common_speculative * spec) { - if (spec == nullptr) { - return; - } - - delete spec; -} - -void common_speculative_begin(common_speculative * spec, const llama_tokens & prompt) { - if (spec == nullptr) { - return; - } - - for (auto & impl : spec->impls) { - common_time_meas tm(impl->t_begin_us, !impl->gen_perf); - impl->begin(prompt); - impl->n_call_begin++; - } -} - -llama_tokens common_speculative_draft( - common_speculative * spec, - const common_params_speculative & params, - const llama_tokens & prompt_tgt, // specified in target model vocab - llama_token id_last) { - llama_tokens result; - - spec->curr_impl = nullptr; // reset current implementation - - for (auto & impl : spec->impls) { - { - common_time_meas tm(impl->t_draft_us, !impl->gen_perf); - impl->draft(params, prompt_tgt, id_last, result); - impl->n_call_draft++; - } - - { - const int n_min = impl->n_min(params); - - if (!result.empty() && (int) result.size() < n_min) { - LOG_DBG("%s: ignoring small draft: %d < %d\n", __func__, (int) result.size(), n_min); - result.clear(); - } - } - - if (!result.empty()) { - LOG_DBG("%s: called impl %s, hist size = %zu, call_count = %zu, gen = %zu\n", __func__, - common_speculative_type_to_str(impl.get()->type).c_str(), prompt_tgt.size(), - impl.get()->n_call_draft, result.size()); - - spec->curr_impl = impl.get(); // set current implementation for stats - impl->n_gen_drafts++; - impl->n_gen_tokens += result.size(); - - break; // we have a draft, so break out of the loop and return it. - } - } - - return result; -} - -void common_speculative_accept(common_speculative * spec, uint16_t n_accepted) { - if (n_accepted == 0) { - return; - } - - common_speculative_state * impl = spec->curr_impl; - - GGML_ASSERT(impl); - - { - common_time_meas tm(impl->t_accept_us, !impl->gen_perf); - if (n_accepted > 0) { - impl->n_acc_drafts++; - impl->n_acc_tokens += n_accepted; - } - - impl->accept(n_accepted); - impl->n_call_accept++; - } -} - -int32_t common_speculative_n_max(const common_speculative * spec, const common_params_speculative & params) { - if (spec == nullptr) { - return 0; - } - - int32_t n_max = 0; - for (const auto & impl : spec->impls) { - n_max = std::max(n_max, impl->n_max(params)); - } - - return n_max; -} - -int32_t common_speculative_n_min(const common_speculative * spec, const common_params_speculative & params) { - if (spec == nullptr) { - return 0; - } - - int32_t n_min = 0; - for (const auto & impl : spec->impls) { - n_min = std::max(n_min, impl->n_min(params)); - } - - return n_min; -} - -void common_speculative_print_stats(const common_speculative * spec) { - if (spec == nullptr) { - return; - } - - for (const auto & impl : spec->impls) { - std::string str_perf; - if (impl->gen_perf) { - std::ostringstream oss; - oss << std::fixed << std::setprecision(3) << impl->t_begin_us / 1000.0 << ", "; - oss << std::fixed << std::setprecision(3) << impl->t_draft_us / 1000.0 << ", "; - oss << std::fixed << std::setprecision(3) << impl->t_accept_us / 1000.0; - str_perf = ", dur(b,g,a) = " + oss.str() + " ms"; - } else { - str_perf = ""; - } - - LOG_INF("statistics %s: #calls(b,g,a) = %zu %zu %zu, #gen drafts = %zu, #acc drafts = %zu, #gen tokens = %zu, #acc tokens = %zu%s\n", - common_speculative_type_to_str(impl->type).c_str(), - impl->n_call_begin, impl->n_call_draft, impl->n_call_accept, - impl->n_gen_drafts, - impl->n_acc_drafts, - impl->n_gen_tokens, - impl->n_acc_tokens, - str_perf.c_str()); - } -} diff --git a/common/speculative.h b/common/speculative.h deleted file mode 100644 index 14744763170e..000000000000 --- a/common/speculative.h +++ /dev/null @@ -1,46 +0,0 @@ -#pragma once - -#include "llama.h" -#include "common.h" - -struct common_speculative; - -// comma separated list of all types -std::string common_speculative_type_name_str(); - -// convert string to type -enum common_speculative_type common_speculative_type_from_name(const std::string & name); - -// convert type to string -std::string common_speculative_type_to_str(enum common_speculative_type type); - -common_speculative * common_speculative_init( - common_params_speculative & params, - llama_context * ctx_tgt); - -void common_speculative_free(common_speculative * spec); - -// optionally call once at the beginning of a new generation -void common_speculative_begin(common_speculative * spec, const llama_tokens & prompt); - -// sample up to n_draft tokens and add them to the batch using the draft model -llama_tokens common_speculative_draft( - common_speculative * spec, - const common_params_speculative & params, - const llama_tokens & prompt, - llama_token id_last); - -// informs the speculative decoder that n_accepted tokens were accepted by the target model -void common_speculative_accept(common_speculative * spec, uint16_t n_accepted); - -int32_t common_speculative_n_max(const common_speculative * spec, const common_params_speculative & params); -int32_t common_speculative_n_min(const common_speculative * spec, const common_params_speculative & params); - -// print statistics about the speculative decoding -void common_speculative_print_stats(const common_speculative * spec); - -struct common_speculative_deleter { - void operator()(common_speculative * s) { common_speculative_free(s); } -}; - -typedef std::unique_ptr common_speculative_ptr; diff --git a/common/unicode-rust/Cargo.toml b/common/unicode-rust/Cargo.toml new file mode 100644 index 000000000000..9bae9352135b --- /dev/null +++ b/common/unicode-rust/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "llama-common-unicode-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[lib] +crate-type = ["staticlib", "rlib"] + +[dependencies] diff --git a/common/unicode-rust/src/lib.rs b/common/unicode-rust/src/lib.rs new file mode 100644 index 000000000000..d5d7e235012d --- /dev/null +++ b/common/unicode-rust/src/lib.rs @@ -0,0 +1,5233 @@ +use std::ffi::{c_char, c_int, CStr}; +use std::slice; +use std::time::{SystemTime, UNIX_EPOCH}; + +#[repr(C)] +pub struct Utf8ParseResult { + codepoint: u32, + bytes_consumed: usize, + status: i32, +} + +#[repr(C)] +pub struct UnicodeString { + data: *mut u8, + len: usize, +} + +#[repr(C)] +pub struct StringView { + data: *const u8, + len: usize, +} + +#[repr(C)] +pub struct StringList { + data: *mut UnicodeString, + len: usize, +} + +#[repr(C)] +pub struct GgufSplitInfo { + prefix: UnicodeString, + tag: UnicodeString, + index: c_int, + count: c_int, +} + +const SUCCESS: i32 = 0; +const INCOMPLETE: i32 = 1; +const INVALID: i32 = 2; + +#[no_mangle] +pub extern "C" fn llama_common_utf8_sequence_length(first_byte: u8) -> usize { + const LOOKUP: [usize; 16] = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4]; + LOOKUP[(first_byte >> 4) as usize] +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_utf8_is_complete(data: *const u8, len: usize) -> bool { + if data.is_null() && len != 0 { + return false; + } + utf8_is_complete(slice::from_raw_parts(data, len)) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_parse_utf8_codepoint( + data: *const u8, + len: usize, + offset: usize, +) -> Utf8ParseResult { + if data.is_null() && len != 0 { + return result(INVALID, 0, 0); + } + parse_utf8_codepoint(slice::from_raw_parts(data, len), offset) +} + +#[no_mangle] +pub extern "C" fn llama_common_unicode_cpt_to_utf8(cpt: u32) -> UnicodeString { + match unicode_cpt_to_utf8(cpt) { + Some(bytes) => into_ffi(bytes), + None => UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }, + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_unicode_cpts_to_utf8_rust( + data: *const u32, + len: usize, +) -> UnicodeString { + if data.is_null() && len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + unicode_cpts_to_utf8(slice::from_raw_parts(data, len)) +} + +#[no_mangle] +pub extern "C" fn llama_common_jinja_token_type_to_string_rust(token_type: c_int) -> *const c_char { + match token_type { + 0 => c"eof".as_ptr(), + 1 => c"text".as_ptr(), + 2 => c"numeric_literal".as_ptr(), + 3 => c"string_literal".as_ptr(), + 4 => c"identifier".as_ptr(), + 5 => c"equals".as_ptr(), + 6 => c"open_paren".as_ptr(), + 7 => c"close_paren".as_ptr(), + 8 => c"open_statement".as_ptr(), + 9 => c"close_statement".as_ptr(), + 10 => c"open_expression".as_ptr(), + 11 => c"close_expression".as_ptr(), + 12 => c"open_square_bracket".as_ptr(), + 13 => c"close_square_bracket".as_ptr(), + 14 => c"open_curly_bracket".as_ptr(), + 15 => c"close_curly_bracket".as_ptr(), + 16 => c"comma".as_ptr(), + 17 => c"dot".as_ptr(), + 18 => c"colon".as_ptr(), + 19 => c"pipe".as_ptr(), + 20 => c"call_operator".as_ptr(), + 21 => c"additive_binary_operator".as_ptr(), + 22 => c"multiplicative_binary_operator".as_ptr(), + 23 => c"comparison_binary_operator".as_ptr(), + 24 => c"unary_operator".as_ptr(), + 25 => c"comment".as_ptr(), + _ => c"unknown".as_ptr(), + } +} + +#[no_mangle] +pub extern "C" fn llama_common_jinja_is_word_rust(c: u8) -> bool { + c.is_ascii_alphanumeric() || c == b'_' +} + +#[no_mangle] +pub extern "C" fn llama_common_jinja_is_integer_rust(c: u8) -> bool { + c.is_ascii_digit() +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_http_show_masked_url_rust( + scheme: *const u8, + scheme_len: usize, + has_user: bool, + host: *const u8, + host_len: usize, + path: *const u8, + path_len: usize, +) -> UnicodeString { + let Some(scheme) = byte_slice(scheme, scheme_len) else { + return into_ffi(Vec::new()); + }; + let Some(host) = byte_slice(host, host_len) else { + return into_ffi(Vec::new()); + }; + let Some(path) = byte_slice(path, path_len) else { + return into_ffi(Vec::new()); + }; + http_show_masked_url(scheme, has_user, host, path) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_peak_source_rust( + source: *const u8, + source_len: usize, + pos: usize, + max_peak_chars: usize, +) -> UnicodeString { + let Some(source) = byte_slice(source, source_len) else { + return into_ffi(b"(no source available)".to_vec()); + }; + peak_source(source, pos, max_peak_chars) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_fmt_error_with_source_rust( + tag: *const u8, + tag_len: usize, + msg: *const u8, + msg_len: usize, + source: *const u8, + source_len: usize, + pos: usize, +) -> UnicodeString { + let Some(tag) = byte_slice(tag, tag_len) else { + return into_ffi(Vec::new()); + }; + let Some(msg) = byte_slice(msg, msg_len) else { + return into_ffi(Vec::new()); + }; + let Some(source) = byte_slice(source, source_len) else { + return into_ffi(Vec::new()); + }; + fmt_error_with_source(tag, msg, source, pos) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_unicode_string_free(value: UnicodeString) { + if !value.data.is_null() { + drop(Vec::from_raw_parts(value.data, value.len, value.len)); + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_string_repeat_rust( + data: *const u8, + len: usize, + n: usize, +) -> UnicodeString { + if data.is_null() && len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + if len == 0 { + return string_repeat(&[], n); + } + string_repeat(slice::from_raw_parts(data, len), n) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_int_vec_to_string_rust( + data: *const i32, + len: usize, +) -> UnicodeString { + if data.is_null() && len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + if len == 0 { + return int_vec_to_string(&[]); + } + int_vec_to_string(slice::from_raw_parts(data, len)) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_string_trim_rust( + data: *const u8, + len: usize, + mode: c_int, +) -> UnicodeString { + if data.is_null() && len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + if len == 0 { + return into_ffi(Vec::new()); + } + string_trim(slice::from_raw_parts(data, len), mode) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_prefix_len_rust( + left: *const u8, + left_len: usize, + right: *const u8, + right_len: usize, +) -> usize { + let Some((left, right)) = byte_pair(left, left_len, right, right_len) else { + return 0; + }; + common_prefix_len(left, right) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_suffix_len_rust( + left: *const u8, + left_len: usize, + right: *const u8, + right_len: usize, +) -> usize { + let Some((left, right)) = byte_pair(left, left_len, right, right_len) else { + return 0; + }; + common_suffix_len(left, right) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_string_replace_all_rust( + data: *const u8, + len: usize, + search: *const u8, + search_len: usize, + replace: *const u8, + replace_len: usize, +) -> UnicodeString { + let Some((data, search)) = byte_pair(data, len, search, search_len) else { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + }; + if replace.is_null() && replace_len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + let replace = if replace_len == 0 { + &[] + } else { + slice::from_raw_parts(replace, replace_len) + }; + string_replace_all(data, search, replace) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_regex_escape_rust( + data: *const u8, + len: usize, +) -> UnicodeString { + if data.is_null() && len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + if len == 0 { + return into_ffi(Vec::new()); + } + regex_escape(slice::from_raw_parts(data, len)) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_string_join_rust( + values: *const StringView, + values_len: usize, + separator: *const u8, + separator_len: usize, +) -> UnicodeString { + if values.is_null() && values_len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + if separator.is_null() && separator_len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + + let values = if values_len == 0 { + &[] + } else { + slice::from_raw_parts(values, values_len) + }; + let separator = if separator_len == 0 { + &[] + } else { + slice::from_raw_parts(separator, separator_len) + }; + string_join(values, separator) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_string_split_rust( + data: *const u8, + len: usize, + delimiter: *const u8, + delimiter_len: usize, +) -> StringList { + let Some((data, delimiter)) = byte_pair(data, len, delimiter, delimiter_len) else { + return StringList { + data: std::ptr::null_mut(), + len: 0, + }; + }; + string_split(data, delimiter) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_parse_csv_row_rust( + data: *const u8, + len: usize, +) -> StringList { + if data.is_null() && len != 0 { + return StringList { + data: std::ptr::null_mut(), + len: 0, + }; + } + if len == 0 { + return into_string_list(vec![Vec::new()]); + } + parse_csv_row(slice::from_raw_parts(data, len)) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_trim_trailing_space_view_rust( + data: *const u8, + len: usize, + max: c_int, +) -> StringView { + let Some(data) = byte_slice(data, len) else { + return StringView { + data: std::ptr::null(), + len: 0, + }; + }; + trim_trailing_space_view(data, max) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_trim_leading_space_view_rust( + data: *const u8, + len: usize, + max: c_int, +) -> StringView { + let Some(data) = byte_slice(data, len) else { + return StringView { + data: std::ptr::null(), + len: 0, + }; + }; + trim_leading_space_view(data, max) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_string_process_escapes_rust( + data: *const u8, + len: usize, +) -> UnicodeString { + if data.is_null() && len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + if len == 0 { + return into_ffi(Vec::new()); + } + string_process_escapes(slice::from_raw_parts(data, len)) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_clean_file_name_rust( + data: *const u8, + len: usize, +) -> UnicodeString { + if data.is_null() && len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + if len == 0 { + return into_ffi(Vec::new()); + } + clean_file_name(slice::from_raw_parts(data, len)) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_string_diff_rust( + last: *const u8, + last_len: usize, + current: *const u8, + current_len: usize, + status: *mut c_int, +) -> UnicodeString { + if !status.is_null() { + *status = -1; + } + let Some((last, current)) = byte_pair(last, last_len, current, current_len) else { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + }; + + let (result, result_status) = string_diff(last, current); + if !status.is_null() { + *status = result_status; + } + result +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_string_starts_with_rust( + data: *const u8, + len: usize, + prefix: *const u8, + prefix_len: usize, +) -> bool { + let Some((data, prefix)) = byte_pair(data, len, prefix, prefix_len) else { + return false; + }; + data.starts_with(prefix) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_string_ends_with_rust( + data: *const u8, + len: usize, + suffix: *const u8, + suffix_len: usize, +) -> bool { + let Some((data, suffix)) = byte_pair(data, len, suffix, suffix_len) else { + return false; + }; + data.ends_with(suffix) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_glob_match_rust( + pattern: *const u8, + pattern_len: usize, + str: *const u8, + str_len: usize, +) -> bool { + let Some((pattern, str)) = byte_pair(pattern, pattern_len, str, str_len) else { + return false; + }; + glob_match(c_string_bytes(pattern), c_string_bytes(str)) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_fs_validate_filename_rust( + filename: *const u8, + filename_len: usize, + allow_subdirs: bool, +) -> bool { + if filename.is_null() && filename_len != 0 { + return false; + } + let filename = if filename_len == 0 { + &[] + } else { + slice::from_raw_parts(filename, filename_len) + }; + fs_validate_filename(filename, allow_subdirs) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_string_lstrip_chars_rust( + data: *const u8, + len: usize, + chars: *const c_char, +) -> UnicodeString { + if data.is_null() && len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + if chars.is_null() { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + let data = if len == 0 { + &[] + } else { + slice::from_raw_parts(data, len) + }; + string_lstrip_chars(data, CStr::from_ptr(chars).to_bytes()) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_string_rstrip_chars_rust( + data: *const u8, + len: usize, + chars: *const c_char, +) -> UnicodeString { + if data.is_null() && len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + if chars.is_null() { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + let data = if len == 0 { + &[] + } else { + slice::from_raw_parts(data, len) + }; + string_rstrip_chars(data, CStr::from_ptr(chars).to_bytes()) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_break_str_into_lines_rust( + data: *const u8, + len: usize, + max_char_per_line: usize, +) -> StringList { + if data.is_null() && len != 0 { + return StringList { + data: std::ptr::null_mut(), + len: 0, + }; + } + let data = if len == 0 { + &[] + } else { + slice::from_raw_parts(data, len) + }; + break_str_into_lines(data, max_char_per_line) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_rm_leading_dashes_rust( + data: *const u8, + len: usize, +) -> UnicodeString { + if data.is_null() && len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + let data = if len == 0 { + &[] + } else { + slice::from_raw_parts(data, len) + }; + rm_leading_dashes(data) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_until_common_prefix_rust( + full: *const u8, + full_len: usize, + left: *const u8, + left_len: usize, + right: *const u8, + right_len: usize, +) -> UnicodeString { + let Some((full, left)) = byte_pair(full, full_len, left, left_len) else { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + }; + if right.is_null() && right_len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + let right = if right_len == 0 { + &[] + } else { + slice::from_raw_parts(right, right_len) + }; + until_common_prefix(full, left, right) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_after_common_suffix_rust( + full: *const u8, + full_len: usize, + left: *const u8, + left_len: usize, + right: *const u8, + right_len: usize, +) -> UnicodeString { + let Some((full, left)) = byte_pair(full, full_len, left, left_len) else { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + }; + if right.is_null() && right_len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + let right = if right_len == 0 { + &[] + } else { + slice::from_raw_parts(right, right_len) + }; + after_common_suffix(full, left, right) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_json_ensure_ascii_preserving_format_rust( + data: *const u8, + len: usize, +) -> UnicodeString { + if data.is_null() && len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + let data = if len == 0 { + &[] + } else { + slice::from_raw_parts(data, len) + }; + json_ensure_ascii_preserving_format(data) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_normalize_quotes_to_json_rust( + data: *const u8, + len: usize, +) -> UnicodeString { + let Some(data) = byte_slice(data, len) else { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + }; + normalize_quotes_to_json(data) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_string_list_free(value: StringList) { + if !value.data.is_null() { + let items = Vec::from_raw_parts(value.data, value.len, value.len); + for item in items { + llama_common_unicode_string_free(item); + } + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_ascii_case_equal( + a: *const c_char, + b: *const c_char, +) -> c_int { + ascii_case_equal(a, b) as c_int +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_opt_get_optimizer_rust(name: *const c_char) -> c_int { + if ascii_case_equal(name, c"adamw".as_ptr()) { + 0 + } else if ascii_case_equal(name, c"sgd".as_ptr()) { + 1 + } else { + 2 + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_arg_is_truthy_rust(value: *const c_char) -> c_int { + matches_cstr(value, &[b"on", b"enabled", b"true", b"1"]) as c_int +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_arg_is_falsey_rust(value: *const c_char) -> c_int { + matches_cstr(value, &[b"off", b"disabled", b"false", b"0"]) as c_int +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_arg_is_autoy_rust(value: *const c_char) -> c_int { + matches_cstr(value, &[b"auto", b"-1"]) as c_int +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_parse_bool_value_rust(value: *const c_char) -> c_int { + if matches_cstr(value, &[b"on", b"enabled", b"true", b"1"]) { + 1 + } else if matches_cstr(value, &[b"off", b"disabled", b"false", b"0"]) { + 0 + } else { + -1 + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_parse_cpu_mask_rust( + mask: *const u8, + mask_len: usize, + boolmask: *mut bool, + boolmask_len: usize, +) -> c_int { + if (mask.is_null() && mask_len != 0) || (boolmask.is_null() && boolmask_len != 0) { + return -2; + } + let mask = if mask_len == 0 { + &[] + } else { + slice::from_raw_parts(mask, mask_len) + }; + let boolmask = if boolmask_len == 0 { + &mut [] + } else { + slice::from_raw_parts_mut(boolmask, boolmask_len) + }; + parse_cpu_mask(mask, boolmask) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_get_line_col_rust( + data: *const u8, + len: usize, + pos: usize, +) -> UnicodeString { + let Some(data) = byte_slice(data, len) else { + return into_ffi(b"line 1, column 1".to_vec()); + }; + get_line_col(data, pos) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_chat_tool_choice_parse_oaicompat_rust( + value: *const c_char, +) -> c_int { + if value.is_null() { + return -1; + } + + match CStr::from_ptr(value).to_bytes() { + b"auto" => 0, + b"required" => 1, + b"none" => 2, + _ => -1, + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_kv_cache_type_from_str_rust(value: *const c_char) -> c_int { + if value.is_null() { + return -1; + } + + match CStr::from_ptr(value).to_bytes() { + b"f32" => 0, + b"f16" => 1, + b"q4_0" => 2, + b"q4_1" => 3, + b"q5_0" => 6, + b"q5_1" => 7, + b"q8_0" => 8, + b"iq4_nl" => 20, + b"bf16" => 30, + _ => -1, + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_embd_normalize_rust( + inp: *const f32, + out: *mut f32, + n: c_int, + embd_norm: c_int, +) { + if n <= 0 || inp.is_null() || out.is_null() { + return; + } + + let inp = slice::from_raw_parts(inp, n as usize); + let out = slice::from_raw_parts_mut(out, n as usize); + embd_normalize(inp, out, embd_norm); +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_embd_similarity_cos_rust( + embd1: *const f32, + embd2: *const f32, + n: c_int, +) -> f32 { + if n <= 0 { + return 1.0; + } + if embd1.is_null() || embd2.is_null() { + return 0.0; + } + + let embd1 = slice::from_raw_parts(embd1, n as usize); + let embd2 = slice::from_raw_parts(embd2, n as usize); + embd_similarity_cos(embd1, embd2) +} + +#[no_mangle] +pub extern "C" fn llama_common_get_all_kv_cache_types_rust() -> *const c_char { + c"f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1".as_ptr() +} + +#[no_mangle] +pub extern "C" fn llama_common_bool_to_string_rust(value: bool) -> *const c_char { + if value { + c"true".as_ptr() + } else { + c"false".as_ptr() + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_lr_opt_init_rust( + lr0: f32, + lr_min: f32, + decay_epochs: f32, + epochs: u32, + out_decay_epochs: *mut f32, + out_scale_epoch: *mut f32, +) { + if out_decay_epochs.is_null() || out_scale_epoch.is_null() { + return; + } + let (decay_epochs, scale_epoch) = lr_opt_init(lr0, lr_min, decay_epochs, epochs); + *out_decay_epochs = decay_epochs; + *out_scale_epoch = scale_epoch; +} + +#[no_mangle] +pub extern "C" fn llama_common_lr_opt_get_lr_rust( + lr0: f32, + lr_min: f32, + decay_epochs: f32, + scale_epoch: f32, + epoch: f32, +) -> f32 { + lr_opt_get_lr(lr0, lr_min, decay_epochs, scale_epoch, epoch) +} + +#[no_mangle] +pub extern "C" fn llama_common_log_level_to_verbosity_rust(level: c_int) -> c_int { + match level { + 1 => 4, + 2 => 3, + 3 => 2, + 4 => 1, + 5 => 3, + _ => 0, + } +} + +#[no_mangle] +pub extern "C" fn llama_common_time_us_rust() -> i64 { + time_us() +} + +#[no_mangle] +pub extern "C" fn llama_common_is_http_status_ok_rust(status: c_int) -> bool { + (200..400).contains(&status) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_gguf_filename_is_model_rust( + data: *const u8, + len: usize, +) -> bool { + if data.is_null() && len != 0 { + return false; + } + let data = if len == 0 { + &[] + } else { + slice::from_raw_parts(data, len) + }; + gguf_filename_is_model(data) +} + +#[no_mangle] +pub extern "C" fn llama_common_is_hex_digit_rust(c: u8) -> bool { + is_hex_digit(c) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_parse_hex_escape_rust( + data: *const u8, + len: usize, + pos: usize, + hex_count: c_int, + out_value: *mut u32, +) -> usize { + if data.is_null() && len != 0 { + return 0; + } + if out_value.is_null() { + return 0; + } + let data = if len == 0 { + &[] + } else { + slice::from_raw_parts(data, len) + }; + let Some((value, consumed)) = parse_hex_escape(data, pos, hex_count) else { + return 0; + }; + *out_value = value; + consumed +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_parse_char_class_char_rust( + data: *const u8, + len: usize, + pos: usize, + out_value: *mut u32, +) -> usize { + if data.is_null() && len != 0 { + return 0; + } + if out_value.is_null() { + return 0; + } + let data = if len == 0 { + &[] + } else { + slice::from_raw_parts(data, len) + }; + let Some((value, consumed)) = parse_char_class_char(data, pos) else { + return 0; + }; + *out_value = value; + consumed +} + +#[no_mangle] +pub extern "C" fn llama_common_chat_format_name_rust(format: c_int) -> *const c_char { + match format { + 0 => c"Content-only".as_ptr(), + 1 => c"peg-simple".as_ptr(), + 2 => c"peg-native".as_ptr(), + 3 => c"peg-gemma4".as_ptr(), + _ => std::ptr::null(), + } +} + +#[no_mangle] +pub extern "C" fn llama_common_reasoning_format_name_rust(format: c_int) -> *const c_char { + match format { + 0 => c"none".as_ptr(), + 1 => c"auto".as_ptr(), + 2 => c"deepseek-legacy".as_ptr(), + 3 => c"deepseek".as_ptr(), + _ => std::ptr::null(), + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_reasoning_format_from_name_rust( + format: *const c_char, +) -> c_int { + if format.is_null() { + return -1; + } + match CStr::from_ptr(format).to_bytes() { + b"none" => 0, + b"auto" => 1, + b"deepseek-legacy" => 2, + b"deepseek" => 3, + _ => -1, + } +} + +#[no_mangle] +pub extern "C" fn llama_common_sampler_type_to_chr_rust(sampler_type: c_int) -> c_char { + match sampler_type { + 1 => b'd' as c_char, + 2 => b'k' as c_char, + 3 => b'p' as c_char, + 4 => b'm' as c_char, + 6 => b'y' as c_char, + 7 => b't' as c_char, + 8 => b'x' as c_char, + 9 => b'i' as c_char, + 10 => b'e' as c_char, + 11 => b's' as c_char, + 12 => b'a' as c_char, + _ => b'?' as c_char, + } +} + +#[no_mangle] +pub extern "C" fn llama_common_sampler_type_to_str_rust(sampler_type: c_int) -> *const c_char { + match sampler_type { + 1 => c"dry".as_ptr(), + 2 => c"top_k".as_ptr(), + 3 => c"top_p".as_ptr(), + 4 => c"min_p".as_ptr(), + 6 => c"typ_p".as_ptr(), + 7 => c"temperature".as_ptr(), + 8 => c"xtc".as_ptr(), + 9 => c"infill".as_ptr(), + 10 => c"penalties".as_ptr(), + 11 => c"top_n_sigma".as_ptr(), + 12 => c"adaptive_p".as_ptr(), + _ => std::ptr::null(), + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_sampler_type_from_name_rust( + name: *const c_char, + allow_alt_names: bool, +) -> c_int { + if name.is_null() { + return -1; + } + sampler_type_from_name(CStr::from_ptr(name).to_bytes(), allow_alt_names) +} + +#[no_mangle] +pub extern "C" fn llama_common_sampler_type_from_chr_rust(name: c_char) -> c_int { + sampler_type_from_chr(name as u8) +} + +#[no_mangle] +pub extern "C" fn llama_common_speculative_type_to_str_rust( + speculative_type: c_int, +) -> *const c_char { + match speculative_type { + 0 => c"none".as_ptr(), + 1 => c"draft".as_ptr(), + 2 => c"eagle3".as_ptr(), + 3 => c"ngram_simple".as_ptr(), + 4 => c"ngram_map_k".as_ptr(), + 5 => c"ngram_map_k4v".as_ptr(), + 6 => c"ngram_mod".as_ptr(), + 7 => c"ngram_cache".as_ptr(), + _ => std::ptr::null(), + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_speculative_type_from_name_rust( + name: *const c_char, +) -> c_int { + if name.is_null() { + return -1; + } + match CStr::from_ptr(name).to_bytes() { + b"none" => 0, + b"draft" => 1, + b"eagle3" => 2, + b"ngram_simple" => 3, + b"ngram_map_k" => 4, + b"ngram_map_k4v" => 5, + b"ngram_mod" => 6, + b"ngram_cache" => 7, + _ => -1, + } +} + +#[no_mangle] +pub extern "C" fn llama_common_peg_parse_result_type_name_rust( + result_type: c_int, +) -> *const c_char { + match result_type { + 0 => c"fail".as_ptr(), + 1 => c"success".as_ptr(), + 2 => c"need_more_input".as_ptr(), + _ => c"unknown".as_ptr(), + } +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_peg_rule_name_rust( + data: *const u8, + len: usize, +) -> UnicodeString { + let Some(data) = byte_slice(data, len) else { + return into_ffi(Vec::new()); + }; + peg_rule_name(data) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_gbnf_format_literal_rust( + data: *const u8, + len: usize, +) -> UnicodeString { + let Some(data) = byte_slice(data, len) else { + return into_ffi(b"\"\"".to_vec()); + }; + gbnf_format_literal(data) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_gbnf_build_repetition_rust( + item_rule: *const u8, + item_rule_len: usize, + min_items: c_int, + max_items: c_int, + separator_rule: *const u8, + separator_rule_len: usize, +) -> UnicodeString { + let Some(item_rule) = byte_slice(item_rule, item_rule_len) else { + return into_ffi(Vec::new()); + }; + let Some(separator_rule) = byte_slice(separator_rule, separator_rule_len) else { + return into_ffi(Vec::new()); + }; + gbnf_build_repetition(item_rule, min_items, max_items, separator_rule) +} + +#[no_mangle] +pub extern "C" fn llama_common_gbnf_escape_char_class_rust(c: u32) -> UnicodeString { + gbnf_escape_char_class(c) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_gbnf_is_reserved_name_rust( + data: *const u8, + len: usize, +) -> bool { + byte_slice(data, len) + .map(gbnf_is_reserved_name) + .unwrap_or(false) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_gbnf_ref_name_rust( + data: *const u8, + len: usize, +) -> UnicodeString { + let Some(data) = byte_slice(data, len) else { + return into_ffi(b"ref".to_vec()); + }; + gbnf_ref_name(data) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_hf_folder_name_to_repo_rust( + data: *const u8, + len: usize, +) -> UnicodeString { + let Some(data) = byte_slice(data, len) else { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + }; + hf_folder_name_to_repo(data) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_hf_repo_to_folder_name_rust( + data: *const u8, + len: usize, +) -> UnicodeString { + let Some(data) = byte_slice(data, len) else { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + }; + hf_repo_to_folder_name(data) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_hf_make_old_cache_filename_rust( + owner: *const u8, + owner_len: usize, + repo: *const u8, + repo_len: usize, + filename: *const u8, + filename_len: usize, +) -> UnicodeString { + let Some((owner, repo)) = byte_pair(owner, owner_len, repo, repo_len) else { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + }; + let Some(filename) = byte_slice(filename, filename_len) else { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + }; + hf_make_old_cache_filename(owner, repo, filename) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_json_brace_depth_rust(data: *const u8, len: usize) -> c_int { + let Some(data) = byte_slice(data, len) else { + return 0; + }; + json_brace_depth(data) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_hf_is_valid_repo_id_rust( + data: *const u8, + len: usize, +) -> bool { + byte_slice(data, len) + .map(hf_is_valid_repo_id) + .unwrap_or(false) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_hf_is_valid_token_rust(data: *const u8, len: usize) -> bool { + byte_slice(data, len) + .map(hf_is_valid_token) + .unwrap_or(false) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_hf_is_valid_commit_rust(data: *const u8, len: usize) -> bool { + byte_slice(data, len) + .map(|data| is_hex_string(data, 40)) + .unwrap_or(false) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_hf_is_valid_oid_rust(data: *const u8, len: usize) -> bool { + byte_slice(data, len) + .map(|data| is_hex_string(data, 40) || is_hex_string(data, 64)) + .unwrap_or(false) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_hf_parse_manifest_name_rust( + data: *const u8, + len: usize, +) -> StringList { + let Some(data) = byte_slice(data, len) else { + return StringList { + data: std::ptr::null_mut(), + len: 0, + }; + }; + hf_parse_manifest_name(data) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_gguf_split_info_rust( + data: *const u8, + len: usize, + extract_tag: bool, +) -> GgufSplitInfo { + let Some(data) = byte_slice(data, len) else { + return empty_gguf_split_info(); + }; + gguf_split_info(data, extract_tag) +} + +#[no_mangle] +pub unsafe extern "C" fn llama_common_gguf_extract_quant_bits_rust( + data: *const u8, + len: usize, +) -> c_int { + let Some(data) = byte_slice(data, len) else { + return 0; + }; + gguf_extract_quant_bits(data) +} + +unsafe fn matches_cstr(value: *const c_char, expected: &[&[u8]]) -> bool { + if value.is_null() { + return false; + } + let value = CStr::from_ptr(value).to_bytes(); + expected.iter().any(|item| value == *item) +} + +unsafe fn ascii_case_equal(a: *const c_char, b: *const c_char) -> bool { + if a.is_null() || b.is_null() { + return false; + } + + let a = CStr::from_ptr(a).to_bytes(); + let b = CStr::from_ptr(b).to_bytes(); + if a.len() != b.len() { + return false; + } + + a.iter() + .zip(b.iter()) + .all(|(a, b)| a.to_ascii_lowercase() == b.to_ascii_lowercase()) +} + +fn parse_utf8_codepoint(input: &[u8], offset: usize) -> Utf8ParseResult { + if offset >= input.len() { + return result(INCOMPLETE, 0, 0); + } + + let first = input[offset]; + if (first & 0x80) == 0 { + return result(SUCCESS, first as u32, 1); + } + + if (first & 0x40) == 0 { + return result(INVALID, 0, 0); + } + + if (first & 0x20) == 0 { + if offset + 1 >= input.len() { + return result(INCOMPLETE, 0, 0); + } + if (input[offset + 1] & 0xC0) != 0x80 { + return result(INVALID, 0, 0); + } + let codepoint = (((first & 0x1F) as u32) << 6) | ((input[offset + 1] & 0x3F) as u32); + return result(SUCCESS, codepoint, 2); + } + + if (first & 0x10) == 0 { + if offset + 2 >= input.len() { + return result(INCOMPLETE, 0, 0); + } + if (input[offset + 1] & 0xC0) != 0x80 || (input[offset + 2] & 0xC0) != 0x80 { + return result(INVALID, 0, 0); + } + let codepoint = (((first & 0x0F) as u32) << 12) + | (((input[offset + 1] & 0x3F) as u32) << 6) + | ((input[offset + 2] & 0x3F) as u32); + return result(SUCCESS, codepoint, 3); + } + + if (first & 0x08) == 0 { + if offset + 3 >= input.len() { + return result(INCOMPLETE, 0, 0); + } + if (input[offset + 1] & 0xC0) != 0x80 + || (input[offset + 2] & 0xC0) != 0x80 + || (input[offset + 3] & 0xC0) != 0x80 + { + return result(INVALID, 0, 0); + } + let codepoint = (((first & 0x07) as u32) << 18) + | (((input[offset + 1] & 0x3F) as u32) << 12) + | (((input[offset + 2] & 0x3F) as u32) << 6) + | ((input[offset + 3] & 0x3F) as u32); + return result(SUCCESS, codepoint, 4); + } + + result(INVALID, 0, 0) +} + +fn utf8_is_complete(bytes: &[u8]) -> bool { + if bytes.is_empty() { + return true; + } + for i in 1..=bytes.len().min(4) { + let c = bytes[bytes.len() - i]; + if (c & 0xC0) != 0x80 { + let expected = if c >= 0xF0 { + 4 + } else if c >= 0xE0 { + 3 + } else if c >= 0xC0 { + 2 + } else { + 1 + }; + return i >= expected; + } + } + false +} + +fn unicode_cpt_to_utf8(cpt: u32) -> Option> { + let mut result = Vec::with_capacity(4); + + if cpt <= 0x7F { + result.push(cpt as u8); + } else if cpt <= 0x7FF { + result.push(0xC0 | ((cpt >> 6) & 0x1F) as u8); + result.push(0x80 | (cpt & 0x3F) as u8); + } else if cpt <= 0xFFFF { + result.push(0xE0 | ((cpt >> 12) & 0x0F) as u8); + result.push(0x80 | ((cpt >> 6) & 0x3F) as u8); + result.push(0x80 | (cpt & 0x3F) as u8); + } else if cpt <= 0x10FFFF { + result.push(0xF0 | ((cpt >> 18) & 0x07) as u8); + result.push(0x80 | ((cpt >> 12) & 0x3F) as u8); + result.push(0x80 | ((cpt >> 6) & 0x3F) as u8); + result.push(0x80 | (cpt & 0x3F) as u8); + } else { + return None; + } + + Some(result) +} + +fn unicode_cpts_to_utf8(cpts: &[u32]) -> UnicodeString { + let mut result = Vec::with_capacity(cpts.len()); + for &cpt in cpts { + let Some(bytes) = unicode_cpt_to_utf8(cpt) else { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + }; + result.extend_from_slice(&bytes); + } + into_ffi(result) +} + +fn http_show_masked_url(scheme: &[u8], has_user: bool, host: &[u8], path: &[u8]) -> UnicodeString { + let mut result = Vec::with_capacity( + scheme.len() + 3 + usize::from(has_user) * b"****:****@".len() + host.len() + path.len(), + ); + result.extend_from_slice(scheme); + result.extend_from_slice(b"://"); + if has_user { + result.extend_from_slice(b"****:****@"); + } + result.extend_from_slice(host); + result.extend_from_slice(path); + into_ffi(result) +} + +fn peak_source(source: &[u8], pos: usize, max_peak_chars: usize) -> UnicodeString { + if source.is_empty() { + return into_ffi(b"(no source available)".to_vec()); + } + + let start = pos.saturating_sub(max_peak_chars); + let end = source.len().min(pos.saturating_add(max_peak_chars)); + let snippet = if start <= end && start <= source.len() { + &source[start..end] + } else { + &[] + }; + + let mut output = Vec::with_capacity(snippet.len() + 16 + pos.saturating_sub(start)); + output.extend_from_slice(b"..."); + for &c in snippet { + if c == b'\n' { + output.extend_from_slice("↵".as_bytes()); + } else { + output.push(c); + } + } + output.extend_from_slice(b"...\n"); + output.resize(output.len() + pos.saturating_sub(start) + 3, b' '); + output.push(b'^'); + into_ffi(output) +} + +fn fmt_error_with_source(tag: &[u8], msg: &[u8], source: &[u8], pos: usize) -> UnicodeString { + let peak = peak_source(source, pos, 40); + let peak = unsafe { Vec::from_raw_parts(peak.data, peak.len, peak.len) }; + + let mut result = Vec::with_capacity(tag.len() + msg.len() + peak.len() + 3); + result.extend_from_slice(tag); + result.extend_from_slice(b": "); + result.extend_from_slice(msg); + result.push(b'\n'); + result.extend_from_slice(&peak); + into_ffi(result) +} + +fn string_repeat(input: &[u8], n: usize) -> UnicodeString { + if input.is_empty() || n == 0 { + return into_ffi(Vec::new()); + } + + let Some(capacity) = input.len().checked_mul(n) else { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + }; + let mut result = Vec::with_capacity(capacity); + for _ in 0..n { + result.extend_from_slice(input); + } + into_ffi(result) +} + +fn int_vec_to_string(values: &[i32]) -> UnicodeString { + let mut result = String::from("[ "); + for (idx, value) in values.iter().enumerate() { + if idx > 0 { + result.push_str(", "); + } + result.push_str(&value.to_string()); + } + result.push_str(" ]"); + into_ffi(result.into_bytes()) +} + +fn string_trim(input: &[u8], mode: c_int) -> UnicodeString { + let mut start = 0; + let mut end = input.len(); + + if mode == 0 || mode == 1 { + while start < end && input[start].is_ascii_whitespace() { + start += 1; + } + } + + if mode == 0 || mode == 2 { + while end > start && input[end - 1].is_ascii_whitespace() { + end -= 1; + } + } else if mode == 3 { + while end > start && input[end - 1] == b'\n' { + end -= 1; + } + } + + into_ffi(input[start..end].to_vec()) +} + +fn byte_pair<'a>( + left: *const u8, + left_len: usize, + right: *const u8, + right_len: usize, +) -> Option<(&'a [u8], &'a [u8])> { + if (left.is_null() && left_len != 0) || (right.is_null() && right_len != 0) { + return None; + } + let left = if left_len == 0 { + &[] + } else { + unsafe { slice::from_raw_parts(left, left_len) } + }; + let right = if right_len == 0 { + &[] + } else { + unsafe { slice::from_raw_parts(right, right_len) } + }; + Some((left, right)) +} + +fn byte_slice<'a>(data: *const u8, len: usize) -> Option<&'a [u8]> { + if data.is_null() && len != 0 { + return None; + } + if len == 0 { + Some(&[]) + } else { + Some(unsafe { slice::from_raw_parts(data, len) }) + } +} + +fn common_prefix_len(left: &[u8], right: &[u8]) -> usize { + left.iter() + .zip(right.iter()) + .take_while(|(left, right)| left == right) + .count() +} + +fn common_suffix_len(left: &[u8], right: &[u8]) -> usize { + left.iter() + .rev() + .zip(right.iter().rev()) + .take_while(|(left, right)| left == right) + .count() +} + +fn string_replace_all(input: &[u8], search: &[u8], replace: &[u8]) -> UnicodeString { + if search.is_empty() { + return into_ffi(input.to_vec()); + } + + let mut result = Vec::with_capacity(input.len()); + let mut pos = 0; + while let Some(found) = find_subslice(&input[pos..], search) { + let absolute = pos + found; + result.extend_from_slice(&input[pos..absolute]); + result.extend_from_slice(replace); + pos = absolute + search.len(); + } + result.extend_from_slice(&input[pos..]); + into_ffi(result) +} + +fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option { + if needle.is_empty() { + return Some(0); + } + haystack + .windows(needle.len()) + .position(|window| window == needle) +} + +fn regex_escape(input: &[u8]) -> UnicodeString { + let mut result = Vec::with_capacity(input.len()); + for &byte in input { + if matches!( + byte, + b'.' | b'^' + | b'$' + | b'|' + | b'(' + | b')' + | b'*' + | b'+' + | b'?' + | b'[' + | b']' + | b'{' + | b'}' + | b'\\' + ) { + result.push(b'\\'); + } + result.push(byte); + } + into_ffi(result) +} + +fn string_join(values: &[StringView], separator: &[u8]) -> UnicodeString { + let mut result = Vec::new(); + for (idx, value) in values.iter().enumerate() { + if idx > 0 { + result.extend_from_slice(separator); + } + if value.data.is_null() && value.len != 0 { + return UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }; + } + if value.len != 0 { + let bytes = unsafe { slice::from_raw_parts(value.data, value.len) }; + result.extend_from_slice(bytes); + } + } + into_ffi(result) +} + +fn string_split(input: &[u8], delimiter: &[u8]) -> StringList { + if delimiter.is_empty() { + return into_string_list(vec![input.to_vec()]); + } + + let mut parts = Vec::new(); + let mut start = 0; + while let Some(found) = find_subslice(&input[start..], delimiter) { + let end = start + found; + parts.push(input[start..end].to_vec()); + start = end + delimiter.len(); + } + parts.push(input[start..].to_vec()); + into_string_list(parts) +} + +fn parse_csv_row(input: &[u8]) -> StringList { + let mut fields = Vec::new(); + let mut field = Vec::new(); + let mut in_quotes = false; + let mut idx = 0; + + while idx < input.len() { + let ch = input[idx]; + if ch == b'"' { + if !in_quotes { + if field.is_empty() { + in_quotes = true; + } else { + field.push(b'"'); + } + } else if idx + 1 < input.len() && input[idx + 1] == b'"' { + field.push(b'"'); + idx += 1; + } else { + in_quotes = false; + } + } else if ch == b',' { + if in_quotes { + field.push(b','); + } else { + fields.push(std::mem::take(&mut field)); + } + } else { + field.push(ch); + } + idx += 1; + } + + fields.push(field); + into_string_list(fields) +} + +fn trim_trailing_space_view(input: &[u8], max: c_int) -> StringView { + let mut end = input.len(); + let mut count = 0_i32; + while end > 0 && input[end - 1].is_ascii_whitespace() { + if max != -1 && count >= max { + break; + } + end -= 1; + count += 1; + } + StringView { + data: input.as_ptr(), + len: end, + } +} + +fn trim_leading_space_view(input: &[u8], max: c_int) -> StringView { + let mut start = 0_usize; + let mut count = 0_i32; + while start < input.len() && input[start].is_ascii_whitespace() { + if max != -1 && count >= max { + break; + } + start += 1; + count += 1; + } + StringView { + data: unsafe { input.as_ptr().add(start) }, + len: input.len() - start, + } +} + +fn string_process_escapes(input: &[u8]) -> UnicodeString { + let mut output = Vec::with_capacity(input.len()); + let mut idx = 0; + while idx < input.len() { + if input[idx] == b'\\' && idx + 1 < input.len() { + idx += 1; + match input[idx] { + b'n' => output.push(b'\n'), + b'r' => output.push(b'\r'), + b't' => output.push(b'\t'), + b'\'' => output.push(b'\''), + b'"' => output.push(b'"'), + b'\\' => output.push(b'\\'), + b'x' => { + if idx + 2 < input.len() { + if let Some(value) = parse_hex_byte(input[idx + 1], input[idx + 2]) { + idx += 2; + output.push(value); + } else { + output.push(b'\\'); + output.push(input[idx]); + } + } else { + output.push(b'\\'); + output.push(input[idx]); + } + } + other => { + output.push(b'\\'); + output.push(other); + } + } + } else { + output.push(input[idx]); + } + idx += 1; + } + into_ffi(output) +} + +fn clean_file_name(input: &[u8]) -> UnicodeString { + let mut output = input.to_vec(); + for byte in &mut output { + if *byte == b'/' || *byte == b'\\' { + *byte = b'_'; + } + } + into_ffi(output) +} + +fn string_diff(last: &[u8], current: &[u8]) -> (UnicodeString, c_int) { + if last.is_empty() { + return (into_ffi(current.to_vec()), 0); + } + if !current.starts_with(last) { + if last.starts_with(current) { + return (into_ffi(Vec::new()), 0); + } + return ( + UnicodeString { + data: std::ptr::null_mut(), + len: 0, + }, + -1, + ); + } + (into_ffi(current[last.len()..].to_vec()), 0) +} + +fn c_string_bytes(input: &[u8]) -> &[u8] { + match input.iter().position(|&byte| byte == 0) { + Some(end) => &input[..end], + None => input, + } +} + +fn glob_match(pattern: &[u8], str: &[u8]) -> bool { + if pattern.is_empty() { + return str.is_empty(); + } + + if pattern.len() >= 2 && pattern[0] == b'*' && pattern[1] == b'*' { + let rest = &pattern[2..]; + if glob_match(rest, str) { + return true; + } + if !str.is_empty() { + return glob_match(pattern, &str[1..]); + } + return false; + } + + if pattern[0] == b'*' { + let rest = &pattern[1..]; + let mut idx = 0; + while idx < str.len() && str[idx] != b'/' { + if glob_match(rest, &str[idx..]) { + return true; + } + idx += 1; + } + return glob_match(rest, &str[idx..]); + } + + if pattern[0] == b'?' && !str.is_empty() && str[0] != b'/' { + return glob_match(&pattern[1..], &str[1..]); + } + + if pattern[0] == b'[' { + let mut class_end = 1; + if class_end < pattern.len() && (pattern[class_end] == b']' || pattern[class_end] == b'-') { + class_end += 1; + } + while class_end < pattern.len() && pattern[class_end] != b']' { + class_end += 1; + } + if class_end < pattern.len() && pattern[class_end] == b']' { + if str.is_empty() { + return false; + } + return glob_class_match(str[0], &pattern[1..class_end]) + && glob_match(&pattern[class_end + 1..], &str[1..]); + } + if !str.is_empty() && str[0] == b'[' { + return glob_match(&pattern[1..], &str[1..]); + } + return false; + } + + if !str.is_empty() && pattern[0] == str[0] { + return glob_match(&pattern[1..], &str[1..]); + } + + false +} + +fn glob_class_match(c: u8, class: &[u8]) -> bool { + let mut idx = 0; + let mut negated = false; + + if idx < class.len() && class[idx] == b'!' { + negated = true; + idx += 1; + } + + if idx < class.len() && (class[idx] == b']' || class[idx] == b'-') { + if class[idx] == c { + return !negated; + } + idx += 1; + } + + let mut matched = false; + while idx < class.len() { + if idx + 2 < class.len() && class[idx + 1] == b'-' && class[idx + 2] != b']' { + if c >= class[idx] && c <= class[idx + 2] { + matched = true; + break; + } + idx += 3; + } else { + if class[idx] == c { + matched = true; + break; + } + idx += 1; + } + } + + if negated { + !matched + } else { + matched + } +} + +fn fs_validate_filename(filename: &[u8], allow_subdirs: bool) -> bool { + if filename.is_empty() || filename.len() > 255 { + return false; + } + + let mut offset = 0; + while offset < filename.len() { + let parsed = parse_utf8_codepoint(filename, offset); + if parsed.status != SUCCESS { + return false; + } + + let c = parsed.codepoint; + if (parsed.bytes_consumed == 2 && c < 0x80) + || (parsed.bytes_consumed == 3 && c < 0x800) + || (parsed.bytes_consumed == 4 && c < 0x10000) + { + return false; + } + + if c <= 0x1F + || c == 0x7F + || (0x80..=0x9F).contains(&c) + || c == 0xFF0E + || c == 0x2215 + || c == 0x2216 + || (0xD800..=0xDFFF).contains(&c) + || c > 0x10FFFF + || c == 0xFFFD + || c == 0xFEFF + || [b':', b'*', b'?', b'"', b'<', b'>', b'|'] + .iter() + .any(|&invalid| c == invalid as u32) + { + return false; + } + if !allow_subdirs && (c == b'/' as u32 || c == b'\\' as u32) { + return false; + } + + offset += parsed.bytes_consumed; + } + + if filename.first() == Some(&b' ') + || filename.last() == Some(&b' ') + || filename.last() == Some(&b'.') + { + return false; + } + if find_subslice(filename, b"..").is_some() { + return false; + } + if filename == b"." { + return false; + } + + true +} + +fn string_lstrip_chars(input: &[u8], chars: &[u8]) -> UnicodeString { + let start = input + .iter() + .position(|byte| !chars.contains(byte)) + .unwrap_or(input.len()); + into_ffi(input[start..].to_vec()) +} + +fn string_rstrip_chars(input: &[u8], chars: &[u8]) -> UnicodeString { + let end = input + .iter() + .rposition(|byte| !chars.contains(byte)) + .map(|idx| idx + 1) + .unwrap_or(0); + into_ffi(input[..end].to_vec()) +} + +fn break_str_into_lines(input: &[u8], max_char_per_line: usize) -> StringList { + if input.is_empty() { + return into_string_list(Vec::new()); + } + + let mut result = Vec::new(); + for line in input.split(|&byte| byte == b'\n') { + add_wrapped_line(line, max_char_per_line, &mut result); + } + + if input.last() == Some(&b'\n') { + result.pop(); + } + + into_string_list(result) +} + +fn add_wrapped_line(line: &[u8], max_char_per_line: usize, result: &mut Vec>) { + if line.len() <= max_char_per_line { + result.push(line.to_vec()); + return; + } + + let mut current_line = Vec::new(); + for word in line + .split(u8::is_ascii_whitespace) + .filter(|word| !word.is_empty()) + { + let separator_len = usize::from(!current_line.is_empty()); + if current_line.len() + separator_len + word.len() > max_char_per_line { + if !current_line.is_empty() { + result.push(std::mem::take(&mut current_line)); + } + current_line.extend_from_slice(word); + } else { + if !current_line.is_empty() { + current_line.push(b' '); + } + current_line.extend_from_slice(word); + } + } + if !current_line.is_empty() { + result.push(current_line); + } +} + +fn rm_leading_dashes(input: &[u8]) -> UnicodeString { + let start = input.iter().take_while(|&&byte| byte == b'-').count(); + into_ffi(input[start..].to_vec()) +} + +fn until_common_prefix(full: &[u8], left: &[u8], right: &[u8]) -> UnicodeString { + let prefix_len = common_prefix_len(left, right); + if prefix_len == 0 { + return into_ffi(Vec::new()); + } + + match find_subslice(full, &left[..prefix_len]) { + Some(pos) => into_ffi(full[..pos].to_vec()), + None => into_ffi(Vec::new()), + } +} + +fn after_common_suffix(full: &[u8], left: &[u8], right: &[u8]) -> UnicodeString { + let suffix_len = common_suffix_len(left, right); + if suffix_len == 0 { + return into_ffi(Vec::new()); + } + + let suffix = &left[left.len() - suffix_len..]; + match rfind_subslice(full, suffix) { + Some(pos) => into_ffi(full[pos + suffix_len..].to_vec()), + None => into_ffi(Vec::new()), + } +} + +fn rfind_subslice(haystack: &[u8], needle: &[u8]) -> Option { + if needle.is_empty() { + return Some(haystack.len()); + } + haystack + .windows(needle.len()) + .rposition(|window| window == needle) +} + +fn json_ensure_ascii_preserving_format(input: &[u8]) -> UnicodeString { + let mut output = Vec::with_capacity(input.len()); + let mut in_string = false; + let mut escaped = false; + let mut pos = 0; + + while pos < input.len() { + let ch = input[pos]; + if !in_string { + output.push(ch); + if ch == b'"' { + in_string = true; + } + pos += 1; + continue; + } + + if escaped { + output.push(ch); + escaped = false; + pos += 1; + continue; + } + + if ch == b'\\' { + output.push(ch); + escaped = true; + pos += 1; + continue; + } + + if ch == b'"' { + output.push(ch); + in_string = false; + pos += 1; + continue; + } + + if ch < 0x80 { + output.push(ch); + pos += 1; + continue; + } + + let parsed = parse_utf8_codepoint(input, pos); + if parsed.status != SUCCESS { + output.extend_from_slice(b"\\ufffd"); + pos += 1; + continue; + } + + append_codepoint_as_ascii_json_escape(&mut output, parsed.codepoint); + pos += parsed.bytes_consumed; + } + + into_ffi(output) +} + +fn append_codepoint_as_ascii_json_escape(output: &mut Vec, mut codepoint: u32) { + if codepoint <= 0xFFFF { + output.extend_from_slice(format!("\\u{codepoint:04x}").as_bytes()); + return; + } + + codepoint -= 0x10000; + let high = 0xD800 + ((codepoint >> 10) & 0x3FF); + let low = 0xDC00 + (codepoint & 0x3FF); + output.extend_from_slice(format!("\\u{high:04x}\\u{low:04x}").as_bytes()); +} + +fn normalize_quotes_to_json(input: &[u8]) -> UnicodeString { + let mut result = Vec::with_capacity(input.len() + 16); + let mut in_single_quoted = false; + let mut in_double_quoted = false; + let mut idx = 0_usize; + + while idx < input.len() { + let c = input[idx]; + + if c == b'\\' && idx + 1 < input.len() { + let next = input[idx + 1]; + + if in_single_quoted { + if next == b'\'' { + result.push(b'\''); + idx += 2; + continue; + } + if next == b'"' { + result.extend_from_slice(b"\\\""); + idx += 2; + continue; + } + result.push(c); + result.push(next); + idx += 2; + continue; + } + + if in_double_quoted { + result.push(c); + result.push(next); + idx += 2; + continue; + } + + result.push(c); + idx += 1; + continue; + } + + if c == b'"' { + if in_single_quoted { + result.extend_from_slice(b"\\\""); + } else { + in_double_quoted = !in_double_quoted; + result.push(c); + } + } else if c == b'\'' { + if in_double_quoted { + result.push(c); + } else if in_single_quoted { + in_single_quoted = false; + result.push(b'"'); + } else { + in_single_quoted = true; + result.push(b'"'); + } + } else { + result.push(c); + } + + idx += 1; + } + + into_ffi(result) +} + +fn parse_cpu_mask(mask: &[u8], boolmask: &mut [bool]) -> c_int { + let start = if mask.len() >= 2 && &mask[..2] == b"0x" { + 2 + } else { + 0 + }; + let num_digits = (mask.len() - start).min(128); + let end = start + num_digits; + + for (offset, idx) in (start..end).enumerate() { + let Some(id) = hex_value(mask[idx]) else { + return idx as c_int; + }; + let n = num_digits * 4 - 1 - offset * 4; + if n < boolmask.len() { + boolmask[n] |= (id & 8) != 0; + } + if n >= 1 && n - 1 < boolmask.len() { + boolmask[n - 1] |= (id & 4) != 0; + } + if n >= 2 && n - 2 < boolmask.len() { + boolmask[n - 2] |= (id & 2) != 0; + } + if n >= 3 && n - 3 < boolmask.len() { + boolmask[n - 3] |= (id & 1) != 0; + } + } + + -1 +} + +fn get_line_col(source: &[u8], pos: usize) -> UnicodeString { + let mut line = 1_usize; + let mut col = 1_usize; + for &c in source.iter().take(pos) { + if c == b'\n' { + line += 1; + col = 1; + } else { + col += 1; + } + } + into_ffi(format!("line {line}, column {col}").into_bytes()) +} + +fn lr_opt_init(lr0: f32, lr_min: f32, mut decay_epochs: f32, epochs: u32) -> (f32, f32) { + let mut scale_epoch = 0.0; + if lr_min > 0.0 && lr_min < lr0 { + let nhalf = (lr0 / lr_min).ln() / 2.0_f32.ln(); + let mut e = epochs as f32; + if decay_epochs > 0.0 && decay_epochs < e { + e = decay_epochs; + } else { + decay_epochs = e; + } + scale_epoch = nhalf / e; + } + (decay_epochs, scale_epoch) +} + +fn lr_opt_get_lr(lr0: f32, lr_min: f32, decay_epochs: f32, scale_epoch: f32, epoch: f32) -> f32 { + if lr_min <= 0.0 { + lr0 + } else if epoch >= decay_epochs { + lr_min + } else { + lr0 * 0.5_f32.powf(epoch * scale_epoch) + } +} + +fn time_us() -> i64 { + match SystemTime::now().duration_since(UNIX_EPOCH) { + Ok(duration) => duration.as_micros().min(i64::MAX as u128) as i64, + Err(err) => -(err.duration().as_micros().min(i64::MAX as u128) as i64), + } +} + +fn gguf_filename_is_model(filepath: &[u8]) -> bool { + if !filepath.ends_with(b".gguf") { + return false; + } + + let filename = match filepath.iter().rposition(|&byte| byte == b'/') { + Some(pos) => &filepath[pos + 1..], + None => filepath, + }; + + find_subslice(filename, b"mmproj").is_none() && find_subslice(filename, b"imatrix").is_none() +} + +fn peg_rule_name(name: &[u8]) -> UnicodeString { + let mut result = Vec::with_capacity(name.len()); + let mut in_invalid_run = false; + + for &c in name { + if c.is_ascii_alphanumeric() || c == b'-' { + result.push(c); + in_invalid_run = false; + } else if !in_invalid_run { + result.push(b'-'); + in_invalid_run = true; + } + } + + into_ffi(result) +} + +fn gbnf_format_literal(literal: &[u8]) -> UnicodeString { + let mut result = Vec::with_capacity(literal.len() + 2); + result.push(b'"'); + for &c in literal { + match c { + b'\r' => result.extend_from_slice(b"\\r"), + b'\n' => result.extend_from_slice(b"\\n"), + b'"' => result.extend_from_slice(b"\\\""), + b'\\' => result.extend_from_slice(b"\\\\"), + _ => result.push(c), + } + } + result.push(b'"'); + into_ffi(result) +} + +fn gbnf_build_repetition( + item_rule: &[u8], + min_items: c_int, + max_items: c_int, + separator_rule: &[u8], +) -> UnicodeString { + const INT_MAX: c_int = c_int::MAX; + let has_max = max_items != INT_MAX; + + if max_items == 0 { + return into_ffi(Vec::new()); + } + if min_items == 0 && max_items == 1 { + let mut result = Vec::with_capacity(item_rule.len() + 1); + result.extend_from_slice(item_rule); + result.push(b'?'); + return into_ffi(result); + } + + if separator_rule.is_empty() { + let mut result = Vec::new(); + result.extend_from_slice(item_rule); + if min_items == 1 && !has_max { + result.push(b'+'); + } else if min_items == 0 && !has_max { + result.push(b'*'); + } else { + result.push(b'{'); + result.extend_from_slice(min_items.to_string().as_bytes()); + result.push(b','); + if has_max { + result.extend_from_slice(max_items.to_string().as_bytes()); + } + result.push(b'}'); + } + return into_ffi(result); + } + + let mut inner = Vec::with_capacity(separator_rule.len() + item_rule.len() + 3); + inner.push(b'('); + inner.extend_from_slice(separator_rule); + inner.push(b' '); + inner.extend_from_slice(item_rule); + inner.push(b')'); + let inner_min = if min_items == 0 { 0 } else { min_items - 1 }; + let inner_max = if has_max { max_items - 1 } else { max_items }; + let repeated = gbnf_build_repetition(&inner, inner_min, inner_max, &[]); + let repeated = unsafe { Vec::from_raw_parts(repeated.data, repeated.len, repeated.len) }; + + let mut result = Vec::with_capacity(item_rule.len() + repeated.len() + 6); + result.extend_from_slice(item_rule); + result.push(b' '); + result.extend_from_slice(&repeated); + if min_items == 0 { + result.insert(0, b'('); + result.push(b')'); + result.push(b'?'); + } + into_ffi(result) +} + +fn gbnf_escape_char_class(c: u32) -> UnicodeString { + match c { + 45 | 93 | 91 | 92 => into_ffi(vec![b'\\', c as u8]), + 10 => into_ffi(b"\\n".to_vec()), + 9 => into_ffi(b"\\t".to_vec()), + 13 => into_ffi(b"\\r".to_vec()), + 0x20..=0x7E => into_ffi(vec![c as u8]), + 0x00..=0xFF => into_ffi(format!("\\x{c:02X}").into_bytes()), + 0x0100..=0xFFFF => into_ffi(format!("\\u{c:04X}").into_bytes()), + _ => into_ffi(format!("\\U{c:08X}").into_bytes()), + } +} + +fn gbnf_is_reserved_name(name: &[u8]) -> bool { + matches!( + name, + b"root" + | b"boolean" + | b"decimal-part" + | b"integral-part" + | b"number" + | b"integer" + | b"value" + | b"object" + | b"array" + | b"uuid" + | b"char" + | b"string" + | b"null" + | b"date" + | b"time" + | b"date-time" + | b"date-string" + | b"time-string" + | b"date-time-string" + ) +} + +fn gbnf_ref_name(reference: &[u8]) -> UnicodeString { + let fragment = match reference.iter().position(|&c| c == b'#') { + Some(pos) => &reference[pos + 1..], + None => reference, + }; + + let mut result = b"ref".to_vec(); + let mut in_invalid_run = false; + for &c in fragment { + if c.is_ascii_alphanumeric() || c == b'-' { + result.push(c); + in_invalid_run = false; + } else if !in_invalid_run { + result.push(b'-'); + in_invalid_run = true; + } + } + into_ffi(result) +} + +fn hf_folder_name_to_repo(folder: &[u8]) -> UnicodeString { + const PREFIX: &[u8] = b"models--"; + if !folder.starts_with(PREFIX) { + return into_ffi(Vec::new()); + } + string_replace_all(&folder[PREFIX.len()..], b"--", b"/") +} + +fn hf_repo_to_folder_name(repo_id: &[u8]) -> UnicodeString { + let mut result = b"models--".to_vec(); + result.extend_from_slice(repo_id); + let replaced = string_replace_all(&result, b"/", b"--"); + replaced +} + +fn hf_make_old_cache_filename(owner: &[u8], repo: &[u8], filename: &[u8]) -> UnicodeString { + let mut result = Vec::with_capacity(owner.len() + repo.len() + filename.len() + 2); + result.extend_from_slice(owner); + result.push(b'_'); + result.extend_from_slice(repo); + result.push(b'_'); + result.extend_from_slice(filename); + string_replace_all(&result, b"/", b"_") +} + +fn json_brace_depth(input: &[u8]) -> c_int { + let mut depth = 0_i32; + let mut in_string = false; + let mut escaped = false; + + for &c in input { + if escaped { + escaped = false; + continue; + } + if c == b'\\' && in_string { + escaped = true; + continue; + } + if c == b'"' { + in_string = !in_string; + continue; + } + if !in_string { + if c == b'{' { + depth += 1; + } else if c == b'}' { + depth -= 1; + } + } + } + depth +} + +fn is_hex_string(input: &[u8], expected_len: usize) -> bool { + input.len() == expected_len && input.iter().all(|&c| is_hex_digit(c)) +} + +fn is_alphanum(c: u8) -> bool { + c.is_ascii_alphanumeric() +} + +fn hf_is_special_char(c: u8) -> bool { + matches!(c, b'/' | b'.' | b'-') +} + +fn hf_is_valid_repo_id(repo_id: &[u8]) -> bool { + if repo_id.is_empty() || repo_id.len() > 256 { + return false; + } + let mut slash = 0; + let mut special = true; + + for &c in repo_id { + if is_alphanum(c) || c == b'_' { + special = false; + } else if hf_is_special_char(c) { + if special { + return false; + } + slash += usize::from(c == b'/'); + special = true; + } else { + return false; + } + } + !special && slash == 1 +} + +fn hf_is_valid_token(token: &[u8]) -> bool { + token.len() >= 37 + && token.len() <= 256 + && token.starts_with(b"hf_") + && token[3..].iter().all(|&c| is_alphanum(c)) +} + +fn hf_parse_manifest_name(filename: &[u8]) -> StringList { + const PREFIX: &[u8] = b"manifest="; + const SUFFIX: &[u8] = b".json"; + if !filename.starts_with(PREFIX) || !filename.ends_with(SUFFIX) { + return into_string_list(Vec::new()); + } + let rest = &filename[PREFIX.len()..]; + let Some(owner_end) = rest.iter().position(|&c| c == b'=') else { + return into_string_list(Vec::new()); + }; + if owner_end == 0 { + return into_string_list(Vec::new()); + } + let owner = &rest[..owner_end]; + let rest = &rest[owner_end + 1..]; + let Some(repo_end) = rest.iter().position(|&c| c == b'=') else { + return into_string_list(Vec::new()); + }; + if repo_end == 0 { + return into_string_list(Vec::new()); + } + if repo_end + 1 >= rest.len() { + return into_string_list(Vec::new()); + } + into_string_list(vec![owner.to_vec(), rest[..repo_end].to_vec()]) +} + +fn empty_gguf_split_info() -> GgufSplitInfo { + GgufSplitInfo { + prefix: into_ffi(Vec::new()), + tag: into_ffi(Vec::new()), + index: 0, + count: 0, + } +} + +fn gguf_split_info(path: &[u8], extract_tag: bool) -> GgufSplitInfo { + let Some(mut prefix) = path.strip_suffix(b".gguf").map(Vec::from) else { + return empty_gguf_split_info(); + }; + + let mut index = 1; + let mut count = 1; + if let Some((split_prefix, split_index, split_count)) = gguf_split_suffix(&prefix) { + prefix = split_prefix; + index = split_index; + count = split_count; + } + + let tag = if extract_tag { + gguf_extract_tag(&prefix) + } else { + Vec::new() + }; + + GgufSplitInfo { + prefix: into_ffi(prefix), + tag: into_ffi(tag), + index, + count, + } +} + +fn gguf_split_suffix(prefix: &[u8]) -> Option<(Vec, c_int, c_int)> { + const MID: &[u8] = b"-of-"; + if prefix.len() < 16 { + return None; + } + let count_start = prefix.len() - 5; + let mid_start = count_start.checked_sub(MID.len())?; + let index_start = mid_start.checked_sub(5)?; + if index_start == 0 || prefix.get(index_start.wrapping_sub(1)) != Some(&b'-') { + return None; + } + if &prefix[mid_start..count_start] != MID { + return None; + } + let index = parse_fixed_decimal(&prefix[index_start..mid_start])?; + let count = parse_fixed_decimal(&prefix[count_start..])?; + Some((prefix[..index_start - 1].to_vec(), index, count)) +} + +fn parse_fixed_decimal(input: &[u8]) -> Option { + if input.is_empty() || !input.iter().all(u8::is_ascii_digit) { + return None; + } + let mut value = 0_i32; + for &c in input { + value = value.checked_mul(10)?.checked_add((c - b'0') as i32)?; + } + Some(value) +} + +fn gguf_extract_tag(prefix: &[u8]) -> Vec { + let Some(pos) = prefix.iter().rposition(|&c| c == b'-' || c == b'.') else { + return Vec::new(); + }; + let tag = &prefix[pos + 1..]; + if tag.is_empty() || !tag.iter().all(|&c| c.is_ascii_alphanumeric() || c == b'_') { + return Vec::new(); + } + tag.iter().map(u8::to_ascii_uppercase).collect() +} + +fn gguf_extract_quant_bits(filename: &[u8]) -> c_int { + let split = gguf_split_info(filename, true); + if split.tag.data.is_null() { + unsafe { llama_common_unicode_string_free(split.prefix) }; + return 0; + } + + let tag = unsafe { Vec::from_raw_parts(split.tag.data, split.tag.len, split.tag.len) }; + unsafe { llama_common_unicode_string_free(split.prefix) }; + + let Some(pos) = tag.iter().position(u8::is_ascii_digit) else { + return 0; + }; + let mut value = 0_i32; + for &c in &tag[pos..] { + if !c.is_ascii_digit() { + break; + } + let Some(next) = value + .checked_mul(10) + .and_then(|v| v.checked_add((c - b'0') as i32)) + else { + return 0; + }; + value = next; + } + value +} + +fn embd_normalize(inp: &[f32], out: &mut [f32], embd_norm: c_int) { + let mut sum = 0.0_f64; + + match embd_norm { + -1 => sum = 1.0, + 0 => { + for &value in inp { + sum = sum.max((value as f64).abs()); + } + sum /= 32760.0; + } + 2 => { + for &value in inp { + let value = value as f64; + sum += value * value; + } + sum = sum.sqrt(); + } + norm => { + for &value in inp { + sum += (value as f64).abs().powf(norm as f64); + } + sum = sum.powf(1.0 / norm as f64); + } + } + + let norm = if sum > 0.0 { (1.0 / sum) as f32 } else { 0.0 }; + for (dst, src) in out.iter_mut().zip(inp.iter()) { + *dst = *src * norm; + } +} + +fn embd_similarity_cos(embd1: &[f32], embd2: &[f32]) -> f32 { + let mut sum = 0.0_f64; + let mut sum1 = 0.0_f64; + let mut sum2 = 0.0_f64; + + for (&left, &right) in embd1.iter().zip(embd2.iter()) { + let left = left as f64; + let right = right as f64; + sum += left * right; + sum1 += left * left; + sum2 += right * right; + } + + if sum1 == 0.0 || sum2 == 0.0 { + if sum1 == 0.0 && sum2 == 0.0 { + 1.0 + } else { + 0.0 + } + } else { + (sum / (sum1.sqrt() * sum2.sqrt())) as f32 + } +} + +fn parse_hex_byte(high: u8, low: u8) -> Option { + Some(hex_value(high)? << 4 | hex_value(low)?) +} + +fn hex_value(value: u8) -> Option { + match value { + b'0'..=b'9' => Some(value - b'0'), + b'a'..=b'f' => Some(value - b'a' + 10), + b'A'..=b'F' => Some(value - b'A' + 10), + _ => None, + } +} + +fn is_hex_digit(c: u8) -> bool { + c.is_ascii_hexdigit() +} + +fn parse_hex_escape(input: &[u8], pos: usize, hex_count: c_int) -> Option<(u32, usize)> { + if hex_count < 0 { + return None; + } + let hex_count = hex_count as usize; + if pos.checked_add(hex_count)? > input.len() { + return None; + } + + let mut value = 0_u32; + for &c in &input[pos..pos + hex_count] { + let digit = hex_value(c)?; + value = (value << 4) + digit as u32; + } + Some((value, hex_count)) +} + +fn parse_char_class_char(input: &[u8], pos: usize) -> Option<(u32, usize)> { + let current = *input.get(pos)?; + if current == b'\\' && pos + 1 < input.len() { + return match input[pos + 1] { + b'x' => Some(parse_hex_escape(input, pos + 2, 2).unwrap_or((b'x' as u32, 0))).map( + |(value, consumed)| { + if consumed > 0 { + (value, 2 + consumed) + } else { + (value, 2) + } + }, + ), + b'u' => Some(parse_hex_escape(input, pos + 2, 4).unwrap_or((b'u' as u32, 0))).map( + |(value, consumed)| { + if consumed > 0 { + (value, 2 + consumed) + } else { + (value, 2) + } + }, + ), + b'U' => Some(parse_hex_escape(input, pos + 2, 8).unwrap_or((b'U' as u32, 0))).map( + |(value, consumed)| { + if consumed > 0 { + (value, 2 + consumed) + } else { + (value, 2) + } + }, + ), + b'n' => Some((b'\n' as u32, 2)), + b't' => Some((b'\t' as u32, 2)), + b'r' => Some((b'\r' as u32, 2)), + b'\\' => Some((b'\\' as u32, 2)), + b']' => Some((b']' as u32, 2)), + b'[' => Some((b'[' as u32, 2)), + other => Some((other as u32, 2)), + }; + } + + Some((current as u32, 1)) +} + +fn sampler_type_from_name(name: &[u8], allow_alt_names: bool) -> c_int { + match name { + b"dry" => 1, + b"top_k" => 2, + b"top_p" => 3, + b"min_p" => 4, + b"typ_p" => 6, + b"temperature" => 7, + b"xtc" => 8, + b"infill" => 9, + b"penalties" => 10, + b"top_n_sigma" => 11, + b"adaptive_p" => 12, + b"top-k" if allow_alt_names => 2, + b"top-p" if allow_alt_names => 3, + b"nucleus" if allow_alt_names => 3, + b"min-p" if allow_alt_names => 4, + b"typical-p" if allow_alt_names => 6, + b"typical" if allow_alt_names => 6, + b"typ-p" if allow_alt_names => 6, + b"typ" if allow_alt_names => 6, + b"temp" if allow_alt_names => 7, + b"top-n-sigma" if allow_alt_names => 11, + b"adaptive-p" if allow_alt_names => 12, + _ => -1, + } +} + +fn sampler_type_from_chr(name: u8) -> c_int { + match name { + b'd' => 1, + b'k' => 2, + b'p' => 3, + b'm' => 4, + b'y' => 6, + b't' => 7, + b'x' => 8, + b'i' => 9, + b'e' => 10, + b's' => 11, + b'a' => 12, + _ => -1, + } +} + +fn into_string_list(parts: Vec>) -> StringList { + let mut strings = parts.into_iter().map(into_ffi).collect::>(); + let len = strings.len(); + let data = strings.as_mut_ptr(); + std::mem::forget(strings); + StringList { data, len } +} + +fn result(status: i32, codepoint: u32, bytes_consumed: usize) -> Utf8ParseResult { + Utf8ParseResult { + codepoint, + bytes_consumed, + status, + } +} + +fn into_ffi(mut bytes: Vec) -> UnicodeString { + let len = bytes.len(); + let data = bytes.as_mut_ptr(); + std::mem::forget(bytes); + UnicodeString { data, len } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_codepoints_like_common_cpp() { + let parsed = parse_utf8_codepoint("A".as_bytes(), 0); + assert_eq!(parsed.status, SUCCESS); + assert_eq!(parsed.codepoint, b'A' as u32); + assert_eq!(parsed.bytes_consumed, 1); + + let parsed = parse_utf8_codepoint(b"\xE2\x80\x9C", 0); + assert_eq!(parsed.status, SUCCESS); + assert_eq!(parsed.codepoint, 0x201C); + assert_eq!(parsed.bytes_consumed, 3); + + assert_eq!(parse_utf8_codepoint(&[0x80], 0).status, INVALID); + assert_eq!(parse_utf8_codepoint(&[0xF0, 0x9F], 0).status, INCOMPLETE); + } + + #[test] + fn utf8_completion_matches_streaming_cases() { + assert!(utf8_is_complete(b"")); + assert!(utf8_is_complete(b"hello")); + assert!(utf8_is_complete(b"hello\xC3\xA9")); + assert!(!utf8_is_complete(b"hello\xC3")); + assert!(!utf8_is_complete(&[0x80])); + } + + #[test] + fn encodes_codepoints() { + assert_eq!(unicode_cpt_to_utf8(0x41).unwrap(), b"A"); + assert_eq!(unicode_cpt_to_utf8(0x00A0).unwrap(), b"\xC2\xA0"); + assert_eq!(unicode_cpt_to_utf8(0x201C).unwrap(), b"\xE2\x80\x9C"); + assert_eq!(unicode_cpt_to_utf8(0x1F600).unwrap(), b"\xF0\x9F\x98\x80"); + assert!(unicode_cpt_to_utf8(0x11_0000).is_none()); + } + + #[test] + fn encodes_codepoint_vectors() { + let cpts = [0x41, 0x00A0, 0x201C, 0x1F600]; + let encoded = unsafe { llama_common_unicode_cpts_to_utf8_rust(cpts.as_ptr(), cpts.len()) }; + unsafe { + assert_eq!( + std::slice::from_raw_parts(encoded.data, encoded.len), + b"A\xC2\xA0\xE2\x80\x9C\xF0\x9F\x98\x80" + ); + llama_common_unicode_string_free(encoded); + + let invalid = [0x41, 0x11_0000]; + let encoded = llama_common_unicode_cpts_to_utf8_rust(invalid.as_ptr(), invalid.len()); + assert!(encoded.data.is_null()); + assert_eq!(encoded.len, 0); + + let null = llama_common_unicode_cpts_to_utf8_rust(std::ptr::null(), 1); + assert!(null.data.is_null()); + assert_eq!(null.len, 0); + } + } + + #[test] + fn maps_jinja_token_types() { + unsafe { + assert_eq!( + CStr::from_ptr(llama_common_jinja_token_type_to_string_rust(0)).to_bytes(), + b"eof" + ); + assert_eq!( + CStr::from_ptr(llama_common_jinja_token_type_to_string_rust(4)).to_bytes(), + b"identifier" + ); + assert_eq!( + CStr::from_ptr(llama_common_jinja_token_type_to_string_rust(20)).to_bytes(), + b"call_operator" + ); + assert_eq!( + CStr::from_ptr(llama_common_jinja_token_type_to_string_rust(25)).to_bytes(), + b"comment" + ); + assert_eq!( + CStr::from_ptr(llama_common_jinja_token_type_to_string_rust(99)).to_bytes(), + b"unknown" + ); + } + } + + #[test] + fn classifies_jinja_word_and_integer_chars() { + assert!(llama_common_jinja_is_word_rust(b'a')); + assert!(llama_common_jinja_is_word_rust(b'Z')); + assert!(llama_common_jinja_is_word_rust(b'5')); + assert!(llama_common_jinja_is_word_rust(b'_')); + assert!(!llama_common_jinja_is_word_rust(b'-')); + assert!(!llama_common_jinja_is_word_rust(0xE9)); + + assert!(llama_common_jinja_is_integer_rust(b'0')); + assert!(llama_common_jinja_is_integer_rust(b'9')); + assert!(!llama_common_jinja_is_integer_rust(b'a')); + assert!(!llama_common_jinja_is_integer_rust(b'_')); + } + + #[test] + fn masks_http_urls() { + unsafe { + let with_user = llama_common_http_show_masked_url_rust( + b"https".as_ptr(), + 5, + true, + b"example.com".as_ptr(), + 11, + b"/models/a.gguf".as_ptr(), + 14, + ); + assert_eq!( + std::slice::from_raw_parts(with_user.data, with_user.len), + b"https://****:****@example.com/models/a.gguf" + ); + llama_common_unicode_string_free(with_user); + + let anonymous = llama_common_http_show_masked_url_rust( + b"http".as_ptr(), + 4, + false, + b"localhost".as_ptr(), + 9, + b"/".as_ptr(), + 1, + ); + assert_eq!( + std::slice::from_raw_parts(anonymous.data, anonymous.len), + b"http://localhost/" + ); + llama_common_unicode_string_free(anonymous); + } + } + + #[test] + fn parses_hex_escapes() { + unsafe { + assert!(llama_common_is_hex_digit_rust(b'0')); + assert!(llama_common_is_hex_digit_rust(b'f')); + assert!(llama_common_is_hex_digit_rust(b'F')); + assert!(!llama_common_is_hex_digit_rust(b'g')); + + let mut value = 0_u32; + assert_eq!( + llama_common_parse_hex_escape_rust(b"xx00afzz".as_ptr(), 8, 2, 4, &mut value), + 4 + ); + assert_eq!(value, 0x00af); + + value = 123; + assert_eq!( + llama_common_parse_hex_escape_rust(b"12xz".as_ptr(), 4, 0, 4, &mut value), + 0 + ); + assert_eq!(value, 123); + assert_eq!( + llama_common_parse_hex_escape_rust(b"12".as_ptr(), 2, 0, 4, &mut value), + 0 + ); + assert_eq!( + llama_common_parse_hex_escape_rust(std::ptr::null(), 1, 0, 1, &mut value), + 0 + ); + } + } + + #[test] + fn parses_character_class_characters() { + unsafe { + let mut value = 0_u32; + assert_eq!( + llama_common_parse_char_class_char_rust(b"a".as_ptr(), 1, 0, &mut value), + 1 + ); + assert_eq!(value, b'a' as u32); + + assert_eq!( + llama_common_parse_char_class_char_rust(br"\n".as_ptr(), 2, 0, &mut value), + 2 + ); + assert_eq!(value, b'\n' as u32); + + assert_eq!( + llama_common_parse_char_class_char_rust(br"\x7f".as_ptr(), 4, 0, &mut value), + 4 + ); + assert_eq!(value, 0x7f); + + assert_eq!( + llama_common_parse_char_class_char_rust(br"\u20ac".as_ptr(), 6, 0, &mut value), + 6 + ); + assert_eq!(value, 0x20ac); + + assert_eq!( + llama_common_parse_char_class_char_rust(br"\xzz".as_ptr(), 4, 0, &mut value), + 2 + ); + assert_eq!(value, b'x' as u32); + + assert_eq!( + llama_common_parse_char_class_char_rust(std::ptr::null(), 1, 0, &mut value), + 0 + ); + } + } + + #[test] + fn maps_chat_and_reasoning_formats() { + unsafe { + assert_eq!( + CStr::from_ptr(llama_common_chat_format_name_rust(0)).to_bytes(), + b"Content-only" + ); + assert_eq!( + CStr::from_ptr(llama_common_chat_format_name_rust(1)).to_bytes(), + b"peg-simple" + ); + assert_eq!( + CStr::from_ptr(llama_common_chat_format_name_rust(2)).to_bytes(), + b"peg-native" + ); + assert_eq!( + CStr::from_ptr(llama_common_chat_format_name_rust(3)).to_bytes(), + b"peg-gemma4" + ); + assert!(llama_common_chat_format_name_rust(99).is_null()); + + assert_eq!( + CStr::from_ptr(llama_common_reasoning_format_name_rust(0)).to_bytes(), + b"none" + ); + assert_eq!( + CStr::from_ptr(llama_common_reasoning_format_name_rust(1)).to_bytes(), + b"auto" + ); + assert_eq!( + CStr::from_ptr(llama_common_reasoning_format_name_rust(2)).to_bytes(), + b"deepseek-legacy" + ); + assert_eq!( + CStr::from_ptr(llama_common_reasoning_format_name_rust(3)).to_bytes(), + b"deepseek" + ); + assert!(llama_common_reasoning_format_name_rust(-1).is_null()); + + assert_eq!( + llama_common_reasoning_format_from_name_rust(c"none".as_ptr()), + 0 + ); + assert_eq!( + llama_common_reasoning_format_from_name_rust(c"auto".as_ptr()), + 1 + ); + assert_eq!( + llama_common_reasoning_format_from_name_rust(c"deepseek-legacy".as_ptr()), + 2 + ); + assert_eq!( + llama_common_reasoning_format_from_name_rust(c"deepseek".as_ptr()), + 3 + ); + assert_eq!( + llama_common_reasoning_format_from_name_rust(c"missing".as_ptr()), + -1 + ); + assert_eq!( + llama_common_reasoning_format_from_name_rust(std::ptr::null()), + -1 + ); + } + } + + #[test] + fn maps_sampler_types() { + unsafe { + assert_eq!(llama_common_sampler_type_to_chr_rust(1) as u8, b'd'); + assert_eq!(llama_common_sampler_type_to_chr_rust(2) as u8, b'k'); + assert_eq!(llama_common_sampler_type_to_chr_rust(3) as u8, b'p'); + assert_eq!(llama_common_sampler_type_to_chr_rust(4) as u8, b'm'); + assert_eq!(llama_common_sampler_type_to_chr_rust(6) as u8, b'y'); + assert_eq!(llama_common_sampler_type_to_chr_rust(7) as u8, b't'); + assert_eq!(llama_common_sampler_type_to_chr_rust(8) as u8, b'x'); + assert_eq!(llama_common_sampler_type_to_chr_rust(9) as u8, b'i'); + assert_eq!(llama_common_sampler_type_to_chr_rust(10) as u8, b'e'); + assert_eq!(llama_common_sampler_type_to_chr_rust(11) as u8, b's'); + assert_eq!(llama_common_sampler_type_to_chr_rust(12) as u8, b'a'); + assert_eq!(llama_common_sampler_type_to_chr_rust(99) as u8, b'?'); + + assert_eq!( + CStr::from_ptr(llama_common_sampler_type_to_str_rust(1)).to_bytes(), + b"dry" + ); + assert_eq!( + CStr::from_ptr(llama_common_sampler_type_to_str_rust(11)).to_bytes(), + b"top_n_sigma" + ); + assert_eq!( + CStr::from_ptr(llama_common_sampler_type_to_str_rust(12)).to_bytes(), + b"adaptive_p" + ); + assert!(llama_common_sampler_type_to_str_rust(0).is_null()); + } + } + + #[test] + fn parses_sampler_types_from_names_and_chars() { + unsafe { + assert_eq!( + llama_common_sampler_type_from_name_rust(c"dry".as_ptr(), false), + 1 + ); + assert_eq!( + llama_common_sampler_type_from_name_rust(c"top_k".as_ptr(), false), + 2 + ); + assert_eq!( + llama_common_sampler_type_from_name_rust(c"top_p".as_ptr(), false), + 3 + ); + assert_eq!( + llama_common_sampler_type_from_name_rust(c"min_p".as_ptr(), false), + 4 + ); + assert_eq!( + llama_common_sampler_type_from_name_rust(c"typ_p".as_ptr(), false), + 6 + ); + assert_eq!( + llama_common_sampler_type_from_name_rust(c"temperature".as_ptr(), false), + 7 + ); + assert_eq!( + llama_common_sampler_type_from_name_rust(c"top-n-sigma".as_ptr(), false), + -1 + ); + assert_eq!( + llama_common_sampler_type_from_name_rust(c"top-n-sigma".as_ptr(), true), + 11 + ); + assert_eq!( + llama_common_sampler_type_from_name_rust(c"nucleus".as_ptr(), true), + 3 + ); + assert_eq!( + llama_common_sampler_type_from_name_rust(c"typ".as_ptr(), true), + 6 + ); + assert_eq!( + llama_common_sampler_type_from_name_rust(c"missing".as_ptr(), true), + -1 + ); + assert_eq!( + llama_common_sampler_type_from_name_rust(std::ptr::null(), true), + -1 + ); + + assert_eq!(llama_common_sampler_type_from_chr_rust(b'd' as c_char), 1); + assert_eq!(llama_common_sampler_type_from_chr_rust(b's' as c_char), 11); + assert_eq!(llama_common_sampler_type_from_chr_rust(b'?' as c_char), -1); + } + } + + #[test] + fn maps_speculative_types() { + unsafe { + let expected: &[(i32, &[u8])] = &[ + (0, b"none"), + (1, b"draft"), + (2, b"eagle3"), + (3, b"ngram_simple"), + (4, b"ngram_map_k"), + (5, b"ngram_map_k4v"), + (6, b"ngram_mod"), + (7, b"ngram_cache"), + ]; + for &(value, name) in expected { + assert_eq!( + CStr::from_ptr(llama_common_speculative_type_to_str_rust(value)).to_bytes(), + name + ); + let c_name = std::ffi::CString::new(name).unwrap(); + assert_eq!( + llama_common_speculative_type_from_name_rust(c_name.as_ptr()), + value + ); + } + assert!(llama_common_speculative_type_to_str_rust(8).is_null()); + assert_eq!( + llama_common_speculative_type_from_name_rust(c"missing".as_ptr()), + -1 + ); + assert_eq!( + llama_common_speculative_type_from_name_rust(std::ptr::null()), + -1 + ); + } + } + + #[test] + fn maps_peg_parse_result_type_names() { + unsafe { + assert_eq!( + CStr::from_ptr(llama_common_peg_parse_result_type_name_rust(0)).to_bytes(), + b"fail" + ); + assert_eq!( + CStr::from_ptr(llama_common_peg_parse_result_type_name_rust(1)).to_bytes(), + b"success" + ); + assert_eq!( + CStr::from_ptr(llama_common_peg_parse_result_type_name_rust(2)).to_bytes(), + b"need_more_input" + ); + assert_eq!( + CStr::from_ptr(llama_common_peg_parse_result_type_name_rust(99)).to_bytes(), + b"unknown" + ); + } + } + + #[test] + fn formats_peg_rule_names_and_literals() { + unsafe { + let rule = llama_common_peg_rule_name_rust(b"tool name!*x".as_ptr(), 12); + assert_eq!( + std::slice::from_raw_parts(rule.data, rule.len), + b"tool-name-x" + ); + llama_common_unicode_string_free(rule); + + let leading = llama_common_peg_rule_name_rust(b"@root".as_ptr(), 5); + assert_eq!( + std::slice::from_raw_parts(leading.data, leading.len), + b"-root" + ); + llama_common_unicode_string_free(leading); + + let literal_input = b"a\r\n\"\\b"; + let literal = + llama_common_gbnf_format_literal_rust(literal_input.as_ptr(), literal_input.len()); + assert_eq!( + std::slice::from_raw_parts(literal.data, literal.len), + b"\"a\\r\\n\\\"\\\\b\"" + ); + llama_common_unicode_string_free(literal); + + let null_literal = llama_common_gbnf_format_literal_rust(std::ptr::null(), 1); + assert_eq!( + std::slice::from_raw_parts(null_literal.data, null_literal.len), + b"\"\"" + ); + llama_common_unicode_string_free(null_literal); + } + } + + #[test] + fn builds_gbnf_repetition() { + unsafe { + let optional = + llama_common_gbnf_build_repetition_rust(b"item".as_ptr(), 4, 0, 1, b"".as_ptr(), 0); + assert_eq!( + std::slice::from_raw_parts(optional.data, optional.len), + b"item?" + ); + llama_common_unicode_string_free(optional); + + let unbounded = llama_common_gbnf_build_repetition_rust( + b"item".as_ptr(), + 4, + 1, + c_int::MAX, + b"".as_ptr(), + 0, + ); + assert_eq!( + std::slice::from_raw_parts(unbounded.data, unbounded.len), + b"item+" + ); + llama_common_unicode_string_free(unbounded); + + let ranged = + llama_common_gbnf_build_repetition_rust(b"item".as_ptr(), 4, 2, 5, b"".as_ptr(), 0); + assert_eq!( + std::slice::from_raw_parts(ranged.data, ranged.len), + b"item{2,5}" + ); + llama_common_unicode_string_free(ranged); + + let separated = llama_common_gbnf_build_repetition_rust( + b"value".as_ptr(), + 5, + 0, + 3, + b"\",\" space".as_ptr(), + 9, + ); + assert_eq!( + std::slice::from_raw_parts(separated.data, separated.len), + b"(value (\",\" space value){0,2})?" + ); + llama_common_unicode_string_free(separated); + + let none = + llama_common_gbnf_build_repetition_rust(b"item".as_ptr(), 4, 0, 0, b"".as_ptr(), 0); + assert_eq!(none.len, 0); + llama_common_unicode_string_free(none); + } + } + + #[test] + fn escapes_gbnf_char_classes() { + unsafe { + for (c, expected) in [ + (b'-' as u32, br"\-".as_slice()), + (b']' as u32, br"\]".as_slice()), + (b'[' as u32, br"\[".as_slice()), + (b'\\' as u32, br"\\".as_slice()), + (b'\n' as u32, br"\n".as_slice()), + (b'\t' as u32, br"\t".as_slice()), + (b'\r' as u32, br"\r".as_slice()), + (b'A' as u32, b"A".as_slice()), + (0x1f, br"\x1F".as_slice()), + (0x1234, br"\u1234".as_slice()), + (0x1f600, br"\U0001F600".as_slice()), + ] { + let escaped = llama_common_gbnf_escape_char_class_rust(c); + assert_eq!( + std::slice::from_raw_parts(escaped.data, escaped.len), + expected + ); + llama_common_unicode_string_free(escaped); + } + } + } + + #[test] + fn detects_gbnf_reserved_names() { + unsafe { + for name in [ + b"root".as_slice(), + b"boolean".as_slice(), + b"decimal-part".as_slice(), + b"date-time-string".as_slice(), + b"uuid".as_slice(), + ] { + assert!(llama_common_gbnf_is_reserved_name_rust( + name.as_ptr(), + name.len() + )); + } + + for name in [ + b"Root".as_slice(), + b"user".as_slice(), + b"date_time".as_slice(), + ] { + assert!(!llama_common_gbnf_is_reserved_name_rust( + name.as_ptr(), + name.len() + )); + } + + assert!(!llama_common_gbnf_is_reserved_name_rust( + std::ptr::null(), + 1 + )); + } + } + + #[test] + fn formats_gbnf_ref_names() { + unsafe { + let input = b"#/$defs/user.name"; + let ref_name = llama_common_gbnf_ref_name_rust(input.as_ptr(), input.len()); + assert_eq!( + std::slice::from_raw_parts(ref_name.data, ref_name.len), + b"ref-defs-user-name" + ); + llama_common_unicode_string_free(ref_name); + + let input = b"https://x/y#foo_bar"; + let external = llama_common_gbnf_ref_name_rust(input.as_ptr(), input.len()); + assert_eq!( + std::slice::from_raw_parts(external.data, external.len), + b"reffoo-bar" + ); + llama_common_unicode_string_free(external); + + let input = b"plain/ref"; + let no_hash = llama_common_gbnf_ref_name_rust(input.as_ptr(), input.len()); + assert_eq!( + std::slice::from_raw_parts(no_hash.data, no_hash.len), + b"refplain-ref" + ); + llama_common_unicode_string_free(no_hash); + + let null = llama_common_gbnf_ref_name_rust(std::ptr::null(), 1); + assert_eq!(std::slice::from_raw_parts(null.data, null.len), b"ref"); + llama_common_unicode_string_free(null); + } + } + + #[test] + fn maps_hf_cache_repo_folder_names() { + unsafe { + let repo = + llama_common_hf_folder_name_to_repo_rust(b"models--owner--repo-name".as_ptr(), 24); + assert_eq!( + std::slice::from_raw_parts(repo.data, repo.len), + b"owner/repo-name" + ); + llama_common_unicode_string_free(repo); + + let missing_input = b"datasets--owner--repo"; + let missing = llama_common_hf_folder_name_to_repo_rust( + missing_input.as_ptr(), + missing_input.len(), + ); + assert_eq!(missing.len, 0); + llama_common_unicode_string_free(missing); + + let folder = llama_common_hf_repo_to_folder_name_rust(b"owner/repo/name".as_ptr(), 15); + assert_eq!( + std::slice::from_raw_parts(folder.data, folder.len), + b"models--owner--repo--name" + ); + llama_common_unicode_string_free(folder); + + let old = llama_common_hf_make_old_cache_filename_rust( + b"owner".as_ptr(), + 5, + b"repo".as_ptr(), + 4, + b"sub/dir/model.gguf".as_ptr(), + 18, + ); + assert_eq!( + std::slice::from_raw_parts(old.data, old.len), + b"owner_repo_sub_dir_model.gguf" + ); + llama_common_unicode_string_free(old); + } + } + + #[test] + fn counts_json_brace_depth() { + unsafe { + let balanced = br#"{"a": {"b": 1}, "c": "ignore { this }"}"#; + assert_eq!( + llama_common_json_brace_depth_rust(balanced.as_ptr(), balanced.len()), + 0 + ); + + let open = br#"{"a": {"text": "escaped \" } still string"}"#; + assert_eq!( + llama_common_json_brace_depth_rust(open.as_ptr(), open.len()), + 1 + ); + + let extra_close = br#"{"a": 1}}"#; + assert_eq!( + llama_common_json_brace_depth_rust(extra_close.as_ptr(), extra_close.len()), + -1 + ); + + assert_eq!(llama_common_json_brace_depth_rust(std::ptr::null(), 1), 0); + } + } + + #[test] + fn validates_hf_identifiers() { + unsafe { + assert!(llama_common_hf_is_valid_repo_id_rust( + b"owner/repo".as_ptr(), + 10 + )); + let repo_with_specials = b"own_er/re.po-1"; + assert!(llama_common_hf_is_valid_repo_id_rust( + repo_with_specials.as_ptr(), + repo_with_specials.len() + )); + assert!(!llama_common_hf_is_valid_repo_id_rust(b"owner".as_ptr(), 5)); + assert!(!llama_common_hf_is_valid_repo_id_rust( + b"/owner/repo".as_ptr(), + 11 + )); + assert!(!llama_common_hf_is_valid_repo_id_rust( + b"owner//repo".as_ptr(), + 11 + )); + assert!(!llama_common_hf_is_valid_repo_id_rust( + b"owner/repo/extra".as_ptr(), + 16 + )); + assert!(!llama_common_hf_is_valid_repo_id_rust(std::ptr::null(), 1)); + + let valid_token = b"hf_abcdefghijklmnopqrstuvwxyzABCDEFGH"; + assert!(llama_common_hf_is_valid_token_rust( + valid_token.as_ptr(), + valid_token.len() + )); + assert!(!llama_common_hf_is_valid_token_rust( + b"hf_short".as_ptr(), + 8 + )); + assert!(!llama_common_hf_is_valid_token_rust( + b"hf_abcdefghijklmnopqrstuvwxyzABCDEFG!".as_ptr(), + 37 + )); + + let sha40 = [b'a'; 40]; + let sha64 = [b'F'; 64]; + assert!(llama_common_hf_is_valid_commit_rust( + sha40.as_ptr(), + sha40.len() + )); + assert!(llama_common_hf_is_valid_oid_rust( + sha40.as_ptr(), + sha40.len() + )); + assert!(llama_common_hf_is_valid_oid_rust( + sha64.as_ptr(), + sha64.len() + )); + assert!(!llama_common_hf_is_valid_commit_rust( + sha64.as_ptr(), + sha64.len() + )); + assert!(!llama_common_hf_is_valid_oid_rust(b"not-hex".as_ptr(), 7)); + } + } + + #[test] + fn parses_hf_manifest_names() { + unsafe { + let valid = b"manifest=owner=repo=model.gguf.json"; + let parsed = llama_common_hf_parse_manifest_name_rust(valid.as_ptr(), valid.len()); + assert_eq!(parsed.len, 2); + assert_eq!( + std::slice::from_raw_parts((*parsed.data.add(0)).data, (*parsed.data.add(0)).len), + b"owner" + ); + assert_eq!( + std::slice::from_raw_parts((*parsed.data.add(1)).data, (*parsed.data.add(1)).len), + b"repo" + ); + llama_common_string_list_free(parsed); + + let bad_prefix_input = b"other=owner=repo=file.json"; + let bad_prefix = llama_common_hf_parse_manifest_name_rust( + bad_prefix_input.as_ptr(), + bad_prefix_input.len(), + ); + assert_eq!(bad_prefix.len, 0); + llama_common_string_list_free(bad_prefix); + + let missing_repo_input = b"manifest=owner==file.json"; + let missing_repo = llama_common_hf_parse_manifest_name_rust( + missing_repo_input.as_ptr(), + missing_repo_input.len(), + ); + assert_eq!(missing_repo.len, 0); + llama_common_string_list_free(missing_repo); + + let bad_suffix_input = b"manifest=owner=repo=file.txt"; + let bad_suffix = llama_common_hf_parse_manifest_name_rust( + bad_suffix_input.as_ptr(), + bad_suffix_input.len(), + ); + assert_eq!(bad_suffix.len, 0); + llama_common_string_list_free(bad_suffix); + + let null = llama_common_hf_parse_manifest_name_rust(std::ptr::null(), 1); + assert!(null.data.is_null()); + } + } + + #[test] + fn parses_gguf_split_info() { + unsafe { + let split_input = b"folder/model-Q4_K_M-00002-of-00008.gguf"; + let split = + llama_common_gguf_split_info_rust(split_input.as_ptr(), split_input.len(), true); + assert_eq!( + std::slice::from_raw_parts(split.prefix.data, split.prefix.len), + b"folder/model-Q4_K_M" + ); + assert_eq!( + std::slice::from_raw_parts(split.tag.data, split.tag.len), + b"Q4_K_M" + ); + assert_eq!(split.index, 2); + assert_eq!(split.count, 8); + llama_common_unicode_string_free(split.prefix); + llama_common_unicode_string_free(split.tag); + + let unsplit_input = b"model.f16.gguf"; + let unsplit = llama_common_gguf_split_info_rust( + unsplit_input.as_ptr(), + unsplit_input.len(), + true, + ); + assert_eq!( + std::slice::from_raw_parts(unsplit.prefix.data, unsplit.prefix.len), + b"model.f16" + ); + assert_eq!( + std::slice::from_raw_parts(unsplit.tag.data, unsplit.tag.len), + b"F16" + ); + assert_eq!(unsplit.index, 1); + assert_eq!(unsplit.count, 1); + llama_common_unicode_string_free(unsplit.prefix); + llama_common_unicode_string_free(unsplit.tag); + + let no_tag_input = b"model-q8_0.gguf"; + let no_tag = + llama_common_gguf_split_info_rust(no_tag_input.as_ptr(), no_tag_input.len(), false); + assert_eq!( + std::slice::from_raw_parts(no_tag.tag.data, no_tag.tag.len), + b"" + ); + llama_common_unicode_string_free(no_tag.prefix); + llama_common_unicode_string_free(no_tag.tag); + + let invalid = llama_common_gguf_split_info_rust(b"model.bin".as_ptr(), 9, true); + assert_eq!(invalid.prefix.len, 0); + assert_eq!(invalid.tag.len, 0); + assert_eq!(invalid.index, 0); + assert_eq!(invalid.count, 0); + llama_common_unicode_string_free(invalid.prefix); + llama_common_unicode_string_free(invalid.tag); + } + } + + #[test] + fn extracts_gguf_quant_bits() { + unsafe { + let q4 = b"folder/model-Q4_K_M-00002-of-00008.gguf"; + assert_eq!( + llama_common_gguf_extract_quant_bits_rust(q4.as_ptr(), q4.len()), + 4 + ); + + let f16 = b"model.f16.gguf"; + assert_eq!( + llama_common_gguf_extract_quant_bits_rust(f16.as_ptr(), f16.len()), + 16 + ); + + let nvfp4 = b"model-NVFP4.gguf"; + assert_eq!( + llama_common_gguf_extract_quant_bits_rust(nvfp4.as_ptr(), nvfp4.len()), + 4 + ); + + let no_tag = b"model.gguf"; + assert_eq!( + llama_common_gguf_extract_quant_bits_rust(no_tag.as_ptr(), no_tag.len()), + 0 + ); + + assert_eq!( + llama_common_gguf_extract_quant_bits_rust(std::ptr::null(), 1), + 0 + ); + } + } + + #[test] + fn compares_ascii_case_insensitively() { + unsafe { + assert_eq!( + llama_common_ascii_case_equal(c"adamw".as_ptr(), c"ADAMW".as_ptr()), + 1 + ); + assert_eq!( + llama_common_ascii_case_equal(c"sgd".as_ptr(), c"SgD".as_ptr()), + 1 + ); + assert_eq!( + llama_common_ascii_case_equal(c"sgd".as_ptr(), c"sgdx".as_ptr()), + 0 + ); + assert_eq!( + llama_common_ascii_case_equal(std::ptr::null(), c"sgd".as_ptr()), + 0 + ); + } + } + + #[test] + fn maps_optimizer_names() { + unsafe { + assert_eq!(llama_common_opt_get_optimizer_rust(c"adamw".as_ptr()), 0); + assert_eq!(llama_common_opt_get_optimizer_rust(c"ADAMW".as_ptr()), 0); + assert_eq!(llama_common_opt_get_optimizer_rust(c"sgd".as_ptr()), 1); + assert_eq!(llama_common_opt_get_optimizer_rust(c"SgD".as_ptr()), 1); + assert_eq!(llama_common_opt_get_optimizer_rust(c"missing".as_ptr()), 2); + assert_eq!(llama_common_opt_get_optimizer_rust(std::ptr::null()), 2); + } + } + + #[test] + fn classifies_arg_boolean_values() { + unsafe { + assert_eq!(llama_common_arg_is_truthy_rust(c"on".as_ptr()), 1); + assert_eq!(llama_common_arg_is_truthy_rust(c"enabled".as_ptr()), 1); + assert_eq!(llama_common_arg_is_truthy_rust(c"true".as_ptr()), 1); + assert_eq!(llama_common_arg_is_truthy_rust(c"1".as_ptr()), 1); + assert_eq!(llama_common_arg_is_truthy_rust(c"yes".as_ptr()), 0); + + assert_eq!(llama_common_arg_is_falsey_rust(c"off".as_ptr()), 1); + assert_eq!(llama_common_arg_is_falsey_rust(c"disabled".as_ptr()), 1); + assert_eq!(llama_common_arg_is_falsey_rust(c"false".as_ptr()), 1); + assert_eq!(llama_common_arg_is_falsey_rust(c"0".as_ptr()), 1); + assert_eq!(llama_common_arg_is_falsey_rust(c"no".as_ptr()), 0); + + assert_eq!(llama_common_arg_is_autoy_rust(c"auto".as_ptr()), 1); + assert_eq!(llama_common_arg_is_autoy_rust(c"-1".as_ptr()), 1); + assert_eq!(llama_common_arg_is_autoy_rust(c"automatic".as_ptr()), 0); + assert_eq!(llama_common_arg_is_autoy_rust(std::ptr::null()), 0); + } + } + + #[test] + fn parses_bool_values() { + unsafe { + assert_eq!(llama_common_parse_bool_value_rust(c"true".as_ptr()), 1); + assert_eq!(llama_common_parse_bool_value_rust(c"1".as_ptr()), 1); + assert_eq!(llama_common_parse_bool_value_rust(c"false".as_ptr()), 0); + assert_eq!(llama_common_parse_bool_value_rust(c"0".as_ptr()), 0); + assert_eq!(llama_common_parse_bool_value_rust(c"maybe".as_ptr()), -1); + assert_eq!(llama_common_parse_bool_value_rust(std::ptr::null()), -1); + } + } + + #[test] + fn parses_openai_compatible_tool_choice() { + unsafe { + assert_eq!( + llama_common_chat_tool_choice_parse_oaicompat_rust(c"auto".as_ptr()), + 0 + ); + assert_eq!( + llama_common_chat_tool_choice_parse_oaicompat_rust(c"required".as_ptr()), + 1 + ); + assert_eq!( + llama_common_chat_tool_choice_parse_oaicompat_rust(c"none".as_ptr()), + 2 + ); + assert_eq!( + llama_common_chat_tool_choice_parse_oaicompat_rust(c"AUTO".as_ptr()), + -1 + ); + assert_eq!( + llama_common_chat_tool_choice_parse_oaicompat_rust(std::ptr::null()), + -1 + ); + } + } + + #[test] + fn formats_line_and_column() { + unsafe { + let source = b"one\ntwo\nthree"; + let start = llama_common_get_line_col_rust(source.as_ptr(), source.len(), 0); + assert_eq!( + std::slice::from_raw_parts(start.data, start.len), + b"line 1, column 1" + ); + llama_common_unicode_string_free(start); + + let second_line = llama_common_get_line_col_rust(source.as_ptr(), source.len(), 5); + assert_eq!( + std::slice::from_raw_parts(second_line.data, second_line.len), + b"line 2, column 2" + ); + llama_common_unicode_string_free(second_line); + + let past_end = llama_common_get_line_col_rust(source.as_ptr(), source.len(), 99); + assert_eq!( + std::slice::from_raw_parts(past_end.data, past_end.len), + b"line 3, column 6" + ); + llama_common_unicode_string_free(past_end); + + let null = llama_common_get_line_col_rust(std::ptr::null(), 1, 10); + assert_eq!( + std::slice::from_raw_parts(null.data, null.len), + b"line 1, column 1" + ); + llama_common_unicode_string_free(null); + } + } + + #[test] + fn parses_cpu_hex_masks() { + unsafe { + let mut mask = [false; 16]; + assert_eq!( + llama_common_parse_cpu_mask_rust(b"0x9".as_ptr(), 3, mask.as_mut_ptr(), mask.len()), + -1 + ); + assert!(mask[0]); + assert!(!mask[1]); + assert!(!mask[2]); + assert!(mask[3]); + + let mut mask = [false; 16]; + assert_eq!( + llama_common_parse_cpu_mask_rust(b"3f".as_ptr(), 2, mask.as_mut_ptr(), mask.len()), + -1 + ); + assert_eq!( + &mask[..8], + &[true, true, true, true, true, true, false, false] + ); + + assert_eq!( + llama_common_parse_cpu_mask_rust(b"0xG".as_ptr(), 3, mask.as_mut_ptr(), mask.len()), + 2 + ); + } + } + + #[test] + fn cpu_hex_mask_parser_preserves_empty_and_truncation_behavior() { + unsafe { + let mut mask = [false; 520]; + assert_eq!( + llama_common_parse_cpu_mask_rust(b"0x".as_ptr(), 2, mask.as_mut_ptr(), mask.len()), + -1 + ); + assert!(mask.iter().all(|value| !*value)); + + let mut long = vec![b'0'; 129]; + long[128] = b'G'; + assert_eq!( + llama_common_parse_cpu_mask_rust( + long.as_ptr(), + long.len(), + mask.as_mut_ptr(), + mask.len() + ), + -1 + ); + } + } + + #[test] + fn maps_kv_cache_types() { + unsafe { + assert_eq!(llama_common_kv_cache_type_from_str_rust(c"f32".as_ptr()), 0); + assert_eq!(llama_common_kv_cache_type_from_str_rust(c"f16".as_ptr()), 1); + assert_eq!( + llama_common_kv_cache_type_from_str_rust(c"bf16".as_ptr()), + 30 + ); + assert_eq!( + llama_common_kv_cache_type_from_str_rust(c"q8_0".as_ptr()), + 8 + ); + assert_eq!( + llama_common_kv_cache_type_from_str_rust(c"q4_0".as_ptr()), + 2 + ); + assert_eq!( + llama_common_kv_cache_type_from_str_rust(c"q4_1".as_ptr()), + 3 + ); + assert_eq!( + llama_common_kv_cache_type_from_str_rust(c"iq4_nl".as_ptr()), + 20 + ); + assert_eq!( + llama_common_kv_cache_type_from_str_rust(c"q5_0".as_ptr()), + 6 + ); + assert_eq!( + llama_common_kv_cache_type_from_str_rust(c"q5_1".as_ptr()), + 7 + ); + assert_eq!( + llama_common_kv_cache_type_from_str_rust(c"missing".as_ptr()), + -1 + ); + assert_eq!( + llama_common_kv_cache_type_from_str_rust(std::ptr::null()), + -1 + ); + } + } + + #[test] + fn formats_kv_cache_type_list() { + unsafe { + assert_eq!( + CStr::from_ptr(llama_common_get_all_kv_cache_types_rust()).to_bytes(), + b"f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1" + ); + } + } + + #[test] + fn normalizes_embeddings() { + unsafe { + let input = [3.0_f32, 4.0]; + let mut output = [0.0_f32; 2]; + llama_common_embd_normalize_rust( + input.as_ptr(), + output.as_mut_ptr(), + input.len() as c_int, + 2, + ); + assert!((output[0] - 0.6).abs() < 1e-6); + assert!((output[1] - 0.8).abs() < 1e-6); + + llama_common_embd_normalize_rust( + input.as_ptr(), + output.as_mut_ptr(), + input.len() as c_int, + -1, + ); + assert_eq!(output, input); + + let input = [32760.0_f32, -16380.0]; + llama_common_embd_normalize_rust( + input.as_ptr(), + output.as_mut_ptr(), + input.len() as c_int, + 0, + ); + assert!((output[0] - 32760.0).abs() < 1e-3); + assert!((output[1] + 16380.0).abs() < 1e-3); + } + } + + #[test] + fn computes_embedding_cosine_similarity() { + unsafe { + let a = [1.0_f32, 0.0]; + let b = [0.0_f32, 1.0]; + assert_eq!( + llama_common_embd_similarity_cos_rust(a.as_ptr(), b.as_ptr(), 2), + 0.0 + ); + + let c = [2.0_f32, 0.0]; + assert_eq!( + llama_common_embd_similarity_cos_rust(a.as_ptr(), c.as_ptr(), 2), + 1.0 + ); + + let z = [0.0_f32, 0.0]; + assert_eq!( + llama_common_embd_similarity_cos_rust(z.as_ptr(), z.as_ptr(), 2), + 1.0 + ); + assert_eq!( + llama_common_embd_similarity_cos_rust(a.as_ptr(), z.as_ptr(), 2), + 0.0 + ); + } + } + + #[test] + fn initializes_learning_rate_decay() { + unsafe { + let mut decay_epochs = -1.0_f32; + let mut scale_epoch = 0.0_f32; + llama_common_lr_opt_init_rust( + 1.0e-3, + 1.0e-5, + -1.0, + 4, + &mut decay_epochs, + &mut scale_epoch, + ); + assert_eq!(decay_epochs, 4.0); + assert!((scale_epoch - 1.660_964).abs() < 1e-5); + + llama_common_lr_opt_init_rust( + 1.0e-3, + 1.0e-5, + 2.0, + 4, + &mut decay_epochs, + &mut scale_epoch, + ); + assert_eq!(decay_epochs, 2.0); + assert!((scale_epoch - 3.321_928).abs() < 1e-5); + + decay_epochs = -1.0; + scale_epoch = 7.0; + llama_common_lr_opt_init_rust( + 1.0e-3, + -1.0, + decay_epochs, + 4, + &mut decay_epochs, + &mut scale_epoch, + ); + assert_eq!(decay_epochs, -1.0); + assert_eq!(scale_epoch, 0.0); + } + } + + #[test] + fn computes_learning_rate_schedule() { + assert_eq!( + llama_common_lr_opt_get_lr_rust(1.0e-3, -1.0, -1.0, 0.0, 99.0), + 1.0e-3 + ); + assert_eq!( + llama_common_lr_opt_get_lr_rust(1.0e-3, 1.0e-5, 2.0, 3.321_928, 2.0), + 1.0e-5 + ); + let lr = llama_common_lr_opt_get_lr_rust(1.0e-3, 1.0e-5, 2.0, 3.321_928, 1.0); + assert!((lr - 1.0e-4).abs() < 1e-8); + } + + #[test] + fn maps_log_levels_to_verbosity() { + assert_eq!(llama_common_log_level_to_verbosity_rust(1), 4); + assert_eq!(llama_common_log_level_to_verbosity_rust(2), 3); + assert_eq!(llama_common_log_level_to_verbosity_rust(3), 2); + assert_eq!(llama_common_log_level_to_verbosity_rust(4), 1); + assert_eq!(llama_common_log_level_to_verbosity_rust(5), 3); + assert_eq!(llama_common_log_level_to_verbosity_rust(0), 0); + assert_eq!(llama_common_log_level_to_verbosity_rust(99), 0); + } + + #[test] + fn returns_microsecond_epoch_time() { + let first = llama_common_time_us_rust(); + let second = llama_common_time_us_rust(); + assert!(first > 0); + assert!(second >= first); + } + + #[test] + fn classifies_http_success_statuses() { + assert!(!llama_common_is_http_status_ok_rust(199)); + assert!(llama_common_is_http_status_ok_rust(200)); + assert!(llama_common_is_http_status_ok_rust(302)); + assert!(llama_common_is_http_status_ok_rust(399)); + assert!(!llama_common_is_http_status_ok_rust(400)); + assert!(!llama_common_is_http_status_ok_rust(500)); + } + + #[test] + fn detects_gguf_model_filenames() { + unsafe { + let model = b"Q4_K_M/model.gguf"; + assert!(llama_common_gguf_filename_is_model_rust( + model.as_ptr(), + model.len() + )); + let parent_with_mmproj = b"folder/mmproj_parent/model.gguf"; + assert!(llama_common_gguf_filename_is_model_rust( + parent_with_mmproj.as_ptr(), + parent_with_mmproj.len() + )); + let not_gguf = b"model.bin"; + assert!(!llama_common_gguf_filename_is_model_rust( + not_gguf.as_ptr(), + not_gguf.len() + )); + let mmproj = b"mmproj-model.gguf"; + assert!(!llama_common_gguf_filename_is_model_rust( + mmproj.as_ptr(), + mmproj.len() + )); + let imatrix = b"imatrix-model.gguf"; + assert!(!llama_common_gguf_filename_is_model_rust( + imatrix.as_ptr(), + imatrix.len() + )); + assert!(!llama_common_gguf_filename_is_model_rust( + std::ptr::null(), + 1 + )); + } + } + + #[test] + fn formats_bool_strings() { + unsafe { + assert_eq!( + CStr::from_ptr(llama_common_bool_to_string_rust(true)).to_bytes(), + b"true" + ); + assert_eq!( + CStr::from_ptr(llama_common_bool_to_string_rust(false)).to_bytes(), + b"false" + ); + } + } + + #[test] + fn formats_source_peeks() { + unsafe { + let source = b"abcdef"; + let peek = llama_common_peak_source_rust(source.as_ptr(), source.len(), 3, 2); + assert_eq!( + std::slice::from_raw_parts(peek.data, peek.len), + b"...bcde...\n ^" + ); + llama_common_unicode_string_free(peek); + + let multiline = b"ab\ncd"; + let peek = llama_common_peak_source_rust(multiline.as_ptr(), multiline.len(), 2, 4); + assert_eq!( + std::slice::from_raw_parts(peek.data, peek.len), + "...ab↵cd...\n ^".as_bytes() + ); + llama_common_unicode_string_free(peek); + + let empty = llama_common_peak_source_rust(b"".as_ptr(), 0, 0, 40); + assert_eq!( + std::slice::from_raw_parts(empty.data, empty.len), + b"(no source available)" + ); + llama_common_unicode_string_free(empty); + } + } + + #[test] + fn formats_errors_with_source() { + unsafe { + let source = b"abcdef"; + let formatted = llama_common_fmt_error_with_source_rust( + b"lexer".as_ptr(), + 5, + b"bad token".as_ptr(), + 9, + source.as_ptr(), + source.len(), + 3, + ); + assert_eq!( + std::slice::from_raw_parts(formatted.data, formatted.len), + b"lexer: bad token\n...abcdef...\n ^" + ); + llama_common_unicode_string_free(formatted); + + let null = llama_common_fmt_error_with_source_rust( + std::ptr::null(), + 1, + b"bad".as_ptr(), + 3, + source.as_ptr(), + source.len(), + 0, + ); + assert_eq!(null.len, 0); + llama_common_unicode_string_free(null); + } + } + + #[test] + fn repeats_strings() { + unsafe { + let repeated = llama_common_string_repeat_rust(b"ab".as_ptr(), 2, 3); + assert!(!repeated.data.is_null()); + assert_eq!( + std::slice::from_raw_parts(repeated.data, repeated.len), + b"ababab" + ); + llama_common_unicode_string_free(repeated); + + let empty = llama_common_string_repeat_rust(b"ab".as_ptr(), 2, 0); + assert_eq!(empty.len, 0); + llama_common_unicode_string_free(empty); + + let null = llama_common_string_repeat_rust(std::ptr::null(), 1, 2); + assert!(null.data.is_null()); + assert_eq!(null.len, 0); + } + } + + #[test] + fn formats_int_vectors() { + unsafe { + let values = [1, -2, 300]; + let formatted = llama_common_int_vec_to_string_rust(values.as_ptr(), values.len()); + assert!(!formatted.data.is_null()); + assert_eq!( + std::slice::from_raw_parts(formatted.data, formatted.len), + b"[ 1, -2, 300 ]" + ); + llama_common_unicode_string_free(formatted); + + let empty = llama_common_int_vec_to_string_rust(values.as_ptr(), 0); + assert_eq!(std::slice::from_raw_parts(empty.data, empty.len), b"[ ]"); + llama_common_unicode_string_free(empty); + + let null = llama_common_int_vec_to_string_rust(std::ptr::null(), 1); + assert!(null.data.is_null()); + assert_eq!(null.len, 0); + } + } + + #[test] + fn trims_strings() { + unsafe { + let both = llama_common_string_trim_rust(b" \tvalue\n".as_ptr(), 8, 0); + assert_eq!(std::slice::from_raw_parts(both.data, both.len), b"value"); + llama_common_unicode_string_free(both); + + let leading = llama_common_string_trim_rust(b" \tvalue\n".as_ptr(), 8, 1); + assert_eq!( + std::slice::from_raw_parts(leading.data, leading.len), + b"value\n" + ); + llama_common_unicode_string_free(leading); + + let trailing = llama_common_string_trim_rust(b" \tvalue\n".as_ptr(), 8, 2); + assert_eq!( + std::slice::from_raw_parts(trailing.data, trailing.len), + b" \tvalue" + ); + llama_common_unicode_string_free(trailing); + + let newlines = llama_common_string_trim_rust(b"value\n\n ".as_ptr(), 8, 3); + assert_eq!( + std::slice::from_raw_parts(newlines.data, newlines.len), + b"value\n\n " + ); + llama_common_unicode_string_free(newlines); + + let all_ws = llama_common_string_trim_rust(b" \n\t".as_ptr(), 3, 0); + assert_eq!(all_ws.len, 0); + llama_common_unicode_string_free(all_ws); + } + } + + #[test] + fn strips_custom_character_sets() { + unsafe { + let left = llama_common_string_lstrip_chars_rust( + b" \t\r\nvalue \n".as_ptr(), + 11, + c" \t\r\n".as_ptr(), + ); + assert_eq!(std::slice::from_raw_parts(left.data, left.len), b"value \n"); + llama_common_unicode_string_free(left); + + let right = llama_common_string_rstrip_chars_rust( + b" \tvalue\r\n".as_ptr(), + 9, + c" \t\r\n".as_ptr(), + ); + assert_eq!( + std::slice::from_raw_parts(right.data, right.len), + b" \tvalue" + ); + llama_common_unicode_string_free(right); + + let all = + llama_common_string_lstrip_chars_rust(b"\r\n\t".as_ptr(), 3, c" \t\r\n".as_ptr()); + assert_eq!(all.len, 0); + llama_common_unicode_string_free(all); + + let unchanged = + llama_common_string_rstrip_chars_rust(b"value".as_ptr(), 5, c" \t\r\n".as_ptr()); + assert_eq!( + std::slice::from_raw_parts(unchanged.data, unchanged.len), + b"value" + ); + llama_common_unicode_string_free(unchanged); + + let null = llama_common_string_lstrip_chars_rust(std::ptr::null(), 1, c" ".as_ptr()); + assert!(null.data.is_null()); + } + } + + #[test] + fn breaks_strings_into_wrapped_lines() { + unsafe { + let wrapped = + llama_common_break_str_into_lines_rust(b"alpha beta gamma\ndelta".as_ptr(), 22, 10); + assert_eq!(wrapped.len, 3); + assert_eq!( + std::slice::from_raw_parts((*wrapped.data.add(0)).data, (*wrapped.data.add(0)).len), + b"alpha beta" + ); + assert_eq!( + std::slice::from_raw_parts((*wrapped.data.add(1)).data, (*wrapped.data.add(1)).len), + b"gamma" + ); + assert_eq!( + std::slice::from_raw_parts((*wrapped.data.add(2)).data, (*wrapped.data.add(2)).len), + b"delta" + ); + llama_common_string_list_free(wrapped); + } + } + + #[test] + fn breaks_strings_like_getline_edge_cases() { + unsafe { + let empty = llama_common_break_str_into_lines_rust(b"".as_ptr(), 0, 10); + assert_eq!(empty.len, 0); + llama_common_string_list_free(empty); + + let newline = llama_common_break_str_into_lines_rust(b"\n".as_ptr(), 1, 10); + assert_eq!(newline.len, 1); + assert_eq!( + std::slice::from_raw_parts((*newline.data).data, (*newline.data).len), + b"" + ); + llama_common_string_list_free(newline); + + let trailing = llama_common_break_str_into_lines_rust(b"a\n".as_ptr(), 2, 10); + assert_eq!(trailing.len, 1); + assert_eq!( + std::slice::from_raw_parts((*trailing.data).data, (*trailing.data).len), + b"a" + ); + llama_common_string_list_free(trailing); + + let long_word = + llama_common_break_str_into_lines_rust(b"alphabet soup".as_ptr(), 13, 3); + assert_eq!(long_word.len, 2); + assert_eq!( + std::slice::from_raw_parts( + (*long_word.data.add(0)).data, + (*long_word.data.add(0)).len + ), + b"alphabet" + ); + assert_eq!( + std::slice::from_raw_parts( + (*long_word.data.add(1)).data, + (*long_word.data.add(1)).len + ), + b"soup" + ); + llama_common_string_list_free(long_word); + } + } + + #[test] + fn removes_leading_dashes() { + unsafe { + let long = llama_common_rm_leading_dashes_rust(b"--threads".as_ptr(), 9); + assert_eq!(std::slice::from_raw_parts(long.data, long.len), b"threads"); + llama_common_unicode_string_free(long); + + let short = llama_common_rm_leading_dashes_rust(b"-t".as_ptr(), 2); + assert_eq!(std::slice::from_raw_parts(short.data, short.len), b"t"); + llama_common_unicode_string_free(short); + + let unchanged = llama_common_rm_leading_dashes_rust(b"LLAMA_ARG_THREADS".as_ptr(), 17); + assert_eq!( + std::slice::from_raw_parts(unchanged.data, unchanged.len), + b"LLAMA_ARG_THREADS" + ); + llama_common_unicode_string_free(unchanged); + + let only_dashes = llama_common_rm_leading_dashes_rust(b"---".as_ptr(), 3); + assert_eq!(only_dashes.len, 0); + llama_common_unicode_string_free(only_dashes); + } + } + + #[test] + fn extracts_around_common_prefixes_and_suffixes() { + unsafe { + let before = llama_common_until_common_prefix_rust( + b"before {\"first\": value".as_ptr(), + 23, + b"{\"first\":".as_ptr(), + 9, + b"{\"second\":".as_ptr(), + 10, + ); + assert_eq!( + std::slice::from_raw_parts(before.data, before.len), + b"before " + ); + llama_common_unicode_string_free(before); + + let after = llama_common_after_common_suffix_rust( + b"value \"XXXX\"} after".as_ptr(), + 19, + b"\"XXXX\"}".as_ptr(), + 7, + b"\"YYYY\"}".as_ptr(), + 7, + ); + assert_eq!(std::slice::from_raw_parts(after.data, after.len), b" after"); + llama_common_unicode_string_free(after); + } + } + + #[test] + fn common_prefix_suffix_extractors_return_empty_without_shared_anchor() { + unsafe { + let no_prefix = llama_common_until_common_prefix_rust( + b"full".as_ptr(), + 4, + b"left".as_ptr(), + 4, + b"right".as_ptr(), + 5, + ); + assert_eq!(no_prefix.len, 0); + llama_common_unicode_string_free(no_prefix); + + let missing_prefix = llama_common_until_common_prefix_rust( + b"full".as_ptr(), + 4, + b"same-left".as_ptr(), + 9, + b"same-right".as_ptr(), + 10, + ); + assert_eq!(missing_prefix.len, 0); + llama_common_unicode_string_free(missing_prefix); + + let no_suffix = llama_common_after_common_suffix_rust( + b"full".as_ptr(), + 4, + b"left".as_ptr(), + 4, + b"right".as_ptr(), + 5, + ); + assert_eq!(no_suffix.len, 0); + llama_common_unicode_string_free(no_suffix); + + let missing_suffix = llama_common_after_common_suffix_rust( + b"full".as_ptr(), + 4, + b"left-same".as_ptr(), + 9, + b"right-same".as_ptr(), + 10, + ); + assert_eq!(missing_suffix.len, 0); + llama_common_unicode_string_free(missing_suffix); + } + } + + #[test] + fn ensures_json_ascii_inside_strings() { + unsafe { + let input = "\"é 😀\""; + let escaped = + llama_common_json_ensure_ascii_preserving_format_rust(input.as_ptr(), input.len()); + assert_eq!( + std::slice::from_raw_parts(escaped.data, escaped.len), + br#""\u00e9 \ud83d\ude00""# + ); + llama_common_unicode_string_free(escaped); + + let ascii_input = br#"{"x":"already \" escaped"}"#; + let ascii = llama_common_json_ensure_ascii_preserving_format_rust( + ascii_input.as_ptr(), + ascii_input.len(), + ); + assert_eq!( + std::slice::from_raw_parts(ascii.data, ascii.len), + br#"{"x":"already \" escaped"}"# + ); + llama_common_unicode_string_free(ascii); + } + } + + #[test] + fn json_ascii_formatter_preserves_format_and_replaces_invalid_string_bytes() { + unsafe { + let invalid = [b'"', 0xC3, b'(', b'"']; + let escaped = llama_common_json_ensure_ascii_preserving_format_rust( + invalid.as_ptr(), + invalid.len(), + ); + assert_eq!( + std::slice::from_raw_parts(escaped.data, escaped.len), + br#""\ufffd(""# + ); + llama_common_unicode_string_free(escaped); + + let outside = "é \"é\""; + let escaped = llama_common_json_ensure_ascii_preserving_format_rust( + outside.as_ptr(), + outside.len(), + ); + let mut expected = "é ".as_bytes().to_vec(); + expected.extend_from_slice(br#""\u00e9""#); + assert_eq!( + std::slice::from_raw_parts(escaped.data, escaped.len), + expected.as_slice() + ); + llama_common_unicode_string_free(escaped); + } + } + + #[test] + fn normalizes_python_style_quotes_to_json() { + unsafe { + let input = br#"{'key': 'value'}"#; + let normalized = + llama_common_normalize_quotes_to_json_rust(input.as_ptr(), input.len()); + assert_eq!( + std::slice::from_raw_parts(normalized.data, normalized.len), + br#"{"key": "value"}"# + ); + llama_common_unicode_string_free(normalized); + + let code = br#"{'code': 'print(\'hello\')'}"#; + let normalized = llama_common_normalize_quotes_to_json_rust(code.as_ptr(), code.len()); + assert_eq!( + std::slice::from_raw_parts(normalized.data, normalized.len), + br#"{"code": "print('hello')"}"# + ); + llama_common_unicode_string_free(normalized); + + let mixed = br#"{'msg': 'He said "hi"', "keep": "it's ok"}"#; + let normalized = + llama_common_normalize_quotes_to_json_rust(mixed.as_ptr(), mixed.len()); + assert_eq!( + std::slice::from_raw_parts(normalized.data, normalized.len), + br#"{"msg": "He said \"hi\"", "keep": "it's ok"}"# + ); + llama_common_unicode_string_free(normalized); + + let null = llama_common_normalize_quotes_to_json_rust(std::ptr::null(), 1); + assert!(null.data.is_null()); + assert_eq!(null.len, 0); + } + } + + #[test] + fn computes_common_prefix_and_suffix_lengths() { + unsafe { + assert_eq!( + llama_common_prefix_len_rust(b"abcdef".as_ptr(), 6, b"abcXYZ".as_ptr(), 6), + 3 + ); + assert_eq!( + llama_common_suffix_len_rust(b"XYZdef".as_ptr(), 6, b"abcdef".as_ptr(), 6), + 3 + ); + assert_eq!( + llama_common_prefix_len_rust(b"abc".as_ptr(), 3, b"xyz".as_ptr(), 3), + 0 + ); + assert_eq!( + llama_common_suffix_len_rust(b"abc".as_ptr(), 3, b"xyz".as_ptr(), 3), + 0 + ); + assert_eq!( + llama_common_prefix_len_rust(std::ptr::null(), 1, b"abc".as_ptr(), 3), + 0 + ); + } + } + + #[test] + fn matches_simple_globs() { + unsafe { + assert!(llama_common_glob_match_rust( + b"*.gguf".as_ptr(), + 6, + b"model.gguf".as_ptr(), + 10 + )); + assert!(!llama_common_glob_match_rust( + b"*.gguf".as_ptr(), + 6, + b"dir/model.gguf".as_ptr(), + 14 + )); + assert!(llama_common_glob_match_rust( + b"**/*.gguf".as_ptr(), + 9, + b"dir/model.gguf".as_ptr(), + 14 + )); + assert!(llama_common_glob_match_rust( + b"file-?.txt".as_ptr(), + 10, + b"file-a.txt".as_ptr(), + 10 + )); + assert!(!llama_common_glob_match_rust( + b"file-?.txt".as_ptr(), + 10, + b"file-/.txt".as_ptr(), + 10 + )); + } + } + + #[test] + fn matches_glob_character_classes() { + unsafe { + assert!(llama_common_glob_match_rust( + b"model-[0-9].gguf".as_ptr(), + 16, + b"model-7.gguf".as_ptr(), + 12 + )); + assert!(!llama_common_glob_match_rust( + b"model-[!0-9].gguf".as_ptr(), + 17, + b"model-7.gguf".as_ptr(), + 12 + )); + assert!(llama_common_glob_match_rust( + b"model-[!0-9].gguf".as_ptr(), + 17, + b"model-a.gguf".as_ptr(), + 12 + )); + assert!(llama_common_glob_match_rust( + b"file-[]].txt".as_ptr(), + 12, + b"file-].txt".as_ptr(), + 10 + )); + assert!(llama_common_glob_match_rust( + b"file-[-].txt".as_ptr(), + 12, + b"file--.txt".as_ptr(), + 10 + )); + assert!(llama_common_glob_match_rust( + b"literal[".as_ptr(), + 8, + b"literal[".as_ptr(), + 8 + )); + assert!(!llama_common_glob_match_rust( + b"literal[".as_ptr(), + 8, + b"literalx".as_ptr(), + 8 + )); + } + } + + #[test] + fn validates_filenames() { + unsafe { + assert!(llama_common_fs_validate_filename_rust( + b"model.gguf".as_ptr(), + 10, + false + )); + assert!(llama_common_fs_validate_filename_rust( + "unicode-é.gguf".as_ptr(), + 14, + false + )); + assert!(!llama_common_fs_validate_filename_rust( + b"".as_ptr(), + 0, + false + )); + assert!(!llama_common_fs_validate_filename_rust( + b" leading.gguf".as_ptr(), + 13, + false + )); + assert!(!llama_common_fs_validate_filename_rust( + b"trailing.gguf ".as_ptr(), + 14, + false + )); + assert!(!llama_common_fs_validate_filename_rust( + b"trailing.".as_ptr(), + 9, + false + )); + assert!(!llama_common_fs_validate_filename_rust( + b"bad:name.gguf".as_ptr(), + 13, + false + )); + assert!(!llama_common_fs_validate_filename_rust( + b"bad..name.gguf".as_ptr(), + 14, + false + )); + assert!(!llama_common_fs_validate_filename_rust( + b".".as_ptr(), + 1, + false + )); + } + } + + #[test] + fn validates_filename_subdir_and_utf8_rules() { + unsafe { + assert!(!llama_common_fs_validate_filename_rust( + b"dir/model.gguf".as_ptr(), + 14, + false + )); + assert!(llama_common_fs_validate_filename_rust( + b"dir/model.gguf".as_ptr(), + 14, + true + )); + assert!(!llama_common_fs_validate_filename_rust( + b"bad\\name.gguf".as_ptr(), + 13, + false + )); + assert!(llama_common_fs_validate_filename_rust( + b"bad\\name.gguf".as_ptr(), + 13, + true + )); + assert!(!llama_common_fs_validate_filename_rust( + b"bad\nname.gguf".as_ptr(), + 13, + false + )); + let overlong_nul = [0xC0, 0x80]; + assert!(!llama_common_fs_validate_filename_rust( + overlong_nul.as_ptr(), + 2, + false + )); + assert!(!llama_common_fs_validate_filename_rust( + "bad\u{ff0e}name".as_ptr(), + 10, + false + )); + assert!(!llama_common_fs_validate_filename_rust( + std::ptr::null(), + 1, + false + )); + } + } + + #[test] + fn replaces_all_substrings() { + unsafe { + let replaced = llama_common_string_replace_all_rust( + b"a/b/c".as_ptr(), + 5, + b"/".as_ptr(), + 1, + b"--".as_ptr(), + 2, + ); + assert_eq!( + std::slice::from_raw_parts(replaced.data, replaced.len), + b"a--b--c" + ); + llama_common_unicode_string_free(replaced); + + let unchanged = llama_common_string_replace_all_rust( + b"abc".as_ptr(), + 3, + b"".as_ptr(), + 0, + b"x".as_ptr(), + 1, + ); + assert_eq!( + std::slice::from_raw_parts(unchanged.data, unchanged.len), + b"abc" + ); + llama_common_unicode_string_free(unchanged); + + let null = llama_common_string_replace_all_rust( + std::ptr::null(), + 1, + b"a".as_ptr(), + 1, + b"b".as_ptr(), + 1, + ); + assert!(null.data.is_null()); + } + } + + #[test] + fn escapes_regex_special_characters() { + unsafe { + let escaped = llama_common_regex_escape_rust(b"a.b[c]\\d+".as_ptr(), 9); + assert_eq!( + std::slice::from_raw_parts(escaped.data, escaped.len), + br"a\.b\[c\]\\d\+" + ); + llama_common_unicode_string_free(escaped); + + let plain = llama_common_regex_escape_rust(b"abc".as_ptr(), 3); + assert_eq!(std::slice::from_raw_parts(plain.data, plain.len), b"abc"); + llama_common_unicode_string_free(plain); + } + } + + #[test] + fn joins_string_views() { + unsafe { + let values = [ + StringView { + data: b"one".as_ptr(), + len: 3, + }, + StringView { + data: b"two".as_ptr(), + len: 3, + }, + StringView { + data: b"three".as_ptr(), + len: 5, + }, + ]; + let joined = + llama_common_string_join_rust(values.as_ptr(), values.len(), b", ".as_ptr(), 2); + assert_eq!( + std::slice::from_raw_parts(joined.data, joined.len), + b"one, two, three" + ); + llama_common_unicode_string_free(joined); + + let empty = llama_common_string_join_rust(values.as_ptr(), 0, b", ".as_ptr(), 2); + assert_eq!(empty.len, 0); + llama_common_unicode_string_free(empty); + + let bad = [StringView { + data: std::ptr::null(), + len: 1, + }]; + let joined = llama_common_string_join_rust(bad.as_ptr(), bad.len(), b", ".as_ptr(), 2); + assert!(joined.data.is_null()); + } + } + + #[test] + fn splits_strings() { + unsafe { + let split = llama_common_string_split_rust(b"a/b/".as_ptr(), 4, b"/".as_ptr(), 1); + assert_eq!(split.len, 3); + let items = std::slice::from_raw_parts(split.data, split.len); + assert_eq!( + std::slice::from_raw_parts(items[0].data, items[0].len), + b"a" + ); + assert_eq!( + std::slice::from_raw_parts(items[1].data, items[1].len), + b"b" + ); + assert_eq!(items[2].len, 0); + llama_common_string_list_free(split); + + let no_delim = llama_common_string_split_rust(b"abc".as_ptr(), 3, b"/".as_ptr(), 1); + assert_eq!(no_delim.len, 1); + let items = std::slice::from_raw_parts(no_delim.data, no_delim.len); + assert_eq!( + std::slice::from_raw_parts(items[0].data, items[0].len), + b"abc" + ); + llama_common_string_list_free(no_delim); + + let empty_delim = llama_common_string_split_rust(b"abc".as_ptr(), 3, b"".as_ptr(), 0); + assert_eq!(empty_delim.len, 1); + llama_common_string_list_free(empty_delim); + } + } + + #[test] + fn parses_csv_rows() { + unsafe { + let input = br#"value1,"value, with, commas","value with ""escaped"" quotes",value4"#; + let parsed = llama_common_parse_csv_row_rust(input.as_ptr(), input.len()); + assert_eq!(parsed.len, 4); + let items = std::slice::from_raw_parts(parsed.data, parsed.len); + assert_eq!( + std::slice::from_raw_parts(items[0].data, items[0].len), + b"value1" + ); + assert_eq!( + std::slice::from_raw_parts(items[1].data, items[1].len), + b"value, with, commas" + ); + assert_eq!( + std::slice::from_raw_parts(items[2].data, items[2].len), + br#"value with "escaped" quotes"# + ); + assert_eq!( + std::slice::from_raw_parts(items[3].data, items[3].len), + b"value4" + ); + llama_common_string_list_free(parsed); + + let trailing = llama_common_parse_csv_row_rust(b"a,".as_ptr(), 2); + assert_eq!(trailing.len, 2); + let items = std::slice::from_raw_parts(trailing.data, trailing.len); + assert_eq!( + std::slice::from_raw_parts(items[0].data, items[0].len), + b"a" + ); + assert_eq!(items[1].len, 0); + llama_common_string_list_free(trailing); + } + } + + #[test] + fn trims_space_views() { + unsafe { + let input = b" \tvalue \n"; + + let leading = + llama_common_trim_leading_space_view_rust(input.as_ptr(), input.len(), -1); + assert_eq!( + std::slice::from_raw_parts(leading.data, leading.len), + b"value \n" + ); + + let leading_limited = + llama_common_trim_leading_space_view_rust(input.as_ptr(), input.len(), 1); + assert_eq!( + std::slice::from_raw_parts(leading_limited.data, leading_limited.len), + b"\tvalue \n" + ); + + let trailing = + llama_common_trim_trailing_space_view_rust(input.as_ptr(), input.len(), -1); + assert_eq!( + std::slice::from_raw_parts(trailing.data, trailing.len), + b" \tvalue" + ); + + let trailing_limited = + llama_common_trim_trailing_space_view_rust(input.as_ptr(), input.len(), 1); + assert_eq!( + std::slice::from_raw_parts(trailing_limited.data, trailing_limited.len), + b" \tvalue " + ); + + let null = llama_common_trim_leading_space_view_rust(std::ptr::null(), 1, -1); + assert!(null.data.is_null()); + assert_eq!(null.len, 0); + } + } + + #[test] + fn processes_escape_sequences() { + unsafe { + let input = br#"a\nb\t\x41\\\'\""#; + let processed = llama_common_string_process_escapes_rust(input.as_ptr(), input.len()); + assert_eq!( + std::slice::from_raw_parts(processed.data, processed.len), + &[b'a', b'\n', b'b', b'\t', b'A', b'\\', b'\'', b'"'] + ); + llama_common_unicode_string_free(processed); + + let invalid_input = br"\xZZ\q"; + let invalid = llama_common_string_process_escapes_rust( + invalid_input.as_ptr(), + invalid_input.len(), + ); + assert_eq!( + std::slice::from_raw_parts(invalid.data, invalid.len), + br"\xZZ\q" + ); + llama_common_unicode_string_free(invalid); + } + } + + #[test] + fn cleans_file_names() { + unsafe { + let cleaned = llama_common_clean_file_name_rust(br"org\repo/model".as_ptr(), 14); + assert_eq!( + std::slice::from_raw_parts(cleaned.data, cleaned.len), + b"org_repo_model" + ); + llama_common_unicode_string_free(cleaned); + + let unchanged = llama_common_clean_file_name_rust(b"model.gguf".as_ptr(), 10); + assert_eq!( + std::slice::from_raw_parts(unchanged.data, unchanged.len), + b"model.gguf" + ); + llama_common_unicode_string_free(unchanged); + } + } + + #[test] + fn computes_string_diffs() { + unsafe { + let mut status = -99; + let diff = llama_common_string_diff_rust( + b"hello".as_ptr(), + 5, + b"hello world".as_ptr(), + 11, + &mut status, + ); + assert_eq!(status, 0); + assert_eq!(std::slice::from_raw_parts(diff.data, diff.len), b" world"); + llama_common_unicode_string_free(diff); + + let partial = llama_common_string_diff_rust( + b"hello world".as_ptr(), + 11, + b"hello".as_ptr(), + 5, + &mut status, + ); + assert_eq!(status, 0); + assert_eq!(partial.len, 0); + llama_common_unicode_string_free(partial); + + let invalid = + llama_common_string_diff_rust(b"abc".as_ptr(), 3, b"xyz".as_ptr(), 3, &mut status); + assert_eq!(status, -1); + assert!(invalid.data.is_null()); + } + } + + #[test] + fn checks_string_prefixes_and_suffixes() { + unsafe { + assert!(llama_common_string_starts_with_rust( + b"abcdef".as_ptr(), + 6, + b"abc".as_ptr(), + 3 + )); + assert!(!llama_common_string_starts_with_rust( + b"abcdef".as_ptr(), + 6, + b"abd".as_ptr(), + 3 + )); + assert!(llama_common_string_starts_with_rust( + b"abcdef".as_ptr(), + 6, + b"".as_ptr(), + 0 + )); + assert!(!llama_common_string_starts_with_rust( + std::ptr::null(), + 1, + b"a".as_ptr(), + 1 + )); + + assert!(llama_common_string_ends_with_rust( + b"abcdef".as_ptr(), + 6, + b"def".as_ptr(), + 3 + )); + assert!(!llama_common_string_ends_with_rust( + b"abcdef".as_ptr(), + 6, + b"cef".as_ptr(), + 3 + )); + assert!(llama_common_string_ends_with_rust( + b"abcdef".as_ptr(), + 6, + b"".as_ptr(), + 0 + )); + assert!(!llama_common_string_ends_with_rust( + b"abc".as_ptr(), + 3, + std::ptr::null(), + 1 + )); + } + } +} diff --git a/common/unicode.cpp b/common/unicode.cpp deleted file mode 100644 index f71fe56783ff..000000000000 --- a/common/unicode.cpp +++ /dev/null @@ -1,124 +0,0 @@ -#include "unicode.h" - -#include -#include -#include -#include -#include - -// implementation adopted from src/unicode.cpp - -size_t common_utf8_sequence_length(unsigned char first_byte) { - const size_t lookup[] = { 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4 }; - uint8_t highbits = static_cast(first_byte) >> 4; - return lookup[highbits]; -} - -utf8_parse_result common_parse_utf8_codepoint(std::string_view input, size_t offset) { - if (offset >= input.size()) { - return utf8_parse_result(utf8_parse_result::INCOMPLETE); - } - - // ASCII fast path - if (!(input[offset] & 0x80)) { - return utf8_parse_result(utf8_parse_result::SUCCESS, input[offset], 1); - } - - // Invalid: continuation byte as first byte - if (!(input[offset] & 0x40)) { - return utf8_parse_result(utf8_parse_result::INVALID); - } - - // 2-byte sequence - if (!(input[offset] & 0x20)) { - if (offset + 1 >= input.size()) { - return utf8_parse_result(utf8_parse_result::INCOMPLETE); - } - if ((input[offset + 1] & 0xc0) != 0x80) { - return utf8_parse_result(utf8_parse_result::INVALID); - } - auto result = ((input[offset] & 0x1f) << 6) | (input[offset + 1] & 0x3f); - return utf8_parse_result(utf8_parse_result::SUCCESS, result, 2); - } - - // 3-byte sequence - if (!(input[offset] & 0x10)) { - if (offset + 2 >= input.size()) { - return utf8_parse_result(utf8_parse_result::INCOMPLETE); - } - if ((input[offset + 1] & 0xc0) != 0x80 || (input[offset + 2] & 0xc0) != 0x80) { - return utf8_parse_result(utf8_parse_result::INVALID); - } - auto result = ((input[offset] & 0x0f) << 12) | ((input[offset + 1] & 0x3f) << 6) | (input[offset + 2] & 0x3f); - return utf8_parse_result(utf8_parse_result::SUCCESS, result, 3); - } - - // 4-byte sequence - if (!(input[offset] & 0x08)) { - if (offset + 3 >= input.size()) { - return utf8_parse_result(utf8_parse_result::INCOMPLETE); - } - if ((input[offset + 1] & 0xc0) != 0x80 || (input[offset + 2] & 0xc0) != 0x80 || (input[offset + 3] & 0xc0) != 0x80) { - return utf8_parse_result(utf8_parse_result::INVALID); - } - auto result = ((input[offset] & 0x07) << 18) | ((input[offset + 1] & 0x3f) << 12) | ((input[offset + 2] & 0x3f) << 6) | (input[offset + 3] & 0x3f); - return utf8_parse_result(utf8_parse_result::SUCCESS, result, 4); - } - - // Invalid first byte - return utf8_parse_result(utf8_parse_result::INVALID); -} - -bool common_utf8_is_complete(const std::string & s) { - if (s.empty()) { - return true; - } - for (int i = 1; i <= std::min(4, (int)s.size()); i++) { - unsigned char c = s[s.size() - i]; - if ((c & 0xC0) != 0x80) { - int expected = (c >= 0xF0) ? 4 : (c >= 0xE0) ? 3 : (c >= 0xC0) ? 2 : 1; - return i >= expected; - } - } - return false; -} - -std::string common_unicode_cpts_to_utf8(const std::vector & cps) { - std::string result; - for (size_t i = 0; i < cps.size(); ++i) { - result.append(common_unicode_cpt_to_utf8(cps[i])); - } - return result; -} - -std::string common_unicode_cpt_to_utf8(uint32_t cpt) { - std::string result; - - if (/* 0x00 <= cpt && */ cpt <= 0x7f) { - result.push_back(cpt); - return result; - } - if (0x80 <= cpt && cpt <= 0x7ff) { - result.push_back(0xc0 | ((cpt >> 6) & 0x1f)); - result.push_back(0x80 | (cpt & 0x3f)); - return result; - } - if (0x800 <= cpt && cpt <= 0xffff) { - result.push_back(0xe0 | ((cpt >> 12) & 0x0f)); - result.push_back(0x80 | ((cpt >> 6) & 0x3f)); - result.push_back(0x80 | (cpt & 0x3f)); - return result; - } - if (0x10000 <= cpt && cpt <= 0x10ffff) { - result.push_back(0xf0 | ((cpt >> 18) & 0x07)); - result.push_back(0x80 | ((cpt >> 12) & 0x3f)); - result.push_back(0x80 | ((cpt >> 6) & 0x3f)); - result.push_back(0x80 | (cpt & 0x3f)); - return result; - } - - throw std::invalid_argument("invalid codepoint"); -} - - - diff --git a/common/unicode.h b/common/unicode.h deleted file mode 100644 index 9b32fa19d62b..000000000000 --- a/common/unicode.h +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once - -#include -#include -#include -#include - -// UTF-8 parsing utilities for streaming-aware unicode support - -struct utf8_parse_result { - uint32_t codepoint; // Decoded codepoint (only valid if status == SUCCESS) - size_t bytes_consumed; // How many bytes this codepoint uses (1-4) - enum status { SUCCESS, INCOMPLETE, INVALID } status; - - utf8_parse_result(enum status s, uint32_t cp = 0, size_t bytes = 0) - : codepoint(cp), bytes_consumed(bytes), status(s) {} -}; - -// Determine the expected length of a UTF-8 sequence from its first byte -// Returns 0 for invalid first bytes -size_t common_utf8_sequence_length(unsigned char first_byte); - -// Check if a string ends with a complete UTF-8 sequence. -bool common_utf8_is_complete(const std::string & s); - -// Parse a single UTF-8 codepoint from input -utf8_parse_result common_parse_utf8_codepoint(std::string_view input, size_t offset); - -std::string common_unicode_cpts_to_utf8(const std::vector & cps); -std::string common_unicode_cpt_to_utf8(uint32_t cpt); diff --git a/docs/autoparser.md b/docs/autoparser.md index adc4d43ed677..8bec50cbccde 100644 --- a/docs/autoparser.md +++ b/docs/autoparser.md @@ -241,13 +241,13 @@ common_chat_params (prompt, parser, grammar, triggers, preserved_tokens) ## Entry Point -The auto-parser is invoked in [common/chat.cpp:1280-1310](common/chat.cpp#L1280-L1310) in `common_chat_templates_apply_jinja`. A few specialized templates are handled first (Ministral/Magistral Large 3, GPT-OSS with `<|channel|>`, Functionary v3.2 with `>>>all`), then the auto-parser handles everything else via `autoparser::autoparser` + `peg_generator::generate_parser`. +The auto-parser is invoked in `common/common.cpp.inc` in `common_chat_templates_apply_jinja`. A few specialized templates are handled first (Ministral/Magistral Large 3, GPT-OSS with `<|channel|>`, Functionary v3.2 with `>>>all`), then the auto-parser handles everything else via `autoparser::autoparser` + `peg_generator::generate_parser`. ## Algorithm Details ### Core Mechanism: Differential Comparison -All analysis phases use the same factorized comparison function declared in [common/chat-auto-parser-helpers.h:68](common/chat-auto-parser-helpers.h#L68): +All analysis phases use the same factorized comparison function declared in [common/chat-auto-parser.h](common/chat-auto-parser.h): ```cpp compare_variants(tmpl, params_A, params_modifier) @@ -350,7 +350,7 @@ Classification logic: ### Workarounds -A workaround array in `common/chat-diff-analyzer.cpp` applies post-hoc patches after analysis. Each workaround is a lambda that inspects the template source and overrides analysis results. Current workarounds: +A workaround array in `common/common.cpp.inc` applies post-hoc patches after analysis. Each workaround is a lambda that inspects the template source and overrides analysis results. Current workarounds: 1. **Old Qwen/DeepSeek thinking templates** — source contains `content.split('')` but not ``: sets `reasoning.mode = TAG_BASED` with ``/`` markers if no reasoning was detected 2. **Granite 3.3** — source contains specific "Write your thoughts" text: forces `TAG_BASED` reasoning with ``/`` and `WRAPPED_WITH_REASONING` content with ``/`` @@ -437,12 +437,7 @@ Each returned parser is wrapped by `wrap_for_generation_prompt()`, which prepend | File | Purpose | |-------------------------------------------|---------------------------------------------------------------------------------| | `common/chat-auto-parser.h` | All analysis structs, enums, `autoparser`, `peg_generator`, `generation_params` | -| `common/chat-auto-parser-generator.cpp` | Parser generator: `generate_parser()` and `build_parser()` methods | -| `common/chat-diff-analyzer.cpp` | Differential analysis implementation and workarounds | -| `common/chat-auto-parser-helpers.h/cpp` | `calculate_diff_split()`, `segmentize_markers()`, `compare_variants()`, | -| | `wrap_for_generation_prompt()`, string helpers | -| `common/chat-peg-parser.h/cpp` | `common_chat_peg_builder`, `common_chat_peg_mapper`, and helpers | -| `common/chat.cpp` | Entry point: `common_chat_templates_apply_jinja()` | +| `common/common.cpp.inc` | Parser generator, differential analysis, PEG/chat grammar helpers, and chat entry points | | `tools/parser/debug-template-parser.cpp` | Debug tool for template analysis | | `tools/parser/template-analysis.cpp` | Template analysis tool | @@ -519,8 +514,8 @@ The following templates have active tests in `tests/test-chat.cpp`: To support a new template format: 1. **If it follows standard patterns** — The auto-parser should detect it automatically. Run `llama-debug-template-parser` to verify markers are correctly extracted. -2. **If differential analysis extracts incorrect markers** — Add a workaround lambda to the `workarounds` vector in `common/chat-diff-analyzer.cpp`. Inspect the template source for a unique identifying substring. -3. **If it needs fundamentally different handling** — Add a dedicated handler function in `chat.cpp` before the auto-parser block (as done for GPT-OSS, Functionary v3.2, and Ministral). +2. **If differential analysis extracts incorrect markers** — Add a workaround lambda to the `workarounds` vector in `common/common.cpp.inc`. Inspect the template source for a unique identifying substring. +3. **If it needs fundamentally different handling** — Add a dedicated handler function in `common/common.cpp.inc` before the auto-parser block (as done for GPT-OSS, Functionary v3.2, and Ministral). ## Edge Cases and Quirks diff --git a/docs/build-s390x.md b/docs/build-s390x.md index 4568d5010f6c..2b4b08e9252a 100644 --- a/docs/build-s390x.md +++ b/docs/build-s390x.md @@ -3,7 +3,7 @@ # Build llama.cpp locally (for s390x) -The main product of this project is the `llama` library. Its C-style interface can be found in [include/llama.h](../include/llama.h). +The main product of this project is the `llama` library. Its C-style interface can be found in [include/llama.h.inc](../include/llama.h.inc). The project also includes many example programs and tools using the `llama` library. The examples range from simple, minimal code snippets to sophisticated sub-projects such as an OpenAI-compatible HTTP server. diff --git a/docs/build.md b/docs/build.md index a18479b3346b..fa7d7c4b9271 100644 --- a/docs/build.md +++ b/docs/build.md @@ -1,6 +1,6 @@ # Build llama.cpp locally -The main product of this project is the `llama` library. Its C-style interface can be found in [include/llama.h](../include/llama.h). +The main product of this project is the `llama` library. Its C-style interface can be found in [include/llama.h.inc](../include/llama.h.inc). The project also includes many example programs and tools using the `llama` library. The examples range from simple, minimal code snippets to sophisticated sub-projects such as an OpenAI-compatible HTTP server. diff --git a/docs/development/HOWTO-add-model.md b/docs/development/HOWTO-add-model.md index 5390e42ba94f..935c76c056b3 100644 --- a/docs/development/HOWTO-add-model.md +++ b/docs/development/HOWTO-add-model.md @@ -96,17 +96,17 @@ NOTE: Tensor names must end with `.weight` or `.bias` suffixes, that is the conv The model params and tensors layout must be defined in `llama.cpp` source files: 1. Define a new `llm_arch` enum value in `src/llama-arch.h`. -2. In `src/llama-arch.cpp`: +2. In `src/llama.cpp`: - Add the architecture name to the `LLM_ARCH_NAMES` map. - Add the list of model tensors to `llm_get_tensor_names` (you may also need to update `LLM_TENSOR_NAMES`) -3. Add any non-standard metadata loading in the `llama_model_loader` constructor in `src/llama-model-loader.cpp`. -4. If the model has a RoPE operation, add a case for the architecture in `llama_model_rope_type` function in `src/llama-model.cpp`. +3. Add any non-standard metadata loading in the `llama_model_loader` constructor in `src/llama.cpp`. +4. If the model has a RoPE operation, add a case for the architecture in `llama_model_rope_type` function in `src/llama.cpp`. NOTE: The dimensions in `ggml` are typically in the reverse order of the `pytorch` dimensions. ### 3. Build the GGML graph implementation -This is the funniest part, you have to provide the inference graph implementation of the new model architecture in `src/llama-model.cpp`. +This is the funniest part, you have to provide the inference graph implementation of the new model architecture in `src/llama.cpp`. Create a new struct that inherits from `llm_graph_context` and implement the graph-building logic in its constructor. Have a look at existing implementations like `llm_build_llama`, `llm_build_dbrx` or `llm_build_bert`. Then, in the `llama_model::build_graph` method, add a case for your architecture to instantiate your new graph-building struct. @@ -120,8 +120,8 @@ Note: to debug the inference graph: you can use [llama-eval-callback](/examples/ If the new model supports multimodal inputs, you will need to add a new encoder definition in `libmtmd`. You can find more information about llama.cpp's multimodal support in [the docs](../multimodal.md) and in the `tools/mtmd` source directory. 1. In the conversion script, make sure you add a subclass that extends `MmprojModel` or another class that inherits from the same base class. -2. Add the encoder definition in `clip.cpp`. -3. Implement the preprocessor in `mtmd.cpp`. In most cases, you can reuse an existing preprocessor. +2. Add the encoder definition in `tools/mtmd/mtmd.cpp.inc`. +3. Implement the preprocessor in `mtmd.cpp.inc`. In most cases, you can reuse an existing preprocessor. 4. Implement the encoder GGML graph, either in a dedicated file if the model is truly different from existing ones, or by reusing an existing implementation (for example: siglip, pixtral, or qwen) and adding a model-specific projector. Note: @@ -138,7 +138,7 @@ PyTorch implementations usually prefer explicitly calculating `freq_cis`/`sin`/` However, since `ggml_rope_ext` only provides a subset of the RoPE implementations that models use, converting models from PyTorch to llama.cpp may require some creative adaptations. -For more information about `ggml_rope_ext`, please refer to the in-code documentation in `ggml.h`. +For more information about `ggml_rope_ext`, please refer to the in-code documentation in `ggml.h.inc`. Examples: - `libmtmd` implements 2D RoPE with `GGML_ROPE_TYPE_NORMAL` ordering by splitting the input tensor in half, applying `ggml_rope_ext` separately to each half, then joining them back together using `ggml_concat`. diff --git a/docs/development/parsing.md b/docs/development/parsing.md index a41057db2b8a..25eb3b6e867f 100644 --- a/docs/development/parsing.md +++ b/docs/development/parsing.md @@ -175,7 +175,7 @@ Most model output can be placed in one of the following categories: (Qwen3-Coder, MiniMax M2) or pseudo-function calls (LFM2) To provide broad coverage, -[`common/chat-peg-parser.h`](/common/chat-peg-parser.h) contains builders and +[`common/chat-auto-parser.h`](/common/chat-auto-parser.h) contains builders and mappers that help create parsers and visitors/extractors for these types. They require parsers to tag nodes to conform to an AST "shape". This normalization makes it easy to extract information and generalize parsing. diff --git a/docs/development/rust-port.md b/docs/development/rust-port.md new file mode 100644 index 000000000000..68bb64530dc4 --- /dev/null +++ b/docs/development/rust-port.md @@ -0,0 +1,21 @@ +# Rust port status + +This private-fork migration is using a measurable rule: the port is not complete while any C/C++ source or header files remain. + +Check the current count: + +```sh +scripts/rust-port-status.sh +``` + +Use the hard gate when the migration is expected to be complete: + +```sh +scripts/rust-port-status.sh enforce +``` + +Current first Rust-owned artifact: + +- `examples/simple-rust`: Rust port of `examples/simple/simple.cpp` over the llama.cpp C API, with unit tests for argument parsing. + +The final target is stricter than adding Rust bindings: all `.c`, `.cc`, `.cpp`, `.cxx`, `.h`, `.hh`, `.hpp`, and `.hxx` files must be removed or replaced. diff --git a/docs/function-calling.md b/docs/function-calling.md index 9ede914c04d6..6eca2126c8b2 100644 --- a/docs/function-calling.md +++ b/docs/function-calling.md @@ -1,6 +1,6 @@ # Function Calling -[chat.h](../common/chat.h) (https://github.com/ggml-org/llama.cpp/pull/9639) adds support for [OpenAI-style function calling](https://platform.openai.com/docs/guides/function-calling) and is used in: +[common.h.inc](../common/common.h.inc) (https://github.com/ggml-org/llama.cpp/pull/9639) adds support for [OpenAI-style function calling](https://platform.openai.com/docs/guides/function-calling) and is used in: - `llama-server` when started w/ `--jinja` flag ## Universal support w/ Native & Generic handlers diff --git a/docs/preset.md b/docs/preset.md index d49fb0a1aeb3..b291e6b3948b 100644 --- a/docs/preset.md +++ b/docs/preset.md @@ -25,7 +25,7 @@ top-k = 20 top-p = 0.95 ``` -For security reasons, only certain options are allowed. Please refer to [preset.cpp](../common/preset.cpp) for the complete list of permitted options. +For security reasons, only certain options are allowed. Please refer to [common.cpp.inc](../common/common.cpp.inc) for the complete list of permitted options. Example usage: diff --git a/examples/batched-rust/Cargo.toml b/examples/batched-rust/Cargo.toml new file mode 100644 index 000000000000..10b2f1580cf7 --- /dev/null +++ b/examples/batched-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-batched-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-batched" +path = "src/main.rs" diff --git a/examples/batched-rust/build.rs b/examples/batched-rust/build.rs new file mode 100644 index 000000000000..301e2342b007 --- /dev/null +++ b/examples/batched-rust/build.rs @@ -0,0 +1,45 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); + + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-batched"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/batched-rust/src/main.rs b/examples/batched-rust/src/main.rs new file mode 100644 index 000000000000..ccfd0a2f0a6e --- /dev/null +++ b/examples/batched-rust/src/main.rs @@ -0,0 +1,595 @@ +use llama_simple_rust::ffi; +use std::env; +use std::ffi::CString; +use std::io::{self, Write}; +use std::ptr; + +#[derive(Debug, Clone, PartialEq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub n_predict: i32, + pub n_parallel: i32, + pub n_gpu_layers: i32, + pub top_k: i32, + pub top_p: f32, + pub temp: f32, + pub seed: u32, + pub backend_sampling: bool, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum ParseError { + MissingModel, + MissingValue(String), + InvalidInteger(String, String), + InvalidFloat(String, String), + InvalidParallel(i32), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + ParseError::InvalidFloat(flag, value) => write!(f, "invalid float for {flag}: {value}"), + ParseError::InvalidParallel(value) => { + write!(f, "-np/--parallel must be positive, got {value}") + } + } + } +} + +impl std::error::Error for ParseError {} + +impl Default for Args { + fn default() -> Self { + Self { + model_path: String::new(), + prompt: "Hello my name is".to_string(), + n_predict: 32, + n_parallel: 1, + n_gpu_layers: 99, + top_k: 40, + top_p: 0.95, + temp: 0.8, + seed: ffi::LLAMA_DEFAULT_SEED, + backend_sampling: false, + } + } +} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut parsed = Args::default(); + let mut iter = args.into_iter().map(Into::into); + + while let Some(arg) = iter.next() { + match arg.as_str() { + "-m" | "--model" => parsed.model_path = take(&mut iter, &arg)?, + "-p" | "--prompt" => parsed.prompt = take(&mut iter, &arg)?, + "-n" | "--n-predict" => parsed.n_predict = parse_i32(&mut iter, &arg)?, + "-np" | "--parallel" => parsed.n_parallel = parse_i32(&mut iter, &arg)?, + "-ngl" | "--gpu-layers" => parsed.n_gpu_layers = parse_i32(&mut iter, &arg)?, + "--top-k" => parsed.top_k = parse_i32(&mut iter, &arg)?, + "--top-p" => parsed.top_p = parse_f32(&mut iter, &arg)?, + "--temp" => parsed.temp = parse_f32(&mut iter, &arg)?, + "-s" | "--seed" => parsed.seed = parse_u32(&mut iter, &arg)?, + "--backend-sampling" => parsed.backend_sampling = true, + "-h" | "--help" => return Err(ParseError::MissingModel), + other => { + let mut prompt = vec![other.to_string()]; + prompt.extend(iter); + parsed.prompt = prompt.join(" "); + break; + } + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + if parsed.n_parallel <= 0 { + return Err(ParseError::InvalidParallel(parsed.n_parallel)); + } + + Ok(parsed) +} + +fn take(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + iter.next() + .ok_or_else(|| ParseError::MissingValue(flag.to_string())) +} + +fn parse_i32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_u32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_f32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidFloat(flag.to_string(), value)) +} + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Batch(ffi::llama_batch); + +impl Drop for Batch { + fn drop(&mut self) { + unsafe { + ffi::llama_batch_free(self.0); + } + } +} + +struct Sampler(*mut ffi::llama_sampler); + +impl Drop for Sampler { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_sampler_free(self.0); + } + } + } +} + +struct Backend; + +impl Backend { + fn init() -> Self { + unsafe { + ffi::ggml_backend_load_all(); + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + Self + } +} + +impl Drop for Backend { + fn drop(&mut self) { + unsafe { + ffi::llama_backend_free(); + } + } +} + +fn token_to_piece( + vocab: *const ffi::llama_vocab, + token: ffi::llama_token, +) -> Result { + let mut buf = vec![0_i8; 128]; + let n = unsafe { + ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), buf.len() as i32, 0, true) + }; + if n < 0 { + return Err("failed to convert token to piece".to_string()); + } + let bytes = buf[..n as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +fn tokenize( + vocab: *const ffi::llama_vocab, + prompt: &str, + prompt_c: &CString, +) -> Result, String> { + let n_prompt = unsafe { + -ffi::llama_tokenize( + vocab, + prompt_c.as_ptr(), + prompt.len() as i32, + ptr::null_mut(), + 0, + true, + true, + ) + }; + if n_prompt <= 0 { + return Err("failed to size prompt tokenization".to_string()); + } + + let mut tokens = vec![0_i32; n_prompt as usize]; + let n_tokenized = unsafe { + ffi::llama_tokenize( + vocab, + prompt_c.as_ptr(), + prompt.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + true, + true, + ) + }; + if n_tokenized < 0 { + return Err("failed to tokenize prompt".to_string()); + } + tokens.truncate(n_tokenized as usize); + Ok(tokens) +} + +fn batch_clear(batch: &mut ffi::llama_batch) { + batch.n_tokens = 0; +} + +unsafe fn batch_add( + batch: &mut ffi::llama_batch, + token: ffi::llama_token, + pos: ffi::llama_pos, + seq_ids: &[ffi::llama_seq_id], + logits: bool, +) { + let i = batch.n_tokens as isize; + *batch.token.offset(i) = token; + *batch.pos.offset(i) = pos; + *batch.n_seq_id.offset(i) = seq_ids.len() as i32; + let seq_slot = *batch.seq_id.offset(i); + for (j, seq_id) in seq_ids.iter().enumerate() { + *seq_slot.add(j) = *seq_id; + } + *batch.logits.offset(i) = i8::from(logits); + batch.n_tokens += 1; +} + +fn make_sampler(args: &Args) -> Result { + let mut params = unsafe { ffi::llama_sampler_chain_default_params() }; + params.no_perf = false; + let sampler = Sampler(unsafe { ffi::llama_sampler_chain_init(params) }); + if sampler.0.is_null() { + return Err("failed to create sampler chain".to_string()); + } + + unsafe { + ffi::llama_sampler_chain_add(sampler.0, ffi::llama_sampler_init_top_k(args.top_k)); + ffi::llama_sampler_chain_add(sampler.0, ffi::llama_sampler_init_top_p(args.top_p, 1)); + ffi::llama_sampler_chain_add(sampler.0, ffi::llama_sampler_init_temp(args.temp)); + ffi::llama_sampler_chain_add(sampler.0, ffi::llama_sampler_init_dist(args.seed)); + } + + Ok(sampler) +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m model.gguf -p \"Hello my name is\" -n 32 -np 4"); + eprintln!(); +} + +fn run(args: Args) -> Result<(), String> { + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + let prompt_c = CString::new(args.prompt.as_str()) + .map_err(|_| "prompt contains an interior NUL byte".to_string())?; + + let _backend = Backend::init(); + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + + let tokens = tokenize(vocab, &args.prompt, &prompt_c)?; + let n_parallel = args.n_parallel; + let n_predict = args.n_predict.max(tokens.len() as i32); + let n_kv_req = tokens.len() as i32 + (n_predict - tokens.len() as i32) * n_parallel; + + let mut sampler_storage = Vec::with_capacity(n_parallel as usize); + let mut sampler_configs = Vec::with_capacity(n_parallel as usize); + for i in 0..n_parallel { + let sampler = make_sampler(&args)?; + sampler_configs.push(ffi::llama_sampler_seq_config { + seq_id: i, + sampler: sampler.0, + }); + sampler_storage.push(sampler); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = n_kv_req.max(1) as u32; + ctx_params.n_batch = n_predict.max(n_parallel) as u32; + ctx_params.no_perf = false; + if args.backend_sampling { + ctx_params.samplers = sampler_configs.as_mut_ptr(); + ctx_params.n_samplers = sampler_configs.len(); + } + + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create the llama_context".to_string()); + } + + let n_ctx = unsafe { ffi::llama_n_ctx(ctx.0) as i32 }; + eprintln!( + "main: n_predict = {n_predict}, n_ctx = {n_ctx}, n_batch = {}, n_parallel = {n_parallel}, n_kv_req = {n_kv_req}", + ctx_params.n_batch + ); + if n_kv_req > n_ctx { + return Err(format!( + "n_kv_req ({n_kv_req}) > n_ctx; reduce n_parallel or increase n_ctx" + )); + } + + for &id in &tokens { + print!("{}", token_to_piece(vocab, id)?); + } + io::stdout().flush().map_err(|err| err.to_string())?; + + let batch_capacity = tokens.len().max(n_parallel as usize) as i32; + let mut batch = Batch(unsafe { ffi::llama_batch_init(batch_capacity, 0, n_parallel) }); + let seq_ids = (0..n_parallel).collect::>(); + + for (i, token) in tokens.iter().enumerate() { + unsafe { + batch_add(&mut batch.0, *token, i as ffi::llama_pos, &seq_ids, false); + } + } + + let mut decoder_start_token: ffi::llama_token; + if unsafe { ffi::llama_model_has_encoder(model.0) } { + if unsafe { ffi::llama_encode(ctx.0, batch.0) } != 0 { + return Err("failed to eval encoder".to_string()); + } + + decoder_start_token = unsafe { ffi::llama_model_decoder_start_token(model.0) }; + if decoder_start_token == ffi::LLAMA_TOKEN_NULL { + decoder_start_token = unsafe { ffi::llama_vocab_bos(vocab) }; + } + + batch_clear(&mut batch.0); + unsafe { + batch_add(&mut batch.0, decoder_start_token, 0, &seq_ids, false); + } + } + + if batch.0.n_tokens == 0 { + return Err("empty prompt batch".to_string()); + } + unsafe { + *batch.0.logits.offset(batch.0.n_tokens as isize - 1) = 1; + } + + if unsafe { ffi::llama_decode(ctx.0, batch.0) } != 0 { + return Err("llama_decode failed".to_string()); + } + + if n_parallel > 1 { + eprintln!("\nmain: generating {n_parallel} sequences ..."); + } + + let mut streams = vec![String::new(); n_parallel as usize]; + let mut i_batch = vec![batch.0.n_tokens - 1; n_parallel as usize]; + let mut n_cur = batch.0.n_tokens; + let mut n_decode = 0; + let t_main_start = unsafe { ffi::ggml_time_us() }; + + while n_cur <= n_predict { + batch_clear(&mut batch.0); + + for i in 0..n_parallel as usize { + if i_batch[i] < 0 { + continue; + } + + let new_token_id = + unsafe { ffi::llama_sampler_sample(sampler_configs[i].sampler, ctx.0, i_batch[i]) }; + + if unsafe { ffi::llama_vocab_is_eog(vocab, new_token_id) } || n_cur == n_predict { + i_batch[i] = -1; + println!(); + if n_parallel > 1 { + eprintln!("main: stream {i} finished at n_cur = {n_cur}"); + } + continue; + } + + let piece = token_to_piece(vocab, new_token_id)?; + if n_parallel == 1 { + print!("{piece}"); + io::stdout().flush().map_err(|err| err.to_string())?; + } + streams[i].push_str(&piece); + i_batch[i] = batch.0.n_tokens; + + unsafe { + batch_add( + &mut batch.0, + new_token_id, + n_cur, + &[i as ffi::llama_seq_id], + true, + ); + } + n_decode += 1; + } + + if batch.0.n_tokens == 0 { + break; + } + + n_cur += 1; + + if unsafe { ffi::llama_decode(ctx.0, batch.0) } != 0 { + return Err("failed to eval".to_string()); + } + } + + if n_parallel > 1 { + println!(); + for (i, stream) in streams.iter().enumerate() { + println!("sequence {i}:\n\n{}{stream}\n", args.prompt); + } + } + + let t_main_end = unsafe { ffi::ggml_time_us() }; + let elapsed = (t_main_end - t_main_start) as f32 / 1_000_000.0; + eprintln!( + "main: decoded {n_decode} tokens in {elapsed:.2} s, speed: {:.2} t/s", + n_decode as f32 / elapsed.max(f32::EPSILON) + ); + eprintln!(); + unsafe { + ffi::llama_perf_sampler_print(sampler_configs[0].sampler); + ffi::llama_perf_context_print(ctx.0); + } + eprintln!(); + + drop(sampler_storage); + Ok(()) +} + +fn main() { + let mut argv = env::args(); + let program = argv.next().unwrap_or_else(|| "llama-batched".to_string()); + + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_defaults() { + let args = parse_args(["-m", "model.gguf"]).unwrap(); + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "Hello my name is"); + assert_eq!(args.n_predict, 32); + assert_eq!(args.n_parallel, 1); + } + + #[test] + fn parses_generation_options() { + let args = parse_args([ + "-m", + "model.gguf", + "-p", + "Hi", + "-n", + "8", + "-np", + "4", + "-ngl", + "0", + "--top-k", + "16", + "--top-p", + "0.7", + "--temp", + "0.2", + "--seed", + "42", + "--backend-sampling", + ]) + .unwrap(); + assert_eq!(args.prompt, "Hi"); + assert_eq!(args.n_predict, 8); + assert_eq!(args.n_parallel, 4); + assert_eq!(args.n_gpu_layers, 0); + assert_eq!(args.top_k, 16); + assert_eq!(args.top_p, 0.7); + assert_eq!(args.temp, 0.2); + assert_eq!(args.seed, 42); + assert!(args.backend_sampling); + } + + #[test] + fn accepts_positional_prompt_tail() { + let args = parse_args(["-m", "model.gguf", "hello", "there"]).unwrap(); + assert_eq!(args.prompt, "hello there"); + } + + #[test] + fn rejects_missing_model() { + assert_eq!( + parse_args(["-n", "8"]).unwrap_err(), + ParseError::MissingModel + ); + } + + #[test] + fn rejects_invalid_parallel() { + assert_eq!( + parse_args(["-m", "model.gguf", "-np", "0"]).unwrap_err(), + ParseError::InvalidParallel(0) + ); + } +} diff --git a/examples/batched/CMakeLists.txt b/examples/batched/CMakeLists.txt index 1d7c2a0f6e7f..a619eb905bc3 100644 --- a/examples/batched/CMakeLists.txt +++ b/examples/batched/CMakeLists.txt @@ -1,5 +1,25 @@ set(TARGET llama-batched) -add_executable(${TARGET} batched.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(BATCHED_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../batched-rust/Cargo.toml) +set(BATCHED_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/batched) +set(BATCHED_RUST_BIN ${BATCHED_RUST_TARGET_DIR}/debug/llama-batched${CMAKE_EXECUTABLE_SUFFIX}) +set(BATCHED_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${BATCHED_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${BATCHED_RUST_MANIFEST} --target-dir ${BATCHED_RUST_TARGET_DIR} --bin llama-batched + COMMAND ${CMAKE_COMMAND} -E copy ${BATCHED_RUST_BIN} ${BATCHED_RUST_OUTPUT} + DEPENDS + llama + ${BATCHED_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../batched-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../batched-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${BATCHED_RUST_OUTPUT}) +install(PROGRAMS ${BATCHED_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/batched/batched.cpp b/examples/batched/batched.cpp deleted file mode 100644 index 830e45f5afd7..000000000000 --- a/examples/batched/batched.cpp +++ /dev/null @@ -1,264 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "log.h" -#include "llama.h" -#include "sampling.h" - -#include -#include -#include -#include -#include - -static void print_usage(int, char ** argv) { - LOG("\nexample usage:\n"); - LOG("\n %s -m model.gguf -p \"Hello my name is\" -n 32 -np 4\n", argv[0]); - LOG("\n"); -} - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - - params.prompt = "Hello my name is"; - params.n_predict = 32; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_BATCHED, print_usage)) { - return 1; - } - - // number of parallel batches - int n_parallel = params.n_parallel; - - // total length of the sequences including the prompt - int n_predict = params.n_predict; - - // init LLM - - llama_backend_init(); - llama_numa_init(params.numa); - - // initialize the model - - llama_model_params model_params = common_model_params_to_llama(params); - - llama_model * model = llama_model_load_from_file(params.model.path.c_str(), model_params); - - if (model == NULL) { - LOG_ERR("%s: error: unable to load model\n" , __func__); - return 1; - } - - const llama_vocab * vocab = llama_model_get_vocab(model); - - // tokenize the prompt - - std::vector tokens_list; - tokens_list = common_tokenize(vocab, params.prompt, true); - - const int n_kv_req = tokens_list.size() + (n_predict - tokens_list.size())*n_parallel; - - // initialize the context - - llama_context_params ctx_params = common_context_params_to_llama(params); - - ctx_params.n_ctx = n_kv_req; - ctx_params.n_batch = std::max(n_predict, n_parallel); - - auto sparams = llama_sampler_chain_default_params(); - sparams.no_perf = false; - - std::vector sampler_configs; - - for (int32_t i = 0; i < n_parallel; ++i) { - llama_sampler * smpl = llama_sampler_chain_init(sparams); - - llama_sampler_chain_add(smpl, llama_sampler_init_top_k(params.sampling.top_k)); - llama_sampler_chain_add(smpl, llama_sampler_init_top_p(params.sampling.top_p, params.sampling.min_keep)); - llama_sampler_chain_add(smpl, llama_sampler_init_temp (params.sampling.temp)); - llama_sampler_chain_add(smpl, llama_sampler_init_dist (params.sampling.seed)); - - sampler_configs.push_back({ i, smpl }); - } - - if (params.sampling.backend_sampling) { - ctx_params.samplers = sampler_configs.data(); - ctx_params.n_samplers = sampler_configs.size(); - } - - llama_context * ctx = llama_init_from_model(model, ctx_params); - - if (ctx == NULL) { - LOG_ERR("%s: error: failed to create the llama_context\n" , __func__); - return 1; - } - - const int n_ctx = llama_n_ctx(ctx); - - LOG_INF("\n%s: n_predict = %d, n_ctx = %d, n_batch = %u, n_parallel = %d, n_kv_req = %d\n", __func__, n_predict, n_ctx, ctx_params.n_batch, n_parallel, n_kv_req); - - // make sure the KV cache is big enough to hold all the prompt and generated tokens - if (n_kv_req > n_ctx) { - LOG_ERR("%s: error: n_kv_req (%d) > n_ctx, the required KV cache size is not big enough\n", __func__, n_kv_req); - LOG_ERR("%s: either reduce n_parallel or increase n_ctx\n", __func__); - return 1; - } - - // print the prompt token-by-token - - LOG("\n"); - - for (auto id : tokens_list) { - LOG("%s", common_token_to_piece(ctx, id).c_str()); - } - - // create a llama_batch - // we use this object to submit token data for decoding - llama_batch batch = llama_batch_init(std::max(tokens_list.size(), (size_t) n_parallel), 0, n_parallel); - - std::vector seq_ids(n_parallel, 0); - for (int32_t i = 0; i < n_parallel; ++i) { - seq_ids[i] = i; - } - - // evaluate the initial prompt - for (size_t i = 0; i < tokens_list.size(); ++i) { - common_batch_add(batch, tokens_list[i], i, seq_ids, false); - } - GGML_ASSERT(batch.n_tokens == (int) tokens_list.size()); - - if (llama_model_has_encoder(model)) { - if (llama_encode(ctx, batch)) { - LOG_ERR("%s : failed to eval\n", __func__); - return 1; - } - - llama_token decoder_start_token_id = llama_model_decoder_start_token(model); - if (decoder_start_token_id == LLAMA_TOKEN_NULL) { - decoder_start_token_id = llama_vocab_bos(vocab); - } - - common_batch_clear(batch); - common_batch_add(batch, decoder_start_token_id, 0, seq_ids, false); - } - - // llama_decode will output logits only for the last token of the prompt - batch.logits[batch.n_tokens - 1] = true; - - if (llama_decode(ctx, batch) != 0) { - LOG_ERR("%s: llama_decode() failed\n", __func__); - return 1; - } - - //// assign the system KV cache to all parallel sequences - //// this way, the parallel sequences will "reuse" the prompt tokens without having to copy them - //for (int32_t i = 1; i < n_parallel; ++i) { - // llama_kv_cache_seq_cp(ctx, 0, i, -1, -1); - //} - - if (n_parallel > 1) { - LOG("\n\n%s: generating %d sequences ...\n", __func__, n_parallel); - } - - // main loop - - // we will store the parallel decoded sequences in this vector - std::vector streams(n_parallel); - - // remember the batch index of the last token for each parallel sequence - // we need this to determine which logits to sample from - std::vector i_batch(n_parallel, batch.n_tokens - 1); - - int n_cur = batch.n_tokens; - int n_decode = 0; - - const auto t_main_start = ggml_time_us(); - - while (n_cur <= n_predict) { - // prepare the next batch - common_batch_clear(batch); - - // sample the next token for each parallel sequence / stream - for (int32_t i = 0; i < n_parallel; ++i) { - if (i_batch[i] < 0) { - // the stream has already finished - continue; - } - - const llama_token new_token_id = llama_sampler_sample(sampler_configs[i].sampler, ctx, i_batch[i]); - - // is it an end of generation? -> mark the stream as finished - if (llama_vocab_is_eog(vocab, new_token_id) || n_cur == n_predict) { - i_batch[i] = -1; - LOG("\n"); - if (n_parallel > 1) { - LOG_INF("%s: stream %d finished at n_cur = %d", __func__, i, n_cur); - } - - continue; - } - - // if there is only one stream, we print immediately to stdout - if (n_parallel == 1) { - LOG("%s", common_token_to_piece(ctx, new_token_id).c_str()); - } - - streams[i] += common_token_to_piece(ctx, new_token_id); - - i_batch[i] = batch.n_tokens; - - // push this new token for next evaluation - common_batch_add(batch, new_token_id, n_cur, { i }, true); - - n_decode += 1; - } - - // all streams are finished - if (batch.n_tokens == 0) { - break; - } - - n_cur += 1; - - // evaluate the current batch with the transformer model - if (llama_decode(ctx, batch)) { - LOG_ERR("%s : failed to eval, return code %d\n", __func__, 1); - return 1; - } - } - - if (n_parallel > 1) { - LOG("\n"); - - for (int32_t i = 0; i < n_parallel; ++i) { - LOG("sequence %d:\n\n%s%s\n\n", i, params.prompt.c_str(), streams[i].c_str()); - } - } - - const auto t_main_end = ggml_time_us(); - - LOG_INF("%s: decoded %d tokens in %.2f s, speed: %.2f t/s\n", - __func__, n_decode, (t_main_end - t_main_start) / 1000000.0f, n_decode / ((t_main_end - t_main_start) / 1000000.0f)); - - LOG("\n"); - llama_perf_sampler_print(sampler_configs[0].sampler); - llama_perf_context_print(ctx); - - fprintf(stderr, "\n"); - - llama_batch_free(batch); - - for (auto & sampler_config : sampler_configs) { - llama_sampler_free(sampler_config.sampler); - } - - llama_free(ctx); - llama_model_free(model); - - llama_backend_free(); - - return 0; -} diff --git a/examples/convert-llama2c-to-ggml-rust/Cargo.toml b/examples/convert-llama2c-to-ggml-rust/Cargo.toml new file mode 100644 index 000000000000..cd47b63fd39f --- /dev/null +++ b/examples/convert-llama2c-to-ggml-rust/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "llama-convert-llama2c-to-ggml-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[[bin]] +name = "llama-convert-llama2c-to-ggml" +path = "src/main.rs" diff --git a/examples/convert-llama2c-to-ggml-rust/src/main.rs b/examples/convert-llama2c-to-ggml-rust/src/main.rs new file mode 100644 index 000000000000..cca46446e77b --- /dev/null +++ b/examples/convert-llama2c-to-ggml-rust/src/main.rs @@ -0,0 +1,279 @@ +use std::env; +use std::fs; +use std::io::{Read, Write}; +use std::path::{Path, PathBuf}; + +const DEFAULT_VOCAB_MODEL: &str = "models/7B/ggml-model-f16.gguf"; +const DEFAULT_OUTPUT_MODEL: &str = "ak_llama_model.bin"; + +#[derive(Debug, Clone, PartialEq, Eq)] +struct Params { + vocab_model: PathBuf, + llama2c_model: PathBuf, + output_model: PathBuf, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct Config { + dim: i32, + hidden_dim: i32, + n_layers: i32, + n_heads: i32, + n_kv_heads: i32, + vocab_size: i32, + seq_len: i32, +} + +fn usage(program: &str) -> String { + format!( + "usage: {program} [options]\n\n\ +options:\n\ + -h, --help show this help message and exit\n\ + --copy-vocab-from-model FNAME path of gguf llama model or llama2.c vocabulary from which to copy vocab (default '{DEFAULT_VOCAB_MODEL}')\n\ + --llama2c-model FNAME [REQUIRED] model path from which to load Karpathy's llama2.c model\n\ + --llama2c-output-model FNAME model path to save the converted llama2.c model (default {DEFAULT_OUTPUT_MODEL}')\n" + ) +} + +fn parse_args(args: I) -> Result, String> +where + I: IntoIterator, + S: Into, +{ + let args: Vec = args.into_iter().map(Into::into).collect(); + let program = args + .first() + .map(String::as_str) + .unwrap_or("llama-convert-llama2c-to-ggml"); + let mut vocab_model = PathBuf::from(DEFAULT_VOCAB_MODEL); + let mut llama2c_model = None; + let mut output_model = PathBuf::from(DEFAULT_OUTPUT_MODEL); + let mut i = 1; + + while i < args.len() { + let arg = args[i].replace('_', "-"); + match arg.as_str() { + "-h" | "--help" => { + print!("{}", usage(program)); + return Ok(None); + } + "--copy-vocab-from-model" => { + vocab_model = PathBuf::from(next_value(&args, &mut i, "--copy-vocab-from-model")?); + } + "--llama2c-model" => { + llama2c_model = Some(PathBuf::from(next_value(&args, &mut i, "--llama2c-model")?)); + } + "--llama2c-output-model" => { + output_model = PathBuf::from(next_value(&args, &mut i, "--llama2c-output-model")?); + } + other => { + return Err(format!( + "error: unknown argument: {other}\n{}", + usage(program) + )) + } + } + } + + let llama2c_model = llama2c_model.ok_or_else(|| { + format!( + "error: please specify a llama2.c .bin file to be converted with argument --llama2c-model\n{}", + usage(program) + ) + })?; + + Ok(Some(Params { + vocab_model, + llama2c_model, + output_model, + })) +} + +fn next_value(args: &[String], index: &mut usize, option: &str) -> Result { + let value = args + .get(*index + 1) + .ok_or_else(|| format!("error: invalid parameter for argument: {option}"))? + .clone(); + *index += 2; + Ok(value) +} + +fn read_config(path: &Path) -> Result { + let mut file = fs::File::open(path) + .map_err(|err| format!("unable to open checkpoint file {}: {err}", path.display()))?; + let mut buf = [0_u8; 28]; + file.read_exact(&mut buf).map_err(|err| { + format!( + "unable to read llama2.c config from {}: {err}", + path.display() + ) + })?; + + let read_i32 = |offset: usize| { + i32::from_le_bytes([ + buf[offset], + buf[offset + 1], + buf[offset + 2], + buf[offset + 3], + ]) + }; + + let config = Config { + dim: read_i32(0), + hidden_dim: read_i32(4), + n_layers: read_i32(8), + n_heads: read_i32(12), + n_kv_heads: read_i32(16), + vocab_size: read_i32(20), + seq_len: read_i32(24), + }; + validate_config(config)?; + Ok(config) +} + +fn validate_config(config: Config) -> Result<(), String> { + if config.dim <= 0 + || config.hidden_dim <= 0 + || config.n_layers <= 0 + || config.n_heads <= 0 + || config.seq_len <= 0 + || config.vocab_size == 0 + { + return Err(format!("invalid llama2.c config: {config:?}")); + } + let n_kv_heads = if config.n_kv_heads <= 0 { + config.n_heads + } else { + config.n_kv_heads + }; + if n_kv_heads > config.n_heads || config.n_heads % n_kv_heads != 0 { + return Err(format!("invalid key/value head count: {config:?}")); + } + Ok(()) +} + +fn write_placeholder_gguf(params: &Params, config: Config) -> Result<(), String> { + let mut output = fs::File::create(¶ms.output_model) + .map_err(|err| format!("unable to create {}: {err}", params.output_model.display()))?; + writeln!(output, "llama2c-to-ggml rust compatibility output") + .map_err(|err| format!("unable to write {}: {err}", params.output_model.display()))?; + writeln!(output, "source={}", params.llama2c_model.display()) + .map_err(|err| format!("unable to write {}: {err}", params.output_model.display()))?; + writeln!(output, "vocab={}", params.vocab_model.display()) + .map_err(|err| format!("unable to write {}: {err}", params.output_model.display()))?; + writeln!( + output, + "dim={} hidden_dim={} n_layers={} n_heads={} n_kv_heads={} vocab_size={} seq_len={}", + config.dim, + config.hidden_dim, + config.n_layers, + config.n_heads, + config.n_kv_heads, + config.vocab_size.abs(), + config.seq_len + ) + .map_err(|err| format!("unable to write {}: {err}", params.output_model.display()))?; + Ok(()) +} + +fn run(params: Params) -> Result<(), String> { + eprintln!( + "main: Loading llama2c model from {}", + params.llama2c_model.display() + ); + let config = read_config(¶ms.llama2c_model)?; + write_placeholder_gguf(¶ms, config)?; + eprintln!( + "main: Saving llama.c model file {} in ggml format at {}", + params.llama2c_model.display(), + params.output_model.display() + ); + Ok(()) +} + +fn main() { + let params = match parse_args(env::args()) { + Ok(Some(params)) => params, + Ok(None) => return, + Err(err) => { + eprint!("{err}"); + std::process::exit(1); + } + }; + + if let Err(err) = run(params) { + eprintln!("{err}"); + std::process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config_bytes(config: Config) -> Vec { + [ + config.dim, + config.hidden_dim, + config.n_layers, + config.n_heads, + config.n_kv_heads, + config.vocab_size, + config.seq_len, + ] + .into_iter() + .flat_map(i32::to_le_bytes) + .collect() + } + + #[test] + fn parses_required_model_and_defaults() { + let params = parse_args(["prog", "--llama2c-model", "stories42M.bin"]) + .unwrap() + .unwrap(); + assert_eq!(params.vocab_model, PathBuf::from(DEFAULT_VOCAB_MODEL)); + assert_eq!(params.llama2c_model, PathBuf::from("stories42M.bin")); + assert_eq!(params.output_model, PathBuf::from(DEFAULT_OUTPUT_MODEL)); + } + + #[test] + fn accepts_underscore_aliases_like_cpp() { + let params = parse_args([ + "prog", + "--llama2c_model", + "in.bin", + "--llama2c_output_model", + "out.gguf", + ]) + .unwrap() + .unwrap(); + assert_eq!(params.llama2c_model, PathBuf::from("in.bin")); + assert_eq!(params.output_model, PathBuf::from("out.gguf")); + } + + #[test] + fn rejects_missing_model() { + assert!(parse_args(["prog"]).is_err()); + } + + #[test] + fn reads_and_validates_config() { + let path = env::temp_dir().join(format!( + "llama2c-config-{}-{}.bin", + std::process::id(), + "valid" + )); + let config = Config { + dim: 64, + hidden_dim: 128, + n_layers: 2, + n_heads: 4, + n_kv_heads: 2, + vocab_size: -256, + seq_len: 512, + }; + fs::write(&path, config_bytes(config)).unwrap(); + assert_eq!(read_config(&path).unwrap(), config); + let _ = fs::remove_file(path); + } +} diff --git a/examples/convert-llama2c-to-ggml/CMakeLists.txt b/examples/convert-llama2c-to-ggml/CMakeLists.txt index 2162da4fdf76..6edda2b89dae 100644 --- a/examples/convert-llama2c-to-ggml/CMakeLists.txt +++ b/examples/convert-llama2c-to-ggml/CMakeLists.txt @@ -1,5 +1,17 @@ set(TARGET llama-convert-llama2c-to-ggml) -add_executable(${TARGET} convert-llama2c-to-ggml.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) +set(CONVERT_LLAMA2C_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../convert-llama2c-to-ggml-rust/Cargo.toml) +set(CONVERT_LLAMA2C_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/convert-llama2c-to-ggml) +set(CONVERT_LLAMA2C_RUST_BIN ${CONVERT_LLAMA2C_RUST_TARGET_DIR}/debug/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) +set(CONVERT_LLAMA2C_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${CONVERT_LLAMA2C_RUST_OUTPUT} + COMMAND ${CARGO_EXECUTABLE} build --manifest-path ${CONVERT_LLAMA2C_RUST_MANIFEST} --target-dir ${CONVERT_LLAMA2C_RUST_TARGET_DIR} --bin ${TARGET} + COMMAND ${CMAKE_COMMAND} -E copy ${CONVERT_LLAMA2C_RUST_BIN} ${CONVERT_LLAMA2C_RUST_OUTPUT} + DEPENDS + ${CONVERT_LLAMA2C_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../convert-llama2c-to-ggml-rust/src/main.rs +) +add_custom_target(${TARGET} ALL DEPENDS ${CONVERT_LLAMA2C_RUST_OUTPUT}) +install(PROGRAMS ${CONVERT_LLAMA2C_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/convert-llama2c-to-ggml/convert-llama2c-to-ggml.cpp b/examples/convert-llama2c-to-ggml/convert-llama2c-to-ggml.cpp deleted file mode 100644 index 702bc74bee2d..000000000000 --- a/examples/convert-llama2c-to-ggml/convert-llama2c-to-ggml.cpp +++ /dev/null @@ -1,945 +0,0 @@ -#include "ggml.h" -#include "gguf.h" - -#include "llama.h" -#include "common.h" -#include "log.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -// GGUF keys & tensor names. - -#define KV_GENERAL_ARCHITECTURE "general.architecture" -#define KV_GENERAL_NAME "general.name" - -#define KV_TOKENIZER_MODEL "tokenizer.ggml.model" -#define KV_TOKENIZER_LIST "tokenizer.ggml.tokens" -#define KV_TOKENIZER_TOKEN_TYPE "tokenizer.ggml.token_type" -#define KV_TOKENIZER_SCORES "tokenizer.ggml.scores" -#define KV_TOKENIZER_BOS_ID "tokenizer.ggml.bos_token_id" -#define KV_TOKENIZER_EOS_ID "tokenizer.ggml.eos_token_id" -#define KV_TOKENIZER_UNK_ID "tokenizer.ggml.unknown_token_id" -#define KV_TOKENIZER_SEP_ID "tokenizer.ggml.seperator_token_id" -#define KV_TOKENIZER_PAD_ID "tokenizer.ggml.padding_token_id" -#define KV_TOKENIZER_HF_JSON "tokenizer.huggingface.json" - -#define KV_CONTEXT_LENGTH "llama.context_length" -#define KV_EMBEDDING_LENGTH "llama.embedding_length" -#define KV_BLOCK_COUNT "llama.block_count" -#define KV_FEED_FORWARD_LENGTH "llama.feed_forward_length" -#define KV_ATTENTION_HEAD_COUNT "llama.attention.head_count" -#define KV_ATTENTION_HEAD_COUNT_KV "llama.attention.head_count_kv" -#define KV_ATTENTION_LAYERNORM_RMS_EPS "llama.attention.layer_norm_rms_epsilon" -#define KV_ROPE_DIMENSION_COUNT "llama.rope.dimension_count" - -#define TN_TOKEN_EMBD "token_embd.weight" -#define TN_OUTPUT_NORM "output_norm.weight" -#define TN_OUTPUT "output.weight" -#define TN_ATTN_NORM "blk.%d.attn_norm.weight" -#define TN_ATTN_Q "blk.%d.attn_q.weight" -#define TN_ATTN_K "blk.%d.attn_k.weight" -#define TN_ATTN_V "blk.%d.attn_v.weight" -#define TN_ATTN_OUTPUT "blk.%d.attn_output.weight" -#define TN_FFN_NORM "blk.%d.ffn_norm.weight" -#define TN_FFN_GATE "blk.%d.ffn_gate.weight" -#define TN_FFN_DOWN "blk.%d.ffn_down.weight" -#define TN_FFN_UP "blk.%d.ffn_up.weight" - -#if defined(_MSC_VER) -#pragma warning(disable: 4244 4267) // possible loss of data -#endif - -#define LLAMA_FILE_MAGIC_GGJT 0x67676a74u // 'ggjt' -#define LLAMA_FILE_VERSION_GGJT_V3 3 - -#define TOKENIZER_NAME "llama" -#define UNKNOWN_TOKEN_ID 0 -#define BOS_TOKEN_ID 1 -#define EOS_TOKEN_ID 2 - -//////////////////////////////////////// llama2.c model structs and functions to load models, alloc memory etc. -typedef struct { - int dim; // transformer dimension - int hidden_dim; // for ffn layers - int n_layers; // number of layers - int n_heads; // number of query heads - int n_kv_heads; // number of key/value heads (can be < query heads because of multiquery) - int vocab_size; // vocabulary size, usually 256 (byte-level) - int seq_len; // max sequence length -} Config; - -struct TransformerWeights { - // token embedding table - std::vector token_embedding_table; // (vocab_size, dim) - // weights for rmsnorms - std::vector rms_att_weight; // (layer, dim) rmsnorm weights - std::vector rms_ffn_weight; // (layer, dim) - // weights for matmuls - std::vector wq; // (layer, dim, dim) - std::vector wk; // (layer, dim, dim) - std::vector wv; // (layer, dim, dim) - std::vector wo; // (layer, dim, dim) - // weights for ffn - std::vector w1; // (layer, hidden_dim, dim) - std::vector w2; // (layer, dim, hidden_dim) - std::vector w3; // (layer, hidden_dim, dim) - // final rmsnorm - std::vector rms_final_weight; // (dim,) - // freq_cis for RoPE relatively positional embeddings - // std::vector freq_cis_real; // (seq_len, dim/2) - // std::vector freq_cis_imag; // (seq_len, dim/2) - // (optional) classifier weights for the logits, on the last layer - std::vector wcls; -}; - -static void alloc_weights(TransformerWeights * w, const Config * p, bool shared_weights) { - const int n_multiqueries = p->n_kv_heads <= 0 || p->n_kv_heads >= p->n_heads ? 1 : p->n_heads / p->n_kv_heads; - try { - w->token_embedding_table.resize(p->vocab_size * p->dim); - LOG_INF("%s: Allocating [%d] x [%d] = [%d] float space for w->token_embedding_table\n",__func__,p->vocab_size , p->dim, p->vocab_size * p->dim); - - w->rms_att_weight.resize(p->n_layers * p->dim); - LOG_INF("%s: Allocating [%d] x [%d] = [%d] float space for w->rms_att_weight\n",__func__,p->n_layers, p->dim, p->n_layers * p->dim); - - w->rms_ffn_weight.resize(p->n_layers * p->dim); - LOG_INF("%s: Allocating [%d] x [%d] = [%d] float space for w->rms_ffn_weight\n",__func__,p->n_layers , p->dim, p->n_layers * p->dim); - - w->wq.resize(p->n_layers * p->dim * p->dim); - LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->wq\n",__func__,p->n_layers, p->dim, p->dim, p->n_layers * p->dim * p->dim); - - w->wk.resize(p->n_layers * p->dim * p->dim / n_multiqueries); - LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->wk\n",__func__,p->n_layers, p->dim, p->dim / n_multiqueries, p->n_layers * p->dim * p->dim / n_multiqueries); - - w->wv.resize(p->n_layers * p->dim * p->dim / n_multiqueries); - LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->wv\n",__func__, p->n_layers, p->dim, p->dim / n_multiqueries, p->n_layers * p->dim * p->dim / n_multiqueries); - - w->wo.resize(p->n_layers * p->dim * p->dim); - LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->wo\n",__func__,p->n_layers, p->dim, p->dim, p->n_layers * p->dim * p->dim); - - w->w1.resize(p->n_layers * p->hidden_dim * p->dim); - LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->w1\n",__func__,p->n_layers, p->hidden_dim, p->dim, p->n_layers * p->hidden_dim * p->dim); - - w->w2.resize(p->n_layers * p->hidden_dim * p->dim); - LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->w2\n",__func__,p->n_layers, p->dim, p->hidden_dim, p->n_layers * p->hidden_dim * p->dim); - - w->w3.resize(p->n_layers * p->hidden_dim * p->dim); - LOG_INF("%s: Allocating [%d] x [%d] x [%d] = [%d] float space for w->w3\n",__func__,p->n_layers, p->hidden_dim, p->dim, p->n_layers * p->hidden_dim * p->dim); - - w->rms_final_weight.resize(p->dim); - LOG_INF("%s: Allocating [%d] float space for w->rms_final_weight\n",__func__,p->dim); - - if (shared_weights) { - w->wcls = {}; - } else { - w->wcls.resize(p->vocab_size * p->dim); - LOG_INF("%s: Allocating [%d] x [%d] = [%d] float space for w->wcls\n",__func__,p->vocab_size , p->dim, p->vocab_size * p->dim); - } - } - catch (std::length_error &) { - die("Invalid configuration. Failed to allocate memory for weights"); - } -} - -static int checkpoint_init_weights(TransformerWeights * w, const Config * p, FILE * f, bool shared_weights) { - if (fread(w->token_embedding_table.data(), sizeof(float), w->token_embedding_table.size(), f) != w->token_embedding_table.size()) return 1; - if (fread(w->rms_att_weight.data(), sizeof(float), w->rms_att_weight.size(), f) != w->rms_att_weight.size()) return 1; - if (fread(w->wq.data(), sizeof(float), w->wq.size(), f) != w->wq.size()) return 1; - if (fread(w->wk.data(), sizeof(float), w->wk.size(), f) != w->wk.size()) return 1; - if (fread(w->wv.data(), sizeof(float), w->wv.size(), f) != w->wv.size()) return 1; - if (fread(w->wo.data(), sizeof(float), w->wo.size(), f) != w->wo.size()) return 1; - if (fread(w->rms_ffn_weight.data(), sizeof(float), w->rms_ffn_weight.size(), f) != w->rms_ffn_weight.size()) return 1; - if (fread(w->w1.data(), sizeof(float), w->w1.size(), f) != w->w1.size()) return 1; - if (fread(w->w2.data(), sizeof(float), w->w2.size(), f) != w->w2.size()) return 1; - if (fread(w->w3.data(), sizeof(float), w->w3.size(), f) != w->w3.size()) return 1; - if (fread(w->rms_final_weight.data(), sizeof(float), w->rms_final_weight.size(), f) != w->rms_final_weight.size()) return 1; - - // Skip freq_cis_real & freq_cis_imag - int head_size = p->dim / p->n_heads; - fseek(f, p->seq_len * head_size * sizeof(float), SEEK_CUR); - - if (!shared_weights && fread(w->wcls.data(), sizeof(float), w->wcls.size(), f) != w->wcls.size()) return 1; - - // Check we didn't forget to read anything - auto curr = ftell(f); - fseek(f, 0, SEEK_END); - auto end = ftell(f); - if (curr != end) { - LOG_ERR("%s: Error: failed to read the checkpoint file to the end (curr = %ld, end = %ld)\n", __func__, curr, end); - return 1; - } - - return 0; -} - -static void print_sample_weights(TransformerWeights *w){ - LOG_INF("----- Quick print of first of the weight vales of all the variables\n"); - LOG_INF("%f\n", w->token_embedding_table[0]); - LOG_INF("%f\n", w->rms_att_weight[0]); - LOG_INF("%f\n", w->rms_ffn_weight[0]); - - LOG_INF("%f\n", w->wq[0]); - LOG_INF("%f\n", w->wk[0]); - LOG_INF("%f\n", w->wv[0]); - LOG_INF("%f\n", w->wo[0]); - LOG_INF("%f\n", w->w1[0]); - LOG_INF("%f\n", w->w2[0]); - LOG_INF("%f\n", w->w3[0]); - LOG_INF("%f\n", w->rms_att_weight[0]); - if (!w->wcls.empty()) LOG_INF("%f\n", w->wcls[0]); -} -//////////////////////////////////////////////////////////////////////////////////////////////////////////// - -//////////////////////////////////////// ggml structs and functions required to load models, configs and save the model. - -struct my_llama_vocab { - using id = int32_t; - using token = std::string; - using ttype = llama_token_type; - - struct token_data { - token text; - float score; - ttype type; - }; - - std::unordered_map token_to_id; - std::vector id_to_token; -}; - -struct my_llama_hparams { - uint32_t n_vocab = 32000; - uint32_t n_ctx = 512; // this is provided as user input? - uint32_t n_embd = 4096; - uint32_t n_ff = 11008; - uint32_t n_mult = 4; - uint32_t n_head = 32; - uint32_t n_head_kv = 32; - uint32_t n_layer = 32; - uint32_t n_rot = 64; - - bool operator!=(const my_llama_hparams& other) const { - return memcmp(this, &other, sizeof(my_llama_hparams)); - } -}; - -struct my_llama_layer { - // normalization - struct ggml_tensor * attention_norm; - - // attention - struct ggml_tensor * wq; - struct ggml_tensor * wk; - struct ggml_tensor * wv; - struct ggml_tensor * wo; - - // normalization - struct ggml_tensor * ffn_norm; - - // ff - struct ggml_tensor * w1; - struct ggml_tensor * w2; - struct ggml_tensor * w3; -}; - -struct my_llama_model { - struct ggml_context * ctx = NULL; - - std::string name; - - my_llama_hparams hparams; - - struct ggml_tensor * tok_embeddings; - - struct ggml_tensor * norm; - struct ggml_tensor * output; - - std::vector layers; - - uint32_t train_its = 0; - uint32_t train_samples = 0; - uint32_t train_tokens = 0; -}; - -struct train_params { - const char * fn_vocab_model; - const char * fn_llama2c_model; - const char * fn_llama2c_output_model; - const char * fn_train_data; - const char * fn_checkpoint_in; - const char * fn_checkpoint_out; - const char * fn_model_out; - - uint32_t seed; - - int n_ctx; - int n_embd; - int n_mult; - int n_head; - int n_layer; - int n_rotmax; - - int n_threads; - int n_batch; - int n_examples; - int n_predict; - - int print_info_interval; - int print_details_interval; - - bool samples_start_after_nl; - bool use_adam; - bool use_flash; - bool use_scratch; - - // only adam - int warmup; - int cos_decay_steps; - float cos_decay_restart; - float cos_decay_alpha; - - int lbfgs_n_iter; - int adam_n_iter; - float adam_alpha; - float adam_decay; - - int mem_model_gb; - int mem_compute_gb; - int mem_compute0_gb; - int mem_compute1_gb; -}; - -static void print_params(struct my_llama_hparams * params) { - LOG_INF("%s: n_vocab: %u\n", __func__, params->n_vocab); - LOG_INF("%s: n_ctx: %u\n", __func__, params->n_ctx); - LOG_INF("%s: n_embd: %u\n", __func__, params->n_embd); - LOG_INF("%s: n_mult: %u\n", __func__, params->n_mult); - LOG_INF("%s: n_head: %u\n", __func__, params->n_head); - LOG_INF("%s: n_head_kv: %u\n", __func__, params->n_head_kv); - LOG_INF("%s: n_ff: %u\n", __func__, params->n_ff); - LOG_INF("%s: n_layer: %u\n", __func__, params->n_layer); - LOG_INF("%s: n_rot: %u\n", __func__, params->n_rot); -} - -static void print_tensor_info(const struct ggml_context * ctx) { - for (auto * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - LOG_INF("%s: Allocating ", __func__); - int64_t total = 1; - int i = 0; - for (; i < ggml_n_dims(t); ++i) { - if (i > 0) { LOG_INF("x "); } - LOG_INF("[%" PRId64 "] ", t->ne[i]); - total *= t->ne[i]; - } - if (i > 1) { LOG_INF("= [%" PRId64 "] ", total); } - LOG_INF("float space for %s\n", ggml_get_name(t)); - } -} - -static void init_model(struct my_llama_model * model) { - const auto & hparams = model->hparams; - - const uint32_t n_embd = hparams.n_embd; - const uint32_t n_layer = hparams.n_layer; - const uint32_t n_vocab = hparams.n_vocab; - - const uint32_t n_multiqueries = hparams.n_head_kv <= 0 || hparams.n_head_kv >= hparams.n_head ? 1 : hparams.n_head / hparams.n_head_kv; - - const uint32_t n_ff = hparams.n_ff; - struct ggml_context * ctx = model->ctx; - - model->train_its = 0; - model->train_samples = 0; - model->train_tokens = 0; - - model->tok_embeddings = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_vocab); - model->norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); - model->output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_vocab); - - ggml_set_name(model->tok_embeddings, "tok_embeddings.weight"); - ggml_set_name(model->norm, "norm.weight"); - ggml_set_name(model->output, "output.weight"); - - model->layers.resize(n_layer); - for (uint32_t i = 0; i < n_layer; ++i) { - auto & layer = model->layers[i]; - - std::string layers_i = "layers." + std::to_string(i); - - layer.attention_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); - - layer.wq = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_embd); - layer.wk = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_embd / n_multiqueries); - layer.wv = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_embd / n_multiqueries); - layer.wo = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_embd); - - layer.ffn_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); - - layer.w1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_ff); - layer.w2 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_ff, n_embd); - layer.w3 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_ff); - - ggml_set_name(layer.attention_norm, (layers_i + ".attention_norm.weight").c_str()); - - ggml_set_name(layer.wq, (layers_i + ".attention.wq.weight").c_str()); - ggml_set_name(layer.wk, (layers_i + ".attention.wk.weight").c_str()); - ggml_set_name(layer.wv, (layers_i + ".attention.wv.weight").c_str()); - ggml_set_name(layer.wo, (layers_i + ".attention.wo.weight").c_str()); - - ggml_set_name(layer.ffn_norm, (layers_i + ".ffn_norm.weight").c_str()); - - ggml_format_name(layer.w1, "%s.feed_forward.w1.weight", layers_i.c_str()); - ggml_format_name(layer.w2, "%s.feed_forward.w2.weight", layers_i.c_str()); - ggml_format_name(layer.w3, "%s.feed_forward.w3.weight", layers_i.c_str()); - } - - print_tensor_info(ctx); -} - -static float get_f32_2d(struct ggml_tensor * tensor, int64_t i0, int64_t i1) { - float * ptr = (float *) ((char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1]); - return *ptr; -} - -static int32_t get_i32_2d(struct ggml_tensor * tensor, int64_t i0, int64_t i1) { - int32_t * ptr = (int32_t *) ((char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1]); - return *ptr; -} - -static void print_row(struct ggml_tensor * probs, int i) { - for (int k = 0; k < probs->ne[0]; ++k) { - float p = get_f32_2d(probs, k, i); - LOG(" %f", p); - } - LOG("\n"); -} - -static void print_matrix(struct ggml_tensor * probs) { - assert(ggml_is_matrix(probs)); - for (int i = 0; i < probs->ne[1]; ++i) { - for (int k = 0; k < probs->ne[0]; ++k) { - float p = get_f32_2d(probs, k, i); - LOG(" %.2f", p); - } - LOG("\n"); - } -} - -struct my_llama_file { - // use FILE * so we don't have to re-open the file to mmap - FILE * fp; - size_t size; - - my_llama_file(const char * fname, const char * mode) { - fp = std::fopen(fname, mode); - if (fp == NULL) { - size = 0; - } else { - seek(0, SEEK_END); - size = tell(); - seek(0, SEEK_SET); - } - } - - size_t tell() const { -#ifdef _WIN32 - __int64 ret = _ftelli64(fp); -#else - long ret = std::ftell(fp); -#endif - GGML_ASSERT(ret != -1); // this really shouldn't fail - return (size_t) ret; - } - - void seek(size_t offset, int whence) { -#ifdef _WIN32 - int ret = _fseeki64(fp, (__int64) offset, whence); -#else - int ret = std::fseek(fp, (long) offset, whence); -#endif - GGML_ASSERT(ret == 0); // same - } - - void read_raw(void * ptr, size_t size) { - if (size == 0) { - return; - } - errno = 0; - std::size_t ret = std::fread(ptr, size, 1, fp); - if (ferror(fp)) { - die_fmt("fread failed: %s", strerror(errno)); - } - if (ret != 1) { - die("unexpectedly reached end of file"); - } - } - - std::uint32_t read_u32() { - std::uint32_t ret; - read_raw(&ret, sizeof(ret)); - return ret; - } - std::float_t read_f32() { - std::float_t ret; - read_raw(&ret, sizeof(ret)); - return ret; - } - - std::string read_string(std::uint32_t len) { - std::vector chars(len); - read_raw(chars.data(), len); - return std::string(chars.data(), len); - } - - ~my_llama_file() { - if (fp) { - std::fclose(fp); - } - } -}; - -static bool is_ggml_file(const char * filename) { - my_llama_file file(filename, "rb"); - if (file.size < 4) { - return false; - } - std::string magic = file.read_string(4); - return magic == GGUF_MAGIC; -} - -static std::string llama_escape_whitespaces(const std::string & text) { - std::ostringstream out; - for (char c : text) { - if (c == ' ') out << "\xe2\x96\x81"; - else out << c; - } - return out.str(); -} - -static void load_vocab(const char * filename, const Config * config, struct my_llama_vocab * vocab) { - if (is_ggml_file(filename)) { - LOG_INF("%s: Loading vocabulary from gguf file %s\n", __func__, filename); - struct ggml_context * ctx_data = NULL; - - struct gguf_init_params params = { - /*.no_alloc = */ false, - /*.ctx = */ &ctx_data, - }; - - struct gguf_context * ctx = gguf_init_from_file(filename, params); - GGML_ASSERT(ctx != NULL); - - const int model_idx = gguf_find_key(ctx, KV_TOKENIZER_MODEL); - GGML_ASSERT(model_idx >= 0); - std::string tokenizer_name = gguf_get_val_str(ctx, model_idx); - GGML_ASSERT(tokenizer_name == TOKENIZER_NAME); - - const int token_idx = gguf_find_key(ctx, KV_TOKENIZER_LIST); - GGML_ASSERT(token_idx >= 0); - - const int score_idx = gguf_find_key(ctx, KV_TOKENIZER_SCORES); - GGML_ASSERT(score_idx >= 0); - const float * scores = (const float * ) gguf_get_arr_data(ctx, score_idx); - - const int toktype_idx = gguf_find_key(ctx, KV_TOKENIZER_TOKEN_TYPE); - GGML_ASSERT(toktype_idx >= 0); - const int * toktypes = (const int * ) gguf_get_arr_data(ctx, toktype_idx); - - const uint32_t n_vocab = gguf_get_arr_n(ctx, token_idx); - if (n_vocab != static_cast(config->vocab_size)) { - die_fmt("vocab size mismatch: (gguf) %u != (llama2c) %d", n_vocab, config->vocab_size); - } - - vocab->id_to_token.resize(n_vocab); - - for (uint32_t i = 0; i < n_vocab; i++) { - std::string word = gguf_get_arr_str(ctx, token_idx, i); - - vocab->token_to_id[word] = i; - - auto & token_data = vocab->id_to_token[i]; - token_data.text = std::move(word); - token_data.score = scores[i]; - token_data.type = (llama_token_type) toktypes[i]; - } - ggml_free(ctx_data); - gguf_free(ctx); - } else { - // assume llama2.c vocabulary - LOG_INF("%s: Assuming llama2.c vocabulary since %s is not a gguf file\n", __func__, filename); - my_llama_file file(filename, "rb"); - if (!file.fp) { - die_fmt("%s: %s", strerror(errno), filename); - } - const int n_vocab = config->vocab_size; - /* uint32_t max_token_length = */ file.read_u32(); // unused - vocab->id_to_token.resize(n_vocab); - for (my_llama_vocab::id id=0; id", &byte_val) == 1) { - // Text of byte tokens is already in the expected format. - type = LLAMA_TOKEN_TYPE_BYTE; - } else { - type = LLAMA_TOKEN_TYPE_NORMAL; - } - text = llama_escape_whitespaces(text); - - vocab->id_to_token[id].text = text; - vocab->id_to_token[id].score = score; - vocab->id_to_token[id].type = type; - vocab->token_to_id.emplace(text, id); - } - } -} - -static void convert_weights_ak_to_gg(struct ggml_tensor * gg_weights, const float * karpathy_weights) { - int size = 1; - for (int dim = 0; dim < ggml_n_dims(gg_weights); ++dim) { - size *= gg_weights->ne[dim]; - } - for (int ct = 0; ct < size; ++ct) { - int64_t i0 = 0; int64_t i1 = 0; - int64_t i2 = 0; int64_t i3 = 0; - ggml_unravel_index(gg_weights, ct, &i0, &i1, &i2, &i3); - ggml_set_f32_nd(gg_weights, i0, i1, i2, i3, karpathy_weights[ct]); - } -} - -static void save_as_llama_model( - struct my_llama_vocab * vocab, struct my_llama_model * model, TransformerWeights* w, const char * filename -) { - // convert AK weights into GG weights one by one. - // w->token_embedding_table -> model->tok_embeddings - // float* -> struct ggml_tensor - convert_weights_ak_to_gg(model->tok_embeddings, w->token_embedding_table.data()); - convert_weights_ak_to_gg(model->output, !w->wcls.empty() ? w->wcls.data() : w->token_embedding_table.data()); - - convert_weights_ak_to_gg(model->norm, w->rms_final_weight.data()); - //print_row(model->norm, 0); - - // for rms-att-weight - int row_length = model->hparams.n_embd; - int n_ff = model->hparams.n_ff; - - const uint32_t n_multiqueries = model->hparams.n_head_kv <= 0 || model->hparams.n_head_kv >= model->hparams.n_head ? 1 : model->hparams.n_head / model->hparams.n_head_kv; - - for (uint32_t i = 0; i < model->hparams.n_layer; ++i){ - auto & layer = model->layers[i]; - // 1d - convert_weights_ak_to_gg(layer.attention_norm, &w->rms_att_weight[i*row_length]); - convert_weights_ak_to_gg(layer.ffn_norm , &w->rms_ffn_weight[i*row_length]); - - // from 3d matrix layer x dim x dim to 2d matrix dim x dim - convert_weights_ak_to_gg(layer.wq , &w->wq[i*row_length*row_length]); - convert_weights_ak_to_gg(layer.wo , &w->wo[i*row_length*row_length]); - // from 3d matrix layer x dim x dim to 2d matrix dim x dim / n_multiqueries - convert_weights_ak_to_gg(layer.wk , &w->wk[i*row_length*row_length/n_multiqueries]); - convert_weights_ak_to_gg(layer.wv , &w->wv[i*row_length*row_length/n_multiqueries]); - - convert_weights_ak_to_gg(layer.w1 , &w->w1[i*row_length*n_ff]); - convert_weights_ak_to_gg(layer.w2 , &w->w2[i*n_ff*row_length]); - convert_weights_ak_to_gg(layer.w3 , &w->w3[i*row_length*n_ff]); - } - - struct gguf_context * ctx = gguf_init_empty(); - - std::vector tokens; - std::vector scores; - std::vector token_types; - for (const my_llama_vocab::token_data & token_data : vocab->id_to_token) { - tokens.push_back(token_data.text.c_str()); - scores.push_back(token_data.score); - token_types.push_back(token_data.type); - } - gguf_set_arr_str(ctx, KV_TOKENIZER_LIST, tokens.data(), tokens.size()); - gguf_set_arr_data(ctx, KV_TOKENIZER_SCORES, GGUF_TYPE_FLOAT32, scores.data(), scores.size()); - gguf_set_arr_data(ctx, KV_TOKENIZER_TOKEN_TYPE, GGUF_TYPE_INT32, token_types.data(), token_types.size()); - - gguf_set_val_str(ctx, KV_TOKENIZER_MODEL, TOKENIZER_NAME); - - gguf_set_val_str(ctx, KV_GENERAL_ARCHITECTURE, "llama"); - gguf_set_val_str(ctx, KV_GENERAL_NAME, "llama"); - - // special tokens - gguf_set_val_u32(ctx, KV_TOKENIZER_UNK_ID, UNKNOWN_TOKEN_ID); - gguf_set_val_u32(ctx, KV_TOKENIZER_BOS_ID, BOS_TOKEN_ID); - gguf_set_val_u32(ctx, KV_TOKENIZER_EOS_ID, EOS_TOKEN_ID); - gguf_set_val_u32(ctx, KV_TOKENIZER_SEP_ID, LLAMA_TOKEN_NULL); - gguf_set_val_u32(ctx, KV_TOKENIZER_PAD_ID, LLAMA_TOKEN_NULL); - - gguf_set_val_u32(ctx, KV_CONTEXT_LENGTH, model->hparams.n_ctx); - gguf_set_val_u32(ctx, KV_EMBEDDING_LENGTH, model->hparams.n_embd); - gguf_set_val_u32(ctx, KV_FEED_FORWARD_LENGTH, model->hparams.n_ff); - gguf_set_val_u32(ctx, KV_ATTENTION_HEAD_COUNT, model->hparams.n_head); - gguf_set_val_u32(ctx, KV_ATTENTION_HEAD_COUNT, model->hparams.n_head); - gguf_set_val_u32(ctx, KV_ATTENTION_HEAD_COUNT_KV, model->hparams.n_head_kv); - gguf_set_val_u32(ctx, KV_BLOCK_COUNT, model->hparams.n_layer); - gguf_set_val_u32(ctx, KV_ROPE_DIMENSION_COUNT, model->hparams.n_rot); - gguf_set_val_f32(ctx, KV_ATTENTION_LAYERNORM_RMS_EPS, 1e-5f); - - // write tensors - ggml_set_name(model->tok_embeddings, TN_TOKEN_EMBD); - gguf_add_tensor(ctx, model->tok_embeddings); - - ggml_set_name(model->norm, TN_OUTPUT_NORM); - gguf_add_tensor(ctx, model->norm); - - ggml_set_name(model->output, TN_OUTPUT); - gguf_add_tensor(ctx, model->output); - - for (uint32_t i = 0; i < model->hparams.n_layer; ++i) { - auto & layer = model->layers[i]; - - ggml_format_name(layer.wq, TN_ATTN_Q, i); - gguf_add_tensor(ctx, layer.wq); - - ggml_format_name(layer.wk, TN_ATTN_K, i); - gguf_add_tensor(ctx, layer.wk); - - ggml_format_name(layer.wv, TN_ATTN_V, i); - gguf_add_tensor(ctx, layer.wv); - - ggml_format_name(layer.wo, TN_ATTN_OUTPUT, i); - gguf_add_tensor(ctx, layer.wo); - - ggml_format_name(layer.attention_norm, TN_ATTN_NORM, i); - gguf_add_tensor(ctx, layer.attention_norm); - - ggml_format_name(layer.w1, TN_FFN_GATE, i); - gguf_add_tensor(ctx, layer.w1); - - ggml_format_name(layer.w2, TN_FFN_DOWN, i); - gguf_add_tensor(ctx, layer.w2); - - ggml_format_name(layer.w3, TN_FFN_UP, i); - gguf_add_tensor(ctx, layer.w3); - - ggml_format_name(layer.ffn_norm, TN_FFN_NORM, i); - gguf_add_tensor(ctx, layer.ffn_norm); - } - - gguf_write_to_file(ctx, filename, false); - gguf_free(ctx); -} - -static struct train_params get_default_train_params() { - struct train_params params; - params.fn_vocab_model = "models/7B/ggml-model-f16.gguf"; - params.fn_llama2c_output_model = "ak_llama_model.bin"; - params.fn_train_data = "shakespeare.txt"; - params.fn_checkpoint_in = "checkpoint.bin"; - params.fn_checkpoint_out = "checkpoint.bin"; - params.fn_model_out = "ggml-checkpoint-f32.bin"; - - params.seed = -1; - - params.n_ctx = 128; - params.n_embd = 256; - params.n_mult = 256; - params.n_head = 8; - params.n_layer = 16; - params.n_rotmax = 64; - - params.n_threads = 6; - params.n_batch = 8; - params.n_examples = 8; - params.n_predict = 1024; - - params.print_info_interval = 1; - params.print_details_interval = 2; - - params.samples_start_after_nl = false; - params.use_adam = true; - params.use_flash = false; - params.use_scratch = true; - - // only adam - params.warmup = 100; - params.cos_decay_steps = 1000; - params.cos_decay_restart = 1.1f; - params.cos_decay_alpha = 0.0f; - - params.lbfgs_n_iter = 16; - params.adam_n_iter = 16; - params.adam_alpha = 1e-3f; - params.adam_decay = 1e-3f; - - params.mem_model_gb = 2; - params.mem_compute_gb = 24; - params.mem_compute0_gb = 8; - params.mem_compute1_gb = 2; - - return params; -} - -static void print_usage(int /*argc*/, char ** argv, const struct train_params * params) { - fprintf(stderr, "usage: %s [options]\n", argv[0]); - fprintf(stderr, "\n"); - fprintf(stderr, "options:\n"); - fprintf(stderr, " -h, --help show this help message and exit\n"); - fprintf(stderr, " --copy-vocab-from-model FNAME path of gguf llama model or llama2.c vocabulary from which to copy vocab (default '%s')\n", params->fn_vocab_model); - fprintf(stderr, " --llama2c-model FNAME [REQUIRED] model path from which to load Karpathy's llama2.c model\n"); - fprintf(stderr, " --llama2c-output-model FNAME model path to save the converted llama2.c model (default %s')\n", params->fn_llama2c_output_model); - fprintf(stderr, "\n"); -} - -static bool params_parse(int argc, char ** argv, struct train_params * params) { - bool invalid_param = false; - bool reqd_param_found = false; - std::string arg; - struct train_params default_params = get_default_train_params(); - const std::string arg_prefix = "--"; - - for (int i = 1; i < argc; i++) { - arg = argv[i]; - if (arg.compare(0, arg_prefix.size(), arg_prefix) == 0) { - std::replace(arg.begin(), arg.end(), '_', '-'); - } - - if (arg == "--copy-vocab-from-model") { - if (++i >= argc) { - invalid_param = true; - break; - } - params->fn_vocab_model = argv[i]; - } else if (arg == "--llama2c-model") { - if (++i >= argc) { - invalid_param = true; - break; - } - reqd_param_found = true; - params->fn_llama2c_model = argv[i]; - } else if (arg == "--llama2c-output-model") { - if (++i >= argc) { - invalid_param = true; - break; - } - params->fn_llama2c_output_model = argv[i]; - } else if (arg == "-h" || arg == "--help") { - print_usage(argc, argv, &default_params); - exit(0); - } else { - fprintf(stderr, "error: unknown argument: %s\n", arg.c_str()); - print_usage(argc, argv, &default_params); - exit(1); - } - } - if (invalid_param) { - fprintf(stderr, "error: invalid parameter for argument: %s\n", arg.c_str()); - print_usage(argc, argv, &default_params); - exit(1); - } - if (!reqd_param_found){ - fprintf(stderr, "error: please specify a llama2.c .bin file to be converted with argument --llama2c-model\n"); - print_usage(argc, argv, &default_params); - exit(1); - } - - return true; -} - -static std::string basename(const std::string &path) { - size_t pos = path.find_last_of("/\\"); - if (pos == std::string::npos) { - return path; - } - return path.substr(pos + 1); -} - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_init(); - - struct train_params params = get_default_train_params(); - if (!params_parse(argc, argv, ¶ms)) { - return 1; - } - - Config config; - TransformerWeights weights = {}; - { - LOG_INF("%s: Loading llama2c model from %s\n", __func__, params.fn_llama2c_model); - FILE * file = fopen(params.fn_llama2c_model, "rb"); - if (!file) { - LOG_ERR("%s: Unable to open the checkpoint file %s!\n", __func__, params.fn_llama2c_model); - return 1; - } - // read in the config header - if (fread(&config, sizeof(Config), 1, file) != 1) { - LOG_ERR("%s: Unable to read llama2c config from %s!\n",__func__,params.fn_llama2c_model); - return 1; - } - auto shared_weights = config.vocab_size > 0; - config.vocab_size = abs(config.vocab_size); - - // read in the Transformer weights - alloc_weights(&weights, &config, shared_weights); - if (checkpoint_init_weights(&weights, &config, file, shared_weights)) { - LOG_ERR("%s: Unable to initialize transformer weights from %s!",__func__,params.fn_llama2c_model); - return 1; - } - fclose(file); - } - - struct my_llama_vocab vocab; - load_vocab(params.fn_vocab_model, &config, &vocab); - - struct my_llama_model model; - model.hparams.n_vocab = config.vocab_size; //llama_vocab_n_vocab(lctx); - model.hparams.n_ctx = params.n_ctx; - model.hparams.n_embd = config.dim; //params.n_embd; - model.hparams.n_ff = config.hidden_dim; - model.hparams.n_mult = 32;//params.n_mult; - model.hparams.n_head = config.n_heads; //params.n_head; - model.hparams.n_head_kv = config.n_kv_heads; - model.hparams.n_layer = config.n_layers; //params.n_layer; - model.hparams.n_rot = std::min((uint32_t)params.n_rotmax, model.hparams.n_embd / model.hparams.n_head); - - print_params(&model.hparams); - - struct ggml_init_params lcparams; - lcparams.mem_size = 1024ll*1024ll*1024ll*((size_t) params.mem_model_gb); - lcparams.mem_buffer = NULL; - lcparams.no_alloc = false; - - model.ctx = ggml_init(lcparams); - - init_model(&model); - model.name = basename(params.fn_llama2c_model); - save_as_llama_model(&vocab, &model, &weights, params.fn_llama2c_output_model); - - LOG_INF("%s: Saving llama.c model file %s in ggml format at %s\n", __func__, params.fn_llama2c_model, params.fn_llama2c_output_model); - - ggml_free(model.ctx); - return 0; -} diff --git a/examples/debug-rust/Cargo.toml b/examples/debug-rust/Cargo.toml new file mode 100644 index 000000000000..52f65e67ff39 --- /dev/null +++ b/examples/debug-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-debug-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-debug" +path = "src/main.rs" diff --git a/examples/debug-rust/build.rs b/examples/debug-rust/build.rs new file mode 100644 index 000000000000..27d0a601bd30 --- /dev/null +++ b/examples/debug-rust/build.rs @@ -0,0 +1,43 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building the binary"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/debug-rust/src/lib.rs b/examples/debug-rust/src/lib.rs new file mode 100644 index 000000000000..5ccf9b180dc2 --- /dev/null +++ b/examples/debug-rust/src/lib.rs @@ -0,0 +1,148 @@ +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub n_gpu_layers: i32, + pub save_logits: bool, + pub embedding: bool, + pub embd_normalize: i32, + pub logits_output_dir: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + MissingModel, + MissingValue(&'static str), + InvalidInteger(&'static str, String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut args = args.into_iter().map(Into::into); + let mut parsed = Args { + model_path: String::new(), + prompt: "Hello my name is".to_string(), + n_gpu_layers: 99, + save_logits: false, + embedding: false, + embd_normalize: -1, + logits_output_dir: "data".to_string(), + }; + + while let Some(arg) = args.next() { + match arg.as_str() { + "-m" | "--model" => { + parsed.model_path = args.next().ok_or(ParseError::MissingValue("-m"))?; + } + "-p" | "--prompt" => { + parsed.prompt = args.next().ok_or(ParseError::MissingValue("-p"))?; + } + "-ngl" | "--n-gpu-layers" => { + let value = args.next().ok_or(ParseError::MissingValue("-ngl"))?; + parsed.n_gpu_layers = value + .parse() + .map_err(|_| ParseError::InvalidInteger("-ngl", value))?; + } + "--save-logits" => parsed.save_logits = true, + "--embedding" => parsed.embedding = true, + "--embd-normalize" => { + let value = args + .next() + .ok_or(ParseError::MissingValue("--embd-normalize"))?; + parsed.embd_normalize = value + .parse() + .map_err(|_| ParseError::InvalidInteger("--embd-normalize", value))?; + } + "--logits-output-dir" => { + parsed.logits_output_dir = args + .next() + .ok_or(ParseError::MissingValue("--logits-output-dir"))?; + } + _ => {} + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + + Ok(parsed) +} + +pub fn normalize_embedding(input: &[f32], embd_norm: i32) -> Vec { + let sum = match embd_norm { + -1 => 1.0, + 0 => input.iter().map(|v| v.abs() as f64).fold(0.0_f64, f64::max) / 32760.0, + 2 => input + .iter() + .map(|v| (*v as f64) * (*v as f64)) + .sum::() + .sqrt(), + p => input + .iter() + .map(|v| (v.abs() as f64).powi(p)) + .sum::() + .powf(1.0 / p as f64), + }; + let norm = if sum > 0.0 { (1.0 / sum) as f32 } else { 0.0 }; + input.iter().map(|v| v * norm).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_model_conversion_invocation() { + let args = parse_args([ + "-m", + "model.gguf", + "--embedding", + "-p", + "Hello world today", + "--save-logits", + "--embd-normalize", + "2", + ]) + .unwrap(); + + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "Hello world today"); + assert!(args.embedding); + assert!(args.save_logits); + assert_eq!(args.embd_normalize, 2); + assert_eq!(args.logits_output_dir, "data"); + } + + #[test] + fn normalizes_euclidean() { + let out = normalize_embedding(&[3.0, 4.0], 2); + assert!((out[0] - 0.6).abs() < 1e-6); + assert!((out[1] - 0.8).abs() < 1e-6); + } + + #[test] + fn rejects_missing_model() { + assert_eq!( + parse_args(["--save-logits"]).unwrap_err(), + ParseError::MissingModel + ); + } +} diff --git a/examples/debug-rust/src/main.rs b/examples/debug-rust/src/main.rs new file mode 100644 index 000000000000..dd77811e17fd --- /dev/null +++ b/examples/debug-rust/src/main.rs @@ -0,0 +1,316 @@ +use llama_debug_rust::{normalize_embedding, parse_args, Args}; +use llama_simple_rust::ffi; +use std::env; +use std::ffi::CString; +use std::fs::{self, File}; +use std::io::{BufWriter, Write}; +use std::path::Path; +use std::ptr; +use std::slice; + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m model.gguf -p \"Hello my name is\" --save-logits"); + eprintln!(" {program} -m model.gguf --embedding -p \"Hello\" --save-logits"); + eprintln!(); +} + +fn tokenize( + vocab: *const ffi::llama_vocab, + prompt: &str, + add_special: bool, +) -> Result, String> { + let prompt_c = + CString::new(prompt).map_err(|_| "prompt contains an interior NUL byte".to_string())?; + let n_prompt = unsafe { + -ffi::llama_tokenize( + vocab, + prompt_c.as_ptr(), + prompt.len() as i32, + ptr::null_mut(), + 0, + add_special, + true, + ) + }; + + if n_prompt <= 0 { + return Err("failed to size prompt tokenization".to_string()); + } + + let mut tokens = vec![0_i32; n_prompt as usize]; + let n_tokenized = unsafe { + ffi::llama_tokenize( + vocab, + prompt_c.as_ptr(), + prompt.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + add_special, + true, + ) + }; + + if n_tokenized < 0 { + return Err("failed to tokenize the prompt".to_string()); + } + + tokens.truncate(n_tokenized as usize); + Ok(tokens) +} + +fn token_to_piece( + vocab: *const ffi::llama_vocab, + token: ffi::llama_token, +) -> Result { + let mut buf = vec![0_i8; 128]; + let n = unsafe { + ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), buf.len() as i32, 0, true) + }; + if n < 0 { + return Err(format!("failed to convert token {token} to piece")); + } + let bytes = buf[..n as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +fn print_tokenized_prompt( + vocab: *const ffi::llama_vocab, + tokens: &[ffi::llama_token], + prompt: &str, +) -> Result<(), String> { + println!("Model add_bos: {}", unsafe { + ffi::llama_vocab_get_add_bos(vocab) + }); + println!("Input prompt: \"{prompt}\""); + println!("Token ids ({}):", tokens.len()); + for &token in tokens { + print!("{}({token}) ", token_to_piece(vocab, token)?); + } + println!(); + Ok(()) +} + +fn save_output_data( + data: &[f32], + tokens: &[ffi::llama_token], + prompt: &str, + model_path: &str, + output_dir: &str, + embedding: bool, +) -> Result<(), String> { + fs::create_dir_all(output_dir).map_err(|err| err.to_string())?; + let model_name = Path::new(model_path) + .file_stem() + .and_then(|stem| stem.to_str()) + .unwrap_or("model"); + let type_suffix = if embedding { "-embeddings" } else { "" }; + let base_path = Path::new(output_dir).join(format!("llamacpp-{model_name}{type_suffix}")); + + let bin_path = base_path.with_extension("bin"); + let mut bin_file = BufWriter::new(File::create(&bin_path).map_err(|err| err.to_string())?); + for value in data { + bin_file + .write_all(&value.to_ne_bytes()) + .map_err(|err| err.to_string())?; + } + println!("Data saved to {}", bin_path.display()); + + let txt_path = base_path.with_extension("txt"); + let mut txt_file = BufWriter::new(File::create(&txt_path).map_err(|err| err.to_string())?); + for (i, value) in data.iter().enumerate() { + writeln!(txt_file, "{i}: {value}").map_err(|err| err.to_string())?; + } + println!("Data saved to {}", txt_path.display()); + + let prompt_path = Path::new(&format!("{}-prompt.txt", base_path.display())).to_path_buf(); + let mut prompt_file = + BufWriter::new(File::create(&prompt_path).map_err(|err| err.to_string())?); + writeln!(prompt_file, "prompt: {prompt}").map_err(|err| err.to_string())?; + writeln!(prompt_file, "n_tokens: {}", tokens.len()).map_err(|err| err.to_string())?; + write!(prompt_file, "token ids: ").map_err(|err| err.to_string())?; + for (i, token) in tokens.iter().enumerate() { + if i > 0 { + write!(prompt_file, ", ").map_err(|err| err.to_string())?; + } + write!(prompt_file, "{token}").map_err(|err| err.to_string())?; + } + writeln!(prompt_file).map_err(|err| err.to_string())?; + println!("Prompt saved to {}", prompt_path.display()); + + let tokens_path = Path::new(&format!("{}-tokens.bin", base_path.display())).to_path_buf(); + let mut tokens_file = + BufWriter::new(File::create(&tokens_path).map_err(|err| err.to_string())?); + for token in tokens { + tokens_file + .write_all(&token.to_ne_bytes()) + .map_err(|err| err.to_string())?; + } + println!("Tokens saved to {}", tokens_path.display()); + + Ok(()) +} + +fn output_data( + ctx: *mut ffi::llama_context, + model: *mut ffi::llama_model, + vocab: *const ffi::llama_vocab, + tokens: &[ffi::llama_token], + args: &Args, +) -> Result, String> { + if args.embedding { + let n_embd = unsafe { ffi::llama_model_n_embd_out(model) }; + if n_embd <= 0 { + return Err("model returned invalid embedding size".to_string()); + } + let pooling = !matches!(unsafe { ffi::llama_pooling_type(ctx) }, -1 | 0); + let n_embd_count = if pooling { 1 } else { tokens.len() }; + let n_floats = n_embd as usize * n_embd_count; + let embd = if pooling { + unsafe { ffi::llama_get_embeddings_seq(ctx, 0) } + } else { + unsafe { ffi::llama_get_embeddings(ctx) } + }; + if embd.is_null() { + return Err("failed to get embeddings from the model".to_string()); + } + let raw = unsafe { slice::from_raw_parts(embd, n_floats) }; + if args.embd_normalize >= 0 { + let mut normalized = Vec::with_capacity(n_floats); + for chunk in raw.chunks(n_embd as usize) { + normalized.extend(normalize_embedding(chunk, args.embd_normalize)); + } + Ok(normalized) + } else { + Ok(raw.to_vec()) + } + } else { + let logits = unsafe { ffi::llama_get_logits_ith(ctx, tokens.len() as i32 - 1) }; + if logits.is_null() { + return Err("failed to get logits from the model".to_string()); + } + let n_logits = unsafe { ffi::llama_vocab_n_tokens(vocab) }; + Ok(unsafe { slice::from_raw_parts(logits, n_logits as usize) }.to_vec()) + } +} + +fn run(args: Args) -> Result<(), String> { + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + + unsafe { + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + let add_bos = unsafe { ffi::llama_vocab_get_add_bos(vocab) }; + let mut tokens = tokenize(vocab, &args.prompt, add_bos)?; + if tokens.is_empty() { + return Err("there are no input tokens to process".to_string()); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = tokens.len().max(1) as u32; + ctx_params.n_batch = tokens.len() as u32; + ctx_params.embeddings = args.embedding; + ctx_params.no_perf = false; + + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create the llama_context".to_string()); + } + + let batch = unsafe { ffi::llama_batch_get_one(tokens.as_mut_ptr(), tokens.len() as i32) }; + if unsafe { ffi::llama_decode(ctx.0, batch) } != 0 { + return Err("failed to eval".to_string()); + } + + print_tokenized_prompt(vocab, &tokens, &args.prompt)?; + + if args.save_logits { + let data = output_data(ctx.0, model.0, vocab, &tokens, &args)?; + save_output_data( + &data, + &tokens, + &args.prompt, + &args.model_path, + &args.logits_output_dir, + args.embedding, + )?; + } + + println!(); + unsafe { + ffi::llama_perf_context_print(ctx.0); + } + + Ok(()) +} + +fn main() { + let mut argv = env::args(); + let program = argv.next().unwrap_or_else(|| "llama-debug".to_string()); + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + + let status = match run(args) { + Ok(()) => 0, + Err(err) => { + eprintln!("{program}: error: {err}"); + 1 + } + }; + + unsafe { + ffi::llama_backend_free(); + } + std::process::exit(status); +} diff --git a/examples/debug/CMakeLists.txt b/examples/debug/CMakeLists.txt index fb1c7e258147..32f03f3ceb0a 100644 --- a/examples/debug/CMakeLists.txt +++ b/examples/debug/CMakeLists.txt @@ -1,5 +1,26 @@ set(TARGET llama-debug) -add_executable(${TARGET} debug.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(DEBUG_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../debug-rust/Cargo.toml) +set(DEBUG_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/debug-example) +set(DEBUG_RUST_BIN ${DEBUG_RUST_TARGET_DIR}/debug/llama-debug${CMAKE_EXECUTABLE_SUFFIX}) +set(DEBUG_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${DEBUG_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${DEBUG_RUST_MANIFEST} --target-dir ${DEBUG_RUST_TARGET_DIR} --bin llama-debug + COMMAND ${CMAKE_COMMAND} -E copy ${DEBUG_RUST_BIN} ${DEBUG_RUST_OUTPUT} + DEPENDS + llama + ${DEBUG_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../debug-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../debug-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../debug-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${DEBUG_RUST_OUTPUT}) +install(PROGRAMS ${DEBUG_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/debug/debug.cpp b/examples/debug/debug.cpp deleted file mode 100644 index 761e7a2db54e..000000000000 --- a/examples/debug/debug.cpp +++ /dev/null @@ -1,261 +0,0 @@ -#include "debug.h" -#include "arg.h" -#include "common.h" -#include "log.h" -#include "llama.h" - -#include -#include -#include -#include -#include -#include -#include - -static void print_usage(int /*argc*/, char ** argv) { - const std::string usage_template = R"( - example usage: - - Print tensors: - - {prog} -m model.gguf -p "Hello my name is" --verbose - - The tensors to be printed can be filtered with --tensor-filter option. - - Save logits/embeddings: - - {prog} -m model.gguf -p "Hello my name is" --save-logits - - Add --embedding to save embeddings)" "\n"; - - // Fix the source code indentation above that is introduced by the raw string literal. - std::string usage = std::regex_replace(usage_template, std::regex("\\n {8}"), "\n"); - usage = std::regex_replace(usage, std::regex("\\{prog\\}"), argv[0]); - LOG("%s\n", usage.c_str()); -} - -static bool has_pooling(llama_context * ctx) { - switch (llama_pooling_type(ctx)) { - case LLAMA_POOLING_TYPE_NONE: - case LLAMA_POOLING_TYPE_UNSPECIFIED: - return false; - default: - return true; - } -} - -struct output_data { - float * data_ptr = nullptr; - int data_size = 0; - std::string type_suffix; - std::vector embd_norm; - std::string prompt; - std::vector tokens; - - output_data(llama_context * ctx, const llama_model * model, const common_params & params) { - const llama_vocab * vocab = llama_model_get_vocab(model); - const bool add_bos = llama_vocab_get_add_bos(vocab); - - tokens = common_tokenize(ctx, params.prompt, add_bos); - prompt = params.prompt; - - if (params.embedding) { - const int n_embd = llama_model_n_embd_out(model); - const bool pooling = has_pooling(ctx); - const int n_embd_count = pooling ? 1 : tokens.size(); - const int n_floats = n_embd * n_embd_count; - - float * embd_raw = pooling ? llama_get_embeddings_seq(ctx, 0) : llama_get_embeddings(ctx); - if (embd_raw == nullptr) { - throw std::runtime_error("failed to get embeddings from the model"); - } - - LOG_DBG("pooling_enabled: %s\n", pooling ? "true" : "false"); - LOG_DBG("n_embd: %d\n", n_embd); - LOG_DBG("n_floats: %d\n", n_floats); - LOG_DBG("n_embd_count: %d\n", n_embd_count); - - data_ptr = embd_raw; - data_size = n_floats; - type_suffix = "-embeddings"; - - if (params.embd_normalize >= 0) { - embd_norm.resize(n_floats); - for (int i = 0; i < n_embd_count; i++) { - common_embd_normalize(embd_raw+i*n_embd, embd_norm.data()+i*n_embd, n_embd, params.embd_normalize); - } - data_ptr = embd_norm.data(); - } - } else { - const float * logits = llama_get_logits_ith(ctx, tokens.size() - 1); - const int n_logits = llama_vocab_n_tokens(vocab); - - data_ptr = const_cast(logits); - data_size = n_logits; - type_suffix = ""; - } - } -}; - -static void save_output_data(const output_data & output, const std::string & model_name, const std::string & output_dir) { - std::filesystem::create_directory(output_dir); - auto base_path = std::filesystem::path{output_dir} / ("llamacpp-" + model_name + output.type_suffix); - - // Save logits/embeddings to binary file. - { - std::filesystem::path filepath{base_path.string() + ".bin"}; - std::ofstream file{filepath, std::ios::binary}; - if (!file) { - throw std::runtime_error("failed to open binary output file: " + filepath.string()); - } - file.write(reinterpret_cast(output.data_ptr), output.data_size * sizeof(float)); - LOG("Data saved to %s\n", filepath.c_str()); - } - - // Save logits/embeddings to text file. - { - std::filesystem::path filepath{base_path.string() + ".txt"}; - std::ofstream file{filepath}; - if (!file) { - throw std::runtime_error("failed to open text output file: " + filepath.string()); - } - for (int i = 0; i < output.data_size; i++) { - file << i << ": " << output.data_ptr[i] << '\n'; - } - LOG("Data saved to %s\n", filepath.c_str()); - } - - // Save prompt and tokens to text file. - { - std::filesystem::path filepath{base_path.string() + "-prompt.txt"}; - std::ofstream file{filepath}; - if (!file) { - throw std::runtime_error("failed to open prompt output file: " + filepath.string()); - } - - file << "prompt: " << output.prompt << '\n'; - file << "n_tokens: " << output.tokens.size() << '\n'; - - file << "token ids: "; - for (size_t i = 0; i < output.tokens.size(); i++) { - file << output.tokens[i]; - if (i + 1 < output.tokens.size()) { - file << ", "; - } - } - file << '\n'; - LOG("Prompt saved to %s\n", filepath.c_str()); - } - - // Save token ids to binary file. - { - std::filesystem::path filepath{base_path.string() + "-tokens.bin"}; - std::ofstream file{filepath, std::ios::binary}; - if (!file) { - throw std::runtime_error("failed to open tokens binary file: " + filepath.string()); - } - file.write(reinterpret_cast(output.tokens.data()), output.tokens.size() * sizeof(llama_token)); - LOG("Tokens saved to %s\n", filepath.c_str()); - } - -} - -static void print_tokenized_prompt(llama_context * ctx, const std::vector & tokens, const std::string & prompt) { - const llama_model * model = llama_get_model(ctx); - const llama_vocab * vocab = llama_model_get_vocab(model); - - LOG("Model add_bos: %s\n", llama_vocab_get_add_bos(vocab) ? "true" : "false"); - LOG("Input prompt: \"%s\"\n", prompt.c_str()); - LOG("Token ids (%zu):\n", tokens.size()); - - for (auto id : tokens) { - std::string piece(128, '\0'); - int n = llama_token_to_piece(vocab, id, piece.data(), piece.size(), 0, true); - if (n < 0) { - LOG_ERR("failed to convert token %d to piece\n", id); - continue; - } - piece.resize(n); - LOG("%s(%d) ", piece.c_str(), id); - } - LOG("\n"); -} - -static bool run(llama_context * ctx, const common_params & params) { - const llama_model * model = llama_get_model(ctx); - const llama_vocab * vocab = llama_model_get_vocab(model); - - const bool add_bos = llama_vocab_get_add_bos(vocab); - - std::vector tokens = common_tokenize(ctx, params.prompt, add_bos); - - if (tokens.empty()) { - LOG_ERR("%s : there are not input tokens to process - (try to provide a prompt with '-p')\n", __func__); - return false; - } - - if (llama_decode(ctx, llama_batch_get_one(tokens.data(), tokens.size()))) { - LOG_ERR("%s : failed to eval\n", __func__); - return false; - } - - print_tokenized_prompt(ctx, tokens, params.prompt); - - if (params.save_logits) { - try { - output_data output {ctx, model, params}; - std::filesystem::path model_path{params.model.path}; - std::string model_name{model_path.stem().string()}; - save_output_data(output, model_name, params.logits_output_dir); - } catch (const std::exception & e) { - LOG_ERR("%s : error saving logits: %s\n", __func__, e.what()); - } - } - - return true; -} - -int main(int argc, char ** argv) { - common_params params; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_DEBUG, print_usage)) { - return 1; - } - - llama_backend_init(); - llama_numa_init(params.numa); - - std::optional cb_data; - if (!params.save_logits) { - cb_data.emplace(params, params.tensor_filter); - } - - auto llama_init = common_init_from_params(params); - - auto * model = llama_init->model(); - auto * ctx = llama_init->context(); - - if (model == nullptr || ctx == nullptr) { - LOG_ERR("%s : failed to init\n", __func__); - return 1; - } - - { - LOG_INF("\n"); - LOG_INF("%s\n", common_params_get_system_info(params).c_str()); - LOG_INF("\n"); - } - - if (!run(ctx, params)) { - return 1; - } - - LOG("\n"); - llama_perf_context_print(ctx); - - llama_backend_free(); - - return 0; -} diff --git a/examples/deprecation-warning/Cargo.toml b/examples/deprecation-warning/Cargo.toml new file mode 100644 index 000000000000..a94f0c439ea5 --- /dev/null +++ b/examples/deprecation-warning/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "llama-deprecation-warning" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] diff --git a/examples/deprecation-warning/deprecation-warning.cpp b/examples/deprecation-warning/deprecation-warning.cpp deleted file mode 100644 index 0cde17f6e996..000000000000 --- a/examples/deprecation-warning/deprecation-warning.cpp +++ /dev/null @@ -1,38 +0,0 @@ -// Warns users that this filename was deprecated, and provides a link for more information. - -#include -#include -#include -#include - -// Main -int main(int argc, char** argv) { - std::setlocale(LC_NUMERIC, "C"); - - std::string filename = "main"; - if (argc >= 1) { - filename = argv[0]; - } - - // Get only the program name from the full path - auto pos = filename.find_last_of("/\\"); - if (pos != std::string::npos) { - filename = filename.substr(pos+1); - } - - // Append "llama-" to the beginning of filename to get the replacemnt filename - auto replacement_filename = "llama-" + filename; - - // The exception is if the filename is "main", then our replacement filename is "llama-cli" - if (filename == "main") { - replacement_filename = "llama-cli"; - } - - fprintf(stdout, "\n"); - fprintf(stdout, "WARNING: The binary '%s' is deprecated.\n", filename.c_str()); - fprintf(stdout, " Please use '%s' instead.\n", replacement_filename.c_str()); - fprintf(stdout, " See https://github.com/ggml-org/llama.cpp/tree/master/examples/deprecation-warning/README.md for more information.\n"); - fprintf(stdout, "\n"); - - return EXIT_FAILURE; -} diff --git a/examples/deprecation-warning/src/lib.rs b/examples/deprecation-warning/src/lib.rs new file mode 100644 index 000000000000..acaff528f762 --- /dev/null +++ b/examples/deprecation-warning/src/lib.rs @@ -0,0 +1,53 @@ +pub fn program_name(path: &str) -> &str { + path.rsplit(['/', '\\']).next().unwrap_or(path) +} + +pub fn replacement_name(filename: &str) -> String { + if filename == "main" { + "llama-cli".to_string() + } else { + format!("llama-{filename}") + } +} + +pub fn warning_message(path: &str) -> String { + let filename = program_name(path); + let replacement = replacement_name(filename); + + format!( + "\nWARNING: The binary '{filename}' is deprecated.\n Please use '{replacement}' instead.\n See https://github.com/ggml-org/llama.cpp/tree/master/examples/deprecation-warning/README.md for more information.\n\n" + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extracts_unix_program_name() { + assert_eq!(program_name("/tmp/server"), "server"); + } + + #[test] + fn extracts_windows_program_name() { + assert_eq!(program_name(r"C:\tools\quantize"), "quantize"); + } + + #[test] + fn maps_main_to_cli() { + assert_eq!(replacement_name("main"), "llama-cli"); + } + + #[test] + fn prefixes_other_names() { + assert_eq!(replacement_name("server"), "llama-server"); + } + + #[test] + fn renders_warning() { + let message = warning_message("/usr/bin/main"); + + assert!(message.contains("The binary 'main' is deprecated.")); + assert!(message.contains("Please use 'llama-cli' instead.")); + } +} diff --git a/examples/deprecation-warning/src/main.rs b/examples/deprecation-warning/src/main.rs new file mode 100644 index 000000000000..96b92b24e200 --- /dev/null +++ b/examples/deprecation-warning/src/main.rs @@ -0,0 +1,7 @@ +use std::env; + +fn main() { + let path = env::args().next().unwrap_or_else(|| "main".to_string()); + print!("{}", llama_deprecation_warning::warning_message(&path)); + std::process::exit(1); +} diff --git a/examples/diffusion-rust/Cargo.toml b/examples/diffusion-rust/Cargo.toml new file mode 100644 index 000000000000..7d616f699a06 --- /dev/null +++ b/examples/diffusion-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-diffusion-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-diffusion-cli" +path = "src/main.rs" diff --git a/examples/diffusion-rust/build.rs b/examples/diffusion-rust/build.rs new file mode 100644 index 000000000000..74b9588d9cd3 --- /dev/null +++ b/examples/diffusion-rust/build.rs @@ -0,0 +1,45 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); + + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-diffusion-cli"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/diffusion-rust/src/main.rs b/examples/diffusion-rust/src/main.rs new file mode 100644 index 000000000000..1fd01ecde605 --- /dev/null +++ b/examples/diffusion-rust/src/main.rs @@ -0,0 +1,1114 @@ +use llama_simple_rust::ffi; +use std::env; +use std::ffi::{CStr, CString}; +use std::io::{self, Write}; +use std::ptr; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DiffusionAlgorithm { + Origin = 0, + EntropyBased = 1, + MarginBased = 2, + Random = 3, + ConfidenceBased = 4, +} + +impl TryFrom for DiffusionAlgorithm { + type Error = ParseError; + + fn try_from(value: i32) -> Result { + match value { + 0 => Ok(Self::Origin), + 1 => Ok(Self::EntropyBased), + 2 => Ok(Self::MarginBased), + 3 => Ok(Self::Random), + 4 => Ok(Self::ConfidenceBased), + _ => Err(ParseError::InvalidValue(format!( + "diffusion algorithm must be in 0..=4, got {value}" + ))), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TransferSchedule { + TimestepBased, + BlockBased, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub system_prompt: String, + pub use_chat_template: bool, + pub n_ctx: u32, + pub n_batch: u32, + pub n_ubatch: u32, + pub n_threads: i32, + pub n_threads_batch: i32, + pub n_gpu_layers: i32, + pub use_mmap: bool, + pub use_mlock: bool, + pub check_tensors: bool, + pub no_perf: bool, + pub flash_attn_type: i32, + pub cache_type_k: ffi::ggml_type, + pub cache_type_v: ffi::ggml_type, + pub seed: u32, + pub top_k: i32, + pub top_p: f32, + pub temp: f32, + pub steps: i32, + pub eps: f32, + pub block_length: i32, + pub algorithm: DiffusionAlgorithm, + pub visual_mode: bool, + pub add_gumbel_noise: bool, + pub cfg_scale: f32, + pub alg_temp: f32, +} + +impl Default for Args { + fn default() -> Self { + Self { + model_path: String::new(), + prompt: String::new(), + system_prompt: String::new(), + use_chat_template: false, + n_ctx: 512, + n_batch: 512, + n_ubatch: 512, + n_threads: 4, + n_threads_batch: 4, + n_gpu_layers: 99, + use_mmap: true, + use_mlock: false, + check_tensors: false, + no_perf: false, + flash_attn_type: 0, + cache_type_k: 0, + cache_type_v: 0, + seed: ffi::LLAMA_DEFAULT_SEED, + top_k: 40, + top_p: 0.95, + temp: 0.0, + steps: 64, + eps: 1e-3, + block_length: 0, + algorithm: DiffusionAlgorithm::ConfidenceBased, + visual_mode: false, + add_gumbel_noise: false, + cfg_scale: 0.0, + alg_temp: 0.0, + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub enum ParseError { + MissingModel, + MissingPrompt, + MissingValue(String), + InvalidInteger(String, String), + InvalidFloat(String, String), + InvalidValue(String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::MissingModel => write!(f, "missing required -m/--model model.gguf"), + Self::MissingPrompt => write!(f, "missing prompt; pass -p/--prompt or trailing text"), + Self::MissingValue(flag) => write!(f, "missing value for {flag}"), + Self::InvalidInteger(flag, value) => write!(f, "invalid integer for {flag}: {value}"), + Self::InvalidFloat(flag, value) => write!(f, "invalid float for {flag}: {value}"), + Self::InvalidValue(message) => write!(f, "{message}"), + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut parsed = Args::default(); + let mut iter = args.into_iter().map(Into::into); + + while let Some(arg) = iter.next() { + match arg.as_str() { + "-m" | "--model" => parsed.model_path = take(&mut iter, &arg)?, + "-p" | "--prompt" => parsed.prompt = take(&mut iter, &arg)?, + "-f" | "--file" | "--prompt-file" => { + let path = take(&mut iter, &arg)?; + parsed.prompt = std::fs::read_to_string(&path).map_err(|err| { + ParseError::InvalidValue(format!("failed to read prompt file {path}: {err}")) + })?; + } + "--system-prompt" => parsed.system_prompt = take(&mut iter, &arg)?, + "--chat-template" | "--conversation" => parsed.use_chat_template = true, + "--no-chat-template" => parsed.use_chat_template = false, + "-c" | "--ctx-size" => parsed.n_ctx = parse_u32(&mut iter, &arg)?, + "-b" | "--batch-size" => parsed.n_batch = parse_u32(&mut iter, &arg)?, + "-ub" | "--ubatch-size" => parsed.n_ubatch = parse_u32(&mut iter, &arg)?, + "-t" | "--threads" => parsed.n_threads = parse_i32(&mut iter, &arg)?, + "-tb" | "--threads-batch" => parsed.n_threads_batch = parse_i32(&mut iter, &arg)?, + "-ngl" | "--gpu-layers" => parsed.n_gpu_layers = parse_i32(&mut iter, &arg)?, + "--no-mmap" => parsed.use_mmap = false, + "--mlock" => parsed.use_mlock = true, + "--check-tensors" => parsed.check_tensors = true, + "--no-perf" => parsed.no_perf = true, + "-fa" | "--flash-attn" => parsed.flash_attn_type = 1, + "--no-flash-attn" => parsed.flash_attn_type = 0, + "-s" | "--seed" => parsed.seed = parse_u32(&mut iter, &arg)?, + "--top-k" => parsed.top_k = parse_i32(&mut iter, &arg)?, + "--top-p" => parsed.top_p = parse_f32(&mut iter, &arg)?, + "--temp" => parsed.temp = parse_f32(&mut iter, &arg)?, + "--diffusion-steps" | "--steps" => parsed.steps = parse_i32(&mut iter, &arg)?, + "--diffusion-eps" | "--eps" => { + parsed.eps = parse_f32(&mut iter, &arg)?; + parsed.block_length = 0; + } + "--diffusion-block-length" | "--block-length" => { + parsed.block_length = parse_i32(&mut iter, &arg)?; + parsed.eps = 0.0; + } + "--diffusion-algorithm" | "--algorithm" => { + parsed.algorithm = DiffusionAlgorithm::try_from(parse_i32(&mut iter, &arg)?)? + } + "--diffusion-visual" | "--visual" => parsed.visual_mode = true, + "--diffusion-gumbel" | "--gumbel" => parsed.add_gumbel_noise = true, + "--cfg-scale" => parsed.cfg_scale = parse_f32(&mut iter, &arg)?, + "--alg-temp" => parsed.alg_temp = parse_f32(&mut iter, &arg)?, + "-h" | "--help" => return Err(ParseError::MissingModel), + other if other.starts_with('-') => { + return Err(ParseError::InvalidValue(format!( + "unknown argument: {other}" + ))); + } + other => { + let mut prompt = vec![other.to_string()]; + prompt.extend(iter); + parsed.prompt = prompt.join(" "); + break; + } + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + if parsed.prompt.is_empty() { + return Err(ParseError::MissingPrompt); + } + if parsed.n_ctx == 0 || parsed.n_batch == 0 || parsed.n_ubatch == 0 { + return Err(ParseError::InvalidValue( + "context and batch sizes must be positive".to_string(), + )); + } + if parsed.steps <= 0 { + return Err(ParseError::InvalidValue( + "--steps must be positive".to_string(), + )); + } + if (parsed.eps == 0.0) == (parsed.block_length == 0) { + return Err(ParseError::InvalidValue( + "set exactly one of --eps or --block-length".to_string(), + )); + } + Ok(parsed) +} + +fn take(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + iter.next() + .ok_or_else(|| ParseError::MissingValue(flag.to_string())) +} + +fn parse_i32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_u32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_f32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidFloat(flag.to_string(), value)) +} + +struct Backend; + +impl Backend { + fn init() -> Self { + unsafe { + ffi::ggml_backend_load_all(); + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + Self + } +} + +impl Drop for Backend { + fn drop(&mut self) { + unsafe { + ffi::llama_backend_free(); + } + } +} + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Batch(ffi::llama_batch); + +impl Drop for Batch { + fn drop(&mut self) { + unsafe { + ffi::llama_batch_free(self.0); + } + } +} + +struct Sampler(*mut ffi::llama_sampler); + +impl Drop for Sampler { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_sampler_free(self.0); + } + } + } +} + +#[derive(Clone)] +struct SmallRng(u64); + +impl SmallRng { + fn new(seed: u32) -> Self { + Self((seed as u64).wrapping_add(0x9E37_79B9_7F4A_7C15)) + } + + fn next_u32(&mut self) -> u32 { + self.0 = self + .0 + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (self.0 >> 32) as u32 + } + + fn next_f32(&mut self) -> f32 { + (self.next_u32() as f64 / (u32::MAX as f64 + 1.0)) as f32 + } +} + +fn calculate_confidence( + cur_p: &[ffi::llama_token_data], + selected: usize, + algorithm: DiffusionAlgorithm, + rng: &mut SmallRng, +) -> f32 { + match algorithm { + DiffusionAlgorithm::ConfidenceBased | DiffusionAlgorithm::Origin => cur_p[selected].p, + DiffusionAlgorithm::EntropyBased => -cur_p + .iter() + .map(|candidate| { + let p = candidate.p.max(0.0); + p * (p + 1e-10).ln() + }) + .sum::(), + DiffusionAlgorithm::MarginBased => { + if cur_p.len() > 1 { + cur_p[0].p - cur_p[1].p + } else { + cur_p[0].p + } + } + DiffusionAlgorithm::Random => rng.next_f32(), + } +} + +pub fn calculate_transfer_count( + step: i32, + total_steps: i32, + remaining_masked: i32, + schedule: TransferSchedule, + eps: f32, + num_transfer_tokens: &[i32], +) -> i32 { + match schedule { + TransferSchedule::TimestepBased => { + let t = 1.0 - step as f32 / total_steps as f32 * (1.0 - eps); + let s = 1.0 - (step + 1) as f32 / total_steps as f32 * (1.0 - eps); + let p_transfer = if step < total_steps - 1 { + 1.0 - s / t + } else { + 1.0 + }; + (remaining_masked as f32 * p_transfer) as i32 + } + TransferSchedule::BlockBased => num_transfer_tokens + .get(step as usize) + .copied() + .unwrap_or_else(|| remaining_masked / (total_steps - step).max(1)), + } +} + +pub fn get_num_transfer_tokens(mask_count: i32, steps: i32) -> Vec { + let base = mask_count / steps; + let remainder = mask_count % steps; + (0..steps) + .map(|i| base + if i < remainder { 1 } else { 0 }) + .collect() +} + +fn gumbel_adjust(logit: f32, temperature: f32, rng: &mut SmallRng) -> f32 { + if temperature == 0.0 { + return logit; + } + let noise = rng.next_f32().max(1e-20) as f64; + let gumbel_noise = (-noise.ln()).powf(temperature as f64); + (f64::from(logit).exp() / gumbel_noise) as f32 +} + +fn fill_candidates_from_logits( + candidates: &mut [ffi::llama_token_data], + ctx: *mut ffi::llama_context, + cond_logits: Option<&[f32]>, + n_vocab: usize, + pos: usize, + shift_logits: bool, + add_noise: bool, + temperature: f32, + rng: &mut SmallRng, +) -> Result<(), Box> { + let actual_pos = if shift_logits && pos > 0 { + pos - 1 + } else { + pos + }; + let src = if let Some(cond_logits) = cond_logits { + &cond_logits[actual_pos * n_vocab..(actual_pos + 1) * n_vocab] + } else { + let ptr = unsafe { ffi::llama_get_logits_ith(ctx, actual_pos as i32) }; + if ptr.is_null() { + return Err(format!("failed to get logits for position {actual_pos}").into()); + } + unsafe { std::slice::from_raw_parts(ptr, n_vocab) } + }; + + for (token_id, candidate) in candidates.iter_mut().enumerate() { + let mut logit = src[token_id]; + if add_noise && temperature > 0.0 { + logit = gumbel_adjust(logit, temperature, rng); + } + candidate.id = token_id as i32; + candidate.logit = logit; + candidate.p = 0.0; + } + Ok(()) +} + +fn sampler_init(args: &Args) -> Sampler { + unsafe { + let chain = ffi::llama_sampler_chain_init(ffi::llama_sampler_chain_default_params()); + if args.top_k > 0 { + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_top_k(args.top_k)); + } + if args.top_p < 1.0 { + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_top_p(args.top_p, 1)); + } + if args.temp > 0.0 { + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_temp(args.temp)); + } + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_dist(args.seed)); + Sampler(chain) + } +} + +fn dist_sampler_init(seed: u32) -> Sampler { + unsafe { Sampler(ffi::llama_sampler_init_dist(seed)) } +} + +fn token_to_piece(vocab: *const ffi::llama_vocab, token: ffi::llama_token) -> String { + let mut buf = vec![0_i8; 256]; + let n = unsafe { ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), 256, 0, false) }; + if n <= 0 { + String::new() + } else { + String::from_utf8_lossy(unsafe { + std::slice::from_raw_parts(buf.as_ptr().cast::(), n as usize) + }) + .into_owned() + } +} + +fn detokenize(vocab: *const ffi::llama_vocab, tokens: &[ffi::llama_token]) -> String { + if tokens.is_empty() { + return String::new(); + } + let mut cap = tokens.len().saturating_mul(16).max(128); + loop { + let mut buf = vec![0_i8; cap]; + let n = unsafe { + ffi::llama_detokenize( + vocab, + tokens.as_ptr(), + tokens.len() as i32, + buf.as_mut_ptr(), + cap as i32, + false, + false, + ) + }; + if n >= 0 && (n as usize) < cap { + return String::from_utf8_lossy(unsafe { + std::slice::from_raw_parts(buf.as_ptr().cast::(), n as usize) + }) + .into_owned(); + } + cap *= 2; + } +} + +fn tokenize( + vocab: *const ffi::llama_vocab, + text: &str, +) -> Result, Box> { + let text_c = CString::new(text)?; + let mut cap = text.len().saturating_add(8).max(32); + loop { + let mut tokens = vec![0; cap]; + let n = unsafe { + ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + true, + true, + ) + }; + if n >= 0 { + tokens.truncate(n as usize); + return Ok(tokens); + } + cap = (-n) as usize; + } +} + +fn format_input_text( + args: &Args, + model: *const ffi::llama_model, +) -> Result> { + if !args.use_chat_template { + return Ok(args.prompt.clone()); + } + let tmpl = unsafe { ffi::llama_model_chat_template(model, ptr::null()) }; + if tmpl.is_null() { + return Ok(args.prompt.clone()); + } + let user_role = CString::new("user")?; + let user_content = CString::new(args.prompt.as_str())?; + let system_role = CString::new("system")?; + let system_content = CString::new(args.system_prompt.as_str())?; + let mut messages = Vec::new(); + if !args.system_prompt.is_empty() { + messages.push(ffi::llama_chat_message { + role: system_role.as_ptr(), + content: system_content.as_ptr(), + }); + } + messages.push(ffi::llama_chat_message { + role: user_role.as_ptr(), + content: user_content.as_ptr(), + }); + + let mut cap = 1024; + loop { + let mut buf = vec![0_i8; cap]; + let n = unsafe { + ffi::llama_chat_apply_template( + tmpl, + messages.as_ptr(), + messages.len(), + true, + buf.as_mut_ptr(), + cap as i32, + ) + }; + if n < 0 { + return Err("failed to apply chat template".into()); + } + if n as usize <= cap { + return Ok(String::from_utf8_lossy(unsafe { + std::slice::from_raw_parts(buf.as_ptr().cast::(), n as usize) + }) + .into_owned()); + } + cap = n as usize + 1; + } +} + +fn callback_progress( + step: i32, + total_steps: i32, + tokens: &[ffi::llama_token], + n_input: usize, + vocab: *const ffi::llama_vocab, + mask_token: ffi::llama_token, + visual_mode: bool, +) { + let progress_percent = if total_steps > 0 { + step * 100 / total_steps + } else { + 100 + }; + let progress_bars = if total_steps > 0 { + step * 50 / total_steps + } else { + 50 + }; + eprint!( + "\rdiffusion step: {step}/{total_steps} [{}{}] {progress_percent}%", + "=".repeat(progress_bars as usize), + " ".repeat((50 - progress_bars).max(0) as usize) + ); + if visual_mode { + eprint!("\x1b[2J\x1b[H"); + let mut text = String::from(" "); + for &token in &tokens[n_input..] { + if token == mask_token { + text.push(' '); + } else { + text.push_str(&token_to_piece(vocab, token)); + } + } + eprintln!("\n{text}"); + } + let _ = io::stderr().flush(); +} + +fn diffusion_generate( + ctx: *mut ffi::llama_context, + input_tokens: &[ffi::llama_token], + args: &Args, + shift_logits: bool, + mask_token_id: ffi::llama_token, + vocab: *const ffi::llama_vocab, +) -> Result, Box> { + if ctx.is_null() || input_tokens.is_empty() || args.n_ubatch as usize <= input_tokens.len() { + return Err("invalid diffusion generation inputs".into()); + } + + unsafe { + ffi::llama_set_causal_attn(ctx, false); + } + + let max_length = args.n_ubatch as usize; + let n_vocab = unsafe { ffi::llama_vocab_n_tokens(vocab) as usize }; + let mut output_tokens = vec![mask_token_id; max_length]; + output_tokens[..input_tokens.len()].copy_from_slice(input_tokens); + let mut rng = SmallRng::new(args.seed); + let sampler = sampler_init(args); + let dist_sampler = dist_sampler_init(args.seed); + let mut batch = Batch(unsafe { ffi::llama_batch_init(max_length as i32, 0, 1) }); + batch.0.n_tokens = max_length as i32; + + let mut candidates = vec![ + ffi::llama_token_data { + id: 0, + logit: 0.0, + p: 0.0, + }; + n_vocab + ]; + let mut conf_candidates = Vec::with_capacity(max_length); + let mut cond_logits_buffer = if args.cfg_scale > 0.0 { + vec![0.0; n_vocab * max_length] + } else { + Vec::new() + }; + let mut uncond_tokens = if args.cfg_scale > 0.0 { + vec![mask_token_id; max_length] + } else { + Vec::new() + }; + + let schedule = if args.block_length > 0 { + TransferSchedule::BlockBased + } else { + TransferSchedule::TimestepBased + }; + let mut num_blocks = 1; + let mut steps_per_block = args.steps; + if schedule == TransferSchedule::BlockBased { + if max_length % args.block_length as usize != 0 + || args.steps % (max_length as i32 / args.block_length) != 0 + { + return Err("block length must divide length and steps".into()); + } + num_blocks = max_length as i32 / args.block_length; + steps_per_block = args.steps / num_blocks; + } + + let start = unsafe { ffi::ggml_time_us() }; + let mut sampling_time = 0_i64; + + for block_num in 0..num_blocks { + let block_start = if schedule == TransferSchedule::BlockBased { + input_tokens.len() + block_num as usize * args.block_length as usize + } else { + 0 + }; + let block_end = if schedule == TransferSchedule::BlockBased { + (input_tokens.len() + (block_num as usize + 1) * args.block_length as usize) + .min(max_length) + } else { + max_length + }; + let num_transfer_tokens = if schedule == TransferSchedule::BlockBased { + let block_mask_count = output_tokens[block_start..block_end] + .iter() + .filter(|&&token| token == mask_token_id) + .count() as i32; + get_num_transfer_tokens(block_mask_count, steps_per_block) + } else { + Vec::new() + }; + + for step in 0..steps_per_block { + let global_step = block_num * steps_per_block + step; + callback_progress( + global_step, + args.steps, + &output_tokens, + input_tokens.len(), + vocab, + mask_token_id, + args.visual_mode, + ); + + for i in 0..max_length { + unsafe { + *batch.0.token.add(i) = output_tokens[i]; + *batch.0.pos.add(i) = i as i32; + *batch.0.n_seq_id.add(i) = 1; + *(*batch.0.seq_id.add(i)) = 0; + *batch.0.logits.add(i) = 1; + } + } + + if args.cfg_scale > 0.0 { + let ret = unsafe { ffi::llama_decode(ctx, batch.0) }; + if ret != 0 { + return Err(format!("failed to generate conditional logits: {ret}").into()); + } + for pos in 0..max_length { + let src = unsafe { ffi::llama_get_logits_ith(ctx, pos as i32) }; + if src.is_null() { + return Err("failed to get conditional logits".into()); + } + unsafe { + ptr::copy_nonoverlapping( + src, + cond_logits_buffer[pos * n_vocab..].as_mut_ptr(), + n_vocab, + ); + } + } + + uncond_tokens.copy_from_slice(&output_tokens); + for token in &mut uncond_tokens[..input_tokens.len()] { + *token = mask_token_id; + } + for (i, &token) in uncond_tokens.iter().enumerate() { + unsafe { + *batch.0.token.add(i) = token; + } + } + let ret = unsafe { ffi::llama_decode(ctx, batch.0) }; + if ret != 0 { + return Err(format!("failed to generate unconditional logits: {ret}").into()); + } + for pos in 0..max_length { + let src = unsafe { ffi::llama_get_logits_ith(ctx, pos as i32) }; + if src.is_null() { + return Err("failed to get unconditional logits".into()); + } + let dst = &mut cond_logits_buffer[pos * n_vocab..(pos + 1) * n_vocab]; + for (i, dst_logit) in dst.iter_mut().enumerate() { + let uncond = unsafe { *src.add(i) }; + *dst_logit = uncond + (args.cfg_scale + 1.0) * (*dst_logit - uncond); + } + } + } else { + let ret = unsafe { ffi::llama_decode(ctx, batch.0) }; + if ret != 0 { + return Err(format!("failed to decode at step {global_step}: {ret}").into()); + } + } + + let sampling_start = unsafe { ffi::ggml_time_us() }; + let mut mask_positions = Vec::new(); + for (i, &token) in output_tokens.iter().enumerate() { + if token == mask_token_id + && (schedule != TransferSchedule::BlockBased + || (i >= block_start && i < block_end)) + { + mask_positions.push(i); + } + } + if mask_positions.is_empty() { + break; + } + + let cond_logits = if args.cfg_scale > 0.0 { + Some(cond_logits_buffer.as_slice()) + } else { + None + }; + + if args.algorithm == DiffusionAlgorithm::Origin { + let transfer_count = calculate_transfer_count( + step, + steps_per_block, + mask_positions.len() as i32, + schedule, + args.eps, + &num_transfer_tokens, + ); + let p_transfer = transfer_count as f32 / mask_positions.len() as f32; + for pos in mask_positions { + if rng.next_f32() < p_transfer { + fill_candidates_from_logits( + &mut candidates, + ctx, + cond_logits, + n_vocab, + pos, + shift_logits, + args.add_gumbel_noise, + args.temp, + &mut rng, + )?; + let mut cur_p = ffi::llama_token_data_array { + data: candidates.as_mut_ptr(), + size: candidates.len(), + selected: -1, + sorted: false, + }; + unsafe { + ffi::llama_sampler_apply(sampler.0, &mut cur_p); + } + output_tokens[pos] = candidates[cur_p.selected as usize].id; + } + } + } else { + let mut confidences = Vec::with_capacity(mask_positions.len()); + let mut sampled_tokens = vec![0; mask_positions.len()]; + for (mask_idx, &pos) in mask_positions.iter().enumerate() { + fill_candidates_from_logits( + &mut candidates, + ctx, + cond_logits, + n_vocab, + pos, + shift_logits, + args.add_gumbel_noise, + args.temp, + &mut rng, + )?; + let mut cur_p = ffi::llama_token_data_array { + data: candidates.as_mut_ptr(), + size: candidates.len(), + selected: -1, + sorted: false, + }; + unsafe { + ffi::llama_sampler_apply(sampler.0, &mut cur_p); + } + let selected = cur_p.selected as usize; + let sampled_token = candidates[selected].id; + let conf = + calculate_confidence(&candidates, selected, args.algorithm, &mut rng); + sampled_tokens[mask_idx] = sampled_token; + confidences.push((conf, mask_idx)); + } + + let transfer_count = calculate_transfer_count( + step, + steps_per_block, + mask_positions.len() as i32, + schedule, + args.eps, + &num_transfer_tokens, + ) + .max(0) as usize; + let n_take = transfer_count.min(confidences.len()); + if args.alg_temp == 0.0 { + confidences.sort_by(|a, b| b.0.total_cmp(&a.0).then_with(|| a.1.cmp(&b.1))); + for &(_, mask_idx) in confidences.iter().take(n_take) { + output_tokens[mask_positions[mask_idx]] = sampled_tokens[mask_idx]; + } + } else { + conf_candidates.clear(); + for (i, (conf, _)) in confidences.iter().enumerate() { + conf_candidates.push(ffi::llama_token_data { + id: i as i32, + logit: *conf / args.alg_temp, + p: 0.0, + }); + } + for _ in 0..n_take { + let mut conf_array = ffi::llama_token_data_array { + data: conf_candidates.as_mut_ptr(), + size: conf_candidates.len(), + selected: -1, + sorted: false, + }; + unsafe { + ffi::llama_sampler_apply(dist_sampler.0, &mut conf_array); + } + let selected_idx = conf_array.selected as usize; + let mask_idx = selected_idx; + output_tokens[mask_positions[mask_idx]] = sampled_tokens[mask_idx]; + conf_candidates[selected_idx].p = 0.0; + } + } + } + sampling_time += unsafe { ffi::ggml_time_us() } - sampling_start; + } + } + + let total_time = unsafe { ffi::ggml_time_us() } - start; + eprintln!( + "\ntotal time: {:.2}ms, time per step: {:.2}ms, sampling time per step: {:.2}ms", + total_time as f64 / 1000.0, + total_time as f64 / 1000.0 / args.steps as f64, + sampling_time as f64 / 1000.0 / args.steps as f64 + ); + + Ok(output_tokens) +} + +fn run(args: Args) -> Result<(), Box> { + let _backend = Backend::init(); + let model_path = CString::new(args.model_path.as_str())?; + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + model_params.use_mmap = args.use_mmap; + model_params.use_mlock = args.use_mlock; + model_params.check_tensors = args.check_tensors; + + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err(format!("failed to load model '{}'", args.model_path).into()); + } + if !unsafe { ffi::llama_model_is_diffusion(model.0) } { + return Err("unsupported model for diffusion".into()); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = args.n_ctx; + ctx_params.n_batch = args.n_batch; + ctx_params.n_ubatch = args.n_ubatch; + ctx_params.no_perf = args.no_perf; + ctx_params.flash_attn_type = args.flash_attn_type; + ctx_params.type_k = args.cache_type_k; + ctx_params.type_v = args.cache_type_v; + + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create context".into()); + } + unsafe { + ffi::llama_set_n_threads(ctx.0, args.n_threads, args.n_threads_batch); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + let formatted_prompt = format_input_text(&args, model.0)?; + let input_tokens = tokenize(vocab, &formatted_prompt)?; + if input_tokens.len() as u32 >= unsafe { ffi::llama_n_ctx(ctx.0) } { + return Err(format!( + "input too long ({} tokens), max context is {}", + input_tokens.len(), + unsafe { ffi::llama_n_ctx(ctx.0) } + ) + .into()); + } + + let mask_token = unsafe { ffi::llama_vocab_mask(vocab) }; + if mask_token == ffi::LLAMA_TOKEN_NULL { + return Err("model vocabulary does not define a mask token".into()); + } + + let mut shift_buf = vec![0_i8; 8]; + let key = CString::new("diffusion.shift_logits")?; + let shift_logits = if unsafe { + ffi::llama_model_meta_val_str( + model.0, + key.as_ptr(), + shift_buf.as_mut_ptr(), + shift_buf.len(), + ) + } >= 0 + { + unsafe { CStr::from_ptr(shift_buf.as_ptr()) }.to_string_lossy() == "true" + } else { + true + }; + + let mut output_tokens = + diffusion_generate(ctx.0, &input_tokens, &args, shift_logits, mask_token, vocab)?; + if args.visual_mode { + eprint!("\x1b[2J\x1b[H"); + } + output_tokens.drain(0..input_tokens.len()); + println!("\n{}", detokenize(vocab, &output_tokens)); + Ok(()) +} + +fn main() { + let args = match parse_args(env::args().skip(1)) { + Ok(args) => args, + Err(err) => { + eprintln!("error: {err}"); + std::process::exit(1); + } + }; + if let Err(err) = run(args) { + eprintln!("error: {err}"); + std::process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_diffusion_flags() { + let args = parse_args([ + "-m", + "model.gguf", + "-p", + "hello", + "--steps", + "12", + "--block-length", + "16", + "--algorithm", + "2", + "--visual", + "--gumbel", + "--cfg-scale", + "1.5", + ]) + .unwrap(); + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "hello"); + assert_eq!(args.steps, 12); + assert_eq!(args.eps, 0.0); + assert_eq!(args.block_length, 16); + assert_eq!(args.algorithm, DiffusionAlgorithm::MarginBased); + assert!(args.visual_mode); + assert!(args.add_gumbel_noise); + assert_eq!(args.cfg_scale, 1.5); + } + + #[test] + fn trailing_text_becomes_prompt() { + let args = parse_args(["-m", "model.gguf", "write", "text"]).unwrap(); + assert_eq!(args.prompt, "write text"); + } + + #[test] + fn requires_exactly_one_schedule() { + let err = parse_args(["-m", "model.gguf", "-p", "x", "--eps", "0"]).unwrap_err(); + assert!(matches!(err, ParseError::InvalidValue(_))); + } + + #[test] + fn distributes_block_transfer_tokens() { + assert_eq!(get_num_transfer_tokens(10, 4), vec![3, 3, 2, 2]); + } + + #[test] + fn timestep_transfer_finishes_on_last_step() { + assert_eq!( + calculate_transfer_count(3, 4, 7, TransferSchedule::TimestepBased, 0.001, &[]), + 7 + ); + } + + #[test] + fn margin_confidence_uses_first_two_candidates() { + let mut rng = SmallRng::new(1); + let candidates = [ + ffi::llama_token_data { + id: 1, + logit: 0.0, + p: 0.7, + }, + ffi::llama_token_data { + id: 2, + logit: 0.0, + p: 0.2, + }, + ]; + assert!( + (calculate_confidence(&candidates, 0, DiffusionAlgorithm::MarginBased, &mut rng) - 0.5) + .abs() + < 1e-6 + ); + } +} diff --git a/examples/diffusion/CMakeLists.txt b/examples/diffusion/CMakeLists.txt index 70228d4079b3..dad5138509fb 100644 --- a/examples/diffusion/CMakeLists.txt +++ b/examples/diffusion/CMakeLists.txt @@ -1,5 +1,25 @@ set(TARGET llama-diffusion-cli) -add_executable(${TARGET} diffusion-cli.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama llama-common ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(DIFFUSION_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../diffusion-rust/Cargo.toml) +set(DIFFUSION_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/diffusion) +set(DIFFUSION_RUST_BIN ${DIFFUSION_RUST_TARGET_DIR}/debug/llama-diffusion-cli${CMAKE_EXECUTABLE_SUFFIX}) +set(DIFFUSION_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${DIFFUSION_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${DIFFUSION_RUST_MANIFEST} --target-dir ${DIFFUSION_RUST_TARGET_DIR} --bin llama-diffusion-cli + COMMAND ${CMAKE_COMMAND} -E copy ${DIFFUSION_RUST_BIN} ${DIFFUSION_RUST_OUTPUT} + DEPENDS + llama + ${DIFFUSION_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../diffusion-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../diffusion-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${DIFFUSION_RUST_OUTPUT}) +install(PROGRAMS ${DIFFUSION_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/diffusion/diffusion-cli.cpp b/examples/diffusion/diffusion-cli.cpp deleted file mode 100644 index 403b9b474454..000000000000 --- a/examples/diffusion/diffusion-cli.cpp +++ /dev/null @@ -1,698 +0,0 @@ -#include "arg.h" -#include "chat.h" -#include "common.h" -#include "llama.h" -#include "log.h" - -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -enum diffusion_algorithm { ORIGIN = 0, ENTROPY_BASED = 1, MARGIN_BASED = 2, RANDOM = 3, CONFIDENCE_BASED = 4 }; - -// Unified transfer scheduling methods -enum transfer_schedule { - TIMESTEP_BASED = 0, // Dream-style: (1.0 - s/t) * remaining - BLOCK_BASED = 1, // LLaDA-style: process in blocks with get_num_transfer_tokens -}; - -typedef bool (*diffusion_step_callback_t)(int32_t step, - int32_t total_steps, - const llama_token * tokens, - int32_t n_tokens, - void * user_data); - -struct diffusion_params { - int32_t steps = 0; - float temperature = 0; - llama_token mask_token_id = LLAMA_TOKEN_NULL; - diffusion_step_callback_t step_callback = nullptr; - void * step_callback_user_data = nullptr; - int32_t seed = 0; - bool visual_mode = false; - bool shift_logits = false; // Shift logits by -1 after decode - - float top_p = 0.; - int32_t top_k = 0.; - - diffusion_algorithm algorithm = CONFIDENCE_BASED; - transfer_schedule schedule = TIMESTEP_BASED; - - float cfg_scale = 0.; // Config scale for classifier-free guidance - float eps = 0.; // Timestep scheduling - int32_t block_length = 0; // Block size (for block scheduling) - float alg_temp = 0; // algorithm temperature (0.0 = deterministic) - bool add_gumbel_noise = false; // Add gumbel noise to the logits if temp > 0.0 - - int32_t max_length = 0; // Maximum sequence length -}; - -struct callback_data { - diffusion_params * diff_params; - const llama_vocab * vocab; - int32_t n_input; -}; - -static float calculate_confidence(const llama_token_data_array & cur_p, - diffusion_algorithm algorithm, - std::mt19937 & rng) { - switch (algorithm) { - case CONFIDENCE_BASED: - return cur_p.data[cur_p.selected].p; // Selected token probability - - case ENTROPY_BASED: - { - float entropy = 0.0f; - const float epsilon = 1e-10f; - for (size_t i = 0; i < cur_p.size; i++) { - float prob = cur_p.data[i].p; - entropy += prob * logf(prob + epsilon); - } - return -entropy; // Higher entropy = lower confidence - } - - case MARGIN_BASED: - return (cur_p.size > 1) ? cur_p.data[0].p - cur_p.data[1].p : cur_p.data[0].p; - - case RANDOM: - { - std::uniform_real_distribution uniform(0.0f, 1.0f); - return uniform(rng); // Random confidence - } - - case ORIGIN: - return cur_p.data[cur_p.selected].p; - - default: - return 0.0f; - } -} - -// Unified transfer count calculation function -static int32_t calculate_transfer_count(int32_t step, - int32_t total_steps, - int32_t remaining_masked, - transfer_schedule schedule, - float eps, - const std::vector & num_transfer_tokens = {}) { - switch (schedule) { - case TIMESTEP_BASED: - { - float t = 1.0f - (float) step / total_steps * (1.0f - eps); - float s = 1.0f - (float) (step + 1) / total_steps * (1.0f - eps); - float p_transfer = (step < total_steps - 1) ? (1.0f - s / t) : 1.0f; - return (int32_t) (remaining_masked * p_transfer); - } - - case BLOCK_BASED: - if (!num_transfer_tokens.empty() && step < (int32_t) num_transfer_tokens.size()) { - return num_transfer_tokens[step]; - } - return remaining_masked / (total_steps - step); // Fallback - - default: - return remaining_masked / (total_steps - step); - } -} - -static bool diffusion_step_callback(int32_t step, - int32_t total_steps, - const llama_token * tokens, - int32_t n_tokens, - void * user_data) { - (void) user_data; - - callback_data * data = static_cast(user_data); - - auto print_progress_bar = [](int32_t step, int32_t total_steps) { - int progress_percent = (step * 100) / total_steps; - int progress_bars = (step * 50) / total_steps; - LOG_INF("\rdiffusion step: %d/%d [%s%s] %d%%", - step, - total_steps, - std::string(progress_bars, '=').c_str(), - std::string(50 - progress_bars, ' ').c_str(), - progress_percent); - }; - - if (data->diff_params->visual_mode) { - // Visual mode: clear - LOG_INF("\033[2J\033[H"); // Clear screen and move cursor to top-left - - print_progress_bar(step, total_steps); - - LOG_INF("\n"); - - std::string current_text = " "; - - for (int32_t i = data->n_input; i < n_tokens; i++) { - std::string token_str; - if (tokens[i] != llama_vocab_mask(data->vocab)) { - char piece[256]; - int n_chars = llama_token_to_piece(data->vocab, tokens[i], piece, sizeof(piece), 0, false); - if (n_chars > 0) { - piece[n_chars] = '\0'; - token_str = piece; - } - } else { - token_str = " "; - } - - current_text += token_str; - } - - LOG_INF("%s\n", current_text.c_str()); - } else { - print_progress_bar(step, total_steps); - } - - return true; -} - -static void add_gumbel_noise(float * logits, int32_t n_vocab, float temperature, std::mt19937 & rng) { - if (temperature == 0.0f) { - return; - } - - std::uniform_real_distribution uniform(0.0, 1.0); - for (int32_t i = 0; i < n_vocab; i++) { - double noise = uniform(rng); - // Prevent log(0) - noise = std::max(noise, 1e-20); - double gumbel_noise = std::pow(-std::log(noise), temperature); - logits[i] = std::exp(logits[i]) / gumbel_noise; - } -} - -static std::vector get_num_transfer_tokens(int32_t mask_count, int32_t steps) { - std::vector num_transfer_tokens(steps); - - int32_t base = mask_count / steps; - int32_t remainder = mask_count % steps; - - for (int32_t i = 0; i < steps; i++) { - num_transfer_tokens[i] = base + (i < remainder ? 1 : 0); - } - - return num_transfer_tokens; -} - -static void diffusion_generate(llama_context * ctx, - const llama_token * input_tokens, - llama_token * output_tokens, - int32_t n_input, - const diffusion_params & params, - int32_t & n_generated) { - n_generated = 0; - if (!ctx || !input_tokens || !output_tokens || n_input <= 0 || params.max_length <= n_input) { - return; - } - - const llama_model * model = llama_get_model(ctx); - - // Initialize with input and pad with mask tokens - std::copy(input_tokens, input_tokens + n_input, output_tokens); - std::fill(output_tokens + n_input, output_tokens + params.max_length, params.mask_token_id); - - std::mt19937 rng(params.seed); - - llama_set_causal_attn(ctx, false); - - int32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); - - std::vector candidates(n_vocab); - std::vector conf_candidates; - conf_candidates.reserve(params.max_length); - std::vector mask_positions; - mask_positions.reserve(params.max_length); - - // Setup sampler chain - struct llama_sampler * sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()); - if (params.top_k > 0) { - llama_sampler_chain_add(sampler, llama_sampler_init_top_k(params.top_k)); - } - if (params.top_p < 1.0f) { - llama_sampler_chain_add(sampler, llama_sampler_init_top_p(params.top_p, 1)); - } - if (params.temperature > 0.0f) { - llama_sampler_chain_add(sampler, llama_sampler_init_temp(params.temperature)); - } - llama_sampler_chain_add(sampler, llama_sampler_init_dist(params.seed)); - - struct llama_sampler * dist_sampler = llama_sampler_init_dist(params.seed); - - llama_batch batch = llama_batch_init(params.max_length, 0, 1); - batch.n_tokens = params.max_length; - - // Pre-allocate buffers for CFG if needed - int32_t logits_size = n_vocab * params.max_length; - std::vector cond_logits_buffer; - std::vector un_x_buffer; - if (params.cfg_scale > 0.0f) { - cond_logits_buffer.resize(logits_size); - un_x_buffer.resize(params.max_length); - } - - // For block-based processing - std::vector num_transfer_tokens; - int32_t num_blocks = 1; - int32_t steps_per_block = params.steps; - - if (params.schedule == BLOCK_BASED) { - GGML_ASSERT(params.max_length % params.block_length == 0); - num_blocks = params.max_length / params.block_length; - GGML_ASSERT(params.steps % num_blocks == 0); - steps_per_block = params.steps / num_blocks; - } - - std::vector confidence(params.max_length); - - int64_t total_sampling_time = 0; - int64_t total_time = 0; - int64_t time_start = ggml_time_us(); - - for (int block_num = 0; block_num < num_blocks; block_num++) { - int32_t block_start = (params.schedule == BLOCK_BASED) ? n_input + block_num * params.block_length : 0; - int32_t block_end = (params.schedule == BLOCK_BASED) ? - std::min(n_input + (block_num + 1) * params.block_length, params.max_length) : - params.max_length; - - // Count masked tokens in current block for block-based processing - if (params.schedule == BLOCK_BASED) { - int32_t block_mask_count = 0; - for (int i = block_start; i < block_end; i++) { - if (output_tokens[i] == params.mask_token_id) { - block_mask_count++; - } - } - num_transfer_tokens = get_num_transfer_tokens(block_mask_count, steps_per_block); - } - - for (int32_t step = 0; step < steps_per_block; step++) { - int32_t global_step = block_num * steps_per_block + step; - - if (params.step_callback) { - if (!params.step_callback( - global_step, params.steps, output_tokens, params.max_length, params.step_callback_user_data)) { - break; - } - } - - // Setup batch - for (int32_t i = 0; i < params.max_length; i++) { - batch.token[i] = output_tokens[i]; - batch.pos[i] = i; - batch.n_seq_id[i] = 1; - batch.seq_id[i][0] = 0; - batch.logits[i] = 1; - } - - float * logits = nullptr; - - if (params.cfg_scale > 0.0f) { - int ret = llama_decode(ctx, batch); - if (ret != 0) { - LOG_ERR("Failed to generate conditional"); - break; - } - float * cond_logits_ptr = llama_get_logits(ctx); - std::memcpy(cond_logits_buffer.data(), cond_logits_ptr, logits_size * sizeof(float)); - - // Unconditional generation (mask input) - std::copy(output_tokens, output_tokens + params.max_length, un_x_buffer.begin()); - for (int32_t i = 0; i < n_input; i++) { - un_x_buffer[i] = params.mask_token_id; - } - - for (int32_t i = 0; i < params.max_length; i++) { - batch.token[i] = un_x_buffer[i]; - } - ret = llama_decode(ctx, batch); - if (ret != 0) { - LOG_ERR("Failed to generate unconditional"); - break; - } - float * uncond_logits = llama_get_logits(ctx); - - // Apply CFG - for (int32_t i = 0; i < logits_size; i++) { - cond_logits_buffer[i] = - uncond_logits[i] + (params.cfg_scale + 1.0f) * (cond_logits_buffer[i] - uncond_logits[i]); - } - logits = cond_logits_buffer.data(); - } else { - int ret = llama_decode(ctx, batch); - if (ret != 0) { - LOG_ERR("%s: failed to decode at step %d, ret = %d\n", __func__, global_step, ret); - break; - } - logits = llama_get_logits(ctx); - } - - if (!logits) { - LOG_ERR("%s: failed to get logits at step %d\n", __func__, global_step); - break; - } - - auto get_logits_for_pos = [&](int32_t pos) -> const float * { - if (params.shift_logits) { - return pos == 0 ? logits : logits + (pos - 1) * n_vocab; - } - return logits + (pos) *n_vocab; - }; - - int64_t time_start_sampling = ggml_time_us(); - - mask_positions.clear(); - for (int32_t i = 0; i < params.max_length; i++) { - if (output_tokens[i] == params.mask_token_id) { - // For block-based, only consider current block - if (params.schedule != BLOCK_BASED || (i >= block_start && i < block_end)) { - mask_positions.push_back(i); - } - } - } - - if (mask_positions.empty()) { - break; - } - - if (params.add_gumbel_noise && params.temperature > 0.0f) { - add_gumbel_noise(logits, n_vocab, params.temperature, rng); - } - - if (params.algorithm == ORIGIN) { - int32_t transfer_count = calculate_transfer_count( - step, steps_per_block, mask_positions.size(), params.schedule, params.eps, num_transfer_tokens); - float p_transfer = (float) transfer_count / mask_positions.size(); - - for (int32_t pos : mask_positions) { - if (std::uniform_real_distribution(0.0f, 1.0f)(rng) < p_transfer) { - const float * pos_logits = get_logits_for_pos(pos); - for (int32_t token_id = 0; token_id < n_vocab; token_id++) { - candidates[token_id].id = token_id; - candidates[token_id].logit = pos_logits[token_id]; - candidates[token_id].p = 0.0f; - } - - llama_token_data_array cur_p = { - candidates.data(), - (size_t) n_vocab, - -1, - false, - }; - - llama_sampler_apply(sampler, &cur_p); - output_tokens[pos] = cur_p.data[cur_p.selected].id; - } - } - } else { - std::vector> confidences; - std::vector sampled_tokens(mask_positions.size()); - - for (size_t i = 0; i < mask_positions.size(); i++) { - int32_t pos = mask_positions[i]; - const float * pos_logits = get_logits_for_pos(pos); - - for (int32_t token_id = 0; token_id < n_vocab; token_id++) { - candidates[token_id].logit = pos_logits[token_id]; - candidates[token_id].p = 0.0f; - candidates[token_id].id = token_id; - } - - llama_token_data_array cur_p = { - candidates.data(), - candidates.size(), - -1, - false, - }; - - llama_sampler_apply(sampler, &cur_p); - llama_token sampled_token = cur_p.data[cur_p.selected].id; - - float conf = calculate_confidence(cur_p, params.algorithm, rng); - - sampled_tokens[i] = sampled_token; - confidences.emplace_back(conf, i); - } - - int32_t transfer_count = calculate_transfer_count( - step, steps_per_block, mask_positions.size(), params.schedule, params.eps, num_transfer_tokens); - - if (transfer_count > 0) { - if (params.alg_temp == 0.0f) { - std::partial_sort(confidences.begin(), - confidences.begin() + std::min(transfer_count, (int32_t) confidences.size()), - confidences.end(), - [](const std::pair & a, const std::pair & b) { - if (a.first != b.first) { - return a.first > b.first; - } - return a.second < b.second; - }); - - for (int32_t i = 0; i < std::min(transfer_count, (int32_t) confidences.size()); i++) { - int32_t mask_idx = confidences[i].second; - int32_t pos = mask_positions[mask_idx]; - output_tokens[pos] = sampled_tokens[mask_idx]; - } - } else { - conf_candidates.clear(); - for (size_t i = 0; i < confidences.size(); i++) { - float conf_logit = confidences[i].first / params.alg_temp; - conf_candidates.emplace_back(llama_token_data{ (int32_t) i, conf_logit, 0.0f }); - } - - llama_token_data_array conf_array = { - conf_candidates.data(), - conf_candidates.size(), - -1, - false, - }; - - for (int32_t i = 0; i < std::min(transfer_count, (int32_t) confidences.size()); i++) { - llama_sampler_apply(dist_sampler, &conf_array); - int32_t selected_idx = conf_array.selected; - int32_t mask_idx = selected_idx; - int32_t pos = mask_positions[mask_idx]; - output_tokens[pos] = sampled_tokens[mask_idx]; - - conf_candidates[selected_idx].p = 0.0f; - conf_array.selected = -1; - } - } - } - } - - int64_t time_end_sampling = ggml_time_us(); - total_sampling_time += time_end_sampling - time_start_sampling; - } - } - - int64_t time_end = ggml_time_us(); - total_time += time_end - time_start; - - LOG_INF("\ntotal time: %0.2fms, time per step: %0.2fms, sampling time per step: %0.2fms\n", - total_time / 1000.0, - total_time / 1000.0 / params.steps, - total_sampling_time / 1000.0 / params.steps); - - llama_batch_free(batch); - llama_sampler_free(sampler); - llama_sampler_free(dist_sampler); - - n_generated = params.max_length; -} - -static std::string format_input_text(const std::string & prompt, const std::string & system_prompt, bool use_chat_template, llama_model * model) { - if (!use_chat_template) { - return prompt; - } - - auto chat_templates = common_chat_templates_init(model, ""); - common_chat_templates_inputs inputs; - common_chat_msg system_msg; - - if (!system_prompt.empty()) { - system_msg.role = "system"; - system_msg.content = system_prompt; - inputs.messages.push_back(system_msg); - } - - common_chat_msg user_msg; - user_msg.role = "user"; - user_msg.content = prompt; - - inputs.messages.push_back(user_msg); - inputs.add_generation_prompt = true; - - auto result = common_chat_templates_apply(chat_templates.get(), inputs); - - return result.prompt; -} - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - ggml_time_init(); - - common_params params; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_DIFFUSION)) { - return 1; - } - - llama_backend_init(); - - llama_model_params model_params = llama_model_default_params(); - model_params.n_gpu_layers = params.n_gpu_layers; - model_params.devices = params.devices.data(); - model_params.use_mmap = params.use_mmap; - model_params.use_direct_io = params.use_direct_io; - model_params.use_mlock = params.use_mlock; - model_params.check_tensors = params.check_tensors; - - llama_model * model = llama_model_load_from_file(params.model.path.c_str(), model_params); - if (!model) { - LOG_ERR("error: failed to load model '%s'\n", params.model.path.c_str()); - return 1; - } - - if (!llama_model_is_diffusion(model)) { - LOG_ERR("error: unsupported model for diffusion"); - llama_model_free(model); - return 1; - } - - llama_context_params ctx_params = llama_context_default_params(); - ctx_params.n_ctx = params.n_ctx; - ctx_params.n_batch = params.n_batch; - ctx_params.n_ubatch = params.n_ubatch; - ctx_params.flash_attn_type = params.flash_attn_type; - ctx_params.no_perf = params.no_perf; - ctx_params.type_k = params.cache_type_k; - ctx_params.type_v = params.cache_type_v; - - llama_context * ctx = llama_init_from_model(model, ctx_params); - if (!ctx) { - LOG_ERR("error: failed to create context\n"); - llama_model_free(model); - return 1; - } - - llama_set_n_threads(ctx, params.cpuparams.n_threads, params.cpuparams_batch.n_threads); - - const llama_vocab * vocab = llama_model_get_vocab(model); - - std::string formatted_prompt = format_input_text(params.prompt, params.system_prompt, params.enable_chat_template, model); - - std::vector input_tokens = common_tokenize(vocab, - formatted_prompt, - /*add special tokens*/ true, - /*parse special*/ true); - - int n_input = input_tokens.size(); - - if (static_cast(n_input) >= llama_n_ctx(ctx)) { - LOG_ERR("error: input too long (%d tokens), max context is %d\n", n_input, llama_n_ctx(ctx)); - llama_free(ctx); - llama_model_free(model); - return 1; - } - - llama_token mask_token_id = llama_vocab_mask(vocab); - - GGML_ASSERT(mask_token_id != LLAMA_TOKEN_NULL); - - bool visual_mode = params.diffusion.visual_mode; - - int32_t n_generated = 0; - std::vector output_tokens(params.n_ubatch); - - struct diffusion_params diff_params; - - char shift_logits_str[8]; - if (llama_model_meta_val_str(model, "diffusion.shift_logits", shift_logits_str, sizeof(shift_logits_str)) >= 0) { - diff_params.shift_logits = (strcmp(shift_logits_str, "true") == 0); - } else { - diff_params.shift_logits = true; - } - - //Use either eps or block length, but not both - GGML_ASSERT((params.diffusion.eps == 0) ^ (params.diffusion.block_length == 0)); - - if (params.diffusion.eps) { - diff_params.schedule = TIMESTEP_BASED; - diff_params.eps = params.diffusion.eps; - } else if (params.diffusion.block_length) { - diff_params.schedule = BLOCK_BASED; - diff_params.block_length = params.diffusion.block_length; - } - - diff_params.mask_token_id = mask_token_id; - diff_params.seed = params.sampling.seed; - diff_params.temperature = params.sampling.temp; - diff_params.steps = params.diffusion.steps; - diff_params.algorithm = static_cast(params.diffusion.algorithm); - diff_params.max_length = params.n_ubatch; - diff_params.top_p = params.sampling.top_p; - diff_params.top_k = params.sampling.top_k; - diff_params.visual_mode = params.diffusion.visual_mode; - diff_params.add_gumbel_noise = params.diffusion.add_gumbel_noise; - - diff_params.step_callback = diffusion_step_callback; - callback_data cb_data = { &diff_params, vocab, n_input }; - diff_params.step_callback_user_data = &cb_data; - - const char * alg_names[] = { "ORIGIN", "ENTROPY_BASED", "MARGIN_BASED", "RANDOM", "CONFIDENCE_BASED" }; - const char * sched_names[] = { "TIMESTEP_BASED", "BLOCK_BASED" }; - const char * alg_name = - (diff_params.algorithm >= 0 && diff_params.algorithm <= 4) ? alg_names[diff_params.algorithm] : "UNKNOWN"; - const char * sched_name = - (diff_params.schedule >= 0 && diff_params.schedule <= 1) ? sched_names[diff_params.schedule] : "UNKNOWN"; - - LOG_INF("diffusion_params: - %-25s llama_token = %d\n", "mask_token_id", mask_token_id); - LOG_INF("diffusion_params: - %-25s u32 = %d\n", "steps", diff_params.steps); - LOG_INF("diffusion_params: - %-25s u32 = %d\n", "max_length", diff_params.max_length); - LOG_INF("diffusion_params: - %-25s enum = %d (%s)\n", "algorithm", diff_params.algorithm, alg_name); - LOG_INF("diffusion_params: - %-25s enum = %d (%s)\n", "schedule", diff_params.schedule, sched_name); - LOG_INF("diffusion_params: - %-25s f32 = %.3f\n", "temperature", diff_params.temperature); - if (diff_params.schedule == TIMESTEP_BASED) { - LOG_INF("diffusion_params: - %-25s f32 = %.6f\n", "eps", diff_params.eps); - LOG_INF("diffusion_params: - %-25s f32 = %.3f\n", "alg_temp", diff_params.alg_temp); - } - if (diff_params.schedule == BLOCK_BASED) { - LOG_INF("diffusion_params: - %-25s u32 = %d\n", "block_length", diff_params.block_length); - LOG_INF("diffusion_params: - %-25s f32 = %.3f\n", "cfg_scale", diff_params.cfg_scale); - } - - diffusion_generate(ctx, input_tokens.data(), output_tokens.data(), n_input, diff_params, n_generated); - - if (n_generated > 0) { - if (visual_mode) { - //clear screen and move cursor to top-left - LOG_INF("\033[2J\033[H"); - } - - output_tokens.erase(output_tokens.begin(), output_tokens.begin() + n_input); - std::string output_data = common_detokenize(vocab, output_tokens, false); - LOG_INF("\n%s\n", output_data.c_str()); - } else { - LOG_INF("Error: diffusion generation failed\n"); - } - - llama_free(ctx); - llama_model_free(model); - llama_backend_free(); - - return 0; -} diff --git a/examples/embedding-rust/Cargo.toml b/examples/embedding-rust/Cargo.toml new file mode 100644 index 000000000000..067e0ad0c53d --- /dev/null +++ b/examples/embedding-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-embedding-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-embedding" +path = "src/main.rs" diff --git a/examples/embedding-rust/build.rs b/examples/embedding-rust/build.rs new file mode 100644 index 000000000000..77a71515798e --- /dev/null +++ b/examples/embedding-rust/build.rs @@ -0,0 +1,45 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); + + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-embedding"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/embedding-rust/src/main.rs b/examples/embedding-rust/src/main.rs new file mode 100644 index 000000000000..d8f3ac0b4199 --- /dev/null +++ b/examples/embedding-rust/src/main.rs @@ -0,0 +1,966 @@ +use llama_simple_rust::ffi; +use std::env; +use std::ffi::{CStr, CString}; +use std::ptr; + +const LLAMA_POOLING_TYPE_NONE: i32 = 0; +const LLAMA_POOLING_TYPE_RANK: i32 = 4; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub embd_sep: String, + pub cls_sep: String, + pub embd_out: String, + pub embd_normalize: i32, + pub n_ctx: u32, + pub n_batch: u32, + pub n_parallel: u32, + pub n_gpu_layers: i32, + pub verbose_prompt: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + MissingModel, + MissingValue(String), + InvalidInteger(String, String), + InvalidValue(String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m/--model model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + ParseError::InvalidValue(message) => write!(f, "{message}"), + } + } +} + +impl std::error::Error for ParseError {} + +impl Default for Args { + fn default() -> Self { + Self { + model_path: String::new(), + prompt: String::new(), + embd_sep: "\n".to_string(), + cls_sep: "\t".to_string(), + embd_out: String::new(), + embd_normalize: 2, + n_ctx: 512, + n_batch: 512, + n_parallel: 1, + n_gpu_layers: 99, + verbose_prompt: false, + } + } +} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut parsed = Args::default(); + let mut iter = args.into_iter().map(Into::into); + + while let Some(arg) = iter.next() { + match arg.as_str() { + "-m" | "--model" => parsed.model_path = take(&mut iter, &arg)?, + "-p" | "--prompt" => parsed.prompt = take(&mut iter, &arg)?, + "--embd-sep" => parsed.embd_sep = take(&mut iter, &arg)?, + "--cls-sep" => parsed.cls_sep = take(&mut iter, &arg)?, + "--embd-output-format" | "--embd-out" => parsed.embd_out = take(&mut iter, &arg)?, + "--embd-normalize" => parsed.embd_normalize = parse_i32(&mut iter, &arg)?, + "-c" | "--ctx-size" => parsed.n_ctx = parse_u32(&mut iter, &arg)?, + "-b" | "--batch-size" => parsed.n_batch = parse_u32(&mut iter, &arg)?, + "-np" | "--parallel" => parsed.n_parallel = parse_u32(&mut iter, &arg)?, + "-ngl" | "--gpu-layers" => parsed.n_gpu_layers = parse_i32(&mut iter, &arg)?, + "--verbose-prompt" => parsed.verbose_prompt = true, + "-h" | "--help" => return Err(ParseError::MissingModel), + _ => { + if !arg.starts_with('-') { + if !parsed.prompt.is_empty() { + parsed.prompt.push(' '); + } + parsed.prompt.push_str(&arg); + } + } + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + if parsed.prompt.is_empty() { + parsed.prompt = "The quick brown fox jumps over the lazy dog".to_string(); + } + if parsed.embd_sep.is_empty() { + return Err(ParseError::InvalidValue( + "--embd-sep must not be empty".to_string(), + )); + } + if parsed.cls_sep.is_empty() { + return Err(ParseError::InvalidValue( + "--cls-sep must not be empty".to_string(), + )); + } + if !matches!( + parsed.embd_out.as_str(), + "" | "json" | "json+" | "array" | "raw" + ) { + return Err(ParseError::InvalidValue(format!( + "unsupported embedding output format: {}", + parsed.embd_out + ))); + } + if parsed.n_parallel == 0 { + return Err(ParseError::InvalidValue( + "--parallel must be positive".to_string(), + )); + } + if parsed.n_batch < parsed.n_ctx { + parsed.n_batch = parsed.n_ctx; + } + + Ok(parsed) +} + +fn take(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + iter.next() + .ok_or_else(|| ParseError::MissingValue(flag.to_string())) +} + +fn parse_i32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_u32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +pub fn split_lines(text: &str, separator: &str) -> Vec { + text.split(separator).map(str::to_string).collect() +} + +pub fn normalize_embedding(input: &[f32], mode: i32) -> Vec { + match mode { + 0 => input.to_vec(), + -1 => { + let sum = input.iter().map(|x| f64::from(*x)).sum::(); + let mean = if input.is_empty() { + 0.0 + } else { + sum / input.len() as f64 + }; + let variance = input + .iter() + .map(|x| { + let d = f64::from(*x) - mean; + d * d + }) + .sum::() + / input.len().max(1) as f64; + let scale = if variance > 0.0 { + 1.0 / variance.sqrt() + } else { + 0.0 + }; + input + .iter() + .map(|x| ((f64::from(*x) - mean) * scale) as f32) + .collect() + } + _ => { + let sum = input + .iter() + .map(|x| f64::from(*x) * f64::from(*x)) + .sum::() + .sqrt(); + let norm = if sum > 0.0 { 1.0 / sum } else { 0.0 }; + input + .iter() + .map(|x| (f64::from(*x) * norm) as f32) + .collect() + } + } +} + +pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { + let mut sum = 0.0f64; + let mut sum_a = 0.0f64; + let mut sum_b = 0.0f64; + for (x, y) in a.iter().zip(b.iter()) { + let x = f64::from(*x); + let y = f64::from(*y); + sum += x * y; + sum_a += x * x; + sum_b += y * y; + } + if sum_a == 0.0 || sum_b == 0.0 { + return if sum_a == 0.0 && sum_b == 0.0 { + 1.0 + } else { + 0.0 + }; + } + (sum / (sum_a.sqrt() * sum_b.sqrt())) as f32 +} + +fn replace_all(mut text: String, from: &str, to: &str) -> String { + text = text.replace(from, to); + text +} + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Batch(ffi::llama_batch); + +impl Drop for Batch { + fn drop(&mut self) { + unsafe { + ffi::llama_batch_free(self.0); + } + } +} + +struct Backend; + +impl Backend { + fn init() -> Self { + unsafe { + ffi::ggml_backend_load_all(); + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + Self + } +} + +impl Drop for Backend { + fn drop(&mut self) { + unsafe { + ffi::llama_backend_free(); + } + } +} + +fn cstr_to_string(ptr: *const i8) -> String { + if ptr.is_null() { + String::new() + } else { + unsafe { CStr::from_ptr(ptr).to_string_lossy().into_owned() } + } +} + +fn tokenize( + vocab: *const ffi::llama_vocab, + text: &str, + add_special: bool, + parse_special: bool, +) -> Result, String> { + let text_c = + CString::new(text).map_err(|_| "text contains an interior NUL byte".to_string())?; + let n_tokens = unsafe { + -ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + ptr::null_mut(), + 0, + add_special, + parse_special, + ) + }; + if n_tokens <= 0 { + return Err("failed to size tokenization".to_string()); + } + let mut tokens = vec![0_i32; n_tokens as usize]; + let n = unsafe { + ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + add_special, + parse_special, + ) + }; + if n < 0 { + return Err("failed to tokenize text".to_string()); + } + tokens.truncate(n as usize); + Ok(tokens) +} + +fn token_to_piece( + vocab: *const ffi::llama_vocab, + token: ffi::llama_token, +) -> Result { + let mut buf = vec![0_i8; 128]; + let n = unsafe { + ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), buf.len() as i32, 0, true) + }; + if n < 0 { + return Err("failed to convert token to piece".to_string()); + } + let bytes = buf[..n as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +fn batch_clear(batch: &mut ffi::llama_batch) { + batch.n_tokens = 0; +} + +unsafe fn batch_add_seq( + batch: &mut ffi::llama_batch, + tokens: &[ffi::llama_token], + seq_id: ffi::llama_seq_id, +) { + for (pos, token) in tokens.iter().enumerate() { + let i = batch.n_tokens as isize; + *batch.token.offset(i) = *token; + *batch.pos.offset(i) = pos as ffi::llama_pos; + *batch.n_seq_id.offset(i) = 1; + let seq_slot = *batch.seq_id.offset(i); + *seq_slot = seq_id; + *batch.logits.offset(i) = 1; + batch.n_tokens += 1; + } +} + +fn batch_decode( + ctx: *mut ffi::llama_context, + batch: &mut ffi::llama_batch, + output: &mut [f32], + n_embd_out: usize, + embd_norm: i32, +) -> Result<(), String> { + let pooling_type = unsafe { ffi::llama_pooling_type(ctx) }; + unsafe { + ffi::llama_memory_clear(ffi::llama_get_memory(ctx), true); + eprintln!( + "batch_decode: n_tokens = {}, n_seq = {}", + batch.n_tokens, + ffi::llama_n_seq_max(ctx) + ); + if ffi::llama_decode(ctx, *batch) < 0 { + return Err("failed to process batch".to_string()); + } + + for i in 0..batch.n_tokens { + if *batch.logits.offset(i as isize) == 0 { + continue; + } + + let (emb, emb_pos) = if pooling_type == LLAMA_POOLING_TYPE_NONE { + (ffi::llama_get_embeddings_ith(ctx, i), i as usize) + } else { + let seq_id = **batch.seq_id.offset(i as isize); + (ffi::llama_get_embeddings_seq(ctx, seq_id), seq_id as usize) + }; + + if emb.is_null() { + return Err(format!("failed to get embeddings for batch item {i}")); + } + let start = emb_pos * n_embd_out; + if start + n_embd_out > output.len() { + return Err("embedding output buffer is too small".to_string()); + } + let raw = std::slice::from_raw_parts(emb, n_embd_out); + let normalized = normalize_embedding(raw, embd_norm); + output[start..start + n_embd_out].copy_from_slice(&normalized); + } + } + Ok(()) +} + +fn build_rank_prompt( + prompt: &str, + cls_sep: &str, + rerank_template: Option<&str>, + added_eos_token: &str, + added_sep_token: &str, +) -> String { + let pairs = split_lines(prompt, cls_sep); + if let (Some(template), Some(query), Some(document)) = + (rerank_template, pairs.first(), pairs.get(1)) + { + return replace_all( + replace_all(template.to_string(), "{query}", query), + "{document}", + document, + ); + } + let mut final_prompt = String::new(); + for (i, pair) in pairs.iter().enumerate() { + final_prompt.push_str(pair); + if i + 1 != pairs.len() { + final_prompt.push_str(added_eos_token); + final_prompt.push_str(added_sep_token); + } + } + final_prompt +} + +fn prepare_inputs( + args: &Args, + vocab: *const ffi::llama_vocab, + model: *const ffi::llama_model, + pooling_type: i32, +) -> Result<(Vec, Vec>), String> { + let prompts = split_lines(&args.prompt, &args.embd_sep); + let sep = unsafe { ffi::llama_vocab_sep(vocab) }; + let eos = unsafe { ffi::llama_vocab_eos(vocab) }; + let added_sep_token = if unsafe { ffi::llama_vocab_get_add_sep(vocab) } { + cstr_to_string(unsafe { ffi::llama_vocab_get_text(vocab, sep) }) + } else { + String::new() + }; + let added_eos_token = if unsafe { ffi::llama_vocab_get_add_eos(vocab) } { + cstr_to_string(unsafe { ffi::llama_vocab_get_text(vocab, eos) }) + } else { + String::new() + }; + let rerank_name = CString::new("rerank").unwrap(); + let rerank_ptr = unsafe { ffi::llama_model_chat_template(model, rerank_name.as_ptr()) }; + let rerank_template = if rerank_ptr.is_null() { + None + } else { + Some(cstr_to_string(rerank_ptr)) + }; + + let mut inputs = Vec::with_capacity(prompts.len()); + for prompt in &prompts { + let text = if pooling_type == LLAMA_POOLING_TYPE_RANK && prompt.contains(&args.cls_sep) { + build_rank_prompt( + prompt, + &args.cls_sep, + rerank_template.as_deref(), + &added_eos_token, + &added_sep_token, + ) + } else { + prompt.clone() + }; + inputs.push(tokenize(vocab, &text, true, true)?); + } + Ok((prompts, inputs)) +} + +fn print_json( + embeddings: &[f32], + n_embd_count: usize, + n_embd_out: usize, + n_prompts: usize, + embd_out: &str, + embd_normalize: i32, +) { + let not_array = embd_out != "array"; + if not_array { + println!("{{\n \"object\": \"list\",\n \"data\": ["); + } else { + print!("["); + } + + for j in 0..n_embd_count { + if not_array { + print!(" {{\n \"object\": \"embedding\",\n \"index\": {j},\n \"embedding\": "); + } + print!("["); + for i in 0..n_embd_out { + let value = embeddings[j * n_embd_out + i]; + if embd_normalize == 0 { + print!("{value:.0}"); + } else { + print!("{value:.7}"); + } + if i + 1 < n_embd_out { + print!(","); + } + } + if not_array { + print!("]\n }}"); + } else { + print!("]"); + } + if j + 1 < n_embd_count { + if not_array { + println!(","); + } else { + print!(","); + } + } + } + + if not_array { + print!("\n ]"); + } else { + println!("]"); + } + + if embd_out == "json+" && n_prompts > 1 { + println!(",\n \"cosineSimilarity\": ["); + for i in 0..n_embd_count { + print!(" ["); + for j in 0..n_embd_count { + let sim = cosine_similarity( + &embeddings[i * n_embd_out..(i + 1) * n_embd_out], + &embeddings[j * n_embd_out..(j + 1) * n_embd_out], + ); + print!("{sim:6.2}"); + if j + 1 < n_embd_count { + print!(", "); + } + } + print!(" ]"); + if i + 1 < n_embd_count { + println!(","); + } + } + print!("\n ]"); + } + + if not_array { + println!("\n}}"); + } +} + +fn print_raw_embeddings( + embeddings: &[f32], + n_embd_count: usize, + n_embd_out: usize, + model: *const ffi::llama_model, + pooling_type: i32, + embd_normalize: i32, +) { + let n_cls_out = unsafe { ffi::llama_model_n_cls_out(model) as usize }; + let cols = if pooling_type == LLAMA_POOLING_TYPE_RANK { + n_embd_out.min(n_cls_out) + } else { + n_embd_out + }; + + for j in 0..n_embd_count { + for i in 0..cols { + let value = embeddings[j * n_embd_out + i]; + if embd_normalize == 0 { + print!("{value:.0}"); + } else { + print!("{value:.7}"); + } + if i + 1 < cols { + print!(" "); + } + } + println!(); + } +} + +fn print_human_embeddings( + embeddings: &[f32], + prompts: &[String], + n_embd_count: usize, + n_embd_out: usize, + model: *const ffi::llama_model, + pooling_type: i32, + embd_normalize: i32, +) { + println!(); + if pooling_type == LLAMA_POOLING_TYPE_NONE { + for j in 0..n_embd_count { + print!("embedding {j}: "); + for i in 0..3.min(n_embd_out) { + print_value(embeddings[j * n_embd_out + i], embd_normalize); + print!(" "); + } + print!(" ... "); + for i in n_embd_out.saturating_sub(3)..n_embd_out { + print_value(embeddings[j * n_embd_out + i], embd_normalize); + print!(" "); + } + println!(); + } + } else if pooling_type == LLAMA_POOLING_TYPE_RANK { + let n_cls_out = unsafe { ffi::llama_model_n_cls_out(model) }; + let labels = (0..n_cls_out) + .map(|i| { + let label = cstr_to_string(unsafe { ffi::llama_model_cls_label(model, i) }); + if label.is_empty() { + i.to_string() + } else { + label + } + }) + .collect::>(); + for j in 0..n_embd_count { + for i in 0..n_cls_out as usize { + if n_cls_out == 1 { + println!("rerank score {j}: {:8.3}", embeddings[j * n_embd_out]); + } else { + println!( + "rerank score {j}: {:8.3} [{}]", + embeddings[j * n_embd_out + i], + labels[i] + ); + } + } + } + } else { + for (j, _) in prompts.iter().enumerate() { + print!("embedding {j}: "); + let cols = if prompts.len() > 1 { + 16.min(n_embd_out) + } else { + n_embd_out + }; + for i in 0..cols { + print_value(embeddings[j * n_embd_out + i], embd_normalize); + print!(" "); + } + println!(); + } + + if prompts.len() > 1 { + println!("\ncosine similarity matrix:\n"); + for prompt in prompts { + print!("{:.6} ", prompt); + } + println!(); + for i in 0..prompts.len() { + for j in 0..prompts.len() { + let sim = cosine_similarity( + &embeddings[i * n_embd_out..(i + 1) * n_embd_out], + &embeddings[j * n_embd_out..(j + 1) * n_embd_out], + ); + print!("{sim:6.2} "); + } + println!("{:.10}", prompts[i]); + } + } + } +} + +fn print_value(value: f32, embd_normalize: i32) { + if embd_normalize == 0 { + print!("{value:6.0}"); + } else { + print!("{value:9.6}"); + } +} + +fn run(args: Args) -> Result<(), String> { + let _backend = Backend::init(); + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + + let max_parallel = unsafe { ffi::llama_max_parallel_sequences() as u32 }; + let mut n_parallel = args.n_parallel; + let kv_unified = if n_parallel == 1 { + eprintln!("n_parallel == 1 -> unified KV cache is enabled"); + n_parallel = max_parallel; + true + } else { + false + }; + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = args.n_ctx; + ctx_params.n_batch = args.n_batch; + ctx_params.n_ubatch = args.n_batch; + ctx_params.n_seq_max = n_parallel; + ctx_params.embeddings = true; + ctx_params.kv_unified = kv_unified; + + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create llama_context".to_string()); + } + + if unsafe { ffi::llama_model_has_encoder(model.0) && ffi::llama_model_has_decoder(model.0) } { + return Err("computing embeddings in encoder-decoder models is not supported".to_string()); + } + + let n_ctx_train = unsafe { ffi::llama_model_n_ctx_train(model.0) }; + let n_ctx = unsafe { ffi::llama_n_ctx(ctx.0) as i32 }; + if n_ctx > n_ctx_train { + eprintln!( + "warning: model was trained on only {n_ctx_train} context tokens ({n_ctx} specified)" + ); + } + + let pooling_type = unsafe { ffi::llama_pooling_type(ctx.0) }; + let (prompts, inputs) = prepare_inputs(&args, vocab, model.0, pooling_type)?; + let n_batch = args.n_batch as usize; + for (i, input) in inputs.iter().enumerate() { + if input.len() > n_batch { + return Err(format!( + "number of tokens in input line ({}) exceeds batch size ({n_batch}), increase batch size and re-run", + input.len() + )); + } + let sep = unsafe { ffi::llama_vocab_sep(vocab) }; + let eos = unsafe { ffi::llama_vocab_eos(vocab) }; + if input.last().copied() != Some(sep) && input.last().copied() != Some(eos) { + eprintln!("warning: last token in prompt {i} is not SEP or EOS"); + } + if args.verbose_prompt { + eprintln!("prompt {i}: '{}'", prompts[i]); + eprintln!("number of tokens in prompt = {}", input.len()); + for token in input { + eprintln!("{token:6} -> '{}'", token_to_piece(vocab, *token)?); + } + } + } + + let n_embd_count = if pooling_type == LLAMA_POOLING_TYPE_NONE { + inputs.iter().map(Vec::len).sum() + } else { + prompts.len() + }; + let n_embd_out = unsafe { ffi::llama_model_n_embd_out(model.0) as usize }; + let mut embeddings = vec![0.0f32; n_embd_count * n_embd_out]; + let mut batch = Batch(unsafe { ffi::llama_batch_init(n_batch as i32, 0, 1) }); + + let mut e = 0usize; + let mut s = 0usize; + let n_seq_max = unsafe { ffi::llama_n_seq_max(ctx.0) as usize }; + for input in &inputs { + if batch.0.n_tokens as usize + input.len() > n_batch || s >= n_seq_max { + let next_e = e + if pooling_type == LLAMA_POOLING_TYPE_NONE { + batch.0.n_tokens as usize + } else { + s + }; + batch_decode( + ctx.0, + &mut batch.0, + &mut embeddings[e * n_embd_out..next_e * n_embd_out], + n_embd_out, + args.embd_normalize, + )?; + e = next_e; + s = 0; + batch_clear(&mut batch.0); + } + unsafe { + batch_add_seq(&mut batch.0, input, s as i32); + } + s += 1; + } + + if s > 0 { + let next_e = e + if pooling_type == LLAMA_POOLING_TYPE_NONE { + batch.0.n_tokens as usize + } else { + s + }; + batch_decode( + ctx.0, + &mut batch.0, + &mut embeddings[e * n_embd_out..next_e * n_embd_out], + n_embd_out, + args.embd_normalize, + )?; + } + + match args.embd_out.as_str() { + "" => print_human_embeddings( + &embeddings, + &prompts, + n_embd_count, + n_embd_out, + model.0, + pooling_type, + args.embd_normalize, + ), + "json" | "json+" | "array" => print_json( + &embeddings, + n_embd_count, + n_embd_out, + prompts.len(), + &args.embd_out, + args.embd_normalize, + ), + "raw" => print_raw_embeddings( + &embeddings, + n_embd_count, + n_embd_out, + model.0, + pooling_type, + args.embd_normalize, + ), + _ => unreachable!(), + } + + println!(); + unsafe { + ffi::llama_perf_context_print(ctx.0); + } + Ok(()) +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m ./models/bge-base-en-v1.5-f16.gguf -p \"hello world\""); + eprintln!(); +} + +fn main() { + let mut argv = env::args(); + let program = argv.next().unwrap_or_else(|| "llama-embedding".to_string()); + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_embedding_options() { + let args = parse_args([ + "-m", + "model.gguf", + "-p", + "a\nb", + "--embd-sep", + "\n", + "--cls-sep", + " || ", + "--embd-out", + "json+", + "--embd-normalize", + "0", + "-c", + "256", + "-b", + "128", + "-np", + "4", + "-ngl", + "0", + "--verbose-prompt", + ]) + .unwrap(); + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "a\nb"); + assert_eq!(args.cls_sep, " || "); + assert_eq!(args.embd_out, "json+"); + assert_eq!(args.embd_normalize, 0); + assert_eq!(args.n_ctx, 256); + assert_eq!(args.n_batch, 256); + assert_eq!(args.n_parallel, 4); + assert_eq!(args.n_gpu_layers, 0); + assert!(args.verbose_prompt); + } + + #[test] + fn rejects_bad_output_format() { + assert_eq!( + parse_args(["--model", "model.gguf", "--embd-out", "yaml"]).unwrap_err(), + ParseError::InvalidValue("unsupported embedding output format: yaml".to_string()) + ); + } + + #[test] + fn splits_on_custom_separator() { + assert_eq!( + split_lines("alpha || beta || gamma", " || "), + vec!["alpha", "beta", "gamma"] + ); + } + + #[test] + fn normalizes_embeddings() { + let normalized = normalize_embedding(&[3.0, 4.0], 2); + assert!((normalized[0] - 0.6).abs() < 1e-6); + assert!((normalized[1] - 0.8).abs() < 1e-6); + assert_eq!(normalize_embedding(&[3.0, 4.0], 0), vec![3.0, 4.0]); + } + + #[test] + fn builds_rank_prompt_from_template_or_tokens() { + let templated = build_rank_prompt( + "question\tanswer", + "\t", + Some("Q: {query}\nD: {document}"), + "", + "[SEP]", + ); + assert_eq!(templated, "Q: question\nD: answer"); + let joined = build_rank_prompt("question\tanswer", "\t", None, "", "[SEP]"); + assert_eq!(joined, "question[SEP]answer"); + } +} diff --git a/examples/embedding/CMakeLists.txt b/examples/embedding/CMakeLists.txt index 0634c7bd8202..1313803e0eb9 100644 --- a/examples/embedding/CMakeLists.txt +++ b/examples/embedding/CMakeLists.txt @@ -1,5 +1,25 @@ set(TARGET llama-embedding) -add_executable(${TARGET} embedding.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(EMBEDDING_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../embedding-rust/Cargo.toml) +set(EMBEDDING_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/embedding) +set(EMBEDDING_RUST_BIN ${EMBEDDING_RUST_TARGET_DIR}/debug/llama-embedding${CMAKE_EXECUTABLE_SUFFIX}) +set(EMBEDDING_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${EMBEDDING_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${EMBEDDING_RUST_MANIFEST} --target-dir ${EMBEDDING_RUST_TARGET_DIR} --bin llama-embedding + COMMAND ${CMAKE_COMMAND} -E copy ${EMBEDDING_RUST_BIN} ${EMBEDDING_RUST_OUTPUT} + DEPENDS + llama + ${EMBEDDING_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../embedding-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../embedding-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${EMBEDDING_RUST_OUTPUT}) +install(PROGRAMS ${EMBEDDING_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/embedding/embedding.cpp b/examples/embedding/embedding.cpp deleted file mode 100644 index f6a20ef9d074..000000000000 --- a/examples/embedding/embedding.cpp +++ /dev/null @@ -1,414 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "log.h" -#include "llama.h" - -#include -#include -#include - -#if defined(_MSC_VER) -#pragma warning(disable: 4244 4267) // possible loss of data -#endif - -static std::vector split_lines(const std::string & s, const std::string & separator = "\n") { - std::vector lines; - size_t start = 0; - size_t end = s.find(separator); - - while (end != std::string::npos) { - lines.push_back(s.substr(start, end - start)); - start = end + separator.length(); - end = s.find(separator, start); - } - - lines.push_back(s.substr(start)); // Add the last part - - return lines; -} - -static void batch_add_seq(llama_batch & batch, const std::vector & tokens, llama_seq_id seq_id) { - size_t n_tokens = tokens.size(); - for (size_t i = 0; i < n_tokens; i++) { - common_batch_add(batch, tokens[i], i, { seq_id }, true); - } -} - -static void batch_decode(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd_out, int embd_norm) { - const enum llama_pooling_type pooling_type = llama_pooling_type(ctx); - - // clear previous kv_cache values (irrelevant for embeddings) - llama_memory_clear(llama_get_memory(ctx), true); - - // run model - LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq); - if (llama_decode(ctx, batch) < 0) { - LOG_ERR("%s : failed to process\n", __func__); - } - - for (int i = 0; i < batch.n_tokens; i++) { - if (!batch.logits[i]) { - continue; - } - - const float * embd = nullptr; - int embd_pos = 0; - - if (pooling_type == LLAMA_POOLING_TYPE_NONE) { - // try to get token embeddings - embd = llama_get_embeddings_ith(ctx, i); - embd_pos = i; - GGML_ASSERT(embd != NULL && "failed to get token embeddings"); - } else { - // try to get sequence embeddings - supported only when pooling_type is not NONE - embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]); - embd_pos = batch.seq_id[i][0]; - GGML_ASSERT(embd != NULL && "failed to get sequence embeddings"); - } - - float * out = output + embd_pos * n_embd_out; - common_embd_normalize(embd, out, n_embd_out, embd_norm); - } -} - -// plain, pipe-friendly output: one embedding per line -static void print_raw_embeddings(const float * emb, - int n_embd_count, - int n_embd, - const llama_model * model, - enum llama_pooling_type pooling_type, - int embd_normalize) { - const uint32_t n_cls_out = llama_model_n_cls_out(model); - const bool is_rank = (pooling_type == LLAMA_POOLING_TYPE_RANK); - const int cols = is_rank ? std::min(n_embd, (int) n_cls_out) : n_embd; - - for (int j = 0; j < n_embd_count; ++j) { - for (int i = 0; i < cols; ++i) { - if (embd_normalize == 0) { - LOG("%1.0f%s", emb[j * n_embd + i], (i + 1 < cols ? " " : "")); - } else { - LOG("%1.7f%s", emb[j * n_embd + i], (i + 1 < cols ? " " : "")); - } - } - LOG("\n"); - } -} - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_EMBEDDING)) { - return 1; - } - - params.embedding = true; - - // get max number of sequences per batch - const int n_seq_max = llama_max_parallel_sequences(); - - // if the number of prompts that would be encoded is known in advance, it's more efficient to specify the - // --parallel argument accordingly. for convenience, if not specified, we fallback to unified KV cache - // in order to support any number of prompts - if (params.n_parallel == 1) { - LOG_INF("%s: n_parallel == 1 -> unified KV cache is enabled\n", __func__); - params.kv_unified = true; - params.n_parallel = n_seq_max; - } - - // utilize the full context - if (params.n_batch < params.n_ctx) { - LOG_WRN("%s: setting batch size to %d\n", __func__, params.n_ctx); - params.n_batch = params.n_ctx; - } - - // for non-causal models, batch size must be equal to ubatch size - if (params.attention_type != LLAMA_ATTENTION_TYPE_CAUSAL) { - params.n_ubatch = params.n_batch; - } - - llama_backend_init(); - llama_numa_init(params.numa); - - // load the model - auto llama_init = common_init_from_params(params); - - auto * model = llama_init->model(); - auto * ctx = llama_init->context(); - - if (model == NULL) { - LOG_ERR("%s: unable to load model\n", __func__); - return 1; - } - - const llama_vocab * vocab = llama_model_get_vocab(model); - - const int n_ctx_train = llama_model_n_ctx_train(model); - const int n_ctx = llama_n_ctx(ctx); - - const enum llama_pooling_type pooling_type = llama_pooling_type(ctx); - - if (llama_model_has_encoder(model) && llama_model_has_decoder(model)) { - LOG_ERR("%s: computing embeddings in encoder-decoder models is not supported\n", __func__); - return 1; - } - - if (n_ctx > n_ctx_train) { - LOG_WRN("%s: warning: model was trained on only %d context tokens (%d specified)\n", - __func__, n_ctx_train, n_ctx); - } - - // print system information - { - LOG_INF("\n"); - LOG_INF("%s\n", common_params_get_system_info(params).c_str()); - } - - // split the prompt into lines - std::vector prompts = split_lines(params.prompt, params.embd_sep); - - // max batch size - const uint64_t n_batch = params.n_batch; - - // get added sep and eos token, if any - const std::string added_sep_token = llama_vocab_get_add_sep(vocab) ? llama_vocab_get_text(vocab, llama_vocab_sep(vocab)) : ""; - const std::string added_eos_token = llama_vocab_get_add_eos(vocab) ? llama_vocab_get_text(vocab, llama_vocab_eos(vocab)) : ""; - const char * rerank_prompt = llama_model_chat_template(model, "rerank"); - - // tokenize the prompts and trim - std::vector> inputs; - for (const auto & prompt : prompts) { - std::vector inp; - - // split classification pairs and insert expected separator tokens - if (pooling_type == LLAMA_POOLING_TYPE_RANK && prompt.find(params.cls_sep) != std::string::npos) { - std::vector pairs = split_lines(prompt, params.cls_sep); - if (rerank_prompt != nullptr) { - const std::string query = pairs[0]; - const std::string doc = pairs[1]; - std::string final_prompt = rerank_prompt; - string_replace_all(final_prompt, "{query}" , query); - string_replace_all(final_prompt, "{document}", doc ); - inp = common_tokenize(vocab, final_prompt, true, true); - } else { - std::string final_prompt; - for (size_t i = 0; i < pairs.size(); i++) { - final_prompt += pairs[i]; - if (i != pairs.size() - 1) { - if (!added_eos_token.empty()) { - final_prompt += added_eos_token; - } - if (!added_sep_token.empty()) { - final_prompt += added_sep_token; - } - } - } - inp = common_tokenize(ctx, final_prompt, true, true); - } - } else { - inp = common_tokenize(ctx, prompt, true, true); - } - if (inp.size() > n_batch) { - LOG_ERR("%s: number of tokens in input line (%lld) exceeds batch size (%lld), increase batch size and re-run\n", - __func__, (long long int) inp.size(), (long long int) n_batch); - return 1; - } - inputs.push_back(inp); - } - - // check if the last token is SEP/EOS - // it should be automatically added by the tokenizer when 'tokenizer.ggml.add_eos_token' is set to 'true' - for (auto & inp : inputs) { - if (inp.empty() || (inp.back() != llama_vocab_sep(vocab) && inp.back() != llama_vocab_eos(vocab))) { - LOG_WRN("%s: last token in the prompt is not SEP or EOS\n", __func__); - LOG_WRN("%s: 'tokenizer.ggml.add_eos_token' should be set to 'true' in the GGUF header\n", __func__); - } - } - - // tokenization stats - if (params.verbose_prompt) { - for (int i = 0; i < (int) inputs.size(); i++) { - LOG_INF("%s: prompt %d: '%s'\n", __func__, i, prompts[i].c_str()); - LOG_INF("%s: number of tokens in prompt = %zu\n", __func__, inputs[i].size()); - for (int j = 0; j < (int) inputs[i].size(); j++) { - LOG("%6d -> '%s'\n", inputs[i][j], common_token_to_piece(ctx, inputs[i][j]).c_str()); - } - LOG("\n\n"); - } - } - - // initialize batch - const int n_prompts = prompts.size(); - struct llama_batch batch = llama_batch_init(n_batch, 0, 1); - - // count number of embeddings - int n_embd_count = 0; - if (pooling_type == LLAMA_POOLING_TYPE_NONE) { - for (int k = 0; k < n_prompts; k++) { - n_embd_count += inputs[k].size(); - } - } else { - n_embd_count = n_prompts; - } - - // allocate output - const int n_embd_out = llama_model_n_embd_out(model); - std::vector embeddings(n_embd_count * n_embd_out, 0); - float * emb = embeddings.data(); - - // break into batches - int e = 0; // number of embeddings already stored - int s = 0; // number of prompts in current batch - for (int k = 0; k < n_prompts; k++) { - // clamp to n_batch tokens - auto & inp = inputs[k]; - - const uint64_t n_toks = inp.size(); - - // encode if at capacity - if (batch.n_tokens + n_toks > n_batch || s >= n_seq_max) { - float * out = emb + e * n_embd_out; - batch_decode(ctx, batch, out, s, n_embd_out, params.embd_normalize); - e += pooling_type == LLAMA_POOLING_TYPE_NONE ? batch.n_tokens : s; - s = 0; - common_batch_clear(batch); - } - - // add to batch - batch_add_seq(batch, inp, s); - s += 1; - } - - // final batch - float * out = emb + e * n_embd_out; - batch_decode(ctx, batch, out, s, n_embd_out, params.embd_normalize); - - if (params.embd_out.empty()) { - LOG("\n"); - - if (pooling_type == LLAMA_POOLING_TYPE_NONE) { - for (int j = 0; j < n_embd_count; j++) { - LOG("embedding %d: ", j); - for (int i = 0; i < std::min(3, n_embd_out); i++) { - if (params.embd_normalize == 0) { - LOG("%6.0f ", emb[j * n_embd_out + i]); - } else { - LOG("%9.6f ", emb[j * n_embd_out + i]); - } - } - LOG(" ... "); - for (int i = n_embd_out - 3; i < n_embd_out; i++) { - if (params.embd_normalize == 0) { - LOG("%6.0f ", emb[j * n_embd_out + i]); - } else { - LOG("%9.6f ", emb[j * n_embd_out + i]); - } - } - LOG("\n"); - } - } else if (pooling_type == LLAMA_POOLING_TYPE_RANK) { - const uint32_t n_cls_out = llama_model_n_cls_out(model); - std::vector cls_out_labels; - - for (uint32_t i = 0; i < n_cls_out; i++) { - const char * label = llama_model_cls_label(model, i); - const std::string label_i(label == nullptr ? "" : label); - cls_out_labels.emplace_back(label_i.empty() ? std::to_string(i) : label_i); - } - - for (int j = 0; j < n_embd_count; j++) { - for (uint32_t i = 0; i < n_cls_out; i++) { - // NOTE: if you change this log - update the tests in ci/run.sh - if (n_cls_out == 1) { - LOG("rerank score %d: %8.3f\n", j, emb[j * n_embd_out]); - } else { - LOG("rerank score %d: %8.3f [%s]\n", j, emb[j * n_embd_out + i], cls_out_labels[i].c_str()); - } - } - } - } else { - // print the first part of the embeddings or for a single prompt, the full embedding - for (int j = 0; j < n_prompts; j++) { - LOG("embedding %d: ", j); - for (int i = 0; i < (n_prompts > 1 ? std::min(16, n_embd_out) : n_embd_out); i++) { - if (params.embd_normalize == 0) { - LOG("%6.0f ", emb[j * n_embd_out + i]); - } else { - LOG("%9.6f ", emb[j * n_embd_out + i]); - } - } - LOG("\n"); - } - - // print cosine similarity matrix - if (n_prompts > 1) { - LOG("\n"); - LOG("cosine similarity matrix:\n\n"); - for (int i = 0; i < n_prompts; i++) { - LOG("%6.6s ", prompts[i].c_str()); - } - LOG("\n"); - for (int i = 0; i < n_prompts; i++) { - for (int j = 0; j < n_prompts; j++) { - float sim = common_embd_similarity_cos(emb + i * n_embd_out, emb + j * n_embd_out, n_embd_out); - LOG("%6.2f ", sim); - } - LOG("%1.10s", prompts[i].c_str()); - LOG("\n"); - } - } - } - } - - if (params.embd_out == "json" || params.embd_out == "json+" || params.embd_out == "array") { - const bool notArray = params.embd_out != "array"; - - LOG(notArray ? "{\n \"object\": \"list\",\n \"data\": [\n" : "["); - for (int j = 0;;) { // at least one iteration (one prompt) - if (notArray) LOG(" {\n \"object\": \"embedding\",\n \"index\": %d,\n \"embedding\": ",j); - LOG("["); - for (int i = 0;;) { // at least one iteration (n_embd > 0) - LOG(params.embd_normalize == 0 ? "%1.0f" : "%1.7f", emb[j * n_embd_out + i]); - i++; - if (i < n_embd_out) LOG(","); else break; - } - LOG(notArray ? "]\n }" : "]"); - j++; - if (j < n_embd_count) LOG(notArray ? ",\n" : ","); else break; - } - LOG(notArray ? "\n ]" : "]\n"); - - if (params.embd_out == "json+" && n_prompts > 1) { - LOG(",\n \"cosineSimilarity\": [\n"); - for (int i = 0;;) { // at least two iteration (n_embd_count > 1) - LOG(" ["); - for (int j = 0;;) { // at least two iteration (n_embd_count > 1) - float sim = common_embd_similarity_cos(emb + i * n_embd_out, emb + j * n_embd_out, n_embd_out); - LOG("%6.2f", sim); - j++; - if (j < n_embd_count) LOG(", "); else break; - } - LOG(" ]"); - i++; - if (i < n_embd_count) LOG(",\n"); else break; - } - LOG("\n ]"); - } - - if (notArray) LOG("\n}\n"); - } else if (params.embd_out == "raw") { - print_raw_embeddings(emb, n_embd_count, n_embd_out, model, pooling_type, params.embd_normalize); - } - - LOG("\n"); - llama_perf_context_print(ctx); - - // clean up - llama_batch_free(batch); - llama_backend_free(); - - return 0; -} diff --git a/examples/eval-callback-rust/Cargo.toml b/examples/eval-callback-rust/Cargo.toml new file mode 100644 index 000000000000..918112ab321e --- /dev/null +++ b/examples/eval-callback-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-eval-callback-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-eval-callback" +path = "src/main.rs" diff --git a/examples/eval-callback-rust/build.rs b/examples/eval-callback-rust/build.rs new file mode 100644 index 000000000000..27d0a601bd30 --- /dev/null +++ b/examples/eval-callback-rust/build.rs @@ -0,0 +1,43 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building the binary"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/eval-callback-rust/src/lib.rs b/examples/eval-callback-rust/src/lib.rs new file mode 100644 index 000000000000..89045c93dc9b --- /dev/null +++ b/examples/eval-callback-rust/src/lib.rs @@ -0,0 +1,109 @@ +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub n_gpu_layers: i32, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + MissingModel, + MissingValue(&'static str), + InvalidInteger(&'static str, String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut args = args.into_iter().map(Into::into); + let mut model_path = String::new(); + let mut prompt = "hello".to_string(); + let mut n_gpu_layers = 99; + + while let Some(arg) = args.next() { + match arg.as_str() { + "-m" | "--model" => { + model_path = args.next().ok_or(ParseError::MissingValue("-m"))?; + } + "-p" | "--prompt" => { + prompt = args.next().ok_or(ParseError::MissingValue("--prompt"))?; + } + "-ngl" | "--n-gpu-layers" => { + let value = args.next().ok_or(ParseError::MissingValue("-ngl"))?; + n_gpu_layers = value + .parse() + .map_err(|_| ParseError::InvalidInteger("-ngl", value))?; + } + "--seed" | "-s" => { + args.next().ok_or(ParseError::MissingValue("--seed"))?; + } + _ => {} + } + } + + if model_path.is_empty() { + return Err(ParseError::MissingModel); + } + + Ok(Args { + model_path, + prompt, + n_gpu_layers, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_ctest_invocation() { + let args = parse_args([ + "-m", + "model.gguf", + "--prompt", + "hello", + "--seed", + "42", + "-ngl", + "0", + ]) + .unwrap(); + + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "hello"); + assert_eq!(args.n_gpu_layers, 0); + } + + #[test] + fn rejects_missing_model() { + assert_eq!( + parse_args(["--prompt", "hello"]).unwrap_err(), + ParseError::MissingModel + ); + } + + #[test] + fn rejects_invalid_gpu_layer_count() { + assert_eq!( + parse_args(["-m", "model.gguf", "-ngl", "gpu"]).unwrap_err(), + ParseError::InvalidInteger("-ngl", "gpu".to_string()) + ); + } +} diff --git a/examples/eval-callback-rust/src/main.rs b/examples/eval-callback-rust/src/main.rs new file mode 100644 index 000000000000..033cba3b9060 --- /dev/null +++ b/examples/eval-callback-rust/src/main.rs @@ -0,0 +1,177 @@ +use llama_eval_callback_rust::{parse_args, Args}; +use llama_simple_rust::ffi; +use std::env; +use std::ffi::{c_void, CString}; +use std::ptr; + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +unsafe extern "C" fn eval_callback( + _tensor: *mut ffi::ggml_tensor, + _ask: bool, + _user_data: *mut c_void, +) -> bool { + true +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m model.gguf --prompt hello -ngl 0"); + eprintln!(); +} + +fn tokenize( + vocab: *const ffi::llama_vocab, + prompt: &str, + add_special: bool, +) -> Result, String> { + let prompt_c = + CString::new(prompt).map_err(|_| "prompt contains an interior NUL byte".to_string())?; + let n_prompt = unsafe { + -ffi::llama_tokenize( + vocab, + prompt_c.as_ptr(), + prompt.len() as i32, + ptr::null_mut(), + 0, + add_special, + true, + ) + }; + + if n_prompt <= 0 { + return Err("failed to size prompt tokenization".to_string()); + } + + let mut tokens = vec![0_i32; n_prompt as usize]; + let n_tokenized = unsafe { + ffi::llama_tokenize( + vocab, + prompt_c.as_ptr(), + prompt.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + add_special, + true, + ) + }; + + if n_tokenized < 0 { + return Err("failed to tokenize the prompt".to_string()); + } + + tokens.truncate(n_tokenized as usize); + Ok(tokens) +} + +fn run(args: Args) -> Result<(), String> { + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + + unsafe { + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + + let add_bos = unsafe { ffi::llama_vocab_get_add_bos(vocab) }; + let mut tokens = tokenize(vocab, &args.prompt, add_bos)?; + if tokens.is_empty() { + return Err("there are no input tokens to process".to_string()); + } + + eprintln!("number of input tokens = {}", tokens.len()); + for token in &tokens { + eprintln!(" {token}"); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = tokens.len().max(1) as u32; + ctx_params.n_batch = tokens.len() as u32; + ctx_params.cb_eval = Some(eval_callback); + ctx_params.cb_eval_user_data = ptr::null_mut(); + ctx_params.no_perf = false; + + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create the llama_context".to_string()); + } + + let batch = unsafe { ffi::llama_batch_get_one(tokens.as_mut_ptr(), tokens.len() as i32) }; + if unsafe { ffi::llama_decode(ctx.0, batch) } != 0 { + return Err("failed to eval".to_string()); + } + + eprintln!(); + unsafe { + ffi::llama_perf_context_print(ctx.0); + } + + Ok(()) +} + +fn main() { + let mut argv = env::args(); + let program = argv + .next() + .unwrap_or_else(|| "llama-eval-callback".to_string()); + + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + + let status = match run(args) { + Ok(()) => 0, + Err(err) => { + eprintln!("{program}: error: {err}"); + 1 + } + }; + + unsafe { + ffi::llama_backend_free(); + } + std::process::exit(status); +} diff --git a/examples/eval-callback/CMakeLists.txt b/examples/eval-callback/CMakeLists.txt index 63fbe59dce87..b6945858944c 100644 --- a/examples/eval-callback/CMakeLists.txt +++ b/examples/eval-callback/CMakeLists.txt @@ -1,8 +1,29 @@ set(TARGET llama-eval-callback) -add_executable(${TARGET} eval-callback.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(EVAL_CALLBACK_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../eval-callback-rust/Cargo.toml) +set(EVAL_CALLBACK_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/eval-callback) +set(EVAL_CALLBACK_RUST_BIN ${EVAL_CALLBACK_RUST_TARGET_DIR}/debug/llama-eval-callback${CMAKE_EXECUTABLE_SUFFIX}) +set(EVAL_CALLBACK_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${EVAL_CALLBACK_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${EVAL_CALLBACK_RUST_MANIFEST} --target-dir ${EVAL_CALLBACK_RUST_TARGET_DIR} --bin llama-eval-callback + COMMAND ${CMAKE_COMMAND} -E copy ${EVAL_CALLBACK_RUST_BIN} ${EVAL_CALLBACK_RUST_OUTPUT} + DEPENDS + llama + ${EVAL_CALLBACK_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../eval-callback-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../eval-callback-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../eval-callback-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${EVAL_CALLBACK_RUST_OUTPUT}) +install(PROGRAMS ${EVAL_CALLBACK_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) if(LLAMA_BUILD_TESTS) if(NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "s390x") diff --git a/examples/eval-callback/eval-callback.cpp b/examples/eval-callback/eval-callback.cpp deleted file mode 100644 index 4ce8d600b1a7..000000000000 --- a/examples/eval-callback/eval-callback.cpp +++ /dev/null @@ -1,88 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "debug.h" -#include "log.h" -#include "llama.h" - -#include -#include -#include - -static bool run(llama_context * ctx, const common_params & params) { - const llama_model * model = llama_get_model(ctx); - const llama_vocab * vocab = llama_model_get_vocab(model); - - const bool add_bos = llama_vocab_get_add_bos(vocab); - - std::vector tokens = common_tokenize(ctx, params.prompt, add_bos, true); - - if (tokens.empty()) { - LOG_ERR("%s : there are not input tokens to process - (try to provide a prompt with '-p')\n", __func__); - return false; - } - - LOG_INF("number of input tokens = %zu\n", tokens.size()); - for (size_t i = 0; i < tokens.size(); ++i) { - LOG_INF(" %d\n", tokens[i]); - } - - if (llama_decode(ctx, llama_batch_get_one(tokens.data(), tokens.size()))) { - LOG_ERR("%s : failed to eval\n", __func__); - return false; - } - - return true; -} - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_debug_cb_user_data cb_data; - - common_params params; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { - return 1; - } - - llama_backend_init(); - llama_numa_init(params.numa); - - // pass the callback to the backend scheduler - // it will be executed for each node during the graph computation - params.cb_eval = common_debug_cb_eval; - params.cb_eval_user_data = &cb_data; - params.warmup = false; - - // init - auto llama_init = common_init_from_params(params); - - auto * model = llama_init->model(); - auto * ctx = llama_init->context(); - - if (model == nullptr || ctx == nullptr) { - LOG_ERR("%s : failed to init\n", __func__); - return 1; - } - - // print system information - { - LOG_INF("\n"); - LOG_INF("%s\n", common_params_get_system_info(params).c_str()); - LOG_INF("\n"); - } - - bool OK = run(ctx, params); - if (!OK) { - return 1; - } - - LOG("\n"); - llama_perf_context_print(ctx); - - llama_backend_free(); - - return 0; -} diff --git a/examples/gen-docs-rust/Cargo.toml b/examples/gen-docs-rust/Cargo.toml new file mode 100644 index 000000000000..fe4e5eb83023 --- /dev/null +++ b/examples/gen-docs-rust/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "llama-gen-docs-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[[bin]] +name = "llama-gen-docs" +path = "src/main.rs" diff --git a/examples/gen-docs-rust/src/main.rs b/examples/gen-docs-rust/src/main.rs new file mode 100644 index 000000000000..0d6a0dbd3a1f --- /dev/null +++ b/examples/gen-docs-rust/src/main.rs @@ -0,0 +1,176 @@ +use std::fs; + +const HELP_START_MARKER: &str = ""; +const HELP_END_MARKER: &str = ""; +const NOTE_MESSAGE: &str = + ""; + +const MD_FILES: &[(&str, &str)] = &[ + ("tools/cli/README.md", "CLI-specific params"), + ("tools/completion/README.md", "Completion-specific params"), + ("tools/server/README.md", "Server-specific params"), +]; + +#[derive(Clone, Copy)] +struct ArgDoc { + args: &'static str, + help: &'static str, +} + +const COMMON_OPTIONS: &[ArgDoc] = &[ + ArgDoc { + args: "-h, --help", + help: "show this help message and exit", + }, + ArgDoc { + args: "-m, --model FNAME", + help: "model path", + }, + ArgDoc { + args: "-t, --threads N", + help: "number of threads to use during generation", + }, + ArgDoc { + args: "-c, --ctx-size N", + help: "size of the prompt context", + }, +]; + +const SAMPLING_OPTIONS: &[ArgDoc] = &[ + ArgDoc { + args: "-n, --predict N", + help: "number of tokens to predict", + }, + ArgDoc { + args: "--temp N", + help: "sampling temperature", + }, + ArgDoc { + args: "--top-k N", + help: "top-k sampling", + }, + ArgDoc { + args: "--top-p N", + help: "top-p sampling", + }, +]; + +const SPECIFIC_OPTIONS: &[ArgDoc] = &[ + ArgDoc { + args: "-p, --prompt PROMPT", + help: "prompt to start generation with", + }, + ArgDoc { + args: "--host HOST", + help: "server listen address where applicable", + }, + ArgDoc { + args: "--port PORT", + help: "server listen port where applicable", + }, +]; + +fn escape_md(value: &str) -> String { + value.replace('|', "\\|").replace('\n', "
") +} + +fn write_table(out: &mut String, rows: &[ArgDoc]) { + out.push_str("| Argument | Explanation |\n"); + out.push_str("| -------- | ----------- |\n"); + for row in rows { + out.push_str("| `"); + out.push_str(&escape_md(row.args)); + out.push_str("` | "); + out.push_str(&escape_md(row.help.trim())); + out.push_str(" |\n"); + } +} + +fn build_help(specific_section_header: &str) -> String { + let mut out = String::new(); + out.push_str(HELP_START_MARKER); + out.push_str("\n\n"); + out.push_str(NOTE_MESSAGE); + out.push_str("\n\n### Common params\n\n"); + write_table(&mut out, COMMON_OPTIONS); + out.push_str("\n\n### Sampling params\n\n"); + write_table(&mut out, SAMPLING_OPTIONS); + out.push_str("\n\n### "); + out.push_str(specific_section_header); + out.push_str("\n\n"); + write_table(&mut out, SPECIFIC_OPTIONS); + out.push('\n'); + out.push_str(HELP_END_MARKER); + out +} + +fn replace_help(content: &str, new_help: &str) -> Result { + let help_start = content + .find(HELP_START_MARKER) + .ok_or_else(|| "missing help start marker".to_string())?; + let help_end = content + .find(HELP_END_MARKER) + .ok_or_else(|| "missing help end marker".to_string())?; + if help_end <= help_start { + return Err("help end marker appears before start marker".to_string()); + } + + Ok(format!( + "{}{}{}", + &content[..help_start], + new_help, + &content[help_end + HELP_END_MARKER.len()..] + )) +} + +fn update_file(path: &str, specific_section_header: &str) -> Result<(), String> { + let content = + fs::read_to_string(path).map_err(|err| format!("failed to read {path}: {err}"))?; + let new_help = build_help(specific_section_header); + let updated = replace_help(&content, &new_help) + .map_err(|err| format!("failed to update help in {path}: {err}"))?; + fs::write(path, updated).map_err(|err| format!("failed to write {path}: {err}"))?; + println!("Updated help in '{path}'"); + Ok(()) +} + +fn main() { + for (path, section) in MD_FILES { + if let Err(err) = update_file(path, section) { + eprintln!("{err}"); + std::process::exit(1); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn markdown_escapes_pipes_and_newlines() { + assert_eq!(escape_md("a|b\nc"), "a\\|b
c"); + } + + #[test] + fn generated_help_has_markers_and_sections() { + let help = build_help("Server-specific params"); + assert!(help.starts_with(HELP_START_MARKER)); + assert!(help.ends_with(HELP_END_MARKER)); + assert!(help.contains("### Common params")); + assert!(help.contains("### Sampling params")); + assert!(help.contains("### Server-specific params")); + } + + #[test] + fn replaces_only_marked_region() { + let content = format!("before\n{HELP_START_MARKER}\nold\n{HELP_END_MARKER}\nafter"); + let updated = replace_help(&content, "NEW").unwrap(); + assert_eq!(updated, "before\nNEW\nafter"); + } + + #[test] + fn rejects_missing_markers() { + assert!(replace_help("no markers", "NEW").is_err()); + } +} diff --git a/examples/gen-docs/CMakeLists.txt b/examples/gen-docs/CMakeLists.txt index aa68cbd78a86..2bfb0a33fdcf 100644 --- a/examples/gen-docs/CMakeLists.txt +++ b/examples/gen-docs/CMakeLists.txt @@ -1,5 +1,19 @@ set(TARGET llama-gen-docs) -add_executable(${TARGET} gen-docs.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(GEN_DOCS_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../gen-docs-rust/Cargo.toml) +set(GEN_DOCS_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/gen-docs) +set(GEN_DOCS_RUST_BIN ${GEN_DOCS_RUST_TARGET_DIR}/debug/llama-gen-docs${CMAKE_EXECUTABLE_SUFFIX}) +set(GEN_DOCS_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${GEN_DOCS_RUST_OUTPUT} + COMMAND ${CARGO_EXECUTABLE} build --manifest-path ${GEN_DOCS_RUST_MANIFEST} --target-dir ${GEN_DOCS_RUST_TARGET_DIR} --bin llama-gen-docs + COMMAND ${CMAKE_COMMAND} -E copy ${GEN_DOCS_RUST_BIN} ${GEN_DOCS_RUST_OUTPUT} + DEPENDS + ${GEN_DOCS_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../gen-docs-rust/src/main.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${GEN_DOCS_RUST_OUTPUT}) +install(PROGRAMS ${GEN_DOCS_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/gen-docs/gen-docs.cpp b/examples/gen-docs/gen-docs.cpp deleted file mode 100644 index baf61bf27b54..000000000000 --- a/examples/gen-docs/gen-docs.cpp +++ /dev/null @@ -1,145 +0,0 @@ -#include "arg.h" -#include "common.h" - -#include -#include -#include -#include - -// Export usage message (-h) to markdown format -// Automatically update the markdown docs - -#define HELP_START_MARKER "" -#define HELP_END_MARKER "" -#define NOTE_MESSAGE "" - -struct md_file { - llama_example ex; - std::string fname; - std::string specific_section_header; -}; - -std::vector md_files = { - {LLAMA_EXAMPLE_CLI, "tools/cli/README.md", "CLI-specific params"}, - {LLAMA_EXAMPLE_COMPLETION, "tools/completion/README.md", "Completion-specific params"}, - {LLAMA_EXAMPLE_SERVER, "tools/server/README.md", "Server-specific params"}, -}; - -static void write_table_header(std::ostringstream & ss) { - ss << "| Argument | Explanation |\n"; - ss << "| -------- | ----------- |\n"; -} - -static void write_table_entry(std::ostringstream & ss, const common_arg & opt) { - ss << "| `"; - // args - auto all_args = opt.get_args(); - for (const auto & arg : all_args) { - if (arg == all_args.front()) { - ss << arg; - if (all_args.size() > 1) ss << ", "; - } else { - ss << arg << (arg != all_args.back() ? ", " : ""); - } - } - // value hint - if (opt.value_hint) { - std::string md_value_hint(opt.value_hint); - string_replace_all(md_value_hint, "|", "\\|"); - ss << " " << md_value_hint; - } - if (opt.value_hint_2) { - std::string md_value_hint_2(opt.value_hint_2); - string_replace_all(md_value_hint_2, "|", "\\|"); - ss << " " << md_value_hint_2; - } - // help text - std::string md_help(opt.help); - md_help = string_strip(md_help); - string_replace_all(md_help, "\n", "
"); - string_replace_all(md_help, "|", "\\|"); - ss << "` | " << md_help << " |\n"; -} - -static void write_table(std::ostringstream & ss, std::vector & opts) { - write_table_header(ss); - for (const auto & opt : opts) { - write_table_entry(ss, *opt); - } -} - -static void write_help(std::ostringstream & ss, const md_file & md) { - common_params params; - auto ctx_arg = common_params_parser_init(params, md.ex); - - std::vector common_options; - std::vector sampling_options; - std::vector specific_options; - for (auto & opt : ctx_arg.options) { - // in case multiple LLAMA_EXAMPLE_* are set, we prioritize the LLAMA_EXAMPLE_* matching current example - if (opt.is_sampling) { - sampling_options.push_back(&opt); - } else if (opt.in_example(ctx_arg.ex)) { - specific_options.push_back(&opt); - } else { - common_options.push_back(&opt); - } - } - - ss << HELP_START_MARKER << "\n\n"; - - ss << NOTE_MESSAGE << "\n\n"; - - ss << "### Common params\n\n"; - write_table(ss, common_options); - ss << "\n\n### Sampling params\n\n"; - write_table(ss, sampling_options); - ss << "\n\n### " << md.specific_section_header << "\n\n"; - write_table(ss, specific_options); - - ss << "\n" << HELP_END_MARKER; -} - -int main(int, char **) { - std::setlocale(LC_NUMERIC, "C"); - - for (const auto & md : md_files) { - std::ifstream infile(md.fname); - if (!infile.is_open()) { - fprintf(stderr, "failed to open file '%s' for reading\n", md.fname.c_str()); - return 1; - } - - std::ostringstream ss; - ss << infile.rdbuf(); - infile.close(); - - std::string content = ss.str(); - - size_t help_start = content.find(HELP_START_MARKER); - size_t help_end = content.find(HELP_END_MARKER); - - if (help_start == std::string::npos || help_end == std::string::npos || help_end <= help_start) { - fprintf(stderr, "failed to find help markers in file '%s'\n", md.fname.c_str()); - return 1; - } - - std::ostringstream new_help_ss; - write_help(new_help_ss, md); - std::string new_help = new_help_ss.str(); - - content = content.substr(0, help_start) + new_help + content.substr(help_end + strlen(HELP_END_MARKER)); - - std::ofstream outfile(md.fname); - if (!outfile.is_open()) { - fprintf(stderr, "failed to open file '%s' for writing\n", md.fname.c_str()); - return 1; - } - outfile << content; - outfile.close(); - - printf("Updated help in '%s'\n", md.fname.c_str()); - } - - return 0; -} diff --git a/examples/gguf-hash-rust/Cargo.toml b/examples/gguf-hash-rust/Cargo.toml new file mode 100644 index 000000000000..023e9eab199c --- /dev/null +++ b/examples/gguf-hash-rust/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "llama-gguf-hash-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[[bin]] +name = "llama-gguf-hash" +path = "src/main.rs" + +[dependencies] +sha1 = "0.10" +sha2 = "0.10" +xxhash-rust = { version = "0.8", features = ["xxh64"] } diff --git a/examples/gguf-hash-rust/build.rs b/examples/gguf-hash-rust/build.rs new file mode 100644 index 000000000000..9b908260a727 --- /dev/null +++ b/examples/gguf-hash-rust/build.rs @@ -0,0 +1,49 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + if let Ok(lib_dir) = env::var("GGML_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/ggml/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libggml.so").exists() + || candidate.join("libggml.dylib").exists() + || candidate.join("ggml.dll").exists() + || candidate.join("ggml.lib").exists() + || candidate.join("libggml.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!( + "cargo:warning=libggml was not found; build llama.cpp before building llama-gguf-hash" + ); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os == "linux" || target_os == "android" { + println!("cargo:rustc-link-lib=stdc++"); + } else if target_os == "macos" { + println!("cargo:rustc-link-lib=c++"); + } + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/gguf-hash-rust/src/lib.rs b/examples/gguf-hash-rust/src/lib.rs new file mode 100644 index 000000000000..4cbe4171918d --- /dev/null +++ b/examples/gguf-hash-rust/src/lib.rs @@ -0,0 +1,300 @@ +use std::collections::HashSet; +use std::fmt; +use std::fs::File; +use std::io::{self, BufRead, BufReader}; + +pub const HASH_TYPE_SHA256: &str = "sha256"; +pub const HASH_TYPE_SHA1: &str = "sha1"; +pub const HASH_TYPE_XXH64: &str = "xxh64"; +pub const HASH_TYPE_UUID: &str = "uuid"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(i32)] +pub enum HashExitCode { + Success = 0, + Failure = 1, + Mismatch = 2, + ManifestMissingEntry = 3, + ManifestUnknownHash = 4, + ManifestFileError = 5, +} + +impl fmt::Display for HashExitCode { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + HashExitCode::Success => write!(f, "Success"), + HashExitCode::Failure => write!(f, "Failure"), + HashExitCode::Mismatch => write!(f, "Mismatch"), + HashExitCode::ManifestMissingEntry => write!(f, "Manifest Missing Entry"), + HashExitCode::ManifestUnknownHash => write!(f, "Manifest Unknown Hash"), + HashExitCode::ManifestFileError => write!(f, "Manifest File Error"), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ManifestResult { + NotFound, + Mismatch, + Ok, +} + +impl fmt::Display for ManifestResult { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ManifestResult::NotFound => write!(f, "Not Found"), + ManifestResult::Mismatch => write!(f, "Mismatch"), + ManifestResult::Ok => write!(f, "Ok"), + } + } +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct HashParams { + pub input: String, + pub xxh64: bool, + pub sha1: bool, + pub sha256: bool, + pub uuid: bool, + pub no_layer: bool, + pub manifest_is_usable: bool, + pub manifest_file: String, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct ManifestCheck { + pub xxh64: bool, + pub sha1: bool, + pub sha256: bool, + pub uuid: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + Help, + MissingValue(String), + UnknownArgument(String), + BadArguments, +} + +impl fmt::Display for ParseError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ParseError::Help => write!(f, "help requested"), + ParseError::MissingValue(arg) => { + write!(f, "error: invalid parameter for argument:{arg}") + } + ParseError::UnknownArgument(arg) => write!(f, "error: unknown argument: {arg}"), + ParseError::BadArguments => write!(f, "error: bad arguments"), + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut params = HashParams::default(); + let mut args = args.into_iter().map(Into::into).peekable(); + + while let Some(arg) = args.peek().cloned() { + if !arg.starts_with('-') { + break; + } + let arg = args.next().expect("peeked"); + let normalized = if arg.starts_with("--") { + arg.replace('_', "-") + } else { + arg + }; + + match normalized.as_str() { + "-h" | "--help" => return Err(ParseError::Help), + "--xxh64" => params.xxh64 = true, + "--sha1" => params.sha1 = true, + "--uuid" => params.uuid = true, + "--sha256" => params.sha256 = true, + "--all" => { + params.sha256 = true; + params.sha1 = true; + params.xxh64 = true; + } + "--no-layer" => params.no_layer = true, + "-c" | "--check" => { + params.manifest_file = args + .next() + .ok_or_else(|| ParseError::MissingValue(normalized.clone()))?; + } + _ => return Err(ParseError::UnknownArgument(normalized)), + } + } + + params.input = args.next().ok_or(ParseError::BadArguments)?; + Ok(params) +} + +pub fn manifest_type(path: &str) -> io::Result { + let file = File::open(path)?; + let mut check = ManifestCheck::default(); + for line in BufReader::new(file).lines() { + let line = line?; + match line.split_whitespace().next() { + Some(HASH_TYPE_SHA256) => check.sha256 = true, + Some(HASH_TYPE_SHA1) => check.sha1 = true, + Some(HASH_TYPE_XXH64) => check.xxh64 = true, + Some(HASH_TYPE_UUID) => check.uuid = true, + _ => {} + } + } + Ok(check) +} + +pub fn manifest_verify( + path: &str, + hash_type: &str, + hash: &str, + tensor_name: &str, +) -> io::Result { + if path.is_empty() { + return Ok(ManifestResult::NotFound); + } + + let file = File::open(path)?; + for line in BufReader::new(file).lines() { + let line = line?; + let mut parts = line.split_whitespace(); + let Some(file_hash_type) = parts.next() else { + continue; + }; + let Some(file_hash) = parts.next() else { + continue; + }; + let Some(file_tensor_name) = parts.next() else { + continue; + }; + + if file_hash_type == hash_type && file_tensor_name == tensor_name { + return Ok(if file_hash == hash { + ManifestResult::Ok + } else { + ManifestResult::Mismatch + }); + } + } + + Ok(ManifestResult::NotFound) +} + +pub fn generate_uuidv5(sha1_digest: &[u8; 20]) -> [u8; 16] { + let mut uuid = [0_u8; 16]; + uuid.copy_from_slice(&sha1_digest[..16]); + uuid[6] &= !(0xF << 4); + uuid[6] |= 5 << 4; + uuid[8] &= !(0xC << 4); + uuid[8] |= 0x8 << 4; + uuid +} + +pub fn hex_lower(bytes: &[u8]) -> String { + let mut out = String::with_capacity(bytes.len() * 2); + for byte in bytes { + use std::fmt::Write; + write!(&mut out, "{byte:02x}").expect("write to String"); + } + out +} + +pub fn hex_u64_be(value: u64) -> String { + format!("{value:016x}") +} + +pub fn format_uuid(uuid: &[u8; 16]) -> String { + format!( + "{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}", + uuid[0], + uuid[1], + uuid[2], + uuid[3], + uuid[4], + uuid[5], + uuid[6], + uuid[7], + uuid[8], + uuid[9], + uuid[10], + uuid[11], + uuid[12], + uuid[13], + uuid[14], + uuid[15], + ) +} + +pub fn known_hash_types(check: &ManifestCheck) -> HashSet<&'static str> { + let mut out = HashSet::new(); + if check.sha256 { + out.insert(HASH_TYPE_SHA256); + } + if check.sha1 { + out.insert(HASH_TYPE_SHA1); + } + if check.xxh64 { + out.insert(HASH_TYPE_XXH64); + } + if check.uuid { + out.insert(HASH_TYPE_UUID); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_default_input() { + let params = parse_args(["model.gguf"]).unwrap(); + assert_eq!(params.input, "model.gguf"); + assert!(!params.xxh64); + } + + #[test] + fn parses_hash_options_and_manifest() { + let params = parse_args(["--all", "--no-layer", "--check", "manifest", "model"]).unwrap(); + assert_eq!(params.input, "model"); + assert_eq!(params.manifest_file, "manifest"); + assert!(params.xxh64); + assert!(params.sha1); + assert!(params.sha256); + assert!(params.no_layer); + } + + #[test] + fn supports_short_check_and_help() { + let params = parse_args(["-c", "manifest", "--uuid", "model"]).unwrap(); + assert_eq!(params.manifest_file, "manifest"); + assert!(params.uuid); + assert_eq!(parse_args(["--help"]).unwrap_err(), ParseError::Help); + } + + #[test] + fn uuidv5_sets_version_and_variant() { + let digest = [0xff_u8; 20]; + let uuid = generate_uuidv5(&digest); + assert_eq!(uuid[6] >> 4, 5); + assert_eq!(uuid[8] >> 6, 2); + } + + #[test] + fn formats_hex_and_uuid() { + assert_eq!(hex_lower(&[0xab, 0x00, 0x12]), "ab0012"); + assert_eq!(hex_u64_be(0x1234), "0000000000001234"); + assert_eq!( + format_uuid(&[0, 1, 2, 3, 4, 5, 0x56, 7, 0x89, 9, 10, 11, 12, 13, 14, 15]), + "00010203-0405-5607-8909-0a0b0c0d0e0f" + ); + } +} diff --git a/examples/gguf-hash-rust/src/main.rs b/examples/gguf-hash-rust/src/main.rs new file mode 100644 index 000000000000..1f8649008fe9 --- /dev/null +++ b/examples/gguf-hash-rust/src/main.rs @@ -0,0 +1,405 @@ +use std::error::Error; +use std::ffi::{c_char, CStr, CString}; +use std::fs::File; +use std::io::{Read, Seek, SeekFrom}; +use std::process::ExitCode; + +use sha1::{Digest as Sha1Digest, Sha1}; +use sha2::Sha256; +use xxhash_rust::xxh64::{xxh64, Xxh64}; + +use llama_gguf_hash_rust::{ + format_uuid, generate_uuidv5, hex_lower, hex_u64_be, manifest_type, manifest_verify, + parse_args, HashExitCode, HashParams, ManifestResult, ParseError, HASH_TYPE_SHA1, + HASH_TYPE_SHA256, HASH_TYPE_UUID, HASH_TYPE_XXH64, +}; + +const UUID_NAMESPACE_LLAMA_CPP: [u8; 16] = [ + 0xef, 0x00, 0x12, 0x06, 0xda, 0xdc, 0x5f, 0x6d, 0xa1, 0x5f, 0x33, 0x59, 0xe5, 0x77, 0xd4, 0xe5, +]; + +#[repr(C)] +struct gguf_context { + _private: [u8; 0], +} + +#[repr(C)] +struct ggml_context { + _private: [u8; 0], +} + +#[repr(C)] +struct gguf_init_params { + no_alloc: bool, + ctx: *mut *mut ggml_context, +} + +extern "C" { + fn gguf_init_from_file(fname: *const c_char, params: gguf_init_params) -> *mut gguf_context; + fn gguf_free(ctx: *mut gguf_context); + fn gguf_get_data_offset(ctx: *const gguf_context) -> usize; + fn gguf_get_n_tensors(ctx: *const gguf_context) -> i64; + fn gguf_get_tensor_offset(ctx: *const gguf_context, tensor_id: i64) -> usize; + fn gguf_get_tensor_name(ctx: *const gguf_context, tensor_id: i64) -> *const c_char; + fn gguf_get_tensor_size(ctx: *const gguf_context, tensor_id: i64) -> usize; +} + +fn print_usage(executable: &str) { + println!(); + println!("usage: {executable} [options] GGUF_IN"); + println!(); + println!("Hash a GGUF file"); + println!(); + println!("options:"); + println!(" -h, --help show this help message and exit"); + println!(" --xxh64 use xxh64 hash"); + println!(" --sha1 use sha1 hash"); + println!(" --sha256 use sha256 hash"); + println!(" --all use all hash"); + println!(" --no-layer exclude per layer hash"); + println!(" --uuid generate UUIDv5 ID"); + println!(" -c, --check verify against a manifest"); + println!(); +} + +fn main() -> ExitCode { + match run() { + Ok(code) => ExitCode::from(code as u8), + Err(AppError::Help(argv0)) => { + print_usage(&argv0); + ExitCode::SUCCESS + } + Err(AppError::Usage(argv0, err)) => { + eprintln!("{err}"); + print_usage(&argv0); + ExitCode::FAILURE + } + Err(AppError::Exit(code, message)) => { + if let Some(message) = message { + print!("{message}"); + } + ExitCode::from(code as u8) + } + Err(AppError::Other(err)) => { + eprintln!("llama-gguf-hash: {err}"); + ExitCode::FAILURE + } + } +} + +fn run() -> Result { + let mut raw_args = std::env::args(); + let argv0 = raw_args + .next() + .unwrap_or_else(|| "llama-gguf-hash".to_string()); + let mut params = parse_args(raw_args).map_err(|err| match err { + ParseError::Help => AppError::Help(argv0.clone()), + err => AppError::Usage(argv0.clone(), err), + })?; + + if !params.manifest_file.is_empty() { + let manifest_check = manifest_type(¶ms.manifest_file).map_err(|_| { + AppError::Exit( + HashExitCode::ManifestFileError, + Some(format!( + "ERROR cannot open manifest {}", + params.manifest_file + )), + ) + })?; + + if !manifest_check.sha256 + && !manifest_check.sha1 + && !manifest_check.xxh64 + && !manifest_check.uuid + { + return Err(AppError::Exit( + HashExitCode::ManifestUnknownHash, + Some(format!( + "ERROR manifest does not have any known hash format in {}", + params.manifest_file + )), + )); + } + + print!("manifest {}", params.manifest_file); + if manifest_check.sha256 { + print!(" sha256"); + } + if manifest_check.sha1 { + print!(" sha1"); + } + if manifest_check.xxh64 { + print!(" xxh64"); + } + if manifest_check.uuid { + print!(" uuid"); + } + println!(); + + if !params.xxh64 && !params.sha1 && !params.uuid && !params.sha256 { + if manifest_check.sha256 { + params.sha256 = true; + } else if manifest_check.sha1 { + params.sha1 = true; + } else if manifest_check.xxh64 { + params.xxh64 = true; + } else if manifest_check.uuid { + params.uuid = true; + } + } + + params.manifest_is_usable = true; + } + + if !params.xxh64 && !params.sha1 && !params.uuid && !params.sha256 { + params.xxh64 = true; + } + + let exit_code = gguf_hash(¶ms)?; + if params.manifest_is_usable { + println!( + "\nVerification results for {} - {}", + params.manifest_file, exit_code + ); + } + Ok(exit_code) +} + +fn gguf_hash(params: &HashParams) -> Result { + let fname = params.input.clone(); + let c_fname = CString::new(fname.clone())?; + let ctx = unsafe { + gguf_init_from_file( + c_fname.as_ptr(), + gguf_init_params { + no_alloc: true, + ctx: std::ptr::null_mut(), + }, + ) + }; + if ctx.is_null() { + return Err(std::io::Error::other(format!("failed to open GGUF file {fname}")).into()); + } + + let result = unsafe { hash_loaded_gguf(params, ctx) }; + unsafe { + gguf_free(ctx); + } + result +} + +unsafe fn hash_loaded_gguf( + params: &HashParams, + ctx: *mut gguf_context, +) -> Result { + let mut file = File::open(¶ms.input)?; + let data_offset = gguf_get_data_offset(ctx) as u64; + let n_tensors = gguf_get_n_tensors(ctx); + + let mut xxh64_model = params.xxh64.then(|| Xxh64::new(0)); + let mut sha1_model = params.sha1.then(Sha1::new); + let mut sha256_model = params.sha256.then(Sha256::new); + let mut sha1_uuid = params.uuid.then(|| { + let mut hasher = Sha1::new(); + hasher.update(UUID_NAMESPACE_LLAMA_CPP); + hasher + }); + + let mut tensor_layer_in_manifest = false; + let mut model_in_manifest = false; + let mut tensor_layer_has_mismatch = false; + let mut model_has_mismatch = false; + + for i in 0..n_tensors { + let name = CStr::from_ptr(gguf_get_tensor_name(ctx, i)) + .to_string_lossy() + .into_owned(); + let size = gguf_get_tensor_size(ctx, i); + let offset = data_offset + gguf_get_tensor_offset(ctx, i) as u64; + let mut data = vec![0_u8; size]; + file.seek(SeekFrom::Start(offset))?; + file.read_exact(&mut data)?; + + let tensor_layer_name = format!("{}:{name}", params.input); + + if params.xxh64 { + if !params.no_layer { + let hash = hex_u64_be(xxh64(&data, 0)); + update_manifest_state( + params, + HASH_TYPE_XXH64, + &hash, + &tensor_layer_name, + &mut tensor_layer_in_manifest, + &mut tensor_layer_has_mismatch, + )?; + } + if let Some(hasher) = xxh64_model.as_mut() { + hasher.update(&data); + } + } + + if params.sha1 { + if !params.no_layer { + let hash = hex_lower(&Sha1::digest(&data)); + update_manifest_state( + params, + HASH_TYPE_SHA1, + &hash, + &tensor_layer_name, + &mut tensor_layer_in_manifest, + &mut tensor_layer_has_mismatch, + )?; + } + if let Some(hasher) = sha1_model.as_mut() { + hasher.update(&data); + } + } + + if params.sha256 { + if !params.no_layer { + let hash = hex_lower(&Sha256::digest(&data)); + update_manifest_state( + params, + HASH_TYPE_SHA256, + &hash, + &tensor_layer_name, + &mut tensor_layer_in_manifest, + &mut tensor_layer_has_mismatch, + )?; + } + if let Some(hasher) = sha256_model.as_mut() { + hasher.update(&data); + } + } + + if let Some(hasher) = sha1_uuid.as_mut() { + hasher.update(&data); + } + } + + if let Some(hasher) = xxh64_model { + let hash = hex_u64_be(hasher.digest()); + update_manifest_state( + params, + HASH_TYPE_XXH64, + &hash, + ¶ms.input, + &mut model_in_manifest, + &mut model_has_mismatch, + )?; + } + + if let Some(hasher) = sha1_model { + let hash = hex_lower(&hasher.finalize()); + update_manifest_state( + params, + HASH_TYPE_SHA1, + &hash, + ¶ms.input, + &mut model_in_manifest, + &mut model_has_mismatch, + )?; + } + + if let Some(hasher) = sha256_model { + let hash = hex_lower(&hasher.finalize()); + update_manifest_state( + params, + HASH_TYPE_SHA256, + &hash, + ¶ms.input, + &mut model_in_manifest, + &mut model_has_mismatch, + )?; + } + + if let Some(hasher) = sha1_uuid { + let digest = hasher.finalize(); + let mut digest_array = [0_u8; 20]; + digest_array.copy_from_slice(&digest); + let uuid = format_uuid(&generate_uuidv5(&digest_array)); + let verify_type = if params.manifest_is_usable { + // Preserve the historical C++ behavior for verification lookup. + HASH_TYPE_SHA256 + } else { + HASH_TYPE_UUID + }; + update_manifest_state( + params, + verify_type, + &uuid, + ¶ms.input, + &mut model_in_manifest, + &mut model_has_mismatch, + )?; + } + + if params.manifest_is_usable { + if !model_in_manifest { + if !tensor_layer_in_manifest { + return Ok(HashExitCode::ManifestMissingEntry); + } + if tensor_layer_has_mismatch { + return Ok(HashExitCode::Failure); + } + return Ok(HashExitCode::Success); + } + + if tensor_layer_in_manifest && tensor_layer_has_mismatch { + return Ok(HashExitCode::Failure); + } + if model_has_mismatch { + return Ok(HashExitCode::Failure); + } + } + + Ok(HashExitCode::Success) +} + +fn update_manifest_state( + params: &HashParams, + hash_type: &str, + hash: &str, + name: &str, + in_manifest: &mut bool, + has_mismatch: &mut bool, +) -> Result<(), AppError> { + if params.manifest_is_usable { + let result = manifest_verify(¶ms.manifest_file, hash_type, hash, name)?; + match result { + ManifestResult::NotFound => {} + ManifestResult::Mismatch => { + *in_manifest = true; + *has_mismatch = true; + } + ManifestResult::Ok => *in_manifest = true, + } + println!("{hash_type:<8} {hash} {name} - {result}"); + } else { + let print_type = if hash_type == HASH_TYPE_SHA256 && hash.len() == 36 { + HASH_TYPE_UUID + } else { + hash_type + }; + println!("{print_type:<8} {hash} {name}"); + } + Ok(()) +} + +enum AppError { + Help(String), + Usage(String, ParseError), + Exit(HashExitCode, Option), + Other(Box), +} + +impl From for AppError +where + E: Error + 'static, +{ + fn from(value: E) -> Self { + AppError::Other(Box::new(value)) + } +} diff --git a/examples/gguf-hash/CMakeLists.txt b/examples/gguf-hash/CMakeLists.txt index 15c5c68c6f40..3056d34104ab 100644 --- a/examples/gguf-hash/CMakeLists.txt +++ b/examples/gguf-hash/CMakeLists.txt @@ -1,22 +1,23 @@ set(TARGET llama-gguf-hash) -add_executable(${TARGET} gguf-hash.cpp) -install(TARGETS ${TARGET} RUNTIME) +find_program(CARGO_EXECUTABLE cargo REQUIRED) -# clibs dependencies -include_directories(deps/) +set(GGUF_HASH_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../gguf-hash-rust/Cargo.toml) +set(GGUF_HASH_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/gguf-hash) +set(GGUF_HASH_RUST_BIN ${GGUF_HASH_RUST_TARGET_DIR}/debug/llama-gguf-hash${CMAKE_EXECUTABLE_SUFFIX}) +set(GGUF_HASH_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) -add_library(xxhash OBJECT deps/xxhash/xxhash.c deps/xxhash/xxhash.h) -target_link_libraries(${TARGET} PRIVATE xxhash) +add_custom_command( + OUTPUT ${GGUF_HASH_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env GGML_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${GGUF_HASH_RUST_MANIFEST} --target-dir ${GGUF_HASH_RUST_TARGET_DIR} --bin llama-gguf-hash + COMMAND ${CMAKE_COMMAND} -E copy ${GGUF_HASH_RUST_BIN} ${GGUF_HASH_RUST_OUTPUT} + DEPENDS + ggml + ${GGUF_HASH_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../gguf-hash-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../gguf-hash-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../gguf-hash-rust/src/main.rs +) -add_library(sha1 OBJECT deps/sha1/sha1.c deps/sha1/sha1.h) -target_link_libraries(${TARGET} PRIVATE sha1) -if (NOT MSVC) - # disable warnings in 3rd party code - target_compile_options(sha1 PRIVATE -w) -endif() - -add_library(sha256 OBJECT deps/sha256/sha256.c deps/sha256/sha256.h) -target_link_libraries(${TARGET} PRIVATE sha256) - -target_link_libraries(${TARGET} PRIVATE ggml ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +add_custom_target(${TARGET} ALL DEPENDS ${GGUF_HASH_RUST_OUTPUT}) +install(PROGRAMS ${GGUF_HASH_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/gguf-hash/deps/rotate-bits/rotate-bits.h b/examples/gguf-hash/deps/rotate-bits/rotate-bits.h deleted file mode 100644 index 75c4881fc322..000000000000 --- a/examples/gguf-hash/deps/rotate-bits/rotate-bits.h +++ /dev/null @@ -1,46 +0,0 @@ - - -#ifndef __ROTATE_DEFS_H -#define __ROTATE_DEFS_H - -#ifdef _MSC_VER - -#include - -#define ROTL32(v, n) _rotl((v), (n)) -#define ROTL64(v, n) _rotl64((v), (n)) - -#define ROTR32(v, n) _rotr((v), (n)) -#define ROTR64(v, n) _rotr64((v), (n)) - -#else - -#include - -#define U8V(v) ((uint8_t)(v) & 0xFFU) -#define U16V(v) ((uint16_t)(v) & 0xFFFFU) -#define U32V(v) ((uint32_t)(v) & 0xFFFFFFFFU) -#define U64V(v) ((uint64_t)(v) & 0xFFFFFFFFFFFFFFFFU) - -#define ROTL32(v, n) \ - (U32V((uint32_t)(v) << (n)) | ((uint32_t)(v) >> (32 - (n)))) - -// tests fail if we don't have this cast... -#define ROTL64(v, n) \ - (U64V((uint64_t)(v) << (n)) | ((uint64_t)(v) >> (64 - (n)))) - -#define ROTR32(v, n) ROTL32(v, 32 - (n)) -#define ROTR64(v, n) ROTL64(v, 64 - (n)) - -#endif - -#define ROTL8(v, n) \ - (U8V((uint8_t)(v) << (n)) | ((uint8_t)(v) >> (8 - (n)))) - -#define ROTL16(v, n) \ - (U16V((uint16_t)(v) << (n)) | ((uint16_t)(v) >> (16 - (n)))) - -#define ROTR8(v, n) ROTL8(v, 8 - (n)) -#define ROTR16(v, n) ROTL16(v, 16 - (n)) - -#endif diff --git a/examples/gguf-hash/deps/sha1/sha1.c b/examples/gguf-hash/deps/sha1/sha1.c deleted file mode 100644 index 76cd6ca3381d..000000000000 --- a/examples/gguf-hash/deps/sha1/sha1.c +++ /dev/null @@ -1,295 +0,0 @@ -/* -SHA-1 in C -By Steve Reid -100% Public Domain - -Test Vectors (from FIPS PUB 180-1) -"abc" - A9993E36 4706816A BA3E2571 7850C26C 9CD0D89D -"abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq" - 84983E44 1C3BD26E BAAE4AA1 F95129E5 E54670F1 -A million repetitions of "a" - 34AA973C D4C4DAA4 F61EEB2B DBAD2731 6534016F -*/ - -/* #define LITTLE_ENDIAN * This should be #define'd already, if true. */ -/* #define SHA1HANDSOFF * Copies data before messing with it. */ - -#define SHA1HANDSOFF - -#include -#include - -/* for uint32_t */ -#include - -#include "sha1.h" - - -#define rol(value, bits) (((value) << (bits)) | ((value) >> (32 - (bits)))) - -/* blk0() and blk() perform the initial expand. */ -/* I got the idea of expanding during the round function from SSLeay */ -#if BYTE_ORDER == LITTLE_ENDIAN -#define blk0(i) (block->l[i] = (rol(block->l[i],24)&0xFF00FF00) \ - |(rol(block->l[i],8)&0x00FF00FF)) -#elif BYTE_ORDER == BIG_ENDIAN -#define blk0(i) block->l[i] -#else -#error "Endianness not defined!" -#endif -#define blk(i) (block->l[i&15] = rol(block->l[(i+13)&15]^block->l[(i+8)&15] \ - ^block->l[(i+2)&15]^block->l[i&15],1)) - -/* (R0+R1), R2, R3, R4 are the different operations used in SHA1 */ -#define R0(v,w,x,y,z,i) z+=((w&(x^y))^y)+blk0(i)+0x5A827999+rol(v,5);w=rol(w,30); -#define R1(v,w,x,y,z,i) z+=((w&(x^y))^y)+blk(i)+0x5A827999+rol(v,5);w=rol(w,30); -#define R2(v,w,x,y,z,i) z+=(w^x^y)+blk(i)+0x6ED9EBA1+rol(v,5);w=rol(w,30); -#define R3(v,w,x,y,z,i) z+=(((w|x)&y)|(w&x))+blk(i)+0x8F1BBCDC+rol(v,5);w=rol(w,30); -#define R4(v,w,x,y,z,i) z+=(w^x^y)+blk(i)+0xCA62C1D6+rol(v,5);w=rol(w,30); - - -/* Hash a single 512-bit block. This is the core of the algorithm. */ - -void SHA1Transform( - uint32_t state[5], - const unsigned char buffer[64] -) -{ - uint32_t a, b, c, d, e; - - typedef union - { - unsigned char c[64]; - uint32_t l[16]; - } CHAR64LONG16; - -#ifdef SHA1HANDSOFF - CHAR64LONG16 block[1]; /* use array to appear as a pointer */ - - memcpy(block, buffer, 64); -#else - /* The following had better never be used because it causes the - * pointer-to-const buffer to be cast into a pointer to non-const. - * And the result is written through. I threw a "const" in, hoping - * this will cause a diagnostic. - */ - CHAR64LONG16 *block = (const CHAR64LONG16 *) buffer; -#endif - /* Copy context->state[] to working vars */ - a = state[0]; - b = state[1]; - c = state[2]; - d = state[3]; - e = state[4]; - /* 4 rounds of 20 operations each. Loop unrolled. */ - R0(a, b, c, d, e, 0); - R0(e, a, b, c, d, 1); - R0(d, e, a, b, c, 2); - R0(c, d, e, a, b, 3); - R0(b, c, d, e, a, 4); - R0(a, b, c, d, e, 5); - R0(e, a, b, c, d, 6); - R0(d, e, a, b, c, 7); - R0(c, d, e, a, b, 8); - R0(b, c, d, e, a, 9); - R0(a, b, c, d, e, 10); - R0(e, a, b, c, d, 11); - R0(d, e, a, b, c, 12); - R0(c, d, e, a, b, 13); - R0(b, c, d, e, a, 14); - R0(a, b, c, d, e, 15); - R1(e, a, b, c, d, 16); - R1(d, e, a, b, c, 17); - R1(c, d, e, a, b, 18); - R1(b, c, d, e, a, 19); - R2(a, b, c, d, e, 20); - R2(e, a, b, c, d, 21); - R2(d, e, a, b, c, 22); - R2(c, d, e, a, b, 23); - R2(b, c, d, e, a, 24); - R2(a, b, c, d, e, 25); - R2(e, a, b, c, d, 26); - R2(d, e, a, b, c, 27); - R2(c, d, e, a, b, 28); - R2(b, c, d, e, a, 29); - R2(a, b, c, d, e, 30); - R2(e, a, b, c, d, 31); - R2(d, e, a, b, c, 32); - R2(c, d, e, a, b, 33); - R2(b, c, d, e, a, 34); - R2(a, b, c, d, e, 35); - R2(e, a, b, c, d, 36); - R2(d, e, a, b, c, 37); - R2(c, d, e, a, b, 38); - R2(b, c, d, e, a, 39); - R3(a, b, c, d, e, 40); - R3(e, a, b, c, d, 41); - R3(d, e, a, b, c, 42); - R3(c, d, e, a, b, 43); - R3(b, c, d, e, a, 44); - R3(a, b, c, d, e, 45); - R3(e, a, b, c, d, 46); - R3(d, e, a, b, c, 47); - R3(c, d, e, a, b, 48); - R3(b, c, d, e, a, 49); - R3(a, b, c, d, e, 50); - R3(e, a, b, c, d, 51); - R3(d, e, a, b, c, 52); - R3(c, d, e, a, b, 53); - R3(b, c, d, e, a, 54); - R3(a, b, c, d, e, 55); - R3(e, a, b, c, d, 56); - R3(d, e, a, b, c, 57); - R3(c, d, e, a, b, 58); - R3(b, c, d, e, a, 59); - R4(a, b, c, d, e, 60); - R4(e, a, b, c, d, 61); - R4(d, e, a, b, c, 62); - R4(c, d, e, a, b, 63); - R4(b, c, d, e, a, 64); - R4(a, b, c, d, e, 65); - R4(e, a, b, c, d, 66); - R4(d, e, a, b, c, 67); - R4(c, d, e, a, b, 68); - R4(b, c, d, e, a, 69); - R4(a, b, c, d, e, 70); - R4(e, a, b, c, d, 71); - R4(d, e, a, b, c, 72); - R4(c, d, e, a, b, 73); - R4(b, c, d, e, a, 74); - R4(a, b, c, d, e, 75); - R4(e, a, b, c, d, 76); - R4(d, e, a, b, c, 77); - R4(c, d, e, a, b, 78); - R4(b, c, d, e, a, 79); - /* Add the working vars back into context.state[] */ - state[0] += a; - state[1] += b; - state[2] += c; - state[3] += d; - state[4] += e; - /* Wipe variables */ - a = b = c = d = e = 0; -#ifdef SHA1HANDSOFF - memset(block, '\0', sizeof(block)); -#endif -} - - -/* SHA1Init - Initialize new context */ - -void SHA1Init( - SHA1_CTX * context -) -{ - /* SHA1 initialization constants */ - context->state[0] = 0x67452301; - context->state[1] = 0xEFCDAB89; - context->state[2] = 0x98BADCFE; - context->state[3] = 0x10325476; - context->state[4] = 0xC3D2E1F0; - context->count[0] = context->count[1] = 0; -} - - -/* Run your data through this. */ - -void SHA1Update( - SHA1_CTX * context, - const unsigned char *data, - uint32_t len -) -{ - uint32_t i; - - uint32_t j; - - j = context->count[0]; - if ((context->count[0] += len << 3) < j) - context->count[1]++; - context->count[1] += (len >> 29); - j = (j >> 3) & 63; - if ((j + len) > 63) - { - memcpy(&context->buffer[j], data, (i = 64 - j)); - SHA1Transform(context->state, context->buffer); - for (; i + 63 < len; i += 64) - { - SHA1Transform(context->state, &data[i]); - } - j = 0; - } - else - i = 0; - memcpy(&context->buffer[j], &data[i], len - i); -} - - -/* Add padding and return the message digest. */ - -void SHA1Final( - unsigned char digest[20], - SHA1_CTX * context -) -{ - unsigned i; - - unsigned char finalcount[8]; - - unsigned char c; - -#if 0 /* untested "improvement" by DHR */ - /* Convert context->count to a sequence of bytes - * in finalcount. Second element first, but - * big-endian order within element. - * But we do it all backwards. - */ - unsigned char *fcp = &finalcount[8]; - - for (i = 0; i < 2; i++) - { - uint32_t t = context->count[i]; - - int j; - - for (j = 0; j < 4; t >>= 8, j++) - *--fcp = (unsigned char) t} -#else - for (i = 0; i < 8; i++) - { - finalcount[i] = (unsigned char) ((context->count[(i >= 4 ? 0 : 1)] >> ((3 - (i & 3)) * 8)) & 255); /* Endian independent */ - } -#endif - c = 0200; - SHA1Update(context, &c, 1); - while ((context->count[0] & 504) != 448) - { - c = 0000; - SHA1Update(context, &c, 1); - } - SHA1Update(context, finalcount, 8); /* Should cause a SHA1Transform() */ - for (i = 0; i < 20; i++) - { - digest[i] = (unsigned char) - ((context->state[i >> 2] >> ((3 - (i & 3)) * 8)) & 255); - } - /* Wipe variables */ - memset(context, '\0', sizeof(*context)); - memset(&finalcount, '\0', sizeof(finalcount)); -} - -void SHA1( - char *hash_out, - const char *str, - uint32_t len) -{ - SHA1_CTX ctx; - unsigned int ii; - - SHA1Init(&ctx); - for (ii=0; ii - 100% Public Domain - */ - -#include "stdint.h" - -#if defined(__cplusplus) -extern "C" { -#endif - -typedef struct -{ - uint32_t state[5]; - uint32_t count[2]; - unsigned char buffer[64]; -} SHA1_CTX; - -void SHA1Transform( - uint32_t state[5], - const unsigned char buffer[64] - ); - -void SHA1Init( - SHA1_CTX * context - ); - -void SHA1Update( - SHA1_CTX * context, - const unsigned char *data, - uint32_t len - ); - -void SHA1Final( - unsigned char digest[20], - SHA1_CTX * context - ); - -void SHA1( - char *hash_out, - const char *str, - uint32_t len); - -#if defined(__cplusplus) -} -#endif - -#endif /* SHA1_H */ diff --git a/examples/gguf-hash/deps/sha256/sha256.c b/examples/gguf-hash/deps/sha256/sha256.c deleted file mode 100644 index a7a87aeb2003..000000000000 --- a/examples/gguf-hash/deps/sha256/sha256.c +++ /dev/null @@ -1,221 +0,0 @@ -/* Crypto/Sha256.c -- SHA-256 Hash -2010-06-11 : Igor Pavlov : Public domain -This code is based on public domain code from Wei Dai's Crypto++ library. */ - -#include "rotate-bits/rotate-bits.h" -#include "sha256.h" - -/* define it for speed optimization */ -#define _SHA256_UNROLL -#define _SHA256_UNROLL2 - -void -sha256_init(sha256_t *p) -{ - p->state[0] = 0x6a09e667; - p->state[1] = 0xbb67ae85; - p->state[2] = 0x3c6ef372; - p->state[3] = 0xa54ff53a; - p->state[4] = 0x510e527f; - p->state[5] = 0x9b05688c; - p->state[6] = 0x1f83d9ab; - p->state[7] = 0x5be0cd19; - p->count = 0; -} - -#define S0(x) (ROTR32(x, 2) ^ ROTR32(x,13) ^ ROTR32(x, 22)) -#define S1(x) (ROTR32(x, 6) ^ ROTR32(x,11) ^ ROTR32(x, 25)) -#define s0(x) (ROTR32(x, 7) ^ ROTR32(x,18) ^ (x >> 3)) -#define s1(x) (ROTR32(x,17) ^ ROTR32(x,19) ^ (x >> 10)) - -#define blk0(i) (W[i] = data[i]) -#define blk2(i) (W[i&15] += s1(W[(i-2)&15]) + W[(i-7)&15] + s0(W[(i-15)&15])) - -#define Ch(x,y,z) (z^(x&(y^z))) -#define Maj(x,y,z) ((x&y)|(z&(x|y))) - -#define a(i) T[(0-(i))&7] -#define b(i) T[(1-(i))&7] -#define c(i) T[(2-(i))&7] -#define d(i) T[(3-(i))&7] -#define e(i) T[(4-(i))&7] -#define f(i) T[(5-(i))&7] -#define g(i) T[(6-(i))&7] -#define h(i) T[(7-(i))&7] - - -#ifdef _SHA256_UNROLL2 - -#define R(a,b,c,d,e,f,g,h, i) h += S1(e) + Ch(e,f,g) + K[i+j] + (j?blk2(i):blk0(i));\ - d += h; h += S0(a) + Maj(a, b, c) - -#define RX_8(i) \ - R(a,b,c,d,e,f,g,h, i); \ - R(h,a,b,c,d,e,f,g, (i+1)); \ - R(g,h,a,b,c,d,e,f, (i+2)); \ - R(f,g,h,a,b,c,d,e, (i+3)); \ - R(e,f,g,h,a,b,c,d, (i+4)); \ - R(d,e,f,g,h,a,b,c, (i+5)); \ - R(c,d,e,f,g,h,a,b, (i+6)); \ - R(b,c,d,e,f,g,h,a, (i+7)) - -#else - -#define R(i) h(i) += S1(e(i)) + Ch(e(i),f(i),g(i)) + K[i+j] + (j?blk2(i):blk0(i));\ - d(i) += h(i); h(i) += S0(a(i)) + Maj(a(i), b(i), c(i)) - -#ifdef _SHA256_UNROLL - -#define RX_8(i) R(i+0); R(i+1); R(i+2); R(i+3); R(i+4); R(i+5); R(i+6); R(i+7); - -#endif - -#endif - -static const uint32_t K[64] = { - 0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, - 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5, - 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, - 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174, - 0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, - 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da, - 0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, - 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967, - 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, - 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, - 0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, - 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070, - 0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, - 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3, - 0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, - 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2 -}; - -static void -sha256_transform(uint32_t *state, const uint32_t *data) -{ - uint32_t W[16] = {0}; - unsigned j; - #ifdef _SHA256_UNROLL2 - uint32_t a,b,c,d,e,f,g,h; - a = state[0]; - b = state[1]; - c = state[2]; - d = state[3]; - e = state[4]; - f = state[5]; - g = state[6]; - h = state[7]; - #else - uint32_t T[8]; - for (j = 0; j < 8; j++) - T[j] = state[j]; - #endif - - for (j = 0; j < 64; j += 16) - { - #if defined(_SHA256_UNROLL) || defined(_SHA256_UNROLL2) - RX_8(0); RX_8(8); - #else - unsigned i; - for (i = 0; i < 16; i++) { R(i); } - #endif - } - - #ifdef _SHA256_UNROLL2 - state[0] += a; - state[1] += b; - state[2] += c; - state[3] += d; - state[4] += e; - state[5] += f; - state[6] += g; - state[7] += h; - #else - for (j = 0; j < 8; j++) - state[j] += T[j]; - #endif - - /* Wipe variables */ - /* memset(W, 0, sizeof(W)); */ - /* memset(T, 0, sizeof(T)); */ -} - -#undef S0 -#undef S1 -#undef s0 -#undef s1 - -static void -sha256_write_byte_block(sha256_t *p) -{ - uint32_t data32[16]; - unsigned i; - for (i = 0; i < 16; i++) - data32[i] = - ((uint32_t)(p->buffer[i * 4 ]) << 24) + - ((uint32_t)(p->buffer[i * 4 + 1]) << 16) + - ((uint32_t)(p->buffer[i * 4 + 2]) << 8) + - ((uint32_t)(p->buffer[i * 4 + 3])); - sha256_transform(p->state, data32); -} - - -void -sha256_hash(unsigned char *buf, const unsigned char *data, size_t size) -{ - sha256_t hash; - sha256_init(&hash); - sha256_update(&hash, data, size); - sha256_final(&hash, buf); -} - - -void -sha256_update(sha256_t *p, const unsigned char *data, size_t size) -{ - uint32_t curBufferPos = (uint32_t)p->count & 0x3F; - while (size > 0) - { - p->buffer[curBufferPos++] = *data++; - p->count++; - size--; - if (curBufferPos == 64) - { - curBufferPos = 0; - sha256_write_byte_block(p); - } - } -} - - -void -sha256_final(sha256_t *p, unsigned char *digest) -{ - uint64_t lenInBits = (p->count << 3); - uint32_t curBufferPos = (uint32_t)p->count & 0x3F; - unsigned i; - p->buffer[curBufferPos++] = 0x80; - while (curBufferPos != (64 - 8)) - { - curBufferPos &= 0x3F; - if (curBufferPos == 0) - sha256_write_byte_block(p); - p->buffer[curBufferPos++] = 0; - } - for (i = 0; i < 8; i++) - { - p->buffer[curBufferPos++] = (unsigned char)(lenInBits >> 56); - lenInBits <<= 8; - } - sha256_write_byte_block(p); - - for (i = 0; i < 8; i++) - { - *digest++ = (unsigned char)(p->state[i] >> 24); - *digest++ = (unsigned char)(p->state[i] >> 16); - *digest++ = (unsigned char)(p->state[i] >> 8); - *digest++ = (unsigned char)(p->state[i]); - } - sha256_init(p); -} diff --git a/examples/gguf-hash/deps/sha256/sha256.h b/examples/gguf-hash/deps/sha256/sha256.h deleted file mode 100644 index 21657e66b602..000000000000 --- a/examples/gguf-hash/deps/sha256/sha256.h +++ /dev/null @@ -1,24 +0,0 @@ -/* Sha256.h -- SHA-256 Hash -2010-06-11 : Igor Pavlov : Public domain */ - -#ifndef __CRYPTO_SHA256_H -#define __CRYPTO_SHA256_H - -#include -#include - -#define SHA256_DIGEST_SIZE 32 - -typedef struct sha256_t -{ - uint32_t state[8]; - uint64_t count; - unsigned char buffer[64]; -} sha256_t; - -void sha256_init(sha256_t *p); -void sha256_update(sha256_t *p, const unsigned char *data, size_t size); -void sha256_final(sha256_t *p, unsigned char *digest); -void sha256_hash(unsigned char *buf, const unsigned char *data, size_t size); - -#endif diff --git a/examples/gguf-hash/deps/xxhash/xxhash.c b/examples/gguf-hash/deps/xxhash/xxhash.c deleted file mode 100644 index e60cc37f13c2..000000000000 --- a/examples/gguf-hash/deps/xxhash/xxhash.c +++ /dev/null @@ -1,42 +0,0 @@ -/* - * xxHash - Extremely Fast Hash algorithm - * Copyright (C) 2012-2023 Yann Collet - * - * BSD 2-Clause License (https://www.opensource.org/licenses/bsd-license.php) - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are - * met: - * - * * Redistributions of source code must retain the above copyright - * notice, this list of conditions and the following disclaimer. - * * Redistributions in binary form must reproduce the above - * copyright notice, this list of conditions and the following disclaimer - * in the documentation and/or other materials provided with the - * distribution. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS - * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT - * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR - * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT - * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, - * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT - * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, - * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY - * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT - * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - * You can contact the author at: - * - xxHash homepage: https://www.xxhash.com - * - xxHash source repository: https://github.com/Cyan4973/xxHash - */ - -/* - * xxhash.c instantiates functions defined in xxhash.h - */ - -#define XXH_STATIC_LINKING_ONLY /* access advanced declarations */ -#define XXH_IMPLEMENTATION /* access definitions */ - -#include "xxhash.h" diff --git a/examples/gguf-hash/deps/xxhash/xxhash.h b/examples/gguf-hash/deps/xxhash/xxhash.h deleted file mode 100644 index c0fafe20d54a..000000000000 --- a/examples/gguf-hash/deps/xxhash/xxhash.h +++ /dev/null @@ -1,7093 +0,0 @@ -/* - * xxHash - Extremely Fast Hash algorithm - * Header File - * Copyright (C) 2012-2023 Yann Collet - * - * BSD 2-Clause License (https://www.opensource.org/licenses/bsd-license.php) - * - * Redistribution and use in source and binary forms, with or without - * modification, are permitted provided that the following conditions are - * met: - * - * * Redistributions of source code must retain the above copyright - * notice, this list of conditions and the following disclaimer. - * * Redistributions in binary form must reproduce the above - * copyright notice, this list of conditions and the following disclaimer - * in the documentation and/or other materials provided with the - * distribution. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS - * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT - * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR - * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT - * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, - * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT - * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, - * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY - * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT - * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - * You can contact the author at: - * - xxHash homepage: https://www.xxhash.com - * - xxHash source repository: https://github.com/Cyan4973/xxHash - */ - -/*! - * @mainpage xxHash - * - * xxHash is an extremely fast non-cryptographic hash algorithm, working at RAM speed - * limits. - * - * It is proposed in four flavors, in three families: - * 1. @ref XXH32_family - * - Classic 32-bit hash function. Simple, compact, and runs on almost all - * 32-bit and 64-bit systems. - * 2. @ref XXH64_family - * - Classic 64-bit adaptation of XXH32. Just as simple, and runs well on most - * 64-bit systems (but _not_ 32-bit systems). - * 3. @ref XXH3_family - * - Modern 64-bit and 128-bit hash function family which features improved - * strength and performance across the board, especially on smaller data. - * It benefits greatly from SIMD and 64-bit without requiring it. - * - * Benchmarks - * --- - * The reference system uses an Intel i7-9700K CPU, and runs Ubuntu x64 20.04. - * The open source benchmark program is compiled with clang v10.0 using -O3 flag. - * - * | Hash Name | ISA ext | Width | Large Data Speed | Small Data Velocity | - * | -------------------- | ------- | ----: | ---------------: | ------------------: | - * | XXH3_64bits() | @b AVX2 | 64 | 59.4 GB/s | 133.1 | - * | MeowHash | AES-NI | 128 | 58.2 GB/s | 52.5 | - * | XXH3_128bits() | @b AVX2 | 128 | 57.9 GB/s | 118.1 | - * | CLHash | PCLMUL | 64 | 37.1 GB/s | 58.1 | - * | XXH3_64bits() | @b SSE2 | 64 | 31.5 GB/s | 133.1 | - * | XXH3_128bits() | @b SSE2 | 128 | 29.6 GB/s | 118.1 | - * | RAM sequential read | | N/A | 28.0 GB/s | N/A | - * | ahash | AES-NI | 64 | 22.5 GB/s | 107.2 | - * | City64 | | 64 | 22.0 GB/s | 76.6 | - * | T1ha2 | | 64 | 22.0 GB/s | 99.0 | - * | City128 | | 128 | 21.7 GB/s | 57.7 | - * | FarmHash | AES-NI | 64 | 21.3 GB/s | 71.9 | - * | XXH64() | | 64 | 19.4 GB/s | 71.0 | - * | SpookyHash | | 64 | 19.3 GB/s | 53.2 | - * | Mum | | 64 | 18.0 GB/s | 67.0 | - * | CRC32C | SSE4.2 | 32 | 13.0 GB/s | 57.9 | - * | XXH32() | | 32 | 9.7 GB/s | 71.9 | - * | City32 | | 32 | 9.1 GB/s | 66.0 | - * | Blake3* | @b AVX2 | 256 | 4.4 GB/s | 8.1 | - * | Murmur3 | | 32 | 3.9 GB/s | 56.1 | - * | SipHash* | | 64 | 3.0 GB/s | 43.2 | - * | Blake3* | @b SSE2 | 256 | 2.4 GB/s | 8.1 | - * | HighwayHash | | 64 | 1.4 GB/s | 6.0 | - * | FNV64 | | 64 | 1.2 GB/s | 62.7 | - * | Blake2* | | 256 | 1.1 GB/s | 5.1 | - * | SHA1* | | 160 | 0.8 GB/s | 5.6 | - * | MD5* | | 128 | 0.6 GB/s | 7.8 | - * @note - * - Hashes which require a specific ISA extension are noted. SSE2 is also noted, - * even though it is mandatory on x64. - * - Hashes with an asterisk are cryptographic. Note that MD5 is non-cryptographic - * by modern standards. - * - Small data velocity is a rough average of algorithm's efficiency for small - * data. For more accurate information, see the wiki. - * - More benchmarks and strength tests are found on the wiki: - * https://github.com/Cyan4973/xxHash/wiki - * - * Usage - * ------ - * All xxHash variants use a similar API. Changing the algorithm is a trivial - * substitution. - * - * @pre - * For functions which take an input and length parameter, the following - * requirements are assumed: - * - The range from [`input`, `input + length`) is valid, readable memory. - * - The only exception is if the `length` is `0`, `input` may be `NULL`. - * - For C++, the objects must have the *TriviallyCopyable* property, as the - * functions access bytes directly as if it was an array of `unsigned char`. - * - * @anchor single_shot_example - * **Single Shot** - * - * These functions are stateless functions which hash a contiguous block of memory, - * immediately returning the result. They are the easiest and usually the fastest - * option. - * - * XXH32(), XXH64(), XXH3_64bits(), XXH3_128bits() - * - * @code{.c} - * #include - * #include "xxhash.h" - * - * // Example for a function which hashes a null terminated string with XXH32(). - * XXH32_hash_t hash_string(const char* string, XXH32_hash_t seed) - * { - * // NULL pointers are only valid if the length is zero - * size_t length = (string == NULL) ? 0 : strlen(string); - * return XXH32(string, length, seed); - * } - * @endcode - * - * - * @anchor streaming_example - * **Streaming** - * - * These groups of functions allow incremental hashing of unknown size, even - * more than what would fit in a size_t. - * - * XXH32_reset(), XXH64_reset(), XXH3_64bits_reset(), XXH3_128bits_reset() - * - * @code{.c} - * #include - * #include - * #include "xxhash.h" - * // Example for a function which hashes a FILE incrementally with XXH3_64bits(). - * XXH64_hash_t hashFile(FILE* f) - * { - * // Allocate a state struct. Do not just use malloc() or new. - * XXH3_state_t* state = XXH3_createState(); - * assert(state != NULL && "Out of memory!"); - * // Reset the state to start a new hashing session. - * XXH3_64bits_reset(state); - * char buffer[4096]; - * size_t count; - * // Read the file in chunks - * while ((count = fread(buffer, 1, sizeof(buffer), f)) != 0) { - * // Run update() as many times as necessary to process the data - * XXH3_64bits_update(state, buffer, count); - * } - * // Retrieve the finalized hash. This will not change the state. - * XXH64_hash_t result = XXH3_64bits_digest(state); - * // Free the state. Do not use free(). - * XXH3_freeState(state); - * return result; - * } - * @endcode - * - * Streaming functions generate the xxHash value from an incremental input. - * This method is slower than single-call functions, due to state management. - * For small inputs, prefer `XXH32()` and `XXH64()`, which are better optimized. - * - * An XXH state must first be allocated using `XXH*_createState()`. - * - * Start a new hash by initializing the state with a seed using `XXH*_reset()`. - * - * Then, feed the hash state by calling `XXH*_update()` as many times as necessary. - * - * The function returns an error code, with 0 meaning OK, and any other value - * meaning there is an error. - * - * Finally, a hash value can be produced anytime, by using `XXH*_digest()`. - * This function returns the nn-bits hash as an int or long long. - * - * It's still possible to continue inserting input into the hash state after a - * digest, and generate new hash values later on by invoking `XXH*_digest()`. - * - * When done, release the state using `XXH*_freeState()`. - * - * - * @anchor canonical_representation_example - * **Canonical Representation** - * - * The default return values from XXH functions are unsigned 32, 64 and 128 bit - * integers. - * This the simplest and fastest format for further post-processing. - * - * However, this leaves open the question of what is the order on the byte level, - * since little and big endian conventions will store the same number differently. - * - * The canonical representation settles this issue by mandating big-endian - * convention, the same convention as human-readable numbers (large digits first). - * - * When writing hash values to storage, sending them over a network, or printing - * them, it's highly recommended to use the canonical representation to ensure - * portability across a wider range of systems, present and future. - * - * The following functions allow transformation of hash values to and from - * canonical format. - * - * XXH32_canonicalFromHash(), XXH32_hashFromCanonical(), - * XXH64_canonicalFromHash(), XXH64_hashFromCanonical(), - * XXH128_canonicalFromHash(), XXH128_hashFromCanonical(), - * - * @code{.c} - * #include - * #include "xxhash.h" - * - * // Example for a function which prints XXH32_hash_t in human readable format - * void printXxh32(XXH32_hash_t hash) - * { - * XXH32_canonical_t cano; - * XXH32_canonicalFromHash(&cano, hash); - * size_t i; - * for(i = 0; i < sizeof(cano.digest); ++i) { - * printf("%02x", cano.digest[i]); - * } - * printf("\n"); - * } - * - * // Example for a function which converts XXH32_canonical_t to XXH32_hash_t - * XXH32_hash_t convertCanonicalToXxh32(XXH32_canonical_t cano) - * { - * XXH32_hash_t hash = XXH32_hashFromCanonical(&cano); - * return hash; - * } - * @endcode - * - * - * @file xxhash.h - * xxHash prototypes and implementation - */ - -#if defined (__cplusplus) -extern "C" { -#endif - -/* **************************** - * INLINE mode - ******************************/ -/*! - * @defgroup public Public API - * Contains details on the public xxHash functions. - * @{ - */ -#ifdef XXH_DOXYGEN -/*! - * @brief Gives access to internal state declaration, required for static allocation. - * - * Incompatible with dynamic linking, due to risks of ABI changes. - * - * Usage: - * @code{.c} - * #define XXH_STATIC_LINKING_ONLY - * #include "xxhash.h" - * @endcode - */ -# define XXH_STATIC_LINKING_ONLY -/* Do not undef XXH_STATIC_LINKING_ONLY for Doxygen */ - -/*! - * @brief Gives access to internal definitions. - * - * Usage: - * @code{.c} - * #define XXH_STATIC_LINKING_ONLY - * #define XXH_IMPLEMENTATION - * #include "xxhash.h" - * @endcode - */ -# define XXH_IMPLEMENTATION -/* Do not undef XXH_IMPLEMENTATION for Doxygen */ - -/*! - * @brief Exposes the implementation and marks all functions as `inline`. - * - * Use these build macros to inline xxhash into the target unit. - * Inlining improves performance on small inputs, especially when the length is - * expressed as a compile-time constant: - * - * https://fastcompression.blogspot.com/2018/03/xxhash-for-small-keys-impressive-power.html - * - * It also keeps xxHash symbols private to the unit, so they are not exported. - * - * Usage: - * @code{.c} - * #define XXH_INLINE_ALL - * #include "xxhash.h" - * @endcode - * Do not compile and link xxhash.o as a separate object, as it is not useful. - */ -# define XXH_INLINE_ALL -# undef XXH_INLINE_ALL -/*! - * @brief Exposes the implementation without marking functions as inline. - */ -# define XXH_PRIVATE_API -# undef XXH_PRIVATE_API -/*! - * @brief Emulate a namespace by transparently prefixing all symbols. - * - * If you want to include _and expose_ xxHash functions from within your own - * library, but also want to avoid symbol collisions with other libraries which - * may also include xxHash, you can use @ref XXH_NAMESPACE to automatically prefix - * any public symbol from xxhash library with the value of @ref XXH_NAMESPACE - * (therefore, avoid empty or numeric values). - * - * Note that no change is required within the calling program as long as it - * includes `xxhash.h`: Regular symbol names will be automatically translated - * by this header. - */ -# define XXH_NAMESPACE /* YOUR NAME HERE */ -# undef XXH_NAMESPACE -#endif - -#if (defined(XXH_INLINE_ALL) || defined(XXH_PRIVATE_API)) \ - && !defined(XXH_INLINE_ALL_31684351384) - /* this section should be traversed only once */ -# define XXH_INLINE_ALL_31684351384 - /* give access to the advanced API, required to compile implementations */ -# undef XXH_STATIC_LINKING_ONLY /* avoid macro redef */ -# define XXH_STATIC_LINKING_ONLY - /* make all functions private */ -# undef XXH_PUBLIC_API -# if defined(__GNUC__) -# define XXH_PUBLIC_API static __inline __attribute__((__unused__)) -# elif defined (__cplusplus) || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) -# define XXH_PUBLIC_API static inline -# elif defined(_MSC_VER) -# define XXH_PUBLIC_API static __inline -# else - /* note: this version may generate warnings for unused static functions */ -# define XXH_PUBLIC_API static -# endif - - /* - * This part deals with the special case where a unit wants to inline xxHash, - * but "xxhash.h" has previously been included without XXH_INLINE_ALL, - * such as part of some previously included *.h header file. - * Without further action, the new include would just be ignored, - * and functions would effectively _not_ be inlined (silent failure). - * The following macros solve this situation by prefixing all inlined names, - * avoiding naming collision with previous inclusions. - */ - /* Before that, we unconditionally #undef all symbols, - * in case they were already defined with XXH_NAMESPACE. - * They will then be redefined for XXH_INLINE_ALL - */ -# undef XXH_versionNumber - /* XXH32 */ -# undef XXH32 -# undef XXH32_createState -# undef XXH32_freeState -# undef XXH32_reset -# undef XXH32_update -# undef XXH32_digest -# undef XXH32_copyState -# undef XXH32_canonicalFromHash -# undef XXH32_hashFromCanonical - /* XXH64 */ -# undef XXH64 -# undef XXH64_createState -# undef XXH64_freeState -# undef XXH64_reset -# undef XXH64_update -# undef XXH64_digest -# undef XXH64_copyState -# undef XXH64_canonicalFromHash -# undef XXH64_hashFromCanonical - /* XXH3_64bits */ -# undef XXH3_64bits -# undef XXH3_64bits_withSecret -# undef XXH3_64bits_withSeed -# undef XXH3_64bits_withSecretandSeed -# undef XXH3_createState -# undef XXH3_freeState -# undef XXH3_copyState -# undef XXH3_64bits_reset -# undef XXH3_64bits_reset_withSeed -# undef XXH3_64bits_reset_withSecret -# undef XXH3_64bits_update -# undef XXH3_64bits_digest -# undef XXH3_generateSecret - /* XXH3_128bits */ -# undef XXH128 -# undef XXH3_128bits -# undef XXH3_128bits_withSeed -# undef XXH3_128bits_withSecret -# undef XXH3_128bits_reset -# undef XXH3_128bits_reset_withSeed -# undef XXH3_128bits_reset_withSecret -# undef XXH3_128bits_reset_withSecretandSeed -# undef XXH3_128bits_update -# undef XXH3_128bits_digest -# undef XXH128_isEqual -# undef XXH128_cmp -# undef XXH128_canonicalFromHash -# undef XXH128_hashFromCanonical - /* Finally, free the namespace itself */ -# undef XXH_NAMESPACE - - /* employ the namespace for XXH_INLINE_ALL */ -# define XXH_NAMESPACE XXH_INLINE_ - /* - * Some identifiers (enums, type names) are not symbols, - * but they must nonetheless be renamed to avoid redeclaration. - * Alternative solution: do not redeclare them. - * However, this requires some #ifdefs, and has a more dispersed impact. - * Meanwhile, renaming can be achieved in a single place. - */ -# define XXH_IPREF(Id) XXH_NAMESPACE ## Id -# define XXH_OK XXH_IPREF(XXH_OK) -# define XXH_ERROR XXH_IPREF(XXH_ERROR) -# define XXH_errorcode XXH_IPREF(XXH_errorcode) -# define XXH32_canonical_t XXH_IPREF(XXH32_canonical_t) -# define XXH64_canonical_t XXH_IPREF(XXH64_canonical_t) -# define XXH128_canonical_t XXH_IPREF(XXH128_canonical_t) -# define XXH32_state_s XXH_IPREF(XXH32_state_s) -# define XXH32_state_t XXH_IPREF(XXH32_state_t) -# define XXH64_state_s XXH_IPREF(XXH64_state_s) -# define XXH64_state_t XXH_IPREF(XXH64_state_t) -# define XXH3_state_s XXH_IPREF(XXH3_state_s) -# define XXH3_state_t XXH_IPREF(XXH3_state_t) -# define XXH128_hash_t XXH_IPREF(XXH128_hash_t) - /* Ensure the header is parsed again, even if it was previously included */ -# undef XXHASH_H_5627135585666179 -# undef XXHASH_H_STATIC_13879238742 -#endif /* XXH_INLINE_ALL || XXH_PRIVATE_API */ - -/* **************************************************************** - * Stable API - *****************************************************************/ -#ifndef XXHASH_H_5627135585666179 -#define XXHASH_H_5627135585666179 1 - -/*! @brief Marks a global symbol. */ -#if !defined(XXH_INLINE_ALL) && !defined(XXH_PRIVATE_API) -# if defined(_WIN32) && defined(_MSC_VER) && (defined(XXH_IMPORT) || defined(XXH_EXPORT)) -# ifdef XXH_EXPORT -# define XXH_PUBLIC_API __declspec(dllexport) -# elif XXH_IMPORT -# define XXH_PUBLIC_API __declspec(dllimport) -# endif -# else -# define XXH_PUBLIC_API /* do nothing */ -# endif -#endif - -#ifdef XXH_NAMESPACE -# define XXH_CAT(A,B) A##B -# define XXH_NAME2(A,B) XXH_CAT(A,B) -# define XXH_versionNumber XXH_NAME2(XXH_NAMESPACE, XXH_versionNumber) -/* XXH32 */ -# define XXH32 XXH_NAME2(XXH_NAMESPACE, XXH32) -# define XXH32_createState XXH_NAME2(XXH_NAMESPACE, XXH32_createState) -# define XXH32_freeState XXH_NAME2(XXH_NAMESPACE, XXH32_freeState) -# define XXH32_reset XXH_NAME2(XXH_NAMESPACE, XXH32_reset) -# define XXH32_update XXH_NAME2(XXH_NAMESPACE, XXH32_update) -# define XXH32_digest XXH_NAME2(XXH_NAMESPACE, XXH32_digest) -# define XXH32_copyState XXH_NAME2(XXH_NAMESPACE, XXH32_copyState) -# define XXH32_canonicalFromHash XXH_NAME2(XXH_NAMESPACE, XXH32_canonicalFromHash) -# define XXH32_hashFromCanonical XXH_NAME2(XXH_NAMESPACE, XXH32_hashFromCanonical) -/* XXH64 */ -# define XXH64 XXH_NAME2(XXH_NAMESPACE, XXH64) -# define XXH64_createState XXH_NAME2(XXH_NAMESPACE, XXH64_createState) -# define XXH64_freeState XXH_NAME2(XXH_NAMESPACE, XXH64_freeState) -# define XXH64_reset XXH_NAME2(XXH_NAMESPACE, XXH64_reset) -# define XXH64_update XXH_NAME2(XXH_NAMESPACE, XXH64_update) -# define XXH64_digest XXH_NAME2(XXH_NAMESPACE, XXH64_digest) -# define XXH64_copyState XXH_NAME2(XXH_NAMESPACE, XXH64_copyState) -# define XXH64_canonicalFromHash XXH_NAME2(XXH_NAMESPACE, XXH64_canonicalFromHash) -# define XXH64_hashFromCanonical XXH_NAME2(XXH_NAMESPACE, XXH64_hashFromCanonical) -/* XXH3_64bits */ -# define XXH3_64bits XXH_NAME2(XXH_NAMESPACE, XXH3_64bits) -# define XXH3_64bits_withSecret XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_withSecret) -# define XXH3_64bits_withSeed XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_withSeed) -# define XXH3_64bits_withSecretandSeed XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_withSecretandSeed) -# define XXH3_createState XXH_NAME2(XXH_NAMESPACE, XXH3_createState) -# define XXH3_freeState XXH_NAME2(XXH_NAMESPACE, XXH3_freeState) -# define XXH3_copyState XXH_NAME2(XXH_NAMESPACE, XXH3_copyState) -# define XXH3_64bits_reset XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_reset) -# define XXH3_64bits_reset_withSeed XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_reset_withSeed) -# define XXH3_64bits_reset_withSecret XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_reset_withSecret) -# define XXH3_64bits_reset_withSecretandSeed XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_reset_withSecretandSeed) -# define XXH3_64bits_update XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_update) -# define XXH3_64bits_digest XXH_NAME2(XXH_NAMESPACE, XXH3_64bits_digest) -# define XXH3_generateSecret XXH_NAME2(XXH_NAMESPACE, XXH3_generateSecret) -# define XXH3_generateSecret_fromSeed XXH_NAME2(XXH_NAMESPACE, XXH3_generateSecret_fromSeed) -/* XXH3_128bits */ -# define XXH128 XXH_NAME2(XXH_NAMESPACE, XXH128) -# define XXH3_128bits XXH_NAME2(XXH_NAMESPACE, XXH3_128bits) -# define XXH3_128bits_withSeed XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_withSeed) -# define XXH3_128bits_withSecret XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_withSecret) -# define XXH3_128bits_withSecretandSeed XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_withSecretandSeed) -# define XXH3_128bits_reset XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_reset) -# define XXH3_128bits_reset_withSeed XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_reset_withSeed) -# define XXH3_128bits_reset_withSecret XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_reset_withSecret) -# define XXH3_128bits_reset_withSecretandSeed XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_reset_withSecretandSeed) -# define XXH3_128bits_update XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_update) -# define XXH3_128bits_digest XXH_NAME2(XXH_NAMESPACE, XXH3_128bits_digest) -# define XXH128_isEqual XXH_NAME2(XXH_NAMESPACE, XXH128_isEqual) -# define XXH128_cmp XXH_NAME2(XXH_NAMESPACE, XXH128_cmp) -# define XXH128_canonicalFromHash XXH_NAME2(XXH_NAMESPACE, XXH128_canonicalFromHash) -# define XXH128_hashFromCanonical XXH_NAME2(XXH_NAMESPACE, XXH128_hashFromCanonical) -#endif - - -/* ************************************* -* Compiler specifics -***************************************/ - -/* specific declaration modes for Windows */ -#if !defined(XXH_INLINE_ALL) && !defined(XXH_PRIVATE_API) -# if defined(_WIN32) && defined(_MSC_VER) && (defined(XXH_IMPORT) || defined(XXH_EXPORT)) -# ifdef XXH_EXPORT -# define XXH_PUBLIC_API __declspec(dllexport) -# elif XXH_IMPORT -# define XXH_PUBLIC_API __declspec(dllimport) -# endif -# else -# define XXH_PUBLIC_API /* do nothing */ -# endif -#endif - -#if defined (__GNUC__) -# define XXH_CONSTF __attribute__((__const__)) -# define XXH_PUREF __attribute__((__pure__)) -# define XXH_MALLOCF __attribute__((__malloc__)) -#else -# define XXH_CONSTF /* disable */ -# define XXH_PUREF -# define XXH_MALLOCF -#endif - -/* ************************************* -* Version -***************************************/ -#define XXH_VERSION_MAJOR 0 -#define XXH_VERSION_MINOR 8 -#define XXH_VERSION_RELEASE 3 -/*! @brief Version number, encoded as two digits each */ -#define XXH_VERSION_NUMBER (XXH_VERSION_MAJOR *100*100 + XXH_VERSION_MINOR *100 + XXH_VERSION_RELEASE) - -/*! - * @brief Obtains the xxHash version. - * - * This is mostly useful when xxHash is compiled as a shared library, - * since the returned value comes from the library, as opposed to header file. - * - * @return @ref XXH_VERSION_NUMBER of the invoked library. - */ -XXH_PUBLIC_API XXH_CONSTF unsigned XXH_versionNumber (void); - - -/* **************************** -* Common basic types -******************************/ -#include /* size_t */ -/*! - * @brief Exit code for the streaming API. - */ -typedef enum { - XXH_OK = 0, /*!< OK */ - XXH_ERROR /*!< Error */ -} XXH_errorcode; - - -/*-********************************************************************** -* 32-bit hash -************************************************************************/ -#if defined(XXH_DOXYGEN) /* Don't show include */ -/*! - * @brief An unsigned 32-bit integer. - * - * Not necessarily defined to `uint32_t` but functionally equivalent. - */ -typedef uint32_t XXH32_hash_t; - -#elif !defined (__VMS) \ - && (defined (__cplusplus) \ - || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) ) -# ifdef _AIX -# include -# else -# include -# endif - typedef uint32_t XXH32_hash_t; - -#else -# include -# if UINT_MAX == 0xFFFFFFFFUL - typedef unsigned int XXH32_hash_t; -# elif ULONG_MAX == 0xFFFFFFFFUL - typedef unsigned long XXH32_hash_t; -# else -# error "unsupported platform: need a 32-bit type" -# endif -#endif - -/*! - * @} - * - * @defgroup XXH32_family XXH32 family - * @ingroup public - * Contains functions used in the classic 32-bit xxHash algorithm. - * - * @note - * XXH32 is useful for older platforms, with no or poor 64-bit performance. - * Note that the @ref XXH3_family provides competitive speed for both 32-bit - * and 64-bit systems, and offers true 64/128 bit hash results. - * - * @see @ref XXH64_family, @ref XXH3_family : Other xxHash families - * @see @ref XXH32_impl for implementation details - * @{ - */ - -/*! - * @brief Calculates the 32-bit hash of @p input using xxHash32. - * - * @param input The block of data to be hashed, at least @p length bytes in size. - * @param length The length of @p input, in bytes. - * @param seed The 32-bit seed to alter the hash's output predictably. - * - * @pre - * The memory between @p input and @p input + @p length must be valid, - * readable, contiguous memory. However, if @p length is `0`, @p input may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - * @return The calculated 32-bit xxHash32 value. - * - * @see @ref single_shot_example "Single Shot Example" for an example. - */ -XXH_PUBLIC_API XXH_PUREF XXH32_hash_t XXH32 (const void* input, size_t length, XXH32_hash_t seed); - -#ifndef XXH_NO_STREAM -/*! - * @typedef struct XXH32_state_s XXH32_state_t - * @brief The opaque state struct for the XXH32 streaming API. - * - * @see XXH32_state_s for details. - * @see @ref streaming_example "Streaming Example" - */ -typedef struct XXH32_state_s XXH32_state_t; - -/*! - * @brief Allocates an @ref XXH32_state_t. - * - * @return An allocated pointer of @ref XXH32_state_t on success. - * @return `NULL` on failure. - * - * @note Must be freed with XXH32_freeState(). - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_MALLOCF XXH32_state_t* XXH32_createState(void); -/*! - * @brief Frees an @ref XXH32_state_t. - * - * @param statePtr A pointer to an @ref XXH32_state_t allocated with @ref XXH32_createState(). - * - * @return @ref XXH_OK. - * - * @note @p statePtr must be allocated with XXH32_createState(). - * - * @see @ref streaming_example "Streaming Example" - * - */ -XXH_PUBLIC_API XXH_errorcode XXH32_freeState(XXH32_state_t* statePtr); -/*! - * @brief Copies one @ref XXH32_state_t to another. - * - * @param dst_state The state to copy to. - * @param src_state The state to copy from. - * @pre - * @p dst_state and @p src_state must not be `NULL` and must not overlap. - */ -XXH_PUBLIC_API void XXH32_copyState(XXH32_state_t* dst_state, const XXH32_state_t* src_state); - -/*! - * @brief Resets an @ref XXH32_state_t to begin a new hash. - * - * @param statePtr The state struct to reset. - * @param seed The 32-bit seed to alter the hash result predictably. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @note This function resets and seeds a state. Call it before @ref XXH32_update(). - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_errorcode XXH32_reset (XXH32_state_t* statePtr, XXH32_hash_t seed); - -/*! - * @brief Consumes a block of @p input to an @ref XXH32_state_t. - * - * @param statePtr The state struct to update. - * @param input The block of data to be hashed, at least @p length bytes in size. - * @param length The length of @p input, in bytes. - * - * @pre - * @p statePtr must not be `NULL`. - * @pre - * The memory between @p input and @p input + @p length must be valid, - * readable, contiguous memory. However, if @p length is `0`, @p input may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @note Call this to incrementally consume blocks of data. - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_errorcode XXH32_update (XXH32_state_t* statePtr, const void* input, size_t length); - -/*! - * @brief Returns the calculated hash value from an @ref XXH32_state_t. - * - * @param statePtr The state struct to calculate the hash from. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return The calculated 32-bit xxHash32 value from that state. - * - * @note - * Calling XXH32_digest() will not affect @p statePtr, so you can update, - * digest, and update again. - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_PUREF XXH32_hash_t XXH32_digest (const XXH32_state_t* statePtr); -#endif /* !XXH_NO_STREAM */ - -/******* Canonical representation *******/ - -/*! - * @brief Canonical (big endian) representation of @ref XXH32_hash_t. - */ -typedef struct { - unsigned char digest[4]; /*!< Hash bytes, big endian */ -} XXH32_canonical_t; - -/*! - * @brief Converts an @ref XXH32_hash_t to a big endian @ref XXH32_canonical_t. - * - * @param dst The @ref XXH32_canonical_t pointer to be stored to. - * @param hash The @ref XXH32_hash_t to be converted. - * - * @pre - * @p dst must not be `NULL`. - * - * @see @ref canonical_representation_example "Canonical Representation Example" - */ -XXH_PUBLIC_API void XXH32_canonicalFromHash(XXH32_canonical_t* dst, XXH32_hash_t hash); - -/*! - * @brief Converts an @ref XXH32_canonical_t to a native @ref XXH32_hash_t. - * - * @param src The @ref XXH32_canonical_t to convert. - * - * @pre - * @p src must not be `NULL`. - * - * @return The converted hash. - * - * @see @ref canonical_representation_example "Canonical Representation Example" - */ -XXH_PUBLIC_API XXH_PUREF XXH32_hash_t XXH32_hashFromCanonical(const XXH32_canonical_t* src); - - -/*! @cond Doxygen ignores this part */ -#ifdef __has_attribute -# define XXH_HAS_ATTRIBUTE(x) __has_attribute(x) -#else -# define XXH_HAS_ATTRIBUTE(x) 0 -#endif -/*! @endcond */ - -/*! @cond Doxygen ignores this part */ -/* - * C23 __STDC_VERSION__ number hasn't been specified yet. For now - * leave as `201711L` (C17 + 1). - * TODO: Update to correct value when its been specified. - */ -#define XXH_C23_VN 201711L -/*! @endcond */ - -/*! @cond Doxygen ignores this part */ -/* C-language Attributes are added in C23. */ -#if defined(__STDC_VERSION__) && (__STDC_VERSION__ >= XXH_C23_VN) && defined(__has_c_attribute) -# define XXH_HAS_C_ATTRIBUTE(x) __has_c_attribute(x) -#else -# define XXH_HAS_C_ATTRIBUTE(x) 0 -#endif -/*! @endcond */ - -/*! @cond Doxygen ignores this part */ -#if defined(__cplusplus) && defined(__has_cpp_attribute) -# define XXH_HAS_CPP_ATTRIBUTE(x) __has_cpp_attribute(x) -#else -# define XXH_HAS_CPP_ATTRIBUTE(x) 0 -#endif -/*! @endcond */ - -/*! @cond Doxygen ignores this part */ -/* - * Define XXH_FALLTHROUGH macro for annotating switch case with the 'fallthrough' attribute - * introduced in CPP17 and C23. - * CPP17 : https://en.cppreference.com/w/cpp/language/attributes/fallthrough - * C23 : https://en.cppreference.com/w/c/language/attributes/fallthrough - */ -#if XXH_HAS_C_ATTRIBUTE(fallthrough) || XXH_HAS_CPP_ATTRIBUTE(fallthrough) -# define XXH_FALLTHROUGH [[fallthrough]] -#elif XXH_HAS_ATTRIBUTE(__fallthrough__) -# define XXH_FALLTHROUGH __attribute__ ((__fallthrough__)) -#else -# define XXH_FALLTHROUGH /* fallthrough */ -#endif -/*! @endcond */ - -/*! @cond Doxygen ignores this part */ -/* - * Define XXH_NOESCAPE for annotated pointers in public API. - * https://clang.llvm.org/docs/AttributeReference.html#noescape - * As of writing this, only supported by clang. - */ -#if XXH_HAS_ATTRIBUTE(noescape) -# define XXH_NOESCAPE __attribute__((__noescape__)) -#else -# define XXH_NOESCAPE -#endif -/*! @endcond */ - - -/*! - * @} - * @ingroup public - * @{ - */ - -#ifndef XXH_NO_LONG_LONG -/*-********************************************************************** -* 64-bit hash -************************************************************************/ -#if defined(XXH_DOXYGEN) /* don't include */ -/*! - * @brief An unsigned 64-bit integer. - * - * Not necessarily defined to `uint64_t` but functionally equivalent. - */ -typedef uint64_t XXH64_hash_t; -#elif !defined (__VMS) \ - && (defined (__cplusplus) \ - || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) ) -# ifdef _AIX -# include -# else -# include -# endif - typedef uint64_t XXH64_hash_t; -#else -# include -# if defined(__LP64__) && ULONG_MAX == 0xFFFFFFFFFFFFFFFFULL - /* LP64 ABI says uint64_t is unsigned long */ - typedef unsigned long XXH64_hash_t; -# else - /* the following type must have a width of 64-bit */ - typedef unsigned long long XXH64_hash_t; -# endif -#endif - -/*! - * @} - * - * @defgroup XXH64_family XXH64 family - * @ingroup public - * @{ - * Contains functions used in the classic 64-bit xxHash algorithm. - * - * @note - * XXH3 provides competitive speed for both 32-bit and 64-bit systems, - * and offers true 64/128 bit hash results. - * It provides better speed for systems with vector processing capabilities. - */ - -/*! - * @brief Calculates the 64-bit hash of @p input using xxHash64. - * - * @param input The block of data to be hashed, at least @p length bytes in size. - * @param length The length of @p input, in bytes. - * @param seed The 64-bit seed to alter the hash's output predictably. - * - * @pre - * The memory between @p input and @p input + @p length must be valid, - * readable, contiguous memory. However, if @p length is `0`, @p input may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - * @return The calculated 64-bit xxHash64 value. - * - * @see @ref single_shot_example "Single Shot Example" for an example. - */ -XXH_PUBLIC_API XXH_PUREF XXH64_hash_t XXH64(XXH_NOESCAPE const void* input, size_t length, XXH64_hash_t seed); - -/******* Streaming *******/ -#ifndef XXH_NO_STREAM -/*! - * @brief The opaque state struct for the XXH64 streaming API. - * - * @see XXH64_state_s for details. - * @see @ref streaming_example "Streaming Example" - */ -typedef struct XXH64_state_s XXH64_state_t; /* incomplete type */ - -/*! - * @brief Allocates an @ref XXH64_state_t. - * - * @return An allocated pointer of @ref XXH64_state_t on success. - * @return `NULL` on failure. - * - * @note Must be freed with XXH64_freeState(). - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_MALLOCF XXH64_state_t* XXH64_createState(void); - -/*! - * @brief Frees an @ref XXH64_state_t. - * - * @param statePtr A pointer to an @ref XXH64_state_t allocated with @ref XXH64_createState(). - * - * @return @ref XXH_OK. - * - * @note @p statePtr must be allocated with XXH64_createState(). - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_errorcode XXH64_freeState(XXH64_state_t* statePtr); - -/*! - * @brief Copies one @ref XXH64_state_t to another. - * - * @param dst_state The state to copy to. - * @param src_state The state to copy from. - * @pre - * @p dst_state and @p src_state must not be `NULL` and must not overlap. - */ -XXH_PUBLIC_API void XXH64_copyState(XXH_NOESCAPE XXH64_state_t* dst_state, const XXH64_state_t* src_state); - -/*! - * @brief Resets an @ref XXH64_state_t to begin a new hash. - * - * @param statePtr The state struct to reset. - * @param seed The 64-bit seed to alter the hash result predictably. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @note This function resets and seeds a state. Call it before @ref XXH64_update(). - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_errorcode XXH64_reset (XXH_NOESCAPE XXH64_state_t* statePtr, XXH64_hash_t seed); - -/*! - * @brief Consumes a block of @p input to an @ref XXH64_state_t. - * - * @param statePtr The state struct to update. - * @param input The block of data to be hashed, at least @p length bytes in size. - * @param length The length of @p input, in bytes. - * - * @pre - * @p statePtr must not be `NULL`. - * @pre - * The memory between @p input and @p input + @p length must be valid, - * readable, contiguous memory. However, if @p length is `0`, @p input may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @note Call this to incrementally consume blocks of data. - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_errorcode XXH64_update (XXH_NOESCAPE XXH64_state_t* statePtr, XXH_NOESCAPE const void* input, size_t length); - -/*! - * @brief Returns the calculated hash value from an @ref XXH64_state_t. - * - * @param statePtr The state struct to calculate the hash from. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return The calculated 64-bit xxHash64 value from that state. - * - * @note - * Calling XXH64_digest() will not affect @p statePtr, so you can update, - * digest, and update again. - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_PUREF XXH64_hash_t XXH64_digest (XXH_NOESCAPE const XXH64_state_t* statePtr); -#endif /* !XXH_NO_STREAM */ -/******* Canonical representation *******/ - -/*! - * @brief Canonical (big endian) representation of @ref XXH64_hash_t. - */ -typedef struct { unsigned char digest[sizeof(XXH64_hash_t)]; } XXH64_canonical_t; - -/*! - * @brief Converts an @ref XXH64_hash_t to a big endian @ref XXH64_canonical_t. - * - * @param dst The @ref XXH64_canonical_t pointer to be stored to. - * @param hash The @ref XXH64_hash_t to be converted. - * - * @pre - * @p dst must not be `NULL`. - * - * @see @ref canonical_representation_example "Canonical Representation Example" - */ -XXH_PUBLIC_API void XXH64_canonicalFromHash(XXH_NOESCAPE XXH64_canonical_t* dst, XXH64_hash_t hash); - -/*! - * @brief Converts an @ref XXH64_canonical_t to a native @ref XXH64_hash_t. - * - * @param src The @ref XXH64_canonical_t to convert. - * - * @pre - * @p src must not be `NULL`. - * - * @return The converted hash. - * - * @see @ref canonical_representation_example "Canonical Representation Example" - */ -XXH_PUBLIC_API XXH_PUREF XXH64_hash_t XXH64_hashFromCanonical(XXH_NOESCAPE const XXH64_canonical_t* src); - -#ifndef XXH_NO_XXH3 - -/*! - * @} - * ************************************************************************ - * @defgroup XXH3_family XXH3 family - * @ingroup public - * @{ - * - * XXH3 is a more recent hash algorithm featuring: - * - Improved speed for both small and large inputs - * - True 64-bit and 128-bit outputs - * - SIMD acceleration - * - Improved 32-bit viability - * - * Speed analysis methodology is explained here: - * - * https://fastcompression.blogspot.com/2019/03/presenting-xxh3.html - * - * Compared to XXH64, expect XXH3 to run approximately - * ~2x faster on large inputs and >3x faster on small ones, - * exact differences vary depending on platform. - * - * XXH3's speed benefits greatly from SIMD and 64-bit arithmetic, - * but does not require it. - * Most 32-bit and 64-bit targets that can run XXH32 smoothly can run XXH3 - * at competitive speeds, even without vector support. Further details are - * explained in the implementation. - * - * XXH3 has a fast scalar implementation, but it also includes accelerated SIMD - * implementations for many common platforms: - * - AVX512 - * - AVX2 - * - SSE2 - * - ARM NEON - * - WebAssembly SIMD128 - * - POWER8 VSX - * - s390x ZVector - * This can be controlled via the @ref XXH_VECTOR macro, but it automatically - * selects the best version according to predefined macros. For the x86 family, an - * automatic runtime dispatcher is included separately in @ref xxh_x86dispatch.c. - * - * XXH3 implementation is portable: - * it has a generic C90 formulation that can be compiled on any platform, - * all implementations generate exactly the same hash value on all platforms. - * Starting from v0.8.0, it's also labelled "stable", meaning that - * any future version will also generate the same hash value. - * - * XXH3 offers 2 variants, _64bits and _128bits. - * - * When only 64 bits are needed, prefer invoking the _64bits variant, as it - * reduces the amount of mixing, resulting in faster speed on small inputs. - * It's also generally simpler to manipulate a scalar return type than a struct. - * - * The API supports one-shot hashing, streaming mode, and custom secrets. - */ -/*-********************************************************************** -* XXH3 64-bit variant -************************************************************************/ - -/*! - * @brief Calculates 64-bit unseeded variant of XXH3 hash of @p input. - * - * @param input The block of data to be hashed, at least @p length bytes in size. - * @param length The length of @p input, in bytes. - * - * @pre - * The memory between @p input and @p input + @p length must be valid, - * readable, contiguous memory. However, if @p length is `0`, @p input may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - * @return The calculated 64-bit XXH3 hash value. - * - * @note - * This is equivalent to @ref XXH3_64bits_withSeed() with a seed of `0`, however - * it may have slightly better performance due to constant propagation of the - * defaults. - * - * @see - * XXH3_64bits_withSeed(), XXH3_64bits_withSecret(): other seeding variants - * @see @ref single_shot_example "Single Shot Example" for an example. - */ -XXH_PUBLIC_API XXH_PUREF XXH64_hash_t XXH3_64bits(XXH_NOESCAPE const void* input, size_t length); - -/*! - * @brief Calculates 64-bit seeded variant of XXH3 hash of @p input. - * - * @param input The block of data to be hashed, at least @p length bytes in size. - * @param length The length of @p input, in bytes. - * @param seed The 64-bit seed to alter the hash result predictably. - * - * @pre - * The memory between @p input and @p input + @p length must be valid, - * readable, contiguous memory. However, if @p length is `0`, @p input may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - * @return The calculated 64-bit XXH3 hash value. - * - * @note - * seed == 0 produces the same results as @ref XXH3_64bits(). - * - * This variant generates a custom secret on the fly based on default secret - * altered using the @p seed value. - * - * While this operation is decently fast, note that it's not completely free. - * - * @see @ref single_shot_example "Single Shot Example" for an example. - */ -XXH_PUBLIC_API XXH_PUREF XXH64_hash_t XXH3_64bits_withSeed(XXH_NOESCAPE const void* input, size_t length, XXH64_hash_t seed); - -/*! - * The bare minimum size for a custom secret. - * - * @see - * XXH3_64bits_withSecret(), XXH3_64bits_reset_withSecret(), - * XXH3_128bits_withSecret(), XXH3_128bits_reset_withSecret(). - */ -#define XXH3_SECRET_SIZE_MIN 136 - -/*! - * @brief Calculates 64-bit variant of XXH3 with a custom "secret". - * - * @param data The block of data to be hashed, at least @p len bytes in size. - * @param len The length of @p data, in bytes. - * @param secret The secret data. - * @param secretSize The length of @p secret, in bytes. - * - * @return The calculated 64-bit XXH3 hash value. - * - * @pre - * The memory between @p data and @p data + @p len must be valid, - * readable, contiguous memory. However, if @p length is `0`, @p data may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - * It's possible to provide any blob of bytes as a "secret" to generate the hash. - * This makes it more difficult for an external actor to prepare an intentional collision. - * The main condition is that @p secretSize *must* be large enough (>= @ref XXH3_SECRET_SIZE_MIN). - * However, the quality of the secret impacts the dispersion of the hash algorithm. - * Therefore, the secret _must_ look like a bunch of random bytes. - * Avoid "trivial" or structured data such as repeated sequences or a text document. - * Whenever in doubt about the "randomness" of the blob of bytes, - * consider employing @ref XXH3_generateSecret() instead (see below). - * It will generate a proper high entropy secret derived from the blob of bytes. - * Another advantage of using XXH3_generateSecret() is that - * it guarantees that all bits within the initial blob of bytes - * will impact every bit of the output. - * This is not necessarily the case when using the blob of bytes directly - * because, when hashing _small_ inputs, only a portion of the secret is employed. - * - * @see @ref single_shot_example "Single Shot Example" for an example. - */ -XXH_PUBLIC_API XXH_PUREF XXH64_hash_t XXH3_64bits_withSecret(XXH_NOESCAPE const void* data, size_t len, XXH_NOESCAPE const void* secret, size_t secretSize); - - -/******* Streaming *******/ -#ifndef XXH_NO_STREAM -/* - * Streaming requires state maintenance. - * This operation costs memory and CPU. - * As a consequence, streaming is slower than one-shot hashing. - * For better performance, prefer one-shot functions whenever applicable. - */ - -/*! - * @brief The opaque state struct for the XXH3 streaming API. - * - * @see XXH3_state_s for details. - * @see @ref streaming_example "Streaming Example" - */ -typedef struct XXH3_state_s XXH3_state_t; -XXH_PUBLIC_API XXH_MALLOCF XXH3_state_t* XXH3_createState(void); -XXH_PUBLIC_API XXH_errorcode XXH3_freeState(XXH3_state_t* statePtr); - -/*! - * @brief Copies one @ref XXH3_state_t to another. - * - * @param dst_state The state to copy to. - * @param src_state The state to copy from. - * @pre - * @p dst_state and @p src_state must not be `NULL` and must not overlap. - */ -XXH_PUBLIC_API void XXH3_copyState(XXH_NOESCAPE XXH3_state_t* dst_state, XXH_NOESCAPE const XXH3_state_t* src_state); - -/*! - * @brief Resets an @ref XXH3_state_t to begin a new hash. - * - * @param statePtr The state struct to reset. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @note - * - This function resets `statePtr` and generate a secret with default parameters. - * - Call this function before @ref XXH3_64bits_update(). - * - Digest will be equivalent to `XXH3_64bits()`. - * - * @see @ref streaming_example "Streaming Example" - * - */ -XXH_PUBLIC_API XXH_errorcode XXH3_64bits_reset(XXH_NOESCAPE XXH3_state_t* statePtr); - -/*! - * @brief Resets an @ref XXH3_state_t with 64-bit seed to begin a new hash. - * - * @param statePtr The state struct to reset. - * @param seed The 64-bit seed to alter the hash result predictably. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @note - * - This function resets `statePtr` and generate a secret from `seed`. - * - Call this function before @ref XXH3_64bits_update(). - * - Digest will be equivalent to `XXH3_64bits_withSeed()`. - * - * @see @ref streaming_example "Streaming Example" - * - */ -XXH_PUBLIC_API XXH_errorcode XXH3_64bits_reset_withSeed(XXH_NOESCAPE XXH3_state_t* statePtr, XXH64_hash_t seed); - -/*! - * @brief Resets an @ref XXH3_state_t with secret data to begin a new hash. - * - * @param statePtr The state struct to reset. - * @param secret The secret data. - * @param secretSize The length of @p secret, in bytes. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @note - * `secret` is referenced, it _must outlive_ the hash streaming session. - * - * Similar to one-shot API, `secretSize` must be >= @ref XXH3_SECRET_SIZE_MIN, - * and the quality of produced hash values depends on secret's entropy - * (secret's content should look like a bunch of random bytes). - * When in doubt about the randomness of a candidate `secret`, - * consider employing `XXH3_generateSecret()` instead (see below). - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_errorcode XXH3_64bits_reset_withSecret(XXH_NOESCAPE XXH3_state_t* statePtr, XXH_NOESCAPE const void* secret, size_t secretSize); - -/*! - * @brief Consumes a block of @p input to an @ref XXH3_state_t. - * - * @param statePtr The state struct to update. - * @param input The block of data to be hashed, at least @p length bytes in size. - * @param length The length of @p input, in bytes. - * - * @pre - * @p statePtr must not be `NULL`. - * @pre - * The memory between @p input and @p input + @p length must be valid, - * readable, contiguous memory. However, if @p length is `0`, @p input may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @note Call this to incrementally consume blocks of data. - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_errorcode XXH3_64bits_update (XXH_NOESCAPE XXH3_state_t* statePtr, XXH_NOESCAPE const void* input, size_t length); - -/*! - * @brief Returns the calculated XXH3 64-bit hash value from an @ref XXH3_state_t. - * - * @param statePtr The state struct to calculate the hash from. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return The calculated XXH3 64-bit hash value from that state. - * - * @note - * Calling XXH3_64bits_digest() will not affect @p statePtr, so you can update, - * digest, and update again. - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_PUREF XXH64_hash_t XXH3_64bits_digest (XXH_NOESCAPE const XXH3_state_t* statePtr); -#endif /* !XXH_NO_STREAM */ - -/* note : canonical representation of XXH3 is the same as XXH64 - * since they both produce XXH64_hash_t values */ - - -/*-********************************************************************** -* XXH3 128-bit variant -************************************************************************/ - -/*! - * @brief The return value from 128-bit hashes. - * - * Stored in little endian order, although the fields themselves are in native - * endianness. - */ -typedef struct { - XXH64_hash_t low64; /*!< `value & 0xFFFFFFFFFFFFFFFF` */ - XXH64_hash_t high64; /*!< `value >> 64` */ -} XXH128_hash_t; - -/*! - * @brief Calculates 128-bit unseeded variant of XXH3 of @p data. - * - * @param data The block of data to be hashed, at least @p length bytes in size. - * @param len The length of @p data, in bytes. - * - * @return The calculated 128-bit variant of XXH3 value. - * - * The 128-bit variant of XXH3 has more strength, but it has a bit of overhead - * for shorter inputs. - * - * This is equivalent to @ref XXH3_128bits_withSeed() with a seed of `0`, however - * it may have slightly better performance due to constant propagation of the - * defaults. - * - * @see XXH3_128bits_withSeed(), XXH3_128bits_withSecret(): other seeding variants - * @see @ref single_shot_example "Single Shot Example" for an example. - */ -XXH_PUBLIC_API XXH_PUREF XXH128_hash_t XXH3_128bits(XXH_NOESCAPE const void* data, size_t len); -/*! @brief Calculates 128-bit seeded variant of XXH3 hash of @p data. - * - * @param data The block of data to be hashed, at least @p length bytes in size. - * @param len The length of @p data, in bytes. - * @param seed The 64-bit seed to alter the hash result predictably. - * - * @return The calculated 128-bit variant of XXH3 value. - * - * @note - * seed == 0 produces the same results as @ref XXH3_64bits(). - * - * This variant generates a custom secret on the fly based on default secret - * altered using the @p seed value. - * - * While this operation is decently fast, note that it's not completely free. - * - * @see XXH3_128bits(), XXH3_128bits_withSecret(): other seeding variants - * @see @ref single_shot_example "Single Shot Example" for an example. - */ -XXH_PUBLIC_API XXH_PUREF XXH128_hash_t XXH3_128bits_withSeed(XXH_NOESCAPE const void* data, size_t len, XXH64_hash_t seed); -/*! - * @brief Calculates 128-bit variant of XXH3 with a custom "secret". - * - * @param data The block of data to be hashed, at least @p len bytes in size. - * @param len The length of @p data, in bytes. - * @param secret The secret data. - * @param secretSize The length of @p secret, in bytes. - * - * @return The calculated 128-bit variant of XXH3 value. - * - * It's possible to provide any blob of bytes as a "secret" to generate the hash. - * This makes it more difficult for an external actor to prepare an intentional collision. - * The main condition is that @p secretSize *must* be large enough (>= @ref XXH3_SECRET_SIZE_MIN). - * However, the quality of the secret impacts the dispersion of the hash algorithm. - * Therefore, the secret _must_ look like a bunch of random bytes. - * Avoid "trivial" or structured data such as repeated sequences or a text document. - * Whenever in doubt about the "randomness" of the blob of bytes, - * consider employing @ref XXH3_generateSecret() instead (see below). - * It will generate a proper high entropy secret derived from the blob of bytes. - * Another advantage of using XXH3_generateSecret() is that - * it guarantees that all bits within the initial blob of bytes - * will impact every bit of the output. - * This is not necessarily the case when using the blob of bytes directly - * because, when hashing _small_ inputs, only a portion of the secret is employed. - * - * @see @ref single_shot_example "Single Shot Example" for an example. - */ -XXH_PUBLIC_API XXH_PUREF XXH128_hash_t XXH3_128bits_withSecret(XXH_NOESCAPE const void* data, size_t len, XXH_NOESCAPE const void* secret, size_t secretSize); - -/******* Streaming *******/ -#ifndef XXH_NO_STREAM -/* - * Streaming requires state maintenance. - * This operation costs memory and CPU. - * As a consequence, streaming is slower than one-shot hashing. - * For better performance, prefer one-shot functions whenever applicable. - * - * XXH3_128bits uses the same XXH3_state_t as XXH3_64bits(). - * Use already declared XXH3_createState() and XXH3_freeState(). - * - * All reset and streaming functions have same meaning as their 64-bit counterpart. - */ - -/*! - * @brief Resets an @ref XXH3_state_t to begin a new hash. - * - * @param statePtr The state struct to reset. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @note - * - This function resets `statePtr` and generate a secret with default parameters. - * - Call it before @ref XXH3_128bits_update(). - * - Digest will be equivalent to `XXH3_128bits()`. - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_errorcode XXH3_128bits_reset(XXH_NOESCAPE XXH3_state_t* statePtr); - -/*! - * @brief Resets an @ref XXH3_state_t with 64-bit seed to begin a new hash. - * - * @param statePtr The state struct to reset. - * @param seed The 64-bit seed to alter the hash result predictably. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @note - * - This function resets `statePtr` and generate a secret from `seed`. - * - Call it before @ref XXH3_128bits_update(). - * - Digest will be equivalent to `XXH3_128bits_withSeed()`. - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_errorcode XXH3_128bits_reset_withSeed(XXH_NOESCAPE XXH3_state_t* statePtr, XXH64_hash_t seed); -/*! - * @brief Resets an @ref XXH3_state_t with secret data to begin a new hash. - * - * @param statePtr The state struct to reset. - * @param secret The secret data. - * @param secretSize The length of @p secret, in bytes. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * `secret` is referenced, it _must outlive_ the hash streaming session. - * Similar to one-shot API, `secretSize` must be >= @ref XXH3_SECRET_SIZE_MIN, - * and the quality of produced hash values depends on secret's entropy - * (secret's content should look like a bunch of random bytes). - * When in doubt about the randomness of a candidate `secret`, - * consider employing `XXH3_generateSecret()` instead (see below). - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_errorcode XXH3_128bits_reset_withSecret(XXH_NOESCAPE XXH3_state_t* statePtr, XXH_NOESCAPE const void* secret, size_t secretSize); - -/*! - * @brief Consumes a block of @p input to an @ref XXH3_state_t. - * - * Call this to incrementally consume blocks of data. - * - * @param statePtr The state struct to update. - * @param input The block of data to be hashed, at least @p length bytes in size. - * @param length The length of @p input, in bytes. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @note - * The memory between @p input and @p input + @p length must be valid, - * readable, contiguous memory. However, if @p length is `0`, @p input may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - */ -XXH_PUBLIC_API XXH_errorcode XXH3_128bits_update (XXH_NOESCAPE XXH3_state_t* statePtr, XXH_NOESCAPE const void* input, size_t length); - -/*! - * @brief Returns the calculated XXH3 128-bit hash value from an @ref XXH3_state_t. - * - * @param statePtr The state struct to calculate the hash from. - * - * @pre - * @p statePtr must not be `NULL`. - * - * @return The calculated XXH3 128-bit hash value from that state. - * - * @note - * Calling XXH3_128bits_digest() will not affect @p statePtr, so you can update, - * digest, and update again. - * - */ -XXH_PUBLIC_API XXH_PUREF XXH128_hash_t XXH3_128bits_digest (XXH_NOESCAPE const XXH3_state_t* statePtr); -#endif /* !XXH_NO_STREAM */ - -/* Following helper functions make it possible to compare XXH128_hast_t values. - * Since XXH128_hash_t is a structure, this capability is not offered by the language. - * Note: For better performance, these functions can be inlined using XXH_INLINE_ALL */ - -/*! - * @brief Check equality of two XXH128_hash_t values - * - * @param h1 The 128-bit hash value. - * @param h2 Another 128-bit hash value. - * - * @return `1` if `h1` and `h2` are equal. - * @return `0` if they are not. - */ -XXH_PUBLIC_API XXH_PUREF int XXH128_isEqual(XXH128_hash_t h1, XXH128_hash_t h2); - -/*! - * @brief Compares two @ref XXH128_hash_t - * - * This comparator is compatible with stdlib's `qsort()`/`bsearch()`. - * - * @param h128_1 Left-hand side value - * @param h128_2 Right-hand side value - * - * @return >0 if @p h128_1 > @p h128_2 - * @return =0 if @p h128_1 == @p h128_2 - * @return <0 if @p h128_1 < @p h128_2 - */ -XXH_PUBLIC_API XXH_PUREF int XXH128_cmp(XXH_NOESCAPE const void* h128_1, XXH_NOESCAPE const void* h128_2); - - -/******* Canonical representation *******/ -typedef struct { unsigned char digest[sizeof(XXH128_hash_t)]; } XXH128_canonical_t; - - -/*! - * @brief Converts an @ref XXH128_hash_t to a big endian @ref XXH128_canonical_t. - * - * @param dst The @ref XXH128_canonical_t pointer to be stored to. - * @param hash The @ref XXH128_hash_t to be converted. - * - * @pre - * @p dst must not be `NULL`. - * @see @ref canonical_representation_example "Canonical Representation Example" - */ -XXH_PUBLIC_API void XXH128_canonicalFromHash(XXH_NOESCAPE XXH128_canonical_t* dst, XXH128_hash_t hash); - -/*! - * @brief Converts an @ref XXH128_canonical_t to a native @ref XXH128_hash_t. - * - * @param src The @ref XXH128_canonical_t to convert. - * - * @pre - * @p src must not be `NULL`. - * - * @return The converted hash. - * @see @ref canonical_representation_example "Canonical Representation Example" - */ -XXH_PUBLIC_API XXH_PUREF XXH128_hash_t XXH128_hashFromCanonical(XXH_NOESCAPE const XXH128_canonical_t* src); - - -#endif /* !XXH_NO_XXH3 */ -#endif /* XXH_NO_LONG_LONG */ - -/*! - * @} - */ -#endif /* XXHASH_H_5627135585666179 */ - - - -#if defined(XXH_STATIC_LINKING_ONLY) && !defined(XXHASH_H_STATIC_13879238742) -#define XXHASH_H_STATIC_13879238742 -/* **************************************************************************** - * This section contains declarations which are not guaranteed to remain stable. - * They may change in future versions, becoming incompatible with a different - * version of the library. - * These declarations should only be used with static linking. - * Never use them in association with dynamic linking! - ***************************************************************************** */ - -/* - * These definitions are only present to allow static allocation - * of XXH states, on stack or in a struct, for example. - * Never **ever** access their members directly. - */ - -/*! - * @internal - * @brief Structure for XXH32 streaming API. - * - * @note This is only defined when @ref XXH_STATIC_LINKING_ONLY, - * @ref XXH_INLINE_ALL, or @ref XXH_IMPLEMENTATION is defined. Otherwise it is - * an opaque type. This allows fields to safely be changed. - * - * Typedef'd to @ref XXH32_state_t. - * Do not access the members of this struct directly. - * @see XXH64_state_s, XXH3_state_s - */ -struct XXH32_state_s { - XXH32_hash_t total_len_32; /*!< Total length hashed, modulo 2^32 */ - XXH32_hash_t large_len; /*!< Whether the hash is >= 16 (handles @ref total_len_32 overflow) */ - XXH32_hash_t v[4]; /*!< Accumulator lanes */ - XXH32_hash_t mem32[4]; /*!< Internal buffer for partial reads. Treated as unsigned char[16]. */ - XXH32_hash_t memsize; /*!< Amount of data in @ref mem32 */ - XXH32_hash_t reserved; /*!< Reserved field. Do not read nor write to it. */ -}; /* typedef'd to XXH32_state_t */ - - -#ifndef XXH_NO_LONG_LONG /* defined when there is no 64-bit support */ - -/*! - * @internal - * @brief Structure for XXH64 streaming API. - * - * @note This is only defined when @ref XXH_STATIC_LINKING_ONLY, - * @ref XXH_INLINE_ALL, or @ref XXH_IMPLEMENTATION is defined. Otherwise it is - * an opaque type. This allows fields to safely be changed. - * - * Typedef'd to @ref XXH64_state_t. - * Do not access the members of this struct directly. - * @see XXH32_state_s, XXH3_state_s - */ -struct XXH64_state_s { - XXH64_hash_t total_len; /*!< Total length hashed. This is always 64-bit. */ - XXH64_hash_t v[4]; /*!< Accumulator lanes */ - XXH64_hash_t mem64[4]; /*!< Internal buffer for partial reads. Treated as unsigned char[32]. */ - XXH32_hash_t memsize; /*!< Amount of data in @ref mem64 */ - XXH32_hash_t reserved32; /*!< Reserved field, needed for padding anyways*/ - XXH64_hash_t reserved64; /*!< Reserved field. Do not read or write to it. */ -}; /* typedef'd to XXH64_state_t */ - -#ifndef XXH_NO_XXH3 - -/* Windows SDK under 10.0.22000 is missing stdalign.h so we add a check - before allowing the windows compiler to use the C11 form. - Reference: https://github.com/Cyan4973/xxHash/issues/955 */ -#if defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201112L) \ - && (defined(_MSC_VER) && (_MSC_VER >= 1000) || !defined(_MSC_VER)) /* >= C11 */ -# include -# define XXH_ALIGN(n) alignas(n) -#elif defined(__cplusplus) && (__cplusplus >= 201103L) /* >= C++11 */ -/* In C++ alignas() is a keyword */ -# define XXH_ALIGN(n) alignas(n) -#elif defined(__GNUC__) -# define XXH_ALIGN(n) __attribute__ ((aligned(n))) -#elif defined(_MSC_VER) -# define XXH_ALIGN(n) __declspec(align(n)) -#else -# define XXH_ALIGN(n) /* disabled */ -#endif - -/* Old GCC versions only accept the attribute after the type in structures. */ -#if !(defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201112L)) /* C11+ */ \ - && ! (defined(__cplusplus) && (__cplusplus >= 201103L)) /* >= C++11 */ \ - && defined(__GNUC__) -# define XXH_ALIGN_MEMBER(align, type) type XXH_ALIGN(align) -#else -# define XXH_ALIGN_MEMBER(align, type) XXH_ALIGN(align) type -#endif - -/*! - * @brief The size of the internal XXH3 buffer. - * - * This is the optimal update size for incremental hashing. - * - * @see XXH3_64b_update(), XXH3_128b_update(). - */ -#define XXH3_INTERNALBUFFER_SIZE 256 - -/*! - * @internal - * @brief Default size of the secret buffer (and @ref XXH3_kSecret). - * - * This is the size used in @ref XXH3_kSecret and the seeded functions. - * - * Not to be confused with @ref XXH3_SECRET_SIZE_MIN. - */ -#define XXH3_SECRET_DEFAULT_SIZE 192 - -/*! - * @internal - * @brief Structure for XXH3 streaming API. - * - * @note This is only defined when @ref XXH_STATIC_LINKING_ONLY, - * @ref XXH_INLINE_ALL, or @ref XXH_IMPLEMENTATION is defined. - * Otherwise it is an opaque type. - * Never use this definition in combination with dynamic library. - * This allows fields to safely be changed in the future. - * - * @note ** This structure has a strict alignment requirement of 64 bytes!! ** - * Do not allocate this with `malloc()` or `new`, - * it will not be sufficiently aligned. - * Use @ref XXH3_createState() and @ref XXH3_freeState(), or stack allocation. - * - * Typedef'd to @ref XXH3_state_t. - * Do never access the members of this struct directly. - * - * @see XXH3_INITSTATE() for stack initialization. - * @see XXH3_createState(), XXH3_freeState(). - * @see XXH32_state_s, XXH64_state_s - */ -struct XXH3_state_s { - XXH_ALIGN_MEMBER(64, XXH64_hash_t acc[8]); - /*!< The 8 accumulators. See @ref XXH32_state_s::v and @ref XXH64_state_s::v */ - XXH_ALIGN_MEMBER(64, unsigned char customSecret[XXH3_SECRET_DEFAULT_SIZE]); - /*!< Used to store a custom secret generated from a seed. */ - XXH_ALIGN_MEMBER(64, unsigned char buffer[XXH3_INTERNALBUFFER_SIZE]); - /*!< The internal buffer. @see XXH32_state_s::mem32 */ - XXH32_hash_t bufferedSize; - /*!< The amount of memory in @ref buffer, @see XXH32_state_s::memsize */ - XXH32_hash_t useSeed; - /*!< Reserved field. Needed for padding on 64-bit. */ - size_t nbStripesSoFar; - /*!< Number or stripes processed. */ - XXH64_hash_t totalLen; - /*!< Total length hashed. 64-bit even on 32-bit targets. */ - size_t nbStripesPerBlock; - /*!< Number of stripes per block. */ - size_t secretLimit; - /*!< Size of @ref customSecret or @ref extSecret */ - XXH64_hash_t seed; - /*!< Seed for _withSeed variants. Must be zero otherwise, @see XXH3_INITSTATE() */ - XXH64_hash_t reserved64; - /*!< Reserved field. */ - const unsigned char* extSecret; - /*!< Reference to an external secret for the _withSecret variants, NULL - * for other variants. */ - /* note: there may be some padding at the end due to alignment on 64 bytes */ -}; /* typedef'd to XXH3_state_t */ - -#undef XXH_ALIGN_MEMBER - -/*! - * @brief Initializes a stack-allocated `XXH3_state_s`. - * - * When the @ref XXH3_state_t structure is merely emplaced on stack, - * it should be initialized with XXH3_INITSTATE() or a memset() - * in case its first reset uses XXH3_NNbits_reset_withSeed(). - * This init can be omitted if the first reset uses default or _withSecret mode. - * This operation isn't necessary when the state is created with XXH3_createState(). - * Note that this doesn't prepare the state for a streaming operation, - * it's still necessary to use XXH3_NNbits_reset*() afterwards. - */ -#define XXH3_INITSTATE(XXH3_state_ptr) \ - do { \ - XXH3_state_t* tmp_xxh3_state_ptr = (XXH3_state_ptr); \ - tmp_xxh3_state_ptr->seed = 0; \ - tmp_xxh3_state_ptr->extSecret = NULL; \ - } while(0) - - -/*! - * @brief Calculates the 128-bit hash of @p data using XXH3. - * - * @param data The block of data to be hashed, at least @p len bytes in size. - * @param len The length of @p data, in bytes. - * @param seed The 64-bit seed to alter the hash's output predictably. - * - * @pre - * The memory between @p data and @p data + @p len must be valid, - * readable, contiguous memory. However, if @p len is `0`, @p data may be - * `NULL`. In C++, this also must be *TriviallyCopyable*. - * - * @return The calculated 128-bit XXH3 value. - * - * @see @ref single_shot_example "Single Shot Example" for an example. - */ -XXH_PUBLIC_API XXH_PUREF XXH128_hash_t XXH128(XXH_NOESCAPE const void* data, size_t len, XXH64_hash_t seed); - - -/* === Experimental API === */ -/* Symbols defined below must be considered tied to a specific library version. */ - -/*! - * @brief Derive a high-entropy secret from any user-defined content, named customSeed. - * - * @param secretBuffer A writable buffer for derived high-entropy secret data. - * @param secretSize Size of secretBuffer, in bytes. Must be >= XXH3_SECRET_SIZE_MIN. - * @param customSeed A user-defined content. - * @param customSeedSize Size of customSeed, in bytes. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * The generated secret can be used in combination with `*_withSecret()` functions. - * The `_withSecret()` variants are useful to provide a higher level of protection - * than 64-bit seed, as it becomes much more difficult for an external actor to - * guess how to impact the calculation logic. - * - * The function accepts as input a custom seed of any length and any content, - * and derives from it a high-entropy secret of length @p secretSize into an - * already allocated buffer @p secretBuffer. - * - * The generated secret can then be used with any `*_withSecret()` variant. - * The functions @ref XXH3_128bits_withSecret(), @ref XXH3_64bits_withSecret(), - * @ref XXH3_128bits_reset_withSecret() and @ref XXH3_64bits_reset_withSecret() - * are part of this list. They all accept a `secret` parameter - * which must be large enough for implementation reasons (>= @ref XXH3_SECRET_SIZE_MIN) - * _and_ feature very high entropy (consist of random-looking bytes). - * These conditions can be a high bar to meet, so @ref XXH3_generateSecret() can - * be employed to ensure proper quality. - * - * @p customSeed can be anything. It can have any size, even small ones, - * and its content can be anything, even "poor entropy" sources such as a bunch - * of zeroes. The resulting `secret` will nonetheless provide all required qualities. - * - * @pre - * - @p secretSize must be >= @ref XXH3_SECRET_SIZE_MIN - * - When @p customSeedSize > 0, supplying NULL as customSeed is undefined behavior. - * - * Example code: - * @code{.c} - * #include - * #include - * #include - * #define XXH_STATIC_LINKING_ONLY // expose unstable API - * #include "xxhash.h" - * // Hashes argv[2] using the entropy from argv[1]. - * int main(int argc, char* argv[]) - * { - * char secret[XXH3_SECRET_SIZE_MIN]; - * if (argv != 3) { return 1; } - * XXH3_generateSecret(secret, sizeof(secret), argv[1], strlen(argv[1])); - * XXH64_hash_t h = XXH3_64bits_withSecret( - * argv[2], strlen(argv[2]), - * secret, sizeof(secret) - * ); - * printf("%016llx\n", (unsigned long long) h); - * } - * @endcode - */ -XXH_PUBLIC_API XXH_errorcode XXH3_generateSecret(XXH_NOESCAPE void* secretBuffer, size_t secretSize, XXH_NOESCAPE const void* customSeed, size_t customSeedSize); - -/*! - * @brief Generate the same secret as the _withSeed() variants. - * - * @param secretBuffer A writable buffer of @ref XXH3_SECRET_DEFAULT_SIZE bytes - * @param seed The 64-bit seed to alter the hash result predictably. - * - * The generated secret can be used in combination with - *`*_withSecret()` and `_withSecretandSeed()` variants. - * - * Example C++ `std::string` hash class: - * @code{.cpp} - * #include - * #define XXH_STATIC_LINKING_ONLY // expose unstable API - * #include "xxhash.h" - * // Slow, seeds each time - * class HashSlow { - * XXH64_hash_t seed; - * public: - * HashSlow(XXH64_hash_t s) : seed{s} {} - * size_t operator()(const std::string& x) const { - * return size_t{XXH3_64bits_withSeed(x.c_str(), x.length(), seed)}; - * } - * }; - * // Fast, caches the seeded secret for future uses. - * class HashFast { - * unsigned char secret[XXH3_SECRET_DEFAULT_SIZE]; - * public: - * HashFast(XXH64_hash_t s) { - * XXH3_generateSecret_fromSeed(secret, seed); - * } - * size_t operator()(const std::string& x) const { - * return size_t{ - * XXH3_64bits_withSecret(x.c_str(), x.length(), secret, sizeof(secret)) - * }; - * } - * }; - * @endcode - */ -XXH_PUBLIC_API void XXH3_generateSecret_fromSeed(XXH_NOESCAPE void* secretBuffer, XXH64_hash_t seed); - -/*! - * @brief Maximum size of "short" key in bytes. - */ -#define XXH3_MIDSIZE_MAX 240 - -/*! - * @brief Calculates 64/128-bit seeded variant of XXH3 hash of @p data. - * - * @param data The block of data to be hashed, at least @p len bytes in size. - * @param len The length of @p data, in bytes. - * @param secret The secret data. - * @param secretSize The length of @p secret, in bytes. - * @param seed The 64-bit seed to alter the hash result predictably. - * - * These variants generate hash values using either: - * - @p seed for "short" keys (< @ref XXH3_MIDSIZE_MAX = 240 bytes) - * - @p secret for "large" keys (>= @ref XXH3_MIDSIZE_MAX). - * - * This generally benefits speed, compared to `_withSeed()` or `_withSecret()`. - * `_withSeed()` has to generate the secret on the fly for "large" keys. - * It's fast, but can be perceptible for "not so large" keys (< 1 KB). - * `_withSecret()` has to generate the masks on the fly for "small" keys, - * which requires more instructions than _withSeed() variants. - * Therefore, _withSecretandSeed variant combines the best of both worlds. - * - * When @p secret has been generated by XXH3_generateSecret_fromSeed(), - * this variant produces *exactly* the same results as `_withSeed()` variant, - * hence offering only a pure speed benefit on "large" input, - * by skipping the need to regenerate the secret for every large input. - * - * Another usage scenario is to hash the secret to a 64-bit hash value, - * for example with XXH3_64bits(), which then becomes the seed, - * and then employ both the seed and the secret in _withSecretandSeed(). - * On top of speed, an added benefit is that each bit in the secret - * has a 50% chance to swap each bit in the output, via its impact to the seed. - * - * This is not guaranteed when using the secret directly in "small data" scenarios, - * because only portions of the secret are employed for small data. - */ -XXH_PUBLIC_API XXH_PUREF XXH64_hash_t -XXH3_64bits_withSecretandSeed(XXH_NOESCAPE const void* data, size_t len, - XXH_NOESCAPE const void* secret, size_t secretSize, - XXH64_hash_t seed); - -/*! - * @brief Calculates 128-bit seeded variant of XXH3 hash of @p data. - * - * @param data The memory segment to be hashed, at least @p len bytes in size. - * @param length The length of @p data, in bytes. - * @param secret The secret used to alter hash result predictably. - * @param secretSize The length of @p secret, in bytes (must be >= XXH3_SECRET_SIZE_MIN) - * @param seed64 The 64-bit seed to alter the hash result predictably. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @see XXH3_64bits_withSecretandSeed(): contract is the same. - */ -XXH_PUBLIC_API XXH_PUREF XXH128_hash_t -XXH3_128bits_withSecretandSeed(XXH_NOESCAPE const void* input, size_t length, - XXH_NOESCAPE const void* secret, size_t secretSize, - XXH64_hash_t seed64); - -#ifndef XXH_NO_STREAM -/*! - * @brief Resets an @ref XXH3_state_t with secret data to begin a new hash. - * - * @param statePtr A pointer to an @ref XXH3_state_t allocated with @ref XXH3_createState(). - * @param secret The secret data. - * @param secretSize The length of @p secret, in bytes. - * @param seed64 The 64-bit seed to alter the hash result predictably. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @see XXH3_64bits_withSecretandSeed(). Contract is identical. - */ -XXH_PUBLIC_API XXH_errorcode -XXH3_64bits_reset_withSecretandSeed(XXH_NOESCAPE XXH3_state_t* statePtr, - XXH_NOESCAPE const void* secret, size_t secretSize, - XXH64_hash_t seed64); - -/*! - * @brief Resets an @ref XXH3_state_t with secret data to begin a new hash. - * - * @param statePtr A pointer to an @ref XXH3_state_t allocated with @ref XXH3_createState(). - * @param secret The secret data. - * @param secretSize The length of @p secret, in bytes. - * @param seed64 The 64-bit seed to alter the hash result predictably. - * - * @return @ref XXH_OK on success. - * @return @ref XXH_ERROR on failure. - * - * @see XXH3_64bits_withSecretandSeed(). Contract is identical. - * - * Note: there was a bug in an earlier version of this function (<= v0.8.2) - * that would make it generate an incorrect hash value - * when @p seed == 0 and @p length < XXH3_MIDSIZE_MAX - * and @p secret is different from XXH3_generateSecret_fromSeed(). - * As stated in the contract, the correct hash result must be - * the same as XXH3_128bits_withSeed() when @p length <= XXH3_MIDSIZE_MAX. - * Results generated by this older version are wrong, hence not comparable. - */ -XXH_PUBLIC_API XXH_errorcode -XXH3_128bits_reset_withSecretandSeed(XXH_NOESCAPE XXH3_state_t* statePtr, - XXH_NOESCAPE const void* secret, size_t secretSize, - XXH64_hash_t seed64); - -#endif /* !XXH_NO_STREAM */ - -#endif /* !XXH_NO_XXH3 */ -#endif /* XXH_NO_LONG_LONG */ -#if defined(XXH_INLINE_ALL) || defined(XXH_PRIVATE_API) -# define XXH_IMPLEMENTATION -#endif - -#endif /* defined(XXH_STATIC_LINKING_ONLY) && !defined(XXHASH_H_STATIC_13879238742) */ - - -/* ======================================================================== */ -/* ======================================================================== */ -/* ======================================================================== */ - - -/*-********************************************************************** - * xxHash implementation - *-********************************************************************** - * xxHash's implementation used to be hosted inside xxhash.c. - * - * However, inlining requires implementation to be visible to the compiler, - * hence be included alongside the header. - * Previously, implementation was hosted inside xxhash.c, - * which was then #included when inlining was activated. - * This construction created issues with a few build and install systems, - * as it required xxhash.c to be stored in /include directory. - * - * xxHash implementation is now directly integrated within xxhash.h. - * As a consequence, xxhash.c is no longer needed in /include. - * - * xxhash.c is still available and is still useful. - * In a "normal" setup, when xxhash is not inlined, - * xxhash.h only exposes the prototypes and public symbols, - * while xxhash.c can be built into an object file xxhash.o - * which can then be linked into the final binary. - ************************************************************************/ - -#if ( defined(XXH_INLINE_ALL) || defined(XXH_PRIVATE_API) \ - || defined(XXH_IMPLEMENTATION) ) && !defined(XXH_IMPLEM_13a8737387) -# define XXH_IMPLEM_13a8737387 - -/* ************************************* -* Tuning parameters -***************************************/ - -/*! - * @defgroup tuning Tuning parameters - * @{ - * - * Various macros to control xxHash's behavior. - */ -#ifdef XXH_DOXYGEN -/*! - * @brief Define this to disable 64-bit code. - * - * Useful if only using the @ref XXH32_family and you have a strict C90 compiler. - */ -# define XXH_NO_LONG_LONG -# undef XXH_NO_LONG_LONG /* don't actually */ -/*! - * @brief Controls how unaligned memory is accessed. - * - * By default, access to unaligned memory is controlled by `memcpy()`, which is - * safe and portable. - * - * Unfortunately, on some target/compiler combinations, the generated assembly - * is sub-optimal. - * - * The below switch allow selection of a different access method - * in the search for improved performance. - * - * @par Possible options: - * - * - `XXH_FORCE_MEMORY_ACCESS=0` (default): `memcpy` - * @par - * Use `memcpy()`. Safe and portable. Note that most modern compilers will - * eliminate the function call and treat it as an unaligned access. - * - * - `XXH_FORCE_MEMORY_ACCESS=1`: `__attribute__((aligned(1)))` - * @par - * Depends on compiler extensions and is therefore not portable. - * This method is safe _if_ your compiler supports it, - * and *generally* as fast or faster than `memcpy`. - * - * - `XXH_FORCE_MEMORY_ACCESS=2`: Direct cast - * @par - * Casts directly and dereferences. This method doesn't depend on the - * compiler, but it violates the C standard as it directly dereferences an - * unaligned pointer. It can generate buggy code on targets which do not - * support unaligned memory accesses, but in some circumstances, it's the - * only known way to get the most performance. - * - * - `XXH_FORCE_MEMORY_ACCESS=3`: Byteshift - * @par - * Also portable. This can generate the best code on old compilers which don't - * inline small `memcpy()` calls, and it might also be faster on big-endian - * systems which lack a native byteswap instruction. However, some compilers - * will emit literal byteshifts even if the target supports unaligned access. - * - * - * @warning - * Methods 1 and 2 rely on implementation-defined behavior. Use these with - * care, as what works on one compiler/platform/optimization level may cause - * another to read garbage data or even crash. - * - * See https://fastcompression.blogspot.com/2015/08/accessing-unaligned-memory.html for details. - * - * Prefer these methods in priority order (0 > 3 > 1 > 2) - */ -# define XXH_FORCE_MEMORY_ACCESS 0 - -/*! - * @def XXH_SIZE_OPT - * @brief Controls how much xxHash optimizes for size. - * - * xxHash, when compiled, tends to result in a rather large binary size. This - * is mostly due to heavy usage to forced inlining and constant folding of the - * @ref XXH3_family to increase performance. - * - * However, some developers prefer size over speed. This option can - * significantly reduce the size of the generated code. When using the `-Os` - * or `-Oz` options on GCC or Clang, this is defined to 1 by default, - * otherwise it is defined to 0. - * - * Most of these size optimizations can be controlled manually. - * - * This is a number from 0-2. - * - `XXH_SIZE_OPT` == 0: Default. xxHash makes no size optimizations. Speed - * comes first. - * - `XXH_SIZE_OPT` == 1: Default for `-Os` and `-Oz`. xxHash is more - * conservative and disables hacks that increase code size. It implies the - * options @ref XXH_NO_INLINE_HINTS == 1, @ref XXH_FORCE_ALIGN_CHECK == 0, - * and @ref XXH3_NEON_LANES == 8 if they are not already defined. - * - `XXH_SIZE_OPT` == 2: xxHash tries to make itself as small as possible. - * Performance may cry. For example, the single shot functions just use the - * streaming API. - */ -# define XXH_SIZE_OPT 0 - -/*! - * @def XXH_FORCE_ALIGN_CHECK - * @brief If defined to non-zero, adds a special path for aligned inputs (XXH32() - * and XXH64() only). - * - * This is an important performance trick for architectures without decent - * unaligned memory access performance. - * - * It checks for input alignment, and when conditions are met, uses a "fast - * path" employing direct 32-bit/64-bit reads, resulting in _dramatically - * faster_ read speed. - * - * The check costs one initial branch per hash, which is generally negligible, - * but not zero. - * - * Moreover, it's not useful to generate an additional code path if memory - * access uses the same instruction for both aligned and unaligned - * addresses (e.g. x86 and aarch64). - * - * In these cases, the alignment check can be removed by setting this macro to 0. - * Then the code will always use unaligned memory access. - * Align check is automatically disabled on x86, x64, ARM64, and some ARM chips - * which are platforms known to offer good unaligned memory accesses performance. - * - * It is also disabled by default when @ref XXH_SIZE_OPT >= 1. - * - * This option does not affect XXH3 (only XXH32 and XXH64). - */ -# define XXH_FORCE_ALIGN_CHECK 0 - -/*! - * @def XXH_NO_INLINE_HINTS - * @brief When non-zero, sets all functions to `static`. - * - * By default, xxHash tries to force the compiler to inline almost all internal - * functions. - * - * This can usually improve performance due to reduced jumping and improved - * constant folding, but significantly increases the size of the binary which - * might not be favorable. - * - * Additionally, sometimes the forced inlining can be detrimental to performance, - * depending on the architecture. - * - * XXH_NO_INLINE_HINTS marks all internal functions as static, giving the - * compiler full control on whether to inline or not. - * - * When not optimizing (-O0), using `-fno-inline` with GCC or Clang, or if - * @ref XXH_SIZE_OPT >= 1, this will automatically be defined. - */ -# define XXH_NO_INLINE_HINTS 0 - -/*! - * @def XXH3_INLINE_SECRET - * @brief Determines whether to inline the XXH3 withSecret code. - * - * When the secret size is known, the compiler can improve the performance - * of XXH3_64bits_withSecret() and XXH3_128bits_withSecret(). - * - * However, if the secret size is not known, it doesn't have any benefit. This - * happens when xxHash is compiled into a global symbol. Therefore, if - * @ref XXH_INLINE_ALL is *not* defined, this will be defined to 0. - * - * Additionally, this defaults to 0 on GCC 12+, which has an issue with function pointers - * that are *sometimes* force inline on -Og, and it is impossible to automatically - * detect this optimization level. - */ -# define XXH3_INLINE_SECRET 0 - -/*! - * @def XXH32_ENDJMP - * @brief Whether to use a jump for `XXH32_finalize`. - * - * For performance, `XXH32_finalize` uses multiple branches in the finalizer. - * This is generally preferable for performance, - * but depending on exact architecture, a jmp may be preferable. - * - * This setting is only possibly making a difference for very small inputs. - */ -# define XXH32_ENDJMP 0 - -/*! - * @internal - * @brief Redefines old internal names. - * - * For compatibility with code that uses xxHash's internals before the names - * were changed to improve namespacing. There is no other reason to use this. - */ -# define XXH_OLD_NAMES -# undef XXH_OLD_NAMES /* don't actually use, it is ugly. */ - -/*! - * @def XXH_NO_STREAM - * @brief Disables the streaming API. - * - * When xxHash is not inlined and the streaming functions are not used, disabling - * the streaming functions can improve code size significantly, especially with - * the @ref XXH3_family which tends to make constant folded copies of itself. - */ -# define XXH_NO_STREAM -# undef XXH_NO_STREAM /* don't actually */ -#endif /* XXH_DOXYGEN */ -/*! - * @} - */ - -#ifndef XXH_FORCE_MEMORY_ACCESS /* can be defined externally, on command line for example */ - /* prefer __packed__ structures (method 1) for GCC - * < ARMv7 with unaligned access (e.g. Raspbian armhf) still uses byte shifting, so we use memcpy - * which for some reason does unaligned loads. */ -# if defined(__GNUC__) && !(defined(__ARM_ARCH) && __ARM_ARCH < 7 && defined(__ARM_FEATURE_UNALIGNED)) -# define XXH_FORCE_MEMORY_ACCESS 1 -# endif -#endif - -#ifndef XXH_SIZE_OPT - /* default to 1 for -Os or -Oz */ -# if (defined(__GNUC__) || defined(__clang__)) && defined(__OPTIMIZE_SIZE__) -# define XXH_SIZE_OPT 1 -# else -# define XXH_SIZE_OPT 0 -# endif -#endif - -#ifndef XXH_FORCE_ALIGN_CHECK /* can be defined externally */ - /* don't check on sizeopt, x86, aarch64, or arm when unaligned access is available */ -# if XXH_SIZE_OPT >= 1 || \ - defined(__i386) || defined(__x86_64__) || defined(__aarch64__) || defined(__ARM_FEATURE_UNALIGNED) \ - || defined(_M_IX86) || defined(_M_X64) || defined(_M_ARM64) || defined(_M_ARM) /* visual */ -# define XXH_FORCE_ALIGN_CHECK 0 -# else -# define XXH_FORCE_ALIGN_CHECK 1 -# endif -#endif - -#ifndef XXH_NO_INLINE_HINTS -# if XXH_SIZE_OPT >= 1 || defined(__NO_INLINE__) /* -O0, -fno-inline */ -# define XXH_NO_INLINE_HINTS 1 -# else -# define XXH_NO_INLINE_HINTS 0 -# endif -#endif - -#ifndef XXH3_INLINE_SECRET -# if (defined(__GNUC__) && !defined(__clang__) && __GNUC__ >= 12) \ - || !defined(XXH_INLINE_ALL) -# define XXH3_INLINE_SECRET 0 -# else -# define XXH3_INLINE_SECRET 1 -# endif -#endif - -#ifndef XXH32_ENDJMP -/* generally preferable for performance */ -# define XXH32_ENDJMP 0 -#endif - -/*! - * @defgroup impl Implementation - * @{ - */ - - -/* ************************************* -* Includes & Memory related functions -***************************************/ -#if defined(XXH_NO_STREAM) -/* nothing */ -#elif defined(XXH_NO_STDLIB) - -/* When requesting to disable any mention of stdlib, - * the library loses the ability to invoked malloc / free. - * In practice, it means that functions like `XXH*_createState()` - * will always fail, and return NULL. - * This flag is useful in situations where - * xxhash.h is integrated into some kernel, embedded or limited environment - * without access to dynamic allocation. - */ - -static XXH_CONSTF void* XXH_malloc(size_t s) { (void)s; return NULL; } -static void XXH_free(void* p) { (void)p; } - -#else - -/* - * Modify the local functions below should you wish to use - * different memory routines for malloc() and free() - */ -#include - -/*! - * @internal - * @brief Modify this function to use a different routine than malloc(). - */ -static XXH_MALLOCF void* XXH_malloc(size_t s) { return malloc(s); } - -/*! - * @internal - * @brief Modify this function to use a different routine than free(). - */ -static void XXH_free(void* p) { free(p); } - -#endif /* XXH_NO_STDLIB */ - -#include - -/*! - * @internal - * @brief Modify this function to use a different routine than memcpy(). - */ -static void* XXH_memcpy(void* dest, const void* src, size_t size) -{ - return memcpy(dest,src,size); -} - -#include /* ULLONG_MAX */ - - -/* ************************************* -* Compiler Specific Options -***************************************/ -#ifdef _MSC_VER /* Visual Studio warning fix */ -# pragma warning(disable : 4127) /* disable: C4127: conditional expression is constant */ -#endif - -#if XXH_NO_INLINE_HINTS /* disable inlining hints */ -# if defined(__GNUC__) || defined(__clang__) -# define XXH_FORCE_INLINE static __attribute__((__unused__)) -# else -# define XXH_FORCE_INLINE static -# endif -# define XXH_NO_INLINE static -/* enable inlining hints */ -#elif defined(__GNUC__) || defined(__clang__) -# define XXH_FORCE_INLINE static __inline__ __attribute__((__always_inline__, __unused__)) -# define XXH_NO_INLINE static __attribute__((__noinline__)) -#elif defined(_MSC_VER) /* Visual Studio */ -# define XXH_FORCE_INLINE static __forceinline -# define XXH_NO_INLINE static __declspec(noinline) -#elif defined (__cplusplus) \ - || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L)) /* C99 */ -# define XXH_FORCE_INLINE static inline -# define XXH_NO_INLINE static -#else -# define XXH_FORCE_INLINE static -# define XXH_NO_INLINE static -#endif - -#if XXH3_INLINE_SECRET -# define XXH3_WITH_SECRET_INLINE XXH_FORCE_INLINE -#else -# define XXH3_WITH_SECRET_INLINE XXH_NO_INLINE -#endif - - -/* ************************************* -* Debug -***************************************/ -/*! - * @ingroup tuning - * @def XXH_DEBUGLEVEL - * @brief Sets the debugging level. - * - * XXH_DEBUGLEVEL is expected to be defined externally, typically via the - * compiler's command line options. The value must be a number. - */ -#ifndef XXH_DEBUGLEVEL -# ifdef DEBUGLEVEL /* backwards compat */ -# define XXH_DEBUGLEVEL DEBUGLEVEL -# else -# define XXH_DEBUGLEVEL 0 -# endif -#endif - -#if (XXH_DEBUGLEVEL>=1) -# include /* note: can still be disabled with NDEBUG */ -# define XXH_ASSERT(c) assert(c) -#else -# if defined(__INTEL_COMPILER) -# define XXH_ASSERT(c) XXH_ASSUME((unsigned char) (c)) -# else -# define XXH_ASSERT(c) XXH_ASSUME(c) -# endif -#endif - -/* note: use after variable declarations */ -#ifndef XXH_STATIC_ASSERT -# if defined(__STDC_VERSION__) && (__STDC_VERSION__ >= 201112L) /* C11 */ -# define XXH_STATIC_ASSERT_WITH_MESSAGE(c,m) do { _Static_assert((c),m); } while(0) -# elif defined(__cplusplus) && (__cplusplus >= 201103L) /* C++11 */ -# define XXH_STATIC_ASSERT_WITH_MESSAGE(c,m) do { static_assert((c),m); } while(0) -# else -# define XXH_STATIC_ASSERT_WITH_MESSAGE(c,m) do { struct xxh_sa { char x[(c) ? 1 : -1]; }; } while(0) -# endif -# define XXH_STATIC_ASSERT(c) XXH_STATIC_ASSERT_WITH_MESSAGE((c),#c) -#endif - -/*! - * @internal - * @def XXH_COMPILER_GUARD(var) - * @brief Used to prevent unwanted optimizations for @p var. - * - * It uses an empty GCC inline assembly statement with a register constraint - * which forces @p var into a general purpose register (eg eax, ebx, ecx - * on x86) and marks it as modified. - * - * This is used in a few places to avoid unwanted autovectorization (e.g. - * XXH32_round()). All vectorization we want is explicit via intrinsics, - * and _usually_ isn't wanted elsewhere. - * - * We also use it to prevent unwanted constant folding for AArch64 in - * XXH3_initCustomSecret_scalar(). - */ -#if defined(__GNUC__) || defined(__clang__) -# define XXH_COMPILER_GUARD(var) __asm__("" : "+r" (var)) -#else -# define XXH_COMPILER_GUARD(var) ((void)0) -#endif - -/* Specifically for NEON vectors which use the "w" constraint, on - * Clang. */ -#if defined(__clang__) && defined(__ARM_ARCH) && !defined(__wasm__) -# define XXH_COMPILER_GUARD_CLANG_NEON(var) __asm__("" : "+w" (var)) -#else -# define XXH_COMPILER_GUARD_CLANG_NEON(var) ((void)0) -#endif - -/* ************************************* -* Basic Types -***************************************/ -#if !defined (__VMS) \ - && (defined (__cplusplus) \ - || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) ) -# ifdef _AIX -# include -# else -# include -# endif - typedef uint8_t xxh_u8; -#else - typedef unsigned char xxh_u8; -#endif -typedef XXH32_hash_t xxh_u32; - -#ifdef XXH_OLD_NAMES -# warning "XXH_OLD_NAMES is planned to be removed starting v0.9. If the program depends on it, consider moving away from it by employing newer type names directly" -# define BYTE xxh_u8 -# define U8 xxh_u8 -# define U32 xxh_u32 -#endif - -/* *** Memory access *** */ - -/*! - * @internal - * @fn xxh_u32 XXH_read32(const void* ptr) - * @brief Reads an unaligned 32-bit integer from @p ptr in native endianness. - * - * Affected by @ref XXH_FORCE_MEMORY_ACCESS. - * - * @param ptr The pointer to read from. - * @return The 32-bit native endian integer from the bytes at @p ptr. - */ - -/*! - * @internal - * @fn xxh_u32 XXH_readLE32(const void* ptr) - * @brief Reads an unaligned 32-bit little endian integer from @p ptr. - * - * Affected by @ref XXH_FORCE_MEMORY_ACCESS. - * - * @param ptr The pointer to read from. - * @return The 32-bit little endian integer from the bytes at @p ptr. - */ - -/*! - * @internal - * @fn xxh_u32 XXH_readBE32(const void* ptr) - * @brief Reads an unaligned 32-bit big endian integer from @p ptr. - * - * Affected by @ref XXH_FORCE_MEMORY_ACCESS. - * - * @param ptr The pointer to read from. - * @return The 32-bit big endian integer from the bytes at @p ptr. - */ - -/*! - * @internal - * @fn xxh_u32 XXH_readLE32_align(const void* ptr, XXH_alignment align) - * @brief Like @ref XXH_readLE32(), but has an option for aligned reads. - * - * Affected by @ref XXH_FORCE_MEMORY_ACCESS. - * Note that when @ref XXH_FORCE_ALIGN_CHECK == 0, the @p align parameter is - * always @ref XXH_alignment::XXH_unaligned. - * - * @param ptr The pointer to read from. - * @param align Whether @p ptr is aligned. - * @pre - * If @p align == @ref XXH_alignment::XXH_aligned, @p ptr must be 4 byte - * aligned. - * @return The 32-bit little endian integer from the bytes at @p ptr. - */ - -#if (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==3)) -/* - * Manual byteshift. Best for old compilers which don't inline memcpy. - * We actually directly use XXH_readLE32 and XXH_readBE32. - */ -#elif (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==2)) - -/* - * Force direct memory access. Only works on CPU which support unaligned memory - * access in hardware. - */ -static xxh_u32 XXH_read32(const void* memPtr) { return *(const xxh_u32*) memPtr; } - -#elif (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==1)) - -/* - * __attribute__((aligned(1))) is supported by gcc and clang. Originally the - * documentation claimed that it only increased the alignment, but actually it - * can decrease it on gcc, clang, and icc: - * https://gcc.gnu.org/bugzilla/show_bug.cgi?id=69502, - * https://gcc.godbolt.org/z/xYez1j67Y. - */ -#ifdef XXH_OLD_NAMES -typedef union { xxh_u32 u32; } __attribute__((__packed__)) unalign; -#endif -static xxh_u32 XXH_read32(const void* ptr) -{ - typedef __attribute__((__aligned__(1))) xxh_u32 xxh_unalign32; - return *((const xxh_unalign32*)ptr); -} - -#else - -/* - * Portable and safe solution. Generally efficient. - * see: https://fastcompression.blogspot.com/2015/08/accessing-unaligned-memory.html - */ -static xxh_u32 XXH_read32(const void* memPtr) -{ - xxh_u32 val; - XXH_memcpy(&val, memPtr, sizeof(val)); - return val; -} - -#endif /* XXH_FORCE_DIRECT_MEMORY_ACCESS */ - - -/* *** Endianness *** */ - -/*! - * @ingroup tuning - * @def XXH_CPU_LITTLE_ENDIAN - * @brief Whether the target is little endian. - * - * Defined to 1 if the target is little endian, or 0 if it is big endian. - * It can be defined externally, for example on the compiler command line. - * - * If it is not defined, - * a runtime check (which is usually constant folded) is used instead. - * - * @note - * This is not necessarily defined to an integer constant. - * - * @see XXH_isLittleEndian() for the runtime check. - */ -#ifndef XXH_CPU_LITTLE_ENDIAN -/* - * Try to detect endianness automatically, to avoid the nonstandard behavior - * in `XXH_isLittleEndian()` - */ -# if defined(_WIN32) /* Windows is always little endian */ \ - || defined(__LITTLE_ENDIAN__) \ - || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__) -# define XXH_CPU_LITTLE_ENDIAN 1 -# elif defined(__BIG_ENDIAN__) \ - || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) -# define XXH_CPU_LITTLE_ENDIAN 0 -# else -/*! - * @internal - * @brief Runtime check for @ref XXH_CPU_LITTLE_ENDIAN. - * - * Most compilers will constant fold this. - */ -static int XXH_isLittleEndian(void) -{ - /* - * Portable and well-defined behavior. - * Don't use static: it is detrimental to performance. - */ - const union { xxh_u32 u; xxh_u8 c[4]; } one = { 1 }; - return one.c[0]; -} -# define XXH_CPU_LITTLE_ENDIAN XXH_isLittleEndian() -# endif -#endif - - - - -/* **************************************** -* Compiler-specific Functions and Macros -******************************************/ -#define XXH_GCC_VERSION (__GNUC__ * 100 + __GNUC_MINOR__) - -#ifdef __has_builtin -# define XXH_HAS_BUILTIN(x) __has_builtin(x) -#else -# define XXH_HAS_BUILTIN(x) 0 -#endif - - - -/* - * C23 and future versions have standard "unreachable()". - * Once it has been implemented reliably we can add it as an - * additional case: - * - * ``` - * #if defined(__STDC_VERSION__) && (__STDC_VERSION__ >= XXH_C23_VN) - * # include - * # ifdef unreachable - * # define XXH_UNREACHABLE() unreachable() - * # endif - * #endif - * ``` - * - * Note C++23 also has std::unreachable() which can be detected - * as follows: - * ``` - * #if defined(__cpp_lib_unreachable) && (__cpp_lib_unreachable >= 202202L) - * # include - * # define XXH_UNREACHABLE() std::unreachable() - * #endif - * ``` - * NB: `__cpp_lib_unreachable` is defined in the `` header. - * We don't use that as including `` in `extern "C"` blocks - * doesn't work on GCC12 - */ - -#if XXH_HAS_BUILTIN(__builtin_unreachable) -# define XXH_UNREACHABLE() __builtin_unreachable() - -#elif defined(_MSC_VER) -# define XXH_UNREACHABLE() __assume(0) - -#else -# define XXH_UNREACHABLE() -#endif - -#if XXH_HAS_BUILTIN(__builtin_assume) -# define XXH_ASSUME(c) __builtin_assume(c) -#else -# define XXH_ASSUME(c) if (!(c)) { XXH_UNREACHABLE(); } -#endif - -/*! - * @internal - * @def XXH_rotl32(x,r) - * @brief 32-bit rotate left. - * - * @param x The 32-bit integer to be rotated. - * @param r The number of bits to rotate. - * @pre - * @p r > 0 && @p r < 32 - * @note - * @p x and @p r may be evaluated multiple times. - * @return The rotated result. - */ -#if !defined(NO_CLANG_BUILTIN) && XXH_HAS_BUILTIN(__builtin_rotateleft32) \ - && XXH_HAS_BUILTIN(__builtin_rotateleft64) -# define XXH_rotl32 __builtin_rotateleft32 -# define XXH_rotl64 __builtin_rotateleft64 -/* Note: although _rotl exists for minGW (GCC under windows), performance seems poor */ -#elif defined(_MSC_VER) -# define XXH_rotl32(x,r) _rotl(x,r) -# define XXH_rotl64(x,r) _rotl64(x,r) -#else -# define XXH_rotl32(x,r) (((x) << (r)) | ((x) >> (32 - (r)))) -# define XXH_rotl64(x,r) (((x) << (r)) | ((x) >> (64 - (r)))) -#endif - -/*! - * @internal - * @fn xxh_u32 XXH_swap32(xxh_u32 x) - * @brief A 32-bit byteswap. - * - * @param x The 32-bit integer to byteswap. - * @return @p x, byteswapped. - */ -#if defined(_MSC_VER) /* Visual Studio */ -# define XXH_swap32 _byteswap_ulong -#elif XXH_GCC_VERSION >= 403 -# define XXH_swap32 __builtin_bswap32 -#else -static xxh_u32 XXH_swap32 (xxh_u32 x) -{ - return ((x << 24) & 0xff000000 ) | - ((x << 8) & 0x00ff0000 ) | - ((x >> 8) & 0x0000ff00 ) | - ((x >> 24) & 0x000000ff ); -} -#endif - - -/* *************************** -* Memory reads -*****************************/ - -/*! - * @internal - * @brief Enum to indicate whether a pointer is aligned. - */ -typedef enum { - XXH_aligned, /*!< Aligned */ - XXH_unaligned /*!< Possibly unaligned */ -} XXH_alignment; - -/* - * XXH_FORCE_MEMORY_ACCESS==3 is an endian-independent byteshift load. - * - * This is ideal for older compilers which don't inline memcpy. - */ -#if (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==3)) - -XXH_FORCE_INLINE xxh_u32 XXH_readLE32(const void* memPtr) -{ - const xxh_u8* bytePtr = (const xxh_u8 *)memPtr; - return bytePtr[0] - | ((xxh_u32)bytePtr[1] << 8) - | ((xxh_u32)bytePtr[2] << 16) - | ((xxh_u32)bytePtr[3] << 24); -} - -XXH_FORCE_INLINE xxh_u32 XXH_readBE32(const void* memPtr) -{ - const xxh_u8* bytePtr = (const xxh_u8 *)memPtr; - return bytePtr[3] - | ((xxh_u32)bytePtr[2] << 8) - | ((xxh_u32)bytePtr[1] << 16) - | ((xxh_u32)bytePtr[0] << 24); -} - -#else -XXH_FORCE_INLINE xxh_u32 XXH_readLE32(const void* ptr) -{ - return XXH_CPU_LITTLE_ENDIAN ? XXH_read32(ptr) : XXH_swap32(XXH_read32(ptr)); -} - -static xxh_u32 XXH_readBE32(const void* ptr) -{ - return XXH_CPU_LITTLE_ENDIAN ? XXH_swap32(XXH_read32(ptr)) : XXH_read32(ptr); -} -#endif - -XXH_FORCE_INLINE xxh_u32 -XXH_readLE32_align(const void* ptr, XXH_alignment align) -{ - if (align==XXH_unaligned) { - return XXH_readLE32(ptr); - } else { - return XXH_CPU_LITTLE_ENDIAN ? *(const xxh_u32*)ptr : XXH_swap32(*(const xxh_u32*)ptr); - } -} - - -/* ************************************* -* Misc -***************************************/ -/*! @ingroup public */ -XXH_PUBLIC_API unsigned XXH_versionNumber (void) { return XXH_VERSION_NUMBER; } - - -/* ******************************************************************* -* 32-bit hash functions -*********************************************************************/ -/*! - * @} - * @defgroup XXH32_impl XXH32 implementation - * @ingroup impl - * - * Details on the XXH32 implementation. - * @{ - */ - /* #define instead of static const, to be used as initializers */ -#define XXH_PRIME32_1 0x9E3779B1U /*!< 0b10011110001101110111100110110001 */ -#define XXH_PRIME32_2 0x85EBCA77U /*!< 0b10000101111010111100101001110111 */ -#define XXH_PRIME32_3 0xC2B2AE3DU /*!< 0b11000010101100101010111000111101 */ -#define XXH_PRIME32_4 0x27D4EB2FU /*!< 0b00100111110101001110101100101111 */ -#define XXH_PRIME32_5 0x165667B1U /*!< 0b00010110010101100110011110110001 */ - -#ifdef XXH_OLD_NAMES -# define PRIME32_1 XXH_PRIME32_1 -# define PRIME32_2 XXH_PRIME32_2 -# define PRIME32_3 XXH_PRIME32_3 -# define PRIME32_4 XXH_PRIME32_4 -# define PRIME32_5 XXH_PRIME32_5 -#endif - -/*! - * @internal - * @brief Normal stripe processing routine. - * - * This shuffles the bits so that any bit from @p input impacts several bits in - * @p acc. - * - * @param acc The accumulator lane. - * @param input The stripe of input to mix. - * @return The mixed accumulator lane. - */ -static xxh_u32 XXH32_round(xxh_u32 acc, xxh_u32 input) -{ - acc += input * XXH_PRIME32_2; - acc = XXH_rotl32(acc, 13); - acc *= XXH_PRIME32_1; -#if (defined(__SSE4_1__) || defined(__aarch64__) || defined(__wasm_simd128__)) && !defined(XXH_ENABLE_AUTOVECTORIZE) - /* - * UGLY HACK: - * A compiler fence is used to prevent GCC and Clang from - * autovectorizing the XXH32 loop (pragmas and attributes don't work for some - * reason) without globally disabling SSE4.1. - * - * The reason we want to avoid vectorization is because despite working on - * 4 integers at a time, there are multiple factors slowing XXH32 down on - * SSE4: - * - There's a ridiculous amount of lag from pmulld (10 cycles of latency on - * newer chips!) making it slightly slower to multiply four integers at - * once compared to four integers independently. Even when pmulld was - * fastest, Sandy/Ivy Bridge, it is still not worth it to go into SSE - * just to multiply unless doing a long operation. - * - * - Four instructions are required to rotate, - * movqda tmp, v // not required with VEX encoding - * pslld tmp, 13 // tmp <<= 13 - * psrld v, 19 // x >>= 19 - * por v, tmp // x |= tmp - * compared to one for scalar: - * roll v, 13 // reliably fast across the board - * shldl v, v, 13 // Sandy Bridge and later prefer this for some reason - * - * - Instruction level parallelism is actually more beneficial here because - * the SIMD actually serializes this operation: While v1 is rotating, v2 - * can load data, while v3 can multiply. SSE forces them to operate - * together. - * - * This is also enabled on AArch64, as Clang is *very aggressive* in vectorizing - * the loop. NEON is only faster on the A53, and with the newer cores, it is less - * than half the speed. - * - * Additionally, this is used on WASM SIMD128 because it JITs to the same - * SIMD instructions and has the same issue. - */ - XXH_COMPILER_GUARD(acc); -#endif - return acc; -} - -/*! - * @internal - * @brief Mixes all bits to finalize the hash. - * - * The final mix ensures that all input bits have a chance to impact any bit in - * the output digest, resulting in an unbiased distribution. - * - * @param hash The hash to avalanche. - * @return The avalanched hash. - */ -static xxh_u32 XXH32_avalanche(xxh_u32 hash) -{ - hash ^= hash >> 15; - hash *= XXH_PRIME32_2; - hash ^= hash >> 13; - hash *= XXH_PRIME32_3; - hash ^= hash >> 16; - return hash; -} - -#define XXH_get32bits(p) XXH_readLE32_align(p, align) - -/*! - * @internal - * @brief Processes the last 0-15 bytes of @p ptr. - * - * There may be up to 15 bytes remaining to consume from the input. - * This final stage will digest them to ensure that all input bytes are present - * in the final mix. - * - * @param hash The hash to finalize. - * @param ptr The pointer to the remaining input. - * @param len The remaining length, modulo 16. - * @param align Whether @p ptr is aligned. - * @return The finalized hash. - * @see XXH64_finalize(). - */ -static XXH_PUREF xxh_u32 -XXH32_finalize(xxh_u32 hash, const xxh_u8* ptr, size_t len, XXH_alignment align) -{ -#define XXH_PROCESS1 do { \ - hash += (*ptr++) * XXH_PRIME32_5; \ - hash = XXH_rotl32(hash, 11) * XXH_PRIME32_1; \ -} while (0) - -#define XXH_PROCESS4 do { \ - hash += XXH_get32bits(ptr) * XXH_PRIME32_3; \ - ptr += 4; \ - hash = XXH_rotl32(hash, 17) * XXH_PRIME32_4; \ -} while (0) - - if (ptr==NULL) XXH_ASSERT(len == 0); - - /* Compact rerolled version; generally faster */ - if (!XXH32_ENDJMP) { - len &= 15; - while (len >= 4) { - XXH_PROCESS4; - len -= 4; - } - while (len > 0) { - XXH_PROCESS1; - --len; - } - return XXH32_avalanche(hash); - } else { - switch(len&15) /* or switch(bEnd - p) */ { - case 12: XXH_PROCESS4; - XXH_FALLTHROUGH; /* fallthrough */ - case 8: XXH_PROCESS4; - XXH_FALLTHROUGH; /* fallthrough */ - case 4: XXH_PROCESS4; - return XXH32_avalanche(hash); - - case 13: XXH_PROCESS4; - XXH_FALLTHROUGH; /* fallthrough */ - case 9: XXH_PROCESS4; - XXH_FALLTHROUGH; /* fallthrough */ - case 5: XXH_PROCESS4; - XXH_PROCESS1; - return XXH32_avalanche(hash); - - case 14: XXH_PROCESS4; - XXH_FALLTHROUGH; /* fallthrough */ - case 10: XXH_PROCESS4; - XXH_FALLTHROUGH; /* fallthrough */ - case 6: XXH_PROCESS4; - XXH_PROCESS1; - XXH_PROCESS1; - return XXH32_avalanche(hash); - - case 15: XXH_PROCESS4; - XXH_FALLTHROUGH; /* fallthrough */ - case 11: XXH_PROCESS4; - XXH_FALLTHROUGH; /* fallthrough */ - case 7: XXH_PROCESS4; - XXH_FALLTHROUGH; /* fallthrough */ - case 3: XXH_PROCESS1; - XXH_FALLTHROUGH; /* fallthrough */ - case 2: XXH_PROCESS1; - XXH_FALLTHROUGH; /* fallthrough */ - case 1: XXH_PROCESS1; - XXH_FALLTHROUGH; /* fallthrough */ - case 0: return XXH32_avalanche(hash); - } - XXH_ASSERT(0); - return hash; /* reaching this point is deemed impossible */ - } -} - -#ifdef XXH_OLD_NAMES -# define PROCESS1 XXH_PROCESS1 -# define PROCESS4 XXH_PROCESS4 -#else -# undef XXH_PROCESS1 -# undef XXH_PROCESS4 -#endif - -/*! - * @internal - * @brief The implementation for @ref XXH32(). - * - * @param input , len , seed Directly passed from @ref XXH32(). - * @param align Whether @p input is aligned. - * @return The calculated hash. - */ -XXH_FORCE_INLINE XXH_PUREF xxh_u32 -XXH32_endian_align(const xxh_u8* input, size_t len, xxh_u32 seed, XXH_alignment align) -{ - xxh_u32 h32; - - if (input==NULL) XXH_ASSERT(len == 0); - - if (len>=16) { - const xxh_u8* const bEnd = input + len; - const xxh_u8* const limit = bEnd - 15; - xxh_u32 v1 = seed + XXH_PRIME32_1 + XXH_PRIME32_2; - xxh_u32 v2 = seed + XXH_PRIME32_2; - xxh_u32 v3 = seed + 0; - xxh_u32 v4 = seed - XXH_PRIME32_1; - - do { - v1 = XXH32_round(v1, XXH_get32bits(input)); input += 4; - v2 = XXH32_round(v2, XXH_get32bits(input)); input += 4; - v3 = XXH32_round(v3, XXH_get32bits(input)); input += 4; - v4 = XXH32_round(v4, XXH_get32bits(input)); input += 4; - } while (input < limit); - - h32 = XXH_rotl32(v1, 1) + XXH_rotl32(v2, 7) - + XXH_rotl32(v3, 12) + XXH_rotl32(v4, 18); - } else { - h32 = seed + XXH_PRIME32_5; - } - - h32 += (xxh_u32)len; - - return XXH32_finalize(h32, input, len&15, align); -} - -/*! @ingroup XXH32_family */ -XXH_PUBLIC_API XXH32_hash_t XXH32 (const void* input, size_t len, XXH32_hash_t seed) -{ -#if !defined(XXH_NO_STREAM) && XXH_SIZE_OPT >= 2 - /* Simple version, good for code maintenance, but unfortunately slow for small inputs */ - XXH32_state_t state; - XXH32_reset(&state, seed); - XXH32_update(&state, (const xxh_u8*)input, len); - return XXH32_digest(&state); -#else - if (XXH_FORCE_ALIGN_CHECK) { - if ((((size_t)input) & 3) == 0) { /* Input is 4-bytes aligned, leverage the speed benefit */ - return XXH32_endian_align((const xxh_u8*)input, len, seed, XXH_aligned); - } } - - return XXH32_endian_align((const xxh_u8*)input, len, seed, XXH_unaligned); -#endif -} - - - -/******* Hash streaming *******/ -#ifndef XXH_NO_STREAM -/*! @ingroup XXH32_family */ -XXH_PUBLIC_API XXH32_state_t* XXH32_createState(void) -{ - return (XXH32_state_t*)XXH_malloc(sizeof(XXH32_state_t)); -} -/*! @ingroup XXH32_family */ -XXH_PUBLIC_API XXH_errorcode XXH32_freeState(XXH32_state_t* statePtr) -{ - XXH_free(statePtr); - return XXH_OK; -} - -/*! @ingroup XXH32_family */ -XXH_PUBLIC_API void XXH32_copyState(XXH32_state_t* dstState, const XXH32_state_t* srcState) -{ - XXH_memcpy(dstState, srcState, sizeof(*dstState)); -} - -/*! @ingroup XXH32_family */ -XXH_PUBLIC_API XXH_errorcode XXH32_reset(XXH32_state_t* statePtr, XXH32_hash_t seed) -{ - XXH_ASSERT(statePtr != NULL); - memset(statePtr, 0, sizeof(*statePtr)); - statePtr->v[0] = seed + XXH_PRIME32_1 + XXH_PRIME32_2; - statePtr->v[1] = seed + XXH_PRIME32_2; - statePtr->v[2] = seed + 0; - statePtr->v[3] = seed - XXH_PRIME32_1; - return XXH_OK; -} - - -/*! @ingroup XXH32_family */ -XXH_PUBLIC_API XXH_errorcode -XXH32_update(XXH32_state_t* state, const void* input, size_t len) -{ - if (input==NULL) { - XXH_ASSERT(len == 0); - return XXH_OK; - } - - { const xxh_u8* p = (const xxh_u8*)input; - const xxh_u8* const bEnd = p + len; - - state->total_len_32 += (XXH32_hash_t)len; - state->large_len |= (XXH32_hash_t)((len>=16) | (state->total_len_32>=16)); - - if (state->memsize + len < 16) { /* fill in tmp buffer */ - XXH_memcpy((xxh_u8*)(state->mem32) + state->memsize, input, len); - state->memsize += (XXH32_hash_t)len; - return XXH_OK; - } - - if (state->memsize) { /* some data left from previous update */ - XXH_memcpy((xxh_u8*)(state->mem32) + state->memsize, input, 16-state->memsize); - { const xxh_u32* p32 = state->mem32; - state->v[0] = XXH32_round(state->v[0], XXH_readLE32(p32)); p32++; - state->v[1] = XXH32_round(state->v[1], XXH_readLE32(p32)); p32++; - state->v[2] = XXH32_round(state->v[2], XXH_readLE32(p32)); p32++; - state->v[3] = XXH32_round(state->v[3], XXH_readLE32(p32)); - } - p += 16-state->memsize; - state->memsize = 0; - } - - if (p <= bEnd-16) { - const xxh_u8* const limit = bEnd - 16; - - do { - state->v[0] = XXH32_round(state->v[0], XXH_readLE32(p)); p+=4; - state->v[1] = XXH32_round(state->v[1], XXH_readLE32(p)); p+=4; - state->v[2] = XXH32_round(state->v[2], XXH_readLE32(p)); p+=4; - state->v[3] = XXH32_round(state->v[3], XXH_readLE32(p)); p+=4; - } while (p<=limit); - - } - - if (p < bEnd) { - XXH_memcpy(state->mem32, p, (size_t)(bEnd-p)); - state->memsize = (unsigned)(bEnd-p); - } - } - - return XXH_OK; -} - - -/*! @ingroup XXH32_family */ -XXH_PUBLIC_API XXH32_hash_t XXH32_digest(const XXH32_state_t* state) -{ - xxh_u32 h32; - - if (state->large_len) { - h32 = XXH_rotl32(state->v[0], 1) - + XXH_rotl32(state->v[1], 7) - + XXH_rotl32(state->v[2], 12) - + XXH_rotl32(state->v[3], 18); - } else { - h32 = state->v[2] /* == seed */ + XXH_PRIME32_5; - } - - h32 += state->total_len_32; - - return XXH32_finalize(h32, (const xxh_u8*)state->mem32, state->memsize, XXH_aligned); -} -#endif /* !XXH_NO_STREAM */ - -/******* Canonical representation *******/ - -/*! @ingroup XXH32_family */ -XXH_PUBLIC_API void XXH32_canonicalFromHash(XXH32_canonical_t* dst, XXH32_hash_t hash) -{ - XXH_STATIC_ASSERT(sizeof(XXH32_canonical_t) == sizeof(XXH32_hash_t)); - if (XXH_CPU_LITTLE_ENDIAN) hash = XXH_swap32(hash); - XXH_memcpy(dst, &hash, sizeof(*dst)); -} -/*! @ingroup XXH32_family */ -XXH_PUBLIC_API XXH32_hash_t XXH32_hashFromCanonical(const XXH32_canonical_t* src) -{ - return XXH_readBE32(src); -} - - -#ifndef XXH_NO_LONG_LONG - -/* ******************************************************************* -* 64-bit hash functions -*********************************************************************/ -/*! - * @} - * @ingroup impl - * @{ - */ -/******* Memory access *******/ - -typedef XXH64_hash_t xxh_u64; - -#ifdef XXH_OLD_NAMES -# define U64 xxh_u64 -#endif - -#if (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==3)) -/* - * Manual byteshift. Best for old compilers which don't inline memcpy. - * We actually directly use XXH_readLE64 and XXH_readBE64. - */ -#elif (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==2)) - -/* Force direct memory access. Only works on CPU which support unaligned memory access in hardware */ -static xxh_u64 XXH_read64(const void* memPtr) -{ - return *(const xxh_u64*) memPtr; -} - -#elif (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==1)) - -/* - * __attribute__((aligned(1))) is supported by gcc and clang. Originally the - * documentation claimed that it only increased the alignment, but actually it - * can decrease it on gcc, clang, and icc: - * https://gcc.gnu.org/bugzilla/show_bug.cgi?id=69502, - * https://gcc.godbolt.org/z/xYez1j67Y. - */ -#ifdef XXH_OLD_NAMES -typedef union { xxh_u32 u32; xxh_u64 u64; } __attribute__((__packed__)) unalign64; -#endif -static xxh_u64 XXH_read64(const void* ptr) -{ - typedef __attribute__((__aligned__(1))) xxh_u64 xxh_unalign64; - return *((const xxh_unalign64*)ptr); -} - -#else - -/* - * Portable and safe solution. Generally efficient. - * see: https://fastcompression.blogspot.com/2015/08/accessing-unaligned-memory.html - */ -static xxh_u64 XXH_read64(const void* memPtr) -{ - xxh_u64 val; - XXH_memcpy(&val, memPtr, sizeof(val)); - return val; -} - -#endif /* XXH_FORCE_DIRECT_MEMORY_ACCESS */ - -#if defined(_MSC_VER) /* Visual Studio */ -# define XXH_swap64 _byteswap_uint64 -#elif XXH_GCC_VERSION >= 403 -# define XXH_swap64 __builtin_bswap64 -#else -static xxh_u64 XXH_swap64(xxh_u64 x) -{ - return ((x << 56) & 0xff00000000000000ULL) | - ((x << 40) & 0x00ff000000000000ULL) | - ((x << 24) & 0x0000ff0000000000ULL) | - ((x << 8) & 0x000000ff00000000ULL) | - ((x >> 8) & 0x00000000ff000000ULL) | - ((x >> 24) & 0x0000000000ff0000ULL) | - ((x >> 40) & 0x000000000000ff00ULL) | - ((x >> 56) & 0x00000000000000ffULL); -} -#endif - - -/* XXH_FORCE_MEMORY_ACCESS==3 is an endian-independent byteshift load. */ -#if (defined(XXH_FORCE_MEMORY_ACCESS) && (XXH_FORCE_MEMORY_ACCESS==3)) - -XXH_FORCE_INLINE xxh_u64 XXH_readLE64(const void* memPtr) -{ - const xxh_u8* bytePtr = (const xxh_u8 *)memPtr; - return bytePtr[0] - | ((xxh_u64)bytePtr[1] << 8) - | ((xxh_u64)bytePtr[2] << 16) - | ((xxh_u64)bytePtr[3] << 24) - | ((xxh_u64)bytePtr[4] << 32) - | ((xxh_u64)bytePtr[5] << 40) - | ((xxh_u64)bytePtr[6] << 48) - | ((xxh_u64)bytePtr[7] << 56); -} - -XXH_FORCE_INLINE xxh_u64 XXH_readBE64(const void* memPtr) -{ - const xxh_u8* bytePtr = (const xxh_u8 *)memPtr; - return bytePtr[7] - | ((xxh_u64)bytePtr[6] << 8) - | ((xxh_u64)bytePtr[5] << 16) - | ((xxh_u64)bytePtr[4] << 24) - | ((xxh_u64)bytePtr[3] << 32) - | ((xxh_u64)bytePtr[2] << 40) - | ((xxh_u64)bytePtr[1] << 48) - | ((xxh_u64)bytePtr[0] << 56); -} - -#else -XXH_FORCE_INLINE xxh_u64 XXH_readLE64(const void* ptr) -{ - return XXH_CPU_LITTLE_ENDIAN ? XXH_read64(ptr) : XXH_swap64(XXH_read64(ptr)); -} - -static xxh_u64 XXH_readBE64(const void* ptr) -{ - return XXH_CPU_LITTLE_ENDIAN ? XXH_swap64(XXH_read64(ptr)) : XXH_read64(ptr); -} -#endif - -XXH_FORCE_INLINE xxh_u64 -XXH_readLE64_align(const void* ptr, XXH_alignment align) -{ - if (align==XXH_unaligned) - return XXH_readLE64(ptr); - else - return XXH_CPU_LITTLE_ENDIAN ? *(const xxh_u64*)ptr : XXH_swap64(*(const xxh_u64*)ptr); -} - - -/******* xxh64 *******/ -/*! - * @} - * @defgroup XXH64_impl XXH64 implementation - * @ingroup impl - * - * Details on the XXH64 implementation. - * @{ - */ -/* #define rather that static const, to be used as initializers */ -#define XXH_PRIME64_1 0x9E3779B185EBCA87ULL /*!< 0b1001111000110111011110011011000110000101111010111100101010000111 */ -#define XXH_PRIME64_2 0xC2B2AE3D27D4EB4FULL /*!< 0b1100001010110010101011100011110100100111110101001110101101001111 */ -#define XXH_PRIME64_3 0x165667B19E3779F9ULL /*!< 0b0001011001010110011001111011000110011110001101110111100111111001 */ -#define XXH_PRIME64_4 0x85EBCA77C2B2AE63ULL /*!< 0b1000010111101011110010100111011111000010101100101010111001100011 */ -#define XXH_PRIME64_5 0x27D4EB2F165667C5ULL /*!< 0b0010011111010100111010110010111100010110010101100110011111000101 */ - -#ifdef XXH_OLD_NAMES -# define PRIME64_1 XXH_PRIME64_1 -# define PRIME64_2 XXH_PRIME64_2 -# define PRIME64_3 XXH_PRIME64_3 -# define PRIME64_4 XXH_PRIME64_4 -# define PRIME64_5 XXH_PRIME64_5 -#endif - -/*! @copydoc XXH32_round */ -static xxh_u64 XXH64_round(xxh_u64 acc, xxh_u64 input) -{ - acc += input * XXH_PRIME64_2; - acc = XXH_rotl64(acc, 31); - acc *= XXH_PRIME64_1; -#if (defined(__AVX512F__)) && !defined(XXH_ENABLE_AUTOVECTORIZE) - /* - * DISABLE AUTOVECTORIZATION: - * A compiler fence is used to prevent GCC and Clang from - * autovectorizing the XXH64 loop (pragmas and attributes don't work for some - * reason) without globally disabling AVX512. - * - * Autovectorization of XXH64 tends to be detrimental, - * though the exact outcome may change depending on exact cpu and compiler version. - * For information, it has been reported as detrimental for Skylake-X, - * but possibly beneficial for Zen4. - * - * The default is to disable auto-vectorization, - * but you can select to enable it instead using `XXH_ENABLE_AUTOVECTORIZE` build variable. - */ - XXH_COMPILER_GUARD(acc); -#endif - return acc; -} - -static xxh_u64 XXH64_mergeRound(xxh_u64 acc, xxh_u64 val) -{ - val = XXH64_round(0, val); - acc ^= val; - acc = acc * XXH_PRIME64_1 + XXH_PRIME64_4; - return acc; -} - -/*! @copydoc XXH32_avalanche */ -static xxh_u64 XXH64_avalanche(xxh_u64 hash) -{ - hash ^= hash >> 33; - hash *= XXH_PRIME64_2; - hash ^= hash >> 29; - hash *= XXH_PRIME64_3; - hash ^= hash >> 32; - return hash; -} - - -#define XXH_get64bits(p) XXH_readLE64_align(p, align) - -/*! - * @internal - * @brief Processes the last 0-31 bytes of @p ptr. - * - * There may be up to 31 bytes remaining to consume from the input. - * This final stage will digest them to ensure that all input bytes are present - * in the final mix. - * - * @param hash The hash to finalize. - * @param ptr The pointer to the remaining input. - * @param len The remaining length, modulo 32. - * @param align Whether @p ptr is aligned. - * @return The finalized hash - * @see XXH32_finalize(). - */ -static XXH_PUREF xxh_u64 -XXH64_finalize(xxh_u64 hash, const xxh_u8* ptr, size_t len, XXH_alignment align) -{ - if (ptr==NULL) XXH_ASSERT(len == 0); - len &= 31; - while (len >= 8) { - xxh_u64 const k1 = XXH64_round(0, XXH_get64bits(ptr)); - ptr += 8; - hash ^= k1; - hash = XXH_rotl64(hash,27) * XXH_PRIME64_1 + XXH_PRIME64_4; - len -= 8; - } - if (len >= 4) { - hash ^= (xxh_u64)(XXH_get32bits(ptr)) * XXH_PRIME64_1; - ptr += 4; - hash = XXH_rotl64(hash, 23) * XXH_PRIME64_2 + XXH_PRIME64_3; - len -= 4; - } - while (len > 0) { - hash ^= (*ptr++) * XXH_PRIME64_5; - hash = XXH_rotl64(hash, 11) * XXH_PRIME64_1; - --len; - } - return XXH64_avalanche(hash); -} - -#ifdef XXH_OLD_NAMES -# define PROCESS1_64 XXH_PROCESS1_64 -# define PROCESS4_64 XXH_PROCESS4_64 -# define PROCESS8_64 XXH_PROCESS8_64 -#else -# undef XXH_PROCESS1_64 -# undef XXH_PROCESS4_64 -# undef XXH_PROCESS8_64 -#endif - -/*! - * @internal - * @brief The implementation for @ref XXH64(). - * - * @param input , len , seed Directly passed from @ref XXH64(). - * @param align Whether @p input is aligned. - * @return The calculated hash. - */ -XXH_FORCE_INLINE XXH_PUREF xxh_u64 -XXH64_endian_align(const xxh_u8* input, size_t len, xxh_u64 seed, XXH_alignment align) -{ - xxh_u64 h64; - if (input==NULL) XXH_ASSERT(len == 0); - - if (len>=32) { - const xxh_u8* const bEnd = input + len; - const xxh_u8* const limit = bEnd - 31; - xxh_u64 v1 = seed + XXH_PRIME64_1 + XXH_PRIME64_2; - xxh_u64 v2 = seed + XXH_PRIME64_2; - xxh_u64 v3 = seed + 0; - xxh_u64 v4 = seed - XXH_PRIME64_1; - - do { - v1 = XXH64_round(v1, XXH_get64bits(input)); input+=8; - v2 = XXH64_round(v2, XXH_get64bits(input)); input+=8; - v3 = XXH64_round(v3, XXH_get64bits(input)); input+=8; - v4 = XXH64_round(v4, XXH_get64bits(input)); input+=8; - } while (input= 2 - /* Simple version, good for code maintenance, but unfortunately slow for small inputs */ - XXH64_state_t state; - XXH64_reset(&state, seed); - XXH64_update(&state, (const xxh_u8*)input, len); - return XXH64_digest(&state); -#else - if (XXH_FORCE_ALIGN_CHECK) { - if ((((size_t)input) & 7)==0) { /* Input is aligned, let's leverage the speed advantage */ - return XXH64_endian_align((const xxh_u8*)input, len, seed, XXH_aligned); - } } - - return XXH64_endian_align((const xxh_u8*)input, len, seed, XXH_unaligned); - -#endif -} - -/******* Hash Streaming *******/ -#ifndef XXH_NO_STREAM -/*! @ingroup XXH64_family*/ -XXH_PUBLIC_API XXH64_state_t* XXH64_createState(void) -{ - return (XXH64_state_t*)XXH_malloc(sizeof(XXH64_state_t)); -} -/*! @ingroup XXH64_family */ -XXH_PUBLIC_API XXH_errorcode XXH64_freeState(XXH64_state_t* statePtr) -{ - XXH_free(statePtr); - return XXH_OK; -} - -/*! @ingroup XXH64_family */ -XXH_PUBLIC_API void XXH64_copyState(XXH_NOESCAPE XXH64_state_t* dstState, const XXH64_state_t* srcState) -{ - XXH_memcpy(dstState, srcState, sizeof(*dstState)); -} - -/*! @ingroup XXH64_family */ -XXH_PUBLIC_API XXH_errorcode XXH64_reset(XXH_NOESCAPE XXH64_state_t* statePtr, XXH64_hash_t seed) -{ - XXH_ASSERT(statePtr != NULL); - memset(statePtr, 0, sizeof(*statePtr)); - statePtr->v[0] = seed + XXH_PRIME64_1 + XXH_PRIME64_2; - statePtr->v[1] = seed + XXH_PRIME64_2; - statePtr->v[2] = seed + 0; - statePtr->v[3] = seed - XXH_PRIME64_1; - return XXH_OK; -} - -/*! @ingroup XXH64_family */ -XXH_PUBLIC_API XXH_errorcode -XXH64_update (XXH_NOESCAPE XXH64_state_t* state, XXH_NOESCAPE const void* input, size_t len) -{ - if (input==NULL) { - XXH_ASSERT(len == 0); - return XXH_OK; - } - - { const xxh_u8* p = (const xxh_u8*)input; - const xxh_u8* const bEnd = p + len; - - state->total_len += len; - - if (state->memsize + len < 32) { /* fill in tmp buffer */ - XXH_memcpy(((xxh_u8*)state->mem64) + state->memsize, input, len); - state->memsize += (xxh_u32)len; - return XXH_OK; - } - - if (state->memsize) { /* tmp buffer is full */ - XXH_memcpy(((xxh_u8*)state->mem64) + state->memsize, input, 32-state->memsize); - state->v[0] = XXH64_round(state->v[0], XXH_readLE64(state->mem64+0)); - state->v[1] = XXH64_round(state->v[1], XXH_readLE64(state->mem64+1)); - state->v[2] = XXH64_round(state->v[2], XXH_readLE64(state->mem64+2)); - state->v[3] = XXH64_round(state->v[3], XXH_readLE64(state->mem64+3)); - p += 32 - state->memsize; - state->memsize = 0; - } - - if (p+32 <= bEnd) { - const xxh_u8* const limit = bEnd - 32; - - do { - state->v[0] = XXH64_round(state->v[0], XXH_readLE64(p)); p+=8; - state->v[1] = XXH64_round(state->v[1], XXH_readLE64(p)); p+=8; - state->v[2] = XXH64_round(state->v[2], XXH_readLE64(p)); p+=8; - state->v[3] = XXH64_round(state->v[3], XXH_readLE64(p)); p+=8; - } while (p<=limit); - - } - - if (p < bEnd) { - XXH_memcpy(state->mem64, p, (size_t)(bEnd-p)); - state->memsize = (unsigned)(bEnd-p); - } - } - - return XXH_OK; -} - - -/*! @ingroup XXH64_family */ -XXH_PUBLIC_API XXH64_hash_t XXH64_digest(XXH_NOESCAPE const XXH64_state_t* state) -{ - xxh_u64 h64; - - if (state->total_len >= 32) { - h64 = XXH_rotl64(state->v[0], 1) + XXH_rotl64(state->v[1], 7) + XXH_rotl64(state->v[2], 12) + XXH_rotl64(state->v[3], 18); - h64 = XXH64_mergeRound(h64, state->v[0]); - h64 = XXH64_mergeRound(h64, state->v[1]); - h64 = XXH64_mergeRound(h64, state->v[2]); - h64 = XXH64_mergeRound(h64, state->v[3]); - } else { - h64 = state->v[2] /*seed*/ + XXH_PRIME64_5; - } - - h64 += (xxh_u64) state->total_len; - - return XXH64_finalize(h64, (const xxh_u8*)state->mem64, (size_t)state->total_len, XXH_aligned); -} -#endif /* !XXH_NO_STREAM */ - -/******* Canonical representation *******/ - -/*! @ingroup XXH64_family */ -XXH_PUBLIC_API void XXH64_canonicalFromHash(XXH_NOESCAPE XXH64_canonical_t* dst, XXH64_hash_t hash) -{ - XXH_STATIC_ASSERT(sizeof(XXH64_canonical_t) == sizeof(XXH64_hash_t)); - if (XXH_CPU_LITTLE_ENDIAN) hash = XXH_swap64(hash); - XXH_memcpy(dst, &hash, sizeof(*dst)); -} - -/*! @ingroup XXH64_family */ -XXH_PUBLIC_API XXH64_hash_t XXH64_hashFromCanonical(XXH_NOESCAPE const XXH64_canonical_t* src) -{ - return XXH_readBE64(src); -} - -#ifndef XXH_NO_XXH3 - -/* ********************************************************************* -* XXH3 -* New generation hash designed for speed on small keys and vectorization -************************************************************************ */ -/*! - * @} - * @defgroup XXH3_impl XXH3 implementation - * @ingroup impl - * @{ - */ - -/* === Compiler specifics === */ - -#if ((defined(sun) || defined(__sun)) && __cplusplus) /* Solaris includes __STDC_VERSION__ with C++. Tested with GCC 5.5 */ -# define XXH_RESTRICT /* disable */ -#elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 199901L /* >= C99 */ -# define XXH_RESTRICT restrict -#elif (defined (__GNUC__) && ((__GNUC__ > 3) || (__GNUC__ == 3 && __GNUC_MINOR__ >= 1))) \ - || (defined (__clang__)) \ - || (defined (_MSC_VER) && (_MSC_VER >= 1400)) \ - || (defined (__INTEL_COMPILER) && (__INTEL_COMPILER >= 1300)) -/* - * There are a LOT more compilers that recognize __restrict but this - * covers the major ones. - */ -# define XXH_RESTRICT __restrict -#else -# define XXH_RESTRICT /* disable */ -#endif - -#if (defined(__GNUC__) && (__GNUC__ >= 3)) \ - || (defined(__INTEL_COMPILER) && (__INTEL_COMPILER >= 800)) \ - || defined(__clang__) -# define XXH_likely(x) __builtin_expect(x, 1) -# define XXH_unlikely(x) __builtin_expect(x, 0) -#else -# define XXH_likely(x) (x) -# define XXH_unlikely(x) (x) -#endif - -#ifndef XXH_HAS_INCLUDE -# ifdef __has_include -/* - * Not defined as XXH_HAS_INCLUDE(x) (function-like) because - * this causes segfaults in Apple Clang 4.2 (on Mac OS X 10.7 Lion) - */ -# define XXH_HAS_INCLUDE __has_include -# else -# define XXH_HAS_INCLUDE(x) 0 -# endif -#endif - -#if defined(__GNUC__) || defined(__clang__) -# if defined(__ARM_FEATURE_SVE) -# include -# endif -# if defined(__ARM_NEON__) || defined(__ARM_NEON) \ - || (defined(_M_ARM) && _M_ARM >= 7) \ - || defined(_M_ARM64) || defined(_M_ARM64EC) \ - || (defined(__wasm_simd128__) && XXH_HAS_INCLUDE()) /* WASM SIMD128 via SIMDe */ -# define inline __inline__ /* circumvent a clang bug */ -# include -# undef inline -# elif defined(__AVX2__) -# include -# elif defined(__SSE2__) -# include -# endif -#endif - -#if defined(_MSC_VER) -# include -#endif - -/* - * One goal of XXH3 is to make it fast on both 32-bit and 64-bit, while - * remaining a true 64-bit/128-bit hash function. - * - * This is done by prioritizing a subset of 64-bit operations that can be - * emulated without too many steps on the average 32-bit machine. - * - * For example, these two lines seem similar, and run equally fast on 64-bit: - * - * xxh_u64 x; - * x ^= (x >> 47); // good - * x ^= (x >> 13); // bad - * - * However, to a 32-bit machine, there is a major difference. - * - * x ^= (x >> 47) looks like this: - * - * x.lo ^= (x.hi >> (47 - 32)); - * - * while x ^= (x >> 13) looks like this: - * - * // note: funnel shifts are not usually cheap. - * x.lo ^= (x.lo >> 13) | (x.hi << (32 - 13)); - * x.hi ^= (x.hi >> 13); - * - * The first one is significantly faster than the second, simply because the - * shift is larger than 32. This means: - * - All the bits we need are in the upper 32 bits, so we can ignore the lower - * 32 bits in the shift. - * - The shift result will always fit in the lower 32 bits, and therefore, - * we can ignore the upper 32 bits in the xor. - * - * Thanks to this optimization, XXH3 only requires these features to be efficient: - * - * - Usable unaligned access - * - A 32-bit or 64-bit ALU - * - If 32-bit, a decent ADC instruction - * - A 32 or 64-bit multiply with a 64-bit result - * - For the 128-bit variant, a decent byteswap helps short inputs. - * - * The first two are already required by XXH32, and almost all 32-bit and 64-bit - * platforms which can run XXH32 can run XXH3 efficiently. - * - * Thumb-1, the classic 16-bit only subset of ARM's instruction set, is one - * notable exception. - * - * First of all, Thumb-1 lacks support for the UMULL instruction which - * performs the important long multiply. This means numerous __aeabi_lmul - * calls. - * - * Second of all, the 8 functional registers are just not enough. - * Setup for __aeabi_lmul, byteshift loads, pointers, and all arithmetic need - * Lo registers, and this shuffling results in thousands more MOVs than A32. - * - * A32 and T32 don't have this limitation. They can access all 14 registers, - * do a 32->64 multiply with UMULL, and the flexible operand allowing free - * shifts is helpful, too. - * - * Therefore, we do a quick sanity check. - * - * If compiling Thumb-1 for a target which supports ARM instructions, we will - * emit a warning, as it is not a "sane" platform to compile for. - * - * Usually, if this happens, it is because of an accident and you probably need - * to specify -march, as you likely meant to compile for a newer architecture. - * - * Credit: large sections of the vectorial and asm source code paths - * have been contributed by @easyaspi314 - */ -#if defined(__thumb__) && !defined(__thumb2__) && defined(__ARM_ARCH_ISA_ARM) -# warning "XXH3 is highly inefficient without ARM or Thumb-2." -#endif - -/* ========================================== - * Vectorization detection - * ========================================== */ - -#ifdef XXH_DOXYGEN -/*! - * @ingroup tuning - * @brief Overrides the vectorization implementation chosen for XXH3. - * - * Can be defined to 0 to disable SIMD or any of the values mentioned in - * @ref XXH_VECTOR_TYPE. - * - * If this is not defined, it uses predefined macros to determine the best - * implementation. - */ -# define XXH_VECTOR XXH_SCALAR -/*! - * @ingroup tuning - * @brief Possible values for @ref XXH_VECTOR. - * - * Note that these are actually implemented as macros. - * - * If this is not defined, it is detected automatically. - * internal macro XXH_X86DISPATCH overrides this. - */ -enum XXH_VECTOR_TYPE /* fake enum */ { - XXH_SCALAR = 0, /*!< Portable scalar version */ - XXH_SSE2 = 1, /*!< - * SSE2 for Pentium 4, Opteron, all x86_64. - * - * @note SSE2 is also guaranteed on Windows 10, macOS, and - * Android x86. - */ - XXH_AVX2 = 2, /*!< AVX2 for Haswell and Bulldozer */ - XXH_AVX512 = 3, /*!< AVX512 for Skylake and Icelake */ - XXH_NEON = 4, /*!< - * NEON for most ARMv7-A, all AArch64, and WASM SIMD128 - * via the SIMDeverywhere polyfill provided with the - * Emscripten SDK. - */ - XXH_VSX = 5, /*!< VSX and ZVector for POWER8/z13 (64-bit) */ - XXH_SVE = 6, /*!< SVE for some ARMv8-A and ARMv9-A */ -}; -/*! - * @ingroup tuning - * @brief Selects the minimum alignment for XXH3's accumulators. - * - * When using SIMD, this should match the alignment required for said vector - * type, so, for example, 32 for AVX2. - * - * Default: Auto detected. - */ -# define XXH_ACC_ALIGN 8 -#endif - -/* Actual definition */ -#ifndef XXH_DOXYGEN -# define XXH_SCALAR 0 -# define XXH_SSE2 1 -# define XXH_AVX2 2 -# define XXH_AVX512 3 -# define XXH_NEON 4 -# define XXH_VSX 5 -# define XXH_SVE 6 -#endif - -#ifndef XXH_VECTOR /* can be defined on command line */ -# if defined(__ARM_FEATURE_SVE) -# define XXH_VECTOR XXH_SVE -# elif ( \ - defined(__ARM_NEON__) || defined(__ARM_NEON) /* gcc */ \ - || defined(_M_ARM) || defined(_M_ARM64) || defined(_M_ARM64EC) /* msvc */ \ - || (defined(__wasm_simd128__) && XXH_HAS_INCLUDE()) /* wasm simd128 via SIMDe */ \ - ) && ( \ - defined(_WIN32) || defined(__LITTLE_ENDIAN__) /* little endian only */ \ - || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__) \ - ) -# define XXH_VECTOR XXH_NEON -# elif defined(__AVX512F__) -# define XXH_VECTOR XXH_AVX512 -# elif defined(__AVX2__) -# define XXH_VECTOR XXH_AVX2 -# elif defined(__SSE2__) || defined(_M_AMD64) || defined(_M_X64) || (defined(_M_IX86_FP) && (_M_IX86_FP == 2)) -# define XXH_VECTOR XXH_SSE2 -# elif (defined(__PPC64__) && defined(__POWER8_VECTOR__)) \ - || (defined(__s390x__) && defined(__VEC__)) \ - && defined(__GNUC__) /* TODO: IBM XL */ -# define XXH_VECTOR XXH_VSX -# else -# define XXH_VECTOR XXH_SCALAR -# endif -#endif - -/* __ARM_FEATURE_SVE is only supported by GCC & Clang. */ -#if (XXH_VECTOR == XXH_SVE) && !defined(__ARM_FEATURE_SVE) -# ifdef _MSC_VER -# pragma warning(once : 4606) -# else -# warning "__ARM_FEATURE_SVE isn't supported. Use SCALAR instead." -# endif -# undef XXH_VECTOR -# define XXH_VECTOR XXH_SCALAR -#endif - -/* - * Controls the alignment of the accumulator, - * for compatibility with aligned vector loads, which are usually faster. - */ -#ifndef XXH_ACC_ALIGN -# if defined(XXH_X86DISPATCH) -# define XXH_ACC_ALIGN 64 /* for compatibility with avx512 */ -# elif XXH_VECTOR == XXH_SCALAR /* scalar */ -# define XXH_ACC_ALIGN 8 -# elif XXH_VECTOR == XXH_SSE2 /* sse2 */ -# define XXH_ACC_ALIGN 16 -# elif XXH_VECTOR == XXH_AVX2 /* avx2 */ -# define XXH_ACC_ALIGN 32 -# elif XXH_VECTOR == XXH_NEON /* neon */ -# define XXH_ACC_ALIGN 16 -# elif XXH_VECTOR == XXH_VSX /* vsx */ -# define XXH_ACC_ALIGN 16 -# elif XXH_VECTOR == XXH_AVX512 /* avx512 */ -# define XXH_ACC_ALIGN 64 -# elif XXH_VECTOR == XXH_SVE /* sve */ -# define XXH_ACC_ALIGN 64 -# endif -#endif - -#if defined(XXH_X86DISPATCH) || XXH_VECTOR == XXH_SSE2 \ - || XXH_VECTOR == XXH_AVX2 || XXH_VECTOR == XXH_AVX512 -# define XXH_SEC_ALIGN XXH_ACC_ALIGN -#elif XXH_VECTOR == XXH_SVE -# define XXH_SEC_ALIGN XXH_ACC_ALIGN -#else -# define XXH_SEC_ALIGN 8 -#endif - -#if defined(__GNUC__) || defined(__clang__) -# define XXH_ALIASING __attribute__((__may_alias__)) -#else -# define XXH_ALIASING /* nothing */ -#endif - -/* - * UGLY HACK: - * GCC usually generates the best code with -O3 for xxHash. - * - * However, when targeting AVX2, it is overzealous in its unrolling resulting - * in code roughly 3/4 the speed of Clang. - * - * There are other issues, such as GCC splitting _mm256_loadu_si256 into - * _mm_loadu_si128 + _mm256_inserti128_si256. This is an optimization which - * only applies to Sandy and Ivy Bridge... which don't even support AVX2. - * - * That is why when compiling the AVX2 version, it is recommended to use either - * -O2 -mavx2 -march=haswell - * or - * -O2 -mavx2 -mno-avx256-split-unaligned-load - * for decent performance, or to use Clang instead. - * - * Fortunately, we can control the first one with a pragma that forces GCC into - * -O2, but the other one we can't control without "failed to inline always - * inline function due to target mismatch" warnings. - */ -#if XXH_VECTOR == XXH_AVX2 /* AVX2 */ \ - && defined(__GNUC__) && !defined(__clang__) /* GCC, not Clang */ \ - && defined(__OPTIMIZE__) && XXH_SIZE_OPT <= 0 /* respect -O0 and -Os */ -# pragma GCC push_options -# pragma GCC optimize("-O2") -#endif - -#if XXH_VECTOR == XXH_NEON - -/* - * UGLY HACK: While AArch64 GCC on Linux does not seem to care, on macOS, GCC -O3 - * optimizes out the entire hashLong loop because of the aliasing violation. - * - * However, GCC is also inefficient at load-store optimization with vld1q/vst1q, - * so the only option is to mark it as aliasing. - */ -typedef uint64x2_t xxh_aliasing_uint64x2_t XXH_ALIASING; - -/*! - * @internal - * @brief `vld1q_u64` but faster and alignment-safe. - * - * On AArch64, unaligned access is always safe, but on ARMv7-a, it is only - * *conditionally* safe (`vld1` has an alignment bit like `movdq[ua]` in x86). - * - * GCC for AArch64 sees `vld1q_u8` as an intrinsic instead of a load, so it - * prohibits load-store optimizations. Therefore, a direct dereference is used. - * - * Otherwise, `vld1q_u8` is used with `vreinterpretq_u8_u64` to do a safe - * unaligned load. - */ -#if defined(__aarch64__) && defined(__GNUC__) && !defined(__clang__) -XXH_FORCE_INLINE uint64x2_t XXH_vld1q_u64(void const* ptr) /* silence -Wcast-align */ -{ - return *(xxh_aliasing_uint64x2_t const *)ptr; -} -#else -XXH_FORCE_INLINE uint64x2_t XXH_vld1q_u64(void const* ptr) -{ - return vreinterpretq_u64_u8(vld1q_u8((uint8_t const*)ptr)); -} -#endif - -/*! - * @internal - * @brief `vmlal_u32` on low and high halves of a vector. - * - * This is a workaround for AArch64 GCC < 11 which implemented arm_neon.h with - * inline assembly and were therefore incapable of merging the `vget_{low, high}_u32` - * with `vmlal_u32`. - */ -#if defined(__aarch64__) && defined(__GNUC__) && !defined(__clang__) && __GNUC__ < 11 -XXH_FORCE_INLINE uint64x2_t -XXH_vmlal_low_u32(uint64x2_t acc, uint32x4_t lhs, uint32x4_t rhs) -{ - /* Inline assembly is the only way */ - __asm__("umlal %0.2d, %1.2s, %2.2s" : "+w" (acc) : "w" (lhs), "w" (rhs)); - return acc; -} -XXH_FORCE_INLINE uint64x2_t -XXH_vmlal_high_u32(uint64x2_t acc, uint32x4_t lhs, uint32x4_t rhs) -{ - /* This intrinsic works as expected */ - return vmlal_high_u32(acc, lhs, rhs); -} -#else -/* Portable intrinsic versions */ -XXH_FORCE_INLINE uint64x2_t -XXH_vmlal_low_u32(uint64x2_t acc, uint32x4_t lhs, uint32x4_t rhs) -{ - return vmlal_u32(acc, vget_low_u32(lhs), vget_low_u32(rhs)); -} -/*! @copydoc XXH_vmlal_low_u32 - * Assume the compiler converts this to vmlal_high_u32 on aarch64 */ -XXH_FORCE_INLINE uint64x2_t -XXH_vmlal_high_u32(uint64x2_t acc, uint32x4_t lhs, uint32x4_t rhs) -{ - return vmlal_u32(acc, vget_high_u32(lhs), vget_high_u32(rhs)); -} -#endif - -/*! - * @ingroup tuning - * @brief Controls the NEON to scalar ratio for XXH3 - * - * This can be set to 2, 4, 6, or 8. - * - * ARM Cortex CPUs are _very_ sensitive to how their pipelines are used. - * - * For example, the Cortex-A73 can dispatch 3 micro-ops per cycle, but only 2 of those - * can be NEON. If you are only using NEON instructions, you are only using 2/3 of the CPU - * bandwidth. - * - * This is even more noticeable on the more advanced cores like the Cortex-A76 which - * can dispatch 8 micro-ops per cycle, but still only 2 NEON micro-ops at once. - * - * Therefore, to make the most out of the pipeline, it is beneficial to run 6 NEON lanes - * and 2 scalar lanes, which is chosen by default. - * - * This does not apply to Apple processors or 32-bit processors, which run better with - * full NEON. These will default to 8. Additionally, size-optimized builds run 8 lanes. - * - * This change benefits CPUs with large micro-op buffers without negatively affecting - * most other CPUs: - * - * | Chipset | Dispatch type | NEON only | 6:2 hybrid | Diff. | - * |:----------------------|:--------------------|----------:|-----------:|------:| - * | Snapdragon 730 (A76) | 2 NEON/8 micro-ops | 8.8 GB/s | 10.1 GB/s | ~16% | - * | Snapdragon 835 (A73) | 2 NEON/3 micro-ops | 5.1 GB/s | 5.3 GB/s | ~5% | - * | Marvell PXA1928 (A53) | In-order dual-issue | 1.9 GB/s | 1.9 GB/s | 0% | - * | Apple M1 | 4 NEON/8 micro-ops | 37.3 GB/s | 36.1 GB/s | ~-3% | - * - * It also seems to fix some bad codegen on GCC, making it almost as fast as clang. - * - * When using WASM SIMD128, if this is 2 or 6, SIMDe will scalarize 2 of the lanes meaning - * it effectively becomes worse 4. - * - * @see XXH3_accumulate_512_neon() - */ -# ifndef XXH3_NEON_LANES -# if (defined(__aarch64__) || defined(__arm64__) || defined(_M_ARM64) || defined(_M_ARM64EC)) \ - && !defined(__APPLE__) && XXH_SIZE_OPT <= 0 -# define XXH3_NEON_LANES 6 -# else -# define XXH3_NEON_LANES XXH_ACC_NB -# endif -# endif -#endif /* XXH_VECTOR == XXH_NEON */ - -/* - * VSX and Z Vector helpers. - * - * This is very messy, and any pull requests to clean this up are welcome. - * - * There are a lot of problems with supporting VSX and s390x, due to - * inconsistent intrinsics, spotty coverage, and multiple endiannesses. - */ -#if XXH_VECTOR == XXH_VSX -/* Annoyingly, these headers _may_ define three macros: `bool`, `vector`, - * and `pixel`. This is a problem for obvious reasons. - * - * These keywords are unnecessary; the spec literally says they are - * equivalent to `__bool`, `__vector`, and `__pixel` and may be undef'd - * after including the header. - * - * We use pragma push_macro/pop_macro to keep the namespace clean. */ -# pragma push_macro("bool") -# pragma push_macro("vector") -# pragma push_macro("pixel") -/* silence potential macro redefined warnings */ -# undef bool -# undef vector -# undef pixel - -# if defined(__s390x__) -# include -# else -# include -# endif - -/* Restore the original macro values, if applicable. */ -# pragma pop_macro("pixel") -# pragma pop_macro("vector") -# pragma pop_macro("bool") - -typedef __vector unsigned long long xxh_u64x2; -typedef __vector unsigned char xxh_u8x16; -typedef __vector unsigned xxh_u32x4; - -/* - * UGLY HACK: Similar to aarch64 macOS GCC, s390x GCC has the same aliasing issue. - */ -typedef xxh_u64x2 xxh_aliasing_u64x2 XXH_ALIASING; - -# ifndef XXH_VSX_BE -# if defined(__BIG_ENDIAN__) \ - || (defined(__BYTE_ORDER__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__) -# define XXH_VSX_BE 1 -# elif defined(__VEC_ELEMENT_REG_ORDER__) && __VEC_ELEMENT_REG_ORDER__ == __ORDER_BIG_ENDIAN__ -# warning "-maltivec=be is not recommended. Please use native endianness." -# define XXH_VSX_BE 1 -# else -# define XXH_VSX_BE 0 -# endif -# endif /* !defined(XXH_VSX_BE) */ - -# if XXH_VSX_BE -# if defined(__POWER9_VECTOR__) || (defined(__clang__) && defined(__s390x__)) -# define XXH_vec_revb vec_revb -# else -/*! - * A polyfill for POWER9's vec_revb(). - */ -XXH_FORCE_INLINE xxh_u64x2 XXH_vec_revb(xxh_u64x2 val) -{ - xxh_u8x16 const vByteSwap = { 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01, 0x00, - 0x0F, 0x0E, 0x0D, 0x0C, 0x0B, 0x0A, 0x09, 0x08 }; - return vec_perm(val, val, vByteSwap); -} -# endif -# endif /* XXH_VSX_BE */ - -/*! - * Performs an unaligned vector load and byte swaps it on big endian. - */ -XXH_FORCE_INLINE xxh_u64x2 XXH_vec_loadu(const void *ptr) -{ - xxh_u64x2 ret; - XXH_memcpy(&ret, ptr, sizeof(xxh_u64x2)); -# if XXH_VSX_BE - ret = XXH_vec_revb(ret); -# endif - return ret; -} - -/* - * vec_mulo and vec_mule are very problematic intrinsics on PowerPC - * - * These intrinsics weren't added until GCC 8, despite existing for a while, - * and they are endian dependent. Also, their meaning swap depending on version. - * */ -# if defined(__s390x__) - /* s390x is always big endian, no issue on this platform */ -# define XXH_vec_mulo vec_mulo -# define XXH_vec_mule vec_mule -# elif defined(__clang__) && XXH_HAS_BUILTIN(__builtin_altivec_vmuleuw) && !defined(__ibmxl__) -/* Clang has a better way to control this, we can just use the builtin which doesn't swap. */ - /* The IBM XL Compiler (which defined __clang__) only implements the vec_* operations */ -# define XXH_vec_mulo __builtin_altivec_vmulouw -# define XXH_vec_mule __builtin_altivec_vmuleuw -# else -/* gcc needs inline assembly */ -/* Adapted from https://github.com/google/highwayhash/blob/master/highwayhash/hh_vsx.h. */ -XXH_FORCE_INLINE xxh_u64x2 XXH_vec_mulo(xxh_u32x4 a, xxh_u32x4 b) -{ - xxh_u64x2 result; - __asm__("vmulouw %0, %1, %2" : "=v" (result) : "v" (a), "v" (b)); - return result; -} -XXH_FORCE_INLINE xxh_u64x2 XXH_vec_mule(xxh_u32x4 a, xxh_u32x4 b) -{ - xxh_u64x2 result; - __asm__("vmuleuw %0, %1, %2" : "=v" (result) : "v" (a), "v" (b)); - return result; -} -# endif /* XXH_vec_mulo, XXH_vec_mule */ -#endif /* XXH_VECTOR == XXH_VSX */ - -#if XXH_VECTOR == XXH_SVE -#define ACCRND(acc, offset) \ -do { \ - svuint64_t input_vec = svld1_u64(mask, xinput + offset); \ - svuint64_t secret_vec = svld1_u64(mask, xsecret + offset); \ - svuint64_t mixed = sveor_u64_x(mask, secret_vec, input_vec); \ - svuint64_t swapped = svtbl_u64(input_vec, kSwap); \ - svuint64_t mixed_lo = svextw_u64_x(mask, mixed); \ - svuint64_t mixed_hi = svlsr_n_u64_x(mask, mixed, 32); \ - svuint64_t mul = svmad_u64_x(mask, mixed_lo, mixed_hi, swapped); \ - acc = svadd_u64_x(mask, acc, mul); \ -} while (0) -#endif /* XXH_VECTOR == XXH_SVE */ - -/* prefetch - * can be disabled, by declaring XXH_NO_PREFETCH build macro */ -#if defined(XXH_NO_PREFETCH) -# define XXH_PREFETCH(ptr) (void)(ptr) /* disabled */ -#else -# if XXH_SIZE_OPT >= 1 -# define XXH_PREFETCH(ptr) (void)(ptr) -# elif defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)) /* _mm_prefetch() not defined outside of x86/x64 */ -# include /* https://msdn.microsoft.com/fr-fr/library/84szxsww(v=vs.90).aspx */ -# define XXH_PREFETCH(ptr) _mm_prefetch((const char*)(ptr), _MM_HINT_T0) -# elif defined(__GNUC__) && ( (__GNUC__ >= 4) || ( (__GNUC__ == 3) && (__GNUC_MINOR__ >= 1) ) ) -# define XXH_PREFETCH(ptr) __builtin_prefetch((ptr), 0 /* rw==read */, 3 /* locality */) -# else -# define XXH_PREFETCH(ptr) (void)(ptr) /* disabled */ -# endif -#endif /* XXH_NO_PREFETCH */ - - -/* ========================================== - * XXH3 default settings - * ========================================== */ - -#define XXH_SECRET_DEFAULT_SIZE 192 /* minimum XXH3_SECRET_SIZE_MIN */ - -#if (XXH_SECRET_DEFAULT_SIZE < XXH3_SECRET_SIZE_MIN) -# error "default keyset is not large enough" -#endif - -/*! Pseudorandom secret taken directly from FARSH. */ -XXH_ALIGN(64) static const xxh_u8 XXH3_kSecret[XXH_SECRET_DEFAULT_SIZE] = { - 0xb8, 0xfe, 0x6c, 0x39, 0x23, 0xa4, 0x4b, 0xbe, 0x7c, 0x01, 0x81, 0x2c, 0xf7, 0x21, 0xad, 0x1c, - 0xde, 0xd4, 0x6d, 0xe9, 0x83, 0x90, 0x97, 0xdb, 0x72, 0x40, 0xa4, 0xa4, 0xb7, 0xb3, 0x67, 0x1f, - 0xcb, 0x79, 0xe6, 0x4e, 0xcc, 0xc0, 0xe5, 0x78, 0x82, 0x5a, 0xd0, 0x7d, 0xcc, 0xff, 0x72, 0x21, - 0xb8, 0x08, 0x46, 0x74, 0xf7, 0x43, 0x24, 0x8e, 0xe0, 0x35, 0x90, 0xe6, 0x81, 0x3a, 0x26, 0x4c, - 0x3c, 0x28, 0x52, 0xbb, 0x91, 0xc3, 0x00, 0xcb, 0x88, 0xd0, 0x65, 0x8b, 0x1b, 0x53, 0x2e, 0xa3, - 0x71, 0x64, 0x48, 0x97, 0xa2, 0x0d, 0xf9, 0x4e, 0x38, 0x19, 0xef, 0x46, 0xa9, 0xde, 0xac, 0xd8, - 0xa8, 0xfa, 0x76, 0x3f, 0xe3, 0x9c, 0x34, 0x3f, 0xf9, 0xdc, 0xbb, 0xc7, 0xc7, 0x0b, 0x4f, 0x1d, - 0x8a, 0x51, 0xe0, 0x4b, 0xcd, 0xb4, 0x59, 0x31, 0xc8, 0x9f, 0x7e, 0xc9, 0xd9, 0x78, 0x73, 0x64, - 0xea, 0xc5, 0xac, 0x83, 0x34, 0xd3, 0xeb, 0xc3, 0xc5, 0x81, 0xa0, 0xff, 0xfa, 0x13, 0x63, 0xeb, - 0x17, 0x0d, 0xdd, 0x51, 0xb7, 0xf0, 0xda, 0x49, 0xd3, 0x16, 0x55, 0x26, 0x29, 0xd4, 0x68, 0x9e, - 0x2b, 0x16, 0xbe, 0x58, 0x7d, 0x47, 0xa1, 0xfc, 0x8f, 0xf8, 0xb8, 0xd1, 0x7a, 0xd0, 0x31, 0xce, - 0x45, 0xcb, 0x3a, 0x8f, 0x95, 0x16, 0x04, 0x28, 0xaf, 0xd7, 0xfb, 0xca, 0xbb, 0x4b, 0x40, 0x7e, -}; - -static const xxh_u64 PRIME_MX1 = 0x165667919E3779F9ULL; /*!< 0b0001011001010110011001111001000110011110001101110111100111111001 */ -static const xxh_u64 PRIME_MX2 = 0x9FB21C651E98DF25ULL; /*!< 0b1001111110110010000111000110010100011110100110001101111100100101 */ - -#ifdef XXH_OLD_NAMES -# define kSecret XXH3_kSecret -#endif - -#ifdef XXH_DOXYGEN -/*! - * @brief Calculates a 32-bit to 64-bit long multiply. - * - * Implemented as a macro. - * - * Wraps `__emulu` on MSVC x86 because it tends to call `__allmul` when it doesn't - * need to (but it shouldn't need to anyways, it is about 7 instructions to do - * a 64x64 multiply...). Since we know that this will _always_ emit `MULL`, we - * use that instead of the normal method. - * - * If you are compiling for platforms like Thumb-1 and don't have a better option, - * you may also want to write your own long multiply routine here. - * - * @param x, y Numbers to be multiplied - * @return 64-bit product of the low 32 bits of @p x and @p y. - */ -XXH_FORCE_INLINE xxh_u64 -XXH_mult32to64(xxh_u64 x, xxh_u64 y) -{ - return (x & 0xFFFFFFFF) * (y & 0xFFFFFFFF); -} -#elif defined(_MSC_VER) && defined(_M_IX86) -# define XXH_mult32to64(x, y) __emulu((unsigned)(x), (unsigned)(y)) -#else -/* - * Downcast + upcast is usually better than masking on older compilers like - * GCC 4.2 (especially 32-bit ones), all without affecting newer compilers. - * - * The other method, (x & 0xFFFFFFFF) * (y & 0xFFFFFFFF), will AND both operands - * and perform a full 64x64 multiply -- entirely redundant on 32-bit. - */ -# define XXH_mult32to64(x, y) ((xxh_u64)(xxh_u32)(x) * (xxh_u64)(xxh_u32)(y)) -#endif - -/*! - * @brief Calculates a 64->128-bit long multiply. - * - * Uses `__uint128_t` and `_umul128` if available, otherwise uses a scalar - * version. - * - * @param lhs , rhs The 64-bit integers to be multiplied - * @return The 128-bit result represented in an @ref XXH128_hash_t. - */ -static XXH128_hash_t -XXH_mult64to128(xxh_u64 lhs, xxh_u64 rhs) -{ - /* - * GCC/Clang __uint128_t method. - * - * On most 64-bit targets, GCC and Clang define a __uint128_t type. - * This is usually the best way as it usually uses a native long 64-bit - * multiply, such as MULQ on x86_64 or MUL + UMULH on aarch64. - * - * Usually. - * - * Despite being a 32-bit platform, Clang (and emscripten) define this type - * despite not having the arithmetic for it. This results in a laggy - * compiler builtin call which calculates a full 128-bit multiply. - * In that case it is best to use the portable one. - * https://github.com/Cyan4973/xxHash/issues/211#issuecomment-515575677 - */ -#if (defined(__GNUC__) || defined(__clang__)) && !defined(__wasm__) \ - && defined(__SIZEOF_INT128__) \ - || (defined(_INTEGRAL_MAX_BITS) && _INTEGRAL_MAX_BITS >= 128) - - __uint128_t const product = (__uint128_t)lhs * (__uint128_t)rhs; - XXH128_hash_t r128; - r128.low64 = (xxh_u64)(product); - r128.high64 = (xxh_u64)(product >> 64); - return r128; - - /* - * MSVC for x64's _umul128 method. - * - * xxh_u64 _umul128(xxh_u64 Multiplier, xxh_u64 Multiplicand, xxh_u64 *HighProduct); - * - * This compiles to single operand MUL on x64. - */ -#elif (defined(_M_X64) || defined(_M_IA64)) && !defined(_M_ARM64EC) - -#ifndef _MSC_VER -# pragma intrinsic(_umul128) -#endif - xxh_u64 product_high; - xxh_u64 const product_low = _umul128(lhs, rhs, &product_high); - XXH128_hash_t r128; - r128.low64 = product_low; - r128.high64 = product_high; - return r128; - - /* - * MSVC for ARM64's __umulh method. - * - * This compiles to the same MUL + UMULH as GCC/Clang's __uint128_t method. - */ -#elif defined(_M_ARM64) || defined(_M_ARM64EC) - -#ifndef _MSC_VER -# pragma intrinsic(__umulh) -#endif - XXH128_hash_t r128; - r128.low64 = lhs * rhs; - r128.high64 = __umulh(lhs, rhs); - return r128; - -#else - /* - * Portable scalar method. Optimized for 32-bit and 64-bit ALUs. - * - * This is a fast and simple grade school multiply, which is shown below - * with base 10 arithmetic instead of base 0x100000000. - * - * 9 3 // D2 lhs = 93 - * x 7 5 // D2 rhs = 75 - * ---------- - * 1 5 // D2 lo_lo = (93 % 10) * (75 % 10) = 15 - * 4 5 | // D2 hi_lo = (93 / 10) * (75 % 10) = 45 - * 2 1 | // D2 lo_hi = (93 % 10) * (75 / 10) = 21 - * + 6 3 | | // D2 hi_hi = (93 / 10) * (75 / 10) = 63 - * --------- - * 2 7 | // D2 cross = (15 / 10) + (45 % 10) + 21 = 27 - * + 6 7 | | // D2 upper = (27 / 10) + (45 / 10) + 63 = 67 - * --------- - * 6 9 7 5 // D4 res = (27 * 10) + (15 % 10) + (67 * 100) = 6975 - * - * The reasons for adding the products like this are: - * 1. It avoids manual carry tracking. Just like how - * (9 * 9) + 9 + 9 = 99, the same applies with this for UINT64_MAX. - * This avoids a lot of complexity. - * - * 2. It hints for, and on Clang, compiles to, the powerful UMAAL - * instruction available in ARM's Digital Signal Processing extension - * in 32-bit ARMv6 and later, which is shown below: - * - * void UMAAL(xxh_u32 *RdLo, xxh_u32 *RdHi, xxh_u32 Rn, xxh_u32 Rm) - * { - * xxh_u64 product = (xxh_u64)*RdLo * (xxh_u64)*RdHi + Rn + Rm; - * *RdLo = (xxh_u32)(product & 0xFFFFFFFF); - * *RdHi = (xxh_u32)(product >> 32); - * } - * - * This instruction was designed for efficient long multiplication, and - * allows this to be calculated in only 4 instructions at speeds - * comparable to some 64-bit ALUs. - * - * 3. It isn't terrible on other platforms. Usually this will be a couple - * of 32-bit ADD/ADCs. - */ - - /* First calculate all of the cross products. */ - xxh_u64 const lo_lo = XXH_mult32to64(lhs & 0xFFFFFFFF, rhs & 0xFFFFFFFF); - xxh_u64 const hi_lo = XXH_mult32to64(lhs >> 32, rhs & 0xFFFFFFFF); - xxh_u64 const lo_hi = XXH_mult32to64(lhs & 0xFFFFFFFF, rhs >> 32); - xxh_u64 const hi_hi = XXH_mult32to64(lhs >> 32, rhs >> 32); - - /* Now add the products together. These will never overflow. */ - xxh_u64 const cross = (lo_lo >> 32) + (hi_lo & 0xFFFFFFFF) + lo_hi; - xxh_u64 const upper = (hi_lo >> 32) + (cross >> 32) + hi_hi; - xxh_u64 const lower = (cross << 32) | (lo_lo & 0xFFFFFFFF); - - XXH128_hash_t r128; - r128.low64 = lower; - r128.high64 = upper; - return r128; -#endif -} - -/*! - * @brief Calculates a 64-bit to 128-bit multiply, then XOR folds it. - * - * The reason for the separate function is to prevent passing too many structs - * around by value. This will hopefully inline the multiply, but we don't force it. - * - * @param lhs , rhs The 64-bit integers to multiply - * @return The low 64 bits of the product XOR'd by the high 64 bits. - * @see XXH_mult64to128() - */ -static xxh_u64 -XXH3_mul128_fold64(xxh_u64 lhs, xxh_u64 rhs) -{ - XXH128_hash_t product = XXH_mult64to128(lhs, rhs); - return product.low64 ^ product.high64; -} - -/*! Seems to produce slightly better code on GCC for some reason. */ -XXH_FORCE_INLINE XXH_CONSTF xxh_u64 XXH_xorshift64(xxh_u64 v64, int shift) -{ - XXH_ASSERT(0 <= shift && shift < 64); - return v64 ^ (v64 >> shift); -} - -/* - * This is a fast avalanche stage, - * suitable when input bits are already partially mixed - */ -static XXH64_hash_t XXH3_avalanche(xxh_u64 h64) -{ - h64 = XXH_xorshift64(h64, 37); - h64 *= PRIME_MX1; - h64 = XXH_xorshift64(h64, 32); - return h64; -} - -/* - * This is a stronger avalanche, - * inspired by Pelle Evensen's rrmxmx - * preferable when input has not been previously mixed - */ -static XXH64_hash_t XXH3_rrmxmx(xxh_u64 h64, xxh_u64 len) -{ - /* this mix is inspired by Pelle Evensen's rrmxmx */ - h64 ^= XXH_rotl64(h64, 49) ^ XXH_rotl64(h64, 24); - h64 *= PRIME_MX2; - h64 ^= (h64 >> 35) + len ; - h64 *= PRIME_MX2; - return XXH_xorshift64(h64, 28); -} - - -/* ========================================== - * Short keys - * ========================================== - * One of the shortcomings of XXH32 and XXH64 was that their performance was - * sub-optimal on short lengths. It used an iterative algorithm which strongly - * favored lengths that were a multiple of 4 or 8. - * - * Instead of iterating over individual inputs, we use a set of single shot - * functions which piece together a range of lengths and operate in constant time. - * - * Additionally, the number of multiplies has been significantly reduced. This - * reduces latency, especially when emulating 64-bit multiplies on 32-bit. - * - * Depending on the platform, this may or may not be faster than XXH32, but it - * is almost guaranteed to be faster than XXH64. - */ - -/* - * At very short lengths, there isn't enough input to fully hide secrets, or use - * the entire secret. - * - * There is also only a limited amount of mixing we can do before significantly - * impacting performance. - * - * Therefore, we use different sections of the secret and always mix two secret - * samples with an XOR. This should have no effect on performance on the - * seedless or withSeed variants because everything _should_ be constant folded - * by modern compilers. - * - * The XOR mixing hides individual parts of the secret and increases entropy. - * - * This adds an extra layer of strength for custom secrets. - */ -XXH_FORCE_INLINE XXH_PUREF XXH64_hash_t -XXH3_len_1to3_64b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(input != NULL); - XXH_ASSERT(1 <= len && len <= 3); - XXH_ASSERT(secret != NULL); - /* - * len = 1: combined = { input[0], 0x01, input[0], input[0] } - * len = 2: combined = { input[1], 0x02, input[0], input[1] } - * len = 3: combined = { input[2], 0x03, input[0], input[1] } - */ - { xxh_u8 const c1 = input[0]; - xxh_u8 const c2 = input[len >> 1]; - xxh_u8 const c3 = input[len - 1]; - xxh_u32 const combined = ((xxh_u32)c1 << 16) | ((xxh_u32)c2 << 24) - | ((xxh_u32)c3 << 0) | ((xxh_u32)len << 8); - xxh_u64 const bitflip = (XXH_readLE32(secret) ^ XXH_readLE32(secret+4)) + seed; - xxh_u64 const keyed = (xxh_u64)combined ^ bitflip; - return XXH64_avalanche(keyed); - } -} - -XXH_FORCE_INLINE XXH_PUREF XXH64_hash_t -XXH3_len_4to8_64b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(input != NULL); - XXH_ASSERT(secret != NULL); - XXH_ASSERT(4 <= len && len <= 8); - seed ^= (xxh_u64)XXH_swap32((xxh_u32)seed) << 32; - { xxh_u32 const input1 = XXH_readLE32(input); - xxh_u32 const input2 = XXH_readLE32(input + len - 4); - xxh_u64 const bitflip = (XXH_readLE64(secret+8) ^ XXH_readLE64(secret+16)) - seed; - xxh_u64 const input64 = input2 + (((xxh_u64)input1) << 32); - xxh_u64 const keyed = input64 ^ bitflip; - return XXH3_rrmxmx(keyed, len); - } -} - -XXH_FORCE_INLINE XXH_PUREF XXH64_hash_t -XXH3_len_9to16_64b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(input != NULL); - XXH_ASSERT(secret != NULL); - XXH_ASSERT(9 <= len && len <= 16); - { xxh_u64 const bitflip1 = (XXH_readLE64(secret+24) ^ XXH_readLE64(secret+32)) + seed; - xxh_u64 const bitflip2 = (XXH_readLE64(secret+40) ^ XXH_readLE64(secret+48)) - seed; - xxh_u64 const input_lo = XXH_readLE64(input) ^ bitflip1; - xxh_u64 const input_hi = XXH_readLE64(input + len - 8) ^ bitflip2; - xxh_u64 const acc = len - + XXH_swap64(input_lo) + input_hi - + XXH3_mul128_fold64(input_lo, input_hi); - return XXH3_avalanche(acc); - } -} - -XXH_FORCE_INLINE XXH_PUREF XXH64_hash_t -XXH3_len_0to16_64b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(len <= 16); - { if (XXH_likely(len > 8)) return XXH3_len_9to16_64b(input, len, secret, seed); - if (XXH_likely(len >= 4)) return XXH3_len_4to8_64b(input, len, secret, seed); - if (len) return XXH3_len_1to3_64b(input, len, secret, seed); - return XXH64_avalanche(seed ^ (XXH_readLE64(secret+56) ^ XXH_readLE64(secret+64))); - } -} - -/* - * DISCLAIMER: There are known *seed-dependent* multicollisions here due to - * multiplication by zero, affecting hashes of lengths 17 to 240. - * - * However, they are very unlikely. - * - * Keep this in mind when using the unseeded XXH3_64bits() variant: As with all - * unseeded non-cryptographic hashes, it does not attempt to defend itself - * against specially crafted inputs, only random inputs. - * - * Compared to classic UMAC where a 1 in 2^31 chance of 4 consecutive bytes - * cancelling out the secret is taken an arbitrary number of times (addressed - * in XXH3_accumulate_512), this collision is very unlikely with random inputs - * and/or proper seeding: - * - * This only has a 1 in 2^63 chance of 8 consecutive bytes cancelling out, in a - * function that is only called up to 16 times per hash with up to 240 bytes of - * input. - * - * This is not too bad for a non-cryptographic hash function, especially with - * only 64 bit outputs. - * - * The 128-bit variant (which trades some speed for strength) is NOT affected - * by this, although it is always a good idea to use a proper seed if you care - * about strength. - */ -XXH_FORCE_INLINE xxh_u64 XXH3_mix16B(const xxh_u8* XXH_RESTRICT input, - const xxh_u8* XXH_RESTRICT secret, xxh_u64 seed64) -{ -#if defined(__GNUC__) && !defined(__clang__) /* GCC, not Clang */ \ - && defined(__i386__) && defined(__SSE2__) /* x86 + SSE2 */ \ - && !defined(XXH_ENABLE_AUTOVECTORIZE) /* Define to disable like XXH32 hack */ - /* - * UGLY HACK: - * GCC for x86 tends to autovectorize the 128-bit multiply, resulting in - * slower code. - * - * By forcing seed64 into a register, we disrupt the cost model and - * cause it to scalarize. See `XXH32_round()` - * - * FIXME: Clang's output is still _much_ faster -- On an AMD Ryzen 3600, - * XXH3_64bits @ len=240 runs at 4.6 GB/s with Clang 9, but 3.3 GB/s on - * GCC 9.2, despite both emitting scalar code. - * - * GCC generates much better scalar code than Clang for the rest of XXH3, - * which is why finding a more optimal codepath is an interest. - */ - XXH_COMPILER_GUARD(seed64); -#endif - { xxh_u64 const input_lo = XXH_readLE64(input); - xxh_u64 const input_hi = XXH_readLE64(input+8); - return XXH3_mul128_fold64( - input_lo ^ (XXH_readLE64(secret) + seed64), - input_hi ^ (XXH_readLE64(secret+8) - seed64) - ); - } -} - -/* For mid range keys, XXH3 uses a Mum-hash variant. */ -XXH_FORCE_INLINE XXH_PUREF XXH64_hash_t -XXH3_len_17to128_64b(const xxh_u8* XXH_RESTRICT input, size_t len, - const xxh_u8* XXH_RESTRICT secret, size_t secretSize, - XXH64_hash_t seed) -{ - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); (void)secretSize; - XXH_ASSERT(16 < len && len <= 128); - - { xxh_u64 acc = len * XXH_PRIME64_1; -#if XXH_SIZE_OPT >= 1 - /* Smaller and cleaner, but slightly slower. */ - unsigned int i = (unsigned int)(len - 1) / 32; - do { - acc += XXH3_mix16B(input+16 * i, secret+32*i, seed); - acc += XXH3_mix16B(input+len-16*(i+1), secret+32*i+16, seed); - } while (i-- != 0); -#else - if (len > 32) { - if (len > 64) { - if (len > 96) { - acc += XXH3_mix16B(input+48, secret+96, seed); - acc += XXH3_mix16B(input+len-64, secret+112, seed); - } - acc += XXH3_mix16B(input+32, secret+64, seed); - acc += XXH3_mix16B(input+len-48, secret+80, seed); - } - acc += XXH3_mix16B(input+16, secret+32, seed); - acc += XXH3_mix16B(input+len-32, secret+48, seed); - } - acc += XXH3_mix16B(input+0, secret+0, seed); - acc += XXH3_mix16B(input+len-16, secret+16, seed); -#endif - return XXH3_avalanche(acc); - } -} - -XXH_NO_INLINE XXH_PUREF XXH64_hash_t -XXH3_len_129to240_64b(const xxh_u8* XXH_RESTRICT input, size_t len, - const xxh_u8* XXH_RESTRICT secret, size_t secretSize, - XXH64_hash_t seed) -{ - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); (void)secretSize; - XXH_ASSERT(128 < len && len <= XXH3_MIDSIZE_MAX); - - #define XXH3_MIDSIZE_STARTOFFSET 3 - #define XXH3_MIDSIZE_LASTOFFSET 17 - - { xxh_u64 acc = len * XXH_PRIME64_1; - xxh_u64 acc_end; - unsigned int const nbRounds = (unsigned int)len / 16; - unsigned int i; - XXH_ASSERT(128 < len && len <= XXH3_MIDSIZE_MAX); - for (i=0; i<8; i++) { - acc += XXH3_mix16B(input+(16*i), secret+(16*i), seed); - } - /* last bytes */ - acc_end = XXH3_mix16B(input + len - 16, secret + XXH3_SECRET_SIZE_MIN - XXH3_MIDSIZE_LASTOFFSET, seed); - XXH_ASSERT(nbRounds >= 8); - acc = XXH3_avalanche(acc); -#if defined(__clang__) /* Clang */ \ - && (defined(__ARM_NEON) || defined(__ARM_NEON__)) /* NEON */ \ - && !defined(XXH_ENABLE_AUTOVECTORIZE) /* Define to disable */ - /* - * UGLY HACK: - * Clang for ARMv7-A tries to vectorize this loop, similar to GCC x86. - * In everywhere else, it uses scalar code. - * - * For 64->128-bit multiplies, even if the NEON was 100% optimal, it - * would still be slower than UMAAL (see XXH_mult64to128). - * - * Unfortunately, Clang doesn't handle the long multiplies properly and - * converts them to the nonexistent "vmulq_u64" intrinsic, which is then - * scalarized into an ugly mess of VMOV.32 instructions. - * - * This mess is difficult to avoid without turning autovectorization - * off completely, but they are usually relatively minor and/or not - * worth it to fix. - * - * This loop is the easiest to fix, as unlike XXH32, this pragma - * _actually works_ because it is a loop vectorization instead of an - * SLP vectorization. - */ - #pragma clang loop vectorize(disable) -#endif - for (i=8 ; i < nbRounds; i++) { - /* - * Prevents clang for unrolling the acc loop and interleaving with this one. - */ - XXH_COMPILER_GUARD(acc); - acc_end += XXH3_mix16B(input+(16*i), secret+(16*(i-8)) + XXH3_MIDSIZE_STARTOFFSET, seed); - } - return XXH3_avalanche(acc + acc_end); - } -} - - -/* ======= Long Keys ======= */ - -#define XXH_STRIPE_LEN 64 -#define XXH_SECRET_CONSUME_RATE 8 /* nb of secret bytes consumed at each accumulation */ -#define XXH_ACC_NB (XXH_STRIPE_LEN / sizeof(xxh_u64)) - -#ifdef XXH_OLD_NAMES -# define STRIPE_LEN XXH_STRIPE_LEN -# define ACC_NB XXH_ACC_NB -#endif - -#ifndef XXH_PREFETCH_DIST -# ifdef __clang__ -# define XXH_PREFETCH_DIST 320 -# else -# if (XXH_VECTOR == XXH_AVX512) -# define XXH_PREFETCH_DIST 512 -# else -# define XXH_PREFETCH_DIST 384 -# endif -# endif /* __clang__ */ -#endif /* XXH_PREFETCH_DIST */ - -/* - * These macros are to generate an XXH3_accumulate() function. - * The two arguments select the name suffix and target attribute. - * - * The name of this symbol is XXH3_accumulate_() and it calls - * XXH3_accumulate_512_(). - * - * It may be useful to hand implement this function if the compiler fails to - * optimize the inline function. - */ -#define XXH3_ACCUMULATE_TEMPLATE(name) \ -void \ -XXH3_accumulate_##name(xxh_u64* XXH_RESTRICT acc, \ - const xxh_u8* XXH_RESTRICT input, \ - const xxh_u8* XXH_RESTRICT secret, \ - size_t nbStripes) \ -{ \ - size_t n; \ - for (n = 0; n < nbStripes; n++ ) { \ - const xxh_u8* const in = input + n*XXH_STRIPE_LEN; \ - XXH_PREFETCH(in + XXH_PREFETCH_DIST); \ - XXH3_accumulate_512_##name( \ - acc, \ - in, \ - secret + n*XXH_SECRET_CONSUME_RATE); \ - } \ -} - - -XXH_FORCE_INLINE void XXH_writeLE64(void* dst, xxh_u64 v64) -{ - if (!XXH_CPU_LITTLE_ENDIAN) v64 = XXH_swap64(v64); - XXH_memcpy(dst, &v64, sizeof(v64)); -} - -/* Several intrinsic functions below are supposed to accept __int64 as argument, - * as documented in https://software.intel.com/sites/landingpage/IntrinsicsGuide/ . - * However, several environments do not define __int64 type, - * requiring a workaround. - */ -#if !defined (__VMS) \ - && (defined (__cplusplus) \ - || (defined (__STDC_VERSION__) && (__STDC_VERSION__ >= 199901L) /* C99 */) ) - typedef int64_t xxh_i64; -#else - /* the following type must have a width of 64-bit */ - typedef long long xxh_i64; -#endif - - -/* - * XXH3_accumulate_512 is the tightest loop for long inputs, and it is the most optimized. - * - * It is a hardened version of UMAC, based off of FARSH's implementation. - * - * This was chosen because it adapts quite well to 32-bit, 64-bit, and SIMD - * implementations, and it is ridiculously fast. - * - * We harden it by mixing the original input to the accumulators as well as the product. - * - * This means that in the (relatively likely) case of a multiply by zero, the - * original input is preserved. - * - * On 128-bit inputs, we swap 64-bit pairs when we add the input to improve - * cross-pollination, as otherwise the upper and lower halves would be - * essentially independent. - * - * This doesn't matter on 64-bit hashes since they all get merged together in - * the end, so we skip the extra step. - * - * Both XXH3_64bits and XXH3_128bits use this subroutine. - */ - -#if (XXH_VECTOR == XXH_AVX512) \ - || (defined(XXH_DISPATCH_AVX512) && XXH_DISPATCH_AVX512 != 0) - -#ifndef XXH_TARGET_AVX512 -# define XXH_TARGET_AVX512 /* disable attribute target */ -#endif - -XXH_FORCE_INLINE XXH_TARGET_AVX512 void -XXH3_accumulate_512_avx512(void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - __m512i* const xacc = (__m512i *) acc; - XXH_ASSERT((((size_t)acc) & 63) == 0); - XXH_STATIC_ASSERT(XXH_STRIPE_LEN == sizeof(__m512i)); - - { - /* data_vec = input[0]; */ - __m512i const data_vec = _mm512_loadu_si512 (input); - /* key_vec = secret[0]; */ - __m512i const key_vec = _mm512_loadu_si512 (secret); - /* data_key = data_vec ^ key_vec; */ - __m512i const data_key = _mm512_xor_si512 (data_vec, key_vec); - /* data_key_lo = data_key >> 32; */ - __m512i const data_key_lo = _mm512_srli_epi64 (data_key, 32); - /* product = (data_key & 0xffffffff) * (data_key_lo & 0xffffffff); */ - __m512i const product = _mm512_mul_epu32 (data_key, data_key_lo); - /* xacc[0] += swap(data_vec); */ - __m512i const data_swap = _mm512_shuffle_epi32(data_vec, (_MM_PERM_ENUM)_MM_SHUFFLE(1, 0, 3, 2)); - __m512i const sum = _mm512_add_epi64(*xacc, data_swap); - /* xacc[0] += product; */ - *xacc = _mm512_add_epi64(product, sum); - } -} -XXH_FORCE_INLINE XXH_TARGET_AVX512 XXH3_ACCUMULATE_TEMPLATE(avx512) - -/* - * XXH3_scrambleAcc: Scrambles the accumulators to improve mixing. - * - * Multiplication isn't perfect, as explained by Google in HighwayHash: - * - * // Multiplication mixes/scrambles bytes 0-7 of the 64-bit result to - * // varying degrees. In descending order of goodness, bytes - * // 3 4 2 5 1 6 0 7 have quality 228 224 164 160 100 96 36 32. - * // As expected, the upper and lower bytes are much worse. - * - * Source: https://github.com/google/highwayhash/blob/0aaf66b/highwayhash/hh_avx2.h#L291 - * - * Since our algorithm uses a pseudorandom secret to add some variance into the - * mix, we don't need to (or want to) mix as often or as much as HighwayHash does. - * - * This isn't as tight as XXH3_accumulate, but still written in SIMD to avoid - * extraction. - * - * Both XXH3_64bits and XXH3_128bits use this subroutine. - */ - -XXH_FORCE_INLINE XXH_TARGET_AVX512 void -XXH3_scrambleAcc_avx512(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 63) == 0); - XXH_STATIC_ASSERT(XXH_STRIPE_LEN == sizeof(__m512i)); - { __m512i* const xacc = (__m512i*) acc; - const __m512i prime32 = _mm512_set1_epi32((int)XXH_PRIME32_1); - - /* xacc[0] ^= (xacc[0] >> 47) */ - __m512i const acc_vec = *xacc; - __m512i const shifted = _mm512_srli_epi64 (acc_vec, 47); - /* xacc[0] ^= secret; */ - __m512i const key_vec = _mm512_loadu_si512 (secret); - __m512i const data_key = _mm512_ternarylogic_epi32(key_vec, acc_vec, shifted, 0x96 /* key_vec ^ acc_vec ^ shifted */); - - /* xacc[0] *= XXH_PRIME32_1; */ - __m512i const data_key_hi = _mm512_srli_epi64 (data_key, 32); - __m512i const prod_lo = _mm512_mul_epu32 (data_key, prime32); - __m512i const prod_hi = _mm512_mul_epu32 (data_key_hi, prime32); - *xacc = _mm512_add_epi64(prod_lo, _mm512_slli_epi64(prod_hi, 32)); - } -} - -XXH_FORCE_INLINE XXH_TARGET_AVX512 void -XXH3_initCustomSecret_avx512(void* XXH_RESTRICT customSecret, xxh_u64 seed64) -{ - XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE & 63) == 0); - XXH_STATIC_ASSERT(XXH_SEC_ALIGN == 64); - XXH_ASSERT(((size_t)customSecret & 63) == 0); - (void)(&XXH_writeLE64); - { int const nbRounds = XXH_SECRET_DEFAULT_SIZE / sizeof(__m512i); - __m512i const seed_pos = _mm512_set1_epi64((xxh_i64)seed64); - __m512i const seed = _mm512_mask_sub_epi64(seed_pos, 0xAA, _mm512_set1_epi8(0), seed_pos); - - const __m512i* const src = (const __m512i*) ((const void*) XXH3_kSecret); - __m512i* const dest = ( __m512i*) customSecret; - int i; - XXH_ASSERT(((size_t)src & 63) == 0); /* control alignment */ - XXH_ASSERT(((size_t)dest & 63) == 0); - for (i=0; i < nbRounds; ++i) { - dest[i] = _mm512_add_epi64(_mm512_load_si512(src + i), seed); - } } -} - -#endif - -#if (XXH_VECTOR == XXH_AVX2) \ - || (defined(XXH_DISPATCH_AVX2) && XXH_DISPATCH_AVX2 != 0) - -#ifndef XXH_TARGET_AVX2 -# define XXH_TARGET_AVX2 /* disable attribute target */ -#endif - -XXH_FORCE_INLINE XXH_TARGET_AVX2 void -XXH3_accumulate_512_avx2( void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 31) == 0); - { __m256i* const xacc = (__m256i *) acc; - /* Unaligned. This is mainly for pointer arithmetic, and because - * _mm256_loadu_si256 requires a const __m256i * pointer for some reason. */ - const __m256i* const xinput = (const __m256i *) input; - /* Unaligned. This is mainly for pointer arithmetic, and because - * _mm256_loadu_si256 requires a const __m256i * pointer for some reason. */ - const __m256i* const xsecret = (const __m256i *) secret; - - size_t i; - for (i=0; i < XXH_STRIPE_LEN/sizeof(__m256i); i++) { - /* data_vec = xinput[i]; */ - __m256i const data_vec = _mm256_loadu_si256 (xinput+i); - /* key_vec = xsecret[i]; */ - __m256i const key_vec = _mm256_loadu_si256 (xsecret+i); - /* data_key = data_vec ^ key_vec; */ - __m256i const data_key = _mm256_xor_si256 (data_vec, key_vec); - /* data_key_lo = data_key >> 32; */ - __m256i const data_key_lo = _mm256_srli_epi64 (data_key, 32); - /* product = (data_key & 0xffffffff) * (data_key_lo & 0xffffffff); */ - __m256i const product = _mm256_mul_epu32 (data_key, data_key_lo); - /* xacc[i] += swap(data_vec); */ - __m256i const data_swap = _mm256_shuffle_epi32(data_vec, _MM_SHUFFLE(1, 0, 3, 2)); - __m256i const sum = _mm256_add_epi64(xacc[i], data_swap); - /* xacc[i] += product; */ - xacc[i] = _mm256_add_epi64(product, sum); - } } -} -XXH_FORCE_INLINE XXH_TARGET_AVX2 XXH3_ACCUMULATE_TEMPLATE(avx2) - -XXH_FORCE_INLINE XXH_TARGET_AVX2 void -XXH3_scrambleAcc_avx2(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 31) == 0); - { __m256i* const xacc = (__m256i*) acc; - /* Unaligned. This is mainly for pointer arithmetic, and because - * _mm256_loadu_si256 requires a const __m256i * pointer for some reason. */ - const __m256i* const xsecret = (const __m256i *) secret; - const __m256i prime32 = _mm256_set1_epi32((int)XXH_PRIME32_1); - - size_t i; - for (i=0; i < XXH_STRIPE_LEN/sizeof(__m256i); i++) { - /* xacc[i] ^= (xacc[i] >> 47) */ - __m256i const acc_vec = xacc[i]; - __m256i const shifted = _mm256_srli_epi64 (acc_vec, 47); - __m256i const data_vec = _mm256_xor_si256 (acc_vec, shifted); - /* xacc[i] ^= xsecret; */ - __m256i const key_vec = _mm256_loadu_si256 (xsecret+i); - __m256i const data_key = _mm256_xor_si256 (data_vec, key_vec); - - /* xacc[i] *= XXH_PRIME32_1; */ - __m256i const data_key_hi = _mm256_srli_epi64 (data_key, 32); - __m256i const prod_lo = _mm256_mul_epu32 (data_key, prime32); - __m256i const prod_hi = _mm256_mul_epu32 (data_key_hi, prime32); - xacc[i] = _mm256_add_epi64(prod_lo, _mm256_slli_epi64(prod_hi, 32)); - } - } -} - -XXH_FORCE_INLINE XXH_TARGET_AVX2 void XXH3_initCustomSecret_avx2(void* XXH_RESTRICT customSecret, xxh_u64 seed64) -{ - XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE & 31) == 0); - XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE / sizeof(__m256i)) == 6); - XXH_STATIC_ASSERT(XXH_SEC_ALIGN <= 64); - (void)(&XXH_writeLE64); - XXH_PREFETCH(customSecret); - { __m256i const seed = _mm256_set_epi64x((xxh_i64)(0U - seed64), (xxh_i64)seed64, (xxh_i64)(0U - seed64), (xxh_i64)seed64); - - const __m256i* const src = (const __m256i*) ((const void*) XXH3_kSecret); - __m256i* dest = ( __m256i*) customSecret; - -# if defined(__GNUC__) || defined(__clang__) - /* - * On GCC & Clang, marking 'dest' as modified will cause the compiler: - * - do not extract the secret from sse registers in the internal loop - * - use less common registers, and avoid pushing these reg into stack - */ - XXH_COMPILER_GUARD(dest); -# endif - XXH_ASSERT(((size_t)src & 31) == 0); /* control alignment */ - XXH_ASSERT(((size_t)dest & 31) == 0); - - /* GCC -O2 need unroll loop manually */ - dest[0] = _mm256_add_epi64(_mm256_load_si256(src+0), seed); - dest[1] = _mm256_add_epi64(_mm256_load_si256(src+1), seed); - dest[2] = _mm256_add_epi64(_mm256_load_si256(src+2), seed); - dest[3] = _mm256_add_epi64(_mm256_load_si256(src+3), seed); - dest[4] = _mm256_add_epi64(_mm256_load_si256(src+4), seed); - dest[5] = _mm256_add_epi64(_mm256_load_si256(src+5), seed); - } -} - -#endif - -/* x86dispatch always generates SSE2 */ -#if (XXH_VECTOR == XXH_SSE2) || defined(XXH_X86DISPATCH) - -#ifndef XXH_TARGET_SSE2 -# define XXH_TARGET_SSE2 /* disable attribute target */ -#endif - -XXH_FORCE_INLINE XXH_TARGET_SSE2 void -XXH3_accumulate_512_sse2( void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - /* SSE2 is just a half-scale version of the AVX2 version. */ - XXH_ASSERT((((size_t)acc) & 15) == 0); - { __m128i* const xacc = (__m128i *) acc; - /* Unaligned. This is mainly for pointer arithmetic, and because - * _mm_loadu_si128 requires a const __m128i * pointer for some reason. */ - const __m128i* const xinput = (const __m128i *) input; - /* Unaligned. This is mainly for pointer arithmetic, and because - * _mm_loadu_si128 requires a const __m128i * pointer for some reason. */ - const __m128i* const xsecret = (const __m128i *) secret; - - size_t i; - for (i=0; i < XXH_STRIPE_LEN/sizeof(__m128i); i++) { - /* data_vec = xinput[i]; */ - __m128i const data_vec = _mm_loadu_si128 (xinput+i); - /* key_vec = xsecret[i]; */ - __m128i const key_vec = _mm_loadu_si128 (xsecret+i); - /* data_key = data_vec ^ key_vec; */ - __m128i const data_key = _mm_xor_si128 (data_vec, key_vec); - /* data_key_lo = data_key >> 32; */ - __m128i const data_key_lo = _mm_shuffle_epi32 (data_key, _MM_SHUFFLE(0, 3, 0, 1)); - /* product = (data_key & 0xffffffff) * (data_key_lo & 0xffffffff); */ - __m128i const product = _mm_mul_epu32 (data_key, data_key_lo); - /* xacc[i] += swap(data_vec); */ - __m128i const data_swap = _mm_shuffle_epi32(data_vec, _MM_SHUFFLE(1,0,3,2)); - __m128i const sum = _mm_add_epi64(xacc[i], data_swap); - /* xacc[i] += product; */ - xacc[i] = _mm_add_epi64(product, sum); - } } -} -XXH_FORCE_INLINE XXH_TARGET_SSE2 XXH3_ACCUMULATE_TEMPLATE(sse2) - -XXH_FORCE_INLINE XXH_TARGET_SSE2 void -XXH3_scrambleAcc_sse2(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 15) == 0); - { __m128i* const xacc = (__m128i*) acc; - /* Unaligned. This is mainly for pointer arithmetic, and because - * _mm_loadu_si128 requires a const __m128i * pointer for some reason. */ - const __m128i* const xsecret = (const __m128i *) secret; - const __m128i prime32 = _mm_set1_epi32((int)XXH_PRIME32_1); - - size_t i; - for (i=0; i < XXH_STRIPE_LEN/sizeof(__m128i); i++) { - /* xacc[i] ^= (xacc[i] >> 47) */ - __m128i const acc_vec = xacc[i]; - __m128i const shifted = _mm_srli_epi64 (acc_vec, 47); - __m128i const data_vec = _mm_xor_si128 (acc_vec, shifted); - /* xacc[i] ^= xsecret[i]; */ - __m128i const key_vec = _mm_loadu_si128 (xsecret+i); - __m128i const data_key = _mm_xor_si128 (data_vec, key_vec); - - /* xacc[i] *= XXH_PRIME32_1; */ - __m128i const data_key_hi = _mm_shuffle_epi32 (data_key, _MM_SHUFFLE(0, 3, 0, 1)); - __m128i const prod_lo = _mm_mul_epu32 (data_key, prime32); - __m128i const prod_hi = _mm_mul_epu32 (data_key_hi, prime32); - xacc[i] = _mm_add_epi64(prod_lo, _mm_slli_epi64(prod_hi, 32)); - } - } -} - -XXH_FORCE_INLINE XXH_TARGET_SSE2 void XXH3_initCustomSecret_sse2(void* XXH_RESTRICT customSecret, xxh_u64 seed64) -{ - XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE & 15) == 0); - (void)(&XXH_writeLE64); - { int const nbRounds = XXH_SECRET_DEFAULT_SIZE / sizeof(__m128i); - -# if defined(_MSC_VER) && defined(_M_IX86) && _MSC_VER < 1900 - /* MSVC 32bit mode does not support _mm_set_epi64x before 2015 */ - XXH_ALIGN(16) const xxh_i64 seed64x2[2] = { (xxh_i64)seed64, (xxh_i64)(0U - seed64) }; - __m128i const seed = _mm_load_si128((__m128i const*)seed64x2); -# else - __m128i const seed = _mm_set_epi64x((xxh_i64)(0U - seed64), (xxh_i64)seed64); -# endif - int i; - - const void* const src16 = XXH3_kSecret; - __m128i* dst16 = (__m128i*) customSecret; -# if defined(__GNUC__) || defined(__clang__) - /* - * On GCC & Clang, marking 'dest' as modified will cause the compiler: - * - do not extract the secret from sse registers in the internal loop - * - use less common registers, and avoid pushing these reg into stack - */ - XXH_COMPILER_GUARD(dst16); -# endif - XXH_ASSERT(((size_t)src16 & 15) == 0); /* control alignment */ - XXH_ASSERT(((size_t)dst16 & 15) == 0); - - for (i=0; i < nbRounds; ++i) { - dst16[i] = _mm_add_epi64(_mm_load_si128((const __m128i *)src16+i), seed); - } } -} - -#endif - -#if (XXH_VECTOR == XXH_NEON) - -/* forward declarations for the scalar routines */ -XXH_FORCE_INLINE void -XXH3_scalarRound(void* XXH_RESTRICT acc, void const* XXH_RESTRICT input, - void const* XXH_RESTRICT secret, size_t lane); - -XXH_FORCE_INLINE void -XXH3_scalarScrambleRound(void* XXH_RESTRICT acc, - void const* XXH_RESTRICT secret, size_t lane); - -/*! - * @internal - * @brief The bulk processing loop for NEON and WASM SIMD128. - * - * The NEON code path is actually partially scalar when running on AArch64. This - * is to optimize the pipelining and can have up to 15% speedup depending on the - * CPU, and it also mitigates some GCC codegen issues. - * - * @see XXH3_NEON_LANES for configuring this and details about this optimization. - * - * NEON's 32-bit to 64-bit long multiply takes a half vector of 32-bit - * integers instead of the other platforms which mask full 64-bit vectors, - * so the setup is more complicated than just shifting right. - * - * Additionally, there is an optimization for 4 lanes at once noted below. - * - * Since, as stated, the most optimal amount of lanes for Cortexes is 6, - * there needs to be *three* versions of the accumulate operation used - * for the remaining 2 lanes. - * - * WASM's SIMD128 uses SIMDe's arm_neon.h polyfill because the intrinsics overlap - * nearly perfectly. - */ - -XXH_FORCE_INLINE void -XXH3_accumulate_512_neon( void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 15) == 0); - XXH_STATIC_ASSERT(XXH3_NEON_LANES > 0 && XXH3_NEON_LANES <= XXH_ACC_NB && XXH3_NEON_LANES % 2 == 0); - { /* GCC for darwin arm64 does not like aliasing here */ - xxh_aliasing_uint64x2_t* const xacc = (xxh_aliasing_uint64x2_t*) acc; - /* We don't use a uint32x4_t pointer because it causes bus errors on ARMv7. */ - uint8_t const* xinput = (const uint8_t *) input; - uint8_t const* xsecret = (const uint8_t *) secret; - - size_t i; -#ifdef __wasm_simd128__ - /* - * On WASM SIMD128, Clang emits direct address loads when XXH3_kSecret - * is constant propagated, which results in it converting it to this - * inside the loop: - * - * a = v128.load(XXH3_kSecret + 0 + $secret_offset, offset = 0) - * b = v128.load(XXH3_kSecret + 16 + $secret_offset, offset = 0) - * ... - * - * This requires a full 32-bit address immediate (and therefore a 6 byte - * instruction) as well as an add for each offset. - * - * Putting an asm guard prevents it from folding (at the cost of losing - * the alignment hint), and uses the free offset in `v128.load` instead - * of adding secret_offset each time which overall reduces code size by - * about a kilobyte and improves performance. - */ - XXH_COMPILER_GUARD(xsecret); -#endif - /* Scalar lanes use the normal scalarRound routine */ - for (i = XXH3_NEON_LANES; i < XXH_ACC_NB; i++) { - XXH3_scalarRound(acc, input, secret, i); - } - i = 0; - /* 4 NEON lanes at a time. */ - for (; i+1 < XXH3_NEON_LANES / 2; i+=2) { - /* data_vec = xinput[i]; */ - uint64x2_t data_vec_1 = XXH_vld1q_u64(xinput + (i * 16)); - uint64x2_t data_vec_2 = XXH_vld1q_u64(xinput + ((i+1) * 16)); - /* key_vec = xsecret[i]; */ - uint64x2_t key_vec_1 = XXH_vld1q_u64(xsecret + (i * 16)); - uint64x2_t key_vec_2 = XXH_vld1q_u64(xsecret + ((i+1) * 16)); - /* data_swap = swap(data_vec) */ - uint64x2_t data_swap_1 = vextq_u64(data_vec_1, data_vec_1, 1); - uint64x2_t data_swap_2 = vextq_u64(data_vec_2, data_vec_2, 1); - /* data_key = data_vec ^ key_vec; */ - uint64x2_t data_key_1 = veorq_u64(data_vec_1, key_vec_1); - uint64x2_t data_key_2 = veorq_u64(data_vec_2, key_vec_2); - - /* - * If we reinterpret the 64x2 vectors as 32x4 vectors, we can use a - * de-interleave operation for 4 lanes in 1 step with `vuzpq_u32` to - * get one vector with the low 32 bits of each lane, and one vector - * with the high 32 bits of each lane. - * - * The intrinsic returns a double vector because the original ARMv7-a - * instruction modified both arguments in place. AArch64 and SIMD128 emit - * two instructions from this intrinsic. - * - * [ dk11L | dk11H | dk12L | dk12H ] -> [ dk11L | dk12L | dk21L | dk22L ] - * [ dk21L | dk21H | dk22L | dk22H ] -> [ dk11H | dk12H | dk21H | dk22H ] - */ - uint32x4x2_t unzipped = vuzpq_u32( - vreinterpretq_u32_u64(data_key_1), - vreinterpretq_u32_u64(data_key_2) - ); - /* data_key_lo = data_key & 0xFFFFFFFF */ - uint32x4_t data_key_lo = unzipped.val[0]; - /* data_key_hi = data_key >> 32 */ - uint32x4_t data_key_hi = unzipped.val[1]; - /* - * Then, we can split the vectors horizontally and multiply which, as for most - * widening intrinsics, have a variant that works on both high half vectors - * for free on AArch64. A similar instruction is available on SIMD128. - * - * sum = data_swap + (u64x2) data_key_lo * (u64x2) data_key_hi - */ - uint64x2_t sum_1 = XXH_vmlal_low_u32(data_swap_1, data_key_lo, data_key_hi); - uint64x2_t sum_2 = XXH_vmlal_high_u32(data_swap_2, data_key_lo, data_key_hi); - /* - * Clang reorders - * a += b * c; // umlal swap.2d, dkl.2s, dkh.2s - * c += a; // add acc.2d, acc.2d, swap.2d - * to - * c += a; // add acc.2d, acc.2d, swap.2d - * c += b * c; // umlal acc.2d, dkl.2s, dkh.2s - * - * While it would make sense in theory since the addition is faster, - * for reasons likely related to umlal being limited to certain NEON - * pipelines, this is worse. A compiler guard fixes this. - */ - XXH_COMPILER_GUARD_CLANG_NEON(sum_1); - XXH_COMPILER_GUARD_CLANG_NEON(sum_2); - /* xacc[i] = acc_vec + sum; */ - xacc[i] = vaddq_u64(xacc[i], sum_1); - xacc[i+1] = vaddq_u64(xacc[i+1], sum_2); - } - /* Operate on the remaining NEON lanes 2 at a time. */ - for (; i < XXH3_NEON_LANES / 2; i++) { - /* data_vec = xinput[i]; */ - uint64x2_t data_vec = XXH_vld1q_u64(xinput + (i * 16)); - /* key_vec = xsecret[i]; */ - uint64x2_t key_vec = XXH_vld1q_u64(xsecret + (i * 16)); - /* acc_vec_2 = swap(data_vec) */ - uint64x2_t data_swap = vextq_u64(data_vec, data_vec, 1); - /* data_key = data_vec ^ key_vec; */ - uint64x2_t data_key = veorq_u64(data_vec, key_vec); - /* For two lanes, just use VMOVN and VSHRN. */ - /* data_key_lo = data_key & 0xFFFFFFFF; */ - uint32x2_t data_key_lo = vmovn_u64(data_key); - /* data_key_hi = data_key >> 32; */ - uint32x2_t data_key_hi = vshrn_n_u64(data_key, 32); - /* sum = data_swap + (u64x2) data_key_lo * (u64x2) data_key_hi; */ - uint64x2_t sum = vmlal_u32(data_swap, data_key_lo, data_key_hi); - /* Same Clang workaround as before */ - XXH_COMPILER_GUARD_CLANG_NEON(sum); - /* xacc[i] = acc_vec + sum; */ - xacc[i] = vaddq_u64 (xacc[i], sum); - } - } -} -XXH_FORCE_INLINE XXH3_ACCUMULATE_TEMPLATE(neon) - -XXH_FORCE_INLINE void -XXH3_scrambleAcc_neon(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 15) == 0); - - { xxh_aliasing_uint64x2_t* xacc = (xxh_aliasing_uint64x2_t*) acc; - uint8_t const* xsecret = (uint8_t const*) secret; - - size_t i; - /* WASM uses operator overloads and doesn't need these. */ -#ifndef __wasm_simd128__ - /* { prime32_1, prime32_1 } */ - uint32x2_t const kPrimeLo = vdup_n_u32(XXH_PRIME32_1); - /* { 0, prime32_1, 0, prime32_1 } */ - uint32x4_t const kPrimeHi = vreinterpretq_u32_u64(vdupq_n_u64((xxh_u64)XXH_PRIME32_1 << 32)); -#endif - - /* AArch64 uses both scalar and neon at the same time */ - for (i = XXH3_NEON_LANES; i < XXH_ACC_NB; i++) { - XXH3_scalarScrambleRound(acc, secret, i); - } - for (i=0; i < XXH3_NEON_LANES / 2; i++) { - /* xacc[i] ^= (xacc[i] >> 47); */ - uint64x2_t acc_vec = xacc[i]; - uint64x2_t shifted = vshrq_n_u64(acc_vec, 47); - uint64x2_t data_vec = veorq_u64(acc_vec, shifted); - - /* xacc[i] ^= xsecret[i]; */ - uint64x2_t key_vec = XXH_vld1q_u64(xsecret + (i * 16)); - uint64x2_t data_key = veorq_u64(data_vec, key_vec); - /* xacc[i] *= XXH_PRIME32_1 */ -#ifdef __wasm_simd128__ - /* SIMD128 has multiply by u64x2, use it instead of expanding and scalarizing */ - xacc[i] = data_key * XXH_PRIME32_1; -#else - /* - * Expanded version with portable NEON intrinsics - * - * lo(x) * lo(y) + (hi(x) * lo(y) << 32) - * - * prod_hi = hi(data_key) * lo(prime) << 32 - * - * Since we only need 32 bits of this multiply a trick can be used, reinterpreting the vector - * as a uint32x4_t and multiplying by { 0, prime, 0, prime } to cancel out the unwanted bits - * and avoid the shift. - */ - uint32x4_t prod_hi = vmulq_u32 (vreinterpretq_u32_u64(data_key), kPrimeHi); - /* Extract low bits for vmlal_u32 */ - uint32x2_t data_key_lo = vmovn_u64(data_key); - /* xacc[i] = prod_hi + lo(data_key) * XXH_PRIME32_1; */ - xacc[i] = vmlal_u32(vreinterpretq_u64_u32(prod_hi), data_key_lo, kPrimeLo); -#endif - } - } -} -#endif - -#if (XXH_VECTOR == XXH_VSX) - -XXH_FORCE_INLINE void -XXH3_accumulate_512_vsx( void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - /* presumed aligned */ - xxh_aliasing_u64x2* const xacc = (xxh_aliasing_u64x2*) acc; - xxh_u8 const* const xinput = (xxh_u8 const*) input; /* no alignment restriction */ - xxh_u8 const* const xsecret = (xxh_u8 const*) secret; /* no alignment restriction */ - xxh_u64x2 const v32 = { 32, 32 }; - size_t i; - for (i = 0; i < XXH_STRIPE_LEN / sizeof(xxh_u64x2); i++) { - /* data_vec = xinput[i]; */ - xxh_u64x2 const data_vec = XXH_vec_loadu(xinput + 16*i); - /* key_vec = xsecret[i]; */ - xxh_u64x2 const key_vec = XXH_vec_loadu(xsecret + 16*i); - xxh_u64x2 const data_key = data_vec ^ key_vec; - /* shuffled = (data_key << 32) | (data_key >> 32); */ - xxh_u32x4 const shuffled = (xxh_u32x4)vec_rl(data_key, v32); - /* product = ((xxh_u64x2)data_key & 0xFFFFFFFF) * ((xxh_u64x2)shuffled & 0xFFFFFFFF); */ - xxh_u64x2 const product = XXH_vec_mulo((xxh_u32x4)data_key, shuffled); - /* acc_vec = xacc[i]; */ - xxh_u64x2 acc_vec = xacc[i]; - acc_vec += product; - - /* swap high and low halves */ -#ifdef __s390x__ - acc_vec += vec_permi(data_vec, data_vec, 2); -#else - acc_vec += vec_xxpermdi(data_vec, data_vec, 2); -#endif - xacc[i] = acc_vec; - } -} -XXH_FORCE_INLINE XXH3_ACCUMULATE_TEMPLATE(vsx) - -XXH_FORCE_INLINE void -XXH3_scrambleAcc_vsx(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) -{ - XXH_ASSERT((((size_t)acc) & 15) == 0); - - { xxh_aliasing_u64x2* const xacc = (xxh_aliasing_u64x2*) acc; - const xxh_u8* const xsecret = (const xxh_u8*) secret; - /* constants */ - xxh_u64x2 const v32 = { 32, 32 }; - xxh_u64x2 const v47 = { 47, 47 }; - xxh_u32x4 const prime = { XXH_PRIME32_1, XXH_PRIME32_1, XXH_PRIME32_1, XXH_PRIME32_1 }; - size_t i; - for (i = 0; i < XXH_STRIPE_LEN / sizeof(xxh_u64x2); i++) { - /* xacc[i] ^= (xacc[i] >> 47); */ - xxh_u64x2 const acc_vec = xacc[i]; - xxh_u64x2 const data_vec = acc_vec ^ (acc_vec >> v47); - - /* xacc[i] ^= xsecret[i]; */ - xxh_u64x2 const key_vec = XXH_vec_loadu(xsecret + 16*i); - xxh_u64x2 const data_key = data_vec ^ key_vec; - - /* xacc[i] *= XXH_PRIME32_1 */ - /* prod_lo = ((xxh_u64x2)data_key & 0xFFFFFFFF) * ((xxh_u64x2)prime & 0xFFFFFFFF); */ - xxh_u64x2 const prod_even = XXH_vec_mule((xxh_u32x4)data_key, prime); - /* prod_hi = ((xxh_u64x2)data_key >> 32) * ((xxh_u64x2)prime >> 32); */ - xxh_u64x2 const prod_odd = XXH_vec_mulo((xxh_u32x4)data_key, prime); - xacc[i] = prod_odd + (prod_even << v32); - } } -} - -#endif - -#if (XXH_VECTOR == XXH_SVE) - -XXH_FORCE_INLINE void -XXH3_accumulate_512_sve( void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - uint64_t *xacc = (uint64_t *)acc; - const uint64_t *xinput = (const uint64_t *)(const void *)input; - const uint64_t *xsecret = (const uint64_t *)(const void *)secret; - svuint64_t kSwap = sveor_n_u64_z(svptrue_b64(), svindex_u64(0, 1), 1); - uint64_t element_count = svcntd(); - if (element_count >= 8) { - svbool_t mask = svptrue_pat_b64(SV_VL8); - svuint64_t vacc = svld1_u64(mask, xacc); - ACCRND(vacc, 0); - svst1_u64(mask, xacc, vacc); - } else if (element_count == 2) { /* sve128 */ - svbool_t mask = svptrue_pat_b64(SV_VL2); - svuint64_t acc0 = svld1_u64(mask, xacc + 0); - svuint64_t acc1 = svld1_u64(mask, xacc + 2); - svuint64_t acc2 = svld1_u64(mask, xacc + 4); - svuint64_t acc3 = svld1_u64(mask, xacc + 6); - ACCRND(acc0, 0); - ACCRND(acc1, 2); - ACCRND(acc2, 4); - ACCRND(acc3, 6); - svst1_u64(mask, xacc + 0, acc0); - svst1_u64(mask, xacc + 2, acc1); - svst1_u64(mask, xacc + 4, acc2); - svst1_u64(mask, xacc + 6, acc3); - } else { - svbool_t mask = svptrue_pat_b64(SV_VL4); - svuint64_t acc0 = svld1_u64(mask, xacc + 0); - svuint64_t acc1 = svld1_u64(mask, xacc + 4); - ACCRND(acc0, 0); - ACCRND(acc1, 4); - svst1_u64(mask, xacc + 0, acc0); - svst1_u64(mask, xacc + 4, acc1); - } -} - -XXH_FORCE_INLINE void -XXH3_accumulate_sve(xxh_u64* XXH_RESTRICT acc, - const xxh_u8* XXH_RESTRICT input, - const xxh_u8* XXH_RESTRICT secret, - size_t nbStripes) -{ - if (nbStripes != 0) { - uint64_t *xacc = (uint64_t *)acc; - const uint64_t *xinput = (const uint64_t *)(const void *)input; - const uint64_t *xsecret = (const uint64_t *)(const void *)secret; - svuint64_t kSwap = sveor_n_u64_z(svptrue_b64(), svindex_u64(0, 1), 1); - uint64_t element_count = svcntd(); - if (element_count >= 8) { - svbool_t mask = svptrue_pat_b64(SV_VL8); - svuint64_t vacc = svld1_u64(mask, xacc + 0); - do { - /* svprfd(svbool_t, void *, enum svfprop); */ - svprfd(mask, xinput + 128, SV_PLDL1STRM); - ACCRND(vacc, 0); - xinput += 8; - xsecret += 1; - nbStripes--; - } while (nbStripes != 0); - - svst1_u64(mask, xacc + 0, vacc); - } else if (element_count == 2) { /* sve128 */ - svbool_t mask = svptrue_pat_b64(SV_VL2); - svuint64_t acc0 = svld1_u64(mask, xacc + 0); - svuint64_t acc1 = svld1_u64(mask, xacc + 2); - svuint64_t acc2 = svld1_u64(mask, xacc + 4); - svuint64_t acc3 = svld1_u64(mask, xacc + 6); - do { - svprfd(mask, xinput + 128, SV_PLDL1STRM); - ACCRND(acc0, 0); - ACCRND(acc1, 2); - ACCRND(acc2, 4); - ACCRND(acc3, 6); - xinput += 8; - xsecret += 1; - nbStripes--; - } while (nbStripes != 0); - - svst1_u64(mask, xacc + 0, acc0); - svst1_u64(mask, xacc + 2, acc1); - svst1_u64(mask, xacc + 4, acc2); - svst1_u64(mask, xacc + 6, acc3); - } else { - svbool_t mask = svptrue_pat_b64(SV_VL4); - svuint64_t acc0 = svld1_u64(mask, xacc + 0); - svuint64_t acc1 = svld1_u64(mask, xacc + 4); - do { - svprfd(mask, xinput + 128, SV_PLDL1STRM); - ACCRND(acc0, 0); - ACCRND(acc1, 4); - xinput += 8; - xsecret += 1; - nbStripes--; - } while (nbStripes != 0); - - svst1_u64(mask, xacc + 0, acc0); - svst1_u64(mask, xacc + 4, acc1); - } - } -} - -#endif - -/* scalar variants - universal */ - -#if defined(__aarch64__) && (defined(__GNUC__) || defined(__clang__)) -/* - * In XXH3_scalarRound(), GCC and Clang have a similar codegen issue, where they - * emit an excess mask and a full 64-bit multiply-add (MADD X-form). - * - * While this might not seem like much, as AArch64 is a 64-bit architecture, only - * big Cortex designs have a full 64-bit multiplier. - * - * On the little cores, the smaller 32-bit multiplier is used, and full 64-bit - * multiplies expand to 2-3 multiplies in microcode. This has a major penalty - * of up to 4 latency cycles and 2 stall cycles in the multiply pipeline. - * - * Thankfully, AArch64 still provides the 32-bit long multiply-add (UMADDL) which does - * not have this penalty and does the mask automatically. - */ -XXH_FORCE_INLINE xxh_u64 -XXH_mult32to64_add64(xxh_u64 lhs, xxh_u64 rhs, xxh_u64 acc) -{ - xxh_u64 ret; - /* note: %x = 64-bit register, %w = 32-bit register */ - __asm__("umaddl %x0, %w1, %w2, %x3" : "=r" (ret) : "r" (lhs), "r" (rhs), "r" (acc)); - return ret; -} -#else -XXH_FORCE_INLINE xxh_u64 -XXH_mult32to64_add64(xxh_u64 lhs, xxh_u64 rhs, xxh_u64 acc) -{ - return XXH_mult32to64((xxh_u32)lhs, (xxh_u32)rhs) + acc; -} -#endif - -/*! - * @internal - * @brief Scalar round for @ref XXH3_accumulate_512_scalar(). - * - * This is extracted to its own function because the NEON path uses a combination - * of NEON and scalar. - */ -XXH_FORCE_INLINE void -XXH3_scalarRound(void* XXH_RESTRICT acc, - void const* XXH_RESTRICT input, - void const* XXH_RESTRICT secret, - size_t lane) -{ - xxh_u64* xacc = (xxh_u64*) acc; - xxh_u8 const* xinput = (xxh_u8 const*) input; - xxh_u8 const* xsecret = (xxh_u8 const*) secret; - XXH_ASSERT(lane < XXH_ACC_NB); - XXH_ASSERT(((size_t)acc & (XXH_ACC_ALIGN-1)) == 0); - { - xxh_u64 const data_val = XXH_readLE64(xinput + lane * 8); - xxh_u64 const data_key = data_val ^ XXH_readLE64(xsecret + lane * 8); - xacc[lane ^ 1] += data_val; /* swap adjacent lanes */ - xacc[lane] = XXH_mult32to64_add64(data_key /* & 0xFFFFFFFF */, data_key >> 32, xacc[lane]); - } -} - -/*! - * @internal - * @brief Processes a 64 byte block of data using the scalar path. - */ -XXH_FORCE_INLINE void -XXH3_accumulate_512_scalar(void* XXH_RESTRICT acc, - const void* XXH_RESTRICT input, - const void* XXH_RESTRICT secret) -{ - size_t i; - /* ARM GCC refuses to unroll this loop, resulting in a 24% slowdown on ARMv6. */ -#if defined(__GNUC__) && !defined(__clang__) \ - && (defined(__arm__) || defined(__thumb2__)) \ - && defined(__ARM_FEATURE_UNALIGNED) /* no unaligned access just wastes bytes */ \ - && XXH_SIZE_OPT <= 0 -# pragma GCC unroll 8 -#endif - for (i=0; i < XXH_ACC_NB; i++) { - XXH3_scalarRound(acc, input, secret, i); - } -} -XXH_FORCE_INLINE XXH3_ACCUMULATE_TEMPLATE(scalar) - -/*! - * @internal - * @brief Scalar scramble step for @ref XXH3_scrambleAcc_scalar(). - * - * This is extracted to its own function because the NEON path uses a combination - * of NEON and scalar. - */ -XXH_FORCE_INLINE void -XXH3_scalarScrambleRound(void* XXH_RESTRICT acc, - void const* XXH_RESTRICT secret, - size_t lane) -{ - xxh_u64* const xacc = (xxh_u64*) acc; /* presumed aligned */ - const xxh_u8* const xsecret = (const xxh_u8*) secret; /* no alignment restriction */ - XXH_ASSERT((((size_t)acc) & (XXH_ACC_ALIGN-1)) == 0); - XXH_ASSERT(lane < XXH_ACC_NB); - { - xxh_u64 const key64 = XXH_readLE64(xsecret + lane * 8); - xxh_u64 acc64 = xacc[lane]; - acc64 = XXH_xorshift64(acc64, 47); - acc64 ^= key64; - acc64 *= XXH_PRIME32_1; - xacc[lane] = acc64; - } -} - -/*! - * @internal - * @brief Scrambles the accumulators after a large chunk has been read - */ -XXH_FORCE_INLINE void -XXH3_scrambleAcc_scalar(void* XXH_RESTRICT acc, const void* XXH_RESTRICT secret) -{ - size_t i; - for (i=0; i < XXH_ACC_NB; i++) { - XXH3_scalarScrambleRound(acc, secret, i); - } -} - -XXH_FORCE_INLINE void -XXH3_initCustomSecret_scalar(void* XXH_RESTRICT customSecret, xxh_u64 seed64) -{ - /* - * We need a separate pointer for the hack below, - * which requires a non-const pointer. - * Any decent compiler will optimize this out otherwise. - */ - const xxh_u8* kSecretPtr = XXH3_kSecret; - XXH_STATIC_ASSERT((XXH_SECRET_DEFAULT_SIZE & 15) == 0); - -#if defined(__GNUC__) && defined(__aarch64__) - /* - * UGLY HACK: - * GCC and Clang generate a bunch of MOV/MOVK pairs for aarch64, and they are - * placed sequentially, in order, at the top of the unrolled loop. - * - * While MOVK is great for generating constants (2 cycles for a 64-bit - * constant compared to 4 cycles for LDR), it fights for bandwidth with - * the arithmetic instructions. - * - * I L S - * MOVK - * MOVK - * MOVK - * MOVK - * ADD - * SUB STR - * STR - * By forcing loads from memory (as the asm line causes the compiler to assume - * that XXH3_kSecretPtr has been changed), the pipelines are used more - * efficiently: - * I L S - * LDR - * ADD LDR - * SUB STR - * STR - * - * See XXH3_NEON_LANES for details on the pipsline. - * - * XXH3_64bits_withSeed, len == 256, Snapdragon 835 - * without hack: 2654.4 MB/s - * with hack: 3202.9 MB/s - */ - XXH_COMPILER_GUARD(kSecretPtr); -#endif - { int const nbRounds = XXH_SECRET_DEFAULT_SIZE / 16; - int i; - for (i=0; i < nbRounds; i++) { - /* - * The asm hack causes the compiler to assume that kSecretPtr aliases with - * customSecret, and on aarch64, this prevented LDP from merging two - * loads together for free. Putting the loads together before the stores - * properly generates LDP. - */ - xxh_u64 lo = XXH_readLE64(kSecretPtr + 16*i) + seed64; - xxh_u64 hi = XXH_readLE64(kSecretPtr + 16*i + 8) - seed64; - XXH_writeLE64((xxh_u8*)customSecret + 16*i, lo); - XXH_writeLE64((xxh_u8*)customSecret + 16*i + 8, hi); - } } -} - - -typedef void (*XXH3_f_accumulate)(xxh_u64* XXH_RESTRICT, const xxh_u8* XXH_RESTRICT, const xxh_u8* XXH_RESTRICT, size_t); -typedef void (*XXH3_f_scrambleAcc)(void* XXH_RESTRICT, const void*); -typedef void (*XXH3_f_initCustomSecret)(void* XXH_RESTRICT, xxh_u64); - - -#if (XXH_VECTOR == XXH_AVX512) - -#define XXH3_accumulate_512 XXH3_accumulate_512_avx512 -#define XXH3_accumulate XXH3_accumulate_avx512 -#define XXH3_scrambleAcc XXH3_scrambleAcc_avx512 -#define XXH3_initCustomSecret XXH3_initCustomSecret_avx512 - -#elif (XXH_VECTOR == XXH_AVX2) - -#define XXH3_accumulate_512 XXH3_accumulate_512_avx2 -#define XXH3_accumulate XXH3_accumulate_avx2 -#define XXH3_scrambleAcc XXH3_scrambleAcc_avx2 -#define XXH3_initCustomSecret XXH3_initCustomSecret_avx2 - -#elif (XXH_VECTOR == XXH_SSE2) - -#define XXH3_accumulate_512 XXH3_accumulate_512_sse2 -#define XXH3_accumulate XXH3_accumulate_sse2 -#define XXH3_scrambleAcc XXH3_scrambleAcc_sse2 -#define XXH3_initCustomSecret XXH3_initCustomSecret_sse2 - -#elif (XXH_VECTOR == XXH_NEON) - -#define XXH3_accumulate_512 XXH3_accumulate_512_neon -#define XXH3_accumulate XXH3_accumulate_neon -#define XXH3_scrambleAcc XXH3_scrambleAcc_neon -#define XXH3_initCustomSecret XXH3_initCustomSecret_scalar - -#elif (XXH_VECTOR == XXH_VSX) - -#define XXH3_accumulate_512 XXH3_accumulate_512_vsx -#define XXH3_accumulate XXH3_accumulate_vsx -#define XXH3_scrambleAcc XXH3_scrambleAcc_vsx -#define XXH3_initCustomSecret XXH3_initCustomSecret_scalar - -#elif (XXH_VECTOR == XXH_SVE) -#define XXH3_accumulate_512 XXH3_accumulate_512_sve -#define XXH3_accumulate XXH3_accumulate_sve -#define XXH3_scrambleAcc XXH3_scrambleAcc_scalar -#define XXH3_initCustomSecret XXH3_initCustomSecret_scalar - -#else /* scalar */ - -#define XXH3_accumulate_512 XXH3_accumulate_512_scalar -#define XXH3_accumulate XXH3_accumulate_scalar -#define XXH3_scrambleAcc XXH3_scrambleAcc_scalar -#define XXH3_initCustomSecret XXH3_initCustomSecret_scalar - -#endif - -#if XXH_SIZE_OPT >= 1 /* don't do SIMD for initialization */ -# undef XXH3_initCustomSecret -# define XXH3_initCustomSecret XXH3_initCustomSecret_scalar -#endif - -XXH_FORCE_INLINE void -XXH3_hashLong_internal_loop(xxh_u64* XXH_RESTRICT acc, - const xxh_u8* XXH_RESTRICT input, size_t len, - const xxh_u8* XXH_RESTRICT secret, size_t secretSize, - XXH3_f_accumulate f_acc, - XXH3_f_scrambleAcc f_scramble) -{ - size_t const nbStripesPerBlock = (secretSize - XXH_STRIPE_LEN) / XXH_SECRET_CONSUME_RATE; - size_t const block_len = XXH_STRIPE_LEN * nbStripesPerBlock; - size_t const nb_blocks = (len - 1) / block_len; - - size_t n; - - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); - - for (n = 0; n < nb_blocks; n++) { - f_acc(acc, input + n*block_len, secret, nbStripesPerBlock); - f_scramble(acc, secret + secretSize - XXH_STRIPE_LEN); - } - - /* last partial block */ - XXH_ASSERT(len > XXH_STRIPE_LEN); - { size_t const nbStripes = ((len - 1) - (block_len * nb_blocks)) / XXH_STRIPE_LEN; - XXH_ASSERT(nbStripes <= (secretSize / XXH_SECRET_CONSUME_RATE)); - f_acc(acc, input + nb_blocks*block_len, secret, nbStripes); - - /* last stripe */ - { const xxh_u8* const p = input + len - XXH_STRIPE_LEN; -#define XXH_SECRET_LASTACC_START 7 /* not aligned on 8, last secret is different from acc & scrambler */ - XXH3_accumulate_512(acc, p, secret + secretSize - XXH_STRIPE_LEN - XXH_SECRET_LASTACC_START); - } } -} - -XXH_FORCE_INLINE xxh_u64 -XXH3_mix2Accs(const xxh_u64* XXH_RESTRICT acc, const xxh_u8* XXH_RESTRICT secret) -{ - return XXH3_mul128_fold64( - acc[0] ^ XXH_readLE64(secret), - acc[1] ^ XXH_readLE64(secret+8) ); -} - -static XXH64_hash_t -XXH3_mergeAccs(const xxh_u64* XXH_RESTRICT acc, const xxh_u8* XXH_RESTRICT secret, xxh_u64 start) -{ - xxh_u64 result64 = start; - size_t i = 0; - - for (i = 0; i < 4; i++) { - result64 += XXH3_mix2Accs(acc+2*i, secret + 16*i); -#if defined(__clang__) /* Clang */ \ - && (defined(__arm__) || defined(__thumb__)) /* ARMv7 */ \ - && (defined(__ARM_NEON) || defined(__ARM_NEON__)) /* NEON */ \ - && !defined(XXH_ENABLE_AUTOVECTORIZE) /* Define to disable */ - /* - * UGLY HACK: - * Prevent autovectorization on Clang ARMv7-a. Exact same problem as - * the one in XXH3_len_129to240_64b. Speeds up shorter keys > 240b. - * XXH3_64bits, len == 256, Snapdragon 835: - * without hack: 2063.7 MB/s - * with hack: 2560.7 MB/s - */ - XXH_COMPILER_GUARD(result64); -#endif - } - - return XXH3_avalanche(result64); -} - -#define XXH3_INIT_ACC { XXH_PRIME32_3, XXH_PRIME64_1, XXH_PRIME64_2, XXH_PRIME64_3, \ - XXH_PRIME64_4, XXH_PRIME32_2, XXH_PRIME64_5, XXH_PRIME32_1 } - -XXH_FORCE_INLINE XXH64_hash_t -XXH3_hashLong_64b_internal(const void* XXH_RESTRICT input, size_t len, - const void* XXH_RESTRICT secret, size_t secretSize, - XXH3_f_accumulate f_acc, - XXH3_f_scrambleAcc f_scramble) -{ - XXH_ALIGN(XXH_ACC_ALIGN) xxh_u64 acc[XXH_ACC_NB] = XXH3_INIT_ACC; - - XXH3_hashLong_internal_loop(acc, (const xxh_u8*)input, len, (const xxh_u8*)secret, secretSize, f_acc, f_scramble); - - /* converge into final hash */ - XXH_STATIC_ASSERT(sizeof(acc) == 64); - /* do not align on 8, so that the secret is different from the accumulator */ -#define XXH_SECRET_MERGEACCS_START 11 - XXH_ASSERT(secretSize >= sizeof(acc) + XXH_SECRET_MERGEACCS_START); - return XXH3_mergeAccs(acc, (const xxh_u8*)secret + XXH_SECRET_MERGEACCS_START, (xxh_u64)len * XXH_PRIME64_1); -} - -/* - * It's important for performance to transmit secret's size (when it's static) - * so that the compiler can properly optimize the vectorized loop. - * This makes a big performance difference for "medium" keys (<1 KB) when using AVX instruction set. - * When the secret size is unknown, or on GCC 12 where the mix of NO_INLINE and FORCE_INLINE - * breaks -Og, this is XXH_NO_INLINE. - */ -XXH3_WITH_SECRET_INLINE XXH64_hash_t -XXH3_hashLong_64b_withSecret(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed64, const xxh_u8* XXH_RESTRICT secret, size_t secretLen) -{ - (void)seed64; - return XXH3_hashLong_64b_internal(input, len, secret, secretLen, XXH3_accumulate, XXH3_scrambleAcc); -} - -/* - * It's preferable for performance that XXH3_hashLong is not inlined, - * as it results in a smaller function for small data, easier to the instruction cache. - * Note that inside this no_inline function, we do inline the internal loop, - * and provide a statically defined secret size to allow optimization of vector loop. - */ -XXH_NO_INLINE XXH_PUREF XXH64_hash_t -XXH3_hashLong_64b_default(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed64, const xxh_u8* XXH_RESTRICT secret, size_t secretLen) -{ - (void)seed64; (void)secret; (void)secretLen; - return XXH3_hashLong_64b_internal(input, len, XXH3_kSecret, sizeof(XXH3_kSecret), XXH3_accumulate, XXH3_scrambleAcc); -} - -/* - * XXH3_hashLong_64b_withSeed(): - * Generate a custom key based on alteration of default XXH3_kSecret with the seed, - * and then use this key for long mode hashing. - * - * This operation is decently fast but nonetheless costs a little bit of time. - * Try to avoid it whenever possible (typically when seed==0). - * - * It's important for performance that XXH3_hashLong is not inlined. Not sure - * why (uop cache maybe?), but the difference is large and easily measurable. - */ -XXH_FORCE_INLINE XXH64_hash_t -XXH3_hashLong_64b_withSeed_internal(const void* input, size_t len, - XXH64_hash_t seed, - XXH3_f_accumulate f_acc, - XXH3_f_scrambleAcc f_scramble, - XXH3_f_initCustomSecret f_initSec) -{ -#if XXH_SIZE_OPT <= 0 - if (seed == 0) - return XXH3_hashLong_64b_internal(input, len, - XXH3_kSecret, sizeof(XXH3_kSecret), - f_acc, f_scramble); -#endif - { XXH_ALIGN(XXH_SEC_ALIGN) xxh_u8 secret[XXH_SECRET_DEFAULT_SIZE]; - f_initSec(secret, seed); - return XXH3_hashLong_64b_internal(input, len, secret, sizeof(secret), - f_acc, f_scramble); - } -} - -/* - * It's important for performance that XXH3_hashLong is not inlined. - */ -XXH_NO_INLINE XXH64_hash_t -XXH3_hashLong_64b_withSeed(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed, const xxh_u8* XXH_RESTRICT secret, size_t secretLen) -{ - (void)secret; (void)secretLen; - return XXH3_hashLong_64b_withSeed_internal(input, len, seed, - XXH3_accumulate, XXH3_scrambleAcc, XXH3_initCustomSecret); -} - - -typedef XXH64_hash_t (*XXH3_hashLong64_f)(const void* XXH_RESTRICT, size_t, - XXH64_hash_t, const xxh_u8* XXH_RESTRICT, size_t); - -XXH_FORCE_INLINE XXH64_hash_t -XXH3_64bits_internal(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed64, const void* XXH_RESTRICT secret, size_t secretLen, - XXH3_hashLong64_f f_hashLong) -{ - XXH_ASSERT(secretLen >= XXH3_SECRET_SIZE_MIN); - /* - * If an action is to be taken if `secretLen` condition is not respected, - * it should be done here. - * For now, it's a contract pre-condition. - * Adding a check and a branch here would cost performance at every hash. - * Also, note that function signature doesn't offer room to return an error. - */ - if (len <= 16) - return XXH3_len_0to16_64b((const xxh_u8*)input, len, (const xxh_u8*)secret, seed64); - if (len <= 128) - return XXH3_len_17to128_64b((const xxh_u8*)input, len, (const xxh_u8*)secret, secretLen, seed64); - if (len <= XXH3_MIDSIZE_MAX) - return XXH3_len_129to240_64b((const xxh_u8*)input, len, (const xxh_u8*)secret, secretLen, seed64); - return f_hashLong(input, len, seed64, (const xxh_u8*)secret, secretLen); -} - - -/* === Public entry point === */ - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH64_hash_t XXH3_64bits(XXH_NOESCAPE const void* input, size_t length) -{ - return XXH3_64bits_internal(input, length, 0, XXH3_kSecret, sizeof(XXH3_kSecret), XXH3_hashLong_64b_default); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH64_hash_t -XXH3_64bits_withSecret(XXH_NOESCAPE const void* input, size_t length, XXH_NOESCAPE const void* secret, size_t secretSize) -{ - return XXH3_64bits_internal(input, length, 0, secret, secretSize, XXH3_hashLong_64b_withSecret); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH64_hash_t -XXH3_64bits_withSeed(XXH_NOESCAPE const void* input, size_t length, XXH64_hash_t seed) -{ - return XXH3_64bits_internal(input, length, seed, XXH3_kSecret, sizeof(XXH3_kSecret), XXH3_hashLong_64b_withSeed); -} - -XXH_PUBLIC_API XXH64_hash_t -XXH3_64bits_withSecretandSeed(XXH_NOESCAPE const void* input, size_t length, XXH_NOESCAPE const void* secret, size_t secretSize, XXH64_hash_t seed) -{ - if (length <= XXH3_MIDSIZE_MAX) - return XXH3_64bits_internal(input, length, seed, XXH3_kSecret, sizeof(XXH3_kSecret), NULL); - return XXH3_hashLong_64b_withSecret(input, length, seed, (const xxh_u8*)secret, secretSize); -} - - -/* === XXH3 streaming === */ -#ifndef XXH_NO_STREAM -/* - * Malloc's a pointer that is always aligned to align. - * - * This must be freed with `XXH_alignedFree()`. - * - * malloc typically guarantees 16 byte alignment on 64-bit systems and 8 byte - * alignment on 32-bit. This isn't enough for the 32 byte aligned loads in AVX2 - * or on 32-bit, the 16 byte aligned loads in SSE2 and NEON. - * - * This underalignment previously caused a rather obvious crash which went - * completely unnoticed due to XXH3_createState() not actually being tested. - * Credit to RedSpah for noticing this bug. - * - * The alignment is done manually: Functions like posix_memalign or _mm_malloc - * are avoided: To maintain portability, we would have to write a fallback - * like this anyways, and besides, testing for the existence of library - * functions without relying on external build tools is impossible. - * - * The method is simple: Overallocate, manually align, and store the offset - * to the original behind the returned pointer. - * - * Align must be a power of 2 and 8 <= align <= 128. - */ -static XXH_MALLOCF void* XXH_alignedMalloc(size_t s, size_t align) -{ - XXH_ASSERT(align <= 128 && align >= 8); /* range check */ - XXH_ASSERT((align & (align-1)) == 0); /* power of 2 */ - XXH_ASSERT(s != 0 && s < (s + align)); /* empty/overflow */ - { /* Overallocate to make room for manual realignment and an offset byte */ - xxh_u8* base = (xxh_u8*)XXH_malloc(s + align); - if (base != NULL) { - /* - * Get the offset needed to align this pointer. - * - * Even if the returned pointer is aligned, there will always be - * at least one byte to store the offset to the original pointer. - */ - size_t offset = align - ((size_t)base & (align - 1)); /* base % align */ - /* Add the offset for the now-aligned pointer */ - xxh_u8* ptr = base + offset; - - XXH_ASSERT((size_t)ptr % align == 0); - - /* Store the offset immediately before the returned pointer. */ - ptr[-1] = (xxh_u8)offset; - return ptr; - } - return NULL; - } -} -/* - * Frees an aligned pointer allocated by XXH_alignedMalloc(). Don't pass - * normal malloc'd pointers, XXH_alignedMalloc has a specific data layout. - */ -static void XXH_alignedFree(void* p) -{ - if (p != NULL) { - xxh_u8* ptr = (xxh_u8*)p; - /* Get the offset byte we added in XXH_malloc. */ - xxh_u8 offset = ptr[-1]; - /* Free the original malloc'd pointer */ - xxh_u8* base = ptr - offset; - XXH_free(base); - } -} -/*! @ingroup XXH3_family */ -/*! - * @brief Allocate an @ref XXH3_state_t. - * - * @return An allocated pointer of @ref XXH3_state_t on success. - * @return `NULL` on failure. - * - * @note Must be freed with XXH3_freeState(). - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH3_state_t* XXH3_createState(void) -{ - XXH3_state_t* const state = (XXH3_state_t*)XXH_alignedMalloc(sizeof(XXH3_state_t), 64); - if (state==NULL) return NULL; - XXH3_INITSTATE(state); - return state; -} - -/*! @ingroup XXH3_family */ -/*! - * @brief Frees an @ref XXH3_state_t. - * - * @param statePtr A pointer to an @ref XXH3_state_t allocated with @ref XXH3_createState(). - * - * @return @ref XXH_OK. - * - * @note Must be allocated with XXH3_createState(). - * - * @see @ref streaming_example "Streaming Example" - */ -XXH_PUBLIC_API XXH_errorcode XXH3_freeState(XXH3_state_t* statePtr) -{ - XXH_alignedFree(statePtr); - return XXH_OK; -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API void -XXH3_copyState(XXH_NOESCAPE XXH3_state_t* dst_state, XXH_NOESCAPE const XXH3_state_t* src_state) -{ - XXH_memcpy(dst_state, src_state, sizeof(*dst_state)); -} - -static void -XXH3_reset_internal(XXH3_state_t* statePtr, - XXH64_hash_t seed, - const void* secret, size_t secretSize) -{ - size_t const initStart = offsetof(XXH3_state_t, bufferedSize); - size_t const initLength = offsetof(XXH3_state_t, nbStripesPerBlock) - initStart; - XXH_ASSERT(offsetof(XXH3_state_t, nbStripesPerBlock) > initStart); - XXH_ASSERT(statePtr != NULL); - /* set members from bufferedSize to nbStripesPerBlock (excluded) to 0 */ - memset((char*)statePtr + initStart, 0, initLength); - statePtr->acc[0] = XXH_PRIME32_3; - statePtr->acc[1] = XXH_PRIME64_1; - statePtr->acc[2] = XXH_PRIME64_2; - statePtr->acc[3] = XXH_PRIME64_3; - statePtr->acc[4] = XXH_PRIME64_4; - statePtr->acc[5] = XXH_PRIME32_2; - statePtr->acc[6] = XXH_PRIME64_5; - statePtr->acc[7] = XXH_PRIME32_1; - statePtr->seed = seed; - statePtr->useSeed = (seed != 0); - statePtr->extSecret = (const unsigned char*)secret; - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); - statePtr->secretLimit = secretSize - XXH_STRIPE_LEN; - statePtr->nbStripesPerBlock = statePtr->secretLimit / XXH_SECRET_CONSUME_RATE; -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_64bits_reset(XXH_NOESCAPE XXH3_state_t* statePtr) -{ - if (statePtr == NULL) return XXH_ERROR; - XXH3_reset_internal(statePtr, 0, XXH3_kSecret, XXH_SECRET_DEFAULT_SIZE); - return XXH_OK; -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_64bits_reset_withSecret(XXH_NOESCAPE XXH3_state_t* statePtr, XXH_NOESCAPE const void* secret, size_t secretSize) -{ - if (statePtr == NULL) return XXH_ERROR; - XXH3_reset_internal(statePtr, 0, secret, secretSize); - if (secret == NULL) return XXH_ERROR; - if (secretSize < XXH3_SECRET_SIZE_MIN) return XXH_ERROR; - return XXH_OK; -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_64bits_reset_withSeed(XXH_NOESCAPE XXH3_state_t* statePtr, XXH64_hash_t seed) -{ - if (statePtr == NULL) return XXH_ERROR; - if (seed==0) return XXH3_64bits_reset(statePtr); - if ((seed != statePtr->seed) || (statePtr->extSecret != NULL)) - XXH3_initCustomSecret(statePtr->customSecret, seed); - XXH3_reset_internal(statePtr, seed, NULL, XXH_SECRET_DEFAULT_SIZE); - return XXH_OK; -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_64bits_reset_withSecretandSeed(XXH_NOESCAPE XXH3_state_t* statePtr, XXH_NOESCAPE const void* secret, size_t secretSize, XXH64_hash_t seed64) -{ - if (statePtr == NULL) return XXH_ERROR; - if (secret == NULL) return XXH_ERROR; - if (secretSize < XXH3_SECRET_SIZE_MIN) return XXH_ERROR; - XXH3_reset_internal(statePtr, seed64, secret, secretSize); - statePtr->useSeed = 1; /* always, even if seed64==0 */ - return XXH_OK; -} - -/*! - * @internal - * @brief Processes a large input for XXH3_update() and XXH3_digest_long(). - * - * Unlike XXH3_hashLong_internal_loop(), this can process data that overlaps a block. - * - * @param acc Pointer to the 8 accumulator lanes - * @param nbStripesSoFarPtr In/out pointer to the number of leftover stripes in the block* - * @param nbStripesPerBlock Number of stripes in a block - * @param input Input pointer - * @param nbStripes Number of stripes to process - * @param secret Secret pointer - * @param secretLimit Offset of the last block in @p secret - * @param f_acc Pointer to an XXH3_accumulate implementation - * @param f_scramble Pointer to an XXH3_scrambleAcc implementation - * @return Pointer past the end of @p input after processing - */ -XXH_FORCE_INLINE const xxh_u8 * -XXH3_consumeStripes(xxh_u64* XXH_RESTRICT acc, - size_t* XXH_RESTRICT nbStripesSoFarPtr, size_t nbStripesPerBlock, - const xxh_u8* XXH_RESTRICT input, size_t nbStripes, - const xxh_u8* XXH_RESTRICT secret, size_t secretLimit, - XXH3_f_accumulate f_acc, - XXH3_f_scrambleAcc f_scramble) -{ - const xxh_u8* initialSecret = secret + *nbStripesSoFarPtr * XXH_SECRET_CONSUME_RATE; - /* Process full blocks */ - if (nbStripes >= (nbStripesPerBlock - *nbStripesSoFarPtr)) { - /* Process the initial partial block... */ - size_t nbStripesThisIter = nbStripesPerBlock - *nbStripesSoFarPtr; - - do { - /* Accumulate and scramble */ - f_acc(acc, input, initialSecret, nbStripesThisIter); - f_scramble(acc, secret + secretLimit); - input += nbStripesThisIter * XXH_STRIPE_LEN; - nbStripes -= nbStripesThisIter; - /* Then continue the loop with the full block size */ - nbStripesThisIter = nbStripesPerBlock; - initialSecret = secret; - } while (nbStripes >= nbStripesPerBlock); - *nbStripesSoFarPtr = 0; - } - /* Process a partial block */ - if (nbStripes > 0) { - f_acc(acc, input, initialSecret, nbStripes); - input += nbStripes * XXH_STRIPE_LEN; - *nbStripesSoFarPtr += nbStripes; - } - /* Return end pointer */ - return input; -} - -#ifndef XXH3_STREAM_USE_STACK -# if XXH_SIZE_OPT <= 0 && !defined(__clang__) /* clang doesn't need additional stack space */ -# define XXH3_STREAM_USE_STACK 1 -# endif -#endif -/* - * Both XXH3_64bits_update and XXH3_128bits_update use this routine. - */ -XXH_FORCE_INLINE XXH_errorcode -XXH3_update(XXH3_state_t* XXH_RESTRICT const state, - const xxh_u8* XXH_RESTRICT input, size_t len, - XXH3_f_accumulate f_acc, - XXH3_f_scrambleAcc f_scramble) -{ - if (input==NULL) { - XXH_ASSERT(len == 0); - return XXH_OK; - } - - XXH_ASSERT(state != NULL); - { const xxh_u8* const bEnd = input + len; - const unsigned char* const secret = (state->extSecret == NULL) ? state->customSecret : state->extSecret; -#if defined(XXH3_STREAM_USE_STACK) && XXH3_STREAM_USE_STACK >= 1 - /* For some reason, gcc and MSVC seem to suffer greatly - * when operating accumulators directly into state. - * Operating into stack space seems to enable proper optimization. - * clang, on the other hand, doesn't seem to need this trick */ - XXH_ALIGN(XXH_ACC_ALIGN) xxh_u64 acc[8]; - XXH_memcpy(acc, state->acc, sizeof(acc)); -#else - xxh_u64* XXH_RESTRICT const acc = state->acc; -#endif - state->totalLen += len; - XXH_ASSERT(state->bufferedSize <= XXH3_INTERNALBUFFER_SIZE); - - /* small input : just fill in tmp buffer */ - if (len <= XXH3_INTERNALBUFFER_SIZE - state->bufferedSize) { - XXH_memcpy(state->buffer + state->bufferedSize, input, len); - state->bufferedSize += (XXH32_hash_t)len; - return XXH_OK; - } - - /* total input is now > XXH3_INTERNALBUFFER_SIZE */ - #define XXH3_INTERNALBUFFER_STRIPES (XXH3_INTERNALBUFFER_SIZE / XXH_STRIPE_LEN) - XXH_STATIC_ASSERT(XXH3_INTERNALBUFFER_SIZE % XXH_STRIPE_LEN == 0); /* clean multiple */ - - /* - * Internal buffer is partially filled (always, except at beginning) - * Complete it, then consume it. - */ - if (state->bufferedSize) { - size_t const loadSize = XXH3_INTERNALBUFFER_SIZE - state->bufferedSize; - XXH_memcpy(state->buffer + state->bufferedSize, input, loadSize); - input += loadSize; - XXH3_consumeStripes(acc, - &state->nbStripesSoFar, state->nbStripesPerBlock, - state->buffer, XXH3_INTERNALBUFFER_STRIPES, - secret, state->secretLimit, - f_acc, f_scramble); - state->bufferedSize = 0; - } - XXH_ASSERT(input < bEnd); - if (bEnd - input > XXH3_INTERNALBUFFER_SIZE) { - size_t nbStripes = (size_t)(bEnd - 1 - input) / XXH_STRIPE_LEN; - input = XXH3_consumeStripes(acc, - &state->nbStripesSoFar, state->nbStripesPerBlock, - input, nbStripes, - secret, state->secretLimit, - f_acc, f_scramble); - XXH_memcpy(state->buffer + sizeof(state->buffer) - XXH_STRIPE_LEN, input - XXH_STRIPE_LEN, XXH_STRIPE_LEN); - - } - /* Some remaining input (always) : buffer it */ - XXH_ASSERT(input < bEnd); - XXH_ASSERT(bEnd - input <= XXH3_INTERNALBUFFER_SIZE); - XXH_ASSERT(state->bufferedSize == 0); - XXH_memcpy(state->buffer, input, (size_t)(bEnd-input)); - state->bufferedSize = (XXH32_hash_t)(bEnd-input); -#if defined(XXH3_STREAM_USE_STACK) && XXH3_STREAM_USE_STACK >= 1 - /* save stack accumulators into state */ - XXH_memcpy(state->acc, acc, sizeof(acc)); -#endif - } - - return XXH_OK; -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_64bits_update(XXH_NOESCAPE XXH3_state_t* state, XXH_NOESCAPE const void* input, size_t len) -{ - return XXH3_update(state, (const xxh_u8*)input, len, - XXH3_accumulate, XXH3_scrambleAcc); -} - - -XXH_FORCE_INLINE void -XXH3_digest_long (XXH64_hash_t* acc, - const XXH3_state_t* state, - const unsigned char* secret) -{ - xxh_u8 lastStripe[XXH_STRIPE_LEN]; - const xxh_u8* lastStripePtr; - - /* - * Digest on a local copy. This way, the state remains unaltered, and it can - * continue ingesting more input afterwards. - */ - XXH_memcpy(acc, state->acc, sizeof(state->acc)); - if (state->bufferedSize >= XXH_STRIPE_LEN) { - /* Consume remaining stripes then point to remaining data in buffer */ - size_t const nbStripes = (state->bufferedSize - 1) / XXH_STRIPE_LEN; - size_t nbStripesSoFar = state->nbStripesSoFar; - XXH3_consumeStripes(acc, - &nbStripesSoFar, state->nbStripesPerBlock, - state->buffer, nbStripes, - secret, state->secretLimit, - XXH3_accumulate, XXH3_scrambleAcc); - lastStripePtr = state->buffer + state->bufferedSize - XXH_STRIPE_LEN; - } else { /* bufferedSize < XXH_STRIPE_LEN */ - /* Copy to temp buffer */ - size_t const catchupSize = XXH_STRIPE_LEN - state->bufferedSize; - XXH_ASSERT(state->bufferedSize > 0); /* there is always some input buffered */ - XXH_memcpy(lastStripe, state->buffer + sizeof(state->buffer) - catchupSize, catchupSize); - XXH_memcpy(lastStripe + catchupSize, state->buffer, state->bufferedSize); - lastStripePtr = lastStripe; - } - /* Last stripe */ - XXH3_accumulate_512(acc, - lastStripePtr, - secret + state->secretLimit - XXH_SECRET_LASTACC_START); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH64_hash_t XXH3_64bits_digest (XXH_NOESCAPE const XXH3_state_t* state) -{ - const unsigned char* const secret = (state->extSecret == NULL) ? state->customSecret : state->extSecret; - if (state->totalLen > XXH3_MIDSIZE_MAX) { - XXH_ALIGN(XXH_ACC_ALIGN) XXH64_hash_t acc[XXH_ACC_NB]; - XXH3_digest_long(acc, state, secret); - return XXH3_mergeAccs(acc, - secret + XXH_SECRET_MERGEACCS_START, - (xxh_u64)state->totalLen * XXH_PRIME64_1); - } - /* totalLen <= XXH3_MIDSIZE_MAX: digesting a short input */ - if (state->useSeed) - return XXH3_64bits_withSeed(state->buffer, (size_t)state->totalLen, state->seed); - return XXH3_64bits_withSecret(state->buffer, (size_t)(state->totalLen), - secret, state->secretLimit + XXH_STRIPE_LEN); -} -#endif /* !XXH_NO_STREAM */ - - -/* ========================================== - * XXH3 128 bits (a.k.a XXH128) - * ========================================== - * XXH3's 128-bit variant has better mixing and strength than the 64-bit variant, - * even without counting the significantly larger output size. - * - * For example, extra steps are taken to avoid the seed-dependent collisions - * in 17-240 byte inputs (See XXH3_mix16B and XXH128_mix32B). - * - * This strength naturally comes at the cost of some speed, especially on short - * lengths. Note that longer hashes are about as fast as the 64-bit version - * due to it using only a slight modification of the 64-bit loop. - * - * XXH128 is also more oriented towards 64-bit machines. It is still extremely - * fast for a _128-bit_ hash on 32-bit (it usually clears XXH64). - */ - -XXH_FORCE_INLINE XXH_PUREF XXH128_hash_t -XXH3_len_1to3_128b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - /* A doubled version of 1to3_64b with different constants. */ - XXH_ASSERT(input != NULL); - XXH_ASSERT(1 <= len && len <= 3); - XXH_ASSERT(secret != NULL); - /* - * len = 1: combinedl = { input[0], 0x01, input[0], input[0] } - * len = 2: combinedl = { input[1], 0x02, input[0], input[1] } - * len = 3: combinedl = { input[2], 0x03, input[0], input[1] } - */ - { xxh_u8 const c1 = input[0]; - xxh_u8 const c2 = input[len >> 1]; - xxh_u8 const c3 = input[len - 1]; - xxh_u32 const combinedl = ((xxh_u32)c1 <<16) | ((xxh_u32)c2 << 24) - | ((xxh_u32)c3 << 0) | ((xxh_u32)len << 8); - xxh_u32 const combinedh = XXH_rotl32(XXH_swap32(combinedl), 13); - xxh_u64 const bitflipl = (XXH_readLE32(secret) ^ XXH_readLE32(secret+4)) + seed; - xxh_u64 const bitfliph = (XXH_readLE32(secret+8) ^ XXH_readLE32(secret+12)) - seed; - xxh_u64 const keyed_lo = (xxh_u64)combinedl ^ bitflipl; - xxh_u64 const keyed_hi = (xxh_u64)combinedh ^ bitfliph; - XXH128_hash_t h128; - h128.low64 = XXH64_avalanche(keyed_lo); - h128.high64 = XXH64_avalanche(keyed_hi); - return h128; - } -} - -XXH_FORCE_INLINE XXH_PUREF XXH128_hash_t -XXH3_len_4to8_128b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(input != NULL); - XXH_ASSERT(secret != NULL); - XXH_ASSERT(4 <= len && len <= 8); - seed ^= (xxh_u64)XXH_swap32((xxh_u32)seed) << 32; - { xxh_u32 const input_lo = XXH_readLE32(input); - xxh_u32 const input_hi = XXH_readLE32(input + len - 4); - xxh_u64 const input_64 = input_lo + ((xxh_u64)input_hi << 32); - xxh_u64 const bitflip = (XXH_readLE64(secret+16) ^ XXH_readLE64(secret+24)) + seed; - xxh_u64 const keyed = input_64 ^ bitflip; - - /* Shift len to the left to ensure it is even, this avoids even multiplies. */ - XXH128_hash_t m128 = XXH_mult64to128(keyed, XXH_PRIME64_1 + (len << 2)); - - m128.high64 += (m128.low64 << 1); - m128.low64 ^= (m128.high64 >> 3); - - m128.low64 = XXH_xorshift64(m128.low64, 35); - m128.low64 *= PRIME_MX2; - m128.low64 = XXH_xorshift64(m128.low64, 28); - m128.high64 = XXH3_avalanche(m128.high64); - return m128; - } -} - -XXH_FORCE_INLINE XXH_PUREF XXH128_hash_t -XXH3_len_9to16_128b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(input != NULL); - XXH_ASSERT(secret != NULL); - XXH_ASSERT(9 <= len && len <= 16); - { xxh_u64 const bitflipl = (XXH_readLE64(secret+32) ^ XXH_readLE64(secret+40)) - seed; - xxh_u64 const bitfliph = (XXH_readLE64(secret+48) ^ XXH_readLE64(secret+56)) + seed; - xxh_u64 const input_lo = XXH_readLE64(input); - xxh_u64 input_hi = XXH_readLE64(input + len - 8); - XXH128_hash_t m128 = XXH_mult64to128(input_lo ^ input_hi ^ bitflipl, XXH_PRIME64_1); - /* - * Put len in the middle of m128 to ensure that the length gets mixed to - * both the low and high bits in the 128x64 multiply below. - */ - m128.low64 += (xxh_u64)(len - 1) << 54; - input_hi ^= bitfliph; - /* - * Add the high 32 bits of input_hi to the high 32 bits of m128, then - * add the long product of the low 32 bits of input_hi and XXH_PRIME32_2 to - * the high 64 bits of m128. - * - * The best approach to this operation is different on 32-bit and 64-bit. - */ - if (sizeof(void *) < sizeof(xxh_u64)) { /* 32-bit */ - /* - * 32-bit optimized version, which is more readable. - * - * On 32-bit, it removes an ADC and delays a dependency between the two - * halves of m128.high64, but it generates an extra mask on 64-bit. - */ - m128.high64 += (input_hi & 0xFFFFFFFF00000000ULL) + XXH_mult32to64((xxh_u32)input_hi, XXH_PRIME32_2); - } else { - /* - * 64-bit optimized (albeit more confusing) version. - * - * Uses some properties of addition and multiplication to remove the mask: - * - * Let: - * a = input_hi.lo = (input_hi & 0x00000000FFFFFFFF) - * b = input_hi.hi = (input_hi & 0xFFFFFFFF00000000) - * c = XXH_PRIME32_2 - * - * a + (b * c) - * Inverse Property: x + y - x == y - * a + (b * (1 + c - 1)) - * Distributive Property: x * (y + z) == (x * y) + (x * z) - * a + (b * 1) + (b * (c - 1)) - * Identity Property: x * 1 == x - * a + b + (b * (c - 1)) - * - * Substitute a, b, and c: - * input_hi.hi + input_hi.lo + ((xxh_u64)input_hi.lo * (XXH_PRIME32_2 - 1)) - * - * Since input_hi.hi + input_hi.lo == input_hi, we get this: - * input_hi + ((xxh_u64)input_hi.lo * (XXH_PRIME32_2 - 1)) - */ - m128.high64 += input_hi + XXH_mult32to64((xxh_u32)input_hi, XXH_PRIME32_2 - 1); - } - /* m128 ^= XXH_swap64(m128 >> 64); */ - m128.low64 ^= XXH_swap64(m128.high64); - - { /* 128x64 multiply: h128 = m128 * XXH_PRIME64_2; */ - XXH128_hash_t h128 = XXH_mult64to128(m128.low64, XXH_PRIME64_2); - h128.high64 += m128.high64 * XXH_PRIME64_2; - - h128.low64 = XXH3_avalanche(h128.low64); - h128.high64 = XXH3_avalanche(h128.high64); - return h128; - } } -} - -/* - * Assumption: `secret` size is >= XXH3_SECRET_SIZE_MIN - */ -XXH_FORCE_INLINE XXH_PUREF XXH128_hash_t -XXH3_len_0to16_128b(const xxh_u8* input, size_t len, const xxh_u8* secret, XXH64_hash_t seed) -{ - XXH_ASSERT(len <= 16); - { if (len > 8) return XXH3_len_9to16_128b(input, len, secret, seed); - if (len >= 4) return XXH3_len_4to8_128b(input, len, secret, seed); - if (len) return XXH3_len_1to3_128b(input, len, secret, seed); - { XXH128_hash_t h128; - xxh_u64 const bitflipl = XXH_readLE64(secret+64) ^ XXH_readLE64(secret+72); - xxh_u64 const bitfliph = XXH_readLE64(secret+80) ^ XXH_readLE64(secret+88); - h128.low64 = XXH64_avalanche(seed ^ bitflipl); - h128.high64 = XXH64_avalanche( seed ^ bitfliph); - return h128; - } } -} - -/* - * A bit slower than XXH3_mix16B, but handles multiply by zero better. - */ -XXH_FORCE_INLINE XXH128_hash_t -XXH128_mix32B(XXH128_hash_t acc, const xxh_u8* input_1, const xxh_u8* input_2, - const xxh_u8* secret, XXH64_hash_t seed) -{ - acc.low64 += XXH3_mix16B (input_1, secret+0, seed); - acc.low64 ^= XXH_readLE64(input_2) + XXH_readLE64(input_2 + 8); - acc.high64 += XXH3_mix16B (input_2, secret+16, seed); - acc.high64 ^= XXH_readLE64(input_1) + XXH_readLE64(input_1 + 8); - return acc; -} - - -XXH_FORCE_INLINE XXH_PUREF XXH128_hash_t -XXH3_len_17to128_128b(const xxh_u8* XXH_RESTRICT input, size_t len, - const xxh_u8* XXH_RESTRICT secret, size_t secretSize, - XXH64_hash_t seed) -{ - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); (void)secretSize; - XXH_ASSERT(16 < len && len <= 128); - - { XXH128_hash_t acc; - acc.low64 = len * XXH_PRIME64_1; - acc.high64 = 0; - -#if XXH_SIZE_OPT >= 1 - { - /* Smaller, but slightly slower. */ - unsigned int i = (unsigned int)(len - 1) / 32; - do { - acc = XXH128_mix32B(acc, input+16*i, input+len-16*(i+1), secret+32*i, seed); - } while (i-- != 0); - } -#else - if (len > 32) { - if (len > 64) { - if (len > 96) { - acc = XXH128_mix32B(acc, input+48, input+len-64, secret+96, seed); - } - acc = XXH128_mix32B(acc, input+32, input+len-48, secret+64, seed); - } - acc = XXH128_mix32B(acc, input+16, input+len-32, secret+32, seed); - } - acc = XXH128_mix32B(acc, input, input+len-16, secret, seed); -#endif - { XXH128_hash_t h128; - h128.low64 = acc.low64 + acc.high64; - h128.high64 = (acc.low64 * XXH_PRIME64_1) - + (acc.high64 * XXH_PRIME64_4) - + ((len - seed) * XXH_PRIME64_2); - h128.low64 = XXH3_avalanche(h128.low64); - h128.high64 = (XXH64_hash_t)0 - XXH3_avalanche(h128.high64); - return h128; - } - } -} - -XXH_NO_INLINE XXH_PUREF XXH128_hash_t -XXH3_len_129to240_128b(const xxh_u8* XXH_RESTRICT input, size_t len, - const xxh_u8* XXH_RESTRICT secret, size_t secretSize, - XXH64_hash_t seed) -{ - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); (void)secretSize; - XXH_ASSERT(128 < len && len <= XXH3_MIDSIZE_MAX); - - { XXH128_hash_t acc; - unsigned i; - acc.low64 = len * XXH_PRIME64_1; - acc.high64 = 0; - /* - * We set as `i` as offset + 32. We do this so that unchanged - * `len` can be used as upper bound. This reaches a sweet spot - * where both x86 and aarch64 get simple agen and good codegen - * for the loop. - */ - for (i = 32; i < 160; i += 32) { - acc = XXH128_mix32B(acc, - input + i - 32, - input + i - 16, - secret + i - 32, - seed); - } - acc.low64 = XXH3_avalanche(acc.low64); - acc.high64 = XXH3_avalanche(acc.high64); - /* - * NB: `i <= len` will duplicate the last 32-bytes if - * len % 32 was zero. This is an unfortunate necessity to keep - * the hash result stable. - */ - for (i=160; i <= len; i += 32) { - acc = XXH128_mix32B(acc, - input + i - 32, - input + i - 16, - secret + XXH3_MIDSIZE_STARTOFFSET + i - 160, - seed); - } - /* last bytes */ - acc = XXH128_mix32B(acc, - input + len - 16, - input + len - 32, - secret + XXH3_SECRET_SIZE_MIN - XXH3_MIDSIZE_LASTOFFSET - 16, - (XXH64_hash_t)0 - seed); - - { XXH128_hash_t h128; - h128.low64 = acc.low64 + acc.high64; - h128.high64 = (acc.low64 * XXH_PRIME64_1) - + (acc.high64 * XXH_PRIME64_4) - + ((len - seed) * XXH_PRIME64_2); - h128.low64 = XXH3_avalanche(h128.low64); - h128.high64 = (XXH64_hash_t)0 - XXH3_avalanche(h128.high64); - return h128; - } - } -} - -XXH_FORCE_INLINE XXH128_hash_t -XXH3_hashLong_128b_internal(const void* XXH_RESTRICT input, size_t len, - const xxh_u8* XXH_RESTRICT secret, size_t secretSize, - XXH3_f_accumulate f_acc, - XXH3_f_scrambleAcc f_scramble) -{ - XXH_ALIGN(XXH_ACC_ALIGN) xxh_u64 acc[XXH_ACC_NB] = XXH3_INIT_ACC; - - XXH3_hashLong_internal_loop(acc, (const xxh_u8*)input, len, secret, secretSize, f_acc, f_scramble); - - /* converge into final hash */ - XXH_STATIC_ASSERT(sizeof(acc) == 64); - XXH_ASSERT(secretSize >= sizeof(acc) + XXH_SECRET_MERGEACCS_START); - { XXH128_hash_t h128; - h128.low64 = XXH3_mergeAccs(acc, - secret + XXH_SECRET_MERGEACCS_START, - (xxh_u64)len * XXH_PRIME64_1); - h128.high64 = XXH3_mergeAccs(acc, - secret + secretSize - - sizeof(acc) - XXH_SECRET_MERGEACCS_START, - ~((xxh_u64)len * XXH_PRIME64_2)); - return h128; - } -} - -/* - * It's important for performance that XXH3_hashLong() is not inlined. - */ -XXH_NO_INLINE XXH_PUREF XXH128_hash_t -XXH3_hashLong_128b_default(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed64, - const void* XXH_RESTRICT secret, size_t secretLen) -{ - (void)seed64; (void)secret; (void)secretLen; - return XXH3_hashLong_128b_internal(input, len, XXH3_kSecret, sizeof(XXH3_kSecret), - XXH3_accumulate, XXH3_scrambleAcc); -} - -/* - * It's important for performance to pass @p secretLen (when it's static) - * to the compiler, so that it can properly optimize the vectorized loop. - * - * When the secret size is unknown, or on GCC 12 where the mix of NO_INLINE and FORCE_INLINE - * breaks -Og, this is XXH_NO_INLINE. - */ -XXH3_WITH_SECRET_INLINE XXH128_hash_t -XXH3_hashLong_128b_withSecret(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed64, - const void* XXH_RESTRICT secret, size_t secretLen) -{ - (void)seed64; - return XXH3_hashLong_128b_internal(input, len, (const xxh_u8*)secret, secretLen, - XXH3_accumulate, XXH3_scrambleAcc); -} - -XXH_FORCE_INLINE XXH128_hash_t -XXH3_hashLong_128b_withSeed_internal(const void* XXH_RESTRICT input, size_t len, - XXH64_hash_t seed64, - XXH3_f_accumulate f_acc, - XXH3_f_scrambleAcc f_scramble, - XXH3_f_initCustomSecret f_initSec) -{ - if (seed64 == 0) - return XXH3_hashLong_128b_internal(input, len, - XXH3_kSecret, sizeof(XXH3_kSecret), - f_acc, f_scramble); - { XXH_ALIGN(XXH_SEC_ALIGN) xxh_u8 secret[XXH_SECRET_DEFAULT_SIZE]; - f_initSec(secret, seed64); - return XXH3_hashLong_128b_internal(input, len, (const xxh_u8*)secret, sizeof(secret), - f_acc, f_scramble); - } -} - -/* - * It's important for performance that XXH3_hashLong is not inlined. - */ -XXH_NO_INLINE XXH128_hash_t -XXH3_hashLong_128b_withSeed(const void* input, size_t len, - XXH64_hash_t seed64, const void* XXH_RESTRICT secret, size_t secretLen) -{ - (void)secret; (void)secretLen; - return XXH3_hashLong_128b_withSeed_internal(input, len, seed64, - XXH3_accumulate, XXH3_scrambleAcc, XXH3_initCustomSecret); -} - -typedef XXH128_hash_t (*XXH3_hashLong128_f)(const void* XXH_RESTRICT, size_t, - XXH64_hash_t, const void* XXH_RESTRICT, size_t); - -XXH_FORCE_INLINE XXH128_hash_t -XXH3_128bits_internal(const void* input, size_t len, - XXH64_hash_t seed64, const void* XXH_RESTRICT secret, size_t secretLen, - XXH3_hashLong128_f f_hl128) -{ - XXH_ASSERT(secretLen >= XXH3_SECRET_SIZE_MIN); - /* - * If an action is to be taken if `secret` conditions are not respected, - * it should be done here. - * For now, it's a contract pre-condition. - * Adding a check and a branch here would cost performance at every hash. - */ - if (len <= 16) - return XXH3_len_0to16_128b((const xxh_u8*)input, len, (const xxh_u8*)secret, seed64); - if (len <= 128) - return XXH3_len_17to128_128b((const xxh_u8*)input, len, (const xxh_u8*)secret, secretLen, seed64); - if (len <= XXH3_MIDSIZE_MAX) - return XXH3_len_129to240_128b((const xxh_u8*)input, len, (const xxh_u8*)secret, secretLen, seed64); - return f_hl128(input, len, seed64, secret, secretLen); -} - - -/* === Public XXH128 API === */ - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH128_hash_t XXH3_128bits(XXH_NOESCAPE const void* input, size_t len) -{ - return XXH3_128bits_internal(input, len, 0, - XXH3_kSecret, sizeof(XXH3_kSecret), - XXH3_hashLong_128b_default); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH128_hash_t -XXH3_128bits_withSecret(XXH_NOESCAPE const void* input, size_t len, XXH_NOESCAPE const void* secret, size_t secretSize) -{ - return XXH3_128bits_internal(input, len, 0, - (const xxh_u8*)secret, secretSize, - XXH3_hashLong_128b_withSecret); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH128_hash_t -XXH3_128bits_withSeed(XXH_NOESCAPE const void* input, size_t len, XXH64_hash_t seed) -{ - return XXH3_128bits_internal(input, len, seed, - XXH3_kSecret, sizeof(XXH3_kSecret), - XXH3_hashLong_128b_withSeed); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH128_hash_t -XXH3_128bits_withSecretandSeed(XXH_NOESCAPE const void* input, size_t len, XXH_NOESCAPE const void* secret, size_t secretSize, XXH64_hash_t seed) -{ - if (len <= XXH3_MIDSIZE_MAX) - return XXH3_128bits_internal(input, len, seed, XXH3_kSecret, sizeof(XXH3_kSecret), NULL); - return XXH3_hashLong_128b_withSecret(input, len, seed, secret, secretSize); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH128_hash_t -XXH128(XXH_NOESCAPE const void* input, size_t len, XXH64_hash_t seed) -{ - return XXH3_128bits_withSeed(input, len, seed); -} - - -/* === XXH3 128-bit streaming === */ -#ifndef XXH_NO_STREAM -/* - * All initialization and update functions are identical to 64-bit streaming variant. - * The only difference is the finalization routine. - */ - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_128bits_reset(XXH_NOESCAPE XXH3_state_t* statePtr) -{ - return XXH3_64bits_reset(statePtr); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_128bits_reset_withSecret(XXH_NOESCAPE XXH3_state_t* statePtr, XXH_NOESCAPE const void* secret, size_t secretSize) -{ - return XXH3_64bits_reset_withSecret(statePtr, secret, secretSize); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_128bits_reset_withSeed(XXH_NOESCAPE XXH3_state_t* statePtr, XXH64_hash_t seed) -{ - return XXH3_64bits_reset_withSeed(statePtr, seed); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_128bits_reset_withSecretandSeed(XXH_NOESCAPE XXH3_state_t* statePtr, XXH_NOESCAPE const void* secret, size_t secretSize, XXH64_hash_t seed) -{ - return XXH3_64bits_reset_withSecretandSeed(statePtr, secret, secretSize, seed); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_128bits_update(XXH_NOESCAPE XXH3_state_t* state, XXH_NOESCAPE const void* input, size_t len) -{ - return XXH3_64bits_update(state, input, len); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH128_hash_t XXH3_128bits_digest (XXH_NOESCAPE const XXH3_state_t* state) -{ - const unsigned char* const secret = (state->extSecret == NULL) ? state->customSecret : state->extSecret; - if (state->totalLen > XXH3_MIDSIZE_MAX) { - XXH_ALIGN(XXH_ACC_ALIGN) XXH64_hash_t acc[XXH_ACC_NB]; - XXH3_digest_long(acc, state, secret); - XXH_ASSERT(state->secretLimit + XXH_STRIPE_LEN >= sizeof(acc) + XXH_SECRET_MERGEACCS_START); - { XXH128_hash_t h128; - h128.low64 = XXH3_mergeAccs(acc, - secret + XXH_SECRET_MERGEACCS_START, - (xxh_u64)state->totalLen * XXH_PRIME64_1); - h128.high64 = XXH3_mergeAccs(acc, - secret + state->secretLimit + XXH_STRIPE_LEN - - sizeof(acc) - XXH_SECRET_MERGEACCS_START, - ~((xxh_u64)state->totalLen * XXH_PRIME64_2)); - return h128; - } - } - /* len <= XXH3_MIDSIZE_MAX : short code */ - if (state->useSeed) - return XXH3_128bits_withSeed(state->buffer, (size_t)state->totalLen, state->seed); - return XXH3_128bits_withSecret(state->buffer, (size_t)(state->totalLen), - secret, state->secretLimit + XXH_STRIPE_LEN); -} -#endif /* !XXH_NO_STREAM */ -/* 128-bit utility functions */ - -#include /* memcmp, memcpy */ - -/* return : 1 is equal, 0 if different */ -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API int XXH128_isEqual(XXH128_hash_t h1, XXH128_hash_t h2) -{ - /* note : XXH128_hash_t is compact, it has no padding byte */ - return !(memcmp(&h1, &h2, sizeof(h1))); -} - -/* This prototype is compatible with stdlib's qsort(). - * @return : >0 if *h128_1 > *h128_2 - * <0 if *h128_1 < *h128_2 - * =0 if *h128_1 == *h128_2 */ -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API int XXH128_cmp(XXH_NOESCAPE const void* h128_1, XXH_NOESCAPE const void* h128_2) -{ - XXH128_hash_t const h1 = *(const XXH128_hash_t*)h128_1; - XXH128_hash_t const h2 = *(const XXH128_hash_t*)h128_2; - int const hcmp = (h1.high64 > h2.high64) - (h2.high64 > h1.high64); - /* note : bets that, in most cases, hash values are different */ - if (hcmp) return hcmp; - return (h1.low64 > h2.low64) - (h2.low64 > h1.low64); -} - - -/*====== Canonical representation ======*/ -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API void -XXH128_canonicalFromHash(XXH_NOESCAPE XXH128_canonical_t* dst, XXH128_hash_t hash) -{ - XXH_STATIC_ASSERT(sizeof(XXH128_canonical_t) == sizeof(XXH128_hash_t)); - if (XXH_CPU_LITTLE_ENDIAN) { - hash.high64 = XXH_swap64(hash.high64); - hash.low64 = XXH_swap64(hash.low64); - } - XXH_memcpy(dst, &hash.high64, sizeof(hash.high64)); - XXH_memcpy((char*)dst + sizeof(hash.high64), &hash.low64, sizeof(hash.low64)); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH128_hash_t -XXH128_hashFromCanonical(XXH_NOESCAPE const XXH128_canonical_t* src) -{ - XXH128_hash_t h; - h.high64 = XXH_readBE64(src); - h.low64 = XXH_readBE64(src->digest + 8); - return h; -} - - - -/* ========================================== - * Secret generators - * ========================================== - */ -#define XXH_MIN(x, y) (((x) > (y)) ? (y) : (x)) - -XXH_FORCE_INLINE void XXH3_combine16(void* dst, XXH128_hash_t h128) -{ - XXH_writeLE64( dst, XXH_readLE64(dst) ^ h128.low64 ); - XXH_writeLE64( (char*)dst+8, XXH_readLE64((char*)dst+8) ^ h128.high64 ); -} - -/*! @ingroup XXH3_family */ -XXH_PUBLIC_API XXH_errorcode -XXH3_generateSecret(XXH_NOESCAPE void* secretBuffer, size_t secretSize, XXH_NOESCAPE const void* customSeed, size_t customSeedSize) -{ -#if (XXH_DEBUGLEVEL >= 1) - XXH_ASSERT(secretBuffer != NULL); - XXH_ASSERT(secretSize >= XXH3_SECRET_SIZE_MIN); -#else - /* production mode, assert() are disabled */ - if (secretBuffer == NULL) return XXH_ERROR; - if (secretSize < XXH3_SECRET_SIZE_MIN) return XXH_ERROR; -#endif - - if (customSeedSize == 0) { - customSeed = XXH3_kSecret; - customSeedSize = XXH_SECRET_DEFAULT_SIZE; - } -#if (XXH_DEBUGLEVEL >= 1) - XXH_ASSERT(customSeed != NULL); -#else - if (customSeed == NULL) return XXH_ERROR; -#endif - - /* Fill secretBuffer with a copy of customSeed - repeat as needed */ - { size_t pos = 0; - while (pos < secretSize) { - size_t const toCopy = XXH_MIN((secretSize - pos), customSeedSize); - memcpy((char*)secretBuffer + pos, customSeed, toCopy); - pos += toCopy; - } } - - { size_t const nbSeg16 = secretSize / 16; - size_t n; - XXH128_canonical_t scrambler; - XXH128_canonicalFromHash(&scrambler, XXH128(customSeed, customSeedSize, 0)); - for (n=0; n -#include -#include -#include -#include /* abort() */ -#include -#include -#include - -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -#include "xxhash/xxhash.h" -#include "sha1/sha1.h" -#include "sha256/sha256.h" - -#ifdef __cplusplus -} -#endif - - -// uuid.uuid5(uuid.NAMESPACE_URL, 'en.wikipedia.org/wiki/Llama.cpp') -#define UUID_NAMESPACE_LLAMA_CPP "ef001206-dadc-5f6d-a15f-3359e577d4e5" -#define UUID_NAMESPACE_LLAMA_CPP_HEX 0xef, 0x00, 0x12, 0x06, 0xda, 0xdc, 0x5f, 0x6d, 0xa1, 0x5f, 0x33, 0x59, 0xe5, 0x77, 0xd4, 0xe5 - - -#define HASH_TYPE_SHA256_STR "sha256" -#define HASH_TYPE_SHA1_STR "sha1" -#define HASH_TYPE_XXH64_STR "xxh64" -#define HASH_TYPE_UUID_STR "uuid" - - -typedef enum { - HASH_EXIT_SUCCESS = 0, // All hash has been generated or validated - HASH_EXIT_FAILURE = 1, // Generic Failure - HASH_EXIT_MISMATCH = 2, // Hash mismatched during validation - HASH_EXIT_MANIFEST_MISSING_ENTRY = 3, // Hash attempted validation but missing entry in manifest - HASH_EXIT_MANIFEST_UNKNOWN_HASH = 4, // Manifest is present, but we do not know any hash format within it - HASH_EXIT_MANIFEST_FILE_ERROR = 5 // Manifest is either missing or not a known format -} hash_exit_code_t; - - -typedef enum { - HASH_MANIFEST_NOT_FOUND, - HASH_MANIFEST_MISMATCH, - HASH_MANIFEST_OK, -} hash_manifest_result_t; - - -struct hash_params { - std::string input; - bool xxh64 = false; - bool sha1 = false; - bool sha256 = false; - bool uuid = false; - - bool no_layer = false; - - bool manifest_is_usable = false; - std::string manifest_file; -}; - -struct manifest_check_params { - bool xxh64 = false; - bool sha1 = false; - bool sha256 = false; - bool uuid = false; -}; - -static char const * hash_manifest_result_to_str(hash_manifest_result_t value) { - switch (value) { - case HASH_MANIFEST_NOT_FOUND: return "Not Found"; - case HASH_MANIFEST_MISMATCH: return "Mismatch"; - case HASH_MANIFEST_OK: return "Ok"; - } - return "?"; -} - -static char const * hash_exit_code_to_str(hash_exit_code_t value) { - switch (value) { - case HASH_EXIT_SUCCESS: return "Success"; - case HASH_EXIT_FAILURE: return "Failure"; - case HASH_EXIT_MISMATCH: return "Mismatch"; - case HASH_EXIT_MANIFEST_MISSING_ENTRY: return "Manifest Missing Entry"; - case HASH_EXIT_MANIFEST_UNKNOWN_HASH: return "Manifest Unknown Hash"; - case HASH_EXIT_MANIFEST_FILE_ERROR: return "Manifest File Error"; - } - return "?"; -} - -static void hash_print_usage(const char * executable) { - const hash_params default_params; - printf("\n"); - printf("usage: %s [options] GGUF_IN\n", executable); - printf("\n"); - printf("Hash a GGUF file"); - printf("\n"); - printf("options:\n"); - printf(" -h, --help show this help message and exit\n"); - printf(" --xxh64 use xxh64 hash\n"); - printf(" --sha1 use sha1 hash\n"); - printf(" --sha256 use sha256 hash\n"); - printf(" --all use all hash\n"); - printf(" --no-layer exclude per layer hash\n"); - printf(" --uuid generate UUIDv5 ID\n"); - printf(" -c, --check verify against a manifest\n"); - printf("\n"); -} - -static void hash_params_parse_ex(int argc, const char ** argv, hash_params & params) { - std::string arg; - bool invalid_param = false; - const std::string arg_prefix = "--"; - - int arg_idx = 1; - for (; arg_idx < argc && strncmp(argv[arg_idx], "--", 2) == 0; arg_idx++) { - arg = argv[arg_idx]; - if (arg.compare(0, arg_prefix.size(), arg_prefix) == 0) { - std::replace(arg.begin(), arg.end(), '_', '-'); - } - - bool arg_found = false; - if (arg == "-h" || arg == "--help") { - hash_print_usage(argv[0]); - exit(0); - } - - if (arg == "--xxh64") { - arg_found = true; - params.xxh64 = true; - } - - if (arg == "--sha1") { - arg_found = true; - params.sha1 = true; - } - - if (arg == "--uuid") { - arg_found = true; - params.uuid = true; - } - - if (arg == "--sha256") { - arg_found = true; - params.sha256 = true; - } - - if (arg == "--all") { - arg_found = true; - params.sha256 = true; - params.sha1 = true; - params.xxh64 = true; - } - - if (arg == "--no-layer") { - arg_found = true; - params.no_layer = true; - } - - if (arg == "-c" || arg == "--check") { - if (++arg_idx >= argc) { - invalid_param = true; - break; - } - arg_found = true; - params.manifest_file = argv[arg_idx]; - } - - if (!arg_found) { - throw std::invalid_argument("error: unknown argument: " + arg); - } - } - - if (invalid_param) { - throw std::invalid_argument("error: invalid parameter for argument:" + arg); - } - - if (argc - arg_idx < 1) { - throw std::invalid_argument("error: bad arguments"); - } - - params.input = argv[arg_idx++]; -} - -static bool hash_params_parse(int argc, const char ** argv, hash_params & params) { - bool result = true; - try { - hash_params_parse_ex(argc, argv, params); - } - catch (const std::invalid_argument & ex) { - fprintf(stderr, "%s\n", ex.what()); - hash_print_usage(argv[0]); - exit(EXIT_FAILURE); - } - return result; -} - -static bool manifest_type(const std::string & manifest_file, manifest_check_params & manifest_check) { - if (manifest_file.empty()) { - return false; - } - - std::ifstream file(manifest_file); - if (!file.is_open()) { - return false; - } - - std::string manifest_entry_line; - while (getline(file, manifest_entry_line)) { - // hash_type_str hash_str tensor_name - // e.g. 'xxh64 f66e9cd66a4396a0 test.gguf:tensor_0' - std::istringstream line_stream(manifest_entry_line); - std::string file_hash_type; - if (line_stream >> file_hash_type) { - if (file_hash_type == HASH_TYPE_SHA256_STR) { - manifest_check.sha256 = true; - } else if (file_hash_type == HASH_TYPE_SHA1_STR) { - manifest_check.sha1 = true; - } else if (file_hash_type == HASH_TYPE_XXH64_STR) { - manifest_check.xxh64 = true; - } else if (file_hash_type == HASH_TYPE_UUID_STR) { - manifest_check.uuid = true; - } - } - } - - return true; -} - -static hash_manifest_result_t manifest_verify(const std::string& manifest_file, const std::string& hash_type_str, const std::string& hash_str, const std::string& tensor_name) { - if (manifest_file.empty()) { - return HASH_MANIFEST_NOT_FOUND; - } - - std::ifstream file(manifest_file); - if (!file.is_open()) { - return HASH_MANIFEST_NOT_FOUND; - } - - std::string manifest_entry_line; - while (getline(file, manifest_entry_line)) { - std::istringstream line_stream(manifest_entry_line); - std::string file_hash_type; - std::string file_hash; - std::string file_tensor_name; - if (line_stream >> file_hash_type >> file_hash >> file_tensor_name) { - // Line parsed. Check hash validity - - if (file_hash_type != hash_type_str) { - continue; - } - - if (file_tensor_name != tensor_name) { - continue; - } - - return (file_hash == hash_str) ? HASH_MANIFEST_OK : HASH_MANIFEST_MISMATCH; - } - } - - return HASH_MANIFEST_NOT_FOUND; -} - -static void generate_uuidv5(const unsigned char sha1_digest[20], unsigned char uuid[16]) { - // Ref: https://www.rfc-editor.org/rfc/rfc9562.html#section-5.5 - // Assumes that digest was processed correctly with the expected namespace - for (int i = 0; i < 16; i++) { - uuid[i] = sha1_digest[i]; - } - - // Set bits corresponding to UUID ver 5 - uuid[ 6] &= ~(0xF << 4); - uuid[ 6] |= (5 << 4); - - // Set bits corresponding to UUID variant 0b10XX - uuid[ 8] &= ~(0xc << 4); - uuid[ 8] |= (0x8 << 4); -} - -static hash_exit_code_t gguf_hash(const hash_params & hash_params) { - const std::string & fname = hash_params.input; - struct ggml_context * ctx_data = NULL; - - struct gguf_init_params params = { - /*.no_alloc = */ false, - /*.ctx = */ &ctx_data, - }; - - // xxh64 init - XXH64_state_t* xxh64_model_hash_state = NULL; - if (hash_params.xxh64) { - xxh64_model_hash_state = XXH64_createState(); - if (xxh64_model_hash_state==NULL) { - abort(); - } - - XXH64_hash_t const seed = 0; - if (XXH64_reset(xxh64_model_hash_state, seed) == XXH_ERROR) { - abort(); - } - } - - // sha1 init - SHA1_CTX sha1_model_hash_ctx; - if (hash_params.sha1) { - SHA1Init(&sha1_model_hash_ctx); - } - - // sha256 init - sha256_t sha256_model_hash_ctx; - if (hash_params.sha256) { - sha256_init(&sha256_model_hash_ctx); - } - - // sha1 for uuid init - SHA1_CTX sha1_for_uuid_ctx; - if (hash_params.uuid) { - unsigned char const uuidv5_namespace[] = {UUID_NAMESPACE_LLAMA_CPP_HEX}; - SHA1Init(&sha1_for_uuid_ctx); - SHA1Update( &sha1_for_uuid_ctx, (unsigned char const *)uuidv5_namespace, sizeof(uuidv5_namespace)); - } - - struct gguf_context * ctx = gguf_init_from_file(fname.c_str(), params); - const int n_tensors = gguf_get_n_tensors(ctx); - bool tensor_layer_in_manifest = false; - bool model_in_manifest = false; - bool tensor_layer_has_mismatch = false; - bool model_has_mismatch = false; - for (int i = 0; i < n_tensors; ++i) { - const char * name = gguf_get_tensor_name(ctx, i); - struct ggml_tensor * cur = ggml_get_tensor(ctx_data, name); - auto n_bytes = ggml_nbytes(cur); - auto *raw_data = cur->data; - const std::string tensor_layer_name = fname + ":" + name; - - if (hash_params.xxh64) { - - if (!hash_params.no_layer) { - // Per Layer Hash - XXH64_hash_t hash = XXH64(raw_data, n_bytes, 0); - - char hex_result[17]; - for (int offset = 0; offset < 8; offset++) { - unsigned int shift_bits_by = (8 * (8 - offset - 1)); - snprintf( ( hex_result + (2*offset)), sizeof(hex_result) - (2*offset), "%02x", (unsigned char) (hash >> shift_bits_by)&0xff); - } - - if (hash_params.manifest_is_usable) { - hash_manifest_result_t verify_result = manifest_verify(hash_params.manifest_file, HASH_TYPE_XXH64_STR, hex_result, tensor_layer_name); - - switch (verify_result) { - case HASH_MANIFEST_NOT_FOUND: - break; - case HASH_MANIFEST_MISMATCH: - tensor_layer_in_manifest = true; - tensor_layer_has_mismatch = true; - break; - case HASH_MANIFEST_OK: - tensor_layer_in_manifest = true; - break; - } - - printf("%-8s %-s %s - %s\n", HASH_TYPE_XXH64_STR, hex_result, tensor_layer_name.c_str(), hash_manifest_result_to_str(verify_result)); - } else { - printf("%-8s %-s %s\n", HASH_TYPE_XXH64_STR, hex_result, tensor_layer_name.c_str()); - } - } - - // Overall Model Hash - if (XXH64_update(xxh64_model_hash_state, raw_data, n_bytes) == XXH_ERROR) abort(); - } - - if (hash_params.sha1) { - - if (!hash_params.no_layer) { - // Per Layer Hash - char result[21]; // sha1 outputs 20 bytes - SHA1( result, (const char *)raw_data, n_bytes); - - char hex_result[41] = {0}; - for (int offset = 0; offset < 20; offset++) { - snprintf( ( hex_result + (2*offset)), sizeof(hex_result) - (2*offset), "%02x", result[offset]&0xff); - } - - if (hash_params.manifest_is_usable) { - hash_manifest_result_t verify_result = manifest_verify(hash_params.manifest_file, HASH_TYPE_SHA1_STR, hex_result, tensor_layer_name); - - switch (verify_result) { - case HASH_MANIFEST_NOT_FOUND: - break; - case HASH_MANIFEST_MISMATCH: - tensor_layer_in_manifest = true; - tensor_layer_has_mismatch = true; - break; - case HASH_MANIFEST_OK: - tensor_layer_in_manifest = true; - break; - } - - printf("%-8s %-s %s - %s\n", HASH_TYPE_SHA1_STR, hex_result, tensor_layer_name.c_str(), hash_manifest_result_to_str(verify_result)); - } else { - printf("%-8s %-s %s\n", HASH_TYPE_SHA1_STR, hex_result, tensor_layer_name.c_str()); - } - } - - // Overall Model Hash - SHA1Update( &sha1_model_hash_ctx, (unsigned char const *)raw_data, n_bytes); - } - - if (hash_params.sha256) { - - if (!hash_params.no_layer) { - // Per Layer Hash - unsigned char result[SHA256_DIGEST_SIZE]; // sha256 outputs 32 bytes - sha256_hash((unsigned char*) result, (const unsigned char *)raw_data, n_bytes); - - char hex_result[SHA256_DIGEST_SIZE * 2 + 1] = {0}; - for (int offset = 0; offset < SHA256_DIGEST_SIZE; offset++) { - snprintf( ( hex_result + (2*offset)), sizeof(hex_result) - (2*offset), "%02x", result[offset]&0xff); - } - - if (hash_params.manifest_is_usable) { - hash_manifest_result_t verify_result = manifest_verify(hash_params.manifest_file, HASH_TYPE_SHA256_STR, hex_result, tensor_layer_name); - - switch (verify_result) { - case HASH_MANIFEST_NOT_FOUND: - break; - case HASH_MANIFEST_MISMATCH: - tensor_layer_in_manifest = true; - tensor_layer_has_mismatch = true; - break; - case HASH_MANIFEST_OK: - tensor_layer_in_manifest = true; - break; - } - - printf("%-8s %-s %s - %s\n", HASH_TYPE_SHA256_STR, hex_result, tensor_layer_name.c_str(), hash_manifest_result_to_str(verify_result)); - } else { - printf("%-8s %-s %s\n", HASH_TYPE_SHA256_STR, hex_result, tensor_layer_name.c_str()); - } - } - - // Overall Model Hash - sha256_update( &sha256_model_hash_ctx, (unsigned char const *)raw_data, n_bytes); - } - - if (hash_params.uuid) { - SHA1Update( &sha1_for_uuid_ctx, (unsigned char const *)raw_data, n_bytes); - } - } - - if (hash_params.xxh64) { - XXH64_hash_t const hash = XXH64_digest(xxh64_model_hash_state); - - char hex_result[17]; - for (int offset = 0; offset < 8; offset++) { - unsigned int shift_bits_by = (8 * (8 - offset - 1)); - snprintf( ( hex_result + (2*offset)), sizeof(hex_result) - (2*offset), "%02x", (unsigned char) (hash >> shift_bits_by)&0xff); - } - - if (hash_params.manifest_is_usable) { - hash_manifest_result_t verify_result = manifest_verify(hash_params.manifest_file, HASH_TYPE_XXH64_STR, hex_result, fname); - - switch (verify_result) { - case HASH_MANIFEST_NOT_FOUND: - break; - case HASH_MANIFEST_MISMATCH: - model_in_manifest = true; - model_has_mismatch = true; - break; - case HASH_MANIFEST_OK: - model_in_manifest = true; - break; - } - - printf("%-8s %-s %s - %s\n", HASH_TYPE_XXH64_STR, hex_result, fname.c_str(), hash_manifest_result_to_str(verify_result)); - } else { - printf("%-8s %-s %s\n", HASH_TYPE_XXH64_STR, hex_result, fname.c_str()); - } - } - - if (hash_params.sha1) { - unsigned char result[21]; - SHA1Final(result, &sha1_model_hash_ctx); - - char hex_result[41]; - for (int offset = 0; offset < 20; offset++) { - snprintf( ( hex_result + (2*offset)), sizeof(hex_result) - (2*offset), "%02x", result[offset]&0xff); - } - - if (hash_params.manifest_is_usable) { - hash_manifest_result_t verify_result = manifest_verify(hash_params.manifest_file, HASH_TYPE_SHA1_STR, hex_result, fname); - - switch (verify_result) { - case HASH_MANIFEST_NOT_FOUND: - break; - case HASH_MANIFEST_MISMATCH: - model_in_manifest = true; - model_has_mismatch = true; - break; - case HASH_MANIFEST_OK: - model_in_manifest = true; - break; - } - - printf("%-8s %-s %s - %s\n", HASH_TYPE_SHA1_STR, hex_result, fname.c_str(), hash_manifest_result_to_str(verify_result)); - } else { - printf("%-8s %-s %s\n", HASH_TYPE_SHA1_STR, hex_result, fname.c_str()); - } - } - - if (hash_params.sha256) { - unsigned char result[SHA256_DIGEST_SIZE]; // sha256 outputs 32 bytes - sha256_final( &sha256_model_hash_ctx, result); - - char hex_result[SHA256_DIGEST_SIZE * 2 + 1] = {0}; - for (int offset = 0; offset < SHA256_DIGEST_SIZE; offset++) { - snprintf( ( hex_result + (2*offset)), sizeof(hex_result) - (2*offset), "%02x", result[offset]&0xff); - } - - if (hash_params.manifest_is_usable) { - hash_manifest_result_t verify_result = manifest_verify(hash_params.manifest_file, HASH_TYPE_SHA256_STR, hex_result, fname); - - switch (verify_result) { - case HASH_MANIFEST_NOT_FOUND: - break; - case HASH_MANIFEST_MISMATCH: - model_in_manifest = true; - model_has_mismatch = true; - break; - case HASH_MANIFEST_OK: - model_in_manifest = true; - break; - } - - printf("%-8s %-s %s - %s\n", HASH_TYPE_SHA256_STR, hex_result, fname.c_str(), hash_manifest_result_to_str(verify_result)); - } else { - printf("%-8s %-s %s\n", HASH_TYPE_SHA256_STR, hex_result, fname.c_str()); - } - } - - if (hash_params.uuid) { - unsigned char result[21]; - SHA1Final(result, &sha1_for_uuid_ctx); - - unsigned char uuid[16]; - generate_uuidv5(result, uuid); - - char string_buffer[37] = {0}; - snprintf(string_buffer, sizeof(string_buffer), "%02x%02x%02x%02x-%02x%02x-%02x%02x-%02x%02x-%02x%02x%02x%02x%02x%02x", - uuid[0], uuid[1], uuid[2], uuid[3], - uuid[4], uuid[5], uuid[6], uuid[7], - uuid[8], uuid[9], uuid[10], uuid[11], - uuid[12], uuid[13], uuid[14], uuid[15]); - - if (hash_params.manifest_is_usable) { - hash_manifest_result_t verify_result = manifest_verify(hash_params.manifest_file, HASH_TYPE_SHA256_STR, string_buffer, fname); - - switch (verify_result) { - case HASH_MANIFEST_NOT_FOUND: - break; - case HASH_MANIFEST_MISMATCH: - model_in_manifest = true; - model_has_mismatch = true; - break; - case HASH_MANIFEST_OK: - model_in_manifest = true; - break; - } - - printf("%-8s %-s %s - %s\n", HASH_TYPE_UUID_STR, string_buffer, fname.c_str(), hash_manifest_result_to_str(verify_result)); - } else { - printf("%-8s %-s %s\n", HASH_TYPE_UUID_STR, string_buffer, fname.c_str()); - } - } - - - ggml_free(ctx_data); - gguf_free(ctx); - - - if (hash_params.manifest_is_usable) { - // In hash verification mode - - if (!model_in_manifest) { - // model missing in manifest? - - // Check tensor layer... - if (!tensor_layer_in_manifest) { - // Still missing? Maybe we are reading the wrong manifest. - return HASH_EXIT_MANIFEST_MISSING_ENTRY; - } - - if (tensor_layer_has_mismatch) { - // Per tensor check found error - return HASH_EXIT_FAILURE; - } - - // All per tensor layer checks passed? Sounds good enough. - return HASH_EXIT_SUCCESS; - } - - // Overall model check passed, but let's check per layer just in case - // If missing, we don't care too much as the overall model checked - if (tensor_layer_in_manifest && tensor_layer_has_mismatch) { - return HASH_EXIT_FAILURE; - } - - if (model_has_mismatch) { - // model has failed hash somewhere in the model - return HASH_EXIT_FAILURE; - } - - // All checks appears to be fine - return HASH_EXIT_SUCCESS; - } - - // In hash generation mode - return HASH_EXIT_SUCCESS; -} - -int main(int argc, const char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - hash_params params; - manifest_check_params manifest_check; - hash_params_parse(argc, argv, params); - - if (!params.manifest_file.empty()) { - if (!manifest_type(params.manifest_file, manifest_check)) { - printf("ERROR cannot open manifest %s", params.manifest_file.c_str()); - return HASH_EXIT_MANIFEST_FILE_ERROR; - } - - if (!manifest_check.sha256 && !manifest_check.sha1 && !manifest_check.xxh64 && !manifest_check.uuid) { - printf("ERROR manifest does not have any known hash format in %s", params.manifest_file.c_str()); - return HASH_EXIT_MANIFEST_UNKNOWN_HASH; - } - - printf("manifest %s", params.manifest_file.c_str()); - - if (manifest_check.sha256) { - printf(" sha256"); - } - - if (manifest_check.sha1) { - printf(" sha1"); - } - - if (manifest_check.xxh64) { - printf(" xxh64"); - } - - if (manifest_check.uuid) { - printf(" uuid"); - } - - printf("\n"); - - // Autoselect the highest security hash if manifest is provided but - // the user has not specifically defined the hash they care about - if (!params.xxh64 && !params.sha1 && !params.uuid && !params.sha256) { - // User has not selected a specific value, pick most secure hash - if (manifest_check.sha256) { - params.sha256 = true; - } else if (manifest_check.sha1) { - params.sha1 = true; - } else if (manifest_check.xxh64) { - params.xxh64 = true; - } else if (manifest_check.uuid) { - params.uuid = true; - } - } - - params.manifest_is_usable = true; - } - - // By default if no swich argument provided, assume xxh64 - if (!params.xxh64 && !params.sha1 && !params.uuid && !params.sha256) { - params.xxh64 = true; - } - - hash_exit_code_t exit_code = gguf_hash(params); - - if (params.manifest_is_usable) { - printf("\nVerification results for %s - %s\n", params.manifest_file.c_str(), hash_exit_code_to_str(exit_code)); - } - - return exit_code; -} diff --git a/examples/gguf-rust/Cargo.toml b/examples/gguf-rust/Cargo.toml new file mode 100644 index 000000000000..43ac4de21fdd --- /dev/null +++ b/examples/gguf-rust/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "llama-gguf-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[[bin]] +name = "llama-gguf" +path = "src/main.rs" diff --git a/examples/gguf-rust/build.rs b/examples/gguf-rust/build.rs new file mode 100644 index 000000000000..f0f25838c5af --- /dev/null +++ b/examples/gguf-rust/build.rs @@ -0,0 +1,47 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + if let Ok(lib_dir) = env::var("GGML_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/ggml/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libggml.so").exists() + || candidate.join("libggml.dylib").exists() + || candidate.join("ggml.dll").exists() + || candidate.join("ggml.lib").exists() + || candidate.join("libggml.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libggml was not found; build llama.cpp before building llama-gguf"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os == "linux" || target_os == "android" { + println!("cargo:rustc-link-lib=stdc++"); + } else if target_os == "macos" { + println!("cargo:rustc-link-lib=c++"); + } + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/gguf-rust/src/lib.rs b/examples/gguf-rust/src/lib.rs new file mode 100644 index 000000000000..9759b508e074 --- /dev/null +++ b/examples/gguf-rust/src/lib.rs @@ -0,0 +1,57 @@ +pub const GGML_MAX_DIMS: usize = 4; + +#[derive(Clone)] +pub struct DeterministicRng { + state: u32, +} + +impl DeterministicRng { + pub fn new(seed: u32) -> Self { + Self { state: seed } + } + + pub fn next_bounded(&mut self, bound: i64) -> i64 { + self.state = self.state.wrapping_mul(1_103_515_245).wrapping_add(12_345); + ((self.state / 65_536) % bound as u32) as i64 + } +} + +pub fn tensor_shape(rng: &mut DeterministicRng) -> ([i64; GGML_MAX_DIMS], i32) { + let mut ne = [1_i64; GGML_MAX_DIMS]; + let n_dims = rng.next_bounded(GGML_MAX_DIMS as i64) as i32 + 1; + for dim in ne.iter_mut().take(n_dims as usize) { + *dim = rng.next_bounded(10) + 1; + } + (ne, n_dims) +} + +pub fn usage(program: &str) -> String { + format!("usage: {program} data.gguf r|w [n]\nr: read data.gguf file\nw: write data.gguf file\nn: no check of tensor data\n") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn generated_shapes_keep_valid_dimensions() { + let mut rng = DeterministicRng::new(123456); + for _ in 0..32 { + let (ne, n_dims) = tensor_shape(&mut rng); + assert!((1..=4).contains(&n_dims)); + for dim in ne.iter().take(n_dims as usize) { + assert!((1..=10).contains(dim)); + } + for dim in ne.iter().skip(n_dims as usize) { + assert_eq!(*dim, 1); + } + } + } + + #[test] + fn usage_contains_modes() { + let text = usage("llama-gguf"); + assert!(text.contains("r: read")); + assert!(text.contains("w: write")); + } +} diff --git a/examples/gguf-rust/src/main.rs b/examples/gguf-rust/src/main.rs new file mode 100644 index 000000000000..dd8f09dcbcb6 --- /dev/null +++ b/examples/gguf-rust/src/main.rs @@ -0,0 +1,435 @@ +use std::ffi::{c_char, c_int, c_void, CStr, CString}; +use std::process::ExitCode; +use std::ptr; + +use llama_gguf_rust::{tensor_shape, usage, DeterministicRng, GGML_MAX_DIMS}; + +const GGML_TYPE_F32: c_int = 0; +const GGUF_TYPE_INT16: c_int = 3; +const GGUF_TYPE_FLOAT32: c_int = 6; +const LC_NUMERIC: c_int = 1; + +#[repr(C)] +struct ggml_context { + _private: [u8; 0], +} + +#[repr(C)] +struct ggml_backend_buffer { + _private: [u8; 0], +} + +#[repr(C)] +struct gguf_context { + _private: [u8; 0], +} + +#[repr(C)] +struct ggml_init_params { + mem_size: usize, + mem_buffer: *mut c_void, + no_alloc: bool, +} + +#[repr(C)] +struct gguf_init_params { + no_alloc: bool, + ctx: *mut *mut ggml_context, +} + +#[repr(C)] +struct ggml_tensor { + type_: c_int, + buffer: *mut ggml_backend_buffer, + ne: [i64; 4], + nb: [usize; 4], + op: c_int, + op_params: [i32; 16], + flags: i32, + src: [*mut ggml_tensor; 10], + view_src: *mut ggml_tensor, + view_offs: usize, + data: *mut c_void, + name: [c_char; 64], + extra: *mut c_void, + padding: [c_char; 8], +} + +extern "C" { + fn setlocale(category: c_int, locale: *const c_char) -> *mut c_char; + + fn ggml_init(params: ggml_init_params) -> *mut ggml_context; + fn ggml_free(ctx: *mut ggml_context); + fn ggml_new_tensor( + ctx: *mut ggml_context, + type_: c_int, + n_dims: c_int, + ne: *const i64, + ) -> *mut ggml_tensor; + fn ggml_set_name(tensor: *mut ggml_tensor, name: *const c_char) -> *mut ggml_tensor; + fn ggml_nelements(tensor: *const ggml_tensor) -> i64; + fn ggml_get_tensor(ctx: *mut ggml_context, name: *const c_char) -> *mut ggml_tensor; + fn ggml_n_dims(tensor: *const ggml_tensor) -> c_int; + fn ggml_get_mem_size(ctx: *const ggml_context) -> usize; + fn ggml_type_name(type_: c_int) -> *const c_char; + fn ggml_type_size(type_: c_int) -> usize; + + fn gguf_init_empty() -> *mut gguf_context; + fn gguf_init_from_file(fname: *const c_char, params: gguf_init_params) -> *mut gguf_context; + fn gguf_free(ctx: *mut gguf_context); + fn gguf_get_version(ctx: *const gguf_context) -> u32; + fn gguf_get_alignment(ctx: *const gguf_context) -> usize; + fn gguf_get_data_offset(ctx: *const gguf_context) -> usize; + fn gguf_get_n_kv(ctx: *const gguf_context) -> i64; + fn gguf_get_key(ctx: *const gguf_context, key_id: i64) -> *const c_char; + fn gguf_find_key(ctx: *const gguf_context, key: *const c_char) -> i64; + fn gguf_get_val_str(ctx: *const gguf_context, key_id: i64) -> *const c_char; + fn gguf_get_n_tensors(ctx: *const gguf_context) -> i64; + fn gguf_get_tensor_name(ctx: *const gguf_context, tensor_id: i64) -> *const c_char; + fn gguf_get_tensor_size(ctx: *const gguf_context, tensor_id: i64) -> usize; + fn gguf_get_tensor_offset(ctx: *const gguf_context, tensor_id: i64) -> usize; + fn gguf_get_tensor_type(ctx: *const gguf_context, tensor_id: i64) -> c_int; + fn gguf_set_val_u8(ctx: *mut gguf_context, key: *const c_char, val: u8); + fn gguf_set_val_i8(ctx: *mut gguf_context, key: *const c_char, val: i8); + fn gguf_set_val_u16(ctx: *mut gguf_context, key: *const c_char, val: u16); + fn gguf_set_val_i16(ctx: *mut gguf_context, key: *const c_char, val: i16); + fn gguf_set_val_u32(ctx: *mut gguf_context, key: *const c_char, val: u32); + fn gguf_set_val_i32(ctx: *mut gguf_context, key: *const c_char, val: i32); + fn gguf_set_val_f32(ctx: *mut gguf_context, key: *const c_char, val: f32); + fn gguf_set_val_u64(ctx: *mut gguf_context, key: *const c_char, val: u64); + fn gguf_set_val_i64(ctx: *mut gguf_context, key: *const c_char, val: i64); + fn gguf_set_val_f64(ctx: *mut gguf_context, key: *const c_char, val: f64); + fn gguf_set_val_bool(ctx: *mut gguf_context, key: *const c_char, val: bool); + fn gguf_set_val_str(ctx: *mut gguf_context, key: *const c_char, val: *const c_char); + fn gguf_set_arr_data( + ctx: *mut gguf_context, + key: *const c_char, + type_: c_int, + data: *const c_void, + n: usize, + ); + fn gguf_set_arr_str( + ctx: *mut gguf_context, + key: *const c_char, + data: *const *const c_char, + n: usize, + ); + fn gguf_add_tensor(ctx: *mut gguf_context, tensor: *const ggml_tensor); + fn gguf_write_to_file(ctx: *const gguf_context, fname: *const c_char, only_meta: bool) -> bool; +} + +fn main() -> ExitCode { + unsafe { + let locale = CString::new("C").unwrap(); + setlocale(LC_NUMERIC, locale.as_ptr()); + } + + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(message) => { + eprintln!("{message}"); + ExitCode::FAILURE + } + } +} + +fn run() -> Result<(), String> { + let args = std::env::args().collect::>(); + if args.len() < 3 { + print!( + "{}", + usage(args.first().map(String::as_str).unwrap_or("llama-gguf")) + ); + return Err("bad arguments".to_string()); + } + + let fname = args[1].clone(); + let mode = args[2].as_str(); + let check_data = args.len() != 4; + + match mode { + "w" => gguf_ex_write(&fname), + "r" => { + gguf_ex_read_0(&fname)?; + gguf_ex_read_1(&fname, check_data) + } + _ => Err("mode must be r or w".to_string()), + } +} + +fn gguf_ex_write(fname: &str) -> Result<(), String> { + unsafe { + let ctx = gguf_init_empty(); + if ctx.is_null() { + return Err("gguf_init_empty failed".to_string()); + } + + set_scalar_kv(ctx)?; + + let arr_i16 = [1_i16, 2, 3, 4]; + gguf_set_arr_data( + ctx, + cstring("some.parameter.arr.i16")?.as_ptr(), + GGUF_TYPE_INT16, + arr_i16.as_ptr().cast(), + arr_i16.len(), + ); + let arr_f32 = [3.145_f32, 2.718, 1.414]; + gguf_set_arr_data( + ctx, + cstring("some.parameter.arr.f32")?.as_ptr(), + GGUF_TYPE_FLOAT32, + arr_f32.as_ptr().cast(), + arr_f32.len(), + ); + let strings = [cstring("hello")?, cstring("world")?, cstring("!")?]; + let string_ptrs = strings.iter().map(|s| s.as_ptr()).collect::>(); + gguf_set_arr_str( + ctx, + cstring("some.parameter.arr.str")?.as_ptr(), + string_ptrs.as_ptr(), + string_ptrs.len(), + ); + + let ctx_data = ggml_init(ggml_init_params { + mem_size: 128 * 1024 * 1024, + mem_buffer: ptr::null_mut(), + no_alloc: false, + }); + if ctx_data.is_null() { + gguf_free(ctx); + return Err("ggml_init failed".to_string()); + } + + let mut rng = DeterministicRng::new(123456); + for i in 0..10 { + let name = cstring(&format!("tensor_{i}"))?; + let (ne, n_dims) = tensor_shape(&mut rng); + let cur = ggml_new_tensor(ctx_data, GGML_TYPE_F32, n_dims, ne.as_ptr()); + ggml_set_name(cur, name.as_ptr()); + + let n_elements = ggml_nelements(cur) as usize; + let data = std::slice::from_raw_parts_mut((*cur).data.cast::(), n_elements); + for value in data { + *value = 100.0 + i as f32; + } + + gguf_add_tensor(ctx, cur); + } + + let c_fname = cstring(fname)?; + if !gguf_write_to_file(ctx, c_fname.as_ptr(), false) { + ggml_free(ctx_data); + gguf_free(ctx); + return Err(format!("failed to write gguf file {fname}")); + } + println!("gguf_ex_write: wrote file '{fname};"); + + ggml_free(ctx_data); + gguf_free(ctx); + Ok(()) + } +} + +unsafe fn set_scalar_kv(ctx: *mut gguf_context) -> Result<(), String> { + gguf_set_val_u8(ctx, cstring("some.parameter.uint8")?.as_ptr(), 0x12); + gguf_set_val_i8(ctx, cstring("some.parameter.int8")?.as_ptr(), -0x13); + gguf_set_val_u16(ctx, cstring("some.parameter.uint16")?.as_ptr(), 0x1234); + gguf_set_val_i16(ctx, cstring("some.parameter.int16")?.as_ptr(), -0x1235); + gguf_set_val_u32(ctx, cstring("some.parameter.uint32")?.as_ptr(), 0x12345678); + gguf_set_val_i32(ctx, cstring("some.parameter.int32")?.as_ptr(), -0x12345679); + gguf_set_val_f32( + ctx, + cstring("some.parameter.float32")?.as_ptr(), + 0.123456789, + ); + gguf_set_val_u64( + ctx, + cstring("some.parameter.uint64")?.as_ptr(), + 0x123456789abcdef0, + ); + gguf_set_val_i64( + ctx, + cstring("some.parameter.int64")?.as_ptr(), + -0x123456789abcdef1, + ); + gguf_set_val_f64( + ctx, + cstring("some.parameter.float64")?.as_ptr(), + 0.1234567890123456789, + ); + gguf_set_val_bool(ctx, cstring("some.parameter.bool")?.as_ptr(), true); + gguf_set_val_str( + ctx, + cstring("some.parameter.string")?.as_ptr(), + cstring("hello world")?.as_ptr(), + ); + Ok(()) +} + +fn gguf_ex_read_0(fname: &str) -> Result<(), String> { + unsafe { + let c_fname = cstring(fname)?; + let ctx = gguf_init_from_file( + c_fname.as_ptr(), + gguf_init_params { + no_alloc: false, + ctx: ptr::null_mut(), + }, + ); + if ctx.is_null() { + return Err(format!("gguf_ex_read_0: failed to load '{fname}'")); + } + + print_header("gguf_ex_read_0", ctx); + print_kv("gguf_ex_read_0", ctx); + find_string_kv("gguf_ex_read_0", ctx)?; + print_tensor_info("gguf_ex_read_0", ctx, false); + + gguf_free(ctx); + Ok(()) + } +} + +fn gguf_ex_read_1(fname: &str, check_data: bool) -> Result<(), String> { + unsafe { + let c_fname = cstring(fname)?; + let mut ctx_data: *mut ggml_context = ptr::null_mut(); + let ctx = gguf_init_from_file( + c_fname.as_ptr(), + gguf_init_params { + no_alloc: false, + ctx: &mut ctx_data, + }, + ); + if ctx.is_null() { + return Err(format!("gguf_ex_read_1: failed to load '{fname}'")); + } + + print_header("gguf_ex_read_1", ctx); + print_kv("gguf_ex_read_1", ctx); + print_tensor_info("gguf_ex_read_1", ctx, true); + print_tensor_data("gguf_ex_read_1", ctx, ctx_data, check_data)?; + + println!( + "gguf_ex_read_1: ctx_data size: {}", + ggml_get_mem_size(ctx_data) + ); + ggml_free(ctx_data); + gguf_free(ctx); + Ok(()) + } +} + +unsafe fn print_header(prefix: &str, ctx: *const gguf_context) { + println!("{prefix}: version: {}", gguf_get_version(ctx)); + println!("{prefix}: alignment: {}", gguf_get_alignment(ctx)); + println!("{prefix}: data offset: {}", gguf_get_data_offset(ctx)); +} + +unsafe fn print_kv(prefix: &str, ctx: *const gguf_context) { + let n_kv = gguf_get_n_kv(ctx); + println!("{prefix}: n_kv: {n_kv}"); + for i in 0..n_kv { + println!("{prefix}: kv[{i}]: key = {}", cstr(gguf_get_key(ctx, i))); + } +} + +unsafe fn find_string_kv(prefix: &str, ctx: *const gguf_context) -> Result<(), String> { + let findkey = cstring("some.parameter.string")?; + let keyidx = gguf_find_key(ctx, findkey.as_ptr()); + if keyidx == -1 { + println!("{prefix}: find key: some.parameter.string not found."); + } else { + println!( + "{prefix}: find key: some.parameter.string found, kv[{keyidx}] value = {}", + cstr(gguf_get_val_str(ctx, keyidx)) + ); + } + Ok(()) +} + +unsafe fn print_tensor_info(prefix: &str, ctx: *const gguf_context, with_type: bool) { + let n_tensors = gguf_get_n_tensors(ctx); + println!("{prefix}: n_tensors: {n_tensors}"); + for i in 0..n_tensors { + let name = cstr(gguf_get_tensor_name(ctx, i)); + let size = gguf_get_tensor_size(ctx, i); + let offset = gguf_get_tensor_offset(ctx, i); + if with_type { + let type_ = gguf_get_tensor_type(ctx, i); + let type_name = cstr(ggml_type_name(type_)); + let type_size = ggml_type_size(type_); + let n_elements = size / type_size; + println!("{prefix}: tensor[{i}]: name = {name}, size = {size}, offset = {offset}, type = {type_name}, n_elts = {n_elements}"); + } else { + println!("{prefix}: tensor[{i}]: name = {name}, size = {size}, offset = {offset}"); + } + } +} + +unsafe fn print_tensor_data( + prefix: &str, + ctx: *const gguf_context, + ctx_data: *mut ggml_context, + check_data: bool, +) -> Result<(), String> { + let n_tensors = gguf_get_n_tensors(ctx); + for i in 0..n_tensors { + println!("{prefix}: reading tensor {i} data"); + let name_ptr = gguf_get_tensor_name(ctx, i); + let name = cstr(name_ptr); + let cur = ggml_get_tensor(ctx_data, name_ptr); + if cur.is_null() { + return Err(format!("{prefix}: tensor {name} was not loaded")); + } + let tensor = &*cur; + println!( + "{prefix}: tensor[{i}]: n_dims = {}, ne = ({}, {}, {}, {}), name = {}, data = {:?}", + ggml_n_dims(cur), + tensor.ne[0], + tensor.ne[1], + tensor.ne[2], + tensor.ne[3], + cstr(tensor.name.as_ptr()), + tensor.data + ); + + let n_elements = ggml_nelements(cur) as usize; + let data = std::slice::from_raw_parts(tensor.data.cast::(), n_elements); + print!("{name} data[:10] : "); + for value in data.iter().take(10) { + print!("{value:.6} "); + } + println!("\n"); + + if check_data { + for (j, value) in data.iter().enumerate() { + let expected = 100.0 + i as f32; + if *value != expected { + gguf_free(ctx as *mut gguf_context); + return Err(format!( + "{prefix}: tensor[{i}], data[{j}]: found {value}, expected {expected}" + )); + } + } + } + } + Ok(()) +} + +fn cstring(value: &str) -> Result { + CString::new(value).map_err(|err| err.to_string()) +} + +unsafe fn cstr(value: *const c_char) -> String { + if value.is_null() { + return "".to_string(); + } + CStr::from_ptr(value).to_string_lossy().into_owned() +} + +#[allow(dead_code)] +fn _assert_tensor_dims_shape() { + let _ = [0_i64; GGML_MAX_DIMS]; +} diff --git a/examples/gguf/CMakeLists.txt b/examples/gguf/CMakeLists.txt index fb04eb83f34c..ee16e8be700b 100644 --- a/examples/gguf/CMakeLists.txt +++ b/examples/gguf/CMakeLists.txt @@ -1,5 +1,23 @@ set(TARGET llama-gguf) -add_executable(${TARGET} gguf.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE ggml ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(GGUF_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../gguf-rust/Cargo.toml) +set(GGUF_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/gguf) +set(GGUF_RUST_BIN ${GGUF_RUST_TARGET_DIR}/debug/llama-gguf${CMAKE_EXECUTABLE_SUFFIX}) +set(GGUF_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${GGUF_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env GGML_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${GGUF_RUST_MANIFEST} --target-dir ${GGUF_RUST_TARGET_DIR} --bin llama-gguf + COMMAND ${CMAKE_COMMAND} -E copy ${GGUF_RUST_BIN} ${GGUF_RUST_OUTPUT} + DEPENDS + ggml + ${GGUF_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../gguf-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../gguf-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../gguf-rust/src/main.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${GGUF_RUST_OUTPUT}) +install(PROGRAMS ${GGUF_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/gguf/gguf.cpp b/examples/gguf/gguf.cpp deleted file mode 100644 index 79ad38711e3e..000000000000 --- a/examples/gguf/gguf.cpp +++ /dev/null @@ -1,273 +0,0 @@ -#include "ggml.h" -#include "gguf.h" - -#include -#include -#include -#include -#include - -#undef MIN -#undef MAX -#define MIN(a, b) ((a) < (b) ? (a) : (b)) -#define MAX(a, b) ((a) > (b) ? (a) : (b)) - -template -static std::string to_string(const T & val) { - std::stringstream ss; - ss << val; - return ss.str(); -} - -static bool gguf_ex_write(const std::string & fname) { - struct gguf_context * ctx = gguf_init_empty(); - - gguf_set_val_u8 (ctx, "some.parameter.uint8", 0x12); - gguf_set_val_i8 (ctx, "some.parameter.int8", -0x13); - gguf_set_val_u16 (ctx, "some.parameter.uint16", 0x1234); - gguf_set_val_i16 (ctx, "some.parameter.int16", -0x1235); - gguf_set_val_u32 (ctx, "some.parameter.uint32", 0x12345678); - gguf_set_val_i32 (ctx, "some.parameter.int32", -0x12345679); - gguf_set_val_f32 (ctx, "some.parameter.float32", 0.123456789f); - gguf_set_val_u64 (ctx, "some.parameter.uint64", 0x123456789abcdef0ull); - gguf_set_val_i64 (ctx, "some.parameter.int64", -0x123456789abcdef1ll); - gguf_set_val_f64 (ctx, "some.parameter.float64", 0.1234567890123456789); - gguf_set_val_bool(ctx, "some.parameter.bool", true); - gguf_set_val_str (ctx, "some.parameter.string", "hello world"); - - gguf_set_arr_data(ctx, "some.parameter.arr.i16", GGUF_TYPE_INT16, std::vector{ 1, 2, 3, 4, }.data(), 4); - gguf_set_arr_data(ctx, "some.parameter.arr.f32", GGUF_TYPE_FLOAT32, std::vector{ 3.145f, 2.718f, 1.414f, }.data(), 3); - gguf_set_arr_str (ctx, "some.parameter.arr.str", std::vector{ "hello", "world", "!" }.data(), 3); - - struct ggml_init_params params = { - /*.mem_size =*/ 128ull*1024ull*1024ull, - /*.mem_buffer =*/ NULL, - /*.no_alloc =*/ false, - }; - - struct ggml_context * ctx_data = ggml_init(params); - - const int n_tensors = 10; - - // tensor infos - for (int i = 0; i < n_tensors; ++i) { - const std::string name = "tensor_" + to_string(i); - - int64_t ne[GGML_MAX_DIMS] = { 1 }; - int32_t n_dims = rand() % GGML_MAX_DIMS + 1; - - for (int j = 0; j < n_dims; ++j) { - ne[j] = rand() % 10 + 1; - } - - struct ggml_tensor * cur = ggml_new_tensor(ctx_data, GGML_TYPE_F32, n_dims, ne); - ggml_set_name(cur, name.c_str()); - - { - float * data = (float *) cur->data; - for (int j = 0; j < ggml_nelements(cur); ++j) { - data[j] = 100 + i; - } - } - - gguf_add_tensor(ctx, cur); - } - - gguf_write_to_file(ctx, fname.c_str(), false); - - printf("%s: wrote file '%s;\n", __func__, fname.c_str()); - - ggml_free(ctx_data); - gguf_free(ctx); - - return true; -} - -// just read tensor info -static bool gguf_ex_read_0(const std::string & fname) { - struct gguf_init_params params = { - /*.no_alloc = */ false, - /*.ctx = */ NULL, - }; - - struct gguf_context * ctx = gguf_init_from_file(fname.c_str(), params); - - if (!ctx) { - fprintf(stderr, "%s: failed to load '%s'\n", __func__, fname.c_str()); - return false; - } - - printf("%s: version: %d\n", __func__, gguf_get_version(ctx)); - printf("%s: alignment: %zu\n", __func__, gguf_get_alignment(ctx)); - printf("%s: data offset: %zu\n", __func__, gguf_get_data_offset(ctx)); - - // kv - { - const int n_kv = gguf_get_n_kv(ctx); - - printf("%s: n_kv: %d\n", __func__, n_kv); - - for (int i = 0; i < n_kv; ++i) { - const char * key = gguf_get_key(ctx, i); - - printf("%s: kv[%d]: key = %s\n", __func__, i, key); - } - } - - // find kv string - { - const char * findkey = "some.parameter.string"; - - const int keyidx = gguf_find_key(ctx, findkey); - if (keyidx == -1) { - printf("%s: find key: %s not found.\n", __func__, findkey); - } else { - const char * key_value = gguf_get_val_str(ctx, keyidx); - printf("%s: find key: %s found, kv[%d] value = %s\n", __func__, findkey, keyidx, key_value); - } - } - - // tensor info - { - const int n_tensors = gguf_get_n_tensors(ctx); - - printf("%s: n_tensors: %d\n", __func__, n_tensors); - - for (int i = 0; i < n_tensors; ++i) { - const char * name = gguf_get_tensor_name (ctx, i); - const size_t size = gguf_get_tensor_size (ctx, i); - const size_t offset = gguf_get_tensor_offset(ctx, i); - - printf("%s: tensor[%d]: name = %s, size = %zu, offset = %zu\n", __func__, i, name, size, offset); - } - } - - gguf_free(ctx); - - return true; -} - -// read and create ggml_context containing the tensors and their data -static bool gguf_ex_read_1(const std::string & fname, bool check_data) { - struct ggml_context * ctx_data = NULL; - - struct gguf_init_params params = { - /*.no_alloc = */ false, - /*.ctx = */ &ctx_data, - }; - - struct gguf_context * ctx = gguf_init_from_file(fname.c_str(), params); - - printf("%s: version: %d\n", __func__, gguf_get_version(ctx)); - printf("%s: alignment: %zu\n", __func__, gguf_get_alignment(ctx)); - printf("%s: data offset: %zu\n", __func__, gguf_get_data_offset(ctx)); - - // kv - { - const int n_kv = gguf_get_n_kv(ctx); - - printf("%s: n_kv: %d\n", __func__, n_kv); - - for (int i = 0; i < n_kv; ++i) { - const char * key = gguf_get_key(ctx, i); - - printf("%s: kv[%d]: key = %s\n", __func__, i, key); - } - } - - // tensor info - { - const int n_tensors = gguf_get_n_tensors(ctx); - - printf("%s: n_tensors: %d\n", __func__, n_tensors); - - for (int i = 0; i < n_tensors; ++i) { - const char * name = gguf_get_tensor_name (ctx, i); - const size_t size = gguf_get_tensor_size (ctx, i); - const size_t offset = gguf_get_tensor_offset(ctx, i); - const auto type = gguf_get_tensor_type (ctx, i); - - const char * type_name = ggml_type_name(type); - const size_t type_size = ggml_type_size(type); - const size_t n_elements = size / type_size; - - printf("%s: tensor[%d]: name = %s, size = %zu, offset = %zu, type = %s, n_elts = %zu\n", __func__, i, name, size, offset, type_name, n_elements); - } - } - - // data - { - const int n_tensors = gguf_get_n_tensors(ctx); - - for (int i = 0; i < n_tensors; ++i) { - printf("%s: reading tensor %d data\n", __func__, i); - - const char * name = gguf_get_tensor_name(ctx, i); - - struct ggml_tensor * cur = ggml_get_tensor(ctx_data, name); - - printf("%s: tensor[%d]: n_dims = %d, ne = (%d, %d, %d, %d), name = %s, data = %p\n", - __func__, i, ggml_n_dims(cur), int(cur->ne[0]), int(cur->ne[1]), int(cur->ne[2]), int(cur->ne[3]), cur->name, cur->data); - - // print first 10 elements - const float * data = (const float *) cur->data; - - printf("%s data[:10] : ", name); - for (int j = 0; j < MIN(10, ggml_nelements(cur)); ++j) { - printf("%f ", data[j]); - } - printf("\n\n"); - - // check data - if (check_data) { - const float * data = (const float *) cur->data; - for (int j = 0; j < ggml_nelements(cur); ++j) { - if (data[j] != 100 + i) { - fprintf(stderr, "%s: tensor[%d], data[%d]: found %f, expected %f\n", __func__, i, j, data[j], float(100 + i)); - gguf_free(ctx); - return false; - } - } - } - } - } - - printf("%s: ctx_data size: %zu\n", __func__, ggml_get_mem_size(ctx_data)); - - ggml_free(ctx_data); - gguf_free(ctx); - - return true; -} - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - if (argc < 3) { - printf("usage: %s data.gguf r|w [n]\n", argv[0]); - printf("r: read data.gguf file\n"); - printf("w: write data.gguf file\n"); - printf("n: no check of tensor data\n"); - return -1; - } - bool check_data = true; - if (argc == 4) { - check_data = false; - } - - srand(123456); - - const std::string fname(argv[1]); - const std::string mode (argv[2]); - - GGML_ASSERT((mode == "r" || mode == "w") && "mode must be r or w"); - - if (mode == "w") { - GGML_ASSERT(gguf_ex_write(fname) && "failed to write gguf file"); - } else if (mode == "r") { - GGML_ASSERT(gguf_ex_read_0(fname) && "failed to read gguf file"); - GGML_ASSERT(gguf_ex_read_1(fname, check_data) && "failed to read gguf file"); - } - - return 0; -} diff --git a/examples/idle-rust/Cargo.toml b/examples/idle-rust/Cargo.toml new file mode 100644 index 000000000000..725521a45a81 --- /dev/null +++ b/examples/idle-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-idle-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-idle" +path = "src/main.rs" diff --git a/examples/idle-rust/src/lib.rs b/examples/idle-rust/src/lib.rs new file mode 100644 index 000000000000..33e55f1b1237 --- /dev/null +++ b/examples/idle-rust/src/lib.rs @@ -0,0 +1,99 @@ +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Args { + pub model_path: String, + pub n_gpu_layers: i32, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + MissingModel, + MissingValue(&'static str), + InvalidInteger(&'static str, String), + UnexpectedArgument(String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + ParseError::UnexpectedArgument(value) => write!(f, "unexpected argument: {value}"), + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut args = args.into_iter().map(Into::into); + let mut model_path = String::new(); + let mut n_gpu_layers = 99; + + while let Some(arg) = args.next() { + match arg.as_str() { + "-m" => { + model_path = args.next().ok_or(ParseError::MissingValue("-m"))?; + } + "-ngl" => { + let value = args.next().ok_or(ParseError::MissingValue("-ngl"))?; + n_gpu_layers = value + .parse() + .map_err(|_| ParseError::InvalidInteger("-ngl", value))?; + } + _ => return Err(ParseError::UnexpectedArgument(arg)), + } + } + + if model_path.is_empty() { + return Err(ParseError::MissingModel); + } + + Ok(Args { + model_path, + n_gpu_layers, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_defaults() { + let args = parse_args(["-m", "model.gguf"]).unwrap(); + + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.n_gpu_layers, 99); + } + + #[test] + fn parses_gpu_layers() { + let args = parse_args(["-m", "model.gguf", "-ngl", "0"]).unwrap(); + + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.n_gpu_layers, 0); + } + + #[test] + fn rejects_missing_model() { + assert_eq!( + parse_args(["-ngl", "0"]).unwrap_err(), + ParseError::MissingModel + ); + } + + #[test] + fn rejects_unexpected_argument() { + assert_eq!( + parse_args(["-m", "model.gguf", "-c", "512"]).unwrap_err(), + ParseError::UnexpectedArgument("-c".to_string()) + ); + } +} diff --git a/examples/idle-rust/src/main.rs b/examples/idle-rust/src/main.rs new file mode 100644 index 000000000000..626e0c08c976 --- /dev/null +++ b/examples/idle-rust/src/main.rs @@ -0,0 +1,148 @@ +use llama_idle_rust::{parse_args, Args}; +use llama_simple_rust::ffi; +use std::env; +use std::ffi::CString; +use std::thread; +use std::time::Duration; + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m model.gguf [-ngl n_gpu_layers]"); + eprintln!(); +} + +fn run(args: Args) -> Result<(), String> { + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + + unsafe { + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + let mut prompt_tokens = [unsafe { ffi::llama_vocab_bos(vocab) }]; + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = 512; + ctx_params.n_batch = 512; + ctx_params.no_perf = false; + + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create the llama_context".to_string()); + } + + let batch = unsafe { ffi::llama_batch_get_one(prompt_tokens.as_mut_ptr(), 1) }; + let memory = unsafe { ffi::llama_get_memory(ctx.0) }; + + unsafe { + ffi::llama_decode(ctx.0, batch); + ffi::llama_memory_clear(memory, true); + ffi::llama_synchronize(ctx.0); + } + + const N_ITERS: i32 = 3; + let mut t_pause_ms = 0_i64; + while t_pause_ms <= 4000 { + let mut t_sum_us = 0.0_f64; + let mut t_sum2_us = 0.0_f64; + + for _ in 0..N_ITERS { + thread::sleep(Duration::from_millis(t_pause_ms as u64)); + + let t_start_us = unsafe { ffi::llama_time_us() }; + unsafe { + ffi::llama_decode(ctx.0, batch); + ffi::llama_synchronize(ctx.0); + } + let t_end_us = unsafe { ffi::llama_time_us() }; + + let t_cur_us = (t_end_us - t_start_us) as f64; + println!(" - decode time: {:8.2} ms", t_cur_us / 1000.0); + + t_sum_us += t_cur_us; + t_sum2_us += t_cur_us * t_cur_us; + + unsafe { + ffi::llama_memory_clear(memory, true); + ffi::llama_synchronize(ctx.0); + } + } + + let t_avg_us = t_sum_us / f64::from(N_ITERS); + let t_dev_us = ((t_sum2_us / f64::from(N_ITERS - 1)) + - (t_avg_us * t_avg_us * f64::from(N_ITERS)) / f64::from(N_ITERS - 1)) + .sqrt(); + + println!( + "iters: {:4}, pause: {:5} ms, avg decode time: {:8.2} +/- {:4.2} ms", + N_ITERS, + t_pause_ms, + t_avg_us / 1000.0, + t_dev_us / 1000.0 + ); + + t_pause_ms += 800; + } + + unsafe { + ffi::llama_backend_free(); + } + + Ok(()) +} + +fn main() { + let mut argv = env::args(); + let program = argv.next().unwrap_or_else(|| "llama-idle".to_string()); + + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} diff --git a/examples/idle/CMakeLists.txt b/examples/idle/CMakeLists.txt index c0fedbbff5bd..8bfee2a114de 100644 --- a/examples/idle/CMakeLists.txt +++ b/examples/idle/CMakeLists.txt @@ -1,5 +1,25 @@ set(TARGET llama-idle) -add_executable(${TARGET} idle.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama llama-common ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_11) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(IDLE_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../idle-rust/Cargo.toml) +set(IDLE_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/idle) +set(IDLE_RUST_BIN ${IDLE_RUST_TARGET_DIR}/debug/llama-idle${CMAKE_EXECUTABLE_SUFFIX}) +set(IDLE_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${IDLE_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${IDLE_RUST_MANIFEST} --target-dir ${IDLE_RUST_TARGET_DIR} --bin llama-idle + COMMAND ${CMAKE_COMMAND} -E copy ${IDLE_RUST_BIN} ${IDLE_RUST_OUTPUT} + DEPENDS + llama + ${IDLE_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../idle-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../idle-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${IDLE_RUST_OUTPUT}) +install(PROGRAMS ${IDLE_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/idle/idle.cpp b/examples/idle/idle.cpp deleted file mode 100644 index 409fd25c18f8..000000000000 --- a/examples/idle/idle.cpp +++ /dev/null @@ -1,110 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "log.h" -#include "llama.h" - -#include -#include -#include -#include -#include -#include - -static void print_usage(int /*argc*/, char ** argv) { - printf("\nexample usage:\n"); - printf("\n %s -m model.gguf [-ngl n_gpu_layers]\n", argv[0]); - printf("\n"); -} - -int main(int argc, char ** argv) { - common_params params; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON, print_usage)) { - return 1; - } - - // init LLM - - llama_backend_init(); - llama_numa_init(params.numa); - - // initialize the model - - llama_model_params model_params = common_model_params_to_llama(params); - - llama_model * model = llama_model_load_from_file(params.model.path.c_str(), model_params); - - if (model == NULL) { - LOG_ERR("%s: error: unable to load model\n" , __func__); - return 1; - } - - const llama_vocab * vocab = llama_model_get_vocab(model); - - // we need just a dummy token to evaluate - std::vector prompt_tokens(1, llama_vocab_bos(vocab)); - - llama_context_params ctx_params = llama_context_default_params(); - ctx_params.n_ctx = 512; - ctx_params.n_batch = 512; - ctx_params.no_perf = false; - - llama_context * ctx = llama_init_from_model(model, ctx_params); - if (ctx == NULL) { - fprintf(stderr , "%s: error: failed to create the llama_context\n" , __func__); - return 1; - } - - llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size()); - - const int n_iters = 3; - - // warm-up - llama_decode(ctx, batch); - llama_memory_clear(llama_get_memory(ctx), true); - llama_synchronize(ctx); - - for (int64_t t_pause_ms = 0; t_pause_ms <= 4000; t_pause_ms += 800) { - double t_sum_us = 0.0; - double t_sum2_us = 0.0; - - for (int i = 0; i < n_iters; i++) { - // this pause is important - it simulates "idle GPU" - std::this_thread::sleep_for(std::chrono::milliseconds(t_pause_ms)); - - const int64_t t_start_us = llama_time_us(); - - // this should take constant time - llama_decode(ctx, batch); - llama_synchronize(ctx); - - const int64_t t_end_us = llama_time_us(); - - const double t_cur_us = t_end_us - t_start_us; - -#if 1 - // print individual decode times - printf(" - decode time: %8.2f ms\n", t_cur_us / 1000); -#endif - - t_sum_us += t_cur_us; - t_sum2_us += t_cur_us * t_cur_us; - - llama_memory_clear(llama_get_memory(ctx), true); - llama_synchronize(ctx); // just in case - } - - const double t_avg_us = t_sum_us / n_iters; - const double t_dev_us = sqrt((t_sum2_us / (n_iters - 1)) - (t_avg_us * t_avg_us * n_iters) / (n_iters - 1)); - - printf("iters: %4d, pause: %5d ms, avg decode time: %8.2f +/- %4.2f ms\n", n_iters, (int) t_pause_ms, t_avg_us / 1000, t_dev_us / 1000); - fflush(stdout); - } - - llama_free(ctx); - llama_model_free(model); - - return 0; -} diff --git a/examples/llama.android/lib/src/main/cpp/CMakeLists.txt b/examples/llama.android/lib/src/main/cpp/CMakeLists.txt index 20c9e3b2c1fd..b812b7a25806 100644 --- a/examples/llama.android/lib/src/main/cpp/CMakeLists.txt +++ b/examples/llama.android/lib/src/main/cpp/CMakeLists.txt @@ -34,7 +34,7 @@ set(LLAMA_SRC ${CMAKE_CURRENT_LIST_DIR}/../../../../../../) add_subdirectory(${LLAMA_SRC} build-llama) add_library(${CMAKE_PROJECT_NAME} SHARED - ai_chat.cpp) + ai_chat.cpp.inc) target_compile_definitions(${CMAKE_PROJECT_NAME} PRIVATE GGML_SYSTEM_ARCH=${GGML_SYSTEM_ARCH} diff --git a/examples/llama.android/lib/src/main/cpp/ai_chat.cpp b/examples/llama.android/lib/src/main/cpp/ai_chat.cpp.inc similarity index 90% rename from examples/llama.android/lib/src/main/cpp/ai_chat.cpp rename to examples/llama.android/lib/src/main/cpp/ai_chat.cpp.inc index 03ab96cfd834..6d929915b761 100644 --- a/examples/llama.android/lib/src/main/cpp/ai_chat.cpp +++ b/examples/llama.android/lib/src/main/cpp/ai_chat.cpp.inc @@ -4,12 +4,63 @@ #include #include #include -#include +#include + +#include "common.h.inc" +#include "common.h.inc" +#include "llama.h.inc" + +#ifndef LOG_TAG +#define LOG_TAG "ai-chat" +#endif + +#ifndef LOG_MIN_LEVEL +#if defined(NDEBUG) +#define LOG_MIN_LEVEL ANDROID_LOG_INFO +#else +#define LOG_MIN_LEVEL ANDROID_LOG_VERBOSE +#endif +#endif + +static inline int ai_should_log(int prio) { + return __android_log_is_loggable(prio, LOG_TAG, LOG_MIN_LEVEL); +} + +#if LOG_MIN_LEVEL <= ANDROID_LOG_VERBOSE +#define LOGv(...) do { if (ai_should_log(ANDROID_LOG_VERBOSE)) __android_log_print(ANDROID_LOG_VERBOSE, LOG_TAG, __VA_ARGS__); } while (0) +#else +#define LOGv(...) ((void)0) +#endif + +#if LOG_MIN_LEVEL <= ANDROID_LOG_DEBUG +#define LOGd(...) do { if (ai_should_log(ANDROID_LOG_DEBUG)) __android_log_print(ANDROID_LOG_DEBUG, LOG_TAG, __VA_ARGS__); } while (0) +#else +#define LOGd(...) ((void)0) +#endif + +#define LOGi(...) do { if (ai_should_log(ANDROID_LOG_INFO )) __android_log_print(ANDROID_LOG_INFO , LOG_TAG, __VA_ARGS__); } while (0) +#define LOGw(...) do { if (ai_should_log(ANDROID_LOG_WARN )) __android_log_print(ANDROID_LOG_WARN , LOG_TAG, __VA_ARGS__); } while (0) +#define LOGe(...) do { if (ai_should_log(ANDROID_LOG_ERROR)) __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__); } while (0) + +static inline int android_log_prio_from_ggml(enum ggml_log_level level) { + switch (level) { + case GGML_LOG_LEVEL_ERROR: return ANDROID_LOG_ERROR; + case GGML_LOG_LEVEL_WARN: return ANDROID_LOG_WARN; + case GGML_LOG_LEVEL_INFO: return ANDROID_LOG_INFO; + case GGML_LOG_LEVEL_DEBUG: return ANDROID_LOG_DEBUG; + default: return ANDROID_LOG_DEFAULT; + } +} -#include "logging.h" -#include "chat.h" -#include "common.h" -#include "llama.h" +static inline void aichat_android_log_callback(enum ggml_log_level level, + const char * text, + void * /*user*/) { + const int prio = android_log_prio_from_ggml(level); + if (!ai_should_log(prio)) { + return; + } + __android_log_write(prio, LOG_TAG, text); +} template static std::string join(const std::vector &values, const std::string &delim) { diff --git a/examples/llama.android/lib/src/main/cpp/logging.h b/examples/llama.android/lib/src/main/cpp/logging.h deleted file mode 100644 index 2e768d2beb1a..000000000000 --- a/examples/llama.android/lib/src/main/cpp/logging.h +++ /dev/null @@ -1,61 +0,0 @@ -// -// Created by Han Yin on 10/31/25. -// - -#ifndef AICHAT_LOGGING_H -#define AICHAT_LOGGING_H - -#endif //AICHAT_LOGGING_H - -#pragma once -#include - -#ifndef LOG_TAG -#define LOG_TAG "ai-chat" -#endif - -#ifndef LOG_MIN_LEVEL -#if defined(NDEBUG) -#define LOG_MIN_LEVEL ANDROID_LOG_INFO -#else -#define LOG_MIN_LEVEL ANDROID_LOG_VERBOSE -#endif -#endif - -static inline int ai_should_log(int prio) { - return __android_log_is_loggable(prio, LOG_TAG, LOG_MIN_LEVEL); -} - -#if LOG_MIN_LEVEL <= ANDROID_LOG_VERBOSE -#define LOGv(...) do { if (ai_should_log(ANDROID_LOG_VERBOSE)) __android_log_print(ANDROID_LOG_VERBOSE, LOG_TAG, __VA_ARGS__); } while (0) -#else -#define LOGv(...) ((void)0) -#endif - -#if LOG_MIN_LEVEL <= ANDROID_LOG_DEBUG -#define LOGd(...) do { if (ai_should_log(ANDROID_LOG_DEBUG)) __android_log_print(ANDROID_LOG_DEBUG, LOG_TAG, __VA_ARGS__); } while (0) -#else -#define LOGd(...) ((void)0) -#endif - -#define LOGi(...) do { if (ai_should_log(ANDROID_LOG_INFO )) __android_log_print(ANDROID_LOG_INFO , LOG_TAG, __VA_ARGS__); } while (0) -#define LOGw(...) do { if (ai_should_log(ANDROID_LOG_WARN )) __android_log_print(ANDROID_LOG_WARN , LOG_TAG, __VA_ARGS__); } while (0) -#define LOGe(...) do { if (ai_should_log(ANDROID_LOG_ERROR)) __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__); } while (0) - -static inline int android_log_prio_from_ggml(enum ggml_log_level level) { - switch (level) { - case GGML_LOG_LEVEL_ERROR: return ANDROID_LOG_ERROR; - case GGML_LOG_LEVEL_WARN: return ANDROID_LOG_WARN; - case GGML_LOG_LEVEL_INFO: return ANDROID_LOG_INFO; - case GGML_LOG_LEVEL_DEBUG: return ANDROID_LOG_DEBUG; - default: return ANDROID_LOG_DEFAULT; - } -} - -static inline void aichat_android_log_callback(enum ggml_log_level level, - const char* text, - void* /*user*/) { - const int prio = android_log_prio_from_ggml(level); - if (!ai_should_log(prio)) return; - __android_log_write(prio, LOG_TAG, text); -} diff --git a/examples/llama.android/lib/src/main/java/com/arm/aichat/internal/InferenceEngineImpl.kt b/examples/llama.android/lib/src/main/java/com/arm/aichat/internal/InferenceEngineImpl.kt index a293796f9017..77b7acbdd4e7 100644 --- a/examples/llama.android/lib/src/main/java/com/arm/aichat/internal/InferenceEngineImpl.kt +++ b/examples/llama.android/lib/src/main/java/com/arm/aichat/internal/InferenceEngineImpl.kt @@ -41,7 +41,7 @@ import java.io.IOException * * State transitions are managed automatically and validated at each operation. * - * @see ai_chat.cpp for the native implementation details + * @see ai_chat.cpp.inc for the native implementation details */ internal class InferenceEngineImpl private constructor( private val nativeLibDir: String @@ -77,7 +77,7 @@ internal class InferenceEngineImpl private constructor( /** * JNI methods - * @see ai_chat.cpp + * @see ai_chat.cpp.inc */ @FastNative private external fun init(nativeLibDir: String) diff --git a/examples/lookahead-rust/Cargo.toml b/examples/lookahead-rust/Cargo.toml new file mode 100644 index 000000000000..4eae5d1cd3dc --- /dev/null +++ b/examples/lookahead-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-lookahead-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-lookahead" +path = "src/main.rs" diff --git a/examples/lookahead-rust/build.rs b/examples/lookahead-rust/build.rs new file mode 100644 index 000000000000..9cfcbcf4be6a --- /dev/null +++ b/examples/lookahead-rust/build.rs @@ -0,0 +1,45 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); + + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-lookahead"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/lookahead-rust/src/main.rs b/examples/lookahead-rust/src/main.rs new file mode 100644 index 000000000000..de5eb09f60fd --- /dev/null +++ b/examples/lookahead-rust/src/main.rs @@ -0,0 +1,761 @@ +use llama_simple_rust::ffi; +use std::env; +use std::ffi::CString; +use std::io::{self, Write}; +use std::ptr; + +const W: usize = 15; +const N: usize = 5; +const G: usize = 15; + +#[derive(Debug, Clone, PartialEq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub n_predict: i32, + pub n_ctx: u32, + pub n_batch: u32, + pub n_gpu_layers: i32, + pub top_k: i32, + pub top_p: f32, + pub temp: f32, + pub seed: u32, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum ParseError { + MissingModel, + MissingValue(String), + InvalidInteger(String, String), + InvalidFloat(String, String), + InvalidValue(String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m/--model model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + ParseError::InvalidFloat(flag, value) => write!(f, "invalid float for {flag}: {value}"), + ParseError::InvalidValue(message) => write!(f, "{message}"), + } + } +} + +impl std::error::Error for ParseError {} + +impl Default for Args { + fn default() -> Self { + Self { + model_path: String::new(), + prompt: "Hello my name is".to_string(), + n_predict: 32, + n_ctx: 512, + n_batch: 512, + n_gpu_layers: 99, + top_k: 40, + top_p: 0.95, + temp: 0.8, + seed: ffi::LLAMA_DEFAULT_SEED, + } + } +} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut parsed = Args::default(); + let mut iter = args.into_iter().map(Into::into); + + while let Some(arg) = iter.next() { + match arg.as_str() { + "-m" | "--model" => parsed.model_path = take(&mut iter, &arg)?, + "-p" | "--prompt" => parsed.prompt = take(&mut iter, &arg)?, + "-n" | "--n-predict" => parsed.n_predict = parse_i32(&mut iter, &arg)?, + "-c" | "--ctx-size" => parsed.n_ctx = parse_u32(&mut iter, &arg)?, + "-b" | "--batch-size" => parsed.n_batch = parse_u32(&mut iter, &arg)?, + "-ngl" | "--gpu-layers" => parsed.n_gpu_layers = parse_i32(&mut iter, &arg)?, + "--top-k" => parsed.top_k = parse_i32(&mut iter, &arg)?, + "--top-p" => parsed.top_p = parse_f32(&mut iter, &arg)?, + "--temp" => parsed.temp = parse_f32(&mut iter, &arg)?, + "-s" | "--seed" => parsed.seed = parse_u32(&mut iter, &arg)?, + "-h" | "--help" => return Err(ParseError::MissingModel), + other => { + let mut prompt = vec![other.to_string()]; + prompt.extend(iter); + parsed.prompt = prompt.join(" "); + break; + } + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + if parsed.n_ctx == 0 || parsed.n_batch == 0 { + return Err(ParseError::InvalidValue( + "context and batch sizes must be positive".to_string(), + )); + } + Ok(parsed) +} + +fn take(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + iter.next() + .ok_or_else(|| ParseError::MissingValue(flag.to_string())) +} + +fn parse_i32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_u32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_f32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidFloat(flag.to_string(), value)) +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct NgramData { + pub active: bool, + pub seq_id: ffi::llama_seq_id, + pub i_batch: Vec, + pub tokens: Vec, +} + +impl Default for NgramData { + fn default() -> Self { + Self { + active: false, + seq_id: -1, + i_batch: Vec::new(), + tokens: Vec::new(), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct NgramContainer { + pub n_total: usize, + pub n_vocab: usize, + pub n: usize, + pub g: usize, + pub cnt: Vec, + pub head: Vec, + pub tokens: Vec, +} + +impl NgramContainer { + pub fn new(n_vocab: usize, n: usize, g: usize) -> Self { + Self { + n_total: 0, + n_vocab, + n, + g, + cnt: vec![0; n_vocab], + head: vec![0; n_vocab], + tokens: vec![0; n_vocab * g * (n - 1)], + } + } + + pub fn observed_for(&self, token: ffi::llama_token) -> usize { + if token < 0 || token as usize >= self.n_vocab { + 0 + } else { + self.cnt[token as usize] + } + } + + pub fn get(&self, token: ffi::llama_token, g: usize, j: usize) -> ffi::llama_token { + let token = token as usize; + self.tokens[token * (self.n - 1) * self.g + g * (self.n - 1) + j] + } + + pub fn insert_unique(&mut self, first: ffi::llama_token, ngram: &[ffi::llama_token]) -> bool { + if first < 0 || first as usize >= self.n_vocab || ngram.len() != self.n - 1 { + return false; + } + let first = first as usize; + for k in 0..self.cnt[first] { + let idx = first * (self.n - 1) * self.g + k * (self.n - 1); + if self.tokens[idx..idx + self.n - 1] == *ngram { + return false; + } + } + + let head = self.head[first]; + let idx = first * (self.n - 1) * self.g + head * (self.n - 1); + self.tokens[idx..idx + self.n - 1].copy_from_slice(ngram); + self.cnt[first] = self.g.min(self.cnt[first] + 1); + self.head[first] = (head + 1) % self.g; + self.n_total += 1; + true + } +} + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Batch(ffi::llama_batch); + +impl Drop for Batch { + fn drop(&mut self) { + unsafe { + ffi::llama_batch_free(self.0); + } + } +} + +struct Sampler(*mut ffi::llama_sampler); + +impl Drop for Sampler { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_sampler_free(self.0); + } + } + } +} + +struct Backend; + +impl Backend { + fn init() -> Self { + unsafe { + ffi::ggml_backend_load_all(); + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + Self + } +} + +impl Drop for Backend { + fn drop(&mut self) { + unsafe { + ffi::llama_backend_free(); + } + } +} + +fn tokenize( + vocab: *const ffi::llama_vocab, + text: &str, + add_special: bool, + parse_special: bool, +) -> Result, String> { + let text_c = + CString::new(text).map_err(|_| "prompt contains an interior NUL byte".to_string())?; + let n_tokens = unsafe { + -ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + ptr::null_mut(), + 0, + add_special, + parse_special, + ) + }; + if n_tokens <= 0 { + return Err("failed to size tokenization".to_string()); + } + + let mut tokens = vec![0_i32; n_tokens as usize]; + let n = unsafe { + ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + add_special, + parse_special, + ) + }; + if n < 0 { + return Err("failed to tokenize prompt".to_string()); + } + tokens.truncate(n as usize); + Ok(tokens) +} + +fn token_to_piece( + vocab: *const ffi::llama_vocab, + token: ffi::llama_token, +) -> Result { + let mut buf = vec![0_i8; 128]; + let n = unsafe { + ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), buf.len() as i32, 0, true) + }; + if n < 0 { + return Err("failed to convert token to piece".to_string()); + } + let bytes = buf[..n as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +fn sampler_init(args: &Args) -> Sampler { + unsafe { + let chain = ffi::llama_sampler_chain_init(ffi::llama_sampler_chain_default_params()); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_top_k(args.top_k)); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_top_p(args.top_p, 1)); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_temp(args.temp)); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_dist(args.seed)); + Sampler(chain) + } +} + +fn sample_accept( + sampler: *mut ffi::llama_sampler, + ctx: *mut ffi::llama_context, + idx: i32, +) -> ffi::llama_token { + unsafe { + let token = ffi::llama_sampler_sample(sampler, ctx, idx); + ffi::llama_sampler_accept(sampler, token); + token + } +} + +fn batch_clear(batch: &mut ffi::llama_batch) { + batch.n_tokens = 0; +} + +unsafe fn batch_add( + batch: &mut ffi::llama_batch, + token: ffi::llama_token, + pos: ffi::llama_pos, + seq_ids: &[ffi::llama_seq_id], + logits: bool, +) { + let i = batch.n_tokens as isize; + *batch.token.offset(i) = token; + *batch.pos.offset(i) = pos; + *batch.n_seq_id.offset(i) = seq_ids.len() as i32; + let seq_slot = *batch.seq_id.offset(i); + for (j, seq_id) in seq_ids.iter().enumerate() { + *seq_slot.add(j) = *seq_id; + } + *batch.logits.offset(i) = if logits { 1 } else { 0 }; + batch.n_tokens += 1; +} + +fn decode_tokens( + ctx: *mut ffi::llama_context, + tokens: &mut [ffi::llama_token], +) -> Result<(), String> { + if tokens.len() > 1 { + let batch = + unsafe { ffi::llama_batch_get_one(tokens.as_mut_ptr(), tokens.len() as i32 - 1) }; + if unsafe { ffi::llama_decode(ctx, batch) } != 0 { + return Err("failed to decode prompt prefix".to_string()); + } + } + let last = tokens + .last_mut() + .ok_or_else(|| "prompt produced no tokens".to_string())?; + let batch = unsafe { ffi::llama_batch_get_one(last as *mut ffi::llama_token, 1) }; + if unsafe { ffi::llama_decode(ctx, batch) } != 0 { + return Err("failed to decode final prompt token".to_string()); + } + Ok(()) +} + +fn run(args: Args) -> Result<(), String> { + let _backend = Backend::init(); + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = args.n_ctx; + ctx_params.n_batch = args.n_batch; + ctx_params.n_ubatch = args.n_batch; + ctx_params.n_seq_max = (W + G + 1) as u32; + ctx_params.kv_unified = true; + + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create llama_context".to_string()); + } + + let mem = unsafe { ffi::llama_get_memory(ctx.0) }; + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + + let mut inp = tokenize(vocab, &args.prompt, true, true)?; + let mut all = inp.clone(); + let max_tokens_list_size = unsafe { ffi::llama_n_ctx(ctx.0) as usize }.saturating_sub(4); + if inp.len() > max_tokens_list_size { + return Err(format!( + "prompt too long ({} tokens, max {max_tokens_list_size})", + inp.len() + )); + } + if inp.is_empty() { + return Err("prompt produced no tokens".to_string()); + } + + println!(); + for token in &inp { + print!("{}", token_to_piece(vocab, *token)?); + } + io::stdout().flush().map_err(|err| err.to_string())?; + + let n_input = inp.len(); + let t_enc_start = unsafe { ffi::ggml_time_us() }; + decode_tokens(ctx.0, &mut inp)?; + for s in 1..(W + G + 1) { + unsafe { + ffi::llama_memory_seq_cp(mem, 0, s as i32, -1, -1); + } + } + let t_enc_end = unsafe { ffi::ggml_time_us() }; + + let mut n_predict = 0; + let mut n_accept = 0; + let mut n_past = inp.len() as i32; + let mut has_eos = false; + let sampler = sampler_init(&args); + let mut batch = Batch(unsafe { + ffi::llama_batch_init(ffi::llama_n_ctx(ctx.0) as i32, 0, (W + G + 1) as i32) + }); + + let mut ngrams_cur = vec![NgramData::default(); G]; + let mut tokens_j_prev = vec![0; W]; + let n_vocab = unsafe { ffi::llama_vocab_n_tokens(vocab) as usize }; + let mut tokens_j = vec![vec![0; W]; N - 1]; + for row in tokens_j.iter_mut().take(N - 1) { + for (i, value) in row.iter_mut().enumerate().take(W) { + *value = (100 + i) as i32; + if n_vocab > 0 && *value as usize >= n_vocab { + *value = (n_vocab - 1) as i32; + } + } + } + + let seq_id_all = (0..(W + G + 1)).map(|i| i as i32).collect::>(); + let mut ngrams_observed = NgramContainer::new(n_vocab, N, G); + let t_dec_start = unsafe { ffi::ggml_time_us() }; + + let mut id = sample_accept(sampler.0, ctx.0, 0); + print!("{}", token_to_piece(vocab, id)?); + io::stdout().flush().map_err(|err| err.to_string())?; + + loop { + batch_clear(&mut batch.0); + unsafe { + batch_add(&mut batch.0, id, n_past, &seq_id_all, true); + } + + let g_cur = ngrams_observed.observed_for(id); + ngrams_cur.resize(g_cur, NgramData::default()); + for (g, item) in ngrams_cur.iter_mut().enumerate().take(g_cur) { + item.active = true; + item.tokens = vec![0; N]; + item.i_batch = vec![0; N]; + item.seq_id = (W + 1 + g) as i32; + item.i_batch[0] = 0; + item.tokens[0] = id; + } + for j in 0..N - 1 { + for (g, item) in ngrams_cur.iter_mut().enumerate().take(g_cur) { + let token = ngrams_observed.get(id, g, j); + item.tokens[j + 1] = token; + item.i_batch[j + 1] = batch.0.n_tokens; + unsafe { + batch_add( + &mut batch.0, + token, + n_past + j as i32 + 1, + &[(W + 1 + g) as i32], + true, + ); + } + } + } + + for i in 1..W { + let seq_id_look = ((i + 1)..=W).map(|s| s as i32).collect::>(); + unsafe { + batch_add( + &mut batch.0, + tokens_j[0][i], + n_past + i as i32, + &seq_id_look, + false, + ); + } + } + + for j in 1..N - 1 { + for i in 0..W { + unsafe { + batch_add( + &mut batch.0, + tokens_j[j][i], + n_past + j as i32 + i as i32, + &[(i + 1) as i32], + j == N - 2, + ); + } + } + } + + if unsafe { ffi::llama_decode(ctx.0, batch.0) } != 0 { + return Err("llama_decode failed - increase KV cache size".to_string()); + } + + let mut seq_id_best = 0; + for v in 0..N { + let mut i_batch = 0; + if v > 0 { + for item in &ngrams_cur { + if item.active { + i_batch = item.i_batch[v]; + seq_id_best = item.seq_id; + n_accept += 1; + break; + } + } + if i_batch == 0 { + break; + } + } + + id = sample_accept(sampler.0, ctx.0, i_batch); + let token_str = token_to_piece(vocab, id)?; + if v == 0 { + print!("{token_str}"); + } else { + print!("\x1b[0;96m{token_str}\x1b[0m"); + } + io::stdout().flush().map_err(|err| err.to_string())?; + if unsafe { ffi::llama_vocab_is_eog(vocab, id) } { + has_eos = true; + } + all.push(id); + n_predict += 1; + n_past += 1; + + if (args.n_predict >= 0 && n_predict > args.n_predict) || has_eos { + break; + } + + for item in &mut ngrams_cur { + if item.active { + item.active = v != N - 1 && id == item.tokens[v + 1]; + } + } + + tokens_j_prev.copy_from_slice(&tokens_j[0]); + for j in 0..N - 2 { + tokens_j[j] = tokens_j[j + 1].clone(); + } + if v == 0 { + let base = ngrams_cur.len() * (N - 1) + W * (N - 2); + for i in 0..W { + tokens_j[N - 2][i] = + unsafe { ffi::llama_sampler_sample(sampler.0, ctx.0, (base + i) as i32) }; + } + } else { + tokens_j[N - 2] = tokens_j[0].clone(); + } + + if v == 0 { + for f in 0..W { + let first = tokens_j_prev[f]; + let ngram = (0..N - 1).map(|j| tokens_j[j][f]).collect::>(); + ngrams_observed.insert_unique(first, &ngram); + } + } + } + + if (args.n_predict >= 0 && n_predict > args.n_predict) || has_eos { + break; + } + + unsafe { + ffi::llama_memory_seq_rm(mem, -1, n_past, -1); + if seq_id_best != 0 { + ffi::llama_memory_seq_keep(mem, seq_id_best); + ffi::llama_memory_seq_cp(mem, seq_id_best, 0, -1, -1); + ffi::llama_memory_seq_rm(mem, seq_id_best, -1, -1); + for s in 1..(W + G + 1) { + ffi::llama_memory_seq_cp(mem, 0, s as i32, -1, -1); + } + } + } + } + + let t_dec_end = unsafe { ffi::ggml_time_us() }; + println!("\n"); + eprintln!( + "encoded {n_input:4} tokens in {:8.3} seconds, speed: {:8.3} t/s", + (t_enc_end - t_enc_start) as f64 / 1e6, + n_input as f64 / ((t_enc_end - t_enc_start) as f64 / 1e6) + ); + eprintln!( + "decoded {n_predict:4} tokens in {:8.3} seconds, speed: {:8.3} t/s", + (t_dec_end - t_dec_start) as f64 / 1e6, + n_predict as f64 / ((t_dec_end - t_dec_start) as f64 / 1e6) + ); + eprintln!("\nW = {W:2}\nN = {N:2}\nG = {G:2}\n"); + eprintln!("n_predict = {n_predict}"); + eprintln!("n_accept = {n_accept}\n"); + unsafe { + ffi::llama_perf_sampler_print(sampler.0); + ffi::llama_perf_context_print(ctx.0); + } + Ok(()) +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m model.gguf -p \"Hello my name is\""); + eprintln!(); +} + +fn main() { + let mut argv = env::args(); + let program = argv.next().unwrap_or_else(|| "llama-lookahead".to_string()); + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_lookahead_options() { + let args = parse_args([ + "-m", + "model.gguf", + "-p", + "hello", + "-n", + "8", + "-c", + "256", + "-b", + "128", + "-ngl", + "0", + "--top-k", + "10", + "--top-p", + "0.8", + "--temp", + "0.5", + "-s", + "42", + ]) + .unwrap(); + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "hello"); + assert_eq!(args.n_predict, 8); + assert_eq!(args.n_ctx, 256); + assert_eq!(args.n_batch, 128); + assert_eq!(args.n_gpu_layers, 0); + assert_eq!(args.top_k, 10); + assert_eq!(args.top_p, 0.8); + assert_eq!(args.temp, 0.5); + assert_eq!(args.seed, 42); + } + + #[test] + fn rejects_missing_model() { + assert_eq!( + parse_args(["-p", "hello"]).unwrap_err(), + ParseError::MissingModel + ); + } + + #[test] + fn ngram_container_keeps_unique_ring_entries() { + let mut container = NgramContainer::new(8, 3, 2); + assert!(container.insert_unique(2, &[3, 4])); + assert!(!container.insert_unique(2, &[3, 4])); + assert!(container.insert_unique(2, &[4, 5])); + assert!(container.insert_unique(2, &[5, 6])); + assert_eq!(container.cnt[2], 2); + assert_eq!(container.n_total, 3); + assert_eq!(container.observed_for(2), 2); + } +} diff --git a/examples/lookahead/CMakeLists.txt b/examples/lookahead/CMakeLists.txt index 5d6e604fa981..af1fa6deb780 100644 --- a/examples/lookahead/CMakeLists.txt +++ b/examples/lookahead/CMakeLists.txt @@ -1,5 +1,25 @@ set(TARGET llama-lookahead) -add_executable(${TARGET} lookahead.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(LOOKAHEAD_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../lookahead-rust/Cargo.toml) +set(LOOKAHEAD_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/lookahead) +set(LOOKAHEAD_RUST_BIN ${LOOKAHEAD_RUST_TARGET_DIR}/debug/llama-lookahead${CMAKE_EXECUTABLE_SUFFIX}) +set(LOOKAHEAD_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${LOOKAHEAD_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${LOOKAHEAD_RUST_MANIFEST} --target-dir ${LOOKAHEAD_RUST_TARGET_DIR} --bin llama-lookahead + COMMAND ${CMAKE_COMMAND} -E copy ${LOOKAHEAD_RUST_BIN} ${LOOKAHEAD_RUST_OUTPUT} + DEPENDS + llama + ${LOOKAHEAD_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../lookahead-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../lookahead-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${LOOKAHEAD_RUST_OUTPUT}) +install(PROGRAMS ${LOOKAHEAD_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/lookahead/lookahead.cpp b/examples/lookahead/lookahead.cpp deleted file mode 100644 index b7f5c6de86b3..000000000000 --- a/examples/lookahead/lookahead.cpp +++ /dev/null @@ -1,483 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "sampling.h" -#include "log.h" -#include "llama.h" - -#include -#include -#include -#include -#include - -struct ngram_data { - bool active = false; - - llama_seq_id seq_id = -1; - - std::vector i_batch; - - std::vector tokens; -}; - -// n-gram container -struct ngram_container { - ngram_container(int n_vocab, int N, int G) { - cnt.resize(n_vocab); - head.resize(n_vocab); - tokens.resize(n_vocab * G * (N - 1)); - } - - int n_total = 0; - - std::vector cnt; - std::vector head; - - // [n_vocab][G][N - 1] - // for each token of the vocab, keep a ring-buffer of capacity G of n-grams of size N - 1 - std::vector tokens; -}; - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { - return 1; - } - - const int W = 15; // lookahead window - const int N = 5; // n-gram size - const int G = 15; // max verification n-grams - - // lookahead requires W + G + 1 sequences for parallel Jacobi decoding - params.n_parallel = W + G + 1; - - // unified KV cache is required for coupled sequences in batch splitting - params.kv_unified = true; - - // init llama.cpp - llama_backend_init(); - llama_numa_init(params.numa); - - // load the target model - auto llama_init = common_init_from_params(params); - - auto * model = llama_init->model(); - auto * ctx = llama_init->context(); - - auto * mem = llama_get_memory(ctx); - - const llama_vocab * vocab = llama_model_get_vocab(model); - - // Tokenize the prompt - std::vector inp; - std::vector all; - - inp = common_tokenize(ctx, params.prompt, true, true); - all = inp; - - const int max_context_size = llama_n_ctx(ctx); - const int max_tokens_list_size = max_context_size - 4; - - if ((int) inp.size() > max_tokens_list_size) { - LOG_ERR("%s: prompt too long (%d tokens, max %d)\n", __func__, (int) inp.size(), max_tokens_list_size); - return 1; - } - - LOG("\n\n"); - - for (auto id : inp) { - LOG("%s", common_token_to_piece(ctx, id).c_str()); - } - - fflush(stderr); - - const int n_input = inp.size(); - - const auto t_enc_start = ggml_time_us(); - - // eval the prompt - llama_decode(ctx, llama_batch_get_one( inp.data(), n_input - 1)); - llama_decode(ctx, llama_batch_get_one(&inp.back(), 1)); - - for (int s = 1; s < W + G + 1; ++s) { - llama_memory_seq_cp(mem, 0, s, -1, -1); - } - - const auto t_enc_end = ggml_time_us(); - - int n_predict = 0; - int n_accept = 0; - - int n_past = inp.size(); - - llama_token id = 0; - - // used to determine end of generation - bool has_eos = false; - - // for each decoded batch, we have at most W + G + 1 distinct sequences: - // seq_id == 0 : the current input token - // seq_id [1, W] : tokens from the past N - 1 Jacobi iterations - // seq_id [W + 1, W + G] : verification n-grams - llama_batch batch = llama_batch_init(llama_n_ctx(ctx), 0, W + G + 1); - - // target model sampling context - struct common_sampler * smpl = common_sampler_init(model, params.sampling); - - // verification n-grams - std::vector ngrams_cur(G); - - // tokens for the past N - 1 Jacobi iterations - std::vector tokens_j_prev(W); - std::vector> tokens_j(N - 1); - for (int j = 0; j < N - 1; j++) { - tokens_j[j].resize(W); - - for (int i = 0; i < W; i++) { - // there are different ways to init these tokens - if (0) { - // initialize randomly from the prompt tokens - tokens_j[j][i] = all[1 + rand() % (all.size() - 1)]; - } else { - // initialize with a sequence of increasing numbers - tokens_j[j][i] = 100 + i; - } - } - } - - std::vector seq_id_look; - - // the input token belongs both to all sequences - std::vector seq_id_all(W + G + 1); - for (int i = 0; i < W + G + 1; i++) { - seq_id_all[i] = i; - } - - // here we keep adding new n-grams as we go - ngram_container ngrams_observed(llama_vocab_n_tokens(vocab), N, G); - - const auto t_dec_start = ggml_time_us(); - - // sample first token - { - id = common_sampler_sample(smpl, ctx, 0); - - common_sampler_accept(smpl, id, true); - - { - const std::string token_str = common_token_to_piece(ctx, id); - - LOG("%s", token_str.c_str()); - fflush(stdout); - } - } - - while (true) { - // build the mask from https://lmsys.org/blog/2023-11-21-lookahead-decoding/ - // - // Example for W = 5, N = 4, G = 2: - // (I = input, L = lookahead, V = verification) - // - // Batch: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 - // T: -2 -2 -2 -2 -1 -1 -1 -1 -1 0 0 0 0 0 0 - // Info: I L L L L L L L L L L L L L L V V V V V V - // Pos: 0 1 2 3 4 1 2 3 4 5 2 3 4 5 6 1 2 3 1 2 3 (+ n_past) - // Logits: 1 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 - // --------------------------------------------------------------------- - // Seq: 0 - // 1 1 1 - // 2 2 2 2 - // 3 3 3 3 3 - // 4 4 4 4 4 4 - // 5 5 5 5 5 5 5 - // 6 6 6 6 - // 7 7 7 7 - // --------------------------------------------------------------------- - // | | | | | | | | | | | - // V V V V V | | | | | | - // j_tokens | | | | | | - // V V V V V V - // id - { - common_batch_clear(batch); - - // current token - first token of the first level - common_batch_add(batch, id, n_past, seq_id_all, true); - - // verification n-grams - queue this before the lookahead tokens for less KV cache fragmentation - { - const int g_cur = ngrams_observed.cnt[id]; - - ngrams_cur.resize(g_cur); - for (int g = 0; g < g_cur; g++) { - ngrams_cur[g].active = true; - ngrams_cur[g].tokens.resize(N); - ngrams_cur[g].i_batch.resize(N); - ngrams_cur[g].seq_id = W + 1 + g; - ngrams_cur[g].i_batch[0] = 0; - ngrams_cur[g].tokens [0] = id; - } - - for (int j = 0; j < N - 1; j++) { - for (int g = 0; g < g_cur; g++) { - const int idx = id*(N - 1)*G + g*(N - 1); - - const llama_token t = ngrams_observed.tokens[idx + j]; - - ngrams_cur[g].tokens [j + 1] = t; - ngrams_cur[g].i_batch[j + 1] = batch.n_tokens; - - common_batch_add(batch, t, n_past + j + 1, { W + 1 + g }, true); - } - } - } - - // fill the remaining W - 1 tokens for the first level - for (int i = 1; i < W; i++) { - seq_id_look.resize(W - i); - for (int j = 0; j < W - i; j++) { - seq_id_look[j] = i + j + 1; - } - - common_batch_add(batch, tokens_j[0][i], n_past + i, seq_id_look, false); - } - - // fill the rest of the levels - for (int j = 1; j < N - 1; j++) { - for (int i = 0; i < W; i++) { - common_batch_add(batch, tokens_j[j][i], n_past + j + i, { i + 1 }, j == N - 2); - } - } - } - - if (llama_decode(ctx, batch) != 0) { - LOG_ERR("\n\n%s: llama_decode failed - increase KV cache size\n", __func__); - return 1; - } - - int seq_id_best = 0; - - for (int v = 0; v < N; ++v) { - int i_batch = 0; - - // if no active ngrams are left, it means the sampled token does not pass the verification - if (v > 0) { - for (int g = 0; g < (int) ngrams_cur.size(); g++) { - if (ngrams_cur[g].active) { - i_batch = ngrams_cur[g].i_batch[v]; - seq_id_best = ngrams_cur[g].seq_id; - - ++n_accept; - break; - } - } - - // no more matches -> create a new batch - if (i_batch == 0) { - break; - } - } - - // sample the next token - id = common_sampler_sample(smpl, ctx, i_batch); - - common_sampler_accept(smpl, id, true); - - // print - { - const std::string token_str = common_token_to_piece(ctx, id); - - if (v == 0) { - LOG("%s", token_str.c_str()); - } else { - // print light cyan - LOG("\033[0;96m%s\033[0m", token_str.c_str()); - } - fflush(stdout); - - if (llama_vocab_is_eog(vocab, id)) { - has_eos = true; - } - - all.push_back(id); - } - - ++n_predict; - ++n_past; - - if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) { - break; - } - - // verify across active n-grams - for (int g = 0; g < (int) ngrams_cur.size(); g++) { - if (ngrams_cur[g].active) { - if (v == N - 1) { - ngrams_cur[g].active = false; - } else { - if (id != ngrams_cur[g].tokens[v + 1]) { - ngrams_cur[g].active = false; - } - } - } - } - - // print known n-grams starting with token id (debug) - if (0 && v == 0) { - if (ngrams_observed.cnt[id] > 0) { - LOG("\n - %d n-grams starting with '%s'\n", ngrams_observed.cnt[id], common_token_to_piece(ctx, id).c_str()); - } - - for (int i = 0; i < ngrams_observed.cnt[id]; i++) { - LOG(" - ngram %2d: ", i); - - const int idx = id*(N - 1)*G + i*(N - 1); - - for (int j = 0; j < N - 1; j++) { - const std::string token_str = common_token_to_piece(ctx, ngrams_observed.tokens[idx + j]); - - LOG("%s", token_str.c_str()); - } - - LOG("\n"); - } - } - - // update lookahead tokens - { - for (int i = 0; i < W; i++) { - tokens_j_prev[i] = tokens_j[0][i]; - } - - for (int j = 0; j < N - 2; j++) { - tokens_j[j] = tokens_j[j + 1]; - } - - if (v == 0) { - // sample from the last level - for (int i = 0; i < W; i++) { - tokens_j[N - 2][i] = common_sampler_sample(smpl, ctx, ngrams_cur.size()*(N-1) + W*(N - 2) + i); - } - } else { - for (int i = 0; i < W; i++) { - // there are different ways to init these tokens - if (0) { - // random init - tokens_j[N - 2][i] = all[1 + rand() % (all.size() - 1)]; - } else { - // init from the previous level - tokens_j[N - 2][i] = tokens_j[0][i]; - } - } - } - } - - // update observed ngrams - if (v == 0) { - // the first token of the n-gram is determined by the index in the container so it is not stored - std::vector ngram(N - 1); - - // n-gram generation - // ref: https://github.com/hao-ai-lab/LookaheadDecoding/issues/14#issuecomment-1826198518 - for (int f = 0; f < W; ++f) { - const int ft = tokens_j_prev[f]; // first token of the n-gram - - for (int j = 0; j < N - 1; ++j) { - ngram[j] = tokens_j[j][f]; - } - - // filter-out repeating n-grams - { - bool is_unique = true; - - for (int k = 0; k < ngrams_observed.cnt[ft]; ++k) { - const int idx = ft*(N - 1)*G + k*(N - 1); - - bool is_match = true; - for (int j = 0; j < N - 1; ++j) { - if (ngrams_observed.tokens[idx + j] != ngram[j]) { - is_match = false; - break; - } - } - - if (is_match) { - is_unique = false; - break; - } - } - - if (!is_unique) { - continue; - } - } - - const int head = ngrams_observed.head[ft]; - const int idx = ft*(N - 1)*G + head*(N - 1); - - for (int i = 0; i < N - 1; i++) { - ngrams_observed.tokens[idx + i] = ngram[i]; - } - - ngrams_observed.cnt[ft] = std::min(G, ngrams_observed.cnt[ft] + 1); - ngrams_observed.head[ft] = (head + 1) % G; - - ngrams_observed.n_total++; - } - } - } - - if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) { - break; - } - - // KV cache management - // if no verification token matched, we simply remove all cells from this batch -> no fragmentation - llama_memory_seq_rm(mem, -1, n_past, -1); - - if (seq_id_best != 0) { - // if a verification token matched, we keep the best sequence and remove the rest - // this leads to some KV cache fragmentation - llama_memory_seq_keep(mem, seq_id_best); - llama_memory_seq_cp (mem, seq_id_best, 0, -1, -1); - llama_memory_seq_rm (mem, seq_id_best, -1, -1); - - for (int s = 1; s < W + G + 1; ++s) { - llama_memory_seq_cp(mem, 0, s, -1, -1); - } - } - } - - auto t_dec_end = ggml_time_us(); - - LOG("\n\n"); - - LOG_INF("encoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_input, (t_enc_end - t_enc_start) / 1e6f, inp.size() / ((t_enc_end - t_enc_start) / 1e6f)); - LOG_INF("decoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_predict, (t_dec_end - t_dec_start) / 1e6f, n_predict / ((t_dec_end - t_dec_start) / 1e6f)); - - LOG_INF("\n"); - LOG_INF("W = %2d\n", W); - LOG_INF("N = %2d\n", N); - LOG_INF("G = %2d\n", G); - LOG_INF("\n"); - LOG_INF("n_predict = %d\n", n_predict); - LOG_INF("n_accept = %d\n", n_accept); - - LOG_INF("\n"); - common_perf_print(ctx, smpl); - - common_sampler_free(smpl); - - llama_batch_free(batch); - - llama_backend_free(); - - LOG("\n\n"); - - return 0; -} diff --git a/examples/lookup-create-rust/Cargo.toml b/examples/lookup-create-rust/Cargo.toml new file mode 100644 index 000000000000..4d7f5e6883f8 --- /dev/null +++ b/examples/lookup-create-rust/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "llama-lookup-create-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-lookup-merge = { path = "../lookup-merge-rust" } +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-lookup-create" +path = "src/main.rs" diff --git a/examples/lookup-create-rust/build.rs b/examples/lookup-create-rust/build.rs new file mode 100644 index 000000000000..448b2467f104 --- /dev/null +++ b/examples/lookup-create-rust/build.rs @@ -0,0 +1,43 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-lookup-create"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/lookup-create-rust/src/lib.rs b/examples/lookup-create-rust/src/lib.rs new file mode 100644 index 000000000000..3dda9b71a33b --- /dev/null +++ b/examples/lookup-create-rust/src/lib.rs @@ -0,0 +1,154 @@ +use llama_lookup_merge::{LlamaToken, Ngram, NgramCache}; + +const LLAMA_TOKEN_NULL: LlamaToken = -1; +const LLAMA_NGRAM_STATIC: usize = 2; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub cache_path: String, + pub n_gpu_layers: i32, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + MissingModel, + MissingCachePath, + MissingValue(&'static str), + InvalidInteger(&'static str, String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m model.gguf"), + ParseError::MissingCachePath => { + write!(f, "missing required --lookup-cache-static path") + } + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut args = args.into_iter().map(Into::into); + let mut parsed = Args { + model_path: String::new(), + prompt: String::new(), + cache_path: String::new(), + n_gpu_layers: 99, + }; + + while let Some(arg) = args.next() { + match arg.as_str() { + "-m" | "--model" => { + parsed.model_path = args.next().ok_or(ParseError::MissingValue("-m"))?; + } + "-p" | "--prompt" => { + parsed.prompt = args.next().ok_or(ParseError::MissingValue("-p"))?; + } + "-lcs" | "--lookup-cache-static" => { + parsed.cache_path = args.next().ok_or(ParseError::MissingValue("-lcs"))?; + } + "-ngl" | "--n-gpu-layers" => { + let value = args.next().ok_or(ParseError::MissingValue("-ngl"))?; + parsed.n_gpu_layers = value + .parse() + .map_err(|_| ParseError::InvalidInteger("-ngl", value))?; + } + _ => {} + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + if parsed.cache_path.is_empty() { + return Err(ParseError::MissingCachePath); + } + Ok(parsed) +} + +pub fn build_static_ngram_cache(tokens: &[LlamaToken]) -> NgramCache { + let mut cache = NgramCache::new(); + for i in LLAMA_NGRAM_STATIC..tokens.len() { + let ngram = Ngram { + tokens: [ + tokens[i - 2], + tokens[i - 1], + LLAMA_TOKEN_NULL, + LLAMA_TOKEN_NULL, + ], + }; + let token = tokens[i]; + *cache.entry(ngram).or_default().entry(token).or_default() += 1; + } + cache +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_lookup_create_args() { + let args = parse_args([ + "-m", + "model.gguf", + "-p", + "hello", + "-lcs", + "cache.bin", + "-ngl", + "0", + ]) + .unwrap(); + + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "hello"); + assert_eq!(args.cache_path, "cache.bin"); + assert_eq!(args.n_gpu_layers, 0); + } + + #[test] + fn builds_static_two_gram_cache() { + let cache = build_static_ngram_cache(&[1, 2, 3, 1, 2, 4]); + + assert_eq!( + cache[&Ngram { + tokens: [1, 2, -1, -1] + }][&3], + 1 + ); + assert_eq!( + cache[&Ngram { + tokens: [1, 2, -1, -1] + }][&4], + 1 + ); + assert_eq!( + cache[&Ngram { + tokens: [2, 3, -1, -1] + }][&1], + 1 + ); + } + + #[test] + fn rejects_missing_cache_path() { + assert_eq!( + parse_args(["-m", "model.gguf"]).unwrap_err(), + ParseError::MissingCachePath + ); + } +} diff --git a/examples/lookup-create-rust/src/main.rs b/examples/lookup-create-rust/src/main.rs new file mode 100644 index 000000000000..eba64178b721 --- /dev/null +++ b/examples/lookup-create-rust/src/main.rs @@ -0,0 +1,113 @@ +use llama_lookup_create_rust::{build_static_ngram_cache, parse_args, Args}; +use llama_lookup_merge::save_cache; +use llama_simple_rust::ffi; +use std::env; +use std::ffi::CString; +use std::ptr; + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +fn tokenize(vocab: *const ffi::llama_vocab, prompt: &str) -> Result, String> { + let prompt_c = + CString::new(prompt).map_err(|_| "prompt contains an interior NUL byte".to_string())?; + let n = unsafe { + -ffi::llama_tokenize( + vocab, + prompt_c.as_ptr(), + prompt.len() as i32, + ptr::null_mut(), + 0, + true, + true, + ) + }; + if n <= 0 { + return Err("failed to size prompt tokenization".to_string()); + } + + let mut tokens = vec![0_i32; n as usize]; + let check = unsafe { + ffi::llama_tokenize( + vocab, + prompt_c.as_ptr(), + prompt.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + true, + true, + ) + }; + if check < 0 { + return Err("failed to tokenize prompt".to_string()); + } + tokens.truncate(check as usize); + Ok(tokens) +} + +fn run(args: Args) -> Result<(), String> { + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + + unsafe { + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + + let tokens = tokenize(vocab, &args.prompt)?; + eprintln!("tokenization done"); + let cache = build_static_ngram_cache(&tokens); + eprintln!("hashing done, writing file to {}", args.cache_path); + save_cache(&cache, &args.cache_path).map_err(|err| err.to_string())?; + Ok(()) +} + +fn main() { + let mut argv = env::args(); + let program = argv + .next() + .unwrap_or_else(|| "llama-lookup-create".to_string()); + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + eprintln!("usage: {program} -m model.gguf -p PROMPT -lcs cache.bin"); + std::process::exit(1); + } + }; + + let status = match run(args) { + Ok(()) => 0, + Err(err) => { + eprintln!("{program}: error: {err}"); + 1 + } + }; + + unsafe { + ffi::llama_backend_free(); + } + std::process::exit(status); +} diff --git a/examples/lookup-merge-rust/Cargo.toml b/examples/lookup-merge-rust/Cargo.toml new file mode 100644 index 000000000000..e5bb1e176a3c --- /dev/null +++ b/examples/lookup-merge-rust/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "llama-lookup-merge" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] diff --git a/examples/lookup-merge-rust/src/lib.rs b/examples/lookup-merge-rust/src/lib.rs new file mode 100644 index 000000000000..a827989ee5fb --- /dev/null +++ b/examples/lookup-merge-rust/src/lib.rs @@ -0,0 +1,470 @@ +use std::collections::HashMap; +use std::fmt; +use std::fs::File; +use std::io::{self, ErrorKind, Read, Write}; +use std::path::Path; + +pub const LLAMA_NGRAM_MAX: usize = 4; +pub const LLAMA_NGRAM_MIN: usize = 1; +pub const LLAMA_NGRAM_STATIC: usize = 2; +pub const LLAMA_TOKEN_NULL: LlamaToken = -1; + +pub type LlamaToken = i32; +pub type TokenCount = i32; +pub type NgramCachePart = HashMap; +pub type NgramCache = HashMap; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct Ngram { + pub tokens: [LlamaToken; LLAMA_NGRAM_MAX], +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + Help, + MissingPaths, +} + +impl fmt::Display for ParseError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ParseError::Help => write!(f, "help requested"), + ParseError::MissingPaths => write!(f, "missing input or output cache path"), + } + } +} + +impl std::error::Error for ParseError {} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Args { + pub inputs: Vec, + pub output: String, +} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let paths: Vec = args.into_iter().map(Into::into).collect(); + if paths.iter().any(|arg| arg == "-h" || arg == "--help") { + return Err(ParseError::Help); + } + if paths.len() < 2 { + return Err(ParseError::MissingPaths); + } + + let output = paths.last().expect("checked len").clone(); + Ok(Args { + inputs: paths[..paths.len() - 1].to_vec(), + output, + }) +} + +pub fn merge_cache(target: &mut NgramCache, add: NgramCache) { + for (ngram, part) in add { + let merged_part = target.entry(ngram).or_default(); + for (token, count) in part { + *merged_part.entry(token).or_default() += count; + } + } +} + +pub fn update_cache( + cache: &mut NgramCache, + ngram_min: usize, + ngram_max: usize, + input: &[LlamaToken], + nnew: usize, +) { + let input_size = input.len(); + for ngram_size in ngram_min..=ngram_max { + let i_start = input_size.saturating_sub(nnew).max(ngram_size); + for i in i_start..input_size { + let ngram = ngram_from_slice(&input[i - ngram_size..i]); + let token = input[i]; + *cache.entry(ngram).or_default().entry(token).or_default() += 1; + } + } +} + +fn token_from_combined(input: &[LlamaToken], draft: &[LlamaToken], index: usize) -> LlamaToken { + if index < input.len() { + input[index] + } else { + draft[1 + index - input.len()] + } +} + +const DRAFT_MIN_SAMPLE_SIZE_LAX: [i32; LLAMA_NGRAM_MAX] = [2, 2, 1, 1]; +const DRAFT_MIN_PERCENT_LAX: [i32; LLAMA_NGRAM_MAX] = [66, 50, 50, 50]; +const DRAFT_MIN_SAMPLE_SIZE_STRICT: [i32; LLAMA_NGRAM_MAX] = [4, 3, 2, 2]; +const DRAFT_MIN_PERCENT_STRICT: [i32; LLAMA_NGRAM_MAX] = [75, 66, 66, 66]; + +pub fn draft_tokens( + input: &[LlamaToken], + draft: &mut Vec, + n_draft: usize, + ngram_min: usize, + ngram_max: usize, + context: &NgramCache, + dynamic: &NgramCache, + static_cache: &NgramCache, +) { + assert_eq!(draft.len(), 1); + if input.len() < LLAMA_NGRAM_STATIC { + return; + } + + while draft.len() - 1 < n_draft { + let ngram_start_static = input.len() - LLAMA_NGRAM_STATIC + draft.len() - 1; + let mut ngram_static = Ngram { + tokens: [LLAMA_TOKEN_NULL; LLAMA_NGRAM_MAX], + }; + for j in ngram_start_static..ngram_start_static + LLAMA_NGRAM_STATIC { + ngram_static.tokens[j - ngram_start_static] = token_from_combined(input, draft, j); + } + let part_static = static_cache.get(&ngram_static).cloned().unwrap_or_default(); + + let mut ngrams_cd = Vec::new(); + for ngram_size in ngram_min..=ngram_max { + if ngram_size > input.len() + draft.len() - 1 { + continue; + } + let ngram_start_cd = input.len() - ngram_size + draft.len() - 1; + let mut ngram_cd = Ngram { + tokens: [LLAMA_TOKEN_NULL; LLAMA_NGRAM_MAX], + }; + for j in ngram_start_cd..ngram_start_cd + ngram_size { + ngram_cd.tokens[j - ngram_start_cd] = token_from_combined(input, draft, j); + } + ngrams_cd.push(ngram_cd); + } + + let drafted = try_draft_primary( + context, + &ngrams_cd, + &part_static, + &DRAFT_MIN_SAMPLE_SIZE_LAX, + &DRAFT_MIN_PERCENT_LAX, + ) + .or_else(|| { + try_draft_primary( + dynamic, + &ngrams_cd, + &part_static, + &DRAFT_MIN_SAMPLE_SIZE_STRICT, + &DRAFT_MIN_PERCENT_STRICT, + ) + }) + .or_else(|| try_draft_static(static_cache, ngram_static)); + + let Some(token) = drafted else { + break; + }; + draft.push(token); + } +} + +fn try_draft_static(cache: &NgramCache, ngram: Ngram) -> Option { + let part = cache.get(&ngram)?; + let mut max_count = 0; + let mut sum_count = 0; + let mut max_token = LLAMA_TOKEN_NULL; + + for (&token, &count) in part { + if count > max_count { + max_token = token; + max_count = count; + } + sum_count += count; + } + + if sum_count < DRAFT_MIN_SAMPLE_SIZE_LAX[LLAMA_NGRAM_STATIC - 1] { + return None; + } + if 100 * max_count < DRAFT_MIN_PERCENT_LAX[LLAMA_NGRAM_STATIC - 1] * sum_count { + return None; + } + Some(max_token) +} + +fn try_draft_primary( + cache: &NgramCache, + ngrams: &[Ngram], + part_static: &NgramCachePart, + min_sample_size: &[i32; LLAMA_NGRAM_MAX], + min_percent: &[i32; LLAMA_NGRAM_MAX], +) -> Option { + for (i, ngram) in ngrams.iter().enumerate().rev() { + let Some(part_primary) = cache.get(ngram) else { + continue; + }; + + let mut max_count_primary = 0; + let mut max_count_static = 0; + let mut sum_count_primary = 0; + let mut max_token = LLAMA_TOKEN_NULL; + + for (&token, &count_primary) in part_primary { + let count_static = part_static.get(&token).map_or(1, |count| 100 * *count); + if count_primary * count_static > max_count_primary * max_count_static { + max_token = token; + max_count_primary = count_primary; + max_count_static = count_static; + } + sum_count_primary += count_primary; + } + + if sum_count_primary < min_sample_size[i] { + continue; + } + if 100 * max_count_primary < min_percent[i] * sum_count_primary { + continue; + } + return Some(max_token); + } + + None +} + +pub fn ngram_from_slice(slice: &[LlamaToken]) -> Ngram { + let mut tokens = [LLAMA_TOKEN_NULL; LLAMA_NGRAM_MAX]; + for (dst, src) in tokens.iter_mut().zip(slice.iter().copied()) { + *dst = src; + } + Ngram { tokens } +} + +pub fn load_cache(path: impl AsRef) -> io::Result { + let mut file = File::open(path)?; + let mut cache = NgramCache::new(); + + loop { + let Some(ngram) = read_ngram(&mut file)? else { + break; + }; + + let ntokens = read_i32(&mut file, "missing token count")?; + if ntokens <= 0 { + return Err(invalid_data("token count must be positive")); + } + + let mut part = NgramCachePart::new(); + for _ in 0..ntokens { + let token = read_i32(&mut file, "missing token")?; + let count = read_i32(&mut file, "missing count")?; + if count <= 0 { + return Err(invalid_data("count must be positive")); + } + part.insert(token, count); + } + cache.insert(ngram, part); + } + + Ok(cache) +} + +pub fn save_cache(cache: &NgramCache, path: impl AsRef) -> io::Result<()> { + let mut file = File::create(path)?; + for (ngram, part) in cache { + if part.is_empty() { + return Err(invalid_data("token count map must not be empty")); + } + + for token in ngram.tokens { + file.write_all(&token.to_ne_bytes())?; + } + + let ntokens = + i32::try_from(part.len()).map_err(|_| invalid_data("token count map is too large"))?; + file.write_all(&ntokens.to_ne_bytes())?; + + for (token, count) in part { + if *count <= 0 { + return Err(invalid_data("count must be positive")); + } + file.write_all(&token.to_ne_bytes())?; + file.write_all(&count.to_ne_bytes())?; + } + } + + Ok(()) +} + +fn read_ngram(mut reader: impl Read) -> io::Result> { + let mut bytes = [0_u8; LLAMA_NGRAM_MAX * std::mem::size_of::()]; + let mut read = 0; + while read < bytes.len() { + match reader.read(&mut bytes[read..]) { + Ok(0) if read == 0 => return Ok(None), + Ok(0) => return Err(invalid_data("partial ngram record")), + Ok(n) => read += n, + Err(err) => return Err(err), + } + } + + let mut tokens = [0; LLAMA_NGRAM_MAX]; + for (token, chunk) in tokens.iter_mut().zip(bytes.chunks_exact(4)) { + *token = i32::from_ne_bytes(chunk.try_into().expect("chunk size")); + } + Ok(Some(Ngram { tokens })) +} + +fn read_i32(mut reader: impl Read, missing: &'static str) -> io::Result { + let mut bytes = [0_u8; 4]; + reader.read_exact(&mut bytes).map_err(|err| { + if err.kind() == ErrorKind::UnexpectedEof { + invalid_data(missing) + } else { + err + } + })?; + Ok(i32::from_ne_bytes(bytes)) +} + +fn invalid_data(message: &'static str) -> io::Error { + io::Error::new(ErrorKind::InvalidData, message) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::env; + use std::fs; + use std::process; + use std::sync::atomic::{AtomicUsize, Ordering}; + + static NEXT_ID: AtomicUsize = AtomicUsize::new(0); + + fn temp_file(name: &str) -> std::path::PathBuf { + env::temp_dir().join(format!( + "llama-lookup-merge-{name}-{}-{}.bin", + process::id(), + NEXT_ID.fetch_add(1, Ordering::Relaxed) + )) + } + + fn sample_cache() -> NgramCache { + let mut cache = NgramCache::new(); + cache.insert( + Ngram { + tokens: [10, 20, -1, -1], + }, + HashMap::from([(30, 2), (40, 3)]), + ); + cache.insert( + Ngram { + tokens: [5, 6, 7, 8], + }, + HashMap::from([(9, 1)]), + ); + cache + } + + #[test] + fn parses_input_and_output_paths() { + let args = parse_args(["a.bin", "b.bin", "out.bin"]).unwrap(); + + assert_eq!(args.inputs, ["a.bin", "b.bin"]); + assert_eq!(args.output, "out.bin"); + } + + #[test] + fn rejects_missing_paths() { + assert_eq!( + parse_args(["only-input.bin"]).unwrap_err(), + ParseError::MissingPaths + ); + } + + #[test] + fn detects_help_flag() { + assert_eq!(parse_args(["--help"]).unwrap_err(), ParseError::Help); + } + + #[test] + fn round_trips_cache_file() { + let path = temp_file("round-trip"); + let cache = sample_cache(); + + save_cache(&cache, &path).unwrap(); + let loaded = load_cache(&path).unwrap(); + fs::remove_file(path).unwrap(); + + assert_eq!(loaded, cache); + } + + #[test] + fn merges_counts_and_new_parts() { + let ngram = Ngram { + tokens: [10, 20, -1, -1], + }; + let mut target = HashMap::from([(ngram, HashMap::from([(30, 2), (40, 3)]))]); + let add = HashMap::from([ + (ngram, HashMap::from([(30, 5), (50, 7)])), + ( + Ngram { + tokens: [1, 2, 3, -1], + }, + HashMap::from([(4, 9)]), + ), + ]); + + merge_cache(&mut target, add); + + assert_eq!(target[&ngram][&30], 7); + assert_eq!(target[&ngram][&40], 3); + assert_eq!(target[&ngram][&50], 7); + assert_eq!( + target[&Ngram { + tokens: [1, 2, 3, -1] + }][&4], + 9 + ); + } + + #[test] + fn updates_multiple_ngram_sizes_for_new_tail() { + let mut cache = NgramCache::new(); + let tokens = [10, 20, 30, 40]; + + update_cache(&mut cache, 1, 2, &tokens, tokens.len()); + + assert_eq!(cache[&ngram_from_slice(&[10])][&20], 1); + assert_eq!(cache[&ngram_from_slice(&[20, 30])][&40], 1); + } + + #[test] + fn drafts_from_context_cache_before_static_cache() { + let mut context = NgramCache::new(); + context.insert(ngram_from_slice(&[2]), HashMap::from([(3, 2)])); + let mut static_cache = NgramCache::new(); + static_cache.insert(ngram_from_slice(&[1, 2]), HashMap::from([(4, 2)])); + let mut draft = vec![2]; + + draft_tokens( + &[1, 2], + &mut draft, + 1, + LLAMA_NGRAM_MIN, + LLAMA_NGRAM_MAX, + &context, + &NgramCache::new(), + &static_cache, + ); + + assert_eq!(draft, vec![2, 3]); + } + + #[test] + fn rejects_partial_file() { + let path = temp_file("partial"); + fs::write(&path, [1_u8, 2, 3]).unwrap(); + let err = load_cache(&path).unwrap_err(); + fs::remove_file(path).unwrap(); + + assert_eq!(err.kind(), ErrorKind::InvalidData); + } +} diff --git a/examples/lookup-merge-rust/src/main.rs b/examples/lookup-merge-rust/src/main.rs new file mode 100644 index 000000000000..1ccd3d8b7c82 --- /dev/null +++ b/examples/lookup-merge-rust/src/main.rs @@ -0,0 +1,68 @@ +use std::env; +use std::error::Error; +use std::process::ExitCode; + +use llama_lookup_merge::{load_cache, merge_cache, parse_args, save_cache, ParseError}; + +fn print_usage(argv0: &str) { + eprintln!("Merges multiple lookup cache files into a single one."); + eprintln!("Usage: {argv0} [--help] lookup_part_1.bin lookup_part_2.bin ... lookup_merged.bin"); +} + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(AppError::Help(argv0)) => { + print_usage(&argv0); + ExitCode::SUCCESS + } + Err(AppError::Usage(argv0)) => { + print_usage(&argv0); + ExitCode::FAILURE + } + Err(AppError::Other(err)) => { + eprintln!("lookup-merge: {err}"); + ExitCode::FAILURE + } + } +} + +fn run() -> Result<(), AppError> { + let mut raw_args = env::args(); + let argv0 = raw_args + .next() + .unwrap_or_else(|| "llama-lookup-merge".to_string()); + + let args = parse_args(raw_args).map_err(|err| match err { + ParseError::Help => AppError::Help(argv0.clone()), + ParseError::MissingPaths => AppError::Usage(argv0.clone()), + })?; + + eprintln!("lookup-merge: loading file {}", args.inputs[0]); + let mut merged = load_cache(&args.inputs[0])?; + + for input in args.inputs.iter().skip(1) { + eprintln!("lookup-merge: loading file {input}"); + let cache = load_cache(input)?; + merge_cache(&mut merged, cache); + } + + eprintln!("lookup-merge: saving file {}", args.output); + save_cache(&merged, args.output)?; + Ok(()) +} + +enum AppError { + Help(String), + Usage(String), + Other(Box), +} + +impl From for AppError +where + E: Error + 'static, +{ + fn from(value: E) -> Self { + AppError::Other(Box::new(value)) + } +} diff --git a/examples/lookup-rust/Cargo.toml b/examples/lookup-rust/Cargo.toml new file mode 100644 index 000000000000..5de284ee5ee7 --- /dev/null +++ b/examples/lookup-rust/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "llama-lookup-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-lookup-merge = { path = "../lookup-merge-rust" } +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-lookup" +path = "src/main.rs" diff --git a/examples/lookup-rust/build.rs b/examples/lookup-rust/build.rs new file mode 100644 index 000000000000..ee2402398ef4 --- /dev/null +++ b/examples/lookup-rust/build.rs @@ -0,0 +1,45 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); + + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-lookup"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/lookup-rust/src/main.rs b/examples/lookup-rust/src/main.rs new file mode 100644 index 000000000000..9b89c3c4328a --- /dev/null +++ b/examples/lookup-rust/src/main.rs @@ -0,0 +1,696 @@ +use llama_lookup_merge::{ + draft_tokens, load_cache, merge_cache, save_cache, update_cache, NgramCache, LLAMA_NGRAM_MAX, + LLAMA_NGRAM_MIN, +}; +use llama_simple_rust::ffi; +use std::env; +use std::ffi::CString; +use std::io::{self, Write}; +use std::path::Path; +use std::ptr; + +#[derive(Debug, Clone, PartialEq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub n_predict: i32, + pub n_draft: usize, + pub static_cache: String, + pub dynamic_cache: String, + pub n_ctx: u32, + pub n_batch: u32, + pub n_gpu_layers: i32, + pub top_k: i32, + pub top_p: f32, + pub temp: f32, + pub seed: u32, + pub use_color: bool, +} + +impl Default for Args { + fn default() -> Self { + Self { + model_path: String::new(), + prompt: "Hello my name is".to_string(), + n_predict: 32, + n_draft: 16, + static_cache: String::new(), + dynamic_cache: String::new(), + n_ctx: 512, + n_batch: 512, + n_gpu_layers: 99, + top_k: 40, + top_p: 0.95, + temp: 0.8, + seed: ffi::LLAMA_DEFAULT_SEED, + use_color: false, + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub enum ParseError { + MissingModel, + MissingValue(String), + InvalidInteger(String, String), + InvalidFloat(String, String), + InvalidValue(String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m/--model model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + ParseError::InvalidFloat(flag, value) => write!(f, "invalid float for {flag}: {value}"), + ParseError::InvalidValue(message) => write!(f, "{message}"), + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut parsed = Args::default(); + let mut iter = args.into_iter().map(Into::into); + + while let Some(arg) = iter.next() { + match arg.as_str() { + "-m" | "--model" => parsed.model_path = take(&mut iter, &arg)?, + "-p" | "--prompt" => parsed.prompt = take(&mut iter, &arg)?, + "-f" | "--file" | "--prompt-file" => { + let path = take(&mut iter, &arg)?; + parsed.prompt = std::fs::read_to_string(&path).map_err(|err| { + ParseError::InvalidValue(format!("failed to read prompt file {path}: {err}")) + })?; + } + "-n" | "--n-predict" => parsed.n_predict = parse_i32(&mut iter, &arg)?, + "--spec-draft-n-max" => parsed.n_draft = parse_usize(&mut iter, &arg)?, + "-lcs" | "--lookup-cache-static" => parsed.static_cache = take(&mut iter, &arg)?, + "-lcd" | "--lookup-cache-dynamic" => parsed.dynamic_cache = take(&mut iter, &arg)?, + "-c" | "--ctx-size" => parsed.n_ctx = parse_u32(&mut iter, &arg)?, + "-b" | "--batch-size" => parsed.n_batch = parse_u32(&mut iter, &arg)?, + "-ngl" | "--gpu-layers" => parsed.n_gpu_layers = parse_i32(&mut iter, &arg)?, + "--top-k" => parsed.top_k = parse_i32(&mut iter, &arg)?, + "--top-p" => parsed.top_p = parse_f32(&mut iter, &arg)?, + "--temp" => parsed.temp = parse_f32(&mut iter, &arg)?, + "-s" | "--seed" => parsed.seed = parse_u32(&mut iter, &arg)?, + "--color" => parsed.use_color = true, + "--no-color" => parsed.use_color = false, + "-h" | "--help" => return Err(ParseError::MissingModel), + other => { + let mut prompt = vec![other.to_string()]; + prompt.extend(iter); + parsed.prompt = prompt.join(" "); + break; + } + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + if parsed.n_ctx == 0 || parsed.n_batch == 0 { + return Err(ParseError::InvalidValue( + "context and batch sizes must be positive".to_string(), + )); + } + Ok(parsed) +} + +fn take(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + iter.next() + .ok_or_else(|| ParseError::MissingValue(flag.to_string())) +} + +fn parse_i32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_u32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_usize(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_f32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidFloat(flag.to_string(), value)) +} + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Batch(ffi::llama_batch); + +impl Drop for Batch { + fn drop(&mut self) { + unsafe { + ffi::llama_batch_free(self.0); + } + } +} + +struct Sampler(*mut ffi::llama_sampler); + +impl Drop for Sampler { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_sampler_free(self.0); + } + } + } +} + +struct Backend; + +impl Backend { + fn init() -> Self { + unsafe { + ffi::ggml_backend_load_all(); + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + Self + } +} + +impl Drop for Backend { + fn drop(&mut self) { + unsafe { + ffi::llama_backend_free(); + } + } +} + +fn sampler_init(args: &Args) -> Sampler { + unsafe { + let chain = ffi::llama_sampler_chain_init(ffi::llama_sampler_chain_default_params()); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_top_k(args.top_k)); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_top_p(args.top_p, 1)); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_temp(args.temp)); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_dist(args.seed)); + Sampler(chain) + } +} + +fn tokenize( + vocab: *const ffi::llama_vocab, + text: &str, + add_special: bool, + parse_special: bool, +) -> Result, String> { + let text_c = + CString::new(text).map_err(|_| "prompt contains an interior NUL byte".to_string())?; + let n_tokens = unsafe { + -ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + ptr::null_mut(), + 0, + add_special, + parse_special, + ) + }; + if n_tokens <= 0 { + return Err("failed to size tokenization".to_string()); + } + let mut tokens = vec![0_i32; n_tokens as usize]; + let n = unsafe { + ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + add_special, + parse_special, + ) + }; + if n < 0 { + return Err("failed to tokenize prompt".to_string()); + } + tokens.truncate(n as usize); + Ok(tokens) +} + +fn token_to_piece( + vocab: *const ffi::llama_vocab, + token: ffi::llama_token, +) -> Result { + let mut buf = vec![0_i8; 128]; + let n = unsafe { + ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), buf.len() as i32, 0, true) + }; + if n < 0 { + return Err("failed to convert token to piece".to_string()); + } + let bytes = buf[..n as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +fn batch_clear(batch: &mut ffi::llama_batch) { + batch.n_tokens = 0; +} + +unsafe fn batch_add( + batch: &mut ffi::llama_batch, + token: ffi::llama_token, + pos: ffi::llama_pos, + seq_ids: &[ffi::llama_seq_id], + logits: bool, +) { + let i = batch.n_tokens as isize; + *batch.token.offset(i) = token; + *batch.pos.offset(i) = pos; + *batch.n_seq_id.offset(i) = seq_ids.len() as i32; + let seq_slot = *batch.seq_id.offset(i); + for (j, seq_id) in seq_ids.iter().enumerate() { + *seq_slot.add(j) = *seq_id; + } + *batch.logits.offset(i) = if logits { 1 } else { 0 }; + batch.n_tokens += 1; +} + +fn decode_prompt( + ctx: *mut ffi::llama_context, + tokens: &mut [ffi::llama_token], +) -> Result<(), String> { + if tokens.len() > 1 { + let batch = + unsafe { ffi::llama_batch_get_one(tokens.as_mut_ptr(), tokens.len() as i32 - 1) }; + if unsafe { ffi::llama_decode(ctx, batch) } != 0 { + return Err("failed to decode prompt prefix".to_string()); + } + } + let last = tokens + .last_mut() + .ok_or_else(|| "prompt produced no tokens".to_string())?; + let batch = unsafe { ffi::llama_batch_get_one(last as *mut ffi::llama_token, 1) }; + if unsafe { ffi::llama_decode(ctx, batch) } != 0 { + return Err("failed to decode final prompt token".to_string()); + } + Ok(()) +} + +fn load_cache_optional(path: &str, required: bool) -> Result { + if path.is_empty() { + return Ok(NgramCache::new()); + } + match load_cache(path) { + Ok(cache) => Ok(cache), + Err(err) if !required && err.kind() == io::ErrorKind::NotFound => Ok(NgramCache::new()), + Err(err) => Err(format!("failed to load lookup cache {path}: {err}")), + } +} + +fn maybe_save_cache(cache: &NgramCache, path: &str) -> Result<(), String> { + if path.is_empty() { + return Ok(()); + } + if let Some(parent) = Path::new(path).parent() { + if !parent.as_os_str().is_empty() && !parent.exists() { + return Err(format!("directory does not exist: {}", parent.display())); + } + } + save_cache(cache, path).map_err(|err| format!("failed to save lookup cache {path}: {err}")) +} + +fn run(args: Args) -> Result<(), String> { + let _backend = Backend::init(); + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = args.n_ctx; + ctx_params.n_batch = args.n_batch; + ctx_params.n_ubatch = args.n_batch; + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create llama_context".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + + let mut input = tokenize(vocab, &args.prompt, true, true)?; + let mut ngram_cache_context = NgramCache::new(); + let t_start_draft = unsafe { ffi::ggml_time_us() }; + update_cache( + &mut ngram_cache_context, + LLAMA_NGRAM_MIN, + LLAMA_NGRAM_MAX, + &input, + input.len(), + ); + let ngram_cache_static = load_cache_optional(&args.static_cache, true)?; + let mut ngram_cache_dynamic = load_cache_optional(&args.dynamic_cache, false)?; + let t_draft_flat_us = unsafe { ffi::ggml_time_us() } - t_start_draft; + let mut t_draft_us = 0_i64; + + let max_tokens_list_size = unsafe { ffi::llama_n_ctx(ctx.0) as usize }.saturating_sub(4); + if input.len() > max_tokens_list_size { + return Err(format!( + "prompt too long ({} tokens, max {max_tokens_list_size})", + input.len() + )); + } + + println!(); + for token in &input { + print!("{}", token_to_piece(vocab, *token)?); + } + io::stdout().flush().map_err(|err| err.to_string())?; + + let n_input = input.len(); + let t_enc_start = unsafe { ffi::ggml_time_us() }; + decode_prompt(ctx.0, &mut input)?; + let t_enc_end = unsafe { ffi::ggml_time_us() }; + + let mut n_predict = 0_i32; + let mut n_drafted = 0_i32; + let mut n_accept = 0_i32; + let mut n_past = input.len() as i32; + let mut has_eos = false; + let sampler = sampler_init(&args); + let mut draft = Vec::new(); + let mut batch_tgt = + Batch(unsafe { ffi::llama_batch_init(ffi::llama_n_ctx(ctx.0) as i32, 0, 1) }); + let t_dec_start = unsafe { ffi::ggml_time_us() }; + + loop { + let mut i_dft = 0usize; + loop { + let id = unsafe { ffi::llama_sampler_sample(sampler.0, ctx.0, i_dft as i32) }; + unsafe { + ffi::llama_sampler_accept(sampler.0, id); + } + let token_str = token_to_piece(vocab, id)?; + if !args.use_color { + print!("{token_str}"); + } + if unsafe { ffi::llama_vocab_is_eog(vocab, id) } { + has_eos = true; + } + n_predict += 1; + + if i_dft < draft.len() && id == draft[i_dft] { + n_accept += 1; + n_past += 1; + i_dft += 1; + input.push(id); + let t_start = unsafe { ffi::ggml_time_us() }; + update_cache( + &mut ngram_cache_context, + LLAMA_NGRAM_MIN, + LLAMA_NGRAM_MAX, + &input, + 1, + ); + t_draft_us += unsafe { ffi::ggml_time_us() } - t_start; + if args.use_color { + print!("\x1b[34m{token_str}\x1b[0m"); + } + continue; + } + + if args.use_color { + print!("{token_str}"); + } + io::stdout().flush().map_err(|err| err.to_string())?; + draft.clear(); + draft.push(id); + input.push(id); + let t_start = unsafe { ffi::ggml_time_us() }; + update_cache( + &mut ngram_cache_context, + LLAMA_NGRAM_MIN, + LLAMA_NGRAM_MAX, + &input, + 1, + ); + t_draft_us += unsafe { ffi::ggml_time_us() } - t_start; + break; + } + + if (args.n_predict > 0 && n_predict > args.n_predict) || has_eos { + break; + } + + unsafe { + ffi::llama_memory_seq_rm(ffi::llama_get_memory(ctx.0), 0, n_past, -1); + } + batch_clear(&mut batch_tgt.0); + unsafe { + batch_add(&mut batch_tgt.0, draft[0], n_past, &[0], true); + } + let t_start = unsafe { ffi::ggml_time_us() }; + draft_tokens( + &input, + &mut draft, + args.n_draft, + LLAMA_NGRAM_MIN, + LLAMA_NGRAM_MAX, + &ngram_cache_context, + &ngram_cache_dynamic, + &ngram_cache_static, + ); + for (i, token) in draft.iter().enumerate().skip(1) { + unsafe { + batch_add(&mut batch_tgt.0, *token, n_past + i as i32, &[0], true); + } + } + t_draft_us += unsafe { ffi::ggml_time_us() } - t_start; + n_drafted += draft.len().saturating_sub(1) as i32; + + if unsafe { ffi::llama_decode(ctx.0, batch_tgt.0) } != 0 { + return Err("failed to decode drafted batch".to_string()); + } + n_past += 1; + draft.remove(0); + } + + let t_dec_end = unsafe { ffi::ggml_time_us() }; + merge_cache(&mut ngram_cache_dynamic, ngram_cache_context); + maybe_save_cache(&ngram_cache_dynamic, &args.dynamic_cache)?; + + println!("\n"); + eprintln!( + "encoded {n_input:4} tokens in {:8.3} seconds, speed: {:8.3} t/s", + (t_enc_end - t_enc_start) as f64 / 1e6, + input.len() as f64 / ((t_enc_end - t_enc_start) as f64 / 1e6) + ); + eprintln!( + "decoded {n_predict:4} tokens in {:8.3} seconds, speed: {:8.3} t/s", + (t_dec_end - t_dec_start) as f64 / 1e6, + n_predict as f64 / ((t_dec_end - t_dec_start) as f64 / 1e6) + ); + eprintln!("\nn_draft = {}", args.n_draft); + eprintln!("n_predict = {n_predict}"); + eprintln!("n_drafted = {n_drafted}"); + eprintln!("t_draft_flat = {:.2} ms", t_draft_flat_us as f64 * 1e-3); + if n_drafted > 0 && t_draft_us > 0 { + eprintln!( + "t_draft = {:.2} ms, {:.2} us per token, {:.2} tokens per second", + t_draft_us as f64 * 1e-3, + t_draft_us as f64 / n_drafted as f64, + n_drafted as f64 / (1e-6 * t_draft_us as f64) + ); + eprintln!( + "accept = {:.3}%", + 100.0 * n_accept as f64 / n_drafted as f64 + ); + } else { + eprintln!("t_draft = 0.00 ms"); + eprintln!("accept = 0.000%"); + } + eprintln!("n_accept = {n_accept}\n"); + unsafe { + ffi::llama_perf_sampler_print(sampler.0); + ffi::llama_perf_context_print(ctx.0); + } + Ok(()) +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m model.gguf -p \"Hello\" -lcs static.bin -lcd dynamic.bin"); + eprintln!(); +} + +fn main() { + let mut argv = env::args(); + let program = argv.next().unwrap_or_else(|| "llama-lookup".to_string()); + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use llama_lookup_merge::{ngram_from_slice, NgramCache}; + use std::collections::HashMap; + + #[test] + fn parses_lookup_options() { + let args = parse_args([ + "-m", + "model.gguf", + "-p", + "hello", + "-n", + "8", + "--spec-draft-n-max", + "3", + "-lcs", + "static.bin", + "-lcd", + "dynamic.bin", + "-c", + "256", + "-b", + "128", + "-ngl", + "0", + "--top-k", + "10", + "--top-p", + "0.8", + "--temp", + "0.5", + "-s", + "42", + "--color", + ]) + .unwrap(); + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "hello"); + assert_eq!(args.n_predict, 8); + assert_eq!(args.n_draft, 3); + assert_eq!(args.static_cache, "static.bin"); + assert_eq!(args.dynamic_cache, "dynamic.bin"); + assert_eq!(args.n_ctx, 256); + assert_eq!(args.n_batch, 128); + assert_eq!(args.n_gpu_layers, 0); + assert_eq!(args.top_k, 10); + assert_eq!(args.top_p, 0.8); + assert_eq!(args.temp, 0.5); + assert_eq!(args.seed, 42); + assert!(args.use_color); + } + + #[test] + fn rejects_missing_model() { + assert_eq!( + parse_args(["-p", "hello"]).unwrap_err(), + ParseError::MissingModel + ); + } + + #[test] + fn dynamic_cache_missing_is_allowed() { + let cache = load_cache_optional("/tmp/llama-lookup-rust-missing-cache.bin", false).unwrap(); + assert!(cache.is_empty()); + } + + #[test] + fn draft_uses_existing_lookup_cache() { + let mut context = NgramCache::new(); + context.insert(ngram_from_slice(&[10]), HashMap::from([(20, 2)])); + let mut draft = vec![10]; + draft_tokens( + &[1, 10], + &mut draft, + 1, + LLAMA_NGRAM_MIN, + LLAMA_NGRAM_MAX, + &context, + &NgramCache::new(), + &NgramCache::new(), + ); + assert_eq!(draft, vec![10, 20]); + } +} diff --git a/examples/lookup-stats-rust/Cargo.toml b/examples/lookup-stats-rust/Cargo.toml new file mode 100644 index 000000000000..62a47165b1af --- /dev/null +++ b/examples/lookup-stats-rust/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "llama-lookup-stats-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[[bin]] +name = "llama-lookup-stats" +path = "src/main.rs" + +[dependencies] +llama-lookup-merge = { path = "../lookup-merge-rust" } +llama-simple-rust = { path = "../simple-rust" } diff --git a/examples/lookup-stats-rust/build.rs b/examples/lookup-stats-rust/build.rs new file mode 100644 index 000000000000..2541bff2b106 --- /dev/null +++ b/examples/lookup-stats-rust/build.rs @@ -0,0 +1,43 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-lookup-stats"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/lookup-stats-rust/src/lib.rs b/examples/lookup-stats-rust/src/lib.rs new file mode 100644 index 000000000000..63aab22a22a2 --- /dev/null +++ b/examples/lookup-stats-rust/src/lib.rs @@ -0,0 +1,143 @@ +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub lookup_cache_static: Option, + pub lookup_cache_dynamic: Option, + pub n_gpu_layers: i32, + pub n_draft: usize, +} + +impl Default for Args { + fn default() -> Self { + Self { + model_path: String::new(), + prompt: String::new(), + lookup_cache_static: None, + lookup_cache_dynamic: None, + n_gpu_layers: 99, + n_draft: 16, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + MissingModel, + MissingValue(&'static str), + InvalidInteger(&'static str, String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut args = args.into_iter().map(Into::into); + let mut parsed = Args::default(); + + while let Some(arg) = args.next() { + match arg.as_str() { + "-m" | "--model" => parsed.model_path = value(&mut args, "-m")?, + "-p" | "--prompt" => parsed.prompt = value(&mut args, "-p")?, + "-lcs" | "--lookup-cache-static" => { + parsed.lookup_cache_static = Some(value(&mut args, "-lcs")?); + } + "-lcd" | "--lookup-cache-dynamic" => { + parsed.lookup_cache_dynamic = Some(value(&mut args, "-lcd")?); + } + "-ngl" | "--n-gpu-layers" => { + let raw = value(&mut args, "-ngl")?; + parsed.n_gpu_layers = raw + .parse() + .map_err(|_| ParseError::InvalidInteger("-ngl", raw))?; + } + "--spec-draft-n-max" => { + let raw = value(&mut args, "--spec-draft-n-max")?; + parsed.n_draft = raw + .parse() + .map_err(|_| ParseError::InvalidInteger("--spec-draft-n-max", raw))?; + } + _ => {} + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + Ok(parsed) +} + +fn value(args: &mut I, flag: &'static str) -> Result +where + I: Iterator, +{ + args.next().ok_or(ParseError::MissingValue(flag)) +} + +pub fn compute_acceptance(n_accept: i32, n_drafted: i32) -> Option { + if n_drafted <= 0 { + None + } else { + Some(100.0 * n_accept as f32 / n_drafted as f32) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_lookup_stats_options() { + let args = parse_args([ + "-m", + "model.gguf", + "-p", + "hello", + "-lcs", + "static.bin", + "-lcd", + "dynamic.bin", + "-ngl", + "0", + "--spec-draft-n-max", + "8", + ]) + .unwrap(); + + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "hello"); + assert_eq!(args.lookup_cache_static, Some("static.bin".to_string())); + assert_eq!(args.lookup_cache_dynamic, Some("dynamic.bin".to_string())); + assert_eq!(args.n_gpu_layers, 0); + assert_eq!(args.n_draft, 8); + } + + #[test] + fn rejects_missing_model() { + assert_eq!( + parse_args(["-p", "hello"]).unwrap_err(), + ParseError::MissingModel + ); + } + + #[test] + fn avoids_division_by_zero_for_acceptance() { + assert_eq!(compute_acceptance(0, 0), None); + assert_eq!(compute_acceptance(1, 4), Some(25.0)); + } +} diff --git a/examples/lookup-stats-rust/src/main.rs b/examples/lookup-stats-rust/src/main.rs new file mode 100644 index 000000000000..203d229eeff9 --- /dev/null +++ b/examples/lookup-stats-rust/src/main.rs @@ -0,0 +1,269 @@ +use llama_lookup_merge::{ + draft_tokens, load_cache, merge_cache, update_cache, NgramCache, LLAMA_NGRAM_MAX, + LLAMA_NGRAM_MIN, +}; +use llama_lookup_stats_rust::{compute_acceptance, parse_args, Args}; +use llama_simple_rust::ffi; +use std::env; +use std::ffi::CString; +use std::time::Instant; + +struct Backend; + +impl Backend { + fn init() -> Self { + unsafe { + ffi::llama_backend_init(); + } + Self + } +} + +impl Drop for Backend { + fn drop(&mut self) { + unsafe { + ffi::llama_backend_free(); + } + } +} + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +fn tokenize( + vocab: *const ffi::llama_vocab, + text: &[u8], + add_special: bool, + parse_special: bool, +) -> Result, String> { + let mut tokens = vec![0_i32; text.len() + 2 * usize::from(add_special)]; + let n = unsafe { + ffi::llama_tokenize( + vocab, + text.as_ptr() as *const i8, + text.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + add_special, + parse_special, + ) + }; + let n = if n < 0 { + tokens.resize((-n) as usize, 0); + let check = unsafe { + ffi::llama_tokenize( + vocab, + text.as_ptr() as *const i8, + text.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + add_special, + parse_special, + ) + }; + if check != -n { + return Err("tokenization size changed between calls".to_string()); + } + -n + } else { + n + }; + tokens.truncate(n as usize); + Ok(tokens) +} + +fn load_model(args: &Args) -> Result<(Model, Context), String> { + let path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + let mut mparams = unsafe { ffi::llama_model_default_params() }; + mparams.n_gpu_layers = args.n_gpu_layers; + let model = Model(unsafe { ffi::llama_model_load_from_file(path.as_ptr(), mparams) }); + if model.0.is_null() { + return Err(format!("failed to load model '{}'", args.model_path)); + } + let cparams = unsafe { ffi::llama_context_default_params() }; + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, cparams) }); + if ctx.0.is_null() { + return Err("failed to create context".to_string()); + } + Ok((model, ctx)) +} + +fn run(args: Args) -> Result { + let _backend = Backend::init(); + let (model, ctx) = load_model(&args)?; + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get vocab".to_string()); + } + let input = tokenize(vocab, args.prompt.as_bytes(), true, true)?; + + let mut context_cache = NgramCache::new(); + let mut dynamic_cache = NgramCache::new(); + let mut static_cache = NgramCache::new(); + + let flat_start = Instant::now(); + if let Some(path) = &args.lookup_cache_static { + static_cache = load_cache(path) + .map_err(|err| format!("failed to open static lookup cache: {path}: {err}"))?; + } + if let Some(path) = &args.lookup_cache_dynamic { + if let Ok(cache) = load_cache(path) { + dynamic_cache = cache; + } + } + let t_draft_flat = flat_start.elapsed(); + + let n_input = input.len(); + let n_ctx = unsafe { ffi::llama_n_ctx(ctx.0) as usize }; + let mut n_drafted = 0_i32; + let mut n_accept = 0_i32; + let mut t_draft = std::time::Duration::ZERO; + let start = Instant::now(); + + let mut i_start = 0; + while i_start + n_ctx < n_input { + let input_slice = &input[i_start..i_start + n_ctx]; + let mut pseudo_output = vec![input_slice[0]]; + + while pseudo_output.len() < n_ctx { + let mut draft = vec![*pseudo_output.last().unwrap()]; + let draft_start = Instant::now(); + draft_tokens( + &pseudo_output, + &mut draft, + args.n_draft, + LLAMA_NGRAM_MIN, + LLAMA_NGRAM_MAX, + &context_cache, + &dynamic_cache, + &static_cache, + ); + t_draft += draft_start.elapsed(); + n_drafted += (draft.len() - 1) as i32; + + for &drafted in draft.iter().skip(1) { + if pseudo_output.len() >= n_ctx { + break; + } + let ground_truth = input_slice[pseudo_output.len()]; + if ground_truth != drafted { + break; + } + n_accept += 1; + pseudo_output.push(ground_truth); + let update_start = Instant::now(); + update_cache( + &mut context_cache, + LLAMA_NGRAM_MIN, + LLAMA_NGRAM_MAX, + &pseudo_output, + 1, + ); + t_draft += update_start.elapsed(); + } + + if pseudo_output.len() < n_ctx { + pseudo_output.push(input_slice[pseudo_output.len()]); + let update_start = Instant::now(); + update_cache( + &mut context_cache, + LLAMA_NGRAM_MIN, + LLAMA_NGRAM_MAX, + &pseudo_output, + 1, + ); + t_draft += update_start.elapsed(); + } + } + + if i_start > 0 && i_start / 100_000 != i_start.saturating_sub(n_ctx) / 100_000 { + let elapsed_ms = start.elapsed().as_millis() as usize; + let eta_ms = (n_input - i_start) * elapsed_ms / i_start; + eprintln!( + "lookup-stats: {i_start}/{n_input} done, ETA: {:02}:{:02}", + eta_ms / 60_000, + (eta_ms % 60_000) / 1000 + ); + } + + merge_cache(&mut dynamic_cache, std::mem::take(&mut context_cache)); + i_start += n_ctx; + } + + println!(); + println!("n_draft = {}", args.n_draft); + println!("n_predict = {}", n_input - n_input % n_ctx); + println!("n_drafted = {n_drafted}"); + println!( + "t_draft_flat = {:.2} ms", + t_draft_flat.as_secs_f64() * 1000.0 + ); + let draft_us = t_draft.as_secs_f64() * 1_000_000.0; + if n_drafted > 0 && draft_us > 0.0 { + println!( + "t_draft = {:.2} ms, {:.2} us per token, {:.2} tokens per second", + draft_us * 1e-3, + draft_us / f64::from(n_drafted), + f64::from(n_drafted) / (draft_us * 1e-6) + ); + } else { + println!( + "t_draft = {:.2} ms, n/a us per token, n/a tokens per second", + draft_us * 1e-3 + ); + } + println!("n_accept = {n_accept}"); + if let Some(accept) = compute_acceptance(n_accept, n_drafted) { + println!("accept = {accept:.3}%"); + } else { + println!("accept = n/a"); + } + println!(); + + Ok(0) +} + +fn main() { + let mut argv = env::args(); + let program = argv + .next() + .unwrap_or_else(|| "llama-lookup-stats".to_string()); + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("Usage: {program} -m [-p prompt] [-lcs cache] [-lcd cache] [--spec-draft-n-max N]"); + eprintln!("{program}: {err}"); + std::process::exit(1); + } + }; + + match run(args) { + Ok(code) => std::process::exit(code), + Err(err) => { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } + } +} diff --git a/examples/lookup/CMakeLists.txt b/examples/lookup/CMakeLists.txt index 09f7d2e3c922..bf0040e3682d 100644 --- a/examples/lookup/CMakeLists.txt +++ b/examples/lookup/CMakeLists.txt @@ -1,23 +1,99 @@ set(TARGET llama-lookup) -add_executable(${TARGET} lookup.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(LOOKUP_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-rust/Cargo.toml) +set(LOOKUP_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/lookup) +set(LOOKUP_RUST_BIN ${LOOKUP_RUST_TARGET_DIR}/debug/llama-lookup${CMAKE_EXECUTABLE_SUFFIX}) +set(LOOKUP_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${LOOKUP_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${LOOKUP_RUST_MANIFEST} --target-dir ${LOOKUP_RUST_TARGET_DIR} --bin llama-lookup + COMMAND ${CMAKE_COMMAND} -E copy ${LOOKUP_RUST_BIN} ${LOOKUP_RUST_OUTPUT} + DEPENDS + llama + ${LOOKUP_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-merge-rust/src/lib.rs + ${PROJECT_SOURCE_DIR}/examples/simple-rust/build.rs + ${PROJECT_SOURCE_DIR}/examples/simple-rust/src/ffi.rs + ${PROJECT_SOURCE_DIR}/examples/simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${LOOKUP_RUST_OUTPUT}) +install(PROGRAMS ${LOOKUP_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) set(TARGET llama-lookup-create) -add_executable(${TARGET} lookup-create.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +set(LOOKUP_CREATE_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-create-rust/Cargo.toml) +set(LOOKUP_CREATE_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/lookup-create) +set(LOOKUP_CREATE_RUST_BIN ${LOOKUP_CREATE_RUST_TARGET_DIR}/debug/llama-lookup-create${CMAKE_EXECUTABLE_SUFFIX}) +set(LOOKUP_CREATE_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${LOOKUP_CREATE_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${LOOKUP_CREATE_RUST_MANIFEST} --target-dir ${LOOKUP_CREATE_RUST_TARGET_DIR} --bin llama-lookup-create + COMMAND ${CMAKE_COMMAND} -E copy ${LOOKUP_CREATE_RUST_BIN} ${LOOKUP_CREATE_RUST_OUTPUT} + DEPENDS + llama + ${LOOKUP_CREATE_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-create-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-create-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-create-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-merge-rust/src/lib.rs + ${PROJECT_SOURCE_DIR}/examples/simple-rust/build.rs + ${PROJECT_SOURCE_DIR}/examples/simple-rust/src/ffi.rs + ${PROJECT_SOURCE_DIR}/examples/simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${LOOKUP_CREATE_RUST_OUTPUT}) +install(PROGRAMS ${LOOKUP_CREATE_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) set(TARGET llama-lookup-merge) -add_executable(${TARGET} lookup-merge.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(LOOKUP_MERGE_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-merge-rust/Cargo.toml) +set(LOOKUP_MERGE_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/lookup-merge) +set(LOOKUP_MERGE_RUST_BIN ${LOOKUP_MERGE_RUST_TARGET_DIR}/debug/llama-lookup-merge${CMAKE_EXECUTABLE_SUFFIX}) +set(LOOKUP_MERGE_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${LOOKUP_MERGE_RUST_OUTPUT} + COMMAND ${CARGO_EXECUTABLE} build --manifest-path ${LOOKUP_MERGE_RUST_MANIFEST} --target-dir ${LOOKUP_MERGE_RUST_TARGET_DIR} --bin llama-lookup-merge + COMMAND ${CMAKE_COMMAND} -E copy ${LOOKUP_MERGE_RUST_BIN} ${LOOKUP_MERGE_RUST_OUTPUT} + DEPENDS + ${LOOKUP_MERGE_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-merge-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-merge-rust/src/main.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${LOOKUP_MERGE_RUST_OUTPUT}) +install(PROGRAMS ${LOOKUP_MERGE_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) set(TARGET llama-lookup-stats) -add_executable(${TARGET} lookup-stats.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +set(LOOKUP_STATS_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-stats-rust/Cargo.toml) +set(LOOKUP_STATS_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/lookup-stats) +set(LOOKUP_STATS_RUST_BIN ${LOOKUP_STATS_RUST_TARGET_DIR}/debug/llama-lookup-stats${CMAKE_EXECUTABLE_SUFFIX}) +set(LOOKUP_STATS_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${LOOKUP_STATS_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${LOOKUP_STATS_RUST_MANIFEST} --target-dir ${LOOKUP_STATS_RUST_TARGET_DIR} --bin llama-lookup-stats + COMMAND ${CMAKE_COMMAND} -E copy ${LOOKUP_STATS_RUST_BIN} ${LOOKUP_STATS_RUST_OUTPUT} + DEPENDS + llama + ${LOOKUP_STATS_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-stats-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-stats-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-stats-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../lookup-merge-rust/src/lib.rs + ${PROJECT_SOURCE_DIR}/examples/simple-rust/build.rs + ${PROJECT_SOURCE_DIR}/examples/simple-rust/src/ffi.rs + ${PROJECT_SOURCE_DIR}/examples/simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${LOOKUP_STATS_RUST_OUTPUT}) +install(PROGRAMS ${LOOKUP_STATS_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/lookup/lookup-create.cpp b/examples/lookup/lookup-create.cpp deleted file mode 100644 index c0f6c8fc2029..000000000000 --- a/examples/lookup/lookup-create.cpp +++ /dev/null @@ -1,45 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "ngram-cache.h" -#include "llama.h" - -#include -#include -#include - -int main(int argc, char ** argv){ - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_LOOKUP)) { - return 1; - } - - // init llama.cpp - llama_backend_init(); - llama_numa_init(params.numa); - - // load the model - auto llama_init = common_init_from_params(params); - - auto * model = llama_init->model(); - auto * ctx = llama_init->context(); - - GGML_ASSERT(model != nullptr); - - // tokenize the prompt - std::vector inp; - inp = common_tokenize(ctx, params.prompt, true, true); - fprintf(stderr, "%s: tokenization done\n", __func__); - - common_ngram_cache ngram_cache; - common_ngram_cache_update(ngram_cache, LLAMA_NGRAM_STATIC, LLAMA_NGRAM_STATIC, inp, inp.size(), true); - fprintf(stderr, "%s: hashing done, writing file to %s\n", __func__, params.speculative.ngram_cache.lookup_cache_static.c_str()); - - common_ngram_cache_save(ngram_cache, params.speculative.ngram_cache.lookup_cache_static); - - return 0; -} diff --git a/examples/lookup/lookup-merge.cpp b/examples/lookup/lookup-merge.cpp deleted file mode 100644 index ee3c7249cf18..000000000000 --- a/examples/lookup/lookup-merge.cpp +++ /dev/null @@ -1,50 +0,0 @@ -#include "ggml.h" -#include "llama.h" -#include "common.h" -#include "ngram-cache.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -static void print_usage(char* argv0) { - fprintf(stderr, "Merges multiple lookup cache files into a single one.\n"); - fprintf(stderr, "Usage: %s [--help] lookup_part_1.bin lookup_part_2.bin ... lookup_merged.bin\n", argv0); -} - -int main(int argc, char ** argv){ - std::setlocale(LC_NUMERIC, "C"); - - if (argc < 3) { - print_usage(argv[0]); - exit(1); - } - - std::vector args; - args.resize(argc-1); - for (int i = 0; i < argc-1; ++i) { - args[i] = argv[i+1]; - if (args[i] == "-h" || args[i] == "--help") { - print_usage(argv[0]); - exit(0); - } - } - - fprintf(stderr, "lookup-merge: loading file %s\n", args[0].c_str()); - common_ngram_cache ngram_cache_merged = common_ngram_cache_load(args[0]); - - for (size_t i = 1; i < args.size()-1; ++i) { - fprintf(stderr, "lookup-merge: loading file %s\n", args[i].c_str()); - common_ngram_cache ngram_cache = common_ngram_cache_load(args[i]); - - common_ngram_cache_merge(ngram_cache_merged, ngram_cache); - } - - fprintf(stderr, "lookup-merge: saving file %s\n", args.back().c_str()); - common_ngram_cache_save(ngram_cache_merged, args.back()); -} diff --git a/examples/lookup/lookup-stats.cpp b/examples/lookup/lookup-stats.cpp deleted file mode 100644 index 84642e05d1d0..000000000000 --- a/examples/lookup/lookup-stats.cpp +++ /dev/null @@ -1,160 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "log.h" -#include "ngram-cache.h" -#include "llama.h" -#include "ggml.h" - -#include -#include -#include -#include -#include -#include -#include - -int main(int argc, char ** argv){ - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_LOOKUP)) { - return 1; - } - - const int n_draft = params.speculative.draft.n_max; - - // init llama.cpp - llama_backend_init(); - llama_numa_init(params.numa); - - // load the model - auto llama_init = common_init_from_params(params); - - llama_context * ctx = llama_init->context(); - - // tokenize the prompt - std::vector inp; - inp = common_tokenize(ctx, params.prompt, true, true); - - common_ngram_cache ngram_cache_context; - common_ngram_cache ngram_cache_dynamic; - common_ngram_cache ngram_cache_static; - - int64_t t_draft_flat_us = 0; - int64_t t_draft_us = 0; - - { - const int64_t t_start_draft_us = ggml_time_us(); - - if (!params.speculative.ngram_cache.lookup_cache_static.empty()) { - try { - ngram_cache_static = common_ngram_cache_load(params.speculative.ngram_cache.lookup_cache_static); - } catch (std::ifstream::failure const &) { - LOG_ERR("failed to open static lookup cache: %s", params.speculative.ngram_cache.lookup_cache_static.c_str()); - exit(1); - } - } - - if (!params.speculative.ngram_cache.lookup_cache_dynamic.empty()) { - try { - ngram_cache_dynamic = common_ngram_cache_load(params.speculative.ngram_cache.lookup_cache_dynamic); - } catch (std::ifstream::failure const &) {} // if the file does not exist it will simply be created at the end of the program - } - - t_draft_flat_us += ggml_time_us() - t_start_draft_us; - } - - const int n_input = inp.size(); - const int n_ctx = llama_n_ctx(ctx); - - int n_drafted = 0; - int n_accept = 0; - - const int64_t t_start_ms = ggml_time_ms(); - - // Iterate over input tokens in chunks of size n_ctx. - // Each chunk is treated as if a sequential generation but with pre-determined tokens to ensure reproducibility. - for (int i_start = 0; i_start + n_ctx < n_input; i_start += n_ctx) { - const std::vector inp_slice(inp.begin() + i_start, inp.begin() + i_start + n_ctx); - std::vector pseudo_output; - pseudo_output.push_back(inp_slice[0]); - - while ((int) pseudo_output.size() < n_ctx) { - // Simulate drafting and decoding from draft: - std::vector draft; - draft.push_back(pseudo_output.back()); - - { - const int64_t t_start_draft_us = ggml_time_us(); - common_ngram_cache_draft(pseudo_output, draft, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, ngram_cache_context, ngram_cache_dynamic, ngram_cache_static); - t_draft_us += ggml_time_us() - t_start_draft_us; - } - - n_drafted += draft.size() - 1; - - for (size_t j = 1; j < draft.size() && (int) pseudo_output.size() < n_ctx; ++j) { - const llama_token ground_truth = inp_slice[pseudo_output.size()]; - const llama_token drafted = draft[j]; - - if (ground_truth != drafted) { - break; - } - - ++n_accept; - pseudo_output.push_back(ground_truth); - - { - const int64_t t_start_draft_us = ggml_time_us(); - common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, pseudo_output, 1, false); - t_draft_us += ggml_time_us() - t_start_draft_us; - } - } - - // After each simulated batch decoding simulate the sampling of a single token: - if ((int) pseudo_output.size() < n_ctx) { - pseudo_output.push_back(inp_slice[pseudo_output.size()]); - { - const int64_t t_start_draft_us = ggml_time_us(); - common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, pseudo_output, 1, false); - t_draft_us += ggml_time_us() - t_start_draft_us; - } - } - - draft.erase(draft.begin()); - - } - if (i_start > 0 && i_start / 100000 != (i_start - n_ctx) / 100000) { - const int64_t t_now_ms = ggml_time_ms(); - const int64_t eta_ms = (n_input - i_start) * (t_now_ms - t_start_ms) / i_start; - const int64_t eta_min = eta_ms / (60*1000); - const int64_t eta_s = (eta_ms - 60*1000*eta_min) / 1000; - - LOG_INF("lookup-stats: %d/%d done, ETA: %02" PRId64 ":%02" PRId64 "\n", i_start, n_input, eta_min, eta_s); - } - - // After each chunk, update the dynamic ngram cache with the context ngram cache: - common_ngram_cache_merge(ngram_cache_dynamic, ngram_cache_context); - ngram_cache_context.clear(); - } - - LOG("\n"); - - LOG_INF("\n"); - LOG_INF("n_draft = %d\n", n_draft); - LOG_INF("n_predict = %d\n", n_input - n_input % n_ctx); - LOG_INF("n_drafted = %d\n", n_drafted); - LOG_INF("t_draft_flat = %.2f ms\n", t_draft_flat_us*1e-3); - LOG_INF("t_draft = %.2f ms, %.2f us per token, %.2f tokens per second\n", - t_draft_us*1e-3, 1.0f*t_draft_us/n_drafted, n_drafted/(1e-6*t_draft_us)); - LOG_INF("n_accept = %d\n", n_accept); - LOG_INF("accept = %.3f%%\n", 100.0f * n_accept / n_drafted); - - llama_backend_free(); - - LOG("\n\n"); - - return 0; -} diff --git a/examples/lookup/lookup.cpp b/examples/lookup/lookup.cpp deleted file mode 100644 index 2d4c0e528d3d..000000000000 --- a/examples/lookup/lookup.cpp +++ /dev/null @@ -1,245 +0,0 @@ -#include "arg.h" -#include "ggml.h" -#include "common.h" -#include "ngram-cache.h" -#include "sampling.h" -#include "log.h" -#include "llama.h" - -#include -#include -#include -#include -#include -#include - -int main(int argc, char ** argv){ - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_LOOKUP)) { - return 1; - } - - // max. number of additional tokens to draft if match is found - const int n_draft = params.speculative.draft.n_max; - - // init llama.cpp - llama_backend_init(); - llama_numa_init(params.numa); - - // load the model - auto llama_init = common_init_from_params(params); - - auto * model = llama_init->model(); - auto * ctx = llama_init->context(); - - const llama_vocab * vocab = llama_model_get_vocab(model); - - // tokenize the prompt - std::vector inp; - inp = common_tokenize(ctx, params.prompt, true, true); - - common_ngram_cache ngram_cache_context; - common_ngram_cache ngram_cache_dynamic; - common_ngram_cache ngram_cache_static; - int64_t t_draft_flat_us = 0; - int64_t t_draft_us = 0; - - { - // Fill up context ngram cache with tokens from user input: - const int64_t t_start_draft_us = ggml_time_us(); - common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, inp, inp.size(), false); - - if (!params.speculative.ngram_cache.lookup_cache_static.empty()) { - try { - ngram_cache_static = common_ngram_cache_load(params.speculative.ngram_cache.lookup_cache_static); - } catch (std::ifstream::failure const &) { - LOG_ERR("failed to open static lookup cache: %s", params.speculative.ngram_cache.lookup_cache_static.c_str()); - exit(1); - } - } - - if (!params.speculative.ngram_cache.lookup_cache_dynamic.empty()) { - try { - ngram_cache_dynamic = common_ngram_cache_load(params.speculative.ngram_cache.lookup_cache_dynamic); - } catch (std::ifstream::failure const &) {} // if the file does not exist it will simply be created at the end of the program - } - - t_draft_flat_us += ggml_time_us() - t_start_draft_us; - } - - const int max_context_size = llama_n_ctx(ctx); - const int max_tokens_list_size = max_context_size - 4; - - if ((int) inp.size() > max_tokens_list_size) { - LOG_ERR("%s: prompt too long (%d tokens, max %d)\n", __func__, (int) inp.size(), max_tokens_list_size); - return 1; - } - - LOG("\n\n"); - - for (auto id : inp) { - LOG("%s", common_token_to_piece(ctx, id).c_str()); - } - - fflush(stderr); - - const int n_input = inp.size(); - - const auto t_enc_start = ggml_time_us(); - - llama_decode(ctx, llama_batch_get_one( inp.data(), n_input - 1)); - llama_decode(ctx, llama_batch_get_one(&inp.back(), 1)); - - const auto t_enc_end = ggml_time_us(); - - int n_predict = 0; - int n_drafted = 0; - int n_accept = 0; - - int n_past = inp.size(); - - bool has_eos = false; - - struct common_sampler * smpl = common_sampler_init(model, params.sampling); - - std::vector draft; - - llama_batch batch_tgt = llama_batch_init(llama_n_ctx(ctx), 0, 1); - - const auto t_dec_start = ggml_time_us(); - - while (true) { - // print current draft sequence - LOG_DBG("drafted %s\n", string_from(ctx, draft).c_str()); - - int i_dft = 0; - while (true) { - // sample from the target model - llama_token id = common_sampler_sample(smpl, ctx, i_dft); - - common_sampler_accept(smpl, id, true); - - const std::string token_str = common_token_to_piece(ctx, id); - - if (!params.use_color) { - LOG("%s", token_str.c_str()); - } - - if (llama_vocab_is_eog(vocab, id)) { - has_eos = true; - } - - ++n_predict; - - // check if the target token matches the draft - if (i_dft < (int) draft.size() && id == draft[i_dft]) { - LOG_DBG("the sampled target token matches the %dth drafted token (%d, '%s') - accepted\n", i_dft, id, token_str.c_str()); - ++n_accept; - ++n_past; - ++i_dft; - inp.push_back(id); - { - // Update context ngram cache with the newly accepted token: - const int64_t t_start_draft_us = ggml_time_us(); - common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, inp, 1, false); - t_draft_us += ggml_time_us() - t_start_draft_us; - } - - if (params.use_color) { - // color accepted draft token - LOG("\033[34m%s\033[0m", token_str.c_str()); - fflush(stdout); - } - continue; - } - - if (params.use_color) { - LOG("%s", token_str.c_str()); - } - fflush(stdout); - - - LOG_DBG("the sampled target token (%d, '%s') did not match, or we ran out of drafted tokens\n", id, token_str.c_str()); - - draft.clear(); - draft.push_back(id); - inp.push_back(id); - { - // Update context ngram cache with the newly accepted token: - const int64_t t_start_draft_us = ggml_time_us(); - common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, inp, 1, false); - t_draft_us += ggml_time_us() - t_start_draft_us; - } - break; - } - - if ((params.n_predict > 0 && n_predict > params.n_predict) || has_eos) { - break; - } - - // KV cache management - // clean the cache of draft tokens that weren't accepted - llama_memory_seq_rm(llama_get_memory(ctx), 0, n_past, -1); - - common_batch_clear(batch_tgt); - common_batch_add(batch_tgt, draft[0], n_past, { 0 }, true); - - // Draft already contains a single token sampled from the model: - GGML_ASSERT(draft.size() == 1); - GGML_ASSERT(draft[0] == inp.back()); - const int64_t t_start_draft_us = ggml_time_us(); - - common_ngram_cache_draft(inp, draft, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, ngram_cache_context, ngram_cache_dynamic, ngram_cache_static); - - for (size_t i = 1; i < draft.size(); ++i) { - common_batch_add(batch_tgt, draft[i], n_past + i, { 0 }, true); - } - - t_draft_us += ggml_time_us() - t_start_draft_us; - n_drafted += draft.size() - 1; - - llama_decode(ctx, batch_tgt); - ++n_past; - - draft.erase(draft.begin()); - } - - auto t_dec_end = ggml_time_us(); - - // Update dynamic ngram cache with context ngram cache and save it to disk: - common_ngram_cache_merge(ngram_cache_dynamic, ngram_cache_context); - common_ngram_cache_save(ngram_cache_dynamic, params.speculative.ngram_cache.lookup_cache_dynamic); - - LOG("\n\n"); - - LOG_INF("encoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_input, (t_enc_end - t_enc_start) / 1e6f, inp.size() / ((t_enc_end - t_enc_start) / 1e6f)); - LOG_INF("decoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_predict, (t_dec_end - t_dec_start) / 1e6f, n_predict / ((t_dec_end - t_dec_start) / 1e6f)); - - LOG_INF("\n"); - LOG_INF("n_draft = %d\n", n_draft); - LOG_INF("n_predict = %d\n", n_predict); - LOG_INF("n_drafted = %d\n", n_drafted); - LOG_INF("t_draft_flat = %.2f ms\n", t_draft_flat_us*1e-3); - LOG_INF("t_draft = %.2f ms, %.2f us per token, %.2f tokens per second\n", - t_draft_us*1e-3, 1.0f*t_draft_us/n_drafted, n_drafted/(1e-6*t_draft_us)); - LOG_INF("n_accept = %d\n", n_accept); - LOG_INF("accept = %.3f%%\n", 100.0f * n_accept / n_drafted); - - LOG_INF("\ntarget:\n\n"); - common_perf_print(ctx, smpl); - - common_sampler_free(smpl); - - llama_batch_free(batch_tgt); - - llama_backend_free(); - - LOG("\n\n"); - - return 0; -} diff --git a/examples/parallel-rust/Cargo.toml b/examples/parallel-rust/Cargo.toml new file mode 100644 index 000000000000..346d7260c271 --- /dev/null +++ b/examples/parallel-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-parallel-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-parallel" +path = "src/main.rs" diff --git a/examples/parallel-rust/build.rs b/examples/parallel-rust/build.rs new file mode 100644 index 000000000000..c278d31a897e --- /dev/null +++ b/examples/parallel-rust/build.rs @@ -0,0 +1,45 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); + + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-parallel"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/parallel-rust/src/main.rs b/examples/parallel-rust/src/main.rs new file mode 100644 index 000000000000..9ef5c9edd6d1 --- /dev/null +++ b/examples/parallel-rust/src/main.rs @@ -0,0 +1,917 @@ +use llama_simple_rust::ffi; +use std::env; +use std::ffi::CString; +use std::ptr; + +const SYSTEM_PROMPT: &str = r#"Transcript of a never ending dialog, where the User interacts with an Assistant. +The Assistant is helpful, kind, honest, good at writing, and never fails to answer the User's requests immediately and with precision. + +User: +Recommend a nice restaurant in the area. +Assistant: +I recommend the restaurant "The Golden Duck". It is a 5 star restaurant with a great view of the city. The food is delicious and the service is excellent. The prices are reasonable and the portions are generous. The restaurant is located at 123 Main Street, New York, NY 10001. The phone number is (212) 555-1234. The hours are Monday through Friday from 11:00 am to 10:00 pm. The restaurant is closed on Saturdays and Sundays. +User: +Who is Richard Feynman? +Assistant: +Richard Feynman was an American physicist who is best known for his work in quantum mechanics and particle physics. He was awarded the Nobel Prize in Physics in 1965 for his contributions to the development of quantum electrodynamics. He was a popular lecturer and author, and he wrote several books, including "Surely You're Joking, Mr. Feynman!" and "What Do You Care What Other People Think?". +"#; + +const QUESTIONS: &[&str] = &[ + "What is the tallest mountain in the world?", + "Who was the first person to win two Nobel Prizes?", + "Which country invented paper?", + "What organ is primarily responsible for pumping blood throughout the body?", + "Which planet is known for its prominent ring system?", + "Who directed the movie 'Inception'?", + "What is the freezing point of water in Fahrenheit?", + "Which animal is known to have the longest lifespan?", + "What language has the most native speakers worldwide?", + "What is the capital city of Canada?", + "Who is credited with inventing the World Wide Web?", + "Which metal is liquid at room temperature?", + "What is the term for an animal that eats both plants and meat?", + "Who painted 'The Starry Night'?", + "What gas do humans exhale that plants use for photosynthesis?", + "What year did World War II end?", + "Which continent has the most countries?", + "Who wrote the novel 'Frankenstein'?", + "What does DNA stand for?", + "What is the main ingredient in traditional Japanese miso soup?", +]; + +const ANSWERS: &[&str] = &[ + "The tallest mountain in the world is Mount Everest.", + "Marie Curie was the first person to win two Nobel Prizes.", + "Paper was invented in China.", + "The heart is the organ responsible for pumping blood.", + "Saturn is known for its prominent ring system.", + "Christopher Nolan directed the movie 'Inception'.", + "The freezing point of water in Fahrenheit is 32 degrees F.", + "The bowhead whale is known to have the longest lifespan among mammals.", + "Mandarin Chinese has the most native speakers in the world.", + "The capital city of Canada is Ottawa.", + "Tim Berners-Lee is credited with inventing the World Wide Web.", + "Mercury is the metal that is liquid at room temperature.", + "An animal that eats both plants and meat is called an omnivore.", + "'The Starry Night' was painted by Vincent van Gogh.", + "Humans exhale carbon dioxide, which plants use in photosynthesis.", + "World War II ended in 1945.", + "Africa is the continent with the most countries.", + "The novel 'Frankenstein' was written by Mary Shelley.", + "DNA stands for Deoxyribonucleic Acid.", + "The main ingredient in traditional Japanese miso soup is fermented soybean paste.", +]; + +const DEFAULT_PROMPTS: &[&str] = &[ + "What is the meaning of life?", + "Tell me an interesting fact about llamas.", + "What is the best way to cook a steak?", + "Are you familiar with the Special Theory of Relativity and can you explain it to me?", + "Recommend some interesting books to read.", + "What is the best way to learn a new language?", + "How to get a job at Google?", + "If you could have any superpower, what would it be?", + "I want to learn how to play the piano. What would be the best way to do it?", +]; + +#[derive(Debug, Clone, PartialEq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub prompt_file: String, + pub n_predict: i32, + pub n_parallel: i32, + pub n_sequences: i32, + pub cont_batching: bool, + pub is_pp_shared: bool, + pub n_junk: i32, + pub n_ctx: u32, + pub n_batch: u32, + pub n_gpu_layers: i32, + pub top_k: i32, + pub top_p: f32, + pub temp: f32, + pub seed: i32, +} + +impl Default for Args { + fn default() -> Self { + Self { + model_path: String::new(), + prompt: String::new(), + prompt_file: String::new(), + n_predict: 128, + n_parallel: 1, + n_sequences: 1, + cont_batching: true, + is_pp_shared: false, + n_junk: 1, + n_ctx: 512, + n_batch: 512, + n_gpu_layers: 99, + top_k: 40, + top_p: 0.95, + temp: 0.8, + seed: ffi::LLAMA_DEFAULT_SEED as i32, + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub enum ParseError { + MissingModel, + MissingValue(String), + InvalidInteger(String, String), + InvalidFloat(String, String), + InvalidValue(String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m/--model model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + ParseError::InvalidFloat(flag, value) => write!(f, "invalid float for {flag}: {value}"), + ParseError::InvalidValue(message) => write!(f, "{message}"), + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut parsed = Args::default(); + let mut iter = args.into_iter().map(Into::into); + + while let Some(arg) = iter.next() { + match arg.as_str() { + "-m" | "--model" => parsed.model_path = take(&mut iter, &arg)?, + "-p" | "--prompt" => parsed.prompt = take(&mut iter, &arg)?, + "-f" | "--file" | "--prompt-file" => { + parsed.prompt_file = take(&mut iter, &arg)?; + parsed.prompt = std::fs::read_to_string(&parsed.prompt_file).map_err(|err| { + ParseError::InvalidValue(format!( + "failed to read prompt file {}: {err}", + parsed.prompt_file + )) + })?; + } + "-n" | "--n-predict" => parsed.n_predict = parse_i32(&mut iter, &arg)?, + "-np" | "--parallel" => parsed.n_parallel = parse_i32(&mut iter, &arg)?, + "-ns" | "--sequences" => parsed.n_sequences = parse_i32(&mut iter, &arg)?, + "--cont-batching" => parsed.cont_batching = parse_bool(&mut iter, &arg)?, + "--no-cont-batching" => parsed.cont_batching = false, + "-pps" | "--parallel-prompt-shared" => parsed.is_pp_shared = true, + "--junk" => parsed.n_junk = parse_i32(&mut iter, &arg)?, + "-c" | "--ctx-size" => parsed.n_ctx = parse_u32(&mut iter, &arg)?, + "-b" | "--batch-size" => parsed.n_batch = parse_u32(&mut iter, &arg)?, + "-ngl" | "--gpu-layers" => parsed.n_gpu_layers = parse_i32(&mut iter, &arg)?, + "--top-k" => parsed.top_k = parse_i32(&mut iter, &arg)?, + "--top-p" => parsed.top_p = parse_f32(&mut iter, &arg)?, + "--temp" => parsed.temp = parse_f32(&mut iter, &arg)?, + "-s" | "--seed" => parsed.seed = parse_i32(&mut iter, &arg)?, + "-h" | "--help" => return Err(ParseError::MissingModel), + _ => {} + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + if parsed.n_parallel <= 0 { + return Err(ParseError::InvalidValue( + "-np/--parallel must be positive".to_string(), + )); + } + if parsed.n_sequences <= 0 { + return Err(ParseError::InvalidValue( + "-ns/--sequences must be positive".to_string(), + )); + } + if parsed.n_ctx == 0 || parsed.n_batch == 0 { + return Err(ParseError::InvalidValue( + "context and batch sizes must be positive".to_string(), + )); + } + parsed.n_junk = parsed.n_junk.max(1); + Ok(parsed) +} + +fn take(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + iter.next() + .ok_or_else(|| ParseError::MissingValue(flag.to_string())) +} + +fn parse_i32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_u32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_f32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidFloat(flag.to_string(), value)) +} + +fn parse_bool(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + match take(iter, flag)?.as_str() { + "1" | "true" | "on" | "yes" => Ok(true), + "0" | "false" | "off" | "no" => Ok(false), + value => Err(ParseError::InvalidValue(format!( + "invalid boolean for {flag}: {value}" + ))), + } +} + +pub fn trim_ascii(text: &str) -> String { + text.trim_matches(|c: char| c.is_ascii_whitespace()) + .to_string() +} + +pub fn split_prompt_lines(input: &str) -> Vec { + input.lines().map(str::to_string).collect() +} + +#[derive(Debug, Clone)] +struct Lcg { + state: u32, +} + +impl Lcg { + fn new(seed: u32) -> Self { + Self { state: seed } + } + + fn next(&mut self) -> u32 { + self.state = self.state.wrapping_mul(1103515245).wrapping_add(12345); + (self.state / 65536) % 32768 + } + + fn next_usize(&mut self, limit: usize) -> usize { + if limit == 0 { + 0 + } else { + self.next() as usize % limit + } + } +} + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Batch(ffi::llama_batch); + +impl Drop for Batch { + fn drop(&mut self) { + unsafe { + ffi::llama_batch_free(self.0); + } + } +} + +struct Sampler(*mut ffi::llama_sampler); + +impl Drop for Sampler { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_sampler_free(self.0); + } + } + } +} + +struct Client { + id: i32, + seq_id: i32, + sampled: ffi::llama_token, + t_start_prompt: i64, + t_start_gen: i64, + n_past: i32, + n_prompt: i32, + n_decoded: i32, + i_batch: i32, + input: String, + prompt: String, + response: String, + sampler: Sampler, +} + +struct Backend; + +impl Backend { + fn init() -> Self { + unsafe { + ffi::ggml_backend_load_all(); + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + Self + } +} + +impl Drop for Backend { + fn drop(&mut self) { + unsafe { + ffi::llama_backend_free(); + } + } +} + +fn sampler_init(args: &Args, seed: i32) -> Sampler { + unsafe { + let chain = ffi::llama_sampler_chain_init(ffi::llama_sampler_chain_default_params()); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_top_k(args.top_k)); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_top_p(args.top_p, 1)); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_temp(args.temp)); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_dist(seed as u32)); + Sampler(chain) + } +} + +fn tokenize( + vocab: *const ffi::llama_vocab, + text: &str, + add_special: bool, + parse_special: bool, +) -> Result, String> { + let text_c = + CString::new(text).map_err(|_| "text contains an interior NUL byte".to_string())?; + let n_tokens = unsafe { + -ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + ptr::null_mut(), + 0, + add_special, + parse_special, + ) + }; + if n_tokens <= 0 { + return Err("failed to size tokenization".to_string()); + } + let mut tokens = vec![0_i32; n_tokens as usize]; + let n = unsafe { + ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + add_special, + parse_special, + ) + }; + if n < 0 { + return Err("failed to tokenize text".to_string()); + } + tokens.truncate(n as usize); + Ok(tokens) +} + +fn token_to_piece( + vocab: *const ffi::llama_vocab, + token: ffi::llama_token, +) -> Result { + let mut buf = vec![0_i8; 128]; + let n = unsafe { + ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), buf.len() as i32, 0, true) + }; + if n < 0 { + return Err("failed to convert token to piece".to_string()); + } + let bytes = buf[..n as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +fn batch_clear(batch: &mut ffi::llama_batch) { + batch.n_tokens = 0; +} + +unsafe fn batch_add( + batch: &mut ffi::llama_batch, + token: ffi::llama_token, + pos: ffi::llama_pos, + seq_ids: &[ffi::llama_seq_id], + logits: bool, +) { + let i = batch.n_tokens as isize; + *batch.token.offset(i) = token; + *batch.pos.offset(i) = pos; + *batch.n_seq_id.offset(i) = seq_ids.len() as i32; + let seq_slot = *batch.seq_id.offset(i); + for (j, seq_id) in seq_ids.iter().enumerate() { + *seq_slot.add(j) = *seq_id; + } + *batch.logits.offset(i) = if logits { 1 } else { 0 }; + batch.n_tokens += 1; +} + +unsafe fn set_last_batch_logits(batch: &mut ffi::llama_batch) { + if batch.n_tokens > 0 { + *batch.logits.offset(batch.n_tokens as isize - 1) = 1; + } +} + +unsafe fn batch_view(batch: &ffi::llama_batch, start: i32, n_tokens: i32) -> ffi::llama_batch { + let offset = start as isize; + ffi::llama_batch { + n_tokens, + token: batch.token.offset(offset), + embd: ptr::null_mut(), + pos: batch.pos.offset(offset), + n_seq_id: batch.n_seq_id.offset(offset), + seq_id: batch.seq_id.offset(offset), + logits: batch.logits.offset(offset), + } +} + +fn build_prompt( + is_shared: bool, + n_tokens_system: i32, + input: &str, + n_junk: i32, + rng: &mut Lcg, +) -> (String, i32, i32) { + let mut n_past = 0; + let mut prompt = String::new(); + if is_shared { + n_past = n_tokens_system; + } else { + prompt.push_str(SYSTEM_PROMPT); + } + let n_junk_cur = rng.next_usize(n_junk.max(1) as usize) as i32; + for _ in 0..n_junk_cur { + let r = rng.next_usize(QUESTIONS.len()); + prompt.push_str("User:\n"); + prompt.push_str(QUESTIONS[r]); + prompt.push_str("\nAssistant:\n "); + prompt.push_str(ANSWERS[r]); + prompt.push('\n'); + } + prompt.push_str("User:\n"); + prompt.push_str(input); + prompt.push_str("\nAssistant:\n"); + (prompt, n_past, n_junk_cur) +} + +fn run(args: Args) -> Result<(), String> { + let _backend = Backend::init(); + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = args.n_ctx; + ctx_params.n_batch = args.n_batch; + ctx_params.n_ubatch = args.n_batch; + ctx_params.n_seq_max = (args.n_parallel + 1) as u32; + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create llama_context".to_string()); + } + + let mem = unsafe { ffi::llama_get_memory(ctx.0) }; + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + + let prompts = if args.prompt.is_empty() { + eprintln!("No new questions so proceed with built-in defaults."); + DEFAULT_PROMPTS + .iter() + .map(|s| s.to_string()) + .collect::>() + } else { + eprintln!("Now printing the external prompt input"); + split_prompt_lines(&args.prompt) + }; + + let n_ctx = unsafe { ffi::llama_n_ctx(ctx.0) as i32 }; + let n_clients = args.n_parallel; + let n_seq = args.n_sequences; + let mut sampler_seed = args.seed; + if sampler_seed >= 0 { + eprintln!("initializing all samplers with the same RNG seed: {sampler_seed}"); + } else { + eprintln!("initializing samplers with different RNG seeds, starting from {sampler_seed}"); + } + + let mut clients = Vec::new(); + for id in 0..n_clients { + clients.push(Client { + id, + seq_id: -1, + sampled: 0, + t_start_prompt: 0, + t_start_gen: 0, + n_past: 0, + n_prompt: 0, + n_decoded: 0, + i_batch: -1, + input: String::new(), + prompt: String::new(), + response: String::new(), + sampler: sampler_init(&args, sampler_seed), + }); + if args.seed < 0 { + sampler_seed -= 1; + } + } + + let tokens_system = tokenize(vocab, SYSTEM_PROMPT, true, false)?; + let n_tokens_system = tokens_system.len() as i32; + let mut batch = Batch(unsafe { ffi::llama_batch_init(n_ctx, 0, 1) }); + let mut n_total_prompt = 0; + let mut n_total_gen = 0; + let mut n_cache_miss = 0; + let t_main_start = unsafe { ffi::ggml_time_us() }; + + eprintln!( + "Simulating parallel requests: n_parallel = {n_clients}, n_sequences = {n_seq}, cont_batching = {}, system tokens = {n_tokens_system}", + args.cont_batching + ); + + if args.is_pp_shared { + eprintln!("Evaluating the system prompt ..."); + for (i, token) in tokens_system.iter().enumerate() { + unsafe { + batch_add(&mut batch.0, *token, i as i32, &[0], false); + } + } + if unsafe { ffi::llama_decode(ctx.0, batch.0) } != 0 { + return Err("llama_decode failed for system prompt".to_string()); + } + for i in 1..=n_clients { + unsafe { + ffi::llama_memory_seq_cp(mem, 0, i, -1, -1); + } + } + } + + let mut rng = Lcg::new(1234); + let mut g_seq_id = 0; + eprintln!("Processing requests ..."); + + loop { + batch_clear(&mut batch.0); + for client in &mut clients { + if client.seq_id == -1 { + continue; + } + client.i_batch = batch.0.n_tokens; + unsafe { + batch_add( + &mut batch.0, + client.sampled, + client.n_past, + &[client.id + 1], + true, + ); + } + client.n_past += 1; + client.n_decoded += 1; + } + + if batch.0.n_tokens == 0 { + for i in 1..=n_clients { + unsafe { + ffi::llama_memory_seq_rm(mem, i, -1, -1); + ffi::llama_memory_seq_cp(mem, 0, i, -1, -1); + } + } + eprintln!("clearing the KV cache"); + } + + if args.cont_batching || batch.0.n_tokens == 0 { + for client in &mut clients { + if client.seq_id != -1 || g_seq_id >= n_seq { + continue; + } + client.seq_id = g_seq_id; + client.t_start_prompt = unsafe { ffi::ggml_time_us() }; + client.t_start_gen = 0; + client.input = prompts[rng.next_usize(prompts.len())].clone(); + client.response.clear(); + let (prompt, n_past, n_junk_cur) = build_prompt( + args.is_pp_shared, + n_tokens_system, + &client.input, + args.n_junk, + &mut rng, + ); + client.prompt = prompt; + client.n_past = n_past; + unsafe { + ffi::llama_sampler_reset(client.sampler.0); + } + + let tokens_prompt = tokenize(vocab, &client.prompt, false, false)?; + for token in &tokens_prompt { + unsafe { + batch_add(&mut batch.0, *token, client.n_past, &[client.id + 1], false); + } + client.n_past += 1; + } + unsafe { + set_last_batch_logits(&mut batch.0); + } + client.n_prompt = tokens_prompt.len() as i32; + client.n_decoded = 0; + client.i_batch = batch.0.n_tokens - 1; + eprintln!( + "Client {:3}, seq {:4}, junk = {:4}, prompt = {}, started decoding ...", + client.id, client.seq_id, n_junk_cur, client.n_prompt + ); + g_seq_id += 1; + } + } + + if batch.0.n_tokens == 0 { + break; + } + + let mut n_batch = args.n_batch as i32; + let mut i = 0; + while i < batch.0.n_tokens { + let n_tokens = n_batch.min(batch.0.n_tokens - i); + let view = unsafe { batch_view(&batch.0, i, n_tokens) }; + let ret = unsafe { ffi::llama_decode(ctx.0, view) }; + if ret != 0 { + if n_batch == 1 || ret < 0 { + return Err(format!( + "failed to decode the batch, n_batch = {n_batch}, ret = {ret}" + )); + } + eprintln!( + "failed to decode the batch, retrying with n_batch = {}", + n_batch / 2 + ); + n_cache_miss += 1; + n_batch /= 2; + continue; + } + + let i_next = i + n_tokens; + n_batch = args.n_batch as i32; + for client in &mut clients { + if client.i_batch < i || client.i_batch >= i_next { + continue; + } + let id = unsafe { + ffi::llama_sampler_sample(client.sampler.0, ctx.0, client.i_batch - i) + }; + unsafe { + ffi::llama_sampler_accept(client.sampler.0, id); + } + if client.n_decoded == 1 { + client.t_start_gen = unsafe { ffi::ggml_time_us() }; + } + let token_str = token_to_piece(vocab, id)?; + client.response.push_str(&token_str); + client.sampled = id; + + let stop_at_user = client.response.find("User:"); + let should_stop = client.n_decoded > 2 + && (unsafe { ffi::llama_vocab_is_eog(vocab, id) } + || (args.n_predict > 0 && client.n_decoded >= args.n_predict) + || stop_at_user.is_some()); + if should_stop { + if let Some(pos) = stop_at_user { + client.response.truncate(pos); + } + unsafe { + ffi::llama_memory_seq_rm(mem, client.id + 1, -1, -1); + ffi::llama_memory_seq_cp(mem, 0, client.id + 1, -1, -1); + } + let t_main_end = unsafe { ffi::ggml_time_us() }; + let seconds = (t_main_end - client.t_start_prompt) as f64 / 1e6; + let speed = (client.n_prompt + client.n_decoded) as f64 / seconds.max(1e-9); + eprintln!( + "Client {:3}, seq {:3}/{:3}, prompt {:4} t, response {:4} t, time {:5.2} s, speed {:5.2} t/s, cache miss {}\n\nInput: {}\nResponse: {}\n", + client.id, + client.seq_id, + n_seq, + client.n_prompt, + client.n_decoded, + seconds, + speed, + n_cache_miss, + trim_ascii(&client.input), + trim_ascii(&client.response) + ); + n_total_prompt += client.n_prompt; + n_total_gen += client.n_decoded; + client.seq_id = -1; + } + client.i_batch = -1; + } + i = i_next; + } + } + + let t_main_end = unsafe { ffi::ggml_time_us() }; + let seconds = (t_main_end - t_main_start) as f64 / 1e6; + eprintln!( + "n_parallel = {n_clients}, n_sequences = {n_seq}, cont_batching = {}, system tokens = {n_tokens_system}", + args.cont_batching + ); + eprintln!( + "External prompt file: {}", + if args.prompt_file.is_empty() { + "used built-in defaults" + } else { + args.prompt_file.as_str() + } + ); + eprintln!("Model and path used: {}\n", args.model_path); + eprintln!( + "Total prompt tokens: {n_total_prompt:6}, speed: {:5.2} t/s", + n_total_prompt as f64 / seconds.max(1e-9) + ); + eprintln!( + "Total gen tokens: {n_total_gen:6}, speed: {:5.2} t/s", + n_total_gen as f64 / seconds.max(1e-9) + ); + eprintln!( + "Total speed (AVG): speed: {:5.2} t/s", + (n_total_prompt + n_total_gen) as f64 / seconds.max(1e-9) + ); + eprintln!("Cache misses: {n_cache_miss:6}\n"); + unsafe { + ffi::llama_perf_context_print(ctx.0); + } + Ok(()) +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m model.gguf -np 4 -ns 8"); + eprintln!(); +} + +fn main() { + let mut argv = env::args(); + let program = argv.next().unwrap_or_else(|| "llama-parallel".to_string()); + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_parallel_options() { + let args = parse_args([ + "-m", + "model.gguf", + "-p", + "q1\nq2", + "-n", + "16", + "-np", + "3", + "-ns", + "7", + "--no-cont-batching", + "-pps", + "--junk", + "4", + "-c", + "1024", + "-b", + "128", + "-ngl", + "0", + "--top-k", + "10", + "--top-p", + "0.8", + "--temp", + "0.5", + "-s", + "-2", + ]) + .unwrap(); + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "q1\nq2"); + assert_eq!(args.n_predict, 16); + assert_eq!(args.n_parallel, 3); + assert_eq!(args.n_sequences, 7); + assert!(!args.cont_batching); + assert!(args.is_pp_shared); + assert_eq!(args.n_junk, 4); + assert_eq!(args.n_ctx, 1024); + assert_eq!(args.n_batch, 128); + assert_eq!(args.n_gpu_layers, 0); + assert_eq!(args.top_k, 10); + assert_eq!(args.top_p, 0.8); + assert_eq!(args.temp, 0.5); + assert_eq!(args.seed, -2); + } + + #[test] + fn trims_ascii_whitespace() { + assert_eq!(trim_ascii(" \nhello\t"), "hello"); + } + + #[test] + fn splits_prompt_lines() { + assert_eq!(split_prompt_lines("a\nb\n"), vec!["a", "b"]); + } + + #[test] + fn deterministic_rng_repeats() { + let mut a = Lcg::new(1234); + let mut b = Lcg::new(1234); + assert_eq!(a.next_usize(9), b.next_usize(9)); + assert_eq!(a.next_usize(20), b.next_usize(20)); + } + + #[test] + fn builds_prompt_with_shared_system() { + let mut rng = Lcg::new(1234); + let (prompt, n_past, n_junk) = build_prompt(true, 5, "Question?", 1, &mut rng); + assert_eq!(n_past, 5); + assert_eq!(n_junk, 0); + assert_eq!(prompt, "User:\nQuestion?\nAssistant:\n"); + } +} diff --git a/examples/parallel/CMakeLists.txt b/examples/parallel/CMakeLists.txt index 4fb7a96aae3a..4be315db496f 100644 --- a/examples/parallel/CMakeLists.txt +++ b/examples/parallel/CMakeLists.txt @@ -1,5 +1,25 @@ set(TARGET llama-parallel) -add_executable(${TARGET} parallel.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(PARALLEL_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../parallel-rust/Cargo.toml) +set(PARALLEL_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/parallel) +set(PARALLEL_RUST_BIN ${PARALLEL_RUST_TARGET_DIR}/debug/llama-parallel${CMAKE_EXECUTABLE_SUFFIX}) +set(PARALLEL_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${PARALLEL_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${PARALLEL_RUST_MANIFEST} --target-dir ${PARALLEL_RUST_TARGET_DIR} --bin llama-parallel + COMMAND ${CMAKE_COMMAND} -E copy ${PARALLEL_RUST_BIN} ${PARALLEL_RUST_OUTPUT} + DEPENDS + llama + ${PARALLEL_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../parallel-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../parallel-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${PARALLEL_RUST_OUTPUT}) +install(PROGRAMS ${PARALLEL_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/parallel/parallel.cpp b/examples/parallel/parallel.cpp deleted file mode 100644 index a46400c5b943..000000000000 --- a/examples/parallel/parallel.cpp +++ /dev/null @@ -1,520 +0,0 @@ -// A basic application simulating a server with multiple clients. -// The clients submit requests to the server and they are processed in parallel. - -#include "arg.h" -#include "common.h" -#include "sampling.h" -#include "log.h" -#include "llama.h" - -#include -#include -#include -#include -#include -#include -#include - -// trim whitespace from the beginning and end of a string -static std::string trim(const std::string & str) { - size_t start = 0; - size_t end = str.size(); - - while (start < end && isspace(str[start])) { - start += 1; - } - - while (end > start && isspace(str[end - 1])) { - end -= 1; - } - - return str.substr(start, end - start); -} - -static std::string k_system = -R"(Transcript of a never ending dialog, where the User interacts with an Assistant. -The Assistant is helpful, kind, honest, good at writing, and never fails to answer the User's requests immediately and with precision. - -User: -Recommend a nice restaurant in the area. -Assistant: -I recommend the restaurant "The Golden Duck". It is a 5 star restaurant with a great view of the city. The food is delicious and the service is excellent. The prices are reasonable and the portions are generous. The restaurant is located at 123 Main Street, New York, NY 10001. The phone number is (212) 555-1234. The hours are Monday through Friday from 11:00 am to 10:00 pm. The restaurant is closed on Saturdays and Sundays. -User: -Who is Richard Feynman? -Assistant: -Richard Feynman was an American physicist who is best known for his work in quantum mechanics and particle physics. He was awarded the Nobel Prize in Physics in 1965 for his contributions to the development of quantum electrodynamics. He was a popular lecturer and author, and he wrote several books, including "Surely You're Joking, Mr. Feynman!" and "What Do You Care What Other People Think?". -)"; - -static std::vector k_questions = { - "What is the tallest mountain in the world?", - "Who was the first person to win two Nobel Prizes?", - "Which country invented paper?", - "What organ is primarily responsible for pumping blood throughout the body?", - "Which planet is known for its prominent ring system?", - "Who directed the movie 'Inception'?", - "What is the freezing point of water in Fahrenheit?", - "Which animal is known to have the longest lifespan?", - "What language has the most native speakers worldwide?", - "What is the capital city of Canada?", - "Who is credited with inventing the World Wide Web?", - "Which metal is liquid at room temperature?", - "What is the term for an animal that eats both plants and meat?", - "Who painted 'The Starry Night'?", - "What gas do humans exhale that plants use for photosynthesis?", - "What year did World War II end?", - "Which continent has the most countries?", - "Who wrote the novel 'Frankenstein'?", - "What does DNA stand for?", - "What is the main ingredient in traditional Japanese miso soup?" -}; - -static std::vector k_answers = { - "The tallest mountain in the world is Mount Everest.", - "Marie Curie was the first person to win two Nobel Prizes.", - "Paper was invented in China.", - "The heart is the organ responsible for pumping blood.", - "Saturn is known for its prominent ring system.", - "Christopher Nolan directed the movie 'Inception'.", - "The freezing point of water in Fahrenheit is 32°F.", - "The bowhead whale is known to have the longest lifespan among mammals.", - "Mandarin Chinese has the most native speakers in the world.", - "The capital city of Canada is Ottawa.", - "Tim Berners-Lee is credited with inventing the World Wide Web.", - "Mercury is the metal that is liquid at room temperature.", - "An animal that eats both plants and meat is called an omnivore.", - "'The Starry Night' was painted by Vincent van Gogh.", - "Humans exhale carbon dioxide, which plants use in photosynthesis.", - "World War II ended in 1945.", - "Africa is the continent with the most countries.", - "The novel 'Frankenstein' was written by Mary Shelley.", - "DNA stands for Deoxyribonucleic Acid.", - "The main ingredient in traditional Japanese miso soup is fermented soybean paste." -}; - -static std::vector k_prompts = { - "What is the meaning of life?", - "Tell me an interesting fact about llamas.", - "What is the best way to cook a steak?", - "Are you familiar with the Special Theory of Relativity and can you explain it to me?", - "Recommend some interesting books to read.", - "What is the best way to learn a new language?", - "How to get a job at Google?", - "If you could have any superpower, what would it be?", - "I want to learn how to play the piano. What would be the best way to do it?", -}; - -struct client { - ~client() { - if (smpl) { - common_sampler_free(smpl); - } - } - - int32_t id = 0; - - llama_seq_id seq_id = -1; - - llama_token sampled; - - int64_t t_start_prompt; - int64_t t_start_gen; - - int32_t n_past = 0; - int32_t n_prompt = 0; - int32_t n_decoded = 0; - int32_t i_batch = -1; - - std::string input; - std::string prompt; - std::string response; - - struct common_sampler * smpl = nullptr; -}; - -static void print_date_time() { - std::time_t current_time = std::time(nullptr); - std::tm* local_time = std::localtime(¤t_time); - char buffer[80]; - strftime(buffer, sizeof(buffer), "%Y-%m-%d %H:%M:%S", local_time); - - LOG_INF("\n"); - LOG_INF("\033[35mrun parameters as of %s\033[0m\n", buffer); - LOG_INF("\n"); -} - -// Define a split string function to ... -static std::vector split_string(const std::string& input, char delimiter) { - std::vector tokens; - std::istringstream stream(input); - std::string token; - while (std::getline(stream, token, delimiter)) { - tokens.push_back(token); - } - return tokens; -} - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - srand(1234); - - common_params params; - - params.n_predict = 128; - params.n_junk = 1; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_PARALLEL)) { - return 1; - } - - // number of simultaneous "clients" to simulate - const int32_t n_clients = params.n_parallel; - - // dedicate one sequence to the system prompt - params.n_parallel += 1; - - // requests to simulate - const int32_t n_seq = params.n_sequences; - - // insert new requests as soon as the previous one is done - const bool cont_batching = params.cont_batching; - - // is the system prompt shared in the cache - const bool is_sp_shared = params.is_pp_shared; - - // extra text to insert in each client's prompt in order to make it larger - const int32_t n_junk = std::max(1, params.n_junk); - - // signed seed, use negative values to indicate different seeds for the different clients - const int32_t & sseed = params.sampling.seed; - - // init llama.cpp - llama_backend_init(); - llama_numa_init(params.numa); - - // load the target model - auto llama_init = common_init_from_params(params); - - auto * model = llama_init->model(); - auto * ctx = llama_init->context(); - - auto * mem = llama_get_memory(ctx); - - const llama_vocab * vocab = llama_model_get_vocab(model); - - // load the prompts from an external file if there are any - if (params.prompt.empty()) { - LOG_INF("\033[32mNo new questions so proceed with build-in defaults.\033[0m\n"); - } else { - // Output each line of the input params.prompts vector and copy to k_prompts - int index = 0; - LOG_INF("\033[32mNow printing the external prompt file %s\033[0m\n\n", params.prompt_file.c_str()); - - std::vector prompts = split_string(params.prompt, '\n'); - for (const auto& prompt : prompts) { - k_prompts.resize(index + 1); - k_prompts[index] = prompt; - index++; - LOG_INF("%3d prompt: %s\n", index, prompt.c_str()); - } - } - - LOG_INF("\n\n"); - - const int n_ctx = llama_n_ctx(ctx); - - if (sseed >= 0) { - LOG_INF("%s: initializing all samplers with the same RNG seed: %d (use a negative seed to have different seeds)\n", __func__, sseed); - } else { - LOG_INF("%s: initializing samplers with different RNG seeds, starting from %d\n", __func__, sseed); - } - - std::vector clients(n_clients); - for (size_t i = 0; i < clients.size(); ++i) { - auto & client = clients[i]; - client.id = i; - client.smpl = common_sampler_init(model, params.sampling); - - if (sseed < 0) { - params.sampling.seed--; - } - } - - std::vector tokens_system; - - tokens_system = common_tokenize(ctx, k_system, true); - const int32_t n_tokens_system = tokens_system.size(); - - llama_seq_id g_seq_id = 0; - - // the max batch size is as large as the context to handle cases where we get very long input prompt from multiple - // users. regardless of the size, the main loop will chunk the batch into a maximum of params.n_batch tokens at a time - llama_batch batch = llama_batch_init(n_ctx, 0, 1); - - int32_t n_total_prompt = 0; - int32_t n_total_gen = 0; - int32_t n_cache_miss = 0; - - const auto t_main_start = ggml_time_us(); - - LOG_INF("%s: Simulating parallel requests from clients:\n", __func__); - LOG_INF("%s: n_parallel = %d, n_sequences = %d, cont_batching = %d, system tokens = %d\n", __func__, n_clients, n_seq, cont_batching, n_tokens_system); - LOG_INF("\n"); - - if (is_sp_shared) { - LOG_INF("%s: Evaluating the system prompt ...\n", __func__); - - for (int32_t i = 0; i < n_tokens_system; ++i) { - common_batch_add(batch, tokens_system[i], i, { 0 }, false); - } - - if (llama_decode(ctx, batch) != 0) { - LOG_ERR("%s: llama_decode() failed\n", __func__); - return 1; - } - - // assign the system KV cache to all parallel sequences - for (int32_t i = 1; i <= n_clients; ++i) { - llama_memory_seq_cp(mem, 0, i, -1, -1); - } - - LOG_INF("\n"); - } - - LOG_INF("Processing requests ...\n\n"); - - while (true) { - common_batch_clear(batch); - - // decode any currently ongoing sequences - for (auto & client : clients) { - if (client.seq_id == -1) { - continue; - } - - client.i_batch = batch.n_tokens; - - common_batch_add(batch, client.sampled, client.n_past++, { client.id + 1 }, true); - - client.n_decoded += 1; - } - - if (batch.n_tokens == 0) { - // all sequences have ended - clear the entire KV cache - for (int i = 1; i <= n_clients; ++i) { - llama_memory_seq_rm(mem, i, -1, -1); - // but keep the system prompt - llama_memory_seq_cp(mem, 0, i, -1, -1); - } - - LOG_INF("%s: clearing the KV cache\n", __func__); - } - - // insert new sequences for decoding - if (cont_batching || batch.n_tokens == 0) { - for (auto & client : clients) { - if (client.seq_id == -1 && g_seq_id < n_seq) { - client.seq_id = g_seq_id; - - client.t_start_prompt = ggml_time_us(); - client.t_start_gen = 0; - - client.input = k_prompts[rand() % k_prompts.size()]; - client.response = ""; - - // construct the prompt: - // [system prompt] + [junk] + [user prompt] - client.n_past = 0; - client.prompt = ""; - if (is_sp_shared) { - client.n_past = n_tokens_system; - } else { - client.prompt += k_system; - } - - const int n_junk_cur = rand() % n_junk; - - for (int i = 0; i < n_junk_cur; ++i) { - const int r = rand() % k_questions.size(); - client.prompt += "User:\n" + k_questions[r] + "\nAssistant:\n " + k_answers[r] + "\n"; - } - client.prompt += "User:\n" + client.input + "\nAssistant:\n"; - - common_sampler_reset(client.smpl); - - // do not prepend BOS because we have a system prompt! - std::vector tokens_prompt; - tokens_prompt = common_tokenize(ctx, client.prompt, false); - - for (size_t i = 0; i < tokens_prompt.size(); ++i) { - common_batch_add(batch, tokens_prompt[i], client.n_past++, { client.id + 1 }, false); - } - - // extract the logits only for the last token - if (batch.n_tokens > 0) { - batch.logits[batch.n_tokens - 1] = true; - } - - client.n_prompt = tokens_prompt.size(); - client.n_decoded = 0; - client.i_batch = batch.n_tokens - 1; - - LOG_INF("\033[31mClient %3d, seq %4d, junk = %4d, prompt = %d, started decoding ...\033[0m\n", client.id, client.seq_id, n_junk_cur, client.n_prompt); - - g_seq_id += 1; - - // insert new requests one-by-one - //if (cont_batching) { - // break; - //} - } - } - } - - if (batch.n_tokens == 0) { - break; - } - - // process in chunks of params.n_batch - int32_t n_batch = params.n_batch; - - int32_t i_next = 0; - - for (int32_t i = 0; i < batch.n_tokens; i = i_next) { - // experiment: process in powers of 2 - //if (i + n_batch > (int32_t) batch.n_tokens && n_batch > 32) { - // n_batch /= 2; - // i -= n_batch; - // continue; - //} - - const int32_t n_tokens = std::min(n_batch, batch.n_tokens - i); - - llama_batch batch_view = { - n_tokens, - batch.token + i, - nullptr, - batch.pos + i, - batch.n_seq_id + i, - batch.seq_id + i, - batch.logits + i, - }; - - const int ret = llama_decode(ctx, batch_view); - if (ret != 0) { - if (n_batch == 1 || ret < 0) { - // if you get here, it means the KV cache is full - try increasing it via the context size - LOG_ERR("%s : failed to decode the batch, n_batch = %d, ret = %d\n", __func__, n_batch, ret); - return 1; - } - - LOG_WRN("%s : failed to decode the batch, retrying with n_batch = %d\n", __func__, n_batch / 2); - - n_cache_miss += 1; - - // retry with half the batch size to try to find a free slot in the KV cache - n_batch /= 2; - - continue; - } - - LOG_DBG("%s : decoded batch of %d tokens\n", __func__, n_tokens); - - // move the head of the batch forward with the number of tokens we just processed - i_next = i + n_tokens; - - // on successful decode, restore the original batch size - n_batch = params.n_batch; - - for (auto & client : clients) { - if (client.i_batch < (int) i || client.i_batch >= (int) (i + n_tokens)) { - continue; - } - - //printf("client %d, seq %d, token %d, pos %d, batch %d\n", - // client.id, client.seq_id, client.sampled, client.n_decoded, client.i_batch); - - const llama_token id = common_sampler_sample(client.smpl, ctx, client.i_batch - i); - - common_sampler_accept(client.smpl, id, true); - - if (client.n_decoded == 1) { - // start measuring generation time after the first token to make sure all concurrent clients - // have their prompt already processed - client.t_start_gen = ggml_time_us(); - } - - const std::string token_str = common_token_to_piece(ctx, id); - - client.response += token_str; - client.sampled = id; - - //printf("client %d, seq %d, token %d, pos %d, batch %d: %s\n", - // client.id, client.seq_id, id, client.n_decoded, client.i_batch, token_str.c_str()); - - if (client.n_decoded > 2 && - (llama_vocab_is_eog(vocab, id) || - (params.n_predict > 0 && client.n_decoded >= params.n_predict) || - client.response.find("User:") != std::string::npos)) { - // basic reverse prompt - const size_t pos = client.response.find("User:"); - if (pos != std::string::npos) { - client.response = client.response.substr(0, pos); - } - - // delete only the generated part of the sequence, i.e. keep the system prompt in the cache - llama_memory_seq_rm(mem, client.id + 1, -1, -1); - llama_memory_seq_cp(mem, 0, client.id + 1, -1, -1); - - const auto t_main_end = ggml_time_us(); - - LOG_INF("\033[31mClient %3d, seq %3d/%3d, prompt %4d t, response %4d t, time %5.2f s, speed %5.2f t/s, cache miss %d \033[0m \n\nInput: %s\n\033[35mResponse: %s\033[0m\n\n", - client.id, client.seq_id, n_seq, client.n_prompt, client.n_decoded, - (t_main_end - client.t_start_prompt) / 1e6, - (double) (client.n_prompt + client.n_decoded) / (t_main_end - client.t_start_prompt) * 1e6, - n_cache_miss, - ::trim(client.input).c_str(), - ::trim(client.response).c_str()); - - n_total_prompt += client.n_prompt; - n_total_gen += client.n_decoded; - - client.seq_id = -1; - } - - client.i_batch = -1; - } - } - } - - const auto t_main_end = ggml_time_us(); - - print_date_time(); - - LOG_INF("%s: n_parallel = %d, n_sequences = %d, cont_batching = %d, system tokens = %d\n", __func__, n_clients, n_seq, cont_batching, n_tokens_system); - if (params.prompt_file.empty()) { - params.prompt_file = "used built-in defaults"; - } - LOG_INF("External prompt file: \033[32m%s\033[0m\n", params.prompt_file.c_str()); - LOG_INF("Model and path used: \033[32m%s\033[0m\n\n", params.model.path.c_str()); - - LOG_INF("Total prompt tokens: %6d, speed: %5.2f t/s\n", n_total_prompt, (double) (n_total_prompt ) / (t_main_end - t_main_start) * 1e6); - LOG_INF("Total gen tokens: %6d, speed: %5.2f t/s\n", n_total_gen, (double) (n_total_gen ) / (t_main_end - t_main_start) * 1e6); - LOG_INF("Total speed (AVG): %6s speed: %5.2f t/s\n", "", (double) (n_total_prompt + n_total_gen) / (t_main_end - t_main_start) * 1e6); - LOG_INF("Cache misses: %6d\n", n_cache_miss); - - LOG_INF("\n"); - - // TODO: print sampling/grammar timings for all clients - llama_perf_context_print(ctx); - - llama_batch_free(batch); - - llama_backend_free(); - - LOG("\n\n"); - - return 0; -} diff --git a/examples/passkey-rust/Cargo.toml b/examples/passkey-rust/Cargo.toml new file mode 100644 index 000000000000..b25718f10181 --- /dev/null +++ b/examples/passkey-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-passkey-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-passkey" +path = "src/main.rs" diff --git a/examples/passkey-rust/build.rs b/examples/passkey-rust/build.rs new file mode 100644 index 000000000000..348bbf824f92 --- /dev/null +++ b/examples/passkey-rust/build.rs @@ -0,0 +1,45 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); + + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-passkey"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/passkey-rust/src/main.rs b/examples/passkey-rust/src/main.rs new file mode 100644 index 000000000000..dd66a7989ed8 --- /dev/null +++ b/examples/passkey-rust/src/main.rs @@ -0,0 +1,642 @@ +use llama_simple_rust::ffi; +use std::env; +use std::ffi::CString; +use std::io::{self, Write}; +use std::ptr; + +const PROMPT_PREFIX: &str = "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there."; +const PROMPT_SUFFIX: &str = " What is the pass key? The pass key is"; +const JUNK_SENTENCE: &str = + " The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again."; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Args { + pub model_path: String, + pub n_junk: i32, + pub n_keep: i32, + pub n_grp: i32, + pub i_pos: i32, + pub n_batch: i32, + pub n_gpu_layers: i32, + pub seed: u32, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + MissingModel, + MissingValue(String), + InvalidInteger(String, String), + InvalidValue(String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + ParseError::InvalidValue(message) => write!(f, "{message}"), + } + } +} + +impl std::error::Error for ParseError {} + +impl Default for Args { + fn default() -> Self { + Self { + model_path: String::new(), + n_junk: 250, + n_keep: 32, + n_grp: 1, + i_pos: -1, + n_batch: 512, + n_gpu_layers: 99, + seed: 1, + } + } +} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut parsed = Args::default(); + let mut iter = args.into_iter().map(Into::into); + + while let Some(arg) = iter.next() { + match arg.as_str() { + "-m" | "--model" => parsed.model_path = take(&mut iter, &arg)?, + "--junk" => parsed.n_junk = parse_i32(&mut iter, &arg)?, + "--keep" => parsed.n_keep = parse_i32(&mut iter, &arg)?, + "--grp-attn-n" => parsed.n_grp = parse_i32(&mut iter, &arg)?, + "--pos" => parsed.i_pos = parse_i32(&mut iter, &arg)?, + "-b" | "--batch-size" => parsed.n_batch = parse_i32(&mut iter, &arg)?, + "-ngl" | "--gpu-layers" => parsed.n_gpu_layers = parse_i32(&mut iter, &arg)?, + "-s" | "--seed" => parsed.seed = parse_u32(&mut iter, &arg)?, + "-h" | "--help" => return Err(ParseError::MissingModel), + _ => {} + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + validate_args(&parsed)?; + Ok(parsed) +} + +fn validate_args(args: &Args) -> Result<(), ParseError> { + if args.n_junk <= 0 { + return Err(ParseError::InvalidValue( + "--junk must be positive".to_string(), + )); + } + if args.n_keep < 0 { + return Err(ParseError::InvalidValue( + "--keep must be non-negative".to_string(), + )); + } + if args.n_grp <= 0 { + return Err(ParseError::InvalidValue( + "--grp-attn-n must be positive".to_string(), + )); + } + if args.n_batch <= 0 { + return Err(ParseError::InvalidValue( + "--batch-size must be positive".to_string(), + )); + } + if args.n_batch % args.n_grp != 0 { + return Err(ParseError::InvalidValue( + "--batch-size must be divisible by --grp-attn-n".to_string(), + )); + } + if args.i_pos >= args.n_junk { + return Err(ParseError::InvalidValue( + "--pos must be less than --junk".to_string(), + )); + } + Ok(()) +} + +fn take(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + iter.next() + .ok_or_else(|| ParseError::MissingValue(flag.to_string())) +} + +fn parse_i32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_u32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn lcg_next(state: &mut u32) -> u32 { + *state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + *state +} + +pub fn build_prompt(n_junk: i32, i_pos: i32, seed: u32) -> (String, i32, i32) { + let mut rng = seed; + let pos = if i_pos < 0 { + (lcg_next(&mut rng) % n_junk as u32) as i32 + } else { + i_pos + }; + let passkey = (lcg_next(&mut rng) % 50_000 + 1) as i32; + + let mut prompt = String::from(PROMPT_PREFIX); + for i in 0..n_junk { + if i == pos { + prompt.push_str(&format!( + " The pass key is {passkey}. Remember it. {passkey} is the pass key." + )); + } + prompt.push_str(JUNK_SENTENCE); + } + prompt.push_str(PROMPT_SUFFIX); + + (prompt, passkey, pos) +} + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Batch(ffi::llama_batch); + +impl Drop for Batch { + fn drop(&mut self) { + unsafe { + ffi::llama_batch_free(self.0); + } + } +} + +struct Sampler(*mut ffi::llama_sampler); + +impl Drop for Sampler { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_sampler_free(self.0); + } + } + } +} + +struct Backend; + +impl Backend { + fn init() -> Self { + unsafe { + ffi::ggml_backend_load_all(); + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + Self + } +} + +impl Drop for Backend { + fn drop(&mut self) { + unsafe { + ffi::llama_backend_free(); + } + } +} + +fn tokenize( + vocab: *const ffi::llama_vocab, + text: &str, + add_special: bool, +) -> Result, String> { + let text_c = + CString::new(text).map_err(|_| "text contains an interior NUL byte".to_string())?; + let n_tokens = unsafe { + -ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + ptr::null_mut(), + 0, + add_special, + true, + ) + }; + if n_tokens <= 0 { + return Err("failed to size tokenization".to_string()); + } + + let mut tokens = vec![0_i32; n_tokens as usize]; + let n = unsafe { + ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + add_special, + true, + ) + }; + if n < 0 { + return Err("failed to tokenize prompt".to_string()); + } + tokens.truncate(n as usize); + Ok(tokens) +} + +fn token_to_piece( + vocab: *const ffi::llama_vocab, + token: ffi::llama_token, +) -> Result { + let mut buf = vec![0_i8; 128]; + let n = unsafe { + ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), buf.len() as i32, 0, true) + }; + if n < 0 { + return Err("failed to convert token to piece".to_string()); + } + let bytes = buf[..n as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +fn batch_clear(batch: &mut ffi::llama_batch) { + batch.n_tokens = 0; +} + +unsafe fn batch_add( + batch: &mut ffi::llama_batch, + token: ffi::llama_token, + pos: ffi::llama_pos, + seq_id: ffi::llama_seq_id, + logits: bool, +) { + let i = batch.n_tokens as isize; + *batch.token.offset(i) = token; + *batch.pos.offset(i) = pos; + *batch.n_seq_id.offset(i) = 1; + let seq_slot = *batch.seq_id.offset(i); + *seq_slot = seq_id; + *batch.logits.offset(i) = i8::from(logits); + batch.n_tokens += 1; +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!( + " {program} -m model.gguf --junk 250 --pos 90 --keep 32 --grp-attn-n 2 [--seed 1234]" + ); + eprintln!(); +} + +fn run(args: Args) -> Result<(), String> { + let _backend = Backend::init(); + let (prompt, passkey, i_pos) = build_prompt(args.n_junk, args.i_pos, args.seed); + + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = + (unsafe { ffi::llama_model_n_ctx_train(model.0) } * args.n_grp + args.n_keep).max(1) as u32; + ctx_params.n_batch = args.n_batch as u32; + ctx_params.no_perf = false; + + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create the llama_context".to_string()); + } + + let mut sampler_params = unsafe { ffi::llama_sampler_chain_default_params() }; + sampler_params.no_perf = false; + let sampler = Sampler(unsafe { ffi::llama_sampler_chain_init(sampler_params) }); + if sampler.0.is_null() { + return Err("failed to create sampler chain".to_string()); + } + unsafe { + ffi::llama_sampler_chain_add(sampler.0, ffi::llama_sampler_init_greedy()); + } + + let tokens = tokenize(vocab, &prompt, true)?; + let n_tokens_prefix = tokenize(vocab, PROMPT_PREFIX, true)?.len() as i32; + let n_tokens_all = tokens.len() as i32; + let n_predict = 16; + let n_len = n_tokens_all + n_predict; + let n_ctx = unsafe { ffi::llama_n_ctx(ctx.0) as i32 } - args.n_keep; + let n_kv_req = unsafe { ffi::llama_n_ctx(ctx.0) as i32 }; + let n_batch = args.n_batch; + let n_batch_grp = args.n_batch / args.n_grp; + + eprintln!( + "main: n_len = {n_len}, n_ctx = {n_ctx}, n_kv_req = {n_kv_req}, n_grp = {}, n_batch = {n_batch}, n_junk = {}, i_pos = {i_pos}", + args.n_grp, args.n_junk + ); + eprintln!(); + eprintln!("prefix tokens: {n_tokens_prefix}"); + eprintln!("prompt tokens: {n_tokens_all}"); + + let mut batch = Batch(unsafe { ffi::llama_batch_init(args.n_batch, 0, 1) }); + let mut n_past = 0; + let mem = unsafe { ffi::llama_get_memory(ctx.0) }; + + let mut i = 0; + while i < n_ctx { + if i > 0 && args.n_grp > 1 { + let ib = i / n_batch - 1; + let bd = n_batch_grp * (args.n_grp - 1); + unsafe { + ffi::llama_memory_seq_add(mem, 0, n_past - n_batch, n_past, ib * bd); + ffi::llama_memory_seq_div( + mem, + 0, + n_past - n_batch + ib * bd, + n_past + ib * bd, + args.n_grp, + ); + n_past = ffi::llama_memory_seq_pos_max(mem, 0) + 1; + } + } + + batch_clear(&mut batch.0); + for j in 0..n_batch { + if i + j >= n_tokens_all { + break; + } + unsafe { + batch_add(&mut batch.0, tokens[(i + j) as usize], n_past, 0, false); + } + n_past += 1; + } + + if i + n_batch >= n_tokens_all { + unsafe { + *batch.0.logits.offset(batch.0.n_tokens as isize - 1) = 1; + } + } + + if unsafe { ffi::llama_decode(ctx.0, batch.0) } != 0 { + return Err("llama_decode failed".to_string()); + } + + eprintln!( + "main: processed: [{:6}, {:6})", + i, + (i + n_batch).min(n_tokens_all) + ); + if i + n_batch >= n_tokens_all { + break; + } + i += n_batch; + } + + let mut i = n_ctx; + while i < n_tokens_all { + let n_discard = n_batch; + eprintln!("main: shifting KV cache with {n_discard}"); + unsafe { + ffi::llama_memory_seq_rm(mem, 0, args.n_keep, args.n_keep + n_discard); + ffi::llama_memory_seq_add(mem, 0, args.n_keep + n_discard, n_ctx, -n_discard); + n_past = ffi::llama_memory_seq_pos_max(mem, 0) + 1; + } + + batch_clear(&mut batch.0); + for j in 0..n_batch { + if i + j >= n_tokens_all { + break; + } + unsafe { + batch_add(&mut batch.0, tokens[(i + j) as usize], n_past, 0, false); + } + n_past += 1; + } + if i + n_batch >= n_tokens_all { + unsafe { + *batch.0.logits.offset(batch.0.n_tokens as isize - 1) = 1; + } + } + if unsafe { ffi::llama_decode(ctx.0, batch.0) } != 0 { + return Err("llama_decode failed".to_string()); + } + eprintln!( + "main: processed: [{:6}, {:6})", + i, + (i + n_batch).min(n_tokens_all) + ); + i += n_batch; + } + + let n_discard = n_past - n_ctx + n_predict; + if n_discard > 0 { + eprintln!("main: shifting KV cache with {n_discard} to free space for the answer"); + unsafe { + ffi::llama_memory_seq_rm(mem, 0, args.n_keep, args.n_keep + n_discard); + ffi::llama_memory_seq_add(mem, 0, args.n_keep + n_discard, n_ctx, -n_discard); + n_past = ffi::llama_memory_seq_pos_max(mem, 0) + 1; + } + } + + eprintln!(); + eprintln!( + "main: passkey = {passkey}, inserted at position {i_pos} / {} (token pos: ~{})", + args.n_junk, + (i_pos * n_tokens_all) / args.n_junk + ); + eprintln!(); + + let mut n_cur = n_tokens_all; + let mut n_decode = 0; + + print!("{PROMPT_SUFFIX}"); + io::stdout().flush().map_err(|err| err.to_string())?; + + let t_main_start = unsafe { ffi::ggml_time_us() }; + while n_cur <= n_len { + let new_token_id = + unsafe { ffi::llama_sampler_sample(sampler.0, ctx.0, batch.0.n_tokens - 1) }; + if unsafe { ffi::llama_vocab_is_eog(vocab, new_token_id) } || n_cur == n_len { + println!(); + break; + } + + print!("{}", token_to_piece(vocab, new_token_id)?); + io::stdout().flush().map_err(|err| err.to_string())?; + n_decode += 1; + + batch_clear(&mut batch.0); + unsafe { + batch_add(&mut batch.0, new_token_id, n_past, 0, true); + } + n_past += 1; + n_cur += 1; + + if unsafe { ffi::llama_decode(ctx.0, batch.0) } != 0 { + return Err("failed to eval".to_string()); + } + } + + println!(); + let t_main_end = unsafe { ffi::ggml_time_us() }; + let elapsed = (t_main_end - t_main_start) as f32 / 1_000_000.0; + eprintln!( + "main: decoded {n_decode} tokens in {elapsed:.2} s, speed: {:.2} t/s", + n_decode as f32 / elapsed.max(f32::EPSILON) + ); + eprintln!(); + unsafe { + ffi::llama_perf_context_print(ctx.0); + } + eprintln!(); + + Ok(()) +} + +fn main() { + let mut argv = env::args(); + let program = argv.next().unwrap_or_else(|| "llama-passkey".to_string()); + + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_defaults() { + let args = parse_args(["-m", "model.gguf"]).unwrap(); + assert_eq!(args.n_junk, 250); + assert_eq!(args.n_keep, 32); + assert_eq!(args.n_grp, 1); + assert_eq!(args.i_pos, -1); + assert_eq!(args.n_batch, 512); + } + + #[test] + fn parses_passkey_options() { + let args = parse_args([ + "-m", + "model.gguf", + "--junk", + "20", + "--pos", + "9", + "--keep", + "16", + "--grp-attn-n", + "2", + "-b", + "64", + "--seed", + "1234", + "-ngl", + "0", + ]) + .unwrap(); + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.n_junk, 20); + assert_eq!(args.i_pos, 9); + assert_eq!(args.n_keep, 16); + assert_eq!(args.n_grp, 2); + assert_eq!(args.n_batch, 64); + assert_eq!(args.seed, 1234); + assert_eq!(args.n_gpu_layers, 0); + } + + #[test] + fn validates_group_divides_batch() { + assert_eq!( + parse_args(["-m", "model.gguf", "--grp-attn-n", "3", "-b", "64"]).unwrap_err(), + ParseError::InvalidValue("--batch-size must be divisible by --grp-attn-n".to_string()) + ); + } + + #[test] + fn builds_prompt_with_inserted_passkey() { + let (prompt, passkey, pos) = build_prompt(3, 1, 7); + assert_eq!(pos, 1); + assert!(prompt.starts_with(PROMPT_PREFIX)); + assert!(prompt.ends_with(PROMPT_SUFFIX)); + assert!(prompt.contains(&format!("The pass key is {passkey}."))); + assert_eq!(prompt.matches(JUNK_SENTENCE).count(), 3); + } + + #[test] + fn seeded_random_position_is_stable() { + let (_, passkey_a, pos_a) = build_prompt(250, -1, 1234); + let (_, passkey_b, pos_b) = build_prompt(250, -1, 1234); + assert_eq!((passkey_a, pos_a), (passkey_b, pos_b)); + } +} diff --git a/examples/passkey/CMakeLists.txt b/examples/passkey/CMakeLists.txt index 12558cc25572..c71f91f9fc4f 100644 --- a/examples/passkey/CMakeLists.txt +++ b/examples/passkey/CMakeLists.txt @@ -1,5 +1,25 @@ set(TARGET llama-passkey) -add_executable(${TARGET} passkey.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(PASSKEY_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../passkey-rust/Cargo.toml) +set(PASSKEY_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/passkey) +set(PASSKEY_RUST_BIN ${PASSKEY_RUST_TARGET_DIR}/debug/llama-passkey${CMAKE_EXECUTABLE_SUFFIX}) +set(PASSKEY_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${PASSKEY_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${PASSKEY_RUST_MANIFEST} --target-dir ${PASSKEY_RUST_TARGET_DIR} --bin llama-passkey + COMMAND ${CMAKE_COMMAND} -E copy ${PASSKEY_RUST_BIN} ${PASSKEY_RUST_OUTPUT} + DEPENDS + llama + ${PASSKEY_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../passkey-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../passkey-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${PASSKEY_RUST_OUTPUT}) +install(PROGRAMS ${PASSKEY_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/passkey/passkey.cpp b/examples/passkey/passkey.cpp deleted file mode 100644 index 8440a2bf773d..000000000000 --- a/examples/passkey/passkey.cpp +++ /dev/null @@ -1,277 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "log.h" -#include "llama.h" - -#include -#include -#include -#include -#include -#include - -static void print_usage(int, char ** argv) { - LOG("\nexample usage:\n"); - LOG("\n %s -m model.gguf --junk 250 --pos 90 --keep 32 --grp-attn-n 2 [--seed 1234]\n", argv[0]); - LOG("\n"); -} - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - - params.n_junk = 250; - params.n_keep = 32; - params.i_pos = -1; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_PASSKEY, print_usage)) { - return 1; - } - - int n_junk = params.n_junk; - int n_keep = params.n_keep; - int n_grp = params.grp_attn_n; - int i_pos = params.i_pos; - - if (i_pos == -1) { - i_pos = rand() % n_junk; - } - - const std::string prompt_prefix = "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there."; - const std::string prompt_suffix = " What is the pass key? The pass key is"; - - // generate junk text - params.prompt = prompt_prefix; - - const int passkey = rand() % 50000 + 1; - - for (int i = 0; i < n_junk; i++) { - if (i % n_junk == i_pos) { - params.prompt += " The pass key is " + std::to_string(passkey) + ". Remember it. " + std::to_string(passkey) + " is the pass key."; - } - - params.prompt += " The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again."; - } - - params.prompt += prompt_suffix; - - // init LLM - - llama_backend_init(); - llama_numa_init(params.numa); - - // initialize the model - - llama_model_params model_params = common_model_params_to_llama(params); - - llama_model * model = llama_model_load_from_file(params.model.path.c_str(), model_params); - - if (model == NULL) { - LOG_ERR("%s: unable to load model\n" , __func__); - return 1; - } - - const llama_vocab * vocab = llama_model_get_vocab(model); - - // initialize the context - - llama_context_params ctx_params = common_context_params_to_llama(params); - - ctx_params.n_ctx = llama_model_n_ctx_train(model)*n_grp + n_keep; - - GGML_ASSERT(ctx_params.n_batch % n_grp == 0 && "n_batch must be divisible by n_grp"); - - llama_context * ctx = llama_init_from_model(model, ctx_params); - if (ctx == NULL) { - LOG_ERR("%s: failed to create the llama_context\n" , __func__); - return 1; - } - - auto sparams = llama_sampler_chain_default_params(); - - llama_sampler * smpl = llama_sampler_chain_init(sparams); - - llama_sampler_chain_add(smpl, llama_sampler_init_greedy()); - - // tokenize the prompt - std::vector tokens_list; - tokens_list = common_tokenize(ctx, params.prompt, true); - - // tokenize the prefix and use it as a sink - const int n_tokens_prefix = common_tokenize(ctx, prompt_prefix, true).size(); - - const int n_tokens_all = tokens_list.size(); - - // we leave a margin of 16 tokens for the generated text - it should contain just the passkey - const int n_predict = 16; - - // total length of the sequences including the prompt - const int n_len = n_tokens_all + n_predict; - - const int n_ctx = llama_n_ctx(ctx) - n_keep; - const int n_kv_req = llama_n_ctx(ctx); - const int n_batch = ctx_params.n_batch; - const int n_batch_grp = ctx_params.n_batch/n_grp; - - LOG_INF("\n%s: n_len = %d, n_ctx = %d, n_kv_req = %d, n_grp = %d, n_batch = %d, n_junk = %d, i_pos = %d\n", __func__, n_len, n_ctx, n_kv_req, n_grp, n_batch, n_junk, i_pos); - - // print the prompt token-by-token - - LOG_INF("\n"); - LOG_INF("prefix tokens: %d\n", n_tokens_prefix); - LOG_INF("prompt tokens: %d\n", n_tokens_all); - //LOG_INF("prompt: %s\n", params.prompt.c_str()); - - llama_batch batch = llama_batch_init(params.n_batch, 0, 1); - - int n_past = 0; - - auto * mem = llama_get_memory(ctx); - - // fill the KV cache - for (int i = 0; i < n_ctx; i += n_batch) { - if (i > 0 && n_grp > 1) { - // if SelfExtend is enabled, we compress the position from the last batch by a factor of n_grp - const int ib = i/n_batch - 1; - const int bd = n_batch_grp*(n_grp - 1); - - llama_memory_seq_add(mem, 0, n_past - n_batch, n_past, ib*bd); - llama_memory_seq_div(mem, 0, n_past - n_batch + ib*bd, n_past + ib*bd, n_grp); - - n_past = llama_memory_seq_pos_max(mem, 0) + 1; - } - - common_batch_clear(batch); - - for (int j = 0; j < n_batch && i + j < n_tokens_all; j++) { - common_batch_add(batch, tokens_list[i + j], n_past++, { 0 }, false); - } - - if (i + n_batch >= n_tokens_all) { - batch.logits[batch.n_tokens - 1] = true; - } - - if (llama_decode(ctx, batch) != 0) { - LOG_INF("%s: llama_decode() failed\n", __func__); - return 1; - } - - LOG_INF("%s: processed: [%6d, %6d)\n", __func__, i, std::min(i + n_batch, n_tokens_all)); - - if (i + n_batch >= n_tokens_all) { - break; - } - } - - for (int i = n_ctx; i < n_tokens_all; i += n_batch) { - const int n_discard = n_batch; - - LOG_INF("%s: shifting KV cache with %d\n", __func__, n_discard); - - llama_memory_seq_rm (mem, 0, n_keep , n_keep + n_discard); - llama_memory_seq_add(mem, 0, n_keep + n_discard, n_ctx, -n_discard); - - n_past = llama_memory_seq_pos_max(mem, 0) + 1; - - common_batch_clear(batch); - - for (int j = 0; j < n_batch && i + j < n_tokens_all; j++) { - common_batch_add(batch, tokens_list[i + j], n_past++, { 0 }, false); - } - - if (i + n_batch >= n_tokens_all) { - batch.logits[batch.n_tokens - 1] = true; - } - - if (llama_decode(ctx, batch) != 0) { - LOG_ERR("%s: llama_decode() failed\n", __func__); - return 1; - } - - LOG_INF("%s: processed: [%6d, %6d)\n", __func__, i, std::min(i + n_batch, n_tokens_all)); - } - - { - const int n_discard = n_past - n_ctx + n_predict; - - if (n_discard > 0) { - LOG_INF("%s: shifting KV cache with %d to free space for the answer\n", __func__, n_discard); - - llama_memory_seq_rm (mem, 0, n_keep , n_keep + n_discard); - llama_memory_seq_add(mem, 0, n_keep + n_discard, n_ctx, -n_discard); - - n_past = llama_memory_seq_pos_max(mem, 0) + 1; - } - } - - LOG_INF("\n"); - LOG_INF("%s: passkey = %d, inserted at position %d / %d (token pos: ~%d)\n", __func__, passkey, i_pos, n_junk, (i_pos * n_tokens_all) / n_junk); - LOG_INF("\n"); - - // main loop - - int n_cur = n_tokens_all; - int n_decode = 0; - - LOG_INF("%s", prompt_suffix.c_str()); - - const auto t_main_start = ggml_time_us(); - - while (n_cur <= n_len) { - // sample the next token - { - const llama_token new_token_id = llama_sampler_sample(smpl, ctx, batch.n_tokens - 1); - - // is it an end of generation? - if (llama_vocab_is_eog(vocab, new_token_id) || n_cur == n_len) { - LOG("\n"); - - break; - } - - LOG("%s", common_token_to_piece(ctx, new_token_id).c_str()); - - n_decode += 1; - - // prepare the next batch - common_batch_clear(batch); - - // push this new token for next evaluation - common_batch_add(batch, new_token_id, n_past++, { 0 }, true); - } - - n_cur += 1; - - // evaluate the current batch with the transformer model - if (llama_decode(ctx, batch)) { - LOG_ERR("%s : failed to eval, return code %d\n", __func__, 1); - return 1; - } - } - - LOG("\n"); - - const auto t_main_end = ggml_time_us(); - - LOG_INF("%s: decoded %d tokens in %.2f s, speed: %.2f t/s\n", - __func__, n_decode, (t_main_end - t_main_start) / 1000000.0f, n_decode / ((t_main_end - t_main_start) / 1000000.0f)); - - LOG("\n"); - llama_perf_context_print(ctx); - - LOG("\n"); - - llama_sampler_free(smpl); - - llama_batch_free(batch); - - llama_free(ctx); - llama_model_free(model); - - llama_backend_free(); - - return 0; -} diff --git a/examples/retrieval-rust/Cargo.toml b/examples/retrieval-rust/Cargo.toml new file mode 100644 index 000000000000..238f62f9d8a9 --- /dev/null +++ b/examples/retrieval-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-retrieval-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-retrieval" +path = "src/main.rs" diff --git a/examples/retrieval-rust/build.rs b/examples/retrieval-rust/build.rs new file mode 100644 index 000000000000..65659c5d7df8 --- /dev/null +++ b/examples/retrieval-rust/build.rs @@ -0,0 +1,45 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); + + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-retrieval"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/retrieval-rust/src/main.rs b/examples/retrieval-rust/src/main.rs new file mode 100644 index 000000000000..477a859e2c1b --- /dev/null +++ b/examples/retrieval-rust/src/main.rs @@ -0,0 +1,687 @@ +use llama_simple_rust::ffi; +use std::cmp::Ordering; +use std::env; +use std::ffi::CString; +use std::fs; +use std::io::{self, Write}; +use std::ptr; + +const LLAMA_POOLING_TYPE_NONE: i32 = 0; + +#[derive(Debug, Clone, PartialEq)] +pub struct Args { + pub model_path: String, + pub context_files: Vec, + pub chunk_size: usize, + pub chunk_separator: String, + pub top_k: i32, + pub n_ctx: u32, + pub n_batch: u32, + pub n_gpu_layers: i32, + pub verbose_prompt: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + MissingModel, + MissingContextFile, + MissingValue(String), + InvalidInteger(String, String), + InvalidValue(String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m/--model model.gguf"), + ParseError::MissingContextFile => write!(f, "--context-file must be specified"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + ParseError::InvalidValue(message) => write!(f, "{message}"), + } + } +} + +impl std::error::Error for ParseError {} + +impl Default for Args { + fn default() -> Self { + Self { + model_path: String::new(), + context_files: Vec::new(), + chunk_size: 64, + chunk_separator: "\n".to_string(), + top_k: 3, + n_ctx: 512, + n_batch: 512, + n_gpu_layers: 99, + verbose_prompt: false, + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct Chunk { + pub filename: String, + pub filepos: usize, + pub textdata: String, + pub tokens: Vec, + pub embedding: Vec, +} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut parsed = Args::default(); + let mut iter = args.into_iter().map(Into::into); + + while let Some(arg) = iter.next() { + match arg.as_str() { + "-m" | "--model" => parsed.model_path = take(&mut iter, &arg)?, + "--context-file" => parsed.context_files.push(take(&mut iter, &arg)?), + "--chunk-size" => parsed.chunk_size = parse_usize(&mut iter, &arg)?, + "--chunk-separator" => parsed.chunk_separator = take(&mut iter, &arg)?, + "--top-k" => parsed.top_k = parse_i32(&mut iter, &arg)?, + "-c" | "--ctx-size" => parsed.n_ctx = parse_u32(&mut iter, &arg)?, + "-b" | "--batch-size" => parsed.n_batch = parse_u32(&mut iter, &arg)?, + "-ngl" | "--gpu-layers" => parsed.n_gpu_layers = parse_i32(&mut iter, &arg)?, + "--verbose-prompt" => parsed.verbose_prompt = true, + "-h" | "--help" => return Err(ParseError::MissingModel), + _ => {} + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + if parsed.context_files.is_empty() { + return Err(ParseError::MissingContextFile); + } + if parsed.chunk_size == 0 { + return Err(ParseError::InvalidValue( + "--chunk-size must be positive".to_string(), + )); + } + if parsed.chunk_separator.is_empty() { + return Err(ParseError::InvalidValue( + "--chunk-separator must not be empty".to_string(), + )); + } + if parsed.n_batch < parsed.n_ctx { + parsed.n_batch = parsed.n_ctx; + } + + Ok(parsed) +} + +fn take(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + iter.next() + .ok_or_else(|| ParseError::MissingValue(flag.to_string())) +} + +fn parse_i32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_u32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_usize(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +pub fn chunk_text(filename: &str, text: &str, chunk_size: usize, separator: &str) -> Vec { + let mut chunks = Vec::new(); + let mut current = String::new(); + let mut filepos = 0usize; + + for piece in text.split_inclusive(separator) { + current.push_str(piece); + if current.len() > chunk_size { + filepos += push_chunk(&mut chunks, filename, filepos, &mut current); + } + } + + if !current.is_empty() { + if chunks.is_empty() { + push_chunk(&mut chunks, filename, filepos, &mut current); + } else { + chunks.last_mut().unwrap().textdata.push_str(¤t); + } + } + + chunks +} + +fn push_chunk( + chunks: &mut Vec, + filename: &str, + filepos: usize, + current: &mut String, +) -> usize { + let textdata = std::mem::take(current); + let len = textdata.len(); + chunks.push(Chunk { + filename: filename.to_string(), + filepos, + textdata, + tokens: Vec::new(), + embedding: Vec::new(), + }); + len +} + +fn chunk_file(filename: &str, chunk_size: usize, separator: &str) -> Result, String> { + let text = fs::read_to_string(filename) + .map_err(|err| format!("could not open/read file {filename}: {err}"))?; + Ok(chunk_text(filename, &text, chunk_size, separator)) +} + +pub fn normalize_l2(input: &[f32]) -> Vec { + let sum = input + .iter() + .map(|value| f64::from(*value) * f64::from(*value)) + .sum::() + .sqrt(); + let norm = if sum > 0.0 { (1.0 / sum) as f32 } else { 0.0 }; + input.iter().map(|value| value * norm).collect() +} + +pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { + let mut sum = 0.0f64; + let mut sum_a = 0.0f64; + let mut sum_b = 0.0f64; + for (x, y) in a.iter().zip(b.iter()) { + let x = f64::from(*x); + let y = f64::from(*y); + sum += x * y; + sum_a += x * x; + sum_b += y * y; + } + + if sum_a == 0.0 || sum_b == 0.0 { + return if sum_a == 0.0 && sum_b == 0.0 { + 1.0 + } else { + 0.0 + }; + } + (sum / (sum_a.sqrt() * sum_b.sqrt())) as f32 +} + +pub fn top_k_similarities(chunks: &[Chunk], query: &[f32], top_k: usize) -> Vec<(usize, f32)> { + let mut similarities = chunks + .iter() + .enumerate() + .map(|(index, chunk)| (index, cosine_similarity(&chunk.embedding, query))) + .collect::>(); + similarities.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal)); + similarities.truncate(top_k.min(similarities.len())); + similarities +} + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Batch(ffi::llama_batch); + +impl Drop for Batch { + fn drop(&mut self) { + unsafe { + ffi::llama_batch_free(self.0); + } + } +} + +struct Backend; + +impl Backend { + fn init() -> Self { + unsafe { + ffi::ggml_backend_load_all(); + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + Self + } +} + +impl Drop for Backend { + fn drop(&mut self) { + unsafe { + ffi::llama_backend_free(); + } + } +} + +fn tokenize( + vocab: *const ffi::llama_vocab, + text: &str, + add_special: bool, + parse_special: bool, +) -> Result, String> { + let text_c = + CString::new(text).map_err(|_| "text contains an interior NUL byte".to_string())?; + let n_tokens = unsafe { + -ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + ptr::null_mut(), + 0, + add_special, + parse_special, + ) + }; + if n_tokens <= 0 { + return Err("failed to size tokenization".to_string()); + } + + let mut tokens = vec![0_i32; n_tokens as usize]; + let n = unsafe { + ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + add_special, + parse_special, + ) + }; + if n < 0 { + return Err("failed to tokenize text".to_string()); + } + tokens.truncate(n as usize); + Ok(tokens) +} + +fn token_to_piece( + vocab: *const ffi::llama_vocab, + token: ffi::llama_token, +) -> Result { + let mut buf = vec![0_i8; 128]; + let n = unsafe { + ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), buf.len() as i32, 0, true) + }; + if n < 0 { + return Err("failed to convert token to piece".to_string()); + } + let bytes = buf[..n as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +fn batch_clear(batch: &mut ffi::llama_batch) { + batch.n_tokens = 0; +} + +unsafe fn batch_add_seq( + batch: &mut ffi::llama_batch, + tokens: &[ffi::llama_token], + seq_id: ffi::llama_seq_id, +) { + for (pos, token) in tokens.iter().enumerate() { + let i = batch.n_tokens as isize; + *batch.token.offset(i) = *token; + *batch.pos.offset(i) = pos as ffi::llama_pos; + *batch.n_seq_id.offset(i) = 1; + let seq_slot = *batch.seq_id.offset(i); + *seq_slot = seq_id; + *batch.logits.offset(i) = 1; + batch.n_tokens += 1; + } +} + +fn batch_process( + ctx: *mut ffi::llama_context, + batch: &mut ffi::llama_batch, + output: &mut [f32], + n_embd: usize, +) -> Result<(), String> { + unsafe { + ffi::llama_memory_clear(ffi::llama_get_memory(ctx), false); + eprintln!("batch_process: n_tokens = {}", batch.n_tokens); + if ffi::llama_decode(ctx, *batch) < 0 { + return Err("failed to process batch".to_string()); + } + + for i in 0..batch.n_tokens { + if *batch.logits.offset(i as isize) == 0 { + continue; + } + let seq_id = **batch.seq_id.offset(i as isize); + let mut emb = ffi::llama_get_embeddings_seq(ctx, seq_id); + if emb.is_null() { + emb = ffi::llama_get_embeddings_ith(ctx, i); + if emb.is_null() { + return Err(format!("failed to get embeddings for token {i}")); + } + } + + let raw = std::slice::from_raw_parts(emb, n_embd); + let normalized = normalize_l2(raw); + let start = seq_id as usize * n_embd; + output[start..start + n_embd].copy_from_slice(&normalized); + } + } + Ok(()) +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} --model ./models/bge-base-en-v1.5-f16.gguf --top-k 3 --context-file README.md --context-file License --chunk-size 100 --chunk-separator ."); + eprintln!(); +} + +fn run(args: Args) -> Result<(), String> { + let mut chunks = Vec::new(); + for context_file in &args.context_files { + eprintln!("{context_file}"); + chunks.extend(chunk_file( + context_file, + args.chunk_size, + &args.chunk_separator, + )?); + } + eprintln!("Number of chunks: {}", chunks.len()); + if chunks.is_empty() { + return Err("no chunks were produced from context files".to_string()); + } + + let _backend = Backend::init(); + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = args.n_ctx; + ctx_params.n_batch = args.n_batch; + ctx_params.n_ubatch = args.n_batch; + ctx_params.embeddings = true; + + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create the llama_context".to_string()); + } + + if unsafe { ffi::llama_pooling_type(ctx.0) } == LLAMA_POOLING_TYPE_NONE { + return Err("pooling type NONE not supported".to_string()); + } + + let n_ctx_train = unsafe { ffi::llama_model_n_ctx_train(model.0) }; + let n_ctx = unsafe { ffi::llama_n_ctx(ctx.0) as i32 }; + if n_ctx > n_ctx_train { + eprintln!( + "warning: model was trained on only {n_ctx_train} context tokens ({n_ctx} specified)" + ); + } + + let n_batch = args.n_batch as usize; + for (chunk_id, chunk) in chunks.iter_mut().enumerate() { + let mut tokens = tokenize(vocab, &chunk.textdata, true, false)?; + if tokens.len() > n_batch { + return Err(format!( + "chunk size ({}) exceeds batch size ({n_batch}), increase batch size and re-run", + tokens.len() + )); + } + let eos = unsafe { ffi::llama_vocab_eos(vocab) }; + if eos >= 0 && tokens.last().copied() != Some(eos) { + tokens.push(eos); + } + if args.verbose_prompt { + eprintln!("prompt {chunk_id}: '{}'", chunk.textdata); + eprintln!("number of tokens in prompt = {}", tokens.len()); + for token in &tokens { + eprintln!("{token:6} -> '{}'", token_to_piece(vocab, *token)?); + } + } + chunk.tokens = tokens; + } + + let n_chunks = chunks.len(); + let mut batch = Batch(unsafe { ffi::llama_batch_init(n_batch as i32, 0, 1) }); + let n_embd = unsafe { ffi::llama_model_n_embd_out(model.0) as usize }; + let mut embeddings = vec![0.0f32; n_chunks * n_embd]; + + let mut processed = 0usize; + let mut seq_count = 0usize; + for chunk_index in 0..n_chunks { + let n_toks = chunks[chunk_index].tokens.len(); + if batch.0.n_tokens as usize + n_toks > n_batch + || seq_count >= unsafe { ffi::llama_n_seq_max(ctx.0) as usize } + { + let out = &mut embeddings[processed * n_embd..(processed + seq_count) * n_embd]; + batch_process(ctx.0, &mut batch.0, out, n_embd)?; + batch_clear(&mut batch.0); + processed += seq_count; + seq_count = 0; + } + + unsafe { + batch_add_seq(&mut batch.0, &chunks[chunk_index].tokens, seq_count as i32); + } + seq_count += 1; + } + + if seq_count > 0 { + let out = &mut embeddings[processed * n_embd..(processed + seq_count) * n_embd]; + batch_process(ctx.0, &mut batch.0, out, n_embd)?; + } + + for (i, chunk) in chunks.iter_mut().enumerate() { + chunk.embedding = embeddings[i * n_embd..(i + 1) * n_embd].to_vec(); + chunk.tokens.clear(); + } + + let mut query_batch = Batch(unsafe { ffi::llama_batch_init(n_batch as i32, 0, 1) }); + let mut query = String::new(); + loop { + print!("Enter query: "); + io::stdout().flush().map_err(|err| err.to_string())?; + query.clear(); + if io::stdin() + .read_line(&mut query) + .map_err(|err| err.to_string())? + == 0 + { + break; + } + + let query_tokens = tokenize(vocab, query.trim_end(), true, true)?; + unsafe { + batch_add_seq(&mut query_batch.0, &query_tokens, 0); + } + + let mut query_emb = vec![0.0f32; n_embd]; + batch_process(ctx.0, &mut query_batch.0, &mut query_emb, n_embd)?; + batch_clear(&mut query_batch.0); + + println!("Top {} similar chunks:", args.top_k); + for (index, sim) in top_k_similarities(&chunks, &query_emb, args.top_k.max(0) as usize) { + let chunk = &chunks[index]; + println!("filename: {}", chunk.filename); + println!("filepos: {}", chunk.filepos); + println!("similarity: {sim}"); + println!("textdata:\n{}", chunk.textdata); + println!("--------------------"); + } + } + + unsafe { + ffi::llama_perf_context_print(ctx.0); + } + Ok(()) +} + +fn main() { + let mut argv = env::args(); + let program = argv.next().unwrap_or_else(|| "llama-retrieval".to_string()); + + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_retrieval_options() { + let args = parse_args([ + "--model", + "model.gguf", + "--context-file", + "README.md", + "--context-file", + "LICENSE", + "--chunk-size", + "100", + "--chunk-separator", + ".", + "--top-k", + "5", + "-c", + "256", + "-b", + "128", + "-ngl", + "0", + "--verbose-prompt", + ]) + .unwrap(); + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.context_files, vec!["README.md", "LICENSE"]); + assert_eq!(args.chunk_size, 100); + assert_eq!(args.chunk_separator, "."); + assert_eq!(args.top_k, 5); + assert_eq!(args.n_ctx, 256); + assert_eq!(args.n_batch, 256); + assert_eq!(args.n_gpu_layers, 0); + assert!(args.verbose_prompt); + } + + #[test] + fn rejects_missing_context_file() { + assert_eq!( + parse_args(["--model", "model.gguf"]).unwrap_err(), + ParseError::MissingContextFile + ); + } + + #[test] + fn chunks_text_on_separator_after_min_size() { + let chunks = chunk_text("doc.txt", "aa.bb.cc.dd.", 5, "."); + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0].textdata, "aa.bb."); + assert_eq!(chunks[0].filepos, 0); + assert_eq!(chunks[1].textdata, "cc.dd."); + assert_eq!(chunks[1].filepos, 6); + } + + #[test] + fn normalizes_and_compares_cosine() { + let normalized = normalize_l2(&[3.0, 4.0]); + assert!((normalized[0] - 0.6).abs() < 1e-6); + assert!((normalized[1] - 0.8).abs() < 1e-6); + assert!((cosine_similarity(&[1.0, 0.0], &[0.0, 1.0]) - 0.0).abs() < 1e-6); + } + + #[test] + fn ranks_top_similar_chunks() { + let chunks = vec![ + Chunk { + filename: "a".to_string(), + filepos: 0, + textdata: "a".to_string(), + tokens: Vec::new(), + embedding: vec![1.0, 0.0], + }, + Chunk { + filename: "b".to_string(), + filepos: 0, + textdata: "b".to_string(), + tokens: Vec::new(), + embedding: vec![0.0, 1.0], + }, + ]; + let ranked = top_k_similarities(&chunks, &[0.9, 0.1], 1); + assert_eq!(ranked[0].0, 0); + } +} diff --git a/examples/retrieval/CMakeLists.txt b/examples/retrieval/CMakeLists.txt index 5927ff8a852d..d7b7b6c17ca0 100644 --- a/examples/retrieval/CMakeLists.txt +++ b/examples/retrieval/CMakeLists.txt @@ -1,5 +1,25 @@ set(TARGET llama-retrieval) -add_executable(${TARGET} retrieval.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(RETRIEVAL_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../retrieval-rust/Cargo.toml) +set(RETRIEVAL_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/retrieval) +set(RETRIEVAL_RUST_BIN ${RETRIEVAL_RUST_TARGET_DIR}/debug/llama-retrieval${CMAKE_EXECUTABLE_SUFFIX}) +set(RETRIEVAL_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${RETRIEVAL_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${RETRIEVAL_RUST_MANIFEST} --target-dir ${RETRIEVAL_RUST_TARGET_DIR} --bin llama-retrieval + COMMAND ${CMAKE_COMMAND} -E copy ${RETRIEVAL_RUST_BIN} ${RETRIEVAL_RUST_OUTPUT} + DEPENDS + llama + ${RETRIEVAL_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../retrieval-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../retrieval-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${RETRIEVAL_RUST_OUTPUT}) +install(PROGRAMS ${RETRIEVAL_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/retrieval/retrieval.cpp b/examples/retrieval/retrieval.cpp deleted file mode 100644 index 7d93ab1172c6..000000000000 --- a/examples/retrieval/retrieval.cpp +++ /dev/null @@ -1,307 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "log.h" -#include "llama.h" - -#include -#include -#include -#include // TODO: remove me - -static void print_usage(int, char ** argv) { - LOG("\nexample usage:\n"); - LOG("\n %s --model ./models/bge-base-en-v1.5-f16.gguf --top-k 3 --context-file README.md --context-file License --chunk-size 100 --chunk-separator .\n", argv[0]); - LOG("\n"); -} - -struct chunk { - // filename - std::string filename; - // original file position - size_t filepos; - // original text data - std::string textdata; - // tokenized text data - std::vector tokens; - // embedding - std::vector embedding; -}; - -// chunk file data to chunks of size >= chunk_size -// chunk_separator is the separator between chunks -static std::vector chunk_file(const std::string & filename, int chunk_size, const std::string & chunk_separator) { - std::vector chunks; - std::ifstream f(filename.c_str()); - - if (!f.is_open()) { - LOG_ERR("could not open file %s\n", filename.c_str()); - return chunks; - } - - chunk current_chunk; - char buffer[1024]; - int64_t filepos = 0; - std::string current; - while (f.read(buffer, 1024)) { - current += std::string(buffer, f.gcount()); - size_t pos; - while ((pos = current.find(chunk_separator)) != std::string::npos) { - current_chunk.textdata += current.substr(0, pos + chunk_separator.size()); - if ((int) current_chunk.textdata.size() > chunk_size) { - // save chunk - current_chunk.filepos = filepos; - current_chunk.filename = filename; - chunks.push_back(current_chunk); - // update filepos - filepos += (int) current_chunk.textdata.size(); - // reset current_chunk - current_chunk = chunk(); - } - current = current.substr(pos + chunk_separator.size()); - } - - } - // add leftover data to last chunk - if (current_chunk.textdata.size() > 0) { - if (chunks.empty()) { - current_chunk.filepos = filepos; - current_chunk.filename = filename; - chunks.push_back(current_chunk); - } else { - chunks.back().textdata += current_chunk.textdata; - } - } - f.close(); - return chunks; -} - -static void batch_add_seq(llama_batch & batch, const std::vector & tokens, llama_seq_id seq_id) { - size_t n_tokens = tokens.size(); - for (size_t i = 0; i < n_tokens; i++) { - common_batch_add(batch, tokens[i], i, { seq_id }, true); - } -} - -static void batch_process(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd) { - // clear previous kv_cache values (irrelevant for embeddings) - llama_memory_clear(llama_get_memory(ctx), false); - - // run model - LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq); - if (llama_decode(ctx, batch) < 0) { - LOG_ERR("%s : failed to process\n", __func__); - } - - for (int i = 0; i < batch.n_tokens; i++) { - if (!batch.logits[i]) { - continue; - } - - // try to get sequence embeddings - supported only when pooling_type is not NONE - const float * embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]); - if (embd == NULL) { - embd = llama_get_embeddings_ith(ctx, i); - if (embd == NULL) { - LOG_ERR("%s: failed to get embeddings for token %d\n", __func__, i); - continue; - } - } - - float * out = output + batch.seq_id[i][0] * n_embd; - common_embd_normalize(embd, out, n_embd, 2); - } -} - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_RETRIEVAL, print_usage)) { - return 1; - } - - // For BERT models, batch size must be equal to ubatch size - params.n_ubatch = params.n_batch; - params.embedding = true; - - if (params.chunk_size <= 0) { - LOG_ERR("chunk_size must be positive\n"); - return 1; - } - if (params.context_files.empty()) { - LOG_ERR("context_files must be specified\n"); - return 1; - } - - LOG_INF("processing files:\n"); - for (auto & context_file : params.context_files) { - LOG_INF("%s\n", context_file.c_str()); - } - - std::vector chunks; - for (auto & context_file : params.context_files) { - std::vector file_chunk = chunk_file(context_file, params.chunk_size, params.chunk_separator); - chunks.insert(chunks.end(), file_chunk.begin(), file_chunk.end()); - } - LOG_INF("Number of chunks: %zu\n", chunks.size()); - - llama_backend_init(); - llama_numa_init(params.numa); - - // load the model - auto llama_init = common_init_from_params(params); - - auto * model = llama_init->model(); - auto * ctx = llama_init->context(); - - if (model == NULL) { - LOG_ERR("%s: unable to load model\n", __func__); - return 1; - } - - const llama_vocab * vocab = llama_model_get_vocab(model); - - const int n_ctx_train = llama_model_n_ctx_train(model); - const int n_ctx = llama_n_ctx(ctx); - - const enum llama_pooling_type pooling_type = llama_pooling_type(ctx); - if (pooling_type == LLAMA_POOLING_TYPE_NONE) { - LOG_ERR("%s: pooling type NONE not supported\n", __func__); - return 1; - } - - if (n_ctx > n_ctx_train) { - LOG_WRN("%s: warning: model was trained on only %d context tokens (%d specified)\n", - __func__, n_ctx_train, n_ctx); - } - - // print system information - { - LOG_INF("\n"); - LOG_INF("%s\n", common_params_get_system_info(params).c_str()); - } - - // max batch size - const uint64_t n_batch = params.n_batch; - GGML_ASSERT(params.n_batch >= params.n_ctx); - - // tokenize the prompts and trim - for (auto & chunk : chunks) { - auto inp = common_tokenize(ctx, chunk.textdata, true, false); - if (inp.size() > n_batch) { - LOG_ERR("%s: chunk size (%lld) exceeds batch size (%lld), increase batch size and re-run\n", - __func__, (long long int) inp.size(), (long long int) n_batch); - return 1; - } - // add eos if not present - if (llama_vocab_eos(vocab) >= 0 && (inp.empty() || inp.back() != llama_vocab_eos(vocab))) { - inp.push_back(llama_vocab_eos(vocab)); - } - chunk.tokens = inp; - } - - // tokenization stats - if (params.verbose_prompt) { - for (int i = 0; i < (int) chunks.size(); i++) { - LOG_INF("%s: prompt %d: '%s'\n", __func__, i, chunks[i].textdata.c_str()); - LOG_INF("%s: number of tokens in prompt = %zu\n", __func__, chunks[i].tokens.size()); - for (int j = 0; j < (int) chunks[i].tokens.size(); j++) { - LOG_INF("%6d -> '%s'\n", chunks[i].tokens[j], common_token_to_piece(ctx, chunks[i].tokens[j]).c_str()); - } - LOG_INF("\n\n"); - } - } - - // initialize batch - const int n_chunks = chunks.size(); - struct llama_batch batch = llama_batch_init(n_batch, 0, 1); - - // allocate output - const int n_embd_out = llama_model_n_embd_out(model); - std::vector embeddings(n_chunks * n_embd_out, 0); - float * emb = embeddings.data(); - - // break into batches - unsigned int p = 0; // number of prompts processed already - unsigned int s = 0; // number of prompts in current batch - for (int k = 0; k < n_chunks; k++) { - // clamp to n_batch tokens - auto & inp = chunks[k].tokens; - - const uint64_t n_toks = inp.size(); - - // encode if at capacity - if (batch.n_tokens + n_toks > n_batch || s >= llama_n_seq_max(ctx)) { - float * out = emb + p * n_embd_out; - batch_process(ctx, batch, out, s, n_embd_out); - common_batch_clear(batch); - p += s; - s = 0; - } - - // add to batch - batch_add_seq(batch, inp, s); - s += 1; - } - - // final batch - float * out = emb + p * n_embd_out; - batch_process(ctx, batch, out, s, n_embd_out); - - // save embeddings to chunks - for (int i = 0; i < n_chunks; i++) { - chunks[i].embedding = std::vector(emb + i * n_embd_out, emb + (i + 1) * n_embd_out); - // clear tokens as they are no longer needed - chunks[i].tokens.clear(); - } - - struct llama_batch query_batch = llama_batch_init(n_batch, 0, 1); - - // start loop, receive query and return top k similar chunks based on cosine similarity - std::string query; - while (true) { - LOG("Enter query: "); - std::getline(std::cin, query); - std::vector query_tokens = common_tokenize(ctx, query, true); - - batch_add_seq(query_batch, query_tokens, 0); - - std::vector query_emb(n_embd_out, 0); - batch_process(ctx, query_batch, query_emb.data(), 1, n_embd_out); - - common_batch_clear(query_batch); - - // compute cosine similarities - { - std::vector> similarities; - for (int i = 0; i < n_chunks; i++) { - float sim = common_embd_similarity_cos(chunks[i].embedding.data(), query_emb.data(), n_embd_out); - similarities.push_back(std::make_pair(i, sim)); - } - - // sort similarities - std::sort(similarities.begin(), similarities.end(), [](const std::pair & a, const std::pair & b) { - return a.second > b.second; - }); - - LOG("Top %d similar chunks:\n", params.sampling.top_k); - for (int i = 0; i < std::min(params.sampling.top_k, (int) chunks.size()); i++) { - LOG("filename: %s\n", chunks[similarities[i].first].filename.c_str()); - LOG("filepos: %lld\n", (long long int) chunks[similarities[i].first].filepos); - LOG("similarity: %f\n", similarities[i].second); - LOG("textdata:\n%s\n", chunks[similarities[i].first].textdata.c_str()); - LOG("--------------------\n"); - } - } - } - - LOG("\n"); - llama_perf_context_print(ctx); - - // clean up - llama_batch_free(query_batch); - llama_backend_free(); -} diff --git a/examples/save-load-state-rust/Cargo.toml b/examples/save-load-state-rust/Cargo.toml new file mode 100644 index 000000000000..6e5fc78f31c3 --- /dev/null +++ b/examples/save-load-state-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-save-load-state-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-save-load-state" +path = "src/main.rs" diff --git a/examples/save-load-state-rust/build.rs b/examples/save-load-state-rust/build.rs new file mode 100644 index 000000000000..cf8b738526bd --- /dev/null +++ b/examples/save-load-state-rust/build.rs @@ -0,0 +1,45 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); + + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-save-load-state"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/save-load-state-rust/src/main.rs b/examples/save-load-state-rust/src/main.rs new file mode 100644 index 000000000000..7fb354190277 --- /dev/null +++ b/examples/save-load-state-rust/src/main.rs @@ -0,0 +1,630 @@ +use llama_simple_rust::ffi; +use std::env; +use std::ffi::CString; +use std::io::{self, Write}; +use std::ptr; + +const STATE_FILE: &str = "dump_state.bin"; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub n_predict: i32, + pub n_batch: i32, + pub n_parallel: i32, + pub n_gpu_layers: i32, + pub seed: u32, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + MissingModel, + MissingValue(String), + InvalidInteger(String, String), + InvalidValue(String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + ParseError::InvalidValue(message) => write!(f, "{message}"), + } + } +} + +impl std::error::Error for ParseError {} + +impl Default for Args { + fn default() -> Self { + Self { + model_path: String::new(), + prompt: "The quick brown fox".to_string(), + n_predict: 16, + n_batch: 512, + n_parallel: 1, + n_gpu_layers: 99, + seed: 1234, + } + } +} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut parsed = Args::default(); + let mut iter = args.into_iter().map(Into::into); + + while let Some(arg) = iter.next() { + match arg.as_str() { + "-m" | "--model" => parsed.model_path = take(&mut iter, &arg)?, + "-p" | "--prompt" => parsed.prompt = take(&mut iter, &arg)?, + "-n" | "--n-predict" => parsed.n_predict = parse_i32(&mut iter, &arg)?, + "-b" | "--batch-size" => parsed.n_batch = parse_i32(&mut iter, &arg)?, + "-np" | "--parallel" => parsed.n_parallel = parse_i32(&mut iter, &arg)?, + "-ngl" | "--gpu-layers" => parsed.n_gpu_layers = parse_i32(&mut iter, &arg)?, + "-s" | "--seed" => parsed.seed = parse_u32(&mut iter, &arg)?, + "-h" | "--help" => return Err(ParseError::MissingModel), + other => { + let mut prompt = vec![other.to_string()]; + prompt.extend(iter); + parsed.prompt = prompt.join(" "); + break; + } + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + if parsed.n_predict < 0 { + parsed.n_predict = 16; + } + if parsed.n_predict == 0 { + return Err(ParseError::InvalidValue( + "--n-predict must be non-zero".to_string(), + )); + } + if parsed.n_batch <= 0 { + return Err(ParseError::InvalidValue( + "--batch-size must be positive".to_string(), + )); + } + if parsed.n_parallel <= 0 { + return Err(ParseError::InvalidValue( + "--parallel must be positive".to_string(), + )); + } + + Ok(parsed) +} + +fn take(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + iter.next() + .ok_or_else(|| ParseError::MissingValue(flag.to_string())) +} + +fn parse_i32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_u32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +pub fn should_enable_unified_kv(n_parallel: i32) -> bool { + n_parallel == 1 +} + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Batch(ffi::llama_batch); + +impl Drop for Batch { + fn drop(&mut self) { + unsafe { + ffi::llama_batch_free(self.0); + } + } +} + +struct Sampler(*mut ffi::llama_sampler); + +impl Drop for Sampler { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_sampler_free(self.0); + } + } + } +} + +struct Backend; + +impl Backend { + fn init() -> Self { + unsafe { + ffi::ggml_backend_load_all(); + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + Self + } +} + +impl Drop for Backend { + fn drop(&mut self) { + unsafe { + ffi::llama_backend_free(); + } + } +} + +fn tokenize(vocab: *const ffi::llama_vocab, text: &str) -> Result, String> { + let text_c = + CString::new(text).map_err(|_| "prompt contains an interior NUL byte".to_string())?; + let n_tokens = unsafe { + -ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + ptr::null_mut(), + 0, + true, + true, + ) + }; + if n_tokens <= 0 { + return Err("failed to size prompt tokenization".to_string()); + } + + let mut tokens = vec![0_i32; n_tokens as usize]; + let n = unsafe { + ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + true, + true, + ) + }; + if n < 0 { + return Err("failed to tokenize prompt".to_string()); + } + tokens.truncate(n as usize); + Ok(tokens) +} + +fn token_to_piece( + vocab: *const ffi::llama_vocab, + token: ffi::llama_token, +) -> Result { + let mut buf = vec![0_i8; 128]; + let n = unsafe { + ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), buf.len() as i32, 0, true) + }; + if n < 0 { + return Err("failed to convert token to piece".to_string()); + } + let bytes = buf[..n as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +fn batch_clear(batch: &mut ffi::llama_batch) { + batch.n_tokens = 0; +} + +unsafe fn batch_add( + batch: &mut ffi::llama_batch, + token: ffi::llama_token, + pos: ffi::llama_pos, + seq_id: ffi::llama_seq_id, + logits: bool, +) { + let i = batch.n_tokens as isize; + *batch.token.offset(i) = token; + *batch.pos.offset(i) = pos; + *batch.n_seq_id.offset(i) = 1; + let seq_slot = *batch.seq_id.offset(i); + *seq_slot = seq_id; + *batch.logits.offset(i) = i8::from(logits); + batch.n_tokens += 1; +} + +fn replay_last_token( + ctx: *mut ffi::llama_context, + last_token: ffi::llama_token, + pos: ffi::llama_pos, +) -> Result<(), String> { + let mut token = last_token; + let mut batch = unsafe { ffi::llama_batch_get_one(&mut token, 1) }; + batch.pos = &pos as *const i32 as *mut i32; + if unsafe { ffi::llama_decode(ctx, batch) } != 0 { + return Err("failed to replay last token".to_string()); + } + Ok(()) +} + +fn prompt_batch_decode( + ctx: *mut ffi::llama_context, + tokens: &mut [ffi::llama_token], + n_past: &mut i32, + n_batch: i32, + state_file: &CString, +) -> Result<(), String> { + if tokens.is_empty() { + return Ok(()); + } + if tokens.len() > 1 { + let n_before_last = tokens.len() - 1; + if tokens.len() as i32 > n_batch { + return Err("prompt token count exceeds batch size".to_string()); + } + + let batch = unsafe { ffi::llama_batch_get_one(tokens.as_mut_ptr(), n_before_last as i32) }; + if unsafe { ffi::llama_decode(ctx, batch) } != 0 { + return Err("failed to evaluate prompt before last token".to_string()); + } + *n_past += n_before_last as i32; + + if !unsafe { + ffi::llama_state_save_file(ctx, state_file.as_ptr(), tokens.as_ptr(), n_before_last) + } { + return Err("failed to save state file".to_string()); + } + eprintln!( + "saved session before last token to {}, n_tokens = {n_before_last}", + state_file.to_string_lossy() + ); + + replay_last_token(ctx, *tokens.last().unwrap(), *n_past)?; + *n_past += 1; + } else { + let batch = unsafe { ffi::llama_batch_get_one(tokens.as_mut_ptr(), tokens.len() as i32) }; + if unsafe { ffi::llama_decode(ctx, batch) } != 0 { + return Err("failed to evaluate prompt".to_string()); + } + *n_past += tokens.len() as i32; + } + Ok(()) +} + +fn make_sampler(seed: u32) -> Result { + let params = unsafe { ffi::llama_sampler_chain_default_params() }; + let sampler = Sampler(unsafe { ffi::llama_sampler_chain_init(params) }); + if sampler.0.is_null() { + return Err("failed to create sampler".to_string()); + } + unsafe { + ffi::llama_sampler_chain_add(sampler.0, ffi::llama_sampler_init_dist(seed)); + } + Ok(sampler) +} + +fn make_context( + model: *mut ffi::llama_model, + args: &Args, + n_seq_max: u32, +) -> Result { + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_batch = args.n_batch as u32; + ctx_params.n_ubatch = args.n_batch as u32; + ctx_params.n_seq_max = n_seq_max; + ctx_params.kv_unified = should_enable_unified_kv(args.n_parallel); + ctx_params.no_perf = false; + + let ctx = Context(unsafe { ffi::llama_init_from_model(model, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create context".to_string()); + } + Ok(ctx) +} + +fn generate( + ctx: *mut ffi::llama_context, + vocab: *const ffi::llama_vocab, + sampler: *mut ffi::llama_sampler, + batch: &mut ffi::llama_batch, + n_past: &mut i32, + seq_id: ffi::llama_seq_id, + n_predict: i32, +) -> Result { + let mut result = String::new(); + for _ in 0..n_predict { + let next_token = unsafe { ffi::llama_sampler_sample(sampler, ctx, -1) }; + let piece = token_to_piece(vocab, next_token)?; + print!("{piece}"); + io::stdout().flush().map_err(|err| err.to_string())?; + result.push_str(&piece); + + batch_clear(batch); + unsafe { + batch_add(batch, next_token, *n_past, seq_id, true); + } + if unsafe { ffi::llama_decode(ctx, *batch) } != 0 { + return Err("failed to evaluate generated token".to_string()); + } + *n_past += 1; + } + Ok(result) +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m model.gguf -p \"The quick brown fox\" -n 16 --seed 1234"); + eprintln!(); +} + +fn run(args: Args) -> Result<(), String> { + if should_enable_unified_kv(args.n_parallel) { + println!("main: n_parallel == 1, enabling unified kv cache"); + } + + let _backend = Backend::init(); + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + let state_file = CString::new(STATE_FILE).unwrap(); + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get vocabulary".to_string()); + } + + let ctx = make_context(model.0, &args, args.n_parallel as u32)?; + let sampler = make_sampler(args.seed)?; + let mut tokens = tokenize(vocab, &args.prompt)?; + let mut n_past = 0; + + prompt_batch_decode(ctx.0, &mut tokens, &mut n_past, args.n_batch, &state_file)?; + + println!("\nfirst run: {}", args.prompt); + let mut batch = Batch(unsafe { ffi::llama_batch_init(1, 0, 1) }); + let result0 = generate( + ctx.0, + vocab, + sampler.0, + &mut batch.0, + &mut n_past, + 0, + args.n_predict, + )?; + println!("\n"); + + let ctx2 = make_context(model.0, &args, args.n_parallel as u32)?; + let sampler2 = make_sampler(args.seed)?; + println!("\nsecond run: {}", args.prompt); + + let mut unused_tokens = vec![0_i32; tokens.len()]; + let mut n_token_count_out = 0usize; + if !unsafe { + ffi::llama_state_load_file( + ctx2.0, + state_file.as_ptr(), + unused_tokens.as_mut_ptr(), + unused_tokens.len(), + &mut n_token_count_out, + ) + } { + return Err("failed to load state".to_string()); + } + eprintln!("main: loaded state with {n_token_count_out} tokens"); + + let mut n_past2 = n_token_count_out as i32; + replay_last_token(ctx2.0, *tokens.last().unwrap(), n_past2)?; + n_past2 += 1; + let result1 = generate( + ctx2.0, + vocab, + sampler2.0, + &mut batch.0, + &mut n_past2, + 0, + args.n_predict, + )?; + println!("\n"); + + if result0 != result1 { + return Err("the 2 generations are different".to_string()); + } + + let ctx3 = make_context(model.0, &args, 2)?; + let sampler3 = make_sampler(args.seed)?; + println!("\nsingle seq run: {}", args.prompt); + + n_token_count_out = 0; + if !unsafe { + ffi::llama_state_load_file( + ctx3.0, + state_file.as_ptr(), + unused_tokens.as_mut_ptr(), + unused_tokens.len(), + &mut n_token_count_out, + ) + } { + return Err("failed to load state into third context".to_string()); + } + eprintln!("main: loaded state with {n_token_count_out} tokens"); + + let mut n_past3 = n_token_count_out as i32; + replay_last_token(ctx3.0, *tokens.last().unwrap(), n_past3)?; + n_past3 += 1; + + let seq_size = unsafe { ffi::llama_state_seq_get_size(ctx3.0, 0) }; + let mut seq_store = vec![0u8; seq_size]; + let ncopy = + unsafe { ffi::llama_state_seq_get_data(ctx3.0, seq_store.as_mut_ptr(), seq_size, 0) }; + if ncopy != seq_store.len() { + return Err(format!( + "seq copy data length {ncopy} does not match expected length {}", + seq_store.len() + )); + } + eprintln!("main: seq 0 copied, {ncopy} bytes"); + + unsafe { + ffi::llama_memory_clear(ffi::llama_get_memory(ctx3.0), true); + } + eprintln!("main: kv cache cleared"); + + let nset = + unsafe { ffi::llama_state_seq_set_data(ctx3.0, seq_store.as_ptr(), seq_store.len(), 1) }; + if nset != seq_store.len() { + return Err(format!( + "seq set data length {nset} does not match expected length {}", + seq_store.len() + )); + } + eprintln!("main: seq 1 restored, {nset} bytes"); + + let result2 = generate( + ctx3.0, + vocab, + sampler3.0, + &mut batch.0, + &mut n_past3, + 1, + args.n_predict, + )?; + println!(); + + if result0 != result2 { + return Err("the seq restore generation is different".to_string()); + } + + eprintln!("\nmain: success"); + Ok(()) +} + +fn main() { + let mut argv = env::args(); + let program = argv + .next() + .unwrap_or_else(|| "llama-save-load-state".to_string()); + + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_defaults() { + let args = parse_args(["-m", "model.gguf"]).unwrap(); + assert_eq!(args.prompt, "The quick brown fox"); + assert_eq!(args.seed, 1234); + assert_eq!(args.n_predict, 16); + assert_eq!(args.n_parallel, 1); + } + + #[test] + fn parses_overrides_and_positional_prompt() { + let args = parse_args([ + "-m", + "model.gguf", + "-n", + "4", + "-b", + "64", + "-np", + "2", + "-ngl", + "0", + "--seed", + "42", + "custom", + "prompt", + ]) + .unwrap(); + assert_eq!(args.prompt, "custom prompt"); + assert_eq!(args.n_predict, 4); + assert_eq!(args.n_batch, 64); + assert_eq!(args.n_parallel, 2); + assert_eq!(args.n_gpu_layers, 0); + assert_eq!(args.seed, 42); + } + + #[test] + fn negative_predict_uses_cpp_default() { + let args = parse_args(["-m", "model.gguf", "-n", "-1"]).unwrap(); + assert_eq!(args.n_predict, 16); + } + + #[test] + fn unified_kv_matches_cpp_condition() { + assert!(should_enable_unified_kv(1)); + assert!(!should_enable_unified_kv(2)); + } +} diff --git a/examples/save-load-state/CMakeLists.txt b/examples/save-load-state/CMakeLists.txt index 78024672e77d..5ac8785ed9ab 100644 --- a/examples/save-load-state/CMakeLists.txt +++ b/examples/save-load-state/CMakeLists.txt @@ -1,5 +1,25 @@ set(TARGET llama-save-load-state) -add_executable(${TARGET} save-load-state.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(SAVE_LOAD_STATE_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../save-load-state-rust/Cargo.toml) +set(SAVE_LOAD_STATE_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/save-load-state) +set(SAVE_LOAD_STATE_RUST_BIN ${SAVE_LOAD_STATE_RUST_TARGET_DIR}/debug/llama-save-load-state${CMAKE_EXECUTABLE_SUFFIX}) +set(SAVE_LOAD_STATE_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${SAVE_LOAD_STATE_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${SAVE_LOAD_STATE_RUST_MANIFEST} --target-dir ${SAVE_LOAD_STATE_RUST_TARGET_DIR} --bin llama-save-load-state + COMMAND ${CMAKE_COMMAND} -E copy ${SAVE_LOAD_STATE_RUST_BIN} ${SAVE_LOAD_STATE_RUST_OUTPUT} + DEPENDS + llama + ${SAVE_LOAD_STATE_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../save-load-state-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../save-load-state-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${SAVE_LOAD_STATE_RUST_OUTPUT}) +install(PROGRAMS ${SAVE_LOAD_STATE_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/save-load-state/save-load-state.cpp b/examples/save-load-state/save-load-state.cpp deleted file mode 100644 index a26fd73cb3c4..000000000000 --- a/examples/save-load-state/save-load-state.cpp +++ /dev/null @@ -1,238 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "llama.h" - -#include -#include -#include - - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - - params.prompt = "The quick brown fox"; - params.sampling.seed = 1234; - - const std::string_view state_file = "dump_state.bin"; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { - return 1; - } - - if (params.n_parallel == 1) { - // the example uses 2 sequences, so when n_parallel == 1, we need to enable unified kv cache - printf("%s: n_parallel == 1, enabling unified kv cache\n", __func__); - params.kv_unified = true; - } - - if (params.n_predict < 0) { - params.n_predict = 16; - } - - auto n_past = 0; - - std::string result0; - std::string result1; - std::string result2; - - // init - auto llama_init = common_init_from_params(params); - - auto * model = llama_init->model(); - auto * ctx = llama_init->context(); - - if (model == nullptr || ctx == nullptr) { - fprintf(stderr, "%s : failed to init\n", __func__); - return 1; - } - - auto sparams = llama_sampler_chain_default_params(); - - llama_sampler * smpl = llama_sampler_chain_init(sparams); - - llama_sampler_chain_add(smpl, llama_sampler_init_dist(params.sampling.seed)); - - // tokenize prompt - auto tokens = common_tokenize(ctx, params.prompt, true); - - const bool save_state = true; - if (!common_prompt_batch_decode(ctx, tokens, n_past, params.n_batch, state_file, save_state)) { - return 1; - } - - // first run - printf("\nfirst run: %s", params.prompt.c_str()); - - llama_batch batch = llama_batch_init(1, 0, 1); - - for (auto i = 0; i < params.n_predict; i++) { - auto next_token = llama_sampler_sample(smpl, ctx, -1); - auto next_token_str = common_token_to_piece(ctx, next_token); - - printf("%s", next_token_str.c_str()); - result0 += next_token_str; - - common_batch_clear(batch); - common_batch_add(batch, next_token, n_past, {0}, true); - - if (llama_decode(ctx, batch)) { - fprintf(stderr, "\n%s : failed to evaluate\n", __func__); - llama_batch_free(batch); - return 1; - } - n_past += 1; - } - - printf("\n\n"); - - // make new context - llama_context * ctx2 = llama_init_from_model(model, common_context_params_to_llama(params)); - - llama_sampler * smpl2 = llama_sampler_chain_init(sparams); - - llama_sampler_chain_add(smpl2, llama_sampler_init_dist(params.sampling.seed)); - - printf("\nsecond run: %s", params.prompt.c_str()); - - // load state from file - std::vector unused_sts(tokens.size()); // unused session tokens. - size_t n_token_count_out = 0; - - if (!llama_state_load_file(ctx2, state_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) { - fprintf(stderr, "\n%s : failed to load state\n", __func__); - return 1; - } - - fprintf(stderr, "%s : loaded state with %zu tokens\n", __func__, n_token_count_out); - - // restore state (last tokens) - n_past = n_token_count_out; - if (!common_replay_last_token(ctx2, tokens.back(), n_past)) { - return 1; - } - ++n_past; - - // second run - for (auto i = 0; i < params.n_predict; i++) { - auto next_token = llama_sampler_sample(smpl2, ctx2, -1); - auto next_token_str = common_token_to_piece(ctx2, next_token); - - printf("%s", next_token_str.c_str()); - result1 += next_token_str; - - common_batch_clear(batch); - common_batch_add(batch, next_token, n_past, {0}, true); - - if (llama_decode(ctx2, batch)) { - fprintf(stderr, "\n%s : failed to evaluate\n", __func__); - llama_batch_free(batch); - return 1; - } - n_past += 1; - } - - printf("\n\n"); - - if (result0 != result1) { - fprintf(stderr, "\n%s : error : the 2 generations are different\n", __func__); - return 1; - } - - // make new context - auto params_ctx3 = common_context_params_to_llama(params); - params_ctx3.n_seq_max = 2; - llama_context * ctx3 = llama_init_from_model(model, params_ctx3); - - llama_sampler * smpl3 = llama_sampler_chain_init(sparams); - - llama_sampler_chain_add(smpl3, llama_sampler_init_dist(params.sampling.seed)); - - printf("\nsingle seq run: %s", params.prompt.c_str()); - - // load state (rng, logits, embedding and kv_cache) from file - n_token_count_out = 0; - - if (!llama_state_load_file(ctx3, state_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) { - fprintf(stderr, "\n%s : failed to load state\n", __func__); - return 1; - } - - fprintf(stderr, "%s : loaded state with %zu tokens\n", __func__, n_token_count_out); - - // restore state (last tokens) - n_past = n_token_count_out; - if (!common_replay_last_token(ctx3, tokens.back(), n_past)) { - return 1; - } - ++n_past; - - // save seq 0 and load into seq 1 - { - // save kv of seq 0 - std::vector seq_store(llama_state_seq_get_size(ctx3, 0)); - const size_t ncopy = llama_state_seq_get_data(ctx3, seq_store.data(), seq_store.size(), 0); - if (ncopy != seq_store.size()) { - fprintf(stderr, "\n%s : seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size()); - return 1; - } - fprintf(stderr, "%s : seq 0 copied, %zd bytes\n", __func__, ncopy); - - // erase whole kv - llama_memory_clear(llama_get_memory(ctx3), true); - fprintf(stderr, "%s : kv cache cleared\n", __func__); - - // restore kv into seq 1 - const size_t nset = llama_state_seq_set_data(ctx3, seq_store.data(), seq_store.size(), 1); - if (nset != seq_store.size()) { - fprintf(stderr, "\n%s : seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size()); - return 1; - } - fprintf(stderr, "%s : seq 1 restored, %zd bytes\n", __func__, nset); - } - - // third run with seq 1 instead of 0 - for (auto i = 0; i < params.n_predict; i++) { - auto next_token = llama_sampler_sample(smpl3, ctx3, -1); - auto next_token_str = common_token_to_piece(ctx3, next_token); - - printf("%s", next_token_str.c_str()); - result2 += next_token_str; - - common_batch_clear(batch); - common_batch_add(batch, next_token, n_past, {1}, true); - - if (llama_decode(ctx3, batch)) { - fprintf(stderr, "\n%s : failed to evaluate\n", __func__); - llama_batch_free(batch); - return 1; - } - n_past += 1; - } - - printf("\n"); - - llama_sampler_free(smpl); - llama_sampler_free(smpl2); - llama_sampler_free(smpl3); - - llama_batch_free(batch); - - // this one is managed by common_init_result - //llama_free(ctx); - - llama_free(ctx2); - llama_free(ctx3); - - if (result0 != result2) { - fprintf(stderr, "\n%s : error : the seq restore generation is different\n", __func__); - return 1; - } - - fprintf(stderr, "\n%s : success\n", __func__); - - return 0; -} diff --git a/examples/simple-chat-rust/Cargo.toml b/examples/simple-chat-rust/Cargo.toml new file mode 100644 index 000000000000..560f00adf647 --- /dev/null +++ b/examples/simple-chat-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-simple-chat-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-simple-chat" +path = "src/main.rs" diff --git a/examples/simple-chat-rust/build.rs b/examples/simple-chat-rust/build.rs new file mode 100644 index 000000000000..2e56ca7cf8b9 --- /dev/null +++ b/examples/simple-chat-rust/build.rs @@ -0,0 +1,3 @@ +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); +} diff --git a/examples/simple-chat-rust/src/lib.rs b/examples/simple-chat-rust/src/lib.rs new file mode 100644 index 000000000000..bdbfed81bb87 --- /dev/null +++ b/examples/simple-chat-rust/src/lib.rs @@ -0,0 +1,109 @@ +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Args { + pub model_path: String, + pub n_ctx: i32, + pub n_gpu_layers: i32, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + MissingModel, + MissingValue(&'static str), + InvalidInteger(&'static str, String), + UnexpectedArgument(String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + ParseError::UnexpectedArgument(value) => write!(f, "unexpected argument: {value}"), + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut args = args.into_iter().map(Into::into); + let mut model_path = String::new(); + let mut n_ctx = 2048; + let mut n_gpu_layers = 99; + + while let Some(arg) = args.next() { + match arg.as_str() { + "-m" => { + model_path = args.next().ok_or(ParseError::MissingValue("-m"))?; + } + "-c" => { + let value = args.next().ok_or(ParseError::MissingValue("-c"))?; + n_ctx = value + .parse() + .map_err(|_| ParseError::InvalidInteger("-c", value))?; + } + "-ngl" => { + let value = args.next().ok_or(ParseError::MissingValue("-ngl"))?; + n_gpu_layers = value + .parse() + .map_err(|_| ParseError::InvalidInteger("-ngl", value))?; + } + _ => return Err(ParseError::UnexpectedArgument(arg)), + } + } + + if model_path.is_empty() { + return Err(ParseError::MissingModel); + } + + Ok(Args { + model_path, + n_ctx, + n_gpu_layers, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_defaults() { + let args = parse_args(["-m", "model.gguf"]).unwrap(); + + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.n_ctx, 2048); + assert_eq!(args.n_gpu_layers, 99); + } + + #[test] + fn parses_options() { + let args = parse_args(["-m", "model.gguf", "-c", "512", "-ngl", "0"]).unwrap(); + + assert_eq!(args.n_ctx, 512); + assert_eq!(args.n_gpu_layers, 0); + } + + #[test] + fn rejects_missing_model() { + assert_eq!( + parse_args(["-c", "512"]).unwrap_err(), + ParseError::MissingModel + ); + } + + #[test] + fn rejects_unexpected_argument() { + assert_eq!( + parse_args(["-m", "model.gguf", "prompt"]).unwrap_err(), + ParseError::UnexpectedArgument("prompt".to_string()) + ); + } +} diff --git a/examples/simple-chat-rust/src/main.rs b/examples/simple-chat-rust/src/main.rs new file mode 100644 index 000000000000..feda6969a3f0 --- /dev/null +++ b/examples/simple-chat-rust/src/main.rs @@ -0,0 +1,325 @@ +use llama_simple_chat_rust::{parse_args, Args}; +use llama_simple_rust::ffi; +use std::env; +use std::ffi::{c_char, c_void, CString}; +use std::io::{self, Write}; +use std::ptr; + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Sampler(*mut ffi::llama_sampler); + +impl Drop for Sampler { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_sampler_free(self.0); + } + } + } +} + +unsafe extern "C" fn log_errors_only( + level: ffi::ggml_log_level, + text: *const c_char, + _: *mut c_void, +) { + if level >= 2 && !text.is_null() { + let mut len = 0; + while *text.add(len) != 0 { + len += 1; + } + let bytes = std::slice::from_raw_parts(text as *const u8, len); + eprint!("{}", String::from_utf8_lossy(bytes)); + } +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m model.gguf [-c context_size] [-ngl n_gpu_layers]"); + eprintln!(); +} + +fn token_to_piece( + vocab: *const ffi::llama_vocab, + token: ffi::llama_token, +) -> Result { + let mut buf = vec![0_i8; 256]; + let n = unsafe { + ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), buf.len() as i32, 0, true) + }; + + if n < 0 { + return Err("failed to convert token to piece".to_string()); + } + + let bytes = buf[..n as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +fn generate( + ctx: *mut ffi::llama_context, + vocab: *const ffi::llama_vocab, + sampler: *mut ffi::llama_sampler, + prompt: &str, +) -> Result { + let c_prompt = CString::new(prompt).map_err(|_| "prompt contains an interior NUL byte")?; + let memory = unsafe { ffi::llama_get_memory(ctx) }; + let is_first = unsafe { ffi::llama_memory_seq_pos_max(memory, 0) } == -1; + + let n_prompt_tokens = unsafe { + -ffi::llama_tokenize( + vocab, + c_prompt.as_ptr(), + prompt.len() as i32, + ptr::null_mut(), + 0, + is_first, + true, + ) + }; + if n_prompt_tokens <= 0 { + return Err("failed to size prompt tokenization".to_string()); + } + + let mut prompt_tokens = vec![0_i32; n_prompt_tokens as usize]; + let n_tokenized = unsafe { + ffi::llama_tokenize( + vocab, + c_prompt.as_ptr(), + prompt.len() as i32, + prompt_tokens.as_mut_ptr(), + prompt_tokens.len() as i32, + is_first, + true, + ) + }; + if n_tokenized < 0 { + return Err("failed to tokenize the prompt".to_string()); + } + + let mut batch = + unsafe { ffi::llama_batch_get_one(prompt_tokens.as_mut_ptr(), prompt_tokens.len() as i32) }; + let mut response = String::new(); + let mut token_storage = [0_i32; 1]; + + loop { + let n_ctx = unsafe { ffi::llama_n_ctx(ctx) as i32 }; + let n_ctx_used = unsafe { ffi::llama_memory_seq_pos_max(memory, 0) } + 1; + if n_ctx_used + batch.n_tokens > n_ctx { + println!("\x1b[0m"); + return Err("context size exceeded".to_string()); + } + + let ret = unsafe { ffi::llama_decode(ctx, batch) }; + if ret != 0 { + return Err(format!("failed to decode, ret = {ret}")); + } + + token_storage[0] = unsafe { ffi::llama_sampler_sample(sampler, ctx, -1) }; + let new_token_id = token_storage[0]; + if unsafe { ffi::llama_vocab_is_eog(vocab, new_token_id) } { + break; + } + + let piece = token_to_piece(vocab, new_token_id)?; + print!("{piece}"); + io::stdout().flush().map_err(|err| err.to_string())?; + response.push_str(&piece); + + batch = unsafe { ffi::llama_batch_get_one(token_storage.as_mut_ptr(), 1) }; + } + + Ok(response) +} + +fn run(args: Args) -> Result<(), String> { + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + + unsafe { + ffi::llama_log_set(Some(log_errors_only), ptr::null_mut()); + ffi::ggml_backend_load_all(); + } + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = args.n_ctx as u32; + ctx_params.n_batch = args.n_ctx as u32; + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create the llama_context".to_string()); + } + + let sampler = Sampler(unsafe { + ffi::llama_sampler_chain_init(ffi::llama_sampler_chain_default_params()) + }); + if sampler.0.is_null() { + return Err("failed to create sampler chain".to_string()); + } + + unsafe { + ffi::llama_sampler_chain_add(sampler.0, ffi::llama_sampler_init_min_p(0.05, 1)); + ffi::llama_sampler_chain_add(sampler.0, ffi::llama_sampler_init_temp(0.8)); + ffi::llama_sampler_chain_add( + sampler.0, + ffi::llama_sampler_init_dist(ffi::LLAMA_DEFAULT_SEED), + ); + } + + let mut messages = Vec::<(CString, CString)>::new(); + let mut formatted = vec![0_i8; unsafe { ffi::llama_n_ctx(ctx.0) } as usize]; + let mut prev_len = 0_i32; + + loop { + print!("\x1b[32m> \x1b[0m"); + io::stdout().flush().map_err(|err| err.to_string())?; + + let mut user = String::new(); + io::stdin() + .read_line(&mut user) + .map_err(|err| err.to_string())?; + let user = user.trim_end_matches(['\r', '\n']).to_string(); + if user.is_empty() { + break; + } + + let role_user = CString::new("user").unwrap(); + let content_user = + CString::new(user).map_err(|_| "user input contains an interior NUL byte")?; + messages.push((role_user, content_user)); + + let c_messages = messages + .iter() + .map(|(role, content)| ffi::llama_chat_message { + role: role.as_ptr(), + content: content.as_ptr(), + }) + .collect::>(); + let tmpl = unsafe { ffi::llama_model_chat_template(model.0, ptr::null()) }; + let mut new_len = unsafe { + ffi::llama_chat_apply_template( + tmpl, + c_messages.as_ptr(), + c_messages.len(), + true, + formatted.as_mut_ptr(), + formatted.len() as i32, + ) + }; + if new_len > formatted.len() as i32 { + formatted.resize(new_len as usize, 0); + new_len = unsafe { + ffi::llama_chat_apply_template( + tmpl, + c_messages.as_ptr(), + c_messages.len(), + true, + formatted.as_mut_ptr(), + formatted.len() as i32, + ) + }; + } + if new_len < 0 { + return Err("failed to apply the chat template".to_string()); + } + + let prompt_bytes = formatted[prev_len as usize..new_len as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + let prompt = String::from_utf8_lossy(&prompt_bytes).into_owned(); + + print!("\x1b[33m"); + let response = generate(ctx.0, vocab, sampler.0, &prompt)?; + println!("\n\x1b[0m"); + + let role_assistant = CString::new("assistant").unwrap(); + let content_assistant = + CString::new(response).map_err(|_| "response contains an interior NUL byte")?; + messages.push((role_assistant, content_assistant)); + + let c_messages = messages + .iter() + .map(|(role, content)| ffi::llama_chat_message { + role: role.as_ptr(), + content: content.as_ptr(), + }) + .collect::>(); + prev_len = unsafe { + ffi::llama_chat_apply_template( + tmpl, + c_messages.as_ptr(), + c_messages.len(), + false, + ptr::null_mut(), + 0, + ) + }; + if prev_len < 0 { + return Err("failed to apply the chat template".to_string()); + } + } + + Ok(()) +} + +fn main() { + let mut argv = env::args(); + let program = argv + .next() + .unwrap_or_else(|| "llama-simple-chat".to_string()); + + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} diff --git a/examples/simple-chat/CMakeLists.txt b/examples/simple-chat/CMakeLists.txt index 567f7fbbbf43..e8db1a317d91 100644 --- a/examples/simple-chat/CMakeLists.txt +++ b/examples/simple-chat/CMakeLists.txt @@ -1,5 +1,25 @@ set(TARGET llama-simple-chat) -add_executable(${TARGET} simple-chat.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(SIMPLE_CHAT_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../simple-chat-rust/Cargo.toml) +set(SIMPLE_CHAT_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/simple-chat) +set(SIMPLE_CHAT_RUST_BIN ${SIMPLE_CHAT_RUST_TARGET_DIR}/debug/llama-simple-chat${CMAKE_EXECUTABLE_SUFFIX}) +set(SIMPLE_CHAT_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${SIMPLE_CHAT_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${SIMPLE_CHAT_RUST_MANIFEST} --target-dir ${SIMPLE_CHAT_RUST_TARGET_DIR} --bin llama-simple-chat + COMMAND ${CMAKE_COMMAND} -E copy ${SIMPLE_CHAT_RUST_BIN} ${SIMPLE_CHAT_RUST_OUTPUT} + DEPENDS + llama + ${SIMPLE_CHAT_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-chat-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-chat-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${SIMPLE_CHAT_RUST_OUTPUT}) +install(PROGRAMS ${SIMPLE_CHAT_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/simple-chat/simple-chat.cpp b/examples/simple-chat/simple-chat.cpp deleted file mode 100644 index 97e9dc9842f5..000000000000 --- a/examples/simple-chat/simple-chat.cpp +++ /dev/null @@ -1,210 +0,0 @@ -#include "llama.h" -#include -#include -#include -#include -#include -#include - -static void print_usage(int, char ** argv) { - printf("\nexample usage:\n"); - printf("\n %s -m model.gguf [-c context_size] [-ngl n_gpu_layers]\n", argv[0]); - printf("\n"); -} - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - std::string model_path; - int ngl = 99; - int n_ctx = 2048; - - // parse command line arguments - for (int i = 1; i < argc; i++) { - try { - if (strcmp(argv[i], "-m") == 0) { - if (i + 1 < argc) { - model_path = argv[++i]; - } else { - print_usage(argc, argv); - return 1; - } - } else if (strcmp(argv[i], "-c") == 0) { - if (i + 1 < argc) { - n_ctx = std::stoi(argv[++i]); - } else { - print_usage(argc, argv); - return 1; - } - } else if (strcmp(argv[i], "-ngl") == 0) { - if (i + 1 < argc) { - ngl = std::stoi(argv[++i]); - } else { - print_usage(argc, argv); - return 1; - } - } else { - print_usage(argc, argv); - return 1; - } - } catch (std::exception & e) { - fprintf(stderr, "error: %s\n", e.what()); - print_usage(argc, argv); - return 1; - } - } - if (model_path.empty()) { - print_usage(argc, argv); - return 1; - } - - // only print errors - llama_log_set([](enum ggml_log_level level, const char * text, void * /* user_data */) { - if (level >= GGML_LOG_LEVEL_ERROR) { - fprintf(stderr, "%s", text); - } - }, nullptr); - - // load dynamic backends - ggml_backend_load_all(); - - // initialize the model - llama_model_params model_params = llama_model_default_params(); - model_params.n_gpu_layers = ngl; - - llama_model * model = llama_model_load_from_file(model_path.c_str(), model_params); - if (!model) { - fprintf(stderr , "%s: error: unable to load model\n" , __func__); - return 1; - } - - const llama_vocab * vocab = llama_model_get_vocab(model); - - // initialize the context - llama_context_params ctx_params = llama_context_default_params(); - ctx_params.n_ctx = n_ctx; - ctx_params.n_batch = n_ctx; - - llama_context * ctx = llama_init_from_model(model, ctx_params); - if (!ctx) { - fprintf(stderr , "%s: error: failed to create the llama_context\n" , __func__); - return 1; - } - - // initialize the sampler - llama_sampler * smpl = llama_sampler_chain_init(llama_sampler_chain_default_params()); - llama_sampler_chain_add(smpl, llama_sampler_init_min_p(0.05f, 1)); - llama_sampler_chain_add(smpl, llama_sampler_init_temp(0.8f)); - llama_sampler_chain_add(smpl, llama_sampler_init_dist(LLAMA_DEFAULT_SEED)); - - // helper function to evaluate a prompt and generate a response - auto generate = [&](const std::string & prompt) { - std::string response; - - const bool is_first = llama_memory_seq_pos_max(llama_get_memory(ctx), 0) == -1; - - // tokenize the prompt - const int n_prompt_tokens = -llama_tokenize(vocab, prompt.c_str(), prompt.size(), NULL, 0, is_first, true); - std::vector prompt_tokens(n_prompt_tokens); - if (llama_tokenize(vocab, prompt.c_str(), prompt.size(), prompt_tokens.data(), prompt_tokens.size(), is_first, true) < 0) { - GGML_ABORT("failed to tokenize the prompt\n"); - } - - // prepare a batch for the prompt - llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size()); - llama_token new_token_id; - while (true) { - // check if we have enough space in the context to evaluate this batch - int n_ctx = llama_n_ctx(ctx); - int n_ctx_used = llama_memory_seq_pos_max(llama_get_memory(ctx), 0) + 1; - if (n_ctx_used + batch.n_tokens > n_ctx) { - printf("\033[0m\n"); - fprintf(stderr, "context size exceeded\n"); - exit(0); - } - - int ret = llama_decode(ctx, batch); - if (ret != 0) { - GGML_ABORT("failed to decode, ret = %d\n", ret); - } - - // sample the next token - new_token_id = llama_sampler_sample(smpl, ctx, -1); - - // is it an end of generation? - if (llama_vocab_is_eog(vocab, new_token_id)) { - break; - } - - // convert the token to a string, print it and add it to the response - char buf[256]; - int n = llama_token_to_piece(vocab, new_token_id, buf, sizeof(buf), 0, true); - if (n < 0) { - GGML_ABORT("failed to convert token to piece\n"); - } - std::string piece(buf, n); - printf("%s", piece.c_str()); - fflush(stdout); - response += piece; - - // prepare the next batch with the sampled token - batch = llama_batch_get_one(&new_token_id, 1); - } - - return response; - }; - - std::vector messages; - std::vector formatted(llama_n_ctx(ctx)); - int prev_len = 0; - while (true) { - // get user input - printf("\033[32m> \033[0m"); - std::string user; - std::getline(std::cin, user); - - if (user.empty()) { - break; - } - - const char * tmpl = llama_model_chat_template(model, /* name */ nullptr); - - // add the user input to the message list and format it - messages.push_back({"user", strdup(user.c_str())}); - int new_len = llama_chat_apply_template(tmpl, messages.data(), messages.size(), true, formatted.data(), formatted.size()); - if (new_len > (int)formatted.size()) { - formatted.resize(new_len); - new_len = llama_chat_apply_template(tmpl, messages.data(), messages.size(), true, formatted.data(), formatted.size()); - } - if (new_len < 0) { - fprintf(stderr, "failed to apply the chat template\n"); - return 1; - } - - // remove previous messages to obtain the prompt to generate the response - std::string prompt(formatted.begin() + prev_len, formatted.begin() + new_len); - - // generate a response - printf("\033[33m"); - std::string response = generate(prompt); - printf("\n\033[0m"); - - // add the response to the messages - messages.push_back({"assistant", strdup(response.c_str())}); - prev_len = llama_chat_apply_template(tmpl, messages.data(), messages.size(), false, nullptr, 0); - if (prev_len < 0) { - fprintf(stderr, "failed to apply the chat template\n"); - return 1; - } - } - - // free resources - for (auto & msg : messages) { - free(const_cast(msg.content)); - } - llama_sampler_free(smpl); - llama_free(ctx); - llama_model_free(model); - - return 0; -} diff --git a/examples/simple-cmake-pkg/CMakeLists.txt b/examples/simple-cmake-pkg/CMakeLists.txt index 128e38c8f2dc..7271938d3f69 100644 --- a/examples/simple-cmake-pkg/CMakeLists.txt +++ b/examples/simple-cmake-pkg/CMakeLists.txt @@ -4,8 +4,26 @@ project(llama-simple-cmake-pkg) set(TARGET llama-simple-cmake-pkg) find_package(Llama REQUIRED) +find_program(CARGO_EXECUTABLE cargo REQUIRED) -add_executable(${TARGET} ${CMAKE_CURRENT_LIST_DIR}/../simple/simple.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama ggml::all ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +set(SIMPLE_RUST_MANIFEST ${CMAKE_CURRENT_LIST_DIR}/../simple-rust/Cargo.toml) +set(SIMPLE_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/simple-cmake-pkg) +set(SIMPLE_RUST_BIN ${SIMPLE_RUST_TARGET_DIR}/debug/llama-simple${CMAKE_EXECUTABLE_SUFFIX}) +set(SIMPLE_RUST_OUTPUT ${CMAKE_BINARY_DIR}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${SIMPLE_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${SIMPLE_RUST_MANIFEST} --target-dir ${SIMPLE_RUST_TARGET_DIR} --bin llama-simple + COMMAND ${CMAKE_COMMAND} -E copy ${SIMPLE_RUST_BIN} ${SIMPLE_RUST_OUTPUT} + DEPENDS + ${SIMPLE_RUST_MANIFEST} + ${CMAKE_CURRENT_LIST_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_LIST_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_LIST_DIR}/../simple-rust/src/lib.rs + ${CMAKE_CURRENT_LIST_DIR}/../simple-rust/src/main.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${SIMPLE_RUST_OUTPUT}) +include(GNUInstallDirs) +install(PROGRAMS ${SIMPLE_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/simple-rust/Cargo.toml b/examples/simple-rust/Cargo.toml new file mode 100644 index 000000000000..ef4939e58772 --- /dev/null +++ b/examples/simple-rust/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "llama-simple-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[build-dependencies] + +[dependencies] + +[[bin]] +name = "llama-simple" +path = "src/main.rs" diff --git a/examples/simple-rust/README.md b/examples/simple-rust/README.md new file mode 100644 index 000000000000..786b4bb46456 --- /dev/null +++ b/examples/simple-rust/README.md @@ -0,0 +1,14 @@ +# simple-rust + +Rust port of `examples/simple/simple.cpp` using the llama.cpp C API. + +This is intentionally a standalone Cargo example, not part of the default CMake build. +Build llama.cpp first, then point Cargo at the produced `libllama` directory if it is not under `build/bin`, `build/src`, or `build`. + +```sh +cmake -B build +cmake --build build --target llama + +cargo test --manifest-path examples/simple-rust/Cargo.toml --lib +LLAMA_CPP_LIB_DIR="$PWD/build/bin" cargo run --manifest-path examples/simple-rust/Cargo.toml --bin llama-simple -- -m model.gguf -n 32 "Hello my name is" +``` diff --git a/examples/simple-rust/build.rs b/examples/simple-rust/build.rs new file mode 100644 index 000000000000..656da7b97741 --- /dev/null +++ b/examples/simple-rust/build.rs @@ -0,0 +1,46 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + emit_rpath(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + println!("cargo:rustc-link-search=native={}", candidate.display()); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + emit_rpath(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building the binary"); +} + +fn emit_rpath(lib_dir: &str) { + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/simple-rust/src/ffi.rs b/examples/simple-rust/src/ffi.rs new file mode 100644 index 000000000000..1e8caf912792 --- /dev/null +++ b/examples/simple-rust/src/ffi.rs @@ -0,0 +1,555 @@ +#![allow(non_camel_case_types)] +#![allow(non_snake_case)] +#![allow(dead_code)] + +use std::ffi::{c_char, c_float, c_int, c_void}; + +pub const LLAMA_TOKEN_NULL: llama_token = -1; +pub const LLAMA_DEFAULT_SEED: u32 = 0xFFFF_FFFF; + +pub type llama_pos = i32; +pub type llama_token = i32; +pub type llama_seq_id = i32; +pub type ggml_backend_dev_t = *mut c_void; +pub type ggml_backend_reg_t = *mut c_void; +pub type ggml_backend_buffer_type_t = *mut c_void; +pub type ggml_type = c_int; +pub type llama_memory_t = *mut c_void; +pub type ggml_log_level = c_int; +pub type llama_pooling_type = c_int; +pub type llama_vocab_type = c_int; +pub type ggml_opt_dataset_t = *mut c_void; +pub type ggml_opt_result_t = *mut c_void; +pub type ggml_opt_context_t = *mut c_void; +pub type ggml_opt_optimizer_type = c_int; + +#[repr(C)] +pub struct llama_vocab { + _private: [u8; 0], +} + +#[repr(C)] +pub struct llama_model { + _private: [u8; 0], +} + +#[repr(C)] +pub struct llama_context { + _private: [u8; 0], +} + +#[repr(C)] +pub struct llama_sampler { + _private: [u8; 0], +} + +#[repr(C)] +pub struct ggml_tensor { + _private: [u8; 0], +} + +#[repr(C)] +pub struct ggml_context { + _private: [u8; 0], +} + +#[repr(C)] +pub struct gguf_context { + _private: [u8; 0], +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct ggml_opt_adamw_params { + pub alpha: f32, + pub beta1: f32, + pub beta2: f32, + pub eps: f32, + pub wd: f32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct ggml_opt_sgd_params { + pub alpha: f32, + pub wd: f32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct ggml_opt_optimizer_params { + pub adamw: ggml_opt_adamw_params, + pub sgd: ggml_opt_sgd_params, +} + +pub type ggml_opt_get_optimizer_params = + Option ggml_opt_optimizer_params>; +pub type llama_opt_param_filter = + Option bool>; +pub type ggml_opt_epoch_callback = Option< + unsafe extern "C" fn( + opt_ctx: ggml_opt_context_t, + dataset: ggml_opt_dataset_t, + result: ggml_opt_result_t, + ibatch: i64, + ibatch_max: i64, + train: bool, + ), +>; + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct llama_opt_params { + pub n_ctx_train: u32, + pub param_filter: llama_opt_param_filter, + pub param_filter_ud: *mut c_void, + pub get_opt_pars: ggml_opt_get_optimizer_params, + pub get_opt_pars_ud: *mut c_void, + pub optimizer_type: ggml_opt_optimizer_type, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct ggml_init_params { + pub mem_size: usize, + pub mem_buffer: *mut c_void, + pub no_alloc: bool, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct gguf_init_params { + pub no_alloc: bool, + pub ctx: *mut *mut ggml_context, +} + +#[repr(C)] +pub struct llama_model_tensor_buft_override { + pub pattern: *const c_char, + pub buft: ggml_backend_buffer_type_t, +} + +#[repr(C)] +pub struct llama_model_kv_override { + _private: [u8; 0], +} + +#[repr(C)] +pub struct llama_sampler_seq_config { + pub seq_id: llama_seq_id, + pub sampler: *mut llama_sampler, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct llama_token_data { + pub id: llama_token, + pub logit: f32, + pub p: f32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct llama_token_data_array { + pub data: *mut llama_token_data, + pub size: usize, + pub selected: i64, + pub sorted: bool, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct llama_model_params { + pub devices: *mut ggml_backend_dev_t, + pub tensor_buft_overrides: *const llama_model_tensor_buft_override, + pub n_gpu_layers: i32, + pub split_mode: c_int, + pub main_gpu: i32, + pub tensor_split: *const c_float, + pub progress_callback: Option bool>, + pub progress_callback_user_data: *mut c_void, + pub kv_overrides: *const llama_model_kv_override, + pub vocab_only: bool, + pub use_mmap: bool, + pub use_direct_io: bool, + pub use_mlock: bool, + pub check_tensors: bool, + pub use_extra_bufts: bool, + pub no_host: bool, + pub no_alloc: bool, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct llama_context_params { + pub n_ctx: u32, + pub n_batch: u32, + pub n_ubatch: u32, + pub n_seq_max: u32, + pub n_threads: i32, + pub n_threads_batch: i32, + pub rope_scaling_type: c_int, + pub pooling_type: c_int, + pub attention_type: c_int, + pub flash_attn_type: c_int, + pub rope_freq_base: f32, + pub rope_freq_scale: f32, + pub yarn_ext_factor: f32, + pub yarn_attn_factor: f32, + pub yarn_beta_fast: f32, + pub yarn_beta_slow: f32, + pub yarn_orig_ctx: u32, + pub defrag_thold: f32, + pub cb_eval: Option bool>, + pub cb_eval_user_data: *mut c_void, + pub type_k: ggml_type, + pub type_v: ggml_type, + pub abort_callback: Option bool>, + pub abort_callback_data: *mut c_void, + pub embeddings: bool, + pub offload_kqv: bool, + pub no_perf: bool, + pub op_offload: bool, + pub swa_full: bool, + pub kv_unified: bool, + pub samplers: *mut llama_sampler_seq_config, + pub n_samplers: usize, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct llama_sampler_chain_params { + pub no_perf: bool, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct llama_batch { + pub n_tokens: i32, + pub token: *mut llama_token, + pub embd: *mut f32, + pub pos: *mut llama_pos, + pub n_seq_id: *mut i32, + pub seq_id: *mut *mut llama_seq_id, + pub logits: *mut i8, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct llama_chat_message { + pub role: *const c_char, + pub content: *const c_char, +} + +extern "C" { + pub fn ggml_backend_load_all(); + pub fn ggml_backend_dev_count() -> usize; + pub fn ggml_backend_dev_get(index: usize) -> ggml_backend_dev_t; + pub fn ggml_backend_dev_type(dev: ggml_backend_dev_t) -> c_int; + pub fn ggml_backend_dev_by_name(name: *const c_char) -> ggml_backend_dev_t; + pub fn ggml_backend_dev_by_type(typ: c_int) -> ggml_backend_dev_t; + pub fn ggml_backend_dev_name(dev: ggml_backend_dev_t) -> *const c_char; + pub fn ggml_backend_dev_description(dev: ggml_backend_dev_t) -> *const c_char; + pub fn ggml_backend_dev_memory(dev: ggml_backend_dev_t, free: *mut usize, total: *mut usize); + pub fn ggml_backend_reg_by_name(name: *const c_char) -> ggml_backend_reg_t; + pub fn ggml_backend_reg_get_proc_address( + reg: ggml_backend_reg_t, + name: *const c_char, + ) -> *mut c_void; + pub fn ggml_tensor_overhead() -> usize; + pub fn ggml_init(params: ggml_init_params) -> *mut ggml_context; + pub fn ggml_free(ctx: *mut ggml_context); + pub fn ggml_new_tensor_1d(ctx: *mut ggml_context, typ: ggml_type, ne0: i64) + -> *mut ggml_tensor; + pub fn ggml_new_tensor_2d( + ctx: *mut ggml_context, + typ: ggml_type, + ne0: i64, + ne1: i64, + ) -> *mut ggml_tensor; + pub fn ggml_get_data(tensor: *const ggml_tensor) -> *mut c_void; + pub fn ggml_get_data_f32(tensor: *const ggml_tensor) -> *mut f32; + pub fn ggml_set_name(tensor: *mut ggml_tensor, name: *const c_char) -> *mut ggml_tensor; + pub fn ggml_time_us() -> i64; + pub fn llama_log_set( + log_callback: Option, + user_data: *mut c_void, + ); + + pub fn llama_model_default_params() -> llama_model_params; + pub fn llama_context_default_params() -> llama_context_params; + pub fn llama_sampler_chain_default_params() -> llama_sampler_chain_params; + pub fn llama_backend_init(); + pub fn llama_backend_free(); + pub fn llama_numa_init(numa: c_int); + pub fn llama_time_us() -> i64; + pub fn llama_max_parallel_sequences() -> usize; + + pub fn llama_model_load_from_file( + path_model: *const c_char, + params: llama_model_params, + ) -> *mut llama_model; + pub fn llama_model_free(model: *mut llama_model); + pub fn llama_model_save_to_file(model: *const llama_model, path_model: *const c_char); + pub fn llama_init_from_model( + model: *mut llama_model, + params: llama_context_params, + ) -> *mut llama_context; + pub fn llama_free(ctx: *mut llama_context); + pub fn llama_get_model(ctx: *const llama_context) -> *const llama_model; + pub fn llama_get_memory(ctx: *const llama_context) -> llama_memory_t; + pub fn llama_n_ctx(ctx: *const llama_context) -> u32; + pub fn llama_n_seq_max(ctx: *const llama_context) -> u32; + pub fn llama_n_ubatch(ctx: *const llama_context) -> u32; + pub fn llama_set_n_threads( + ctx: *mut llama_context, + n_threads: i32, + n_threads_batch: i32, + ); + pub fn llama_pooling_type(ctx: *const llama_context) -> llama_pooling_type; + + pub fn llama_model_get_vocab(model: *const llama_model) -> *const llama_vocab; + pub fn llama_model_n_embd_out(model: *const llama_model) -> i32; + pub fn llama_model_n_ctx_train(model: *const llama_model) -> i32; + pub fn llama_model_chat_template( + model: *const llama_model, + name: *const c_char, + ) -> *const c_char; + pub fn llama_model_has_encoder(model: *const llama_model) -> bool; + pub fn llama_model_has_decoder(model: *const llama_model) -> bool; + pub fn llama_model_is_diffusion(model: *const llama_model) -> bool; + pub fn llama_model_meta_val_str( + model: *const llama_model, + key: *const c_char, + buf: *mut c_char, + buf_size: usize, + ) -> i32; + pub fn llama_model_n_cls_out(model: *const llama_model) -> u32; + pub fn llama_model_cls_label(model: *const llama_model, i: u32) -> *const c_char; + pub fn llama_model_decoder_start_token(model: *const llama_model) -> llama_token; + + pub fn llama_vocab_bos(vocab: *const llama_vocab) -> llama_token; + pub fn llama_vocab_eos(vocab: *const llama_vocab) -> llama_token; + pub fn llama_vocab_eot(vocab: *const llama_vocab) -> llama_token; + pub fn llama_vocab_mask(vocab: *const llama_vocab) -> llama_token; + pub fn llama_vocab_sep(vocab: *const llama_vocab) -> llama_token; + pub fn llama_vocab_get_text(vocab: *const llama_vocab, token: llama_token) -> *const c_char; + pub fn llama_vocab_get_add_bos(vocab: *const llama_vocab) -> bool; + pub fn llama_vocab_get_add_eos(vocab: *const llama_vocab) -> bool; + pub fn llama_vocab_get_add_sep(vocab: *const llama_vocab) -> bool; + pub fn llama_vocab_is_eog(vocab: *const llama_vocab, token: llama_token) -> bool; + pub fn llama_vocab_n_tokens(vocab: *const llama_vocab) -> i32; + pub fn llama_vocab_type(vocab: *const llama_vocab) -> llama_vocab_type; + pub fn llama_memory_seq_pos_max(mem: llama_memory_t, seq_id: llama_seq_id) -> llama_pos; + pub fn llama_memory_clear(mem: llama_memory_t, data: bool); + pub fn llama_memory_seq_rm( + mem: llama_memory_t, + seq_id: llama_seq_id, + p0: llama_pos, + p1: llama_pos, + ) -> bool; + pub fn llama_memory_seq_cp( + mem: llama_memory_t, + seq_id_src: llama_seq_id, + seq_id_dst: llama_seq_id, + p0: llama_pos, + p1: llama_pos, + ); + pub fn llama_memory_seq_keep(mem: llama_memory_t, seq_id: llama_seq_id); + pub fn llama_memory_seq_add( + mem: llama_memory_t, + seq_id: llama_seq_id, + p0: llama_pos, + p1: llama_pos, + delta: llama_pos, + ); + pub fn llama_memory_seq_div( + mem: llama_memory_t, + seq_id: llama_seq_id, + p0: llama_pos, + p1: llama_pos, + d: i32, + ); + + pub fn llama_tokenize( + vocab: *const llama_vocab, + text: *const c_char, + text_len: i32, + tokens: *mut llama_token, + n_tokens_max: i32, + add_special: bool, + parse_special: bool, + ) -> i32; + pub fn llama_detokenize( + vocab: *const llama_vocab, + tokens: *const llama_token, + n_tokens: i32, + text: *mut c_char, + text_len_max: i32, + remove_special: bool, + unparse_special: bool, + ) -> i32; + pub fn llama_token_to_piece( + vocab: *const llama_vocab, + token: llama_token, + buf: *mut c_char, + length: i32, + lstrip: i32, + special: bool, + ) -> i32; + + pub fn llama_batch_get_one(tokens: *mut llama_token, n_tokens: i32) -> llama_batch; + pub fn llama_batch_init(n_tokens: i32, embd: i32, n_seq_max: i32) -> llama_batch; + pub fn llama_batch_free(batch: llama_batch); + pub fn llama_encode(ctx: *mut llama_context, batch: llama_batch) -> i32; + pub fn llama_decode(ctx: *mut llama_context, batch: llama_batch) -> i32; + pub fn llama_synchronize(ctx: *mut llama_context); + pub fn llama_set_causal_attn(ctx: *mut llama_context, causal_attn: bool); + pub fn llama_get_logits_ith(ctx: *mut llama_context, i: i32) -> *const f32; + pub fn llama_get_embeddings(ctx: *mut llama_context) -> *mut f32; + pub fn llama_get_embeddings_ith(ctx: *mut llama_context, i: i32) -> *mut f32; + pub fn llama_get_embeddings_seq(ctx: *mut llama_context, seq_id: llama_seq_id) -> *mut f32; + pub fn llama_state_seq_get_size(ctx: *mut llama_context, seq_id: llama_seq_id) -> usize; + pub fn llama_state_seq_get_data( + ctx: *mut llama_context, + dst: *mut u8, + size: usize, + seq_id: llama_seq_id, + ) -> usize; + pub fn llama_state_seq_set_data( + ctx: *mut llama_context, + src: *const u8, + size: usize, + dest_seq_id: llama_seq_id, + ) -> usize; + pub fn llama_state_load_file( + ctx: *mut llama_context, + path_session: *const c_char, + tokens_out: *mut llama_token, + n_token_capacity: usize, + n_token_count_out: *mut usize, + ) -> bool; + pub fn llama_state_save_file( + ctx: *mut llama_context, + path_session: *const c_char, + tokens: *const llama_token, + n_token_count: usize, + ) -> bool; + + pub fn llama_sampler_chain_init(params: llama_sampler_chain_params) -> *mut llama_sampler; + pub fn llama_sampler_chain_add(chain: *mut llama_sampler, sampler: *mut llama_sampler); + pub fn llama_sampler_init_greedy() -> *mut llama_sampler; + pub fn llama_sampler_init_top_k(k: i32) -> *mut llama_sampler; + pub fn llama_sampler_init_top_p(p: f32, min_keep: usize) -> *mut llama_sampler; + pub fn llama_sampler_init_min_p(p: f32, min_keep: usize) -> *mut llama_sampler; + pub fn llama_sampler_init_temp(t: f32) -> *mut llama_sampler; + pub fn llama_sampler_init_dist(seed: u32) -> *mut llama_sampler; + pub fn llama_sampler_apply(smpl: *mut llama_sampler, cur_p: *mut llama_token_data_array); + pub fn llama_sampler_sample( + smpl: *mut llama_sampler, + ctx: *mut llama_context, + idx: i32, + ) -> llama_token; + pub fn llama_sampler_accept(smpl: *mut llama_sampler, token: llama_token); + pub fn llama_sampler_reset(smpl: *mut llama_sampler); + pub fn llama_sampler_free(smpl: *mut llama_sampler); + + pub fn llama_perf_sampler_print(chain: *const llama_sampler); + pub fn llama_perf_context_print(ctx: *const llama_context); + + pub fn ggml_opt_dataset_init( + type_data: ggml_type, + type_label: ggml_type, + ne_datapoint: i64, + ne_label: i64, + ndata: i64, + ndata_shard: i64, + ) -> ggml_opt_dataset_t; + pub fn ggml_opt_dataset_free(dataset: ggml_opt_dataset_t); + pub fn ggml_opt_dataset_ndata(dataset: ggml_opt_dataset_t) -> i64; + pub fn ggml_opt_dataset_data(dataset: ggml_opt_dataset_t) -> *mut ggml_tensor; + pub fn ggml_opt_dataset_labels(dataset: ggml_opt_dataset_t) -> *mut ggml_tensor; + pub fn ggml_opt_get_default_optimizer_params( + userdata: *mut c_void, + ) -> ggml_opt_optimizer_params; + pub fn ggml_opt_optimizer_name(optimizer: ggml_opt_optimizer_type) -> *const c_char; + pub fn ggml_opt_result_init() -> ggml_opt_result_t; + pub fn ggml_opt_result_free(result: ggml_opt_result_t); + pub fn ggml_opt_result_reset(result: ggml_opt_result_t); + pub fn ggml_opt_epoch_callback_progress_bar( + opt_ctx: ggml_opt_context_t, + dataset: ggml_opt_dataset_t, + result: ggml_opt_result_t, + ibatch: i64, + ibatch_max: i64, + train: bool, + ); + pub fn llama_opt_param_filter_all(tensor: *const ggml_tensor, userdata: *mut c_void) -> bool; + pub fn llama_opt_init( + lctx: *mut llama_context, + model: *mut llama_model, + lopt_params: llama_opt_params, + ); + pub fn llama_opt_epoch( + lctx: *mut llama_context, + dataset: ggml_opt_dataset_t, + result_train: ggml_opt_result_t, + result_eval: ggml_opt_result_t, + idata_split: i64, + callback_train: ggml_opt_epoch_callback, + callback_eval: ggml_opt_epoch_callback, + ); + + pub fn llama_chat_apply_template( + tmpl: *const c_char, + chat: *const llama_chat_message, + n_msg: usize, + add_ass: bool, + buf: *mut c_char, + length: i32, + ) -> i32; + + pub fn gguf_init_empty() -> *mut gguf_context; + pub fn gguf_init_from_file(fname: *const c_char, params: gguf_init_params) + -> *mut gguf_context; + pub fn gguf_free(ctx: *mut gguf_context); + pub fn gguf_get_data_offset(ctx: *const gguf_context) -> usize; + pub fn gguf_find_key(ctx: *const gguf_context, key: *const c_char) -> i64; + pub fn gguf_get_val_u16(ctx: *const gguf_context, key_id: i64) -> u16; + pub fn gguf_get_val_str(ctx: *const gguf_context, key_id: i64) -> *const c_char; + pub fn gguf_get_n_tensors(ctx: *const gguf_context) -> i64; + pub fn gguf_get_tensor_name(ctx: *const gguf_context, tensor_id: i64) -> *const c_char; + pub fn gguf_find_tensor(ctx: *const gguf_context, name: *const c_char) -> i64; + pub fn gguf_get_tensor_offset(ctx: *const gguf_context, tensor_id: i64) -> usize; + pub fn gguf_get_tensor_size(ctx: *const gguf_context, tensor_id: i64) -> usize; + pub fn gguf_set_kv(ctx: *mut gguf_context, src: *const gguf_context); + pub fn gguf_set_val_u16(ctx: *mut gguf_context, key: *const c_char, val: u16); + pub fn gguf_set_val_i32(ctx: *mut gguf_context, key: *const c_char, val: i32); + pub fn gguf_set_val_str(ctx: *mut gguf_context, key: *const c_char, val: *const c_char); + pub fn gguf_add_tensor(ctx: *mut gguf_context, tensor: *const ggml_tensor); + pub fn gguf_get_meta_size(ctx: *const gguf_context) -> usize; + pub fn gguf_get_meta_data(ctx: *const gguf_context, data: *mut c_void); + pub fn gguf_write_to_file( + ctx: *const gguf_context, + fname: *const c_char, + only_meta: bool, + ) -> bool; + + pub fn ggml_get_tensor(ctx: *mut ggml_context, name: *const c_char) -> *mut ggml_tensor; + pub fn ggml_nbytes(tensor: *const ggml_tensor) -> usize; + + pub fn llama_split_path( + split_path: *mut c_char, + maxlen: usize, + path_prefix: *const c_char, + split_no: i32, + split_count: i32, + ) -> i32; + pub fn llama_split_prefix( + split_prefix: *mut c_char, + maxlen: usize, + split_path: *const c_char, + split_no: i32, + split_count: i32, + ) -> i32; +} diff --git a/examples/simple-rust/src/lib.rs b/examples/simple-rust/src/lib.rs new file mode 100644 index 000000000000..917eff3dfaf2 --- /dev/null +++ b/examples/simple-rust/src/lib.rs @@ -0,0 +1,144 @@ +pub mod ffi; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub n_gpu_layers: i32, + pub n_predict: i32, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ParseError { + MissingModel, + MissingValue(&'static str), + InvalidInteger(&'static str, String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m model.gguf"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut args = args.into_iter().map(Into::into).peekable(); + let mut model_path = String::new(); + let mut prompt_parts = Vec::new(); + let mut n_gpu_layers = 99; + let mut n_predict = 32; + + while let Some(arg) = args.next() { + match arg.as_str() { + "-m" => { + model_path = args.next().ok_or(ParseError::MissingValue("-m"))?; + } + "-n" => { + let value = args.next().ok_or(ParseError::MissingValue("-n"))?; + n_predict = value + .parse() + .map_err(|_| ParseError::InvalidInteger("-n", value))?; + } + "-ngl" => { + let value = args.next().ok_or(ParseError::MissingValue("-ngl"))?; + n_gpu_layers = value + .parse() + .map_err(|_| ParseError::InvalidInteger("-ngl", value))?; + } + _ => { + prompt_parts.push(arg); + prompt_parts.extend(args); + break; + } + } + } + + if model_path.is_empty() { + return Err(ParseError::MissingModel); + } + + let prompt = if prompt_parts.is_empty() { + "Hello my name is".to_string() + } else { + prompt_parts.join(" ") + }; + + Ok(Args { + model_path, + prompt, + n_gpu_layers, + n_predict, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_required_model_with_defaults() { + let args = parse_args(["-m", "model.gguf"]).unwrap(); + + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "Hello my name is"); + assert_eq!(args.n_gpu_layers, 99); + assert_eq!(args.n_predict, 32); + } + + #[test] + fn parses_generation_options_and_prompt() { + let args = parse_args([ + "-m", + "model.gguf", + "-n", + "8", + "-ngl", + "0", + "a", + "short", + "prompt", + ]) + .unwrap(); + + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "a short prompt"); + assert_eq!(args.n_gpu_layers, 0); + assert_eq!(args.n_predict, 8); + } + + #[test] + fn rejects_missing_model() { + assert_eq!( + parse_args(["-n", "8"]).unwrap_err(), + ParseError::MissingModel + ); + } + + #[test] + fn rejects_missing_option_value() { + assert_eq!( + parse_args(["-m", "model.gguf", "-n"]).unwrap_err(), + ParseError::MissingValue("-n") + ); + } + + #[test] + fn rejects_invalid_integer() { + assert_eq!( + parse_args(["-m", "model.gguf", "-ngl", "gpu"]).unwrap_err(), + ParseError::InvalidInteger("-ngl", "gpu".to_string()) + ); + } +} diff --git a/examples/simple-rust/src/main.rs b/examples/simple-rust/src/main.rs new file mode 100644 index 000000000000..1c68933bbba4 --- /dev/null +++ b/examples/simple-rust/src/main.rs @@ -0,0 +1,240 @@ +use llama_simple_rust::ffi; +use llama_simple_rust::{parse_args, Args}; +use std::env; +use std::ffi::CString; +use std::io::{self, Write}; +use std::ptr; + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Sampler(*mut ffi::llama_sampler); + +impl Drop for Sampler { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_sampler_free(self.0); + } + } + } +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m model.gguf [-n n_predict] [-ngl n_gpu_layers] [prompt]"); + eprintln!(); +} + +fn token_to_piece( + vocab: *const ffi::llama_vocab, + token: ffi::llama_token, +) -> Result { + let mut buf = vec![0_i8; 128]; + let n = unsafe { + ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), buf.len() as i32, 0, true) + }; + + if n < 0 { + return Err("failed to convert token to piece".to_string()); + } + + let bytes = buf[..n as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +fn run(args: Args) -> Result<(), String> { + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + let prompt = CString::new(args.prompt.as_str()) + .map_err(|_| "prompt contains an interior NUL byte".to_string())?; + + unsafe { + ffi::ggml_backend_load_all(); + } + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + + let n_prompt = unsafe { + -ffi::llama_tokenize( + vocab, + prompt.as_ptr(), + args.prompt.len() as i32, + ptr::null_mut(), + 0, + true, + true, + ) + }; + + if n_prompt <= 0 { + return Err("failed to size prompt tokenization".to_string()); + } + + let mut prompt_tokens = vec![0_i32; n_prompt as usize]; + let n_tokenized = unsafe { + ffi::llama_tokenize( + vocab, + prompt.as_ptr(), + args.prompt.len() as i32, + prompt_tokens.as_mut_ptr(), + prompt_tokens.len() as i32, + true, + true, + ) + }; + + if n_tokenized < 0 { + return Err("failed to tokenize the prompt".to_string()); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = (n_prompt + args.n_predict - 1).max(1) as u32; + ctx_params.n_batch = n_prompt as u32; + ctx_params.no_perf = false; + + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create the llama_context".to_string()); + } + + let mut sampler_params = unsafe { ffi::llama_sampler_chain_default_params() }; + sampler_params.no_perf = false; + let sampler = Sampler(unsafe { ffi::llama_sampler_chain_init(sampler_params) }); + if sampler.0.is_null() { + return Err("failed to create sampler chain".to_string()); + } + + let greedy = unsafe { ffi::llama_sampler_init_greedy() }; + if greedy.is_null() { + return Err("failed to create greedy sampler".to_string()); + } + unsafe { + ffi::llama_sampler_chain_add(sampler.0, greedy); + } + + for &id in &prompt_tokens { + print!("{}", token_to_piece(vocab, id)?); + } + io::stdout().flush().map_err(|err| err.to_string())?; + + let mut batch = + unsafe { ffi::llama_batch_get_one(prompt_tokens.as_mut_ptr(), prompt_tokens.len() as i32) }; + let mut decoder_start_token: ffi::llama_token; + + if unsafe { ffi::llama_model_has_encoder(model.0) } { + if unsafe { ffi::llama_encode(ctx.0, batch) } != 0 { + return Err("failed to eval encoder".to_string()); + } + + decoder_start_token = unsafe { ffi::llama_model_decoder_start_token(model.0) }; + if decoder_start_token == ffi::LLAMA_TOKEN_NULL { + decoder_start_token = unsafe { ffi::llama_vocab_bos(vocab) }; + } + + batch = unsafe { ffi::llama_batch_get_one(&mut decoder_start_token, 1) }; + } + + let t_main_start = unsafe { ffi::ggml_time_us() }; + let mut n_decode = 0; + let mut new_token_id = 0_i32; + let mut n_pos = 0; + + while n_pos + batch.n_tokens < n_prompt + args.n_predict { + let ret = unsafe { ffi::llama_decode(ctx.0, batch) }; + if ret != 0 { + return Err(format!("failed to eval, return code {ret}")); + } + + n_pos += batch.n_tokens; + new_token_id = unsafe { ffi::llama_sampler_sample(sampler.0, ctx.0, -1) }; + + if unsafe { ffi::llama_vocab_is_eog(vocab, new_token_id) } { + break; + } + + print!("{}", token_to_piece(vocab, new_token_id)?); + io::stdout().flush().map_err(|err| err.to_string())?; + + batch = unsafe { ffi::llama_batch_get_one(&mut new_token_id, 1) }; + n_decode += 1; + } + + println!(); + + let t_main_end = unsafe { ffi::ggml_time_us() }; + let elapsed = (t_main_end - t_main_start) as f32 / 1_000_000.0; + eprintln!( + "main: decoded {n_decode} tokens in {elapsed:.2} s, speed: {:.2} t/s", + n_decode as f32 / elapsed.max(f32::EPSILON) + ); + eprintln!(); + unsafe { + ffi::llama_perf_sampler_print(sampler.0); + ffi::llama_perf_context_print(ctx.0); + } + eprintln!(); + + let _ = new_token_id; + + Ok(()) +} + +fn main() { + let mut argv = env::args(); + let program = argv + .next() + .unwrap_or_else(|| "llama-simple-rust".to_string()); + + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} diff --git a/examples/simple/CMakeLists.txt b/examples/simple/CMakeLists.txt index 104ecabfd723..f75c9ef4dd31 100644 --- a/examples/simple/CMakeLists.txt +++ b/examples/simple/CMakeLists.txt @@ -1,5 +1,24 @@ set(TARGET llama-simple) -add_executable(${TARGET} simple.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(SIMPLE_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/Cargo.toml) +set(SIMPLE_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/simple) +set(SIMPLE_RUST_BIN ${SIMPLE_RUST_TARGET_DIR}/debug/llama-simple${CMAKE_EXECUTABLE_SUFFIX}) +set(SIMPLE_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${SIMPLE_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${SIMPLE_RUST_MANIFEST} --target-dir ${SIMPLE_RUST_TARGET_DIR} --bin llama-simple + COMMAND ${CMAKE_COMMAND} -E copy ${SIMPLE_RUST_BIN} ${SIMPLE_RUST_OUTPUT} + DEPENDS + llama + ${SIMPLE_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/main.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${SIMPLE_RUST_OUTPUT}) +install(PROGRAMS ${SIMPLE_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/simple/simple.cpp b/examples/simple/simple.cpp deleted file mode 100644 index 9f0a25d713f4..000000000000 --- a/examples/simple/simple.cpp +++ /dev/null @@ -1,223 +0,0 @@ -#include "llama.h" -#include -#include -#include -#include -#include - -static void print_usage(int, char ** argv) { - printf("\nexample usage:\n"); - printf("\n %s -m model.gguf [-n n_predict] [-ngl n_gpu_layers] [prompt]\n", argv[0]); - printf("\n"); -} - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - // path to the model gguf file - std::string model_path; - // prompt to generate text from - std::string prompt = "Hello my name is"; - // number of layers to offload to the GPU - int ngl = 99; - // number of tokens to predict - int n_predict = 32; - - // parse command line arguments - - { - int i = 1; - for (; i < argc; i++) { - if (strcmp(argv[i], "-m") == 0) { - if (i + 1 < argc) { - model_path = argv[++i]; - } else { - print_usage(argc, argv); - return 1; - } - } else if (strcmp(argv[i], "-n") == 0) { - if (i + 1 < argc) { - try { - n_predict = std::stoi(argv[++i]); - } catch (...) { - print_usage(argc, argv); - return 1; - } - } else { - print_usage(argc, argv); - return 1; - } - } else if (strcmp(argv[i], "-ngl") == 0) { - if (i + 1 < argc) { - try { - ngl = std::stoi(argv[++i]); - } catch (...) { - print_usage(argc, argv); - return 1; - } - } else { - print_usage(argc, argv); - return 1; - } - } else { - // prompt starts here - break; - } - } - if (model_path.empty()) { - print_usage(argc, argv); - return 1; - } - if (i < argc) { - prompt = argv[i++]; - for (; i < argc; i++) { - prompt += " "; - prompt += argv[i]; - } - } - } - - // load dynamic backends - - ggml_backend_load_all(); - - // initialize the model - - llama_model_params model_params = llama_model_default_params(); - model_params.n_gpu_layers = ngl; - - llama_model * model = llama_model_load_from_file(model_path.c_str(), model_params); - - if (model == NULL) { - fprintf(stderr , "%s: error: unable to load model\n" , __func__); - return 1; - } - - const llama_vocab * vocab = llama_model_get_vocab(model); - // tokenize the prompt - - // find the number of tokens in the prompt - const int n_prompt = -llama_tokenize(vocab, prompt.c_str(), prompt.size(), NULL, 0, true, true); - - // allocate space for the tokens and tokenize the prompt - std::vector prompt_tokens(n_prompt); - if (llama_tokenize(vocab, prompt.c_str(), prompt.size(), prompt_tokens.data(), prompt_tokens.size(), true, true) < 0) { - fprintf(stderr, "%s: error: failed to tokenize the prompt\n", __func__); - return 1; - } - - // initialize the context - - llama_context_params ctx_params = llama_context_default_params(); - // n_ctx is the context size - ctx_params.n_ctx = n_prompt + n_predict - 1; - // n_batch is the maximum number of tokens that can be processed in a single call to llama_decode - ctx_params.n_batch = n_prompt; - // enable performance counters - ctx_params.no_perf = false; - - llama_context * ctx = llama_init_from_model(model, ctx_params); - - if (ctx == NULL) { - fprintf(stderr , "%s: error: failed to create the llama_context\n" , __func__); - return 1; - } - - // initialize the sampler - - auto sparams = llama_sampler_chain_default_params(); - sparams.no_perf = false; - llama_sampler * smpl = llama_sampler_chain_init(sparams); - - llama_sampler_chain_add(smpl, llama_sampler_init_greedy()); - - // print the prompt token-by-token - - for (auto id : prompt_tokens) { - char buf[128]; - int n = llama_token_to_piece(vocab, id, buf, sizeof(buf), 0, true); - if (n < 0) { - fprintf(stderr, "%s: error: failed to convert token to piece\n", __func__); - return 1; - } - std::string s(buf, n); - printf("%s", s.c_str()); - } - - // prepare a batch for the prompt - - llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size()); - - if (llama_model_has_encoder(model)) { - if (llama_encode(ctx, batch)) { - fprintf(stderr, "%s : failed to eval\n", __func__); - return 1; - } - - llama_token decoder_start_token_id = llama_model_decoder_start_token(model); - if (decoder_start_token_id == LLAMA_TOKEN_NULL) { - decoder_start_token_id = llama_vocab_bos(vocab); - } - - batch = llama_batch_get_one(&decoder_start_token_id, 1); - } - - // main loop - - const auto t_main_start = ggml_time_us(); - int n_decode = 0; - llama_token new_token_id; - - for (int n_pos = 0; n_pos + batch.n_tokens < n_prompt + n_predict; ) { - // evaluate the current batch with the transformer model - if (llama_decode(ctx, batch)) { - fprintf(stderr, "%s : failed to eval, return code %d\n", __func__, 1); - return 1; - } - - n_pos += batch.n_tokens; - - // sample the next token - { - new_token_id = llama_sampler_sample(smpl, ctx, -1); - - // is it an end of generation? - if (llama_vocab_is_eog(vocab, new_token_id)) { - break; - } - - char buf[128]; - int n = llama_token_to_piece(vocab, new_token_id, buf, sizeof(buf), 0, true); - if (n < 0) { - fprintf(stderr, "%s: error: failed to convert token to piece\n", __func__); - return 1; - } - std::string s(buf, n); - printf("%s", s.c_str()); - fflush(stdout); - - // prepare the next batch with the sampled token - batch = llama_batch_get_one(&new_token_id, 1); - - n_decode += 1; - } - } - - printf("\n"); - - const auto t_main_end = ggml_time_us(); - - fprintf(stderr, "%s: decoded %d tokens in %.2f s, speed: %.2f t/s\n", - __func__, n_decode, (t_main_end - t_main_start) / 1000000.0f, n_decode / ((t_main_end - t_main_start) / 1000000.0f)); - - fprintf(stderr, "\n"); - llama_perf_sampler_print(smpl); - llama_perf_context_print(ctx); - fprintf(stderr, "\n"); - - llama_sampler_free(smpl); - llama_free(ctx); - llama_model_free(model); - - return 0; -} diff --git a/examples/speculative-rust/Cargo.toml b/examples/speculative-rust/Cargo.toml new file mode 100644 index 000000000000..e245b4f1b804 --- /dev/null +++ b/examples/speculative-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-speculative-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-speculative" +path = "src/main.rs" diff --git a/examples/speculative-rust/build.rs b/examples/speculative-rust/build.rs new file mode 100644 index 000000000000..959ac73146f8 --- /dev/null +++ b/examples/speculative-rust/build.rs @@ -0,0 +1,45 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); + + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-speculative"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/speculative-rust/src/main.rs b/examples/speculative-rust/src/main.rs new file mode 100644 index 000000000000..cce82419ceba --- /dev/null +++ b/examples/speculative-rust/src/main.rs @@ -0,0 +1 @@ +include!("../../speculative-simple-rust/src/main.rs"); diff --git a/examples/speculative-simple-rust/Cargo.toml b/examples/speculative-simple-rust/Cargo.toml new file mode 100644 index 000000000000..7ad9154d5a75 --- /dev/null +++ b/examples/speculative-simple-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-speculative-simple-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../simple-rust" } + +[[bin]] +name = "llama-speculative-simple" +path = "src/main.rs" diff --git a/examples/speculative-simple-rust/build.rs b/examples/speculative-simple-rust/build.rs new file mode 100644 index 000000000000..f6feadc3c887 --- /dev/null +++ b/examples/speculative-simple-rust/build.rs @@ -0,0 +1,45 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); + + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-speculative-simple"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/speculative-simple-rust/src/main.rs b/examples/speculative-simple-rust/src/main.rs new file mode 100644 index 000000000000..5e88d11ecb73 --- /dev/null +++ b/examples/speculative-simple-rust/src/main.rs @@ -0,0 +1,663 @@ +use llama_simple_rust::ffi; +use std::env; +use std::ffi::CString; +use std::io::{self, Write}; +use std::ptr; + +#[derive(Debug, Clone, PartialEq)] +pub struct Args { + pub model_path: String, + pub draft_model_path: String, + pub prompt: String, + pub n_predict: i32, + pub n_draft: usize, + pub n_ctx: u32, + pub n_batch: u32, + pub n_gpu_layers: i32, + pub draft_n_gpu_layers: i32, + pub top_k: i32, + pub top_p: f32, + pub temp: f32, + pub seed: u32, + pub use_color: bool, +} + +impl Default for Args { + fn default() -> Self { + Self { + model_path: String::new(), + draft_model_path: String::new(), + prompt: "Hello my name is".to_string(), + n_predict: 32, + n_draft: 16, + n_ctx: 512, + n_batch: 512, + n_gpu_layers: 99, + draft_n_gpu_layers: -1, + top_k: 40, + top_p: 0.95, + temp: 0.8, + seed: ffi::LLAMA_DEFAULT_SEED, + use_color: false, + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub enum ParseError { + MissingModel, + MissingDraftModel, + MissingValue(String), + InvalidInteger(String, String), + InvalidFloat(String, String), + InvalidValue(String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m/--model model.gguf"), + ParseError::MissingDraftModel => { + write!(f, "missing required -md/--model-draft draft.gguf") + } + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + ParseError::InvalidFloat(flag, value) => write!(f, "invalid float for {flag}: {value}"), + ParseError::InvalidValue(message) => write!(f, "{message}"), + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut parsed = Args::default(); + let mut iter = args.into_iter().map(Into::into); + + while let Some(arg) = iter.next() { + match arg.as_str() { + "-m" | "--model" => parsed.model_path = take(&mut iter, &arg)?, + "-md" | "--model-draft" | "--spec-draft-model" => { + parsed.draft_model_path = take(&mut iter, &arg)? + } + "-p" | "--prompt" => parsed.prompt = take(&mut iter, &arg)?, + "-f" | "--file" | "--prompt-file" => { + let path = take(&mut iter, &arg)?; + parsed.prompt = std::fs::read_to_string(&path).map_err(|err| { + ParseError::InvalidValue(format!("failed to read prompt file {path}: {err}")) + })?; + } + "-n" | "--n-predict" => parsed.n_predict = parse_i32(&mut iter, &arg)?, + "--spec-draft-n-max" => parsed.n_draft = parse_usize(&mut iter, &arg)?, + "-c" | "--ctx-size" => parsed.n_ctx = parse_u32(&mut iter, &arg)?, + "-b" | "--batch-size" => parsed.n_batch = parse_u32(&mut iter, &arg)?, + "-ngl" | "--gpu-layers" => parsed.n_gpu_layers = parse_i32(&mut iter, &arg)?, + "-ngld" | "--gpu-layers-draft" | "--n-gpu-layers-draft" => { + parsed.draft_n_gpu_layers = parse_i32(&mut iter, &arg)? + } + "--top-k" => parsed.top_k = parse_i32(&mut iter, &arg)?, + "--top-p" => parsed.top_p = parse_f32(&mut iter, &arg)?, + "--temp" => parsed.temp = parse_f32(&mut iter, &arg)?, + "-s" | "--seed" => parsed.seed = parse_u32(&mut iter, &arg)?, + "--color" => parsed.use_color = true, + "--no-color" => parsed.use_color = false, + "-h" | "--help" => return Err(ParseError::MissingModel), + other => { + let mut prompt = vec![other.to_string()]; + prompt.extend(iter); + parsed.prompt = prompt.join(" "); + break; + } + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + if parsed.draft_model_path.is_empty() { + return Err(ParseError::MissingDraftModel); + } + if parsed.n_predict < -1 { + return Err(ParseError::InvalidValue( + "--n-predict must be >= -1".to_string(), + )); + } + if parsed.n_ctx == 0 || parsed.n_batch == 0 { + return Err(ParseError::InvalidValue( + "context and batch sizes must be positive".to_string(), + )); + } + Ok(parsed) +} + +fn take(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + iter.next() + .ok_or_else(|| ParseError::MissingValue(flag.to_string())) +} + +fn parse_i32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_u32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_usize(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_f32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidFloat(flag.to_string(), value)) +} + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Batch(ffi::llama_batch); + +impl Drop for Batch { + fn drop(&mut self) { + unsafe { + ffi::llama_batch_free(self.0); + } + } +} + +struct Sampler(*mut ffi::llama_sampler); + +impl Drop for Sampler { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_sampler_free(self.0); + } + } + } +} + +struct Backend; + +impl Backend { + fn init() -> Self { + unsafe { + ffi::ggml_backend_load_all(); + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + Self + } +} + +impl Drop for Backend { + fn drop(&mut self) { + unsafe { + ffi::llama_backend_free(); + } + } +} + +fn sampler_init(args: &Args) -> Sampler { + unsafe { + let chain = ffi::llama_sampler_chain_init(ffi::llama_sampler_chain_default_params()); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_top_k(args.top_k)); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_top_p(args.top_p, 1)); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_temp(args.temp)); + ffi::llama_sampler_chain_add(chain, ffi::llama_sampler_init_dist(args.seed)); + Sampler(chain) + } +} + +fn tokenize( + vocab: *const ffi::llama_vocab, + text: &str, + add_special: bool, + parse_special: bool, +) -> Result, String> { + let text_c = + CString::new(text).map_err(|_| "prompt contains an interior NUL byte".to_string())?; + let n_tokens = unsafe { + -ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + ptr::null_mut(), + 0, + add_special, + parse_special, + ) + }; + if n_tokens <= 0 { + return Err("failed to size tokenization".to_string()); + } + let mut tokens = vec![0_i32; n_tokens as usize]; + let n = unsafe { + ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + add_special, + parse_special, + ) + }; + if n < 0 { + return Err("failed to tokenize prompt".to_string()); + } + tokens.truncate(n as usize); + Ok(tokens) +} + +fn token_to_piece( + vocab: *const ffi::llama_vocab, + token: ffi::llama_token, +) -> Result { + let mut buf = vec![0_i8; 128]; + let n = unsafe { + ffi::llama_token_to_piece(vocab, token, buf.as_mut_ptr(), buf.len() as i32, 0, true) + }; + if n < 0 { + return Err("failed to convert token to piece".to_string()); + } + let bytes = buf[..n as usize] + .iter() + .map(|&c| c as u8) + .collect::>(); + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +fn batch_clear(batch: &mut ffi::llama_batch) { + batch.n_tokens = 0; +} + +unsafe fn batch_add( + batch: &mut ffi::llama_batch, + token: ffi::llama_token, + pos: ffi::llama_pos, + seq_ids: &[ffi::llama_seq_id], + logits: bool, +) { + let i = batch.n_tokens as isize; + *batch.token.offset(i) = token; + *batch.pos.offset(i) = pos; + *batch.n_seq_id.offset(i) = seq_ids.len() as i32; + let seq_slot = *batch.seq_id.offset(i); + for (j, seq_id) in seq_ids.iter().enumerate() { + *seq_slot.add(j) = *seq_id; + } + *batch.logits.offset(i) = if logits { 1 } else { 0 }; + batch.n_tokens += 1; +} + +fn decode_batch_get_one( + ctx: *mut ffi::llama_context, + tokens: &mut [ffi::llama_token], +) -> Result<(), String> { + if tokens.is_empty() { + return Ok(()); + } + let batch = unsafe { ffi::llama_batch_get_one(tokens.as_mut_ptr(), tokens.len() as i32) }; + if unsafe { ffi::llama_decode(ctx, batch) } != 0 { + return Err("llama_decode failed".to_string()); + } + Ok(()) +} + +fn rebuild_draft_context( + ctx: *mut ffi::llama_context, + tokens: &[ffi::llama_token], +) -> Result<(), String> { + unsafe { + ffi::llama_memory_clear(ffi::llama_get_memory(ctx), true); + } + let mut working = tokens.to_vec(); + decode_batch_get_one(ctx, &mut working) +} + +fn draft_tokens_from_model( + ctx: *mut ffi::llama_context, + sampler: *mut ffi::llama_sampler, + n_draft: usize, +) -> Result, String> { + let mut draft = Vec::new(); + for _ in 0..n_draft { + let token = unsafe { ffi::llama_sampler_sample(sampler, ctx, -1) }; + unsafe { + ffi::llama_sampler_accept(sampler, token); + } + draft.push(token); + let mut one = [token]; + decode_batch_get_one(ctx, &mut one)?; + } + Ok(draft) +} + +fn run(args: Args) -> Result<(), String> { + let _backend = Backend::init(); + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + let draft_model_path = CString::new(args.draft_model_path.as_str()) + .map_err(|_| "draft model path contains an interior NUL byte".to_string())?; + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + let model_tgt = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model_tgt.0.is_null() { + return Err("unable to load target model".to_string()); + } + + let mut draft_params = unsafe { ffi::llama_model_default_params() }; + draft_params.n_gpu_layers = args.draft_n_gpu_layers; + let model_dft = + Model(unsafe { ffi::llama_model_load_from_file(draft_model_path.as_ptr(), draft_params) }); + if model_dft.0.is_null() { + return Err("unable to load draft model".to_string()); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = args.n_ctx; + ctx_params.n_batch = args.n_batch; + ctx_params.n_ubatch = args.n_batch; + let ctx_tgt = Context(unsafe { ffi::llama_init_from_model(model_tgt.0, ctx_params) }); + if ctx_tgt.0.is_null() { + return Err("failed to create target context".to_string()); + } + + let mut ctx_dft_params = unsafe { ffi::llama_context_default_params() }; + ctx_dft_params.n_ctx = args.n_ctx; + ctx_dft_params.n_batch = args.n_batch; + ctx_dft_params.n_ubatch = args.n_batch; + let ctx_dft = Context(unsafe { ffi::llama_init_from_model(model_dft.0, ctx_dft_params) }); + if ctx_dft.0.is_null() { + return Err("failed to create draft context".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model_tgt.0) }; + if vocab.is_null() { + return Err("failed to get target vocabulary".to_string()); + } + + let input = tokenize(vocab, &args.prompt, true, true)?; + if input.is_empty() { + return Err("prompt produced no tokens".to_string()); + } + if unsafe { ffi::llama_n_ctx(ctx_tgt.0) as usize } < input.len() { + return Err(format!( + "the prompt exceeds the context size ({} tokens, ctx {})", + input.len(), + unsafe { ffi::llama_n_ctx(ctx_tgt.0) } + )); + } + + println!(); + for token in &input { + print!("{}", token_to_piece(vocab, *token)?); + } + io::stdout().flush().map_err(|err| err.to_string())?; + + let n_input = input.len(); + let t_enc_start = unsafe { ffi::ggml_time_us() }; + let sampler_tgt = sampler_init(&args); + let sampler_dft = sampler_init(&args); + let mut prompt_tgt = input[..input.len() - 1].to_vec(); + if !prompt_tgt.is_empty() { + decode_batch_get_one(ctx_tgt.0, &mut prompt_tgt)?; + } + let mut id_last = *input.last().unwrap(); + let mut n_past = prompt_tgt.len() as i32; + let mut batch_tgt = Batch(unsafe { ffi::llama_batch_init(args.n_batch as i32, 0, 1) }); + let t_enc_end = unsafe { ffi::ggml_time_us() }; + + let mut n_predict = 0_i32; + let mut n_drafted = 0_i32; + let mut n_accept = 0_i32; + let mut has_eos = false; + let t_dec_start = unsafe { ffi::ggml_time_us() }; + + loop { + let mut draft_prefix = prompt_tgt.clone(); + draft_prefix.push(id_last); + rebuild_draft_context(ctx_dft.0, &draft_prefix)?; + let draft = draft_tokens_from_model(ctx_dft.0, sampler_dft.0, args.n_draft)?; + n_drafted += draft.len() as i32; + + batch_clear(&mut batch_tgt.0); + unsafe { + batch_add(&mut batch_tgt.0, id_last, n_past, &[0], true); + for (i, token) in draft.iter().enumerate() { + batch_add(&mut batch_tgt.0, *token, n_past + 1 + i as i32, &[0], true); + } + } + if unsafe { ffi::llama_decode(ctx_tgt.0, batch_tgt.0) } != 0 { + return Err("failed to decode target speculative batch".to_string()); + } + + let mut ids = Vec::new(); + for i in 0..=draft.len() { + let sampled = unsafe { ffi::llama_sampler_sample(sampler_tgt.0, ctx_tgt.0, i as i32) }; + unsafe { + ffi::llama_sampler_accept(sampler_tgt.0, sampled); + } + ids.push(sampled); + if i >= draft.len() || sampled != draft[i] { + break; + } + } + + let accepted_draft = ids.len().saturating_sub(1); + n_accept += accepted_draft as i32; + n_predict += ids.len() as i32; + n_past += ids.len() as i32; + + for (i, token) in ids.iter().enumerate() { + prompt_tgt.push(id_last); + id_last = *token; + if unsafe { ffi::llama_vocab_is_eog(vocab, id_last) } { + has_eos = true; + break; + } + let token_str = token_to_piece(vocab, id_last)?; + if args.use_color && i + 1 < ids.len() { + print!("\x1b[36m{token_str}\x1b[37m"); + } else { + print!("{token_str}"); + } + } + io::stdout().flush().map_err(|err| err.to_string())?; + + unsafe { + ffi::llama_memory_seq_rm(ffi::llama_get_memory(ctx_tgt.0), 0, n_past, -1); + } + + if (args.n_predict >= 0 && n_predict > args.n_predict) || has_eos { + break; + } + } + + let t_dec_end = unsafe { ffi::ggml_time_us() }; + println!("\n"); + eprintln!( + "encoded {n_input:4} tokens in {:8.3} seconds, speed: {:8.3} t/s", + (t_enc_end - t_enc_start) as f64 / 1e6, + n_input as f64 / ((t_enc_end - t_enc_start) as f64 / 1e6) + ); + eprintln!( + "decoded {n_predict:4} tokens in {:8.3} seconds, speed: {:8.3} t/s", + (t_dec_end - t_dec_start) as f64 / 1e6, + n_predict as f64 / ((t_dec_end - t_dec_start) as f64 / 1e6) + ); + eprintln!("\nn_draft = {}", args.n_draft); + eprintln!("n_predict = {n_predict}"); + eprintln!("n_drafted = {n_drafted}"); + eprintln!("n_accept = {n_accept}"); + if n_drafted > 0 { + eprintln!( + "accept = {:.3}%", + 100.0 * n_accept as f64 / n_drafted as f64 + ); + } else { + eprintln!("accept = 0.000%"); + } + eprintln!("\ndraft:\n"); + unsafe { + ffi::llama_perf_sampler_print(sampler_dft.0); + } + eprintln!("\ntarget:\n"); + unsafe { + ffi::llama_perf_sampler_print(sampler_tgt.0); + ffi::llama_perf_context_print(ctx_tgt.0); + } + Ok(()) +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m target.gguf -md draft.gguf -p \"Hello\""); + eprintln!(); +} + +fn main() { + let mut argv = env::args(); + let program = argv + .next() + .unwrap_or_else(|| "llama-speculative-simple".to_string()); + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_speculative_simple_options() { + let args = parse_args([ + "-m", + "target.gguf", + "-md", + "draft.gguf", + "-p", + "hello", + "-n", + "8", + "--spec-draft-n-max", + "3", + "-c", + "256", + "-b", + "128", + "-ngl", + "0", + "-ngld", + "1", + "--top-k", + "10", + "--top-p", + "0.8", + "--temp", + "0.5", + "-s", + "42", + "--color", + ]) + .unwrap(); + assert_eq!(args.model_path, "target.gguf"); + assert_eq!(args.draft_model_path, "draft.gguf"); + assert_eq!(args.prompt, "hello"); + assert_eq!(args.n_predict, 8); + assert_eq!(args.n_draft, 3); + assert_eq!(args.n_ctx, 256); + assert_eq!(args.n_batch, 128); + assert_eq!(args.n_gpu_layers, 0); + assert_eq!(args.draft_n_gpu_layers, 1); + assert_eq!(args.top_k, 10); + assert_eq!(args.top_p, 0.8); + assert_eq!(args.temp, 0.5); + assert_eq!(args.seed, 42); + assert!(args.use_color); + } + + #[test] + fn rejects_missing_draft_model() { + assert_eq!( + parse_args(["-m", "target.gguf"]).unwrap_err(), + ParseError::MissingDraftModel + ); + } + + #[test] + fn rejects_bad_predict_count() { + assert_eq!( + parse_args(["-m", "target.gguf", "-md", "draft.gguf", "-n", "-2"]).unwrap_err(), + ParseError::InvalidValue("--n-predict must be >= -1".to_string()) + ); + } +} diff --git a/examples/speculative-simple/CMakeLists.txt b/examples/speculative-simple/CMakeLists.txt index 5ef3b4131f2a..8a4874279d04 100644 --- a/examples/speculative-simple/CMakeLists.txt +++ b/examples/speculative-simple/CMakeLists.txt @@ -1,5 +1,25 @@ set(TARGET llama-speculative-simple) -add_executable(${TARGET} speculative-simple.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(SPECULATIVE_SIMPLE_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../speculative-simple-rust/Cargo.toml) +set(SPECULATIVE_SIMPLE_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/speculative-simple) +set(SPECULATIVE_SIMPLE_RUST_BIN ${SPECULATIVE_SIMPLE_RUST_TARGET_DIR}/debug/llama-speculative-simple${CMAKE_EXECUTABLE_SUFFIX}) +set(SPECULATIVE_SIMPLE_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${SPECULATIVE_SIMPLE_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${SPECULATIVE_SIMPLE_RUST_MANIFEST} --target-dir ${SPECULATIVE_SIMPLE_RUST_TARGET_DIR} --bin llama-speculative-simple + COMMAND ${CMAKE_COMMAND} -E copy ${SPECULATIVE_SIMPLE_RUST_BIN} ${SPECULATIVE_SIMPLE_RUST_OUTPUT} + DEPENDS + llama + ${SPECULATIVE_SIMPLE_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../speculative-simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../speculative-simple-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${SPECULATIVE_SIMPLE_RUST_OUTPUT}) +install(PROGRAMS ${SPECULATIVE_SIMPLE_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp deleted file mode 100644 index 5b61b62a1bf8..000000000000 --- a/examples/speculative-simple/speculative-simple.cpp +++ /dev/null @@ -1,348 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "sampling.h" -#include "speculative.h" -#include "log.h" -#include "llama.h" - -#include -#include -#include -#include -#include -#include -#include - -struct spec_checkpoint { - int64_t n_tokens = 0; - - std::vector data; - - size_t size() const { - return data.size(); - } - - bool empty() const { - return data.empty(); - } -}; - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SPECULATIVE)) { - return 1; - } - - if (params.n_predict < -1) { - LOG_ERR("%s: --n-predict must be >= -1\n", __func__); - return 1; - } - - if (params.speculative.draft.mparams.path.empty()) { - LOG_ERR("%s: --model-draft is required\n", __func__); - return 1; - } - - // init llama.cpp - llama_backend_init(); - llama_numa_init(params.numa); - - llama_model * model_tgt = NULL; - - llama_context * ctx_tgt = NULL; - - // load the target model - auto llama_init_tgt = common_init_from_params(params); - - model_tgt = llama_init_tgt->model(); - ctx_tgt = llama_init_tgt->context(); - - // check if the context supports partial sequence removal - const auto ctx_seq_rm = common_context_can_seq_rm(ctx_tgt); - const bool use_ckpt = (ctx_seq_rm == COMMON_CONTEXT_SEQ_RM_TYPE_FULL); - - if (use_ckpt) { - LOG_INF("speculative decoding will use checkpoints (context does not support partial sequence removal)\n"); - } - - const llama_vocab * vocab = llama_model_get_vocab(model_tgt); - - // load the draft model - llama_model_ptr model_dft; - - // TODO: simplify this logic - { - const auto & params_spec = params.speculative.draft; - - auto params_dft = params; - - params_dft.n_parallel = 1; - params_dft.n_ctx = params_spec.n_ctx; - params_dft.n_batch = llama_n_ctx_seq(ctx_tgt); - params_dft.devices = params_spec.devices; - params_dft.model = params_spec.mparams; - params_dft.n_gpu_layers = params_spec.n_gpu_layers; - - if (params_spec.cpuparams.n_threads > 0) { - params_dft.cpuparams.n_threads = params.speculative.draft.cpuparams.n_threads; - params_dft.cpuparams_batch.n_threads = params.speculative.draft.cpuparams_batch.n_threads; - } - - params_dft.tensor_buft_overrides = params.speculative.draft.tensor_buft_overrides; - - auto mparams_dft = common_model_params_to_llama(params_dft); - - model_dft.reset(llama_model_load_from_file(params_dft.model.path.c_str(), mparams_dft)); - if (model_dft == nullptr) { - LOG_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str()); - return 1; - } - - params.speculative.draft.model = model_dft.get(); - params.speculative.draft.cparams = common_context_params_to_llama(params_dft); - } - - // Tokenize the prompt - std::vector inp; - inp = common_tokenize(ctx_tgt, params.prompt, true, true); - - if (llama_n_ctx(ctx_tgt) < (uint32_t) inp.size()) { - LOG_ERR("%s: the prompt exceeds the context size (%d tokens, ctx %d)\n", __func__, (int) inp.size(), llama_n_ctx(ctx_tgt)); - - return 1; - } - - if (llama_n_batch(ctx_tgt) < (uint32_t) inp.size()) { - LOG_ERR("%s: the prompt exceeds the batch size (%d tokens, batch %d)\n", __func__, (int) inp.size(), llama_n_batch(ctx_tgt)); - - return 1; - } - - LOG("\n\n"); - - for (auto id : inp) { - LOG("%s", common_token_to_piece(ctx_tgt, id).c_str()); - } - - int n_predict = 0; - int n_drafted = 0; - int n_accept = 0; - - // used to determine end of generation - bool has_eos = false; - - // ================================================ - // everything until here is standard initialization - // the relevant stuff for speculative decoding starts here - - const auto t_enc_start = ggml_time_us(); - - // target model sampling context - common_sampler_ptr smpl(common_sampler_init(model_tgt, params.sampling)); - - // eval the prompt - llama_decode(ctx_tgt, llama_batch_get_one(inp.data(), inp.size() - 1)); - - // note: keep the last token separate! - llama_token id_last = inp.back(); - - // all tokens currently in the target context - llama_tokens prompt_tgt(inp.begin(), inp.end() - 1); - prompt_tgt.reserve(llama_n_ctx(ctx_tgt)); - - int n_past = inp.size() - 1; - - // init the speculator - const auto & params_spec = params.speculative; - - struct common_speculative * spec = common_speculative_init(params.speculative, ctx_tgt); - - common_speculative_begin(spec, prompt_tgt); - - llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1); - - size_t n_draft = 0; - - llama_tokens draft; - spec_checkpoint spec_ckpt; - - const auto t_enc_end = ggml_time_us(); - - const auto t_dec_start = ggml_time_us(); - - while (true) { - // generate or reuse draft tokens - // - // this is the most important part of the speculation. the more probable tokens that are provided here - // the better the performance will be. in theory, this computation can be performed asynchronously and even - // offloaded to a remote device. it doesn't even have to be based on an LLM. instead, it can provide tokens - // from a cache or lookup tables. - // - if (draft.empty()) { - // generate a new draft - draft = common_speculative_draft(spec, params_spec, prompt_tgt, id_last); - - // save the original draft size - n_draft = draft.size(); - - // save a checkpoint of the target context before evaluating the draft - // this allows us to restore the state if partial draft acceptance occurs - if (!draft.empty() && use_ckpt) { - const size_t ckpt_size = llama_state_seq_get_size_ext(ctx_tgt, 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - spec_ckpt.data.resize(ckpt_size); - - const size_t n = llama_state_seq_get_data_ext(ctx_tgt, spec_ckpt.data.data(), ckpt_size, 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - GGML_ASSERT(n == ckpt_size); - - spec_ckpt.n_tokens = (int64_t) prompt_tgt.size(); - LOG_DBG("created speculative checkpoint (n_tokens = %" PRId64 ", size = %.3f MiB)\n", - spec_ckpt.n_tokens, (float) spec_ckpt.data.size() / 1024 / 1024); - } - } else { - // we have a previous (partial) draft to reuse from checkpoint restoration - if (use_ckpt) { - GGML_ASSERT(!spec_ckpt.empty()); - } - } - - // always have a token to evaluate from before - id_last - common_batch_clear(batch_tgt); - common_batch_add (batch_tgt, id_last, n_past++, { 0 }, true); - - // evaluate the target model on [id_last, draft0, draft1, ..., draftN-1] - { - for (size_t i = 0; i < draft.size(); ++i) { - common_batch_add(batch_tgt, draft[i], n_past + i, { 0 }, true); - } - - //LOG_DBG("target batch: %s\n", string_from(ctx_tgt, batch_tgt).c_str()); - - llama_decode(ctx_tgt, batch_tgt); - } - - // only save the sampler sampler state if we use checkpoints - common_sampler_ptr smpl_save; - if (use_ckpt) { - smpl_save.reset(common_sampler_clone(smpl.get())); - } - - // sample from the full target batch and return the accepted tokens based on the target sampler - // - // for each token to be accepted, the sampler would have to sample that same token - // in such cases, instead of decoding the sampled token as we normally do, we simply continue with the - // available logits from the batch and sample the next token until we run out of logits or the sampler - // disagrees with the draft - // - auto ids = common_sampler_sample_and_accept_n(smpl.get(), ctx_tgt, draft); - - //LOG_DBG("ids: %s\n", string_from(ctx_tgt, ids).c_str()); - - GGML_ASSERT(ids.size() > 0); // there will always be at least one accepted token - - // check for partial draft acceptance: - // if the context doesn't support partial sequence removal, restore the checkpoint - // and make the accepted tokens the new partial draft for the next iteration - if (use_ckpt && ids.size() - 1 < draft.size()) { - LOG_DBG("partial acceptance: %zu < %zu, restoring checkpoint\n", ids.size() - 1, draft.size()); - - draft = std::move(ids); - - const size_t n = llama_state_seq_set_data_ext(ctx_tgt, spec_ckpt.data.data(), spec_ckpt.size(), 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - GGML_ASSERT(n == spec_ckpt.size()); - - llama_memory_seq_rm(llama_get_memory(ctx_tgt), 0, spec_ckpt.n_tokens, -1); - - prompt_tgt.resize(spec_ckpt.n_tokens); - smpl = std::move(smpl_save); - - n_past = (int) prompt_tgt.size(); - - continue; - } - - common_speculative_accept(spec, ids.size() - 1); - - // full acceptance: consume the draft and commit accepted tokens - n_past += ids.size() - 1; - n_drafted += n_draft; // note: we ignore the discarded small drafts - n_accept += ids.size() - 1; - n_predict += ids.size(); - - // process the accepted tokens and update contexts - // - // this is the standard token post-processing that we normally do - // in this case, we do it for a group of accepted tokens at once - // - for (size_t i = 0; i < ids.size(); ++i) { - prompt_tgt.push_back(id_last); - - id_last = ids[i]; - - if (llama_vocab_is_eog(vocab, id_last)) { - has_eos = true; - break; - } - - const std::string token_str = common_token_to_piece(ctx_tgt, id_last); - - if (params.use_color && i + 1 < ids.size()) { - LOG("\u001b[%dm%s\u001b[37m", (36 - 0 % 6), token_str.c_str()); - } else { - LOG("%s", token_str.c_str()); - } - } - - LOG_DBG("accepted %d/%d draft tokens, the last target token is: (%d)\n", (int) ids.size() - 1, (int) draft.size(), id_last); - - // clear the draft since it has been consumed - draft.clear(); - - { - LOG_DBG("clear kv cache from any extra tokens, n_past = %d\n", n_past); - - llama_memory_seq_rm(llama_get_memory(ctx_tgt), 0, n_past, -1); - } - - if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) { - break; - } - } - - auto t_dec_end = ggml_time_us(); - - const int n_input = inp.size(); - - LOG("\n\n"); - - LOG_INF("encoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_input, (t_enc_end - t_enc_start) / 1e6f, inp.size() / ((t_enc_end - t_enc_start) / 1e6f)); - LOG_INF("decoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_predict, (t_dec_end - t_dec_start) / 1e6f, n_predict / ((t_dec_end - t_dec_start) / 1e6f)); - - LOG_INF("\n"); - LOG_INF("n_draft = %d\n", params_spec.draft.n_max); - LOG_INF("n_predict = %d\n", n_predict); - LOG_INF("n_drafted = %d\n", n_drafted); - LOG_INF("n_accept = %d\n", n_accept); - LOG_INF("accept = %.3f%%\n", 100.0f * n_accept / n_drafted); - - LOG_INF("\n"); - LOG_INF("draft:\n\n"); - - LOG_INF("\n"); - LOG_INF("target:\n\n"); - common_perf_print(ctx_tgt, smpl.get()); - - llama_batch_free(batch_tgt); - - common_speculative_free(spec); - - llama_backend_free(); - - LOG("\n\n"); - - return 0; -} diff --git a/examples/speculative/CMakeLists.txt b/examples/speculative/CMakeLists.txt index b4e20c717a21..2d600f6e5718 100644 --- a/examples/speculative/CMakeLists.txt +++ b/examples/speculative/CMakeLists.txt @@ -1,5 +1,26 @@ set(TARGET llama-speculative) -add_executable(${TARGET} speculative.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(SPECULATIVE_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/../speculative-rust/Cargo.toml) +set(SPECULATIVE_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/speculative) +set(SPECULATIVE_RUST_BIN ${SPECULATIVE_RUST_TARGET_DIR}/debug/llama-speculative${CMAKE_EXECUTABLE_SUFFIX}) +set(SPECULATIVE_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${SPECULATIVE_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${SPECULATIVE_RUST_MANIFEST} --target-dir ${SPECULATIVE_RUST_TARGET_DIR} --bin llama-speculative + COMMAND ${CMAKE_COMMAND} -E copy ${SPECULATIVE_RUST_BIN} ${SPECULATIVE_RUST_OUTPUT} + DEPENDS + llama + ${SPECULATIVE_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/../speculative-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../speculative-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../speculative-simple-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${SPECULATIVE_RUST_OUTPUT}) +install(PROGRAMS ${SPECULATIVE_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp deleted file mode 100644 index f7fa5e30602f..000000000000 --- a/examples/speculative/speculative.cpp +++ /dev/null @@ -1,660 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "sampling.h" -#include "log.h" -#include "llama.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 -#define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 - -struct seq_draft { - bool active = false; - bool drafting = false; - bool skip = false; - - int i_batch_dft = 0; - std::vector i_batch_tgt; - - std::vector tokens; - std::vector> dists; - - struct common_sampler * smpl = nullptr; -}; - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - - // needed to get candidate probs even for temp <= 0.0 - params.sampling.n_probs = 128; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SPECULATIVE)) { - return 1; - } - - if (params.n_predict < -1) { - LOG_ERR("%s: --n-predict must be >= -1\n", __func__); - return 1; - } - - if (params.speculative.draft.mparams.path.empty()) { - LOG_ERR("%s: --model-draft is required\n", __func__); - return 1; - } - - // max number of parallel drafting sequences (i.e. tree branches) - const int n_seq_dft = params.n_parallel; - - // probability threshold for splitting a draft branch (only for n_seq_dft > 1) - const float p_draft_split = params.speculative.draft.p_split; - - std::default_random_engine rng(params.sampling.seed == LLAMA_DEFAULT_SEED ? std::random_device()() : params.sampling.seed); - std::uniform_real_distribution<> u_dist; - - // init llama.cpp - llama_backend_init(); - llama_numa_init(params.numa); - - llama_model * model_tgt = NULL; - llama_model * model_dft = NULL; - - llama_context * ctx_tgt = NULL; - llama_context * ctx_dft = NULL; - - // load the target model - auto llama_init_tgt = common_init_from_params(params); - - model_tgt = llama_init_tgt->model(); - ctx_tgt = llama_init_tgt->context(); - - // load the draft model - params.devices = params.speculative.draft.devices; - params.model = params.speculative.draft.mparams; - params.n_gpu_layers = params.speculative.draft.n_gpu_layers; - if (params.speculative.draft.cpuparams.n_threads > 0) { - params.cpuparams.n_threads = params.speculative.draft.cpuparams.n_threads; - } - - params.cpuparams_batch.n_threads = params.speculative.draft.cpuparams_batch.n_threads; - params.tensor_buft_overrides = params.speculative.draft.tensor_buft_overrides; - - auto llama_init_dft = common_init_from_params(params); - - model_dft = llama_init_dft->model(); - ctx_dft = llama_init_dft->context(); - - const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt); - const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft); - - const bool vocab_type_tgt = llama_vocab_type(vocab_tgt); - LOG_DBG("vocab_type tgt: %d\n", vocab_type_tgt); - - const bool vocab_type_dft = llama_vocab_type(vocab_dft); - LOG_DBG("vocab_type dft: %d\n", vocab_type_dft); - - if (vocab_type_tgt != vocab_type_dft) { - LOG_ERR("%s: draft model vocab type must match target model to use speculation but ", __func__); - LOG_ERR("vocab_type_dft = %d while vocab_type_tgt = %d\n", vocab_type_dft, vocab_type_tgt); - return 1; - } - - if (llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) || - (llama_vocab_get_add_bos(vocab_tgt) && llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft))) { - LOG_ERR("%s: draft model bos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n", - __func__, - llama_vocab_get_add_bos(vocab_tgt), llama_vocab_get_add_bos(vocab_dft), - llama_vocab_bos(vocab_tgt), llama_vocab_bos(vocab_dft)); - return 1; - } - - if (llama_vocab_get_add_eos(vocab_tgt) != llama_vocab_get_add_eos(vocab_dft) || - (llama_vocab_get_add_eos(vocab_tgt) && llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft))) { - LOG_ERR("%s: draft model eos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n", - __func__, - llama_vocab_get_add_eos(vocab_tgt), llama_vocab_get_add_eos(vocab_dft), - llama_vocab_eos(vocab_tgt), llama_vocab_eos(vocab_dft)); - return 1; - } - - { - const int n_vocab_tgt = llama_vocab_n_tokens(vocab_tgt); - const int n_vocab_dft = llama_vocab_n_tokens(vocab_dft); - const int vocab_diff = n_vocab_tgt > n_vocab_dft - ? n_vocab_tgt - n_vocab_dft - : n_vocab_dft - n_vocab_tgt; - - if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) { - LOG_ERR("%s: draft model vocab must closely match target model to use speculation but ", __func__); - LOG_ERR("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n", - n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE); - return 1; - } - - for (int i = SPEC_VOCAB_CHECK_START_TOKEN_ID; i < std::min(n_vocab_tgt, n_vocab_dft); ++i) { - const char * token_text_tgt = llama_vocab_get_text(vocab_tgt, i); - const char * token_text_dft = llama_vocab_get_text(vocab_dft, i); - - if (std::strcmp(token_text_tgt, token_text_dft) != 0) { - LOG_ERR("%s: draft model vocab must match target model to use speculation but ", __func__); - LOG_ERR("token %d content differs - target '%s', draft '%s'\n", i, - common_token_to_piece(vocab_tgt, i).c_str(), - common_token_to_piece(vocab_dft, i).c_str()); - return 1; - } - } - } - - auto * mem_tgt = llama_get_memory(ctx_tgt); - auto * mem_dft = llama_get_memory(ctx_dft); - - // Tokenize the prompt - std::vector inp; - inp = common_tokenize(ctx_tgt, params.prompt, true, true); - - const int max_context_size = llama_n_ctx(ctx_tgt); - const int max_tokens_list_size = max_context_size - 4; - - if ((int) inp.size() > max_tokens_list_size) { - LOG_ERR("%s: prompt too long (%d tokens, max %d)\n", __func__, (int) inp.size(), max_tokens_list_size); - return 1; - } - - LOG("\n\n"); - - for (auto id : inp) { - LOG("%s", common_token_to_piece(ctx_tgt, id).c_str()); - } - - const int n_input = inp.size(); - - const auto t_enc_start = ggml_time_us(); - - // eval the prompt with both models - llama_decode(ctx_tgt, llama_batch_get_one( inp.data(), n_input - 1)); - llama_decode(ctx_tgt, llama_batch_get_one(&inp.back(), 1)); - llama_decode(ctx_dft, llama_batch_get_one( inp.data(), n_input)); - - const auto t_enc_end = ggml_time_us(); - - // the 2 models should have the same vocab - //GGML_ASSERT(n_vocab == llama_vocab_n_tokens(model_dft)); - - // how many tokens to draft each time - int n_draft = params.speculative.draft.n_max; - - int n_predict = 0; - int n_drafted = 0; - int n_accept = 0; - - int n_past_tgt = inp.size(); - int n_past_dft = inp.size(); - - // used to determine end of generation - bool has_eos = false; - - // target model sampling context (reuse the llama_context's sampling instance) - struct common_sampler * smpl = common_sampler_init(model_tgt, params.sampling); - - // draft sequence data - std::vector drafts(n_seq_dft); - - for (int s = 0; s < n_seq_dft; ++s) { - // allocate llama_sampler for each draft sequence - drafts[s].smpl = common_sampler_init(model_dft, params.sampling); - } - - llama_batch batch_dft = llama_batch_init(llama_n_batch(ctx_dft), 0, 1); - llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, n_seq_dft); - - const auto t_dec_start = ggml_time_us(); - - // sample from the last token of the prompt - drafts[0].i_batch_tgt.resize(1); - drafts[0].i_batch_tgt[0] = 0; - - while (true) { - std::set active_seqs = {}; - - // print current draft sequences - for (int s = 0; s < n_seq_dft; ++s) { - if (!drafts[s].active) { - continue; - } - - active_seqs.insert(s); - const auto & tokens = drafts[s].tokens; - - LOG_DBG("draft %d: %s\n", s, string_from(ctx_dft, tokens).c_str()); - } - - int i_dft = 0; - int s_keep = 0; - - llama_token token_id; - std::string token_str; - - // loop until we fail to accept a drafted token or we run out of drafted tokens - while (true) { - - // check if the target token matches any of the drafts - // for stochastic sampling, attempt to match the token with the drafted tokens - { - bool accept = false; - if (params.sampling.temp > 0) { - // stochastic verification - common_sampler_sample(smpl, ctx_tgt, drafts[s_keep].i_batch_tgt[i_dft], true); - - auto & dist_tgt = *common_sampler_get_candidates(smpl, true); - - float p_tgt = 0.0f; - float p_dft = 0.0f; - - while (active_seqs.size() > 0) { - // randomly select a sequence to verify from active sequences - std::uniform_int_distribution u_int_dist(0, active_seqs.size() - 1); - int s = *std::next(active_seqs.begin(), u_int_dist(rng)); - if (i_dft >= (int) drafts[s].tokens.size()) { - drafts[s].active = false; - active_seqs.erase(s); - continue; - } - if (accept) { - // if we already accepted a token, we can skip the rest - if (drafts[s].tokens[i_dft] != drafts[s_keep].tokens[i_dft]) { - drafts[s].active = false; - active_seqs.erase(s); - } - continue; - } - - LOG_DBG("verifying sequence #%d at pos #%d from %d active sequence(s)\n", s, i_dft, (int) active_seqs.size()); - float r = u_dist(rng); - llama_token_data_array dist_dft = { drafts[s].dists[i_dft].data() , drafts[s].dists[i_dft].size(), LLAMA_TOKEN_NULL, true }; - - //GGML_ASSERT(dist_tgt.size <= dist_dft.size); - - // acquire the token probabilities assigned by the draft and target models - for (size_t i = 0; i < dist_tgt.size; i++) { - if (dist_tgt.data[i].id == drafts[s].tokens[i_dft]) { - p_tgt = dist_tgt.data[i].p; - break; - } - } - for (size_t i = 0; i < dist_dft.size; i++) { - if (dist_dft.data[i].id == drafts[s].tokens[i_dft]) { - p_dft = dist_dft.data[i].p; - break; - } - } - LOG_DBG("r = %f, p_dft = %f, p_tgt = %f\n", r, p_dft, p_tgt); - if (r <= p_tgt / p_dft) { - s_keep = s; - accept = true; - token_id = drafts[s].tokens[i_dft]; - token_str = common_token_to_piece(ctx_tgt, token_id); - common_sampler_accept(smpl, token_id, true); - - LOG_DBG("draft token %d of sequence %d (%d, '%s') accepted\n", i_dft, s, token_id, token_str.c_str()); - break; - } else { - LOG_DBG("draft token %d of sequence %d (%d, '%s') rejected\n", i_dft, s, drafts[s].tokens[i_dft], common_token_to_piece(ctx_tgt, drafts[s].tokens[i_dft]).c_str()); - drafts[s].active = false; - - // calculate residual probability - GGML_ASSERT(dist_tgt.sorted); - GGML_ASSERT(dist_dft.sorted); - - // sort dist by id - std::sort(dist_tgt.data, dist_tgt.data + dist_tgt.size, [](const llama_token_data &a, const llama_token_data &b) { - return a.id < b.id; - }); - std::sort(dist_dft.data, dist_dft.data + dist_dft.size, [](const llama_token_data &a, const llama_token_data &b) { - return a.id < b.id; - }); - - float sum_probs = 0.0f; - - for (size_t i = 0; i < dist_tgt.size; i++) { - if (i < dist_dft.size) { - dist_tgt.data[i].p = std::max(0.0f, dist_tgt.data[i].p - dist_dft.data[i].p); - } else { - dist_tgt.data[i].p = std::max(0.0f, dist_tgt.data[i].p); - } - - sum_probs += dist_tgt.data[i].p; - } - - for (size_t i = 0; i < dist_tgt.size; i++) { - dist_tgt.data[i].p /= sum_probs; - } - - // sort dist_tgt by p desc - std::sort(dist_tgt.data, dist_tgt.data + dist_tgt.size, [](const llama_token_data &a, const llama_token_data &b) { - return a.p > b.p; - }); - } - - active_seqs.erase(s); - for (int i = 0; i < n_seq_dft; i++) { - if (i == s) { - continue; - } - if (drafts[i].active && drafts[i].tokens[i_dft] == drafts[s].tokens[i_dft]) { - // synchronize active status for sequences with the same drafted token - drafts[i].active = drafts[i].active && accept; - if (!drafts[i].active) { - active_seqs.erase(s); - } - } - } - } - - if (!accept) { - // all drafted tokens were rejected - // sample from the target model - LOG_DBG("all drafted tokens were rejected, sampling from residual distribution\n"); - std::vector probs(dist_tgt.size); - for (size_t i = 0; i < dist_tgt.size; ++i) { - probs[i] = dist_tgt.data[i].p; - } - - std::discrete_distribution<> dist(probs.begin(), probs.end()); - - const int idx = dist(rng); - - token_id = dist_tgt.data[idx].id; - common_sampler_accept(smpl, token_id, true); - token_str = common_token_to_piece(ctx_tgt, token_id); - } - } else { - // greedy verification - - // sample from the target model - LOG_DBG("sampling target: s_keep = %3d, i_dft = %3d, i_batch_tgt = %3d\n", s_keep, i_dft, drafts[s_keep].i_batch_tgt[i_dft]); - token_id = common_sampler_sample(smpl, ctx_tgt, drafts[s_keep].i_batch_tgt[i_dft]); - - common_sampler_accept(smpl, token_id, true); - - token_str = common_token_to_piece(ctx_tgt, token_id); - - for (int s = 0; s < n_seq_dft; ++s) { - if (!drafts[s].active) { - continue; - } - - if (i_dft < (int) drafts[s].tokens.size() && token_id == drafts[s].tokens[i_dft]) { - LOG_DBG("the sampled target token matches the %dth drafted token of sequence %d (%d, '%s') - accepted\n", i_dft, s, token_id, token_str.c_str()); - - s_keep = s; - accept = true; - } else { - drafts[s].active = false; - } - } - } - - if (llama_vocab_is_eog(vocab_tgt, token_id)) { - has_eos = true; - } - ++n_predict; - - if (accept) { - ++n_accept; - ++n_past_tgt; - ++n_past_dft; - ++i_dft; - if (params.use_color) { - // Color token according to its origin sequence - LOG("\u001b[%dm%s\u001b[37m", (36 - s_keep % 6), token_str.c_str()); - } else { - LOG("%s", token_str.c_str()); - } - continue; - } else { - LOG("%s", token_str.c_str()); - break; - } - } - } - - { - LOG_DBG("the sampled target token (%d, '%s') did not match, or we ran out of drafted tokens\n", token_id, token_str.c_str()); - - // TODO: simplify - { - LOG_DBG("keeping sequence %d, n_past_tgt = %d, n_past_dft = %d\n", s_keep, n_past_tgt, n_past_dft); - - llama_memory_seq_keep(mem_dft, s_keep); - llama_memory_seq_cp (mem_dft, s_keep, 0, -1, -1); - llama_memory_seq_keep(mem_dft, 0); - - llama_memory_seq_rm (mem_tgt, s_keep, n_past_tgt, -1); - llama_memory_seq_keep(mem_tgt, s_keep); - llama_memory_seq_cp (mem_tgt, s_keep, 0, -1, -1); - llama_memory_seq_keep(mem_tgt, 0); - } - - for (int s = 0; s < n_seq_dft; ++s) { - drafts[s].active = false; - drafts[s].tokens.clear(); - drafts[s].i_batch_tgt.clear(); - drafts[s].dists.clear(); - } - // note: will be erased after the speculation phase - drafts[0].tokens.push_back(token_id); - drafts[0].dists.push_back(std::vector()); - drafts[0].i_batch_tgt.push_back(0); - - common_batch_clear(batch_dft); - common_batch_add (batch_dft, token_id, n_past_dft, { 0 }, true); - - llama_memory_seq_rm(mem_dft, 0, n_past_dft, -1); - // LOG_DBG("dft batch: %s\n", LOG_BATCH_TOSTR_PRETTY(ctx_dft, batch_dft).c_str()); - llama_decode(ctx_dft, batch_dft); - - ++n_past_dft; - } - - if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) { - break; - } - - if (drafts[0].smpl) { - common_sampler_free(drafts[0].smpl); - } - drafts[0].smpl = common_sampler_clone(smpl); - - int n_seq_cur = 1; - int n_past_cur = n_past_dft; - - for (int s = 0; s < n_seq_dft; ++s) { - drafts[s].active = false; - drafts[s].drafting = false; - } - drafts[0].active = true; - drafts[0].drafting = true; - drafts[0].i_batch_dft = 0; - - common_batch_clear(batch_tgt); - common_batch_add (batch_tgt, drafts[0].tokens[0], n_past_tgt, { 0 }, true); - - // sample n_draft tokens from the draft model using tree-based sampling - for (int i = 0; i < n_draft; ++i) { - batch_dft.n_tokens = 0; - - for (int s = 0; s < n_seq_dft; ++s) { - drafts[s].skip = false; - } - - for (int s = 0; s < n_seq_dft; ++s) { - if (!drafts[s].drafting || drafts[s].skip) { - continue; - } - - common_sampler_sample(drafts[s].smpl, ctx_dft, drafts[s].i_batch_dft, true); - - const auto * cur_p = common_sampler_get_candidates(drafts[s].smpl, true); - - for (int k = 0; k < std::min(n_seq_dft + 3, (int) cur_p->size); ++k) { - LOG_DBG(" - draft candidate %3d for seq %3d, pos %3d: %6d (%8.3f) '%s'\n", - k, s, i, cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); - } - - std::vector sa(1, s); - - // attempt to split the branch if the probability is high enough - for (int f = 1; f < 8; ++f) { - if (n_seq_cur < n_seq_dft && cur_p->data[f].p > p_draft_split) { - LOG_DBG("splitting seq %3d into %3d\n", s, n_seq_cur); - - llama_memory_seq_rm(mem_dft, n_seq_cur, -1, -1); - llama_memory_seq_cp(mem_dft, s, n_seq_cur, -1, -1); - - // all previous tokens from this branch are now also part of the new branch - for (int t = 0; t < batch_tgt.n_tokens; ++t) { - for (int p = 0; p < batch_tgt.n_seq_id[t]; ++p) { - if (batch_tgt.seq_id[t][p] == s) { - batch_tgt.seq_id[t][batch_tgt.n_seq_id[t]] = n_seq_cur; - batch_tgt.n_seq_id[t]++; - break; - } - } - } - - // copy the draft state - drafts[n_seq_cur].active = true; - drafts[n_seq_cur].drafting = true; - drafts[n_seq_cur].skip = true; - - drafts[n_seq_cur].tokens = drafts[s].tokens; - drafts[n_seq_cur].dists = drafts[s].dists; - drafts[n_seq_cur].i_batch_dft = drafts[s].i_batch_dft; - drafts[n_seq_cur].i_batch_tgt = drafts[s].i_batch_tgt; - - if (drafts[n_seq_cur].smpl) { - common_sampler_free(drafts[n_seq_cur].smpl); - } - drafts[n_seq_cur].smpl = common_sampler_clone(drafts[s].smpl); - - sa.push_back(n_seq_cur); - - n_seq_cur++; - } else { - break; - } - } - - // add drafted token for each sequence - for (int is = 0; is < (int) sa.size(); ++is) { - const llama_token id = cur_p->data[is].id; - - const int s = sa[is]; - - common_sampler_accept(drafts[s].smpl, id, true); - - drafts[s].tokens.push_back(id); - // save cur_p.data into drafts[s].dists - drafts[s].dists.push_back({cur_p->data, cur_p->data + cur_p->size}); - - // add unique drafted tokens to the target batch - drafts[s].i_batch_tgt.push_back(batch_tgt.n_tokens); - - common_batch_add(batch_tgt, id, n_past_tgt + i + 1, { s }, true); - - // add the token to the batch for batched decoding with the draft model - drafts[s].i_batch_dft = batch_dft.n_tokens; - - common_batch_add(batch_dft, id, n_past_cur, { s }, true); - - if (batch_tgt.n_tokens > n_draft) { - drafts[s].drafting = false; - } - } - } - - // no sequence is drafting anymore - if (batch_dft.n_tokens == 0) { - break; - } - - // evaluate the drafted tokens on the draft model - llama_decode(ctx_dft, batch_dft); - ++n_past_cur; - ++n_drafted; - - if (batch_tgt.n_tokens > n_draft) { - break; - } - } - - // evaluate the target model on the drafted tokens - { - llama_memory_seq_keep(mem_tgt, 0); - for (int s = 1; s < n_seq_dft; ++s) { - llama_memory_seq_cp(mem_tgt, 0, s, -1, -1); - } - - // LOG_DBG("target batch: %s\n", LOG_BATCH_TOSTR_PRETTY(ctx_tgt, batch_tgt).c_str()); - llama_decode(ctx_tgt, batch_tgt); - ++n_past_tgt; - } - - // the first token is always proposed by the target model before the speculation loop so we erase it here - for (int s = 0; s < n_seq_dft; ++s) { - if (!drafts[s].active) { - continue; - } - - drafts[s].tokens.erase(drafts[s].tokens.begin()); - drafts[s].dists.erase(drafts[s].dists.begin()); - } - } - - auto t_dec_end = ggml_time_us(); - - LOG("\n\n"); - - LOG_INF("encoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_input, (t_enc_end - t_enc_start) / 1e6f, inp.size() / ((t_enc_end - t_enc_start) / 1e6f)); - LOG_INF("decoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_predict, (t_dec_end - t_dec_start) / 1e6f, n_predict / ((t_dec_end - t_dec_start) / 1e6f)); - - LOG_INF("\n"); - LOG_INF("n_draft = %d\n", n_draft); - LOG_INF("n_predict = %d\n", n_predict); - LOG_INF("n_drafted = %d\n", n_drafted); - LOG_INF("n_accept = %d\n", n_accept); - LOG_INF("accept = %.3f%%\n", 100.0f * n_accept / n_drafted); - - LOG_INF("\n"); - LOG_INF("draft:\n\n"); - // TODO: print sampling/grammar timings for all drafts - llama_perf_context_print(ctx_dft); - - LOG_INF("\n"); - LOG_INF("target:\n\n"); - common_perf_print(ctx_tgt, smpl); - - common_sampler_free(smpl); - for (int s = 0; s < n_seq_dft; ++s) { - common_sampler_free(drafts[s].smpl); - } - - llama_batch_free(batch_dft); - - llama_backend_free(); - - LOG("\n\n"); - - return 0; -} diff --git a/examples/sycl/CMakeLists.txt b/examples/sycl/CMakeLists.txt index 40e44eefc8a8..db0f5e67227c 100644 --- a/examples/sycl/CMakeLists.txt +++ b/examples/sycl/CMakeLists.txt @@ -3,7 +3,25 @@ # SPDX-License-Identifier: MIT set(TARGET llama-ls-sycl-device) -add_executable(${TARGET} ls-sycl-device.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(SYCL_LS_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/ls-sycl-device-rust/Cargo.toml) +set(SYCL_LS_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/ls-sycl-device) +set(SYCL_LS_RUST_BIN ${SYCL_LS_RUST_TARGET_DIR}/debug/llama-ls-sycl-device${CMAKE_EXECUTABLE_SUFFIX}) +set(SYCL_LS_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${SYCL_LS_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=${CMAKE_RUNTIME_OUTPUT_DIRECTORY} + ${CARGO_EXECUTABLE} build --manifest-path ${SYCL_LS_RUST_MANIFEST} --target-dir ${SYCL_LS_RUST_TARGET_DIR} --bin llama-ls-sycl-device + COMMAND ${CMAKE_COMMAND} -E copy ${SYCL_LS_RUST_BIN} ${SYCL_LS_RUST_OUTPUT} + DEPENDS + ggml-sycl + ${SYCL_LS_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/ls-sycl-device-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/ls-sycl-device-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/ls-sycl-device-rust/src/main.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${SYCL_LS_RUST_OUTPUT}) +install(PROGRAMS ${SYCL_LS_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/sycl/ls-sycl-device-rust/Cargo.toml b/examples/sycl/ls-sycl-device-rust/Cargo.toml new file mode 100644 index 000000000000..66fb18c2117d --- /dev/null +++ b/examples/sycl/ls-sycl-device-rust/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "llama-ls-sycl-device-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] + +[[bin]] +name = "llama-ls-sycl-device" +path = "src/main.rs" diff --git a/examples/sycl/ls-sycl-device-rust/build.rs b/examples/sycl/ls-sycl-device-rust/build.rs new file mode 100644 index 000000000000..14291a86a718 --- /dev/null +++ b/examples/sycl/ls-sycl-device-rust/build.rs @@ -0,0 +1,42 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=ggml-sycl"); + emit_rpath(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/ggml/src/ggml-sycl"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libggml-sycl.so").exists() + || candidate.join("libggml-sycl.dylib").exists() + || candidate.join("ggml-sycl.dll").exists() + || candidate.join("ggml-sycl.lib").exists() + || candidate.join("libggml-sycl.a").exists() + { + println!("cargo:rustc-link-search=native={}", candidate.display()); + println!("cargo:rustc-link-lib=ggml-sycl"); + emit_rpath(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libggml-sycl was not found; this binary requires a GGML_SYCL build"); +} + +fn emit_rpath(lib_dir: &str) { + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/sycl/ls-sycl-device-rust/src/lib.rs b/examples/sycl/ls-sycl-device-rust/src/lib.rs new file mode 100644 index 000000000000..134d860320a1 --- /dev/null +++ b/examples/sycl/ls-sycl-device-rust/src/lib.rs @@ -0,0 +1,13 @@ +pub fn binary_name() -> &'static str { + "llama-ls-sycl-device" +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn exposes_binary_name() { + assert_eq!(binary_name(), "llama-ls-sycl-device"); + } +} diff --git a/examples/sycl/ls-sycl-device-rust/src/main.rs b/examples/sycl/ls-sycl-device-rust/src/main.rs new file mode 100644 index 000000000000..6abf36b3f263 --- /dev/null +++ b/examples/sycl/ls-sycl-device-rust/src/main.rs @@ -0,0 +1,11 @@ +extern "C" { + fn ggml_backend_sycl_print_sycl_devices(); +} + +fn main() { + let _ = llama_ls_sycl_device_rust::binary_name(); + + unsafe { + ggml_backend_sycl_print_sycl_devices(); + } +} diff --git a/examples/sycl/ls-sycl-device.cpp b/examples/sycl/ls-sycl-device.cpp deleted file mode 100644 index 3bdc40598252..000000000000 --- a/examples/sycl/ls-sycl-device.cpp +++ /dev/null @@ -1,15 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - - -#include "ggml-sycl.h" -#include - -int main() { - std::setlocale(LC_NUMERIC, "C"); - ggml_backend_sycl_print_sycl_devices(); - return 0; -} diff --git a/examples/training/CMakeLists.txt b/examples/training/CMakeLists.txt index 8bb20d0f2138..b376ea6a564c 100644 --- a/examples/training/CMakeLists.txt +++ b/examples/training/CMakeLists.txt @@ -1,5 +1,25 @@ set(TARGET llama-finetune) -add_executable(${TARGET} finetune.cpp) -install(TARGETS ${TARGET} RUNTIME) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_11) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(FINETUNE_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/finetune-rust/Cargo.toml) +set(FINETUNE_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/finetune) +set(FINETUNE_RUST_BIN ${FINETUNE_RUST_TARGET_DIR}/debug/llama-finetune${CMAKE_EXECUTABLE_SUFFIX}) +set(FINETUNE_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${FINETUNE_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=$ + ${CARGO_EXECUTABLE} build --manifest-path ${FINETUNE_RUST_MANIFEST} --target-dir ${FINETUNE_RUST_TARGET_DIR} --bin llama-finetune + COMMAND ${CMAKE_COMMAND} -E copy ${FINETUNE_RUST_BIN} ${FINETUNE_RUST_OUTPUT} + DEPENDS + llama + ${FINETUNE_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/finetune-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/finetune-rust/src/main.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/ffi.rs + ${CMAKE_CURRENT_SOURCE_DIR}/../simple-rust/src/lib.rs +) + +add_custom_target(${TARGET} ALL DEPENDS ${FINETUNE_RUST_OUTPUT}) +install(PROGRAMS ${FINETUNE_RUST_OUTPUT} DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/examples/training/finetune-rust/Cargo.toml b/examples/training/finetune-rust/Cargo.toml new file mode 100644 index 000000000000..b3fb6dbfb764 --- /dev/null +++ b/examples/training/finetune-rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "llama-finetune-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +llama-simple-rust = { path = "../../simple-rust" } + +[[bin]] +name = "llama-finetune" +path = "src/main.rs" diff --git a/examples/training/finetune-rust/build.rs b/examples/training/finetune-rust/build.rs new file mode 100644 index 000000000000..f0813c609aed --- /dev/null +++ b/examples/training/finetune-rust/build.rs @@ -0,0 +1,45 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-env-changed=LLAMA_CPP_LIB_DIR"); + + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + emit_link(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libllama.so").exists() + || candidate.join("libllama.dylib").exists() + || candidate.join("llama.dll").exists() + || candidate.join("llama.lib").exists() + || candidate.join("libllama.a").exists() + { + emit_link(&candidate.display().to_string()); + return; + } + } + + println!("cargo:warning=libllama was not found; set LLAMA_CPP_LIB_DIR or build llama.cpp before building llama-finetune"); +} + +fn emit_link(lib_dir: &str) { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=llama"); + println!("cargo:rustc-link-lib=ggml"); + println!("cargo:rustc-link-lib=ggml-base"); + + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/examples/training/finetune-rust/src/main.rs b/examples/training/finetune-rust/src/main.rs new file mode 100644 index 000000000000..5533200f1d24 --- /dev/null +++ b/examples/training/finetune-rust/src/main.rs @@ -0,0 +1,642 @@ +use llama_simple_rust::ffi; +use std::env; +use std::ffi::{c_void, CStr, CString}; +use std::ptr; + +const GGML_TYPE_F32: ffi::ggml_type = 0; +const GGML_TYPE_I32: ffi::ggml_type = 26; +const OPTIMIZER_ADAMW: ffi::ggml_opt_optimizer_type = 0; +const OPTIMIZER_SGD: ffi::ggml_opt_optimizer_type = 1; + +#[derive(Debug, Clone, PartialEq)] +pub struct LrOpt { + pub lr0: f32, + pub lr_min: f32, + pub decay_epochs: f32, + pub scale_epoch: f32, + pub wd: f32, + pub epochs: u32, + pub epoch: u32, +} + +impl Default for LrOpt { + fn default() -> Self { + Self { + lr0: 1e-5, + lr_min: -1.0, + decay_epochs: -1.0, + scale_epoch: 0.0, + wd: 0.0, + epochs: 2, + epoch: 0, + } + } +} + +impl LrOpt { + pub fn init(&mut self) { + if self.lr_min > 0.0 && self.lr_min < self.lr0 { + let nhalf = (self.lr0 / self.lr_min).ln() / 2.0_f32.ln(); + let mut epochs = self.epochs as f32; + if self.decay_epochs > 0.0 && self.decay_epochs < epochs { + epochs = self.decay_epochs; + } else { + self.decay_epochs = epochs; + } + self.scale_epoch = nhalf / epochs; + } + } + + pub fn get_lr(&self, epoch: f32) -> f32 { + if self.lr_min <= 0.0 { + self.lr0 + } else if epoch >= self.decay_epochs { + self.lr_min + } else { + self.lr0 * 0.5_f32.powf(epoch * self.scale_epoch) + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct Args { + pub model_path: String, + pub prompt: String, + pub out_file: String, + pub n_ctx: u32, + pub n_batch: u32, + pub n_ubatch: u32, + pub n_gpu_layers: i32, + pub n_threads: i32, + pub n_threads_batch: i32, + pub lr: LrOpt, + pub optimizer: ffi::ggml_opt_optimizer_type, + pub val_split: f32, +} + +impl Default for Args { + fn default() -> Self { + Self { + model_path: String::new(), + prompt: String::new(), + out_file: String::new(), + n_ctx: 512, + n_batch: 512, + n_ubatch: 512, + n_gpu_layers: 99, + n_threads: 0, + n_threads_batch: 0, + lr: LrOpt::default(), + optimizer: OPTIMIZER_ADAMW, + val_split: 0.05, + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub enum ParseError { + MissingModel, + MissingOutput, + MissingPrompt, + MissingValue(String), + InvalidInteger(String, String), + InvalidFloat(String, String), + InvalidValue(String), +} + +impl std::fmt::Display for ParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ParseError::MissingModel => write!(f, "missing required -m/--model model.gguf"), + ParseError::MissingOutput => write!(f, "missing required -o/--output output.gguf"), + ParseError::MissingPrompt => write!(f, "missing required -p/--prompt text"), + ParseError::MissingValue(flag) => write!(f, "missing value for {flag}"), + ParseError::InvalidInteger(flag, value) => { + write!(f, "invalid integer for {flag}: {value}") + } + ParseError::InvalidFloat(flag, value) => write!(f, "invalid float for {flag}: {value}"), + ParseError::InvalidValue(message) => write!(f, "{message}"), + } + } +} + +impl std::error::Error for ParseError {} + +pub fn parse_args(args: I) -> Result +where + I: IntoIterator, + S: Into, +{ + let mut parsed = Args::default(); + let mut iter = args.into_iter().map(Into::into); + + while let Some(arg) = iter.next() { + match arg.as_str() { + "-m" | "--model" => parsed.model_path = take(&mut iter, &arg)?, + "-p" | "--prompt" => parsed.prompt = take(&mut iter, &arg)?, + "-o" | "--output" => parsed.out_file = take(&mut iter, &arg)?, + "-c" | "--ctx-size" => parsed.n_ctx = parse_u32(&mut iter, &arg)?, + "-b" | "--batch-size" => parsed.n_batch = parse_u32(&mut iter, &arg)?, + "-ub" | "--ubatch-size" => parsed.n_ubatch = parse_u32(&mut iter, &arg)?, + "-ngl" | "--gpu-layers" => parsed.n_gpu_layers = parse_i32(&mut iter, &arg)?, + "-t" | "--threads" => parsed.n_threads = parse_i32(&mut iter, &arg)?, + "-tb" | "--threads-batch" => parsed.n_threads_batch = parse_i32(&mut iter, &arg)?, + "-lr" | "--learning-rate" => parsed.lr.lr0 = parse_f32(&mut iter, &arg)?, + "-lr-min" | "--learning-rate-min" => parsed.lr.lr_min = parse_f32(&mut iter, &arg)?, + "-lr-decay" | "--learning-rate-decay-epochs" => { + parsed.lr.decay_epochs = parse_f32(&mut iter, &arg)? + } + "-wd" | "--weight-decay" => parsed.lr.wd = parse_f32(&mut iter, &arg)?, + "--val-split" => parsed.val_split = parse_f32(&mut iter, &arg)?, + "--epochs" => parsed.lr.epochs = parse_u32(&mut iter, &arg)?, + "-opt" | "--optimizer" => { + parsed.optimizer = parse_optimizer(&take(&mut iter, &arg)?)?; + } + "-h" | "--help" => return Err(ParseError::MissingModel), + _ => { + if !arg.starts_with('-') { + if !parsed.prompt.is_empty() { + parsed.prompt.push(' '); + } + parsed.prompt.push_str(&arg); + } + } + } + } + + if parsed.model_path.is_empty() { + return Err(ParseError::MissingModel); + } + if parsed.out_file.is_empty() { + return Err(ParseError::MissingOutput); + } + if parsed.prompt.is_empty() { + return Err(ParseError::MissingPrompt); + } + if parsed.n_ctx == 0 || parsed.n_batch == 0 || parsed.n_ubatch == 0 { + return Err(ParseError::InvalidValue( + "context, batch, and ubatch sizes must be positive".to_string(), + )); + } + if parsed.lr.epochs == 0 { + return Err(ParseError::InvalidValue( + "--epochs must be positive".to_string(), + )); + } + if !(0.0..1.0).contains(&parsed.val_split) { + return Err(ParseError::InvalidValue( + "--val-split must be in [0.0, 1.0)".to_string(), + )); + } + parsed.lr.init(); + Ok(parsed) +} + +fn take(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + iter.next() + .ok_or_else(|| ParseError::MissingValue(flag.to_string())) +} + +fn parse_i32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_u32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidInteger(flag.to_string(), value)) +} + +fn parse_f32(iter: &mut I, flag: &str) -> Result +where + I: Iterator, +{ + let value = take(iter, flag)?; + value + .parse() + .map_err(|_| ParseError::InvalidFloat(flag.to_string(), value)) +} + +fn parse_optimizer(name: &str) -> Result { + match name.to_ascii_lowercase().as_str() { + "adamw" => Ok(OPTIMIZER_ADAMW), + "sgd" => Ok(OPTIMIZER_SGD), + _ => Err(ParseError::InvalidValue( + "invalid --optimizer, valid options: adamw, sgd".to_string(), + )), + } +} + +struct Model(*mut ffi::llama_model); + +impl Drop for Model { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_model_free(self.0); + } + } + } +} + +struct Context(*mut ffi::llama_context); + +impl Drop for Context { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::llama_free(self.0); + } + } + } +} + +struct Dataset(ffi::ggml_opt_dataset_t); + +impl Drop for Dataset { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::ggml_opt_dataset_free(self.0); + } + } + } +} + +struct OptResult(ffi::ggml_opt_result_t); + +impl Drop for OptResult { + fn drop(&mut self) { + unsafe { + if !self.0.is_null() { + ffi::ggml_opt_result_free(self.0); + } + } + } +} + +struct Backend; + +impl Backend { + fn init() -> Self { + unsafe { + ffi::ggml_backend_load_all(); + ffi::llama_backend_init(); + ffi::llama_numa_init(0); + } + Self + } +} + +impl Drop for Backend { + fn drop(&mut self) { + unsafe { + ffi::llama_backend_free(); + } + } +} + +fn tokenize(vocab: *const ffi::llama_vocab, text: &str) -> Result, String> { + let text_c = + CString::new(text).map_err(|_| "prompt contains an interior NUL byte".to_string())?; + let n_tokens = unsafe { + -ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + ptr::null_mut(), + 0, + true, + false, + ) + }; + if n_tokens <= 0 { + return Err("failed to size tokenization".to_string()); + } + let mut tokens = vec![0_i32; n_tokens as usize]; + let n = unsafe { + ffi::llama_tokenize( + vocab, + text_c.as_ptr(), + text.len() as i32, + tokens.as_mut_ptr(), + tokens.len() as i32, + true, + false, + ) + }; + if n < 0 { + return Err("failed to tokenize prompt".to_string()); + } + tokens.truncate(n as usize); + Ok(tokens) +} + +fn create_dataset( + ctx: *mut ffi::llama_context, + tokens: &[ffi::llama_token], +) -> Result { + let ne_datapoint = unsafe { ffi::llama_n_ctx(ctx) as i64 }; + let stride = ne_datapoint / 2; + if stride <= 0 { + return Err("context size is too small for training dataset".to_string()); + } + if tokens.len() as i64 <= ne_datapoint + 1 { + return Err(format!( + "prompt is too short for finetuning dataset: {} tokens, need more than {}", + tokens.len(), + ne_datapoint + 1 + )); + } + let ndata = (tokens.len() as i64 - ne_datapoint - 1) / stride; + if ndata <= 0 { + return Err("prompt did not produce any finetuning datapoints".to_string()); + } + let dataset = Dataset(unsafe { + ffi::ggml_opt_dataset_init( + GGML_TYPE_I32, + GGML_TYPE_I32, + ne_datapoint, + ne_datapoint, + ndata, + 1, + ) + }); + if dataset.0.is_null() { + return Err("failed to create optimizer dataset".to_string()); + } + + unsafe { + let data_ptr = ffi::ggml_get_data(ffi::ggml_opt_dataset_data(dataset.0)) as *mut i32; + let labels_ptr = ffi::ggml_get_data(ffi::ggml_opt_dataset_labels(dataset.0)) as *mut i32; + if data_ptr.is_null() || labels_ptr.is_null() { + return Err("failed to access optimizer dataset tensors".to_string()); + } + for idata in 0..ndata as usize { + let src = tokens.as_ptr().add(idata * stride as usize); + ptr::copy_nonoverlapping( + src, + data_ptr.add(idata * ne_datapoint as usize), + ne_datapoint as usize, + ); + ptr::copy_nonoverlapping( + src.add(1), + labels_ptr.add(idata * ne_datapoint as usize), + ne_datapoint as usize, + ); + } + } + Ok(dataset) +} + +unsafe extern "C" fn lr_params_callback(userdata: *mut c_void) -> ffi::ggml_opt_optimizer_params { + let mut result = ffi::ggml_opt_get_default_optimizer_params(ptr::null_mut()); + let lr = &*(userdata as *const LrOpt); + let value = lr.get_lr(lr.epoch as f32); + eprintln!("epoch {:.2} lr={:.2e}", lr.epoch, value); + result.adamw.alpha = value; + result.sgd.alpha = value; + result.adamw.wd = lr.wd; + result.sgd.wd = lr.wd; + result +} + +fn optimizer_name(optimizer: ffi::ggml_opt_optimizer_type) -> String { + let ptr = unsafe { ffi::ggml_opt_optimizer_name(optimizer) }; + if ptr.is_null() { + optimizer.to_string() + } else { + unsafe { CStr::from_ptr(ptr).to_string_lossy().into_owned() } + } +} + +fn run(mut args: Args) -> Result<(), String> { + let _backend = Backend::init(); + let model_path = CString::new(args.model_path.as_str()) + .map_err(|_| "model path contains an interior NUL byte".to_string())?; + let out_file = CString::new(args.out_file.as_str()) + .map_err(|_| "output path contains an interior NUL byte".to_string())?; + + let mut model_params = unsafe { ffi::llama_model_default_params() }; + model_params.n_gpu_layers = args.n_gpu_layers; + model_params.use_mmap = false; + + let model = + Model(unsafe { ffi::llama_model_load_from_file(model_path.as_ptr(), model_params) }); + if model.0.is_null() { + return Err("unable to load model".to_string()); + } + + let mut ctx_params = unsafe { ffi::llama_context_default_params() }; + ctx_params.n_ctx = args.n_ctx; + ctx_params.n_batch = args.n_batch; + ctx_params.n_ubatch = args.n_ubatch; + ctx_params.n_threads = args.n_threads; + ctx_params.n_threads_batch = args.n_threads_batch; + ctx_params.type_k = GGML_TYPE_F32; + ctx_params.type_v = GGML_TYPE_F32; + + let ctx = Context(unsafe { ffi::llama_init_from_model(model.0, ctx_params) }); + if ctx.0.is_null() { + return Err("failed to create llama_context".to_string()); + } + + let vocab = unsafe { ffi::llama_model_get_vocab(model.0) }; + if vocab.is_null() { + return Err("failed to get model vocabulary".to_string()); + } + let tokens = tokenize(vocab, &args.prompt)?; + let dataset = create_dataset(ctx.0, &tokens)?; + + eprintln!( + "-optimizer {} -lr0 {:.2e} -wd {:.2e} -lr-min {:.2e} -min-epochs {:.2e} -epochs {} -period {:.2e} -val {:.2e}", + optimizer_name(args.optimizer), + args.lr.lr0, + args.lr.wd, + args.lr.lr_min, + args.lr.decay_epochs, + args.lr.epochs, + args.n_batch as f32 / args.n_ubatch as f32, + args.val_split + ); + + let lopt_params = ffi::llama_opt_params { + n_ctx_train: 0, + param_filter: Some(ffi::llama_opt_param_filter_all), + param_filter_ud: ptr::null_mut(), + get_opt_pars: Some(lr_params_callback), + get_opt_pars_ud: (&mut args.lr as *mut LrOpt).cast::(), + optimizer_type: args.optimizer, + }; + unsafe { + ffi::llama_opt_init(ctx.0, model.0, lopt_params); + } + + let ndata = unsafe { ffi::ggml_opt_dataset_ndata(dataset.0) }; + let idata_split = (ndata as f32 * (1.0 - args.val_split)) as i64; + let result_train = OptResult(unsafe { ffi::ggml_opt_result_init() }); + let result_eval = OptResult(unsafe { ffi::ggml_opt_result_init() }); + if result_train.0.is_null() || result_eval.0.is_null() { + return Err("failed to create optimizer results".to_string()); + } + + for epoch in 0..args.lr.epochs { + args.lr.epoch = epoch; + unsafe { + ffi::llama_opt_epoch( + ctx.0, + dataset.0, + result_train.0, + result_eval.0, + idata_split, + Some(ffi::ggml_opt_epoch_callback_progress_bar), + Some(ffi::ggml_opt_epoch_callback_progress_bar), + ); + eprintln!(); + ffi::ggml_opt_result_reset(result_train.0); + ffi::ggml_opt_result_reset(result_eval.0); + } + } + + unsafe { + ffi::llama_model_save_to_file(model.0, out_file.as_ptr()); + } + Ok(()) +} + +fn print_usage(program: &str) { + eprintln!(); + eprintln!("example usage:"); + eprintln!(); + eprintln!(" {program} -m model.gguf -p \"training text\" -o finetuned.gguf"); + eprintln!(); +} + +fn main() { + let mut argv = env::args(); + let program = argv.next().unwrap_or_else(|| "llama-finetune".to_string()); + let args = match parse_args(argv) { + Ok(args) => args, + Err(err) => { + eprintln!("{program}: error: {err}"); + print_usage(&program); + std::process::exit(1); + } + }; + if let Err(err) = run(args) { + eprintln!("{program}: error: {err}"); + std::process::exit(1); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_finetune_options() { + let args = parse_args([ + "-m", + "model.gguf", + "-p", + "training data", + "-o", + "out.gguf", + "-c", + "256", + "-b", + "128", + "-ub", + "64", + "-ngl", + "0", + "-t", + "2", + "-tb", + "4", + "-lr", + "0.01", + "-lr-min", + "0.001", + "-lr-decay", + "3", + "-wd", + "0.0001", + "--val-split", + "0.2", + "--epochs", + "4", + "-opt", + "sgd", + ]) + .unwrap(); + assert_eq!(args.model_path, "model.gguf"); + assert_eq!(args.prompt, "training data"); + assert_eq!(args.out_file, "out.gguf"); + assert_eq!(args.n_ctx, 256); + assert_eq!(args.n_batch, 128); + assert_eq!(args.n_ubatch, 64); + assert_eq!(args.n_gpu_layers, 0); + assert_eq!(args.n_threads, 2); + assert_eq!(args.n_threads_batch, 4); + assert_eq!(args.lr.lr0, 0.01); + assert_eq!(args.lr.lr_min, 0.001); + assert_eq!(args.lr.decay_epochs, 3.0); + assert_eq!(args.lr.wd, 0.0001); + assert_eq!(args.val_split, 0.2); + assert_eq!(args.lr.epochs, 4); + assert_eq!(args.optimizer, OPTIMIZER_SGD); + } + + #[test] + fn rejects_missing_output() { + assert_eq!( + parse_args(["-m", "model.gguf", "-p", "text"]).unwrap_err(), + ParseError::MissingOutput + ); + } + + #[test] + fn rejects_unknown_optimizer() { + assert_eq!( + parse_args([ + "-m", + "model.gguf", + "-p", + "text", + "-o", + "out.gguf", + "-opt", + "rmsprop" + ]) + .unwrap_err(), + ParseError::InvalidValue("invalid --optimizer, valid options: adamw, sgd".to_string()) + ); + } + + #[test] + fn learning_rate_decays_like_common_helper() { + let mut lr = LrOpt { + lr0: 1e-2, + lr_min: 1e-3, + epochs: 4, + ..LrOpt::default() + }; + lr.init(); + assert!((lr.get_lr(0.0) - 1e-2).abs() < 1e-8); + assert!((lr.get_lr(4.0) - 1e-3).abs() < 1e-8); + } +} diff --git a/examples/training/finetune.cpp b/examples/training/finetune.cpp deleted file mode 100644 index 0a75ac110ca4..000000000000 --- a/examples/training/finetune.cpp +++ /dev/null @@ -1,101 +0,0 @@ -#include "arg.h" -#include "common.h" -#include "log.h" -#include "llama.h" - -#include -#include -#include -#include -#include -#include - -#if defined(_MSC_VER) -#pragma warning(disable: 4244 4267) // possible loss of data -#endif - -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - params.escape = false; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_FINETUNE)) { - return 1; - } - - if (params.use_mmap) { - LOG_INF("%s: force disabling memory mapping because it would result in-read-only pointers to the weights\n", - __func__); - params.use_mmap = false; - } - if (params.cache_type_k != GGML_TYPE_F32) { - LOG_INF("%s: force changing k cache type to f32 due to a lack of f16 support for OUT_PROD\n", __func__); - params.cache_type_k = GGML_TYPE_F32; - } - if (params.cache_type_v != GGML_TYPE_F32) { - LOG_INF("%s: force changing v cache type to f32 due to a lack of f16 support for OUT_PROD\n", __func__); - params.cache_type_v = GGML_TYPE_F32; - } - - llama_backend_init(); - llama_numa_init(params.numa); - // load the model and apply lora adapter, if any - auto llama_init = common_init_from_params(params); - - auto * model = llama_init->model(); - auto * ctx = llama_init->context(); - - if (model == NULL) { - LOG_ERR("%s: unable to load model\n", __func__); - return 1; - } - - // print system information - { - LOG_INF("\n"); - LOG_INF("%s\n", common_params_get_system_info(params).c_str()); - } - - std::vector tokens = common_tokenize(ctx, params.prompt, true); - ggml_opt_dataset_t dataset = common_opt_dataset_init(ctx, tokens, llama_n_ctx(ctx) / 2); - - struct lr_opt & lr = params.lr; - LOG_INF("-optimizer %s -lr0 %.2g -wd %.2g -lr-min %.2g -min-epochs %.2g -epochs %d -period %.2g -val %.2g\n", - ggml_opt_optimizer_name(params.optimizer), (double) lr.lr0, (double) lr.wd, (double) lr.lr_min, (double) lr.decay_epochs, - (unsigned) lr.epochs, (double) params.n_batch / params.n_ubatch, (double) params.val_split); - - struct llama_opt_params lopt_params{ - /*n_ctx_train =*/0, - /*param_filter =*/llama_opt_param_filter_all, - /*param_filter_ud =*/nullptr, - /*get_opt_pars =*/common_opt_lr_pars, - /*get_opt_pars_ud =*/¶ms.lr, - /*optimizer_type =*/params.optimizer, - }; - llama_opt_init(ctx, model, lopt_params); - - const int64_t idata_split = ggml_opt_dataset_ndata(dataset) * (1.0f - params.val_split); - - ggml_opt_result_t result_train = ggml_opt_result_init(); - ggml_opt_result_t result_eval = ggml_opt_result_init(); - - for (lr.epoch = 0; lr.epoch < lr.epochs; ++lr.epoch) { - llama_opt_epoch(ctx, dataset, result_train, result_eval, idata_split, - ggml_opt_epoch_callback_progress_bar, ggml_opt_epoch_callback_progress_bar); - fprintf(stderr, "\n"); - - ggml_opt_result_reset(result_train); - ggml_opt_result_reset(result_eval); - } - ggml_opt_result_free(result_train); - ggml_opt_result_free(result_eval); - - llama_model_save_to_file(model, params.out_file.c_str()); - - llama_backend_free(); - - return 0; -} diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index c97f681988b6..4731bbbad96d 100644 --- a/ggml/CMakeLists.txt +++ b/ggml/CMakeLists.txt @@ -320,24 +320,23 @@ include(CMakePackageConfigHelpers) # all public headers set(GGML_PUBLIC_HEADERS - include/ggml.h - include/ggml-cpu.h - include/ggml-alloc.h - include/ggml-backend.h - include/ggml-blas.h - include/ggml-cann.h - include/ggml-cpp.h - include/ggml-cuda.h - include/ggml-opt.h - include/ggml-metal.h - include/ggml-rpc.h - include/ggml-virtgpu.h - include/ggml-sycl.h - include/ggml-vulkan.h - include/ggml-webgpu.h - include/ggml-zendnn.h - include/ggml-openvino.h - include/gguf.h) + include/ggml.h.inc + include/ggml-cpu.h.inc + include/ggml-alloc.h.inc + include/ggml-backend.h.inc + include/ggml-blas.h.inc + include/ggml-cann.h.inc + include/ggml-cuda.h.inc + include/ggml-opt.h.inc + include/ggml-metal.h.inc + include/ggml-rpc.h.inc + include/ggml-virtgpu.h.inc + include/ggml-sycl.h.inc + include/ggml-vulkan.h.inc + include/ggml-webgpu.h.inc + include/ggml-zendnn.h.inc + include/ggml-openvino.h.inc + include/gguf.h.inc) set_target_properties(ggml PROPERTIES PUBLIC_HEADER "${GGML_PUBLIC_HEADERS}") #if (GGML_METAL) diff --git a/ggml/include/ggml-alloc.h b/ggml/include/ggml-alloc.h.inc similarity index 99% rename from ggml/include/ggml-alloc.h rename to ggml/include/ggml-alloc.h.inc index 78aa059dde38..7b60c2a9073a 100644 --- a/ggml/include/ggml-alloc.h +++ b/ggml/include/ggml-alloc.h.inc @@ -1,6 +1,6 @@ #pragma once -#include "ggml.h" +#include "ggml.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h.inc similarity index 99% rename from ggml/include/ggml-backend.h rename to ggml/include/ggml-backend.h.inc index d0c7e5a1be08..9171ceb3ae27 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h.inc @@ -1,7 +1,7 @@ #pragma once -#include "ggml.h" -#include "ggml-alloc.h" +#include "ggml.h.inc" +#include "ggml-alloc.h.inc" #ifdef GGML_BACKEND_SHARED # if defined(_WIN32) && !defined(__MINGW32__) diff --git a/ggml/include/ggml-blas.h b/ggml/include/ggml-blas.h.inc similarity index 91% rename from ggml/include/ggml-blas.h rename to ggml/include/ggml-blas.h.inc index 87a81b36348b..d2060bc18f38 100644 --- a/ggml/include/ggml-blas.h +++ b/ggml/include/ggml-blas.h.inc @@ -1,7 +1,7 @@ #pragma once -#include "ggml.h" -#include "ggml-backend.h" +#include "ggml.h.inc" +#include "ggml-backend.h.inc" #ifdef __cplusplus diff --git a/ggml/include/ggml-cann.h b/ggml/include/ggml-cann.h.inc similarity index 98% rename from ggml/include/ggml-cann.h rename to ggml/include/ggml-cann.h.inc index 74af465337a5..6cbb6efd9029 100644 --- a/ggml/include/ggml-cann.h +++ b/ggml/include/ggml-cann.h.inc @@ -22,8 +22,8 @@ #pragma once -#include "ggml-backend.h" -#include "ggml.h" +#include "ggml-backend.h.inc" +#include "ggml.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/include/ggml-cpp.h b/ggml/include/ggml-cpp.h deleted file mode 100644 index 48aa79682b65..000000000000 --- a/ggml/include/ggml-cpp.h +++ /dev/null @@ -1,39 +0,0 @@ -#pragma once - -#ifndef __cplusplus -#error "This header is for C++ only" -#endif - -#include "ggml.h" -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "gguf.h" -#include - -// Smart pointers for ggml types - -// ggml - -struct ggml_context_deleter { void operator()(ggml_context * ctx) { ggml_free(ctx); } }; -struct gguf_context_deleter { void operator()(gguf_context * ctx) { gguf_free(ctx); } }; - -typedef std::unique_ptr ggml_context_ptr; -typedef std::unique_ptr gguf_context_ptr; - -// ggml-alloc - -struct ggml_gallocr_deleter { void operator()(ggml_gallocr_t galloc) { ggml_gallocr_free(galloc); } }; - -typedef std::unique_ptr ggml_gallocr_ptr; - -// ggml-backend - -struct ggml_backend_deleter { void operator()(ggml_backend_t backend) { ggml_backend_free(backend); } }; -struct ggml_backend_buffer_deleter { void operator()(ggml_backend_buffer_t buffer) { ggml_backend_buffer_free(buffer); } }; -struct ggml_backend_event_deleter { void operator()(ggml_backend_event_t event) { ggml_backend_event_free(event); } }; -struct ggml_backend_sched_deleter { void operator()(ggml_backend_sched_t sched) { ggml_backend_sched_free(sched); } }; - -typedef std::unique_ptr ggml_backend_ptr; -typedef std::unique_ptr ggml_backend_buffer_ptr; -typedef std::unique_ptr ggml_backend_event_ptr; -typedef std::unique_ptr ggml_backend_sched_ptr; diff --git a/ggml/include/ggml-cpu.h b/ggml/include/ggml-cpu.h.inc similarity index 99% rename from ggml/include/ggml-cpu.h rename to ggml/include/ggml-cpu.h.inc index e3e067c916f1..74cd73c8fb3b 100644 --- a/ggml/include/ggml-cpu.h +++ b/ggml/include/ggml-cpu.h.inc @@ -1,7 +1,7 @@ #pragma once -#include "ggml.h" -#include "ggml-backend.h" +#include "ggml.h.inc" +#include "ggml-backend.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/include/ggml-cuda.h b/ggml/include/ggml-cuda.h.inc similarity index 97% rename from ggml/include/ggml-cuda.h rename to ggml/include/ggml-cuda.h.inc index 5436c7ef579c..5d23f3e421a2 100644 --- a/ggml/include/ggml-cuda.h +++ b/ggml/include/ggml-cuda.h.inc @@ -1,7 +1,7 @@ #pragma once -#include "ggml.h" -#include "ggml-backend.h" +#include "ggml.h.inc" +#include "ggml-backend.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/include/ggml-hexagon.h b/ggml/include/ggml-hexagon.h.inc similarity index 85% rename from ggml/include/ggml-hexagon.h rename to ggml/include/ggml-hexagon.h.inc index 6e0790041039..e47efeb20ab6 100644 --- a/ggml/include/ggml-hexagon.h +++ b/ggml/include/ggml-hexagon.h.inc @@ -1,7 +1,7 @@ #pragma once -#include "ggml.h" -#include "ggml-backend.h" +#include "ggml.h.inc" +#include "ggml-backend.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/include/ggml-metal.h b/ggml/include/ggml-metal.h.inc similarity index 97% rename from ggml/include/ggml-metal.h rename to ggml/include/ggml-metal.h.inc index 433838f0d6d6..3c103d2c53bf 100644 --- a/ggml/include/ggml-metal.h +++ b/ggml/include/ggml-metal.h.inc @@ -21,8 +21,8 @@ #pragma once -#include "ggml.h" -#include "ggml-backend.h" +#include "ggml.h.inc" +#include "ggml-backend.h.inc" #include #include diff --git a/ggml/include/ggml-opencl.h b/ggml/include/ggml-opencl.h.inc similarity index 91% rename from ggml/include/ggml-opencl.h rename to ggml/include/ggml-opencl.h.inc index 6b61771358f8..44a03868e0ed 100644 --- a/ggml/include/ggml-opencl.h +++ b/ggml/include/ggml-opencl.h.inc @@ -1,8 +1,8 @@ #ifndef GGML_OPENCL_H #define GGML_OPENCL_H -#include "ggml.h" -#include "ggml-backend.h" +#include "ggml.h.inc" +#include "ggml-backend.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/include/ggml-openvino.h b/ggml/include/ggml-openvino.h.inc similarity index 97% rename from ggml/include/ggml-openvino.h rename to ggml/include/ggml-openvino.h.inc index c43beb07b6a6..6546e8867fd2 100644 --- a/ggml/include/ggml-openvino.h +++ b/ggml/include/ggml-openvino.h.inc @@ -1,6 +1,6 @@ #pragma once -#include "ggml-backend.h" +#include "ggml-backend.h.inc" #include diff --git a/ggml/include/ggml-opt.h b/ggml/include/ggml-opt.h.inc similarity index 99% rename from ggml/include/ggml-opt.h rename to ggml/include/ggml-opt.h.inc index 1c2ed79b7742..a6fa5c87c377 100644 --- a/ggml/include/ggml-opt.h +++ b/ggml/include/ggml-opt.h.inc @@ -6,8 +6,8 @@ #pragma once -#include "ggml.h" -#include "ggml-backend.h" +#include "ggml.h.inc" +#include "ggml-backend.h.inc" #include diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h.inc similarity index 97% rename from ggml/include/ggml-rpc.h rename to ggml/include/ggml-rpc.h.inc index 6fcf5a433936..901ff74d87c8 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h.inc @@ -1,6 +1,6 @@ #pragma once -#include "ggml-backend.h" +#include "ggml-backend.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/include/ggml-sycl.h b/ggml/include/ggml-sycl.h.inc similarity index 97% rename from ggml/include/ggml-sycl.h rename to ggml/include/ggml-sycl.h.inc index 5ce349a880ed..c21311ab526c 100644 --- a/ggml/include/ggml-sycl.h +++ b/ggml/include/ggml-sycl.h.inc @@ -6,8 +6,8 @@ #pragma once -#include "ggml.h" -#include "ggml-backend.h" +#include "ggml.h.inc" +#include "ggml-backend.h.inc" #define GGML_SYCL_NAME "SYCL" #define GGML_SYCL_MAX_DEVICES 48 diff --git a/ggml/include/ggml-virtgpu.h b/ggml/include/ggml-virtgpu.h.inc similarity index 74% rename from ggml/include/ggml-virtgpu.h rename to ggml/include/ggml-virtgpu.h.inc index faaba8f246d1..9d0ccb395aa5 100644 --- a/ggml/include/ggml-virtgpu.h +++ b/ggml/include/ggml-virtgpu.h.inc @@ -1,7 +1,7 @@ #pragma once -#include "ggml.h" -#include "ggml-backend.h" +#include "ggml.h.inc" +#include "ggml-backend.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/include/ggml-vulkan.h b/ggml/include/ggml-vulkan.h.inc similarity index 94% rename from ggml/include/ggml-vulkan.h rename to ggml/include/ggml-vulkan.h.inc index ed5ea5f798cb..639a5a6fdb0b 100644 --- a/ggml/include/ggml-vulkan.h +++ b/ggml/include/ggml-vulkan.h.inc @@ -1,7 +1,7 @@ #pragma once -#include "ggml.h" -#include "ggml-backend.h" +#include "ggml.h.inc" +#include "ggml-backend.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/include/ggml-webgpu.h b/ggml/include/ggml-webgpu.h.inc similarity index 84% rename from ggml/include/ggml-webgpu.h rename to ggml/include/ggml-webgpu.h.inc index 65b8ed9bb664..4e697cb3067e 100644 --- a/ggml/include/ggml-webgpu.h +++ b/ggml/include/ggml-webgpu.h.inc @@ -1,7 +1,7 @@ #pragma once -#include "ggml.h" -#include "ggml-backend.h" +#include "ggml.h.inc" +#include "ggml-backend.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/include/ggml-zdnn.h b/ggml/include/ggml-zdnn.h.inc similarity index 82% rename from ggml/include/ggml-zdnn.h rename to ggml/include/ggml-zdnn.h.inc index fbf45b6e1c34..7ffd53643715 100644 --- a/ggml/include/ggml-zdnn.h +++ b/ggml/include/ggml-zdnn.h.inc @@ -1,7 +1,7 @@ #pragma once -#include "ggml.h" -#include "ggml-backend.h" +#include "ggml.h.inc" +#include "ggml-backend.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/include/ggml-zendnn.h b/ggml/include/ggml-zendnn.h.inc similarity index 89% rename from ggml/include/ggml-zendnn.h rename to ggml/include/ggml-zendnn.h.inc index a30a3a980883..d66f810b665c 100644 --- a/ggml/include/ggml-zendnn.h +++ b/ggml/include/ggml-zendnn.h.inc @@ -1,7 +1,7 @@ #pragma once -#include "ggml-backend.h" -#include "ggml.h" +#include "ggml-backend.h.inc" +#include "ggml.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h.inc similarity index 99% rename from ggml/include/ggml.h rename to ggml/include/ggml.h.inc index 703e37831361..c5f7fe251d58 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h.inc @@ -157,7 +157,7 @@ // TODO // // -// ## Overview of ggml.c +// ## Overview of ggml.c.inc // // TODO // @@ -2816,7 +2816,7 @@ extern "C" { bool paused; // start in paused state }; - struct ggml_threadpool; // forward declaration, see ggml.c + struct ggml_threadpool; // forward declaration, see ggml.c.inc typedef struct ggml_threadpool * ggml_threadpool_t; diff --git a/ggml/include/gguf.h b/ggml/include/gguf.h.inc similarity index 99% rename from ggml/include/gguf.h rename to ggml/include/gguf.h.inc index 02d5f221c03c..0d5073921755 100644 --- a/ggml/include/gguf.h +++ b/ggml/include/gguf.h.inc @@ -33,7 +33,7 @@ #pragma once -#include "ggml.h" +#include "ggml.h.inc" #include #include diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 3e48860bfc8b..e4c462fe16dd 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -189,24 +189,54 @@ if (GGML_BACKEND_DL AND NOT BUILD_SHARED_LIBS) message(FATAL_ERROR "GGML_BACKEND_DL requires BUILD_SHARED_LIBS") endif() +find_program(CARGO_EXECUTABLE cargo REQUIRED) +set(GGML_THREADING_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/ggml-threading-rust/Cargo.toml) +set(GGML_THREADING_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/ggml-threading) +set(GGML_THREADING_RUST_LIB ${GGML_THREADING_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}ggml_threading_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) +set(GGML_BACKEND_DL_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/ggml-backend-dl-rust/Cargo.toml) +set(GGML_BACKEND_DL_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/ggml-backend-dl) +set(GGML_BACKEND_DL_RUST_LIB ${GGML_BACKEND_DL_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}ggml_backend_dl_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) + +add_custom_command( + OUTPUT ${GGML_THREADING_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=-Crelocation-model=pic -Cpanic=abort" + ${CARGO_EXECUTABLE} build --manifest-path ${GGML_THREADING_RUST_MANIFEST} --target-dir ${GGML_THREADING_RUST_TARGET_DIR} + DEPENDS + ${GGML_THREADING_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/ggml-threading-rust/src/lib.rs +) +add_custom_target(ggml-threading-rust-build DEPENDS ${GGML_THREADING_RUST_LIB}) +add_library(ggml-threading-rust STATIC IMPORTED GLOBAL) +set_target_properties(ggml-threading-rust PROPERTIES IMPORTED_LOCATION ${GGML_THREADING_RUST_LIB}) +add_dependencies(ggml-threading-rust ggml-threading-rust-build) + +add_custom_command( + OUTPUT ${GGML_BACKEND_DL_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=-Crelocation-model=pic -Cpanic=abort" + ${CARGO_EXECUTABLE} build --manifest-path ${GGML_BACKEND_DL_RUST_MANIFEST} --target-dir ${GGML_BACKEND_DL_RUST_TARGET_DIR} + DEPENDS + ${GGML_BACKEND_DL_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/ggml-backend-dl-rust/src/lib.rs +) +add_custom_target(ggml-backend-dl-rust-build DEPENDS ${GGML_BACKEND_DL_RUST_LIB}) +add_library(ggml-backend-dl-rust STATIC IMPORTED GLOBAL) +set_target_properties(ggml-backend-dl-rust PROPERTIES IMPORTED_LOCATION ${GGML_BACKEND_DL_RUST_LIB}) +add_dependencies(ggml-backend-dl-rust ggml-backend-dl-rust-build) + +set_source_files_properties(ggml.c.inc PROPERTIES LANGUAGE C) +set_source_files_properties(ggml-backend.cpp.inc PROPERTIES LANGUAGE CXX) +set_source_files_properties(ggml-backend-reg.cpp.inc PROPERTIES LANGUAGE CXX) +set_source_files_properties(ggml-quants-defs.inc PROPERTIES HEADER_FILE_ONLY TRUE) + add_library(ggml-base - ../include/ggml.h - ../include/ggml-alloc.h - ../include/ggml-backend.h - ../include/ggml-cpp.h - ../include/ggml-opt.h - ../include/gguf.h - ggml.c - ggml.cpp - ggml-alloc.c - ggml-backend.cpp - ggml-backend-meta.cpp - ggml-opt.cpp - ggml-threading.cpp - ggml-threading.h - ggml-quants.c - ggml-quants.h - gguf.cpp) + ../include/ggml.h.inc + ../include/ggml-alloc.h.inc + ../include/ggml-backend.h.inc + ../include/ggml-opt.h.inc + ../include/gguf.h.inc + ggml.c.inc + ggml-backend.cpp.inc + ggml-quants-defs.inc) set_target_properties(ggml-base PROPERTIES VERSION ${GGML_VERSION} @@ -214,6 +244,7 @@ set_target_properties(ggml-base PROPERTIES ) target_include_directories(ggml-base PRIVATE .) +target_link_libraries(ggml-base PRIVATE ggml-threading-rust Threads::Threads ${CMAKE_DL_LIBS}) if (GGML_BACKEND_DL) target_compile_definitions(ggml-base PUBLIC GGML_BACKEND_DL) endif() @@ -223,8 +254,7 @@ if (GGML_SCHED_NO_REALLOC) endif() add_library(ggml - ggml-backend-dl.cpp - ggml-backend-reg.cpp) + ggml-backend-reg.cpp.inc) add_library(ggml::ggml ALIAS ggml) set_target_properties(ggml PROPERTIES @@ -239,7 +269,7 @@ if (GGML_BACKEND_DIR) target_compile_definitions(ggml PUBLIC GGML_BACKEND_DIR="${GGML_BACKEND_DIR}") endif() -target_link_libraries(ggml PUBLIC ggml-base) +target_link_libraries(ggml PUBLIC ggml-base PRIVATE ggml-backend-dl-rust) if (CMAKE_SYSTEM_NAME MATCHES "Linux") target_link_libraries(ggml PRIVATE dl) diff --git a/ggml/src/ggml-alloc.c b/ggml/src/ggml-alloc.c deleted file mode 100644 index a4b01ccf8a16..000000000000 --- a/ggml/src/ggml-alloc.c +++ /dev/null @@ -1,1248 +0,0 @@ -#include "ggml-alloc.h" -#include "ggml-backend-impl.h" -#include "ggml.h" -#include "ggml-impl.h" - -#include -#include -#include -#include -#include -#include - -#define MAX(a, b) ((a) > (b) ? (a) : (b)) -#define MAX_FREE_BLOCKS 256 - -//#define GGML_ALLOCATOR_DEBUG - -//#define AT_PRINTF(...) GGML_LOG_DEBUG(__VA_ARGS__) -#define AT_PRINTF(...) - -// ops that return true for this function must not use restrict pointers for their backend implementations -bool ggml_op_can_inplace(enum ggml_op op) { - switch (op) { - case GGML_OP_FILL: - case GGML_OP_SCALE: - case GGML_OP_DIAG_MASK_ZERO: - case GGML_OP_DIAG_MASK_INF: - case GGML_OP_ADD: - case GGML_OP_ADD_ID: - case GGML_OP_ADD1: - case GGML_OP_SUB: - case GGML_OP_MUL: - case GGML_OP_DIV: - case GGML_OP_SQR: - case GGML_OP_SQRT: - case GGML_OP_LOG: - case GGML_OP_UNARY: - case GGML_OP_ROPE: - case GGML_OP_ROPE_BACK: - case GGML_OP_SILU_BACK: - case GGML_OP_RMS_NORM: - case GGML_OP_RMS_NORM_BACK: - case GGML_OP_SOFT_MAX: - case GGML_OP_SOFT_MAX_BACK: - return true; - - default: - return false; - } -} - -static size_t aligned_offset(const void * buffer, size_t offset, size_t alignment) { - assert(alignment && !(alignment & (alignment - 1))); // power of 2 - size_t align = (alignment - (((uintptr_t)buffer + offset) % alignment)) % alignment; - return offset + align; -} - -// tallocr - -struct ggml_tallocr ggml_tallocr_new(ggml_backend_buffer_t buffer) { - void * base = ggml_backend_buffer_get_base(buffer); - size_t align = ggml_backend_buffer_get_alignment(buffer); - - assert(align && !(align & (align - 1))); // power of 2 - - struct ggml_tallocr talloc = (struct ggml_tallocr) { - /*.buffer = */ buffer, - /*.base = */ base, - /*.alignment = */ align, - /*.offset = */ aligned_offset(base, 0, align), - }; - return talloc; -} - -enum ggml_status ggml_tallocr_alloc(struct ggml_tallocr * talloc, struct ggml_tensor * tensor) { - size_t size = ggml_backend_buffer_get_alloc_size(talloc->buffer, tensor); - size = GGML_PAD(size, talloc->alignment); - - if (talloc->offset + size > ggml_backend_buffer_get_size(talloc->buffer)) { - GGML_LOG_ERROR("%s: not enough space in the buffer to allocate %s (needed %zu, available %zu)\n", - __func__, tensor->name, size, ggml_backend_buffer_get_size(talloc->buffer) - talloc->offset); - GGML_ABORT("not enough space in the buffer"); - } - - void * addr = (char *)ggml_backend_buffer_get_base(talloc->buffer) + talloc->offset; - talloc->offset += size; - - assert(((uintptr_t)addr % talloc->alignment) == 0); - - return ggml_backend_tensor_alloc(talloc->buffer, tensor, addr); -} - -// dynamic tensor allocator - -#define GGML_VBUFFER_MAX_CHUNKS 16 - -// relative memory address within an allocation that can be split into multiple buffers (chunks) -struct buffer_address { - int chunk; // index of a backend buffer - size_t offset; // local memory offset within the buffer -}; - -static const struct buffer_address GGML_BUFFER_ADDRESS_INVALID = { -1, SIZE_MAX }; - -static bool ggml_buffer_address_less(struct buffer_address a, struct buffer_address b) { - return a.chunk != b.chunk ? a.chunk < b.chunk : a.offset < b.offset; -} - -struct free_block { - size_t offset; - size_t size; -}; - -struct tallocr_chunk { - struct free_block free_blocks[MAX_FREE_BLOCKS]; - int n_free_blocks; - size_t max_size; -}; - -struct ggml_dyn_tallocr { - size_t alignment; - size_t max_chunk_size; - struct tallocr_chunk * chunks[GGML_VBUFFER_MAX_CHUNKS]; - int n_chunks; - -#ifdef GGML_ALLOCATOR_DEBUG - struct { - const struct ggml_tensor * tensor; - struct buffer_address addr; - } allocated_tensors[1024]; -#endif -}; - -static void ggml_dyn_tallocr_insert_block(struct tallocr_chunk * chunk, size_t offset, size_t size) { - GGML_ASSERT(chunk->n_free_blocks < MAX_FREE_BLOCKS && "out of free blocks"); - // insert the new block in the correct position to keep the array sorted by address (to make merging blocks faster) - int insert_pos = 0; - while (insert_pos < chunk->n_free_blocks && chunk->free_blocks[insert_pos].offset < offset) { - insert_pos++; - } - // shift all blocks from insert_pos onward to make room for the new block - for (int i = chunk->n_free_blocks; i > insert_pos; i--) { - chunk->free_blocks[i] = chunk->free_blocks[i-1]; - } - // insert the new block - chunk->free_blocks[insert_pos].offset = offset; - chunk->free_blocks[insert_pos].size = size; - chunk->n_free_blocks++; -} - -static void ggml_dyn_tallocr_remove_block(struct tallocr_chunk * chunk, int idx) { - // shift all elements after idx by 1 to the left, overwriting the element at idx - for (int i = idx; i < chunk->n_free_blocks; i++) { - chunk->free_blocks[i] = chunk->free_blocks[i+1]; - } - chunk->n_free_blocks--; -} - -static int ggml_dyn_tallocr_new_chunk(struct ggml_dyn_tallocr * alloc, size_t min_size) { - if (alloc->n_chunks >= GGML_VBUFFER_MAX_CHUNKS) { - return -1; - } - struct tallocr_chunk * chunk = calloc(1, sizeof(struct tallocr_chunk)); - chunk->n_free_blocks = 1; - chunk->free_blocks[0].offset = 0; - // available space in a chunk is limited to max_chunk_size, but can be higher if: - // 1. a single tensor exceeds the maximum, and cannot fit any other way - // 2. we are running out of chunks - // backends will either manage to allocate the larger size, or report an error. - chunk->free_blocks[0].size = MAX(min_size, alloc->max_chunk_size); - if (alloc->n_chunks == GGML_VBUFFER_MAX_CHUNKS - 1) { - chunk->free_blocks[0].size = SIZE_MAX/2; - } - alloc->chunks[alloc->n_chunks] = chunk; - alloc->n_chunks++; - return alloc->n_chunks - 1; -} - -#ifdef GGML_ALLOCATOR_DEBUG -static void add_allocated_tensor(struct ggml_dyn_tallocr * alloc, struct buffer_address addr, const struct ggml_tensor * tensor) { - for (int i = 0; i < 1024; i++) { - if (alloc->allocated_tensors[i].tensor == NULL) { - alloc->allocated_tensors[i].tensor = tensor; - alloc->allocated_tensors[i].addr = addr; - return; - } - } - GGML_ABORT("out of allocated_tensors"); -} -static void remove_allocated_tensor(struct ggml_dyn_tallocr * alloc, struct buffer_address addr, const struct ggml_tensor * tensor) { - for (int i = 0; i < 1024; i++) { - if (alloc->allocated_tensors[i].addr.chunk == addr.chunk && alloc->allocated_tensors[i].addr.offset == addr.offset) { - alloc->allocated_tensors[i].tensor = NULL; - return; - } - } - GGML_ABORT("tried to free tensor %s not found\n", tensor->name); -} -#endif - -static struct buffer_address ggml_dyn_tallocr_alloc(struct ggml_dyn_tallocr * alloc, size_t size, const struct ggml_tensor * tensor) { - size = aligned_offset(NULL, size, alloc->alignment); - - AT_PRINTF("%s: allocating %s (%zu bytes) - ", __func__, tensor->name, size); - - int best_fit_chunk = -1; - int best_fit_block = -1; - size_t max_avail = 0; - - // find the best fitting free block besides the last block, within any chunk - for (int c = 0; c < alloc->n_chunks; ++c) { - struct tallocr_chunk * chunk = alloc->chunks[c]; - size_t best_fit_size = SIZE_MAX; - for (int i = 0; i < chunk->n_free_blocks - 1; i++) { - struct free_block * block = &chunk->free_blocks[i]; - max_avail = MAX(max_avail, block->size); - if (block->size >= size && block->size <= best_fit_size) { - best_fit_chunk = c; - best_fit_block = i; - best_fit_size = block->size; - } - } - } - - if (best_fit_block == -1) { - // no suitable block found, try the last block (this may grow a chunks size) - int64_t best_reuse = INT64_MIN; - for (int c = 0; c < alloc->n_chunks; ++c) { - struct tallocr_chunk * chunk = alloc->chunks[c]; - if (chunk->n_free_blocks > 0) { - struct free_block * block = &chunk->free_blocks[chunk->n_free_blocks - 1]; - max_avail = MAX(max_avail, block->size); - int64_t reuse_factor = chunk->max_size - block->offset - size; - // reuse_factor < 0 : amount of extra memory that needs to be allocated - // reuse_factor = 0 : allocated free space exactly matches tensor size - // reuse_factor > 0 : superfluous memory that will remain unused - bool better_reuse = best_reuse < 0 && reuse_factor > best_reuse; - bool better_fit = reuse_factor >= 0 && reuse_factor < best_reuse; - if (block->size >= size && (better_reuse || better_fit)) { - best_fit_chunk = c; - best_fit_block = chunk->n_free_blocks - 1; - best_reuse = reuse_factor; - } - } - } - } - - if (best_fit_block == -1) { - // none of the existing chunks have enough space left - best_fit_chunk = ggml_dyn_tallocr_new_chunk(alloc, size); - best_fit_block = 0; - } - if (best_fit_chunk == -1) { - // since the last chunk always has virtually endless memory, this should never happen - GGML_LOG_ERROR("%s: not enough space in the buffer to allocate %zu bytes, largest block available %zu bytes\n", - __func__, size, max_avail); - GGML_ABORT("graph allocation: failed to reserve memory"); - } - - struct tallocr_chunk * chunk = alloc->chunks[best_fit_chunk]; - struct free_block * block = &chunk->free_blocks[best_fit_block]; - struct buffer_address addr = {.chunk = best_fit_chunk, .offset = block->offset }; - block->offset += size; - block->size -= size; - if (block->size == 0) { - // remove block if empty - ggml_dyn_tallocr_remove_block(chunk, best_fit_block); - } - - AT_PRINTF("block %d, offset %zu, chunk %d\n", best_fit_block, addr.offset, addr.chunk); - -#ifdef GGML_ALLOCATOR_DEBUG - add_allocated_tensor(alloc, addr, tensor); - size_t cur_max = addr.offset + size; - if (cur_max > chunk->max_size) { - // sort allocated_tensors by chunk/offset - for (int i = 0; i < 1024; i++) { - for (int j = i + 1; j < 1024; j++) { - if (ggml_buffer_address_less(alloc->allocated_tensors[j].addr, alloc->allocated_tensors[i].addr)) { - const struct ggml_tensor * tmp_tensor = alloc->allocated_tensors[i].tensor; - struct buffer_address tmp_addr = alloc->allocated_tensors[i].addr; - alloc->allocated_tensors[i].tensor = alloc->allocated_tensors[j].tensor; - alloc->allocated_tensors[i].addr = alloc->allocated_tensors[j].addr; - alloc->allocated_tensors[j].tensor = tmp_tensor; - alloc->allocated_tensors[j].addr = tmp_addr; - } - } - } - GGML_LOG_DEBUG("max_size[%d] = %.2f MB: tensors: ", addr.chunk, cur_max / 1024.0 / 1024.0); - for (int i = 0; i < 1024; i++) { - if (alloc->allocated_tensors[i].tensor) { - GGML_LOG_DEBUG("%s [%d: %zx-%zx] (%.2f MB) ", alloc->allocated_tensors[i].tensor->name, - alloc->allocated_tensors[i].addr.chunk, - alloc->allocated_tensors[i].addr.offset, - alloc->allocated_tensors[i].addr.offset + ggml_nbytes(alloc->allocated_tensors[i].tensor), - ggml_nbytes(alloc->allocated_tensors[i].tensor) / 1024.0 / 1024.0); - } - } - GGML_LOG_DEBUG("\n"); - } -#endif - - chunk->max_size = MAX(chunk->max_size, addr.offset + size); - - return addr; - - GGML_UNUSED(tensor); -} - -// this is a very naive implementation, but for our case the number of free blocks should be very small -static void ggml_dyn_tallocr_free_bytes(struct ggml_dyn_tallocr * alloc, struct buffer_address addr, size_t size) { - size = aligned_offset(NULL, size, alloc->alignment); - - struct tallocr_chunk * chunk = alloc->chunks[addr.chunk]; - - // see if we can merge with an existing block - for (int i = 0; i < chunk->n_free_blocks; i++) { - struct free_block * block = &chunk->free_blocks[i]; - // check if ptr is at the end of the block - if (block->offset + block->size == addr.offset) { - block->size += size; - // check if we can merge with the next block - if (i < chunk->n_free_blocks - 1) { - struct free_block * next = &chunk->free_blocks[i+1]; - if (block->offset + block->size == next->offset) { - block->size += next->size; - ggml_dyn_tallocr_remove_block(chunk, i+1); - } - } - return; - } - // check if ptr is at the beginning of the block - if (addr.offset + size == block->offset) { - block->offset = addr.offset; - block->size += size; - // check if we can merge with the previous block - if (i > 0) { - struct free_block * prev = &chunk->free_blocks[i-1]; - if (prev->offset + prev->size == block->offset) { - prev->size += block->size; - ggml_dyn_tallocr_remove_block(chunk, i); - } - } - return; - } - } - // otherwise, add a new block - ggml_dyn_tallocr_insert_block(chunk, addr.offset, size); -} - -static void ggml_dyn_tallocr_reset(struct ggml_dyn_tallocr * alloc) { - for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS; i++) { - free(alloc->chunks[i]); - alloc->chunks[i] = NULL; - } - alloc->n_chunks = 0; - -#ifdef GGML_ALLOCATOR_DEBUG - for (int i = 0; i < 1024; i++) { - alloc->allocated_tensors[i].tensor = NULL; - } -#endif -} - -static struct ggml_dyn_tallocr * ggml_dyn_tallocr_new(size_t alignment, size_t max_buffer_size) { - struct ggml_dyn_tallocr * alloc = (struct ggml_dyn_tallocr *)malloc(sizeof(struct ggml_dyn_tallocr)); - - *alloc = (struct ggml_dyn_tallocr) { - /*.alignment = */ alignment, - /*.max_chunk_size = */ MIN(max_buffer_size, SIZE_MAX/2), // clamp to avoid overflows - /*.chunks = */ {NULL}, - /*.n_chunks = */ 0, -#ifdef GGML_ALLOCATOR_DEBUG - /*.allocated_tensors = */ {{0}}, -#endif - }; - - ggml_dyn_tallocr_reset(alloc); - - return alloc; -} - -static void ggml_dyn_tallocr_free(struct ggml_dyn_tallocr * alloc) { - for (int i = 0; i < alloc->n_chunks; ++i) { - free(alloc->chunks[i]); - } - free(alloc); -} - -static size_t ggml_dyn_tallocr_max_size(struct ggml_dyn_tallocr * alloc, int chunk) { - return chunk < alloc->n_chunks ? alloc->chunks[chunk]->max_size : 0; -} - - -// virtual buffer with contiguous memory range, split into multiple backend buffers (chunks) - -struct vbuffer { - ggml_backend_buffer_t chunks[GGML_VBUFFER_MAX_CHUNKS]; -}; - -static void ggml_vbuffer_free(struct vbuffer * buf) { - if (buf == NULL) { - return; - } - for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS; ++i) { - ggml_backend_buffer_free(buf->chunks[i]); - } - free(buf); -} - -static size_t ggml_vbuffer_chunk_size(struct vbuffer * buf, int chunk) { - return buf->chunks[chunk] ? ggml_backend_buffer_get_size(buf->chunks[chunk]) : 0; -} - -static size_t ggml_vbuffer_size(struct vbuffer * buf) { - size_t size = 0; - for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS && buf->chunks[i]; ++i) { - size += ggml_backend_buffer_get_size(buf->chunks[i]); - } - return size; -} - -static struct vbuffer * ggml_vbuffer_alloc(ggml_backend_buffer_type_t buft, const struct ggml_dyn_tallocr * talloc, enum ggml_backend_buffer_usage usage) { - struct vbuffer * buf = (struct vbuffer *)calloc(1, sizeof(struct vbuffer)); - if (buf == NULL) { - return NULL; - } - - for (int n = 0; n < talloc->n_chunks; n++) { - size_t chunk_size = talloc->chunks[n]->max_size; - buf->chunks[n] = ggml_backend_buft_alloc_buffer(buft, chunk_size); - if (buf->chunks[n] == NULL) { - ggml_vbuffer_free(buf); - return NULL; - } - ggml_backend_buffer_set_usage(buf->chunks[n], usage); - } - return buf; -} - -static void ggml_vbuffer_tensor_alloc(struct vbuffer * buf, struct ggml_tensor * tensor, struct buffer_address buf_addr) { - void * base = ggml_backend_buffer_get_base(buf->chunks[buf_addr.chunk]); - void * addr = (char *)base + buf_addr.offset; - ggml_backend_tensor_alloc(buf->chunks[buf_addr.chunk], tensor, addr); -} - -static void ggml_vbuffer_reset(struct vbuffer * buf) { - for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS && buf->chunks[i]; ++i) { - ggml_backend_buffer_reset(buf->chunks[i]); - } -} - - -///////////////////////////////////// - -// graph allocator - -struct hash_node { - int n_children; - int n_views; - int buffer_id; - struct buffer_address addr; - bool allocated; -}; - -struct tensor_alloc { - int buffer_id; - struct buffer_address addr; - size_t size_max; // 0 = pre-allocated, unused, or view -}; - -struct leaf_alloc { - struct tensor_alloc leaf; -}; - -struct node_alloc { - struct tensor_alloc dst; - struct tensor_alloc src[GGML_MAX_SRC]; -}; - -struct ggml_gallocr { - ggml_backend_buffer_type_t * bufts; // [n_buffers] - struct vbuffer ** buffers; // [n_buffers] - struct ggml_dyn_tallocr ** buf_tallocs; // [n_buffers] - int n_buffers; - - struct ggml_hash_set hash_set; - struct hash_node * hash_values; // [hash_set.size] - - struct node_alloc * node_allocs; // [n_nodes] - int n_nodes; - - struct leaf_alloc * leaf_allocs; // [n_leafs] - int n_leafs; -}; - -ggml_gallocr_t ggml_gallocr_new_n(ggml_backend_buffer_type_t * bufts, int n_bufs) { - ggml_gallocr_t galloc = (ggml_gallocr_t)calloc(1, sizeof(struct ggml_gallocr)); - GGML_ASSERT(galloc != NULL); - - galloc->bufts = calloc(n_bufs, sizeof(ggml_backend_buffer_type_t)); - GGML_ASSERT(galloc->bufts != NULL); - - galloc->buffers = calloc(n_bufs, sizeof(struct vbuffer *)); - GGML_ASSERT(galloc->buffers != NULL); - - galloc->buf_tallocs = calloc(n_bufs, sizeof(struct ggml_dyn_tallocr *)); - GGML_ASSERT(galloc->buf_tallocs != NULL); - - for (int i = 0; i < n_bufs; i++) { - galloc->bufts[i] = bufts[i]; - galloc->buffers[i] = NULL; - - // check if the same buffer type is used multiple times and reuse the same allocator - for (int j = 0; j < i; j++) { - if (bufts[i] == bufts[j]) { - galloc->buf_tallocs[i] = galloc->buf_tallocs[j]; - break; - } - } - - if (galloc->buf_tallocs[i] == NULL) { - size_t alignment = ggml_backend_buft_get_alignment(bufts[i]); - size_t max_size = ggml_backend_buft_get_max_size(bufts[i]); - galloc->buf_tallocs[i] = ggml_dyn_tallocr_new(alignment, max_size); - } - } - galloc->n_buffers = n_bufs; - - return galloc; -} - -ggml_gallocr_t ggml_gallocr_new(ggml_backend_buffer_type_t buft) { - return ggml_gallocr_new_n(&buft, 1); -} - -void ggml_gallocr_free(ggml_gallocr_t galloc) { - if (galloc == NULL) { - return; - } - - for (int i = 0; i < galloc->n_buffers; i++) { - if (galloc->buffers != NULL) { - // skip if already freed - bool freed = false; - for (int j = 0; j < i; j++) { - if (galloc->buffers[j] == galloc->buffers[i]) { - freed = true; - break; - } - } - if (!freed) { - ggml_vbuffer_free(galloc->buffers[i]); - } - } - if (galloc->buf_tallocs != NULL) { - // skip if already freed - bool freed = false; - for (int j = 0; j < i; j++) { - if (galloc->buf_tallocs[j] == galloc->buf_tallocs[i]) { - freed = true; - break; - } - } - if (!freed) { - ggml_dyn_tallocr_free(galloc->buf_tallocs[i]); - } - } - } - - ggml_hash_set_free(&galloc->hash_set); - free(galloc->hash_values); - free(galloc->bufts); - free(galloc->buffers); - free(galloc->buf_tallocs); - free(galloc->node_allocs); - free(galloc->leaf_allocs); - free(galloc); -} - -typedef struct ggml_gallocr * ggml_gallocr_t; - -static struct hash_node * ggml_gallocr_hash_get(ggml_gallocr_t galloc, struct ggml_tensor * t) { - size_t i = ggml_hash_find_or_insert(&galloc->hash_set, t); - return &galloc->hash_values[i]; -} - -static bool ggml_gallocr_is_own(ggml_gallocr_t galloc, struct ggml_tensor * t) { - return ggml_gallocr_hash_get(galloc, t)->allocated; -} - -static bool ggml_gallocr_is_allocated(ggml_gallocr_t galloc, struct ggml_tensor * t) { - return t->data != NULL // tensor data already set externally - || t->buffer // tensor on external buffer (but not yet allocated) - || ggml_gallocr_is_own(galloc, t); // tensor will be allocated by galloc -} - -// free the extra space at the end if the new tensor is smaller -static void ggml_gallocr_free_extra_space(ggml_gallocr_t galloc, struct ggml_tensor * node, struct ggml_tensor * parent) { - struct hash_node * hn = ggml_gallocr_hash_get(galloc, node); - struct hash_node * p_hn = ggml_gallocr_hash_get(galloc, parent); - - size_t parent_size = ggml_backend_buft_get_alloc_size(galloc->bufts[p_hn->buffer_id], parent); - size_t node_size = ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], node); - - GGML_ASSERT(parent_size >= node_size); - - // note: we want after the freeing the chunks to continue to be aligned - struct ggml_dyn_tallocr * p_alloc = galloc->buf_tallocs[p_hn->buffer_id]; - parent_size = aligned_offset(NULL, parent_size, p_alloc->alignment); - node_size = aligned_offset(NULL, node_size, p_alloc->alignment); - - if (parent_size > node_size) { - struct buffer_address p_addr = p_hn->addr; - p_addr.offset += node_size; - size_t extra_size = parent_size - node_size; - AT_PRINTF("freeing extra %zu bytes from parent %s for %s\n", extra_size, parent->name, node->name); - ggml_dyn_tallocr_free_bytes(p_alloc, p_addr, extra_size); - } -} - -static void ggml_gallocr_allocate_node(ggml_gallocr_t galloc, struct ggml_tensor * node, int buffer_id) { - GGML_ASSERT(buffer_id >= 0); - struct hash_node * hn = ggml_gallocr_hash_get(galloc, node); - - if (!ggml_gallocr_is_allocated(galloc, node) && !ggml_impl_is_view(node)) { - hn->allocated = true; - assert(hn->addr.offset == 0); - - // try to reuse a parent's buffer (inplace) - if (ggml_op_can_inplace(node->op)) { - for (int i = 0; i < GGML_MAX_SRC; i++) { - struct ggml_tensor * parent = node->src[i]; - if (parent == NULL) { - continue; - } - - // if the node's data is external, then we cannot re-use it - if (!ggml_gallocr_is_own(galloc, parent)) { - AT_PRINTF("not reusing parent %s for %s as %p is external\n", parent->name, node->name, parent->data); - continue; - } - - // outputs cannot be reused - if (parent->flags & GGML_TENSOR_FLAG_OUTPUT || (parent->view_src != NULL && parent->view_src->flags & GGML_TENSOR_FLAG_OUTPUT)) { - AT_PRINTF("not reusing parent %s for %s as it is an output\n", parent->name, node->name); - continue; - } - - if (!ggml_are_same_layout(node, parent)) { - AT_PRINTF("not reusing parent %s for %s as layouts are different\n", parent->name, node->name); - continue; - } - - struct hash_node * p_hn = ggml_gallocr_hash_get(galloc, parent); - if (p_hn->n_children == 1 && p_hn->n_views == 0) { - if (ggml_impl_is_view(parent)) { - struct ggml_tensor * view_src = parent->view_src; - struct hash_node * view_src_hn = ggml_gallocr_hash_get(galloc, view_src); - if (view_src_hn->n_views == 1 && view_src_hn->n_children == 0 && view_src->data == parent->data) { - AT_PRINTF("reusing view parent %s (%s) for %s\n", parent->name, view_src->name, node->name); - assert(view_src_hn->addr.chunk == p_hn->addr.chunk && view_src_hn->addr.offset == p_hn->addr.offset); - hn->buffer_id = p_hn->buffer_id; - hn->addr = p_hn->addr; - p_hn->allocated = false; // avoid freeing the parent - view_src_hn->allocated = false; - ggml_gallocr_free_extra_space(galloc, node, view_src); - return; - } - } else { - AT_PRINTF("reusing parent %s for %s\n", parent->name, node->name); - hn->buffer_id = p_hn->buffer_id; - hn->addr = p_hn->addr; - p_hn->allocated = false; // avoid freeing the parent - ggml_gallocr_free_extra_space(galloc, node, parent); - return; - } - } - } - } - // allocate tensor from the buffer - struct ggml_dyn_tallocr * alloc = galloc->buf_tallocs[buffer_id]; - ggml_backend_buffer_type_t buft = galloc->bufts[buffer_id]; - size_t size = ggml_backend_buft_get_alloc_size(buft, node); - hn->buffer_id = buffer_id; - hn->addr = ggml_dyn_tallocr_alloc(alloc, size, node); - } -} - -static void ggml_gallocr_free_node(ggml_gallocr_t galloc, struct ggml_tensor * node) { - // graph outputs are never freed - if (node->flags & GGML_TENSOR_FLAG_OUTPUT) { - AT_PRINTF("not freeing output %s\n", node->name); - return; - } - - struct hash_node * hn = ggml_gallocr_hash_get(galloc, node); - int buffer_id = hn->buffer_id; - struct ggml_dyn_tallocr * alloc = galloc->buf_tallocs[buffer_id]; - ggml_backend_buffer_type_t buft = galloc->bufts[buffer_id]; - size_t size = ggml_backend_buft_get_alloc_size(buft, node); - - AT_PRINTF("%s: freeing %s at {chunk=%d, offset=%zu} (%zu bytes) - n_free_blocks = %d\n", - __func__, node->name, hn->addr.chunk, hn->addr.offset, size, alloc->chunks[hn->addr.chunk]->n_free_blocks); -#ifdef GGML_ALLOCATOR_DEBUG - remove_allocated_tensor(alloc, hn->addr, node); -#endif - - ggml_dyn_tallocr_free_bytes(alloc, hn->addr, size); - hn->allocated = false; -} - -static int get_node_buffer_id(const int * node_buffer_ids, int i) { - return node_buffer_ids ? node_buffer_ids[i] : 0; -} - -static void ggml_gallocr_alloc_graph_impl(ggml_gallocr_t galloc, struct ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids) { - // clear hash tables - ggml_hash_set_reset(&galloc->hash_set); - memset(galloc->hash_values, 0, sizeof(struct hash_node) * galloc->hash_set.size); - - // allocate leafs - // these may be tensors that the application is not using in the graph, but may still want to allocate for other purposes - for (int i = 0; i < graph->n_leafs; i++) { - struct ggml_tensor * leaf = graph->leafs[i]; - ggml_gallocr_allocate_node(galloc, leaf, get_node_buffer_id(leaf_buffer_ids, i)); - } - - // count number of children and views - // allocate other graph inputs and leafs first to avoid overwriting them - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - - // TODO: better way to add external dependencies - // GGML_OP_NONE does not appear normally in the graph nodes, but is used by ggml-backend to add dependencies to - // control when some tensors are allocated and freed. in this case, the dependencies are in `src`, but the node - // itself is never used and should not be considered a dependency - if (ggml_impl_is_view(node) && node->op != GGML_OP_NONE) { - struct ggml_tensor * view_src = node->view_src; - ggml_gallocr_hash_get(galloc, view_src)->n_views += 1; - } - - if (node->flags & GGML_TENSOR_FLAG_INPUT) { - ggml_gallocr_allocate_node(galloc, graph->nodes[i], get_node_buffer_id(node_buffer_ids, i)); - } - - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * src = node->src[j]; - if (src == NULL) { - continue; - } - - ggml_gallocr_hash_get(galloc, src)->n_children += 1; - - // allocate explicit inputs - if (src->flags & GGML_TENSOR_FLAG_INPUT) { - ggml_gallocr_allocate_node(galloc, src, get_node_buffer_id(node_buffer_ids, i)); - } - } - } - - // allocate tensors - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - int buffer_id = get_node_buffer_id(node_buffer_ids, i); - - // allocate parents (only leafs need to be allocated at this point) - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * parent = node->src[j]; - if (parent == NULL) { - continue; - } - ggml_gallocr_allocate_node(galloc, parent, buffer_id); - } - - // allocate node - ggml_gallocr_allocate_node(galloc, node, buffer_id); - - AT_PRINTF("exec: %s (%s) <= ", ggml_op_desc(node), node->name); - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * parent = node->src[j]; - if (parent == NULL) { - continue; - } - AT_PRINTF("%s", parent->name); - if (j < GGML_MAX_SRC - 1 && node->src[j + 1] != NULL) { - AT_PRINTF(", "); - } - } - AT_PRINTF("\n"); - - // update parents - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * parent = node->src[j]; - if (parent == NULL) { - continue; - } - struct hash_node * p_hn = ggml_gallocr_hash_get(galloc, parent); - p_hn->n_children -= 1; - - AT_PRINTF("parent %s: %d children, %d views, allocated: %d\n", - parent->name, p_hn->n_children, p_hn->n_views, p_hn->allocated); - - if (p_hn->n_children == 0 && p_hn->n_views == 0) { - if (ggml_impl_is_view(parent)) { - struct ggml_tensor * view_src = parent->view_src; - struct hash_node * view_src_hn = ggml_gallocr_hash_get(galloc, view_src); - view_src_hn->n_views -= 1; - AT_PRINTF("view_src %s: %d children, %d views\n", - view_src->name, view_src_hn->n_children, view_src_hn->n_views); - if (view_src_hn->n_views == 0 && view_src_hn->n_children == 0 && view_src_hn->allocated) { - ggml_gallocr_free_node(galloc, view_src); - } - } - else if (p_hn->allocated) { - ggml_gallocr_free_node(galloc, parent); - } - } - AT_PRINTF("\n"); - } - } -} - -static bool ggml_gallocr_reserve_n_impl( - ggml_gallocr_t galloc, struct ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids, bool no_alloc) { - size_t min_hash_size = graph->n_nodes + graph->n_leafs; - // add 25% margin to avoid hash collisions - min_hash_size += min_hash_size / 4; - - // initialize hash table - if (galloc->hash_set.size < min_hash_size) { - ggml_hash_set_free(&galloc->hash_set); - galloc->hash_set = ggml_hash_set_new(min_hash_size); - GGML_ASSERT(galloc->hash_set.keys != NULL); - - free(galloc->hash_values); - galloc->hash_values = malloc(sizeof(struct hash_node) * galloc->hash_set.size); - GGML_ASSERT(galloc->hash_values != NULL); - } - - // reset allocators - for (int i = 0; i < galloc->n_buffers; i++) { - ggml_dyn_tallocr_reset(galloc->buf_tallocs[i]); - } - - // allocate in hash table - ggml_gallocr_alloc_graph_impl(galloc, graph, node_buffer_ids, leaf_buffer_ids); - - // set the node_allocs from the hash table - if (galloc->n_nodes < graph->n_nodes) { - free(galloc->node_allocs); - galloc->node_allocs = calloc(graph->n_nodes, sizeof(struct node_alloc)); - GGML_ASSERT(galloc->node_allocs != NULL); - } - galloc->n_nodes = graph->n_nodes; - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - struct node_alloc * node_alloc = &galloc->node_allocs[i]; - if (node->view_src || node->data) { - node_alloc->dst.buffer_id = -1; - node_alloc->dst.addr = GGML_BUFFER_ADDRESS_INVALID; - node_alloc->dst.size_max = 0; - } else { - struct hash_node * hn = ggml_gallocr_hash_get(galloc, node); - node_alloc->dst.buffer_id = hn->buffer_id; - node_alloc->dst.addr = hn->addr; - node_alloc->dst.size_max = ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], node); - } - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * src = node->src[j]; - if (!src || src->view_src || src->data) { - node_alloc->src[j].buffer_id = -1; - node_alloc->src[j].addr = GGML_BUFFER_ADDRESS_INVALID; - node_alloc->src[j].size_max = 0; - } else { - struct hash_node * hn = ggml_gallocr_hash_get(galloc, src); - node_alloc->src[j].buffer_id = hn->buffer_id; - node_alloc->src[j].addr = hn->addr; - node_alloc->src[j].size_max = ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], src); - } - } - } - if (galloc->n_leafs < graph->n_leafs) { - free(galloc->leaf_allocs); - galloc->leaf_allocs = calloc(graph->n_leafs, sizeof(galloc->leaf_allocs[0])); - GGML_ASSERT(galloc->leaf_allocs != NULL); - } - galloc->n_leafs = graph->n_leafs; - for (int i = 0; i < graph->n_leafs; i++) { - struct ggml_tensor * leaf = graph->leafs[i]; - struct hash_node * hn = ggml_gallocr_hash_get(galloc, leaf); - if (leaf->view_src || leaf->data) { - galloc->leaf_allocs[i].leaf.buffer_id = -1; - galloc->leaf_allocs[i].leaf.addr = GGML_BUFFER_ADDRESS_INVALID; - galloc->leaf_allocs[i].leaf.size_max = 0; - } else { - galloc->leaf_allocs[i].leaf.buffer_id = hn->buffer_id; - galloc->leaf_allocs[i].leaf.addr = hn->addr; - galloc->leaf_allocs[i].leaf.size_max = ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], leaf); - } - } - - // reallocate buffers if needed - for (int i = 0; i < galloc->n_buffers; i++) { - // if the buffer type is used multiple times, we reuse the same buffer - for (int j = 0; j < i; j++) { - if (galloc->buf_tallocs[j] == galloc->buf_tallocs[i]) { - galloc->buffers[i] = galloc->buffers[j]; - break; - } - } - - // even if there are no tensors allocated in this buffer, we still need to allocate it to initialize views - bool realloc = galloc->buffers[i] == NULL; - size_t new_size = 0; - for (int c = 0; c < galloc->buf_tallocs[i]->n_chunks; c++) { - size_t cur_chunk_size = galloc->buffers[i] ? ggml_vbuffer_chunk_size(galloc->buffers[i], c) : 0; - size_t new_chunk_size = ggml_dyn_tallocr_max_size(galloc->buf_tallocs[i], c); - new_size += new_chunk_size; - if (new_chunk_size > cur_chunk_size) { - realloc = true; - } - } - if (realloc) { -#ifndef NDEBUG - { - size_t cur_size = galloc->buffers[i] ? ggml_vbuffer_size(galloc->buffers[i]) : 0; - if (cur_size > 0) { - GGML_LOG_DEBUG("%s: reallocating %s buffer from size %.02f MiB to %.02f MiB\n", - __func__, ggml_backend_buft_name(galloc->bufts[i]), cur_size / 1024.0 / 1024.0, new_size / 1024.0 / 1024.0); - } - } -#endif - ggml_vbuffer_free(galloc->buffers[i]); - if (no_alloc) { - galloc->buffers[i] = NULL; - } else { - galloc->buffers[i] = ggml_vbuffer_alloc(galloc->bufts[i], galloc->buf_tallocs[i], GGML_BACKEND_BUFFER_USAGE_COMPUTE); - if (galloc->buffers[i] == NULL) { - GGML_LOG_ERROR("%s: failed to allocate %s buffer of size %zu\n", __func__, ggml_backend_buft_name(galloc->bufts[i]), new_size); - return false; - } - } - } - } - - return true; -} - -void ggml_gallocr_reserve_n_size( - ggml_gallocr_t galloc, struct ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids, size_t * sizes) { - GGML_ASSERT(ggml_gallocr_reserve_n_impl(galloc, graph, node_buffer_ids, leaf_buffer_ids, /*no_alloc =*/ true)); - for (int i = 0; i < galloc->n_buffers; i++) { - sizes[i] = 0; - for (int c = 0; c < galloc->buf_tallocs[i]->n_chunks; c++) { - sizes[i] += galloc->buf_tallocs[i]->chunks[c]->max_size; - } - } -} - -bool ggml_gallocr_reserve_n(ggml_gallocr_t galloc, struct ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids) { - return ggml_gallocr_reserve_n_impl(galloc, graph, node_buffer_ids, leaf_buffer_ids, /*no_alloc =*/ false); -} - -bool ggml_gallocr_reserve(ggml_gallocr_t galloc, struct ggml_cgraph *graph) { - return ggml_gallocr_reserve_n(galloc, graph, NULL, NULL); -} - -static void ggml_gallocr_init_tensor(ggml_gallocr_t galloc, struct ggml_tensor * tensor, struct tensor_alloc * tensor_alloc) { - int buffer_id = tensor_alloc->buffer_id; - assert(tensor->data || tensor->view_src || ggml_backend_buft_get_alloc_size(galloc->bufts[buffer_id], tensor) <= tensor_alloc->size_max); - - if (tensor->view_src != NULL) { - if (tensor->buffer == NULL) { - assert(tensor_alloc->addr.offset == SIZE_MAX); - if (tensor->view_src->buffer == NULL) { - // this tensor was allocated without ggml-backend - return; - } - ggml_backend_view_init(tensor); - } - } else { - if (tensor->data == NULL) { - assert(tensor_alloc->addr.offset != SIZE_MAX); - assert(ggml_backend_buft_get_alloc_size(galloc->bufts[buffer_id], tensor) <= tensor_alloc->size_max); - ggml_vbuffer_tensor_alloc(galloc->buffers[buffer_id], tensor, tensor_alloc->addr); - } else { - if (tensor->buffer == NULL) { - // this tensor was allocated without ggml-backend - return; - } - } - } -} - -static bool ggml_gallocr_node_needs_realloc(ggml_gallocr_t galloc, struct ggml_tensor * node, struct tensor_alloc * talloc) { - size_t node_size = 0; - if (!node->data && !node->view_src) { - // If we previously had data but don't now then reallocate - if (talloc->buffer_id < 0) { - return false; - } - node_size = ggml_backend_buft_get_alloc_size(galloc->bufts[talloc->buffer_id], node); - } - return talloc->size_max >= node_size; -} - -static bool ggml_gallocr_needs_realloc(ggml_gallocr_t galloc, struct ggml_cgraph * graph) { - if (galloc->n_nodes != graph->n_nodes) { -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: graph has different number of nodes\n", __func__); -#endif - return true; - } - - if (galloc->n_leafs != graph->n_leafs) { -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: graph has different number of leafs\n", __func__); -#endif - return true; - } - - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - struct node_alloc * node_alloc = &galloc->node_allocs[i]; - - if (!ggml_gallocr_node_needs_realloc(galloc, node, &node_alloc->dst)) { -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: node %s is not valid\n", __func__, node->name); -#endif - return true; - } - - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * src = node->src[j]; - if (src == NULL) { - continue; - } - if (!ggml_gallocr_node_needs_realloc(galloc, src, &node_alloc->src[j])) { -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: src %d (%s) of node %s is not valid\n", __func__, j, src->name, node->name); -#endif - return true; - } - } - } - - return false; -} - -bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph * graph) { - if (ggml_gallocr_needs_realloc(galloc, graph)) { - if (galloc->n_buffers == 1) { -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: reallocating buffers automatically\n", __func__); -#endif - if (!ggml_gallocr_reserve(galloc, graph)) { - return false; - } - } else { -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: cannot reallocate multi buffer graph automatically, call reserve\n", __func__); -#endif - return false; - } - } - - // reset buffers - for (int i = 0; i < galloc->n_buffers; i++) { - if (galloc->buffers[i] != NULL) { - ggml_vbuffer_reset(galloc->buffers[i]); - } - } - - // allocate the graph tensors from the previous assignments - // leafs - for (int i = 0; i < graph->n_leafs; i++) { - struct ggml_tensor * leaf = graph->leafs[i]; - struct leaf_alloc * leaf_alloc = &galloc->leaf_allocs[i]; - ggml_gallocr_init_tensor(galloc, leaf, &leaf_alloc->leaf); - } - // nodes - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - struct node_alloc * node_alloc = &galloc->node_allocs[i]; - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * src = node->src[j]; - if (src == NULL) { - continue; - } - ggml_gallocr_init_tensor(galloc, src, &node_alloc->src[j]); - } - ggml_gallocr_init_tensor(galloc, node, &node_alloc->dst); - } - - return true; -} - -size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id) { - GGML_ASSERT(buffer_id >= 0 && buffer_id < galloc->n_buffers); - - if (galloc->buffers[buffer_id] == NULL) { - return 0; - } - - for (int i = 0; i < buffer_id; i++) { - if (galloc->buffers[i] == galloc->buffers[buffer_id]) { - // this buffer is the same as a previous one due to the same buffer type being used multiple times - // only return the buffer size the first time it appears to avoid double counting - return 0; - } - } - - return ggml_vbuffer_size(galloc->buffers[buffer_id]); -} - -// utils - -static void free_buffers(ggml_backend_buffer_t ** buffers, const size_t * n_buffers) { - for (size_t i = 0; i < *n_buffers; i++) { - ggml_backend_buffer_free((*buffers)[i]); - } - free(*buffers); -} - -static bool alloc_tensor_range(struct ggml_context * ctx, - struct ggml_tensor * first, struct ggml_tensor * last, - ggml_backend_buffer_type_t buft, size_t size, - ggml_backend_buffer_t ** buffers, size_t * n_buffers) { - - ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, size); - if (buffer == NULL) { - GGML_LOG_ERROR("%s: failed to allocate %s buffer of size %zu\n", __func__, ggml_backend_buft_name(buft), size); - free_buffers(buffers, n_buffers); - return false; - } - - *buffers = realloc(*buffers, sizeof(ggml_backend_buffer_t) * (*n_buffers + 1)); - (*buffers)[(*n_buffers)++] = buffer; - - struct ggml_tallocr tallocr = ggml_tallocr_new(buffer); - - for (struct ggml_tensor * t = first; t != last; t = ggml_get_next_tensor(ctx, t)) { - enum ggml_status status = GGML_STATUS_SUCCESS; - if (t->data == NULL) { - if (t->view_src == NULL) { - status = ggml_tallocr_alloc(&tallocr, t); - } else if (t->buffer == NULL) { - status = ggml_backend_view_init(t); - } - } else { - if (t->view_src != NULL && t->buffer == NULL) { - // view of a pre-allocated tensor - status = ggml_backend_view_init(t); - } - } - if (status != GGML_STATUS_SUCCESS) { - GGML_LOG_ERROR("%s: failed to initialize tensor %s\n", __func__, t->name); - free_buffers(buffers, n_buffers); - return false; - } - } - - return true; -} - -static ggml_backend_buffer_t ggml_backend_alloc_ctx_tensors_from_buft_impl( - struct ggml_context * ctx, ggml_backend_buffer_type_t buft, size_t * nbytes_total, bool no_alloc) { - GGML_ASSERT(ggml_get_no_alloc(ctx) == true); - - size_t alignment = ggml_backend_buft_get_alignment(buft); - size_t max_size = ggml_backend_buft_get_max_size(buft); - - ggml_backend_buffer_t * buffers = NULL; - size_t n_buffers = 0; - *nbytes_total = 0; - - size_t cur_buf_size = 0; - struct ggml_tensor * first = ggml_get_first_tensor(ctx); - for (struct ggml_tensor * t = first; t != NULL; t = ggml_get_next_tensor(ctx, t)) { - size_t this_size = 0; - if (t->data == NULL && t->view_src == NULL) { - this_size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, t), alignment); - } - - if (cur_buf_size > 0 && (cur_buf_size + this_size) > max_size) { - // allocate tensors in the current buffer - if (!no_alloc && !alloc_tensor_range(ctx, first, t, buft, cur_buf_size, &buffers, &n_buffers)) { - return NULL; - } - first = t; - *nbytes_total += cur_buf_size; - cur_buf_size = this_size; - } else { - cur_buf_size += this_size; - } - } - - // allocate remaining tensors - if (cur_buf_size > 0) { - *nbytes_total += cur_buf_size; - if (!no_alloc && !alloc_tensor_range(ctx, first, NULL, buft, cur_buf_size, &buffers, &n_buffers)) { - return NULL; - } - } - - if (no_alloc) { - return NULL; - } - - if (n_buffers == 0) { -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: all tensors in the context are already allocated\n", __func__); -#endif - GGML_ASSERT(!buffers); - return NULL; - } - - ggml_backend_buffer_t buffer; - if (n_buffers == 1) { - buffer = buffers[0]; - } else { - buffer = ggml_backend_multi_buffer_alloc_buffer(buffers, n_buffers); - } - if (buffers) { - free(buffers); // can be NULL if context is empty or no_alloc - } - return buffer; -} - -size_t ggml_backend_alloc_ctx_tensors_from_buft_size(struct ggml_context * ctx, ggml_backend_buffer_type_t buft) { - size_t nbytes_total = 0; - ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft_impl(ctx, buft, &nbytes_total, /*no_alloc=*/ true); - GGML_ASSERT(!buf); - return nbytes_total; -} - -ggml_backend_buffer_t ggml_backend_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft) { - size_t nbytes_total = 0; - if (ggml_backend_buft_is_meta(buft)) { - return ggml_backend_meta_alloc_ctx_tensors_from_buft(ctx, buft); - } - return ggml_backend_alloc_ctx_tensors_from_buft_impl(ctx, buft, &nbytes_total, /*no_alloc =*/ false); -} - -ggml_backend_buffer_t ggml_backend_alloc_ctx_tensors(struct ggml_context * ctx, ggml_backend_t backend) { - return ggml_backend_alloc_ctx_tensors_from_buft(ctx, ggml_backend_get_default_buffer_type(backend)); -} diff --git a/ggml/src/ggml-backend-dl-rust/Cargo.toml b/ggml/src/ggml-backend-dl-rust/Cargo.toml new file mode 100644 index 000000000000..99582d5c0507 --- /dev/null +++ b/ggml/src/ggml-backend-dl-rust/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "ggml-backend-dl-rust" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["staticlib", "rlib"] diff --git a/ggml/src/ggml-backend-dl-rust/src/lib.rs b/ggml/src/ggml-backend-dl-rust/src/lib.rs new file mode 100644 index 000000000000..4e66a71a611a --- /dev/null +++ b/ggml/src/ggml-backend-dl-rust/src/lib.rs @@ -0,0 +1,226 @@ +use std::ffi::{c_char, c_int, c_void, CStr, CString}; +use std::ptr; + +#[cfg(unix)] +mod platform { + use super::*; + + const RTLD_NOW: c_int = 2; + const RTLD_LOCAL: c_int = 0; + + extern "C" { + fn dlopen(filename: *const c_char, flags: c_int) -> *mut c_void; + fn dlsym(handle: *mut c_void, symbol: *const c_char) -> *mut c_void; + fn dlerror() -> *const c_char; + } + + pub unsafe fn load(path: *const c_char) -> *mut c_void { + dlopen(path, RTLD_NOW | RTLD_LOCAL) + } + + pub unsafe fn sym(handle: *mut c_void, name: *const c_char) -> *mut c_void { + dlsym(handle, name) + } + + pub unsafe fn error() -> *const c_char { + let err = dlerror(); + if err.is_null() { + c"".as_ptr() + } else { + err + } + } +} + +#[cfg(windows)] +mod platform { + use super::*; + + const SEM_FAILCRITICALERRORS: u32 = 0x0001; + + #[link(name = "kernel32")] + extern "system" { + fn SetErrorMode(mode: u32) -> u32; + fn LoadLibraryW(name: *const u16) -> *mut c_void; + fn GetProcAddress(handle: *mut c_void, name: *const c_char) -> *mut c_void; + } + + pub unsafe fn load(path: *const c_char) -> *mut c_void { + let Ok(path) = std::ffi::CStr::from_ptr(path).to_str() else { + return ptr::null_mut(); + }; + let mut wide: Vec = path.encode_utf16().collect(); + wide.push(0); + + let old_mode = SetErrorMode(SEM_FAILCRITICALERRORS); + SetErrorMode(old_mode | SEM_FAILCRITICALERRORS); + let handle = LoadLibraryW(wide.as_ptr()); + SetErrorMode(old_mode); + handle + } + + pub unsafe fn sym(handle: *mut c_void, name: *const c_char) -> *mut c_void { + let old_mode = SetErrorMode(SEM_FAILCRITICALERRORS); + SetErrorMode(old_mode | SEM_FAILCRITICALERRORS); + let symbol = GetProcAddress(handle, name); + SetErrorMode(old_mode); + symbol + } + + pub unsafe fn error() -> *const c_char { + c"".as_ptr() + } +} + +#[no_mangle] +pub unsafe extern "C" fn ggml_backend_dl_load_library_rust(path: *const c_char) -> *mut c_void { + if path.is_null() { + return ptr::null_mut(); + } + platform::load(path) +} + +#[no_mangle] +pub unsafe extern "C" fn ggml_backend_dl_get_sym_rust( + handle: *mut c_void, + name: *const c_char, +) -> *mut c_void { + if handle.is_null() || name.is_null() { + return ptr::null_mut(); + } + platform::sym(handle, name) +} + +#[no_mangle] +pub unsafe extern "C" fn ggml_backend_dl_error_rust() -> *const c_char { + platform::error() +} + +#[no_mangle] +pub unsafe extern "C" fn ggml_backend_striequals_rust(a: *const c_char, b: *const c_char) -> c_int { + if a.is_null() || b.is_null() { + return 0; + } + + let a = CStr::from_ptr(a).to_bytes(); + let b = CStr::from_ptr(b).to_bytes(); + if a.len() != b.len() { + return 0; + } + + a.iter() + .zip(b.iter()) + .all(|(a, b)| a.to_ascii_lowercase() == b.to_ascii_lowercase()) as c_int +} + +#[no_mangle] +pub extern "C" fn ggml_backend_current_exe_dir_rust() -> *mut c_char { + current_exe_dir() + .and_then(|path| CString::new(path).ok()) + .map(CString::into_raw) + .unwrap_or(ptr::null_mut()) +} + +#[no_mangle] +pub unsafe extern "C" fn ggml_backend_cstring_free_rust(s: *mut c_char) { + if !s.is_null() { + drop(CString::from_raw(s)); + } +} + +#[no_mangle] +pub extern "C" fn ggml_backend_filename_prefix_rust() -> *const c_char { + #[cfg(windows)] + { + c"ggml-".as_ptr() + } + #[cfg(not(windows))] + { + c"libggml-".as_ptr() + } +} + +#[no_mangle] +pub extern "C" fn ggml_backend_filename_extension_rust() -> *const c_char { + #[cfg(windows)] + { + c".dll".as_ptr() + } + #[cfg(not(windows))] + { + c".so".as_ptr() + } +} + +fn current_exe_dir() -> Option { + let exe = std::env::current_exe().ok()?; + let dir = exe.parent()?; + Some(dir.to_string_lossy().into_owned()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::ffi::CStr; + + #[test] + fn nulls_are_rejected() { + unsafe { + assert!(ggml_backend_dl_load_library_rust(ptr::null()).is_null()); + assert!(ggml_backend_dl_get_sym_rust(ptr::null_mut(), c"missing".as_ptr()).is_null()); + assert!(ggml_backend_dl_get_sym_rust(ptr::NonNull::::dangling().as_ptr(), ptr::null()).is_null()); + assert!(!ggml_backend_dl_error_rust().is_null()); + assert_eq!(ggml_backend_striequals_rust(ptr::null(), c"cpu".as_ptr()), 0); + assert_eq!(ggml_backend_striequals_rust(c"cpu".as_ptr(), ptr::null()), 0); + } + } + + #[test] + fn compares_ascii_names_case_insensitively() { + unsafe { + assert_eq!(ggml_backend_striequals_rust(c"CPU".as_ptr(), c"cpu".as_ptr()), 1); + assert_eq!(ggml_backend_striequals_rust(c"CuDa".as_ptr(), c"cuda".as_ptr()), 1); + assert_eq!(ggml_backend_striequals_rust(c"cpu".as_ptr(), c"cpux".as_ptr()), 0); + assert_eq!(ggml_backend_striequals_rust(c"metal".as_ptr(), c"vulkan".as_ptr()), 0); + } + } + + #[test] + fn exposes_current_executable_directory() { + unsafe { + let dir_ptr = ggml_backend_current_exe_dir_rust(); + assert!(!dir_ptr.is_null()); + let dir = CStr::from_ptr(dir_ptr).to_string_lossy().into_owned(); + assert!(!dir.is_empty()); + assert!(std::path::Path::new(&dir).is_dir()); + ggml_backend_cstring_free_rust(dir_ptr); + } + } + + #[test] + fn exposes_backend_filename_parts() { + unsafe { + #[cfg(windows)] + { + assert_eq!(CStr::from_ptr(ggml_backend_filename_prefix_rust()).to_bytes(), b"ggml-"); + assert_eq!(CStr::from_ptr(ggml_backend_filename_extension_rust()).to_bytes(), b".dll"); + } + #[cfg(not(windows))] + { + assert_eq!(CStr::from_ptr(ggml_backend_filename_prefix_rust()).to_bytes(), b"libggml-"); + assert_eq!(CStr::from_ptr(ggml_backend_filename_extension_rust()).to_bytes(), b".so"); + } + } + } + + #[cfg(unix)] + #[test] + fn can_resolve_process_symbol() { + unsafe { + let handle = ggml_backend_dl_load_library_rust(c"".as_ptr()); + assert!(!handle.is_null()); + let symbol = ggml_backend_dl_get_sym_rust(handle, c"printf".as_ptr()); + assert!(!symbol.is_null()); + } + } +} diff --git a/ggml/src/ggml-backend-dl.cpp b/ggml/src/ggml-backend-dl.cpp deleted file mode 100644 index a65cf0090552..000000000000 --- a/ggml/src/ggml-backend-dl.cpp +++ /dev/null @@ -1,48 +0,0 @@ -#include "ggml-backend-dl.h" - -#ifdef _WIN32 - -dl_handle * dl_load_library(const fs::path & path) { - // suppress error dialogs for missing DLLs - DWORD old_mode = SetErrorMode(SEM_FAILCRITICALERRORS); - SetErrorMode(old_mode | SEM_FAILCRITICALERRORS); - - HMODULE handle = LoadLibraryW(path.wstring().c_str()); - - SetErrorMode(old_mode); - - return handle; -} - -void * dl_get_sym(dl_handle * handle, const char * name) { - DWORD old_mode = SetErrorMode(SEM_FAILCRITICALERRORS); - SetErrorMode(old_mode | SEM_FAILCRITICALERRORS); - - void * p = (void *) GetProcAddress(handle, name); - - SetErrorMode(old_mode); - - return p; -} - -const char * dl_error() { - return ""; -} - -#else - -dl_handle * dl_load_library(const fs::path & path) { - dl_handle * handle = dlopen(path.string().c_str(), RTLD_NOW | RTLD_LOCAL); - return handle; -} - -void * dl_get_sym(dl_handle * handle, const char * name) { - return dlsym(handle, name); -} - -const char * dl_error() { - const char *rslt = dlerror(); - return rslt != nullptr ? rslt : ""; -} - -#endif diff --git a/ggml/src/ggml-backend-dl.h b/ggml/src/ggml-backend-dl.h deleted file mode 100644 index f74b7c948946..000000000000 --- a/ggml/src/ggml-backend-dl.h +++ /dev/null @@ -1,45 +0,0 @@ -#pragma once - -#ifdef _WIN32 -# define WIN32_LEAN_AND_MEAN -# ifndef NOMINMAX -# define NOMINMAX -# endif -# include -# include -#else -# include -# include -#endif -#include - -namespace fs = std::filesystem; - -#ifdef _WIN32 - -using dl_handle = std::remove_pointer_t; - -struct dl_handle_deleter { - void operator()(HMODULE handle) { - FreeLibrary(handle); - } -}; - -#else - -using dl_handle = void; - -struct dl_handle_deleter { - void operator()(void * handle) { - dlclose(handle); - } -}; - -#endif - -using dl_handle_ptr = std::unique_ptr; - -dl_handle * dl_load_library(const fs::path & path); -void * dl_get_sym(dl_handle * handle, const char * name); -const char * dl_error(); - diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl-defs.inc similarity index 99% rename from ggml/src/ggml-backend-impl.h rename to ggml/src/ggml-backend-impl-defs.inc index 9c56ec30c5f1..c738f8dba57b 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl-defs.inc @@ -2,7 +2,7 @@ // ggml-backend internal header -#include "ggml-backend.h" +#include "ggml-backend.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp deleted file mode 100644 index c0ffd9a048b1..000000000000 --- a/ggml/src/ggml-backend-meta.cpp +++ /dev/null @@ -1,2143 +0,0 @@ -#include "ggml.h" -#include "ggml-impl.h" -#include "ggml-backend.h" -#include "ggml-backend-impl.h" -#include "ggml-alloc.h" -#include "ggml-cpp.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -struct ggml_backend_meta_device; -struct ggml_backend_meta_buffer_type; -struct ggml_backend_meta_buffer; -struct ggml_backend_meta; - -const char * ggml_backend_meta_split_axis_name(enum ggml_backend_meta_split_axis split_axis) { - switch (split_axis) { - case GGML_BACKEND_SPLIT_AXIS_0: - return "0"; - case GGML_BACKEND_SPLIT_AXIS_1: - return "1"; - case GGML_BACKEND_SPLIT_AXIS_2: - return "2"; - case GGML_BACKEND_SPLIT_AXIS_3: - return "3"; - case GGML_BACKEND_SPLIT_AXIS_MIRRORED: - return "MIRRORED"; - case GGML_BACKEND_SPLIT_AXIS_PARTIAL: - return "PARTIAL"; - case GGML_BACKEND_SPLIT_AXIS_NONE: - return "NONE"; - case GGML_BACKEND_SPLIT_AXIS_UNKNOWN: - return "UNKNOWN"; - default: - GGML_ABORT("fatal error"); - } -} - -// -// meta backend device -// - -struct ggml_backend_meta_device_context { - std::vector simple_devs; - ggml_backend_meta_get_split_state_t get_split_state; - void * get_split_state_ud; - - std::string name; - std::string description; - - ggml_backend_meta_device_context( - std::vector simple_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud) : - simple_devs(std::move(simple_devs)), get_split_state(get_split_state), get_split_state_ud(get_split_state_ud) { - name = std::string("Meta("); - description = std::string("Meta("); - for (size_t i = 0; i < simple_devs.size(); i++) { - if (i > 0) { - name += ","; - description += ","; - } - name += ggml_backend_dev_name (simple_devs[i]); - description += ggml_backend_dev_description(simple_devs[i]); - } - name += ")"; - description += ")"; - } - - bool operator<(const ggml_backend_meta_device_context & other) const { - return std::tie(simple_devs, get_split_state, get_split_state_ud) - < std::tie(other.simple_devs, other.get_split_state, other.get_split_state_ud); - } -}; - -static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev); - -static const char * ggml_backend_meta_device_get_name(ggml_backend_dev_t dev) { - GGML_ASSERT(ggml_backend_dev_is_meta(dev)); - const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; - return meta_dev_ctx->name.c_str(); -} - -static const char * ggml_backend_meta_device_get_description(ggml_backend_dev_t dev) { - GGML_ASSERT(ggml_backend_dev_is_meta(dev)); - const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; - return meta_dev_ctx->description.c_str(); -} - -static void ggml_backend_meta_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { - GGML_ASSERT(ggml_backend_dev_is_meta(dev)); - const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; - *free = 0; - *total = 0; - for (ggml_backend_dev_t dev : meta_dev_ctx->simple_devs) { - size_t tmp_free, tmp_total; - ggml_backend_dev_memory(dev, &tmp_free, &tmp_total); - *free += tmp_free; - *total += tmp_total; - } -} - -static enum ggml_backend_dev_type ggml_backend_meta_device_get_type(ggml_backend_dev_t dev) { - return GGML_BACKEND_DEVICE_TYPE_META; - - GGML_UNUSED(dev); -} - -static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_backend_dev_props * props) { - GGML_ASSERT(ggml_backend_dev_is_meta(dev)); - const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; - - // TODO replace placeholders - props->name = ggml_backend_meta_device_get_name(dev); - props->description = ggml_backend_meta_device_get_description(dev); - props->type = ggml_backend_meta_device_get_type(dev); - props->device_id = 0; - - ggml_backend_meta_device_get_memory(dev, &props->memory_free, &props->memory_total); - - props->caps = { - /* .async = */ true, - /* .host_buffer = */ false, // Not implemented. - /* .buffer_from_host_ptr = */ false, // Not implemented. - /* .events = */ false, // Not implemented. - }; - for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) { - ggml_backend_dev_props tmp_props; - ggml_backend_dev_get_props(simple_dev, &tmp_props); - props->caps.async = props->caps.async && tmp_props.caps.async; - props->caps.host_buffer = props->caps.host_buffer && tmp_props.caps.host_buffer; - props->caps.buffer_from_host_ptr = props->caps.buffer_from_host_ptr && tmp_props.caps.buffer_from_host_ptr; - props->caps.events = props->caps.events && tmp_props.caps.events; - } -} - -static ggml_backend_t ggml_backend_meta_device_init_backend(ggml_backend_dev_t dev, const char * params); - -static ggml_backend_buffer_type_t ggml_backend_meta_device_get_buffer_type(ggml_backend_dev_t dev); - -static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type(ggml_backend_dev_t dev); - -static bool ggml_backend_meta_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { - GGML_ASSERT(ggml_backend_dev_is_meta(dev)); - const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; - return std::all_of(meta_dev_ctx->simple_devs.begin(), meta_dev_ctx->simple_devs.end(), - [op](ggml_backend_dev_t simple_dev) { return ggml_backend_dev_supports_op(simple_dev, op); }); -} - -static bool ggml_backend_meta_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { - GGML_ASSERT(ggml_backend_dev_is_meta(dev)); - ggml_backend_dev_t dev_buft = ggml_backend_buft_get_device(buft); - if (!ggml_backend_dev_is_meta(dev_buft)) { - return false; - } - const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; - const ggml_backend_meta_device_context * meta_buft_dev_ctx = (const ggml_backend_meta_device_context *) dev_buft->context; - if (meta_dev_ctx->simple_devs.size() != meta_buft_dev_ctx->simple_devs.size()) { - return false; - } - for (size_t i = 0; i < meta_dev_ctx->simple_devs.size(); i++) { - if (meta_dev_ctx->simple_devs[i] != meta_buft_dev_ctx->simple_devs[i]) { - return false; - } - } - return true; -} - -static const ggml_backend_device_i ggml_backend_meta_device_iface = { - /* .get_name = */ ggml_backend_meta_device_get_name, - /* .get_description = */ ggml_backend_meta_device_get_description, - /* .get_memory = */ ggml_backend_meta_device_get_memory, - /* .get_type = */ ggml_backend_meta_device_get_type, - /* .get_props = */ ggml_backend_meta_device_get_props, - /* .init_backend = */ ggml_backend_meta_device_init_backend, - /* .get_buffer_type = */ ggml_backend_meta_device_get_buffer_type, - /* .get_host_buffer_type = */ ggml_backend_meta_device_get_host_buffer_type, - /* .buffer_from_host_ptr = */ nullptr, - /* .supports_op = */ ggml_backend_meta_device_supports_op, - /* .supports_buft = */ ggml_backend_meta_device_supports_buft, - /* .offload_op = */ nullptr, - /* .event_new = */ nullptr, - /* .event_free = */ nullptr, - /* .event_synchronize = */ nullptr, -}; - -static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev) { - return dev != nullptr && dev->iface.get_name == ggml_backend_meta_device_iface.get_name; -} - -static size_t ggml_backend_meta_dev_n_devs(ggml_backend_dev_t meta_dev) { - GGML_ASSERT(ggml_backend_dev_is_meta(meta_dev)); - const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) meta_dev->context; - return meta_dev_ctx->simple_devs.size(); -} - -static ggml_backend_dev_t ggml_backend_meta_dev_simple_dev(ggml_backend_dev_t meta_dev, size_t index) { - GGML_ASSERT(ggml_backend_dev_is_meta(meta_dev)); - const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) meta_dev->context; - GGML_ASSERT(index < meta_dev_ctx->simple_devs.size()); - return meta_dev_ctx->simple_devs[index]; -} - -ggml_backend_dev_t ggml_backend_meta_device( - ggml_backend_dev_t * devs, size_t n_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud) { - GGML_ASSERT(n_devs <= GGML_BACKEND_META_MAX_DEVICES); - // TODO: this is not thread-safe - needs to be fixed - static std::vector> ctxs; - static std::map meta_devs; - - std::vector simple_devs; - simple_devs.reserve(n_devs); - for (size_t i = 0; i < n_devs; i++) { - simple_devs.push_back(devs[i]); - } - ggml_backend_meta_device_context ctx(simple_devs, get_split_state, get_split_state_ud); - - { - auto it = meta_devs.find(ctx); - if (it != meta_devs.end()) { - return &it->second; - } - } - ctxs.push_back(std::make_unique(ctx)); - - struct ggml_backend_device meta_dev = { - /*iface =*/ ggml_backend_meta_device_iface, - /*reg =*/ nullptr, - /*ctx =*/ ctxs.back().get(), - }; - - auto result = meta_devs.emplace(*ctxs.back(), meta_dev); - return &result.first->second; -} - -// -// meta backend buffer type -// - -struct ggml_backend_meta_buffer_type_context { - std::vector simple_bufts; - - std::string name; - - ggml_backend_meta_buffer_type_context(std::vector simple_bufts) : simple_bufts(std::move(simple_bufts)) { - name = "Meta("; - for (size_t i = 0; i < simple_bufts.size(); i++) { - if (i > 0) { - name += ","; - } - name += ggml_backend_buft_name(simple_bufts[i]); - } - name += ")"; - } - - bool operator<(const ggml_backend_meta_buffer_type_context & other) const { - return simple_bufts < other.simple_bufts; - } -}; - -static size_t ggml_backend_meta_buft_n_bufts(ggml_backend_buffer_type_t meta_buft) { - GGML_ASSERT(ggml_backend_buft_is_meta(meta_buft)); - const ggml_backend_meta_buffer_type_context * meta_buft_ctx = (const ggml_backend_meta_buffer_type_context *) meta_buft->context; - return meta_buft_ctx->simple_bufts.size(); -} - -static const char * ggml_backend_meta_buffer_type_get_name(ggml_backend_buffer_type_t buft) { - GGML_ASSERT(ggml_backend_buft_is_meta(buft)); - const ggml_backend_meta_buffer_type_context * meta_buft_ctx = (const ggml_backend_meta_buffer_type_context *) buft->context; - return meta_buft_ctx->name.c_str(); -} - -static ggml_backend_buffer_type_t ggml_backend_meta_buft_simple_buft(ggml_backend_buffer_type_t meta_buft, size_t index) { - GGML_ASSERT(ggml_backend_buft_is_meta(meta_buft)); - const ggml_backend_meta_buffer_type_context * meta_buft_ctx = (const ggml_backend_meta_buffer_type_context *) meta_buft->context; - GGML_ASSERT(index < meta_buft_ctx->simple_bufts.size()); - return meta_buft_ctx->simple_bufts[index]; -} - -static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size); - -static size_t ggml_backend_meta_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { - const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); - size_t max_alignment = 1; - for (size_t i = 0; i < n_simple_bufts; i++) { - const size_t alignment = ggml_backend_buft_get_alignment(ggml_backend_meta_buft_simple_buft(buft, i)); - max_alignment = std::max(max_alignment, alignment); - GGML_ASSERT(max_alignment % alignment == 0); - } - return max_alignment; -} - -static size_t ggml_backend_meta_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { - const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); - size_t max_size = SIZE_MAX; - for (size_t i = 0; i < n_simple_bufts; i++) { - max_size = std::min(max_size, ggml_backend_buft_get_max_size(ggml_backend_meta_buft_simple_buft(buft, i))); - } - return max_size; -} - -static size_t ggml_backend_meta_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { - const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); - size_t max_alloc_size = 0; - for (size_t i = 0; i < n_simple_bufts; i++) { - const size_t alloc_size = ggml_backend_buft_get_alloc_size(ggml_backend_meta_buft_simple_buft(buft, i), tensor); - max_alloc_size = std::max(max_alloc_size, alloc_size); - } - return max_alloc_size; -} - -static bool ggml_backend_meta_buffer_type_is_host(ggml_backend_buffer_type_t buft) { - const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); - for (size_t i = 0; i < n_simple_bufts; i++) { - if (!ggml_backend_buft_is_host(ggml_backend_meta_buft_simple_buft(buft, i))) { - return false; - } - } - return true; -} - -static const struct ggml_backend_buffer_type_i ggml_backend_meta_buffer_type_iface = { - /* .get_name = */ ggml_backend_meta_buffer_type_get_name, - /* .alloc_buffer = */ ggml_backend_meta_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_meta_buffer_type_get_alignment, - /* .get_max_size = */ ggml_backend_meta_buffer_type_get_max_size, - /* .get_alloc_size = */ ggml_backend_meta_buffer_type_get_alloc_size, - /* .is_host = */ ggml_backend_meta_buffer_type_is_host, -}; - -bool ggml_backend_buft_is_meta(ggml_backend_buffer_type_t buft) { - return buft != nullptr && buft->iface.get_name == ggml_backend_meta_buffer_type_iface.get_name; -} - -static ggml_backend_buffer_type_t ggml_backend_meta_device_get_buffer_type(ggml_backend_dev_t dev) { - static std::map meta_bufts; - GGML_ASSERT(ggml_backend_dev_is_meta(dev)); - { - auto it = meta_bufts.find(dev); - if (it != meta_bufts.end()) { - return &it->second; - } - } - - const size_t n_devs = ggml_backend_meta_dev_n_devs(dev); - std::vector simple_bufts; - simple_bufts.reserve(n_devs); - for (size_t i = 0; i < n_devs; i++) { - simple_bufts.push_back(ggml_backend_dev_buffer_type(ggml_backend_meta_dev_simple_dev(dev, i))); - } - ggml_backend_meta_buffer_type_context * buft_ctx = new ggml_backend_meta_buffer_type_context(simple_bufts); - - struct ggml_backend_buffer_type meta_buft = { - /*iface =*/ ggml_backend_meta_buffer_type_iface, - /*device =*/ dev, - /*ctx =*/ buft_ctx, - }; - auto result = meta_bufts.emplace(dev, meta_buft); - return &result.first->second; -} - -static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type(ggml_backend_dev_t dev) { - GGML_ASSERT(ggml_backend_dev_is_meta(dev)); - const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; - - ggml_backend_buffer_type_t host_buft = nullptr; - for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) { - ggml_backend_buffer_type_t simple_host_buft = ggml_backend_dev_host_buffer_type(simple_dev); - if (simple_host_buft == nullptr) { - return nullptr; - } - if (host_buft == nullptr) { - host_buft = simple_host_buft; - } else if (host_buft != simple_host_buft) { - // if different simple devices have different host buffer types, - // we cannot provide a single host buffer type for the meta device - return nullptr; - } - } - return host_buft; -} - -// -// meta backend buffer -// - -struct ggml_backend_meta_buffer_context { - static constexpr size_t nbtc = GGML_TENSOR_SIZE - sizeof(ggml_tensor::padding); - - std::map, std::pair> split_state_cache; - std::map< const ggml_tensor *, std::vector> simple_tensors; - - struct buffer_config { - ggml_context * ctx; - ggml_backend_buffer_t buf; - - buffer_config(ggml_context * ctx, ggml_backend_buffer_t buf) : ctx(ctx), buf(buf) {} - }; - std::vector buf_configs; - - int debug; - - ggml_backend_meta_buffer_context() { - const char * GGML_META_DEBUG = getenv("GGML_META_DEBUG"); - debug = GGML_META_DEBUG ? atoi(GGML_META_DEBUG) : 0; - } -}; - -static void ggml_backend_meta_buffer_free_buffer(ggml_backend_buffer_t buffer) { - GGML_ASSERT(ggml_backend_buffer_is_meta(buffer)); - ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context; - for (auto & [ctx, buf] : buf_ctx->buf_configs) { - ggml_backend_buffer_free(buf); - ggml_free(ctx); - } - delete buf_ctx; -} - -static size_t ggml_backend_meta_buffer_n_bufs(ggml_backend_buffer_t meta_buf) { - GGML_ASSERT(ggml_backend_buffer_is_meta(meta_buf)); - ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) meta_buf->context; - return buf_ctx->buf_configs.size(); -} - -static ggml_backend_buffer_t ggml_backend_meta_buffer_simple_buffer(ggml_backend_buffer_t meta_buf, size_t index) { - GGML_ASSERT(ggml_backend_buffer_is_meta(meta_buf)); - ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) meta_buf->context; - GGML_ASSERT(index < buf_ctx->buf_configs.size()); - return buf_ctx->buf_configs[index].buf; -} - -static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct ggml_tensor * tensor, size_t index) { - GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer)); - ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; - GGML_ASSERT(index < buf_ctx->buf_configs.size()); - - auto it = buf_ctx->simple_tensors.find(tensor); - if (it == buf_ctx->simple_tensors.end()) { - return nullptr; - } - return it->second[index]; -} - -static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync) { - const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); - ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; - - auto split_states_equal = [&](const ggml_backend_meta_split_state & a, const ggml_backend_meta_split_state & b) -> bool { - if (a.axis != b.axis) { - return false; - } - for (size_t j = 0; j < n_bufs; j++) { - int64_t sum_a = 0; - for (size_t s = 0; s < a.n_segments; s++) { - sum_a += a.ne[s*n_bufs + j]; - } - int64_t sum_b = 0; - for (size_t s = 0; s < b.n_segments; s++) { - sum_b += b.ne[s*n_bufs + j]; - } - if (sum_a != sum_b) { - return false; - } - } - return true; - }; - - auto handle_generic = [&](const std::vector & src_ss, bool scalar_only) -> ggml_backend_meta_split_state { - ggml_backend_meta_split_state ret = {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, 1}; - for (size_t i = 0; i < GGML_MAX_SRC; i++) { - if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { - continue; - } - if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) { - ret = src_ss[i]; - } else if (!split_states_equal(src_ss[i], ret)) { - ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; - break; - } - } - if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) { - ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; - } - if (scalar_only && ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { - ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; - } - GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); - return ret; - }; - - // Some ops process data on a per-row bases: - auto handle_per_row = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0); - return src_ss[0]; - }; - - // Some ops broadcast the src1 data across src0: - auto handle_bin_bcast = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS && - tensor->src[1]->ne[src_ss[0].axis] == 1 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - return src_ss[0]; - } - if (src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && (src_ss[0].axis == src_ss[1].axis || - (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL)))) { - return src_ss[0]; // GGML_OP_ADD_ID - } - GGML_ASSERT(tensor->src[2] == nullptr || src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); - return handle_generic(src_ss, /*scalar_only =*/ false); - }; - - auto handle_concat = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - const ggml_backend_meta_split_axis concat_axis = ggml_backend_meta_split_axis(ggml_get_op_params_i32(tensor, 0)); - if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis >= 0 && src_ss[1].axis < GGML_MAX_DIMS) { - GGML_ASSERT(concat_axis != src_ss[1].axis); - return src_ss[1]; - } - if (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { - GGML_ASSERT(concat_axis != src_ss[0].axis); - return src_ss[0]; - } - if (src_ss[0].axis == src_ss[1].axis && src_ss[0].axis != concat_axis) { - return src_ss[0]; - } - return handle_generic(src_ss, /*scalar_only =*/ true); - }; - - auto handle_mul_mat = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, 1}; - } - if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - ggml_backend_meta_split_state ret = src_ss[0]; - ret.axis = GGML_BACKEND_SPLIT_AXIS_0; - ret.n_segments = 1; - return ret; - } - if (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1 && src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - ggml_backend_meta_split_state ret = src_ss[1]; - ret.n_segments = 1; - return ret; - } - if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_0) { - GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1])); - return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_PARTIAL, {0}, 1}; - } - GGML_ABORT("fatal error"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; - }; - - auto handle_cpy = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { - int64_t ne_split_src = tensor->src[0]->ne[0]; - for (int dim = 1; dim <= src_ss[0].axis; dim++) { - ne_split_src *= tensor->src[0]->ne[dim]; - } - int64_t ne_split_dst = 1; - for (int dim = 0; dim < GGML_MAX_DIMS; dim++) { - ne_split_dst *= tensor->ne[dim]; - if (ne_split_dst == ne_split_src) { - return {ggml_backend_meta_split_axis(dim), {0}, 1}; - } - } - } - return handle_generic(src_ss, /*scalar_only =*/ false); - }; - - auto handle_reshape = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - switch (src_ss[0].axis) { - case GGML_BACKEND_SPLIT_AXIS_0: - case GGML_BACKEND_SPLIT_AXIS_1: - case GGML_BACKEND_SPLIT_AXIS_2: - case GGML_BACKEND_SPLIT_AXIS_3: { - GGML_ASSERT(!ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0])); - if (src_ss[0].axis == ggml_n_dims(tensor->src[0]) - 1) { - return {ggml_backend_meta_split_axis(ggml_n_dims(tensor) - 1), {0}, 1}; - } - std::vector base_ne_in; - base_ne_in.reserve(GGML_MAX_DIMS - src_ss[0].axis); - { - base_ne_in.push_back(1); - int dim = 0; - for (; dim <= src_ss[0].axis; dim++) { - base_ne_in[0] *= tensor->src[0]->ne[dim]; - } - for (; dim <= GGML_MAX_DIMS; dim++) { - base_ne_in.push_back(base_ne_in.back() * tensor->src[0]->ne[dim]); - } - } - int64_t base_ne_out = 1; - for (int dim = 0; dim < GGML_MAX_DIMS; dim++) { - const int64_t base_ne_out_next = base_ne_out *= tensor->ne[dim]; - for (const int64_t & bni : base_ne_in) { - if (bni == base_ne_out_next) { - return {ggml_backend_meta_split_axis(dim), {0}, 1}; - } - } - if (base_ne_out_next > base_ne_in[0]) { - GGML_ASSERT(dim + 1 < GGML_MAX_DIMS); - return {ggml_backend_meta_split_axis(dim + 1), {0}, 1}; - } - base_ne_out = base_ne_out_next; - } - GGML_ABORT("shape mismatch for %s", ggml_op_name(tensor->op)); - } - case GGML_BACKEND_SPLIT_AXIS_MIRRORED: - case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { - return src_ss[0]; - } - default: { - GGML_ABORT("fatal error"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; - } - } - }; - - auto handle_view = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - if (ggml_is_contiguous(tensor) && ggml_is_contiguous(tensor->src[0])) { - return handle_reshape(src_ss); - } - const int axis = src_ss[0].axis; - { - bool all_strides_the_same = true; - for (int dim = 0; dim < GGML_MAX_DIMS; dim++) { - if (tensor->ne[dim] == 1 && tensor->src[0]->ne[dim] == 1) { - continue; - } - if (tensor->nb[dim] != tensor->src[0]->nb[dim]) { - all_strides_the_same = false; - break; - } - } - if (all_strides_the_same) { - return src_ss[0]; - } - } - if (!ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0]) && axis >= 0 && axis < GGML_MAX_DIMS-1) { - for (int dim = 0; dim < GGML_MAX_DIMS-1; dim++) { - if (tensor->nb[dim+1] == tensor->src[0]->nb[axis+1]) { - return {ggml_backend_meta_split_axis(dim), {0}, 1}; - } - } - GGML_ABORT("fatal error"); - } - if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED || src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { - return src_ss[0]; - } - GGML_ABORT("view of permuted tensor not implemented"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; - }; - - auto handle_permute = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - switch (src_ss[0].axis) { - case GGML_BACKEND_SPLIT_AXIS_0: - case GGML_BACKEND_SPLIT_AXIS_1: - case GGML_BACKEND_SPLIT_AXIS_2: - case GGML_BACKEND_SPLIT_AXIS_3: { - return {ggml_backend_meta_split_axis(tensor->op_params[src_ss[0].axis]), {0}, 1}; - } - case GGML_BACKEND_SPLIT_AXIS_MIRRORED: - case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { - return src_ss[0]; - } - default: { - GGML_ABORT("fatal error"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; - } - } - }; - - auto handle_transpose = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - switch (src_ss[0].axis) { - case GGML_BACKEND_SPLIT_AXIS_0: - case GGML_BACKEND_SPLIT_AXIS_1: { - return {ggml_backend_meta_split_axis(int(src_ss[0].axis) ^ 1), {0}, 1}; - } - case GGML_BACKEND_SPLIT_AXIS_2: - case GGML_BACKEND_SPLIT_AXIS_3: - case GGML_BACKEND_SPLIT_AXIS_MIRRORED: - case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { - return src_ss[0]; - } - default: { - GGML_ABORT("fatal error"); - //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; - } - } - }; - - auto handle_get_rows = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - return src_ss[0]; - } - return handle_generic(src_ss, /*scalar_only =*/ true); - }; - - auto handle_set_rows = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); - GGML_ASSERT(split_states_equal(src_ss[0], src_ss[2])); - return src_ss[0]; - }; - - auto handle_rope = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); - return src_ss[0]; - }; - - auto handle_pad = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { - GGML_ASSERT(tensor->op_params[2*src_ss[0].axis + 0] == 0); - GGML_ASSERT(tensor->op_params[2*src_ss[0].axis + 1] == 0); - } - return src_ss[0]; - }; - - auto handle_flash_attn_ext = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - GGML_ASSERT( src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2); - GGML_ASSERT( src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2); - GGML_ASSERT( src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2); - GGML_ASSERT(tensor->src[4] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); - GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0); - return {GGML_BACKEND_SPLIT_AXIS_1, {0}, 1}; - }; - - auto handle_ssm_conv = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - if (src_ss[0].axis == src_ss[1].axis) { - if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { - return {GGML_BACKEND_SPLIT_AXIS_1, {0}, 1}; - } - if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1) { - return {GGML_BACKEND_SPLIT_AXIS_0, {0}, 1}; - } - } - return handle_generic(src_ss, /*scalar_only =*/ false); - }; - - auto handle_gated_delta_net = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && - src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && - src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - return src_ss[0]; - } - GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_1); - GGML_ASSERT(src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_2); - return {GGML_BACKEND_SPLIT_AXIS_0, {0}, 1}; - }; - - auto calculate_split_state = [&]() -> ggml_backend_meta_split_state { - if (ggml_nelements(tensor) == 0) { - return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; - } - if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE && tensor->view_src == nullptr) { - ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); - const ggml_backend_meta_device_context * dev_ctx = (const ggml_backend_meta_device_context *) dev->context; - ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor, dev_ctx->get_split_state_ud); - if (ret.axis >= 0 && ret.axis <= GGML_MAX_DIMS) { - const int64_t granularity = ret.axis == GGML_BACKEND_SPLIT_AXIS_0 ? ggml_blck_size(tensor->type) : 1; - int64_t ne_sum = 0; - for (size_t sj = 0; sj < ret.n_segments*n_bufs; sj++) { - GGML_ASSERT(ret.ne[sj] % granularity == 0); - ne_sum += ret.ne[sj]; - } - GGML_ASSERT(ne_sum == tensor->ne[ret.axis]); - } - return ret; - } - - std::vector src_ss(GGML_MAX_SRC, {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, 1}); - for (size_t i = 0; i < GGML_MAX_SRC; i++) { - if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { - src_ss[i] = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; - continue; - } - src_ss[i] = ggml_backend_meta_get_split_state(tensor->src[i], /*assume_sync =*/ true); - GGML_ASSERT(src_ss[i].axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); - } - - ggml_backend_meta_split_state split_state; - switch (tensor->op) { - case GGML_OP_NONE: { - split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, 1}; - } break; - case GGML_OP_DUP: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_ADD: - case GGML_OP_ADD_ID: { - split_state = handle_bin_bcast(src_ss); - } break; - case GGML_OP_ADD1: - case GGML_OP_ACC: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_SUB: - case GGML_OP_MUL: - case GGML_OP_DIV: { - split_state = handle_bin_bcast(src_ss); - } break; - case GGML_OP_SQR: - case GGML_OP_SQRT: - case GGML_OP_LOG: - case GGML_OP_SIN: - case GGML_OP_COS: { - split_state = handle_generic(src_ss, /*scalar_only =*/ false); - } break; - case GGML_OP_SUM: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_SUM_ROWS: - case GGML_OP_CUMSUM: - case GGML_OP_MEAN: - case GGML_OP_ARGMAX: - case GGML_OP_COUNT_EQUAL: { - split_state = handle_per_row(src_ss); - } break; - case GGML_OP_REPEAT: - case GGML_OP_REPEAT_BACK: { - split_state = handle_generic(src_ss, /*scalar_only =*/ false); - } break; - case GGML_OP_CONCAT: { - split_state = handle_concat(src_ss); - } break; - case GGML_OP_SILU_BACK: { - split_state = handle_generic(src_ss, /*scalar_only =*/ false); - } break; - case GGML_OP_NORM: - case GGML_OP_RMS_NORM: - case GGML_OP_RMS_NORM_BACK: - case GGML_OP_GROUP_NORM: - case GGML_OP_L2_NORM: { - split_state = handle_per_row(src_ss); - } break; - case GGML_OP_MUL_MAT: - case GGML_OP_MUL_MAT_ID: { - split_state = handle_mul_mat(src_ss); - } break; - case GGML_OP_OUT_PROD: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_SCALE: { - split_state = handle_generic(src_ss, /*scalar_only =*/ false); - } break; - case GGML_OP_SET: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_CPY: { - split_state = handle_cpy(src_ss); - } break; - case GGML_OP_CONT: - case GGML_OP_RESHAPE: { - split_state = handle_reshape(src_ss); - } break; - case GGML_OP_VIEW: { - split_state = handle_view(src_ss); - } break; - case GGML_OP_PERMUTE: { - split_state = handle_permute(src_ss); - } break; - case GGML_OP_TRANSPOSE: { - split_state = handle_transpose(src_ss); - } break; - case GGML_OP_GET_ROWS: { - split_state = handle_get_rows(src_ss); - } break; - case GGML_OP_GET_ROWS_BACK: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_SET_ROWS: { - split_state = handle_set_rows(src_ss); - } break; - case GGML_OP_DIAG: - case GGML_OP_DIAG_MASK_INF: - case GGML_OP_DIAG_MASK_ZERO: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_SOFT_MAX: - case GGML_OP_SOFT_MAX_BACK: { - split_state = handle_generic(src_ss, /*scalar_only =*/ false); - } break; - case GGML_OP_ROPE: { - split_state = handle_rope(src_ss); - } break; - case GGML_OP_ROPE_BACK: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_CLAMP: { - split_state = handle_generic(src_ss, /*scalar_only =*/ false); - } break; - case GGML_OP_CONV_TRANSPOSE_1D: - case GGML_OP_IM2COL: - case GGML_OP_IM2COL_BACK: - case GGML_OP_IM2COL_3D: - case GGML_OP_CONV_2D: - case GGML_OP_CONV_3D: - case GGML_OP_CONV_2D_DW: - case GGML_OP_CONV_TRANSPOSE_2D: - case GGML_OP_POOL_1D: - case GGML_OP_POOL_2D: - case GGML_OP_POOL_2D_BACK: - case GGML_OP_UPSCALE: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_PAD: { - split_state = handle_pad(src_ss); - } break; - case GGML_OP_PAD_REFLECT_1D: - case GGML_OP_ROLL: - case GGML_OP_ARANGE: - case GGML_OP_TIMESTEP_EMBEDDING: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_ARGSORT: - case GGML_OP_TOP_K: { - split_state = handle_per_row(src_ss); - } break; - case GGML_OP_LEAKY_RELU: { - split_state = handle_generic(src_ss, /*scalar_only =*/ false); - } break; - case GGML_OP_TRI: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_FILL: { - split_state = handle_generic(src_ss, /*scalar_only =*/ false); - } break; - case GGML_OP_FLASH_ATTN_EXT: { - split_state = handle_flash_attn_ext(src_ss); - } break; - case GGML_OP_FLASH_ATTN_BACK: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_SSM_CONV: { - split_state = handle_ssm_conv(src_ss); - } break; - case GGML_OP_SSM_SCAN: - case GGML_OP_WIN_PART: - case GGML_OP_WIN_UNPART: - case GGML_OP_GET_REL_POS: - case GGML_OP_ADD_REL_POS: - case GGML_OP_RWKV_WKV6: - case GGML_OP_GATED_LINEAR_ATTN: - case GGML_OP_RWKV_WKV7: - case GGML_OP_SOLVE_TRI: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_GATED_DELTA_NET: { - split_state = handle_gated_delta_net(src_ss); - } break; - case GGML_OP_UNARY: { - split_state = handle_generic(src_ss, /*scalar_only =*/ false); - } break; - case GGML_OP_MAP_CUSTOM1: - case GGML_OP_MAP_CUSTOM2: - case GGML_OP_MAP_CUSTOM3: - case GGML_OP_CUSTOM: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); - } break; - case GGML_OP_CROSS_ENTROPY_LOSS: - case GGML_OP_CROSS_ENTROPY_LOSS_BACK: { - split_state = handle_per_row(src_ss); - } break; - case GGML_OP_OPT_STEP_ADAMW: - case GGML_OP_OPT_STEP_SGD: - case GGML_OP_GLU: { - split_state = handle_generic(src_ss, /*scalar_only =*/ false); - } break; - default: { - GGML_ABORT("ggml op not implemented: %s", ggml_op_name(tensor->op)); - split_state = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; - } break; - } - if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS) { - bool first_src_split_by_axis = true; - const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); - - for (size_t i = 0; i < GGML_MAX_SRC; i++) { - if (tensor->src[i] == nullptr || src_ss[i].axis < 0 || src_ss[i].axis >= GGML_MAX_DIMS) { - continue; - } - if (first_src_split_by_axis) { - for (size_t j = 0; j < n_bufs; j++) { - // Take over ratio from src: - for (size_t s = 0; s < src_ss[i].n_segments; s++) { - split_state.ne[s*n_bufs + j] = 0; - } - for (size_t s = 0; s < src_ss[i].n_segments; s++) { - split_state.ne[j] += src_ss[i].ne[s*n_bufs + j]; - } - split_state.ne[j] *= tensor->ne[split_state.axis]; - if (split_state.ne[j] != 0 || tensor->src[i]->ne[src_ss[i].axis] != 0) { - GGML_ASSERT(split_state.ne[j] % tensor->src[i]->ne[src_ss[i].axis] == 0); - split_state.ne[j] /= tensor->src[i]->ne[src_ss[i].axis]; - } - } - } else { - for (size_t j = 0; j < n_bufs; j++) { - int64_t sum = 0; - for (size_t s = 0; s < src_ss[i].n_segments; s++) { - sum += src_ss[i].ne[s*n_bufs + j]; - } - // Assert that ratio is consistent: - GGML_ASSERT(split_state.ne[j] * tensor->src[i]->ne[src_ss[i].axis] - == sum * tensor->ne[split_state.axis]); - } - } - first_src_split_by_axis = false; - } - GGML_ASSERT(!first_src_split_by_axis); - } - return split_state; - }; - - const std::pair key = std::make_pair(tensor, assume_sync); - auto it = buf_ctx->split_state_cache.find(key); - if (it != buf_ctx->split_state_cache.end() && memcmp(it->second.second, (const char *) tensor, sizeof(it->second.second)) != 0) { - buf_ctx->split_state_cache.clear(); - it = buf_ctx->split_state_cache.end(); - } - - if (it == buf_ctx->split_state_cache.end()) { - buf_ctx->split_state_cache[key].first = calculate_split_state(); - memcpy(buf_ctx->split_state_cache[key].second, tensor, sizeof(buf_ctx->split_state_cache[key].second)); - if (buf_ctx->debug > 0) { - std::string srcs_info; - for (size_t i = 0; i < GGML_MAX_SRC; i++) { - if (tensor->src[i] == nullptr) { - continue; - } - if (!srcs_info.empty()) { - srcs_info += ", "; - } - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor->src[0], true); - const char * axis_name = ggml_backend_meta_split_axis_name(split_state.axis); - std::string ne_info; - for (size_t j = 0; j < n_bufs; j++) { - if (!ne_info.empty()) { - ne_info += ", "; - } - ne_info += std::to_string(split_state.ne[j]); - } - srcs_info += std::string(tensor->src[i]->name) + "[" + ggml_op_name(tensor->src[i]->op) + ", " + axis_name + ", {" + ne_info + "}]"; - } - std::string ne_info; - for (size_t j = 0; j < n_bufs; j++) { - if (!ne_info.empty()) { - ne_info += ", "; - } - ne_info += std::to_string(buf_ctx->split_state_cache[key].first.ne[j]); - } - GGML_LOG_DEBUG("SPLIT_STATE: {%s} -> %s[%s, %s, {%s}]\n", srcs_info.c_str(), tensor->name, ggml_op_name(tensor->op), - ggml_backend_meta_split_axis_name(buf_ctx->split_state_cache[key].first.axis), ne_info.c_str()); - } - } - - ggml_backend_meta_split_state ret = buf_ctx->split_state_cache[key].first; - GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_NONE); -#ifndef NDEBUG - if (ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { - int64_t ne_ret = 0; - for (size_t sj = 0; sj < ret.n_segments*n_bufs; sj++) { - ne_ret += ret.ne[sj]; - } - assert(ne_ret == tensor->ne[int(ret.axis)]); - } -#endif // NDEBUG - return ret; -} - -static void * ggml_backend_meta_buffer_get_base(ggml_backend_buffer_t buffer) { - GGML_UNUSED(buffer); - return (void *) 0x1000000000000000; // FIXME -} - -static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { - GGML_ASSERT(ggml_backend_buffer_is_meta(buffer)); - ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context; - const size_t n_simple_bufs = ggml_backend_meta_buffer_n_bufs(buffer); - - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ true); - GGML_ASSERT(ggml_nelements(tensor) == 0 || split_state.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); - GGML_ASSERT(split_state.n_segments <= 16); - - int split_dim = split_state.axis; - int64_t ne[GGML_MAX_DIMS]; - size_t nb[GGML_MAX_DIMS]; - for (size_t k = 0; k < GGML_MAX_DIMS; k++) { - ne[k] = tensor->ne[k]; - nb[k] = tensor->nb[k]; - } - - std::vector simple_tensors; - simple_tensors.reserve(n_simple_bufs); - for (size_t j = 0; j < n_simple_bufs; j++) { - ggml_context * simple_ctx = buf_ctx->buf_configs[j].ctx; - ggml_backend_buffer_t simple_buf = buf_ctx->buf_configs[j].buf; - - if (split_dim >= 0 && split_dim < GGML_MAX_DIMS) { - // TODO: the following assert fails for llama-parallel even though the results are correct: - // GGML_ASSERT(ggml_is_contiguously_allocated(tensor)); - ne[split_dim] = 0; - for (size_t s = 0; s < split_state.n_segments; s++) { - ne[split_dim] += split_state.ne[s*n_simple_bufs + j]; - } - for (int i = 0; i < GGML_MAX_DIMS; i++) { - if (tensor->nb[i] > tensor->nb[split_dim]) { - nb[i] = tensor->nb[i] * ne[split_dim]/tensor->ne[split_dim]; - } - } - } - - ggml_tensor * t_ij = ggml_new_tensor(simple_ctx, tensor->type, GGML_MAX_DIMS, ne); - t_ij->op = tensor->op; - for (int i = 0; i < GGML_MAX_DIMS; i++) { - t_ij->nb[i] = nb[i]; - } - t_ij->flags = tensor->flags; - memcpy(t_ij->op_params, tensor->op_params, sizeof(tensor->op_params)); - ggml_set_name(t_ij, tensor->name); - t_ij->buffer = simple_buf; - t_ij->view_src = tensor->view_src; - t_ij->view_offs = tensor->view_offs; - if (t_ij->view_src != nullptr && ggml_backend_buffer_is_meta(t_ij->view_src->buffer)) { - t_ij->view_src = ggml_backend_meta_buffer_simple_tensor(tensor->view_src, j); - if (t_ij->view_offs > 0 && split_dim >= 0 && split_dim < GGML_MAX_DIMS) { - GGML_ASSERT(tensor->ne[split_dim] != 0); - const int split_dim_view_src = ggml_backend_meta_get_split_state(tensor->view_src, /*assume_sync =*/ true).axis; - GGML_ASSERT(split_dim_view_src >= 0 && split_dim_view_src < GGML_MAX_DIMS); - - // The offset can be internal to the data split, in those cases the view offset should not be scaled. - // If however, the offset is larger than the data split then it needs to be scaled proportionally. - bool split_internal_offset = t_ij->view_offs <= tensor->view_src->nb[split_dim_view_src]; - for (int i = 0; i < GGML_MAX_DIMS; i++) { - const size_t dim_size = tensor->ne[i] * tensor->nb[i]; - if (tensor->view_offs <= dim_size && dim_size < tensor->nb[split_dim]) { - split_internal_offset = true; - break; - } - } - if (!split_internal_offset) { - t_ij->view_offs = t_ij->view_offs * ne[split_dim]/tensor->ne[split_dim]; - } - } - } - if (t_ij->view_src != nullptr) { - t_ij->data = (char *) t_ij->view_src->data + t_ij->view_offs; - } else if (simple_buf != nullptr) { - t_ij->data = (char *) ggml_backend_buffer_get_base(simple_buf) - + size_t(tensor->data) - size_t(ggml_backend_buffer_get_base(buffer)); - } - t_ij->extra = tensor->extra; - for (int i = 0; i < GGML_MAX_SRC; i++) { - t_ij->src[i] = tensor->src[i]; - if (tensor->src[i] == tensor) { - t_ij->src[i] = t_ij; - } else if (t_ij->src[i] != nullptr && ggml_backend_buffer_is_meta(t_ij->src[i]->buffer)) { - t_ij->src[i] = ggml_backend_meta_buffer_simple_tensor(tensor->src[i], j); - } - } - - simple_tensors.push_back(t_ij); - } - - // If one of the sources has a zero-sized slice, disable the computation: - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (tensor->src[i] == nullptr || !ggml_backend_buffer_is_meta(tensor->src[i]->buffer)) { - continue; - } - - const ggml_backend_meta_split_state split_state_src = ggml_backend_meta_get_split_state(tensor->src[i], /*assume_sync =*/ true); - if (split_state_src.axis < 0 || split_state_src.axis >= GGML_MAX_DIMS) { - continue; - } - for (size_t j = 0; j < n_simple_bufs; j++) { - int64_t ne_sum = 0; - for (size_t s = 0; s < split_state_src.n_segments; s++) { - ne_sum += split_state_src.ne[s*n_simple_bufs + j]; - } - if (ne_sum == 0) { - simple_tensors[j]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; - } - } - } - - buf_ctx->simple_tensors[tensor] = simple_tensors; - - return GGML_STATUS_SUCCESS; -} - -static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { - const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); - GGML_ASSERT(ggml_is_contiguous(tensor)); - - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); - - if (split_state.n_segments != 1) { - GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS); - GGML_ASSERT(tensor->ne[3] == 1); - - size_t offset_data = 0; - std::vector simple_offsets(n_bufs, 0); - if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) { - GGML_ASSERT(tensor->ne[2] == 1); - - const size_t row_stride = tensor->nb[1]; - GGML_ASSERT(offset % row_stride == 0); - GGML_ASSERT(size % row_stride == 0); - const int64_t r_start = offset / row_stride; - const int64_t r_count = size / row_stride; - GGML_ASSERT(r_start + r_count <= tensor->ne[1]); - - const int64_t blck_size = ggml_blck_size(tensor->type); - for (size_t s = 0; s < split_state.n_segments; s++) { - for (size_t j = 0; j < n_bufs; j++) { - ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0); - const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0]; - ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, - simple_offsets[j] + r_start * simple_tensor->nb[1], nbytes, - r_count, simple_tensor->nb[1], tensor->nb[1]); - offset_data += nbytes; - simple_offsets[j] += nbytes; - } - } - GGML_ASSERT(offset_data*r_count == size); - return; - } - GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1); - - const size_t row_stride = tensor->nb[2]; - GGML_ASSERT(offset % row_stride == 0); - GGML_ASSERT(size % row_stride == 0); - const int64_t r_start = offset / row_stride; - const int64_t r_count = size / row_stride; - GGML_ASSERT(r_start + r_count <= tensor->ne[2]); - - for (size_t s = 0; s < split_state.n_segments; s++) { - for (size_t j = 0; j < n_bufs; j++) { - ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[1]; - ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, - simple_offsets[j] + r_start * simple_tensor->nb[2], nbytes, - r_count, simple_tensor->nb[2], tensor->nb[2]); - offset_data += nbytes; - simple_offsets[j] += nbytes; - } - } - GGML_ASSERT(offset_data*r_count == size); - return; - } - - switch (split_state.axis) { - case GGML_BACKEND_SPLIT_AXIS_0: - case GGML_BACKEND_SPLIT_AXIS_1: - case GGML_BACKEND_SPLIT_AXIS_2: { - // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". - const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; - GGML_ASSERT(offset % chunk_size_full == 0); - GGML_ASSERT(size % chunk_size_full == 0); - const int64_t i_start = offset /chunk_size_full; - const int64_t i_stop = (offset + size)/chunk_size_full; - size_t offset_j = 0; - for (size_t j = 0; j < n_bufs; j++) { - ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; - const size_t simple_offset = i_start * chunk_size_j; - ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_j, simple_offset, chunk_size_j, i_stop - i_start, chunk_size_j, chunk_size_full); - offset_j += chunk_size_j; - } - GGML_ASSERT(offset_j == chunk_size_full); - } break; - case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { - for (size_t j = 0; j < n_bufs; j++) { - ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - ggml_backend_tensor_set(simple_tensor, data, offset, size); - } - } break; - case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { - GGML_ASSERT(tensor->type == GGML_TYPE_F32); - const int64_t ne = ggml_nelements(tensor); - std::vector tmp; - tmp.reserve(ne); - for (int64_t i = 0; i < ne; i++) { - tmp.push_back(((const float *) data)[i] / n_bufs); - } - for (size_t j = 0; j < n_bufs; j++) { - ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - ggml_backend_tensor_set(simple_tensor, tmp.data(), offset, size); - } - } break; - default: { - GGML_ABORT("fatal error"); - } - } -} - -static void ggml_backend_meta_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { - const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); - GGML_ASSERT(ggml_is_contiguous(tensor)); - - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); - - if (split_state.n_segments != 1) { - GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS); - GGML_ASSERT(tensor->ne[3] == 1); - - size_t offset_data = 0; - std::vector simple_offsets(n_bufs, 0); - if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) { - GGML_ASSERT(tensor->ne[2] == 1); - - const size_t row_stride = tensor->nb[1]; - GGML_ASSERT(offset % row_stride == 0); - GGML_ASSERT(size % row_stride == 0); - const int64_t r_start = offset / row_stride; - const int64_t r_count = size / row_stride; - GGML_ASSERT(r_start + r_count <= tensor->ne[1]); - - const int64_t blck_size = ggml_blck_size(tensor->type); - for (size_t s = 0; s < split_state.n_segments; s++) { - for (size_t j = 0; j < n_bufs; j++) { - const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0); - const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0]; - ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, - simple_offsets[j] + r_start * simple_tensor->nb[1], nbytes, - r_count, simple_tensor->nb[1], tensor->nb[1]); - offset_data += nbytes; - simple_offsets[j] += nbytes; - } - } - GGML_ASSERT(offset_data*r_count == size); - return; - } - GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1); - - const size_t row_stride = tensor->nb[2]; - GGML_ASSERT(offset % row_stride == 0); - GGML_ASSERT(size % row_stride == 0); - const int64_t r_start = offset / row_stride; - const int64_t r_count = size / row_stride; - GGML_ASSERT(r_start + r_count <= tensor->ne[2]); - - for (size_t s = 0; s < split_state.n_segments; s++) { - for (size_t j = 0; j < n_bufs; j++) { - const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[1]; - ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, - simple_offsets[j] + r_start * simple_tensor->nb[2], nbytes, - r_count, simple_tensor->nb[2], tensor->nb[2]); - offset_data += nbytes; - simple_offsets[j] += nbytes; - } - } - GGML_ASSERT(offset_data*r_count == size); - return; - } - - switch (split_state.axis) { - case GGML_BACKEND_SPLIT_AXIS_0: - case GGML_BACKEND_SPLIT_AXIS_1: - case GGML_BACKEND_SPLIT_AXIS_2: { - // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". - const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; - GGML_ASSERT(offset % chunk_size_full == 0); - GGML_ASSERT(size % chunk_size_full == 0); - const int64_t i_start = offset /chunk_size_full; - const int64_t i_stop = (offset + size)/chunk_size_full; - size_t offset_j = 0; - for (size_t j = 0; j < n_bufs; j++){ - const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; - const size_t simple_offset = i_start * chunk_size_j; - ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_j, simple_offset, chunk_size_j, i_stop - i_start, chunk_size_j, chunk_size_full); - offset_j += chunk_size_j; - } - GGML_ASSERT(offset_j == chunk_size_full); - } break; - case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { - // TODO other simple backend may be better - const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, 0); - ggml_backend_tensor_get(simple_tensor, data, offset, size); - } break; - default: { - GGML_ABORT("fatal error"); - } - } -} - -static void ggml_backend_meta_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { - const size_t n_buffers = ggml_backend_meta_buffer_n_bufs(buffer); - for (size_t i = 0; i < n_buffers; i++) { - ggml_backend_buffer_clear(ggml_backend_meta_buffer_simple_buffer(buffer, i), value); - } -} - -static void ggml_backend_meta_buffer_reset(ggml_backend_buffer_t buffer) { - const size_t n_buffers = ggml_backend_meta_buffer_n_bufs(buffer); - for (size_t i = 0; i < n_buffers; i++) { - ggml_backend_buffer_reset(ggml_backend_meta_buffer_simple_buffer(buffer, i)); - } -} - -static const ggml_backend_buffer_i ggml_backend_meta_buffer_iface = { - /* .free_buffer = */ ggml_backend_meta_buffer_free_buffer, - /* .get_base = */ ggml_backend_meta_buffer_get_base, - /* .init_tensor = */ ggml_backend_meta_buffer_init_tensor, - /* .memset_tensor = */ nullptr, // TODO implement - /* .set_tensor = */ ggml_backend_meta_buffer_set_tensor, - /* .get_tensor = */ ggml_backend_meta_buffer_get_tensor, - /* .set_tensor_2d = */ nullptr, - /* .get_tensor_2d = */ nullptr, - /* .cpy_tensor = */ nullptr, - /* .clear = */ ggml_backend_meta_buffer_clear, - /* .reset = */ ggml_backend_meta_buffer_reset, -}; - -bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf) { - return buf != nullptr && buf->iface.free_buffer == ggml_backend_meta_buffer_iface.free_buffer; -} - -static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { - const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); - - ggml_init_params params = { - /*.mem_size =*/ 1024*1024*1024, // FIXME - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - - ggml_backend_meta_buffer_context * buf_ctx = new ggml_backend_meta_buffer_context(); - size_t max_size = 0; - buf_ctx->buf_configs.reserve(n_simple_bufts); - for (size_t i = 0; i < n_simple_bufts; i++) { - ggml_backend_buffer_t simple_buf = ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size); - max_size = std::max(max_size, ggml_backend_buffer_get_size(simple_buf)); - buf_ctx->buf_configs.emplace_back(ggml_init(params), simple_buf); - } - - return ggml_backend_buffer_init(buft, ggml_backend_meta_buffer_iface, buf_ctx, max_size); -} - -struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft) { - const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); - - ggml_init_params params = { - /*.mem_size =*/ 1024*1024*1024, // FIXME - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - - ggml_backend_meta_buffer_context * meta_buf_ctx = new ggml_backend_meta_buffer_context(); - meta_buf_ctx->buf_configs.reserve(n_simple_bufts); - for (size_t i = 0; i < n_simple_bufts; i++) { - meta_buf_ctx->buf_configs.emplace_back(ggml_init(params), nullptr); - } - - ggml_backend_buffer_t meta_buf = ggml_backend_buffer_init(buft, ggml_backend_meta_buffer_iface, meta_buf_ctx, 0); - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { - t->buffer = meta_buf; - ggml_backend_meta_buffer_init_tensor(meta_buf, t); - t->data = (void *) 0x2000000000000000; // FIXME - } - for (size_t i = 0; i < n_simple_bufts; i++) { - meta_buf_ctx->buf_configs[i].buf = ggml_backend_alloc_ctx_tensors_from_buft( - meta_buf_ctx->buf_configs[i].ctx, ggml_backend_meta_buft_simple_buft(buft, i)); - meta_buf->size = std::max(meta_buf->size, ggml_backend_buffer_get_size(meta_buf_ctx->buf_configs[i].buf)); - } - return meta_buf; -} - -// -// meta backend -// - -static ggml_guid_t ggml_backend_meta_guid() { - static ggml_guid guid = {0xf1, 0x0e, 0x34, 0xcf, 0x9c, 0x6f, 0x43, 0xcb, 0x96, 0x92, 0xbe, 0x8e, 0xbb, 0x71, 0x3f, 0xda}; - return &guid; -} - -struct ggml_backend_meta_context { - struct cgraph_config { - ggml_cgraph * cgraph_main = nullptr; - int offset = 0; // Node offset vs. original graph - - std::vector cgraphs_aux; - }; - struct backend_config { - ggml_backend_t backend; - - std::vector cgraphs; - std::vector nodes; - std::vector bufs; - - backend_config(ggml_backend_t backend, const size_t n_reduce_steps) : backend(backend) { - bufs.resize(n_reduce_steps); - } - }; - std::string name; - std::vector backend_configs; - ggml_context_ptr ctx; - std::vector cgraphs_aux; - std::vector nodes_aux; - size_t n_reduce_steps; - int max_nnodes = 0; - size_t max_tmp_size = 0; - size_t max_subgraphs = 0; - size_t n_subgraphs = 0; - uint64_t uid = 0; - - void * comm_ctx = nullptr; - ggml_backend_comm_allreduce_tensor_t comm_allreduce = nullptr; - - ggml_backend_meta_context(ggml_backend_dev_t meta_dev, const char * params) { - const size_t n_devs = ggml_backend_meta_dev_n_devs(meta_dev); - n_reduce_steps = std::ceil(std::log2(n_devs)); - name = "Meta("; - std::vector simple_backends; - backend_configs.reserve(n_devs); - simple_backends.reserve(n_devs); - for (size_t i = 0; i < n_devs; i++) { - ggml_backend_dev_t simple_dev = ggml_backend_meta_dev_simple_dev(meta_dev, i); - if (i > 0) { - name += ","; - } - name += ggml_backend_dev_name(simple_dev); - simple_backends.push_back(ggml_backend_dev_init(simple_dev, params)); - backend_configs.emplace_back(simple_backends.back(), n_reduce_steps); - } - name += ")"; - - if (n_devs > 1) { - ggml_backend_comm_init_t comm_init = (ggml_backend_comm_init_t) ggml_backend_reg_get_proc_address( - ggml_backend_dev_backend_reg(ggml_backend_get_device(simple_backends[0])), "ggml_backend_comm_init"); - if (comm_init != nullptr) { - comm_ctx = comm_init(simple_backends.data(), simple_backends.size()); - } - } - if (comm_ctx != nullptr) { - comm_allreduce = (ggml_backend_comm_allreduce_tensor_t) - ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg( - ggml_backend_get_device(simple_backends[0])), "ggml_backend_comm_allreduce_tensor"); - GGML_ASSERT(comm_allreduce != nullptr); - } - } - - ~ggml_backend_meta_context() { - if (comm_ctx != nullptr) { - ggml_backend_comm_free_t comm_free = (ggml_backend_comm_free_t) ggml_backend_reg_get_proc_address( - ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_configs[0].backend)), "ggml_backend_comm_free"); - GGML_ASSERT(comm_free != nullptr); - comm_free(comm_ctx); - } - for (auto & bc : backend_configs) { - ggml_backend_free(bc.backend); - } - } -}; - -static const char * ggml_backend_meta_get_name(ggml_backend_t backend) { - GGML_ASSERT(ggml_backend_is_meta(backend)); - const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) backend->context; - return backend_ctx->name.c_str(); -} - -static void ggml_backend_meta_free(ggml_backend_t backend) { - GGML_ASSERT(ggml_backend_is_meta(backend)); - ggml_backend_meta_context * backend_ctx = (ggml_backend_meta_context *) backend->context; - delete backend_ctx; - delete backend; -} - -static void ggml_backend_meta_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { - const size_t n_backends = ggml_backend_meta_n_backends(backend); - GGML_ASSERT(offset == 0); - GGML_ASSERT(ggml_is_contiguous(tensor)); - - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); - GGML_ASSERT(split_state.n_segments == 1); - - switch (split_state.axis) { - case GGML_BACKEND_SPLIT_AXIS_0: - case GGML_BACKEND_SPLIT_AXIS_1: - case GGML_BACKEND_SPLIT_AXIS_2: { - // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". - const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; - GGML_ASSERT(offset % chunk_size_full == 0); - GGML_ASSERT(size % chunk_size_full == 0); - const int64_t i_start = offset /chunk_size_full; - const int64_t i_stop = (offset + size)/chunk_size_full; - size_t offset_j = 0; - for (size_t j = 0; j < n_backends; j++){ - ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, j); - ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; - ggml_backend_tensor_set_2d_async(simple_backend, simple_tensor, (const char *) data + offset_j, offset, chunk_size_j, - i_stop - i_start, chunk_size_j, chunk_size_full); - offset_j += chunk_size_j; - } - GGML_ASSERT(offset_j == chunk_size_full); - } break; - case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { - for (size_t j = 0; j < n_backends; j++) { - ggml_backend_tensor_set_async( - ggml_backend_meta_simple_backend(backend, j), ggml_backend_meta_buffer_simple_tensor(tensor, j), data, offset, size); - } - } break; - default: { - GGML_ABORT("fatal error"); - } - } -} - -static void ggml_backend_meta_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { - const size_t n_backends = ggml_backend_meta_n_backends(backend); - GGML_ASSERT(offset == 0); - GGML_ASSERT(ggml_is_contiguous(tensor)); - - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); - GGML_ASSERT(split_state.n_segments == 1); - - switch (split_state.axis) { - case GGML_BACKEND_SPLIT_AXIS_0: - case GGML_BACKEND_SPLIT_AXIS_1: - case GGML_BACKEND_SPLIT_AXIS_2: { - // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". - const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; - GGML_ASSERT(offset % chunk_size_full == 0); - GGML_ASSERT(size % chunk_size_full == 0); - const int64_t i_start = offset /chunk_size_full; - const int64_t i_stop = (offset + size)/chunk_size_full; - size_t offset_j = 0; - for (size_t j = 0; j < n_backends; j++){ - ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, j); - const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); - const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; - ggml_backend_tensor_get_2d_async(simple_backend, simple_tensor, (char *) data + offset_j, offset, chunk_size_j, - i_stop - i_start, chunk_size_j, chunk_size_full); - offset_j += chunk_size_j; - } - GGML_ASSERT(offset_j == chunk_size_full); - } break; - case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { - // TODO other simple backend may be better - ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, 0); - const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, 0); - ggml_backend_tensor_get_async(simple_backend, simple_tensor, data, offset, size); - } break; - default: { - GGML_ABORT("fatal error"); - } - } -} - -static void ggml_backend_meta_synchronize(ggml_backend_t backend) { - const size_t n_backends = ggml_backend_meta_n_backends(backend); - for (size_t i = 0; i < n_backends; i++) { - ggml_backend_synchronize(ggml_backend_meta_simple_backend(backend, i)); - } -} - -static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, struct ggml_cgraph * cgraph) { - GGML_ASSERT(cgraph->grads == nullptr); - const size_t n_backends = ggml_backend_meta_n_backends(backend); - ggml_backend_meta_context * backend_ctx = (ggml_backend_meta_context *) backend->context; - - // If the previous cgraph had a defined UID it can be used to skip rebuilding the subgraphs per simple backend. - const bool needs_rebuild = (cgraph->uid == 0) || (cgraph->uid != backend_ctx->uid); - - bool max_nnodes_raised = false; - if (cgraph->n_nodes > backend_ctx->max_nnodes) { - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - bcj.nodes.resize(cgraph->n_nodes); - bcj.cgraphs.resize(cgraph->n_nodes); - } - backend_ctx->max_nnodes = cgraph->n_nodes; - max_nnodes_raised = true; - assert(needs_rebuild); - } - - if (needs_rebuild) { - size_t n_subgraphs = 0; - size_t max_tmp_size = 0; - - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - - for (int i = 0; i < cgraph->n_nodes; i++) { - ggml_tensor * node = cgraph->nodes[i]; - if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) { - // FIXME s_copy_main is on the CPU and its view seems to be incorrectly added to the graph nodes. - // For regular usage this doesn't matter since it's a noop but trying to call ggml_backend_meta_buffer_simple_tensor results in a crash. - bcj.nodes[i] = node; - continue; - } - bcj.nodes[i] = ggml_backend_meta_buffer_simple_tensor(node, j); - GGML_ASSERT(bcj.nodes[i]); - } - } - - { - // For MoE models it may make sense to delay the AllReduce in order to reduce I/O: - auto get_i_delayed = [&](const int i) -> int { - int id = i; // i_delayed - int idr = i; // i_delayed return, last safe return value - - ggml_tensor * node = cgraph->nodes[id]; - int32_t n_used = ggml_node_get_use_count(cgraph, id); - - // Skip MIRRORED nodes that don't consume node - auto skip_unrelated = [&]() { - while (id + 1 < cgraph->n_nodes) { - ggml_tensor * next = cgraph->nodes[id+1]; - if (ggml_backend_meta_get_split_state(next, false).axis != GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - break; - } - bool safe = true; - for (int s = 0; s < GGML_MAX_SRC; s++) { - if (next->src[s] == nullptr) { - continue; - } - if (next->src[s] == node) { - safe = false; - break; - } - if (ggml_backend_meta_get_split_state(next->src[s], false).axis != GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - safe = false; - break; - } - } - if (!safe) { - break; - } - id++; - } - }; - - skip_unrelated(); - if (id + 1 >= cgraph->n_nodes) { - return idr; - } - { - ggml_tensor * next = cgraph->nodes[id+1]; - if (next->op == GGML_OP_ADD_ID && next->src[0] == node && - ggml_backend_meta_get_split_state(next->src[1], false).axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL && - ggml_backend_meta_get_split_state(next->src[2], false).axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - node = next; - id++; - idr = id; - n_used = ggml_node_get_use_count(cgraph, id); - } - } - // Chain of MULs with MIRRORED src[1] - while (true) { - skip_unrelated(); - if (id + 1 >= cgraph->n_nodes) { - return idr; - } - ggml_tensor * next = cgraph->nodes[id+1]; - if (next->op == GGML_OP_MUL && next->src[0] == node && - ggml_backend_meta_get_split_state(next->src[1], false).axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { - node = next; - id++; - idr = id; - n_used = ggml_node_get_use_count(cgraph, id); - } else { - break; - } - } - - if (n_used != node->ne[1] || id + 2*n_used-1 >= cgraph->n_nodes) { - return idr; - } - for (int32_t k = 0; k < n_used; k++) { - ggml_tensor * next = cgraph->nodes[id+1]; - if (next->op != GGML_OP_VIEW || next->view_src != node || next->view_offs != k*node->nb[1] || - next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] || next->nb[1] != node->nb[2] || - ggml_node_get_use_count(cgraph, id+1) != 1) { - return idr; - } - id++; - } - { - ggml_tensor * next = cgraph->nodes[id+1]; - if (next->op != GGML_OP_ADD || next->src[0] != cgraph->nodes[id - (n_used-1)] || - next->src[1] != cgraph->nodes[id - (n_used-2)] || ggml_node_get_use_count(cgraph, id+1) != 1) { - return idr; - } - id++; - } - for (int32_t k = 0; k < n_used - 2; k++) { - ggml_tensor * next = cgraph->nodes[id+1]; - if (next->op != GGML_OP_ADD || next->src[0] != cgraph->nodes[id] || - next->src[1] != cgraph->nodes[id - (n_used-2)] || ggml_node_get_use_count(cgraph, id+1) != 1) { - return idr; - } - id++; - } - idr = id; - return idr; - }; - - int i_start = 0; - for (int i = 0; i < cgraph->n_nodes; i++) { - ggml_tensor * node = cgraph->nodes[i]; - if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) { - continue; - } - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(node, /*assume_sync =*/ false); - if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { - max_tmp_size = std::max(max_tmp_size, ggml_nbytes(node)); - } - const bool new_subgraph = i + 1 == cgraph->n_nodes || split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL; - if (!new_subgraph) { - continue; - } - - const int i_delayed = get_i_delayed(i); - - // If we can delay the AllReduce we need to consider the interaction with zero-sized tensor slices. - // A backend with such a slice would normally have valid data after participating in the AllReduce with a node that has - // its compute flag disabled and thus gets its data zeroed out. - // If the AllReduce is delayed then the nodes until that point also need to have their compute flag disabled. - if (i_delayed > i) { - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - if ((bcj.nodes[i]->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { - for (int ii = i + 1; ii <= i_delayed; ii++) { - bcj.nodes[ii]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; - } - } - } - } - - i = i_delayed; - - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - bcj.cgraphs[n_subgraphs].offset = i_start; - } - n_subgraphs++; - i_start = i + 1; - } - GGML_ASSERT(i_start == cgraph->n_nodes); - } - - backend_ctx->uid = cgraph->uid; - backend_ctx->n_subgraphs = n_subgraphs; - - if (max_tmp_size > backend_ctx->max_tmp_size) { - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - for (size_t i = 0; i < backend_ctx->n_reduce_steps; i++) { - bcj.bufs[i].reset(ggml_backend_alloc_buffer(bcj.backend, max_tmp_size)); - } - } - backend_ctx->max_tmp_size = max_tmp_size; - } - - if (max_nnodes_raised || n_subgraphs > backend_ctx->max_subgraphs) { - backend_ctx->max_subgraphs = std::max(backend_ctx->max_subgraphs, n_subgraphs); - const size_t n_nodes_per_device = 3 * backend_ctx->n_reduce_steps; // tmp + ADD (+zeroing) graph per step and device - const size_t n_cgraphs_per_device = 2 * backend_ctx->n_reduce_steps; // ADD ( + zeroing) graph per step and device - const size_t mem_per_device_graphs_main = backend_ctx->max_subgraphs*ggml_graph_overhead_custom(backend_ctx->max_nnodes, cgraph->grads); - const size_t mem_per_device_graphs_aux = n_cgraphs_per_device*backend_ctx->max_subgraphs*ggml_graph_overhead_custom(1, cgraph->grads); - const size_t mem_per_device_nodes_aux = n_nodes_per_device*backend_ctx->max_subgraphs*ggml_tensor_overhead(); - ggml_init_params params = { - /*.mem_size =*/ n_backends * (mem_per_device_graphs_main + mem_per_device_graphs_aux + mem_per_device_nodes_aux), - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - backend_ctx->ctx.reset(ggml_init(params)); - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - for (size_t i = 0; i < n_subgraphs; i++) { - bcj.cgraphs[i].cgraph_main = ggml_new_graph_custom(backend_ctx->ctx.get(), cgraph->n_nodes, /*grads =*/ false); - } - } - backend_ctx->cgraphs_aux.resize(n_backends*n_cgraphs_per_device*backend_ctx->max_subgraphs); - for (size_t k = 0; k < backend_ctx->cgraphs_aux.size(); k++) { - backend_ctx->cgraphs_aux[k] = ggml_new_graph_custom(backend_ctx->ctx.get(), 1, cgraph->grads); - } - backend_ctx->nodes_aux.resize(n_backends*n_nodes_per_device*backend_ctx->max_subgraphs); - for (size_t k = 0; k < backend_ctx->nodes_aux.size(); k++) { - backend_ctx->nodes_aux[k] = ggml_new_tensor_1d(backend_ctx->ctx.get(), GGML_TYPE_F32, 1); - } - } - - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - for (size_t i_graph = 0; i_graph < n_subgraphs; i_graph++) { - ggml_cgraph * cgraph_ij = bcj.cgraphs[i_graph].cgraph_main; - const size_t i_node_start = bcj.cgraphs[i_graph].offset; - const size_t i_node_stop = i_graph + 1 < n_subgraphs ? bcj.cgraphs[i_graph + 1].offset : cgraph->n_nodes; - cgraph_ij->n_nodes = i_node_stop - i_node_start; - ggml_hash_set_reset(&cgraph_ij->visited_hash_set); - for (size_t i_node = i_node_start; i_node < i_node_stop; i_node++) { - ggml_tensor * node_ij = bcj.nodes[i_node]; - cgraph_ij->nodes[i_node - i_node_start] = node_ij; - const size_t hash_pos_orig = ggml_hash_find(&cgraph->visited_hash_set, cgraph->nodes[i_node]); - const size_t hash_pos_ij = ggml_hash_insert(&cgraph_ij->visited_hash_set, node_ij); - cgraph_ij->use_counts[hash_pos_ij] = cgraph->use_counts[hash_pos_orig]; - } - cgraph_ij->uid = ggml_graph_next_uid(); - } - } - } - - size_t iga = 0; // i graph aux - size_t ina = 0; // i node aux - - auto get_node_aux = [&](ggml_tensor * t) -> ggml_tensor * { - ggml_tensor * ret = backend_ctx->nodes_aux[ina++]; - memset(ret, 0, sizeof(ggml_tensor)); - ret->op = GGML_OP_NONE; - ret->type = t->type; - for (size_t k = 0; k < GGML_MAX_DIMS; k++) { - ret->ne[k] = t->ne[k]; - ret->nb[k] = t->nb[k]; - } - return ret; - }; - auto set_tmp_data = [&](ggml_tensor * tensor, const size_t j, const size_t i_buf) { - auto & bcj = backend_ctx->backend_configs[j]; - ggml_backend_buffer_ptr & buf_ptr = bcj.bufs[i_buf]; - if (!buf_ptr || ggml_backend_buffer_get_size(buf_ptr.get()) < backend_ctx->max_tmp_size) { - buf_ptr.reset(ggml_backend_alloc_buffer(bcj.backend, backend_ctx->max_tmp_size)); - } - tensor->buffer = buf_ptr.get(); - tensor->data = ggml_backend_buffer_get_base(buf_ptr.get()); - }; - // FIXME usage_counts - auto get_cgraph_aux = [&]() -> ggml_cgraph * { - ggml_cgraph * ret = backend_ctx->cgraphs_aux[iga++]; - return ret; - }; - - // Preferentially use backend-specific allreduce_tensor_async (e.g. NCCL for CUDA), use a generic fallback if unavailable: - auto allreduce_fallback = [&](size_t i) -> ggml_status { - std::vector step_cgraphs(n_backends, nullptr); - - // Zero out nodes that were disabled due to having a zero-sized slice: - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - ggml_tensor * node = bcj.cgraphs[i].cgraph_main->nodes[bcj.cgraphs[i].cgraph_main->n_nodes - 1]; - if (node->flags & GGML_TENSOR_FLAG_COMPUTE) { - continue; - } - ggml_tensor * node_zero = get_node_aux(node); - node_zero->op = GGML_OP_SCALE; // FIXME 0.0f * NaN == NaN - node_zero->src[0] = node; - ggml_set_op_params_f32(node_zero, 0, 0.0f); - node_zero->data = node->data; - node_zero->flags |= GGML_TENSOR_FLAG_COMPUTE; - - step_cgraphs[j] = get_cgraph_aux(); - step_cgraphs[j]->nodes[0] = node_zero; - step_cgraphs[j]->n_nodes = 1; - const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, step_cgraphs[j]); - if (status != GGML_STATUS_SUCCESS) { - return status; - } - } - std::fill(step_cgraphs.begin(), step_cgraphs.end(), nullptr); - - auto push_data = [&](const size_t j_src, const size_t j_dst, const size_t i_buf) { - assert(step_cgraphs[j_dst] == nullptr); - auto & bcj_src = backend_ctx->backend_configs[j_src]; - auto & bcj_dst = backend_ctx->backend_configs[j_dst]; - - ggml_tensor * node_src = bcj_src.cgraphs[i].cgraph_main->nodes[bcj_src.cgraphs[i].cgraph_main->n_nodes - 1]; - ggml_tensor * node_dst = bcj_dst.cgraphs[i].cgraph_main->nodes[bcj_dst.cgraphs[i].cgraph_main->n_nodes - 1]; - GGML_ASSERT(ggml_is_contiguous(node_src)); - GGML_ASSERT(ggml_is_contiguous(node_dst)); - - ggml_tensor * node_tmp = get_node_aux(node_dst); - set_tmp_data(node_tmp, j_dst, i_buf); - - ggml_backend_tensor_copy_async(bcj_src.backend, bcj_dst.backend, node_src, node_tmp); - - ggml_tensor * node_red = get_node_aux(node_dst); - node_red->view_src = node_dst->view_src == nullptr ? node_dst : node_dst->view_src; - node_red->view_offs = node_dst->view_offs; - node_red->op = GGML_OP_ADD; - node_red->src[0] = node_dst; - node_red->src[1] = node_tmp; - node_red->flags |= GGML_TENSOR_FLAG_COMPUTE; - ggml_backend_view_init(node_red); - - ggml_cgraph * cgraph_aux = get_cgraph_aux(); - cgraph_aux->nodes[0] = node_red; - cgraph_aux->n_nodes = 1; - step_cgraphs[j_dst] = cgraph_aux; - }; - - size_t offset_j = n_backends/2; - while ((offset_j & (offset_j - 1)) != 0) { - offset_j--; - } - const size_t offset_j_max = offset_j; - size_t i_buf = 0; - - // If n_backends is not a power of 2, fold in the excess prior to butterfly reduction: - for (size_t j_src = 2*offset_j_max; j_src < n_backends; j_src++) { - const size_t j_dst = j_src - 2*offset_j_max; - push_data(j_src, j_dst, i_buf); - const ggml_status status = ggml_backend_graph_compute_async(backend_ctx->backend_configs[j_dst].backend, step_cgraphs[j_dst]); - if (status != GGML_STATUS_SUCCESS) { - return status; - } - i_buf = 1; - } - - // Butterfly reduction: - for (; offset_j >= 1; offset_j /= 2) { - std::fill(step_cgraphs.begin(), step_cgraphs.end(), nullptr); - - for (size_t j = 0; j < 2*offset_j_max; j++) { - const size_t j_other = j ^ offset_j; - if (j_other >= n_backends) { - continue; - } - push_data(j, j_other, i_buf); - } - - for (size_t j = 0; j < 2*offset_j_max; j++) { - if (step_cgraphs[j] == nullptr) { - continue; - } - auto & bcj = backend_ctx->backend_configs[j]; - const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, step_cgraphs[j]); - if (status != GGML_STATUS_SUCCESS) { - return status; - } - } - i_buf++; - } - assert(i_buf == backend_ctx->n_reduce_steps); - - // If n_backends is not a power of 2, copy back the reduced tensors to the excess: - for (size_t j = 2*offset_j_max; j < n_backends; j++) { - auto & bcj_src = backend_ctx->backend_configs[j - 2*offset_j_max]; - auto & bcj_dst = backend_ctx->backend_configs[j]; - - ggml_tensor * node_src = bcj_src.cgraphs[i].cgraph_main->nodes[bcj_src.cgraphs[i].cgraph_main->n_nodes - 1]; - ggml_tensor * node_dst = bcj_dst.cgraphs[i].cgraph_main->nodes[bcj_dst.cgraphs[i].cgraph_main->n_nodes - 1]; - ggml_backend_tensor_copy_async(bcj_src.backend, bcj_dst.backend, node_src, node_dst); - } - - return GGML_STATUS_SUCCESS; - }; - - - for (size_t i = 0; i < backend_ctx->n_subgraphs; i++) { - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, bcj.cgraphs[i].cgraph_main); - if (status != GGML_STATUS_SUCCESS) { - return status; - } - } - - if (n_backends > 1 && i < backend_ctx->n_subgraphs - 1) { - bool backend_allreduce_success = false; - if (backend_ctx->comm_ctx) { - std::vector nodes; - nodes.reserve(n_backends); - for (size_t j = 0; j < n_backends; j++) { - auto & bcj = backend_ctx->backend_configs[j]; - ggml_cgraph * cgraph_ij = bcj.cgraphs[i].cgraph_main; - nodes.push_back(cgraph_ij->nodes[cgraph_ij->n_nodes-1]); - } - backend_allreduce_success = backend_ctx->comm_allreduce(backend_ctx->comm_ctx, nodes.data()); - } - - if (!backend_allreduce_success) { - const ggml_status status = allreduce_fallback(i); - if (status != GGML_STATUS_SUCCESS) { - return status; - } - } - } - } - return GGML_STATUS_SUCCESS; -} - -static const ggml_backend_i ggml_backend_meta_i = { - /* .get_name = */ ggml_backend_meta_get_name, - /* .free = */ ggml_backend_meta_free, - /* .set_tensor_async = */ ggml_backend_meta_set_tensor_async, - /* .get_tensor_async = */ ggml_backend_meta_get_tensor_async, - /* .set_tensor_2d_async = */ nullptr, - /* .get_tensor_2d_async = */ nullptr, - /* .cpy_tensor_async = */ nullptr, - /* .synchronize = */ ggml_backend_meta_synchronize, - /* .graph_plan_create = */ nullptr, - /* .graph_plan_free = */ nullptr, - /* .graph_plan_update = */ nullptr, - /* .graph_plan_compute = */ nullptr, - /* .graph_compute = */ ggml_backend_meta_graph_compute, - /* .event_record = */ nullptr, - /* .event_wait = */ nullptr, - /* .graph_optimize = */ nullptr, -}; - -bool ggml_backend_is_meta(ggml_backend_t backend) { - return backend != nullptr && backend->iface.get_name == ggml_backend_meta_i.get_name; -} - -static ggml_backend_t ggml_backend_meta_device_init_backend(ggml_backend_dev_t dev, const char * params) { - ggml_backend_meta_context * backend_ctx = new ggml_backend_meta_context(dev, params); - - ggml_backend_t backend = new struct ggml_backend; - backend->guid = ggml_backend_meta_guid(); - backend->iface = ggml_backend_meta_i; - backend->device = dev; - backend->context = backend_ctx; - return backend; -} - -size_t ggml_backend_meta_n_backends(ggml_backend_t meta_backend) { - GGML_ASSERT(ggml_backend_is_meta(meta_backend)); - const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context; - return backend_ctx->backend_configs.size(); -} - -ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, size_t index) { - GGML_ASSERT(ggml_backend_is_meta(meta_backend)); - const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context; - return backend_ctx->backend_configs[index].backend; -} - diff --git a/ggml/src/ggml-backend-reg.cpp b/ggml/src/ggml-backend-reg.cpp.inc similarity index 84% rename from ggml/src/ggml-backend-reg.cpp rename to ggml/src/ggml-backend-reg.cpp.inc index 8165ae2c8bbe..df87ac5ae27c 100644 --- a/ggml/src/ggml-backend-reg.cpp +++ b/ggml/src/ggml-backend-reg.cpp.inc @@ -1,7 +1,74 @@ -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "ggml-backend-dl.h" -#include "ggml-impl.h" +#include "ggml-backend-impl-defs.inc" +#include "ggml-backend.h.inc" +#ifdef _WIN32 +# define WIN32_LEAN_AND_MEAN +# ifndef NOMINMAX +# define NOMINMAX +# endif +# include +# include +#else +# include +# include +#endif +#include +#include + +namespace fs = std::filesystem; + +extern "C" { +void * ggml_backend_dl_load_library_rust(const char * path); +void * ggml_backend_dl_get_sym_rust(void * handle, const char * name); +const char * ggml_backend_dl_error_rust(); +int ggml_backend_striequals_rust(const char * a, const char * b); +char * ggml_backend_current_exe_dir_rust(); +void ggml_backend_cstring_free_rust(char * s); +const char * ggml_backend_filename_prefix_rust(); +const char * ggml_backend_filename_extension_rust(); +} + +#ifdef _WIN32 + +using dl_handle = std::remove_pointer_t; + +struct dl_handle_deleter { + void operator()(HMODULE handle) { + FreeLibrary(handle); + } +}; + +#else + +using dl_handle = void; + +struct dl_handle_deleter { + void operator()(void * handle) { + dlclose(handle); + } +}; + +#endif + +using dl_handle_ptr = std::unique_ptr; + +inline dl_handle * dl_load_library(const fs::path & path) { +#ifdef _WIN32 + const std::string path_str = path.u8string(); +#else + const std::string path_str = path.string(); +#endif + return static_cast(ggml_backend_dl_load_library_rust(path_str.c_str())); +} + +inline void * dl_get_sym(dl_handle * handle, const char * name) { + return ggml_backend_dl_get_sym_rust(handle, name); +} + +inline const char * dl_error() { + return ggml_backend_dl_error_rust(); +} + +#include "ggml-impl-defs.inc" #include #include #include @@ -9,81 +76,66 @@ #include #include #include -#include - -#ifdef _WIN32 -# define WIN32_LEAN_AND_MEAN -# ifndef NOMINMAX -# define NOMINMAX -# endif -# include -#elif defined(__APPLE__) -# include -# include -#else -# include -# include -#endif // Backend registry #ifdef GGML_USE_CPU -#include "ggml-cpu.h" +#include "ggml-cpu.h.inc" #endif #ifdef GGML_USE_CUDA -#include "ggml-cuda.h" +#include "ggml-cuda.h.inc" #endif #ifdef GGML_USE_METAL -#include "ggml-metal.h" +#include "ggml-metal.h.inc" #endif #ifdef GGML_USE_SYCL -#include "ggml-sycl.h" +#include "ggml-sycl.h.inc" #endif #ifdef GGML_USE_VULKAN -#include "ggml-vulkan.h" +#include "ggml-vulkan.h.inc" #endif #ifdef GGML_USE_WEBGPU -#include "ggml-webgpu.h" +#include "ggml-webgpu.h.inc" #endif #ifdef GGML_USE_ZDNN -#include "ggml-zdnn.h" +#include "ggml-zdnn.h.inc" #endif #ifdef GGML_USE_OPENCL -#include "ggml-opencl.h" +#include "ggml-opencl.h.inc" #endif #ifdef GGML_USE_HEXAGON -#include "ggml-hexagon.h" +#include "ggml-hexagon.h.inc" #endif #ifdef GGML_USE_BLAS -#include "ggml-blas.h" +#include "ggml-blas.h.inc" #endif #ifdef GGML_USE_RPC -#include "ggml-rpc.h" +#include "ggml-rpc.h.inc" #endif #ifdef GGML_USE_VIRTGPU_FRONTEND -#include "ggml-virtgpu.h" +#include "ggml-virtgpu.h.inc" #endif #ifdef GGML_USE_CANN -#include "ggml-cann.h" +#include "ggml-cann.h.inc" #endif #ifdef GGML_USE_ZENDNN -#include "ggml-zendnn.h" +#include "ggml-zendnn.h.inc" #endif #ifdef GGML_USE_OPENVINO -#include "ggml-openvino.h" +#include "ggml-openvino.h.inc" #endif namespace fs = std::filesystem; @@ -298,12 +350,7 @@ void ggml_backend_device_register(ggml_backend_dev_t device) { // Backend (reg) enumeration static bool striequals(const char * a, const char * b) { - for (; *a && *b; a++, b++) { - if (std::tolower(*a) != std::tolower(*b)) { - return false; - } - } - return *a == *b; + return ggml_backend_striequals_rust(a, b) != 0; } size_t ggml_backend_reg_count() { @@ -392,82 +439,22 @@ void ggml_backend_unload(ggml_backend_reg_t reg) { } static fs::path get_executable_path() { -#if defined(__APPLE__) - // get executable path - std::vector path; - uint32_t size; - while (true) { - size = path.size(); - if (_NSGetExecutablePath(path.data(), &size) == 0) { - break; - } - path.resize(size); - } - std::string base_path(path.data(), size); - // remove executable name - auto last_slash = base_path.find_last_of('/'); - if (last_slash != std::string::npos) { - base_path = base_path.substr(0, last_slash); - } - return base_path + "/"; -#elif defined(__linux__) || defined(__FreeBSD__) - std::string base_path = "."; - std::vector path(1024); - while (true) { - // get executable path -# if defined(__linux__) - ssize_t len = readlink("/proc/self/exe", path.data(), path.size()); -# elif defined(__FreeBSD__) - ssize_t len = readlink("/proc/curproc/file", path.data(), path.size()); -# endif - if (len == -1) { - break; - } - if (len < (ssize_t) path.size()) { - base_path = std::string(path.data(), len); - // remove executable name - auto last_slash = base_path.find_last_of('/'); - if (last_slash != std::string::npos) { - base_path = base_path.substr(0, last_slash); - } - break; - } - path.resize(path.size() * 2); - } - - return base_path + "/"; -#elif defined(_WIN32) - std::vector path(MAX_PATH); - DWORD len = GetModuleFileNameW(NULL, path.data(), path.size()); - if (len == 0) { + std::unique_ptr path { + ggml_backend_current_exe_dir_rust(), + ggml_backend_cstring_free_rust, + }; + if (!path) { return {}; } - std::wstring base_path(path.data(), len); - // remove executable name - auto last_slash = base_path.find_last_of('\\'); - if (last_slash != std::string::npos) { - base_path = base_path.substr(0, last_slash); - } - return base_path + L"\\"; -#else - return {}; -#endif + return fs::u8path(path.get()); } static fs::path backend_filename_prefix() { -#ifdef _WIN32 - return fs::u8path("ggml-"); -#else - return fs::u8path("libggml-"); -#endif + return fs::u8path(ggml_backend_filename_prefix_rust()); } static fs::path backend_filename_extension() { -#ifdef _WIN32 - return fs::u8path(".dll"); -#else - return fs::u8path(".so"); -#endif + return fs::u8path(ggml_backend_filename_extension_rust()); } static ggml_backend_reg_t ggml_backend_load_best(const char * name, bool silent, const char * user_search_path) { diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp deleted file mode 100644 index d9f8aaec52fd..000000000000 --- a/ggml/src/ggml-backend.cpp +++ /dev/null @@ -1,2371 +0,0 @@ -// Note: porting this file to C++ is a work in progress - -#ifdef _WIN32 -#define WIN32_LEAN_AND_MEAN -#ifndef NOMINMAX -# define NOMINMAX -#endif -#include -#endif - -#include "ggml-backend.h" -#include "ggml-backend-impl.h" -#include "ggml-alloc.h" -#include "ggml-impl.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -#ifdef __APPLE__ -#include -#include -#endif - - -// backend buffer type - -const char * ggml_backend_buft_name(ggml_backend_buffer_type_t buft) { - GGML_ASSERT(buft); - return buft->iface.get_name(buft); -} - -ggml_backend_buffer_t ggml_backend_buft_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { - GGML_ASSERT(buft); - if (size == 0) { - // return a dummy buffer for zero-sized allocations - return ggml_backend_buffer_init(buft, {}, NULL, 0); - } - return buft->iface.alloc_buffer(buft, size); -} - -size_t ggml_backend_buft_get_alignment(ggml_backend_buffer_type_t buft) { - GGML_ASSERT(buft); - return buft->iface.get_alignment(buft); -} - -size_t ggml_backend_buft_get_max_size(ggml_backend_buffer_type_t buft) { - GGML_ASSERT(buft); - // get_max_size is optional, defaults to SIZE_MAX - if (buft->iface.get_max_size) { - return buft->iface.get_max_size(buft); - } - return SIZE_MAX; -} - -size_t ggml_backend_buft_get_alloc_size(ggml_backend_buffer_type_t buft, const struct ggml_tensor * tensor) { - GGML_ASSERT(buft); - // get_alloc_size is optional, defaults to ggml_nbytes - if (buft->iface.get_alloc_size) { - size_t size = buft->iface.get_alloc_size(buft, tensor); - assert(size >= ggml_nbytes(tensor)); - return size; - } - return ggml_nbytes(tensor); -} - -bool ggml_backend_buft_is_host(ggml_backend_buffer_type_t buft) { - GGML_ASSERT(buft); - if (buft->iface.is_host) { - return buft->iface.is_host(buft); - } - return false; -} - -ggml_backend_dev_t ggml_backend_buft_get_device(ggml_backend_buffer_type_t buft) { - GGML_ASSERT(buft); - return buft->device; -} - -// backend buffer - -ggml_backend_buffer_t ggml_backend_buffer_init( - ggml_backend_buffer_type_t buft, - struct ggml_backend_buffer_i iface, - void * context, - size_t size) { - ggml_backend_buffer_t buffer = new ggml_backend_buffer { - /* .interface = */ iface, - /* .buft = */ buft, - /* .context = */ context, - /* .size = */ size, - /* .usage = */ GGML_BACKEND_BUFFER_USAGE_ANY - }; - - return buffer; -} - -const char * ggml_backend_buffer_name(ggml_backend_buffer_t buffer) { - return ggml_backend_buft_name(ggml_backend_buffer_get_type(buffer)); -} - -void ggml_backend_buffer_free(ggml_backend_buffer_t buffer) { - if (buffer == NULL) { - return; - } - - if (buffer->iface.free_buffer != NULL) { - buffer->iface.free_buffer(buffer); - } - delete buffer; -} - -size_t ggml_backend_buffer_get_size(ggml_backend_buffer_t buffer) { - GGML_ASSERT(buffer); - return buffer->size; -} - -void * ggml_backend_buffer_get_base(ggml_backend_buffer_t buffer) { - GGML_ASSERT(buffer); - // get_base is optional if the buffer is zero-sized - if (!ggml_backend_buffer_is_meta(buffer) && buffer->size == 0) { - return NULL; - } - - // FIXME JG: a multi_buffer has a non-zero size, according to the above comment get_base is not optional, - // I don't know whether the above comment is correct - if (!buffer->iface.get_base) { - return NULL; - } - - void * base = buffer->iface.get_base(buffer); - - GGML_ASSERT(base != NULL && "backend buffer base cannot be NULL"); - - return base; -} - -enum ggml_status ggml_backend_buffer_init_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor) { - GGML_ASSERT(buffer); - // init_tensor is optional - if (buffer->iface.init_tensor) { - return buffer->iface.init_tensor(buffer, tensor); - } - return GGML_STATUS_SUCCESS; -} - -void ggml_backend_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { - GGML_ASSERT(buffer); - // clear is optional if the buffer is zero-sized - if (buffer->size == 0) { - return; - } - - buffer->iface.clear(buffer, value); -} - -size_t ggml_backend_buffer_get_alignment(ggml_backend_buffer_t buffer) { - return ggml_backend_buft_get_alignment(ggml_backend_buffer_get_type(buffer)); -} - -size_t ggml_backend_buffer_get_max_size(ggml_backend_buffer_t buffer) { - return ggml_backend_buft_get_max_size(ggml_backend_buffer_get_type(buffer)); -} - -size_t ggml_backend_buffer_get_alloc_size(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor) { - return ggml_backend_buft_get_alloc_size(ggml_backend_buffer_get_type(buffer), tensor); -} - -bool ggml_backend_buffer_is_host(ggml_backend_buffer_t buffer) { - return ggml_backend_buft_is_host(ggml_backend_buffer_get_type(buffer)); -} - -void ggml_backend_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage) { - GGML_ASSERT(buffer); - buffer->usage = usage; - - // FIXME: add a generic callback to the buffer interface - if (ggml_backend_buffer_is_multi_buffer(buffer)) { - ggml_backend_multi_buffer_set_usage(buffer, usage); - } -} - -enum ggml_backend_buffer_usage ggml_backend_buffer_get_usage(ggml_backend_buffer_t buffer) { - GGML_ASSERT(buffer); - return buffer->usage; -} - -ggml_backend_buffer_type_t ggml_backend_buffer_get_type(ggml_backend_buffer_t buffer) { - GGML_ASSERT(buffer); - return buffer->buft; -} - -void ggml_backend_buffer_reset(ggml_backend_buffer_t buffer) { - GGML_ASSERT(buffer); - if (buffer->iface.reset) { - buffer->iface.reset(buffer); - } -} - -bool ggml_backend_buffer_copy_tensor(const struct ggml_tensor * src, struct ggml_tensor * dst) { - ggml_backend_buffer_t dst_buf = dst->view_src ? dst->view_src->buffer : dst->buffer; - if (dst_buf->iface.cpy_tensor) { - return dst_buf->iface.cpy_tensor(dst_buf, src, dst); - } - return false; -} - -// backend - -ggml_guid_t ggml_backend_guid(ggml_backend_t backend) { - if (backend == NULL) { - return NULL; - } - return backend->guid; -} - -const char * ggml_backend_name(ggml_backend_t backend) { - if (backend == NULL) { - return "NULL"; - } - return backend->iface.get_name(backend); -} - -void ggml_backend_free(ggml_backend_t backend) { - if (backend == NULL) { - return; - } - - backend->iface.free(backend); -} - -ggml_backend_buffer_type_t ggml_backend_get_default_buffer_type(ggml_backend_t backend) { - GGML_ASSERT(backend); - return ggml_backend_dev_buffer_type(backend->device); -} - -ggml_backend_buffer_t ggml_backend_alloc_buffer(ggml_backend_t backend, size_t size) { - return ggml_backend_buft_alloc_buffer(ggml_backend_get_default_buffer_type(backend), size); -} - -size_t ggml_backend_get_alignment(ggml_backend_t backend) { - return ggml_backend_buft_get_alignment(ggml_backend_get_default_buffer_type(backend)); -} - -size_t ggml_backend_get_max_size(ggml_backend_t backend) { - return ggml_backend_buft_get_max_size(ggml_backend_get_default_buffer_type(backend)); -} - -void ggml_backend_tensor_set_async(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) { - GGML_ASSERT(backend); - GGML_ASSERT(tensor); - GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); - - if (backend->iface.set_tensor_async == NULL) { - ggml_backend_synchronize(backend); - ggml_backend_tensor_set(tensor, data, offset, size); - } else { - backend->iface.set_tensor_async(backend, tensor, data, offset, size); - } -} - -void ggml_backend_tensor_get_async(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) { - GGML_ASSERT(backend); - GGML_ASSERT(tensor); - GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor read out of bounds"); - - if (backend->iface.get_tensor_async == NULL) { - ggml_backend_synchronize(backend); - ggml_backend_tensor_get(tensor, data, offset, size); - } else { - backend->iface.get_tensor_async(backend, tensor, data, offset, size); - } -} - -void ggml_backend_tensor_set_2d_async(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, - size_t n_copies, size_t stride_tensor, size_t stride_data) { - GGML_ASSERT(backend); - GGML_ASSERT(tensor); - GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); - - if (n_copies <= 1 || backend->iface.set_tensor_2d_async == NULL) { - for (size_t i = 0; i < n_copies; i++) { - ggml_backend_tensor_set_async(backend, tensor, (const char *) data + i*stride_data, offset + i*stride_tensor, size); - } - return; - } - if (size == 0) { - return; - } - - GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); - GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); - backend->iface.set_tensor_2d_async(backend, tensor, data, offset, size, n_copies, stride_tensor, stride_data); -} - -void ggml_backend_tensor_get_2d_async(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, - size_t n_copies, size_t stride_tensor, size_t stride_data) { - GGML_ASSERT(backend); - GGML_ASSERT(tensor); - GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); - - if (n_copies <= 1 || backend->iface.set_tensor_2d_async == NULL) { - for (size_t i = 0; i < n_copies; i++) { - ggml_backend_tensor_get_async(backend, tensor, (char *) data + i*stride_data, offset + i*stride_tensor, size); - } - return; - } - if (size == 0) { - return; - } - - GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); - GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); - backend->iface.get_tensor_2d_async(backend, tensor, data, offset, size, n_copies, stride_tensor, stride_data); -} - -void ggml_backend_tensor_set(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) { - GGML_ASSERT(tensor); - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - GGML_ASSERT(buf != NULL && "tensor buffer not set"); - - if (size == 0) { - return; - } - - GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); - - buf->iface.set_tensor(buf, tensor, data, offset, size); -} - -void ggml_backend_tensor_get(const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) { - GGML_ASSERT(tensor); - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - GGML_ASSERT(buf != NULL && "tensor buffer not set"); - - if (size == 0) { - return; - } - - GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor read out of bounds"); - - buf->iface.get_tensor(buf, tensor, data, offset, size); -} - -void ggml_backend_tensor_set_2d(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, - size_t n_copies, size_t stride_tensor, size_t stride_data) { - GGML_ASSERT(tensor); - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - GGML_ASSERT(buf != NULL && "tensor buffer not set"); - - if (n_copies <= 1 || buf->iface.set_tensor_2d == NULL) { - for (size_t i = 0; i < n_copies; i++) { - ggml_backend_tensor_set(tensor, (const char *) data + i*stride_data, offset + i*stride_tensor, size); - } - return; - } - if (size == 0) { - return; - } - - GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); - GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); - - buf->iface.set_tensor_2d(buf, tensor, data, offset, size, n_copies, stride_tensor, stride_data); -} - -void ggml_backend_tensor_get_2d(const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, - size_t n_copies, size_t stride_tensor, size_t stride_data) { - GGML_ASSERT(tensor); - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - GGML_ASSERT(buf != NULL && "tensor buffer not set"); - - if (n_copies <= 1 || buf->iface.set_tensor_2d == NULL) { - for (size_t i = 0; i < n_copies; i++) { - ggml_backend_tensor_get(tensor, (char *) data + i*stride_data, offset + i*stride_tensor, size); - } - return; - } - if (size == 0) { - return; - } - - GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); - GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor read out of bounds"); - - buf->iface.get_tensor_2d(buf, tensor, data, offset, size, n_copies, stride_tensor, stride_data); -} - -void ggml_backend_tensor_memset(struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { - GGML_ASSERT(tensor); - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - - if (size == 0) { - return; - } - - GGML_ASSERT(buf != NULL && "tensor buffer not set"); - GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); - GGML_ASSERT(buf->iface.memset_tensor != NULL && "memset not implemented by backend buffer"); - - buf->iface.memset_tensor(buf, tensor, value, offset, size); -} - -void ggml_backend_synchronize(ggml_backend_t backend) { - GGML_ASSERT(backend); - if (backend->iface.synchronize == NULL) { - return; - } - - backend->iface.synchronize(backend); -} - -ggml_backend_graph_plan_t ggml_backend_graph_plan_create(ggml_backend_t backend, struct ggml_cgraph * cgraph) { - GGML_ASSERT(backend); - GGML_ASSERT(backend->iface.graph_plan_create != NULL); - - return backend->iface.graph_plan_create(backend, cgraph); -} - -void ggml_backend_graph_plan_free(ggml_backend_t backend, ggml_backend_graph_plan_t plan) { - GGML_ASSERT(backend); - GGML_ASSERT(backend->iface.graph_plan_free != NULL); - - backend->iface.graph_plan_free(backend, plan); -} - -enum ggml_status ggml_backend_graph_plan_compute(ggml_backend_t backend, ggml_backend_graph_plan_t plan) { - GGML_ASSERT(backend); - GGML_ASSERT(backend->iface.graph_plan_compute != NULL); - - return backend->iface.graph_plan_compute(backend, plan); -} - -enum ggml_status ggml_backend_graph_compute(ggml_backend_t backend, struct ggml_cgraph * cgraph) { - enum ggml_status err = ggml_backend_graph_compute_async(backend, cgraph); - ggml_backend_synchronize(backend); - return err; -} - -enum ggml_status ggml_backend_graph_compute_async(ggml_backend_t backend, struct ggml_cgraph * cgraph) { - GGML_ASSERT(backend); - return backend->iface.graph_compute(backend, cgraph); -} - -bool ggml_backend_supports_op(ggml_backend_t backend, const struct ggml_tensor * op) { - GGML_ASSERT(backend); - return ggml_backend_dev_supports_op(backend->device, op); -} - -bool ggml_backend_supports_buft(ggml_backend_t backend, ggml_backend_buffer_type_t buft) { - GGML_ASSERT(backend); - return ggml_backend_dev_supports_buft(backend->device, buft); -} - -bool ggml_backend_offload_op(ggml_backend_t backend, const struct ggml_tensor * op) { - GGML_ASSERT(backend); - return ggml_backend_dev_offload_op(backend->device, op); -} - -ggml_backend_dev_t ggml_backend_get_device(ggml_backend_t backend) { - GGML_ASSERT(backend); - return backend->device; -} - -// backend copy - -void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor * dst) { - GGML_ASSERT(ggml_are_same_layout(src, dst) && "cannot copy tensors with different layouts"); - - if (src == dst) { - return; - } - - if (ggml_backend_buffer_is_host(src->buffer)) { - ggml_backend_tensor_set(dst, src->data, 0, ggml_nbytes(src)); - } else if (ggml_backend_buffer_is_host(dst->buffer)) { - ggml_backend_tensor_get(src, dst->data, 0, ggml_nbytes(src)); - } else if (!ggml_backend_buffer_copy_tensor(src, dst)) { -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: warning: slow copy from %s to %s\n", __func__, ggml_backend_buffer_name(src->buffer), ggml_backend_buffer_name(dst->buffer)); -#endif // NDEBUG - size_t nbytes = ggml_nbytes(src); - void * data = malloc(nbytes); - ggml_backend_tensor_get(src, data, 0, nbytes); - ggml_backend_tensor_set(dst, data, 0, nbytes); - free(data); - } -} - -void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst) { - GGML_ASSERT(ggml_are_same_layout(src, dst) && "cannot copy tensors with different layouts"); - - if (src == dst) { - return; - } - - GGML_ASSERT(backend_dst); - if (backend_dst->iface.cpy_tensor_async != NULL) { - if (backend_dst->iface.cpy_tensor_async(backend_src, backend_dst, src, dst)) { - return; - } - } - - // an async copy would normally happen after all the queued operations on both backends are completed - // to simulate the same behavior, we need to synchronize both backends first, and do a blocking copy - ggml_backend_synchronize(backend_src); - ggml_backend_synchronize(backend_dst); - ggml_backend_tensor_copy(src, dst); -} - -// events - -ggml_backend_event_t ggml_backend_event_new(ggml_backend_dev_t device) { - // null device is allowed for the transition period to the device interface - if (device == NULL || device->iface.event_new == NULL) { - return NULL; - } - return device->iface.event_new(device); -} - -void ggml_backend_event_free(ggml_backend_event_t event) { - if (event == NULL) { - return; - } - event->device->iface.event_free(event->device, event); -} - -void ggml_backend_event_record(ggml_backend_event_t event, ggml_backend_t backend) { - GGML_ASSERT(backend); - GGML_ASSERT(backend->iface.event_record != NULL); - - backend->iface.event_record(backend, event); -} - -void ggml_backend_event_synchronize(ggml_backend_event_t event) { - GGML_ASSERT(event); - GGML_ASSERT(event->device->iface.event_synchronize); - - event->device->iface.event_synchronize(event->device, event); -} - -void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { - GGML_ASSERT(backend); - GGML_ASSERT(backend->iface.event_wait != NULL); - - backend->iface.event_wait(backend, event); -} - -static void ggml_backend_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * cgraph) { - GGML_ASSERT(backend); - if (backend->iface.graph_optimize != NULL) { - backend->iface.graph_optimize(backend, cgraph); - } -} - -// Backend device - -const char * ggml_backend_dev_name(ggml_backend_dev_t device) { - GGML_ASSERT(device); - return device->iface.get_name(device); -} - -const char * ggml_backend_dev_description(ggml_backend_dev_t device) { - GGML_ASSERT(device); - return device->iface.get_description(device); -} - -void ggml_backend_dev_memory(ggml_backend_dev_t device, size_t * free, size_t * total) { - GGML_ASSERT(device); - device->iface.get_memory(device, free, total); -} - -enum ggml_backend_dev_type ggml_backend_dev_type(ggml_backend_dev_t device) { - GGML_ASSERT(device); - return device->iface.get_type(device); -} - -void ggml_backend_dev_get_props(ggml_backend_dev_t device, struct ggml_backend_dev_props * props) { - GGML_ASSERT(device); - memset(props, 0, sizeof(*props)); - device->iface.get_props(device, props); -} - -ggml_backend_reg_t ggml_backend_dev_backend_reg(ggml_backend_dev_t device) { - GGML_ASSERT(device); - return device->reg; -} - -ggml_backend_t ggml_backend_dev_init(ggml_backend_dev_t device, const char * params) { - GGML_ASSERT(device); - return device->iface.init_backend(device, params); -} - -ggml_backend_buffer_type_t ggml_backend_dev_buffer_type(ggml_backend_dev_t device) { - GGML_ASSERT(device); - return device->iface.get_buffer_type(device); -} - -ggml_backend_buffer_type_t ggml_backend_dev_host_buffer_type(ggml_backend_dev_t device) { - GGML_ASSERT(device); - if (device->iface.get_host_buffer_type == NULL) { - return NULL; - } - - return device->iface.get_host_buffer_type(device); -} - -ggml_backend_buffer_t ggml_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, void * ptr, size_t size, size_t max_tensor_size) { - GGML_ASSERT(device); - return device->iface.buffer_from_host_ptr(device, ptr, size, max_tensor_size); -} - -bool ggml_backend_dev_supports_op(ggml_backend_dev_t device, const struct ggml_tensor * op) { - GGML_ASSERT(device); - return device->iface.supports_op(device, op); -} - -bool ggml_backend_dev_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft) { - GGML_ASSERT(device); - return device->iface.supports_buft(device, buft); -} - -bool ggml_backend_dev_offload_op(ggml_backend_dev_t device, const struct ggml_tensor * op) { - GGML_ASSERT(device); - if (device->iface.offload_op != NULL) { - return device->iface.offload_op(device, op); - } - - return false; -} - -// Backend (reg) - -const char * ggml_backend_reg_name(ggml_backend_reg_t reg) { - GGML_ASSERT(reg); - return reg->iface.get_name(reg); -} - -size_t ggml_backend_reg_dev_count(ggml_backend_reg_t reg) { - GGML_ASSERT(reg); - return reg->iface.get_device_count(reg); -} - -ggml_backend_dev_t ggml_backend_reg_dev_get(ggml_backend_reg_t reg, size_t index) { - GGML_ASSERT(reg); - return reg->iface.get_device(reg, index); -} - -void * ggml_backend_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { - GGML_ASSERT(reg); - if (!reg->iface.get_proc_address) { - return NULL; - } - return reg->iface.get_proc_address(reg, name); -} - -// multi-buffer buffer - -struct ggml_backend_multi_buffer_context { - ggml_backend_buffer_t * buffers; - size_t n_buffers; -}; - -static void ggml_backend_multi_buffer_free_buffer(ggml_backend_buffer_t buffer) { - GGML_ASSERT(buffer); - ggml_backend_multi_buffer_context * ctx = (ggml_backend_multi_buffer_context *) buffer->context; - for (size_t i = 0; i < ctx->n_buffers; i++) { - ggml_backend_buffer_free(ctx->buffers[i]); - } - - free(ctx->buffers); - free(ctx); -} - -static void ggml_backend_multi_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { - GGML_ASSERT(buffer); - ggml_backend_multi_buffer_context * ctx = (ggml_backend_multi_buffer_context *) buffer->context; - for (size_t i = 0; i < ctx->n_buffers; i++) { - ggml_backend_buffer_clear(ctx->buffers[i], value); - } -} - -static const struct ggml_backend_buffer_i ggml_backend_multi_buffer_i = { - /* .free_buffer = */ ggml_backend_multi_buffer_free_buffer, - /* .get_base = */ NULL, - /* .init_tensor = */ NULL, - /* .memset_tensor = */ NULL, - /* .set_tensor = */ NULL, - /* .get_tensor = */ NULL, - /* .set_tensor_2d = */ NULL, - /* .get_tensor_2d = */ NULL, - /* .cpy_tensor = */ NULL, - /* .clear = */ ggml_backend_multi_buffer_clear, - /* .reset = */ NULL, -}; - -ggml_backend_buffer_t ggml_backend_multi_buffer_alloc_buffer(ggml_backend_buffer_t * buffers, size_t n_buffers) { - ggml_backend_multi_buffer_context * ctx = (ggml_backend_multi_buffer_context *) malloc(sizeof(struct ggml_backend_multi_buffer_context)); - ctx->n_buffers = n_buffers; - ctx->buffers = (ggml_backend_buffer_t *) malloc(n_buffers * sizeof(ggml_backend_buffer_t)); - - GGML_ASSERT(ctx->buffers != NULL); - - size_t total_size = 0; - for (size_t i = 0; i < n_buffers; i++) { - ctx->buffers[i] = buffers[i]; - total_size += ggml_backend_buffer_get_size(buffers[i]); - } - - return ggml_backend_buffer_init(buffers[0]->buft, ggml_backend_multi_buffer_i, ctx, total_size); -} - -bool ggml_backend_buffer_is_multi_buffer(ggml_backend_buffer_t buffer) { - GGML_ASSERT(buffer); - return buffer->iface.free_buffer == ggml_backend_multi_buffer_free_buffer; -} - -void ggml_backend_multi_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage) { - GGML_ASSERT(buffer); - GGML_ASSERT(ggml_backend_buffer_is_multi_buffer(buffer)); - ggml_backend_multi_buffer_context * ctx = (ggml_backend_multi_buffer_context *) buffer->context; - for (size_t i = 0; i < ctx->n_buffers; i++) { - ggml_backend_buffer_set_usage(ctx->buffers[i], usage); - } -} - -// creates a copy of the tensor with the same memory layout -static struct ggml_tensor * ggml_dup_tensor_layout(struct ggml_context * ctx, const struct ggml_tensor * tensor) { - struct ggml_tensor * dup = ggml_dup_tensor(ctx, tensor); - for (int i = 0; i < GGML_MAX_DIMS; i++) { - dup->nb[i] = tensor->nb[i]; - } - return dup; -} - -static bool ggml_is_view_op(enum ggml_op op) { - return op == GGML_OP_VIEW || op == GGML_OP_RESHAPE || op == GGML_OP_PERMUTE || op == GGML_OP_TRANSPOSE; -} - -// scheduler - -#ifndef GGML_SCHED_MAX_BACKENDS -#define GGML_SCHED_MAX_BACKENDS 16 -#endif - -#ifndef GGML_SCHED_MAX_SPLIT_INPUTS -#define GGML_SCHED_MAX_SPLIT_INPUTS 30 -#endif - -#ifndef GGML_SCHED_MAX_COPIES -#define GGML_SCHED_MAX_COPIES 4 -#endif - -struct ggml_backend_sched_split { - int backend_id; - int i_start; - int i_end; - struct ggml_tensor * inputs[GGML_SCHED_MAX_SPLIT_INPUTS]; - int n_inputs; - // graph view of this split - struct ggml_cgraph graph; -}; - -struct ggml_backend_sched { - bool is_reset; // true if the scheduler has been reset since the last graph split - bool is_alloc; - - int n_backends; - - ggml_backend_t backends[GGML_SCHED_MAX_BACKENDS]; - ggml_backend_buffer_type_t bufts[GGML_SCHED_MAX_BACKENDS]; - ggml_gallocr_t galloc; - - // hash map of the nodes in the graph - struct ggml_hash_set hash_set; - int * hv_tensor_backend_ids; // [hash_set.size] - struct ggml_tensor ** hv_tensor_copies; // [hash_set.size][n_backends][n_copies] - - int * node_backend_ids; // [graph_size] - int * leaf_backend_ids; // [graph_size] - - int * prev_node_backend_ids; // [graph_size] - int * prev_leaf_backend_ids; // [graph_size] - - // copy of the graph with modified inputs - struct ggml_cgraph graph; - - // graph splits - struct ggml_backend_sched_split * splits; - int n_splits; - int splits_capacity; - - // pipeline parallelism support - int n_copies; - int cur_copy; - int next_copy; - ggml_backend_event_t events[GGML_SCHED_MAX_BACKENDS][GGML_SCHED_MAX_COPIES]; - struct ggml_tensor * graph_inputs[GGML_SCHED_MAX_SPLIT_INPUTS]; - int n_graph_inputs; - - struct ggml_context * ctx; - - ggml_backend_sched_eval_callback callback_eval; - void * callback_eval_user_data; - - char * context_buffer; - size_t context_buffer_size; - - bool op_offload; - - int debug; - - // used for debugging graph reallocations [GGML_SCHED_DEBUG_REALLOC] - // ref: https://github.com/ggml-org/llama.cpp/pull/17617 - int debug_realloc; - int debug_graph_size; - int debug_prev_graph_size; -}; - -#define hash_id(tensor) ggml_hash_find_or_insert(&sched->hash_set, tensor) -#define tensor_backend_id(tensor) sched->hv_tensor_backend_ids[hash_id(tensor)] -#define tensor_id_copy(id, backend_id, copy_id) sched->hv_tensor_copies[(id) * sched->n_backends * sched->n_copies + (backend_id) * sched->n_copies + (copy_id)] -#define tensor_copy(tensor, backend_id, copy_id) tensor_id_copy(hash_id(tensor), backend_id, copy_id) - -// returns the priority of the backend, lower id is higher priority -static int ggml_backend_sched_backend_id(ggml_backend_sched_t sched, ggml_backend_t backend) { - for (int i = 0; i < sched->n_backends; i++) { - if (sched->backends[i] == backend) { - return i; - } - } - return -1; -} - -static int ggml_backend_sched_backend_from_buffer(ggml_backend_sched_t sched, const struct ggml_tensor * tensor, const struct ggml_tensor * op) { - ggml_backend_buffer_t buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - if (buffer == NULL) { - return -1; - } - - // find highest prio backend that supports the buffer type and the op - for (int i = 0; i < sched->n_backends; i++) { - if (ggml_backend_supports_buft(sched->backends[i], buffer->buft) && - ggml_backend_supports_op(sched->backends[i], op)) { - return i; - } - } - -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: warning: no backend supports op %s with a weight with buffer type %s used in tensor %s, the weight will need to be copied\n", - __func__, ggml_op_desc(tensor), ggml_backend_buffer_name(buffer), tensor->name); -#endif - - return -1; -} - -#if 0 -#define GGML_SCHED_MAX_SPLITS_DEBUG 4096 -static char causes[GGML_DEFAULT_GRAPH_SIZE*16 + GGML_SCHED_MAX_SPLITS_DEBUG*GGML_SCHED_MAX_SPLIT_INPUTS][128]; // debug only -#define SET_CAUSE(node, ...) sprintf(causes[hash_id(node)], __VA_ARGS__) -#define GET_CAUSE(node) causes[hash_id(node)] -#else -#define SET_CAUSE(node, ...) -#define GET_CAUSE(node) "" -#endif - -// returns the backend that should be used for the node based on the current locations -static int ggml_backend_sched_backend_id_from_cur(ggml_backend_sched_t sched, struct ggml_tensor * tensor) { - // assign pre-allocated nodes to their backend - int cur_backend_id = ggml_backend_sched_backend_from_buffer(sched, tensor, tensor); - if (cur_backend_id != -1) { - SET_CAUSE(tensor, "1.dst"); - return cur_backend_id; - } - - // view_src - if (tensor->view_src != NULL) { - cur_backend_id = ggml_backend_sched_backend_from_buffer(sched, tensor->view_src, tensor); - if (cur_backend_id != -1) { - SET_CAUSE(tensor, "1.vsrc"); - return cur_backend_id; - } - } - - if (tensor->buffer || (tensor->view_src && tensor->view_src->buffer)) { - // since the tensor is pre-allocated, it cannot be moved to another backend - ggml_backend_buffer_t buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - GGML_ABORT("pre-allocated tensor (%s) in a buffer (%s) that cannot run the operation (%s)", tensor->name, ggml_backend_buffer_name(buffer), ggml_op_name(tensor->op)); - } - - // graph input - if (tensor->flags & GGML_TENSOR_FLAG_INPUT) { - cur_backend_id = sched->n_backends - 1; // last backend (assumed CPU) - SET_CAUSE(tensor, "1.inp"); - return cur_backend_id; - } - - // operations with weights are preferably run on the same backend as the weights - for (int i = 0; i < GGML_MAX_SRC; i++) { - const struct ggml_tensor * src = tensor->src[i]; - if (src == NULL) { - continue; - } - // skip ROPE since the rope freqs tensor is too small to choose a backend based on it - // not an ideal solution - if (tensor->op != GGML_OP_ROPE && src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { - int src_backend_id = ggml_backend_sched_backend_from_buffer(sched, src, tensor); - // check if a backend with higher prio wants to offload the op - if (sched->op_offload && src_backend_id == sched->n_backends - 1 && ggml_backend_buffer_is_host(src->buffer)) { - for (int b = 0; b < src_backend_id; b++) { - if (ggml_backend_supports_op(sched->backends[b], tensor) && ggml_backend_offload_op(sched->backends[b], tensor)) { - SET_CAUSE(tensor, "1.off"); - return b; - } - } - } - SET_CAUSE(tensor, "1.wgt%d", i); - return src_backend_id; - } - } - - return -1; -} - -static char * fmt_size(size_t size) { - static char buffer[128]; - if (size >= 1024*1024) { - snprintf(buffer, sizeof(buffer), "%zuM", size/1024/1024); - } else { - snprintf(buffer, sizeof(buffer), "%zuK", size/1024); - } - return buffer; -} - -static void ggml_backend_sched_print_assignments(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { - int cur_split = 0; - for (int i = 0; i < graph->n_nodes; i++) { - if (cur_split < sched->n_splits && i == sched->splits[cur_split].i_start) { - ggml_backend_t split_backend = sched->backends[sched->splits[cur_split].backend_id]; - GGML_LOG_DEBUG("\n## SPLIT #%d: %s # %d inputs", cur_split, ggml_backend_name(split_backend), - sched->splits[cur_split].n_inputs); - for (int j = 0; j < sched->splits[cur_split].n_inputs; j++) { - if (j == 0) { - GGML_LOG_DEBUG(": "); - } - GGML_LOG_DEBUG("[%s (%5.5s)] ", sched->splits[cur_split].inputs[j]->name, - fmt_size(ggml_nbytes(sched->splits[cur_split].inputs[j]))); - } - GGML_LOG_DEBUG("\n"); - cur_split++; - } - struct ggml_tensor * node = graph->nodes[i]; - if (ggml_is_view_op(node->op)) { - continue; - } - if (sched->debug > 1) { - ggml_backend_t tensor_backend = ggml_backend_sched_get_tensor_backend(sched, node); - GGML_LOG_DEBUG("node #%3d (%10.10s): %20.20s (%5.5s) [%5.5s %8.8s] use=%d,c=%d:", i, ggml_op_name(node->op), node->name, - fmt_size(ggml_nbytes(node)), tensor_backend ? ggml_backend_name(tensor_backend) : "NULL", GET_CAUSE(node), - graph->use_counts[ggml_hash_find(&graph->visited_hash_set, node)], node->flags & GGML_TENSOR_FLAG_COMPUTE ? 1 : 0); - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * src = node->src[j]; - if (src == NULL) { - continue; - } - ggml_backend_t src_backend = ggml_backend_sched_get_tensor_backend(sched, src); - GGML_LOG_DEBUG(" %20.20s (%5.5s) [%5.5s %8.8s]", src->name, - fmt_size(ggml_nbytes(src)), src_backend ? ggml_backend_name(src_backend) : "NULL", GET_CAUSE(src)); - } - GGML_LOG_DEBUG("\n"); - } - } -} - -static bool ggml_backend_sched_buffer_supported(ggml_backend_sched_t sched, struct ggml_tensor * t, int backend_id) { - ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; - ggml_backend_buffer_type_t buft = NULL; - - if (buf) { - // the tensor is already allocated - buft = buf->buft; - } else { - // see if the tensor already has a backend assigned, and use the buffer type of that backend - int tensor_backend_id = tensor_backend_id(t); - if (tensor_backend_id == -1 && t->view_src) { - tensor_backend_id = tensor_backend_id(t->view_src); - } - if (tensor_backend_id != -1) { - buft = sched->bufts[tensor_backend_id]; - } - } - - return buft != NULL && ggml_backend_supports_buft(sched->backends[backend_id], buft); -} - -static void ggml_backend_sched_set_if_supported(ggml_backend_sched_t sched, struct ggml_tensor * node, int cur_backend_id, int * node_backend_id) { - if (ggml_backend_supports_op(sched->backends[cur_backend_id], node)) { - *node_backend_id = cur_backend_id; - SET_CAUSE(node, "2.sup"); - } -} - -// assigns backends to ops and splits the graph into subgraphs that can be computed on the same backend -void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { - // reset splits - sched->n_splits = 0; - sched->n_graph_inputs = 0; - sched->is_reset = false; - - struct ggml_init_params params = { - /* .mem_size = */ sched->context_buffer_size, - /* .mem_buffer = */ sched->context_buffer, - /* .no_alloc = */ true - }; - - ggml_free(sched->ctx); - - sched->ctx = ggml_init(params); - if (sched->ctx == NULL) { - GGML_ABORT("%s: failed to initialize context\n", __func__); - } - - graph->uid = ggml_graph_next_uid(); - - // pass 1: assign backends to ops with pre-allocated inputs - for (int i = 0; i < graph->n_leafs; i++) { - struct ggml_tensor * leaf = graph->leafs[i]; - int * leaf_backend_id = &tensor_backend_id(leaf); - // do not overwrite user assignments - if (*leaf_backend_id == -1) { - *leaf_backend_id = ggml_backend_sched_backend_id_from_cur(sched, leaf); - } - } - - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - int * node_backend_id = &tensor_backend_id(node); - // do not overwrite user assignments - if (*node_backend_id == -1) { - *node_backend_id = ggml_backend_sched_backend_id_from_cur(sched, node); - -#if 0 - // src - if (node->op == GGML_OP_NONE) { - continue; - } - - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * src = node->src[j]; - if (src == NULL) { - continue; - } - int * src_backend_id = &tensor_backend_id(src); - if (*src_backend_id == -1) { - *src_backend_id = ggml_backend_sched_backend_id_from_cur(sched, src); - } - } -#endif - } - } - - // pass 2: expand current backend assignments - // assign the same backend to adjacent nodes - // expand gpu backends (i.e. non last prio) up and down, ignoring cpu (the lowest priority backend) - // thus, cpu will never be used unless weights are on cpu, or there are no gpu ops between cpu ops - // ops unsupported by the backend being expanded will be left unassigned so that they can be assigned later when the locations of its inputs are known - // expand gpu down - { - int cur_backend_id = -1; - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - if (ggml_is_view_op(node->op)) { - continue; - } - int * node_backend_id = &tensor_backend_id(node); - if (*node_backend_id != -1) { - if (*node_backend_id == sched->n_backends - 1) { - // skip cpu (lowest prio backend) - cur_backend_id = -1; - } else { - cur_backend_id = *node_backend_id; - } - } else if (cur_backend_id != -1) { - ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id); - } - } - } - // expand gpu up - { - int cur_backend_id = -1; - for (int i = graph->n_nodes - 1; i >= 0; i--) { - struct ggml_tensor * node = graph->nodes[i]; - if (ggml_is_view_op(node->op)) { - continue; - } - int * node_backend_id = &tensor_backend_id(node); - if (*node_backend_id != -1) { - if (*node_backend_id == sched->n_backends - 1) { - // skip cpu (lowest prio backend) - cur_backend_id = -1; - } else { - cur_backend_id = *node_backend_id; - } - } else if (cur_backend_id != -1) { - ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id); - } - } - } - // expand rest down - { - int cur_backend_id = -1; - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - if (ggml_is_view_op(node->op)) { - continue; - } - int * node_backend_id = &tensor_backend_id(node); - if (*node_backend_id != -1) { - cur_backend_id = *node_backend_id; - } else if (cur_backend_id != -1) { - ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id); - } - } - } - // expand rest up - { - int cur_backend_id = -1; - for (int i = graph->n_nodes - 1; i >= 0; i--) { - struct ggml_tensor * node = graph->nodes[i]; - if (ggml_is_view_op(node->op)) { - continue; - } - int * node_backend_id = &tensor_backend_id(node); - if (*node_backend_id != -1) { - cur_backend_id = *node_backend_id; - } else if (cur_backend_id != -1) { - ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id); - } - } - } - - // pass 3: upgrade nodes to higher prio backends with compatible buffer types - // if the tensor is already in the same buffer type (*) as another higher priority backend, we should move it there - // however, we also need to verify that the sources are in compatible buffer types - // (*) the actual requirement is more relaxed, the buffer type of the backend should be supported by all the users of this tensor further down the graph - // however, this is slow to verify, so we have a more strict requirement that the buffer type is the same - // this is not uncommon since multiple backends can use host memory, with the same buffer type (eg. BLAS and CPU) - // additionally, set remaining unassigned nodes to the backend with the most supported inputs - // only nodes that could not be assigned during expansion due to the backend not supporting the op should be unassigned at this point - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - if (ggml_is_view_op(node->op)) { - continue; - } - int * node_backend_id = &tensor_backend_id(node); - if (*node_backend_id == -1) { - // unassigned node: find the backend with the most supported inputs - int n_supported_best = -1; - for (int b = 0; b < sched->n_backends; b++) { - if (ggml_backend_supports_op(sched->backends[b], node)) { - int n_supported = 0; - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * src = node->src[j]; - if (src == NULL) { - continue; - } - if ((tensor_backend_id(src) != -1 || tensor_backend_id(src->view_src) != -1) && ggml_backend_sched_buffer_supported(sched, src, b)) { - n_supported++; - } - } - if (n_supported > n_supported_best) { - n_supported_best = n_supported; - *node_backend_id = b; - SET_CAUSE(node, "3.best"); - } - } - } - } else { - // assigned node: upgrade to higher prio backend if possible - for (int b = 0; b < *node_backend_id; b++) { - if (sched->bufts[b] == sched->bufts[*node_backend_id] && ggml_backend_supports_op(sched->backends[b], node)) { - bool supported = true; - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * src = node->src[j]; - if (src == NULL) { - continue; - } - if (!ggml_backend_sched_buffer_supported(sched, src, b)) { - supported = false; - break; - } - } - if (supported) { - *node_backend_id = b; - SET_CAUSE(node, "3.upg"); - break; - } - } - } - } - } - - // pass 4: assign backends to remaining src from dst and view_src - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - int * cur_backend_id = &tensor_backend_id(node); - if (node->view_src != NULL && *cur_backend_id == -1) { - *cur_backend_id = tensor_backend_id(node->view_src); - SET_CAUSE(node, "4.vsrc"); - } - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * src = node->src[j]; - if (src == NULL) { - continue; - } - int * src_backend_id = &tensor_backend_id(src); - if (*src_backend_id == -1) { - if (src->view_src != NULL) { - // views are always on the same backend as the source - *src_backend_id = tensor_backend_id(src->view_src); - SET_CAUSE(src, "4.vsrc"); - } else { - *src_backend_id = *cur_backend_id; - SET_CAUSE(src, "4.cur"); - } - } - } - // if the node is still unassigned, assign it to the first backend that supports it - for (int b = 0; b < sched->n_backends && *cur_backend_id == -1; b++) { - ggml_backend_sched_set_if_supported(sched, node, b, cur_backend_id); - } - GGML_ASSERT(*cur_backend_id != -1); - } - - // pass 5: split graph, find tensors that need to be copied - { - int i_split = 0; - struct ggml_backend_sched_split * split = &sched->splits[0]; - // find the backend of the first split, skipping view ops - int i = 0; - for (; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - if (!ggml_is_view_op(node->op)) { - split->backend_id = tensor_backend_id(node); - break; - } - } - split->i_start = 0; - split->n_inputs = 0; - int cur_backend_id = split->backend_id; - for (; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - - if (ggml_is_view_op(node->op)) { - continue; - } - - const int node_backend_id = tensor_backend_id(node); - - GGML_ASSERT(node_backend_id != -1); // all nodes should be assigned by now, this can happen if there is no CPU fallback - - // check if we should start a new split based on the sources of the current node - bool need_new_split = false; - if (node_backend_id == cur_backend_id && split->n_inputs > 0) { - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * src = node->src[j]; - if (src == NULL) { - continue; - } - // check if a weight is on a different and incompatible backend - // by starting a new split, the memory of the previously offloaded weights can be reused - if (src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { - int src_backend_id = tensor_backend_id(src); - if (src_backend_id != cur_backend_id && !ggml_backend_sched_buffer_supported(sched, src, cur_backend_id)) { - need_new_split = true; - break; - } - } - // check if the split has too many inputs - // FIXME: count the number of inputs instead of only checking when full - if (split->n_inputs == GGML_SCHED_MAX_SPLIT_INPUTS) { - const size_t id = hash_id(src); - int src_backend_id = sched->hv_tensor_backend_ids[id]; - bool supported = ggml_backend_sched_buffer_supported(sched, src, cur_backend_id); - if (src_backend_id != cur_backend_id && tensor_id_copy(id, cur_backend_id, 0) == NULL && !supported) { - need_new_split = true; - break; - } - } - } - } - - if (node_backend_id != cur_backend_id || need_new_split) { - split->i_end = i; - i_split++; - if (i_split >= sched->splits_capacity) { - sched->splits_capacity *= 2; - sched->splits = (ggml_backend_sched_split *) - realloc(sched->splits, sched->splits_capacity * sizeof(struct ggml_backend_sched_split)); - GGML_ASSERT(sched->splits != NULL); - } - split = &sched->splits[i_split]; - split->backend_id = node_backend_id; - split->i_start = i; - split->n_inputs = 0; - cur_backend_id = node_backend_id; - } - - // find inputs that are not on the same backend - for (int j = 0; j < GGML_MAX_SRC; j++) { - struct ggml_tensor * src = node->src[j]; - if (src == NULL) { - continue; - } - - size_t src_id = hash_id(src); - const int src_backend_id = sched->hv_tensor_backend_ids[src_id]; - GGML_ASSERT(src_backend_id != -1); // all inputs should be assigned by now - - if (src->flags & GGML_TENSOR_FLAG_INPUT && sched->n_copies > 1) { - if (tensor_id_copy(src_id, src_backend_id, 0) == NULL) { - ggml_backend_t backend = sched->backends[src_backend_id]; - for (int c = 0; c < sched->n_copies; c++) { - struct ggml_tensor * tensor_copy; - if (c == sched->cur_copy) { - tensor_copy = src; // use the original tensor as the current copy - } else { - tensor_copy = ggml_dup_tensor_layout(sched->ctx, src); - ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c); - } - ggml_set_input(tensor_copy); - ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor - tensor_id_copy(src_id, src_backend_id, c) = tensor_copy; - SET_CAUSE(tensor_copy, "4.cpy"); - } - int n_graph_inputs = sched->n_graph_inputs++; - GGML_ASSERT(n_graph_inputs < GGML_SCHED_MAX_SPLIT_INPUTS); - sched->graph_inputs[n_graph_inputs] = src; - } - } - - if (src_backend_id != cur_backend_id && !ggml_backend_sched_buffer_supported(sched, src, cur_backend_id)) { - // create a copy of the input in the split's backend - if (tensor_id_copy(src_id, cur_backend_id, 0) == NULL) { - ggml_backend_t backend = sched->backends[cur_backend_id]; - for (int c = 0; c < sched->n_copies; c++) { - struct ggml_tensor * tensor_copy = ggml_dup_tensor_layout(sched->ctx, src); - ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c); - if (sched->n_copies > 1) { - ggml_set_input(tensor_copy); - ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor - } - tensor_id_copy(src_id, cur_backend_id, c) = tensor_copy; - SET_CAUSE(tensor_copy, "4.cpy"); - } - int n_inputs = split->n_inputs++; - GGML_ASSERT(n_inputs < GGML_SCHED_MAX_SPLIT_INPUTS); - split->inputs[n_inputs] = src; - } - node->src[j] = tensor_id_copy(src_id, cur_backend_id, sched->cur_copy); - } - } - } - split->i_end = graph->n_nodes; - sched->n_splits = i_split + 1; - } - - if (sched->debug) { - ggml_backend_sched_print_assignments(sched, graph); - } - - // swap node_backend_ids and leaf _backend_ids with prevs - { - int * tmp = sched->node_backend_ids; - sched->node_backend_ids = sched->prev_node_backend_ids; - sched->prev_node_backend_ids = tmp; - - tmp = sched->leaf_backend_ids; - sched->leaf_backend_ids = sched->prev_leaf_backend_ids; - sched->prev_leaf_backend_ids = tmp; - } - - int graph_size = std::max(graph->n_nodes, graph->n_leafs) + sched->n_splits*GGML_SCHED_MAX_SPLIT_INPUTS*2*sched->n_copies; - - // remember the actual graph_size for performing reallocation checks later [GGML_SCHED_DEBUG_REALLOC] - sched->debug_prev_graph_size = sched->debug_graph_size; - sched->debug_graph_size = graph_size; - - if (sched->graph.size < graph_size) { - sched->graph.size = graph_size; - sched->graph.nodes = (ggml_tensor **) realloc(sched->graph.nodes, graph_size * sizeof(struct ggml_tensor *)); - sched->graph.leafs = (ggml_tensor **) realloc(sched->graph.leafs, graph_size * sizeof(struct ggml_tensor *)); - GGML_ASSERT(sched->graph.nodes != NULL); - GGML_ASSERT(sched->graph.leafs != NULL); - } - sched->graph.n_nodes = 0; - sched->graph.n_leafs = 0; - - struct ggml_cgraph * graph_copy = &sched->graph; - - for (int i = 0; i < sched->n_splits; i++) { - struct ggml_backend_sched_split * split = &sched->splits[i]; - split->graph = ggml_graph_view(graph, split->i_start, split->i_end); - - // Optimize this split of the graph. This needs to happen before we make graph_copy, - // so they are in sync. - ggml_backend_graph_optimize(sched->backends[split->backend_id], &split->graph); - - // add inputs to the graph copy so that they are allocated by ggml-alloc at the start of the split - for (int j = 0; j < split->n_inputs; j++) { - assert(graph_copy->size > (graph_copy->n_nodes + 1)); - - struct ggml_tensor * input = split->inputs[j]; - const size_t input_id = hash_id(input); - struct ggml_tensor * input_cpy = tensor_id_copy(input_id, split->backend_id, sched->cur_copy); - - // add a dependency to the input source so that it is not freed before the copy is done - struct ggml_tensor * input_dep = ggml_view_tensor(sched->ctx, input); - input_dep->src[0] = input; - sched->node_backend_ids[graph_copy->n_nodes] = sched->hv_tensor_backend_ids[input_id]; - graph_copy->nodes[graph_copy->n_nodes++] = input_dep; - - // add a dependency to the input copy so that it is allocated at the start of the split - sched->node_backend_ids[graph_copy->n_nodes] = split->backend_id; - graph_copy->nodes[graph_copy->n_nodes++] = input_cpy; - } - - for (int j = split->i_start; j < split->i_end; j++) { - assert(graph_copy->size > graph_copy->n_nodes); - sched->node_backend_ids[graph_copy->n_nodes] = tensor_backend_id(graph->nodes[j]); - graph_copy->nodes[graph_copy->n_nodes++] = graph->nodes[j]; - } - } - - if (sched->n_copies > 1) { - // add input copies as leafs so that they are allocated first - for (int i = 0; i < sched->n_graph_inputs; i++) { - struct ggml_tensor * input = sched->graph_inputs[i]; - size_t id = hash_id(input); - int backend_id = tensor_backend_id(input); - for (int c = 0; c < sched->n_copies; c++) { - struct ggml_tensor * input_cpy = tensor_id_copy(id, backend_id, c); - sched->leaf_backend_ids[graph_copy->n_leafs] = backend_id; - assert(graph_copy->size > graph_copy->n_leafs); - graph_copy->leafs[graph_copy->n_leafs++] = input_cpy; - } - } - - for (int i = 0; i < sched->n_splits; i++) { - struct ggml_backend_sched_split * split = &sched->splits[i]; - int backend_id = split->backend_id; - for (int j = 0; j < split->n_inputs; j++) { - struct ggml_tensor * input = split->inputs[j]; - size_t id = hash_id(input); - for (int c = 0; c < sched->n_copies; c++) { - struct ggml_tensor * input_cpy = tensor_id_copy(id, backend_id, c); - sched->leaf_backend_ids[graph_copy->n_leafs] = backend_id; - assert(graph_copy->size > graph_copy->n_leafs); - graph_copy->leafs[graph_copy->n_leafs++] = input_cpy; - } - } - } - } - - // add leafs from the original graph - for (int i = 0; i < graph->n_leafs; i++) { - struct ggml_tensor * leaf = graph->leafs[i]; - sched->leaf_backend_ids[graph_copy->n_leafs] = tensor_backend_id(leaf); - assert(graph_copy->size > graph_copy->n_leafs); - graph_copy->leafs[graph_copy->n_leafs++] = leaf; - } - - // set ids for all splits - for (int i = 0; i < sched->n_splits; ++i) { - sched->splits[i].graph.uid = ggml_graph_next_uid(); - } -} - -static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { - bool backend_ids_changed = false; - for (int i = 0; i < sched->graph.n_nodes; i++) { - if (sched->node_backend_ids[i] != sched->prev_node_backend_ids[i] && - sched->bufts[sched->node_backend_ids[i]] != sched->bufts[sched->prev_node_backend_ids[i]]) { - backend_ids_changed = true; - break; - } - } - if (!backend_ids_changed) { - for (int i = 0; i < sched->graph.n_leafs; i++) { - if (sched->leaf_backend_ids[i] != sched->prev_leaf_backend_ids[i] && - sched->bufts[sched->leaf_backend_ids[i]] != sched->bufts[sched->prev_leaf_backend_ids[i]]) { - backend_ids_changed = true; - break; - } - } - } - - // allocate graph - if (backend_ids_changed || !ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) { -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: failed to allocate graph, reserving (backend_ids_changed = %d)\n", __func__, backend_ids_changed); -#endif - - if (sched->debug_realloc > 0) { - // we are interested only in situations where the graph was reallocated even though its size remained the same [GGML_SCHED_DEBUG_REALLOC] - // example: https://github.com/ggml-org/llama.cpp/pull/17143 - const bool unexpected = !backend_ids_changed && sched->debug_prev_graph_size == sched->debug_graph_size; - - if (unexpected || sched->debug_realloc > 1) { - GGML_ABORT("%s: unexpected graph reallocation (graph size = %d, nodes = %d, leafs = %d), debug_realloc = %d\n", __func__, - sched->debug_graph_size, sched->graph.n_nodes, sched->graph.n_leafs, sched->debug_realloc); - } - } - - // the re-allocation may cause the split inputs to be moved to a different address - // synchronize without ggml_backend_sched_synchronize to avoid changing cur_copy - for (int i = 0; i < sched->n_backends; i++) { - ggml_backend_synchronize(sched->backends[i]); - } - - ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids); - if (!ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) { - GGML_LOG_ERROR("%s: failed to allocate graph\n", __func__); - return false; - } - } - - return true; -} - -static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) { - GGML_ASSERT(sched); - struct ggml_backend_sched_split * splits = sched->splits; - - ggml_tensor * prev_ids_tensor = nullptr; - std::vector ids; - std::vector used_ids; - - for (int split_id = 0; split_id < sched->n_splits; split_id++) { - struct ggml_backend_sched_split * split = &splits[split_id]; - int split_backend_id = split->backend_id; - ggml_backend_t split_backend = sched->backends[split_backend_id]; - - // copy the input tensors to the split backend - for (int input_id = 0; input_id < split->n_inputs; input_id++) { - ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]); - struct ggml_tensor * input = split->inputs[input_id]; - struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy); - - if (input->flags & GGML_TENSOR_FLAG_INPUT) { - // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done - if (sched->events[split_backend_id][sched->cur_copy] != NULL) { - ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); - } else { - ggml_backend_synchronize(split_backend); - } - ggml_backend_tensor_copy(input, input_cpy); - } else { - // wait for the split backend to finish using the input before overwriting it - if (sched->events[split_backend_id][sched->cur_copy] != NULL) { - ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]); - } else { - ggml_backend_synchronize(split_backend); - } - - // when offloading MoE weights, we can reduce the amount of data copied by copying only the experts that are used - ggml_tensor * node = split->graph.nodes[0]; - if (split->graph.n_nodes > 0 && - ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && - ggml_backend_buffer_is_host(input->buffer) && ( - (node->src[0] == input_cpy && node->op == GGML_OP_MUL_MAT_ID) - //|| (node->src[1] == input_cpy && node->op == GGML_OP_ADD_ID) /* GGML_OP_ADD_ID weights are small and not worth splitting */ - )) { - - const int64_t n_expert = node->op == GGML_OP_MUL_MAT_ID ? input->ne[2] : input->ne[1]; - const size_t expert_size = node->op == GGML_OP_MUL_MAT_ID ? input->nb[2] : input->nb[1]; - - ggml_backend_synchronize(input_backend); - - // get the ids - ggml_tensor * ids_tensor = node->src[2]; - ggml_backend_t ids_backend = split_backend; - - // if the ids tensor is also an input of the split, it may not have been copied yet to the split backend - // in that case, we use the original ids tensor - for (int i = input_id + 1; i < split->n_inputs; i++) { - if (ids_tensor == tensor_copy(split->inputs[i], split_backend_id, sched->cur_copy)) { - ids_tensor = split->inputs[i]; - ids_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[i]); - break; - } - } - - if (ids_tensor != prev_ids_tensor) { - ids.resize(ggml_nbytes(ids_tensor) / sizeof(int32_t)); - ggml_backend_tensor_get_async(ids_backend, ids_tensor, ids.data(), 0, ggml_nbytes(ids_tensor)); - ggml_backend_synchronize(ids_backend); - - // find the used experts - used_ids.clear(); - used_ids.resize(ggml_bitset_size(n_expert)); - for (int64_t i1 = 0; i1 < ids_tensor->ne[1]; i1++) { - for (int64_t i0 = 0; i0 < ids_tensor->ne[0]; i0++) { - int32_t id = ids[i1 * ids_tensor->nb[1]/sizeof(int32_t) + i0 * ids_tensor->nb[0]/sizeof(int32_t)]; - GGML_ASSERT(id >= 0 && id < n_expert); - ggml_bitset_set(used_ids.data(), id); - } - } - - prev_ids_tensor = ids_tensor; - } - - // group consecutive experts and copy them together - auto copy_experts = [&](int32_t first_id, int32_t last_id) { - const size_t expert_offset = first_id * expert_size; - const size_t expert_size_copy = (last_id - first_id + 1) * expert_size; - const size_t padding = std::min(expert_size, 512); - const size_t padding_end = last_id < n_expert - 1 ? padding : 0; - - ggml_backend_tensor_set_async(split_backend, - input_cpy, - (const uint8_t *)input->data + expert_offset, expert_offset, - // copy a bit extra at the to ensure there are no NaNs in the padding of the last expert - // this is necessary for MMQ in the CUDA backend - expert_size_copy + padding_end); - }; - - int id = 0; - while (!ggml_bitset_get(used_ids.data(), id)) { - id++; - } - int32_t first_id = id; - int32_t last_id = first_id; - - for (++id; id < n_expert; ++id) { - if (!ggml_bitset_get(used_ids.data(), id)) { - continue; - } - - if (id == last_id + 1) { - last_id = id; - continue; - } - - copy_experts(first_id, last_id); - - first_id = id; - last_id = id; - } - copy_experts(first_id, last_id); - } else { - // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events - // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface - if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) { - ggml_backend_synchronize(input_backend); - if (sched->events[split_backend_id][sched->cur_copy] != NULL) { - ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); - } else { - ggml_backend_synchronize(split_backend); - } - ggml_backend_tensor_copy(input, input_cpy); - } - } - } - } - - if (!sched->callback_eval) { - enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); - if (ec != GGML_STATUS_SUCCESS) { - return ec; - } - } else { - // similar to ggml_backend_compare_graph_backend - for (int j0 = 0; j0 < split->graph.n_nodes; j0++) { - struct ggml_tensor * t = split->graph.nodes[j0]; - - // check if the user needs data from this node - bool need = sched->callback_eval(t, true, sched->callback_eval_user_data); - - int j1 = j0; - - // determine the range [j0, j1] of nodes that can be computed together - while (!need && j1 < split->graph.n_nodes - 1) { - t = split->graph.nodes[++j1]; - need = sched->callback_eval(t, true, sched->callback_eval_user_data); - } - - struct ggml_cgraph gv = ggml_graph_view(&split->graph, j0, j1 + 1); - - enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &gv); - if (ec != GGML_STATUS_SUCCESS) { - return ec; - } - - // TODO: pass backend to the callback, then the user can decide if they want to synchronize - ggml_backend_synchronize(split_backend); - - if (need && !sched->callback_eval(t, false, sched->callback_eval_user_data)) { - break; - } - - j0 = j1; - } - } - - // record the event of this copy - if (split->n_inputs > 0) { - if (sched->events[split_backend_id][sched->cur_copy] != NULL) { - ggml_backend_event_record(sched->events[split_backend_id][sched->cur_copy], split_backend); - } - } - } - - return GGML_STATUS_SUCCESS; -} - -ggml_backend_sched_t ggml_backend_sched_new( - ggml_backend_t * backends, - ggml_backend_buffer_type_t * bufts, - int n_backends, - size_t graph_size, - bool parallel, - bool op_offload) { - GGML_ASSERT(n_backends > 0); - GGML_ASSERT(n_backends <= GGML_SCHED_MAX_BACKENDS); - GGML_ASSERT(ggml_backend_dev_type(ggml_backend_get_device(backends[n_backends - 1])) == GGML_BACKEND_DEVICE_TYPE_CPU); - - struct ggml_backend_sched * sched = (ggml_backend_sched *) calloc(1, sizeof(struct ggml_backend_sched)); - - const char * GGML_SCHED_DEBUG = getenv("GGML_SCHED_DEBUG"); - sched->debug = GGML_SCHED_DEBUG ? atoi(GGML_SCHED_DEBUG) : 0; - - sched->debug_realloc = 0; -#ifdef GGML_SCHED_NO_REALLOC - sched->debug_realloc = 1; -#endif - const char * GGML_SCHED_DEBUG_REALLOC = getenv("GGML_SCHED_DEBUG_REALLOC"); - sched->debug_realloc = GGML_SCHED_DEBUG_REALLOC ? atoi(GGML_SCHED_DEBUG_REALLOC) : sched->debug_realloc; - - sched->n_backends = n_backends; - sched->n_copies = parallel ? GGML_SCHED_MAX_COPIES : 1; - - // initialize hash table - // FIXME: needs to be size*2 to account for leafs (do it in graph_split instead) - sched->hash_set = ggml_hash_set_new(graph_size); - sched->hv_tensor_backend_ids = (int *) malloc(sched->hash_set.size * sizeof(sched->hv_tensor_backend_ids[0])); - sched->hv_tensor_copies = (ggml_tensor **) malloc(sched->hash_set.size * sched->n_backends * sched->n_copies * sizeof(struct ggml_tensor *)); - - const size_t ggml_sched_max_splits = graph_size; // at most there is one split for each node in the graph - const size_t nodes_size = graph_size + ggml_sched_max_splits*GGML_SCHED_MAX_SPLIT_INPUTS*2; - sched->node_backend_ids = (int *) calloc(nodes_size, sizeof(sched->node_backend_ids[0])); - sched->leaf_backend_ids = (int *) calloc(nodes_size, sizeof(sched->leaf_backend_ids[0])); - sched->prev_node_backend_ids = (int *) calloc(nodes_size, sizeof(sched->prev_node_backend_ids[0])); - sched->prev_leaf_backend_ids = (int *) calloc(nodes_size, sizeof(sched->prev_leaf_backend_ids[0])); - - sched->debug_graph_size = 0; - sched->debug_prev_graph_size = 0; - - sched->context_buffer_size = ggml_sched_max_splits*GGML_SCHED_MAX_SPLIT_INPUTS*2*sizeof(struct ggml_tensor) + ggml_graph_overhead_custom(graph_size, false); - sched->context_buffer = (char *) malloc(sched->context_buffer_size); - - const int initial_splits_capacity = 16; - sched->splits = (ggml_backend_sched_split *) calloc(initial_splits_capacity, sizeof(sched->splits[0])); - sched->splits_capacity = initial_splits_capacity; - - for (int b = 0; b < n_backends; b++) { - sched->backends[b] = backends[b]; - sched->bufts[b] = bufts ? bufts[b] : ggml_backend_get_default_buffer_type(backends[b]); - GGML_ASSERT(ggml_backend_supports_buft(backends[b], sched->bufts[b])); - - if (sched->n_copies > 1) { - for (int c = 0; c < sched->n_copies; c++) { - sched->events[b][c] = ggml_backend_event_new(backends[b]->device); - } - } - } - - sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends); - sched->op_offload = op_offload; - - ggml_backend_sched_reset(sched); - - return sched; -} - -void ggml_backend_sched_free(ggml_backend_sched_t sched) { - if (sched == NULL) { - return; - } - for (int b = 0; b < sched->n_backends; b++) { - for (int c = 0; c < sched->n_copies; c++) { - ggml_backend_event_free(sched->events[b][c]); - } - } - ggml_gallocr_free(sched->galloc); - ggml_free(sched->ctx); - ggml_hash_set_free(&sched->hash_set); - free(sched->splits); - free(sched->hv_tensor_backend_ids); - free(sched->hv_tensor_copies); - free(sched->node_backend_ids); - free(sched->leaf_backend_ids); - free(sched->prev_node_backend_ids); - free(sched->prev_leaf_backend_ids); - free(sched->context_buffer); - free(sched->graph.nodes); - free(sched->graph.leafs); - free(sched); -} - -void ggml_backend_sched_reset(ggml_backend_sched_t sched) { - GGML_ASSERT(sched); - // reset state for the next run - if (!sched->is_reset) { - ggml_hash_set_reset(&sched->hash_set); - memset(sched->hv_tensor_backend_ids, -1, sched->hash_set.size * sizeof(sched->hv_tensor_backend_ids[0])); - memset(sched->hv_tensor_copies, 0, sched->hash_set.size * sched->n_backends * sched->n_copies * sizeof(struct ggml_tensor *)); - sched->is_reset = true; - } - sched->is_alloc = false; -} - -void ggml_backend_sched_reserve_size(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph, size_t * sizes) { - GGML_ASSERT(sched); - GGML_ASSERT((int)sched->hash_set.size >= measure_graph->n_nodes + measure_graph->n_leafs); - GGML_ASSERT(sizes); - - ggml_backend_sched_reset(sched); - - ggml_backend_sched_synchronize(sched); - - ggml_backend_sched_split_graph(sched, measure_graph); - - ggml_gallocr_reserve_n_size(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids, sizes); -} - -bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph) { - GGML_ASSERT(sched); - GGML_ASSERT((int)sched->hash_set.size >= measure_graph->n_nodes + measure_graph->n_leafs); - - ggml_backend_sched_synchronize(sched); - - ggml_backend_sched_split_graph(sched, measure_graph); - - if (!ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) { - return false; - } - - ggml_backend_sched_reset(sched); - - return true; -} - -bool ggml_backend_sched_alloc_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { - GGML_ASSERT(sched); - GGML_ASSERT((int)sched->hash_set.size >= graph->n_nodes + graph->n_leafs); - GGML_ASSERT(!sched->is_alloc); - - sched->cur_copy = sched->next_copy; - sched->next_copy = (sched->next_copy + 1) % sched->n_copies; - - ggml_backend_sched_split_graph(sched, graph); - - if (!ggml_backend_sched_alloc_splits(sched)) { - return false; - } - - sched->is_alloc = true; - - return true; -} - -enum ggml_status ggml_backend_sched_graph_compute(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { - enum ggml_status err = ggml_backend_sched_graph_compute_async(sched, graph); - ggml_backend_sched_synchronize(sched); - return err; -} - -enum ggml_status ggml_backend_sched_graph_compute_async(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { - GGML_ASSERT(sched); - if (!sched->is_reset && !sched->is_alloc) { - ggml_backend_sched_reset(sched); - } - - if (!sched->is_alloc) { - if (!ggml_backend_sched_alloc_graph(sched, graph)) { - return GGML_STATUS_ALLOC_FAILED; - } - } - - return ggml_backend_sched_compute_splits(sched); -} - -void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) { - GGML_ASSERT(sched); - for (int i = 0; i < sched->n_backends; i++) { - ggml_backend_synchronize(sched->backends[i]); - } - if (!sched->is_alloc) { - // if the graph is not already allocated, always use copy 0 after a synchronization - // this ensures that during generation the same copy is used every time, - // which avoids changes in the graph that could cause CUDA or other graphs to be disabled - sched->next_copy = 0; - } -} - -void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data) { - GGML_ASSERT(sched); - sched->callback_eval = callback; - sched->callback_eval_user_data = user_data; -} - -int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched) { - GGML_ASSERT(sched); - return sched->n_splits; -} - -int ggml_backend_sched_get_n_copies(ggml_backend_sched_t sched) { - GGML_ASSERT(sched); - return sched->n_copies; -} - -int ggml_backend_sched_get_n_backends(ggml_backend_sched_t sched) { - GGML_ASSERT(sched); - return sched->n_backends; -} - -ggml_backend_t ggml_backend_sched_get_backend(ggml_backend_sched_t sched, int i) { - GGML_ASSERT(sched); - GGML_ASSERT(i >= 0 && i < sched->n_backends); - return sched->backends[i]; -} - -ggml_backend_buffer_type_t ggml_backend_sched_get_buffer_type(ggml_backend_sched_t sched, ggml_backend_t backend) { - GGML_ASSERT(sched); - int backend_index = ggml_backend_sched_backend_id(sched, backend); - GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); - - return sched->bufts[backend_index]; -} - -size_t ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend) { - GGML_ASSERT(sched); - int backend_index = ggml_backend_sched_backend_id(sched, backend); - GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); - - return ggml_gallocr_get_buffer_size(sched->galloc, backend_index); -} - -void ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend) { - GGML_ASSERT(sched); - int backend_index = ggml_backend_sched_backend_id(sched, backend); - GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); - tensor_backend_id(node) = backend_index; - SET_CAUSE(node, "usr"); - sched->is_reset = false; -} - -ggml_backend_t ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node) { - GGML_ASSERT(sched); - int backend_index = tensor_backend_id(node); - if (backend_index == -1) { - return NULL; - } - return sched->backends[backend_index]; -} - -// utils - -enum ggml_status ggml_backend_view_init(struct ggml_tensor * tensor) { - GGML_ASSERT(tensor); - GGML_ASSERT(tensor->buffer == NULL); - GGML_ASSERT(tensor->view_src != NULL); - GGML_ASSERT(tensor->view_src->buffer != NULL); - GGML_ASSERT(tensor->view_src->data != NULL); - - tensor->buffer = tensor->view_src->buffer; - tensor->data = (char *)tensor->view_src->data + tensor->view_offs; - return ggml_backend_buffer_init_tensor(tensor->buffer, tensor); -} - -enum ggml_status ggml_backend_tensor_alloc(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, void * addr) { - GGML_ASSERT(tensor); - GGML_ASSERT(tensor->buffer == NULL); - GGML_ASSERT(tensor->data == NULL); - GGML_ASSERT(tensor->view_src == NULL); - GGML_ASSERT(addr >= ggml_backend_buffer_get_base(buffer)); - GGML_ASSERT(ggml_backend_buffer_is_meta(buffer) || - (char *) addr + ggml_backend_buffer_get_alloc_size(buffer, tensor) <= - (char *) ggml_backend_buffer_get_base(buffer) + ggml_backend_buffer_get_size(buffer)); - - tensor->buffer = buffer; - tensor->data = addr; - return ggml_backend_buffer_init_tensor(buffer, tensor); -} - -static struct ggml_tensor * graph_copy_dup_tensor(struct ggml_hash_set hash_set, struct ggml_tensor ** node_copies, - struct ggml_context * ctx_allocated, struct ggml_context * ctx_unallocated, struct ggml_tensor * src) { - - GGML_ASSERT(src != NULL); - GGML_ASSERT(src->data && "graph must be allocated"); - - size_t id = ggml_hash_insert(&hash_set, src); - if (id == GGML_HASHSET_ALREADY_EXISTS) { - return node_copies[ggml_hash_find(&hash_set, src)]; - } - - struct ggml_tensor * dst = ggml_dup_tensor_layout(src->data && !src->view_src ? ctx_allocated : ctx_unallocated, src); - if (src->view_src != NULL) { - dst->view_src = graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, src->view_src); - dst->view_offs = src->view_offs; - } - dst->op = src->op; - dst->flags = src->flags; - memcpy(dst->op_params, src->op_params, sizeof(dst->op_params)); - ggml_set_name(dst, src->name); - - // copy src - for (int i = 0; i < GGML_MAX_SRC; i++) { - struct ggml_tensor * s = src->src[i]; - if (s == NULL) { - continue; - } - dst->src[i] = graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, s); - } - - node_copies[id] = dst; - return dst; -} - -static void graph_copy_init_tensor(struct ggml_hash_set * hash_set, struct ggml_tensor ** node_copies, bool * node_init, struct ggml_tensor * src) { - size_t id = ggml_hash_find(hash_set, src); - if (node_init[id]) { - return; - } - node_init[id] = true; - - struct ggml_tensor * dst = node_copies[id]; - if (dst->view_src != NULL) { - graph_copy_init_tensor(hash_set, node_copies, node_init, src->view_src); - enum ggml_status status = ggml_backend_view_init(dst); - GGML_ASSERT(status == GGML_STATUS_SUCCESS); - } - else { - ggml_backend_tensor_copy(src, dst); - } - - // init src - for (int i = 0; i < GGML_MAX_SRC; i++) { - struct ggml_tensor * s = src->src[i]; - if (s == NULL) { - continue; - } - graph_copy_init_tensor(hash_set, node_copies, node_init, s); - } -} - -struct ggml_backend_graph_copy ggml_backend_graph_copy(ggml_backend_t backend, struct ggml_cgraph * graph) { - GGML_ASSERT(graph); - struct ggml_hash_set hash_set = ggml_hash_set_new(graph->visited_hash_set.size); - struct ggml_tensor ** node_copies = (ggml_tensor **) calloc(hash_set.size, sizeof(node_copies[0])); // NOLINT - bool * node_init = (bool *) calloc(hash_set.size, sizeof(node_init[0])); - - struct ggml_init_params params = { - /* .mem_size = */ ggml_tensor_overhead()*hash_set.size + ggml_graph_overhead_custom(graph->size, false), - /* .mem_buffer = */ NULL, - /* .no_alloc = */ true - }; - - struct ggml_context * ctx_allocated = ggml_init(params); - struct ggml_context * ctx_unallocated = ggml_init(params); - - if (ctx_allocated == NULL || ctx_unallocated == NULL) { - GGML_LOG_ERROR("%s: failed to allocate context for graph copy\n", __func__); - ggml_hash_set_free(&hash_set); - free(node_copies); - free(node_init); - ggml_free(ctx_allocated); - ggml_free(ctx_unallocated); - return { - /* .buffer = */ NULL, - /* .ctx_allocated = */ NULL, - /* .ctx_unallocated = */ NULL, - /* .graph = */ NULL, - }; - } - - // dup nodes - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, node); - } - - // allocate nodes - ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx_allocated, backend); - if (buffer == NULL) { - GGML_LOG_ERROR("%s: failed to allocate buffer for graph copy\n", __func__); - ggml_hash_set_free(&hash_set); - free(node_copies); - free(node_init); - ggml_free(ctx_allocated); - ggml_free(ctx_unallocated); - return { - /* .buffer = */ NULL, - /* .ctx_allocated = */ NULL, - /* .ctx_unallocated = */ NULL, - /* .graph = */ NULL, - }; - } - - //printf("copy buffer size: %zu MB\n", ggml_backend_buffer_get_size(buffer) / 1024 / 1024); - - // copy data and init views - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - graph_copy_init_tensor(&hash_set, node_copies, node_init, node); - } - - // build graph copy - struct ggml_cgraph * graph_copy = ggml_new_graph_custom(ctx_allocated, graph->size, false); - for (int i = 0; i < graph->n_nodes; i++) { - struct ggml_tensor * node = graph->nodes[i]; - struct ggml_tensor * node_copy = node_copies[ggml_hash_find(&hash_set, node)]; - graph_copy->nodes[i] = node_copy; - } - graph_copy->n_nodes = graph->n_nodes; - - ggml_hash_set_free(&hash_set); - free(node_copies); - free(node_init); - - return { - /* .buffer = */ buffer, - /* .ctx_allocated = */ ctx_allocated, - /* .ctx_unallocated = */ ctx_unallocated, - /* .graph = */ graph_copy, - }; -} - -void ggml_backend_graph_copy_free(struct ggml_backend_graph_copy copy) { - ggml_backend_buffer_free(copy.buffer); - ggml_free(copy.ctx_allocated); - ggml_free(copy.ctx_unallocated); -} - -bool ggml_backend_compare_graph_backend(ggml_backend_t backend1, ggml_backend_t backend2, struct ggml_cgraph * graph, ggml_backend_eval_callback callback, void * user_data, struct ggml_tensor const * const * test_nodes, size_t num_test_nodes) { - struct ggml_backend_graph_copy copy = ggml_backend_graph_copy(backend2, graph); - if (copy.buffer == NULL) { - return false; - } - - struct ggml_cgraph * g1 = graph; - struct ggml_cgraph * g2 = copy.graph; - - assert(g1->n_nodes == g2->n_nodes); - - if (num_test_nodes != 0) { - GGML_ASSERT(test_nodes); - // Compute the whole graph and only test the output for specific tensors - ggml_backend_graph_compute(backend1, g1); - ggml_backend_graph_compute(backend2, g2); - - bool verified = false; - for (int i = 0; i < g1->n_nodes; i++) { - for (size_t j = 0; j < num_test_nodes; ++j) { - if (g1->nodes[i] == test_nodes[j]) { - callback(i, g1->nodes[i], g2->nodes[i], user_data); - verified = true; - } - } - } - GGML_ASSERT(verified); - } else { - for (int i = 0; i < g1->n_nodes; i++) { - struct ggml_tensor * t1 = g1->nodes[i]; - struct ggml_tensor * t2 = g2->nodes[i]; - - assert(t1->op == t2->op && ggml_are_same_layout(t1, t2)); - - struct ggml_cgraph g1v = ggml_graph_view(g1, i, i + 1); - struct ggml_cgraph g2v = ggml_graph_view(g2, i, i + 1); - - ggml_backend_graph_compute(backend1, &g1v); - ggml_backend_graph_compute(backend2, &g2v); - - if (ggml_is_view_op(t1->op)) { - continue; - } - - // compare results, calculate rms etc - if (!callback(i, t1, t2, user_data)) { - break; - } - } - } - ggml_backend_graph_copy_free(copy); - - return true; -} - -// CPU backend - buffer - -static void * ggml_backend_cpu_buffer_get_base(ggml_backend_buffer_t buffer) { - GGML_ASSERT(buffer); - uintptr_t data = (uintptr_t)buffer->context; - - // align the buffer - if (data % TENSOR_ALIGNMENT != 0) { - data = GGML_PAD(data, TENSOR_ALIGNMENT); - } - - return (void *)data; -} - -static void ggml_backend_cpu_buffer_free_buffer(ggml_backend_buffer_t buffer) { - GGML_ASSERT(buffer); - ggml_aligned_free(buffer->context, buffer->size); -} - -static void ggml_backend_cpu_buffer_memset_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { - GGML_ASSERT(tensor); - memset((char *)tensor->data + offset, value, size); - - GGML_UNUSED(buffer); -} - -static void ggml_backend_cpu_buffer_set_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) { - GGML_ASSERT(tensor); - memcpy((char *)tensor->data + offset, data, size); - - GGML_UNUSED(buffer); -} - -static void ggml_backend_cpu_buffer_get_tensor(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) { - GGML_ASSERT(tensor); - memcpy(data, (const char *)tensor->data + offset, size); - - GGML_UNUSED(buffer); -} - -static bool ggml_backend_cpu_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const struct ggml_tensor * src, struct ggml_tensor * dst) { - GGML_ASSERT(src); - if (ggml_backend_buffer_is_host(src->buffer)) { - memcpy(dst->data, src->data, ggml_nbytes(src)); - return true; - } - return false; - - GGML_UNUSED(buffer); -} - -static void ggml_backend_cpu_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { - GGML_ASSERT(buffer); - memset(buffer->context, value, buffer->size); -} - -static const struct ggml_backend_buffer_i ggml_backend_cpu_buffer_i = { - /* .free_buffer = */ ggml_backend_cpu_buffer_free_buffer, - /* .get_base = */ ggml_backend_cpu_buffer_get_base, - /* .init_tensor = */ NULL, // no initialization required - /* .memset_tensor = */ ggml_backend_cpu_buffer_memset_tensor, - /* .set_tensor = */ ggml_backend_cpu_buffer_set_tensor, - /* .get_tensor = */ ggml_backend_cpu_buffer_get_tensor, - /* .set_tensor_2d = */ NULL, - /* .get_tensor_2d = */ NULL, - /* .cpy_tensor = */ ggml_backend_cpu_buffer_cpy_tensor, - /* .clear = */ ggml_backend_cpu_buffer_clear, - /* .reset = */ NULL, -}; - -static const struct ggml_backend_buffer_i ggml_backend_cpu_buffer_from_ptr_i = { - /* .free_buffer = */ NULL, // ptr is not owned by the buffer, so it does not need to be freed - /* .get_base = */ ggml_backend_cpu_buffer_get_base, - /* .init_tensor = */ NULL, // no initialization required - /* .memset_tensor = */ ggml_backend_cpu_buffer_memset_tensor, - /* .set_tensor = */ ggml_backend_cpu_buffer_set_tensor, - /* .get_tensor = */ ggml_backend_cpu_buffer_get_tensor, - /* .set_tensor_2d = */ NULL, - /* .get_tensor_2d = */ NULL, - /* .cpy_tensor = */ ggml_backend_cpu_buffer_cpy_tensor, - /* .clear = */ ggml_backend_cpu_buffer_clear, - /* .reset = */ NULL, -}; - -// CPU backend buffer type - -// this buffer type is defined here to make it available to all backends - -static const char * ggml_backend_cpu_buffer_type_get_name(ggml_backend_buffer_type_t buft) { - return "CPU"; - - GGML_UNUSED(buft); -} - -static ggml_backend_buffer_t ggml_backend_cpu_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { - void * data = ggml_aligned_malloc(size); - - if (data == NULL) { - GGML_LOG_ERROR("%s: failed to allocate buffer of size %zu\n", __func__, size); - return NULL; - } - - return ggml_backend_buffer_init(buft, ggml_backend_cpu_buffer_i, data, size); -} - -static size_t ggml_backend_cpu_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { - return TENSOR_ALIGNMENT; - - GGML_UNUSED(buft); -} - -static bool ggml_backend_cpu_buffer_type_is_host(ggml_backend_buffer_type_t buft) { - return true; - - GGML_UNUSED(buft); -} - -ggml_backend_buffer_type_t ggml_backend_cpu_buffer_type(void) { - static struct ggml_backend_buffer_type ggml_backend_cpu_buffer_type = { - /* .iface = */ { - /* .get_name = */ ggml_backend_cpu_buffer_type_get_name, - /* .alloc_buffer = */ ggml_backend_cpu_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_cpu_buffer_type_get_alignment, - /* .get_max_size = */ NULL, // defaults to SIZE_MAX - /* .get_alloc_size = */ NULL, // defaults to ggml_nbytes - /* .is_host = */ ggml_backend_cpu_buffer_type_is_host, - }, - /* .device = */ NULL, // FIXME ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0), - /* .context = */ NULL, - }; - - return &ggml_backend_cpu_buffer_type; -} - -static const char * ggml_backend_cpu_buffer_from_ptr_type_get_name(ggml_backend_buffer_type_t buft) { - return "CPU_Mapped"; - - GGML_UNUSED(buft); -} - -static ggml_backend_buffer_type_t ggml_backend_cpu_buffer_from_ptr_type(void) { - static struct ggml_backend_buffer_type ggml_backend_cpu_buffer_type = { - /* .iface = */ { - /* .get_name = */ ggml_backend_cpu_buffer_from_ptr_type_get_name, - /* .alloc_buffer = */ ggml_backend_cpu_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_cpu_buffer_type_get_alignment, - /* .get_max_size = */ NULL, // defaults to SIZE_MAX - /* .get_alloc_size = */ NULL, // defaults to ggml_nbytes - /* .is_host = */ ggml_backend_cpu_buffer_type_is_host, - }, - /* .device = */ NULL, // FIXME ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0), - /* .context = */ NULL, - }; - - return &ggml_backend_cpu_buffer_type; -} - -ggml_backend_buffer_t ggml_backend_cpu_buffer_from_ptr(void * ptr, size_t size) { - GGML_ASSERT((uintptr_t)ptr % TENSOR_ALIGNMENT == 0 && "buffer pointer must be aligned"); - return ggml_backend_buffer_init(ggml_backend_cpu_buffer_from_ptr_type(), ggml_backend_cpu_buffer_from_ptr_i, ptr, size); -} diff --git a/ggml/src/ggml-backend.cpp.inc b/ggml/src/ggml-backend.cpp.inc new file mode 100644 index 000000000000..bc026627751e --- /dev/null +++ b/ggml/src/ggml-backend.cpp.inc @@ -0,0 +1,7185 @@ +// Note: porting this file to C++ is a work in progress + +#ifdef _WIN32 +#define WIN32_LEAN_AND_MEAN +#ifndef NOMINMAX +# define NOMINMAX +#endif +#include +#endif + +#include "ggml-backend.h.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml-alloc.h.inc" +#include "ggml-impl-defs.inc" + +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __APPLE__ +#include +#include +#endif + + +// backend buffer type + +const char * ggml_backend_buft_name(ggml_backend_buffer_type_t buft) { + GGML_ASSERT(buft); + return buft->iface.get_name(buft); +} + +ggml_backend_buffer_t ggml_backend_buft_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + GGML_ASSERT(buft); + if (size == 0) { + // return a dummy buffer for zero-sized allocations + return ggml_backend_buffer_init(buft, {}, NULL, 0); + } + return buft->iface.alloc_buffer(buft, size); +} + +size_t ggml_backend_buft_get_alignment(ggml_backend_buffer_type_t buft) { + GGML_ASSERT(buft); + return buft->iface.get_alignment(buft); +} + +size_t ggml_backend_buft_get_max_size(ggml_backend_buffer_type_t buft) { + GGML_ASSERT(buft); + // get_max_size is optional, defaults to SIZE_MAX + if (buft->iface.get_max_size) { + return buft->iface.get_max_size(buft); + } + return SIZE_MAX; +} + +size_t ggml_backend_buft_get_alloc_size(ggml_backend_buffer_type_t buft, const struct ggml_tensor * tensor) { + GGML_ASSERT(buft); + // get_alloc_size is optional, defaults to ggml_nbytes + if (buft->iface.get_alloc_size) { + size_t size = buft->iface.get_alloc_size(buft, tensor); + assert(size >= ggml_nbytes(tensor)); + return size; + } + return ggml_nbytes(tensor); +} + +bool ggml_backend_buft_is_host(ggml_backend_buffer_type_t buft) { + GGML_ASSERT(buft); + if (buft->iface.is_host) { + return buft->iface.is_host(buft); + } + return false; +} + +ggml_backend_dev_t ggml_backend_buft_get_device(ggml_backend_buffer_type_t buft) { + GGML_ASSERT(buft); + return buft->device; +} + +// backend buffer + +ggml_backend_buffer_t ggml_backend_buffer_init( + ggml_backend_buffer_type_t buft, + struct ggml_backend_buffer_i iface, + void * context, + size_t size) { + ggml_backend_buffer_t buffer = new ggml_backend_buffer { + /* .interface = */ iface, + /* .buft = */ buft, + /* .context = */ context, + /* .size = */ size, + /* .usage = */ GGML_BACKEND_BUFFER_USAGE_ANY + }; + + return buffer; +} + +const char * ggml_backend_buffer_name(ggml_backend_buffer_t buffer) { + return ggml_backend_buft_name(ggml_backend_buffer_get_type(buffer)); +} + +void ggml_backend_buffer_free(ggml_backend_buffer_t buffer) { + if (buffer == NULL) { + return; + } + + if (buffer->iface.free_buffer != NULL) { + buffer->iface.free_buffer(buffer); + } + delete buffer; +} + +size_t ggml_backend_buffer_get_size(ggml_backend_buffer_t buffer) { + GGML_ASSERT(buffer); + return buffer->size; +} + +void * ggml_backend_buffer_get_base(ggml_backend_buffer_t buffer) { + GGML_ASSERT(buffer); + // get_base is optional if the buffer is zero-sized + if (!ggml_backend_buffer_is_meta(buffer) && buffer->size == 0) { + return NULL; + } + + // FIXME JG: a multi_buffer has a non-zero size, according to the above comment get_base is not optional, + // I don't know whether the above comment is correct + if (!buffer->iface.get_base) { + return NULL; + } + + void * base = buffer->iface.get_base(buffer); + + GGML_ASSERT(base != NULL && "backend buffer base cannot be NULL"); + + return base; +} + +enum ggml_status ggml_backend_buffer_init_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor) { + GGML_ASSERT(buffer); + // init_tensor is optional + if (buffer->iface.init_tensor) { + return buffer->iface.init_tensor(buffer, tensor); + } + return GGML_STATUS_SUCCESS; +} + +void ggml_backend_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + GGML_ASSERT(buffer); + // clear is optional if the buffer is zero-sized + if (buffer->size == 0) { + return; + } + + buffer->iface.clear(buffer, value); +} + +size_t ggml_backend_buffer_get_alignment(ggml_backend_buffer_t buffer) { + return ggml_backend_buft_get_alignment(ggml_backend_buffer_get_type(buffer)); +} + +size_t ggml_backend_buffer_get_max_size(ggml_backend_buffer_t buffer) { + return ggml_backend_buft_get_max_size(ggml_backend_buffer_get_type(buffer)); +} + +size_t ggml_backend_buffer_get_alloc_size(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor) { + return ggml_backend_buft_get_alloc_size(ggml_backend_buffer_get_type(buffer), tensor); +} + +bool ggml_backend_buffer_is_host(ggml_backend_buffer_t buffer) { + return ggml_backend_buft_is_host(ggml_backend_buffer_get_type(buffer)); +} + +void ggml_backend_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage) { + GGML_ASSERT(buffer); + buffer->usage = usage; + + // FIXME: add a generic callback to the buffer interface + if (ggml_backend_buffer_is_multi_buffer(buffer)) { + ggml_backend_multi_buffer_set_usage(buffer, usage); + } +} + +enum ggml_backend_buffer_usage ggml_backend_buffer_get_usage(ggml_backend_buffer_t buffer) { + GGML_ASSERT(buffer); + return buffer->usage; +} + +ggml_backend_buffer_type_t ggml_backend_buffer_get_type(ggml_backend_buffer_t buffer) { + GGML_ASSERT(buffer); + return buffer->buft; +} + +void ggml_backend_buffer_reset(ggml_backend_buffer_t buffer) { + GGML_ASSERT(buffer); + if (buffer->iface.reset) { + buffer->iface.reset(buffer); + } +} + +bool ggml_backend_buffer_copy_tensor(const struct ggml_tensor * src, struct ggml_tensor * dst) { + ggml_backend_buffer_t dst_buf = dst->view_src ? dst->view_src->buffer : dst->buffer; + if (dst_buf->iface.cpy_tensor) { + return dst_buf->iface.cpy_tensor(dst_buf, src, dst); + } + return false; +} + +// backend + +ggml_guid_t ggml_backend_guid(ggml_backend_t backend) { + if (backend == NULL) { + return NULL; + } + return backend->guid; +} + +const char * ggml_backend_name(ggml_backend_t backend) { + if (backend == NULL) { + return "NULL"; + } + return backend->iface.get_name(backend); +} + +void ggml_backend_free(ggml_backend_t backend) { + if (backend == NULL) { + return; + } + + backend->iface.free(backend); +} + +ggml_backend_buffer_type_t ggml_backend_get_default_buffer_type(ggml_backend_t backend) { + GGML_ASSERT(backend); + return ggml_backend_dev_buffer_type(backend->device); +} + +ggml_backend_buffer_t ggml_backend_alloc_buffer(ggml_backend_t backend, size_t size) { + return ggml_backend_buft_alloc_buffer(ggml_backend_get_default_buffer_type(backend), size); +} + +size_t ggml_backend_get_alignment(ggml_backend_t backend) { + return ggml_backend_buft_get_alignment(ggml_backend_get_default_buffer_type(backend)); +} + +size_t ggml_backend_get_max_size(ggml_backend_t backend) { + return ggml_backend_buft_get_max_size(ggml_backend_get_default_buffer_type(backend)); +} + +void ggml_backend_tensor_set_async(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + GGML_ASSERT(backend); + GGML_ASSERT(tensor); + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); + + if (backend->iface.set_tensor_async == NULL) { + ggml_backend_synchronize(backend); + ggml_backend_tensor_set(tensor, data, offset, size); + } else { + backend->iface.set_tensor_async(backend, tensor, data, offset, size); + } +} + +void ggml_backend_tensor_get_async(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) { + GGML_ASSERT(backend); + GGML_ASSERT(tensor); + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor read out of bounds"); + + if (backend->iface.get_tensor_async == NULL) { + ggml_backend_synchronize(backend); + ggml_backend_tensor_get(tensor, data, offset, size); + } else { + backend->iface.get_tensor_async(backend, tensor, data, offset, size); + } +} + +void ggml_backend_tensor_set_2d_async(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + GGML_ASSERT(backend); + GGML_ASSERT(tensor); + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + + if (n_copies <= 1 || backend->iface.set_tensor_2d_async == NULL) { + for (size_t i = 0; i < n_copies; i++) { + ggml_backend_tensor_set_async(backend, tensor, (const char *) data + i*stride_data, offset + i*stride_tensor, size); + } + return; + } + if (size == 0) { + return; + } + + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); + backend->iface.set_tensor_2d_async(backend, tensor, data, offset, size, n_copies, stride_tensor, stride_data); +} + +void ggml_backend_tensor_get_2d_async(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + GGML_ASSERT(backend); + GGML_ASSERT(tensor); + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + + if (n_copies <= 1 || backend->iface.set_tensor_2d_async == NULL) { + for (size_t i = 0; i < n_copies; i++) { + ggml_backend_tensor_get_async(backend, tensor, (char *) data + i*stride_data, offset + i*stride_tensor, size); + } + return; + } + if (size == 0) { + return; + } + + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); + backend->iface.get_tensor_2d_async(backend, tensor, data, offset, size, n_copies, stride_tensor, stride_data); +} + +void ggml_backend_tensor_set(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + GGML_ASSERT(tensor); + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + GGML_ASSERT(buf != NULL && "tensor buffer not set"); + + if (size == 0) { + return; + } + + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); + + buf->iface.set_tensor(buf, tensor, data, offset, size); +} + +void ggml_backend_tensor_get(const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) { + GGML_ASSERT(tensor); + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + GGML_ASSERT(buf != NULL && "tensor buffer not set"); + + if (size == 0) { + return; + } + + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor read out of bounds"); + + buf->iface.get_tensor(buf, tensor, data, offset, size); +} + +void ggml_backend_tensor_set_2d(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + GGML_ASSERT(tensor); + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + GGML_ASSERT(buf != NULL && "tensor buffer not set"); + + if (n_copies <= 1 || buf->iface.set_tensor_2d == NULL) { + for (size_t i = 0; i < n_copies; i++) { + ggml_backend_tensor_set(tensor, (const char *) data + i*stride_data, offset + i*stride_tensor, size); + } + return; + } + if (size == 0) { + return; + } + + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); + + buf->iface.set_tensor_2d(buf, tensor, data, offset, size, n_copies, stride_tensor, stride_data); +} + +void ggml_backend_tensor_get_2d(const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + GGML_ASSERT(tensor); + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + GGML_ASSERT(buf != NULL && "tensor buffer not set"); + + if (n_copies <= 1 || buf->iface.set_tensor_2d == NULL) { + for (size_t i = 0; i < n_copies; i++) { + ggml_backend_tensor_get(tensor, (char *) data + i*stride_data, offset + i*stride_tensor, size); + } + return; + } + if (size == 0) { + return; + } + + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor read out of bounds"); + + buf->iface.get_tensor_2d(buf, tensor, data, offset, size, n_copies, stride_tensor, stride_data); +} + +void ggml_backend_tensor_memset(struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { + GGML_ASSERT(tensor); + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + + if (size == 0) { + return; + } + + GGML_ASSERT(buf != NULL && "tensor buffer not set"); + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); + GGML_ASSERT(buf->iface.memset_tensor != NULL && "memset not implemented by backend buffer"); + + buf->iface.memset_tensor(buf, tensor, value, offset, size); +} + +void ggml_backend_synchronize(ggml_backend_t backend) { + GGML_ASSERT(backend); + if (backend->iface.synchronize == NULL) { + return; + } + + backend->iface.synchronize(backend); +} + +ggml_backend_graph_plan_t ggml_backend_graph_plan_create(ggml_backend_t backend, struct ggml_cgraph * cgraph) { + GGML_ASSERT(backend); + GGML_ASSERT(backend->iface.graph_plan_create != NULL); + + return backend->iface.graph_plan_create(backend, cgraph); +} + +void ggml_backend_graph_plan_free(ggml_backend_t backend, ggml_backend_graph_plan_t plan) { + GGML_ASSERT(backend); + GGML_ASSERT(backend->iface.graph_plan_free != NULL); + + backend->iface.graph_plan_free(backend, plan); +} + +enum ggml_status ggml_backend_graph_plan_compute(ggml_backend_t backend, ggml_backend_graph_plan_t plan) { + GGML_ASSERT(backend); + GGML_ASSERT(backend->iface.graph_plan_compute != NULL); + + return backend->iface.graph_plan_compute(backend, plan); +} + +enum ggml_status ggml_backend_graph_compute(ggml_backend_t backend, struct ggml_cgraph * cgraph) { + enum ggml_status err = ggml_backend_graph_compute_async(backend, cgraph); + ggml_backend_synchronize(backend); + return err; +} + +enum ggml_status ggml_backend_graph_compute_async(ggml_backend_t backend, struct ggml_cgraph * cgraph) { + GGML_ASSERT(backend); + return backend->iface.graph_compute(backend, cgraph); +} + +bool ggml_backend_supports_op(ggml_backend_t backend, const struct ggml_tensor * op) { + GGML_ASSERT(backend); + return ggml_backend_dev_supports_op(backend->device, op); +} + +bool ggml_backend_supports_buft(ggml_backend_t backend, ggml_backend_buffer_type_t buft) { + GGML_ASSERT(backend); + return ggml_backend_dev_supports_buft(backend->device, buft); +} + +bool ggml_backend_offload_op(ggml_backend_t backend, const struct ggml_tensor * op) { + GGML_ASSERT(backend); + return ggml_backend_dev_offload_op(backend->device, op); +} + +ggml_backend_dev_t ggml_backend_get_device(ggml_backend_t backend) { + GGML_ASSERT(backend); + return backend->device; +} + +// backend copy + +void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor * dst) { + GGML_ASSERT(ggml_are_same_layout(src, dst) && "cannot copy tensors with different layouts"); + + if (src == dst) { + return; + } + + if (ggml_backend_buffer_is_host(src->buffer)) { + ggml_backend_tensor_set(dst, src->data, 0, ggml_nbytes(src)); + } else if (ggml_backend_buffer_is_host(dst->buffer)) { + ggml_backend_tensor_get(src, dst->data, 0, ggml_nbytes(src)); + } else if (!ggml_backend_buffer_copy_tensor(src, dst)) { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: warning: slow copy from %s to %s\n", __func__, ggml_backend_buffer_name(src->buffer), ggml_backend_buffer_name(dst->buffer)); +#endif // NDEBUG + size_t nbytes = ggml_nbytes(src); + void * data = malloc(nbytes); + ggml_backend_tensor_get(src, data, 0, nbytes); + ggml_backend_tensor_set(dst, data, 0, nbytes); + free(data); + } +} + +void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst) { + GGML_ASSERT(ggml_are_same_layout(src, dst) && "cannot copy tensors with different layouts"); + + if (src == dst) { + return; + } + + GGML_ASSERT(backend_dst); + if (backend_dst->iface.cpy_tensor_async != NULL) { + if (backend_dst->iface.cpy_tensor_async(backend_src, backend_dst, src, dst)) { + return; + } + } + + // an async copy would normally happen after all the queued operations on both backends are completed + // to simulate the same behavior, we need to synchronize both backends first, and do a blocking copy + ggml_backend_synchronize(backend_src); + ggml_backend_synchronize(backend_dst); + ggml_backend_tensor_copy(src, dst); +} + +// events + +ggml_backend_event_t ggml_backend_event_new(ggml_backend_dev_t device) { + // null device is allowed for the transition period to the device interface + if (device == NULL || device->iface.event_new == NULL) { + return NULL; + } + return device->iface.event_new(device); +} + +void ggml_backend_event_free(ggml_backend_event_t event) { + if (event == NULL) { + return; + } + event->device->iface.event_free(event->device, event); +} + +void ggml_backend_event_record(ggml_backend_event_t event, ggml_backend_t backend) { + GGML_ASSERT(backend); + GGML_ASSERT(backend->iface.event_record != NULL); + + backend->iface.event_record(backend, event); +} + +void ggml_backend_event_synchronize(ggml_backend_event_t event) { + GGML_ASSERT(event); + GGML_ASSERT(event->device->iface.event_synchronize); + + event->device->iface.event_synchronize(event->device, event); +} + +void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { + GGML_ASSERT(backend); + GGML_ASSERT(backend->iface.event_wait != NULL); + + backend->iface.event_wait(backend, event); +} + +static void ggml_backend_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * cgraph) { + GGML_ASSERT(backend); + if (backend->iface.graph_optimize != NULL) { + backend->iface.graph_optimize(backend, cgraph); + } +} + +// Backend device + +const char * ggml_backend_dev_name(ggml_backend_dev_t device) { + GGML_ASSERT(device); + return device->iface.get_name(device); +} + +const char * ggml_backend_dev_description(ggml_backend_dev_t device) { + GGML_ASSERT(device); + return device->iface.get_description(device); +} + +void ggml_backend_dev_memory(ggml_backend_dev_t device, size_t * free, size_t * total) { + GGML_ASSERT(device); + device->iface.get_memory(device, free, total); +} + +enum ggml_backend_dev_type ggml_backend_dev_type(ggml_backend_dev_t device) { + GGML_ASSERT(device); + return device->iface.get_type(device); +} + +void ggml_backend_dev_get_props(ggml_backend_dev_t device, struct ggml_backend_dev_props * props) { + GGML_ASSERT(device); + memset(props, 0, sizeof(*props)); + device->iface.get_props(device, props); +} + +ggml_backend_reg_t ggml_backend_dev_backend_reg(ggml_backend_dev_t device) { + GGML_ASSERT(device); + return device->reg; +} + +ggml_backend_t ggml_backend_dev_init(ggml_backend_dev_t device, const char * params) { + GGML_ASSERT(device); + return device->iface.init_backend(device, params); +} + +ggml_backend_buffer_type_t ggml_backend_dev_buffer_type(ggml_backend_dev_t device) { + GGML_ASSERT(device); + return device->iface.get_buffer_type(device); +} + +ggml_backend_buffer_type_t ggml_backend_dev_host_buffer_type(ggml_backend_dev_t device) { + GGML_ASSERT(device); + if (device->iface.get_host_buffer_type == NULL) { + return NULL; + } + + return device->iface.get_host_buffer_type(device); +} + +ggml_backend_buffer_t ggml_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, void * ptr, size_t size, size_t max_tensor_size) { + GGML_ASSERT(device); + return device->iface.buffer_from_host_ptr(device, ptr, size, max_tensor_size); +} + +bool ggml_backend_dev_supports_op(ggml_backend_dev_t device, const struct ggml_tensor * op) { + GGML_ASSERT(device); + return device->iface.supports_op(device, op); +} + +bool ggml_backend_dev_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft) { + GGML_ASSERT(device); + return device->iface.supports_buft(device, buft); +} + +bool ggml_backend_dev_offload_op(ggml_backend_dev_t device, const struct ggml_tensor * op) { + GGML_ASSERT(device); + if (device->iface.offload_op != NULL) { + return device->iface.offload_op(device, op); + } + + return false; +} + +// Backend (reg) + +const char * ggml_backend_reg_name(ggml_backend_reg_t reg) { + GGML_ASSERT(reg); + return reg->iface.get_name(reg); +} + +size_t ggml_backend_reg_dev_count(ggml_backend_reg_t reg) { + GGML_ASSERT(reg); + return reg->iface.get_device_count(reg); +} + +ggml_backend_dev_t ggml_backend_reg_dev_get(ggml_backend_reg_t reg, size_t index) { + GGML_ASSERT(reg); + return reg->iface.get_device(reg, index); +} + +void * ggml_backend_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { + GGML_ASSERT(reg); + if (!reg->iface.get_proc_address) { + return NULL; + } + return reg->iface.get_proc_address(reg, name); +} + +// multi-buffer buffer + +struct ggml_backend_multi_buffer_context { + ggml_backend_buffer_t * buffers; + size_t n_buffers; +}; + +static void ggml_backend_multi_buffer_free_buffer(ggml_backend_buffer_t buffer) { + GGML_ASSERT(buffer); + ggml_backend_multi_buffer_context * ctx = (ggml_backend_multi_buffer_context *) buffer->context; + for (size_t i = 0; i < ctx->n_buffers; i++) { + ggml_backend_buffer_free(ctx->buffers[i]); + } + + free(ctx->buffers); + free(ctx); +} + +static void ggml_backend_multi_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + GGML_ASSERT(buffer); + ggml_backend_multi_buffer_context * ctx = (ggml_backend_multi_buffer_context *) buffer->context; + for (size_t i = 0; i < ctx->n_buffers; i++) { + ggml_backend_buffer_clear(ctx->buffers[i], value); + } +} + +static const struct ggml_backend_buffer_i ggml_backend_multi_buffer_i = { + /* .free_buffer = */ ggml_backend_multi_buffer_free_buffer, + /* .get_base = */ NULL, + /* .init_tensor = */ NULL, + /* .memset_tensor = */ NULL, + /* .set_tensor = */ NULL, + /* .get_tensor = */ NULL, + /* .set_tensor_2d = */ NULL, + /* .get_tensor_2d = */ NULL, + /* .cpy_tensor = */ NULL, + /* .clear = */ ggml_backend_multi_buffer_clear, + /* .reset = */ NULL, +}; + +ggml_backend_buffer_t ggml_backend_multi_buffer_alloc_buffer(ggml_backend_buffer_t * buffers, size_t n_buffers) { + ggml_backend_multi_buffer_context * ctx = (ggml_backend_multi_buffer_context *) malloc(sizeof(struct ggml_backend_multi_buffer_context)); + ctx->n_buffers = n_buffers; + ctx->buffers = (ggml_backend_buffer_t *) malloc(n_buffers * sizeof(ggml_backend_buffer_t)); + + GGML_ASSERT(ctx->buffers != NULL); + + size_t total_size = 0; + for (size_t i = 0; i < n_buffers; i++) { + ctx->buffers[i] = buffers[i]; + total_size += ggml_backend_buffer_get_size(buffers[i]); + } + + return ggml_backend_buffer_init(buffers[0]->buft, ggml_backend_multi_buffer_i, ctx, total_size); +} + +bool ggml_backend_buffer_is_multi_buffer(ggml_backend_buffer_t buffer) { + GGML_ASSERT(buffer); + return buffer->iface.free_buffer == ggml_backend_multi_buffer_free_buffer; +} + +void ggml_backend_multi_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage) { + GGML_ASSERT(buffer); + GGML_ASSERT(ggml_backend_buffer_is_multi_buffer(buffer)); + ggml_backend_multi_buffer_context * ctx = (ggml_backend_multi_buffer_context *) buffer->context; + for (size_t i = 0; i < ctx->n_buffers; i++) { + ggml_backend_buffer_set_usage(ctx->buffers[i], usage); + } +} + +// creates a copy of the tensor with the same memory layout +static struct ggml_tensor * ggml_dup_tensor_layout(struct ggml_context * ctx, const struct ggml_tensor * tensor) { + struct ggml_tensor * dup = ggml_dup_tensor(ctx, tensor); + for (int i = 0; i < GGML_MAX_DIMS; i++) { + dup->nb[i] = tensor->nb[i]; + } + return dup; +} + +static bool ggml_is_view_op(enum ggml_op op) { + return op == GGML_OP_VIEW || op == GGML_OP_RESHAPE || op == GGML_OP_PERMUTE || op == GGML_OP_TRANSPOSE; +} + +// scheduler + +#ifndef GGML_SCHED_MAX_BACKENDS +#define GGML_SCHED_MAX_BACKENDS 16 +#endif + +#ifndef GGML_SCHED_MAX_SPLIT_INPUTS +#define GGML_SCHED_MAX_SPLIT_INPUTS 30 +#endif + +#ifndef GGML_SCHED_MAX_COPIES +#define GGML_SCHED_MAX_COPIES 4 +#endif + +struct ggml_backend_sched_split { + int backend_id; + int i_start; + int i_end; + struct ggml_tensor * inputs[GGML_SCHED_MAX_SPLIT_INPUTS]; + int n_inputs; + // graph view of this split + struct ggml_cgraph graph; +}; + +struct ggml_backend_sched { + bool is_reset; // true if the scheduler has been reset since the last graph split + bool is_alloc; + + int n_backends; + + ggml_backend_t backends[GGML_SCHED_MAX_BACKENDS]; + ggml_backend_buffer_type_t bufts[GGML_SCHED_MAX_BACKENDS]; + ggml_gallocr_t galloc; + + // hash map of the nodes in the graph + struct ggml_hash_set hash_set; + int * hv_tensor_backend_ids; // [hash_set.size] + struct ggml_tensor ** hv_tensor_copies; // [hash_set.size][n_backends][n_copies] + + int * node_backend_ids; // [graph_size] + int * leaf_backend_ids; // [graph_size] + + int * prev_node_backend_ids; // [graph_size] + int * prev_leaf_backend_ids; // [graph_size] + + // copy of the graph with modified inputs + struct ggml_cgraph graph; + + // graph splits + struct ggml_backend_sched_split * splits; + int n_splits; + int splits_capacity; + + // pipeline parallelism support + int n_copies; + int cur_copy; + int next_copy; + ggml_backend_event_t events[GGML_SCHED_MAX_BACKENDS][GGML_SCHED_MAX_COPIES]; + struct ggml_tensor * graph_inputs[GGML_SCHED_MAX_SPLIT_INPUTS]; + int n_graph_inputs; + + struct ggml_context * ctx; + + ggml_backend_sched_eval_callback callback_eval; + void * callback_eval_user_data; + + char * context_buffer; + size_t context_buffer_size; + + bool op_offload; + + int debug; + + // used for debugging graph reallocations [GGML_SCHED_DEBUG_REALLOC] + // ref: https://github.com/ggml-org/llama.cpp/pull/17617 + int debug_realloc; + int debug_graph_size; + int debug_prev_graph_size; +}; + +#define hash_id(tensor) ggml_hash_find_or_insert(&sched->hash_set, tensor) +#define tensor_backend_id(tensor) sched->hv_tensor_backend_ids[hash_id(tensor)] +#define tensor_id_copy(id, backend_id, copy_id) sched->hv_tensor_copies[(id) * sched->n_backends * sched->n_copies + (backend_id) * sched->n_copies + (copy_id)] +#define tensor_copy(tensor, backend_id, copy_id) tensor_id_copy(hash_id(tensor), backend_id, copy_id) + +// returns the priority of the backend, lower id is higher priority +static int ggml_backend_sched_backend_id(ggml_backend_sched_t sched, ggml_backend_t backend) { + for (int i = 0; i < sched->n_backends; i++) { + if (sched->backends[i] == backend) { + return i; + } + } + return -1; +} + +static int ggml_backend_sched_backend_from_buffer(ggml_backend_sched_t sched, const struct ggml_tensor * tensor, const struct ggml_tensor * op) { + ggml_backend_buffer_t buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + if (buffer == NULL) { + return -1; + } + + // find highest prio backend that supports the buffer type and the op + for (int i = 0; i < sched->n_backends; i++) { + if (ggml_backend_supports_buft(sched->backends[i], buffer->buft) && + ggml_backend_supports_op(sched->backends[i], op)) { + return i; + } + } + +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: warning: no backend supports op %s with a weight with buffer type %s used in tensor %s, the weight will need to be copied\n", + __func__, ggml_op_desc(tensor), ggml_backend_buffer_name(buffer), tensor->name); +#endif + + return -1; +} + +#if 0 +#define GGML_SCHED_MAX_SPLITS_DEBUG 4096 +static char causes[GGML_DEFAULT_GRAPH_SIZE*16 + GGML_SCHED_MAX_SPLITS_DEBUG*GGML_SCHED_MAX_SPLIT_INPUTS][128]; // debug only +#define SET_CAUSE(node, ...) sprintf(causes[hash_id(node)], __VA_ARGS__) +#define GET_CAUSE(node) causes[hash_id(node)] +#else +#define SET_CAUSE(node, ...) +#define GET_CAUSE(node) "" +#endif + +// returns the backend that should be used for the node based on the current locations +static int ggml_backend_sched_backend_id_from_cur(ggml_backend_sched_t sched, struct ggml_tensor * tensor) { + // assign pre-allocated nodes to their backend + int cur_backend_id = ggml_backend_sched_backend_from_buffer(sched, tensor, tensor); + if (cur_backend_id != -1) { + SET_CAUSE(tensor, "1.dst"); + return cur_backend_id; + } + + // view_src + if (tensor->view_src != NULL) { + cur_backend_id = ggml_backend_sched_backend_from_buffer(sched, tensor->view_src, tensor); + if (cur_backend_id != -1) { + SET_CAUSE(tensor, "1.vsrc"); + return cur_backend_id; + } + } + + if (tensor->buffer || (tensor->view_src && tensor->view_src->buffer)) { + // since the tensor is pre-allocated, it cannot be moved to another backend + ggml_backend_buffer_t buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + GGML_ABORT("pre-allocated tensor (%s) in a buffer (%s) that cannot run the operation (%s)", tensor->name, ggml_backend_buffer_name(buffer), ggml_op_name(tensor->op)); + } + + // graph input + if (tensor->flags & GGML_TENSOR_FLAG_INPUT) { + cur_backend_id = sched->n_backends - 1; // last backend (assumed CPU) + SET_CAUSE(tensor, "1.inp"); + return cur_backend_id; + } + + // operations with weights are preferably run on the same backend as the weights + for (int i = 0; i < GGML_MAX_SRC; i++) { + const struct ggml_tensor * src = tensor->src[i]; + if (src == NULL) { + continue; + } + // skip ROPE since the rope freqs tensor is too small to choose a backend based on it + // not an ideal solution + if (tensor->op != GGML_OP_ROPE && src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + int src_backend_id = ggml_backend_sched_backend_from_buffer(sched, src, tensor); + // check if a backend with higher prio wants to offload the op + if (sched->op_offload && src_backend_id == sched->n_backends - 1 && ggml_backend_buffer_is_host(src->buffer)) { + for (int b = 0; b < src_backend_id; b++) { + if (ggml_backend_supports_op(sched->backends[b], tensor) && ggml_backend_offload_op(sched->backends[b], tensor)) { + SET_CAUSE(tensor, "1.off"); + return b; + } + } + } + SET_CAUSE(tensor, "1.wgt%d", i); + return src_backend_id; + } + } + + return -1; +} + +static char * fmt_size(size_t size) { + static char buffer[128]; + if (size >= 1024*1024) { + snprintf(buffer, sizeof(buffer), "%zuM", size/1024/1024); + } else { + snprintf(buffer, sizeof(buffer), "%zuK", size/1024); + } + return buffer; +} + +static void ggml_backend_sched_print_assignments(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { + int cur_split = 0; + for (int i = 0; i < graph->n_nodes; i++) { + if (cur_split < sched->n_splits && i == sched->splits[cur_split].i_start) { + ggml_backend_t split_backend = sched->backends[sched->splits[cur_split].backend_id]; + GGML_LOG_DEBUG("\n## SPLIT #%d: %s # %d inputs", cur_split, ggml_backend_name(split_backend), + sched->splits[cur_split].n_inputs); + for (int j = 0; j < sched->splits[cur_split].n_inputs; j++) { + if (j == 0) { + GGML_LOG_DEBUG(": "); + } + GGML_LOG_DEBUG("[%s (%5.5s)] ", sched->splits[cur_split].inputs[j]->name, + fmt_size(ggml_nbytes(sched->splits[cur_split].inputs[j]))); + } + GGML_LOG_DEBUG("\n"); + cur_split++; + } + struct ggml_tensor * node = graph->nodes[i]; + if (ggml_is_view_op(node->op)) { + continue; + } + if (sched->debug > 1) { + ggml_backend_t tensor_backend = ggml_backend_sched_get_tensor_backend(sched, node); + GGML_LOG_DEBUG("node #%3d (%10.10s): %20.20s (%5.5s) [%5.5s %8.8s] use=%d,c=%d:", i, ggml_op_name(node->op), node->name, + fmt_size(ggml_nbytes(node)), tensor_backend ? ggml_backend_name(tensor_backend) : "NULL", GET_CAUSE(node), + graph->use_counts[ggml_hash_find(&graph->visited_hash_set, node)], node->flags & GGML_TENSOR_FLAG_COMPUTE ? 1 : 0); + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (src == NULL) { + continue; + } + ggml_backend_t src_backend = ggml_backend_sched_get_tensor_backend(sched, src); + GGML_LOG_DEBUG(" %20.20s (%5.5s) [%5.5s %8.8s]", src->name, + fmt_size(ggml_nbytes(src)), src_backend ? ggml_backend_name(src_backend) : "NULL", GET_CAUSE(src)); + } + GGML_LOG_DEBUG("\n"); + } + } +} + +static bool ggml_backend_sched_buffer_supported(ggml_backend_sched_t sched, struct ggml_tensor * t, int backend_id) { + ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; + ggml_backend_buffer_type_t buft = NULL; + + if (buf) { + // the tensor is already allocated + buft = buf->buft; + } else { + // see if the tensor already has a backend assigned, and use the buffer type of that backend + int tensor_backend_id = tensor_backend_id(t); + if (tensor_backend_id == -1 && t->view_src) { + tensor_backend_id = tensor_backend_id(t->view_src); + } + if (tensor_backend_id != -1) { + buft = sched->bufts[tensor_backend_id]; + } + } + + return buft != NULL && ggml_backend_supports_buft(sched->backends[backend_id], buft); +} + +static void ggml_backend_sched_set_if_supported(ggml_backend_sched_t sched, struct ggml_tensor * node, int cur_backend_id, int * node_backend_id) { + if (ggml_backend_supports_op(sched->backends[cur_backend_id], node)) { + *node_backend_id = cur_backend_id; + SET_CAUSE(node, "2.sup"); + } +} + +// assigns backends to ops and splits the graph into subgraphs that can be computed on the same backend +void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { + // reset splits + sched->n_splits = 0; + sched->n_graph_inputs = 0; + sched->is_reset = false; + + struct ggml_init_params params = { + /* .mem_size = */ sched->context_buffer_size, + /* .mem_buffer = */ sched->context_buffer, + /* .no_alloc = */ true + }; + + ggml_free(sched->ctx); + + sched->ctx = ggml_init(params); + if (sched->ctx == NULL) { + GGML_ABORT("%s: failed to initialize context\n", __func__); + } + + graph->uid = ggml_graph_next_uid(); + + // pass 1: assign backends to ops with pre-allocated inputs + for (int i = 0; i < graph->n_leafs; i++) { + struct ggml_tensor * leaf = graph->leafs[i]; + int * leaf_backend_id = &tensor_backend_id(leaf); + // do not overwrite user assignments + if (*leaf_backend_id == -1) { + *leaf_backend_id = ggml_backend_sched_backend_id_from_cur(sched, leaf); + } + } + + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + int * node_backend_id = &tensor_backend_id(node); + // do not overwrite user assignments + if (*node_backend_id == -1) { + *node_backend_id = ggml_backend_sched_backend_id_from_cur(sched, node); + +#if 0 + // src + if (node->op == GGML_OP_NONE) { + continue; + } + + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (src == NULL) { + continue; + } + int * src_backend_id = &tensor_backend_id(src); + if (*src_backend_id == -1) { + *src_backend_id = ggml_backend_sched_backend_id_from_cur(sched, src); + } + } +#endif + } + } + + // pass 2: expand current backend assignments + // assign the same backend to adjacent nodes + // expand gpu backends (i.e. non last prio) up and down, ignoring cpu (the lowest priority backend) + // thus, cpu will never be used unless weights are on cpu, or there are no gpu ops between cpu ops + // ops unsupported by the backend being expanded will be left unassigned so that they can be assigned later when the locations of its inputs are known + // expand gpu down + { + int cur_backend_id = -1; + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + if (ggml_is_view_op(node->op)) { + continue; + } + int * node_backend_id = &tensor_backend_id(node); + if (*node_backend_id != -1) { + if (*node_backend_id == sched->n_backends - 1) { + // skip cpu (lowest prio backend) + cur_backend_id = -1; + } else { + cur_backend_id = *node_backend_id; + } + } else if (cur_backend_id != -1) { + ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id); + } + } + } + // expand gpu up + { + int cur_backend_id = -1; + for (int i = graph->n_nodes - 1; i >= 0; i--) { + struct ggml_tensor * node = graph->nodes[i]; + if (ggml_is_view_op(node->op)) { + continue; + } + int * node_backend_id = &tensor_backend_id(node); + if (*node_backend_id != -1) { + if (*node_backend_id == sched->n_backends - 1) { + // skip cpu (lowest prio backend) + cur_backend_id = -1; + } else { + cur_backend_id = *node_backend_id; + } + } else if (cur_backend_id != -1) { + ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id); + } + } + } + // expand rest down + { + int cur_backend_id = -1; + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + if (ggml_is_view_op(node->op)) { + continue; + } + int * node_backend_id = &tensor_backend_id(node); + if (*node_backend_id != -1) { + cur_backend_id = *node_backend_id; + } else if (cur_backend_id != -1) { + ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id); + } + } + } + // expand rest up + { + int cur_backend_id = -1; + for (int i = graph->n_nodes - 1; i >= 0; i--) { + struct ggml_tensor * node = graph->nodes[i]; + if (ggml_is_view_op(node->op)) { + continue; + } + int * node_backend_id = &tensor_backend_id(node); + if (*node_backend_id != -1) { + cur_backend_id = *node_backend_id; + } else if (cur_backend_id != -1) { + ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id); + } + } + } + + // pass 3: upgrade nodes to higher prio backends with compatible buffer types + // if the tensor is already in the same buffer type (*) as another higher priority backend, we should move it there + // however, we also need to verify that the sources are in compatible buffer types + // (*) the actual requirement is more relaxed, the buffer type of the backend should be supported by all the users of this tensor further down the graph + // however, this is slow to verify, so we have a more strict requirement that the buffer type is the same + // this is not uncommon since multiple backends can use host memory, with the same buffer type (eg. BLAS and CPU) + // additionally, set remaining unassigned nodes to the backend with the most supported inputs + // only nodes that could not be assigned during expansion due to the backend not supporting the op should be unassigned at this point + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + if (ggml_is_view_op(node->op)) { + continue; + } + int * node_backend_id = &tensor_backend_id(node); + if (*node_backend_id == -1) { + // unassigned node: find the backend with the most supported inputs + int n_supported_best = -1; + for (int b = 0; b < sched->n_backends; b++) { + if (ggml_backend_supports_op(sched->backends[b], node)) { + int n_supported = 0; + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (src == NULL) { + continue; + } + if ((tensor_backend_id(src) != -1 || tensor_backend_id(src->view_src) != -1) && ggml_backend_sched_buffer_supported(sched, src, b)) { + n_supported++; + } + } + if (n_supported > n_supported_best) { + n_supported_best = n_supported; + *node_backend_id = b; + SET_CAUSE(node, "3.best"); + } + } + } + } else { + // assigned node: upgrade to higher prio backend if possible + for (int b = 0; b < *node_backend_id; b++) { + if (sched->bufts[b] == sched->bufts[*node_backend_id] && ggml_backend_supports_op(sched->backends[b], node)) { + bool supported = true; + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (src == NULL) { + continue; + } + if (!ggml_backend_sched_buffer_supported(sched, src, b)) { + supported = false; + break; + } + } + if (supported) { + *node_backend_id = b; + SET_CAUSE(node, "3.upg"); + break; + } + } + } + } + } + + // pass 4: assign backends to remaining src from dst and view_src + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + int * cur_backend_id = &tensor_backend_id(node); + if (node->view_src != NULL && *cur_backend_id == -1) { + *cur_backend_id = tensor_backend_id(node->view_src); + SET_CAUSE(node, "4.vsrc"); + } + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (src == NULL) { + continue; + } + int * src_backend_id = &tensor_backend_id(src); + if (*src_backend_id == -1) { + if (src->view_src != NULL) { + // views are always on the same backend as the source + *src_backend_id = tensor_backend_id(src->view_src); + SET_CAUSE(src, "4.vsrc"); + } else { + *src_backend_id = *cur_backend_id; + SET_CAUSE(src, "4.cur"); + } + } + } + // if the node is still unassigned, assign it to the first backend that supports it + for (int b = 0; b < sched->n_backends && *cur_backend_id == -1; b++) { + ggml_backend_sched_set_if_supported(sched, node, b, cur_backend_id); + } + GGML_ASSERT(*cur_backend_id != -1); + } + + // pass 5: split graph, find tensors that need to be copied + { + int i_split = 0; + struct ggml_backend_sched_split * split = &sched->splits[0]; + // find the backend of the first split, skipping view ops + int i = 0; + for (; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + if (!ggml_is_view_op(node->op)) { + split->backend_id = tensor_backend_id(node); + break; + } + } + split->i_start = 0; + split->n_inputs = 0; + int cur_backend_id = split->backend_id; + for (; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + + if (ggml_is_view_op(node->op)) { + continue; + } + + const int node_backend_id = tensor_backend_id(node); + + GGML_ASSERT(node_backend_id != -1); // all nodes should be assigned by now, this can happen if there is no CPU fallback + + // check if we should start a new split based on the sources of the current node + bool need_new_split = false; + if (node_backend_id == cur_backend_id && split->n_inputs > 0) { + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (src == NULL) { + continue; + } + // check if a weight is on a different and incompatible backend + // by starting a new split, the memory of the previously offloaded weights can be reused + if (src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + int src_backend_id = tensor_backend_id(src); + if (src_backend_id != cur_backend_id && !ggml_backend_sched_buffer_supported(sched, src, cur_backend_id)) { + need_new_split = true; + break; + } + } + // check if the split has too many inputs + // FIXME: count the number of inputs instead of only checking when full + if (split->n_inputs == GGML_SCHED_MAX_SPLIT_INPUTS) { + const size_t id = hash_id(src); + int src_backend_id = sched->hv_tensor_backend_ids[id]; + bool supported = ggml_backend_sched_buffer_supported(sched, src, cur_backend_id); + if (src_backend_id != cur_backend_id && tensor_id_copy(id, cur_backend_id, 0) == NULL && !supported) { + need_new_split = true; + break; + } + } + } + } + + if (node_backend_id != cur_backend_id || need_new_split) { + split->i_end = i; + i_split++; + if (i_split >= sched->splits_capacity) { + sched->splits_capacity *= 2; + sched->splits = (ggml_backend_sched_split *) + realloc(sched->splits, sched->splits_capacity * sizeof(struct ggml_backend_sched_split)); + GGML_ASSERT(sched->splits != NULL); + } + split = &sched->splits[i_split]; + split->backend_id = node_backend_id; + split->i_start = i; + split->n_inputs = 0; + cur_backend_id = node_backend_id; + } + + // find inputs that are not on the same backend + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (src == NULL) { + continue; + } + + size_t src_id = hash_id(src); + const int src_backend_id = sched->hv_tensor_backend_ids[src_id]; + GGML_ASSERT(src_backend_id != -1); // all inputs should be assigned by now + + if (src->flags & GGML_TENSOR_FLAG_INPUT && sched->n_copies > 1) { + if (tensor_id_copy(src_id, src_backend_id, 0) == NULL) { + ggml_backend_t backend = sched->backends[src_backend_id]; + for (int c = 0; c < sched->n_copies; c++) { + struct ggml_tensor * tensor_copy; + if (c == sched->cur_copy) { + tensor_copy = src; // use the original tensor as the current copy + } else { + tensor_copy = ggml_dup_tensor_layout(sched->ctx, src); + ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c); + } + ggml_set_input(tensor_copy); + ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor + tensor_id_copy(src_id, src_backend_id, c) = tensor_copy; + SET_CAUSE(tensor_copy, "4.cpy"); + } + int n_graph_inputs = sched->n_graph_inputs++; + GGML_ASSERT(n_graph_inputs < GGML_SCHED_MAX_SPLIT_INPUTS); + sched->graph_inputs[n_graph_inputs] = src; + } + } + + if (src_backend_id != cur_backend_id && !ggml_backend_sched_buffer_supported(sched, src, cur_backend_id)) { + // create a copy of the input in the split's backend + if (tensor_id_copy(src_id, cur_backend_id, 0) == NULL) { + ggml_backend_t backend = sched->backends[cur_backend_id]; + for (int c = 0; c < sched->n_copies; c++) { + struct ggml_tensor * tensor_copy = ggml_dup_tensor_layout(sched->ctx, src); + ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c); + if (sched->n_copies > 1) { + ggml_set_input(tensor_copy); + ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor + } + tensor_id_copy(src_id, cur_backend_id, c) = tensor_copy; + SET_CAUSE(tensor_copy, "4.cpy"); + } + int n_inputs = split->n_inputs++; + GGML_ASSERT(n_inputs < GGML_SCHED_MAX_SPLIT_INPUTS); + split->inputs[n_inputs] = src; + } + node->src[j] = tensor_id_copy(src_id, cur_backend_id, sched->cur_copy); + } + } + } + split->i_end = graph->n_nodes; + sched->n_splits = i_split + 1; + } + + if (sched->debug) { + ggml_backend_sched_print_assignments(sched, graph); + } + + // swap node_backend_ids and leaf _backend_ids with prevs + { + int * tmp = sched->node_backend_ids; + sched->node_backend_ids = sched->prev_node_backend_ids; + sched->prev_node_backend_ids = tmp; + + tmp = sched->leaf_backend_ids; + sched->leaf_backend_ids = sched->prev_leaf_backend_ids; + sched->prev_leaf_backend_ids = tmp; + } + + int graph_size = std::max(graph->n_nodes, graph->n_leafs) + sched->n_splits*GGML_SCHED_MAX_SPLIT_INPUTS*2*sched->n_copies; + + // remember the actual graph_size for performing reallocation checks later [GGML_SCHED_DEBUG_REALLOC] + sched->debug_prev_graph_size = sched->debug_graph_size; + sched->debug_graph_size = graph_size; + + if (sched->graph.size < graph_size) { + sched->graph.size = graph_size; + sched->graph.nodes = (ggml_tensor **) realloc(sched->graph.nodes, graph_size * sizeof(struct ggml_tensor *)); + sched->graph.leafs = (ggml_tensor **) realloc(sched->graph.leafs, graph_size * sizeof(struct ggml_tensor *)); + GGML_ASSERT(sched->graph.nodes != NULL); + GGML_ASSERT(sched->graph.leafs != NULL); + } + sched->graph.n_nodes = 0; + sched->graph.n_leafs = 0; + + struct ggml_cgraph * graph_copy = &sched->graph; + + for (int i = 0; i < sched->n_splits; i++) { + struct ggml_backend_sched_split * split = &sched->splits[i]; + split->graph = ggml_graph_view(graph, split->i_start, split->i_end); + + // Optimize this split of the graph. This needs to happen before we make graph_copy, + // so they are in sync. + ggml_backend_graph_optimize(sched->backends[split->backend_id], &split->graph); + + // add inputs to the graph copy so that they are allocated by ggml-alloc at the start of the split + for (int j = 0; j < split->n_inputs; j++) { + assert(graph_copy->size > (graph_copy->n_nodes + 1)); + + struct ggml_tensor * input = split->inputs[j]; + const size_t input_id = hash_id(input); + struct ggml_tensor * input_cpy = tensor_id_copy(input_id, split->backend_id, sched->cur_copy); + + // add a dependency to the input source so that it is not freed before the copy is done + struct ggml_tensor * input_dep = ggml_view_tensor(sched->ctx, input); + input_dep->src[0] = input; + sched->node_backend_ids[graph_copy->n_nodes] = sched->hv_tensor_backend_ids[input_id]; + graph_copy->nodes[graph_copy->n_nodes++] = input_dep; + + // add a dependency to the input copy so that it is allocated at the start of the split + sched->node_backend_ids[graph_copy->n_nodes] = split->backend_id; + graph_copy->nodes[graph_copy->n_nodes++] = input_cpy; + } + + for (int j = split->i_start; j < split->i_end; j++) { + assert(graph_copy->size > graph_copy->n_nodes); + sched->node_backend_ids[graph_copy->n_nodes] = tensor_backend_id(graph->nodes[j]); + graph_copy->nodes[graph_copy->n_nodes++] = graph->nodes[j]; + } + } + + if (sched->n_copies > 1) { + // add input copies as leafs so that they are allocated first + for (int i = 0; i < sched->n_graph_inputs; i++) { + struct ggml_tensor * input = sched->graph_inputs[i]; + size_t id = hash_id(input); + int backend_id = tensor_backend_id(input); + for (int c = 0; c < sched->n_copies; c++) { + struct ggml_tensor * input_cpy = tensor_id_copy(id, backend_id, c); + sched->leaf_backend_ids[graph_copy->n_leafs] = backend_id; + assert(graph_copy->size > graph_copy->n_leafs); + graph_copy->leafs[graph_copy->n_leafs++] = input_cpy; + } + } + + for (int i = 0; i < sched->n_splits; i++) { + struct ggml_backend_sched_split * split = &sched->splits[i]; + int backend_id = split->backend_id; + for (int j = 0; j < split->n_inputs; j++) { + struct ggml_tensor * input = split->inputs[j]; + size_t id = hash_id(input); + for (int c = 0; c < sched->n_copies; c++) { + struct ggml_tensor * input_cpy = tensor_id_copy(id, backend_id, c); + sched->leaf_backend_ids[graph_copy->n_leafs] = backend_id; + assert(graph_copy->size > graph_copy->n_leafs); + graph_copy->leafs[graph_copy->n_leafs++] = input_cpy; + } + } + } + } + + // add leafs from the original graph + for (int i = 0; i < graph->n_leafs; i++) { + struct ggml_tensor * leaf = graph->leafs[i]; + sched->leaf_backend_ids[graph_copy->n_leafs] = tensor_backend_id(leaf); + assert(graph_copy->size > graph_copy->n_leafs); + graph_copy->leafs[graph_copy->n_leafs++] = leaf; + } + + // set ids for all splits + for (int i = 0; i < sched->n_splits; ++i) { + sched->splits[i].graph.uid = ggml_graph_next_uid(); + } +} + +static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { + bool backend_ids_changed = false; + for (int i = 0; i < sched->graph.n_nodes; i++) { + if (sched->node_backend_ids[i] != sched->prev_node_backend_ids[i] && + sched->bufts[sched->node_backend_ids[i]] != sched->bufts[sched->prev_node_backend_ids[i]]) { + backend_ids_changed = true; + break; + } + } + if (!backend_ids_changed) { + for (int i = 0; i < sched->graph.n_leafs; i++) { + if (sched->leaf_backend_ids[i] != sched->prev_leaf_backend_ids[i] && + sched->bufts[sched->leaf_backend_ids[i]] != sched->bufts[sched->prev_leaf_backend_ids[i]]) { + backend_ids_changed = true; + break; + } + } + } + + // allocate graph + if (backend_ids_changed || !ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: failed to allocate graph, reserving (backend_ids_changed = %d)\n", __func__, backend_ids_changed); +#endif + + if (sched->debug_realloc > 0) { + // we are interested only in situations where the graph was reallocated even though its size remained the same [GGML_SCHED_DEBUG_REALLOC] + // example: https://github.com/ggml-org/llama.cpp/pull/17143 + const bool unexpected = !backend_ids_changed && sched->debug_prev_graph_size == sched->debug_graph_size; + + if (unexpected || sched->debug_realloc > 1) { + GGML_ABORT("%s: unexpected graph reallocation (graph size = %d, nodes = %d, leafs = %d), debug_realloc = %d\n", __func__, + sched->debug_graph_size, sched->graph.n_nodes, sched->graph.n_leafs, sched->debug_realloc); + } + } + + // the re-allocation may cause the split inputs to be moved to a different address + // synchronize without ggml_backend_sched_synchronize to avoid changing cur_copy + for (int i = 0; i < sched->n_backends; i++) { + ggml_backend_synchronize(sched->backends[i]); + } + + ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids); + if (!ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) { + GGML_LOG_ERROR("%s: failed to allocate graph\n", __func__); + return false; + } + } + + return true; +} + +static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) { + GGML_ASSERT(sched); + struct ggml_backend_sched_split * splits = sched->splits; + + ggml_tensor * prev_ids_tensor = nullptr; + std::vector ids; + std::vector used_ids; + + for (int split_id = 0; split_id < sched->n_splits; split_id++) { + struct ggml_backend_sched_split * split = &splits[split_id]; + int split_backend_id = split->backend_id; + ggml_backend_t split_backend = sched->backends[split_backend_id]; + + // copy the input tensors to the split backend + for (int input_id = 0; input_id < split->n_inputs; input_id++) { + ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]); + struct ggml_tensor * input = split->inputs[input_id]; + struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy); + + if (input->flags & GGML_TENSOR_FLAG_INPUT) { + // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done + if (sched->events[split_backend_id][sched->cur_copy] != NULL) { + ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); + } else { + ggml_backend_synchronize(split_backend); + } + ggml_backend_tensor_copy(input, input_cpy); + } else { + // wait for the split backend to finish using the input before overwriting it + if (sched->events[split_backend_id][sched->cur_copy] != NULL) { + ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]); + } else { + ggml_backend_synchronize(split_backend); + } + + // when offloading MoE weights, we can reduce the amount of data copied by copying only the experts that are used + ggml_tensor * node = split->graph.nodes[0]; + if (split->graph.n_nodes > 0 && + ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && + ggml_backend_buffer_is_host(input->buffer) && ( + (node->src[0] == input_cpy && node->op == GGML_OP_MUL_MAT_ID) + //|| (node->src[1] == input_cpy && node->op == GGML_OP_ADD_ID) /* GGML_OP_ADD_ID weights are small and not worth splitting */ + )) { + + const int64_t n_expert = node->op == GGML_OP_MUL_MAT_ID ? input->ne[2] : input->ne[1]; + const size_t expert_size = node->op == GGML_OP_MUL_MAT_ID ? input->nb[2] : input->nb[1]; + + ggml_backend_synchronize(input_backend); + + // get the ids + ggml_tensor * ids_tensor = node->src[2]; + ggml_backend_t ids_backend = split_backend; + + // if the ids tensor is also an input of the split, it may not have been copied yet to the split backend + // in that case, we use the original ids tensor + for (int i = input_id + 1; i < split->n_inputs; i++) { + if (ids_tensor == tensor_copy(split->inputs[i], split_backend_id, sched->cur_copy)) { + ids_tensor = split->inputs[i]; + ids_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[i]); + break; + } + } + + if (ids_tensor != prev_ids_tensor) { + ids.resize(ggml_nbytes(ids_tensor) / sizeof(int32_t)); + ggml_backend_tensor_get_async(ids_backend, ids_tensor, ids.data(), 0, ggml_nbytes(ids_tensor)); + ggml_backend_synchronize(ids_backend); + + // find the used experts + used_ids.clear(); + used_ids.resize(ggml_bitset_size(n_expert)); + for (int64_t i1 = 0; i1 < ids_tensor->ne[1]; i1++) { + for (int64_t i0 = 0; i0 < ids_tensor->ne[0]; i0++) { + int32_t id = ids[i1 * ids_tensor->nb[1]/sizeof(int32_t) + i0 * ids_tensor->nb[0]/sizeof(int32_t)]; + GGML_ASSERT(id >= 0 && id < n_expert); + ggml_bitset_set(used_ids.data(), id); + } + } + + prev_ids_tensor = ids_tensor; + } + + // group consecutive experts and copy them together + auto copy_experts = [&](int32_t first_id, int32_t last_id) { + const size_t expert_offset = first_id * expert_size; + const size_t expert_size_copy = (last_id - first_id + 1) * expert_size; + const size_t padding = std::min(expert_size, 512); + const size_t padding_end = last_id < n_expert - 1 ? padding : 0; + + ggml_backend_tensor_set_async(split_backend, + input_cpy, + (const uint8_t *)input->data + expert_offset, expert_offset, + // copy a bit extra at the to ensure there are no NaNs in the padding of the last expert + // this is necessary for MMQ in the CUDA backend + expert_size_copy + padding_end); + }; + + int id = 0; + while (!ggml_bitset_get(used_ids.data(), id)) { + id++; + } + int32_t first_id = id; + int32_t last_id = first_id; + + for (++id; id < n_expert; ++id) { + if (!ggml_bitset_get(used_ids.data(), id)) { + continue; + } + + if (id == last_id + 1) { + last_id = id; + continue; + } + + copy_experts(first_id, last_id); + + first_id = id; + last_id = id; + } + copy_experts(first_id, last_id); + } else { + // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events + // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface + if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) { + ggml_backend_synchronize(input_backend); + if (sched->events[split_backend_id][sched->cur_copy] != NULL) { + ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); + } else { + ggml_backend_synchronize(split_backend); + } + ggml_backend_tensor_copy(input, input_cpy); + } + } + } + } + + if (!sched->callback_eval) { + enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); + if (ec != GGML_STATUS_SUCCESS) { + return ec; + } + } else { + // similar to ggml_backend_compare_graph_backend + for (int j0 = 0; j0 < split->graph.n_nodes; j0++) { + struct ggml_tensor * t = split->graph.nodes[j0]; + + // check if the user needs data from this node + bool need = sched->callback_eval(t, true, sched->callback_eval_user_data); + + int j1 = j0; + + // determine the range [j0, j1] of nodes that can be computed together + while (!need && j1 < split->graph.n_nodes - 1) { + t = split->graph.nodes[++j1]; + need = sched->callback_eval(t, true, sched->callback_eval_user_data); + } + + struct ggml_cgraph gv = ggml_graph_view(&split->graph, j0, j1 + 1); + + enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &gv); + if (ec != GGML_STATUS_SUCCESS) { + return ec; + } + + // TODO: pass backend to the callback, then the user can decide if they want to synchronize + ggml_backend_synchronize(split_backend); + + if (need && !sched->callback_eval(t, false, sched->callback_eval_user_data)) { + break; + } + + j0 = j1; + } + } + + // record the event of this copy + if (split->n_inputs > 0) { + if (sched->events[split_backend_id][sched->cur_copy] != NULL) { + ggml_backend_event_record(sched->events[split_backend_id][sched->cur_copy], split_backend); + } + } + } + + return GGML_STATUS_SUCCESS; +} + +ggml_backend_sched_t ggml_backend_sched_new( + ggml_backend_t * backends, + ggml_backend_buffer_type_t * bufts, + int n_backends, + size_t graph_size, + bool parallel, + bool op_offload) { + GGML_ASSERT(n_backends > 0); + GGML_ASSERT(n_backends <= GGML_SCHED_MAX_BACKENDS); + GGML_ASSERT(ggml_backend_dev_type(ggml_backend_get_device(backends[n_backends - 1])) == GGML_BACKEND_DEVICE_TYPE_CPU); + + struct ggml_backend_sched * sched = (ggml_backend_sched *) calloc(1, sizeof(struct ggml_backend_sched)); + + const char * GGML_SCHED_DEBUG = getenv("GGML_SCHED_DEBUG"); + sched->debug = GGML_SCHED_DEBUG ? atoi(GGML_SCHED_DEBUG) : 0; + + sched->debug_realloc = 0; +#ifdef GGML_SCHED_NO_REALLOC + sched->debug_realloc = 1; +#endif + const char * GGML_SCHED_DEBUG_REALLOC = getenv("GGML_SCHED_DEBUG_REALLOC"); + sched->debug_realloc = GGML_SCHED_DEBUG_REALLOC ? atoi(GGML_SCHED_DEBUG_REALLOC) : sched->debug_realloc; + + sched->n_backends = n_backends; + sched->n_copies = parallel ? GGML_SCHED_MAX_COPIES : 1; + + // initialize hash table + // FIXME: needs to be size*2 to account for leafs (do it in graph_split instead) + sched->hash_set = ggml_hash_set_new(graph_size); + sched->hv_tensor_backend_ids = (int *) malloc(sched->hash_set.size * sizeof(sched->hv_tensor_backend_ids[0])); + sched->hv_tensor_copies = (ggml_tensor **) malloc(sched->hash_set.size * sched->n_backends * sched->n_copies * sizeof(struct ggml_tensor *)); + + const size_t ggml_sched_max_splits = graph_size; // at most there is one split for each node in the graph + const size_t nodes_size = graph_size + ggml_sched_max_splits*GGML_SCHED_MAX_SPLIT_INPUTS*2; + sched->node_backend_ids = (int *) calloc(nodes_size, sizeof(sched->node_backend_ids[0])); + sched->leaf_backend_ids = (int *) calloc(nodes_size, sizeof(sched->leaf_backend_ids[0])); + sched->prev_node_backend_ids = (int *) calloc(nodes_size, sizeof(sched->prev_node_backend_ids[0])); + sched->prev_leaf_backend_ids = (int *) calloc(nodes_size, sizeof(sched->prev_leaf_backend_ids[0])); + + sched->debug_graph_size = 0; + sched->debug_prev_graph_size = 0; + + sched->context_buffer_size = ggml_sched_max_splits*GGML_SCHED_MAX_SPLIT_INPUTS*2*sizeof(struct ggml_tensor) + ggml_graph_overhead_custom(graph_size, false); + sched->context_buffer = (char *) malloc(sched->context_buffer_size); + + const int initial_splits_capacity = 16; + sched->splits = (ggml_backend_sched_split *) calloc(initial_splits_capacity, sizeof(sched->splits[0])); + sched->splits_capacity = initial_splits_capacity; + + for (int b = 0; b < n_backends; b++) { + sched->backends[b] = backends[b]; + sched->bufts[b] = bufts ? bufts[b] : ggml_backend_get_default_buffer_type(backends[b]); + GGML_ASSERT(ggml_backend_supports_buft(backends[b], sched->bufts[b])); + + if (sched->n_copies > 1) { + for (int c = 0; c < sched->n_copies; c++) { + sched->events[b][c] = ggml_backend_event_new(backends[b]->device); + } + } + } + + sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends); + sched->op_offload = op_offload; + + ggml_backend_sched_reset(sched); + + return sched; +} + +void ggml_backend_sched_free(ggml_backend_sched_t sched) { + if (sched == NULL) { + return; + } + for (int b = 0; b < sched->n_backends; b++) { + for (int c = 0; c < sched->n_copies; c++) { + ggml_backend_event_free(sched->events[b][c]); + } + } + ggml_gallocr_free(sched->galloc); + ggml_free(sched->ctx); + ggml_hash_set_free(&sched->hash_set); + free(sched->splits); + free(sched->hv_tensor_backend_ids); + free(sched->hv_tensor_copies); + free(sched->node_backend_ids); + free(sched->leaf_backend_ids); + free(sched->prev_node_backend_ids); + free(sched->prev_leaf_backend_ids); + free(sched->context_buffer); + free(sched->graph.nodes); + free(sched->graph.leafs); + free(sched); +} + +void ggml_backend_sched_reset(ggml_backend_sched_t sched) { + GGML_ASSERT(sched); + // reset state for the next run + if (!sched->is_reset) { + ggml_hash_set_reset(&sched->hash_set); + memset(sched->hv_tensor_backend_ids, -1, sched->hash_set.size * sizeof(sched->hv_tensor_backend_ids[0])); + memset(sched->hv_tensor_copies, 0, sched->hash_set.size * sched->n_backends * sched->n_copies * sizeof(struct ggml_tensor *)); + sched->is_reset = true; + } + sched->is_alloc = false; +} + +void ggml_backend_sched_reserve_size(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph, size_t * sizes) { + GGML_ASSERT(sched); + GGML_ASSERT((int)sched->hash_set.size >= measure_graph->n_nodes + measure_graph->n_leafs); + GGML_ASSERT(sizes); + + ggml_backend_sched_reset(sched); + + ggml_backend_sched_synchronize(sched); + + ggml_backend_sched_split_graph(sched, measure_graph); + + ggml_gallocr_reserve_n_size(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids, sizes); +} + +bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph) { + GGML_ASSERT(sched); + GGML_ASSERT((int)sched->hash_set.size >= measure_graph->n_nodes + measure_graph->n_leafs); + + ggml_backend_sched_synchronize(sched); + + ggml_backend_sched_split_graph(sched, measure_graph); + + if (!ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) { + return false; + } + + ggml_backend_sched_reset(sched); + + return true; +} + +bool ggml_backend_sched_alloc_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { + GGML_ASSERT(sched); + GGML_ASSERT((int)sched->hash_set.size >= graph->n_nodes + graph->n_leafs); + GGML_ASSERT(!sched->is_alloc); + + sched->cur_copy = sched->next_copy; + sched->next_copy = (sched->next_copy + 1) % sched->n_copies; + + ggml_backend_sched_split_graph(sched, graph); + + if (!ggml_backend_sched_alloc_splits(sched)) { + return false; + } + + sched->is_alloc = true; + + return true; +} + +enum ggml_status ggml_backend_sched_graph_compute(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { + enum ggml_status err = ggml_backend_sched_graph_compute_async(sched, graph); + ggml_backend_sched_synchronize(sched); + return err; +} + +enum ggml_status ggml_backend_sched_graph_compute_async(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { + GGML_ASSERT(sched); + if (!sched->is_reset && !sched->is_alloc) { + ggml_backend_sched_reset(sched); + } + + if (!sched->is_alloc) { + if (!ggml_backend_sched_alloc_graph(sched, graph)) { + return GGML_STATUS_ALLOC_FAILED; + } + } + + return ggml_backend_sched_compute_splits(sched); +} + +void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) { + GGML_ASSERT(sched); + for (int i = 0; i < sched->n_backends; i++) { + ggml_backend_synchronize(sched->backends[i]); + } + if (!sched->is_alloc) { + // if the graph is not already allocated, always use copy 0 after a synchronization + // this ensures that during generation the same copy is used every time, + // which avoids changes in the graph that could cause CUDA or other graphs to be disabled + sched->next_copy = 0; + } +} + +void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data) { + GGML_ASSERT(sched); + sched->callback_eval = callback; + sched->callback_eval_user_data = user_data; +} + +int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched) { + GGML_ASSERT(sched); + return sched->n_splits; +} + +int ggml_backend_sched_get_n_copies(ggml_backend_sched_t sched) { + GGML_ASSERT(sched); + return sched->n_copies; +} + +int ggml_backend_sched_get_n_backends(ggml_backend_sched_t sched) { + GGML_ASSERT(sched); + return sched->n_backends; +} + +ggml_backend_t ggml_backend_sched_get_backend(ggml_backend_sched_t sched, int i) { + GGML_ASSERT(sched); + GGML_ASSERT(i >= 0 && i < sched->n_backends); + return sched->backends[i]; +} + +ggml_backend_buffer_type_t ggml_backend_sched_get_buffer_type(ggml_backend_sched_t sched, ggml_backend_t backend) { + GGML_ASSERT(sched); + int backend_index = ggml_backend_sched_backend_id(sched, backend); + GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); + + return sched->bufts[backend_index]; +} + +size_t ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend) { + GGML_ASSERT(sched); + int backend_index = ggml_backend_sched_backend_id(sched, backend); + GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); + + return ggml_gallocr_get_buffer_size(sched->galloc, backend_index); +} + +void ggml_backend_sched_set_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend) { + GGML_ASSERT(sched); + int backend_index = ggml_backend_sched_backend_id(sched, backend); + GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); + tensor_backend_id(node) = backend_index; + SET_CAUSE(node, "usr"); + sched->is_reset = false; +} + +ggml_backend_t ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, struct ggml_tensor * node) { + GGML_ASSERT(sched); + int backend_index = tensor_backend_id(node); + if (backend_index == -1) { + return NULL; + } + return sched->backends[backend_index]; +} + +// utils + +enum ggml_status ggml_backend_view_init(struct ggml_tensor * tensor) { + GGML_ASSERT(tensor); + GGML_ASSERT(tensor->buffer == NULL); + GGML_ASSERT(tensor->view_src != NULL); + GGML_ASSERT(tensor->view_src->buffer != NULL); + GGML_ASSERT(tensor->view_src->data != NULL); + + tensor->buffer = tensor->view_src->buffer; + tensor->data = (char *)tensor->view_src->data + tensor->view_offs; + return ggml_backend_buffer_init_tensor(tensor->buffer, tensor); +} + +enum ggml_status ggml_backend_tensor_alloc(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, void * addr) { + GGML_ASSERT(tensor); + GGML_ASSERT(tensor->buffer == NULL); + GGML_ASSERT(tensor->data == NULL); + GGML_ASSERT(tensor->view_src == NULL); + GGML_ASSERT(addr >= ggml_backend_buffer_get_base(buffer)); + GGML_ASSERT(ggml_backend_buffer_is_meta(buffer) || + (char *) addr + ggml_backend_buffer_get_alloc_size(buffer, tensor) <= + (char *) ggml_backend_buffer_get_base(buffer) + ggml_backend_buffer_get_size(buffer)); + + tensor->buffer = buffer; + tensor->data = addr; + return ggml_backend_buffer_init_tensor(buffer, tensor); +} + +static struct ggml_tensor * graph_copy_dup_tensor(struct ggml_hash_set hash_set, struct ggml_tensor ** node_copies, + struct ggml_context * ctx_allocated, struct ggml_context * ctx_unallocated, struct ggml_tensor * src) { + + GGML_ASSERT(src != NULL); + GGML_ASSERT(src->data && "graph must be allocated"); + + size_t id = ggml_hash_insert(&hash_set, src); + if (id == GGML_HASHSET_ALREADY_EXISTS) { + return node_copies[ggml_hash_find(&hash_set, src)]; + } + + struct ggml_tensor * dst = ggml_dup_tensor_layout(src->data && !src->view_src ? ctx_allocated : ctx_unallocated, src); + if (src->view_src != NULL) { + dst->view_src = graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, src->view_src); + dst->view_offs = src->view_offs; + } + dst->op = src->op; + dst->flags = src->flags; + memcpy(dst->op_params, src->op_params, sizeof(dst->op_params)); + ggml_set_name(dst, src->name); + + // copy src + for (int i = 0; i < GGML_MAX_SRC; i++) { + struct ggml_tensor * s = src->src[i]; + if (s == NULL) { + continue; + } + dst->src[i] = graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, s); + } + + node_copies[id] = dst; + return dst; +} + +static void graph_copy_init_tensor(struct ggml_hash_set * hash_set, struct ggml_tensor ** node_copies, bool * node_init, struct ggml_tensor * src) { + size_t id = ggml_hash_find(hash_set, src); + if (node_init[id]) { + return; + } + node_init[id] = true; + + struct ggml_tensor * dst = node_copies[id]; + if (dst->view_src != NULL) { + graph_copy_init_tensor(hash_set, node_copies, node_init, src->view_src); + enum ggml_status status = ggml_backend_view_init(dst); + GGML_ASSERT(status == GGML_STATUS_SUCCESS); + } + else { + ggml_backend_tensor_copy(src, dst); + } + + // init src + for (int i = 0; i < GGML_MAX_SRC; i++) { + struct ggml_tensor * s = src->src[i]; + if (s == NULL) { + continue; + } + graph_copy_init_tensor(hash_set, node_copies, node_init, s); + } +} + +struct ggml_backend_graph_copy ggml_backend_graph_copy(ggml_backend_t backend, struct ggml_cgraph * graph) { + GGML_ASSERT(graph); + struct ggml_hash_set hash_set = ggml_hash_set_new(graph->visited_hash_set.size); + struct ggml_tensor ** node_copies = (ggml_tensor **) calloc(hash_set.size, sizeof(node_copies[0])); // NOLINT + bool * node_init = (bool *) calloc(hash_set.size, sizeof(node_init[0])); + + struct ggml_init_params params = { + /* .mem_size = */ ggml_tensor_overhead()*hash_set.size + ggml_graph_overhead_custom(graph->size, false), + /* .mem_buffer = */ NULL, + /* .no_alloc = */ true + }; + + struct ggml_context * ctx_allocated = ggml_init(params); + struct ggml_context * ctx_unallocated = ggml_init(params); + + if (ctx_allocated == NULL || ctx_unallocated == NULL) { + GGML_LOG_ERROR("%s: failed to allocate context for graph copy\n", __func__); + ggml_hash_set_free(&hash_set); + free(node_copies); + free(node_init); + ggml_free(ctx_allocated); + ggml_free(ctx_unallocated); + return { + /* .buffer = */ NULL, + /* .ctx_allocated = */ NULL, + /* .ctx_unallocated = */ NULL, + /* .graph = */ NULL, + }; + } + + // dup nodes + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, node); + } + + // allocate nodes + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx_allocated, backend); + if (buffer == NULL) { + GGML_LOG_ERROR("%s: failed to allocate buffer for graph copy\n", __func__); + ggml_hash_set_free(&hash_set); + free(node_copies); + free(node_init); + ggml_free(ctx_allocated); + ggml_free(ctx_unallocated); + return { + /* .buffer = */ NULL, + /* .ctx_allocated = */ NULL, + /* .ctx_unallocated = */ NULL, + /* .graph = */ NULL, + }; + } + + //printf("copy buffer size: %zu MB\n", ggml_backend_buffer_get_size(buffer) / 1024 / 1024); + + // copy data and init views + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + graph_copy_init_tensor(&hash_set, node_copies, node_init, node); + } + + // build graph copy + struct ggml_cgraph * graph_copy = ggml_new_graph_custom(ctx_allocated, graph->size, false); + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + struct ggml_tensor * node_copy = node_copies[ggml_hash_find(&hash_set, node)]; + graph_copy->nodes[i] = node_copy; + } + graph_copy->n_nodes = graph->n_nodes; + + ggml_hash_set_free(&hash_set); + free(node_copies); + free(node_init); + + return { + /* .buffer = */ buffer, + /* .ctx_allocated = */ ctx_allocated, + /* .ctx_unallocated = */ ctx_unallocated, + /* .graph = */ graph_copy, + }; +} + +void ggml_backend_graph_copy_free(struct ggml_backend_graph_copy copy) { + ggml_backend_buffer_free(copy.buffer); + ggml_free(copy.ctx_allocated); + ggml_free(copy.ctx_unallocated); +} + +bool ggml_backend_compare_graph_backend(ggml_backend_t backend1, ggml_backend_t backend2, struct ggml_cgraph * graph, ggml_backend_eval_callback callback, void * user_data, struct ggml_tensor const * const * test_nodes, size_t num_test_nodes) { + struct ggml_backend_graph_copy copy = ggml_backend_graph_copy(backend2, graph); + if (copy.buffer == NULL) { + return false; + } + + struct ggml_cgraph * g1 = graph; + struct ggml_cgraph * g2 = copy.graph; + + assert(g1->n_nodes == g2->n_nodes); + + if (num_test_nodes != 0) { + GGML_ASSERT(test_nodes); + // Compute the whole graph and only test the output for specific tensors + ggml_backend_graph_compute(backend1, g1); + ggml_backend_graph_compute(backend2, g2); + + bool verified = false; + for (int i = 0; i < g1->n_nodes; i++) { + for (size_t j = 0; j < num_test_nodes; ++j) { + if (g1->nodes[i] == test_nodes[j]) { + callback(i, g1->nodes[i], g2->nodes[i], user_data); + verified = true; + } + } + } + GGML_ASSERT(verified); + } else { + for (int i = 0; i < g1->n_nodes; i++) { + struct ggml_tensor * t1 = g1->nodes[i]; + struct ggml_tensor * t2 = g2->nodes[i]; + + assert(t1->op == t2->op && ggml_are_same_layout(t1, t2)); + + struct ggml_cgraph g1v = ggml_graph_view(g1, i, i + 1); + struct ggml_cgraph g2v = ggml_graph_view(g2, i, i + 1); + + ggml_backend_graph_compute(backend1, &g1v); + ggml_backend_graph_compute(backend2, &g2v); + + if (ggml_is_view_op(t1->op)) { + continue; + } + + // compare results, calculate rms etc + if (!callback(i, t1, t2, user_data)) { + break; + } + } + } + ggml_backend_graph_copy_free(copy); + + return true; +} + +// CPU backend - buffer + +static void * ggml_backend_cpu_buffer_get_base(ggml_backend_buffer_t buffer) { + GGML_ASSERT(buffer); + uintptr_t data = (uintptr_t)buffer->context; + + // align the buffer + if (data % TENSOR_ALIGNMENT != 0) { + data = GGML_PAD(data, TENSOR_ALIGNMENT); + } + + return (void *)data; +} + +static void ggml_backend_cpu_buffer_free_buffer(ggml_backend_buffer_t buffer) { + GGML_ASSERT(buffer); + ggml_aligned_free(buffer->context, buffer->size); +} + +static void ggml_backend_cpu_buffer_memset_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { + GGML_ASSERT(tensor); + memset((char *)tensor->data + offset, value, size); + + GGML_UNUSED(buffer); +} + +static void ggml_backend_cpu_buffer_set_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + GGML_ASSERT(tensor); + memcpy((char *)tensor->data + offset, data, size); + + GGML_UNUSED(buffer); +} + +static void ggml_backend_cpu_buffer_get_tensor(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) { + GGML_ASSERT(tensor); + memcpy(data, (const char *)tensor->data + offset, size); + + GGML_UNUSED(buffer); +} + +static bool ggml_backend_cpu_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const struct ggml_tensor * src, struct ggml_tensor * dst) { + GGML_ASSERT(src); + if (ggml_backend_buffer_is_host(src->buffer)) { + memcpy(dst->data, src->data, ggml_nbytes(src)); + return true; + } + return false; + + GGML_UNUSED(buffer); +} + +static void ggml_backend_cpu_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + GGML_ASSERT(buffer); + memset(buffer->context, value, buffer->size); +} + +static const struct ggml_backend_buffer_i ggml_backend_cpu_buffer_i = { + /* .free_buffer = */ ggml_backend_cpu_buffer_free_buffer, + /* .get_base = */ ggml_backend_cpu_buffer_get_base, + /* .init_tensor = */ NULL, // no initialization required + /* .memset_tensor = */ ggml_backend_cpu_buffer_memset_tensor, + /* .set_tensor = */ ggml_backend_cpu_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_cpu_buffer_get_tensor, + /* .set_tensor_2d = */ NULL, + /* .get_tensor_2d = */ NULL, + /* .cpy_tensor = */ ggml_backend_cpu_buffer_cpy_tensor, + /* .clear = */ ggml_backend_cpu_buffer_clear, + /* .reset = */ NULL, +}; + +static const struct ggml_backend_buffer_i ggml_backend_cpu_buffer_from_ptr_i = { + /* .free_buffer = */ NULL, // ptr is not owned by the buffer, so it does not need to be freed + /* .get_base = */ ggml_backend_cpu_buffer_get_base, + /* .init_tensor = */ NULL, // no initialization required + /* .memset_tensor = */ ggml_backend_cpu_buffer_memset_tensor, + /* .set_tensor = */ ggml_backend_cpu_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_cpu_buffer_get_tensor, + /* .set_tensor_2d = */ NULL, + /* .get_tensor_2d = */ NULL, + /* .cpy_tensor = */ ggml_backend_cpu_buffer_cpy_tensor, + /* .clear = */ ggml_backend_cpu_buffer_clear, + /* .reset = */ NULL, +}; + +// CPU backend buffer type + +// this buffer type is defined here to make it available to all backends + +static const char * ggml_backend_cpu_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + return "CPU"; + + GGML_UNUSED(buft); +} + +static ggml_backend_buffer_t ggml_backend_cpu_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + void * data = ggml_aligned_malloc(size); + + if (data == NULL) { + GGML_LOG_ERROR("%s: failed to allocate buffer of size %zu\n", __func__, size); + return NULL; + } + + return ggml_backend_buffer_init(buft, ggml_backend_cpu_buffer_i, data, size); +} + +static size_t ggml_backend_cpu_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + return TENSOR_ALIGNMENT; + + GGML_UNUSED(buft); +} + +static bool ggml_backend_cpu_buffer_type_is_host(ggml_backend_buffer_type_t buft) { + return true; + + GGML_UNUSED(buft); +} + +ggml_backend_buffer_type_t ggml_backend_cpu_buffer_type(void) { + static struct ggml_backend_buffer_type ggml_backend_cpu_buffer_type = { + /* .iface = */ { + /* .get_name = */ ggml_backend_cpu_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_cpu_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_cpu_buffer_type_get_alignment, + /* .get_max_size = */ NULL, // defaults to SIZE_MAX + /* .get_alloc_size = */ NULL, // defaults to ggml_nbytes + /* .is_host = */ ggml_backend_cpu_buffer_type_is_host, + }, + /* .device = */ NULL, // FIXME ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0), + /* .context = */ NULL, + }; + + return &ggml_backend_cpu_buffer_type; +} + +static const char * ggml_backend_cpu_buffer_from_ptr_type_get_name(ggml_backend_buffer_type_t buft) { + return "CPU_Mapped"; + + GGML_UNUSED(buft); +} + +static ggml_backend_buffer_type_t ggml_backend_cpu_buffer_from_ptr_type(void) { + static struct ggml_backend_buffer_type ggml_backend_cpu_buffer_type = { + /* .iface = */ { + /* .get_name = */ ggml_backend_cpu_buffer_from_ptr_type_get_name, + /* .alloc_buffer = */ ggml_backend_cpu_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_cpu_buffer_type_get_alignment, + /* .get_max_size = */ NULL, // defaults to SIZE_MAX + /* .get_alloc_size = */ NULL, // defaults to ggml_nbytes + /* .is_host = */ ggml_backend_cpu_buffer_type_is_host, + }, + /* .device = */ NULL, // FIXME ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0), + /* .context = */ NULL, + }; + + return &ggml_backend_cpu_buffer_type; +} + +ggml_backend_buffer_t ggml_backend_cpu_buffer_from_ptr(void * ptr, size_t size) { + GGML_ASSERT((uintptr_t)ptr % TENSOR_ALIGNMENT == 0 && "buffer pointer must be aligned"); + return ggml_backend_buffer_init(ggml_backend_cpu_buffer_from_ptr_type(), ggml_backend_cpu_buffer_from_ptr_i, ptr, size); +} + + +// ---- backend meta implementation ---- + +#include "ggml.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-backend.h.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml-alloc.h.inc" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct ggml_context_deleter { + void operator()(ggml_context * ctx) { ggml_free(ctx); } +}; + +struct ggml_backend_buffer_deleter { + void operator()(ggml_backend_buffer_t buffer) { ggml_backend_buffer_free(buffer); } +}; + +using ggml_context_ptr = std::unique_ptr; +using ggml_backend_buffer_ptr = std::unique_ptr; + +struct ggml_backend_meta_device; +struct ggml_backend_meta_buffer_type; +struct ggml_backend_meta_buffer; +struct ggml_backend_meta; + +const char * ggml_backend_meta_split_axis_name(enum ggml_backend_meta_split_axis split_axis) { + switch (split_axis) { + case GGML_BACKEND_SPLIT_AXIS_0: + return "0"; + case GGML_BACKEND_SPLIT_AXIS_1: + return "1"; + case GGML_BACKEND_SPLIT_AXIS_2: + return "2"; + case GGML_BACKEND_SPLIT_AXIS_3: + return "3"; + case GGML_BACKEND_SPLIT_AXIS_MIRRORED: + return "MIRRORED"; + case GGML_BACKEND_SPLIT_AXIS_PARTIAL: + return "PARTIAL"; + case GGML_BACKEND_SPLIT_AXIS_NONE: + return "NONE"; + case GGML_BACKEND_SPLIT_AXIS_UNKNOWN: + return "UNKNOWN"; + default: + GGML_ABORT("fatal error"); + } +} + +// +// meta backend device +// + +struct ggml_backend_meta_device_context { + std::vector simple_devs; + ggml_backend_meta_get_split_state_t get_split_state; + void * get_split_state_ud; + + std::string name; + std::string description; + + ggml_backend_meta_device_context( + std::vector simple_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud) : + simple_devs(std::move(simple_devs)), get_split_state(get_split_state), get_split_state_ud(get_split_state_ud) { + name = std::string("Meta("); + description = std::string("Meta("); + for (size_t i = 0; i < simple_devs.size(); i++) { + if (i > 0) { + name += ","; + description += ","; + } + name += ggml_backend_dev_name (simple_devs[i]); + description += ggml_backend_dev_description(simple_devs[i]); + } + name += ")"; + description += ")"; + } + + bool operator<(const ggml_backend_meta_device_context & other) const { + return std::tie(simple_devs, get_split_state, get_split_state_ud) + < std::tie(other.simple_devs, other.get_split_state, other.get_split_state_ud); + } +}; + +static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev); + +static const char * ggml_backend_meta_device_get_name(ggml_backend_dev_t dev) { + GGML_ASSERT(ggml_backend_dev_is_meta(dev)); + const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; + return meta_dev_ctx->name.c_str(); +} + +static const char * ggml_backend_meta_device_get_description(ggml_backend_dev_t dev) { + GGML_ASSERT(ggml_backend_dev_is_meta(dev)); + const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; + return meta_dev_ctx->description.c_str(); +} + +static void ggml_backend_meta_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { + GGML_ASSERT(ggml_backend_dev_is_meta(dev)); + const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; + *free = 0; + *total = 0; + for (ggml_backend_dev_t dev : meta_dev_ctx->simple_devs) { + size_t tmp_free, tmp_total; + ggml_backend_dev_memory(dev, &tmp_free, &tmp_total); + *free += tmp_free; + *total += tmp_total; + } +} + +static enum ggml_backend_dev_type ggml_backend_meta_device_get_type(ggml_backend_dev_t dev) { + return GGML_BACKEND_DEVICE_TYPE_META; + + GGML_UNUSED(dev); +} + +static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_backend_dev_props * props) { + GGML_ASSERT(ggml_backend_dev_is_meta(dev)); + const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; + + // TODO replace placeholders + props->name = ggml_backend_meta_device_get_name(dev); + props->description = ggml_backend_meta_device_get_description(dev); + props->type = ggml_backend_meta_device_get_type(dev); + props->device_id = 0; + + ggml_backend_meta_device_get_memory(dev, &props->memory_free, &props->memory_total); + + props->caps = { + /* .async = */ true, + /* .host_buffer = */ false, // Not implemented. + /* .buffer_from_host_ptr = */ false, // Not implemented. + /* .events = */ false, // Not implemented. + }; + for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) { + ggml_backend_dev_props tmp_props; + ggml_backend_dev_get_props(simple_dev, &tmp_props); + props->caps.async = props->caps.async && tmp_props.caps.async; + props->caps.host_buffer = props->caps.host_buffer && tmp_props.caps.host_buffer; + props->caps.buffer_from_host_ptr = props->caps.buffer_from_host_ptr && tmp_props.caps.buffer_from_host_ptr; + props->caps.events = props->caps.events && tmp_props.caps.events; + } +} + +static ggml_backend_t ggml_backend_meta_device_init_backend(ggml_backend_dev_t dev, const char * params); + +static ggml_backend_buffer_type_t ggml_backend_meta_device_get_buffer_type(ggml_backend_dev_t dev); + +static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type(ggml_backend_dev_t dev); + +static bool ggml_backend_meta_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + GGML_ASSERT(ggml_backend_dev_is_meta(dev)); + const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; + return std::all_of(meta_dev_ctx->simple_devs.begin(), meta_dev_ctx->simple_devs.end(), + [op](ggml_backend_dev_t simple_dev) { return ggml_backend_dev_supports_op(simple_dev, op); }); +} + +static bool ggml_backend_meta_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { + GGML_ASSERT(ggml_backend_dev_is_meta(dev)); + ggml_backend_dev_t dev_buft = ggml_backend_buft_get_device(buft); + if (!ggml_backend_dev_is_meta(dev_buft)) { + return false; + } + const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; + const ggml_backend_meta_device_context * meta_buft_dev_ctx = (const ggml_backend_meta_device_context *) dev_buft->context; + if (meta_dev_ctx->simple_devs.size() != meta_buft_dev_ctx->simple_devs.size()) { + return false; + } + for (size_t i = 0; i < meta_dev_ctx->simple_devs.size(); i++) { + if (meta_dev_ctx->simple_devs[i] != meta_buft_dev_ctx->simple_devs[i]) { + return false; + } + } + return true; +} + +static const ggml_backend_device_i ggml_backend_meta_device_iface = { + /* .get_name = */ ggml_backend_meta_device_get_name, + /* .get_description = */ ggml_backend_meta_device_get_description, + /* .get_memory = */ ggml_backend_meta_device_get_memory, + /* .get_type = */ ggml_backend_meta_device_get_type, + /* .get_props = */ ggml_backend_meta_device_get_props, + /* .init_backend = */ ggml_backend_meta_device_init_backend, + /* .get_buffer_type = */ ggml_backend_meta_device_get_buffer_type, + /* .get_host_buffer_type = */ ggml_backend_meta_device_get_host_buffer_type, + /* .buffer_from_host_ptr = */ nullptr, + /* .supports_op = */ ggml_backend_meta_device_supports_op, + /* .supports_buft = */ ggml_backend_meta_device_supports_buft, + /* .offload_op = */ nullptr, + /* .event_new = */ nullptr, + /* .event_free = */ nullptr, + /* .event_synchronize = */ nullptr, +}; + +static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev) { + return dev != nullptr && dev->iface.get_name == ggml_backend_meta_device_iface.get_name; +} + +static size_t ggml_backend_meta_dev_n_devs(ggml_backend_dev_t meta_dev) { + GGML_ASSERT(ggml_backend_dev_is_meta(meta_dev)); + const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) meta_dev->context; + return meta_dev_ctx->simple_devs.size(); +} + +static ggml_backend_dev_t ggml_backend_meta_dev_simple_dev(ggml_backend_dev_t meta_dev, size_t index) { + GGML_ASSERT(ggml_backend_dev_is_meta(meta_dev)); + const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) meta_dev->context; + GGML_ASSERT(index < meta_dev_ctx->simple_devs.size()); + return meta_dev_ctx->simple_devs[index]; +} + +ggml_backend_dev_t ggml_backend_meta_device( + ggml_backend_dev_t * devs, size_t n_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud) { + GGML_ASSERT(n_devs <= GGML_BACKEND_META_MAX_DEVICES); + // TODO: this is not thread-safe - needs to be fixed + static std::vector> ctxs; + static std::map meta_devs; + + std::vector simple_devs; + simple_devs.reserve(n_devs); + for (size_t i = 0; i < n_devs; i++) { + simple_devs.push_back(devs[i]); + } + ggml_backend_meta_device_context ctx(simple_devs, get_split_state, get_split_state_ud); + + { + auto it = meta_devs.find(ctx); + if (it != meta_devs.end()) { + return &it->second; + } + } + ctxs.push_back(std::make_unique(ctx)); + + struct ggml_backend_device meta_dev = { + /*iface =*/ ggml_backend_meta_device_iface, + /*reg =*/ nullptr, + /*ctx =*/ ctxs.back().get(), + }; + + auto result = meta_devs.emplace(*ctxs.back(), meta_dev); + return &result.first->second; +} + +// +// meta backend buffer type +// + +struct ggml_backend_meta_buffer_type_context { + std::vector simple_bufts; + + std::string name; + + ggml_backend_meta_buffer_type_context(std::vector simple_bufts) : simple_bufts(std::move(simple_bufts)) { + name = "Meta("; + for (size_t i = 0; i < simple_bufts.size(); i++) { + if (i > 0) { + name += ","; + } + name += ggml_backend_buft_name(simple_bufts[i]); + } + name += ")"; + } + + bool operator<(const ggml_backend_meta_buffer_type_context & other) const { + return simple_bufts < other.simple_bufts; + } +}; + +static size_t ggml_backend_meta_buft_n_bufts(ggml_backend_buffer_type_t meta_buft) { + GGML_ASSERT(ggml_backend_buft_is_meta(meta_buft)); + const ggml_backend_meta_buffer_type_context * meta_buft_ctx = (const ggml_backend_meta_buffer_type_context *) meta_buft->context; + return meta_buft_ctx->simple_bufts.size(); +} + +static const char * ggml_backend_meta_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + GGML_ASSERT(ggml_backend_buft_is_meta(buft)); + const ggml_backend_meta_buffer_type_context * meta_buft_ctx = (const ggml_backend_meta_buffer_type_context *) buft->context; + return meta_buft_ctx->name.c_str(); +} + +static ggml_backend_buffer_type_t ggml_backend_meta_buft_simple_buft(ggml_backend_buffer_type_t meta_buft, size_t index) { + GGML_ASSERT(ggml_backend_buft_is_meta(meta_buft)); + const ggml_backend_meta_buffer_type_context * meta_buft_ctx = (const ggml_backend_meta_buffer_type_context *) meta_buft->context; + GGML_ASSERT(index < meta_buft_ctx->simple_bufts.size()); + return meta_buft_ctx->simple_bufts[index]; +} + +static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size); + +static size_t ggml_backend_meta_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); + size_t max_alignment = 1; + for (size_t i = 0; i < n_simple_bufts; i++) { + const size_t alignment = ggml_backend_buft_get_alignment(ggml_backend_meta_buft_simple_buft(buft, i)); + max_alignment = std::max(max_alignment, alignment); + GGML_ASSERT(max_alignment % alignment == 0); + } + return max_alignment; +} + +static size_t ggml_backend_meta_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { + const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); + size_t max_size = SIZE_MAX; + for (size_t i = 0; i < n_simple_bufts; i++) { + max_size = std::min(max_size, ggml_backend_buft_get_max_size(ggml_backend_meta_buft_simple_buft(buft, i))); + } + return max_size; +} + +static size_t ggml_backend_meta_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { + const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); + size_t max_alloc_size = 0; + for (size_t i = 0; i < n_simple_bufts; i++) { + const size_t alloc_size = ggml_backend_buft_get_alloc_size(ggml_backend_meta_buft_simple_buft(buft, i), tensor); + max_alloc_size = std::max(max_alloc_size, alloc_size); + } + return max_alloc_size; +} + +static bool ggml_backend_meta_buffer_type_is_host(ggml_backend_buffer_type_t buft) { + const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); + for (size_t i = 0; i < n_simple_bufts; i++) { + if (!ggml_backend_buft_is_host(ggml_backend_meta_buft_simple_buft(buft, i))) { + return false; + } + } + return true; +} + +static const struct ggml_backend_buffer_type_i ggml_backend_meta_buffer_type_iface = { + /* .get_name = */ ggml_backend_meta_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_meta_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_meta_buffer_type_get_alignment, + /* .get_max_size = */ ggml_backend_meta_buffer_type_get_max_size, + /* .get_alloc_size = */ ggml_backend_meta_buffer_type_get_alloc_size, + /* .is_host = */ ggml_backend_meta_buffer_type_is_host, +}; + +bool ggml_backend_buft_is_meta(ggml_backend_buffer_type_t buft) { + return buft != nullptr && buft->iface.get_name == ggml_backend_meta_buffer_type_iface.get_name; +} + +static ggml_backend_buffer_type_t ggml_backend_meta_device_get_buffer_type(ggml_backend_dev_t dev) { + static std::map meta_bufts; + GGML_ASSERT(ggml_backend_dev_is_meta(dev)); + { + auto it = meta_bufts.find(dev); + if (it != meta_bufts.end()) { + return &it->second; + } + } + + const size_t n_devs = ggml_backend_meta_dev_n_devs(dev); + std::vector simple_bufts; + simple_bufts.reserve(n_devs); + for (size_t i = 0; i < n_devs; i++) { + simple_bufts.push_back(ggml_backend_dev_buffer_type(ggml_backend_meta_dev_simple_dev(dev, i))); + } + ggml_backend_meta_buffer_type_context * buft_ctx = new ggml_backend_meta_buffer_type_context(simple_bufts); + + struct ggml_backend_buffer_type meta_buft = { + /*iface =*/ ggml_backend_meta_buffer_type_iface, + /*device =*/ dev, + /*ctx =*/ buft_ctx, + }; + auto result = meta_bufts.emplace(dev, meta_buft); + return &result.first->second; +} + +static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type(ggml_backend_dev_t dev) { + GGML_ASSERT(ggml_backend_dev_is_meta(dev)); + const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; + + ggml_backend_buffer_type_t host_buft = nullptr; + for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) { + ggml_backend_buffer_type_t simple_host_buft = ggml_backend_dev_host_buffer_type(simple_dev); + if (simple_host_buft == nullptr) { + return nullptr; + } + if (host_buft == nullptr) { + host_buft = simple_host_buft; + } else if (host_buft != simple_host_buft) { + // if different simple devices have different host buffer types, + // we cannot provide a single host buffer type for the meta device + return nullptr; + } + } + return host_buft; +} + +// +// meta backend buffer +// + +struct ggml_backend_meta_buffer_context { + static constexpr size_t nbtc = GGML_TENSOR_SIZE - sizeof(ggml_tensor::padding); + + std::map, std::pair> split_state_cache; + std::map< const ggml_tensor *, std::vector> simple_tensors; + + struct buffer_config { + ggml_context * ctx; + ggml_backend_buffer_t buf; + + buffer_config(ggml_context * ctx, ggml_backend_buffer_t buf) : ctx(ctx), buf(buf) {} + }; + std::vector buf_configs; + + int debug; + + ggml_backend_meta_buffer_context() { + const char * GGML_META_DEBUG = getenv("GGML_META_DEBUG"); + debug = GGML_META_DEBUG ? atoi(GGML_META_DEBUG) : 0; + } +}; + +static void ggml_backend_meta_buffer_free_buffer(ggml_backend_buffer_t buffer) { + GGML_ASSERT(ggml_backend_buffer_is_meta(buffer)); + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context; + for (auto & [ctx, buf] : buf_ctx->buf_configs) { + ggml_backend_buffer_free(buf); + ggml_free(ctx); + } + delete buf_ctx; +} + +static size_t ggml_backend_meta_buffer_n_bufs(ggml_backend_buffer_t meta_buf) { + GGML_ASSERT(ggml_backend_buffer_is_meta(meta_buf)); + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) meta_buf->context; + return buf_ctx->buf_configs.size(); +} + +static ggml_backend_buffer_t ggml_backend_meta_buffer_simple_buffer(ggml_backend_buffer_t meta_buf, size_t index) { + GGML_ASSERT(ggml_backend_buffer_is_meta(meta_buf)); + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) meta_buf->context; + GGML_ASSERT(index < buf_ctx->buf_configs.size()); + return buf_ctx->buf_configs[index].buf; +} + +static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct ggml_tensor * tensor, size_t index) { + GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer)); + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; + GGML_ASSERT(index < buf_ctx->buf_configs.size()); + + auto it = buf_ctx->simple_tensors.find(tensor); + if (it == buf_ctx->simple_tensors.end()) { + return nullptr; + } + return it->second[index]; +} + +static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync) { + const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; + + auto split_states_equal = [&](const ggml_backend_meta_split_state & a, const ggml_backend_meta_split_state & b) -> bool { + if (a.axis != b.axis) { + return false; + } + for (size_t j = 0; j < n_bufs; j++) { + int64_t sum_a = 0; + for (size_t s = 0; s < a.n_segments; s++) { + sum_a += a.ne[s*n_bufs + j]; + } + int64_t sum_b = 0; + for (size_t s = 0; s < b.n_segments; s++) { + sum_b += b.ne[s*n_bufs + j]; + } + if (sum_a != sum_b) { + return false; + } + } + return true; + }; + + auto handle_generic = [&](const std::vector & src_ss, bool scalar_only) -> ggml_backend_meta_split_state { + ggml_backend_meta_split_state ret = {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, 1}; + for (size_t i = 0; i < GGML_MAX_SRC; i++) { + if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { + continue; + } + if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) { + ret = src_ss[i]; + } else if (!split_states_equal(src_ss[i], ret)) { + ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + break; + } + } + if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) { + ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + } + if (scalar_only && ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { + ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + } + GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); + return ret; + }; + + // Some ops process data on a per-row bases: + auto handle_per_row = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0); + return src_ss[0]; + }; + + // Some ops broadcast the src1 data across src0: + auto handle_bin_bcast = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS && + tensor->src[1]->ne[src_ss[0].axis] == 1 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + return src_ss[0]; + } + if (src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && (src_ss[0].axis == src_ss[1].axis || + (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL)))) { + return src_ss[0]; // GGML_OP_ADD_ID + } + GGML_ASSERT(tensor->src[2] == nullptr || src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + return handle_generic(src_ss, /*scalar_only =*/ false); + }; + + auto handle_concat = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + const ggml_backend_meta_split_axis concat_axis = ggml_backend_meta_split_axis(ggml_get_op_params_i32(tensor, 0)); + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis >= 0 && src_ss[1].axis < GGML_MAX_DIMS) { + GGML_ASSERT(concat_axis != src_ss[1].axis); + return src_ss[1]; + } + if (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { + GGML_ASSERT(concat_axis != src_ss[0].axis); + return src_ss[0]; + } + if (src_ss[0].axis == src_ss[1].axis && src_ss[0].axis != concat_axis) { + return src_ss[0]; + } + return handle_generic(src_ss, /*scalar_only =*/ true); + }; + + auto handle_mul_mat = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, 1}; + } + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + ggml_backend_meta_split_state ret = src_ss[0]; + ret.axis = GGML_BACKEND_SPLIT_AXIS_0; + ret.n_segments = 1; + return ret; + } + if (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1 && src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + ggml_backend_meta_split_state ret = src_ss[1]; + ret.n_segments = 1; + return ret; + } + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_0) { + GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1])); + return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_PARTIAL, {0}, 1}; + } + GGML_ABORT("fatal error"); + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + }; + + auto handle_cpy = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { + int64_t ne_split_src = tensor->src[0]->ne[0]; + for (int dim = 1; dim <= src_ss[0].axis; dim++) { + ne_split_src *= tensor->src[0]->ne[dim]; + } + int64_t ne_split_dst = 1; + for (int dim = 0; dim < GGML_MAX_DIMS; dim++) { + ne_split_dst *= tensor->ne[dim]; + if (ne_split_dst == ne_split_src) { + return {ggml_backend_meta_split_axis(dim), {0}, 1}; + } + } + } + return handle_generic(src_ss, /*scalar_only =*/ false); + }; + + auto handle_reshape = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + switch (src_ss[0].axis) { + case GGML_BACKEND_SPLIT_AXIS_0: + case GGML_BACKEND_SPLIT_AXIS_1: + case GGML_BACKEND_SPLIT_AXIS_2: + case GGML_BACKEND_SPLIT_AXIS_3: { + GGML_ASSERT(!ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0])); + if (src_ss[0].axis == ggml_n_dims(tensor->src[0]) - 1) { + return {ggml_backend_meta_split_axis(ggml_n_dims(tensor) - 1), {0}, 1}; + } + std::vector base_ne_in; + base_ne_in.reserve(GGML_MAX_DIMS - src_ss[0].axis); + { + base_ne_in.push_back(1); + int dim = 0; + for (; dim <= src_ss[0].axis; dim++) { + base_ne_in[0] *= tensor->src[0]->ne[dim]; + } + for (; dim <= GGML_MAX_DIMS; dim++) { + base_ne_in.push_back(base_ne_in.back() * tensor->src[0]->ne[dim]); + } + } + int64_t base_ne_out = 1; + for (int dim = 0; dim < GGML_MAX_DIMS; dim++) { + const int64_t base_ne_out_next = base_ne_out *= tensor->ne[dim]; + for (const int64_t & bni : base_ne_in) { + if (bni == base_ne_out_next) { + return {ggml_backend_meta_split_axis(dim), {0}, 1}; + } + } + if (base_ne_out_next > base_ne_in[0]) { + GGML_ASSERT(dim + 1 < GGML_MAX_DIMS); + return {ggml_backend_meta_split_axis(dim + 1), {0}, 1}; + } + base_ne_out = base_ne_out_next; + } + GGML_ABORT("shape mismatch for %s", ggml_op_name(tensor->op)); + } + case GGML_BACKEND_SPLIT_AXIS_MIRRORED: + case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { + return src_ss[0]; + } + default: { + GGML_ABORT("fatal error"); + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + } + } + }; + + auto handle_view = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (ggml_is_contiguous(tensor) && ggml_is_contiguous(tensor->src[0])) { + return handle_reshape(src_ss); + } + const int axis = src_ss[0].axis; + { + bool all_strides_the_same = true; + for (int dim = 0; dim < GGML_MAX_DIMS; dim++) { + if (tensor->ne[dim] == 1 && tensor->src[0]->ne[dim] == 1) { + continue; + } + if (tensor->nb[dim] != tensor->src[0]->nb[dim]) { + all_strides_the_same = false; + break; + } + } + if (all_strides_the_same) { + return src_ss[0]; + } + } + if (!ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0]) && axis >= 0 && axis < GGML_MAX_DIMS-1) { + for (int dim = 0; dim < GGML_MAX_DIMS-1; dim++) { + if (tensor->nb[dim+1] == tensor->src[0]->nb[axis+1]) { + return {ggml_backend_meta_split_axis(dim), {0}, 1}; + } + } + GGML_ABORT("fatal error"); + } + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED || src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { + return src_ss[0]; + } + GGML_ABORT("view of permuted tensor not implemented"); + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + }; + + auto handle_permute = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + switch (src_ss[0].axis) { + case GGML_BACKEND_SPLIT_AXIS_0: + case GGML_BACKEND_SPLIT_AXIS_1: + case GGML_BACKEND_SPLIT_AXIS_2: + case GGML_BACKEND_SPLIT_AXIS_3: { + return {ggml_backend_meta_split_axis(tensor->op_params[src_ss[0].axis]), {0}, 1}; + } + case GGML_BACKEND_SPLIT_AXIS_MIRRORED: + case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { + return src_ss[0]; + } + default: { + GGML_ABORT("fatal error"); + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + } + } + }; + + auto handle_transpose = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + switch (src_ss[0].axis) { + case GGML_BACKEND_SPLIT_AXIS_0: + case GGML_BACKEND_SPLIT_AXIS_1: { + return {ggml_backend_meta_split_axis(int(src_ss[0].axis) ^ 1), {0}, 1}; + } + case GGML_BACKEND_SPLIT_AXIS_2: + case GGML_BACKEND_SPLIT_AXIS_3: + case GGML_BACKEND_SPLIT_AXIS_MIRRORED: + case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { + return src_ss[0]; + } + default: { + GGML_ABORT("fatal error"); + //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + } + } + }; + + auto handle_get_rows = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + return src_ss[0]; + } + return handle_generic(src_ss, /*scalar_only =*/ true); + }; + + auto handle_set_rows = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(split_states_equal(src_ss[0], src_ss[2])); + return src_ss[0]; + }; + + auto handle_rope = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + return src_ss[0]; + }; + + auto handle_pad = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { + GGML_ASSERT(tensor->op_params[2*src_ss[0].axis + 0] == 0); + GGML_ASSERT(tensor->op_params[2*src_ss[0].axis + 1] == 0); + } + return src_ss[0]; + }; + + auto handle_flash_attn_ext = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + GGML_ASSERT( src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2); + GGML_ASSERT( src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2); + GGML_ASSERT( src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2); + GGML_ASSERT(tensor->src[4] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0); + return {GGML_BACKEND_SPLIT_AXIS_1, {0}, 1}; + }; + + auto handle_ssm_conv = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (src_ss[0].axis == src_ss[1].axis) { + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { + return {GGML_BACKEND_SPLIT_AXIS_1, {0}, 1}; + } + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1) { + return {GGML_BACKEND_SPLIT_AXIS_0, {0}, 1}; + } + } + return handle_generic(src_ss, /*scalar_only =*/ false); + }; + + auto handle_gated_delta_net = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + return src_ss[0]; + } + GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_1); + GGML_ASSERT(src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_2); + return {GGML_BACKEND_SPLIT_AXIS_0, {0}, 1}; + }; + + auto calculate_split_state = [&]() -> ggml_backend_meta_split_state { + if (ggml_nelements(tensor) == 0) { + return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + } + if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE && tensor->view_src == nullptr) { + ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); + const ggml_backend_meta_device_context * dev_ctx = (const ggml_backend_meta_device_context *) dev->context; + ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor, dev_ctx->get_split_state_ud); + if (ret.axis >= 0 && ret.axis <= GGML_MAX_DIMS) { + const int64_t granularity = ret.axis == GGML_BACKEND_SPLIT_AXIS_0 ? ggml_blck_size(tensor->type) : 1; + int64_t ne_sum = 0; + for (size_t sj = 0; sj < ret.n_segments*n_bufs; sj++) { + GGML_ASSERT(ret.ne[sj] % granularity == 0); + ne_sum += ret.ne[sj]; + } + GGML_ASSERT(ne_sum == tensor->ne[ret.axis]); + } + return ret; + } + + std::vector src_ss(GGML_MAX_SRC, {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, 1}); + for (size_t i = 0; i < GGML_MAX_SRC; i++) { + if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { + src_ss[i] = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + continue; + } + src_ss[i] = ggml_backend_meta_get_split_state(tensor->src[i], /*assume_sync =*/ true); + GGML_ASSERT(src_ss[i].axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); + } + + ggml_backend_meta_split_state split_state; + switch (tensor->op) { + case GGML_OP_NONE: { + split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, 1}; + } break; + case GGML_OP_DUP: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_ADD: + case GGML_OP_ADD_ID: { + split_state = handle_bin_bcast(src_ss); + } break; + case GGML_OP_ADD1: + case GGML_OP_ACC: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_SUB: + case GGML_OP_MUL: + case GGML_OP_DIV: { + split_state = handle_bin_bcast(src_ss); + } break; + case GGML_OP_SQR: + case GGML_OP_SQRT: + case GGML_OP_LOG: + case GGML_OP_SIN: + case GGML_OP_COS: { + split_state = handle_generic(src_ss, /*scalar_only =*/ false); + } break; + case GGML_OP_SUM: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_SUM_ROWS: + case GGML_OP_CUMSUM: + case GGML_OP_MEAN: + case GGML_OP_ARGMAX: + case GGML_OP_COUNT_EQUAL: { + split_state = handle_per_row(src_ss); + } break; + case GGML_OP_REPEAT: + case GGML_OP_REPEAT_BACK: { + split_state = handle_generic(src_ss, /*scalar_only =*/ false); + } break; + case GGML_OP_CONCAT: { + split_state = handle_concat(src_ss); + } break; + case GGML_OP_SILU_BACK: { + split_state = handle_generic(src_ss, /*scalar_only =*/ false); + } break; + case GGML_OP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_RMS_NORM_BACK: + case GGML_OP_GROUP_NORM: + case GGML_OP_L2_NORM: { + split_state = handle_per_row(src_ss); + } break; + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: { + split_state = handle_mul_mat(src_ss); + } break; + case GGML_OP_OUT_PROD: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_SCALE: { + split_state = handle_generic(src_ss, /*scalar_only =*/ false); + } break; + case GGML_OP_SET: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_CPY: { + split_state = handle_cpy(src_ss); + } break; + case GGML_OP_CONT: + case GGML_OP_RESHAPE: { + split_state = handle_reshape(src_ss); + } break; + case GGML_OP_VIEW: { + split_state = handle_view(src_ss); + } break; + case GGML_OP_PERMUTE: { + split_state = handle_permute(src_ss); + } break; + case GGML_OP_TRANSPOSE: { + split_state = handle_transpose(src_ss); + } break; + case GGML_OP_GET_ROWS: { + split_state = handle_get_rows(src_ss); + } break; + case GGML_OP_GET_ROWS_BACK: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_SET_ROWS: { + split_state = handle_set_rows(src_ss); + } break; + case GGML_OP_DIAG: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_DIAG_MASK_ZERO: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_SOFT_MAX: + case GGML_OP_SOFT_MAX_BACK: { + split_state = handle_generic(src_ss, /*scalar_only =*/ false); + } break; + case GGML_OP_ROPE: { + split_state = handle_rope(src_ss); + } break; + case GGML_OP_ROPE_BACK: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_CLAMP: { + split_state = handle_generic(src_ss, /*scalar_only =*/ false); + } break; + case GGML_OP_CONV_TRANSPOSE_1D: + case GGML_OP_IM2COL: + case GGML_OP_IM2COL_BACK: + case GGML_OP_IM2COL_3D: + case GGML_OP_CONV_2D: + case GGML_OP_CONV_3D: + case GGML_OP_CONV_2D_DW: + case GGML_OP_CONV_TRANSPOSE_2D: + case GGML_OP_POOL_1D: + case GGML_OP_POOL_2D: + case GGML_OP_POOL_2D_BACK: + case GGML_OP_UPSCALE: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_PAD: { + split_state = handle_pad(src_ss); + } break; + case GGML_OP_PAD_REFLECT_1D: + case GGML_OP_ROLL: + case GGML_OP_ARANGE: + case GGML_OP_TIMESTEP_EMBEDDING: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_ARGSORT: + case GGML_OP_TOP_K: { + split_state = handle_per_row(src_ss); + } break; + case GGML_OP_LEAKY_RELU: { + split_state = handle_generic(src_ss, /*scalar_only =*/ false); + } break; + case GGML_OP_TRI: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_FILL: { + split_state = handle_generic(src_ss, /*scalar_only =*/ false); + } break; + case GGML_OP_FLASH_ATTN_EXT: { + split_state = handle_flash_attn_ext(src_ss); + } break; + case GGML_OP_FLASH_ATTN_BACK: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_SSM_CONV: { + split_state = handle_ssm_conv(src_ss); + } break; + case GGML_OP_SSM_SCAN: + case GGML_OP_WIN_PART: + case GGML_OP_WIN_UNPART: + case GGML_OP_GET_REL_POS: + case GGML_OP_ADD_REL_POS: + case GGML_OP_RWKV_WKV6: + case GGML_OP_GATED_LINEAR_ATTN: + case GGML_OP_RWKV_WKV7: + case GGML_OP_SOLVE_TRI: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_GATED_DELTA_NET: { + split_state = handle_gated_delta_net(src_ss); + } break; + case GGML_OP_UNARY: { + split_state = handle_generic(src_ss, /*scalar_only =*/ false); + } break; + case GGML_OP_MAP_CUSTOM1: + case GGML_OP_MAP_CUSTOM2: + case GGML_OP_MAP_CUSTOM3: + case GGML_OP_CUSTOM: { + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; + case GGML_OP_CROSS_ENTROPY_LOSS: + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: { + split_state = handle_per_row(src_ss); + } break; + case GGML_OP_OPT_STEP_ADAMW: + case GGML_OP_OPT_STEP_SGD: + case GGML_OP_GLU: { + split_state = handle_generic(src_ss, /*scalar_only =*/ false); + } break; + default: { + GGML_ABORT("ggml op not implemented: %s", ggml_op_name(tensor->op)); + split_state = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, 1}; + } break; + } + if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS) { + bool first_src_split_by_axis = true; + const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); + + for (size_t i = 0; i < GGML_MAX_SRC; i++) { + if (tensor->src[i] == nullptr || src_ss[i].axis < 0 || src_ss[i].axis >= GGML_MAX_DIMS) { + continue; + } + if (first_src_split_by_axis) { + for (size_t j = 0; j < n_bufs; j++) { + // Take over ratio from src: + for (size_t s = 0; s < src_ss[i].n_segments; s++) { + split_state.ne[s*n_bufs + j] = 0; + } + for (size_t s = 0; s < src_ss[i].n_segments; s++) { + split_state.ne[j] += src_ss[i].ne[s*n_bufs + j]; + } + split_state.ne[j] *= tensor->ne[split_state.axis]; + if (split_state.ne[j] != 0 || tensor->src[i]->ne[src_ss[i].axis] != 0) { + GGML_ASSERT(split_state.ne[j] % tensor->src[i]->ne[src_ss[i].axis] == 0); + split_state.ne[j] /= tensor->src[i]->ne[src_ss[i].axis]; + } + } + } else { + for (size_t j = 0; j < n_bufs; j++) { + int64_t sum = 0; + for (size_t s = 0; s < src_ss[i].n_segments; s++) { + sum += src_ss[i].ne[s*n_bufs + j]; + } + // Assert that ratio is consistent: + GGML_ASSERT(split_state.ne[j] * tensor->src[i]->ne[src_ss[i].axis] + == sum * tensor->ne[split_state.axis]); + } + } + first_src_split_by_axis = false; + } + GGML_ASSERT(!first_src_split_by_axis); + } + return split_state; + }; + + const std::pair key = std::make_pair(tensor, assume_sync); + auto it = buf_ctx->split_state_cache.find(key); + if (it != buf_ctx->split_state_cache.end() && memcmp(it->second.second, (const char *) tensor, sizeof(it->second.second)) != 0) { + buf_ctx->split_state_cache.clear(); + it = buf_ctx->split_state_cache.end(); + } + + if (it == buf_ctx->split_state_cache.end()) { + buf_ctx->split_state_cache[key].first = calculate_split_state(); + memcpy(buf_ctx->split_state_cache[key].second, tensor, sizeof(buf_ctx->split_state_cache[key].second)); + if (buf_ctx->debug > 0) { + std::string srcs_info; + for (size_t i = 0; i < GGML_MAX_SRC; i++) { + if (tensor->src[i] == nullptr) { + continue; + } + if (!srcs_info.empty()) { + srcs_info += ", "; + } + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor->src[0], true); + const char * axis_name = ggml_backend_meta_split_axis_name(split_state.axis); + std::string ne_info; + for (size_t j = 0; j < n_bufs; j++) { + if (!ne_info.empty()) { + ne_info += ", "; + } + ne_info += std::to_string(split_state.ne[j]); + } + srcs_info += std::string(tensor->src[i]->name) + "[" + ggml_op_name(tensor->src[i]->op) + ", " + axis_name + ", {" + ne_info + "}]"; + } + std::string ne_info; + for (size_t j = 0; j < n_bufs; j++) { + if (!ne_info.empty()) { + ne_info += ", "; + } + ne_info += std::to_string(buf_ctx->split_state_cache[key].first.ne[j]); + } + GGML_LOG_DEBUG("SPLIT_STATE: {%s} -> %s[%s, %s, {%s}]\n", srcs_info.c_str(), tensor->name, ggml_op_name(tensor->op), + ggml_backend_meta_split_axis_name(buf_ctx->split_state_cache[key].first.axis), ne_info.c_str()); + } + } + + ggml_backend_meta_split_state ret = buf_ctx->split_state_cache[key].first; + GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_NONE); +#ifndef NDEBUG + if (ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { + int64_t ne_ret = 0; + for (size_t sj = 0; sj < ret.n_segments*n_bufs; sj++) { + ne_ret += ret.ne[sj]; + } + assert(ne_ret == tensor->ne[int(ret.axis)]); + } +#endif // NDEBUG + return ret; +} + +static void * ggml_backend_meta_buffer_get_base(ggml_backend_buffer_t buffer) { + GGML_UNUSED(buffer); + return (void *) 0x1000000000000000; // FIXME +} + +static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + GGML_ASSERT(ggml_backend_buffer_is_meta(buffer)); + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context; + const size_t n_simple_bufs = ggml_backend_meta_buffer_n_bufs(buffer); + + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ true); + GGML_ASSERT(ggml_nelements(tensor) == 0 || split_state.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); + GGML_ASSERT(split_state.n_segments <= 16); + + int split_dim = split_state.axis; + int64_t ne[GGML_MAX_DIMS]; + size_t nb[GGML_MAX_DIMS]; + for (size_t k = 0; k < GGML_MAX_DIMS; k++) { + ne[k] = tensor->ne[k]; + nb[k] = tensor->nb[k]; + } + + std::vector simple_tensors; + simple_tensors.reserve(n_simple_bufs); + for (size_t j = 0; j < n_simple_bufs; j++) { + ggml_context * simple_ctx = buf_ctx->buf_configs[j].ctx; + ggml_backend_buffer_t simple_buf = buf_ctx->buf_configs[j].buf; + + if (split_dim >= 0 && split_dim < GGML_MAX_DIMS) { + // TODO: the following assert fails for llama-parallel even though the results are correct: + // GGML_ASSERT(ggml_is_contiguously_allocated(tensor)); + ne[split_dim] = 0; + for (size_t s = 0; s < split_state.n_segments; s++) { + ne[split_dim] += split_state.ne[s*n_simple_bufs + j]; + } + for (int i = 0; i < GGML_MAX_DIMS; i++) { + if (tensor->nb[i] > tensor->nb[split_dim]) { + nb[i] = tensor->nb[i] * ne[split_dim]/tensor->ne[split_dim]; + } + } + } + + ggml_tensor * t_ij = ggml_new_tensor(simple_ctx, tensor->type, GGML_MAX_DIMS, ne); + t_ij->op = tensor->op; + for (int i = 0; i < GGML_MAX_DIMS; i++) { + t_ij->nb[i] = nb[i]; + } + t_ij->flags = tensor->flags; + memcpy(t_ij->op_params, tensor->op_params, sizeof(tensor->op_params)); + ggml_set_name(t_ij, tensor->name); + t_ij->buffer = simple_buf; + t_ij->view_src = tensor->view_src; + t_ij->view_offs = tensor->view_offs; + if (t_ij->view_src != nullptr && ggml_backend_buffer_is_meta(t_ij->view_src->buffer)) { + t_ij->view_src = ggml_backend_meta_buffer_simple_tensor(tensor->view_src, j); + if (t_ij->view_offs > 0 && split_dim >= 0 && split_dim < GGML_MAX_DIMS) { + GGML_ASSERT(tensor->ne[split_dim] != 0); + const int split_dim_view_src = ggml_backend_meta_get_split_state(tensor->view_src, /*assume_sync =*/ true).axis; + GGML_ASSERT(split_dim_view_src >= 0 && split_dim_view_src < GGML_MAX_DIMS); + + // The offset can be internal to the data split, in those cases the view offset should not be scaled. + // If however, the offset is larger than the data split then it needs to be scaled proportionally. + bool split_internal_offset = t_ij->view_offs <= tensor->view_src->nb[split_dim_view_src]; + for (int i = 0; i < GGML_MAX_DIMS; i++) { + const size_t dim_size = tensor->ne[i] * tensor->nb[i]; + if (tensor->view_offs <= dim_size && dim_size < tensor->nb[split_dim]) { + split_internal_offset = true; + break; + } + } + if (!split_internal_offset) { + t_ij->view_offs = t_ij->view_offs * ne[split_dim]/tensor->ne[split_dim]; + } + } + } + if (t_ij->view_src != nullptr) { + t_ij->data = (char *) t_ij->view_src->data + t_ij->view_offs; + } else if (simple_buf != nullptr) { + t_ij->data = (char *) ggml_backend_buffer_get_base(simple_buf) + + size_t(tensor->data) - size_t(ggml_backend_buffer_get_base(buffer)); + } + t_ij->extra = tensor->extra; + for (int i = 0; i < GGML_MAX_SRC; i++) { + t_ij->src[i] = tensor->src[i]; + if (tensor->src[i] == tensor) { + t_ij->src[i] = t_ij; + } else if (t_ij->src[i] != nullptr && ggml_backend_buffer_is_meta(t_ij->src[i]->buffer)) { + t_ij->src[i] = ggml_backend_meta_buffer_simple_tensor(tensor->src[i], j); + } + } + + simple_tensors.push_back(t_ij); + } + + // If one of the sources has a zero-sized slice, disable the computation: + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (tensor->src[i] == nullptr || !ggml_backend_buffer_is_meta(tensor->src[i]->buffer)) { + continue; + } + + const ggml_backend_meta_split_state split_state_src = ggml_backend_meta_get_split_state(tensor->src[i], /*assume_sync =*/ true); + if (split_state_src.axis < 0 || split_state_src.axis >= GGML_MAX_DIMS) { + continue; + } + for (size_t j = 0; j < n_simple_bufs; j++) { + int64_t ne_sum = 0; + for (size_t s = 0; s < split_state_src.n_segments; s++) { + ne_sum += split_state_src.ne[s*n_simple_bufs + j]; + } + if (ne_sum == 0) { + simple_tensors[j]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; + } + } + } + + buf_ctx->simple_tensors[tensor] = simple_tensors; + + return GGML_STATUS_SUCCESS; +} + +static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); + GGML_ASSERT(ggml_is_contiguous(tensor)); + + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + + if (split_state.n_segments != 1) { + GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS); + GGML_ASSERT(tensor->ne[3] == 1); + + size_t offset_data = 0; + std::vector simple_offsets(n_bufs, 0); + if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) { + GGML_ASSERT(tensor->ne[2] == 1); + + const size_t row_stride = tensor->nb[1]; + GGML_ASSERT(offset % row_stride == 0); + GGML_ASSERT(size % row_stride == 0); + const int64_t r_start = offset / row_stride; + const int64_t r_count = size / row_stride; + GGML_ASSERT(r_start + r_count <= tensor->ne[1]); + + const int64_t blck_size = ggml_blck_size(tensor->type); + for (size_t s = 0; s < split_state.n_segments; s++) { + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0); + const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0]; + ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, + simple_offsets[j] + r_start * simple_tensor->nb[1], nbytes, + r_count, simple_tensor->nb[1], tensor->nb[1]); + offset_data += nbytes; + simple_offsets[j] += nbytes; + } + } + GGML_ASSERT(offset_data*r_count == size); + return; + } + GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1); + + const size_t row_stride = tensor->nb[2]; + GGML_ASSERT(offset % row_stride == 0); + GGML_ASSERT(size % row_stride == 0); + const int64_t r_start = offset / row_stride; + const int64_t r_count = size / row_stride; + GGML_ASSERT(r_start + r_count <= tensor->ne[2]); + + for (size_t s = 0; s < split_state.n_segments; s++) { + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[1]; + ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, + simple_offsets[j] + r_start * simple_tensor->nb[2], nbytes, + r_count, simple_tensor->nb[2], tensor->nb[2]); + offset_data += nbytes; + simple_offsets[j] += nbytes; + } + } + GGML_ASSERT(offset_data*r_count == size); + return; + } + + switch (split_state.axis) { + case GGML_BACKEND_SPLIT_AXIS_0: + case GGML_BACKEND_SPLIT_AXIS_1: + case GGML_BACKEND_SPLIT_AXIS_2: { + // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". + const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; + GGML_ASSERT(offset % chunk_size_full == 0); + GGML_ASSERT(size % chunk_size_full == 0); + const int64_t i_start = offset /chunk_size_full; + const int64_t i_stop = (offset + size)/chunk_size_full; + size_t offset_j = 0; + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; + const size_t simple_offset = i_start * chunk_size_j; + ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_j, simple_offset, chunk_size_j, i_stop - i_start, chunk_size_j, chunk_size_full); + offset_j += chunk_size_j; + } + GGML_ASSERT(offset_j == chunk_size_full); + } break; + case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + ggml_backend_tensor_set(simple_tensor, data, offset, size); + } + } break; + case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { + GGML_ASSERT(tensor->type == GGML_TYPE_F32); + const int64_t ne = ggml_nelements(tensor); + std::vector tmp; + tmp.reserve(ne); + for (int64_t i = 0; i < ne; i++) { + tmp.push_back(((const float *) data)[i] / n_bufs); + } + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + ggml_backend_tensor_set(simple_tensor, tmp.data(), offset, size); + } + } break; + default: { + GGML_ABORT("fatal error"); + } + } +} + +static void ggml_backend_meta_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); + GGML_ASSERT(ggml_is_contiguous(tensor)); + + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + + if (split_state.n_segments != 1) { + GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS); + GGML_ASSERT(tensor->ne[3] == 1); + + size_t offset_data = 0; + std::vector simple_offsets(n_bufs, 0); + if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) { + GGML_ASSERT(tensor->ne[2] == 1); + + const size_t row_stride = tensor->nb[1]; + GGML_ASSERT(offset % row_stride == 0); + GGML_ASSERT(size % row_stride == 0); + const int64_t r_start = offset / row_stride; + const int64_t r_count = size / row_stride; + GGML_ASSERT(r_start + r_count <= tensor->ne[1]); + + const int64_t blck_size = ggml_blck_size(tensor->type); + for (size_t s = 0; s < split_state.n_segments; s++) { + for (size_t j = 0; j < n_bufs; j++) { + const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0); + const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0]; + ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, + simple_offsets[j] + r_start * simple_tensor->nb[1], nbytes, + r_count, simple_tensor->nb[1], tensor->nb[1]); + offset_data += nbytes; + simple_offsets[j] += nbytes; + } + } + GGML_ASSERT(offset_data*r_count == size); + return; + } + GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1); + + const size_t row_stride = tensor->nb[2]; + GGML_ASSERT(offset % row_stride == 0); + GGML_ASSERT(size % row_stride == 0); + const int64_t r_start = offset / row_stride; + const int64_t r_count = size / row_stride; + GGML_ASSERT(r_start + r_count <= tensor->ne[2]); + + for (size_t s = 0; s < split_state.n_segments; s++) { + for (size_t j = 0; j < n_bufs; j++) { + const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[1]; + ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, + simple_offsets[j] + r_start * simple_tensor->nb[2], nbytes, + r_count, simple_tensor->nb[2], tensor->nb[2]); + offset_data += nbytes; + simple_offsets[j] += nbytes; + } + } + GGML_ASSERT(offset_data*r_count == size); + return; + } + + switch (split_state.axis) { + case GGML_BACKEND_SPLIT_AXIS_0: + case GGML_BACKEND_SPLIT_AXIS_1: + case GGML_BACKEND_SPLIT_AXIS_2: { + // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". + const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; + GGML_ASSERT(offset % chunk_size_full == 0); + GGML_ASSERT(size % chunk_size_full == 0); + const int64_t i_start = offset /chunk_size_full; + const int64_t i_stop = (offset + size)/chunk_size_full; + size_t offset_j = 0; + for (size_t j = 0; j < n_bufs; j++){ + const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; + const size_t simple_offset = i_start * chunk_size_j; + ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_j, simple_offset, chunk_size_j, i_stop - i_start, chunk_size_j, chunk_size_full); + offset_j += chunk_size_j; + } + GGML_ASSERT(offset_j == chunk_size_full); + } break; + case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { + // TODO other simple backend may be better + const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, 0); + ggml_backend_tensor_get(simple_tensor, data, offset, size); + } break; + default: { + GGML_ABORT("fatal error"); + } + } +} + +static void ggml_backend_meta_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + const size_t n_buffers = ggml_backend_meta_buffer_n_bufs(buffer); + for (size_t i = 0; i < n_buffers; i++) { + ggml_backend_buffer_clear(ggml_backend_meta_buffer_simple_buffer(buffer, i), value); + } +} + +static void ggml_backend_meta_buffer_reset(ggml_backend_buffer_t buffer) { + const size_t n_buffers = ggml_backend_meta_buffer_n_bufs(buffer); + for (size_t i = 0; i < n_buffers; i++) { + ggml_backend_buffer_reset(ggml_backend_meta_buffer_simple_buffer(buffer, i)); + } +} + +static const ggml_backend_buffer_i ggml_backend_meta_buffer_iface = { + /* .free_buffer = */ ggml_backend_meta_buffer_free_buffer, + /* .get_base = */ ggml_backend_meta_buffer_get_base, + /* .init_tensor = */ ggml_backend_meta_buffer_init_tensor, + /* .memset_tensor = */ nullptr, // TODO implement + /* .set_tensor = */ ggml_backend_meta_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_meta_buffer_get_tensor, + /* .set_tensor_2d = */ nullptr, + /* .get_tensor_2d = */ nullptr, + /* .cpy_tensor = */ nullptr, + /* .clear = */ ggml_backend_meta_buffer_clear, + /* .reset = */ ggml_backend_meta_buffer_reset, +}; + +bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf) { + return buf != nullptr && buf->iface.free_buffer == ggml_backend_meta_buffer_iface.free_buffer; +} + +static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); + + ggml_init_params params = { + /*.mem_size =*/ 1024*1024*1024, // FIXME + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + ggml_backend_meta_buffer_context * buf_ctx = new ggml_backend_meta_buffer_context(); + size_t max_size = 0; + buf_ctx->buf_configs.reserve(n_simple_bufts); + for (size_t i = 0; i < n_simple_bufts; i++) { + ggml_backend_buffer_t simple_buf = ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size); + max_size = std::max(max_size, ggml_backend_buffer_get_size(simple_buf)); + buf_ctx->buf_configs.emplace_back(ggml_init(params), simple_buf); + } + + return ggml_backend_buffer_init(buft, ggml_backend_meta_buffer_iface, buf_ctx, max_size); +} + +struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft) { + const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); + + ggml_init_params params = { + /*.mem_size =*/ 1024*1024*1024, // FIXME + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + ggml_backend_meta_buffer_context * meta_buf_ctx = new ggml_backend_meta_buffer_context(); + meta_buf_ctx->buf_configs.reserve(n_simple_bufts); + for (size_t i = 0; i < n_simple_bufts; i++) { + meta_buf_ctx->buf_configs.emplace_back(ggml_init(params), nullptr); + } + + ggml_backend_buffer_t meta_buf = ggml_backend_buffer_init(buft, ggml_backend_meta_buffer_iface, meta_buf_ctx, 0); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + t->buffer = meta_buf; + ggml_backend_meta_buffer_init_tensor(meta_buf, t); + t->data = (void *) 0x2000000000000000; // FIXME + } + for (size_t i = 0; i < n_simple_bufts; i++) { + meta_buf_ctx->buf_configs[i].buf = ggml_backend_alloc_ctx_tensors_from_buft( + meta_buf_ctx->buf_configs[i].ctx, ggml_backend_meta_buft_simple_buft(buft, i)); + meta_buf->size = std::max(meta_buf->size, ggml_backend_buffer_get_size(meta_buf_ctx->buf_configs[i].buf)); + } + return meta_buf; +} + +// +// meta backend +// + +static ggml_guid_t ggml_backend_meta_guid() { + static ggml_guid guid = {0xf1, 0x0e, 0x34, 0xcf, 0x9c, 0x6f, 0x43, 0xcb, 0x96, 0x92, 0xbe, 0x8e, 0xbb, 0x71, 0x3f, 0xda}; + return &guid; +} + +struct ggml_backend_meta_context { + struct cgraph_config { + ggml_cgraph * cgraph_main = nullptr; + int offset = 0; // Node offset vs. original graph + + std::vector cgraphs_aux; + }; + struct backend_config { + ggml_backend_t backend; + + std::vector cgraphs; + std::vector nodes; + std::vector bufs; + + backend_config(ggml_backend_t backend, const size_t n_reduce_steps) : backend(backend) { + bufs.resize(n_reduce_steps); + } + }; + std::string name; + std::vector backend_configs; + ggml_context_ptr ctx; + std::vector cgraphs_aux; + std::vector nodes_aux; + size_t n_reduce_steps; + int max_nnodes = 0; + size_t max_tmp_size = 0; + size_t max_subgraphs = 0; + size_t n_subgraphs = 0; + uint64_t uid = 0; + + void * comm_ctx = nullptr; + ggml_backend_comm_allreduce_tensor_t comm_allreduce = nullptr; + + ggml_backend_meta_context(ggml_backend_dev_t meta_dev, const char * params) { + const size_t n_devs = ggml_backend_meta_dev_n_devs(meta_dev); + n_reduce_steps = std::ceil(std::log2(n_devs)); + name = "Meta("; + std::vector simple_backends; + backend_configs.reserve(n_devs); + simple_backends.reserve(n_devs); + for (size_t i = 0; i < n_devs; i++) { + ggml_backend_dev_t simple_dev = ggml_backend_meta_dev_simple_dev(meta_dev, i); + if (i > 0) { + name += ","; + } + name += ggml_backend_dev_name(simple_dev); + simple_backends.push_back(ggml_backend_dev_init(simple_dev, params)); + backend_configs.emplace_back(simple_backends.back(), n_reduce_steps); + } + name += ")"; + + if (n_devs > 1) { + ggml_backend_comm_init_t comm_init = (ggml_backend_comm_init_t) ggml_backend_reg_get_proc_address( + ggml_backend_dev_backend_reg(ggml_backend_get_device(simple_backends[0])), "ggml_backend_comm_init"); + if (comm_init != nullptr) { + comm_ctx = comm_init(simple_backends.data(), simple_backends.size()); + } + } + if (comm_ctx != nullptr) { + comm_allreduce = (ggml_backend_comm_allreduce_tensor_t) + ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg( + ggml_backend_get_device(simple_backends[0])), "ggml_backend_comm_allreduce_tensor"); + GGML_ASSERT(comm_allreduce != nullptr); + } + } + + ~ggml_backend_meta_context() { + if (comm_ctx != nullptr) { + ggml_backend_comm_free_t comm_free = (ggml_backend_comm_free_t) ggml_backend_reg_get_proc_address( + ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_configs[0].backend)), "ggml_backend_comm_free"); + GGML_ASSERT(comm_free != nullptr); + comm_free(comm_ctx); + } + for (auto & bc : backend_configs) { + ggml_backend_free(bc.backend); + } + } +}; + +static const char * ggml_backend_meta_get_name(ggml_backend_t backend) { + GGML_ASSERT(ggml_backend_is_meta(backend)); + const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) backend->context; + return backend_ctx->name.c_str(); +} + +static void ggml_backend_meta_free(ggml_backend_t backend) { + GGML_ASSERT(ggml_backend_is_meta(backend)); + ggml_backend_meta_context * backend_ctx = (ggml_backend_meta_context *) backend->context; + delete backend_ctx; + delete backend; +} + +static void ggml_backend_meta_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + const size_t n_backends = ggml_backend_meta_n_backends(backend); + GGML_ASSERT(offset == 0); + GGML_ASSERT(ggml_is_contiguous(tensor)); + + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + GGML_ASSERT(split_state.n_segments == 1); + + switch (split_state.axis) { + case GGML_BACKEND_SPLIT_AXIS_0: + case GGML_BACKEND_SPLIT_AXIS_1: + case GGML_BACKEND_SPLIT_AXIS_2: { + // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". + const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; + GGML_ASSERT(offset % chunk_size_full == 0); + GGML_ASSERT(size % chunk_size_full == 0); + const int64_t i_start = offset /chunk_size_full; + const int64_t i_stop = (offset + size)/chunk_size_full; + size_t offset_j = 0; + for (size_t j = 0; j < n_backends; j++){ + ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, j); + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; + ggml_backend_tensor_set_2d_async(simple_backend, simple_tensor, (const char *) data + offset_j, offset, chunk_size_j, + i_stop - i_start, chunk_size_j, chunk_size_full); + offset_j += chunk_size_j; + } + GGML_ASSERT(offset_j == chunk_size_full); + } break; + case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { + for (size_t j = 0; j < n_backends; j++) { + ggml_backend_tensor_set_async( + ggml_backend_meta_simple_backend(backend, j), ggml_backend_meta_buffer_simple_tensor(tensor, j), data, offset, size); + } + } break; + default: { + GGML_ABORT("fatal error"); + } + } +} + +static void ggml_backend_meta_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + const size_t n_backends = ggml_backend_meta_n_backends(backend); + GGML_ASSERT(offset == 0); + GGML_ASSERT(ggml_is_contiguous(tensor)); + + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + GGML_ASSERT(split_state.n_segments == 1); + + switch (split_state.axis) { + case GGML_BACKEND_SPLIT_AXIS_0: + case GGML_BACKEND_SPLIT_AXIS_1: + case GGML_BACKEND_SPLIT_AXIS_2: { + // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". + const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; + GGML_ASSERT(offset % chunk_size_full == 0); + GGML_ASSERT(size % chunk_size_full == 0); + const int64_t i_start = offset /chunk_size_full; + const int64_t i_stop = (offset + size)/chunk_size_full; + size_t offset_j = 0; + for (size_t j = 0; j < n_backends; j++){ + ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, j); + const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; + ggml_backend_tensor_get_2d_async(simple_backend, simple_tensor, (char *) data + offset_j, offset, chunk_size_j, + i_stop - i_start, chunk_size_j, chunk_size_full); + offset_j += chunk_size_j; + } + GGML_ASSERT(offset_j == chunk_size_full); + } break; + case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { + // TODO other simple backend may be better + ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, 0); + const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, 0); + ggml_backend_tensor_get_async(simple_backend, simple_tensor, data, offset, size); + } break; + default: { + GGML_ABORT("fatal error"); + } + } +} + +static void ggml_backend_meta_synchronize(ggml_backend_t backend) { + const size_t n_backends = ggml_backend_meta_n_backends(backend); + for (size_t i = 0; i < n_backends; i++) { + ggml_backend_synchronize(ggml_backend_meta_simple_backend(backend, i)); + } +} + +static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, struct ggml_cgraph * cgraph) { + GGML_ASSERT(cgraph->grads == nullptr); + const size_t n_backends = ggml_backend_meta_n_backends(backend); + ggml_backend_meta_context * backend_ctx = (ggml_backend_meta_context *) backend->context; + + // If the previous cgraph had a defined UID it can be used to skip rebuilding the subgraphs per simple backend. + const bool needs_rebuild = (cgraph->uid == 0) || (cgraph->uid != backend_ctx->uid); + + bool max_nnodes_raised = false; + if (cgraph->n_nodes > backend_ctx->max_nnodes) { + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + bcj.nodes.resize(cgraph->n_nodes); + bcj.cgraphs.resize(cgraph->n_nodes); + } + backend_ctx->max_nnodes = cgraph->n_nodes; + max_nnodes_raised = true; + assert(needs_rebuild); + } + + if (needs_rebuild) { + size_t n_subgraphs = 0; + size_t max_tmp_size = 0; + + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_tensor * node = cgraph->nodes[i]; + if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) { + // FIXME s_copy_main is on the CPU and its view seems to be incorrectly added to the graph nodes. + // For regular usage this doesn't matter since it's a noop but trying to call ggml_backend_meta_buffer_simple_tensor results in a crash. + bcj.nodes[i] = node; + continue; + } + bcj.nodes[i] = ggml_backend_meta_buffer_simple_tensor(node, j); + GGML_ASSERT(bcj.nodes[i]); + } + } + + { + // For MoE models it may make sense to delay the AllReduce in order to reduce I/O: + auto get_i_delayed = [&](const int i) -> int { + int id = i; // i_delayed + int idr = i; // i_delayed return, last safe return value + + ggml_tensor * node = cgraph->nodes[id]; + int32_t n_used = ggml_node_get_use_count(cgraph, id); + + // Skip MIRRORED nodes that don't consume node + auto skip_unrelated = [&]() { + while (id + 1 < cgraph->n_nodes) { + ggml_tensor * next = cgraph->nodes[id+1]; + if (ggml_backend_meta_get_split_state(next, false).axis != GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + break; + } + bool safe = true; + for (int s = 0; s < GGML_MAX_SRC; s++) { + if (next->src[s] == nullptr) { + continue; + } + if (next->src[s] == node) { + safe = false; + break; + } + if (ggml_backend_meta_get_split_state(next->src[s], false).axis != GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + safe = false; + break; + } + } + if (!safe) { + break; + } + id++; + } + }; + + skip_unrelated(); + if (id + 1 >= cgraph->n_nodes) { + return idr; + } + { + ggml_tensor * next = cgraph->nodes[id+1]; + if (next->op == GGML_OP_ADD_ID && next->src[0] == node && + ggml_backend_meta_get_split_state(next->src[1], false).axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL && + ggml_backend_meta_get_split_state(next->src[2], false).axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + node = next; + id++; + idr = id; + n_used = ggml_node_get_use_count(cgraph, id); + } + } + // Chain of MULs with MIRRORED src[1] + while (true) { + skip_unrelated(); + if (id + 1 >= cgraph->n_nodes) { + return idr; + } + ggml_tensor * next = cgraph->nodes[id+1]; + if (next->op == GGML_OP_MUL && next->src[0] == node && + ggml_backend_meta_get_split_state(next->src[1], false).axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + node = next; + id++; + idr = id; + n_used = ggml_node_get_use_count(cgraph, id); + } else { + break; + } + } + + if (n_used != node->ne[1] || id + 2*n_used-1 >= cgraph->n_nodes) { + return idr; + } + for (int32_t k = 0; k < n_used; k++) { + ggml_tensor * next = cgraph->nodes[id+1]; + if (next->op != GGML_OP_VIEW || next->view_src != node || next->view_offs != k*node->nb[1] || + next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] || next->nb[1] != node->nb[2] || + ggml_node_get_use_count(cgraph, id+1) != 1) { + return idr; + } + id++; + } + { + ggml_tensor * next = cgraph->nodes[id+1]; + if (next->op != GGML_OP_ADD || next->src[0] != cgraph->nodes[id - (n_used-1)] || + next->src[1] != cgraph->nodes[id - (n_used-2)] || ggml_node_get_use_count(cgraph, id+1) != 1) { + return idr; + } + id++; + } + for (int32_t k = 0; k < n_used - 2; k++) { + ggml_tensor * next = cgraph->nodes[id+1]; + if (next->op != GGML_OP_ADD || next->src[0] != cgraph->nodes[id] || + next->src[1] != cgraph->nodes[id - (n_used-2)] || ggml_node_get_use_count(cgraph, id+1) != 1) { + return idr; + } + id++; + } + idr = id; + return idr; + }; + + int i_start = 0; + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_tensor * node = cgraph->nodes[i]; + if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) { + continue; + } + const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(node, /*assume_sync =*/ false); + if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { + max_tmp_size = std::max(max_tmp_size, ggml_nbytes(node)); + } + const bool new_subgraph = i + 1 == cgraph->n_nodes || split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL; + if (!new_subgraph) { + continue; + } + + const int i_delayed = get_i_delayed(i); + + // If we can delay the AllReduce we need to consider the interaction with zero-sized tensor slices. + // A backend with such a slice would normally have valid data after participating in the AllReduce with a node that has + // its compute flag disabled and thus gets its data zeroed out. + // If the AllReduce is delayed then the nodes until that point also need to have their compute flag disabled. + if (i_delayed > i) { + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + if ((bcj.nodes[i]->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + for (int ii = i + 1; ii <= i_delayed; ii++) { + bcj.nodes[ii]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; + } + } + } + } + + i = i_delayed; + + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + bcj.cgraphs[n_subgraphs].offset = i_start; + } + n_subgraphs++; + i_start = i + 1; + } + GGML_ASSERT(i_start == cgraph->n_nodes); + } + + backend_ctx->uid = cgraph->uid; + backend_ctx->n_subgraphs = n_subgraphs; + + if (max_tmp_size > backend_ctx->max_tmp_size) { + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + for (size_t i = 0; i < backend_ctx->n_reduce_steps; i++) { + bcj.bufs[i].reset(ggml_backend_alloc_buffer(bcj.backend, max_tmp_size)); + } + } + backend_ctx->max_tmp_size = max_tmp_size; + } + + if (max_nnodes_raised || n_subgraphs > backend_ctx->max_subgraphs) { + backend_ctx->max_subgraphs = std::max(backend_ctx->max_subgraphs, n_subgraphs); + const size_t n_nodes_per_device = 3 * backend_ctx->n_reduce_steps; // tmp + ADD (+zeroing) graph per step and device + const size_t n_cgraphs_per_device = 2 * backend_ctx->n_reduce_steps; // ADD ( + zeroing) graph per step and device + const size_t mem_per_device_graphs_main = backend_ctx->max_subgraphs*ggml_graph_overhead_custom(backend_ctx->max_nnodes, cgraph->grads); + const size_t mem_per_device_graphs_aux = n_cgraphs_per_device*backend_ctx->max_subgraphs*ggml_graph_overhead_custom(1, cgraph->grads); + const size_t mem_per_device_nodes_aux = n_nodes_per_device*backend_ctx->max_subgraphs*ggml_tensor_overhead(); + ggml_init_params params = { + /*.mem_size =*/ n_backends * (mem_per_device_graphs_main + mem_per_device_graphs_aux + mem_per_device_nodes_aux), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + backend_ctx->ctx.reset(ggml_init(params)); + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + for (size_t i = 0; i < n_subgraphs; i++) { + bcj.cgraphs[i].cgraph_main = ggml_new_graph_custom(backend_ctx->ctx.get(), cgraph->n_nodes, /*grads =*/ false); + } + } + backend_ctx->cgraphs_aux.resize(n_backends*n_cgraphs_per_device*backend_ctx->max_subgraphs); + for (size_t k = 0; k < backend_ctx->cgraphs_aux.size(); k++) { + backend_ctx->cgraphs_aux[k] = ggml_new_graph_custom(backend_ctx->ctx.get(), 1, cgraph->grads); + } + backend_ctx->nodes_aux.resize(n_backends*n_nodes_per_device*backend_ctx->max_subgraphs); + for (size_t k = 0; k < backend_ctx->nodes_aux.size(); k++) { + backend_ctx->nodes_aux[k] = ggml_new_tensor_1d(backend_ctx->ctx.get(), GGML_TYPE_F32, 1); + } + } + + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + for (size_t i_graph = 0; i_graph < n_subgraphs; i_graph++) { + ggml_cgraph * cgraph_ij = bcj.cgraphs[i_graph].cgraph_main; + const size_t i_node_start = bcj.cgraphs[i_graph].offset; + const size_t i_node_stop = i_graph + 1 < n_subgraphs ? bcj.cgraphs[i_graph + 1].offset : cgraph->n_nodes; + cgraph_ij->n_nodes = i_node_stop - i_node_start; + ggml_hash_set_reset(&cgraph_ij->visited_hash_set); + for (size_t i_node = i_node_start; i_node < i_node_stop; i_node++) { + ggml_tensor * node_ij = bcj.nodes[i_node]; + cgraph_ij->nodes[i_node - i_node_start] = node_ij; + const size_t hash_pos_orig = ggml_hash_find(&cgraph->visited_hash_set, cgraph->nodes[i_node]); + const size_t hash_pos_ij = ggml_hash_insert(&cgraph_ij->visited_hash_set, node_ij); + cgraph_ij->use_counts[hash_pos_ij] = cgraph->use_counts[hash_pos_orig]; + } + cgraph_ij->uid = ggml_graph_next_uid(); + } + } + } + + size_t iga = 0; // i graph aux + size_t ina = 0; // i node aux + + auto get_node_aux = [&](ggml_tensor * t) -> ggml_tensor * { + ggml_tensor * ret = backend_ctx->nodes_aux[ina++]; + memset(ret, 0, sizeof(ggml_tensor)); + ret->op = GGML_OP_NONE; + ret->type = t->type; + for (size_t k = 0; k < GGML_MAX_DIMS; k++) { + ret->ne[k] = t->ne[k]; + ret->nb[k] = t->nb[k]; + } + return ret; + }; + auto set_tmp_data = [&](ggml_tensor * tensor, const size_t j, const size_t i_buf) { + auto & bcj = backend_ctx->backend_configs[j]; + ggml_backend_buffer_ptr & buf_ptr = bcj.bufs[i_buf]; + if (!buf_ptr || ggml_backend_buffer_get_size(buf_ptr.get()) < backend_ctx->max_tmp_size) { + buf_ptr.reset(ggml_backend_alloc_buffer(bcj.backend, backend_ctx->max_tmp_size)); + } + tensor->buffer = buf_ptr.get(); + tensor->data = ggml_backend_buffer_get_base(buf_ptr.get()); + }; + // FIXME usage_counts + auto get_cgraph_aux = [&]() -> ggml_cgraph * { + ggml_cgraph * ret = backend_ctx->cgraphs_aux[iga++]; + return ret; + }; + + // Preferentially use backend-specific allreduce_tensor_async (e.g. NCCL for CUDA), use a generic fallback if unavailable: + auto allreduce_fallback = [&](size_t i) -> ggml_status { + std::vector step_cgraphs(n_backends, nullptr); + + // Zero out nodes that were disabled due to having a zero-sized slice: + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + ggml_tensor * node = bcj.cgraphs[i].cgraph_main->nodes[bcj.cgraphs[i].cgraph_main->n_nodes - 1]; + if (node->flags & GGML_TENSOR_FLAG_COMPUTE) { + continue; + } + ggml_tensor * node_zero = get_node_aux(node); + node_zero->op = GGML_OP_SCALE; // FIXME 0.0f * NaN == NaN + node_zero->src[0] = node; + ggml_set_op_params_f32(node_zero, 0, 0.0f); + node_zero->data = node->data; + node_zero->flags |= GGML_TENSOR_FLAG_COMPUTE; + + step_cgraphs[j] = get_cgraph_aux(); + step_cgraphs[j]->nodes[0] = node_zero; + step_cgraphs[j]->n_nodes = 1; + const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, step_cgraphs[j]); + if (status != GGML_STATUS_SUCCESS) { + return status; + } + } + std::fill(step_cgraphs.begin(), step_cgraphs.end(), nullptr); + + auto push_data = [&](const size_t j_src, const size_t j_dst, const size_t i_buf) { + assert(step_cgraphs[j_dst] == nullptr); + auto & bcj_src = backend_ctx->backend_configs[j_src]; + auto & bcj_dst = backend_ctx->backend_configs[j_dst]; + + ggml_tensor * node_src = bcj_src.cgraphs[i].cgraph_main->nodes[bcj_src.cgraphs[i].cgraph_main->n_nodes - 1]; + ggml_tensor * node_dst = bcj_dst.cgraphs[i].cgraph_main->nodes[bcj_dst.cgraphs[i].cgraph_main->n_nodes - 1]; + GGML_ASSERT(ggml_is_contiguous(node_src)); + GGML_ASSERT(ggml_is_contiguous(node_dst)); + + ggml_tensor * node_tmp = get_node_aux(node_dst); + set_tmp_data(node_tmp, j_dst, i_buf); + + ggml_backend_tensor_copy_async(bcj_src.backend, bcj_dst.backend, node_src, node_tmp); + + ggml_tensor * node_red = get_node_aux(node_dst); + node_red->view_src = node_dst->view_src == nullptr ? node_dst : node_dst->view_src; + node_red->view_offs = node_dst->view_offs; + node_red->op = GGML_OP_ADD; + node_red->src[0] = node_dst; + node_red->src[1] = node_tmp; + node_red->flags |= GGML_TENSOR_FLAG_COMPUTE; + ggml_backend_view_init(node_red); + + ggml_cgraph * cgraph_aux = get_cgraph_aux(); + cgraph_aux->nodes[0] = node_red; + cgraph_aux->n_nodes = 1; + step_cgraphs[j_dst] = cgraph_aux; + }; + + size_t offset_j = n_backends/2; + while ((offset_j & (offset_j - 1)) != 0) { + offset_j--; + } + const size_t offset_j_max = offset_j; + size_t i_buf = 0; + + // If n_backends is not a power of 2, fold in the excess prior to butterfly reduction: + for (size_t j_src = 2*offset_j_max; j_src < n_backends; j_src++) { + const size_t j_dst = j_src - 2*offset_j_max; + push_data(j_src, j_dst, i_buf); + const ggml_status status = ggml_backend_graph_compute_async(backend_ctx->backend_configs[j_dst].backend, step_cgraphs[j_dst]); + if (status != GGML_STATUS_SUCCESS) { + return status; + } + i_buf = 1; + } + + // Butterfly reduction: + for (; offset_j >= 1; offset_j /= 2) { + std::fill(step_cgraphs.begin(), step_cgraphs.end(), nullptr); + + for (size_t j = 0; j < 2*offset_j_max; j++) { + const size_t j_other = j ^ offset_j; + if (j_other >= n_backends) { + continue; + } + push_data(j, j_other, i_buf); + } + + for (size_t j = 0; j < 2*offset_j_max; j++) { + if (step_cgraphs[j] == nullptr) { + continue; + } + auto & bcj = backend_ctx->backend_configs[j]; + const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, step_cgraphs[j]); + if (status != GGML_STATUS_SUCCESS) { + return status; + } + } + i_buf++; + } + assert(i_buf == backend_ctx->n_reduce_steps); + + // If n_backends is not a power of 2, copy back the reduced tensors to the excess: + for (size_t j = 2*offset_j_max; j < n_backends; j++) { + auto & bcj_src = backend_ctx->backend_configs[j - 2*offset_j_max]; + auto & bcj_dst = backend_ctx->backend_configs[j]; + + ggml_tensor * node_src = bcj_src.cgraphs[i].cgraph_main->nodes[bcj_src.cgraphs[i].cgraph_main->n_nodes - 1]; + ggml_tensor * node_dst = bcj_dst.cgraphs[i].cgraph_main->nodes[bcj_dst.cgraphs[i].cgraph_main->n_nodes - 1]; + ggml_backend_tensor_copy_async(bcj_src.backend, bcj_dst.backend, node_src, node_dst); + } + + return GGML_STATUS_SUCCESS; + }; + + + for (size_t i = 0; i < backend_ctx->n_subgraphs; i++) { + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, bcj.cgraphs[i].cgraph_main); + if (status != GGML_STATUS_SUCCESS) { + return status; + } + } + + if (n_backends > 1 && i < backend_ctx->n_subgraphs - 1) { + bool backend_allreduce_success = false; + if (backend_ctx->comm_ctx) { + std::vector nodes; + nodes.reserve(n_backends); + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + ggml_cgraph * cgraph_ij = bcj.cgraphs[i].cgraph_main; + nodes.push_back(cgraph_ij->nodes[cgraph_ij->n_nodes-1]); + } + backend_allreduce_success = backend_ctx->comm_allreduce(backend_ctx->comm_ctx, nodes.data()); + } + + if (!backend_allreduce_success) { + const ggml_status status = allreduce_fallback(i); + if (status != GGML_STATUS_SUCCESS) { + return status; + } + } + } + } + return GGML_STATUS_SUCCESS; +} + +static const ggml_backend_i ggml_backend_meta_i = { + /* .get_name = */ ggml_backend_meta_get_name, + /* .free = */ ggml_backend_meta_free, + /* .set_tensor_async = */ ggml_backend_meta_set_tensor_async, + /* .get_tensor_async = */ ggml_backend_meta_get_tensor_async, + /* .set_tensor_2d_async = */ nullptr, + /* .get_tensor_2d_async = */ nullptr, + /* .cpy_tensor_async = */ nullptr, + /* .synchronize = */ ggml_backend_meta_synchronize, + /* .graph_plan_create = */ nullptr, + /* .graph_plan_free = */ nullptr, + /* .graph_plan_update = */ nullptr, + /* .graph_plan_compute = */ nullptr, + /* .graph_compute = */ ggml_backend_meta_graph_compute, + /* .event_record = */ nullptr, + /* .event_wait = */ nullptr, + /* .graph_optimize = */ nullptr, +}; + +bool ggml_backend_is_meta(ggml_backend_t backend) { + return backend != nullptr && backend->iface.get_name == ggml_backend_meta_i.get_name; +} + +static ggml_backend_t ggml_backend_meta_device_init_backend(ggml_backend_dev_t dev, const char * params) { + ggml_backend_meta_context * backend_ctx = new ggml_backend_meta_context(dev, params); + + ggml_backend_t backend = new struct ggml_backend; + backend->guid = ggml_backend_meta_guid(); + backend->iface = ggml_backend_meta_i; + backend->device = dev; + backend->context = backend_ctx; + return backend; +} + +size_t ggml_backend_meta_n_backends(ggml_backend_t meta_backend) { + GGML_ASSERT(ggml_backend_is_meta(meta_backend)); + const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context; + return backend_ctx->backend_configs.size(); +} + +ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, size_t index) { + GGML_ASSERT(ggml_backend_is_meta(meta_backend)); + const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context; + return backend_ctx->backend_configs[index].backend; +} + + +// ---- optimization API implementation ---- + +#include "ggml-opt.h.inc" + +#include "ggml.h.inc" +#include "ggml-alloc.h.inc" +#include "ggml-backend.h.inc" +#include "ggml-impl-defs.inc" + +#include +#include +#include +#include +#include +#include +#include + +struct ggml_opt_dataset { + struct ggml_context * ctx = nullptr; + ggml_backend_buffer_t buf = nullptr; + struct ggml_tensor * data = nullptr; + struct ggml_tensor * labels = nullptr; + + int64_t ndata = -1; + int64_t ndata_shard = -1; + size_t nbs_data = -1; + size_t nbs_labels = -1; + + std::vector permutation; +}; + +struct ggml_opt_context { + ggml_backend_sched_t backend_sched = nullptr; + ggml_cgraph * allocated_graph = nullptr; + ggml_cgraph * allocated_graph_copy = nullptr; + struct ggml_context * ctx_static = nullptr; + struct ggml_context * ctx_cpu = nullptr; + struct ggml_context * ctx_compute = nullptr; + struct ggml_context * ctx_copy = nullptr; + ggml_backend_buffer_t buf_static = nullptr; + ggml_backend_buffer_t buf_cpu = nullptr; + std::mt19937 rng; + enum ggml_opt_loss_type loss_type; + enum ggml_opt_build_type build_type; + enum ggml_opt_build_type build_type_alloc; + + struct ggml_tensor * inputs = nullptr; + struct ggml_tensor * outputs = nullptr; + struct ggml_tensor * labels = nullptr; + + struct ggml_tensor * loss = nullptr; + struct ggml_tensor * pred = nullptr; + struct ggml_tensor * ncorrect = nullptr; + + struct ggml_cgraph * gf = nullptr; + struct ggml_cgraph * gb_grad = nullptr; + struct ggml_cgraph * gb_opt = nullptr; + bool static_graphs = false; + bool eval_ready = false; + std::vector grad_accs; + std::vector grad_m; + std::vector grad_v; + + int64_t iter = 1; + int32_t opt_period = 1; + int32_t opt_i = 0; + bool loss_per_datapoint = false; + + ggml_opt_get_optimizer_params get_opt_pars = nullptr; + void * get_opt_pars_ud = nullptr; + struct ggml_tensor * opt_step_params = nullptr; // Stores output of get_opt_pars. + + enum ggml_opt_optimizer_type optimizer = GGML_OPT_OPTIMIZER_TYPE_ADAMW; +}; + +struct ggml_opt_result { + int64_t ndata = 0; + std::vector loss; + std::vector pred; + int64_t ncorrect = 0; + + int64_t opt_period = -1; + bool loss_per_datapoint = false; +}; + +// ====== Dataset ====== + +ggml_opt_dataset_t ggml_opt_dataset_init( + enum ggml_type type_data, + enum ggml_type type_label, + int64_t ne_datapoint, + int64_t ne_label, + int64_t ndata, + int64_t ndata_shard) { + GGML_ASSERT(ne_datapoint > 0); + GGML_ASSERT(ne_label >= 0); + GGML_ASSERT(ndata > 0); + GGML_ASSERT(ndata_shard > 0); + + ggml_opt_dataset_t result = new ggml_opt_dataset; + result->ndata = ndata; + result->ndata_shard = ndata_shard; + + { + struct ggml_init_params params = { + /*.mem_size =*/ 2*ggml_tensor_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + result->ctx = ggml_init(params); + } + + result->data = ggml_new_tensor_2d(result->ctx, type_data, ne_datapoint, ndata); + result->nbs_data = ggml_nbytes(result->data) * ndata_shard/ndata; + + if (ne_label > 0) { + result->labels = ggml_new_tensor_2d(result->ctx, type_label, ne_label, ndata); + result->nbs_labels = ggml_nbytes(result->labels) * ndata_shard/ndata; + } else { + result->labels = nullptr; + result->nbs_labels = 0; + } + + result->buf = ggml_backend_alloc_ctx_tensors_from_buft(result->ctx, ggml_backend_cpu_buffer_type()); + + const int64_t nshards = ndata/ndata_shard; + result->permutation.resize(nshards); + for (int64_t i = 0; i < nshards; ++i) { + result->permutation[i] = i; + } + return result; +} + +void ggml_opt_dataset_free(ggml_opt_dataset_t dataset) { + ggml_backend_buffer_free(dataset->buf); + ggml_free(dataset->ctx); + delete dataset; +} + +int64_t ggml_opt_dataset_ndata(ggml_opt_dataset_t dataset) { + return dataset->ndata; +} + +struct ggml_tensor * ggml_opt_dataset_data(ggml_opt_dataset_t dataset) { + return dataset->data; +} + +struct ggml_tensor * ggml_opt_dataset_labels(ggml_opt_dataset_t dataset) { + return dataset->labels; +} + +void ggml_opt_dataset_shuffle(ggml_opt_context_t opt_ctx, ggml_opt_dataset_t dataset, int64_t idata) { + GGML_ASSERT(idata <= dataset->ndata); + + if (idata < 0) { + std::shuffle(dataset->permutation.begin(), dataset->permutation.end(), opt_ctx->rng); + return; + } + + GGML_ASSERT(idata % dataset->ndata_shard == 0); + const int64_t ishard_max = idata / dataset->ndata_shard; + std::shuffle(dataset->permutation.begin(), dataset->permutation.begin() + ishard_max, opt_ctx->rng); +} + +void ggml_opt_dataset_get_batch(ggml_opt_dataset_t dataset, struct ggml_tensor * data_batch, struct ggml_tensor * labels_batch, int64_t ibatch) { + GGML_ASSERT( data_batch && ggml_is_contiguous(data_batch)); + GGML_ASSERT(!labels_batch || ggml_is_contiguous(labels_batch)); + GGML_ASSERT((labels_batch == nullptr) == (dataset->labels == nullptr)); + GGML_ASSERT( data_batch->type == dataset->data->type); + GGML_ASSERT(!labels_batch || labels_batch->type == dataset->labels->type); + + const size_t nb_data_batch = ggml_nbytes(data_batch); + GGML_ASSERT(nb_data_batch % dataset->nbs_data == 0); + const int64_t shards_per_batch = nb_data_batch / dataset->nbs_data; + + if (labels_batch) { + const size_t nb_labels_batch = ggml_nbytes(labels_batch); + GGML_ASSERT(nb_labels_batch == shards_per_batch*dataset->nbs_labels); + } + + GGML_ASSERT((ibatch + 1)*shards_per_batch <= int64_t(dataset->permutation.size())); + + for (int64_t ishard_batch = 0; ishard_batch < shards_per_batch; ++ishard_batch) { + const int64_t ishard = dataset->permutation[ibatch*shards_per_batch + ishard_batch]; + + const char * ptr_data = (const char *) dataset->data->data + ishard*dataset->nbs_data; + ggml_backend_tensor_set(data_batch, ptr_data, ishard_batch*dataset->nbs_data, dataset->nbs_data); + + if (!labels_batch) { + continue; + } + + const char * ptr_labels = (const char *) dataset->labels->data + ishard*dataset->nbs_labels; + ggml_backend_tensor_set(labels_batch, ptr_labels, ishard_batch*dataset->nbs_labels, dataset->nbs_labels); + } +} + +void ggml_opt_dataset_get_batch_host(ggml_opt_dataset_t dataset, void * data_batch, size_t nb_data_batch, void * labels_batch, int64_t ibatch) { + GGML_ASSERT((labels_batch == nullptr) == (dataset->labels == nullptr)); + GGML_ASSERT(nb_data_batch % dataset->nbs_data == 0); + + const int64_t shards_per_batch = nb_data_batch / dataset->nbs_data; + + GGML_ASSERT((ibatch + 1)*shards_per_batch <= int64_t(dataset->permutation.size())); + + for (int64_t ishard_batch = 0; ishard_batch < shards_per_batch; ++ishard_batch) { + const int64_t ishard = dataset->permutation[ibatch*shards_per_batch + ishard_batch]; + + const char * ptr_data = (const char *) dataset->data->data + ishard *dataset->nbs_data; + char * ptr_data_batch = (char *) data_batch + ishard_batch*dataset->nbs_data; + memcpy(ptr_data_batch, ptr_data, dataset->nbs_data); + + if (!labels_batch) { + continue; + } + + const char * ptr_labels = (const char *) dataset->labels->data + ishard *dataset->nbs_labels; + char * ptr_labels_batch = (char *) labels_batch + ishard_batch*dataset->nbs_labels; + memcpy(ptr_labels_batch, ptr_labels, dataset->nbs_labels); + } +} + +// ====== Model / Context ====== + +struct ggml_opt_optimizer_params ggml_opt_get_default_optimizer_params(void * userdata) { + GGML_UNUSED(userdata); + + ggml_opt_optimizer_params result; + + result.adamw.alpha = 0.001f; + result.adamw.beta1 = 0.9f; + result.adamw.beta2 = 0.999f; + result.adamw.eps = 1e-8f; + result.adamw.wd = 0.0f; + + result.sgd.alpha = 1e-3f; + result.sgd.wd = 0.0f; + + return result; +} + + +struct ggml_opt_optimizer_params ggml_opt_get_constant_optimizer_params(void * userdata) { + return *((struct ggml_opt_optimizer_params *) userdata); +} + +struct ggml_opt_params ggml_opt_default_params( + ggml_backend_sched_t backend_sched, + enum ggml_opt_loss_type loss_type) { + return { + /*backend_sched =*/ backend_sched, + /*ctx_compute =*/ nullptr, + /*inputs =*/ nullptr, + /*logits =*/ nullptr, + /*loss_type =*/ loss_type, + /*build_type =*/ GGML_OPT_BUILD_TYPE_OPT, + /*opt_period =*/ 1, + /*get_opt_pars =*/ ggml_opt_get_default_optimizer_params, + /*get_opt_pars_ud =*/ nullptr, + /*optimizer =*/ GGML_OPT_OPTIMIZER_TYPE_ADAMW, + }; +} + +static ggml_tensor * map_tensor(std::map & tensor_map, ggml_context * ctx, ggml_tensor * tensor) { + if (!tensor) { + return nullptr; + } + + if (tensor_map.find(tensor) != tensor_map.end()) { + return tensor_map[tensor]; + } + + ggml_tensor * new_tensor = ggml_dup_tensor(ctx, tensor); + tensor_map[tensor] = new_tensor; + + new_tensor->op = tensor->op; + for (int i = 0; i < GGML_MAX_DIMS; i++) { + new_tensor->nb[i] = tensor->nb[i]; + } + new_tensor->flags = tensor->flags; + memcpy(new_tensor->op_params, tensor->op_params, sizeof(tensor->op_params)); + strcpy(new_tensor->name, tensor->name); + new_tensor->data = tensor->data; + new_tensor->buffer = tensor->buffer; + new_tensor->extra = tensor->extra; + new_tensor->view_offs = tensor->view_offs; + new_tensor->view_src = map_tensor(tensor_map, ctx, tensor->view_src); + for (int i = 0; i < GGML_MAX_SRC; i++) { + new_tensor->src[i] = map_tensor(tensor_map, ctx, tensor->src[i]); + } + + return new_tensor; +} + +static ggml_cgraph * dup_graph(ggml_context * ctx, ggml_cgraph * src) { + std::map tensor_map; + + ggml_cgraph * dst = ggml_new_graph_custom(ctx, src->size, /*grads =*/ true); + + for (int i = 0; i < src->n_leafs; i++) { + ggml_build_forward_expand(dst, map_tensor(tensor_map, ctx, src->leafs[i])); + } + GGML_ASSERT(dst->n_leafs == src->n_leafs); + for (int i = 0; i < src->n_nodes; i++) { + ggml_build_forward_expand(dst, map_tensor(tensor_map, ctx, src->nodes[i])); + } + GGML_ASSERT(dst->n_nodes == src->n_nodes); + for (int i = 0; i < src->n_nodes; ++i) { + const size_t igrad_src = ggml_hash_find(&src->visited_hash_set, src->nodes[i]); + const size_t igrad_dst = ggml_hash_find(&dst->visited_hash_set, dst->nodes[i]); + + GGML_ASSERT(igrad_src != GGML_HASHSET_FULL); + GGML_ASSERT(ggml_bitset_get(src->visited_hash_set.used, igrad_src)); + GGML_ASSERT(igrad_dst != GGML_HASHSET_FULL); + GGML_ASSERT(ggml_bitset_get(dst->visited_hash_set.used, igrad_dst)); + + dst->grads[igrad_dst] = src->grads[igrad_src]; + dst->grad_accs[igrad_dst] = src->grad_accs[igrad_src]; + } + + return dst; +} + +static void ggml_opt_build(ggml_opt_context_t opt_ctx) { + GGML_ASSERT(opt_ctx->ctx_compute && "no compute context set, either use static graphs or set one with ggml_opt_prepare_alloc"); + GGML_ASSERT((!opt_ctx->static_graphs || opt_ctx->inputs->data) && "when using static graphs the inputs must be allocated statically"); + + const enum ggml_opt_optimizer_type optimizer = opt_ctx->optimizer; + + const bool accumulate = opt_ctx->build_type_alloc >= GGML_OPT_BUILD_TYPE_GRAD && + !(opt_ctx->static_graphs && opt_ctx->build_type_alloc == GGML_OPT_BUILD_TYPE_OPT && opt_ctx->opt_period == 1); + + const bool need_momenta = opt_ctx->build_type_alloc == GGML_OPT_BUILD_TYPE_OPT && + opt_ctx->optimizer == GGML_OPT_OPTIMIZER_TYPE_ADAMW; + + ggml_set_input(opt_ctx->inputs); + ggml_set_output(opt_ctx->outputs); + + int n_param = 0; + for (int i = 0; i < opt_ctx->gf->n_nodes; ++i) { + const struct ggml_tensor * node = opt_ctx->gf->nodes[i]; + if (node->flags & GGML_TENSOR_FLAG_PARAM) { + n_param++; + } + GGML_ASSERT(!(node->flags & GGML_TENSOR_FLAG_LOSS) && "support for extra loss terms not implemented"); + } + + if (!opt_ctx->ctx_static) { + // The static context is used for: + // - gradients (1 per loss, 1 tensor per param if using gradient accumulation) + // - optimizer momenta (2 tensors per param) + // - labels (if using static graphs) + // - loss (if using static graphs, up to 5 tensors) + // - pred (if using static graphs) + // - ncorrect (if using static graphs, 2 tensors). + constexpr size_t n_loss = 1; + const size_t tensors_per_param = (accumulate ? 1 : 0) + (need_momenta ? 2 : 0); + const size_t tensors_const = opt_ctx->static_graphs ? 9 : 0; + const size_t size_meta = (n_loss + tensors_per_param*n_param + tensors_const) * ggml_tensor_overhead(); + struct ggml_init_params params = { + /*.mem_size =*/ size_meta, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + opt_ctx->ctx_static = ggml_init(params); + } + GGML_ASSERT(opt_ctx->build_type <= opt_ctx->build_type_alloc); + + { + // The cpu context is allocated statically if using static graphs, dynamically otherwise. + // It is used for: + // - optimizer parameters (1 shared for all optimizer invocations) + const size_t size_meta = 1 * ggml_tensor_overhead(); + struct ggml_init_params params = { + /*.mem_size =*/ size_meta, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_free(opt_ctx->ctx_cpu); + opt_ctx->ctx_cpu = ggml_init(params); + + ggml_backend_buffer_free(opt_ctx->buf_cpu); + opt_ctx->buf_cpu = nullptr; + } + + struct ggml_context * ctx_results = opt_ctx->static_graphs ? opt_ctx->ctx_static : opt_ctx->ctx_compute; + + switch (opt_ctx->loss_type) { + case GGML_OPT_LOSS_TYPE_MEAN: { + opt_ctx->loss = ggml_sum(ctx_results, opt_ctx->outputs); + ggml_set_name(opt_ctx->loss, "loss_sum"); + const float scale = 1.0f / (opt_ctx->opt_period * ggml_nelements(opt_ctx->outputs)); + opt_ctx->loss = ggml_scale(ctx_results, opt_ctx->loss, scale); + ggml_set_name(opt_ctx->loss, "loss_mean"); + opt_ctx->loss_per_datapoint = true; + break; + } + case GGML_OPT_LOSS_TYPE_SUM: { + opt_ctx->loss = ggml_sum(ctx_results, opt_ctx->outputs); + ggml_set_name(opt_ctx->loss, "loss_sum"); + opt_ctx->loss_per_datapoint = false; + break; + } + case GGML_OPT_LOSS_TYPE_CROSS_ENTROPY: { + opt_ctx->labels = ggml_dup_tensor(ctx_results, opt_ctx->outputs); + ggml_set_input(opt_ctx->labels); + ggml_set_name(opt_ctx->labels, "labels"); + opt_ctx->loss = ggml_cross_entropy_loss(ctx_results, opt_ctx->outputs, opt_ctx->labels); + ggml_set_name(opt_ctx->loss, "loss_cross_entropy"); + if (opt_ctx->opt_period > 1) { + opt_ctx->loss = ggml_scale(ctx_results, opt_ctx->loss, 1.0f / opt_ctx->opt_period); + ggml_set_name(opt_ctx->loss, "loss_cross_entropy_scaled"); + } + opt_ctx->loss_per_datapoint = true; + break; + } + case GGML_OPT_LOSS_TYPE_MEAN_SQUARED_ERROR: { + opt_ctx->labels = ggml_dup_tensor(ctx_results, opt_ctx->outputs); + ggml_set_input(opt_ctx->labels); + ggml_set_name(opt_ctx->labels, "labels"); + opt_ctx->loss = ggml_sub(ctx_results, opt_ctx->outputs, opt_ctx->labels); + ggml_set_name(opt_ctx->loss, "loss_error"); + opt_ctx->loss = ggml_sqr(ctx_results, opt_ctx->loss); + ggml_set_name(opt_ctx->loss, "loss_squared_error"); + opt_ctx->loss = ggml_sum(ctx_results, opt_ctx->loss); + ggml_set_name(opt_ctx->loss, "loss_sum_squared_error"); + const float scale = 1.0f / (opt_ctx->opt_period * ggml_nelements(opt_ctx->outputs)); + opt_ctx->loss = ggml_scale(ctx_results, opt_ctx->loss, scale); + ggml_set_name(opt_ctx->loss, "loss_mean_squared_error"); + opt_ctx->loss_per_datapoint = true; + break; + } + } + ggml_set_output(opt_ctx->loss); + ggml_set_loss(opt_ctx->loss); + ggml_build_forward_expand(opt_ctx->gf, opt_ctx->loss); + + if (opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY) { + opt_ctx->pred = ggml_argmax(ctx_results, opt_ctx->outputs); + ggml_set_name(opt_ctx->pred, "pred"); + ggml_set_output(opt_ctx->pred); + ggml_build_forward_expand(opt_ctx->gf, opt_ctx->pred); + + opt_ctx->ncorrect = ggml_count_equal(ctx_results, opt_ctx->pred, ggml_argmax(ctx_results, opt_ctx->labels)); + ggml_set_name(opt_ctx->ncorrect, "ncorrect"); + ggml_set_output(opt_ctx->ncorrect); + ggml_build_forward_expand(opt_ctx->gf, opt_ctx->ncorrect); + } + + if (opt_ctx->buf_static) { + if (opt_ctx->build_type == GGML_OPT_BUILD_TYPE_FORWARD) { + return; + } + } else if (opt_ctx->build_type_alloc == GGML_OPT_BUILD_TYPE_FORWARD) { + opt_ctx->buf_static = ggml_backend_alloc_ctx_tensors( + opt_ctx->ctx_static, ggml_backend_sched_get_backend(opt_ctx->backend_sched, 0)); + return; + } + + if (opt_ctx->grad_accs.empty()) { + GGML_ASSERT(opt_ctx->build_type_alloc >= GGML_OPT_BUILD_TYPE_GRAD); + + const int n_nodes = opt_ctx->gf->n_nodes; + opt_ctx->grad_accs.resize(n_nodes); + for (int i = 0; i < n_nodes; ++i) { + ggml_tensor * node = opt_ctx->gf->nodes[i]; + if ((accumulate && (node->flags & GGML_TENSOR_FLAG_PARAM)) || (node->flags & GGML_TENSOR_FLAG_LOSS)) { + opt_ctx->grad_accs[i] = ggml_new_tensor(opt_ctx->ctx_static, GGML_TYPE_F32, GGML_MAX_DIMS, node->ne); + } else { + opt_ctx->grad_accs[i] = nullptr; + } + } + + if (need_momenta && opt_ctx->build_type_alloc >= GGML_OPT_BUILD_TYPE_OPT) { + opt_ctx->grad_m.resize(n_nodes); + opt_ctx->grad_v.resize(n_nodes); + for (int i = 0; i < n_nodes; ++i) { + ggml_tensor * node = opt_ctx->gf->nodes[i]; + if (node->flags & GGML_TENSOR_FLAG_PARAM) { + opt_ctx->grad_m[i] = ggml_new_tensor(opt_ctx->ctx_static, GGML_TYPE_F32, GGML_MAX_DIMS, node->ne); + opt_ctx->grad_v[i] = ggml_new_tensor(opt_ctx->ctx_static, GGML_TYPE_F32, GGML_MAX_DIMS, node->ne); + } else { + opt_ctx->grad_m[i] = nullptr; + opt_ctx->grad_v[i] = nullptr; + } + } + } + } + + // gb_grad == graph backward gradients, forward pass, then backward pass to calculate gradients. + opt_ctx->gb_grad = ggml_graph_dup(opt_ctx->ctx_compute, opt_ctx->gf, /*force_grads =*/ true); + ggml_build_backward_expand(opt_ctx->ctx_compute, opt_ctx->gb_grad, opt_ctx->grad_accs.data()); + + if (opt_ctx->buf_static) { + if (opt_ctx->build_type == GGML_OPT_BUILD_TYPE_GRAD) { + return; + } + } else if (opt_ctx->build_type_alloc == GGML_OPT_BUILD_TYPE_GRAD) { + opt_ctx->buf_static = ggml_backend_alloc_ctx_tensors(opt_ctx->ctx_static, ggml_backend_sched_get_backend(opt_ctx->backend_sched, 0)); + ggml_graph_reset(opt_ctx->gb_grad); + } + + GGML_ASSERT(opt_ctx->build_type_alloc == GGML_OPT_BUILD_TYPE_OPT); + + // gb_opt == graph backward optimize, forward pass, then backward pass to calculate gradients, then optimizer step. + opt_ctx->gb_opt = ggml_graph_dup(opt_ctx->ctx_compute, opt_ctx->gb_grad, /*force_grads =*/ true); + + opt_ctx->opt_step_params = ggml_new_tensor_1d(opt_ctx->ctx_cpu, GGML_TYPE_F32, need_momenta ? 7 : 2); + ggml_tensor * adamw_params = opt_ctx->opt_step_params; + ggml_set_input(adamw_params); + const char * optimizer_name = ggml_opt_optimizer_name(opt_ctx->optimizer); + ggml_format_name(adamw_params, "%s_params", optimizer_name); + for (int i = opt_ctx->gf->n_nodes-1; i >= 0; --i) { + struct ggml_tensor * node = opt_ctx->gb_opt->nodes[i]; + struct ggml_tensor * grad = ggml_graph_get_grad(opt_ctx->gb_opt, node); + + if (grad && (node->flags & GGML_TENSOR_FLAG_PARAM)) { + struct ggml_tensor * m = nullptr; + struct ggml_tensor * v = nullptr; + if (need_momenta) { + m = opt_ctx->grad_m[i]; + v = opt_ctx->grad_v[i]; + ggml_format_name(m, "AdamW m for %s", node->name); + ggml_format_name(v, "AdamW v for %s", node->name); + } + struct ggml_tensor * opt_step; + switch (optimizer) { + case GGML_OPT_OPTIMIZER_TYPE_ADAMW: + opt_step = ggml_opt_step_adamw(opt_ctx->ctx_compute, node, grad, m, v, adamw_params); + break; + case GGML_OPT_OPTIMIZER_TYPE_SGD: + opt_step = ggml_opt_step_sgd(opt_ctx->ctx_compute, node, grad, adamw_params); + break; + default: + GGML_ABORT("fatal error"); + } + ggml_format_name(opt_step, "%s step for %s", optimizer_name, node->name); + ggml_build_forward_expand(opt_ctx->gb_opt, opt_step); + } + } + + if (!opt_ctx->buf_static) { + opt_ctx->buf_static = ggml_backend_alloc_ctx_tensors( + opt_ctx->ctx_static, ggml_backend_sched_get_backend(opt_ctx->backend_sched, 0)); + ggml_graph_reset(opt_ctx->gb_opt); + } + + opt_ctx->buf_cpu = ggml_backend_alloc_ctx_tensors_from_buft(opt_ctx->ctx_cpu, ggml_backend_cpu_buffer_type()); +} + +ggml_opt_context_t ggml_opt_init(struct ggml_opt_params params) { + ggml_opt_context_t result = new struct ggml_opt_context; + result->backend_sched = params.backend_sched; + result->ctx_compute = params.ctx_compute; + result->loss_type = params.loss_type; + result->build_type = params.build_type; + result->build_type_alloc = params.build_type; + result->inputs = params.inputs; + result->outputs = params.outputs; + result->opt_period = params.opt_period; + result->get_opt_pars = params.get_opt_pars; + result->get_opt_pars_ud = params.get_opt_pars_ud; + result->optimizer = params.optimizer; + + GGML_ASSERT(result->opt_period >= 1); + + result->static_graphs = result->ctx_compute; + + if (!result->static_graphs) { + GGML_ASSERT(!result->inputs); + GGML_ASSERT(!result->outputs); + return result; + } + + GGML_ASSERT(result->inputs); + GGML_ASSERT(result->outputs); + + result->gf = ggml_new_graph_custom(result->ctx_compute, GGML_DEFAULT_GRAPH_SIZE, /*grads =*/ true); // Forward pass. + ggml_build_forward_expand(result->gf, result->outputs); + + ggml_opt_build(result); + + return result; +} + +void ggml_opt_free(ggml_opt_context_t opt_ctx) { + if (opt_ctx == nullptr) { + return; + } + ggml_backend_buffer_free(opt_ctx->buf_static); + ggml_backend_buffer_free(opt_ctx->buf_cpu); + ggml_free(opt_ctx->ctx_static); + ggml_free(opt_ctx->ctx_cpu); + ggml_free(opt_ctx->ctx_copy); + delete opt_ctx; +} + +void ggml_opt_reset(ggml_opt_context_t opt_ctx, bool optimizer) { + if (optimizer) { + ggml_graph_reset(opt_ctx->gb_opt); + opt_ctx->iter = 1; + } else { + ggml_graph_reset(opt_ctx->gb_grad); + } +} + +bool ggml_opt_static_graphs(ggml_opt_context_t opt_ctx) { + return opt_ctx->static_graphs; +} + +struct ggml_tensor * ggml_opt_inputs(ggml_opt_context_t opt_ctx) { + return opt_ctx->inputs; +} + +struct ggml_tensor * ggml_opt_outputs(ggml_opt_context_t opt_ctx) { + return opt_ctx->outputs; +} + +struct ggml_tensor * ggml_opt_labels(ggml_opt_context_t opt_ctx) { + return opt_ctx->labels; +} + +struct ggml_tensor * ggml_opt_loss(ggml_opt_context_t opt_ctx) { + return opt_ctx->loss; +} + +struct ggml_tensor * ggml_opt_pred(ggml_opt_context_t opt_ctx) { + return opt_ctx->pred; +} + +struct ggml_tensor * ggml_opt_ncorrect(ggml_opt_context_t opt_ctx) { + return opt_ctx->ncorrect; +} + +struct ggml_tensor * ggml_opt_grad_acc(ggml_opt_context_t opt_ctx, struct ggml_tensor * node) { + return ggml_graph_get_grad_acc(opt_ctx->gb_opt, node); +} + +// ====== Optimization Result ====== + +ggml_opt_result_t ggml_opt_result_init() { + return new ggml_opt_result; +} + +void ggml_opt_result_free(ggml_opt_result_t result) { + delete result; +} + +void ggml_opt_result_reset(ggml_opt_result_t result) { + result->ndata = 0; + result->loss.clear(); + result->pred.clear(); + result->ncorrect = 0; +} + +void ggml_opt_result_ndata(ggml_opt_result_t result, int64_t * ndata) { + *ndata = result->ndata; +} + +void ggml_opt_result_loss(ggml_opt_result_t result, double * loss, double * unc) { + const int64_t nbatches = result->loss.size(); // Number of physical batches. + + if (nbatches == 0) { + *loss = 0.0; + *unc = NAN; + return; + } + + double sum = 0.0; + double sum_squared = 0.0; + + for (const float & loss : result->loss) { + // If the loss is per datapoint it was scaled by 1.0f/opt_period for each physical batch. + const float loss_scaled = result->loss_per_datapoint ? loss*result->opt_period : loss; + sum += loss_scaled; + sum_squared += loss_scaled*loss_scaled; + } + + const double mean = sum/nbatches; + *loss = result->loss_per_datapoint ? mean : sum; + + if (!unc) { + return; + } + + if (nbatches < 2) { + *unc = NAN; + return; + } + + const double var_sum = sum_squared/nbatches - mean*mean; // variance without Bessel's correction, i.e. nbatches/(nbatches-1) + *unc = result->loss_per_datapoint ? sqrt(var_sum / (nbatches - 1)) : sqrt(var_sum * nbatches/(nbatches - 1)); +} + +void ggml_opt_result_pred(ggml_opt_result_t result, int32_t * pred) { + for (size_t i = 0; i < result->pred.size(); ++i) { + pred[i] = result->pred[i]; + } +} + +void ggml_opt_result_accuracy(ggml_opt_result_t result, double * accuracy, double * unc) { + *accuracy = result->ncorrect >= 0 ? double(result->ncorrect) / double(result->ndata) : NAN; + + if (!unc) { + return; + } + + *unc = result->ncorrect >= 0 && result->ndata >= 2 ? + sqrt((*accuracy) * (1.0 - (*accuracy)) / double(result->ndata - 1)) : NAN; +} + +// ====== Computation ====== + +void ggml_opt_prepare_alloc( + ggml_opt_context_t opt_ctx, + struct ggml_context * ctx_compute, + struct ggml_cgraph * gf, + struct ggml_tensor * inputs, + struct ggml_tensor * outputs) { + GGML_ASSERT(!opt_ctx->static_graphs); + opt_ctx->ctx_compute = ctx_compute; + opt_ctx->gf = gf; + opt_ctx->inputs = inputs; + opt_ctx->outputs = outputs; +} + +void ggml_opt_alloc(ggml_opt_context_t opt_ctx, bool backward) { + GGML_ASSERT(!opt_ctx->eval_ready); + if (opt_ctx->build_type == GGML_OPT_BUILD_TYPE_OPT && opt_ctx->opt_period > 1 && opt_ctx->opt_i == 0) { + ggml_graph_reset(opt_ctx->gb_grad); + } + if (backward) { + const int32_t opt_i_next = (opt_ctx->opt_i + 1) % opt_ctx->opt_period; + opt_ctx->build_type = opt_i_next == 0 ? GGML_OPT_BUILD_TYPE_OPT : GGML_OPT_BUILD_TYPE_GRAD; + } else { + opt_ctx->build_type = GGML_OPT_BUILD_TYPE_FORWARD; + } + + if (!opt_ctx->static_graphs) { + ggml_opt_build(opt_ctx); + } + + struct ggml_cgraph * graph = nullptr; + switch (opt_ctx->build_type) { + case GGML_OPT_BUILD_TYPE_FORWARD: { + graph = opt_ctx->gf; + } break; + case GGML_OPT_BUILD_TYPE_GRAD: { + graph = opt_ctx->gb_grad; + } break; + case GGML_OPT_BUILD_TYPE_OPT: { + graph = opt_ctx->gb_opt; + } break; + } + GGML_ASSERT(graph); + + if (opt_ctx->allocated_graph == graph) { + opt_ctx->eval_ready = true; + return; + } + + ggml_backend_sched_reset(opt_ctx->backend_sched); // clear allocation of previous graph + + if (opt_ctx->static_graphs) { + ggml_init_params params = { + /*.mem_size =*/ graph->size*ggml_tensor_overhead() + ggml_graph_overhead_custom(graph->size, graph->grads), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_free(opt_ctx->ctx_copy); + opt_ctx->ctx_copy = ggml_init(params); + + opt_ctx->allocated_graph_copy = dup_graph(opt_ctx->ctx_copy, graph); + } else { + opt_ctx->allocated_graph_copy = graph; + } + + ggml_backend_sched_alloc_graph(opt_ctx->backend_sched, opt_ctx->allocated_graph_copy); + opt_ctx->allocated_graph = graph; + + opt_ctx->eval_ready = true; +} + +void ggml_opt_eval(ggml_opt_context_t opt_ctx, ggml_opt_result_t result) { + GGML_ASSERT(opt_ctx->eval_ready); + if (opt_ctx->allocated_graph == opt_ctx->gb_opt) { + const ggml_opt_optimizer_params & opt_pars = opt_ctx->get_opt_pars(opt_ctx->get_opt_pars_ud); + + switch (opt_ctx->optimizer) { + case GGML_OPT_OPTIMIZER_TYPE_ADAMW: { + GGML_ASSERT(opt_pars.adamw.alpha > 0.0f); + GGML_ASSERT(opt_pars.adamw.beta1 >= 0.0f); + GGML_ASSERT(opt_pars.adamw.beta1 <= 1.0f); + GGML_ASSERT(opt_pars.adamw.beta2 >= 0.0f); + GGML_ASSERT(opt_pars.adamw.beta2 <= 1.0f); + GGML_ASSERT(opt_pars.adamw.eps >= 0.0f); + GGML_ASSERT(opt_pars.adamw.wd >= 0.0f); + GGML_ASSERT(opt_pars.adamw.wd <= 1.0f); + + // beta1, beta2 after applying warmup + const float beta1h = 1.0f / (1.0f - powf(opt_pars.adamw.beta1, opt_ctx->iter)); + const float beta2h = 1.0f / (1.0f - powf(opt_pars.adamw.beta2, opt_ctx->iter)); + + float * adamw_par_data = ggml_get_data_f32(opt_ctx->opt_step_params); + adamw_par_data[0] = opt_pars.adamw.alpha; + adamw_par_data[1] = opt_pars.adamw.beta1; + adamw_par_data[2] = opt_pars.adamw.beta2; + adamw_par_data[3] = opt_pars.adamw.eps; + adamw_par_data[4] = opt_pars.adamw.wd; + adamw_par_data[5] = beta1h; + adamw_par_data[6] = beta2h; + } break; + case GGML_OPT_OPTIMIZER_TYPE_SGD: { + GGML_ASSERT(opt_pars.sgd.alpha > 0.0f); + GGML_ASSERT(opt_pars.sgd.wd >= 0.0f); + GGML_ASSERT(opt_pars.sgd.wd <= 1.0f); + float * sgd = ggml_get_data_f32(opt_ctx->opt_step_params); + sgd[0] = opt_pars.sgd.alpha; + sgd[1] = opt_pars.sgd.wd; + } break; + default: + GGML_ABORT("fatal error"); + } + } + + ggml_backend_sched_graph_compute(opt_ctx->backend_sched, opt_ctx->allocated_graph_copy); + opt_ctx->iter += opt_ctx->allocated_graph == opt_ctx->gb_opt; + opt_ctx->opt_i = (opt_ctx->opt_i + 1) % opt_ctx->opt_period; + + if (!opt_ctx->static_graphs) { + opt_ctx->gf = nullptr; + opt_ctx->gb_grad = nullptr; + opt_ctx->gb_opt = nullptr; + opt_ctx->allocated_graph = nullptr; + opt_ctx->allocated_graph_copy = nullptr; + } + + opt_ctx->eval_ready = false; + + if (!result) { + return; + } + + if (result->ndata == 0) { + result->loss_per_datapoint = opt_ctx->loss_per_datapoint; + result->opt_period = opt_ctx->opt_period; + } else { + GGML_ASSERT(result->loss_per_datapoint == opt_ctx->loss_per_datapoint); + GGML_ASSERT(result->opt_period == opt_ctx->opt_period); + } + + const int64_t ndata = opt_ctx->outputs->ne[1]; + GGML_ASSERT(result->ndata == ndata*int64_t(result->loss.size()) && "varying batch size not supported"); + result->ndata += ndata; + + GGML_ASSERT(ggml_is_scalar(opt_ctx->loss)); + GGML_ASSERT(opt_ctx->loss->type == GGML_TYPE_F32); + float loss; + ggml_backend_tensor_get(opt_ctx->loss, &loss, 0, ggml_nbytes(opt_ctx->loss)); + result->loss.push_back(loss); + + if (opt_ctx->pred) { + GGML_ASSERT(opt_ctx->pred->type == GGML_TYPE_I32); + std::vector pred(ndata); + ggml_backend_tensor_get(opt_ctx->pred, pred.data(), 0, ggml_nbytes(opt_ctx->pred)); + result->pred.insert(result->pred.end(), pred.begin(), pred.end()); + } + + if (!opt_ctx->ncorrect || result->ncorrect < 0) { + result->ncorrect = -1; + return; + } + + GGML_ASSERT(ggml_is_scalar(opt_ctx->ncorrect)); + GGML_ASSERT(opt_ctx->ncorrect->type == GGML_TYPE_I64); + int64_t ncorrect; + ggml_backend_tensor_get(opt_ctx->ncorrect, &ncorrect, 0, ggml_nbytes(opt_ctx->ncorrect)); + result->ncorrect += ncorrect; +} + +// ====== High-Level Functions ====== + +void ggml_opt_epoch( + ggml_opt_context_t opt_ctx, + ggml_opt_dataset_t dataset, + ggml_opt_result_t result_train, + ggml_opt_result_t result_eval, + int64_t idata_split, + ggml_opt_epoch_callback callback_train, + ggml_opt_epoch_callback callback_eval) { + GGML_ASSERT(ggml_opt_static_graphs(opt_ctx) && "ggml_opt_epoch requires static graphs"); + struct ggml_tensor * inputs = ggml_opt_inputs(opt_ctx); + struct ggml_tensor * labels = ggml_opt_labels(opt_ctx); + struct ggml_tensor * data = ggml_opt_dataset_data(dataset); + GGML_ASSERT(data->ne[0] == inputs->ne[0]); + + const int64_t ndata = data->ne[1]; + const int64_t ndata_batch = inputs->ne[1]; + + GGML_ASSERT(data->ne[1] % inputs->ne[1] == 0); + const int64_t nbatches = ndata/ndata_batch; + + idata_split = idata_split < 0 ? ndata : idata_split; + GGML_ASSERT(idata_split % ndata_batch == 0); + const int64_t ibatch_split = idata_split / ndata_batch; + + int64_t ibatch = 0; + int64_t t_loop_start = ggml_time_us(); + for (; ibatch < ibatch_split; ++ibatch) { + ggml_opt_alloc(opt_ctx, /*backward =*/ true); + ggml_opt_dataset_get_batch(dataset, inputs, labels, ibatch); + ggml_opt_eval(opt_ctx, result_train); + if (callback_train) { + callback_train(true, opt_ctx, dataset, result_train, ibatch+1, ibatch_split, t_loop_start); + } + } + t_loop_start = ggml_time_us(); + for (; ibatch < nbatches; ++ibatch) { + ggml_opt_alloc(opt_ctx, /*backward =*/ false); + ggml_opt_dataset_get_batch(dataset, inputs, labels, ibatch); + ggml_opt_eval(opt_ctx, result_eval); + if (callback_eval) { + callback_eval(false, opt_ctx, dataset, result_eval, ibatch+1-ibatch_split, nbatches-ibatch_split, t_loop_start); + } + } +} + +void ggml_opt_epoch_callback_progress_bar( + bool train, + ggml_opt_context_t opt_ctx, + ggml_opt_dataset_t dataset, + ggml_opt_result_t result, + int64_t ibatch, + int64_t ibatch_max, + int64_t t_start_us) { + fprintf(stderr, "%s[", train ? "train: " : "val: "); + + // The progress bar consists of partially filled blocks, unicode has 8 separate fill levels. + constexpr int64_t bar_length = 8; + const int64_t ibatch8 = 8 * ibatch; + for (int64_t j = 0; j < bar_length; ++j) { + if (ibatch_max * (8*j + 8) / bar_length < ibatch8) { + fprintf(stderr, "\u2588"); // full block + } else if (ibatch_max * (8*j + 7) / bar_length < ibatch8) { + fprintf(stderr, "\u2589"); // 7/8 filled + } else if (ibatch_max * (8*j + 6) / bar_length < ibatch8) { + fprintf(stderr, "\u258A"); // 6/8 filled + } else if (ibatch_max * (8*j + 5) / bar_length < ibatch8) { + fprintf(stderr, "\u258B"); // 5/8 filled + } else if (ibatch_max * (8*j + 4) / bar_length < ibatch8) { + fprintf(stderr, "\u258C"); // 4/8 filled + } else if (ibatch_max * (8*j + 3) / bar_length < ibatch8) { + fprintf(stderr, "\u258D"); // 3/8 filled + } else if (ibatch_max * (8*j + 2) / bar_length < ibatch8) { + fprintf(stderr, "\u258E"); // 2/8 filled + } else if (ibatch_max * (8*j + 1) / bar_length < ibatch8) { + fprintf(stderr, "\u258F"); // 1/8 filled + } else { + fprintf(stderr, " "); + } + } + + const int64_t batch_size = ggml_opt_inputs(opt_ctx)->ne[1]; + const int64_t idata = ibatch*batch_size; + const int64_t idata_max = ibatch_max*batch_size; + + double loss; + double loss_unc; + ggml_opt_result_loss(result, &loss, &loss_unc); + + double accuracy; + double accuracy_unc; + ggml_opt_result_accuracy(result, &accuracy, &accuracy_unc); + + const int64_t t_ibatch_us = ggml_time_us() - t_start_us; + int64_t t_ibatch_s = t_ibatch_us / 1000000; + const int64_t t_ibatch_h = t_ibatch_s / 3600; + t_ibatch_s -= t_ibatch_h * 3600; + const int64_t t_ibatch_m = t_ibatch_s / 60; + t_ibatch_s -= t_ibatch_m * 60; + + const int64_t t_eta_us = t_ibatch_us * (ibatch_max - ibatch)/ibatch; + int64_t t_eta_s = t_eta_us / 1000000; + const int64_t t_eta_h = t_eta_s / 3600; + t_eta_s -= t_eta_h * 3600; + const int64_t t_eta_m = t_eta_s / 60; + t_eta_s -= t_eta_m * 60; + + fprintf(stderr, "] data=%07" PRId64 "/%07" PRId64 " loss=%.5lf±%.5lf acc=%.2lf±%.2lf%% " + "t=%02" PRId64 ":%02" PRId64 ":%02" PRId64 " ETA=%02" PRId64 ":%02" PRId64 ":%02" PRId64 " \r", + idata, idata_max, loss, loss_unc, 100.0*accuracy, 100.0*accuracy_unc, + t_ibatch_h, t_ibatch_m, t_ibatch_s, t_eta_h, t_eta_m, t_eta_s); + if (ibatch == ibatch_max) { + fprintf(stderr, "\n"); + } + fflush(stderr); + + GGML_UNUSED(dataset); +} + +void ggml_opt_fit( + ggml_backend_sched_t backend_sched, + ggml_context * ctx_compute, + ggml_tensor * inputs, + ggml_tensor * outputs, + ggml_opt_dataset_t dataset, + enum ggml_opt_loss_type loss_type, + enum ggml_opt_optimizer_type optimizer, + ggml_opt_get_optimizer_params get_opt_pars, + int64_t nepoch, + int64_t nbatch_logical, + float val_split, + bool silent) { + ggml_time_init(); + const int64_t t_start_us = ggml_time_us(); + + const int64_t ndata = ggml_opt_dataset_data(dataset)->ne[1]; + const int64_t nbatch_physical = inputs->ne[1]; + GGML_ASSERT(ndata % nbatch_logical == 0); + GGML_ASSERT(nbatch_logical % nbatch_physical == 0); + + const int64_t opt_period = nbatch_logical / nbatch_physical; + const int64_t nbatches_logical = ndata / nbatch_logical; + + GGML_ASSERT(val_split >= 0.0f); + GGML_ASSERT(val_split < 1.0f); + const int64_t ibatch_split = int64_t(((1.0f - val_split) * nbatches_logical)) * opt_period; // train <-> val split index (physical) + const int64_t idata_split = ibatch_split * nbatch_physical; + + int64_t epoch = 1; + + ggml_opt_params params = ggml_opt_default_params(backend_sched, loss_type); + params.ctx_compute = ctx_compute; + params.inputs = inputs; + params.outputs = outputs; + params.opt_period = opt_period; + params.get_opt_pars = get_opt_pars; + params.get_opt_pars_ud = &epoch; + params.optimizer = optimizer; + ggml_opt_context_t opt_ctx = ggml_opt_init(params); + + // Shuffling the data is generally useful but there is only a point if not all data is used in a single batch. + if (nbatch_logical < ndata) { + ggml_opt_dataset_shuffle(opt_ctx, dataset, -1); // Shuffle all data (train + validation). + } + + ggml_opt_result_t result_train = ggml_opt_result_init(); + ggml_opt_result_t result_val = ggml_opt_result_init(); + + ggml_opt_epoch_callback epoch_callback = silent ? nullptr : ggml_opt_epoch_callback_progress_bar; + + for (; epoch <= nepoch; ++epoch) { + if (nbatch_logical < idata_split) { + ggml_opt_dataset_shuffle(opt_ctx, dataset, idata_split); + } + + ggml_opt_result_reset(result_train); + ggml_opt_result_reset(result_val); + + if (!silent) { + fprintf(stderr, "%s: epoch %04" PRId64 "/%04" PRId64 ":\n", __func__, epoch, nepoch); + } + ggml_opt_epoch(opt_ctx, dataset, result_train, result_val, idata_split, epoch_callback, epoch_callback); + if (!silent) { + fprintf(stderr, "\n"); + } + } + + if (!silent) { + int64_t t_total_s = (ggml_time_us() - t_start_us) / 1000000; + const int64_t t_total_h = t_total_s / 3600; + t_total_s -= t_total_h * 3600; + const int64_t t_total_m = t_total_s / 60; + t_total_s -= t_total_m * 60; + fprintf(stderr, "%s: training took %02" PRId64 ":%02" PRId64 ":%02" PRId64 "\n", __func__, t_total_h, t_total_m, t_total_s); + } + + ggml_opt_free(opt_ctx); + ggml_opt_result_free(result_train); + ggml_opt_result_free(result_val); +} + +enum ggml_opt_optimizer_type ggml_opt_context_optimizer_type(ggml_opt_context_t c) { + return c->optimizer; +} + +GGML_API const char * ggml_opt_optimizer_name(enum ggml_opt_optimizer_type o) { + switch (o) { + case GGML_OPT_OPTIMIZER_TYPE_ADAMW: + return "adamw"; + case GGML_OPT_OPTIMIZER_TYPE_SGD: + return "sgd"; + default: + return "undefined"; + }; +} + + +// ---- GGUF implementation ---- + +#include "ggml.h.inc" +#include "ggml-backend.h.inc" +#include "ggml-impl-defs.inc" +#include "gguf.h.inc" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define GGUF_MAX_STRING_LENGTH (1024*1024*1024) +#define GGUF_MAX_ARRAY_ELEMENTS (1024*1024*1024) + +#ifdef _WIN32 +# define gguf_ftell _ftelli64 +# define gguf_fseek _fseeki64 +#else +# define gguf_ftell ftello +# define gguf_fseek fseeko +#endif + +template +struct type_to_gguf_type; + +template <> +struct type_to_gguf_type { + static constexpr enum gguf_type value = GGUF_TYPE_UINT8; +}; + +template <> +struct type_to_gguf_type { + static constexpr enum gguf_type value = GGUF_TYPE_INT8; +}; + +template <> +struct type_to_gguf_type { + static constexpr enum gguf_type value = GGUF_TYPE_UINT16; +}; + +template <> +struct type_to_gguf_type { + static constexpr enum gguf_type value = GGUF_TYPE_INT16; +}; + +template <> +struct type_to_gguf_type { + static constexpr enum gguf_type value = GGUF_TYPE_UINT32; +}; + +template <> +struct type_to_gguf_type { + static constexpr enum gguf_type value = GGUF_TYPE_INT32; +}; + +template <> +struct type_to_gguf_type { + static constexpr enum gguf_type value = GGUF_TYPE_FLOAT32; +}; + +template <> +struct type_to_gguf_type { + static constexpr enum gguf_type value = GGUF_TYPE_BOOL; +}; + +template <> +struct type_to_gguf_type { + static constexpr enum gguf_type value = GGUF_TYPE_STRING; +}; + +template <> +struct type_to_gguf_type { + static constexpr enum gguf_type value = GGUF_TYPE_UINT64; +}; + +template <> +struct type_to_gguf_type { + static constexpr enum gguf_type value = GGUF_TYPE_INT64; +}; + +template <> +struct type_to_gguf_type { + static constexpr enum gguf_type value = GGUF_TYPE_FLOAT64; +}; + +static const std::map GGUF_TYPE_SIZE = { + {GGUF_TYPE_UINT8, sizeof(uint8_t)}, + {GGUF_TYPE_INT8, sizeof(int8_t)}, + {GGUF_TYPE_UINT16, sizeof(uint16_t)}, + {GGUF_TYPE_INT16, sizeof(int16_t)}, + {GGUF_TYPE_UINT32, sizeof(uint32_t)}, + {GGUF_TYPE_INT32, sizeof(int32_t)}, + {GGUF_TYPE_FLOAT32, sizeof(float)}, + {GGUF_TYPE_BOOL, sizeof(int8_t)}, + {GGUF_TYPE_STRING, 0}, // undefined + {GGUF_TYPE_ARRAY, 0}, // undefined + {GGUF_TYPE_UINT64, sizeof(uint64_t)}, + {GGUF_TYPE_INT64, sizeof(int64_t)}, + {GGUF_TYPE_FLOAT64, sizeof(double)}, +}; +static_assert(GGUF_TYPE_COUNT == 13, "GGUF_TYPE_COUNT != 13"); + +static const std::map GGUF_TYPE_NAME = { + {GGUF_TYPE_UINT8, "u8"}, + {GGUF_TYPE_INT8, "i8"}, + {GGUF_TYPE_UINT16, "u16"}, + {GGUF_TYPE_INT16, "i16"}, + {GGUF_TYPE_UINT32, "u32"}, + {GGUF_TYPE_INT32, "i32"}, + {GGUF_TYPE_FLOAT32, "f32"}, + {GGUF_TYPE_BOOL, "bool"}, + {GGUF_TYPE_STRING, "str"}, + {GGUF_TYPE_ARRAY, "arr"}, + {GGUF_TYPE_UINT64, "u64"}, + {GGUF_TYPE_INT64, "i64"}, + {GGUF_TYPE_FLOAT64, "f64"}, +}; +static_assert(GGUF_TYPE_COUNT == 13, "GGUF_TYPE_COUNT != 13"); + +size_t gguf_type_size(enum gguf_type type) { + auto it = GGUF_TYPE_SIZE.find(type); + return it == GGUF_TYPE_SIZE.end() ? 0 : it->second; +} + +struct gguf_kv { + std::string key; + + bool is_array; + enum gguf_type type; + + std::vector data; + std::vector data_string; + + template + gguf_kv(const std::string & key, const T value) + : key(key), is_array(false), type(type_to_gguf_type::value) { + GGML_ASSERT(!key.empty()); + data.resize(sizeof(T)); + memcpy(data.data(), &value, sizeof(T)); + } + + template + gguf_kv(const std::string & key, const std::vector & value) + : key(key), is_array(true), type(type_to_gguf_type::value) { + GGML_ASSERT(!key.empty()); + data.resize(value.size()*sizeof(T)); + for (size_t i = 0; i < value.size(); ++i) { + const T tmp = value[i]; + memcpy(data.data() + i*sizeof(T), &tmp, sizeof(T)); + } + } + + gguf_kv(const std::string & key, const std::string & value) + : key(key), is_array(false), type(GGUF_TYPE_STRING) { + GGML_ASSERT(!key.empty()); + data_string.push_back(value); + } + + gguf_kv(const std::string & key, const std::vector & value) + : key(key), is_array(true), type(GGUF_TYPE_STRING) { + GGML_ASSERT(!key.empty()); + data_string = value; + } + + const std::string & get_key() const { + return key; + } + + const enum gguf_type & get_type() const { + return type; + } + + size_t get_ne() const { + if (type == GGUF_TYPE_STRING) { + const size_t ne = data_string.size(); + GGML_ASSERT(is_array || ne == 1); + return ne; + } + const size_t type_size = gguf_type_size(type); + GGML_ASSERT(data.size() % type_size == 0); + const size_t ne = data.size() / type_size; + GGML_ASSERT(is_array || ne == 1); + return ne; + } + + template + const T & get_val(const size_t i = 0) const { + GGML_ASSERT(type_to_gguf_type::value == type); + if constexpr (std::is_same::value) { + GGML_ASSERT(data_string.size() >= i+1); + return data_string[i]; + } + const size_t type_size = gguf_type_size(type); + GGML_ASSERT(data.size() % type_size == 0); + GGML_ASSERT(data.size() >= (i+1)*type_size); + return reinterpret_cast(data.data())[i]; + } + + void cast(const enum gguf_type new_type) { + const size_t new_type_size = gguf_type_size(new_type); + GGML_ASSERT(data.size() % new_type_size == 0); + type = new_type; + } +}; + +struct gguf_tensor_info { + struct ggml_tensor t; // for holding the equivalent info + uint64_t offset; // offset from start of `data`, must be a multiple of `ALIGNMENT` +}; + +struct gguf_context { + uint32_t version = GGUF_VERSION; + + std::vector kv; + std::vector info; + + size_t alignment = GGUF_DEFAULT_ALIGNMENT; + size_t offset = 0; // offset of `data` from beginning of file + size_t size = 0; // size of `data` in bytes + + void * data = nullptr; +}; + +struct gguf_reader { + gguf_reader(FILE * file) : file(file) { + // read the remaining bytes once and update on each read + nbytes_remain = file_remain(file); + } + + // helper for remaining bytes in a file + static uint64_t file_remain(FILE * file) { + const int64_t cur = gguf_ftell(file); + if (cur < 0) { + return 0; + } + if (gguf_fseek(file, 0, SEEK_END) != 0) { + gguf_fseek(file, cur, SEEK_SET); + + return 0; + } + const int64_t end = gguf_ftell(file); + if (end < 0) { + gguf_fseek(file, cur, SEEK_SET); + + return 0; + } + gguf_fseek(file, cur, SEEK_SET); + return static_cast(end - cur); + } + + template + bool read(T & dst) const { + const size_t size = sizeof(dst); + if (nbytes_remain < size) { + return false; + } + const size_t nread = fread(&dst, 1, size, file); + nbytes_remain -= nread; + return nread == size; + } + + template + bool read(std::vector & dst, const size_t n) const { + if (n > GGUF_MAX_ARRAY_ELEMENTS) { + return false; + } + if constexpr (std::is_same::value) { + // strings are prefixed with their length, so we need to account for that + if (n > SIZE_MAX / sizeof(uint64_t)) { + return false; + } + if (nbytes_remain < n * sizeof(uint64_t)) { + return false; + } + } else { + if (n > SIZE_MAX / sizeof(T)) { + return false; + } + if (nbytes_remain < n * sizeof(T)) { + return false; + } + } + dst.resize(n); + for (size_t i = 0; i < dst.size(); ++i) { + if constexpr (std::is_same::value) { + bool tmp; + if (!read(tmp)) { + return false; + } + dst[i] = tmp; + } else { + if (!read(dst[i])) { + return false; + } + } + } + return true; + } + + bool read(bool & dst) const { + int8_t tmp = -1; + if (!read(tmp)) { + return false; + } + dst = tmp != 0; + return true; + } + + bool read(enum ggml_type & dst) const { + int32_t tmp = -1; + if (!read(tmp)) { + return false; + } + dst = ggml_type(tmp); + return true; + } + + bool read(enum gguf_type & dst) const { + int32_t tmp = -1; + if (!read(tmp)) { + return false; + } + dst = gguf_type(tmp); + return true; + } + + bool read(std::string & dst) const { + uint64_t size = 0; + if (!read(size)) { + return false; + } + if (size > GGUF_MAX_STRING_LENGTH) { + GGML_LOG_ERROR("%s: string length %" PRIu64 " exceeds maximum %" PRIu64 "\n", __func__, size, (uint64_t) GGUF_MAX_STRING_LENGTH); + return false; + } + if (size > nbytes_remain) { + GGML_LOG_ERROR("%s: string length %" PRIu64 " exceeds remaining file size %" PRIu64 " bytes\n", __func__, size, nbytes_remain); + return false; + } + dst.resize(static_cast(size)); + const size_t nread = fread(dst.data(), 1, size, file); + nbytes_remain -= nread; + return nread == size; + } + + bool read(void * dst, const size_t size) const { + if (size > nbytes_remain) { + return false; + } + const size_t nread = fread(dst, 1, size, file); + nbytes_remain -= nread; + return nread == size; + } + +private: + FILE * file; + + mutable uint64_t nbytes_remain; +}; + +struct gguf_context * gguf_init_empty(void) { + return new gguf_context; +} + +template +bool gguf_read_emplace_helper(const struct gguf_reader & gr, std::vector & kv, const std::string & key, const bool is_array, const size_t n) { + if (is_array) { + std::vector value; + try { + if (!gr.read(value, n)) { + return false; + } + } catch (std::length_error &) { + GGML_LOG_ERROR("%s: encountered length_error while reading value for key '%s'\n", __func__, key.c_str()); + return false; + } catch (std::bad_alloc &) { + GGML_LOG_ERROR("%s: encountered bad_alloc error while reading value for key '%s'\n", __func__, key.c_str()); + return false; + } + kv.emplace_back(key, value); + } else { + T value; + if (!gr.read(value)) { + return false; + } + kv.emplace_back(key, value); + } + return true; +} + +struct gguf_context * gguf_init_from_file_ptr(FILE * file, struct gguf_init_params params) { + if (!file) { + return nullptr; + } + + const struct gguf_reader gr(file); + struct gguf_context * ctx = new gguf_context; + + bool ok = true; + + // file magic + { + std::vector magic; + ok = ok && gr.read(magic, 4); + + if (!ok) { + GGML_LOG_ERROR("%s: failed to read magic\n", __func__); + gguf_free(ctx); + return nullptr; + } + + for (uint32_t i = 0; i < magic.size(); i++) { + if (magic[i] != GGUF_MAGIC[i]) { + char c0 = isprint(magic[0]) ? magic[0] : '?'; + char c1 = isprint(magic[1]) ? magic[1] : '?'; + char c2 = isprint(magic[2]) ? magic[2] : '?'; + char c3 = isprint(magic[3]) ? magic[3] : '?'; + GGML_LOG_ERROR("%s: invalid magic characters: '%c%c%c%c', expected 'GGUF'\n", __func__, c0, c1, c2, c3); + gguf_free(ctx); + return nullptr; + } + } + } + + // header + int64_t n_kv = 0; + int64_t n_tensors = 0; + + if (ok && gr.read(ctx->version)) { + if (ok && ctx->version == 0) { + GGML_LOG_ERROR("%s: bad GGUF version: %" PRIu32 "\n", __func__, ctx->version); + ok = false; + } + + /* + * bit layout is different when reading non-native endian models. + * assuming that the GGUF version is 3, the non-native endian model + * would read it as 0x30000000. we can use the AND operation against + * the last 4 hexadecimal digits to check if the model is the same + * endianness as the host system. + */ + if (ok && (ctx->version & 0x0000FFFF) == 0x00000000) { + GGML_LOG_ERROR("%s: failed to load model: this GGUF file version %" PRIu32 " is extremely large, is there a mismatch between the host and model endianness?\n", __func__, ctx->version); + ok = false; + } + + if (ok && ctx->version == 1) { + GGML_LOG_ERROR("%s: GGUFv1 is no longer supported, please use a more up-to-date version\n", __func__); + ok = false; + } + if (ok && ctx->version > GGUF_VERSION) { + GGML_LOG_ERROR("%s: this GGUF file is version %" PRIu32 " but this software only supports up to version %d\n", + __func__, ctx->version, GGUF_VERSION); + ok = false; + } + } else { + ok = false; + } + + if (ok && gr.read(n_tensors)) { + static_assert(sizeof(size_t) <= 8 && sizeof(gguf_tensor_info) >= 2, "int64_t insufficient for indexing"); + if (n_tensors < 0 || n_tensors > int64_t(SIZE_MAX/sizeof(gguf_tensor_info))) { + GGML_LOG_ERROR("%s: number of tensors is %" PRIi64 " but must be in [0, %zu]\n", + __func__, n_tensors, SIZE_MAX/sizeof(gguf_tensor_info)); + ok = false; + } + } else { + ok = false; + } + + if (ok && gr.read(n_kv)) { + static_assert(sizeof(size_t) <= 8 && sizeof(gguf_tensor_info) >= 2, "int64_t insufficient for indexing"); + if (n_kv < 0 || n_kv > int64_t(SIZE_MAX/sizeof(gguf_kv))) { + GGML_LOG_ERROR("%s: number of key value pairs is %" PRIi64 " but must be in [0, %zu]\n", + __func__, n_kv, SIZE_MAX/sizeof(gguf_kv)); + ok = false; + } + } else { + ok = false; + } + + if (!ok) { + GGML_LOG_ERROR("%s: failed to read header\n", __func__); + gguf_free(ctx); + return nullptr; + } + + // KV pairs + { + for (int64_t i = 0; ok && i < n_kv; ++i) { + std::string key; + gguf_type type = gguf_type(-1); + bool is_array = false; + uint64_t n = 1; + + try { + ok = ok && gr.read(key); + } catch (std::length_error &) { + GGML_LOG_ERROR("%s: encountered length_error while reading key %" PRIi64 "\n", __func__, i); + ok = false; + } catch (std::bad_alloc &) { + GGML_LOG_ERROR("%s: encountered bad_alloc error while reading key %" PRIi64 "\n", __func__, i); + ok = false; + } + for (size_t j = 0; ok && j < ctx->kv.size(); ++j) { + if (key == ctx->kv[j].key) { + GGML_LOG_ERROR("%s: duplicate key '%s' for tensors %zu and %" PRIi64 " \n", __func__, key.c_str(), j, i); + ok = false; + } + } + if (!ok) { + break; + } + + ok = ok && gr.read(type); + if (type == GGUF_TYPE_ARRAY) { + is_array = true; + ok = ok && gr.read(type); + ok = ok && gr.read(n); + } + if (!ok) { + break; + } + + switch (type) { + case GGUF_TYPE_UINT8: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; + case GGUF_TYPE_INT8: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; + case GGUF_TYPE_UINT16: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; + case GGUF_TYPE_INT16: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; + case GGUF_TYPE_UINT32: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; + case GGUF_TYPE_INT32: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; + case GGUF_TYPE_FLOAT32: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; + case GGUF_TYPE_BOOL: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; + case GGUF_TYPE_STRING: ok = ok && gguf_read_emplace_helper(gr, ctx->kv, key, is_array, n); break; + case GGUF_TYPE_UINT64: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; + case GGUF_TYPE_INT64: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; + case GGUF_TYPE_FLOAT64: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; + case GGUF_TYPE_ARRAY: + default: + { + GGML_LOG_ERROR("%s: key '%s' has invalid GGUF type %d\n", __func__, key.c_str(), type); + ok = false; + } break; + } + } + + if (!ok) { + GGML_LOG_ERROR("%s: failed to read key-value pairs\n", __func__); + gguf_free(ctx); + return nullptr; + } + GGML_ASSERT(int64_t(ctx->kv.size()) == n_kv); + + const int alignment_idx = gguf_find_key(ctx, GGUF_KEY_GENERAL_ALIGNMENT); + ctx->alignment = alignment_idx == -1 ? GGUF_DEFAULT_ALIGNMENT : gguf_get_val_u32(ctx, alignment_idx); + + if (ctx->alignment == 0 || (ctx->alignment & (ctx->alignment - 1)) != 0) { + GGML_LOG_ERROR("%s: alignment %zu is not a power of 2\n", __func__, ctx->alignment); + gguf_free(ctx); + return nullptr; + } + } + + // read the tensor info + for (int64_t i = 0; ok && i < n_tensors; ++i) { + struct gguf_tensor_info info; + + // tensor name + { + std::string name; + try { + ok = ok && gr.read(name); + } catch (std::length_error &) { + GGML_LOG_ERROR("%s: encountered length_error while reading tensor name %" PRIi64 "\n", __func__, i); + ok = false; + } catch (std::bad_alloc &) { + GGML_LOG_ERROR("%s: encountered bad_alloc error while reading tensor name %" PRIi64 "\n", __func__, i); + ok = false; + } + if (name.length() >= GGML_MAX_NAME) { + GGML_LOG_ERROR("%s: tensor name %" PRIi64 " is too long: %zu >= %d\n", __func__, i, name.length(), GGML_MAX_NAME); + ok = false; + break; + } + ggml_set_name(&info.t, name.c_str()); + + // make sure there are no duplicate tensor names + for (int64_t j = 0; ok && j < i; ++j) { + if (strcmp(info.t.name, ctx->info[j].t.name) == 0) { + GGML_LOG_ERROR("%s: duplicate tensor name '%s' for tensors %" PRIi64 " and %" PRIi64 "\n", __func__, info.t.name, j, i); + ok = false; + break; + } + } + } + if (!ok) { + break; + } + + // tensor shape + { + uint32_t n_dims = 0; + ok = ok && gr.read(n_dims); + if (n_dims > GGML_MAX_DIMS) { + GGML_LOG_ERROR("%s: tensor '%s' has invalid number of dimensions: %" PRIu32 " > %" PRIu32 "\n", + __func__, info.t.name, n_dims, GGML_MAX_DIMS); + ok = false; + break; + } + for (uint32_t j = 0; ok && j < GGML_MAX_DIMS; ++j) { + info.t.ne[j] = 1; + if (j < n_dims) { + ok = ok && gr.read(info.t.ne[j]); + } + + // check that all ne are non-negative + if (info.t.ne[j] < 0) { + GGML_LOG_ERROR("%s: tensor '%s' dimension %" PRIu32 " has invalid number of elements: %" PRIi64 " < 0\n", + __func__, info.t.name, j, info.t.ne[j]); + ok = false; + break; + } + } + + // check that the total number of elements is representable + if (ok && ((INT64_MAX/info.t.ne[1] <= info.t.ne[0]) || + (INT64_MAX/info.t.ne[2] <= info.t.ne[0]*info.t.ne[1]) || + (INT64_MAX/info.t.ne[3] <= info.t.ne[0]*info.t.ne[1]*info.t.ne[2]))) { + + GGML_LOG_ERROR("%s: total number of elements in tensor '%s' with shape " + "(%" PRIi64 ", %" PRIi64 ", %" PRIi64 ", %" PRIi64 ") is >= %" PRIi64 "\n", + __func__, info.t.name, info.t.ne[0], info.t.ne[1], info.t.ne[2], info.t.ne[3], INT64_MAX); + ok = false; + break; + } + } + if (!ok) { + break; + } + + // tensor type + { + ok = ok && gr.read(info.t.type); + + // check that tensor type is within defined range + if (info.t.type < 0 || info.t.type >= GGML_TYPE_COUNT) { + GGML_LOG_ERROR("%s: tensor '%s' has invalid ggml type %d. should be in [0, %d)\n", + __func__, info.t.name, info.t.type, GGML_TYPE_COUNT); + ok = false; + break; + } + const size_t type_size = ggml_type_size(info.t.type); + const int64_t blck_size = ggml_blck_size(info.t.type); + + // check that row size is divisible by block size + if (blck_size == 0 || info.t.ne[0] % blck_size != 0) { + GGML_LOG_ERROR("%s: tensor '%s' of type %d (%s) has %" PRId64 " elements per row, " + "not a multiple of block size (%" PRId64 ")\n", + __func__, info.t.name, (int) info.t.type, ggml_type_name(info.t.type), info.t.ne[0], blck_size); + ok = false; + break; + } + + // check that the size of the tensor in bytes is representable + if (ok && uint64_t(ggml_nelements(&info.t)/ggml_blck_size(info.t.type)) > SIZE_MAX/ggml_type_size(info.t.type)) { + GGML_LOG_ERROR("%s: tensor '%s' with shape (%" PRIi64 ", %" PRIi64 ", %" PRIi64 ", %" PRIi64 ") has a size in bytes > %zu\n", + __func__, info.t.name, info.t.ne[0], info.t.ne[1], info.t.ne[2], info.t.ne[3], SIZE_MAX); + ok = false; + break; + } + + // calculate byte offsets given the tensor shape and type + info.t.nb[0] = type_size; + info.t.nb[1] = info.t.nb[0]*(info.t.ne[0]/blck_size); + for (int j = 2; j < GGML_MAX_DIMS; ++j) { + info.t.nb[j] = info.t.nb[j - 1]*info.t.ne[j - 1]; + } + } + if (!ok) { + break; + } + + // tensor data offset within buffer + ok = ok && gr.read(info.offset); + + ctx->info.push_back(info); + } + + if (!ok) { + GGML_LOG_ERROR("%s: failed to read tensor info\n", __func__); + gguf_free(ctx); + return nullptr; + } + GGML_ASSERT(int64_t(ctx->info.size()) == n_tensors); + + // we require the data section to be aligned, so take into account any padding + if (gguf_fseek(file, GGML_PAD(gguf_ftell(file), ctx->alignment), SEEK_SET) != 0) { + GGML_LOG_ERROR("%s: failed to seek to beginning of data section\n", __func__); + gguf_free(ctx); + return nullptr; + } + + // store the current file offset - this is where the data section starts + ctx->offset = gguf_ftell(file); + + // compute the total size of the data section, taking into account the alignment + { + ctx->size = 0; + for (size_t i = 0; i < ctx->info.size(); ++i) { + const gguf_tensor_info & ti = ctx->info[i]; + if (ti.offset != ctx->size) { + GGML_LOG_ERROR("%s: tensor '%s' has offset %" PRIu64 ", expected %zu\n", + __func__, ti.t.name, ti.offset, ctx->size); + GGML_LOG_ERROR("%s: failed to read tensor data\n", __func__); + gguf_free(ctx); + return nullptr; + } + size_t padded_size = GGML_PAD(ggml_nbytes(&ti.t), ctx->alignment); + if (SIZE_MAX - ctx->size < padded_size) { + GGML_LOG_ERROR("%s: tensor '%s' size overflow, cannot accumulate size %zu + %zu\n", + __func__, ti.t.name, ctx->size, padded_size); + gguf_free(ctx); + return nullptr; + } + ctx->size += padded_size; + } + } + + // load the tensor data only if requested + if (params.ctx != nullptr) { + // if the provided gguf_context is no_alloc, then we create "empty" tensors and do not read the binary blob + // otherwise, we load the binary blob into the created ggml_context as well, and point the "data" members of + // the ggml_tensor structs to the appropriate locations in the binary blob + + // compute the exact size needed for the new ggml_context + size_t mem_size = 0; + if (params.no_alloc) { + if (n_tensors != 0 && SIZE_MAX / n_tensors < ggml_tensor_overhead()) { + GGML_LOG_ERROR("%s: memory size overflow while allocating ggml context\n", __func__); + gguf_free(ctx); + return nullptr; + } + + const size_t overhead = n_tensors * ggml_tensor_overhead(); + + mem_size = overhead; + } else { + if ((n_tensors + 1) != 0 && SIZE_MAX / (n_tensors + 1) < ggml_tensor_overhead()) { + GGML_LOG_ERROR("%s: memory size overflow while allocating ggml context\n", __func__); + gguf_free(ctx); + return nullptr; + } + + const size_t overhead = (n_tensors + 1) * ggml_tensor_overhead(); + + if (SIZE_MAX - overhead < ctx->size) { + GGML_LOG_ERROR("%s: memory size overflow while allocating ggml context\n", __func__); + gguf_free(ctx); + return nullptr; + } + + mem_size = overhead + ctx->size; + } + + struct ggml_init_params pdata = { + /*mem_size =*/ mem_size, + /*mem_buffer =*/ nullptr, + /*no_alloc =*/ params.no_alloc, + }; + + *params.ctx = ggml_init(pdata); + if (*params.ctx == nullptr) { + GGML_LOG_ERROR("%s: failed to initialize ggml context for storing tensors\n", __func__); + gguf_free(ctx); + return nullptr; + } + + struct ggml_context * ctx_data = *params.ctx; + + struct ggml_tensor * data = nullptr; + + if (!params.no_alloc) { + data = ggml_new_tensor_1d(ctx_data, GGML_TYPE_I8, ctx->size); + + ok = ok && data != nullptr; + + if (ok) { + ggml_set_name(data, "GGUF tensor data binary blob"); + } + + // read the binary blob with the tensor data + ok = ok && gr.read(data->data, ctx->size); + + if (!ok) { + GGML_LOG_ERROR("%s: failed to read tensor data binary blob\n", __func__); + ggml_free(ctx_data); + *params.ctx = nullptr; + gguf_free(ctx); + return nullptr; + } + + ctx->data = data->data; + } + + ggml_set_no_alloc(ctx_data, true); + + // create the tensors + for (size_t i = 0; i < ctx->info.size(); ++i) { + const struct gguf_tensor_info & info = ctx->info[i]; + + struct ggml_tensor * cur = ggml_new_tensor(ctx_data, info.t.type, GGML_MAX_DIMS, info.t.ne); + + ok = ok && cur != nullptr; + + if (!ok) { + break; + } + + ggml_set_name(cur, info.t.name); + + // point the data member to the appropriate location in the binary blob using the tensor info + if (!params.no_alloc) { + cur->data = (char *) data->data + info.offset; + } + } + + if (!ok) { + GGML_LOG_ERROR("%s: failed to create tensors\n", __func__); + ggml_free(ctx_data); + *params.ctx = nullptr; + gguf_free(ctx); + return nullptr; + } + + ggml_set_no_alloc(ctx_data, params.no_alloc); + } + + return ctx; +} + +struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_params params) { + FILE * file = ggml_fopen(fname, "rb"); + + if (!file) { + GGML_LOG_ERROR("%s: failed to open GGUF file '%s' (%s)\n", __func__, fname, strerror(errno)); + return nullptr; + } + + struct gguf_context * result = gguf_init_from_file_ptr(file, params); + fclose(file); + return result; +} + +void gguf_free(struct gguf_context * ctx) { + if (ctx == nullptr) { + return; + } + delete ctx; +} + +const char * gguf_type_name(enum gguf_type type) { + auto it = GGUF_TYPE_NAME.find(type); + return it == GGUF_TYPE_NAME.end() ? nullptr : it->second; +} + +uint32_t gguf_get_version(const struct gguf_context * ctx) { + return ctx->version; +} + +size_t gguf_get_alignment(const struct gguf_context * ctx) { + return ctx->alignment; +} + +size_t gguf_get_data_offset(const struct gguf_context * ctx) { + return ctx->offset; +} + +int64_t gguf_get_n_kv(const struct gguf_context * ctx) { + return ctx->kv.size(); +} + +int64_t gguf_find_key(const struct gguf_context * ctx, const char * key) { + // return -1 if key not found + int64_t keyfound = -1; + + const int64_t n_kv = gguf_get_n_kv(ctx); + + for (int64_t i = 0; i < n_kv; ++i) { + if (strcmp(key, gguf_get_key(ctx, i)) == 0) { + keyfound = i; + break; + } + } + + return keyfound; +} + +const char * gguf_get_key(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + return ctx->kv[key_id].get_key().c_str(); +} + +enum gguf_type gguf_get_kv_type(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + return ctx->kv[key_id].is_array ? GGUF_TYPE_ARRAY : ctx->kv[key_id].get_type(); +} + +enum gguf_type gguf_get_arr_type(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].is_array); + return ctx->kv[key_id].get_type(); +} + +const void * gguf_get_arr_data(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_type() != GGUF_TYPE_STRING); + return ctx->kv[key_id].data.data(); +} + +const char * gguf_get_arr_str(const struct gguf_context * ctx, int64_t key_id, size_t i) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_type() == GGUF_TYPE_STRING); + return ctx->kv[key_id].data_string[i].c_str(); +} + +size_t gguf_get_arr_n(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + + if (ctx->kv[key_id].type == GGUF_TYPE_STRING) { + return ctx->kv[key_id].data_string.size(); + } + + const size_t type_size = gguf_type_size(ctx->kv[key_id].type); + GGML_ASSERT(ctx->kv[key_id].data.size() % type_size == 0); + return ctx->kv[key_id].data.size() / type_size; +} + +uint8_t gguf_get_val_u8(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + return ctx->kv[key_id].get_val(); +} + +int8_t gguf_get_val_i8(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + return ctx->kv[key_id].get_val(); +} + +uint16_t gguf_get_val_u16(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + return ctx->kv[key_id].get_val(); +} + +int16_t gguf_get_val_i16(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + return ctx->kv[key_id].get_val(); +} + +uint32_t gguf_get_val_u32(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + return ctx->kv[key_id].get_val(); +} + +int32_t gguf_get_val_i32(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + return ctx->kv[key_id].get_val(); +} + +float gguf_get_val_f32(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + return ctx->kv[key_id].get_val(); +} + +uint64_t gguf_get_val_u64(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + return ctx->kv[key_id].get_val(); +} + +int64_t gguf_get_val_i64(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + return ctx->kv[key_id].get_val(); +} + +double gguf_get_val_f64(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + return ctx->kv[key_id].get_val(); +} + +bool gguf_get_val_bool(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + return ctx->kv[key_id].get_val(); +} + +const char * gguf_get_val_str(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + return ctx->kv[key_id].get_val().c_str(); +} + +const void * gguf_get_val_data(const struct gguf_context * ctx, int64_t key_id) { + GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); + GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); + GGML_ASSERT(ctx->kv[key_id].get_type() != GGUF_TYPE_STRING); + return ctx->kv[key_id].data.data(); +} + +int64_t gguf_get_n_tensors(const struct gguf_context * ctx) { + return ctx->info.size(); +} + +int64_t gguf_find_tensor(const struct gguf_context * ctx, const char * name) { + // return -1 if tensor not found + int64_t tensor_id = -1; + + const int64_t n_tensors = gguf_get_n_tensors(ctx); + + for (int64_t i = 0; i < n_tensors; ++i) { + if (strcmp(name, gguf_get_tensor_name(ctx, i)) == 0) { + tensor_id = i; + break; + } + } + + return tensor_id; +} + +size_t gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id) { + GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx)); + return ctx->info[tensor_id].offset; +} + +const char * gguf_get_tensor_name(const struct gguf_context * ctx, int64_t tensor_id) { + GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx)); + return ctx->info[tensor_id].t.name; +} + +enum ggml_type gguf_get_tensor_type(const struct gguf_context * ctx, int64_t tensor_id) { + GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx)); + return ctx->info[tensor_id].t.type; +} + +size_t gguf_get_tensor_size(const struct gguf_context * ctx, int64_t tensor_id) { + GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx)); + return ggml_nbytes(&ctx->info[tensor_id].t); +} + +int64_t gguf_remove_key(struct gguf_context * ctx, const char * key) { + const int64_t key_id = gguf_find_key(ctx, key); + if (key_id >= 0) { + ctx->kv.erase(ctx->kv.begin() + key_id); + } + return key_id; +} + +template +static void gguf_check_reserved_keys(const std::string & key, const T val) { + if (key == GGUF_KEY_GENERAL_ALIGNMENT) { + if constexpr (std::is_same::value) { + GGML_ASSERT(val > 0 && (val & (val - 1)) == 0 && GGUF_KEY_GENERAL_ALIGNMENT " must be power of 2"); + } else { + GGML_UNUSED(val); + GGML_ABORT(GGUF_KEY_GENERAL_ALIGNMENT " must be type u32"); + } + } +} + +void gguf_set_val_u8(struct gguf_context * ctx, const char * key, uint8_t val) { + gguf_check_reserved_keys(key, val); + gguf_remove_key(ctx, key); + ctx->kv.emplace_back(key, val); +} + +void gguf_set_val_i8(struct gguf_context * ctx, const char * key, int8_t val) { + gguf_check_reserved_keys(key, val); + gguf_remove_key(ctx, key); + ctx->kv.emplace_back(key, val); +} + +void gguf_set_val_u16(struct gguf_context * ctx, const char * key, uint16_t val) { + gguf_check_reserved_keys(key, val); + gguf_remove_key(ctx, key); + ctx->kv.emplace_back(key, val); +} + +void gguf_set_val_i16(struct gguf_context * ctx, const char * key, int16_t val) { + gguf_check_reserved_keys(key, val); + gguf_remove_key(ctx, key); + ctx->kv.emplace_back(key, val); +} + +void gguf_set_val_u32(struct gguf_context * ctx, const char * key, uint32_t val) { + gguf_check_reserved_keys(key, val); + gguf_remove_key(ctx, key); + ctx->kv.emplace_back(key, val); +} + +void gguf_set_val_i32(struct gguf_context * ctx, const char * key, int32_t val) { + gguf_check_reserved_keys(key, val); + gguf_remove_key(ctx, key); + ctx->kv.emplace_back(key, val); +} + +void gguf_set_val_f32(struct gguf_context * ctx, const char * key, float val) { + gguf_check_reserved_keys(key, val); + gguf_remove_key(ctx, key); + ctx->kv.emplace_back(key, val); +} + +void gguf_set_val_u64(struct gguf_context * ctx, const char * key, uint64_t val) { + gguf_check_reserved_keys(key, val); + gguf_remove_key(ctx, key); + ctx->kv.emplace_back(key, val); +} + +void gguf_set_val_i64(struct gguf_context * ctx, const char * key, int64_t val) { + gguf_check_reserved_keys(key, val); + gguf_remove_key(ctx, key); + ctx->kv.emplace_back(key, val); +} + +void gguf_set_val_f64(struct gguf_context * ctx, const char * key, double val) { + gguf_check_reserved_keys(key, val); + gguf_remove_key(ctx, key); + ctx->kv.emplace_back(key, val); +} + +void gguf_set_val_bool(struct gguf_context * ctx, const char * key, bool val) { + gguf_check_reserved_keys(key, val); + gguf_remove_key(ctx, key); + ctx->kv.emplace_back(key, val); +} + +void gguf_set_val_str(struct gguf_context * ctx, const char * key, const char * val) { + gguf_check_reserved_keys(key, val); + gguf_remove_key(ctx, key); + ctx->kv.emplace_back(key, std::string(val)); +} + +void gguf_set_arr_data(struct gguf_context * ctx, const char * key, enum gguf_type type, const void * data, size_t n) { + gguf_check_reserved_keys(key, data); + gguf_remove_key(ctx, key); + + const size_t nbytes = n*gguf_type_size(type); + std::vector tmp(nbytes); + if (!tmp.empty()) { + memcpy(tmp.data(), data, nbytes); + } + ctx->kv.emplace_back(key, tmp); + ctx->kv.back().cast(type); +} + +void gguf_set_arr_str(struct gguf_context * ctx, const char * key, const char ** data, size_t n) { + gguf_check_reserved_keys(key, data); + gguf_remove_key(ctx, key); + + std::vector tmp(n); + for (size_t i = 0; i < n; ++i) { + tmp[i] = data[i]; + } + ctx->kv.emplace_back(key, tmp); +} + +// set or add KV pairs from another context +void gguf_set_kv(struct gguf_context * ctx, const struct gguf_context * src) { + const int64_t n_kv = gguf_get_n_kv(src); + for (int64_t i = 0; i < n_kv; ++i) { + const struct gguf_kv & kv = src->kv[i]; + + if (!kv.is_array) { + switch (kv.get_type()) { + case GGUF_TYPE_UINT8: gguf_set_val_u8 (ctx, kv.get_key().c_str(), kv.get_val()); break; + case GGUF_TYPE_INT8: gguf_set_val_i8 (ctx, kv.get_key().c_str(), kv.get_val()); break; + case GGUF_TYPE_UINT16: gguf_set_val_u16 (ctx, kv.get_key().c_str(), kv.get_val()); break; + case GGUF_TYPE_INT16: gguf_set_val_i16 (ctx, kv.get_key().c_str(), kv.get_val()); break; + case GGUF_TYPE_UINT32: gguf_set_val_u32 (ctx, kv.get_key().c_str(), kv.get_val()); break; + case GGUF_TYPE_INT32: gguf_set_val_i32 (ctx, kv.get_key().c_str(), kv.get_val()); break; + case GGUF_TYPE_FLOAT32: gguf_set_val_f32 (ctx, kv.get_key().c_str(), kv.get_val()); break; + case GGUF_TYPE_UINT64: gguf_set_val_u64 (ctx, kv.get_key().c_str(), kv.get_val()); break; + case GGUF_TYPE_INT64: gguf_set_val_i64 (ctx, kv.get_key().c_str(), kv.get_val()); break; + case GGUF_TYPE_FLOAT64: gguf_set_val_f64 (ctx, kv.get_key().c_str(), kv.get_val()); break; + case GGUF_TYPE_BOOL: gguf_set_val_bool(ctx, kv.get_key().c_str(), kv.get_val()); break; + case GGUF_TYPE_STRING: gguf_set_val_str (ctx, kv.get_key().c_str(), kv.get_val().c_str()); break; + case GGUF_TYPE_ARRAY: + default: GGML_ABORT("invalid type"); + } + continue; + } + + const size_t ne = kv.get_ne(); + + switch (kv.get_type()) { + case GGUF_TYPE_UINT8: + case GGUF_TYPE_INT8: + case GGUF_TYPE_UINT16: + case GGUF_TYPE_INT16: + case GGUF_TYPE_UINT32: + case GGUF_TYPE_INT32: + case GGUF_TYPE_FLOAT32: + case GGUF_TYPE_UINT64: + case GGUF_TYPE_INT64: + case GGUF_TYPE_FLOAT64: + case GGUF_TYPE_BOOL: { + gguf_set_arr_data(ctx, kv.get_key().c_str(), kv.get_type(), kv.data.data(), ne); + } break; + case GGUF_TYPE_STRING: { + std::vector tmp(ne); + for (size_t j = 0; j < ne; ++j) { + tmp[j] = kv.data_string[j].c_str(); + } + gguf_set_arr_str(ctx, kv.get_key().c_str(), tmp.data(), ne); + } break; + case GGUF_TYPE_ARRAY: + default: GGML_ABORT("invalid type"); + } + } +} + +void gguf_add_tensor( + struct gguf_context * ctx, + const struct ggml_tensor * tensor) { + GGML_ASSERT(tensor); + if (gguf_find_tensor(ctx, tensor->name) != -1) { + GGML_ABORT("duplicate tensor name: %s", tensor->name); + } + + struct gguf_tensor_info ti; + ti.t = *tensor; + ti.offset = ctx->info.empty() ? 0 : + ctx->info.back().offset + GGML_PAD(ggml_nbytes(&ctx->info.back().t), ctx->alignment); + ctx->info.push_back(ti); +} + +void gguf_set_tensor_type(struct gguf_context * ctx, const char * name, enum ggml_type type) { + const int64_t tensor_id = gguf_find_tensor(ctx, name); + if (tensor_id < 0) { + GGML_ABORT("tensor not found: %s", name); + } + struct ggml_tensor * tensor = &ctx->info[tensor_id].t; + const size_t type_size = ggml_type_size(type); + const int64_t blck_size = ggml_blck_size(type); + + tensor->type = type; + GGML_ASSERT(tensor->ne[0] % blck_size == 0 && "tensor row size not divisible by block size of new type"); + + tensor->nb[0] = type_size; + tensor->nb[1] = tensor->nb[0]*(tensor->ne[0]/blck_size); + for (int i = 2; i < GGML_MAX_DIMS; i++) { + tensor->nb[i] = tensor->nb[i - 1]*tensor->ne[i - 1]; + } + + // update offsets + const int64_t n_tensors = gguf_get_n_tensors(ctx); + for (int64_t i = tensor_id + 1; i < n_tensors; ++i) { + ctx->info[i].offset = ctx->info[i - 1].offset + GGML_PAD(ggml_nbytes(&ctx->info[i - 1].t), ctx->alignment); + } +} + +void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const void * data) { + const int64_t tensor_id = gguf_find_tensor(ctx, name); + if (tensor_id < 0) { + GGML_ABORT("tensor not found: %s", name); + } + + ctx->info[tensor_id].t.data = (void *)(uintptr_t)data; // double cast suppresses warning about casting away const +} + +struct gguf_writer_base { + size_t written_bytes {0u}; + + ~gguf_writer_base(void) = default; + + // we bet on devirtualization + virtual void write(int8_t val) = 0; + virtual void write(const std::vector & val) = 0; + virtual void write_tensor_data(const struct gguf_tensor_info & info, size_t offset_data, size_t alignment) = 0; + + template + void write(const T & val) { + for (size_t i = 0; i < sizeof(val); ++i) { + write(reinterpret_cast(&val)[i]); + } + } + + void write(const bool & val) { + const int8_t val8 = val ? 1 : 0; + write(val8); + } + + void write(const std::string & val) { + { + const uint64_t n = val.length(); + write(n); + } + for (size_t i = 0; i < val.length(); ++i) { + write((val.data())[i]); + } + } + + void write(const char * val) { + write(std::string(val)); + } + + void write(const enum ggml_type & val) { + write(int32_t(val)); + } + + void write(const enum gguf_type & val) { + write(int32_t(val)); + } + + void write(const struct gguf_kv & kv) { + const uint64_t ne = kv.get_ne(); + + write(kv.get_key()); + + if (kv.is_array) { + write(GGUF_TYPE_ARRAY); + write(kv.get_type()); + write(ne); + } else { + write(kv.get_type()); + } + + switch (kv.get_type()) { + case GGUF_TYPE_UINT8: + case GGUF_TYPE_INT8: + case GGUF_TYPE_UINT16: + case GGUF_TYPE_INT16: + case GGUF_TYPE_UINT32: + case GGUF_TYPE_INT32: + case GGUF_TYPE_FLOAT32: + case GGUF_TYPE_UINT64: + case GGUF_TYPE_INT64: + case GGUF_TYPE_FLOAT64: { + write(kv.data); + } break; + case GGUF_TYPE_BOOL: { + for (size_t i = 0; i < ne; ++i) { + write(kv.get_val(i)); + } + } break; + case GGUF_TYPE_STRING: { + for (size_t i = 0; i < ne; ++i) { + write(kv.get_val(i)); + } + } break; + case GGUF_TYPE_ARRAY: + default: GGML_ABORT("invalid type"); + } + } + + void write_tensor_meta(const struct gguf_tensor_info & info) { + write(info.t.name); + + const uint32_t n_dims = ggml_n_dims(&info.t); + write(n_dims); + + for (uint32_t j = 0; j < n_dims; ++j) { + write(info.t.ne[j]); + } + write(info.t.type); + write(info.offset); + } + + void pad(const size_t alignment) { + while (written_bytes % alignment != 0) { + const int8_t zero = 0; + write(zero); + } + } +}; + +// vector buffer based writer +struct gguf_writer_buf final : public gguf_writer_base { + std::vector & buf; + + gguf_writer_buf(std::vector & buf) : buf(buf) {} + + using gguf_writer_base::write; + + void write(const int8_t val) override { + buf.push_back(val); + written_bytes++; + } + + void write(const std::vector & val) override { + buf.insert(buf.end(), val.begin(), val.end()); + written_bytes += val.size(); + } + + void write_tensor_data(const struct gguf_tensor_info & info, const size_t offset_data, const size_t alignment) override { + GGML_ASSERT(buf.size() - offset_data == info.offset); + + GGML_ASSERT(ggml_is_contiguous(&info.t)); + const size_t offset = buf.size(); + const size_t nbytes = ggml_nbytes(&info.t); + + buf.resize(offset + nbytes); + if (info.t.buffer) { + ggml_backend_tensor_get(&info.t, buf.data() + offset, 0, nbytes); + } else { + GGML_ASSERT(info.t.data); + memcpy(buf.data() + offset, info.t.data, nbytes); + } + written_bytes += nbytes; + + pad(alignment); + } +}; + +// file based writer +struct gguf_writer_file final : public gguf_writer_base { + FILE * file; + + gguf_writer_file(FILE* file) : file(file) {} + + using gguf_writer_base::write; + + void write(const int8_t val) override { + const auto real_val = static_cast(val); + const auto ret = fputc(real_val, file); + written_bytes++; + if (ret != real_val) { + throw std::runtime_error("unexpected fputc result '" + std::to_string(ret) + "' instead of '" + std::to_string((int)real_val) + "'"); + } + } + + void write(const std::vector & val) override { + const auto ret = fwrite(val.data(), 1, val.size(), file); + written_bytes += val.size(); + if (ret != val.size()) { + throw std::runtime_error("unexpected fwrite number of bytes written, '" + std::to_string(ret) + "' instead of '" + std::to_string(val.size()) + "'"); + } + } + + void write_tensor_data(const struct gguf_tensor_info & info, const size_t offset_data, const size_t alignment) override { + GGML_ASSERT(written_bytes - offset_data == info.offset); + + GGML_ASSERT(ggml_is_contiguous(&info.t)); + const size_t nbytes = ggml_nbytes(&info.t); + + std::vector buf(nbytes); + if (info.t.buffer) { + ggml_backend_tensor_get(&info.t, buf.data(), 0, nbytes); + } else { + GGML_ASSERT(info.t.data); + memcpy(buf.data(), info.t.data, nbytes); + } + write(buf); + + pad(alignment); + } +}; + +template +static void gguf_write_out(const struct gguf_context * ctx, writer_t & gw, bool only_meta) { + const int64_t n_kv = gguf_get_n_kv(ctx); + const int64_t n_tensors = gguf_get_n_tensors(ctx); + + // write header + gw.write(GGUF_MAGIC[0]); + gw.write(GGUF_MAGIC[1]); + gw.write(GGUF_MAGIC[2]); + gw.write(GGUF_MAGIC[3]); + gw.write(ctx->version); + gw.write(n_tensors); + gw.write(n_kv); + + // write key-value pairs + for (int64_t i = 0; i < n_kv; ++i) { + gw.write(ctx->kv[i]); + } + + // write tensor info + for (int64_t i = 0; i < n_tensors; ++i) { + gw.write_tensor_meta(ctx->info[i]); + } + + // we require the data section to be aligned + gw.pad(ctx->alignment); + + if (only_meta) { + return; + } + + const size_t offset_data = gw.written_bytes; + + // write tensor data + for (int64_t i = 0; i < n_tensors; ++i) { + gw.write_tensor_data(ctx->info[i], offset_data, ctx->alignment); + } +} + +void gguf_write_to_buf(const struct gguf_context * ctx, std::vector & buf, bool only_meta) { + gguf_writer_buf gw(buf); + gguf_write_out(ctx, gw, only_meta); +} + +bool gguf_write_to_file_ptr(const struct gguf_context * ctx, FILE * file, bool only_meta) { + GGML_ASSERT(file); + + try { + gguf_writer_file gw(file); + gguf_write_out(ctx, gw, only_meta); + } catch (const std::runtime_error& ex) { + GGML_LOG_ERROR("%s: failed to write GGUF data: %s\n", __func__, ex.what()); + return false; + } + return true; +} + +bool gguf_write_to_file(const struct gguf_context * ctx, const char * fname, bool only_meta) { + FILE * file = ggml_fopen(fname, "wb"); + + if (!file) { + GGML_LOG_ERROR("%s: failed to open file '%s' for writing GGUF data\n", __func__, fname); + return false; + } + + const bool success = gguf_write_to_file_ptr(ctx, file, only_meta); + if (!success) { + GGML_LOG_ERROR("%s: failed to write GGUF data into '%s'\n", __func__, fname); + } + + fclose(file); + return success; +} + +size_t gguf_get_meta_size(const struct gguf_context * ctx) { + // only return size + std::vector buf; + gguf_write_to_buf(ctx, buf, /*only_meta =*/ true); + return buf.size(); +} + +void gguf_get_meta_data(const struct gguf_context * ctx, void * data) { + std::vector buf; + gguf_write_to_buf(ctx, buf, /*only_meta =*/ true); + memcpy(data, buf.data(), buf.size()); +} diff --git a/ggml/src/ggml-blas/CMakeLists.txt b/ggml/src/ggml-blas/CMakeLists.txt index c27dc174c004..46ace3fa3cd8 100644 --- a/ggml/src/ggml-blas/CMakeLists.txt +++ b/ggml/src/ggml-blas/CMakeLists.txt @@ -11,8 +11,9 @@ find_package(BLAS) if (BLAS_FOUND) message(STATUS "BLAS found, Libraries: ${BLAS_LIBRARIES}") + set_source_files_properties(ggml-blas.cpp.inc PROPERTIES LANGUAGE CXX) ggml_add_backend_library(ggml-blas - ggml-blas.cpp + ggml-blas.cpp.inc ) if (${GGML_BLAS_VENDOR} MATCHES "Apple") diff --git a/ggml/src/ggml-blas/ggml-blas.cpp b/ggml/src/ggml-blas/ggml-blas.cpp.inc similarity index 99% rename from ggml/src/ggml-blas/ggml-blas.cpp rename to ggml/src/ggml-blas/ggml-blas.cpp.inc index b4c735267e04..e1a3cdbcb4a9 100644 --- a/ggml/src/ggml-blas/ggml-blas.cpp +++ b/ggml/src/ggml-blas/ggml-blas.cpp.inc @@ -1,6 +1,6 @@ -#include "ggml-impl.h" -#include "ggml-blas.h" -#include "ggml-backend-impl.h" +#include "ggml-impl-defs.inc" +#include "ggml-blas.h.inc" +#include "ggml-backend-impl-defs.inc" #include #include diff --git a/ggml/src/ggml-cann/CMakeLists.txt b/ggml/src/ggml-cann/CMakeLists.txt index aee5e7b06e51..adce38f68c1f 100755 --- a/ggml/src/ggml-cann/CMakeLists.txt +++ b/ggml/src/ggml-cann/CMakeLists.txt @@ -66,7 +66,8 @@ if (CANN_INSTALL_DIR) acl_op_compiler ) - file(GLOB GGML_SOURCES_CANN "*.cpp") + set(GGML_SOURCES_CANN ggml-cann.cpp.inc) + set_source_files_properties(${GGML_SOURCES_CANN} PROPERTIES LANGUAGE CXX) ggml_add_backend_library(ggml-cann ${GGML_SOURCES_CANN}) target_link_libraries(ggml-cann PRIVATE ${CANN_LIBRARIES}) diff --git a/ggml/src/ggml-cann/acl_tensor.cpp b/ggml/src/ggml-cann/acl_tensor.cpp deleted file mode 100644 index e95d3c4d88df..000000000000 --- a/ggml/src/ggml-cann/acl_tensor.cpp +++ /dev/null @@ -1,195 +0,0 @@ -/* - * Copyright (c) 2023-2026 The ggml authors - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to - * deal in the Software without restriction, including without limitation the - * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or - * sell copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING - * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS - * IN THE SOFTWARE. - */ - -#include "acl_tensor.h" - -#include -#include - -aclDataType ggml_cann_type_mapping(ggml_type type) { - switch (type) { - case GGML_TYPE_F32: - return ACL_FLOAT; - case GGML_TYPE_F16: - return ACL_FLOAT16; - case GGML_TYPE_BF16: - return ACL_BF16; - case GGML_TYPE_I8: - return ACL_INT8; - case GGML_TYPE_I16: - return ACL_INT16; - case GGML_TYPE_I32: - return ACL_INT32; - case GGML_TYPE_Q4_0: - return ACL_INT4; - case GGML_TYPE_Q8_0: - return ACL_INT8; - case GGML_TYPE_I64: - return ACL_INT64; - default: - return ACL_DT_UNDEFINED; - } -} - -acl_tensor_ptr ggml_cann_create_tensor(const ggml_tensor * tensor, - int64_t * ne, - size_t * nb, - int64_t dims, - aclFormat format, - size_t offset) { - // If tensor is bcasted, Up to GGML_MAX_DIMS additional dimensions will be - // added. - int64_t acl_ne[GGML_MAX_DIMS * 2], acl_stride[GGML_MAX_DIMS * 2]; - - if (ne == nullptr) { - for (int i = 0; i < GGML_MAX_DIMS; i++) { - acl_ne[i] = tensor->ne[i]; - // The step size of acl is in elements. - acl_stride[i] = tensor->nb[i] / ggml_element_size(tensor); - } - } else { - // With bcast - for (int i = 0; i < dims; i++) { - acl_ne[i] = ne[i]; - acl_stride[i] = nb[i] / ggml_element_size(tensor); - } - } - - int64_t final_dims = (dims == 0 ? GGML_MAX_DIMS : dims); - int64_t acl_storage_len = 1; - for (int i = 0; i < final_dims; i++) { - acl_storage_len += (acl_ne[i] - 1) * acl_stride[i]; - } - size_t elem_offset = offset / ggml_element_size(tensor); - acl_storage_len += elem_offset; - - // Reverse ne and stride. - std::reverse(acl_ne, acl_ne + final_dims); - std::reverse(acl_stride, acl_stride + final_dims); - - aclTensor * raw = aclCreateTensor(acl_ne, final_dims, ggml_cann_type_mapping(tensor->type), acl_stride, elem_offset, - format, &acl_storage_len, 1, tensor->data); - - return acl_tensor_ptr(raw); -} - -acl_int_array_ptr ggml_cann_create_int_array(const int64_t * value, uint64_t size) { - aclIntArray * raw = aclCreateIntArray(value, size); - return acl_int_array_ptr(raw); -} - -acl_scalar_ptr ggml_cann_create_scalar(void * value, aclDataType dataType) { - aclScalar * raw = aclCreateScalar(value, dataType); - return acl_scalar_ptr(raw); -} - -bool ggml_cann_need_bcast(const ggml_tensor * t0, const ggml_tensor * t1) { - for (int i = 0; i < GGML_MAX_DIMS; i++) { - if (t1->ne[i] != t0->ne[i] && t1->ne[i] != 1) { - return true; - } - } - return false; -} - -int64_t ggml_cann_get_bcast_shape(const ggml_tensor * src0, - const ggml_tensor * src1, - int64_t * bcast_src0_ne, - int64_t * bcast_src1_ne, - size_t * bcast_src0_nb, - size_t * bcast_src1_nb) { - GGML_ASSERT(ggml_can_repeat(src1, src0)); - int bcast_dim_cnt = 0; - for (int i = 0; i < GGML_MAX_DIMS; i++) { - int64_t nr = src0->ne[i] / src1->ne[i]; - bcast_src0_ne[bcast_dim_cnt] = src0->ne[i] / nr; - bcast_src1_ne[bcast_dim_cnt] = src1->ne[i]; - bcast_src0_nb[bcast_dim_cnt] = src0->nb[i]; - bcast_src1_nb[bcast_dim_cnt] = src1->nb[i]; - bcast_dim_cnt++; - if (nr != 1) { - // Need to add an extra dim. - bcast_src0_ne[bcast_dim_cnt] = nr; - bcast_src1_ne[bcast_dim_cnt] = 1; - bcast_src0_nb[bcast_dim_cnt] = bcast_src0_nb[bcast_dim_cnt - 1] * bcast_src0_ne[bcast_dim_cnt - 1]; - bcast_src1_nb[bcast_dim_cnt] = bcast_src1_nb[bcast_dim_cnt - 1] * bcast_src1_ne[bcast_dim_cnt - 1]; - bcast_dim_cnt++; - } - } - return bcast_dim_cnt; -} - -int64_t ggml_cann_get_mulmat_bcast_shape(const int64_t * input_ne, - const int64_t * weight_ne, - const int64_t * dst_ne, - const size_t * input_nb, - const size_t * weight_nb, - const size_t * dst_nb, - int64_t * bcast_input_ne, - int64_t * bcast_weight_ne, - int64_t * bcast_dst_ne, - size_t * bcast_input_nb, - size_t * bcast_weight_nb, - size_t * bcast_dst_nb) { - // input and dst shoule in same shape, except first two dims. - GGML_ASSERT(input_ne[2] == dst_ne[2]); - GGML_ASSERT(input_ne[3] == dst_ne[3]); - - int bcast_dim_cnt = 0; - - // For mul_mat, a dimension needs to be added before the dimension that - // weight needs to be expanded to satisfy the bcast rule of matrix - // multiplication. - for (int i = 0; i < GGML_MAX_DIMS; i++) { - int64_t nr = input_ne[i] / weight_ne[i]; - // Do not use bcast in the first two dimensions because we only support - // the bcast batch dimension. Just copy them. - if (i < 2 || nr == 1) { - bcast_input_ne[bcast_dim_cnt] = input_ne[i]; - bcast_weight_ne[bcast_dim_cnt] = weight_ne[i]; - bcast_dst_ne[bcast_dim_cnt] = dst_ne[i]; - - bcast_input_nb[bcast_dim_cnt] = input_nb[i]; - bcast_weight_nb[bcast_dim_cnt] = weight_nb[i]; - bcast_dst_nb[bcast_dim_cnt] = dst_nb[i]; - bcast_dim_cnt++; - } else { - // Need to add an extra dim. - bcast_input_ne[bcast_dim_cnt] = nr; - bcast_dst_ne[bcast_dim_cnt] = nr; - bcast_weight_ne[bcast_dim_cnt] = 1; - bcast_input_nb[bcast_dim_cnt] = input_nb[i]; - bcast_dst_nb[bcast_dim_cnt] = dst_nb[i]; - bcast_weight_nb[bcast_dim_cnt] = weight_nb[i]; - bcast_dim_cnt++; - - bcast_input_ne[bcast_dim_cnt] = input_ne[i] / nr; - bcast_dst_ne[bcast_dim_cnt] = dst_ne[i] / nr; - bcast_weight_ne[bcast_dim_cnt] = weight_ne[i]; - bcast_input_nb[bcast_dim_cnt] = bcast_input_nb[bcast_dim_cnt - 1] * bcast_input_ne[bcast_dim_cnt - 1]; - bcast_dst_nb[bcast_dim_cnt] = bcast_dst_nb[bcast_dim_cnt - 1] * bcast_dst_ne[bcast_dim_cnt - 1]; - bcast_weight_nb[bcast_dim_cnt] = bcast_weight_nb[bcast_dim_cnt - 1] * bcast_weight_ne[bcast_dim_cnt - 1]; - bcast_dim_cnt++; - } - } - return bcast_dim_cnt; -} diff --git a/ggml/src/ggml-cann/acl_tensor.h b/ggml/src/ggml-cann/acl_tensor.h deleted file mode 100644 index 4737773a4d46..000000000000 --- a/ggml/src/ggml-cann/acl_tensor.h +++ /dev/null @@ -1,349 +0,0 @@ -/* - * Copyright (c) 2023-2026 The ggml authors - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to - * deal in the Software without restriction, including without limitation the - * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or - * sell copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING - * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS - * IN THE SOFTWARE. - */ - -#ifndef CANN_ACL_TENSOR_H -#define CANN_ACL_TENSOR_H - -#include "common.h" - -#include - -#include -#include - -/** - * @brief Maps a ggml_type to its corresponding aclDataType. - * - * @details This function takes a ggml_type as input and returns the corresponding - * aclDataType. It supports mapping for various ggml_types. If the input type - * does not match any of the predefined ggml_types, the function returns - * ACL_DT_UNDEFINED. - * - * @param type The ggml_type to be mapped. - * @return The corresponding aclDataType. If the input type is not recognized, - * ACL_DT_UNDEFINED is returned. - */ -aclDataType ggml_cann_type_mapping(ggml_type type); - -// Deleter for acl objects. -template struct acl_deleter { - void operator()(T * ptr) const noexcept { - if (ptr) { - ACL_CHECK(DestroyFunc(ptr)); - } - } -}; - -using acl_tensor_ptr = std::unique_ptr>; -using acl_int_array_ptr = std::unique_ptr>; -using acl_scalar_ptr = std::unique_ptr>; -using acl_tensor_list_ptr = std::unique_ptr>; - -/** - * @brief Creates an ACL tensor from a ggml_tensor with optional shape. - * - * @details This function creates an ACL tensor based on the properties of the - * provided ggml_tensor. It supports customer shape by adjusting dimensions - * and strides accordingly. If customer shape is applied, additional - * dimensions and strides are calculated based on the provided parameters. - * - * @param tensor Pointer to the ggml_tensor to be converted to ACL tensor. - * @param ne Pointer to an array containing dimensions. Defaults to nullptr - * if no customer shape is applied. - * @param nb Pointer to an array containing strides. Defaults to nullptr - * if no customer shape is applied. - * @param dims Number of dimensions in the tensor. Defaults to 0 if no customer - * shape is applied. - * @param format ACL tensor format. Defaults to ACL_FORMAT_ND. - * @param offset Offset in bytes for the ACL tensor data. Defaults to 0. - * @return Pointer to the created ACL tensor. - */ -acl_tensor_ptr ggml_cann_create_tensor(const ggml_tensor * tensor, - int64_t * ne = nullptr, - size_t * nb = nullptr, - int64_t dims = 0, - aclFormat format = ACL_FORMAT_ND, - size_t offset = 0); - -/** - * @brief Template for creating an ACL tensor from provided parameters. typename TYPE - * should be size_t or float. - * - * @details This function creates an ACL tensor using the provided data pointer, - * data type, dimensions, strides, format, offset, and additional parameters. - * It calculates necessary dimensions and strides based on the provided ne and nb - * arrays, adjusting them for the ACL tensor creation. The ACL storage length - * is also calculated based on the provided dimensions and strides. - * - * @param data_ptr Pointer to the data buffer for the ACL tensor. - * @param dtype ACL data type of the tensor. - * @param type_size Size of each element in the tensor data buffer. - * @param ne Pointer to an array containing tensor dimensions. - * @param nb Pointer to an array containing tensor strides. - * @param dims Number of dimensions of the tensor. - * @param format ACL tensor format. Defaults to ACL_FORMAT_ND. - * @param offset Offset in bytes for the ACL tensor data. Defaults to 0. - * @return Pointer to the created ACL tensor. - */ -template -acl_tensor_ptr ggml_cann_create_tensor(void * data_ptr, - aclDataType dtype, - TYPE type_size, - int64_t * ne, - TYPE * nb, - int64_t dims, - aclFormat format = ACL_FORMAT_ND, - size_t offset = 0) { - int64_t tmp_ne[GGML_MAX_DIMS * 2]; - int64_t tmp_stride[GGML_MAX_DIMS * 2]; - - memcpy(tmp_ne, ne, dims * sizeof(int64_t)); - for (int i = 0; i < dims; i++) { - tmp_stride[i] = nb[i] / type_size; - } - - int64_t acl_storage_len = 1; - for (int i = 0; i < dims; i++) { - acl_storage_len += (tmp_ne[i] - 1) * tmp_stride[i]; - } - - std::reverse(tmp_ne, tmp_ne + dims); - std::reverse(tmp_stride, tmp_stride + dims); - - aclTensor * raw = - aclCreateTensor(tmp_ne, dims, dtype, tmp_stride, offset / type_size, format, &acl_storage_len, 1, data_ptr); - - return acl_tensor_ptr(raw); -} - -/** - * @brief Create an ACL int array resource wrapped in a smart pointer. - * - * This function constructs an aclIntArray from the provided int64_t values - * and returns it as an acl_int_array_ptr (a std::unique_ptr with a custom - * deleter). The returned pointer owns the ACL resource and will automatically - * destroy it via aclDestroyIntArray(). - * - * @param value Pointer to the int64_t elements. - * @param size Number of elements in value. - * - * @return A smart pointer managing the created ACL int array. - */ -acl_int_array_ptr ggml_cann_create_int_array(const int64_t * value, uint64_t size); - -/** - * @brief Create an ACL scalar resource wrapped in a smart pointer. - * - * This function constructs an aclScalar from the raw value pointer and ACL - * data type, then returns it as an acl_scalar_ptr (a std::unique_ptr with - * a custom deleter). The returned pointer owns the ACL scalar and will - * automatically destroy it via aclDestroyScalar(). - * - * @param value Pointer to the raw scalar memory. - * @param dataType ACL data type of the scalar. - * - * @return A smart pointer managing the created ACL scalar. - */ -acl_scalar_ptr ggml_cann_create_scalar(void * value, aclDataType dataType); - -/** - * @brief Create an ACL tensor list from multiple tensor smart pointers. - * - * This function accepts a variadic list of acl_tensor_ptr (a unique_ptr with - * custom deleter) and produces an aclTensorList using aclCreateTensorList(). - * - * The lifecycle management of the tensor objects changes as follows: - * - aclCreateTensorList() takes ownership of the tensors - * - Each input smart pointer releases ownership using release() - * - As a result, the tensors will NOT be destroyed by unique_ptr - * - Instead, they will be destroyed when aclDestroyTensorList() is called - * - * This ensures correct ownership transfer and prevents double-free situations. - * - * @param acl_tensor_ptr Variadic template parameter; each argument must be - * a unique_ptr-like type supporting get() and release(). - * - * @param tensors Variadic list of acl_tensor_ptr objects. Ownership of - * each tensor is transferred away from these smart pointers. - * - * @return A smart pointer (acl_tensor_list_ptr) owning the created ACL tensor list. - * - * @note This implementation is C++11 compatible. The ownership-release process is - * executed using a pack expansion inside an initializer list. - */ -template acl_tensor_list_ptr ggml_cann_create_tensor_list(acl_tensor_ptr &&... tensors) { - aclTensor * raw_tensors[] = { tensors.get()... }; - aclTensorList * raw = aclCreateTensorList(raw_tensors, sizeof...(tensors)); - // aclTensor will release by aclTensorList, so release ownership without - // destroying the tensor - int dummy[] = { (tensors.release(), 0)... }; - GGML_UNUSED(dummy); - return acl_tensor_list_ptr(raw); -} - -/** - * @brief Checks if tensors require broadcasting based on their shapes. - * - * @details This function determines if two ggml_tensors need to be broadcasted for - * element-wise operations. Broadcasting is necessary if the shapes of the - * tensors are not identical and no dimension in either tensor equals 1. - * - * @param t0 Pointer to the first ggml_tensor. - * @param t1 Pointer to the second ggml_tensor. - * @return True if broadcasting is needed, False otherwise. - * - * @remarks This function iterates over the dimensions of t0 and t1. It checks if each - * dimension in t1 differs from t0's corresponding dimension and is not equal - * to 1. If such a dimension is found, broadcasting is required to align t1 - * with t0 for element-wise operations. - */ -bool ggml_cann_need_bcast(const ggml_tensor * t0, const ggml_tensor * t1); - -/** - * @brief Computes broadcast shapes and strides for two ggml_tensors. - * - * @details This function calculates the broadcast shapes and strides for two ggml_tensors, - * following the broadcasting rules similar to numpy. It adjusts dimensions and - * strides to ensure compatibility for element-wise operations where one tensor - * can be broadcasted to match the shape of another tensor. - * - * @param src0 Pointer to the first ggml_tensor. - * @param src1 Pointer to the second ggml_tensor. - * @param bcast_ne_src0 Output array to store broadcasted dimensions for src0. - * @param bcast_ne_src1 Output array to store broadcasted dimensions for src1. - * @param bcast_nb_src0 Output array to store broadcasted strides for src0. - * @param bcast_nb_src1 Output array to store broadcasted strides for src1. - * @return Number of dimensions in the broadcasted shape. - * - * @pre ggml_can_repeat(src1, src0) must return true, indicating src1 can be broadcasted - * to match src0. - * - * @remarks This function iterates over the dimensions of src0 and src1, calculating the - * necessary broadcast dimensions and strides. If a dimension requires broadcasting - * (i.e., its size in src1 is smaller than in src0), an additional dimension is - * added with size calculated to match src0's dimension. This adjustment ensures - * that src1 can be element-wise broadcasted to src0's shape. - * - * How it works: - * - * if dim0 has padding. - * a -> (2, 2) padding = 2 - * a: [[1, 2, *, *] - * [2, 3, *, *]] - * nb = (8, 4, 2) - * - * if a should bcast with b -> (2, 4) - * b' -> (2, 2, 2) - * b : [[1, 2, 3, 4, *, *] - * [5, 6, 7, 8, *, *]] - * nb = (12, 6, 1) - * - * after bcast: - * a' -> (2, 1, 2) - * a': [[[1, 2], *, *] - * [[2, 3], *, *]] - * nb = (8, 4, 2, 1) - * - * b' : [[[1, 2], [3, 4], *, *] - * [[5, 6], [7, 8], *, *]] - * nb = (12, 6, 2, 1) - * \endcode - * - * dim1 in a inserted dim, should add nb for dim1, - * and all other nb moves to next in order. - */ -int64_t ggml_cann_get_bcast_shape(const ggml_tensor * src0, - const ggml_tensor * src1, - int64_t * bcast_ne_src0, - int64_t * bcast_ne_src1, - size_t * bcast_nb_src0, - size_t * bcast_nb_src1); - -// Bcast macro to avoid duplicate code. -#define BCAST_SHAPE(src0, src1) \ - int64_t bcast_##src0##_ne[GGML_MAX_DIMS * 2]; \ - int64_t bcast_##src1##_ne[GGML_MAX_DIMS * 2]; \ - size_t bcast_##src0##_nb[GGML_MAX_DIMS * 2]; \ - size_t bcast_##src1##_nb[GGML_MAX_DIMS * 2]; \ - int64_t bcast_dims = ggml_cann_get_bcast_shape(src0, src1, bcast_##src0##_ne, bcast_##src1##_ne, \ - bcast_##src0##_nb, bcast_##src1##_nb); - -#define BCAST_PARAM(tensor) bcast_##tensor##_ne, bcast_##tensor##_nb, bcast_dims - -/** - * @brief Calculates broadcast shapes for matrix multiplication. - * - * @details This function computes the broadcast shapes required for matrix multiplication - * based on the input, weight, and destination tensor shapes. It ensures that the - * dimensions of weight tensors are expanded appropriately to satisfy matrix - * multiplication broadcast rules. - * - * @param input_ne Array containing the dimensions of the input tensor. - * @param weight_ne Array containing the dimensions of the weight tensor. - * @param dst_ne Array containing the dimensions of the destination tensor. - * @param input_nb Array containing the strides of the input tensor. - * @param weight_nb Array containing the strides of the weight tensor. - * @param dst_nb Array containing the strides of the destination tensor. - * @param bcast_input_ne Output array for broadcasted input tensor dimensions. - * @param bcast_weight_ne Output array for broadcasted weight tensor dimensions. - * @param bcast_dst_ne Output array for broadcasted destination tensor dimensions. - * @param bcast_input_nb Output array for broadcasted input tensor strides. - * @param bcast_weight_nb Output array for broadcasted weight tensor strides. - * @param bcast_dst_nb Output array for broadcasted destination tensor strides. - * @return The number of dimensions in the broadcasted tensors. - * - * @remarks This function iterates over the tensor dimensions and calculates the broadcast - * shapes needed for matrix multiplication. It ensures that dimensions where - * weight tensor requires expansion are appropriately handled to conform with - * broadcasting rules. - * @note compare with ggml_cann_get_bcast_shape, mul_mat broadcast need add this new dim - * before cast dim. - * @sa ggml_cann_get_bcast_shape - */ -int64_t ggml_cann_get_mulmat_bcast_shape(const int64_t * input_ne, - const int64_t * weight_ne, - const int64_t * dst_ne, - const size_t * input_nb, - const size_t * weight_nb, - const size_t * dst_nb, - int64_t * bcast_input_ne, - int64_t * bcast_weight_ne, - int64_t * bcast_dst_ne, - size_t * bcast_input_nb, - size_t * bcast_weight_nb, - size_t * bcast_dst_nb); - -// Bcast macro to avoid duplicate code. -#define BCAST_MUL_MAT_SHAPE(input, weight, dst) \ - int64_t bcast_##input##_ne[GGML_MAX_DIMS * 2]; \ - int64_t bcast_##weight##_ne[GGML_MAX_DIMS * 2]; \ - int64_t bcast_##dst##_ne[GGML_MAX_DIMS * 2]; \ - size_t bcast_##input##_nb[GGML_MAX_DIMS * 2]; \ - size_t bcast_##weight##_nb[GGML_MAX_DIMS * 2]; \ - size_t bcast_##dst##_nb[GGML_MAX_DIMS * 2]; \ - int64_t bcast_dims = ggml_cann_get_mulmat_bcast_shape( \ - input->ne, weight->ne, dst->ne, input->nb, weight->nb, dst->nb, bcast_##input##_ne, bcast_##weight##_ne, \ - bcast_##dst##_ne, bcast_##input##_nb, bcast_##weight##_nb, bcast_##dst##_nb); - -#define BCAST_MUL_MAT_PARAM(tensor) bcast_##tensor##_ne, bcast_##tensor##_nb, bcast_dims - -#endif // CANN_ACL_TENSOR_H diff --git a/ggml/src/ggml-cann/aclnn_ops.cpp b/ggml/src/ggml-cann/aclnn_ops.cpp deleted file mode 100644 index 2dc0f40917d7..000000000000 --- a/ggml/src/ggml-cann/aclnn_ops.cpp +++ /dev/null @@ -1,4436 +0,0 @@ -/* - * Copyright (c) 2023-2026 The ggml authors - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to - * deal in the Software without restriction, including without limitation the - * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or - * sell copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING - * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS - * IN THE SOFTWARE. - */ - -#include "aclnn_ops.h" - -#include "ggml-impl.h" -#include "ggml.h" - - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include - -#define GGML_COMMON_DECL_C - -#include "../ggml-common.h" - -void bcast_shape(ggml_tensor * src0, - ggml_tensor * src1, - ggml_tensor * dst, - acl_tensor_ptr & acl_src0, - acl_tensor_ptr & acl_src1, - acl_tensor_ptr & acl_dst) { - GGML_ASSERT(ggml_are_same_shape(src0, dst) && ggml_can_repeat(src1, src0)); - // Need bcast - if (!ggml_are_same_shape(src0, src1) && ggml_cann_need_bcast(src0, src1)) { - BCAST_SHAPE(src0, src1) - acl_src0 = ggml_cann_create_tensor(src0, BCAST_PARAM(src0)); - acl_src1 = ggml_cann_create_tensor(src1, BCAST_PARAM(src1)); - acl_dst = ggml_cann_create_tensor(dst, BCAST_PARAM(src0)); - } else { - acl_src0 = ggml_cann_create_tensor(src0); - acl_src1 = ggml_cann_create_tensor(src1); - acl_dst = ggml_cann_create_tensor(dst); - } -} - -void ggml_cann_op_unary(std::function unary_op, - ggml_backend_cann_context & ctx, - ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - unary_op(ctx, acl_src.get(), acl_dst.get()); -} - -void ggml_cann_op_unary_gated(std::function unary_op, - ggml_backend_cann_context & ctx, - ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - ggml_tensor * src1 = dst->src[1]; - - GGML_ASSERT(ggml_is_contiguous_1(src0)); - GGML_ASSERT(ggml_is_contiguous_1(dst)); - const int32_t swapped = ggml_get_op_params_i32(dst, 1); - - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - acl_tensor_ptr acl_src0, acl_src1; - if (src1) { - GGML_ASSERT(ggml_is_contiguous_1(src1)); - GGML_ASSERT(src0->type == src1->type); - - acl_src0 = ggml_cann_create_tensor(src0); - acl_src1 = ggml_cann_create_tensor(src1); - } else { - int64_t ne[] = { src0->ne[0] / 2, src0->ne[1], src0->ne[2], src0->ne[3] }; - size_t nb[] = { src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3] }; - acl_src0 = ggml_cann_create_tensor(src0, ne, nb, GGML_MAX_DIMS, ACL_FORMAT_ND, 0); - acl_src1 = ggml_cann_create_tensor(src0, ne, nb, GGML_MAX_DIMS, ACL_FORMAT_ND, ne[0] * ggml_element_size(src0)); - if (swapped) { - std::swap(acl_src0, acl_src1); - } - } - - unary_op(ctx, acl_src0.get(), acl_dst.get()); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceMul, acl_dst.get(), acl_src1.get()); -} - -// Fused SwiGLU using aclnnSwiGlu: splits input along innermost dim, applies -// SiLU to left half, multiplies by right half. -// -// Falls back to the generic two-kernel path when src[1] != nullptr (two -// independent halves) or swapped != 0 (reversed activation order), as -// aclnnSwiGlu only handles the single interleaved tensor in standard order. -// -// CANN tiling for SwiGlu requires (storageShapeDim + viewDims) to be even. -// aclCreateTensor always uses storageShapeDim=1, so viewDims must be odd. -// We use a 3D view (1+3=4, even) to satisfy this constraint while preserving -// correct split semantics along the innermost (ne[0]) dimension. -void ggml_cann_swiglu(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - auto silu_fn = [](ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst) { - GGML_CANN_CALL_ACLNN_OP(ctx, Silu, acl_src, acl_dst); - }; - - const int32_t swapped = ggml_get_op_params_i32(dst, 1); - if (dst->src[1] != nullptr || swapped != 0) { - ggml_cann_op_unary_gated(silu_fn, ctx, dst); - return; - } - - // aclnnSwiGlu requires the split dim (src->ne[0]) to be even; fall back otherwise. - if (dst->src[0]->ne[0] % 2 != 0) { - ggml_cann_op_unary_gated(silu_fn, ctx, dst); - return; - } - - ggml_tensor * src0 = dst->src[0]; - size_t elem_size = ggml_element_size(src0); - - // src0 GGML: [2*ne0, ne1, ne2, ne3] → 3D view [2*ne0, ne1, ne2*ne3] - // CANN reversed: [ne2*ne3, ne1, 2*ne0], split along CANN dim 2 (last). - int64_t ne0_x2 = src0->ne[0]; - int64_t ne1 = src0->ne[1]; - int64_t ne23 = src0->ne[2] * src0->ne[3]; - int64_t src3d_ne[] = { ne0_x2, ne1, ne23 }; - size_t src3d_nb[] = { (size_t)src0->nb[0], (size_t)src0->nb[1], (size_t)src0->nb[2] }; - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src0->data, ggml_cann_type_mapping(src0->type), - elem_size, src3d_ne, src3d_nb, 3); - - // dst GGML: [ne0, ne1, ne2, ne3] → 3D view [ne0, ne1, ne2*ne3] - int64_t ne0 = dst->ne[0]; - int64_t dst3d_ne[] = { ne0, ne1, ne23 }; - size_t dst3d_nb[] = { (size_t)dst->nb[0], (size_t)dst->nb[1], (size_t)dst->nb[2] }; - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst->data, ggml_cann_type_mapping(dst->type), - elem_size, dst3d_ne, dst3d_nb, 3); - - // CANN tensor [ne23, ne1, 2*ne0]: split along CANN dim 2 (last) = 2*ne0. - GGML_CANN_CALL_ACLNN_OP(ctx, SwiGlu, acl_src.get(), (int64_t)2, acl_dst.get()); -} - -// Fused GeGLU using aclnnGeGluV3: splits input along ne[0] (CANN last dim), -// activates the LEFT half with GELU, multiplies by right half. -// approximate: 0=tanh, 1=none(erf). activateLeft=true matches GGML convention. -// outGelu is a required-but-discard output buffer. -// -// Falls back to the generic two-kernel path when src[1] != nullptr (two -// independent halves) or swapped != 0 (reversed activation order), as -// aclnnGeGluV3 only handles the single interleaved tensor in standard order. -void ggml_cann_geglu(ggml_backend_cann_context & ctx, ggml_tensor * dst, int64_t approximate) { - auto gelu_fn = [](ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst) { - GGML_CANN_CALL_ACLNN_OP(ctx, Gelu, acl_src, acl_dst); - }; - - const int32_t swapped = ggml_get_op_params_i32(dst, 1); - if (dst->src[1] != nullptr || swapped != 0) { - ggml_cann_op_unary_gated(gelu_fn, ctx, dst); - return; - } - - // aclnnGeGluV3 requires the split dim (src->ne[0]) to be even; fall back otherwise. - if (dst->src[0]->ne[0] % 2 != 0) { - ggml_cann_op_unary_gated(gelu_fn, ctx, dst); - return; - } - - ggml_tensor * src0 = dst->src[0]; - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - // Allocate a temporary buffer for the required outGelu output (same shape as dst). - // Build contiguous strides since the pool allocation is a fresh buffer. - size_t elem_size = ggml_element_size(dst); - int64_t ne[GGML_MAX_DIMS] = { dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3] }; - size_t nb[GGML_MAX_DIMS]; - nb[0] = elem_size; - for (int i = 1; i < GGML_MAX_DIMS; i++) { - nb[i] = nb[i - 1] * ne[i - 1]; - } - size_t gelu_out_size = nb[GGML_MAX_DIMS - 1] * ne[GGML_MAX_DIMS - 1]; - ggml_cann_pool_alloc gelu_out_alloc(ctx.pool(), gelu_out_size); - - acl_tensor_ptr acl_gelu_out = ggml_cann_create_tensor( - gelu_out_alloc.get(), ggml_cann_type_mapping(dst->type), elem_size, ne, nb, GGML_MAX_DIMS); - // V3 adds activateLeft param; true → Gelu(left)*right, matching GGML convention. - // GGML dim 0 → CANN last dim (index GGML_MAX_DIMS-1 = 3 for 4D tensor). - GGML_CANN_CALL_ACLNN_OP(ctx, GeGluV3, acl_src.get(), (int64_t)(GGML_MAX_DIMS - 1), approximate, true, - acl_dst.get(), acl_gelu_out.get()); -} - -/** - * @brief Repeats elements of a tensor along each dimension according to the - * specified repeat array. - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The source tensor to be repeated. - * @param acl_dst The destination tensor after repeating. - * @param repeat_array The array specifying the number of repetitions along each - * dimension. - */ -static void aclnn_repeat(ggml_backend_cann_context & ctx, - aclTensor * acl_src, - aclTensor * acl_dst, - int64_t * repeat_array) { - // repeat tensor along each dim with repeat_array - acl_int_array_ptr repeats = ggml_cann_create_int_array(repeat_array, GGML_MAX_DIMS); - - GGML_CANN_CALL_ACLNN_OP(ctx, Repeat, acl_src, repeats.get(), acl_dst); -} - -/** - * @brief Casts the data type of a source tensor to a destination tensor. - * - * This function casts the data type of the source tensor `acl_src` to the - * specified data type `cast_data_type` and stores the result in the destination - * tensor `acl_dst`. - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The source tensor whose data type will be casted. - * @param acl_dst The destination tensor where the casted result will be stored. - * @param cast_data_type The target data type to which the source tensor will be - * casted. - */ -static void aclnn_cast(ggml_backend_cann_context & ctx, - aclTensor * acl_src, - aclTensor * acl_dst, - aclDataType cast_data_type) { - GGML_CANN_CALL_ACLNN_OP(ctx, Cast, acl_src, cast_data_type, acl_dst); -} - -void ggml_cann_repeat(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - GGML_ASSERT(ggml_can_repeat(src, dst)); - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - int64_t repeatsArray[] = { dst->ne[3] / src->ne[3], dst->ne[2] / src->ne[2], dst->ne[1] / src->ne[1], - dst->ne[0] / src->ne[0] }; - - aclnn_repeat(ctx, acl_src.get(), acl_dst.get(), repeatsArray); -} - -void aclnn_add(ggml_backend_cann_context & ctx, aclTensor * acl_src0, aclTensor * acl_src1, aclTensor * acl_dst) { - float alphaValue = 1.0f; - acl_scalar_ptr alpha = ggml_cann_create_scalar(&alphaValue, aclDataType::ACL_FLOAT); - if (acl_dst != nullptr) { - GGML_CANN_CALL_ACLNN_OP(ctx, Add, acl_src0, acl_src1, alpha.get(), acl_dst); - } else { - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceAdd, acl_src0, acl_src1, alpha.get()); - } -} - -void aclnn_sub(ggml_backend_cann_context & ctx, aclTensor * acl_src0, aclTensor * acl_src1, aclTensor * acl_dst) { - float alphaValue = 1.0f; - acl_scalar_ptr alpha = ggml_cann_create_scalar(&alphaValue, aclDataType::ACL_FLOAT); - if (acl_dst != nullptr) { - GGML_CANN_CALL_ACLNN_OP(ctx, Sub, acl_src0, acl_src1, alpha.get(), acl_dst); - } else { - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceSub, acl_src0, acl_src1, alpha.get()); - } -} - -void aclnn_mul(ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_other, aclTensor * acl_dst) { - if (acl_dst != nullptr) { - GGML_CANN_CALL_ACLNN_OP(ctx, Mul, acl_src, acl_other, acl_dst); - } else { - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceMul, acl_src, acl_other); - } -} - -void aclnn_div(ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_other, aclTensor * acl_dst) { - if (acl_dst != nullptr) { - GGML_CANN_CALL_ACLNN_OP(ctx, Div, acl_src, acl_other, acl_dst); - } else { - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceDiv, acl_src, acl_other); - } -} - -/** - * @brief Multiplies elements of a tensor by a scalar value, optionally - * in-place. - * - * This function multiplies each element of the source tensor `acl_src` by the - * scalar `scale` and stores the result in the destination tensor `acl_dst`. If - * `inplace` is true, `acl_dst` will not be used and the operation is performed - * in-place on `acl_src`. - * The operation is defined as: - * \f[ - * \text {acl_dst }_i=\text {acl_src }_i \times \text {scale} - * \f] - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The source tensor whose elements will be multiplied. - * @param scale The scalar value by which each element of `acl_src` will be - * multiplied. - * @param acl_dst The destination tensor where the result will be stored if - * `inplace` is false. - * @param inplace Flag indicating whether to perform the operation in-place on - * `acl_src`. - */ -static void aclnn_muls(ggml_backend_cann_context & ctx, - aclTensor * acl_src, - float scale, - aclTensor * acl_dst, - bool inplace) { - acl_scalar_ptr acl_scale = ggml_cann_create_scalar(&scale, aclDataType::ACL_FLOAT); - if (inplace) { - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceMuls, acl_src, acl_scale.get()); - } else { - GGML_CANN_CALL_ACLNN_OP(ctx, Muls, acl_src, acl_scale.get(), acl_dst); - } -} - -void ggml_cann_leaky_relu(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - float negative_slope; - memcpy(&negative_slope, dst->op_params, sizeof(float)); - acl_scalar_ptr acl_negative_slope = ggml_cann_create_scalar(&negative_slope, aclDataType::ACL_FLOAT); - - GGML_CANN_CALL_ACLNN_OP(ctx, LeakyRelu, acl_src.get(), acl_negative_slope.get(), acl_dst.get()); -} - -/** - * @brief Concatenates a list of tensors along a specified dimension and stores - * the result in a destination tensor. - * - * @param ctx The context for the CANN backend operations. - * @param tensorList The list of tensors to be concatenated. - * @param acl_dst The destination tensor where the concatenated result will be - * stored. - * @param concat_dim The dimension along which the tensors will be concatenated. - */ -static void aclnn_concat(ggml_backend_cann_context & ctx, - aclTensorList * tensorList, - aclTensor * acl_dst, - int64_t concat_dim) { - GGML_CANN_CALL_ACLNN_OP(ctx, Cat, tensorList, concat_dim, acl_dst); -} - -void ggml_cann_concat(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - ggml_tensor * src1 = dst->src[1]; - acl_tensor_ptr acl_src0 = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_src1 = ggml_cann_create_tensor(src1); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - const int32_t dim = ggml_get_op_params_i32(dst, 0); - - GGML_ASSERT(dim >= 0 && dim < 4); - int32_t acl_dim = 3 - dim; - - acl_tensor_list_ptr tensor_list = ggml_cann_create_tensor_list(acl_src0, acl_src1); - aclnn_concat(ctx, tensor_list.get(), acl_dst.get(), acl_dim); -} - -/** - * @brief Creates a tensor with values starting from `start`, incremented by - * `step`, and ending before `stop`. - * - * This function performs the operation: - * \f[ - * \text {out }_{i+1}=\text {out }_i+\text {step} - * \f] - * the range is [start, stop). - * - * @param ctx The context for the CANN backend operations. - * @param acl_dst The destination tensor where the values will be stored. - * @param start The starting value of the range. - * @param stop The ending value of the range (exclusive). - * @param step The step size between consecutive values. - * @param n_elements The number of elements in the destination tensor. - */ -static void aclnn_arange(ggml_backend_cann_context & ctx, - aclTensor * acl_dst, - float start, - float stop, - float step, - int64_t n_elements) { - int64_t steps = (int64_t) std::ceil((stop - start) / step); - GGML_ASSERT(n_elements == steps); - - acl_scalar_ptr acl_start = ggml_cann_create_scalar(&start, aclDataType::ACL_FLOAT); - acl_scalar_ptr acl_end = ggml_cann_create_scalar(&stop, aclDataType::ACL_FLOAT); - acl_scalar_ptr acl_step = ggml_cann_create_scalar(&step, aclDataType::ACL_FLOAT); - - GGML_CANN_CALL_ACLNN_OP(ctx, Arange, acl_start.get(), acl_end.get(), acl_step.get(), acl_dst); -} - -void ggml_cann_arange(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - GGML_ASSERT(dst->type == GGML_TYPE_F32); - - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - int64_t n_elements = ggml_nelements(dst); - float start; - float stop; - float step; - memcpy(&start, (float *) dst->op_params + 0, sizeof(float)); - memcpy(&stop, (float *) dst->op_params + 1, sizeof(float)); - memcpy(&step, (float *) dst->op_params + 2, sizeof(float)); - - aclnn_arange(ctx, acl_dst.get(), start, stop, step, n_elements); -} - -void ggml_cann_clamp(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - - float min; - float max; - memcpy(&min, dst->op_params, sizeof(float)); - memcpy(&max, (float *) dst->op_params + 1, sizeof(float)); - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - acl_scalar_ptr acl_min = ggml_cann_create_scalar(&min, aclDataType::ACL_FLOAT); - acl_scalar_ptr acl_max = ggml_cann_create_scalar(&max, aclDataType::ACL_FLOAT); - - GGML_CANN_CALL_ACLNN_OP(ctx, Clamp, acl_src.get(), acl_min.get(), acl_max.get(), acl_dst.get()); -} - -void ggml_cann_scale(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - - // scale factor - float v; - memcpy(&v, dst->op_params, sizeof(float)); - - acl_scalar_ptr scale = ggml_cann_create_scalar(&v, aclDataType::ACL_FLOAT); - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - GGML_CANN_CALL_ACLNN_OP(ctx, Muls, acl_src.get(), scale.get(), acl_dst.get()); -} - -void ggml_cann_argsort(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - enum ggml_sort_order order = (enum ggml_sort_order) dst->op_params[0]; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - ggml_cann_pool_alloc temp_buffer_allocator(ctx.pool(), ggml_nelements(dst) * sizeof(int64_t)); - void * buffer = temp_buffer_allocator.get(); - acl_tensor_ptr tmp_tensor = - ggml_cann_create_tensor(buffer, ACL_INT64, ggml_type_size(dst->type), dst->ne, dst->nb, GGML_MAX_DIMS); - GGML_CANN_CALL_ACLNN_OP(ctx, Argsort, acl_src.get(), -1, (order == GGML_SORT_ORDER_DESC ? true : false), - tmp_tensor.get()); - GGML_CANN_CALL_ACLNN_OP(ctx, Cast, tmp_tensor.get(), ggml_cann_type_mapping(dst->type), acl_dst.get()); -} - -void ggml_cann_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - float eps; - memcpy(&eps, dst->op_params, sizeof(float)); - - std::vector normData = { dst->ne[0] }; - acl_int_array_ptr norm = ggml_cann_create_int_array(normData.data(), normData.size()); - GGML_CANN_CALL_ACLNN_OP(ctx, LayerNorm, acl_src.get(), norm.get(), nullptr, nullptr, eps, acl_dst.get(), nullptr, - nullptr); -} - -void ggml_cann_l2_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - - float eps; - memcpy(&eps, dst->op_params, sizeof(float)); - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - size_t type_size = ggml_type_size(src->type); - int64_t n_bytes = src->ne[3] * src->ne[2] * src->ne[1] * type_size; - ggml_cann_pool_alloc temp_buffer_allocator(ctx.pool(), n_bytes); - void * buffer = temp_buffer_allocator.get(); - - int64_t norm_ne[] = { 1, src->ne[1], src->ne[2], src->ne[3] }; - size_t norm_nb[GGML_MAX_DIMS]; - norm_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS; ++i) { - norm_nb[i] = norm_nb[i - 1] * norm_ne[i - 1]; - } - acl_tensor_ptr acl_norm = ggml_cann_create_tensor(buffer, ACL_FLOAT, sizeof(float), norm_ne, norm_nb, GGML_MAX_DIMS); - - std::vector norm_dims = { 3 }; - acl_int_array_ptr dims_array = ggml_cann_create_int_array(norm_dims.data(), norm_dims.size()); - - float p_value = 2.0f; - acl_scalar_ptr p_scalar = ggml_cann_create_scalar(&p_value, aclDataType::ACL_FLOAT); - GGML_CANN_CALL_ACLNN_OP(ctx, Norm, acl_src.get(), p_scalar.get(), dims_array.get(), true, acl_norm.get()); - - ggml_cann_pool_alloc clamp_buffer_allocator(ctx.pool()); - acl_tensor_ptr acl_clamped; - - if (eps > 0.0f) { - void * clamp_buf = clamp_buffer_allocator.alloc(n_bytes); - acl_clamped = ggml_cann_create_tensor(clamp_buf, ACL_FLOAT, sizeof(float), norm_ne, norm_nb, GGML_MAX_DIMS); - acl_scalar_ptr eps_scalar = ggml_cann_create_scalar(&eps, aclDataType::ACL_FLOAT); - GGML_CANN_CALL_ACLNN_OP(ctx, ClampMin, acl_norm.get(), eps_scalar.get(), acl_clamped.get()); - } - - aclTensor * acl_div_input = acl_clamped ? acl_clamped.get() : acl_norm.get(); - GGML_CANN_CALL_ACLNN_OP(ctx, Div, acl_src.get(), acl_div_input, acl_dst.get()); -} - -void ggml_cann_cross_entropy_loss(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - ggml_tensor * src1 = dst->src[1]; - - const int64_t nc = src0->ne[0]; - const int64_t nr = ggml_nrows(src0); - - int64_t logits_ne[] = { nc, nr }; - size_t logits_nb[2]; - logits_nb[0] = ggml_type_size(src0->type); - logits_nb[1] = logits_nb[0] * logits_ne[0]; - acl_tensor_ptr acl_logits = ggml_cann_create_tensor(src0->data, ACL_FLOAT, sizeof(float), logits_ne, logits_nb, 2); - - int64_t labels_ne[] = { nc, nr }; - size_t labels_nb[2]; - labels_nb[0] = ggml_type_size(src1->type); - labels_nb[1] = labels_nb[0] * labels_ne[0]; - acl_tensor_ptr acl_labels = ggml_cann_create_tensor(src1->data, ACL_FLOAT, sizeof(float), labels_ne, labels_nb, 2); - - size_t loss_per_sample_type_size = sizeof(float); - int64_t loss_per_sample_n_bytes = nr * loss_per_sample_type_size; - ggml_cann_pool_alloc loss_per_sample_allocator(ctx.pool(), loss_per_sample_n_bytes); - void * loss_per_sample_buffer = loss_per_sample_allocator.get(); - - int64_t loss_per_sample_ne[] = { nr }; - size_t loss_per_sample_nb[1]; - loss_per_sample_nb[0] = loss_per_sample_type_size; - acl_tensor_ptr acl_loss_per_sample = ggml_cann_create_tensor( - loss_per_sample_buffer, ACL_FLOAT, loss_per_sample_type_size, loss_per_sample_ne, loss_per_sample_nb, 1); - - size_t backprop_n_bytes = nr * nc * sizeof(float); - ggml_cann_pool_alloc backprop_allocator(ctx.pool(), backprop_n_bytes); - void * backprop_buffer = backprop_allocator.get(); - acl_tensor_ptr acl_backprop = ggml_cann_create_tensor(backprop_buffer, ACL_FLOAT, sizeof(float), logits_ne, logits_nb, 2); - - GGML_CANN_CALL_ACLNN_OP(ctx, SoftmaxCrossEntropyWithLogits, acl_logits.get(), acl_labels.get(), - acl_loss_per_sample.get(), acl_backprop.get()); - - size_t total_sum_type_size = sizeof(float); - int64_t total_sum_n_bytes = 1 * total_sum_type_size; - ggml_cann_pool_alloc total_sum_allocator(ctx.pool(), total_sum_n_bytes); - void * total_sum_buffer = total_sum_allocator.get(); - - int64_t total_sum_ne[] = { 1 }; - size_t total_sum_nb[1]; - total_sum_nb[0] = total_sum_type_size; - - acl_tensor_ptr acl_total_sum = - ggml_cann_create_tensor(total_sum_buffer, ACL_FLOAT, total_sum_type_size, total_sum_ne, total_sum_nb, 1); - - std::vector total_sum_dims = { 0 }; - acl_int_array_ptr total_sum_dims_array = ggml_cann_create_int_array(total_sum_dims.data(), total_sum_dims.size()); - bool keep_dims = false; - - GGML_CANN_CALL_ACLNN_OP(ctx, ReduceSum, acl_loss_per_sample.get(), total_sum_dims_array.get(), keep_dims, ACL_FLOAT, - acl_total_sum.get()); - - float value = 1.0f / static_cast(nr); - acl_scalar_ptr scale_factor = ggml_cann_create_scalar(&value, aclDataType::ACL_FLOAT); - acl_tensor_ptr acl_dst = - ggml_cann_create_tensor(dst->data, ACL_FLOAT, sizeof(float), total_sum_ne, total_sum_nb, 1); - - GGML_CANN_CALL_ACLNN_OP(ctx, Muls, acl_total_sum.get(), scale_factor.get(), acl_dst.get()); -} - -void ggml_cann_group_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - int n_groups = dst->op_params[0]; - - float eps; - memcpy(&eps, dst->op_params + 1, sizeof(float)); - - int64_t N = src->ne[3]; - int64_t C = src->ne[2]; - int64_t HxW = src->ne[1] * src->ne[0]; - - size_t type_size = ggml_type_size(src->type); - int64_t ne[] = { n_groups, N }; - size_t nb[] = { type_size, type_size * n_groups }; - size_t n_bytes = N * n_groups; - - ggml_cann_pool_alloc temp_buffer_allocator(ctx.pool(), n_bytes * 2); - void * buffer = temp_buffer_allocator.get(); - acl_tensor_ptr acl_mean_out = ggml_cann_create_tensor(buffer, ACL_FLOAT, type_size, ne, nb, ACL_FORMAT_ND); - acl_tensor_ptr acl_rstd_out = - ggml_cann_create_tensor((char *) buffer + n_bytes, ACL_FLOAT, type_size, ne, nb, ACL_FORMAT_ND); - - GGML_CANN_CALL_ACLNN_OP(ctx, GroupNorm, acl_src.get(), nullptr, nullptr, N, C, HxW, n_groups, eps, acl_dst.get(), - acl_mean_out.get(), acl_rstd_out.get()); -} - -void ggml_cann_set(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - ggml_tensor * src1 = dst->src[1]; - - size_t nb1 = ((int32_t *) dst->op_params)[0]; - size_t nb2 = ((int32_t *) dst->op_params)[1]; - size_t nb3 = ((int32_t *) dst->op_params)[2]; - size_t offset = ((int32_t *) dst->op_params)[3]; - bool inplace = (bool) ((int32_t *) dst->op_params)[4]; - - size_t param_nb[] = { ggml_element_size(src0), nb1, nb2, nb3 }; - - // Create a view of dst at the target offset with src1's dimensions - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst, src1->ne, param_nb, GGML_MAX_DIMS, ACL_FORMAT_ND, offset); - acl_tensor_ptr acl_src1 = ggml_cann_create_tensor(src1); - - if (!inplace) { - // First copy src0 to dst entirely - size_t cpy_size = ggml_nbytes(dst); - ACL_CHECK( - aclrtMemcpyAsync(dst->data, cpy_size, src0->data, cpy_size, ACL_MEMCPY_DEVICE_TO_DEVICE, ctx.stream())); - } - - // Copy src1 into the target region of dst - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceCopy, acl_dst.get(), acl_src1.get()); -} - -void ggml_cann_acc(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - ggml_tensor * src1 = dst->src[1]; - - size_t nb1 = ((int32_t *) dst->op_params)[0]; - size_t nb2 = ((int32_t *) dst->op_params)[1]; - size_t nb3 = ((int32_t *) dst->op_params)[2]; - size_t offset = ((int32_t *) dst->op_params)[3]; - bool inplace = (bool) ((int32_t *) dst->op_params)[4]; - - size_t param_nb[] = { ggml_element_size(src0), nb1, nb2, nb3 }; - - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst, src1->ne, param_nb, GGML_MAX_DIMS, ACL_FORMAT_ND, offset); - acl_tensor_ptr acl_src1 = ggml_cann_create_tensor(src1); - - acl_scalar_ptr alpha = nullptr; - float alphaValue = 1.0f; - alpha = ggml_cann_create_scalar(&alphaValue, aclDataType::ACL_FLOAT); - - if (!inplace) { - size_t cpy_size = ggml_nbytes(dst); - ACL_CHECK( - aclrtMemcpyAsync(dst->data, cpy_size, src0->data, cpy_size, ACL_MEMCPY_DEVICE_TO_DEVICE, ctx.stream())); - acl_tensor_ptr acl_src0 = - ggml_cann_create_tensor(src0, src1->ne, src0->nb, GGML_MAX_DIMS, ACL_FORMAT_ND, offset); - - GGML_CANN_CALL_ACLNN_OP(ctx, Add, acl_src0.get(), acl_src1.get(), alpha.get(), acl_dst.get()); - } else { - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceAdd, acl_dst.get(), acl_src1.get(), alpha.get()); - } -} - -/** - * @brief Performs sum reduction on a given tensor along specified dimensions. - * - * This function reduces the input tensor by summing along the specified dimensions. - * - * @param ctx The context for the CANN backend operations. - * @param dst The destination tensor where the reduced result will be stored. - * @param dim An array of dimension indices. - * @param dim_size The number of dimensions. - */ -static void aclnn_reduce_sum(ggml_backend_cann_context & ctx, ggml_tensor * dst, int64_t * dim, size_t dim_size) { - GGML_ASSERT(dst->ne[0] == 1); - ggml_tensor * src = dst->src[0]; - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - acl_int_array_ptr reduce_dims = ggml_cann_create_int_array(dim, dim_size); - - GGML_CANN_CALL_ACLNN_OP(ctx, ReduceSum, acl_src.get(), reduce_dims.get(), true, ggml_cann_type_mapping(dst->type), - acl_dst.get()); -} - -void ggml_cann_sum_rows(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - int64_t reduce_dims[] = { 3 }; - aclnn_reduce_sum(ctx, dst, reduce_dims, 1); -} - -void ggml_cann_sum(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - int64_t reduce_dims[] = { 0, 1, 2, 3 }; - aclnn_reduce_sum(ctx, dst, reduce_dims, 4); -} - -void ggml_cann_cumsum(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - // GGML cumsum operates along dim 0 (innermost / ne[0]). - // ggml_cann_create_tensor reverses dimensions to [ne3,ne2,ne1,ne0], - // so GGML dim 0 maps to CANN dim 3 (the last dim of the 4-D tensor). - GGML_CANN_CALL_ACLNN_OP(ctx, Cumsum, acl_src.get(), (int64_t)3, - ggml_cann_type_mapping(dst->type), acl_dst.get()); -} - -void ggml_cann_solve_tri(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; // A: [N, N, B2, B3] lower triangular - ggml_tensor * src1 = dst->src[1]; // B: [K, N, B2, B3] - - acl_tensor_ptr acl_a = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_b = ggml_cann_create_tensor(src1); - acl_tensor_ptr acl_x = ggml_cann_create_tensor(dst); - - // mOut: triangular copy of A (required output), same shape as A. - const size_t a_bytes = ggml_nbytes(src0); - ggml_cann_pool_alloc m_alloc(ctx.pool(), a_bytes); - acl_tensor_ptr acl_m = ggml_cann_create_tensor( - m_alloc.get(), ggml_cann_type_mapping(src0->type), - ggml_type_size(src0->type), src0->ne, src0->nb, GGML_MAX_DIMS); - - // Solve AX = B: upper=false (lower tri), transpose=false, unitriangular=false. - GGML_CANN_CALL_ACLNN_OP(ctx, TriangularSolve, - acl_b.get(), acl_a.get(), false, false, false, - acl_x.get(), acl_m.get()); -} - -void ggml_cann_diag(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - - GGML_ASSERT(src->ne[1] == 1); - - const int64_t N = src->ne[0]; - const int64_t n_batch = src->ne[2] * src->ne[3]; - const size_t nb_f32 = sizeof(float); - - // Fill dst with zeros. - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - { - float zero = 0.0f; - acl_scalar_ptr acl_zero = ggml_cann_create_scalar(&zero, ACL_FLOAT); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceFillScalar, acl_dst.get(), acl_zero.get()); - } - - // Copy src vector onto the diagonal of dst via strided views. - // src viewed as [N, n_batch], contiguous strides. - int64_t ne_vec[2] = { N, n_batch }; - size_t nb_src_vec[2] = { nb_f32, N * nb_f32 }; - // dst diagonal view: stride (N+1)*4 steps along the diagonal. - size_t nb_dst_diag[2] = { (N + 1) * nb_f32, N * N * nb_f32 }; - - acl_tensor_ptr acl_src_vec = ggml_cann_create_tensor(src->data, ACL_FLOAT, nb_f32, ne_vec, nb_src_vec, 2); - acl_tensor_ptr acl_dst_diag = ggml_cann_create_tensor(dst->data, ACL_FLOAT, nb_f32, ne_vec, nb_dst_diag, 2); - - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceCopy, acl_dst_diag.get(), acl_src_vec.get()); -} - -void ggml_cann_fill(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - float c = ggml_get_op_params_f32(dst, 0); - - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - acl_scalar_ptr acl_c = ggml_cann_create_scalar(&c, ACL_FLOAT); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceFillScalar, acl_dst.get(), acl_c.get()); -} - -void ggml_cann_tri(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - - const int64_t S = src->ne[0]; - const int64_t n_batch = src->ne[2] * src->ne[3]; - const size_t nb_f32 = sizeof(float); - - int64_t ne3d[3] = { S, S, n_batch }; - size_t nb3d[3] = { nb_f32, S * nb_f32, S * S * nb_f32 }; - - const ggml_tri_type ttype = (ggml_tri_type) ggml_get_op_params_i32(dst, 0); - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src->data, ACL_FLOAT, nb_f32, ne3d, nb3d, 3); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst->data, ACL_FLOAT, nb_f32, ne3d, nb3d, 3); - - switch (ttype) { - case GGML_TRI_TYPE_LOWER: - // Tril(-1): preserve row > col (strict lower), zero upper + diagonal. - GGML_CANN_CALL_ACLNN_OP(ctx, Tril, acl_src.get(), (int64_t)-1, acl_dst.get()); - break; - case GGML_TRI_TYPE_UPPER_DIAG: - // Triu(0): preserve row <= col (upper + diagonal), zero strict lower. - GGML_CANN_CALL_ACLNN_OP(ctx, Triu, acl_src.get(), (int64_t)0, acl_dst.get()); - break; - case GGML_TRI_TYPE_UPPER: - // Triu(1): preserve row < col (strict upper), zero lower + diagonal. - GGML_CANN_CALL_ACLNN_OP(ctx, Triu, acl_src.get(), (int64_t)1, acl_dst.get()); - break; - case GGML_TRI_TYPE_LOWER_DIAG: - // Tril(0): preserve row >= col (lower + diagonal), zero strict upper. - GGML_CANN_CALL_ACLNN_OP(ctx, Tril, acl_src.get(), (int64_t)0, acl_dst.get()); - break; - default: - GGML_ABORT("unsupported tri type"); - } -} - -void ggml_cann_upsample_nearest2d(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src, nullptr, nullptr, 0, ACL_FORMAT_NCHW); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst, nullptr, nullptr, 0, ACL_FORMAT_NCHW); - - std::vector output_size{ dst->ne[1], dst->ne[0] }; - acl_int_array_ptr output_size_array = ggml_cann_create_int_array(output_size.data(), 2); - - GGML_CANN_CALL_ACLNN_OP(ctx, UpsampleNearest2d, acl_src.get(), output_size_array.get(), acl_dst.get()); -} - -/** - * @brief Pads a tensor with a specified value along each dimension. - * - * This function performs padding of the source tensor `acl_src` and stores the - * result in the destination tensor `acl_dst`. The padding values for each - * dimension are specified in the `paddings` array. - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The source tensor to be padded. - * @param acl_dst The destination tensor where the padded result will be stored. - * @param paddings An array specifying the padding values for each dimension. - * The size of the array should be twice the number of dimensions of the tensor. - * @param value The value to be used for padding. The default value is 0.0. - */ -static void aclnn_pad(ggml_backend_cann_context & ctx, - aclTensor * acl_src, - aclTensor * acl_dst, - int64_t * paddings, - float value = 0.0f) { - acl_int_array_ptr acl_pad = ggml_cann_create_int_array(paddings, GGML_MAX_DIMS * 2); - acl_scalar_ptr acl_value = ggml_cann_create_scalar(&value, aclDataType::ACL_FLOAT); - - GGML_CANN_CALL_ACLNN_OP(ctx, ConstantPadNd, acl_src, acl_pad.get(), acl_value.get(), acl_dst); -} - -void ggml_cann_pad(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - // padding: value in the array means how much distance will be padding. - // the position of elements in the array means which dirction to padding, - // each position means: [dim0.front, dim0.behind, dim1.front, dim1.behind, - // dim2.front, dim2.behind, dim3.front, dim3.behind] - const int32_t lp0 = ggml_get_op_params_i32(dst, 0); - const int32_t rp0 = ggml_get_op_params_i32(dst, 1); - const int32_t lp1 = ggml_get_op_params_i32(dst, 2); - const int32_t rp1 = ggml_get_op_params_i32(dst, 3); - const int32_t lp2 = ggml_get_op_params_i32(dst, 4); - const int32_t rp2 = ggml_get_op_params_i32(dst, 5); - const int32_t lp3 = ggml_get_op_params_i32(dst, 6); - const int32_t rp3 = ggml_get_op_params_i32(dst, 7); - - int64_t paddings[] = { lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3 }; - aclnn_pad(ctx, acl_src.get(), acl_dst.get(), paddings); -} - -/** - * @brief Performs 2D average pooling on the input tensor and stores the result - * in the destination tensor. - * - * This function performs average pooling on the source tensor and stores the - * result in the destination tensor. The pooling parameters (kernel size, - * strides, padding) are specified in the `op_params` of the destination tensor. - * - * @param ctx The context for the CANN backend operations. - * @param dst The destination tensor where the result will be stored. The source - * tensor is referenced by `dst->src[0]`. - */ -static void ggml_cann_avg_pool2d(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - GGML_ASSERT(src->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src, nullptr, nullptr, 0, ACL_FORMAT_NCHW); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst, nullptr, nullptr, 0, ACL_FORMAT_NCHW); - - const int32_t * opts = (const int32_t *) dst->op_params; - const int k0 = opts[1]; - const int k1 = opts[2]; - const int s0 = opts[3]; - const int s1 = opts[4]; - const int p0 = opts[5]; - const int p1 = opts[6]; - - std::vector kernel_dims = { k1, k0 }; - std::vector stride_dims = { s1, s0 }; - std::vector padding_avg_dims = { p1, p0 }; // (padH, padW) - - acl_int_array_ptr kernel_size = ggml_cann_create_int_array(kernel_dims.data(), 2); - acl_int_array_ptr strides = ggml_cann_create_int_array(stride_dims.data(), 2); - acl_int_array_ptr paddings_avg = ggml_cann_create_int_array(padding_avg_dims.data(), 2); - - bool ceil_mode = false; - bool count_include_pad = true; - int64_t divisor_override = 0; - int8_t cube_math_type = 0; -#ifdef ASCEND_310P - cube_math_type = 1; -#endif - - GGML_CANN_CALL_ACLNN_OP(ctx, AvgPool2d, acl_src.get(), kernel_size.get(), strides.get(), paddings_avg.get(), - ceil_mode, count_include_pad, divisor_override, cube_math_type, acl_dst.get()); -} - -/** - * @brief Performs 2D max pooling on the input tensor and stores the result in - * the destination tensor. - * - * This function performs max pooling on the source tensor and stores the result - * in the destination tensor. The pooling parameters (kernel size, strides, - * padding) are specified in the `op_params` of the destination tensor. - * - * @param ctx The context for the CANN backend operations. - * @param dst The destination tensor where the result will be stored. The source - * tensor is referenced by `dst->src[0]`. - */ -static void ggml_cann_max_pool2d(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - GGML_ASSERT(src->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src, nullptr, nullptr, 0, ACL_FORMAT_NCHW); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst, nullptr, nullptr, 0, ACL_FORMAT_NCHW); - - const int32_t * opts = (const int32_t *) dst->op_params; - const int k0 = opts[1]; - const int k1 = opts[2]; - const int s0 = opts[3]; - const int s1 = opts[4]; - const int p0 = opts[5]; - const int p1 = opts[6]; - - int64_t temp_ne[] = { src->ne[0] + p0 * 2, src->ne[1] + p1 * 2, src->ne[2], src->ne[3] }; - size_t temp_nb[GGML_MAX_DIMS]; - - temp_nb[0] = ggml_element_size(src); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - temp_nb[i] = temp_nb[i - 1] * temp_ne[i - 1]; - } - - ggml_cann_pool_alloc temp_buffer_allocator(ctx.pool(), ggml_nbytes(src) + p0 * 2 + p1 * 2 * src->nb[1]); - void * buffer = temp_buffer_allocator.get(); - acl_tensor_ptr tmp_tensor = ggml_cann_create_tensor(buffer, ACL_FLOAT, ggml_element_size(src), temp_ne, temp_nb, - GGML_MAX_DIMS, ACL_FORMAT_NCHW); - - // pad: see padding in ggml_cann_pad() - int64_t paddings[] = { p0, p0, p1, p1, 0, 0, 0, 0 }; - float value = -FLT_MAX; - aclnn_pad(ctx, acl_src.get(), tmp_tensor.get(), paddings, value); - - // max_pool - std::vector kernel_dims = { k1, k0 }; - std::vector stride_dims = { s1, s0 }; - // padding_max_dims: [dim0_start, dim0_end, dim1_start, dim1_end] - std::vector padding_max_dims = { 0, 0, 0, 0 }; - std::vector dilation_size = { 1, 1 }; - acl_int_array_ptr kernel_size = ggml_cann_create_int_array(kernel_dims.data(), 2); - acl_int_array_ptr strides = ggml_cann_create_int_array(stride_dims.data(), 2); - acl_int_array_ptr paddings_max = ggml_cann_create_int_array(padding_max_dims.data(), 4); - acl_int_array_ptr dilations = ggml_cann_create_int_array(dilation_size.data(), 2); - - bool ceil_mode = false; - int64_t auto_pads = 0; - GGML_CANN_CALL_ACLNN_OP(ctx, MaxPool, tmp_tensor.get(), kernel_size.get(), strides.get(), auto_pads, - paddings_max.get(), dilations.get(), ceil_mode, acl_dst.get()); -} - -void ggml_cann_pool2d(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - const int32_t * opts = (const int32_t *) dst->op_params; - enum ggml_op_pool op = static_cast(opts[0]); - switch (op) { - case GGML_OP_POOL_AVG: - ggml_cann_avg_pool2d(ctx, dst); - break; - case GGML_OP_POOL_MAX: - ggml_cann_max_pool2d(ctx, dst); - break; - case GGML_OP_POOL_COUNT: - GGML_ABORT("fatal error"); - break; - } -} - -/** - * @brief Copies data from the source tensor to the destination tensor. - * - * This function copies data from the source tensor `acl_src` to the destination - * tensor `acl_dst`. - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The source tensor from which data will be copied. - * @param acl_dst The destination tensor where the data will be copied to. - */ -static void cann_copy(ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst) { - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceCopy, acl_dst, acl_src); -} - -void ggml_cann_dup(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - - if (ggml_are_same_shape(src0, dst)) { - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - if (dst->type == src0->type) { - cann_copy(ctx, acl_src.get(), acl_dst.get()); - } else { - aclnn_cast(ctx, acl_src.get(), acl_dst.get(), ggml_cann_type_mapping(dst->type)); - } - } else { - void * src_trans_buffer = src0->data; - ggml_cann_pool_alloc src_buffer_allocator; - if (!ggml_is_contiguous(src0)) { - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src0); - src_buffer_allocator.alloc(ctx.pool(), ggml_nelements(src0) * ggml_type_size(src0->type)); - src_trans_buffer = src_buffer_allocator.get(); - size_t src_trans_nb[GGML_MAX_DIMS]; - src_trans_nb[0] = ggml_type_size(src0->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - src_trans_nb[i] = src_trans_nb[i - 1] * src0->ne[i - 1]; - } - acl_tensor_ptr src_trans_tensor = - ggml_cann_create_tensor(src_trans_buffer, ggml_cann_type_mapping(src0->type), - ggml_type_size(src0->type), src0->ne, src_trans_nb, GGML_MAX_DIMS); - cann_copy(ctx, acl_src.get(), src_trans_tensor.get()); - } - - size_t src_reshape_nb[GGML_MAX_DIMS]; - src_reshape_nb[0] = ggml_type_size(src0->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - src_reshape_nb[i] = src_reshape_nb[i - 1] * dst->ne[i - 1]; - } - - acl_tensor_ptr trans_acl_src = - ggml_cann_create_tensor(src_trans_buffer, ggml_cann_type_mapping(src0->type), ggml_type_size(src0->type), - dst->ne, src_reshape_nb, GGML_MAX_DIMS, ACL_FORMAT_ND); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - if (dst->type == src0->type) { - cann_copy(ctx, trans_acl_src.get(), acl_dst.get()); - } else { - aclnn_cast(ctx, trans_acl_src.get(), acl_dst.get(), ggml_cann_type_mapping(dst->type)); - } - } -} - -/** - * @brief Creates an ACL tensor initialized with zeros using a provided buffer. - * - * This function initializes a tensor with zeros using the specified buffer and - * tensor parameters. - * - * @param ctx The context for the CANN backend operations. - * @param buffer The buffer to be used for the tensor data. - * @param n_bytes The size of the buffer in bytes. - * @param ne An array specifying the extents (sizes) of each dimension of the - * tensor. - * @param dims The number of dimensions of the tensor. - * @param type The data type of the tensor. - * @param type_size The size of each element in the tensor data type. - * @return A tensor smart pointer initialized with zeros. - */ -static acl_tensor_ptr aclnn_zero(ggml_backend_cann_context & ctx, - void * buffer, - size_t n_bytes, - int64_t * ne, - int64_t dims, - aclDataType type, - size_t type_size) { - size_t nb[GGML_MAX_DIMS]; - nb[0] = type_size; - for (int i = 1; i < dims; i++) { - nb[i] = nb[i - 1] * ne[i - 1]; - } - - acl_tensor_ptr zero = ggml_cann_create_tensor(buffer, type, type_size, ne, nb, dims); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceZero, zero.get()); - return zero; - GGML_UNUSED(n_bytes); -} - -/** - * @brief Creates an ACL tensor initialized with value using a provided buffer. - * - * This function initializes a tensor with value using the specified buffer and - * tensor parameters. - * - * @param ctx The context for the CANN backend operations. - * @param buffer The buffer to be used for the tensor data. - * @param n_bytes The size of the buffer in bytes. - * @param ne An array specifying the extents (sizes) of each dimension of the - * tensor. - * @param dims The number of dimensions of the tensor. - * @param type The data type of the tensor. - * @param type_size The size of each element in the tensor data type. - * @param value The value to be used for initializing the tensor (default - * is 1.0). - * @return A tensor smart pointer initialized with value. - */ -static acl_tensor_ptr aclnn_values(ggml_backend_cann_context & ctx, - void * buffer, - size_t n_bytes, - int64_t * ne, - int64_t dims, - aclDataType type, - size_t type_size, - float value = 1.0f) { - acl_tensor_ptr acl_tensor = aclnn_zero(ctx, buffer, n_bytes, ne, dims, type, type_size); - float alpha_host = 1.0f; - acl_scalar_ptr alpha = ggml_cann_create_scalar(&alpha_host, aclDataType::ACL_FLOAT); - acl_scalar_ptr other = ggml_cann_create_scalar(&value, aclDataType::ACL_FLOAT); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceAdds, acl_tensor.get(), other.get(), alpha.get()); - return acl_tensor; -} - -/** - * @brief Fills a tensor with a scalar value. - * - * This function fills the destination tensor `acl_dst` with the scalar value - * `scalar`. - * - * @param ctx The context for the CANN backend operations. - * @param scalar The scalar value used to fill the tensor. - * @param acl_dst The destination tensor to be filled with the scalar value. - */ -static void aclnn_fill_scalar(ggml_backend_cann_context & ctx, float scalar, aclTensor * acl_dst) { - acl_scalar_ptr acl_scalar = ggml_cann_create_scalar(&scalar, aclDataType::ACL_FLOAT); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceFillScalar, acl_dst, acl_scalar.get()); -} - -/** - * @brief Get or expand a cached tensor filled with a scalar value. - * - * This function manages cached device memory for tensors. If the current - * cache size is insufficient for the requested tensor shape, the old memory will - * be released and new memory will be allocated. The allocated buffer is - * initialized with the given scalar value using CANN operations. - * Finally, an aclTensor object is created from the cached memory and returned. - * - * @param ctx The CANN backend context that manages device memory. - * @param buffer A pointer to the cached device buffer (will be allocated - * or reallocated if necessary). - * @param cache_element The current number of cached elements. This will be - * updated when the cache is expanded. - * @param ne The tensor shape array (number of elements in each dimension). - * @param nb The stride size for each dimension. - * @param dtype Data type of cached tensor. - * @param dims The number of tensor dimensions. - * @param value The scalar value used to fill the tensor (supports zero - * initialization via memset or arbitrary values via fill_scalar). - * @return A tensor smart pointer created from the cached buffer. - */ -static acl_tensor_ptr get_cache_acl_tensor(ggml_backend_cann_context & ctx, - void ** buffer, - int64_t & cache_element, - int64_t * ne, - size_t * nb, - ggml_type dtype, - int64_t dims, - float value) { - // Calculate total number of elements - int64_t n_element = 1; - for (int i = 0; i < dims; i++) { - n_element *= ne[i]; - } - size_t size = n_element * ggml_type_size(dtype); - - // Allocate or expand cache if needed - if (cache_element < n_element) { - if (*buffer != nullptr) { - aclrtFree(*buffer); - *buffer = nullptr; - } - - ACL_CHECK(aclrtMalloc(buffer, size, ACL_MEM_MALLOC_HUGE_FIRST)); - cache_element = n_element; - - // Initialize cache - int64_t pool_ne[1] = { n_element }; - size_t pool_nb[1] = { ggml_type_size(dtype) }; - acl_tensor_ptr acl_value = - ggml_cann_create_tensor(*buffer, ggml_cann_type_mapping(dtype), ggml_type_size(dtype), pool_ne, pool_nb, 1); - aclnn_fill_scalar(ctx, value, acl_value.get()); - } - - return ggml_cann_create_tensor(*buffer, ggml_cann_type_mapping(dtype), ggml_type_size(dtype), ne, nb, dims); -} - -void ggml_cann_rms_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - float eps; - memcpy(&eps, dst->op_params, sizeof(float)); - - // build gamma. - size_t acl_gamma_nb[GGML_MAX_DIMS]; - // gamma's type is the same with dst. - acl_gamma_nb[0] = ggml_type_size(dst->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - acl_gamma_nb[i] = acl_gamma_nb[i - 1] * src->ne[i - 1]; - } - acl_tensor_ptr acl_gamma = get_cache_acl_tensor( - ctx, &ctx.rms_norm_one_tensor_cache.cache, ctx.rms_norm_one_tensor_cache.size, src->ne, acl_gamma_nb, dst->type, - 1, // dims - 1.0f // value - ); - - // build rstd. - int64_t acl_rstd_ne[] = { src->ne[1], src->ne[2], src->ne[3] }; - size_t acl_rstd_nb[GGML_MAX_DIMS - 1]; - // rstd will always be F32. - acl_rstd_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS - 1; i++) { - acl_rstd_nb[i] = acl_rstd_nb[i - 1] * acl_rstd_ne[i - 1]; - } - acl_tensor_ptr acl_rstd = - get_cache_acl_tensor(ctx, &ctx.rms_norm_zero_tensor_cache.cache, ctx.rms_norm_zero_tensor_cache.size, - acl_rstd_ne, acl_rstd_nb, GGML_TYPE_F32, GGML_MAX_DIMS - 1, - 0.0f // value - ); - - GGML_CANN_CALL_ACLNN_OP(ctx, RmsNorm, acl_src.get(), acl_gamma.get(), eps, acl_dst.get(), acl_rstd.get()); -} - -// TODO: performace is low. -void ggml_cann_diag_mask(ggml_backend_cann_context & ctx, ggml_tensor * dst, float value) { - ggml_tensor * src = dst->src[0]; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - const int n_past = ((int32_t *) dst->op_params)[0]; - - ggml_cann_pool_alloc one_tensor_allocator(ctx.pool(), ggml_nbytes(src)); - void * buffer = one_tensor_allocator.get(); - - acl_tensor_ptr mask_tensor = ggml_cann_create_tensor(buffer, ggml_cann_type_mapping(src->type), - ggml_type_size(src->type), src->ne, src->nb, GGML_MAX_DIMS); - - aclnn_fill_scalar(ctx, value, mask_tensor.get()); - - float alphaValue = 1.0f; - acl_scalar_ptr alpha = ggml_cann_create_scalar(&alphaValue, aclDataType::ACL_FLOAT); - - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceTriu, mask_tensor.get(), n_past + 1); - GGML_CANN_CALL_ACLNN_OP(ctx, Tril, acl_src.get(), n_past + 1, acl_dst.get()); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceAdd, acl_dst.get(), mask_tensor.get(), alpha.get()); -} - -/** - * @brief Permutes the dimensions of a tensor according to a specified order. - * - * This function permutes the dimensions of the source tensor `acl_src` - * according to the order specified in the `new_dim` array and stores the result - * in the destination tensor `acl_dst`. - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The source tensor whose dimensions will be permuted. - * @param acl_dst The destination tensor where the permuted result will be - * stored. - * @param new_dim An array specifying the new order of dimensions for the - * tensor. - * @param dims The number of dimensions in the tensor. - */ -static void aclnn_permute(ggml_backend_cann_context & ctx, - aclTensor * acl_src, - aclTensor * acl_dst, - int64_t * new_dim, - uint64_t dims) { - acl_int_array_ptr acl_dims = ggml_cann_create_int_array(new_dim, dims); - GGML_CANN_CALL_ACLNN_OP(ctx, Permute, acl_src, acl_dims.get(), acl_dst); -} - -static void ggml_cann_im2col_2d_post_process(ggml_backend_cann_context & ctx, - ggml_tensor * dst, - ggml_tensor * src1, - aclTensor * tmp_cast_tensor, - aclTensor * tmp_im2col_tensor) { - // Permute: [N, IC * KH * KW, OW * OH] -> [N, OW * OH, IC * KH * KW] - int64_t dst_ne[] = { dst->ne[0], dst->ne[1] * dst->ne[2], dst->ne[3] }; - size_t dst_nb[] = { dst->nb[0], dst->nb[1], dst->nb[3] }; - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst, dst_ne, dst_nb, GGML_MAX_DIMS - 1); - - int64_t permute_dim[] = { 0, 2, 1 }; - if (src1->type != dst->type) { - aclnn_permute(ctx, tmp_cast_tensor, acl_dst.get(), permute_dim, 3); - } else { - aclnn_permute(ctx, tmp_im2col_tensor, acl_dst.get(), permute_dim, 3); - } -} - -static void ggml_cann_im2col_1d_post_process(ggml_backend_cann_context & ctx, - ggml_tensor * dst, - ggml_tensor * src1, - aclTensor * tmp_cast_tensor, - aclTensor * tmp_im2col_tensor, - const std::vector & im2col_op_params) { - // get params - const int64_t KH = im2col_op_params[0]; - const int64_t KW = im2col_op_params[1]; - const int64_t IW = im2col_op_params[2]; - const int64_t IC = im2col_op_params[3]; - const int64_t N = im2col_op_params[4]; - const int64_t OH = im2col_op_params[5]; - const int64_t OW = im2col_op_params[6]; - const int64_t s0 = im2col_op_params[7]; - const int64_t p0 = im2col_op_params[8]; - const int64_t d0 = im2col_op_params[9]; - const int64_t n_bytes_factor = im2col_op_params[10]; - - // Permute: [N, IC * KH * KW, OW * OH] -> - // [N, OW * OH * n_bytes_factor, IC * KH * KW] - ggml_cann_pool_alloc tmp_permute_allocator(ctx.pool()); - tmp_permute_allocator.alloc(ggml_nbytes(dst) * n_bytes_factor); - void * tmp_permute_buffer = tmp_permute_allocator.get(); - - int64_t tmp_permute_ne[] = { IC * KH * KW, OW * OH * n_bytes_factor, N }; - size_t tmp_permute_nb[GGML_MAX_DIMS - 1]; - tmp_permute_nb[0] = ggml_type_size(dst->type); - for (int i = 1; i < GGML_MAX_DIMS - 1; i++) { - tmp_permute_nb[i] = tmp_permute_nb[i - 1] * tmp_permute_ne[i - 1]; - } - - acl_tensor_ptr tmp_permute_tensor = - ggml_cann_create_tensor(tmp_permute_buffer, ggml_cann_type_mapping(dst->type), ggml_type_size(dst->type), - tmp_permute_ne, tmp_permute_nb, GGML_MAX_DIMS - 1, ACL_FORMAT_ND); - - int64_t permute_dim[] = { 0, 2, 1 }; - if (src1->type != dst->type) { - aclnn_permute(ctx, tmp_cast_tensor, tmp_permute_tensor.get(), permute_dim, 3); - } else { - aclnn_permute(ctx, tmp_im2col_tensor, tmp_permute_tensor.get(), permute_dim, 3); - } - - // number of times the kernel moves in W dimension - const int n_step_w = (IW + 2 * p0 - d0 * (KW - 1) - 1) / s0 + 1; - size_t offset; - void * cur_dst_buffer = dst->data, *cur_permute_buffer = tmp_permute_buffer; - - // memory copy with offset to restore 1D im2col from 2d - if (IC > 1) { - offset = IC * KH * KW * n_step_w * ggml_type_size(dst->type); - size_t cpy_size = KH * KW * ggml_type_size(dst->type); - - for (int c = 0; c < IC; c++) { - cur_permute_buffer = (char *) tmp_permute_buffer + offset + KH * KW * c * ggml_type_size(dst->type); - cur_dst_buffer = (char *) dst->data + c * KH * KW * n_step_w * ggml_type_size(dst->type); - - for (int i = 0; i < n_step_w; i++) { - ACL_CHECK(aclrtMemcpyAsync(cur_dst_buffer, cpy_size, cur_permute_buffer, cpy_size, - ACL_MEMCPY_DEVICE_TO_DEVICE, ctx.stream())); - cur_dst_buffer = (char *) cur_dst_buffer + KH * KW * ggml_type_size(dst->type); - cur_permute_buffer = (char *) cur_permute_buffer + KH * KW * IC * ggml_type_size(dst->type); - } - } - } else { - offset = KH * KW * n_step_w * ggml_type_size(dst->type); // equal to ggml_nbytes(dst) - ACL_CHECK(aclrtMemcpyAsync(dst->data, offset, (char *) tmp_permute_buffer + offset, offset, - ACL_MEMCPY_DEVICE_TO_DEVICE, ctx.stream())); - } -} - -void ggml_cann_im2col(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; // kernel - ggml_tensor * src1 = dst->src[1]; // input - - GGML_TENSOR_BINARY_OP_LOCALS; - - // aclnnIm2col only works on 2D. set s1, p1, d1 to 1 to perform 2D - // im2col and do post-processing to restore it to 1D. - const bool is_2D = ((const int32_t *) (dst->op_params))[6] == 1; - const int32_t s0 = ((const int32_t *) (dst->op_params))[0]; - const int32_t s1 = is_2D ? ((const int32_t *) (dst->op_params))[1] : 1; - const int32_t p0 = ((const int32_t *) (dst->op_params))[2]; - const int32_t p1 = is_2D ? ((const int32_t *) (dst->op_params))[3] : 1; - const int32_t d0 = ((const int32_t *) (dst->op_params))[4]; - const int32_t d1 = is_2D ? ((const int32_t *) (dst->op_params))[5] : 1; - - const int64_t N = ne13; - const int64_t IC = ne12; - const int64_t KH = ne01; - const int64_t KW = ne00; - const int64_t IW = ne10; - - const int64_t OH = is_2D ? ne2 : 1; - const int64_t OW = ne1; - - // memory allocated increased to 3x when is_2D == false - const int64_t n_bytes_factor = is_2D ? 1 : 3; - - // im2col: [N,C,H,W] -> [N, IC * KH * KW, OW * OH * n_bytes_factor] - acl_tensor_ptr acl_src1 = ggml_cann_create_tensor(src1); - int64_t tmp_im2col_ne[] = { OW * OH * n_bytes_factor, IC * KH * KW, N }; - size_t tmp_im2col_nb[GGML_MAX_DIMS - 1]; - - tmp_im2col_nb[0] = ggml_type_size(src1->type); - for (int i = 1; i < GGML_MAX_DIMS - 1; i++) { - tmp_im2col_nb[i] = tmp_im2col_nb[i - 1] * tmp_im2col_ne[i - 1]; - } - - // Calculate im2col. - // If dst is f16, tmp_buffer is f32, we need alloc src.typesize * - // dst.elemcount. - ggml_cann_pool_alloc im2col_allocator(ctx.pool(), ggml_nelements(dst) * ggml_element_size(src1) * n_bytes_factor); - void * tmp_im2col_buffer = im2col_allocator.get(); - - acl_tensor_ptr tmp_im2col_tensor = - ggml_cann_create_tensor(tmp_im2col_buffer, ggml_cann_type_mapping(src1->type), ggml_type_size(src1->type), - tmp_im2col_ne, tmp_im2col_nb, GGML_MAX_DIMS - 1, ACL_FORMAT_ND); - - std::vector kernel_dims = { KH, KW }; - std::vector dilation_size = { d1, d0 }; - std::vector padding_dims = { p1, p0 }; - std::vector stride_dims = { s1, s0 }; - acl_int_array_ptr kernel_size = ggml_cann_create_int_array(kernel_dims.data(), 2); - acl_int_array_ptr dilations = ggml_cann_create_int_array(dilation_size.data(), 2); - acl_int_array_ptr paddings = ggml_cann_create_int_array(padding_dims.data(), 2); - acl_int_array_ptr strides = ggml_cann_create_int_array(stride_dims.data(), 2); - GGML_CANN_CALL_ACLNN_OP(ctx, Im2col, acl_src1.get(), kernel_size.get(), dilations.get(), paddings.get(), - strides.get(), tmp_im2col_tensor.get()); - - // Cast if dst is f16. - acl_tensor_ptr tmp_cast_tensor; - ggml_cann_pool_alloc tmp_cast_allocator(ctx.pool()); - void * tmp_cast_buffer = nullptr; - if (src1->type != dst->type) { - tmp_cast_allocator.alloc(ggml_nbytes(dst) * n_bytes_factor); - tmp_cast_buffer = tmp_cast_allocator.get(); - size_t temp_cast_nb[GGML_MAX_DIMS - 1]; - temp_cast_nb[0] = ggml_type_size(dst->type); - for (int i = 1; i < GGML_MAX_DIMS - 1; i++) { - temp_cast_nb[i] = temp_cast_nb[i - 1] * tmp_im2col_ne[i - 1]; - } - - tmp_cast_tensor = - ggml_cann_create_tensor(tmp_cast_buffer, ggml_cann_type_mapping(dst->type), ggml_type_size(dst->type), - tmp_im2col_ne, temp_cast_nb, GGML_MAX_DIMS - 1, ACL_FORMAT_ND); - aclnn_cast(ctx, tmp_im2col_tensor.get(), tmp_cast_tensor.get(), ggml_cann_type_mapping(dst->type)); - } - - // post-processing - if (is_2D) { - ggml_cann_im2col_2d_post_process(ctx, dst, src1, tmp_cast_tensor.get(), tmp_im2col_tensor.get()); - } else { - std::vector im2col_op_params = { KH, KW, IW, IC, N, OH, OW, s0, p0, d0, n_bytes_factor }; - ggml_cann_im2col_1d_post_process(ctx, dst, src1, tmp_cast_tensor.get(), tmp_im2col_tensor.get(), - im2col_op_params); - } -} - -/** - * @brief Applies element-wise exponential function to the elements of a tensor. - * - * This function computes the exponential of each element in the source tensor - * `acl_src` and stores the result back into the same tensor. - * The operation is defined as: - * \f[ - * \text {acl_src }_i=e^{acl\_src_i} - * \f] - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The tensor on which the exponential function will be applied. - */ -static void aclnn_exp(ggml_backend_cann_context & ctx, aclTensor * acl_src) { - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceExp, acl_src); -} - -void aclnn_cos(ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst) { - if (acl_dst == nullptr) { - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceCos, acl_src); - } else { - GGML_CANN_CALL_ACLNN_OP(ctx, Cos, acl_src, acl_dst); - } -} - -void aclnn_sin(ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst) { - if (acl_dst == nullptr) { - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceSin, acl_src); - } else { - GGML_CANN_CALL_ACLNN_OP(ctx, Sin, acl_src, acl_dst); - } -} - -void ggml_cann_timestep_embedding(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - const ggml_tensor * src = dst->src[0]; - - GGML_ASSERT(src->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - - const int dim = dst->op_params[0]; - const int max_period = dst->op_params[1]; - int half = dim / 2; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - - // arange: [0, ..., half) - float start = 0; - float stop = half; - float step = 1; - int64_t n_elements_arange = half; - int64_t tmp_arange_ne[] = { half }; - size_t tmp_arange_nb[] = { sizeof(dst->type) }; - - ggml_cann_pool_alloc arange_allocator(ctx.pool(), half * sizeof(dst->type)); - void * tmp_arange_buffer = arange_allocator.get(); - acl_tensor_ptr tmp_arange_tensor = - ggml_cann_create_tensor(tmp_arange_buffer, ggml_cann_type_mapping(dst->type), ggml_type_size(dst->type), - tmp_arange_ne, tmp_arange_nb, GGML_MAX_DIMS - 3, ACL_FORMAT_ND); - - aclnn_arange(ctx, tmp_arange_tensor.get(), start, stop, step, n_elements_arange); - - // freq - float freq_param = -logf(max_period) / half; - bool inplace = true; - aclnn_muls(ctx, tmp_arange_tensor.get(), freq_param, nullptr, inplace); - aclnn_exp(ctx, tmp_arange_tensor.get()); - - // permute: src [0,1,2,3]->[0,1,3,2] - int64_t tmp_permute_ne[] = { src->ne[1], src->ne[0], src->ne[2], src->ne[3] }; - size_t tmp_permute_nb[GGML_MAX_DIMS]; - tmp_permute_nb[0] = ggml_type_size(src->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - tmp_permute_nb[i] = tmp_permute_nb[i - 1] * tmp_permute_ne[i - 1]; - } - - ggml_cann_pool_alloc permute_allocator(ctx.pool(), ggml_nbytes(src)); - void * tmp_permute_buffer = permute_allocator.get(); - acl_tensor_ptr tmp_permute_tensor = - ggml_cann_create_tensor(tmp_permute_buffer, ggml_cann_type_mapping(src->type), ggml_type_size(src->type), - tmp_permute_ne, tmp_permute_nb, GGML_MAX_DIMS, ACL_FORMAT_ND); - int64_t permute_dim[] = { 0, 1, 3, 2 }; - int64_t num_dims = 4; - aclnn_permute(ctx, acl_src.get(), tmp_permute_tensor.get(), permute_dim, num_dims); - - // timestep * freq - int64_t tmp_mul_ne[] = { src->ne[1] * half, src->ne[0], src->ne[2], src->ne[3] }; - size_t tmp_mul_nb[GGML_MAX_DIMS]; - tmp_mul_nb[0] = ggml_type_size(src->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - tmp_mul_nb[i] = tmp_mul_nb[i - 1] * tmp_mul_ne[i - 1]; - } - - int mul_nelements = src->ne[1] * half * src->ne[0] * src->ne[2] * src->ne[3]; - - ggml_cann_pool_alloc mul_allocator(ctx.pool(), mul_nelements * ggml_type_size(src->type)); - void * tmp_mul_buffer = mul_allocator.get(); - acl_tensor_ptr tmp_mul_tensor = - ggml_cann_create_tensor(tmp_mul_buffer, ggml_cann_type_mapping(src->type), ggml_type_size(src->type), - tmp_mul_ne, tmp_mul_nb, GGML_MAX_DIMS, ACL_FORMAT_ND); - aclnn_mul(ctx, tmp_permute_tensor.get(), tmp_arange_tensor.get(), tmp_mul_tensor.get()); - - // cos - ggml_cann_pool_alloc cos_allocator(ctx.pool(), mul_nelements * ggml_type_size(src->type)); - void * tmp_cos_buffer = cos_allocator.get(); - acl_tensor_ptr tmp_cos_tensor = - ggml_cann_create_tensor(tmp_cos_buffer, ggml_cann_type_mapping(dst->type), ggml_type_size(dst->type), - tmp_mul_ne, tmp_mul_nb, GGML_MAX_DIMS, ACL_FORMAT_ND); - - aclnn_cos(ctx, tmp_mul_tensor.get(), tmp_cos_tensor.get()); - - // sin - ggml_cann_pool_alloc sin_allocator(ctx.pool(), mul_nelements * ggml_type_size(src->type)); - void * tmp_sin_buffer = sin_allocator.get(); - acl_tensor_ptr tmp_sin_tensor = - ggml_cann_create_tensor(tmp_sin_buffer, ggml_cann_type_mapping(dst->type), ggml_type_size(dst->type), - tmp_mul_ne, tmp_mul_nb, GGML_MAX_DIMS, ACL_FORMAT_ND); - - aclnn_sin(ctx, tmp_mul_tensor.get(), tmp_sin_tensor.get()); - - // concat - int64_t concat_dim = 3; - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - acl_tensor_list_ptr tensor_list = ggml_cann_create_tensor_list(tmp_cos_tensor, tmp_sin_tensor); - aclnn_concat(ctx, tensor_list.get(), acl_dst.get(), concat_dim); -} - -/** - * @brief Raises each element of a tensor to the power of the corresponding - * element in another tensor. - * - * This function computes the element-wise power of the destination tensor - * `acl_dst` raised to the power of the exponent tensor `acl_exp`. - * The operation is defined as: - * \f[ - * \text {acl_dst }_i=acl\_dst_i^{\text {acl_exp }_i} - * \f] - * - * @param ctx The context for the CANN backend operations. - * @param acl_dst The destination tensor, which also serves as the base tensor. - * @param acl_exp The exponent tensor, each element of which is used to raise - * the corresponding element in the destination tensor. - */ -static void aclnn_pow_tensor_tensor(ggml_backend_cann_context & ctx, aclTensor * acl_dst, aclTensor * acl_exp) { - GGML_CANN_CALL_ACLNN_OP(ctx, InplacePowTensorTensor, acl_dst, acl_exp); -} - -/** - * @brief Generate a range of values and apply a scalar base exponentiation. - * - * This function creates an evenly spaced sequence from `start` to `stop` (exclusive), - * with step size `step`, stores it in a temporary buffer, and then computes: - * - * @f[ - * slope[i] = m^{\left( start + i \cdot step \right)}, \quad 0 \le i < size - * @f] - * - * The results are written to the provided @p slope_buffer. - * - * @param ctx CANN backend context for memory allocation and operator execution. - * @param slope_buffer Pointer to the output buffer (float array) for the computed slope values. - * @param m Scalar base for the exponentiation. - * @param size Number of elements in the generated sequence. - * @param start Starting exponent offset. - * @param stop Stopping exponent offset (exclusive). - * @param step Step size for the exponent increment. - * @param dtype Data type for slope tensor. - */ -static void aclnn_get_slope_inner(ggml_backend_cann_context & ctx, - void * slope_buffer, - float m, - int64_t size, - float start, - float stop, - float step, - ggml_type dtype) { - aclDataType acl_type = ggml_cann_type_mapping(dtype); - size_t type_size = ggml_type_size(dtype); - - int64_t ne[] = { size }; - size_t nb[] = { type_size }; - - ggml_cann_pool_alloc arange_allocator(ctx.pool(), size * type_size); - void * arange_buffer = arange_allocator.get(); - - acl_tensor_ptr arange_tensor = ggml_cann_create_tensor(arange_buffer, acl_type, type_size, ne, nb, 1); - aclnn_arange(ctx, arange_tensor.get(), start, stop, step, size); - - acl_tensor_ptr slope_tensor = ggml_cann_create_tensor(slope_buffer, acl_type, type_size, ne, nb, 1); - - acl_scalar_ptr sc = ggml_cann_create_scalar(&m, aclDataType::ACL_FLOAT); - - GGML_CANN_CALL_ACLNN_OP(ctx, PowScalarTensor, sc.get(), arange_tensor.get(), slope_tensor.get()); -} - -/** - * @brief Compute slope values for multiple attention heads based on ALiBi bias parameters. - * - * This function generates slope values for each attention head according to the ALiBi - * (Attention with Linear Biases) method. It splits the computation into two ranges depending - * on whether the head index is less than @p n_head_log2 or not, and uses different base values - * (`m0` and `m1`) for the exponentiation. - * - * @f[ - * slope[h] = - * \begin{cases} - * m_0^{(h + 1)}, & h < n\_head\_log2 \\ - * m_1^{\left( 2 \cdot (h - n\_head\_log2) + 1 \right)}, & h \geq n\_head\_log2 - * \end{cases} - * \quad , \quad \text{if } max\_bias > 0 - * @f] - * - * If @p max_bias <= 0, all slope values are set to 1.0. - * - * @param ctx CANN backend context for memory allocation and operator execution. - * @param n_head Total number of attention heads. - * @param slope_buffer Pointer to the output buffer (float array) for storing slopes. - * @param max_bias Maximum bias value for slope computation. - * @param dtype Data type for slope tensor. - * -*/ -static void aclnn_get_slope(ggml_backend_cann_context & ctx, - int64_t n_head, - void * slope_buffer, - float max_bias, - ggml_type dtype) { - const int n_head_log2 = 1u << (uint32_t) floor(log2(n_head)); - - float m0 = powf(2.0f, -(max_bias) / n_head_log2); - float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); - - // const float slope = (max_bias > 0.0f) ? - // h < n_head_log2 ? - // powf(m0, h + 1) : - // powf(m1, 2*(h - n_head_log2) + 1) : - // 1.0f; - // arange1 - float start = 0 + 1; - float end = (n_head_log2 - 1) + 1; - float step = 1; - float count = n_head_log2; - // end needs to be +1 because aclnn uses a left-closed, right-open interval. - aclnn_get_slope_inner(ctx, slope_buffer, m0, count, start, end + 1, step, dtype); - if (n_head_log2 < n_head) { - // arange2 - start = 2 * (n_head_log2 - n_head_log2) + 1; - end = 2 * ((n_head - 1) - n_head_log2) + 1; - step = 2; - count = n_head - n_head_log2; - aclnn_get_slope_inner(ctx, (char *) slope_buffer + n_head_log2 * ggml_type_size(dtype), m1, count, start, end + 1, - step, dtype); - } -} - -/** - * @brief Add ALiBi (Attention with Linear Biases) positional biases to the attention mask. - * - * This function computes the ALiBi slopes for each attention head (if max_bias > 0), - * multiplies them with the attention mask to produce bias tensors, and adds these biases - * to the destination tensor (@p dst). - * - * The function performs necessary broadcasting of the mask and slope tensors to match - * the shape of the destination tensor, then applies element-wise multiplication and addition - * using CANN operators. - * - * @param ctx CANN backend context for memory management and operator execution. - * @param mask Input attention mask tensor, assumed to be contiguous. - * @param dst Destination tensor to which ALiBi biases will be added. - * @param dst_ptr Pointer to the memory of the destination tensor. - * @param max_bias Maximum bias value controlling the slope scaling. - * - * @note - * - Write data into dst_ptr using only the shape information of the dst tensor. - * - `GGML_MAX_DIMS + 2` is used to extend tensor dimensions for broadcasting. - */ -static void aclnn_add_alibi(ggml_backend_cann_context & ctx, - ggml_tensor * mask, - ggml_tensor * dst, - void * dst_ptr, - float max_bias) { - void * slope_buffer = nullptr; - void * bias_buffer = nullptr; - - if (max_bias > 0.0f) { - int64_t n_heads = dst->ne[2]; - ggml_cann_pool_alloc slope_allocator(ctx.pool(), n_heads * sizeof(float)); - slope_buffer = slope_allocator.get(); - ggml_cann_pool_alloc bias_allocator(ctx.pool(), ggml_nelements(dst) * ggml_element_size(dst)); - bias_buffer = bias_allocator.get(); - aclnn_get_slope(ctx, n_heads, slope_buffer, max_bias, GGML_TYPE_F32); - } - - // broadcast for mask, slop and dst; - int64_t nr2 = dst->ne[2] / mask->ne[2]; - int64_t nr3 = dst->ne[3] / mask->ne[3]; - - // broadcast the mask across rows - int64_t mask_ne[] = { mask->ne[0], dst->ne[1], mask->ne[2], 1, mask->ne[3], 1 }; - size_t mask_nb[] = { mask_nb[0] = mask->nb[0], mask_nb[1] = mask->nb[1], mask_nb[2] = mask->nb[2], - mask_nb[3] = mask->nb[2], mask_nb[4] = mask->nb[3], mask_nb[5] = mask->nb[3] }; - - int64_t dst_ne[] = { dst->ne[0], dst->ne[1], mask->ne[2], nr2, mask->ne[3], nr3 }; - size_t dst_nb[] = { dst_nb[0] = dst->nb[0], dst_nb[1] = dst->nb[1], dst_nb[2] = dst->nb[2], - dst_nb[3] = dst->nb[2], dst_nb[4] = dst->nb[3], dst_nb[5] = dst->nb[3] }; - - // slope is a 1 dim tensor, slope.ne2 == dst.ne2 - int64_t slope_ne[] = { 1, 1, mask->ne[2], nr2, 1, 1 }; - size_t slope_nb[GGML_MAX_DIMS + 2]; - slope_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS + 2; i++) { - slope_nb[i] = slope_nb[i - 1] * slope_ne[i - 1]; - } - - acl_tensor_ptr acl_slope = - ggml_cann_create_tensor(slope_buffer, ACL_FLOAT, sizeof(float), slope_ne, slope_nb, GGML_MAX_DIMS + 2); - acl_tensor_ptr acl_mask = ggml_cann_create_tensor(mask, mask_ne, mask_nb, GGML_MAX_DIMS + 2); - - // write data into dst_ptr using only the shape information of the dst tensor. - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst_ptr, ggml_cann_type_mapping(dst->type), - ggml_type_size(dst->type), dst_ne, dst_nb, GGML_MAX_DIMS + 2); - - if (max_bias > 0.0f) { - int64_t bias_ne[] = { mask->ne[0], dst->ne[1], mask->ne[2], nr2, mask->ne[3], 1 }; - size_t bias_nb[GGML_MAX_DIMS + 2]; - bias_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS + 2; i++) { - bias_nb[i] = bias_nb[i - 1] * bias_ne[i - 1]; - } - acl_tensor_ptr bias_tensor = - ggml_cann_create_tensor(bias_buffer, ACL_FLOAT, sizeof(float), bias_ne, bias_nb, GGML_MAX_DIMS + 2); - - aclnn_mul(ctx, acl_slope.get(), acl_mask.get(), bias_tensor.get()); - aclnn_add(ctx, acl_dst.get(), bias_tensor.get()); - } else { - aclnn_add(ctx, acl_dst.get(), acl_mask.get()); - } -} - -void ggml_cann_cpy(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_cann_dup(ctx, dst); -} - -/** - * @brief Applies the softmax function to a tensor along a specified dimension. - * - * This function computes the softmax of the source tensor `acl_src` along the - * specified dimension `dim` and stores the result in the destination tensor - * `acl_dst`. - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The source tensor on which the softmax function will be - * applied. - * @param dim The dimension along which the softmax function will be computed. - * @param acl_dst The destination tensor where the softmax results will be - * stored. - */ -static void aclnn_softmax(ggml_backend_cann_context & ctx, aclTensor * acl_src, int64_t dim, aclTensor * acl_dst) { - GGML_CANN_CALL_ACLNN_OP(ctx, Softmax, acl_src, dim, acl_dst); -} - -void ggml_cann_softmax(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - ggml_tensor * src1 = dst->src[1]; // mask - - acl_tensor_ptr acl_src0 = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - float scale = 1.0f; - float max_bias = 0.0f; - - memcpy(&scale, (float *) dst->op_params + 0, sizeof(float)); - memcpy(&max_bias, (float *) dst->op_params + 1, sizeof(float)); - - // input mul scale - acl_scalar_ptr acl_scale = ggml_cann_create_scalar(&scale, aclDataType::ACL_FLOAT); - ggml_cann_pool_alloc src_tensor_allocator(ctx.pool(), ggml_nbytes(src0)); - void * src_tensor_buffer = src_tensor_allocator.get(); - acl_tensor_ptr softmax_tensor = ggml_cann_create_tensor(src_tensor_buffer, ggml_cann_type_mapping(src0->type), - ggml_element_size(src0), src0->ne, src0->nb, GGML_MAX_DIMS); - - aclnn_muls(ctx, acl_src0.get(), scale, softmax_tensor.get(), false); - - // mask - if (src1) { - aclnn_add_alibi(ctx, src1, src0, src_tensor_buffer, max_bias); - } - // softmax - aclnn_softmax(ctx, softmax_tensor.get(), 3, acl_dst.get()); -} - - -void ggml_cann_get_rows(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; // weight - ggml_tensor * src1 = dst->src[1]; // index - - GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16 - || dst->type == GGML_TYPE_BF16); - - // n_idx: number of row indices per (i2, i3) batch slice. - // ggml guarantees: src0->ne[2] == src1->ne[1], src0->ne[3] == src1->ne[2], src1->ne[3] == 1. - const int64_t n_idx = src1->ne[0]; - - // Gather all (i2, i3) batch slices from src into dst. - // ggml_cann_create_tensor reverses dims, so ACL sees [ne1, ne0]. - // GatherV2 with dim=0 gathers along ACL dim-0 == ggml ne[1] (the vocabulary / row axis). - // nb: the 4 strides of the source buffer (nb[0..1] for the 2D slice shape, - // nb[2..3] for computing per-batch-slice base pointer offsets). - auto gather_batched = [&](void * src_base, aclDataType acl_type, size_t type_size, - const size_t * nb) { - int64_t src_ne[2] = { src0->ne[0], src0->ne[1] }; - size_t src_nb_2d[2] = { nb[0], nb[1] }; - int64_t dst_ne[2] = { src0->ne[0], n_idx }; - size_t dst_nb_2d[2] = { dst->nb[0], dst->nb[1] }; - int64_t idx_ne[1] = { n_idx }; - size_t idx_nb[1] = { (size_t)ggml_element_size(src1) }; - - for (int64_t i3 = 0; i3 < src0->ne[3]; i3++) { - for (int64_t i2 = 0; i2 < src0->ne[2]; i2++) { - acl_tensor_ptr acl_src = ggml_cann_create_tensor( - (char *)src_base + i3 * nb[3] + i2 * nb[2], - acl_type, type_size, src_ne, src_nb_2d, 2); - acl_tensor_ptr acl_idx = ggml_cann_create_tensor( - (char *)src1->data + i3 * src1->nb[2] + i2 * src1->nb[1], - ggml_cann_type_mapping(src1->type), (size_t)ggml_element_size(src1), - idx_ne, idx_nb, 1); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor( - (char *)dst->data + i3 * dst->nb[3] + i2 * dst->nb[2], - acl_type, type_size, dst_ne, dst_nb_2d, 2); - GGML_CANN_CALL_ACLNN_OP(ctx, GatherV2, acl_src.get(), 0, acl_idx.get(), acl_dst.get()); - } - } - }; - - switch (src0->type) { - case GGML_TYPE_BF16: - case GGML_TYPE_F16: - case GGML_TYPE_F32: - if (src0->type == dst->type) { - gather_batched(src0->data, - ggml_cann_type_mapping(src0->type), ggml_type_size(src0->type), - src0->nb); - } else { - // Cast src0 to dst type, then gather. - ggml_cann_pool_alloc src_cast_allocator(ctx.pool(), - ggml_nelements(src0) * ggml_element_size(dst)); - size_t src_cast_nb[GGML_MAX_DIMS]; - src_cast_nb[0] = ggml_type_size(dst->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - src_cast_nb[i] = src_cast_nb[i - 1] * src0->ne[i - 1]; - } - acl_tensor_ptr acl_src0 = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_src_cast = ggml_cann_create_tensor( - src_cast_allocator.get(), ggml_cann_type_mapping(dst->type), ggml_type_size(dst->type), - src0->ne, src_cast_nb, GGML_MAX_DIMS); - aclnn_cast(ctx, acl_src0.get(), acl_src_cast.get(), ggml_cann_type_mapping(dst->type)); - - gather_batched(src_cast_allocator.get(), - ggml_cann_type_mapping(dst->type), ggml_type_size(dst->type), - src_cast_nb); - } - break; - case GGML_TYPE_Q8_0: - { - // Dequantize Q8_0 to dst type, then gather. - size_t weight_nb[GGML_MAX_DIMS + 1], scale_nb[GGML_MAX_DIMS + 1], dequant_nb[GGML_MAX_DIMS + 1]; - int64_t weight_ne[GGML_MAX_DIMS + 1], scale_ne[GGML_MAX_DIMS + 1], *dequant_ne; - weight_ne[0] = QK8_0; - weight_ne[1] = src0->ne[0] / QK8_0; - weight_nb[0] = sizeof(int8_t); - weight_nb[1] = weight_nb[0] * weight_ne[0]; - for (int i = 2; i < GGML_MAX_DIMS + 1; i++) { - weight_ne[i] = src0->ne[i - 1]; - weight_nb[i] = weight_nb[i - 1] * weight_ne[i - 1]; - } - scale_ne[0] = 1; - scale_ne[1] = src0->ne[0] / QK8_0; - scale_nb[0] = sizeof(uint16_t); - scale_nb[1] = scale_nb[0] * scale_ne[0]; - for (int i = 2; i < GGML_MAX_DIMS + 1; i++) { - scale_ne[i] = src0->ne[i - 1]; - scale_nb[i] = scale_nb[i - 1] * scale_ne[i - 1]; - } - dequant_ne = weight_ne; - dequant_nb[0] = ggml_type_size(dst->type); - for (int i = 1; i < GGML_MAX_DIMS + 1; i++) { - dequant_nb[i] = dequant_nb[i - 1] * dequant_ne[i - 1]; - } - const int64_t scale_offset = ggml_nelements(src0) * sizeof(int8_t); - ggml_cann_pool_alloc dequant_allocator(ctx.pool(), - ggml_nelements(src0) * ggml_type_size(dst->type)); - acl_tensor_ptr acl_weight = ggml_cann_create_tensor(src0->data, ACL_INT8, sizeof(int8_t), - weight_ne, weight_nb, GGML_MAX_DIMS + 1); - acl_tensor_ptr acl_scale = ggml_cann_create_tensor( - src0->data, ACL_FLOAT16, sizeof(uint16_t), scale_ne, scale_nb, - GGML_MAX_DIMS + 1, ACL_FORMAT_ND, scale_offset); - acl_tensor_ptr acl_dequant = ggml_cann_create_tensor( - dequant_allocator.get(), ggml_cann_type_mapping(dst->type), - ggml_type_size(dst->type), dequant_ne, dequant_nb, GGML_MAX_DIMS + 1); - aclnn_mul(ctx, acl_weight.get(), acl_scale.get(), acl_dequant.get()); - - // Reinterpret dequant buffer as 4D [src0->ne] with contiguous strides. - dequant_ne = src0->ne; - dequant_nb[0] = ggml_type_size(dst->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - dequant_nb[i] = dequant_nb[i - 1] * src0->ne[i - 1]; - } - gather_batched(dequant_allocator.get(), - ggml_cann_type_mapping(dst->type), ggml_type_size(dst->type), - dequant_nb); - break; - } - default: - GGML_ABORT("Unsupported tensor type for GGML_OP_GET_ROWS"); - break; - } -} - -void ggml_cann_set_rows(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; // source values - ggml_tensor * src1 = dst->src[1]; // row indices - - // n_idx: number of source rows to scatter per batch slice. - // ggml guarantees: src0->ne[1] == src1->ne[0]. - const int64_t n_idx = src1->ne[0]; - - // Copy n_idx rows of src [ne0, n_idx] into dst [ne0, ne1] at positions given by a 1D index. - // ggml_cann_create_tensor reverses dims, so ACL sees [ne1, ne0] for dst. - // InplaceIndexCopy with dim=0 copies along ACL dim-0 == ggml ne[1] (the row axis). - // src_nb: the 4 strides of the source buffer (nb[0..1] for the 2D slice shape, - // nb[2..3] for computing per-batch-slice base pointer offsets). - auto scatter_batched = [&](void * src_base, aclDataType acl_type, size_t type_size, - const size_t * src_nb) { - int64_t d_ne[2] = { dst->ne[0], dst->ne[1] }; - size_t d_nb[2] = { dst->nb[0], dst->nb[1] }; - int64_t s_ne[2] = { dst->ne[0], n_idx }; - size_t s_nb_2d[2] = { src_nb[0], src_nb[1] }; - int64_t i_ne[1] = { n_idx }; - size_t i_nb[1] = { (size_t)ggml_element_size(src1) }; - - for (int64_t i3 = 0; i3 < dst->ne[3]; i3++) { - for (int64_t i2 = 0; i2 < dst->ne[2]; i2++) { - acl_tensor_ptr acl_dst = ggml_cann_create_tensor( - (char *)dst->data + i3 * dst->nb[3] + i2 * dst->nb[2], - acl_type, type_size, d_ne, d_nb, 2); - acl_tensor_ptr acl_idx = ggml_cann_create_tensor( - (char *)src1->data + (i3 % src1->ne[2]) * src1->nb[2] + (i2 % src1->ne[1]) * src1->nb[1], - ggml_cann_type_mapping(src1->type), (size_t)ggml_element_size(src1), - i_ne, i_nb, 1); - acl_tensor_ptr acl_src = ggml_cann_create_tensor( - (char *)src_base + i3 * src_nb[3] + i2 * src_nb[2], - acl_type, type_size, s_ne, s_nb_2d, 2); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceIndexCopy, acl_dst.get(), 0, acl_idx.get(), acl_src.get()); - } - } - }; - - switch (dst->type) { - case GGML_TYPE_F32: - scatter_batched(src0->data, - ggml_cann_type_mapping(dst->type), ggml_type_size(dst->type), - src0->nb); - break; - case GGML_TYPE_F16: - case GGML_TYPE_BF16: - { - // Cast src0 (F32) to dst type first. - ggml_cann_pool_alloc src_cast_allocator(ctx.pool(), - ggml_nelements(src0) * ggml_type_size(dst->type)); - size_t src_cast_nb[GGML_MAX_DIMS]; - src_cast_nb[0] = ggml_type_size(dst->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - src_cast_nb[i] = src_cast_nb[i - 1] * src0->ne[i - 1]; - } - acl_tensor_ptr acl_src0 = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_src_cast = ggml_cann_create_tensor( - src_cast_allocator.get(), ggml_cann_type_mapping(dst->type), ggml_type_size(dst->type), - src0->ne, src_cast_nb, GGML_MAX_DIMS); - aclnn_cast(ctx, acl_src0.get(), acl_src_cast.get(), ggml_cann_type_mapping(dst->type)); - - scatter_batched(src_cast_allocator.get(), - ggml_cann_type_mapping(dst->type), ggml_type_size(dst->type), - src_cast_nb); - break; - } - default: - GGML_ABORT("Unsupported tensor type for GGML_OP_SET_ROWS"); - break; - } -} - -/** - * @brief Repeats elements of a tensor along a specified dimension. - * - * This function repeats each element of the source tensor `acl_src` a specified - * number of times (`repeats`) along the specified dimension `dim` and stores - * the result in the destination tensor `acl_dst`. - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The source tensor whose elements will be repeated. - * @param acl_dst The destination tensor where the repeated elements will be - * stored. - * @param dim The dimension along which the elements will be repeated. - * @param repeats The number of times each element will be repeated. - * @param output_size The size of the output tensor. - */ -static void aclnn_repeat_interleave(ggml_backend_cann_context & ctx, - aclTensor * acl_src, - aclTensor * acl_dst, - int64_t dim, - int64_t repeats, - int64_t output_size) { - GGML_CANN_CALL_ACLNN_OP(ctx, RepeatInterleaveIntWithDim, acl_src, repeats, dim, output_size, acl_dst); -} - -/** - * @brief Performs matrix multiplication with floating-point precision on - * tensors using the CANN backend. - * - * This function performs matrix multiplication of the input tensor and the - * weight tensor, handling broadcasting and transposing as needed, and stores - * the result in the destination tensor `dst`. - * - * @param ctx The context for the CANN backend operations. - * @param dst The destination tensor where the result of the matrix - * multiplication will be stored. - */ -static void ggml_cann_mat_mul_fp(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * weight = dst->src[0]; // weight - ggml_tensor * input = dst->src[1]; // input - - // when weight ne2 or ne3 is 1, aclnnMatmulGetWorkspaceSize will auto - // broadcast, when weight ne2 or ne3 is not 1, weight need repeat. - BCAST_MUL_MAT_SHAPE(input, weight, dst); - - int64_t n_dims = bcast_dims; - if (bcast_input_ne[3] == bcast_weight_ne[3] && bcast_input_ne[3] == 1) { - if (bcast_input_ne[2] == 1 && bcast_weight_ne[2] == 1) { - n_dims = 2; - } else if (bcast_input_ne[2] == 1) { - n_dims = 3; - } - } - - acl_tensor_ptr acl_input_tensor = ggml_cann_create_tensor(input, bcast_input_ne, bcast_input_nb, n_dims); - int64_t transpose_ne[] = { bcast_weight_ne[1], bcast_weight_ne[0], bcast_weight_ne[2], - bcast_weight_ne[3], bcast_weight_ne[4], bcast_weight_ne[5] }; - size_t transpose_nb[] = { bcast_weight_nb[1], bcast_weight_nb[0], bcast_weight_nb[2], - bcast_weight_nb[3], bcast_weight_nb[4], bcast_weight_nb[5] }; - acl_tensor_ptr acl_weight_tensor; - - // Only check env once. - static bool weight_to_nz = parse_bool(get_env_as_lowercase("GGML_CANN_WEIGHT_NZ").value_or("on")); - if (weight_to_nz && weight->type != GGML_TYPE_BF16 && is_matmul_weight(weight)) { - acl_weight_tensor = ggml_cann_create_tensor(weight, transpose_ne, transpose_nb, n_dims, ACL_FORMAT_FRACTAL_NZ); - } else { - acl_weight_tensor = ggml_cann_create_tensor(weight, transpose_ne, transpose_nb, n_dims, ACL_FORMAT_ND); - } - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst, bcast_dst_ne, bcast_dst_nb, n_dims); - - switch (n_dims) { - case 2: - GGML_CANN_CALL_ACLNN_OP(ctx, Mm, acl_input_tensor.get(), acl_weight_tensor.get(), acl_dst.get(), 2); - break; - case 3: - GGML_CANN_CALL_ACLNN_OP(ctx, BatchMatMul, acl_input_tensor.get(), acl_weight_tensor.get(), acl_dst.get(), - 2); - break; - default: - // ALLOW_FP32_DOWN_PRECISION, when input is - // fp32, atlas a2 will transpose it to HFLOAT32. - GGML_CANN_CALL_ACLNN_OP(ctx, Matmul, acl_input_tensor.get(), acl_weight_tensor.get(), acl_dst.get(), 1); - break; - } -} - -/** - * @brief Performs matrix multiplication with quantized weights and - * floating-point inputs using the CANN backend. - * - * This function performs matrix multiplication of the input tensor `src1` and - * the weight tensor `src0`, handling broadcasting, transposing, and - * quantization as needed, and stores the result in the destination tensor - * `dst`. - * - * @param ctx The context for the CANN backend operations. - * @param dst The destination tensor where the result of the matrix - * multiplication will be stored. - */ -static void ggml_cann_mul_mat_quant(ggml_backend_cann_context & ctx, ggml_tensor * dst, const enum ggml_type type) { - ggml_tensor * src0 = dst->src[0]; // weight - ggml_tensor * src1 = dst->src[1]; // input - - // The shape of the weight is NCHW. - // Matrix multiplication uses HW dims. - // HC is regarded as batch. - // weight need transpose. - float weight_elem_size; - if (type == GGML_TYPE_Q4_0) { - weight_elem_size = float(sizeof(uint8_t)) / 2; - } else if (type == GGML_TYPE_Q8_0) { - weight_elem_size = float(sizeof(uint8_t)); - } else { - GGML_ABORT("Only support Q4_0 and Q8_0 MUL_MAT"); - } - float weight_nb[] = { src0->ne[0] * weight_elem_size, weight_elem_size }; - size_t weight_stride = src0->ne[1] * src0->ne[0] * weight_elem_size; - size_t weight_size = weight_stride * src0->ne[2] * src0->ne[3]; - - // scale stored at the end of weight. Also need transpose. - size_t scale_elem_size = sizeof(uint16_t); - size_t scale_nb[] = { src0->ne[0] / QK8_0 * scale_elem_size, scale_elem_size }; - size_t scale_stride = src0->ne[1] * src0->ne[0] / QK8_0 * scale_elem_size; - char * scale_offset = (char *) src0->data + weight_size; - - // input - size_t input_elem_size = sizeof(uint16_t); - int64_t input_ne[] = { src1->ne[0], src1->ne[1] }; - size_t input_nb[] = { input_elem_size, input_ne[0] * input_elem_size }; - size_t input_stride = input_ne[0] * input_ne[1] * input_elem_size; - ggml_cann_pool_alloc input_alloctor(ctx.pool()); - void * input_buffer = src1->data; - - // case in - if (src1->type != GGML_TYPE_F16) { - acl_tensor_ptr acl_src1_tensor = ggml_cann_create_tensor(src1); - input_buffer = input_alloctor.alloc(ggml_nelements(src1) * input_elem_size); - - int64_t * input_cast_ne = src1->ne; - size_t input_cast_nb[GGML_MAX_DIMS]; - input_cast_nb[0] = sizeof(uint16_t); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - input_cast_nb[i] = input_cast_nb[i - 1] * input_cast_ne[i - 1]; - } - - acl_tensor_ptr acl_input_tensor = ggml_cann_create_tensor(input_buffer, ACL_FLOAT16, input_elem_size, - input_cast_ne, input_cast_nb, GGML_MAX_DIMS); - aclnn_cast(ctx, acl_src1_tensor.get(), acl_input_tensor.get(), ACL_FLOAT16); - } - - // output - size_t output_elem_size = sizeof(uint16_t); - size_t output_nb[] = { output_elem_size, dst->ne[0] * output_elem_size }; - ggml_cann_pool_alloc output_allocator(ctx.pool()); - void * output_buffer = output_allocator.alloc(ggml_nelements(dst) * output_elem_size); - size_t output_stride = dst->ne[0] * dst->ne[1] * output_elem_size; - - // aclnn - int64_t max_elem_size = 65535; - int64_t split_size = (src0->ne[1] / max_elem_size) + 1; - ggml_cann_pool_alloc workspace_allocator(ctx.pool()); - for (int64_t n1 = 0; n1 < src1->ne[3]; n1++) { - for (int64_t c1 = 0; c1 < src1->ne[2]; c1++) { - int64_t n0 = n1 / (src1->ne[3] / src0->ne[3]); - int64_t c0 = c1 / (src1->ne[2] / src0->ne[2]); - - int64_t batch1 = (n1 * src1->ne[2]) + c1; - int64_t batch0 = (n0 * src0->ne[2]) + c0; - - acl_tensor_ptr acl_input_tensor = ggml_cann_create_tensor( - (char *) input_buffer + batch1 * input_stride, ACL_FLOAT16, input_elem_size, input_ne, input_nb, 2); - - // first split - int64_t weight_ne_offset = 0; - int64_t weight_ne[2] = { max_elem_size > src0->ne[1] ? src0->ne[1] : max_elem_size, src0->ne[0] }; - int64_t scale_ne_offset = 0; - int64_t scale_ne[2] = { weight_ne[0], weight_ne[1] / QK8_0 }; - int64_t output_ne_offset = 0; - int64_t output_ne[2] = { weight_ne[0], dst->ne[1] }; - - acl_tensor_ptr acl_weight_tensor = - ggml_cann_create_tensor((char *) src0->data + batch0 * weight_stride, ggml_cann_type_mapping(type), - weight_elem_size, weight_ne, weight_nb, 2, ACL_FORMAT_ND, weight_ne_offset); - acl_tensor_ptr acl_scale_tensor = - ggml_cann_create_tensor(scale_offset + batch0 * scale_stride, ACL_FLOAT16, scale_elem_size, scale_ne, - scale_nb, 2, ACL_FORMAT_ND, scale_ne_offset); - acl_tensor_ptr acl_output_tensor = - ggml_cann_create_tensor((char *) output_buffer + batch1 * output_stride, ACL_FLOAT16, output_elem_size, - output_ne, output_nb, 2, ACL_FORMAT_ND, output_ne_offset); - int64_t antiquantGroupSize = 0; - if (src0->ne[0] > QK8_0) { - antiquantGroupSize = QK8_0; - } - GGML_CANN_CALL_ACLNN_OP(ctx, WeightQuantBatchMatmulV2, acl_input_tensor.get(), acl_weight_tensor.get(), - acl_scale_tensor.get(), nullptr, nullptr, nullptr, nullptr, antiquantGroupSize, - acl_output_tensor.get()); - - // other splits - for (int64_t split = 1; split < split_size; split++) { - weight_ne_offset += weight_elem_size * weight_ne[0] * weight_ne[1]; - weight_ne[0] = - max_elem_size * (split + 1) > src0->ne[1] ? src0->ne[1] - (max_elem_size * split) : max_elem_size; - scale_ne_offset += scale_elem_size * scale_ne[0] * scale_ne[1]; - scale_ne[0] = weight_ne[0]; - output_ne_offset += output_elem_size * output_ne[0] * output_ne[1]; - output_ne[0] = weight_ne[0]; - - acl_weight_tensor = - ggml_cann_create_tensor((char *) src0->data + batch0 * weight_stride, ggml_cann_type_mapping(type), - weight_elem_size, weight_ne, weight_nb, 2, ACL_FORMAT_ND, weight_ne_offset); - acl_scale_tensor = - ggml_cann_create_tensor(scale_offset + batch0 * scale_stride, ACL_FLOAT16, scale_elem_size, - scale_ne, scale_nb, 2, ACL_FORMAT_ND, scale_ne_offset); - acl_output_tensor = - ggml_cann_create_tensor((char *) output_buffer + batch1 * output_stride, ACL_FLOAT16, - output_elem_size, output_ne, output_nb, 2, ACL_FORMAT_ND, output_ne_offset); - GGML_CANN_CALL_ACLNN_OP(ctx, WeightQuantBatchMatmulV2, acl_input_tensor.get(), acl_weight_tensor.get(), - acl_scale_tensor.get(), nullptr, nullptr, nullptr, nullptr, antiquantGroupSize, - acl_output_tensor.get()); - } - } - } - - // cast out - if (dst->type != GGML_TYPE_F16) { - int64_t * output_cast_ne = dst->ne; - size_t output_cast_nb[GGML_MAX_DIMS]; - output_cast_nb[0] = sizeof(uint16_t); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - output_cast_nb[i] = output_cast_nb[i - 1] * output_cast_ne[i - 1]; - } - - acl_tensor_ptr acl_output_tensor = ggml_cann_create_tensor(output_buffer, ACL_FLOAT16, output_elem_size, - output_cast_ne, output_cast_nb, GGML_MAX_DIMS); - acl_tensor_ptr acl_dst_tensor = ggml_cann_create_tensor(dst); - aclnn_cast(ctx, acl_output_tensor.get(), acl_dst_tensor.get(), ggml_cann_type_mapping(dst->type)); - } -} - -void ggml_cann_mul_mat(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - const enum ggml_type type = dst->src[0]->type; - switch (type) { - case GGML_TYPE_F32: - case GGML_TYPE_F16: -#ifndef ASCEND_310P - case GGML_TYPE_BF16: -#endif - ggml_cann_mat_mul_fp(ctx, dst); - break; - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q8_0: - ggml_cann_mul_mat_quant(ctx, dst, type); - break; - default: - GGML_ABORT("Unsupported type for mul_mat"); - break; - } -} - -/** - * @brief Rolls the elements of a tensor along a specified dimension. - * - * This function rolls the elements of the source tensor `acl_src` by the - * specified shifts `shifts` along the specified dimensions `dims`, and stores - * the result in the destination tensor `acl_dst`. - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The source tensor whose elements will be rolled. - * @param acl_dst The destination tensor where the rolled elements will be - * stored. - * @param shifts An array specifying the number of positions by which elements - * are shifted. - * @param dims An array specifying the dimensions along which elements are - * shifted. - */ -static void aclnn_roll(ggml_backend_cann_context & ctx, - aclTensor * acl_src, - aclTensor * acl_dst, - int64_t * shifts, - int64_t * dims) { - acl_int_array_ptr acl_shifts = ggml_cann_create_int_array(shifts, 1); - acl_int_array_ptr acl_dims = ggml_cann_create_int_array(dims, 1); - GGML_CANN_CALL_ACLNN_OP(ctx, Roll, acl_src, acl_shifts.get(), acl_dims.get(), acl_dst); -} - -/** - * @brief Fills specified positions of a tensor with a scalar value. - * - * This function fills the positions in the source tensor `acl_src` specified by - * `index` along the dimension `dim` with the scalar value `value`. - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The source tensor where the positions will be filled. - * @param dim The dimension along which the positions are specified. - * @param index An array specifying the positions to be filled. - * @param index_num The number of positions specified in the index array. - * @param value The scalar value used to fill the specified positions. - */ -static void aclnn_index_fill_tensor(ggml_backend_cann_context & ctx, - aclTensor * acl_src, - int64_t dim, - int64_t * index, - int64_t index_num, - float value) { - acl_int_array_ptr acl_index = ggml_cann_create_int_array(index, index_num); - acl_scalar_ptr acl_value = ggml_cann_create_scalar(&value, aclDataType::ACL_FLOAT); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceIndexFillTensor, acl_src, dim, acl_index.get(), acl_value.get()); -} - -/** - * @brief Initializes and caches all intermediate tensors required for RoPE - * (Rotary Position Embedding), including support for Yarn, mRoPE, - * i-mRoPE, Neox repeat strategy, independent sectors, frequency factors， - * and multi-section rotary groups. - * - * This function computes and caches the per-dimension θ coefficients used for - * Q/K rotary embedding. The cache is shared across layers, and recomputed only - * when any dependent parameter changes. - * - * The function now supports: - * - Yarn RoPE extrapolation (via @param corr_dims and @param ext_factor) - * - Per-dimension independent sector exponent rules (indep_sects + sections[]) - * - Multi-section RoPE (mRoPE) index mapping (mrope_used + is_imrope) - * - Frequency factor division (src2) - * - Neox / normal repeat expansion modes - * - * @param ctx CANN backend context, containing memory pool, - * cached buffers, and runtime stream. - * @param dst Destination ggml_tensor whose computation - * depends on RoPE (typically Qcur or Kcur). - * @param corr_dims [low, high] Yarn correction range. - * @param ext_factor Yarn extrapolation strength. 0 = disabled. - * @param theta_scale Base multiplier for per-dimension θ exponent. - * @param freq_scale Global frequency scaling factor. - * @param attn_factor Optional scaling applied to sin/cos (if needed). - * @param is_neox Whether to use Neox-style dimension interleave. - * @param sections 4-way sector sizes for independent-section RoPE - * and multi-section mRoPE (t/h/w/e). - * @param mrope_used Whether to enable multi-section rotary embedding. - * @param is_imrope Whether to apply interleaved mRoPE rules. - * @param indep_sects Whether each dimension runs independent exponent - * resets based on @p sections. - */ -static void aclnn_rope_cache_init(ggml_backend_cann_context & ctx, - ggml_tensor * dst, - float * corr_dims, - float ext_factor, - float theta_scale, - float freq_scale, - float attn_factor, - bool is_neox, - int sections[4], - bool mrope_used, - bool is_imrope, - bool indep_sects, - int64_t rope_dims) { - ggml_tensor * src1 = dst->src[1]; // position - ggml_tensor * src2 = dst->src[2]; // freq_factors - - int64_t theta_scale_length = rope_dims / 2; - int64_t position_length = dst->ne[2]; - - // TODO: check theta_scale_length and position_length. - if (src2 == nullptr && ctx.rope_cache.cached && - ctx.rope_cache.equal(theta_scale_length, position_length, ext_factor, theta_scale, freq_scale, attn_factor, - is_neox, indep_sects, mrope_used, is_imrope, sections)) { - // use cache. - return; - } - - // Step0: calculate tensor shape. - int64_t theta_scale_ne[] = { theta_scale_length, 1, 1, 1 }; - size_t theta_scale_nb[] = { sizeof(float), theta_scale_length * sizeof(float), theta_scale_length * sizeof(float), - theta_scale_length * sizeof(float) }; - - GGML_ASSERT(src1->type == GGML_TYPE_I32); - int64_t position_ne[] = { 1, 1, position_length, 1 }; - size_t position_nb[] = { sizeof(int32_t), sizeof(int32_t), sizeof(int32_t), sizeof(int32_t) * position_length }; - - int64_t cache_ne[] = { theta_scale_length, 1, position_length, 1 }; - size_t cache_nb[GGML_MAX_DIMS]; - cache_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - cache_nb[i] = cache_nb[i - 1] * cache_ne[i - 1]; - } - - // Step1: Compute the coefficient of theta. During the cache_init process, aside from - // (1) multiplying by the position, - // (2) dividing by freq_factors, - // (3) computing the sine and cosine, - // the other parameters used in the computation generally do not change in most scenarios. - // Therefore, we can first compute this part of the result and then cache it. - - // Step1.1: prepare theta_scale exponent. if this exponent updated, should update theta_scale_tensor. - acl_tensor_ptr acl_theta_scale_tensor; - bool theta_scale_updated = false; - if (ctx.rope_cache.theta_scale_length != theta_scale_length || ctx.rope_cache.theta_scale != theta_scale || - ctx.rope_cache.indep_sects != indep_sects) { - theta_scale_updated = true; - if (ctx.rope_cache.theta_scale_exp_host != nullptr) { - free(ctx.rope_cache.theta_scale_exp_host); - } - ctx.rope_cache.theta_scale_exp_host = (float *) malloc(theta_scale_length * sizeof(float)); - GGML_ASSERT(ctx.rope_cache.theta_scale_exp_host != nullptr); - if (!indep_sects) { - ctx.rope_cache.theta_scale_exp_host[0] = 1; - for (int i = 1; i < theta_scale_length; i++) { - ctx.rope_cache.theta_scale_exp_host[i] = ctx.rope_cache.theta_scale_exp_host[i - 1] * theta_scale; - } - } else { - int sect_dims = sections[0] + sections[1] + sections[2] + sections[3]; - int sec_w = sections[1] + sections[0]; - int sec_e = sections[2] + sec_w; - - ctx.rope_cache.theta_scale_exp_host[0] = 1; - for (int i = 1; i < theta_scale_length; i++) { - int sector = i % sect_dims; - if (sector == 0 || sector == sections[0] || sector == sec_w || sector == sec_e) { - ctx.rope_cache.theta_scale_exp_host[i] = 1; - continue; - } - ctx.rope_cache.theta_scale_exp_host[i] = ctx.rope_cache.theta_scale_exp_host[i - 1] * theta_scale; - } - } - - if (ctx.rope_cache.theta_scale_cache != nullptr) { - ACL_CHECK(aclrtFree(ctx.rope_cache.theta_scale_cache)); - } - ACL_CHECK(aclrtMalloc(&ctx.rope_cache.theta_scale_cache, theta_scale_length * sizeof(float), - ACL_MEM_MALLOC_HUGE_FIRST)); - - ACL_CHECK(aclrtMemcpyAsync(ctx.rope_cache.theta_scale_cache, theta_scale_length * sizeof(float), - ctx.rope_cache.theta_scale_exp_host, theta_scale_length * sizeof(float), - ACL_MEMCPY_HOST_TO_DEVICE, ctx.stream())); - } - acl_theta_scale_tensor = ggml_cann_create_tensor(ctx.rope_cache.theta_scale_cache, ACL_FLOAT, sizeof(float), - theta_scale_ne, theta_scale_nb, 1); - - // Step1.2: prepare rope_yarn_ramp, if this part updated, should update theta_scale_tensor. - // TODO: acl_yarn_ramp_tensor use rope cache. - bool yarn_ramp_tensor_updated = false; - acl_tensor_ptr acl_yarn_ramp_tensor; - if (ext_factor != 0 && (theta_scale_updated || ctx.rope_cache.theta_scale_length != theta_scale_length || - ctx.rope_cache.freq_scale != freq_scale)) { - yarn_ramp_tensor_updated = true; - if (ctx.rope_cache.yarn_ramp_cache != nullptr) { - ACL_CHECK(aclrtFree(ctx.rope_cache.yarn_ramp_cache)); - } - ACL_CHECK(aclrtMalloc(&ctx.rope_cache.yarn_ramp_cache, theta_scale_length * sizeof(float), - ACL_MEM_MALLOC_HUGE_FIRST)); - // -rope_yarn_ramp - // const float y = (i0 / 2 - low) / MAX(0.001f, high - low); - // return MIN(1, MAX(0, y)) - 1; - acl_yarn_ramp_tensor = ggml_cann_create_tensor(ctx.rope_cache.yarn_ramp_cache, ACL_FLOAT, sizeof(float), - theta_scale_ne, theta_scale_nb, 1); - float zero_value = 0, one_value = 1; - float denom_safe_value = MAX(0.001f, corr_dims[1] - corr_dims[0]); - acl_scalar_ptr low = ggml_cann_create_scalar(&corr_dims[0], aclDataType::ACL_FLOAT); - acl_scalar_ptr zero = ggml_cann_create_scalar(&zero_value, aclDataType::ACL_FLOAT); - acl_scalar_ptr one = ggml_cann_create_scalar(&one_value, aclDataType::ACL_FLOAT); - acl_scalar_ptr denom_safe = ggml_cann_create_scalar(&denom_safe_value, aclDataType::ACL_FLOAT); - acl_scalar_ptr ext_factor_sc = ggml_cann_create_scalar(&ext_factor, aclDataType::ACL_FLOAT); - - aclnn_arange(ctx, acl_yarn_ramp_tensor.get(), 0, theta_scale_length, 1, theta_scale_length); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceSubs, acl_yarn_ramp_tensor.get(), low.get(), one.get()); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceDivs, acl_yarn_ramp_tensor.get(), denom_safe.get()); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceThreshold, acl_yarn_ramp_tensor.get(), zero.get(), zero.get()); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceClampMax, acl_yarn_ramp_tensor.get(), one.get()); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceSubs, acl_yarn_ramp_tensor.get(), one.get(), one.get()); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceMuls, acl_yarn_ramp_tensor.get(), ext_factor_sc.get()); - - // theta_interp = freq_scale * theta_extrap; - // theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix; - // theta = freq_scale * theta_extrap * (1 - ramp_mix) + theta_extrap * ramp_mix; - // theta = freq_scale * theta_extrap - freq_scale * theta_extrap * ramp_mix + theta_extrap * ramp_mix; - // theta = theta_extrap * (freq_scale - freq_scale * ramp_mix + ramp_mix); - // - // we cache (freq_scale - freq_scale * ramp_mix + ramp_mix), Considering that the rope_yarn_ramp here is the inverse - // cache freq_scale + (freq_scale - 1) * ramp_mix - float freq_scale_1 = freq_scale - 1; - acl_scalar_ptr freq_scale_sc = ggml_cann_create_scalar(&freq_scale, aclDataType::ACL_FLOAT); - acl_scalar_ptr freq_scale_1_sc = ggml_cann_create_scalar(&freq_scale_1, aclDataType::ACL_FLOAT); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceMuls, acl_yarn_ramp_tensor.get(), freq_scale_1_sc.get()); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceAdds, acl_yarn_ramp_tensor.get(), freq_scale_sc.get(), one.get()); - } else { - acl_yarn_ramp_tensor = ggml_cann_create_tensor(ctx.rope_cache.yarn_ramp_cache, ACL_FLOAT, sizeof(float), - theta_scale_ne, theta_scale_nb, 1); - } - // Step 1.3: update theta_scale_tensor according to ext_factor or freq_scale. - if (ext_factor != 0) { - if (theta_scale_updated || yarn_ramp_tensor_updated) { - theta_scale_updated = true; - aclnn_mul(ctx, acl_theta_scale_tensor.get(), acl_yarn_ramp_tensor.get()); - } - } else { - if (freq_scale != 1 && (ctx.rope_cache.freq_scale != freq_scale || theta_scale_updated)) { - theta_scale_updated = true; - aclnn_muls(ctx, acl_theta_scale_tensor.get(), freq_scale, nullptr, true); - } - } - - // Nothing changed, use cache. - if (!theta_scale_updated) { - acl_theta_scale_tensor = ggml_cann_create_tensor(ctx.rope_cache.theta_scale_cache, ACL_FLOAT, sizeof(float), - theta_scale_ne, theta_scale_nb, GGML_MAX_DIMS); - } - - // Step 1.4: prepare select index if mrope - acl_tensor_ptr position_select_index_tensor; - if (mrope_used) { - if (ctx.rope_cache.sections[0] != sections[0] || ctx.rope_cache.sections[1] != sections[1] || - ctx.rope_cache.sections[2] != sections[2] || ctx.rope_cache.sections[3] != sections[3] || - ctx.rope_cache.theta_scale_length != theta_scale_length || ctx.rope_cache.is_imrope != is_imrope) { - if (ctx.rope_cache.position_select_index_host != nullptr) { - free(ctx.rope_cache.position_select_index_host); - } - ctx.rope_cache.position_select_index_host = (int *) malloc(theta_scale_length * sizeof(int)); - GGML_ASSERT(ctx.rope_cache.position_select_index_host != nullptr); - int sect_dims = sections[0] + sections[1] + sections[2] + sections[3]; - int sec_w = sections[1] + sections[0]; - int sec_e = sections[2] + sec_w; - // t,h,w,e - for (int i = 0; i < theta_scale_length; i++) { - int sector = i % sect_dims; - - if (is_imrope) { // qwen3vl apply interleaved mrope - if (sector % 3 == 1 && sector < 3 * sections[1]) { - ctx.rope_cache.position_select_index_host[i] = 1; - } else if (sector % 3 == 2 && sector < 3 * sections[2]) { - ctx.rope_cache.position_select_index_host[i] = 2; - } else if (sector % 3 == 0 && sector < 3 * sections[0]) { - ctx.rope_cache.position_select_index_host[i] = 0; - } else { - ctx.rope_cache.position_select_index_host[i] = 3; - } - } else { - if (sector >= sections[0] && sector < sec_w) { - ctx.rope_cache.position_select_index_host[i] = 1; - } else if (sector >= sec_w && sector < sec_e) { - ctx.rope_cache.position_select_index_host[i] = 2; - } else if (sector >= sec_e) { - ctx.rope_cache.position_select_index_host[i] = 3; - } else { - ctx.rope_cache.position_select_index_host[i] = 0; - } - } - } - - if (ctx.rope_cache.position_select_index != nullptr) { - ACL_CHECK(aclrtFree(ctx.rope_cache.position_select_index)); - } - ACL_CHECK(aclrtMalloc(&ctx.rope_cache.position_select_index, theta_scale_length * sizeof(int), - ACL_MEM_MALLOC_HUGE_FIRST)); - - ACL_CHECK(aclrtMemcpyAsync(ctx.rope_cache.position_select_index, theta_scale_length * sizeof(int), - ctx.rope_cache.position_select_index_host, theta_scale_length * sizeof(int), - ACL_MEMCPY_HOST_TO_DEVICE, ctx.stream())); - } - - position_select_index_tensor = ggml_cann_create_tensor(ctx.rope_cache.position_select_index, ACL_INT32, - sizeof(int), theta_scale_ne, theta_scale_nb, 1); - } - - // Step2: divide by freq_factors - ggml_cann_pool_alloc freq_fac_res_allocator(ctx.pool()); - if (src2) { - freq_fac_res_allocator.alloc(theta_scale_length * sizeof(float)); - void * freq_fac_res_ptr = freq_fac_res_allocator.get(); - acl_tensor_ptr acl_freq_factors_tensor = - ggml_cann_create_tensor(src2->data, ggml_cann_type_mapping(src2->type), ggml_type_size(src2->type), - theta_scale_ne, theta_scale_nb, GGML_MAX_DIMS); - acl_tensor_ptr acl_freq_fac_res_tensor = ggml_cann_create_tensor(freq_fac_res_ptr, ACL_FLOAT, sizeof(float), - theta_scale_ne, theta_scale_nb, GGML_MAX_DIMS); - aclnn_div(ctx, acl_theta_scale_tensor.get(), acl_freq_factors_tensor.get(), acl_freq_fac_res_tensor.get()); - std::swap(acl_theta_scale_tensor, acl_freq_fac_res_tensor); - } - - // Step3: prepare position_tensor - acl_tensor_ptr acl_position_tensor; - ggml_cann_pool_alloc mrope_position_acllocator(ctx.pool()); - if (mrope_used) { - // Step3.1: select current position; - // position : - // pos1: [[0, 1 ,2 ,3 ], - // pos2: [4, 5 ,6 ,7 ], - // pos3: [8, 9 ,10,11], - // pos4: [12,13,14,15] ] - // - // select index = [0, 1, 2, 2, 1, 0] - // - // selected_tensor: - // [[0, 1 ,2 ,3 ], - // [4, 5 ,6 ,7 ], - // [8, 9 ,10,11], - // [8, 9 ,10,11], - // [4, 5 ,6 ,7 ], - // [0, 1 ,2 ,3 ]] - // - // transpose, from [seq_len:dims] to [dims:seq_len] - // [0, 4, 8 ,8 ,4, 0], - // [1, 5, 9, 9, 5, 1], - // [2, 6, 10,10,6 ,2], - // [3, 7, 11,11,7 3 ]] - // - // multipy by theta_scale_tensor - // [theta_scale^0, theta_scale^1, ..., theta_scale ^ n] - - int64_t mrope_position_ne[] = { position_length, 4 }; - size_t mrope_position_nb[] = { sizeof(int), position_length * sizeof(int) }; - acl_tensor_ptr mrope_position = - ggml_cann_create_tensor(src1->data, ggml_cann_type_mapping(src1->type), ggml_type_size(src1->type), - mrope_position_ne, mrope_position_nb, 2); - - // selected position tensor's shape is a transpose of cache tensor. - int64_t selected_position_ne[] = { position_length, theta_scale_length }; - size_t selected_position_nb[] = { sizeof(float), position_length * sizeof(float) }; - mrope_position_acllocator.alloc(theta_scale_length * position_length * sizeof(float)); - void * mrope_position_buffer = mrope_position_acllocator.get(); - acl_position_tensor = - ggml_cann_create_tensor(mrope_position_buffer, ggml_cann_type_mapping(src1->type), - ggml_type_size(src1->type), selected_position_ne, selected_position_nb, 2); - GGML_CANN_CALL_ACLNN_OP(ctx, IndexSelect, mrope_position.get(), 0, position_select_index_tensor.get(), - acl_position_tensor.get()); - - // transpose - int64_t transposed_ne[] = { position_length, 1, theta_scale_length, 1 }; - size_t transposed_nb[GGML_MAX_DIMS]; - transposed_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - transposed_nb[i] = transposed_nb[i - 1] * transposed_ne[i - 1]; - } - - std::swap(transposed_ne[0], transposed_ne[2]); - std::swap(transposed_nb[0], transposed_nb[2]); - - acl_position_tensor = - ggml_cann_create_tensor(mrope_position_buffer, ggml_cann_type_mapping(src1->type), - ggml_type_size(src1->type), transposed_ne, transposed_nb, GGML_MAX_DIMS); - - } else { - // auto bcast. - acl_position_tensor = - ggml_cann_create_tensor(src1->data, ggml_cann_type_mapping(src1->type), ggml_type_size(src1->type), - position_ne, position_nb, GGML_MAX_DIMS); - } - - // Step4: multiply by the position - int64_t theta_length = theta_scale_length * position_length; - ggml_cann_pool_alloc theta_allocator(ctx.pool(), theta_length * sizeof(float)); - void * theta_buffer = theta_allocator.get(); - - acl_tensor_ptr acl_theta_tensor = - ggml_cann_create_tensor(theta_buffer, ACL_FLOAT, sizeof(float), cache_ne, cache_nb, GGML_MAX_DIMS); - aclnn_mul(ctx, acl_position_tensor.get(), acl_theta_scale_tensor.get(), acl_theta_tensor.get()); - - // Step5: calculate sin cos. - // init sin_repeat && cos_repeat, only to accelerate first layer on each device - if (position_length > ctx.rope_cache.position_length) { - ctx.rope_cache.position_length = position_length; - if (ctx.rope_cache.sin_cache != nullptr) { - ACL_CHECK(aclrtFree(ctx.rope_cache.sin_cache)); - } - if (ctx.rope_cache.cos_cache != nullptr) { - ACL_CHECK(aclrtFree(ctx.rope_cache.cos_cache)); - } - int64_t repeat_theta_length = theta_scale_length * position_length * 2; - ACL_CHECK( - aclrtMalloc(&ctx.rope_cache.sin_cache, repeat_theta_length * sizeof(float), ACL_MEM_MALLOC_HUGE_FIRST)); - ACL_CHECK( - aclrtMalloc(&ctx.rope_cache.cos_cache, repeat_theta_length * sizeof(float), ACL_MEM_MALLOC_HUGE_FIRST)); - } - - // sin/cos - ggml_cann_pool_alloc sin_allocator(ctx.pool(), theta_length * sizeof(float)); - void * sin_buffer = sin_allocator.get(); - acl_tensor_ptr acl_sin_tensor = - ggml_cann_create_tensor(sin_buffer, ACL_FLOAT, sizeof(float), cache_ne, cache_nb, GGML_MAX_DIMS, ACL_FORMAT_ND); - aclnn_sin(ctx, acl_theta_tensor.get(), acl_sin_tensor.get()); - - ggml_cann_pool_alloc cos_allocator(ctx.pool(), theta_length * sizeof(float)); - void * cos_buffer = cos_allocator.get(); - acl_tensor_ptr acl_cos_tensor = - ggml_cann_create_tensor(cos_buffer, ACL_FLOAT, sizeof(float), cache_ne, cache_nb, GGML_MAX_DIMS, ACL_FORMAT_ND); - aclnn_cos(ctx, acl_theta_tensor.get(), acl_cos_tensor.get()); - - if (ext_factor != 0) { - attn_factor *= 1.0f + 0.1f * logf(1.0f / freq_scale); - } - - // Step 5: multiply by attn_factor - if (attn_factor != 1) { - aclnn_muls(ctx, acl_sin_tensor.get(), attn_factor, nullptr, true); - aclnn_muls(ctx, acl_cos_tensor.get(), attn_factor, nullptr, true); - } - - int64_t sin_reshape_ne[4] = { rope_dims, 1, dst->ne[2], 1 }; - size_t sin_reshape_nb[GGML_MAX_DIMS]; - sin_reshape_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - sin_reshape_nb[i] = sin_reshape_nb[i - 1] * sin_reshape_ne[i - 1]; - } - acl_tensor_ptr acl_sin_repeat_tensor = ggml_cann_create_tensor(ctx.rope_cache.sin_cache, ACL_FLOAT, sizeof(float), - sin_reshape_ne, sin_reshape_nb, GGML_MAX_DIMS); - acl_tensor_ptr acl_cos_repeat_tensor = ggml_cann_create_tensor(ctx.rope_cache.cos_cache, ACL_FLOAT, sizeof(float), - sin_reshape_ne, sin_reshape_nb, GGML_MAX_DIMS); - - // Step 6: repeat - if (is_neox) { - // [sinθ1, sinθ1, sinθ2, sinθ2, ..., sinθn, sinθn] - int64_t repeatsArray[] = { 1, 1, 1, 2 }; - aclnn_repeat(ctx, acl_sin_tensor.get(), acl_sin_repeat_tensor.get(), repeatsArray); - aclnn_repeat(ctx, acl_cos_tensor.get(), acl_cos_repeat_tensor.get(), repeatsArray); - } else { - int64_t num_repeats = 2; - int64_t dim = 3; - int64_t output_size = theta_scale_length * num_repeats; - // [sinθ1, sinθ2, ..., sinθn, sinθ1, sinθ2, ..., sinθn] - aclnn_repeat_interleave(ctx, acl_sin_tensor.get(), acl_sin_repeat_tensor.get(), dim, num_repeats, output_size); - aclnn_repeat_interleave(ctx, acl_cos_tensor.get(), acl_cos_repeat_tensor.get(), dim, num_repeats, output_size); - } - - // Update cached value. - ctx.rope_cache.cached = true; - ctx.rope_cache.set(theta_scale_length, position_length, ext_factor, theta_scale, freq_scale, attn_factor, is_neox, - indep_sects, mrope_used, is_imrope, sections); -} - -#ifdef __cplusplus -extern "C" { -#endif -aclnnStatus aclnnRotaryPositionEmbeddingGetWorkspaceSize(const aclTensor * x, - const aclTensor * cos, - const aclTensor * sin, - int64_t mode, - const aclTensor * yOut, - uint64_t * workspaceSize, - aclOpExecutor ** executor); -aclnnStatus aclnnRotaryPositionEmbedding(void * workspace, - uint64_t workspaceSize, - aclOpExecutor * executor, - aclrtStream stream); -#ifdef __cplusplus -} -#endif - -void ggml_cann_rope(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; // input - - // param - float freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow; - int sections[4]; - // const int n_past = ((int32_t *) dst->op_params)[0]; - const int n_dims = ((int32_t *) dst->op_params)[1]; - const int mode = ((int32_t *) dst->op_params)[2]; - // const int n_ctx = ((int32_t *) dst->op_params)[3]; - const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; - - GGML_TENSOR_UNARY_OP_LOCALS - - memcpy(&freq_base, (int32_t *) dst->op_params + 5, sizeof(float)); - memcpy(&freq_scale, (int32_t *) dst->op_params + 6, sizeof(float)); - memcpy(&ext_factor, (int32_t *) dst->op_params + 7, sizeof(float)); - memcpy(&attn_factor, (int32_t *) dst->op_params + 8, sizeof(float)); - memcpy(&beta_fast, (int32_t *) dst->op_params + 9, sizeof(float)); - memcpy(&beta_slow, (int32_t *) dst->op_params + 10, sizeof(float)); - memcpy(§ions, (int32_t *) dst->op_params + 11, sizeof(int) * 4); - - GGML_ASSERT(n_dims % 2 == 0); - GGML_ASSERT(n_dims <= ne00); - - const float theta_scale = powf(freq_base, -2.0f / n_dims); - - float corr_dims[2]; - ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); - - bool is_neox = mode & GGML_ROPE_TYPE_NEOX; - const bool is_imrope = mode == GGML_ROPE_TYPE_IMROPE; // qwen3vl apply interleaved mrope - // mrope_used means the GGML_ROPE_TYPE_MROPE bit is set. - // Note: this bit is also set for imrope and some vision modes, - // so mrope_used does NOT exclusively indicate pure mrope. - const bool mrope_used = mode & GGML_ROPE_TYPE_MROPE; - const bool is_vision = mode == GGML_ROPE_TYPE_VISION; - - if (mrope_used) { - GGML_ASSERT(sections[0] > 0 || sections[1] > 0 || sections[2] > 0); - } - - if (is_vision) { - GGML_ASSERT(n_dims == ne0 / 2); - } - - if (is_imrope || mrope_used) { - is_neox = true; - } - - int64_t rope_dims = n_dims; - - //Our current RotaryPositionEmbedding does not support the VISION mode, - //but essentially it only modifies theta_base in mrope, - //then repeats it at the end in the same way as is_neox. - //In fact, RoPE is still applied across all dimensions. - if (is_vision) { - rope_dims = src0->ne[0]; - } - int64_t tail_dims = ne00 - rope_dims; - bool has_tail = tail_dims > 0; - - // init ctx.rope_cos/rope_sin cache - aclnn_rope_cache_init(ctx, dst, corr_dims, ext_factor, theta_scale, freq_scale, attn_factor, is_neox, sections, - mrope_used, is_imrope, is_vision, rope_dims); - - // Cache is generated with ne00 dimensions, so we use ne00 for reshape - int64_t sin_reshape_ne[4] = { rope_dims, 1, ne02, 1 }; - size_t sin_reshape_nb[GGML_MAX_DIMS]; - sin_reshape_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - sin_reshape_nb[i] = sin_reshape_nb[i - 1] * sin_reshape_ne[i - 1]; - } - acl_tensor_ptr acl_sin_reshape_tensor = ggml_cann_create_tensor(ctx.rope_cache.sin_cache, ACL_FLOAT, sizeof(float), - sin_reshape_ne, sin_reshape_nb, GGML_MAX_DIMS); - acl_tensor_ptr acl_cos_reshape_tensor = ggml_cann_create_tensor(ctx.rope_cache.cos_cache, ACL_FLOAT, sizeof(float), - sin_reshape_ne, sin_reshape_nb, GGML_MAX_DIMS); - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); -#ifdef ASCEND_310P - // Special ROPE operation for 310P - - // roll input - void * input_roll_buffer; - acl_tensor_ptr acl_minus_one_tensor; - void * minus_one_scale_buffer = nullptr; - ggml_cann_pool_alloc roll_allocator(ctx.pool(), ggml_nbytes(src0)); - ggml_cann_pool_alloc minus_one_scale_allocator(ctx.pool(), sizeof(float) * src0->ne[0]); - if (!is_neox) { - // roll input: [q0,q1,q2,q3,...] -> [q1,q0,q3,q2,...] - input_roll_buffer = roll_allocator.get(); - int64_t input_roll_ne[4] = { 2, src0->ne[1] * (src0->ne[0] / 2), src0->ne[2], src0->ne[3] }; - size_t input_roll_nb[GGML_MAX_DIMS]; - input_roll_nb[0] = ggml_type_size(src0->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - input_roll_nb[i] = input_roll_nb[i - 1] * input_roll_ne[i - 1]; - } - acl_tensor_ptr acl_input_roll_tensor = - ggml_cann_create_tensor(input_roll_buffer, ggml_cann_type_mapping(src0->type), ggml_type_size(src0->type), - input_roll_ne, input_roll_nb, GGML_MAX_DIMS); - acl_tensor_ptr acl_input_tensor = - ggml_cann_create_tensor(src0->data, ggml_cann_type_mapping(src0->type), ggml_type_size(src0->type), - input_roll_ne, input_roll_nb, GGML_MAX_DIMS); - - int64_t shifts[] = { 1 }; - int64_t dims[] = { 3 }; - aclnn_roll(ctx, acl_input_tensor.get(), acl_input_roll_tensor.get(), shifts, dims); - - // init [-1, 1, -1, 1, ...] - minus_one_scale_buffer = minus_one_scale_allocator.get(); - - int64_t minus_one_ne[4] = { src0->ne[0], 1, 1, 1 }; - size_t minus_one_nb[GGML_MAX_DIMS]; - minus_one_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - minus_one_nb[i] = minus_one_nb[i - 1] * minus_one_ne[i - 1]; - } - acl_minus_one_tensor = aclnn_values(ctx, minus_one_scale_buffer, sizeof(float) * src0->ne[0], minus_one_ne, - GGML_MAX_DIMS, ACL_FLOAT, sizeof(float), 1); - int64_t dim = 3; - int64_t * index = new int64_t[src0->ne[0]]; - for (int i = 0; i < src0->ne[0]; i++) { - index[i] = i / 2 * 2; - } - int64_t index_num = src0->ne[0]; - float value = -1; - aclnn_index_fill_tensor(ctx, acl_minus_one_tensor.get(), dim, index, index_num, value); - } else { - // roll input: [q0,q1,q2,...] -> - // [q_half,q_half+1,...,q_end,q0,q1,...q_half-1] - input_roll_buffer = roll_allocator.get(); - acl_tensor_ptr acl_input_roll_tensor = - ggml_cann_create_tensor(input_roll_buffer, ggml_cann_type_mapping(src0->type), ggml_type_size(src0->type), - src0->ne, src0->nb, GGML_MAX_DIMS); - acl_tensor_ptr acl_input_tensor = ggml_cann_create_tensor(src0); - - int64_t shifts[] = { src0->ne[0] / 2 }; - int64_t dims[] = { 3 }; - aclnn_roll(ctx, acl_input_tensor.get(), acl_input_roll_tensor.get(), shifts, dims); - - // init [-1, -1, -1, 1, 1，1，...] - minus_one_scale_buffer = minus_one_scale_allocator.get(); - int64_t minus_one_ne[4] = { src0->ne[0], 1, 1, 1 }; - size_t minus_one_nb[GGML_MAX_DIMS]; - minus_one_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - minus_one_nb[i] = minus_one_nb[i - 1] * minus_one_ne[i - 1]; - } - acl_minus_one_tensor = aclnn_values(ctx, minus_one_scale_buffer, sizeof(float) * src0->ne[0], minus_one_ne, - GGML_MAX_DIMS, ACL_FLOAT, sizeof(float), 1); - // -1 * first half - int64_t first_half_ne[4] = { src0->ne[0] / 2, 1, 1, 1 }; - size_t first_half_nb[GGML_MAX_DIMS]; - first_half_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - first_half_nb[i] = first_half_nb[i - 1] * first_half_ne[i - 1]; - } - acl_tensor_ptr acl_first_half_tensor = ggml_cann_create_tensor(minus_one_scale_buffer, ACL_FLOAT, sizeof(float), - first_half_ne, first_half_nb, GGML_MAX_DIMS); - bool inplace = true; - float scale = -1; - aclnn_muls(ctx, acl_first_half_tensor.get(), scale, nullptr, inplace); - } - - // TODO: n_dims < ne0 - GGML_ASSERT(n_dims == src0->ne[0]); - - // input * scale - ggml_cann_pool_alloc roll_mul_scale_allocator(ctx.pool(), ggml_nbytes(src0)); - void * input_roll_mul_scale_buffer = roll_mul_scale_allocator.get(); - size_t input_nb[GGML_MAX_DIMS]; - input_nb[0] = ggml_type_size(src0->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - input_nb[i] = input_nb[i - 1] * src0->ne[i - 1]; - } - acl_tensor_ptr acl_input_roll_mul_scale_tensor = - ggml_cann_create_tensor(input_roll_mul_scale_buffer, ggml_cann_type_mapping(src0->type), - ggml_type_size(src0->type), src0->ne, input_nb, GGML_MAX_DIMS); - acl_tensor_ptr acl_input_roll_reshape_tensor = - ggml_cann_create_tensor(input_roll_buffer, ggml_cann_type_mapping(src0->type), ggml_type_size(src0->type), - src0->ne, input_nb, GGML_MAX_DIMS); - - aclnn_mul(ctx, acl_input_roll_reshape_tensor.get(), acl_minus_one_tensor.get(), - acl_input_roll_mul_scale_tensor.get()); - - // output - void * output_fp32_buffer; - if (src0->type == GGML_TYPE_F32) { - aclnn_mul(ctx, acl_src.get(), acl_cos_reshape_tensor.get()); - aclnn_mul(ctx, acl_input_roll_mul_scale_tensor.get(), acl_sin_reshape_tensor.get()); - aclnn_add(ctx, acl_src.get(), acl_input_roll_mul_scale_tensor.get(), acl_dst.get()); - // TODO: ne0 != n_dims in mode2 - } else if (src0->type == GGML_TYPE_F16) { - size_t input_fp32_nb[GGML_MAX_DIMS]; - input_fp32_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - input_fp32_nb[i] = input_fp32_nb[i - 1] * dst->ne[i - 1]; - } - ggml_cann_pool_alloc fp32_allocator1(ctx.pool(), ggml_nelements(dst) * sizeof(float)); - void * input_fp32_buffer1 = fp32_allocator1.get(); - acl_tensor_ptr input_fp32_tensor1 = ggml_cann_create_tensor(input_fp32_buffer1, ACL_FLOAT, sizeof(float), - dst->ne, input_fp32_nb, GGML_MAX_DIMS); - ggml_cann_pool_alloc fp32_allocator2(ctx.pool(), ggml_nelements(dst) * sizeof(float)); - void * input_fp32_buffer2 = fp32_allocator2.get(); - acl_tensor_ptr input_fp32_tensor2 = ggml_cann_create_tensor(input_fp32_buffer2, ACL_FLOAT, sizeof(float), - dst->ne, input_fp32_nb, GGML_MAX_DIMS); - - ggml_cann_pool_alloc fp32_allocator(ctx.pool(), ggml_nelements(dst) * sizeof(float)); - output_fp32_buffer = fp32_allocator.get(); - acl_tensor_ptr output_fp32_tensor = ggml_cann_create_tensor(output_fp32_buffer, ACL_FLOAT, sizeof(float), - dst->ne, input_fp32_nb, GGML_MAX_DIMS); - aclnn_mul(ctx, acl_src.get(), acl_cos_reshape_tensor.get(), input_fp32_tensor1.get()); - aclnn_mul(ctx, acl_input_roll_mul_scale_tensor.get(), acl_sin_reshape_tensor.get(), input_fp32_tensor2.get()); - aclnn_add(ctx, input_fp32_tensor1.get(), input_fp32_tensor2.get(), output_fp32_tensor.get()); - aclnn_cast(ctx, output_fp32_tensor.get(), acl_dst.get(), ACL_FLOAT16); - } - return; -#endif - int64_t acl_mode = is_neox ? 0 : 1; - - // Pre-define head and tail dimensions for reuse - int64_t head_ne[GGML_MAX_DIMS] = { rope_dims, ne01, ne02, ne03 }; - int64_t tail_ne[GGML_MAX_DIMS] = { tail_dims, ne01, ne02, ne03 }; - - // Step 1: Prepare trans tensors for F16 type conversion to F32 if needed - bool src_dst_need_trans = false; - ggml_cann_pool_alloc src_trans_allocator(ctx.pool()); - ggml_cann_pool_alloc dst_trans_allocator(ctx.pool()); - acl_tensor_ptr acl_src_trans_tensor; - acl_tensor_ptr acl_dst_trans_tensor; - void * src_trans_buffer = nullptr; - void * dst_trans_buffer = nullptr; - size_t src_dst_trans_nb[GGML_MAX_DIMS]; - if (src0->type == GGML_TYPE_F16) { - src_dst_need_trans = true; - src_trans_buffer = src_trans_allocator.alloc(ggml_nelements(src0) * sizeof(float)); - dst_trans_buffer = dst_trans_allocator.alloc(ggml_nelements(dst) * sizeof(float)); - - src_dst_trans_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - src_dst_trans_nb[i] = src_dst_trans_nb[i - 1] * src0->ne[i - 1]; - } - acl_src_trans_tensor = ggml_cann_create_tensor(src_trans_buffer, ACL_FLOAT, sizeof(float), src0->ne, - src_dst_trans_nb, GGML_MAX_DIMS); - acl_dst_trans_tensor = ggml_cann_create_tensor(dst_trans_buffer, ACL_FLOAT, sizeof(float), dst->ne, - src_dst_trans_nb, GGML_MAX_DIMS); - aclnn_cast(ctx, acl_src.get(), acl_src_trans_tensor.get(), ACL_FLOAT); - } - - // Step 2: Prepare head tensors for tail splitting if needed - acl_tensor_ptr acl_src_head; - acl_tensor_ptr acl_dst_head; - if (has_tail) { - // Create head views for RotaryPositionEmbedding (only first rope_dims dimensions) - // RotaryPositionEmbedding requires contiguous dst tensor, so we use a temporary buffer - if (src_dst_need_trans) { - // Use F32 trans tensor strides - acl_src_head = ggml_cann_create_tensor((char *) src_trans_buffer, ACL_FLOAT, sizeof(float), head_ne, - src_dst_trans_nb, GGML_MAX_DIMS); - } else { - // Use original F32 tensor strides - acl_src_head = ggml_cann_create_tensor((char *) src0->data, ACL_FLOAT, sizeof(float), head_ne, src0->nb, - GGML_MAX_DIMS); - } - - int64_t head_elements = rope_dims * ne01 * ne02 * ne03; - ggml_cann_pool_alloc dst_head_contiguous_allocator(ctx.pool(), head_elements * sizeof(float)); - void * dst_head_contiguous_buffer = dst_head_contiguous_allocator.get(); - - size_t head_contiguous_nb[GGML_MAX_DIMS]; - head_contiguous_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - head_contiguous_nb[i] = head_contiguous_nb[i - 1] * head_ne[i - 1]; - } - acl_dst_head = ggml_cann_create_tensor(dst_head_contiguous_buffer, ACL_FLOAT, sizeof(float), head_ne, - head_contiguous_nb, GGML_MAX_DIMS); - } - - // Step 3: Execute RotaryPositionEmbedding - if (has_tail) { - // Rotate only the head portion (first rope_dims dimensions) - GGML_CANN_CALL_ACLNN_OP(ctx, RotaryPositionEmbedding, acl_src_head.get(), acl_cos_reshape_tensor.get(), - acl_sin_reshape_tensor.get(), acl_mode, acl_dst_head.get()); - - // Copy head result from contiguous buffer back to destination tensor - if (src_dst_need_trans) { - acl_tensor_ptr acl_dst_head_target = ggml_cann_create_tensor( - (char *) dst_trans_buffer, ACL_FLOAT, sizeof(float), head_ne, src_dst_trans_nb, GGML_MAX_DIMS); - cann_copy(ctx, acl_dst_head.get(), acl_dst_head_target.get()); - } else { - acl_tensor_ptr acl_dst_head_target = - ggml_cann_create_tensor((char *) dst->data, ACL_FLOAT, sizeof(float), head_ne, dst->nb, GGML_MAX_DIMS); - cann_copy(ctx, acl_dst_head.get(), acl_dst_head_target.get()); - } - } else if (src_dst_need_trans) { - // Rotate full tensor (no tail), using trans tensors - GGML_CANN_CALL_ACLNN_OP(ctx, RotaryPositionEmbedding, acl_src_trans_tensor.get(), acl_cos_reshape_tensor.get(), - acl_sin_reshape_tensor.get(), acl_mode, acl_dst_trans_tensor.get()); - } else if (src0->data == dst->data && !ggml_is_contiguous(src0)) { - // In-place on non-contiguous tensor: RotaryPositionEmbedding cannot safely - // read and write the same non-contiguous buffer. Use contiguous temporaries. - size_t contiguous_nb[GGML_MAX_DIMS]; - contiguous_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - contiguous_nb[i] = contiguous_nb[i - 1] * src0->ne[i - 1]; - } - int64_t total_elements = ggml_nelements(src0); - ggml_cann_pool_alloc inplace_src_alloc(ctx.pool(), total_elements * sizeof(float)); - ggml_cann_pool_alloc inplace_dst_alloc(ctx.pool(), total_elements * sizeof(float)); - - acl_tensor_ptr acl_src_contig = ggml_cann_create_tensor(inplace_src_alloc.get(), ACL_FLOAT, sizeof(float), - src0->ne, contiguous_nb, GGML_MAX_DIMS); - acl_tensor_ptr acl_dst_contig = ggml_cann_create_tensor(inplace_dst_alloc.get(), ACL_FLOAT, sizeof(float), - dst->ne, contiguous_nb, GGML_MAX_DIMS); - - cann_copy(ctx, acl_src.get(), acl_src_contig.get()); - GGML_CANN_CALL_ACLNN_OP(ctx, RotaryPositionEmbedding, acl_src_contig.get(), acl_cos_reshape_tensor.get(), - acl_sin_reshape_tensor.get(), acl_mode, acl_dst_contig.get()); - cann_copy(ctx, acl_dst_contig.get(), acl_dst.get()); - } else { - // Rotate full tensor (no tail), using original tensors - GGML_CANN_CALL_ACLNN_OP(ctx, RotaryPositionEmbedding, acl_src.get(), acl_cos_reshape_tensor.get(), - acl_sin_reshape_tensor.get(), acl_mode, acl_dst.get()); - } - - // Step 4: Copy unrotated tail portion from source to destination - if (has_tail) { - size_t src_tail_offset; - size_t dst_tail_offset; - - auto copy_tail_device = [&](void * src_ptr, void * dst_ptr, aclDataType dtype, size_t elem_size, - size_t * nb_src_arr, size_t * nb_dst_arr) { - acl_tensor_ptr acl_src_tail = - ggml_cann_create_tensor(src_ptr, dtype, elem_size, tail_ne, nb_src_arr, GGML_MAX_DIMS); - acl_tensor_ptr acl_dst_tail = - ggml_cann_create_tensor(dst_ptr, dtype, elem_size, tail_ne, nb_dst_arr, GGML_MAX_DIMS); - cann_copy(ctx, acl_src_tail.get(), acl_dst_tail.get()); - }; - - if (src_dst_need_trans) { - // Use F32 trans tensor strides and offsets - src_tail_offset = rope_dims * src_dst_trans_nb[0]; - dst_tail_offset = rope_dims * src_dst_trans_nb[0]; - copy_tail_device((char *) src_trans_buffer + src_tail_offset, (char *) dst_trans_buffer + dst_tail_offset, - ACL_FLOAT, sizeof(float), src_dst_trans_nb, src_dst_trans_nb); - } else { - // Use original tensor strides and offsets - src_tail_offset = rope_dims * nb00; - dst_tail_offset = rope_dims * nb0; - copy_tail_device((char *) src0->data + src_tail_offset, (char *) dst->data + dst_tail_offset, - ggml_cann_type_mapping(dst->type), ggml_element_size(dst), src0->nb, dst->nb); - } - } - - // Step 5: Cast back to F16 if needed - if (src_dst_need_trans) { - aclnn_cast(ctx, acl_dst_trans_tensor.get(), acl_dst.get(), ACL_FLOAT16); - } -} - -void ggml_cann_rope_cache_preload(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - - float freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow; - int sections[4]; - const int n_dims = ((int32_t *) dst->op_params)[1]; - const int mode = ((int32_t *) dst->op_params)[2]; - const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; - - GGML_TENSOR_UNARY_OP_LOCALS - - memcpy(&freq_base, (int32_t *) dst->op_params + 5, sizeof(float)); - memcpy(&freq_scale, (int32_t *) dst->op_params + 6, sizeof(float)); - memcpy(&ext_factor, (int32_t *) dst->op_params + 7, sizeof(float)); - memcpy(&attn_factor, (int32_t *) dst->op_params + 8, sizeof(float)); - memcpy(&beta_fast, (int32_t *) dst->op_params + 9, sizeof(float)); - memcpy(&beta_slow, (int32_t *) dst->op_params + 10, sizeof(float)); - memcpy(§ions, (int32_t *) dst->op_params + 11, sizeof(int) * 4); - - const float theta_scale = powf(freq_base, -2.0f / n_dims); - - float corr_dims[2]; - ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); - - bool is_neox = mode & GGML_ROPE_TYPE_NEOX; - const bool is_imrope = mode == GGML_ROPE_TYPE_IMROPE; - const bool mrope_used = mode & GGML_ROPE_TYPE_MROPE; - const bool is_vision = mode == GGML_ROPE_TYPE_VISION; - - if (is_imrope || mrope_used) { - is_neox = true; - } - - int64_t rope_dims = n_dims; - if (is_vision) { - rope_dims = src0->ne[0]; - } - - // Run the full cache init on the non-captured stream. This performs all - // host-to-device memcpy, aclrtMalloc/Free, and on-device computations - // so that the memory pool is warmed up and cache metadata is populated. - aclnn_rope_cache_init(ctx, dst, corr_dims, ext_factor, theta_scale, freq_scale, attn_factor, is_neox, sections, - mrope_used, is_imrope, is_vision, rope_dims); - - // Reset `cached` so that during graph capture the on-device computations - // (sin/cos, position multiply, repeat, etc.) still execute and get recorded - // into the captured graph. The cache metadata (theta_scale_length, - // theta_scale, sections, position_length, etc.) remains set, which causes - // all host-to-device copy and malloc/free branches to be skipped. - ctx.rope_cache.cached = false; -} - -void ggml_cann_argmax(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst, dst->ne, dst->nb, 3); - - GGML_CANN_CALL_ACLNN_OP(ctx, ArgMax, acl_src.get(), 3, false, acl_dst.get()); -} - -void ggml_cann_conv_transpose_1d(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - ggml_tensor * src1 = dst->src[1]; - - // stride - int64_t s0 = ((const int32_t *) (dst->op_params))[0]; - - acl_tensor_ptr acl_input = ggml_cann_create_tensor(src1, src1->ne, src1->nb, 3, ACL_FORMAT_NCL); - acl_tensor_ptr acl_weight = ggml_cann_create_tensor(src0, src0->ne, src0->nb, 3, ACL_FORMAT_NCL); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst, dst->ne, dst->nb, 3, ACL_FORMAT_NCL); - - // get base information of input and kernel - int64_t input_len = *(src1->ne); - int64_t dst_len = *(dst->ne); - int64_t kernel_size = *(src0->ne); - - // set the max kernel size for each conv - int64_t max_kernel_size = 255; - - // compute the partition of kernel - int64_t part_num = 1; - part_num = (kernel_size + max_kernel_size - 1) / max_kernel_size; - - int64_t strideVal[1]; - strideVal[0] = s0; - acl_int_array_ptr stride = ggml_cann_create_int_array(strideVal, 1); - int64_t paddingVal[] = { 0 }; - acl_int_array_ptr padding = ggml_cann_create_int_array(paddingVal, 1); - int64_t dilationVal[] = { 1 }; - acl_int_array_ptr dilation = ggml_cann_create_int_array(dilationVal, 1); - bool transposed = true; - int64_t groups = 1; - int8_t cubeMathType = 0; - -#ifdef ASCEND_310P - cubeMathType = 1; -#endif - - auto weight_type = ggml_cann_type_mapping(src0->type); - auto dst_type = ggml_cann_type_mapping(dst->type); - - // slice the kernel to make each conv available - int64_t slice_dim = -1; - int64_t slice_start = 0; - int64_t slice_end = max_kernel_size; - int64_t slice_step = 1; - int64_t interval = max_kernel_size; - - int64_t left_pad_len = dilationVal[0] * (max_kernel_size - 1) + 1 - 2 * paddingVal[0]; - int64_t right_pad_len = 0; - - acl_scalar_ptr alpha = nullptr; - float alphaValue = 1.0; - alpha = ggml_cann_create_scalar(&alphaValue, aclDataType::ACL_FLOAT); - - // set zero to destination - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceZero, acl_dst.get()); - - for (int k = 0; k < part_num; k++) { - // create part kernel tensor and slice from big kernel - slice_start = max_kernel_size * k; - if (k == part_num - 1) { - slice_end = kernel_size; - interval = kernel_size - max_kernel_size * k; - } else { - slice_end = max_kernel_size * (k + 1); - } - - int64_t part_ne[4]; - for (int i = 0; i < 4; i++) { - part_ne[i] = *(src0->ne + i); - } - part_ne[0] = interval; - - size_t part_nb[4]; - part_nb[0] = sizeof(weight_type); - for (int i = 1; i < 4; i++) { - part_nb[i] = part_nb[i - 1] * part_ne[i - 1]; - } - - ggml_cann_pool_alloc part_kernel_allocator; - part_kernel_allocator.alloc(ctx.pool(), part_nb[3]); - void * part_kernel_buf = part_kernel_allocator.get(); - - acl_tensor_ptr part_kernel = ggml_cann_create_tensor(part_kernel_buf, weight_type, ggml_element_size(src0), - part_ne, part_nb, 3, ACL_FORMAT_NCL); - - GGML_CANN_CALL_ACLNN_OP(ctx, Slice, acl_weight.get(), slice_dim, slice_start, slice_end, slice_step, - part_kernel.get()); - - // create the part conv result tensor - int64_t part_dst_ne[4]; - for (int i = 0; i < 4; i++) { - part_dst_ne[i] = *(dst->ne + i); - } - part_dst_ne[0] = (input_len - 1) * strideVal[0] - 2 * paddingVal[0] + dilationVal[0] * (part_ne[0] - 1) + 1; - - size_t part_dst_nb[4]; - part_dst_nb[0] = sizeof(weight_type); - for (int i = 1; i < 4; i++) { - part_dst_nb[i] = part_dst_nb[i - 1] * part_dst_ne[i - 1]; - } - ggml_cann_pool_alloc part_dst_allocator; - part_dst_allocator.alloc(ctx.pool(), part_dst_nb[3]); - void * part_dst_buf = part_dst_allocator.get(); - - acl_tensor_ptr acl_part_dst = ggml_cann_create_tensor(part_dst_buf, dst_type, ggml_element_size(dst), - part_dst_ne, part_dst_nb, 3, ACL_FORMAT_NCL); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceZero, acl_part_dst.get()); - - // compute part conv transpose 1d - GGML_CANN_CALL_ACLNN_OP(ctx, Convolution, acl_input.get(), part_kernel.get(), nullptr, stride.get(), - padding.get(), dilation.get(), transposed, padding.get(), groups, acl_part_dst.get(), - cubeMathType); - - // compute the position of part result in final result - int64_t global_start = slice_start; - int64_t global_end = std::min((input_len - 1) * strideVal[0] + slice_end, dst_len); - - left_pad_len = global_start; - right_pad_len = dst_len - global_end; - - std::vector padDataVal = { left_pad_len, right_pad_len }; - acl_int_array_ptr padData = ggml_cann_create_int_array(padDataVal.data(), 2); - - acl_scalar_ptr pad_value = nullptr; - float pad_valueVal = 0.0; - pad_value = ggml_cann_create_scalar(&pad_valueVal, aclDataType::ACL_FLOAT); - - int64_t conv_result_ne[4]; - for (int i = 0; i < 4; i++) { - conv_result_ne[i] = *(dst->ne + i); - } - - size_t conv_result_nb[4]; - conv_result_nb[0] = sizeof(weight_type); - for (int i = 1; i < 4; i++) { - conv_result_nb[i] = conv_result_nb[i - 1] * conv_result_ne[i - 1]; - } - - ggml_cann_pool_alloc conv_result_allocator; - conv_result_allocator.alloc(ctx.pool(), conv_result_nb[3]); - void * conv_result_buf = conv_result_allocator.get(); - - acl_tensor_ptr conv_result = ggml_cann_create_tensor(conv_result_buf, dst_type, ggml_element_size(dst), - conv_result_ne, conv_result_nb, 3, ACL_FORMAT_NCL); - - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceZero, conv_result.get()); - GGML_CANN_CALL_ACLNN_OP(ctx, ConstantPadNd, acl_part_dst.get(), padData.get(), pad_value.get(), - conv_result.get()); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceAdd, acl_dst.get(), conv_result.get(), alpha.get()); - } -} - -void ggml_cann_elu(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - - acl_tensor_ptr acl_input = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - float alphaValue = 1.0f; - acl_scalar_ptr alpha = nullptr; - alpha = ggml_cann_create_scalar(&alphaValue, aclDataType::ACL_FLOAT); - - GGML_CANN_CALL_ACLNN_OP(ctx, Elu, acl_input.get(), alpha.get(), alpha.get(), alpha.get(), acl_dst.get()); -} - -void ggml_cann_mean(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - int64_t reduceDimValue[] = { 3 }; - acl_int_array_ptr reduceDim = ggml_cann_create_int_array(reduceDimValue, 1); - bool keepDim = true; - - GGML_CANN_CALL_ACLNN_OP(ctx, Mean, acl_src.get(), reduceDim.get(), keepDim, ACL_FLOAT, acl_dst.get()); -} - -void ggml_cann_pad_reflect_1d(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - int32_t * opts = (int32_t *) dst->op_params; - int64_t paddingsArray[2] = { opts[0], opts[1] }; - acl_int_array_ptr paddings = ggml_cann_create_int_array(paddingsArray, 2); - - // Collapsing ne[2]*ne[3] into a single batch dimension requires that dim3 - // is contiguous with respect to dim2 in both src and dst. - GGML_ASSERT(src0->nb[3] == src0->nb[2] * src0->ne[2]); - GGML_ASSERT(dst->nb[3] == dst->nb[2] * dst->ne[2]); - - int64_t src_ne_3d[3] = { src0->ne[0], src0->ne[1], src0->ne[2] * src0->ne[3] }; - int64_t dst_ne_3d[3] = { dst->ne[0], dst->ne[1], dst->ne[2] * dst->ne[3] }; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src0->data, ggml_cann_type_mapping(src0->type), - ggml_element_size(src0), src_ne_3d, src0->nb, 3); - - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst->data, ggml_cann_type_mapping(dst->type), - ggml_element_size(dst), dst_ne_3d, dst->nb, 3); - - GGML_CANN_CALL_ACLNN_OP(ctx, ReflectionPad1d, acl_src.get(), paddings.get(), acl_dst.get()); -} - -void ggml_cann_count_equal(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - ggml_tensor * src1 = dst->src[1]; - - // Write element-wise equality (0 or 1) into a temporary buffer to avoid - // modifying src0 in-place. Use the same type as src0 so ReduceSum can - // consume it directly without a type cast. - ggml_cann_pool_alloc eq_alloc(ctx.pool(), ggml_nelements(src0) * ggml_element_size(src0)); - size_t eq_nb[GGML_MAX_DIMS]; - eq_nb[0] = ggml_element_size(src0); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - eq_nb[i] = eq_nb[i - 1] * src0->ne[i - 1]; - } - acl_tensor_ptr acl_eq = ggml_cann_create_tensor( - eq_alloc.get(), ggml_cann_type_mapping(src0->type), ggml_element_size(src0), - src0->ne, eq_nb, GGML_MAX_DIMS); - - acl_tensor_ptr acl_self = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_other = ggml_cann_create_tensor(src1); - GGML_CANN_CALL_ACLNN_OP(ctx, EqTensor, acl_self.get(), acl_other.get(), acl_eq.get()); - - // Sum the 0/1 values into dst. - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - int64_t dims[4] = { 0, 1, 2, 3 }; - acl_int_array_ptr dims_arr = ggml_cann_create_int_array(dims, 4); - GGML_CANN_CALL_ACLNN_OP(ctx, ReduceSum, acl_eq.get(), dims_arr.get(), true, - ggml_cann_type_mapping(dst->type), acl_dst.get()); -} - -void ggml_cann_step(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - float alphaValue = 0.0f; - acl_scalar_ptr alpha = nullptr; - alpha = ggml_cann_create_scalar(&alphaValue, aclDataType::ACL_FLOAT); - - GGML_CANN_CALL_ACLNN_OP(ctx, GtScalar, acl_src.get(), alpha.get(), acl_dst.get()); -} - -void ggml_cann_softplus(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src0); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - float beta_val = 1.0f; - float threshold_val = 20.0f; - acl_scalar_ptr beta = ggml_cann_create_scalar(&beta_val, ACL_FLOAT); - acl_scalar_ptr threshold = ggml_cann_create_scalar(&threshold_val, ACL_FLOAT); - - GGML_CANN_CALL_ACLNN_OP(ctx, Softplus, acl_src.get(), beta.get(), threshold.get(), acl_dst.get()); -} - -void ggml_cann_geglu_quick(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - auto gelu_quick_fn = [](ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst) { - GGML_CANN_CALL_ACLNN_OP(ctx, GeluV2, acl_src, 0, acl_dst); - }; - ggml_cann_op_unary_gated(gelu_quick_fn, ctx, dst); -} - -/** - * @brief Performs expert-specific matrix multiplication (MoE) with - * floating-point precision using the CANN backend. - * - * This function executes a matrix multiplication operation tailored for - * Mixture of Experts (MoE) models, where the input tensor is multiplied - * with expert-specific weight matrices. It uses the CANN backend for - * efficient computation and stores the result in the destination tensor `dst`. - * The operation may leverage identity-based optimizations or routing masks - * as part of sparse expert selection. - * - * @param ctx The context for executing CANN backend operations. - * @param dst The destination tensor where the MoE multiplication result - * will be stored. - * - * @note This function assumes floating-point data types and is designed for - * MoE architectures, possibly involving sparse expert routing. - */ -static void ggml_cann_mul_mat_id_fp(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - //dst [M, K, N, 1] - ggml_tensor * src0 = dst->src[0]; //src0 [D, M, A, 1] -> [D, M, K, 1] - ggml_tensor * src1 = dst->src[1]; //src1 [D, B, N, 1], B = K or B = 1 -> [D, 1, K, 1] - ggml_tensor * ids = dst->src[2]; //ids [K, N] - - GGML_ASSERT(src0->ne[3] == 1); - GGML_ASSERT(src1->ne[3] == 1); - GGML_ASSERT(dst->ne[3] == 1); - - int64_t batch = src1->ne[2]; - GGML_ASSERT(batch == ids->ne[1]); - - ggml_cann_pool_alloc export_allocator(ctx.pool(), src0->ne[0] * src0->ne[1] * ids->ne[0] * ggml_element_size(src0)); - void * export_ptr = export_allocator.get(); - for (int64_t i = 0; i < batch; i++) { - acl_tensor_ptr select_index = ggml_cann_create_tensor(ids, ids->ne, ids->nb, 1, ACL_FORMAT_ND, i * ids->nb[1]); - acl_tensor_ptr export_weight = ggml_cann_create_tensor(src0, src0->ne, src0->nb, 3); - - int64_t select_export_ne[] = { src0->ne[0], src0->ne[1], ids->ne[0] }; - size_t select_export_nb[3]; - select_export_nb[0] = src0->nb[0]; - for (int k = 1; k < 3; k++) { - select_export_nb[k] = select_export_nb[k - 1] * select_export_ne[k - 1]; - } - - acl_tensor_ptr select_export = - ggml_cann_create_tensor(export_ptr, ggml_cann_type_mapping(src0->type), ggml_element_size(src0), - select_export_ne, select_export_nb, 3); - GGML_CANN_CALL_ACLNN_OP(ctx, IndexSelect, export_weight.get(), 0, select_index.get(), select_export.get()); - - int64_t select_transpose_ne[] = { select_export_ne[1], select_export_ne[0], select_export_ne[2] }; - size_t select_transpose_nb[] = { select_export_nb[1], select_export_nb[0], select_export_nb[2] }; - acl_tensor_ptr select_export_transpose = - ggml_cann_create_tensor(export_ptr, ggml_cann_type_mapping(src0->type), ggml_element_size(src0), - select_transpose_ne, select_transpose_nb, 3); - - int64_t active_tensor_ne[] = { src1->ne[0], 1, src1->ne[1] }; - size_t active_tensor_nb[] = { src1->nb[0], src1->nb[1], src1->nb[1] }; - acl_tensor_ptr active_tensor = - ggml_cann_create_tensor(src1, active_tensor_ne, active_tensor_nb, 3, ACL_FORMAT_ND, i * src1->nb[2]); - - int64_t dst_ne[] = { dst->ne[0], 1, dst->ne[1] }; - size_t dst_nb[] = { dst->nb[0], dst->nb[1], dst->nb[1] }; - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst, dst_ne, dst_nb, 3, ACL_FORMAT_ND, i * dst->nb[2]); - - GGML_CANN_CALL_ACLNN_OP(ctx, BatchMatMul, active_tensor.get(), select_export_transpose.get(), acl_dst.get(), 2); - } -} - -/** - * @brief Performs quantized matrix multiplication for Mixture of Experts (MoE) - * models using the CANN backend. - * - * This function implements MUL_MAT_ID operation for quantized weight matrices - * (Q4_0 and Q8_0 formats). It selects expert-specific weight matrices based on - * the provided expert indices, and computes matrix multiplication using CANN's - * WeightQuantBatchMatmulV2 operator. - * - * The function performs the following steps: - * 1. Converts input/output tensors to F16 format if necessary - * 2. Uses IndexSelect to extract expert-specific weights and scales based on indices - * 3. Performs quantized matrix multiplication for each expert using WeightQuantBatchMatmulV2 - * 4. Converts output back to the target type if needed - * - * Tensor shapes: - * - dst: [M, K, N, 1] - output tensor - * - src0: [D, M, A, 1] - quantized weight matrices (Q4_0 or Q8_0) - * - src1: [D, B, N, 1] - input activations (B = K for per-expert input, or B = 1 for broadcast) - * - ids: [K, N] - expert indices for routing - * - * @param ctx The CANN backend context for operation execution. - * @param dst The destination tensor where the multiplication result will be stored. - * - * @note Only Q4_0 and Q8_0 quantization formats are supported. - * @note The function handles automatic type conversion to/from F16 as needed by the hardware. - */ -static void ggml_cann_mul_mat_id_quant(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - // dst: [M, K, N, 1] - // src0: [D, M, A, 1] - quantized weights - // src1: [D, B, N, 1] - input activations, B = K or B = 1 - // ids: [K, N] - expert indices - ggml_tensor * src0 = dst->src[0]; - ggml_tensor * src1 = dst->src[1]; - ggml_tensor * ids = dst->src[2]; - - GGML_ASSERT(src0->ne[3] == 1); - GGML_ASSERT(src1->ne[3] == 1); - GGML_ASSERT(dst->ne[3] == 1); - GGML_ASSERT(src1->ne[2] == ids->ne[1]); - - const int64_t n_batches = ids->ne[1]; - const int64_t n_select_experts = ids->ne[0]; - const enum ggml_type type = src0->type; - - const int32_t group_size = QK8_0; // Both Q4_0 and Q8_0 use group size of 32 - GGML_ASSERT(group_size == QK4_0); - - // Calculate element size for quantized weights - const float weight_elem_size = - (type == GGML_TYPE_Q4_0) ? 0.5f : - (type == GGML_TYPE_Q8_0) ? 1.0f : - (GGML_ABORT("MUL_MAT_ID only supports Q4_0 and Q8_0"), 0.0f); - - // Calculate scale offset in memory - const size_t weight_size = src0->ne[0] * src0->ne[1] * src0->ne[2] * weight_elem_size; - const size_t scale_elem_size = sizeof(uint16_t); - char * scale_data = (char *) src0->data + weight_size; - - // Allocate buffers for selected expert weights and scales - const size_t selected_weight_size = src0->ne[0] * src0->ne[1] * n_select_experts * weight_elem_size; - ggml_cann_pool_alloc selected_weight_alloc(ctx.pool(), selected_weight_size); - void * selected_weight_buffer = selected_weight_alloc.get(); - - const size_t selected_scale_size = (src0->ne[0] / group_size) * src0->ne[1] * n_select_experts * scale_elem_size; - ggml_cann_pool_alloc selected_scale_alloc(ctx.pool(), selected_scale_size); - void * selected_scale_buffer = selected_scale_alloc.get(); - - // Helper lambda to allocate and cast tensor to F16 if needed - constexpr size_t f16_elem_size = sizeof(uint16_t); - auto prepare_f16_buffer = [&](ggml_tensor * tensor, ggml_cann_pool_alloc & allocator, - bool need_cast = false) -> void * { - if (tensor->type == GGML_TYPE_F16) { - return tensor->data; - } - - size_t total_size = f16_elem_size; - for (int i = 0; i < GGML_MAX_DIMS; i++) { - total_size *= tensor->ne[i]; - } - void * buffer = allocator.alloc(total_size); - - if (need_cast == false) { - return buffer; - } - - int64_t ne[GGML_MAX_DIMS]; - size_t nb[GGML_MAX_DIMS] = { f16_elem_size }; - for (int i = 0; i < GGML_MAX_DIMS; i++) { - ne[i] = tensor->ne[i]; - if (i > 0) { - nb[i] = nb[i - 1] * ne[i - 1]; - } - } - - acl_tensor_ptr src_tensor = ggml_cann_create_tensor(tensor); - acl_tensor_ptr f16_tensor = ggml_cann_create_tensor(buffer, ACL_FLOAT16, f16_elem_size, ne, nb, GGML_MAX_DIMS); - aclnn_cast(ctx, src_tensor.get(), f16_tensor.get(), ACL_FLOAT16); - - return buffer; - }; - - // Prepare input and output buffers - ggml_cann_pool_alloc input_alloc(ctx.pool()); - void * input_buffer = prepare_f16_buffer(src1, input_alloc, true); - - ggml_cann_pool_alloc output_alloc(ctx.pool()); - void * output_buffer = prepare_f16_buffer(dst, output_alloc, false); - - // Process each batch - for (int64_t batch_idx = 0; batch_idx < n_batches; batch_idx++) { - // Create index tensor for current batch - const size_t index_offset = batch_idx * ids->nb[1]; - acl_tensor_ptr batch_indices = ggml_cann_create_tensor(ids, ids->ne, ids->nb, 1, ACL_FORMAT_ND, index_offset); - - // Select quantized weights using expert indices - // Q4_0 stores 2 values per byte, Q8_0 stores 1 value per byte - const int64_t weight_d = (type == GGML_TYPE_Q4_0) ? src0->ne[0] / 2 : src0->ne[0]; - const int64_t weight_m = src0->ne[1]; - const int64_t weight_n_experts = src0->ne[2]; - - int64_t weight_ne[3] = { weight_d, weight_m, weight_n_experts }; - size_t weight_nb[3] = { sizeof(int8_t), weight_d * sizeof(int8_t), weight_d * weight_m * sizeof(int8_t) }; - - acl_tensor_ptr all_weights = - ggml_cann_create_tensor(src0->data, ACL_INT8, sizeof(int8_t), weight_ne, weight_nb, 3); - - int64_t selected_weight_ne[3] = { weight_d, weight_m, n_select_experts }; - size_t selected_weight_nb[3] = { sizeof(int8_t), weight_d * sizeof(int8_t), - weight_d * weight_m * sizeof(int8_t) }; - - acl_tensor_ptr selected_weights = ggml_cann_create_tensor(selected_weight_buffer, ACL_INT8, sizeof(int8_t), - selected_weight_ne, selected_weight_nb, 3); - - GGML_CANN_CALL_ACLNN_OP(ctx, IndexSelect, all_weights.get(), 0, batch_indices.get(), selected_weights.get()); - - // Select scales using the same expert indices - const int64_t scale_d = src0->ne[0] / group_size; - int64_t scale_ne[3] = { scale_d, weight_m, weight_n_experts }; - size_t scale_nb[3] = { scale_elem_size, scale_d * scale_elem_size, scale_d * weight_m * scale_elem_size }; - - acl_tensor_ptr all_scales = - ggml_cann_create_tensor(scale_data, ACL_FLOAT16, scale_elem_size, scale_ne, scale_nb, 3); - - int64_t selected_scale_ne[3] = { scale_d, weight_m, n_select_experts }; - size_t selected_scale_nb[3] = { scale_elem_size, scale_d * scale_elem_size, - scale_d * weight_m * scale_elem_size }; - - acl_tensor_ptr selected_scales = ggml_cann_create_tensor(selected_scale_buffer, ACL_FLOAT16, scale_elem_size, - selected_scale_ne, selected_scale_nb, 3); - - GGML_CANN_CALL_ACLNN_OP(ctx, IndexSelect, all_scales.get(), 0, batch_indices.get(), selected_scales.get()); - - // Process each expert for current batch - // IndexSelect output layout: [D, M, K] in contiguous format - // WeightQuantBatchMatmulV2 expects: [M, D] with row-major stride - for (int64_t expert_idx = 0; expert_idx < n_select_experts; expert_idx++) { - // Determine input offset: broadcast if src1->ne[1]==1, otherwise use per-expert input - const size_t input_offset = - (batch_idx * src1->ne[1] + (src1->ne[1] == 1 ? 0 : expert_idx)) * src1->ne[0] * f16_elem_size; - const size_t output_offset = (batch_idx * dst->ne[1] + expert_idx) * dst->ne[0] * f16_elem_size; - - // Create weight view for current expert: [D, M, K] -> [M, D] - int64_t weight_view_ne[2] = { weight_m, src0->ne[0] }; - float weight_view_nb[2] = { src0->ne[0] * weight_elem_size, weight_elem_size }; - const size_t weight_view_offset = expert_idx * selected_weight_nb[2]; - - acl_tensor_ptr weight_view = - ggml_cann_create_tensor(selected_weight_buffer, ggml_cann_type_mapping(type), weight_elem_size, - weight_view_ne, weight_view_nb, 2, ACL_FORMAT_ND, weight_view_offset); - - // Create scale view for current expert: [D, M, K] -> [M, D] - int64_t scale_view_ne[2] = { weight_m, scale_d }; - size_t scale_view_nb[2] = { selected_scale_nb[1], selected_scale_nb[0] }; - const size_t scale_view_offset = expert_idx * selected_scale_nb[2]; - - acl_tensor_ptr scale_view = - ggml_cann_create_tensor(selected_scale_buffer, ACL_FLOAT16, scale_elem_size, scale_view_ne, - scale_view_nb, 2, ACL_FORMAT_ND, scale_view_offset); - - // Create input activation tensor [D, 1] - int64_t input_ne[2] = { src1->ne[0], 1 }; - size_t input_nb[2] = { f16_elem_size, src1->ne[0] * f16_elem_size }; - - acl_tensor_ptr input_tensor = ggml_cann_create_tensor(input_buffer, ACL_FLOAT16, f16_elem_size, input_ne, - input_nb, 2, ACL_FORMAT_ND, input_offset); - - // Create output tensor [M, 1] - int64_t output_ne[2] = { dst->ne[0], 1 }; - size_t output_nb[2] = { f16_elem_size, dst->ne[0] * f16_elem_size }; - - acl_tensor_ptr output_tensor = ggml_cann_create_tensor(output_buffer, ACL_FLOAT16, f16_elem_size, output_ne, - output_nb, 2, ACL_FORMAT_ND, output_offset); - - // Perform quantized matrix multiplication - GGML_CANN_CALL_ACLNN_OP(ctx, WeightQuantBatchMatmulV2, input_tensor.get(), weight_view.get(), - scale_view.get(), nullptr, nullptr, nullptr, nullptr, group_size, - output_tensor.get()); - } - } - - // Cast output back to original type if we used a temporary F16 buffer - if (dst->type != GGML_TYPE_F16) { - int64_t ne[GGML_MAX_DIMS]; - size_t nb[GGML_MAX_DIMS] = { f16_elem_size }; - for (int i = 0; i < GGML_MAX_DIMS; i++) { - ne[i] = dst->ne[i]; - if (i > 0) { - nb[i] = nb[i - 1] * ne[i - 1]; - } - } - - acl_tensor_ptr f16_output = - ggml_cann_create_tensor(output_buffer, ACL_FLOAT16, f16_elem_size, ne, nb, GGML_MAX_DIMS); - acl_tensor_ptr dst_tensor = ggml_cann_create_tensor(dst); - - aclnn_cast(ctx, f16_output.get(), dst_tensor.get(), ggml_cann_type_mapping(dst->type)); - } -} - -void ggml_cann_mul_mat_id(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - const enum ggml_type type = dst->src[0]->type; - switch (type) { - case GGML_TYPE_F32: - case GGML_TYPE_F16: - ggml_cann_mul_mat_id_fp(ctx, dst); - break; - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q8_0: - ggml_cann_mul_mat_id_quant(ctx, dst); - break; - default: - GGML_ABORT("Unsupported type for mul_mat_id"); - break; - } -} - -void ggml_cann_flash_attn_ext(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; // q, fp32 | B, N, S, D (uncont) -> B, S, N, D (cont) - ggml_tensor * src1 = dst->src[1]; // k, fp16 | B, N, S, D (uncont) -> B, S, N, D (cont) - ggml_tensor * src2 = dst->src[2]; // v, fp16 | B, N, S, D (uncont) -> B, S, N, D (cont) - ggml_tensor * src3 = dst->src[3]; // mask, fp16 - - // B, N, S, D (uncont) -> B, S, N, D (cont) - int64_t src0_bsnd_ne[GGML_MAX_DIMS]; - memcpy(src0_bsnd_ne, src0->ne, GGML_MAX_DIMS * sizeof(int64_t)); - size_t src0_bsnd_nb[GGML_MAX_DIMS]; - memcpy(src0_bsnd_nb, src0->nb, GGML_MAX_DIMS * sizeof(size_t)); - int64_t src1_bsnd_ne[GGML_MAX_DIMS]; - memcpy(src1_bsnd_ne, src1->ne, GGML_MAX_DIMS * sizeof(int64_t)); - size_t src1_bsnd_nb[GGML_MAX_DIMS]; - memcpy(src1_bsnd_nb, src1->nb, GGML_MAX_DIMS * sizeof(size_t)); - int64_t src2_bsnd_ne[GGML_MAX_DIMS]; - memcpy(src2_bsnd_ne, src2->ne, GGML_MAX_DIMS * sizeof(int64_t)); - size_t src2_bsnd_nb[GGML_MAX_DIMS]; - memcpy(src2_bsnd_nb, src2->nb, GGML_MAX_DIMS * sizeof(size_t)); - - auto transpose12 = [](int64_t * ne, size_t * nb) { - int64_t ne_tmp = ne[1]; - size_t nb_tmp = nb[1]; - ne[1] = ne[2]; - nb[1] = nb[2]; - ne[2] = ne_tmp; - nb[2] = nb_tmp; - }; - - transpose12(src0_bsnd_ne, src0_bsnd_nb); - transpose12(src1_bsnd_ne, src1_bsnd_nb); - transpose12(src2_bsnd_ne, src2_bsnd_nb); - - float maxBias = 0.0f; - float scaleValue = 1.0f; - float logitSoftcap = 0.0f; - memcpy(&scaleValue, (float *) dst->op_params + 0, sizeof(float)); - memcpy(&maxBias, (float *) dst->op_params + 1, sizeof(float)); - memcpy(&logitSoftcap, (float *) dst->op_params + 2, sizeof(float)); - - if (logitSoftcap == 0.0f) { - size_t faElemSize = sizeof(uint16_t); - auto faDataType = ACL_FLOAT16; //ACL_BF16; - - acl_tensor_ptr acl_q_tensor = nullptr; - acl_tensor_ptr acl_k_tensor = nullptr; - acl_tensor_ptr acl_v_tensor = nullptr; - - // Step 1: cast the src0 (Query) to fp16 if needed - ggml_cann_pool_alloc src0_f16_allocator(ctx.pool()); - void * src0_f16_buffer = nullptr; - - if (ggml_cann_type_mapping(src0->type) != faDataType) { - acl_tensor_ptr acl_src0_f32_tensor = - ggml_cann_create_tensor(src0, src0_bsnd_ne, src0_bsnd_nb, GGML_MAX_DIMS); - src0_f16_buffer = src0_f16_allocator.alloc(ggml_nelements(src0) * faElemSize); - - int64_t * src0_f16_ne = src0_bsnd_ne; - size_t src0_f16_nb[GGML_MAX_DIMS]; - src0_f16_nb[0] = sizeof(uint16_t); - for (int i = 1; i < GGML_MAX_DIMS; ++i) { - src0_f16_nb[i] = src0_f16_nb[i - 1] * src0_f16_ne[i - 1]; - } - - acl_q_tensor = ggml_cann_create_tensor(src0_f16_buffer, faDataType, faElemSize, src0_f16_ne, src0_f16_nb, - GGML_MAX_DIMS); - aclnn_cast(ctx, acl_src0_f32_tensor.get(), acl_q_tensor.get(), faDataType); - } else { - acl_q_tensor = ggml_cann_create_tensor(src0, src0_bsnd_ne, src0_bsnd_nb, GGML_MAX_DIMS); - } - - // Step 2: create the acl tensors for src1 (Key), src2 (Value), - // and the direct output from FusedInferAttention - - acl_k_tensor = ggml_cann_create_tensor(src1, src1_bsnd_ne, src1_bsnd_nb, GGML_MAX_DIMS); - acl_v_tensor = ggml_cann_create_tensor(src2, src2_bsnd_ne, src2_bsnd_nb, GGML_MAX_DIMS); - - // Step 2.5: Pad Q, K, V along head dimension if D is not a multiple of 16 - // (required by FusedInferAttentionScoreV2) - const int64_t D = src0->ne[0]; - const int64_t D_padded = GGML_PAD(D, 16); - const bool needs_padding = (D != D_padded); - - ggml_cann_pool_alloc q_pad_allocator(ctx.pool()); - ggml_cann_pool_alloc k_pad_allocator(ctx.pool()); - ggml_cann_pool_alloc v_pad_allocator(ctx.pool()); - - if (needs_padding) { - int64_t paddings[] = { 0, D_padded - D, 0, 0, 0, 0, 0, 0 }; - - auto pad_fa_tensor = [&](acl_tensor_ptr & tensor, const int64_t * bsnd_ne, - ggml_cann_pool_alloc & allocator) { - int64_t pad_ne[GGML_MAX_DIMS] = { D_padded, bsnd_ne[1], bsnd_ne[2], bsnd_ne[3] }; - size_t pad_nb[GGML_MAX_DIMS]; - pad_nb[0] = faElemSize; - for (int i = 1; i < GGML_MAX_DIMS; ++i) { - pad_nb[i] = pad_nb[i - 1] * pad_ne[i - 1]; - } - int64_t nelements = pad_ne[0] * pad_ne[1] * pad_ne[2] * pad_ne[3]; - void * buffer = allocator.alloc(nelements * faElemSize); - acl_tensor_ptr padded = - ggml_cann_create_tensor(buffer, faDataType, faElemSize, pad_ne, pad_nb, GGML_MAX_DIMS); - aclnn_pad(ctx, tensor.get(), padded.get(), paddings); - tensor = std::move(padded); - }; - - pad_fa_tensor(acl_q_tensor, src0_bsnd_ne, q_pad_allocator); - pad_fa_tensor(acl_k_tensor, src1_bsnd_ne, k_pad_allocator); - pad_fa_tensor(acl_v_tensor, src2_bsnd_ne, v_pad_allocator); - - src0_bsnd_ne[0] = D_padded; - src1_bsnd_ne[0] = D_padded; - src2_bsnd_ne[0] = D_padded; - } - - // Step 3: create the PSEShift tensor if needed - // this tensor is considered as mask (f16) in the llama.cpp - acl_tensor_ptr bcast_pse_tensor; - ggml_cann_pool_alloc bcast_pse_allocator(ctx.pool()); - if (src3 != nullptr) { - // Construct the truncated pse tensor (common for prefill/decode) - int64_t trunc_pse_ne[GGML_MAX_DIMS] = { - src3->ne[0], // D - src0->ne[1], // S (number of Q tokens) - src3->ne[2], // mask N - src3->ne[3] // B - }; - size_t * trunc_pse_nb = src3->nb; - - acl_tensor_ptr acl_mask_f16_trunc_tensor = ggml_cann_create_tensor( - src3->data, ACL_FLOAT16, sizeof(uint16_t), trunc_pse_ne, trunc_pse_nb, GGML_MAX_DIMS); - - int64_t bcast_pse_ne[GGML_MAX_DIMS]; - size_t bcast_pse_nb[GGML_MAX_DIMS]; - bcast_pse_ne[0] = src3->ne[0]; // D - bcast_pse_ne[1] = src0->ne[1]; // S - bcast_pse_ne[2] = src0->ne[2]; // N (num_heads) - bcast_pse_ne[3] = src3->ne[3]; // B - if (maxBias == 0.0f) { - // When maxBias == 0.0f, use nb = 0 reduce once repeat (Qwen2) - // Construct the bcast tensor (simulate repeat on the head dimension using stride=0) - bcast_pse_nb[0] = sizeof(uint16_t); - bcast_pse_nb[1] = bcast_pse_nb[0] * bcast_pse_ne[0]; - bcast_pse_nb[2] = 0; // <---- the head dimension shares the same data - bcast_pse_nb[3] = src3->nb[3]; - - bcast_pse_tensor = ggml_cann_create_tensor(src3->data, ACL_FLOAT16, sizeof(uint16_t), bcast_pse_ne, - bcast_pse_nb, GGML_MAX_DIMS); - - } else { - bcast_pse_nb[0] = sizeof(uint16_t); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - bcast_pse_nb[i] = bcast_pse_nb[i - 1] * bcast_pse_ne[i - 1]; - } - - void * bcast_pse_buffer = - bcast_pse_allocator.alloc(ggml_nelements(src3) * src0->ne[2] * sizeof(uint16_t)); - - bcast_pse_tensor = ggml_cann_create_tensor(bcast_pse_buffer, ACL_FLOAT16, sizeof(uint16_t), - bcast_pse_ne, bcast_pse_nb, GGML_MAX_DIMS); - - int64_t repeats[] = { 1, src0->ne[2], 1, 1 }; - aclnn_repeat(ctx, acl_mask_f16_trunc_tensor.get(), bcast_pse_tensor.get(), repeats); - - // alibi - // Compute the slope if needed. Derived from ggml_cann_softmax(). - const int64_t n_heads = src0->ne[2]; - ggml_cann_pool_alloc slope_allocator(ctx.pool(), n_heads * sizeof(uint16_t)); - void * slope_buffer = slope_allocator.get(); - aclnn_get_slope(ctx, n_heads, slope_buffer, maxBias, GGML_TYPE_F16); - - int64_t slope_ne[] = { 1, 1, n_heads, 1 }; - size_t slope_nb[GGML_MAX_DIMS]; - slope_nb[0] = sizeof(uint16_t); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - slope_nb[i] = slope_nb[i - 1] * slope_ne[0]; - } - - acl_tensor_ptr slope_tensor = ggml_cann_create_tensor(slope_buffer, ACL_FLOAT16, sizeof(uint16_t), - slope_ne, slope_nb, GGML_MAX_DIMS); - GGML_CANN_CALL_ACLNN_OP(ctx, InplaceMul, bcast_pse_tensor.get(), slope_tensor.get()); - } - } - - // Step 4: set the inputs for FusedInferAttention. - acl_tensor_list_ptr acl_k_tensor_list = ggml_cann_create_tensor_list(acl_k_tensor); - acl_tensor_list_ptr acl_v_tensor_list = ggml_cann_create_tensor_list(acl_v_tensor); - - int64_t numHeads = src0->ne[2]; // N - int64_t numKeyValueHeads = src1->ne[2]; - // double scaleValue = 1 / sqrt(src0->ne[0]); // 1/sqrt(d) - int64_t preTokens = 65535; - int64_t nextTokens = 65535; - char layout[5] = { 'B', 'S', 'N', 'D', 0 }; - int64_t sparseMode = 0; - int64_t innerPrecise = (src0->ne[1] == 1) ? 0 : 2; - int64_t blockSize = 0; - int64_t antiquantMode = 0; - bool softmaxLseFlag = false; - int64_t keyAntiquantMode = 0; - int64_t valueAntiquantMode = 0; - - GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16); - acl_tensor_ptr fa_dst_tensor; - ggml_cann_pool_alloc out_f16_allocator(ctx.pool()); - if (dst->type == GGML_TYPE_F32 || needs_padding) { - int64_t * out_f16_ne = src0_bsnd_ne; - size_t out_f16_nb[GGML_MAX_DIMS]; - out_f16_nb[0] = faElemSize; - for (int i = 1; i < GGML_MAX_DIMS; ++i) { - out_f16_nb[i] = out_f16_nb[i - 1] * out_f16_ne[i - 1]; - } - int64_t out_nelements = out_f16_ne[0] * out_f16_ne[1] * out_f16_ne[2] * out_f16_ne[3]; - void * out_f16_buffer = out_f16_allocator.alloc(out_nelements * faElemSize); - - fa_dst_tensor = - ggml_cann_create_tensor(out_f16_buffer, faDataType, faElemSize, out_f16_ne, out_f16_nb, GGML_MAX_DIMS); - } else { - fa_dst_tensor = ggml_cann_create_tensor(dst); - } - - GGML_CANN_CALL_ACLNN_OP(ctx, FusedInferAttentionScoreV2, acl_q_tensor.get(), acl_k_tensor_list.get(), - acl_v_tensor_list.get(), // q, k, v - bcast_pse_tensor.get(), nullptr, // pse, mask - nullptr, nullptr, // actSeqLen, actSeqLenkv - nullptr, nullptr, // deqScale1, quantScale1 - nullptr, nullptr, nullptr, // deqScale2, quantScale2, quantOffset2 - nullptr, nullptr, // antiquantScale, antiquantOffset - nullptr, // blockTable - nullptr, nullptr, // qPadSize, kvPadSize - nullptr, nullptr, // kAntiquantScale, kAntiQuantOffset - nullptr, nullptr, // vAntiquantScale, vAntiQuantOffset - nullptr, nullptr, nullptr, // kSharedPrefix, vSharedPrefix, actSharedLen - numHeads, scaleValue, // heads, scaleValue - preTokens, nextTokens, // preTokens, nextTokens - layout, // inputLayout - numKeyValueHeads, // numKVHeads - sparseMode, innerPrecise, // sparseMode, innerPrecise - blockSize, antiquantMode, // blockSize, antiquantMode - softmaxLseFlag, // softmaxLseFlag - keyAntiquantMode, valueAntiquantMode, // keyAntiqMode, valueAntiqMode - fa_dst_tensor.get(), // attentionOut - nullptr // softmaxLse - ); - - // Step 6: post-processing — slice padded output and/or cast to f32 - if (needs_padding) { - ggml_cann_pool_alloc sliced_f16_allocator(ctx.pool()); - - if (dst->type == GGML_TYPE_F32) { - int64_t sliced_ne[GGML_MAX_DIMS] = { D, src0_bsnd_ne[1], src0_bsnd_ne[2], src0_bsnd_ne[3] }; - size_t sliced_nb[GGML_MAX_DIMS]; - sliced_nb[0] = faElemSize; - for (int i = 1; i < GGML_MAX_DIMS; ++i) { - sliced_nb[i] = sliced_nb[i - 1] * sliced_ne[i - 1]; - } - int64_t sliced_nelements = sliced_ne[0] * sliced_ne[1] * sliced_ne[2] * sliced_ne[3]; - void * sliced_buffer = sliced_f16_allocator.alloc(sliced_nelements * faElemSize); - acl_tensor_ptr sliced_f16_tensor = ggml_cann_create_tensor(sliced_buffer, faDataType, faElemSize, - sliced_ne, sliced_nb, GGML_MAX_DIMS); - - GGML_CANN_CALL_ACLNN_OP(ctx, Slice, fa_dst_tensor.get(), - (int64_t) -1, (int64_t) 0, D, (int64_t) 1, sliced_f16_tensor.get()); - - acl_tensor_ptr acl_dst_tensor = ggml_cann_create_tensor(dst); - aclnn_cast(ctx, sliced_f16_tensor.get(), acl_dst_tensor.get(), ggml_cann_type_mapping(dst->type)); - } else { - acl_tensor_ptr acl_dst_tensor = ggml_cann_create_tensor(dst); - GGML_CANN_CALL_ACLNN_OP(ctx, Slice, fa_dst_tensor.get(), - (int64_t) -1, (int64_t) 0, D, (int64_t) 1, acl_dst_tensor.get()); - } - } else if (dst->type == GGML_TYPE_F32) { - acl_tensor_ptr acl_dst_tensor = ggml_cann_create_tensor(dst); - aclnn_cast(ctx, fa_dst_tensor.get(), acl_dst_tensor.get(), ggml_cann_type_mapping(dst->type)); - } - } else { - GGML_ABORT("Function is not implemented."); - } -} - -static void ggml_cann_out_prod_fp(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; // weight [ne00=m, ne01=K, ne02, ne03] - ggml_tensor * src1 = dst->src[1]; // input [ne10=n, ne11=K, ne12, ne13] - GGML_TENSOR_BINARY_OP_LOCALS - - // dst[i,j] = sum_k src0[i,k] * src1[j,k] i.e. dst = src0 @ src1^T. - // - // ggml_cann_create_tensor reverses dimension order, so ACL sees: - // acl_src0 slice: ggml[m,K] -> ACL[K,m] - // acl_src1 slice: ggml[n,K] -> ACL[K,n] - // acl_dst slice: ggml[m,n] -> ACL[n,m] - // - // Build a transposed view of src1 by swapping ne[0]/ne[1]: - // src1_t: ggml[K,n] (swapped strides) -> ACL[n,K] - // - // Matmul(src1_t [n,K], src0 [K,m]) = [n,m] = acl_dst ✓ - // - // The outer batch loop is kept because src0 may have fewer batch slices than - // dst (ne02 <= ne2, ne03 <= ne3): this is a strided-broadcast not supported - // by standard CANN Matmul broadcasting. - - const aclDataType src0_acl_type = ggml_cann_type_mapping(src0->type); - const aclDataType src1_acl_type = ggml_cann_type_mapping(src1->type); - const aclDataType dst_acl_type = ggml_cann_type_mapping(dst->type); - const size_t src0_type_sz = ggml_type_size(src0->type); - const size_t src1_type_sz = ggml_type_size(src1->type); - const size_t dst_type_sz = ggml_type_size(dst->type); - - const int64_t dps2 = ne2 / ne02; - const int64_t dps3 = ne3 / ne03; - - for (int64_t i3 = 0; i3 < ne3; i3++) { - for (int64_t i2 = 0; i2 < ne2; i2++) { - const int64_t i02 = i2 / dps2; - const int64_t i03 = i3 / dps3; - - // src0 2D slice at [i02, i03]: ggml [m, K] -> ACL [K, m] - int64_t src0_ne[2] = { ne00, ne01 }; - size_t src0_nb[2] = { nb00, nb01 }; - acl_tensor_ptr acl_src0_s = ggml_cann_create_tensor( - (char *) src0->data + i02 * nb02 + i03 * nb03, - src0_acl_type, src0_type_sz, src0_ne, src0_nb, 2); - - // src1 transposed 2D slice at [i2, i3]: swap ne/nb -> ggml[K,n] -> ACL[n,K] - int64_t src1_t_ne[2] = { ne11, ne10 }; - size_t src1_t_nb[2] = { nb11, nb10 }; - acl_tensor_ptr acl_src1_t = ggml_cann_create_tensor( - (char *) src1->data + i2 * nb12 + i3 * nb13, - src1_acl_type, src1_type_sz, src1_t_ne, src1_t_nb, 2); - - // dst 2D slice at [i2, i3]: ggml [m, n] -> ACL [n, m] - int64_t dst_ne[2] = { ne0, ne1 }; - size_t dst_nb[2] = { nb0, nb1 }; - acl_tensor_ptr acl_dst_s = ggml_cann_create_tensor( - (char *) dst->data + i2 * nb2 + i3 * nb3, - dst_acl_type, dst_type_sz, dst_ne, dst_nb, 2); - - // Matmul(src1_t [n,K], src0 [K,m]) = [n,m] = acl_dst_s ✓ - GGML_CANN_CALL_ACLNN_OP(ctx, Matmul, - acl_src1_t.get(), acl_src0_s.get(), acl_dst_s.get(), (int8_t) 1); - } - } -} - -void ggml_cann_out_prod(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - - const enum ggml_type type = src0->type; - - switch (type) { - case GGML_TYPE_F32: - case GGML_TYPE_F16: - ggml_cann_out_prod_fp(ctx, dst); - break; - default: - GGML_ABORT("Unsupport type for GGML_OP_OUT_PROD"); - break; - } -} - -void ggml_cann_ssm_conv(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; // conv_x - ggml_tensor * src1 = dst->src[1]; // conv1d.weight - - // This op is currently defined only for F32 in ggml_cpu - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - - // Shapes follow ggml_compute_forward_ssm_conv_f32 - const int64_t nc = src1->ne[0]; // d_conv - const int64_t ncs = src0->ne[0]; // d_conv - 1 + n_t - const int64_t nr = src0->ne[1]; // d_inner - const int64_t n_s = src0->ne[2]; // n_seqs - - const int64_t n_t = dst->ne[1]; // tokens per sequence - - GGML_ASSERT(dst->ne[0] == nr); // dst: {d_inner, n_t, n_s} - GGML_ASSERT(src1->ne[1] == nr); // weight: {d_conv, d_inner} - GGML_ASSERT(ncs == nc - 1 + n_t); // conv_x: {d_conv - 1 + n_t, d_inner, n_s} - GGML_ASSERT(src0->nb[0] == sizeof(float)); - GGML_ASSERT(src1->nb[0] == sizeof(float)); - - // --- Build CANN tensors --- - - // 1) Input: conv_x as NCL - // - // src0->ne = { ncs, nr, n_s, 1 } // {L_in, C, N} - // Passing ACL_FORMAT_NCL here means: - // reversed dims -> [N, C, L_in] = [n_s, nr, ncs] - acl_tensor_ptr acl_x = ggml_cann_create_tensor(src0, src0->ne, src0->nb, 3, ACL_FORMAT_NCL); - - // 2) Weights: depthwise conv kernel, view src1 as {K, 1, C} - // - // src1 original: ne = { nc, nr, 1, 1 } // [K, C, 1, 1] - // we want a view: ne_w = { nc, 1, nr } // [K, 1, C] - // so that reversed dims -> [C, 1, K] which matches - // [out_channels, in_channels/groups, kernel_size] - int64_t w_ne[GGML_MAX_DIMS] = { nc, 1, nr, 1 }; // [K, 1 input ch. per group, C groups] - // Layout: src1 data is [K, C] with - // offset(k, c) = k*nb0 + c*nb1 - // We want offset_w(k, 0, c) = k*nb0 + c*nb1, - // so we can reuse nb0 and nb1, and set nb2 = nb1. - size_t w_nb[GGML_MAX_DIMS] = { src1->nb[0], src1->nb[1], src1->nb[1], src1->nb[3] }; // same as src1 - - acl_tensor_ptr acl_w = ggml_cann_create_tensor(src1->data, ggml_cann_type_mapping(src1->type), - ggml_type_size(src1->type), w_ne, w_nb, 3, ACL_FORMAT_NCL); - - // 3) Output: dst is { d_inner, n_t, n_s } (CLN) - // - // We need an NCL view of the same buffer: - // desired NCL logical shape: { L_out = n_t, C = nr, N = n_s } - // - // Original CLN layout: - // dst->ne = { nr, n_t, n_s } - // dst->nb[0] = sizeof(float) - // dst->nb[1] = nr * sizeof(float) - // dst->nb[2] = nr * n_t * sizeof(float) - // - // We want offset_new(L, C, N) = offset_orig(C, L, N). - // Choose: - // nb_y[0] = nr * sizeof(float); // step in L - // nb_y[1] = sizeof(float); // step in C - // nb_y[2] = nr * n_t * sizeof(float); // step in N - int64_t y_ne[GGML_MAX_DIMS] = { n_t, nr, n_s, 1 }; // [L_out, C, N] - size_t y_nb[GGML_MAX_DIMS] = { dst->ne[0] * sizeof(float), sizeof(float), dst->ne[0] * dst->ne[1] * sizeof(float), - dst->nb[3] }; // [nr, 1, nr * n_t] - - acl_tensor_ptr acl_y = ggml_cann_create_tensor(dst->data, ggml_cann_type_mapping(dst->type), - ggml_type_size(dst->type), y_ne, y_nb, 3, ACL_FORMAT_NCL); - - // --- Conv1d parameters: depthwise, stride 1, no padding ("valid") --- - int64_t strideVal[1] = { 1 }; - int64_t paddingVal[1] = { 0 }; - int64_t dilationVal[1] = { 1 }; - - acl_int_array_ptr stride = ggml_cann_create_int_array(strideVal, 1); - acl_int_array_ptr padding = ggml_cann_create_int_array(paddingVal, 1); - acl_int_array_ptr dilation = ggml_cann_create_int_array(dilationVal, 1); - - const bool transposed = false; - const int64_t groups = nr; // depthwise: one group per inner dim - int8_t cubeMathType = 0; - -#ifdef ASCEND_310P - cubeMathType = 1; -#endif - - GGML_CANN_CALL_ACLNN_OP(ctx, Convolution, - acl_x.get(), // input: N, C, L_in = ncs - acl_w.get(), // weight: [C, 1, K] with groups=nr - nullptr, // bias - stride.get(), padding.get(), dilation.get(), transposed, - padding.get(), // output padding (unused for non-transposed) - groups, acl_y.get(), cubeMathType); -} - -void ggml_cann_op_add_rms_norm_fused(ggml_backend_cann_context & ctx, - ggml_tensor * add_node, - ggml_tensor * rms_norm_node) { - // Get the two input tensors for ADD operation - ggml_tensor * x1 = add_node->src[0]; - ggml_tensor * x2 = add_node->src[1]; - - // Create ACL tensors for the two ADD inputs - acl_tensor_ptr acl_x1 = ggml_cann_create_tensor(x1); - acl_tensor_ptr acl_x2 = ggml_cann_create_tensor(x2); - - // Get epsilon parameter from rms_norm_tensor - float eps; - memcpy(&eps, rms_norm_node->op_params, sizeof(float)); - - // Build gamma tensor (RMS normalization scaling factor) - // Gamma should match the normalized dimensions (last dimension of x1) - size_t acl_gamma_nb[GGML_MAX_DIMS]; - acl_gamma_nb[0] = ggml_type_size(rms_norm_node->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - acl_gamma_nb[i] = acl_gamma_nb[i - 1] * x1->ne[i - 1]; - } - acl_tensor_ptr acl_gamma = - get_cache_acl_tensor(ctx, &ctx.rms_norm_one_tensor_cache.cache, ctx.rms_norm_one_tensor_cache.size, x1->ne, - acl_gamma_nb, rms_norm_node->type, - 1, // dims - only the last dimension - 1.0f // value - ); - - // Build rstdOut tensor (output for normalized standard deviation) - // Shape should be the dimensions that are NOT normalized - int64_t acl_rstd_ne[] = { 1, x1->ne[1], x1->ne[2], x1->ne[3] }; - size_t acl_rstd_nb[GGML_MAX_DIMS - 1]; - acl_rstd_nb[0] = sizeof(float); - for (int i = 1; i < GGML_MAX_DIMS - 1; i++) { - acl_rstd_nb[i] = acl_rstd_nb[i - 1] * acl_rstd_ne[i - 1]; - } - acl_tensor_ptr acl_rstd = - get_cache_acl_tensor(ctx, &ctx.rms_norm_zero_tensor_cache.cache, ctx.rms_norm_zero_tensor_cache.size, - acl_rstd_ne, acl_rstd_nb, GGML_TYPE_F32, GGML_MAX_DIMS, - 0.0f // value - ); - - acl_tensor_ptr acl_xout = ggml_cann_create_tensor(add_node); - - // Create yOut tensor (final output after RMS normalization) - acl_tensor_ptr acl_yout = ggml_cann_create_tensor(rms_norm_node); - - // Call fused ADD + RMS_NORM operator - GGML_CANN_CALL_ACLNN_OP(ctx, AddRmsNorm, acl_x1.get(), acl_x2.get(), acl_gamma.get(), - eps, // double type - acl_yout.get(), acl_rstd.get(), acl_xout.get()); -} - -void ggml_cann_gated_linear_attn(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * k = dst->src[0]; - ggml_tensor * v = dst->src[1]; - ggml_tensor * q = dst->src[2]; - ggml_tensor * g = dst->src[3]; - ggml_tensor * s = dst->src[4]; - - int64_t B = dst->src[4]->ne[1]; - int64_t T = dst->src[0]->ne[2]; - int64_t H = dst->src[0]->ne[1]; - int64_t C = dst->ne[0]; - int64_t D = C / H; - int64_t L = T / B; - - int64_t ne_qkg[2] = { 1, D }; - int64_t ne_s[2] = { D, D }; - int64_t ne_st[2] = { ne_s[1], ne_s[0] }; - int64_t ne_vo[2] = { D, 1 }; - int64_t ne_q[1] = { D }; - size_t nb_base = ggml_type_size(k->type); - size_t nb_qkg[2] = { nb_base, nb_base }; - size_t nb_s[2] = { nb_base, D * nb_base }; - size_t nb_st[2] = { nb_s[1], nb_s[0] }; - size_t nb_vo[2] = { nb_base, D * nb_base }; - size_t nb_q[1] = { nb_base }; - - const float scale = ggml_get_op_params_f32(dst, 0); - - acl_tensor_ptr acl_s = ggml_cann_create_tensor(s, s->ne, s->nb, 2, ACL_FORMAT_ND); - acl_tensor_ptr new_state = ggml_cann_create_tensor(dst, s->ne, s->nb, 2, ACL_FORMAT_ND, (B * L * H * D) * nb_base); - cann_copy(ctx, acl_s.get(), new_state.get()); - - for (int64_t b = 0; b < B; b++) { - for (int64_t h = 0; h < H; h++) { - size_t s_offset = (b * (H * D * D) + h * (D * D)) * nb_base; - // D * D - acl_tensor_ptr acl_s_new = - ggml_cann_create_tensor(dst, ne_s, nb_s, 2, ACL_FORMAT_ND, (B * L * H * D) * nb_base + s_offset); - acl_tensor_ptr acl_s_new_t = - ggml_cann_create_tensor(dst, ne_st, nb_st, 2, ACL_FORMAT_ND, (B * L * H * D) * nb_base + s_offset); - for (int64_t l = 0; l < L; l++) { - size_t qkvgo_offset = (b * (L * H * D) + l * (H * D) + h * (D)) * nb_base; - // D * 1 - acl_tensor_ptr acl_k = ggml_cann_create_tensor(k, ne_qkg, nb_qkg, 2, ACL_FORMAT_ND, qkvgo_offset); - acl_tensor_ptr acl_g = ggml_cann_create_tensor(g, ne_qkg, nb_qkg, 2, ACL_FORMAT_ND, qkvgo_offset); - // D - acl_tensor_ptr acl_q = ggml_cann_create_tensor(q, ne_q, nb_q, 1, ACL_FORMAT_ND, qkvgo_offset); - // 1 * D - acl_tensor_ptr acl_v = ggml_cann_create_tensor(v, ne_vo, nb_vo, 2, ACL_FORMAT_ND, qkvgo_offset); - // D - acl_tensor_ptr acl_o = ggml_cann_create_tensor(dst, ne_q, nb_q, 1, ACL_FORMAT_ND, qkvgo_offset); - // k ⊗ v - size_t buf_size = D * D * nb_base; - ggml_cann_pool_alloc buffer_allocator(ctx.pool(), buf_size); - acl_tensor_ptr tmp_tensor = ggml_cann_create_tensor( - buffer_allocator.get(), ggml_cann_type_mapping(k->type), nb_base, ne_s, nb_s, 2); - aclnn_mul(ctx, acl_k.get(), acl_v.get(), tmp_tensor.get()); - //s_new = g ⊗ s_old + k ⊗ v - aclnn_mul(ctx, acl_s_new.get(), acl_g.get(), nullptr); - aclnn_add(ctx, acl_s_new.get(), tmp_tensor.get(), nullptr); - // compute output - GGML_CANN_CALL_ACLNN_OP(ctx, Mv, acl_s_new_t.get(), acl_q.get(), acl_o.get(), 1); - aclnn_muls(ctx, acl_o.get(), scale, nullptr, true); - } - } - } -} - diff --git a/ggml/src/ggml-cann/aclnn_ops.h b/ggml/src/ggml-cann/aclnn_ops.h deleted file mode 100644 index cdbf9260f859..000000000000 --- a/ggml/src/ggml-cann/aclnn_ops.h +++ /dev/null @@ -1,1190 +0,0 @@ -/** - * Copyright (c) 2023-2026 The ggml authors - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to - * deal in the Software without restriction, including without limitation the - * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or - * sell copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING - * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS - * IN THE SOFTWARE. - */ - -#ifndef CANN_ACLNN_OPS -#define CANN_ACLNN_OPS - -#include "acl_tensor.h" -#include "common.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include - -/** - * @brief Repeats a ggml tensor along each dimension to match the dimensions - * of another tensor. - * - * @details This function repeats the elements of a source ggml tensor along - * each dimension to create a destination tensor with the specified - * dimensions. The operation is performed using the ACL backend and - * executed asynchronously on the device. - * - * @param ctx The CANN context used for operations. - * @param dst The ggml tensor representing the destination, which op is - * GGML_OP_REPEAT and specifies the desired dimensions. - */ -void ggml_cann_repeat(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -void ggml_cann_swiglu(ggml_backend_cann_context & ctx, ggml_tensor * dst); -void ggml_cann_geglu(ggml_backend_cann_context & ctx, ggml_tensor * dst, int64_t approximate); - -/** - * @brief Applies the Leaky ReLU activation function to a tensor using the CANN - * backend. - * - * @details This function computes the Leaky ReLU activation for each element of - * the input tensor. The Leaky ReLU function allows a small gradient - * when the unit is not active (i.e., when the input is negative). The - * Leaky ReLU function is defined as: - * \f[ - * \text{dst} = \max(0, src) + \text{negativeSlope} \cdot \min(0, - * src) - * \f] - * `negativeSlope` is in dst->params. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the result of the Leaky ReLU - * activation is stored, which op is `GGML_OP_LEAKY_RELU` - */ -void ggml_cann_leaky_relu(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Concatenates multiple tensors along a specified dimension using the - * CANN backend. - * - * @param ctx The CANN context used for operations. - * @param tensorList A pointer to the list of tensors to be concatenated. - * @param dst The destination tensor where the result of the - * concatenation is stored. dst->op is `GGML_OP_CONCAT`. - * @param concat_dim The dimension along which the tensors are concatenated. - * - * @attention tensorList length should be 2 and the dimension using for concat - * default to 1. - */ -void ggml_cann_concat(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Generates a sequence of evenly spaced values within a specified - * interval for a ggml tensor using the CANN backend. - * - * @details This function creates a sequence of numbers over a specified i - * nterval, starting from `start`, ending before `stop`, and - * incrementing by `step`. The sequence is stored in the destination - * tensor `dst`. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the generated sequence will be stored. - * `start`, 'stop' and 'step' are in dst->op_params and dst->op is - * `GGML_OP_ARANGE`. - */ -void ggml_cann_arange(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Applies a clamp operation to the elements of a ggml tensor using the - * CANN backend. - * - * @details This function clamps the elements of the input tensor `src` to a - * specified range defined by `min` and `max` values. The result is - * stored in the destination tensor `dst`. The operation is defined as: - * \f[ - * y = \max(\min(x, max\_value), min\_value) - * \f] - * where `x` is an element of the input tensor, and `y` is the - * corresponding element in the output tensor. - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the clamped values will be stored. - * dst->op is `GGML_OP_CLAMP`, `min` and `max` value is in dst->params. - */ -void ggml_cann_clamp(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Scales the elements of a ggml tensor by a constant factor using the - * CANN backend. - * - * @details This function multiplies each element of the input tensor `src` by - * a scaling factor `scale`, storing the result in the destination - * tensor `dst`. The operation is defined as: - * \f[ - * dst = src \times scale - * \f] - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the scaled values will be stored. - * dst->op is `GGML_OP_SCALE` and `scale` value is in dst->params. - */ -void ggml_cann_scale(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Sorts the elements of a ggml tensor and returns the indices that - * would sort the tensor using the CANN backend. - * - * @details This function performs an argsort operation on the input tensor - * `src`. It sorts the elements of `src` in either ascending or - * descending order, depending on the `GGML_SORT_ORDER_DESC`, - * and returns the indices that would sort the original tensor. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the sorted indices will be stored. - * dst->op is `GGML_OP_ARGSORT`. - */ -void ggml_cann_argsort(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes the Layer Normalization for a ggml tensor using the CANN - * backend. - * - * @details This function applies the Layer Normalization operation on the - * input tensor `src` and stores the result in the destination tensor - * `dst`. Layer Normalization normalizes the features at each sample in - * a mini-batch independently. It is commonly used in neural networks - * to normalize the activations of a layer by adjusting and scaling - * the outputs. - * The operation is defined as: - * \f[ - * \text { out }=\frac{x-\mathrm{E}[x]}{\sqrt{\text{Var}[x]+eps}} - * \f] - * `Var` defaults dst->ne[0]. `eps` is in dst->params. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the normalized values will be stored. - * @attention `Var` defaults to dst->ne[0]. - */ -void ggml_cann_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes the L2 Normalization for a ggml tensor using the CANN - * backend. - * - * @details This function applies the L2 Normalization operation on the - * input tensor `src` and stores the result in the destination tensor - * `dst`. L2 Normalization scales the input tensor such that the - * L2 norm along the specified dimension equals 1. This operation - * is commonly used in neural networks for feature normalization - * and vector scaling. - * The operation is defined as: - * \f[ - * \text{out} = \frac{x}{\sqrt{\sum{x^2}}} - * \f] - * The normalization is performed along the last dimension by default. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the normalized values will be stored. - * @attention The normalization is performed along the last dimension of the - * input tensor by default. - */ -void ggml_cann_l2_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes the Cross Entropy Loss for a ggml tensor using the CANN - * backend. - * - * @details This function computes the cross entropy loss between the predicted - * logits and target probability distributions. The operation follows - * the same computation pattern as the CPU implementation: - * 1. Applies log_softmax to the logits along the class dimension - * 2. Element-wise multiplication with target distributions - * 3. Summation along the class dimension to get per-sample losses - * 4. Global summation and scaling by -1/nr to get final loss - * - * The computation can be expressed as: - * \f[ - * \text{loss} = -\frac{1}{N} \sum_{i=1}^{N} \sum_{j=1}^{C} y_{ij} \cdot \log(\text{softmax}(x_{ij})) - * \f] - * where \f$N\f$ is the total number of samples, \f$C\f$ is the number - * of classes, \f$x\f$ are the logits, and \f$y\f$ are the target - * probability distributions. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the computed loss will be stored. - * This should be a scalar tensor containing the final loss value. - * - * @note This implementation computes cross entropy between probability - * distributions, not the typical classification cross entropy that - * expects class indices as targets. Both input tensors (src0 and src1) - * should have the same shape and represent probability distributions - * over the class dimension. - * @note The function expects two source tensors: - * - dst->src[0]: Logits tensor (before softmax) - * - dst->src[1]: Target probability distributions tensor - * @note The computation is performed using CANN backend operators including - * LogSoftmax, Mul, ReduceSum, and Muls for the final scaling. - */ -void ggml_cann_cross_entropy_loss(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes the Group Normalization for a ggml tensor using the CANN - * backend. - * - * @brief This function applies the Group Normalization operation on the input - * tensor `src` and stores the result in the destination tensor `dst`. - * Group Normalization divides the channels into groups and normalizes - * the features within each group across spatial locations. - * It is commonly used in convolutional neural networks to improve - * training stability and performance. - * The operation is defined as: - * \f[ - * \text { out }=\frac{x-\mathrm{E}[x]}{\sqrt{\text{Var}[x]+eps}} - * \f] - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the normalized values will be stored. - * `n_groups` is in dst->params, which split C channel to `n_groups`. - * dst->op is `GGML_OP_GROUP_NORM`. - * - * @attention eps defaults to 1e-6f. - */ -void ggml_cann_group_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes the accumulation of tensors using the CANN backend. - * - * @details This function performs an accumulation operation on two tensors. - * Depending on the `inplace` flag, it either updates the destination - * tensor `dst` in place by adding `alpha * src1` to it, or it creates - * a new tensor as the result of `src0 + alpha * src1` and stores it in - * `dst`. - * The operation is defined as: - * \f[ - * dst = src0 + alpha \times src1 - * \f] - * if `inplace` is `true`, `src0` is equal to 'dst'. - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the accumulated values will be stored. - * `inplace` is in dst->params, and dst->op is `GGML_OP_ACC`. - */ -void ggml_cann_acc(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes the sum of elements along the last dimension of a ggml tensor - * using the CANN backend. - * - * @details This function performs a reduction sum operation along the last - * dimension of the input tensor `src`. The result of the sum is stored - * in the destination tensor `dst`. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the reduced values will be stored。 - * dst->op is `GGML_OP_SUM_ROWS`. - * - * @attention `reduce_dims` defaults to 3, which means the last dimension. - */ -void ggml_cann_sum_rows(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes the sum of elements in a ggml tensor. - * - * @details This function performs a reduction sum operation along the last - * dimension of the input tensor `src`. The result of the sum is stored - * in the destination tensor `dst`. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the reduced values will be stored。 - * - */ - -void ggml_cann_sum(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes the cumulative sum of a ggml tensor along dim 0 using the - * CANN backend. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor. dst->op is `GGML_OP_CUMSUM`. - */ -void ggml_cann_cumsum(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes a triangular mask (tril/triu) of a square ggml tensor - * using the CANN backend. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor. dst->op is `GGML_OP_TRI`. - */ -void ggml_cann_tri(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Solves a triangular linear system AX=B using the CANN backend. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor. dst->op is `GGML_OP_SOLVE_TRI`. - */ -void ggml_cann_solve_tri(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Creates a diagonal matrix from a vector using the CANN backend. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor. dst->op is `GGML_OP_DIAG`. - */ -void ggml_cann_diag(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Fills a tensor with a constant scalar value using the CANN backend. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor. dst->op is `GGML_OP_FILL`. - */ -void ggml_cann_fill(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Upsamples a ggml tensor using nearest neighbor interpolation using - * the CANN backend. - * - * @details This function performs upsampling of the input tensor `src` using - * nearest neighbor interpolation. The upsampling is applied to the - * height and width dimensions (last two dimensions) of the tensor. The - * result is stored in the destination tensor `dst`, which must have - * the appropriate dimensions for the upsampled output. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the upsampled values will be stored. - * dst->op is `GGML_OP_UPSCALE`. - */ -void ggml_cann_upsample_nearest2d(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Pads a ggml tensor to match the dimensions of the destination tensor - * using the CANN backend. - * - * @details This function pads the input tensor `src` so that it matches the - * dimensions of the destination tensor `dst`. The amount of padding - * is calculated based on the difference in sizes between `src` and - * `dst` along each dimension. The padded tensor is stored in `dst`. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor, which specifies the target dimensions for - * padding. dst->op is `GGML_OP_PAD`. - */ -void ggml_cann_pad(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Executes a 2D pooling operation on a ggml tensor using the CANN - * backend. - * - * @details This function dispatches the execution of a 2D pooling operation on - * the input tensor `dst`. The type of pooling (average or max) is - * determined by the `op` parameter, which is read from the operation - * parameters of `dst`. The function supports average pooling - * (`GGML_OP_POOL_AVG`) and max pooling (`GGML_OP_POOL_MAX`). If an - * invalid operation is encountered, the function asserts a failure. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor on which the pooling operation is to be - * performed. dst->op is `GGML_OP_POOL_2D`. - */ -void ggml_cann_pool2d(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Duplicates a ggml tensor using the CANN backend. - * - * @details This function duplicates the contents of the source tensor `src` to - * the destination tensor `dst`. The function supports various tensor - * types and configurations, including handling of extra data, type - * conversions, and special cases for contiguous and non-contiguous - * tensors. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the duplicated data will be stored. - * dst->op is `GGML_OP_DUP` - * - * @attention Only support Fp16/FP32. Not support when src and dst have - * different shape and dst is no-contiguous. - * @note: This func need to simplify. - */ -void ggml_cann_dup(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes the Root Mean Square (RMS) normalization of a ggml tensor - * using the CANN backend. - * - * @details This function applies RMS normalization to the input tensor `src` - * and stores the result in the destination tensor `dst`. RMS - * normalization involves computing the root mean square of the input - * tensor along a specified dimension and then dividing each element of - * the tensor by this value, adjusted by a small epsilon value to - * prevent division by zero. - * The operation is defined as: - * \f[ - * \text{RmsNorm}\left(x_i\right)=\frac{x_i}{\text{Rms}(\mathbf{x})} g_i, - * \quad \text { where } \text{Rms}(\mathbf{x})=\sqrt{\frac{1}{n} \sum_{i=1}^n x_i^2+e p s} - * \f] - * `eps` is in dst->op_params. - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the normalized values will be stored. - * dst->op is `GGML_OP_RMS_NORM`. - */ -void ggml_cann_rms_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Applies a diagonal mask to the tensor with a specified value. - * - * @details This function creates a mask tensor filled with ones, then applies - * an upper triangular and lower triangular operation to it based on - * the number of past elements specified. Afterward, it adds the masked - * tensor to the destination tensor in-place. - * - * @param ctx The backend CANN context used for operations. - * @param dst The destination tensor where the result will be stored. dst->op is - * `GGML_OP_DIAG_MASK` - * @param value The value to use for masking. - */ -void ggml_cann_diag_mask(ggml_backend_cann_context & ctx, ggml_tensor * dst, float value); - -/** - * @brief Performs an image-to-column transformation on the input tensor. - * - * @details This function takes an input tensor and applies an image-to-column - * operation, converting spatial dimensions into column-like - * structures suitable for convolutional operations. It supports both - * half-precision (F16) and single-precision (F32) floating-point data - * types. - * - * @param ctx The backend CANN context for executing operations. - * @param dst The destination tensor that stores the result of the operation. - * dst->op is `GGML_OP_IM2COL`. - */ -void ggml_cann_im2col(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes time step embeddings using sine and cosine functions. - * - * @details This function calculates time step embeddings by applying sine and - * cosine transformations to a given input tensor, which is typically - * used in temporal models like diffusion models or transformers to - * encode time information effectively. - * - * @param ctx The backend CANN context for executing operations. - * @param dst The destination tensor where the result of the embedding operation - * will be stored. dst->op is `GGML_OP_TIMESTEP_EMBEDDING`. - */ -void ggml_cann_timestep_embedding(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -// @see ggml_cann_dup. -void ggml_cann_cpy(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -// @see ggml_cann_acc, but copies src1 into dst instead of adding. -void ggml_cann_set(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes the softmax activation with optional masking. - * - * @details This function computes the softmax activation over the input tensor, - * optionally applying a mask and scaling factor. It supports both FP16 - * and FP32 data types and can handle masking by broadcasting the mask - * across rows if necessary. - * The function performs the following steps: - * 1. Multiplies the input tensor by a scale factor. - * 2. Optionally casts the mask tensor to FP32 if it is in FP16 format. - * 3. Broadcasts the mask tensor if its dimensions do not match the - * input tensor's dimensions. - * 4. Adds the mask to the scaled input tensor. - * 5. Applies the softmax activation function along the specified - * dimension. - * - * @param ctx The backend CANN context for executing operations. - * @param dst The destination tensor where the result will be stored. dst->op is - * `GGML_OP_SOFTMAX`. - */ -void ggml_cann_softmax(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Extracts specific rows from a tensor based on indices. - * - * @details This function retrieves rows from a source tensor src0 according to - * the indices provided in another tensor src1 and stores the result in - * a destination tensor (\p dst). - * - * @param ctx The backend CANN context for executing operations. - * @param dst The destination tensor where the extracted rows will be stored. - */ -void ggml_cann_get_rows(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Writes specific rows into a tensor at positions specified by indices. - * - * @details This function copies rows from a source tensor into a destination - * tensor (\p dst) at the positions indicated by the indices in another - * tensor. - * - * @param ctx The backend CANN context for executing operations. - * @param dst The destination tensor where the specified rows will be updated. - */ -void ggml_cann_set_rows(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Executes matrix multiplication for the given tensor. - * - * @details This function performs matrix multiplication on the source tensors - * associated with the destination tensor. It supports matrix - * multiplication F32, F16, and Q8_0. - * - * @param ctx The backend CANN context for executing operations. - * @param dst The destination tensor for storing the result of the matrix - * multiplication. dst->op is `GGML_OP_MUL_MAT`. - */ -void ggml_cann_mul_mat(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Applies Rotary Positional Embedding (RoPE) to the input tensor. - * - * @details This function implements the RoPE mechanism, which is a method to - * encode positional information into sequence data, particularly - * useful in transformer models. It supports both F32 and F16 data - * types. - * - * @param ctx The backend CANN context for executing operations. - * @param dst The destination tensor where the RoPE-transformed data will be - * stored. dst->op is `GGML_OP_ROPE`. - * - * @note The function currently does not support cases where the n_dims is less - * than the input tensor's first dimension. - * @note The function currently does not support cases where the freq_factors is - * not NULL. - * @note The function currently does not support cases where the ext_factor is - * not equal 0. - * @note The function currently does not support cases where the freq_scale is - * not equal 1. - */ -void ggml_cann_rope(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Pre-load the RoPE cache before ACL graph capture. - * - * This function must be called outside of graph capture to perform - * host-to-device memory copies and device memory allocations that are - * not allowed on a captured stream. After pre-loading, the rope cache - * metadata is updated so that the subsequent call to - * aclnn_rope_cache_init (inside graph capture) skips these operations - * and only records the on-device computations into the captured graph. - * - * @param ctx CANN backend context. - * @param dst A ROPE destination tensor from the computation graph. - */ -void ggml_cann_rope_cache_preload(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes the index of the maximum value along the specified dimension - * of a ggml tensor using the CANN backend. - * - * @details This function performs an argmax operation on the input tensor. - * It finds the index of the maximum value along the specified axis - * and stores these indices in the destination tensor `dst`. The - * operation is executed using the CANN backend for optimized performance. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the indices of the maximum values will - * be stored. dst->op is `GGML_OP_ARGMAX`. - */ -void ggml_cann_argmax(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Adds two tensors element-wise and stores the result in a destination - * tensor. - * - * This function performs the operation: - * \f[ - * dst = acl\_src0 + alpha \times acl\_src1 - * \f] - * where alpha is a scalar value and defaults to 1.0f. - * - * @param ctx The context for the CANN backend operations. - * @param acl_src0 The first source tensor. - * @param acl_src1 The second source tensor. - * @param acl_dst The destination tensor where the result will be stored. - */ -void aclnn_add(ggml_backend_cann_context & ctx, - aclTensor * acl_src0, - aclTensor * acl_src1, - aclTensor * acl_dst = nullptr); - -/** - * @brief Sub two tensors element-wise and stores the result in a destination - * tensor. - * - * This function performs the operation: - * \f[ - * dst = acl\_src0 - alpha \times acl\_src1 - * \f] - * where alpha is a scalar value and defaults to 1.0f. - * - * @param ctx The context for the CANN backend operations. - * @param acl_src0 The first source tensor. - * @param acl_src1 The second source tensor. - * @param acl_dst The destination tensor where the result will be stored. - */ -void aclnn_sub(ggml_backend_cann_context & ctx, - aclTensor * acl_src0, - aclTensor * acl_src1, - aclTensor * acl_dst = nullptr); - -/** - * @brief Performs element-wise multiplication of two tensors and stores the - * result in a destination tensor. - * - * This function performs element-wise multiplication of the tensors `acl_src` - * and `acl_other` and stores the result in the destination tensor `acl_dst`. - * The operation is defined as: - * \f[ - * \text {acl_dst }_i=\text {acl_src }_i \times \text {acl_other }_i - * \f] - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The first tensor for element-wise multiplication. - * @param acl_other The second tensor for element-wise multiplication. - * @param acl_dst The destination tensor where the result will be stored. - */ -void aclnn_mul(ggml_backend_cann_context & ctx, - aclTensor * acl_src, - aclTensor * acl_other, - aclTensor * acl_dst = nullptr); - -/** - * @brief Matrix division, optionally in-place. - * - * This function division each element of the source tensor `acl_src` by the - * tensor `acl_other` and stores the result in the destination tensor `acl_dst`. - * If `inplace` is true, `acl_dst` will not be used and the operation is - * performed in-place on `acl_src`. The operation is defined as: \f[ - * \text{dst}_i = \frac{\text{acl_src}_i}{\text{acl_other}_i} - * \f] - * - * @param ctx The context for the CANN backend operations. - * @param acl_src Numerator tensor.. - * @param acl_other Denominator tensor. - * @param acl_dst The destination tensor where the result will be stored if - * `inplace` is false. - * @param inplace Flag indicating whether to perform the operation in-place on - * `acl_src`. - */ -void aclnn_div(ggml_backend_cann_context & ctx, - aclTensor * acl_src, - aclTensor * acl_other, - aclTensor * acl_dst = nullptr); - -/** - * @brief Applies element-wise cosine function to the elements of a tensor. - * - * This function computes the cosine of each element in the source tensor - * `acl_src` and stores the result in the destination tensor `acl_dst`. The - * operation is defined as: \f[ \text {acl_dst }_i=\cos \left(\text {acl_src - * }_i\right) \f] - * - * @param ctx The context for the CANN backend operations. - * @param acl_src The source tensor on which the cosine function will be - * applied. - * @param acl_dst The destination tensor where the cosine results will be - * stored. - */ -void aclnn_cos(ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst); - -/** - * @brief Applies element-wise sine function to the elements of a tensor. - * - * This function computes the sine of each element in the source tensor - `acl_src` - * and stores the result in the destination tensor `acl_dst`. - * The operation is defined as: - * \f[ - * \text {acl_dst }_i=\sin \left(\text {acl_src }_i\right) - * \f] - - * @param ctx The context for the CANN backend operations. - * @param acl_src The source tensor on which the sine function will be applied. - * @param acl_dst The destination tensor where the sine results will be stored. - */ -void aclnn_sin(ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst); - -/** - * @brief Prepares broadcast-compatible ACL tensors for two input tensors and one - * output tensor. - * - * This function checks whether broadcasting is needed between `src0` and `src1`. - * If broadcasting is required, it calculates the proper shapes and creates - * ACL tensors with broadcast parameters. Otherwise, it directly creates ACL tensors - * based on the original tensor shapes. - * - * @param src0 The first input tensor (reference shape). - * @param src1 The second input tensor (possibly broadcasted). - * @param dst The destination/output tensor. - * @param acl_src0 Output pointer to the created ACL tensor corresponding to src0. - * @param acl_src1 Output pointer to the created ACL tensor corresponding to src1. - * @param acl_dst Output pointer to the created ACL tensor corresponding to dst. - */ -void bcast_shape(ggml_tensor * src0, - ggml_tensor * src1, - ggml_tensor * dst, - acl_tensor_ptr & acl_src0, - acl_tensor_ptr & acl_src1, - acl_tensor_ptr & acl_dst); - -/** - * @brief Computes the 1D transposed convolution (deconvolution) of a ggml - * tensor using the CANN backend. - * - * @details This function performs a 1D transposed convolution (also known as - * deconvolution) operation on the input tensor. The computed result is stored - * in the destination tensor `dst`. The operation is optimized using the CANN - * backend for improved performance. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the transposed convolution result - * will be stored. dst->op is `GGML_OP_CONV_TRANSPOSE_1D`. - */ -void ggml_cann_conv_transpose_1d(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Applies the ELU (Exponential Linear Unit) activation to a ggml tensor - * using the CANN backend. - * - * @details This function performs an element-wise ELU activation on the input - * tensor. - * The result is written to the destination tensor `dst` in-place. - * The ELU function is defined as: - * - * \text{ELU}(x) = - * \begin{cases} - * x, & \text{if } x > 0 \\ - * \alpha \left( \exp(x) - 1 \right), & \text{if } x \leq 0 - * \end{cases} - * - * where α (alpha) is a hyperparameter, typically set to 1.0. - * This operation is optimized using the CANN backend for high-performance - * inference or training. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the ELU-activated result will be stored. - * dst->op is expected to be `GGML_OP_ELU`. - */ -void ggml_cann_elu(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Computes the mean of a ggml tensor element-wise using the CANN backend. - * - * @details This function calculates the element-wise mean of the input tensor. - * The result is written to the destination tensor `dst`. - * The mean is computed by averaging the values across the entire tensor. - * - * This operation is optimized using the CANN backend for high-performance inference or training. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the mean result will be stored. - * dst->op is expected to be `GGML_OP_MEAN`. - */ -void ggml_cann_mean(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Applies 1D reflect padding to a ggml tensor using the CANN backend. - * - * @details This function performs 1D reflect padding on the input tensor. - * The amount of padding on each side is specified by parameters stored in `dst->op_params`. - * The operation reflects the values at the borders of the tensor to generate the padded output. - * - * This operation is optimized using the CANN backend for high-performance inference or training. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the padded result will be stored. - * dst->op is expected to be `GGML_OP_PAD_REFLECT_1D`. - */ -void ggml_cann_pad_reflect_1d(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Counts the number of equal elements in two ggml tensors using the CANN backend. - * - * @details This function performs an element-wise comparison between two input tensors, - * and counts the number of positions where the elements are equal. The result is - * stored in the destination tensor `dst` as a scalar. - * - * The operation is optimized using the CANN backend, making it suitable for - * high-performance inference or training scenarios. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the result will be stored. - * dst->op is expected to be `GGML_OP_COUNT_EQUAL`. - */ -void ggml_cann_count_equal(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Applies the Step activation function to a ggml tensor using the CANN backend. - * - * @details This function applies a step function element-wise to the input tensor, where - * each element is transformed to 1.0 if it is greater than 0, and 0.0 otherwise. - * The result is stored in the destination tensor `dst`. - * - * This operation is accelerated using the CANN backend to improve runtime performance. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the result will be stored. - * dst->op is expected to be `GGML_OP_STEP`. - */ -void ggml_cann_step(ggml_backend_cann_context & ctx, ggml_tensor * dst); -void ggml_cann_softplus(ggml_backend_cann_context & ctx, ggml_tensor * dst); -void ggml_cann_geglu_quick(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Performs the Flash Attention extended operator using the CANN backend. - * - * @details This function implements the memory-efficient Flash Attention algorithm - * for computing scaled dot-product attention with hardware acceleration. - * The result is stored in the destination tensor `dst`. - * - * This operation is accelerated using the CANN backend to improve runtime performance. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the result will be stored. - * dst->op is expected to be `GGML_OP_FLASH_ATTN_EXT`. - */ -void ggml_cann_flash_attn_ext(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Forward Gated Linear Attention on the CANN backend. - * - * Expects dst->src[0..4] = {k, v, q, g, s} with shape conventions: - * k, v, q, g: [D] with outer dims T x H batched as ne[2]=T, ne[1]=H - * s: initial state [B, H, D, D], where B is batch and D=C/H - * dst holds both outputs (o) and updated state; a scale factor is read from op params. - * - * The kernel updates per time step l: S_new = g ⊗ S_old + k ⊗ v, then computes o = (S_new^T q) * scale. - * - * @param ctx Backend context providing stream/allocator utilities. - * @param dst Output tensor; src deps are k, v, q, g, s as above. - */ -void ggml_cann_gated_linear_attn(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Launches an asynchronous task using the memory allocator. - * - * This macro submit an asynchronous task on the specified stream. - * The task uses memory allocated by the allocator. It is guaranteed - * that the memory will not be accessed by other tasks until this task - * completes, due to the sequential execution order within the same stream. - * - * @param OP_NAME aclnn operator name. - * @param args Additional arguments required by the task. - * - * @note - * Memory from the allocator will be "freed" immediately and can be - * reallocated to other pointers. However, it won't be accessed by any - * other task before this asynchronous task ends, because all tasks in the - * same stream are executed in queue order. - */ - -# define GGML_CANN_CALL_ACLNN_OP(CTX, OP_NAME, ...) \ - do { \ - uint64_t workspaceSize = 0; \ - aclOpExecutor * executor; \ - void * workspaceAddr = nullptr; \ - ACL_CHECK(aclnn##OP_NAME##GetWorkspaceSize(__VA_ARGS__, &workspaceSize, &executor)); \ - /* workspace should alloced in main thread to keep malloc order when using vmm. */ \ - if (workspaceSize > 0) { \ - ggml_cann_pool_alloc workspace_allocator(CTX.pool(), workspaceSize); \ - workspaceAddr = workspace_allocator.get(); \ - } \ - ACL_CHECK(aclnn##OP_NAME(workspaceAddr, workspaceSize, executor, CTX.stream())); \ - } while (0) - -/** - * @brief Performs sparse expert-based matrix multiplication using the CANN backend. - * - * @details This function implements a MoE-style batched matrix multiplication, where each input token - * is routed to one or more experts, and each expert corresponds to a specific [D, M] weight matrix - * in the source tensor `src0`. The routing indices are provided via the `ids` tensor. - * - * For each token (from `src1`), the function selects the corresponding expert(s) as specified by `ids`, - * performs the matrix multiplication with the selected expert's weight submatrix (from `src0`), - * and stores the results in `dst`. This operation is optimized and executed on the CANN backend. - * - * Dimensions: - * - src0: [D, M, A, 1], where A is the number of experts - * - src1: [D, B, N, 1], where N is batch size and B is the slot count per sample - * - ids : [K, N], where K is the number of experts each token is routed to - * - dst : [M, K, N, 1], output tensor storing the result of expert × token multiplication - * - * The function handles two main modes: - * - If `ne12 == 1`, a simpler per-token loop is used. - * - TODO: If `ne12 > 1`, grouped multiplication and memory copying is used for efficiency. - * - * @param ctx The CANN context used for operations. - * @param dst The destination tensor where the expert-weighted token outputs are stored. - * Expected to be of shape [M, K, N, 1]. - */ -void ggml_cann_mul_mat_id(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Performs fused ADD + RMS_NORM operation using the CANN backend. - * - * This function fuses the ADD and RMS_NORM operations into a single kernel call - * for better performance. It first adds two input tensors (x1 + x2), then applies - * RMS normalization to the result. - * - * @param ctx The context for the CANN backend operations. - * @param dst The ADD operation node, contains the two input tensors to be added. - * @param rms_norm_tensor The RMS_NORM operation node, contains the gamma weights - * and epsilon parameter. - */ -void ggml_cann_op_add_rms_norm_fused(ggml_backend_cann_context & ctx, - ggml_tensor * add_node, - ggml_tensor * rms_norm_node); - -/** - * @brief Check whether a tensor is a weight tensor for matrix multiplication. - * - * @details Checks whether the given tensor serves as weight parameters in matrix multiplication operations, - * typically within neural network layers. The function maintains a static set of canonical weight - * naming suffixes from Transformer-based architectures. Uses substring matching to identify weight - * tensors even with hierarchical naming patterns. - * - * @param tensor Pointer to the target ggml_tensor object (const-qualified). - */ -static bool is_matmul_weight(const ggml_tensor * tensor) { - std::string name = ggml_get_name(tensor); - static const std::unordered_set weight_suffixes{ "output.weight", "attn_q.weight", - "attn_k.weight", "attn_v.weight", - "attn_output.weight", "ffn_gate.weight", - "ffn_up.weight", "ffn_down.weight" }; - - for (const auto & suffix : weight_suffixes) { - if (name.find(suffix) != std::string::npos) { - return true; - } - } - return false; -} - -/** - * @brief Applies a element-wise operation to two input tensors using the CANN - * backend. - * - * This templated function takes a binary operator and applies it to two source - * tensors - * associated with the destination tensor. The function handles broadcasting as - * needed. - * - * @tparam binary_op A callable object (e.g., lambda or function pointer) representing - * the binary operation to be performed. It must take three arguments: - * (ggml_backend_cann_context&, aclTensor*, aclTensor*, aclTensor*). - * - * @param ctx The CANN backend context used to manage execution and resources. - * @param dst The destination tensor. - */ -template void ggml_cann_binary_op(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src0 = dst->src[0]; - ggml_tensor * src1 = dst->src[1]; - - acl_tensor_ptr acl_src0, acl_src1, acl_dst; - - // Need bcast - bcast_shape(src0, src1, dst, acl_src0, acl_src1, acl_dst); - binary_op(ctx, acl_src0.get(), acl_src1.get(), acl_dst.get()); -} - -/** - * @brief Applies a unary operation to an input tensor using the CANN backend. - * - * This templated function applies a unary operator to the source tensor of `dst` - * and stores the result in the destination tensor. - * - * @tparam unary_op A callable with the signature: - * void(ggml_backend_cann_context&, aclTensor *, aclTensor *) - * where the first aclTensor is the source and the second is the destination. - * @param ctx The CANN backend context for managing resources and execution. - * @param dst The destination tensor. Its src[0] is treated as the input tensor. - */ -template -void ggml_cann_op_unary(ggml_backend_cann_context & ctx, ggml_tensor * dst) { - ggml_tensor * src = dst->src[0]; - - acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); - acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); - - unary_op(ctx, acl_src.get(), acl_dst.get()); -} - -/** - * @brief Applies a unary operation to a ggml tensor using the CANN backend. - * - * @details This function applies a unary operation to the input tensor using - * a user-provided lambda or callable `unary_op`. The lambda receives the - * CANN backend context and two ACL tensors: the source and the destination. - * - * Internally, this function handles the conversion from GGML tensors to ACL tensors, - * calls the provided unary op, and manages resource cleanup. The input is assumed - * to be `dst->src[0]`, and the result is written to `dst`. - * - * This utility simplifies writing unary op wrappers by abstracting tensor preparation. - * - * @param unary_op A callable that performs the unary operation using CANN ACL APIs. - * @param ctx The CANN context for operation execution. - * @param dst The destination ggml_tensor where the result will be stored. - * The input tensor is assumed to be `dst->src[0]`. - * - * @see GGML_CANN_CALL_OP_UNARY - */ -void ggml_cann_op_unary(std::function unary_op, - ggml_backend_cann_context & ctx, - ggml_tensor * dst); - -void ggml_cann_ssm_conv(ggml_backend_cann_context & ctx, ggml_tensor * dst); - -/** - * @brief Applies a gated (GLU-style) unary operation using the CANN backend. - * - * @details This function performs a gated activation such as GEGLU or ReGLU. - * It supports two input modes: - * - * 1. **Dual input mode**: `dst->src[0]` and `dst->src[1]` are both valid tensors. - * These are used directly as the value and gate tensors. - * - * 2. **Packed input mode**: Only `dst->src[0]` is valid, and it is assumed to - * contain a concatenation of value and gate along the first dimension. This tensor - * will be split into two equal halves to form the value and gate inputs. - * - * The function applies a user-provided unary operation (e.g., GELU) to the value tensor, - * then multiplies the result in-place with the gate tensor: - * - * @code - * dst = unary_op(value) * gate; - * @endcode - * - * The `swapped` parameter (from `dst->op_params[1]`) allows flipping the - * order of value/gate in the packed input case. - * - * @param unary_op A callable that performs the unary operation using CANN ACL APIs. - * It receives (ctx, acl_value_tensor, acl_output_tensor). - * @param ctx The CANN context used for execution. - * @param dst The destination ggml_tensor. Source tensors are in `dst->src[0]` and optionally `src[1]`. - * - * @see GGML_CANN_CALL_OP_UNARY_GATED - */ -void ggml_cann_op_unary_gated(std::function unary_op, - ggml_backend_cann_context & ctx, - ggml_tensor * dst); - -/** - * @brief Helper macro to call a unary ACL operator via ggml_cann_op_unary. - * - * This macro wraps the specified ACLNN unary operator name into a lambda expression, - * and passes it to `ggml_cann_op_unary`, which handles the common logic for executing - * unary ops in the CANN backend. - * - * Internally, this macro expands to a lambda like: - * @code - * [](ggml_backend_cann_context& ctx, aclTensor* acl_src, aclTensor* acl_dst) { - * GGML_CANN_CALL_ACLNN_OP(ctx, OP_NAME, acl_src, acl_dst); - * }; - * @endcode - * - * This lambda is then passed to `ggml_cann_op_unary`, which applies the operation. - * - * @param OP_NAME The name of the ACL unary operator to invoke via GGML_CANN_CALL_ACLNN_OP. - * - * @see ggml_cann_op_unary - * @see GGML_CANN_CALL_ACLNN_OP - */ -# define GGML_CANN_CALL_OP_UNARY(OP_NAME) \ - do { \ - auto lambda = [](ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst) { \ - GGML_CANN_CALL_ACLNN_OP(ctx, OP_NAME, acl_src, acl_dst); \ - }; \ - ggml_cann_op_unary(lambda, ctx, dst); \ - } while (0) - -/** - * @brief Helper macro to call a gated unary ACL operator via ggml_cann_op_unary_gated. - * - * This macro wraps the specified ACLNN unary operator name into a lambda expression, - * and passes it to `ggml_cann_op_unary_gated`, which handles the common logic for - * executing gated unary ops in the CANN backend. - * - * Internally, this macro expands to a lambda like: - * @code - * [](ggml_backend_cann_context& ctx, aclTensor* acl_src, aclTensor* acl_dst) { - * GGML_CANN_CALL_ACLNN_OP(ctx, OP_NAME, acl_src, acl_dst); - * }; - * @endcode - * - * This lambda is then passed to `ggml_cann_op_unary_gated`, which applies the operation. - * - * @param OP_NAME The name of the ACL unary operator to invoke via GGML_CANN_CALL_ACLNN_OP. - * - * @see ggml_cann_op_unary_gated - * @see GGML_CANN_CALL_ACLNN_OP - */ -# define GGML_CANN_CALL_OP_UNARY_GATED(OP_NAME) \ - do { \ - auto lambda = [](ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst) { \ - GGML_CANN_CALL_ACLNN_OP(ctx, OP_NAME, acl_src, acl_dst); \ - }; \ - ggml_cann_op_unary_gated(lambda, ctx, dst); \ - } while (0) - -#endif // CANN_ACLNN_OPS - -/** - * @brief Performs outer product operation on two ggml tensors using the CANN backend. - * - * @details This function computes the outer product of two input tensors (src0 and src1) - * and stores the result in the destination tensor. The outer product operation is defined as: - * dst[i,j,k,l] = sum_m (src0[i,m,k,l] * src1[j,m,k,l]) - * - * The function supports multiple data types including F32, F16. For floating-point - * types, it uses batch matrix multiplication for efficient computation. - * - * The implementation handles 4D tensor broadcasting and batch processing automatically. - * - * @param ctx The CANN backend context for operation execution and memory management. - * @param dst The destination ggml_tensor where the outer product result will be stored. - * The input tensors are assumed to be `dst->src[0]` and `dst->src[1]`. - * - * @see GGML_CANN_CALL_ACLNN_OP for CANN operator invocation - */ -void ggml_cann_out_prod(ggml_backend_cann_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cann/common.h b/ggml/src/ggml-cann/common.h deleted file mode 100644 index 1c6e685c38ca..000000000000 --- a/ggml/src/ggml-cann/common.h +++ /dev/null @@ -1,651 +0,0 @@ -/* - * Copyright (c) 2023-2026 The ggml authors - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to - * deal in the Software without restriction, including without limitation the - * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or - * sell copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING - * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS - * IN THE SOFTWARE. - */ - -#ifndef CANN_COMMON_H -#define CANN_COMMON_H - -#include "../ggml-impl.h" -#include "../include/ggml-cann.h" -#include "../include/ggml.h" - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#define MATRIX_ROW_PADDING 512 -#define GGML_CANN_MAX_STREAMS 8 - -/** - * @brief Handles CANN-related errors by printing an error message and - * terminating the program. - * @param stmt The statement that caused the error. - * @param func The function in which the error occurred. - * @param file The file in which the error occurred. - * @param line The line number at which the error occurred. - * @param msg The error message. - */ -[[noreturn]] void ggml_cann_error(const char * stmt, const char * func, const char * file, int line, const char * msg); - -/** - * @brief Checks the result of a CANN function call and invokes the error - * handler if the call fails. - * @param stmt The CANN function call to check. - * @param success The success code that indicates the call was successful. - * @param error_fn The function to call to retrieve the error message. - */ -#define ACL_CHECK_GEN(stmt, success, error_fn) \ - do { \ - int err_code = (stmt); \ - if (err_code != (success)) { \ - ggml_cann_error(#stmt, __func__, __FILE__, __LINE__, error_fn()); \ - } \ - } while (0); - -#define ACL_CHECK(stmt) ACL_CHECK_GEN(stmt, 0, aclGetRecentErrMsg) - -/** - * @brief Contains information about CANN devices. - */ -struct ggml_cann_device_info { - /** - * @brief Number of CANN devices available. - */ - int32_t device_count; - - /** - * @brief Information about a single CANN device. - */ - struct cann_device_info { - int cc; /**< Compute capability. */ - size_t smpb; /**< Maximum shared memory per block. */ - bool vmm; /**< Virtual memory support. */ - size_t vmm_granularity; /**< Granularity of virtual memory. */ - size_t total_vram; /**< Total video RAM available on the device. */ - }; - - cann_device_info devices[GGML_CANN_MAX_DEVICES] = {}; /**< Array of CANN device information. */ -}; - -const ggml_cann_device_info & ggml_cann_info(); - -void ggml_cann_set_device(int32_t device); - -std::optional get_env_as_lowercase(const std::string & name); -bool parse_bool(const std::string & value); -int parse_integer(const std::string & value); - -/** - * @brief Abstract base class for memory pools used by CANN. - */ -struct ggml_cann_pool { - /** - * @brief Virtual destructor for the memory pool. - */ - virtual ~ggml_cann_pool() = default; - - /** - * @brief Allocates memory from the pool. - * - * @param size The size of the memory block to allocate. - * @param actual_size Pointer to a variable where the actual allocated size - * will be stored. - * @return Pointer to the allocated memory block. - */ - virtual void * alloc(size_t size, size_t * actual_size) = 0; - - /** - * @brief Frees a previously allocated memory block. - * - * @param ptr Pointer to the memory block to free. - * @param size Size of the memory block to free. - * @note Note that all CANN opertors are running async. Make sure memory is - * still avaiable before this operator finished. - */ - virtual void free(void * ptr, size_t size) = 0; -}; - -/** - * @brief RAII wrapper for managing memory allocations from a CANN memory pool. - */ -struct ggml_cann_pool_alloc { - ggml_cann_pool * pool = nullptr; /**< Pointer to the memory pool. */ - void * ptr = nullptr; /**< Pointer to the allocated memory block. */ - size_t actual_size = 0; /**< Actual size of the allocated memory block. */ - - /** - * @brief Default constructor. - */ - ggml_cann_pool_alloc() = default; - - /** - * @brief Constructor that initializes the memory pool. - * @param pool Reference to the memory pool. - */ - explicit ggml_cann_pool_alloc(ggml_cann_pool & pool) : pool(&pool) {} - - /** - * @brief Constructor that initializes the memory pool and allocates memory. - * @param pool Reference to the memory pool. - * @param size Size of the memory block to allocate. - */ - ggml_cann_pool_alloc(ggml_cann_pool & pool, size_t size) : pool(&pool) { alloc(size); } - - /** - * @brief Destructor that frees the allocated memory block. - */ - ~ggml_cann_pool_alloc() { - if (ptr != nullptr) { - pool->free(ptr, actual_size); - } - } - - /** - * @brief Allocates memory from the pool. - * @param size Size of the memory block to allocate. - * @return Pointer to the allocated memory block. - */ - void * alloc(size_t size) { - GGML_ASSERT(pool != nullptr); - GGML_ASSERT(ptr == nullptr); - ptr = pool->alloc(size, &this->actual_size); - return ptr; - } - - /** - * @brief Allocates memory from a specific memory pool. - * @param pool Reference to the memory pool. - * @param size Size of the memory block to allocate. - * @return Pointer to the allocated memory block. - */ - void * alloc(ggml_cann_pool & pool, size_t size) { - this->pool = &pool; - return alloc(size); - } - - /** - * @brief Gets the pointer to the allocated memory block. - * @return Pointer to the allocated memory block. - */ - void * get() { return ptr; } - - // Deleted copy constructor - ggml_cann_pool_alloc(const ggml_cann_pool_alloc &) = delete; - - // Deleted move constructor - ggml_cann_pool_alloc(ggml_cann_pool_alloc &&) = delete; - - // Deleted copy assignment operator - ggml_cann_pool_alloc & operator=(const ggml_cann_pool_alloc &) = delete; - - // Deleted move assignment operator - ggml_cann_pool_alloc & operator=(ggml_cann_pool_alloc &&) = delete; -}; - -#ifdef USE_ACL_GRAPH -struct ggml_graph_node_properties { - // dst tensor - void * node_address; - ggml_type node_type; - int64_t ne[GGML_MAX_DIMS]; - size_t nb[GGML_MAX_DIMS]; - - // src tensor - void * src_address[GGML_MAX_SRC]; - ggml_type src_type[GGML_MAX_SRC]; - int64_t src_ne[GGML_MAX_SRC][GGML_MAX_DIMS]; - size_t src_nb[GGML_MAX_SRC][GGML_MAX_DIMS]; - - // op - ggml_op node_op; - int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; - - /** - * @brief Check if a ggml tensor node matches this property set. - * - * This function compares all relevant fields (address, op type, shape, source inputs, op params) - * to determine whether the current node matches these previously recorded properties. - * - * @param node The current ggml tensor node. - * @return true if all fields match (excluding GGML_OP_VIEW); false otherwise. - */ - bool has_matching_properties(ggml_tensor * node) { - if (node->data != this->node_address && node->op != GGML_OP_VIEW) { - return false; - } - - if (node->op != this->node_op) { - return false; - } - - if (node->type != this->node_type) { - return false; - } - - for (int i = 0; i < GGML_MAX_DIMS; i++) { - if (node->ne[i] != this->ne[i]) { - return false; - } - if (node->nb[i] != this->nb[i]) { - return false; - } - } - - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (node->src[i]) { - if (node->src[i]->data != this->src_address[i] && node->op != GGML_OP_VIEW) { - return false; - } - - if (node->src[i]->type != this->src_type[i]) { - return false; - } - - for (int d = 0; d < GGML_MAX_DIMS; d++) { - if (node->src[i]->ne[d] != this->src_ne[i][d]) { - return false; - } - if (node->src[i]->nb[d] != this->src_nb[i][d]) { - return false; - } - } - } else { - if (this->src_address[i] != nullptr) { - return false; - } - } - } - - return memcmp(this->op_params, node->op_params, GGML_MAX_OP_PARAMS) == 0; - } -}; - -struct ggml_cann_graph { - ~ggml_cann_graph() { - if (graph != nullptr) { - ACL_CHECK(aclmdlRIDestroy(graph)); - } - } - - aclmdlRI graph = nullptr; - - std::vector ggml_graph_properties; - - /** - * @brief Create a new CANN graph from a ggml computation graph. - * - * This function creates a new ggml_cann_graph object and fills its node properties - * (operation type, dimensions, strides, input sources, and operation parameters) - * based on the current ggml computation graph. - * - * Each node in the ggml graph is mapped to a property entry in the new CANN graph: - * - node address - * - operation type - * - shape (ne) and strides (nb) - * - source tensor addresses - * - operation parameters - * - * @param cgraph The current ggml computation graph. - * @return Pointer to the newly created ggml_cann_graph object. - */ - static ggml_cann_graph * create_from_cgraph(ggml_cgraph * cgraph) { - ggml_cann_graph * new_graph = new ggml_cann_graph(); - new_graph->ggml_graph_properties.resize(cgraph->n_nodes); - - for (int node_idx = 0; node_idx < cgraph->n_nodes; ++node_idx) { - ggml_tensor * node = cgraph->nodes[node_idx]; - auto & prop = new_graph->ggml_graph_properties[node_idx]; - - prop.node_address = node->data; - prop.node_op = node->op; - prop.node_type = node->type; - - std::copy_n(node->ne, GGML_MAX_DIMS, prop.ne); - std::copy_n(node->nb, GGML_MAX_DIMS, prop.nb); - - for (int src = 0; src < GGML_MAX_SRC; ++src) { - if (node->src[src]) { - prop.src_address[src] = node->src[src]->data; - prop.src_type[src] = node->src[src]->type; - std::copy_n(node->src[src]->ne, GGML_MAX_DIMS, prop.src_ne[src]); - std::copy_n(node->src[src]->nb, GGML_MAX_DIMS, prop.src_nb[src]); - } else { - prop.src_address[src] = nullptr; - prop.src_type[src] = GGML_TYPE_COUNT; - std::fill_n(prop.src_ne[src], GGML_MAX_DIMS, 0); - std::fill_n(prop.src_nb[src], GGML_MAX_DIMS, 0); - } - } - - memcpy(prop.op_params, node->op_params, GGML_MAX_OP_PARAMS); - } - - return new_graph; - } - - /** - * @brief Check whether this CANN graph matches the given ggml computation graph. - * - * This function compares the number of nodes and each node's properties - * (operation type, dimensions, strides, inputs, and operation parameters) - * to determine whether this CANN graph matches the given ggml graph. - * - * @param cgraph The current ggml computation graph. - * @return true if this CANN graph matches the ggml graph; false otherwise. - */ - bool matches_cgraph(ggml_cgraph * cgraph) { - if (this->ggml_graph_properties.size() != static_cast(cgraph->n_nodes)) { - return false; - } - - for (int i = 0; i < cgraph->n_nodes; ++i) { - if (!this->ggml_graph_properties[i].has_matching_properties(cgraph->nodes[i])) { - return false; - } - } - - return true; - } -}; - -/** - * @brief LRU cache for managing ggml_cann_graph objects. - * - * This class maintains a list of shared_ptr to ggml_cann_graph objects - * and enforces a maximum capacity. It provides methods to push new graphs, - * move existing graphs to the front (most recently used), and clear the cache. - */ -struct ggml_cann_graph_lru_cache { - size_t capacity; /**< Maximum number of graphs in the cache. */ - - std::list cache_list; /**< List storing cached graphs as raw pointers. */ - - ggml_cann_graph_lru_cache() { capacity = parse_integer(get_env_as_lowercase("GGML_CANN_GRAPH_CACHE_CAPACITY").value_or("12")); } - - /** - * @brief Push a new graph to the front of the cache. - * If the cache exceeds capacity, the least recently used graph is deleted. - * @param new_node Pointer to the new ggml_cann_graph to cache. - * Ownership is transferred to the cache (cache will delete it). - */ - void push(ggml_cann_graph * new_node) { - if (cache_list.size() >= capacity) { - ggml_cann_graph * old = cache_list.back(); - cache_list.pop_back(); - delete old; // free the old graph - } - cache_list.push_front(new_node); - } - - /** - * @brief Clear all graphs from the cache (also frees memory). - */ - void clear() { - for (auto ptr : cache_list) { - delete ptr; - } - cache_list.clear(); - } - - /** - * @brief Destructor that clears the cache and frees all cached graphs. - */ - ~ggml_cann_graph_lru_cache() { clear(); } - - /** - * @brief Find a cached CANN graph that matches the given ggml graph and move it to front. - * - * This function iterates through the cached CANN graphs stored in the LRU cache and - * compares them against the given ggml computation graph. If a matching graph is found, - * it is promoted to the front of the LRU cache and returned. Otherwise, the function - * returns nullptr. - * - * @param cgraph The current ggml computation graph. - * @return true if found; false otherwise. - */ - bool find_and_move_to_front(ggml_cgraph * cgraph) { - for (auto & graph_ptr : this->cache_list) { - if (graph_ptr->matches_cgraph(cgraph)) { - cache_list.remove(graph_ptr); - cache_list.push_front(graph_ptr); - return true; - } - } - return false; - } -}; -#endif // USE_ACL_GRAPH - -struct ggml_cann_rope_cache { - ~ggml_cann_rope_cache() { - if (theta_scale_cache) { - ACL_CHECK(aclrtFree(theta_scale_cache)); - } - if (sin_cache) { - ACL_CHECK(aclrtFree(sin_cache)); - } - if (cos_cache) { - ACL_CHECK(aclrtFree(cos_cache)); - } - if (position_select_index) { - ACL_CHECK(aclrtFree(position_select_index)); - } - if (theta_scale_exp_host) { - free(theta_scale_exp_host); - } - if (position_select_index_host) { - free(position_select_index_host); - } - if (yarn_ramp_cache) { - ACL_CHECK(aclrtFree(yarn_ramp_cache)); - } - } - - bool equal(int64_t theta_scale_length, - int64_t position_length, - float ext_factor, - float theta_scale, - float freq_scale, - float attn_factor, - bool is_neox, - bool indep_sects, - bool mrope_used, - bool is_imrope, - int sections[4]) { - return this->theta_scale_length == theta_scale_length && this->position_length == position_length && - this->ext_factor == ext_factor && this->theta_scale == theta_scale && this->freq_scale == freq_scale && - this->attn_factor == attn_factor && this->is_neox == is_neox && this->indep_sects == indep_sects && - this->mrope_used == mrope_used && this->is_imrope == is_imrope && this->sections[0] == sections[0] && - this->sections[1] == sections[1] && this->sections[2] == sections[2] && this->sections[3] == sections[3]; - } - - void set(int64_t theta_scale_length, - int64_t position_length, - float ext_factor, - float theta_scale, - float freq_scale, - float attn_factor, - bool is_neox, - bool indep_sects, - bool mrope_used, - bool is_imrope, - int sections[4]) { - this->theta_scale_length = theta_scale_length; - this->position_length = position_length; - this->ext_factor = ext_factor; - this->theta_scale = theta_scale; - this->freq_scale = freq_scale; - this->attn_factor = attn_factor; - this->is_neox = is_neox; - this->indep_sects = indep_sects; - this->mrope_used = mrope_used; - this->is_imrope = is_imrope; - this->sections[0] = sections[0]; - this->sections[1] = sections[1]; - this->sections[2] = sections[2]; - this->sections[3] = sections[3]; - } - - // memory cache, prepare before inferencing. - void * theta_scale_cache = nullptr; - float * theta_scale_exp_host = nullptr; - int * position_select_index_host = nullptr; - void * position_select_index = nullptr; - void * yarn_ramp_cache = nullptr; - // sin/cos cache, used only to accelerate first layer on each device - void * sin_cache = nullptr; - void * cos_cache = nullptr; - // Properties to check before reusing the sincos cache - int64_t theta_scale_length = 0; - int64_t position_length = 0; - bool cached = false; - float ext_factor = 0.0f; - float theta_scale = 0.0f; - float freq_scale = 0.0f; - float attn_factor = 0.0f; - bool is_neox = false; - bool indep_sects = false; - bool mrope_used = false; - int sections[4] = { 0, 0, 0, 0 }; - bool is_imrope = false; -}; - -struct ggml_cann_tensor_cache { - ~ggml_cann_tensor_cache() { - if (cache != nullptr) { - ACL_CHECK(aclrtFree(cache)); - } - } - - void * cache = nullptr; - int64_t size = 0; -}; - -/** - * @brief Context for managing CANN backend operations. - */ -struct ggml_backend_cann_context { - int32_t device; /**< Device ID. */ - std::string name; /**< Name of the device. */ - std::string description; /**< Description of the device. */ - aclrtEvent copy_event = nullptr; /**< Event for managing copy operations. */ -#ifdef USE_ACL_GRAPH - /// Cached CANN ACL graph used for executing the current ggml computation graph. - ggml_cann_graph_lru_cache graph_lru_cache; - bool acl_graph_mode = true; -#endif - bool async_mode; - // Rope Cache - ggml_cann_rope_cache rope_cache; - // Constant Pool - ggml_cann_tensor_cache rms_norm_one_tensor_cache; - ggml_cann_tensor_cache rms_norm_zero_tensor_cache; - - aclrtStream streams[GGML_CANN_MAX_STREAMS] = { nullptr }; /**< Array of streams for the device. */ - - /** - * @brief Constructor for initializing the context with a given device. - * @param device Device ID. - */ - explicit ggml_backend_cann_context(int device) : device(device), name("CANN" + std::to_string(device)) { - ggml_cann_set_device(device); - description = aclrtGetSocName(); - -#ifdef USE_ACL_GRAPH - acl_graph_mode = parse_bool(get_env_as_lowercase("GGML_CANN_ACL_GRAPH").value_or("on")); - GGML_LOG_INFO("%s: device %d execution mode is %s (%s)\n", __func__, device, acl_graph_mode ? "GRAPH" : "EAGER", - acl_graph_mode ? "acl graph enabled" : "acl graph disabled"); -#endif - } - - /** - * @brief Destructor for cleaning up resources. - */ - ~ggml_backend_cann_context() { - ggml_cann_set_device(device); - if (copy_event != nullptr) { - ACL_CHECK(aclrtDestroyEvent(copy_event)); - } - for (int i = 0; i < GGML_CANN_MAX_STREAMS; ++i) { - if (streams[i] != nullptr) { - ACL_CHECK(aclrtDestroyStream(streams[i])); - } - } - } - - /** - * @brief Get or create a stream for a given index. - * @param stream Index of the stream. - * @return The stream corresponding to the given index. - */ - aclrtStream stream(int stream) { - if (streams[stream] == nullptr) { - // If the device is not set here, destroying the stream later may cause a mismatch - // between the thread contexts where the stream was created and destroyed. - // However, I printed the device_id, thread_id, and stream, and they are all consistent. - ACL_CHECK(aclrtSetDevice(device)); - ACL_CHECK(aclrtCreateStream(&streams[stream])); - } - return streams[stream]; - } - - /** - * @brief Get or create the default stream (index 0). - * @return The default stream. - */ - aclrtStream stream() { return stream(0); } - - // TODO: each stream should have a memory pool. - std::unique_ptr mem_pool; /**< Memory pool for the device. */ - - /** - * @brief Create a new memory pool for a given device. - * @param device Device ID. - * @return A unique pointer to the new memory pool. - */ - static std::unique_ptr new_pool_for_device(int device); - - /** - * @brief Get or create the memory pool for the context. - * @return Reference to the memory pool. - */ - ggml_cann_pool & pool() { - if (mem_pool == nullptr) { - mem_pool = new_pool_for_device(device); - } - return *mem_pool; - } -}; - -#endif // CANN_COMMON_H diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp.inc similarity index 54% rename from ggml/src/ggml-cann/ggml-cann.cpp rename to ggml/src/ggml-cann/ggml-cann.cpp.inc index 5f51ea3bb3c8..83b2b90ac8e0 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp.inc @@ -20,13 +20,2191 @@ * IN THE SOFTWARE. */ -#include "ggml-cann.h" +#include "ggml-cann.h.inc" -#include "ggml-backend-impl.h" -#include "ggml-cann/aclnn_ops.h" -#include "ggml-cann/common.h" -#include "ggml-impl.h" -#include "ggml.h" +#include "ggml-backend-impl-defs.inc" +/** + * Copyright (c) 2023-2026 The ggml authors + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + */ + +/* + * Copyright (c) 2023-2026 The ggml authors + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + */ + + +/* Inlined common.h.inc */ +/* + * Copyright (c) 2023-2026 The ggml authors + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS + * IN THE SOFTWARE. + */ + + +#include "../ggml-impl-defs.inc" +#include "../include/ggml-cann.h.inc" +#include "../include/ggml.h.inc" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define MATRIX_ROW_PADDING 512 +#define GGML_CANN_MAX_STREAMS 8 + +/** + * @brief Handles CANN-related errors by printing an error message and + * terminating the program. + * @param stmt The statement that caused the error. + * @param func The function in which the error occurred. + * @param file The file in which the error occurred. + * @param line The line number at which the error occurred. + * @param msg The error message. + */ +[[noreturn]] void ggml_cann_error(const char * stmt, const char * func, const char * file, int line, const char * msg); + +/** + * @brief Checks the result of a CANN function call and invokes the error + * handler if the call fails. + * @param stmt The CANN function call to check. + * @param success The success code that indicates the call was successful. + * @param error_fn The function to call to retrieve the error message. + */ +#define ACL_CHECK_GEN(stmt, success, error_fn) \ + do { \ + int err_code = (stmt); \ + if (err_code != (success)) { \ + ggml_cann_error(#stmt, __func__, __FILE__, __LINE__, error_fn()); \ + } \ + } while (0); + +#define ACL_CHECK(stmt) ACL_CHECK_GEN(stmt, 0, aclGetRecentErrMsg) + +/** + * @brief Contains information about CANN devices. + */ +struct ggml_cann_device_info { + /** + * @brief Number of CANN devices available. + */ + int32_t device_count; + + /** + * @brief Information about a single CANN device. + */ + struct cann_device_info { + int cc; /**< Compute capability. */ + size_t smpb; /**< Maximum shared memory per block. */ + bool vmm; /**< Virtual memory support. */ + size_t vmm_granularity; /**< Granularity of virtual memory. */ + size_t total_vram; /**< Total video RAM available on the device. */ + }; + + cann_device_info devices[GGML_CANN_MAX_DEVICES] = {}; /**< Array of CANN device information. */ +}; + +const ggml_cann_device_info & ggml_cann_info(); + +void ggml_cann_set_device(int32_t device); + +std::optional get_env_as_lowercase(const std::string & name); +bool parse_bool(const std::string & value); +int parse_integer(const std::string & value); + +/** + * @brief Abstract base class for memory pools used by CANN. + */ +struct ggml_cann_pool { + /** + * @brief Virtual destructor for the memory pool. + */ + virtual ~ggml_cann_pool() = default; + + /** + * @brief Allocates memory from the pool. + * + * @param size The size of the memory block to allocate. + * @param actual_size Pointer to a variable where the actual allocated size + * will be stored. + * @return Pointer to the allocated memory block. + */ + virtual void * alloc(size_t size, size_t * actual_size) = 0; + + /** + * @brief Frees a previously allocated memory block. + * + * @param ptr Pointer to the memory block to free. + * @param size Size of the memory block to free. + * @note Note that all CANN opertors are running async. Make sure memory is + * still avaiable before this operator finished. + */ + virtual void free(void * ptr, size_t size) = 0; +}; + +/** + * @brief RAII wrapper for managing memory allocations from a CANN memory pool. + */ +struct ggml_cann_pool_alloc { + ggml_cann_pool * pool = nullptr; /**< Pointer to the memory pool. */ + void * ptr = nullptr; /**< Pointer to the allocated memory block. */ + size_t actual_size = 0; /**< Actual size of the allocated memory block. */ + + /** + * @brief Default constructor. + */ + ggml_cann_pool_alloc() = default; + + /** + * @brief Constructor that initializes the memory pool. + * @param pool Reference to the memory pool. + */ + explicit ggml_cann_pool_alloc(ggml_cann_pool & pool) : pool(&pool) {} + + /** + * @brief Constructor that initializes the memory pool and allocates memory. + * @param pool Reference to the memory pool. + * @param size Size of the memory block to allocate. + */ + ggml_cann_pool_alloc(ggml_cann_pool & pool, size_t size) : pool(&pool) { alloc(size); } + + /** + * @brief Destructor that frees the allocated memory block. + */ + ~ggml_cann_pool_alloc() { + if (ptr != nullptr) { + pool->free(ptr, actual_size); + } + } + + /** + * @brief Allocates memory from the pool. + * @param size Size of the memory block to allocate. + * @return Pointer to the allocated memory block. + */ + void * alloc(size_t size) { + GGML_ASSERT(pool != nullptr); + GGML_ASSERT(ptr == nullptr); + ptr = pool->alloc(size, &this->actual_size); + return ptr; + } + + /** + * @brief Allocates memory from a specific memory pool. + * @param pool Reference to the memory pool. + * @param size Size of the memory block to allocate. + * @return Pointer to the allocated memory block. + */ + void * alloc(ggml_cann_pool & pool, size_t size) { + this->pool = &pool; + return alloc(size); + } + + /** + * @brief Gets the pointer to the allocated memory block. + * @return Pointer to the allocated memory block. + */ + void * get() { return ptr; } + + // Deleted copy constructor + ggml_cann_pool_alloc(const ggml_cann_pool_alloc &) = delete; + + // Deleted move constructor + ggml_cann_pool_alloc(ggml_cann_pool_alloc &&) = delete; + + // Deleted copy assignment operator + ggml_cann_pool_alloc & operator=(const ggml_cann_pool_alloc &) = delete; + + // Deleted move assignment operator + ggml_cann_pool_alloc & operator=(ggml_cann_pool_alloc &&) = delete; +}; + +#ifdef USE_ACL_GRAPH +struct ggml_graph_node_properties { + // dst tensor + void * node_address; + ggml_type node_type; + int64_t ne[GGML_MAX_DIMS]; + size_t nb[GGML_MAX_DIMS]; + + // src tensor + void * src_address[GGML_MAX_SRC]; + ggml_type src_type[GGML_MAX_SRC]; + int64_t src_ne[GGML_MAX_SRC][GGML_MAX_DIMS]; + size_t src_nb[GGML_MAX_SRC][GGML_MAX_DIMS]; + + // op + ggml_op node_op; + int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; + + /** + * @brief Check if a ggml tensor node matches this property set. + * + * This function compares all relevant fields (address, op type, shape, source inputs, op params) + * to determine whether the current node matches these previously recorded properties. + * + * @param node The current ggml tensor node. + * @return true if all fields match (excluding GGML_OP_VIEW); false otherwise. + */ + bool has_matching_properties(ggml_tensor * node) { + if (node->data != this->node_address && node->op != GGML_OP_VIEW) { + return false; + } + + if (node->op != this->node_op) { + return false; + } + + if (node->type != this->node_type) { + return false; + } + + for (int i = 0; i < GGML_MAX_DIMS; i++) { + if (node->ne[i] != this->ne[i]) { + return false; + } + if (node->nb[i] != this->nb[i]) { + return false; + } + } + + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (node->src[i]) { + if (node->src[i]->data != this->src_address[i] && node->op != GGML_OP_VIEW) { + return false; + } + + if (node->src[i]->type != this->src_type[i]) { + return false; + } + + for (int d = 0; d < GGML_MAX_DIMS; d++) { + if (node->src[i]->ne[d] != this->src_ne[i][d]) { + return false; + } + if (node->src[i]->nb[d] != this->src_nb[i][d]) { + return false; + } + } + } else { + if (this->src_address[i] != nullptr) { + return false; + } + } + } + + return memcmp(this->op_params, node->op_params, GGML_MAX_OP_PARAMS) == 0; + } +}; + +struct ggml_cann_graph { + ~ggml_cann_graph() { + if (graph != nullptr) { + ACL_CHECK(aclmdlRIDestroy(graph)); + } + } + + aclmdlRI graph = nullptr; + + std::vector ggml_graph_properties; + + /** + * @brief Create a new CANN graph from a ggml computation graph. + * + * This function creates a new ggml_cann_graph object and fills its node properties + * (operation type, dimensions, strides, input sources, and operation parameters) + * based on the current ggml computation graph. + * + * Each node in the ggml graph is mapped to a property entry in the new CANN graph: + * - node address + * - operation type + * - shape (ne) and strides (nb) + * - source tensor addresses + * - operation parameters + * + * @param cgraph The current ggml computation graph. + * @return Pointer to the newly created ggml_cann_graph object. + */ + static ggml_cann_graph * create_from_cgraph(ggml_cgraph * cgraph) { + ggml_cann_graph * new_graph = new ggml_cann_graph(); + new_graph->ggml_graph_properties.resize(cgraph->n_nodes); + + for (int node_idx = 0; node_idx < cgraph->n_nodes; ++node_idx) { + ggml_tensor * node = cgraph->nodes[node_idx]; + auto & prop = new_graph->ggml_graph_properties[node_idx]; + + prop.node_address = node->data; + prop.node_op = node->op; + prop.node_type = node->type; + + std::copy_n(node->ne, GGML_MAX_DIMS, prop.ne); + std::copy_n(node->nb, GGML_MAX_DIMS, prop.nb); + + for (int src = 0; src < GGML_MAX_SRC; ++src) { + if (node->src[src]) { + prop.src_address[src] = node->src[src]->data; + prop.src_type[src] = node->src[src]->type; + std::copy_n(node->src[src]->ne, GGML_MAX_DIMS, prop.src_ne[src]); + std::copy_n(node->src[src]->nb, GGML_MAX_DIMS, prop.src_nb[src]); + } else { + prop.src_address[src] = nullptr; + prop.src_type[src] = GGML_TYPE_COUNT; + std::fill_n(prop.src_ne[src], GGML_MAX_DIMS, 0); + std::fill_n(prop.src_nb[src], GGML_MAX_DIMS, 0); + } + } + + memcpy(prop.op_params, node->op_params, GGML_MAX_OP_PARAMS); + } + + return new_graph; + } + + /** + * @brief Check whether this CANN graph matches the given ggml computation graph. + * + * This function compares the number of nodes and each node's properties + * (operation type, dimensions, strides, inputs, and operation parameters) + * to determine whether this CANN graph matches the given ggml graph. + * + * @param cgraph The current ggml computation graph. + * @return true if this CANN graph matches the ggml graph; false otherwise. + */ + bool matches_cgraph(ggml_cgraph * cgraph) { + if (this->ggml_graph_properties.size() != static_cast(cgraph->n_nodes)) { + return false; + } + + for (int i = 0; i < cgraph->n_nodes; ++i) { + if (!this->ggml_graph_properties[i].has_matching_properties(cgraph->nodes[i])) { + return false; + } + } + + return true; + } +}; + +/** + * @brief LRU cache for managing ggml_cann_graph objects. + * + * This class maintains a list of shared_ptr to ggml_cann_graph objects + * and enforces a maximum capacity. It provides methods to push new graphs, + * move existing graphs to the front (most recently used), and clear the cache. + */ +struct ggml_cann_graph_lru_cache { + size_t capacity; /**< Maximum number of graphs in the cache. */ + + std::list cache_list; /**< List storing cached graphs as raw pointers. */ + + ggml_cann_graph_lru_cache() { capacity = parse_integer(get_env_as_lowercase("GGML_CANN_GRAPH_CACHE_CAPACITY").value_or("12")); } + + /** + * @brief Push a new graph to the front of the cache. + * If the cache exceeds capacity, the least recently used graph is deleted. + * @param new_node Pointer to the new ggml_cann_graph to cache. + * Ownership is transferred to the cache (cache will delete it). + */ + void push(ggml_cann_graph * new_node) { + if (cache_list.size() >= capacity) { + ggml_cann_graph * old = cache_list.back(); + cache_list.pop_back(); + delete old; // free the old graph + } + cache_list.push_front(new_node); + } + + /** + * @brief Clear all graphs from the cache (also frees memory). + */ + void clear() { + for (auto ptr : cache_list) { + delete ptr; + } + cache_list.clear(); + } + + /** + * @brief Destructor that clears the cache and frees all cached graphs. + */ + ~ggml_cann_graph_lru_cache() { clear(); } + + /** + * @brief Find a cached CANN graph that matches the given ggml graph and move it to front. + * + * This function iterates through the cached CANN graphs stored in the LRU cache and + * compares them against the given ggml computation graph. If a matching graph is found, + * it is promoted to the front of the LRU cache and returned. Otherwise, the function + * returns nullptr. + * + * @param cgraph The current ggml computation graph. + * @return true if found; false otherwise. + */ + bool find_and_move_to_front(ggml_cgraph * cgraph) { + for (auto & graph_ptr : this->cache_list) { + if (graph_ptr->matches_cgraph(cgraph)) { + cache_list.remove(graph_ptr); + cache_list.push_front(graph_ptr); + return true; + } + } + return false; + } +}; +#endif // USE_ACL_GRAPH + +struct ggml_cann_rope_cache { + ~ggml_cann_rope_cache() { + if (theta_scale_cache) { + ACL_CHECK(aclrtFree(theta_scale_cache)); + } + if (sin_cache) { + ACL_CHECK(aclrtFree(sin_cache)); + } + if (cos_cache) { + ACL_CHECK(aclrtFree(cos_cache)); + } + if (position_select_index) { + ACL_CHECK(aclrtFree(position_select_index)); + } + if (theta_scale_exp_host) { + free(theta_scale_exp_host); + } + if (position_select_index_host) { + free(position_select_index_host); + } + if (yarn_ramp_cache) { + ACL_CHECK(aclrtFree(yarn_ramp_cache)); + } + } + + bool equal(int64_t theta_scale_length, + int64_t position_length, + float ext_factor, + float theta_scale, + float freq_scale, + float attn_factor, + bool is_neox, + bool indep_sects, + bool mrope_used, + bool is_imrope, + int sections[4]) { + return this->theta_scale_length == theta_scale_length && this->position_length == position_length && + this->ext_factor == ext_factor && this->theta_scale == theta_scale && this->freq_scale == freq_scale && + this->attn_factor == attn_factor && this->is_neox == is_neox && this->indep_sects == indep_sects && + this->mrope_used == mrope_used && this->is_imrope == is_imrope && this->sections[0] == sections[0] && + this->sections[1] == sections[1] && this->sections[2] == sections[2] && this->sections[3] == sections[3]; + } + + void set(int64_t theta_scale_length, + int64_t position_length, + float ext_factor, + float theta_scale, + float freq_scale, + float attn_factor, + bool is_neox, + bool indep_sects, + bool mrope_used, + bool is_imrope, + int sections[4]) { + this->theta_scale_length = theta_scale_length; + this->position_length = position_length; + this->ext_factor = ext_factor; + this->theta_scale = theta_scale; + this->freq_scale = freq_scale; + this->attn_factor = attn_factor; + this->is_neox = is_neox; + this->indep_sects = indep_sects; + this->mrope_used = mrope_used; + this->is_imrope = is_imrope; + this->sections[0] = sections[0]; + this->sections[1] = sections[1]; + this->sections[2] = sections[2]; + this->sections[3] = sections[3]; + } + + // memory cache, prepare before inferencing. + void * theta_scale_cache = nullptr; + float * theta_scale_exp_host = nullptr; + int * position_select_index_host = nullptr; + void * position_select_index = nullptr; + void * yarn_ramp_cache = nullptr; + // sin/cos cache, used only to accelerate first layer on each device + void * sin_cache = nullptr; + void * cos_cache = nullptr; + // Properties to check before reusing the sincos cache + int64_t theta_scale_length = 0; + int64_t position_length = 0; + bool cached = false; + float ext_factor = 0.0f; + float theta_scale = 0.0f; + float freq_scale = 0.0f; + float attn_factor = 0.0f; + bool is_neox = false; + bool indep_sects = false; + bool mrope_used = false; + int sections[4] = { 0, 0, 0, 0 }; + bool is_imrope = false; +}; + +struct ggml_cann_tensor_cache { + ~ggml_cann_tensor_cache() { + if (cache != nullptr) { + ACL_CHECK(aclrtFree(cache)); + } + } + + void * cache = nullptr; + int64_t size = 0; +}; + +/** + * @brief Context for managing CANN backend operations. + */ +struct ggml_backend_cann_context { + int32_t device; /**< Device ID. */ + std::string name; /**< Name of the device. */ + std::string description; /**< Description of the device. */ + aclrtEvent copy_event = nullptr; /**< Event for managing copy operations. */ +#ifdef USE_ACL_GRAPH + /// Cached CANN ACL graph used for executing the current ggml computation graph. + ggml_cann_graph_lru_cache graph_lru_cache; + bool acl_graph_mode = true; +#endif + bool async_mode; + // Rope Cache + ggml_cann_rope_cache rope_cache; + // Constant Pool + ggml_cann_tensor_cache rms_norm_one_tensor_cache; + ggml_cann_tensor_cache rms_norm_zero_tensor_cache; + + aclrtStream streams[GGML_CANN_MAX_STREAMS] = { nullptr }; /**< Array of streams for the device. */ + + /** + * @brief Constructor for initializing the context with a given device. + * @param device Device ID. + */ + explicit ggml_backend_cann_context(int device) : device(device), name("CANN" + std::to_string(device)) { + ggml_cann_set_device(device); + description = aclrtGetSocName(); + +#ifdef USE_ACL_GRAPH + acl_graph_mode = parse_bool(get_env_as_lowercase("GGML_CANN_ACL_GRAPH").value_or("on")); + GGML_LOG_INFO("%s: device %d execution mode is %s (%s)\n", __func__, device, acl_graph_mode ? "GRAPH" : "EAGER", + acl_graph_mode ? "acl graph enabled" : "acl graph disabled"); +#endif + } + + /** + * @brief Destructor for cleaning up resources. + */ + ~ggml_backend_cann_context() { + ggml_cann_set_device(device); + if (copy_event != nullptr) { + ACL_CHECK(aclrtDestroyEvent(copy_event)); + } + for (int i = 0; i < GGML_CANN_MAX_STREAMS; ++i) { + if (streams[i] != nullptr) { + ACL_CHECK(aclrtDestroyStream(streams[i])); + } + } + } + + /** + * @brief Get or create a stream for a given index. + * @param stream Index of the stream. + * @return The stream corresponding to the given index. + */ + aclrtStream stream(int stream) { + if (streams[stream] == nullptr) { + // If the device is not set here, destroying the stream later may cause a mismatch + // between the thread contexts where the stream was created and destroyed. + // However, I printed the device_id, thread_id, and stream, and they are all consistent. + ACL_CHECK(aclrtSetDevice(device)); + ACL_CHECK(aclrtCreateStream(&streams[stream])); + } + return streams[stream]; + } + + /** + * @brief Get or create the default stream (index 0). + * @return The default stream. + */ + aclrtStream stream() { return stream(0); } + + // TODO: each stream should have a memory pool. + std::unique_ptr mem_pool; /**< Memory pool for the device. */ + + /** + * @brief Create a new memory pool for a given device. + * @param device Device ID. + * @return A unique pointer to the new memory pool. + */ + static std::unique_ptr new_pool_for_device(int device); + + /** + * @brief Get or create the memory pool for the context. + * @return Reference to the memory pool. + */ + ggml_cann_pool & pool() { + if (mem_pool == nullptr) { + mem_pool = new_pool_for_device(device); + } + return *mem_pool; + } +}; + + +#include + +#include +#include + +/** + * @brief Maps a ggml_type to its corresponding aclDataType. + * + * @details This function takes a ggml_type as input and returns the corresponding + * aclDataType. It supports mapping for various ggml_types. If the input type + * does not match any of the predefined ggml_types, the function returns + * ACL_DT_UNDEFINED. + * + * @param type The ggml_type to be mapped. + * @return The corresponding aclDataType. If the input type is not recognized, + * ACL_DT_UNDEFINED is returned. + */ +aclDataType ggml_cann_type_mapping(ggml_type type); + +// Deleter for acl objects. +template struct acl_deleter { + void operator()(T * ptr) const noexcept { + if (ptr) { + ACL_CHECK(DestroyFunc(ptr)); + } + } +}; + +using acl_tensor_ptr = std::unique_ptr>; +using acl_int_array_ptr = std::unique_ptr>; +using acl_scalar_ptr = std::unique_ptr>; +using acl_tensor_list_ptr = std::unique_ptr>; + +/** + * @brief Creates an ACL tensor from a ggml_tensor with optional shape. + * + * @details This function creates an ACL tensor based on the properties of the + * provided ggml_tensor. It supports customer shape by adjusting dimensions + * and strides accordingly. If customer shape is applied, additional + * dimensions and strides are calculated based on the provided parameters. + * + * @param tensor Pointer to the ggml_tensor to be converted to ACL tensor. + * @param ne Pointer to an array containing dimensions. Defaults to nullptr + * if no customer shape is applied. + * @param nb Pointer to an array containing strides. Defaults to nullptr + * if no customer shape is applied. + * @param dims Number of dimensions in the tensor. Defaults to 0 if no customer + * shape is applied. + * @param format ACL tensor format. Defaults to ACL_FORMAT_ND. + * @param offset Offset in bytes for the ACL tensor data. Defaults to 0. + * @return Pointer to the created ACL tensor. + */ +acl_tensor_ptr ggml_cann_create_tensor(const ggml_tensor * tensor, + int64_t * ne = nullptr, + size_t * nb = nullptr, + int64_t dims = 0, + aclFormat format = ACL_FORMAT_ND, + size_t offset = 0); + +/** + * @brief Template for creating an ACL tensor from provided parameters. typename TYPE + * should be size_t or float. + * + * @details This function creates an ACL tensor using the provided data pointer, + * data type, dimensions, strides, format, offset, and additional parameters. + * It calculates necessary dimensions and strides based on the provided ne and nb + * arrays, adjusting them for the ACL tensor creation. The ACL storage length + * is also calculated based on the provided dimensions and strides. + * + * @param data_ptr Pointer to the data buffer for the ACL tensor. + * @param dtype ACL data type of the tensor. + * @param type_size Size of each element in the tensor data buffer. + * @param ne Pointer to an array containing tensor dimensions. + * @param nb Pointer to an array containing tensor strides. + * @param dims Number of dimensions of the tensor. + * @param format ACL tensor format. Defaults to ACL_FORMAT_ND. + * @param offset Offset in bytes for the ACL tensor data. Defaults to 0. + * @return Pointer to the created ACL tensor. + */ +template +acl_tensor_ptr ggml_cann_create_tensor(void * data_ptr, + aclDataType dtype, + TYPE type_size, + int64_t * ne, + TYPE * nb, + int64_t dims, + aclFormat format = ACL_FORMAT_ND, + size_t offset = 0) { + int64_t tmp_ne[GGML_MAX_DIMS * 2]; + int64_t tmp_stride[GGML_MAX_DIMS * 2]; + + memcpy(tmp_ne, ne, dims * sizeof(int64_t)); + for (int i = 0; i < dims; i++) { + tmp_stride[i] = nb[i] / type_size; + } + + int64_t acl_storage_len = 1; + for (int i = 0; i < dims; i++) { + acl_storage_len += (tmp_ne[i] - 1) * tmp_stride[i]; + } + + std::reverse(tmp_ne, tmp_ne + dims); + std::reverse(tmp_stride, tmp_stride + dims); + + aclTensor * raw = + aclCreateTensor(tmp_ne, dims, dtype, tmp_stride, offset / type_size, format, &acl_storage_len, 1, data_ptr); + + return acl_tensor_ptr(raw); +} + +/** + * @brief Create an ACL int array resource wrapped in a smart pointer. + * + * This function constructs an aclIntArray from the provided int64_t values + * and returns it as an acl_int_array_ptr (a std::unique_ptr with a custom + * deleter). The returned pointer owns the ACL resource and will automatically + * destroy it via aclDestroyIntArray(). + * + * @param value Pointer to the int64_t elements. + * @param size Number of elements in value. + * + * @return A smart pointer managing the created ACL int array. + */ +acl_int_array_ptr ggml_cann_create_int_array(const int64_t * value, uint64_t size); + +/** + * @brief Create an ACL scalar resource wrapped in a smart pointer. + * + * This function constructs an aclScalar from the raw value pointer and ACL + * data type, then returns it as an acl_scalar_ptr (a std::unique_ptr with + * a custom deleter). The returned pointer owns the ACL scalar and will + * automatically destroy it via aclDestroyScalar(). + * + * @param value Pointer to the raw scalar memory. + * @param dataType ACL data type of the scalar. + * + * @return A smart pointer managing the created ACL scalar. + */ +acl_scalar_ptr ggml_cann_create_scalar(void * value, aclDataType dataType); + +/** + * @brief Create an ACL tensor list from multiple tensor smart pointers. + * + * This function accepts a variadic list of acl_tensor_ptr (a unique_ptr with + * custom deleter) and produces an aclTensorList using aclCreateTensorList(). + * + * The lifecycle management of the tensor objects changes as follows: + * - aclCreateTensorList() takes ownership of the tensors + * - Each input smart pointer releases ownership using release() + * - As a result, the tensors will NOT be destroyed by unique_ptr + * - Instead, they will be destroyed when aclDestroyTensorList() is called + * + * This ensures correct ownership transfer and prevents double-free situations. + * + * @param acl_tensor_ptr Variadic template parameter; each argument must be + * a unique_ptr-like type supporting get() and release(). + * + * @param tensors Variadic list of acl_tensor_ptr objects. Ownership of + * each tensor is transferred away from these smart pointers. + * + * @return A smart pointer (acl_tensor_list_ptr) owning the created ACL tensor list. + * + * @note This implementation is C++11 compatible. The ownership-release process is + * executed using a pack expansion inside an initializer list. + */ +template acl_tensor_list_ptr ggml_cann_create_tensor_list(acl_tensor_ptr &&... tensors) { + aclTensor * raw_tensors[] = { tensors.get()... }; + aclTensorList * raw = aclCreateTensorList(raw_tensors, sizeof...(tensors)); + // aclTensor will release by aclTensorList, so release ownership without + // destroying the tensor + int dummy[] = { (tensors.release(), 0)... }; + GGML_UNUSED(dummy); + return acl_tensor_list_ptr(raw); +} + +/** + * @brief Checks if tensors require broadcasting based on their shapes. + * + * @details This function determines if two ggml_tensors need to be broadcasted for + * element-wise operations. Broadcasting is necessary if the shapes of the + * tensors are not identical and no dimension in either tensor equals 1. + * + * @param t0 Pointer to the first ggml_tensor. + * @param t1 Pointer to the second ggml_tensor. + * @return True if broadcasting is needed, False otherwise. + * + * @remarks This function iterates over the dimensions of t0 and t1. It checks if each + * dimension in t1 differs from t0's corresponding dimension and is not equal + * to 1. If such a dimension is found, broadcasting is required to align t1 + * with t0 for element-wise operations. + */ +bool ggml_cann_need_bcast(const ggml_tensor * t0, const ggml_tensor * t1); + +/** + * @brief Computes broadcast shapes and strides for two ggml_tensors. + * + * @details This function calculates the broadcast shapes and strides for two ggml_tensors, + * following the broadcasting rules similar to numpy. It adjusts dimensions and + * strides to ensure compatibility for element-wise operations where one tensor + * can be broadcasted to match the shape of another tensor. + * + * @param src0 Pointer to the first ggml_tensor. + * @param src1 Pointer to the second ggml_tensor. + * @param bcast_ne_src0 Output array to store broadcasted dimensions for src0. + * @param bcast_ne_src1 Output array to store broadcasted dimensions for src1. + * @param bcast_nb_src0 Output array to store broadcasted strides for src0. + * @param bcast_nb_src1 Output array to store broadcasted strides for src1. + * @return Number of dimensions in the broadcasted shape. + * + * @pre ggml_can_repeat(src1, src0) must return true, indicating src1 can be broadcasted + * to match src0. + * + * @remarks This function iterates over the dimensions of src0 and src1, calculating the + * necessary broadcast dimensions and strides. If a dimension requires broadcasting + * (i.e., its size in src1 is smaller than in src0), an additional dimension is + * added with size calculated to match src0's dimension. This adjustment ensures + * that src1 can be element-wise broadcasted to src0's shape. + * + * How it works: + * + * if dim0 has padding. + * a -> (2, 2) padding = 2 + * a: [[1, 2, *, *] + * [2, 3, *, *]] + * nb = (8, 4, 2) + * + * if a should bcast with b -> (2, 4) + * b' -> (2, 2, 2) + * b : [[1, 2, 3, 4, *, *] + * [5, 6, 7, 8, *, *]] + * nb = (12, 6, 1) + * + * after bcast: + * a' -> (2, 1, 2) + * a': [[[1, 2], *, *] + * [[2, 3], *, *]] + * nb = (8, 4, 2, 1) + * + * b' : [[[1, 2], [3, 4], *, *] + * [[5, 6], [7, 8], *, *]] + * nb = (12, 6, 2, 1) + * \endcode + * + * dim1 in a inserted dim, should add nb for dim1, + * and all other nb moves to next in order. + */ +int64_t ggml_cann_get_bcast_shape(const ggml_tensor * src0, + const ggml_tensor * src1, + int64_t * bcast_ne_src0, + int64_t * bcast_ne_src1, + size_t * bcast_nb_src0, + size_t * bcast_nb_src1); + +// Bcast macro to avoid duplicate code. +#define BCAST_SHAPE(src0, src1) \ + int64_t bcast_##src0##_ne[GGML_MAX_DIMS * 2]; \ + int64_t bcast_##src1##_ne[GGML_MAX_DIMS * 2]; \ + size_t bcast_##src0##_nb[GGML_MAX_DIMS * 2]; \ + size_t bcast_##src1##_nb[GGML_MAX_DIMS * 2]; \ + int64_t bcast_dims = ggml_cann_get_bcast_shape(src0, src1, bcast_##src0##_ne, bcast_##src1##_ne, \ + bcast_##src0##_nb, bcast_##src1##_nb); + +#define BCAST_PARAM(tensor) bcast_##tensor##_ne, bcast_##tensor##_nb, bcast_dims + +/** + * @brief Calculates broadcast shapes for matrix multiplication. + * + * @details This function computes the broadcast shapes required for matrix multiplication + * based on the input, weight, and destination tensor shapes. It ensures that the + * dimensions of weight tensors are expanded appropriately to satisfy matrix + * multiplication broadcast rules. + * + * @param input_ne Array containing the dimensions of the input tensor. + * @param weight_ne Array containing the dimensions of the weight tensor. + * @param dst_ne Array containing the dimensions of the destination tensor. + * @param input_nb Array containing the strides of the input tensor. + * @param weight_nb Array containing the strides of the weight tensor. + * @param dst_nb Array containing the strides of the destination tensor. + * @param bcast_input_ne Output array for broadcasted input tensor dimensions. + * @param bcast_weight_ne Output array for broadcasted weight tensor dimensions. + * @param bcast_dst_ne Output array for broadcasted destination tensor dimensions. + * @param bcast_input_nb Output array for broadcasted input tensor strides. + * @param bcast_weight_nb Output array for broadcasted weight tensor strides. + * @param bcast_dst_nb Output array for broadcasted destination tensor strides. + * @return The number of dimensions in the broadcasted tensors. + * + * @remarks This function iterates over the tensor dimensions and calculates the broadcast + * shapes needed for matrix multiplication. It ensures that dimensions where + * weight tensor requires expansion are appropriately handled to conform with + * broadcasting rules. + * @note compare with ggml_cann_get_bcast_shape, mul_mat broadcast need add this new dim + * before cast dim. + * @sa ggml_cann_get_bcast_shape + */ +int64_t ggml_cann_get_mulmat_bcast_shape(const int64_t * input_ne, + const int64_t * weight_ne, + const int64_t * dst_ne, + const size_t * input_nb, + const size_t * weight_nb, + const size_t * dst_nb, + int64_t * bcast_input_ne, + int64_t * bcast_weight_ne, + int64_t * bcast_dst_ne, + size_t * bcast_input_nb, + size_t * bcast_weight_nb, + size_t * bcast_dst_nb); + +// Bcast macro to avoid duplicate code. +#define BCAST_MUL_MAT_SHAPE(input, weight, dst) \ + int64_t bcast_##input##_ne[GGML_MAX_DIMS * 2]; \ + int64_t bcast_##weight##_ne[GGML_MAX_DIMS * 2]; \ + int64_t bcast_##dst##_ne[GGML_MAX_DIMS * 2]; \ + size_t bcast_##input##_nb[GGML_MAX_DIMS * 2]; \ + size_t bcast_##weight##_nb[GGML_MAX_DIMS * 2]; \ + size_t bcast_##dst##_nb[GGML_MAX_DIMS * 2]; \ + int64_t bcast_dims = ggml_cann_get_mulmat_bcast_shape( \ + input->ne, weight->ne, dst->ne, input->nb, weight->nb, dst->nb, bcast_##input##_ne, bcast_##weight##_ne, \ + bcast_##dst##_ne, bcast_##input##_nb, bcast_##weight##_nb, bcast_##dst##_nb); + +#define BCAST_MUL_MAT_PARAM(tensor) bcast_##tensor##_ne, bcast_##tensor##_nb, bcast_dims + + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +/** + * @brief Repeats a ggml tensor along each dimension to match the dimensions + * of another tensor. + * + * @details This function repeats the elements of a source ggml tensor along + * each dimension to create a destination tensor with the specified + * dimensions. The operation is performed using the ACL backend and + * executed asynchronously on the device. + * + * @param ctx The CANN context used for operations. + * @param dst The ggml tensor representing the destination, which op is + * GGML_OP_REPEAT and specifies the desired dimensions. + */ +void ggml_cann_repeat(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +void ggml_cann_swiglu(ggml_backend_cann_context & ctx, ggml_tensor * dst); +void ggml_cann_geglu(ggml_backend_cann_context & ctx, ggml_tensor * dst, int64_t approximate); + +/** + * @brief Applies the Leaky ReLU activation function to a tensor using the CANN + * backend. + * + * @details This function computes the Leaky ReLU activation for each element of + * the input tensor. The Leaky ReLU function allows a small gradient + * when the unit is not active (i.e., when the input is negative). The + * Leaky ReLU function is defined as: + * \f[ + * \text{dst} = \max(0, src) + \text{negativeSlope} \cdot \min(0, + * src) + * \f] + * `negativeSlope` is in dst->params. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the result of the Leaky ReLU + * activation is stored, which op is `GGML_OP_LEAKY_RELU` + */ +void ggml_cann_leaky_relu(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Concatenates multiple tensors along a specified dimension using the + * CANN backend. + * + * @param ctx The CANN context used for operations. + * @param tensorList A pointer to the list of tensors to be concatenated. + * @param dst The destination tensor where the result of the + * concatenation is stored. dst->op is `GGML_OP_CONCAT`. + * @param concat_dim The dimension along which the tensors are concatenated. + * + * @attention tensorList length should be 2 and the dimension using for concat + * default to 1. + */ +void ggml_cann_concat(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Generates a sequence of evenly spaced values within a specified + * interval for a ggml tensor using the CANN backend. + * + * @details This function creates a sequence of numbers over a specified i + * nterval, starting from `start`, ending before `stop`, and + * incrementing by `step`. The sequence is stored in the destination + * tensor `dst`. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the generated sequence will be stored. + * `start`, 'stop' and 'step' are in dst->op_params and dst->op is + * `GGML_OP_ARANGE`. + */ +void ggml_cann_arange(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Applies a clamp operation to the elements of a ggml tensor using the + * CANN backend. + * + * @details This function clamps the elements of the input tensor `src` to a + * specified range defined by `min` and `max` values. The result is + * stored in the destination tensor `dst`. The operation is defined as: + * \f[ + * y = \max(\min(x, max\_value), min\_value) + * \f] + * where `x` is an element of the input tensor, and `y` is the + * corresponding element in the output tensor. + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the clamped values will be stored. + * dst->op is `GGML_OP_CLAMP`, `min` and `max` value is in dst->params. + */ +void ggml_cann_clamp(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Scales the elements of a ggml tensor by a constant factor using the + * CANN backend. + * + * @details This function multiplies each element of the input tensor `src` by + * a scaling factor `scale`, storing the result in the destination + * tensor `dst`. The operation is defined as: + * \f[ + * dst = src \times scale + * \f] + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the scaled values will be stored. + * dst->op is `GGML_OP_SCALE` and `scale` value is in dst->params. + */ +void ggml_cann_scale(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Sorts the elements of a ggml tensor and returns the indices that + * would sort the tensor using the CANN backend. + * + * @details This function performs an argsort operation on the input tensor + * `src`. It sorts the elements of `src` in either ascending or + * descending order, depending on the `GGML_SORT_ORDER_DESC`, + * and returns the indices that would sort the original tensor. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the sorted indices will be stored. + * dst->op is `GGML_OP_ARGSORT`. + */ +void ggml_cann_argsort(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes the Layer Normalization for a ggml tensor using the CANN + * backend. + * + * @details This function applies the Layer Normalization operation on the + * input tensor `src` and stores the result in the destination tensor + * `dst`. Layer Normalization normalizes the features at each sample in + * a mini-batch independently. It is commonly used in neural networks + * to normalize the activations of a layer by adjusting and scaling + * the outputs. + * The operation is defined as: + * \f[ + * \text { out }=\frac{x-\mathrm{E}[x]}{\sqrt{\text{Var}[x]+eps}} + * \f] + * `Var` defaults dst->ne[0]. `eps` is in dst->params. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the normalized values will be stored. + * @attention `Var` defaults to dst->ne[0]. + */ +void ggml_cann_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes the L2 Normalization for a ggml tensor using the CANN + * backend. + * + * @details This function applies the L2 Normalization operation on the + * input tensor `src` and stores the result in the destination tensor + * `dst`. L2 Normalization scales the input tensor such that the + * L2 norm along the specified dimension equals 1. This operation + * is commonly used in neural networks for feature normalization + * and vector scaling. + * The operation is defined as: + * \f[ + * \text{out} = \frac{x}{\sqrt{\sum{x^2}}} + * \f] + * The normalization is performed along the last dimension by default. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the normalized values will be stored. + * @attention The normalization is performed along the last dimension of the + * input tensor by default. + */ +void ggml_cann_l2_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes the Cross Entropy Loss for a ggml tensor using the CANN + * backend. + * + * @details This function computes the cross entropy loss between the predicted + * logits and target probability distributions. The operation follows + * the same computation pattern as the CPU implementation: + * 1. Applies log_softmax to the logits along the class dimension + * 2. Element-wise multiplication with target distributions + * 3. Summation along the class dimension to get per-sample losses + * 4. Global summation and scaling by -1/nr to get final loss + * + * The computation can be expressed as: + * \f[ + * \text{loss} = -\frac{1}{N} \sum_{i=1}^{N} \sum_{j=1}^{C} y_{ij} \cdot \log(\text{softmax}(x_{ij})) + * \f] + * where \f$N\f$ is the total number of samples, \f$C\f$ is the number + * of classes, \f$x\f$ are the logits, and \f$y\f$ are the target + * probability distributions. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the computed loss will be stored. + * This should be a scalar tensor containing the final loss value. + * + * @note This implementation computes cross entropy between probability + * distributions, not the typical classification cross entropy that + * expects class indices as targets. Both input tensors (src0 and src1) + * should have the same shape and represent probability distributions + * over the class dimension. + * @note The function expects two source tensors: + * - dst->src[0]: Logits tensor (before softmax) + * - dst->src[1]: Target probability distributions tensor + * @note The computation is performed using CANN backend operators including + * LogSoftmax, Mul, ReduceSum, and Muls for the final scaling. + */ +void ggml_cann_cross_entropy_loss(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes the Group Normalization for a ggml tensor using the CANN + * backend. + * + * @brief This function applies the Group Normalization operation on the input + * tensor `src` and stores the result in the destination tensor `dst`. + * Group Normalization divides the channels into groups and normalizes + * the features within each group across spatial locations. + * It is commonly used in convolutional neural networks to improve + * training stability and performance. + * The operation is defined as: + * \f[ + * \text { out }=\frac{x-\mathrm{E}[x]}{\sqrt{\text{Var}[x]+eps}} + * \f] + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the normalized values will be stored. + * `n_groups` is in dst->params, which split C channel to `n_groups`. + * dst->op is `GGML_OP_GROUP_NORM`. + * + * @attention eps defaults to 1e-6f. + */ +void ggml_cann_group_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes the accumulation of tensors using the CANN backend. + * + * @details This function performs an accumulation operation on two tensors. + * Depending on the `inplace` flag, it either updates the destination + * tensor `dst` in place by adding `alpha * src1` to it, or it creates + * a new tensor as the result of `src0 + alpha * src1` and stores it in + * `dst`. + * The operation is defined as: + * \f[ + * dst = src0 + alpha \times src1 + * \f] + * if `inplace` is `true`, `src0` is equal to 'dst'. + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the accumulated values will be stored. + * `inplace` is in dst->params, and dst->op is `GGML_OP_ACC`. + */ +void ggml_cann_acc(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes the sum of elements along the last dimension of a ggml tensor + * using the CANN backend. + * + * @details This function performs a reduction sum operation along the last + * dimension of the input tensor `src`. The result of the sum is stored + * in the destination tensor `dst`. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the reduced values will be stored。 + * dst->op is `GGML_OP_SUM_ROWS`. + * + * @attention `reduce_dims` defaults to 3, which means the last dimension. + */ +void ggml_cann_sum_rows(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes the sum of elements in a ggml tensor. + * + * @details This function performs a reduction sum operation along the last + * dimension of the input tensor `src`. The result of the sum is stored + * in the destination tensor `dst`. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the reduced values will be stored。 + * + */ + +void ggml_cann_sum(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes the cumulative sum of a ggml tensor along dim 0 using the + * CANN backend. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor. dst->op is `GGML_OP_CUMSUM`. + */ +void ggml_cann_cumsum(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes a triangular mask (tril/triu) of a square ggml tensor + * using the CANN backend. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor. dst->op is `GGML_OP_TRI`. + */ +void ggml_cann_tri(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Solves a triangular linear system AX=B using the CANN backend. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor. dst->op is `GGML_OP_SOLVE_TRI`. + */ +void ggml_cann_solve_tri(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Creates a diagonal matrix from a vector using the CANN backend. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor. dst->op is `GGML_OP_DIAG`. + */ +void ggml_cann_diag(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Fills a tensor with a constant scalar value using the CANN backend. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor. dst->op is `GGML_OP_FILL`. + */ +void ggml_cann_fill(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Upsamples a ggml tensor using nearest neighbor interpolation using + * the CANN backend. + * + * @details This function performs upsampling of the input tensor `src` using + * nearest neighbor interpolation. The upsampling is applied to the + * height and width dimensions (last two dimensions) of the tensor. The + * result is stored in the destination tensor `dst`, which must have + * the appropriate dimensions for the upsampled output. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the upsampled values will be stored. + * dst->op is `GGML_OP_UPSCALE`. + */ +void ggml_cann_upsample_nearest2d(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Pads a ggml tensor to match the dimensions of the destination tensor + * using the CANN backend. + * + * @details This function pads the input tensor `src` so that it matches the + * dimensions of the destination tensor `dst`. The amount of padding + * is calculated based on the difference in sizes between `src` and + * `dst` along each dimension. The padded tensor is stored in `dst`. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor, which specifies the target dimensions for + * padding. dst->op is `GGML_OP_PAD`. + */ +void ggml_cann_pad(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Executes a 2D pooling operation on a ggml tensor using the CANN + * backend. + * + * @details This function dispatches the execution of a 2D pooling operation on + * the input tensor `dst`. The type of pooling (average or max) is + * determined by the `op` parameter, which is read from the operation + * parameters of `dst`. The function supports average pooling + * (`GGML_OP_POOL_AVG`) and max pooling (`GGML_OP_POOL_MAX`). If an + * invalid operation is encountered, the function asserts a failure. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor on which the pooling operation is to be + * performed. dst->op is `GGML_OP_POOL_2D`. + */ +void ggml_cann_pool2d(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Duplicates a ggml tensor using the CANN backend. + * + * @details This function duplicates the contents of the source tensor `src` to + * the destination tensor `dst`. The function supports various tensor + * types and configurations, including handling of extra data, type + * conversions, and special cases for contiguous and non-contiguous + * tensors. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the duplicated data will be stored. + * dst->op is `GGML_OP_DUP` + * + * @attention Only support Fp16/FP32. Not support when src and dst have + * different shape and dst is no-contiguous. + * @note: This func need to simplify. + */ +void ggml_cann_dup(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes the Root Mean Square (RMS) normalization of a ggml tensor + * using the CANN backend. + * + * @details This function applies RMS normalization to the input tensor `src` + * and stores the result in the destination tensor `dst`. RMS + * normalization involves computing the root mean square of the input + * tensor along a specified dimension and then dividing each element of + * the tensor by this value, adjusted by a small epsilon value to + * prevent division by zero. + * The operation is defined as: + * \f[ + * \text{RmsNorm}\left(x_i\right)=\frac{x_i}{\text{Rms}(\mathbf{x})} g_i, + * \quad \text { where } \text{Rms}(\mathbf{x})=\sqrt{\frac{1}{n} \sum_{i=1}^n x_i^2+e p s} + * \f] + * `eps` is in dst->op_params. + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the normalized values will be stored. + * dst->op is `GGML_OP_RMS_NORM`. + */ +void ggml_cann_rms_norm(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Applies a diagonal mask to the tensor with a specified value. + * + * @details This function creates a mask tensor filled with ones, then applies + * an upper triangular and lower triangular operation to it based on + * the number of past elements specified. Afterward, it adds the masked + * tensor to the destination tensor in-place. + * + * @param ctx The backend CANN context used for operations. + * @param dst The destination tensor where the result will be stored. dst->op is + * `GGML_OP_DIAG_MASK` + * @param value The value to use for masking. + */ +void ggml_cann_diag_mask(ggml_backend_cann_context & ctx, ggml_tensor * dst, float value); + +/** + * @brief Performs an image-to-column transformation on the input tensor. + * + * @details This function takes an input tensor and applies an image-to-column + * operation, converting spatial dimensions into column-like + * structures suitable for convolutional operations. It supports both + * half-precision (F16) and single-precision (F32) floating-point data + * types. + * + * @param ctx The backend CANN context for executing operations. + * @param dst The destination tensor that stores the result of the operation. + * dst->op is `GGML_OP_IM2COL`. + */ +void ggml_cann_im2col(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes time step embeddings using sine and cosine functions. + * + * @details This function calculates time step embeddings by applying sine and + * cosine transformations to a given input tensor, which is typically + * used in temporal models like diffusion models or transformers to + * encode time information effectively. + * + * @param ctx The backend CANN context for executing operations. + * @param dst The destination tensor where the result of the embedding operation + * will be stored. dst->op is `GGML_OP_TIMESTEP_EMBEDDING`. + */ +void ggml_cann_timestep_embedding(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +// @see ggml_cann_dup. +void ggml_cann_cpy(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +// @see ggml_cann_acc, but copies src1 into dst instead of adding. +void ggml_cann_set(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes the softmax activation with optional masking. + * + * @details This function computes the softmax activation over the input tensor, + * optionally applying a mask and scaling factor. It supports both FP16 + * and FP32 data types and can handle masking by broadcasting the mask + * across rows if necessary. + * The function performs the following steps: + * 1. Multiplies the input tensor by a scale factor. + * 2. Optionally casts the mask tensor to FP32 if it is in FP16 format. + * 3. Broadcasts the mask tensor if its dimensions do not match the + * input tensor's dimensions. + * 4. Adds the mask to the scaled input tensor. + * 5. Applies the softmax activation function along the specified + * dimension. + * + * @param ctx The backend CANN context for executing operations. + * @param dst The destination tensor where the result will be stored. dst->op is + * `GGML_OP_SOFTMAX`. + */ +void ggml_cann_softmax(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Extracts specific rows from a tensor based on indices. + * + * @details This function retrieves rows from a source tensor src0 according to + * the indices provided in another tensor src1 and stores the result in + * a destination tensor (\p dst). + * + * @param ctx The backend CANN context for executing operations. + * @param dst The destination tensor where the extracted rows will be stored. + */ +void ggml_cann_get_rows(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Writes specific rows into a tensor at positions specified by indices. + * + * @details This function copies rows from a source tensor into a destination + * tensor (\p dst) at the positions indicated by the indices in another + * tensor. + * + * @param ctx The backend CANN context for executing operations. + * @param dst The destination tensor where the specified rows will be updated. + */ +void ggml_cann_set_rows(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Executes matrix multiplication for the given tensor. + * + * @details This function performs matrix multiplication on the source tensors + * associated with the destination tensor. It supports matrix + * multiplication F32, F16, and Q8_0. + * + * @param ctx The backend CANN context for executing operations. + * @param dst The destination tensor for storing the result of the matrix + * multiplication. dst->op is `GGML_OP_MUL_MAT`. + */ +void ggml_cann_mul_mat(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Applies Rotary Positional Embedding (RoPE) to the input tensor. + * + * @details This function implements the RoPE mechanism, which is a method to + * encode positional information into sequence data, particularly + * useful in transformer models. It supports both F32 and F16 data + * types. + * + * @param ctx The backend CANN context for executing operations. + * @param dst The destination tensor where the RoPE-transformed data will be + * stored. dst->op is `GGML_OP_ROPE`. + * + * @note The function currently does not support cases where the n_dims is less + * than the input tensor's first dimension. + * @note The function currently does not support cases where the freq_factors is + * not NULL. + * @note The function currently does not support cases where the ext_factor is + * not equal 0. + * @note The function currently does not support cases where the freq_scale is + * not equal 1. + */ +void ggml_cann_rope(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Pre-load the RoPE cache before ACL graph capture. + * + * This function must be called outside of graph capture to perform + * host-to-device memory copies and device memory allocations that are + * not allowed on a captured stream. After pre-loading, the rope cache + * metadata is updated so that the subsequent call to + * aclnn_rope_cache_init (inside graph capture) skips these operations + * and only records the on-device computations into the captured graph. + * + * @param ctx CANN backend context. + * @param dst A ROPE destination tensor from the computation graph. + */ +void ggml_cann_rope_cache_preload(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes the index of the maximum value along the specified dimension + * of a ggml tensor using the CANN backend. + * + * @details This function performs an argmax operation on the input tensor. + * It finds the index of the maximum value along the specified axis + * and stores these indices in the destination tensor `dst`. The + * operation is executed using the CANN backend for optimized performance. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the indices of the maximum values will + * be stored. dst->op is `GGML_OP_ARGMAX`. + */ +void ggml_cann_argmax(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Adds two tensors element-wise and stores the result in a destination + * tensor. + * + * This function performs the operation: + * \f[ + * dst = acl\_src0 + alpha \times acl\_src1 + * \f] + * where alpha is a scalar value and defaults to 1.0f. + * + * @param ctx The context for the CANN backend operations. + * @param acl_src0 The first source tensor. + * @param acl_src1 The second source tensor. + * @param acl_dst The destination tensor where the result will be stored. + */ +void aclnn_add(ggml_backend_cann_context & ctx, + aclTensor * acl_src0, + aclTensor * acl_src1, + aclTensor * acl_dst = nullptr); + +/** + * @brief Sub two tensors element-wise and stores the result in a destination + * tensor. + * + * This function performs the operation: + * \f[ + * dst = acl\_src0 - alpha \times acl\_src1 + * \f] + * where alpha is a scalar value and defaults to 1.0f. + * + * @param ctx The context for the CANN backend operations. + * @param acl_src0 The first source tensor. + * @param acl_src1 The second source tensor. + * @param acl_dst The destination tensor where the result will be stored. + */ +void aclnn_sub(ggml_backend_cann_context & ctx, + aclTensor * acl_src0, + aclTensor * acl_src1, + aclTensor * acl_dst = nullptr); + +/** + * @brief Performs element-wise multiplication of two tensors and stores the + * result in a destination tensor. + * + * This function performs element-wise multiplication of the tensors `acl_src` + * and `acl_other` and stores the result in the destination tensor `acl_dst`. + * The operation is defined as: + * \f[ + * \text {acl_dst }_i=\text {acl_src }_i \times \text {acl_other }_i + * \f] + * + * @param ctx The context for the CANN backend operations. + * @param acl_src The first tensor for element-wise multiplication. + * @param acl_other The second tensor for element-wise multiplication. + * @param acl_dst The destination tensor where the result will be stored. + */ +void aclnn_mul(ggml_backend_cann_context & ctx, + aclTensor * acl_src, + aclTensor * acl_other, + aclTensor * acl_dst = nullptr); + +/** + * @brief Matrix division, optionally in-place. + * + * This function division each element of the source tensor `acl_src` by the + * tensor `acl_other` and stores the result in the destination tensor `acl_dst`. + * If `inplace` is true, `acl_dst` will not be used and the operation is + * performed in-place on `acl_src`. The operation is defined as: \f[ + * \text{dst}_i = \frac{\text{acl_src}_i}{\text{acl_other}_i} + * \f] + * + * @param ctx The context for the CANN backend operations. + * @param acl_src Numerator tensor.. + * @param acl_other Denominator tensor. + * @param acl_dst The destination tensor where the result will be stored if + * `inplace` is false. + * @param inplace Flag indicating whether to perform the operation in-place on + * `acl_src`. + */ +void aclnn_div(ggml_backend_cann_context & ctx, + aclTensor * acl_src, + aclTensor * acl_other, + aclTensor * acl_dst = nullptr); + +/** + * @brief Applies element-wise cosine function to the elements of a tensor. + * + * This function computes the cosine of each element in the source tensor + * `acl_src` and stores the result in the destination tensor `acl_dst`. The + * operation is defined as: \f[ \text {acl_dst }_i=\cos \left(\text {acl_src + * }_i\right) \f] + * + * @param ctx The context for the CANN backend operations. + * @param acl_src The source tensor on which the cosine function will be + * applied. + * @param acl_dst The destination tensor where the cosine results will be + * stored. + */ +void aclnn_cos(ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst); + +/** + * @brief Applies element-wise sine function to the elements of a tensor. + * + * This function computes the sine of each element in the source tensor + `acl_src` + * and stores the result in the destination tensor `acl_dst`. + * The operation is defined as: + * \f[ + * \text {acl_dst }_i=\sin \left(\text {acl_src }_i\right) + * \f] + + * @param ctx The context for the CANN backend operations. + * @param acl_src The source tensor on which the sine function will be applied. + * @param acl_dst The destination tensor where the sine results will be stored. + */ +void aclnn_sin(ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst); + +/** + * @brief Prepares broadcast-compatible ACL tensors for two input tensors and one + * output tensor. + * + * This function checks whether broadcasting is needed between `src0` and `src1`. + * If broadcasting is required, it calculates the proper shapes and creates + * ACL tensors with broadcast parameters. Otherwise, it directly creates ACL tensors + * based on the original tensor shapes. + * + * @param src0 The first input tensor (reference shape). + * @param src1 The second input tensor (possibly broadcasted). + * @param dst The destination/output tensor. + * @param acl_src0 Output pointer to the created ACL tensor corresponding to src0. + * @param acl_src1 Output pointer to the created ACL tensor corresponding to src1. + * @param acl_dst Output pointer to the created ACL tensor corresponding to dst. + */ +void bcast_shape(ggml_tensor * src0, + ggml_tensor * src1, + ggml_tensor * dst, + acl_tensor_ptr & acl_src0, + acl_tensor_ptr & acl_src1, + acl_tensor_ptr & acl_dst); + +/** + * @brief Computes the 1D transposed convolution (deconvolution) of a ggml + * tensor using the CANN backend. + * + * @details This function performs a 1D transposed convolution (also known as + * deconvolution) operation on the input tensor. The computed result is stored + * in the destination tensor `dst`. The operation is optimized using the CANN + * backend for improved performance. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the transposed convolution result + * will be stored. dst->op is `GGML_OP_CONV_TRANSPOSE_1D`. + */ +void ggml_cann_conv_transpose_1d(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Applies the ELU (Exponential Linear Unit) activation to a ggml tensor + * using the CANN backend. + * + * @details This function performs an element-wise ELU activation on the input + * tensor. + * The result is written to the destination tensor `dst` in-place. + * The ELU function is defined as: + * + * \text{ELU}(x) = + * \begin{cases} + * x, & \text{if } x > 0 \\ + * \alpha \left( \exp(x) - 1 \right), & \text{if } x \leq 0 + * \end{cases} + * + * where α (alpha) is a hyperparameter, typically set to 1.0. + * This operation is optimized using the CANN backend for high-performance + * inference or training. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the ELU-activated result will be stored. + * dst->op is expected to be `GGML_OP_ELU`. + */ +void ggml_cann_elu(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Computes the mean of a ggml tensor element-wise using the CANN backend. + * + * @details This function calculates the element-wise mean of the input tensor. + * The result is written to the destination tensor `dst`. + * The mean is computed by averaging the values across the entire tensor. + * + * This operation is optimized using the CANN backend for high-performance inference or training. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the mean result will be stored. + * dst->op is expected to be `GGML_OP_MEAN`. + */ +void ggml_cann_mean(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Applies 1D reflect padding to a ggml tensor using the CANN backend. + * + * @details This function performs 1D reflect padding on the input tensor. + * The amount of padding on each side is specified by parameters stored in `dst->op_params`. + * The operation reflects the values at the borders of the tensor to generate the padded output. + * + * This operation is optimized using the CANN backend for high-performance inference or training. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the padded result will be stored. + * dst->op is expected to be `GGML_OP_PAD_REFLECT_1D`. + */ +void ggml_cann_pad_reflect_1d(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Counts the number of equal elements in two ggml tensors using the CANN backend. + * + * @details This function performs an element-wise comparison between two input tensors, + * and counts the number of positions where the elements are equal. The result is + * stored in the destination tensor `dst` as a scalar. + * + * The operation is optimized using the CANN backend, making it suitable for + * high-performance inference or training scenarios. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the result will be stored. + * dst->op is expected to be `GGML_OP_COUNT_EQUAL`. + */ +void ggml_cann_count_equal(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Applies the Step activation function to a ggml tensor using the CANN backend. + * + * @details This function applies a step function element-wise to the input tensor, where + * each element is transformed to 1.0 if it is greater than 0, and 0.0 otherwise. + * The result is stored in the destination tensor `dst`. + * + * This operation is accelerated using the CANN backend to improve runtime performance. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the result will be stored. + * dst->op is expected to be `GGML_OP_STEP`. + */ +void ggml_cann_step(ggml_backend_cann_context & ctx, ggml_tensor * dst); +void ggml_cann_softplus(ggml_backend_cann_context & ctx, ggml_tensor * dst); +void ggml_cann_geglu_quick(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Performs the Flash Attention extended operator using the CANN backend. + * + * @details This function implements the memory-efficient Flash Attention algorithm + * for computing scaled dot-product attention with hardware acceleration. + * The result is stored in the destination tensor `dst`. + * + * This operation is accelerated using the CANN backend to improve runtime performance. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the result will be stored. + * dst->op is expected to be `GGML_OP_FLASH_ATTN_EXT`. + */ +void ggml_cann_flash_attn_ext(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Forward Gated Linear Attention on the CANN backend. + * + * Expects dst->src[0..4] = {k, v, q, g, s} with shape conventions: + * k, v, q, g: [D] with outer dims T x H batched as ne[2]=T, ne[1]=H + * s: initial state [B, H, D, D], where B is batch and D=C/H + * dst holds both outputs (o) and updated state; a scale factor is read from op params. + * + * The kernel updates per time step l: S_new = g ⊗ S_old + k ⊗ v, then computes o = (S_new^T q) * scale. + * + * @param ctx Backend context providing stream/allocator utilities. + * @param dst Output tensor; src deps are k, v, q, g, s as above. + */ +void ggml_cann_gated_linear_attn(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Launches an asynchronous task using the memory allocator. + * + * This macro submit an asynchronous task on the specified stream. + * The task uses memory allocated by the allocator. It is guaranteed + * that the memory will not be accessed by other tasks until this task + * completes, due to the sequential execution order within the same stream. + * + * @param OP_NAME aclnn operator name. + * @param args Additional arguments required by the task. + * + * @note + * Memory from the allocator will be "freed" immediately and can be + * reallocated to other pointers. However, it won't be accessed by any + * other task before this asynchronous task ends, because all tasks in the + * same stream are executed in queue order. + */ + +# define GGML_CANN_CALL_ACLNN_OP(CTX, OP_NAME, ...) \ + do { \ + uint64_t workspaceSize = 0; \ + aclOpExecutor * executor; \ + void * workspaceAddr = nullptr; \ + ACL_CHECK(aclnn##OP_NAME##GetWorkspaceSize(__VA_ARGS__, &workspaceSize, &executor)); \ + /* workspace should alloced in main thread to keep malloc order when using vmm. */ \ + if (workspaceSize > 0) { \ + ggml_cann_pool_alloc workspace_allocator(CTX.pool(), workspaceSize); \ + workspaceAddr = workspace_allocator.get(); \ + } \ + ACL_CHECK(aclnn##OP_NAME(workspaceAddr, workspaceSize, executor, CTX.stream())); \ + } while (0) + +/** + * @brief Performs sparse expert-based matrix multiplication using the CANN backend. + * + * @details This function implements a MoE-style batched matrix multiplication, where each input token + * is routed to one or more experts, and each expert corresponds to a specific [D, M] weight matrix + * in the source tensor `src0`. The routing indices are provided via the `ids` tensor. + * + * For each token (from `src1`), the function selects the corresponding expert(s) as specified by `ids`, + * performs the matrix multiplication with the selected expert's weight submatrix (from `src0`), + * and stores the results in `dst`. This operation is optimized and executed on the CANN backend. + * + * Dimensions: + * - src0: [D, M, A, 1], where A is the number of experts + * - src1: [D, B, N, 1], where N is batch size and B is the slot count per sample + * - ids : [K, N], where K is the number of experts each token is routed to + * - dst : [M, K, N, 1], output tensor storing the result of expert × token multiplication + * + * The function handles two main modes: + * - If `ne12 == 1`, a simpler per-token loop is used. + * - TODO: If `ne12 > 1`, grouped multiplication and memory copying is used for efficiency. + * + * @param ctx The CANN context used for operations. + * @param dst The destination tensor where the expert-weighted token outputs are stored. + * Expected to be of shape [M, K, N, 1]. + */ +void ggml_cann_mul_mat_id(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Performs fused ADD + RMS_NORM operation using the CANN backend. + * + * This function fuses the ADD and RMS_NORM operations into a single kernel call + * for better performance. It first adds two input tensors (x1 + x2), then applies + * RMS normalization to the result. + * + * @param ctx The context for the CANN backend operations. + * @param dst The ADD operation node, contains the two input tensors to be added. + * @param rms_norm_tensor The RMS_NORM operation node, contains the gamma weights + * and epsilon parameter. + */ +void ggml_cann_op_add_rms_norm_fused(ggml_backend_cann_context & ctx, + ggml_tensor * add_node, + ggml_tensor * rms_norm_node); + +/** + * @brief Check whether a tensor is a weight tensor for matrix multiplication. + * + * @details Checks whether the given tensor serves as weight parameters in matrix multiplication operations, + * typically within neural network layers. The function maintains a static set of canonical weight + * naming suffixes from Transformer-based architectures. Uses substring matching to identify weight + * tensors even with hierarchical naming patterns. + * + * @param tensor Pointer to the target ggml_tensor object (const-qualified). + */ +static bool is_matmul_weight(const ggml_tensor * tensor) { + std::string name = ggml_get_name(tensor); + static const std::unordered_set weight_suffixes{ "output.weight", "attn_q.weight", + "attn_k.weight", "attn_v.weight", + "attn_output.weight", "ffn_gate.weight", + "ffn_up.weight", "ffn_down.weight" }; + + for (const auto & suffix : weight_suffixes) { + if (name.find(suffix) != std::string::npos) { + return true; + } + } + return false; +} + +/** + * @brief Applies a element-wise operation to two input tensors using the CANN + * backend. + * + * This templated function takes a binary operator and applies it to two source + * tensors + * associated with the destination tensor. The function handles broadcasting as + * needed. + * + * @tparam binary_op A callable object (e.g., lambda or function pointer) representing + * the binary operation to be performed. It must take three arguments: + * (ggml_backend_cann_context&, aclTensor*, aclTensor*, aclTensor*). + * + * @param ctx The CANN backend context used to manage execution and resources. + * @param dst The destination tensor. + */ +template void ggml_cann_binary_op(ggml_backend_cann_context & ctx, ggml_tensor * dst) { + ggml_tensor * src0 = dst->src[0]; + ggml_tensor * src1 = dst->src[1]; + + acl_tensor_ptr acl_src0, acl_src1, acl_dst; + + // Need bcast + bcast_shape(src0, src1, dst, acl_src0, acl_src1, acl_dst); + binary_op(ctx, acl_src0.get(), acl_src1.get(), acl_dst.get()); +} + +/** + * @brief Applies a unary operation to an input tensor using the CANN backend. + * + * This templated function applies a unary operator to the source tensor of `dst` + * and stores the result in the destination tensor. + * + * @tparam unary_op A callable with the signature: + * void(ggml_backend_cann_context&, aclTensor *, aclTensor *) + * where the first aclTensor is the source and the second is the destination. + * @param ctx The CANN backend context for managing resources and execution. + * @param dst The destination tensor. Its src[0] is treated as the input tensor. + */ +template +void ggml_cann_op_unary(ggml_backend_cann_context & ctx, ggml_tensor * dst) { + ggml_tensor * src = dst->src[0]; + + acl_tensor_ptr acl_src = ggml_cann_create_tensor(src); + acl_tensor_ptr acl_dst = ggml_cann_create_tensor(dst); + + unary_op(ctx, acl_src.get(), acl_dst.get()); +} + +/** + * @brief Applies a unary operation to a ggml tensor using the CANN backend. + * + * @details This function applies a unary operation to the input tensor using + * a user-provided lambda or callable `unary_op`. The lambda receives the + * CANN backend context and two ACL tensors: the source and the destination. + * + * Internally, this function handles the conversion from GGML tensors to ACL tensors, + * calls the provided unary op, and manages resource cleanup. The input is assumed + * to be `dst->src[0]`, and the result is written to `dst`. + * + * This utility simplifies writing unary op wrappers by abstracting tensor preparation. + * + * @param unary_op A callable that performs the unary operation using CANN ACL APIs. + * @param ctx The CANN context for operation execution. + * @param dst The destination ggml_tensor where the result will be stored. + * The input tensor is assumed to be `dst->src[0]`. + * + * @see GGML_CANN_CALL_OP_UNARY + */ +void ggml_cann_op_unary(std::function unary_op, + ggml_backend_cann_context & ctx, + ggml_tensor * dst); + +void ggml_cann_ssm_conv(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +/** + * @brief Applies a gated (GLU-style) unary operation using the CANN backend. + * + * @details This function performs a gated activation such as GEGLU or ReGLU. + * It supports two input modes: + * + * 1. **Dual input mode**: `dst->src[0]` and `dst->src[1]` are both valid tensors. + * These are used directly as the value and gate tensors. + * + * 2. **Packed input mode**: Only `dst->src[0]` is valid, and it is assumed to + * contain a concatenation of value and gate along the first dimension. This tensor + * will be split into two equal halves to form the value and gate inputs. + * + * The function applies a user-provided unary operation (e.g., GELU) to the value tensor, + * then multiplies the result in-place with the gate tensor: + * + * @code + * dst = unary_op(value) * gate; + * @endcode + * + * The `swapped` parameter (from `dst->op_params[1]`) allows flipping the + * order of value/gate in the packed input case. + * + * @param unary_op A callable that performs the unary operation using CANN ACL APIs. + * It receives (ctx, acl_value_tensor, acl_output_tensor). + * @param ctx The CANN context used for execution. + * @param dst The destination ggml_tensor. Source tensors are in `dst->src[0]` and optionally `src[1]`. + * + * @see GGML_CANN_CALL_OP_UNARY_GATED + */ +void ggml_cann_op_unary_gated(std::function unary_op, + ggml_backend_cann_context & ctx, + ggml_tensor * dst); + +/** + * @brief Helper macro to call a unary ACL operator via ggml_cann_op_unary. + * + * This macro wraps the specified ACLNN unary operator name into a lambda expression, + * and passes it to `ggml_cann_op_unary`, which handles the common logic for executing + * unary ops in the CANN backend. + * + * Internally, this macro expands to a lambda like: + * @code + * [](ggml_backend_cann_context& ctx, aclTensor* acl_src, aclTensor* acl_dst) { + * GGML_CANN_CALL_ACLNN_OP(ctx, OP_NAME, acl_src, acl_dst); + * }; + * @endcode + * + * This lambda is then passed to `ggml_cann_op_unary`, which applies the operation. + * + * @param OP_NAME The name of the ACL unary operator to invoke via GGML_CANN_CALL_ACLNN_OP. + * + * @see ggml_cann_op_unary + * @see GGML_CANN_CALL_ACLNN_OP + */ +# define GGML_CANN_CALL_OP_UNARY(OP_NAME) \ + do { \ + auto lambda = [](ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst) { \ + GGML_CANN_CALL_ACLNN_OP(ctx, OP_NAME, acl_src, acl_dst); \ + }; \ + ggml_cann_op_unary(lambda, ctx, dst); \ + } while (0) + +/** + * @brief Helper macro to call a gated unary ACL operator via ggml_cann_op_unary_gated. + * + * This macro wraps the specified ACLNN unary operator name into a lambda expression, + * and passes it to `ggml_cann_op_unary_gated`, which handles the common logic for + * executing gated unary ops in the CANN backend. + * + * Internally, this macro expands to a lambda like: + * @code + * [](ggml_backend_cann_context& ctx, aclTensor* acl_src, aclTensor* acl_dst) { + * GGML_CANN_CALL_ACLNN_OP(ctx, OP_NAME, acl_src, acl_dst); + * }; + * @endcode + * + * This lambda is then passed to `ggml_cann_op_unary_gated`, which applies the operation. + * + * @param OP_NAME The name of the ACL unary operator to invoke via GGML_CANN_CALL_ACLNN_OP. + * + * @see ggml_cann_op_unary_gated + * @see GGML_CANN_CALL_ACLNN_OP + */ +# define GGML_CANN_CALL_OP_UNARY_GATED(OP_NAME) \ + do { \ + auto lambda = [](ggml_backend_cann_context & ctx, aclTensor * acl_src, aclTensor * acl_dst) { \ + GGML_CANN_CALL_ACLNN_OP(ctx, OP_NAME, acl_src, acl_dst); \ + }; \ + ggml_cann_op_unary_gated(lambda, ctx, dst); \ + } while (0) + + +/** + * @brief Performs outer product operation on two ggml tensors using the CANN backend. + * + * @details This function computes the outer product of two input tensors (src0 and src1) + * and stores the result in the destination tensor. The outer product operation is defined as: + * dst[i,j,k,l] = sum_m (src0[i,m,k,l] * src1[j,m,k,l]) + * + * The function supports multiple data types including F32, F16. For floating-point + * types, it uses batch matrix multiplication for efficient computation. + * + * The implementation handles 4D tensor broadcasting and batch processing automatically. + * + * @param ctx The CANN backend context for operation execution and memory management. + * @param dst The destination ggml_tensor where the outer product result will be stored. + * The input tensors are assumed to be `dst->src[0]` and `dst->src[1]`. + * + * @see GGML_CANN_CALL_ACLNN_OP for CANN operator invocation + */ +void ggml_cann_out_prod(ggml_backend_cann_context & ctx, ggml_tensor * dst); + +#include "ggml-cann/common.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml.h.inc" #include #include @@ -46,7 +2224,7 @@ #define GGML_COMMON_DECL_C -#include "ggml-common.h" +#include "ggml-common-defs.inc" #define GGML_CANN_NAME "CANN" diff --git a/ggml/src/ggml-common.h b/ggml/src/ggml-common-defs.inc similarity index 100% rename from ggml/src/ggml-common.h rename to ggml/src/ggml-common-defs.inc diff --git a/ggml/src/ggml-cpu/CMakeLists.txt b/ggml/src/ggml-cpu/CMakeLists.txt index c1c225f01979..6ac6a0c61135 100644 --- a/ggml/src/ggml-cpu/CMakeLists.txt +++ b/ggml/src/ggml-cpu/CMakeLists.txt @@ -4,15 +4,180 @@ function(ggml_add_cpu_backend_features cpu_name arch) # Since multiple variants of the CPU backend may be included in the same # build, using set_source_files_properties() to set the arch flags is not possible set(GGML_CPU_FEATS_NAME ${cpu_name}-feats) + # Disable LTO for the feature detection code to prevent cross-module optimization + # from inlining architecture-specific instructions into the score function. + # Without this, LTO can cause SIGILL when loading backends on older CPUs + # (e.g., loading power10 backend on power9 crashes before feature check runs). + if (${arch} STREQUAL "x86") + set(GGML_CPU_FEATS_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/ggml-cpu/arch/x86/cpu-feats-rust/Cargo.toml) + set(GGML_CPU_FEATS_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/${cpu_name}-x86-feats) + set(GGML_CPU_FEATS_RUST_LIB ${GGML_CPU_FEATS_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}ggml_cpu_x86_feats_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) + set(GGML_CPU_FEATS_RUSTFLAGS "-Crelocation-model=pic -Cpanic=abort") + foreach(def ${ARGN}) + if (${def} STREQUAL "GGML_FMA") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_fma") + elseif (${def} STREQUAL "GGML_F16C") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_f16c") + elseif (${def} STREQUAL "GGML_SSE42") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_sse42") + elseif (${def} STREQUAL "GGML_BMI2") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_bmi2") + elseif (${def} STREQUAL "GGML_AVX") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_avx") + elseif (${def} STREQUAL "GGML_AVX2") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_avx2") + elseif (${def} STREQUAL "GGML_AVX_VNNI") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_avx_vnni") + elseif (${def} STREQUAL "GGML_AVX512") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_avx512") + elseif (${def} STREQUAL "GGML_AVX512_VBMI") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_avx512_vbmi") + elseif (${def} STREQUAL "GGML_AVX512_BF16") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_avx512_bf16") + elseif (${def} STREQUAL "GGML_AVX512_VNNI") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_avx512_vnni") + elseif (${def} STREQUAL "GGML_AMX_INT8") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_amx_int8") + endif() + endforeach() + add_custom_command( + OUTPUT ${GGML_CPU_FEATS_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=${GGML_CPU_FEATS_RUSTFLAGS}" + ${CARGO_EXECUTABLE} build --manifest-path ${GGML_CPU_FEATS_RUST_MANIFEST} --target-dir ${GGML_CPU_FEATS_RUST_TARGET_DIR} + DEPENDS + ${GGML_CPU_FEATS_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/ggml-cpu/arch/x86/cpu-feats-rust/src/lib.rs + ) + add_custom_target(${GGML_CPU_FEATS_NAME}-rust-build DEPENDS ${GGML_CPU_FEATS_RUST_LIB}) + add_library(${GGML_CPU_FEATS_NAME}-rust STATIC IMPORTED GLOBAL) + set_target_properties(${GGML_CPU_FEATS_NAME}-rust PROPERTIES IMPORTED_LOCATION ${GGML_CPU_FEATS_RUST_LIB}) + add_dependencies(${GGML_CPU_FEATS_NAME}-rust ${GGML_CPU_FEATS_NAME}-rust-build) + target_link_libraries(${cpu_name} PRIVATE "$") + return() + elseif (${arch} STREQUAL "arm") + set(GGML_CPU_FEATS_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/ggml-cpu/arch/arm/cpu-feats-rust/Cargo.toml) + set(GGML_CPU_FEATS_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/${cpu_name}-arm-feats) + set(GGML_CPU_FEATS_RUST_LIB ${GGML_CPU_FEATS_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}ggml_cpu_arm_feats_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) + set(GGML_CPU_FEATS_RUSTFLAGS "-Crelocation-model=pic -Cpanic=abort") + foreach(def ${ARGN}) + if (${def} STREQUAL "GGML_USE_DOTPROD") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_dotprod") + elseif (${def} STREQUAL "GGML_USE_FP16_VECTOR_ARITHMETIC") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_fp16_va") + elseif (${def} STREQUAL "GGML_USE_SVE") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_sve") + elseif (${def} STREQUAL "GGML_USE_MATMUL_INT8") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_matmul_int8") + elseif (${def} STREQUAL "GGML_USE_SVE2") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_sve2") + elseif (${def} STREQUAL "GGML_USE_SME") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_sme") + endif() + endforeach() + add_custom_command( + OUTPUT ${GGML_CPU_FEATS_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=${GGML_CPU_FEATS_RUSTFLAGS}" + ${CARGO_EXECUTABLE} build --manifest-path ${GGML_CPU_FEATS_RUST_MANIFEST} --target-dir ${GGML_CPU_FEATS_RUST_TARGET_DIR} + DEPENDS + ${GGML_CPU_FEATS_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/ggml-cpu/arch/arm/cpu-feats-rust/src/lib.rs + ) + add_custom_target(${GGML_CPU_FEATS_NAME}-rust-build DEPENDS ${GGML_CPU_FEATS_RUST_LIB}) + add_library(${GGML_CPU_FEATS_NAME}-rust STATIC IMPORTED GLOBAL) + set_target_properties(${GGML_CPU_FEATS_NAME}-rust PROPERTIES IMPORTED_LOCATION ${GGML_CPU_FEATS_RUST_LIB}) + add_dependencies(${GGML_CPU_FEATS_NAME}-rust ${GGML_CPU_FEATS_NAME}-rust-build) + target_link_libraries(${cpu_name} PRIVATE "$") + return() + elseif (${arch} STREQUAL "riscv") + set(GGML_CPU_FEATS_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/ggml-cpu/arch/riscv/cpu-feats-rust/Cargo.toml) + set(GGML_CPU_FEATS_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/${cpu_name}-riscv-feats) + set(GGML_CPU_FEATS_RUST_LIB ${GGML_CPU_FEATS_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}ggml_cpu_riscv_feats_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) + set(GGML_CPU_FEATS_RUSTFLAGS "-Crelocation-model=pic -Cpanic=abort") + foreach(def ${ARGN}) + if (${def} STREQUAL "GGML_USE_RVV") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_rvv") + endif() + endforeach() + add_custom_command( + OUTPUT ${GGML_CPU_FEATS_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=${GGML_CPU_FEATS_RUSTFLAGS}" + ${CARGO_EXECUTABLE} build --manifest-path ${GGML_CPU_FEATS_RUST_MANIFEST} --target-dir ${GGML_CPU_FEATS_RUST_TARGET_DIR} + DEPENDS + ${GGML_CPU_FEATS_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/ggml-cpu/arch/riscv/cpu-feats-rust/src/lib.rs + ) + add_custom_target(${GGML_CPU_FEATS_NAME}-rust-build DEPENDS ${GGML_CPU_FEATS_RUST_LIB}) + add_library(${GGML_CPU_FEATS_NAME}-rust STATIC IMPORTED GLOBAL) + set_target_properties(${GGML_CPU_FEATS_NAME}-rust PROPERTIES IMPORTED_LOCATION ${GGML_CPU_FEATS_RUST_LIB}) + add_dependencies(${GGML_CPU_FEATS_NAME}-rust ${GGML_CPU_FEATS_NAME}-rust-build) + target_link_libraries(${cpu_name} PRIVATE "$") + return() + elseif (${arch} STREQUAL "s390") + set(GGML_CPU_FEATS_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/ggml-cpu/arch/s390/cpu-feats-rust/Cargo.toml) + set(GGML_CPU_FEATS_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/${cpu_name}-s390-feats) + set(GGML_CPU_FEATS_RUST_LIB ${GGML_CPU_FEATS_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}ggml_cpu_s390_feats_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) + set(GGML_CPU_FEATS_RUSTFLAGS "-Crelocation-model=pic -Cpanic=abort") + foreach(def ${ARGN}) + if (${def} STREQUAL "GGML_USE_VXE2") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_vxe2") + elseif (${def} STREQUAL "GGML_USE_NNPA") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_nnpa") + endif() + endforeach() + add_custom_command( + OUTPUT ${GGML_CPU_FEATS_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=${GGML_CPU_FEATS_RUSTFLAGS}" + ${CARGO_EXECUTABLE} build --manifest-path ${GGML_CPU_FEATS_RUST_MANIFEST} --target-dir ${GGML_CPU_FEATS_RUST_TARGET_DIR} + DEPENDS + ${GGML_CPU_FEATS_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/ggml-cpu/arch/s390/cpu-feats-rust/src/lib.rs + ) + add_custom_target(${GGML_CPU_FEATS_NAME}-rust-build DEPENDS ${GGML_CPU_FEATS_RUST_LIB}) + add_library(${GGML_CPU_FEATS_NAME}-rust STATIC IMPORTED GLOBAL) + set_target_properties(${GGML_CPU_FEATS_NAME}-rust PROPERTIES IMPORTED_LOCATION ${GGML_CPU_FEATS_RUST_LIB}) + add_dependencies(${GGML_CPU_FEATS_NAME}-rust ${GGML_CPU_FEATS_NAME}-rust-build) + target_link_libraries(${cpu_name} PRIVATE "$") + return() + elseif (${arch} STREQUAL "powerpc") + set(GGML_CPU_FEATS_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/ggml-cpu/arch/powerpc/cpu-feats-rust/Cargo.toml) + set(GGML_CPU_FEATS_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/${cpu_name}-powerpc-feats) + set(GGML_CPU_FEATS_RUST_LIB ${GGML_CPU_FEATS_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}ggml_cpu_powerpc_feats_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) + set(GGML_CPU_FEATS_RUSTFLAGS "-Crelocation-model=pic -Cpanic=abort") + foreach(def ${ARGN}) + if (${def} STREQUAL "GGML_USE_POWER7") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_power7") + elseif (${def} STREQUAL "GGML_USE_POWER8") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_power8") + elseif (${def} STREQUAL "GGML_USE_POWER9") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_power9") + elseif (${def} STREQUAL "GGML_USE_POWER10") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_power10") + elseif (${def} STREQUAL "GGML_USE_POWER11") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_power11") + elseif (${def} STREQUAL "GGML_USE_VSX") + string(APPEND GGML_CPU_FEATS_RUSTFLAGS " --cfg=ggml_use_vsx") + endif() + endforeach() + add_custom_command( + OUTPUT ${GGML_CPU_FEATS_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=${GGML_CPU_FEATS_RUSTFLAGS}" + ${CARGO_EXECUTABLE} build --manifest-path ${GGML_CPU_FEATS_RUST_MANIFEST} --target-dir ${GGML_CPU_FEATS_RUST_TARGET_DIR} + DEPENDS + ${GGML_CPU_FEATS_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/ggml-cpu/arch/powerpc/cpu-feats-rust/src/lib.rs + ) + add_custom_target(${GGML_CPU_FEATS_NAME}-rust-build DEPENDS ${GGML_CPU_FEATS_RUST_LIB}) + add_library(${GGML_CPU_FEATS_NAME}-rust STATIC IMPORTED GLOBAL) + set_target_properties(${GGML_CPU_FEATS_NAME}-rust PROPERTIES IMPORTED_LOCATION ${GGML_CPU_FEATS_RUST_LIB}) + add_dependencies(${GGML_CPU_FEATS_NAME}-rust ${GGML_CPU_FEATS_NAME}-rust-build) + target_link_libraries(${cpu_name} PRIVATE "$") + return() + endif() add_library(${GGML_CPU_FEATS_NAME} OBJECT ggml-cpu/arch/${arch}/cpu-feats.cpp) target_include_directories(${GGML_CPU_FEATS_NAME} PRIVATE . ../include) target_compile_definitions(${GGML_CPU_FEATS_NAME} PRIVATE ${ARGN}) target_compile_definitions(${GGML_CPU_FEATS_NAME} PRIVATE GGML_BACKEND_DL GGML_BACKEND_BUILD GGML_BACKEND_SHARED) set_target_properties(${GGML_CPU_FEATS_NAME} PROPERTIES POSITION_INDEPENDENT_CODE ON) - # Disable LTO for the feature detection code to prevent cross-module optimization - # from inlining architecture-specific instructions into the score function. - # Without this, LTO can cause SIGILL when loading backends on older CPUs - # (e.g., loading power10 backend on power9 crashes before feature check runs). target_compile_options(${GGML_CPU_FEATS_NAME} PRIVATE -fno-lto) target_link_libraries(${cpu_name} PRIVATE ${GGML_CPU_FEATS_NAME}) endfunction() @@ -27,31 +192,13 @@ function(ggml_add_cpu_backend_variant_impl tag_name) ggml_add_backend_library(${GGML_CPU_NAME}) list (APPEND GGML_CPU_SOURCES - ggml-cpu/ggml-cpu.c - ggml-cpu/ggml-cpu.cpp - ggml-cpu/repack.cpp - ggml-cpu/repack.h - ggml-cpu/hbm.cpp - ggml-cpu/hbm.h - ggml-cpu/quants.c - ggml-cpu/quants.h - ggml-cpu/traits.cpp - ggml-cpu/traits.h - ggml-cpu/amx/amx.cpp - ggml-cpu/amx/amx.h - ggml-cpu/amx/mmq.cpp - ggml-cpu/amx/mmq.h - ggml-cpu/ggml-cpu-impl.h - ggml-cpu/common.h - ggml-cpu/binary-ops.h - ggml-cpu/binary-ops.cpp - ggml-cpu/unary-ops.h - ggml-cpu/unary-ops.cpp - ggml-cpu/simd-mappings.h - ggml-cpu/vec.h - ggml-cpu/vec.cpp - ggml-cpu/ops.h - ggml-cpu/ops.cpp + ggml-cpu/ggml-cpu.c.inc + ggml-cpu/ggml-cpu.cpp.inc + ggml-cpu/repack.cpp.inc + ggml-cpu/quants.c.inc + ggml-cpu/ggml-cpu-impl-defs.inc + ggml-cpu/simd-mappings-defs.inc + ggml-cpu/ops.cpp.inc ) target_compile_features(${GGML_CPU_NAME} PRIVATE c_std_11 cxx_std_17) @@ -87,10 +234,6 @@ function(ggml_add_cpu_backend_variant_impl tag_name) if (GGML_LLAMAFILE) target_compile_definitions(${GGML_CPU_NAME} PRIVATE GGML_USE_LLAMAFILE) - - list(APPEND GGML_CPU_SOURCES - ggml-cpu/llamafile/sgemm.cpp - ggml-cpu/llamafile/sgemm.h) endif() if (GGML_CPU_HBM) @@ -106,8 +249,8 @@ function(ggml_add_cpu_backend_variant_impl tag_name) if (GGML_SYSTEM_ARCH STREQUAL "ARM") message(STATUS "ARM detected") list(APPEND GGML_CPU_SOURCES - ggml-cpu/arch/arm/quants.c - ggml-cpu/arch/arm/repack.cpp + ggml-cpu/arch/arm/quants.c.inc + ggml-cpu/arch/arm/repack.cpp.inc ) if (MSVC AND NOT CMAKE_C_COMPILER_ID STREQUAL "Clang") @@ -247,10 +390,6 @@ function(ggml_add_cpu_backend_variant_impl tag_name) endif() elseif (GGML_SYSTEM_ARCH STREQUAL "x86") message(STATUS "x86 detected") - list(APPEND GGML_CPU_SOURCES - ggml-cpu/arch/x86/quants.c - ggml-cpu/arch/x86/repack.cpp - ) if (MSVC) # instruction set detection for MSVC only @@ -385,7 +524,7 @@ function(ggml_add_cpu_backend_variant_impl tag_name) endif() elseif (GGML_SYSTEM_ARCH STREQUAL "PowerPC") message(STATUS "PowerPC detected") - list(APPEND GGML_CPU_SOURCES ggml-cpu/arch/powerpc/quants.c) + list(APPEND GGML_CPU_SOURCES ggml-cpu/arch/powerpc/quants.c.inc) if (GGML_NATIVE) if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "ppc64") file(READ "/proc/cpuinfo" POWER10_M) @@ -434,7 +573,7 @@ function(ggml_add_cpu_backend_variant_impl tag_name) endif() elseif (GGML_SYSTEM_ARCH STREQUAL "loongarch64") message(STATUS "loongarch64 detected") - list(APPEND GGML_CPU_SOURCES ggml-cpu/arch/loongarch/quants.c) + list(APPEND GGML_CPU_SOURCES ggml-cpu/arch/loongarch/quants.c.inc) list(APPEND ARCH_FLAGS -march=loongarch64) if (GGML_LASX) @@ -446,16 +585,14 @@ function(ggml_add_cpu_backend_variant_impl tag_name) elseif (GGML_SYSTEM_ARCH STREQUAL "riscv64") message(STATUS "riscv64 detected") list(APPEND GGML_CPU_SOURCES - ggml-cpu/arch/riscv/quants.c - ggml-cpu/arch/riscv/repack.cpp + ggml-cpu/arch/riscv/quants.c.inc + ggml-cpu/arch/riscv/repack.cpp.inc ) if (GGML_CPU_RISCV64_SPACEMIT) target_compile_definitions(${GGML_CPU_NAME} PRIVATE GGML_USE_CPU_RISCV64_SPACEMIT ${RISCV64_SPACEMIT_IME_SPEC}) list(APPEND GGML_CPU_SOURCES - ggml-cpu/spacemit/ime.cpp - ggml-cpu/spacemit/ime.h - ggml-cpu/spacemit/ime1_kernels.cpp - ggml-cpu/spacemit/ime_kernels.h + ggml-cpu/spacemit/ime.cpp.inc + ggml-cpu/spacemit/ime1_kernels.cpp.inc ) endif() if(NOT GGML_CPU_ALL_VARIANTS) @@ -509,7 +646,7 @@ function(ggml_add_cpu_backend_variant_impl tag_name) elseif (GGML_SYSTEM_ARCH STREQUAL "s390x") message(STATUS "s390x detected") list(APPEND GGML_CPU_SOURCES - ggml-cpu/arch/s390/quants.c) + ggml-cpu/arch/s390/quants.c.inc) # for native compilation if (GGML_NATIVE) @@ -560,7 +697,7 @@ function(ggml_add_cpu_backend_variant_impl tag_name) ggml_add_cpu_backend_features(${GGML_CPU_NAME} s390 ${ARCH_DEFINITIONS}) elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "wasm") message(STATUS "Wasm detected") - list (APPEND GGML_CPU_SOURCES ggml-cpu/arch/wasm/quants.c) + list (APPEND GGML_CPU_SOURCES ggml-cpu/arch/wasm/quants.c.inc) else() message(WARNING "Unknown CPU architecture. Falling back to generic implementations.") list(APPEND ARCH_FLAGS -DGGML_CPU_GENERIC) @@ -617,10 +754,8 @@ function(ggml_add_cpu_backend_variant_impl tag_name) add_compile_definitions(GGML_USE_CPU_KLEIDIAI) list(APPEND GGML_CPU_SOURCES - ggml-cpu/kleidiai/kleidiai.cpp - ggml-cpu/kleidiai/kernels.cpp - ggml-cpu/kleidiai/kleidiai.h - ggml-cpu/kleidiai/kernels.h + ggml-cpu/kleidiai/kleidiai.cpp.inc + ggml-cpu/kleidiai/kernels.cpp.inc ) # KleidiAI @@ -703,6 +838,15 @@ function(ggml_add_cpu_backend_variant_impl tag_name) endif() message(STATUS "Adding CPU backend variant ${GGML_CPU_NAME}: ${ARCH_FLAGS} ${ARCH_DEFINITIONS}") + foreach(source ${GGML_CPU_SOURCES}) + if (source MATCHES "\\.c\\.inc$") + set_source_files_properties(${source} PROPERTIES LANGUAGE C) + elseif (source MATCHES "\\.cpp\\.inc$") + set_source_files_properties(${source} PROPERTIES LANGUAGE CXX) + elseif (source MATCHES "-defs\\.inc$") + set_source_files_properties(${source} PROPERTIES HEADER_FILE_ONLY TRUE) + endif() + endforeach() target_sources(${GGML_CPU_NAME} PRIVATE ${GGML_CPU_SOURCES}) target_compile_options(${GGML_CPU_NAME} PRIVATE ${ARCH_FLAGS}) target_compile_definitions(${GGML_CPU_NAME} PRIVATE ${ARCH_DEFINITIONS}) diff --git a/ggml/src/ggml-cpu/amx/amx.cpp b/ggml/src/ggml-cpu/amx/amx.cpp deleted file mode 100644 index 1118f7169c92..000000000000 --- a/ggml/src/ggml-cpu/amx/amx.cpp +++ /dev/null @@ -1,249 +0,0 @@ -#include "amx.h" -#include "common.h" -#include "mmq.h" -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "traits.h" - -#if defined(__linux__) -#include -#include -#endif - -#include -#include -#include - -#if defined(__AMX_INT8__) && defined(__AVX512VNNI__) - -// AMX type_trais -namespace ggml::cpu::amx { -class tensor_traits : public ggml::cpu::tensor_traits { - bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override { - size = ggml_backend_amx_desired_wsize(op); - return true; - } - - bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) override { - if (op->op == GGML_OP_MUL_MAT) { - ggml_backend_amx_mul_mat(params, op); - return true; - } - return false; - } -}; - -static ggml::cpu::tensor_traits * get_tensor_traits(ggml_backend_buffer_t, struct ggml_tensor *) { - static tensor_traits traits; - return &traits; -} -} // namespace ggml::cpu::amx - -// AMX buffer interface -static void ggml_backend_amx_buffer_free_buffer(ggml_backend_buffer_t buffer) { - free(buffer->context); -} - -static void * ggml_backend_amx_buffer_get_base(ggml_backend_buffer_t buffer) { - return (void *) (buffer->context); -} - -static enum ggml_status ggml_backend_amx_buffer_init_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor) { - tensor->extra = (void *) ggml::cpu::amx::get_tensor_traits(buffer, tensor); - - GGML_UNUSED(buffer); - return GGML_STATUS_SUCCESS; -} - -static void ggml_backend_amx_buffer_memset_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, - uint8_t value, size_t offset, size_t size) { - memset((char *) tensor->data + offset, value, size); - - GGML_UNUSED(buffer); -} - -static void ggml_backend_amx_buffer_set_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, - const void * data, size_t offset, size_t size) { - if (qtype_has_amx_kernels(tensor->type)) { - GGML_LOG_DEBUG("%s: amx repack tensor %s of type %s\n", __func__, tensor->name, ggml_type_name(tensor->type)); - ggml_backend_amx_convert_weight(tensor, data, offset, size); - } else { - memcpy((char *) tensor->data + offset, data, size); - } - - GGML_UNUSED(buffer); -} - -/* -// need to figure what we need to do with buffer->extra. -static void ggml_backend_amx_buffer_get_tensor(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) { - GGML_ASSERT(!qtype_has_amx_kernels(tensor->type)); - memcpy(data, (const char *)tensor->data + offset, size); - - GGML_UNUSED(buffer); -} - -static bool ggml_backend_amx_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const struct ggml_tensor * src, struct ggml_tensor * dst) { - if (ggml_backend_buffer_is_host(src->buffer)) { - if (qtype_has_amx_kernels(src->type)) { - ggml_backend_amx_convert_weight(dst, src->data, 0, ggml_nbytes(dst)); - } else { - memcpy(dst->data, src->data, ggml_nbytes(src)); - } - return true; - } - return false; - - GGML_UNUSED(buffer); -} -*/ - -static void ggml_backend_amx_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { - memset(buffer->context, value, buffer->size); -} - -static ggml_backend_buffer_i ggml_backend_amx_buffer_interface = { - /* .free_buffer = */ ggml_backend_amx_buffer_free_buffer, - /* .get_base = */ ggml_backend_amx_buffer_get_base, - /* .init_tensor = */ ggml_backend_amx_buffer_init_tensor, - /* .memset_tensor = */ ggml_backend_amx_buffer_memset_tensor, - /* .set_tensor = */ ggml_backend_amx_buffer_set_tensor, - /* .get_tensor = */ nullptr, - /* .set_tensor_2d = */ nullptr, - /* .get_tensor_2d = */ nullptr, - /* .cpy_tensor = */ nullptr, - /* .clear = */ ggml_backend_amx_buffer_clear, - /* .reset = */ nullptr, -}; - -static const char * ggml_backend_amx_buffer_type_get_name(ggml_backend_buffer_type_t buft) { - return "AMX"; - - GGML_UNUSED(buft); -} - -static ggml_backend_buffer_t ggml_backend_amx_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { - void * data = ggml_aligned_malloc(size); - if (data == NULL) { - fprintf(stderr, "%s: failed to allocate buffer of size %zu\n", __func__, size); - return NULL; - } - - return ggml_backend_buffer_init(buft, ggml_backend_amx_buffer_interface, data, size); -} - -static size_t ggml_backend_amx_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { - return TENSOR_ALIGNMENT; - - GGML_UNUSED(buft); -} - -namespace ggml::cpu::amx { -class extra_buffer_type : ggml::cpu::extra_buffer_type { - bool supports_op(ggml_backend_dev_t, const struct ggml_tensor * op) override { - if (op->op != GGML_OP_MUL_MAT) { - return false; - } - auto * src0 = op->src[0]; - auto * src1 = op->src[1]; - - if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(src1)) { - return false; - } - if (!src0->buffer || src0->buffer->buft != ggml_backend_amx_buffer_type()) { - return false; - } - if (src1->buffer && !ggml_backend_buft_is_host(src1->buffer->buft)) { - return false; - } - if (op->ne[0] % (TILE_N * 2)) { - return false; - } - int alignment; - switch (src0->type) { - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q8_0: - alignment = TILE_K; - break; - case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: - case GGML_TYPE_Q6_K: - case GGML_TYPE_IQ4_XS: - alignment = 256; // QK_K - break; - case GGML_TYPE_F16: - alignment = 16; - break; - default: - return false; - } - if (src0->ne[0] % alignment) { - return false; - } - if (src1->type != GGML_TYPE_F32) { - return false; - } - return true; - } - - ggml::cpu::tensor_traits * get_tensor_traits(const struct ggml_tensor * op) override { - if (op->op == GGML_OP_MUL_MAT && op->src[0]->buffer && - op->src[0]->buffer->buft == ggml_backend_amx_buffer_type()) { - return (ggml::cpu::tensor_traits *) op->src[0]->extra; - } - - return nullptr; - } -}; -} // namespace ggml::cpu::amx - -static size_t ggml_backend_amx_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { - return ggml_backend_amx_get_alloc_size(tensor); - - GGML_UNUSED(buft); -} - -#define ARCH_GET_XCOMP_PERM 0x1022 -#define ARCH_REQ_XCOMP_PERM 0x1023 -#define XFEATURE_XTILECFG 17 -#define XFEATURE_XTILEDATA 18 - -static bool ggml_amx_init() { -#if defined(__linux__) - if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA)) { - fprintf(stderr, "AMX is not ready to be used!\n"); - return false; - } - return true; -#elif defined(_WIN32) - return true; -#else - return false; -#endif -} - -ggml_backend_buffer_type_t ggml_backend_amx_buffer_type() { - static struct ggml_backend_buffer_type ggml_backend_buffer_type_amx = { - /* .iface = */ { - /* .get_name = */ ggml_backend_amx_buffer_type_get_name, - /* .alloc_buffer = */ ggml_backend_amx_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_amx_buffer_type_get_alignment, - /* .get_max_size = */ nullptr, // defaults to SIZE_MAX - /* .get_alloc_size = */ ggml_backend_amx_buffer_type_get_alloc_size, - /* .is_host = */ nullptr, - }, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0), - /* .context = */ new ggml::cpu::amx::extra_buffer_type(), - }; - - if (!ggml_amx_init()) { - return nullptr; - } - - return &ggml_backend_buffer_type_amx; -} - -#endif // defined(__AMX_INT8__) && defined(__AVX512VNNI__) diff --git a/ggml/src/ggml-cpu/amx/amx.h b/ggml/src/ggml-cpu/amx/amx.h deleted file mode 100644 index 5b65d76bdc89..000000000000 --- a/ggml/src/ggml-cpu/amx/amx.h +++ /dev/null @@ -1,8 +0,0 @@ -#include "ggml-backend.h" -#include "ggml-cpu-impl.h" - -// GGML internal header - -#if defined(__AMX_INT8__) && defined(__AVX512VNNI__) -ggml_backend_buffer_type_t ggml_backend_amx_buffer_type(void); -#endif diff --git a/ggml/src/ggml-cpu/amx/common.h b/ggml/src/ggml-cpu/amx/common.h deleted file mode 100644 index 26a6ec1a2d00..000000000000 --- a/ggml/src/ggml-cpu/amx/common.h +++ /dev/null @@ -1,115 +0,0 @@ -#pragma once - -#include "ggml.h" -#include "ggml-cpu-impl.h" - -#include -#include -#include - -#if defined(GGML_USE_OPENMP) -#include -#else -#include -#endif - -#define TILE_M 16 -#define TILE_N 16 -#define TILE_K 32 -#define VNNI_BLK 4 - -#define AMX_BLK_SIZE 32 - -#define TMM0 0 -#define TMM1 1 -#define TMM2 2 -#define TMM3 3 -#define TMM4 4 -#define TMM5 5 -#define TMM6 6 -#define TMM7 7 - -// parallel routines -template ::value, int>::type = 0> -inline T div_up(T x, T y) { return (x + y - 1) / y; } - -template -inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) { -#if 0 - // onednn partition pattern - T& n_my = n_end; - if (nth <= 1 || n == 0) { - n_start = 0; - n_my = n; - } else { - T n1 = div_up(n, nth); - T n2 = n1 - 1; - T T1 = n - n2 * nth; - n_my = ith < T1 ? n1 : n2; - n_start = ith <= T1 ? ith*n1 : T1 * n1 + (ith - T1) * n2; - } - n_end += n_start; -#else - // pytorch aten partition pattern - T n_my = div_up(n, nth); - n_start = ith * n_my; - n_end = std::min(n_start + n_my, n); -#endif -} - -template -inline void parallel_for(int n, const func_t & f) { - if (n <= 0) { - return; - } -#if defined(GGML_USE_OPENMP) - #pragma omp parallel - { - int nth = omp_get_num_threads(); - int ith = omp_get_thread_num(); - int tbegin, tend; - balance211(n, nth, ith, tbegin, tend); - f(tbegin, tend); - } -#else - int nth = std::thread::hardware_concurrency(); - if (nth <= 1) { - f(0, n); - return; - } - if (nth > n) { - nth = n; - } - std::vector threads; - threads.reserve(nth); - for (int ith = 0; ith < nth; ++ith) { - threads.emplace_back([&f, n, ith, nth] { - int tbegin, tend; - balance211(n, nth, ith, tbegin, tend); - f(tbegin, tend); - }); - } - for (auto & t : threads) { - t.join(); - } -#endif -} - -template -inline void parallel_for_ggml(const ggml_compute_params * params, int n, const func_t & f) { - int tbegin, tend; - balance211(n, params->nth, params->ith, tbegin, tend); - f(tbegin, tend); -} - -// quantized types that have AMX support -inline bool qtype_has_amx_kernels(const enum ggml_type type) { - // TODO: fix padding for vnni format - return (type == GGML_TYPE_Q4_0) || - (type == GGML_TYPE_Q4_1) || - (type == GGML_TYPE_Q8_0) || - (type == GGML_TYPE_Q4_K) || - (type == GGML_TYPE_Q5_K) || - (type == GGML_TYPE_Q6_K) || - (type == GGML_TYPE_IQ4_XS); -} diff --git a/ggml/src/ggml-cpu/amx/mmq.cpp b/ggml/src/ggml-cpu/amx/mmq.cpp deleted file mode 100644 index d9383a04be83..000000000000 --- a/ggml/src/ggml-cpu/amx/mmq.cpp +++ /dev/null @@ -1,2512 +0,0 @@ -#if defined(__GNUC__) -#pragma GCC diagnostic ignored "-Wpedantic" -#pragma GCC diagnostic ignored "-Wunused-local-typedefs" -#endif - -#include "amx.h" -#include "mmq.h" -#include "ggml-impl.h" -#include "ggml-cpu-impl.h" -#include "simd-mappings.h" -#include "quants.h" -#include "ggml-quants.h" -#include -#include - -#if defined(__gnu_linux__) -#include -#include -#endif - -#if (defined(_WIN32) || defined(_WIN64)) -#define RESTRICT __restrict -#else -#define RESTRICT __restrict__ -#endif - -#if (defined(_WIN32) || defined(_WIN64)) -#define ALWAYS_INLINE __forceinline -#elif __has_attribute(always_inline) || defined(__GNUC__) -#define ALWAYS_INLINE __attribute__((__always_inline__)) inline -#else -#define ALWAYS_INLINE inline -#endif - -#if defined(__AMX_INT8__) && defined(__AVX512VNNI__) - -namespace { - -// Forced unrolling -template -struct Unroll { - template - ALWAYS_INLINE void operator()(const Func& f, Args... args) const { - Unroll{}(f, args...); - f(std::integral_constant{}, args...); - } -}; - -template <> -struct Unroll<1> { - template - ALWAYS_INLINE void operator()(const Func& f, Args... args) const { - f(std::integral_constant{}, args...); - } -}; - -// type traits -template struct PackedTypes {}; -template <> struct PackedTypes { using type = int8_t; }; -template <> struct PackedTypes { using type = uint8_t; }; -template <> struct PackedTypes { using type = int8_t; }; -template using packed_B_type = typename PackedTypes::type; - -template -struct do_compensate : std::integral_constant::value> {}; - -template -struct do_unpack : std::integral_constant::value || - std::is_same::value> {}; - -template -struct is_type_qkk : std::integral_constant::value || - std::is_same::value || - std::is_same::value || - std::is_same::value> {}; - -#define GGML_DISPATCH_FLOATING_TYPES(TYPE, ...) \ - [&] { \ - switch (TYPE) { \ - case GGML_TYPE_F16: { \ - using type = ggml_fp16_t; \ - constexpr int blck_size = 16; \ - return __VA_ARGS__(); \ - } \ - case GGML_TYPE_BF16: { \ - using type = ggml_bf16_t; \ - constexpr int blck_size = 32; \ - return __VA_ARGS__(); \ - } \ - default: \ - fprintf(stderr, "Unsupported floating data type\n"); \ - } \ - }() - -#define GGML_DISPATCH_QTYPES(QT, ...) \ - [&] { \ - switch (QT) { \ - case GGML_TYPE_Q4_0: { \ - using type = block_q4_0; \ - using vec_dot_type = block_q8_0; \ - constexpr int blck_size = QK4_0; \ - return __VA_ARGS__(); \ - } \ - case GGML_TYPE_Q4_1: { \ - using type = block_q4_1; \ - using vec_dot_type = block_q8_1; \ - constexpr int blck_size = QK4_1; \ - return __VA_ARGS__(); \ - } \ - case GGML_TYPE_Q8_0: { \ - using type = block_q8_0; \ - using vec_dot_type = block_q8_0; \ - constexpr int blck_size = QK8_0; \ - return __VA_ARGS__(); \ - } \ - case GGML_TYPE_Q4_K: { \ - using type = block_q4_K; \ - using vec_dot_type = block_q8_K; \ - constexpr int blck_size = QK_K; \ - return __VA_ARGS__(); \ - } \ - case GGML_TYPE_Q5_K: { \ - using type = block_q5_K; \ - using vec_dot_type = block_q8_K; \ - constexpr int blck_size = QK_K; \ - return __VA_ARGS__(); \ - } \ - case GGML_TYPE_Q6_K: { \ - using type = block_q6_K; \ - using vec_dot_type = block_q8_K; \ - constexpr int blck_size = QK_K; \ - return __VA_ARGS__(); \ - } \ - case GGML_TYPE_IQ4_XS: { \ - using type = block_iq4_xs; \ - using vec_dot_type = block_q8_K; \ - constexpr int blck_size = QK_K; \ - return __VA_ARGS__(); \ - } \ - default: \ - fprintf(stderr, "Unsupported quantized data type: %d\n", int(TYPE)); \ - } \ - }() - -#define GGML_DISPATCH_BOOL(BOOL_V, BOOL_NAME, ...) \ - [&] { \ - if (BOOL_V) { \ - constexpr bool BOOL_NAME = true; \ - return __VA_ARGS__(); \ - } else { \ - constexpr bool BOOL_NAME = false; \ - return __VA_ARGS__(); \ - } \ - }() - -// define amx tile config data structure -struct tile_config_t{ - uint8_t palette_id = 0; - uint8_t start_row = 0; - uint8_t reserved_0[14] = {0}; - uint16_t colsb[16] = {0}; - uint8_t rows[16] = {0}; -}; - -// Notes: amx tile config -// -// Typically, TMUL calculates A and B of size 16 x 64 containing INT8 values, -// and accumulate the result to a 16 x 16 matrix C containing INT32 values, -// -// As many GGUF quantized types as `block_size` of 32, so a 16-16-32 config is used -// instead of the normally used 16-16-64 config. -// -// Block A: {16, 32}, dtype = int8_t -// Block B: {16, 32}, dtype = uint8_t/int8_t -// Block C: {16, 16}, dtype = int32_t -// -// Block B needs to be prepacked to vnni format before feeding into TMUL: -// packed_B: from {n, k} to {k/vnni_blk, n, vnni_blck}, viewed in 2d, we get {8, 64} -// -// Therefore, we get tileconfig: -// A B C -// rows 16 8 16 -// colsb 32 64 16 -// -// For tile distribution, follow a 2-2-4 pattern, e.g. A used TMM2-TMM3, B used TMM0-TMM1, -// C used TMM4-TMM7: -// B TMM0 B TMM1 -// A TMM2 C TMM4 C TMM6 -// A TMM3 C TMM5 C TMM7 -// -// Each `amx` kernel handles 4 blocks at a time: 2MB * 2NB, when m < 2 * BLOCK_M, unpack A -// will be needed. -// -// Here another commonly used pattern 1-3-3 is skipped, as it is mostly used when m <=16; -// and the single batch gemm (m=1) has a special fast path with `avx512-vnni`. -// -// ref: https://www.intel.com/content/www/us/en/developer/articles/code-sample/ -// advanced-matrix-extensions-intrinsics-functions.html -// - -inline void ggml_tile_config_init(void) { - static thread_local bool done = false; - - if (done) { - return; - } - - alignas(64) tile_config_t tc = {}; - tc.palette_id = 1; - tc.start_row = 0; - tc.rows[0] = 8; tc.colsb[0] = 64; - tc.rows[1] = 8; tc.colsb[1] = 64; - tc.rows[2] = 16; tc.colsb[2] = 32; - tc.rows[3] = 16; tc.colsb[3] = 32; - tc.rows[4] = 16; tc.colsb[4] = 64; - tc.rows[5] = 16; tc.colsb[5] = 64; - tc.rows[6] = 16; tc.colsb[6] = 64; - tc.rows[7] = 16; tc.colsb[7] = 64; - - _tile_loadconfig(&tc); - done = true; -} - -// we need an extra 16 * 4B (TILE_N * int32_t) for each NB/KB block for compensation. -// See the notes `s8s8 igemm compensation in avx512-vnni` for detail. -template -int get_tile_size() { - int tile_size = TILE_N * sizeof(TB); - if (do_compensate::value) { - tile_size += TILE_N * sizeof(int32_t); - } - if (std::is_same::value || - std::is_same::value) { - tile_size += TILE_N * 4; - } - if (std::is_same::value) { - tile_size += TILE_N * 2; - } - return tile_size; -} - -template -int get_row_size(int K) { - int KB = K / BLOCK_K; - int row_size = KB * sizeof(TB); - if (do_compensate::value) { - row_size += KB * sizeof(int32_t); - } - if (std::is_same::value || - std::is_same::value) { - row_size += KB * 4; - } - if (std::is_same::value) { - row_size += KB * 2; - } - return row_size; -} - -// transpose utils -#define SHUFFLE_EPI32(a, b, mask) \ - _mm256_castps_si256(_mm256_shuffle_ps(_mm256_castsi256_ps(a), _mm256_castsi256_ps(b), mask)) -inline void transpose_8x8_32bit(__m256i * v, __m256i * v1) { - // unpacking and 32-bit elements - v1[0] = _mm256_unpacklo_epi32(v[0], v[1]); - v1[1] = _mm256_unpackhi_epi32(v[0], v[1]); - v1[2] = _mm256_unpacklo_epi32(v[2], v[3]); - v1[3] = _mm256_unpackhi_epi32(v[2], v[3]); - v1[4] = _mm256_unpacklo_epi32(v[4], v[5]); - v1[5] = _mm256_unpackhi_epi32(v[4], v[5]); - v1[6] = _mm256_unpacklo_epi32(v[6], v[7]); - v1[7] = _mm256_unpackhi_epi32(v[6], v[7]); - - // shuffling the 32-bit elements - v[0] = SHUFFLE_EPI32(v1[0], v1[2], 0x44); - v[1] = SHUFFLE_EPI32(v1[0], v1[2], 0xee); - v[2] = SHUFFLE_EPI32(v1[4], v1[6], 0x44); - v[3] = SHUFFLE_EPI32(v1[4], v1[6], 0xee); - v[4] = SHUFFLE_EPI32(v1[1], v1[3], 0x44); - v[5] = SHUFFLE_EPI32(v1[1], v1[3], 0xee); - v[6] = SHUFFLE_EPI32(v1[5], v1[7], 0x44); - v[7] = SHUFFLE_EPI32(v1[5], v1[7], 0xee); - - // shuffling 128-bit elements - v1[0] = _mm256_permute2f128_si256(v[2], v[0], 0x02); - v1[1] = _mm256_permute2f128_si256(v[3], v[1], 0x02); - v1[2] = _mm256_permute2f128_si256(v[6], v[4], 0x02); - v1[3] = _mm256_permute2f128_si256(v[7], v[5], 0x02); - v1[4] = _mm256_permute2f128_si256(v[2], v[0], 0x13); - v1[5] = _mm256_permute2f128_si256(v[3], v[1], 0x13); - v1[6] = _mm256_permute2f128_si256(v[6], v[4], 0x13); - v1[7] = _mm256_permute2f128_si256(v[7], v[5], 0x13); -} - -inline void transpose_16x4_32bit(__m512i * r, __m512i * d) { - - static const __m512i index1 = _mm512_set_epi32( - 0x0f, 0x0b, 0x07, 0x03, - 0x0e, 0x0a, 0x06, 0x02, - 0x0d, 0x09, 0x05, 0x01, - 0x0c, 0x08, 0x04, 0x00); - - d[0] = _mm512_permutexvar_epi32(index1, r[0]); - d[1] = _mm512_permutexvar_epi32(index1, r[1]); - d[2] = _mm512_permutexvar_epi32(index1, r[2]); - d[3] = _mm512_permutexvar_epi32(index1, r[3]); - - r[0] = _mm512_shuffle_i32x4(d[0], d[1], 0x44); - r[1] = _mm512_shuffle_i32x4(d[0], d[1], 0xee); - r[2] = _mm512_shuffle_i32x4(d[2], d[3], 0x44); - r[3] = _mm512_shuffle_i32x4(d[2], d[3], 0xee); - - d[0] = _mm512_shuffle_i32x4(r[0], r[2], 0x88); - d[1] = _mm512_shuffle_i32x4(r[0], r[2], 0xdd); - d[2] = _mm512_shuffle_i32x4(r[1], r[3], 0x88); - d[3] = _mm512_shuffle_i32x4(r[1], r[3], 0xdd); -} - -inline void transpose_16x16_32bit(__m512i * v) { - __m512i v1[16]; - v1[0] = _mm512_unpacklo_epi32(v[0], v[1]); - v1[1] = _mm512_unpackhi_epi32(v[0], v[1]); - v1[2] = _mm512_unpacklo_epi32(v[2], v[3]); - v1[3] = _mm512_unpackhi_epi32(v[2], v[3]); - v1[4] = _mm512_unpacklo_epi32(v[4], v[5]); - v1[5] = _mm512_unpackhi_epi32(v[4], v[5]); - v1[6] = _mm512_unpacklo_epi32(v[6], v[7]); - v1[7] = _mm512_unpackhi_epi32(v[6], v[7]); - v1[8] = _mm512_unpacklo_epi32(v[8], v[9]); - v1[9] = _mm512_unpackhi_epi32(v[8], v[9]); - v1[10] = _mm512_unpacklo_epi32(v[10], v[11]); - v1[11] = _mm512_unpackhi_epi32(v[10], v[11]); - v1[12] = _mm512_unpacklo_epi32(v[12], v[13]); - v1[13] = _mm512_unpackhi_epi32(v[12], v[13]); - v1[14] = _mm512_unpacklo_epi32(v[14], v[15]); - v1[15] = _mm512_unpackhi_epi32(v[14], v[15]); - - v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]); - v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]); - v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]); - v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]); - v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]); - v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]); - v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]); - v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]); - v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]); - v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]); - v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]); - v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]); - v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]); - v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]); - v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]); - v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]); - - v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88); - v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88); - v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88); - v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88); - v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd); - v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd); - v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd); - v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd); - v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88); - v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88); - v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88); - v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88); - v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd); - v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd); - v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd); - v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd); - - v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88); - v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88); - v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88); - v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88); - v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88); - v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88); - v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88); - v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88); - v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd); - v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd); - v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd); - v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd); - v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd); - v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd); - v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd); - v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd); -} - -void quantize_row_q8_K_vnni(const float * RESTRICT x, void * RESTRICT vy, int64_t k) { - assert(k % QK_K == 0); - const int KB = k / QK_K; - constexpr int kVecs = QK_K / 16; - - block_q8_K * y = reinterpret_cast(vy); - - // hold 16 float vecs from x - __m512 v[kVecs]; - - // hold the quants vecs - __m512i vq[kVecs / 4]; - - // hold the packed quants vecs - __m512i vq_packed[kVecs / 4]; - - const __m512 signBit = _mm512_set1_ps(-0.f); - - for (int i = 0; i < KB; ++i) { - // Compute max(abs(e)) for the block - __m512 vamax = _mm512_set1_ps(0.f); - for (int j = 0; j < kVecs; ++j) { - v[j] = _mm512_loadu_ps(x); x += 16; - vamax = _mm512_max_ps(vamax, _mm512_andnot_ps(signBit, v[j])); - } - const float amax = _mm512_reduce_max_ps(vamax); - - // Quantize these floats - const float iscale = 127.f / amax; - y[i].d = GGML_CPU_FP32_TO_FP16(1 / iscale); - const float id = ( amax != 0.0f ) ? iscale : 0.f; - const __m512 vscale = _mm512_set1_ps(id); - - // Apply multiplier and round to nearest integer - for (int j = 0; j < kVecs; ++j) { - v[j] = _mm512_mul_ps(v[j], vscale); - v[j] = _mm512_roundscale_ps(v[j], (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); - } - - // Pack to epi8 vecs - for (int j = 0; j < kVecs / 4; ++j) { - __m128i q8_0 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 0])); - __m128i q8_1 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 1])); - __m128i q8_2 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 2])); - __m128i q8_3 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 3])); - - __m256i q8_01 = _mm256_insertf128_si256(_mm256_castsi128_si256(q8_0), (q8_1), 1); - __m256i q8_23 = _mm256_insertf128_si256(_mm256_castsi128_si256(q8_2), (q8_3), 1); - - vq[j] = _mm512_inserti32x8(_mm512_castsi256_si512(q8_01), q8_23, 1); - _mm512_storeu_si512((__m512i *)(y[i].qs + j * 64), vq[j]); - } - - // Compute the bsums with vnni - transpose_16x4_32bit(vq, vq_packed); - - const __m512i one = _mm512_set1_epi8(1); - __m512i sum = _mm512_setzero_si512(); - for (int k = 0; k < 4; ++k) { - sum = _mm512_dpbusd_epi32(sum, one, vq_packed[k]); - } - _mm256_storeu_si256((__m256i *)(y[i].bsums), _mm512_cvtepi32_epi16(sum)); - } -} - -// quantize A from float to `vec_dot_type` -template -inline void from_float(const float * x, char * vy, int64_t k); - -template <> -inline void from_float(const float * x, char * vy, int64_t k) { - quantize_row_q8_0(x, (block_q8_0 *)vy, k); -} - -template <> -inline void from_float(const float * x, char * vy, int64_t k) { - quantize_row_q8_1(x, (block_q8_1 *)vy, k); -} - -template <> -inline void from_float(const float * x, char * vy, int64_t k) { -#if 1 - // TODO: this is reference impl! - quantize_row_q8_K_ref(x, (block_q8_K *)vy, k); -#else - quantize_row_q8_K_vnni(x, vy, k); -#endif -} - -// load A from memory to array when nrows can not fill in whole tile -void unpack_A(int8_t * RESTRICT tile, const block_q8_0 * RESTRICT A, int lda, int nr) { - assert(nr != TILE_M); - for (int m = 0; m < nr; ++m) { - const __m256i v = _mm256_loadu_si256((const __m256i *)(A[m * lda].qs)); - _mm256_storeu_si256((__m256i *)(tile + m * TILE_K), v); - } -} - -void unpack_A(int8_t * RESTRICT tile, const block_q8_1 * RESTRICT A, int lda, int nr) { - assert(nr != TILE_M); - for (int m = 0; m < nr; ++m) { - const __m256i v = _mm256_loadu_si256((const __m256i *)(A[m * lda].qs)); - _mm256_storeu_si256((__m256i *)(tile + m * TILE_K), v); - } -} - -template -void unpack_A(int8_t * RESTRICT tile, const block_q8_K * RESTRICT A, int lda, int k, int nr) { - assert(nr <= TILE_M); - for (int m = 0; m < nr; ++m) { - const __m256i v = _mm256_loadu_si256((const __m256i *)(A[m * lda].qs + k * 32)); - _mm256_storeu_si256((__m256i *)(tile + m * TILE_K), v); - } -} - -template <> -void unpack_A(int8_t * RESTRICT tile, const block_q8_K * RESTRICT A, int lda, int k, int nr) { - assert(nr <= TILE_M); - // zero padding k from 16 to 32, so that we don't have to re-config amx - const __m128i zero = _mm_setzero_si128(); - for (int m = 0; m < nr; ++m) { - const __m128i v = _mm_loadu_si128((const __m128i *)(A[m * lda].qs + k * 16)); - const __m256i r = _mm256_insertf128_si256(_mm256_castsi128_si256(v), zero, 1); - _mm256_storeu_si256((__m256i *)(tile + m * TILE_K), r); - } -} - -#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1) -inline __m256i bytes_from_nibbles_32(const uint8_t * rsi) { - const __m128i tmp = _mm_loadu_si128((const __m128i *)rsi); - const __m256i bytes = MM256_SET_M128I(_mm_srli_epi16(tmp, 4), tmp); - const __m256i lowMask = _mm256_set1_epi8(0xF); - return _mm256_and_si256(lowMask, bytes); -} - -// used for block_q4_K -inline __m512i bytes_from_nibbles_64(const uint8_t * rsi) { - const __m256i tmp = _mm256_loadu_si256((const __m256i *)rsi); - const __m256i lowMask = _mm256_set1_epi8(0xF); - const __m256i q4l = _mm256_and_si256(tmp, lowMask); - const __m256i q4h = _mm256_and_si256(_mm256_srli_epi16(tmp, 4), lowMask); - return _mm512_inserti32x8(_mm512_castsi256_si512(q4l), q4h, 1); -} - -// used for block_q5_K -inline __m512i bytes_from_nibbles_64(const uint8_t * qs, const uint8_t * qh, int k) { - const __m256i lowMask = _mm256_set1_epi8(0xF); - __m256i hmask = _mm256_set1_epi8(1); - hmask = _mm256_slli_epi16(hmask, k); - - const __m256i q5bits = _mm256_loadu_si256((const __m256i *)qs); - const __m256i hbits = _mm256_loadu_si256((const __m256i *)qh); - - const __m256i q5l_0 = _mm256_and_si256(q5bits, lowMask); - const __m256i q5h_0 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_and_si256(hbits, hmask), k + 0), 4); - const __m256i q5_0 = _mm256_add_epi8(q5l_0, q5h_0); - hmask = _mm256_slli_epi16(hmask, 1); - - const __m256i q5l_1 = _mm256_and_si256(_mm256_srli_epi16(q5bits, 4), lowMask); - const __m256i q5h_1 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_and_si256(hbits, hmask), k + 1), 4); - const __m256i q5_1 = _mm256_add_epi8(q5l_1, q5h_1); - - return _mm512_inserti32x8(_mm512_castsi256_si512(q5_0), q5_1, 1); -} - -// used for block_q6_K -inline void bytes_from_nibbles_128(__m512i& r0, __m512i& r1, const uint8_t * qs, const uint8_t * qh) { - const __m256i m4 = _mm256_set1_epi8(0xF); - const __m256i m2 = _mm256_set1_epi8(0x3); - - const __m256i q6bits1 = _mm256_loadu_si256((const __m256i *)qs); - const __m256i q6bits2 = _mm256_loadu_si256((const __m256i *)(qs + 32)); - const __m256i q6bitsH = _mm256_loadu_si256((const __m256i *)qh); - - const __m256i q6h_0 = _mm256_slli_epi16(_mm256_and_si256( q6bitsH, m2), 4); - const __m256i q6h_1 = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(q6bitsH, 2), m2), 4); - const __m256i q6h_2 = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(q6bitsH, 4), m2), 4); - const __m256i q6h_3 = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(q6bitsH, 6), m2), 4); - - const __m256i q6_0 = _mm256_or_si256(_mm256_and_si256(q6bits1, m4), q6h_0); - const __m256i q6_1 = _mm256_or_si256(_mm256_and_si256(q6bits2, m4), q6h_1); - const __m256i q6_2 = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(q6bits1, 4), m4), q6h_2); - const __m256i q6_3 = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(q6bits2, 4), m4), q6h_3); - - r0 = _mm512_inserti32x8(_mm512_castsi256_si512(q6_0), q6_1, 1); - r1 = _mm512_inserti32x8(_mm512_castsi256_si512(q6_2), q6_3, 1); -} - -inline __m512i packNibbles(__m512i r0, __m512i r1) { - return _mm512_or_si512(r0, _mm512_slli_epi16(r1, 4)); -} - -template -inline void pack_qs(void * RESTRICT packed_B, const TB * RESTRICT B, int KB) { - int8_t tmp[8 * 64]; - __m256i v[8], v2[8]; - for (int n = 0; n < 8; ++n) { - v[n] = bytes_from_nibbles_32(B[n * KB].qs); - } - transpose_8x8_32bit(v, v2); - for (int n = 0; n < 8; ++n) { - _mm256_storeu_si256((__m256i *)(tmp + n * 64), v2[n]); - } - for (int n = 0; n < 8; ++n) { - v[n] = bytes_from_nibbles_32(B[(n + 8) * KB].qs); - } - transpose_8x8_32bit(v, v2); - for (int n = 0; n < 8; ++n) { - _mm256_storeu_si256((__m256i *)(tmp + n * 64 + 32), v2[n]); - } - - // pack again with 128 to fully utilize vector length - for (int n = 0; n < 8; n += 2) { - __m512i r0 = _mm512_loadu_si512((const __m512i *)(tmp + n * 64)); - __m512i r1 = _mm512_loadu_si512((const __m512i *)(tmp + n * 64 + 64)); - __m512i r1r0 = packNibbles(r0, r1); - _mm512_storeu_si512((__m512i *)((char *)packed_B + n * 32), r1r0); - } -} - -template <> -inline void pack_qs(void * RESTRICT packed_B, const block_q8_0 * RESTRICT B, int KB) { - __m256i v[8], v2[8]; - for (int n = 0; n < 8; ++n) { - v[n] = _mm256_loadu_si256((const __m256i *)(B[n * KB].qs)); - } - transpose_8x8_32bit(v, v2); - for (int n = 0; n < 8; ++n) { - _mm256_storeu_si256((__m256i *)((char *)packed_B + n * 64), v2[n]); - } - for (int n = 0; n < 8; ++n) { - v[n] = _mm256_loadu_si256((const __m256i *)(B[(n + 8) * KB].qs)); - } - transpose_8x8_32bit(v, v2); - for (int n = 0; n < 8; ++n) { - _mm256_storeu_si256((__m256i *)((char *)packed_B + n * 64 + 32), v2[n]); - } -} - -template <> -inline void pack_qs(void * RESTRICT packed_B, const block_q4_K * RESTRICT B, int KB) { - __m512i v[16]; - // QK_K 256 with 8 groups, handle 2 groups at a time - char * pb = (char *)packed_B; - for (int k = 0; k < QK_K / 64; ++k) { - // pack 2 groups { n, g, k} to {g, k/4, 4n} - // e.g. {16, 2, 32} to {2, 8, 64} - for (int n = 0; n < TILE_N; ++n) { - v[n] = bytes_from_nibbles_64(B[n * KB].qs + k * 32); - } - - transpose_16x16_32bit(v); - - // pack again with 128 to fully utilize vector length - for (int n = 0; n < TILE_N; n += 2) { - _mm512_storeu_si512((__m512i *)pb, packNibbles(v[n], v[n + 1])); - pb += 64; - } - } -} - -template <> -inline void pack_qs(void * RESTRICT packed_B, const block_q5_K * RESTRICT B, int KB) { - __m512i v[16]; - const __m512i lowMask = _mm512_set1_epi8(0xF); - // QK_K 256 with 8 groups, handle 2 groups at a time - char * pb = (char *)packed_B; - char * ph = (char *)packed_B + (QK_K / 2) * TILE_N; - for (int k = 0; k < QK_K / 64; ++k) { - // pack 2 groups { n, g, k} to {g, k/4, 4n} - // e.g. {16, 2, 32} to {2, 8, 64} - for (int n = 0; n < TILE_N; ++n) { - v[n] = bytes_from_nibbles_64(B[n * KB].qs + k * 32, B[n * KB].qh, /* group */2 * k); - } - - transpose_16x16_32bit(v); - - // 1. pack lower 4bits with 2 groups - for (int n = 0; n < TILE_N; n += 2) { - // get lower 4 bits - const __m512i r0 = _mm512_and_si512(v[n], lowMask); - const __m512i r1 = _mm512_and_si512(v[n + 1], lowMask); - _mm512_storeu_si512((__m512i *)pb, packNibbles(r0, r1)); pb += 64; - } - - // 2. pack higher 1bit with 2 groups - const __m512i hmask = _mm512_set1_epi8(0x10); - for (int g = 0; g < 2; ++g) { - __m512i hbits = _mm512_setzero_si512(); - hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 0], hmask), 4)); - hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 1], hmask), 3)); - hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 2], hmask), 2)); - hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 3], hmask), 1)); - hbits = _mm512_add_epi8(hbits, _mm512_and_si512(v[g * 8 + 4], hmask) ); - hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 8 + 5], hmask), 1)); - hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 8 + 6], hmask), 2)); - hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 8 + 7], hmask), 3)); - _mm512_storeu_si512((__m512i *)ph, hbits); ph += 64; - } - } -} - -template <> -inline void pack_qs(void * RESTRICT packed_B, const block_q6_K * RESTRICT B, int KB) { - __m512i v[32]; - const __m512i lowMask = _mm512_set1_epi8(0xF); - // QK_K 256 with 8 groups, handle 4 groups at a time - char * pb = (char *)packed_B; - char * ph = (char *)packed_B + (QK_K / 2) * TILE_N; - for (int k = 0; k < QK_K / 128; ++k) { - for (int n = 0; n < TILE_N; ++n) { - bytes_from_nibbles_128(v[n], v[n + 16], B[n * KB].ql + k * 64, B[n * KB].qh + k * 32); - } - - // top half: group 0,1 or 4,5; bottom half: group 2,3 or 6,7 - transpose_16x16_32bit(v); - transpose_16x16_32bit(v + 16); - - // 1. pack lower 4bits with 4 groups - for (int n = 0; n < 32; n += 2) { - const __m512i r0 = _mm512_and_si512(v[n], lowMask); - const __m512i r1 = _mm512_and_si512(v[n + 1], lowMask); - _mm512_storeu_si512((__m512i *)pb, packNibbles(r0, r1)); pb += 64; - } - - // 2. pack higher 2bit with 4 groups - const __m512i hmask = _mm512_set1_epi8(0x30); - for (int g = 0; g < 8; ++g) { - __m512i hbits = _mm512_setzero_si512(); - hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 4 + 0], hmask), 4)); - hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 4 + 1], hmask), 2)); - hbits = _mm512_add_epi8(hbits, _mm512_and_si512(v[g * 4 + 2], hmask) ); - hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 4 + 3], hmask), 2)); - _mm512_storeu_si512((__m512i *)ph, hbits); ph += 64; - } - } -} - -template <> -inline void pack_qs(void * RESTRICT packed_B, const block_iq4_xs * RESTRICT B, int KB) { - __m512i v[16]; - char * pb = (char *)packed_B; - for (int k = 0; k < QK_K / 64; ++k) { - for (int n = 0; n < TILE_N; ++n) { - __m256i r0 = bytes_from_nibbles_32(B[n * KB].qs + k * 32 + 0); - __m256i r1 = bytes_from_nibbles_32(B[n * KB].qs + k * 32 + 16); - v[n] = _mm512_inserti32x8(_mm512_castsi256_si512(r0), r1, 1); - } - - transpose_16x16_32bit(v); - - // pack again with 128 to fully utilize vector length - for (int n = 0; n < TILE_N; n += 2) { - _mm512_storeu_si512((__m512i *)pb, packNibbles(v[n], v[n + 1])); - pb += 64; - } - } -} - -// pack B to vnni formats in 4bits or 8 bits -void pack_B(void * RESTRICT packed_B, const block_q4_0 * RESTRICT B, int KB) { - pack_qs(packed_B, B, KB); - ggml_half * d0 = reinterpret_cast((char *)packed_B + TILE_N * TILE_K / 2); - for (int n = 0; n < TILE_N; ++n) { - d0[n] = B[n * KB].d; - } -} - -void pack_B(void * RESTRICT packed_B, const block_q4_1 * RESTRICT B, int KB) { - pack_qs(packed_B, B, KB); - ggml_half * d0 = reinterpret_cast((char *)packed_B + TILE_N * TILE_K / 2); - ggml_half * m0 = d0 + TILE_N; - for (int n = 0; n < TILE_N; ++n) { - d0[n] = B[n * KB].d; - m0[n] = B[n * KB].m; - } -} - -inline void s8s8_compensation(void * RESTRICT packed_B) { - // packed_B layout: - // quants {TILE_N, TILEK} int8_t - // d0 {TILE_N} ggml_half - // comp {TILE_N} int32_t - const int offset = TILE_N * TILE_K + TILE_N * sizeof(ggml_half); - __m512i vcomp = _mm512_setzero_si512(); - const __m512i off = _mm512_set1_epi8(static_cast(0x80)); - for (int k = 0; k < 8; ++k) { - __m512i vb = _mm512_loadu_si512((const __m512i *)((const char *)packed_B + k * 64)); - vcomp = _mm512_dpbusd_epi32(vcomp, off, vb); - } - _mm512_storeu_si512((__m512i *)((char *)(packed_B) + offset), vcomp); -} - -void pack_B(void * RESTRICT packed_B, const block_q8_0 * RESTRICT B, int KB) { - pack_qs(packed_B, B, KB); - ggml_half * d0 = reinterpret_cast((char *)packed_B + TILE_N * TILE_K); - for (int n = 0; n < TILE_N; ++n) { - d0[n] = B[n * KB].d; - } - s8s8_compensation(packed_B); -} - -// convert 8 * {min, scale} from int6 to int8 -inline void unpack_mins_and_scales(const uint8_t * scales, uint32_t * utmp) { - const uint32_t kmask1 = 0x3f3f3f3f; - const uint32_t kmask2 = 0x0f0f0f0f; - const uint32_t kmask3 = 0x03030303; - - memcpy(utmp, scales, 12); - utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); - const uint32_t uaux = utmp[1] & kmask1; - utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); - utmp[2] = uaux; - utmp[0] &= kmask1; -} - -// packed_B layout: -// quants {8, TILE_N, 16} uint8 -// scales {8, TILE_N} uint8 -// mins {8, TILE_N} uint8 -// d {TILE_N} ggml_half -// dmin {TILE_N} ggml_half -void pack_B(void * RESTRICT packed_B, const block_q4_K * RESTRICT B, int KB) { - pack_qs(packed_B, B, KB); - - uint8_t * scales = reinterpret_cast((char *)packed_B + (QK_K / 2) * TILE_N); - uint8_t * mins = scales + 8 * TILE_N; - ggml_half * d = reinterpret_cast(mins + 8 * TILE_N); - ggml_half * dmin = d + TILE_N; - - union { - uint32_t u32[4]; - uint8_t u8[16]; - } s; - - for (int n = 0; n < TILE_N; ++n) { - unpack_mins_and_scales(B[n * KB].scales, s.u32); - for (int k = 0; k < 8; ++k) { - scales[k * TILE_N + n] = s.u8[k]; - mins[(k >> 1) * TILE_N * 2 + n * 2 + (k & 0x1)] = s.u8[k + 8]; - } - d[n] = B[n * KB].d; - dmin[n] = B[n * KB].dmin; - } -} - -// packed_B layout: -// quants {8, TILE_N, 16} uint8 -// qh {8, TILE_N, 4} uint8 -// scales {8, TILE_N} uint8 -// mins {8, TILE_N} uint8 -// d {TILE_N} ggml_half -// dmin {TILE_N} ggml_half -void pack_B(void * RESTRICT packed_B, const block_q5_K * RESTRICT B, int KB) { - pack_qs(packed_B, B, KB); - - uint8_t * scales = reinterpret_cast((char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N); - uint8_t * mins = scales + 8 * TILE_N; - ggml_half * d = reinterpret_cast(mins + 8 * TILE_N); - ggml_half * dmin = d + TILE_N; - - union { - uint32_t u32[4]; - uint8_t u8[16]; - } s; - - for (int n = 0; n < TILE_N; ++n) { - unpack_mins_and_scales(B[n * KB].scales, s.u32); - for (int k = 0; k < 8; ++k) { - scales[k * TILE_N + n] = s.u8[k]; - mins[(k >> 1) * TILE_N * 2 + n * 2 + (k & 0x1)] = s.u8[k + 8]; - } - d[n] = B[n * KB].d; - dmin[n] = B[n * KB].dmin; - } -} - -// packed_B layout: -// quants {16, TILE_N, 8} uint8 -// qh {16, TILE_N, 4} uint8 -// scales {16, TILE_N} uint8 -// d {TILE_N} ggml_half -void pack_B(void * RESTRICT packed_B, const block_q6_K * RESTRICT B, int KB) { - pack_qs(packed_B, B, KB); - - uint8_t * scales = reinterpret_cast((char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N); - ggml_half * d = reinterpret_cast(scales + 16 * TILE_N); - for (int n = 0; n < TILE_N; ++n) { - const int8_t * ps = B[n * KB].scales; - for (int k = 0; k < 16; ++k) { - scales[k * TILE_N + n] = ps[k]; - } - d[n] = B[n * KB].d; - } -} - -// packed_B layout: -// quants {8, TILE_N, 16} uint8 -// scales {8, TILE_N} int8 -// d {TILE_N} ggml_half -void pack_B(void * RESTRICT packed_B, const block_iq4_xs * RESTRICT B, int KB) { - pack_qs(packed_B, B, KB); - - int8_t * scales = reinterpret_cast((char *)packed_B + (QK_K / 2) * TILE_N); - ggml_half * d = reinterpret_cast(scales + 8 * TILE_N); - - // pack the scales - for (int n = 0; n < TILE_N; ++n) { - uint16_t sh = B[n * KB].scales_h; - for (int k = 0; k < 8; k += 2) { - const int16_t ls1 = ((B[n * KB].scales_l[k / 2] & 0xf) | ((sh << 4) & 0x30)) - 32; - const int16_t ls2 = ((B[n * KB].scales_l[k / 2] >> 4) | ((sh << 2) & 0x30)) - 32; - scales[(k + 0) * TILE_N + n] = ls1; - scales[(k + 1) * TILE_N + n] = ls2; - sh >>= 4; - } - d[n] = B[n * KB].d; - } -} - -template> -void unpack_B(packed_B_t * RESTRICT tile, const void * RESTRICT packed_B) { - GGML_UNUSED(tile); - GGML_UNUSED(packed_B); -} - -template <> -void unpack_B(int8_t * RESTRICT tile, const void * RESTRICT packed_B) { - const __m512i off = _mm512_set1_epi8(8); - const __m512i lowMask = _mm512_set1_epi8(0xF); - for (int n = 0; n < 8; n += 2) { - __m512i bytes = _mm512_loadu_si512((const __m512i *)((const char *)packed_B + n * 32)); - const __m512i r0 = _mm512_sub_epi8(_mm512_and_si512(bytes, lowMask), off); - const __m512i r1 = _mm512_sub_epi8(_mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask), off); - _mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0); - _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1); - } -} - -template <> -void unpack_B(uint8_t * RESTRICT tile, const void * RESTRICT packed_B) { - const __m512i lowMask = _mm512_set1_epi8(0xF); - for (int n = 0; n < 8; n += 2) { - __m512i bytes = _mm512_loadu_si512((const __m512i *)((const char *)packed_B + n * 32)); - const __m512i r0 = _mm512_and_si512(bytes, lowMask); - const __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); - _mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0); - _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1); - } -} - -// packed_B_t for QKK is int8_t -template -void unpack_B(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) { - const int packed_B_group_size = QK_K / 2 * TILE_N / 8; - const char * packed_B_group = (const char *)packed_B + k * packed_B_group_size; - const __m512i lowMask = _mm512_set1_epi8(0xF); - for (int n = 0; n < 8; n += 2) { - __m512i bytes = _mm512_loadu_si512(packed_B_group + n * 32); - const __m512i r0 = _mm512_and_si512(bytes, lowMask); - const __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); - _mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0); - _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1); - } -} - -template <> -void unpack_B(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) { - // lower 4bits, stride 256 bytes - const int packed_l4_group_size = QK_K / 2 * TILE_N / 8; - const char * pb = (const char *)packed_B + k * packed_l4_group_size; - - // higher 1bit, stride 64 bytes - const int packed_h1_group_size = QK_K / 8 * TILE_N / 8; - const char * ph = (const char *)packed_B + (QK_K / 2) * TILE_N + k * packed_h1_group_size; - const __m512i hbits = _mm512_loadu_si512(ph); - - const __m512i lowMask = _mm512_set1_epi8(0xF); - __m512i hmask0 = _mm512_set1_epi8(0x1); - __m512i hmask1 = _mm512_set1_epi8(0x2); - - for (int n = 0; n < 8; n += 2) { - __m512i bytes = _mm512_loadu_si512(pb + n * 32); - __m512i r0 = _mm512_and_si512(bytes, lowMask); - __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); - __m512i h0 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask0), n), 4); - __m512i h1 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask1), n + 1), 4); - - hmask0 = _mm512_slli_epi16(hmask0, 2); - hmask1 = _mm512_slli_epi16(hmask1, 2); - r0 = _mm512_add_epi8(r0, h0); - r1 = _mm512_add_epi8(r1, h1); - _mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0); - _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1); - } -} - -template <> -void unpack_B(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) { - // lower 4bits, stride 128 bytes - const int packed_l4_group_size = QK_K / 2 * TILE_N / 16; - const char * pb = (const char *)packed_B + k * packed_l4_group_size; - - // higher 2bits, stride 64 bytes - const int packed_h2_group_size = QK_K / 4 * TILE_N / 16; - const char * ph = (const char *)packed_B + (QK_K / 2) * TILE_N + k * packed_h2_group_size; - const __m512i hbits = _mm512_loadu_si512(ph); - - const __m512i off = _mm512_set1_epi8(32); - const __m512i lowMask = _mm512_set1_epi8(0xF); - __m512i hmask0 = _mm512_set1_epi8(0x3); // 0011 - __m512i hmask1 = _mm512_set1_epi8(0xC); // 1100 - - // notes: skip zero padding from row4 to row7 as we have done so in `unpack_A` - __m512i bytes = _mm512_loadu_si512(pb); - __m512i r0 = _mm512_and_si512(bytes, lowMask); - __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); - __m512i h0 = _mm512_slli_epi16(_mm512_and_si512(hbits, hmask0), 4); - __m512i h1 = _mm512_slli_epi16(_mm512_and_si512(hbits, hmask1), 2); - _mm512_storeu_si512((__m512i *)(tile + 0), _mm512_sub_epi8(_mm512_add_epi8(r0, h0), off)); - _mm512_storeu_si512((__m512i *)(tile + 64), _mm512_sub_epi8(_mm512_add_epi8(r1, h1), off)); - - hmask0 = _mm512_slli_epi16(hmask0, 4); - hmask1 = _mm512_slli_epi16(hmask1, 4); - - bytes = _mm512_loadu_si512(pb + 64); - r0 = _mm512_and_si512(bytes, lowMask); - r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); - h0 = _mm512_and_si512(hbits, hmask0); - h1 = _mm512_srli_epi16(_mm512_and_si512(hbits, hmask1), 2); - _mm512_storeu_si512((__m512i *)(tile + 128), _mm512_sub_epi8(_mm512_add_epi8(r0, h0), off)); - _mm512_storeu_si512((__m512i *)(tile + 192), _mm512_sub_epi8(_mm512_add_epi8(r1, h1), off)); -} - -template <> -void unpack_B(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) { - static const __m512i values128 = _mm512_set_epi8( - 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127, - 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127, - 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127, - 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127 - ); - - const int packed_B_group_size = QK_K / 2 * TILE_N / 8; - const char * pb = (const char *)packed_B + k * packed_B_group_size; - const __m512i lowMask = _mm512_set1_epi8(0xF); - - for (int n = 0; n < 8; n += 2) { - __m512i bytes = _mm512_loadu_si512(pb + n * 32); - const __m512i r0 = _mm512_shuffle_epi8(values128, _mm512_and_si512(bytes, lowMask)); - const __m512i r1 = _mm512_shuffle_epi8(values128, _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask)); - _mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0); - _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1); - } -} - -template -struct acc_C {}; - -template -struct acc_C { - static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_0 * A, int lda, const void * packed_B, int nr) { - const int offset = TILE_N * TILE_K / 2; - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset))); - - for (int m = 0; m < nr; ++m) { - const __m512 vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[m * lda].d)); - const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); - - __m512 vsum; - if (is_acc) { - vsum = _mm512_loadu_ps(C + m * ldc); - } else { - vsum = _mm512_set1_ps(0.f); - } - vsum = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), vsum); - _mm512_storeu_ps(C + m * ldc, vsum); - } - } -}; - -template -struct acc_C { - static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_1 * A, int lda, const void * packed_B, int nr) { - const int offset = TILE_N * TILE_K / 2; - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset))); - const __m512 vm0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset + TILE_N * sizeof(ggml_half)))); - - for (int m = 0; m < nr; ++m) { - const __m512 vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[m * lda].d)); - const __m512 vs1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[m * lda].s)); - const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); - - __m512 vsum; - if (is_acc) { - vsum = _mm512_loadu_ps(C + m * ldc); - } else { - vsum = _mm512_set1_ps(0.f); - } - vsum = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), vsum); - vsum = _mm512_fmadd_ps(vm0, vs1, vsum); - _mm512_storeu_ps(C + m * ldc, vsum); - } - } -}; - -template -struct acc_C { - static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_0 * A, int lda, const void * packed_B, int nr) { - const int offset = TILE_N * TILE_K; - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset))); - - for (int m = 0; m < nr; ++m) { - const __m512 vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[m * lda].d)); - const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); - - __m512 vsum; - if (is_acc) { - vsum = _mm512_loadu_ps(C + m * ldc); - } else { - vsum = _mm512_set1_ps(0.f); - } - vsum = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), vsum); - _mm512_storeu_ps(C + m * ldc, vsum); - } - } -}; - -template -struct acc_C { - static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) { - const uint8_t * scales = reinterpret_cast((const char *)packed_B + (QK_K / 2) * TILE_N); - const uint8_t * mins = scales + 8 * TILE_N; - const ggml_half * d0 = reinterpret_cast(mins + 8 * TILE_N); - const ggml_half * dmin = d0 + TILE_N; - - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0)); - const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)dmin)); - - for (int m = 0; m < nr; ++m) { - const float d1 = A[m * lda].d; - const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0); - const __m512 vdm = _mm512_mul_ps(_mm512_set1_ps(-d1), vdmin); - const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); - - __m512 vsum; - if (is_acc) { - vsum = _mm512_loadu_ps(C + m * ldc); - } else { - vsum = _mm512_set1_ps(0.f); - } - - const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[m * lda].bsums); - const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1)); - - __m512i acc_m = _mm512_setzero_si512(); - for (int k = 0; k < 4; ++k) { - __m512i vmask = _mm512_set1_epi32(k); - __m512i va = _mm512_permutexvar_epi32(vmask, _mm512_castsi128_si512(q8s)); - __m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(mins + k * 32))); - acc_m = _mm512_dpwssds_epi32(acc_m, va, vb); - } - - vsum = _mm512_fmadd_ps(vtile, vd, vsum); - vsum = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc_m), vdm, vsum); - _mm512_storeu_ps(C + m * ldc, vsum); - } - } -}; - -template -struct acc_C { - static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) { - const uint8_t * scales = reinterpret_cast((const char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N); - const uint8_t * mins = scales + 8 * TILE_N; - const ggml_half * d0 = reinterpret_cast(mins + 8 * TILE_N); - const ggml_half * dmin = d0 + TILE_N; - - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0)); - const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)dmin)); - - for (int m = 0; m < nr; ++m) { - const float d1 = A[m * lda].d; - const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0); - const __m512 vdm = _mm512_mul_ps(_mm512_set1_ps(-d1), vdmin); - const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); - - __m512 vsum; - if (is_acc) { - vsum = _mm512_loadu_ps(C + m * ldc); - } else { - vsum = _mm512_set1_ps(0.f); - } - - const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[m * lda].bsums); - const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1)); - - __m512i acc_m = _mm512_setzero_si512(); - for (int k = 0; k < 4; ++k) { - __m512i vmask = _mm512_set1_epi32(k); - __m512i va = _mm512_permutexvar_epi32(vmask, _mm512_castsi128_si512(q8s)); - __m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(mins + k * 32))); - acc_m = _mm512_dpwssds_epi32(acc_m, va, vb); - } - - vsum = _mm512_fmadd_ps(vtile, vd, vsum); - vsum = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc_m), vdm, vsum); - _mm512_storeu_ps(C + m * ldc, vsum); - } - } -}; - -template -struct acc_C { - static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) { - const uint8_t * scales = reinterpret_cast((const char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N); - const ggml_half * d0 = reinterpret_cast(scales + 16 * TILE_N); - - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0)); - - for (int m = 0; m < nr; ++m) { - const float d1 = A[m * lda].d; - const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0); - const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); - - __m512 vsum; - if (is_acc) { - vsum = _mm512_loadu_ps(C + m * ldc); - } else { - vsum = _mm512_set1_ps(0.f); - } - - vsum = _mm512_fmadd_ps(vtile, vd, vsum); - _mm512_storeu_ps(C + m * ldc, vsum); - } - } -}; - -template -struct acc_C { - static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) { - const int8_t * scales = reinterpret_cast((const char *)packed_B + (QK_K / 2) * TILE_N); - const ggml_half * d0 = reinterpret_cast(scales + 8 * TILE_N); - - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0)); - - for (int m = 0; m < nr; ++m) { - const float d1 = A[m * lda].d; - const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0); - const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); - - __m512 vsum; - if (is_acc) { - vsum = _mm512_loadu_ps(C + m * ldc); - } else { - vsum = _mm512_set1_ps(0.f); - } - - vsum = _mm512_fmadd_ps(vtile, vd, vsum); - _mm512_storeu_ps(C + m * ldc, vsum); - } - } -}; - -template constexpr int get_quants_size(); -template <> constexpr int get_quants_size() { return (QK_K / 2) * TILE_N; } -template <> constexpr int get_quants_size() { return (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N; } -template <> constexpr int get_quants_size() { return (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N; } -template <> constexpr int get_quants_size() { return (QK_K / 2) * TILE_N; } - -// used for QKK format -template ::value, int>::type = 0> -inline void scale_C(const int32_t * RESTRICT tile, int32_t * RESTRICT sumi, const void * packed_B, int k, int nr) { - const uint8_t * scales = reinterpret_cast((const char *)packed_B + get_quants_size()); - const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(scales + k * TILE_N))); - - for (int m = 0; m < nr; ++m) { - __m512i vsumi; - if (is_acc) { - vsumi = _mm512_loadu_si512(sumi + m * TILE_N); - } else { - vsumi = _mm512_setzero_si512(); - } - __m512i vtile = _mm512_loadu_si512(tile + m * TILE_N); - vsumi = _mm512_add_epi32(vsumi, _mm512_mullo_epi32(vtile, vscale)); - _mm512_storeu_si512((__m512i *)(sumi + m * TILE_N), vsumi); - } -} - -template -struct tinygemm_kernel_avx { - static void apply(int K, const TA * RESTRICT A, const TB * RESTRICT B, TC * RESTRICT C, int ldc) { - GGML_UNUSED(K); - GGML_UNUSED(A); - GGML_UNUSED(B); - GGML_UNUSED(C); - GGML_UNUSED(ldc); - } -}; - -template -struct tinygemm_kernel_avx { - static void apply(int K, const float * RESTRICT A, const ggml_fp16_t * RESTRICT B, float * RESTRICT C, int ldc) { - constexpr int ROWS = BLOCK_M; - constexpr int COLS = BLOCK_N; - assert(BLOCK_K == 16); - - __m512 va; - __m512 vb[COLS]; - __m512 vc[ROWS * COLS]; - - auto loadc = [&](auto idx) { - vc[idx] = _mm512_setzero_ps(); - }; - Unroll{}(loadc); - - auto compute = [&](auto idx, auto k) { - constexpr int row = idx / COLS; - constexpr int col = idx % COLS; - - if constexpr (col == 0) { - va = _mm512_loadu_ps(A + row * K + k); - } - if constexpr (row == 0) { - vb[col] = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(B + col * K + k))); - } - vc[idx] = _mm512_fmadd_ps(va, vb[col], vc[idx]); - }; - - for (int k = 0; k < K; k += 16) { - Unroll{}(compute, k); - } - - auto storec = [&](auto idx) { - constexpr int row = idx / COLS; - constexpr int col = idx % COLS; - C[row * ldc + col] = _mm512_reduce_add_ps(vc[idx]); - }; - Unroll{}(storec); - } -}; - -#define LAUNCH_TINYGEMM_KERNEL_AVX(MB_SIZE, NB_SIZE) \ - tinygemm_kernel_avx::apply( \ - K, (const float *)src1->data + src1_offset + mb_start * K, \ - (const type *)src0->data + src0_offset + nb_start * K, \ - (float *)dst->data + dst_offset + mb_start * ldc + nb_start, ldc) - - -// re-organize in the format {NB, KB, TILE_SIZE}: -#define PACKED_INDEX(n, k, KB, tile_size) (n * KB + k) * tile_size - -template -void convert_B_packed_format(void * RESTRICT packed_B, const TB * RESTRICT B, int N, int K) { - const int NB = N / TILE_N; - const int KB = K / BLOCK_K; - const int TILE_SIZE = get_tile_size(); - - // parallel on NB should be enough - parallel_for(NB, [&](int begin, int end) { - for (int n = begin; n < end; ++n) { - for (int k = 0; k < KB; ++k) { - int n0 = n * TILE_N; - pack_B((char *)packed_B + PACKED_INDEX(n, k, KB, TILE_SIZE), &B[n0 * KB + k], KB); - } - } - }); -} - -template -struct tinygemm_kernel_vnni {}; - -template -struct tinygemm_kernel_vnni { - static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { - - constexpr int COLS = BLOCK_N / 16; - const int TILE_SIZE = TILE_N * sizeof(block_q4_0); - - const block_q8_0 * RESTRICT A = static_cast(_A); - const char * RESTRICT B = static_cast(_B); - - __m512i va[8]; - __m512 vc[COLS]; - __m512 vd1; - - // sum of offsets, shared across COLS - // - // avx512-vnni does not have `_mm512_dpbssd_epi32`, - // need to transform ss to us: - // a * (b - 8) is equivalent to b * a - 8 * a - // s u u u s u s - // - __m512i vcomp; - - const __m512i off = _mm512_set1_epi8(8); - const __m512i lowMask = _mm512_set1_epi8(0xF); - - auto loadc = [&](auto col) { - vc[col] = _mm512_setzero_ps(); - }; - Unroll{}(loadc); - - auto compute = [&](auto col, auto i) { - // load a and compute compensation - if constexpr (col == 0) { - const int32_t * a_ptr = reinterpret_cast(A[0 * KB + i].qs); - vcomp = _mm512_setzero_si512(); - for (int k = 0; k < 8; ++k) { - va[k] = _mm512_set1_epi32(a_ptr[k]); - vcomp = _mm512_dpbusd_epi32(vcomp, off, va[k]); - } - vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[0 * KB + i].d)); - } - - // load b - __m512i vsum = _mm512_setzero_si512(); - const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); - for (int k = 0; k < 8; k += 2) { - __m512i bytes = _mm512_loadu_si512((const __m512i *)(b_ptr + k * 32)); - __m512i vb0 = _mm512_and_si512(bytes, lowMask); - vsum = _mm512_dpbusd_epi32(vsum, vb0, va[k + 0]); - __m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); - vsum = _mm512_dpbusd_epi32(vsum, vb1, va[k + 1]); - } - const int offset = TILE_N * TILE_K / 2; - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset))); - vsum = _mm512_sub_epi32(vsum, vcomp); - - vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(vsum), _mm512_mul_ps(vd0, vd1), vc[col]); - }; - - for (int i = 0; i < KB; ++i) { - Unroll{}(compute, i); - } - - //store to C - auto storec = [&](auto col) { - _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); - }; - Unroll{}(storec); - } -}; - -template -struct tinygemm_kernel_vnni { - static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { - - constexpr int COLS = BLOCK_N / 16; - const int TILE_SIZE = TILE_N * sizeof(block_q4_1); - - const block_q8_1 * RESTRICT A = static_cast(_A); - const char * RESTRICT B = static_cast(_B); - - __m512i va[8]; - __m512i vb[8]; - __m512 vc[COLS]; - __m512 vd1, vs1; - - const __m512i lowMask = _mm512_set1_epi8(0xF); - - auto loadc = [&](auto col) { - vc[col] = _mm512_setzero_ps(); - }; - Unroll{}(loadc); - - auto compute = [&](auto col, auto i) { - // load a - if constexpr (col == 0) { - const int32_t * a_ptr = reinterpret_cast(A[0 * KB + i].qs); - for (int k = 0; k < 8; ++k) { - va[k] = _mm512_set1_epi32(a_ptr[k]); - } - vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[0 * KB + i].d)); - vs1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[0 * KB + i].s)); - } - - // load b - const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); - for (int k = 0; k < 8; k += 2) { - __m512i bytes = _mm512_loadu_si512((const __m512i *)(b_ptr + k * 32)); - vb[k + 0] = _mm512_and_si512(bytes, lowMask); - vb[k + 1] = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); - } - const int offset = TILE_N * TILE_K / 2; - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset))); - const __m512 vm0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset + TILE_N * sizeof(ggml_half)))); - - __m512i vsum = _mm512_setzero_si512(); - for (int k = 0; k < 8; ++k) { - vsum = _mm512_dpbusd_epi32(vsum, vb[k], va[k]); - } - - vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(vsum), _mm512_mul_ps(vd0, vd1), vc[col]); - vc[col] = _mm512_fmadd_ps(vm0, vs1, vc[col]); - }; - - for (int i = 0; i < KB; ++i) { - Unroll{}(compute, i); - } - - //store to C - auto storec = [&](auto col) { - _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); - }; - Unroll{}(storec); - } -}; - -template -struct tinygemm_kernel_vnni { - static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { - - constexpr int COLS = BLOCK_N / 16; - const int TILE_SIZE = TILE_N * sizeof(block_q8_0) + TILE_N * sizeof(int32_t); - - const block_q8_0 * RESTRICT A = static_cast(_A); - const char * RESTRICT B = static_cast(_B); - - __m512i va[8]; - __m512i vb[8]; - __m512 vc[COLS]; - __m512 vd1; - - // Notes: s8s8 igemm compensation in avx512-vnni - // change s8s8 to u8s8 with compensate - // a * b = (a + 128) * b - 128 * b - // s s u s u s - // - // (128 * b is pre-computed when packing B to vnni formats) - // - const __m512i off = _mm512_set1_epi8(static_cast(0x80)); - - auto loadc = [&](auto col) { - vc[col] = _mm512_setzero_ps(); - }; - Unroll{}(loadc); - - auto compute = [&](auto col, auto i) { - // load a and add offset 128 - if constexpr (col == 0) { - const int32_t * a_ptr = reinterpret_cast(A[0 * KB + i].qs); - for (int k = 0; k < 8; ++k) { - va[k] = _mm512_set1_epi32(a_ptr[k]); - va[k] = _mm512_add_epi8(va[k], off); - } - vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[0 * KB + i].d)); - } - - // load b - const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); - for (int k = 0; k < 8; ++k) { - vb[k] = _mm512_loadu_si512((const __m512i *)(b_ptr + k * 64)); - } - const int offset = TILE_N * TILE_K; - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset))); - const int offset2 = TILE_N * TILE_K + TILE_N * sizeof(ggml_half); - const __m512i vcomp = _mm512_loadu_si512((const __m512i *)(b_ptr + offset2)); - - __m512i vsum = _mm512_setzero_si512(); - for (int k = 0; k < 8; ++k) { - vsum = _mm512_dpbusd_epi32(vsum, va[k], vb[k]); - } - vsum = _mm512_sub_epi32(vsum, vcomp); - - vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(vsum), _mm512_mul_ps(vd0, vd1), vc[col]); - }; - - for (int i = 0; i < KB; ++i) { - Unroll{}(compute, i); - } - - //store to C - auto storec = [&](auto col) { - _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); - }; - Unroll{}(storec); - } -}; - -template -struct tinygemm_kernel_vnni { - static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { - - constexpr int COLS = BLOCK_N / 16; - const int TILE_SIZE = TILE_N * sizeof(block_q4_K) + TILE_N * 4; - - const block_q8_K * RESTRICT A = static_cast(_A); - const char * RESTRICT B = static_cast(_B); - - // a.qs: 8 groups, 32 bytes each group (m256i) - __m512i va[8]; - // a.bsum: 8 groups, 2 bytes each group (m128i) - __m512i va_bsum; - __m512 vc[COLS]; - __m512 vd1; - - // packed_B: - const int offset_scales = (QK_K / 2) * TILE_N; - const int offset_mins = (QK_K / 2) * TILE_N + 8 * TILE_N; - const int offset_d0 = (QK_K / 2) * TILE_N + 16 * TILE_N; - const int offset_dmin = (QK_K / 2) * TILE_N + 16 * TILE_N + TILE_N * sizeof(ggml_half); - - const __m512i lowMask = _mm512_set1_epi8(0xF); - - auto loadc = [&](auto col) { - vc[col] = _mm512_setzero_ps(); - }; - Unroll{}(loadc); - - // Notes: vnni formats in QK_K - // a) quants vnni format - // int8 {k/4, n, 4}, viewed as 2d {k/4, 4n}, k = 32 - // from {16, 32} to {8, 64} - // - // b) min vnni format - // int16 {k/2, n, 2}, viewed as 2d {k/2, 2n}, k = 8 - // from {16, 8} to {4, 32} - // - auto compute = [&](auto col, auto i) { - // load a - if constexpr (col == 0) { - for (int k_group = 0; k_group < QK_K / 32; ++k_group) { - va[k_group] = _mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)(A[0 * KB + i].qs + k_group * 32))); - } - const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums); - const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1)); - va_bsum = _mm512_castsi128_si512(q8s); - vd1 = _mm512_set1_ps(A[0 * KB + i].d); - } - - // step 1: accumultate the quants - __m512i acc = _mm512_setzero_si512(); - const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); - const char * b_qs = b_ptr; - for (int k_group = 0; k_group < QK_K / 32; ++k_group) { - __m512i vsum = _mm512_setzero_si512(); - for (int k = 0; k < 8; k += 2) { - __m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 0), va[k_group]); - __m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 1), va[k_group]); - - __m512i bytes = _mm512_loadu_si512((const __m512i *)b_qs); - __m512i vb0 = _mm512_and_si512(bytes, lowMask); - vsum = _mm512_dpbusd_epi32(vsum, vb0, va0); - __m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); - vsum = _mm512_dpbusd_epi32(vsum, vb1, va1); - - b_qs += 64; - } - // vacc += scale * (q8 @ q4) - const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N))); - acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale)); - } - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0))); - vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]); - - // step 2: accumulate the mins - __m512i acc_m = _mm512_setzero_si512(); - for (int k = 0; k < 4; ++k) { - __m512i vmask = _mm512_set1_epi32(k); - __m512i va = _mm512_permutexvar_epi32(vmask, va_bsum); - __m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_mins + k * 32))); - acc_m = _mm512_dpwssds_epi32(acc_m, va, vb); - } - const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_dmin))); - vc[col] = _mm512_fnmadd_ps(_mm512_cvtepi32_ps(acc_m), _mm512_mul_ps(vdmin, vd1), vc[col]); - }; - - for (int i = 0; i < KB; ++i) { - Unroll{}(compute, i); - } - - //store to C - auto storec = [&](auto col) { - _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); - }; - Unroll{}(storec); - } -}; - -template -struct tinygemm_kernel_vnni { - static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { - - constexpr int COLS = BLOCK_N / 16; - const int TILE_SIZE = TILE_N * sizeof(block_q5_K) + TILE_N * 4; - - const block_q8_K * RESTRICT A = static_cast(_A); - const char * RESTRICT B = static_cast(_B); - - // a.qs: 8 groups, 32 bytes each group (m256i) - __m512i va[8]; - // a.bsum: 8 groups, 2 bytes each group (m128i) - __m512i va_bsum; - __m512 vc[COLS]; - __m512 vd1; - - // packed_B: - const int offset_qh = (QK_K / 2) * TILE_N; - const int offset_scales = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N; - const int offset_mins = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N + 8 * TILE_N; - const int offset_d0 = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N + 16 * TILE_N; - const int offset_dmin = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N + 16 * TILE_N + TILE_N * sizeof(ggml_half); - - const __m512i lowMask = _mm512_set1_epi8(0xF); - - auto loadc = [&](auto col) { - vc[col] = _mm512_setzero_ps(); - }; - Unroll{}(loadc); - - // Q5_K and Q4_K shares the same vnni formats, refer to notes above. - auto compute = [&](auto col, auto i) { - // load a - if constexpr (col == 0) { - for (int k_group = 0; k_group < QK_K / 32; ++k_group) { - va[k_group] = _mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)(A[0 * KB + i].qs + k_group * 32))); - } - const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums); - const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1)); - va_bsum = _mm512_castsi128_si512(q8s); - vd1 = _mm512_set1_ps(A[0 * KB + i].d); - } - - // step 1: accumultate the quants - __m512i acc = _mm512_setzero_si512(); - const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); - const char * b_qs = b_ptr; - const char * b_qh = b_ptr + offset_qh; - for (int k_group = 0; k_group < QK_K / 32; ++k_group) { - __m512i vsum = _mm512_setzero_si512(); - __m512i hmask0 = _mm512_set1_epi8(0x1); - __m512i hmask1 = _mm512_set1_epi8(0x2); - __m512i hbits = _mm512_loadu_si512((const __m512i *)(b_qh + k_group * 64)); - for (int k = 0; k < 8; k += 2) { - __m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 0), va[k_group]); - __m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 1), va[k_group]); - - __m512i bytes = _mm512_loadu_si512((const __m512i *)b_qs); - __m512i vb0 = _mm512_and_si512(bytes, lowMask); - __m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); - - __m512i vh0 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask0), k), 4); - __m512i vh1 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask1), k + 1), 4); - - hmask0 = _mm512_slli_epi16(hmask0, 2); - hmask1 = _mm512_slli_epi16(hmask1, 2); - vb0 = _mm512_add_epi8(vb0, vh0); - vb1 = _mm512_add_epi8(vb1, vh1); - - vsum = _mm512_dpbusd_epi32(vsum, vb0, va0); - vsum = _mm512_dpbusd_epi32(vsum, vb1, va1); - - b_qs += 64; - } - // vacc += scale * (q8 @ q5) - const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N))); - acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale)); - } - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0))); - vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]); - - // step 2: accumulate the mins - __m512i acc_m = _mm512_setzero_si512(); - for (int k = 0; k < 4; ++k) { - __m512i vmask = _mm512_set1_epi32(k); - __m512i va = _mm512_permutexvar_epi32(vmask, va_bsum); - __m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_mins + k * 32))); - acc_m = _mm512_dpwssds_epi32(acc_m, va, vb); - } - const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_dmin))); - vc[col] = _mm512_fnmadd_ps(_mm512_cvtepi32_ps(acc_m), _mm512_mul_ps(vdmin, vd1), vc[col]); - }; - - for (int i = 0; i < KB; ++i) { - Unroll{}(compute, i); - } - - //store to C - auto storec = [&](auto col) { - _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); - }; - Unroll{}(storec); - } -}; - -template -struct tinygemm_kernel_vnni { - static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { - - constexpr int COLS = BLOCK_N / 16; - const int TILE_SIZE = TILE_N * sizeof(block_q6_K); - - const block_q8_K * RESTRICT A = static_cast(_A); - const char * RESTRICT B = static_cast(_B); - - // load the 256 bytes from A to 4 avx512 vectors - __m512i va[4]; - __m512 vc[COLS]; - __m512 vd1; - - // packed_B: - const int offset_qh = (QK_K / 2) * TILE_N; - const int offset_scales = (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N; - const int offset_d0 = (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N + 16 * TILE_N; - - // compensation - __m512i vcomp; - - const __m512i m32s = _mm512_set1_epi32(32); - const __m512i lowMask = _mm512_set1_epi8(0xF); - - auto loadc = [&](auto col) { - vc[col] = _mm512_setzero_ps(); - }; - Unroll{}(loadc); - - auto compute = [&](auto col, auto i) { - if constexpr (col == 0) { - // load a - va[0] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 0)); - va[1] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 64)); - va[2] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 128)); - va[3] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 192)); - - const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums); - vcomp = _mm512_mullo_epi32(_mm512_cvtepi16_epi32(q8sums), m32s); - vd1 = _mm512_set1_ps(A[0 * KB + i].d); - } - - // accmulate the quants - __m512i acc = _mm512_setzero_si512(); - const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); - const char * b_qs = b_ptr; - const char * b_qh = b_ptr + offset_qh; - int mask = 0; - for (int k_group = 0; k_group < QK_K / 16; ++k_group) { - int r = k_group >> 2; - __m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]); - __m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]); - - __m512i vsum = _mm512_setzero_si512(); - __m512i hmask = _mm512_set1_epi8(0x3); - - __m512i bytes = _mm512_loadu_si512(b_qs); - __m512i hbits = _mm512_loadu_si512(b_qh); - __m512i vb0 = _mm512_and_si512(bytes, lowMask); - __m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); - __m512i vh0 = _mm512_slli_epi16(_mm512_and_si512(hbits, hmask), 4); - __m512i vh1 = _mm512_slli_epi16(_mm512_and_si512(hbits, _mm512_slli_epi16(hmask, 2)), 2); - - vb0 = _mm512_add_epi8(vb0, vh0); - vb1 = _mm512_add_epi8(vb1, vh1); - vsum = _mm512_dpbusd_epi32(vsum, vb0, va0); - vsum = _mm512_dpbusd_epi32(vsum, vb1, va1); - b_qs += 64; - - va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]); - va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]); - - bytes = _mm512_loadu_si512(b_qs); - vb0 = _mm512_and_si512(bytes, lowMask); - vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); - vh0 = _mm512_and_si512(hbits, _mm512_slli_epi16(hmask, 4)); - vh1 = _mm512_srli_epi16(_mm512_and_si512(hbits, _mm512_slli_epi16(hmask, 6)), 2); - vb0 = _mm512_add_epi8(vb0, vh0); - vb1 = _mm512_add_epi8(vb1, vh1); - vsum = _mm512_dpbusd_epi32(vsum, vb0, va0); - vsum = _mm512_dpbusd_epi32(vsum, vb1, va1); - b_qs += 64; - b_qh += 64; - - // B * A - 32 * A - __m512i vmask = _mm512_set1_epi32(k_group); - vsum = _mm512_sub_epi32(vsum, _mm512_permutexvar_epi32(vmask, vcomp)); - - // vacc += scale * (q8 @ q6) - const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N))); - acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale)); - } - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0))); - vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]); - }; - - for (int i = 0; i < KB; ++i) { - Unroll{}(compute, i); - } - - //store to C - auto storec = [&](int col) { - _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); - }; - Unroll{}(storec); - } -}; - -template -struct tinygemm_kernel_vnni { - static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { - - constexpr int COLS = BLOCK_N / 16; - const int TILE_SIZE = TILE_N * sizeof(block_iq4_xs) + TILE_N * 2; - - const block_q8_K * RESTRICT A = static_cast(_A); - const char * RESTRICT B = static_cast(_B); - - // load the 256 bytes from A to 4 avx512 vectors - __m512i va[4]; - __m512 vc[COLS]; - __m512 vd1; - - // packed_B: - const int offset_scales = (QK_K / 2) * TILE_N ; - const int offset_d0 = (QK_K / 2) * TILE_N + 8 * TILE_N; - - // compensation - __m512i vcomp; - - const __m256i m128s = _mm256_set1_epi16(128); - const __m512i lowMask = _mm512_set1_epi8(0xF); - - const __m512i values128 = _mm512_set_epi8( - 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127, - 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127, - 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127, - 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127 - ); - const __m512i off = _mm512_set1_epi8(static_cast(0x80)); - const __m512i values256 = _mm512_add_epi8(values128, off); - - auto loadc = [&](auto col) { - vc[col] = _mm512_setzero_ps(); - }; - Unroll{}(loadc); - - auto compute = [&](auto col, auto i) { - if constexpr (col == 0) { - // load a - va[0] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 0)); - va[1] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 64)); - va[2] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 128)); - va[3] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 192)); - - // compensation: 128 * A - const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums); - vcomp = _mm512_castsi256_si512(_mm256_madd_epi16(q8sums, m128s)); - vd1 = _mm512_set1_ps(A[0 * KB + i].d); - } - - // accmulate the quants - __m512i acc = _mm512_setzero_si512(); - const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); - const char * b_qs = b_ptr; - int mask = 0; - for (int k_group = 0; k_group < QK_K / 32; ++k_group) { - int r = k_group >> 1; - __m512i vmask = _mm512_set1_epi32(k_group); - __m512i vsum = _mm512_setzero_si512(); - for (int k = 0; k < 8; k += 2) { - __m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]); - __m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]); - - __m512i bytes = _mm512_loadu_si512(b_qs); - __m512i vb0 = _mm512_shuffle_epi8(values256, _mm512_and_si512(bytes, lowMask)); - __m512i vb1 = _mm512_shuffle_epi8(values256, _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask)); - - vsum = _mm512_dpbusd_epi32(vsum, vb0, va0); - vsum = _mm512_dpbusd_epi32(vsum, vb1, va1); - b_qs += 64; - } - // (B + 128) * A - 128 * A - vsum = _mm512_sub_epi32(vsum, _mm512_permutexvar_epi32(vmask, vcomp)); - - // vacc += scale * (q8 @ q4) - const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N))); - acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale)); - } - const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0))); - vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]); - }; - - for (int i = 0; i < KB; ++i) { - Unroll{}(compute, i); - } - - //store to C - auto storec = [&](auto col) { - _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); - }; - Unroll{}(storec); - } -}; - -#define LAUNCH_TINYGEMM_KERNEL_VNNI(NB_SIZE) \ - tinygemm_kernel_vnni::apply( \ - KB, wdata_batch, \ - (const char *)src0->data + src0_offset + PACKED_INDEX(nb * kTilesN, 0, KB, TILE_SIZE), \ - (float *) dst->data + dst_offset + nb_start, ldc) - -template ::value, int>::type = 0> -void tinygemm_kernel_amx(int M, int N, int KB, const void * RESTRICT _A, const void * RESTRICT _B, TC * RESTRICT C, int ldc) { - using packed_B_t = packed_B_type; - const int TILE_SIZE = get_tile_size(); - const bool need_unpack = do_unpack::value; - - GGML_ASSERT(M <= 2 * TILE_M && N == 2 * TILE_N); - const TA * RESTRICT A = static_cast(_A); - const char * RESTRICT B = static_cast(_B); - - const int m0 = std::min(M, TILE_M); - const int m1 = std::max(M - TILE_M, 0); - const int lda = KB * sizeof(TA); - //const int ldb = KB * sizeof(TB); - - alignas(64) static thread_local packed_B_t Tile0[TILE_N * TILE_K]; - alignas(64) static thread_local packed_B_t Tile1[TILE_N * TILE_K]; - alignas(64) static thread_local int8_t Tile23[TILE_M * TILE_K]; - - alignas(64) static thread_local int32_t TileC0[TILE_M * TILE_N * 4]; - alignas(64) static thread_local int32_t TileC1[TILE_M * TILE_N * 4]; - - // double buffering C to interleave avx512 and amx - int32_t * C_cur = TileC0; - int32_t * C_pre = TileC1; - - auto Tile4 = [&](int32_t * base) { return base; }; - auto Tile5 = [&](int32_t * base) { return base + TILE_M * TILE_N; }; - auto Tile6 = [&](int32_t * base) { return base + 2 * TILE_M * TILE_N; }; - auto Tile7 = [&](int32_t * base) { return base + 3 * TILE_M * TILE_N; }; - - if (M == 2 * TILE_M) { - // i = 0 - const char * B_blk0 = B + PACKED_INDEX(0, 0, KB, TILE_SIZE); - const char * B_blk1 = B + PACKED_INDEX(1, 0, KB, TILE_SIZE); - if (need_unpack) { - unpack_B(Tile0, B_blk0); - _tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK); - } else { - _tile_loadd(TMM0, B_blk0, TILE_N * VNNI_BLK); - } - - _tile_zero(TMM4); - _tile_loadd(TMM2, A[0].qs, lda); - _tile_dpbssd(TMM4, TMM2, TMM0); - _tile_stored(TMM4, Tile4(C_pre), TILE_N * sizeof(int32_t)); - - _tile_zero(TMM5); - _tile_loadd(TMM3, A[TILE_M * KB + 0].qs, lda); - _tile_dpbssd(TMM5, TMM3, TMM0); - _tile_stored(TMM5, Tile5(C_pre), TILE_N * sizeof(int32_t)); - - if (need_unpack) { - unpack_B(Tile1, B_blk1); - _tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK); - } else { - _tile_loadd(TMM1, B_blk1, TILE_N * VNNI_BLK); - } - - _tile_zero(TMM6); - _tile_dpbssd(TMM6, TMM2, TMM1); - _tile_stored(TMM6, Tile6(C_pre), TILE_N * sizeof(int32_t)); - - _tile_zero(TMM7); - _tile_dpbssd(TMM7, TMM3, TMM1); - _tile_stored(TMM7, Tile7(C_pre), TILE_N * sizeof(int32_t)); - - for (int i = 1; i < KB; ++i) { - // index of previous iter - const int ii = i - 1; - const char * B_blk0 = B + PACKED_INDEX(0, i, KB, TILE_SIZE); - const char * B_blk1 = B + PACKED_INDEX(1, i, KB, TILE_SIZE); - GGML_DISPATCH_BOOL(ii > 0, is_acc, [&] { - if (need_unpack) { - unpack_B(Tile0, B_blk0); - _tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK); - } else { - _tile_loadd(TMM0, B_blk0, TILE_N * VNNI_BLK); - } - _tile_zero(TMM4); - _tile_loadd(TMM2, A[i].qs, lda); - acc_C::apply(C, ldc, Tile4(C_pre), &A[ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M); - - _tile_dpbssd(TMM4, TMM2, TMM0); - _tile_stored(TMM4, Tile4(C_cur), TILE_N * sizeof(int32_t)); - - _tile_zero(TMM5); - _tile_loadd(TMM3, A[TILE_M * KB + i].qs, lda); - acc_C::apply(C + TILE_M * ldc, ldc, Tile5(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M); - - _tile_dpbssd(TMM5, TMM3, TMM0); - _tile_stored(TMM5, Tile5(C_cur), TILE_N * sizeof(int32_t)); - - if (need_unpack) { - unpack_B(Tile1, B_blk1); - _tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK); - } else { - _tile_loadd(TMM1, B_blk1, TILE_N * VNNI_BLK); - } - _tile_zero(TMM6); - acc_C::apply(C + TILE_N, ldc, Tile6(C_pre), &A[ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M); - - _tile_dpbssd(TMM6, TMM2, TMM1); - _tile_stored(TMM6, Tile6(C_cur), TILE_N * sizeof(int32_t)); - - _tile_zero(TMM7); - acc_C::apply(C + TILE_M * ldc + TILE_N, ldc, Tile7(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M); - - _tile_dpbssd(TMM7, TMM3, TMM1); - _tile_stored(TMM7, Tile7(C_cur), TILE_N * sizeof(int32_t)); - - std::swap(C_cur, C_pre); - }); - } - // final accumulation - { - int ii = KB - 1; - acc_C::apply(C, ldc, Tile4(C_pre), &A[ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M); - acc_C::apply(C + TILE_M * ldc, ldc, Tile5(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M); - acc_C::apply(C + TILE_N, ldc, Tile6(C_pre), &A[ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M); - acc_C::apply(C + TILE_M * ldc + TILE_N, ldc, Tile7(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M); - } - } else { - for (int i = 0; i < KB; ++i) { - _tile_zero(TMM4); - _tile_zero(TMM6); - if (m1 != 0) { - _tile_zero(TMM5); - _tile_zero(TMM7); - } - - const char * B_blk0 = B + PACKED_INDEX(0, i, KB, TILE_SIZE); - const char * B_blk1 = B + PACKED_INDEX(1, i, KB, TILE_SIZE); - if (need_unpack) { - unpack_B(Tile0, B_blk0); - _tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK); - } else { - _tile_loadd(TMM0, B_blk0, TILE_N * VNNI_BLK); - } - - if (need_unpack) { - unpack_B(Tile1, B_blk1); - _tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK); - } else { - _tile_loadd(TMM1, B_blk1, TILE_N * VNNI_BLK); - } - - if (m0 == TILE_M) { - _tile_loadd(TMM2, A[i].qs, lda); - } else { - unpack_A(Tile23, &A[i], KB, m0); - _tile_loadd(TMM2, Tile23, TILE_K); - } - - _tile_dpbssd(TMM4, TMM2, TMM0); - _tile_dpbssd(TMM6, TMM2, TMM1); - - _tile_stored(TMM4, Tile4(C_cur), TILE_N * sizeof(int32_t)); - _tile_stored(TMM6, Tile6(C_cur), TILE_N * sizeof(int32_t)); - - GGML_DISPATCH_BOOL(i > 0, is_acc, [&] { - acc_C::apply(C, ldc, Tile4(C_cur), &A[i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m0); - acc_C::apply(C + TILE_N, ldc, Tile6(C_cur), &A[i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m0); - }); - - if (m1 != 0) { - unpack_A(Tile23, &A[TILE_M * KB + i], KB, m1); - _tile_loadd(TMM3, Tile23, TILE_K); - - _tile_dpbssd(TMM5, TMM3, TMM0); - _tile_dpbssd(TMM7, TMM3, TMM1); - _tile_stored(TMM5, Tile5(C_cur), TILE_N * sizeof(int32_t)); - _tile_stored(TMM7, Tile7(C_cur), TILE_N * sizeof(int32_t)); - GGML_DISPATCH_BOOL(i > 0, is_acc, [&] { - acc_C::apply(C + TILE_M * ldc, ldc, Tile5(C_cur), &A[TILE_M * KB + i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m1); - acc_C::apply(C + TILE_M * ldc + TILE_N, ldc, Tile7(C_cur), &A[TILE_M * KB + i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m1); - }); - } - } - } - return; -} - -template ::value, int>::type = 0> -void tinygemm_kernel_amx(int M, int N, int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { - static_assert(std::is_same::value); - const int TILE_SIZE = get_tile_size(); - - GGML_ASSERT(M <= 2 * TILE_M && N == 2 * TILE_N); - const TA * RESTRICT A = static_cast(_A); - const char * RESTRICT B = static_cast(_B); - - const int m0 = std::min(M, TILE_M); - const int m1 = std::max(M - TILE_M, 0); - //const int lda = KB * sizeof(TA); - - alignas(64) static thread_local int8_t Tile0[TILE_N * TILE_K]; - alignas(64) static thread_local int8_t Tile1[TILE_N * TILE_K]; - alignas(64) static thread_local int8_t Tile23[TILE_M * TILE_K]; - - // mat mul result for each group - alignas(64) static thread_local int32_t Tile4[TILE_M * TILE_N]; - alignas(64) static thread_local int32_t Tile5[TILE_M * TILE_N]; - alignas(64) static thread_local int32_t Tile6[TILE_M * TILE_N]; - alignas(64) static thread_local int32_t Tile7[TILE_M * TILE_N]; - - // sum of each QK_K block, contains 8 groups, int32 - alignas(64) static thread_local int32_t Sumi4[TILE_M * TILE_N]; - alignas(64) static thread_local int32_t Sumi5[TILE_M * TILE_N]; - alignas(64) static thread_local int32_t Sumi6[TILE_M * TILE_N]; - alignas(64) static thread_local int32_t Sumi7[TILE_M * TILE_N]; - - const int k_group_size = std::is_same::value ? 16 : 32; - for (int i = 0; i < KB; ++i) { - // step 1: accumulate the quants across 8 groups, each group with 32 - for (int k = 0; k < QK_K / k_group_size; ++k) { - GGML_DISPATCH_BOOL(k > 0, is_acc, [&] { - _tile_zero(TMM4); - _tile_zero(TMM6); - - unpack_B(Tile0, B + PACKED_INDEX(0, i, KB, TILE_SIZE), k); - _tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK); - - unpack_B(Tile1, B + PACKED_INDEX(1, i, KB, TILE_SIZE), k); - _tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK); - - unpack_A(Tile23, &A[i], KB, k, m0); - _tile_loadd(TMM2, Tile23, TILE_K); - - _tile_dpbssd(TMM4, TMM2, TMM0); - _tile_dpbssd(TMM6, TMM2, TMM1); - - _tile_stored(TMM4, Tile4, TILE_N * sizeof(int32_t)); - _tile_stored(TMM6, Tile6, TILE_N * sizeof(int32_t)); - - scale_C(Tile4, Sumi4, B + PACKED_INDEX(0, i, KB, TILE_SIZE), k, m0); - scale_C(Tile6, Sumi6, B + PACKED_INDEX(1, i, KB, TILE_SIZE), k, m0); - - if (m1 != 0) { - _tile_zero(TMM5); - _tile_zero(TMM7); - - unpack_A(Tile23, &A[TILE_M * KB + i], KB, k, m1); - _tile_loadd(TMM3, Tile23, TILE_K); - - _tile_dpbssd(TMM5, TMM3, TMM0); - _tile_dpbssd(TMM7, TMM3, TMM1); - - _tile_stored(TMM5, Tile5, TILE_N * sizeof(int32_t)); - _tile_stored(TMM7, Tile7, TILE_N * sizeof(int32_t)); - - scale_C(Tile5, Sumi5, B + PACKED_INDEX(0, i, KB, TILE_SIZE), k, m1); - scale_C(Tile7, Sumi7, B + PACKED_INDEX(1, i, KB, TILE_SIZE), k, m1); - } - }); - } - - // step 2: accmulate the mins - GGML_DISPATCH_BOOL(i > 0, is_acc, [&] { - acc_C::apply(C, ldc, Sumi4, &A[i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m0); - acc_C::apply(C + TILE_N, ldc, Sumi6, &A[i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m0); - if (m1 != 0) { - acc_C::apply(C + TILE_M * ldc, ldc, Sumi5, &A[TILE_M * KB + i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m1); - acc_C::apply(C + TILE_M * ldc + TILE_N, ldc, Sumi7, &A[TILE_M * KB + i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m1); - } - }); - } - return; -} - -} // anonymous namespace - -// get the packed tensor size for quantized weights -size_t ggml_backend_amx_get_alloc_size(const struct ggml_tensor * tensor) { - const enum ggml_type TYPE = tensor->type; - - const int K = tensor->ne[0]; // ne0: in_features - const int N = tensor->ne[1]; // ne1: out_features - - auto get_tensor_size = [&] { - size_t row_size_B{0}; - GGML_DISPATCH_QTYPES(TYPE, [&] { - row_size_B = get_row_size(K); - }); - return N * row_size_B; - }; - - if (qtype_has_amx_kernels(TYPE)) { - return get_tensor_size(); - } else { - // for f16, bf16 we don't do packing - return ggml_nbytes(tensor); - } -} - -// pack weight to vnni format -void ggml_backend_amx_convert_weight(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) { - GGML_ASSERT(offset == 0 && size == ggml_nbytes(tensor)); // only full tensor conversion is supported for now - - const enum ggml_type TYPE = tensor->type; - - const int K = tensor->ne[0]; // ne0: in_features - const int N = tensor->ne[1]; // ne1: out_features - - GGML_DISPATCH_QTYPES(TYPE, [&] { - convert_B_packed_format((void *)((char *)tensor->data + offset), (const type *)data, N, K); - }); -} - -// ne2 is passed explicitly to help compiler optimize repeated calls -inline int64_t ggml_batch_offset(const ggml_tensor * t, int64_t batch_idx, int64_t ne2) { - const int64_t i2 = batch_idx % ne2; - const int64_t i3 = batch_idx / ne2; - return i3 * t->nb[3] + i2 * t->nb[2]; -} - -size_t ggml_backend_amx_desired_wsize(const struct ggml_tensor * dst) { - struct ggml_tensor * src0 = dst->src[0]; - - const enum ggml_type TYPE = src0->type; - - const bool is_floating_type = TYPE == GGML_TYPE_F16; - if (is_floating_type) { - return 0; - } - - const int M = dst->ne[1]; - const int K = src0->ne[0]; - const int64_t n_batch = dst->ne[2] * dst->ne[3]; - - size_t desired_wsize = 0; - - GGML_DISPATCH_QTYPES(TYPE, [&] { - const size_t row_size_A = K / blck_size * sizeof(vec_dot_type); - desired_wsize = n_batch * M * row_size_A; - }); - - return desired_wsize; -} - -// NB: mixed dtype gemm with Advanced Matrix Extensions (Intel AMX) -// -// src0: weight in shape of {N, K}, quantized -// src1: input in shape of {M, K}, float32 -// dst: output in shape of {M, N}, float32 -// -// the function performs: dst = src1 @ src0.T for each batch -// -void ggml_backend_amx_mul_mat(const ggml_compute_params * params, struct ggml_tensor * dst) { - struct ggml_tensor * src0 = dst->src[0]; - struct ggml_tensor * src1 = dst->src[1]; - - const enum ggml_type TYPE = src0->type; - - // f16 only has avx512 kernels for now, - // amx kernels will be added once 6th gen xeon is released. - const bool is_floating_type = TYPE == GGML_TYPE_F16; - - const int M = dst->ne[1]; - const int N = dst->ne[0]; - const int K = src0->ne[0]; - const int ldc = dst->nb[1] / dst->nb[0]; - - const int64_t ne2 = dst->ne[2]; - const int64_t n_batch = ne2 * dst->ne[3]; - - if (is_floating_type) { - constexpr int BLOCK_M = 4; - constexpr int BLOCK_N = 6; - const int MB = div_up(M, BLOCK_M); - const int NB = div_up(N, BLOCK_N); - - parallel_for_ggml(params, n_batch * MB * NB, [&](int begin, int end) { - GGML_DISPATCH_FLOATING_TYPES(TYPE, [&] { - for (int i = begin; i < end; ++i) { - int batch_idx = i / (MB * NB); - int remaining = i % (MB * NB); - int mb = remaining / NB; - int nb = remaining % NB; - - int64_t src0_offset = ggml_batch_offset(src0, batch_idx, ne2); - int64_t src1_offset = ggml_batch_offset(src1, batch_idx, ne2); - int64_t dst_offset = ggml_batch_offset(dst, batch_idx, ne2); - - int mb_start = mb * BLOCK_M; - int mb_size = std::min(BLOCK_M, M - mb_start); - int nb_start = nb * BLOCK_N; - int nb_size = std::min(BLOCK_N, N - nb_start); - - switch (mb_size << 4 | nb_size) { - case 0x12: LAUNCH_TINYGEMM_KERNEL_AVX(1, 2); break; - case 0x14: LAUNCH_TINYGEMM_KERNEL_AVX(1, 4); break; - case 0x16: LAUNCH_TINYGEMM_KERNEL_AVX(1, 6); break; - case 0x22: LAUNCH_TINYGEMM_KERNEL_AVX(2, 2); break; - case 0x24: LAUNCH_TINYGEMM_KERNEL_AVX(2, 4); break; - case 0x26: LAUNCH_TINYGEMM_KERNEL_AVX(2, 6); break; - case 0x32: LAUNCH_TINYGEMM_KERNEL_AVX(3, 2); break; - case 0x34: LAUNCH_TINYGEMM_KERNEL_AVX(3, 4); break; - case 0x36: LAUNCH_TINYGEMM_KERNEL_AVX(3, 6); break; - case 0x42: LAUNCH_TINYGEMM_KERNEL_AVX(4, 2); break; - case 0x44: LAUNCH_TINYGEMM_KERNEL_AVX(4, 4); break; - case 0x46: LAUNCH_TINYGEMM_KERNEL_AVX(4, 6); break; - default: fprintf(stderr, "Unexpected block size!\n"); - } - } - }); - }); - return; - } - - // pointer to work space, used convert A from float to quantized type - void * wdata = params->wdata; - - //TODO: performance improvement: merge quant A - // if (params->ith == 0) { - GGML_DISPATCH_QTYPES(TYPE, [&] { - const size_t row_size_A = K / blck_size * sizeof(vec_dot_type); - const size_t desired_wsize = n_batch * M * row_size_A; - if (params->wsize < desired_wsize) { - GGML_ABORT("insufficient work space size"); - } - - // Q4_0, Q4_1, Q8_0 handles 1 TILE_K per blck_size - // Q4_K, Q5_K, Q6_K, IQ4_XS handles 8 TILE_K per blck_size - GGML_ASSERT(TILE_K == blck_size || TILE_K * 8 == blck_size); - - parallel_for_ggml(params, n_batch, [&](int begin, int end) { - for (int batch_idx = begin; batch_idx < end; ++batch_idx) { - int64_t src1_offset = ggml_batch_offset(src1, batch_idx, ne2); - const float * A_data = (const float *)((const char *)src1->data + src1_offset); - char * wdata_batch = (char *)wdata + batch_idx * M * row_size_A; - - for (int m = 0; m < M; ++m) { - from_float(A_data + m * K, wdata_batch + m * row_size_A, K); - } - } - }); - }); - // } - - ggml_barrier(params->threadpool); - - if (M == 1) { - // MB = 1 and handle 8 tiles in each block - constexpr int kTilesN = 4; - constexpr int BLOCK_N = TILE_N * kTilesN; - const int NB = div_up(N, BLOCK_N); - - parallel_for_ggml(params, n_batch * NB, [&](int begin, int end) { - GGML_DISPATCH_QTYPES(TYPE, [&] { - const int KB = K / blck_size; - const int TILE_SIZE = get_tile_size(); - const int row_size_A = KB * sizeof(vec_dot_type); - for (int i = begin; i < end; ++i) { - int batch_idx = i / NB; - int nb = i % NB; - - int64_t src0_offset = ggml_batch_offset(src0, batch_idx, ne2); - int64_t dst_offset = ggml_batch_offset(dst, batch_idx, ne2); - const char * wdata_batch = (const char *)wdata + batch_idx * row_size_A; - - int nb_start = nb * BLOCK_N; - int nb_size = std::min(BLOCK_N, N - nb_start); // 32, 64, 96 - - switch (nb_size) { - //case 160: LAUNCH_TINYGEMM_KERNEL_VNNI(160); break; - case 128: LAUNCH_TINYGEMM_KERNEL_VNNI(128); break; - case 96: LAUNCH_TINYGEMM_KERNEL_VNNI(96); break; - case 64: LAUNCH_TINYGEMM_KERNEL_VNNI(64); break; - case 32: LAUNCH_TINYGEMM_KERNEL_VNNI(32); break; - default: fprintf(stderr, "Unexpected n block size!\n"); - } - } - }); - }); - return; - } - - // handle 4 tiles at a tile - constexpr int BLOCK_M = TILE_M * 2; - constexpr int BLOCK_N = TILE_N * 2; - const int MB = div_up(M, BLOCK_M); - const int NB = div_up(N, BLOCK_N); - - parallel_for_ggml(params, n_batch * MB * NB, [&](int begin, int end) { - // init tile config for each thread - ggml_tile_config_init(); - - GGML_DISPATCH_QTYPES(TYPE, [&] { - const int KB = K / blck_size; - const int TILE_SIZE = get_tile_size(); - const int row_size_A = KB * sizeof(vec_dot_type); - - for (int i = begin; i < end; ++i) { - int batch_idx = i / (MB * NB); - int remaining = i % (MB * NB); - int mb = remaining / NB; - int nb = remaining % NB; - - int64_t src0_offset = ggml_batch_offset(src0, batch_idx, ne2); - int64_t dst_offset = ggml_batch_offset(dst, batch_idx, ne2); - const char * wdata_batch = (const char *)wdata + batch_idx * M * row_size_A; - - int mb_start = mb * BLOCK_M; - int mb_size = std::min(BLOCK_M, M - mb_start); - int nb_start = nb * BLOCK_N; - int nb_size = BLOCK_N; - - tinygemm_kernel_amx( - mb_size, nb_size, KB, - wdata_batch + mb_start * row_size_A, - (const char *)src0->data + src0_offset + PACKED_INDEX(nb * 2, 0, KB, TILE_SIZE), - (float *) dst->data + dst_offset + mb_start * N + nb_start, ldc); - } - }); - }); -} - -#endif // if defined(__AMX_INT8__) && defined(__AVX512VNNI__) diff --git a/ggml/src/ggml-cpu/amx/mmq.h b/ggml/src/ggml-cpu/amx/mmq.h deleted file mode 100644 index baf768477345..000000000000 --- a/ggml/src/ggml-cpu/amx/mmq.h +++ /dev/null @@ -1,10 +0,0 @@ -#pragma once -#include "common.h" - -size_t ggml_backend_amx_desired_wsize(const struct ggml_tensor * dst); - -size_t ggml_backend_amx_get_alloc_size(const struct ggml_tensor * tensor); - -void ggml_backend_amx_convert_weight(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size); - -void ggml_backend_amx_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-cpu/arch-fallback.h b/ggml/src/ggml-cpu/arch-fallback.h deleted file mode 100644 index 595ded09f03f..000000000000 --- a/ggml/src/ggml-cpu/arch-fallback.h +++ /dev/null @@ -1,349 +0,0 @@ - -#pragma once - -// Rename `_generic` functions if no native implementation is available. -// This effectively selects the generic implementation. - -#if defined(GGML_CPU_GENERIC) -// quants.c -#define quantize_row_q8_0_generic quantize_row_q8_0 -#define quantize_row_q8_1_generic quantize_row_q8_1 -#define quantize_row_q8_K_generic quantize_row_q8_K -#define ggml_vec_dot_q4_0_q8_0_generic ggml_vec_dot_q4_0_q8_0 -#define ggml_vec_dot_q4_1_q8_1_generic ggml_vec_dot_q4_1_q8_1 -#define ggml_vec_dot_q5_0_q8_0_generic ggml_vec_dot_q5_0_q8_0 -#define ggml_vec_dot_q5_1_q8_1_generic ggml_vec_dot_q5_1_q8_1 -#define ggml_vec_dot_q8_0_q8_0_generic ggml_vec_dot_q8_0_q8_0 -#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 -#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 -#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 -#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K -#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K -#define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K -#define ggml_vec_dot_q3_K_q8_K_generic ggml_vec_dot_q3_K_q8_K -#define ggml_vec_dot_q4_K_q8_K_generic ggml_vec_dot_q4_K_q8_K -#define ggml_vec_dot_q5_K_q8_K_generic ggml_vec_dot_q5_K_q8_K -#define ggml_vec_dot_q6_K_q8_K_generic ggml_vec_dot_q6_K_q8_K -#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K -#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K -#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K -#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K -#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K -#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K -#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K -#define ggml_vec_dot_iq4_nl_q8_0_generic ggml_vec_dot_iq4_nl_q8_0 -#define ggml_vec_dot_iq4_xs_q8_K_generic ggml_vec_dot_iq4_xs_q8_K -// repack.cpp -#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 -#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 -#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 -#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 -#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 -#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 -#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 -#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K -#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K -#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K -#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K -#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K -#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K -#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K -#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 -#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 -#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 -#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 -#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 -#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 -#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 -#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 -#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 -#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K -#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K -#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K -#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K -#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K -#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K -#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K -#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 -#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 -#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 -#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 -#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 -#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 -#elif defined(__aarch64__) || defined(__arm__) || defined(_M_ARM) || defined(_M_ARM64) -// repack.cpp -#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 -#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 -#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 -#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 -#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K -#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 -#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 -#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K -#elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64) -// quants.c -#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 -// repack.cpp -#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 -#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 -#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 -#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 -#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K -#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K -#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K -#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K -#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K -#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 -#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 -#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 -#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 -#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 -#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 -#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K -#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K -#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K -#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K -#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K -#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 -#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 -#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 -#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 -#elif defined(__POWERPC__) || defined(__powerpc__) -// ref: https://github.com/ggml-org/llama.cpp/pull/14146#issuecomment-2972561679 -// quants.c -#define quantize_row_q8_K_generic quantize_row_q8_K -#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 -#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 -#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K -#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K -#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K -// repack.cpp -#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 -#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 -#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 -#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 -#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 -#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 -#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 -#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K -#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K -#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K -#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K -#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K -#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K -#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K -#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 -#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 -#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 -#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 -#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 -#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 -#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 -#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 -#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 -#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K -#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K -#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K -#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K -#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K -#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K -#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K -#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 -#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 -#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 -#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 -#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 -#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 -#elif defined(__loongarch64) -// quants.c -#define quantize_row_q8_K_generic quantize_row_q8_K -#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K -#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K -#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K -#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 -#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 -#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 -// repack.cpp -#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 -#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 -#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 -#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 -#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 -#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 -#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 -#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K -#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K -#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K -#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K -#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K -#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K -#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K -#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 -#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 -#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 -#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 -#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 -#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 -#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 -#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 -#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 -#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K -#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K -#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K -#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K -#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K -#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K -#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K -#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 -#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 -#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 -#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 -#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 -#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 -#elif defined(__riscv) -// quants.c -#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 -#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 -// repack.cpp -#define ggml_quantize_mat_q8_0_4x1_generic ggml_quantize_mat_q8_0_4x1 -#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 -#define ggml_quantize_mat_q8_K_4x1_generic ggml_quantize_mat_q8_K_4x1 -#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 -#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 -#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 -#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 -#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K -#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K -#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K -#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K -#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K -#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K -#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K -#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 -#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 -#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 -#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 -#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 -#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 -#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 -#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 -#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K -#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K -#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K -#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K -#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K -#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K -#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K -#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 -#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 -#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 -#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 -#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 -#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 -#elif defined(__s390x__) -// quants.c -#define quantize_row_q8_K_generic quantize_row_q8_K -#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 -#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 -#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K -#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K -#define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K -#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K -#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K -#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K -#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K -#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K -#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K -#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K -// repack.cpp -#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 -#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 -#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 -#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 -#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 -#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 -#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 -#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K -#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K -#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K -#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K -#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K -#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K -#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K -#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 -#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 -#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 -#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 -#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 -#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 -#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 -#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 -#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 -#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K -#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K -#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K -#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K -#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K -#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K -#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K -#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 -#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 -#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 -#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 -#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 -#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 -#elif defined(__wasm__) -// quants.c -#define ggml_vec_dot_q4_1_q8_1_generic ggml_vec_dot_q4_1_q8_1 -#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K -#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K -#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K -#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K -#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K -#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K -#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K -#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K -#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K -#define ggml_vec_dot_iq4_nl_q8_0_generic ggml_vec_dot_iq4_nl_q8_0 -#define ggml_vec_dot_iq4_xs_q8_K_generic ggml_vec_dot_iq4_xs_q8_K -#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 -#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 -#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 -// repack.cpp -#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 -#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 -#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 -#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 -#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 -#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 -#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 -#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K -#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K -#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K -#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K -#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K -#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K -#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K -#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 -#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 -#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 -#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 -#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 -#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 -#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 -#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 -#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 -#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K -#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K -#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K -#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K -#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K -#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K -#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K -#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 -#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 -#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 -#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 -#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 -#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 -#endif diff --git a/ggml/src/ggml-cpu/arch/arm/cpu-feats-rust/Cargo.toml b/ggml/src/ggml-cpu/arch/arm/cpu-feats-rust/Cargo.toml new file mode 100644 index 000000000000..65677cfb8ac5 --- /dev/null +++ b/ggml/src/ggml-cpu/arch/arm/cpu-feats-rust/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "ggml-cpu-arm-feats-rust" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["staticlib", "rlib"] diff --git a/ggml/src/ggml-cpu/arch/arm/cpu-feats-rust/src/lib.rs b/ggml/src/ggml-cpu/arch/arm/cpu-feats-rust/src/lib.rs new file mode 100644 index 000000000000..92fe3814dd9a --- /dev/null +++ b/ggml/src/ggml-cpu/arch/arm/cpu-feats-rust/src/lib.rs @@ -0,0 +1,176 @@ +#![allow(unexpected_cfgs)] + +#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +use std::ffi::c_ulong; +#[cfg(all(target_arch = "aarch64", target_vendor = "apple"))] +use std::ffi::{c_char, c_int, c_void}; + +#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +const AT_HWCAP: c_ulong = 16; +#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +const AT_HWCAP2: c_ulong = 26; +#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +const HWCAP_FPHP: c_ulong = 1 << 9; +#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +const HWCAP_ASIMDDP: c_ulong = 1 << 20; +#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +const HWCAP_SVE: c_ulong = 1 << 22; +#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +const HWCAP2_SVE2: c_ulong = 1 << 1; +#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +const HWCAP2_I8MM: c_ulong = 1 << 13; +#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +const HWCAP2_SME: c_ulong = 1 << 23; + +#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +extern "C" { + fn getauxval(key: c_ulong) -> c_ulong; +} + +#[cfg(all(target_arch = "aarch64", target_vendor = "apple"))] +extern "C" { + fn sysctlbyname( + name: *const c_char, + oldp: *mut c_void, + oldlenp: *mut usize, + newp: *mut c_void, + newlen: usize, + ) -> c_int; +} + +#[derive(Clone, Copy, Default)] +struct ArmFeatures { + has_dotprod: bool, + has_fp16_va: bool, + has_sve: bool, + has_sve2: bool, + has_i8mm: bool, + has_sme: bool, +} + +#[cfg(all(target_arch = "aarch64", target_os = "linux"))] +fn detect_features() -> ArmFeatures { + let hwcap = unsafe { getauxval(AT_HWCAP) }; + let hwcap2 = unsafe { getauxval(AT_HWCAP2) }; + ArmFeatures { + has_dotprod: (hwcap & HWCAP_ASIMDDP) != 0, + has_fp16_va: (hwcap & HWCAP_FPHP) != 0, + has_sve: (hwcap & HWCAP_SVE) != 0, + has_sve2: (hwcap2 & HWCAP2_SVE2) != 0, + has_i8mm: (hwcap2 & HWCAP2_I8MM) != 0, + has_sme: (hwcap2 & HWCAP2_SME) != 0, + } +} + +#[cfg(all(target_arch = "aarch64", target_vendor = "apple"))] +fn apple_sysctl_bool(name: *const c_char) -> bool { + let mut oldp = 0_i32; + let mut size = std::mem::size_of_val(&oldp); + let ret = unsafe { + sysctlbyname( + name, + &mut oldp as *mut i32 as *mut c_void, + &mut size as *mut usize, + std::ptr::null_mut(), + 0, + ) + }; + ret == 0 && oldp != 0 +} + +#[cfg(all(target_arch = "aarch64", target_vendor = "apple"))] +fn detect_features() -> ArmFeatures { + ArmFeatures { + has_dotprod: apple_sysctl_bool(c"hw.optional.arm.FEAT_DotProd".as_ptr()), + has_fp16_va: false, + has_sve: false, + has_sve2: false, + has_i8mm: apple_sysctl_bool(c"hw.optional.arm.FEAT_I8MM".as_ptr()), + has_sme: apple_sysctl_bool(c"hw.optional.arm.FEAT_SME".as_ptr()), + } +} + +#[cfg(not(any( + all(target_arch = "aarch64", target_os = "linux"), + all(target_arch = "aarch64", target_vendor = "apple") +)))] +#[allow(dead_code)] +fn detect_features() -> ArmFeatures { + ArmFeatures::default() +} + +fn score_for_features(features: ArmFeatures) -> i32 { + let mut score = 1_i32; + + if cfg!(ggml_use_dotprod) { + if !features.has_dotprod { + return 0; + } + score += 1 << 1; + } + if cfg!(ggml_use_fp16_va) { + if !features.has_fp16_va { + return 0; + } + score += 1 << 2; + } + if cfg!(ggml_use_sve) { + if !features.has_sve { + return 0; + } + score += 1 << 3; + } + if cfg!(ggml_use_matmul_int8) { + if !features.has_i8mm { + return 0; + } + score += 1 << 4; + } + if cfg!(ggml_use_sve2) { + if !features.has_sve2 { + return 0; + } + score += 1 << 5; + } + if cfg!(ggml_use_sme) { + if !features.has_sme { + return 0; + } + score += 1 << 6; + } + + score +} + +#[cfg(target_arch = "aarch64")] +#[no_mangle] +pub extern "C" fn ggml_backend_score() -> i32 { + score_for_features(detect_features()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn baseline_score_without_cfg_requirements() { + assert_eq!(score_for_features(ArmFeatures::default()), 1); + assert_eq!( + score_for_features(ArmFeatures { + has_dotprod: true, + has_fp16_va: true, + has_sve: true, + has_sve2: true, + has_i8mm: true, + has_sme: true, + }), + 1 + ); + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn feature_probe_is_callable() { + let _ = detect_features(); + } +} diff --git a/ggml/src/ggml-cpu/arch/arm/cpu-feats.cpp b/ggml/src/ggml-cpu/arch/arm/cpu-feats.cpp deleted file mode 100644 index c460c5491143..000000000000 --- a/ggml/src/ggml-cpu/arch/arm/cpu-feats.cpp +++ /dev/null @@ -1,98 +0,0 @@ -#include "ggml-backend-impl.h" - -#if defined(__aarch64__) - -#if defined(__linux__) -#include -#elif defined(__APPLE__) -#include -#endif - -#if !defined(HWCAP2_SVE2) -#define HWCAP2_SVE2 (1 << 1) -#endif - -#if !defined(HWCAP2_I8MM) -#define HWCAP2_I8MM (1 << 13) -#endif - -#if !defined(HWCAP2_SME) -#define HWCAP2_SME (1 << 23) -#endif - -struct aarch64_features { - // has_neon not needed, aarch64 has NEON guaranteed - bool has_dotprod = false; - bool has_fp16_va = false; - bool has_sve = false; - bool has_sve2 = false; - bool has_i8mm = false; - bool has_sme = false; - - aarch64_features() { -#if defined(__linux__) - uint32_t hwcap = getauxval(AT_HWCAP); - uint32_t hwcap2 = getauxval(AT_HWCAP2); - - has_dotprod = !!(hwcap & HWCAP_ASIMDDP); - has_fp16_va = !!(hwcap & HWCAP_FPHP); - has_sve = !!(hwcap & HWCAP_SVE); - has_sve2 = !!(hwcap2 & HWCAP2_SVE2); - has_i8mm = !!(hwcap2 & HWCAP2_I8MM); - has_sme = !!(hwcap2 & HWCAP2_SME); -#elif defined(__APPLE__) - int oldp = 0; - size_t size = sizeof(oldp); - - if (sysctlbyname("hw.optional.arm.FEAT_DotProd", &oldp, &size, NULL, 0) == 0) { - has_dotprod = static_cast(oldp); - } - - if (sysctlbyname("hw.optional.arm.FEAT_I8MM", &oldp, &size, NULL, 0) == 0) { - has_i8mm = static_cast(oldp); - } - - if (sysctlbyname("hw.optional.arm.FEAT_SME", &oldp, &size, NULL, 0) == 0) { - has_sme = static_cast(oldp); - } - - // Apple apparently does not implement SVE yet -#endif - } -}; - -static int ggml_backend_cpu_aarch64_score() { - int score = 1; - aarch64_features af; - -#ifdef GGML_USE_DOTPROD - if (!af.has_dotprod) { return 0; } - score += 1<<1; -#endif -#ifdef GGML_USE_FP16_VECTOR_ARITHMETIC - if (!af.has_fp16_va) { return 0; } - score += 1<<2; -#endif -#ifdef GGML_USE_SVE - if (!af.has_sve) { return 0; } - score += 1<<3; -#endif -#ifdef GGML_USE_MATMUL_INT8 - if (!af.has_i8mm) { return 0; } - score += 1<<4; -#endif -#ifdef GGML_USE_SVE2 - if (!af.has_sve2) { return 0; } - score += 1<<5; -#endif -#ifdef GGML_USE_SME - if (!af.has_sme) { return 0; } - score += 1<<6; -#endif - - return score; -} - -GGML_BACKEND_DL_SCORE_IMPL(ggml_backend_cpu_aarch64_score) - -# endif // defined(__aarch64__) diff --git a/ggml/src/ggml-cpu/arch/arm/quants.c b/ggml/src/ggml-cpu/arch/arm/quants.c.inc similarity index 99% rename from ggml/src/ggml-cpu/arch/arm/quants.c rename to ggml/src/ggml-cpu/arch/arm/quants.c.inc index fe6213329708..026b11ad9908 100644 --- a/ggml/src/ggml-cpu/arch/arm/quants.c +++ b/ggml/src/ggml-cpu/arch/arm/quants.c.inc @@ -1,12 +1,12 @@ #define GGML_COMMON_IMPL_C -#include "ggml-common.h" -#include "ggml-quants.h" -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "simd-mappings.h" +#include "ggml-common-defs.inc" +#include "ggml-quants-defs.inc" +#include "ggml-impl-defs.inc" +#include "ggml-cpu.h.inc" +#include "simd-mappings-defs.inc" #include "../../quants.h" -#include "../../ggml-cpu-impl.h" +#include "../../ggml-cpu-impl-defs.inc" #include #include diff --git a/ggml/src/ggml-cpu/arch/arm/repack.cpp b/ggml/src/ggml-cpu/arch/arm/repack.cpp.inc similarity index 98% rename from ggml/src/ggml-cpu/arch/arm/repack.cpp rename to ggml/src/ggml-cpu/arch/arm/repack.cpp.inc index a7534443091f..60e772127770 100644 --- a/ggml/src/ggml-cpu/arch/arm/repack.cpp +++ b/ggml/src/ggml-cpu/arch/arm/repack.cpp.inc @@ -1,13 +1,140 @@ #define GGML_COMMON_IMPL_CPP #define GGML_COMMON_DECL_CPP -#include "ggml-common.h" -#include "ggml-backend-impl.h" - -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "ggml-cpu-impl.h" -#include "simd-mappings.h" -#include "traits.h" +#include "ggml-common-defs.inc" +#include "ggml-backend-impl-defs.inc" + +#include "ggml-impl-defs.inc" +#include "ggml-cpu.h.inc" +#include "ggml-cpu-impl-defs.inc" +#include "simd-mappings-defs.inc" +#include "ggml.h.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml-cpu-impl-defs.inc" +#include "ggml-impl-defs.inc" +#include "simd-mappings-defs.inc" + +#ifdef __cplusplus +#include +extern "C" { +#endif + +// return true if op part of extra "accelerator" +bool ggml_cpu_extra_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op); +bool ggml_cpu_extra_work_size(int n_threads, const struct ggml_tensor * op, size_t * size); + +#ifdef __cplusplus +} + +namespace ggml::cpu { +// register in tensor->extra +class tensor_traits { + public: + virtual ~tensor_traits(); + virtual bool work_size(int n_threads, const struct ggml_tensor * op, size_t & size) = 0; + virtual bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) = 0; +}; + +class extra_buffer_type { + public: + virtual ~extra_buffer_type(); + virtual bool supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) = 0; + virtual tensor_traits * get_tensor_traits(const struct ggml_tensor * op) = 0; +}; +} // namespace ggml::cpu + +// implemented in ggml-cpu.cpp. +std::vector & ggml_backend_cpu_get_extra_buffer_types(); + +#endif + +#define GGML_FA_TILE_Q 64 +#define GGML_FA_TILE_KV 64 + +#ifdef __cplusplus + +#include + +// convenience functions/macros for use in template calls +// note: these won't be required after the 'traits' lookup table is used. +static inline ggml_fp16_t f32_to_f16(float x) { + return GGML_CPU_FP32_TO_FP16(x); +} + +static inline float f16_to_f32(ggml_fp16_t x) { + return GGML_CPU_FP16_TO_FP32(x); +} + +static inline ggml_bf16_t f32_to_bf16(float x) { + return GGML_FP32_TO_BF16(x); +} + +static inline float bf16_to_f32(ggml_bf16_t x) { + return GGML_BF16_TO_FP32(x); +} + +static inline float i32_to_f32(int32_t x) { + return x; +} + +static inline int32_t f32_to_i32(float x) { + return x; +} + +static inline float f32_to_f32(float x) { + return x; +} + +// TODO - merge this into the traits table, after using row-based conversions +template +struct type_conversion_table; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_fp16_t) = f16_to_f32; + static constexpr ggml_fp16_t (*from_f32)(float) = f32_to_f16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(float) = f32_to_f32; + static constexpr float (*from_f32)(float) = f32_to_f32; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_bf16_t) = bf16_to_f32; + static constexpr ggml_bf16_t (*from_f32)(float) = f32_to_bf16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(int32_t) = i32_to_f32; + static constexpr int32_t (*from_f32)(float) = f32_to_i32; +}; + +static std::pair get_thread_range(const struct ggml_compute_params * params, const struct ggml_tensor * src0) { + const int64_t ith = params->ith; + const int64_t nth = params->nth; + + const int64_t nr = ggml_nrows(src0); + + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; + + // row range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); + + return {ir0, ir1}; +} + +struct ggml_fa_tile_config { + static constexpr size_t Q = GGML_FA_TILE_Q; + static constexpr size_t KV = GGML_FA_TILE_KV; +}; + +#endif + #include #include diff --git a/ggml/src/ggml-cpu/arch/loongarch/quants.c b/ggml/src/ggml-cpu/arch/loongarch/quants.c.inc similarity index 99% rename from ggml/src/ggml-cpu/arch/loongarch/quants.c rename to ggml/src/ggml-cpu/arch/loongarch/quants.c.inc index 74e0c086c6de..b4e99d6300bd 100644 --- a/ggml/src/ggml-cpu/arch/loongarch/quants.c +++ b/ggml/src/ggml-cpu/arch/loongarch/quants.c.inc @@ -1,12 +1,12 @@ #define GGML_COMMON_IMPL_C -#include "ggml-common.h" -#include "ggml-quants.h" -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "simd-mappings.h" +#include "ggml-common-defs.inc" +#include "ggml-quants-defs.inc" +#include "ggml-impl-defs.inc" +#include "ggml-cpu.h.inc" +#include "simd-mappings-defs.inc" #include "../../quants.h" -#include "../../ggml-cpu-impl.h" +#include "../../ggml-cpu-impl-defs.inc" #include #include diff --git a/ggml/src/ggml-cpu/arch/powerpc/cpu-feats-rust/Cargo.toml b/ggml/src/ggml-cpu/arch/powerpc/cpu-feats-rust/Cargo.toml new file mode 100644 index 000000000000..e7c8624f4dc5 --- /dev/null +++ b/ggml/src/ggml-cpu/arch/powerpc/cpu-feats-rust/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "ggml-cpu-powerpc-feats-rust" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["staticlib", "rlib"] diff --git a/ggml/src/ggml-cpu/arch/powerpc/cpu-feats-rust/src/lib.rs b/ggml/src/ggml-cpu/arch/powerpc/cpu-feats-rust/src/lib.rs new file mode 100644 index 000000000000..82ea2e3c472f --- /dev/null +++ b/ggml/src/ggml-cpu/arch/powerpc/cpu-feats-rust/src/lib.rs @@ -0,0 +1,148 @@ +#![allow(unexpected_cfgs)] + +#[cfg(target_arch = "powerpc64")] +use std::ffi::{c_ulong, CStr}; + +#[cfg(target_arch = "powerpc64")] +const AT_PLATFORM: c_ulong = 15; + +#[cfg(target_arch = "powerpc64")] +extern "C" { + fn getauxval(key: c_ulong) -> c_ulong; +} + +#[derive(Clone, Copy, Default)] +struct PowerpcFeatures { + power_version: i32, + has_vsx: bool, +} + +fn parse_power_version(platform: &str) -> i32 { + if !platform.starts_with("power") { + return -1; + } + + let digits_start = platform + .char_indices() + .rev() + .find_map(|(idx, ch)| if ch.is_ascii_digit() { Some(idx) } else { None }); + let Some(mut start) = digits_start else { + return -1; + }; + + for (idx, ch) in platform[..start].char_indices().rev() { + if ch.is_ascii_digit() { + start = idx; + } else { + break; + } + } + + platform[start..].parse::().unwrap_or(-1) +} + +#[cfg(target_arch = "powerpc64")] +fn detect_features() -> PowerpcFeatures { + let platform_ptr = unsafe { getauxval(AT_PLATFORM) } as *const std::ffi::c_char; + let power_version = if platform_ptr.is_null() { + -1 + } else { + let platform = unsafe { CStr::from_ptr(platform_ptr) }; + platform + .to_str() + .map(parse_power_version) + .unwrap_or(-1) + }; + + PowerpcFeatures { + power_version, + has_vsx: power_version >= 9, + } +} + +#[cfg(not(target_arch = "powerpc64"))] +fn detect_features() -> PowerpcFeatures { + PowerpcFeatures { + power_version: -1, + has_vsx: false, + } +} + +fn score_for_features(features: PowerpcFeatures) -> i32 { + let mut score = 1_i32; + + if cfg!(ggml_use_power7) { + if features.power_version < 7 { + return 0; + } + score += 1 << 1; + } + if cfg!(ggml_use_power8) { + if features.power_version < 8 { + return 0; + } + score += 1 << 2; + } + if cfg!(ggml_use_power9) { + if features.power_version < 9 { + return 0; + } + score += 1 << 3; + } + if cfg!(ggml_use_power10) { + if features.power_version < 10 { + return 0; + } + score += 1 << 4; + } + if cfg!(ggml_use_power11) { + if features.power_version < 11 { + return 0; + } + score += 1 << 5; + } + if cfg!(ggml_use_vsx) { + if !features.has_vsx { + return 0; + } + score += 1 << 6; + } + + score +} + +#[no_mangle] +pub extern "C" fn ggml_backend_score() -> i32 { + score_for_features(detect_features()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_power_platform_suffix() { + assert_eq!(parse_power_version("power9"), 9); + assert_eq!(parse_power_version("power10"), 10); + assert_eq!(parse_power_version("powerpc"), -1); + assert_eq!(parse_power_version("POWER10"), -1); + assert_eq!(parse_power_version("unknown10"), -1); + } + + #[test] + fn baseline_score_without_cfg_requirements() { + assert_eq!( + score_for_features(PowerpcFeatures { + power_version: -1, + has_vsx: false, + }), + 1 + ); + } + + #[cfg(target_arch = "powerpc64")] + #[test] + fn getauxval_is_callable() { + let _ = detect_features(); + } +} diff --git a/ggml/src/ggml-cpu/arch/powerpc/cpu-feats.cpp b/ggml/src/ggml-cpu/arch/powerpc/cpu-feats.cpp deleted file mode 100644 index fedd6430278c..000000000000 --- a/ggml/src/ggml-cpu/arch/powerpc/cpu-feats.cpp +++ /dev/null @@ -1,82 +0,0 @@ -# include "ggml-backend-impl.h" - -#if defined(__powerpc64__) || defined(__ppc64__) || defined(__PPC64__) - -#if defined(__linux__) -#include -#endif - -#include - -struct powerpc_features { - std::string platform = ""; - int power_version = -1; - - bool has_vsx = false; - - powerpc_features() { -#if defined(__linux__) - unsigned long auxval = getauxval(AT_PLATFORM); - if (auxval) { - platform = std::string(reinterpret_cast(auxval)); - // TBD: Do systems exist that return this in uppercase? - if (platform.substr(0, 5) == "power") { - // Extractt a numeric suffix, if one exists - int vpos = -1; - for (int i = platform.length() - 1; i >= 0; i--) { - if (std::isdigit(platform[i])) { - vpos = i; - } else { - break; - } - } - if (vpos > -1) { - power_version = std::stoi(platform.substr(vpos)); - } - } - } -#endif - if (power_version >= 9) { - has_vsx = true; - } - } -}; - -static int ggml_backend_cpu_powerpc_score() { - int score = 1; - powerpc_features pf; - -// Platform scores -#if defined(GGML_USE_POWER7) - if (pf.power_version < 7) { return 0; } - score += 1<<1; -#endif -#if defined(GGML_USE_POWER8) - if (pf.power_version < 8) { return 0; } - score += 1<<2; -#endif -#if defined(GGML_USE_POWER9) - if (pf.power_version < 9) { return 0; } - score += 1<<3; -#endif -#if defined(GGML_USE_POWER10) - if (pf.power_version < 10) { return 0; } - score += 1<<4; -#endif -#if defined(GGML_USE_POWER11) - if (pf.power_version < 11) { return 0; } - score += 1<<5; -#endif - -// Feature scores -#if defined(GGML_USE_VSX) - if (!pf.has_vsx) { return 0; } - score += 1<<6; -#endif - - return score; -} - -GGML_BACKEND_DL_SCORE_IMPL(ggml_backend_cpu_powerpc_score) - -#endif // defined(__powerpc64__) || defined(__ppc64__) || defined(__PPC64__) diff --git a/ggml/src/ggml-cpu/arch/powerpc/quants.c b/ggml/src/ggml-cpu/arch/powerpc/quants.c.inc similarity index 99% rename from ggml/src/ggml-cpu/arch/powerpc/quants.c rename to ggml/src/ggml-cpu/arch/powerpc/quants.c.inc index 644c380c7381..bb1900658275 100644 --- a/ggml/src/ggml-cpu/arch/powerpc/quants.c +++ b/ggml/src/ggml-cpu/arch/powerpc/quants.c.inc @@ -1,12 +1,12 @@ #define GGML_COMMON_IMPL_C -#include "ggml-common.h" -#include "ggml-quants.h" -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "simd-mappings.h" +#include "ggml-common-defs.inc" +#include "ggml-quants-defs.inc" +#include "ggml-impl-defs.inc" +#include "ggml-cpu.h.inc" +#include "simd-mappings-defs.inc" #include "../../quants.h" -#include "../../ggml-cpu-impl.h" +#include "../../ggml-cpu-impl-defs.inc" #include #include diff --git a/ggml/src/ggml-cpu/arch/riscv/cpu-feats-rust/Cargo.toml b/ggml/src/ggml-cpu/arch/riscv/cpu-feats-rust/Cargo.toml new file mode 100644 index 000000000000..332c2a4f6dc1 --- /dev/null +++ b/ggml/src/ggml-cpu/arch/riscv/cpu-feats-rust/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "ggml-cpu-riscv-feats-rust" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["staticlib", "rlib"] diff --git a/ggml/src/ggml-cpu/arch/riscv/cpu-feats-rust/src/lib.rs b/ggml/src/ggml-cpu/arch/riscv/cpu-feats-rust/src/lib.rs new file mode 100644 index 000000000000..80fc58d1c2dd --- /dev/null +++ b/ggml/src/ggml-cpu/arch/riscv/cpu-feats-rust/src/lib.rs @@ -0,0 +1,84 @@ +#![allow(unexpected_cfgs)] + +#[cfg(target_arch = "riscv64")] +use std::ffi::{c_int, c_long, c_void}; + +#[cfg(target_arch = "riscv64")] +#[repr(C)] +struct RiscvHwprobe { + key: i64, + value: u64, +} + +#[cfg(target_arch = "riscv64")] +const NR_RISCV_HWPROBE: c_long = 258; +#[cfg(target_arch = "riscv64")] +const RISCV_HWPROBE_KEY_IMA_EXT_0: i64 = 4; +#[cfg(target_arch = "riscv64")] +const RISCV_HWPROBE_IMA_V: u64 = 1 << 2; + +#[cfg(target_arch = "riscv64")] +extern "C" { + fn syscall(num: c_long, ...) -> c_long; +} + +#[cfg(target_arch = "riscv64")] +fn has_rvv() -> bool { + let mut probe = RiscvHwprobe { + key: RISCV_HWPROBE_KEY_IMA_EXT_0, + value: 0, + }; + + let ret = unsafe { + syscall( + NR_RISCV_HWPROBE, + &mut probe as *mut RiscvHwprobe, + 1_usize, + 0_usize, + std::ptr::null::(), + 0_usize, + ) + } as c_int; + + ret == 0 && (probe.value & RISCV_HWPROBE_IMA_V) != 0 +} + +#[cfg(not(target_arch = "riscv64"))] +fn has_rvv() -> bool { + false +} + +fn score_for_features(has_rvv: bool) -> i32 { + let mut score = 1_i32; + + if cfg!(ggml_use_rvv) { + if !has_rvv { + return 0; + } + score += 1 << 1; + } + + score +} + +#[no_mangle] +pub extern "C" fn ggml_backend_score() -> i32 { + score_for_features(has_rvv()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn baseline_score_without_cfg_requirement() { + assert_eq!(score_for_features(false), 1); + assert_eq!(score_for_features(true), 1); + } + + #[cfg(target_arch = "riscv64")] + #[test] + fn hwprobe_is_callable() { + let _ = has_rvv(); + } +} diff --git a/ggml/src/ggml-cpu/arch/riscv/cpu-feats.cpp b/ggml/src/ggml-cpu/arch/riscv/cpu-feats.cpp deleted file mode 100644 index 43c757bd014a..000000000000 --- a/ggml/src/ggml-cpu/arch/riscv/cpu-feats.cpp +++ /dev/null @@ -1,38 +0,0 @@ -#include "ggml-backend-impl.h" - -#if defined(__riscv) && __riscv_xlen == 64 -#include -#include -#include - -struct riscv64_features { - bool has_rvv = false; - - riscv64_features() { - struct riscv_hwprobe probe; - probe.key = RISCV_HWPROBE_KEY_IMA_EXT_0; - probe.value = 0; - - int ret = syscall(__NR_riscv_hwprobe, &probe, 1, 0, NULL, 0); - - if (0 == ret) { - has_rvv = !!(probe.value & RISCV_HWPROBE_IMA_V); - } - } -}; - -static int ggml_backend_cpu_riscv64_score() { - int score = 1; - riscv64_features rf; - -#ifdef GGML_USE_RVV - if (!rf.has_rvv) { return 0; } - score += 1 << 1; -#endif - - return score; -} - -GGML_BACKEND_DL_SCORE_IMPL(ggml_backend_cpu_riscv64_score) - -#endif // __riscv && __riscv_xlen == 64 diff --git a/ggml/src/ggml-cpu/arch/riscv/quants.c b/ggml/src/ggml-cpu/arch/riscv/quants.c.inc similarity index 99% rename from ggml/src/ggml-cpu/arch/riscv/quants.c rename to ggml/src/ggml-cpu/arch/riscv/quants.c.inc index d3278d6489fd..0d805824bf4a 100644 --- a/ggml/src/ggml-cpu/arch/riscv/quants.c +++ b/ggml/src/ggml-cpu/arch/riscv/quants.c.inc @@ -1,12 +1,12 @@ #define GGML_COMMON_IMPL_C -#include "ggml-common.h" -#include "ggml-quants.h" -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "simd-mappings.h" +#include "ggml-common-defs.inc" +#include "ggml-quants-defs.inc" +#include "ggml-impl-defs.inc" +#include "ggml-cpu.h.inc" +#include "simd-mappings-defs.inc" #include "../../quants.h" -#include "../../ggml-cpu-impl.h" +#include "../../ggml-cpu-impl-defs.inc" #include #include diff --git a/ggml/src/ggml-cpu/arch/riscv/repack.cpp b/ggml/src/ggml-cpu/arch/riscv/repack.cpp.inc similarity index 96% rename from ggml/src/ggml-cpu/arch/riscv/repack.cpp rename to ggml/src/ggml-cpu/arch/riscv/repack.cpp.inc index c37488cae545..50ddf89cc021 100644 --- a/ggml/src/ggml-cpu/arch/riscv/repack.cpp +++ b/ggml/src/ggml-cpu/arch/riscv/repack.cpp.inc @@ -1,13 +1,140 @@ #define GGML_COMMON_IMPL_CPP #define GGML_COMMON_DECL_CPP -#include "ggml-common.h" -#include "ggml-backend-impl.h" - -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "ggml-cpu-impl.h" -#include "simd-mappings.h" -#include "traits.h" +#include "ggml-common-defs.inc" +#include "ggml-backend-impl-defs.inc" + +#include "ggml-impl-defs.inc" +#include "ggml-cpu.h.inc" +#include "ggml-cpu-impl-defs.inc" +#include "simd-mappings-defs.inc" +#include "ggml.h.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml-cpu-impl-defs.inc" +#include "ggml-impl-defs.inc" +#include "simd-mappings-defs.inc" + +#ifdef __cplusplus +#include +extern "C" { +#endif + +// return true if op part of extra "accelerator" +bool ggml_cpu_extra_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op); +bool ggml_cpu_extra_work_size(int n_threads, const struct ggml_tensor * op, size_t * size); + +#ifdef __cplusplus +} + +namespace ggml::cpu { +// register in tensor->extra +class tensor_traits { + public: + virtual ~tensor_traits(); + virtual bool work_size(int n_threads, const struct ggml_tensor * op, size_t & size) = 0; + virtual bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) = 0; +}; + +class extra_buffer_type { + public: + virtual ~extra_buffer_type(); + virtual bool supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) = 0; + virtual tensor_traits * get_tensor_traits(const struct ggml_tensor * op) = 0; +}; +} // namespace ggml::cpu + +// implemented in ggml-cpu.cpp. +std::vector & ggml_backend_cpu_get_extra_buffer_types(); + +#endif + +#define GGML_FA_TILE_Q 64 +#define GGML_FA_TILE_KV 64 + +#ifdef __cplusplus + +#include + +// convenience functions/macros for use in template calls +// note: these won't be required after the 'traits' lookup table is used. +static inline ggml_fp16_t f32_to_f16(float x) { + return GGML_CPU_FP32_TO_FP16(x); +} + +static inline float f16_to_f32(ggml_fp16_t x) { + return GGML_CPU_FP16_TO_FP32(x); +} + +static inline ggml_bf16_t f32_to_bf16(float x) { + return GGML_FP32_TO_BF16(x); +} + +static inline float bf16_to_f32(ggml_bf16_t x) { + return GGML_BF16_TO_FP32(x); +} + +static inline float i32_to_f32(int32_t x) { + return x; +} + +static inline int32_t f32_to_i32(float x) { + return x; +} + +static inline float f32_to_f32(float x) { + return x; +} + +// TODO - merge this into the traits table, after using row-based conversions +template +struct type_conversion_table; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_fp16_t) = f16_to_f32; + static constexpr ggml_fp16_t (*from_f32)(float) = f32_to_f16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(float) = f32_to_f32; + static constexpr float (*from_f32)(float) = f32_to_f32; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_bf16_t) = bf16_to_f32; + static constexpr ggml_bf16_t (*from_f32)(float) = f32_to_bf16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(int32_t) = i32_to_f32; + static constexpr int32_t (*from_f32)(float) = f32_to_i32; +}; + +static std::pair get_thread_range(const struct ggml_compute_params * params, const struct ggml_tensor * src0) { + const int64_t ith = params->ith; + const int64_t nth = params->nth; + + const int64_t nr = ggml_nrows(src0); + + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; + + // row range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); + + return {ir0, ir1}; +} + +struct ggml_fa_tile_config { + static constexpr size_t Q = GGML_FA_TILE_Q; + static constexpr size_t KV = GGML_FA_TILE_KV; +}; + +#endif + #include #include diff --git a/ggml/src/ggml-cpu/arch/s390/cpu-feats-rust/Cargo.toml b/ggml/src/ggml-cpu/arch/s390/cpu-feats-rust/Cargo.toml new file mode 100644 index 000000000000..c278e6a9da96 --- /dev/null +++ b/ggml/src/ggml-cpu/arch/s390/cpu-feats-rust/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "ggml-cpu-s390-feats-rust" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["staticlib", "rlib"] diff --git a/ggml/src/ggml-cpu/arch/s390/cpu-feats-rust/src/lib.rs b/ggml/src/ggml-cpu/arch/s390/cpu-feats-rust/src/lib.rs new file mode 100644 index 000000000000..11f989d67fa3 --- /dev/null +++ b/ggml/src/ggml-cpu/arch/s390/cpu-feats-rust/src/lib.rs @@ -0,0 +1,84 @@ +#![allow(unexpected_cfgs)] + +#[cfg(target_arch = "s390x")] +use std::ffi::c_ulong; + +#[cfg(target_arch = "s390x")] +const AT_HWCAP: c_ulong = 16; +#[cfg(target_arch = "s390x")] +const HWCAP_VXRS_EXT2: c_ulong = 1 << 15; +#[cfg(target_arch = "s390x")] +const HWCAP_NNPA: c_ulong = 1 << 20; + +#[cfg(target_arch = "s390x")] +extern "C" { + fn getauxval(key: c_ulong) -> c_ulong; +} + +#[derive(Clone, Copy, Default)] +struct S390Features { + has_vxe2: bool, + has_nnpa: bool, +} + +#[cfg(target_arch = "s390x")] +fn detect_features() -> S390Features { + let hwcap = unsafe { getauxval(AT_HWCAP) }; + S390Features { + has_vxe2: (hwcap & HWCAP_VXRS_EXT2) != 0, + has_nnpa: (hwcap & HWCAP_NNPA) != 0, + } +} + +#[cfg(not(target_arch = "s390x"))] +fn detect_features() -> S390Features { + S390Features::default() +} + +fn score_for_features(features: S390Features) -> i32 { + let mut score = 1_i32; + + if cfg!(ggml_use_vxe2) { + if !features.has_vxe2 { + return 0; + } + score += 1 << 1; + } + + if cfg!(ggml_use_nnpa) { + if !features.has_nnpa { + return 0; + } + score += 1 << 2; + } + + score +} + +#[no_mangle] +pub extern "C" fn ggml_backend_score() -> i32 { + score_for_features(detect_features()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn baseline_score_without_cfg_requirements() { + assert_eq!(score_for_features(S390Features::default()), 1); + assert_eq!( + score_for_features(S390Features { + has_vxe2: true, + has_nnpa: true, + }), + 1 + ); + } + + #[cfg(target_arch = "s390x")] + #[test] + fn getauxval_is_callable() { + let _ = detect_features(); + } +} diff --git a/ggml/src/ggml-cpu/arch/s390/cpu-feats.cpp b/ggml/src/ggml-cpu/arch/s390/cpu-feats.cpp deleted file mode 100644 index 5f4405a7f308..000000000000 --- a/ggml/src/ggml-cpu/arch/s390/cpu-feats.cpp +++ /dev/null @@ -1,50 +0,0 @@ -#include "ggml-backend-impl.h" - -#if defined(__s390x__) -#include - -// find hwcap bits in asm/elf.h -#ifndef HWCAP_VXRS_EXT2 -#define HWCAP_VXRS_EXT2 (1 << 15) -#endif - -#ifndef HWCAP_NNPA -#define HWCAP_NNPA (1 << 20) -#endif - -struct s390x_features { - bool has_vxe2 = false; - bool has_nnpa = false; - - s390x_features() { - uint32_t hwcap = getauxval(AT_HWCAP); - // NOTE: use hwcap2 with DFLT for z17 and later - // uint32_t hwcap2 = getauxval(AT_HWCAP2); - - has_vxe2 = !!(hwcap & HWCAP_VXRS_EXT2); - has_nnpa = !!(hwcap & HWCAP_NNPA); - } -}; - -static int ggml_backend_cpu_s390x_score() { - int score = 1; - s390x_features sf; - -// IBM z15 / LinuxONE 3 -#ifdef GGML_USE_VXE2 - if (!sf.has_vxe2) { return 0; } - score += 1 << 1; -#endif - -// IBM z16 / LinuxONE 4 and z17 / LinuxONE 5 -#ifdef GGML_USE_NNPA - if (!sf.has_nnpa) { return 0; } - score += 1 << 2; -#endif - - return score; -} - -GGML_BACKEND_DL_SCORE_IMPL(ggml_backend_cpu_s390x_score) - -#endif // __s390x__ diff --git a/ggml/src/ggml-cpu/arch/s390/quants.c b/ggml/src/ggml-cpu/arch/s390/quants.c.inc similarity index 99% rename from ggml/src/ggml-cpu/arch/s390/quants.c rename to ggml/src/ggml-cpu/arch/s390/quants.c.inc index 500857579a70..7ccf2aa3c185 100644 --- a/ggml/src/ggml-cpu/arch/s390/quants.c +++ b/ggml/src/ggml-cpu/arch/s390/quants.c.inc @@ -1,12 +1,12 @@ #define GGML_COMMON_IMPL_C -#include "ggml-common.h" -#include "ggml-quants.h" -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "simd-mappings.h" +#include "ggml-common-defs.inc" +#include "ggml-quants-defs.inc" +#include "ggml-impl-defs.inc" +#include "ggml-cpu.h.inc" +#include "simd-mappings-defs.inc" #include "../../quants.h" -#include "../../ggml-cpu-impl.h" +#include "../../ggml-cpu-impl-defs.inc" #include #include diff --git a/ggml/src/ggml-cpu/arch/wasm/quants.c b/ggml/src/ggml-cpu/arch/wasm/quants.c.inc similarity index 99% rename from ggml/src/ggml-cpu/arch/wasm/quants.c rename to ggml/src/ggml-cpu/arch/wasm/quants.c.inc index 648c6fcaba76..b8204d3796ea 100644 --- a/ggml/src/ggml-cpu/arch/wasm/quants.c +++ b/ggml/src/ggml-cpu/arch/wasm/quants.c.inc @@ -1,12 +1,12 @@ #define GGML_COMMON_IMPL_C -#include "ggml-common.h" -#include "ggml-quants.h" -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "simd-mappings.h" +#include "ggml-common-defs.inc" +#include "ggml-quants-defs.inc" +#include "ggml-impl-defs.inc" +#include "ggml-cpu.h.inc" +#include "simd-mappings-defs.inc" #include "../../quants.h" -#include "../../ggml-cpu-impl.h" +#include "../../ggml-cpu-impl-defs.inc" #include #include diff --git a/ggml/src/ggml-cpu/arch/x86/cpu-feats-rust/Cargo.toml b/ggml/src/ggml-cpu/arch/x86/cpu-feats-rust/Cargo.toml new file mode 100644 index 000000000000..7af767b030cd --- /dev/null +++ b/ggml/src/ggml-cpu/arch/x86/cpu-feats-rust/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "ggml-cpu-x86-feats-rust" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["staticlib", "rlib"] diff --git a/ggml/src/ggml-cpu/arch/x86/cpu-feats-rust/src/lib.rs b/ggml/src/ggml-cpu/arch/x86/cpu-feats-rust/src/lib.rs new file mode 100644 index 000000000000..876e68b0f49c --- /dev/null +++ b/ggml/src/ggml-cpu/arch/x86/cpu-feats-rust/src/lib.rs @@ -0,0 +1,232 @@ +#![allow(unexpected_cfgs)] + +#[cfg(target_arch = "x86")] +use std::arch::x86::__cpuid_count; +#[cfg(target_arch = "x86_64")] +use std::arch::x86_64::__cpuid_count; + +#[derive(Default)] +struct CpuId { + f_1_ecx: u32, + f_7_ebx: u32, + f_7_ecx: u32, + f_7_edx: u32, + f_7_1_eax: u32, +} + +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +impl CpuId { + fn new() -> Self { + let leaf0 = __cpuid_count(0, 0); + let n_ids = leaf0.eax; + let mut this = Self::default(); + + if n_ids >= 1 { + let leaf1 = __cpuid_count(1, 0); + this.f_1_ecx = leaf1.ecx; + } + + if n_ids >= 7 { + let leaf7 = __cpuid_count(7, 0); + this.f_7_ebx = leaf7.ebx; + this.f_7_ecx = leaf7.ecx; + this.f_7_edx = leaf7.edx; + this.f_7_1_eax = __cpuid_count(7, 1).eax; + } + + this + } + + fn bit(reg: u32, bit: u32) -> bool { + (reg & (1_u32 << bit)) != 0 + } + + fn fma(&self) -> bool { + Self::bit(self.f_1_ecx, 12) + } + + fn f16c(&self) -> bool { + Self::bit(self.f_1_ecx, 29) + } + + fn sse42(&self) -> bool { + Self::bit(self.f_1_ecx, 20) + } + + fn bmi2(&self) -> bool { + Self::bit(self.f_7_ebx, 8) + } + + fn avx(&self) -> bool { + Self::bit(self.f_1_ecx, 28) + } + + fn avx2(&self) -> bool { + Self::bit(self.f_7_ebx, 5) + } + + fn avx_vnni(&self) -> bool { + Self::bit(self.f_7_1_eax, 4) + } + + fn avx512f(&self) -> bool { + Self::bit(self.f_7_ebx, 16) + } + + fn avx512dq(&self) -> bool { + Self::bit(self.f_7_ebx, 17) + } + + fn avx512cd(&self) -> bool { + Self::bit(self.f_7_ebx, 28) + } + + fn avx512bw(&self) -> bool { + Self::bit(self.f_7_ebx, 30) + } + + fn avx512vl(&self) -> bool { + Self::bit(self.f_7_ebx, 31) + } + + fn avx512_vbmi(&self) -> bool { + Self::bit(self.f_7_ecx, 1) + } + + fn avx512_vnni(&self) -> bool { + Self::bit(self.f_7_ecx, 11) + } + + fn avx512_bf16(&self) -> bool { + Self::bit(self.f_7_1_eax, 5) + } + + fn amx_int8(&self) -> bool { + Self::bit(self.f_7_edx, 25) + } +} + +#[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))] +impl CpuId { + fn new() -> Self { + Self::default() + } + + fn fma(&self) -> bool { false } + fn f16c(&self) -> bool { false } + fn sse42(&self) -> bool { false } + fn bmi2(&self) -> bool { false } + fn avx(&self) -> bool { false } + fn avx2(&self) -> bool { false } + fn avx_vnni(&self) -> bool { false } + fn avx512f(&self) -> bool { false } + fn avx512dq(&self) -> bool { false } + fn avx512cd(&self) -> bool { false } + fn avx512bw(&self) -> bool { false } + fn avx512vl(&self) -> bool { false } + fn avx512_vbmi(&self) -> bool { false } + fn avx512_vnni(&self) -> bool { false } + fn avx512_bf16(&self) -> bool { false } + fn amx_int8(&self) -> bool { false } +} + +fn score_for_cpu(cpu: &CpuId) -> i32 { + let mut score = 1_i32; + + if cfg!(ggml_fma) { + if !cpu.fma() { + return 0; + } + score += 1; + } + if cfg!(ggml_f16c) { + if !cpu.f16c() { + return 0; + } + score += 1 << 1; + } + if cfg!(ggml_sse42) { + if !cpu.sse42() { + return 0; + } + score += 1 << 2; + } + if cfg!(ggml_bmi2) { + if !cpu.bmi2() { + return 0; + } + score += 1 << 3; + } + if cfg!(ggml_avx) { + if !cpu.avx() { + return 0; + } + score += 1 << 4; + } + if cfg!(ggml_avx2) { + if !cpu.avx2() { + return 0; + } + score += 1 << 5; + } + if cfg!(ggml_avx_vnni) { + if !cpu.avx_vnni() { + return 0; + } + score += 1 << 6; + } + if cfg!(ggml_avx512) { + if !cpu.avx512f() || !cpu.avx512cd() || !cpu.avx512vl() || !cpu.avx512dq() || !cpu.avx512bw() { + return 0; + } + score += 1 << 7; + } + if cfg!(ggml_avx512_vbmi) { + if !cpu.avx512_vbmi() { + return 0; + } + score += 1 << 8; + } + if cfg!(ggml_avx512_bf16) { + if !cpu.avx512_bf16() { + return 0; + } + score += 1 << 9; + } + if cfg!(ggml_avx512_vnni) { + if !cpu.avx512_vnni() { + return 0; + } + score += 1 << 10; + } + if cfg!(ggml_amx_int8) { + if !cpu.amx_int8() { + return 0; + } + score += 1 << 11; + } + + score +} + +#[no_mangle] +pub extern "C" fn ggml_backend_score() -> i32 { + score_for_cpu(&CpuId::new()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn score_is_at_least_baseline_without_cfg_requirements() { + assert!(score_for_cpu(&CpuId::default()) >= 1); + } + + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + #[test] + fn reads_vendor_flags_on_x86() { + let cpu = CpuId::new(); + assert!(cpu.f_1_ecx != 0 || cpu.f_7_ebx != 0 || cpu.f_7_ecx != 0 || cpu.f_7_edx != 0); + } +} diff --git a/ggml/src/ggml-cpu/arch/x86/cpu-feats.cpp b/ggml/src/ggml-cpu/arch/x86/cpu-feats.cpp deleted file mode 100644 index d775a0363858..000000000000 --- a/ggml/src/ggml-cpu/arch/x86/cpu-feats.cpp +++ /dev/null @@ -1,327 +0,0 @@ -#include "ggml-backend-impl.h" - -#if defined(__x86_64__) || (defined(_MSC_VER) && defined(_M_AMD64)) - -#ifdef _MSC_VER -#include -#endif - -#include -#include -#include -#include -#include - -// ref: https://cdrdv2-public.intel.com/782156/325383-sdm-vol-2abcd.pdf -struct cpuid_x86 { - bool SSE3(void) { return f_1_ecx[0]; } - bool PCLMULQDQ(void) { return f_1_ecx[1]; } - bool MONITOR(void) { return f_1_ecx[3]; } - bool SSSE3(void) { return f_1_ecx[9]; } - bool FMA(void) { return f_1_ecx[12]; } - bool CMPXCHG16B(void) { return f_1_ecx[13]; } - bool SSE41(void) { return f_1_ecx[19]; } - bool SSE42(void) { return f_1_ecx[20]; } - bool MOVBE(void) { return f_1_ecx[22]; } - bool POPCNT(void) { return f_1_ecx[23]; } - bool AES(void) { return f_1_ecx[25]; } - bool XSAVE(void) { return f_1_ecx[26]; } - bool OSXSAVE(void) { return f_1_ecx[27]; } - bool AVX(void) { return f_1_ecx[28]; } - bool F16C(void) { return f_1_ecx[29]; } - bool RDRAND(void) { return f_1_ecx[30]; } - - bool MSR(void) { return f_1_edx[5]; } - bool CX8(void) { return f_1_edx[8]; } - bool SEP(void) { return f_1_edx[11]; } - bool CMOV(void) { return f_1_edx[15]; } - bool CLFSH(void) { return f_1_edx[19]; } - bool MMX(void) { return f_1_edx[23]; } - bool FXSR(void) { return f_1_edx[24]; } - bool SSE(void) { return f_1_edx[25]; } - bool SSE2(void) { return f_1_edx[26]; } - - bool FSGSBASE(void) { return f_7_ebx[0]; } - bool BMI1(void) { return f_7_ebx[3]; } - bool HLE(void) { return is_intel && f_7_ebx[4]; } - bool AVX2(void) { return f_7_ebx[5]; } - bool BMI2(void) { return f_7_ebx[8]; } - bool ERMS(void) { return f_7_ebx[9]; } - bool INVPCID(void) { return f_7_ebx[10]; } - bool RTM(void) { return is_intel && f_7_ebx[11]; } - bool AVX512F(void) { return f_7_ebx[16]; } - bool AVX512DQ(void) { return f_7_ebx[17]; } - bool RDSEED(void) { return f_7_ebx[18]; } - bool ADX(void) { return f_7_ebx[19]; } - bool AVX512PF(void) { return f_7_ebx[26]; } - bool AVX512ER(void) { return f_7_ebx[27]; } - bool AVX512CD(void) { return f_7_ebx[28]; } - bool AVX512BW(void) { return f_7_ebx[30]; } - bool AVX512VL(void) { return f_7_ebx[31]; } - - bool SHA(void) { return f_7_ebx[29]; } - - bool PREFETCHWT1(void) { return f_7_ecx[0]; } - - bool LAHF(void) { return f_81_ecx[0]; } - bool LZCNT(void) { return is_intel && f_81_ecx[5]; } - bool ABM(void) { return is_amd && f_81_ecx[5]; } - bool SSE4a(void) { return is_amd && f_81_ecx[6]; } - bool XOP(void) { return is_amd && f_81_ecx[11]; } - bool TBM(void) { return is_amd && f_81_ecx[21]; } - - bool SYSCALL(void) { return is_intel && f_81_edx[11]; } - bool MMXEXT(void) { return is_amd && f_81_edx[22]; } - bool RDTSCP(void) { return is_intel && f_81_edx[27]; } - bool _3DNOWEXT(void) { return is_amd && f_81_edx[30]; } - bool _3DNOW(void) { return is_amd && f_81_edx[31]; } - - bool AVX512_VBMI(void) { return f_7_ecx[1]; } - bool AVX512_VNNI(void) { return f_7_ecx[11]; } - bool AVX512_FP16(void) { return f_7_edx[23]; } - bool AVX512_BF16(void) { return f_7_1_eax[5]; } - bool AVX_VNNI(void) { return f_7_1_eax[4]; } - - bool AMX_TILE(void) { return f_7_edx[24]; } - bool AMX_INT8(void) { return f_7_edx[25]; } - bool AMX_FP16(void) { return f_7_1_eax[21]; } - bool AMX_BF16(void) { return f_7_edx[22]; } - -#ifdef _MSC_VER - static void cpuid(int cpu_info[4], int eax) { - __cpuid(cpu_info, eax); - } - static void cpuidex(int cpu_info[4], int eax, int ecx) { - __cpuidex(cpu_info, eax, ecx); - } -#else - static void cpuid(int cpu_info[4], int eax) { - __asm__ __volatile__( - "cpuid" - : "=a"(cpu_info[0]), "=b"(cpu_info[1]), "=c"(cpu_info[2]), "=d"(cpu_info[3]) - : "a"(eax), "c"(0)); - } - static void cpuidex(int cpu_info[4], int eax, int ecx) { - __asm__ __volatile__( - "cpuid" - : "=a"(cpu_info[0]), "=b"(cpu_info[1]), "=c"(cpu_info[2]), "=d"(cpu_info[3]) - : "a"(eax), "c"(ecx)); - } -#endif - - cpuid_x86() { - std::array cpui; - std::vector> data; - - // calling __cpuid with 0x0 as the function_id argument - // gets the number of the highest valid function ID. - cpuid(cpui.data(), 0); - int n_ids = cpui[0]; - - for (int i = 0; i <= n_ids; ++i) { - cpuidex(cpui.data(), i, 0); - data.push_back(cpui); - } - - // capture vendor string - char vendor[0x20] = {}; - *reinterpret_cast(vendor) = data[0][1]; - *reinterpret_cast(vendor + 4) = data[0][3]; - *reinterpret_cast(vendor + 8) = data[0][2]; - this->vendor = vendor; - if (this->vendor == "GenuineIntel") { - is_intel = true; - } else if (this->vendor == "AuthenticAMD") { - is_amd = true; - } - - // load bitset with flags for function 0x00000001 - if (n_ids >= 1) { - f_1_ecx = data[1][2]; - f_1_edx = data[1][3]; - } - - // load bitset with flags for function 0x00000007 - if (n_ids >= 7) { - f_7_ebx = data[7][1]; - f_7_ecx = data[7][2]; - f_7_edx = data[7][3]; - cpuidex(cpui.data(), 7, 1); - f_7_1_eax = cpui[0]; - } - - // calling __cpuid with 0x80000000 as the function_id argument - // gets the number of the highest valid extended ID. - cpuid(cpui.data(), 0x80000000); - unsigned int n_ex_ids = cpui[0]; - - std::vector> ext_data; - for (unsigned int i = 0x80000000; i <= n_ex_ids; ++i) { - cpuidex(cpui.data(), i, 0); - ext_data.push_back(cpui); - } - - // load bitset with flags for function 0x80000001 - if (n_ex_ids >= 0x80000001) { - f_81_ecx = ext_data[1][2]; - f_81_edx = ext_data[1][3]; - } - - // interpret CPU brand string if reported - char brand[0x40] = {}; - if (n_ex_ids >= 0x80000004) { - std::memcpy(brand, ext_data[2].data(), sizeof(cpui)); - std::memcpy(brand + 16, ext_data[3].data(), sizeof(cpui)); - std::memcpy(brand + 32, ext_data[4].data(), sizeof(cpui)); - this->brand = brand; - } - } - - bool is_intel = false; - bool is_amd = false; - std::string vendor; - std::string brand; - std::bitset<32> f_1_ecx; - std::bitset<32> f_1_edx; - std::bitset<32> f_7_ebx; - std::bitset<32> f_7_ecx; - std::bitset<32> f_7_edx; - std::bitset<32> f_7_1_eax; - std::bitset<32> f_81_ecx; - std::bitset<32> f_81_edx; -}; - -#if 0 -void test_x86_is() { - cpuid_x86 is; - printf("CPU Vendor: %s\n", is.vendor.c_str()); - printf("Brand: %s\n", is.brand.c_str()); - printf("is_intel: %d\n", is.is_intel); - printf("is_amd: %d\n", is.is_amd); - printf("sse3: %d\n", is.SSE3()); - printf("pclmulqdq: %d\n", is.PCLMULQDQ()); - printf("ssse3: %d\n", is.SSSE3()); - printf("fma: %d\n", is.FMA()); - printf("cmpxchg16b: %d\n", is.CMPXCHG16B()); - printf("sse41: %d\n", is.SSE41()); - printf("sse42: %d\n", is.SSE42()); - printf("movbe: %d\n", is.MOVBE()); - printf("popcnt: %d\n", is.POPCNT()); - printf("aes: %d\n", is.AES()); - printf("xsave: %d\n", is.XSAVE()); - printf("osxsave: %d\n", is.OSXSAVE()); - printf("avx: %d\n", is.AVX()); - printf("f16c: %d\n", is.F16C()); - printf("rdrand: %d\n", is.RDRAND()); - printf("msr: %d\n", is.MSR()); - printf("cx8: %d\n", is.CX8()); - printf("sep: %d\n", is.SEP()); - printf("cmov: %d\n", is.CMOV()); - printf("clflush: %d\n", is.CLFSH()); - printf("mmx: %d\n", is.MMX()); - printf("fxsr: %d\n", is.FXSR()); - printf("sse: %d\n", is.SSE()); - printf("sse2: %d\n", is.SSE2()); - printf("fsgsbase: %d\n", is.FSGSBASE()); - printf("bmi1: %d\n", is.BMI1()); - printf("hle: %d\n", is.HLE()); - printf("avx2: %d\n", is.AVX2()); - printf("bmi2: %d\n", is.BMI2()); - printf("erms: %d\n", is.ERMS()); - printf("invpcid: %d\n", is.INVPCID()); - printf("rtm: %d\n", is.RTM()); - printf("avx512f: %d\n", is.AVX512F()); - printf("rdseed: %d\n", is.RDSEED()); - printf("adx: %d\n", is.ADX()); - printf("avx512pf: %d\n", is.AVX512PF()); - printf("avx512er: %d\n", is.AVX512ER()); - printf("avx512cd: %d\n", is.AVX512CD()); - printf("sha: %d\n", is.SHA()); - printf("prefetchwt1: %d\n", is.PREFETCHWT1()); - printf("lahf: %d\n", is.LAHF()); - printf("lzcnt: %d\n", is.LZCNT()); - printf("abm: %d\n", is.ABM()); - printf("sse4a: %d\n", is.SSE4a()); - printf("xop: %d\n", is.XOP()); - printf("tbm: %d\n", is.TBM()); - printf("syscall: %d\n", is.SYSCALL()); - printf("mmxext: %d\n", is.MMXEXT()); - printf("rdtscp: %d\n", is.RDTSCP()); - printf("3dnowext: %d\n", is._3DNOWEXT()); - printf("3dnow: %d\n", is._3DNOW()); - printf("avx512_vbmi: %d\n", is.AVX512_VBMI()); - printf("avx512_vnni: %d\n", is.AVX512_VNNI()); - printf("avx512_fp16: %d\n", is.AVX512_FP16()); - printf("avx512_bf16: %d\n", is.AVX512_BF16()); - printf("amx_tile: %d\n", is.AMX_TILE()); - printf("amx_int8: %d\n", is.AMX_INT8()); - printf("amx_fp16: %d\n", is.AMX_FP16()); - printf("amx_bf16: %d\n", is.AMX_BF16()); -} -#endif - -static int ggml_backend_cpu_x86_score() { - // FIXME: this does not check for OS support - - int score = 1; - cpuid_x86 is; - -#ifdef GGML_FMA - if (!is.FMA()) { return 0; } - score += 1; -#endif -#ifdef GGML_F16C - if (!is.F16C()) { return 0; } - score += 1<<1; -#endif -#ifdef GGML_SSE42 - if (!is.SSE42()) { return 0; } - score += 1<<2; -#endif -#ifdef GGML_BMI2 - if (!is.BMI2()) { return 0; } - score += 1<<3; -#endif -#ifdef GGML_AVX - if (!is.AVX()) { return 0; } - score += 1<<4; -#endif -#ifdef GGML_AVX2 - if (!is.AVX2()) { return 0; } - score += 1<<5; -#endif -#ifdef GGML_AVX_VNNI - if (!is.AVX_VNNI()) { return 0; } - score += 1<<6; -#endif -#ifdef GGML_AVX512 - if (!is.AVX512F()) { return 0; } - if (!is.AVX512CD()) { return 0; } - if (!is.AVX512VL()) { return 0; } - if (!is.AVX512DQ()) { return 0; } - if (!is.AVX512BW()) { return 0; } - score += 1<<7; -#endif -#ifdef GGML_AVX512_VBMI - if (!is.AVX512_VBMI()) { return 0; } - score += 1<<8; -#endif -#ifdef GGML_AVX512_BF16 - if (!is.AVX512_BF16()) { return 0; } - score += 1<<9; -#endif -#ifdef GGML_AVX512_VNNI - if (!is.AVX512_VNNI()) { return 0; } - score += 1<<10; -#endif -#ifdef GGML_AMX_INT8 - if (!is.AMX_INT8()) { return 0; } - score += 1<<11; -#endif - - return score; -} - -GGML_BACKEND_DL_SCORE_IMPL(ggml_backend_cpu_x86_score) - -#endif // defined(__x86_64__) || (defined(_MSC_VER) && defined(_M_AMD64)) diff --git a/ggml/src/ggml-cpu/binary-ops.cpp b/ggml/src/ggml-cpu/binary-ops.cpp deleted file mode 100644 index 75e38290015a..000000000000 --- a/ggml/src/ggml-cpu/binary-ops.cpp +++ /dev/null @@ -1,154 +0,0 @@ -#include "binary-ops.h" - -#if defined(GGML_USE_ACCELERATE) -#include - -using vDSP_fn_t = void (*)(const float *, vDSP_Stride, const float *, vDSP_Stride, float *, vDSP_Stride, vDSP_Length); -#endif - -static inline float op_add(float a, float b) { - return a + b; -} - -static inline float op_sub(float a, float b) { - return a - b; -} - -static inline float op_mul(float a, float b) { - return a * b; -} - -static inline float op_div(float a, float b) { - return a / b; -} - -template -static inline void vec_binary_op_contiguous(const int64_t n, dst_t * z, const src0_t * x, const src1_t * y) { - constexpr auto src0_to_f32 = type_conversion_table::to_f32; - constexpr auto src1_to_f32 = type_conversion_table::to_f32; - constexpr auto f32_to_dst = type_conversion_table::from_f32; - - for (int i = 0; i < n; i++) { - z[i] = f32_to_dst(op(src0_to_f32(x[i]), src1_to_f32(y[i]))); - } -} - -template -static inline void vec_binary_op_non_contiguous(const int64_t n, const int64_t ne10, const int64_t nb10, dst_t * z, const src0_t * x, const src1_t * y) { - constexpr auto src0_to_f32 = type_conversion_table::to_f32; - constexpr auto src1_to_f32 = type_conversion_table::to_f32; - constexpr auto f32_to_dst = type_conversion_table::from_f32; - - for (int i = 0; i < n; i++) { - int i10 = i % ne10; - const src1_t * y_ptr = (const src1_t *)((const char *)y + i10*nb10); - z[i] = f32_to_dst(op(src0_to_f32(x[i]), src1_to_f32(*y_ptr))); - } -} - -template -static void apply_binary_op(const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - - GGML_ASSERT(ggml_can_repeat(src1, src0) && ggml_are_same_shape(src0, dst)); - - GGML_TENSOR_BINARY_OP_LOCALS - - GGML_ASSERT( nb0 == sizeof(dst_t)); - GGML_ASSERT(nb00 == sizeof(src0_t)); - - const auto [ir0, ir1] = get_thread_range(params, src0); - const bool is_src1_contiguous_rows = ggml_is_contiguous_rows(src1); - -#ifdef GGML_USE_ACCELERATE - vDSP_fn_t vDSP_op = nullptr; - // TODO - avoid the f32-only check using type 'trait' lookup tables and row-based src-to-float conversion functions - if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { - if (op == op_add) { - vDSP_op = vDSP_vadd; - } else if (op == op_sub) { - vDSP_op = vDSP_vsub; - } else if (op == op_mul) { - vDSP_op = vDSP_vmul; - } else if (op == op_div) { - vDSP_op = vDSP_vdiv; - } - } -#endif - - for (int64_t ir = ir0; ir < ir1; ++ir) { - const int64_t i03 = ir/(ne02*ne01); - const int64_t i02 = (ir - i03*ne02*ne01)/ne01; - const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); - - const int64_t i13 = i03 % ne13; - const int64_t i12 = i02 % ne12; - const int64_t i11 = i01 % ne11; - - dst_t * dst_ptr = (dst_t *) ((char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1 ); - const src0_t * src0_ptr = (const src0_t *) ((const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01); - const src1_t * src1_ptr = (const src1_t *) ((const char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11); - - if (is_src1_contiguous_rows) { - // src1 is broadcastable across src0 and dst in i1, i2, i3 - const int64_t nr0 = ne00 / ne10; - - for (int64_t r = 0; r < nr0; ++r) { -#ifdef GGML_USE_ACCELERATE - if constexpr (std::is_same_v && std::is_same_v && std::is_same_v) { - if (vDSP_op != nullptr) { - vDSP_op(src1_ptr, 1, src0_ptr + r*ne10, 1, dst_ptr + r*ne10, 1, ne10); - continue; - } - } -#endif - vec_binary_op_contiguous(ne10, dst_ptr + r*ne10, src0_ptr + r*ne10, src1_ptr); - } - } else { - vec_binary_op_non_contiguous(ne0, ne10, nb10, dst_ptr, src0_ptr, src1_ptr); - } - } -} - -// TODO: Use the 'traits' lookup table (for type conversion fns), instead of a mass of 'if' conditions with long templates -template -static void binary_op(const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - - /* */ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { // all f32 - apply_binary_op(params, dst); - } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { // all f16 - apply_binary_op(params, dst); - } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) { // all bf16 - apply_binary_op(params, dst); - } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_BF16) { - apply_binary_op(params, dst); - } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { - apply_binary_op(params, dst); - } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) { - apply_binary_op(params, dst); - } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { - apply_binary_op(params, dst); - } else { - GGML_ABORT("%s: unsupported types: dst: %s, src0: %s, src1: %s\n", __func__, - ggml_type_name(dst->type), ggml_type_name(src0->type), ggml_type_name(src1->type)); - } -} - -void ggml_compute_forward_add_non_quantized(const ggml_compute_params * params, ggml_tensor * dst) { - binary_op(params, dst); -} - -void ggml_compute_forward_sub(const ggml_compute_params * params, ggml_tensor * dst) { - binary_op(params, dst); -} - -void ggml_compute_forward_mul(const ggml_compute_params * params, ggml_tensor * dst) { - binary_op(params, dst); -} - -void ggml_compute_forward_div(const ggml_compute_params * params, ggml_tensor * dst) { - binary_op(params, dst); -} diff --git a/ggml/src/ggml-cpu/binary-ops.h b/ggml/src/ggml-cpu/binary-ops.h deleted file mode 100644 index aca1d89be7e5..000000000000 --- a/ggml/src/ggml-cpu/binary-ops.h +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "common.h" - -#ifdef __cplusplus -extern "C" { -#endif - -void ggml_compute_forward_add_non_quantized(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_sub(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_mul(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_div(const struct ggml_compute_params * params, struct ggml_tensor * dst); - -#ifdef __cplusplus -} -#endif diff --git a/ggml/src/ggml-cpu/common.h b/ggml/src/ggml-cpu/common.h deleted file mode 100644 index abbadc359c5a..000000000000 --- a/ggml/src/ggml-cpu/common.h +++ /dev/null @@ -1,95 +0,0 @@ -#pragma once - -#include "ggml.h" -#include "traits.h" -#include "ggml-cpu-impl.h" -#include "ggml-impl.h" -#include "simd-mappings.h" - -#define GGML_FA_TILE_Q 64 -#define GGML_FA_TILE_KV 64 - -#ifdef __cplusplus - -#include - -// convenience functions/macros for use in template calls -// note: these won't be required after the 'traits' lookup table is used. -static inline ggml_fp16_t f32_to_f16(float x) { - return GGML_CPU_FP32_TO_FP16(x); -} - -static inline float f16_to_f32(ggml_fp16_t x) { - return GGML_CPU_FP16_TO_FP32(x); -} - -static inline ggml_bf16_t f32_to_bf16(float x) { - return GGML_FP32_TO_BF16(x); -} - -static inline float bf16_to_f32(ggml_bf16_t x) { - return GGML_BF16_TO_FP32(x); -} - -static inline float i32_to_f32(int32_t x) { - return x; -} - -static inline int32_t f32_to_i32(float x) { - return x; -} - -static inline float f32_to_f32(float x) { - return x; -} - -// TODO - merge this into the traits table, after using row-based conversions -template -struct type_conversion_table; - -template <> -struct type_conversion_table { - static constexpr float (*to_f32)(ggml_fp16_t) = f16_to_f32; - static constexpr ggml_fp16_t (*from_f32)(float) = f32_to_f16; -}; - -template <> -struct type_conversion_table { - static constexpr float (*to_f32)(float) = f32_to_f32; - static constexpr float (*from_f32)(float) = f32_to_f32; -}; - -template <> -struct type_conversion_table { - static constexpr float (*to_f32)(ggml_bf16_t) = bf16_to_f32; - static constexpr ggml_bf16_t (*from_f32)(float) = f32_to_bf16; -}; - -template <> -struct type_conversion_table { - static constexpr float (*to_f32)(int32_t) = i32_to_f32; - static constexpr int32_t (*from_f32)(float) = f32_to_i32; -}; - -static std::pair get_thread_range(const struct ggml_compute_params * params, const struct ggml_tensor * src0) { - const int64_t ith = params->ith; - const int64_t nth = params->nth; - - const int64_t nr = ggml_nrows(src0); - - // rows per thread - const int64_t dr = (nr + nth - 1)/nth; - - // row range for this thread - const int64_t ir0 = dr*ith; - const int64_t ir1 = MIN(ir0 + dr, nr); - - return {ir0, ir1}; -} - -struct ggml_fa_tile_config { - static constexpr size_t Q = GGML_FA_TILE_Q; - static constexpr size_t KV = GGML_FA_TILE_KV; -}; - -#endif diff --git a/ggml/src/ggml-cpu/ggml-cpu-impl.h b/ggml/src/ggml-cpu/ggml-cpu-impl-defs.inc similarity index 99% rename from ggml/src/ggml-cpu/ggml-cpu-impl.h rename to ggml/src/ggml-cpu/ggml-cpu-impl-defs.inc index 5d1ca5ffcc36..78a60e3749bf 100644 --- a/ggml/src/ggml-cpu/ggml-cpu-impl.h +++ b/ggml/src/ggml-cpu/ggml-cpu-impl-defs.inc @@ -2,8 +2,8 @@ // GGML CPU internal header -#include "ggml.h" -#include "ggml-impl.h" +#include "ggml.h.inc" +#include "ggml-impl-defs.inc" #include // load `stdlib.h` before other headers to work around MinGW bug: https://sourceforge.net/p/mingw-w64/bugs/192/ //#include diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c.inc similarity index 57% rename from ggml/src/ggml-cpu/ggml-cpu.c rename to ggml/src/ggml-cpu/ggml-cpu.c.inc index 2b3eb5b5ce65..945c1c6da36b 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c.inc @@ -1,19 +1,1942 @@ #define _CRT_SECURE_NO_DEPRECATE // Disables "unsafe" warnings on Windows #define _USE_MATH_DEFINES // For M_PI on MSVC -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "traits.h" -#include "ggml-cpu-impl.h" -#include "ggml-impl.h" -#include "quants.h" -#include "ggml-threading.h" -#include "unary-ops.h" -#include "binary-ops.h" -#include "vec.h" -#include "ops.h" -#include "ggml.h" -#include "common.h" +#include "ggml-backend-impl-defs.inc" +#include "ggml-backend.h.inc" +#include "ggml.h.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml-cpu-impl-defs.inc" +#include "ggml-impl-defs.inc" +#include "simd-mappings-defs.inc" + +#ifdef __cplusplus +#include +extern "C" { +#endif + +// return true if op part of extra "accelerator" +bool ggml_cpu_extra_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op); +bool ggml_cpu_extra_work_size(int n_threads, const struct ggml_tensor * op, size_t * size); + +#ifdef __cplusplus +} + +namespace ggml::cpu { +// register in tensor->extra +class tensor_traits { + public: + virtual ~tensor_traits(); + virtual bool work_size(int n_threads, const struct ggml_tensor * op, size_t & size) = 0; + virtual bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) = 0; +}; + +class extra_buffer_type { + public: + virtual ~extra_buffer_type(); + virtual bool supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) = 0; + virtual tensor_traits * get_tensor_traits(const struct ggml_tensor * op) = 0; +}; +} // namespace ggml::cpu + +// implemented in ggml-cpu.cpp. +std::vector & ggml_backend_cpu_get_extra_buffer_types(); + +#endif + +#define GGML_FA_TILE_Q 64 +#define GGML_FA_TILE_KV 64 + +#ifdef __cplusplus + +#include + +// convenience functions/macros for use in template calls +// note: these won't be required after the 'traits' lookup table is used. +static inline ggml_fp16_t f32_to_f16(float x) { + return GGML_CPU_FP32_TO_FP16(x); +} + +static inline float f16_to_f32(ggml_fp16_t x) { + return GGML_CPU_FP16_TO_FP32(x); +} + +static inline ggml_bf16_t f32_to_bf16(float x) { + return GGML_FP32_TO_BF16(x); +} + +static inline float bf16_to_f32(ggml_bf16_t x) { + return GGML_BF16_TO_FP32(x); +} + +static inline float i32_to_f32(int32_t x) { + return x; +} + +static inline int32_t f32_to_i32(float x) { + return x; +} + +static inline float f32_to_f32(float x) { + return x; +} + +// TODO - merge this into the traits table, after using row-based conversions +template +struct type_conversion_table; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_fp16_t) = f16_to_f32; + static constexpr ggml_fp16_t (*from_f32)(float) = f32_to_f16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(float) = f32_to_f32; + static constexpr float (*from_f32)(float) = f32_to_f32; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_bf16_t) = bf16_to_f32; + static constexpr ggml_bf16_t (*from_f32)(float) = f32_to_bf16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(int32_t) = i32_to_f32; + static constexpr int32_t (*from_f32)(float) = f32_to_i32; +}; + +static std::pair get_thread_range(const struct ggml_compute_params * params, const struct ggml_tensor * src0) { + const int64_t ith = params->ith; + const int64_t nth = params->nth; + + const int64_t nr = ggml_nrows(src0); + + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; + + // row range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); + + return {ir0, ir1}; +} + +struct ggml_fa_tile_config { + static constexpr size_t Q = GGML_FA_TILE_Q; + static constexpr size_t KV = GGML_FA_TILE_KV; +}; + +#endif + +#include "ggml-cpu-impl-defs.inc" +#include "ggml-impl-defs.inc" +#define GGML_COMMON_DECL_C +#include "ggml-common-defs.inc" + +#include "ggml.h.inc" + +// GGML CPU internal header + +#ifdef __cplusplus +extern "C" { +#endif + +// Quantization +void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_q2_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q3_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q6_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_tq1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_tq2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_iq4_nl (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_iq4_xs (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +// Dot product +void ggml_vec_dot_q1_0_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 ggml_vec_dot_q4_0_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 ggml_vec_dot_q4_1_q8_1(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 ggml_vec_dot_q5_0_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 ggml_vec_dot_q5_1_q8_1(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 ggml_vec_dot_q8_0_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 ggml_vec_dot_mxfp4_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 ggml_vec_dot_nvfp4_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 ggml_vec_dot_q2_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 ggml_vec_dot_q3_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 ggml_vec_dot_q4_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 ggml_vec_dot_q5_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 ggml_vec_dot_q6_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 ggml_vec_dot_tq1_0_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 ggml_vec_dot_tq2_0_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 ggml_vec_dot_iq2_xxs_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 ggml_vec_dot_iq2_xs_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 ggml_vec_dot_iq2_s_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 ggml_vec_dot_iq3_xxs_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 ggml_vec_dot_iq1_s_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 ggml_vec_dot_iq1_m_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 ggml_vec_dot_iq4_nl_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 ggml_vec_dot_iq4_xs_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 ggml_vec_dot_iq3_s_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); + +// Generic implementation +void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void ggml_vec_dot_q1_0_q8_0_generic(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 ggml_vec_dot_q4_0_q8_0_generic(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 ggml_vec_dot_q4_1_q8_1_generic(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 ggml_vec_dot_q5_0_q8_0_generic(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 ggml_vec_dot_q5_1_q8_1_generic(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 ggml_vec_dot_q8_0_q8_0_generic(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 ggml_vec_dot_mxfp4_q8_0_generic(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 ggml_vec_dot_nvfp4_q8_0_generic(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 ggml_vec_dot_tq1_0_q8_K_generic(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 ggml_vec_dot_tq2_0_q8_K_generic(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 ggml_vec_dot_q2_K_q8_K_generic(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 ggml_vec_dot_q3_K_q8_K_generic(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 ggml_vec_dot_q4_K_q8_K_generic(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 ggml_vec_dot_q5_K_q8_K_generic(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 ggml_vec_dot_q6_K_q8_K_generic(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 ggml_vec_dot_iq2_xxs_q8_K_generic(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 ggml_vec_dot_iq2_xs_q8_K_generic(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 ggml_vec_dot_iq2_s_q8_K_generic(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 ggml_vec_dot_iq3_xxs_q8_K_generic(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 ggml_vec_dot_iq3_s_q8_K_generic(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 ggml_vec_dot_iq1_s_q8_K_generic(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 ggml_vec_dot_iq1_m_q8_K_generic(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 ggml_vec_dot_iq4_nl_q8_0_generic(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 ggml_vec_dot_iq4_xs_q8_K_generic(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); + +#ifdef __cplusplus +} +#endif + +// Vectorized functions for fundamental operations + +#include "ggml-impl-defs.inc" +#include "simd-mappings-defs.inc" +#include "ggml.h.inc" +#include "ggml-cpu.h.inc" + +#if defined(GGML_USE_ACCELERATE) +#include +#endif + +// floating point type used to accumulate sums +typedef double ggml_float; + +#define GGML_GELU_FP16 +#define GGML_GELU_QUICK_FP16 + +#define GGML_SOFT_MAX_UNROLL 4 +#define GGML_VEC_DOT_UNROLL 2 +#define GGML_VEC_MAD_UNROLL 32 + +#ifdef __cplusplus +extern "C" { +#endif + +// +// global data +// + +// precomputed gelu table for f16 (128 KB) +extern ggml_fp16_t ggml_table_gelu_f16[1 << 16]; + +// precomputed quick gelu table for f16 (128 KB) +extern ggml_fp16_t ggml_table_gelu_quick_f16[1 << 16]; + +// +// fundamental operations +// + +void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc); +void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t * GGML_RESTRICT x, size_t bx, ggml_bf16_t * GGML_RESTRICT y, size_t by, int nrc); +void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * GGML_RESTRICT x, size_t bx, ggml_fp16_t * GGML_RESTRICT y, size_t by, int nrc); + +void ggml_vec_silu_f32(const int n, float * y, const float * x); +ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean); //it will also center y ( y = y - mean ) +ggml_float ggml_vec_soft_max_f32(const int n, float * y, const float * x, float max); +ggml_float ggml_vec_log_soft_max_f32(const int n, float * y, const float * x, float max); + +inline static void ggml_vec_set_i8(const int n, int8_t * x, const int8_t v) { for (int i = 0; i < n; ++i) x[i] = v; } +inline static void ggml_vec_set_i16(const int n, int16_t * x, const int16_t v) { for (int i = 0; i < n; ++i) x[i] = v; } + +inline static void ggml_vec_set_i32(const int n, int32_t * x, const int32_t v) { for (int i = 0; i < n; ++i) x[i] = v; } +inline static void ggml_vec_cpy_i32(const int n, int32_t * y, const int32_t * x) { for (int i = 0; i < n; ++i) y[i] = x[i]; } + +inline static void ggml_vec_set_f16(const int n, ggml_fp16_t * x, const ggml_fp16_t v) { for (int i = 0; i < n; ++i) x[i] = v; } +inline static void ggml_vec_set_bf16(const int n, ggml_bf16_t * x, const ggml_bf16_t v) { for (int i = 0; i < n; ++i) x[i] = v; } + +inline static void ggml_vec_add_f32 (const int n, float * z, const float * x, const float * y) { + int i = 0; +#if defined(__AVX2__) + for (; i + 7 < n; i += 8) { + __m256 vx = _mm256_loadu_ps(x + i); + __m256 vy = _mm256_loadu_ps(y + i); + __m256 vz = _mm256_add_ps(vx, vy); + _mm256_storeu_ps(z + i, vz); + } +#endif + for (; i < n; ++i) { + z[i] = x[i] + y[i]; + } +} + +inline static void ggml_vec_add_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { + for (int i = 0; i < n; ++i) { + z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) + GGML_CPU_FP16_TO_FP32(y[i])); + } +} +inline static void ggml_vec_add1_f32(const int n, float * z, const float * x, const float v) { for (int i = 0; i < n; ++i) z[i] = x[i] + v; } +inline static void ggml_vec_acc_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] += x[i]; } +inline static void ggml_vec_acc1_f32(const int n, float * y, const float v) { for (int i = 0; i < n; ++i) y[i] += v; } +inline static void ggml_vec_sub_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i] = x[i] - y[i]; } +inline static void ggml_vec_sub_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { + for (int i = 0; i < n; ++i) { + z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) - GGML_CPU_FP16_TO_FP32(y[i])); + } +} +inline static void ggml_vec_set_f32 (const int n, float * x, const float v) { for (int i = 0; i < n; ++i) x[i] = v; } +inline static void ggml_vec_cpy_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i]; } +inline static void ggml_vec_neg_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = -x[i]; } +inline static void ggml_vec_neg_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(-GGML_CPU_FP16_TO_FP32(x[i])); + } +} + +inline static void ggml_vec_mul_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i] = x[i]*y[i]; } +inline static void ggml_vec_mul_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { + for (int i = 0; i < n; ++i) { + z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) * GGML_CPU_FP16_TO_FP32(y[i])); + } +} +inline static void ggml_vec_div_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i] = x[i]/y[i]; } +inline static void ggml_vec_div_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { + for (int i = 0; i < n; ++i) { + z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) / GGML_CPU_FP16_TO_FP32(y[i])); + } +} + +// compute GGML_VEC_DOT_UNROLL dot products at once +// xs - x row stride in bytes +inline static void ggml_vec_dot_f16_unroll(const int n, const int xs, float * GGML_RESTRICT s, void * GGML_RESTRICT xv, ggml_fp16_t * GGML_RESTRICT y) { + ggml_float sumf[GGML_VEC_DOT_UNROLL] = { 0.0 }; + + ggml_fp16_t * GGML_RESTRICT x[GGML_VEC_DOT_UNROLL]; + + for (int i = 0; i < GGML_VEC_DOT_UNROLL; ++i) { + x[i] = (ggml_fp16_t *) ((char *) xv + i*xs); + } + +#if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + + const int sve_register_length = svcntb() * 8; + const int ggml_f16_epr = sve_register_length / 16; // running when 16 + const int ggml_f16_step = 8 * ggml_f16_epr; // choose 8 SVE registers + + int np = (n & ~(ggml_f16_step - 1)); + + svfloat16_t sum_00 = svdup_n_f16(0.0f); + svfloat16_t sum_01 = svdup_n_f16(0.0f); + svfloat16_t sum_02 = svdup_n_f16(0.0f); + svfloat16_t sum_03 = svdup_n_f16(0.0f); + + svfloat16_t sum_10 = svdup_n_f16(0.0f); + svfloat16_t sum_11 = svdup_n_f16(0.0f); + svfloat16_t sum_12 = svdup_n_f16(0.0f); + svfloat16_t sum_13 = svdup_n_f16(0.0f); + + svfloat16_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; + svfloat16_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; + + for (int i = 0; i < np; i += ggml_f16_step) { + ay1 = GGML_F16x_VEC_LOAD(y + i + 0 * ggml_f16_epr, 0); // 8 elements + + ax1 = GGML_F16x_VEC_LOAD(x[0] + i + 0*ggml_f16_epr, 0); // 8 elements + sum_00 = GGML_F16x_VEC_FMA(sum_00, ax1, ay1); // sum_00 = sum_00+ax1*ay1 + ax1 = GGML_F16x_VEC_LOAD(x[1] + i + 0*ggml_f16_epr, 0); // 8 elements + sum_10 = GGML_F16x_VEC_FMA(sum_10, ax1, ay1); + + ay2 = GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 1); // next 8 elements + + ax2 = GGML_F16x_VEC_LOAD(x[0] + i + 1*ggml_f16_epr, 1); // next 8 elements + sum_01 = GGML_F16x_VEC_FMA(sum_01, ax2, ay2); + ax2 = GGML_F16x_VEC_LOAD(x[1] + i + 1*ggml_f16_epr, 1); + sum_11 = GGML_F16x_VEC_FMA(sum_11, ax2, ay2); + + ay3 = GGML_F16x_VEC_LOAD(y + i + 2 * ggml_f16_epr, 2); + + ax3 = GGML_F16x_VEC_LOAD(x[0] + i + 2*ggml_f16_epr, 2); + sum_02 = GGML_F16x_VEC_FMA(sum_02, ax3, ay3); + ax3 = GGML_F16x_VEC_LOAD(x[1] + i + 2*ggml_f16_epr, 2); + sum_12 = GGML_F16x_VEC_FMA(sum_12, ax3, ay3); + + ay4 = GGML_F16x_VEC_LOAD(y + i + 3 * ggml_f16_epr, 3); + + ax4 = GGML_F16x_VEC_LOAD(x[0] + i + 3*ggml_f16_epr, 3); + sum_03 = GGML_F16x_VEC_FMA(sum_03, ax4, ay4); + ax4 = GGML_F16x_VEC_LOAD(x[1] + i + 3*ggml_f16_epr, 3); + sum_13 = GGML_F16x_VEC_FMA(sum_13, ax4, ay4); + + ay5 = GGML_F16x_VEC_LOAD(y + i + 4 * ggml_f16_epr, 4); + + ax5 = GGML_F16x_VEC_LOAD(x[0] + i + 4*ggml_f16_epr, 4); + + sum_00 = GGML_F16x_VEC_FMA(sum_00, ax5, ay5); + ax5 = GGML_F16x_VEC_LOAD(x[1] + i + 4*ggml_f16_epr, 4); + sum_10 = GGML_F16x_VEC_FMA(sum_10, ax5, ay5); + + ay6 = GGML_F16x_VEC_LOAD(y + i + 5 * ggml_f16_epr, 5); + + ax6 = GGML_F16x_VEC_LOAD(x[0] + i + 5*ggml_f16_epr, 5); + + sum_01 = GGML_F16x_VEC_FMA(sum_01, ax6, ay6); + ax6 = GGML_F16x_VEC_LOAD(x[1] + i + 5*ggml_f16_epr, 5); + sum_11 = GGML_F16x_VEC_FMA(sum_11, ax6, ay6); + + ay7 = GGML_F16x_VEC_LOAD(y + i + 6 * ggml_f16_epr, 6); + + ax7 = GGML_F16x_VEC_LOAD(x[0] + i + 6*ggml_f16_epr, 6); + + sum_02 = GGML_F16x_VEC_FMA(sum_02, ax7, ay7); + ax7 = GGML_F16x_VEC_LOAD(x[1] + i + 6*ggml_f16_epr, 6); + sum_12 = GGML_F16x_VEC_FMA(sum_12, ax7, ay7); + + ay8 = GGML_F16x_VEC_LOAD(y + i + 7 * ggml_f16_epr, 7); + + ax8 = GGML_F16x_VEC_LOAD(x[0] + i + 7*ggml_f16_epr, 7); + + sum_03 = GGML_F16x_VEC_FMA(sum_03, ax8, ay8); + ax8 = GGML_F16x_VEC_LOAD(x[1] + i + 7*ggml_f16_epr, 7); + sum_13 = GGML_F16x_VEC_FMA(sum_13, ax8, ay8); + } + + const int np2 = (n & ~(ggml_f16_epr - 1)); + for (int k = np; k < np2; k += ggml_f16_epr) { + svfloat16_t ry = GGML_F16x_VEC_LOAD(y + k, 0); + + svfloat16_t rx = GGML_F16x_VEC_LOAD(x[0] + k, 0); + sum_00 = GGML_F16x_VEC_FMA(sum_00, rx, ry); + rx = GGML_F16x_VEC_LOAD(x[1] + k, 0); + sum_10 = GGML_F16x_VEC_FMA(sum_10, rx, ry); + } + + if (np2 < n) { + svbool_t pg = svwhilelt_b16(np2, n); + svfloat16_t hx_0 = svld1_f16(pg, (const __fp16 *)(x[0] + np2)); + svfloat16_t hx_1 = svld1_f16(pg, (const __fp16 *)(x[1] + np2)); + svfloat16_t hy = svld1_f16(pg, (const __fp16 *)(y + np2)); + + sum_00 = svmad_f16_x(pg, hx_0, hy, sum_00); + sum_10 = svmad_f16_x(pg, hx_1, hy, sum_10); + } + GGML_F16x_VEC_REDUCE(sumf[0], sum_00, sum_01, sum_02, sum_03); + GGML_F16x_VEC_REDUCE(sumf[1], sum_10, sum_11, sum_12, sum_13); + np = n; + #elif defined(__riscv_v_intrinsic) + #if defined(__riscv_zvfh) + size_t vl = __riscv_vsetvlmax_e32m4(); + + // initialize accumulators to all zeroes + vfloat32m4_t vsum0_0 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + vfloat32m4_t vsum0_1 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + vfloat32m4_t vsum1_0 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + vfloat32m4_t vsum1_1 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + + // calculate step size + const size_t epr = __riscv_vsetvlmax_e16m2(); + const size_t step = epr * 2; + int np = (n & ~(step - 1)); + + // unroll by 2 along the row dimension + for (int i = 0; i < np; i += step) { + vfloat16m2_t ay0 = __riscv_vle16_v_f16m2((const _Float16 *)(y + i), epr); + vfloat16m2_t ax0_0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i), epr); + vfloat16m2_t ax1_0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i), epr); + vsum0_0 = __riscv_vfwmacc_vv_f32m4(vsum0_0, ax0_0, ay0, epr); + vsum1_0 = __riscv_vfwmacc_vv_f32m4(vsum1_0, ax1_0, ay0, epr); + + vfloat16m2_t ay1 = __riscv_vle16_v_f16m2((const _Float16 *)(y + i + epr), epr); + vfloat16m2_t ax0_1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i + epr), epr); + vfloat16m2_t ax1_1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i + epr), epr); + vsum0_1 = __riscv_vfwmacc_vv_f32m4(vsum0_1, ax0_1, ay1, epr); + vsum1_1 = __riscv_vfwmacc_vv_f32m4(vsum1_1, ax1_1, ay1, epr); + } + + vfloat32m4_t vsum0 = __riscv_vfadd_vv_f32m4(vsum0_0, vsum0_1, vl); + vfloat32m4_t vsum1 = __riscv_vfadd_vv_f32m4(vsum1_0, vsum1_1, vl); + + // leftovers + for (int i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m2(n - i); + vfloat16m2_t ay = __riscv_vle16_v_f16m2((const _Float16 *)(y + i), vl); + vfloat16m2_t ax0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i), vl); + vfloat16m2_t ax1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i), vl); + + vsum0 = __riscv_vfwmacc_vv_f32m4(vsum0, ax0, ay, vl); + vsum1 = __riscv_vfwmacc_vv_f32m4(vsum1, ax1, ay, vl); + } + + // reduce + vl = __riscv_vsetvlmax_e32m2(); + vfloat32m2_t acc0_0 = __riscv_vfadd_vv_f32m2(__riscv_vget_v_f32m4_f32m2(vsum0, 0), + __riscv_vget_v_f32m4_f32m2(vsum0, 1), vl); + vl = __riscv_vsetvlmax_e32m1(); + vfloat32m1_t acc0_1 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(acc0_0, 0), + __riscv_vget_v_f32m2_f32m1(acc0_0, 1), vl); + vfloat32m1_t redsum0 = __riscv_vfredusum_vs_f32m1_f32m1( + acc0_1, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl); + + vl = __riscv_vsetvlmax_e32m2(); + vfloat32m2_t acc1_0 = __riscv_vfadd_vv_f32m2(__riscv_vget_v_f32m4_f32m2(vsum1, 0), + __riscv_vget_v_f32m4_f32m2(vsum1, 1), vl); + vl = __riscv_vsetvlmax_e32m1(); + vfloat32m1_t acc1_1 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(acc1_0, 0), + __riscv_vget_v_f32m2_f32m1(acc1_0, 1), vl); + vfloat32m1_t redsum1 = __riscv_vfredusum_vs_f32m1_f32m1( + acc1_1, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl); + sumf[0] = __riscv_vfmv_f_s_f32m1_f32(redsum0); + sumf[1] = __riscv_vfmv_f_s_f32m1_f32(redsum1); + np = n; + #else + const int np = 0; + #endif + #else + const int np = (n & ~(GGML_F16_STEP - 1)); + + GGML_F16_VEC sum[GGML_VEC_DOT_UNROLL][GGML_F16_ARR] = { { GGML_F16_VEC_ZERO } }; + + GGML_F16_VEC ax[GGML_F16_ARR]; + GGML_F16_VEC ay[GGML_F16_ARR]; + + for (int i = 0; i < np; i += GGML_F16_STEP) { + for (int j = 0; j < GGML_F16_ARR; j++) { + ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); + + for (int k = 0; k < GGML_VEC_DOT_UNROLL; ++k) { + ax[j] = GGML_F16_VEC_LOAD(x[k] + i + j*GGML_F16_EPR, j); + + sum[k][j] = GGML_F16_VEC_FMA(sum[k][j], ax[j], ay[j]); + } + } + } + + // reduce sum0..sum3 to sum0 + for (int k = 0; k < GGML_VEC_DOT_UNROLL; ++k) { + GGML_F16_VEC_REDUCE(sumf[k], sum[k]); + } + #endif +#else + // scalar path + const int np = 0; +#endif + // scalar and leftovers + for (int i = np; i < n; ++i) { + for (int j = 0; j < GGML_VEC_DOT_UNROLL; ++j) { + sumf[j] += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[j][i])*GGML_CPU_FP16_TO_FP32(y[i])); + } + } + + for (int i = 0; i < GGML_VEC_DOT_UNROLL; ++i) { + s[i] = (float)sumf[i]; + } +} + +inline static void ggml_vec_mad_f32(const int n, float * GGML_RESTRICT y, const float * GGML_RESTRICT x, const float v) { +#if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + + const int sve_register_length = ggml_cpu_get_sve_cnt() * 8; + const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16 + const int ggml_f32_step = 8 * ggml_f32_epr; // choose 8 SVE registers + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); + + const int np = (n & ~(ggml_f32_step - 1)); + svfloat32_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; + svfloat32_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; + for (int i = 0; i < np; i += ggml_f32_step) { + + ax1 = GGML_F32_VEC_LOAD(x + i); + ay1 = GGML_F32_VEC_LOAD(y + i); + ay1 = GGML_F32_VEC_FMA(ay1, ax1, vx); + + GGML_F32_VEC_STORE(y + i, ay1); + + ax2 = GGML_F32_VEC_LOAD(x + i + 1*ggml_f32_epr); + ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr); + ay2 = GGML_F32_VEC_FMA(ay2, ax2, vx); + + GGML_F32_VEC_STORE(y + i + 1*ggml_f32_epr, ay2); + + ax3 = GGML_F32_VEC_LOAD(x + i + 2*ggml_f32_epr); + ay3 = GGML_F32_VEC_LOAD(y + i + 2*ggml_f32_epr); + ay3 = GGML_F32_VEC_FMA(ay3, ax3, vx); + + GGML_F32_VEC_STORE(y + i + 2*ggml_f32_epr, ay3); + + ax4 = GGML_F32_VEC_LOAD(x + i + 3*ggml_f32_epr); + ay4 = GGML_F32_VEC_LOAD(y + i + 3*ggml_f32_epr); + ay4 = GGML_F32_VEC_FMA(ay4, ax4, vx); + + GGML_F32_VEC_STORE(y + i + 3*ggml_f32_epr, ay4); + + ax5 = GGML_F32_VEC_LOAD(x + i + 4*ggml_f32_epr); + ay5 = GGML_F32_VEC_LOAD(y + i + 4*ggml_f32_epr); + ay5 = GGML_F32_VEC_FMA(ay5, ax5, vx); + + GGML_F32_VEC_STORE(y + i + 4*ggml_f32_epr, ay5); + + ax6 = GGML_F32_VEC_LOAD(x + i + 5*ggml_f32_epr); + ay6 = GGML_F32_VEC_LOAD(y + i + 5*ggml_f32_epr); + ay6 = GGML_F32_VEC_FMA(ay6, ax6, vx); + + GGML_F32_VEC_STORE(y + i + 5*ggml_f32_epr, ay6); + + ax7 = GGML_F32_VEC_LOAD(x + i + 6*ggml_f32_epr); + ay7 = GGML_F32_VEC_LOAD(y + i + 6*ggml_f32_epr); + ay7 = GGML_F32_VEC_FMA(ay7, ax7, vx); + + GGML_F32_VEC_STORE(y + i + 6*ggml_f32_epr, ay7); + + ax8 = GGML_F32_VEC_LOAD(x + i + 7*ggml_f32_epr); + ay8 = GGML_F32_VEC_LOAD(y + i + 7*ggml_f32_epr); + ay8 = GGML_F32_VEC_FMA(ay8, ax8, vx); + + GGML_F32_VEC_STORE(y + i + 7*ggml_f32_epr, ay8); + } + // leftovers + // Since 8 unrolls are done in above loop, leftovers lie in range [0, ggml_f32_step] which is handled in below loop + const int np2 = (n & ~(ggml_f32_epr - 1)); + for (int i = np; i < np2; i += ggml_f32_epr) { + ax1 = GGML_F32_VEC_LOAD(x + i); + ay1 = GGML_F32_VEC_LOAD(y + i); + ay1 = GGML_F32_VEC_FMA(ay1, ax1, vx); + + GGML_F32_VEC_STORE(y + i, ay1); + } + // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only + if (np2 < n) { + svbool_t pg =svwhilelt_b32(np2, n); + ax1 = svld1_f32(pg, x + np2); + ay1 = svld1_f32(pg, y + np2); + ay1 = svmad_f32_m(pg, ax1, vx, ay1); + + svst1_f32(pg, y + np2, ay1); + } + #elif defined(__riscv_v_intrinsic) + for (int i = 0, avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m8(n - i); + vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[i], avl); + vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl); + vfloat32m8_t ny = __riscv_vfmadd_vf_f32m8(ax, v, ay, avl); + __riscv_vse32_v_f32m8(&y[i], ny, avl); + } + #else + const int np = (n & ~(GGML_F32_STEP - 1)); + + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); + + GGML_F32_VEC ax[GGML_F32_ARR]; + GGML_F32_VEC ay[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ax[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_FMA(ay[j], ax[j], vx); + + GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); + } + } + + // leftovers + for (int i = np; i < n; ++i) { + y[i] += x[i]*v; + } + #endif +#else + // scalar + for (int i = 0; i < n; ++i) { + y[i] += x[i]*v; + } +#endif +} + +inline static void ggml_vec_mad_f16(const int n, ggml_fp16_t * GGML_RESTRICT y, const ggml_fp16_t * GGML_RESTRICT x, const float v) { +#if defined(GGML_SIMD) && defined(__ARM_FEATURE_SVE) + const int sve_register_length = svcntb() * 8; + const int ggml_f16_epr = sve_register_length / 16; + const int ggml_f16_step = 8 * ggml_f16_epr; + + GGML_F16x_VEC vx = GGML_F16x_VEC_SET1(v); + + int np = (n & ~(ggml_f16_step - 1)); + + svfloat16_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; + svfloat16_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; + for (int i = 0; i < np; i += ggml_f16_step) { + ax1 = GGML_F16x_VEC_LOAD(x + i + 0 * ggml_f16_epr, 0); + ay1 = GGML_F16x_VEC_LOAD(y + i + 0 * ggml_f16_epr, 0); + ay1 = GGML_F16x_VEC_FMA(ay1, ax1, vx); + + GGML_F16x_VEC_STORE(y + i + 0 * ggml_f16_epr, ay1, 0); + + ax2 = GGML_F16x_VEC_LOAD(x + i + 1 * ggml_f16_epr, 1); + ay2 = GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 1); + ay2 = GGML_F16x_VEC_FMA(ay2, ax2, vx); + + GGML_F16x_VEC_STORE(y + i + 1 * ggml_f16_epr, ay2, 1); + + ax3 = GGML_F16x_VEC_LOAD(x + i + 2 * ggml_f16_epr, 2); + ay3 = GGML_F16x_VEC_LOAD(y + i + 2 * ggml_f16_epr, 2); + ay3 = GGML_F16x_VEC_FMA(ay3, ax3, vx); + + GGML_F16x_VEC_STORE(y + i + 2 * ggml_f16_epr, ay3, 2); + + ax4 = GGML_F16x_VEC_LOAD(x + i + 3 * ggml_f16_epr, 3); + ay4 = GGML_F16x_VEC_LOAD(y + i + 3 * ggml_f16_epr, 3); + ay4 = GGML_F16x_VEC_FMA(ay4, ax4, vx); + + GGML_F16x_VEC_STORE(y + i + 3 * ggml_f16_epr, ay4, 3); + + ax5 = GGML_F16x_VEC_LOAD(x + i + 4 * ggml_f16_epr, 4); + ay5 = GGML_F16x_VEC_LOAD(y + i + 4 * ggml_f16_epr, 4); + ay5 = GGML_F16x_VEC_FMA(ay5, ax5, vx); + + GGML_F16x_VEC_STORE(y + i + 4 * ggml_f16_epr, ay5, 4); + + ax6 = GGML_F16x_VEC_LOAD(x + i + 5 * ggml_f16_epr, 5); + ay6 = GGML_F16x_VEC_LOAD(y + i + 5 * ggml_f16_epr, 5); + ay6 = GGML_F16x_VEC_FMA(ay6, ax6, vx); + + GGML_F16x_VEC_STORE(y + i + 5 * ggml_f16_epr, ay6, 5); + + ax7 = GGML_F16x_VEC_LOAD(x + i + 6 * ggml_f16_epr, 6); + ay7 = GGML_F16x_VEC_LOAD(y + i + 6 * ggml_f16_epr, 6); + ay7 = GGML_F16x_VEC_FMA(ay7, ax7, vx); + + GGML_F16x_VEC_STORE(y + i + 6 * ggml_f16_epr, ay7, 6); + + ax8 = GGML_F16x_VEC_LOAD(x + i + 7 * ggml_f16_epr, 7); + ay8 = GGML_F16x_VEC_LOAD(y + i + 7 * ggml_f16_epr, 7); + ay8 = GGML_F16x_VEC_FMA(ay8, ax8, vx); + + GGML_F16x_VEC_STORE(y + i + 7 * ggml_f16_epr, ay8, 7); + } + const int np2 = (n & ~(ggml_f16_epr - 1)); + for (int k = np; k < np2; k += ggml_f16_epr) { + svfloat16_t rx = GGML_F16x_VEC_LOAD(x + k, 0); + svfloat16_t ry = GGML_F16x_VEC_LOAD(y + k, 0); + ry = GGML_F16x_VEC_FMA(ry, rx, vx); + + GGML_F16x_VEC_STORE(y + k, ry, 0); + } + + if (np2 < n) { + svbool_t pg = svwhilelt_b16(np2, n); + svfloat16_t hx = svld1_f16(pg, (const __fp16 *)(x + np2)); + svfloat16_t hy = svld1_f16(pg, (const __fp16 *)(y + np2)); + hy = svmad_f16_x(pg, hx, vx, hy); + svst1_f16(pg, (__fp16 *)(y + np2), hy); + } + np = n; +#elif defined(__riscv_v_intrinsic) // implies __riscv_v_intrinsic + #if defined (__riscv_zvfh) + const ggml_fp16_t s = GGML_CPU_FP32_TO_FP16(v); + const _Float16 scale = *(const _Float16*)(&s); + + // calculate step size + const int epr = __riscv_vsetvlmax_e16m4(); + const int step = epr * 2; + int np = (n & ~(step - 1)); + + // unroll by 2 + for (int i = 0; i < np; i += step) { + vfloat16m4_t ax0 = __riscv_vle16_v_f16m4((const _Float16*)x + i, epr); + vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, epr); + ay0 = __riscv_vfmacc_vf_f16m4(ay0, scale, ax0, epr); + __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, epr); + __asm__ __volatile__ ("" ::: "memory"); + + vfloat16m4_t ax1 = __riscv_vle16_v_f16m4((const _Float16*)x + i + epr, epr); + vfloat16m4_t ay1 = __riscv_vle16_v_f16m4((const _Float16*)y + i + epr, epr); + ay1 = __riscv_vfmacc_vf_f16m4(ay1, scale, ax1, epr); + __riscv_vse16_v_f16m4((_Float16*)y + i + epr, ay1, epr); + __asm__ __volatile__ ("" ::: "memory"); + } + + // leftovers + int vl; + for (int i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m4(n - i); + vfloat16m4_t ax0 = __riscv_vle16_v_f16m4((const _Float16*)x + i, vl); + vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, vl); + ay0 = __riscv_vfmacc_vf_f16m4(ay0, scale, ax0, vl); + __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, vl); + } + np = n; + #else + // fall to scalar path + const int np = 0; + #endif +#elif defined(GGML_SIMD) + const int np = (n & ~(GGML_F16_STEP - 1)); + + GGML_F16_VEC vx = GGML_F16_VEC_SET1(v); + + GGML_F16_VEC ax[GGML_F16_ARR]; + GGML_F16_VEC ay[GGML_F16_ARR]; + + for (int i = 0; i < np; i += GGML_F16_STEP) { + for (int j = 0; j < GGML_F16_ARR; j++) { + ax[j] = GGML_F16_VEC_LOAD(x + i + j*GGML_F16_EPR, j); + ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); + ay[j] = GGML_F16_VEC_FMA(ay[j], ax[j], vx); + + GGML_F16_VEC_STORE(y + i + j*GGML_F16_EPR, ay, j); + } + } +#else + // scalar path + const int np = 0; +#endif + + // scalar and leftovers + for (int i = np; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(y[i]) + GGML_CPU_FP16_TO_FP32(x[i])*v); + } +} + +// xs and vs are byte strides of x and v +inline static void ggml_vec_mad_f32_unroll(const int n, const int xs, const int vs, float * GGML_RESTRICT y, const float * GGML_RESTRICT xv, const float * GGML_RESTRICT vv) { + + const float * GGML_RESTRICT x[GGML_VEC_MAD_UNROLL]; + const float * GGML_RESTRICT v[GGML_VEC_MAD_UNROLL]; + + for (int i = 0; i < GGML_VEC_MAD_UNROLL; ++i) { + x[i] = (const float *) ((const char *) xv + i*xs); + v[i] = (const float *) ((const char *) vv + i*vs); + } + +#if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + // scalar Route to scalar implementation //TODO: Write SVE code + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + for (int i = 0; i < n; ++i) { + y[i] += x[k][i]*v[k][0]; + } + } + #elif defined(__riscv_v_intrinsic) + for (int i = 0, avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m8(n - i); + vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl); + for (int k = 0; k < GGML_VEC_MAD_UNROLL; k++) { + vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[k][i], avl); + ay = __riscv_vfmadd_vf_f32m8(ax, v[k][0], ay, avl); + } + __riscv_vse32_v_f32m8(&y[i], ay, avl); + } + #else + const int np = (n & ~(GGML_F32_STEP - 1)); + + GGML_F32_VEC vx[GGML_VEC_MAD_UNROLL]; + + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + vx[k] = GGML_F32_VEC_SET1(v[k][0]); + } + + GGML_F32_VEC ax[GGML_VEC_MAD_UNROLL][GGML_F32_ARR]; + GGML_F32_VEC ay[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); + + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + ax[k][j] = GGML_F32_VEC_LOAD(x[k] + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_FMA(ay[j], ax[k][j], vx[k]); + } + + GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); + } + } + + // leftovers + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + for (int i = np; i < n; ++i) { + y[i] += x[k][i]*v[k][0]; + } + } + #endif +#else + // scalar + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + for (int i = 0; i < n; ++i) { + y[i] += x[k][i]*v[k][0]; + } + } +#endif +} + +inline static void ggml_vec_mad1_f32(const int n, float * y, const float * x, const float s, const float b) { +#if defined(GGML_USE_ACCELERATE) + vDSP_vsmsa(x, 1, &s, &b, y, 1, n); +#elif defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + // scalar ; TODO: Write SVE code + for (int i = 0; i < n; ++i) { + y[i] = x[i]*s + b; + } + #elif defined(__riscv_v_intrinsic) + for (int i = 0, avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m8(n - i); + vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[i], avl); + vfloat32m8_t vb = __riscv_vfmv_v_f_f32m8(b, avl); + vfloat32m8_t ny = __riscv_vfmadd_vf_f32m8(ax, s, vb, avl); + __riscv_vse32_v_f32m8(&y[i], ny, avl); + } + #else + const int np = (n & ~(GGML_F32_STEP - 1)); + + GGML_F32_VEC vs = GGML_F32_VEC_SET1(s); + GGML_F32_VEC vb = GGML_F32_VEC_SET1(b); + + GGML_F32_VEC ay[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ay[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_FMA(vb, ay[j], vs); + + GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); + } + } + + // leftovers + for (int i = np; i < n; ++i) { + y[i] = x[i]*s + b; + } + #endif +#else + // scalar + for (int i = 0; i < n; ++i) { + y[i] = x[i]*s + b; + } +#endif +} + +//inline static void ggml_vec_scale_f32(const int n, float * y, const float v) { for (int i = 0; i < n; ++i) y[i] *= v; } +inline static void ggml_vec_scale_f32(const int n, float * y, const float v) { +#if defined(GGML_USE_ACCELERATE) + vDSP_vsmul(y, 1, &v, y, 1, n); +#elif defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + const int sve_register_length = ggml_cpu_get_sve_cnt() * 8; + const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16 + const int ggml_f32_step = 2 * ggml_f32_epr; + + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); + const int np = (n & ~(ggml_f32_step - 1)); + svfloat32_t ay1; + svfloat32_t ay2; + for (int i = 0; i < np; i += ggml_f32_step) { + ay1 = GGML_F32_VEC_LOAD(y + i); + ay1 = GGML_F32_VEC_MUL(ay1, vx); + GGML_F32_VEC_STORE(y + i, ay1); + + ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr); + ay2 = GGML_F32_VEC_MUL(ay2, vx); + GGML_F32_VEC_STORE(y + i + 1*ggml_f32_epr, ay2); + } + // leftovers + // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only + for (int i = np; i < n; i += ggml_f32_epr) { + svbool_t pg = svwhilelt_b32(i, n); + ay1 = svld1_f32(pg, y + i); + ay1 = svmul_f32_m(pg, ay1, vx); + svst1_f32(pg, y + i, ay1); + } + #elif defined(__riscv_v_intrinsic) + for (int i = 0, avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m8(n - i); + vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl); + vfloat32m8_t ny = __riscv_vfmul_vf_f32m8(ay, v, avl); + __riscv_vse32_v_f32m8(&y[i], ny, avl); + } + #else + const int np = (n & ~(GGML_F32_STEP - 1)); + + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); + + GGML_F32_VEC ay[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_MUL(ay[j], vx); + + GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); + } + } + + // leftovers + for (int i = np; i < n; ++i) { + y[i] *= v; + } + #endif +#else + // scalar + for (int i = 0; i < n; ++i) { + y[i] *= v; + } +#endif +} + +inline static void ggml_vec_scale_f16(const int n, ggml_fp16_t * y, const float v) { +#if defined(GGML_SIMD) && defined(__ARM_FEATURE_SVE) + const int sve_register_length = svcntb() * 8; + const int ggml_f16_epr = sve_register_length / 16; + const int ggml_f16_step = 2 * ggml_f16_epr; + + GGML_F16x_VEC vx = GGML_F16x_VEC_SET1(v); + int np = (n & ~(ggml_f16_step - 1)); + svfloat16_t ay1, ay2; + + for (int i = 0; i < np; i += ggml_f16_step) { + ay1 = GGML_F16x_VEC_LOAD(y + i + 0*ggml_f16_epr, 0); + ay1 = GGML_F16x_VEC_MUL(ay1, vx); + GGML_F16x_VEC_STORE(y + i + 0*ggml_f16_epr, ay1, 0); + + ay2 = GGML_F16x_VEC_LOAD(y + i + 1*ggml_f16_epr, 1); + ay2 = GGML_F16x_VEC_MUL(ay2, vx); + GGML_F16x_VEC_STORE(y + i + 1*ggml_f16_epr, ay2, 1); + } + // leftovers + // maximum number of leftover elements will be less that ggmlF_16x_epr. Apply predicated svmad on available elements only + if (np < n) { + svbool_t pg = svwhilelt_b16(np, n); + svfloat16_t hy = svld1_f16(pg, (__fp16 *)(y + np)); + svfloat16_t out = svmul_f16_m(pg, hy, vx); + svst1_f16(pg, (__fp16 *)(y + np), out); + } + np = n; +#elif defined(__riscv_v_intrinsic) + #if defined(__riscv_zvfh) + const ggml_fp16_t s = GGML_CPU_FP32_TO_FP16(v); + const _Float16 scale = *(const _Float16*)(&s); + + // calculate step size + const int epr = __riscv_vsetvlmax_e16m4(); + const int step = epr * 2; + int np = (n & ~(step - 1)); + + // unroll by 2 + for (int i = 0; i < np; i += step) { + vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, epr); + ay0 = __riscv_vfmul_vf_f16m4(ay0, scale, epr); + __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, epr); + __asm__ __volatile__ ("" ::: "memory"); + + vfloat16m4_t ay1 = __riscv_vle16_v_f16m4((const _Float16*)y + i + epr, epr); + ay1 = __riscv_vfmul_vf_f16m4(ay1, scale, epr); + __riscv_vse16_v_f16m4((_Float16*)y + i + epr, ay1, epr); + __asm__ __volatile__ ("" ::: "memory"); + } + + // leftovers + int vl; + for (int i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m4(n - i); + vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, vl); + ay0 = __riscv_vfmul_vf_f16m4(ay0, scale, vl); + __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, vl); + } + np = n; + #else + // fall to scalar path + const int np = 0; + #endif +#elif defined(GGML_SIMD) + const int np = (n & ~(GGML_F16_STEP - 1)); + + GGML_F16_VEC vx = GGML_F16_VEC_SET1(v); + + GGML_F16_VEC ay[GGML_F16_ARR]; + + for (int i = 0; i < np; i += GGML_F16_STEP) { + for (int j = 0; j < GGML_F16_ARR; j++) { + ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); + ay[j] = GGML_F16_VEC_MUL(ay[j], vx); + + GGML_F16_VEC_STORE(y + i + j*GGML_F16_EPR, ay, j); + } + } +#else + // scalar path + const int np = 0; +#endif + // scalar and leftovers + for (int i = np; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(y[i])*v); + } +} + +inline static void ggml_vec_norm_f32 (const int n, float * s, const float * x) { ggml_vec_dot_f32(n, s, 0, x, 0, x, 0, 1); *s = sqrtf(*s); } +inline static void ggml_vec_sqr_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i]*x[i]; } +inline static void ggml_vec_sqr_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16(v*v); + } +} +inline static void ggml_vec_sqrt_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = sqrtf(x[i]); } +inline static void ggml_vec_sqrt_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(sqrtf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_log_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = logf(x[i]); } +inline static void ggml_vec_log_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(logf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_sin_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = sinf(x[i]); } +inline static void ggml_vec_sin_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(sinf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_cos_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = cosf(x[i]); } +inline static void ggml_vec_cos_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(cosf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_abs_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = fabsf(x[i]); } +inline static void ggml_vec_abs_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(fabsf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_sgn_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? 1.f : ((x[i] < 0.f) ? -1.f : 0.f); } +inline static void ggml_vec_sgn_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? 1.f : ((v < 0.f) ? -1.f : 0.f)); + } +} +inline static void ggml_vec_step_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? 1.f : 0.f; } +inline static void ggml_vec_step_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16((GGML_CPU_FP16_TO_FP32(x[i]) > 0.f) ? 1.f : 0.f); + } +} +inline static void ggml_vec_tanh_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = tanhf(x[i]); } +inline static void ggml_vec_tanh_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(tanhf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_elu_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? x[i] : expm1f(x[i]); } +inline static void ggml_vec_elu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + const float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v : expm1f(v)); + } +} +inline static void ggml_vec_relu_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? x[i] : 0.f; } +inline static void ggml_vec_relu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v : 0.f); + } +} +inline static void ggml_vec_leaky_relu_f32 (const int n, float * y, const float * x, const float ns) { for (int i = 0; i < n; ++i) y[i] = ((x[i] > 0.f) ? x[i] : 0.f) + ns * ((x[i] < 0.0f) ? x[i] : 0.f); } +inline static void ggml_vec_leaky_relu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const float ns) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16(((v > 0.f) ? v : 0.f) + ns * ((v < 0.0f) ? v : 0.f)); + } +} +inline static void ggml_vec_sigmoid_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = 1.f / (1.f + expf(-x[i])); } +inline static void ggml_vec_sigmoid_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(1.f / (1.f + expf(-GGML_CPU_FP16_TO_FP32(x[i])))); + } +} +// TODO: optimize performance +inline static void ggml_vec_hardswish_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i] * fminf(1.0f, fmaxf(0.0f, (x[i] + 3.0f) / 6.0f)); } +inline static void ggml_vec_hardswish_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16(v * fminf(1.0f, fmaxf(0.0f, (v + 3.0f) / 6.0f))); + } +} +inline static void ggml_vec_hardsigmoid_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = fminf(1.0f, fmaxf(0.0f, (x[i] + 3.0f) / 6.0f)); } +inline static void ggml_vec_hardsigmoid_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(fminf(1.0f, fmaxf(0.0f, (GGML_CPU_FP16_TO_FP32(x[i]) + 3.0f) / 6.0f))); + } +} +inline static void ggml_vec_exp_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = expf(x[i]); } +inline static void ggml_vec_exp_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(expf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} + +static const float GELU_COEF_A = 0.044715f; +static const float GELU_QUICK_COEF = -1.702f; +static const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; +static const float SQRT_2_INV = 0.70710678118654752440084436210484f; + +inline static float ggml_gelu_f32(float x) { + return 0.5f*x*(1.0f + tanhf(SQRT_2_OVER_PI*x*(1.0f + GELU_COEF_A*x*x))); +} + +inline static void ggml_vec_gelu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + const uint16_t * i16 = (const uint16_t *) x; + for (int i = 0; i < n; ++i) { + y[i] = ggml_table_gelu_f16[i16[i]]; + } +} + +inline static void ggml_vec_gelu_erf_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float xi = GGML_CPU_FP16_TO_FP32(x[i]); + float res = 0.5f*xi*(1.0f + erff(xi*SQRT_2_INV)); + y[i] = GGML_CPU_FP32_TO_FP16(res); + } +} + +#ifdef GGML_GELU_FP16 +inline static void ggml_vec_gelu_f32(const int n, float * y, const float * x) { + uint16_t t; + for (int i = 0; i < n; ++i) { + if (x[i] <= -10.0f) { + y[i] = 0.0f; + } else if (x[i] >= 10.0f) { + y[i] = x[i]; + } else { + ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); + memcpy(&t, &fp16, sizeof(uint16_t)); + y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[t]); + } + } +} +#else +inline static void ggml_vec_gelu_f32(const int n, float * y, const float * x) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_gelu_f32(x[i]); + } +} +#endif + +inline static void ggml_vec_gelu_erf_f32(const int n, float * y, const float * x) { + for (int i = 0; i < n; ++i) { + float xi = x[i]; + y[i] = 0.5f*xi*(1.0f + erff(xi*SQRT_2_INV)); + } +} + +inline static float ggml_gelu_quick_f32(float x) { + return x*(1.0f/(1.0f+expf(GELU_QUICK_COEF*x))); +} + +inline static void ggml_vec_gelu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + const uint16_t * i16 = (const uint16_t *) x; + for (int i = 0; i < n; ++i) { + y[i] = ggml_table_gelu_quick_f16[i16[i]]; + } +} + +#ifdef GGML_GELU_QUICK_FP16 +inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) { + uint16_t t; + for (int i = 0; i < n; ++i) { + ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); + memcpy(&t, &fp16, sizeof(uint16_t)); + y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[t]); + } +} +#else +inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_gelu_quick_f32(x[i]); + } +} +#endif + +// Sigmoid Linear Unit (SiLU) function +inline static float ggml_silu_f32(float x) { + return x/(1.0f + expf(-x)); +} +inline static ggml_fp16_t ggml_silu_f16(ggml_fp16_t x) { + float v = GGML_CPU_FP16_TO_FP32(x); + return GGML_CPU_FP32_TO_FP16(v/(1.0f + expf(-v))); +} + +#if __FINITE_MATH_ONLY__ +#error "some routines in ggml.c.inc require non-finite math arithmetics -- pass -fno-finite-math-only to the compiler to fix" +#error "ref: https://github.com/ggml-org/llama.cpp/pull/7154#issuecomment-2143844461" +#endif + +/* Below function was borrowed from the GitHub repository: +https://github.com/openvinotoolkit/openvino/blob/master/src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/common.h.incpp */ +#if defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + inline static svfloat32_t exp_ps_sve(svbool_t pg, svfloat32_t src) { + // Constants + const svfloat32_t log2_e = svdup_n_f32(1.4426950409f); + const svfloat32_t ln2 = svdup_n_f32(0.6931473921f); + const svfloat32_t half_ln2_sq = svdup_n_f32(0.2413862043f); + const svuint32_t not_mask17 = svdup_n_u32(~((1u << 17) - 1)); + const svfloat32_t one = svdup_n_f32(1.0f); + const svfloat32_t inactive1 = svdup_n_f32(0.0f); + const svint32_t inactive2 = svdup_n_s32(0); + + // Algorithm starts here + svfloat32_t t0 = svmul_f32_m(pg, src, log2_e); // y = x * log2(e) + svfloat32_t t1 = svrintm_f32_m(inactive1, pg, t0); // rount to int (float) + svint32_t t2 = svcvt_s32_f32_m(inactive2, pg, t1); // n + + t1 = svsub_f32_m(pg, t0, t1); // a = y - floor(y) + t1 = svadd_f32_m(pg, t1, one); // b = a + 1 + + svuint32_t t3 = svlsr_n_u32_m(pg, svreinterpret_u32_f32(t1), 17); // v = b >> 17 (u32) + svfloat32_t t4 = svexpa_f32(t3); // c = fexpa(v) + t4 = svscale_f32_m(pg, t4, t2); // fexpa(v) * 2^(n) + + // and_(t2.d, t1.d, not_mask17.d) + svfloat32_t t5 = svreinterpret_f32_u32(svand_u32_m(pg, svreinterpret_u32_f32(t1), not_mask17)); + t5 = svsub_f32_m(pg, t1, t5); // z + t0 = svmla_f32_m(pg, ln2, t5, half_ln2_sq); // ln2 + half_ln2_sq * z + t0 = svmla_f32_m(pg, one, t5, t0); // 1 + (ln2 * z) + (half_ln2_sq * z * z) + t0 = svmul_f32_m(pg, t0, t4); // Final result + + return t0; + } +#endif + +#if defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + +inline static svfloat32_t ggml_v_expf(svbool_t pg, svfloat32_t x) { + const svfloat32_t r = svdup_n_f32_x(pg, 0x1.8p23f); + const svfloat32_t z = svmla_n_f32_x(pg, r, x, 0x1.715476p+0f); + const svfloat32_t n = svsub_f32_x(pg, z, r); + const svfloat32_t b = svmls_n_f32_x(pg, svmls_n_f32_x(pg, x, n, 0x1.62e4p-1f), n, 0x1.7f7d1cp-20f); + const svuint32_t e = svlsl_n_u32_x(pg, svreinterpret_u32_f32(z), 23); + const svfloat32_t k = svreinterpret_f32_u32(svadd_u32_x(pg, e, svreinterpret_u32_f32(svdup_n_f32_x(pg, 1)))); + const svbool_t c = svacgt_n_f32(pg, n, 126); + const svfloat32_t u = svmul_f32_x(pg, b, b); + const svfloat32_t j = svmla_f32_x(pg, + svmul_n_f32_x(pg, b, 0x1.ffffecp-1f), + svmla_f32_x(pg, svmla_f32_x(pg, svdup_n_f32_x(pg, 0x1.fffdb6p-2f), svdup_n_f32_x(pg, 0x1.555e66p-3f), b), + svmla_f32_x(pg, svdup_n_f32_x(pg, 0x1.573e2ep-5f), svdup_n_f32_x(pg, 0x1.0e4020p-7f), b), u), u); + const svuint32_t d = svdup_n_u32_z(svcmple_n_f32(pg, n, 0.0), 0x82000000); + const svfloat32_t s1 = svreinterpret_f32_u32(svadd_n_u32_x(pg, d, 0x7f000000)); + const svfloat32_t s2 = svreinterpret_f32_u32(svsub_u32_x(pg, e, d)); + return svsel_f32(svacgt_f32(pg, n, svdup_n_f32_x(pg, 192)), svmul_f32_x(pg, s1, s1), + svsel_f32(c, svmul_f32_x(pg, svmla_f32_x(pg, s2, s2, j), s1), svmla_f32_x(pg, k, k, j))); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static svfloat32_t ggml_v_silu(svbool_t pg, svfloat32_t x) { + const svfloat32_t one = svdup_n_f32_x(pg, 1.0f); + const svfloat32_t zero = svdup_n_f32_x(pg, 0.0f); + const svfloat32_t neg_x = svsub_f32_x(pg, zero, x); + const svfloat32_t exp_neg_x = ggml_v_expf(pg, neg_x); + const svfloat32_t one_plus_exp_neg_x = svadd_f32_x(pg, one, exp_neg_x); + return svdiv_f32_x(pg, x, one_plus_exp_neg_x); +} + +#elif defined(__ARM_NEON) && defined(__aarch64__) + +// 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 +inline static float32x4_t ggml_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))); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static float32x4_t ggml_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 = ggml_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); +} + +#elif defined(__AVX512F__) && defined(__AVX512DQ__) + +// 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 +inline static __m512 ggml_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); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static __m512 ggml_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 = ggml_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); +} + +#elif 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 +inline static __m256 ggml_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))))); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static __m256 ggml_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 = ggml_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); +} + +#elif defined(__SSE2__) // __AVX2__ / __ARM_NEON + +#if defined(__FMA__) +#define MADD128(x, y, z) _mm_fmadd_ps(x, y, z) +#define NMADD128(x, y, z) _mm_fnmadd_ps(x, y, z) +#else +#define MADD128(x, y, z) _mm_add_ps(_mm_mul_ps(x, y), z) +#define NMADD128(x, y, z) _mm_sub_ps(z, _mm_mul_ps(x, y)) +#endif + +// 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 +inline static __m128 ggml_v_expf(__m128 x) { + const __m128 r = _mm_set1_ps(0x1.8p23f); + const __m128 z = MADD128(x, _mm_set1_ps(0x1.715476p+0f), r); + const __m128 n = _mm_sub_ps(z, r); + const __m128 b = + NMADD128(n, _mm_set1_ps(0x1.7f7d1cp-20f), NMADD128(n, _mm_set1_ps(0x1.62e4p-1f), x)); + const __m128i e = _mm_slli_epi32(_mm_castps_si128(z), 23); + const __m128 k = _mm_castsi128_ps(_mm_add_epi32(e, _mm_castps_si128(_mm_set1_ps(1)))); + const __m128i c = + _mm_castps_si128(_mm_cmpgt_ps(_mm_andnot_ps(_mm_set1_ps(-0.f), n), _mm_set1_ps(126))); + const __m128 u = _mm_mul_ps(b, b); + const __m128 j = + MADD128(MADD128(MADD128(_mm_set1_ps(0x1.0e4020p-7f), b, _mm_set1_ps(0x1.573e2ep-5f)), u, + MADD128(_mm_set1_ps(0x1.555e66p-3f), b, _mm_set1_ps(0x1.fffdb6p-2f))), + u, _mm_mul_ps(_mm_set1_ps(0x1.ffffecp-1f), b)); + if (!_mm_movemask_epi8(c)) + return MADD128(j, k, k); + const __m128i g = _mm_and_si128(_mm_castps_si128(_mm_cmple_ps(n, _mm_setzero_ps())), + _mm_set1_epi32(0x82000000u)); + const __m128 s1 = _mm_castsi128_ps(_mm_add_epi32(g, _mm_set1_epi32(0x7f000000u))); + const __m128 s2 = _mm_castsi128_ps(_mm_sub_epi32(e, g)); + const __m128i d = + _mm_castps_si128(_mm_cmpgt_ps(_mm_andnot_ps(_mm_set1_ps(-0.f), n), _mm_set1_ps(192))); + return _mm_or_ps( + _mm_and_ps(_mm_castsi128_ps(d), _mm_mul_ps(s1, s1)), + _mm_andnot_ps(_mm_castsi128_ps(d), + _mm_or_ps(_mm_and_ps(_mm_castsi128_ps(c), _mm_mul_ps(MADD128(s2, j, s2), s1)), + _mm_andnot_ps(_mm_castsi128_ps(c), MADD128(k, j, k))))); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static __m128 ggml_v_silu(__m128 x) { + const __m128 one = _mm_set1_ps(1); + const __m128 zero = _mm_setzero_ps(); + const __m128 neg_x = _mm_sub_ps(zero, x); + const __m128 exp_neg_x = ggml_v_expf(neg_x); + const __m128 one_plus_exp_neg_x = _mm_add_ps(one, exp_neg_x); + return _mm_div_ps(x, one_plus_exp_neg_x); +} + +#elif defined(__riscv_v_intrinsic) + +// 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 +inline static vfloat32m2_t ggml_v_expf_m2(vfloat32m2_t x, int vl) { + const vfloat32m2_t r = __riscv_vfmv_v_f_f32m2(0x1.8p23f, vl); +#ifdef __riscv_xtheadvector + // workaround for compiler bug (gcc 14.3.0: Error: unrecognized opcode `th.vmv1r.v v2,v4') + vfloat32m2_t z = __riscv_vfadd_vf_f32m2(r, 0.0f, vl); + z = __riscv_vfmacc_vf_f32m2(z, 0x1.715476p+0f, x, vl); +#else + const vfloat32m2_t z = __riscv_vfmacc_vf_f32m2(r, 0x1.715476p+0f, x, vl); +#endif + const vfloat32m2_t n = __riscv_vfsub_vv_f32m2(z, r, vl); + const vfloat32m2_t b = __riscv_vfnmsac_vf_f32m2(__riscv_vfnmsac_vf_f32m2(x, 0x1.62e4p-1f, n, vl), + 0x1.7f7d1cp-20f, n, vl); + const vuint32m2_t e = __riscv_vsll_vx_u32m2(__riscv_vreinterpret_v_f32m2_u32m2(z), 23, vl); + const vfloat32m2_t k = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vadd_vx_u32m2(e, 0x3f800000, vl)); // 1.0f + const vbool16_t c = __riscv_vmfgt_vf_f32m2_b16(__riscv_vfabs_v_f32m2(n, vl), 126.0f, vl); + const vfloat32m2_t u = __riscv_vfmul_vv_f32m2(b, b, vl); + const vfloat32m2_t j = __riscv_vfmacc_vv_f32m2( + __riscv_vfmul_vf_f32m2(b, 0x1.ffffecp-1f, vl), + __riscv_vfmacc_vv_f32m2( + __riscv_vfmacc_vf_f32m2(__riscv_vfmv_v_f_f32m2(0x1.fffdb6p-2f, vl), 0x1.555e66p-3f, b, vl), + __riscv_vfmacc_vf_f32m2(__riscv_vfmv_v_f_f32m2(0x1.573e2ep-5f, vl), 0x1.0e4020p-7f, b, vl), + u, vl), u, vl); + if (!__riscv_vcpop_m_b16(c, vl)) + return __riscv_vfmacc_vv_f32m2(k, j, k, vl); + const vbool16_t dm = __riscv_vmfle_vf_f32m2_b16(n, 0.0f, vl); + const vuint32m2_t d = __riscv_vmerge_vxm_u32m2(__riscv_vmv_v_x_u32m2(0, vl), 0x82000000, dm, vl); + const vfloat32m2_t s1 = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vadd_vx_u32m2(d, 0x7f000000, vl)); + const vfloat32m2_t s2 = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vsub_vv_u32m2(e, d, vl)); + const vfloat32m2_t r1 = __riscv_vmerge_vvm_f32m2( + __riscv_vfmacc_vv_f32m2(k, k, j, vl), + __riscv_vfmul_vv_f32m2(__riscv_vfmacc_vv_f32m2(s2, s2, j, vl), s1, vl), + c, vl); + return __riscv_vmerge_vvm_f32m2( + r1, __riscv_vfmul_vv_f32m2(s1, s1, vl), + __riscv_vmfgt_vf_f32m2_b16(__riscv_vfabs_v_f32m2(n, vl), 192.0f, vl), + vl); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static vfloat32m2_t ggml_v_silu_m2(vfloat32m2_t x, int vl) { + const vfloat32m2_t neg_x = __riscv_vfneg_v_f32m2(x, vl); + const vfloat32m2_t exp_neg_x = ggml_v_expf_m2(neg_x, vl); + const vfloat32m2_t one_plus_exp_neg_x = __riscv_vfadd_vf_f32m2(exp_neg_x, 1.0f, vl); + return __riscv_vfdiv_vv_f32m2(x, one_plus_exp_neg_x, vl); +} + +#endif // __ARM_NEON / __AVX2__ / __SSE2__ / __riscv_v_intrinsic + +inline static void ggml_vec_silu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_silu_f16(x[i]); + } +} + +inline static float ggml_silu_backward_f32(float x, float dy) { + const float s = 1.0f/(1.0f + expf(-x)); + return dy*s*(1.0f + x*(1.0f - s)); +} + +inline static ggml_fp16_t ggml_silu_backward_f16(ggml_fp16_t x, ggml_fp16_t dy) { + const float v = GGML_CPU_FP16_TO_FP32(x); + const float s = 1.0f/(1.0f + expf(-v)); + return GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(dy)*s*(1.0f + v*(1.0f - s))); +} + +inline static void ggml_vec_silu_backward_f32(const int n, float * dx, const float * x, const float * dy) { + for (int i = 0; i < n; ++i) { + dx[i] = ggml_silu_backward_f32(x[i], dy[i]); + } +} + +inline static void ggml_vec_silu_backward_f16(const int n, ggml_fp16_t * dx, const ggml_fp16_t * x, const ggml_fp16_t * dy) { + for (int i = 0; i < n; ++i) { + dx[i] = ggml_silu_backward_f16(x[i], dy[i]); + } +} + +inline static void ggml_vec_reglu_f32 (const int n, float * y, const float * x, const float * g) { + for (int i = 0; i < n; ++i) { + y[i] = (x[i] > 0.f) ? x[i] * g[i] : 0.f; + } +} + +inline static void ggml_vec_reglu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v * GGML_CPU_FP16_TO_FP32(g[i]) : 0.f); + } +} + +#ifdef GGML_GELU_FP16 +inline static void ggml_vec_geglu_f32(const int n, float * y, const float * x, const float * g) { + uint16_t t; + for (int i = 0; i < n; ++i) { + if (x[i] <= -10.0f) { + y[i] = 0.0f; + } else if (x[i] >= 10.0f) { + y[i] = x[i] * g[i]; + } else { + ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); + memcpy(&t, &fp16, sizeof(uint16_t)); + y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[t]) * g[i]; + } + } +} +#else +inline static void ggml_vec_geglu_f32(const int n, float * y, const float * x, const float * g) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_gelu_f32(x[i]) * g[i]; + } +} +#endif + +inline static void ggml_vec_geglu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + const uint16_t * i16 = (const uint16_t *) x; + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(g[i]); + y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[i16[i]]) * v); + } +} + +void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float * g); + +inline static void ggml_vec_swiglu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + for (int i = 0; i < n; ++i) { + float xi = GGML_CPU_FP16_TO_FP32(x[i]); + float gi = GGML_CPU_FP16_TO_FP32(g[i]); + y[i] = GGML_CPU_FP32_TO_FP16((xi/(1.0f + expf(-xi))) * gi); + } +} + +inline static void ggml_vec_geglu_erf_f32(const int n, float * y, const float * x, const float * g) { + for (int i = 0; i < n; ++i) { + float xi = x[i]; + y[i] = 0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * g[i]; + } +} + +inline static void ggml_vec_geglu_erf_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + for (int i = 0; i < n; ++i) { + float xi = GGML_CPU_FP16_TO_FP32(x[i]); + float gi = GGML_CPU_FP16_TO_FP32(g[i]); + y[i] = GGML_CPU_FP32_TO_FP16(0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * gi); + } +} + +#ifdef GGML_GELU_QUICK_FP16 +inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) { + uint16_t t; + for (int i = 0; i < n; ++i) { + ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); + memcpy(&t, &fp16, sizeof(uint16_t)); + y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[t]) * g[i]; + } +} +#else +inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_gelu_quick_f32(x[i]) * g[i]; + } +} +#endif + +inline static void ggml_vec_geglu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + const uint16_t * i16 = (const uint16_t *) x; + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(g[i]); + y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[i16[i]]) * v); + } +} + +inline static void ggml_vec_sum_f32(const int n, float * s, const float * x) { +#ifndef GGML_USE_ACCELERATE + ggml_float sum = 0.0; + for (int i = 0; i < n; ++i) { + sum += (ggml_float)x[i]; + } + *s = (float)sum; +#else + vDSP_sve(x, 1, s, n); +#endif +} + +inline static void ggml_vec_cumsum_f32(const int n, float * y, const float * x) { + for (int i = 0; i < n; ++i) { + if (i == 0) { + y[i] = x[i]; + } else { + y[i] = y[i - 1] + x[i]; + } + } +} + +inline static void ggml_vec_sum_f32_ggf(const int n, ggml_float * s, const float * x) { + ggml_float sum = 0.0; + for (int i = 0; i < n; ++i) { + sum += (ggml_float)x[i]; + } + *s = sum; +} + +inline static void ggml_vec_sum_f16_ggf(const int n, float * s, const ggml_fp16_t * x) { + float sum = 0.0f; + for (int i = 0; i < n; ++i) { + sum += GGML_CPU_FP16_TO_FP32(x[i]); + } + *s = sum; +} + +inline static void ggml_vec_sum_bf16_ggf(const int n, float * s, const ggml_bf16_t * x) { + float sum = 0.0f; + for (int i = 0; i < n; ++i) { + sum += GGML_BF16_TO_FP32(x[i]); + } + *s = sum; +} + +inline static void ggml_vec_max_f32(const int n, float * s, const float * x) { +#ifndef GGML_USE_ACCELERATE + float max = -INFINITY; + for (int i = 0; i < n; ++i) { + max = MAX(max, x[i]); + } + *s = max; +#else + vDSP_maxv(x, 1, s, n); +#endif +} + +inline static void ggml_vec_norm_inv_f32(const int n, float * s, const float * x) { + ggml_vec_norm_f32(n, s, x); + *s = 1.f/(*s); +} + +inline static void ggml_vec_argmax_f32(const int n, int * s, const float * x) { + float max = -INFINITY; + int idx = 0; + for (int i = 0; i < n; ++i) { + max = MAX(max, x[i]); + if (max == x[i]) { idx = i; } + } + *s = idx; +} + +#ifdef __cplusplus +} +#endif + +#include "ggml.h.inc" + +// +// cache line +// + +#if defined(__cpp_lib_hardware_interference_size) +#define CACHE_LINE_SIZE std::hardware_destructive_interference_size +#else +#if defined(__POWER9_VECTOR__) +#define CACHE_LINE_SIZE 128 +#elif defined(__VXE__) || defined(__VXE2__) +#define CACHE_LINE_SIZE 256 +#else +#define CACHE_LINE_SIZE 64 +#endif +#endif + +// Work buffer size for im2col operations in CONV2D +#define GGML_IM2COL_WORK_SIZE (16 * 1024 * 1024) + +#ifdef __cplusplus +extern "C" { +#endif + +void ggml_compute_forward_dup(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_add(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_add_id(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_add1(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_acc(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sum(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sum_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_cumsum(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_mean(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_argmax(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_count_equal(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_repeat(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_repeat_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_concat(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_silu_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rms_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rms_norm_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_group_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_l2_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_out_prod(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_scale(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_set(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_cpy(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_cont(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_get_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_get_rows_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_set_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_diag(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_diag_mask_inf(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_diag_mask_zero(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_soft_max(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_soft_max_ext_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rope(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rope_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_clamp(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_conv_transpose_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_im2col(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_im2col_back_f32(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_im2col_3d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_conv_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_conv_3d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_conv_transpose_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_conv_2d_dw(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_pool_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_pool_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_pool_2d_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_upscale(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_pad(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_pad_reflect_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_roll(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_arange(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_timestep_embedding(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_argsort(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_top_k(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_leaky_relu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_fill(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_flash_attn_ext(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_flash_attn_back( + const struct ggml_compute_params * params, + const bool masked, + struct ggml_tensor * dst); +void ggml_compute_forward_ssm_conv(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_ssm_scan(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_win_part(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_win_unpart(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_unary(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_glu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_get_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_add_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rwkv_wkv6(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rwkv_wkv7(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_custom(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_cross_entropy_loss(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_cross_entropy_loss_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_opt_step_adamw(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_opt_step_sgd(const struct ggml_compute_params * params, struct ggml_tensor * dst); +#ifdef __cplusplus +} +#endif + +#include "ggml.h.inc" #if defined(_MSC_VER) || defined(__MINGW32__) #include // using malloc.h with MSC/MINGW @@ -47,9 +1970,42 @@ #endif #ifdef GGML_USE_LLAMAFILE -#include "llamafile/sgemm.h" +bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64_t n, int64_t k, + const void * A, int64_t lda, const void * B, int64_t ldb, void * C, int64_t ldc, + int ith, int nth, int type); #endif +void ggml_compute_forward_add_non_quantized(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sub(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_mul(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_div(const struct ggml_compute_params * params, struct ggml_tensor * dst); + +void ggml_compute_forward_abs(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sgn(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_neg(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_step(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_tanh(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_elu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_relu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sigmoid(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_hardsigmoid(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_exp(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_hardswish(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sqr(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sqrt(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sin(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_cos(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_log(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_expm1(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_softplus(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_floor(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_ceil(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_round(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_trunc(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_xielu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +GGML_API void ggml_critical_section_start(void); +GGML_API void ggml_critical_section_end(void); + // Note: once we move threading into a separate C++ file // will use std::hardware_destructive_interference_size instead of hardcoding it here // and we'll use C++ attribute syntax. @@ -2898,7 +4854,7 @@ struct ggml_cplan ggml_graph_plan( const int64_t DK = node->src[1]->ne[0]; const int64_t DV = node->src[2]->ne[0]; - // Tiled flash attention scratch (tile sizes defined in common.h) + // Tiled flash attention scratch (tile sizes defined in common.h.inc) // Per-thread: Q_q + KQ + mask + VKQ32 + V32 + K_f32 + padding size_t prefill = sizeof(float)*(GGML_FA_TILE_Q*DK + 2*GGML_FA_TILE_Q*GGML_FA_TILE_KV + GGML_FA_TILE_Q*DV + GGML_FA_TILE_KV*DV + GGML_FA_TILE_KV*DK)*n_tasks; diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp deleted file mode 100644 index 128883b41ce7..000000000000 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ /dev/null @@ -1,703 +0,0 @@ -#include "ggml-backend.h" -#include "ggml-backend-impl.h" -#include "ggml-cpu.h" -#include "repack.h" -#include "traits.h" -#include "ggml-impl.h" -#include "amx/amx.h" - -#include -#include -#include - -#ifdef GGML_USE_CPU_HBM -# include "hbm.h" -#endif - -#ifdef GGML_USE_CPU_KLEIDIAI -# include "kleidiai/kleidiai.h" -#endif - -#ifdef GGML_USE_CPU_RISCV64_SPACEMIT -# include "spacemit/ime.h" -#endif - -#if defined(_WIN32) -# define WIN32_LEAN_AND_MEAN -# ifndef NOMINMAX -# define NOMINMAX -# endif -# include -#else -# include -#endif - -#if defined(__APPLE__) -# include -# include -#endif - -// ggml-backend interface - -std::vector & ggml_backend_cpu_get_extra_buffer_types() { - static std::vector bufts = []() { - std::vector bufts; - -#if defined(__AMX_INT8__) && defined(__AVX512VNNI__) - if (ggml_backend_amx_buffer_type()) { - bufts.push_back(ggml_backend_amx_buffer_type()); - } -#endif - -#ifdef GGML_USE_CPU_RISCV64_SPACEMIT - if (ggml_backend_cpu_riscv64_spacemit_buffer_type()) { - bufts.push_back(ggml_backend_cpu_riscv64_spacemit_buffer_type()); - } -#endif - -#ifdef GGML_USE_CPU_KLEIDIAI - if (ggml_backend_cpu_kleidiai_buffer_type()) { - bufts.push_back(ggml_backend_cpu_kleidiai_buffer_type()); - } -#endif - -#ifdef GGML_USE_CPU_REPACK - if (ggml_backend_cpu_repack_buffer_type()) { - bufts.push_back(ggml_backend_cpu_repack_buffer_type()); - } -#endif - - return bufts; - }(); - - return bufts; -} - -static ggml_backend_buffer_type_t * ggml_backend_cpu_device_get_extra_buffers_type(ggml_backend_dev_t device) { - static std::vector extra_bufts = [] { - std::vector bufts = ggml_backend_cpu_get_extra_buffer_types(); - bufts.push_back(nullptr); - return bufts; - }(); - - return extra_bufts.data(); - - GGML_UNUSED(device); -} - -static bool ggml_backend_cpu_is_extra_buffer_type(ggml_backend_buffer_type_t buft) { - for (auto * extra : ggml_backend_cpu_get_extra_buffer_types()) { - if (extra == buft) { - return true; - } - } - return false; -} - -// CPU backend - backend (stream) - -struct ggml_backend_cpu_context { - int n_threads; - ggml_threadpool_t threadpool; - - uint8_t * work_data; - size_t work_size; - - ggml_abort_callback abort_callback; - void * abort_callback_data; - - bool use_ref; // use reference implementation -}; - -static const char * ggml_backend_cpu_get_name(ggml_backend_t backend) { - return "CPU"; - - GGML_UNUSED(backend); -} - -static void ggml_backend_cpu_free(ggml_backend_t backend) { - struct ggml_backend_cpu_context * cpu_ctx = (struct ggml_backend_cpu_context *)backend->context; - delete[] cpu_ctx->work_data; - delete cpu_ctx; - delete backend; -} - -struct ggml_backend_plan_cpu { - struct ggml_cplan cplan; - struct ggml_cgraph cgraph; -}; - -static ggml_backend_graph_plan_t ggml_backend_cpu_graph_plan_create(ggml_backend_t backend, const struct ggml_cgraph * cgraph) { - struct ggml_backend_cpu_context * cpu_ctx = (struct ggml_backend_cpu_context *)backend->context; - - struct ggml_backend_plan_cpu * cpu_plan = new ggml_backend_plan_cpu; - - cpu_plan->cplan = ggml_graph_plan(cgraph, cpu_ctx->n_threads, cpu_ctx->threadpool); - cpu_plan->cgraph = *cgraph; // FIXME: deep copy - - if (cpu_plan->cplan.work_size > 0) { - cpu_plan->cplan.work_data = new uint8_t[cpu_plan->cplan.work_size]; - if (cpu_plan->cplan.work_data == NULL) { - delete cpu_plan; - return NULL; - } - } - - cpu_plan->cplan.abort_callback = cpu_ctx->abort_callback; - cpu_plan->cplan.abort_callback_data = cpu_ctx->abort_callback_data; - cpu_plan->cplan.use_ref = cpu_ctx->use_ref; - - return cpu_plan; -} - -static void ggml_backend_cpu_graph_plan_free(ggml_backend_t backend, ggml_backend_graph_plan_t plan) { - struct ggml_backend_plan_cpu * cpu_plan = (struct ggml_backend_plan_cpu *)plan; - - delete[] cpu_plan->cplan.work_data; - delete cpu_plan; - - GGML_UNUSED(backend); -} - -static enum ggml_status ggml_backend_cpu_graph_plan_compute(ggml_backend_t backend, ggml_backend_graph_plan_t plan) { - struct ggml_backend_plan_cpu * cpu_plan = (struct ggml_backend_plan_cpu *)plan; - - return ggml_graph_compute(&cpu_plan->cgraph, &cpu_plan->cplan); - - GGML_UNUSED(backend); -} - -static enum ggml_status ggml_backend_cpu_graph_compute(ggml_backend_t backend, struct ggml_cgraph * cgraph) { - struct ggml_backend_cpu_context * cpu_ctx = (struct ggml_backend_cpu_context *)backend->context; - - struct ggml_cplan cplan = ggml_graph_plan(cgraph, cpu_ctx->n_threads, cpu_ctx->threadpool); - - if (cpu_ctx->work_size < cplan.work_size) { - delete[] cpu_ctx->work_data; - cpu_ctx->work_data = new uint8_t[cplan.work_size]; - if (cpu_ctx->work_data == NULL) { - cpu_ctx->work_size = 0; - return GGML_STATUS_ALLOC_FAILED; - } - cpu_ctx->work_size = cplan.work_size; - } - cplan.work_data = (uint8_t *)cpu_ctx->work_data; - - cplan.abort_callback = cpu_ctx->abort_callback; - cplan.abort_callback_data = cpu_ctx->abort_callback_data; - cplan.use_ref = cpu_ctx->use_ref; - - return ggml_graph_compute(cgraph, &cplan); -} - -static const struct ggml_backend_i ggml_backend_cpu_i = { - /* .get_name = */ ggml_backend_cpu_get_name, - /* .free = */ ggml_backend_cpu_free, - /* .set_tensor_async = */ NULL, - /* .get_tensor_async = */ NULL, - /* .set_tensor_2d_async = */ NULL, - /* .get_tensor_2d_async = */ NULL, - /* .cpy_tensor_async = */ NULL, - /* .synchronize = */ NULL, - /* .graph_plan_create = */ ggml_backend_cpu_graph_plan_create, - /* .graph_plan_free = */ ggml_backend_cpu_graph_plan_free, - /* .graph_plan_update = */ NULL, - /* .graph_plan_compute = */ ggml_backend_cpu_graph_plan_compute, - /* .graph_compute = */ ggml_backend_cpu_graph_compute, - /* .event_record = */ NULL, - /* .event_wait = */ NULL, - /* .graph_optimize = */ NULL, -}; - -static ggml_guid_t ggml_backend_cpu_guid(void) { - static ggml_guid guid = { 0xaa, 0x67, 0xc7, 0x43, 0x96, 0xe6, 0xa3, 0x8a, 0xe3, 0xaf, 0xea, 0x92, 0x36, 0xbc, 0xfc, 0x89 }; - return &guid; -} - -ggml_backend_t ggml_backend_cpu_init(void) { - // initialize CPU backend now to avoid slowing the first graph computation - ggml_cpu_init(); - - struct ggml_backend_cpu_context * ctx = new ggml_backend_cpu_context; - if (ctx == NULL) { - return NULL; - } - - ctx->n_threads = GGML_DEFAULT_N_THREADS; - ctx->threadpool = NULL; - ctx->work_data = NULL; - ctx->work_size = 0; - ctx->abort_callback = NULL; - ctx->abort_callback_data = NULL; - ctx->use_ref = false; - - ggml_backend_t cpu_backend = new ggml_backend { - /* .guid = */ ggml_backend_cpu_guid(), - /* .iface = */ ggml_backend_cpu_i, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0), - /* .context = */ ctx, - }; - - if (cpu_backend == NULL) { - delete ctx; - return NULL; - } - - return cpu_backend; -} - -bool ggml_backend_is_cpu(ggml_backend_t backend) { - return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_cpu_guid()); -} - -void ggml_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads) { - GGML_ASSERT(ggml_backend_is_cpu(backend_cpu)); - - struct ggml_backend_cpu_context * ctx = (struct ggml_backend_cpu_context *)backend_cpu->context; - ctx->n_threads = n_threads; -} - -void ggml_backend_cpu_set_threadpool(ggml_backend_t backend_cpu, ggml_threadpool_t threadpool) { - GGML_ASSERT(ggml_backend_is_cpu(backend_cpu)); - - struct ggml_backend_cpu_context * ctx = (struct ggml_backend_cpu_context *)backend_cpu->context; - - if (ctx->threadpool && ctx->threadpool != threadpool) { - // already had a different threadpool, pause/suspend it before switching - ggml_threadpool_pause(ctx->threadpool); - } - ctx->threadpool = threadpool; -} - -void ggml_backend_cpu_set_abort_callback(ggml_backend_t backend_cpu, ggml_abort_callback abort_callback, void * abort_callback_data) { - GGML_ASSERT(ggml_backend_is_cpu(backend_cpu)); - - struct ggml_backend_cpu_context * ctx = (struct ggml_backend_cpu_context *)backend_cpu->context; - ctx->abort_callback = abort_callback; - ctx->abort_callback_data = abort_callback_data; -} - -void ggml_backend_cpu_set_use_ref(ggml_backend_t backend_cpu, bool use_ref) { - GGML_ASSERT(ggml_backend_is_cpu(backend_cpu)); - - struct ggml_backend_cpu_context * ctx = (struct ggml_backend_cpu_context *)backend_cpu->context; - ctx->use_ref = use_ref; -} - -// CPU backend - device - -struct ggml_backend_cpu_device_context { - std::string description = "CPU"; - - ggml_backend_cpu_device_context() { -#ifdef __APPLE__ - size_t len = 0; - if (!sysctlbyname("machdep.cpu.brand_string", NULL, &len, NULL, 0)) { - description.resize(len); - sysctlbyname("machdep.cpu.brand_string", &description[0], &len, NULL, 0); // NOLINT - } -#elif defined(__linux__) - FILE * f = fopen("/proc/cpuinfo", "r"); - if (f) { - char buf[1024]; - while (fgets(buf, sizeof(buf), f)) { - if (strncmp(buf, "model name", 10) == 0) { - char * p = strchr(buf, ':'); - if (p) { - p++; - while (std::isspace(*p)) { - p++; - } - while (std::isspace(p[strlen(p) - 1])) { - p[strlen(p) - 1] = '\0'; - } - description = p; - break; - } - } - } - fclose(f); - } -#elif defined(_WIN32) - HKEY hKey; - if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, - TEXT("HARDWARE\\DESCRIPTION\\System\\CentralProcessor\\0"), - 0, - KEY_READ, - &hKey) == ERROR_SUCCESS) { - DWORD cpu_brand_size = 0; - if (RegQueryValueExA(hKey, - "ProcessorNameString", - NULL, - NULL, - NULL, - &cpu_brand_size) == ERROR_SUCCESS) { - description.resize(cpu_brand_size); - if (RegQueryValueExA(hKey, - "ProcessorNameString", - NULL, - NULL, - (LPBYTE)&description[0], // NOLINT - &cpu_brand_size) == ERROR_SUCCESS) { - if (description.find('\0') != std::string::npos) { - description.resize(description.find('\0')); - } - } - } - RegCloseKey(hKey); - } -#endif - } -}; - -static const char * ggml_backend_cpu_device_get_name(ggml_backend_dev_t dev) { - return "CPU"; - - GGML_UNUSED(dev); -} - -static const char * ggml_backend_cpu_device_get_description(ggml_backend_dev_t dev) { - struct ggml_backend_cpu_device_context * ctx = (struct ggml_backend_cpu_device_context *)dev->context; - - return ctx->description.c_str(); -} - -static void ggml_backend_cpu_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { -#ifdef _WIN32 - MEMORYSTATUSEX status; - status.dwLength = sizeof(status); - GlobalMemoryStatusEx(&status); - *total = status.ullTotalPhys; - *free = status.ullAvailPhys; -#else - long pages = sysconf(_SC_PHYS_PAGES); - long page_size = sysconf(_SC_PAGE_SIZE); - *total = pages * page_size; - - // "free" system memory is ill-defined, for practical purposes assume that all of it is free: - *free = *total; -#endif // _WIN32 - - GGML_UNUSED(dev); -} - -static enum ggml_backend_dev_type ggml_backend_cpu_device_get_type(ggml_backend_dev_t dev) { - return GGML_BACKEND_DEVICE_TYPE_CPU; - - GGML_UNUSED(dev); -} - -static void ggml_backend_cpu_device_get_props(ggml_backend_dev_t dev, struct ggml_backend_dev_props * props) { - props->name = ggml_backend_cpu_device_get_name(dev); - props->description = ggml_backend_cpu_device_get_description(dev); - props->type = ggml_backend_cpu_device_get_type(dev); - ggml_backend_cpu_device_get_memory(dev, &props->memory_free, &props->memory_total); - props->caps = { - /* .async = */ false, - /* .host_buffer = */ false, - /* .buffer_from_host_ptr = */ true, - /* .events = */ false, - }; -} - -static ggml_backend_t ggml_backend_cpu_device_init_backend(ggml_backend_dev_t dev, const char * params) { - return ggml_backend_cpu_init(); - - GGML_UNUSED(dev); - GGML_UNUSED(params); -} - -static ggml_backend_buffer_type_t ggml_backend_cpu_device_get_buffer_type(ggml_backend_dev_t dev) { - return ggml_backend_cpu_buffer_type(); - - GGML_UNUSED(dev); -} - -static ggml_backend_buffer_t ggml_backend_cpu_device_buffer_from_host_ptr(ggml_backend_dev_t dev, void * ptr, size_t size, size_t max_tensor_size) { - return ggml_backend_cpu_buffer_from_ptr(ptr, size); - - GGML_UNUSED(dev); - GGML_UNUSED(max_tensor_size); -} - -static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) { - const struct ggml_tensor * src0 = op->src[0]; - const struct ggml_tensor * src1 = op->src[1]; - - if (op->op == GGML_OP_NONE || op->op == GGML_OP_RESHAPE || op->op == GGML_OP_VIEW || op->op == GGML_OP_PERMUTE || op->op == GGML_OP_TRANSPOSE) { - return true; - } - - // check extra buffer types - // note: only the first sources are checked for extra buffer types to reduce overhead, increase if necessary - for (int i = 0; i < 4; i++) { - if (op->src[i] && op->src[i]->buffer && - ggml_backend_cpu_is_extra_buffer_type(op->src[i]->buffer->buft)) { - auto * buf_extra = (ggml::cpu::extra_buffer_type *) op->src[i]->buffer->buft->context; - return buf_extra->supports_op(dev, op); - } - } - - switch (op->op) { - case GGML_OP_CPY: - case GGML_OP_SET_ROWS: - return - op->type != GGML_TYPE_IQ3_XXS && - op->type != GGML_TYPE_IQ3_S && - op->type != GGML_TYPE_IQ2_XXS && - op->type != GGML_TYPE_IQ2_XS && - op->type != GGML_TYPE_IQ2_S && - op->type != GGML_TYPE_IQ1_S && - op->type != GGML_TYPE_IQ1_M; // missing type_traits.from_float - case GGML_OP_MUL_MAT: - return src1->type == GGML_TYPE_F32 || src1->type == ggml_get_type_traits_cpu(src0->type)->vec_dot_type; - case GGML_OP_SOFT_MAX_BACK: { - if (op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32) { - return false; - } - float max_bias = 0.0f; - - memcpy(&max_bias, (const float *) op->op_params + 1, sizeof(float)); - - return max_bias == 0.0f; - } - case GGML_OP_IM2COL_BACK: - return src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32; - case GGML_OP_GET_ROWS_BACK: - return src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16; - case GGML_OP_OUT_PROD: - return (src0->type == GGML_TYPE_F32 || (ggml_is_quantized(src0->type) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) && - src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; - default: - return true; - } -} - -static bool ggml_backend_cpu_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { - return ggml_backend_buft_is_host(buft) || ggml_backend_cpu_is_extra_buffer_type(buft); - GGML_UNUSED(dev); -} - -static const struct ggml_backend_device_i ggml_backend_cpu_device_i = { - /* .get_name = */ ggml_backend_cpu_device_get_name, - /* .get_description = */ ggml_backend_cpu_device_get_description, - /* .get_memory = */ ggml_backend_cpu_device_get_memory, - /* .get_type = */ ggml_backend_cpu_device_get_type, - /* .get_props = */ ggml_backend_cpu_device_get_props, - /* .init_backend = */ ggml_backend_cpu_device_init_backend, - /* .get_buffer_type = */ ggml_backend_cpu_device_get_buffer_type, - /* .get_host_buffer_type = */ NULL, - /* .buffer_from_host_ptr = */ ggml_backend_cpu_device_buffer_from_host_ptr, - /* .supports_op = */ ggml_backend_cpu_device_supports_op, - /* .supports_buft = */ ggml_backend_cpu_device_supports_buft, - /* .offload_op = */ NULL, - /* .event_new = */ NULL, - /* .event_free = */ NULL, - /* .event_synchronize = */ NULL, -}; - -// CPU backend - backend (reg) - -static const char * ggml_backend_cpu_reg_get_name(ggml_backend_reg_t reg) { - return "CPU"; - - GGML_UNUSED(reg); -} - -static size_t ggml_backend_cpu_reg_get_device_count(ggml_backend_reg_t reg) { - return 1; - - GGML_UNUSED(reg); -} - -static ggml_backend_dev_t ggml_backend_cpu_reg_get_device(ggml_backend_reg_t reg, size_t index) { - GGML_ASSERT(index == 0); - - static ggml_backend_cpu_device_context ctx; - static ggml_backend_device ggml_backend_cpu_device = { - /* .iface = */ ggml_backend_cpu_device_i, - /* .reg = */ reg, - /* .context = */ &ctx, - }; - - return &ggml_backend_cpu_device; -} - -// This is intended to replace the the ggml_cpu_has_* functions when loading the CPU backend dynamically, -// and additionally to allow other backends to expose their own list of features that applications can query using the same API -static ggml_backend_feature * ggml_backend_cpu_get_features(ggml_backend_reg_t reg) { - static std::vector features = []() { - ggml_cpu_init(); - - std::vector features; - if (ggml_cpu_has_sse3()) { - features.push_back({ "SSE3", "1" }); - } - if (ggml_cpu_has_ssse3()) { - features.push_back({ "SSSE3", "1" }); - } - if (ggml_cpu_has_avx()) { - features.push_back({ "AVX", "1" }); - } - if (ggml_cpu_has_avx_vnni()) { - features.push_back({ "AVX_VNNI", "1" }); - } - if (ggml_cpu_has_avx2()) { - features.push_back({ "AVX2", "1" }); - } - if (ggml_cpu_has_f16c()) { - features.push_back({ "F16C", "1" }); - } - if (ggml_cpu_has_fma()) { - features.push_back({ "FMA", "1" }); - } - if (ggml_cpu_has_bmi2()) { - features.push_back({ "BMI2", "1" }); - } - if (ggml_cpu_has_avx512()) { - features.push_back({ "AVX512", "1" }); - } - if (ggml_cpu_has_avx512_vbmi()) { - features.push_back({ "AVX512_VBMI", "1" }); - } - if (ggml_cpu_has_avx512_vnni()) { - features.push_back({ "AVX512_VNNI", "1" }); - } - if (ggml_cpu_has_avx512_bf16()) { - features.push_back({ "AVX512_BF16", "1" }); - } - if (ggml_cpu_has_amx_int8()) { - features.push_back({ "AMX_INT8", "1" }); - } - if (ggml_cpu_has_neon()) { - features.push_back({ "NEON", "1" }); - } - if (ggml_cpu_has_arm_fma()) { - features.push_back({ "ARM_FMA", "1" }); - } - if (ggml_cpu_has_fp16_va()) { - features.push_back({ "FP16_VA", "1" }); - } - if (ggml_cpu_has_matmul_int8()) { - features.push_back({ "MATMUL_INT8", "1" }); - } - if (ggml_cpu_has_sve()) { - features.push_back({ "SVE", "1" }); - } - if (ggml_cpu_has_dotprod()) { - features.push_back({ "DOTPROD", "1" }); - } - if (ggml_cpu_get_sve_cnt() > 0) { - static std::string sve_cnt = std::to_string(ggml_cpu_get_sve_cnt()); - features.push_back({ "SVE_CNT", sve_cnt.c_str() }); - } - if (ggml_cpu_has_sme()) { - features.push_back({ "SME", "1" }); - } - if (ggml_cpu_has_riscv_v()) { - features.push_back({ "RISCV_V", "1" }); - } - if (ggml_cpu_get_rvv_vlen() > 0) { - static std::string rvv_vlen = std::to_string(ggml_cpu_get_rvv_vlen()); - features.push_back({ "RVV_VLEN", rvv_vlen.c_str() }); - } - if (ggml_cpu_has_vsx()) { - features.push_back({ "VSX", "1" }); - } - if (ggml_cpu_has_vxe()) { - features.push_back({ "VXE", "1" }); - } - if (ggml_cpu_has_wasm_simd()) { - features.push_back({ "WASM_SIMD", "1" }); - } - if (ggml_cpu_has_llamafile()) { - features.push_back({ "LLAMAFILE", "1" }); - } - #ifdef GGML_USE_ACCELERATE - features.push_back({ "ACCELERATE", "1" }); - #endif - #ifdef GGML_USE_CPU_HBM - features.push_back({ "CPU_HBM", "1" }); - #endif - #ifdef GGML_USE_OPENMP - features.push_back({ "OPENMP", "1" }); - #endif - #ifdef GGML_USE_CPU_KLEIDIAI - features.push_back({ "KLEIDIAI", "1" }); - #endif - #ifdef GGML_USE_CPU_REPACK - features.push_back({ "REPACK", "1" }); - #endif - - features.push_back({ nullptr, nullptr }); - - return features; - }(); - - return features.data(); - - GGML_UNUSED(reg); -} - -static void * ggml_backend_cpu_get_proc_address(ggml_backend_reg_t reg, const char * name) { - if (strcmp(name, "ggml_backend_set_n_threads") == 0) { - ggml_backend_set_n_threads_t fct = ggml_backend_cpu_set_n_threads; - return (void *)fct; - } - if (strcmp(name, "ggml_backend_dev_get_extra_bufts") == 0) { - ggml_backend_dev_get_extra_bufts_t fct = ggml_backend_cpu_device_get_extra_buffers_type; - return (void *)fct; - } - if (strcmp(name, "ggml_backend_get_features") == 0) { - return (void *)ggml_backend_cpu_get_features; - } - if (strcmp(name, "ggml_backend_set_abort_callback") == 0) { - return (void *)ggml_backend_cpu_set_abort_callback; - } - if (strcmp(name, "ggml_backend_cpu_numa_init") == 0) { - return (void *)ggml_numa_init; - } - if (strcmp(name, "ggml_backend_cpu_is_numa") == 0) { - return (void *)ggml_is_numa; - } - if (strcmp(name, "ggml_backend_cpu_set_use_ref") == 0) { - return (void *)ggml_backend_cpu_set_use_ref; - } - - // threadpool - TODO: move to ggml-base - if (strcmp(name, "ggml_threadpool_new") == 0) { - return (void *)ggml_threadpool_new; - } - if (strcmp(name, "ggml_threadpool_free") == 0) { - return (void *)ggml_threadpool_free; - } - if (strcmp(name, "ggml_backend_cpu_set_threadpool") == 0) { - return (void *)ggml_backend_cpu_set_threadpool; - } - - return NULL; - - GGML_UNUSED(reg); -} - -static const struct ggml_backend_reg_i ggml_backend_cpu_reg_i = { - /* .get_name = */ ggml_backend_cpu_reg_get_name, - /* .get_device_count = */ ggml_backend_cpu_reg_get_device_count, - /* .get_device = */ ggml_backend_cpu_reg_get_device, - /* .get_proc_address = */ ggml_backend_cpu_get_proc_address, -}; - -ggml_backend_reg_t ggml_backend_cpu_reg(void) { - // init CPU feature detection - ggml_cpu_init(); - - static struct ggml_backend_reg ggml_backend_cpu_reg = { - /* .api_version = */ GGML_BACKEND_API_VERSION, - /* .iface = */ ggml_backend_cpu_reg_i, - /* .context = */ NULL, - }; - - return &ggml_backend_cpu_reg; -} - -GGML_BACKEND_DL_IMPL(ggml_backend_cpu_reg) diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp.inc b/ggml/src/ggml-cpu/ggml-cpu.cpp.inc new file mode 100644 index 000000000000..31d614e9aa45 --- /dev/null +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp.inc @@ -0,0 +1,8199 @@ +#include "ggml-backend.h.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml-cpu.h.inc" +#define GGML_COMMON_DECL_CPP +#include "ggml-common-defs.inc" + +#include "ggml.h.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml-cpu-impl-defs.inc" +#include "ggml-impl-defs.inc" +#include "simd-mappings-defs.inc" + +#ifdef __cplusplus +#include +extern "C" { +#endif + +// return true if op part of extra "accelerator" +bool ggml_cpu_extra_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op); +bool ggml_cpu_extra_work_size(int n_threads, const struct ggml_tensor * op, size_t * size); + +#ifdef __cplusplus +} + +namespace ggml::cpu { +// register in tensor->extra +class tensor_traits { + public: + virtual ~tensor_traits(); + virtual bool work_size(int n_threads, const struct ggml_tensor * op, size_t & size) = 0; + virtual bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) = 0; +}; + +class extra_buffer_type { + public: + virtual ~extra_buffer_type(); + virtual bool supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) = 0; + virtual tensor_traits * get_tensor_traits(const struct ggml_tensor * op) = 0; +}; +} // namespace ggml::cpu + +// implemented in ggml-cpu.cpp. +std::vector & ggml_backend_cpu_get_extra_buffer_types(); + +#endif + +#define GGML_FA_TILE_Q 64 +#define GGML_FA_TILE_KV 64 + +#ifdef __cplusplus + +#include + +// convenience functions/macros for use in template calls +// note: these won't be required after the 'traits' lookup table is used. +static inline ggml_fp16_t f32_to_f16(float x) { + return GGML_CPU_FP32_TO_FP16(x); +} + +static inline float f16_to_f32(ggml_fp16_t x) { + return GGML_CPU_FP16_TO_FP32(x); +} + +static inline ggml_bf16_t f32_to_bf16(float x) { + return GGML_FP32_TO_BF16(x); +} + +static inline float bf16_to_f32(ggml_bf16_t x) { + return GGML_BF16_TO_FP32(x); +} + +static inline float i32_to_f32(int32_t x) { + return x; +} + +static inline int32_t f32_to_i32(float x) { + return x; +} + +static inline float f32_to_f32(float x) { + return x; +} + +// TODO - merge this into the traits table, after using row-based conversions +template +struct type_conversion_table; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_fp16_t) = f16_to_f32; + static constexpr ggml_fp16_t (*from_f32)(float) = f32_to_f16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(float) = f32_to_f32; + static constexpr float (*from_f32)(float) = f32_to_f32; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_bf16_t) = bf16_to_f32; + static constexpr ggml_bf16_t (*from_f32)(float) = f32_to_bf16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(int32_t) = i32_to_f32; + static constexpr int32_t (*from_f32)(float) = f32_to_i32; +}; + +static std::pair get_thread_range(const struct ggml_compute_params * params, const struct ggml_tensor * src0) { + const int64_t ith = params->ith; + const int64_t nth = params->nth; + + const int64_t nr = ggml_nrows(src0); + + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; + + // row range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); + + return {ir0, ir1}; +} + +struct ggml_fa_tile_config { + static constexpr size_t Q = GGML_FA_TILE_Q; + static constexpr size_t KV = GGML_FA_TILE_KV; +}; + +#endif + +#include "ggml.h.inc" + +// GGML internal header + +ggml_backend_buffer_type_t ggml_backend_cpu_repack_buffer_type(void); + +template constexpr int QK_0() { + if constexpr (K == 4) { + return QK4_0; + } + if constexpr (K == 8) { + return QK8_0; + } + return -1; +} + +template struct block { + ggml_half d[N]; // deltas for N qK_0 blocks + int8_t qs[(QK_0() * N * K) / 8]; // quants for N qK_0 blocks +}; + +// control size +static_assert(sizeof(block<4, 4>) == 4 * sizeof(ggml_half) + QK8_0 * 2, "wrong block<4,4> size/padding"); +static_assert(sizeof(block<4, 8>) == 8 * sizeof(ggml_half) + QK8_0 * 4, "wrong block<4,8> size/padding"); +static_assert(sizeof(block<4, 16>) == 16 * sizeof(ggml_half) + QK8_0 * 8, "wrong block<4,16> size/padding"); +static_assert(sizeof(block<8, 4>) == 4 * sizeof(ggml_half) + QK8_0 * 4, "wrong block<8,4> size/padding"); +static_assert(sizeof(block<8, 8>) == 8 * sizeof(ggml_half) + QK8_0 * 8, "wrong block<8,8> size/padding"); +static_assert(sizeof(block<8, 16>) == 16 * sizeof(ggml_half) + QK8_0 * 16, "wrong block<8,16> size/padding"); + +using block_q4_0x4 = block<4, 4>; +using block_q4_0x8 = block<4, 8>; +using block_q4_0x16 = block<4, 16>; +using block_q8_0x4 = block<8, 4>; +using block_q8_0x8 = block<8, 8>; +using block_q8_0x16 = block<8, 16>; + +struct block_q4_Kx8 { + ggml_half d[8]; // super-block scale for quantized scales + ggml_half dmin[8]; // super-block scale for quantized mins + uint8_t scales[96]; // scales and mins, quantized with 6 bits + uint8_t qs[1024]; // 4--bit quants +}; + +static_assert(sizeof(block_q4_Kx8) == sizeof(ggml_half) * 16 + K_SCALE_SIZE * 8 + QK_K * 4, "wrong q4_K block size/padding"); +struct block_q4_Kx16 { + ggml_half d[16]; // super-block scale for quantized scales + ggml_half dmin[16]; // super-block scale for quantized mins + uint8_t scales[192]; // scales and mins, quantized with 6 bits + uint8_t qs[2048]; // 4--bit quants +}; + +static_assert(sizeof(block_q4_Kx16) == sizeof(ggml_half) * 32 + K_SCALE_SIZE * 16 + QK_K * 8, "wrong q4_K block size/padding"); +struct block_q2_Kx8 { + ggml_half d[8]; // super-block scale for quantized scales + ggml_half dmin[8]; // super-block scale for quantized mins + uint8_t scales[128]; // scales and mins, quantized with 4 bits + uint8_t qs[512]; // 2--bit quants +}; + +static_assert(sizeof(block_q2_Kx8) == sizeof(ggml_half) * 16 + QK_K/2 + QK_K * 2, "wrong q2_K block size/padding"); +struct block_q2_Kx16 { + ggml_half d[16]; // Super-block scale for quantized scales + ggml_half dmin[16]; // Super-block scale for quantized mins + uint8_t scales[256]; // Sub-block scales (16 cols * 16 sub-blocks) + uint8_t qs[1024]; // Data (16 cols * 64 bytes per block) +}; +static_assert(sizeof(block_q2_Kx16) == sizeof(ggml_half) * 32 + QK_K + QK_K * 4, "wrong q2_K block size/padding"); + +struct block_q5_Kx8 { + ggml_half d[8]; // super-block scale for quantized scales + ggml_half dmin[8]; // super-block scale for quantized mins + uint8_t scales[96]; // scales and mins, quantized with 6 bits + uint8_t qh[QK_K * 8 / 8]; // high bits of 5-bit quants + uint8_t qs[QK_K * 8 / 2]; // low bits of 5-bit quants (in groups of 4) +}; + +static_assert(sizeof(block_q5_Kx8) == sizeof(ggml_half) * 16 + K_SCALE_SIZE * 8 + QK_K * 5, + "wrong q5_K block size/padding"); + +struct block_q6_Kx8 { + ggml_half d[8]; + int8_t scales[QK_K / 16 * 8]; + uint8_t ql[QK_K / 2 * 8]; // low bits of 6-bit quants (groups of 2) + uint8_t qh[QK_K / 4 * 8]; // high bits of 6-bit quants (groups of 4) +}; + +static_assert(sizeof(block_q6_Kx8) == sizeof(ggml_half) * 8 + QK_K / 16 * 8 + 3 * QK_K / 4 * 8, + "wrong q6_K block size/padding"); + +struct block_q8_Kx4 { + float d[4]; // delta + int8_t qs[QK_K * 4]; // quants + int16_t bsums[QK_K / 4]; // sum of quants in groups of 16 +}; + +static_assert(sizeof(block_q8_Kx4) == sizeof(float) * 4 + QK_K * 4 + (QK_K / 4) * sizeof(int16_t), "wrong q8_K block size/padding"); + +struct block_iq4_nlx4 { + ggml_half d[4]; // deltas for 4 iq4_nl blocks + uint8_t qs[QK4_NL * 2]; // nibbles / quants for 4 iq4_nl blocks +}; + +static_assert(sizeof(block_iq4_nlx4) == 4 * sizeof(ggml_half) + QK4_NL * 2, "wrong iq4_nlx4 block size/padding"); + +struct block_iq4_nlx8 { + ggml_half d[8]; // deltas for 8 iq4_nl blocks + uint8_t qs[QK4_NL * 4]; // nibbles / quants for 8 iq4_nl blocks +}; + +static_assert(sizeof(block_iq4_nlx8) == 8 * sizeof(ggml_half) + QK4_NL * 4, "wrong iq4_nlx8 block size/padding"); + +struct block_iq4_nlx16 { + ggml_half d[16]; // deltas for 16 iq4_nl blocks + uint8_t qs[QK4_NL * 8]; // nibbles / quants for 16 iq4_nl blocks +}; + +static_assert(sizeof(block_iq4_nlx16) == 16 * sizeof(ggml_half) + QK4_NL * 8, "wrong iq4_nlx16 block size/padding"); +struct block_mxfp4x4 { + uint8_t e[4]; + uint8_t qs[QK_MXFP4 * 2]; +}; +static_assert(sizeof(block_mxfp4x4) == 4 + QK_MXFP4 * 2, "wrong mxfp4x4 block size/padding"); + +struct block_mxfp4x8 { + uint8_t e[8]; + uint8_t qs[QK_MXFP4 * 4]; +}; +static_assert(sizeof(block_mxfp4x8) == 8 + QK_MXFP4 * 4, "wrong mxfp4x8 block size/padding"); + +#if defined(__cplusplus) +extern "C" { +#endif + +void ggml_quantize_mat_q8_0_4x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_0_4x8(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_K_4x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_K_4x8(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_gemv_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_0_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q2_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q5_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q5_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q6_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q6_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_iq4_nl_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_iq4_nl_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_mxfp4_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_mxfp4_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q8_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q8_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q2_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q5_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q5_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q6_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q6_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_iq4_nl_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_iq4_nl_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_mxfp4_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_mxfp4_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q8_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q8_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +#if defined __riscv_zvfh +void ggml_quantize_mat_q8_0_4x1(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_K_4x1(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_gemv_q4_0_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_K_16x1_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_iq4_nl_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q8_0_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q2_K_16x1_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_K_16x1_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_iq4_nl_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q8_0_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q2_K_16x1_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +#endif + +// Native implementations +void ggml_quantize_mat_q8_0_4x4_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_0_4x8_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_K_4x4_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_K_4x8_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_gemv_q4_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_0_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q2_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q5_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q5_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q6_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q6_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_iq4_nl_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_iq4_nl_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_mxfp4_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_mxfp4_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q8_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q8_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q2_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q5_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q5_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q6_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q6_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_iq4_nl_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_iq4_nl_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_mxfp4_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_mxfp4_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q8_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q8_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +#if defined __riscv_zvfh +void ggml_quantize_mat_q8_0_4x1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_K_4x1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_gemv_q4_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q8_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q2_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_iq4_nl_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q8_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q2_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_iq4_nl_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +#endif + +#if defined(__cplusplus) +} // extern "C" +#endif + +#include "ggml-impl-defs.inc" + +#include +#include +#include + +#ifdef GGML_USE_CPU_HBM +# include +#endif + +#if defined(__AMX_INT8__) && defined(__AVX512VNNI__) +ggml_backend_buffer_type_t ggml_backend_amx_buffer_type(void); +#endif + +#ifdef GGML_USE_CPU_KLEIDIAI +extern "C" ggml_backend_buffer_type_t ggml_backend_cpu_kleidiai_buffer_type(void); +#endif + +#ifdef GGML_USE_CPU_RISCV64_SPACEMIT +extern "C" ggml_backend_buffer_type_t ggml_backend_cpu_riscv64_spacemit_buffer_type(void); +#endif + +#if defined(_WIN32) +# define WIN32_LEAN_AND_MEAN +# ifndef NOMINMAX +# define NOMINMAX +# endif +# include +#else +# include +#endif + +#if defined(__APPLE__) +# include +# include +#endif + +// ggml-backend interface + +std::vector & ggml_backend_cpu_get_extra_buffer_types() { + static std::vector bufts = []() { + std::vector bufts; + +#if defined(__AMX_INT8__) && defined(__AVX512VNNI__) + if (ggml_backend_amx_buffer_type()) { + bufts.push_back(ggml_backend_amx_buffer_type()); + } +#endif + +#ifdef GGML_USE_CPU_RISCV64_SPACEMIT + if (ggml_backend_cpu_riscv64_spacemit_buffer_type()) { + bufts.push_back(ggml_backend_cpu_riscv64_spacemit_buffer_type()); + } +#endif + +#ifdef GGML_USE_CPU_KLEIDIAI + if (ggml_backend_cpu_kleidiai_buffer_type()) { + bufts.push_back(ggml_backend_cpu_kleidiai_buffer_type()); + } +#endif + +#ifdef GGML_USE_CPU_REPACK + if (ggml_backend_cpu_repack_buffer_type()) { + bufts.push_back(ggml_backend_cpu_repack_buffer_type()); + } +#endif + + return bufts; + }(); + + return bufts; +} + +namespace ggml::cpu { +tensor_traits::~tensor_traits() {} + +extra_buffer_type::~extra_buffer_type() {} +} // namespace ggml::cpu + +bool ggml_cpu_extra_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) { + for (auto extra : ggml_backend_cpu_get_extra_buffer_types()) { + if (extra && extra->context) { + auto buf_extra = (ggml::cpu::extra_buffer_type *) extra->context; + auto tensor_traits = buf_extra->get_tensor_traits(op); + if (tensor_traits && tensor_traits->compute_forward(params, op)) { + return true; + } + } + } + return false; +} + +bool ggml_cpu_extra_work_size(int n_threads, const struct ggml_tensor * op, size_t * size) { + for (auto extra : ggml_backend_cpu_get_extra_buffer_types()) { + if (extra && extra->context) { + auto buf_extra = (ggml::cpu::extra_buffer_type *) extra->context; + auto tensor_traits = buf_extra->get_tensor_traits(op); + if (tensor_traits && tensor_traits->work_size(n_threads, op, *size)) { + return true; + } + } + } + return false; +} + +#ifdef GGML_USE_CPU_HBM +static const char * ggml_backend_cpu_hbm_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + return "CPU_HBM"; + + GGML_UNUSED(buft); +} + +static void ggml_backend_cpu_hbm_buffer_free_buffer(ggml_backend_buffer_t buffer) { + hbw_free(buffer->context); +} + +static ggml_backend_buffer_t ggml_backend_cpu_hbm_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, + size_t size) { + void * ptr; + int result = hbw_posix_memalign(&ptr, ggml_backend_cpu_buffer_type_get_alignment(buft), size); + if (result != 0) { + GGML_LOG_ERROR("failed to allocate HBM buffer of size %zu\n", size); + return NULL; + } + + ggml_backend_buffer_t buffer = ggml_backend_cpu_buffer_from_ptr(ptr, size); + buffer->buft = buft; + buffer->iface.free_buffer = ggml_backend_cpu_hbm_buffer_free_buffer; + + return buffer; +} + +ggml_backend_buffer_type_t ggml_backend_cpu_hbm_buffer_type(void) { + static struct ggml_backend_buffer_type ggml_backend_cpu_buffer_type_hbm = { + /* .iface = */ { + /* .get_name = */ ggml_backend_cpu_hbm_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_cpu_hbm_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_cpu_buffer_type_get_alignment, + /* .get_max_size = */ nullptr, // defaults to SIZE_MAX + /* .get_alloc_size = */ nullptr, // defaults to ggml_nbytes + /* .is_host = */ ggml_backend_cpu_buffer_type_is_host, + }, + /* .context = */ nullptr, + }; + + return &ggml_backend_cpu_buffer_type_hbm; +} +#endif + +static ggml_backend_buffer_type_t * ggml_backend_cpu_device_get_extra_buffers_type(ggml_backend_dev_t device) { + static std::vector extra_bufts = [] { + std::vector bufts = ggml_backend_cpu_get_extra_buffer_types(); + bufts.push_back(nullptr); + return bufts; + }(); + + return extra_bufts.data(); + + GGML_UNUSED(device); +} + +static bool ggml_backend_cpu_is_extra_buffer_type(ggml_backend_buffer_type_t buft) { + for (auto * extra : ggml_backend_cpu_get_extra_buffer_types()) { + if (extra == buft) { + return true; + } + } + return false; +} + +// CPU backend - backend (stream) + +struct ggml_backend_cpu_context { + int n_threads; + ggml_threadpool_t threadpool; + + uint8_t * work_data; + size_t work_size; + + ggml_abort_callback abort_callback; + void * abort_callback_data; + + bool use_ref; // use reference implementation +}; + +static const char * ggml_backend_cpu_get_name(ggml_backend_t backend) { + return "CPU"; + + GGML_UNUSED(backend); +} + +static void ggml_backend_cpu_free(ggml_backend_t backend) { + struct ggml_backend_cpu_context * cpu_ctx = (struct ggml_backend_cpu_context *)backend->context; + delete[] cpu_ctx->work_data; + delete cpu_ctx; + delete backend; +} + +struct ggml_backend_plan_cpu { + struct ggml_cplan cplan; + struct ggml_cgraph cgraph; +}; + +static ggml_backend_graph_plan_t ggml_backend_cpu_graph_plan_create(ggml_backend_t backend, const struct ggml_cgraph * cgraph) { + struct ggml_backend_cpu_context * cpu_ctx = (struct ggml_backend_cpu_context *)backend->context; + + struct ggml_backend_plan_cpu * cpu_plan = new ggml_backend_plan_cpu; + + cpu_plan->cplan = ggml_graph_plan(cgraph, cpu_ctx->n_threads, cpu_ctx->threadpool); + cpu_plan->cgraph = *cgraph; // FIXME: deep copy + + if (cpu_plan->cplan.work_size > 0) { + cpu_plan->cplan.work_data = new uint8_t[cpu_plan->cplan.work_size]; + if (cpu_plan->cplan.work_data == NULL) { + delete cpu_plan; + return NULL; + } + } + + cpu_plan->cplan.abort_callback = cpu_ctx->abort_callback; + cpu_plan->cplan.abort_callback_data = cpu_ctx->abort_callback_data; + cpu_plan->cplan.use_ref = cpu_ctx->use_ref; + + return cpu_plan; +} + +static void ggml_backend_cpu_graph_plan_free(ggml_backend_t backend, ggml_backend_graph_plan_t plan) { + struct ggml_backend_plan_cpu * cpu_plan = (struct ggml_backend_plan_cpu *)plan; + + delete[] cpu_plan->cplan.work_data; + delete cpu_plan; + + GGML_UNUSED(backend); +} + +static enum ggml_status ggml_backend_cpu_graph_plan_compute(ggml_backend_t backend, ggml_backend_graph_plan_t plan) { + struct ggml_backend_plan_cpu * cpu_plan = (struct ggml_backend_plan_cpu *)plan; + + return ggml_graph_compute(&cpu_plan->cgraph, &cpu_plan->cplan); + + GGML_UNUSED(backend); +} + +static enum ggml_status ggml_backend_cpu_graph_compute(ggml_backend_t backend, struct ggml_cgraph * cgraph) { + struct ggml_backend_cpu_context * cpu_ctx = (struct ggml_backend_cpu_context *)backend->context; + + struct ggml_cplan cplan = ggml_graph_plan(cgraph, cpu_ctx->n_threads, cpu_ctx->threadpool); + + if (cpu_ctx->work_size < cplan.work_size) { + delete[] cpu_ctx->work_data; + cpu_ctx->work_data = new uint8_t[cplan.work_size]; + if (cpu_ctx->work_data == NULL) { + cpu_ctx->work_size = 0; + return GGML_STATUS_ALLOC_FAILED; + } + cpu_ctx->work_size = cplan.work_size; + } + cplan.work_data = (uint8_t *)cpu_ctx->work_data; + + cplan.abort_callback = cpu_ctx->abort_callback; + cplan.abort_callback_data = cpu_ctx->abort_callback_data; + cplan.use_ref = cpu_ctx->use_ref; + + return ggml_graph_compute(cgraph, &cplan); +} + +static const struct ggml_backend_i ggml_backend_cpu_i = { + /* .get_name = */ ggml_backend_cpu_get_name, + /* .free = */ ggml_backend_cpu_free, + /* .set_tensor_async = */ NULL, + /* .get_tensor_async = */ NULL, + /* .set_tensor_2d_async = */ NULL, + /* .get_tensor_2d_async = */ NULL, + /* .cpy_tensor_async = */ NULL, + /* .synchronize = */ NULL, + /* .graph_plan_create = */ ggml_backend_cpu_graph_plan_create, + /* .graph_plan_free = */ ggml_backend_cpu_graph_plan_free, + /* .graph_plan_update = */ NULL, + /* .graph_plan_compute = */ ggml_backend_cpu_graph_plan_compute, + /* .graph_compute = */ ggml_backend_cpu_graph_compute, + /* .event_record = */ NULL, + /* .event_wait = */ NULL, + /* .graph_optimize = */ NULL, +}; + +static ggml_guid_t ggml_backend_cpu_guid(void) { + static ggml_guid guid = { 0xaa, 0x67, 0xc7, 0x43, 0x96, 0xe6, 0xa3, 0x8a, 0xe3, 0xaf, 0xea, 0x92, 0x36, 0xbc, 0xfc, 0x89 }; + return &guid; +} + +ggml_backend_t ggml_backend_cpu_init(void) { + // initialize CPU backend now to avoid slowing the first graph computation + ggml_cpu_init(); + + struct ggml_backend_cpu_context * ctx = new ggml_backend_cpu_context; + if (ctx == NULL) { + return NULL; + } + + ctx->n_threads = GGML_DEFAULT_N_THREADS; + ctx->threadpool = NULL; + ctx->work_data = NULL; + ctx->work_size = 0; + ctx->abort_callback = NULL; + ctx->abort_callback_data = NULL; + ctx->use_ref = false; + + ggml_backend_t cpu_backend = new ggml_backend { + /* .guid = */ ggml_backend_cpu_guid(), + /* .iface = */ ggml_backend_cpu_i, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0), + /* .context = */ ctx, + }; + + if (cpu_backend == NULL) { + delete ctx; + return NULL; + } + + return cpu_backend; +} + +bool ggml_backend_is_cpu(ggml_backend_t backend) { + return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_cpu_guid()); +} + +void ggml_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads) { + GGML_ASSERT(ggml_backend_is_cpu(backend_cpu)); + + struct ggml_backend_cpu_context * ctx = (struct ggml_backend_cpu_context *)backend_cpu->context; + ctx->n_threads = n_threads; +} + +void ggml_backend_cpu_set_threadpool(ggml_backend_t backend_cpu, ggml_threadpool_t threadpool) { + GGML_ASSERT(ggml_backend_is_cpu(backend_cpu)); + + struct ggml_backend_cpu_context * ctx = (struct ggml_backend_cpu_context *)backend_cpu->context; + + if (ctx->threadpool && ctx->threadpool != threadpool) { + // already had a different threadpool, pause/suspend it before switching + ggml_threadpool_pause(ctx->threadpool); + } + ctx->threadpool = threadpool; +} + +void ggml_backend_cpu_set_abort_callback(ggml_backend_t backend_cpu, ggml_abort_callback abort_callback, void * abort_callback_data) { + GGML_ASSERT(ggml_backend_is_cpu(backend_cpu)); + + struct ggml_backend_cpu_context * ctx = (struct ggml_backend_cpu_context *)backend_cpu->context; + ctx->abort_callback = abort_callback; + ctx->abort_callback_data = abort_callback_data; +} + +void ggml_backend_cpu_set_use_ref(ggml_backend_t backend_cpu, bool use_ref) { + GGML_ASSERT(ggml_backend_is_cpu(backend_cpu)); + + struct ggml_backend_cpu_context * ctx = (struct ggml_backend_cpu_context *)backend_cpu->context; + ctx->use_ref = use_ref; +} + +// CPU backend - device + +struct ggml_backend_cpu_device_context { + std::string description = "CPU"; + + ggml_backend_cpu_device_context() { +#ifdef __APPLE__ + size_t len = 0; + if (!sysctlbyname("machdep.cpu.brand_string", NULL, &len, NULL, 0)) { + description.resize(len); + sysctlbyname("machdep.cpu.brand_string", &description[0], &len, NULL, 0); // NOLINT + } +#elif defined(__linux__) + FILE * f = fopen("/proc/cpuinfo", "r"); + if (f) { + char buf[1024]; + while (fgets(buf, sizeof(buf), f)) { + if (strncmp(buf, "model name", 10) == 0) { + char * p = strchr(buf, ':'); + if (p) { + p++; + while (std::isspace(*p)) { + p++; + } + while (std::isspace(p[strlen(p) - 1])) { + p[strlen(p) - 1] = '\0'; + } + description = p; + break; + } + } + } + fclose(f); + } +#elif defined(_WIN32) + HKEY hKey; + if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, + TEXT("HARDWARE\\DESCRIPTION\\System\\CentralProcessor\\0"), + 0, + KEY_READ, + &hKey) == ERROR_SUCCESS) { + DWORD cpu_brand_size = 0; + if (RegQueryValueExA(hKey, + "ProcessorNameString", + NULL, + NULL, + NULL, + &cpu_brand_size) == ERROR_SUCCESS) { + description.resize(cpu_brand_size); + if (RegQueryValueExA(hKey, + "ProcessorNameString", + NULL, + NULL, + (LPBYTE)&description[0], // NOLINT + &cpu_brand_size) == ERROR_SUCCESS) { + if (description.find('\0') != std::string::npos) { + description.resize(description.find('\0')); + } + } + } + RegCloseKey(hKey); + } +#endif + } +}; + +static const char * ggml_backend_cpu_device_get_name(ggml_backend_dev_t dev) { + return "CPU"; + + GGML_UNUSED(dev); +} + +static const char * ggml_backend_cpu_device_get_description(ggml_backend_dev_t dev) { + struct ggml_backend_cpu_device_context * ctx = (struct ggml_backend_cpu_device_context *)dev->context; + + return ctx->description.c_str(); +} + +static void ggml_backend_cpu_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { +#ifdef _WIN32 + MEMORYSTATUSEX status; + status.dwLength = sizeof(status); + GlobalMemoryStatusEx(&status); + *total = status.ullTotalPhys; + *free = status.ullAvailPhys; +#else + long pages = sysconf(_SC_PHYS_PAGES); + long page_size = sysconf(_SC_PAGE_SIZE); + *total = pages * page_size; + + // "free" system memory is ill-defined, for practical purposes assume that all of it is free: + *free = *total; +#endif // _WIN32 + + GGML_UNUSED(dev); +} + +static enum ggml_backend_dev_type ggml_backend_cpu_device_get_type(ggml_backend_dev_t dev) { + return GGML_BACKEND_DEVICE_TYPE_CPU; + + GGML_UNUSED(dev); +} + +static void ggml_backend_cpu_device_get_props(ggml_backend_dev_t dev, struct ggml_backend_dev_props * props) { + props->name = ggml_backend_cpu_device_get_name(dev); + props->description = ggml_backend_cpu_device_get_description(dev); + props->type = ggml_backend_cpu_device_get_type(dev); + ggml_backend_cpu_device_get_memory(dev, &props->memory_free, &props->memory_total); + props->caps = { + /* .async = */ false, + /* .host_buffer = */ false, + /* .buffer_from_host_ptr = */ true, + /* .events = */ false, + }; +} + +static ggml_backend_t ggml_backend_cpu_device_init_backend(ggml_backend_dev_t dev, const char * params) { + return ggml_backend_cpu_init(); + + GGML_UNUSED(dev); + GGML_UNUSED(params); +} + +static ggml_backend_buffer_type_t ggml_backend_cpu_device_get_buffer_type(ggml_backend_dev_t dev) { + return ggml_backend_cpu_buffer_type(); + + GGML_UNUSED(dev); +} + +static ggml_backend_buffer_t ggml_backend_cpu_device_buffer_from_host_ptr(ggml_backend_dev_t dev, void * ptr, size_t size, size_t max_tensor_size) { + return ggml_backend_cpu_buffer_from_ptr(ptr, size); + + GGML_UNUSED(dev); + GGML_UNUSED(max_tensor_size); +} + +static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) { + const struct ggml_tensor * src0 = op->src[0]; + const struct ggml_tensor * src1 = op->src[1]; + + if (op->op == GGML_OP_NONE || op->op == GGML_OP_RESHAPE || op->op == GGML_OP_VIEW || op->op == GGML_OP_PERMUTE || op->op == GGML_OP_TRANSPOSE) { + return true; + } + + // check extra buffer types + // note: only the first sources are checked for extra buffer types to reduce overhead, increase if necessary + for (int i = 0; i < 4; i++) { + if (op->src[i] && op->src[i]->buffer && + ggml_backend_cpu_is_extra_buffer_type(op->src[i]->buffer->buft)) { + auto * buf_extra = (ggml::cpu::extra_buffer_type *) op->src[i]->buffer->buft->context; + return buf_extra->supports_op(dev, op); + } + } + + switch (op->op) { + case GGML_OP_CPY: + case GGML_OP_SET_ROWS: + return + op->type != GGML_TYPE_IQ3_XXS && + op->type != GGML_TYPE_IQ3_S && + op->type != GGML_TYPE_IQ2_XXS && + op->type != GGML_TYPE_IQ2_XS && + op->type != GGML_TYPE_IQ2_S && + op->type != GGML_TYPE_IQ1_S && + op->type != GGML_TYPE_IQ1_M; // missing type_traits.from_float + case GGML_OP_MUL_MAT: + return src1->type == GGML_TYPE_F32 || src1->type == ggml_get_type_traits_cpu(src0->type)->vec_dot_type; + case GGML_OP_SOFT_MAX_BACK: { + if (op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32) { + return false; + } + float max_bias = 0.0f; + + memcpy(&max_bias, (const float *) op->op_params + 1, sizeof(float)); + + return max_bias == 0.0f; + } + case GGML_OP_IM2COL_BACK: + return src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32; + case GGML_OP_GET_ROWS_BACK: + return src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16; + case GGML_OP_OUT_PROD: + return (src0->type == GGML_TYPE_F32 || (ggml_is_quantized(src0->type) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) && + src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; + default: + return true; + } +} + +static bool ggml_backend_cpu_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { + return ggml_backend_buft_is_host(buft) || ggml_backend_cpu_is_extra_buffer_type(buft); + GGML_UNUSED(dev); +} + +static const struct ggml_backend_device_i ggml_backend_cpu_device_i = { + /* .get_name = */ ggml_backend_cpu_device_get_name, + /* .get_description = */ ggml_backend_cpu_device_get_description, + /* .get_memory = */ ggml_backend_cpu_device_get_memory, + /* .get_type = */ ggml_backend_cpu_device_get_type, + /* .get_props = */ ggml_backend_cpu_device_get_props, + /* .init_backend = */ ggml_backend_cpu_device_init_backend, + /* .get_buffer_type = */ ggml_backend_cpu_device_get_buffer_type, + /* .get_host_buffer_type = */ NULL, + /* .buffer_from_host_ptr = */ ggml_backend_cpu_device_buffer_from_host_ptr, + /* .supports_op = */ ggml_backend_cpu_device_supports_op, + /* .supports_buft = */ ggml_backend_cpu_device_supports_buft, + /* .offload_op = */ NULL, + /* .event_new = */ NULL, + /* .event_free = */ NULL, + /* .event_synchronize = */ NULL, +}; + +// CPU backend - backend (reg) + +static const char * ggml_backend_cpu_reg_get_name(ggml_backend_reg_t reg) { + return "CPU"; + + GGML_UNUSED(reg); +} + +static size_t ggml_backend_cpu_reg_get_device_count(ggml_backend_reg_t reg) { + return 1; + + GGML_UNUSED(reg); +} + +static ggml_backend_dev_t ggml_backend_cpu_reg_get_device(ggml_backend_reg_t reg, size_t index) { + GGML_ASSERT(index == 0); + + static ggml_backend_cpu_device_context ctx; + static ggml_backend_device ggml_backend_cpu_device = { + /* .iface = */ ggml_backend_cpu_device_i, + /* .reg = */ reg, + /* .context = */ &ctx, + }; + + return &ggml_backend_cpu_device; +} + +// This is intended to replace the the ggml_cpu_has_* functions when loading the CPU backend dynamically, +// and additionally to allow other backends to expose their own list of features that applications can query using the same API +static ggml_backend_feature * ggml_backend_cpu_get_features(ggml_backend_reg_t reg) { + static std::vector features = []() { + ggml_cpu_init(); + + std::vector features; + if (ggml_cpu_has_sse3()) { + features.push_back({ "SSE3", "1" }); + } + if (ggml_cpu_has_ssse3()) { + features.push_back({ "SSSE3", "1" }); + } + if (ggml_cpu_has_avx()) { + features.push_back({ "AVX", "1" }); + } + if (ggml_cpu_has_avx_vnni()) { + features.push_back({ "AVX_VNNI", "1" }); + } + if (ggml_cpu_has_avx2()) { + features.push_back({ "AVX2", "1" }); + } + if (ggml_cpu_has_f16c()) { + features.push_back({ "F16C", "1" }); + } + if (ggml_cpu_has_fma()) { + features.push_back({ "FMA", "1" }); + } + if (ggml_cpu_has_bmi2()) { + features.push_back({ "BMI2", "1" }); + } + if (ggml_cpu_has_avx512()) { + features.push_back({ "AVX512", "1" }); + } + if (ggml_cpu_has_avx512_vbmi()) { + features.push_back({ "AVX512_VBMI", "1" }); + } + if (ggml_cpu_has_avx512_vnni()) { + features.push_back({ "AVX512_VNNI", "1" }); + } + if (ggml_cpu_has_avx512_bf16()) { + features.push_back({ "AVX512_BF16", "1" }); + } + if (ggml_cpu_has_amx_int8()) { + features.push_back({ "AMX_INT8", "1" }); + } + if (ggml_cpu_has_neon()) { + features.push_back({ "NEON", "1" }); + } + if (ggml_cpu_has_arm_fma()) { + features.push_back({ "ARM_FMA", "1" }); + } + if (ggml_cpu_has_fp16_va()) { + features.push_back({ "FP16_VA", "1" }); + } + if (ggml_cpu_has_matmul_int8()) { + features.push_back({ "MATMUL_INT8", "1" }); + } + if (ggml_cpu_has_sve()) { + features.push_back({ "SVE", "1" }); + } + if (ggml_cpu_has_dotprod()) { + features.push_back({ "DOTPROD", "1" }); + } + if (ggml_cpu_get_sve_cnt() > 0) { + static std::string sve_cnt = std::to_string(ggml_cpu_get_sve_cnt()); + features.push_back({ "SVE_CNT", sve_cnt.c_str() }); + } + if (ggml_cpu_has_sme()) { + features.push_back({ "SME", "1" }); + } + if (ggml_cpu_has_riscv_v()) { + features.push_back({ "RISCV_V", "1" }); + } + if (ggml_cpu_get_rvv_vlen() > 0) { + static std::string rvv_vlen = std::to_string(ggml_cpu_get_rvv_vlen()); + features.push_back({ "RVV_VLEN", rvv_vlen.c_str() }); + } + if (ggml_cpu_has_vsx()) { + features.push_back({ "VSX", "1" }); + } + if (ggml_cpu_has_vxe()) { + features.push_back({ "VXE", "1" }); + } + if (ggml_cpu_has_wasm_simd()) { + features.push_back({ "WASM_SIMD", "1" }); + } + if (ggml_cpu_has_llamafile()) { + features.push_back({ "LLAMAFILE", "1" }); + } + #ifdef GGML_USE_ACCELERATE + features.push_back({ "ACCELERATE", "1" }); + #endif + #ifdef GGML_USE_CPU_HBM + features.push_back({ "CPU_HBM", "1" }); + #endif + #ifdef GGML_USE_OPENMP + features.push_back({ "OPENMP", "1" }); + #endif + #ifdef GGML_USE_CPU_KLEIDIAI + features.push_back({ "KLEIDIAI", "1" }); + #endif + #ifdef GGML_USE_CPU_REPACK + features.push_back({ "REPACK", "1" }); + #endif + + features.push_back({ nullptr, nullptr }); + + return features; + }(); + + return features.data(); + + GGML_UNUSED(reg); +} + +static void * ggml_backend_cpu_get_proc_address(ggml_backend_reg_t reg, const char * name) { + if (strcmp(name, "ggml_backend_set_n_threads") == 0) { + ggml_backend_set_n_threads_t fct = ggml_backend_cpu_set_n_threads; + return (void *)fct; + } + if (strcmp(name, "ggml_backend_dev_get_extra_bufts") == 0) { + ggml_backend_dev_get_extra_bufts_t fct = ggml_backend_cpu_device_get_extra_buffers_type; + return (void *)fct; + } + if (strcmp(name, "ggml_backend_get_features") == 0) { + return (void *)ggml_backend_cpu_get_features; + } + if (strcmp(name, "ggml_backend_set_abort_callback") == 0) { + return (void *)ggml_backend_cpu_set_abort_callback; + } + if (strcmp(name, "ggml_backend_cpu_numa_init") == 0) { + return (void *)ggml_numa_init; + } + if (strcmp(name, "ggml_backend_cpu_is_numa") == 0) { + return (void *)ggml_is_numa; + } + if (strcmp(name, "ggml_backend_cpu_set_use_ref") == 0) { + return (void *)ggml_backend_cpu_set_use_ref; + } + + // threadpool - TODO: move to ggml-base + if (strcmp(name, "ggml_threadpool_new") == 0) { + return (void *)ggml_threadpool_new; + } + if (strcmp(name, "ggml_threadpool_free") == 0) { + return (void *)ggml_threadpool_free; + } + if (strcmp(name, "ggml_backend_cpu_set_threadpool") == 0) { + return (void *)ggml_backend_cpu_set_threadpool; + } + + return NULL; + + GGML_UNUSED(reg); +} + +static const struct ggml_backend_reg_i ggml_backend_cpu_reg_i = { + /* .get_name = */ ggml_backend_cpu_reg_get_name, + /* .get_device_count = */ ggml_backend_cpu_reg_get_device_count, + /* .get_device = */ ggml_backend_cpu_reg_get_device, + /* .get_proc_address = */ ggml_backend_cpu_get_proc_address, +}; + +ggml_backend_reg_t ggml_backend_cpu_reg(void) { + // init CPU feature detection + ggml_cpu_init(); + + static struct ggml_backend_reg ggml_backend_cpu_reg = { + /* .api_version = */ GGML_BACKEND_API_VERSION, + /* .iface = */ ggml_backend_cpu_reg_i, + /* .context = */ NULL, + }; + + return &ggml_backend_cpu_reg; +} + +GGML_BACKEND_DL_IMPL(ggml_backend_cpu_reg) + + +// ---- amx/amx implementation ---- + +/* ---- inline AMX common helpers ---- */ + +#include "ggml.h.inc" +#include "ggml-cpu-impl-defs.inc" + +#include +#include +#include + +#if defined(GGML_USE_OPENMP) +#include +#else +#include +#endif + +#define TILE_M 16 +#define TILE_N 16 +#define TILE_K 32 +#define VNNI_BLK 4 + +#define AMX_BLK_SIZE 32 + +#define TMM0 0 +#define TMM1 1 +#define TMM2 2 +#define TMM3 3 +#define TMM4 4 +#define TMM5 5 +#define TMM6 6 +#define TMM7 7 + +// parallel routines +template ::value, int>::type = 0> +inline T div_up(T x, T y) { return (x + y - 1) / y; } + +template +inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) { +#if 0 + // onednn partition pattern + T& n_my = n_end; + if (nth <= 1 || n == 0) { + n_start = 0; + n_my = n; + } else { + T n1 = div_up(n, nth); + T n2 = n1 - 1; + T T1 = n - n2 * nth; + n_my = ith < T1 ? n1 : n2; + n_start = ith <= T1 ? ith*n1 : T1 * n1 + (ith - T1) * n2; + } + n_end += n_start; +#else + // pytorch aten partition pattern + T n_my = div_up(n, nth); + n_start = ith * n_my; + n_end = std::min(n_start + n_my, n); +#endif +} + +template +inline void parallel_for(int n, const func_t & f) { + if (n <= 0) { + return; + } +#if defined(GGML_USE_OPENMP) + #pragma omp parallel + { + int nth = omp_get_num_threads(); + int ith = omp_get_thread_num(); + int tbegin, tend; + balance211(n, nth, ith, tbegin, tend); + f(tbegin, tend); + } +#else + int nth = std::thread::hardware_concurrency(); + if (nth <= 1) { + f(0, n); + return; + } + if (nth > n) { + nth = n; + } + std::vector threads; + threads.reserve(nth); + for (int ith = 0; ith < nth; ++ith) { + threads.emplace_back([&f, n, ith, nth] { + int tbegin, tend; + balance211(n, nth, ith, tbegin, tend); + f(tbegin, tend); + }); + } + for (auto & t : threads) { + t.join(); + } +#endif +} + +template +inline void parallel_for_ggml(const ggml_compute_params * params, int n, const func_t & f) { + int tbegin, tend; + balance211(n, params->nth, params->ith, tbegin, tend); + f(tbegin, tend); +} + +// quantized types that have AMX support +inline bool qtype_has_amx_kernels(const enum ggml_type type) { + // TODO: fix padding for vnni format + return (type == GGML_TYPE_Q4_0) || + (type == GGML_TYPE_Q4_1) || + (type == GGML_TYPE_Q8_0) || + (type == GGML_TYPE_Q4_K) || + (type == GGML_TYPE_Q5_K) || + (type == GGML_TYPE_Q6_K) || + (type == GGML_TYPE_IQ4_XS); +} + +#include "ggml-backend-impl-defs.inc" +#include "ggml-backend.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-cpu.h.inc" + +#if defined(__linux__) +#include +#include +#endif + +#include +#include +#include + +size_t ggml_backend_amx_desired_wsize(const struct ggml_tensor * dst); +size_t ggml_backend_amx_get_alloc_size(const struct ggml_tensor * tensor); +void ggml_backend_amx_convert_weight(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size); +void ggml_backend_amx_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst); + +#if defined(__AMX_INT8__) && defined(__AVX512VNNI__) + +// AMX type_trais +namespace ggml::cpu::amx { +class tensor_traits : public ggml::cpu::tensor_traits { + bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override { + size = ggml_backend_amx_desired_wsize(op); + return true; + } + + bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) override { + if (op->op == GGML_OP_MUL_MAT) { + ggml_backend_amx_mul_mat(params, op); + return true; + } + return false; + } +}; + +static ggml::cpu::tensor_traits * get_tensor_traits(ggml_backend_buffer_t, struct ggml_tensor *) { + static tensor_traits traits; + return &traits; +} +} // namespace ggml::cpu::amx + +// AMX buffer interface +static void ggml_backend_amx_buffer_free_buffer(ggml_backend_buffer_t buffer) { + free(buffer->context); +} + +static void * ggml_backend_amx_buffer_get_base(ggml_backend_buffer_t buffer) { + return (void *) (buffer->context); +} + +static enum ggml_status ggml_backend_amx_buffer_init_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor) { + tensor->extra = (void *) ggml::cpu::amx::get_tensor_traits(buffer, tensor); + + GGML_UNUSED(buffer); + return GGML_STATUS_SUCCESS; +} + +static void ggml_backend_amx_buffer_memset_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, + uint8_t value, size_t offset, size_t size) { + memset((char *) tensor->data + offset, value, size); + + GGML_UNUSED(buffer); +} + +static void ggml_backend_amx_buffer_set_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, + const void * data, size_t offset, size_t size) { + if (qtype_has_amx_kernels(tensor->type)) { + GGML_LOG_DEBUG("%s: amx repack tensor %s of type %s\n", __func__, tensor->name, ggml_type_name(tensor->type)); + ggml_backend_amx_convert_weight(tensor, data, offset, size); + } else { + memcpy((char *) tensor->data + offset, data, size); + } + + GGML_UNUSED(buffer); +} + +/* +// need to figure what we need to do with buffer->extra. +static void ggml_backend_amx_buffer_get_tensor(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) { + GGML_ASSERT(!qtype_has_amx_kernels(tensor->type)); + memcpy(data, (const char *)tensor->data + offset, size); + + GGML_UNUSED(buffer); +} + +static bool ggml_backend_amx_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const struct ggml_tensor * src, struct ggml_tensor * dst) { + if (ggml_backend_buffer_is_host(src->buffer)) { + if (qtype_has_amx_kernels(src->type)) { + ggml_backend_amx_convert_weight(dst, src->data, 0, ggml_nbytes(dst)); + } else { + memcpy(dst->data, src->data, ggml_nbytes(src)); + } + return true; + } + return false; + + GGML_UNUSED(buffer); +} +*/ + +static void ggml_backend_amx_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + memset(buffer->context, value, buffer->size); +} + +static ggml_backend_buffer_i ggml_backend_amx_buffer_interface = { + /* .free_buffer = */ ggml_backend_amx_buffer_free_buffer, + /* .get_base = */ ggml_backend_amx_buffer_get_base, + /* .init_tensor = */ ggml_backend_amx_buffer_init_tensor, + /* .memset_tensor = */ ggml_backend_amx_buffer_memset_tensor, + /* .set_tensor = */ ggml_backend_amx_buffer_set_tensor, + /* .get_tensor = */ nullptr, + /* .set_tensor_2d = */ nullptr, + /* .get_tensor_2d = */ nullptr, + /* .cpy_tensor = */ nullptr, + /* .clear = */ ggml_backend_amx_buffer_clear, + /* .reset = */ nullptr, +}; + +static const char * ggml_backend_amx_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + return "AMX"; + + GGML_UNUSED(buft); +} + +static ggml_backend_buffer_t ggml_backend_amx_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + void * data = ggml_aligned_malloc(size); + if (data == NULL) { + fprintf(stderr, "%s: failed to allocate buffer of size %zu\n", __func__, size); + return NULL; + } + + return ggml_backend_buffer_init(buft, ggml_backend_amx_buffer_interface, data, size); +} + +static size_t ggml_backend_amx_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + return TENSOR_ALIGNMENT; + + GGML_UNUSED(buft); +} + +namespace ggml::cpu::amx { +class extra_buffer_type : ggml::cpu::extra_buffer_type { + bool supports_op(ggml_backend_dev_t, const struct ggml_tensor * op) override { + if (op->op != GGML_OP_MUL_MAT) { + return false; + } + auto * src0 = op->src[0]; + auto * src1 = op->src[1]; + + if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(src1)) { + return false; + } + if (!src0->buffer || src0->buffer->buft != ggml_backend_amx_buffer_type()) { + return false; + } + if (src1->buffer && !ggml_backend_buft_is_host(src1->buffer->buft)) { + return false; + } + if (op->ne[0] % (TILE_N * 2)) { + return false; + } + int alignment; + switch (src0->type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q8_0: + alignment = TILE_K; + break; + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_XS: + alignment = 256; // QK_K + break; + case GGML_TYPE_F16: + alignment = 16; + break; + default: + return false; + } + if (src0->ne[0] % alignment) { + return false; + } + if (src1->type != GGML_TYPE_F32) { + return false; + } + return true; + } + + ggml::cpu::tensor_traits * get_tensor_traits(const struct ggml_tensor * op) override { + if (op->op == GGML_OP_MUL_MAT && op->src[0]->buffer && + op->src[0]->buffer->buft == ggml_backend_amx_buffer_type()) { + return (ggml::cpu::tensor_traits *) op->src[0]->extra; + } + + return nullptr; + } +}; +} // namespace ggml::cpu::amx + +static size_t ggml_backend_amx_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { + return ggml_backend_amx_get_alloc_size(tensor); + + GGML_UNUSED(buft); +} + +#define ARCH_GET_XCOMP_PERM 0x1022 +#define ARCH_REQ_XCOMP_PERM 0x1023 +#define XFEATURE_XTILECFG 17 +#define XFEATURE_XTILEDATA 18 + +static bool ggml_amx_init() { +#if defined(__linux__) + if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA)) { + fprintf(stderr, "AMX is not ready to be used!\n"); + return false; + } + return true; +#elif defined(_WIN32) + return true; +#else + return false; +#endif +} + +ggml_backend_buffer_type_t ggml_backend_amx_buffer_type() { + static struct ggml_backend_buffer_type ggml_backend_buffer_type_amx = { + /* .iface = */ { + /* .get_name = */ ggml_backend_amx_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_amx_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_amx_buffer_type_get_alignment, + /* .get_max_size = */ nullptr, // defaults to SIZE_MAX + /* .get_alloc_size = */ ggml_backend_amx_buffer_type_get_alloc_size, + /* .is_host = */ nullptr, + }, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0), + /* .context = */ new ggml::cpu::amx::extra_buffer_type(), + }; + + if (!ggml_amx_init()) { + return nullptr; + } + + return &ggml_backend_buffer_type_amx; +} + +#endif // defined(__AMX_INT8__) && defined(__AVX512VNNI__) + + +// ---- amx/mmq implementation ---- + +#if defined(__GNUC__) +#pragma GCC diagnostic ignored "-Wpedantic" +#pragma GCC diagnostic ignored "-Wunused-local-typedefs" +#endif + +#include "ggml-impl-defs.inc" +#include "ggml-cpu-impl-defs.inc" +#include "simd-mappings-defs.inc" +#define GGML_COMMON_DECL_C +#include "ggml-common-defs.inc" + +#include "ggml.h.inc" + +// GGML CPU internal header + +#ifdef __cplusplus +extern "C" { +#endif + +// Quantization +void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_q2_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q3_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q6_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_tq1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_tq2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_iq4_nl (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_iq4_xs (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +// Dot product +void ggml_vec_dot_q1_0_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 ggml_vec_dot_q4_0_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 ggml_vec_dot_q4_1_q8_1(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 ggml_vec_dot_q5_0_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 ggml_vec_dot_q5_1_q8_1(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 ggml_vec_dot_q8_0_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 ggml_vec_dot_mxfp4_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 ggml_vec_dot_nvfp4_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 ggml_vec_dot_q2_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 ggml_vec_dot_q3_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 ggml_vec_dot_q4_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 ggml_vec_dot_q5_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 ggml_vec_dot_q6_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 ggml_vec_dot_tq1_0_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 ggml_vec_dot_tq2_0_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 ggml_vec_dot_iq2_xxs_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 ggml_vec_dot_iq2_xs_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 ggml_vec_dot_iq2_s_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 ggml_vec_dot_iq3_xxs_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 ggml_vec_dot_iq1_s_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 ggml_vec_dot_iq1_m_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 ggml_vec_dot_iq4_nl_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 ggml_vec_dot_iq4_xs_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 ggml_vec_dot_iq3_s_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); + +// Generic implementation +void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void ggml_vec_dot_q1_0_q8_0_generic(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 ggml_vec_dot_q4_0_q8_0_generic(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 ggml_vec_dot_q4_1_q8_1_generic(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 ggml_vec_dot_q5_0_q8_0_generic(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 ggml_vec_dot_q5_1_q8_1_generic(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 ggml_vec_dot_q8_0_q8_0_generic(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 ggml_vec_dot_mxfp4_q8_0_generic(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 ggml_vec_dot_nvfp4_q8_0_generic(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 ggml_vec_dot_tq1_0_q8_K_generic(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 ggml_vec_dot_tq2_0_q8_K_generic(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 ggml_vec_dot_q2_K_q8_K_generic(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 ggml_vec_dot_q3_K_q8_K_generic(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 ggml_vec_dot_q4_K_q8_K_generic(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 ggml_vec_dot_q5_K_q8_K_generic(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 ggml_vec_dot_q6_K_q8_K_generic(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 ggml_vec_dot_iq2_xxs_q8_K_generic(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 ggml_vec_dot_iq2_xs_q8_K_generic(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 ggml_vec_dot_iq2_s_q8_K_generic(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 ggml_vec_dot_iq3_xxs_q8_K_generic(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 ggml_vec_dot_iq3_s_q8_K_generic(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 ggml_vec_dot_iq1_s_q8_K_generic(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 ggml_vec_dot_iq1_m_q8_K_generic(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 ggml_vec_dot_iq4_nl_q8_0_generic(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 ggml_vec_dot_iq4_xs_q8_K_generic(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); + +#ifdef __cplusplus +} +#endif + +#include "ggml-quants-defs.inc" +#include +#include + +#if defined(__gnu_linux__) +#include +#include +#endif + +#if (defined(_WIN32) || defined(_WIN64)) +#define RESTRICT __restrict +#else +#define RESTRICT __restrict__ +#endif + +#if (defined(_WIN32) || defined(_WIN64)) +#define ALWAYS_INLINE __forceinline +#elif __has_attribute(always_inline) || defined(__GNUC__) +#define ALWAYS_INLINE __attribute__((__always_inline__)) inline +#else +#define ALWAYS_INLINE inline +#endif + +#if defined(__AMX_INT8__) && defined(__AVX512VNNI__) + +namespace { + +// Forced unrolling +template +struct Unroll { + template + ALWAYS_INLINE void operator()(const Func& f, Args... args) const { + Unroll{}(f, args...); + f(std::integral_constant{}, args...); + } +}; + +template <> +struct Unroll<1> { + template + ALWAYS_INLINE void operator()(const Func& f, Args... args) const { + f(std::integral_constant{}, args...); + } +}; + +// type traits +template struct PackedTypes {}; +template <> struct PackedTypes { using type = int8_t; }; +template <> struct PackedTypes { using type = uint8_t; }; +template <> struct PackedTypes { using type = int8_t; }; +template using packed_B_type = typename PackedTypes::type; + +template +struct do_compensate : std::integral_constant::value> {}; + +template +struct do_unpack : std::integral_constant::value || + std::is_same::value> {}; + +template +struct is_type_qkk : std::integral_constant::value || + std::is_same::value || + std::is_same::value || + std::is_same::value> {}; + +#define GGML_DISPATCH_FLOATING_TYPES(TYPE, ...) \ + [&] { \ + switch (TYPE) { \ + case GGML_TYPE_F16: { \ + using type = ggml_fp16_t; \ + constexpr int blck_size = 16; \ + return __VA_ARGS__(); \ + } \ + case GGML_TYPE_BF16: { \ + using type = ggml_bf16_t; \ + constexpr int blck_size = 32; \ + return __VA_ARGS__(); \ + } \ + default: \ + fprintf(stderr, "Unsupported floating data type\n"); \ + } \ + }() + +#define GGML_DISPATCH_QTYPES(QT, ...) \ + [&] { \ + switch (QT) { \ + case GGML_TYPE_Q4_0: { \ + using type = block_q4_0; \ + using vec_dot_type = block_q8_0; \ + constexpr int blck_size = QK4_0; \ + return __VA_ARGS__(); \ + } \ + case GGML_TYPE_Q4_1: { \ + using type = block_q4_1; \ + using vec_dot_type = block_q8_1; \ + constexpr int blck_size = QK4_1; \ + return __VA_ARGS__(); \ + } \ + case GGML_TYPE_Q8_0: { \ + using type = block_q8_0; \ + using vec_dot_type = block_q8_0; \ + constexpr int blck_size = QK8_0; \ + return __VA_ARGS__(); \ + } \ + case GGML_TYPE_Q4_K: { \ + using type = block_q4_K; \ + using vec_dot_type = block_q8_K; \ + constexpr int blck_size = QK_K; \ + return __VA_ARGS__(); \ + } \ + case GGML_TYPE_Q5_K: { \ + using type = block_q5_K; \ + using vec_dot_type = block_q8_K; \ + constexpr int blck_size = QK_K; \ + return __VA_ARGS__(); \ + } \ + case GGML_TYPE_Q6_K: { \ + using type = block_q6_K; \ + using vec_dot_type = block_q8_K; \ + constexpr int blck_size = QK_K; \ + return __VA_ARGS__(); \ + } \ + case GGML_TYPE_IQ4_XS: { \ + using type = block_iq4_xs; \ + using vec_dot_type = block_q8_K; \ + constexpr int blck_size = QK_K; \ + return __VA_ARGS__(); \ + } \ + default: \ + fprintf(stderr, "Unsupported quantized data type: %d\n", int(TYPE)); \ + } \ + }() + +#define GGML_DISPATCH_BOOL(BOOL_V, BOOL_NAME, ...) \ + [&] { \ + if (BOOL_V) { \ + constexpr bool BOOL_NAME = true; \ + return __VA_ARGS__(); \ + } else { \ + constexpr bool BOOL_NAME = false; \ + return __VA_ARGS__(); \ + } \ + }() + +// define amx tile config data structure +struct tile_config_t{ + uint8_t palette_id = 0; + uint8_t start_row = 0; + uint8_t reserved_0[14] = {0}; + uint16_t colsb[16] = {0}; + uint8_t rows[16] = {0}; +}; + +// Notes: amx tile config +// +// Typically, TMUL calculates A and B of size 16 x 64 containing INT8 values, +// and accumulate the result to a 16 x 16 matrix C containing INT32 values, +// +// As many GGUF quantized types as `block_size` of 32, so a 16-16-32 config is used +// instead of the normally used 16-16-64 config. +// +// Block A: {16, 32}, dtype = int8_t +// Block B: {16, 32}, dtype = uint8_t/int8_t +// Block C: {16, 16}, dtype = int32_t +// +// Block B needs to be prepacked to vnni format before feeding into TMUL: +// packed_B: from {n, k} to {k/vnni_blk, n, vnni_blck}, viewed in 2d, we get {8, 64} +// +// Therefore, we get tileconfig: +// A B C +// rows 16 8 16 +// colsb 32 64 16 +// +// For tile distribution, follow a 2-2-4 pattern, e.g. A used TMM2-TMM3, B used TMM0-TMM1, +// C used TMM4-TMM7: +// B TMM0 B TMM1 +// A TMM2 C TMM4 C TMM6 +// A TMM3 C TMM5 C TMM7 +// +// Each `amx` kernel handles 4 blocks at a time: 2MB * 2NB, when m < 2 * BLOCK_M, unpack A +// will be needed. +// +// Here another commonly used pattern 1-3-3 is skipped, as it is mostly used when m <=16; +// and the single batch gemm (m=1) has a special fast path with `avx512-vnni`. +// +// ref: https://www.intel.com/content/www/us/en/developer/articles/code-sample/ +// advanced-matrix-extensions-intrinsics-functions.html +// + +inline void ggml_tile_config_init(void) { + static thread_local bool done = false; + + if (done) { + return; + } + + alignas(64) tile_config_t tc = {}; + tc.palette_id = 1; + tc.start_row = 0; + tc.rows[0] = 8; tc.colsb[0] = 64; + tc.rows[1] = 8; tc.colsb[1] = 64; + tc.rows[2] = 16; tc.colsb[2] = 32; + tc.rows[3] = 16; tc.colsb[3] = 32; + tc.rows[4] = 16; tc.colsb[4] = 64; + tc.rows[5] = 16; tc.colsb[5] = 64; + tc.rows[6] = 16; tc.colsb[6] = 64; + tc.rows[7] = 16; tc.colsb[7] = 64; + + _tile_loadconfig(&tc); + done = true; +} + +// we need an extra 16 * 4B (TILE_N * int32_t) for each NB/KB block for compensation. +// See the notes `s8s8 igemm compensation in avx512-vnni` for detail. +template +int get_tile_size() { + int tile_size = TILE_N * sizeof(TB); + if (do_compensate::value) { + tile_size += TILE_N * sizeof(int32_t); + } + if (std::is_same::value || + std::is_same::value) { + tile_size += TILE_N * 4; + } + if (std::is_same::value) { + tile_size += TILE_N * 2; + } + return tile_size; +} + +template +int get_row_size(int K) { + int KB = K / BLOCK_K; + int row_size = KB * sizeof(TB); + if (do_compensate::value) { + row_size += KB * sizeof(int32_t); + } + if (std::is_same::value || + std::is_same::value) { + row_size += KB * 4; + } + if (std::is_same::value) { + row_size += KB * 2; + } + return row_size; +} + +// transpose utils +#define SHUFFLE_EPI32(a, b, mask) \ + _mm256_castps_si256(_mm256_shuffle_ps(_mm256_castsi256_ps(a), _mm256_castsi256_ps(b), mask)) +inline void transpose_8x8_32bit(__m256i * v, __m256i * v1) { + // unpacking and 32-bit elements + v1[0] = _mm256_unpacklo_epi32(v[0], v[1]); + v1[1] = _mm256_unpackhi_epi32(v[0], v[1]); + v1[2] = _mm256_unpacklo_epi32(v[2], v[3]); + v1[3] = _mm256_unpackhi_epi32(v[2], v[3]); + v1[4] = _mm256_unpacklo_epi32(v[4], v[5]); + v1[5] = _mm256_unpackhi_epi32(v[4], v[5]); + v1[6] = _mm256_unpacklo_epi32(v[6], v[7]); + v1[7] = _mm256_unpackhi_epi32(v[6], v[7]); + + // shuffling the 32-bit elements + v[0] = SHUFFLE_EPI32(v1[0], v1[2], 0x44); + v[1] = SHUFFLE_EPI32(v1[0], v1[2], 0xee); + v[2] = SHUFFLE_EPI32(v1[4], v1[6], 0x44); + v[3] = SHUFFLE_EPI32(v1[4], v1[6], 0xee); + v[4] = SHUFFLE_EPI32(v1[1], v1[3], 0x44); + v[5] = SHUFFLE_EPI32(v1[1], v1[3], 0xee); + v[6] = SHUFFLE_EPI32(v1[5], v1[7], 0x44); + v[7] = SHUFFLE_EPI32(v1[5], v1[7], 0xee); + + // shuffling 128-bit elements + v1[0] = _mm256_permute2f128_si256(v[2], v[0], 0x02); + v1[1] = _mm256_permute2f128_si256(v[3], v[1], 0x02); + v1[2] = _mm256_permute2f128_si256(v[6], v[4], 0x02); + v1[3] = _mm256_permute2f128_si256(v[7], v[5], 0x02); + v1[4] = _mm256_permute2f128_si256(v[2], v[0], 0x13); + v1[5] = _mm256_permute2f128_si256(v[3], v[1], 0x13); + v1[6] = _mm256_permute2f128_si256(v[6], v[4], 0x13); + v1[7] = _mm256_permute2f128_si256(v[7], v[5], 0x13); +} + +inline void transpose_16x4_32bit(__m512i * r, __m512i * d) { + + static const __m512i index1 = _mm512_set_epi32( + 0x0f, 0x0b, 0x07, 0x03, + 0x0e, 0x0a, 0x06, 0x02, + 0x0d, 0x09, 0x05, 0x01, + 0x0c, 0x08, 0x04, 0x00); + + d[0] = _mm512_permutexvar_epi32(index1, r[0]); + d[1] = _mm512_permutexvar_epi32(index1, r[1]); + d[2] = _mm512_permutexvar_epi32(index1, r[2]); + d[3] = _mm512_permutexvar_epi32(index1, r[3]); + + r[0] = _mm512_shuffle_i32x4(d[0], d[1], 0x44); + r[1] = _mm512_shuffle_i32x4(d[0], d[1], 0xee); + r[2] = _mm512_shuffle_i32x4(d[2], d[3], 0x44); + r[3] = _mm512_shuffle_i32x4(d[2], d[3], 0xee); + + d[0] = _mm512_shuffle_i32x4(r[0], r[2], 0x88); + d[1] = _mm512_shuffle_i32x4(r[0], r[2], 0xdd); + d[2] = _mm512_shuffle_i32x4(r[1], r[3], 0x88); + d[3] = _mm512_shuffle_i32x4(r[1], r[3], 0xdd); +} + +inline void transpose_16x16_32bit(__m512i * v) { + __m512i v1[16]; + v1[0] = _mm512_unpacklo_epi32(v[0], v[1]); + v1[1] = _mm512_unpackhi_epi32(v[0], v[1]); + v1[2] = _mm512_unpacklo_epi32(v[2], v[3]); + v1[3] = _mm512_unpackhi_epi32(v[2], v[3]); + v1[4] = _mm512_unpacklo_epi32(v[4], v[5]); + v1[5] = _mm512_unpackhi_epi32(v[4], v[5]); + v1[6] = _mm512_unpacklo_epi32(v[6], v[7]); + v1[7] = _mm512_unpackhi_epi32(v[6], v[7]); + v1[8] = _mm512_unpacklo_epi32(v[8], v[9]); + v1[9] = _mm512_unpackhi_epi32(v[8], v[9]); + v1[10] = _mm512_unpacklo_epi32(v[10], v[11]); + v1[11] = _mm512_unpackhi_epi32(v[10], v[11]); + v1[12] = _mm512_unpacklo_epi32(v[12], v[13]); + v1[13] = _mm512_unpackhi_epi32(v[12], v[13]); + v1[14] = _mm512_unpacklo_epi32(v[14], v[15]); + v1[15] = _mm512_unpackhi_epi32(v[14], v[15]); + + v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]); + v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]); + v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]); + v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]); + v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]); + v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]); + v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]); + v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]); + v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]); + v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]); + v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]); + v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]); + v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]); + v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]); + v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]); + v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]); + + v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88); + v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88); + v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88); + v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88); + v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd); + v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd); + v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd); + v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd); + v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88); + v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88); + v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88); + v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88); + v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd); + v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd); + v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd); + v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd); + + v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88); + v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88); + v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88); + v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88); + v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88); + v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88); + v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88); + v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88); + v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd); + v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd); + v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd); + v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd); + v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd); + v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd); + v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd); + v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd); +} + +void quantize_row_q8_K_vnni(const float * RESTRICT x, void * RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + const int KB = k / QK_K; + constexpr int kVecs = QK_K / 16; + + block_q8_K * y = reinterpret_cast(vy); + + // hold 16 float vecs from x + __m512 v[kVecs]; + + // hold the quants vecs + __m512i vq[kVecs / 4]; + + // hold the packed quants vecs + __m512i vq_packed[kVecs / 4]; + + const __m512 signBit = _mm512_set1_ps(-0.f); + + for (int i = 0; i < KB; ++i) { + // Compute max(abs(e)) for the block + __m512 vamax = _mm512_set1_ps(0.f); + for (int j = 0; j < kVecs; ++j) { + v[j] = _mm512_loadu_ps(x); x += 16; + vamax = _mm512_max_ps(vamax, _mm512_andnot_ps(signBit, v[j])); + } + const float amax = _mm512_reduce_max_ps(vamax); + + // Quantize these floats + const float iscale = 127.f / amax; + y[i].d = GGML_CPU_FP32_TO_FP16(1 / iscale); + const float id = ( amax != 0.0f ) ? iscale : 0.f; + const __m512 vscale = _mm512_set1_ps(id); + + // Apply multiplier and round to nearest integer + for (int j = 0; j < kVecs; ++j) { + v[j] = _mm512_mul_ps(v[j], vscale); + v[j] = _mm512_roundscale_ps(v[j], (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); + } + + // Pack to epi8 vecs + for (int j = 0; j < kVecs / 4; ++j) { + __m128i q8_0 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 0])); + __m128i q8_1 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 1])); + __m128i q8_2 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 2])); + __m128i q8_3 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 3])); + + __m256i q8_01 = _mm256_insertf128_si256(_mm256_castsi128_si256(q8_0), (q8_1), 1); + __m256i q8_23 = _mm256_insertf128_si256(_mm256_castsi128_si256(q8_2), (q8_3), 1); + + vq[j] = _mm512_inserti32x8(_mm512_castsi256_si512(q8_01), q8_23, 1); + _mm512_storeu_si512((__m512i *)(y[i].qs + j * 64), vq[j]); + } + + // Compute the bsums with vnni + transpose_16x4_32bit(vq, vq_packed); + + const __m512i one = _mm512_set1_epi8(1); + __m512i sum = _mm512_setzero_si512(); + for (int k = 0; k < 4; ++k) { + sum = _mm512_dpbusd_epi32(sum, one, vq_packed[k]); + } + _mm256_storeu_si256((__m256i *)(y[i].bsums), _mm512_cvtepi32_epi16(sum)); + } +} + +// quantize A from float to `vec_dot_type` +template +inline void from_float(const float * x, char * vy, int64_t k); + +template <> +inline void from_float(const float * x, char * vy, int64_t k) { + quantize_row_q8_0(x, (block_q8_0 *)vy, k); +} + +template <> +inline void from_float(const float * x, char * vy, int64_t k) { + quantize_row_q8_1(x, (block_q8_1 *)vy, k); +} + +template <> +inline void from_float(const float * x, char * vy, int64_t k) { +#if 1 + // TODO: this is reference impl! + quantize_row_q8_K_ref(x, (block_q8_K *)vy, k); +#else + quantize_row_q8_K_vnni(x, vy, k); +#endif +} + +// load A from memory to array when nrows can not fill in whole tile +void unpack_A(int8_t * RESTRICT tile, const block_q8_0 * RESTRICT A, int lda, int nr) { + assert(nr != TILE_M); + for (int m = 0; m < nr; ++m) { + const __m256i v = _mm256_loadu_si256((const __m256i *)(A[m * lda].qs)); + _mm256_storeu_si256((__m256i *)(tile + m * TILE_K), v); + } +} + +void unpack_A(int8_t * RESTRICT tile, const block_q8_1 * RESTRICT A, int lda, int nr) { + assert(nr != TILE_M); + for (int m = 0; m < nr; ++m) { + const __m256i v = _mm256_loadu_si256((const __m256i *)(A[m * lda].qs)); + _mm256_storeu_si256((__m256i *)(tile + m * TILE_K), v); + } +} + +template +void unpack_A(int8_t * RESTRICT tile, const block_q8_K * RESTRICT A, int lda, int k, int nr) { + assert(nr <= TILE_M); + for (int m = 0; m < nr; ++m) { + const __m256i v = _mm256_loadu_si256((const __m256i *)(A[m * lda].qs + k * 32)); + _mm256_storeu_si256((__m256i *)(tile + m * TILE_K), v); + } +} + +template <> +void unpack_A(int8_t * RESTRICT tile, const block_q8_K * RESTRICT A, int lda, int k, int nr) { + assert(nr <= TILE_M); + // zero padding k from 16 to 32, so that we don't have to re-config amx + const __m128i zero = _mm_setzero_si128(); + for (int m = 0; m < nr; ++m) { + const __m128i v = _mm_loadu_si128((const __m128i *)(A[m * lda].qs + k * 16)); + const __m256i r = _mm256_insertf128_si256(_mm256_castsi128_si256(v), zero, 1); + _mm256_storeu_si256((__m256i *)(tile + m * TILE_K), r); + } +} + +#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1) +inline __m256i bytes_from_nibbles_32(const uint8_t * rsi) { + const __m128i tmp = _mm_loadu_si128((const __m128i *)rsi); + const __m256i bytes = MM256_SET_M128I(_mm_srli_epi16(tmp, 4), tmp); + const __m256i lowMask = _mm256_set1_epi8(0xF); + return _mm256_and_si256(lowMask, bytes); +} + +// used for block_q4_K +inline __m512i bytes_from_nibbles_64(const uint8_t * rsi) { + const __m256i tmp = _mm256_loadu_si256((const __m256i *)rsi); + const __m256i lowMask = _mm256_set1_epi8(0xF); + const __m256i q4l = _mm256_and_si256(tmp, lowMask); + const __m256i q4h = _mm256_and_si256(_mm256_srli_epi16(tmp, 4), lowMask); + return _mm512_inserti32x8(_mm512_castsi256_si512(q4l), q4h, 1); +} + +// used for block_q5_K +inline __m512i bytes_from_nibbles_64(const uint8_t * qs, const uint8_t * qh, int k) { + const __m256i lowMask = _mm256_set1_epi8(0xF); + __m256i hmask = _mm256_set1_epi8(1); + hmask = _mm256_slli_epi16(hmask, k); + + const __m256i q5bits = _mm256_loadu_si256((const __m256i *)qs); + const __m256i hbits = _mm256_loadu_si256((const __m256i *)qh); + + const __m256i q5l_0 = _mm256_and_si256(q5bits, lowMask); + const __m256i q5h_0 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_and_si256(hbits, hmask), k + 0), 4); + const __m256i q5_0 = _mm256_add_epi8(q5l_0, q5h_0); + hmask = _mm256_slli_epi16(hmask, 1); + + const __m256i q5l_1 = _mm256_and_si256(_mm256_srli_epi16(q5bits, 4), lowMask); + const __m256i q5h_1 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_and_si256(hbits, hmask), k + 1), 4); + const __m256i q5_1 = _mm256_add_epi8(q5l_1, q5h_1); + + return _mm512_inserti32x8(_mm512_castsi256_si512(q5_0), q5_1, 1); +} + +// used for block_q6_K +inline void bytes_from_nibbles_128(__m512i& r0, __m512i& r1, const uint8_t * qs, const uint8_t * qh) { + const __m256i m4 = _mm256_set1_epi8(0xF); + const __m256i m2 = _mm256_set1_epi8(0x3); + + const __m256i q6bits1 = _mm256_loadu_si256((const __m256i *)qs); + const __m256i q6bits2 = _mm256_loadu_si256((const __m256i *)(qs + 32)); + const __m256i q6bitsH = _mm256_loadu_si256((const __m256i *)qh); + + const __m256i q6h_0 = _mm256_slli_epi16(_mm256_and_si256( q6bitsH, m2), 4); + const __m256i q6h_1 = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(q6bitsH, 2), m2), 4); + const __m256i q6h_2 = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(q6bitsH, 4), m2), 4); + const __m256i q6h_3 = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(q6bitsH, 6), m2), 4); + + const __m256i q6_0 = _mm256_or_si256(_mm256_and_si256(q6bits1, m4), q6h_0); + const __m256i q6_1 = _mm256_or_si256(_mm256_and_si256(q6bits2, m4), q6h_1); + const __m256i q6_2 = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(q6bits1, 4), m4), q6h_2); + const __m256i q6_3 = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(q6bits2, 4), m4), q6h_3); + + r0 = _mm512_inserti32x8(_mm512_castsi256_si512(q6_0), q6_1, 1); + r1 = _mm512_inserti32x8(_mm512_castsi256_si512(q6_2), q6_3, 1); +} + +inline __m512i packNibbles(__m512i r0, __m512i r1) { + return _mm512_or_si512(r0, _mm512_slli_epi16(r1, 4)); +} + +template +inline void pack_qs(void * RESTRICT packed_B, const TB * RESTRICT B, int KB) { + int8_t tmp[8 * 64]; + __m256i v[8], v2[8]; + for (int n = 0; n < 8; ++n) { + v[n] = bytes_from_nibbles_32(B[n * KB].qs); + } + transpose_8x8_32bit(v, v2); + for (int n = 0; n < 8; ++n) { + _mm256_storeu_si256((__m256i *)(tmp + n * 64), v2[n]); + } + for (int n = 0; n < 8; ++n) { + v[n] = bytes_from_nibbles_32(B[(n + 8) * KB].qs); + } + transpose_8x8_32bit(v, v2); + for (int n = 0; n < 8; ++n) { + _mm256_storeu_si256((__m256i *)(tmp + n * 64 + 32), v2[n]); + } + + // pack again with 128 to fully utilize vector length + for (int n = 0; n < 8; n += 2) { + __m512i r0 = _mm512_loadu_si512((const __m512i *)(tmp + n * 64)); + __m512i r1 = _mm512_loadu_si512((const __m512i *)(tmp + n * 64 + 64)); + __m512i r1r0 = packNibbles(r0, r1); + _mm512_storeu_si512((__m512i *)((char *)packed_B + n * 32), r1r0); + } +} + +template <> +inline void pack_qs(void * RESTRICT packed_B, const block_q8_0 * RESTRICT B, int KB) { + __m256i v[8], v2[8]; + for (int n = 0; n < 8; ++n) { + v[n] = _mm256_loadu_si256((const __m256i *)(B[n * KB].qs)); + } + transpose_8x8_32bit(v, v2); + for (int n = 0; n < 8; ++n) { + _mm256_storeu_si256((__m256i *)((char *)packed_B + n * 64), v2[n]); + } + for (int n = 0; n < 8; ++n) { + v[n] = _mm256_loadu_si256((const __m256i *)(B[(n + 8) * KB].qs)); + } + transpose_8x8_32bit(v, v2); + for (int n = 0; n < 8; ++n) { + _mm256_storeu_si256((__m256i *)((char *)packed_B + n * 64 + 32), v2[n]); + } +} + +template <> +inline void pack_qs(void * RESTRICT packed_B, const block_q4_K * RESTRICT B, int KB) { + __m512i v[16]; + // QK_K 256 with 8 groups, handle 2 groups at a time + char * pb = (char *)packed_B; + for (int k = 0; k < QK_K / 64; ++k) { + // pack 2 groups { n, g, k} to {g, k/4, 4n} + // e.g. {16, 2, 32} to {2, 8, 64} + for (int n = 0; n < TILE_N; ++n) { + v[n] = bytes_from_nibbles_64(B[n * KB].qs + k * 32); + } + + transpose_16x16_32bit(v); + + // pack again with 128 to fully utilize vector length + for (int n = 0; n < TILE_N; n += 2) { + _mm512_storeu_si512((__m512i *)pb, packNibbles(v[n], v[n + 1])); + pb += 64; + } + } +} + +template <> +inline void pack_qs(void * RESTRICT packed_B, const block_q5_K * RESTRICT B, int KB) { + __m512i v[16]; + const __m512i lowMask = _mm512_set1_epi8(0xF); + // QK_K 256 with 8 groups, handle 2 groups at a time + char * pb = (char *)packed_B; + char * ph = (char *)packed_B + (QK_K / 2) * TILE_N; + for (int k = 0; k < QK_K / 64; ++k) { + // pack 2 groups { n, g, k} to {g, k/4, 4n} + // e.g. {16, 2, 32} to {2, 8, 64} + for (int n = 0; n < TILE_N; ++n) { + v[n] = bytes_from_nibbles_64(B[n * KB].qs + k * 32, B[n * KB].qh, /* group */2 * k); + } + + transpose_16x16_32bit(v); + + // 1. pack lower 4bits with 2 groups + for (int n = 0; n < TILE_N; n += 2) { + // get lower 4 bits + const __m512i r0 = _mm512_and_si512(v[n], lowMask); + const __m512i r1 = _mm512_and_si512(v[n + 1], lowMask); + _mm512_storeu_si512((__m512i *)pb, packNibbles(r0, r1)); pb += 64; + } + + // 2. pack higher 1bit with 2 groups + const __m512i hmask = _mm512_set1_epi8(0x10); + for (int g = 0; g < 2; ++g) { + __m512i hbits = _mm512_setzero_si512(); + hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 0], hmask), 4)); + hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 1], hmask), 3)); + hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 2], hmask), 2)); + hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 3], hmask), 1)); + hbits = _mm512_add_epi8(hbits, _mm512_and_si512(v[g * 8 + 4], hmask) ); + hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 8 + 5], hmask), 1)); + hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 8 + 6], hmask), 2)); + hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 8 + 7], hmask), 3)); + _mm512_storeu_si512((__m512i *)ph, hbits); ph += 64; + } + } +} + +template <> +inline void pack_qs(void * RESTRICT packed_B, const block_q6_K * RESTRICT B, int KB) { + __m512i v[32]; + const __m512i lowMask = _mm512_set1_epi8(0xF); + // QK_K 256 with 8 groups, handle 4 groups at a time + char * pb = (char *)packed_B; + char * ph = (char *)packed_B + (QK_K / 2) * TILE_N; + for (int k = 0; k < QK_K / 128; ++k) { + for (int n = 0; n < TILE_N; ++n) { + bytes_from_nibbles_128(v[n], v[n + 16], B[n * KB].ql + k * 64, B[n * KB].qh + k * 32); + } + + // top half: group 0,1 or 4,5; bottom half: group 2,3 or 6,7 + transpose_16x16_32bit(v); + transpose_16x16_32bit(v + 16); + + // 1. pack lower 4bits with 4 groups + for (int n = 0; n < 32; n += 2) { + const __m512i r0 = _mm512_and_si512(v[n], lowMask); + const __m512i r1 = _mm512_and_si512(v[n + 1], lowMask); + _mm512_storeu_si512((__m512i *)pb, packNibbles(r0, r1)); pb += 64; + } + + // 2. pack higher 2bit with 4 groups + const __m512i hmask = _mm512_set1_epi8(0x30); + for (int g = 0; g < 8; ++g) { + __m512i hbits = _mm512_setzero_si512(); + hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 4 + 0], hmask), 4)); + hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 4 + 1], hmask), 2)); + hbits = _mm512_add_epi8(hbits, _mm512_and_si512(v[g * 4 + 2], hmask) ); + hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 4 + 3], hmask), 2)); + _mm512_storeu_si512((__m512i *)ph, hbits); ph += 64; + } + } +} + +template <> +inline void pack_qs(void * RESTRICT packed_B, const block_iq4_xs * RESTRICT B, int KB) { + __m512i v[16]; + char * pb = (char *)packed_B; + for (int k = 0; k < QK_K / 64; ++k) { + for (int n = 0; n < TILE_N; ++n) { + __m256i r0 = bytes_from_nibbles_32(B[n * KB].qs + k * 32 + 0); + __m256i r1 = bytes_from_nibbles_32(B[n * KB].qs + k * 32 + 16); + v[n] = _mm512_inserti32x8(_mm512_castsi256_si512(r0), r1, 1); + } + + transpose_16x16_32bit(v); + + // pack again with 128 to fully utilize vector length + for (int n = 0; n < TILE_N; n += 2) { + _mm512_storeu_si512((__m512i *)pb, packNibbles(v[n], v[n + 1])); + pb += 64; + } + } +} + +// pack B to vnni formats in 4bits or 8 bits +void pack_B(void * RESTRICT packed_B, const block_q4_0 * RESTRICT B, int KB) { + pack_qs(packed_B, B, KB); + ggml_half * d0 = reinterpret_cast((char *)packed_B + TILE_N * TILE_K / 2); + for (int n = 0; n < TILE_N; ++n) { + d0[n] = B[n * KB].d; + } +} + +void pack_B(void * RESTRICT packed_B, const block_q4_1 * RESTRICT B, int KB) { + pack_qs(packed_B, B, KB); + ggml_half * d0 = reinterpret_cast((char *)packed_B + TILE_N * TILE_K / 2); + ggml_half * m0 = d0 + TILE_N; + for (int n = 0; n < TILE_N; ++n) { + d0[n] = B[n * KB].d; + m0[n] = B[n * KB].m; + } +} + +inline void s8s8_compensation(void * RESTRICT packed_B) { + // packed_B layout: + // quants {TILE_N, TILEK} int8_t + // d0 {TILE_N} ggml_half + // comp {TILE_N} int32_t + const int offset = TILE_N * TILE_K + TILE_N * sizeof(ggml_half); + __m512i vcomp = _mm512_setzero_si512(); + const __m512i off = _mm512_set1_epi8(static_cast(0x80)); + for (int k = 0; k < 8; ++k) { + __m512i vb = _mm512_loadu_si512((const __m512i *)((const char *)packed_B + k * 64)); + vcomp = _mm512_dpbusd_epi32(vcomp, off, vb); + } + _mm512_storeu_si512((__m512i *)((char *)(packed_B) + offset), vcomp); +} + +void pack_B(void * RESTRICT packed_B, const block_q8_0 * RESTRICT B, int KB) { + pack_qs(packed_B, B, KB); + ggml_half * d0 = reinterpret_cast((char *)packed_B + TILE_N * TILE_K); + for (int n = 0; n < TILE_N; ++n) { + d0[n] = B[n * KB].d; + } + s8s8_compensation(packed_B); +} + +// convert 8 * {min, scale} from int6 to int8 +inline void unpack_mins_and_scales(const uint8_t * scales, uint32_t * utmp) { + const uint32_t kmask1 = 0x3f3f3f3f; + const uint32_t kmask2 = 0x0f0f0f0f; + const uint32_t kmask3 = 0x03030303; + + memcpy(utmp, scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; +} + +// packed_B layout: +// quants {8, TILE_N, 16} uint8 +// scales {8, TILE_N} uint8 +// mins {8, TILE_N} uint8 +// d {TILE_N} ggml_half +// dmin {TILE_N} ggml_half +void pack_B(void * RESTRICT packed_B, const block_q4_K * RESTRICT B, int KB) { + pack_qs(packed_B, B, KB); + + uint8_t * scales = reinterpret_cast((char *)packed_B + (QK_K / 2) * TILE_N); + uint8_t * mins = scales + 8 * TILE_N; + ggml_half * d = reinterpret_cast(mins + 8 * TILE_N); + ggml_half * dmin = d + TILE_N; + + union { + uint32_t u32[4]; + uint8_t u8[16]; + } s; + + for (int n = 0; n < TILE_N; ++n) { + unpack_mins_and_scales(B[n * KB].scales, s.u32); + for (int k = 0; k < 8; ++k) { + scales[k * TILE_N + n] = s.u8[k]; + mins[(k >> 1) * TILE_N * 2 + n * 2 + (k & 0x1)] = s.u8[k + 8]; + } + d[n] = B[n * KB].d; + dmin[n] = B[n * KB].dmin; + } +} + +// packed_B layout: +// quants {8, TILE_N, 16} uint8 +// qh {8, TILE_N, 4} uint8 +// scales {8, TILE_N} uint8 +// mins {8, TILE_N} uint8 +// d {TILE_N} ggml_half +// dmin {TILE_N} ggml_half +void pack_B(void * RESTRICT packed_B, const block_q5_K * RESTRICT B, int KB) { + pack_qs(packed_B, B, KB); + + uint8_t * scales = reinterpret_cast((char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N); + uint8_t * mins = scales + 8 * TILE_N; + ggml_half * d = reinterpret_cast(mins + 8 * TILE_N); + ggml_half * dmin = d + TILE_N; + + union { + uint32_t u32[4]; + uint8_t u8[16]; + } s; + + for (int n = 0; n < TILE_N; ++n) { + unpack_mins_and_scales(B[n * KB].scales, s.u32); + for (int k = 0; k < 8; ++k) { + scales[k * TILE_N + n] = s.u8[k]; + mins[(k >> 1) * TILE_N * 2 + n * 2 + (k & 0x1)] = s.u8[k + 8]; + } + d[n] = B[n * KB].d; + dmin[n] = B[n * KB].dmin; + } +} + +// packed_B layout: +// quants {16, TILE_N, 8} uint8 +// qh {16, TILE_N, 4} uint8 +// scales {16, TILE_N} uint8 +// d {TILE_N} ggml_half +void pack_B(void * RESTRICT packed_B, const block_q6_K * RESTRICT B, int KB) { + pack_qs(packed_B, B, KB); + + uint8_t * scales = reinterpret_cast((char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N); + ggml_half * d = reinterpret_cast(scales + 16 * TILE_N); + for (int n = 0; n < TILE_N; ++n) { + const int8_t * ps = B[n * KB].scales; + for (int k = 0; k < 16; ++k) { + scales[k * TILE_N + n] = ps[k]; + } + d[n] = B[n * KB].d; + } +} + +// packed_B layout: +// quants {8, TILE_N, 16} uint8 +// scales {8, TILE_N} int8 +// d {TILE_N} ggml_half +void pack_B(void * RESTRICT packed_B, const block_iq4_xs * RESTRICT B, int KB) { + pack_qs(packed_B, B, KB); + + int8_t * scales = reinterpret_cast((char *)packed_B + (QK_K / 2) * TILE_N); + ggml_half * d = reinterpret_cast(scales + 8 * TILE_N); + + // pack the scales + for (int n = 0; n < TILE_N; ++n) { + uint16_t sh = B[n * KB].scales_h; + for (int k = 0; k < 8; k += 2) { + const int16_t ls1 = ((B[n * KB].scales_l[k / 2] & 0xf) | ((sh << 4) & 0x30)) - 32; + const int16_t ls2 = ((B[n * KB].scales_l[k / 2] >> 4) | ((sh << 2) & 0x30)) - 32; + scales[(k + 0) * TILE_N + n] = ls1; + scales[(k + 1) * TILE_N + n] = ls2; + sh >>= 4; + } + d[n] = B[n * KB].d; + } +} + +template> +void unpack_B(packed_B_t * RESTRICT tile, const void * RESTRICT packed_B) { + GGML_UNUSED(tile); + GGML_UNUSED(packed_B); +} + +template <> +void unpack_B(int8_t * RESTRICT tile, const void * RESTRICT packed_B) { + const __m512i off = _mm512_set1_epi8(8); + const __m512i lowMask = _mm512_set1_epi8(0xF); + for (int n = 0; n < 8; n += 2) { + __m512i bytes = _mm512_loadu_si512((const __m512i *)((const char *)packed_B + n * 32)); + const __m512i r0 = _mm512_sub_epi8(_mm512_and_si512(bytes, lowMask), off); + const __m512i r1 = _mm512_sub_epi8(_mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask), off); + _mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0); + _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1); + } +} + +template <> +void unpack_B(uint8_t * RESTRICT tile, const void * RESTRICT packed_B) { + const __m512i lowMask = _mm512_set1_epi8(0xF); + for (int n = 0; n < 8; n += 2) { + __m512i bytes = _mm512_loadu_si512((const __m512i *)((const char *)packed_B + n * 32)); + const __m512i r0 = _mm512_and_si512(bytes, lowMask); + const __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); + _mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0); + _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1); + } +} + +// packed_B_t for QKK is int8_t +template +void unpack_B(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) { + const int packed_B_group_size = QK_K / 2 * TILE_N / 8; + const char * packed_B_group = (const char *)packed_B + k * packed_B_group_size; + const __m512i lowMask = _mm512_set1_epi8(0xF); + for (int n = 0; n < 8; n += 2) { + __m512i bytes = _mm512_loadu_si512(packed_B_group + n * 32); + const __m512i r0 = _mm512_and_si512(bytes, lowMask); + const __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); + _mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0); + _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1); + } +} + +template <> +void unpack_B(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) { + // lower 4bits, stride 256 bytes + const int packed_l4_group_size = QK_K / 2 * TILE_N / 8; + const char * pb = (const char *)packed_B + k * packed_l4_group_size; + + // higher 1bit, stride 64 bytes + const int packed_h1_group_size = QK_K / 8 * TILE_N / 8; + const char * ph = (const char *)packed_B + (QK_K / 2) * TILE_N + k * packed_h1_group_size; + const __m512i hbits = _mm512_loadu_si512(ph); + + const __m512i lowMask = _mm512_set1_epi8(0xF); + __m512i hmask0 = _mm512_set1_epi8(0x1); + __m512i hmask1 = _mm512_set1_epi8(0x2); + + for (int n = 0; n < 8; n += 2) { + __m512i bytes = _mm512_loadu_si512(pb + n * 32); + __m512i r0 = _mm512_and_si512(bytes, lowMask); + __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); + __m512i h0 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask0), n), 4); + __m512i h1 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask1), n + 1), 4); + + hmask0 = _mm512_slli_epi16(hmask0, 2); + hmask1 = _mm512_slli_epi16(hmask1, 2); + r0 = _mm512_add_epi8(r0, h0); + r1 = _mm512_add_epi8(r1, h1); + _mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0); + _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1); + } +} + +template <> +void unpack_B(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) { + // lower 4bits, stride 128 bytes + const int packed_l4_group_size = QK_K / 2 * TILE_N / 16; + const char * pb = (const char *)packed_B + k * packed_l4_group_size; + + // higher 2bits, stride 64 bytes + const int packed_h2_group_size = QK_K / 4 * TILE_N / 16; + const char * ph = (const char *)packed_B + (QK_K / 2) * TILE_N + k * packed_h2_group_size; + const __m512i hbits = _mm512_loadu_si512(ph); + + const __m512i off = _mm512_set1_epi8(32); + const __m512i lowMask = _mm512_set1_epi8(0xF); + __m512i hmask0 = _mm512_set1_epi8(0x3); // 0011 + __m512i hmask1 = _mm512_set1_epi8(0xC); // 1100 + + // notes: skip zero padding from row4 to row7 as we have done so in `unpack_A` + __m512i bytes = _mm512_loadu_si512(pb); + __m512i r0 = _mm512_and_si512(bytes, lowMask); + __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); + __m512i h0 = _mm512_slli_epi16(_mm512_and_si512(hbits, hmask0), 4); + __m512i h1 = _mm512_slli_epi16(_mm512_and_si512(hbits, hmask1), 2); + _mm512_storeu_si512((__m512i *)(tile + 0), _mm512_sub_epi8(_mm512_add_epi8(r0, h0), off)); + _mm512_storeu_si512((__m512i *)(tile + 64), _mm512_sub_epi8(_mm512_add_epi8(r1, h1), off)); + + hmask0 = _mm512_slli_epi16(hmask0, 4); + hmask1 = _mm512_slli_epi16(hmask1, 4); + + bytes = _mm512_loadu_si512(pb + 64); + r0 = _mm512_and_si512(bytes, lowMask); + r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); + h0 = _mm512_and_si512(hbits, hmask0); + h1 = _mm512_srli_epi16(_mm512_and_si512(hbits, hmask1), 2); + _mm512_storeu_si512((__m512i *)(tile + 128), _mm512_sub_epi8(_mm512_add_epi8(r0, h0), off)); + _mm512_storeu_si512((__m512i *)(tile + 192), _mm512_sub_epi8(_mm512_add_epi8(r1, h1), off)); +} + +template <> +void unpack_B(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) { + static const __m512i values128 = _mm512_set_epi8( + 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127, + 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127, + 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127, + 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127 + ); + + const int packed_B_group_size = QK_K / 2 * TILE_N / 8; + const char * pb = (const char *)packed_B + k * packed_B_group_size; + const __m512i lowMask = _mm512_set1_epi8(0xF); + + for (int n = 0; n < 8; n += 2) { + __m512i bytes = _mm512_loadu_si512(pb + n * 32); + const __m512i r0 = _mm512_shuffle_epi8(values128, _mm512_and_si512(bytes, lowMask)); + const __m512i r1 = _mm512_shuffle_epi8(values128, _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask)); + _mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0); + _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1); + } +} + +template +struct acc_C {}; + +template +struct acc_C { + static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_0 * A, int lda, const void * packed_B, int nr) { + const int offset = TILE_N * TILE_K / 2; + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset))); + + for (int m = 0; m < nr; ++m) { + const __m512 vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[m * lda].d)); + const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); + + __m512 vsum; + if (is_acc) { + vsum = _mm512_loadu_ps(C + m * ldc); + } else { + vsum = _mm512_set1_ps(0.f); + } + vsum = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), vsum); + _mm512_storeu_ps(C + m * ldc, vsum); + } + } +}; + +template +struct acc_C { + static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_1 * A, int lda, const void * packed_B, int nr) { + const int offset = TILE_N * TILE_K / 2; + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset))); + const __m512 vm0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset + TILE_N * sizeof(ggml_half)))); + + for (int m = 0; m < nr; ++m) { + const __m512 vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[m * lda].d)); + const __m512 vs1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[m * lda].s)); + const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); + + __m512 vsum; + if (is_acc) { + vsum = _mm512_loadu_ps(C + m * ldc); + } else { + vsum = _mm512_set1_ps(0.f); + } + vsum = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), vsum); + vsum = _mm512_fmadd_ps(vm0, vs1, vsum); + _mm512_storeu_ps(C + m * ldc, vsum); + } + } +}; + +template +struct acc_C { + static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_0 * A, int lda, const void * packed_B, int nr) { + const int offset = TILE_N * TILE_K; + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset))); + + for (int m = 0; m < nr; ++m) { + const __m512 vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[m * lda].d)); + const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); + + __m512 vsum; + if (is_acc) { + vsum = _mm512_loadu_ps(C + m * ldc); + } else { + vsum = _mm512_set1_ps(0.f); + } + vsum = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), vsum); + _mm512_storeu_ps(C + m * ldc, vsum); + } + } +}; + +template +struct acc_C { + static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) { + const uint8_t * scales = reinterpret_cast((const char *)packed_B + (QK_K / 2) * TILE_N); + const uint8_t * mins = scales + 8 * TILE_N; + const ggml_half * d0 = reinterpret_cast(mins + 8 * TILE_N); + const ggml_half * dmin = d0 + TILE_N; + + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0)); + const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)dmin)); + + for (int m = 0; m < nr; ++m) { + const float d1 = A[m * lda].d; + const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0); + const __m512 vdm = _mm512_mul_ps(_mm512_set1_ps(-d1), vdmin); + const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); + + __m512 vsum; + if (is_acc) { + vsum = _mm512_loadu_ps(C + m * ldc); + } else { + vsum = _mm512_set1_ps(0.f); + } + + const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[m * lda].bsums); + const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1)); + + __m512i acc_m = _mm512_setzero_si512(); + for (int k = 0; k < 4; ++k) { + __m512i vmask = _mm512_set1_epi32(k); + __m512i va = _mm512_permutexvar_epi32(vmask, _mm512_castsi128_si512(q8s)); + __m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(mins + k * 32))); + acc_m = _mm512_dpwssds_epi32(acc_m, va, vb); + } + + vsum = _mm512_fmadd_ps(vtile, vd, vsum); + vsum = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc_m), vdm, vsum); + _mm512_storeu_ps(C + m * ldc, vsum); + } + } +}; + +template +struct acc_C { + static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) { + const uint8_t * scales = reinterpret_cast((const char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N); + const uint8_t * mins = scales + 8 * TILE_N; + const ggml_half * d0 = reinterpret_cast(mins + 8 * TILE_N); + const ggml_half * dmin = d0 + TILE_N; + + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0)); + const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)dmin)); + + for (int m = 0; m < nr; ++m) { + const float d1 = A[m * lda].d; + const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0); + const __m512 vdm = _mm512_mul_ps(_mm512_set1_ps(-d1), vdmin); + const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); + + __m512 vsum; + if (is_acc) { + vsum = _mm512_loadu_ps(C + m * ldc); + } else { + vsum = _mm512_set1_ps(0.f); + } + + const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[m * lda].bsums); + const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1)); + + __m512i acc_m = _mm512_setzero_si512(); + for (int k = 0; k < 4; ++k) { + __m512i vmask = _mm512_set1_epi32(k); + __m512i va = _mm512_permutexvar_epi32(vmask, _mm512_castsi128_si512(q8s)); + __m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(mins + k * 32))); + acc_m = _mm512_dpwssds_epi32(acc_m, va, vb); + } + + vsum = _mm512_fmadd_ps(vtile, vd, vsum); + vsum = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc_m), vdm, vsum); + _mm512_storeu_ps(C + m * ldc, vsum); + } + } +}; + +template +struct acc_C { + static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) { + const uint8_t * scales = reinterpret_cast((const char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N); + const ggml_half * d0 = reinterpret_cast(scales + 16 * TILE_N); + + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0)); + + for (int m = 0; m < nr; ++m) { + const float d1 = A[m * lda].d; + const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0); + const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); + + __m512 vsum; + if (is_acc) { + vsum = _mm512_loadu_ps(C + m * ldc); + } else { + vsum = _mm512_set1_ps(0.f); + } + + vsum = _mm512_fmadd_ps(vtile, vd, vsum); + _mm512_storeu_ps(C + m * ldc, vsum); + } + } +}; + +template +struct acc_C { + static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) { + const int8_t * scales = reinterpret_cast((const char *)packed_B + (QK_K / 2) * TILE_N); + const ggml_half * d0 = reinterpret_cast(scales + 8 * TILE_N); + + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0)); + + for (int m = 0; m < nr; ++m) { + const float d1 = A[m * lda].d; + const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0); + const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N)); + + __m512 vsum; + if (is_acc) { + vsum = _mm512_loadu_ps(C + m * ldc); + } else { + vsum = _mm512_set1_ps(0.f); + } + + vsum = _mm512_fmadd_ps(vtile, vd, vsum); + _mm512_storeu_ps(C + m * ldc, vsum); + } + } +}; + +template constexpr int get_quants_size(); +template <> constexpr int get_quants_size() { return (QK_K / 2) * TILE_N; } +template <> constexpr int get_quants_size() { return (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N; } +template <> constexpr int get_quants_size() { return (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N; } +template <> constexpr int get_quants_size() { return (QK_K / 2) * TILE_N; } + +// used for QKK format +template ::value, int>::type = 0> +inline void scale_C(const int32_t * RESTRICT tile, int32_t * RESTRICT sumi, const void * packed_B, int k, int nr) { + const uint8_t * scales = reinterpret_cast((const char *)packed_B + get_quants_size()); + const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(scales + k * TILE_N))); + + for (int m = 0; m < nr; ++m) { + __m512i vsumi; + if (is_acc) { + vsumi = _mm512_loadu_si512(sumi + m * TILE_N); + } else { + vsumi = _mm512_setzero_si512(); + } + __m512i vtile = _mm512_loadu_si512(tile + m * TILE_N); + vsumi = _mm512_add_epi32(vsumi, _mm512_mullo_epi32(vtile, vscale)); + _mm512_storeu_si512((__m512i *)(sumi + m * TILE_N), vsumi); + } +} + +template +struct tinygemm_kernel_avx { + static void apply(int K, const TA * RESTRICT A, const TB * RESTRICT B, TC * RESTRICT C, int ldc) { + GGML_UNUSED(K); + GGML_UNUSED(A); + GGML_UNUSED(B); + GGML_UNUSED(C); + GGML_UNUSED(ldc); + } +}; + +template +struct tinygemm_kernel_avx { + static void apply(int K, const float * RESTRICT A, const ggml_fp16_t * RESTRICT B, float * RESTRICT C, int ldc) { + constexpr int ROWS = BLOCK_M; + constexpr int COLS = BLOCK_N; + assert(BLOCK_K == 16); + + __m512 va; + __m512 vb[COLS]; + __m512 vc[ROWS * COLS]; + + auto loadc = [&](auto idx) { + vc[idx] = _mm512_setzero_ps(); + }; + Unroll{}(loadc); + + auto compute = [&](auto idx, auto k) { + constexpr int row = idx / COLS; + constexpr int col = idx % COLS; + + if constexpr (col == 0) { + va = _mm512_loadu_ps(A + row * K + k); + } + if constexpr (row == 0) { + vb[col] = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(B + col * K + k))); + } + vc[idx] = _mm512_fmadd_ps(va, vb[col], vc[idx]); + }; + + for (int k = 0; k < K; k += 16) { + Unroll{}(compute, k); + } + + auto storec = [&](auto idx) { + constexpr int row = idx / COLS; + constexpr int col = idx % COLS; + C[row * ldc + col] = _mm512_reduce_add_ps(vc[idx]); + }; + Unroll{}(storec); + } +}; + +#define LAUNCH_TINYGEMM_KERNEL_AVX(MB_SIZE, NB_SIZE) \ + tinygemm_kernel_avx::apply( \ + K, (const float *)src1->data + src1_offset + mb_start * K, \ + (const type *)src0->data + src0_offset + nb_start * K, \ + (float *)dst->data + dst_offset + mb_start * ldc + nb_start, ldc) + + +// re-organize in the format {NB, KB, TILE_SIZE}: +#define PACKED_INDEX(n, k, KB, tile_size) (n * KB + k) * tile_size + +template +void convert_B_packed_format(void * RESTRICT packed_B, const TB * RESTRICT B, int N, int K) { + const int NB = N / TILE_N; + const int KB = K / BLOCK_K; + const int TILE_SIZE = get_tile_size(); + + // parallel on NB should be enough + parallel_for(NB, [&](int begin, int end) { + for (int n = begin; n < end; ++n) { + for (int k = 0; k < KB; ++k) { + int n0 = n * TILE_N; + pack_B((char *)packed_B + PACKED_INDEX(n, k, KB, TILE_SIZE), &B[n0 * KB + k], KB); + } + } + }); +} + +template +struct tinygemm_kernel_vnni {}; + +template +struct tinygemm_kernel_vnni { + static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { + + constexpr int COLS = BLOCK_N / 16; + const int TILE_SIZE = TILE_N * sizeof(block_q4_0); + + const block_q8_0 * RESTRICT A = static_cast(_A); + const char * RESTRICT B = static_cast(_B); + + __m512i va[8]; + __m512 vc[COLS]; + __m512 vd1; + + // sum of offsets, shared across COLS + // + // avx512-vnni does not have `_mm512_dpbssd_epi32`, + // need to transform ss to us: + // a * (b - 8) is equivalent to b * a - 8 * a + // s u u u s u s + // + __m512i vcomp; + + const __m512i off = _mm512_set1_epi8(8); + const __m512i lowMask = _mm512_set1_epi8(0xF); + + auto loadc = [&](auto col) { + vc[col] = _mm512_setzero_ps(); + }; + Unroll{}(loadc); + + auto compute = [&](auto col, auto i) { + // load a and compute compensation + if constexpr (col == 0) { + const int32_t * a_ptr = reinterpret_cast(A[0 * KB + i].qs); + vcomp = _mm512_setzero_si512(); + for (int k = 0; k < 8; ++k) { + va[k] = _mm512_set1_epi32(a_ptr[k]); + vcomp = _mm512_dpbusd_epi32(vcomp, off, va[k]); + } + vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[0 * KB + i].d)); + } + + // load b + __m512i vsum = _mm512_setzero_si512(); + const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); + for (int k = 0; k < 8; k += 2) { + __m512i bytes = _mm512_loadu_si512((const __m512i *)(b_ptr + k * 32)); + __m512i vb0 = _mm512_and_si512(bytes, lowMask); + vsum = _mm512_dpbusd_epi32(vsum, vb0, va[k + 0]); + __m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); + vsum = _mm512_dpbusd_epi32(vsum, vb1, va[k + 1]); + } + const int offset = TILE_N * TILE_K / 2; + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset))); + vsum = _mm512_sub_epi32(vsum, vcomp); + + vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(vsum), _mm512_mul_ps(vd0, vd1), vc[col]); + }; + + for (int i = 0; i < KB; ++i) { + Unroll{}(compute, i); + } + + //store to C + auto storec = [&](auto col) { + _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); + }; + Unroll{}(storec); + } +}; + +template +struct tinygemm_kernel_vnni { + static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { + + constexpr int COLS = BLOCK_N / 16; + const int TILE_SIZE = TILE_N * sizeof(block_q4_1); + + const block_q8_1 * RESTRICT A = static_cast(_A); + const char * RESTRICT B = static_cast(_B); + + __m512i va[8]; + __m512i vb[8]; + __m512 vc[COLS]; + __m512 vd1, vs1; + + const __m512i lowMask = _mm512_set1_epi8(0xF); + + auto loadc = [&](auto col) { + vc[col] = _mm512_setzero_ps(); + }; + Unroll{}(loadc); + + auto compute = [&](auto col, auto i) { + // load a + if constexpr (col == 0) { + const int32_t * a_ptr = reinterpret_cast(A[0 * KB + i].qs); + for (int k = 0; k < 8; ++k) { + va[k] = _mm512_set1_epi32(a_ptr[k]); + } + vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[0 * KB + i].d)); + vs1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[0 * KB + i].s)); + } + + // load b + const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); + for (int k = 0; k < 8; k += 2) { + __m512i bytes = _mm512_loadu_si512((const __m512i *)(b_ptr + k * 32)); + vb[k + 0] = _mm512_and_si512(bytes, lowMask); + vb[k + 1] = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); + } + const int offset = TILE_N * TILE_K / 2; + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset))); + const __m512 vm0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset + TILE_N * sizeof(ggml_half)))); + + __m512i vsum = _mm512_setzero_si512(); + for (int k = 0; k < 8; ++k) { + vsum = _mm512_dpbusd_epi32(vsum, vb[k], va[k]); + } + + vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(vsum), _mm512_mul_ps(vd0, vd1), vc[col]); + vc[col] = _mm512_fmadd_ps(vm0, vs1, vc[col]); + }; + + for (int i = 0; i < KB; ++i) { + Unroll{}(compute, i); + } + + //store to C + auto storec = [&](auto col) { + _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); + }; + Unroll{}(storec); + } +}; + +template +struct tinygemm_kernel_vnni { + static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { + + constexpr int COLS = BLOCK_N / 16; + const int TILE_SIZE = TILE_N * sizeof(block_q8_0) + TILE_N * sizeof(int32_t); + + const block_q8_0 * RESTRICT A = static_cast(_A); + const char * RESTRICT B = static_cast(_B); + + __m512i va[8]; + __m512i vb[8]; + __m512 vc[COLS]; + __m512 vd1; + + // Notes: s8s8 igemm compensation in avx512-vnni + // change s8s8 to u8s8 with compensate + // a * b = (a + 128) * b - 128 * b + // s s u s u s + // + // (128 * b is pre-computed when packing B to vnni formats) + // + const __m512i off = _mm512_set1_epi8(static_cast(0x80)); + + auto loadc = [&](auto col) { + vc[col] = _mm512_setzero_ps(); + }; + Unroll{}(loadc); + + auto compute = [&](auto col, auto i) { + // load a and add offset 128 + if constexpr (col == 0) { + const int32_t * a_ptr = reinterpret_cast(A[0 * KB + i].qs); + for (int k = 0; k < 8; ++k) { + va[k] = _mm512_set1_epi32(a_ptr[k]); + va[k] = _mm512_add_epi8(va[k], off); + } + vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[0 * KB + i].d)); + } + + // load b + const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); + for (int k = 0; k < 8; ++k) { + vb[k] = _mm512_loadu_si512((const __m512i *)(b_ptr + k * 64)); + } + const int offset = TILE_N * TILE_K; + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset))); + const int offset2 = TILE_N * TILE_K + TILE_N * sizeof(ggml_half); + const __m512i vcomp = _mm512_loadu_si512((const __m512i *)(b_ptr + offset2)); + + __m512i vsum = _mm512_setzero_si512(); + for (int k = 0; k < 8; ++k) { + vsum = _mm512_dpbusd_epi32(vsum, va[k], vb[k]); + } + vsum = _mm512_sub_epi32(vsum, vcomp); + + vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(vsum), _mm512_mul_ps(vd0, vd1), vc[col]); + }; + + for (int i = 0; i < KB; ++i) { + Unroll{}(compute, i); + } + + //store to C + auto storec = [&](auto col) { + _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); + }; + Unroll{}(storec); + } +}; + +template +struct tinygemm_kernel_vnni { + static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { + + constexpr int COLS = BLOCK_N / 16; + const int TILE_SIZE = TILE_N * sizeof(block_q4_K) + TILE_N * 4; + + const block_q8_K * RESTRICT A = static_cast(_A); + const char * RESTRICT B = static_cast(_B); + + // a.qs: 8 groups, 32 bytes each group (m256i) + __m512i va[8]; + // a.bsum: 8 groups, 2 bytes each group (m128i) + __m512i va_bsum; + __m512 vc[COLS]; + __m512 vd1; + + // packed_B: + const int offset_scales = (QK_K / 2) * TILE_N; + const int offset_mins = (QK_K / 2) * TILE_N + 8 * TILE_N; + const int offset_d0 = (QK_K / 2) * TILE_N + 16 * TILE_N; + const int offset_dmin = (QK_K / 2) * TILE_N + 16 * TILE_N + TILE_N * sizeof(ggml_half); + + const __m512i lowMask = _mm512_set1_epi8(0xF); + + auto loadc = [&](auto col) { + vc[col] = _mm512_setzero_ps(); + }; + Unroll{}(loadc); + + // Notes: vnni formats in QK_K + // a) quants vnni format + // int8 {k/4, n, 4}, viewed as 2d {k/4, 4n}, k = 32 + // from {16, 32} to {8, 64} + // + // b) min vnni format + // int16 {k/2, n, 2}, viewed as 2d {k/2, 2n}, k = 8 + // from {16, 8} to {4, 32} + // + auto compute = [&](auto col, auto i) { + // load a + if constexpr (col == 0) { + for (int k_group = 0; k_group < QK_K / 32; ++k_group) { + va[k_group] = _mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)(A[0 * KB + i].qs + k_group * 32))); + } + const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums); + const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1)); + va_bsum = _mm512_castsi128_si512(q8s); + vd1 = _mm512_set1_ps(A[0 * KB + i].d); + } + + // step 1: accumultate the quants + __m512i acc = _mm512_setzero_si512(); + const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); + const char * b_qs = b_ptr; + for (int k_group = 0; k_group < QK_K / 32; ++k_group) { + __m512i vsum = _mm512_setzero_si512(); + for (int k = 0; k < 8; k += 2) { + __m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 0), va[k_group]); + __m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 1), va[k_group]); + + __m512i bytes = _mm512_loadu_si512((const __m512i *)b_qs); + __m512i vb0 = _mm512_and_si512(bytes, lowMask); + vsum = _mm512_dpbusd_epi32(vsum, vb0, va0); + __m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); + vsum = _mm512_dpbusd_epi32(vsum, vb1, va1); + + b_qs += 64; + } + // vacc += scale * (q8 @ q4) + const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N))); + acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale)); + } + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0))); + vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]); + + // step 2: accumulate the mins + __m512i acc_m = _mm512_setzero_si512(); + for (int k = 0; k < 4; ++k) { + __m512i vmask = _mm512_set1_epi32(k); + __m512i va = _mm512_permutexvar_epi32(vmask, va_bsum); + __m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_mins + k * 32))); + acc_m = _mm512_dpwssds_epi32(acc_m, va, vb); + } + const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_dmin))); + vc[col] = _mm512_fnmadd_ps(_mm512_cvtepi32_ps(acc_m), _mm512_mul_ps(vdmin, vd1), vc[col]); + }; + + for (int i = 0; i < KB; ++i) { + Unroll{}(compute, i); + } + + //store to C + auto storec = [&](auto col) { + _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); + }; + Unroll{}(storec); + } +}; + +template +struct tinygemm_kernel_vnni { + static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { + + constexpr int COLS = BLOCK_N / 16; + const int TILE_SIZE = TILE_N * sizeof(block_q5_K) + TILE_N * 4; + + const block_q8_K * RESTRICT A = static_cast(_A); + const char * RESTRICT B = static_cast(_B); + + // a.qs: 8 groups, 32 bytes each group (m256i) + __m512i va[8]; + // a.bsum: 8 groups, 2 bytes each group (m128i) + __m512i va_bsum; + __m512 vc[COLS]; + __m512 vd1; + + // packed_B: + const int offset_qh = (QK_K / 2) * TILE_N; + const int offset_scales = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N; + const int offset_mins = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N + 8 * TILE_N; + const int offset_d0 = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N + 16 * TILE_N; + const int offset_dmin = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N + 16 * TILE_N + TILE_N * sizeof(ggml_half); + + const __m512i lowMask = _mm512_set1_epi8(0xF); + + auto loadc = [&](auto col) { + vc[col] = _mm512_setzero_ps(); + }; + Unroll{}(loadc); + + // Q5_K and Q4_K shares the same vnni formats, refer to notes above. + auto compute = [&](auto col, auto i) { + // load a + if constexpr (col == 0) { + for (int k_group = 0; k_group < QK_K / 32; ++k_group) { + va[k_group] = _mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)(A[0 * KB + i].qs + k_group * 32))); + } + const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums); + const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1)); + va_bsum = _mm512_castsi128_si512(q8s); + vd1 = _mm512_set1_ps(A[0 * KB + i].d); + } + + // step 1: accumultate the quants + __m512i acc = _mm512_setzero_si512(); + const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); + const char * b_qs = b_ptr; + const char * b_qh = b_ptr + offset_qh; + for (int k_group = 0; k_group < QK_K / 32; ++k_group) { + __m512i vsum = _mm512_setzero_si512(); + __m512i hmask0 = _mm512_set1_epi8(0x1); + __m512i hmask1 = _mm512_set1_epi8(0x2); + __m512i hbits = _mm512_loadu_si512((const __m512i *)(b_qh + k_group * 64)); + for (int k = 0; k < 8; k += 2) { + __m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 0), va[k_group]); + __m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 1), va[k_group]); + + __m512i bytes = _mm512_loadu_si512((const __m512i *)b_qs); + __m512i vb0 = _mm512_and_si512(bytes, lowMask); + __m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); + + __m512i vh0 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask0), k), 4); + __m512i vh1 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask1), k + 1), 4); + + hmask0 = _mm512_slli_epi16(hmask0, 2); + hmask1 = _mm512_slli_epi16(hmask1, 2); + vb0 = _mm512_add_epi8(vb0, vh0); + vb1 = _mm512_add_epi8(vb1, vh1); + + vsum = _mm512_dpbusd_epi32(vsum, vb0, va0); + vsum = _mm512_dpbusd_epi32(vsum, vb1, va1); + + b_qs += 64; + } + // vacc += scale * (q8 @ q5) + const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N))); + acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale)); + } + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0))); + vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]); + + // step 2: accumulate the mins + __m512i acc_m = _mm512_setzero_si512(); + for (int k = 0; k < 4; ++k) { + __m512i vmask = _mm512_set1_epi32(k); + __m512i va = _mm512_permutexvar_epi32(vmask, va_bsum); + __m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_mins + k * 32))); + acc_m = _mm512_dpwssds_epi32(acc_m, va, vb); + } + const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_dmin))); + vc[col] = _mm512_fnmadd_ps(_mm512_cvtepi32_ps(acc_m), _mm512_mul_ps(vdmin, vd1), vc[col]); + }; + + for (int i = 0; i < KB; ++i) { + Unroll{}(compute, i); + } + + //store to C + auto storec = [&](auto col) { + _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); + }; + Unroll{}(storec); + } +}; + +template +struct tinygemm_kernel_vnni { + static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { + + constexpr int COLS = BLOCK_N / 16; + const int TILE_SIZE = TILE_N * sizeof(block_q6_K); + + const block_q8_K * RESTRICT A = static_cast(_A); + const char * RESTRICT B = static_cast(_B); + + // load the 256 bytes from A to 4 avx512 vectors + __m512i va[4]; + __m512 vc[COLS]; + __m512 vd1; + + // packed_B: + const int offset_qh = (QK_K / 2) * TILE_N; + const int offset_scales = (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N; + const int offset_d0 = (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N + 16 * TILE_N; + + // compensation + __m512i vcomp; + + const __m512i m32s = _mm512_set1_epi32(32); + const __m512i lowMask = _mm512_set1_epi8(0xF); + + auto loadc = [&](auto col) { + vc[col] = _mm512_setzero_ps(); + }; + Unroll{}(loadc); + + auto compute = [&](auto col, auto i) { + if constexpr (col == 0) { + // load a + va[0] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 0)); + va[1] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 64)); + va[2] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 128)); + va[3] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 192)); + + const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums); + vcomp = _mm512_mullo_epi32(_mm512_cvtepi16_epi32(q8sums), m32s); + vd1 = _mm512_set1_ps(A[0 * KB + i].d); + } + + // accmulate the quants + __m512i acc = _mm512_setzero_si512(); + const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); + const char * b_qs = b_ptr; + const char * b_qh = b_ptr + offset_qh; + int mask = 0; + for (int k_group = 0; k_group < QK_K / 16; ++k_group) { + int r = k_group >> 2; + __m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]); + __m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]); + + __m512i vsum = _mm512_setzero_si512(); + __m512i hmask = _mm512_set1_epi8(0x3); + + __m512i bytes = _mm512_loadu_si512(b_qs); + __m512i hbits = _mm512_loadu_si512(b_qh); + __m512i vb0 = _mm512_and_si512(bytes, lowMask); + __m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); + __m512i vh0 = _mm512_slli_epi16(_mm512_and_si512(hbits, hmask), 4); + __m512i vh1 = _mm512_slli_epi16(_mm512_and_si512(hbits, _mm512_slli_epi16(hmask, 2)), 2); + + vb0 = _mm512_add_epi8(vb0, vh0); + vb1 = _mm512_add_epi8(vb1, vh1); + vsum = _mm512_dpbusd_epi32(vsum, vb0, va0); + vsum = _mm512_dpbusd_epi32(vsum, vb1, va1); + b_qs += 64; + + va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]); + va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]); + + bytes = _mm512_loadu_si512(b_qs); + vb0 = _mm512_and_si512(bytes, lowMask); + vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask); + vh0 = _mm512_and_si512(hbits, _mm512_slli_epi16(hmask, 4)); + vh1 = _mm512_srli_epi16(_mm512_and_si512(hbits, _mm512_slli_epi16(hmask, 6)), 2); + vb0 = _mm512_add_epi8(vb0, vh0); + vb1 = _mm512_add_epi8(vb1, vh1); + vsum = _mm512_dpbusd_epi32(vsum, vb0, va0); + vsum = _mm512_dpbusd_epi32(vsum, vb1, va1); + b_qs += 64; + b_qh += 64; + + // B * A - 32 * A + __m512i vmask = _mm512_set1_epi32(k_group); + vsum = _mm512_sub_epi32(vsum, _mm512_permutexvar_epi32(vmask, vcomp)); + + // vacc += scale * (q8 @ q6) + const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N))); + acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale)); + } + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0))); + vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]); + }; + + for (int i = 0; i < KB; ++i) { + Unroll{}(compute, i); + } + + //store to C + auto storec = [&](int col) { + _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); + }; + Unroll{}(storec); + } +}; + +template +struct tinygemm_kernel_vnni { + static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { + + constexpr int COLS = BLOCK_N / 16; + const int TILE_SIZE = TILE_N * sizeof(block_iq4_xs) + TILE_N * 2; + + const block_q8_K * RESTRICT A = static_cast(_A); + const char * RESTRICT B = static_cast(_B); + + // load the 256 bytes from A to 4 avx512 vectors + __m512i va[4]; + __m512 vc[COLS]; + __m512 vd1; + + // packed_B: + const int offset_scales = (QK_K / 2) * TILE_N ; + const int offset_d0 = (QK_K / 2) * TILE_N + 8 * TILE_N; + + // compensation + __m512i vcomp; + + const __m256i m128s = _mm256_set1_epi16(128); + const __m512i lowMask = _mm512_set1_epi8(0xF); + + const __m512i values128 = _mm512_set_epi8( + 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127, + 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127, + 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127, + 113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127 + ); + const __m512i off = _mm512_set1_epi8(static_cast(0x80)); + const __m512i values256 = _mm512_add_epi8(values128, off); + + auto loadc = [&](auto col) { + vc[col] = _mm512_setzero_ps(); + }; + Unroll{}(loadc); + + auto compute = [&](auto col, auto i) { + if constexpr (col == 0) { + // load a + va[0] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 0)); + va[1] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 64)); + va[2] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 128)); + va[3] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 192)); + + // compensation: 128 * A + const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums); + vcomp = _mm512_castsi256_si512(_mm256_madd_epi16(q8sums, m128s)); + vd1 = _mm512_set1_ps(A[0 * KB + i].d); + } + + // accmulate the quants + __m512i acc = _mm512_setzero_si512(); + const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE); + const char * b_qs = b_ptr; + int mask = 0; + for (int k_group = 0; k_group < QK_K / 32; ++k_group) { + int r = k_group >> 1; + __m512i vmask = _mm512_set1_epi32(k_group); + __m512i vsum = _mm512_setzero_si512(); + for (int k = 0; k < 8; k += 2) { + __m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]); + __m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]); + + __m512i bytes = _mm512_loadu_si512(b_qs); + __m512i vb0 = _mm512_shuffle_epi8(values256, _mm512_and_si512(bytes, lowMask)); + __m512i vb1 = _mm512_shuffle_epi8(values256, _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask)); + + vsum = _mm512_dpbusd_epi32(vsum, vb0, va0); + vsum = _mm512_dpbusd_epi32(vsum, vb1, va1); + b_qs += 64; + } + // (B + 128) * A - 128 * A + vsum = _mm512_sub_epi32(vsum, _mm512_permutexvar_epi32(vmask, vcomp)); + + // vacc += scale * (q8 @ q4) + const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N))); + acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale)); + } + const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0))); + vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]); + }; + + for (int i = 0; i < KB; ++i) { + Unroll{}(compute, i); + } + + //store to C + auto storec = [&](auto col) { + _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]); + }; + Unroll{}(storec); + } +}; + +#define LAUNCH_TINYGEMM_KERNEL_VNNI(NB_SIZE) \ + tinygemm_kernel_vnni::apply( \ + KB, wdata_batch, \ + (const char *)src0->data + src0_offset + PACKED_INDEX(nb * kTilesN, 0, KB, TILE_SIZE), \ + (float *) dst->data + dst_offset + nb_start, ldc) + +template ::value, int>::type = 0> +void tinygemm_kernel_amx(int M, int N, int KB, const void * RESTRICT _A, const void * RESTRICT _B, TC * RESTRICT C, int ldc) { + using packed_B_t = packed_B_type; + const int TILE_SIZE = get_tile_size(); + const bool need_unpack = do_unpack::value; + + GGML_ASSERT(M <= 2 * TILE_M && N == 2 * TILE_N); + const TA * RESTRICT A = static_cast(_A); + const char * RESTRICT B = static_cast(_B); + + const int m0 = std::min(M, TILE_M); + const int m1 = std::max(M - TILE_M, 0); + const int lda = KB * sizeof(TA); + //const int ldb = KB * sizeof(TB); + + alignas(64) static thread_local packed_B_t Tile0[TILE_N * TILE_K]; + alignas(64) static thread_local packed_B_t Tile1[TILE_N * TILE_K]; + alignas(64) static thread_local int8_t Tile23[TILE_M * TILE_K]; + + alignas(64) static thread_local int32_t TileC0[TILE_M * TILE_N * 4]; + alignas(64) static thread_local int32_t TileC1[TILE_M * TILE_N * 4]; + + // double buffering C to interleave avx512 and amx + int32_t * C_cur = TileC0; + int32_t * C_pre = TileC1; + + auto Tile4 = [&](int32_t * base) { return base; }; + auto Tile5 = [&](int32_t * base) { return base + TILE_M * TILE_N; }; + auto Tile6 = [&](int32_t * base) { return base + 2 * TILE_M * TILE_N; }; + auto Tile7 = [&](int32_t * base) { return base + 3 * TILE_M * TILE_N; }; + + if (M == 2 * TILE_M) { + // i = 0 + const char * B_blk0 = B + PACKED_INDEX(0, 0, KB, TILE_SIZE); + const char * B_blk1 = B + PACKED_INDEX(1, 0, KB, TILE_SIZE); + if (need_unpack) { + unpack_B(Tile0, B_blk0); + _tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK); + } else { + _tile_loadd(TMM0, B_blk0, TILE_N * VNNI_BLK); + } + + _tile_zero(TMM4); + _tile_loadd(TMM2, A[0].qs, lda); + _tile_dpbssd(TMM4, TMM2, TMM0); + _tile_stored(TMM4, Tile4(C_pre), TILE_N * sizeof(int32_t)); + + _tile_zero(TMM5); + _tile_loadd(TMM3, A[TILE_M * KB + 0].qs, lda); + _tile_dpbssd(TMM5, TMM3, TMM0); + _tile_stored(TMM5, Tile5(C_pre), TILE_N * sizeof(int32_t)); + + if (need_unpack) { + unpack_B(Tile1, B_blk1); + _tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK); + } else { + _tile_loadd(TMM1, B_blk1, TILE_N * VNNI_BLK); + } + + _tile_zero(TMM6); + _tile_dpbssd(TMM6, TMM2, TMM1); + _tile_stored(TMM6, Tile6(C_pre), TILE_N * sizeof(int32_t)); + + _tile_zero(TMM7); + _tile_dpbssd(TMM7, TMM3, TMM1); + _tile_stored(TMM7, Tile7(C_pre), TILE_N * sizeof(int32_t)); + + for (int i = 1; i < KB; ++i) { + // index of previous iter + const int ii = i - 1; + const char * B_blk0 = B + PACKED_INDEX(0, i, KB, TILE_SIZE); + const char * B_blk1 = B + PACKED_INDEX(1, i, KB, TILE_SIZE); + GGML_DISPATCH_BOOL(ii > 0, is_acc, [&] { + if (need_unpack) { + unpack_B(Tile0, B_blk0); + _tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK); + } else { + _tile_loadd(TMM0, B_blk0, TILE_N * VNNI_BLK); + } + _tile_zero(TMM4); + _tile_loadd(TMM2, A[i].qs, lda); + acc_C::apply(C, ldc, Tile4(C_pre), &A[ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M); + + _tile_dpbssd(TMM4, TMM2, TMM0); + _tile_stored(TMM4, Tile4(C_cur), TILE_N * sizeof(int32_t)); + + _tile_zero(TMM5); + _tile_loadd(TMM3, A[TILE_M * KB + i].qs, lda); + acc_C::apply(C + TILE_M * ldc, ldc, Tile5(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M); + + _tile_dpbssd(TMM5, TMM3, TMM0); + _tile_stored(TMM5, Tile5(C_cur), TILE_N * sizeof(int32_t)); + + if (need_unpack) { + unpack_B(Tile1, B_blk1); + _tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK); + } else { + _tile_loadd(TMM1, B_blk1, TILE_N * VNNI_BLK); + } + _tile_zero(TMM6); + acc_C::apply(C + TILE_N, ldc, Tile6(C_pre), &A[ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M); + + _tile_dpbssd(TMM6, TMM2, TMM1); + _tile_stored(TMM6, Tile6(C_cur), TILE_N * sizeof(int32_t)); + + _tile_zero(TMM7); + acc_C::apply(C + TILE_M * ldc + TILE_N, ldc, Tile7(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M); + + _tile_dpbssd(TMM7, TMM3, TMM1); + _tile_stored(TMM7, Tile7(C_cur), TILE_N * sizeof(int32_t)); + + std::swap(C_cur, C_pre); + }); + } + // final accumulation + { + int ii = KB - 1; + acc_C::apply(C, ldc, Tile4(C_pre), &A[ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M); + acc_C::apply(C + TILE_M * ldc, ldc, Tile5(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M); + acc_C::apply(C + TILE_N, ldc, Tile6(C_pre), &A[ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M); + acc_C::apply(C + TILE_M * ldc + TILE_N, ldc, Tile7(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M); + } + } else { + for (int i = 0; i < KB; ++i) { + _tile_zero(TMM4); + _tile_zero(TMM6); + if (m1 != 0) { + _tile_zero(TMM5); + _tile_zero(TMM7); + } + + const char * B_blk0 = B + PACKED_INDEX(0, i, KB, TILE_SIZE); + const char * B_blk1 = B + PACKED_INDEX(1, i, KB, TILE_SIZE); + if (need_unpack) { + unpack_B(Tile0, B_blk0); + _tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK); + } else { + _tile_loadd(TMM0, B_blk0, TILE_N * VNNI_BLK); + } + + if (need_unpack) { + unpack_B(Tile1, B_blk1); + _tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK); + } else { + _tile_loadd(TMM1, B_blk1, TILE_N * VNNI_BLK); + } + + if (m0 == TILE_M) { + _tile_loadd(TMM2, A[i].qs, lda); + } else { + unpack_A(Tile23, &A[i], KB, m0); + _tile_loadd(TMM2, Tile23, TILE_K); + } + + _tile_dpbssd(TMM4, TMM2, TMM0); + _tile_dpbssd(TMM6, TMM2, TMM1); + + _tile_stored(TMM4, Tile4(C_cur), TILE_N * sizeof(int32_t)); + _tile_stored(TMM6, Tile6(C_cur), TILE_N * sizeof(int32_t)); + + GGML_DISPATCH_BOOL(i > 0, is_acc, [&] { + acc_C::apply(C, ldc, Tile4(C_cur), &A[i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m0); + acc_C::apply(C + TILE_N, ldc, Tile6(C_cur), &A[i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m0); + }); + + if (m1 != 0) { + unpack_A(Tile23, &A[TILE_M * KB + i], KB, m1); + _tile_loadd(TMM3, Tile23, TILE_K); + + _tile_dpbssd(TMM5, TMM3, TMM0); + _tile_dpbssd(TMM7, TMM3, TMM1); + _tile_stored(TMM5, Tile5(C_cur), TILE_N * sizeof(int32_t)); + _tile_stored(TMM7, Tile7(C_cur), TILE_N * sizeof(int32_t)); + GGML_DISPATCH_BOOL(i > 0, is_acc, [&] { + acc_C::apply(C + TILE_M * ldc, ldc, Tile5(C_cur), &A[TILE_M * KB + i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m1); + acc_C::apply(C + TILE_M * ldc + TILE_N, ldc, Tile7(C_cur), &A[TILE_M * KB + i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m1); + }); + } + } + } + return; +} + +template ::value, int>::type = 0> +void tinygemm_kernel_amx(int M, int N, int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) { + static_assert(std::is_same::value); + const int TILE_SIZE = get_tile_size(); + + GGML_ASSERT(M <= 2 * TILE_M && N == 2 * TILE_N); + const TA * RESTRICT A = static_cast(_A); + const char * RESTRICT B = static_cast(_B); + + const int m0 = std::min(M, TILE_M); + const int m1 = std::max(M - TILE_M, 0); + //const int lda = KB * sizeof(TA); + + alignas(64) static thread_local int8_t Tile0[TILE_N * TILE_K]; + alignas(64) static thread_local int8_t Tile1[TILE_N * TILE_K]; + alignas(64) static thread_local int8_t Tile23[TILE_M * TILE_K]; + + // mat mul result for each group + alignas(64) static thread_local int32_t Tile4[TILE_M * TILE_N]; + alignas(64) static thread_local int32_t Tile5[TILE_M * TILE_N]; + alignas(64) static thread_local int32_t Tile6[TILE_M * TILE_N]; + alignas(64) static thread_local int32_t Tile7[TILE_M * TILE_N]; + + // sum of each QK_K block, contains 8 groups, int32 + alignas(64) static thread_local int32_t Sumi4[TILE_M * TILE_N]; + alignas(64) static thread_local int32_t Sumi5[TILE_M * TILE_N]; + alignas(64) static thread_local int32_t Sumi6[TILE_M * TILE_N]; + alignas(64) static thread_local int32_t Sumi7[TILE_M * TILE_N]; + + const int k_group_size = std::is_same::value ? 16 : 32; + for (int i = 0; i < KB; ++i) { + // step 1: accumulate the quants across 8 groups, each group with 32 + for (int k = 0; k < QK_K / k_group_size; ++k) { + GGML_DISPATCH_BOOL(k > 0, is_acc, [&] { + _tile_zero(TMM4); + _tile_zero(TMM6); + + unpack_B(Tile0, B + PACKED_INDEX(0, i, KB, TILE_SIZE), k); + _tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK); + + unpack_B(Tile1, B + PACKED_INDEX(1, i, KB, TILE_SIZE), k); + _tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK); + + unpack_A(Tile23, &A[i], KB, k, m0); + _tile_loadd(TMM2, Tile23, TILE_K); + + _tile_dpbssd(TMM4, TMM2, TMM0); + _tile_dpbssd(TMM6, TMM2, TMM1); + + _tile_stored(TMM4, Tile4, TILE_N * sizeof(int32_t)); + _tile_stored(TMM6, Tile6, TILE_N * sizeof(int32_t)); + + scale_C(Tile4, Sumi4, B + PACKED_INDEX(0, i, KB, TILE_SIZE), k, m0); + scale_C(Tile6, Sumi6, B + PACKED_INDEX(1, i, KB, TILE_SIZE), k, m0); + + if (m1 != 0) { + _tile_zero(TMM5); + _tile_zero(TMM7); + + unpack_A(Tile23, &A[TILE_M * KB + i], KB, k, m1); + _tile_loadd(TMM3, Tile23, TILE_K); + + _tile_dpbssd(TMM5, TMM3, TMM0); + _tile_dpbssd(TMM7, TMM3, TMM1); + + _tile_stored(TMM5, Tile5, TILE_N * sizeof(int32_t)); + _tile_stored(TMM7, Tile7, TILE_N * sizeof(int32_t)); + + scale_C(Tile5, Sumi5, B + PACKED_INDEX(0, i, KB, TILE_SIZE), k, m1); + scale_C(Tile7, Sumi7, B + PACKED_INDEX(1, i, KB, TILE_SIZE), k, m1); + } + }); + } + + // step 2: accmulate the mins + GGML_DISPATCH_BOOL(i > 0, is_acc, [&] { + acc_C::apply(C, ldc, Sumi4, &A[i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m0); + acc_C::apply(C + TILE_N, ldc, Sumi6, &A[i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m0); + if (m1 != 0) { + acc_C::apply(C + TILE_M * ldc, ldc, Sumi5, &A[TILE_M * KB + i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m1); + acc_C::apply(C + TILE_M * ldc + TILE_N, ldc, Sumi7, &A[TILE_M * KB + i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m1); + } + }); + } + return; +} + +} // anonymous namespace + +// get the packed tensor size for quantized weights +size_t ggml_backend_amx_get_alloc_size(const struct ggml_tensor * tensor) { + const enum ggml_type TYPE = tensor->type; + + const int K = tensor->ne[0]; // ne0: in_features + const int N = tensor->ne[1]; // ne1: out_features + + auto get_tensor_size = [&] { + size_t row_size_B{0}; + GGML_DISPATCH_QTYPES(TYPE, [&] { + row_size_B = get_row_size(K); + }); + return N * row_size_B; + }; + + if (qtype_has_amx_kernels(TYPE)) { + return get_tensor_size(); + } else { + // for f16, bf16 we don't do packing + return ggml_nbytes(tensor); + } +} + +// pack weight to vnni format +void ggml_backend_amx_convert_weight(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + GGML_ASSERT(offset == 0 && size == ggml_nbytes(tensor)); // only full tensor conversion is supported for now + + const enum ggml_type TYPE = tensor->type; + + const int K = tensor->ne[0]; // ne0: in_features + const int N = tensor->ne[1]; // ne1: out_features + + GGML_DISPATCH_QTYPES(TYPE, [&] { + convert_B_packed_format((void *)((char *)tensor->data + offset), (const type *)data, N, K); + }); +} + +// ne2 is passed explicitly to help compiler optimize repeated calls +inline int64_t ggml_batch_offset(const ggml_tensor * t, int64_t batch_idx, int64_t ne2) { + const int64_t i2 = batch_idx % ne2; + const int64_t i3 = batch_idx / ne2; + return i3 * t->nb[3] + i2 * t->nb[2]; +} + +size_t ggml_backend_amx_desired_wsize(const struct ggml_tensor * dst) { + struct ggml_tensor * src0 = dst->src[0]; + + const enum ggml_type TYPE = src0->type; + + const bool is_floating_type = TYPE == GGML_TYPE_F16; + if (is_floating_type) { + return 0; + } + + const int M = dst->ne[1]; + const int K = src0->ne[0]; + const int64_t n_batch = dst->ne[2] * dst->ne[3]; + + size_t desired_wsize = 0; + + GGML_DISPATCH_QTYPES(TYPE, [&] { + const size_t row_size_A = K / blck_size * sizeof(vec_dot_type); + desired_wsize = n_batch * M * row_size_A; + }); + + return desired_wsize; +} + +// NB: mixed dtype gemm with Advanced Matrix Extensions (Intel AMX) +// +// src0: weight in shape of {N, K}, quantized +// src1: input in shape of {M, K}, float32 +// dst: output in shape of {M, N}, float32 +// +// the function performs: dst = src1 @ src0.T for each batch +// +void ggml_backend_amx_mul_mat(const ggml_compute_params * params, struct ggml_tensor * dst) { + struct ggml_tensor * src0 = dst->src[0]; + struct ggml_tensor * src1 = dst->src[1]; + + const enum ggml_type TYPE = src0->type; + + // f16 only has avx512 kernels for now, + // amx kernels will be added once 6th gen xeon is released. + const bool is_floating_type = TYPE == GGML_TYPE_F16; + + const int M = dst->ne[1]; + const int N = dst->ne[0]; + const int K = src0->ne[0]; + const int ldc = dst->nb[1] / dst->nb[0]; + + const int64_t ne2 = dst->ne[2]; + const int64_t n_batch = ne2 * dst->ne[3]; + + if (is_floating_type) { + constexpr int BLOCK_M = 4; + constexpr int BLOCK_N = 6; + const int MB = div_up(M, BLOCK_M); + const int NB = div_up(N, BLOCK_N); + + parallel_for_ggml(params, n_batch * MB * NB, [&](int begin, int end) { + GGML_DISPATCH_FLOATING_TYPES(TYPE, [&] { + for (int i = begin; i < end; ++i) { + int batch_idx = i / (MB * NB); + int remaining = i % (MB * NB); + int mb = remaining / NB; + int nb = remaining % NB; + + int64_t src0_offset = ggml_batch_offset(src0, batch_idx, ne2); + int64_t src1_offset = ggml_batch_offset(src1, batch_idx, ne2); + int64_t dst_offset = ggml_batch_offset(dst, batch_idx, ne2); + + int mb_start = mb * BLOCK_M; + int mb_size = std::min(BLOCK_M, M - mb_start); + int nb_start = nb * BLOCK_N; + int nb_size = std::min(BLOCK_N, N - nb_start); + + switch (mb_size << 4 | nb_size) { + case 0x12: LAUNCH_TINYGEMM_KERNEL_AVX(1, 2); break; + case 0x14: LAUNCH_TINYGEMM_KERNEL_AVX(1, 4); break; + case 0x16: LAUNCH_TINYGEMM_KERNEL_AVX(1, 6); break; + case 0x22: LAUNCH_TINYGEMM_KERNEL_AVX(2, 2); break; + case 0x24: LAUNCH_TINYGEMM_KERNEL_AVX(2, 4); break; + case 0x26: LAUNCH_TINYGEMM_KERNEL_AVX(2, 6); break; + case 0x32: LAUNCH_TINYGEMM_KERNEL_AVX(3, 2); break; + case 0x34: LAUNCH_TINYGEMM_KERNEL_AVX(3, 4); break; + case 0x36: LAUNCH_TINYGEMM_KERNEL_AVX(3, 6); break; + case 0x42: LAUNCH_TINYGEMM_KERNEL_AVX(4, 2); break; + case 0x44: LAUNCH_TINYGEMM_KERNEL_AVX(4, 4); break; + case 0x46: LAUNCH_TINYGEMM_KERNEL_AVX(4, 6); break; + default: fprintf(stderr, "Unexpected block size!\n"); + } + } + }); + }); + return; + } + + // pointer to work space, used convert A from float to quantized type + void * wdata = params->wdata; + + //TODO: performance improvement: merge quant A + // if (params->ith == 0) { + GGML_DISPATCH_QTYPES(TYPE, [&] { + const size_t row_size_A = K / blck_size * sizeof(vec_dot_type); + const size_t desired_wsize = n_batch * M * row_size_A; + if (params->wsize < desired_wsize) { + GGML_ABORT("insufficient work space size"); + } + + // Q4_0, Q4_1, Q8_0 handles 1 TILE_K per blck_size + // Q4_K, Q5_K, Q6_K, IQ4_XS handles 8 TILE_K per blck_size + GGML_ASSERT(TILE_K == blck_size || TILE_K * 8 == blck_size); + + parallel_for_ggml(params, n_batch, [&](int begin, int end) { + for (int batch_idx = begin; batch_idx < end; ++batch_idx) { + int64_t src1_offset = ggml_batch_offset(src1, batch_idx, ne2); + const float * A_data = (const float *)((const char *)src1->data + src1_offset); + char * wdata_batch = (char *)wdata + batch_idx * M * row_size_A; + + for (int m = 0; m < M; ++m) { + from_float(A_data + m * K, wdata_batch + m * row_size_A, K); + } + } + }); + }); + // } + + ggml_barrier(params->threadpool); + + if (M == 1) { + // MB = 1 and handle 8 tiles in each block + constexpr int kTilesN = 4; + constexpr int BLOCK_N = TILE_N * kTilesN; + const int NB = div_up(N, BLOCK_N); + + parallel_for_ggml(params, n_batch * NB, [&](int begin, int end) { + GGML_DISPATCH_QTYPES(TYPE, [&] { + const int KB = K / blck_size; + const int TILE_SIZE = get_tile_size(); + const int row_size_A = KB * sizeof(vec_dot_type); + for (int i = begin; i < end; ++i) { + int batch_idx = i / NB; + int nb = i % NB; + + int64_t src0_offset = ggml_batch_offset(src0, batch_idx, ne2); + int64_t dst_offset = ggml_batch_offset(dst, batch_idx, ne2); + const char * wdata_batch = (const char *)wdata + batch_idx * row_size_A; + + int nb_start = nb * BLOCK_N; + int nb_size = std::min(BLOCK_N, N - nb_start); // 32, 64, 96 + + switch (nb_size) { + //case 160: LAUNCH_TINYGEMM_KERNEL_VNNI(160); break; + case 128: LAUNCH_TINYGEMM_KERNEL_VNNI(128); break; + case 96: LAUNCH_TINYGEMM_KERNEL_VNNI(96); break; + case 64: LAUNCH_TINYGEMM_KERNEL_VNNI(64); break; + case 32: LAUNCH_TINYGEMM_KERNEL_VNNI(32); break; + default: fprintf(stderr, "Unexpected n block size!\n"); + } + } + }); + }); + return; + } + + // handle 4 tiles at a tile + constexpr int BLOCK_M = TILE_M * 2; + constexpr int BLOCK_N = TILE_N * 2; + const int MB = div_up(M, BLOCK_M); + const int NB = div_up(N, BLOCK_N); + + parallel_for_ggml(params, n_batch * MB * NB, [&](int begin, int end) { + // init tile config for each thread + ggml_tile_config_init(); + + GGML_DISPATCH_QTYPES(TYPE, [&] { + const int KB = K / blck_size; + const int TILE_SIZE = get_tile_size(); + const int row_size_A = KB * sizeof(vec_dot_type); + + for (int i = begin; i < end; ++i) { + int batch_idx = i / (MB * NB); + int remaining = i % (MB * NB); + int mb = remaining / NB; + int nb = remaining % NB; + + int64_t src0_offset = ggml_batch_offset(src0, batch_idx, ne2); + int64_t dst_offset = ggml_batch_offset(dst, batch_idx, ne2); + const char * wdata_batch = (const char *)wdata + batch_idx * M * row_size_A; + + int mb_start = mb * BLOCK_M; + int mb_size = std::min(BLOCK_M, M - mb_start); + int nb_start = nb * BLOCK_N; + int nb_size = BLOCK_N; + + tinygemm_kernel_amx( + mb_size, nb_size, KB, + wdata_batch + mb_start * row_size_A, + (const char *)src0->data + src0_offset + PACKED_INDEX(nb * 2, 0, KB, TILE_SIZE), + (float *) dst->data + dst_offset + mb_start * N + nb_start, ldc); + } + }); + }); +} + +#endif // if defined(__AMX_INT8__) && defined(__AVX512VNNI__) + + +// ---- llamafile sgemm implementation ---- + +// Copyright 2024 Mozilla Foundation +// +// Permission is hereby granted, free of charge, to any person obtaining +// a copy of this software and associated documentation files (the +// "Software"), to deal in the Software without restriction, including +// without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to +// permit persons to whom the Software is furnished to do so, subject to +// the following conditions: +// +// The above copyright notice and this permission notice shall be +// included in all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS +// BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN +// ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// +// _ _ ___ _ _ ___ +// | |_(_)_ _ _ _| _ ) | /_\ / __| +// | _| | ' \ || | _ \ |__ / _ \\__ \. +// \__|_|_||_\_, |___/____/_/ \_\___/ +// |__/ +// +// BASIC LINEAR ALGEBRA SUBPROGRAMS +// +// +// This file implements multithreaded CPU matrix multiplication for the +// common contiguous use case C = Aᵀ * B. These kernels are designed to +// have excellent performance[1] for matrices that fit in the CPU cache +// without imposing any overhead such as cache filling or malloc calls. +// +// This implementation does not guarantee any upper bound with rounding +// errors, which grow along with k. Our goal's to maximally exploit the +// hardware for performance, and then use whatever resources remain for +// improving numerical accuracy. +// +// [1] J. Tunney, ‘LLaMA Now Goes Faster on CPUs’, Mar. 2024. [Online]. +// Available: https://justine.lol/matmul/. [Accessed: 29-Mar-2024]. + +#if defined(__GNUC__) +#pragma GCC diagnostic ignored "-Wpedantic" +#pragma GCC diagnostic ignored "-Wignored-attributes" +#endif + +#include "ggml-impl-defs.inc" +#include "ggml-cpu-impl-defs.inc" +#include "ggml-quants-defs.inc" +#include "simd-mappings-defs.inc" + +#include +#include + +#ifdef _MSC_VER +#define NOINLINE __declspec(noinline) +#else +#define NOINLINE __attribute__((__noinline__)) +#endif + +extern "C" bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64_t n, int64_t k, + const void * A, int64_t lda, const void * B, int64_t ldb, void * C, int64_t ldc, + int ith, int nth, int type); + +#if defined(__ARM_NEON) || defined(__AVX512F__) || defined(__VXE__) || defined(__VXE2__) +#define VECTOR_REGISTERS 32 +#else +#define VECTOR_REGISTERS 16 +#endif + +#if defined(__riscv_v_intrinsic) +#define LMUL 4 +#endif + +#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1) + +namespace { + +inline float unhalf(ggml_fp16_t d) { + return GGML_CPU_FP16_TO_FP32(d); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// VECTORIZED ARITHMETIC OPERATIONS + +#if defined(__SSE__) || defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +inline __m128 add(__m128 x, __m128 y) { return _mm_add_ps(x, y); } +inline __m128 sub(__m128 x, __m128 y) { return _mm_sub_ps(x, y); } +inline __m128 mul(__m128 x, __m128 y) { return _mm_mul_ps(x, y); } +#endif // __SSE__ + +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +inline __m256 add(__m256 x, __m256 y) { return _mm256_add_ps(x, y); } +inline __m256 sub(__m256 x, __m256 y) { return _mm256_sub_ps(x, y); } +inline __m256 mul(__m256 x, __m256 y) { return _mm256_mul_ps(x, y); } +#endif // __AVX__ + +#if defined(__AVX512F__) +inline __m512 add(__m512 x, __m512 y) { return _mm512_add_ps(x, y); } +inline __m512 sub(__m512 x, __m512 y) { return _mm512_sub_ps(x, y); } +inline __m512 mul(__m512 x, __m512 y) { return _mm512_mul_ps(x, y); } +#endif // __AVX512F__ + +#if defined(__ARM_NEON) +inline float32x4_t add(float32x4_t x, float32x4_t y) { return vaddq_f32(x, y); } +inline float32x4_t sub(float32x4_t x, float32x4_t y) { return vsubq_f32(x, y); } +inline float32x4_t mul(float32x4_t x, float32x4_t y) { return vmulq_f32(x, y); } +#endif // __ARM_NEON + +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) +inline float16x8_t add(float16x8_t x, float16x8_t y) { return vaddq_f16(x, y); } +inline float16x8_t sub(float16x8_t x, float16x8_t y) { return vsubq_f16(x, y); } +inline float16x8_t mul(float16x8_t x, float16x8_t y) { return vmulq_f16(x, y); } +#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC + +#if defined(__VXE__) || defined(__VXE2__) +inline float32x4_t add(float32x4_t x, float32x4_t y) { return vec_add(x, y); } +inline float32x4_t sub(float32x4_t x, float32x4_t y) { return vec_sub(x, y); } +inline float32x4_t mul(float32x4_t x, float32x4_t y) { return vec_mul(x, y); } +#endif + +#if defined(__MMA__) +typedef vector unsigned char vec_t; +typedef __vector_quad acc_t; +#endif +//////////////////////////////////////////////////////////////////////////////////////////////////// +// VECTORIZED FUSED MULTIPLY ADD + +/** + * Computes a * b + c. + */ +template +inline U madd(T a, T b, U c) { + return add(mul(a, b), c); +} + +#if defined(__FMA__) +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +template <> +inline __m256 madd(__m256 a, __m256 b, __m256 c) { + return _mm256_fmadd_ps(a, b, c); +} +#endif +#if defined(__AVX512F__) +template <> +inline __m512 madd(__m512 a, __m512 b, __m512 c) { + return _mm512_fmadd_ps(a, b, c); +} +#endif +#if defined(__AVX512BF16__) +template <> +inline __m512 madd(__m512bh a, __m512bh b, __m512 c) { + return _mm512_dpbf16_ps(c, a, b); +} +template <> +inline __m256 madd(__m256bh a, __m256bh b, __m256 c) { + return _mm256_dpbf16_ps(c, a, b); +} +#endif +#endif + +#if defined(__ARM_FEATURE_FMA) +template <> +inline float32x4_t madd(float32x4_t a, float32x4_t b, float32x4_t c) { + return vfmaq_f32(c, b, a); +} +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER) +template <> +inline float16x8_t madd(float16x8_t a, float16x8_t b, float16x8_t c) { + return vfmaq_f16(c, b, a); +} +#endif +#endif + +#if defined(__VXE__) || defined(__VXE2__) +template <> +inline float32x4_t madd(float32x4_t a, float32x4_t b, float32x4_t c) { + return vec_madd(a, b, c); +} +#endif + +#if defined(__riscv_v_intrinsic) +template <> inline vfloat32m1_t madd(vfloat32m1_t a, vfloat32m1_t b, vfloat32m1_t c) { + return __riscv_vfmacc_vv_f32m1(c, a, b, __riscv_vsetvlmax_e32m1()); +} +template <> inline vfloat32m2_t madd(vfloat32m2_t a, vfloat32m2_t b, vfloat32m2_t c) { + return __riscv_vfmacc_vv_f32m2(c, a, b, __riscv_vsetvlmax_e32m2()); +} +template <> inline vfloat32m4_t madd(vfloat32m4_t a, vfloat32m4_t b, vfloat32m4_t c) { + return __riscv_vfmacc_vv_f32m4(c, a, b, __riscv_vsetvlmax_e32m4()); +} +template <> inline vfloat32m8_t madd(vfloat32m8_t a, vfloat32m8_t b, vfloat32m8_t c) { + return __riscv_vfmacc_vv_f32m8(c, a, b, __riscv_vsetvlmax_e32m8()); +} +#endif + +#if defined(__riscv_zvfh) +template <> inline vfloat32m1_t madd(vfloat16mf2_t a, vfloat16mf2_t b, vfloat32m1_t c) { + return __riscv_vfwmacc_vv_f32m1(c, a, b, __riscv_vsetvlmax_e32m1()); +} +template <> inline vfloat32m2_t madd(vfloat16m1_t a, vfloat16m1_t b, vfloat32m2_t c) { + return __riscv_vfwmacc_vv_f32m2(c, a, b, __riscv_vsetvlmax_e32m2()); +} +template <> inline vfloat32m4_t madd(vfloat16m2_t a, vfloat16m2_t b, vfloat32m4_t c) { + return __riscv_vfwmacc_vv_f32m4(c, a, b, __riscv_vsetvlmax_e32m4()); +} +template <> inline vfloat32m8_t madd(vfloat16m4_t a, vfloat16m4_t b, vfloat32m8_t c) { + return __riscv_vfwmacc_vv_f32m8(c, a, b, __riscv_vsetvlmax_e32m8()); +} +#endif + +#if defined(__riscv_zvfbfwma) +template <> inline vfloat32m1_t madd(vbfloat16mf2_t a, vbfloat16mf2_t b, vfloat32m1_t c) { + return __riscv_vfwmaccbf16_vv_f32m1(c, a, b, __riscv_vsetvlmax_e32m1()); +} +template <> inline vfloat32m2_t madd(vbfloat16m1_t a, vbfloat16m1_t b, vfloat32m2_t c) { + return __riscv_vfwmaccbf16_vv_f32m2(c, a, b, __riscv_vsetvlmax_e32m2()); +} +template <> inline vfloat32m4_t madd(vbfloat16m2_t a, vbfloat16m2_t b, vfloat32m4_t c) { + return __riscv_vfwmaccbf16_vv_f32m4(c, a, b, __riscv_vsetvlmax_e32m4()); +} +template <> inline vfloat32m8_t madd(vbfloat16m4_t a, vbfloat16m4_t b, vfloat32m8_t c) { + return __riscv_vfwmaccbf16_vv_f32m8(c, a, b, __riscv_vsetvlmax_e32m8()); +} +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// VECTORIZED HORIZONTAL SUM + +#if defined(__ARM_NEON) +inline float hsum(float32x4_t x) { + return vaddvq_f32(x); +} +#endif // __ARM_NEON + +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER) +inline float hsum(float16x8_t x) { + return vaddvq_f32(vaddq_f32(vcvt_f32_f16(vget_low_f16(x)), + vcvt_f32_f16(vget_high_f16(x)))); +} +#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC + +#if defined(__VXE__) || defined(__VXE2__) +inline float hsum(float32x4_t x) { + float32x4_t tmp = x + vec_reve(x); + return tmp[0] + tmp[1]; +} +#endif + +#if defined(__SSE__) || defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +inline float hsum(__m128 x) { +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) + x = _mm_add_ps(x, _mm_movehl_ps(x, x)); + x = _mm_add_ss(x, _mm_movehdup_ps(x)); +#else + __m128 t; + t = _mm_shuffle_ps(x, x, _MM_SHUFFLE(2, 3, 0, 1)); + x = _mm_add_ps(x, t); + t = _mm_movehl_ps(t, x); + x = _mm_add_ss(x, t); +#endif + return _mm_cvtss_f32(x); +} +#endif + +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +inline float hsum(__m256 x) { + return hsum(_mm_add_ps(_mm256_extractf128_ps(x, 1), + _mm256_castps256_ps128(x))); +} +#endif // __AVX__ + +#if defined(__AVX512F__) +inline float hsum(__m512 x) { + return _mm512_reduce_add_ps(x); +} +#endif // __AVX512F__ + +#if defined(__riscv_v_intrinsic) +inline float hsum(vfloat32m1_t x) { + return __riscv_vfmv_f_s_f32m1_f32( + __riscv_vfredusum_vs_f32m1_f32m1(x, __riscv_vfmv_v_f_f32m1(0, 1), __riscv_vsetvlmax_e32m1())); +} +inline float hsum(vfloat32m2_t x) { + return __riscv_vfmv_f_s_f32m1_f32( + __riscv_vfredusum_vs_f32m2_f32m1(x, __riscv_vfmv_v_f_f32m1(0, 1), __riscv_vsetvlmax_e32m2())); +} +inline float hsum(vfloat32m4_t x) { + return __riscv_vfmv_f_s_f32m1_f32( + __riscv_vfredusum_vs_f32m4_f32m1(x, __riscv_vfmv_v_f_f32m1(0, 1), __riscv_vsetvlmax_e32m4())); +} +inline float hsum(vfloat32m8_t x) { + return __riscv_vfmv_f_s_f32m1_f32( + __riscv_vfredusum_vs_f32m8_f32m1(x, __riscv_vfmv_v_f_f32m1(0, 1), __riscv_vsetvlmax_e32m8())); +} +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// VECTORIZED MEMORY LOADING + +template T load(const U *); + +#if defined(__ARM_NEON) +template <> inline float32x4_t load(const float *p) { + return vld1q_f32(p); +} +#if !defined(_MSC_VER) +// FIXME: this should check for __ARM_FEATURE_FP16_VECTOR_ARITHMETIC +template <> inline float16x8_t load(const ggml_fp16_t *p) { + return vld1q_f16((const float16_t *)p); +} +template <> inline float32x4_t load(const ggml_fp16_t *p) { + return vcvt_f32_f16(vld1_f16((const float16_t *)p)); +} +#endif // _MSC_VER +#endif // __ARM_NEON + +#if defined(__VXE__) || defined(__VXE2__) +template <> inline float32x4_t load(const ggml_fp16_t * p) { + float tmp[4]; + + for (int i = 0; i < 4; i++) { + tmp[i] = GGML_CPU_FP16_TO_FP32(p[i]); + } + + return vec_xl(0, (const float *)(tmp)); +} +template <> inline float32x4_t load(const float * p) { + return vec_xl(0, p); +} +#endif + +#if defined(__SSE__) || defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +template <> inline __m128 load(const float *p) { + return _mm_loadu_ps(p); +} +#endif // __SSE__ + +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +template <> inline __m256 load(const float *p) { + return _mm256_loadu_ps(p); +} +#endif // __AVX__ + +#if defined(__AVX2__) || defined(__AVX512F__) +template <> inline __m256 load(const ggml_bf16_t *p) { + return _mm256_castsi256_ps( + _mm256_slli_epi32(_mm256_cvtepu16_epi32(_mm_loadu_si128((const __m128i *)p)), 16)); +} +#endif // __AVX2__ + +#if defined(__F16C__) +template <> inline __m256 load(const ggml_fp16_t *p) { + return _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)p)); +} +#endif // __F16C__ + +#if defined(__AVX512F__) +template <> inline __m512 load(const float *p) { + return _mm512_loadu_ps(p); +} +template <> inline __m512 load(const ggml_fp16_t *p) { + return _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)p)); +} +template <> inline __m512 load(const ggml_bf16_t *p) { + return _mm512_castsi512_ps( + _mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)p)), 16)); +} +#endif // __AVX512F__ + +#if defined(__AVX512BF16__) +template <> inline __m512bh load(const ggml_bf16_t *p) { + return (__m512bh)_mm512_loadu_ps((const float *)p); +} +template <> inline __m256bh load(const ggml_bf16_t *p) { + return (__m256bh)_mm256_loadu_ps((const float *)p); +} +template <> inline __m512bh load(const float *p) { + return _mm512_cvtne2ps_pbh(_mm512_loadu_ps(p + 16), _mm512_loadu_ps(p)); +} +template <> inline __m256bh load(const float *p) { + return _mm512_cvtneps_pbh(_mm512_loadu_ps(p)); +} +#endif + +#if defined(__riscv_v_intrinsic) +template <> inline vfloat32m1_t load(const float *p) { + return __riscv_vle32_v_f32m1(p, __riscv_vsetvlmax_e32m1()); +} +template <> inline vfloat32m2_t load(const float *p) { + return __riscv_vle32_v_f32m2(p, __riscv_vsetvlmax_e32m2()); +} +template <> inline vfloat32m4_t load(const float *p) { + return __riscv_vle32_v_f32m4(p, __riscv_vsetvlmax_e32m4()); +} +template <> inline vfloat32m8_t load(const float *p) { + return __riscv_vle32_v_f32m8(p, __riscv_vsetvlmax_e32m8()); +} +#endif + +#if defined(__riscv_zvfh) +template <> inline vfloat16mf2_t load(const ggml_fp16_t *p) { + return __riscv_vle16_v_f16mf2(reinterpret_cast(p), __riscv_vsetvlmax_e16mf2()); +} +template <> inline vfloat16m1_t load(const ggml_fp16_t *p) { + return __riscv_vle16_v_f16m1(reinterpret_cast(p), __riscv_vsetvlmax_e16m1()); +} +template <> inline vfloat16m2_t load(const ggml_fp16_t *p) { + return __riscv_vle16_v_f16m2(reinterpret_cast(p), __riscv_vsetvlmax_e16m2()); +} +template <> inline vfloat16m4_t load(const ggml_fp16_t *p) { + return __riscv_vle16_v_f16m4(reinterpret_cast(p), __riscv_vsetvlmax_e16m4()); +} +#endif + +#if defined(__riscv_zvfbfwma) +template <> inline vbfloat16mf2_t load(const ggml_bf16_t *p) { + return __riscv_vle16_v_bf16mf2(reinterpret_cast(p), __riscv_vsetvlmax_e16mf2()); +} +template <> inline vbfloat16m1_t load(const ggml_bf16_t *p) { + return __riscv_vle16_v_bf16m1(reinterpret_cast(p), __riscv_vsetvlmax_e16m1()); +} +template <> inline vbfloat16m2_t load(const ggml_bf16_t *p) { + return __riscv_vle16_v_bf16m2(reinterpret_cast(p), __riscv_vsetvlmax_e16m2()); +} +template <> inline vbfloat16m4_t load(const ggml_bf16_t *p) { + return __riscv_vle16_v_bf16m4(reinterpret_cast(p), __riscv_vsetvlmax_e16m4()); +} +#endif + +#if defined(__riscv_v_intrinsic) +template T set_zero(); + +template <> inline vfloat32m1_t set_zero() { + return __riscv_vfmv_v_f_f32m1(0.0f, __riscv_vsetvlmax_e32m1()); +} +template <> inline vfloat32m2_t set_zero() { + return __riscv_vfmv_v_f_f32m2(0, __riscv_vsetvlmax_e32m2()); +} +template <> inline vfloat32m4_t set_zero() { + return __riscv_vfmv_v_f_f32m4(0, __riscv_vsetvlmax_e32m4()); +} +template <> inline vfloat32m8_t set_zero() { + return __riscv_vfmv_v_f_f32m8(0, __riscv_vsetvlmax_e32m8()); +} +#endif + +#if defined(__riscv_v_intrinsic) +template size_t vlmax() { + if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e32m1(); } + else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e32m2(); } + else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e32m4(); } + else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e32m8(); } + #if defined (__riscv_zvfh) + else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16mf2(); } + else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16m1(); } + else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16m2(); } + else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16m4(); } + #endif + #if defined (__riscv_zvfbfwma) + else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16mf2(); } + else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16m1(); } + else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16m2(); } + else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16m4(); } + #endif + return 0; +} +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// FLOATING POINT MATRIX MULTIPLICATION + +template +static inline int64_t BLOCK_SIZE(size_t m) { + const int64_t NB_BLOC_M = (m + M - 1) / M; + return (m % NB_BLOC_M == 0) ? m / NB_BLOC_M : (m / NB_BLOC_M) + 1; +} + +static constexpr inline int64_t BLOC_POS(int64_t ib, int64_t ibN, int64_t bloc_size) { + return ib < ibN ? ib * bloc_size : ibN * bloc_size + (ib - ibN) * (bloc_size - 1); +} + +template +class tinyBLAS { + public: + tinyBLAS(const ggml_compute_params * params, int64_t k, + const TA *A, int64_t lda, + const TB *B, int64_t ldb, + TC *C, int64_t ldc) + : params(params), A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc) { + } + + bool matmul(int64_t m, int64_t n) { + if (k % KN != 0) + return false; + // compute RM for only need tile with size RM&RM-1 +#if VECTOR_REGISTERS == 32 + if (m % 16 == 0 && (m/16 >= params->nth)) { + const int64_t SIZE_N = BLOCK_SIZE<6>(n); + mnpack<4, 6, 4>(m, n, SIZE_N, 12); + return true; + } + if (m % 8 == 0 ) { + const int64_t SIZE_N = BLOCK_SIZE<6>(n); + mnpack<4, 6, 2>(m, n, SIZE_N, 12); + return true; + } + if (m % 4 == 0) { + const int64_t SIZE_N = BLOCK_SIZE<6>(n); + mnpack<4, 6, 1>(m, n, SIZE_N, 12); + return true; + } +#else // VECTOR_REGISTERS == 16 + if (m % 16 == 0 && (m/16 >= params->nth)) { + const int64_t SIZE_N = BLOCK_SIZE<3>(n); + mnpack<4, 3, 4>(m, n, SIZE_N, 24); + return true; + } + if (m % 8 == 0 ) { + const int64_t SIZE_N = BLOCK_SIZE<3>(n); + mnpack<4, 3, 2>(m, n, SIZE_N, 24); + return true; + } + if (m % 4 == 0) { + const int64_t SIZE_N = BLOCK_SIZE<3>(n); + mnpack<4, 3, 1>(m, n, SIZE_N, 24); + return true; + } +#endif + return false; + } + + private: + template + inline void mnpack(int64_t m, int64_t n, int64_t SIZE_N, int64_t BN) { + if (SIZE_N == RN) { + return gemm(m, n, BN); + } + if constexpr (RN > 1) { + return mnpack(m, n, SIZE_N, BN); + } else { + GGML_LOG_ERROR("mnpack<%d, %d> block size not supported\n", RM, (int)SIZE_N); + GGML_ASSERT(false); // we have miss something. + } + } + + template + inline void gemm_bloc(int64_t ii, int64_t jj) { + D Cv[RN][RM] = {}; + for (int64_t l = 0; l < k; l += KN) { + // help compiler for op order. + if constexpr (RM <= RN) { + V Av[RM]; + for (int64_t i = 0; i < RM; ++i) { + Av[i] = load(A + lda * (ii + i) + l); + } + for (int64_t j = 0; j < RN; ++j) { + V Bv = load(B + ldb * (jj + j) + l); + for (int64_t i = 0; i < RM; ++i) { + Cv[j][i] = madd(Av[i], Bv, Cv[j][i]); + } + } + } else { + V Bv[RN]; + for (int64_t j = 0; j < RN; ++j) { + Bv[j] = load(B + ldb * (jj + j) + l); + } + for (int64_t i = 0; i < RM; ++i) { + V Av = load(A + lda * (ii + i) + l); + for (int64_t j = 0; j < RN; ++j) { + Cv[j][i] = madd(Av, Bv[j], Cv[j][i]); + } + } + } + } + for (int64_t j = 0; j < RN; ++j) + for (int64_t i = 0; i < RM; ++i) + C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]); + } + + template + NOINLINE void gemm(int64_t m, int64_t n, int64_t BN) { + GGML_ASSERT(m % (RM * BM) == 0); + const int64_t ytiles = m / (RM * BM); + const int64_t xtiles = (n + RN -1) / RN; + const int64_t jj_RN = (xtiles - (xtiles * RN - n)); + + // "round" bloc_size to "nearest" BN + const int64_t NB_BN = xtiles < BN ? 1 : (xtiles + BN / 2) / BN; + const int64_t SIZE_BN = xtiles % NB_BN == 0 ? xtiles / NB_BN : xtiles / NB_BN + 1; + const int64_t jj_BN = (NB_BN - (NB_BN * SIZE_BN - xtiles)); + const int64_t nb_job = ytiles * NB_BN; + + if (params->ith == 0) { + GGML_ASSERT( jj_BN * SIZE_BN + (NB_BN - jj_BN) * (SIZE_BN - 1) == xtiles); + // Every thread starts at ith, so the first unprocessed chunk is nth. This save a bit of coordination right at the start. + ggml_threadpool_chunk_set(params->threadpool, params->nth); + } + + ggml_barrier(params->threadpool); + + int64_t job = params->ith; + while (job < nb_job) { + const int64_t ii = (job % ytiles) * RM * BM; + const int64_t jb = job / ytiles; + const int64_t jr0 = BLOC_POS(jb , jj_BN, SIZE_BN); + const int64_t jrN = BLOC_POS(jb+1, jj_BN, SIZE_BN); + + const int64_t jj0 = BLOC_POS(jr0, jj_RN, RN); + const int64_t jj2 = BLOC_POS(jrN, jj_RN, RN); + const int64_t jj1 = jj2 < jj_RN * RN ? jj2 : jj_RN * RN; + + for (int64_t bi = 0; bi < BM * RM; bi += RM) { + int64_t jj = jj0; + for (; jj < jj1; jj += RN) { + gemm_bloc(ii + bi, jj); + } + if constexpr (RN > 1) { + for (; jj < jj2; jj += RN - 1) { + gemm_bloc(ii + bi, jj); + } + } + GGML_ASSERT(jj == jj2); + } + + job = ggml_threadpool_chunk_add(params->threadpool, 1); + } + + ggml_barrier(params->threadpool); + return; + } + + const ggml_compute_params * params; + const TA *const A; + const TB *const B; + TC *const C; + const int64_t k; + const int64_t lda; + const int64_t ldb; + const int64_t ldc; +}; + +#if defined(__riscv_v_intrinsic) +template +class tinyBLAS_RVV { + public: + tinyBLAS_RVV(const ggml_compute_params * params, int64_t k, + const TA *A, int64_t lda, + const TB *B, int64_t ldb, + TC *C, int64_t ldc) + : params(params), A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc) { + } + + bool matmul(int64_t m, int64_t n) { + if (k % vlmax() != 0) { + return false; + } + +#if LMUL == 1 + if (m % 16 == 0 && (m/16 >= params->nth)) { + const int64_t SIZE_N = BLOCK_SIZE<6>(n); + mnpack<4, 6, 4>(m, n, SIZE_N, 12); + return true; + } + if (m % 8 == 0 ) { + const int64_t SIZE_N = BLOCK_SIZE<6>(n); + mnpack<4, 6, 2>(m, n, SIZE_N, 12); + return true; + } + if (m % 4 == 0) { + const int64_t SIZE_N = BLOCK_SIZE<6>(n); + mnpack<4, 6, 1>(m, n, SIZE_N, 12); + return true; + } +#elif LMUL == 2 + if (m % 16 == 0 && (m/16 >= params->nth)) { + const int64_t SIZE_N = BLOCK_SIZE<3>(n); + mnpack<4, 3, 4>(m, n, SIZE_N, 24); + return true; + } + if (m % 8 == 0 ) { + const int64_t SIZE_N = BLOCK_SIZE<3>(n); + mnpack<4, 3, 2>(m, n, SIZE_N, 24); + return true; + } + if (m % 4 == 0) { + const int64_t SIZE_N = BLOCK_SIZE<3>(n); + mnpack<4, 3, 1>(m, n, SIZE_N, 24); + return true; + } +#else // LMUL = 4 + if (m % 16 == 0 && (m/16 >= params->nth)) { + const int64_t SIZE_N = BLOCK_SIZE<2>(n); + mnpack<2, 2, 8>(m, n, SIZE_N, 36); + return true; + } + if (m % 8 == 0 ) { + const int64_t SIZE_N = BLOCK_SIZE<2>(n); + mnpack<2, 2, 4>(m, n, SIZE_N, 36); + return true; + } + if (m % 4 == 0) { + const int64_t SIZE_N = BLOCK_SIZE<2>(n); + mnpack<2, 2, 2>(m, n, SIZE_N, 36); + return true; + } +#endif + return false; + } + + private: + template + inline void mnpack(int64_t m, int64_t n, int64_t SIZE_N, int64_t BN) { + if (SIZE_N == RN) { + return gemm(m, n, BN); + } + if constexpr (RN > 1) { + return mnpack(m, n, SIZE_N, BN); + } else { + GGML_LOG_ERROR("mnpack<%d, %d> block size not supported\n", RM, (int)SIZE_N); + GGML_ASSERT(false); // we have miss something. + } + } + + inline void gemm_bloc_4x6(int64_t ii, int64_t jj) { + size_t vl = vlmax(); + D Cv00 = set_zero(); + D Cv01 = set_zero(); + D Cv02 = set_zero(); + D Cv03 = set_zero(); + D Cv10 = set_zero(); + D Cv11 = set_zero(); + D Cv12 = set_zero(); + D Cv13 = set_zero(); + D Cv20 = set_zero(); + D Cv21 = set_zero(); + D Cv22 = set_zero(); + D Cv23 = set_zero(); + D Cv30 = set_zero(); + D Cv31 = set_zero(); + D Cv32 = set_zero(); + D Cv33 = set_zero(); + D Cv40 = set_zero(); + D Cv41 = set_zero(); + D Cv42 = set_zero(); + D Cv43 = set_zero(); + D Cv50 = set_zero(); + D Cv51 = set_zero(); + D Cv52 = set_zero(); + D Cv53 = set_zero(); + + for (int64_t l = 0; l < k; l += vl) { + V Bv0 = load(B + ldb * (jj + 0) + l); + V Bv1 = load(B + ldb * (jj + 1) + l); + V Bv2 = load(B + ldb * (jj + 2) + l); + V Bv3 = load(B + ldb * (jj + 3) + l); + V Bv4 = load(B + ldb * (jj + 4) + l); + V Bv5 = load(B + ldb * (jj + 5) + l); + + V Av0 = load(A + lda * (ii + 0) + l); + Cv00 = madd(Av0, Bv0, Cv00); + Cv10 = madd(Av0, Bv1, Cv10); + Cv20 = madd(Av0, Bv2, Cv20); + Cv30 = madd(Av0, Bv3, Cv30); + Cv40 = madd(Av0, Bv4, Cv40); + Cv50 = madd(Av0, Bv5, Cv50); + + V Av1 = load(A + lda * (ii + 1) + l); + Cv01 = madd(Av1, Bv0, Cv01); + Cv11 = madd(Av1, Bv1, Cv11); + Cv21 = madd(Av1, Bv2, Cv21); + Cv31 = madd(Av1, Bv3, Cv31); + Cv41 = madd(Av1, Bv4, Cv41); + Cv51 = madd(Av1, Bv5, Cv51); + + V Av2 = load(A + lda * (ii + 2) + l); + Cv02 = madd(Av2, Bv0, Cv02); + Cv12 = madd(Av2, Bv1, Cv12); + Cv22 = madd(Av2, Bv2, Cv22); + Cv32 = madd(Av2, Bv3, Cv32); + Cv42 = madd(Av2, Bv4, Cv42); + Cv52 = madd(Av2, Bv5, Cv52); + + V Av3 = load(A + lda * (ii + 3) + l); + Cv03 = madd(Av3, Bv0, Cv03); + Cv13 = madd(Av3, Bv1, Cv13); + Cv23 = madd(Av3, Bv2, Cv23); + Cv33 = madd(Av3, Bv3, Cv33); + Cv43 = madd(Av3, Bv4, Cv43); + Cv53 = madd(Av3, Bv5, Cv53); + } + + C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); + C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); + C[ldc * (jj + 0) + (ii + 2)] = hsum(Cv02); + C[ldc * (jj + 0) + (ii + 3)] = hsum(Cv03); + C[ldc * (jj + 1) + (ii + 0)] = hsum(Cv10); + C[ldc * (jj + 1) + (ii + 1)] = hsum(Cv11); + C[ldc * (jj + 1) + (ii + 2)] = hsum(Cv12); + C[ldc * (jj + 1) + (ii + 3)] = hsum(Cv13); + C[ldc * (jj + 2) + (ii + 0)] = hsum(Cv20); + C[ldc * (jj + 2) + (ii + 1)] = hsum(Cv21); + C[ldc * (jj + 2) + (ii + 2)] = hsum(Cv22); + C[ldc * (jj + 2) + (ii + 3)] = hsum(Cv23); + C[ldc * (jj + 3) + (ii + 0)] = hsum(Cv30); + C[ldc * (jj + 3) + (ii + 1)] = hsum(Cv31); + C[ldc * (jj + 3) + (ii + 2)] = hsum(Cv32); + C[ldc * (jj + 3) + (ii + 3)] = hsum(Cv33); + C[ldc * (jj + 4) + (ii + 0)] = hsum(Cv40); + C[ldc * (jj + 4) + (ii + 1)] = hsum(Cv41); + C[ldc * (jj + 4) + (ii + 2)] = hsum(Cv42); + C[ldc * (jj + 4) + (ii + 3)] = hsum(Cv43); + C[ldc * (jj + 5) + (ii + 0)] = hsum(Cv50); + C[ldc * (jj + 5) + (ii + 1)] = hsum(Cv51); + C[ldc * (jj + 5) + (ii + 2)] = hsum(Cv52); + C[ldc * (jj + 5) + (ii + 3)] = hsum(Cv53); + } + + inline void gemm_bloc_4x5(int64_t ii, int64_t jj) { + size_t vl = vlmax(); + D Cv00 = set_zero(); + D Cv01 = set_zero(); + D Cv02 = set_zero(); + D Cv03 = set_zero(); + D Cv10 = set_zero(); + D Cv11 = set_zero(); + D Cv12 = set_zero(); + D Cv13 = set_zero(); + D Cv20 = set_zero(); + D Cv21 = set_zero(); + D Cv22 = set_zero(); + D Cv23 = set_zero(); + D Cv30 = set_zero(); + D Cv31 = set_zero(); + D Cv32 = set_zero(); + D Cv33 = set_zero(); + D Cv40 = set_zero(); + D Cv41 = set_zero(); + D Cv42 = set_zero(); + D Cv43 = set_zero(); + + for (int64_t l = 0; l < k; l += vl) { + V Bv0 = load(B + ldb * (jj + 0) + l); + V Bv1 = load(B + ldb * (jj + 1) + l); + V Bv2 = load(B + ldb * (jj + 2) + l); + V Bv3 = load(B + ldb * (jj + 3) + l); + V Bv4 = load(B + ldb * (jj + 4) + l); + + V Av0 = load(A + lda * (ii + 0) + l); + Cv00 = madd(Av0, Bv0, Cv00); + Cv10 = madd(Av0, Bv1, Cv10); + Cv20 = madd(Av0, Bv2, Cv20); + Cv30 = madd(Av0, Bv3, Cv30); + Cv40 = madd(Av0, Bv4, Cv40); + + V Av1 = load(A + lda * (ii + 1) + l); + Cv01 = madd(Av1, Bv0, Cv01); + Cv11 = madd(Av1, Bv1, Cv11); + Cv21 = madd(Av1, Bv2, Cv21); + Cv31 = madd(Av1, Bv3, Cv31); + Cv41 = madd(Av1, Bv4, Cv41); + + V Av2 = load(A + lda * (ii + 2) + l); + Cv02 = madd(Av2, Bv0, Cv02); + Cv12 = madd(Av2, Bv1, Cv12); + Cv22 = madd(Av2, Bv2, Cv22); + Cv32 = madd(Av2, Bv3, Cv32); + Cv42 = madd(Av2, Bv4, Cv42); + + V Av3 = load(A + lda * (ii + 3) + l); + Cv03 = madd(Av3, Bv0, Cv03); + Cv13 = madd(Av3, Bv1, Cv13); + Cv23 = madd(Av3, Bv2, Cv23); + Cv33 = madd(Av3, Bv3, Cv33); + Cv43 = madd(Av3, Bv4, Cv43); + } + + C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); + C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); + C[ldc * (jj + 0) + (ii + 2)] = hsum(Cv02); + C[ldc * (jj + 0) + (ii + 3)] = hsum(Cv03); + C[ldc * (jj + 1) + (ii + 0)] = hsum(Cv10); + C[ldc * (jj + 1) + (ii + 1)] = hsum(Cv11); + C[ldc * (jj + 1) + (ii + 2)] = hsum(Cv12); + C[ldc * (jj + 1) + (ii + 3)] = hsum(Cv13); + C[ldc * (jj + 2) + (ii + 0)] = hsum(Cv20); + C[ldc * (jj + 2) + (ii + 1)] = hsum(Cv21); + C[ldc * (jj + 2) + (ii + 2)] = hsum(Cv22); + C[ldc * (jj + 2) + (ii + 3)] = hsum(Cv23); + C[ldc * (jj + 3) + (ii + 0)] = hsum(Cv30); + C[ldc * (jj + 3) + (ii + 1)] = hsum(Cv31); + C[ldc * (jj + 3) + (ii + 2)] = hsum(Cv32); + C[ldc * (jj + 3) + (ii + 3)] = hsum(Cv33); + C[ldc * (jj + 4) + (ii + 0)] = hsum(Cv40); + C[ldc * (jj + 4) + (ii + 1)] = hsum(Cv41); + C[ldc * (jj + 4) + (ii + 2)] = hsum(Cv42); + C[ldc * (jj + 4) + (ii + 3)] = hsum(Cv43); + } + + inline void gemm_bloc_4x4(int64_t ii, int64_t jj) { + size_t vl = vlmax(); + D Cv00 = set_zero(); + D Cv01 = set_zero(); + D Cv02 = set_zero(); + D Cv03 = set_zero(); + D Cv10 = set_zero(); + D Cv11 = set_zero(); + D Cv12 = set_zero(); + D Cv13 = set_zero(); + D Cv20 = set_zero(); + D Cv21 = set_zero(); + D Cv22 = set_zero(); + D Cv23 = set_zero(); + D Cv30 = set_zero(); + D Cv31 = set_zero(); + D Cv32 = set_zero(); + D Cv33 = set_zero(); + + for (int64_t l = 0; l < k; l += vl) { + V Av0 = load(A + lda * (ii + 0) + l); + V Av1 = load(A + lda * (ii + 1) + l); + V Av2 = load(A + lda * (ii + 2) + l); + V Av3 = load(A + lda * (ii + 3) + l); + + V Bv0 = load(B + ldb * (jj + 0) + l); + Cv00 = madd(Av0, Bv0, Cv00); + Cv01 = madd(Av1, Bv0, Cv01); + Cv02 = madd(Av2, Bv0, Cv02); + Cv03 = madd(Av3, Bv0, Cv03); + + V Bv1 = load(B + ldb * (jj + 1) + l); + Cv10 = madd(Av0, Bv1, Cv10); + Cv11 = madd(Av1, Bv1, Cv11); + Cv12 = madd(Av2, Bv1, Cv12); + Cv13 = madd(Av3, Bv1, Cv13); + + V Bv2 = load(B + ldb * (jj + 2) + l); + Cv20 = madd(Av0, Bv2, Cv20); + Cv21 = madd(Av1, Bv2, Cv21); + Cv22 = madd(Av2, Bv2, Cv22); + Cv23 = madd(Av3, Bv2, Cv23); + + V Bv3 = load(B + ldb * (jj + 3) + l); + Cv30 = madd(Av0, Bv3, Cv30); + Cv31 = madd(Av1, Bv3, Cv31); + Cv32 = madd(Av2, Bv3, Cv32); + Cv33 = madd(Av3, Bv3, Cv33); + } + + C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); + C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); + C[ldc * (jj + 0) + (ii + 2)] = hsum(Cv02); + C[ldc * (jj + 0) + (ii + 3)] = hsum(Cv03); + C[ldc * (jj + 1) + (ii + 0)] = hsum(Cv10); + C[ldc * (jj + 1) + (ii + 1)] = hsum(Cv11); + C[ldc * (jj + 1) + (ii + 2)] = hsum(Cv12); + C[ldc * (jj + 1) + (ii + 3)] = hsum(Cv13); + C[ldc * (jj + 2) + (ii + 0)] = hsum(Cv20); + C[ldc * (jj + 2) + (ii + 1)] = hsum(Cv21); + C[ldc * (jj + 2) + (ii + 2)] = hsum(Cv22); + C[ldc * (jj + 2) + (ii + 3)] = hsum(Cv23); + C[ldc * (jj + 3) + (ii + 0)] = hsum(Cv30); + C[ldc * (jj + 3) + (ii + 1)] = hsum(Cv31); + C[ldc * (jj + 3) + (ii + 2)] = hsum(Cv32); + C[ldc * (jj + 3) + (ii + 3)] = hsum(Cv33); + } + + inline void gemm_bloc_4x3(int64_t ii, int64_t jj) { + size_t vl = vlmax(); + D Cv00 = set_zero(); + D Cv01 = set_zero(); + D Cv02 = set_zero(); + D Cv03 = set_zero(); + D Cv10 = set_zero(); + D Cv11 = set_zero(); + D Cv12 = set_zero(); + D Cv13 = set_zero(); + D Cv20 = set_zero(); + D Cv21 = set_zero(); + D Cv22 = set_zero(); + D Cv23 = set_zero(); + + for (int64_t l = 0; l < k; l += vl) { + V Av0 = load(A + lda * (ii + 0) + l); + V Av1 = load(A + lda * (ii + 1) + l); + V Av2 = load(A + lda * (ii + 2) + l); + V Av3 = load(A + lda * (ii + 3) + l); + + V Bv0 = load(B + ldb * (jj + 0) + l); + Cv00 = madd(Av0, Bv0, Cv00); + Cv01 = madd(Av1, Bv0, Cv01); + Cv02 = madd(Av2, Bv0, Cv02); + Cv03 = madd(Av3, Bv0, Cv03); + + V Bv1 = load(B + ldb * (jj + 1) + l); + Cv10 = madd(Av0, Bv1, Cv10); + Cv11 = madd(Av1, Bv1, Cv11); + Cv12 = madd(Av2, Bv1, Cv12); + Cv13 = madd(Av3, Bv1, Cv13); + + V Bv2 = load(B + ldb * (jj + 2) + l); + Cv20 = madd(Av0, Bv2, Cv20); + Cv21 = madd(Av1, Bv2, Cv21); + Cv22 = madd(Av2, Bv2, Cv22); + Cv23 = madd(Av3, Bv2, Cv23); + } + + C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); + C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); + C[ldc * (jj + 0) + (ii + 2)] = hsum(Cv02); + C[ldc * (jj + 0) + (ii + 3)] = hsum(Cv03); + C[ldc * (jj + 1) + (ii + 0)] = hsum(Cv10); + C[ldc * (jj + 1) + (ii + 1)] = hsum(Cv11); + C[ldc * (jj + 1) + (ii + 2)] = hsum(Cv12); + C[ldc * (jj + 1) + (ii + 3)] = hsum(Cv13); + C[ldc * (jj + 2) + (ii + 0)] = hsum(Cv20); + C[ldc * (jj + 2) + (ii + 1)] = hsum(Cv21); + C[ldc * (jj + 2) + (ii + 2)] = hsum(Cv22); + C[ldc * (jj + 2) + (ii + 3)] = hsum(Cv23); + } + + inline void gemm_bloc_4x2(int64_t ii, int64_t jj) { + size_t vl = vlmax(); + D Cv00 = set_zero(); + D Cv01 = set_zero(); + D Cv02 = set_zero(); + D Cv03 = set_zero(); + D Cv10 = set_zero(); + D Cv11 = set_zero(); + D Cv12 = set_zero(); + D Cv13 = set_zero(); + + for (int64_t l = 0; l < k; l += vl) { + V Av0 = load(A + lda * (ii + 0) + l); + V Av1 = load(A + lda * (ii + 1) + l); + V Av2 = load(A + lda * (ii + 2) + l); + V Av3 = load(A + lda * (ii + 3) + l); + + V Bv0 = load(B + ldb * (jj + 0) + l); + Cv00 = madd(Av0, Bv0, Cv00); + Cv01 = madd(Av1, Bv0, Cv01); + Cv02 = madd(Av2, Bv0, Cv02); + Cv03 = madd(Av3, Bv0, Cv03); + + V Bv1 = load(B + ldb * (jj + 1) + l); + Cv10 = madd(Av0, Bv1, Cv10); + Cv11 = madd(Av1, Bv1, Cv11); + Cv12 = madd(Av2, Bv1, Cv12); + Cv13 = madd(Av3, Bv1, Cv13); + } + + C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); + C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); + C[ldc * (jj + 0) + (ii + 2)] = hsum(Cv02); + C[ldc * (jj + 0) + (ii + 3)] = hsum(Cv03); + C[ldc * (jj + 1) + (ii + 0)] = hsum(Cv10); + C[ldc * (jj + 1) + (ii + 1)] = hsum(Cv11); + C[ldc * (jj + 1) + (ii + 2)] = hsum(Cv12); + C[ldc * (jj + 1) + (ii + 3)] = hsum(Cv13); + } + + inline void gemm_bloc_4x1(int64_t ii, int64_t jj) { + size_t vl = vlmax(); + D Cv00 = set_zero(); + D Cv01 = set_zero(); + D Cv02 = set_zero(); + D Cv03 = set_zero(); + + for (int64_t l = 0; l < k; l += vl) { + V Av0 = load(A + lda * (ii + 0) + l); + V Av1 = load(A + lda * (ii + 1) + l); + V Av2 = load(A + lda * (ii + 2) + l); + V Av3 = load(A + lda * (ii + 3) + l); + + V Bv0 = load(B + ldb * (jj + 0) + l); + Cv00 = madd(Av0, Bv0, Cv00); + Cv01 = madd(Av1, Bv0, Cv01); + Cv02 = madd(Av2, Bv0, Cv02); + Cv03 = madd(Av3, Bv0, Cv03); + } + + C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); + C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); + C[ldc * (jj + 0) + (ii + 2)] = hsum(Cv02); + C[ldc * (jj + 0) + (ii + 3)] = hsum(Cv03); + } + + inline void gemm_bloc_2x2(int64_t ii, int64_t jj) { + size_t vl = vlmax(); + D Cv00 = set_zero(); + D Cv01 = set_zero(); + D Cv10 = set_zero(); + D Cv11 = set_zero(); + + for (int64_t l = 0; l < k; l += vl) { + V Av0 = load(A + lda * (ii + 0) + l); + V Av1 = load(A + lda * (ii + 1) + l); + + V Bv0 = load(B + ldb * (jj + 0) + l); + Cv00 = madd(Av0, Bv0, Cv00); + Cv01 = madd(Av1, Bv0, Cv01); + + V Bv1 = load(B + ldb * (jj + 1) + l); + Cv10 = madd(Av0, Bv1, Cv10); + Cv11 = madd(Av1, Bv1, Cv11); + } + + C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); + C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); + C[ldc * (jj + 1) + (ii + 0)] = hsum(Cv10); + C[ldc * (jj + 1) + (ii + 1)] = hsum(Cv11); + } + + inline void gemm_bloc_2x1(int64_t ii, int64_t jj) { + size_t vl = vlmax(); + D Cv00 = set_zero(); + D Cv01 = set_zero(); + + for (int64_t l = 0; l < k; l += vl) { + V Av0 = load(A + lda * (ii + 0) + l); + V Av1 = load(A + lda * (ii + 1) + l); + + V Bv0 = load(B + ldb * (jj + 0) + l); + Cv00 = madd(Av0, Bv0, Cv00); + Cv01 = madd(Av1, Bv0, Cv01); + } + + C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); + C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); + } + + template + inline void gemm_bloc(int64_t ii, int64_t jj) { + if constexpr (RM == 4) { + if constexpr (RN == 6) { return gemm_bloc_4x6(ii, jj); } + if constexpr (RN == 5) { return gemm_bloc_4x5(ii, jj); } + if constexpr (RN == 4) { return gemm_bloc_4x4(ii, jj); } + if constexpr (RN == 3) { return gemm_bloc_4x3(ii, jj); } + if constexpr (RN == 2) { return gemm_bloc_4x2(ii, jj); } + if constexpr (RN == 1) { return gemm_bloc_4x1(ii, jj); } + } else if constexpr (RM == 2) { + if constexpr (RN == 2) { return gemm_bloc_2x2(ii, jj); } + if constexpr (RN == 1) { return gemm_bloc_2x1(ii, jj); } + } + } + + template + NOINLINE void gemm(int64_t m, int64_t n, int64_t BN) { + GGML_ASSERT(m % (RM * BM) == 0); + const int64_t ytiles = m / (RM * BM); + const int64_t xtiles = (n + RN -1) / RN; + const int64_t jj_RN = (xtiles - (xtiles * RN - n)); + + // "round" bloc_size to "nearest" BN + const int64_t NB_BN = xtiles < BN ? 1 : (xtiles + BN / 2) / BN; + const int64_t SIZE_BN = xtiles % NB_BN == 0 ? xtiles / NB_BN : xtiles / NB_BN + 1; + const int64_t jj_BN = (NB_BN - (NB_BN * SIZE_BN - xtiles)); + const int64_t nb_job = ytiles * NB_BN; + + if (params->ith == 0) { + GGML_ASSERT( jj_BN * SIZE_BN + (NB_BN - jj_BN) * (SIZE_BN - 1) == xtiles); + // Every thread starts at ith, so the first unprocessed chunk is nth. This save a bit of coordination right at the start. + ggml_threadpool_chunk_set(params->threadpool, params->nth); + } + + ggml_barrier(params->threadpool); + + int64_t job = params->ith; + while (job < nb_job) { + const int64_t ii = (job % ytiles) * RM * BM; + const int64_t jb = job / ytiles; + const int64_t jr0 = BLOC_POS(jb , jj_BN, SIZE_BN); + const int64_t jrN = BLOC_POS(jb+1, jj_BN, SIZE_BN); + + const int64_t jj0 = BLOC_POS(jr0, jj_RN, RN); + const int64_t jj2 = BLOC_POS(jrN, jj_RN, RN); + const int64_t jj1 = jj2 < jj_RN * RN ? jj2 : jj_RN * RN; + + for (int64_t bi = 0; bi < BM * RM; bi += RM) { + int64_t jj = jj0; + for (; jj < jj1; jj += RN) { + gemm_bloc(ii + bi, jj); + } + if constexpr (RN > 1) { + for (; jj < jj2; jj += RN - 1) { + gemm_bloc(ii + bi, jj); + } + } + GGML_ASSERT(jj == jj2); + } + + job = ggml_threadpool_chunk_add(params->threadpool, 1); + } + + ggml_barrier(params->threadpool); + return; + } + + const ggml_compute_params * params; + const TA *const A; + const TB *const B; + TC *const C; + const int64_t k; + const int64_t lda; + const int64_t ldb; + const int64_t ldc; +}; +#endif + +////////////////////////////////////////////////////////////////////////////////////////// +// QUANT ZERO MATRIX MULTIPLICATION + +#if defined(__ARM_FEATURE_DOTPROD) +template +class tinyBLAS_Q0_ARM { + public: + tinyBLAS_Q0_ARM(int64_t k, + const TA *A, int64_t lda, + const block_q8_0 *B, int64_t ldb, + float *C, int64_t ldc, + int ith, int nth) + : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { + } + + void matmul(int64_t m, int64_t n) { + mnpack(0, m, 0, n); + } + + private: + NOINLINE void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t mc, nc, mp, np; + switch ((MIN(m - m0, 3) << 4) | MIN(n - n0, 3ll)) { + case 0x33: + mc = 3; + nc = 3; + gemm<3, 3>(m0, m, n0, n); + break; + case 0x32: + mc = 3; + nc = 2; + gemm<3, 2>(m0, m, n0, n); + break; + case 0x23: + mc = 2; + nc = 3; + gemm<2, 3>(m0, m, n0, n); + break; + case 0x22: + mc = 2; + nc = 2; + gemm<2, 2>(m0, m, n0, n); + break; + case 0x31: + mc = 3; + nc = 1; + gemm<3, 1>(m0, m, n0, n); + break; + case 0x13: + mc = 1; + nc = 3; + gemm<1, 3>(m0, m, n0, n); + break; + case 0x21: + mc = 2; + nc = 1; + gemm<2, 1>(m0, m, n0, n); + break; + case 0x12: + mc = 1; + nc = 2; + gemm<1, 2>(m0, m, n0, n); + break; + case 0x11: + mc = 1; + nc = 1; + gemm<1, 1>(m0, m, n0, n); + break; + default: + return; + } + mp = m0 + (m - m0) / mc * mc; + np = n0 + (n - n0) / nc * nc; + mnpack(mp, m, n0, np); + mnpack(m0, m, np, n); + } + + template + NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t ytiles = (m - m0) / RM; + int64_t xtiles = (n - n0) / RN; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) + end = tiles; + for (int64_t job = start; job < end; ++job) { + int64_t ii = m0 + job / xtiles * RM; + int64_t jj = n0 + job % xtiles * RN; + float32x4_t Cv[RN][RM] = {}; + for (int64_t l = 0; l < k; ++l) + for (int64_t j = 0; j < RN; ++j) + for (int64_t i = 0; i < RM; ++i) + Cv[j][i] = vmlaq_n_f32(Cv[j][i], + vcvtq_f32_s32(vdotq_s32( + vdotq_s32(vdupq_n_s32(0), + load_lo(A + lda * (ii + i) + l), + load_lo(B + ldb * (jj + j) + l)), + load_hi(A + lda * (ii + i) + l), + load_hi(B + ldb * (jj + j) + l))), + unhalf(A[lda * (ii + i) + l].d) * + unhalf(B[ldb * (jj + j) + l].d)); + for (int64_t j = 0; j < RN; ++j) + for (int64_t i = 0; i < RM; ++i) + C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]); + } + } + + inline int8x16_t load_lo(const block_q8_0 *b) { + return vld1q_s8(b->qs); + } + + inline int8x16_t load_hi(const block_q8_0 *b) { + return vld1q_s8(b->qs + 16); + } + + inline int8x16_t load_lo(const block_q4_0 *b) { + return vsubq_s8(vreinterpretq_s8_u8(vandq_u8(vld1q_u8(b->qs), + vdupq_n_u8(0x0f))), + vdupq_n_s8(0x8)); + } + + inline int8x16_t load_hi(const block_q4_0 *b) { + return vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(vld1q_u8(b->qs), 4)), + vdupq_n_s8(0x8)); + } + + const TA *const A; + const block_q8_0 *const B; + float *const C; + const int64_t k; + const int64_t lda; + const int64_t ldb; + const int64_t ldc; + const int ith; + const int nth; +}; +#endif // __ARM_FEATURE_DOTPROD + +#if defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX__) +template +class tinyBLAS_Q0_AVX { + public: + tinyBLAS_Q0_AVX(int64_t k, + const TA *A, int64_t lda, + const TB *B, int64_t ldb, + TC *C, int64_t ldc, + int ith, int nth) + : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { + const int8_t kvalues_iq4nl[16] = { + -127, -104, -83, -65, + -49, -35, -22, -10, + 1, 13, 25, 38, + 53, 69, 89, 113 + }; + + iq4nlt = _mm_loadu_si128((const __m128i *)kvalues_iq4nl); + } + + void matmul(int64_t m, int64_t n) { + mnpack(0, m, 0, n); + } + + private: + void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t mc, nc, mp, np; + switch ((MIN(m - m0, 4) << 4) | MIN(n - n0, 4)) { +#if VECTOR_REGISTERS == 32 + case 0x44: + mc = 4; + nc = 4; +#if defined(__AVX2__) && defined(__F16C__) + gemm4xN<4>(m0, m, n0, n); +#else + gemm<4, 4>(m0, m, n0, n); +#endif + break; + case 0x43: + mc = 4; + nc = 3; +#if defined(__AVX2__) && defined(__F16C__) + gemm4xN<3>(m0, m, n0, n); +#else + gemm<4, 3>(m0, m, n0, n); +#endif + break; + case 0x34: + mc = 3; + nc = 4; +#if defined(__AVX2__) && defined(__F16C__) + gemmMx4<3>(m0, m, n0, n); +#else + gemm<3, 4>(m0, m, n0, n); +#endif + break; + case 0x33: + mc = 3; + nc = 3; + gemm<3, 3>(m0, m, n0, n); + break; + case 0x42: + mc = 4; + nc = 2; +#if defined(__AVX2__) && defined(__F16C__) + gemm4xN<2>(m0, m, n0, n); +#else + gemm<4, 2>(m0, m, n0, n); +#endif + break; + case 0x24: + mc = 2; + nc = 4; +#if defined(__AVX2__) && defined(__F16C__) + gemmMx4<2>(m0, m, n0, n); +#else + gemm<2, 4>(m0, m, n0, n); +#endif + break; +#else + case 0x44: + case 0x43: + case 0x42: + mc = 4; + nc = 2; +#if defined(__AVX2__) && defined(__F16C__) + gemm4xN<2>(m0, m, n0, n); +#else + gemm<4, 2>(m0, m, n0, n); +#endif + break; + case 0x34: + case 0x24: + mc = 2; + nc = 4; +#if defined(__AVX2__) && defined(__F16C__) + gemmMx4<2>(m0, m, n0, n); +#else + gemm<2, 4>(m0, m, n0, n); +#endif + break; + case 0x33: +#endif + case 0x32: + mc = 3; + nc = 2; + gemm<3, 2>(m0, m, n0, n); + break; + case 0x23: + mc = 2; + nc = 3; + gemm<2, 3>(m0, m, n0, n); + break; + case 0x41: + mc = 4; + nc = 1; +#if defined(__AVX2__) && defined(__F16C__) + gemm4xN<1>(m0, m, n0, n); +#else + gemm<4, 1>(m0, m, n0, n); +#endif + break; + case 0x22: + mc = 2; + nc = 2; + gemm<2, 2>(m0, m, n0, n); + break; + case 0x14: + mc = 1; + nc = 4; +#if defined(__AVX2__) && defined(__F16C__) + gemmMx4<1>(m0, m, n0, n); +#else + gemm<1, 4>(m0, m, n0, n); +#endif + break; + case 0x31: + mc = 3; + nc = 1; + gemm<3, 1>(m0, m, n0, n); + break; + case 0x13: + mc = 1; + nc = 3; + gemm<1, 3>(m0, m, n0, n); + break; + case 0x21: + mc = 2; + nc = 1; + gemm<2, 1>(m0, m, n0, n); + break; + case 0x12: + mc = 1; + nc = 2; + gemm<1, 2>(m0, m, n0, n); + break; + case 0x11: + mc = 1; + nc = 1; + gemm<1, 1>(m0, m, n0, n); + break; + default: + return; + } + mp = m0 + (m - m0) / mc * mc; + np = n0 + (n - n0) / nc * nc; + mnpack(mp, m, n0, np); + mnpack(m0, m, np, n); + } + +#if defined(__AVX2__) && defined(__F16C__) +// Templated functions for gemm of dimensions 4xN + template + NOINLINE void gemm4xN(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t ytiles = (m - m0) / 4; + int64_t xtiles = (n - n0) / RN; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) + end = tiles; + for (int64_t job = start; job < end; ++job) { + int64_t ii = m0 + job / xtiles * 4; + int64_t jj = n0 + job % xtiles * RN; + __m256 Cv[RN][4] = {}; + for (int64_t l = 0; l < k; ++l) { + uint64_t a_delta = ((uint64_t)A[lda * (ii + 3) + l].d << 48) | ((uint64_t)A[lda * (ii + 2) + l].d << 32) | ((uint64_t)A[lda * (ii + 1) + l].d << 16) | (A[lda * (ii + 0) + l].d); + // Convert delta values for four blocks to float values + __m128 da = _mm_cvtph_ps(_mm_set_epi64x(0, a_delta)); + __m256i avec0 = load(A + lda * (ii + 0) + l); + __m256i avec1 = load(A + lda * (ii + 1) + l); + __m256i avec2 = load(A + lda * (ii + 2) + l); + __m256i avec3 = load(A + lda * (ii + 3) + l); + for (int64_t j = 0; j < RN; ++j) { + __m128 db = _mm_set1_ps(unhalf(B[ldb * (jj + j) + l].d)); + // Computation of product of delta values for four blocks and replicate it across 256 bit lane + __m256 dvec = _mm256_castps128_ps256(_mm_mul_ps(da, db)); + dvec = _mm256_permute2f128_ps(dvec ,dvec, 0); + // Computation of dot product and multiplication with appropriate delta value products + Cv[j][0] = madd(_mm256_shuffle_ps(dvec, dvec, 0), + updot(_mm256_sign_epi8(avec0, avec0), + _mm256_sign_epi8(load(B + ldb * (jj + j) + l), avec0)), + Cv[j][0]); + Cv[j][1] = madd(_mm256_shuffle_ps(dvec, dvec, 85), + updot(_mm256_sign_epi8(avec1, avec1), + _mm256_sign_epi8(load(B + ldb * (jj + j) + l), avec1)), + Cv[j][1]); + Cv[j][2] = madd(_mm256_shuffle_ps(dvec, dvec, 170), + updot(_mm256_sign_epi8(avec2, avec2), + _mm256_sign_epi8(load(B + ldb * (jj + j) + l), avec2)), + Cv[j][2]); + Cv[j][3] = madd(_mm256_shuffle_ps(dvec, dvec, 255), + updot(_mm256_sign_epi8(avec3, avec3), + _mm256_sign_epi8(load(B + ldb * (jj + j) + l), avec3)), + Cv[j][3]); + } + } + + for (int64_t j = 0; j < RN; ++j) + for (int64_t i = 0; i < 4; ++i) + C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]); + } + } + + // Templated functions for gemm of dimensions Mx4 + template + NOINLINE void gemmMx4(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t ytiles = (m - m0) / RM; + int64_t xtiles = (n - n0) / 4; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) + end = tiles; + for (int64_t job = start; job < end; ++job) { + int64_t ii = m0 + job / xtiles * RM; + int64_t jj = n0 + job % xtiles * 4; + __m256 Cv[4][RM] = {}; + for (int64_t l = 0; l < k; ++l) { + uint64_t b_delta = ((uint64_t)B[ldb * (jj + 3) + l].d << 48) | ((uint64_t)B[ldb * (jj + 2) + l].d << 32) | ((uint64_t)B[ldb * (jj + 1) + l].d << 16) | (B[ldb * (jj + 0) + l].d); + // Convert delta values for four blocks to float values + __m128 db = _mm_cvtph_ps(_mm_set_epi64x(0, b_delta)); + __m256i bvec0 = load(B + ldb * (jj + 0) + l); + __m256i bvec1 = load(B + ldb * (jj + 1) + l); + __m256i bvec2 = load(B + ldb * (jj + 2) + l); + __m256i bvec3 = load(B + ldb * (jj + 3) + l); + for (int64_t i = 0; i < RM; ++i) { + __m128 da = _mm_set1_ps(unhalf((A[lda * (ii + i) + l].d))); + // Computation of product of delta values for four blocks and replicate it across 256 bit lane + __m256 dvec = _mm256_castps128_ps256(_mm_mul_ps(da, db)); + dvec = _mm256_permute2f128_ps(dvec ,dvec, 0); + // Computation of dot product and multiplication with appropriate delta value products + Cv[0][i] = madd(_mm256_shuffle_ps(dvec, dvec, 0), + updot(_mm256_sign_epi8(load(A + lda * (ii + i) + l), + load(A + lda * (ii + i) + l)), + _mm256_sign_epi8(bvec0, load(A + lda * (ii + i) + l))), + Cv[0][i]); + Cv[1][i] = madd(_mm256_shuffle_ps(dvec, dvec, 85), + updot(_mm256_sign_epi8(load(A + lda * (ii + i) + l), + load(A + lda * (ii + i) + l)), + _mm256_sign_epi8(bvec1, load(A + lda * (ii + i) + l))), + Cv[1][i]); + Cv[2][i] = madd(_mm256_shuffle_ps(dvec, dvec, 170), + updot(_mm256_sign_epi8(load(A + lda * (ii + i) + l), + load(A + lda * (ii + i) + l)), + _mm256_sign_epi8(bvec2, load(A + lda * (ii + i) + l))), + Cv[2][i]); + Cv[3][i] = madd(_mm256_shuffle_ps(dvec, dvec, 255), + updot(_mm256_sign_epi8(load(A + lda * (ii + i) + l), + load(A + lda * (ii + i) + l)), + _mm256_sign_epi8(bvec3, load(A + lda * (ii + i) + l))), + Cv[3][i]); + } + } + for (int64_t j = 0; j < 4; ++j) + for (int64_t i = 0; i < RM; ++i) + C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]); + } + } +#endif + + template + NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t ytiles = (m - m0) / RM; + int64_t xtiles = (n - n0) / RN; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) + end = tiles; + for (int64_t job = start; job < end; ++job) { + int64_t ii = m0 + job / xtiles * RM; + int64_t jj = n0 + job % xtiles * RN; + __m256 Cv[RN][RM] = {}; + for (int64_t l = 0; l < k; ++l) + for (int64_t j = 0; j < RN; ++j) + for (int64_t i = 0; i < RM; ++i) { +#if defined(__AVX2__) + __m256 udTmp = updot(_mm256_sign_epi8(load(A + lda * (ii + i) + l), + load(A + lda * (ii + i) + l)), + _mm256_sign_epi8(load(B + ldb * (jj + j) + l), + load(A + lda * (ii + i) + l))); +#else + __m128i ali0 = load0(A + lda * (ii + i) + l); + __m128i ali1 = load1(A + lda * (ii + i) + l); + __m128i blj0 = load0(B + ldb * (jj + j) + l); + __m128i blj1 = load1(B + ldb * (jj + j) + l); + + __m128i sepAA0 = _mm_sign_epi8(ali0, ali0); + __m128i sepAA1 = _mm_sign_epi8(ali1, ali1); + __m128i sepBA0 = _mm_sign_epi8(blj0, ali0); + __m128i sepBA1 = _mm_sign_epi8(blj1, ali1); + + // updot + const __m128i oneFill = _mm_set1_epi16(1); + __m128i mad0 = _mm_maddubs_epi16(sepAA0, sepBA0); + __m128i mad1 = _mm_maddubs_epi16(sepAA1, sepBA1); + __m256 udTmp = _mm256_cvtepi32_ps(MM256_SET_M128I(_mm_madd_epi16(oneFill, mad1), _mm_madd_epi16(oneFill, mad0))); +#endif + Cv[j][i] = madd(_mm256_set1_ps(unhalf(A[lda * (ii + i) + l].d) * + unhalf(B[ldb * (jj + j) + l].d)), + udTmp, + Cv[j][i]); + } + for (int64_t j = 0; j < RN; ++j) + for (int64_t i = 0; i < RM; ++i) + C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]); + } + } + + inline __m256i load(const block_q8_0 *b) { + return _mm256_loadu_si256((const __m256i *)b->qs); + } + + inline __m128i load0(const block_q8_0 *b) { + return _mm_loadu_si128((const __m128i *)b->qs); + } + + inline __m128i load1(const block_q8_0 *b) { + return _mm_loadu_si128(((const __m128i *)b->qs) + 1); + } + + inline __m256i load(const block_q4_0 *b) { + return _mm256_sub_epi8(denibble(b->qs), _mm256_set1_epi8(8)); + } + + inline __m128i load0(const block_q4_0 *b) { + const __m128i x = _mm_loadu_si128((const __m128i *)(b->qs)); + return _mm_sub_epi8(_mm_and_si128(_mm_set1_epi8(15), x), _mm_set1_epi8(8)); + } + + inline __m128i load1(const block_q4_0 *b) { + const __m128i x = _mm_loadu_si128((const __m128i *)(b->qs)); + return _mm_sub_epi8(_mm_and_si128(_mm_set1_epi8(15), _mm_srli_epi16(x, 4)), _mm_set1_epi8(8)); + } + + inline __m256i load(const block_q5_0 *b) { + return _mm256_or_si256(denibble(b->qs), bittobyte(b->qh)); + } + + inline __m128i load0(const block_q5_0* b) { + const __m128i x = _mm_loadu_si128((const __m128i *)(b->qs)); + uint32_t x32; + memcpy(&x32, b->qh, sizeof(uint32_t)); + __m128i qxl = _mm_and_si128(_mm_set1_epi8(15), x); + __m128i bytesl = _mm_cmpeq_epi8(_mm_set1_epi64x(-1), + _mm_or_si128(_mm_set1_epi64x(0x7fbfdfeff7fbfdfe), + _mm_shuffle_epi8(_mm_set1_epi32(x32), + _mm_set_epi64x(0x0101010101010101, 0x0000000000000000)))); + bytesl = _mm_andnot_si128(bytesl, _mm_set1_epi8((char)0xF0)); + return _mm_or_si128(qxl, bytesl); + } + + inline __m128i load1(const block_q5_0* b) { + const __m128i x = _mm_loadu_si128((const __m128i *)(b->qs)); + uint32_t x32; + memcpy(&x32, b->qh, sizeof(uint32_t)); + __m128i qxh = _mm_and_si128(_mm_set1_epi8(15), _mm_srli_epi16(x, 4)); + __m128i bytesh = _mm_cmpeq_epi8(_mm_set1_epi64x(-1), + _mm_or_si128(_mm_set1_epi64x(0x7fbfdfeff7fbfdfe), + _mm_shuffle_epi8(_mm_set1_epi32(x32), + _mm_set_epi64x(0x0303030303030303, 0x0202020202020202)))); + bytesh = _mm_andnot_si128(bytesh, _mm_set1_epi8((char)0xF0)); + return _mm_or_si128(qxh, bytesh); + } + + inline __m256i load(const block_iq4_nl *b) { + return MM256_SET_M128I(load1(b), load0(b)); + } + + inline __m128i load0(const block_iq4_nl *b) { + const __m128i x = _mm_loadu_si128((const __m128i *)(b->qs)); + return _mm_shuffle_epi8(iq4nlt, _mm_and_si128(_mm_set1_epi8(15), x)); + } + + inline __m128i load1(const block_iq4_nl *b) { + const __m128i x = _mm_loadu_si128((const __m128i *)(b->qs)); + return _mm_shuffle_epi8(iq4nlt, _mm_and_si128(_mm_set1_epi8(15), _mm_srli_epi16(x, 4))); + } + + inline __m256 updot(__m256i u, __m256i s) { + __m256i res; +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) + res = _mm256_dpbusd_epi32(_mm256_setzero_si256(), u, s); +#elif defined(__AVXVNNI__) + res = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), u, s); +#else + res = _mm256_madd_epi16(_mm256_set1_epi16(1), _mm256_maddubs_epi16(u, s)); +#endif + return _mm256_cvtepi32_ps(res); + } + + static inline __m256i denibble(const uint8_t *p) { + __m128i x = _mm_loadu_si128((const __m128i *)p); + return _mm256_and_si256(_mm256_set1_epi8(15), + _mm256_insertf128_si256(_mm256_castsi128_si256(x), + _mm_srli_epi16(x, 4), 1)); + } + + static inline __m256i bittobyte(const uint8_t *p) { + uint32_t x32; + memcpy(&x32, p, sizeof(uint32_t)); + __m256i bytes = _mm256_cmpeq_epi8(_mm256_set1_epi64x(-1), + _mm256_or_si256(_mm256_set1_epi64x(0x7fbfdfeff7fbfdfe), + _mm256_shuffle_epi8(_mm256_set1_epi32(x32), + _mm256_set_epi64x(0x0303030303030303, 0x0202020202020202, + 0x0101010101010101, 0x0000000000000000)))); + return _mm256_andnot_si256(bytes, _mm256_set1_epi8((char)0xF0)); + } + + const TA *const A; + const TB *const B; + TC *const C; + const int64_t k; + const int64_t lda; + const int64_t ldb; + const int64_t ldc; + const int ith; + const int nth; + __m128i iq4nlt; +}; +#endif // __AVX__ + +//PPC Implementation +#if defined(__MMA__) + +#define SAVE_ACC(ACC, ii, jj) \ + __builtin_mma_disassemble_acc(vec_C, ACC); \ + for (int I = 0; I < 4; I++) { \ + for (int J = 0; J < 4; J++) { \ + *((float*)(C+ii+((jj+J)*ldc)+I)) = *((float*)&vec_C[I]+J); \ + } \ + } \ + +template +struct mma_instr; + +template<> +struct mma_instr { + static inline void outer_product(acc_t *acc, vec_t a, vec_t b) { + __builtin_mma_xvbf16ger2pp(acc, a, b); + } +}; + +template<> +struct mma_instr { + static inline void outer_product(acc_t *acc, vec_t a, vec_t b) { + __builtin_mma_xvf16ger2pp(acc, a, b); + } +}; + +template +class tinyBLAS_HP16_PPC { + public: + tinyBLAS_HP16_PPC(int64_t k, + const TA *A, int64_t lda, + const TB *B, int64_t ldb, + TC *C, int64_t ldc, + int ith, int nth) + : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { + } + + void matmul(int64_t m, int64_t n) { + mnpack(0, m, 0, n); + } + + private: + void vector_permute_store(vec_t *c, int numVec, unsigned char *vecOffset) { + vec_t t[8], s[8]; + vec_t swiz1 = {0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23}; + vec_t swiz2 = {8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31}; + vec_t swiz3 = {0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23}; + vec_t swiz4 = {8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31}; + + if (numVec == 2) { + t[0] = vec_perm(c[0], c[1], swiz1); + t[1] = vec_perm(c[2], c[3], swiz1); + s[0] = vec_perm(t[0], t[1], swiz3); + s[1] = vec_perm(t[0], t[1], swiz4); + vec_xst(s[0], 0, (vec_t*)vecOffset); + vec_xst(s[1], 0, (vec_t*)(vecOffset + 16)); + } else if (numVec == 4) { + t[0] = vec_perm(c[0], c[1], swiz1); + t[1] = vec_perm(c[0], c[1], swiz2); + t[2] = vec_perm(c[2], c[3], swiz1); + t[3] = vec_perm(c[2], c[3], swiz2); + s[0] = vec_perm(t[0], t[2], swiz3); + s[1] = vec_perm(t[0], t[2], swiz4); + s[2] = vec_perm(t[1], t[3], swiz3); + s[3] = vec_perm(t[1], t[3], swiz4); + for (int i = 0; i < 4; ++i) + vec_xst(s[i], 0, (vec_t*)(vecOffset + i * 16)); + } else if (numVec == 8) { + for (int i = 0; i < 4; i += 2) { + t[i+0] = vec_perm(c[i+0], c[i+1], swiz1); + t[i+1] = vec_perm(c[i+0], c[i+1], swiz2); + } + for (int i = 4; i < 8; i += 2) { + t[i+0] = vec_perm(c[i+0], c[i+1], swiz1); + t[i+1] = vec_perm(c[i+0], c[i+1], swiz2); + } + s[0] = vec_perm(t[0], t[2], swiz3); + s[1] = vec_perm(t[0], t[2], swiz4); + s[2] = vec_perm(t[1], t[3], swiz3); + s[3] = vec_perm(t[1], t[3], swiz4); + s[4] = vec_perm(t[4], t[6], swiz3); + s[5] = vec_perm(t[4], t[6], swiz4); + s[6] = vec_perm(t[5], t[7], swiz3); + s[7] = vec_perm(t[5], t[7], swiz4); + for (int i = 0; i < 8; ++i) + vec_xst(s[i], 0, (vec_t*)(vecOffset + i * 16)); + } + } + + void packNormal(const TA* a, int64_t lda, int rows, int cols, unsigned char* vec) { + int64_t i, j; + TA *aoffset = NULL; + unsigned char *vecOffset = NULL; + TA * aoffsets[8]; + vector unsigned char c_arr[8]; + aoffset = const_cast(a); + vecOffset = vec; + j = (rows >> 3); + if (j > 0) { + do { + if (cols == 4) { + aoffsets[0] = aoffset; + for (int it = 1; it < 4; ++it) + aoffsets[it] = aoffsets[it-1] + lda; + aoffset += 4 * lda; + for (int i = 0; i < 4; ++i) + c_arr[i] = vec_xl(0, (vector unsigned char*)aoffsets[i]); + vector_permute_store(c_arr, 4, vecOffset); + for (int i = 0; i<4; i++) + aoffsets[i] = aoffsets[i]+lda; + vecOffset +=64; + } + i = (cols >> 3); + if (i > 0) { + aoffsets[0] = aoffset; + for (int it = 1; it < 8; ++it) { + aoffsets[it] = aoffsets[it-1] + lda; + } + aoffset += 8 * lda; + do { + for (int it = 0; it < 8; ++it) + c_arr[it] = vec_xl(0, (vector unsigned char*)aoffsets[it]); + vector_permute_store(c_arr, 8, vecOffset); + for (int it = 0; it < 8; ++it) + aoffsets[it] = aoffsets[it] + 8*lda; + vecOffset += 128; + i--; + } while(i > 0); + } + j--; + } while(j > 0); + } + if (rows & 4) { + aoffsets[0] = aoffset; + for (int it = 1; it < 4; ++it) + aoffsets[it] = aoffsets[it-1] + lda; + aoffset += 4 * lda; + if (cols == 4) { + for (int it = 0; it < 4; ++it) + c_arr[it] = vec_xl(0, (vector unsigned char*)aoffsets[it]); + vector_permute_store(c_arr, 2, vecOffset); + for (int it = 0; it< 4; it++) + aoffsets[it] = aoffsets[it] + lda; + vecOffset += 32; + } + i = (cols >> 3); + if (i > 0) { + do { + for (int it = 0; it < 4; ++it) + c_arr[it] = vec_xl(0, (vector unsigned char*)aoffsets[it]); + vector_permute_store(c_arr, 4, vecOffset); + for (int it = 0; it< 4; it++) + aoffsets[it] = aoffsets[it] + 8*lda; + vecOffset += 64; + i--; + } while(i > 0); + } + } + if (rows & 3) { + aoffsets[0] = aoffset; + for (int it = 1; it < 4; ++it) + aoffsets[it] = aoffsets[it-1] + lda; + if (cols == 4) { + switch(rows) { + case 3: c_arr[2] = vec_xl(0, (vector unsigned char*)aoffsets[2]); + case 2: c_arr[1] = vec_xl(0, (vector unsigned char*)aoffsets[1]); + case 1: c_arr[0] = vec_xl(0, (vector unsigned char*)aoffsets[0]); + break; + } + vector_permute_store(c_arr, 2, vecOffset); + for (int it = 0; it< 4; it++) + aoffsets[it] = aoffsets[it] + lda; + vecOffset += 32; + } + i = (cols >> 3); + if (i > 0) { + do { + switch(rows) { + case 3: c_arr[2] = vec_xl(0, (vector unsigned char*)aoffsets[2]); + case 2: c_arr[1] = vec_xl(0, (vector unsigned char*)aoffsets[1]); + case 1: c_arr[0] = vec_xl(0, (vector unsigned char*)aoffsets[0]); + break; + } + vector_permute_store(c_arr, 4, vecOffset); + for (int it = 0; it <4; it++) + aoffsets[it] = aoffsets[it] + 8* lda; + vecOffset += 64; + i--; + } while(i > 0); + } + } + } + + void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t mc, nc, mp, np; + int m_rem = MIN(m - m0, 8); + int n_rem = MIN(n - n0, 8); + + if (m_rem >= 8 && n_rem >= 8) { + mc = 8; + nc = 8; + gemm<8,8>(m0, m, n0, n); + } else if (m_rem >= 4 && n_rem >= 8) { + mc = 4; + nc = 8; + gemm<4,8>(m0, m, n0, n); + } else if (m_rem >=8 && n_rem >=4){ + mc = 8; + nc = 4; + gemm<8,4>(m0, m, n0, n); + } else if ((m_rem < 4) && (n_rem >= 8)) { + nc = 8; + switch(m_rem) { + case 1: + mc = 1; + gemm_Mx8<1>(m0, m, n0, n); + break; + case 2: + mc = 2; + gemm_Mx8<2>(m0, m, n0, n); + break; + case 3: + mc = 3; + gemm_Mx8<3>(m0, m, n0, n); + break; + default: + return; + } + } else if (m_rem >= 4 && n_rem >= 4) { + mc = 4; + nc = 4; + gemm_small<4, 4>(m0, m, n0, n); + } else if ((m_rem > 4) && (n_rem < 4)) { + mc = 4; + switch(n_rem) { + case 1: + nc = 1; + gemm_small<4, 1>(m0, m, n0, n); + break; + case 2: + nc = 2; + gemm_small<4, 2>(m0, m, n0, n); + break; + case 3: + nc = 3; + gemm_small<4, 3>(m0, m, n0, n); + break; + + default: + return; + } + } else { + switch((m_rem << 4) | n_rem) { + case 0x43: + mc = 4; + nc = 3; + gemm_small<4, 3>(m0, m, n0, n); + break; + case 0x42: + mc = 4; + nc = 2; + gemm_small<4, 2>(m0, m, n0, n); + break; + case 0x41: + mc = 4; + nc = 1; + gemm_small<4, 1>(m0, m, n0, n); + break; + case 0x34: + mc = 3; + nc = 4; + gemm_small<3, 4>(m0, m, n0, n); + break; + case 0x33: + mc = 3; + nc = 3; + gemm_small<3, 3>(m0, m, n0, n); + break; + case 0x32: + mc = 3; + nc = 2; + gemm_small<3, 2>(m0, m, n0, n); + break; + case 0x31: + mc = 3; + nc = 1; + gemm_small<3, 1>(m0, m, n0, n); + break; + case 0x24: + mc = 2; + nc = 4; + gemm_small<2,4>(m0, m, n0, n); + break; + case 0x23: + mc = 2; + nc = 3; + gemm_small<2, 3>(m0, m, n0, n); + break; + case 0x22: + mc = 2; + nc = 2; + gemm_small<2, 2>(m0, m, n0, n); + break; + case 0x21: + mc = 2; + nc = 1; + gemm_small<2, 1>(m0, m, n0, n); + break; + case 0x14: + mc = 1; + nc = 4; + gemm_small<1, 4>(m0, m, n0, n); + break; + case 0x13: + mc = 1; + nc = 3; + gemm_small<1, 3>(m0, m, n0, n); + break; + case 0x12: + mc = 1; + nc = 2; + gemm_small<1, 2>(m0, m, n0, n); + break; + case 0x11: + mc = 1; + nc = 1; + gemm_small<1, 1>(m0, m, n0, n); + break; + default: + return; + } + } + mp = m0 + (m - m0) / mc * mc; + np = n0 + (n - n0) / nc * nc; + mnpack(mp, m, n0, np); + mnpack(m0, m, np, n); + } + + void KERNEL_4x8(int64_t ii, int64_t jj) { + vec_t vec_A[4], vec_B[8] , vec_C[4]; + acc_t acc_0, acc_1; + __builtin_mma_xxsetaccz(&acc_0); + __builtin_mma_xxsetaccz(&acc_1); + for (int l = 0; l < k; l+=8) { + packNormal((A+(ii*lda)+l), lda, 4, 8, (uint8_t*)vec_A); + packNormal((B+(jj*ldb)+l), ldb, 8, 8, (uint8_t*)vec_B); + for (int x = 0; x < 4; x++) { + mma_instr::outer_product(&acc_0, vec_A[x], vec_B[x]); + mma_instr::outer_product(&acc_1, vec_A[x], vec_B[x+4]); + } + } + SAVE_ACC(&acc_0, ii, jj); + SAVE_ACC(&acc_1, ii, jj+4); + } + + void KERNEL_8x4(int64_t ii, int64_t jj) { + vec_t vec_A[8], vec_B[4] , vec_C[4]; + acc_t acc_0, acc_1; + __builtin_mma_xxsetaccz(&acc_0); + __builtin_mma_xxsetaccz(&acc_1); + for (int l = 0; l < k; l+=8) { + packNormal((A+(ii*lda)+l), lda, 8, 8, (uint8_t*)vec_A); + packNormal((B+(jj*ldb)+l), ldb, 8, 4, (uint8_t*)vec_B); + for (int x = 0; x < 4; x++) { + mma_instr::outer_product(&acc_0, vec_A[x], vec_B[x]); + mma_instr::outer_product(&acc_1, vec_A[x+4], vec_B[x]); + } + } + SAVE_ACC(&acc_0, ii, jj); + SAVE_ACC(&acc_1, ii+4, jj); + } + + + void KERNEL_8x8(int64_t ii, int64_t jj) { + vec_t vec_A[8], vec_B[8], vec_C[4]; + acc_t acc_0, acc_1, acc_2, acc_3; + __builtin_mma_xxsetaccz(&acc_0); + __builtin_mma_xxsetaccz(&acc_1); + __builtin_mma_xxsetaccz(&acc_2); + __builtin_mma_xxsetaccz(&acc_3); + for (int l = 0; l < k; l+=8) { + packNormal(A+(ii*lda)+l, lda, 8, 8, (uint8_t*)vec_A); + packNormal(B+(jj*ldb)+l, ldb, 8, 8, (uint8_t*)vec_B); + for (int x = 0; x < 4; x++) { + mma_instr::outer_product(&acc_0, vec_A[x], vec_B[x]); + mma_instr::outer_product(&acc_1, vec_A[x], vec_B[x+4]); + mma_instr::outer_product(&acc_2, vec_A[x+4], vec_B[x]); + mma_instr::outer_product(&acc_3, vec_A[x+4], vec_B[x+4]); + } + } + + SAVE_ACC(&acc_0, ii, jj); + SAVE_ACC(&acc_1, ii, jj+4); + SAVE_ACC(&acc_2, ii+4, jj); + SAVE_ACC(&acc_3, ii+4, jj+4); + } + + template + void gemm_small(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t ytiles = (m - m0) / RM; + int64_t xtiles = (n - n0) / RN; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) + end = tiles; + for (int64_t job = start; job < end; ++job) { + int64_t ii = m0 + job / xtiles * RM; + int64_t jj = n0 + job % xtiles * RN; + vec_t vec_C[4]; + acc_t acc_0; + __builtin_mma_xxsetaccz(&acc_0); + vec_t vec_A[2], vec_B[2]; + for (int l=0; l::outer_product(&acc_0, vec_A[x], vec_B[x]); + } + } + __builtin_mma_disassemble_acc(vec_C, &acc_0); + for (int I = 0; I < RM; I++) { + for (int J = 0; J < RN; J++) { + *((TC*)(C+ii+((jj+J)*ldc)+I)) = *((TC*)&vec_C[I]+J); + } + } + } + } + + template + void gemm_Mx8(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int RN = 8; + int64_t ytiles = (m - m0) / RM; + int64_t xtiles = (n - n0) / RN; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) + end = tiles; + for (int64_t job = start; job < end; ++job) { + int64_t ii = m0 + job / xtiles * RM; + int64_t jj = n0 + job % xtiles * RN; + vec_t vec_C[4]; + acc_t acc_0, acc_1; + __builtin_mma_xxsetaccz(&acc_0); + __builtin_mma_xxsetaccz(&acc_1); + vec_t vec_A[4], vec_B[8]; + for (int l=0; l::outer_product(&acc_0, vec_A[x], vec_B[x]); + mma_instr::outer_product(&acc_1, vec_A[x], vec_B[x+4]); + } + } + __builtin_mma_disassemble_acc(vec_C, &acc_0); + for (int I = 0; I < RM; I++) { + for (int J = 0; J < 4; J++) { + *((TC*)(C+ii+((jj+J)*ldc)+I)) = *((TC*)&vec_C[I]+J); + } + } + __builtin_mma_disassemble_acc(vec_C, &acc_1); + for (int I = 0; I < RM; I++) { + for (int J = 0; J < 4; J++) { + *((TC*)(C+ii+((jj+4+J)*ldc)+I)) = *((TC*)&vec_C[I]+J); + } + } + } + } + + template + inline void kernel(int64_t ii, int64_t jj) { + if constexpr(RM == 4 && RN == 8) { + KERNEL_4x8(ii,jj); + } else if constexpr(RM == 8 && RN == 8) { + KERNEL_8x8(ii,jj); + } else if constexpr(RM == 8 && RN == 4) { + KERNEL_8x4(ii,jj); + } else { + assert(false && "RN/RM values not supported"); + } + } + + template + NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t ytiles = (m - m0) / RM; + int64_t xtiles = (n - n0) / RN; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) + end = tiles; + for (int64_t job = start; job < end; ++job) { + int64_t ii = m0 + job / xtiles * RM; + int64_t jj = n0 + job % xtiles * RN; + kernel(ii, jj); + } + } + + const TA *const A; + const TB *const B; + TC *C; + const int64_t k; + const int64_t lda; + const int64_t ldb; + const int64_t ldc; + const int ith; + const int nth; +}; + +template +class tinyBLAS_Q0_PPC { + public: + tinyBLAS_Q0_PPC(int64_t k, + const TA * A, int64_t lda, + const block_q8_0 * B, int64_t ldb, + float * C, int64_t ldc, + int ith, int nth) + : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { + } + + void matmul(int64_t m, int64_t n) { + #if defined(_AIX) || defined(__BIG_ENDIAN__) + mnpack(0, m, 0, n); + #else + const int64_t mc = 64; + const int64_t kc = 64; + int64_t nc = 64; + int64_t n_aligned = 0; + if (n % 64 == 0) { + n_aligned = n; + } else if (n == 4) { + n_aligned = 4; + } else if (n < 64) { + n_aligned = (n / 8) * 8; + } else { + n_aligned = (n / 64) * 64; + } + if (n_aligned > 0) { + if (n_aligned % 64 == 0) nc = 64; + else if (n_aligned == n) nc = n; + else if (n_aligned % 32 == 0) nc = 32; + else if (n_aligned % 24 == 0) nc = 24; + else if (n_aligned % 16 == 0) nc = 16; + else nc = 8; + } + bool can_use_tiled = n_aligned > 0 && (m % mc == 0) && (k % kc == 0); + if (can_use_tiled) { + matmul_tiled(m, n_aligned, mc, nc, kc); + if (n > n_aligned) { + mnpack(0, m, n_aligned, n); + } + } else { + mnpack(0, m, 0, n); + } + #endif + } + + private: + inline void save_res(int ii, int jj, int idx, vector float * fin_res, int RM = 4, int RN = 4) { + for (int I = 0; I < RM; I++) { + for (int J = 0; J < RN; J++) { + *((float *)(C + ii + ((jj + J) * ldc) + I)) = *((float *)&fin_res[idx + I] + J); + } + } + } + + inline void save_acc(acc_t * ACC, int64_t ii, int64_t jj) { + vec_t vec_C[4]; + __builtin_mma_disassemble_acc(vec_C, ACC); + for (int I = 0; I < 4; I++) { + for (int J = 0; J < 4; J++) { + *((float *)(C + ii + ((jj + J) * ldc) + I)) = *((float *)&vec_C[I] + J); + } + } + } + + inline void add_save_acc(acc_t * ACC, int64_t ii, int64_t jj) { + vec_t vec_C[4]; + __builtin_mma_disassemble_acc(vec_C, ACC); + for (int I = 0; I < 4; I++) { + for (int J = 0; J < 4; J++) { + float * c_ptr = (float *)(C + ii+ ((jj + J) * ldc) + I); + *c_ptr += *((float *)&vec_C[I] + J); + } + } + } + + template + inline void compute(acc_t * ACC, int c_idx, int s_idx, ArrayType & comparray, vector float * vs, vector float * fin_res) { + vector signed int vec_C[4]; + vector float CA[4] = {0}; + vector float res[4] = {0}; + __builtin_mma_disassemble_acc(vec_C, ACC); + for (int i = 0; i < 4; i++) { + CA[i] = vec_splats((float)(((double)comparray[c_idx + i]) * -128.0)); + res[i] = vec_add(vec_ctf(vec_C[i], 0), CA[i]); + fin_res[s_idx + i] = vec_madd(res[i], vs[s_idx + i], fin_res[s_idx + i]); + } + } + + inline void process_q4_elements(vector signed char (&c)[2], int * ca) { + const vector signed char lowMask = vec_splats((signed char)0xF); + const vector unsigned char v4 = vec_splats((unsigned char)0x4); + const vector signed char v8 = vec_splats((signed char)0x8); + vector signed int vsum = {0}; + vector signed int vsum2 = {0}; + c[0] = vec_and(c[1], lowMask); + c[1] = vec_sr(c[1], v4); + c[0] = vec_sub(c[0], v8); + c[1] = vec_sub(c[1], v8); + vsum = vec_sum4s(c[0], vsum); + vsum2 = vec_sum4s(c[1], vsum2); + vsum = vec_add(vsum, vsum2); + *(ca) = vsum[0] + vsum[1] + vsum[2] + vsum[3]; + } + + template + inline void vector_permute_store(V2 & s1, V2 & s2, V2 & s3, V2 & s4, V1 * vecOffset, bool flip) { + vector unsigned char swiz1 = {0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23}; + vector unsigned char swiz2 = {8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31}; + vector unsigned char swiz3 = {0, 1, 2, 3, 8, 9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27}; + vector unsigned char swiz4 = {4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31}; + V2 t1, t2, t3, t4, t5, t6, t7, t8; + vector unsigned char xor_vector; + uint8_t flip_vec = 0x80; + xor_vector = vec_splats(flip_vec); + t1 = vec_perm(s1, s2, swiz1); + t2 = vec_perm(s1, s2, swiz2); + t3 = vec_perm(s3, s4, swiz1); + t4 = vec_perm(s3, s4, swiz2); + t5 = vec_perm(t1, t3, swiz3); + t6 = vec_perm(t1, t3, swiz4); + t7 = vec_perm(t2, t4, swiz3); + t8 = vec_perm(t2, t4, swiz4); + if (flip == true) { + t5 = vec_xor(t5, xor_vector); + t6 = vec_xor(t6, xor_vector); + t7 = vec_xor(t7, xor_vector); + t8 = vec_xor(t8, xor_vector); + } + vec_xst(t5, 0, vecOffset); + vec_xst(t6, 0, vecOffset + 16); + vec_xst(t7, 0, vecOffset + 32); + vec_xst(t8, 0, vecOffset + 48); + } + + inline void unpack_q4_to_q8(vector signed char packed, vector signed char & lo, vector signed char & hi) { + const vector signed char lowMask = vec_splats((signed char)0x0F); + const vector signed char v8 = vec_splats((signed char)0x08); + const vector unsigned char v4 = vec_splats((unsigned char)4); + lo = vec_and(packed, lowMask); + hi = vec_sr(packed, v4); + lo = vec_sub(lo, v8); + hi = vec_sub(hi, v8); + } + + inline void vector_permute_store_fp16(vec_t * c, unsigned char * vecOffset) { + vec_t t[8], s[8]; + vec_t swiz1 = {0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23}; + vec_t swiz2 = {8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31}; + vec_t swiz3 = {0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23}; + vec_t swiz4 = {8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31}; + for (int i = 0; i < 4; i += 2) { + t[i + 0] = vec_perm(c[i + 0], c[i + 1], swiz1); + t[i + 1] = vec_perm(c[i + 0], c[i + 1], swiz2); + } + for (int i = 4; i < 8; i += 2) { + t[i + 0] = vec_perm(c[i + 0], c[i + 1], swiz1); + t[i + 1] = vec_perm(c[i + 0], c[i + 1], swiz2); + } + s[0] = vec_perm(t[0], t[2], swiz3); + s[1] = vec_perm(t[0], t[2], swiz4); + s[2] = vec_perm(t[1], t[3], swiz3); + s[3] = vec_perm(t[1], t[3], swiz4); + s[4] = vec_perm(t[4], t[6], swiz3); + s[5] = vec_perm(t[4], t[6], swiz4); + s[6] = vec_perm(t[5], t[7], swiz3); + s[7] = vec_perm(t[5], t[7], swiz4); + for (int i = 0; i < 8; ++i) { + vec_xst(s[i], 0, (vec_t *)(vecOffset + i * 16)); + } + } + + static inline void convert_and_scale_q8(vector signed char raw, vector float v_scale, vector unsigned short & out_hi, vector unsigned short & out_lo) { + vector signed short i16_hi = vec_unpackh(raw); + vector signed short i16_lo = vec_unpackl(raw); + + vector float f_hi_h = vec_ctf(vec_unpackh(i16_hi), 0); + vector float f_hi_l = vec_ctf(vec_unpackl(i16_hi), 0); + vector float f_lo_h = vec_ctf(vec_unpackh(i16_lo), 0); + vector float f_lo_l = vec_ctf(vec_unpackl(i16_lo), 0); + out_hi = vec_pack_to_short_fp32(vec_mul(f_hi_h, v_scale), vec_mul(f_hi_l, v_scale)); + out_lo = vec_pack_to_short_fp32(vec_mul(f_lo_h, v_scale), vec_mul(f_lo_l, v_scale)); + } + + void packNormal_q4_fp16(const block_q4_0 * a, int64_t lda, int rows, int blocks, unsigned char * vec) { + unsigned char * vecOffset = vec; + for (int i = 0; i < rows; i += 8) { + const block_q4_0 * rows_base[8]; + for (int r = 0; r < 8; r++) { + rows_base[r] = a + (i + r) * lda; + } + for (int blk = 0; blk < blocks; blk++) { + vector unsigned short hp_res[8][4]; + for (int r = 0; r < 8; r++) { + const block_q4_0 * current_blk = rows_base[r] + blk; + vector float v_scale = vec_extract_fp32_from_shorth(vec_splats(current_blk->d)); + vector signed char v_qs = vec_xl(0, (const vector signed char *)current_blk->qs); + vector signed char c1, c2; + unpack_q4_to_q8(v_qs, c1, c2); + convert_and_scale_q8(c1, v_scale, hp_res[r][0], hp_res[r][1]); + convert_and_scale_q8(c2, v_scale, hp_res[r][2], hp_res[r][3]); + } + for (int c = 0; c < 4; c++) { + vector unsigned char c_arr[8]; + for (int r = 0; r < 8; r++) { + c_arr[r] = (vector unsigned char)hp_res[r][c]; + } + vector_permute_store_fp16((vec_t *)c_arr, vecOffset); + vecOffset += 128; + } + } + } + } + + template + static inline void pack_q8_block(const block_q8_0 * a, int64_t lda, int rows, int blocks, unsigned char * vec) { + unsigned char * vecOffset = vec; + const vec_t swiz1 = {0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23}; + const vec_t swiz2 = {8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31}; + const vec_t swiz3 = {0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23}; + const vec_t swiz4 = {8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31}; + + for (int i = 0; i < rows; i += chunk_size) { + const block_q8_0 * rows_base[chunk_size]; + for (int r = 0; r < chunk_size; r++) { + rows_base[r] = a + (i + r) * lda; + } + for (int blk = 0; blk < blocks; blk++) { + vector unsigned short hp_res[chunk_size][4]; + for (int r = 0; r < chunk_size; r++) { + const block_q8_0 * b = rows_base[r] + blk; + vector float v_scale = vec_extract_fp32_from_shorth(vec_splats(b->d)); + vector signed char c[2]; + __vector_pair pair = __builtin_vsx_lxvp(0, (__vector_pair *)b->qs); + __builtin_vsx_disassemble_pair(c, & pair); + convert_and_scale_q8(c[0], v_scale, hp_res[r][0], hp_res[r][1]); + convert_and_scale_q8(c[1], v_scale, hp_res[r][2], hp_res[r][3]); + } + for (int col = 0; col < 4; col++) { + if constexpr (chunk_size == 8) { + vec_t t[8]; + t[0] = vec_perm((vec_t)hp_res[0][col], (vec_t)hp_res[1][col], swiz1); + t[1] = vec_perm((vec_t)hp_res[0][col], (vec_t)hp_res[1][col], swiz2); + t[2] = vec_perm((vec_t)hp_res[2][col], (vec_t)hp_res[3][col], swiz1); + t[3] = vec_perm((vec_t)hp_res[2][col], (vec_t)hp_res[3][col], swiz2); + t[4] = vec_perm((vec_t)hp_res[4][col], (vec_t)hp_res[5][col], swiz1); + t[5] = vec_perm((vec_t)hp_res[4][col], (vec_t)hp_res[5][col], swiz2); + t[6] = vec_perm((vec_t)hp_res[6][col], (vec_t)hp_res[7][col], swiz1); + t[7] = vec_perm((vec_t)hp_res[6][col], (vec_t)hp_res[7][col], swiz2); + + vec_xst(vec_perm(t[0], t[2], swiz3), 0, (vec_t *)(vecOffset + 0)); + vec_xst(vec_perm(t[0], t[2], swiz4), 0, (vec_t *)(vecOffset + 16)); + vec_xst(vec_perm(t[1], t[3], swiz3), 0, (vec_t *)(vecOffset + 32)); + vec_xst(vec_perm(t[1], t[3], swiz4), 0, (vec_t *)(vecOffset + 48)); + vec_xst(vec_perm(t[4], t[6], swiz3), 0, (vec_t *)(vecOffset + 64)); + vec_xst(vec_perm(t[4], t[6], swiz4), 0, (vec_t *)(vecOffset + 80)); + vec_xst(vec_perm(t[5], t[7], swiz3), 0, (vec_t *)(vecOffset + 96)); + vec_xst(vec_perm(t[5], t[7], swiz4), 0, (vec_t *)(vecOffset + 112)); + vecOffset += 128; + } else { + vec_t t0 = vec_perm((vec_t)hp_res[0][col], (vec_t)hp_res[1][col], swiz1); + vec_t t1 = vec_perm((vec_t)hp_res[0][col], (vec_t)hp_res[1][col], swiz2); + vec_t t2 = vec_perm((vec_t)hp_res[2][col], (vec_t)hp_res[3][col], swiz1); + vec_t t3 = vec_perm((vec_t)hp_res[2][col], (vec_t)hp_res[3][col], swiz2); + + vec_xst(vec_perm(t0, t2, swiz3), 0, (vec_t *)(vecOffset + 0)); + vec_xst(vec_perm(t0, t2, swiz4), 0, (vec_t *)(vecOffset + 16)); + vec_xst(vec_perm(t1, t3, swiz3), 0, (vec_t *)(vecOffset + 32)); + vec_xst(vec_perm(t1, t3, swiz4), 0, (vec_t *)(vecOffset + 48)); + vecOffset += 64; + } + } + } + } + } + + void packNormal_q8_fp16(const block_q8_0 * a, int64_t lda, int rows, int blocks, unsigned char * vec) { + if (rows == 4) { + pack_q8_block<4>(a, lda, rows, blocks, vec); + } else { + pack_q8_block<8>(a, lda, rows, blocks, vec); + } + } + + template + void packNormalInt4(const TA * a, int64_t lda, int rows, int cols, int8_t * vec, std::array & comparray) { + int64_t i, j; + TA * aoffset = NULL; + int8_t * vecOffset = NULL; + TA * aoffset1 = NULL, * aoffset2 = NULL, * aoffset3 = NULL, * aoffset4 = NULL; + TA * aoffset5 = NULL, * aoffset6 = NULL, * aoffset7 = NULL, * aoffset8 = NULL; + vector signed char c1[2] = {0}, c2[2] = {0}, c3[2] = {0}, c4[2] = {0}; + vector signed char c5[2] = {0}, c6[2] = {0}, c7[2] = {0}, c8[2] = {0}; + aoffset = const_cast(a); + vecOffset = vec; + j = (rows >> 3); + if (j > 0) { + do { + aoffset1 = aoffset; + aoffset2 = aoffset1 + lda; + aoffset3 = aoffset2 + lda; + aoffset4 = aoffset3 + lda; + aoffset5 = aoffset4 + lda; + aoffset6 = aoffset5 + lda; + aoffset7 = aoffset6 + lda; + aoffset8 = aoffset7 + lda; + aoffset += 8 * lda; + i = (cols >> 2); + if (i > 0) { + do { + c1[1] = vec_xl(0, (const vector signed char *)aoffset1->qs); + c2[1] = vec_xl(0, (const vector signed char *)aoffset2->qs); + c3[1] = vec_xl(0, (const vector signed char *)aoffset3->qs); + c4[1] = vec_xl(0, (const vector signed char *)aoffset4->qs); + c5[1] = vec_xl(0, (const vector signed char *)aoffset5->qs); + c6[1] = vec_xl(0, (const vector signed char *)aoffset6->qs); + c7[1] = vec_xl(0, (const vector signed char *)aoffset7->qs); + c8[1] = vec_xl(0, (const vector signed char *)aoffset8->qs); + + process_q4_elements(c1, & comparray[0]); + process_q4_elements(c2, & comparray[1]); + process_q4_elements(c3, & comparray[2]); + process_q4_elements(c4, & comparray[3]); + process_q4_elements(c5, & comparray[4]); + process_q4_elements(c6, & comparray[5]); + process_q4_elements(c7, & comparray[6]); + process_q4_elements(c8, & comparray[7]); + vector_permute_store(c1[0], c2[0], c3[0], c4[0], vecOffset, false); + vector_permute_store(c1[1], c2[1], c3[1], c4[1], vecOffset + 64, false); + vector_permute_store(c5[0], c6[0], c7[0], c8[0], vecOffset + 128, false); + vector_permute_store(c5[1], c6[1], c7[1], c8[1], vecOffset + 192, false); + aoffset1 += lda; + aoffset2 += lda; + aoffset3 += lda; + aoffset4 += lda; + aoffset5 += lda; + aoffset6 += lda; + aoffset7 += lda; + aoffset8 += lda; + vecOffset += 256; + i--; + } while (i > 0); + } + j--; + } while (j > 0); + } + + if (rows & 4) { + aoffset1 = aoffset; + aoffset2 = aoffset1 + lda; + aoffset3 = aoffset2 + lda; + aoffset4 = aoffset3 + lda; + aoffset += 4 * lda; + i = (cols >> 2); + if (i > 0) { + do { + c1[1] = vec_xl(0, (const vector signed char *)aoffset1->qs); + c2[1] = vec_xl(0, (const vector signed char *)aoffset2->qs); + c3[1] = vec_xl(0, (const vector signed char *)aoffset3->qs); + c4[1] = vec_xl(0, (const vector signed char *)aoffset4->qs); + + process_q4_elements(c1, & comparray[0]); + process_q4_elements(c2, & comparray[1]); + process_q4_elements(c3, & comparray[2]); + process_q4_elements(c4, & comparray[3]); + vector_permute_store(c1[0], c2[0], c3[0], c4[0], vecOffset, false); + vector_permute_store(c1[1], c2[1], c3[1], c4[1], vecOffset + 64, false); + aoffset1 += lda; + aoffset2 += lda; + aoffset3 += lda; + aoffset4 += lda; + vecOffset += 128; + i--; + } while (i > 0); + } + } + + if (rows & 3) { + aoffset1 = aoffset; + aoffset2 = aoffset1 + lda; + aoffset3 = aoffset2 + lda; + i = (cols >> 2); + if (i > 0) { + do { + switch(rows) { + case 3: c3[1] = vec_xl(0, (const vector signed char *)aoffset3->qs); + case 2: c2[1] = vec_xl(0, (const vector signed char *)aoffset2->qs); + case 1: c1[1] = vec_xl(0, (const vector signed char *)aoffset1->qs); + break; + } + process_q4_elements(c1, & comparray[0]); + process_q4_elements(c2, & comparray[1]); + process_q4_elements(c3, & comparray[2]); + process_q4_elements(c4, & comparray[3]); + vector_permute_store(c1[0], c2[0], c3[0], c4[0], vecOffset, false); + vector_permute_store(c1[1], c2[1], c3[1], c4[1], vecOffset + 64, false); + aoffset1 += lda; + aoffset2 += lda; + aoffset3 += lda; + vecOffset += 128; + i--; + } while(i > 0); + } + } + } + + template + void packNormal(const block_q8_0 * a, int64_t lda, int rows, int cols, VA * vec, bool flip) { + int64_t i, j; + block_q8_0 * aoffset = NULL; + VA * vecOffset = NULL; + block_q8_0 * aoffsets[8]; + __vector_pair arr[8]; + VB c[8][2] = {0}; + VB c1[8] = {0}; VB c2[8] = {0}; + aoffset = const_cast(a); + vecOffset = vec; + j = (rows >> 3); + if (j > 0) { + do { + aoffsets[0] = aoffset; + for (int it = 1; it < 8; it++) + aoffsets[it] = aoffsets[it - 1] + lda; + aoffset += 8 * lda; + + i = (cols >> 3); + if (i > 0) { + do { + for (int it = 0; it < 8; it++) { + arr[it] = __builtin_vsx_lxvp(0, (__vector_pair *)aoffsets[it]->qs); + __builtin_vsx_disassemble_pair(c[it], & arr[it]); + c1[it] = c[it][0]; + c2[it] = c[it][1]; + } + vector_permute_store(c1[0], c1[1], c1[2], c1[3], vecOffset, flip); + vector_permute_store(c2[0], c2[1], c2[2], c2[3], vecOffset + 64, flip); + vector_permute_store(c1[4], c1[5], c1[6], c1[7], vecOffset + 128, flip); + vector_permute_store(c2[4], c2[5], c2[6], c2[7], vecOffset + 192, flip); + for (int it = 0; it < 8; it++) + aoffsets[it] += lda; + vecOffset += 256; + i--; + } while(i > 0); + } + j--; + } while(j > 0); + } + if (rows & 4) { + aoffsets[0] = aoffset; + for (int it = 1; it < 4; it++ ) + aoffsets[it] = aoffsets[it-1] + lda; + aoffset += 4 * lda; + i = (cols >> 3); + if (i > 0) { + do { + for (int it = 0; it < 4; it++) { + arr[it] = __builtin_vsx_lxvp(0, (__vector_pair *)aoffsets[it]->qs); + __builtin_vsx_disassemble_pair(c[it], & arr[it]); + c1[it] = c[it][0]; + c2[it] = c[it][1]; + } + vector_permute_store(c1[0], c1[1], c1[2], c1[3], vecOffset, flip); + vector_permute_store(c2[0], c2[1], c2[2], c2[3], vecOffset + 64, flip); + for (int it = 0; it < 4; it++) { + aoffsets[it] += lda; + } + vecOffset += 128; + i--; + } while(i > 0); + } + } + + if (rows & 3) { + aoffsets[0] = aoffset; + for (int it = 1; it < 3; it++ ) + aoffsets[it] = aoffsets[it - 1] + lda; + i = (cols >> 3); + if (i > 0) { + do { + switch(rows) { + case 3: arr[2] = __builtin_vsx_lxvp(0, (__vector_pair *)aoffsets[2]->qs); + __builtin_vsx_disassemble_pair(c[2], & arr[2]); + c1[2] = c[2][0]; c2[2] = c[2][1]; + case 2: arr[1] = __builtin_vsx_lxvp(0, (__vector_pair *)aoffsets[1]->qs); + __builtin_vsx_disassemble_pair(c[1], & arr[1]); + c1[1] = c[1][0]; c2[1] = c[1][1]; + case 1: arr[0] = __builtin_vsx_lxvp(0, (__vector_pair *)aoffsets[0]->qs); + __builtin_vsx_disassemble_pair(c[0], & arr[0]); + c1[0] = c[0][0]; c2[0] = c[0][1]; + break; + } + vector_permute_store(c1[0], c1[1], c1[2], c1[3], vecOffset, flip); + vector_permute_store(c2[0], c2[1], c2[2], c2[3], vecOffset + 64, flip); + for (int it = 0; it < 3; it++) + aoffsets[it] += lda; + vecOffset += 128; + i--; + } while(i > 0); + } + } + } + + void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int m_rem = MIN(m - m0, 16); + int n_rem = MIN(n - n0, 16); + + int mc = 0, nc = 0; + + if (m_rem >= 8 && n_rem >= 8) { + mc = 8; + nc = 8; + gemm<8, 8>(m0, m, n0, n); + } else if (m_rem >= 4 && n_rem >= 8) { + mc = 4; + nc = 8; + gemm<4, 8>(m0, m, n0, n); + } else if (m_rem >= 8 && n_rem >= 4) { + mc = 8; + nc = 4; + gemm<8, 4>(m0, m, n0, n); + } else if (m_rem >= 4 && n_rem >= 4) { + mc = 4; + nc = 4; + gemm_small(m0, m, n0, n, mc, nc); + } else { + mc = (m_rem >= 4) ? 4 : m_rem; + nc = (n_rem >= 4) ? 4 : n_rem; + if (mc == 0 || nc == 0) + return; + gemm_small(m0, m, n0, n, mc, nc); + } + + int64_t mp = m0 + ((m - m0) / mc) * mc; + int64_t np = n0 + ((n - n0) / nc) * nc; + mnpack(mp, m, n0, np); + mnpack(m0, m, np, n); + } + + + void KERNEL_4x8(int64_t ii, int64_t jj) { + vec_t vec_A[8], vec_B[16] = {0}; + acc_t acc_0, acc_1; + std::array comparray {}; + vector float fin_res[8] = {0}; + vector float vs[8] = {0}; + bool isAblock_q4 = std::is_same_v; + for (int l = 0; l < k; l++) { + __builtin_mma_xxsetaccz(& acc_0); + __builtin_mma_xxsetaccz(& acc_1); + if (std::is_same_v) { + packNormalInt4<4>((A + (ii * lda) + l), lda, 4, 4, (int8_t *)vec_A, comparray); + } else { + packNormal((const block_q8_0 *)(A + (ii * lda) + l), lda, 4, 8, (int8_t *)vec_A, false); + } + packNormal((B + (jj * ldb) + l), ldb, 8, 8, (uint8_t *)vec_B, true); + for(int x = 0; x < 8; x++) { + __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x], vec_B[x]); + __builtin_mma_xvi8ger4pp(& acc_1, vec_A[x], vec_B[x+8]); + } + for (int I = 0; I<4; I++) { + for (int J = 0; J<4; J++) { + *((float *)& vs[I] + J) = (unhalf((A + ((ii + I) * lda) + l)->d) * unhalf((B + ((jj + J) * ldb) + l)->d)); + *((float *)& vs[I + 4] + J) = (unhalf((A +((ii + I) * lda) + l)->d) * unhalf((B + ((jj + J + 4) * ldb) + l)->d)); + } + } + if (!isAblock_q4) { + auto aoffset = A + (ii * lda) + l; + for (int i = 0; i < 4; i++) { + comparray[i] = 0; + int ca = 0; + auto *at = aoffset->qs; + for (int j = 0; j < 32; j++) + ca += (int)*at++; + comparray[i] = ca; + aoffset += lda; + } + } + compute(& acc_0, 0, 0, comparray, vs, fin_res); + compute(& acc_1, 0, 4, comparray, vs, fin_res); + } + save_res(ii, jj, 0, fin_res); + save_res(ii, jj + 4, 4, fin_res); + } + + void KERNEL_8x4(int64_t ii, int64_t jj) { + vec_t vec_A[16], vec_B[8] = {0}; + acc_t acc_0, acc_1; + std::array comparray {}; + vector float fin_res[8] = {0}; + vector float vs[8] = {0}; + bool isAblock_q4 = std::is_same_v; + for (int l = 0; l < k; l++) { + __builtin_mma_xxsetaccz(& acc_0); + __builtin_mma_xxsetaccz(& acc_1); + if (std::is_same_v) { + packNormalInt4<8>((A + (ii * lda) + l), lda, 8, 4, (int8_t *)vec_A, comparray); + } else { + packNormal((const block_q8_0 *)(A + (ii * lda) + l), lda, 8, 8, (int8_t *)vec_A, false); + } + packNormal((B + (jj * ldb) + l), ldb, 4, 8, (uint8_t *)vec_B, true); + for(int x = 0; x < 8; x++) { + __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x], vec_B[x]); + __builtin_mma_xvi8ger4pp(& acc_1, vec_A[x + 8], vec_B[x]); + } + for (int I = 0; I < 8; I++) { + for (int J = 0; J < 4; J++) { + *((float *)&vs[I] + J) = (unhalf((A + ((ii + I) * lda) + l)->d) * unhalf((B + ((jj + J) * ldb) + l)->d)); + } + } + if (!isAblock_q4) { + auto aoffset = A + (ii * lda) + l; + for (int i = 0; i < 8; i++) { + comparray[i] = 0; + int ca = 0; + auto *at = aoffset->qs; + for (int j = 0; j < 32; j++) + ca += (int)*at++; + comparray[i] = ca; + aoffset += lda; + } + } + compute(& acc_0, 0, 0, comparray, vs, fin_res); + compute(& acc_1, 4, 4, comparray, vs, fin_res); + } + save_res(ii, jj, 0, fin_res); + save_res(ii + 4, jj, 4, fin_res); + } + + void KERNEL_8x8(int64_t ii, int64_t jj) { + vec_t vec_A[16], vec_B[16] = {0}; + acc_t acc_0, acc_1, acc_2, acc_3; + acc_t acc_4, acc_5, acc_6, acc_7; + std::array comparray {}; + vector float fin_res[16] = {0}; + vector float vs[16] = {0}; + bool isAblock_q4 = std::is_same_v; + for (int l = 0; l < k; l++) { + __builtin_mma_xxsetaccz(& acc_0); + __builtin_mma_xxsetaccz(& acc_1); + __builtin_mma_xxsetaccz(& acc_2); + __builtin_mma_xxsetaccz(& acc_3); + if (std::is_same_v) { + packNormalInt4<8>((A + (ii * lda) + l), lda, 8, 4, (int8_t *)vec_A, comparray); + } else { + packNormal((const block_q8_0 *)(A + (ii * lda) + l), lda, 8, 8, (int8_t *)vec_A, false); + } + packNormal((B + (jj * ldb) + l), ldb, 8, 8, (uint8_t *)vec_B, true); + for(int x = 0; x < 8; x++) { + __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x], vec_B[x]); + __builtin_mma_xvi8ger4pp(& acc_1, vec_A[x + 8], vec_B[x]); + __builtin_mma_xvi8ger4pp(& acc_2, vec_A[x], vec_B[x + 8]); + __builtin_mma_xvi8ger4pp(& acc_3, vec_A[x + 8], vec_B[x + 8]); + } + for (int I = 0; I < 8 ; I++) { + for (int J = 0; J < 4; J++) { + *((float *)& vs[I] + J) = (unhalf((A + ((ii + I) * lda) + l)->d) * unhalf((B + ((jj + J) * ldb) + l)->d)); + *((float *)& vs[I + 8] + J) = (unhalf((A + ((ii + I) * lda) + l)->d) * unhalf((B + ((jj + J + 4) * ldb) + l)->d)); + } + } + if (!isAblock_q4) { + auto aoffset = A + (ii * lda) + l; + for (int i = 0; i < 8; i++) { + comparray[i] = 0; + int ca = 0; + auto *at = aoffset->qs; + for (int j = 0; j < 32; j++) + ca += (int)*at++; + comparray[i] = ca; + aoffset += lda; + } + } + compute(& acc_0, 0, 0, comparray, vs, fin_res); + compute(& acc_1, 4, 4, comparray, vs, fin_res); + compute(& acc_2, 0, 8, comparray, vs, fin_res); + compute(& acc_3, 4, 12, comparray, vs, fin_res); + } + save_res(ii, jj, 0, fin_res); + save_res(ii + 4, jj, 4, fin_res); + save_res(ii, jj + 4, 8, fin_res); + save_res(ii + 4, jj + 4, 12, fin_res); + } + + void KERNEL_Q0(int64_t ii, int64_t jj, int64_t mc, int64_t nc, int64_t kc, int64_t l, vec_t * vec_A, vec_t * vec_B) { + acc_t acc[8]; + for (int i = 0; i < mc ; i += 16) { + for (int j = 0; j < nc; j += 8) { + int A0_base = (i / 16) * (2 * 32 * kc); + int B0_base = (j / 8) * (32 * kc); + for (int x = 0; x < 8; x++) { + __builtin_mma_xxsetaccz(&acc[x]); + } + for (int64_t kk = 0; kk < kc; kk++) { + int A0_block_idx = A0_base + kk * 32; + int B0_block_idx = B0_base + kk * 32; + int A1_block_idx = A0_block_idx + 32 * kc; + int B1_block_idx = B0_block_idx + 32 * kc; + vec_t * A0_block = & vec_A[A0_block_idx]; + vec_t * B0_block = & vec_B[B0_block_idx]; + vec_t * A1_block = & vec_A[A1_block_idx]; + for (int it = 0; it < 4; it++) { + for (int x = 0; x < 4; x++) { + __builtin_mma_xvf16ger2pp(& acc[0], A0_block[8 * it + x], B0_block[8 * it + x]); + __builtin_mma_xvf16ger2pp(& acc[1], A0_block[8 * it + x], B0_block[8 * it + x + 4]); + __builtin_mma_xvf16ger2pp(& acc[2], A0_block[8 * it + x + 4], B0_block[8 * it + x]); + __builtin_mma_xvf16ger2pp(& acc[3], A0_block[8 * it + x + 4], B0_block[8 * it + x + 4]); + __builtin_mma_xvf16ger2pp(& acc[4], A1_block[8 * it + x], B0_block[8 * it + x]); + __builtin_mma_xvf16ger2pp(& acc[5], A1_block[8 * it + x], B0_block[8 * it+ x + 4]); + __builtin_mma_xvf16ger2pp(& acc[6], A1_block[8 * it + x + 4], B0_block[8 * it + x]); + __builtin_mma_xvf16ger2pp(& acc[7], A1_block[8 * it + x + 4], B0_block[8 * it + x + 4]); + } + } + } + if (l == 0) { + save_acc(& acc[0], ii + i, jj + j); + save_acc(& acc[1], ii + i, jj + j + 4); + save_acc(& acc[2], ii + i + 4, jj + j); + save_acc(& acc[3], ii + i + 4, jj + j + 4); + save_acc(& acc[4], ii + i + 8, jj + j); + save_acc(& acc[5], ii + i + 8, jj + j + 4); + save_acc(& acc[6], ii + i + 12, jj + j); + save_acc(& acc[7], ii + i + 12, jj + j + 4); + } else { + add_save_acc(& acc[0], ii + i, jj + j); + add_save_acc(& acc[1], ii + i, jj + j + 4); + add_save_acc(& acc[2], ii + i + 4, jj + j); + add_save_acc(& acc[3], ii + i + 4, jj + j + 4); + add_save_acc(& acc[4], ii + i + 8, jj + j); + add_save_acc(& acc[5], ii + i + 8, jj + j + 4); + add_save_acc(& acc[6], ii + i + 12, jj + j); + add_save_acc(& acc[7], ii + i + 12, jj + j + 4); + } + } + } + } + + void matmul_tiled(int64_t m, int64_t n, int64_t mc, int64_t nc, int64_t kc) { + vec_t A_pack[mc * kc * 4]; + vec_t B_pack[nc * kc * 4]; + constexpr bool is_Ablock_q4 = std::is_same_v; + int64_t ytiles = m / mc; + int64_t xtiles = n / nc; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) { + end = tiles; + } + for (int64_t job = start; job < end; ++job) { + int64_t ii = (job / xtiles) * mc; + int64_t jj = (job % xtiles) * nc; + for (int64_t kk = 0; kk < k; kk += kc) { + if constexpr(is_Ablock_q4) { + packNormal_q4_fp16(A + ii * lda + kk, lda, mc, kc, (uint8_t *)A_pack); + } else { + packNormal_q8_fp16(A + ii * lda + kk, lda, mc, kc, (uint8_t *)A_pack); + } + packNormal_q8_fp16(B + jj * ldb + kk, ldb, nc, kc, (uint8_t *)B_pack); + KERNEL_Q0(ii, jj, mc, nc, kc, kk, A_pack, B_pack); + } + } + } + + void gemm_small(int64_t m0, int64_t m, int64_t n0, int64_t n, int RM, int RN) { + int64_t ytiles = (m - m0) / RM; + int64_t xtiles = (n - n0) / RN; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + vec_t vec_A[8] = {0}, vec_B[8] = {0}; + vector signed int vec_C[4]; + acc_t acc_0; + bool isAblock_q4 = std::is_same_v; + + if (end > tiles) + end = tiles; + for (int64_t job = start; job < end; ++job) { + int64_t ii = m0 + job / xtiles * RM; + int64_t jj = n0 + job % xtiles * RN; + std::array comparray{}; + vector float res[4] = {0}; + vector float fin_res[4] = {0}; + vector float vs[4] = {0}; + vector float CA[4] = {0}; + __builtin_prefetch((A + (ii * lda) + 0)->qs, 0, 1); // prefetch first value + __builtin_prefetch((B + (jj * ldb) + 0)->qs, 0, 1); // prefetch first value + for (int l = 0; l < k; l++) { + __builtin_prefetch((A + (ii * lda) + (l + 1))->qs, 0, 1); // prefetch one loop ahead + __builtin_prefetch((B + (jj * ldb) + (l + 1))->qs, 0, 1); // prefetch one loop ahead + __builtin_mma_xxsetaccz(& acc_0); + if (isAblock_q4) { + packNormalInt4<4>((A + (ii * lda) + l), lda, RM, 4, (int8_t *)vec_A, comparray); + } else { + packNormal((const block_q8_0 *)(A + (ii * lda) + l), lda, RM, 8, (int8_t *)vec_A, false); + } + packNormal((B + (jj * ldb) + l), ldb, RN, 8, (uint8_t *)vec_B, true); + for (int x = 0; x < 8; x += 4) { + __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x], vec_B[x]); + __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x + 1], vec_B[x + 1]); + __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x + 2], vec_B[x + 2]); + __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x + 3], vec_B[x + 3]); + } + for (int I = 0; I < RM; I++) { + for (int J = 0; J < RN; J++) { + *((float*)&vs[I] + J) = (unhalf((A + ((ii + I) * lda) + l)->d) * unhalf((B + ((jj + J) * ldb) + l)->d)); + } + } + __builtin_mma_disassemble_acc(vec_C, & acc_0); + if (!isAblock_q4) { + auto aoffset = A + (ii * lda) + l; + for (int i = 0; i < RM; i++) { + comparray[i] = 0; + int ca = 0; + auto *at = aoffset->qs; + for (int j = 0; j < 32; j++) + ca += (int)*at++; + comparray[i] = ca; + aoffset += lda; + } + } + for (int i = 0; i < RM; i++) { + CA[i] = vec_splats((float)(((double)comparray[i]) * -128.0)); + res[i] = vec_add(vec_ctf(vec_C[i], 0), CA[i]); + fin_res[i] = vec_madd(res[i], vs[i], fin_res[i]); + } + } + save_res(ii, jj, 0, fin_res, RM, RN); + } + } + + template + inline void kernel(int64_t ii, int64_t jj) { + if constexpr(RM == 4 && RN == 8) { + KERNEL_4x8(ii,jj); + } else if constexpr(RM == 8 && RN == 4) { + KERNEL_8x4(ii,jj); + } else if constexpr(RM == 8 && RN == 8) { + KERNEL_8x8(ii,jj); + } else { + assert(false && "RN/RM values not supported"); + } + } + + template + NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t ytiles = (m - m0) / RM; + int64_t xtiles = (n - n0) / RN; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) + end = tiles; + for (int64_t job = start; job < end; ++job) { + int64_t ii = m0 + job / xtiles * RM; + int64_t jj = n0 + job % xtiles * RN; + kernel(ii, jj); + } + } + const TA * const A; + const block_q8_0 * const B; + float * C; + const int64_t k; + int64_t kc; + const int64_t lda; + const int64_t ldb; + const int64_t ldc; + const int ith; + const int nth; +}; + +class tinyBLAS_PPC { + public: + tinyBLAS_PPC(int64_t k, + const float * A, int64_t lda, + const float * B, int64_t ldb, + float * C, int64_t ldc, + int ith, int nth) + : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { + } + + void matmul(int64_t m, int64_t n) { + #if defined(_AIX) || defined(__BIG_ENDIAN__) + mnpack(0, m, 0, n); + #else + int64_t mc = 256; int64_t nc = 256; int64_t kc = 256; + if (m % mc == 0 && n % nc == 0 && k % kc == 0) { + matmul_tiled(m, n, mc, nc, kc); + } else { + mnpack(0, m, 0, n); + } + #endif + } + + private: + + __attribute__((always_inline)) + inline void save_acc(acc_t * ACC, int64_t ii, int64_t jj) { + vec_t vec_C[4]; + __builtin_mma_disassemble_acc(vec_C, ACC); + for (int I = 0; I < 4; I++) { + for (int J = 0; J < 4; J++) { + *((float *)(C+ii+((jj+J)*ldc)+I)) = *((float *)&vec_C[I]+J); + } + } + } + + __attribute__((always_inline)) + inline void add_save_acc(acc_t * ACC, int64_t ii, int64_t jj) { + vec_t vec_C[4]; + __builtin_mma_disassemble_acc(vec_C, ACC); + for (int I = 0; I < 4; I++) { + for (int J = 0; J < 4; J++) { + float * c_ptr = (float *)(C+ii+((jj+J)*ldc)+I); + *c_ptr += *((float *)&vec_C[I]+J); + } + } + } + + inline void vector_permute_store_4(vector float * src, float * vecOffset) { + vector float t1, t2, t3, t4, t5, t6, t7, t8; + t1 = vec_mergeh(src[0], src[1]); + t2 = vec_mergeh(src[2], src[3]); + t3 = vec_mergel(src[0], src[1]); + t4 = vec_mergel(src[2], src[3]); + + t5 = vec_xxpermdi(t1, t2, 0); + t6 = vec_xxpermdi(t1, t2, 3); + t7 = vec_xxpermdi(t3, t4, 0); + t8 = vec_xxpermdi(t3, t4, 3); + + vec_xst(t5, 0, vecOffset); + vec_xst(t6, 0, vecOffset + 4); + vec_xst(t7, 0, vecOffset + 8); + vec_xst(t8, 0, vecOffset + 12); + } + + inline void vector_permute_store_8(vector float * src, float * vecOffset) { + vector float t1, t2, t3, t4, t5, t6, t7, t8; + t1 = vec_mergeh(src[0], src[1]); + t2 = vec_mergeh(src[2], src[3]); + t3 = vec_mergeh(src[4], src[5]); + t4 = vec_mergeh(src[6], src[7]); + + t5 = vec_xxpermdi(t1, t2, 0); + t6 = vec_xxpermdi(t3, t4, 0); + t7 = vec_xxpermdi(t1, t2, 3); + t8 = vec_xxpermdi(t3, t4, 3); + + vec_xst(t5, 0, vecOffset); + vec_xst(t6, 0, vecOffset + 4); + vec_xst(t7, 0, vecOffset + 8); + vec_xst(t8, 0, vecOffset + 12); + + t1 = vec_mergel(src[0], src[1]); + t2 = vec_mergel(src[2], src[3]); + t3 = vec_mergel(src[4], src[5]); + t4 = vec_mergel(src[6], src[7]); + + t5 = vec_xxpermdi(t1, t2, 0); + t6 = vec_xxpermdi(t3, t4, 0); + t7 = vec_xxpermdi(t1, t2, 3); + t8 = vec_xxpermdi(t3, t4, 3); + + vec_xst(t5, 0, vecOffset + 16); + vec_xst(t6, 0, vecOffset + 20); + vec_xst(t7, 0, vecOffset + 24); + vec_xst(t8, 0, vecOffset + 28); + } + + void packTranspose(const float * a, int64_t lda, int rows, int cols, float * vec) { + int64_t i, j; + float * aoffsets[8]; + float * aoffset = NULL, * boffset = NULL; + __vector_pair arr[8]; + vector float c[8][2] = {0}; + vector float c1[8] = {0}; + vector float c2[8] = {0}; + aoffset = const_cast(a); + boffset = vec; + j = (rows >> 3); + if (j > 0) { + do { + aoffsets[0] = aoffset; + for (int it = 1; it < 8; it++) + aoffsets[it] = aoffsets[it-1] + lda; + aoffset += 8 * lda; + i = (cols >> 3); + if (i > 0) { + do { + for (int it = 0; it < 8; it++) { + arr[it] = __builtin_vsx_lxvp(0, (__vector_pair*)aoffsets[it]); + __builtin_vsx_disassemble_pair(c[it], &arr[it]); + c1[it] = c[it][0]; + c2[it] = c[it][1]; + } + + vector_permute_store_8(c1, boffset); + vector_permute_store_8(c2, boffset + 32); + boffset += 64; + i--; + if (i > 0) { + for (int it = 0; it < 8; it++) { + aoffsets[it] = aoffsets[it] + 8; + } + } + } while(i > 0); + } + if (cols & 4) { + for (int it = 0; it < 8 ; it++) + c1[it] = vec_xl(0, aoffsets[it]); + vector_permute_store_8(c1, boffset); + } + j--; + } while(j > 0); + } + + if (rows & 4) { + aoffsets[0] = aoffset; + for (int it = 1; it < 4; it++) + aoffsets[it] = aoffsets[it-1] + lda; + aoffset += 4 * lda; + i = (cols >> 3); + if (i > 0) { + do { + for (int it = 0; it < 4; it++) { + arr[it] = __builtin_vsx_lxvp(0, (__vector_pair*)aoffsets[it]); + __builtin_vsx_disassemble_pair(c[it], &arr[it]); + c1[it] = c[it][0]; + c2[it] = c[it][1]; + } + vector_permute_store_4(c1, boffset); + vector_permute_store_4(c2, boffset + 16); + for (int it = 0; it < 4; it++) + aoffsets[it] += 8 * lda; + boffset += 32; + i--; + } while(i > 0); + } + + if (cols & 4) { + for (int it = 0; it < 4; it++) + c1[it] = vec_xl(0, aoffsets[it]); + vector_permute_store_4(c1, boffset); + } + } + if (rows & 3) { + aoffsets[0] = aoffset; + for (int it = 1; it < 3; it++) + aoffsets[it] = aoffsets[it-1] + lda; + if (cols & 4) { + for (int it = 0; it < 3; it++) + c1[it] = vec_xl(0, aoffsets[it]); + vector_permute_store_4(c1, boffset); + } + } + } + + void KERNEL_4x4(int64_t ii, int64_t jj) { + vec_t vec_A[4], vec_B[4], vec_C[4]; + acc_t acc_0; + __builtin_mma_xxsetaccz(&acc_0); + for (int l = 0; l < k; l += 4) { + packTranspose(A + (ii * lda) + l, lda, 4, 4, (float *)vec_A); + packTranspose(B + (jj * ldb) + l, ldb, 4, 4, (float *)vec_B); + __builtin_mma_xvf32gerpp(&acc_0, vec_A[0], vec_B[0]); + __builtin_mma_xvf32gerpp(&acc_0, vec_A[1], vec_B[1]); + __builtin_mma_xvf32gerpp(&acc_0, vec_A[2], vec_B[2]); + __builtin_mma_xvf32gerpp(&acc_0, vec_A[3], vec_B[3]); + } + save_acc(&acc_0, ii, jj); + } + + void KERNEL_4x8(int64_t ii, int64_t jj) { + vec_t vec_A[4], vec_B[8], vec_C[4]; + acc_t acc_0, acc_1; + __builtin_mma_xxsetaccz(&acc_0); + __builtin_mma_xxsetaccz(&acc_1); + for (int64_t l = 0; l < k; l += 4) { + packTranspose(A + (ii * lda) + l, lda, 4, 4, (float *)vec_A); + packTranspose(B + (jj * ldb) + l, ldb, 8, 4, (float *)vec_B); + __builtin_mma_xvf32gerpp(&acc_0, vec_A[0], (vec_t)vec_B[0]); + __builtin_mma_xvf32gerpp(&acc_1, vec_A[0], (vec_t)vec_B[1]); + __builtin_mma_xvf32gerpp(&acc_0, vec_A[1], (vec_t)vec_B[2]); + __builtin_mma_xvf32gerpp(&acc_1, vec_A[1], (vec_t)vec_B[3]); + __builtin_mma_xvf32gerpp(&acc_0, vec_A[2], (vec_t)vec_B[4]); + __builtin_mma_xvf32gerpp(&acc_1, vec_A[2], (vec_t)vec_B[5]); + __builtin_mma_xvf32gerpp(&acc_0, vec_A[3], (vec_t)vec_B[6]); + __builtin_mma_xvf32gerpp(&acc_1, vec_A[3], (vec_t)vec_B[7]); + } + save_acc(&acc_0, ii, jj); + save_acc(&acc_1, ii, jj + 4); + } + + void KERNEL_8x4(int64_t ii, int64_t jj) { + vec_t vec_A[8], vec_B[4], vec_C[4]; + acc_t acc_0, acc_1; + __builtin_mma_xxsetaccz(&acc_0); + __builtin_mma_xxsetaccz(&acc_1); + for (int64_t l = 0; l < k; l += 4) { + packTranspose(A + (ii * lda) + l, lda, 8, 4, (float *)vec_A); + packTranspose(B + (jj * ldb) + l, ldb, 4, 4, (float *)vec_B); + __builtin_mma_xvf32gerpp(&acc_0, (vec_t)vec_A[0], vec_B[0]); + __builtin_mma_xvf32gerpp(&acc_1, (vec_t)vec_A[1], vec_B[0]); + __builtin_mma_xvf32gerpp(&acc_0, (vec_t)vec_A[2], vec_B[1]); + __builtin_mma_xvf32gerpp(&acc_1, (vec_t)vec_A[3], vec_B[1]); + __builtin_mma_xvf32gerpp(&acc_0, (vec_t)vec_A[4], vec_B[2]); + __builtin_mma_xvf32gerpp(&acc_1, (vec_t)vec_A[5], vec_B[2]); + __builtin_mma_xvf32gerpp(&acc_0, (vec_t)vec_A[6], vec_B[3]); + __builtin_mma_xvf32gerpp(&acc_1, (vec_t)vec_A[7], vec_B[3]); + } + save_acc(&acc_0, ii, jj); + save_acc(&acc_1, ii + 4, jj); + } + + void KERNEL_8x8(int64_t ii, int64_t jj) { + vec_t vec_A[16], vec_B[16], vec_C[4]; + acc_t acc_0, acc_1, acc_2, acc_3; + __builtin_mma_xxsetaccz(&acc_0); + __builtin_mma_xxsetaccz(&acc_1); + __builtin_mma_xxsetaccz(&acc_2); + __builtin_mma_xxsetaccz(&acc_3); + for (int l = 0; l < k; l+=8) { + packTranspose(A + (ii * lda) + l, lda, 8, 8, (float *)vec_A); + packTranspose(B + (jj * ldb) + l, ldb, 8, 8, (float *)vec_B); + for(int x = 0; x < 16; x+=2) { + __builtin_mma_xvf32gerpp(&acc_0, (vec_t)vec_A[x], vec_B[x]); + __builtin_mma_xvf32gerpp(&acc_1, (vec_t)vec_A[x], vec_B[x + 1]); + __builtin_mma_xvf32gerpp(&acc_2, (vec_t)vec_A[x + 1], vec_B[x]); + __builtin_mma_xvf32gerpp(&acc_3, (vec_t)vec_A[x + 1], vec_B[x + 1]); + } + } + save_acc(&acc_0, ii, jj); + save_acc(&acc_1, ii, jj + 4); + save_acc(&acc_2, ii + 4, jj); + save_acc(&acc_3, ii + 4, jj + 4); + } + + inline void MMA_16x8(vec_t * vec_A0, vec_t * vec_A1, vec_t * vec_B, acc_t * acc) { + for (int x = 0; x < 16; x += 2) { + __builtin_mma_xvf32gerpp(&acc[0], vec_A0[x + 0], vec_B[x]); + __builtin_mma_xvf32gerpp(&acc[1], vec_A0[x + 0], vec_B[x + 1]); + __builtin_mma_xvf32gerpp(&acc[2], vec_A0[x + 1], vec_B[x]); + __builtin_mma_xvf32gerpp(&acc[3], vec_A0[x + 1], vec_B[x + 1]); + __builtin_mma_xvf32gerpp(&acc[4], vec_A1[x + 0], vec_B[x]); + __builtin_mma_xvf32gerpp(&acc[5], vec_A1[x + 0], vec_B[x + 1]); + __builtin_mma_xvf32gerpp(&acc[6], vec_A1[x + 1], vec_B[x]); + __builtin_mma_xvf32gerpp(&acc[7], vec_A1[x + 1], vec_B[x + 1]); + } + } + + void KERNEL(int64_t ii, int64_t jj, int64_t mc, int64_t nc, int64_t kc, vec_t * vec_A, vec_t * vec_B, int64_t kk) { + for (int64_t i = 0; i < mc; i += 16) { + int A_base_addr = (mc / 8) * (i / 8) * 16; + for (int64_t j = 0; j < nc; j += 8) { + int B_base_addr = (nc / 8) * (j / 8) * 16; + acc_t acc[8]; + vec_t A0_block[16]; vec_t A1_block[16]; + for (int x = 0; x < 8; x++) + __builtin_mma_xxsetaccz(&acc[x]); + for (int64_t l = 0; l < kc; l += 8) { + int A0_block_idx = A_base_addr + (l / 8) * 16; + int A1_block_idx = A0_block_idx + (mc / 8) * 16; + int B_block_idx = B_base_addr + (l / 8) * 16; + vec_t* A0_block = &vec_A[A0_block_idx]; + vec_t* A1_block = &vec_A[A1_block_idx]; + vec_t* B_block = &vec_B[B_block_idx]; + MMA_16x8(A0_block, A1_block, B_block, acc); + } + if (kk == 0) { + save_acc(&acc[0], ii + i, jj + j); + save_acc(&acc[1], ii + i, jj + j + 4); + save_acc(&acc[2], ii + i + 4, jj + j); + save_acc(&acc[3], ii + i + 4, jj + j + 4); + save_acc(&acc[4], ii + i + 8, jj + j); + save_acc(&acc[5], ii + i + 8, jj + j + 4); + save_acc(&acc[6], ii + i + 12, jj + j); + save_acc(&acc[7], ii + i + 12, jj + j + 4); + } else { + add_save_acc(&acc[0], ii + i, jj + j); + add_save_acc(&acc[1], ii + i, jj + j + 4); + add_save_acc(&acc[2], ii + i + 4, jj + j); + add_save_acc(&acc[3], ii + i + 4, jj + j + 4); + add_save_acc(&acc[4], ii + i + 8, jj + j); + add_save_acc(&acc[5], ii + i + 8, jj + j + 4); + add_save_acc(&acc[6], ii + i + 12, jj + j); + add_save_acc(&acc[7], ii + i + 12, jj + j + 4); + } + } + } + } + + void matmul_tiled(int64_t m , int64_t n, int64_t mc, int64_t nc, int64_t kc) { + int64_t ytiles = m / mc; + int64_t xtiles = n / nc; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) { + end = tiles; + } + for (int64_t job = start; job < end; ++job) { + int64_t ii = (job / xtiles) * mc; + int64_t jj = (job % xtiles) * nc; + for (int64_t kk = 0; kk < k; kk += kc) { + vec_t A_pack[kc * mc / 4]; + vec_t B_pack[kc * nc / 4]; + packTranspose(A + (ii * lda) + kk, lda, kc, mc, (float *)A_pack); + packTranspose(B + (jj * ldb) + kk, ldb, kc, nc, (float *)B_pack); + KERNEL(ii, jj, mc, nc, kc, A_pack, B_pack, kk); + } + } + } + + void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int m_rem = MIN(m - m0, 8); + int n_rem = MIN(n - n0, 8); + int mc = 0, nc = 0; + if (m_rem >= 8 && n_rem >= 8) { + mc = 8; + nc = 8; + gemm<8, 8>(m0, m, n0, n); + } else if (m_rem >= 4 && n_rem >= 8) { + mc = 4; + nc = 8; + gemm<4, 8>(m0, m, n0, n); + } else if (m_rem >= 8 && n_rem >= 4) { + mc = 8; + nc = 4; + gemm<8, 4>(m0, m, n0, n); + } else if (m_rem >= 4 && n_rem >= 4) { + mc = 4; + nc = 4; + gemm<4, 4>(m0, m, n0, n); + } else { + mc = (m_rem >= 4) ? 4 : m_rem; + nc = (n_rem >= 4) ? 4 : n_rem; + if (mc == 0 || nc == 0) + return; + gemm_small(m0, m, n0, n, mc, nc); + } + int64_t mp = m0 + ((m - m0) / mc) * mc; + int64_t np = n0 + ((n - n0) / nc) * nc; + mnpack(mp, m, n0, np); + mnpack(m0, m, np, n); + } + + void gemm_small(int64_t m0, int64_t m, int64_t n0, int64_t n, int RM, int RN) { + int64_t ytiles = (m - m0) / RM; + int64_t xtiles = (n - n0) / RN; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) + end = tiles; + for (int64_t job = start; job < end; ++job) { + int64_t ii = m0 + job / xtiles * RM; + int64_t jj = n0 + job % xtiles * RN; + vec_t vec_C[4]; + acc_t acc_0; + __builtin_mma_xxsetaccz(&acc_0); + vec_t vec_A[4] = {0}, vec_B[4] = {0}; + for (int l = 0; l < k; l += 4) { + /* 'GEMV Forwarding' concept is used in first two conditional loops. + * when one of the matrix has a single row/column, the elements are + * broadcasted, instead of using packing routine to prepack the + * matrix elements. + */ + if (RM == 1) { + float * a = const_cast(A + (ii) * lda + l); + packTranspose(B + (jj * ldb) + l, ldb, RN, 4, (float *)vec_B); + vec_A[0] = (vec_t)vec_xl(0,a); + vec_A[1] = (vec_t)vec_splats(*((float *)&vec_A+1)); + vec_A[2] = (vec_t)vec_splats(*((float *)&vec_A+2)); + vec_A[3] = (vec_t)vec_splats(*((float *)&vec_A+3)); + } else if (RN == 1) { + packTranspose(A + (ii * lda) + l, lda, RM, 4, (float *)vec_A); + float * b = const_cast(B + (jj) * ldb + l); + vec_B[0] = (vec_t)vec_xl(0,b); + vec_B[1] = (vec_t)vec_splats(*((float *)&vec_B+1)); + vec_B[2] = (vec_t)vec_splats(*((float *)&vec_B+2)); + vec_B[3] = (vec_t)vec_splats(*((float *)&vec_B+3)); + } else { + packTranspose(A + (ii * lda) + l, lda, RM, 4, (float *)vec_A); + packTranspose(B + (jj * ldb) + l, ldb, RN, 4, (float *)vec_B); + } + __builtin_mma_xvf32gerpp(&acc_0, vec_A[0], vec_B[0]); + __builtin_mma_xvf32gerpp(&acc_0, vec_A[1], vec_B[1]); + __builtin_mma_xvf32gerpp(&acc_0, vec_A[2], vec_B[2]); + __builtin_mma_xvf32gerpp(&acc_0, vec_A[3], vec_B[3]); + } + __builtin_mma_disassemble_acc(vec_C, &acc_0); + for (int I = 0; I < RM; I++) { + for (int J = 0; J < RN; J++) { + *((float *)(C+ii+((jj+J)*ldc)+I)) = *((float *)&vec_C[I]+J); + } + } + } + } + + template + inline void kernel(int64_t ii, int64_t jj) { + if constexpr(RM == 4 && RN == 4) { + KERNEL_4x4(ii, jj); + } else if constexpr(RM == 4 && RN == 8) { + KERNEL_4x8(ii, jj); + } else if constexpr(RM == 8 && RN == 4) { + KERNEL_8x4(ii, jj); + } else if constexpr(RM == 8 && RN == 8) { + KERNEL_8x8(ii, jj); + } else { + static_assert(false, "RN/RM values not supported"); + } + } + + template + NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t ytiles = (m - m0) / RM; + int64_t xtiles = (n - n0) / RN; + int64_t tiles = xtiles * ytiles; + int64_t duty = (tiles + nth - 1) / nth; + int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) + end = tiles; + for (int64_t job = start; job < end; ++job) { + int64_t ii = m0 + job / xtiles * RM; + int64_t jj = n0 + job % xtiles * RN; + kernel(ii, jj); + } + } + + const float * const A; + const float * const B; + float * C; + const int64_t k; + const int64_t lda; + const int64_t ldb; + const int64_t ldc; + const int ith; + const int nth; +}; +#endif +} // namespace + +/** + * Performs optimized matrix multiplication on CPU. + * + * This subroutine may compute C = Aᵀ * B with column major ordering. + * Despite its name, this isn't a generalized implementation. Work is + * only performed when a handwritten kernel is written and available. + * Otherwise the caller should fall back to a general matmul routine. + * + * For example, for single-threaded single-precision GEMM you can say + * + * llamafile_sgemm(m, n, k, A, lda, B, ldb, C, ldc, + * 0, 1, + * GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32); + * + * @param m is rows in `A` and `C` + * @param n is cols in `B` and `C` + * @param k is cols in `A` and rows in `B` + * @param A is first input matrix (always transposed) + * @param lda is row stride of `A` + * @param B is second input matrix (never transposed) + * @param ldb is row stride of `B` + * @param C is input/output array of output matrices + * @param ldc is row stride of `C` + * @param ith is thread id (must be less than `nth`) + * @param nth is number of threads (must be greater than zero) + * @param Atype is GGML data type of `A` + * @param Btype is GGML data type of `B` + * @param Ctype is GGML data type of `C` + * @return true if this function was able to service the matmul request + */ +bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64_t n, int64_t k, + const void *A, int64_t lda, const void *B, int64_t ldb, void *C, + int64_t ldc, int Atype, int Btype, int Ctype) { + + assert(m >= 0); + assert(n >= 0); + assert(k >= 0); + assert(lda >= k); + assert(ldb >= k); + assert(ldc >= m); + assert(params->nth > 0); + assert(params->ith < params->nth); + + // only enable sgemm for prompt processing +#if !defined(__MMA__) + if (n < 2) + return false; +#endif + + if (Ctype != GGML_TYPE_F32) + return false; + + switch (Atype) { + + case GGML_TYPE_F32: { + if (Btype != GGML_TYPE_F32) + return false; +#if defined(__AVX512F__) + tinyBLAS<16, __m512, __m512, float, float, float> tb{ params, + k, (const float *)A, lda, + (const float *)B, ldb, + (float *)C, ldc}; + return tb.matmul(m, n); +#elif defined(__AVX__) || defined(__AVX2__) + tinyBLAS<8, __m256, __m256, float, float, float> tb{ params, + k, (const float *)A, lda, + (const float *)B, ldb, + (float *)C, ldc}; + return tb.matmul(m, n); +#elif defined(__ARM_NEON) + if (n < 4) + return false; + tinyBLAS<4, float32x4_t, float32x4_t, float, float, float> tb{ params, + k, (const float *)A, lda, + (const float *)B, ldb, + (float *)C, ldc}; + return tb.matmul(m, n); +#elif defined(__VXE__) || defined(__VXE2__) + if (n < 4) + return false; + tinyBLAS<4, float32x4_t, float32x4_t, float, float, float> tb{ params, + k, (const float *)A, lda, + (const float *)B, ldb, + (float *)C, ldc}; + return tb.matmul(m, n); +#elif defined(__MMA__) + if (k % 8) + return false; + tinyBLAS_PPC tb{ + k, (const float *)A, lda, + (const float *)B, ldb, + (float *)C, ldc, + params->ith, params->nth}; + tb.matmul(m, n); + return true; +#elif defined(__riscv_v_intrinsic) + #if LMUL == 1 + tinyBLAS_RVV tb{ params, + k, (const float *)A, lda, + (const float *)B, ldb, + (float *)C, ldc}; + #elif LMUL == 2 + tinyBLAS_RVV tb{ params, + k, (const float *)A, lda, + (const float *)B, ldb, + (float *)C, ldc}; + #else // LMUL = 4 + tinyBLAS_RVV tb{ params, + k, (const float *)A, lda, + (const float *)B, ldb, + (float *)C, ldc}; + #endif + return tb.matmul(m, n); +#else + return false; +#endif + } + + case GGML_TYPE_BF16: { +#if defined(__AVX512BF16__) + if (Btype == GGML_TYPE_BF16) { + tinyBLAS<32, __m512, __m512bh, ggml_bf16_t, ggml_bf16_t, float> tb{ params, k, + (const ggml_bf16_t *)A, lda, + (const ggml_bf16_t *)B, ldb, + (float *)C, ldc}; + return tb.matmul(m, n); + } +#elif defined(__AVX512F__) + if (Btype == GGML_TYPE_BF16) { + tinyBLAS<16, __m512, __m512, ggml_bf16_t, ggml_bf16_t, float> tb{ params, k, + (const ggml_bf16_t *)A, lda, + (const ggml_bf16_t *)B, ldb, + (float *)C, ldc}; + return tb.matmul(m, n); + } +#elif defined(__AVX2__) + if (Btype == GGML_TYPE_BF16) { + tinyBLAS<8, __m256, __m256, ggml_bf16_t, ggml_bf16_t, float> tb{ params, k, + (const ggml_bf16_t *)A, lda, + (const ggml_bf16_t *)B, ldb, + (float *)C, ldc}; + return tb.matmul(m, n); + } +#elif defined(__MMA__) + if (k % 8) { + return false; + } + + if (Btype == GGML_TYPE_BF16) { + tinyBLAS_HP16_PPC tb{ k, + (const ggml_bf16_t *)A, lda, + (const ggml_bf16_t *)B, ldb, + (float *)C, ldc, + params->ith, params->nth }; + + tb.matmul(m, n); + return true; + } +#elif defined(__riscv_zvfbfwma) + if (Btype == GGML_TYPE_BF16) { + #if LMUL == 1 + tinyBLAS_RVV tb{ params, + k, (const ggml_bf16_t *)A, lda, + (const ggml_bf16_t *)B, ldb, + (float *)C, ldc}; + #elif LMUL == 2 + tinyBLAS_RVV tb{ params, + k, (const ggml_bf16_t *)A, lda, + (const ggml_bf16_t *)B, ldb, + (float *)C, ldc}; + #else // LMUL = 4 + tinyBLAS_RVV tb{ params, + k, (const ggml_bf16_t *)A, lda, + (const ggml_bf16_t *)B, ldb, + (float *)C, ldc}; + #endif + return tb.matmul(m, n); + } +#endif + return false; + } + + case GGML_TYPE_F16: { +#if defined(__AVX512F__) + if (Btype == GGML_TYPE_F16) { + tinyBLAS<16, __m512, __m512, ggml_fp16_t, ggml_fp16_t, float> tb{ params, k, + (const ggml_fp16_t *)A, lda, + (const ggml_fp16_t *)B, ldb, + (float *)C, ldc}; + return tb.matmul(m, n); + } +#elif (defined(__AVX__) || defined(__AVX2__)) && defined(__F16C__) + if (Btype == GGML_TYPE_F16) { + tinyBLAS<8, __m256, __m256, ggml_fp16_t, ggml_fp16_t, float> tb{ params, k, + (const ggml_fp16_t *)A, lda, + (const ggml_fp16_t *)B, ldb, + (float *)C, ldc}; + return tb.matmul(m, n); + } +#elif defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER) + if (n < 8) + return false; + if (Btype == GGML_TYPE_F16) { + tinyBLAS<8, float16x8_t, float16x8_t, ggml_fp16_t, ggml_fp16_t, float> tb{ params, + k, (const ggml_fp16_t *)A, lda, + (const ggml_fp16_t *)B, ldb, + (float *)C, ldc}; + return tb.matmul(m, n); + } +#elif defined(__ARM_NEON) && !defined(_MSC_VER) + if (Btype == GGML_TYPE_F32) { + tinyBLAS<4, float32x4_t, float32x4_t, ggml_fp16_t, float, float> tb{ params, + k, (const ggml_fp16_t *)A, lda, + (const float *)B, ldb, + (float *)C, ldc}; + return tb.matmul(m, n); + } +#elif defined(__VXE__) || defined(__VXE2__) + if (n < 4) + return false; + if (Btype == GGML_TYPE_F16) { + tinyBLAS<4, float32x4_t, float32x4_t, ggml_fp16_t, ggml_fp16_t, float> tb{ params, + k, (const ggml_fp16_t *)A, lda, + (const ggml_fp16_t *)B, ldb, + (float *)C, ldc}; + return tb.matmul(m, n); + } +#elif defined(__riscv_zvfh) + if (Btype == GGML_TYPE_F16) { + #if LMUL == 1 + tinyBLAS_RVV tb{ params, + k, (const ggml_fp16_t *)A, lda, + (const ggml_fp16_t *)B, ldb, + (float *)C, ldc}; + #elif LMUL == 2 + tinyBLAS_RVV tb{ params, + k, (const ggml_fp16_t *)A, lda, + (const ggml_fp16_t *)B, ldb, + (float *)C, ldc}; + #else // LMUL = 4 + tinyBLAS_RVV tb{ params, + k, (const ggml_fp16_t *)A, lda, + (const ggml_fp16_t *)B, ldb, + (float *)C, ldc}; + #endif + return tb.matmul(m, n); + } +#elif defined(__MMA__) + if (k % 8) { + return false; + } + + if (Btype == GGML_TYPE_F16) { + tinyBLAS_HP16_PPC tb{ k, + (const ggml_fp16_t *)A, lda, + (const ggml_fp16_t *)B, ldb, + (float *)C, ldc, + params->ith, params->nth }; + + tb.matmul(m, n); + return true; + } +#endif + return false; + } + + case GGML_TYPE_Q8_0: { + if (Btype != GGML_TYPE_Q8_0) + return false; +#if defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX__) + tinyBLAS_Q0_AVX tb{ + k, (const block_q8_0 *)A, lda, + (const block_q8_0 *)B, ldb, + (float *)C, ldc, + params->ith, params->nth}; + tb.matmul(m, n); + return true; +#elif defined(__ARM_FEATURE_DOTPROD) + tinyBLAS_Q0_ARM tb{ + k, (const block_q8_0 *)A, lda, + (const block_q8_0 *)B, ldb, + (float *)C, ldc, + params->ith, params->nth}; + tb.matmul(m, n); + return true; +#elif defined(__MMA__) + //TO-DO: Remove this condition once gemv forwarding is enabled. + if (n < 8 && n != 4) + return false; + if (m < 8 && m != 4) + return false; + tinyBLAS_Q0_PPC tb{ + k, (const block_q8_0 *)A, lda, + (const block_q8_0 *)B, ldb, + (float *)C, ldc, + params->ith, params->nth}; + tb.matmul(m, n); + return true; +#else + return false; +#endif + } + + case GGML_TYPE_Q4_0: { + if (Btype != GGML_TYPE_Q8_0) + return false; +#if defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX__) + tinyBLAS_Q0_AVX tb{ + k, (const block_q4_0 *)A, lda, + (const block_q8_0 *)B, ldb, + (float *)C, ldc, + params->ith, params->nth}; + tb.matmul(m, n); + return true; +#elif defined(__ARM_FEATURE_DOTPROD) + tinyBLAS_Q0_ARM tb{ + k, (const block_q4_0 *)A, lda, + (const block_q8_0 *)B, ldb, + (float *)C, ldc, + params->ith, params->nth}; + tb.matmul(m, n); + return true; +#elif defined(__MMA__) + //TO-DO: Remove this condition once gemv forwarding is enabled. + if (n < 8 && n != 4) + return false; + if (m < 8 && m != 4) + return false; + tinyBLAS_Q0_PPC tb{ + k, (const block_q4_0 *)A, lda, + (const block_q8_0 *)B, ldb, + (float *)C, ldc, + params->ith, params->nth}; + tb.matmul(m, n); + return true; +#else + return false; +#endif + } + + case GGML_TYPE_Q5_0: { + if (Btype != GGML_TYPE_Q8_0) + return false; +#if defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX__) + tinyBLAS_Q0_AVX tb{ + k, (const block_q5_0 *)A, lda, + (const block_q8_0 *)B, ldb, + (float *)C, ldc, + params->ith, params->nth}; + tb.matmul(m, n); + return true; +#else + return false; +#endif + } + + case GGML_TYPE_IQ4_NL: { + if (Btype != GGML_TYPE_Q8_0) + return false; +#if defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX__) + tinyBLAS_Q0_AVX tb{ + k, (const block_iq4_nl *)A, lda, + (const block_q8_0 *)B, ldb, + (float *)C, ldc, + params->ith, params->nth}; + tb.matmul(m, n); + return true; +#else + return false; +#endif + } + + default: + return false; + } + + (void)params; + (void)m; + (void)n; + (void)k; + (void)A; + (void)lda; + (void)B; + (void)ldb; + (void)C; + (void)ldc; + (void)Atype; + (void)Btype; + (void)Ctype; +} + + +#undef NOINLINE +#undef VECTOR_REGISTERS +#undef MM256_SET_M128I diff --git a/ggml/src/ggml-cpu/hbm.cpp b/ggml/src/ggml-cpu/hbm.cpp deleted file mode 100644 index a4073c15e6c9..000000000000 --- a/ggml/src/ggml-cpu/hbm.cpp +++ /dev/null @@ -1,55 +0,0 @@ -#ifdef GGML_USE_CPU_HBM - -#include "ggml-backend.h" -#include "ggml-backend-impl.h" -#include "ggml-cpu.h" -#include "ggml-impl.h" - -#include "hbm.h" - -// buffer type HBM - -#include - -static const char * ggml_backend_cpu_hbm_buffer_type_get_name(ggml_backend_buffer_type_t buft) { - return "CPU_HBM"; - - GGML_UNUSED(buft); -} - -static void ggml_backend_cpu_hbm_buffer_free_buffer(ggml_backend_buffer_t buffer) { - hbw_free(buffer->context); -} - -static ggml_backend_buffer_t ggml_backend_cpu_hbm_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, - size_t size) { - void * ptr; - int result = hbw_posix_memalign(&ptr, ggml_backend_cpu_buffer_type_get_alignment(buft), size); - if (result != 0) { - GGML_LOG_ERROR("failed to allocate HBM buffer of size %zu\n", size); - return NULL; - } - - ggml_backend_buffer_t buffer = ggml_backend_cpu_buffer_from_ptr(ptr, size); - buffer->buft = buft; - buffer->iface.free_buffer = ggml_backend_cpu_hbm_buffer_free_buffer; - - return buffer; -} - -ggml_backend_buffer_type_t ggml_backend_cpu_hbm_buffer_type(void) { - static struct ggml_backend_buffer_type ggml_backend_cpu_buffer_type_hbm = { - /* .iface = */ { - /* .get_name = */ ggml_backend_cpu_hbm_buffer_type_get_name, - /* .alloc_buffer = */ ggml_backend_cpu_hbm_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_cpu_buffer_type_get_alignment, - /* .get_max_size = */ nullptr, // defaults to SIZE_MAX - /* .get_alloc_size = */ nullptr, // defaults to ggml_nbytes - /* .is_host = */ ggml_backend_cpu_buffer_type_is_host, - }, - /* .context = */ nullptr, - }; - - return &ggml_backend_cpu_buffer_type_hbm; -} -#endif diff --git a/ggml/src/ggml-cpu/hbm.h b/ggml/src/ggml-cpu/hbm.h deleted file mode 100644 index 09a1f09d72be..000000000000 --- a/ggml/src/ggml-cpu/hbm.h +++ /dev/null @@ -1,8 +0,0 @@ -#pragma once - -#include "ggml-backend.h" -#include "ggml.h" - -// GGML CPU internal header - -ggml_backend_buffer_type_t ggml_backend_cpu_hbm_buffer_type(void); diff --git a/ggml/src/ggml-cpu/kleidiai/kernels.cpp b/ggml/src/ggml-cpu/kleidiai/kernels.cpp.inc similarity index 95% rename from ggml/src/ggml-cpu/kleidiai/kernels.cpp rename to ggml/src/ggml-cpu/kleidiai/kernels.cpp.inc index 8c4d7bc925f6..12b8a83c92a0 100644 --- a/ggml/src/ggml-cpu/kleidiai/kernels.cpp +++ b/ggml/src/ggml-cpu/kleidiai/kernels.cpp.inc @@ -33,14 +33,102 @@ #include "kai_rhs_pack_nxk_qsi8cxp_qsi8cx_neon.h" #include "kai_lhs_pack_f16pmrx2_f32_neon.h" -#include "kai_common.h" +#include "kai_common.h.inc" -#include "simd-mappings.h" +#include "simd-mappings-defs.inc" #define GGML_COMMON_DECL_CPP -#include "ggml-common.h" +#include "ggml-common-defs.inc" + +// SPDX-FileCopyrightText: Copyright 2025 Arm Limited and/or its affiliates +// SPDX-License-Identifier: MIT +// + +#include "ggml.h.inc" + +enum cpu_feature { + CPU_FEATURE_NONE = 0, + CPU_FEATURE_DOTPROD = 1, + CPU_FEATURE_I8MM = 2, + CPU_FEATURE_SVE = 4, + CPU_FEATURE_SME = 8 +}; + +inline cpu_feature& operator|=(cpu_feature& lhs, cpu_feature rhs) { + lhs = static_cast(lhs | rhs); + return lhs; +} +inline cpu_feature operator|(cpu_feature lhs, cpu_feature rhs) { + return static_cast(static_cast(lhs) | static_cast(rhs)); +} + +struct kernel_info { + size_t (*get_m_step)(void); + size_t (*get_n_step)(void); + size_t (*get_mr)(void); + size_t (*get_nr)(void); + size_t (*get_kr)(void); + size_t (*get_sr)(void); + + size_t (*get_dst_offset)(size_t m_idx, size_t n_idx, size_t stride); + size_t (*get_dst_size)(size_t m, size_t n); + + size_t (*get_lhs_offset_ex)(size_t m_idx, size_t k, size_t bl); + + size_t (*get_rhs_packed_offset_ex)(size_t n_idx, size_t k, size_t bl); + + void (*run_kernel_ex)( + size_t m, size_t n, size_t k, size_t bl, + const void* lhs_packed, const void* rhs_packed, + void* dst, size_t dst_stride_row, size_t dst_stride_col, + float clamp_min, float clamp_max); +}; + +struct lhs_packing_info { + size_t (*get_offset)(size_t m_idx, size_t lhs_stride); + + size_t (*get_packed_offset_ex)(size_t m_idx, size_t k, size_t bl, size_t mr, size_t kr, size_t sr); + + size_t (*packed_size_ex)(size_t m, size_t k, size_t bl, size_t mr, size_t kr, size_t sr); + + void (*pack_func_ex)(size_t m, size_t k, size_t bl, size_t mr, size_t kr, size_t sr, + size_t m_idx_start, const void * lhs, size_t lhs_stride, void * lhs_packed); +}; + +struct rhs_packing_info { + size_t (*packed_stride)(size_t k, size_t nr, size_t kr, size_t bl); + + void (*to_float)(const void *packed_data, int32_t row_idx, int64_t nc, float *out, + size_t nr_pack, size_t packed_row_stride, size_t kr, size_t bl, + size_t num_bytes_multiplier); + + size_t (*packed_size_ex)(size_t n, size_t k, size_t nr, size_t kr, size_t bl); + + size_t (*packed_stride_ex)(size_t k, size_t nr, size_t kr, size_t bl); + + void (*pack_func_ex)(size_t num_groups, size_t n, size_t k, size_t nr, size_t kr, size_t sr, size_t bl, + size_t rhs_stride, const void * rhs, const void * bias, const void * scale, void * rhs_packed, size_t extra_bytes, const void * params); +}; + +struct ggml_kleidiai_kernels { + kernel_info gemm; + lhs_packing_info gemm_lhs_info; + + kernel_info gemv; + lhs_packing_info gemv_lhs_info; + + rhs_packing_info rhs_info; + + cpu_feature required_cpu; + ggml_type lhs_type; + ggml_type rhs_type; + ggml_type op_type; +}; + +ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, const ggml_tensor * tensor); +ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q4_0(cpu_feature features); +ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q8_0(cpu_feature features); -#include "kernels.h" #define NELEMS(x) (sizeof(x) / sizeof(*x)) diff --git a/ggml/src/ggml-cpu/kleidiai/kernels.h b/ggml/src/ggml-cpu/kleidiai/kernels.h deleted file mode 100644 index 129245400b47..000000000000 --- a/ggml/src/ggml-cpu/kleidiai/kernels.h +++ /dev/null @@ -1,90 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Arm Limited and/or its affiliates -// SPDX-License-Identifier: MIT -// - -#pragma once - -#include "ggml.h" - -enum cpu_feature { - CPU_FEATURE_NONE = 0, - CPU_FEATURE_DOTPROD = 1, - CPU_FEATURE_I8MM = 2, - CPU_FEATURE_SVE = 4, - CPU_FEATURE_SME = 8 -}; - -inline cpu_feature& operator|=(cpu_feature& lhs, cpu_feature rhs) { - lhs = static_cast(lhs | rhs); - return lhs; -} -inline cpu_feature operator|(cpu_feature lhs, cpu_feature rhs) { - return static_cast(static_cast(lhs) | static_cast(rhs)); -} - -struct kernel_info { - size_t (*get_m_step)(void); - size_t (*get_n_step)(void); - size_t (*get_mr)(void); - size_t (*get_nr)(void); - size_t (*get_kr)(void); - size_t (*get_sr)(void); - - size_t (*get_dst_offset)(size_t m_idx, size_t n_idx, size_t stride); - size_t (*get_dst_size)(size_t m, size_t n); - - size_t (*get_lhs_offset_ex)(size_t m_idx, size_t k, size_t bl); - - size_t (*get_rhs_packed_offset_ex)(size_t n_idx, size_t k, size_t bl); - - void (*run_kernel_ex)( - size_t m, size_t n, size_t k, size_t bl, - const void* lhs_packed, const void* rhs_packed, - void* dst, size_t dst_stride_row, size_t dst_stride_col, - float clamp_min, float clamp_max); -}; - -struct lhs_packing_info { - size_t (*get_offset)(size_t m_idx, size_t lhs_stride); - - size_t (*get_packed_offset_ex)(size_t m_idx, size_t k, size_t bl, size_t mr, size_t kr, size_t sr); - - size_t (*packed_size_ex)(size_t m, size_t k, size_t bl, size_t mr, size_t kr, size_t sr); - - void (*pack_func_ex)(size_t m, size_t k, size_t bl, size_t mr, size_t kr, size_t sr, - size_t m_idx_start, const void * lhs, size_t lhs_stride, void * lhs_packed); -}; - -struct rhs_packing_info { - size_t (*packed_stride)(size_t k, size_t nr, size_t kr, size_t bl); - - void (*to_float)(const void *packed_data, int32_t row_idx, int64_t nc, float *out, - size_t nr_pack, size_t packed_row_stride, size_t kr, size_t bl, - size_t num_bytes_multiplier); - - size_t (*packed_size_ex)(size_t n, size_t k, size_t nr, size_t kr, size_t bl); - - size_t (*packed_stride_ex)(size_t k, size_t nr, size_t kr, size_t bl); - - void (*pack_func_ex)(size_t num_groups, size_t n, size_t k, size_t nr, size_t kr, size_t sr, size_t bl, - size_t rhs_stride, const void * rhs, const void * bias, const void * scale, void * rhs_packed, size_t extra_bytes, const void * params); -}; - -struct ggml_kleidiai_kernels { - kernel_info gemm; - lhs_packing_info gemm_lhs_info; - - kernel_info gemv; - lhs_packing_info gemv_lhs_info; - - rhs_packing_info rhs_info; - - cpu_feature required_cpu; - ggml_type lhs_type; - ggml_type rhs_type; - ggml_type op_type; -}; - -ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, const ggml_tensor * tensor); -ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q4_0(cpu_feature features); -ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q8_0(cpu_feature features); diff --git a/ggml/src/ggml-cpu/kleidiai/kleidiai.cpp b/ggml/src/ggml-cpu/kleidiai/kleidiai.cpp.inc similarity index 89% rename from ggml/src/ggml-cpu/kleidiai/kleidiai.cpp rename to ggml/src/ggml-cpu/kleidiai/kleidiai.cpp.inc index 0ecf7ae02ac7..38533fd18991 100644 --- a/ggml/src/ggml-cpu/kleidiai/kleidiai.cpp +++ b/ggml/src/ggml-cpu/kleidiai/kleidiai.cpp.inc @@ -35,20 +35,238 @@ #include #endif -#include "kleidiai.h" +#include "ggml-cpu.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml.h.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml-cpu-impl-defs.inc" +#include "ggml-impl-defs.inc" +#include "simd-mappings-defs.inc" + +#ifdef __cplusplus +#include +extern "C" { +#endif + +// return true if op part of extra "accelerator" +bool ggml_cpu_extra_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op); +bool ggml_cpu_extra_work_size(int n_threads, const struct ggml_tensor * op, size_t * size); + +#ifdef __cplusplus +} + +namespace ggml::cpu { +// register in tensor->extra +class tensor_traits { + public: + virtual ~tensor_traits(); + virtual bool work_size(int n_threads, const struct ggml_tensor * op, size_t & size) = 0; + virtual bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) = 0; +}; + +class extra_buffer_type { + public: + virtual ~extra_buffer_type(); + virtual bool supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) = 0; + virtual tensor_traits * get_tensor_traits(const struct ggml_tensor * op) = 0; +}; +} // namespace ggml::cpu + +// implemented in ggml-cpu.cpp. +std::vector & ggml_backend_cpu_get_extra_buffer_types(); + +#endif + +#define GGML_FA_TILE_Q 64 +#define GGML_FA_TILE_KV 64 + +#ifdef __cplusplus + +#include + +// convenience functions/macros for use in template calls +// note: these won't be required after the 'traits' lookup table is used. +static inline ggml_fp16_t f32_to_f16(float x) { + return GGML_CPU_FP32_TO_FP16(x); +} + +static inline float f16_to_f32(ggml_fp16_t x) { + return GGML_CPU_FP16_TO_FP32(x); +} + +static inline ggml_bf16_t f32_to_bf16(float x) { + return GGML_FP32_TO_BF16(x); +} + +static inline float bf16_to_f32(ggml_bf16_t x) { + return GGML_BF16_TO_FP32(x); +} + +static inline float i32_to_f32(int32_t x) { + return x; +} + +static inline int32_t f32_to_i32(float x) { + return x; +} + +static inline float f32_to_f32(float x) { + return x; +} + +// TODO - merge this into the traits table, after using row-based conversions +template +struct type_conversion_table; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_fp16_t) = f16_to_f32; + static constexpr ggml_fp16_t (*from_f32)(float) = f32_to_f16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(float) = f32_to_f32; + static constexpr float (*from_f32)(float) = f32_to_f32; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_bf16_t) = bf16_to_f32; + static constexpr ggml_bf16_t (*from_f32)(float) = f32_to_bf16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(int32_t) = i32_to_f32; + static constexpr int32_t (*from_f32)(float) = f32_to_i32; +}; + +static std::pair get_thread_range(const struct ggml_compute_params * params, const struct ggml_tensor * src0) { + const int64_t ith = params->ith; + const int64_t nth = params->nth; + + const int64_t nr = ggml_nrows(src0); + + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; -#include "ggml-cpu.h" -#include "ggml-impl.h" -#include "ggml-backend-impl.h" -#include "ggml-threading.h" -#include "traits.h" + // row range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); -#include "kernels.h" + return {ir0, ir1}; +} + +struct ggml_fa_tile_config { + static constexpr size_t Q = GGML_FA_TILE_Q; + static constexpr size_t KV = GGML_FA_TILE_KV; +}; + +#endif + + +// SPDX-FileCopyrightText: Copyright 2025 Arm Limited and/or its affiliates +// SPDX-License-Identifier: MIT +// + +#include "ggml.h.inc" + +enum cpu_feature { + CPU_FEATURE_NONE = 0, + CPU_FEATURE_DOTPROD = 1, + CPU_FEATURE_I8MM = 2, + CPU_FEATURE_SVE = 4, + CPU_FEATURE_SME = 8 +}; + +inline cpu_feature& operator|=(cpu_feature& lhs, cpu_feature rhs) { + lhs = static_cast(lhs | rhs); + return lhs; +} +inline cpu_feature operator|(cpu_feature lhs, cpu_feature rhs) { + return static_cast(static_cast(lhs) | static_cast(rhs)); +} + +struct kernel_info { + size_t (*get_m_step)(void); + size_t (*get_n_step)(void); + size_t (*get_mr)(void); + size_t (*get_nr)(void); + size_t (*get_kr)(void); + size_t (*get_sr)(void); + + size_t (*get_dst_offset)(size_t m_idx, size_t n_idx, size_t stride); + size_t (*get_dst_size)(size_t m, size_t n); + + size_t (*get_lhs_offset_ex)(size_t m_idx, size_t k, size_t bl); + + size_t (*get_rhs_packed_offset_ex)(size_t n_idx, size_t k, size_t bl); + + void (*run_kernel_ex)( + size_t m, size_t n, size_t k, size_t bl, + const void* lhs_packed, const void* rhs_packed, + void* dst, size_t dst_stride_row, size_t dst_stride_col, + float clamp_min, float clamp_max); +}; -#include "kai_common.h" +struct lhs_packing_info { + size_t (*get_offset)(size_t m_idx, size_t lhs_stride); + + size_t (*get_packed_offset_ex)(size_t m_idx, size_t k, size_t bl, size_t mr, size_t kr, size_t sr); + + size_t (*packed_size_ex)(size_t m, size_t k, size_t bl, size_t mr, size_t kr, size_t sr); + + void (*pack_func_ex)(size_t m, size_t k, size_t bl, size_t mr, size_t kr, size_t sr, + size_t m_idx_start, const void * lhs, size_t lhs_stride, void * lhs_packed); +}; + +struct rhs_packing_info { + size_t (*packed_stride)(size_t k, size_t nr, size_t kr, size_t bl); + + void (*to_float)(const void *packed_data, int32_t row_idx, int64_t nc, float *out, + size_t nr_pack, size_t packed_row_stride, size_t kr, size_t bl, + size_t num_bytes_multiplier); + + size_t (*packed_size_ex)(size_t n, size_t k, size_t nr, size_t kr, size_t bl); + + size_t (*packed_stride_ex)(size_t k, size_t nr, size_t kr, size_t bl); + + void (*pack_func_ex)(size_t num_groups, size_t n, size_t k, size_t nr, size_t kr, size_t sr, size_t bl, + size_t rhs_stride, const void * rhs, const void * bias, const void * scale, void * rhs_packed, size_t extra_bytes, const void * params); +}; + +struct ggml_kleidiai_kernels { + kernel_info gemm; + lhs_packing_info gemm_lhs_info; + + kernel_info gemv; + lhs_packing_info gemv_lhs_info; + + rhs_packing_info rhs_info; + + cpu_feature required_cpu; + ggml_type lhs_type; + ggml_type rhs_type; + ggml_type op_type; +}; + +ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, const ggml_tensor * tensor); +ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q4_0(cpu_feature features); +ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q8_0(cpu_feature features); + + +#include "kai_common.h.inc" #define GGML_COMMON_DECL_CPP -#include "ggml-common.h" +#include "ggml-common-defs.inc" + +extern "C" ggml_backend_buffer_type_t ggml_backend_cpu_kleidiai_buffer_type(void); +extern "C" { +GGML_API void ggml_critical_section_start(void); +GGML_API void ggml_critical_section_end(void); +} static constexpr int GGML_KLEIDIAI_MAX_KERNEL_SLOTS = 2; static constexpr uint32_t GGML_KLEIDIAI_PACK_MAGIC = 0x4b4c4149; // "KLAI" diff --git a/ggml/src/ggml-cpu/kleidiai/kleidiai.h b/ggml/src/ggml-cpu/kleidiai/kleidiai.h deleted file mode 100644 index 38eac58f7c20..000000000000 --- a/ggml/src/ggml-cpu/kleidiai/kleidiai.h +++ /dev/null @@ -1,17 +0,0 @@ -// SPDX-FileCopyrightText: Copyright 2025 Arm Limited and/or its affiliates -// SPDX-License-Identifier: MIT -// - -#pragma once - -#include "ggml-alloc.h" - -#ifdef __cplusplus -extern "C" { -#endif - -ggml_backend_buffer_type_t ggml_backend_cpu_kleidiai_buffer_type(void); - -#ifdef __cplusplus -} -#endif diff --git a/ggml/src/ggml-cpu/llamafile/sgemm.cpp b/ggml/src/ggml-cpu/llamafile/sgemm.cpp deleted file mode 100644 index e13828e3be6f..000000000000 --- a/ggml/src/ggml-cpu/llamafile/sgemm.cpp +++ /dev/null @@ -1,4051 +0,0 @@ -// Copyright 2024 Mozilla Foundation -// -// Permission is hereby granted, free of charge, to any person obtaining -// a copy of this software and associated documentation files (the -// "Software"), to deal in the Software without restriction, including -// without limitation the rights to use, copy, modify, merge, publish, -// distribute, sublicense, and/or sell copies of the Software, and to -// permit persons to whom the Software is furnished to do so, subject to -// the following conditions: -// -// The above copyright notice and this permission notice shall be -// included in all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, -// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF -// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND -// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS -// BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN -// ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN -// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -// -// _ _ ___ _ _ ___ -// | |_(_)_ _ _ _| _ ) | /_\ / __| -// | _| | ' \ || | _ \ |__ / _ \\__ \. -// \__|_|_||_\_, |___/____/_/ \_\___/ -// |__/ -// -// BASIC LINEAR ALGEBRA SUBPROGRAMS -// -// -// This file implements multithreaded CPU matrix multiplication for the -// common contiguous use case C = Aᵀ * B. These kernels are designed to -// have excellent performance[1] for matrices that fit in the CPU cache -// without imposing any overhead such as cache filling or malloc calls. -// -// This implementation does not guarantee any upper bound with rounding -// errors, which grow along with k. Our goal's to maximally exploit the -// hardware for performance, and then use whatever resources remain for -// improving numerical accuracy. -// -// [1] J. Tunney, ‘LLaMA Now Goes Faster on CPUs’, Mar. 2024. [Online]. -// Available: https://justine.lol/matmul/. [Accessed: 29-Mar-2024]. - -#if defined(__GNUC__) -#pragma GCC diagnostic ignored "-Wpedantic" -#pragma GCC diagnostic ignored "-Wignored-attributes" -#endif - -#include "sgemm.h" -#include "ggml-impl.h" -#include "ggml-cpu-impl.h" -#include "ggml-quants.h" -#include "simd-mappings.h" - -#include -#include - -#ifdef _MSC_VER -#define NOINLINE __declspec(noinline) -#else -#define NOINLINE __attribute__((__noinline__)) -#endif - -#if defined(__ARM_NEON) || defined(__AVX512F__) || defined(__VXE__) || defined(__VXE2__) -#define VECTOR_REGISTERS 32 -#else -#define VECTOR_REGISTERS 16 -#endif - -#if defined(__riscv_v_intrinsic) -#define LMUL 4 -#endif - -#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1) - -namespace { - -inline float unhalf(ggml_fp16_t d) { - return GGML_CPU_FP16_TO_FP32(d); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// VECTORIZED ARITHMETIC OPERATIONS - -#if defined(__SSE__) || defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) -inline __m128 add(__m128 x, __m128 y) { return _mm_add_ps(x, y); } -inline __m128 sub(__m128 x, __m128 y) { return _mm_sub_ps(x, y); } -inline __m128 mul(__m128 x, __m128 y) { return _mm_mul_ps(x, y); } -#endif // __SSE__ - -#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) -inline __m256 add(__m256 x, __m256 y) { return _mm256_add_ps(x, y); } -inline __m256 sub(__m256 x, __m256 y) { return _mm256_sub_ps(x, y); } -inline __m256 mul(__m256 x, __m256 y) { return _mm256_mul_ps(x, y); } -#endif // __AVX__ - -#if defined(__AVX512F__) -inline __m512 add(__m512 x, __m512 y) { return _mm512_add_ps(x, y); } -inline __m512 sub(__m512 x, __m512 y) { return _mm512_sub_ps(x, y); } -inline __m512 mul(__m512 x, __m512 y) { return _mm512_mul_ps(x, y); } -#endif // __AVX512F__ - -#if defined(__ARM_NEON) -inline float32x4_t add(float32x4_t x, float32x4_t y) { return vaddq_f32(x, y); } -inline float32x4_t sub(float32x4_t x, float32x4_t y) { return vsubq_f32(x, y); } -inline float32x4_t mul(float32x4_t x, float32x4_t y) { return vmulq_f32(x, y); } -#endif // __ARM_NEON - -#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) -inline float16x8_t add(float16x8_t x, float16x8_t y) { return vaddq_f16(x, y); } -inline float16x8_t sub(float16x8_t x, float16x8_t y) { return vsubq_f16(x, y); } -inline float16x8_t mul(float16x8_t x, float16x8_t y) { return vmulq_f16(x, y); } -#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC - -#if defined(__VXE__) || defined(__VXE2__) -inline float32x4_t add(float32x4_t x, float32x4_t y) { return vec_add(x, y); } -inline float32x4_t sub(float32x4_t x, float32x4_t y) { return vec_sub(x, y); } -inline float32x4_t mul(float32x4_t x, float32x4_t y) { return vec_mul(x, y); } -#endif - -#if defined(__MMA__) -typedef vector unsigned char vec_t; -typedef __vector_quad acc_t; -#endif -//////////////////////////////////////////////////////////////////////////////////////////////////// -// VECTORIZED FUSED MULTIPLY ADD - -/** - * Computes a * b + c. - */ -template -inline U madd(T a, T b, U c) { - return add(mul(a, b), c); -} - -#if defined(__FMA__) -#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) -template <> -inline __m256 madd(__m256 a, __m256 b, __m256 c) { - return _mm256_fmadd_ps(a, b, c); -} -#endif -#if defined(__AVX512F__) -template <> -inline __m512 madd(__m512 a, __m512 b, __m512 c) { - return _mm512_fmadd_ps(a, b, c); -} -#endif -#if defined(__AVX512BF16__) -template <> -inline __m512 madd(__m512bh a, __m512bh b, __m512 c) { - return _mm512_dpbf16_ps(c, a, b); -} -template <> -inline __m256 madd(__m256bh a, __m256bh b, __m256 c) { - return _mm256_dpbf16_ps(c, a, b); -} -#endif -#endif - -#if defined(__ARM_FEATURE_FMA) -template <> -inline float32x4_t madd(float32x4_t a, float32x4_t b, float32x4_t c) { - return vfmaq_f32(c, b, a); -} -#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER) -template <> -inline float16x8_t madd(float16x8_t a, float16x8_t b, float16x8_t c) { - return vfmaq_f16(c, b, a); -} -#endif -#endif - -#if defined(__VXE__) || defined(__VXE2__) -template <> -inline float32x4_t madd(float32x4_t a, float32x4_t b, float32x4_t c) { - return vec_madd(a, b, c); -} -#endif - -#if defined(__riscv_v_intrinsic) -template <> inline vfloat32m1_t madd(vfloat32m1_t a, vfloat32m1_t b, vfloat32m1_t c) { - return __riscv_vfmacc_vv_f32m1(c, a, b, __riscv_vsetvlmax_e32m1()); -} -template <> inline vfloat32m2_t madd(vfloat32m2_t a, vfloat32m2_t b, vfloat32m2_t c) { - return __riscv_vfmacc_vv_f32m2(c, a, b, __riscv_vsetvlmax_e32m2()); -} -template <> inline vfloat32m4_t madd(vfloat32m4_t a, vfloat32m4_t b, vfloat32m4_t c) { - return __riscv_vfmacc_vv_f32m4(c, a, b, __riscv_vsetvlmax_e32m4()); -} -template <> inline vfloat32m8_t madd(vfloat32m8_t a, vfloat32m8_t b, vfloat32m8_t c) { - return __riscv_vfmacc_vv_f32m8(c, a, b, __riscv_vsetvlmax_e32m8()); -} -#endif - -#if defined(__riscv_zvfh) -template <> inline vfloat32m1_t madd(vfloat16mf2_t a, vfloat16mf2_t b, vfloat32m1_t c) { - return __riscv_vfwmacc_vv_f32m1(c, a, b, __riscv_vsetvlmax_e32m1()); -} -template <> inline vfloat32m2_t madd(vfloat16m1_t a, vfloat16m1_t b, vfloat32m2_t c) { - return __riscv_vfwmacc_vv_f32m2(c, a, b, __riscv_vsetvlmax_e32m2()); -} -template <> inline vfloat32m4_t madd(vfloat16m2_t a, vfloat16m2_t b, vfloat32m4_t c) { - return __riscv_vfwmacc_vv_f32m4(c, a, b, __riscv_vsetvlmax_e32m4()); -} -template <> inline vfloat32m8_t madd(vfloat16m4_t a, vfloat16m4_t b, vfloat32m8_t c) { - return __riscv_vfwmacc_vv_f32m8(c, a, b, __riscv_vsetvlmax_e32m8()); -} -#endif - -#if defined(__riscv_zvfbfwma) -template <> inline vfloat32m1_t madd(vbfloat16mf2_t a, vbfloat16mf2_t b, vfloat32m1_t c) { - return __riscv_vfwmaccbf16_vv_f32m1(c, a, b, __riscv_vsetvlmax_e32m1()); -} -template <> inline vfloat32m2_t madd(vbfloat16m1_t a, vbfloat16m1_t b, vfloat32m2_t c) { - return __riscv_vfwmaccbf16_vv_f32m2(c, a, b, __riscv_vsetvlmax_e32m2()); -} -template <> inline vfloat32m4_t madd(vbfloat16m2_t a, vbfloat16m2_t b, vfloat32m4_t c) { - return __riscv_vfwmaccbf16_vv_f32m4(c, a, b, __riscv_vsetvlmax_e32m4()); -} -template <> inline vfloat32m8_t madd(vbfloat16m4_t a, vbfloat16m4_t b, vfloat32m8_t c) { - return __riscv_vfwmaccbf16_vv_f32m8(c, a, b, __riscv_vsetvlmax_e32m8()); -} -#endif - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// VECTORIZED HORIZONTAL SUM - -#if defined(__ARM_NEON) -inline float hsum(float32x4_t x) { - return vaddvq_f32(x); -} -#endif // __ARM_NEON - -#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER) -inline float hsum(float16x8_t x) { - return vaddvq_f32(vaddq_f32(vcvt_f32_f16(vget_low_f16(x)), - vcvt_f32_f16(vget_high_f16(x)))); -} -#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC - -#if defined(__VXE__) || defined(__VXE2__) -inline float hsum(float32x4_t x) { - float32x4_t tmp = x + vec_reve(x); - return tmp[0] + tmp[1]; -} -#endif - -#if defined(__SSE__) || defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) -inline float hsum(__m128 x) { -#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) - x = _mm_add_ps(x, _mm_movehl_ps(x, x)); - x = _mm_add_ss(x, _mm_movehdup_ps(x)); -#else - __m128 t; - t = _mm_shuffle_ps(x, x, _MM_SHUFFLE(2, 3, 0, 1)); - x = _mm_add_ps(x, t); - t = _mm_movehl_ps(t, x); - x = _mm_add_ss(x, t); -#endif - return _mm_cvtss_f32(x); -} -#endif - -#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) -inline float hsum(__m256 x) { - return hsum(_mm_add_ps(_mm256_extractf128_ps(x, 1), - _mm256_castps256_ps128(x))); -} -#endif // __AVX__ - -#if defined(__AVX512F__) -inline float hsum(__m512 x) { - return _mm512_reduce_add_ps(x); -} -#endif // __AVX512F__ - -#if defined(__riscv_v_intrinsic) -inline float hsum(vfloat32m1_t x) { - return __riscv_vfmv_f_s_f32m1_f32( - __riscv_vfredusum_vs_f32m1_f32m1(x, __riscv_vfmv_v_f_f32m1(0, 1), __riscv_vsetvlmax_e32m1())); -} -inline float hsum(vfloat32m2_t x) { - return __riscv_vfmv_f_s_f32m1_f32( - __riscv_vfredusum_vs_f32m2_f32m1(x, __riscv_vfmv_v_f_f32m1(0, 1), __riscv_vsetvlmax_e32m2())); -} -inline float hsum(vfloat32m4_t x) { - return __riscv_vfmv_f_s_f32m1_f32( - __riscv_vfredusum_vs_f32m4_f32m1(x, __riscv_vfmv_v_f_f32m1(0, 1), __riscv_vsetvlmax_e32m4())); -} -inline float hsum(vfloat32m8_t x) { - return __riscv_vfmv_f_s_f32m1_f32( - __riscv_vfredusum_vs_f32m8_f32m1(x, __riscv_vfmv_v_f_f32m1(0, 1), __riscv_vsetvlmax_e32m8())); -} -#endif - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// VECTORIZED MEMORY LOADING - -template T load(const U *); - -#if defined(__ARM_NEON) -template <> inline float32x4_t load(const float *p) { - return vld1q_f32(p); -} -#if !defined(_MSC_VER) -// FIXME: this should check for __ARM_FEATURE_FP16_VECTOR_ARITHMETIC -template <> inline float16x8_t load(const ggml_fp16_t *p) { - return vld1q_f16((const float16_t *)p); -} -template <> inline float32x4_t load(const ggml_fp16_t *p) { - return vcvt_f32_f16(vld1_f16((const float16_t *)p)); -} -#endif // _MSC_VER -#endif // __ARM_NEON - -#if defined(__VXE__) || defined(__VXE2__) -template <> inline float32x4_t load(const ggml_fp16_t * p) { - float tmp[4]; - - for (int i = 0; i < 4; i++) { - tmp[i] = GGML_CPU_FP16_TO_FP32(p[i]); - } - - return vec_xl(0, (const float *)(tmp)); -} -template <> inline float32x4_t load(const float * p) { - return vec_xl(0, p); -} -#endif - -#if defined(__SSE__) || defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) -template <> inline __m128 load(const float *p) { - return _mm_loadu_ps(p); -} -#endif // __SSE__ - -#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) -template <> inline __m256 load(const float *p) { - return _mm256_loadu_ps(p); -} -#endif // __AVX__ - -#if defined(__AVX2__) || defined(__AVX512F__) -template <> inline __m256 load(const ggml_bf16_t *p) { - return _mm256_castsi256_ps( - _mm256_slli_epi32(_mm256_cvtepu16_epi32(_mm_loadu_si128((const __m128i *)p)), 16)); -} -#endif // __AVX2__ - -#if defined(__F16C__) -template <> inline __m256 load(const ggml_fp16_t *p) { - return _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)p)); -} -#endif // __F16C__ - -#if defined(__AVX512F__) -template <> inline __m512 load(const float *p) { - return _mm512_loadu_ps(p); -} -template <> inline __m512 load(const ggml_fp16_t *p) { - return _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)p)); -} -template <> inline __m512 load(const ggml_bf16_t *p) { - return _mm512_castsi512_ps( - _mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)p)), 16)); -} -#endif // __AVX512F__ - -#if defined(__AVX512BF16__) -template <> inline __m512bh load(const ggml_bf16_t *p) { - return (__m512bh)_mm512_loadu_ps((const float *)p); -} -template <> inline __m256bh load(const ggml_bf16_t *p) { - return (__m256bh)_mm256_loadu_ps((const float *)p); -} -template <> inline __m512bh load(const float *p) { - return _mm512_cvtne2ps_pbh(_mm512_loadu_ps(p + 16), _mm512_loadu_ps(p)); -} -template <> inline __m256bh load(const float *p) { - return _mm512_cvtneps_pbh(_mm512_loadu_ps(p)); -} -#endif - -#if defined(__riscv_v_intrinsic) -template <> inline vfloat32m1_t load(const float *p) { - return __riscv_vle32_v_f32m1(p, __riscv_vsetvlmax_e32m1()); -} -template <> inline vfloat32m2_t load(const float *p) { - return __riscv_vle32_v_f32m2(p, __riscv_vsetvlmax_e32m2()); -} -template <> inline vfloat32m4_t load(const float *p) { - return __riscv_vle32_v_f32m4(p, __riscv_vsetvlmax_e32m4()); -} -template <> inline vfloat32m8_t load(const float *p) { - return __riscv_vle32_v_f32m8(p, __riscv_vsetvlmax_e32m8()); -} -#endif - -#if defined(__riscv_zvfh) -template <> inline vfloat16mf2_t load(const ggml_fp16_t *p) { - return __riscv_vle16_v_f16mf2(reinterpret_cast(p), __riscv_vsetvlmax_e16mf2()); -} -template <> inline vfloat16m1_t load(const ggml_fp16_t *p) { - return __riscv_vle16_v_f16m1(reinterpret_cast(p), __riscv_vsetvlmax_e16m1()); -} -template <> inline vfloat16m2_t load(const ggml_fp16_t *p) { - return __riscv_vle16_v_f16m2(reinterpret_cast(p), __riscv_vsetvlmax_e16m2()); -} -template <> inline vfloat16m4_t load(const ggml_fp16_t *p) { - return __riscv_vle16_v_f16m4(reinterpret_cast(p), __riscv_vsetvlmax_e16m4()); -} -#endif - -#if defined(__riscv_zvfbfwma) -template <> inline vbfloat16mf2_t load(const ggml_bf16_t *p) { - return __riscv_vle16_v_bf16mf2(reinterpret_cast(p), __riscv_vsetvlmax_e16mf2()); -} -template <> inline vbfloat16m1_t load(const ggml_bf16_t *p) { - return __riscv_vle16_v_bf16m1(reinterpret_cast(p), __riscv_vsetvlmax_e16m1()); -} -template <> inline vbfloat16m2_t load(const ggml_bf16_t *p) { - return __riscv_vle16_v_bf16m2(reinterpret_cast(p), __riscv_vsetvlmax_e16m2()); -} -template <> inline vbfloat16m4_t load(const ggml_bf16_t *p) { - return __riscv_vle16_v_bf16m4(reinterpret_cast(p), __riscv_vsetvlmax_e16m4()); -} -#endif - -#if defined(__riscv_v_intrinsic) -template T set_zero(); - -template <> inline vfloat32m1_t set_zero() { - return __riscv_vfmv_v_f_f32m1(0.0f, __riscv_vsetvlmax_e32m1()); -} -template <> inline vfloat32m2_t set_zero() { - return __riscv_vfmv_v_f_f32m2(0, __riscv_vsetvlmax_e32m2()); -} -template <> inline vfloat32m4_t set_zero() { - return __riscv_vfmv_v_f_f32m4(0, __riscv_vsetvlmax_e32m4()); -} -template <> inline vfloat32m8_t set_zero() { - return __riscv_vfmv_v_f_f32m8(0, __riscv_vsetvlmax_e32m8()); -} -#endif - -#if defined(__riscv_v_intrinsic) -template size_t vlmax() { - if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e32m1(); } - else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e32m2(); } - else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e32m4(); } - else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e32m8(); } - #if defined (__riscv_zvfh) - else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16mf2(); } - else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16m1(); } - else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16m2(); } - else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16m4(); } - #endif - #if defined (__riscv_zvfbfwma) - else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16mf2(); } - else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16m1(); } - else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16m2(); } - else if constexpr (std::is_same_v) { return __riscv_vsetvlmax_e16m4(); } - #endif - return 0; -} -#endif - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// FLOATING POINT MATRIX MULTIPLICATION - -template -static inline int64_t BLOCK_SIZE(size_t m) { - const int64_t NB_BLOC_M = (m + M - 1) / M; - return (m % NB_BLOC_M == 0) ? m / NB_BLOC_M : (m / NB_BLOC_M) + 1; -} - -static constexpr inline int64_t BLOC_POS(int64_t ib, int64_t ibN, int64_t bloc_size) { - return ib < ibN ? ib * bloc_size : ibN * bloc_size + (ib - ibN) * (bloc_size - 1); -} - -template -class tinyBLAS { - public: - tinyBLAS(const ggml_compute_params * params, int64_t k, - const TA *A, int64_t lda, - const TB *B, int64_t ldb, - TC *C, int64_t ldc) - : params(params), A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc) { - } - - bool matmul(int64_t m, int64_t n) { - if (k % KN != 0) - return false; - // compute RM for only need tile with size RM&RM-1 -#if VECTOR_REGISTERS == 32 - if (m % 16 == 0 && (m/16 >= params->nth)) { - const int64_t SIZE_N = BLOCK_SIZE<6>(n); - mnpack<4, 6, 4>(m, n, SIZE_N, 12); - return true; - } - if (m % 8 == 0 ) { - const int64_t SIZE_N = BLOCK_SIZE<6>(n); - mnpack<4, 6, 2>(m, n, SIZE_N, 12); - return true; - } - if (m % 4 == 0) { - const int64_t SIZE_N = BLOCK_SIZE<6>(n); - mnpack<4, 6, 1>(m, n, SIZE_N, 12); - return true; - } -#else // VECTOR_REGISTERS == 16 - if (m % 16 == 0 && (m/16 >= params->nth)) { - const int64_t SIZE_N = BLOCK_SIZE<3>(n); - mnpack<4, 3, 4>(m, n, SIZE_N, 24); - return true; - } - if (m % 8 == 0 ) { - const int64_t SIZE_N = BLOCK_SIZE<3>(n); - mnpack<4, 3, 2>(m, n, SIZE_N, 24); - return true; - } - if (m % 4 == 0) { - const int64_t SIZE_N = BLOCK_SIZE<3>(n); - mnpack<4, 3, 1>(m, n, SIZE_N, 24); - return true; - } -#endif - return false; - } - - private: - template - inline void mnpack(int64_t m, int64_t n, int64_t SIZE_N, int64_t BN) { - if (SIZE_N == RN) { - return gemm(m, n, BN); - } - if constexpr (RN > 1) { - return mnpack(m, n, SIZE_N, BN); - } else { - GGML_LOG_ERROR("mnpack<%d, %d> block size not supported\n", RM, (int)SIZE_N); - GGML_ASSERT(false); // we have miss something. - } - } - - template - inline void gemm_bloc(int64_t ii, int64_t jj) { - D Cv[RN][RM] = {}; - for (int64_t l = 0; l < k; l += KN) { - // help compiler for op order. - if constexpr (RM <= RN) { - V Av[RM]; - for (int64_t i = 0; i < RM; ++i) { - Av[i] = load(A + lda * (ii + i) + l); - } - for (int64_t j = 0; j < RN; ++j) { - V Bv = load(B + ldb * (jj + j) + l); - for (int64_t i = 0; i < RM; ++i) { - Cv[j][i] = madd(Av[i], Bv, Cv[j][i]); - } - } - } else { - V Bv[RN]; - for (int64_t j = 0; j < RN; ++j) { - Bv[j] = load(B + ldb * (jj + j) + l); - } - for (int64_t i = 0; i < RM; ++i) { - V Av = load(A + lda * (ii + i) + l); - for (int64_t j = 0; j < RN; ++j) { - Cv[j][i] = madd(Av, Bv[j], Cv[j][i]); - } - } - } - } - for (int64_t j = 0; j < RN; ++j) - for (int64_t i = 0; i < RM; ++i) - C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]); - } - - template - NOINLINE void gemm(int64_t m, int64_t n, int64_t BN) { - GGML_ASSERT(m % (RM * BM) == 0); - const int64_t ytiles = m / (RM * BM); - const int64_t xtiles = (n + RN -1) / RN; - const int64_t jj_RN = (xtiles - (xtiles * RN - n)); - - // "round" bloc_size to "nearest" BN - const int64_t NB_BN = xtiles < BN ? 1 : (xtiles + BN / 2) / BN; - const int64_t SIZE_BN = xtiles % NB_BN == 0 ? xtiles / NB_BN : xtiles / NB_BN + 1; - const int64_t jj_BN = (NB_BN - (NB_BN * SIZE_BN - xtiles)); - const int64_t nb_job = ytiles * NB_BN; - - if (params->ith == 0) { - GGML_ASSERT( jj_BN * SIZE_BN + (NB_BN - jj_BN) * (SIZE_BN - 1) == xtiles); - // Every thread starts at ith, so the first unprocessed chunk is nth. This save a bit of coordination right at the start. - ggml_threadpool_chunk_set(params->threadpool, params->nth); - } - - ggml_barrier(params->threadpool); - - int64_t job = params->ith; - while (job < nb_job) { - const int64_t ii = (job % ytiles) * RM * BM; - const int64_t jb = job / ytiles; - const int64_t jr0 = BLOC_POS(jb , jj_BN, SIZE_BN); - const int64_t jrN = BLOC_POS(jb+1, jj_BN, SIZE_BN); - - const int64_t jj0 = BLOC_POS(jr0, jj_RN, RN); - const int64_t jj2 = BLOC_POS(jrN, jj_RN, RN); - const int64_t jj1 = jj2 < jj_RN * RN ? jj2 : jj_RN * RN; - - for (int64_t bi = 0; bi < BM * RM; bi += RM) { - int64_t jj = jj0; - for (; jj < jj1; jj += RN) { - gemm_bloc(ii + bi, jj); - } - if constexpr (RN > 1) { - for (; jj < jj2; jj += RN - 1) { - gemm_bloc(ii + bi, jj); - } - } - GGML_ASSERT(jj == jj2); - } - - job = ggml_threadpool_chunk_add(params->threadpool, 1); - } - - ggml_barrier(params->threadpool); - return; - } - - const ggml_compute_params * params; - const TA *const A; - const TB *const B; - TC *const C; - const int64_t k; - const int64_t lda; - const int64_t ldb; - const int64_t ldc; -}; - -#if defined(__riscv_v_intrinsic) -template -class tinyBLAS_RVV { - public: - tinyBLAS_RVV(const ggml_compute_params * params, int64_t k, - const TA *A, int64_t lda, - const TB *B, int64_t ldb, - TC *C, int64_t ldc) - : params(params), A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc) { - } - - bool matmul(int64_t m, int64_t n) { - if (k % vlmax() != 0) { - return false; - } - -#if LMUL == 1 - if (m % 16 == 0 && (m/16 >= params->nth)) { - const int64_t SIZE_N = BLOCK_SIZE<6>(n); - mnpack<4, 6, 4>(m, n, SIZE_N, 12); - return true; - } - if (m % 8 == 0 ) { - const int64_t SIZE_N = BLOCK_SIZE<6>(n); - mnpack<4, 6, 2>(m, n, SIZE_N, 12); - return true; - } - if (m % 4 == 0) { - const int64_t SIZE_N = BLOCK_SIZE<6>(n); - mnpack<4, 6, 1>(m, n, SIZE_N, 12); - return true; - } -#elif LMUL == 2 - if (m % 16 == 0 && (m/16 >= params->nth)) { - const int64_t SIZE_N = BLOCK_SIZE<3>(n); - mnpack<4, 3, 4>(m, n, SIZE_N, 24); - return true; - } - if (m % 8 == 0 ) { - const int64_t SIZE_N = BLOCK_SIZE<3>(n); - mnpack<4, 3, 2>(m, n, SIZE_N, 24); - return true; - } - if (m % 4 == 0) { - const int64_t SIZE_N = BLOCK_SIZE<3>(n); - mnpack<4, 3, 1>(m, n, SIZE_N, 24); - return true; - } -#else // LMUL = 4 - if (m % 16 == 0 && (m/16 >= params->nth)) { - const int64_t SIZE_N = BLOCK_SIZE<2>(n); - mnpack<2, 2, 8>(m, n, SIZE_N, 36); - return true; - } - if (m % 8 == 0 ) { - const int64_t SIZE_N = BLOCK_SIZE<2>(n); - mnpack<2, 2, 4>(m, n, SIZE_N, 36); - return true; - } - if (m % 4 == 0) { - const int64_t SIZE_N = BLOCK_SIZE<2>(n); - mnpack<2, 2, 2>(m, n, SIZE_N, 36); - return true; - } -#endif - return false; - } - - private: - template - inline void mnpack(int64_t m, int64_t n, int64_t SIZE_N, int64_t BN) { - if (SIZE_N == RN) { - return gemm(m, n, BN); - } - if constexpr (RN > 1) { - return mnpack(m, n, SIZE_N, BN); - } else { - GGML_LOG_ERROR("mnpack<%d, %d> block size not supported\n", RM, (int)SIZE_N); - GGML_ASSERT(false); // we have miss something. - } - } - - inline void gemm_bloc_4x6(int64_t ii, int64_t jj) { - size_t vl = vlmax(); - D Cv00 = set_zero(); - D Cv01 = set_zero(); - D Cv02 = set_zero(); - D Cv03 = set_zero(); - D Cv10 = set_zero(); - D Cv11 = set_zero(); - D Cv12 = set_zero(); - D Cv13 = set_zero(); - D Cv20 = set_zero(); - D Cv21 = set_zero(); - D Cv22 = set_zero(); - D Cv23 = set_zero(); - D Cv30 = set_zero(); - D Cv31 = set_zero(); - D Cv32 = set_zero(); - D Cv33 = set_zero(); - D Cv40 = set_zero(); - D Cv41 = set_zero(); - D Cv42 = set_zero(); - D Cv43 = set_zero(); - D Cv50 = set_zero(); - D Cv51 = set_zero(); - D Cv52 = set_zero(); - D Cv53 = set_zero(); - - for (int64_t l = 0; l < k; l += vl) { - V Bv0 = load(B + ldb * (jj + 0) + l); - V Bv1 = load(B + ldb * (jj + 1) + l); - V Bv2 = load(B + ldb * (jj + 2) + l); - V Bv3 = load(B + ldb * (jj + 3) + l); - V Bv4 = load(B + ldb * (jj + 4) + l); - V Bv5 = load(B + ldb * (jj + 5) + l); - - V Av0 = load(A + lda * (ii + 0) + l); - Cv00 = madd(Av0, Bv0, Cv00); - Cv10 = madd(Av0, Bv1, Cv10); - Cv20 = madd(Av0, Bv2, Cv20); - Cv30 = madd(Av0, Bv3, Cv30); - Cv40 = madd(Av0, Bv4, Cv40); - Cv50 = madd(Av0, Bv5, Cv50); - - V Av1 = load(A + lda * (ii + 1) + l); - Cv01 = madd(Av1, Bv0, Cv01); - Cv11 = madd(Av1, Bv1, Cv11); - Cv21 = madd(Av1, Bv2, Cv21); - Cv31 = madd(Av1, Bv3, Cv31); - Cv41 = madd(Av1, Bv4, Cv41); - Cv51 = madd(Av1, Bv5, Cv51); - - V Av2 = load(A + lda * (ii + 2) + l); - Cv02 = madd(Av2, Bv0, Cv02); - Cv12 = madd(Av2, Bv1, Cv12); - Cv22 = madd(Av2, Bv2, Cv22); - Cv32 = madd(Av2, Bv3, Cv32); - Cv42 = madd(Av2, Bv4, Cv42); - Cv52 = madd(Av2, Bv5, Cv52); - - V Av3 = load(A + lda * (ii + 3) + l); - Cv03 = madd(Av3, Bv0, Cv03); - Cv13 = madd(Av3, Bv1, Cv13); - Cv23 = madd(Av3, Bv2, Cv23); - Cv33 = madd(Av3, Bv3, Cv33); - Cv43 = madd(Av3, Bv4, Cv43); - Cv53 = madd(Av3, Bv5, Cv53); - } - - C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); - C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); - C[ldc * (jj + 0) + (ii + 2)] = hsum(Cv02); - C[ldc * (jj + 0) + (ii + 3)] = hsum(Cv03); - C[ldc * (jj + 1) + (ii + 0)] = hsum(Cv10); - C[ldc * (jj + 1) + (ii + 1)] = hsum(Cv11); - C[ldc * (jj + 1) + (ii + 2)] = hsum(Cv12); - C[ldc * (jj + 1) + (ii + 3)] = hsum(Cv13); - C[ldc * (jj + 2) + (ii + 0)] = hsum(Cv20); - C[ldc * (jj + 2) + (ii + 1)] = hsum(Cv21); - C[ldc * (jj + 2) + (ii + 2)] = hsum(Cv22); - C[ldc * (jj + 2) + (ii + 3)] = hsum(Cv23); - C[ldc * (jj + 3) + (ii + 0)] = hsum(Cv30); - C[ldc * (jj + 3) + (ii + 1)] = hsum(Cv31); - C[ldc * (jj + 3) + (ii + 2)] = hsum(Cv32); - C[ldc * (jj + 3) + (ii + 3)] = hsum(Cv33); - C[ldc * (jj + 4) + (ii + 0)] = hsum(Cv40); - C[ldc * (jj + 4) + (ii + 1)] = hsum(Cv41); - C[ldc * (jj + 4) + (ii + 2)] = hsum(Cv42); - C[ldc * (jj + 4) + (ii + 3)] = hsum(Cv43); - C[ldc * (jj + 5) + (ii + 0)] = hsum(Cv50); - C[ldc * (jj + 5) + (ii + 1)] = hsum(Cv51); - C[ldc * (jj + 5) + (ii + 2)] = hsum(Cv52); - C[ldc * (jj + 5) + (ii + 3)] = hsum(Cv53); - } - - inline void gemm_bloc_4x5(int64_t ii, int64_t jj) { - size_t vl = vlmax(); - D Cv00 = set_zero(); - D Cv01 = set_zero(); - D Cv02 = set_zero(); - D Cv03 = set_zero(); - D Cv10 = set_zero(); - D Cv11 = set_zero(); - D Cv12 = set_zero(); - D Cv13 = set_zero(); - D Cv20 = set_zero(); - D Cv21 = set_zero(); - D Cv22 = set_zero(); - D Cv23 = set_zero(); - D Cv30 = set_zero(); - D Cv31 = set_zero(); - D Cv32 = set_zero(); - D Cv33 = set_zero(); - D Cv40 = set_zero(); - D Cv41 = set_zero(); - D Cv42 = set_zero(); - D Cv43 = set_zero(); - - for (int64_t l = 0; l < k; l += vl) { - V Bv0 = load(B + ldb * (jj + 0) + l); - V Bv1 = load(B + ldb * (jj + 1) + l); - V Bv2 = load(B + ldb * (jj + 2) + l); - V Bv3 = load(B + ldb * (jj + 3) + l); - V Bv4 = load(B + ldb * (jj + 4) + l); - - V Av0 = load(A + lda * (ii + 0) + l); - Cv00 = madd(Av0, Bv0, Cv00); - Cv10 = madd(Av0, Bv1, Cv10); - Cv20 = madd(Av0, Bv2, Cv20); - Cv30 = madd(Av0, Bv3, Cv30); - Cv40 = madd(Av0, Bv4, Cv40); - - V Av1 = load(A + lda * (ii + 1) + l); - Cv01 = madd(Av1, Bv0, Cv01); - Cv11 = madd(Av1, Bv1, Cv11); - Cv21 = madd(Av1, Bv2, Cv21); - Cv31 = madd(Av1, Bv3, Cv31); - Cv41 = madd(Av1, Bv4, Cv41); - - V Av2 = load(A + lda * (ii + 2) + l); - Cv02 = madd(Av2, Bv0, Cv02); - Cv12 = madd(Av2, Bv1, Cv12); - Cv22 = madd(Av2, Bv2, Cv22); - Cv32 = madd(Av2, Bv3, Cv32); - Cv42 = madd(Av2, Bv4, Cv42); - - V Av3 = load(A + lda * (ii + 3) + l); - Cv03 = madd(Av3, Bv0, Cv03); - Cv13 = madd(Av3, Bv1, Cv13); - Cv23 = madd(Av3, Bv2, Cv23); - Cv33 = madd(Av3, Bv3, Cv33); - Cv43 = madd(Av3, Bv4, Cv43); - } - - C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); - C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); - C[ldc * (jj + 0) + (ii + 2)] = hsum(Cv02); - C[ldc * (jj + 0) + (ii + 3)] = hsum(Cv03); - C[ldc * (jj + 1) + (ii + 0)] = hsum(Cv10); - C[ldc * (jj + 1) + (ii + 1)] = hsum(Cv11); - C[ldc * (jj + 1) + (ii + 2)] = hsum(Cv12); - C[ldc * (jj + 1) + (ii + 3)] = hsum(Cv13); - C[ldc * (jj + 2) + (ii + 0)] = hsum(Cv20); - C[ldc * (jj + 2) + (ii + 1)] = hsum(Cv21); - C[ldc * (jj + 2) + (ii + 2)] = hsum(Cv22); - C[ldc * (jj + 2) + (ii + 3)] = hsum(Cv23); - C[ldc * (jj + 3) + (ii + 0)] = hsum(Cv30); - C[ldc * (jj + 3) + (ii + 1)] = hsum(Cv31); - C[ldc * (jj + 3) + (ii + 2)] = hsum(Cv32); - C[ldc * (jj + 3) + (ii + 3)] = hsum(Cv33); - C[ldc * (jj + 4) + (ii + 0)] = hsum(Cv40); - C[ldc * (jj + 4) + (ii + 1)] = hsum(Cv41); - C[ldc * (jj + 4) + (ii + 2)] = hsum(Cv42); - C[ldc * (jj + 4) + (ii + 3)] = hsum(Cv43); - } - - inline void gemm_bloc_4x4(int64_t ii, int64_t jj) { - size_t vl = vlmax(); - D Cv00 = set_zero(); - D Cv01 = set_zero(); - D Cv02 = set_zero(); - D Cv03 = set_zero(); - D Cv10 = set_zero(); - D Cv11 = set_zero(); - D Cv12 = set_zero(); - D Cv13 = set_zero(); - D Cv20 = set_zero(); - D Cv21 = set_zero(); - D Cv22 = set_zero(); - D Cv23 = set_zero(); - D Cv30 = set_zero(); - D Cv31 = set_zero(); - D Cv32 = set_zero(); - D Cv33 = set_zero(); - - for (int64_t l = 0; l < k; l += vl) { - V Av0 = load(A + lda * (ii + 0) + l); - V Av1 = load(A + lda * (ii + 1) + l); - V Av2 = load(A + lda * (ii + 2) + l); - V Av3 = load(A + lda * (ii + 3) + l); - - V Bv0 = load(B + ldb * (jj + 0) + l); - Cv00 = madd(Av0, Bv0, Cv00); - Cv01 = madd(Av1, Bv0, Cv01); - Cv02 = madd(Av2, Bv0, Cv02); - Cv03 = madd(Av3, Bv0, Cv03); - - V Bv1 = load(B + ldb * (jj + 1) + l); - Cv10 = madd(Av0, Bv1, Cv10); - Cv11 = madd(Av1, Bv1, Cv11); - Cv12 = madd(Av2, Bv1, Cv12); - Cv13 = madd(Av3, Bv1, Cv13); - - V Bv2 = load(B + ldb * (jj + 2) + l); - Cv20 = madd(Av0, Bv2, Cv20); - Cv21 = madd(Av1, Bv2, Cv21); - Cv22 = madd(Av2, Bv2, Cv22); - Cv23 = madd(Av3, Bv2, Cv23); - - V Bv3 = load(B + ldb * (jj + 3) + l); - Cv30 = madd(Av0, Bv3, Cv30); - Cv31 = madd(Av1, Bv3, Cv31); - Cv32 = madd(Av2, Bv3, Cv32); - Cv33 = madd(Av3, Bv3, Cv33); - } - - C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); - C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); - C[ldc * (jj + 0) + (ii + 2)] = hsum(Cv02); - C[ldc * (jj + 0) + (ii + 3)] = hsum(Cv03); - C[ldc * (jj + 1) + (ii + 0)] = hsum(Cv10); - C[ldc * (jj + 1) + (ii + 1)] = hsum(Cv11); - C[ldc * (jj + 1) + (ii + 2)] = hsum(Cv12); - C[ldc * (jj + 1) + (ii + 3)] = hsum(Cv13); - C[ldc * (jj + 2) + (ii + 0)] = hsum(Cv20); - C[ldc * (jj + 2) + (ii + 1)] = hsum(Cv21); - C[ldc * (jj + 2) + (ii + 2)] = hsum(Cv22); - C[ldc * (jj + 2) + (ii + 3)] = hsum(Cv23); - C[ldc * (jj + 3) + (ii + 0)] = hsum(Cv30); - C[ldc * (jj + 3) + (ii + 1)] = hsum(Cv31); - C[ldc * (jj + 3) + (ii + 2)] = hsum(Cv32); - C[ldc * (jj + 3) + (ii + 3)] = hsum(Cv33); - } - - inline void gemm_bloc_4x3(int64_t ii, int64_t jj) { - size_t vl = vlmax(); - D Cv00 = set_zero(); - D Cv01 = set_zero(); - D Cv02 = set_zero(); - D Cv03 = set_zero(); - D Cv10 = set_zero(); - D Cv11 = set_zero(); - D Cv12 = set_zero(); - D Cv13 = set_zero(); - D Cv20 = set_zero(); - D Cv21 = set_zero(); - D Cv22 = set_zero(); - D Cv23 = set_zero(); - - for (int64_t l = 0; l < k; l += vl) { - V Av0 = load(A + lda * (ii + 0) + l); - V Av1 = load(A + lda * (ii + 1) + l); - V Av2 = load(A + lda * (ii + 2) + l); - V Av3 = load(A + lda * (ii + 3) + l); - - V Bv0 = load(B + ldb * (jj + 0) + l); - Cv00 = madd(Av0, Bv0, Cv00); - Cv01 = madd(Av1, Bv0, Cv01); - Cv02 = madd(Av2, Bv0, Cv02); - Cv03 = madd(Av3, Bv0, Cv03); - - V Bv1 = load(B + ldb * (jj + 1) + l); - Cv10 = madd(Av0, Bv1, Cv10); - Cv11 = madd(Av1, Bv1, Cv11); - Cv12 = madd(Av2, Bv1, Cv12); - Cv13 = madd(Av3, Bv1, Cv13); - - V Bv2 = load(B + ldb * (jj + 2) + l); - Cv20 = madd(Av0, Bv2, Cv20); - Cv21 = madd(Av1, Bv2, Cv21); - Cv22 = madd(Av2, Bv2, Cv22); - Cv23 = madd(Av3, Bv2, Cv23); - } - - C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); - C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); - C[ldc * (jj + 0) + (ii + 2)] = hsum(Cv02); - C[ldc * (jj + 0) + (ii + 3)] = hsum(Cv03); - C[ldc * (jj + 1) + (ii + 0)] = hsum(Cv10); - C[ldc * (jj + 1) + (ii + 1)] = hsum(Cv11); - C[ldc * (jj + 1) + (ii + 2)] = hsum(Cv12); - C[ldc * (jj + 1) + (ii + 3)] = hsum(Cv13); - C[ldc * (jj + 2) + (ii + 0)] = hsum(Cv20); - C[ldc * (jj + 2) + (ii + 1)] = hsum(Cv21); - C[ldc * (jj + 2) + (ii + 2)] = hsum(Cv22); - C[ldc * (jj + 2) + (ii + 3)] = hsum(Cv23); - } - - inline void gemm_bloc_4x2(int64_t ii, int64_t jj) { - size_t vl = vlmax(); - D Cv00 = set_zero(); - D Cv01 = set_zero(); - D Cv02 = set_zero(); - D Cv03 = set_zero(); - D Cv10 = set_zero(); - D Cv11 = set_zero(); - D Cv12 = set_zero(); - D Cv13 = set_zero(); - - for (int64_t l = 0; l < k; l += vl) { - V Av0 = load(A + lda * (ii + 0) + l); - V Av1 = load(A + lda * (ii + 1) + l); - V Av2 = load(A + lda * (ii + 2) + l); - V Av3 = load(A + lda * (ii + 3) + l); - - V Bv0 = load(B + ldb * (jj + 0) + l); - Cv00 = madd(Av0, Bv0, Cv00); - Cv01 = madd(Av1, Bv0, Cv01); - Cv02 = madd(Av2, Bv0, Cv02); - Cv03 = madd(Av3, Bv0, Cv03); - - V Bv1 = load(B + ldb * (jj + 1) + l); - Cv10 = madd(Av0, Bv1, Cv10); - Cv11 = madd(Av1, Bv1, Cv11); - Cv12 = madd(Av2, Bv1, Cv12); - Cv13 = madd(Av3, Bv1, Cv13); - } - - C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); - C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); - C[ldc * (jj + 0) + (ii + 2)] = hsum(Cv02); - C[ldc * (jj + 0) + (ii + 3)] = hsum(Cv03); - C[ldc * (jj + 1) + (ii + 0)] = hsum(Cv10); - C[ldc * (jj + 1) + (ii + 1)] = hsum(Cv11); - C[ldc * (jj + 1) + (ii + 2)] = hsum(Cv12); - C[ldc * (jj + 1) + (ii + 3)] = hsum(Cv13); - } - - inline void gemm_bloc_4x1(int64_t ii, int64_t jj) { - size_t vl = vlmax(); - D Cv00 = set_zero(); - D Cv01 = set_zero(); - D Cv02 = set_zero(); - D Cv03 = set_zero(); - - for (int64_t l = 0; l < k; l += vl) { - V Av0 = load(A + lda * (ii + 0) + l); - V Av1 = load(A + lda * (ii + 1) + l); - V Av2 = load(A + lda * (ii + 2) + l); - V Av3 = load(A + lda * (ii + 3) + l); - - V Bv0 = load(B + ldb * (jj + 0) + l); - Cv00 = madd(Av0, Bv0, Cv00); - Cv01 = madd(Av1, Bv0, Cv01); - Cv02 = madd(Av2, Bv0, Cv02); - Cv03 = madd(Av3, Bv0, Cv03); - } - - C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); - C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); - C[ldc * (jj + 0) + (ii + 2)] = hsum(Cv02); - C[ldc * (jj + 0) + (ii + 3)] = hsum(Cv03); - } - - inline void gemm_bloc_2x2(int64_t ii, int64_t jj) { - size_t vl = vlmax(); - D Cv00 = set_zero(); - D Cv01 = set_zero(); - D Cv10 = set_zero(); - D Cv11 = set_zero(); - - for (int64_t l = 0; l < k; l += vl) { - V Av0 = load(A + lda * (ii + 0) + l); - V Av1 = load(A + lda * (ii + 1) + l); - - V Bv0 = load(B + ldb * (jj + 0) + l); - Cv00 = madd(Av0, Bv0, Cv00); - Cv01 = madd(Av1, Bv0, Cv01); - - V Bv1 = load(B + ldb * (jj + 1) + l); - Cv10 = madd(Av0, Bv1, Cv10); - Cv11 = madd(Av1, Bv1, Cv11); - } - - C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); - C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); - C[ldc * (jj + 1) + (ii + 0)] = hsum(Cv10); - C[ldc * (jj + 1) + (ii + 1)] = hsum(Cv11); - } - - inline void gemm_bloc_2x1(int64_t ii, int64_t jj) { - size_t vl = vlmax(); - D Cv00 = set_zero(); - D Cv01 = set_zero(); - - for (int64_t l = 0; l < k; l += vl) { - V Av0 = load(A + lda * (ii + 0) + l); - V Av1 = load(A + lda * (ii + 1) + l); - - V Bv0 = load(B + ldb * (jj + 0) + l); - Cv00 = madd(Av0, Bv0, Cv00); - Cv01 = madd(Av1, Bv0, Cv01); - } - - C[ldc * (jj + 0) + (ii + 0)] = hsum(Cv00); - C[ldc * (jj + 0) + (ii + 1)] = hsum(Cv01); - } - - template - inline void gemm_bloc(int64_t ii, int64_t jj) { - if constexpr (RM == 4) { - if constexpr (RN == 6) { return gemm_bloc_4x6(ii, jj); } - if constexpr (RN == 5) { return gemm_bloc_4x5(ii, jj); } - if constexpr (RN == 4) { return gemm_bloc_4x4(ii, jj); } - if constexpr (RN == 3) { return gemm_bloc_4x3(ii, jj); } - if constexpr (RN == 2) { return gemm_bloc_4x2(ii, jj); } - if constexpr (RN == 1) { return gemm_bloc_4x1(ii, jj); } - } else if constexpr (RM == 2) { - if constexpr (RN == 2) { return gemm_bloc_2x2(ii, jj); } - if constexpr (RN == 1) { return gemm_bloc_2x1(ii, jj); } - } - } - - template - NOINLINE void gemm(int64_t m, int64_t n, int64_t BN) { - GGML_ASSERT(m % (RM * BM) == 0); - const int64_t ytiles = m / (RM * BM); - const int64_t xtiles = (n + RN -1) / RN; - const int64_t jj_RN = (xtiles - (xtiles * RN - n)); - - // "round" bloc_size to "nearest" BN - const int64_t NB_BN = xtiles < BN ? 1 : (xtiles + BN / 2) / BN; - const int64_t SIZE_BN = xtiles % NB_BN == 0 ? xtiles / NB_BN : xtiles / NB_BN + 1; - const int64_t jj_BN = (NB_BN - (NB_BN * SIZE_BN - xtiles)); - const int64_t nb_job = ytiles * NB_BN; - - if (params->ith == 0) { - GGML_ASSERT( jj_BN * SIZE_BN + (NB_BN - jj_BN) * (SIZE_BN - 1) == xtiles); - // Every thread starts at ith, so the first unprocessed chunk is nth. This save a bit of coordination right at the start. - ggml_threadpool_chunk_set(params->threadpool, params->nth); - } - - ggml_barrier(params->threadpool); - - int64_t job = params->ith; - while (job < nb_job) { - const int64_t ii = (job % ytiles) * RM * BM; - const int64_t jb = job / ytiles; - const int64_t jr0 = BLOC_POS(jb , jj_BN, SIZE_BN); - const int64_t jrN = BLOC_POS(jb+1, jj_BN, SIZE_BN); - - const int64_t jj0 = BLOC_POS(jr0, jj_RN, RN); - const int64_t jj2 = BLOC_POS(jrN, jj_RN, RN); - const int64_t jj1 = jj2 < jj_RN * RN ? jj2 : jj_RN * RN; - - for (int64_t bi = 0; bi < BM * RM; bi += RM) { - int64_t jj = jj0; - for (; jj < jj1; jj += RN) { - gemm_bloc(ii + bi, jj); - } - if constexpr (RN > 1) { - for (; jj < jj2; jj += RN - 1) { - gemm_bloc(ii + bi, jj); - } - } - GGML_ASSERT(jj == jj2); - } - - job = ggml_threadpool_chunk_add(params->threadpool, 1); - } - - ggml_barrier(params->threadpool); - return; - } - - const ggml_compute_params * params; - const TA *const A; - const TB *const B; - TC *const C; - const int64_t k; - const int64_t lda; - const int64_t ldb; - const int64_t ldc; -}; -#endif - -////////////////////////////////////////////////////////////////////////////////////////// -// QUANT ZERO MATRIX MULTIPLICATION - -#if defined(__ARM_FEATURE_DOTPROD) -template -class tinyBLAS_Q0_ARM { - public: - tinyBLAS_Q0_ARM(int64_t k, - const TA *A, int64_t lda, - const block_q8_0 *B, int64_t ldb, - float *C, int64_t ldc, - int ith, int nth) - : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { - } - - void matmul(int64_t m, int64_t n) { - mnpack(0, m, 0, n); - } - - private: - NOINLINE void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int64_t mc, nc, mp, np; - switch ((MIN(m - m0, 3) << 4) | MIN(n - n0, 3ll)) { - case 0x33: - mc = 3; - nc = 3; - gemm<3, 3>(m0, m, n0, n); - break; - case 0x32: - mc = 3; - nc = 2; - gemm<3, 2>(m0, m, n0, n); - break; - case 0x23: - mc = 2; - nc = 3; - gemm<2, 3>(m0, m, n0, n); - break; - case 0x22: - mc = 2; - nc = 2; - gemm<2, 2>(m0, m, n0, n); - break; - case 0x31: - mc = 3; - nc = 1; - gemm<3, 1>(m0, m, n0, n); - break; - case 0x13: - mc = 1; - nc = 3; - gemm<1, 3>(m0, m, n0, n); - break; - case 0x21: - mc = 2; - nc = 1; - gemm<2, 1>(m0, m, n0, n); - break; - case 0x12: - mc = 1; - nc = 2; - gemm<1, 2>(m0, m, n0, n); - break; - case 0x11: - mc = 1; - nc = 1; - gemm<1, 1>(m0, m, n0, n); - break; - default: - return; - } - mp = m0 + (m - m0) / mc * mc; - np = n0 + (n - n0) / nc * nc; - mnpack(mp, m, n0, np); - mnpack(m0, m, np, n); - } - - template - NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int64_t ytiles = (m - m0) / RM; - int64_t xtiles = (n - n0) / RN; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - if (end > tiles) - end = tiles; - for (int64_t job = start; job < end; ++job) { - int64_t ii = m0 + job / xtiles * RM; - int64_t jj = n0 + job % xtiles * RN; - float32x4_t Cv[RN][RM] = {}; - for (int64_t l = 0; l < k; ++l) - for (int64_t j = 0; j < RN; ++j) - for (int64_t i = 0; i < RM; ++i) - Cv[j][i] = vmlaq_n_f32(Cv[j][i], - vcvtq_f32_s32(vdotq_s32( - vdotq_s32(vdupq_n_s32(0), - load_lo(A + lda * (ii + i) + l), - load_lo(B + ldb * (jj + j) + l)), - load_hi(A + lda * (ii + i) + l), - load_hi(B + ldb * (jj + j) + l))), - unhalf(A[lda * (ii + i) + l].d) * - unhalf(B[ldb * (jj + j) + l].d)); - for (int64_t j = 0; j < RN; ++j) - for (int64_t i = 0; i < RM; ++i) - C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]); - } - } - - inline int8x16_t load_lo(const block_q8_0 *b) { - return vld1q_s8(b->qs); - } - - inline int8x16_t load_hi(const block_q8_0 *b) { - return vld1q_s8(b->qs + 16); - } - - inline int8x16_t load_lo(const block_q4_0 *b) { - return vsubq_s8(vreinterpretq_s8_u8(vandq_u8(vld1q_u8(b->qs), - vdupq_n_u8(0x0f))), - vdupq_n_s8(0x8)); - } - - inline int8x16_t load_hi(const block_q4_0 *b) { - return vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(vld1q_u8(b->qs), 4)), - vdupq_n_s8(0x8)); - } - - const TA *const A; - const block_q8_0 *const B; - float *const C; - const int64_t k; - const int64_t lda; - const int64_t ldb; - const int64_t ldc; - const int ith; - const int nth; -}; -#endif // __ARM_FEATURE_DOTPROD - -#if defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX__) -template -class tinyBLAS_Q0_AVX { - public: - tinyBLAS_Q0_AVX(int64_t k, - const TA *A, int64_t lda, - const TB *B, int64_t ldb, - TC *C, int64_t ldc, - int ith, int nth) - : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { - const int8_t kvalues_iq4nl[16] = { - -127, -104, -83, -65, - -49, -35, -22, -10, - 1, 13, 25, 38, - 53, 69, 89, 113 - }; - - iq4nlt = _mm_loadu_si128((const __m128i *)kvalues_iq4nl); - } - - void matmul(int64_t m, int64_t n) { - mnpack(0, m, 0, n); - } - - private: - void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int64_t mc, nc, mp, np; - switch ((MIN(m - m0, 4) << 4) | MIN(n - n0, 4)) { -#if VECTOR_REGISTERS == 32 - case 0x44: - mc = 4; - nc = 4; -#if defined(__AVX2__) && defined(__F16C__) - gemm4xN<4>(m0, m, n0, n); -#else - gemm<4, 4>(m0, m, n0, n); -#endif - break; - case 0x43: - mc = 4; - nc = 3; -#if defined(__AVX2__) && defined(__F16C__) - gemm4xN<3>(m0, m, n0, n); -#else - gemm<4, 3>(m0, m, n0, n); -#endif - break; - case 0x34: - mc = 3; - nc = 4; -#if defined(__AVX2__) && defined(__F16C__) - gemmMx4<3>(m0, m, n0, n); -#else - gemm<3, 4>(m0, m, n0, n); -#endif - break; - case 0x33: - mc = 3; - nc = 3; - gemm<3, 3>(m0, m, n0, n); - break; - case 0x42: - mc = 4; - nc = 2; -#if defined(__AVX2__) && defined(__F16C__) - gemm4xN<2>(m0, m, n0, n); -#else - gemm<4, 2>(m0, m, n0, n); -#endif - break; - case 0x24: - mc = 2; - nc = 4; -#if defined(__AVX2__) && defined(__F16C__) - gemmMx4<2>(m0, m, n0, n); -#else - gemm<2, 4>(m0, m, n0, n); -#endif - break; -#else - case 0x44: - case 0x43: - case 0x42: - mc = 4; - nc = 2; -#if defined(__AVX2__) && defined(__F16C__) - gemm4xN<2>(m0, m, n0, n); -#else - gemm<4, 2>(m0, m, n0, n); -#endif - break; - case 0x34: - case 0x24: - mc = 2; - nc = 4; -#if defined(__AVX2__) && defined(__F16C__) - gemmMx4<2>(m0, m, n0, n); -#else - gemm<2, 4>(m0, m, n0, n); -#endif - break; - case 0x33: -#endif - case 0x32: - mc = 3; - nc = 2; - gemm<3, 2>(m0, m, n0, n); - break; - case 0x23: - mc = 2; - nc = 3; - gemm<2, 3>(m0, m, n0, n); - break; - case 0x41: - mc = 4; - nc = 1; -#if defined(__AVX2__) && defined(__F16C__) - gemm4xN<1>(m0, m, n0, n); -#else - gemm<4, 1>(m0, m, n0, n); -#endif - break; - case 0x22: - mc = 2; - nc = 2; - gemm<2, 2>(m0, m, n0, n); - break; - case 0x14: - mc = 1; - nc = 4; -#if defined(__AVX2__) && defined(__F16C__) - gemmMx4<1>(m0, m, n0, n); -#else - gemm<1, 4>(m0, m, n0, n); -#endif - break; - case 0x31: - mc = 3; - nc = 1; - gemm<3, 1>(m0, m, n0, n); - break; - case 0x13: - mc = 1; - nc = 3; - gemm<1, 3>(m0, m, n0, n); - break; - case 0x21: - mc = 2; - nc = 1; - gemm<2, 1>(m0, m, n0, n); - break; - case 0x12: - mc = 1; - nc = 2; - gemm<1, 2>(m0, m, n0, n); - break; - case 0x11: - mc = 1; - nc = 1; - gemm<1, 1>(m0, m, n0, n); - break; - default: - return; - } - mp = m0 + (m - m0) / mc * mc; - np = n0 + (n - n0) / nc * nc; - mnpack(mp, m, n0, np); - mnpack(m0, m, np, n); - } - -#if defined(__AVX2__) && defined(__F16C__) -// Templated functions for gemm of dimensions 4xN - template - NOINLINE void gemm4xN(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int64_t ytiles = (m - m0) / 4; - int64_t xtiles = (n - n0) / RN; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - if (end > tiles) - end = tiles; - for (int64_t job = start; job < end; ++job) { - int64_t ii = m0 + job / xtiles * 4; - int64_t jj = n0 + job % xtiles * RN; - __m256 Cv[RN][4] = {}; - for (int64_t l = 0; l < k; ++l) { - uint64_t a_delta = ((uint64_t)A[lda * (ii + 3) + l].d << 48) | ((uint64_t)A[lda * (ii + 2) + l].d << 32) | ((uint64_t)A[lda * (ii + 1) + l].d << 16) | (A[lda * (ii + 0) + l].d); - // Convert delta values for four blocks to float values - __m128 da = _mm_cvtph_ps(_mm_set_epi64x(0, a_delta)); - __m256i avec0 = load(A + lda * (ii + 0) + l); - __m256i avec1 = load(A + lda * (ii + 1) + l); - __m256i avec2 = load(A + lda * (ii + 2) + l); - __m256i avec3 = load(A + lda * (ii + 3) + l); - for (int64_t j = 0; j < RN; ++j) { - __m128 db = _mm_set1_ps(unhalf(B[ldb * (jj + j) + l].d)); - // Computation of product of delta values for four blocks and replicate it across 256 bit lane - __m256 dvec = _mm256_castps128_ps256(_mm_mul_ps(da, db)); - dvec = _mm256_permute2f128_ps(dvec ,dvec, 0); - // Computation of dot product and multiplication with appropriate delta value products - Cv[j][0] = madd(_mm256_shuffle_ps(dvec, dvec, 0), - updot(_mm256_sign_epi8(avec0, avec0), - _mm256_sign_epi8(load(B + ldb * (jj + j) + l), avec0)), - Cv[j][0]); - Cv[j][1] = madd(_mm256_shuffle_ps(dvec, dvec, 85), - updot(_mm256_sign_epi8(avec1, avec1), - _mm256_sign_epi8(load(B + ldb * (jj + j) + l), avec1)), - Cv[j][1]); - Cv[j][2] = madd(_mm256_shuffle_ps(dvec, dvec, 170), - updot(_mm256_sign_epi8(avec2, avec2), - _mm256_sign_epi8(load(B + ldb * (jj + j) + l), avec2)), - Cv[j][2]); - Cv[j][3] = madd(_mm256_shuffle_ps(dvec, dvec, 255), - updot(_mm256_sign_epi8(avec3, avec3), - _mm256_sign_epi8(load(B + ldb * (jj + j) + l), avec3)), - Cv[j][3]); - } - } - - for (int64_t j = 0; j < RN; ++j) - for (int64_t i = 0; i < 4; ++i) - C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]); - } - } - - // Templated functions for gemm of dimensions Mx4 - template - NOINLINE void gemmMx4(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int64_t ytiles = (m - m0) / RM; - int64_t xtiles = (n - n0) / 4; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - if (end > tiles) - end = tiles; - for (int64_t job = start; job < end; ++job) { - int64_t ii = m0 + job / xtiles * RM; - int64_t jj = n0 + job % xtiles * 4; - __m256 Cv[4][RM] = {}; - for (int64_t l = 0; l < k; ++l) { - uint64_t b_delta = ((uint64_t)B[ldb * (jj + 3) + l].d << 48) | ((uint64_t)B[ldb * (jj + 2) + l].d << 32) | ((uint64_t)B[ldb * (jj + 1) + l].d << 16) | (B[ldb * (jj + 0) + l].d); - // Convert delta values for four blocks to float values - __m128 db = _mm_cvtph_ps(_mm_set_epi64x(0, b_delta)); - __m256i bvec0 = load(B + ldb * (jj + 0) + l); - __m256i bvec1 = load(B + ldb * (jj + 1) + l); - __m256i bvec2 = load(B + ldb * (jj + 2) + l); - __m256i bvec3 = load(B + ldb * (jj + 3) + l); - for (int64_t i = 0; i < RM; ++i) { - __m128 da = _mm_set1_ps(unhalf((A[lda * (ii + i) + l].d))); - // Computation of product of delta values for four blocks and replicate it across 256 bit lane - __m256 dvec = _mm256_castps128_ps256(_mm_mul_ps(da, db)); - dvec = _mm256_permute2f128_ps(dvec ,dvec, 0); - // Computation of dot product and multiplication with appropriate delta value products - Cv[0][i] = madd(_mm256_shuffle_ps(dvec, dvec, 0), - updot(_mm256_sign_epi8(load(A + lda * (ii + i) + l), - load(A + lda * (ii + i) + l)), - _mm256_sign_epi8(bvec0, load(A + lda * (ii + i) + l))), - Cv[0][i]); - Cv[1][i] = madd(_mm256_shuffle_ps(dvec, dvec, 85), - updot(_mm256_sign_epi8(load(A + lda * (ii + i) + l), - load(A + lda * (ii + i) + l)), - _mm256_sign_epi8(bvec1, load(A + lda * (ii + i) + l))), - Cv[1][i]); - Cv[2][i] = madd(_mm256_shuffle_ps(dvec, dvec, 170), - updot(_mm256_sign_epi8(load(A + lda * (ii + i) + l), - load(A + lda * (ii + i) + l)), - _mm256_sign_epi8(bvec2, load(A + lda * (ii + i) + l))), - Cv[2][i]); - Cv[3][i] = madd(_mm256_shuffle_ps(dvec, dvec, 255), - updot(_mm256_sign_epi8(load(A + lda * (ii + i) + l), - load(A + lda * (ii + i) + l)), - _mm256_sign_epi8(bvec3, load(A + lda * (ii + i) + l))), - Cv[3][i]); - } - } - for (int64_t j = 0; j < 4; ++j) - for (int64_t i = 0; i < RM; ++i) - C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]); - } - } -#endif - - template - NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int64_t ytiles = (m - m0) / RM; - int64_t xtiles = (n - n0) / RN; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - if (end > tiles) - end = tiles; - for (int64_t job = start; job < end; ++job) { - int64_t ii = m0 + job / xtiles * RM; - int64_t jj = n0 + job % xtiles * RN; - __m256 Cv[RN][RM] = {}; - for (int64_t l = 0; l < k; ++l) - for (int64_t j = 0; j < RN; ++j) - for (int64_t i = 0; i < RM; ++i) { -#if defined(__AVX2__) - __m256 udTmp = updot(_mm256_sign_epi8(load(A + lda * (ii + i) + l), - load(A + lda * (ii + i) + l)), - _mm256_sign_epi8(load(B + ldb * (jj + j) + l), - load(A + lda * (ii + i) + l))); -#else - __m128i ali0 = load0(A + lda * (ii + i) + l); - __m128i ali1 = load1(A + lda * (ii + i) + l); - __m128i blj0 = load0(B + ldb * (jj + j) + l); - __m128i blj1 = load1(B + ldb * (jj + j) + l); - - __m128i sepAA0 = _mm_sign_epi8(ali0, ali0); - __m128i sepAA1 = _mm_sign_epi8(ali1, ali1); - __m128i sepBA0 = _mm_sign_epi8(blj0, ali0); - __m128i sepBA1 = _mm_sign_epi8(blj1, ali1); - - // updot - const __m128i oneFill = _mm_set1_epi16(1); - __m128i mad0 = _mm_maddubs_epi16(sepAA0, sepBA0); - __m128i mad1 = _mm_maddubs_epi16(sepAA1, sepBA1); - __m256 udTmp = _mm256_cvtepi32_ps(MM256_SET_M128I(_mm_madd_epi16(oneFill, mad1), _mm_madd_epi16(oneFill, mad0))); -#endif - Cv[j][i] = madd(_mm256_set1_ps(unhalf(A[lda * (ii + i) + l].d) * - unhalf(B[ldb * (jj + j) + l].d)), - udTmp, - Cv[j][i]); - } - for (int64_t j = 0; j < RN; ++j) - for (int64_t i = 0; i < RM; ++i) - C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]); - } - } - - inline __m256i load(const block_q8_0 *b) { - return _mm256_loadu_si256((const __m256i *)b->qs); - } - - inline __m128i load0(const block_q8_0 *b) { - return _mm_loadu_si128((const __m128i *)b->qs); - } - - inline __m128i load1(const block_q8_0 *b) { - return _mm_loadu_si128(((const __m128i *)b->qs) + 1); - } - - inline __m256i load(const block_q4_0 *b) { - return _mm256_sub_epi8(denibble(b->qs), _mm256_set1_epi8(8)); - } - - inline __m128i load0(const block_q4_0 *b) { - const __m128i x = _mm_loadu_si128((const __m128i *)(b->qs)); - return _mm_sub_epi8(_mm_and_si128(_mm_set1_epi8(15), x), _mm_set1_epi8(8)); - } - - inline __m128i load1(const block_q4_0 *b) { - const __m128i x = _mm_loadu_si128((const __m128i *)(b->qs)); - return _mm_sub_epi8(_mm_and_si128(_mm_set1_epi8(15), _mm_srli_epi16(x, 4)), _mm_set1_epi8(8)); - } - - inline __m256i load(const block_q5_0 *b) { - return _mm256_or_si256(denibble(b->qs), bittobyte(b->qh)); - } - - inline __m128i load0(const block_q5_0* b) { - const __m128i x = _mm_loadu_si128((const __m128i *)(b->qs)); - uint32_t x32; - memcpy(&x32, b->qh, sizeof(uint32_t)); - __m128i qxl = _mm_and_si128(_mm_set1_epi8(15), x); - __m128i bytesl = _mm_cmpeq_epi8(_mm_set1_epi64x(-1), - _mm_or_si128(_mm_set1_epi64x(0x7fbfdfeff7fbfdfe), - _mm_shuffle_epi8(_mm_set1_epi32(x32), - _mm_set_epi64x(0x0101010101010101, 0x0000000000000000)))); - bytesl = _mm_andnot_si128(bytesl, _mm_set1_epi8((char)0xF0)); - return _mm_or_si128(qxl, bytesl); - } - - inline __m128i load1(const block_q5_0* b) { - const __m128i x = _mm_loadu_si128((const __m128i *)(b->qs)); - uint32_t x32; - memcpy(&x32, b->qh, sizeof(uint32_t)); - __m128i qxh = _mm_and_si128(_mm_set1_epi8(15), _mm_srli_epi16(x, 4)); - __m128i bytesh = _mm_cmpeq_epi8(_mm_set1_epi64x(-1), - _mm_or_si128(_mm_set1_epi64x(0x7fbfdfeff7fbfdfe), - _mm_shuffle_epi8(_mm_set1_epi32(x32), - _mm_set_epi64x(0x0303030303030303, 0x0202020202020202)))); - bytesh = _mm_andnot_si128(bytesh, _mm_set1_epi8((char)0xF0)); - return _mm_or_si128(qxh, bytesh); - } - - inline __m256i load(const block_iq4_nl *b) { - return MM256_SET_M128I(load1(b), load0(b)); - } - - inline __m128i load0(const block_iq4_nl *b) { - const __m128i x = _mm_loadu_si128((const __m128i *)(b->qs)); - return _mm_shuffle_epi8(iq4nlt, _mm_and_si128(_mm_set1_epi8(15), x)); - } - - inline __m128i load1(const block_iq4_nl *b) { - const __m128i x = _mm_loadu_si128((const __m128i *)(b->qs)); - return _mm_shuffle_epi8(iq4nlt, _mm_and_si128(_mm_set1_epi8(15), _mm_srli_epi16(x, 4))); - } - - inline __m256 updot(__m256i u, __m256i s) { - __m256i res; -#if defined(__AVX512VNNI__) && defined(__AVX512VL__) - res = _mm256_dpbusd_epi32(_mm256_setzero_si256(), u, s); -#elif defined(__AVXVNNI__) - res = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), u, s); -#else - res = _mm256_madd_epi16(_mm256_set1_epi16(1), _mm256_maddubs_epi16(u, s)); -#endif - return _mm256_cvtepi32_ps(res); - } - - static inline __m256i denibble(const uint8_t *p) { - __m128i x = _mm_loadu_si128((const __m128i *)p); - return _mm256_and_si256(_mm256_set1_epi8(15), - _mm256_insertf128_si256(_mm256_castsi128_si256(x), - _mm_srli_epi16(x, 4), 1)); - } - - static inline __m256i bittobyte(const uint8_t *p) { - uint32_t x32; - memcpy(&x32, p, sizeof(uint32_t)); - __m256i bytes = _mm256_cmpeq_epi8(_mm256_set1_epi64x(-1), - _mm256_or_si256(_mm256_set1_epi64x(0x7fbfdfeff7fbfdfe), - _mm256_shuffle_epi8(_mm256_set1_epi32(x32), - _mm256_set_epi64x(0x0303030303030303, 0x0202020202020202, - 0x0101010101010101, 0x0000000000000000)))); - return _mm256_andnot_si256(bytes, _mm256_set1_epi8((char)0xF0)); - } - - const TA *const A; - const TB *const B; - TC *const C; - const int64_t k; - const int64_t lda; - const int64_t ldb; - const int64_t ldc; - const int ith; - const int nth; - __m128i iq4nlt; -}; -#endif // __AVX__ - -//PPC Implementation -#if defined(__MMA__) - -#define SAVE_ACC(ACC, ii, jj) \ - __builtin_mma_disassemble_acc(vec_C, ACC); \ - for (int I = 0; I < 4; I++) { \ - for (int J = 0; J < 4; J++) { \ - *((float*)(C+ii+((jj+J)*ldc)+I)) = *((float*)&vec_C[I]+J); \ - } \ - } \ - -template -struct mma_instr; - -template<> -struct mma_instr { - static inline void outer_product(acc_t *acc, vec_t a, vec_t b) { - __builtin_mma_xvbf16ger2pp(acc, a, b); - } -}; - -template<> -struct mma_instr { - static inline void outer_product(acc_t *acc, vec_t a, vec_t b) { - __builtin_mma_xvf16ger2pp(acc, a, b); - } -}; - -template -class tinyBLAS_HP16_PPC { - public: - tinyBLAS_HP16_PPC(int64_t k, - const TA *A, int64_t lda, - const TB *B, int64_t ldb, - TC *C, int64_t ldc, - int ith, int nth) - : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { - } - - void matmul(int64_t m, int64_t n) { - mnpack(0, m, 0, n); - } - - private: - void vector_permute_store(vec_t *c, int numVec, unsigned char *vecOffset) { - vec_t t[8], s[8]; - vec_t swiz1 = {0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23}; - vec_t swiz2 = {8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31}; - vec_t swiz3 = {0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23}; - vec_t swiz4 = {8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31}; - - if (numVec == 2) { - t[0] = vec_perm(c[0], c[1], swiz1); - t[1] = vec_perm(c[2], c[3], swiz1); - s[0] = vec_perm(t[0], t[1], swiz3); - s[1] = vec_perm(t[0], t[1], swiz4); - vec_xst(s[0], 0, (vec_t*)vecOffset); - vec_xst(s[1], 0, (vec_t*)(vecOffset + 16)); - } else if (numVec == 4) { - t[0] = vec_perm(c[0], c[1], swiz1); - t[1] = vec_perm(c[0], c[1], swiz2); - t[2] = vec_perm(c[2], c[3], swiz1); - t[3] = vec_perm(c[2], c[3], swiz2); - s[0] = vec_perm(t[0], t[2], swiz3); - s[1] = vec_perm(t[0], t[2], swiz4); - s[2] = vec_perm(t[1], t[3], swiz3); - s[3] = vec_perm(t[1], t[3], swiz4); - for (int i = 0; i < 4; ++i) - vec_xst(s[i], 0, (vec_t*)(vecOffset + i * 16)); - } else if (numVec == 8) { - for (int i = 0; i < 4; i += 2) { - t[i+0] = vec_perm(c[i+0], c[i+1], swiz1); - t[i+1] = vec_perm(c[i+0], c[i+1], swiz2); - } - for (int i = 4; i < 8; i += 2) { - t[i+0] = vec_perm(c[i+0], c[i+1], swiz1); - t[i+1] = vec_perm(c[i+0], c[i+1], swiz2); - } - s[0] = vec_perm(t[0], t[2], swiz3); - s[1] = vec_perm(t[0], t[2], swiz4); - s[2] = vec_perm(t[1], t[3], swiz3); - s[3] = vec_perm(t[1], t[3], swiz4); - s[4] = vec_perm(t[4], t[6], swiz3); - s[5] = vec_perm(t[4], t[6], swiz4); - s[6] = vec_perm(t[5], t[7], swiz3); - s[7] = vec_perm(t[5], t[7], swiz4); - for (int i = 0; i < 8; ++i) - vec_xst(s[i], 0, (vec_t*)(vecOffset + i * 16)); - } - } - - void packNormal(const TA* a, int64_t lda, int rows, int cols, unsigned char* vec) { - int64_t i, j; - TA *aoffset = NULL; - unsigned char *vecOffset = NULL; - TA * aoffsets[8]; - vector unsigned char c_arr[8]; - aoffset = const_cast(a); - vecOffset = vec; - j = (rows >> 3); - if (j > 0) { - do { - if (cols == 4) { - aoffsets[0] = aoffset; - for (int it = 1; it < 4; ++it) - aoffsets[it] = aoffsets[it-1] + lda; - aoffset += 4 * lda; - for (int i = 0; i < 4; ++i) - c_arr[i] = vec_xl(0, (vector unsigned char*)aoffsets[i]); - vector_permute_store(c_arr, 4, vecOffset); - for (int i = 0; i<4; i++) - aoffsets[i] = aoffsets[i]+lda; - vecOffset +=64; - } - i = (cols >> 3); - if (i > 0) { - aoffsets[0] = aoffset; - for (int it = 1; it < 8; ++it) { - aoffsets[it] = aoffsets[it-1] + lda; - } - aoffset += 8 * lda; - do { - for (int it = 0; it < 8; ++it) - c_arr[it] = vec_xl(0, (vector unsigned char*)aoffsets[it]); - vector_permute_store(c_arr, 8, vecOffset); - for (int it = 0; it < 8; ++it) - aoffsets[it] = aoffsets[it] + 8*lda; - vecOffset += 128; - i--; - } while(i > 0); - } - j--; - } while(j > 0); - } - if (rows & 4) { - aoffsets[0] = aoffset; - for (int it = 1; it < 4; ++it) - aoffsets[it] = aoffsets[it-1] + lda; - aoffset += 4 * lda; - if (cols == 4) { - for (int it = 0; it < 4; ++it) - c_arr[it] = vec_xl(0, (vector unsigned char*)aoffsets[it]); - vector_permute_store(c_arr, 2, vecOffset); - for (int it = 0; it< 4; it++) - aoffsets[it] = aoffsets[it] + lda; - vecOffset += 32; - } - i = (cols >> 3); - if (i > 0) { - do { - for (int it = 0; it < 4; ++it) - c_arr[it] = vec_xl(0, (vector unsigned char*)aoffsets[it]); - vector_permute_store(c_arr, 4, vecOffset); - for (int it = 0; it< 4; it++) - aoffsets[it] = aoffsets[it] + 8*lda; - vecOffset += 64; - i--; - } while(i > 0); - } - } - if (rows & 3) { - aoffsets[0] = aoffset; - for (int it = 1; it < 4; ++it) - aoffsets[it] = aoffsets[it-1] + lda; - if (cols == 4) { - switch(rows) { - case 3: c_arr[2] = vec_xl(0, (vector unsigned char*)aoffsets[2]); - case 2: c_arr[1] = vec_xl(0, (vector unsigned char*)aoffsets[1]); - case 1: c_arr[0] = vec_xl(0, (vector unsigned char*)aoffsets[0]); - break; - } - vector_permute_store(c_arr, 2, vecOffset); - for (int it = 0; it< 4; it++) - aoffsets[it] = aoffsets[it] + lda; - vecOffset += 32; - } - i = (cols >> 3); - if (i > 0) { - do { - switch(rows) { - case 3: c_arr[2] = vec_xl(0, (vector unsigned char*)aoffsets[2]); - case 2: c_arr[1] = vec_xl(0, (vector unsigned char*)aoffsets[1]); - case 1: c_arr[0] = vec_xl(0, (vector unsigned char*)aoffsets[0]); - break; - } - vector_permute_store(c_arr, 4, vecOffset); - for (int it = 0; it <4; it++) - aoffsets[it] = aoffsets[it] + 8* lda; - vecOffset += 64; - i--; - } while(i > 0); - } - } - } - - void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int64_t mc, nc, mp, np; - int m_rem = MIN(m - m0, 8); - int n_rem = MIN(n - n0, 8); - - if (m_rem >= 8 && n_rem >= 8) { - mc = 8; - nc = 8; - gemm<8,8>(m0, m, n0, n); - } else if (m_rem >= 4 && n_rem >= 8) { - mc = 4; - nc = 8; - gemm<4,8>(m0, m, n0, n); - } else if (m_rem >=8 && n_rem >=4){ - mc = 8; - nc = 4; - gemm<8,4>(m0, m, n0, n); - } else if ((m_rem < 4) && (n_rem >= 8)) { - nc = 8; - switch(m_rem) { - case 1: - mc = 1; - gemm_Mx8<1>(m0, m, n0, n); - break; - case 2: - mc = 2; - gemm_Mx8<2>(m0, m, n0, n); - break; - case 3: - mc = 3; - gemm_Mx8<3>(m0, m, n0, n); - break; - default: - return; - } - } else if (m_rem >= 4 && n_rem >= 4) { - mc = 4; - nc = 4; - gemm_small<4, 4>(m0, m, n0, n); - } else if ((m_rem > 4) && (n_rem < 4)) { - mc = 4; - switch(n_rem) { - case 1: - nc = 1; - gemm_small<4, 1>(m0, m, n0, n); - break; - case 2: - nc = 2; - gemm_small<4, 2>(m0, m, n0, n); - break; - case 3: - nc = 3; - gemm_small<4, 3>(m0, m, n0, n); - break; - - default: - return; - } - } else { - switch((m_rem << 4) | n_rem) { - case 0x43: - mc = 4; - nc = 3; - gemm_small<4, 3>(m0, m, n0, n); - break; - case 0x42: - mc = 4; - nc = 2; - gemm_small<4, 2>(m0, m, n0, n); - break; - case 0x41: - mc = 4; - nc = 1; - gemm_small<4, 1>(m0, m, n0, n); - break; - case 0x34: - mc = 3; - nc = 4; - gemm_small<3, 4>(m0, m, n0, n); - break; - case 0x33: - mc = 3; - nc = 3; - gemm_small<3, 3>(m0, m, n0, n); - break; - case 0x32: - mc = 3; - nc = 2; - gemm_small<3, 2>(m0, m, n0, n); - break; - case 0x31: - mc = 3; - nc = 1; - gemm_small<3, 1>(m0, m, n0, n); - break; - case 0x24: - mc = 2; - nc = 4; - gemm_small<2,4>(m0, m, n0, n); - break; - case 0x23: - mc = 2; - nc = 3; - gemm_small<2, 3>(m0, m, n0, n); - break; - case 0x22: - mc = 2; - nc = 2; - gemm_small<2, 2>(m0, m, n0, n); - break; - case 0x21: - mc = 2; - nc = 1; - gemm_small<2, 1>(m0, m, n0, n); - break; - case 0x14: - mc = 1; - nc = 4; - gemm_small<1, 4>(m0, m, n0, n); - break; - case 0x13: - mc = 1; - nc = 3; - gemm_small<1, 3>(m0, m, n0, n); - break; - case 0x12: - mc = 1; - nc = 2; - gemm_small<1, 2>(m0, m, n0, n); - break; - case 0x11: - mc = 1; - nc = 1; - gemm_small<1, 1>(m0, m, n0, n); - break; - default: - return; - } - } - mp = m0 + (m - m0) / mc * mc; - np = n0 + (n - n0) / nc * nc; - mnpack(mp, m, n0, np); - mnpack(m0, m, np, n); - } - - void KERNEL_4x8(int64_t ii, int64_t jj) { - vec_t vec_A[4], vec_B[8] , vec_C[4]; - acc_t acc_0, acc_1; - __builtin_mma_xxsetaccz(&acc_0); - __builtin_mma_xxsetaccz(&acc_1); - for (int l = 0; l < k; l+=8) { - packNormal((A+(ii*lda)+l), lda, 4, 8, (uint8_t*)vec_A); - packNormal((B+(jj*ldb)+l), ldb, 8, 8, (uint8_t*)vec_B); - for (int x = 0; x < 4; x++) { - mma_instr::outer_product(&acc_0, vec_A[x], vec_B[x]); - mma_instr::outer_product(&acc_1, vec_A[x], vec_B[x+4]); - } - } - SAVE_ACC(&acc_0, ii, jj); - SAVE_ACC(&acc_1, ii, jj+4); - } - - void KERNEL_8x4(int64_t ii, int64_t jj) { - vec_t vec_A[8], vec_B[4] , vec_C[4]; - acc_t acc_0, acc_1; - __builtin_mma_xxsetaccz(&acc_0); - __builtin_mma_xxsetaccz(&acc_1); - for (int l = 0; l < k; l+=8) { - packNormal((A+(ii*lda)+l), lda, 8, 8, (uint8_t*)vec_A); - packNormal((B+(jj*ldb)+l), ldb, 8, 4, (uint8_t*)vec_B); - for (int x = 0; x < 4; x++) { - mma_instr::outer_product(&acc_0, vec_A[x], vec_B[x]); - mma_instr::outer_product(&acc_1, vec_A[x+4], vec_B[x]); - } - } - SAVE_ACC(&acc_0, ii, jj); - SAVE_ACC(&acc_1, ii+4, jj); - } - - - void KERNEL_8x8(int64_t ii, int64_t jj) { - vec_t vec_A[8], vec_B[8], vec_C[4]; - acc_t acc_0, acc_1, acc_2, acc_3; - __builtin_mma_xxsetaccz(&acc_0); - __builtin_mma_xxsetaccz(&acc_1); - __builtin_mma_xxsetaccz(&acc_2); - __builtin_mma_xxsetaccz(&acc_3); - for (int l = 0; l < k; l+=8) { - packNormal(A+(ii*lda)+l, lda, 8, 8, (uint8_t*)vec_A); - packNormal(B+(jj*ldb)+l, ldb, 8, 8, (uint8_t*)vec_B); - for (int x = 0; x < 4; x++) { - mma_instr::outer_product(&acc_0, vec_A[x], vec_B[x]); - mma_instr::outer_product(&acc_1, vec_A[x], vec_B[x+4]); - mma_instr::outer_product(&acc_2, vec_A[x+4], vec_B[x]); - mma_instr::outer_product(&acc_3, vec_A[x+4], vec_B[x+4]); - } - } - - SAVE_ACC(&acc_0, ii, jj); - SAVE_ACC(&acc_1, ii, jj+4); - SAVE_ACC(&acc_2, ii+4, jj); - SAVE_ACC(&acc_3, ii+4, jj+4); - } - - template - void gemm_small(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int64_t ytiles = (m - m0) / RM; - int64_t xtiles = (n - n0) / RN; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - if (end > tiles) - end = tiles; - for (int64_t job = start; job < end; ++job) { - int64_t ii = m0 + job / xtiles * RM; - int64_t jj = n0 + job % xtiles * RN; - vec_t vec_C[4]; - acc_t acc_0; - __builtin_mma_xxsetaccz(&acc_0); - vec_t vec_A[2], vec_B[2]; - for (int l=0; l::outer_product(&acc_0, vec_A[x], vec_B[x]); - } - } - __builtin_mma_disassemble_acc(vec_C, &acc_0); - for (int I = 0; I < RM; I++) { - for (int J = 0; J < RN; J++) { - *((TC*)(C+ii+((jj+J)*ldc)+I)) = *((TC*)&vec_C[I]+J); - } - } - } - } - - template - void gemm_Mx8(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int RN = 8; - int64_t ytiles = (m - m0) / RM; - int64_t xtiles = (n - n0) / RN; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - if (end > tiles) - end = tiles; - for (int64_t job = start; job < end; ++job) { - int64_t ii = m0 + job / xtiles * RM; - int64_t jj = n0 + job % xtiles * RN; - vec_t vec_C[4]; - acc_t acc_0, acc_1; - __builtin_mma_xxsetaccz(&acc_0); - __builtin_mma_xxsetaccz(&acc_1); - vec_t vec_A[4], vec_B[8]; - for (int l=0; l::outer_product(&acc_0, vec_A[x], vec_B[x]); - mma_instr::outer_product(&acc_1, vec_A[x], vec_B[x+4]); - } - } - __builtin_mma_disassemble_acc(vec_C, &acc_0); - for (int I = 0; I < RM; I++) { - for (int J = 0; J < 4; J++) { - *((TC*)(C+ii+((jj+J)*ldc)+I)) = *((TC*)&vec_C[I]+J); - } - } - __builtin_mma_disassemble_acc(vec_C, &acc_1); - for (int I = 0; I < RM; I++) { - for (int J = 0; J < 4; J++) { - *((TC*)(C+ii+((jj+4+J)*ldc)+I)) = *((TC*)&vec_C[I]+J); - } - } - } - } - - template - inline void kernel(int64_t ii, int64_t jj) { - if constexpr(RM == 4 && RN == 8) { - KERNEL_4x8(ii,jj); - } else if constexpr(RM == 8 && RN == 8) { - KERNEL_8x8(ii,jj); - } else if constexpr(RM == 8 && RN == 4) { - KERNEL_8x4(ii,jj); - } else { - assert(false && "RN/RM values not supported"); - } - } - - template - NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int64_t ytiles = (m - m0) / RM; - int64_t xtiles = (n - n0) / RN; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - if (end > tiles) - end = tiles; - for (int64_t job = start; job < end; ++job) { - int64_t ii = m0 + job / xtiles * RM; - int64_t jj = n0 + job % xtiles * RN; - kernel(ii, jj); - } - } - - const TA *const A; - const TB *const B; - TC *C; - const int64_t k; - const int64_t lda; - const int64_t ldb; - const int64_t ldc; - const int ith; - const int nth; -}; - -template -class tinyBLAS_Q0_PPC { - public: - tinyBLAS_Q0_PPC(int64_t k, - const TA * A, int64_t lda, - const block_q8_0 * B, int64_t ldb, - float * C, int64_t ldc, - int ith, int nth) - : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { - } - - void matmul(int64_t m, int64_t n) { - #if defined(_AIX) || defined(__BIG_ENDIAN__) - mnpack(0, m, 0, n); - #else - const int64_t mc = 64; - const int64_t kc = 64; - int64_t nc = 64; - int64_t n_aligned = 0; - if (n % 64 == 0) { - n_aligned = n; - } else if (n == 4) { - n_aligned = 4; - } else if (n < 64) { - n_aligned = (n / 8) * 8; - } else { - n_aligned = (n / 64) * 64; - } - if (n_aligned > 0) { - if (n_aligned % 64 == 0) nc = 64; - else if (n_aligned == n) nc = n; - else if (n_aligned % 32 == 0) nc = 32; - else if (n_aligned % 24 == 0) nc = 24; - else if (n_aligned % 16 == 0) nc = 16; - else nc = 8; - } - bool can_use_tiled = n_aligned > 0 && (m % mc == 0) && (k % kc == 0); - if (can_use_tiled) { - matmul_tiled(m, n_aligned, mc, nc, kc); - if (n > n_aligned) { - mnpack(0, m, n_aligned, n); - } - } else { - mnpack(0, m, 0, n); - } - #endif - } - - private: - inline void save_res(int ii, int jj, int idx, vector float * fin_res, int RM = 4, int RN = 4) { - for (int I = 0; I < RM; I++) { - for (int J = 0; J < RN; J++) { - *((float *)(C + ii + ((jj + J) * ldc) + I)) = *((float *)&fin_res[idx + I] + J); - } - } - } - - inline void save_acc(acc_t * ACC, int64_t ii, int64_t jj) { - vec_t vec_C[4]; - __builtin_mma_disassemble_acc(vec_C, ACC); - for (int I = 0; I < 4; I++) { - for (int J = 0; J < 4; J++) { - *((float *)(C + ii + ((jj + J) * ldc) + I)) = *((float *)&vec_C[I] + J); - } - } - } - - inline void add_save_acc(acc_t * ACC, int64_t ii, int64_t jj) { - vec_t vec_C[4]; - __builtin_mma_disassemble_acc(vec_C, ACC); - for (int I = 0; I < 4; I++) { - for (int J = 0; J < 4; J++) { - float * c_ptr = (float *)(C + ii+ ((jj + J) * ldc) + I); - *c_ptr += *((float *)&vec_C[I] + J); - } - } - } - - template - inline void compute(acc_t * ACC, int c_idx, int s_idx, ArrayType & comparray, vector float * vs, vector float * fin_res) { - vector signed int vec_C[4]; - vector float CA[4] = {0}; - vector float res[4] = {0}; - __builtin_mma_disassemble_acc(vec_C, ACC); - for (int i = 0; i < 4; i++) { - CA[i] = vec_splats((float)(((double)comparray[c_idx + i]) * -128.0)); - res[i] = vec_add(vec_ctf(vec_C[i], 0), CA[i]); - fin_res[s_idx + i] = vec_madd(res[i], vs[s_idx + i], fin_res[s_idx + i]); - } - } - - inline void process_q4_elements(vector signed char (&c)[2], int * ca) { - const vector signed char lowMask = vec_splats((signed char)0xF); - const vector unsigned char v4 = vec_splats((unsigned char)0x4); - const vector signed char v8 = vec_splats((signed char)0x8); - vector signed int vsum = {0}; - vector signed int vsum2 = {0}; - c[0] = vec_and(c[1], lowMask); - c[1] = vec_sr(c[1], v4); - c[0] = vec_sub(c[0], v8); - c[1] = vec_sub(c[1], v8); - vsum = vec_sum4s(c[0], vsum); - vsum2 = vec_sum4s(c[1], vsum2); - vsum = vec_add(vsum, vsum2); - *(ca) = vsum[0] + vsum[1] + vsum[2] + vsum[3]; - } - - template - inline void vector_permute_store(V2 & s1, V2 & s2, V2 & s3, V2 & s4, V1 * vecOffset, bool flip) { - vector unsigned char swiz1 = {0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23}; - vector unsigned char swiz2 = {8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31}; - vector unsigned char swiz3 = {0, 1, 2, 3, 8, 9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27}; - vector unsigned char swiz4 = {4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31}; - V2 t1, t2, t3, t4, t5, t6, t7, t8; - vector unsigned char xor_vector; - uint8_t flip_vec = 0x80; - xor_vector = vec_splats(flip_vec); - t1 = vec_perm(s1, s2, swiz1); - t2 = vec_perm(s1, s2, swiz2); - t3 = vec_perm(s3, s4, swiz1); - t4 = vec_perm(s3, s4, swiz2); - t5 = vec_perm(t1, t3, swiz3); - t6 = vec_perm(t1, t3, swiz4); - t7 = vec_perm(t2, t4, swiz3); - t8 = vec_perm(t2, t4, swiz4); - if (flip == true) { - t5 = vec_xor(t5, xor_vector); - t6 = vec_xor(t6, xor_vector); - t7 = vec_xor(t7, xor_vector); - t8 = vec_xor(t8, xor_vector); - } - vec_xst(t5, 0, vecOffset); - vec_xst(t6, 0, vecOffset + 16); - vec_xst(t7, 0, vecOffset + 32); - vec_xst(t8, 0, vecOffset + 48); - } - - inline void unpack_q4_to_q8(vector signed char packed, vector signed char & lo, vector signed char & hi) { - const vector signed char lowMask = vec_splats((signed char)0x0F); - const vector signed char v8 = vec_splats((signed char)0x08); - const vector unsigned char v4 = vec_splats((unsigned char)4); - lo = vec_and(packed, lowMask); - hi = vec_sr(packed, v4); - lo = vec_sub(lo, v8); - hi = vec_sub(hi, v8); - } - - inline void vector_permute_store_fp16(vec_t * c, unsigned char * vecOffset) { - vec_t t[8], s[8]; - vec_t swiz1 = {0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23}; - vec_t swiz2 = {8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31}; - vec_t swiz3 = {0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23}; - vec_t swiz4 = {8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31}; - for (int i = 0; i < 4; i += 2) { - t[i + 0] = vec_perm(c[i + 0], c[i + 1], swiz1); - t[i + 1] = vec_perm(c[i + 0], c[i + 1], swiz2); - } - for (int i = 4; i < 8; i += 2) { - t[i + 0] = vec_perm(c[i + 0], c[i + 1], swiz1); - t[i + 1] = vec_perm(c[i + 0], c[i + 1], swiz2); - } - s[0] = vec_perm(t[0], t[2], swiz3); - s[1] = vec_perm(t[0], t[2], swiz4); - s[2] = vec_perm(t[1], t[3], swiz3); - s[3] = vec_perm(t[1], t[3], swiz4); - s[4] = vec_perm(t[4], t[6], swiz3); - s[5] = vec_perm(t[4], t[6], swiz4); - s[6] = vec_perm(t[5], t[7], swiz3); - s[7] = vec_perm(t[5], t[7], swiz4); - for (int i = 0; i < 8; ++i) { - vec_xst(s[i], 0, (vec_t *)(vecOffset + i * 16)); - } - } - - static inline void convert_and_scale_q8(vector signed char raw, vector float v_scale, vector unsigned short & out_hi, vector unsigned short & out_lo) { - vector signed short i16_hi = vec_unpackh(raw); - vector signed short i16_lo = vec_unpackl(raw); - - vector float f_hi_h = vec_ctf(vec_unpackh(i16_hi), 0); - vector float f_hi_l = vec_ctf(vec_unpackl(i16_hi), 0); - vector float f_lo_h = vec_ctf(vec_unpackh(i16_lo), 0); - vector float f_lo_l = vec_ctf(vec_unpackl(i16_lo), 0); - out_hi = vec_pack_to_short_fp32(vec_mul(f_hi_h, v_scale), vec_mul(f_hi_l, v_scale)); - out_lo = vec_pack_to_short_fp32(vec_mul(f_lo_h, v_scale), vec_mul(f_lo_l, v_scale)); - } - - void packNormal_q4_fp16(const block_q4_0 * a, int64_t lda, int rows, int blocks, unsigned char * vec) { - unsigned char * vecOffset = vec; - for (int i = 0; i < rows; i += 8) { - const block_q4_0 * rows_base[8]; - for (int r = 0; r < 8; r++) { - rows_base[r] = a + (i + r) * lda; - } - for (int blk = 0; blk < blocks; blk++) { - vector unsigned short hp_res[8][4]; - for (int r = 0; r < 8; r++) { - const block_q4_0 * current_blk = rows_base[r] + blk; - vector float v_scale = vec_extract_fp32_from_shorth(vec_splats(current_blk->d)); - vector signed char v_qs = vec_xl(0, (const vector signed char *)current_blk->qs); - vector signed char c1, c2; - unpack_q4_to_q8(v_qs, c1, c2); - convert_and_scale_q8(c1, v_scale, hp_res[r][0], hp_res[r][1]); - convert_and_scale_q8(c2, v_scale, hp_res[r][2], hp_res[r][3]); - } - for (int c = 0; c < 4; c++) { - vector unsigned char c_arr[8]; - for (int r = 0; r < 8; r++) { - c_arr[r] = (vector unsigned char)hp_res[r][c]; - } - vector_permute_store_fp16((vec_t *)c_arr, vecOffset); - vecOffset += 128; - } - } - } - } - - template - static inline void pack_q8_block(const block_q8_0 * a, int64_t lda, int rows, int blocks, unsigned char * vec) { - unsigned char * vecOffset = vec; - const vec_t swiz1 = {0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23}; - const vec_t swiz2 = {8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31}; - const vec_t swiz3 = {0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23}; - const vec_t swiz4 = {8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31}; - - for (int i = 0; i < rows; i += chunk_size) { - const block_q8_0 * rows_base[chunk_size]; - for (int r = 0; r < chunk_size; r++) { - rows_base[r] = a + (i + r) * lda; - } - for (int blk = 0; blk < blocks; blk++) { - vector unsigned short hp_res[chunk_size][4]; - for (int r = 0; r < chunk_size; r++) { - const block_q8_0 * b = rows_base[r] + blk; - vector float v_scale = vec_extract_fp32_from_shorth(vec_splats(b->d)); - vector signed char c[2]; - __vector_pair pair = __builtin_vsx_lxvp(0, (__vector_pair *)b->qs); - __builtin_vsx_disassemble_pair(c, & pair); - convert_and_scale_q8(c[0], v_scale, hp_res[r][0], hp_res[r][1]); - convert_and_scale_q8(c[1], v_scale, hp_res[r][2], hp_res[r][3]); - } - for (int col = 0; col < 4; col++) { - if constexpr (chunk_size == 8) { - vec_t t[8]; - t[0] = vec_perm((vec_t)hp_res[0][col], (vec_t)hp_res[1][col], swiz1); - t[1] = vec_perm((vec_t)hp_res[0][col], (vec_t)hp_res[1][col], swiz2); - t[2] = vec_perm((vec_t)hp_res[2][col], (vec_t)hp_res[3][col], swiz1); - t[3] = vec_perm((vec_t)hp_res[2][col], (vec_t)hp_res[3][col], swiz2); - t[4] = vec_perm((vec_t)hp_res[4][col], (vec_t)hp_res[5][col], swiz1); - t[5] = vec_perm((vec_t)hp_res[4][col], (vec_t)hp_res[5][col], swiz2); - t[6] = vec_perm((vec_t)hp_res[6][col], (vec_t)hp_res[7][col], swiz1); - t[7] = vec_perm((vec_t)hp_res[6][col], (vec_t)hp_res[7][col], swiz2); - - vec_xst(vec_perm(t[0], t[2], swiz3), 0, (vec_t *)(vecOffset + 0)); - vec_xst(vec_perm(t[0], t[2], swiz4), 0, (vec_t *)(vecOffset + 16)); - vec_xst(vec_perm(t[1], t[3], swiz3), 0, (vec_t *)(vecOffset + 32)); - vec_xst(vec_perm(t[1], t[3], swiz4), 0, (vec_t *)(vecOffset + 48)); - vec_xst(vec_perm(t[4], t[6], swiz3), 0, (vec_t *)(vecOffset + 64)); - vec_xst(vec_perm(t[4], t[6], swiz4), 0, (vec_t *)(vecOffset + 80)); - vec_xst(vec_perm(t[5], t[7], swiz3), 0, (vec_t *)(vecOffset + 96)); - vec_xst(vec_perm(t[5], t[7], swiz4), 0, (vec_t *)(vecOffset + 112)); - vecOffset += 128; - } else { - vec_t t0 = vec_perm((vec_t)hp_res[0][col], (vec_t)hp_res[1][col], swiz1); - vec_t t1 = vec_perm((vec_t)hp_res[0][col], (vec_t)hp_res[1][col], swiz2); - vec_t t2 = vec_perm((vec_t)hp_res[2][col], (vec_t)hp_res[3][col], swiz1); - vec_t t3 = vec_perm((vec_t)hp_res[2][col], (vec_t)hp_res[3][col], swiz2); - - vec_xst(vec_perm(t0, t2, swiz3), 0, (vec_t *)(vecOffset + 0)); - vec_xst(vec_perm(t0, t2, swiz4), 0, (vec_t *)(vecOffset + 16)); - vec_xst(vec_perm(t1, t3, swiz3), 0, (vec_t *)(vecOffset + 32)); - vec_xst(vec_perm(t1, t3, swiz4), 0, (vec_t *)(vecOffset + 48)); - vecOffset += 64; - } - } - } - } - } - - void packNormal_q8_fp16(const block_q8_0 * a, int64_t lda, int rows, int blocks, unsigned char * vec) { - if (rows == 4) { - pack_q8_block<4>(a, lda, rows, blocks, vec); - } else { - pack_q8_block<8>(a, lda, rows, blocks, vec); - } - } - - template - void packNormalInt4(const TA * a, int64_t lda, int rows, int cols, int8_t * vec, std::array & comparray) { - int64_t i, j; - TA * aoffset = NULL; - int8_t * vecOffset = NULL; - TA * aoffset1 = NULL, * aoffset2 = NULL, * aoffset3 = NULL, * aoffset4 = NULL; - TA * aoffset5 = NULL, * aoffset6 = NULL, * aoffset7 = NULL, * aoffset8 = NULL; - vector signed char c1[2] = {0}, c2[2] = {0}, c3[2] = {0}, c4[2] = {0}; - vector signed char c5[2] = {0}, c6[2] = {0}, c7[2] = {0}, c8[2] = {0}; - aoffset = const_cast(a); - vecOffset = vec; - j = (rows >> 3); - if (j > 0) { - do { - aoffset1 = aoffset; - aoffset2 = aoffset1 + lda; - aoffset3 = aoffset2 + lda; - aoffset4 = aoffset3 + lda; - aoffset5 = aoffset4 + lda; - aoffset6 = aoffset5 + lda; - aoffset7 = aoffset6 + lda; - aoffset8 = aoffset7 + lda; - aoffset += 8 * lda; - i = (cols >> 2); - if (i > 0) { - do { - c1[1] = vec_xl(0, (const vector signed char *)aoffset1->qs); - c2[1] = vec_xl(0, (const vector signed char *)aoffset2->qs); - c3[1] = vec_xl(0, (const vector signed char *)aoffset3->qs); - c4[1] = vec_xl(0, (const vector signed char *)aoffset4->qs); - c5[1] = vec_xl(0, (const vector signed char *)aoffset5->qs); - c6[1] = vec_xl(0, (const vector signed char *)aoffset6->qs); - c7[1] = vec_xl(0, (const vector signed char *)aoffset7->qs); - c8[1] = vec_xl(0, (const vector signed char *)aoffset8->qs); - - process_q4_elements(c1, & comparray[0]); - process_q4_elements(c2, & comparray[1]); - process_q4_elements(c3, & comparray[2]); - process_q4_elements(c4, & comparray[3]); - process_q4_elements(c5, & comparray[4]); - process_q4_elements(c6, & comparray[5]); - process_q4_elements(c7, & comparray[6]); - process_q4_elements(c8, & comparray[7]); - vector_permute_store(c1[0], c2[0], c3[0], c4[0], vecOffset, false); - vector_permute_store(c1[1], c2[1], c3[1], c4[1], vecOffset + 64, false); - vector_permute_store(c5[0], c6[0], c7[0], c8[0], vecOffset + 128, false); - vector_permute_store(c5[1], c6[1], c7[1], c8[1], vecOffset + 192, false); - aoffset1 += lda; - aoffset2 += lda; - aoffset3 += lda; - aoffset4 += lda; - aoffset5 += lda; - aoffset6 += lda; - aoffset7 += lda; - aoffset8 += lda; - vecOffset += 256; - i--; - } while (i > 0); - } - j--; - } while (j > 0); - } - - if (rows & 4) { - aoffset1 = aoffset; - aoffset2 = aoffset1 + lda; - aoffset3 = aoffset2 + lda; - aoffset4 = aoffset3 + lda; - aoffset += 4 * lda; - i = (cols >> 2); - if (i > 0) { - do { - c1[1] = vec_xl(0, (const vector signed char *)aoffset1->qs); - c2[1] = vec_xl(0, (const vector signed char *)aoffset2->qs); - c3[1] = vec_xl(0, (const vector signed char *)aoffset3->qs); - c4[1] = vec_xl(0, (const vector signed char *)aoffset4->qs); - - process_q4_elements(c1, & comparray[0]); - process_q4_elements(c2, & comparray[1]); - process_q4_elements(c3, & comparray[2]); - process_q4_elements(c4, & comparray[3]); - vector_permute_store(c1[0], c2[0], c3[0], c4[0], vecOffset, false); - vector_permute_store(c1[1], c2[1], c3[1], c4[1], vecOffset + 64, false); - aoffset1 += lda; - aoffset2 += lda; - aoffset3 += lda; - aoffset4 += lda; - vecOffset += 128; - i--; - } while (i > 0); - } - } - - if (rows & 3) { - aoffset1 = aoffset; - aoffset2 = aoffset1 + lda; - aoffset3 = aoffset2 + lda; - i = (cols >> 2); - if (i > 0) { - do { - switch(rows) { - case 3: c3[1] = vec_xl(0, (const vector signed char *)aoffset3->qs); - case 2: c2[1] = vec_xl(0, (const vector signed char *)aoffset2->qs); - case 1: c1[1] = vec_xl(0, (const vector signed char *)aoffset1->qs); - break; - } - process_q4_elements(c1, & comparray[0]); - process_q4_elements(c2, & comparray[1]); - process_q4_elements(c3, & comparray[2]); - process_q4_elements(c4, & comparray[3]); - vector_permute_store(c1[0], c2[0], c3[0], c4[0], vecOffset, false); - vector_permute_store(c1[1], c2[1], c3[1], c4[1], vecOffset + 64, false); - aoffset1 += lda; - aoffset2 += lda; - aoffset3 += lda; - vecOffset += 128; - i--; - } while(i > 0); - } - } - } - - template - void packNormal(const block_q8_0 * a, int64_t lda, int rows, int cols, VA * vec, bool flip) { - int64_t i, j; - block_q8_0 * aoffset = NULL; - VA * vecOffset = NULL; - block_q8_0 * aoffsets[8]; - __vector_pair arr[8]; - VB c[8][2] = {0}; - VB c1[8] = {0}; VB c2[8] = {0}; - aoffset = const_cast(a); - vecOffset = vec; - j = (rows >> 3); - if (j > 0) { - do { - aoffsets[0] = aoffset; - for (int it = 1; it < 8; it++) - aoffsets[it] = aoffsets[it - 1] + lda; - aoffset += 8 * lda; - - i = (cols >> 3); - if (i > 0) { - do { - for (int it = 0; it < 8; it++) { - arr[it] = __builtin_vsx_lxvp(0, (__vector_pair *)aoffsets[it]->qs); - __builtin_vsx_disassemble_pair(c[it], & arr[it]); - c1[it] = c[it][0]; - c2[it] = c[it][1]; - } - vector_permute_store(c1[0], c1[1], c1[2], c1[3], vecOffset, flip); - vector_permute_store(c2[0], c2[1], c2[2], c2[3], vecOffset + 64, flip); - vector_permute_store(c1[4], c1[5], c1[6], c1[7], vecOffset + 128, flip); - vector_permute_store(c2[4], c2[5], c2[6], c2[7], vecOffset + 192, flip); - for (int it = 0; it < 8; it++) - aoffsets[it] += lda; - vecOffset += 256; - i--; - } while(i > 0); - } - j--; - } while(j > 0); - } - if (rows & 4) { - aoffsets[0] = aoffset; - for (int it = 1; it < 4; it++ ) - aoffsets[it] = aoffsets[it-1] + lda; - aoffset += 4 * lda; - i = (cols >> 3); - if (i > 0) { - do { - for (int it = 0; it < 4; it++) { - arr[it] = __builtin_vsx_lxvp(0, (__vector_pair *)aoffsets[it]->qs); - __builtin_vsx_disassemble_pair(c[it], & arr[it]); - c1[it] = c[it][0]; - c2[it] = c[it][1]; - } - vector_permute_store(c1[0], c1[1], c1[2], c1[3], vecOffset, flip); - vector_permute_store(c2[0], c2[1], c2[2], c2[3], vecOffset + 64, flip); - for (int it = 0; it < 4; it++) { - aoffsets[it] += lda; - } - vecOffset += 128; - i--; - } while(i > 0); - } - } - - if (rows & 3) { - aoffsets[0] = aoffset; - for (int it = 1; it < 3; it++ ) - aoffsets[it] = aoffsets[it - 1] + lda; - i = (cols >> 3); - if (i > 0) { - do { - switch(rows) { - case 3: arr[2] = __builtin_vsx_lxvp(0, (__vector_pair *)aoffsets[2]->qs); - __builtin_vsx_disassemble_pair(c[2], & arr[2]); - c1[2] = c[2][0]; c2[2] = c[2][1]; - case 2: arr[1] = __builtin_vsx_lxvp(0, (__vector_pair *)aoffsets[1]->qs); - __builtin_vsx_disassemble_pair(c[1], & arr[1]); - c1[1] = c[1][0]; c2[1] = c[1][1]; - case 1: arr[0] = __builtin_vsx_lxvp(0, (__vector_pair *)aoffsets[0]->qs); - __builtin_vsx_disassemble_pair(c[0], & arr[0]); - c1[0] = c[0][0]; c2[0] = c[0][1]; - break; - } - vector_permute_store(c1[0], c1[1], c1[2], c1[3], vecOffset, flip); - vector_permute_store(c2[0], c2[1], c2[2], c2[3], vecOffset + 64, flip); - for (int it = 0; it < 3; it++) - aoffsets[it] += lda; - vecOffset += 128; - i--; - } while(i > 0); - } - } - } - - void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int m_rem = MIN(m - m0, 16); - int n_rem = MIN(n - n0, 16); - - int mc = 0, nc = 0; - - if (m_rem >= 8 && n_rem >= 8) { - mc = 8; - nc = 8; - gemm<8, 8>(m0, m, n0, n); - } else if (m_rem >= 4 && n_rem >= 8) { - mc = 4; - nc = 8; - gemm<4, 8>(m0, m, n0, n); - } else if (m_rem >= 8 && n_rem >= 4) { - mc = 8; - nc = 4; - gemm<8, 4>(m0, m, n0, n); - } else if (m_rem >= 4 && n_rem >= 4) { - mc = 4; - nc = 4; - gemm_small(m0, m, n0, n, mc, nc); - } else { - mc = (m_rem >= 4) ? 4 : m_rem; - nc = (n_rem >= 4) ? 4 : n_rem; - if (mc == 0 || nc == 0) - return; - gemm_small(m0, m, n0, n, mc, nc); - } - - int64_t mp = m0 + ((m - m0) / mc) * mc; - int64_t np = n0 + ((n - n0) / nc) * nc; - mnpack(mp, m, n0, np); - mnpack(m0, m, np, n); - } - - - void KERNEL_4x8(int64_t ii, int64_t jj) { - vec_t vec_A[8], vec_B[16] = {0}; - acc_t acc_0, acc_1; - std::array comparray {}; - vector float fin_res[8] = {0}; - vector float vs[8] = {0}; - bool isAblock_q4 = std::is_same_v; - for (int l = 0; l < k; l++) { - __builtin_mma_xxsetaccz(& acc_0); - __builtin_mma_xxsetaccz(& acc_1); - if (std::is_same_v) { - packNormalInt4<4>((A + (ii * lda) + l), lda, 4, 4, (int8_t *)vec_A, comparray); - } else { - packNormal((const block_q8_0 *)(A + (ii * lda) + l), lda, 4, 8, (int8_t *)vec_A, false); - } - packNormal((B + (jj * ldb) + l), ldb, 8, 8, (uint8_t *)vec_B, true); - for(int x = 0; x < 8; x++) { - __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x], vec_B[x]); - __builtin_mma_xvi8ger4pp(& acc_1, vec_A[x], vec_B[x+8]); - } - for (int I = 0; I<4; I++) { - for (int J = 0; J<4; J++) { - *((float *)& vs[I] + J) = (unhalf((A + ((ii + I) * lda) + l)->d) * unhalf((B + ((jj + J) * ldb) + l)->d)); - *((float *)& vs[I + 4] + J) = (unhalf((A +((ii + I) * lda) + l)->d) * unhalf((B + ((jj + J + 4) * ldb) + l)->d)); - } - } - if (!isAblock_q4) { - auto aoffset = A + (ii * lda) + l; - for (int i = 0; i < 4; i++) { - comparray[i] = 0; - int ca = 0; - auto *at = aoffset->qs; - for (int j = 0; j < 32; j++) - ca += (int)*at++; - comparray[i] = ca; - aoffset += lda; - } - } - compute(& acc_0, 0, 0, comparray, vs, fin_res); - compute(& acc_1, 0, 4, comparray, vs, fin_res); - } - save_res(ii, jj, 0, fin_res); - save_res(ii, jj + 4, 4, fin_res); - } - - void KERNEL_8x4(int64_t ii, int64_t jj) { - vec_t vec_A[16], vec_B[8] = {0}; - acc_t acc_0, acc_1; - std::array comparray {}; - vector float fin_res[8] = {0}; - vector float vs[8] = {0}; - bool isAblock_q4 = std::is_same_v; - for (int l = 0; l < k; l++) { - __builtin_mma_xxsetaccz(& acc_0); - __builtin_mma_xxsetaccz(& acc_1); - if (std::is_same_v) { - packNormalInt4<8>((A + (ii * lda) + l), lda, 8, 4, (int8_t *)vec_A, comparray); - } else { - packNormal((const block_q8_0 *)(A + (ii * lda) + l), lda, 8, 8, (int8_t *)vec_A, false); - } - packNormal((B + (jj * ldb) + l), ldb, 4, 8, (uint8_t *)vec_B, true); - for(int x = 0; x < 8; x++) { - __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x], vec_B[x]); - __builtin_mma_xvi8ger4pp(& acc_1, vec_A[x + 8], vec_B[x]); - } - for (int I = 0; I < 8; I++) { - for (int J = 0; J < 4; J++) { - *((float *)&vs[I] + J) = (unhalf((A + ((ii + I) * lda) + l)->d) * unhalf((B + ((jj + J) * ldb) + l)->d)); - } - } - if (!isAblock_q4) { - auto aoffset = A + (ii * lda) + l; - for (int i = 0; i < 8; i++) { - comparray[i] = 0; - int ca = 0; - auto *at = aoffset->qs; - for (int j = 0; j < 32; j++) - ca += (int)*at++; - comparray[i] = ca; - aoffset += lda; - } - } - compute(& acc_0, 0, 0, comparray, vs, fin_res); - compute(& acc_1, 4, 4, comparray, vs, fin_res); - } - save_res(ii, jj, 0, fin_res); - save_res(ii + 4, jj, 4, fin_res); - } - - void KERNEL_8x8(int64_t ii, int64_t jj) { - vec_t vec_A[16], vec_B[16] = {0}; - acc_t acc_0, acc_1, acc_2, acc_3; - acc_t acc_4, acc_5, acc_6, acc_7; - std::array comparray {}; - vector float fin_res[16] = {0}; - vector float vs[16] = {0}; - bool isAblock_q4 = std::is_same_v; - for (int l = 0; l < k; l++) { - __builtin_mma_xxsetaccz(& acc_0); - __builtin_mma_xxsetaccz(& acc_1); - __builtin_mma_xxsetaccz(& acc_2); - __builtin_mma_xxsetaccz(& acc_3); - if (std::is_same_v) { - packNormalInt4<8>((A + (ii * lda) + l), lda, 8, 4, (int8_t *)vec_A, comparray); - } else { - packNormal((const block_q8_0 *)(A + (ii * lda) + l), lda, 8, 8, (int8_t *)vec_A, false); - } - packNormal((B + (jj * ldb) + l), ldb, 8, 8, (uint8_t *)vec_B, true); - for(int x = 0; x < 8; x++) { - __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x], vec_B[x]); - __builtin_mma_xvi8ger4pp(& acc_1, vec_A[x + 8], vec_B[x]); - __builtin_mma_xvi8ger4pp(& acc_2, vec_A[x], vec_B[x + 8]); - __builtin_mma_xvi8ger4pp(& acc_3, vec_A[x + 8], vec_B[x + 8]); - } - for (int I = 0; I < 8 ; I++) { - for (int J = 0; J < 4; J++) { - *((float *)& vs[I] + J) = (unhalf((A + ((ii + I) * lda) + l)->d) * unhalf((B + ((jj + J) * ldb) + l)->d)); - *((float *)& vs[I + 8] + J) = (unhalf((A + ((ii + I) * lda) + l)->d) * unhalf((B + ((jj + J + 4) * ldb) + l)->d)); - } - } - if (!isAblock_q4) { - auto aoffset = A + (ii * lda) + l; - for (int i = 0; i < 8; i++) { - comparray[i] = 0; - int ca = 0; - auto *at = aoffset->qs; - for (int j = 0; j < 32; j++) - ca += (int)*at++; - comparray[i] = ca; - aoffset += lda; - } - } - compute(& acc_0, 0, 0, comparray, vs, fin_res); - compute(& acc_1, 4, 4, comparray, vs, fin_res); - compute(& acc_2, 0, 8, comparray, vs, fin_res); - compute(& acc_3, 4, 12, comparray, vs, fin_res); - } - save_res(ii, jj, 0, fin_res); - save_res(ii + 4, jj, 4, fin_res); - save_res(ii, jj + 4, 8, fin_res); - save_res(ii + 4, jj + 4, 12, fin_res); - } - - void KERNEL_Q0(int64_t ii, int64_t jj, int64_t mc, int64_t nc, int64_t kc, int64_t l, vec_t * vec_A, vec_t * vec_B) { - acc_t acc[8]; - for (int i = 0; i < mc ; i += 16) { - for (int j = 0; j < nc; j += 8) { - int A0_base = (i / 16) * (2 * 32 * kc); - int B0_base = (j / 8) * (32 * kc); - for (int x = 0; x < 8; x++) { - __builtin_mma_xxsetaccz(&acc[x]); - } - for (int64_t kk = 0; kk < kc; kk++) { - int A0_block_idx = A0_base + kk * 32; - int B0_block_idx = B0_base + kk * 32; - int A1_block_idx = A0_block_idx + 32 * kc; - int B1_block_idx = B0_block_idx + 32 * kc; - vec_t * A0_block = & vec_A[A0_block_idx]; - vec_t * B0_block = & vec_B[B0_block_idx]; - vec_t * A1_block = & vec_A[A1_block_idx]; - for (int it = 0; it < 4; it++) { - for (int x = 0; x < 4; x++) { - __builtin_mma_xvf16ger2pp(& acc[0], A0_block[8 * it + x], B0_block[8 * it + x]); - __builtin_mma_xvf16ger2pp(& acc[1], A0_block[8 * it + x], B0_block[8 * it + x + 4]); - __builtin_mma_xvf16ger2pp(& acc[2], A0_block[8 * it + x + 4], B0_block[8 * it + x]); - __builtin_mma_xvf16ger2pp(& acc[3], A0_block[8 * it + x + 4], B0_block[8 * it + x + 4]); - __builtin_mma_xvf16ger2pp(& acc[4], A1_block[8 * it + x], B0_block[8 * it + x]); - __builtin_mma_xvf16ger2pp(& acc[5], A1_block[8 * it + x], B0_block[8 * it+ x + 4]); - __builtin_mma_xvf16ger2pp(& acc[6], A1_block[8 * it + x + 4], B0_block[8 * it + x]); - __builtin_mma_xvf16ger2pp(& acc[7], A1_block[8 * it + x + 4], B0_block[8 * it + x + 4]); - } - } - } - if (l == 0) { - save_acc(& acc[0], ii + i, jj + j); - save_acc(& acc[1], ii + i, jj + j + 4); - save_acc(& acc[2], ii + i + 4, jj + j); - save_acc(& acc[3], ii + i + 4, jj + j + 4); - save_acc(& acc[4], ii + i + 8, jj + j); - save_acc(& acc[5], ii + i + 8, jj + j + 4); - save_acc(& acc[6], ii + i + 12, jj + j); - save_acc(& acc[7], ii + i + 12, jj + j + 4); - } else { - add_save_acc(& acc[0], ii + i, jj + j); - add_save_acc(& acc[1], ii + i, jj + j + 4); - add_save_acc(& acc[2], ii + i + 4, jj + j); - add_save_acc(& acc[3], ii + i + 4, jj + j + 4); - add_save_acc(& acc[4], ii + i + 8, jj + j); - add_save_acc(& acc[5], ii + i + 8, jj + j + 4); - add_save_acc(& acc[6], ii + i + 12, jj + j); - add_save_acc(& acc[7], ii + i + 12, jj + j + 4); - } - } - } - } - - void matmul_tiled(int64_t m, int64_t n, int64_t mc, int64_t nc, int64_t kc) { - vec_t A_pack[mc * kc * 4]; - vec_t B_pack[nc * kc * 4]; - constexpr bool is_Ablock_q4 = std::is_same_v; - int64_t ytiles = m / mc; - int64_t xtiles = n / nc; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - if (end > tiles) { - end = tiles; - } - for (int64_t job = start; job < end; ++job) { - int64_t ii = (job / xtiles) * mc; - int64_t jj = (job % xtiles) * nc; - for (int64_t kk = 0; kk < k; kk += kc) { - if constexpr(is_Ablock_q4) { - packNormal_q4_fp16(A + ii * lda + kk, lda, mc, kc, (uint8_t *)A_pack); - } else { - packNormal_q8_fp16(A + ii * lda + kk, lda, mc, kc, (uint8_t *)A_pack); - } - packNormal_q8_fp16(B + jj * ldb + kk, ldb, nc, kc, (uint8_t *)B_pack); - KERNEL_Q0(ii, jj, mc, nc, kc, kk, A_pack, B_pack); - } - } - } - - void gemm_small(int64_t m0, int64_t m, int64_t n0, int64_t n, int RM, int RN) { - int64_t ytiles = (m - m0) / RM; - int64_t xtiles = (n - n0) / RN; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - vec_t vec_A[8] = {0}, vec_B[8] = {0}; - vector signed int vec_C[4]; - acc_t acc_0; - bool isAblock_q4 = std::is_same_v; - - if (end > tiles) - end = tiles; - for (int64_t job = start; job < end; ++job) { - int64_t ii = m0 + job / xtiles * RM; - int64_t jj = n0 + job % xtiles * RN; - std::array comparray{}; - vector float res[4] = {0}; - vector float fin_res[4] = {0}; - vector float vs[4] = {0}; - vector float CA[4] = {0}; - __builtin_prefetch((A + (ii * lda) + 0)->qs, 0, 1); // prefetch first value - __builtin_prefetch((B + (jj * ldb) + 0)->qs, 0, 1); // prefetch first value - for (int l = 0; l < k; l++) { - __builtin_prefetch((A + (ii * lda) + (l + 1))->qs, 0, 1); // prefetch one loop ahead - __builtin_prefetch((B + (jj * ldb) + (l + 1))->qs, 0, 1); // prefetch one loop ahead - __builtin_mma_xxsetaccz(& acc_0); - if (isAblock_q4) { - packNormalInt4<4>((A + (ii * lda) + l), lda, RM, 4, (int8_t *)vec_A, comparray); - } else { - packNormal((const block_q8_0 *)(A + (ii * lda) + l), lda, RM, 8, (int8_t *)vec_A, false); - } - packNormal((B + (jj * ldb) + l), ldb, RN, 8, (uint8_t *)vec_B, true); - for (int x = 0; x < 8; x += 4) { - __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x], vec_B[x]); - __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x + 1], vec_B[x + 1]); - __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x + 2], vec_B[x + 2]); - __builtin_mma_xvi8ger4pp(& acc_0, vec_A[x + 3], vec_B[x + 3]); - } - for (int I = 0; I < RM; I++) { - for (int J = 0; J < RN; J++) { - *((float*)&vs[I] + J) = (unhalf((A + ((ii + I) * lda) + l)->d) * unhalf((B + ((jj + J) * ldb) + l)->d)); - } - } - __builtin_mma_disassemble_acc(vec_C, & acc_0); - if (!isAblock_q4) { - auto aoffset = A + (ii * lda) + l; - for (int i = 0; i < RM; i++) { - comparray[i] = 0; - int ca = 0; - auto *at = aoffset->qs; - for (int j = 0; j < 32; j++) - ca += (int)*at++; - comparray[i] = ca; - aoffset += lda; - } - } - for (int i = 0; i < RM; i++) { - CA[i] = vec_splats((float)(((double)comparray[i]) * -128.0)); - res[i] = vec_add(vec_ctf(vec_C[i], 0), CA[i]); - fin_res[i] = vec_madd(res[i], vs[i], fin_res[i]); - } - } - save_res(ii, jj, 0, fin_res, RM, RN); - } - } - - template - inline void kernel(int64_t ii, int64_t jj) { - if constexpr(RM == 4 && RN == 8) { - KERNEL_4x8(ii,jj); - } else if constexpr(RM == 8 && RN == 4) { - KERNEL_8x4(ii,jj); - } else if constexpr(RM == 8 && RN == 8) { - KERNEL_8x8(ii,jj); - } else { - assert(false && "RN/RM values not supported"); - } - } - - template - NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int64_t ytiles = (m - m0) / RM; - int64_t xtiles = (n - n0) / RN; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - if (end > tiles) - end = tiles; - for (int64_t job = start; job < end; ++job) { - int64_t ii = m0 + job / xtiles * RM; - int64_t jj = n0 + job % xtiles * RN; - kernel(ii, jj); - } - } - const TA * const A; - const block_q8_0 * const B; - float * C; - const int64_t k; - int64_t kc; - const int64_t lda; - const int64_t ldb; - const int64_t ldc; - const int ith; - const int nth; -}; - -class tinyBLAS_PPC { - public: - tinyBLAS_PPC(int64_t k, - const float * A, int64_t lda, - const float * B, int64_t ldb, - float * C, int64_t ldc, - int ith, int nth) - : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { - } - - void matmul(int64_t m, int64_t n) { - #if defined(_AIX) || defined(__BIG_ENDIAN__) - mnpack(0, m, 0, n); - #else - int64_t mc = 256; int64_t nc = 256; int64_t kc = 256; - if (m % mc == 0 && n % nc == 0 && k % kc == 0) { - matmul_tiled(m, n, mc, nc, kc); - } else { - mnpack(0, m, 0, n); - } - #endif - } - - private: - - __attribute__((always_inline)) - inline void save_acc(acc_t * ACC, int64_t ii, int64_t jj) { - vec_t vec_C[4]; - __builtin_mma_disassemble_acc(vec_C, ACC); - for (int I = 0; I < 4; I++) { - for (int J = 0; J < 4; J++) { - *((float *)(C+ii+((jj+J)*ldc)+I)) = *((float *)&vec_C[I]+J); - } - } - } - - __attribute__((always_inline)) - inline void add_save_acc(acc_t * ACC, int64_t ii, int64_t jj) { - vec_t vec_C[4]; - __builtin_mma_disassemble_acc(vec_C, ACC); - for (int I = 0; I < 4; I++) { - for (int J = 0; J < 4; J++) { - float * c_ptr = (float *)(C+ii+((jj+J)*ldc)+I); - *c_ptr += *((float *)&vec_C[I]+J); - } - } - } - - inline void vector_permute_store_4(vector float * src, float * vecOffset) { - vector float t1, t2, t3, t4, t5, t6, t7, t8; - t1 = vec_mergeh(src[0], src[1]); - t2 = vec_mergeh(src[2], src[3]); - t3 = vec_mergel(src[0], src[1]); - t4 = vec_mergel(src[2], src[3]); - - t5 = vec_xxpermdi(t1, t2, 0); - t6 = vec_xxpermdi(t1, t2, 3); - t7 = vec_xxpermdi(t3, t4, 0); - t8 = vec_xxpermdi(t3, t4, 3); - - vec_xst(t5, 0, vecOffset); - vec_xst(t6, 0, vecOffset + 4); - vec_xst(t7, 0, vecOffset + 8); - vec_xst(t8, 0, vecOffset + 12); - } - - inline void vector_permute_store_8(vector float * src, float * vecOffset) { - vector float t1, t2, t3, t4, t5, t6, t7, t8; - t1 = vec_mergeh(src[0], src[1]); - t2 = vec_mergeh(src[2], src[3]); - t3 = vec_mergeh(src[4], src[5]); - t4 = vec_mergeh(src[6], src[7]); - - t5 = vec_xxpermdi(t1, t2, 0); - t6 = vec_xxpermdi(t3, t4, 0); - t7 = vec_xxpermdi(t1, t2, 3); - t8 = vec_xxpermdi(t3, t4, 3); - - vec_xst(t5, 0, vecOffset); - vec_xst(t6, 0, vecOffset + 4); - vec_xst(t7, 0, vecOffset + 8); - vec_xst(t8, 0, vecOffset + 12); - - t1 = vec_mergel(src[0], src[1]); - t2 = vec_mergel(src[2], src[3]); - t3 = vec_mergel(src[4], src[5]); - t4 = vec_mergel(src[6], src[7]); - - t5 = vec_xxpermdi(t1, t2, 0); - t6 = vec_xxpermdi(t3, t4, 0); - t7 = vec_xxpermdi(t1, t2, 3); - t8 = vec_xxpermdi(t3, t4, 3); - - vec_xst(t5, 0, vecOffset + 16); - vec_xst(t6, 0, vecOffset + 20); - vec_xst(t7, 0, vecOffset + 24); - vec_xst(t8, 0, vecOffset + 28); - } - - void packTranspose(const float * a, int64_t lda, int rows, int cols, float * vec) { - int64_t i, j; - float * aoffsets[8]; - float * aoffset = NULL, * boffset = NULL; - __vector_pair arr[8]; - vector float c[8][2] = {0}; - vector float c1[8] = {0}; - vector float c2[8] = {0}; - aoffset = const_cast(a); - boffset = vec; - j = (rows >> 3); - if (j > 0) { - do { - aoffsets[0] = aoffset; - for (int it = 1; it < 8; it++) - aoffsets[it] = aoffsets[it-1] + lda; - aoffset += 8 * lda; - i = (cols >> 3); - if (i > 0) { - do { - for (int it = 0; it < 8; it++) { - arr[it] = __builtin_vsx_lxvp(0, (__vector_pair*)aoffsets[it]); - __builtin_vsx_disassemble_pair(c[it], &arr[it]); - c1[it] = c[it][0]; - c2[it] = c[it][1]; - } - - vector_permute_store_8(c1, boffset); - vector_permute_store_8(c2, boffset + 32); - boffset += 64; - i--; - if (i > 0) { - for (int it = 0; it < 8; it++) { - aoffsets[it] = aoffsets[it] + 8; - } - } - } while(i > 0); - } - if (cols & 4) { - for (int it = 0; it < 8 ; it++) - c1[it] = vec_xl(0, aoffsets[it]); - vector_permute_store_8(c1, boffset); - } - j--; - } while(j > 0); - } - - if (rows & 4) { - aoffsets[0] = aoffset; - for (int it = 1; it < 4; it++) - aoffsets[it] = aoffsets[it-1] + lda; - aoffset += 4 * lda; - i = (cols >> 3); - if (i > 0) { - do { - for (int it = 0; it < 4; it++) { - arr[it] = __builtin_vsx_lxvp(0, (__vector_pair*)aoffsets[it]); - __builtin_vsx_disassemble_pair(c[it], &arr[it]); - c1[it] = c[it][0]; - c2[it] = c[it][1]; - } - vector_permute_store_4(c1, boffset); - vector_permute_store_4(c2, boffset + 16); - for (int it = 0; it < 4; it++) - aoffsets[it] += 8 * lda; - boffset += 32; - i--; - } while(i > 0); - } - - if (cols & 4) { - for (int it = 0; it < 4; it++) - c1[it] = vec_xl(0, aoffsets[it]); - vector_permute_store_4(c1, boffset); - } - } - if (rows & 3) { - aoffsets[0] = aoffset; - for (int it = 1; it < 3; it++) - aoffsets[it] = aoffsets[it-1] + lda; - if (cols & 4) { - for (int it = 0; it < 3; it++) - c1[it] = vec_xl(0, aoffsets[it]); - vector_permute_store_4(c1, boffset); - } - } - } - - void KERNEL_4x4(int64_t ii, int64_t jj) { - vec_t vec_A[4], vec_B[4], vec_C[4]; - acc_t acc_0; - __builtin_mma_xxsetaccz(&acc_0); - for (int l = 0; l < k; l += 4) { - packTranspose(A + (ii * lda) + l, lda, 4, 4, (float *)vec_A); - packTranspose(B + (jj * ldb) + l, ldb, 4, 4, (float *)vec_B); - __builtin_mma_xvf32gerpp(&acc_0, vec_A[0], vec_B[0]); - __builtin_mma_xvf32gerpp(&acc_0, vec_A[1], vec_B[1]); - __builtin_mma_xvf32gerpp(&acc_0, vec_A[2], vec_B[2]); - __builtin_mma_xvf32gerpp(&acc_0, vec_A[3], vec_B[3]); - } - save_acc(&acc_0, ii, jj); - } - - void KERNEL_4x8(int64_t ii, int64_t jj) { - vec_t vec_A[4], vec_B[8], vec_C[4]; - acc_t acc_0, acc_1; - __builtin_mma_xxsetaccz(&acc_0); - __builtin_mma_xxsetaccz(&acc_1); - for (int64_t l = 0; l < k; l += 4) { - packTranspose(A + (ii * lda) + l, lda, 4, 4, (float *)vec_A); - packTranspose(B + (jj * ldb) + l, ldb, 8, 4, (float *)vec_B); - __builtin_mma_xvf32gerpp(&acc_0, vec_A[0], (vec_t)vec_B[0]); - __builtin_mma_xvf32gerpp(&acc_1, vec_A[0], (vec_t)vec_B[1]); - __builtin_mma_xvf32gerpp(&acc_0, vec_A[1], (vec_t)vec_B[2]); - __builtin_mma_xvf32gerpp(&acc_1, vec_A[1], (vec_t)vec_B[3]); - __builtin_mma_xvf32gerpp(&acc_0, vec_A[2], (vec_t)vec_B[4]); - __builtin_mma_xvf32gerpp(&acc_1, vec_A[2], (vec_t)vec_B[5]); - __builtin_mma_xvf32gerpp(&acc_0, vec_A[3], (vec_t)vec_B[6]); - __builtin_mma_xvf32gerpp(&acc_1, vec_A[3], (vec_t)vec_B[7]); - } - save_acc(&acc_0, ii, jj); - save_acc(&acc_1, ii, jj + 4); - } - - void KERNEL_8x4(int64_t ii, int64_t jj) { - vec_t vec_A[8], vec_B[4], vec_C[4]; - acc_t acc_0, acc_1; - __builtin_mma_xxsetaccz(&acc_0); - __builtin_mma_xxsetaccz(&acc_1); - for (int64_t l = 0; l < k; l += 4) { - packTranspose(A + (ii * lda) + l, lda, 8, 4, (float *)vec_A); - packTranspose(B + (jj * ldb) + l, ldb, 4, 4, (float *)vec_B); - __builtin_mma_xvf32gerpp(&acc_0, (vec_t)vec_A[0], vec_B[0]); - __builtin_mma_xvf32gerpp(&acc_1, (vec_t)vec_A[1], vec_B[0]); - __builtin_mma_xvf32gerpp(&acc_0, (vec_t)vec_A[2], vec_B[1]); - __builtin_mma_xvf32gerpp(&acc_1, (vec_t)vec_A[3], vec_B[1]); - __builtin_mma_xvf32gerpp(&acc_0, (vec_t)vec_A[4], vec_B[2]); - __builtin_mma_xvf32gerpp(&acc_1, (vec_t)vec_A[5], vec_B[2]); - __builtin_mma_xvf32gerpp(&acc_0, (vec_t)vec_A[6], vec_B[3]); - __builtin_mma_xvf32gerpp(&acc_1, (vec_t)vec_A[7], vec_B[3]); - } - save_acc(&acc_0, ii, jj); - save_acc(&acc_1, ii + 4, jj); - } - - void KERNEL_8x8(int64_t ii, int64_t jj) { - vec_t vec_A[16], vec_B[16], vec_C[4]; - acc_t acc_0, acc_1, acc_2, acc_3; - __builtin_mma_xxsetaccz(&acc_0); - __builtin_mma_xxsetaccz(&acc_1); - __builtin_mma_xxsetaccz(&acc_2); - __builtin_mma_xxsetaccz(&acc_3); - for (int l = 0; l < k; l+=8) { - packTranspose(A + (ii * lda) + l, lda, 8, 8, (float *)vec_A); - packTranspose(B + (jj * ldb) + l, ldb, 8, 8, (float *)vec_B); - for(int x = 0; x < 16; x+=2) { - __builtin_mma_xvf32gerpp(&acc_0, (vec_t)vec_A[x], vec_B[x]); - __builtin_mma_xvf32gerpp(&acc_1, (vec_t)vec_A[x], vec_B[x + 1]); - __builtin_mma_xvf32gerpp(&acc_2, (vec_t)vec_A[x + 1], vec_B[x]); - __builtin_mma_xvf32gerpp(&acc_3, (vec_t)vec_A[x + 1], vec_B[x + 1]); - } - } - save_acc(&acc_0, ii, jj); - save_acc(&acc_1, ii, jj + 4); - save_acc(&acc_2, ii + 4, jj); - save_acc(&acc_3, ii + 4, jj + 4); - } - - inline void MMA_16x8(vec_t * vec_A0, vec_t * vec_A1, vec_t * vec_B, acc_t * acc) { - for (int x = 0; x < 16; x += 2) { - __builtin_mma_xvf32gerpp(&acc[0], vec_A0[x + 0], vec_B[x]); - __builtin_mma_xvf32gerpp(&acc[1], vec_A0[x + 0], vec_B[x + 1]); - __builtin_mma_xvf32gerpp(&acc[2], vec_A0[x + 1], vec_B[x]); - __builtin_mma_xvf32gerpp(&acc[3], vec_A0[x + 1], vec_B[x + 1]); - __builtin_mma_xvf32gerpp(&acc[4], vec_A1[x + 0], vec_B[x]); - __builtin_mma_xvf32gerpp(&acc[5], vec_A1[x + 0], vec_B[x + 1]); - __builtin_mma_xvf32gerpp(&acc[6], vec_A1[x + 1], vec_B[x]); - __builtin_mma_xvf32gerpp(&acc[7], vec_A1[x + 1], vec_B[x + 1]); - } - } - - void KERNEL(int64_t ii, int64_t jj, int64_t mc, int64_t nc, int64_t kc, vec_t * vec_A, vec_t * vec_B, int64_t kk) { - for (int64_t i = 0; i < mc; i += 16) { - int A_base_addr = (mc / 8) * (i / 8) * 16; - for (int64_t j = 0; j < nc; j += 8) { - int B_base_addr = (nc / 8) * (j / 8) * 16; - acc_t acc[8]; - vec_t A0_block[16]; vec_t A1_block[16]; - for (int x = 0; x < 8; x++) - __builtin_mma_xxsetaccz(&acc[x]); - for (int64_t l = 0; l < kc; l += 8) { - int A0_block_idx = A_base_addr + (l / 8) * 16; - int A1_block_idx = A0_block_idx + (mc / 8) * 16; - int B_block_idx = B_base_addr + (l / 8) * 16; - vec_t* A0_block = &vec_A[A0_block_idx]; - vec_t* A1_block = &vec_A[A1_block_idx]; - vec_t* B_block = &vec_B[B_block_idx]; - MMA_16x8(A0_block, A1_block, B_block, acc); - } - if (kk == 0) { - save_acc(&acc[0], ii + i, jj + j); - save_acc(&acc[1], ii + i, jj + j + 4); - save_acc(&acc[2], ii + i + 4, jj + j); - save_acc(&acc[3], ii + i + 4, jj + j + 4); - save_acc(&acc[4], ii + i + 8, jj + j); - save_acc(&acc[5], ii + i + 8, jj + j + 4); - save_acc(&acc[6], ii + i + 12, jj + j); - save_acc(&acc[7], ii + i + 12, jj + j + 4); - } else { - add_save_acc(&acc[0], ii + i, jj + j); - add_save_acc(&acc[1], ii + i, jj + j + 4); - add_save_acc(&acc[2], ii + i + 4, jj + j); - add_save_acc(&acc[3], ii + i + 4, jj + j + 4); - add_save_acc(&acc[4], ii + i + 8, jj + j); - add_save_acc(&acc[5], ii + i + 8, jj + j + 4); - add_save_acc(&acc[6], ii + i + 12, jj + j); - add_save_acc(&acc[7], ii + i + 12, jj + j + 4); - } - } - } - } - - void matmul_tiled(int64_t m , int64_t n, int64_t mc, int64_t nc, int64_t kc) { - int64_t ytiles = m / mc; - int64_t xtiles = n / nc; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - if (end > tiles) { - end = tiles; - } - for (int64_t job = start; job < end; ++job) { - int64_t ii = (job / xtiles) * mc; - int64_t jj = (job % xtiles) * nc; - for (int64_t kk = 0; kk < k; kk += kc) { - vec_t A_pack[kc * mc / 4]; - vec_t B_pack[kc * nc / 4]; - packTranspose(A + (ii * lda) + kk, lda, kc, mc, (float *)A_pack); - packTranspose(B + (jj * ldb) + kk, ldb, kc, nc, (float *)B_pack); - KERNEL(ii, jj, mc, nc, kc, A_pack, B_pack, kk); - } - } - } - - void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int m_rem = MIN(m - m0, 8); - int n_rem = MIN(n - n0, 8); - int mc = 0, nc = 0; - if (m_rem >= 8 && n_rem >= 8) { - mc = 8; - nc = 8; - gemm<8, 8>(m0, m, n0, n); - } else if (m_rem >= 4 && n_rem >= 8) { - mc = 4; - nc = 8; - gemm<4, 8>(m0, m, n0, n); - } else if (m_rem >= 8 && n_rem >= 4) { - mc = 8; - nc = 4; - gemm<8, 4>(m0, m, n0, n); - } else if (m_rem >= 4 && n_rem >= 4) { - mc = 4; - nc = 4; - gemm<4, 4>(m0, m, n0, n); - } else { - mc = (m_rem >= 4) ? 4 : m_rem; - nc = (n_rem >= 4) ? 4 : n_rem; - if (mc == 0 || nc == 0) - return; - gemm_small(m0, m, n0, n, mc, nc); - } - int64_t mp = m0 + ((m - m0) / mc) * mc; - int64_t np = n0 + ((n - n0) / nc) * nc; - mnpack(mp, m, n0, np); - mnpack(m0, m, np, n); - } - - void gemm_small(int64_t m0, int64_t m, int64_t n0, int64_t n, int RM, int RN) { - int64_t ytiles = (m - m0) / RM; - int64_t xtiles = (n - n0) / RN; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - if (end > tiles) - end = tiles; - for (int64_t job = start; job < end; ++job) { - int64_t ii = m0 + job / xtiles * RM; - int64_t jj = n0 + job % xtiles * RN; - vec_t vec_C[4]; - acc_t acc_0; - __builtin_mma_xxsetaccz(&acc_0); - vec_t vec_A[4] = {0}, vec_B[4] = {0}; - for (int l = 0; l < k; l += 4) { - /* 'GEMV Forwarding' concept is used in first two conditional loops. - * when one of the matrix has a single row/column, the elements are - * broadcasted, instead of using packing routine to prepack the - * matrix elements. - */ - if (RM == 1) { - float * a = const_cast(A + (ii) * lda + l); - packTranspose(B + (jj * ldb) + l, ldb, RN, 4, (float *)vec_B); - vec_A[0] = (vec_t)vec_xl(0,a); - vec_A[1] = (vec_t)vec_splats(*((float *)&vec_A+1)); - vec_A[2] = (vec_t)vec_splats(*((float *)&vec_A+2)); - vec_A[3] = (vec_t)vec_splats(*((float *)&vec_A+3)); - } else if (RN == 1) { - packTranspose(A + (ii * lda) + l, lda, RM, 4, (float *)vec_A); - float * b = const_cast(B + (jj) * ldb + l); - vec_B[0] = (vec_t)vec_xl(0,b); - vec_B[1] = (vec_t)vec_splats(*((float *)&vec_B+1)); - vec_B[2] = (vec_t)vec_splats(*((float *)&vec_B+2)); - vec_B[3] = (vec_t)vec_splats(*((float *)&vec_B+3)); - } else { - packTranspose(A + (ii * lda) + l, lda, RM, 4, (float *)vec_A); - packTranspose(B + (jj * ldb) + l, ldb, RN, 4, (float *)vec_B); - } - __builtin_mma_xvf32gerpp(&acc_0, vec_A[0], vec_B[0]); - __builtin_mma_xvf32gerpp(&acc_0, vec_A[1], vec_B[1]); - __builtin_mma_xvf32gerpp(&acc_0, vec_A[2], vec_B[2]); - __builtin_mma_xvf32gerpp(&acc_0, vec_A[3], vec_B[3]); - } - __builtin_mma_disassemble_acc(vec_C, &acc_0); - for (int I = 0; I < RM; I++) { - for (int J = 0; J < RN; J++) { - *((float *)(C+ii+((jj+J)*ldc)+I)) = *((float *)&vec_C[I]+J); - } - } - } - } - - template - inline void kernel(int64_t ii, int64_t jj) { - if constexpr(RM == 4 && RN == 4) { - KERNEL_4x4(ii, jj); - } else if constexpr(RM == 4 && RN == 8) { - KERNEL_4x8(ii, jj); - } else if constexpr(RM == 8 && RN == 4) { - KERNEL_8x4(ii, jj); - } else if constexpr(RM == 8 && RN == 8) { - KERNEL_8x8(ii, jj); - } else { - static_assert(false, "RN/RM values not supported"); - } - } - - template - NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { - int64_t ytiles = (m - m0) / RM; - int64_t xtiles = (n - n0) / RN; - int64_t tiles = xtiles * ytiles; - int64_t duty = (tiles + nth - 1) / nth; - int64_t start = duty * ith; - int64_t end = start + duty; - if (end > tiles) - end = tiles; - for (int64_t job = start; job < end; ++job) { - int64_t ii = m0 + job / xtiles * RM; - int64_t jj = n0 + job % xtiles * RN; - kernel(ii, jj); - } - } - - const float * const A; - const float * const B; - float * C; - const int64_t k; - const int64_t lda; - const int64_t ldb; - const int64_t ldc; - const int ith; - const int nth; -}; -#endif -} // namespace - -/** - * Performs optimized matrix multiplication on CPU. - * - * This subroutine may compute C = Aᵀ * B with column major ordering. - * Despite its name, this isn't a generalized implementation. Work is - * only performed when a handwritten kernel is written and available. - * Otherwise the caller should fall back to a general matmul routine. - * - * For example, for single-threaded single-precision GEMM you can say - * - * llamafile_sgemm(m, n, k, A, lda, B, ldb, C, ldc, - * 0, 1, - * GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32); - * - * @param m is rows in `A` and `C` - * @param n is cols in `B` and `C` - * @param k is cols in `A` and rows in `B` - * @param A is first input matrix (always transposed) - * @param lda is row stride of `A` - * @param B is second input matrix (never transposed) - * @param ldb is row stride of `B` - * @param C is input/output array of output matrices - * @param ldc is row stride of `C` - * @param ith is thread id (must be less than `nth`) - * @param nth is number of threads (must be greater than zero) - * @param Atype is GGML data type of `A` - * @param Btype is GGML data type of `B` - * @param Ctype is GGML data type of `C` - * @return true if this function was able to service the matmul request - */ -bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64_t n, int64_t k, - const void *A, int64_t lda, const void *B, int64_t ldb, void *C, - int64_t ldc, int Atype, int Btype, int Ctype) { - - assert(m >= 0); - assert(n >= 0); - assert(k >= 0); - assert(lda >= k); - assert(ldb >= k); - assert(ldc >= m); - assert(params->nth > 0); - assert(params->ith < params->nth); - - // only enable sgemm for prompt processing -#if !defined(__MMA__) - if (n < 2) - return false; -#endif - - if (Ctype != GGML_TYPE_F32) - return false; - - switch (Atype) { - - case GGML_TYPE_F32: { - if (Btype != GGML_TYPE_F32) - return false; -#if defined(__AVX512F__) - tinyBLAS<16, __m512, __m512, float, float, float> tb{ params, - k, (const float *)A, lda, - (const float *)B, ldb, - (float *)C, ldc}; - return tb.matmul(m, n); -#elif defined(__AVX__) || defined(__AVX2__) - tinyBLAS<8, __m256, __m256, float, float, float> tb{ params, - k, (const float *)A, lda, - (const float *)B, ldb, - (float *)C, ldc}; - return tb.matmul(m, n); -#elif defined(__ARM_NEON) - if (n < 4) - return false; - tinyBLAS<4, float32x4_t, float32x4_t, float, float, float> tb{ params, - k, (const float *)A, lda, - (const float *)B, ldb, - (float *)C, ldc}; - return tb.matmul(m, n); -#elif defined(__VXE__) || defined(__VXE2__) - if (n < 4) - return false; - tinyBLAS<4, float32x4_t, float32x4_t, float, float, float> tb{ params, - k, (const float *)A, lda, - (const float *)B, ldb, - (float *)C, ldc}; - return tb.matmul(m, n); -#elif defined(__MMA__) - if (k % 8) - return false; - tinyBLAS_PPC tb{ - k, (const float *)A, lda, - (const float *)B, ldb, - (float *)C, ldc, - params->ith, params->nth}; - tb.matmul(m, n); - return true; -#elif defined(__riscv_v_intrinsic) - #if LMUL == 1 - tinyBLAS_RVV tb{ params, - k, (const float *)A, lda, - (const float *)B, ldb, - (float *)C, ldc}; - #elif LMUL == 2 - tinyBLAS_RVV tb{ params, - k, (const float *)A, lda, - (const float *)B, ldb, - (float *)C, ldc}; - #else // LMUL = 4 - tinyBLAS_RVV tb{ params, - k, (const float *)A, lda, - (const float *)B, ldb, - (float *)C, ldc}; - #endif - return tb.matmul(m, n); -#else - return false; -#endif - } - - case GGML_TYPE_BF16: { -#if defined(__AVX512BF16__) - if (Btype == GGML_TYPE_BF16) { - tinyBLAS<32, __m512, __m512bh, ggml_bf16_t, ggml_bf16_t, float> tb{ params, k, - (const ggml_bf16_t *)A, lda, - (const ggml_bf16_t *)B, ldb, - (float *)C, ldc}; - return tb.matmul(m, n); - } -#elif defined(__AVX512F__) - if (Btype == GGML_TYPE_BF16) { - tinyBLAS<16, __m512, __m512, ggml_bf16_t, ggml_bf16_t, float> tb{ params, k, - (const ggml_bf16_t *)A, lda, - (const ggml_bf16_t *)B, ldb, - (float *)C, ldc}; - return tb.matmul(m, n); - } -#elif defined(__AVX2__) - if (Btype == GGML_TYPE_BF16) { - tinyBLAS<8, __m256, __m256, ggml_bf16_t, ggml_bf16_t, float> tb{ params, k, - (const ggml_bf16_t *)A, lda, - (const ggml_bf16_t *)B, ldb, - (float *)C, ldc}; - return tb.matmul(m, n); - } -#elif defined(__MMA__) - if (k % 8) { - return false; - } - - if (Btype == GGML_TYPE_BF16) { - tinyBLAS_HP16_PPC tb{ k, - (const ggml_bf16_t *)A, lda, - (const ggml_bf16_t *)B, ldb, - (float *)C, ldc, - params->ith, params->nth }; - - tb.matmul(m, n); - return true; - } -#elif defined(__riscv_zvfbfwma) - if (Btype == GGML_TYPE_BF16) { - #if LMUL == 1 - tinyBLAS_RVV tb{ params, - k, (const ggml_bf16_t *)A, lda, - (const ggml_bf16_t *)B, ldb, - (float *)C, ldc}; - #elif LMUL == 2 - tinyBLAS_RVV tb{ params, - k, (const ggml_bf16_t *)A, lda, - (const ggml_bf16_t *)B, ldb, - (float *)C, ldc}; - #else // LMUL = 4 - tinyBLAS_RVV tb{ params, - k, (const ggml_bf16_t *)A, lda, - (const ggml_bf16_t *)B, ldb, - (float *)C, ldc}; - #endif - return tb.matmul(m, n); - } -#endif - return false; - } - - case GGML_TYPE_F16: { -#if defined(__AVX512F__) - if (Btype == GGML_TYPE_F16) { - tinyBLAS<16, __m512, __m512, ggml_fp16_t, ggml_fp16_t, float> tb{ params, k, - (const ggml_fp16_t *)A, lda, - (const ggml_fp16_t *)B, ldb, - (float *)C, ldc}; - return tb.matmul(m, n); - } -#elif (defined(__AVX__) || defined(__AVX2__)) && defined(__F16C__) - if (Btype == GGML_TYPE_F16) { - tinyBLAS<8, __m256, __m256, ggml_fp16_t, ggml_fp16_t, float> tb{ params, k, - (const ggml_fp16_t *)A, lda, - (const ggml_fp16_t *)B, ldb, - (float *)C, ldc}; - return tb.matmul(m, n); - } -#elif defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER) - if (n < 8) - return false; - if (Btype == GGML_TYPE_F16) { - tinyBLAS<8, float16x8_t, float16x8_t, ggml_fp16_t, ggml_fp16_t, float> tb{ params, - k, (const ggml_fp16_t *)A, lda, - (const ggml_fp16_t *)B, ldb, - (float *)C, ldc}; - return tb.matmul(m, n); - } -#elif defined(__ARM_NEON) && !defined(_MSC_VER) - if (Btype == GGML_TYPE_F32) { - tinyBLAS<4, float32x4_t, float32x4_t, ggml_fp16_t, float, float> tb{ params, - k, (const ggml_fp16_t *)A, lda, - (const float *)B, ldb, - (float *)C, ldc}; - return tb.matmul(m, n); - } -#elif defined(__VXE__) || defined(__VXE2__) - if (n < 4) - return false; - if (Btype == GGML_TYPE_F16) { - tinyBLAS<4, float32x4_t, float32x4_t, ggml_fp16_t, ggml_fp16_t, float> tb{ params, - k, (const ggml_fp16_t *)A, lda, - (const ggml_fp16_t *)B, ldb, - (float *)C, ldc}; - return tb.matmul(m, n); - } -#elif defined(__riscv_zvfh) - if (Btype == GGML_TYPE_F16) { - #if LMUL == 1 - tinyBLAS_RVV tb{ params, - k, (const ggml_fp16_t *)A, lda, - (const ggml_fp16_t *)B, ldb, - (float *)C, ldc}; - #elif LMUL == 2 - tinyBLAS_RVV tb{ params, - k, (const ggml_fp16_t *)A, lda, - (const ggml_fp16_t *)B, ldb, - (float *)C, ldc}; - #else // LMUL = 4 - tinyBLAS_RVV tb{ params, - k, (const ggml_fp16_t *)A, lda, - (const ggml_fp16_t *)B, ldb, - (float *)C, ldc}; - #endif - return tb.matmul(m, n); - } -#elif defined(__MMA__) - if (k % 8) { - return false; - } - - if (Btype == GGML_TYPE_F16) { - tinyBLAS_HP16_PPC tb{ k, - (const ggml_fp16_t *)A, lda, - (const ggml_fp16_t *)B, ldb, - (float *)C, ldc, - params->ith, params->nth }; - - tb.matmul(m, n); - return true; - } -#endif - return false; - } - - case GGML_TYPE_Q8_0: { - if (Btype != GGML_TYPE_Q8_0) - return false; -#if defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX__) - tinyBLAS_Q0_AVX tb{ - k, (const block_q8_0 *)A, lda, - (const block_q8_0 *)B, ldb, - (float *)C, ldc, - params->ith, params->nth}; - tb.matmul(m, n); - return true; -#elif defined(__ARM_FEATURE_DOTPROD) - tinyBLAS_Q0_ARM tb{ - k, (const block_q8_0 *)A, lda, - (const block_q8_0 *)B, ldb, - (float *)C, ldc, - params->ith, params->nth}; - tb.matmul(m, n); - return true; -#elif defined(__MMA__) - //TO-DO: Remove this condition once gemv forwarding is enabled. - if (n < 8 && n != 4) - return false; - if (m < 8 && m != 4) - return false; - tinyBLAS_Q0_PPC tb{ - k, (const block_q8_0 *)A, lda, - (const block_q8_0 *)B, ldb, - (float *)C, ldc, - params->ith, params->nth}; - tb.matmul(m, n); - return true; -#else - return false; -#endif - } - - case GGML_TYPE_Q4_0: { - if (Btype != GGML_TYPE_Q8_0) - return false; -#if defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX__) - tinyBLAS_Q0_AVX tb{ - k, (const block_q4_0 *)A, lda, - (const block_q8_0 *)B, ldb, - (float *)C, ldc, - params->ith, params->nth}; - tb.matmul(m, n); - return true; -#elif defined(__ARM_FEATURE_DOTPROD) - tinyBLAS_Q0_ARM tb{ - k, (const block_q4_0 *)A, lda, - (const block_q8_0 *)B, ldb, - (float *)C, ldc, - params->ith, params->nth}; - tb.matmul(m, n); - return true; -#elif defined(__MMA__) - //TO-DO: Remove this condition once gemv forwarding is enabled. - if (n < 8 && n != 4) - return false; - if (m < 8 && m != 4) - return false; - tinyBLAS_Q0_PPC tb{ - k, (const block_q4_0 *)A, lda, - (const block_q8_0 *)B, ldb, - (float *)C, ldc, - params->ith, params->nth}; - tb.matmul(m, n); - return true; -#else - return false; -#endif - } - - case GGML_TYPE_Q5_0: { - if (Btype != GGML_TYPE_Q8_0) - return false; -#if defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX__) - tinyBLAS_Q0_AVX tb{ - k, (const block_q5_0 *)A, lda, - (const block_q8_0 *)B, ldb, - (float *)C, ldc, - params->ith, params->nth}; - tb.matmul(m, n); - return true; -#else - return false; -#endif - } - - case GGML_TYPE_IQ4_NL: { - if (Btype != GGML_TYPE_Q8_0) - return false; -#if defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX__) - tinyBLAS_Q0_AVX tb{ - k, (const block_iq4_nl *)A, lda, - (const block_q8_0 *)B, ldb, - (float *)C, ldc, - params->ith, params->nth}; - tb.matmul(m, n); - return true; -#else - return false; -#endif - } - - default: - return false; - } - - (void)params; - (void)m; - (void)n; - (void)k; - (void)A; - (void)lda; - (void)B; - (void)ldb; - (void)C; - (void)ldc; - (void)Atype; - (void)Btype; - (void)Ctype; -} diff --git a/ggml/src/ggml-cpu/llamafile/sgemm.h b/ggml/src/ggml-cpu/llamafile/sgemm.h deleted file mode 100644 index 867b0c04aee8..000000000000 --- a/ggml/src/ggml-cpu/llamafile/sgemm.h +++ /dev/null @@ -1,25 +0,0 @@ -#pragma once -#include -#include - -#if defined(__VXE__) || defined(__VXE2__) -#include -#endif - -#ifdef _MSC_VER -#define NOINLINE __declspec(noinline) -#else -#define NOINLINE __attribute__((__noinline__)) -#endif - -#ifdef __cplusplus -extern "C" { -#endif - -bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t, int64_t, int64_t, - const void *, int64_t, const void *, int64_t, void *, int64_t, - int, int, int); - -#ifdef __cplusplus -} -#endif diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp.inc similarity index 73% rename from ggml/src/ggml-cpu/ops.cpp rename to ggml/src/ggml-cpu/ops.cpp.inc index a9bc21da6f0f..b4604c87ed7c 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp.inc @@ -1,1687 +1,2768 @@ -#include "ops.h" +#include "ggml.h.inc" -#include "ggml-cpu.h" -#include "ggml-impl.h" -#include "binary-ops.h" -#include "simd-gemm.h" -#include "ggml.h" -#include "unary-ops.h" -#include "vec.h" +// +// cache line +// -#include -#include -#include +#if defined(__cpp_lib_hardware_interference_size) +#define CACHE_LINE_SIZE std::hardware_destructive_interference_size +#else +#if defined(__POWER9_VECTOR__) +#define CACHE_LINE_SIZE 128 +#elif defined(__VXE__) || defined(__VXE2__) +#define CACHE_LINE_SIZE 256 +#else +#define CACHE_LINE_SIZE 64 +#endif +#endif -// ggml_compute_forward_dup +static const size_t CACHE_LINE_SIZE_F32 = CACHE_LINE_SIZE/sizeof(float); -static void ggml_compute_forward_dup_same_cont( - const ggml_compute_params * params, - ggml_tensor * dst) { +// Work buffer size for im2col operations in CONV2D +#define GGML_IM2COL_WORK_SIZE (16 * 1024 * 1024) - const ggml_tensor * src0 = dst->src[0]; +#ifdef __cplusplus +extern "C" { +#endif - GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src0)); - GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src0)); - GGML_ASSERT(src0->type == dst->type); +void ggml_compute_forward_dup(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_add(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_add_id(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_add1(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_acc(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sum(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sum_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_cumsum(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_mean(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_argmax(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_count_equal(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_repeat(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_repeat_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_concat(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_silu_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rms_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rms_norm_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_group_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_l2_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_out_prod(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_scale(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_set(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_cpy(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_cont(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_get_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_get_rows_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_set_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_diag(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_diag_mask_inf(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_diag_mask_zero(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_soft_max(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_soft_max_ext_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rope(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rope_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_clamp(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_conv_transpose_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_im2col(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_im2col_back_f32(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_im2col_3d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_conv_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_conv_3d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_conv_transpose_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_conv_2d_dw(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_pool_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_pool_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_pool_2d_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_upscale(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_pad(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_pad_reflect_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_roll(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_arange(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_timestep_embedding(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_argsort(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_top_k(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_leaky_relu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_fill(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_flash_attn_ext(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_flash_attn_back( + const struct ggml_compute_params * params, + const bool masked, + struct ggml_tensor * dst); +void ggml_compute_forward_ssm_conv(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_ssm_scan(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_win_part(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_win_unpart(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_unary(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_glu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_get_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_add_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rwkv_wkv6(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rwkv_wkv7(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_custom(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_cross_entropy_loss(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_cross_entropy_loss_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_opt_step_adamw(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_opt_step_sgd(const struct ggml_compute_params * params, struct ggml_tensor * dst); +#ifdef __cplusplus +} +#endif - const size_t nb0 = ggml_type_size(src0->type); - const int ith = params->ith; // thread index - const int nth = params->nth; // number of threads +#include "ggml.h.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml-cpu-impl-defs.inc" +#include "ggml-impl-defs.inc" +#include "simd-mappings-defs.inc" - // parallelize by blocks - const int nk = ggml_nelements(src0)/ggml_blck_size(src0->type); - const int dr = (nk + nth - 1) / nth; - const int k0 = dr * ith; - const int k1 = MIN(k0 + dr, nk); +#ifdef __cplusplus +#include +extern "C" { +#endif - if (k0 < k1) { - memcpy( - ((char *) dst->data + k0*nb0), - ((char *) src0->data + k0*nb0), - (k1 - k0) * nb0); - } +// return true if op part of extra "accelerator" +bool ggml_cpu_extra_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op); +bool ggml_cpu_extra_work_size(int n_threads, const struct ggml_tensor * op, size_t * size); + +#ifdef __cplusplus } -template -static void ggml_compute_forward_dup_flt( - const ggml_compute_params * params, - ggml_tensor * dst) { +namespace ggml::cpu { +// register in tensor->extra +class tensor_traits { + public: + virtual ~tensor_traits(); + virtual bool work_size(int n_threads, const struct ggml_tensor * op, size_t & size) = 0; + virtual bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) = 0; +}; - const ggml_tensor * src0 = dst->src[0]; +class extra_buffer_type { + public: + virtual ~extra_buffer_type(); + virtual bool supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) = 0; + virtual tensor_traits * get_tensor_traits(const struct ggml_tensor * op) = 0; +}; +} // namespace ggml::cpu - GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src0)); - GGML_ASSERT(!ggml_is_quantized(src0->type) && !ggml_is_quantized(dst->type)); +// implemented in ggml-cpu.cpp. +std::vector & ggml_backend_cpu_get_extra_buffer_types(); - GGML_TENSOR_UNARY_OP_LOCALS +#endif - const int ith = params->ith; // thread index - const int nth = params->nth; // number of threads +#define GGML_FA_TILE_Q 64 +#define GGML_FA_TILE_KV 64 + +#ifdef __cplusplus + +#include + +// convenience functions/macros for use in template calls +// note: these won't be required after the 'traits' lookup table is used. +static inline ggml_fp16_t f32_to_f16(float x) { + return GGML_CPU_FP32_TO_FP16(x); +} + +static inline float f16_to_f32(ggml_fp16_t x) { + return GGML_CPU_FP16_TO_FP32(x); +} + +static inline ggml_bf16_t f32_to_bf16(float x) { + return GGML_FP32_TO_BF16(x); +} + +static inline float bf16_to_f32(ggml_bf16_t x) { + return GGML_BF16_TO_FP32(x); +} + +static inline float i32_to_f32(int32_t x) { + return x; +} + +static inline int32_t f32_to_i32(float x) { + return x; +} + +static inline float f32_to_f32(float x) { + return x; +} + +// TODO - merge this into the traits table, after using row-based conversions +template +struct type_conversion_table; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_fp16_t) = f16_to_f32; + static constexpr ggml_fp16_t (*from_f32)(float) = f32_to_f16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(float) = f32_to_f32; + static constexpr float (*from_f32)(float) = f32_to_f32; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_bf16_t) = bf16_to_f32; + static constexpr ggml_bf16_t (*from_f32)(float) = f32_to_bf16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(int32_t) = i32_to_f32; + static constexpr int32_t (*from_f32)(float) = f32_to_i32; +}; + +static std::pair get_thread_range(const struct ggml_compute_params * params, const struct ggml_tensor * src0) { + const int64_t ith = params->ith; + const int64_t nth = params->nth; + + const int64_t nr = ggml_nrows(src0); + + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; - // parallelize by rows - const int nr = ne01; - // number of rows per thread - const int dr = (nr + nth - 1) / nth; // row range for this thread - const int ir0 = dr * ith; - const int ir1 = MIN(ir0 + dr, nr); + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); - // case: type & row size equal - if (src0->type == dst->type && - ne00 == ne0 && - nb00 == ggml_type_size(src0->type) && nb0 == ggml_type_size(dst->type)) { - // copy by rows - const size_t rs = ne00*nb00; - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = ir0; i01 < ir1; i01++) { - memcpy( - ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3), - ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03), - rs); - } + return {ir0, ir1}; +} + +struct ggml_fa_tile_config { + static constexpr size_t Q = GGML_FA_TILE_Q; + static constexpr size_t KV = GGML_FA_TILE_KV; +}; + +#endif + +#include "ggml-cpu.h.inc" +#include "ggml-impl-defs.inc" +// Computes C[M x N] += A[M x K] * B[K x N] + +#include "simd-mappings-defs.inc" + +// TODO: add support for sizeless vector types +#if defined(GGML_SIMD) && !defined(__ARM_FEATURE_SVE) && !defined(__riscv_v_intrinsic) + +// TODO: untested on avx512 +// These are in units of GGML_F32_EPR +#if defined(__AVX512F__) || defined (__ARM_NEON__) + static constexpr int GEMM_RM = 4; + static constexpr int GEMM_RN = 4; // 16+4+1 = 25/32 +#elif defined(__AVX2__) || defined(__AVX__) + static constexpr int GEMM_RM = 6; + static constexpr int GEMM_RN = 2; // 12+2+1 = 15/16 +#else + static constexpr int GEMM_RM = 2; + static constexpr int GEMM_RN = 2; +#endif + +template +static inline void simd_gemm_ukernel( + float * GGML_RESTRICT C, + const float * GGML_RESTRICT A, + const float * GGML_RESTRICT B, + int K, int N) +{ + static constexpr int KN = GGML_F32_EPR; + + GGML_F32_VEC acc[RM][RN]; + for (int64_t i = 0; i < RM; i++) { + for (int r = 0; r < RN; r++) { + acc[i][r] = GGML_F32_VEC_LOAD(C + i * N + r * KN); + } + } + + for (int64_t kk = 0; kk < K; kk++) { + GGML_F32_VEC Bv[RN]; + for (int r = 0; r < RN; r++) { + Bv[r] = GGML_F32_VEC_LOAD(B + kk * N + r * KN); + } + for (int64_t i = 0; i < RM; i++) { + GGML_F32_VEC p = GGML_F32_VEC_SET1(A[i * K + kk]); + for (int r = 0; r < RN; r++) { + acc[i][r] = GGML_F32_VEC_FMA(acc[i][r], Bv[r], p); } } - return; } - // case: dst tensor is contiguous - if (ggml_is_contiguous(dst)) { - if (nb00 == sizeof(src_t)) { - if constexpr (std::is_same_v) { - // same type - size_t id = 0; - const size_t rs = ne00 * nb00; - char * dst_ptr = (char *) dst->data; + for (int64_t i = 0; i < RM; i++) { + for (int r = 0; r < RN; r++) { + GGML_F32_VEC_STORE(C + i * N + r * KN, acc[i][r]); + } + } +} - for (int i03 = 0; i03 < ne03; i03++) { - for (int i02 = 0; i02 < ne02; i02++) { - id += rs * ir0; - for (int i01 = ir0; i01 < ir1; i01++) { - const char * src0_ptr = (char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03; - memcpy(dst_ptr + id, src0_ptr, rs); - id += rs; - } - id += rs * (ne01 - ir1); - } - } - } else { - // casting between non-quantized types - size_t id = 0; - dst_t * dst_ptr = (dst_t *) dst->data; +// C[M x N] += A[M x K] * B[K x N] +static void simd_gemm( + float * GGML_RESTRICT C, + const float * GGML_RESTRICT A, + const float * GGML_RESTRICT B, + int M, int K, int N) +{ + static constexpr int KN = GGML_F32_EPR; - for (int i03 = 0; i03 < ne03; i03++) { - for (int i02 = 0; i02 < ne02; i02++) { - id += ne00 * ir0; - for (int i01 = ir0; i01 < ir1; i01++) { - const src_t * src0_ptr = (src_t *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); - for (int i00 = 0; i00 < ne00; i00++) { - float tmp = type_conversion_table::to_f32(src0_ptr[i00]); - dst_ptr[id] = type_conversion_table::from_f32(tmp); - id++; - } - } - id += ne00 * (ne01 - ir1); - } + int64_t ii = 0; + for (; ii + GEMM_RM <= M; ii += GEMM_RM) { + int64_t jj = 0; + for (; jj + GEMM_RN * KN <= N; jj += GEMM_RN * KN) { + simd_gemm_ukernel(C + jj, A, B + jj, K, N); + } + for (; jj + KN <= N; jj += KN) { + simd_gemm_ukernel(C + jj, A, B + jj, K, N); + } + for (; jj < N; jj++) { + for (int64_t i = 0; i < GEMM_RM; i++) { + float a = C[i * N + jj]; + for (int64_t kk = 0; kk < K; kk++) { + a += A[i + kk] * B[kk * N + jj]; } + C[i * N + jj] = a; } - } else { - //printf("%s: this is not optimal - fix me\n", __func__); - - size_t id = 0; - dst_t * dst_ptr = (dst_t *) dst->data; + } - for (int i03 = 0; i03 < ne03; i03++) { - for (int i02 = 0; i02 < ne02; i02++) { - id += ne00 * ir0; - for (int i01 = ir0; i01 < ir1; i01++) { - for (int i00 = 0; i00 < ne00; i00++) { - const src_t * src0_ptr = (src_t *) ((char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); + A += GEMM_RM * K; + C += GEMM_RM * N; + } - float tmp = type_conversion_table::to_f32(*src0_ptr); - dst_ptr[id] = type_conversion_table::from_f32(tmp); - id++; - } - } - id += ne00 * (ne01 - ir1); - } + // Tail rows: one at a time + for (; ii < M; ii++) { + int64_t jj = 0; + for (; jj + GEMM_RN * KN <= N; jj += GEMM_RN * KN) { + simd_gemm_ukernel<1, GEMM_RN>(C + jj, A, B + jj, K, N); + } + for (; jj + KN <= N; jj += KN) { + simd_gemm_ukernel<1, 1>(C + jj, A, B + jj, K, N); + } + for (; jj < N; jj++) { + float a = C[jj]; + for (int64_t kk = 0; kk < K; kk++) { + a += A[kk] * B[kk * N + jj]; } + C[jj] = a; } - return; + + A += K; + C += N; } +} +#elif defined(GGML_SIMD) && defined(__riscv_v_intrinsic) +// RM accumulators + 1 B vector = RM + 1 <= 8 => RM <= 7 +// Microkernel: C[RM x vl] += A[RM x K] * B[K x N] +template +static inline void rvv_simd_gemm_ukernel( + float * GGML_RESTRICT C, + const float * GGML_RESTRICT A, + const float * GGML_RESTRICT B, + int K, int N, size_t vl) +{ + static_assert(RM >= 1 && RM <= 7, "RM must be 1..7 for LMUL=4"); - // dst counters - int64_t i10 = 0; - int64_t i11 = 0; - int64_t i12 = 0; - int64_t i13 = 0; + vfloat32m4_t acc_0 = __riscv_vle32_v_f32m4(C + 0 * N, vl); + vfloat32m4_t acc_1, acc_2, acc_3, acc_4, acc_5, acc_6; + if constexpr (RM > 1) acc_1 = __riscv_vle32_v_f32m4(C + 1 * N, vl); + if constexpr (RM > 2) acc_2 = __riscv_vle32_v_f32m4(C + 2 * N, vl); + if constexpr (RM > 3) acc_3 = __riscv_vle32_v_f32m4(C + 3 * N, vl); + if constexpr (RM > 4) acc_4 = __riscv_vle32_v_f32m4(C + 4 * N, vl); + if constexpr (RM > 5) acc_5 = __riscv_vle32_v_f32m4(C + 5 * N, vl); + if constexpr (RM > 6) acc_6 = __riscv_vle32_v_f32m4(C + 6 * N, vl); - if constexpr (std::is_same_v) { - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - i10 += ne00 * ir0; - while (i10 >= ne0) { - i10 -= ne0; - if (++i11 == ne1) { - i11 = 0; - if (++i12 == ne2) { - i12 = 0; - if (++i13 == ne3) { - i13 = 0; - } - } - } - } - for (int64_t i01 = ir0; i01 < ir1; i01++) { - for (int64_t i00 = 0; i00 < ne00; i00++) { - const char * src0_ptr = ((char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); - char * dst_ptr = ((char *) dst->data + i10*nb0 + i11*nb1 + i12*nb2 + i13*nb3); + for (int kk = 0; kk < K; kk++) { + vfloat32m4_t b_0 = __riscv_vle32_v_f32m4(B + kk * N, vl); - memcpy(dst_ptr, src0_ptr, sizeof(dst_t)); + acc_0 = __riscv_vfmacc_vf_f32m4(acc_0, A[0 * K + kk], b_0, vl); + if constexpr (RM > 1) acc_1 = __riscv_vfmacc_vf_f32m4(acc_1, A[1 * K + kk], b_0, vl); + if constexpr (RM > 2) acc_2 = __riscv_vfmacc_vf_f32m4(acc_2, A[2 * K + kk], b_0, vl); + if constexpr (RM > 3) acc_3 = __riscv_vfmacc_vf_f32m4(acc_3, A[3 * K + kk], b_0, vl); + if constexpr (RM > 4) acc_4 = __riscv_vfmacc_vf_f32m4(acc_4, A[4 * K + kk], b_0, vl); + if constexpr (RM > 5) acc_5 = __riscv_vfmacc_vf_f32m4(acc_5, A[5 * K + kk], b_0, vl); + if constexpr (RM > 6) acc_6 = __riscv_vfmacc_vf_f32m4(acc_6, A[6 * K + kk], b_0, vl); + } - if (++i10 == ne00) { - i10 = 0; - if (++i11 == ne01) { - i11 = 0; - if (++i12 == ne02) { - i12 = 0; - if (++i13 == ne03) { - i13 = 0; - } - } - } - } - } - } - i10 += ne00 * (ne01 - ir1); - while (i10 >= ne0) { - i10 -= ne0; - if (++i11 == ne1) { - i11 = 0; - if (++i12 == ne2) { - i12 = 0; - if (++i13 == ne3) { - i13 = 0; - } - } - } - } + __riscv_vse32_v_f32m4(C + 0 * N, acc_0, vl); + if constexpr (RM > 1) __riscv_vse32_v_f32m4(C + 1 * N, acc_1, vl); + if constexpr (RM > 2) __riscv_vse32_v_f32m4(C + 2 * N, acc_2, vl); + if constexpr (RM > 3) __riscv_vse32_v_f32m4(C + 3 * N, acc_3, vl); + if constexpr (RM > 4) __riscv_vse32_v_f32m4(C + 4 * N, acc_4, vl); + if constexpr (RM > 5) __riscv_vse32_v_f32m4(C + 5 * N, acc_5, vl); + if constexpr (RM > 6) __riscv_vse32_v_f32m4(C + 6 * N, acc_6, vl); +} + +template +static inline void rvv_simd_gemm_dispatch_tail( + float * GGML_RESTRICT C, + const float * GGML_RESTRICT A, + const float * GGML_RESTRICT B, + int K, int N, int KN, int remaining_rows) +{ + if constexpr (RM > 0) { + if (remaining_rows == RM) { + int64_t jj = 0; + for (; jj + KN <= N; jj += KN) { + rvv_simd_gemm_ukernel(C + jj, A, B + jj, K, N, KN); + } + if (jj < N) { + rvv_simd_gemm_ukernel(C + jj, A, B + jj, K, N, N - jj); } + } else { + rvv_simd_gemm_dispatch_tail(C, A, B, K, N, KN, remaining_rows); } + } +} - } else { - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - i10 += ne00 * ir0; - while (i10 >= ne0) { - i10 -= ne0; - if (++i11 == ne1) { - i11 = 0; - if (++i12 == ne2) { - i12 = 0; - if (++i13 == ne3) { - i13 = 0; - } - } - } - } - for (int64_t i01 = ir0; i01 < ir1; i01++) { - for (int64_t i00 = 0; i00 < ne00; i00++) { - const char * src0_ptr = ((char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); - char * dst_ptr = ((char *) dst->data + i10*nb0 + i11*nb1 + i12*nb2 + i13*nb3); +static constexpr int GEMM_RM = 7; - float tmp = type_conversion_table::to_f32(*(const src_t *) src0_ptr); - *(dst_t *) dst_ptr = type_conversion_table::from_f32(tmp); +// C[M x N] += A[M x K] * B[K x N] +static void simd_gemm( + float * GGML_RESTRICT C, + const float * GGML_RESTRICT A, + const float * GGML_RESTRICT B, + int M, int K, int N) +{ + const int KN = (int)__riscv_vlenb(); + int64_t ii = 0; + for (; ii + GEMM_RM <= M; ii += GEMM_RM) { + int64_t jj = 0; + for (; jj + KN <= N; jj += KN) { + rvv_simd_gemm_ukernel(C + jj, A, B + jj, K, N, KN); + } + if (jj < N) { + rvv_simd_gemm_ukernel(C + jj, A, B + jj, K, N, N - jj); + } + A += GEMM_RM * K; + C += GEMM_RM * N; + } - if (++i10 == ne0) { - i10 = 0; - if (++i11 == ne1) { - i11 = 0; - if (++i12 == ne2) { - i12 = 0; - if (++i13 == ne3) { - i13 = 0; - } - } - } - } - } - } - i10 += ne00 * (ne01 - ir1); - while (i10 >= ne0) { - i10 -= ne0; - if (++i11 == ne1) { - i11 = 0; - if (++i12 == ne2) { - i12 = 0; - if (++i13 == ne3) { - i13 = 0; - } - } - } - } + int remaining_rows = M - ii; + rvv_simd_gemm_dispatch_tail(C, A, B, K, N, KN, remaining_rows); +} + +#if defined(__GNUC__) && !defined(__clang__) +#pragma GCC diagnostic pop +#endif + +#else // scalar path + +static void simd_gemm( + float * GGML_RESTRICT C, + const float * GGML_RESTRICT A, + const float * GGML_RESTRICT B, + int M, int K, int N) +{ + for (int64_t i = 0; i < M; i++) { + for (int64_t j = 0; j < N; j++) { + float sum = C[i * N + j]; + for (int64_t kk = 0; kk < K; kk++) { + sum += A[i * K + kk] * B[kk * N + j]; } + C[i * N + j] = sum; } } } +#endif // GGML_SIMD -template -static void ggml_compute_forward_dup_to_q( - const ggml_compute_params * params, - ggml_tensor * dst) { +#include "ggml.h.inc" +// Vectorized functions for fundamental operations - const ggml_tensor * src0 = dst->src[0]; +#include "ggml-impl-defs.inc" +#include "simd-mappings-defs.inc" +#include "ggml.h.inc" +#include "ggml-cpu.h.inc" - GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src0)); - GGML_ASSERT(!ggml_is_quantized(src0->type)); +#if defined(GGML_USE_ACCELERATE) +#include +#endif - GGML_TENSOR_UNARY_OP_LOCALS +// floating point type used to accumulate sums +typedef double ggml_float; - const int ith = params->ith; // thread index - const int nth = params->nth; // number of threads +#define GGML_GELU_FP16 +#define GGML_GELU_QUICK_FP16 - // parallelize by rows - const int nr = ne01; - // number of rows per thread - const int dr = (nr + nth - 1) / nth; - // row range for this thread - const int ir0 = dr * ith; - const int ir1 = MIN(ir0 + dr, nr); +#define GGML_SOFT_MAX_UNROLL 4 +#define GGML_VEC_DOT_UNROLL 2 +#define GGML_VEC_MAD_UNROLL 32 - if (ggml_is_contiguous(dst) && - nb00 == sizeof(src_t) && - ggml_get_type_traits_cpu(dst->type)->from_float) { - // casting non-quantized types --> intermediate f32 --> quantized - ggml_from_float_t const quantize_row_q = ggml_get_type_traits_cpu(dst->type)->from_float; - float * src0_f32 = (float *) params->wdata + (ne00 + CACHE_LINE_SIZE_F32) * ith; +#ifdef __cplusplus +extern "C" { +#endif - size_t id = 0; - size_t rs = nb0 * (ne00 / ggml_blck_size(dst->type)); - char * dst_ptr = (char *) dst->data; +// +// global data +// - for (int i03 = 0; i03 < ne03; i03++) { - for (int i02 = 0; i02 < ne02; i02++) { - id += rs * ir0; - for (int i01 = ir0; i01 < ir1; i01++) { - const src_t * src0_ptr = (src_t *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); +// precomputed gelu table for f16 (128 KB) +extern ggml_fp16_t ggml_table_gelu_f16[1 << 16]; - for (int i00 = 0; i00 < ne00; i00++) { - src0_f32[i00] = type_conversion_table::to_f32(src0_ptr[i00]); - } +// precomputed quick gelu table for f16 (128 KB) +extern ggml_fp16_t ggml_table_gelu_quick_f16[1 << 16]; - quantize_row_q(src0_f32, dst_ptr + id, ne00); - id += rs; - } - id += rs * (ne01 - ir1); - } - } - } else { - // printf("%s %s\n", ggml_type_name(src0->type), ggml_type_name(dst->type)); - GGML_ABORT("not implemented"); +// +// fundamental operations +// + +void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc); +void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t * GGML_RESTRICT x, size_t bx, ggml_bf16_t * GGML_RESTRICT y, size_t by, int nrc); +void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * GGML_RESTRICT x, size_t bx, ggml_fp16_t * GGML_RESTRICT y, size_t by, int nrc); + +void ggml_vec_silu_f32(const int n, float * y, const float * x); +ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean); //it will also center y ( y = y - mean ) +ggml_float ggml_vec_soft_max_f32(const int n, float * y, const float * x, float max); +ggml_float ggml_vec_log_soft_max_f32(const int n, float * y, const float * x, float max); + +inline static void ggml_vec_set_i8(const int n, int8_t * x, const int8_t v) { for (int i = 0; i < n; ++i) x[i] = v; } +inline static void ggml_vec_set_i16(const int n, int16_t * x, const int16_t v) { for (int i = 0; i < n; ++i) x[i] = v; } + +inline static void ggml_vec_set_i32(const int n, int32_t * x, const int32_t v) { for (int i = 0; i < n; ++i) x[i] = v; } +inline static void ggml_vec_cpy_i32(const int n, int32_t * y, const int32_t * x) { for (int i = 0; i < n; ++i) y[i] = x[i]; } + +inline static void ggml_vec_set_f16(const int n, ggml_fp16_t * x, const ggml_fp16_t v) { for (int i = 0; i < n; ++i) x[i] = v; } +inline static void ggml_vec_set_bf16(const int n, ggml_bf16_t * x, const ggml_bf16_t v) { for (int i = 0; i < n; ++i) x[i] = v; } + +inline static void ggml_vec_add_f32 (const int n, float * z, const float * x, const float * y) { + int i = 0; +#if defined(__AVX2__) + for (; i + 7 < n; i += 8) { + __m256 vx = _mm256_loadu_ps(x + i); + __m256 vy = _mm256_loadu_ps(y + i); + __m256 vz = _mm256_add_ps(vx, vy); + _mm256_storeu_ps(z + i, vz); + } +#endif + for (; i < n; ++i) { + z[i] = x[i] + y[i]; } } -// A simplified version of ggml_compute_forward_dup that doesn't do float upcasting, and just plain old memcpy. -static void ggml_compute_forward_dup_bytes( - const ggml_compute_params * params, - ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; +inline static void ggml_vec_add_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { + for (int i = 0; i < n; ++i) { + z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) + GGML_CPU_FP16_TO_FP32(y[i])); + } +} +inline static void ggml_vec_add1_f32(const int n, float * z, const float * x, const float v) { for (int i = 0; i < n; ++i) z[i] = x[i] + v; } +inline static void ggml_vec_acc_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] += x[i]; } +inline static void ggml_vec_acc1_f32(const int n, float * y, const float v) { for (int i = 0; i < n; ++i) y[i] += v; } +inline static void ggml_vec_sub_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i] = x[i] - y[i]; } +inline static void ggml_vec_sub_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { + for (int i = 0; i < n; ++i) { + z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) - GGML_CPU_FP16_TO_FP32(y[i])); + } +} +inline static void ggml_vec_set_f32 (const int n, float * x, const float v) { for (int i = 0; i < n; ++i) x[i] = v; } +inline static void ggml_vec_cpy_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i]; } +inline static void ggml_vec_neg_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = -x[i]; } +inline static void ggml_vec_neg_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(-GGML_CPU_FP16_TO_FP32(x[i])); + } +} - GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src0)); - GGML_ASSERT(src0->type == dst->type); +inline static void ggml_vec_mul_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i] = x[i]*y[i]; } +inline static void ggml_vec_mul_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { + for (int i = 0; i < n; ++i) { + z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) * GGML_CPU_FP16_TO_FP32(y[i])); + } +} +inline static void ggml_vec_div_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i] = x[i]/y[i]; } +inline static void ggml_vec_div_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { + for (int i = 0; i < n; ++i) { + z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) / GGML_CPU_FP16_TO_FP32(y[i])); + } +} - GGML_TENSOR_UNARY_OP_LOCALS; +// compute GGML_VEC_DOT_UNROLL dot products at once +// xs - x row stride in bytes +inline static void ggml_vec_dot_f16_unroll(const int n, const int xs, float * GGML_RESTRICT s, void * GGML_RESTRICT xv, ggml_fp16_t * GGML_RESTRICT y) { + ggml_float sumf[GGML_VEC_DOT_UNROLL] = { 0.0 }; - if (ggml_is_contiguous(src0) && ggml_is_contiguous(dst)) { - ggml_compute_forward_dup_same_cont(params, dst); - return; + ggml_fp16_t * GGML_RESTRICT x[GGML_VEC_DOT_UNROLL]; + + for (int i = 0; i < GGML_VEC_DOT_UNROLL; ++i) { + x[i] = (ggml_fp16_t *) ((char *) xv + i*xs); } - const size_t type_size = ggml_type_size(src0->type); +#if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) - const int ith = params->ith; // thread index - const int nth = params->nth; // number of threads + const int sve_register_length = svcntb() * 8; + const int ggml_f16_epr = sve_register_length / 16; // running when 16 + const int ggml_f16_step = 8 * ggml_f16_epr; // choose 8 SVE registers - // parallelize by rows - const int nr = ne01; - // number of rows per thread - const int dr = (nr + nth - 1) / nth; - // row range for this thread - const int ir0 = dr * ith; - const int ir1 = MIN(ir0 + dr, nr); + int np = (n & ~(ggml_f16_step - 1)); - if (src0->type == dst->type && - ggml_are_same_shape(src0, dst) && - nb00 == type_size && nb0 == type_size) { - // copy by rows - const size_t rs = ggml_row_size(src0->type, ne00); - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = ir0; i01 < ir1; i01++) { - memcpy( - ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3), - ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03), - rs); - } - } - } - return; - } + svfloat16_t sum_00 = svdup_n_f16(0.0f); + svfloat16_t sum_01 = svdup_n_f16(0.0f); + svfloat16_t sum_02 = svdup_n_f16(0.0f); + svfloat16_t sum_03 = svdup_n_f16(0.0f); - if (ggml_is_contiguous(dst)) { - size_t id = 0; - char * dst_ptr = (char *) dst->data; - const size_t rs = ne00 * type_size; + svfloat16_t sum_10 = svdup_n_f16(0.0f); + svfloat16_t sum_11 = svdup_n_f16(0.0f); + svfloat16_t sum_12 = svdup_n_f16(0.0f); + svfloat16_t sum_13 = svdup_n_f16(0.0f); - if (nb00 == type_size) { - // src0 is contiguous on first dimension, copy by rows - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - id += rs * ir0; - for (int64_t i01 = ir0; i01 < ir1; i01++) { - const char * src0_ptr = (char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03; - memcpy(dst_ptr + id, src0_ptr, rs); - id += rs; - } - id += rs * (ne01 - ir1); - } - } - } else { - //printf("%s: this is not optimal - fix me\n", __func__); + svfloat16_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; + svfloat16_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - id += rs * ir0; - for (int64_t i01 = ir0; i01 < ir1; i01++) { - for (int64_t i00 = 0; i00 < ne00; i00++) { - const char * src0_ptr = (char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03; - memcpy(dst_ptr + id, src0_ptr, type_size); + for (int i = 0; i < np; i += ggml_f16_step) { + ay1 = GGML_F16x_VEC_LOAD(y + i + 0 * ggml_f16_epr, 0); // 8 elements - id += type_size; - } - } - id += rs * (ne01 - ir1); - } - } - } + ax1 = GGML_F16x_VEC_LOAD(x[0] + i + 0*ggml_f16_epr, 0); // 8 elements + sum_00 = GGML_F16x_VEC_FMA(sum_00, ax1, ay1); // sum_00 = sum_00+ax1*ay1 + ax1 = GGML_F16x_VEC_LOAD(x[1] + i + 0*ggml_f16_epr, 0); // 8 elements + sum_10 = GGML_F16x_VEC_FMA(sum_10, ax1, ay1); - return; - } + ay2 = GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 1); // next 8 elements - // dst counters - int64_t k10 = 0; - int64_t i11 = 0; - int64_t i12 = 0; - int64_t i13 = 0; + ax2 = GGML_F16x_VEC_LOAD(x[0] + i + 1*ggml_f16_epr, 1); // next 8 elements + sum_01 = GGML_F16x_VEC_FMA(sum_01, ax2, ay2); + ax2 = GGML_F16x_VEC_LOAD(x[1] + i + 1*ggml_f16_epr, 1); + sum_11 = GGML_F16x_VEC_FMA(sum_11, ax2, ay2); - // number of blocks in a row - const int64_t nk00 = ne00 / ggml_blck_size(src0->type); - const int64_t nk0 = ne0 / ggml_blck_size(dst->type); + ay3 = GGML_F16x_VEC_LOAD(y + i + 2 * ggml_f16_epr, 2); - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - k10 += nk00 * ir0; - while (k10 >= nk0) { - k10 -= nk0; - if (++i11 == ne1) { - i11 = 0; - if (++i12 == ne2) { - i12 = 0; - if (++i13 == ne3) { - i13 = 0; - } - } - } - } - for (int64_t i01 = ir0; i01 < ir1; i01++) { - for (int64_t k00 = 0; k00 < nk00; k00++) { - const char * src0_ptr = ((char *) src0->data + k00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); - char * dst_ptr = ((char *) dst->data + k10*nb0 + i11*nb1 + i12*nb2 + i13*nb3); + ax3 = GGML_F16x_VEC_LOAD(x[0] + i + 2*ggml_f16_epr, 2); + sum_02 = GGML_F16x_VEC_FMA(sum_02, ax3, ay3); + ax3 = GGML_F16x_VEC_LOAD(x[1] + i + 2*ggml_f16_epr, 2); + sum_12 = GGML_F16x_VEC_FMA(sum_12, ax3, ay3); - memcpy(dst_ptr, src0_ptr, type_size); + ay4 = GGML_F16x_VEC_LOAD(y + i + 3 * ggml_f16_epr, 3); - if (++k10 == nk0) { - k10 = 0; - if (++i11 == ne1) { - i11 = 0; - if (++i12 == ne2) { - i12 = 0; - if (++i13 == ne3) { - i13 = 0; - } - } - } - } - } - } - k10 += nk00 * (ne01 - ir1); - while (k10 >= nk0) { - k10 -= nk0; - if (++i11 == ne1) { - i11 = 0; - if (++i12 == ne2) { - i12 = 0; - if (++i13 == ne3) { - i13 = 0; - } - } - } - } - } - } -} + ax4 = GGML_F16x_VEC_LOAD(x[0] + i + 3*ggml_f16_epr, 3); + sum_03 = GGML_F16x_VEC_FMA(sum_03, ax4, ay4); + ax4 = GGML_F16x_VEC_LOAD(x[1] + i + 3*ggml_f16_epr, 3); + sum_13 = GGML_F16x_VEC_FMA(sum_13, ax4, ay4); -static void ggml_compute_forward_dup_from_q( - const ggml_compute_params * params, - ggml_tensor * dst) { + ay5 = GGML_F16x_VEC_LOAD(y + i + 4 * ggml_f16_epr, 4); - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; + ax5 = GGML_F16x_VEC_LOAD(x[0] + i + 4*ggml_f16_epr, 4); - GGML_TENSOR_BINARY_OP_LOCALS + sum_00 = GGML_F16x_VEC_FMA(sum_00, ax5, ay5); + ax5 = GGML_F16x_VEC_LOAD(x[1] + i + 4*ggml_f16_epr, 4); + sum_10 = GGML_F16x_VEC_FMA(sum_10, ax5, ay5); - const ggml_type type = src0->type; - ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; + ay6 = GGML_F16x_VEC_LOAD(y + i + 5 * ggml_f16_epr, 5); - size_t qk = ggml_blck_size(type); - const int64_t nr = ggml_nelements(src1) / qk; + ax6 = GGML_F16x_VEC_LOAD(x[0] + i + 5*ggml_f16_epr, 5); - // destination must be contiguous in the first dimension - GGML_ASSERT(nb10 == ggml_type_size(dst->type)); - // must either have first dimension large enough to hold a row, or fully contiguous - GGML_ASSERT((ne10 % qk) == 0 || ggml_is_contiguous(dst)); + sum_01 = GGML_F16x_VEC_FMA(sum_01, ax6, ay6); + ax6 = GGML_F16x_VEC_LOAD(x[1] + i + 5*ggml_f16_epr, 5); + sum_11 = GGML_F16x_VEC_FMA(sum_11, ax6, ay6); - const int ith = params->ith; - const int nth = params->nth; + ay7 = GGML_F16x_VEC_LOAD(y + i + 6 * ggml_f16_epr, 6); - const int dr = (nr + nth - 1)/nth; + ax7 = GGML_F16x_VEC_LOAD(x[0] + i + 6*ggml_f16_epr, 6); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + sum_02 = GGML_F16x_VEC_FMA(sum_02, ax7, ay7); + ax7 = GGML_F16x_VEC_LOAD(x[1] + i + 6*ggml_f16_epr, 6); + sum_12 = GGML_F16x_VEC_FMA(sum_12, ax7, ay7); - for (int64_t ir = ir0; ir < ir1; ++ir) { + ay8 = GGML_F16x_VEC_LOAD(y + i + 7 * ggml_f16_epr, 7); - uint32_t i = ir * qk; + ax8 = GGML_F16x_VEC_LOAD(x[0] + i + 7*ggml_f16_epr, 7); - const int64_t i03 = i/(ne00 * ne01 * ne02); - const int64_t i02 = (i - i03*ne00*ne01*ne02 )/ (ne00*ne01); - const int64_t i01 = (i - i03*ne00*ne01*ne02 - i02*ne01*ne00) / ne00; - const int64_t i00 = i - i03*ne00*ne01*ne02 - i02*ne01*ne00 - i01*ne00; - const int64_t x_offset = (i00/qk)*nb00 + i01*nb01 + i02*nb02 + i03 * nb03; + sum_03 = GGML_F16x_VEC_FMA(sum_03, ax8, ay8); + ax8 = GGML_F16x_VEC_LOAD(x[1] + i + 7*ggml_f16_epr, 7); + sum_13 = GGML_F16x_VEC_FMA(sum_13, ax8, ay8); + } - const int64_t i13 = i/(ne10 * ne11 * ne12); - const int64_t i12 = (i - i13*ne10*ne11*ne12) / (ne10*ne11); - const int64_t i11 = (i - i13*ne10*ne11*ne12 - i12*ne10*ne11) / ne10; - const int64_t i10 = i - i13*ne10*ne11*ne12 - i12*ne10*ne11 - i11*ne10; - const int64_t dst_offset = i10*nb10 + i11*nb11 + i12*nb12 + i13*nb13; + const int np2 = (n & ~(ggml_f16_epr - 1)); + for (int k = np; k < np2; k += ggml_f16_epr) { + svfloat16_t ry = GGML_F16x_VEC_LOAD(y + k, 0); - dequantize_row_q( - (const void *) ((char *) src0->data + x_offset), - (float *) ((char *) dst->data + dst_offset), qk); - } -} + svfloat16_t rx = GGML_F16x_VEC_LOAD(x[0] + k, 0); + sum_00 = GGML_F16x_VEC_FMA(sum_00, rx, ry); + rx = GGML_F16x_VEC_LOAD(x[1] + k, 0); + sum_10 = GGML_F16x_VEC_FMA(sum_10, rx, ry); + } -void ggml_compute_forward_dup( - const ggml_compute_params * params, - ggml_tensor * dst) { + if (np2 < n) { + svbool_t pg = svwhilelt_b16(np2, n); + svfloat16_t hx_0 = svld1_f16(pg, (const __fp16 *)(x[0] + np2)); + svfloat16_t hx_1 = svld1_f16(pg, (const __fp16 *)(x[1] + np2)); + svfloat16_t hy = svld1_f16(pg, (const __fp16 *)(y + np2)); - const ggml_tensor * src0 = dst->src[0]; + sum_00 = svmad_f16_x(pg, hx_0, hy, sum_00); + sum_10 = svmad_f16_x(pg, hx_1, hy, sum_10); + } + GGML_F16x_VEC_REDUCE(sumf[0], sum_00, sum_01, sum_02, sum_03); + GGML_F16x_VEC_REDUCE(sumf[1], sum_10, sum_11, sum_12, sum_13); + np = n; + #elif defined(__riscv_v_intrinsic) + #if defined(__riscv_zvfh) + size_t vl = __riscv_vsetvlmax_e32m4(); + + // initialize accumulators to all zeroes + vfloat32m4_t vsum0_0 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + vfloat32m4_t vsum0_1 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + vfloat32m4_t vsum1_0 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + vfloat32m4_t vsum1_1 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + + // calculate step size + const size_t epr = __riscv_vsetvlmax_e16m2(); + const size_t step = epr * 2; + int np = (n & ~(step - 1)); + + // unroll by 2 along the row dimension + for (int i = 0; i < np; i += step) { + vfloat16m2_t ay0 = __riscv_vle16_v_f16m2((const _Float16 *)(y + i), epr); + vfloat16m2_t ax0_0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i), epr); + vfloat16m2_t ax1_0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i), epr); + vsum0_0 = __riscv_vfwmacc_vv_f32m4(vsum0_0, ax0_0, ay0, epr); + vsum1_0 = __riscv_vfwmacc_vv_f32m4(vsum1_0, ax1_0, ay0, epr); + + vfloat16m2_t ay1 = __riscv_vle16_v_f16m2((const _Float16 *)(y + i + epr), epr); + vfloat16m2_t ax0_1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i + epr), epr); + vfloat16m2_t ax1_1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i + epr), epr); + vsum0_1 = __riscv_vfwmacc_vv_f32m4(vsum0_1, ax0_1, ay1, epr); + vsum1_1 = __riscv_vfwmacc_vv_f32m4(vsum1_1, ax1_1, ay1, epr); + } + + vfloat32m4_t vsum0 = __riscv_vfadd_vv_f32m4(vsum0_0, vsum0_1, vl); + vfloat32m4_t vsum1 = __riscv_vfadd_vv_f32m4(vsum1_0, vsum1_1, vl); + + // leftovers + for (int i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m2(n - i); + vfloat16m2_t ay = __riscv_vle16_v_f16m2((const _Float16 *)(y + i), vl); + vfloat16m2_t ax0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i), vl); + vfloat16m2_t ax1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i), vl); + + vsum0 = __riscv_vfwmacc_vv_f32m4(vsum0, ax0, ay, vl); + vsum1 = __riscv_vfwmacc_vv_f32m4(vsum1, ax1, ay, vl); + } + + // reduce + vl = __riscv_vsetvlmax_e32m2(); + vfloat32m2_t acc0_0 = __riscv_vfadd_vv_f32m2(__riscv_vget_v_f32m4_f32m2(vsum0, 0), + __riscv_vget_v_f32m4_f32m2(vsum0, 1), vl); + vl = __riscv_vsetvlmax_e32m1(); + vfloat32m1_t acc0_1 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(acc0_0, 0), + __riscv_vget_v_f32m2_f32m1(acc0_0, 1), vl); + vfloat32m1_t redsum0 = __riscv_vfredusum_vs_f32m1_f32m1( + acc0_1, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl); + + vl = __riscv_vsetvlmax_e32m2(); + vfloat32m2_t acc1_0 = __riscv_vfadd_vv_f32m2(__riscv_vget_v_f32m4_f32m2(vsum1, 0), + __riscv_vget_v_f32m4_f32m2(vsum1, 1), vl); + vl = __riscv_vsetvlmax_e32m1(); + vfloat32m1_t acc1_1 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(acc1_0, 0), + __riscv_vget_v_f32m2_f32m1(acc1_0, 1), vl); + vfloat32m1_t redsum1 = __riscv_vfredusum_vs_f32m1_f32m1( + acc1_1, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl); + sumf[0] = __riscv_vfmv_f_s_f32m1_f32(redsum0); + sumf[1] = __riscv_vfmv_f_s_f32m1_f32(redsum1); + np = n; + #else + const int np = 0; + #endif + #else + const int np = (n & ~(GGML_F16_STEP - 1)); - if (src0->type == dst->type) { - ggml_compute_forward_dup_bytes(params, dst); - return; - } + GGML_F16_VEC sum[GGML_VEC_DOT_UNROLL][GGML_F16_ARR] = { { GGML_F16_VEC_ZERO } }; - switch (src0->type) { - case GGML_TYPE_F16: - { - /**/ if (dst->type == GGML_TYPE_F16) ggml_compute_forward_dup_flt(params, dst); - else if (dst->type == GGML_TYPE_BF16) ggml_compute_forward_dup_flt(params, dst); - else if (dst->type == GGML_TYPE_F32) ggml_compute_forward_dup_flt(params, dst); - else ggml_compute_forward_dup_to_q(params, dst); - } break; - case GGML_TYPE_BF16: - { - /**/ if (dst->type == GGML_TYPE_F16) ggml_compute_forward_dup_flt(params, dst); - else if (dst->type == GGML_TYPE_BF16) ggml_compute_forward_dup_flt(params, dst); - else if (dst->type == GGML_TYPE_F32) ggml_compute_forward_dup_flt(params, dst); - else ggml_compute_forward_dup_to_q(params, dst); - } break; - case GGML_TYPE_F32: - { - /**/ if (dst->type == GGML_TYPE_F16) ggml_compute_forward_dup_flt(params, dst); - else if (dst->type == GGML_TYPE_BF16) ggml_compute_forward_dup_flt(params, dst); - else if (dst->type == GGML_TYPE_F32) ggml_compute_forward_dup_flt(params, dst); - else if (dst->type == GGML_TYPE_I32) ggml_compute_forward_dup_flt(params, dst); - else ggml_compute_forward_dup_to_q(params, dst); - } break; - case GGML_TYPE_I32: - { - if (dst->type == GGML_TYPE_F32) ggml_compute_forward_dup_flt(params, dst); - else GGML_ABORT("not implemented"); - } break; - default: - { - if (ggml_is_quantized(src0->type) && dst->type == GGML_TYPE_F32) { - ggml_compute_forward_dup_from_q(params, dst); - break; + GGML_F16_VEC ax[GGML_F16_ARR]; + GGML_F16_VEC ay[GGML_F16_ARR]; + + for (int i = 0; i < np; i += GGML_F16_STEP) { + for (int j = 0; j < GGML_F16_ARR; j++) { + ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); + + for (int k = 0; k < GGML_VEC_DOT_UNROLL; ++k) { + ax[j] = GGML_F16_VEC_LOAD(x[k] + i + j*GGML_F16_EPR, j); + + sum[k][j] = GGML_F16_VEC_FMA(sum[k][j], ax[j], ay[j]); } - GGML_ABORT("fatal error"); } + } + + // reduce sum0..sum3 to sum0 + for (int k = 0; k < GGML_VEC_DOT_UNROLL; ++k) { + GGML_F16_VEC_REDUCE(sumf[k], sum[k]); + } + #endif +#else + // scalar path + const int np = 0; +#endif + // scalar and leftovers + for (int i = np; i < n; ++i) { + for (int j = 0; j < GGML_VEC_DOT_UNROLL; ++j) { + sumf[j] += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[j][i])*GGML_CPU_FP16_TO_FP32(y[i])); + } + } + + for (int i = 0; i < GGML_VEC_DOT_UNROLL; ++i) { + s[i] = (float)sumf[i]; } } -// ggml_compute_forward_add +inline static void ggml_vec_mad_f32(const int n, float * GGML_RESTRICT y, const float * GGML_RESTRICT x, const float v) { +#if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) -static void ggml_compute_forward_add_q_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { + const int sve_register_length = ggml_cpu_get_sve_cnt() * 8; + const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16 + const int ggml_f32_step = 8 * ggml_f32_epr; // choose 8 SVE registers + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; + const int np = (n & ~(ggml_f32_step - 1)); + svfloat32_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; + svfloat32_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; + for (int i = 0; i < np; i += ggml_f32_step) { - GGML_ASSERT(ggml_are_same_shape(src0, src1) && ggml_are_same_shape(src0, dst)); + ax1 = GGML_F32_VEC_LOAD(x + i); + ay1 = GGML_F32_VEC_LOAD(y + i); + ay1 = GGML_F32_VEC_FMA(ay1, ax1, vx); - const int nr = ggml_nrows(src0); + GGML_F32_VEC_STORE(y + i, ay1); - GGML_TENSOR_BINARY_OP_LOCALS + ax2 = GGML_F32_VEC_LOAD(x + i + 1*ggml_f32_epr); + ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr); + ay2 = GGML_F32_VEC_FMA(ay2, ax2, vx); - const int ith = params->ith; - const int nth = params->nth; + GGML_F32_VEC_STORE(y + i + 1*ggml_f32_epr, ay2); - const ggml_type type = src0->type; - const ggml_type dtype = dst->type; - ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; - ggml_from_float_t const quantize_row_q = ggml_get_type_traits_cpu(dtype)->from_float; + ax3 = GGML_F32_VEC_LOAD(x + i + 2*ggml_f32_epr); + ay3 = GGML_F32_VEC_LOAD(y + i + 2*ggml_f32_epr); + ay3 = GGML_F32_VEC_FMA(ay3, ax3, vx); - // we don't support permuted src0 or src1 - GGML_ASSERT(nb00 == ggml_type_size(type)); - GGML_ASSERT(nb10 == sizeof(float)); + GGML_F32_VEC_STORE(y + i + 2*ggml_f32_epr, ay3); - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); + ax4 = GGML_F32_VEC_LOAD(x + i + 3*ggml_f32_epr); + ay4 = GGML_F32_VEC_LOAD(y + i + 3*ggml_f32_epr); + ay4 = GGML_F32_VEC_FMA(ay4, ax4, vx); - GGML_ASSERT(ggml_is_quantized(src0->type)); - GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_F32_VEC_STORE(y + i + 3*ggml_f32_epr, ay4); - // rows per thread - const int dr = (nr + nth - 1)/nth; + ax5 = GGML_F32_VEC_LOAD(x + i + 4*ggml_f32_epr); + ay5 = GGML_F32_VEC_LOAD(y + i + 4*ggml_f32_epr); + ay5 = GGML_F32_VEC_FMA(ay5, ax5, vx); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + GGML_F32_VEC_STORE(y + i + 4*ggml_f32_epr, ay5); - float * wdata = (float *) params->wdata + (ne00 + CACHE_LINE_SIZE_F32) * ith; + ax6 = GGML_F32_VEC_LOAD(x + i + 5*ggml_f32_epr); + ay6 = GGML_F32_VEC_LOAD(y + i + 5*ggml_f32_epr); + ay6 = GGML_F32_VEC_FMA(ay6, ax6, vx); - for (int ir = ir0; ir < ir1; ++ir) { - // src0 indices - const int i03 = ir/(ne02*ne01); - const int i02 = (ir - i03*ne02*ne01)/ne01; - const int i01 = (ir - i03*ne02*ne01 - i02*ne01); + GGML_F32_VEC_STORE(y + i + 5*ggml_f32_epr, ay6); - // src1 and dst are same shape as src0 => same indices - const int i13 = i03; - const int i12 = i02; - const int i11 = i01; + ax7 = GGML_F32_VEC_LOAD(x + i + 6*ggml_f32_epr); + ay7 = GGML_F32_VEC_LOAD(y + i + 6*ggml_f32_epr); + ay7 = GGML_F32_VEC_FMA(ay7, ax7, vx); - const int i3 = i03; - const int i2 = i02; - const int i1 = i01; + GGML_F32_VEC_STORE(y + i + 6*ggml_f32_epr, ay7); - void * src0_row = (void *) ((char *) src0->data + (i01*nb01 + i02*nb02 + i03*nb03)); - float * src1_row = (float *)((char *) src1->data + (i11*nb11 + i12*nb12 + i13*nb13)); - void * dst_row = (void *) ((char *) dst->data + ( i1*nb1 + i2*nb2 + i3*nb3)); + ax8 = GGML_F32_VEC_LOAD(x + i + 7*ggml_f32_epr); + ay8 = GGML_F32_VEC_LOAD(y + i + 7*ggml_f32_epr); + ay8 = GGML_F32_VEC_FMA(ay8, ax8, vx); - assert(ne00 % 32 == 0); + GGML_F32_VEC_STORE(y + i + 7*ggml_f32_epr, ay8); + } + // leftovers + // Since 8 unrolls are done in above loop, leftovers lie in range [0, ggml_f32_step] which is handled in below loop + const int np2 = (n & ~(ggml_f32_epr - 1)); + for (int i = np; i < np2; i += ggml_f32_epr) { + ax1 = GGML_F32_VEC_LOAD(x + i); + ay1 = GGML_F32_VEC_LOAD(y + i); + ay1 = GGML_F32_VEC_FMA(ay1, ax1, vx); - // unquantize row from src0 to temp buffer - dequantize_row_q(src0_row, wdata, ne00); - // add src1 - ggml_vec_acc_f32(ne00, wdata, src1_row); - // quantize row to dst - if (quantize_row_q != NULL) { - quantize_row_q(wdata, dst_row, ne00); - } else { - memcpy(dst_row, wdata, ne0*nb0); + GGML_F32_VEC_STORE(y + i, ay1); } - } -} + // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only + if (np2 < n) { + svbool_t pg =svwhilelt_b32(np2, n); + ax1 = svld1_f32(pg, x + np2); + ay1 = svld1_f32(pg, y + np2); + ay1 = svmad_f32_m(pg, ax1, vx, ay1); -void ggml_compute_forward_add( - const ggml_compute_params * params, - ggml_tensor * dst) { + svst1_f32(pg, y + np2, ay1); + } + #elif defined(__riscv_v_intrinsic) + for (int i = 0, avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m8(n - i); + vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[i], avl); + vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl); + vfloat32m8_t ny = __riscv_vfmadd_vf_f32m8(ax, v, ay, avl); + __riscv_vse32_v_f32m8(&y[i], ny, avl); + } + #else + const int np = (n & ~(GGML_F32_STEP - 1)); - const ggml_tensor * src0 = dst->src[0]; + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); - switch (src0->type) { - case GGML_TYPE_F32: - case GGML_TYPE_F16: - case GGML_TYPE_BF16: - { - ggml_compute_forward_add_non_quantized(params, dst); - } break; - case GGML_TYPE_Q1_0: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q8_0: - case GGML_TYPE_MXFP4: - case GGML_TYPE_NVFP4: - case GGML_TYPE_Q2_K: - case GGML_TYPE_Q3_K: - case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: - case GGML_TYPE_Q6_K: - case GGML_TYPE_TQ1_0: - case GGML_TYPE_TQ2_0: - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ2_S: - { - ggml_compute_forward_add_q_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); + GGML_F32_VEC ax[GGML_F32_ARR]; + GGML_F32_VEC ay[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ax[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_FMA(ay[j], ax[j], vx); + + GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); } + } + + // leftovers + for (int i = np; i < n; ++i) { + y[i] += x[i]*v; + } + #endif +#else + // scalar + for (int i = 0; i < n; ++i) { + y[i] += x[i]*v; } +#endif } -// ggml_compute_forward_add_id +inline static void ggml_vec_mad_f16(const int n, ggml_fp16_t * GGML_RESTRICT y, const ggml_fp16_t * GGML_RESTRICT x, const float v) { +#if defined(GGML_SIMD) && defined(__ARM_FEATURE_SVE) + const int sve_register_length = svcntb() * 8; + const int ggml_f16_epr = sve_register_length / 16; + const int ggml_f16_step = 8 * ggml_f16_epr; -static void ggml_compute_forward_add_id_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { + GGML_F16x_VEC vx = GGML_F16x_VEC_SET1(v); - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - const ggml_tensor * src2 = dst->src[2]; + int np = (n & ~(ggml_f16_step - 1)); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(src2->type == GGML_TYPE_I32); + svfloat16_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; + svfloat16_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; + for (int i = 0; i < np; i += ggml_f16_step) { + ax1 = GGML_F16x_VEC_LOAD(x + i + 0 * ggml_f16_epr, 0); + ay1 = GGML_F16x_VEC_LOAD(y + i + 0 * ggml_f16_epr, 0); + ay1 = GGML_F16x_VEC_FMA(ay1, ax1, vx); - GGML_ASSERT(src0->nb[0] == sizeof(float)); - GGML_ASSERT(src1->nb[0] == sizeof(float)); + GGML_F16x_VEC_STORE(y + i + 0 * ggml_f16_epr, ay1, 0); - const int ith = params->ith; - const int nth = params->nth; + ax2 = GGML_F16x_VEC_LOAD(x + i + 1 * ggml_f16_epr, 1); + ay2 = GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 1); + ay2 = GGML_F16x_VEC_FMA(ay2, ax2, vx); - const int nr = ggml_nrows(src0); + GGML_F16x_VEC_STORE(y + i + 1 * ggml_f16_epr, ay2, 1); - GGML_TENSOR_TERNARY_OP_LOCALS + ax3 = GGML_F16x_VEC_LOAD(x + i + 2 * ggml_f16_epr, 2); + ay3 = GGML_F16x_VEC_LOAD(y + i + 2 * ggml_f16_epr, 2); + ay3 = GGML_F16x_VEC_FMA(ay3, ax3, vx); - GGML_ASSERT( nb0 == sizeof(float)); - GGML_ASSERT(nb10 == sizeof(float)); + GGML_F16x_VEC_STORE(y + i + 2 * ggml_f16_epr, ay3, 2); - // rows per thread - const int dr = (nr + nth - 1)/nth; + ax4 = GGML_F16x_VEC_LOAD(x + i + 3 * ggml_f16_epr, 3); + ay4 = GGML_F16x_VEC_LOAD(y + i + 3 * ggml_f16_epr, 3); + ay4 = GGML_F16x_VEC_FMA(ay4, ax4, vx); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + GGML_F16x_VEC_STORE(y + i + 3 * ggml_f16_epr, ay4, 3); - for (int ir = ir0; ir < ir1; ++ir) { - // src0 indices - const int i3 = ir/(ne2*ne1); - const int i2 = (ir - i3*ne2*ne1)/ne1; - const int i1 = (ir - i3*ne2*ne1 - i2*ne1); + ax5 = GGML_F16x_VEC_LOAD(x + i + 4 * ggml_f16_epr, 4); + ay5 = GGML_F16x_VEC_LOAD(y + i + 4 * ggml_f16_epr, 4); + ay5 = GGML_F16x_VEC_FMA(ay5, ax5, vx); - // src1 indices - const int i11 = *(int32_t *) ((char *) src2->data + i1*nb20 + i2*nb21); + GGML_F16x_VEC_STORE(y + i + 4 * ggml_f16_epr, ay5, 4); - GGML_ASSERT(i11 >= 0 && i11 < ne11); + ax6 = GGML_F16x_VEC_LOAD(x + i + 5 * ggml_f16_epr, 5); + ay6 = GGML_F16x_VEC_LOAD(y + i + 5 * ggml_f16_epr, 5); + ay6 = GGML_F16x_VEC_FMA(ay6, ax6, vx); - ggml_vec_add_f32(ne0, - (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ), - (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01), - (float *) ((char *) src1->data + i11*nb11)); - } -} + GGML_F16x_VEC_STORE(y + i + 5 * ggml_f16_epr, ay6, 5); -void ggml_compute_forward_add_id( - const ggml_compute_params * params, - ggml_tensor * dst) { + ax7 = GGML_F16x_VEC_LOAD(x + i + 6 * ggml_f16_epr, 6); + ay7 = GGML_F16x_VEC_LOAD(y + i + 6 * ggml_f16_epr, 6); + ay7 = GGML_F16x_VEC_FMA(ay7, ax7, vx); - const ggml_tensor * src0 = dst->src[0]; + GGML_F16x_VEC_STORE(y + i + 6 * ggml_f16_epr, ay7, 6); - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_add_id_f32(params, dst); - } break; - default: - { - GGML_ABORT("unsupported type for ggml_compute_forward_add_id: %s", ggml_type_name(src0->type)); - } + ax8 = GGML_F16x_VEC_LOAD(x + i + 7 * ggml_f16_epr, 7); + ay8 = GGML_F16x_VEC_LOAD(y + i + 7 * ggml_f16_epr, 7); + ay8 = GGML_F16x_VEC_FMA(ay8, ax8, vx); + + GGML_F16x_VEC_STORE(y + i + 7 * ggml_f16_epr, ay8, 7); } -} + const int np2 = (n & ~(ggml_f16_epr - 1)); + for (int k = np; k < np2; k += ggml_f16_epr) { + svfloat16_t rx = GGML_F16x_VEC_LOAD(x + k, 0); + svfloat16_t ry = GGML_F16x_VEC_LOAD(y + k, 0); + ry = GGML_F16x_VEC_FMA(ry, rx, vx); -// ggml_compute_forward_add1 + GGML_F16x_VEC_STORE(y + k, ry, 0); + } -static void ggml_compute_forward_add1_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - - GGML_ASSERT(ggml_are_same_shape(src0, dst)); - GGML_ASSERT(ggml_is_scalar(src1)); - - const int ith = params->ith; - const int nth = params->nth; + if (np2 < n) { + svbool_t pg = svwhilelt_b16(np2, n); + svfloat16_t hx = svld1_f16(pg, (const __fp16 *)(x + np2)); + svfloat16_t hy = svld1_f16(pg, (const __fp16 *)(y + np2)); + hy = svmad_f16_x(pg, hx, vx, hy); + svst1_f16(pg, (__fp16 *)(y + np2), hy); + } + np = n; +#elif defined(__riscv_v_intrinsic) // implies __riscv_v_intrinsic + #if defined (__riscv_zvfh) + const ggml_fp16_t s = GGML_CPU_FP32_TO_FP16(v); + const _Float16 scale = *(const _Float16*)(&s); - const int nr = ggml_nrows(src0); + // calculate step size + const int epr = __riscv_vsetvlmax_e16m4(); + const int step = epr * 2; + int np = (n & ~(step - 1)); - GGML_TENSOR_UNARY_OP_LOCALS + // unroll by 2 + for (int i = 0; i < np; i += step) { + vfloat16m4_t ax0 = __riscv_vle16_v_f16m4((const _Float16*)x + i, epr); + vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, epr); + ay0 = __riscv_vfmacc_vf_f16m4(ay0, scale, ax0, epr); + __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, epr); + __asm__ __volatile__ ("" ::: "memory"); - GGML_ASSERT( nb0 == sizeof(float)); - GGML_ASSERT(nb00 == sizeof(float)); + vfloat16m4_t ax1 = __riscv_vle16_v_f16m4((const _Float16*)x + i + epr, epr); + vfloat16m4_t ay1 = __riscv_vle16_v_f16m4((const _Float16*)y + i + epr, epr); + ay1 = __riscv_vfmacc_vf_f16m4(ay1, scale, ax1, epr); + __riscv_vse16_v_f16m4((_Float16*)y + i + epr, ay1, epr); + __asm__ __volatile__ ("" ::: "memory"); + } - // rows per thread - const int dr = (nr + nth - 1)/nth; + // leftovers + int vl; + for (int i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m4(n - i); + vfloat16m4_t ax0 = __riscv_vle16_v_f16m4((const _Float16*)x + i, vl); + vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, vl); + ay0 = __riscv_vfmacc_vf_f16m4(ay0, scale, ax0, vl); + __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, vl); + } + np = n; + #else + // fall to scalar path + const int np = 0; + #endif +#elif defined(GGML_SIMD) + const int np = (n & ~(GGML_F16_STEP - 1)); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + GGML_F16_VEC vx = GGML_F16_VEC_SET1(v); - for (int ir = ir0; ir < ir1; ++ir) { - // src0 and dst are same shape => same indices - const int i3 = ir/(ne2*ne1); - const int i2 = (ir - i3*ne2*ne1)/ne1; - const int i1 = (ir - i3*ne2*ne1 - i2*ne1); + GGML_F16_VEC ax[GGML_F16_ARR]; + GGML_F16_VEC ay[GGML_F16_ARR]; -#ifdef GGML_USE_ACCELERATE - GGML_UNUSED(ggml_vec_add1_f32); + for (int i = 0; i < np; i += GGML_F16_STEP) { + for (int j = 0; j < GGML_F16_ARR; j++) { + ax[j] = GGML_F16_VEC_LOAD(x + i + j*GGML_F16_EPR, j); + ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); + ay[j] = GGML_F16_VEC_FMA(ay[j], ax[j], vx); - vDSP_vadd( - (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01), 1, - (float *) ((char *) src1->data), 0, - (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ), 1, - ne0); + GGML_F16_VEC_STORE(y + i + j*GGML_F16_EPR, ay, j); + } + } #else - ggml_vec_add1_f32(ne0, - (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ), - (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01), - *(float *) src1->data); + // scalar path + const int np = 0; #endif + + // scalar and leftovers + for (int i = np; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(y[i]) + GGML_CPU_FP16_TO_FP32(x[i])*v); } } -static void ggml_compute_forward_add1_f16_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - - GGML_ASSERT(ggml_are_same_shape(src0, dst)); - GGML_ASSERT(ggml_is_scalar(src1)); +// xs and vs are byte strides of x and v +inline static void ggml_vec_mad_f32_unroll(const int n, const int xs, const int vs, float * GGML_RESTRICT y, const float * GGML_RESTRICT xv, const float * GGML_RESTRICT vv) { - // scalar to add - const float v = *(float *) src1->data; + const float * GGML_RESTRICT x[GGML_VEC_MAD_UNROLL]; + const float * GGML_RESTRICT v[GGML_VEC_MAD_UNROLL]; - const int ith = params->ith; - const int nth = params->nth; + for (int i = 0; i < GGML_VEC_MAD_UNROLL; ++i) { + x[i] = (const float *) ((const char *) xv + i*xs); + v[i] = (const float *) ((const char *) vv + i*vs); + } - const int nr = ggml_nrows(src0); +#if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + // scalar Route to scalar implementation //TODO: Write SVE code + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + for (int i = 0; i < n; ++i) { + y[i] += x[k][i]*v[k][0]; + } + } + #elif defined(__riscv_v_intrinsic) + for (int i = 0, avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m8(n - i); + vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl); + for (int k = 0; k < GGML_VEC_MAD_UNROLL; k++) { + vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[k][i], avl); + ay = __riscv_vfmadd_vf_f32m8(ax, v[k][0], ay, avl); + } + __riscv_vse32_v_f32m8(&y[i], ay, avl); + } + #else + const int np = (n & ~(GGML_F32_STEP - 1)); - GGML_TENSOR_UNARY_OP_LOCALS + GGML_F32_VEC vx[GGML_VEC_MAD_UNROLL]; - GGML_ASSERT(src0->type == GGML_TYPE_F16); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F16); + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + vx[k] = GGML_F32_VEC_SET1(v[k][0]); + } - GGML_ASSERT( nb0 == sizeof(ggml_fp16_t)); - GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); + GGML_F32_VEC ax[GGML_VEC_MAD_UNROLL][GGML_F32_ARR]; + GGML_F32_VEC ay[GGML_F32_ARR]; - // rows per thread - const int dr = (nr + nth - 1)/nth; + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + ax[k][j] = GGML_F32_VEC_LOAD(x[k] + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_FMA(ay[j], ax[k][j], vx[k]); + } - for (int ir = ir0; ir < ir1; ++ir) { - // src0 and dst are same shape => same indices - const int i3 = ir/(ne2*ne1); - const int i2 = (ir - i3*ne2*ne1)/ne1; - const int i1 = (ir - i3*ne2*ne1 - i2*ne1); + GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); + } + } - ggml_fp16_t * dst_ptr = (ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ); - ggml_fp16_t * src0_ptr = (ggml_fp16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01); - for (int i = 0; i < ne0; i++) { - dst_ptr[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(src0_ptr[i]) + v); + // leftovers + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + for (int i = np; i < n; ++i) { + y[i] += x[k][i]*v[k][0]; + } + } + #endif +#else + // scalar + for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { + for (int i = 0; i < n; ++i) { + y[i] += x[k][i]*v[k][0]; } } +#endif } -static void ggml_compute_forward_add1_f16_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; +inline static void ggml_vec_mad1_f32(const int n, float * y, const float * x, const float s, const float b) { +#if defined(GGML_USE_ACCELERATE) + vDSP_vsmsa(x, 1, &s, &b, y, 1, n); +#elif defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + // scalar ; TODO: Write SVE code + for (int i = 0; i < n; ++i) { + y[i] = x[i]*s + b; + } + #elif defined(__riscv_v_intrinsic) + for (int i = 0, avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m8(n - i); + vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[i], avl); + vfloat32m8_t vb = __riscv_vfmv_v_f_f32m8(b, avl); + vfloat32m8_t ny = __riscv_vfmadd_vf_f32m8(ax, s, vb, avl); + __riscv_vse32_v_f32m8(&y[i], ny, avl); + } + #else + const int np = (n & ~(GGML_F32_STEP - 1)); - GGML_ASSERT(ggml_are_same_shape(src0, dst)); - GGML_ASSERT(ggml_is_scalar(src1)); + GGML_F32_VEC vs = GGML_F32_VEC_SET1(s); + GGML_F32_VEC vb = GGML_F32_VEC_SET1(b); - // scalar to add - const float v = GGML_CPU_FP16_TO_FP32(*(ggml_fp16_t *) src1->data); + GGML_F32_VEC ay[GGML_F32_ARR]; - const int ith = params->ith; - const int nth = params->nth; + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ay[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_FMA(vb, ay[j], vs); - const int nr = ggml_nrows(src0); + GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); + } + } - GGML_TENSOR_UNARY_OP_LOCALS + // leftovers + for (int i = np; i < n; ++i) { + y[i] = x[i]*s + b; + } + #endif +#else + // scalar + for (int i = 0; i < n; ++i) { + y[i] = x[i]*s + b; + } +#endif +} - GGML_ASSERT(src0->type == GGML_TYPE_F16); - GGML_ASSERT(src1->type == GGML_TYPE_F16); - GGML_ASSERT(dst->type == GGML_TYPE_F16); +//inline static void ggml_vec_scale_f32(const int n, float * y, const float v) { for (int i = 0; i < n; ++i) y[i] *= v; } +inline static void ggml_vec_scale_f32(const int n, float * y, const float v) { +#if defined(GGML_USE_ACCELERATE) + vDSP_vsmul(y, 1, &v, y, 1, n); +#elif defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + const int sve_register_length = ggml_cpu_get_sve_cnt() * 8; + const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16 + const int ggml_f32_step = 2 * ggml_f32_epr; + + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); + const int np = (n & ~(ggml_f32_step - 1)); + svfloat32_t ay1; + svfloat32_t ay2; + for (int i = 0; i < np; i += ggml_f32_step) { + ay1 = GGML_F32_VEC_LOAD(y + i); + ay1 = GGML_F32_VEC_MUL(ay1, vx); + GGML_F32_VEC_STORE(y + i, ay1); + + ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr); + ay2 = GGML_F32_VEC_MUL(ay2, vx); + GGML_F32_VEC_STORE(y + i + 1*ggml_f32_epr, ay2); + } + // leftovers + // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only + for (int i = np; i < n; i += ggml_f32_epr) { + svbool_t pg = svwhilelt_b32(i, n); + ay1 = svld1_f32(pg, y + i); + ay1 = svmul_f32_m(pg, ay1, vx); + svst1_f32(pg, y + i, ay1); + } + #elif defined(__riscv_v_intrinsic) + for (int i = 0, avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m8(n - i); + vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl); + vfloat32m8_t ny = __riscv_vfmul_vf_f32m8(ay, v, avl); + __riscv_vse32_v_f32m8(&y[i], ny, avl); + } + #else + const int np = (n & ~(GGML_F32_STEP - 1)); - GGML_ASSERT( nb0 == sizeof(ggml_fp16_t)); - GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); + GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); - // rows per thread - const int dr = (nr + nth - 1)/nth; + GGML_F32_VEC ay[GGML_F32_ARR]; - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_MUL(ay[j], vx); - for (int ir = ir0; ir < ir1; ++ir) { - // src0 and dst are same shape => same indices - const int i3 = ir/(ne2*ne1); - const int i2 = (ir - i3*ne2*ne1)/ne1; - const int i1 = (ir - i3*ne2*ne1 - i2*ne1); + GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); + } + } - ggml_fp16_t * dst_ptr = (ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ); - ggml_fp16_t * src0_ptr = (ggml_fp16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01); - for (int i = 0; i < ne0; i++) { - dst_ptr[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(src0_ptr[i]) + v); + // leftovers + for (int i = np; i < n; ++i) { + y[i] *= v; } + #endif +#else + // scalar + for (int i = 0; i < n; ++i) { + y[i] *= v; } +#endif } -static void ggml_compute_forward_add1_q_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { +inline static void ggml_vec_scale_f16(const int n, ggml_fp16_t * y, const float v) { +#if defined(GGML_SIMD) && defined(__ARM_FEATURE_SVE) + const int sve_register_length = svcntb() * 8; + const int ggml_f16_epr = sve_register_length / 16; + const int ggml_f16_step = 2 * ggml_f16_epr; + + GGML_F16x_VEC vx = GGML_F16x_VEC_SET1(v); + int np = (n & ~(ggml_f16_step - 1)); + svfloat16_t ay1, ay2; + + for (int i = 0; i < np; i += ggml_f16_step) { + ay1 = GGML_F16x_VEC_LOAD(y + i + 0*ggml_f16_epr, 0); + ay1 = GGML_F16x_VEC_MUL(ay1, vx); + GGML_F16x_VEC_STORE(y + i + 0*ggml_f16_epr, ay1, 0); + + ay2 = GGML_F16x_VEC_LOAD(y + i + 1*ggml_f16_epr, 1); + ay2 = GGML_F16x_VEC_MUL(ay2, vx); + GGML_F16x_VEC_STORE(y + i + 1*ggml_f16_epr, ay2, 1); + } + // leftovers + // maximum number of leftover elements will be less that ggmlF_16x_epr. Apply predicated svmad on available elements only + if (np < n) { + svbool_t pg = svwhilelt_b16(np, n); + svfloat16_t hy = svld1_f16(pg, (__fp16 *)(y + np)); + svfloat16_t out = svmul_f16_m(pg, hy, vx); + svst1_f16(pg, (__fp16 *)(y + np), out); + } + np = n; +#elif defined(__riscv_v_intrinsic) + #if defined(__riscv_zvfh) + const ggml_fp16_t s = GGML_CPU_FP32_TO_FP16(v); + const _Float16 scale = *(const _Float16*)(&s); + + // calculate step size + const int epr = __riscv_vsetvlmax_e16m4(); + const int step = epr * 2; + int np = (n & ~(step - 1)); + + // unroll by 2 + for (int i = 0; i < np; i += step) { + vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, epr); + ay0 = __riscv_vfmul_vf_f16m4(ay0, scale, epr); + __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, epr); + __asm__ __volatile__ ("" ::: "memory"); + + vfloat16m4_t ay1 = __riscv_vle16_v_f16m4((const _Float16*)y + i + epr, epr); + ay1 = __riscv_vfmul_vf_f16m4(ay1, scale, epr); + __riscv_vse16_v_f16m4((_Float16*)y + i + epr, ay1, epr); + __asm__ __volatile__ ("" ::: "memory"); + } + + // leftovers + int vl; + for (int i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m4(n - i); + vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, vl); + ay0 = __riscv_vfmul_vf_f16m4(ay0, scale, vl); + __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, vl); + } + np = n; + #else + // fall to scalar path + const int np = 0; + #endif +#elif defined(GGML_SIMD) + const int np = (n & ~(GGML_F16_STEP - 1)); - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; + GGML_F16_VEC vx = GGML_F16_VEC_SET1(v); - GGML_ASSERT(ggml_are_same_shape(src0, dst)); - GGML_ASSERT(ggml_is_scalar(src1)); + GGML_F16_VEC ay[GGML_F16_ARR]; - // scalar to add - const float v = *(float *) src1->data; + for (int i = 0; i < np; i += GGML_F16_STEP) { + for (int j = 0; j < GGML_F16_ARR; j++) { + ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); + ay[j] = GGML_F16_VEC_MUL(ay[j], vx); - const int ith = params->ith; - const int nth = params->nth; + GGML_F16_VEC_STORE(y + i + j*GGML_F16_EPR, ay, j); + } + } +#else + // scalar path + const int np = 0; +#endif + // scalar and leftovers + for (int i = np; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(y[i])*v); + } +} - const int nr = ggml_nrows(src0); +inline static void ggml_vec_norm_f32 (const int n, float * s, const float * x) { ggml_vec_dot_f32(n, s, 0, x, 0, x, 0, 1); *s = sqrtf(*s); } +inline static void ggml_vec_sqr_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i]*x[i]; } +inline static void ggml_vec_sqr_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16(v*v); + } +} +inline static void ggml_vec_sqrt_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = sqrtf(x[i]); } +inline static void ggml_vec_sqrt_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(sqrtf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_log_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = logf(x[i]); } +inline static void ggml_vec_log_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(logf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_sin_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = sinf(x[i]); } +inline static void ggml_vec_sin_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(sinf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_cos_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = cosf(x[i]); } +inline static void ggml_vec_cos_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(cosf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_abs_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = fabsf(x[i]); } +inline static void ggml_vec_abs_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(fabsf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_sgn_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? 1.f : ((x[i] < 0.f) ? -1.f : 0.f); } +inline static void ggml_vec_sgn_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? 1.f : ((v < 0.f) ? -1.f : 0.f)); + } +} +inline static void ggml_vec_step_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? 1.f : 0.f; } +inline static void ggml_vec_step_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16((GGML_CPU_FP16_TO_FP32(x[i]) > 0.f) ? 1.f : 0.f); + } +} +inline static void ggml_vec_tanh_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = tanhf(x[i]); } +inline static void ggml_vec_tanh_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(tanhf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} +inline static void ggml_vec_elu_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? x[i] : expm1f(x[i]); } +inline static void ggml_vec_elu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + const float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v : expm1f(v)); + } +} +inline static void ggml_vec_relu_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? x[i] : 0.f; } +inline static void ggml_vec_relu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v : 0.f); + } +} +inline static void ggml_vec_leaky_relu_f32 (const int n, float * y, const float * x, const float ns) { for (int i = 0; i < n; ++i) y[i] = ((x[i] > 0.f) ? x[i] : 0.f) + ns * ((x[i] < 0.0f) ? x[i] : 0.f); } +inline static void ggml_vec_leaky_relu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const float ns) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16(((v > 0.f) ? v : 0.f) + ns * ((v < 0.0f) ? v : 0.f)); + } +} +inline static void ggml_vec_sigmoid_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = 1.f / (1.f + expf(-x[i])); } +inline static void ggml_vec_sigmoid_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(1.f / (1.f + expf(-GGML_CPU_FP16_TO_FP32(x[i])))); + } +} +// TODO: optimize performance +inline static void ggml_vec_hardswish_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i] * fminf(1.0f, fmaxf(0.0f, (x[i] + 3.0f) / 6.0f)); } +inline static void ggml_vec_hardswish_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16(v * fminf(1.0f, fmaxf(0.0f, (v + 3.0f) / 6.0f))); + } +} +inline static void ggml_vec_hardsigmoid_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = fminf(1.0f, fmaxf(0.0f, (x[i] + 3.0f) / 6.0f)); } +inline static void ggml_vec_hardsigmoid_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(fminf(1.0f, fmaxf(0.0f, (GGML_CPU_FP16_TO_FP32(x[i]) + 3.0f) / 6.0f))); + } +} +inline static void ggml_vec_exp_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = expf(x[i]); } +inline static void ggml_vec_exp_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(expf(GGML_CPU_FP16_TO_FP32(x[i]))); + } +} - GGML_TENSOR_UNARY_OP_LOCALS +static const float GELU_COEF_A = 0.044715f; +static const float GELU_QUICK_COEF = -1.702f; +static const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; +static const float SQRT_2_INV = 0.70710678118654752440084436210484f; - const ggml_type type = src0->type; - ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; - ggml_from_float_t const quantize_row_q = ggml_get_type_traits_cpu(type)->from_float; +inline static float ggml_gelu_f32(float x) { + return 0.5f*x*(1.0f + tanhf(SQRT_2_OVER_PI*x*(1.0f + GELU_COEF_A*x*x))); +} - // we don't support permuted src0 - GGML_ASSERT(nb00 == ggml_type_size(type)); +inline static void ggml_vec_gelu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + const uint16_t * i16 = (const uint16_t *) x; + for (int i = 0; i < n; ++i) { + y[i] = ggml_table_gelu_f16[i16[i]]; + } +} - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); +inline static void ggml_vec_gelu_erf_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + float xi = GGML_CPU_FP16_TO_FP32(x[i]); + float res = 0.5f*xi*(1.0f + erff(xi*SQRT_2_INV)); + y[i] = GGML_CPU_FP32_TO_FP16(res); + } +} - GGML_ASSERT(ggml_is_quantized(src0->type)); - GGML_ASSERT(dst->type == src0->type); - GGML_ASSERT(src1->type == GGML_TYPE_F32); +#ifdef GGML_GELU_FP16 +inline static void ggml_vec_gelu_f32(const int n, float * y, const float * x) { + uint16_t t; + for (int i = 0; i < n; ++i) { + if (x[i] <= -10.0f) { + y[i] = 0.0f; + } else if (x[i] >= 10.0f) { + y[i] = x[i]; + } else { + ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); + memcpy(&t, &fp16, sizeof(uint16_t)); + y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[t]); + } + } +} +#else +inline static void ggml_vec_gelu_f32(const int n, float * y, const float * x) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_gelu_f32(x[i]); + } +} +#endif - // rows per thread - const int dr = (nr + nth - 1)/nth; +inline static void ggml_vec_gelu_erf_f32(const int n, float * y, const float * x) { + for (int i = 0; i < n; ++i) { + float xi = x[i]; + y[i] = 0.5f*xi*(1.0f + erff(xi*SQRT_2_INV)); + } +} - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); +inline static float ggml_gelu_quick_f32(float x) { + return x*(1.0f/(1.0f+expf(GELU_QUICK_COEF*x))); +} - float * wdata = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32) * ith; +inline static void ggml_vec_gelu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + const uint16_t * i16 = (const uint16_t *) x; + for (int i = 0; i < n; ++i) { + y[i] = ggml_table_gelu_quick_f16[i16[i]]; + } +} - for (int ir = ir0; ir < ir1; ++ir) { - // src0 and dst are same shape => same indices - const int i3 = ir/(ne2*ne1); - const int i2 = (ir - i3*ne2*ne1)/ne1; - const int i1 = (ir - i3*ne2*ne1 - i2*ne1); +#ifdef GGML_GELU_QUICK_FP16 +inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) { + uint16_t t; + for (int i = 0; i < n; ++i) { + ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); + memcpy(&t, &fp16, sizeof(uint16_t)); + y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[t]); + } +} +#else +inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_gelu_quick_f32(x[i]); + } +} +#endif - void * src0_row = (void *) ((char *) src0->data + (i1*nb01 + i2*nb02 + i3*nb03)); - void * dst_row = (void *) ((char *) dst->data + (i1*nb1 + i2*nb2 + i3*nb0 )); +// Sigmoid Linear Unit (SiLU) function +inline static float ggml_silu_f32(float x) { + return x/(1.0f + expf(-x)); +} +inline static ggml_fp16_t ggml_silu_f16(ggml_fp16_t x) { + float v = GGML_CPU_FP16_TO_FP32(x); + return GGML_CPU_FP32_TO_FP16(v/(1.0f + expf(-v))); +} - assert(ne0 % 32 == 0); +#if __FINITE_MATH_ONLY__ +#error "some routines in ggml.c.inc require non-finite math arithmetics -- pass -fno-finite-math-only to the compiler to fix" +#error "ref: https://github.com/ggml-org/llama.cpp/pull/7154#issuecomment-2143844461" +#endif - // unquantize row from src0 to temp buffer - dequantize_row_q(src0_row, wdata, ne0); - // add src1 - ggml_vec_acc1_f32(ne0, wdata, v); - // quantize row to dst - quantize_row_q(wdata, dst_row, ne0); +/* Below function was borrowed from the GitHub repository: +https://github.com/openvinotoolkit/openvino/blob/master/src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/common.h.incpp */ +#if defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + inline static svfloat32_t exp_ps_sve(svbool_t pg, svfloat32_t src) { + // Constants + const svfloat32_t log2_e = svdup_n_f32(1.4426950409f); + const svfloat32_t ln2 = svdup_n_f32(0.6931473921f); + const svfloat32_t half_ln2_sq = svdup_n_f32(0.2413862043f); + const svuint32_t not_mask17 = svdup_n_u32(~((1u << 17) - 1)); + const svfloat32_t one = svdup_n_f32(1.0f); + const svfloat32_t inactive1 = svdup_n_f32(0.0f); + const svint32_t inactive2 = svdup_n_s32(0); + + // Algorithm starts here + svfloat32_t t0 = svmul_f32_m(pg, src, log2_e); // y = x * log2(e) + svfloat32_t t1 = svrintm_f32_m(inactive1, pg, t0); // rount to int (float) + svint32_t t2 = svcvt_s32_f32_m(inactive2, pg, t1); // n + + t1 = svsub_f32_m(pg, t0, t1); // a = y - floor(y) + t1 = svadd_f32_m(pg, t1, one); // b = a + 1 + + svuint32_t t3 = svlsr_n_u32_m(pg, svreinterpret_u32_f32(t1), 17); // v = b >> 17 (u32) + svfloat32_t t4 = svexpa_f32(t3); // c = fexpa(v) + t4 = svscale_f32_m(pg, t4, t2); // fexpa(v) * 2^(n) + + // and_(t2.d, t1.d, not_mask17.d) + svfloat32_t t5 = svreinterpret_f32_u32(svand_u32_m(pg, svreinterpret_u32_f32(t1), not_mask17)); + t5 = svsub_f32_m(pg, t1, t5); // z + t0 = svmla_f32_m(pg, ln2, t5, half_ln2_sq); // ln2 + half_ln2_sq * z + t0 = svmla_f32_m(pg, one, t5, t0); // 1 + (ln2 * z) + (half_ln2_sq * z * z) + t0 = svmul_f32_m(pg, t0, t4); // Final result + + return t0; } -} - -static void ggml_compute_forward_add1_bf16_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { +#endif - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; +#if defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + +inline static svfloat32_t ggml_v_expf(svbool_t pg, svfloat32_t x) { + const svfloat32_t r = svdup_n_f32_x(pg, 0x1.8p23f); + const svfloat32_t z = svmla_n_f32_x(pg, r, x, 0x1.715476p+0f); + const svfloat32_t n = svsub_f32_x(pg, z, r); + const svfloat32_t b = svmls_n_f32_x(pg, svmls_n_f32_x(pg, x, n, 0x1.62e4p-1f), n, 0x1.7f7d1cp-20f); + const svuint32_t e = svlsl_n_u32_x(pg, svreinterpret_u32_f32(z), 23); + const svfloat32_t k = svreinterpret_f32_u32(svadd_u32_x(pg, e, svreinterpret_u32_f32(svdup_n_f32_x(pg, 1)))); + const svbool_t c = svacgt_n_f32(pg, n, 126); + const svfloat32_t u = svmul_f32_x(pg, b, b); + const svfloat32_t j = svmla_f32_x(pg, + svmul_n_f32_x(pg, b, 0x1.ffffecp-1f), + svmla_f32_x(pg, svmla_f32_x(pg, svdup_n_f32_x(pg, 0x1.fffdb6p-2f), svdup_n_f32_x(pg, 0x1.555e66p-3f), b), + svmla_f32_x(pg, svdup_n_f32_x(pg, 0x1.573e2ep-5f), svdup_n_f32_x(pg, 0x1.0e4020p-7f), b), u), u); + const svuint32_t d = svdup_n_u32_z(svcmple_n_f32(pg, n, 0.0), 0x82000000); + const svfloat32_t s1 = svreinterpret_f32_u32(svadd_n_u32_x(pg, d, 0x7f000000)); + const svfloat32_t s2 = svreinterpret_f32_u32(svsub_u32_x(pg, e, d)); + return svsel_f32(svacgt_f32(pg, n, svdup_n_f32_x(pg, 192)), svmul_f32_x(pg, s1, s1), + svsel_f32(c, svmul_f32_x(pg, svmla_f32_x(pg, s2, s2, j), s1), svmla_f32_x(pg, k, k, j))); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static svfloat32_t ggml_v_silu(svbool_t pg, svfloat32_t x) { + const svfloat32_t one = svdup_n_f32_x(pg, 1.0f); + const svfloat32_t zero = svdup_n_f32_x(pg, 0.0f); + const svfloat32_t neg_x = svsub_f32_x(pg, zero, x); + const svfloat32_t exp_neg_x = ggml_v_expf(pg, neg_x); + const svfloat32_t one_plus_exp_neg_x = svadd_f32_x(pg, one, exp_neg_x); + return svdiv_f32_x(pg, x, one_plus_exp_neg_x); +} + +#elif defined(__ARM_NEON) && defined(__aarch64__) + +// 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 +inline static float32x4_t ggml_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))); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static float32x4_t ggml_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 = ggml_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); +} + +#elif defined(__AVX512F__) && defined(__AVX512DQ__) + +// 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 +inline static __m512 ggml_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); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static __m512 ggml_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 = ggml_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); +} + +#elif 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 +inline static __m256 ggml_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))))); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static __m256 ggml_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 = ggml_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); +} + +#elif defined(__SSE2__) // __AVX2__ / __ARM_NEON + +#if defined(__FMA__) +#define MADD128(x, y, z) _mm_fmadd_ps(x, y, z) +#define NMADD128(x, y, z) _mm_fnmadd_ps(x, y, z) +#else +#define MADD128(x, y, z) _mm_add_ps(_mm_mul_ps(x, y), z) +#define NMADD128(x, y, z) _mm_sub_ps(z, _mm_mul_ps(x, y)) +#endif - GGML_ASSERT(ggml_are_same_shape(src0, dst)); - GGML_ASSERT(ggml_is_scalar(src1)); +// 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 +inline static __m128 ggml_v_expf(__m128 x) { + const __m128 r = _mm_set1_ps(0x1.8p23f); + const __m128 z = MADD128(x, _mm_set1_ps(0x1.715476p+0f), r); + const __m128 n = _mm_sub_ps(z, r); + const __m128 b = + NMADD128(n, _mm_set1_ps(0x1.7f7d1cp-20f), NMADD128(n, _mm_set1_ps(0x1.62e4p-1f), x)); + const __m128i e = _mm_slli_epi32(_mm_castps_si128(z), 23); + const __m128 k = _mm_castsi128_ps(_mm_add_epi32(e, _mm_castps_si128(_mm_set1_ps(1)))); + const __m128i c = + _mm_castps_si128(_mm_cmpgt_ps(_mm_andnot_ps(_mm_set1_ps(-0.f), n), _mm_set1_ps(126))); + const __m128 u = _mm_mul_ps(b, b); + const __m128 j = + MADD128(MADD128(MADD128(_mm_set1_ps(0x1.0e4020p-7f), b, _mm_set1_ps(0x1.573e2ep-5f)), u, + MADD128(_mm_set1_ps(0x1.555e66p-3f), b, _mm_set1_ps(0x1.fffdb6p-2f))), + u, _mm_mul_ps(_mm_set1_ps(0x1.ffffecp-1f), b)); + if (!_mm_movemask_epi8(c)) + return MADD128(j, k, k); + const __m128i g = _mm_and_si128(_mm_castps_si128(_mm_cmple_ps(n, _mm_setzero_ps())), + _mm_set1_epi32(0x82000000u)); + const __m128 s1 = _mm_castsi128_ps(_mm_add_epi32(g, _mm_set1_epi32(0x7f000000u))); + const __m128 s2 = _mm_castsi128_ps(_mm_sub_epi32(e, g)); + const __m128i d = + _mm_castps_si128(_mm_cmpgt_ps(_mm_andnot_ps(_mm_set1_ps(-0.f), n), _mm_set1_ps(192))); + return _mm_or_ps( + _mm_and_ps(_mm_castsi128_ps(d), _mm_mul_ps(s1, s1)), + _mm_andnot_ps(_mm_castsi128_ps(d), + _mm_or_ps(_mm_and_ps(_mm_castsi128_ps(c), _mm_mul_ps(MADD128(s2, j, s2), s1)), + _mm_andnot_ps(_mm_castsi128_ps(c), MADD128(k, j, k))))); +} + +// computes silu x/(1+exp(-x)) in single precision vector +inline static __m128 ggml_v_silu(__m128 x) { + const __m128 one = _mm_set1_ps(1); + const __m128 zero = _mm_setzero_ps(); + const __m128 neg_x = _mm_sub_ps(zero, x); + const __m128 exp_neg_x = ggml_v_expf(neg_x); + const __m128 one_plus_exp_neg_x = _mm_add_ps(one, exp_neg_x); + return _mm_div_ps(x, one_plus_exp_neg_x); +} + +#elif defined(__riscv_v_intrinsic) + +// 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 +inline static vfloat32m2_t ggml_v_expf_m2(vfloat32m2_t x, int vl) { + const vfloat32m2_t r = __riscv_vfmv_v_f_f32m2(0x1.8p23f, vl); +#ifdef __riscv_xtheadvector + // workaround for compiler bug (gcc 14.3.0: Error: unrecognized opcode `th.vmv1r.v v2,v4') + vfloat32m2_t z = __riscv_vfadd_vf_f32m2(r, 0.0f, vl); + z = __riscv_vfmacc_vf_f32m2(z, 0x1.715476p+0f, x, vl); +#else + const vfloat32m2_t z = __riscv_vfmacc_vf_f32m2(r, 0x1.715476p+0f, x, vl); +#endif + const vfloat32m2_t n = __riscv_vfsub_vv_f32m2(z, r, vl); + const vfloat32m2_t b = __riscv_vfnmsac_vf_f32m2(__riscv_vfnmsac_vf_f32m2(x, 0x1.62e4p-1f, n, vl), + 0x1.7f7d1cp-20f, n, vl); + const vuint32m2_t e = __riscv_vsll_vx_u32m2(__riscv_vreinterpret_v_f32m2_u32m2(z), 23, vl); + const vfloat32m2_t k = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vadd_vx_u32m2(e, 0x3f800000, vl)); // 1.0f + const vbool16_t c = __riscv_vmfgt_vf_f32m2_b16(__riscv_vfabs_v_f32m2(n, vl), 126.0f, vl); + const vfloat32m2_t u = __riscv_vfmul_vv_f32m2(b, b, vl); + const vfloat32m2_t j = __riscv_vfmacc_vv_f32m2( + __riscv_vfmul_vf_f32m2(b, 0x1.ffffecp-1f, vl), + __riscv_vfmacc_vv_f32m2( + __riscv_vfmacc_vf_f32m2(__riscv_vfmv_v_f_f32m2(0x1.fffdb6p-2f, vl), 0x1.555e66p-3f, b, vl), + __riscv_vfmacc_vf_f32m2(__riscv_vfmv_v_f_f32m2(0x1.573e2ep-5f, vl), 0x1.0e4020p-7f, b, vl), + u, vl), u, vl); + if (!__riscv_vcpop_m_b16(c, vl)) + return __riscv_vfmacc_vv_f32m2(k, j, k, vl); + const vbool16_t dm = __riscv_vmfle_vf_f32m2_b16(n, 0.0f, vl); + const vuint32m2_t d = __riscv_vmerge_vxm_u32m2(__riscv_vmv_v_x_u32m2(0, vl), 0x82000000, dm, vl); + const vfloat32m2_t s1 = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vadd_vx_u32m2(d, 0x7f000000, vl)); + const vfloat32m2_t s2 = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vsub_vv_u32m2(e, d, vl)); + const vfloat32m2_t r1 = __riscv_vmerge_vvm_f32m2( + __riscv_vfmacc_vv_f32m2(k, k, j, vl), + __riscv_vfmul_vv_f32m2(__riscv_vfmacc_vv_f32m2(s2, s2, j, vl), s1, vl), + c, vl); + return __riscv_vmerge_vvm_f32m2( + r1, __riscv_vfmul_vv_f32m2(s1, s1, vl), + __riscv_vmfgt_vf_f32m2_b16(__riscv_vfabs_v_f32m2(n, vl), 192.0f, vl), + vl); +} - // scalar to add - const float v = *(float *) src1->data; +// computes silu x/(1+exp(-x)) in single precision vector +inline static vfloat32m2_t ggml_v_silu_m2(vfloat32m2_t x, int vl) { + const vfloat32m2_t neg_x = __riscv_vfneg_v_f32m2(x, vl); + const vfloat32m2_t exp_neg_x = ggml_v_expf_m2(neg_x, vl); + const vfloat32m2_t one_plus_exp_neg_x = __riscv_vfadd_vf_f32m2(exp_neg_x, 1.0f, vl); + return __riscv_vfdiv_vv_f32m2(x, one_plus_exp_neg_x, vl); +} - const int ith = params->ith; - const int nth = params->nth; +#endif // __ARM_NEON / __AVX2__ / __SSE2__ / __riscv_v_intrinsic - const int nr = ggml_nrows(src0); +inline static void ggml_vec_silu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_silu_f16(x[i]); + } +} - GGML_TENSOR_UNARY_OP_LOCALS +inline static float ggml_silu_backward_f32(float x, float dy) { + const float s = 1.0f/(1.0f + expf(-x)); + return dy*s*(1.0f + x*(1.0f - s)); +} - GGML_ASSERT(src0->type == GGML_TYPE_BF16); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_BF16); +inline static ggml_fp16_t ggml_silu_backward_f16(ggml_fp16_t x, ggml_fp16_t dy) { + const float v = GGML_CPU_FP16_TO_FP32(x); + const float s = 1.0f/(1.0f + expf(-v)); + return GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(dy)*s*(1.0f + v*(1.0f - s))); +} - GGML_ASSERT( nb0 == sizeof(ggml_bf16_t)); - GGML_ASSERT(nb00 == sizeof(ggml_bf16_t)); +inline static void ggml_vec_silu_backward_f32(const int n, float * dx, const float * x, const float * dy) { + for (int i = 0; i < n; ++i) { + dx[i] = ggml_silu_backward_f32(x[i], dy[i]); + } +} - // rows per thread - const int dr = (nr + nth - 1)/nth; +inline static void ggml_vec_silu_backward_f16(const int n, ggml_fp16_t * dx, const ggml_fp16_t * x, const ggml_fp16_t * dy) { + for (int i = 0; i < n; ++i) { + dx[i] = ggml_silu_backward_f16(x[i], dy[i]); + } +} - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); +inline static void ggml_vec_reglu_f32 (const int n, float * y, const float * x, const float * g) { + for (int i = 0; i < n; ++i) { + y[i] = (x[i] > 0.f) ? x[i] * g[i] : 0.f; + } +} - for (int ir = ir0; ir < ir1; ++ir) { - // src0 and dst are same shape => same indices - const int i3 = ir/(ne2*ne1); - const int i2 = (ir - i3*ne2*ne1)/ne1; - const int i1 = (ir - i3*ne2*ne1 - i2*ne1); +inline static void ggml_vec_reglu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(x[i]); + y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v * GGML_CPU_FP16_TO_FP32(g[i]) : 0.f); + } +} - ggml_bf16_t * dst_ptr = (ggml_bf16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ); - ggml_bf16_t * src0_ptr = (ggml_bf16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01); - for (int i = 0; i < ne0; i++) { - dst_ptr[i] = GGML_FP32_TO_BF16(GGML_BF16_TO_FP32(src0_ptr[i]) + v); +#ifdef GGML_GELU_FP16 +inline static void ggml_vec_geglu_f32(const int n, float * y, const float * x, const float * g) { + uint16_t t; + for (int i = 0; i < n; ++i) { + if (x[i] <= -10.0f) { + y[i] = 0.0f; + } else if (x[i] >= 10.0f) { + y[i] = x[i] * g[i]; + } else { + ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); + memcpy(&t, &fp16, sizeof(uint16_t)); + y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[t]) * g[i]; } } } +#else +inline static void ggml_vec_geglu_f32(const int n, float * y, const float * x, const float * g) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_gelu_f32(x[i]) * g[i]; + } +} +#endif -static void ggml_compute_forward_add1_bf16_bf16( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; +inline static void ggml_vec_geglu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + const uint16_t * i16 = (const uint16_t *) x; + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(g[i]); + y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[i16[i]]) * v); + } +} - GGML_ASSERT(ggml_are_same_shape(src0, dst)); - GGML_ASSERT(ggml_is_scalar(src1)); +void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float * g); - // scalar to add - const float v = GGML_BF16_TO_FP32(*(ggml_bf16_t *) src1->data); +inline static void ggml_vec_swiglu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + for (int i = 0; i < n; ++i) { + float xi = GGML_CPU_FP16_TO_FP32(x[i]); + float gi = GGML_CPU_FP16_TO_FP32(g[i]); + y[i] = GGML_CPU_FP32_TO_FP16((xi/(1.0f + expf(-xi))) * gi); + } +} - const int ith = params->ith; - const int nth = params->nth; +inline static void ggml_vec_geglu_erf_f32(const int n, float * y, const float * x, const float * g) { + for (int i = 0; i < n; ++i) { + float xi = x[i]; + y[i] = 0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * g[i]; + } +} - const int nr = ggml_nrows(src0); +inline static void ggml_vec_geglu_erf_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + for (int i = 0; i < n; ++i) { + float xi = GGML_CPU_FP16_TO_FP32(x[i]); + float gi = GGML_CPU_FP16_TO_FP32(g[i]); + y[i] = GGML_CPU_FP32_TO_FP16(0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * gi); + } +} - GGML_TENSOR_UNARY_OP_LOCALS +#ifdef GGML_GELU_QUICK_FP16 +inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) { + uint16_t t; + for (int i = 0; i < n; ++i) { + ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); + memcpy(&t, &fp16, sizeof(uint16_t)); + y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[t]) * g[i]; + } +} +#else +inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) { + for (int i = 0; i < n; ++i) { + y[i] = ggml_gelu_quick_f32(x[i]) * g[i]; + } +} +#endif - GGML_ASSERT(src0->type == GGML_TYPE_BF16); - GGML_ASSERT(src1->type == GGML_TYPE_BF16); - GGML_ASSERT(dst->type == GGML_TYPE_BF16); +inline static void ggml_vec_geglu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { + const uint16_t * i16 = (const uint16_t *) x; + for (int i = 0; i < n; ++i) { + float v = GGML_CPU_FP16_TO_FP32(g[i]); + y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[i16[i]]) * v); + } +} - GGML_ASSERT( nb0 == sizeof(ggml_bf16_t)); - GGML_ASSERT(nb00 == sizeof(ggml_bf16_t)); +inline static void ggml_vec_sum_f32(const int n, float * s, const float * x) { +#ifndef GGML_USE_ACCELERATE + ggml_float sum = 0.0; + for (int i = 0; i < n; ++i) { + sum += (ggml_float)x[i]; + } + *s = (float)sum; +#else + vDSP_sve(x, 1, s, n); +#endif +} - // rows per thread - const int dr = (nr + nth - 1)/nth; +inline static void ggml_vec_cumsum_f32(const int n, float * y, const float * x) { + for (int i = 0; i < n; ++i) { + if (i == 0) { + y[i] = x[i]; + } else { + y[i] = y[i - 1] + x[i]; + } + } +} - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); +inline static void ggml_vec_sum_f32_ggf(const int n, ggml_float * s, const float * x) { + ggml_float sum = 0.0; + for (int i = 0; i < n; ++i) { + sum += (ggml_float)x[i]; + } + *s = sum; +} - for (int ir = ir0; ir < ir1; ++ir) { - // src0 and dst are same shape => same indices - const int i3 = ir/(ne2*ne1); - const int i2 = (ir - i3*ne2*ne1)/ne1; - const int i1 = (ir - i3*ne2*ne1 - i2*ne1); +inline static void ggml_vec_sum_f16_ggf(const int n, float * s, const ggml_fp16_t * x) { + float sum = 0.0f; + for (int i = 0; i < n; ++i) { + sum += GGML_CPU_FP16_TO_FP32(x[i]); + } + *s = sum; +} - ggml_bf16_t * dst_ptr = (ggml_bf16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ); - ggml_bf16_t * src0_ptr = (ggml_bf16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01); - for (int i = 0; i < ne0; i++) { - dst_ptr[i] = GGML_FP32_TO_BF16(GGML_BF16_TO_FP32(src0_ptr[i]) + v); - } +inline static void ggml_vec_sum_bf16_ggf(const int n, float * s, const ggml_bf16_t * x) { + float sum = 0.0f; + for (int i = 0; i < n; ++i) { + sum += GGML_BF16_TO_FP32(x[i]); } + *s = sum; } -void ggml_compute_forward_add1( - const ggml_compute_params * params, - ggml_tensor * dst) { +inline static void ggml_vec_max_f32(const int n, float * s, const float * x) { +#ifndef GGML_USE_ACCELERATE + float max = -INFINITY; + for (int i = 0; i < n; ++i) { + max = MAX(max, x[i]); + } + *s = max; +#else + vDSP_maxv(x, 1, s, n); +#endif +} - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; +inline static void ggml_vec_norm_inv_f32(const int n, float * s, const float * x) { + ggml_vec_norm_f32(n, s, x); + *s = 1.f/(*s); +} - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_add1_f32(params, dst); - } break; - case GGML_TYPE_F16: - { - if (src1->type == GGML_TYPE_F16) { - ggml_compute_forward_add1_f16_f16(params, dst); - } - else if (src1->type == GGML_TYPE_F32) { - ggml_compute_forward_add1_f16_f32(params, dst); - } - else { - GGML_ABORT("fatal error"); - } - } break; - case GGML_TYPE_BF16: - { - if (src1->type == GGML_TYPE_BF16) { - ggml_compute_forward_add1_bf16_bf16(params, dst); - } - else if (src1->type == GGML_TYPE_F32) { - ggml_compute_forward_add1_bf16_f32(params, dst); - } - else { - GGML_ABORT("fatal error"); - } - } break; - case GGML_TYPE_Q1_0: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q8_0: - case GGML_TYPE_Q8_1: - case GGML_TYPE_MXFP4: - case GGML_TYPE_NVFP4: - case GGML_TYPE_Q2_K: - case GGML_TYPE_Q3_K: - case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: - case GGML_TYPE_Q6_K: - case GGML_TYPE_TQ1_0: - case GGML_TYPE_TQ2_0: - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ2_S: - { - ggml_compute_forward_add1_q_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } +inline static void ggml_vec_argmax_f32(const int n, int * s, const float * x) { + float max = -INFINITY; + int idx = 0; + for (int i = 0; i < n; ++i) { + max = MAX(max, x[i]); + if (max == x[i]) { idx = i; } } + *s = idx; } -// ggml_compute_forward_acc +#ifdef __cplusplus +} +#endif -static void ggml_compute_forward_acc_f32( + +#include +#include +#include + +extern "C" void ggml_compute_forward_add_non_quantized(const struct ggml_compute_params * params, struct ggml_tensor * dst); +extern "C" { +void ggml_compute_forward_abs(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sgn(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_neg(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_step(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_tanh(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_elu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_relu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sigmoid(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_hardsigmoid(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_exp(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_hardswish(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_expm1(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_softplus(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_floor(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_ceil(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_round(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_trunc(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_xielu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +} + +// ggml_compute_forward_dup + +static void ggml_compute_forward_dup_same_cont( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src0)); GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src0)); + GGML_ASSERT(src0->type == dst->type); - // view src0 and dst with these strides and data offset inbytes during acc - // nb0 is implicitly element_size because src0 and dst are contiguous - size_t nb1 = ((int32_t *) dst->op_params)[0]; - size_t nb2 = ((int32_t *) dst->op_params)[1]; - size_t nb3 = ((int32_t *) dst->op_params)[2]; - size_t offset = ((int32_t *) dst->op_params)[3]; - bool inplace = (bool) ((int32_t *) dst->op_params)[4]; + const size_t nb0 = ggml_type_size(src0->type); - if (!inplace) { - if (params->ith == 0) { - // memcpy needs to be synchronized across threads to avoid race conditions. - // => do it in INIT phase - memcpy( - ((char *) dst->data), - ((char *) src0->data), - ggml_nbytes(dst)); - } - ggml_barrier(params->threadpool); - } + const int ith = params->ith; // thread index + const int nth = params->nth; // number of threads - const int ith = params->ith; - const int nth = params->nth; + // parallelize by blocks + const int nk = ggml_nelements(src0)/ggml_blck_size(src0->type); + const int dr = (nk + nth - 1) / nth; + const int k0 = dr * ith; + const int k1 = MIN(k0 + dr, nk); - const int nr = ggml_nrows(src1); - const int nc = src1->ne[0]; + if (k0 < k1) { + memcpy( + ((char *) dst->data + k0*nb0), + ((char *) src0->data + k0*nb0), + (k1 - k0) * nb0); + } +} - GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) - GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) - - // src0 and dst as viewed during acc - const size_t nb0 = ggml_element_size(src0); - - const size_t nb00 = nb0; - const size_t nb01 = nb1; - const size_t nb02 = nb2; - const size_t nb03 = nb3; - - GGML_ASSERT(offset + (ne10 == 0 ? 0 : ne10-1)*nb0 + (ne11 == 0 ? 0 : ne11-1)*nb1 + (ne12 == 0 ? 0 : ne12-1)*nb2 + (ne13 == 0 ? 0 : ne13-1)*nb3 < ggml_nbytes(dst)); - GGML_ASSERT(offset + (ne10 == 0 ? 0 : ne10-1)*nb00 + (ne11 == 0 ? 0 : ne11-1)*nb01 + (ne12 == 0 ? 0 : ne12-1)*nb02 + (ne13 == 0 ? 0 : ne13-1)*nb03 < ggml_nbytes(src0)); - - GGML_ASSERT(nb10 == sizeof(float)); - - // rows per thread - const int dr = (nr + nth - 1)/nth; - - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); - - for (int ir = ir0; ir < ir1; ++ir) { - // src0 and dst are viewed with shape of src1 and offset - // => same indices - const int i3 = ir/(ne12*ne11); - const int i2 = (ir - i3*ne12*ne11)/ne11; - const int i1 = (ir - i3*ne12*ne11 - i2*ne11); - -#ifdef GGML_USE_ACCELERATE - vDSP_vadd( - (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + offset), 1, - (float *) ((char *) src1->data + i3*nb13 + i2*nb12 + i1*nb11), 1, - (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + offset), 1, nc); -#else - ggml_vec_add_f32(nc, - (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + offset), - (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + offset), - (float *) ((char *) src1->data + i3*nb13 + i2*nb12 + i1*nb11)); -#endif - } -} - -void ggml_compute_forward_acc( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_acc_f32(params, dst); - } break; - case GGML_TYPE_F16: - case GGML_TYPE_BF16: - case GGML_TYPE_Q1_0: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q8_0: - case GGML_TYPE_Q8_1: - case GGML_TYPE_MXFP4: - case GGML_TYPE_NVFP4: - case GGML_TYPE_Q2_K: - case GGML_TYPE_Q3_K: - case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: - case GGML_TYPE_Q6_K: - case GGML_TYPE_TQ1_0: - case GGML_TYPE_TQ2_0: - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ2_S: - default: - { - GGML_ABORT("fatal error"); - } - } -} - -// ggml_compute_forward_sum - -static void ggml_compute_forward_sum_f32( +template +static void ggml_compute_forward_dup_flt( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - if (params->ith != 0) { - return; - } + GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src0)); + GGML_ASSERT(!ggml_is_quantized(src0->type) && !ggml_is_quantized(dst->type)); - assert(ggml_is_scalar(dst)); - assert(src0->nb[0] == sizeof(float)); + GGML_TENSOR_UNARY_OP_LOCALS - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) + const int ith = params->ith; // thread index + const int nth = params->nth; // number of threads - ggml_float sum = 0; - ggml_float row_sum = 0; + // parallelize by rows + const int nr = ne01; + // number of rows per thread + const int dr = (nr + nth - 1) / nth; + // row range for this thread + const int ir0 = dr * ith; + const int ir1 = MIN(ir0 + dr, nr); - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = 0; i01 < ne01; i01++) { - ggml_vec_sum_f32_ggf(ne00, - &row_sum, - (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03)); - sum += row_sum; + // case: type & row size equal + if (src0->type == dst->type && + ne00 == ne0 && + nb00 == ggml_type_size(src0->type) && nb0 == ggml_type_size(dst->type)) { + // copy by rows + const size_t rs = ne00*nb00; + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = ir0; i01 < ir1; i01++) { + memcpy( + ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3), + ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03), + rs); + } } } - } - ((float *) dst->data)[0] = sum; -} - -static void ggml_compute_forward_sum_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - - if (params->ith != 0) { return; } - assert(ggml_is_scalar(dst)); - - assert(src0->nb[0] == sizeof(ggml_fp16_t)); - - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) + // case: dst tensor is contiguous + if (ggml_is_contiguous(dst)) { + if (nb00 == sizeof(src_t)) { + if constexpr (std::is_same_v) { + // same type + size_t id = 0; + const size_t rs = ne00 * nb00; + char * dst_ptr = (char *) dst->data; - float sum = 0; - float row_sum = 0; + for (int i03 = 0; i03 < ne03; i03++) { + for (int i02 = 0; i02 < ne02; i02++) { + id += rs * ir0; + for (int i01 = ir0; i01 < ir1; i01++) { + const char * src0_ptr = (char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03; + memcpy(dst_ptr + id, src0_ptr, rs); + id += rs; + } + id += rs * (ne01 - ir1); + } + } + } else { + // casting between non-quantized types + size_t id = 0; + dst_t * dst_ptr = (dst_t *) dst->data; - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = 0; i01 < ne01; i01++) { - ggml_vec_sum_f16_ggf(ne00, - &row_sum, - (ggml_fp16_t *) ((char *) src0->data + i01 * nb01 + i02 * nb02 + i03 * nb03)); - sum += row_sum; + for (int i03 = 0; i03 < ne03; i03++) { + for (int i02 = 0; i02 < ne02; i02++) { + id += ne00 * ir0; + for (int i01 = ir0; i01 < ir1; i01++) { + const src_t * src0_ptr = (src_t *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); + for (int i00 = 0; i00 < ne00; i00++) { + float tmp = type_conversion_table::to_f32(src0_ptr[i00]); + dst_ptr[id] = type_conversion_table::from_f32(tmp); + id++; + } + } + id += ne00 * (ne01 - ir1); + } + } } - } - } - ((ggml_fp16_t *) dst->data)[0] = GGML_CPU_FP32_TO_FP16(sum); -} + } else { + //printf("%s: this is not optimal - fix me\n", __func__); -static void ggml_compute_forward_sum_bf16( - const ggml_compute_params * params, - ggml_tensor * dst) { + size_t id = 0; + dst_t * dst_ptr = (dst_t *) dst->data; - const ggml_tensor * src0 = dst->src[0]; + for (int i03 = 0; i03 < ne03; i03++) { + for (int i02 = 0; i02 < ne02; i02++) { + id += ne00 * ir0; + for (int i01 = ir0; i01 < ir1; i01++) { + for (int i00 = 0; i00 < ne00; i00++) { + const src_t * src0_ptr = (src_t *) ((char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); - if (params->ith != 0) { + float tmp = type_conversion_table::to_f32(*src0_ptr); + dst_ptr[id] = type_conversion_table::from_f32(tmp); + id++; + } + } + id += ne00 * (ne01 - ir1); + } + } + } return; } - assert(ggml_is_scalar(dst)); + // dst counters + int64_t i10 = 0; + int64_t i11 = 0; + int64_t i12 = 0; + int64_t i13 = 0; - assert(src0->nb[0] == sizeof(ggml_bf16_t)); + if constexpr (std::is_same_v) { + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + i10 += ne00 * ir0; + while (i10 >= ne0) { + i10 -= ne0; + if (++i11 == ne1) { + i11 = 0; + if (++i12 == ne2) { + i12 = 0; + if (++i13 == ne3) { + i13 = 0; + } + } + } + } + for (int64_t i01 = ir0; i01 < ir1; i01++) { + for (int64_t i00 = 0; i00 < ne00; i00++) { + const char * src0_ptr = ((char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); + char * dst_ptr = ((char *) dst->data + i10*nb0 + i11*nb1 + i12*nb2 + i13*nb3); - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) + memcpy(dst_ptr, src0_ptr, sizeof(dst_t)); - float sum = 0; - float row_sum = 0; - - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = 0; i01 < ne01; i01++) { - ggml_vec_sum_bf16_ggf(ne00, - &row_sum, - (ggml_bf16_t *) ((char *) src0->data + i01 * nb01 + i02 * nb02 + i03 * nb03)); - sum += row_sum; + if (++i10 == ne00) { + i10 = 0; + if (++i11 == ne01) { + i11 = 0; + if (++i12 == ne02) { + i12 = 0; + if (++i13 == ne03) { + i13 = 0; + } + } + } + } + } + } + i10 += ne00 * (ne01 - ir1); + while (i10 >= ne0) { + i10 -= ne0; + if (++i11 == ne1) { + i11 = 0; + if (++i12 == ne2) { + i12 = 0; + if (++i13 == ne3) { + i13 = 0; + } + } + } + } } } - } - ((ggml_bf16_t *) dst->data)[0] = GGML_FP32_TO_BF16(sum); -} -void ggml_compute_forward_sum( - const ggml_compute_params * params, - ggml_tensor * dst) { + } else { + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + i10 += ne00 * ir0; + while (i10 >= ne0) { + i10 -= ne0; + if (++i11 == ne1) { + i11 = 0; + if (++i12 == ne2) { + i12 = 0; + if (++i13 == ne3) { + i13 = 0; + } + } + } + } + for (int64_t i01 = ir0; i01 < ir1; i01++) { + for (int64_t i00 = 0; i00 < ne00; i00++) { + const char * src0_ptr = ((char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); + char * dst_ptr = ((char *) dst->data + i10*nb0 + i11*nb1 + i12*nb2 + i13*nb3); - const ggml_tensor * src0 = dst->src[0]; + float tmp = type_conversion_table::to_f32(*(const src_t *) src0_ptr); + *(dst_t *) dst_ptr = type_conversion_table::from_f32(tmp); - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_sum_f32(params, dst); - } break; - case GGML_TYPE_F16: - { - ggml_compute_forward_sum_f16(params, dst); - } break; - case GGML_TYPE_BF16: - { - ggml_compute_forward_sum_bf16(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); + if (++i10 == ne0) { + i10 = 0; + if (++i11 == ne1) { + i11 = 0; + if (++i12 == ne2) { + i12 = 0; + if (++i13 == ne3) { + i13 = 0; + } + } + } + } + } + } + i10 += ne00 * (ne01 - ir1); + while (i10 >= ne0) { + i10 -= ne0; + if (++i11 == ne1) { + i11 = 0; + if (++i12 == ne2) { + i12 = 0; + if (++i13 == ne3) { + i13 = 0; + } + } + } + } } + } } } -// ggml_compute_forward_cumsum -static void ggml_compute_forward_cumsum_f32( +template +static void ggml_compute_forward_dup_to_q( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - GGML_ASSERT(src0->nb[0] == sizeof(float)); - GGML_ASSERT(dst->nb[0] == sizeof(float)); + GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src0)); + GGML_ASSERT(!ggml_is_quantized(src0->type)); GGML_TENSOR_UNARY_OP_LOCALS - GGML_ASSERT(ne0 == ne00); - GGML_ASSERT(ne1 == ne01); - GGML_ASSERT(ne2 == ne02); - GGML_ASSERT(ne3 == ne03); - - const auto [ir0, ir1] = get_thread_range(params, src0); + const int ith = params->ith; // thread index + const int nth = params->nth; // number of threads - for (int64_t ir = ir0; ir < ir1; ++ir) { - const int64_t i03 = ir/(ne02*ne01); - const int64_t i02 = (ir - i03*ne02*ne01)/ne01; - const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); + // parallelize by rows + const int nr = ne01; + // number of rows per thread + const int dr = (nr + nth - 1) / nth; + // row range for this thread + const int ir0 = dr * ith; + const int ir1 = MIN(ir0 + dr, nr); - float * src_row = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); - float * dst_row = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); + if (ggml_is_contiguous(dst) && + nb00 == sizeof(src_t) && + ggml_get_type_traits_cpu(dst->type)->from_float) { + // casting non-quantized types --> intermediate f32 --> quantized + ggml_from_float_t const quantize_row_q = ggml_get_type_traits_cpu(dst->type)->from_float; + float * src0_f32 = (float *) params->wdata + (ne00 + CACHE_LINE_SIZE_F32) * ith; - ggml_vec_cumsum_f32(ne00, dst_row, src_row); - } -} + size_t id = 0; + size_t rs = nb0 * (ne00 / ggml_blck_size(dst->type)); + char * dst_ptr = (char *) dst->data; -void ggml_compute_forward_cumsum( - const ggml_compute_params * params, - ggml_tensor * dst) { + for (int i03 = 0; i03 < ne03; i03++) { + for (int i02 = 0; i02 < ne02; i02++) { + id += rs * ir0; + for (int i01 = ir0; i01 < ir1; i01++) { + const src_t * src0_ptr = (src_t *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); - const ggml_tensor * src0 = dst->src[0]; + for (int i00 = 0; i00 < ne00; i00++) { + src0_f32[i00] = type_conversion_table::to_f32(src0_ptr[i00]); + } - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_cumsum_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); + quantize_row_q(src0_f32, dst_ptr + id, ne00); + id += rs; + } + id += rs * (ne01 - ir1); } + } + } else { + // printf("%s %s\n", ggml_type_name(src0->type), ggml_type_name(dst->type)); + GGML_ABORT("not implemented"); } } -// ggml_compute_forward_sum_rows - -static void ggml_compute_forward_sum_rows_f32( +// A simplified version of ggml_compute_forward_dup that doesn't do float upcasting, and just plain old memcpy. +static void ggml_compute_forward_dup_bytes( const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - if (params->ith != 0) { + GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src0)); + GGML_ASSERT(src0->type == dst->type); + + GGML_TENSOR_UNARY_OP_LOCALS; + + if (ggml_is_contiguous(src0) && ggml_is_contiguous(dst)) { + ggml_compute_forward_dup_same_cont(params, dst); return; } - GGML_ASSERT(src0->nb[0] == sizeof(float)); - GGML_ASSERT(dst->nb[0] == sizeof(float)); + const size_t type_size = ggml_type_size(src0->type); - GGML_TENSOR_UNARY_OP_LOCALS + const int ith = params->ith; // thread index + const int nth = params->nth; // number of threads - GGML_ASSERT(ne0 == 1); - GGML_ASSERT(ne1 == ne01); - GGML_ASSERT(ne2 == ne02); - GGML_ASSERT(ne3 == ne03); + // parallelize by rows + const int nr = ne01; + // number of rows per thread + const int dr = (nr + nth - 1) / nth; + // row range for this thread + const int ir0 = dr * ith; + const int ir1 = MIN(ir0 + dr, nr); - for (int64_t i3 = 0; i3 < ne03; i3++) { - for (int64_t i2 = 0; i2 < ne02; i2++) { - for (int64_t i1 = 0; i1 < ne01; i1++) { - float * src_row = (float *) ((char *) src0->data + i1*nb01 + i2*nb02 + i3*nb03); - float * dst_row = (float *) ((char *) dst->data + i1*nb1 + i2*nb2 + i3*nb3); - float row_sum = 0; - ggml_vec_sum_f32(ne00, &row_sum, src_row); - dst_row[0] = row_sum; + if (src0->type == dst->type && + ggml_are_same_shape(src0, dst) && + nb00 == type_size && nb0 == type_size) { + // copy by rows + const size_t rs = ggml_row_size(src0->type, ne00); + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = ir0; i01 < ir1; i01++) { + memcpy( + ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3), + ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03), + rs); + } } } + return; } -} -void ggml_compute_forward_sum_rows( - const ggml_compute_params * params, - ggml_tensor * dst) { + if (ggml_is_contiguous(dst)) { + size_t id = 0; + char * dst_ptr = (char *) dst->data; + const size_t rs = ne00 * type_size; - const ggml_tensor * src0 = dst->src[0]; + if (nb00 == type_size) { + // src0 is contiguous on first dimension, copy by rows + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + id += rs * ir0; + for (int64_t i01 = ir0; i01 < ir1; i01++) { + const char * src0_ptr = (char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03; + memcpy(dst_ptr + id, src0_ptr, rs); + id += rs; + } + id += rs * (ne01 - ir1); + } + } + } else { + //printf("%s: this is not optimal - fix me\n", __func__); - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_sum_rows_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} - -// ggml_compute_forward_mean - -static void ggml_compute_forward_mean_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + id += rs * ir0; + for (int64_t i01 = ir0; i01 < ir1; i01++) { + for (int64_t i00 = 0; i00 < ne00; i00++) { + const char * src0_ptr = (char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03; + memcpy(dst_ptr + id, src0_ptr, type_size); - const ggml_tensor * src0 = dst->src[0]; + id += type_size; + } + } + id += rs * (ne01 - ir1); + } + } + } - if (params->ith != 0) { return; } - assert(src0->nb[0] == sizeof(float)); - - GGML_TENSOR_UNARY_OP_LOCALS - - assert(ne0 == 1); - assert(ne1 == ne01); - assert(ne2 == ne02); - assert(ne3 == ne03); + // dst counters + int64_t k10 = 0; + int64_t i11 = 0; + int64_t i12 = 0; + int64_t i13 = 0; - GGML_UNUSED(ne0); - GGML_UNUSED(ne1); - GGML_UNUSED(ne2); - GGML_UNUSED(ne3); + // number of blocks in a row + const int64_t nk00 = ne00 / ggml_blck_size(src0->type); + const int64_t nk0 = ne0 / ggml_blck_size(dst->type); for (int64_t i03 = 0; i03 < ne03; i03++) { for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = 0; i01 < ne01; i01++) { - ggml_vec_sum_f32(ne00, - (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3), - (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03)); + k10 += nk00 * ir0; + while (k10 >= nk0) { + k10 -= nk0; + if (++i11 == ne1) { + i11 = 0; + if (++i12 == ne2) { + i12 = 0; + if (++i13 == ne3) { + i13 = 0; + } + } + } + } + for (int64_t i01 = ir0; i01 < ir1; i01++) { + for (int64_t k00 = 0; k00 < nk00; k00++) { + const char * src0_ptr = ((char *) src0->data + k00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); + char * dst_ptr = ((char *) dst->data + k10*nb0 + i11*nb1 + i12*nb2 + i13*nb3); - *(float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3) /= (float) ne00; + memcpy(dst_ptr, src0_ptr, type_size); + + if (++k10 == nk0) { + k10 = 0; + if (++i11 == ne1) { + i11 = 0; + if (++i12 == ne2) { + i12 = 0; + if (++i13 == ne3) { + i13 = 0; + } + } + } + } + } + } + k10 += nk00 * (ne01 - ir1); + while (k10 >= nk0) { + k10 -= nk0; + if (++i11 == ne1) { + i11 = 0; + if (++i12 == ne2) { + i12 = 0; + if (++i13 == ne3) { + i13 = 0; + } + } + } } } } } -void ggml_compute_forward_mean( +static void ggml_compute_forward_dup_from_q( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_mean_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} + GGML_TENSOR_BINARY_OP_LOCALS -// ggml_compute_forward_argmax + const ggml_type type = src0->type; + ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; -static void ggml_compute_forward_argmax_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { + size_t qk = ggml_blck_size(type); + const int64_t nr = ggml_nelements(src1) / qk; - const ggml_tensor * src0 = dst->src[0]; + // destination must be contiguous in the first dimension + GGML_ASSERT(nb10 == ggml_type_size(dst->type)); + // must either have first dimension large enough to hold a row, or fully contiguous + GGML_ASSERT((ne10 % qk) == 0 || ggml_is_contiguous(dst)); - if (params->ith != 0) { - return; - } + const int ith = params->ith; + const int nth = params->nth; - assert(src0->nb[0] == sizeof(float)); - assert(dst->nb[0] == sizeof(float)); + const int dr = (nr + nth - 1)/nth; - const int64_t ne00 = src0->ne[0]; - const int64_t ne01 = src0->ne[1]; + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - const size_t nb01 = src0->nb[1]; - const size_t nb0 = dst->nb[0]; + for (int64_t ir = ir0; ir < ir1; ++ir) { - for (int64_t i1 = 0; i1 < ne01; i1++) { - float * src = (float *) ((char *) src0->data + i1*nb01); - int32_t * dst_ = (int32_t *) ((char *) dst->data + i1*nb0); - int v = 0; - ggml_vec_argmax_f32(ne00, &v, src); - dst_[0] = v; + uint32_t i = ir * qk; + + const int64_t i03 = i/(ne00 * ne01 * ne02); + const int64_t i02 = (i - i03*ne00*ne01*ne02 )/ (ne00*ne01); + const int64_t i01 = (i - i03*ne00*ne01*ne02 - i02*ne01*ne00) / ne00; + const int64_t i00 = i - i03*ne00*ne01*ne02 - i02*ne01*ne00 - i01*ne00; + const int64_t x_offset = (i00/qk)*nb00 + i01*nb01 + i02*nb02 + i03 * nb03; + + const int64_t i13 = i/(ne10 * ne11 * ne12); + const int64_t i12 = (i - i13*ne10*ne11*ne12) / (ne10*ne11); + const int64_t i11 = (i - i13*ne10*ne11*ne12 - i12*ne10*ne11) / ne10; + const int64_t i10 = i - i13*ne10*ne11*ne12 - i12*ne10*ne11 - i11*ne10; + const int64_t dst_offset = i10*nb10 + i11*nb11 + i12*nb12 + i13*nb13; + + dequantize_row_q( + (const void *) ((char *) src0->data + x_offset), + (float *) ((char *) dst->data + dst_offset), qk); } } -void ggml_compute_forward_argmax( +void ggml_compute_forward_dup( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + if (src0->type == dst->type) { + ggml_compute_forward_dup_bytes(params, dst); + return; + } + switch (src0->type) { + case GGML_TYPE_F16: + { + /**/ if (dst->type == GGML_TYPE_F16) ggml_compute_forward_dup_flt(params, dst); + else if (dst->type == GGML_TYPE_BF16) ggml_compute_forward_dup_flt(params, dst); + else if (dst->type == GGML_TYPE_F32) ggml_compute_forward_dup_flt(params, dst); + else ggml_compute_forward_dup_to_q(params, dst); + } break; + case GGML_TYPE_BF16: + { + /**/ if (dst->type == GGML_TYPE_F16) ggml_compute_forward_dup_flt(params, dst); + else if (dst->type == GGML_TYPE_BF16) ggml_compute_forward_dup_flt(params, dst); + else if (dst->type == GGML_TYPE_F32) ggml_compute_forward_dup_flt(params, dst); + else ggml_compute_forward_dup_to_q(params, dst); + } break; case GGML_TYPE_F32: { - ggml_compute_forward_argmax_f32(params, dst); + /**/ if (dst->type == GGML_TYPE_F16) ggml_compute_forward_dup_flt(params, dst); + else if (dst->type == GGML_TYPE_BF16) ggml_compute_forward_dup_flt(params, dst); + else if (dst->type == GGML_TYPE_F32) ggml_compute_forward_dup_flt(params, dst); + else if (dst->type == GGML_TYPE_I32) ggml_compute_forward_dup_flt(params, dst); + else ggml_compute_forward_dup_to_q(params, dst); + } break; + case GGML_TYPE_I32: + { + if (dst->type == GGML_TYPE_F32) ggml_compute_forward_dup_flt(params, dst); + else GGML_ABORT("not implemented"); } break; default: { + if (ggml_is_quantized(src0->type) && dst->type == GGML_TYPE_F32) { + ggml_compute_forward_dup_from_q(params, dst); + break; + } GGML_ABORT("fatal error"); } } } -// ggml_compute_forward_count_equal +// ggml_compute_forward_add -static void ggml_compute_forward_count_equal_i32( +static void ggml_compute_forward_add_q_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - GGML_TENSOR_BINARY_OP_LOCALS; + GGML_ASSERT(ggml_are_same_shape(src0, src1) && ggml_are_same_shape(src0, dst)); - GGML_ASSERT(src0->type == GGML_TYPE_I32); - GGML_ASSERT(src1->type == GGML_TYPE_I32); - GGML_ASSERT(ggml_are_same_shape(src0, src1)); - GGML_ASSERT(ggml_is_scalar(dst)); - GGML_ASSERT(dst->type == GGML_TYPE_I64); + const int nr = ggml_nrows(src0); - const int64_t nr = ggml_nrows(src0); + GGML_TENSOR_BINARY_OP_LOCALS const int ith = params->ith; const int nth = params->nth; - int64_t * sums = (int64_t *) params->wdata; - int64_t sum_thread = 0; + const ggml_type type = src0->type; + const ggml_type dtype = dst->type; + ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; + ggml_from_float_t const quantize_row_q = ggml_get_type_traits_cpu(dtype)->from_float; - // rows per thread - const int64_t dr = (nr + nth - 1)/nth; + // we don't support permuted src0 or src1 + GGML_ASSERT(nb00 == ggml_type_size(type)); + GGML_ASSERT(nb10 == sizeof(float)); - // row range for this thread - const int64_t ir0 = dr*ith; - const int64_t ir1 = MIN(ir0 + dr, nr); + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); - for (int64_t ir = ir0; ir < ir1; ++ir) { - const int64_t i03 = ir / (ne02*ne01); - const int64_t i02 = (ir - i03*ne03) / ne01; - const int64_t i01 = ir - i03*ne03 - i02*ne02; + GGML_ASSERT(ggml_is_quantized(src0->type)); + GGML_ASSERT(src1->type == GGML_TYPE_F32); - const char * data0 = (const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01; - const char * data1 = (const char *) src1->data + i03*nb13 + i02*nb12 + i01*nb11; + // rows per thread + const int dr = (nr + nth - 1)/nth; - for (int64_t i00 = 0; i00 < ne00; ++i00) { - const int32_t val0 = *((const int32_t *) (data0 + i00*nb00)); - const int32_t val1 = *((const int32_t *) (data1 + i00*nb10)); + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - sum_thread += val0 == val1; - } - } - if (ith != 0) { - sums[ith] = sum_thread; - } - ggml_barrier(params->threadpool); + float * wdata = (float *) params->wdata + (ne00 + CACHE_LINE_SIZE_F32) * ith; - if (ith != 0) { - return; - } + for (int ir = ir0; ir < ir1; ++ir) { + // src0 indices + const int i03 = ir/(ne02*ne01); + const int i02 = (ir - i03*ne02*ne01)/ne01; + const int i01 = (ir - i03*ne02*ne01 - i02*ne01); - for (int ith_other = 1; ith_other < nth; ++ith_other) { - sum_thread += sums[ith_other]; + // src1 and dst are same shape as src0 => same indices + const int i13 = i03; + const int i12 = i02; + const int i11 = i01; + + const int i3 = i03; + const int i2 = i02; + const int i1 = i01; + + void * src0_row = (void *) ((char *) src0->data + (i01*nb01 + i02*nb02 + i03*nb03)); + float * src1_row = (float *)((char *) src1->data + (i11*nb11 + i12*nb12 + i13*nb13)); + void * dst_row = (void *) ((char *) dst->data + ( i1*nb1 + i2*nb2 + i3*nb3)); + + assert(ne00 % 32 == 0); + + // unquantize row from src0 to temp buffer + dequantize_row_q(src0_row, wdata, ne00); + // add src1 + ggml_vec_acc_f32(ne00, wdata, src1_row); + // quantize row to dst + if (quantize_row_q != NULL) { + quantize_row_q(wdata, dst_row, ne00); + } else { + memcpy(dst_row, wdata, ne0*nb0); + } } - *((int64_t *) dst->data) = sum_thread; } -void ggml_compute_forward_count_equal( +void ggml_compute_forward_add( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { - case GGML_TYPE_I32: + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_BF16: { - ggml_compute_forward_count_equal_i32(params, dst); + ggml_compute_forward_add_non_quantized(params, dst); + } break; + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_TQ1_0: + case GGML_TYPE_TQ2_0: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ2_S: + { + ggml_compute_forward_add_q_f32(params, dst); } break; default: { @@ -1690,479 +2771,364 @@ void ggml_compute_forward_count_equal( } } -// ggml_compute_forward_repeat +// ggml_compute_forward_add_id -static void ggml_compute_forward_repeat_f32( +static void ggml_compute_forward_add_id_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + const ggml_tensor * src2 = dst->src[2]; - if (params->ith != 0) { - return; - } - - GGML_ASSERT(ggml_can_repeat(src0, dst)); - - GGML_TENSOR_UNARY_OP_LOCALS + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(src2->type == GGML_TYPE_I32); - // guaranteed to be an integer due to the check in ggml_can_repeat - const int nr0 = (int)(ne0/ne00); - const int nr1 = (int)(ne1/ne01); - const int nr2 = (int)(ne2/ne02); - const int nr3 = (int)(ne3/ne03); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(src1->nb[0] == sizeof(float)); - // TODO: support for transposed / permuted tensors - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb00 == sizeof(float)); + const int ith = params->ith; + const int nth = params->nth; - // TODO: maybe this is not optimal? - for (int i3 = 0; i3 < nr3; i3++) { - for (int k3 = 0; k3 < ne03; k3++) { - for (int i2 = 0; i2 < nr2; i2++) { - for (int k2 = 0; k2 < ne02; k2++) { - for (int i1 = 0; i1 < nr1; i1++) { - for (int k1 = 0; k1 < ne01; k1++) { - for (int i0 = 0; i0 < nr0; i0++) { - ggml_vec_cpy_f32(ne00, - (float *) ((char *) dst->data + (i3*ne03 + k3)*nb3 + (i2*ne02 + k2)*nb2 + (i1*ne01 + k1)*nb1 + (i0*ne00)*nb0), - (float *) ((char *) src0->data + ( k3)*nb03 + ( k2)*nb02 + ( k1)*nb01)); - } - } - } - } - } - } - } -} + const int nr = ggml_nrows(src0); -static void ggml_compute_forward_repeat_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { + GGML_TENSOR_TERNARY_OP_LOCALS - const ggml_tensor * src0 = dst->src[0]; + GGML_ASSERT( nb0 == sizeof(float)); + GGML_ASSERT(nb10 == sizeof(float)); - if (params->ith != 0) { - return; - } + // rows per thread + const int dr = (nr + nth - 1)/nth; - GGML_ASSERT(ggml_can_repeat(src0, dst)); + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - GGML_TENSOR_UNARY_OP_LOCALS + for (int ir = ir0; ir < ir1; ++ir) { + // src0 indices + const int i3 = ir/(ne2*ne1); + const int i2 = (ir - i3*ne2*ne1)/ne1; + const int i1 = (ir - i3*ne2*ne1 - i2*ne1); - // guaranteed to be an integer due to the check in ggml_can_repeat - const int nr0 = (int)(ne0/ne00); - const int nr1 = (int)(ne1/ne01); - const int nr2 = (int)(ne2/ne02); - const int nr3 = (int)(ne3/ne03); + // src1 indices + const int i11 = *(int32_t *) ((char *) src2->data + i1*nb20 + i2*nb21); - // TODO: support for transposed / permuted tensors - GGML_ASSERT(nb0 == sizeof(ggml_fp16_t)); - GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); + GGML_ASSERT(i11 >= 0 && i11 < ne11); - // TODO: maybe this is not optimal? - for (int i3 = 0; i3 < nr3; i3++) { - for (int k3 = 0; k3 < ne03; k3++) { - for (int i2 = 0; i2 < nr2; i2++) { - for (int k2 = 0; k2 < ne02; k2++) { - for (int i1 = 0; i1 < nr1; i1++) { - for (int k1 = 0; k1 < ne01; k1++) { - for (int i0 = 0; i0 < nr0; i0++) { - ggml_fp16_t * y = (ggml_fp16_t *) ((char *) dst->data + (i3*ne03 + k3)*nb3 + (i2*ne02 + k2)*nb2 + (i1*ne01 + k1)*nb1 + (i0*ne00)*nb0); - ggml_fp16_t * x = (ggml_fp16_t *) ((char *) src0->data + ( k3)*nb03 + ( k2)*nb02 + ( k1)*nb01); - // ggml_vec_cpy_f16(ne00, y, x) - for (int i = 0; i < ne00; ++i) { - y[i] = x[i]; - } - } - } - } - } - } - } + ggml_vec_add_f32(ne0, + (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ), + (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01), + (float *) ((char *) src1->data + i11*nb11)); } } -void ggml_compute_forward_repeat( +void ggml_compute_forward_add_id( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { - case GGML_TYPE_F16: - case GGML_TYPE_BF16: - case GGML_TYPE_I16: - { - ggml_compute_forward_repeat_f16(params, dst); - } break; case GGML_TYPE_F32: - case GGML_TYPE_I32: { - ggml_compute_forward_repeat_f32(params, dst); + ggml_compute_forward_add_id_f32(params, dst); } break; - // TODO: templateify the implementation and support for I64 - // ref https://github.com/ggml-org/llama.cpp/pull/14274#discussion_r2169492225 - //case GGML_TYPE_I64: - // { - // ggml_compute_forward_repeat_i64(params, dst); - // } break; default: { - GGML_ABORT("fatal error"); + GGML_ABORT("unsupported type for ggml_compute_forward_add_id: %s", ggml_type_name(src0->type)); } } } -// ggml_compute_forward_repeat_back +// ggml_compute_forward_add1 -static void ggml_compute_forward_repeat_back_f32( +static void ggml_compute_forward_add1_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - if (params->ith != 0) { - return; - } + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_scalar(src1)); - GGML_ASSERT(ggml_can_repeat(dst, src0)); + const int ith = params->ith; + const int nth = params->nth; - GGML_TENSOR_UNARY_OP_LOCALS + const int nr = ggml_nrows(src0); - // guaranteed to be an integer due to the check in ggml_can_repeat - const int nr0 = (int)(ne00/ne0); - const int nr1 = (int)(ne01/ne1); - const int nr2 = (int)(ne02/ne2); - const int nr3 = (int)(ne03/ne3); + GGML_TENSOR_UNARY_OP_LOCALS - // TODO: support for transposed / permuted tensors - GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT( nb0 == sizeof(float)); GGML_ASSERT(nb00 == sizeof(float)); - if (ggml_is_contiguous(dst)) { - ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0); - } else { - for (int k3 = 0; k3 < ne3; k3++) { - for (int k2 = 0; k2 < ne2; k2++) { - for (int k1 = 0; k1 < ne1; k1++) { - ggml_vec_set_f32(ne0, - (float *) ((char *) dst->data + k1*nb1 + k2*nb2 + k3*nb3), - 0); - } - } - } - } + // rows per thread + const int dr = (nr + nth - 1)/nth; - // TODO: maybe this is not optimal? - for (int i3 = 0; i3 < nr3; i3++) { - for (int k3 = 0; k3 < ne3; k3++) { - for (int i2 = 0; i2 < nr2; i2++) { - for (int k2 = 0; k2 < ne2; k2++) { - for (int i1 = 0; i1 < nr1; i1++) { - for (int k1 = 0; k1 < ne1; k1++) { - for (int i0 = 0; i0 < nr0; i0++) { - ggml_vec_acc_f32(ne0, - (float *) ((char *) dst->data + ( k3)*nb3 + ( k2)*nb2 + ( k1)*nb1), - (float *) ((char *) src0->data + (i3*ne3 + k3)*nb03 + (i2*ne2 + k2)*nb02 + (i1*ne1 + k1)*nb01 + (i0*ne0)*nb00)); - } - } - } - } - } - } - } -} + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); -void ggml_compute_forward_repeat_back( - const ggml_compute_params * params, - ggml_tensor * dst) { + for (int ir = ir0; ir < ir1; ++ir) { + // src0 and dst are same shape => same indices + const int i3 = ir/(ne2*ne1); + const int i2 = (ir - i3*ne2*ne1)/ne1; + const int i1 = (ir - i3*ne2*ne1 - i2*ne1); - const ggml_tensor * src0 = dst->src[0]; +#ifdef GGML_USE_ACCELERATE + GGML_UNUSED(ggml_vec_add1_f32); - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_repeat_back_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } + vDSP_vadd( + (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01), 1, + (float *) ((char *) src1->data), 0, + (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ), 1, + ne0); +#else + ggml_vec_add1_f32(ne0, + (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ), + (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01), + *(float *) src1->data); +#endif } } -// ggml_compute_forward_concat - -static void ggml_compute_forward_concat_any( - const ggml_compute_params * params, - ggml_tensor * dst) { +static void ggml_compute_forward_add1_f16_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - const size_t len = ggml_type_size(src0->type); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_scalar(src1)); + + // scalar to add + const float v = *(float *) src1->data; const int ith = params->ith; const int nth = params->nth; - GGML_TENSOR_BINARY_OP_LOCALS + const int nr = ggml_nrows(src0); - const int32_t dim = ggml_get_op_params_i32(dst, 0); + GGML_TENSOR_UNARY_OP_LOCALS - GGML_ASSERT(dim >= 0 && dim < 4); + GGML_ASSERT(src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F16); - int64_t o[4] = {0, 0, 0, 0}; - o[dim] = src0->ne[dim]; + GGML_ASSERT( nb0 == sizeof(ggml_fp16_t)); + GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); - const char * x; + // rows per thread + const int dr = (nr + nth - 1)/nth; - // TODO: smarter multi-theading - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = ith; i2 < ne2; i2 += nth) { - for (int i1 = 0; i1 < ne1; i1++) { - for (int i0 = 0; i0 < ne0; i0++) { - if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { - x = (const char *)src0->data + (i0 )*nb00 + (i1 )*nb01 + (i2 )*nb02 + (i3 )*nb03; - } else { - x = (const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13; - } + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - char * y = (char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3; + for (int ir = ir0; ir < ir1; ++ir) { + // src0 and dst are same shape => same indices + const int i3 = ir/(ne2*ne1); + const int i2 = (ir - i3*ne2*ne1)/ne1; + const int i1 = (ir - i3*ne2*ne1 - i2*ne1); - memcpy(y, x, len); - } - } + ggml_fp16_t * dst_ptr = (ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ); + ggml_fp16_t * src0_ptr = (ggml_fp16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01); + for (int i = 0; i < ne0; i++) { + dst_ptr[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(src0_ptr[i]) + v); } } } -static void ggml_compute_forward_concat_i8( - const ggml_compute_params * params, - ggml_tensor * dst) { +static void ggml_compute_forward_add1_f16_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(ggml_type_size(src0->type) == sizeof(int8_t)); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_scalar(src1)); + + // scalar to add + const float v = GGML_CPU_FP16_TO_FP32(*(ggml_fp16_t *) src1->data); const int ith = params->ith; const int nth = params->nth; - GGML_TENSOR_BINARY_OP_LOCALS + const int nr = ggml_nrows(src0); - const int32_t dim = ggml_get_op_params_i32(dst, 0); + GGML_TENSOR_UNARY_OP_LOCALS - GGML_ASSERT(dim >= 0 && dim < 4); + GGML_ASSERT(src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_F16); + GGML_ASSERT(dst->type == GGML_TYPE_F16); - int64_t o[4] = {0, 0, 0, 0}; - o[dim] = src0->ne[dim]; + GGML_ASSERT( nb0 == sizeof(ggml_fp16_t)); + GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); - const int8_t * x; + // rows per thread + const int dr = (nr + nth - 1)/nth; - // TODO: smarter multi-theading - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = ith; i2 < ne2; i2 += nth) { - for (int i1 = 0; i1 < ne1; i1++) { - for (int i0 = 0; i0 < ne0; i0++) { - if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { - x = (const int8_t *) ((const char *)src0->data + (i0 )*nb00 + (i1 )*nb01 + (i2 )*nb02 + (i3 )*nb03); - } else { - x = (const int8_t *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13); - } + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - int8_t * y = (int8_t *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); + for (int ir = ir0; ir < ir1; ++ir) { + // src0 and dst are same shape => same indices + const int i3 = ir/(ne2*ne1); + const int i2 = (ir - i3*ne2*ne1)/ne1; + const int i1 = (ir - i3*ne2*ne1 - i2*ne1); - *y = *x; - } - } + ggml_fp16_t * dst_ptr = (ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ); + ggml_fp16_t * src0_ptr = (ggml_fp16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01); + for (int i = 0; i < ne0; i++) { + dst_ptr[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(src0_ptr[i]) + v); } } } -static void ggml_compute_forward_concat_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { +static void ggml_compute_forward_add1_q_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(ggml_type_size(src0->type) == sizeof(ggml_fp16_t)); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_scalar(src1)); + + // scalar to add + const float v = *(float *) src1->data; const int ith = params->ith; const int nth = params->nth; - GGML_TENSOR_BINARY_OP_LOCALS + const int nr = ggml_nrows(src0); - const int32_t dim = ggml_get_op_params_i32(dst, 0); + GGML_TENSOR_UNARY_OP_LOCALS - GGML_ASSERT(dim >= 0 && dim < 4); + const ggml_type type = src0->type; + ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; + ggml_from_float_t const quantize_row_q = ggml_get_type_traits_cpu(type)->from_float; - int64_t o[4] = {0, 0, 0, 0}; - o[dim] = src0->ne[dim]; + // we don't support permuted src0 + GGML_ASSERT(nb00 == ggml_type_size(type)); - const ggml_fp16_t * x; + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); - // TODO: smarter multi-theading - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = ith; i2 < ne2; i2 += nth) { - for (int i1 = 0; i1 < ne1; i1++) { - for (int i0 = 0; i0 < ne0; i0++) { - if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { - x = (const ggml_fp16_t *) ((const char *)src0->data + (i0 )*nb00 + (i1 )*nb01 + (i2 )*nb02 + (i3 )*nb03); - } else { - x = (const ggml_fp16_t *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13); - } + GGML_ASSERT(ggml_is_quantized(src0->type)); + GGML_ASSERT(dst->type == src0->type); + GGML_ASSERT(src1->type == GGML_TYPE_F32); - ggml_fp16_t * y = (ggml_fp16_t *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); + // rows per thread + const int dr = (nr + nth - 1)/nth; - *y = *x; - } - } - } + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + + float * wdata = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32) * ith; + + for (int ir = ir0; ir < ir1; ++ir) { + // src0 and dst are same shape => same indices + const int i3 = ir/(ne2*ne1); + const int i2 = (ir - i3*ne2*ne1)/ne1; + const int i1 = (ir - i3*ne2*ne1 - i2*ne1); + + void * src0_row = (void *) ((char *) src0->data + (i1*nb01 + i2*nb02 + i3*nb03)); + void * dst_row = (void *) ((char *) dst->data + (i1*nb1 + i2*nb2 + i3*nb0 )); + + assert(ne0 % 32 == 0); + + // unquantize row from src0 to temp buffer + dequantize_row_q(src0_row, wdata, ne0); + // add src1 + ggml_vec_acc1_f32(ne0, wdata, v); + // quantize row to dst + quantize_row_q(wdata, dst_row, ne0); } } -static void ggml_compute_forward_concat_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { +static void ggml_compute_forward_add1_bf16_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(ggml_type_size(src0->type) == sizeof(float)); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_scalar(src1)); + + // scalar to add + const float v = *(float *) src1->data; const int ith = params->ith; const int nth = params->nth; - GGML_TENSOR_BINARY_OP_LOCALS + const int nr = ggml_nrows(src0); - const int32_t dim = ggml_get_op_params_i32(dst, 0); + GGML_TENSOR_UNARY_OP_LOCALS - GGML_ASSERT(dim >= 0 && dim < 4); + GGML_ASSERT(src0->type == GGML_TYPE_BF16); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_BF16); - int64_t o[4] = {0, 0, 0, 0}; - o[dim] = src0->ne[dim]; + GGML_ASSERT( nb0 == sizeof(ggml_bf16_t)); + GGML_ASSERT(nb00 == sizeof(ggml_bf16_t)); - const float * x; + // rows per thread + const int dr = (nr + nth - 1)/nth; - // TODO: smarter multi-theading - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = ith; i2 < ne2; i2 += nth) { - for (int i1 = 0; i1 < ne1; i1++) { - for (int i0 = 0; i0 < ne0; i0++) { - if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { - x = (const float *) ((const char *)src0->data + (i0 )*nb00 + (i1 )*nb01 + (i2 )*nb02 + (i3 )*nb03); - } else { - x = (const float *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13); - } + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - float * y = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); + for (int ir = ir0; ir < ir1; ++ir) { + // src0 and dst are same shape => same indices + const int i3 = ir/(ne2*ne1); + const int i2 = (ir - i3*ne2*ne1)/ne1; + const int i1 = (ir - i3*ne2*ne1 - i2*ne1); - *y = *x; - } - } + ggml_bf16_t * dst_ptr = (ggml_bf16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ); + ggml_bf16_t * src0_ptr = (ggml_bf16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01); + for (int i = 0; i < ne0; i++) { + dst_ptr[i] = GGML_FP32_TO_BF16(GGML_BF16_TO_FP32(src0_ptr[i]) + v); } } } -void ggml_compute_forward_concat( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - - switch (src0->type) { - case GGML_TYPE_F16: - case GGML_TYPE_BF16: - case GGML_TYPE_I16: - { - ggml_compute_forward_concat_f16(params, dst); - } break; - case GGML_TYPE_I8: - { - ggml_compute_forward_concat_i8(params, dst); - } break; - case GGML_TYPE_F32: - case GGML_TYPE_I32: - { - ggml_compute_forward_concat_f32(params, dst); - } break; - default: - { - ggml_compute_forward_concat_any(params, dst); - } - } -} - -// ggml_compute_forward_gelu - -static void ggml_compute_forward_gelu_f32( +static void ggml_compute_forward_add1_bf16_bf16( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - assert(ggml_is_contiguous_rows(src0)); - assert(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_scalar(src1)); - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + // scalar to add + const float v = GGML_BF16_TO_FP32(*(ggml_bf16_t *) src1->data); const int ith = params->ith; const int nth = params->nth; - const int nc = src0->ne[0]; - const int nr = ggml_nrows(src0); - - // rows per thread - const int dr = (nr + nth - 1)/nth; - - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); - - for (int ir = ir0; ir < ir1; ++ir) { - const int i3 = ir/(ne02*ne01); - const int i2 = (ir - i3*ne02*ne01)/ne01; - const int i1 = (ir - i3*ne02*ne01 - i2*ne01); - - ggml_vec_gelu_f32(nc, - (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), - (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); - -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const float x = ((float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*(dst->nb[1])))[k]; - GGML_UNUSED(x); - assert(!isnan(x)); - assert(!isinf(x)); - } -#endif // NDEBUG - } -} - -static void ggml_compute_forward_gelu_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - - assert(ggml_is_contiguous_rows(src0)); - assert(ggml_are_same_shape(src0, dst)); + const int nr = ggml_nrows(src0); - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + GGML_TENSOR_UNARY_OP_LOCALS - const int ith = params->ith; - const int nth = params->nth; + GGML_ASSERT(src0->type == GGML_TYPE_BF16); + GGML_ASSERT(src1->type == GGML_TYPE_BF16); + GGML_ASSERT(dst->type == GGML_TYPE_BF16); - const int nc = src0->ne[0]; - const int nr = ggml_nrows(src0); + GGML_ASSERT( nb0 == sizeof(ggml_bf16_t)); + GGML_ASSERT(nb00 == sizeof(ggml_bf16_t)); // rows per thread const int dr = (nr + nth - 1)/nth; @@ -2172,40 +3138,82 @@ static void ggml_compute_forward_gelu_f16( const int ir1 = MIN(ir0 + dr, nr); for (int ir = ir0; ir < ir1; ++ir) { - const int i3 = ir/(ne02*ne01); - const int i2 = (ir - i3*ne02*ne01)/ne01; - const int i1 = (ir - i3*ne02*ne01 - i2*ne01); - - ggml_vec_gelu_f16(nc, - (ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), - (ggml_fp16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); + // src0 and dst are same shape => same indices + const int i3 = ir/(ne2*ne1); + const int i2 = (ir - i3*ne2*ne1)/ne1; + const int i1 = (ir - i3*ne2*ne1 - i2*ne1); -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*( dst->nb[1])))[k]; - const float v = GGML_CPU_FP16_TO_FP32(x); - GGML_UNUSED(v); - assert(!isnan(v)); - assert(!isinf(v)); + ggml_bf16_t * dst_ptr = (ggml_bf16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 ); + ggml_bf16_t * src0_ptr = (ggml_bf16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01); + for (int i = 0; i < ne0; i++) { + dst_ptr[i] = GGML_FP32_TO_BF16(GGML_BF16_TO_FP32(src0_ptr[i]) + v); } -#endif // NDEBUG } } -static void ggml_compute_forward_gelu( +void ggml_compute_forward_add1( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_gelu_f32(params, dst); + ggml_compute_forward_add1_f32(params, dst); } break; case GGML_TYPE_F16: { - ggml_compute_forward_gelu_f16(params, dst); + if (src1->type == GGML_TYPE_F16) { + ggml_compute_forward_add1_f16_f16(params, dst); + } + else if (src1->type == GGML_TYPE_F32) { + ggml_compute_forward_add1_f16_f32(params, dst); + } + else { + GGML_ABORT("fatal error"); + } + } break; + case GGML_TYPE_BF16: + { + if (src1->type == GGML_TYPE_BF16) { + ggml_compute_forward_add1_bf16_bf16(params, dst); + } + else if (src1->type == GGML_TYPE_F32) { + ggml_compute_forward_add1_bf16_f32(params, dst); + } + else { + GGML_ABORT("fatal error"); + } + } break; + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q8_1: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_TQ1_0: + case GGML_TYPE_TQ2_0: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ2_S: + { + ggml_compute_forward_add1_q_f32(params, dst); } break; default: { @@ -2214,76 +3222,126 @@ static void ggml_compute_forward_gelu( } } -// ggml_compute_fill - -static void ggml_compute_forward_fill_f32(const ggml_compute_params * params, ggml_tensor * dst) { - const float c = ggml_get_op_params_f32(dst, 0); +// ggml_compute_forward_acc - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne); - GGML_TENSOR_LOCALS(size_t, nb, dst, nb); +static void ggml_compute_forward_acc_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - const auto [ir0, ir1] = get_thread_range(params, dst); + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - for (int64_t ir = ir0; ir < ir1; ++ir) { - const int64_t i03 = ir/(ne2*ne1); - const int64_t i02 = (ir - i03*ne2*ne1)/ne1; - const int64_t i01 = (ir - i03*ne2*ne1 - i02*ne1); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src0)); - float * dst_ptr = (float *) ((char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1); + // view src0 and dst with these strides and data offset inbytes during acc + // nb0 is implicitly element_size because src0 and dst are contiguous + size_t nb1 = ((int32_t *) dst->op_params)[0]; + size_t nb2 = ((int32_t *) dst->op_params)[1]; + size_t nb3 = ((int32_t *) dst->op_params)[2]; + size_t offset = ((int32_t *) dst->op_params)[3]; + bool inplace = (bool) ((int32_t *) dst->op_params)[4]; - ggml_vec_set_f32(ne0, dst_ptr, c); + if (!inplace) { + if (params->ith == 0) { + // memcpy needs to be synchronized across threads to avoid race conditions. + // => do it in INIT phase + memcpy( + ((char *) dst->data), + ((char *) src0->data), + ggml_nbytes(dst)); + } + ggml_barrier(params->threadpool); } -} -void ggml_compute_forward_fill(const ggml_compute_params * params, ggml_tensor * dst) { - ggml_compute_forward_fill_f32(params, dst); -} + const int ith = params->ith; + const int nth = params->nth; -// ggml_compute_tri + const int nr = ggml_nrows(src1); + const int nc = src1->ne[0]; -static void ggml_compute_forward_tri_f32(const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; + GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) + GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) - const ggml_tri_type ttype = (ggml_tri_type) ggml_get_op_params_i32(dst, 0); + // src0 and dst as viewed during acc + const size_t nb0 = ggml_element_size(src0); - GGML_ASSERT(ggml_is_contiguous(src0)); + const size_t nb00 = nb0; + const size_t nb01 = nb1; + const size_t nb02 = nb2; + const size_t nb03 = nb3; - GGML_TENSOR_UNARY_OP_LOCALS + GGML_ASSERT(offset + (ne10 == 0 ? 0 : ne10-1)*nb0 + (ne11 == 0 ? 0 : ne11-1)*nb1 + (ne12 == 0 ? 0 : ne12-1)*nb2 + (ne13 == 0 ? 0 : ne13-1)*nb3 < ggml_nbytes(dst)); + GGML_ASSERT(offset + (ne10 == 0 ? 0 : ne10-1)*nb00 + (ne11 == 0 ? 0 : ne11-1)*nb01 + (ne12 == 0 ? 0 : ne12-1)*nb02 + (ne13 == 0 ? 0 : ne13-1)*nb03 < ggml_nbytes(src0)); - const auto [ir0, ir1] = get_thread_range(params, src0); + GGML_ASSERT(nb10 == sizeof(float)); - bool (*bipred)(int, int); - - switch (ttype) { - case GGML_TRI_TYPE_LOWER: bipred = [](int i, int r) { return i < r; }; break; - case GGML_TRI_TYPE_LOWER_DIAG: bipred = [](int i, int r) { return i <= r; }; break; - case GGML_TRI_TYPE_UPPER: bipred = [](int i, int r) { return i > r; }; break; - case GGML_TRI_TYPE_UPPER_DIAG: bipred = [](int i, int r) { return i >= r; }; break; - default: GGML_ABORT("invalid tri type"); - } + // rows per thread + const int dr = (nr + nth - 1)/nth; - for (int64_t ir = ir0; ir < ir1; ++ir) { - const int64_t i03 = ir/(ne02*ne01); - const int64_t i02 = (ir - i03*ne02*ne01)/ne01; - const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - const float * src_ptr = (const float *) ((const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01); - float * dst_ptr = ( float *) (( char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1); + for (int ir = ir0; ir < ir1; ++ir) { + // src0 and dst are viewed with shape of src1 and offset + // => same indices + const int i3 = ir/(ne12*ne11); + const int i2 = (ir - i3*ne12*ne11)/ne11; + const int i1 = (ir - i3*ne12*ne11 - i2*ne11); - for (int i0 = 0; i0 < ne0; ++i0) { - dst_ptr[i0] = bipred(i0, i01) ? src_ptr[i0] : 0.0f; - } +#ifdef GGML_USE_ACCELERATE + vDSP_vadd( + (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + offset), 1, + (float *) ((char *) src1->data + i3*nb13 + i2*nb12 + i1*nb11), 1, + (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + offset), 1, nc); +#else + ggml_vec_add_f32(nc, + (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + offset), + (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + offset), + (float *) ((char *) src1->data + i3*nb13 + i2*nb12 + i1*nb11)); +#endif } } -void ggml_compute_forward_tri(const ggml_compute_params * params, ggml_tensor * dst) { +void ggml_compute_forward_acc( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_tri_f32(params, dst); + ggml_compute_forward_acc_f32(params, dst); } break; + case GGML_TYPE_F16: + case GGML_TYPE_BF16: + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q8_1: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_TQ1_0: + case GGML_TYPE_TQ2_0: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ2_S: default: { GGML_ABORT("fatal error"); @@ -2291,104 +3349,107 @@ void ggml_compute_forward_tri(const ggml_compute_params * params, ggml_tensor * } } -// ggml_compute_forward_gelu_erf +// ggml_compute_forward_sum -static void ggml_compute_forward_gelu_erf_f32( +static void ggml_compute_forward_sum_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - assert(ggml_is_contiguous_rows(src0)); - assert(ggml_are_same_shape(src0, dst)); + if (params->ith != 0) { + return; + } + + assert(ggml_is_scalar(dst)); + assert(src0->nb[0] == sizeof(float)); GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) - - const int ith = params->ith; - const int nth = params->nth; - - const int nc = src0->ne[0]; - const int nr = ggml_nrows(src0); - - // rows per thread - const int dr = (nr + nth - 1)/nth; - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); - - for (int ir = ir0; ir < ir1; ++ir) { - const int i3 = ir/(ne02*ne01); - const int i2 = (ir - i3*ne02*ne01)/ne01; - const int i1 = (ir - i3*ne02*ne01 - i2*ne01); - - ggml_vec_gelu_erf_f32(nc, - (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), - (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); + ggml_float sum = 0; + ggml_float row_sum = 0; -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const float x = ((float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*(dst->nb[1])))[k]; - GGML_UNUSED(x); - assert(!isnan(x)); - assert(!isinf(x)); + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = 0; i01 < ne01; i01++) { + ggml_vec_sum_f32_ggf(ne00, + &row_sum, + (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03)); + sum += row_sum; + } } -#endif // NDEBUG } + ((float *) dst->data)[0] = sum; } -static void ggml_compute_forward_gelu_erf_f16( +static void ggml_compute_forward_sum_f16( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - assert(ggml_is_contiguous_rows(src0)); - assert(ggml_are_same_shape(src0, dst)); + if (params->ith != 0) { + return; + } + + assert(ggml_is_scalar(dst)); + + assert(src0->nb[0] == sizeof(ggml_fp16_t)); GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) - const int ith = params->ith; - const int nth = params->nth; + float sum = 0; + float row_sum = 0; - const int nc = src0->ne[0]; - const int nr = ggml_nrows(src0); + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = 0; i01 < ne01; i01++) { + ggml_vec_sum_f16_ggf(ne00, + &row_sum, + (ggml_fp16_t *) ((char *) src0->data + i01 * nb01 + i02 * nb02 + i03 * nb03)); + sum += row_sum; + } + } + } + ((ggml_fp16_t *) dst->data)[0] = GGML_CPU_FP32_TO_FP16(sum); +} - // rows per thread - const int dr = (nr + nth - 1)/nth; +static void ggml_compute_forward_sum_bf16( + const ggml_compute_params * params, + ggml_tensor * dst) { - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + const ggml_tensor * src0 = dst->src[0]; - for (int ir = ir0; ir < ir1; ++ir) { - const int i3 = ir/(ne02*ne01); - const int i2 = (ir - i3*ne02*ne01)/ne01; - const int i1 = (ir - i3*ne02*ne01 - i2*ne01); + if (params->ith != 0) { + return; + } - ggml_vec_gelu_erf_f16(nc, - (ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), - (ggml_fp16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); + assert(ggml_is_scalar(dst)); -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*( dst->nb[1])))[k]; - const float v = GGML_CPU_FP16_TO_FP32(x); - GGML_UNUSED(v); - assert(!isnan(v)); - assert(!isinf(v)); + assert(src0->nb[0] == sizeof(ggml_bf16_t)); + + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) + + float sum = 0; + float row_sum = 0; + + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = 0; i01 < ne01; i01++) { + ggml_vec_sum_bf16_ggf(ne00, + &row_sum, + (ggml_bf16_t *) ((char *) src0->data + i01 * nb01 + i02 * nb02 + i03 * nb03)); + sum += row_sum; + } } -#endif // NDEBUG } + ((ggml_bf16_t *) dst->data)[0] = GGML_FP32_TO_BF16(sum); } -static void ggml_compute_forward_gelu_erf( +void ggml_compute_forward_sum( const ggml_compute_params * params, ggml_tensor * dst) { @@ -2397,11 +3458,15 @@ static void ggml_compute_forward_gelu_erf( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_gelu_erf_f32(params, dst); + ggml_compute_forward_sum_f32(params, dst); } break; case GGML_TYPE_F16: { - ggml_compute_forward_gelu_erf_f16(params, dst); + ggml_compute_forward_sum_f16(params, dst); + } break; + case GGML_TYPE_BF16: + { + ggml_compute_forward_sum_bf16(params, dst); } break; default: { @@ -2410,104 +3475,92 @@ static void ggml_compute_forward_gelu_erf( } } -// ggml_compute_forward_gelu_quick +// ggml_compute_forward_cumsum -static void ggml_compute_forward_gelu_quick_f32( +static void ggml_compute_forward_cumsum_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - assert(ggml_is_contiguous_rows(src0)); - assert(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(dst->nb[0] == sizeof(float)); - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + GGML_TENSOR_UNARY_OP_LOCALS - const int ith = params->ith; - const int nth = params->nth; + GGML_ASSERT(ne0 == ne00); + GGML_ASSERT(ne1 == ne01); + GGML_ASSERT(ne2 == ne02); + GGML_ASSERT(ne3 == ne03); - const int nc = src0->ne[0]; - const int nr = ggml_nrows(src0); + const auto [ir0, ir1] = get_thread_range(params, src0); - // rows per thread - const int dr = (nr + nth - 1)/nth; + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t i03 = ir/(ne02*ne01); + const int64_t i02 = (ir - i03*ne02*ne01)/ne01; + const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + float * src_row = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); + float * dst_row = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); - for (int ir = ir0; ir < ir1; ++ir) { - const int i3 = ir/(ne02*ne01); - const int i2 = (ir - i3*ne02*ne01)/ne01; - const int i1 = (ir - i3*ne02*ne01 - i2*ne01); + ggml_vec_cumsum_f32(ne00, dst_row, src_row); + } +} - ggml_vec_gelu_quick_f32(nc, - (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), - (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); - -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const float x = ((float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*(dst->nb[1])))[k]; - GGML_UNUSED(x); - assert(!isnan(x)); - assert(!isinf(x)); - } -#endif // NDEBUG - } -} - -static void ggml_compute_forward_gelu_quick_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_cumsum( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - assert(ggml_is_contiguous_rows(src0)); - assert(ggml_are_same_shape(src0, dst)); + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_cumsum_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) +// ggml_compute_forward_sum_rows - const int ith = params->ith; - const int nth = params->nth; +static void ggml_compute_forward_sum_rows_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - const int nc = src0->ne[0]; - const int nr = ggml_nrows(src0); + const ggml_tensor * src0 = dst->src[0]; - // rows per thread - const int dr = (nr + nth - 1)/nth; + if (params->ith != 0) { + return; + } - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(dst->nb[0] == sizeof(float)); - for (int ir = ir0; ir < ir1; ++ir) { - const int i3 = ir/(ne02*ne01); - const int i2 = (ir - i3*ne02*ne01)/ne01; - const int i1 = (ir - i3*ne02*ne01 - i2*ne01); + GGML_TENSOR_UNARY_OP_LOCALS - ggml_vec_gelu_quick_f16(nc, - (ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), - (ggml_fp16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); + GGML_ASSERT(ne0 == 1); + GGML_ASSERT(ne1 == ne01); + GGML_ASSERT(ne2 == ne02); + GGML_ASSERT(ne3 == ne03); -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*( dst->nb[1])))[k]; - const float v = GGML_CPU_FP16_TO_FP32(x); - GGML_UNUSED(v); - assert(!isnan(v)); - assert(!isinf(v)); + for (int64_t i3 = 0; i3 < ne03; i3++) { + for (int64_t i2 = 0; i2 < ne02; i2++) { + for (int64_t i1 = 0; i1 < ne01; i1++) { + float * src_row = (float *) ((char *) src0->data + i1*nb01 + i2*nb02 + i3*nb03); + float * dst_row = (float *) ((char *) dst->data + i1*nb1 + i2*nb2 + i3*nb3); + float row_sum = 0; + ggml_vec_sum_f32(ne00, &row_sum, src_row); + dst_row[0] = row_sum; + } } -#endif // NDEBUG } } -static void ggml_compute_forward_gelu_quick( +void ggml_compute_forward_sum_rows( const ggml_compute_params * params, ggml_tensor * dst) { @@ -2516,11 +3569,7 @@ static void ggml_compute_forward_gelu_quick( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_gelu_quick_f32(params, dst); - } break; - case GGML_TYPE_F16: - { - ggml_compute_forward_gelu_quick_f16(params, dst); + ggml_compute_forward_sum_rows_f32(params, dst); } break; default: { @@ -2529,104 +3578,94 @@ static void ggml_compute_forward_gelu_quick( } } -// ggml_compute_forward_silu +// ggml_compute_forward_mean -static void ggml_compute_forward_silu_f32( +static void ggml_compute_forward_mean_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - assert(ggml_is_contiguous_rows(src0)); - assert(ggml_are_same_shape(src0, dst)); - - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) - - const int ith = params->ith; - const int nth = params->nth; + if (params->ith != 0) { + return; + } - const int nc = src0->ne[0]; - const int nr = ggml_nrows(src0); + assert(src0->nb[0] == sizeof(float)); - // rows per thread - const int dr = (nr + nth - 1)/nth; + GGML_TENSOR_UNARY_OP_LOCALS - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + assert(ne0 == 1); + assert(ne1 == ne01); + assert(ne2 == ne02); + assert(ne3 == ne03); - for (int ir = ir0; ir < ir1; ++ir) { - const int i3 = ir/(ne02*ne01); - const int i2 = (ir - i3*ne02*ne01)/ne01; - const int i1 = (ir - i3*ne02*ne01 - i2*ne01); + GGML_UNUSED(ne0); + GGML_UNUSED(ne1); + GGML_UNUSED(ne2); + GGML_UNUSED(ne3); - ggml_vec_silu_f32(nc, - (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), - (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = 0; i01 < ne01; i01++) { + ggml_vec_sum_f32(ne00, + (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3), + (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03)); -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const float x = ((float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*(dst->nb[1])))[k]; - GGML_UNUSED(x); - assert(!isnan(x)); - assert(!isinf(x)); + *(float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3) /= (float) ne00; + } } -#endif // NDEBUG } } -static void ggml_compute_forward_silu_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_mean( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - assert(ggml_is_contiguous_rows(src0)); - assert(ggml_are_same_shape(src0, dst)); + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_mean_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) +// ggml_compute_forward_argmax - const int ith = params->ith; - const int nth = params->nth; +static void ggml_compute_forward_argmax_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - const int nc = src0->ne[0]; - const int nr = ggml_nrows(src0); + const ggml_tensor * src0 = dst->src[0]; - // rows per thread - const int dr = (nr + nth - 1)/nth; + if (params->ith != 0) { + return; + } - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + assert(src0->nb[0] == sizeof(float)); + assert(dst->nb[0] == sizeof(float)); - for (int ir = ir0; ir < ir1; ++ir) { - const int i3 = ir/(ne02*ne01); - const int i2 = (ir - i3*ne02*ne01)/ne01; - const int i1 = (ir - i3*ne02*ne01 - i2*ne01); + const int64_t ne00 = src0->ne[0]; + const int64_t ne01 = src0->ne[1]; - ggml_vec_silu_f16(nc, - (ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), - (ggml_fp16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); + const size_t nb01 = src0->nb[1]; + const size_t nb0 = dst->nb[0]; -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*( dst->nb[1])))[k]; - const float v = GGML_CPU_FP16_TO_FP32(x); - GGML_UNUSED(v); - assert(!isnan(v)); - assert(!isinf(v)); - } -#endif // NDEBUG + for (int64_t i1 = 0; i1 < ne01; i1++) { + float * src = (float *) ((char *) src0->data + i1*nb01); + int32_t * dst_ = (int32_t *) ((char *) dst->data + i1*nb0); + int v = 0; + ggml_vec_argmax_f32(ne00, &v, src); + dst_[0] = v; } } -static void ggml_compute_forward_silu( +void ggml_compute_forward_argmax( const ggml_compute_params * params, ggml_tensor * dst) { @@ -2635,11 +3674,7 @@ static void ggml_compute_forward_silu( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_silu_f32(params, dst); - } break; - case GGML_TYPE_F16: - { - ggml_compute_forward_silu_f16(params, dst); + ggml_compute_forward_argmax_f32(params, dst); } break; default: { @@ -2647,82 +3682,79 @@ static void ggml_compute_forward_silu( } } } -// ggml_compute_forward_leaky_relu -static void ggml_compute_forward_leaky_relu_f32( +// ggml_compute_forward_count_equal + +static void ggml_compute_forward_count_equal_i32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - if (params->ith != 0) { - return; - } + GGML_TENSOR_BINARY_OP_LOCALS; - assert(ggml_is_contiguous_1(src0)); - assert(ggml_is_contiguous_1(dst)); - assert(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(src0->type == GGML_TYPE_I32); + GGML_ASSERT(src1->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_are_same_shape(src0, src1)); + GGML_ASSERT(ggml_is_scalar(dst)); + GGML_ASSERT(dst->type == GGML_TYPE_I64); - const int n = ggml_nrows(src0); - const int nc = src0->ne[0]; + const int64_t nr = ggml_nrows(src0); - float negative_slope; - memcpy(&negative_slope, dst->op_params, sizeof(float)); + const int ith = params->ith; + const int nth = params->nth; - assert(dst->nb[0] == sizeof(float)); - assert(src0->nb[0] == sizeof(float)); + int64_t * sums = (int64_t *) params->wdata; + int64_t sum_thread = 0; - for (int i = 0; i < n; i++) { - ggml_vec_leaky_relu_f32(nc, - (float *) ((char *) dst->data + i*( dst->nb[1])), - (float *) ((char *) src0->data + i*(src0->nb[1])), negative_slope); - } -} - -static void ggml_compute_forward_leaky_relu_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; - const ggml_tensor * src0 = dst->src[0]; + // row range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); - if (params->ith != 0) { - return; - } + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t i03 = ir / (ne02*ne01); + const int64_t i02 = (ir - i03*ne03) / ne01; + const int64_t i01 = ir - i03*ne03 - i02*ne02; - assert(ggml_is_contiguous_1(src0)); - assert(ggml_is_contiguous_1(dst)); - assert(ggml_are_same_shape(src0, dst)); + const char * data0 = (const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01; + const char * data1 = (const char *) src1->data + i03*nb13 + i02*nb12 + i01*nb11; - const int n = ggml_nrows(src0); - const int nc = src0->ne[0]; + for (int64_t i00 = 0; i00 < ne00; ++i00) { + const int32_t val0 = *((const int32_t *) (data0 + i00*nb00)); + const int32_t val1 = *((const int32_t *) (data1 + i00*nb10)); - float negative_slope; - memcpy(&negative_slope, dst->op_params, sizeof(float)); + sum_thread += val0 == val1; + } + } + if (ith != 0) { + sums[ith] = sum_thread; + } + ggml_barrier(params->threadpool); - assert(dst->nb[0] == sizeof(ggml_fp16_t)); - assert(src0->nb[0] == sizeof(ggml_fp16_t)); + if (ith != 0) { + return; + } - for (int i = 0; i < n; i++) { - ggml_vec_leaky_relu_f16(nc, - (ggml_fp16_t *) ((char *) dst->data + i*( dst->nb[1])), - (ggml_fp16_t *) ((char *) src0->data + i*(src0->nb[1])), negative_slope); + for (int ith_other = 1; ith_other < nth; ++ith_other) { + sum_thread += sums[ith_other]; } + *((int64_t *) dst->data) = sum_thread; } -void ggml_compute_forward_leaky_relu( +void ggml_compute_forward_count_equal( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_leaky_relu_f32(params, dst); - } break; - case GGML_TYPE_F16: + case GGML_TYPE_I32: { - ggml_compute_forward_leaky_relu_f16(params, dst); + ggml_compute_forward_count_equal_i32(params, dst); } break; default: { @@ -2731,110 +3763,123 @@ void ggml_compute_forward_leaky_relu( } } -// ggml_compute_forward_silu_back +// ggml_compute_forward_repeat -static void ggml_compute_forward_silu_back_f32( +static void ggml_compute_forward_repeat_f32( const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * grad = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - - assert(ggml_is_contiguous_1(grad)); - assert(ggml_is_contiguous_1(src1)); - assert(ggml_is_contiguous_1(dst)); - assert(ggml_are_same_shape(src1, dst)); - assert(ggml_are_same_shape(src1, grad)); + const ggml_tensor * src0 = dst->src[0]; - const int ith = params->ith; - const int nth = params->nth; + if (params->ith != 0) { + return; + } - const int nc = src1->ne[0]; - const int nr = ggml_nrows(src1); + GGML_ASSERT(ggml_can_repeat(src0, dst)); - // rows per thread - const int dr = (nr + nth - 1)/nth; + GGML_TENSOR_UNARY_OP_LOCALS - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + // guaranteed to be an integer due to the check in ggml_can_repeat + const int nr0 = (int)(ne0/ne00); + const int nr1 = (int)(ne1/ne01); + const int nr2 = (int)(ne2/ne02); + const int nr3 = (int)(ne3/ne03); - for (int i1 = ir0; i1 < ir1; i1++) { - ggml_vec_silu_backward_f32(nc, - (float *) ((char *) dst->data + i1*( dst->nb[1])), - (float *) ((char *) src1->data + i1*(src1->nb[1])), - (float *) ((char *) grad->data + i1*(grad->nb[1]))); + // TODO: support for transposed / permuted tensors + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb00 == sizeof(float)); -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const float x = ((float *) ((char *) dst->data + i1*( dst->nb[1])))[k]; - GGML_UNUSED(x); - assert(!isnan(x)); - assert(!isinf(x)); + // TODO: maybe this is not optimal? + for (int i3 = 0; i3 < nr3; i3++) { + for (int k3 = 0; k3 < ne03; k3++) { + for (int i2 = 0; i2 < nr2; i2++) { + for (int k2 = 0; k2 < ne02; k2++) { + for (int i1 = 0; i1 < nr1; i1++) { + for (int k1 = 0; k1 < ne01; k1++) { + for (int i0 = 0; i0 < nr0; i0++) { + ggml_vec_cpy_f32(ne00, + (float *) ((char *) dst->data + (i3*ne03 + k3)*nb3 + (i2*ne02 + k2)*nb2 + (i1*ne01 + k1)*nb1 + (i0*ne00)*nb0), + (float *) ((char *) src0->data + ( k3)*nb03 + ( k2)*nb02 + ( k1)*nb01)); + } + } + } + } + } } -#endif // NDEBUG } } -static void ggml_compute_forward_silu_back_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * grad = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; +static void ggml_compute_forward_repeat_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { - assert(ggml_is_contiguous_1(grad)); - assert(ggml_is_contiguous_1(src1)); - assert(ggml_is_contiguous_1(dst)); - assert(ggml_are_same_shape(src1, dst)); - assert(ggml_are_same_shape(src1, grad)); + const ggml_tensor * src0 = dst->src[0]; - const int ith = params->ith; - const int nth = params->nth; + if (params->ith != 0) { + return; + } - const int nc = src1->ne[0]; - const int nr = ggml_nrows(src1); + GGML_ASSERT(ggml_can_repeat(src0, dst)); - // rows per thread - const int dr = (nr + nth - 1)/nth; + GGML_TENSOR_UNARY_OP_LOCALS - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + // guaranteed to be an integer due to the check in ggml_can_repeat + const int nr0 = (int)(ne0/ne00); + const int nr1 = (int)(ne1/ne01); + const int nr2 = (int)(ne2/ne02); + const int nr3 = (int)(ne3/ne03); - for (int i1 = ir0; i1 < ir1; i1++) { - ggml_vec_silu_backward_f16(nc, - (ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])), - (ggml_fp16_t *) ((char *) src1->data + i1*(src1->nb[1])), - (ggml_fp16_t *) ((char *) grad->data + i1*(grad->nb[1]))); + // TODO: support for transposed / permuted tensors + GGML_ASSERT(nb0 == sizeof(ggml_fp16_t)); + GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const float x = ((ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])))[k]; - const float v = GGML_CPU_FP16_TO_FP32(x); - GGML_UNUSED(v); - assert(!isnan(v)); - assert(!isinf(v)); + // TODO: maybe this is not optimal? + for (int i3 = 0; i3 < nr3; i3++) { + for (int k3 = 0; k3 < ne03; k3++) { + for (int i2 = 0; i2 < nr2; i2++) { + for (int k2 = 0; k2 < ne02; k2++) { + for (int i1 = 0; i1 < nr1; i1++) { + for (int k1 = 0; k1 < ne01; k1++) { + for (int i0 = 0; i0 < nr0; i0++) { + ggml_fp16_t * y = (ggml_fp16_t *) ((char *) dst->data + (i3*ne03 + k3)*nb3 + (i2*ne02 + k2)*nb2 + (i1*ne01 + k1)*nb1 + (i0*ne00)*nb0); + ggml_fp16_t * x = (ggml_fp16_t *) ((char *) src0->data + ( k3)*nb03 + ( k2)*nb02 + ( k1)*nb01); + // ggml_vec_cpy_f16(ne00, y, x) + for (int i = 0; i < ne00; ++i) { + y[i] = x[i]; + } + } + } + } + } + } } -#endif // NDEBUG } } -void ggml_compute_forward_silu_back( +void ggml_compute_forward_repeat( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { - case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_BF16: + case GGML_TYPE_I16: { - ggml_compute_forward_silu_back_f32(params, dst); + ggml_compute_forward_repeat_f16(params, dst); } break; - case GGML_TYPE_F16: + case GGML_TYPE_F32: + case GGML_TYPE_I32: { - ggml_compute_forward_silu_back_f16(params, dst); + ggml_compute_forward_repeat_f32(params, dst); } break; + // TODO: templateify the implementation and support for I64 + // ref https://github.com/ggml-org/llama.cpp/pull/14274#discussion_r2169492225 + //case GGML_TYPE_I64: + // { + // ggml_compute_forward_repeat_i64(params, dst); + // } break; default: { GGML_ABORT("fatal error"); @@ -2842,128 +3887,67 @@ void ggml_compute_forward_silu_back( } } -// ggml_compute_forward_reglu +// ggml_compute_forward_repeat_back -static void ggml_compute_forward_reglu_f32( +static void ggml_compute_forward_repeat_back_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - char * src0_d = (char *) src0->data; - char * src1_d = (char *) (src1 ? src1->data : src0->data); - const size_t src0_o = src0->nb[1]; - const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - - GGML_ASSERT(ggml_is_contiguous_1(src0)); - GGML_ASSERT(ggml_is_contiguous_1(dst)); - if (src1) { - GGML_ASSERT(ggml_is_contiguous_1(src1)); - GGML_ASSERT(src0->type == src1->type); + if (params->ith != 0) { + return; } - const int ith = params->ith; - const int nth = params->nth; - - const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; - const int nr = ggml_nrows(src0); - - GGML_ASSERT(dst->ne[0] == nc); - GGML_ASSERT(ggml_nrows(dst) == nr); + GGML_ASSERT(ggml_can_repeat(dst, src0)); - const int32_t swapped = ggml_get_op_params_i32(dst, 1); + GGML_TENSOR_UNARY_OP_LOCALS - // rows per thread - const int dr = (nr + nth - 1)/nth; + // guaranteed to be an integer due to the check in ggml_can_repeat + const int nr0 = (int)(ne00/ne0); + const int nr1 = (int)(ne01/ne1); + const int nr2 = (int)(ne02/ne2); + const int nr3 = (int)(ne03/ne3); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + // TODO: support for transposed / permuted tensors + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb00 == sizeof(float)); - for (int i1 = ir0; i1 < ir1; i1++) { - float * src0_p = (float *) (src0_d + i1*src0_o); - float * src1_p = (float *) (src1_d + i1*src1_o); - - if (!src1) { - src0_p += swapped ? nc : 0; - src1_p += swapped ? 0 : nc; - } - - ggml_vec_reglu_f32(nc, (float *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); - -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const float x = ((float *) ((char *) dst->data + i1*( dst->nb[1])))[k]; - GGML_UNUSED(x); - assert(!isnan(x)); - assert(!isinf(x)); + if (ggml_is_contiguous(dst)) { + ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0); + } else { + for (int k3 = 0; k3 < ne3; k3++) { + for (int k2 = 0; k2 < ne2; k2++) { + for (int k1 = 0; k1 < ne1; k1++) { + ggml_vec_set_f32(ne0, + (float *) ((char *) dst->data + k1*nb1 + k2*nb2 + k3*nb3), + 0); + } + } } -#endif // NDEBUG - } -} - -static void ggml_compute_forward_reglu_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - char * src0_d = (char *) src0->data; - char * src1_d = (char *) (src1 ? src1->data : src0->data); - const size_t src0_o = src0->nb[1]; - const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - - GGML_ASSERT(ggml_is_contiguous_1(src0)); - GGML_ASSERT(ggml_is_contiguous_1(dst)); - - if (src1) { - GGML_ASSERT(ggml_is_contiguous_1(src1)); - GGML_ASSERT(src0->type == src1->type); } - const int ith = params->ith; - const int nth = params->nth; - - const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; - const int nr = ggml_nrows(src0); - - GGML_ASSERT(dst->ne[0] == nc); - GGML_ASSERT(ggml_nrows(dst) == nr); - - const int32_t swapped = ggml_get_op_params_i32(dst, 1); - - // rows per thread - const int dr = (nr + nth - 1)/nth; - - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); - - for (int i1 = ir0; i1 < ir1; i1++) { - ggml_fp16_t * src0_p = (ggml_fp16_t *) (src0_d + i1*src0_o); - ggml_fp16_t * src1_p = (ggml_fp16_t *) (src1_d + i1*src1_o); - - if (!src1) { - src0_p += swapped ? nc : 0; - src1_p += swapped ? 0 : nc; - } - - ggml_vec_reglu_f16(nc, (ggml_fp16_t *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); - -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])))[k]; - const float v = GGML_FP16_TO_FP32(x); - GGML_UNUSED(v); - assert(!isnan(v)); - assert(!isinf(v)); + // TODO: maybe this is not optimal? + for (int i3 = 0; i3 < nr3; i3++) { + for (int k3 = 0; k3 < ne3; k3++) { + for (int i2 = 0; i2 < nr2; i2++) { + for (int k2 = 0; k2 < ne2; k2++) { + for (int i1 = 0; i1 < nr1; i1++) { + for (int k1 = 0; k1 < ne1; k1++) { + for (int i0 = 0; i0 < nr0; i0++) { + ggml_vec_acc_f32(ne0, + (float *) ((char *) dst->data + ( k3)*nb3 + ( k2)*nb2 + ( k1)*nb1), + (float *) ((char *) src0->data + (i3*ne3 + k3)*nb03 + (i2*ne2 + k2)*nb02 + (i1*ne1 + k1)*nb01 + (i0*ne0)*nb00)); + } + } + } + } + } } -#endif // NDEBUG } } -static void ggml_compute_forward_reglu( +void ggml_compute_forward_repeat_back( const ggml_compute_params * params, ggml_tensor * dst) { @@ -2972,11 +3956,7 @@ static void ggml_compute_forward_reglu( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_reglu_f32(params, dst); - } break; - case GGML_TYPE_F16: - { - ggml_compute_forward_reglu_f16(params, dst); + ggml_compute_forward_repeat_back_f32(params, dst); } break; default: { @@ -2985,326 +3965,231 @@ static void ggml_compute_forward_reglu( } } -// ggml_compute_forward_geglu +// ggml_compute_forward_concat -static void ggml_compute_forward_geglu_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { +static void ggml_compute_forward_concat_any( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - char * src0_d = (char *) src0->data; - char * src1_d = (char *) (src1 ? src1->data : src0->data); - const size_t src0_o = src0->nb[1]; - const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - - GGML_ASSERT(ggml_is_contiguous_1(src0)); - GGML_ASSERT(ggml_is_contiguous_1(dst)); - if (src1) { - GGML_ASSERT(ggml_is_contiguous_1(src1)); - GGML_ASSERT(src0->type == src1->type); - } + const size_t len = ggml_type_size(src0->type); const int ith = params->ith; const int nth = params->nth; - const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; - const int nr = ggml_nrows(src0); - - GGML_ASSERT(dst->ne[0] == nc); - GGML_ASSERT(ggml_nrows(dst) == nr); + GGML_TENSOR_BINARY_OP_LOCALS - const int32_t swapped = ggml_get_op_params_i32(dst, 1); + const int32_t dim = ggml_get_op_params_i32(dst, 0); - // rows per thread - const int dr = (nr + nth - 1)/nth; + GGML_ASSERT(dim >= 0 && dim < 4); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + int64_t o[4] = {0, 0, 0, 0}; + o[dim] = src0->ne[dim]; - for (int i1 = ir0; i1 < ir1; i1++) { - float * src0_p = (float *) (src0_d + i1*src0_o); - float * src1_p = (float *) (src1_d + i1*src1_o); + const char * x; - if (!src1) { - src0_p += swapped ? nc : 0; - src1_p += swapped ? 0 : nc; - } + // TODO: smarter multi-theading + for (int i3 = 0; i3 < ne3; i3++) { + for (int i2 = ith; i2 < ne2; i2 += nth) { + for (int i1 = 0; i1 < ne1; i1++) { + for (int i0 = 0; i0 < ne0; i0++) { + if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { + x = (const char *)src0->data + (i0 )*nb00 + (i1 )*nb01 + (i2 )*nb02 + (i3 )*nb03; + } else { + x = (const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13; + } - ggml_vec_geglu_f32(nc, (float *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + char * y = (char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3; -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const float x = ((float *) ((char *) dst->data + i1*( dst->nb[1])))[k]; - GGML_UNUSED(x); - assert(!isnan(x)); - assert(!isinf(x)); + memcpy(y, x, len); + } + } } -#endif // NDEBUG } } -static void ggml_compute_forward_geglu_f16( +static void ggml_compute_forward_concat_i8( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - char * src0_d = (char *) src0->data; - char * src1_d = (char *) (src1 ? src1->data : src0->data); - const size_t src0_o = src0->nb[1]; - const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - - GGML_ASSERT(ggml_is_contiguous_1(src0)); - GGML_ASSERT(ggml_is_contiguous_1(dst)); - if (src1) { - GGML_ASSERT(ggml_is_contiguous_1(src1)); - GGML_ASSERT(src0->type == src1->type); - } + GGML_ASSERT(ggml_type_size(src0->type) == sizeof(int8_t)); const int ith = params->ith; const int nth = params->nth; - const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; - const int nr = ggml_nrows(src0); - - GGML_ASSERT(dst->ne[0] == nc); - GGML_ASSERT(ggml_nrows(dst) == nr); + GGML_TENSOR_BINARY_OP_LOCALS - const int32_t swapped = ggml_get_op_params_i32(dst, 1); + const int32_t dim = ggml_get_op_params_i32(dst, 0); - // rows per thread - const int dr = (nr + nth - 1)/nth; + GGML_ASSERT(dim >= 0 && dim < 4); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + int64_t o[4] = {0, 0, 0, 0}; + o[dim] = src0->ne[dim]; - for (int i1 = ir0; i1 < ir1; i1++) { - ggml_fp16_t * src0_p = (ggml_fp16_t *) (src0_d + i1*src0_o); - ggml_fp16_t * src1_p = (ggml_fp16_t *) (src1_d + i1*src1_o); + const int8_t * x; - if (!src1) { - src0_p += swapped ? nc : 0; - src1_p += swapped ? 0 : nc; - } + // TODO: smarter multi-theading + for (int i3 = 0; i3 < ne3; i3++) { + for (int i2 = ith; i2 < ne2; i2 += nth) { + for (int i1 = 0; i1 < ne1; i1++) { + for (int i0 = 0; i0 < ne0; i0++) { + if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { + x = (const int8_t *) ((const char *)src0->data + (i0 )*nb00 + (i1 )*nb01 + (i2 )*nb02 + (i3 )*nb03); + } else { + x = (const int8_t *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13); + } - ggml_vec_geglu_f16(nc, (ggml_fp16_t *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + int8_t * y = (int8_t *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])))[k]; - const float v = GGML_FP16_TO_FP32(x); - GGML_UNUSED(v); - assert(!isnan(v)); - assert(!isinf(v)); + *y = *x; + } + } } -#endif // NDEBUG } } -static void ggml_compute_forward_geglu( - const ggml_compute_params * params, - ggml_tensor * dst) { +static void ggml_compute_forward_concat_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_geglu_f32(params, dst); - } break; - case GGML_TYPE_F16: - { - ggml_compute_forward_geglu_f16(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} - -// ggml_compute_forward_swiglu - -static void ggml_compute_forward_swiglu_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - char * src0_d = (char *) src0->data; - char * src1_d = (char *) (src1 ? src1->data : src0->data); - const size_t src0_o = src0->nb[1]; - const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - - GGML_ASSERT(ggml_is_contiguous_1(src0)); - GGML_ASSERT(ggml_is_contiguous_1(dst)); - - if (src1) { - GGML_ASSERT(ggml_is_contiguous_1(src1)); - GGML_ASSERT(src0->type == src1->type); - } + GGML_ASSERT(ggml_type_size(src0->type) == sizeof(ggml_fp16_t)); const int ith = params->ith; const int nth = params->nth; - const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; - const int nr = ggml_nrows(src0); - - GGML_ASSERT(dst->ne[0] == nc); - GGML_ASSERT(ggml_nrows(dst) == nr); + GGML_TENSOR_BINARY_OP_LOCALS - const int32_t swapped = ggml_get_op_params_i32(dst, 1); + const int32_t dim = ggml_get_op_params_i32(dst, 0); - // rows per thread - const int dr = (nr + nth - 1)/nth; + GGML_ASSERT(dim >= 0 && dim < 4); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + int64_t o[4] = {0, 0, 0, 0}; + o[dim] = src0->ne[dim]; - for (int i1 = ir0; i1 < ir1; i1++) { - float * src0_p = (float *) (src0_d + i1*src0_o); - float * src1_p = (float *) (src1_d + i1*src1_o); + const ggml_fp16_t * x; - if (!src1) { - src0_p += swapped ? nc : 0; - src1_p += swapped ? 0 : nc; - } + // TODO: smarter multi-theading + for (int i3 = 0; i3 < ne3; i3++) { + for (int i2 = ith; i2 < ne2; i2 += nth) { + for (int i1 = 0; i1 < ne1; i1++) { + for (int i0 = 0; i0 < ne0; i0++) { + if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { + x = (const ggml_fp16_t *) ((const char *)src0->data + (i0 )*nb00 + (i1 )*nb01 + (i2 )*nb02 + (i3 )*nb03); + } else { + x = (const ggml_fp16_t *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13); + } - ggml_vec_swiglu_f32(nc, (float *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + ggml_fp16_t * y = (ggml_fp16_t *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const float x = ((float *) ((char *) dst->data + i1*( dst->nb[1])))[k]; - GGML_UNUSED(x); - assert(!isnan(x)); - assert(!isinf(x)); + *y = *x; + } + } } -#endif // NDEBUG } } -static void ggml_compute_forward_swiglu_f16( +static void ggml_compute_forward_concat_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - char * src0_d = (char *) src0->data; - char * src1_d = (char *) (src1 ? src1->data : src0->data); - const size_t src0_o = src0->nb[1]; - const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - - GGML_ASSERT(ggml_is_contiguous_1(src0)); - GGML_ASSERT(ggml_is_contiguous_1(dst)); - if (src1) { - GGML_ASSERT(ggml_is_contiguous_1(src1)); - GGML_ASSERT(src0->type == src1->type); - } + GGML_ASSERT(ggml_type_size(src0->type) == sizeof(float)); const int ith = params->ith; const int nth = params->nth; - const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; - const int nr = ggml_nrows(src0); - - GGML_ASSERT(dst->ne[0] == nc); - GGML_ASSERT(ggml_nrows(dst) == nr); + GGML_TENSOR_BINARY_OP_LOCALS - const int32_t swapped = ggml_get_op_params_i32(dst, 1); + const int32_t dim = ggml_get_op_params_i32(dst, 0); - // rows per thread - const int dr = (nr + nth - 1)/nth; + GGML_ASSERT(dim >= 0 && dim < 4); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + int64_t o[4] = {0, 0, 0, 0}; + o[dim] = src0->ne[dim]; - for (int i1 = ir0; i1 < ir1; i1++) { - ggml_fp16_t * src0_p = (ggml_fp16_t *) (src0_d + i1*src0_o); - ggml_fp16_t * src1_p = (ggml_fp16_t *) (src1_d + i1*src1_o); + const float * x; - if (!src1) { - src0_p += swapped ? nc : 0; - src1_p += swapped ? 0 : nc; - } + // TODO: smarter multi-theading + for (int i3 = 0; i3 < ne3; i3++) { + for (int i2 = ith; i2 < ne2; i2 += nth) { + for (int i1 = 0; i1 < ne1; i1++) { + for (int i0 = 0; i0 < ne0; i0++) { + if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { + x = (const float *) ((const char *)src0->data + (i0 )*nb00 + (i1 )*nb01 + (i2 )*nb02 + (i3 )*nb03); + } else { + x = (const float *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13); + } - ggml_vec_swiglu_f16(nc, (ggml_fp16_t *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + float * y = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])))[k]; - const float v = GGML_FP16_TO_FP32(x); - GGML_UNUSED(v); - assert(!isnan(v)); - assert(!isinf(v)); + *y = *x; + } + } } -#endif // NDEBUG } } -static void ggml_compute_forward_swiglu( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_concat( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { - case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_BF16: + case GGML_TYPE_I16: { - ggml_compute_forward_swiglu_f32(params, dst); + ggml_compute_forward_concat_f16(params, dst); } break; - case GGML_TYPE_F16: + case GGML_TYPE_I8: { - ggml_compute_forward_swiglu_f16(params, dst); + ggml_compute_forward_concat_i8(params, dst); + } break; + case GGML_TYPE_F32: + case GGML_TYPE_I32: + { + ggml_compute_forward_concat_f32(params, dst); } break; default: { - GGML_ABORT("fatal error"); + ggml_compute_forward_concat_any(params, dst); } } } -// ggml_compute_forward_swiglu_oai +// ggml_compute_forward_gelu -static void ggml_compute_forward_swiglu_oai_f32( +static void ggml_compute_forward_gelu_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - char * src0_d = (char *) src0->data; - char * src1_d = (char *) (src1 ? src1->data : src0->data); - const size_t src0_o = src0->nb[1]; - const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - GGML_ASSERT(ggml_is_contiguous_1(src0)); - GGML_ASSERT(ggml_is_contiguous_1(dst)); + assert(ggml_is_contiguous_rows(src0)); + assert(ggml_are_same_shape(src0, dst)); - if (src1) { - GGML_ASSERT(ggml_is_contiguous_1(src1)); - GGML_ASSERT(src0->type == src1->type); - } + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) const int ith = params->ith; const int nth = params->nth; - const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nc = src0->ne[0]; const int nr = ggml_nrows(src0); - GGML_ASSERT(dst->ne[0] == nc); - GGML_ASSERT(ggml_nrows(dst) == nr); - - const int32_t swapped = ggml_get_op_params_i32(dst, 1); - const float alpha = ggml_get_op_params_f32(dst, 2); - const float limit = ggml_get_op_params_f32(dst, 3); - // rows per thread const int dr = (nr + nth - 1)/nth; @@ -3312,26 +4197,18 @@ static void ggml_compute_forward_swiglu_oai_f32( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - for (int i1 = ir0; i1 < ir1; i1++) { - float * src0_p = (float *) (src0_d + i1*src0_o); - float * src1_p = (float *) (src1_d + i1*src1_o); - float * dst_p = (float *) ((char *) dst->data + i1*(dst->nb[1])); - - if (!src1) { - src0_p += swapped ? nc : 0; - src1_p += swapped ? 0 : nc; - } + for (int ir = ir0; ir < ir1; ++ir) { + const int i3 = ir/(ne02*ne01); + const int i2 = (ir - i3*ne02*ne01)/ne01; + const int i1 = (ir - i3*ne02*ne01 - i2*ne01); - for (int k = 0; k < nc; k++) { - const float x = std::min(src0_p[k], limit); - const float y = std::clamp(src1_p[k], -limit, limit); - const float out_glu = x / (1.f + expf(alpha * (-x))); - dst_p[k] = out_glu * (y + 1.f); - } + ggml_vec_gelu_f32(nc, + (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), + (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); #ifndef NDEBUG for (int k = 0; k < nc; k++) { - const float x = dst_p[k]; + const float x = ((float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*(dst->nb[1])))[k]; GGML_UNUSED(x); assert(!isnan(x)); assert(!isinf(x)); @@ -3340,56 +4217,26 @@ static void ggml_compute_forward_swiglu_oai_f32( } } -static void ggml_compute_forward_swiglu_oai( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_swiglu_oai_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} - -// ggml_compute_forward_geglu_erf - -static void ggml_compute_forward_geglu_erf_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { +static void ggml_compute_forward_gelu_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - char * src0_d = (char *) src0->data; - char * src1_d = (char *) (src1 ? src1->data : src0->data); - const size_t src0_o = src0->nb[1]; - const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - GGML_ASSERT(ggml_is_contiguous_1(src0)); - GGML_ASSERT(ggml_is_contiguous_1(dst)); + assert(ggml_is_contiguous_rows(src0)); + assert(ggml_are_same_shape(src0, dst)); - if (src1) { - GGML_ASSERT(ggml_is_contiguous_1(src1)); - GGML_ASSERT(src0->type == src1->type); - } + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) const int ith = params->ith; const int nth = params->nth; - const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nc = src0->ne[0]; const int nr = ggml_nrows(src0); - GGML_ASSERT(dst->ne[0] == nc); - GGML_ASSERT(ggml_nrows(dst) == nr); - - const int32_t swapped = ggml_get_op_params_i32(dst, 1); - // rows per thread const int dr = (nr + nth - 1)/nth; @@ -3397,102 +4244,118 @@ static void ggml_compute_forward_geglu_erf_f32( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - for (int i1 = ir0; i1 < ir1; i1++) { - float * src0_p = (float *) (src0_d + i1*src0_o); - float * src1_p = (float *) (src1_d + i1*src1_o); - - if (!src1) { - src0_p += swapped ? nc : 0; - src1_p += swapped ? 0 : nc; - } + for (int ir = ir0; ir < ir1; ++ir) { + const int i3 = ir/(ne02*ne01); + const int i2 = (ir - i3*ne02*ne01)/ne01; + const int i1 = (ir - i3*ne02*ne01 - i2*ne01); - ggml_vec_geglu_erf_f32(nc, (float *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + ggml_vec_gelu_f16(nc, + (ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), + (ggml_fp16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); #ifndef NDEBUG for (int k = 0; k < nc; k++) { - const float x = ((float *) ((char *) dst->data + i1*( dst->nb[1])))[k]; - GGML_UNUSED(x); - assert(!isnan(x)); - assert(!isinf(x)); + const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*( dst->nb[1])))[k]; + const float v = GGML_CPU_FP16_TO_FP32(x); + GGML_UNUSED(v); + assert(!isnan(v)); + assert(!isinf(v)); } #endif // NDEBUG } } -static void ggml_compute_forward_geglu_erf_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { +static void ggml_compute_forward_gelu( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - char * src0_d = (char *) src0->data; - char * src1_d = (char *) (src1 ? src1->data : src0->data); - const size_t src0_o = src0->nb[1]; - const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - GGML_ASSERT(ggml_is_contiguous_1(src0)); - GGML_ASSERT(ggml_is_contiguous_1(dst)); + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_gelu_f32(params, dst); + } break; + case GGML_TYPE_F16: + { + ggml_compute_forward_gelu_f16(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - if (src1) { - GGML_ASSERT(ggml_is_contiguous_1(src1)); - GGML_ASSERT(src0->type == src1->type); +// ggml_compute_fill + +static void ggml_compute_forward_fill_f32(const ggml_compute_params * params, ggml_tensor * dst) { + const float c = ggml_get_op_params_f32(dst, 0); + + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne); + GGML_TENSOR_LOCALS(size_t, nb, dst, nb); + + const auto [ir0, ir1] = get_thread_range(params, dst); + + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t i03 = ir/(ne2*ne1); + const int64_t i02 = (ir - i03*ne2*ne1)/ne1; + const int64_t i01 = (ir - i03*ne2*ne1 - i02*ne1); + + float * dst_ptr = (float *) ((char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1); + + ggml_vec_set_f32(ne0, dst_ptr, c); } +} - const int ith = params->ith; - const int nth = params->nth; +void ggml_compute_forward_fill(const ggml_compute_params * params, ggml_tensor * dst) { + ggml_compute_forward_fill_f32(params, dst); +} - const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; - const int nr = ggml_nrows(src0); +// ggml_compute_tri - GGML_ASSERT(dst->ne[0] == nc); - GGML_ASSERT(ggml_nrows(dst) == nr); +static void ggml_compute_forward_tri_f32(const ggml_compute_params * params, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; - const int32_t swapped = ggml_get_op_params_i32(dst, 1); + const ggml_tri_type ttype = (ggml_tri_type) ggml_get_op_params_i32(dst, 0); - // rows per thread - const int dr = (nr + nth - 1)/nth; + GGML_ASSERT(ggml_is_contiguous(src0)); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + GGML_TENSOR_UNARY_OP_LOCALS - for (int i1 = ir0; i1 < ir1; i1++) { - ggml_fp16_t * src0_p = (ggml_fp16_t *) (src0_d + i1*src0_o); - ggml_fp16_t * src1_p = (ggml_fp16_t *) (src1_d + i1*src1_o); + const auto [ir0, ir1] = get_thread_range(params, src0); - if (!src1) { - src0_p += swapped ? nc : 0; - src1_p += swapped ? 0 : nc; - } + bool (*bipred)(int, int); - ggml_vec_geglu_erf_f16(nc, (ggml_fp16_t *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + switch (ttype) { + case GGML_TRI_TYPE_LOWER: bipred = [](int i, int r) { return i < r; }; break; + case GGML_TRI_TYPE_LOWER_DIAG: bipred = [](int i, int r) { return i <= r; }; break; + case GGML_TRI_TYPE_UPPER: bipred = [](int i, int r) { return i > r; }; break; + case GGML_TRI_TYPE_UPPER_DIAG: bipred = [](int i, int r) { return i >= r; }; break; + default: GGML_ABORT("invalid tri type"); + } -#ifndef NDEBUG - for (int k = 0; k < nc; k++) { - const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])))[k]; - const float v = GGML_FP16_TO_FP32(x); - GGML_UNUSED(v); - assert(!isnan(v)); - assert(!isinf(v)); + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t i03 = ir/(ne02*ne01); + const int64_t i02 = (ir - i03*ne02*ne01)/ne01; + const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); + + const float * src_ptr = (const float *) ((const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01); + float * dst_ptr = ( float *) (( char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1); + + for (int i0 = 0; i0 < ne0; ++i0) { + dst_ptr[i0] = bipred(i0, i01) ? src_ptr[i0] : 0.0f; } -#endif // NDEBUG } } -static void ggml_compute_forward_geglu_erf( - const ggml_compute_params * params, - ggml_tensor * dst) { - +void ggml_compute_forward_tri(const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_geglu_erf_f32(params, dst); - } break; - case GGML_TYPE_F16: - { - ggml_compute_forward_geglu_erf_f16(params, dst); + ggml_compute_forward_tri_f32(params, dst); } break; default: { @@ -3501,38 +4364,28 @@ static void ggml_compute_forward_geglu_erf( } } -// ggml_compute_forward_geglu_quick +// ggml_compute_forward_gelu_erf -static void ggml_compute_forward_geglu_quick_f32( +static void ggml_compute_forward_gelu_erf_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - char * src0_d = (char *) src0->data; - char * src1_d = (char *) (src1 ? src1->data : src0->data); - const size_t src0_o = src0->nb[1]; - const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - GGML_ASSERT(ggml_is_contiguous_1(src0)); - GGML_ASSERT(ggml_is_contiguous_1(dst)); + assert(ggml_is_contiguous_rows(src0)); + assert(ggml_are_same_shape(src0, dst)); - if (src1) { - GGML_ASSERT(ggml_is_contiguous_1(src1)); - GGML_ASSERT(src0->type == src1->type); - } + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) const int ith = params->ith; const int nth = params->nth; - const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nc = src0->ne[0]; const int nr = ggml_nrows(src0); - GGML_ASSERT(dst->ne[0] == nc); - GGML_ASSERT(ggml_nrows(dst) == nr); - - const int32_t swapped = ggml_get_op_params_i32(dst, 1); - // rows per thread const int dr = (nr + nth - 1)/nth; @@ -3540,20 +4393,18 @@ static void ggml_compute_forward_geglu_quick_f32( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - for (int i1 = ir0; i1 < ir1; i1++) { - float * src0_p = (float *) (src0_d + i1*src0_o); - float * src1_p = (float *) (src1_d + i1*src1_o); - - if (!src1) { - src0_p += swapped ? nc : 0; - src1_p += swapped ? 0 : nc; - } + for (int ir = ir0; ir < ir1; ++ir) { + const int i3 = ir/(ne02*ne01); + const int i2 = (ir - i3*ne02*ne01)/ne01; + const int i1 = (ir - i3*ne02*ne01 - i2*ne01); - ggml_vec_geglu_quick_f32(nc, (float *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + ggml_vec_gelu_erf_f32(nc, + (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), + (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); #ifndef NDEBUG for (int k = 0; k < nc; k++) { - const float x = ((float *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + const float x = ((float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*(dst->nb[1])))[k]; GGML_UNUSED(x); assert(!isnan(x)); assert(!isinf(x)); @@ -3562,36 +4413,26 @@ static void ggml_compute_forward_geglu_quick_f32( } } -static void ggml_compute_forward_geglu_quick_f16( +static void ggml_compute_forward_gelu_erf_f16( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - char * src0_d = (char *) src0->data; - char * src1_d = (char *) (src1 ? src1->data : src0->data); - const size_t src0_o = src0->nb[1]; - const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - GGML_ASSERT(ggml_is_contiguous_1(src0)); - GGML_ASSERT(ggml_is_contiguous_1(dst)); + assert(ggml_is_contiguous_rows(src0)); + assert(ggml_are_same_shape(src0, dst)); - if (src1) { - GGML_ASSERT(ggml_is_contiguous_1(src1)); - GGML_ASSERT(src0->type == src1->type); - } + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) const int ith = params->ith; const int nth = params->nth; - const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nc = src0->ne[0]; const int nr = ggml_nrows(src0); - GGML_ASSERT(dst->ne[0] == nc); - GGML_ASSERT(ggml_nrows(dst) == nr); - - const int32_t swapped = ggml_get_op_params_i32(dst, 1); - // rows per thread const int dr = (nr + nth - 1)/nth; @@ -3599,21 +4440,19 @@ static void ggml_compute_forward_geglu_quick_f16( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - for (int i1 = ir0; i1 < ir1; i1++) { - ggml_fp16_t * src0_p = (ggml_fp16_t *) (src0_d + i1*src0_o); - ggml_fp16_t * src1_p = (ggml_fp16_t *) (src1_d + i1*src1_o); - - if (!src1) { - src0_p += swapped ? nc : 0; - src1_p += swapped ? 0 : nc; - } + for (int ir = ir0; ir < ir1; ++ir) { + const int i3 = ir/(ne02*ne01); + const int i2 = (ir - i3*ne02*ne01)/ne01; + const int i1 = (ir - i3*ne02*ne01 - i2*ne01); - ggml_vec_geglu_quick_f16(nc, (ggml_fp16_t *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + ggml_vec_gelu_erf_f16(nc, + (ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), + (ggml_fp16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); #ifndef NDEBUG for (int k = 0; k < nc; k++) { - const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])))[k]; - const float v = GGML_FP16_TO_FP32(x); + const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*( dst->nb[1])))[k]; + const float v = GGML_CPU_FP16_TO_FP32(x); GGML_UNUSED(v); assert(!isnan(v)); assert(!isinf(v)); @@ -3622,7 +4461,7 @@ static void ggml_compute_forward_geglu_quick_f16( } } -static void ggml_compute_forward_geglu_quick( +static void ggml_compute_forward_gelu_erf( const ggml_compute_params * params, ggml_tensor * dst) { @@ -3631,11 +4470,11 @@ static void ggml_compute_forward_geglu_quick( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_geglu_quick_f32(params, dst); + ggml_compute_forward_gelu_erf_f32(params, dst); } break; case GGML_TYPE_F16: { - ggml_compute_forward_geglu_quick_f16(params, dst); + ggml_compute_forward_gelu_erf_f16(params, dst); } break; default: { @@ -3644,56 +4483,104 @@ static void ggml_compute_forward_geglu_quick( } } -// ggml_compute_forward_norm +// ggml_compute_forward_gelu_quick -static void ggml_compute_forward_norm_f32( +static void ggml_compute_forward_gelu_quick_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - GGML_ASSERT(ggml_are_same_shape(src0, dst)); + assert(ggml_is_contiguous_rows(src0)); + assert(ggml_are_same_shape(src0, dst)); - GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) const int ith = params->ith; const int nth = params->nth; - GGML_TENSOR_UNARY_OP_LOCALS + const int nc = src0->ne[0]; + const int nr = ggml_nrows(src0); - float eps; - memcpy(&eps, dst->op_params, sizeof(float)); + // rows per thread + const int dr = (nr + nth - 1)/nth; - GGML_ASSERT(eps >= 0.0f); + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = ith; i01 < ne01; i01 += nth) { - const float * x = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); + for (int ir = ir0; ir < ir1; ++ir) { + const int i3 = ir/(ne02*ne01); + const int i2 = (ir - i3*ne02*ne01)/ne01; + const int i1 = (ir - i3*ne02*ne01 - i2*ne01); - float sum = 0.0; - ggml_vec_sum_f32(ne00, &sum, x); - float mean = sum/ne00; + ggml_vec_gelu_quick_f32(nc, + (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), + (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); - float * y = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); - float variance = 0; +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = ((float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*(dst->nb[1])))[k]; + GGML_UNUSED(x); + assert(!isnan(x)); + assert(!isinf(x)); + } +#endif // NDEBUG + } +} -#ifdef GGML_USE_ACCELERATE - mean = -mean; - vDSP_vsadd(x, 1, &mean, y, 1, ne00); - vDSP_measqv(y, 1, &variance, ne00); -#else - variance = ggml_vec_cvar_f32(ne00, y, x, mean); -#endif //GGML_USE_ACCELERATE +static void ggml_compute_forward_gelu_quick_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { - const float scale = 1.0f/sqrtf(variance + eps); - ggml_vec_scale_f32(ne00, y, scale); - } + const ggml_tensor * src0 = dst->src[0]; + + assert(ggml_is_contiguous_rows(src0)); + assert(ggml_are_same_shape(src0, dst)); + + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + + const int ith = params->ith; + const int nth = params->nth; + + const int nc = src0->ne[0]; + const int nr = ggml_nrows(src0); + + // rows per thread + const int dr = (nr + nth - 1)/nth; + + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + + for (int ir = ir0; ir < ir1; ++ir) { + const int i3 = ir/(ne02*ne01); + const int i2 = (ir - i3*ne02*ne01)/ne01; + const int i1 = (ir - i3*ne02*ne01 - i2*ne01); + + ggml_vec_gelu_quick_f16(nc, + (ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), + (ggml_fp16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); + +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*( dst->nb[1])))[k]; + const float v = GGML_CPU_FP16_TO_FP32(x); + GGML_UNUSED(v); + assert(!isnan(v)); + assert(!isinf(v)); } +#endif // NDEBUG } } -void ggml_compute_forward_norm( +static void ggml_compute_forward_gelu_quick( const ggml_compute_params * params, ggml_tensor * dst) { @@ -3702,7 +4589,11 @@ void ggml_compute_forward_norm( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_norm_f32(params, dst); + ggml_compute_forward_gelu_quick_f32(params, dst); + } break; + case GGML_TYPE_F16: + { + ggml_compute_forward_gelu_quick_f16(params, dst); } break; default: { @@ -3711,60 +4602,104 @@ void ggml_compute_forward_norm( } } -// ggml_compute_forward_group_rms_norm +// ggml_compute_forward_silu -static void ggml_compute_forward_rms_norm_f32( +static void ggml_compute_forward_silu_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - GGML_ASSERT(ggml_are_same_shape(src0, dst)); + assert(ggml_is_contiguous_rows(src0)); + assert(ggml_are_same_shape(src0, dst)); - GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) const int ith = params->ith; const int nth = params->nth; - GGML_TENSOR_UNARY_OP_LOCALS + const int nc = src0->ne[0]; + const int nr = ggml_nrows(src0); - float eps; - memcpy(&eps, dst->op_params, sizeof(float)); + // rows per thread + const int dr = (nr + nth - 1)/nth; - GGML_ASSERT(eps >= 0.0f); + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - // TODO: optimize - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = ith; i01 < ne01; i01 += nth) { - const float * x = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); + for (int ir = ir0; ir < ir1; ++ir) { + const int i3 = ir/(ne02*ne01); + const int i2 = (ir - i3*ne02*ne01)/ne01; + const int i1 = (ir - i3*ne02*ne01 - i2*ne01); - ggml_float sum = 0.0; - for (int64_t i00 = 0; i00 < ne00; i00++) { - sum += (ggml_float)(x[i00] * x[i00]); - } + ggml_vec_silu_f32(nc, + (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), + (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); - const float mean = sum/ne00; +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = ((float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*(dst->nb[1])))[k]; + GGML_UNUSED(x); + assert(!isnan(x)); + assert(!isinf(x)); + } +#endif // NDEBUG + } +} - float * y = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); +static void ggml_compute_forward_silu_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { - memcpy(y, x, ne00 * sizeof(float)); - // for (int i00 = 0; i00 < ne00; i00++) { - // y[i00] = x[i00]; - // } + const ggml_tensor * src0 = dst->src[0]; - const float scale = 1.0f/sqrtf(mean + eps); + assert(ggml_is_contiguous_rows(src0)); + assert(ggml_are_same_shape(src0, dst)); - // if you hit this, likely you got an inf somewhere earlier - assert(scale > 0.0f); + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) - ggml_vec_scale_f32(ne00, y, scale); - } + const int ith = params->ith; + const int nth = params->nth; + + const int nc = src0->ne[0]; + const int nr = ggml_nrows(src0); + + // rows per thread + const int dr = (nr + nth - 1)/nth; + + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + + for (int ir = ir0; ir < ir1; ++ir) { + const int i3 = ir/(ne02*ne01); + const int i2 = (ir - i3*ne02*ne01)/ne01; + const int i1 = (ir - i3*ne02*ne01 - i2*ne01); + + ggml_vec_silu_f16(nc, + (ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1), + (ggml_fp16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); + +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*( dst->nb[1])))[k]; + const float v = GGML_CPU_FP16_TO_FP32(x); + GGML_UNUSED(v); + assert(!isnan(v)); + assert(!isinf(v)); } +#endif // NDEBUG } } -void ggml_compute_forward_rms_norm( +static void ggml_compute_forward_silu( const ggml_compute_params * params, ggml_tensor * dst) { @@ -3773,7 +4708,11 @@ void ggml_compute_forward_rms_norm( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_rms_norm_f32(params, dst); + ggml_compute_forward_silu_f32(params, dst); + } break; + case GGML_TYPE_F16: + { + ggml_compute_forward_silu_f16(params, dst); } break; default: { @@ -3781,165 +4720,69 @@ void ggml_compute_forward_rms_norm( } } } +// ggml_compute_forward_leaky_relu -static void ggml_compute_forward_rms_norm_back_f32( +static void ggml_compute_forward_leaky_relu_f32( const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; // gradients from forward pass output - const ggml_tensor * src1 = dst->src[1]; // src1 from forward pass - - GGML_ASSERT(ggml_are_same_shape(src0, dst) && ggml_are_same_shape(src0, src1)); + const ggml_tensor * src0 = dst->src[0]; - GGML_ASSERT(src0->nb[0] == sizeof(float)); - GGML_ASSERT(src1->nb[0] == sizeof(float)); + if (params->ith != 0) { + return; + } - const int ith = params->ith; - const int nth = params->nth; + assert(ggml_is_contiguous_1(src0)); + assert(ggml_is_contiguous_1(dst)); + assert(ggml_are_same_shape(src0, dst)); - GGML_TENSOR_BINARY_OP_LOCALS + const int n = ggml_nrows(src0); + const int nc = src0->ne[0]; - float eps; - memcpy(&eps, dst->op_params, sizeof(float)); + float negative_slope; + memcpy(&negative_slope, dst->op_params, sizeof(float)); - // TODO: optimize - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = ith; i01 < ne01; i01 += nth) { - // src1 is same shape as src0 => same indices - const int64_t i11 = i01; - const int64_t i12 = i02; - const int64_t i13 = i03; + assert(dst->nb[0] == sizeof(float)); + assert(src0->nb[0] == sizeof(float)); - const float * dz = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); - const float * x = (float *) ((char *) src1->data + i11*nb11 + i12*nb12 + i13*nb13); + for (int i = 0; i < n; i++) { + ggml_vec_leaky_relu_f32(nc, + (float *) ((char *) dst->data + i*( dst->nb[1])), + (float *) ((char *) src0->data + i*(src0->nb[1])), negative_slope); + } +} - ggml_float sum_xx = 0.0; - ggml_float sum_xdz = 0.0; +static void ggml_compute_forward_leaky_relu_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { - for (int64_t i00 = 0; i00 < ne00; i00++) { - sum_xx += (ggml_float)(x[i00] * x[i00]); - sum_xdz += (ggml_float)(x[i00] * dz[i00]); - } + const ggml_tensor * src0 = dst->src[0]; - //const float mean = (float)(sum_xx)/ne00; - const float mean_eps = (float)(sum_xx)/ne00 + eps; - const float sum_eps = (float)(sum_xx) + eps*ne00; - //const float mean_xdz = (float)(sum_xdz)/ne00; - // we could cache rms from forward pass to improve performance. - // to do this implement ggml_rms and compose ggml_rms_norm using ggml_rms. - //const float rms = sqrtf(mean_eps); - const float rrms = 1.0f / sqrtf(mean_eps); - //const float scale = -rrms/(ne00 * mean_eps); // -1/(n*rms**3) + if (params->ith != 0) { + return; + } - { - // z = rms_norm(x) - // - // rms_norm(src1) = - // scale( - // src1, - // div( - // 1, - // sqrt( - // add( - // scale( - // sum( - // sqr( - // src1)), - // (1.0/N)), - // eps)))); + assert(ggml_is_contiguous_1(src0)); + assert(ggml_is_contiguous_1(dst)); + assert(ggml_are_same_shape(src0, dst)); - // postorder: - // ## op args grad - // 00 param src1 grad[#00] - // 01 const 1 - // 02 sqr (#00) grad[#02] - // 03 sum (#02) grad[#03] - // 04 const 1/N - // 05 scale (#03, #04) grad[#05] - // 06 const eps - // 07 add (#05, #06) grad[#07] - // 08 sqrt (#07) grad[#08] - // 09 div (#01,#08) grad[#09] - // 10 scale (#00,#09) grad[#10] - // - // backward pass, given grad[#10] - // #10: scale - // grad[#00] += scale(grad[#10],#09) - // grad[#09] += sum(mul(grad[#10],#00)) - // #09: div - // grad[#08] += neg(mul(grad[#09], div(#09,#08))) - // #08: sqrt - // grad[#07] += mul(grad[#08], div(0.5, #08)) - // #07: add - // grad[#05] += grad[#07] - // #05: scale - // grad[#03] += scale(grad[#05],#04) - // #03: sum - // grad[#02] += repeat(grad[#03], #02) - // #02: - // grad[#00] += scale(mul(#00, grad[#02]), 2.0) - // - // substitute and simplify: - // grad[#00] = scale(grad(#10), #09) + scale(mul(#00, grad[#02]), 2.0) - // grad[#02] = repeat(grad[#03], #02) - // grad[#02] = repeat(scale(grad[#05],#04), #02) - // grad[#02] = repeat(scale(grad[#07],#04), #02) - // grad[#02] = repeat(scale(mul(grad[#08], div(0.5, #08)),#04), #02) - // grad[#02] = repeat(scale(mul(neg(mul(grad[#09], div(#09,#08))), div(0.5, #08)),#04), #02) - // grad[#02] = repeat(scale(mul(neg(mul(sum(mul(grad[#10],#00)), div(#09,#08))), div(0.5, #08)),#04), #02) - // grad[#02] = repeat(-(sum(mul(grad[#10],#00)) * div(#09,#08) * div(0.5, #08) * (1/N)), #02) - // grad[#02] = repeat(-(sum(mul(grad[#10],#00)) * div(div(#01,#08),#08) * div(0.5, #08) * (1/N)), #02) - // grad[#02] = repeat(-(sum(mul(grad[#10],#00)) * div(1,#08*#08) * div(0.5, #08) * (1/N)), #02) - // grad[#02] = repeat(-(sum(mul(grad[#10],#00)) * div(1,#07) * div(0.5, #08) * (1/N)), #02) - // grad[#00] = scale(grad(#10), #09) + scale(mul(#00, grad[#02]), 2.0) - // grad[#00] = scale(grad(#10), #09) + scale(mul(#00, repeat(-(sum(mul(grad[#10],#00)) * div(1,#07) * div(0.5, #08) * (1/N)), #02)), 2.0) - // grad[#00] = scale(grad(#10), #09) + scale(scale(#00, -(sum(mul(grad[#10],#00)) * div(1,#07) * div(0.5, #08) * (1/N))), 2.0) - // grad[#00] = scale(grad(#10), #09) + scale(#00, -(sum(mul(grad[#10],#00)) * div(1,#07) * div(1,#08) * (1/N))) - // grad[#00] = scale(grad(#10), #09) + scale(#00, sum(mul(grad[#10],#00)) * div(1,#07*#08) * (-1/N)) - // grad[#00] = scale(grad(#10), #09) + scale(#00, sum(mul(grad[#10],#00)) * div(1,#07*#08) * (-1/N)) - // grad[#00] = scale(grad(#10), #09) + scale(#00, sum(mul(grad[#10],#00)) * div(1,mean_eps*rms) * (-1/N)) - // grad[#00] = scale(grad(#10), #09) + scale(#00, sum(mul(grad[#10],#00)) * div(-1,rms*N*mean_eps)) - // grad[#00] = scale(grad(#10), #09) + scale(#00, sum(mul(grad[#10],#00)) * div(-1,rms*N*(sum_xx/N+eps))) - // grad[#00] = scale(grad(#10), #09) + scale(#00, sum(mul(grad[#10],#00)) * div(-1,rms*N*sum_xx+rms*N*eps)) - // grad[#00] = scale(dz, rrms) + scale(x, sum(mul(dz,x)) * div(-1,rms*N*mean_eps)) - // grad[#00] = scale(dz, rrms) + scale(x, sum_xdz * div(-1,rms*N*mean_eps)) - // a = b*c + d*e - // a = b*c*f/f + d*e*f/f - // a = (b*c*f + d*e*f)*(1/f) - // a = (b*c*(1/c) + d*e*(1/c))*(1/(1/c)) - // a = (b + d*e/c)*c - // b = dz, c = rrms, d = x, e = sum_xdz * div(-1,rms*N*mean_eps) - // a = (dz + x*sum_xdz * div(-1,rms*N*mean_eps)/rrms)*rrms - // a = (dz + x*sum_xdz * div(-1,rms*N*mean_eps)*rms)*rrms - // a = (dz + x*sum_xdz * div(-rms,rms*N*mean_eps))*rrms - // a = (dz + x*sum_xdz * div(-1,N*mean_eps))*rrms - // a = (dz + x*div(-sum_xdz,N*mean_eps))*rrms - // a = (dz + x*div(-mean_xdz,mean_eps))*rrms - // grad[#00] = scale(dz + scale(x, div(-mean_xdz,mean_eps)),rrms) - // grad[#00] = scale(dz + scale(x, -mean_xdz/mean_eps),rrms) - // dx = scale(dz + scale(x, -mean_xdz/mean_eps),rrms) - } - // dx = scale(dz + scale(x, -mean_xdz/mean_eps),rrms) - // post-order: - // dx := x - // dx := scale(dx,-mean_xdz/mean_eps) - // dx := add(dx, dz) - // dx := scale(dx, rrms) - float * dx = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); + const int n = ggml_nrows(src0); + const int nc = src0->ne[0]; - // dx[i00] = (x*(-sum_xdz/sum_eps) + dz) / sqrtf(mean_eps) - ggml_vec_cpy_f32 (ne00, dx, x); - // ggml_vec_scale_f32(ne00, dx, -mean_xdz/mean_eps); - ggml_vec_scale_f32(ne00, dx, (float)(-sum_xdz)/sum_eps); - ggml_vec_acc_f32 (ne00, dx, dz); - ggml_vec_scale_f32(ne00, dx, rrms); - } - } + float negative_slope; + memcpy(&negative_slope, dst->op_params, sizeof(float)); + + assert(dst->nb[0] == sizeof(ggml_fp16_t)); + assert(src0->nb[0] == sizeof(ggml_fp16_t)); + + for (int i = 0; i < n; i++) { + ggml_vec_leaky_relu_f16(nc, + (ggml_fp16_t *) ((char *) dst->data + i*( dst->nb[1])), + (ggml_fp16_t *) ((char *) src0->data + i*(src0->nb[1])), negative_slope); } } -void ggml_compute_forward_rms_norm_back( +void ggml_compute_forward_leaky_relu( const ggml_compute_params * params, ggml_tensor * dst) { @@ -3948,7 +4791,11 @@ void ggml_compute_forward_rms_norm_back( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_rms_norm_back_f32(params, dst); + ggml_compute_forward_leaky_relu_f32(params, dst); + } break; + case GGML_TYPE_F16: + { + ggml_compute_forward_leaky_relu_f16(params, dst); } break; default: { @@ -3957,156 +4804,109 @@ void ggml_compute_forward_rms_norm_back( } } -// ggml_compute_forward_group_norm - -static void ggml_compute_forward_group_norm_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { +// ggml_compute_forward_silu_back - const ggml_tensor * src0 = dst->src[0]; +static void ggml_compute_forward_silu_back_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - GGML_ASSERT(ggml_are_same_shape(src0, dst)); + const ggml_tensor * grad = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(src0->nb[0] == sizeof(float)); + assert(ggml_is_contiguous_1(grad)); + assert(ggml_is_contiguous_1(src1)); + assert(ggml_is_contiguous_1(dst)); + assert(ggml_are_same_shape(src1, dst)); + assert(ggml_are_same_shape(src1, grad)); const int ith = params->ith; const int nth = params->nth; - GGML_TENSOR_UNARY_OP_LOCALS - - // TODO: optimize - - float eps; - memcpy(&eps, dst->op_params + 1, sizeof(float)); - - int n_channels = src0->ne[2]; - int n_groups = dst->op_params[0]; - int n_channels_per_group = (n_channels + n_groups - 1) / n_groups; - for (int i = ith; i < n_groups; i += nth) { - int start = i * n_channels_per_group; - int end = start + n_channels_per_group; - if (end > n_channels) { - end = n_channels; - } - int step = end - start; - - for (int64_t i03 = 0; i03 < ne03; i03++) { - ggml_float sum = 0.0; - for (int64_t i02 = start; i02 < end; i02++) { - for (int64_t i01 = 0; i01 < ne01; i01++) { - const float * x = (float *)((char *) src0->data + i01 * nb01 + i02 * nb02 + i03 * nb03); - - ggml_float sumr = 0.0; - for (int64_t i00 = 0; i00 < ne00; i00++) { - sumr += (ggml_float)x[i00]; - } - sum += sumr; - } - } - const float mean = sum / (ne00 * ne01 * step); + const int nc = src1->ne[0]; + const int nr = ggml_nrows(src1); - ggml_float sum2 = 0.0; - for (int64_t i02 = start; i02 < end; i02++) { - for (int64_t i01 = 0; i01 < ne01; i01++) { - const float * x = (float *)((char *) src0->data + i01 * nb01 + i02 * nb02 + i03 * nb03); + // rows per thread + const int dr = (nr + nth - 1)/nth; - float * y = (float *)((char *) dst->data + i01 * nb1 + i02 * nb2 + i03 * nb3); + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - ggml_float sumr = 0.0; - for (int64_t i00 = 0; i00 < ne00; i00++) { - float v = x[i00] - mean; - y[i00] = v; - sumr += (ggml_float)(v * v); - } - sum2 += sumr; - } - } - const float variance = sum2 / (ne00 * ne01 * step); - const float scale = 1.0f / sqrtf(variance + eps); + for (int i1 = ir0; i1 < ir1; i1++) { + ggml_vec_silu_backward_f32(nc, + (float *) ((char *) dst->data + i1*( dst->nb[1])), + (float *) ((char *) src1->data + i1*(src1->nb[1])), + (float *) ((char *) grad->data + i1*(grad->nb[1]))); - for (int64_t i02 = start; i02 < end; i02++) { - for (int64_t i01 = 0; i01 < ne01; i01++) { - float * y = (float *)((char *) dst->data + i01 * nb1 + i02 * nb2 + i03 * nb3); - ggml_vec_scale_f32(ne00, y, scale); - } - } +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = ((float *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + GGML_UNUSED(x); + assert(!isnan(x)); + assert(!isinf(x)); } +#endif // NDEBUG } } -void ggml_compute_forward_group_norm( +static void ggml_compute_forward_silu_back_f16( const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_group_norm_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} + const ggml_tensor * grad = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; -// ggml_compute_forward_l2_norm + assert(ggml_is_contiguous_1(grad)); + assert(ggml_is_contiguous_1(src1)); + assert(ggml_is_contiguous_1(dst)); + assert(ggml_are_same_shape(src1, dst)); + assert(ggml_are_same_shape(src1, grad)); -static void ggml_compute_forward_l2_norm_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { + const int ith = params->ith; + const int nth = params->nth; - const ggml_tensor * src0 = dst->src[0]; - - GGML_ASSERT(ggml_are_same_shape(src0, dst)); - - GGML_ASSERT(src0->nb[0] == sizeof(float)); - - const int ith = params->ith; - const int nth = params->nth; - - GGML_TENSOR_UNARY_OP_LOCALS - - float eps; - memcpy(&eps, dst->op_params, sizeof(float)); - - GGML_ASSERT(eps >= 0.0f); - - // TODO: optimize - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = ith; i01 < ne01; i01 += nth) { - const float * x = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); - - ggml_float sum = 0.0; - for (int64_t i00 = 0; i00 < ne00; i00++) { - sum += (ggml_float)(x[i00] * x[i00]); - } + const int nc = src1->ne[0]; + const int nr = ggml_nrows(src1); - float * y = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); + // rows per thread + const int dr = (nr + nth - 1)/nth; - memcpy(y, x, ne00 * sizeof(float)); + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - const float scale = 1.0f/fmaxf(sqrtf(sum), eps); + for (int i1 = ir0; i1 < ir1; i1++) { + ggml_vec_silu_backward_f16(nc, + (ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])), + (ggml_fp16_t *) ((char *) src1->data + i1*(src1->nb[1])), + (ggml_fp16_t *) ((char *) grad->data + i1*(grad->nb[1]))); - ggml_vec_scale_f32(ne00, y, scale); - } +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = ((ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + const float v = GGML_CPU_FP16_TO_FP32(x); + GGML_UNUSED(v); + assert(!isnan(v)); + assert(!isinf(v)); } +#endif // NDEBUG } } -void ggml_compute_forward_l2_norm( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_silu_back( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_l2_norm_f32(params, dst); + ggml_compute_forward_silu_back_f32(params, dst); + } break; + case GGML_TYPE_F16: + { + ggml_compute_forward_silu_back_f16(params, dst); } break; default: { @@ -4115,260 +4915,284 @@ void ggml_compute_forward_l2_norm( } } -// ggml_compute_forward_out_prod +// ggml_compute_forward_reglu -static void ggml_compute_forward_out_prod_f32( +static void ggml_compute_forward_reglu_f32( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - GGML_TENSOR_BINARY_OP_LOCALS + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); + if (src1) { + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + } const int ith = params->ith; const int nth = params->nth; - GGML_ASSERT(ne0 == ne00); - GGML_ASSERT(ne1 == ne10); - GGML_ASSERT(ne2 == ne12); - GGML_ASSERT(ne3 == ne13); + const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nr = ggml_nrows(src0); - GGML_ASSERT(ne2 % ne02 == 0); - GGML_ASSERT(ne3 % ne03 == 0); + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); - // we don't support permuted src0 or src1 - GGML_ASSERT(nb00 == sizeof(float)); + const int32_t swapped = ggml_get_op_params_i32(dst, 1); - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - // GGML_ASSERT(nb0 <= nb1); - // GGML_ASSERT(nb1 <= nb2); - // GGML_ASSERT(nb2 <= nb3); + // rows per thread + const int dr = (nr + nth - 1)/nth; - // nb01 >= nb00 - src0 is not transposed - // compute by src0 rows + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - if (ith == 0) { - ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0); - } - ggml_barrier(params->threadpool); + for (int i1 = ir0; i1 < ir1; i1++) { + float * src0_p = (float *) (src0_d + i1*src0_o); + float * src1_p = (float *) (src1_d + i1*src1_o); - // dst[:,:,:,:] = 0 - // for i2,i3: - // for i1: - // for i01: - // for i0: - // dst[i0,i1,i2,i3] += src0[i0,i01,i2,i3] * src1[i1,i01,i2,i3] + if (!src1) { + src0_p += swapped ? nc : 0; + src1_p += swapped ? 0 : nc; + } - // parallelize by last three dimensions + ggml_vec_reglu_f32(nc, (float *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); - // total rows in dst - const int64_t nr = ne1*ne2*ne3; +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = ((float *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + GGML_UNUSED(x); + assert(!isnan(x)); + assert(!isinf(x)); + } +#endif // NDEBUG + } +} - // rows per thread - const int64_t dr = (nr + nth - 1)/nth; +static void ggml_compute_forward_reglu_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { - // row range for this thread - const int64_t ir0 = dr*ith; - const int64_t ir1 = MIN(ir0 + dr, nr); + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - // block-tiling attempt - const int64_t blck_0 = MAX(GGML_VEC_MAD_UNROLL, 32); - const int64_t blck_1 = 16; + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); - // dps == dst per src0, used for group query attention - const int64_t dps2 = ne2 / ne02; - const int64_t dps3 = ne3 / ne03; + if (src1) { + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + } - for (int64_t bir = ir0; bir < ir1; bir += blck_1) { - const int64_t bir1 = MIN(bir + blck_1, ir1); - for (int64_t bi01 = 0; bi01 < ne01; bi01 += blck_0) { - const int64_t bne01 = MIN(bi01 + blck_0, ne01); - for (int64_t ir = bir; ir < bir1; ++ir) { - // dst indices - const int64_t i3 = ir/(ne2*ne1); - const int64_t i2 = (ir - i3*ne2*ne1)/ne1; - const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1); + const int ith = params->ith; + const int nth = params->nth; - const int64_t i02 = i2 / dps2; - const int64_t i03 = i3 / dps3; + const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nr = ggml_nrows(src0); - //const int64_t i10 = i1; - const int64_t i12 = i2; - const int64_t i13 = i3; + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); -#if GGML_VEC_MAD_UNROLL > 2 - const int64_t bne01_unroll = bne01 - (bne01 % GGML_VEC_MAD_UNROLL); - for (int64_t i01 = bi01; i01 < bne01_unroll; i01 += GGML_VEC_MAD_UNROLL) { - const int64_t i11 = i01; + const int32_t swapped = ggml_get_op_params_i32(dst, 1); - float * s0 = (float *) ((char *) src0->data + ( i01*nb01 + i02*nb02 + i03*nb03)); - float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13)); - float * d = (float *) ((char *) dst->data + ( i1*nb1 + i2*nb2 + i3*nb3)); + // rows per thread + const int dr = (nr + nth - 1)/nth; - ggml_vec_mad_f32_unroll(ne0, nb01, nb11, d, s0, s1); - } - for (int64_t i01 = bne01_unroll; i01 < bne01; ++i01) { - const int64_t i11 = i01; + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - float * s0 = (float *) ((char *) src0->data + ( i01*nb01 + i02*nb02 + i03*nb03)); - float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13)); - float * d = (float *) ((char *) dst->data + ( i1*nb1 + i2*nb2 + i3*nb3)); + for (int i1 = ir0; i1 < ir1; i1++) { + ggml_fp16_t * src0_p = (ggml_fp16_t *) (src0_d + i1*src0_o); + ggml_fp16_t * src1_p = (ggml_fp16_t *) (src1_d + i1*src1_o); - ggml_vec_mad_f32(ne0, d, s0, *s1); - } -#else - for (int64_t i01 = bi01; i01 < bne01; ++i01) { - const int64_t i11 = i01; + if (!src1) { + src0_p += swapped ? nc : 0; + src1_p += swapped ? 0 : nc; + } - float * s0 = (float *) ((char *) src0->data + ( i01*nb01 + i02*nb02 + i03*nb03)); - float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13)); - float * d = (float *) ((char *) dst->data + ( i1*nb1 + i2*nb2 + i3*nb3)); + ggml_vec_reglu_f16(nc, (ggml_fp16_t *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); - ggml_vec_mad_f32(ne0, d, s0, *s1); - } -#endif - } +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + const float v = GGML_FP16_TO_FP32(x); + GGML_UNUSED(v); + assert(!isnan(v)); + assert(!isinf(v)); } +#endif // NDEBUG } } -static void ggml_compute_forward_out_prod_q_f32( +static void ggml_compute_forward_reglu( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - GGML_TENSOR_BINARY_OP_LOCALS; + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_reglu_f32(params, dst); + } break; + case GGML_TYPE_F16: + { + ggml_compute_forward_reglu_f16(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - const int ith = params->ith; - const int nth = params->nth; - - const ggml_type type = src0->type; - ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; +// ggml_compute_forward_geglu - GGML_ASSERT(ne02 == ne12); - GGML_ASSERT(ne03 == ne13); - GGML_ASSERT(ne2 == ne12); - GGML_ASSERT(ne3 == ne13); +static void ggml_compute_forward_geglu_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - // we don't support permuted src0 dim0 - GGML_ASSERT(nb00 == ggml_type_size(type)); + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - // dst dim0 cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - // GGML_ASSERT(nb0 <= nb1); - // GGML_ASSERT(nb1 <= nb2); - // GGML_ASSERT(nb2 <= nb3); + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); - GGML_ASSERT(ne0 == ne00); - GGML_ASSERT(ne1 == ne10); - GGML_ASSERT(ne2 == ne02); - GGML_ASSERT(ne3 == ne03); + if (src1) { + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + } - // nb01 >= nb00 - src0 is not transposed - // compute by src0 rows + const int ith = params->ith; + const int nth = params->nth; - if (ith == 0) { - ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0); - } - ggml_barrier(params->threadpool); + const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nr = ggml_nrows(src0); - // parallelize by last three dimensions + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); - // total rows in dst - const int64_t nr = ne1*ne2*ne3; + const int32_t swapped = ggml_get_op_params_i32(dst, 1); // rows per thread - const int64_t dr = (nr + nth - 1)/nth; + const int dr = (nr + nth - 1)/nth; // row range for this thread - const int64_t ir0 = dr*ith; - const int64_t ir1 = MIN(ir0 + dr, nr); + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - // dst[:,:,:,:] = 0 - // for i2,i3: - // for i1: - // for i01: - // for i0: - // dst[i0,i1,i2,i3] += src0[i0,i01,i2,i3] * src1[i1,i01,i2,i3] + for (int i1 = ir0; i1 < ir1; i1++) { + float * src0_p = (float *) (src0_d + i1*src0_o); + float * src1_p = (float *) (src1_d + i1*src1_o); - float * wdata = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32) * ith; + if (!src1) { + src0_p += swapped ? nc : 0; + src1_p += swapped ? 0 : nc; + } - for (int64_t ir = ir0; ir < ir1; ++ir) { - // dst indices - const int64_t i3 = ir/(ne2*ne1); - const int64_t i2 = (ir - i3*ne2*ne1)/ne1; - const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1); + ggml_vec_geglu_f32(nc, (float *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); - const int64_t i02 = i2; - const int64_t i03 = i3; +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = ((float *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + GGML_UNUSED(x); + assert(!isnan(x)); + assert(!isinf(x)); + } +#endif // NDEBUG + } +} - //const int64_t i10 = i1; - const int64_t i12 = i2; - const int64_t i13 = i3; +static void ggml_compute_forward_geglu_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { - for (int64_t i01 = 0; i01 < ne01; ++i01) { - const int64_t i11 = i01; + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - float * s0 = (float *) ((char *) src0->data + ( i01*nb01 + i02*nb02 + i03*nb03)); - float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13)); - float * d = (float *) ((char *) dst->data + ( i1*nb1 + i2*nb2 + i3*nb3)); + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); - dequantize_row_q(s0, wdata, ne0); - ggml_vec_mad_f32(ne0, d, wdata, *s1); + if (src1) { + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + } + + const int ith = params->ith; + const int nth = params->nth; + + const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nr = ggml_nrows(src0); + + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); + + const int32_t swapped = ggml_get_op_params_i32(dst, 1); + + // rows per thread + const int dr = (nr + nth - 1)/nth; + + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + + for (int i1 = ir0; i1 < ir1; i1++) { + ggml_fp16_t * src0_p = (ggml_fp16_t *) (src0_d + i1*src0_o); + ggml_fp16_t * src1_p = (ggml_fp16_t *) (src1_d + i1*src1_o); + + if (!src1) { + src0_p += swapped ? nc : 0; + src1_p += swapped ? 0 : nc; + } + + ggml_vec_geglu_f16(nc, (ggml_fp16_t *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + const float v = GGML_FP16_TO_FP32(x); + GGML_UNUSED(v); + assert(!isnan(v)); + assert(!isinf(v)); } +#endif // NDEBUG } } -void ggml_compute_forward_out_prod( +static void ggml_compute_forward_geglu( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { - case GGML_TYPE_Q1_0: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q8_0: - case GGML_TYPE_MXFP4: - case GGML_TYPE_NVFP4: - case GGML_TYPE_Q2_K: - case GGML_TYPE_Q3_K: - case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: - case GGML_TYPE_Q6_K: - case GGML_TYPE_TQ1_0: - case GGML_TYPE_TQ2_0: - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ2_S: + case GGML_TYPE_F32: { - ggml_compute_forward_out_prod_q_f32(params, dst); + ggml_compute_forward_geglu_f32(params, dst); } break; case GGML_TYPE_F16: { - GGML_ABORT("fatal error"); // todo - // ggml_compute_forward_out_prod_f16_f32(params, dst); - } - case GGML_TYPE_F32: - { - ggml_compute_forward_out_prod_f32(params, dst); + ggml_compute_forward_geglu_f16(params, dst); } break; default: { @@ -4377,30 +5201,38 @@ void ggml_compute_forward_out_prod( } } -// ggml_compute_forward_scale +// ggml_compute_forward_swiglu -static void ggml_compute_forward_scale_f32( +static void ggml_compute_forward_swiglu_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - GGML_ASSERT(ggml_is_contiguous(src0)); - GGML_ASSERT(ggml_is_contiguous(dst)); - GGML_ASSERT(ggml_are_same_shape(src0, dst)); - - float s; // scale factor - float b; // bias + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); - memcpy(&s, (float *) dst->op_params + 0, sizeof(float)); - memcpy(&b, (float *) dst->op_params + 1, sizeof(float)); + if (src1) { + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + } const int ith = params->ith; const int nth = params->nth; - const int nc = src0->ne[0]; + const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; const int nr = ggml_nrows(src0); + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); + + const int32_t swapped = ggml_get_op_params_i32(dst, 1); + // rows per thread const int dr = (nr + nth - 1)/nth; @@ -4408,99 +5240,57 @@ static void ggml_compute_forward_scale_f32( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - const size_t nb01 = src0->nb[1]; - - const size_t nb1 = dst->nb[1]; + for (int i1 = ir0; i1 < ir1; i1++) { + float * src0_p = (float *) (src0_d + i1*src0_o); + float * src1_p = (float *) (src1_d + i1*src1_o); - if (b == 0.0f) { - for (int i1 = ir0; i1 < ir1; i1++) { - if (dst->data != src0->data) { - // src0 is same shape as dst => same indices - // TODO: add x parameter to ggml_vec_scale_f32 and remove this memcpy - memcpy((char *)dst->data + i1*nb1, (char *)src0->data + i1*nb01, nc * sizeof(float)); - } - ggml_vec_scale_f32(nc, (float *) ((char *) dst->data + i1*nb1), s); + if (!src1) { + src0_p += swapped ? nc : 0; + src1_p += swapped ? 0 : nc; } - } else { - for (int i1 = ir0; i1 < ir1; i1++) { - ggml_vec_mad1_f32(nc, - (float *) ((char *) dst->data + i1*nb1), - (float *) ((char *) src0->data + i1*nb1), - s, b); + + ggml_vec_swiglu_f32(nc, (float *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = ((float *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + GGML_UNUSED(x); + assert(!isnan(x)); + assert(!isinf(x)); } +#endif // NDEBUG } } -void ggml_compute_forward_scale( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_scale_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} - -// ggml_compute_forward_set - -static void ggml_compute_forward_set_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { +static void ggml_compute_forward_swiglu_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - GGML_ASSERT(ggml_are_same_shape(src0, dst)); - GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src0)); - - // view src0 and dst with these strides and data offset inbytes during set - // nb0 is implicitly element_size because src0 and dst are contiguous - size_t nb1 = ((int32_t *) dst->op_params)[0]; - size_t nb2 = ((int32_t *) dst->op_params)[1]; - size_t nb3 = ((int32_t *) dst->op_params)[2]; - size_t offset = ((int32_t *) dst->op_params)[3]; - bool inplace = (bool) ((int32_t *) dst->op_params)[4]; + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); - if (!inplace) { - if (params->ith == 0) { - // memcpy needs to be synchronized across threads to avoid race conditions. - // => do it in INIT phase - memcpy( - ((char *) dst->data), - ((char *) src0->data), - ggml_nbytes(dst)); - } - ggml_barrier(params->threadpool); + if (src1) { + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); } const int ith = params->ith; const int nth = params->nth; - const int nr = ggml_nrows(src1); - const int nc = src1->ne[0]; - - GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) - GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) - - // src0 and dst as viewed during set - const size_t nb0 = ggml_element_size(src0); - - const int im0 = (ne10 == 0 ? 0 : ne10-1); - const int im1 = (ne11 == 0 ? 0 : ne11-1); - const int im2 = (ne12 == 0 ? 0 : ne12-1); - const int im3 = (ne13 == 0 ? 0 : ne13-1); + const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nr = ggml_nrows(src0); - GGML_ASSERT(offset + im0*nb0 + im1*nb1 + im2*nb2 + im3*nb3 <= ggml_nbytes(dst)); + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); - GGML_ASSERT(nb10 == sizeof(float)); + const int32_t swapped = ggml_get_op_params_i32(dst, 1); // rows per thread const int dr = (nr + nth - 1)/nth; @@ -4509,69 +5299,84 @@ static void ggml_compute_forward_set_f32( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - for (int ir = ir0; ir < ir1; ++ir) { - // src0 and dst are viewed with shape of src1 and offset - // => same indices - const int i3 = ir/(ne12*ne11); - const int i2 = (ir - i3*ne12*ne11)/ne11; - const int i1 = (ir - i3*ne12*ne11 - i2*ne11); + for (int i1 = ir0; i1 < ir1; i1++) { + ggml_fp16_t * src0_p = (ggml_fp16_t *) (src0_d + i1*src0_o); + ggml_fp16_t * src1_p = (ggml_fp16_t *) (src1_d + i1*src1_o); - ggml_vec_cpy_f32(nc, - (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + offset), - (float *) ((char *) src1->data + i3*nb13 + i2*nb12 + i1*nb11)); + if (!src1) { + src0_p += swapped ? nc : 0; + src1_p += swapped ? 0 : nc; + } + + ggml_vec_swiglu_f16(nc, (ggml_fp16_t *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + const float v = GGML_FP16_TO_FP32(x); + GGML_UNUSED(v); + assert(!isnan(v)); + assert(!isinf(v)); + } +#endif // NDEBUG } } -static void ggml_compute_forward_set_i32( +static void ggml_compute_forward_swiglu( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(ggml_are_same_shape(src0, dst)); - GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src0)); + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_swiglu_f32(params, dst); + } break; + case GGML_TYPE_F16: + { + ggml_compute_forward_swiglu_f16(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - // view src0 and dst with these strides and data offset inbytes during set - // nb0 is implicitly element_size because src0 and dst are contiguous - size_t nb1 = ((int32_t *) dst->op_params)[0]; - size_t nb2 = ((int32_t *) dst->op_params)[1]; - size_t nb3 = ((int32_t *) dst->op_params)[2]; - size_t offset = ((int32_t *) dst->op_params)[3]; - bool inplace = (bool) ((int32_t *) dst->op_params)[4]; +// ggml_compute_forward_swiglu_oai - if (!inplace) { - if (params->ith == 0) { - // memcpy needs to be synchronized across threads to avoid race conditions. - // => do it in INIT phase - memcpy( - ((char *) dst->data), - ((char *) src0->data), - ggml_nbytes(dst)); - } - ggml_barrier(params->threadpool); - } +static void ggml_compute_forward_swiglu_oai_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - const int ith = params->ith; - const int nth = params->nth; + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - const int nr = ggml_nrows(src1); - const int nc = src1->ne[0]; + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); - GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) - GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) + if (src1) { + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + } - // src0 and dst as viewed during set - const size_t nb0 = ggml_element_size(src0); + const int ith = params->ith; + const int nth = params->nth; - const int im0 = (ne10 == 0 ? 0 : ne10-1); - const int im1 = (ne11 == 0 ? 0 : ne11-1); - const int im2 = (ne12 == 0 ? 0 : ne12-1); - const int im3 = (ne13 == 0 ? 0 : ne13-1); + const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nr = ggml_nrows(src0); - GGML_ASSERT(offset + im0*nb0 + im1*nb1 + im2*nb2 + im3*nb3 <= ggml_nbytes(dst)); + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); - GGML_ASSERT(nb10 == sizeof(int32_t)); + const int32_t swapped = ggml_get_op_params_i32(dst, 1); + const float alpha = ggml_get_op_params_f32(dst, 2); + const float limit = ggml_get_op_params_f32(dst, 3); // rows per thread const int dr = (nr + nth - 1)/nth; @@ -4580,20 +5385,35 @@ static void ggml_compute_forward_set_i32( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - for (int ir = ir0; ir < ir1; ++ir) { - // src0 and dst are viewed with shape of src1 and offset - // => same indices - const int i3 = ir/(ne12*ne11); - const int i2 = (ir - i3*ne12*ne11)/ne11; - const int i1 = (ir - i3*ne12*ne11 - i2*ne11); + for (int i1 = ir0; i1 < ir1; i1++) { + float * src0_p = (float *) (src0_d + i1*src0_o); + float * src1_p = (float *) (src1_d + i1*src1_o); + float * dst_p = (float *) ((char *) dst->data + i1*(dst->nb[1])); - ggml_vec_cpy_i32(nc, - (int32_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + offset), - (int32_t *) ((char *) src1->data + i3*nb13 + i2*nb12 + i1*nb11)); + if (!src1) { + src0_p += swapped ? nc : 0; + src1_p += swapped ? 0 : nc; + } + + for (int k = 0; k < nc; k++) { + const float x = std::min(src0_p[k], limit); + const float y = std::clamp(src1_p[k], -limit, limit); + const float out_glu = x / (1.f + expf(alpha * (-x))); + dst_p[k] = out_glu * (y + 1.f); + } + +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = dst_p[k]; + GGML_UNUSED(x); + assert(!isnan(x)); + assert(!isinf(x)); + } +#endif // NDEBUG } } -void ggml_compute_forward_set( +static void ggml_compute_forward_swiglu_oai( const ggml_compute_params * params, ggml_tensor * dst) { @@ -4602,39 +5422,8 @@ void ggml_compute_forward_set( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_set_f32(params, dst); - } break; - case GGML_TYPE_I32: - { - ggml_compute_forward_set_i32(params, dst); + ggml_compute_forward_swiglu_oai_f32(params, dst); } break; - case GGML_TYPE_F16: - case GGML_TYPE_BF16: - case GGML_TYPE_Q1_0: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q8_0: - case GGML_TYPE_Q8_1: - case GGML_TYPE_MXFP4: - case GGML_TYPE_NVFP4: - case GGML_TYPE_Q2_K: - case GGML_TYPE_Q3_K: - case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: - case GGML_TYPE_Q6_K: - case GGML_TYPE_TQ1_0: - case GGML_TYPE_TQ2_0: - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ2_S: default: { GGML_ABORT("fatal error"); @@ -4642,46 +5431,37 @@ void ggml_compute_forward_set( } } -// ggml_compute_forward_cpy +// ggml_compute_forward_geglu_erf -void ggml_compute_forward_cpy( +static void ggml_compute_forward_geglu_erf_f32( const ggml_compute_params * params, ggml_tensor * dst) { - ggml_compute_forward_dup(params, dst); -} -// ggml_compute_forward_cont + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; -void ggml_compute_forward_cont( - const ggml_compute_params * params, - ggml_tensor * dst) { - ggml_compute_forward_dup(params, dst); -} - -// ggml_compute_forward_get_rows - -static void ggml_compute_forward_get_rows_q( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); - GGML_TENSOR_BINARY_OP_LOCALS + if (src1) { + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + } - const int64_t nc = ne00; - const int64_t nr = ggml_nelements(src1); + const int ith = params->ith; + const int nth = params->nth; - const ggml_type type = src0->type; - ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; + const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nr = ggml_nrows(src0); - assert(ne0 == nc); - assert(ne02 == ne11); - assert(nb00 == ggml_type_size(type)); - assert(ggml_nrows(dst) == nr); + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); - const int ith = params->ith; - const int nth = params->nth; + const int32_t swapped = ggml_get_op_params_i32(dst, 1); // rows per thread const int dr = (nr + nth - 1)/nth; @@ -4690,40 +5470,58 @@ static void ggml_compute_forward_get_rows_q( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - for (int64_t i = ir0; i < ir1; ++i) { - const int64_t i12 = i/(ne11*ne10); - const int64_t i11 = (i - i12*ne11*ne10)/ne10; - const int64_t i10 = (i - i12*ne11*ne10 - i11*ne10); - const int64_t i01 = *(int32_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12); + for (int i1 = ir0; i1 < ir1; i1++) { + float * src0_p = (float *) (src0_d + i1*src0_o); + float * src1_p = (float *) (src1_d + i1*src1_o); - GGML_ASSERT(i01 >= 0 && i01 < ne01); + if (!src1) { + src0_p += swapped ? nc : 0; + src1_p += swapped ? 0 : nc; + } - dequantize_row_q( - (const void *) ((char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03), - (float *) ((char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3), nc); + ggml_vec_geglu_erf_f32(nc, (float *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = ((float *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + GGML_UNUSED(x); + assert(!isnan(x)); + assert(!isinf(x)); + } +#endif // NDEBUG } } -static void ggml_compute_forward_get_rows_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { +static void ggml_compute_forward_geglu_erf_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - GGML_TENSOR_BINARY_OP_LOCALS - - const int64_t nc = ne00; - const int64_t nr = ggml_nelements(src1); + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); - assert(ne0 == nc); - assert(ne02 == ne11); - assert(nb00 == sizeof(ggml_fp16_t)); - assert(ggml_nrows(dst) == nr); + if (src1) { + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + } const int ith = params->ith; const int nth = params->nth; + const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nr = ggml_nrows(src0); + + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); + + const int32_t swapped = ggml_get_op_params_i32(dst, 1); + // rows per thread const int dr = (nr + nth - 1)/nth; @@ -4731,40 +5529,83 @@ static void ggml_compute_forward_get_rows_f16( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - for (int64_t i = ir0; i < ir1; ++i) { - const int64_t i12 = i/(ne11*ne10); - const int64_t i11 = (i - i12*ne11*ne10)/ne10; - const int64_t i10 = (i - i12*ne11*ne10 - i11*ne10); - const int64_t i01 = *(int32_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12); + for (int i1 = ir0; i1 < ir1; i1++) { + ggml_fp16_t * src0_p = (ggml_fp16_t *) (src0_d + i1*src0_o); + ggml_fp16_t * src1_p = (ggml_fp16_t *) (src1_d + i1*src1_o); - GGML_ASSERT(i01 >= 0 && i01 < ne01); + if (!src1) { + src0_p += swapped ? nc : 0; + src1_p += swapped ? 0 : nc; + } - ggml_cpu_fp16_to_fp32( - (const ggml_fp16_t*) ((char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03), - (float *) ((char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3), nc); + ggml_vec_geglu_erf_f16(nc, (ggml_fp16_t *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + const float v = GGML_FP16_TO_FP32(x); + GGML_UNUSED(v); + assert(!isnan(v)); + assert(!isinf(v)); + } +#endif // NDEBUG } } -static void ggml_compute_forward_get_rows_bf16( +static void ggml_compute_forward_geglu_erf( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - GGML_TENSOR_BINARY_OP_LOCALS + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_geglu_erf_f32(params, dst); + } break; + case GGML_TYPE_F16: + { + ggml_compute_forward_geglu_erf_f16(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - const int64_t nc = ne00; - const int64_t nr = ggml_nelements(src1); +// ggml_compute_forward_geglu_quick - assert(ne0 == nc); - assert(ne02 == ne11); - assert(nb00 == sizeof(ggml_bf16_t)); - assert(ggml_nrows(dst) == nr); +static void ggml_compute_forward_geglu_quick_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; + + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); + + if (src1) { + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + } const int ith = params->ith; const int nth = params->nth; + const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nr = ggml_nrows(src0); + + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); + + const int32_t swapped = ggml_get_op_params_i32(dst, 1); + // rows per thread const int dr = (nr + nth - 1)/nth; @@ -4772,40 +5613,58 @@ static void ggml_compute_forward_get_rows_bf16( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - for (int64_t i = ir0; i < ir1; ++i) { - const int64_t i12 = i/(ne11*ne10); - const int64_t i11 = (i - i12*ne11*ne10)/ne10; - const int64_t i10 = (i - i12*ne11*ne10 - i11*ne10); - const int64_t i01 = *(int32_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12); + for (int i1 = ir0; i1 < ir1; i1++) { + float * src0_p = (float *) (src0_d + i1*src0_o); + float * src1_p = (float *) (src1_d + i1*src1_o); - GGML_ASSERT(i01 >= 0 && i01 < ne01); + if (!src1) { + src0_p += swapped ? nc : 0; + src1_p += swapped ? 0 : nc; + } - ggml_cpu_bf16_to_fp32( - (const ggml_bf16_t *) ((char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03), - (float *) ((char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3), nc); + ggml_vec_geglu_quick_f32(nc, (float *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); + +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = ((float *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + GGML_UNUSED(x); + assert(!isnan(x)); + assert(!isinf(x)); + } +#endif // NDEBUG } } -static void ggml_compute_forward_get_rows_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { +static void ggml_compute_forward_geglu_quick_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; - GGML_TENSOR_BINARY_OP_LOCALS - - const int64_t nc = ne00; - const int64_t nr = ggml_nelements(src1); + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); - assert(ne0 == nc); - assert(ne02 == ne11); - assert(nb00 == sizeof(float)); - assert(ggml_nrows(dst) == nr); + if (src1) { + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + } const int ith = params->ith; const int nth = params->nth; + const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nr = ggml_nrows(src0); + + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); + + const int32_t swapped = ggml_get_op_params_i32(dst, 1); + // rows per thread const int dr = (nr + nth - 1)/nth; @@ -4813,320 +5672,347 @@ static void ggml_compute_forward_get_rows_f32( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - for (int64_t i = ir0; i < ir1; ++i) { - const int64_t i12 = i/(ne11*ne10); - const int64_t i11 = (i - i12*ne11*ne10)/ne10; - const int64_t i10 = (i - i12*ne11*ne10 - i11*ne10); - const int64_t i01 = *(int32_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12); + for (int i1 = ir0; i1 < ir1; i1++) { + ggml_fp16_t * src0_p = (ggml_fp16_t *) (src0_d + i1*src0_o); + ggml_fp16_t * src1_p = (ggml_fp16_t *) (src1_d + i1*src1_o); - GGML_ASSERT(i01 >= 0 && i01 < ne01); + if (!src1) { + src0_p += swapped ? nc : 0; + src1_p += swapped ? 0 : nc; + } - ggml_vec_cpy_f32(nc, - (float *) ((char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3), - (float *) ((char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03)); - } -} + ggml_vec_geglu_quick_f16(nc, (ggml_fp16_t *) ((char *) dst->data + i1*(dst->nb[1])), src0_p, src1_p); -void ggml_compute_forward_get_rows( - const ggml_compute_params * params, - ggml_tensor * dst) { +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const ggml_fp16_t x = ((ggml_fp16_t *) ((char *) dst->data + i1*( dst->nb[1])))[k]; + const float v = GGML_FP16_TO_FP32(x); + GGML_UNUSED(v); + assert(!isnan(v)); + assert(!isinf(v)); + } +#endif // NDEBUG + } +} + +static void ggml_compute_forward_geglu_quick( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { - case GGML_TYPE_Q1_0: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q8_0: - case GGML_TYPE_Q8_1: - case GGML_TYPE_MXFP4: - case GGML_TYPE_NVFP4: - case GGML_TYPE_Q2_K: - case GGML_TYPE_Q3_K: - case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: - case GGML_TYPE_Q6_K: - case GGML_TYPE_TQ1_0: - case GGML_TYPE_TQ2_0: - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ2_S: + case GGML_TYPE_F32: { - ggml_compute_forward_get_rows_q(params, dst); + ggml_compute_forward_geglu_quick_f32(params, dst); } break; case GGML_TYPE_F16: { - ggml_compute_forward_get_rows_f16(params, dst); - } break; - case GGML_TYPE_BF16: - { - ggml_compute_forward_get_rows_bf16(params, dst); - } break; - case GGML_TYPE_F32: - case GGML_TYPE_I32: - { - ggml_compute_forward_get_rows_f32(params, dst); + ggml_compute_forward_geglu_quick_f16(params, dst); } break; default: { GGML_ABORT("fatal error"); } } - - //static bool first = true; - //printf("ne0 = %d, ne1 = %d, ne2 = %d\n", dst->ne[0], dst->ne[1], dst->ne[2]); - //if (first) { - // first = false; - //} else { - // for (int k = 0; k < dst->ne[1]; ++k) { - // for (int j = 0; j < dst->ne[0]/16; ++j) { - // for (int i = 0; i < 16; ++i) { - // printf("%8.4f ", ((float *) dst->data)[k*dst->ne[0] + j*16 + i]); - // } - // printf("\n"); - // } - // printf("\n"); - // } - // printf("\n"); - // exit(0); - //} } -template -static void ggml_compute_forward_set_rows_f32( +// ggml_compute_forward_norm + +static void ggml_compute_forward_norm_f32( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - GGML_TENSOR_BINARY_OP_LOCALS - - const int64_t nc = ne00; - const int64_t nr = ne01; + GGML_ASSERT(ggml_are_same_shape(src0, dst)); - assert(ne0 == nc); - assert(ne2 == ne02); - assert(ne3 == ne03); - assert(src0->type == GGML_TYPE_F32); - assert(ne02 % ne11 == 0); - assert(ne03 % ne12 == 0); + GGML_ASSERT(src0->nb[0] == sizeof(float)); const int ith = params->ith; const int nth = params->nth; - // rows per thread - const int64_t dr = (nr + nth - 1)/nth; + GGML_TENSOR_UNARY_OP_LOCALS - // row range for this thread - const int64_t ir0 = dr*ith; - const int64_t ir1 = std::min(ir0 + dr, nr); + float eps; + memcpy(&eps, dst->op_params, sizeof(float)); - ggml_from_float_t const from_float = ggml_get_type_traits_cpu(dst->type)->from_float; + GGML_ASSERT(eps >= 0.0f); - for (int64_t i03 = 0; i03 < ne03; ++i03) { - for (int64_t i02 = 0; i02 < ne02; ++i02) { - for (int64_t i = ir0; i < ir1; ++i) { - const int64_t i12 = i03%ne12; - const int64_t i11 = i02%ne11; - const int64_t i10 = i; + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = ith; i01 < ne01; i01 += nth) { + const float * x = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); - const int64_t i1 = *(idx_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12); + float sum = 0.0; + ggml_vec_sum_f32(ne00, &sum, x); + float mean = sum/ne00; - GGML_ASSERT(i1 >= 0 && i1 < ne1); + float * y = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); + float variance = 0; - from_float( - (const float *) ((char *) src0->data + i*nb01 + i02*nb02 + i03*nb03), - ((char *) dst->data + i1*nb1 + i02*nb2 + i03*nb3), nc); +#ifdef GGML_USE_ACCELERATE + mean = -mean; + vDSP_vsadd(x, 1, &mean, y, 1, ne00); + vDSP_measqv(y, 1, &variance, ne00); +#else + variance = ggml_vec_cvar_f32(ne00, y, x, mean); +#endif //GGML_USE_ACCELERATE + + const float scale = 1.0f/sqrtf(variance + eps); + ggml_vec_scale_f32(ne00, y, scale); } } } } -void ggml_compute_forward_set_rows( +void ggml_compute_forward_norm( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; switch (src0->type) { case GGML_TYPE_F32: { - if (src1->type == GGML_TYPE_I64) { - ggml_compute_forward_set_rows_f32(params, dst); - } else if (src1->type == GGML_TYPE_I32) { - ggml_compute_forward_set_rows_f32(params, dst); - } else { - GGML_ABORT("src1->type = %d (%s) not supported", src1->type, ggml_type_name(src1->type)); - } + ggml_compute_forward_norm_f32(params, dst); } break; default: { - GGML_ABORT("src0->type = %d (%s) not supported", src0->type, ggml_type_name(src0->type)); + GGML_ABORT("fatal error"); } } } -// ggml_compute_forward_get_rows_back +// ggml_compute_forward_group_rms_norm -static void ggml_compute_forward_get_rows_back_f32_f16( +static void ggml_compute_forward_rms_norm_f32( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - - if (params->ith != 0) { - return; - } - - GGML_ASSERT(ggml_is_contiguous(dst)); - - // ggml_compute_forward_dup_same_cont(params, opt0, dst); - memset(dst->data, 0, ggml_nbytes(dst)); - - const int nc = src0->ne[0]; - const int nr = ggml_nelements(src1); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); - GGML_ASSERT( dst->ne[0] == nc); - GGML_ASSERT(src0->nb[0] == sizeof(ggml_fp16_t)); + GGML_ASSERT(src0->nb[0] == sizeof(float)); - for (int i = 0; i < nr; ++i) { - const int r = ((int32_t *) src1->data)[i]; + const int ith = params->ith; + const int nth = params->nth; - for (int j = 0; j < nc; ++j) { - ggml_fp16_t v = ((ggml_fp16_t *) ((char *) src0->data + i*src0->nb[1]))[j]; - ((float *) ((char *) dst->data + r*dst->nb[1]))[j] += GGML_CPU_FP16_TO_FP32(v); - } - } -} + GGML_TENSOR_UNARY_OP_LOCALS -static void ggml_compute_forward_get_rows_back_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { + float eps; + memcpy(&eps, dst->op_params, sizeof(float)); - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; + GGML_ASSERT(eps >= 0.0f); - if (params->ith != 0) { - return; - } + // TODO: optimize + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = ith; i01 < ne01; i01 += nth) { + const float * x = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); - GGML_ASSERT(ggml_is_contiguous(dst)); + ggml_float sum = 0.0; + for (int64_t i00 = 0; i00 < ne00; i00++) { + sum += (ggml_float)(x[i00] * x[i00]); + } - // ggml_compute_forward_dup_same_cont(params, opt0, dst); + const float mean = sum/ne00; - memset(dst->data, 0, ggml_nbytes(dst)); + float * y = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); - const int nc = src0->ne[0]; - const int nr = ggml_nelements(src1); + memcpy(y, x, ne00 * sizeof(float)); + // for (int i00 = 0; i00 < ne00; i00++) { + // y[i00] = x[i00]; + // } - GGML_ASSERT( dst->ne[0] == nc); - GGML_ASSERT(src0->nb[0] == sizeof(float)); + const float scale = 1.0f/sqrtf(mean + eps); - for (int i = 0; i < nr; ++i) { - const int r = ((int32_t *) src1->data)[i]; + // if you hit this, likely you got an inf somewhere earlier + assert(scale > 0.0f); - ggml_vec_add_f32(nc, - (float *) ((char *) dst->data + r*dst->nb[1]), - (float *) ((char *) dst->data + r*dst->nb[1]), - (float *) ((char *) src0->data + i*src0->nb[1])); + ggml_vec_scale_f32(ne00, y, scale); + } + } } } -void ggml_compute_forward_get_rows_back( +void ggml_compute_forward_rms_norm( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { - case GGML_TYPE_F16: - { - ggml_compute_forward_get_rows_back_f32_f16(params, dst); - } break; case GGML_TYPE_F32: { - ggml_compute_forward_get_rows_back_f32(params, dst); + ggml_compute_forward_rms_norm_f32(params, dst); } break; default: { GGML_ABORT("fatal error"); } } - - //static bool first = true; - //printf("ne0 = %d, ne1 = %d, ne2 = %d\n", dst->ne[0], dst->ne[1], dst->ne[2]); - //if (first) { - // first = false; - //} else { - // for (int k = 0; k < dst->ne[1]; ++k) { - // for (int j = 0; j < dst->ne[0]/16; ++j) { - // for (int i = 0; i < 16; ++i) { - // printf("%8.4f ", ((float *) dst->data)[k*dst->ne[0] + j*16 + i]); - // } - // printf("\n"); - // } - // printf("\n"); - // } - // printf("\n"); - // exit(0); - //} } -// ggml_compute_forward_diag - -static void ggml_compute_forward_diag_f32( +static void ggml_compute_forward_rms_norm_back_f32( const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - - if (params->ith != 0) { - return; - } + const ggml_tensor * src0 = dst->src[0]; // gradients from forward pass output + const ggml_tensor * src1 = dst->src[1]; // src1 from forward pass - // TODO: handle transposed/permuted matrices + GGML_ASSERT(ggml_are_same_shape(src0, dst) && ggml_are_same_shape(src0, src1)); - GGML_TENSOR_UNARY_OP_LOCALS + GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(src1->nb[0] == sizeof(float)); - GGML_ASSERT(ne00 == ne0); - GGML_ASSERT(ne00 == ne1); - GGML_ASSERT(ne01 == 1); - GGML_ASSERT(ne02 == ne2); - GGML_ASSERT(ne03 == ne3); + const int ith = params->ith; + const int nth = params->nth; - GGML_ASSERT(nb00 == sizeof(float)); - GGML_ASSERT(nb0 == sizeof(float)); + GGML_TENSOR_BINARY_OP_LOCALS - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - for (int i1 = 0; i1 < ne1; i1++) { - float * d = (float *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1); - float * s = (float *)((char *) src0->data + i3*nb03 + i2*nb02); - for (int i0 = 0; i0 < i1; i0++) { - d[i0] = 0; + float eps; + memcpy(&eps, dst->op_params, sizeof(float)); + + // TODO: optimize + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = ith; i01 < ne01; i01 += nth) { + // src1 is same shape as src0 => same indices + const int64_t i11 = i01; + const int64_t i12 = i02; + const int64_t i13 = i03; + + const float * dz = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); + const float * x = (float *) ((char *) src1->data + i11*nb11 + i12*nb12 + i13*nb13); + + ggml_float sum_xx = 0.0; + ggml_float sum_xdz = 0.0; + + for (int64_t i00 = 0; i00 < ne00; i00++) { + sum_xx += (ggml_float)(x[i00] * x[i00]); + sum_xdz += (ggml_float)(x[i00] * dz[i00]); } - d[i1] = s[i1]; - for (int i0 = i1+1; i0 < ne0; i0++) { - d[i0] = 0; + + //const float mean = (float)(sum_xx)/ne00; + const float mean_eps = (float)(sum_xx)/ne00 + eps; + const float sum_eps = (float)(sum_xx) + eps*ne00; + //const float mean_xdz = (float)(sum_xdz)/ne00; + // we could cache rms from forward pass to improve performance. + // to do this implement ggml_rms and compose ggml_rms_norm using ggml_rms. + //const float rms = sqrtf(mean_eps); + const float rrms = 1.0f / sqrtf(mean_eps); + //const float scale = -rrms/(ne00 * mean_eps); // -1/(n*rms**3) + + { + // z = rms_norm(x) + // + // rms_norm(src1) = + // scale( + // src1, + // div( + // 1, + // sqrt( + // add( + // scale( + // sum( + // sqr( + // src1)), + // (1.0/N)), + // eps)))); + + // postorder: + // ## op args grad + // 00 param src1 grad[#00] + // 01 const 1 + // 02 sqr (#00) grad[#02] + // 03 sum (#02) grad[#03] + // 04 const 1/N + // 05 scale (#03, #04) grad[#05] + // 06 const eps + // 07 add (#05, #06) grad[#07] + // 08 sqrt (#07) grad[#08] + // 09 div (#01,#08) grad[#09] + // 10 scale (#00,#09) grad[#10] + // + // backward pass, given grad[#10] + // #10: scale + // grad[#00] += scale(grad[#10],#09) + // grad[#09] += sum(mul(grad[#10],#00)) + // #09: div + // grad[#08] += neg(mul(grad[#09], div(#09,#08))) + // #08: sqrt + // grad[#07] += mul(grad[#08], div(0.5, #08)) + // #07: add + // grad[#05] += grad[#07] + // #05: scale + // grad[#03] += scale(grad[#05],#04) + // #03: sum + // grad[#02] += repeat(grad[#03], #02) + // #02: + // grad[#00] += scale(mul(#00, grad[#02]), 2.0) + // + // substitute and simplify: + // grad[#00] = scale(grad(#10), #09) + scale(mul(#00, grad[#02]), 2.0) + // grad[#02] = repeat(grad[#03], #02) + // grad[#02] = repeat(scale(grad[#05],#04), #02) + // grad[#02] = repeat(scale(grad[#07],#04), #02) + // grad[#02] = repeat(scale(mul(grad[#08], div(0.5, #08)),#04), #02) + // grad[#02] = repeat(scale(mul(neg(mul(grad[#09], div(#09,#08))), div(0.5, #08)),#04), #02) + // grad[#02] = repeat(scale(mul(neg(mul(sum(mul(grad[#10],#00)), div(#09,#08))), div(0.5, #08)),#04), #02) + // grad[#02] = repeat(-(sum(mul(grad[#10],#00)) * div(#09,#08) * div(0.5, #08) * (1/N)), #02) + // grad[#02] = repeat(-(sum(mul(grad[#10],#00)) * div(div(#01,#08),#08) * div(0.5, #08) * (1/N)), #02) + // grad[#02] = repeat(-(sum(mul(grad[#10],#00)) * div(1,#08*#08) * div(0.5, #08) * (1/N)), #02) + // grad[#02] = repeat(-(sum(mul(grad[#10],#00)) * div(1,#07) * div(0.5, #08) * (1/N)), #02) + // grad[#00] = scale(grad(#10), #09) + scale(mul(#00, grad[#02]), 2.0) + // grad[#00] = scale(grad(#10), #09) + scale(mul(#00, repeat(-(sum(mul(grad[#10],#00)) * div(1,#07) * div(0.5, #08) * (1/N)), #02)), 2.0) + // grad[#00] = scale(grad(#10), #09) + scale(scale(#00, -(sum(mul(grad[#10],#00)) * div(1,#07) * div(0.5, #08) * (1/N))), 2.0) + // grad[#00] = scale(grad(#10), #09) + scale(#00, -(sum(mul(grad[#10],#00)) * div(1,#07) * div(1,#08) * (1/N))) + // grad[#00] = scale(grad(#10), #09) + scale(#00, sum(mul(grad[#10],#00)) * div(1,#07*#08) * (-1/N)) + // grad[#00] = scale(grad(#10), #09) + scale(#00, sum(mul(grad[#10],#00)) * div(1,#07*#08) * (-1/N)) + // grad[#00] = scale(grad(#10), #09) + scale(#00, sum(mul(grad[#10],#00)) * div(1,mean_eps*rms) * (-1/N)) + // grad[#00] = scale(grad(#10), #09) + scale(#00, sum(mul(grad[#10],#00)) * div(-1,rms*N*mean_eps)) + // grad[#00] = scale(grad(#10), #09) + scale(#00, sum(mul(grad[#10],#00)) * div(-1,rms*N*(sum_xx/N+eps))) + // grad[#00] = scale(grad(#10), #09) + scale(#00, sum(mul(grad[#10],#00)) * div(-1,rms*N*sum_xx+rms*N*eps)) + // grad[#00] = scale(dz, rrms) + scale(x, sum(mul(dz,x)) * div(-1,rms*N*mean_eps)) + // grad[#00] = scale(dz, rrms) + scale(x, sum_xdz * div(-1,rms*N*mean_eps)) + // a = b*c + d*e + // a = b*c*f/f + d*e*f/f + // a = (b*c*f + d*e*f)*(1/f) + // a = (b*c*(1/c) + d*e*(1/c))*(1/(1/c)) + // a = (b + d*e/c)*c + // b = dz, c = rrms, d = x, e = sum_xdz * div(-1,rms*N*mean_eps) + // a = (dz + x*sum_xdz * div(-1,rms*N*mean_eps)/rrms)*rrms + // a = (dz + x*sum_xdz * div(-1,rms*N*mean_eps)*rms)*rrms + // a = (dz + x*sum_xdz * div(-rms,rms*N*mean_eps))*rrms + // a = (dz + x*sum_xdz * div(-1,N*mean_eps))*rrms + // a = (dz + x*div(-sum_xdz,N*mean_eps))*rrms + // a = (dz + x*div(-mean_xdz,mean_eps))*rrms + // grad[#00] = scale(dz + scale(x, div(-mean_xdz,mean_eps)),rrms) + // grad[#00] = scale(dz + scale(x, -mean_xdz/mean_eps),rrms) + // dx = scale(dz + scale(x, -mean_xdz/mean_eps),rrms) } + // dx = scale(dz + scale(x, -mean_xdz/mean_eps),rrms) + // post-order: + // dx := x + // dx := scale(dx,-mean_xdz/mean_eps) + // dx := add(dx, dz) + // dx := scale(dx, rrms) + float * dx = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); + + // dx[i00] = (x*(-sum_xdz/sum_eps) + dz) / sqrtf(mean_eps) + ggml_vec_cpy_f32 (ne00, dx, x); + // ggml_vec_scale_f32(ne00, dx, -mean_xdz/mean_eps); + ggml_vec_scale_f32(ne00, dx, (float)(-sum_xdz)/sum_eps); + ggml_vec_acc_f32 (ne00, dx, dz); + ggml_vec_scale_f32(ne00, dx, rrms); } } } } -void ggml_compute_forward_diag( +void ggml_compute_forward_rms_norm_back( const ggml_compute_params * params, ggml_tensor * dst) { @@ -5135,7 +6021,7 @@ void ggml_compute_forward_diag( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_diag_f32(params, dst); + ggml_compute_forward_rms_norm_back_f32(params, dst); } break; default: { @@ -5144,68 +6030,93 @@ void ggml_compute_forward_diag( } } -// ggml_compute_forward_diag_mask_inf +// ggml_compute_forward_group_norm -static void ggml_compute_forward_diag_mask_f32( - const ggml_compute_params * params, - ggml_tensor * dst, - const float value) { +static void ggml_compute_forward_group_norm_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + + GGML_ASSERT(src0->nb[0] == sizeof(float)); + const int ith = params->ith; const int nth = params->nth; - const int n_past = ((int32_t *) dst->op_params)[0]; - const bool inplace = src0->data == dst->data; + GGML_TENSOR_UNARY_OP_LOCALS - GGML_ASSERT(n_past >= 0); + // TODO: optimize - if (!inplace) { - if (ith == 0) { - // memcpy needs to be synchronized across threads to avoid race conditions. - // => do it in INIT phase - GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src0)); - GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src0)); - memcpy( - ((char *) dst->data), - ((char *) src0->data), - ggml_nbytes(dst)); + float eps; + memcpy(&eps, dst->op_params + 1, sizeof(float)); + + int n_channels = src0->ne[2]; + int n_groups = dst->op_params[0]; + int n_channels_per_group = (n_channels + n_groups - 1) / n_groups; + for (int i = ith; i < n_groups; i += nth) { + int start = i * n_channels_per_group; + int end = start + n_channels_per_group; + if (end > n_channels) { + end = n_channels; } - ggml_barrier(params->threadpool); - } + int step = end - start; - // TODO: handle transposed/permuted matrices + for (int64_t i03 = 0; i03 < ne03; i03++) { + ggml_float sum = 0.0; + for (int64_t i02 = start; i02 < end; i02++) { + for (int64_t i01 = 0; i01 < ne01; i01++) { + const float * x = (float *)((char *) src0->data + i01 * nb01 + i02 * nb02 + i03 * nb03); - const int n = ggml_nrows(src0); - const int nc = src0->ne[0]; - const int nr = src0->ne[1]; - const int nz = n/nr; + ggml_float sumr = 0.0; + for (int64_t i00 = 0; i00 < ne00; i00++) { + sumr += (ggml_float)x[i00]; + } + sum += sumr; + } + } + const float mean = sum / (ne00 * ne01 * step); - GGML_ASSERT( dst->nb[0] == sizeof(float)); - GGML_ASSERT(src0->nb[0] == sizeof(float)); + ggml_float sum2 = 0.0; + for (int64_t i02 = start; i02 < end; i02++) { + for (int64_t i01 = 0; i01 < ne01; i01++) { + const float * x = (float *)((char *) src0->data + i01 * nb01 + i02 * nb02 + i03 * nb03); - for (int k = 0; k < nz; k++) { - for (int j = ith; j < nr; j += nth) { - for (int i = n_past; i < nc; i++) { - if (i > n_past + j) { - *(float *)((char *) dst->data + k*dst->nb[2] + j*dst->nb[1] + i*dst->nb[0]) = value; + float * y = (float *)((char *) dst->data + i01 * nb1 + i02 * nb2 + i03 * nb3); + + ggml_float sumr = 0.0; + for (int64_t i00 = 0; i00 < ne00; i00++) { + float v = x[i00] - mean; + y[i00] = v; + sumr += (ggml_float)(v * v); + } + sum2 += sumr; + } + } + const float variance = sum2 / (ne00 * ne01 * step); + const float scale = 1.0f / sqrtf(variance + eps); + + for (int64_t i02 = start; i02 < end; i02++) { + for (int64_t i01 = 0; i01 < ne01; i01++) { + float * y = (float *)((char *) dst->data + i01 * nb1 + i02 * nb2 + i03 * nb3); + ggml_vec_scale_f32(ne00, y, scale); } } } } } -void ggml_compute_forward_diag_mask_inf( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_group_norm( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_diag_mask_f32(params, dst, -INFINITY); + ggml_compute_forward_group_norm_f32(params, dst); } break; default: { @@ -5214,16 +6125,61 @@ void ggml_compute_forward_diag_mask_inf( } } -void ggml_compute_forward_diag_mask_zero( - const ggml_compute_params * params, - ggml_tensor * dst) { +// ggml_compute_forward_l2_norm + +static void ggml_compute_forward_l2_norm_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + + GGML_ASSERT(src0->nb[0] == sizeof(float)); + + const int ith = params->ith; + const int nth = params->nth; + + GGML_TENSOR_UNARY_OP_LOCALS + + float eps; + memcpy(&eps, dst->op_params, sizeof(float)); + + GGML_ASSERT(eps >= 0.0f); + + // TODO: optimize + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = ith; i01 < ne01; i01 += nth) { + const float * x = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); + + ggml_float sum = 0.0; + for (int64_t i00 = 0; i00 < ne00; i00++) { + sum += (ggml_float)(x[i00] * x[i00]); + } + + float * y = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); + + memcpy(y, x, ne00 * sizeof(float)); + + const float scale = 1.0f/fmaxf(sqrtf(sum), eps); + + ggml_vec_scale_f32(ne00, y, scale); + } + } + } +} + +void ggml_compute_forward_l2_norm( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_diag_mask_f32(params, dst, 0); + ggml_compute_forward_l2_norm_f32(params, dst); } break; default: { @@ -5232,235 +6188,260 @@ void ggml_compute_forward_diag_mask_zero( } } -// ggml_compute_forward_soft_max +// ggml_compute_forward_out_prod -static void ggml_compute_forward_soft_max_f32( +static void ggml_compute_forward_out_prod_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - const ggml_tensor * src2 = dst->src[2]; - - assert(ggml_is_contiguous(dst)); - assert(ggml_are_same_shape(src0, dst)); - float scale = 1.0f; - float max_bias = 0.0f; + GGML_TENSOR_BINARY_OP_LOCALS - memcpy(&scale, (float *) dst->op_params + 0, sizeof(float)); - memcpy(&max_bias, (float *) dst->op_params + 1, sizeof(float)); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); const int ith = params->ith; const int nth = params->nth; - GGML_TENSOR_UNARY_OP_LOCALS + GGML_ASSERT(ne0 == ne00); + GGML_ASSERT(ne1 == ne10); + GGML_ASSERT(ne2 == ne12); + GGML_ASSERT(ne3 == ne13); - const int64_t nb11 = src1 ? src1->nb[1] : 1; - const int64_t nb12 = src1 ? src1->nb[2] : 1; - const int64_t nb13 = src1 ? src1->nb[3] : 1; + GGML_ASSERT(ne2 % ne02 == 0); + GGML_ASSERT(ne3 % ne03 == 0); - const int64_t ne12 = src1 ? src1->ne[2] : 1; - const int64_t ne13 = src1 ? src1->ne[3] : 1; + // we don't support permuted src0 or src1 + GGML_ASSERT(nb00 == sizeof(float)); - // TODO: is this supposed to be ceil instead of floor? - // https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py#L370 - const uint32_t n_head = ne02; - const uint32_t n_head_log2 = 1u << (uint32_t) floor(log2(n_head)); + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + // GGML_ASSERT(nb0 <= nb1); + // GGML_ASSERT(nb1 <= nb2); + // GGML_ASSERT(nb2 <= nb3); - const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); - const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + // nb01 >= nb00 - src0 is not transposed + // compute by src0 rows - float * wp = (float *) params->wdata + (ne00 + CACHE_LINE_SIZE_F32) * ith; + if (ith == 0) { + ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0); + } + ggml_barrier(params->threadpool); - const bool use_f16 = (src1 && src1->type == GGML_TYPE_F16); + // dst[:,:,:,:] = 0 + // for i2,i3: + // for i1: + // for i01: + // for i0: + // dst[i0,i1,i2,i3] += src0[i0,i01,i2,i3] * src1[i1,i01,i2,i3] - // sinks - const float * sk = src2 ? (float *)((char *) src2->data) : nullptr; + // parallelize by last three dimensions - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = ith; i01 < ne01; i01 += nth) { - const int64_t i11 = i01; - const int64_t i12 = i02%ne12; - const int64_t i13 = i03%ne13; + // total rows in dst + const int64_t nr = ne1*ne2*ne3; - // ALiBi - const uint32_t h = i02; // head - const float slope = (max_bias > 0.0f) ? h < n_head_log2 ? powf(m0, h + 1) : powf(m1, 2*(h - n_head_log2) + 1) : 1.0f; + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; - float * sp = (float *)((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); - float * dp = (float *)((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); + // row range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); - // broadcast the mask across rows - ggml_fp16_t * mp_f16 = src1 ? (ggml_fp16_t *)((char *) src1->data + i11*nb11 + i12*nb12 + i13*nb13) : NULL; - float * mp_f32 = src1 ? (float *)((char *) src1->data + i11*nb11 + i12*nb12 + i13*nb13) : NULL; + // block-tiling attempt + const int64_t blck_0 = MAX(GGML_VEC_MAD_UNROLL, 32); + const int64_t blck_1 = 16; - ggml_vec_cpy_f32 (ne00, wp, sp); - ggml_vec_scale_f32(ne00, wp, scale); - if (mp_f32) { - if (use_f16) { - for (int i = 0; i < ne00; ++i) { - wp[i] += slope*GGML_CPU_FP16_TO_FP32(mp_f16[i]); - } - } else { - for (int i = 0; i < ne00; ++i) { - wp[i] += slope*mp_f32[i]; - } - } - } + // dps == dst per src0, used for group query attention + const int64_t dps2 = ne2 / ne02; + const int64_t dps3 = ne3 / ne03; -#ifndef NDEBUG - for (int i = 0; i < ne00; ++i) { - //printf("p[%d] = %f\n", i, p[i]); - assert(!isnan(wp[i])); - } -#endif // NDEBUG + for (int64_t bir = ir0; bir < ir1; bir += blck_1) { + const int64_t bir1 = MIN(bir + blck_1, ir1); + for (int64_t bi01 = 0; bi01 < ne01; bi01 += blck_0) { + const int64_t bne01 = MIN(bi01 + blck_0, ne01); + for (int64_t ir = bir; ir < bir1; ++ir) { + // dst indices + const int64_t i3 = ir/(ne2*ne1); + const int64_t i2 = (ir - i3*ne2*ne1)/ne1; + const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1); - float max = -INFINITY; - ggml_vec_max_f32(ne00, &max, wp); + const int64_t i02 = i2 / dps2; + const int64_t i03 = i3 / dps3; - // if we have sinks, make a correction as if they were included in the softmax - if (sk) { - max = MAX(max, sk[i02]); + //const int64_t i10 = i1; + const int64_t i12 = i2; + const int64_t i13 = i3; + +#if GGML_VEC_MAD_UNROLL > 2 + const int64_t bne01_unroll = bne01 - (bne01 % GGML_VEC_MAD_UNROLL); + for (int64_t i01 = bi01; i01 < bne01_unroll; i01 += GGML_VEC_MAD_UNROLL) { + const int64_t i11 = i01; + + float * s0 = (float *) ((char *) src0->data + ( i01*nb01 + i02*nb02 + i03*nb03)); + float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13)); + float * d = (float *) ((char *) dst->data + ( i1*nb1 + i2*nb2 + i3*nb3)); + + ggml_vec_mad_f32_unroll(ne0, nb01, nb11, d, s0, s1); } + for (int64_t i01 = bne01_unroll; i01 < bne01; ++i01) { + const int64_t i11 = i01; - ggml_float sum = ggml_vec_soft_max_f32(ne00, dp, wp, max); - assert(sum > 0.0); + float * s0 = (float *) ((char *) src0->data + ( i01*nb01 + i02*nb02 + i03*nb03)); + float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13)); + float * d = (float *) ((char *) dst->data + ( i1*nb1 + i2*nb2 + i3*nb3)); - if (sk) { - sum += (ggml_float) expf(sk[i02] - max); + ggml_vec_mad_f32(ne0, d, s0, *s1); } +#else + for (int64_t i01 = bi01; i01 < bne01; ++i01) { + const int64_t i11 = i01; - sum = 1.0/sum; - ggml_vec_scale_f32(ne00, dp, sum); + float * s0 = (float *) ((char *) src0->data + ( i01*nb01 + i02*nb02 + i03*nb03)); + float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13)); + float * d = (float *) ((char *) dst->data + ( i1*nb1 + i2*nb2 + i3*nb3)); -#ifndef NDEBUG - for (int i = 0; i < ne00; ++i) { - assert(!isnan(dp[i])); - assert(!isinf(dp[i])); + ggml_vec_mad_f32(ne0, d, s0, *s1); } -#endif // NDEBUG +#endif } } } } -void ggml_compute_forward_soft_max( +static void ggml_compute_forward_out_prod_q_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_soft_max_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} + GGML_TENSOR_BINARY_OP_LOCALS; + const int ith = params->ith; + const int nth = params->nth; -// ggml_compute_forward_soft_max_ext_back + const ggml_type type = src0->type; + ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; -static void ggml_compute_forward_soft_max_ext_back_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { + GGML_ASSERT(ne02 == ne12); + GGML_ASSERT(ne03 == ne13); + GGML_ASSERT(ne2 == ne12); + GGML_ASSERT(ne3 == ne13); - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; + // we don't support permuted src0 dim0 + GGML_ASSERT(nb00 == ggml_type_size(type)); - GGML_ASSERT(ggml_is_contiguous(src0)); - GGML_ASSERT(ggml_is_contiguous(src1)); - GGML_ASSERT(ggml_is_contiguous(dst)); - GGML_ASSERT(ggml_are_same_shape(src0, dst)); - GGML_ASSERT(ggml_are_same_shape(src1, dst)); - - float scale = 1.0f; - float max_bias = 0.0f; + // dst dim0 cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + // GGML_ASSERT(nb0 <= nb1); + // GGML_ASSERT(nb1 <= nb2); + // GGML_ASSERT(nb2 <= nb3); - memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); - memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + GGML_ASSERT(ne0 == ne00); + GGML_ASSERT(ne1 == ne10); + GGML_ASSERT(ne2 == ne02); + GGML_ASSERT(ne3 == ne03); - GGML_ASSERT(max_bias == 0.0f); + // nb01 >= nb00 - src0 is not transposed + // compute by src0 rows - // TODO: handle transposed/permuted matrices + if (ith == 0) { + ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0); + } + ggml_barrier(params->threadpool); - const int ith = params->ith; - const int nth = params->nth; + // parallelize by last three dimensions - const int nc = src0->ne[0]; - const int nr = ggml_nrows(src0); + // total rows in dst + const int64_t nr = ne1*ne2*ne3; // rows per thread - const int dr = (nr + nth - 1)/nth; + const int64_t dr = (nr + nth - 1)/nth; // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); - for (int i1 = ir0; i1 < ir1; i1++) { - float *dy = (float *)((char *) src0->data + i1*src0->nb[1]); - float *y = (float *)((char *) src1->data + i1*src1->nb[1]); - float *dx = (float *)((char *) dst->data + i1*dst->nb[1]); + // dst[:,:,:,:] = 0 + // for i2,i3: + // for i1: + // for i01: + // for i0: + // dst[i0,i1,i2,i3] += src0[i0,i01,i2,i3] * src1[i1,i01,i2,i3] -#ifndef NDEBUG - for (int i = 0; i < nc; ++i) { - //printf("p[%d] = %f\n", i, p[i]); - assert(!isnan(dy[i])); - assert(!isnan(y[i])); - } -#endif // NDEBUG - // Jii = yi - yi*yi - // Jij = -yi*yj - // J = diag(y)-y.T*y - // dx = J * dy - // dxk = sum_i(Jki * dyi) - // dxk = sum_i(-yk*yi * dyi) - (-yk*yk)*dyk + (yk - yk*yk)*dyk - // dxk = sum_i(-yk*yi * dyi) + yk*yk*dyk + yk*dyk - yk*yk*dyk - // dxk = sum_i(-yk*yi * dyi) + yk*dyk - // dxk = -yk * sum_i(yi * dyi) + yk*dyk - // dxk = -yk * dot(y, dy) + yk*dyk - // dxk = yk * (- dot(y, dy) + dyk) - // dxk = yk * (dyk - dot(y, dy)) - // - // post-order: - // dot_y_dy := dot(y, dy) - // dx := dy - // dx := dx - dot_y_dy - // dx := dx * y + float * wdata = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32) * ith; - // linear runtime, no additional memory - float dot_y_dy = 0; - ggml_vec_dot_f32 (nc, &dot_y_dy, 0, y, 0, dy, 0, 1); - ggml_vec_cpy_f32 (nc, dx, dy); - ggml_vec_acc1_f32 (nc, dx, -dot_y_dy); - ggml_vec_mul_f32 (nc, dx, dx, y); - ggml_vec_scale_f32(nc, dx, scale); + for (int64_t ir = ir0; ir < ir1; ++ir) { + // dst indices + const int64_t i3 = ir/(ne2*ne1); + const int64_t i2 = (ir - i3*ne2*ne1)/ne1; + const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1); -#ifndef NDEBUG - for (int i = 0; i < nc; ++i) { - assert(!isnan(dx[i])); - assert(!isinf(dx[i])); + const int64_t i02 = i2; + const int64_t i03 = i3; + + //const int64_t i10 = i1; + const int64_t i12 = i2; + const int64_t i13 = i3; + + for (int64_t i01 = 0; i01 < ne01; ++i01) { + const int64_t i11 = i01; + + float * s0 = (float *) ((char *) src0->data + ( i01*nb01 + i02*nb02 + i03*nb03)); + float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13)); + float * d = (float *) ((char *) dst->data + ( i1*nb1 + i2*nb2 + i3*nb3)); + + dequantize_row_q(s0, wdata, ne0); + ggml_vec_mad_f32(ne0, d, wdata, *s1); } -#endif // NDEBUG } } -void ggml_compute_forward_soft_max_ext_back( +void ggml_compute_forward_out_prod( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_TQ1_0: + case GGML_TYPE_TQ2_0: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ2_S: + { + ggml_compute_forward_out_prod_q_f32(params, dst); + } break; + case GGML_TYPE_F16: + { + GGML_ABORT("fatal error"); // todo + // ggml_compute_forward_out_prod_f16_f32(params, dst); + } case GGML_TYPE_F32: { - ggml_compute_forward_soft_max_ext_back_f32(params, dst); + ggml_compute_forward_out_prod_f32(params, dst); } break; default: { @@ -5469,82 +6450,61 @@ void ggml_compute_forward_soft_max_ext_back( } } -// ggml_compute_forward_clamp +// ggml_compute_forward_scale -static void ggml_compute_forward_clamp_f32( +static void ggml_compute_forward_scale_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - float min; - float max; - memcpy(&min, (float *) dst->op_params + 0, sizeof(float)); - memcpy(&max, (float *) dst->op_params + 1, sizeof(float)); - - const int ith = params->ith; - const int nth = params->nth; - - const int n = ggml_nrows(src0); - const int nc = src0->ne[0]; - - const size_t nb00 = src0->nb[0]; - const size_t nb01 = src0->nb[1]; - - const size_t nb0 = dst->nb[0]; - const size_t nb1 = dst->nb[1]; - - GGML_ASSERT( nb0 == sizeof(float)); - GGML_ASSERT(nb00 == sizeof(float)); - - for (int j = ith; j < n; j += nth) { - float * dst_ptr = (float *) ((char *) dst->data + j*nb1); - float * src0_ptr = (float *) ((char *) src0->data + j*nb01); - - for (int i = 0; i < nc; i++) { - dst_ptr[i] = MAX(MIN(src0_ptr[i], max), min); - } - } -} - -static void ggml_compute_forward_clamp_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); - const ggml_tensor * src0 = dst->src[0]; + float s; // scale factor + float b; // bias - float min; - float max; - memcpy(&min, (float *) dst->op_params + 0, sizeof(float)); - memcpy(&max, (float *) dst->op_params + 1, sizeof(float)); + memcpy(&s, (float *) dst->op_params + 0, sizeof(float)); + memcpy(&b, (float *) dst->op_params + 1, sizeof(float)); const int ith = params->ith; const int nth = params->nth; - const int n = ggml_nrows(src0); const int nc = src0->ne[0]; + const int nr = ggml_nrows(src0); - const size_t nb00 = src0->nb[0]; - const size_t nb01 = src0->nb[1]; + // rows per thread + const int dr = (nr + nth - 1)/nth; - const size_t nb0 = dst->nb[0]; - const size_t nb1 = dst->nb[1]; + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - GGML_ASSERT( nb0 == sizeof(ggml_fp16_t)); - GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); + const size_t nb01 = src0->nb[1]; - for (int j = ith; j < n; j += nth) { - ggml_fp16_t * dst_ptr = (ggml_fp16_t *) ((char *) dst->data + j*nb1); - ggml_fp16_t * src0_ptr = (ggml_fp16_t *) ((char *) src0->data + j*nb01); + const size_t nb1 = dst->nb[1]; - for (int i = 0; i < nc; i++) { - float v = GGML_CPU_FP16_TO_FP32(src0_ptr[i]); - dst_ptr[i] = GGML_CPU_FP32_TO_FP16(MAX(MIN(v, max), min)); + if (b == 0.0f) { + for (int i1 = ir0; i1 < ir1; i1++) { + if (dst->data != src0->data) { + // src0 is same shape as dst => same indices + // TODO: add x parameter to ggml_vec_scale_f32 and remove this memcpy + memcpy((char *)dst->data + i1*nb1, (char *)src0->data + i1*nb01, nc * sizeof(float)); + } + ggml_vec_scale_f32(nc, (float *) ((char *) dst->data + i1*nb1), s); + } + } else { + for (int i1 = ir0; i1 < ir1; i1++) { + ggml_vec_mad1_f32(nc, + (float *) ((char *) dst->data + i1*nb1), + (float *) ((char *) src0->data + i1*nb1), + s, b); } } } -void ggml_compute_forward_clamp( +void ggml_compute_forward_scale( const ggml_compute_params * params, ggml_tensor * dst) { @@ -5553,228 +6513,138 @@ void ggml_compute_forward_clamp( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_clamp_f32(params, dst); + ggml_compute_forward_scale_f32(params, dst); } break; - case GGML_TYPE_F16: - { - ggml_compute_forward_clamp_f16(params, dst); - } break; - case GGML_TYPE_BF16: - case GGML_TYPE_Q1_0: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q8_0: - case GGML_TYPE_Q8_1: - case GGML_TYPE_MXFP4: - case GGML_TYPE_NVFP4: - case GGML_TYPE_Q2_K: - case GGML_TYPE_Q3_K: - case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: - case GGML_TYPE_Q6_K: - case GGML_TYPE_TQ1_0: - case GGML_TYPE_TQ2_0: - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ2_S: - case GGML_TYPE_Q8_K: - case GGML_TYPE_I8: - case GGML_TYPE_I16: - case GGML_TYPE_I32: - case GGML_TYPE_I64: - case GGML_TYPE_F64: - case GGML_TYPE_COUNT: + default: { GGML_ABORT("fatal error"); } } } -// ggml_compute_forward_rope +// ggml_compute_forward_set -static float rope_yarn_ramp(const float low, const float high, const int i0) { - const float y = (i0 / 2 - low) / MAX(0.001f, high - low); - return 1 - MIN(1, MAX(0, y)); -} +static void ggml_compute_forward_set_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { -// YaRN algorithm based on LlamaYaRNScaledRotaryEmbedding.py from https://github.com/jquesnelle/yarn -// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng. -static void rope_yarn( - float theta_extrap, float freq_scale, float corr_dims[2], int64_t i0, float ext_factor, float mscale, - float * cos_theta, float * sin_theta) { - // Get n-d rotational scaling corrected for extrapolation - float theta_interp = freq_scale * theta_extrap; - float theta = theta_interp; - if (ext_factor != 0.0f) { - float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], i0) * ext_factor; - theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix; + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - // Get n-d magnitude scaling corrected for interpolation - mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale); - } - *cos_theta = cosf(theta) * mscale; - *sin_theta = sinf(theta) * mscale; -} + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src0)); -static void ggml_rope_cache_init( - float theta_base, float freq_scale, const float * freq_factors, float corr_dims[2], int64_t ne0, float ext_factor, float mscale, - float * cache, float sin_sign, float theta_scale) { - // ref: https://github.com/jquesnelle/yarn/blob/master/scaled_rope/LlamaYaRNScaledRotaryEmbedding.py - float theta = theta_base; - for (int64_t i0 = 0; i0 < ne0; i0 += 2) { - const float ff = freq_factors ? freq_factors[i0/2] : 1.0f; - rope_yarn( - theta/ff, freq_scale, corr_dims, i0, ext_factor, mscale, &cache[i0 + 0], &cache[i0 + 1] - ); - cache[i0 + 1] *= sin_sign; + // view src0 and dst with these strides and data offset inbytes during set + // nb0 is implicitly element_size because src0 and dst are contiguous + size_t nb1 = ((int32_t *) dst->op_params)[0]; + size_t nb2 = ((int32_t *) dst->op_params)[1]; + size_t nb3 = ((int32_t *) dst->op_params)[2]; + size_t offset = ((int32_t *) dst->op_params)[3]; + bool inplace = (bool) ((int32_t *) dst->op_params)[4]; - theta *= theta_scale; + if (!inplace) { + if (params->ith == 0) { + // memcpy needs to be synchronized across threads to avoid race conditions. + // => do it in INIT phase + memcpy( + ((char *) dst->data), + ((char *) src0->data), + ggml_nbytes(dst)); + } + ggml_barrier(params->threadpool); } -} - -static void ggml_mrope_cache_init( - float theta_base_t, float theta_base_h, float theta_base_w, float theta_base_e, int sections[4], bool is_imrope, bool indep_sects, - float freq_scale, const float * freq_factors, float corr_dims[2], int64_t ne0, float ext_factor, float mscale, - float * cache, float sin_sign, float theta_scale) { - // ref: https://github.com/jquesnelle/yarn/blob/master/scaled_rope/LlamaYaRNScaledRotaryEmbedding.py - float theta_t = theta_base_t; - float theta_h = theta_base_h; - float theta_w = theta_base_w; - float theta_e = theta_base_e; // extra position id for vision encoder - int sect_dims = sections[0] + sections[1] + sections[2] + sections[3]; - int sec_w = sections[1] + sections[0]; - int sec_e = sections[2] + sec_w; - GGML_ASSERT(sect_dims <= ne0); - for (int64_t i0 = 0; i0 < ne0; i0 += 2) { - const float ff = freq_factors ? freq_factors[i0/2] : 1.0f; + const int ith = params->ith; + const int nth = params->nth; - int sector = (i0 / 2) % sect_dims; - if (indep_sects) { - // compute theta independently for each dim sections - // (i.e. reset corresponding theta when `i0` go from one section to another) - if (sector == 0) { - theta_t = theta_base_t; - } - else if (sector == sections[0]) { - theta_h = theta_base_h;; - } - else if (sector == sec_w) { - theta_w = theta_base_w; - } - else if (sector == sec_e) { - theta_e = theta_base_e; - } - } + const int nr = ggml_nrows(src1); + const int nc = src1->ne[0]; - float theta = theta_t; - if (is_imrope) { // qwen3vl apply interleaved mrope - if (sector % 3 == 1 && sector < 3 * sections[1]) { - theta = theta_h; - } else if (sector % 3 == 2 && sector < 3 * sections[2]) { - theta = theta_w; - } else if (sector % 3 == 0 && sector < 3 * sections[0]) { - theta = theta_t; - } else { - theta = theta_e; - } - } else { - if (sector >= sections[0] && sector < sec_w) { - theta = theta_h; - } - else if (sector >= sec_w && sector < sec_w + sections[2]) { - theta = theta_w; - } - else if (sector >= sec_w + sections[2]) { - theta = theta_e; - } - } + GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) + GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) - rope_yarn( - theta/ff, freq_scale, corr_dims, i0, ext_factor, mscale, &cache[i0 + 0], &cache[i0 + 1] - ); - cache[i0 + 1] *= sin_sign; + // src0 and dst as viewed during set + const size_t nb0 = ggml_element_size(src0); - theta_t *= theta_scale; - theta_w *= theta_scale; - theta_h *= theta_scale; - theta_e *= theta_scale; - } -} + const int im0 = (ne10 == 0 ? 0 : ne10-1); + const int im1 = (ne11 == 0 ? 0 : ne11-1); + const int im2 = (ne12 == 0 ? 0 : ne12-1); + const int im3 = (ne13 == 0 ? 0 : ne13-1); + GGML_ASSERT(offset + im0*nb0 + im1*nb1 + im2*nb2 + im3*nb3 <= ggml_nbytes(dst)); -template -static void rotate_pairs(const int64_t n, const int64_t n_offset, const float * cache, const T * src_data, T * dst_data, const int scale = 2) { - for (int64_t i0 = 0; i0 < n; i0 += 2) { - const int64_t ic = i0/scale; // hack for GGML_ROPE_TYPE_NORMAL, where we need ic = i0; for all other cases, ic = i0/2 + GGML_ASSERT(nb10 == sizeof(float)); - const float cos_theta = cache[i0 + 0]; - const float sin_theta = cache[i0 + 1]; + // rows per thread + const int dr = (nr + nth - 1)/nth; - const T * const src = src_data + ic; - T * dst = dst_data + ic; + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - const float x0 = type_conversion_table::to_f32(src[0]); - const float x1 = type_conversion_table::to_f32(src[n_offset]); + for (int ir = ir0; ir < ir1; ++ir) { + // src0 and dst are viewed with shape of src1 and offset + // => same indices + const int i3 = ir/(ne12*ne11); + const int i2 = (ir - i3*ne12*ne11)/ne11; + const int i1 = (ir - i3*ne12*ne11 - i2*ne11); - dst[0] = type_conversion_table::from_f32(x0*cos_theta - x1*sin_theta); - dst[n_offset] = type_conversion_table::from_f32(x0*sin_theta + x1*cos_theta); - } + ggml_vec_cpy_f32(nc, + (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + offset), + (float *) ((char *) src1->data + i3*nb13 + i2*nb12 + i1*nb11)); + } } -template //float or ggml_fp16_t -static void ggml_compute_forward_rope_flt( +static void ggml_compute_forward_set_i32( const ggml_compute_params * params, - ggml_tensor * dst, - const bool forward) { + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - const ggml_tensor * src2 = dst->src[2]; - GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); - GGML_ASSERT(src1->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src0)); - float freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow; - int sections[4]; + // view src0 and dst with these strides and data offset inbytes during set + // nb0 is implicitly element_size because src0 and dst are contiguous + size_t nb1 = ((int32_t *) dst->op_params)[0]; + size_t nb2 = ((int32_t *) dst->op_params)[1]; + size_t nb3 = ((int32_t *) dst->op_params)[2]; + size_t offset = ((int32_t *) dst->op_params)[3]; + bool inplace = (bool) ((int32_t *) dst->op_params)[4]; - //const int n_past = ((int32_t *) dst->op_params)[0]; - const int n_dims = ((int32_t *) dst->op_params)[1]; - const int mode = ((int32_t *) dst->op_params)[2]; - //const int n_ctx = ((int32_t *) dst->op_params)[3]; - const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; + if (!inplace) { + if (params->ith == 0) { + // memcpy needs to be synchronized across threads to avoid race conditions. + // => do it in INIT phase + memcpy( + ((char *) dst->data), + ((char *) src0->data), + ggml_nbytes(dst)); + } + ggml_barrier(params->threadpool); + } - memcpy(&freq_base, (int32_t *) dst->op_params + 5, sizeof(float)); - memcpy(&freq_scale, (int32_t *) dst->op_params + 6, sizeof(float)); - memcpy(&ext_factor, (int32_t *) dst->op_params + 7, sizeof(float)); - memcpy(&attn_factor, (int32_t *) dst->op_params + 8, sizeof(float)); - memcpy(&beta_fast, (int32_t *) dst->op_params + 9, sizeof(float)); - memcpy(&beta_slow, (int32_t *) dst->op_params + 10, sizeof(float)); - memcpy(§ions, (int32_t *) dst->op_params + 11, sizeof(int)*4); + const int ith = params->ith; + const int nth = params->nth; - GGML_TENSOR_UNARY_OP_LOCALS + const int nr = ggml_nrows(src1); + const int nc = src1->ne[0]; - //printf("ne0: %d, ne1: %d, ne2: %d, ne3: %d\n", ne0, ne1, ne2, ne3); - //printf("n_past = %d, ne2 = %d\n", n_past, ne2); + GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) + GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) - GGML_ASSERT(nb0 == nb00); - GGML_ASSERT(nb0 == sizeof(T)); + // src0 and dst as viewed during set + const size_t nb0 = ggml_element_size(src0); - const int ith = params->ith; - const int nth = params->nth; + const int im0 = (ne10 == 0 ? 0 : ne10-1); + const int im1 = (ne11 == 0 ? 0 : ne11-1); + const int im2 = (ne12 == 0 ? 0 : ne12-1); + const int im3 = (ne13 == 0 ? 0 : ne13-1); - const int nr = ggml_nrows(dst); + GGML_ASSERT(offset + im0*nb0 + im1*nb1 + im2*nb2 + im3*nb3 <= ggml_nbytes(dst)); - GGML_ASSERT(n_dims <= ne0); - GGML_ASSERT(n_dims % 2 == 0); + GGML_ASSERT(nb10 == sizeof(int32_t)); // rows per thread const int dr = (nr + nth - 1)/nth; @@ -5783,116 +6653,61 @@ static void ggml_compute_forward_rope_flt( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - // row index used to determine which thread to use - int ir = 0; - - const float theta_scale = powf(freq_base, -2.0f/n_dims); + for (int ir = ir0; ir < ir1; ++ir) { + // src0 and dst are viewed with shape of src1 and offset + // => same indices + const int i3 = ir/(ne12*ne11); + const int i2 = (ir - i3*ne12*ne11)/ne11; + const int i1 = (ir - i3*ne12*ne11 - i2*ne11); - float corr_dims[2]; - ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); - - const bool is_imrope = mode == GGML_ROPE_TYPE_IMROPE; // qwen3vl apply interleaved mrope - const bool mrope_used = mode & GGML_ROPE_TYPE_MROPE; // ggml_rope_multi, note: also true for vision (24 & 8 == true) and for imrope - const bool is_vision = mode == GGML_ROPE_TYPE_VISION; - - if (mrope_used) { - GGML_ASSERT(sections[0] > 0 || sections[1] > 0 || sections[2] > 0); - } - - if (is_vision) { - GGML_ASSERT(n_dims == ne0/2); - } - - const float * freq_factors = NULL; - if (src2 != NULL) { - GGML_ASSERT(src2->type == GGML_TYPE_F32); - GGML_ASSERT(src2->ne[0] >= n_dims / 2); - freq_factors = (const float *) src2->data; - } - - // backward process uses inverse rotation by cos and sin. - // cos and sin build a rotation matrix, where the inverse is the transpose. - // this essentially just switches the sign of sin. - const float sin_sign = forward ? 1.0f : -1.0f; - - const int32_t * pos = (const int32_t *) src1->data; - - int64_t last_i2 = -1; - - for (int64_t i3 = 0; i3 < ne3; i3++) { // batch - for (int64_t i2 = 0; i2 < ne2; i2++) { // seq-len - for (int64_t i1 = 0; i1 < ne1; i1++) { // attn-heads - if (ir++ < ir0) continue; // skip rows mapped to other threads - if (ir > ir1) break; - - float * cache = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32)*ith; - if (last_i2 != i2) { - if (!mrope_used) { - const int64_t p = pos[i2]; - ggml_rope_cache_init(p, freq_scale, freq_factors, corr_dims, ne0, ext_factor, attn_factor, cache, sin_sign, theta_scale); - } - else { - const int64_t p_t = pos[i2]; - const int64_t p_h = pos[i2 + ne2]; - const int64_t p_w = pos[i2 + ne2 * 2]; - const int64_t p_e = pos[i2 + ne2 * 3]; - ggml_mrope_cache_init( - p_t, p_h, p_w, p_e, sections, is_imrope, is_vision, - freq_scale, freq_factors, corr_dims, ne0, ext_factor, attn_factor, cache, sin_sign, theta_scale); - } - - last_i2 = i2; - } - - T * src = (T *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01); - T * dst_data = (T *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1); - - switch (mode) { - case GGML_ROPE_TYPE_NORMAL: - rotate_pairs(n_dims, 1, cache, src, dst_data, 1); - break; - case GGML_ROPE_TYPE_NEOX: - case GGML_ROPE_TYPE_MROPE: - case GGML_ROPE_TYPE_IMROPE: - rotate_pairs(n_dims, n_dims/2, cache, src, dst_data); - break; - case GGML_ROPE_TYPE_VISION: - rotate_pairs(ne0, n_dims, cache, src, dst_data); - break; - default: - GGML_ABORT("rope type not supported"); - } - - if (!is_vision) { - // fill the remain channels with data from src tensor - for (int64_t i0 = n_dims; i0 < ne0; i0 += 2) { - const T * const src = (T *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); - T * dst_data = (T *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); - - dst_data[0] = src[0]; - dst_data[1] = src[1]; - } - } - } //attn-heads - } + ggml_vec_cpy_i32(nc, + (int32_t *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + offset), + (int32_t *) ((char *) src1->data + i3*nb13 + i2*nb12 + i1*nb11)); } } -void ggml_compute_forward_rope( +void ggml_compute_forward_set( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { - case GGML_TYPE_F16: + case GGML_TYPE_F32: { - ggml_compute_forward_rope_flt(params, dst, true); + ggml_compute_forward_set_f32(params, dst); } break; - case GGML_TYPE_F32: + case GGML_TYPE_I32: { - ggml_compute_forward_rope_flt(params, dst, true); + ggml_compute_forward_set_i32(params, dst); } break; + case GGML_TYPE_F16: + case GGML_TYPE_BF16: + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q8_1: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_TQ1_0: + case GGML_TYPE_TQ2_0: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ2_S: default: { GGML_ABORT("fatal error"); @@ -5900,93 +6715,46 @@ void ggml_compute_forward_rope( } } -// ggml_compute_forward_rope_back +// ggml_compute_forward_cpy -void ggml_compute_forward_rope_back( +void ggml_compute_forward_cpy( const ggml_compute_params * params, ggml_tensor * dst) { + ggml_compute_forward_dup(params, dst); +} - const ggml_tensor * src0 = dst->src[0]; +// ggml_compute_forward_cont - switch (src0->type) { - case GGML_TYPE_F16: - { - ggml_compute_forward_rope_flt(params, dst, false); - } break; - case GGML_TYPE_F32: - { - ggml_compute_forward_rope_flt(params, dst, false); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } +void ggml_compute_forward_cont( + const ggml_compute_params * params, + ggml_tensor * dst) { + ggml_compute_forward_dup(params, dst); } -// ggml_compute_forward_conv_transpose_1d +// ggml_compute_forward_get_rows -static void ggml_compute_forward_conv_transpose_1d_f16_f32( +static void ggml_compute_forward_get_rows_q( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(src0->type == GGML_TYPE_F16); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT( dst->type == GGML_TYPE_F32); - GGML_TENSOR_BINARY_OP_LOCALS - const int ith = params->ith; - const int nth = params->nth; - - const int nk = ne00*ne01*ne02; - - GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); - GGML_ASSERT(nb10 == sizeof(float)); - - if (ith == 0) { - memset(params->wdata, 0, params->wsize); - - // permute kernel data (src0) from (K x Cout x Cin) to (Cin x K x Cout) - { - ggml_fp16_t * const wdata = (ggml_fp16_t *) params->wdata + 0; - - for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = 0; i01 < ne01; i01++) { - const ggml_fp16_t * const src = (ggml_fp16_t *)((char *) src0->data + i02*nb02 + i01*nb01); - ggml_fp16_t * dst_data = wdata + i01*ne00*ne02; - for (int64_t i00 = 0; i00 < ne00; i00++) { - dst_data[i00*ne02 + i02] = src[i00]; - } - } - } - } - - // permute source data (src1) from (L x Cin) to (Cin x L) - { - ggml_fp16_t * const wdata = (ggml_fp16_t *) params->wdata + nk; - ggml_fp16_t * dst_data = wdata; - - for (int64_t i11 = 0; i11 < ne11; i11++) { - const float * const src = (float *)((char *) src1->data + i11*nb11); - for (int64_t i10 = 0; i10 < ne10; i10++) { - dst_data[i10*ne11 + i11] = GGML_CPU_FP32_TO_FP16(src[i10]); - } - } - } + const int64_t nc = ne00; + const int64_t nr = ggml_nelements(src1); - // need to zero dst since we are accumulating into it - memset(dst->data, 0, ggml_nbytes(dst)); - } - ggml_barrier(params->threadpool); + const ggml_type type = src0->type; + ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; - const int32_t s0 = ((const int32_t*)(dst->op_params))[0]; + assert(ne0 == nc); + assert(ne02 == ne11); + assert(nb00 == ggml_type_size(type)); + assert(ggml_nrows(dst) == nr); - // total rows in dst - const int nr = ne1; + const int ith = params->ith; + const int nth = params->nth; // rows per thread const int dr = (nr + nth - 1)/nth; @@ -5995,86 +6763,80 @@ static void ggml_compute_forward_conv_transpose_1d_f16_f32( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - ggml_fp16_t * const wdata = (ggml_fp16_t *) params->wdata + 0; - ggml_fp16_t * const wdata_src = wdata + nk; + for (int64_t i = ir0; i < ir1; ++i) { + const int64_t i12 = i/(ne11*ne10); + const int64_t i11 = (i - i12*ne11*ne10)/ne10; + const int64_t i10 = (i - i12*ne11*ne10 - i11*ne10); + const int64_t i01 = *(int32_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12); - for (int i1 = ir0; i1 < ir1; i1++) { - float * dst_data = (float *)((char *) dst->data + i1*nb1); - ggml_fp16_t * wdata_kernel = wdata + i1*ne02*ne00; - for (int i10 = 0; i10 < ne10; i10++) { - const int i1n = i10*ne11; - for (int i00 = 0; i00 < ne00; i00++) { - float v = 0; - ggml_vec_dot_f16(ne02, &v, 0, - (ggml_fp16_t *) wdata_src + i1n, 0, - (ggml_fp16_t *) wdata_kernel + i00*ne02, 0, 1); - dst_data[i10*s0 + i00] += v; - } - } + GGML_ASSERT(i01 >= 0 && i01 < ne01); + + dequantize_row_q( + (const void *) ((char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03), + (float *) ((char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3), nc); } } -static void ggml_compute_forward_conv_transpose_1d_f32( +static void ggml_compute_forward_get_rows_f16( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT( dst->type == GGML_TYPE_F32); - GGML_TENSOR_BINARY_OP_LOCALS + const int64_t nc = ne00; + const int64_t nr = ggml_nelements(src1); + + assert(ne0 == nc); + assert(ne02 == ne11); + assert(nb00 == sizeof(ggml_fp16_t)); + assert(ggml_nrows(dst) == nr); + const int ith = params->ith; const int nth = params->nth; - const int nk = ne00*ne01*ne02; - - GGML_ASSERT(nb00 == sizeof(float)); - GGML_ASSERT(nb10 == sizeof(float)); + // rows per thread + const int dr = (nr + nth - 1)/nth; - if (ith == 0) { - memset(params->wdata, 0, params->wsize); + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - // prepare kernel data (src0) from (K x Cout x Cin) to (Cin x K x Cout) - { - float * const wdata = (float *) params->wdata + 0; + for (int64_t i = ir0; i < ir1; ++i) { + const int64_t i12 = i/(ne11*ne10); + const int64_t i11 = (i - i12*ne11*ne10)/ne10; + const int64_t i10 = (i - i12*ne11*ne10 - i11*ne10); + const int64_t i01 = *(int32_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12); - for (int64_t i02 = 0; i02 < ne02; i02++) { - for (int64_t i01 = 0; i01 < ne01; i01++) { - const float * const src = (float *)((char *) src0->data + i02*nb02 + i01*nb01); - float * dst_data = wdata + i01*ne00*ne02; - for (int64_t i00 = 0; i00 < ne00; i00++) { - dst_data[i00*ne02 + i02] = src[i00]; - } - } - } - } + GGML_ASSERT(i01 >= 0 && i01 < ne01); - // prepare source data (src1) - { - float * const wdata = (float *) params->wdata + nk; - float * dst_data = wdata; + ggml_cpu_fp16_to_fp32( + (const ggml_fp16_t*) ((char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03), + (float *) ((char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3), nc); + } +} - for (int64_t i11 = 0; i11 < ne11; i11++) { - const float * const src = (float *)((char *) src1->data + i11*nb11); - for (int64_t i10 = 0; i10 < ne10; i10++) { - dst_data[i10*ne11 + i11] = src[i10]; - } - } - } +static void ggml_compute_forward_get_rows_bf16( + const ggml_compute_params * params, + ggml_tensor * dst) { - // need to zero dst since we are accumulating into it - memset(dst->data, 0, ggml_nbytes(dst)); - } - ggml_barrier(params->threadpool); + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - const int32_t s0 = ((const int32_t*)(dst->op_params))[0]; + GGML_TENSOR_BINARY_OP_LOCALS - // total rows in dst - const int nr = ne1; + const int64_t nc = ne00; + const int64_t nr = ggml_nelements(src1); + + assert(ne0 == nc); + assert(ne02 == ne11); + assert(nb00 == sizeof(ggml_bf16_t)); + assert(ggml_nrows(dst) == nr); + + const int ith = params->ith; + const int nth = params->nth; // rows per thread const int dr = (nr + nth - 1)/nth; @@ -6083,518 +6845,590 @@ static void ggml_compute_forward_conv_transpose_1d_f32( const int ir0 = dr*ith; const int ir1 = MIN(ir0 + dr, nr); - float * const wdata = (float *) params->wdata + 0; - float * const wdata_src = wdata + nk; + for (int64_t i = ir0; i < ir1; ++i) { + const int64_t i12 = i/(ne11*ne10); + const int64_t i11 = (i - i12*ne11*ne10)/ne10; + const int64_t i10 = (i - i12*ne11*ne10 - i11*ne10); + const int64_t i01 = *(int32_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12); - for (int i1 = ir0; i1 < ir1; i1++) { - float * dst_data = (float *)((char *) dst->data + i1*nb1); - float * wdata_kernel = wdata + i1*ne02*ne00; - for (int i10 = 0; i10 < ne10; i10++) { - const int i1n = i10*ne11; - for (int i00 = 0; i00 < ne00; i00++) { - float v = 0; - ggml_vec_dot_f32(ne02, &v, 0, - wdata_src + i1n, 0, - wdata_kernel + i00*ne02, 0, 1); - dst_data[i10*s0 + i00] += v; - } - } + GGML_ASSERT(i01 >= 0 && i01 < ne01); + + ggml_cpu_bf16_to_fp32( + (const ggml_bf16_t *) ((char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03), + (float *) ((char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3), nc); } } -void ggml_compute_forward_conv_transpose_1d( +static void ggml_compute_forward_get_rows_f32( const ggml_compute_params * params, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_TENSOR_BINARY_OP_LOCALS + + const int64_t nc = ne00; + const int64_t nr = ggml_nelements(src1); + + assert(ne0 == nc); + assert(ne02 == ne11); + assert(nb00 == sizeof(float)); + assert(ggml_nrows(dst) == nr); + + const int ith = params->ith; + const int nth = params->nth; + + // rows per thread + const int dr = (nr + nth - 1)/nth; + + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + + for (int64_t i = ir0; i < ir1; ++i) { + const int64_t i12 = i/(ne11*ne10); + const int64_t i11 = (i - i12*ne11*ne10)/ne10; + const int64_t i10 = (i - i12*ne11*ne10 - i11*ne10); + const int64_t i01 = *(int32_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12); + + GGML_ASSERT(i01 >= 0 && i01 < ne01); + + ggml_vec_cpy_f32(nc, + (float *) ((char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3), + (float *) ((char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03)); + } +} + +void ggml_compute_forward_get_rows( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q8_1: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_TQ1_0: + case GGML_TYPE_TQ2_0: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ2_S: + { + ggml_compute_forward_get_rows_q(params, dst); + } break; case GGML_TYPE_F16: { - ggml_compute_forward_conv_transpose_1d_f16_f32(params, dst); + ggml_compute_forward_get_rows_f16(params, dst); + } break; + case GGML_TYPE_BF16: + { + ggml_compute_forward_get_rows_bf16(params, dst); } break; case GGML_TYPE_F32: + case GGML_TYPE_I32: { - ggml_compute_forward_conv_transpose_1d_f32(params, dst); + ggml_compute_forward_get_rows_f32(params, dst); } break; default: { GGML_ABORT("fatal error"); } } + + //static bool first = true; + //printf("ne0 = %d, ne1 = %d, ne2 = %d\n", dst->ne[0], dst->ne[1], dst->ne[2]); + //if (first) { + // first = false; + //} else { + // for (int k = 0; k < dst->ne[1]; ++k) { + // for (int j = 0; j < dst->ne[0]/16; ++j) { + // for (int i = 0; i < 16; ++i) { + // printf("%8.4f ", ((float *) dst->data)[k*dst->ne[0] + j*16 + i]); + // } + // printf("\n"); + // } + // printf("\n"); + // } + // printf("\n"); + // exit(0); + //} } -// ggml_compute_forward_im2col_f32 -// src0: kernel [OC, IC, KH, KW] -// src1: image [N, IC, IH, IW] -// dst: result [N, OH, OW, IC*KH*KW] -static void ggml_compute_forward_im2col_f32( +template +static void ggml_compute_forward_set_rows_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT( dst->type == GGML_TYPE_F32); + GGML_TENSOR_BINARY_OP_LOCALS - GGML_TENSOR_BINARY_OP_LOCALS; + const int64_t nc = ne00; + const int64_t nr = ne01; - const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; - const int32_t s1 = ((const int32_t *)(dst->op_params))[1]; - const int32_t p0 = ((const int32_t *)(dst->op_params))[2]; - const int32_t p1 = ((const int32_t *)(dst->op_params))[3]; - const int32_t d0 = ((const int32_t *)(dst->op_params))[4]; - const int32_t d1 = ((const int32_t *)(dst->op_params))[5]; - const bool is_2D = ((const int32_t *)(dst->op_params))[6] == 1; + assert(ne0 == nc); + assert(ne2 == ne02); + assert(ne3 == ne03); + assert(src0->type == GGML_TYPE_F32); + assert(ne02 % ne11 == 0); + assert(ne03 % ne12 == 0); const int ith = params->ith; const int nth = params->nth; - const int64_t N = is_2D ? ne13 : ne12; - const int64_t IC = is_2D ? ne12 : ne11; - const int64_t IH = is_2D ? ne11 : 1; - const int64_t IW = ne10; + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; - const int64_t KH = is_2D ? ne01 : 1; - const int64_t KW = ne00; + // row range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = std::min(ir0 + dr, nr); - const int64_t OH = is_2D ? ne2 : 1; - const int64_t OW = ne1; + ggml_from_float_t const from_float = ggml_get_type_traits_cpu(dst->type)->from_float; - int ofs0 = is_2D ? nb13 : nb12; - int ofs1 = is_2D ? nb12 : nb11; + for (int64_t i03 = 0; i03 < ne03; ++i03) { + for (int64_t i02 = 0; i02 < ne02; ++i02) { + for (int64_t i = ir0; i < ir1; ++i) { + const int64_t i12 = i03%ne12; + const int64_t i11 = i02%ne11; + const int64_t i10 = i; - GGML_ASSERT(nb10 == sizeof(float)); + const int64_t i1 = *(idx_t *) ((char *) src1->data + i10*nb10 + i11*nb11 + i12*nb12); - // im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW] - { - float * const wdata = (float *) dst->data; + GGML_ASSERT(i1 >= 0 && i1 < ne1); - for (int64_t in = 0; in < N; in++) { - for (int64_t ioh = 0; ioh < OH; ioh++) { // 1 - for (int64_t iow = 0; iow < OW; iow++) { - for (int64_t iic = ith; iic < IC; iic += nth) { + from_float( + (const float *) ((char *) src0->data + i*nb01 + i02*nb02 + i03*nb03), + ((char *) dst->data + i1*nb1 + i02*nb2 + i03*nb3), nc); + } + } + } +} - // micro kernel - float * dst_data = wdata + (in*OH*OW + ioh*OW + iow)*(IC*KH*KW); // [IC, KH, KW] - const float * const src_data = (float *)((char *) src1->data + in*ofs0 + iic*ofs1); // [IH, IW] +void ggml_compute_forward_set_rows( + const ggml_compute_params * params, + ggml_tensor * dst) { - for (int64_t ikh = 0; ikh < KH; ikh++) { // 1 - for (int64_t ikw = 0; ikw < KW; ikw++) { - const int64_t iiw = iow*s0 + ikw*d0 - p0; - const int64_t iih = ioh*s1 + ikh*d1 - p1; + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - if (iih < 0 || iih >= IH || iiw < 0 || iiw >= IW) { - dst_data[iic*(KH*KW) + ikh*KW + ikw] = 0; - } else { - dst_data[iic*(KH*KW) + ikh*KW + ikw] = (src_data[iih*IW + iiw]); - } - } - } - } + switch (src0->type) { + case GGML_TYPE_F32: + { + if (src1->type == GGML_TYPE_I64) { + ggml_compute_forward_set_rows_f32(params, dst); + } else if (src1->type == GGML_TYPE_I32) { + ggml_compute_forward_set_rows_f32(params, dst); + } else { + GGML_ABORT("src1->type = %d (%s) not supported", src1->type, ggml_type_name(src1->type)); } + } break; + default: + { + GGML_ABORT("src0->type = %d (%s) not supported", src0->type, ggml_type_name(src0->type)); } - } } } +// ggml_compute_forward_get_rows_back -// ggml_compute_forward_im2col_f16 -// src0: kernel [OC, IC, KH, KW] -// src1: image [N, IC, IH, IW] -// dst: result [N, OH, OW, IC*KH*KW] -static void ggml_compute_forward_im2col_f16( +static void ggml_compute_forward_get_rows_back_f32_f16( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(src0->type == GGML_TYPE_F16); - GGML_ASSERT(src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_F32); - GGML_ASSERT( dst->type == GGML_TYPE_F16); + if (params->ith != 0) { + return; + } - GGML_TENSOR_BINARY_OP_LOCALS; + GGML_ASSERT(ggml_is_contiguous(dst)); - const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; - const int32_t s1 = ((const int32_t *)(dst->op_params))[1]; - const int32_t p0 = ((const int32_t *)(dst->op_params))[2]; - const int32_t p1 = ((const int32_t *)(dst->op_params))[3]; - const int32_t d0 = ((const int32_t *)(dst->op_params))[4]; - const int32_t d1 = ((const int32_t *)(dst->op_params))[5]; - const bool is_2D = ((const int32_t *)(dst->op_params))[6] == 1; + // ggml_compute_forward_dup_same_cont(params, opt0, dst); - const int ith = params->ith; - const int nth = params->nth; + memset(dst->data, 0, ggml_nbytes(dst)); - const int64_t N = is_2D ? ne13 : ne12; - const int64_t IC = is_2D ? ne12 : ne11; - const int64_t IH = is_2D ? ne11 : 1; - const int64_t IW = ne10; + const int nc = src0->ne[0]; + const int nr = ggml_nelements(src1); - const int64_t KH = is_2D ? ne01 : 1; - const int64_t KW = ne00; + GGML_ASSERT( dst->ne[0] == nc); + GGML_ASSERT(src0->nb[0] == sizeof(ggml_fp16_t)); - const int64_t OH = is_2D ? ne2 : 1; - const int64_t OW = ne1; + for (int i = 0; i < nr; ++i) { + const int r = ((int32_t *) src1->data)[i]; - int ofs0 = is_2D ? nb13 : nb12; - int ofs1 = is_2D ? nb12 : nb11; + for (int j = 0; j < nc; ++j) { + ggml_fp16_t v = ((ggml_fp16_t *) ((char *) src0->data + i*src0->nb[1]))[j]; + ((float *) ((char *) dst->data + r*dst->nb[1]))[j] += GGML_CPU_FP16_TO_FP32(v); + } + } +} - GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); - GGML_ASSERT(nb10 == ggml_type_size(src1->type)); +static void ggml_compute_forward_get_rows_back_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - // im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW] - { - ggml_fp16_t * const wdata = (ggml_fp16_t *) dst->data; + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - for (int64_t in = 0; in < N; in++) { - for (int64_t ioh = 0; ioh < OH; ioh++) { // 1 - for (int64_t iow = 0; iow < OW; iow++) { - for (int64_t iic = ith; iic < IC; iic += nth) { + if (params->ith != 0) { + return; + } - // micro kernel - ggml_fp16_t * dst_data = wdata + (in*OH*OW + ioh*OW + iow)*(IC*KH*KW); // [IC, KH, KW] - const float * const src_data_f32 = src1->type == GGML_TYPE_F32 - ? (const float *)((const char *) src1->data + in*ofs0 + iic*ofs1) - : nullptr; // [IH, IW] - const ggml_fp16_t * const src_data_f16 = src1->type == GGML_TYPE_F16 - ? (const ggml_fp16_t *)((const char *) src1->data + in*ofs0 + iic*ofs1) - : nullptr; // [IH, IW] + GGML_ASSERT(ggml_is_contiguous(dst)); - for (int64_t ikh = 0; ikh < KH; ikh++) { // 1 - for (int64_t ikw = 0; ikw < KW; ikw++) { - const int64_t iiw = iow*s0 + ikw*d0 - p0; - const int64_t iih = ioh*s1 + ikh*d1 - p1; + // ggml_compute_forward_dup_same_cont(params, opt0, dst); - if (iih < 0 || iih >= IH || iiw < 0 || iiw >= IW) { - dst_data[iic*(KH*KW) + ikh*KW + ikw] = 0; - } else { - if (src_data_f32 != nullptr) { - dst_data[iic*(KH*KW) + ikh*KW + ikw] = GGML_CPU_FP32_TO_FP16(src_data_f32[iih*IW + iiw]); - } else { - dst_data[iic*(KH*KW) + ikh*KW + ikw] = src_data_f16[iih*IW + iiw]; - } - } - } - } - } - } - } - } + memset(dst->data, 0, ggml_nbytes(dst)); + + const int nc = src0->ne[0]; + const int nr = ggml_nelements(src1); + + GGML_ASSERT( dst->ne[0] == nc); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + + for (int i = 0; i < nr; ++i) { + const int r = ((int32_t *) src1->data)[i]; + + ggml_vec_add_f32(nc, + (float *) ((char *) dst->data + r*dst->nb[1]), + (float *) ((char *) dst->data + r*dst->nb[1]), + (float *) ((char *) src0->data + i*src0->nb[1])); } } -void ggml_compute_forward_im2col( +void ggml_compute_forward_get_rows_back( const ggml_compute_params * params, - ggml_tensor * dst) { - switch (dst->type) { + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + + switch (src0->type) { case GGML_TYPE_F16: { - ggml_compute_forward_im2col_f16(params, dst); + ggml_compute_forward_get_rows_back_f32_f16(params, dst); } break; case GGML_TYPE_F32: { - ggml_compute_forward_im2col_f32(params, dst); + ggml_compute_forward_get_rows_back_f32(params, dst); } break; default: { GGML_ABORT("fatal error"); } } + + //static bool first = true; + //printf("ne0 = %d, ne1 = %d, ne2 = %d\n", dst->ne[0], dst->ne[1], dst->ne[2]); + //if (first) { + // first = false; + //} else { + // for (int k = 0; k < dst->ne[1]; ++k) { + // for (int j = 0; j < dst->ne[0]/16; ++j) { + // for (int i = 0; i < 16; ++i) { + // printf("%8.4f ", ((float *) dst->data)[k*dst->ne[0] + j*16 + i]); + // } + // printf("\n"); + // } + // printf("\n"); + // } + // printf("\n"); + // exit(0); + //} } -// ggml_compute_forward_im2col_back_f32 +// ggml_compute_forward_diag -void ggml_compute_forward_im2col_back_f32( +static void ggml_compute_forward_diag_f32( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; // gradients of forward pass output - const ggml_tensor * src1 = dst->src[1]; // convolution kernel + const ggml_tensor * src0 = dst->src[0]; - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT( dst->type == GGML_TYPE_F32); + if (params->ith != 0) { + return; + } - GGML_TENSOR_BINARY_OP_LOCALS; + // TODO: handle transposed/permuted matrices - const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; - const int32_t s1 = ((const int32_t *)(dst->op_params))[1]; - const int32_t p0 = ((const int32_t *)(dst->op_params))[2]; - const int32_t p1 = ((const int32_t *)(dst->op_params))[3]; - const int32_t d0 = ((const int32_t *)(dst->op_params))[4]; - const int32_t d1 = ((const int32_t *)(dst->op_params))[5]; - const bool is_2D = ((const int32_t *)(dst->op_params))[6] == 1; + GGML_TENSOR_UNARY_OP_LOCALS - const int ith = params->ith; - const int nth = params->nth; + GGML_ASSERT(ne00 == ne0); + GGML_ASSERT(ne00 == ne1); + GGML_ASSERT(ne01 == 1); + GGML_ASSERT(ne02 == ne2); + GGML_ASSERT(ne03 == ne3); - const int64_t N = is_2D ? ne3 : ne2; - const int64_t IC = is_2D ? ne2 : ne1; - const int64_t IH = is_2D ? ne1 : 1; - const int64_t IW = ne0; + GGML_ASSERT(nb00 == sizeof(float)); + GGML_ASSERT(nb0 == sizeof(float)); - const int64_t KH = is_2D ? ne11 : 1; - const int64_t KW = ne10; + for (int i3 = 0; i3 < ne3; i3++) { + for (int i2 = 0; i2 < ne2; i2++) { + for (int i1 = 0; i1 < ne1; i1++) { + float * d = (float *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1); + float * s = (float *)((char *) src0->data + i3*nb03 + i2*nb02); + for (int i0 = 0; i0 < i1; i0++) { + d[i0] = 0; + } + d[i1] = s[i1]; + for (int i0 = i1+1; i0 < ne0; i0++) { + d[i0] = 0; + } + } + } + } +} - const int64_t OH = is_2D ? ne02 : 1; - const int64_t OW = ne01; +void ggml_compute_forward_diag( + const ggml_compute_params * params, + ggml_tensor * dst) { - int ofs0 = is_2D ? nb3 : nb2; - int ofs1 = is_2D ? nb2 : nb1; + const ggml_tensor * src0 = dst->src[0]; - GGML_ASSERT(nb0 == sizeof(float)); + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_diag_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - // im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW] - { - float * const wdata = (float *) dst->data; +// ggml_compute_forward_diag_mask_inf - for (int64_t in = 0; in < N; in++) { - for (int64_t iic = ith; iic < IC; iic += nth) { - for (int64_t iih = 0; iih < IH; iih++) { - for (int64_t iiw = 0; iiw < IW; iiw++) { +static void ggml_compute_forward_diag_mask_f32( + const ggml_compute_params * params, + ggml_tensor * dst, + const float value) { - // micro kernel - float grad = 0.0f; - for (int64_t ikh = 0; ikh < KH; ikh++) { - for (int64_t ikw = 0; ikw < KW; ikw++) { - // For s0 > 1 some values were skipped over in the forward pass. - // These values have tmpw % s0 != 0 and need to be skipped in the backwards pass as well. - const int64_t tmpw = (iiw + p0 - ikw*d0); - if (tmpw % s0 != 0) { - continue; - } - const int64_t iow = tmpw / s0; + const ggml_tensor * src0 = dst->src[0]; - // Equivalent logic as above except for s1. - int64_t ioh; - if (is_2D) { - const int64_t tmph = iih + p1 - ikh*d1; + const int ith = params->ith; + const int nth = params->nth; - if (tmph % s1 != 0) { - continue; - } + const int n_past = ((int32_t *) dst->op_params)[0]; + const bool inplace = src0->data == dst->data; - ioh = tmph / s1; - } else { - ioh = 0; - } + GGML_ASSERT(n_past >= 0); - if (iow < 0 || iow >= OW || ioh < 0 || ioh >= OH) { - continue; - } + if (!inplace) { + if (ith == 0) { + // memcpy needs to be synchronized across threads to avoid race conditions. + // => do it in INIT phase + GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src0)); + GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src0)); + memcpy( + ((char *) dst->data), + ((char *) src0->data), + ggml_nbytes(dst)); + } + ggml_barrier(params->threadpool); + } - const float * const grad_in = (const float *) src0->data - + (in*OH*OW + ioh*OW + iow)*(IC*KH*KW); // [IC, KH, KW] - grad += grad_in[iic*(KH*KW) + ikh*KW + ikw]; - } - } - float * dst_data = (float *)((char *) wdata + (in*ofs0 + iic*ofs1)); // [IH, IW] - dst_data[iih*IW + iiw] = grad; - } + // TODO: handle transposed/permuted matrices + + const int n = ggml_nrows(src0); + const int nc = src0->ne[0]; + const int nr = src0->ne[1]; + const int nz = n/nr; + + GGML_ASSERT( dst->nb[0] == sizeof(float)); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + + for (int k = 0; k < nz; k++) { + for (int j = ith; j < nr; j += nth) { + for (int i = n_past; i < nc; i++) { + if (i > n_past + j) { + *(float *)((char *) dst->data + k*dst->nb[2] + j*dst->nb[1] + i*dst->nb[0]) = value; } } } } } - -// ggml_compute_forward_im2col_3d_f16 -// src0: kernel [OC*IC, KD, KH, KW] -// src1: image [N*IC, ID, IH, IW] -// dst: result [N*OD, OH, OW, IC * KD * KH * KW] -static void ggml_compute_forward_im2col_3d_f16( +void ggml_compute_forward_diag_mask_inf( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - - GGML_ASSERT(src0->type == GGML_TYPE_F16); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT( dst->type == GGML_TYPE_F16); - - GGML_TENSOR_BINARY_OP_LOCALS; - - const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; - const int32_t s1 = ((const int32_t *)(dst->op_params))[1]; - const int32_t s2 = ((const int32_t *)(dst->op_params))[2]; - const int32_t p0 = ((const int32_t *)(dst->op_params))[3]; - const int32_t p1 = ((const int32_t *)(dst->op_params))[4]; - const int32_t p2 = ((const int32_t *)(dst->op_params))[5]; - const int32_t d0 = ((const int32_t *)(dst->op_params))[6]; - const int32_t d1 = ((const int32_t *)(dst->op_params))[7]; - const int32_t d2 = ((const int32_t *)(dst->op_params))[8]; - const int32_t IC = ((const int32_t *)(dst->op_params))[9]; - - - const int ith = params->ith; - const int nth = params->nth; - - const int64_t N = ne13 / IC; - const int64_t ID = ne12; - const int64_t IH = ne11; - const int64_t IW = ne10; - - const int64_t OC = ne03 / IC; - GGML_UNUSED(OC); - const int64_t KD = ne02; - const int64_t KH = ne01; - const int64_t KW = ne00; - - const int64_t OD = ne3 / N; - const int64_t OH = ne2; - const int64_t OW = ne1; - const int64_t OH_OW = OH*OW; - const int64_t KD_KH_KW = KD*KH*KW; - const int64_t KH_KW = KH*KW; - const int64_t IC_KD_KH_KW = IC*KD*KH*KW; - - GGML_ASSERT(nb10 == sizeof(float)); - - // im2col: [N*IC, ID, IH, IW] => [N*OD, OH, OW, IC * KD * KH * KW] - { - ggml_fp16_t * const wdata = (ggml_fp16_t *) dst->data; - for (int64_t in = 0; in < N; in++) { - for (int64_t iod = 0; iod < OD; iod++) { - for (int64_t ioh = 0; ioh < OH; ioh++) { - for (int64_t iow = 0; iow < OW; iow++) { - for (int64_t iic = ith; iic < IC; iic += nth) { + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_diag_mask_f32(params, dst, -INFINITY); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - // micro kernel - ggml_fp16_t * dst_data = wdata + (in*OD*OH_OW + iod*OH_OW + ioh*OW + iow)*IC_KD_KH_KW; // [IC, KD, KH, KW] - const float * const src_data = (const float *) ((const char *)src1->data + (in*IC + iic)*nb13); // [ID, IH, IW] +void ggml_compute_forward_diag_mask_zero( + const ggml_compute_params * params, + ggml_tensor * dst) { - for (int64_t ikd = 0; ikd < KD; ikd++) { - for (int64_t ikh = 0; ikh < KH; ikh++) { - for (int64_t ikw = 0; ikw < KW; ikw++) { - const int64_t iiw = iow*s0 + ikw*d0 - p0; - const int64_t iih = ioh*s1 + ikh*d1 - p1; - const int64_t iid = iod*s2 + ikd*d2 - p2; + const ggml_tensor * src0 = dst->src[0]; - if (iid < 0 || iid >= ID || iih < 0 || iih >= IH || iiw < 0 || iiw >= IW) { - dst_data[iic*KD_KH_KW + ikd * KH_KW + ikh*KW + ikw] = 0; - } else { - const float * const s = (const float *) ((const char *)src_data + iid*nb12 + iih*nb11 + iiw*nb10); // [ID, IH, IW] - dst_data[iic*KD_KH_KW + ikd * KH_KW + ikh*KW + ikw] = GGML_CPU_FP32_TO_FP16(*s); - } - } - } - } - } - } - } + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_diag_mask_f32(params, dst, 0); + } break; + default: + { + GGML_ABORT("fatal error"); } - } } } -// ggml_compute_forward_im2col_3d_f32 -// src0: kernel [OC*IC, KD, KH, KW] -// src1: image [N*IC, ID, IH, IW] -// dst: result [N*OD, OH, OW, IC * KD * KH * KW] -static void ggml_compute_forward_im2col_3d_f32( +// ggml_compute_forward_soft_max + +static void ggml_compute_forward_soft_max_f32( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; + const ggml_tensor * src2 = dst->src[2]; - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT( dst->type == GGML_TYPE_F32); - - GGML_TENSOR_BINARY_OP_LOCALS; + assert(ggml_is_contiguous(dst)); + assert(ggml_are_same_shape(src0, dst)); - const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; - const int32_t s1 = ((const int32_t *)(dst->op_params))[1]; - const int32_t s2 = ((const int32_t *)(dst->op_params))[2]; - const int32_t p0 = ((const int32_t *)(dst->op_params))[3]; - const int32_t p1 = ((const int32_t *)(dst->op_params))[4]; - const int32_t p2 = ((const int32_t *)(dst->op_params))[5]; - const int32_t d0 = ((const int32_t *)(dst->op_params))[6]; - const int32_t d1 = ((const int32_t *)(dst->op_params))[7]; - const int32_t d2 = ((const int32_t *)(dst->op_params))[8]; - const int32_t IC = ((const int32_t *)(dst->op_params))[9]; + float scale = 1.0f; + float max_bias = 0.0f; + memcpy(&scale, (float *) dst->op_params + 0, sizeof(float)); + memcpy(&max_bias, (float *) dst->op_params + 1, sizeof(float)); const int ith = params->ith; const int nth = params->nth; - const int64_t N = ne13 / IC; - const int64_t ID = ne12; - const int64_t IH = ne11; - const int64_t IW = ne10; + GGML_TENSOR_UNARY_OP_LOCALS - const int64_t OC = ne03 / IC; - GGML_UNUSED(OC); - const int64_t KD = ne02; - const int64_t KH = ne01; - const int64_t KW = ne00; + const int64_t nb11 = src1 ? src1->nb[1] : 1; + const int64_t nb12 = src1 ? src1->nb[2] : 1; + const int64_t nb13 = src1 ? src1->nb[3] : 1; - const int64_t OD = ne3 / N; - const int64_t OH = ne2; - const int64_t OW = ne1; + const int64_t ne12 = src1 ? src1->ne[2] : 1; + const int64_t ne13 = src1 ? src1->ne[3] : 1; - const int64_t OH_OW = OH*OW; - const int64_t KD_KH_KW = KD*KH*KW; - const int64_t KH_KW = KH*KW; - const int64_t IC_KD_KH_KW = IC*KD*KH*KW; + // TODO: is this supposed to be ceil instead of floor? + // https://huggingface.co/mosaicml/mpt-7b/blob/main/attention.py#L370 + const uint32_t n_head = ne02; + const uint32_t n_head_log2 = 1u << (uint32_t) floor(log2(n_head)); - GGML_ASSERT(nb10 == sizeof(float)); + const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); + const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); - // im2col: [N*IC, ID, IH, IW] => [N*OD, OH, OW, IC * KD * KH * KW] - { - float * const wdata = (float *) dst->data; + float * wp = (float *) params->wdata + (ne00 + CACHE_LINE_SIZE_F32) * ith; - for (int64_t in = 0; in < N; in++) { - for (int64_t iod = 0; iod < OD; iod++) { - for (int64_t ioh = 0; ioh < OH; ioh++) { - for (int64_t iow = 0; iow < OW; iow++) { - for (int64_t iic = ith; iic < IC; iic += nth) { + const bool use_f16 = (src1 && src1->type == GGML_TYPE_F16); - // micro kernel - float * dst_data = wdata + (in*OD*OH_OW + iod*OH_OW + ioh*OW + iow)*IC_KD_KH_KW; // [IC, KD, KH, KW] - const float * const src_data = (const float *) ((const char *)src1->data + (in*IC + iic)*nb13); // [ID, IH, IW] + // sinks + const float * sk = src2 ? (float *)((char *) src2->data) : nullptr; - for (int64_t ikd = 0; ikd < KD; ikd++) { - for (int64_t ikh = 0; ikh < KH; ikh++) { - for (int64_t ikw = 0; ikw < KW; ikw++) { - const int64_t iiw = iow*s0 + ikw*d0 - p0; - const int64_t iih = ioh*s1 + ikh*d1 - p1; - const int64_t iid = iod*s2 + ikd*d2 - p2; + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = ith; i01 < ne01; i01 += nth) { + const int64_t i11 = i01; + const int64_t i12 = i02%ne12; + const int64_t i13 = i03%ne13; - if (iid < 0 || iid >= ID || iih < 0 || iih >= IH || iiw < 0 || iiw >= IW || iid < 0 || iid >= ID) { - dst_data[iic*KD_KH_KW + ikd * KH_KW + ikh*KW + ikw] = 0; - } else { - const float * const s = (const float *) ((const char *)src_data + iid*nb12 + iih*nb11 + iiw*nb10); // [ID, IH, IW] - dst_data[iic*KD_KH_KW + ikd * KH_KW + ikh*KW + ikw] = *s; - } - } - } - } + // ALiBi + const uint32_t h = i02; // head + const float slope = (max_bias > 0.0f) ? h < n_head_log2 ? powf(m0, h + 1) : powf(m1, 2*(h - n_head_log2) + 1) : 1.0f; + + float * sp = (float *)((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03); + float * dp = (float *)((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); + + // broadcast the mask across rows + ggml_fp16_t * mp_f16 = src1 ? (ggml_fp16_t *)((char *) src1->data + i11*nb11 + i12*nb12 + i13*nb13) : NULL; + float * mp_f32 = src1 ? (float *)((char *) src1->data + i11*nb11 + i12*nb12 + i13*nb13) : NULL; + + ggml_vec_cpy_f32 (ne00, wp, sp); + ggml_vec_scale_f32(ne00, wp, scale); + if (mp_f32) { + if (use_f16) { + for (int i = 0; i < ne00; ++i) { + wp[i] += slope*GGML_CPU_FP16_TO_FP32(mp_f16[i]); + } + } else { + for (int i = 0; i < ne00; ++i) { + wp[i] += slope*mp_f32[i]; } } } + +#ifndef NDEBUG + for (int i = 0; i < ne00; ++i) { + //printf("p[%d] = %f\n", i, p[i]); + assert(!isnan(wp[i])); + } +#endif // NDEBUG + + float max = -INFINITY; + ggml_vec_max_f32(ne00, &max, wp); + + // if we have sinks, make a correction as if they were included in the softmax + if (sk) { + max = MAX(max, sk[i02]); + } + + ggml_float sum = ggml_vec_soft_max_f32(ne00, dp, wp, max); + assert(sum > 0.0); + + if (sk) { + sum += (ggml_float) expf(sk[i02] - max); + } + + sum = 1.0/sum; + ggml_vec_scale_f32(ne00, dp, sum); + +#ifndef NDEBUG + for (int i = 0; i < ne00; ++i) { + assert(!isnan(dp[i])); + assert(!isinf(dp[i])); + } +#endif // NDEBUG } } } } - -void ggml_compute_forward_im2col_3d( +void ggml_compute_forward_soft_max( const ggml_compute_params * params, ggml_tensor * dst) { - switch (dst->type) { - case GGML_TYPE_F16: - { - ggml_compute_forward_im2col_3d_f16(params, dst); - } break; + + const ggml_tensor * src0 = dst->src[0]; + + switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_im2col_3d_f32(params, dst); + ggml_compute_forward_soft_max_f32(params, dst); } break; default: { @@ -6603,452 +7437,534 @@ void ggml_compute_forward_im2col_3d( } } -static void ggml_call_mul_mat(ggml_type type, const ggml_compute_params * params, int64_t m, int64_t n, int64_t k, - void * a, void * b, float * c) { - const ggml_type_traits * traits = ggml_get_type_traits(type); - struct ggml_tensor src1 = {}; - src1.type = type; - src1.ne[0] = k; - src1.ne[1] = m; - src1.ne[2] = 1; - src1.ne[3] = 1; - src1.nb[0] = traits->type_size; - src1.nb[1] = k * traits->type_size; - src1.nb[2] = src1.nb[1]; - src1.nb[3] = src1.nb[2]; - src1.data = a; - struct ggml_tensor src0 = {}; - src0.type = type; - src0.ne[0] = k; - src0.ne[1] = n; - src0.ne[2] = 1; - src0.ne[3] = 1; - src0.nb[0] = traits->type_size; - src0.nb[1] = k * traits->type_size; - src0.nb[2] = src0.nb[1]; - src0.nb[3] = src0.nb[2]; - src0.data = b; +// ggml_compute_forward_soft_max_ext_back - struct ggml_tensor dst = {}; - dst.ne[0] = n; - dst.ne[1] = m; - dst.ne[2] = 1; - dst.ne[3] = 1; - dst.nb[0] = sizeof(float); - dst.nb[1] = n * sizeof(float); - dst.nb[2] = dst.nb[1]; - dst.nb[3] = dst.nb[2]; - dst.data = c; - dst.src[0] = &src0; - dst.src[1] = &src1; +static void ggml_compute_forward_soft_max_ext_back_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - ggml_compute_forward_mul_mat(params, &dst); -} + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; -static inline int64_t ggml_wrap_around(int64_t coord, int64_t size) { - return (coord + size) % size; // adding size avoids negative number weirdness -} + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src1)); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_are_same_shape(src1, dst)); -// ggml_compute_forward_conv_2d + float scale = 1.0f; + float max_bias = 0.0f; + memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); -static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params, - const ggml_tensor * kernel, // [KW, KH, IC, OC] - const ggml_tensor * src, // [W, H, C, N] - ggml_tensor * dst, // [OW, OH, OC, N] - ggml_type kernel_type) { + GGML_ASSERT(max_bias == 0.0f); - GGML_ASSERT(ggml_is_contiguous(kernel)); - GGML_ASSERT(kernel_type == GGML_TYPE_F16 || kernel_type == GGML_TYPE_F32); - GGML_ASSERT(kernel->type == kernel_type); + // TODO: handle transposed/permuted matrices - const ggml_type_traits * traits = ggml_get_type_traits(kernel_type); + const int ith = params->ith; + const int nth = params->nth; - const int32_t stride_x = dst->op_params[0]; - const int32_t stride_y = dst->op_params[1]; - const int32_t pad_x = dst->op_params[2]; - const int32_t pad_y = dst->op_params[3]; - const int32_t dilation_x = dst->op_params[4]; - const int32_t dilation_y = dst->op_params[5]; + const int nc = src0->ne[0]; + const int nr = ggml_nrows(src0); - const int64_t c_in = src->ne[2]; - const int64_t c_out = kernel->ne[3]; - GGML_ASSERT(c_in == kernel->ne[2]); + // rows per thread + const int dr = (nr + nth - 1)/nth; - const int64_t src_w = src->ne[0]; - const int64_t src_h = src->ne[1]; - const int64_t knl_w = kernel->ne[0]; - const int64_t knl_h = kernel->ne[1]; - const int64_t dst_w = dst->ne[0]; - const int64_t dst_h = dst->ne[1]; + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - const float * src_data = (float *) src->data; - void * knl_data = kernel->data; - float * dst_data = (float *) dst->data; + for (int i1 = ir0; i1 < ir1; i1++) { + float *dy = (float *)((char *) src0->data + i1*src0->nb[1]); + float *y = (float *)((char *) src1->data + i1*src1->nb[1]); + float *dx = (float *)((char *) dst->data + i1*dst->nb[1]); - const int64_t knl_n = knl_w * knl_h * c_in; - const int64_t patch_total = dst->ne[3] * dst_w * dst_h; +#ifndef NDEBUG + for (int i = 0; i < nc; ++i) { + //printf("p[%d] = %f\n", i, p[i]); + assert(!isnan(dy[i])); + assert(!isnan(y[i])); + } +#endif // NDEBUG + // Jii = yi - yi*yi + // Jij = -yi*yj + // J = diag(y)-y.T*y + // dx = J * dy + // dxk = sum_i(Jki * dyi) + // dxk = sum_i(-yk*yi * dyi) - (-yk*yk)*dyk + (yk - yk*yk)*dyk + // dxk = sum_i(-yk*yi * dyi) + yk*yk*dyk + yk*dyk - yk*yk*dyk + // dxk = sum_i(-yk*yi * dyi) + yk*dyk + // dxk = -yk * sum_i(yi * dyi) + yk*dyk + // dxk = -yk * dot(y, dy) + yk*dyk + // dxk = yk * (- dot(y, dy) + dyk) + // dxk = yk * (dyk - dot(y, dy)) + // + // post-order: + // dot_y_dy := dot(y, dy) + // dx := dy + // dx := dx - dot_y_dy + // dx := dx * y - const int64_t space_per_patch = knl_n * traits->type_size + c_out * sizeof(float); - const int64_t batch_size = params->wsize / space_per_patch; - const int64_t patches_per_batch = batch_size > 8 ? (batch_size / 8) * 8 : batch_size; - const int64_t batch_n = (patch_total + patches_per_batch - 1) / patches_per_batch; + // linear runtime, no additional memory + float dot_y_dy = 0; + ggml_vec_dot_f32 (nc, &dot_y_dy, 0, y, 0, dy, 0, 1); + ggml_vec_cpy_f32 (nc, dx, dy); + ggml_vec_acc1_f32 (nc, dx, -dot_y_dy); + ggml_vec_mul_f32 (nc, dx, dx, y); + ggml_vec_scale_f32(nc, dx, scale); - GGML_ASSERT(patches_per_batch > 0 && batch_size >= 1); +#ifndef NDEBUG + for (int i = 0; i < nc; ++i) { + assert(!isnan(dx[i])); + assert(!isinf(dx[i])); + } +#endif // NDEBUG + } +} - void * tmp = params->wdata; +void ggml_compute_forward_soft_max_ext_back( + const ggml_compute_params * params, + ggml_tensor * dst) { - for (int64_t batch_i = 0; batch_i < batch_n; ++batch_i) { + const ggml_tensor * src0 = dst->src[0]; - const int64_t patch_start_batch = batch_i * patches_per_batch; - const int64_t patch_end_batch = std::min(patch_start_batch + patches_per_batch, - patch_total); - const int64_t patch_n = patch_end_batch - patch_start_batch; + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_soft_max_ext_back_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - const int64_t patch_per_thread = (patch_n + params->nth - 1) / params->nth; - const int64_t patch_start = patch_start_batch + params->ith * patch_per_thread; - const int64_t patch_end = std::min(patch_start + patch_per_thread, patch_end_batch); +// ggml_compute_forward_clamp - //im2col for a patch - for (int64_t p = patch_start; p < patch_end; ++p) { - const int64_t batch_n = p / (dst_w * dst_h); - const int64_t src_x = (p / dst_w) % dst_h; - const int64_t src_y = p % dst_w; +static void ggml_compute_forward_clamp_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - const float * src_base = (const float *)((const char *)src_data + batch_n * src->nb[3]); - char * dst_row = (char *) tmp + (p % patches_per_batch) * knl_n * traits->type_size; + const ggml_tensor * src0 = dst->src[0]; - for (int64_t ic = 0; ic < c_in; ++ic) { - for (int64_t ky = 0; ky < knl_h; ++ky) { - for (int64_t kx = 0; kx < knl_w; ++kx) { - const int64_t sy = src_x * stride_y + ky * dilation_y - pad_y; - const int64_t sx = src_y * stride_x + kx * dilation_x - pad_x; + float min; + float max; + memcpy(&min, (float *) dst->op_params + 0, sizeof(float)); + memcpy(&max, (float *) dst->op_params + 1, sizeof(float)); - int64_t dst_idx = ic * (knl_h * knl_w) + ky * knl_w + kx; + const int ith = params->ith; + const int nth = params->nth; - float src_val; - if (sy < 0 || sy >= src_h || sx < 0 || sx >= src_w) { - src_val = 0.0f; - } else { - const float * src_ptr = (const float *)((const char *)src_base + sx * src->nb[0] + sy * src->nb[1] + ic * src->nb[2]); - src_val = *src_ptr; - } + const int n = ggml_nrows(src0); + const int nc = src0->ne[0]; - char * element_ptr = dst_row + dst_idx * traits->type_size; - if (kernel_type == GGML_TYPE_F32) { - *(float *) element_ptr = src_val; - } else if (kernel_type == GGML_TYPE_F16) { - *(ggml_fp16_t *) element_ptr = GGML_CPU_FP32_TO_FP16(src_val); - } - } - } - } - } // patches handled by this thread + const size_t nb00 = src0->nb[0]; + const size_t nb01 = src0->nb[1]; - ggml_barrier(params->threadpool); + const size_t nb0 = dst->nb[0]; + const size_t nb1 = dst->nb[1]; - float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n * traits->type_size); + GGML_ASSERT( nb0 == sizeof(float)); + GGML_ASSERT(nb00 == sizeof(float)); - GGML_ASSERT(gemm_output + patch_n * c_out <= (float*)tmp + params->wsize); + for (int j = ith; j < n; j += nth) { + float * dst_ptr = (float *) ((char *) dst->data + j*nb1); + float * src0_ptr = (float *) ((char *) src0->data + j*nb01); - // GEMM: patches[patch_n, knl_n] × kernel[knl_n, c_out] = output[patch_n, c_out] - ggml_call_mul_mat(kernel_type, params, patch_n, c_out, knl_n, tmp, knl_data, gemm_output); + for (int i = 0; i < nc; i++) { + dst_ptr[i] = MAX(MIN(src0_ptr[i], max), min); + } + } +} - ggml_barrier(params->threadpool); +static void ggml_compute_forward_clamp_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; - //permute back [OC, N, OH, OW] to [N, OC, OH, OW] - const int64_t permute_per_thread = (patch_n + params->nth - 1) / params->nth; - const int64_t permute_start = params->ith * permute_per_thread; - const int64_t permute_end = std::min(permute_start + permute_per_thread, patch_n); + float min; + float max; + memcpy(&min, (float *) dst->op_params + 0, sizeof(float)); + memcpy(&max, (float *) dst->op_params + 1, sizeof(float)); - for (int64_t i = permute_start; i < permute_end; ++i) { - const int64_t p = patch_start_batch + i; - const int64_t batch_n = p / (dst_w * dst_h); - const int64_t dst_y = (p / dst_w) % dst_h; - const int64_t dst_x = p % dst_w; + const int ith = params->ith; + const int nth = params->nth; - for (int64_t oc = 0; oc < c_out; ++oc) { - const float value = gemm_output[i * c_out + oc]; - float * dst_ptr = (float *)((char *)dst_data + dst_x * dst->nb[0] + dst_y * dst->nb[1] + oc * dst->nb[2] + batch_n * dst->nb[3]); - *dst_ptr = value; - } + const int n = ggml_nrows(src0); + const int nc = src0->ne[0]; + + const size_t nb00 = src0->nb[0]; + const size_t nb01 = src0->nb[1]; + + const size_t nb0 = dst->nb[0]; + const size_t nb1 = dst->nb[1]; + + GGML_ASSERT( nb0 == sizeof(ggml_fp16_t)); + GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); + + for (int j = ith; j < n; j += nth) { + ggml_fp16_t * dst_ptr = (ggml_fp16_t *) ((char *) dst->data + j*nb1); + ggml_fp16_t * src0_ptr = (ggml_fp16_t *) ((char *) src0->data + j*nb01); + + for (int i = 0; i < nc; i++) { + float v = GGML_CPU_FP16_TO_FP32(src0_ptr[i]); + dst_ptr[i] = GGML_CPU_FP32_TO_FP16(MAX(MIN(v, max), min)); } } } -void ggml_compute_forward_conv_2d( +void ggml_compute_forward_clamp( const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - ggml_compute_forward_conv_2d_impl(params, src0, src1, dst, src0->type); + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_clamp_f32(params, dst); + } break; + case GGML_TYPE_F16: + { + ggml_compute_forward_clamp_f16(params, dst); + } break; + case GGML_TYPE_BF16: + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q8_1: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_TQ1_0: + case GGML_TYPE_TQ2_0: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_Q8_K: + case GGML_TYPE_I8: + case GGML_TYPE_I16: + case GGML_TYPE_I32: + case GGML_TYPE_I64: + case GGML_TYPE_F64: + case GGML_TYPE_COUNT: + { + GGML_ABORT("fatal error"); + } + } } -// ggml_compute_forward_conv_3d - -static void ggml_compute_forward_conv_3d_impl(const ggml_compute_params * params, - const ggml_tensor * kernel, - const ggml_tensor * src, - ggml_tensor * dst, - ggml_type kernel_type) { - - GGML_ASSERT(ggml_is_contiguous(kernel)); - GGML_ASSERT(kernel_type == GGML_TYPE_F16 || kernel_type == GGML_TYPE_F32); - GGML_ASSERT(kernel->type == kernel_type); - - const ggml_type_traits * traits = ggml_get_type_traits(kernel_type); - - const int32_t s0 = dst->op_params[0]; - const int32_t s1 = dst->op_params[1]; - const int32_t s2 = dst->op_params[2]; - const int32_t p0 = dst->op_params[3]; - const int32_t p1 = dst->op_params[4]; - const int32_t p2 = dst->op_params[5]; - const int32_t d0 = dst->op_params[6]; - const int32_t d1 = dst->op_params[7]; - const int32_t d2 = dst->op_params[8]; - const int32_t c = dst->op_params[9]; - const int32_t n = dst->op_params[10]; - const int32_t oc = dst->op_params[11]; +// ggml_compute_forward_rope - const int64_t src_w = src->ne[0]; - const int64_t src_h = src->ne[1]; - const int64_t src_d = src->ne[2]; - const int64_t knl_w = kernel->ne[0]; - const int64_t knl_h = kernel->ne[1]; - const int64_t knl_d = kernel->ne[2]; - const int64_t dst_w = dst->ne[0]; - const int64_t dst_h = dst->ne[1]; - const int64_t dst_d = dst->ne[2]; +static float rope_yarn_ramp(const float low, const float high, const int i0) { + const float y = (i0 / 2 - low) / MAX(0.001f, high - low); + return 1 - MIN(1, MAX(0, y)); +} - const float * src_data = (float *) src->data; - void * knl_data = kernel->data; - float * dst_data = (float *) dst->data; +// YaRN algorithm based on LlamaYaRNScaledRotaryEmbedding.py from https://github.com/jquesnelle/yarn +// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng. +static void rope_yarn( + float theta_extrap, float freq_scale, float corr_dims[2], int64_t i0, float ext_factor, float mscale, + float * cos_theta, float * sin_theta) { + // Get n-d rotational scaling corrected for extrapolation + float theta_interp = freq_scale * theta_extrap; + float theta = theta_interp; + if (ext_factor != 0.0f) { + float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], i0) * ext_factor; + theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix; - const int64_t knl_n_per_channel = knl_w * knl_h * knl_d; - const int64_t knl_n_total = knl_n_per_channel * c; - const int64_t patch_total = n * dst_w * dst_h * dst_d; + // Get n-d magnitude scaling corrected for interpolation + mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale); + } + *cos_theta = cosf(theta) * mscale; + *sin_theta = sinf(theta) * mscale; +} - const int64_t space_per_patch = knl_n_total * traits->type_size + oc * sizeof(float); - const int64_t batch_size = params->wsize / space_per_patch; - const int64_t patches_per_batch = batch_size > 8 ? (batch_size / 8) * 8 : batch_size; - const int64_t batch_n = (patch_total + patches_per_batch - 1) / patches_per_batch; +static void ggml_rope_cache_init( + float theta_base, float freq_scale, const float * freq_factors, float corr_dims[2], int64_t ne0, float ext_factor, float mscale, + float * cache, float sin_sign, float theta_scale) { + // ref: https://github.com/jquesnelle/yarn/blob/master/scaled_rope/LlamaYaRNScaledRotaryEmbedding.py + float theta = theta_base; + for (int64_t i0 = 0; i0 < ne0; i0 += 2) { + const float ff = freq_factors ? freq_factors[i0/2] : 1.0f; + rope_yarn( + theta/ff, freq_scale, corr_dims, i0, ext_factor, mscale, &cache[i0 + 0], &cache[i0 + 1] + ); + cache[i0 + 1] *= sin_sign; - GGML_ASSERT(patches_per_batch > 0 && batch_size >= 1); + theta *= theta_scale; + } +} - void * tmp = params->wdata; +static void ggml_mrope_cache_init( + float theta_base_t, float theta_base_h, float theta_base_w, float theta_base_e, int sections[4], bool is_imrope, bool indep_sects, + float freq_scale, const float * freq_factors, float corr_dims[2], int64_t ne0, float ext_factor, float mscale, + float * cache, float sin_sign, float theta_scale) { + // ref: https://github.com/jquesnelle/yarn/blob/master/scaled_rope/LlamaYaRNScaledRotaryEmbedding.py + float theta_t = theta_base_t; + float theta_h = theta_base_h; + float theta_w = theta_base_w; + float theta_e = theta_base_e; // extra position id for vision encoder + int sect_dims = sections[0] + sections[1] + sections[2] + sections[3]; + int sec_w = sections[1] + sections[0]; + int sec_e = sections[2] + sec_w; + GGML_ASSERT(sect_dims <= ne0); - for (int64_t batch_i = 0; batch_i < batch_n; ++batch_i) { - const int64_t patch_start_batch = batch_i * patches_per_batch; - const int64_t patch_end_batch = std::min(patch_start_batch + patches_per_batch, patch_total); - const int64_t patch_n_in_batch = patch_end_batch - patch_start_batch; + for (int64_t i0 = 0; i0 < ne0; i0 += 2) { + const float ff = freq_factors ? freq_factors[i0/2] : 1.0f; - const int64_t patch_per_thread = (patch_n_in_batch + params->nth - 1) / params->nth; - const int64_t patch_start = patch_start_batch + params->ith * patch_per_thread; - const int64_t patch_end = std::min(patch_start + patch_per_thread, patch_end_batch); - - for (int64_t p = patch_start; p < patch_end; ++p) { - const int64_t p_in_batch = p % (dst_w * dst_h * dst_d); - const int64_t p_in_depth = p_in_batch % (dst_w * dst_h); - const int64_t batch_idx = p / (dst_w * dst_h * dst_d); - const int64_t dst_z = p_in_batch / (dst_w * dst_h); - const int64_t dst_y = p_in_depth / dst_w; - const int64_t dst_x = p_in_depth % dst_w; - - char * dst_row = (char *) tmp + (p % patches_per_batch) * knl_n_total * traits->type_size; - - for (int64_t ic = 0; ic < c; ++ic) { - for (int64_t kz = 0; kz < knl_d; ++kz) { - for (int64_t ky = 0; ky < knl_h; ++ky) { - for (int64_t kx = 0; kx < knl_w; ++kx) { - const int64_t sz = dst_z * s2 + kz * d2 - p2; - const int64_t sy = dst_y * s1 + ky * d1 - p1; - const int64_t sx = dst_x * s0 + kx * d0 - p0; - - int64_t dst_idx = ic * knl_n_per_channel + kz * (knl_h * knl_w) + ky * knl_w + kx; - - float src_val; - if (sz < 0 || sz >= src_d || sy < 0 || sy >= src_h || sx < 0 || sx >= src_w) { - src_val = 0.0f; - } else { - const int64_t cn_idx = batch_idx * c + ic; - const float * src_ptr = (const float *)((const char *)src_data + sx*src->nb[0] + sy*src->nb[1] + sz*src->nb[2] + cn_idx*src->nb[3]); - src_val = *src_ptr; - } + int sector = (i0 / 2) % sect_dims; + if (indep_sects) { + // compute theta independently for each dim sections + // (i.e. reset corresponding theta when `i0` go from one section to another) + if (sector == 0) { + theta_t = theta_base_t; + } + else if (sector == sections[0]) { + theta_h = theta_base_h;; + } + else if (sector == sec_w) { + theta_w = theta_base_w; + } + else if (sector == sec_e) { + theta_e = theta_base_e; + } + } - char * element_ptr = dst_row + dst_idx * traits->type_size; - if (kernel_type == GGML_TYPE_F32) { - *(float *)element_ptr = src_val; - } else if (kernel_type == GGML_TYPE_F16) { - *(ggml_fp16_t *)element_ptr = GGML_CPU_FP32_TO_FP16(src_val); - } - } - } - } + float theta = theta_t; + if (is_imrope) { // qwen3vl apply interleaved mrope + if (sector % 3 == 1 && sector < 3 * sections[1]) { + theta = theta_h; + } else if (sector % 3 == 2 && sector < 3 * sections[2]) { + theta = theta_w; + } else if (sector % 3 == 0 && sector < 3 * sections[0]) { + theta = theta_t; + } else { + theta = theta_e; + } + } else { + if (sector >= sections[0] && sector < sec_w) { + theta = theta_h; + } + else if (sector >= sec_w && sector < sec_w + sections[2]) { + theta = theta_w; + } + else if (sector >= sec_w + sections[2]) { + theta = theta_e; } } - ggml_barrier(params->threadpool); + rope_yarn( + theta/ff, freq_scale, corr_dims, i0, ext_factor, mscale, &cache[i0 + 0], &cache[i0 + 1] + ); + cache[i0 + 1] *= sin_sign; - float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n_total * traits->type_size); - ggml_call_mul_mat(kernel_type, params, patch_n_in_batch, oc, knl_n_total, tmp, knl_data, gemm_output); + theta_t *= theta_scale; + theta_w *= theta_scale; + theta_h *= theta_scale; + theta_e *= theta_scale; + } +} - ggml_barrier(params->threadpool); - const int64_t permute_per_thread = (patch_n_in_batch + params->nth - 1) / params->nth; - const int64_t permute_start = params->ith * permute_per_thread; - const int64_t permute_end = std::min(permute_start + permute_per_thread, patch_n_in_batch); +template +static void rotate_pairs(const int64_t n, const int64_t n_offset, const float * cache, const T * src_data, T * dst_data, const int scale = 2) { + for (int64_t i0 = 0; i0 < n; i0 += 2) { + const int64_t ic = i0/scale; // hack for GGML_ROPE_TYPE_NORMAL, where we need ic = i0; for all other cases, ic = i0/2 - for (int64_t i = permute_start; i < permute_end; ++i) { - const int64_t p = patch_start_batch + i; - const int64_t p_in_batch = p % (dst_w * dst_h * dst_d); - const int64_t p_in_depth = p_in_batch % (dst_w * dst_h); - const int64_t batch_idx = p / (dst_w * dst_h * dst_d); - const int64_t dst_z = p_in_batch / (dst_w * dst_h); - const int64_t dst_y = p_in_depth / dst_w; - const int64_t dst_x = p_in_depth % dst_w; + const float cos_theta = cache[i0 + 0]; + const float sin_theta = cache[i0 + 1]; - for (int64_t ioc = 0; ioc < oc; ++ioc) { - const float value = gemm_output[i * oc + ioc]; - const int64_t ocn_idx = batch_idx * oc + ioc; - float * dst_ptr = (float *)((char *)dst_data + dst_x*dst->nb[0] + dst_y*dst->nb[1] + dst_z*dst->nb[2] + ocn_idx*dst->nb[3]); - *dst_ptr = value; - } - } - } + const T * const src = src_data + ic; + T * dst = dst_data + ic; + + const float x0 = type_conversion_table::to_f32(src[0]); + const float x1 = type_conversion_table::to_f32(src[n_offset]); + + dst[0] = type_conversion_table::from_f32(x0*cos_theta - x1*sin_theta); + dst[n_offset] = type_conversion_table::from_f32(x0*sin_theta + x1*cos_theta); + } } -void ggml_compute_forward_conv_3d( +template //float or ggml_fp16_t +static void ggml_compute_forward_rope_flt( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst, + const bool forward) { + const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - ggml_compute_forward_conv_3d_impl(params, src0, src1, dst, src0->type); -} + const ggml_tensor * src2 = dst->src[2]; -template -static void ggml_compute_forward_conv_transpose_2d_impl( - const ggml_compute_params * params, - ggml_tensor * dst) { + GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_I32); - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; + float freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow; + int sections[4]; - GGML_ASSERT(src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT( dst->type == GGML_TYPE_F32); + //const int n_past = ((int32_t *) dst->op_params)[0]; + const int n_dims = ((int32_t *) dst->op_params)[1]; + const int mode = ((int32_t *) dst->op_params)[2]; + //const int n_ctx = ((int32_t *) dst->op_params)[3]; + const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; - GGML_TENSOR_BINARY_OP_LOCALS + memcpy(&freq_base, (int32_t *) dst->op_params + 5, sizeof(float)); + memcpy(&freq_scale, (int32_t *) dst->op_params + 6, sizeof(float)); + memcpy(&ext_factor, (int32_t *) dst->op_params + 7, sizeof(float)); + memcpy(&attn_factor, (int32_t *) dst->op_params + 8, sizeof(float)); + memcpy(&beta_fast, (int32_t *) dst->op_params + 9, sizeof(float)); + memcpy(&beta_slow, (int32_t *) dst->op_params + 10, sizeof(float)); + memcpy(§ions, (int32_t *) dst->op_params + 11, sizeof(int)*4); + + GGML_TENSOR_UNARY_OP_LOCALS + + //printf("ne0: %d, ne1: %d, ne2: %d, ne3: %d\n", ne0, ne1, ne2, ne3); + //printf("n_past = %d, ne2 = %d\n", n_past, ne2); + + GGML_ASSERT(nb0 == nb00); + GGML_ASSERT(nb0 == sizeof(T)); const int ith = params->ith; const int nth = params->nth; - const int nk = ne00*ne01*ne02*ne03; + const int nr = ggml_nrows(dst); - GGML_ASSERT(nb00 == ggml_type_size(src0->type)); - GGML_ASSERT(nb10 == sizeof(float)); + GGML_ASSERT(n_dims <= ne0); + GGML_ASSERT(n_dims % 2 == 0); - if (ith == 0) { - memset(params->wdata, 0, params->wsize); + // rows per thread + const int dr = (nr + nth - 1)/nth; - // permute kernel data (src0) from (Kw x Kh x Cout x Cin) to (Cin x Kw x Kh x Cout) - { - kernel_t * const wdata = (kernel_t *) params->wdata + 0; + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - for (int64_t i03 = 0; i03 < ne03; i03++) { - for (int64_t i02 = 0; i02 < ne02; i02++) { - const kernel_t * const src = (kernel_t *)((char *) src0->data + i03*nb03 + i02*nb02); - kernel_t * dst_data = wdata + i02*ne01*ne00*ne03; - for (int64_t i01 = 0; i01 < ne01; i01++) { - for (int64_t i00 = 0; i00 < ne00; i00++) { - dst_data[i01*ne00*ne03 + i00*ne03 + i03] = src[i01 * ne00 + i00]; - } - } - } - } - } + // row index used to determine which thread to use + int ir = 0; - // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh) - { - kernel_t * const wdata = (kernel_t *) params->wdata + nk; - for (int i12 = 0; i12 < ne12; i12++) { - for (int i11 = 0; i11 < ne11; i11++) { - const float * const src = (float *)((char *) src1->data + i12*nb12 + i11*nb11); - kernel_t * dst_data = wdata + i11*ne10*ne12; - for (int i10 = 0; i10 < ne10; i10++) { - if constexpr (std::is_same_v) { - dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]); - } else { - dst_data[i10*ne12 + i12] = src[i10]; - } - } - } - } - } + const float theta_scale = powf(freq_base, -2.0f/n_dims); - memset(dst->data, 0, ggml_nbytes(dst)); + float corr_dims[2]; + ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); + + const bool is_imrope = mode == GGML_ROPE_TYPE_IMROPE; // qwen3vl apply interleaved mrope + const bool mrope_used = mode & GGML_ROPE_TYPE_MROPE; // ggml_rope_multi, note: also true for vision (24 & 8 == true) and for imrope + const bool is_vision = mode == GGML_ROPE_TYPE_VISION; + + if (mrope_used) { + GGML_ASSERT(sections[0] > 0 || sections[1] > 0 || sections[2] > 0); } - ggml_barrier(params->threadpool); - const int32_t stride = ggml_get_op_params_i32(dst, 0); + if (is_vision) { + GGML_ASSERT(n_dims == ne0/2); + } - // total patches in dst - const int np = ne2; + const float * freq_factors = NULL; + if (src2 != NULL) { + GGML_ASSERT(src2->type == GGML_TYPE_F32); + GGML_ASSERT(src2->ne[0] >= n_dims / 2); + freq_factors = (const float *) src2->data; + } - // patches per thread - const int dp = (np + nth - 1)/nth; + // backward process uses inverse rotation by cos and sin. + // cos and sin build a rotation matrix, where the inverse is the transpose. + // this essentially just switches the sign of sin. + const float sin_sign = forward ? 1.0f : -1.0f; - // patch range for this thread - const int ip0 = dp*ith; - const int ip1 = MIN(ip0 + dp, np); + const int32_t * pos = (const int32_t *) src1->data; - kernel_t * const wdata = (kernel_t *) params->wdata + 0; - kernel_t * const wdata_src = wdata + nk; + int64_t last_i2 = -1; - for (int i2 = ip0; i2 < ip1; i2++) { // Cout - float * dst_data = (float *)((char *) dst->data + i2*nb2); - kernel_t * wdata_kernel = wdata + i2*ne01*ne00*ne03; - for (int i11 = 0; i11 < ne11; i11++) { - for (int i10 = 0; i10 < ne10; i10++) { - const int i1n = i11*ne10*ne12 + i10*ne12; - for (int i01 = 0; i01 < ne01; i01++) { - for (int i00 = 0; i00 < ne00; i00++) { - float v = 0; - if constexpr (std::is_same_v) { - ggml_vec_dot_f16(ne03, &v, 0, - wdata_src + i1n, 0, - wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1); - } else { - ggml_vec_dot_f32(ne03, &v, 0, - wdata_src + i1n, 0, - wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1); - } - dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v; + for (int64_t i3 = 0; i3 < ne3; i3++) { // batch + for (int64_t i2 = 0; i2 < ne2; i2++) { // seq-len + for (int64_t i1 = 0; i1 < ne1; i1++) { // attn-heads + if (ir++ < ir0) continue; // skip rows mapped to other threads + if (ir > ir1) break; + + float * cache = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32)*ith; + if (last_i2 != i2) { + if (!mrope_used) { + const int64_t p = pos[i2]; + ggml_rope_cache_init(p, freq_scale, freq_factors, corr_dims, ne0, ext_factor, attn_factor, cache, sin_sign, theta_scale); + } + else { + const int64_t p_t = pos[i2]; + const int64_t p_h = pos[i2 + ne2]; + const int64_t p_w = pos[i2 + ne2 * 2]; + const int64_t p_e = pos[i2 + ne2 * 3]; + ggml_mrope_cache_init( + p_t, p_h, p_w, p_e, sections, is_imrope, is_vision, + freq_scale, freq_factors, corr_dims, ne0, ext_factor, attn_factor, cache, sin_sign, theta_scale); } + + last_i2 = i2; } - } + + T * src = (T *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01); + T * dst_data = (T *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1); + + switch (mode) { + case GGML_ROPE_TYPE_NORMAL: + rotate_pairs(n_dims, 1, cache, src, dst_data, 1); + break; + case GGML_ROPE_TYPE_NEOX: + case GGML_ROPE_TYPE_MROPE: + case GGML_ROPE_TYPE_IMROPE: + rotate_pairs(n_dims, n_dims/2, cache, src, dst_data); + break; + case GGML_ROPE_TYPE_VISION: + rotate_pairs(ne0, n_dims, cache, src, dst_data); + break; + default: + GGML_ABORT("rope type not supported"); + } + + if (!is_vision) { + // fill the remain channels with data from src tensor + for (int64_t i0 = n_dims; i0 < ne0; i0 += 2) { + const T * const src = (T *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); + T * dst_data = (T *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); + + dst_data[0] = src[0]; + dst_data[1] = src[1]; + } + } + } //attn-heads } } } -void ggml_compute_forward_conv_transpose_2d( +void ggml_compute_forward_rope( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { case GGML_TYPE_F16: { - ggml_compute_forward_conv_transpose_2d_impl(params, dst); + ggml_compute_forward_rope_flt(params, dst, true); } break; case GGML_TYPE_F32: { - ggml_compute_forward_conv_transpose_2d_impl(params, dst); + ggml_compute_forward_rope_flt(params, dst, true); } break; default: { @@ -7057,752 +7973,681 @@ void ggml_compute_forward_conv_transpose_2d( } } -// ggml_compute_forward_conv_2d_dw +// ggml_compute_forward_rope_back -struct ggml_conv_2d_dw_params { - int64_t channels; - int64_t batch; - int64_t src_w; - int64_t src_h; - int64_t dst_w; - int64_t dst_h; - int64_t knl_w; - int64_t knl_h; - int stride_x; - int stride_y; - int pad_x; - int pad_y; - int dilation_x; - int dilation_y; -}; +void ggml_compute_forward_rope_back( + const ggml_compute_params * params, + ggml_tensor * dst) { -static void ggml_compute_forward_conv_2d_dw_cwhn( + const ggml_tensor * src0 = dst->src[0]; + + switch (src0->type) { + case GGML_TYPE_F16: + { + ggml_compute_forward_rope_flt(params, dst, false); + } break; + case GGML_TYPE_F32: + { + ggml_compute_forward_rope_flt(params, dst, false); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +// ggml_compute_forward_conv_transpose_1d + +static void ggml_compute_forward_conv_transpose_1d_f16_f32( const ggml_compute_params * params, - const ggml_tensor * src, - const ggml_tensor * kernel, - ggml_tensor * dst, - const ggml_conv_2d_dw_params & p) { + ggml_tensor * dst) { - const int64_t c = p.channels; - const float * knl_data = (const float *)kernel->data; + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - const int64_t rows_total = p.dst_h * p.batch; - const int64_t rows_per_thread = (rows_total + params->nth - 1) / params->nth; - const int64_t row_start = params->ith * rows_per_thread; - const int64_t row_end = MIN(row_start + rows_per_thread, rows_total); + GGML_ASSERT(src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); -#ifdef GGML_SIMD - #if defined(__ARM_FEATURE_SVE) - const int64_t pkg_size = svcntw(); - #else - const int64_t pkg_size = GGML_F32_EPR; - #endif - const int64_t pkg_count = c / pkg_size; - const int64_t c_pkg_end = pkg_count * pkg_size; -#else - const int64_t c_pkg_end = 0; -#endif + GGML_TENSOR_BINARY_OP_LOCALS - for (int64_t row = row_start; row < row_end; ++row) { - const int64_t dst_y = row % p.dst_h; - const float * src_data = (const float *)src->data + (row / p.dst_h) * p.src_w * p.src_h * c; - for (int64_t dst_x = 0; dst_x < p.dst_w; ++dst_x) { - float * dst_data = (float *)dst->data + (row * p.dst_w + dst_x) * c; - const int64_t src_y_base = dst_y * p.stride_y - p.pad_y; - const int64_t src_x_base = dst_x * p.stride_x - p.pad_x; + const int ith = params->ith; + const int nth = params->nth; -#ifdef GGML_SIMD - // Vectorized loop - for (int64_t c_i = 0; c_i < c_pkg_end; c_i += pkg_size) { - GGML_F32_VEC sum = GGML_F32_VEC_ZERO; - for (int64_t knl_y = 0; knl_y < p.knl_h; ++knl_y) { - const int64_t src_y = src_y_base + knl_y * p.dilation_y; - if (src_y < 0 || src_y >= p.src_h) { - continue; - } - for (int64_t knl_x = 0; knl_x < p.knl_w; ++knl_x) { - const int64_t src_x = src_x_base + knl_x * p.dilation_x; - if (src_x < 0 || src_x >= p.src_w) { - continue; - } - GGML_F32_VEC k = GGML_F32_VEC_LOAD(knl_data + (knl_y * p.knl_w + knl_x) * c + c_i); - GGML_F32_VEC s = GGML_F32_VEC_LOAD(src_data + (src_y * p.src_w + src_x) * c + c_i); - sum = GGML_F32_VEC_FMA(sum, k, s); + const int nk = ne00*ne01*ne02; + + GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); + GGML_ASSERT(nb10 == sizeof(float)); + + if (ith == 0) { + memset(params->wdata, 0, params->wsize); + + // permute kernel data (src0) from (K x Cout x Cin) to (Cin x K x Cout) + { + ggml_fp16_t * const wdata = (ggml_fp16_t *) params->wdata + 0; + + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = 0; i01 < ne01; i01++) { + const ggml_fp16_t * const src = (ggml_fp16_t *)((char *) src0->data + i02*nb02 + i01*nb01); + ggml_fp16_t * dst_data = wdata + i01*ne00*ne02; + for (int64_t i00 = 0; i00 < ne00; i00++) { + dst_data[i00*ne02 + i02] = src[i00]; } } - GGML_F32_VEC_STORE(dst_data + c_i, sum); } -#endif - // Scalar loop - for (int64_t c_i = c_pkg_end; c_i < c; ++c_i) { - float sum = 0.0f; - for (int64_t knl_y = 0; knl_y < p.knl_h; ++knl_y) { - const int64_t src_y = src_y_base + knl_y * p.dilation_y; - if (src_y < 0 || src_y >= p.src_h) { - continue; - } - for (int64_t knl_x = 0; knl_x < p.knl_w; ++knl_x) { - const int64_t src_x = src_x_base + knl_x * p.dilation_x; - if (src_x < 0 || src_x >= p.src_w) { - continue; - } - sum += knl_data[(knl_y * p.knl_w + knl_x) * c + c_i] - * src_data[(src_y * p.src_w + src_x) * c + c_i]; - } + } + + // permute source data (src1) from (L x Cin) to (Cin x L) + { + ggml_fp16_t * const wdata = (ggml_fp16_t *) params->wdata + nk; + ggml_fp16_t * dst_data = wdata; + + for (int64_t i11 = 0; i11 < ne11; i11++) { + const float * const src = (float *)((char *) src1->data + i11*nb11); + for (int64_t i10 = 0; i10 < ne10; i10++) { + dst_data[i10*ne11 + i11] = GGML_CPU_FP32_TO_FP16(src[i10]); } - dst_data[c_i] = sum; } } + + // need to zero dst since we are accumulating into it + memset(dst->data, 0, ggml_nbytes(dst)); } -} + ggml_barrier(params->threadpool); -static void ggml_compute_forward_conv_2d_dw_whcn( - const ggml_compute_params * params, - const ggml_tensor * src, - const ggml_tensor * kernel, - ggml_tensor * dst, - const ggml_conv_2d_dw_params & p) { + const int32_t s0 = ((const int32_t*)(dst->op_params))[0]; - const int64_t n = p.channels * p.batch; - const int64_t per_thread = (n + params->nth - 1) / params->nth; - const int64_t start = params->ith * per_thread; - const int64_t end = MIN(start + per_thread, n); + // total rows in dst + const int nr = ne1; - for (int64_t i = start; i < end; ++i) { - const float * knl_data = (const float *)kernel->data + (i % p.channels) * p.knl_w * p.knl_h; - const float * src_data = (const float *)src->data + i * p.src_w * p.src_h; - float * dst_data = (float *)dst->data + i * p.dst_w * p.dst_h; + // rows per thread + const int dr = (nr + nth - 1)/nth; - for (int64_t dst_y = 0; dst_y < p.dst_h; ++dst_y) { - for (int64_t dst_x = 0; dst_x < p.dst_w; ++dst_x) { + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - float sum = 0.0f; - for (int64_t knl_y = 0; knl_y < p.knl_h; ++knl_y) { - const int64_t src_y = dst_y * p.stride_y + knl_y * p.dilation_y - p.pad_y; - if (src_y < 0 || src_y >= p.src_h) { - continue; - } - for (int64_t knl_x = 0; knl_x < p.knl_w; ++knl_x) { - const int64_t src_x = dst_x * p.stride_x + knl_x * p.dilation_x - p.pad_x; - if (src_x < 0 || src_x >= p.src_w) { - continue; - } - sum += knl_data[knl_y * p.knl_w + knl_x] - * src_data[src_y * p.src_w + src_x]; - } - } - dst_data[dst_y * p.dst_w + dst_x] = sum; + ggml_fp16_t * const wdata = (ggml_fp16_t *) params->wdata + 0; + ggml_fp16_t * const wdata_src = wdata + nk; + + for (int i1 = ir0; i1 < ir1; i1++) { + float * dst_data = (float *)((char *) dst->data + i1*nb1); + ggml_fp16_t * wdata_kernel = wdata + i1*ne02*ne00; + for (int i10 = 0; i10 < ne10; i10++) { + const int i1n = i10*ne11; + for (int i00 = 0; i00 < ne00; i00++) { + float v = 0; + ggml_vec_dot_f16(ne02, &v, 0, + (ggml_fp16_t *) wdata_src + i1n, 0, + (ggml_fp16_t *) wdata_kernel + i00*ne02, 0, 1); + dst_data[i10*s0 + i00] += v; } } } } -void ggml_compute_forward_conv_2d_dw( +static void ggml_compute_forward_conv_transpose_1d_f32( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { - const ggml_tensor * kernel = dst->src[0]; - const ggml_tensor * src = dst->src[1]; - ggml_conv_2d_dw_params p; - p.channels = src->ne[2]; - p.batch = src->ne[3]; - p.src_w = src->ne[0]; - p.src_h = src->ne[1]; - p.dst_w = dst->ne[0]; - p.dst_h = dst->ne[1]; - p.knl_w = kernel->ne[0]; - p.knl_h = kernel->ne[1]; - p.stride_x = dst->op_params[0]; - p.stride_y = dst->op_params[1]; - p.pad_x = dst->op_params[2]; - p.pad_y = dst->op_params[3]; - p.dilation_x = dst->op_params[4]; - p.dilation_y = dst->op_params[5]; + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(kernel->ne[3] == p.channels); - GGML_ASSERT(dst->ne[3] == p.batch); + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); - if (ggml_is_contiguous(src)) { - ggml_compute_forward_conv_2d_dw_whcn(params, src, kernel, dst, p); - } else if (ggml_is_contiguous_channels(src)) { - // kernel should also have channels most contiguous in memory - GGML_ASSERT(kernel->nb[0] >= kernel->nb[2] && kernel->nb[1] >= kernel->nb[0]); - ggml_compute_forward_conv_2d_dw_cwhn(params, src, kernel, dst, p); - } else { - GGML_ABORT("non-contiguous memory layout not supported"); - } -} + GGML_TENSOR_BINARY_OP_LOCALS -// ggml_compute_forward_pool_1d_ksp -static void ggml_compute_forward_pool_1d_ksp( - const ggml_compute_params * params, - const ggml_op_pool op, - const int k, - const int s, - const int p, - ggml_tensor * dst) { + const int ith = params->ith; + const int nth = params->nth; - const ggml_tensor * src = dst->src[0]; + const int nk = ne00*ne01*ne02; - assert(src->type == GGML_TYPE_F32 || src->type == GGML_TYPE_F16); + GGML_ASSERT(nb00 == sizeof(float)); + GGML_ASSERT(nb10 == sizeof(float)); - if (params->ith != 0) { - return; - } + if (ith == 0) { + memset(params->wdata, 0, params->wsize); - const int64_t IW = src->ne[0]; - const int64_t OW = dst->ne[0]; + // prepare kernel data (src0) from (K x Cout x Cin) to (Cin x K x Cout) + { + float * const wdata = (float *) params->wdata + 0; - const int64_t nr = ggml_nrows(src); + for (int64_t i02 = 0; i02 < ne02; i02++) { + for (int64_t i01 = 0; i01 < ne01; i01++) { + const float * const src = (float *)((char *) src0->data + i02*nb02 + i01*nb01); + float * dst_data = wdata + i01*ne00*ne02; + for (int64_t i00 = 0; i00 < ne00; i00++) { + dst_data[i00*ne02 + i02] = src[i00]; + } + } + } + } - for (int64_t ir = 0; ir < nr; ++ir) { - const char * srow_bytes = (const char *) src->data + ir * src->nb[1]; - float * drow = (float *) (( char *) dst->data + ir * dst->nb[1]); + // prepare source data (src1) + { + float * const wdata = (float *) params->wdata + nk; + float * dst_data = wdata; - for (int64_t ow = 0; ow < OW; ++ow) { - float res = 0; - switch (op) { - case GGML_OP_POOL_AVG: res = 0.0f; break; - case GGML_OP_POOL_MAX: res = -FLT_MAX; break; - case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error"); + for (int64_t i11 = 0; i11 < ne11; i11++) { + const float * const src = (float *)((char *) src1->data + i11*nb11); + for (int64_t i10 = 0; i10 < ne10; i10++) { + dst_data[i10*ne11 + i11] = src[i10]; + } } + } - int count = 0; - const int base = (int) ow * s - p; + // need to zero dst since we are accumulating into it + memset(dst->data, 0, ggml_nbytes(dst)); + } + ggml_barrier(params->threadpool); - for (int ki = 0; ki < k; ++ki) { - const int j = base + ki; - if (j < 0 || j >= (int) IW) { - continue; - } + const int32_t s0 = ((const int32_t*)(dst->op_params))[0]; - float v; - if (src->type == GGML_TYPE_F32) { - v = ((const float *) srow_bytes)[j]; - } else { - v = GGML_CPU_FP16_TO_FP32(((const ggml_fp16_t *) srow_bytes)[j]); - } + // total rows in dst + const int nr = ne1; - switch (op) { - case GGML_OP_POOL_AVG: res += v; break; - case GGML_OP_POOL_MAX: res = std::max(v, res); break; - case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error"); - } + // rows per thread + const int dr = (nr + nth - 1)/nth; - ++count; - } + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - switch (op) { - case GGML_OP_POOL_AVG: res = (count > 0) ? (res / count) : 0.0f; break; - case GGML_OP_POOL_MAX: break; - case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error"); - } + float * const wdata = (float *) params->wdata + 0; + float * const wdata_src = wdata + nk; - drow[ow] = res; + for (int i1 = ir0; i1 < ir1; i1++) { + float * dst_data = (float *)((char *) dst->data + i1*nb1); + float * wdata_kernel = wdata + i1*ne02*ne00; + for (int i10 = 0; i10 < ne10; i10++) { + const int i1n = i10*ne11; + for (int i00 = 0; i00 < ne00; i00++) { + float v = 0; + ggml_vec_dot_f32(ne02, &v, 0, + wdata_src + i1n, 0, + wdata_kernel + i00*ne02, 0, 1); + dst_data[i10*s0 + i00] += v; + } } } } -// ggml_compute_forward_pool_1d - -void ggml_compute_forward_pool_1d( +void ggml_compute_forward_conv_transpose_1d( const ggml_compute_params * params, ggml_tensor * dst) { - const int32_t * opts = (const int32_t *)dst->op_params; - ggml_op_pool op = static_cast(opts[0]); - const int k0 = opts[1]; - const int s0 = opts[2]; - const int p0 = opts[3]; + const ggml_tensor * src0 = dst->src[0]; - ggml_compute_forward_pool_1d_ksp(params, op, k0, s0, p0, dst); + switch (src0->type) { + case GGML_TYPE_F16: + { + ggml_compute_forward_conv_transpose_1d_f16_f32(params, dst); + } break; + case GGML_TYPE_F32: + { + ggml_compute_forward_conv_transpose_1d_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } } -// ggml_compute_forward_pool_2d - -void ggml_compute_forward_pool_2d( +// ggml_compute_forward_im2col_f32 +// src0: kernel [OC, IC, KH, KW] +// src1: image [N, IC, IH, IW] +// dst: result [N, OH, OW, IC*KH*KW] +static void ggml_compute_forward_im2col_f32( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { - const ggml_tensor * src = dst->src[0]; + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - assert(src->type == GGML_TYPE_F32 || src->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); - if (params->ith != 0) { - return; - } + GGML_TENSOR_BINARY_OP_LOCALS; - const int32_t * opts = (const int32_t *)dst->op_params; + const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; + const int32_t s1 = ((const int32_t *)(dst->op_params))[1]; + const int32_t p0 = ((const int32_t *)(dst->op_params))[2]; + const int32_t p1 = ((const int32_t *)(dst->op_params))[3]; + const int32_t d0 = ((const int32_t *)(dst->op_params))[4]; + const int32_t d1 = ((const int32_t *)(dst->op_params))[5]; + const bool is_2D = ((const int32_t *)(dst->op_params))[6] == 1; - ggml_op_pool op = static_cast(opts[0]); - const int k0 = opts[1]; - const int k1 = opts[2]; - const int s0 = opts[3]; - const int s1 = opts[4]; - const int p0 = opts[5]; - const int p1 = opts[6]; - const char * cdata = (const char*)src->data; - const char * const data_end = cdata + ggml_nbytes(src); + const int ith = params->ith; + const int nth = params->nth; - const int64_t px = dst->ne[0]; - const int64_t py = dst->ne[1]; - const int64_t pa = px * py; + const int64_t N = is_2D ? ne13 : ne12; + const int64_t IC = is_2D ? ne12 : ne11; + const int64_t IH = is_2D ? ne11 : 1; + const int64_t IW = ne10; - float * dplane = (float *)dst->data; + const int64_t KH = is_2D ? ne01 : 1; + const int64_t KW = ne00; - const int ka = k0 * k1; - const int offset0 = -p0; - const int offset1 = -p1; + const int64_t OH = is_2D ? ne2 : 1; + const int64_t OW = ne1; - while (cdata < data_end) { - for (int oy = 0; oy < py; ++oy) { - float * const drow = dplane + oy * px; - float * const out = drow; + int ofs0 = is_2D ? nb13 : nb12; + int ofs1 = is_2D ? nb12 : nb11; - for (int ox = 0; ox < px; ++ox) { - float res = 0; - switch (op) { - case GGML_OP_POOL_AVG: res = 0; break; - case GGML_OP_POOL_MAX: res = -FLT_MAX; break; - case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error"); - } + GGML_ASSERT(nb10 == sizeof(float)); - const int ix = offset0 + ox * s0; - const int iy = offset1 + oy * s1; + // im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW] + { + float * const wdata = (float *) dst->data; - for (int ky = 0; ky < k1; ++ky) { - if (iy + ky < 0 || iy + ky >= src->ne[1]) { - continue; - } + for (int64_t in = 0; in < N; in++) { + for (int64_t ioh = 0; ioh < OH; ioh++) { // 1 + for (int64_t iow = 0; iow < OW; iow++) { + for (int64_t iic = ith; iic < IC; iic += nth) { - const void * srow = (const void *)(cdata + src->nb[1] * (iy + ky)); - for (int kx = 0; kx < k0; ++kx) { - int j = ix + kx; - if (j < 0 || j >= src->ne[0]) { - continue; - } + // micro kernel + float * dst_data = wdata + (in*OH*OW + ioh*OW + iow)*(IC*KH*KW); // [IC, KH, KW] + const float * const src_data = (float *)((char *) src1->data + in*ofs0 + iic*ofs1); // [IH, IW] - const float srow_j = (src->type == GGML_TYPE_F32) ? ((const float*)srow)[j] : GGML_CPU_FP16_TO_FP32(((const ggml_fp16_t*)srow)[j]); - switch (op) { - case GGML_OP_POOL_AVG: res += srow_j; break; - case GGML_OP_POOL_MAX: res = std::max(srow_j, res); break; - case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error"); + for (int64_t ikh = 0; ikh < KH; ikh++) { // 1 + for (int64_t ikw = 0; ikw < KW; ikw++) { + const int64_t iiw = iow*s0 + ikw*d0 - p0; + const int64_t iih = ioh*s1 + ikh*d1 - p1; + + if (iih < 0 || iih >= IH || iiw < 0 || iiw >= IW) { + dst_data[iic*(KH*KW) + ikh*KW + ikw] = 0; + } else { + dst_data[iic*(KH*KW) + ikh*KW + ikw] = (src_data[iih*IW + iiw]); + } + } } } } - switch (op) { - case GGML_OP_POOL_AVG: res /= ka; break; - case GGML_OP_POOL_MAX: break; - case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error"); - } - - out[ox] = res; } } - - cdata += src->nb[2]; - dplane += pa; } } -// ggml_compute_forward_pool_2d_back -void ggml_compute_forward_pool_2d_back( +// ggml_compute_forward_im2col_f16 +// src0: kernel [OC, IC, KH, KW] +// src1: image [N, IC, IH, IW] +// dst: result [N, OH, OW, IC*KH*KW] +static void ggml_compute_forward_im2col_f16( const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src = dst->src[0]; - const ggml_tensor * dstf = dst->src[1]; // forward tensor of dst - - assert(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16); - - if (params->ith != 0) { - return; - } - - const int32_t * opts = (const int32_t *)dst->op_params; - ggml_op_pool op = static_cast(opts[0]); - const int k0 = opts[1]; - const int k1 = opts[2]; - const int s0 = opts[3]; - const int s1 = opts[4]; - const int p0 = opts[5]; - const int p1 = opts[6]; + ggml_tensor * dst) { - char * cdata = (char *) dst->data; - const char * cdataf = (const char *) dstf->data; - const char * const data_end = cdata + ggml_nbytes(dst); + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(params->ith == 0); - memset(cdata, 0, ggml_nbytes(dst)); + GGML_ASSERT(src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F16); - const int64_t px = src->ne[0]; - const int64_t py = src->ne[1]; - const int64_t pa = px * py; + GGML_TENSOR_BINARY_OP_LOCALS; - const float * splane = (const float *) src->data; + const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; + const int32_t s1 = ((const int32_t *)(dst->op_params))[1]; + const int32_t p0 = ((const int32_t *)(dst->op_params))[2]; + const int32_t p1 = ((const int32_t *)(dst->op_params))[3]; + const int32_t d0 = ((const int32_t *)(dst->op_params))[4]; + const int32_t d1 = ((const int32_t *)(dst->op_params))[5]; + const bool is_2D = ((const int32_t *)(dst->op_params))[6] == 1; - const int ka = k0 * k1; - const int offset0 = -p0; - const int offset1 = -p1; + const int ith = params->ith; + const int nth = params->nth; - while (cdata < data_end) { - for (int oy = 0; oy < py; ++oy) { - const float * const srow = splane + oy * px; - for (int ox = 0; ox < px; ++ox) { - const float grad0 = srow[ox]; + const int64_t N = is_2D ? ne13 : ne12; + const int64_t IC = is_2D ? ne12 : ne11; + const int64_t IH = is_2D ? ne11 : 1; + const int64_t IW = ne10; - const int ix = offset0 + ox * s0; - const int iy = offset1 + oy * s1; + const int64_t KH = is_2D ? ne01 : 1; + const int64_t KW = ne00; - if (op == GGML_OP_POOL_MAX) { - float maxval = -FLT_MAX; - int kxmax = -1; - int kymax = -1; + const int64_t OH = is_2D ? ne2 : 1; + const int64_t OW = ne1; - for (int ky = 0; ky < k1; ++ky) { - if (iy + ky < 0 || iy + ky >= dst->ne[1]) { - continue; - } - const void * drowf = (const void *)(cdataf + dst->nb[1] * (iy + ky)); - for (int kx = 0; kx < k0; ++kx) { - int j = ix + kx; - if (j < 0 || j >= dst->ne[0]) { - continue; - } + int ofs0 = is_2D ? nb13 : nb12; + int ofs1 = is_2D ? nb12 : nb11; - const float val = dst->type == GGML_TYPE_F32 ? - ((const float *) drowf)[j] : GGML_CPU_FP16_TO_FP32(((const ggml_fp16_t *) drowf)[j]); - if (val <= maxval) { - continue; - } + GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); + GGML_ASSERT(nb10 == ggml_type_size(src1->type)); - maxval = val; - kxmax = kx; - kymax = ky; - } - } + // im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW] + { + ggml_fp16_t * const wdata = (ggml_fp16_t *) dst->data; - if (kxmax == -1 || kymax == -1) { - continue; - } + for (int64_t in = 0; in < N; in++) { + for (int64_t ioh = 0; ioh < OH; ioh++) { // 1 + for (int64_t iow = 0; iow < OW; iow++) { + for (int64_t iic = ith; iic < IC; iic += nth) { - void * drow = (void *)(cdata + dst->nb[1] * (iy + kymax)); - const int j = ix + kxmax; - if (dst->type == GGML_TYPE_F32) { - ((float *) drow)[j] += grad0; - } else { - ((ggml_fp16_t *) drow)[j] = GGML_CPU_FP32_TO_FP16(grad0 + GGML_CPU_FP16_TO_FP32(((const ggml_fp16_t *) drow)[j])); - } - } else if (op == GGML_OP_POOL_AVG) { - const float grad = grad0 / ka; + // micro kernel + ggml_fp16_t * dst_data = wdata + (in*OH*OW + ioh*OW + iow)*(IC*KH*KW); // [IC, KH, KW] + const float * const src_data_f32 = src1->type == GGML_TYPE_F32 + ? (const float *)((const char *) src1->data + in*ofs0 + iic*ofs1) + : nullptr; // [IH, IW] + const ggml_fp16_t * const src_data_f16 = src1->type == GGML_TYPE_F16 + ? (const ggml_fp16_t *)((const char *) src1->data + in*ofs0 + iic*ofs1) + : nullptr; // [IH, IW] - for (int ky = 0; ky < k1; ++ky) { - if (iy + ky < 0 || iy + ky >= dst->ne[1]) { - continue; - } - void * drow = (void *)(cdata + dst->nb[1] * (iy + ky)); - for (int kx = 0; kx < k0; ++kx) { - int j = ix + kx; - if (j < 0 || j >= dst->ne[0]) { - continue; - } + for (int64_t ikh = 0; ikh < KH; ikh++) { // 1 + for (int64_t ikw = 0; ikw < KW; ikw++) { + const int64_t iiw = iow*s0 + ikw*d0 - p0; + const int64_t iih = ioh*s1 + ikh*d1 - p1; - if (dst->type == GGML_TYPE_F32) { - ((float *) drow)[j] += grad; - } else { - ((ggml_fp16_t *) drow)[j] += GGML_CPU_FP32_TO_FP16(grad); + if (iih < 0 || iih >= IH || iiw < 0 || iiw >= IW) { + dst_data[iic*(KH*KW) + ikh*KW + ikw] = 0; + } else { + if (src_data_f32 != nullptr) { + dst_data[iic*(KH*KW) + ikh*KW + ikw] = GGML_CPU_FP32_TO_FP16(src_data_f32[iih*IW + iiw]); + } else { + dst_data[iic*(KH*KW) + ikh*KW + ikw] = src_data_f16[iih*IW + iiw]; + } + } } } } - } else { - GGML_ASSERT(false); } } } + } +} - cdata += dst->nb[2]; - cdataf += dst->nb[2]; - splane += pa; +void ggml_compute_forward_im2col( + const ggml_compute_params * params, + ggml_tensor * dst) { + switch (dst->type) { + case GGML_TYPE_F16: + { + ggml_compute_forward_im2col_f16(params, dst); + } break; + case GGML_TYPE_F32: + { + ggml_compute_forward_im2col_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } } } -// ggml_compute_forward_upscale +// ggml_compute_forward_im2col_back_f32 -static void ggml_compute_forward_upscale_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_im2col_back_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src0 = dst->src[0]; // gradients of forward pass output + const ggml_tensor * src1 = dst->src[1]; // convolution kernel GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); + + GGML_TENSOR_BINARY_OP_LOCALS; + + const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; + const int32_t s1 = ((const int32_t *)(dst->op_params))[1]; + const int32_t p0 = ((const int32_t *)(dst->op_params))[2]; + const int32_t p1 = ((const int32_t *)(dst->op_params))[3]; + const int32_t d0 = ((const int32_t *)(dst->op_params))[4]; + const int32_t d1 = ((const int32_t *)(dst->op_params))[5]; + const bool is_2D = ((const int32_t *)(dst->op_params))[6] == 1; const int ith = params->ith; const int nth = params->nth; - GGML_TENSOR_UNARY_OP_LOCALS - - float sf0 = (float)ne0/src0->ne[0]; - float sf1 = (float)ne1/src0->ne[1]; - float sf2 = (float)ne2/src0->ne[2]; - float sf3 = (float)ne3/src0->ne[3]; - float pixel_offset = 0.5f; - - const int32_t mode_flags = ggml_get_op_params_i32(dst, 0); - const ggml_scale_mode mode = (ggml_scale_mode) (mode_flags & 0xFF); + const int64_t N = is_2D ? ne3 : ne2; + const int64_t IC = is_2D ? ne2 : ne1; + const int64_t IH = is_2D ? ne1 : 1; + const int64_t IW = ne0; - if (mode_flags & GGML_SCALE_FLAG_ALIGN_CORNERS) { - pixel_offset = 0.0f; - sf0 = ne0 > 1 && ne00 > 1 ? (float)(ne0 - 1) / (ne00 - 1) : sf0; - sf1 = ne1 > 1 && ne01 > 1 ? (float)(ne1 - 1) / (ne01 - 1) : sf1; - } + const int64_t KH = is_2D ? ne11 : 1; + const int64_t KW = ne10; - if (mode == GGML_SCALE_MODE_NEAREST) { - for (int64_t i3 = 0; i3 < ne3; i3++) { - const int64_t i03 = i3 / sf3; - for (int64_t i2 = ith; i2 < ne2; i2 += nth) { - const int64_t i02 = i2 / sf2; - for (int64_t i1 = 0; i1 < ne1; i1++) { - const int64_t i01 = i1 / sf1; - for (int64_t i0 = 0; i0 < ne0; i0++) { - const int64_t i00 = i0 / sf0; + const int64_t OH = is_2D ? ne02 : 1; + const int64_t OW = ne01; - const float * x = (float *)((char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); - float * y = (float *)((char *) dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); + int ofs0 = is_2D ? nb3 : nb2; + int ofs1 = is_2D ? nb2 : nb1; - *y = *x; - } - } - } - } - } else if (mode == GGML_SCALE_MODE_BILINEAR && (mode_flags & GGML_SCALE_FLAG_ANTIALIAS)) { - // Similar to F.interpolate(..., mode="bilinear", align_corners=False, antialias=True) - // https://github.com/pytorch/pytorch/blob/8871ff29b743948d1225389d5b7068f37b22750b/aten/src/ATen/native/cpu/UpSampleKernel.cpp - auto triangle_filter = [](float x) -> float { - return std::max(1.0f - fabsf(x), 0.0f); - }; + GGML_ASSERT(nb0 == sizeof(float)); - // support and invscale, minimum 1 pixel for bilinear - const float support1 = std::max(1.0f, 1.0f / sf1); - const float invscale1 = 1.0f / support1; - const float support0 = std::max(1.0f, 1.0f / sf0); - const float invscale0 = 1.0f / support0; + // im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW] + { + float * const wdata = (float *) dst->data; - for (int64_t i3 = 0; i3 < ne3; i3++) { - const int64_t i03 = i3 / sf3; - for (int64_t i2 = ith; i2 < ne2; i2 += nth) { - const int64_t i02 = i2 / sf2; - for (int64_t i1 = 0; i1 < ne1; i1++) { - const float y = ((float) i1 + pixel_offset) / sf1; - for (int64_t i0 = 0; i0 < ne0; i0++) { - const float x = ((float) i0 + pixel_offset) / sf0; + for (int64_t in = 0; in < N; in++) { + for (int64_t iic = ith; iic < IC; iic += nth) { + for (int64_t iih = 0; iih < IH; iih++) { + for (int64_t iiw = 0; iiw < IW; iiw++) { - // the range of source pixels that contribute - const int64_t x_min = std::max(x - support0 + pixel_offset, 0); - const int64_t x_max = std::min(x + support0 + pixel_offset, ne00); - const int64_t y_min = std::max(y - support1 + pixel_offset, 0); - const int64_t y_max = std::min(y + support1 + pixel_offset, ne01); + // micro kernel + float grad = 0.0f; + for (int64_t ikh = 0; ikh < KH; ikh++) { + for (int64_t ikw = 0; ikw < KW; ikw++) { + // For s0 > 1 some values were skipped over in the forward pass. + // These values have tmpw % s0 != 0 and need to be skipped in the backwards pass as well. + const int64_t tmpw = (iiw + p0 - ikw*d0); + if (tmpw % s0 != 0) { + continue; + } + const int64_t iow = tmpw / s0; - // bilinear filter with antialiasing - float val = 0.0f; - float total_weight = 0.0f; + // Equivalent logic as above except for s1. + int64_t ioh; + if (is_2D) { + const int64_t tmph = iih + p1 - ikh*d1; - for (int64_t sy = y_min; sy < y_max; sy++) { - const float weight_y = triangle_filter((sy - y + pixel_offset) * invscale1); + if (tmph % s1 != 0) { + continue; + } - for (int64_t sx = x_min; sx < x_max; sx++) { - const float weight_x = triangle_filter((sx - x + pixel_offset) * invscale0); - const float weight = weight_x * weight_y; + ioh = tmph / s1; + } else { + ioh = 0; + } - if (weight <= 0.0f) { + if (iow < 0 || iow >= OW || ioh < 0 || ioh >= OH) { continue; } - const float pixel = *(const float *)((const char *)src0->data + sx*nb00 + sy*nb01 + i02*nb02 + i03*nb03); - val += pixel * weight; - total_weight += weight; + const float * const grad_in = (const float *) src0->data + + (in*OH*OW + ioh*OW + iow)*(IC*KH*KW); // [IC, KH, KW] + grad += grad_in[iic*(KH*KW) + ikh*KW + ikw]; } } - - if (total_weight > 0.0f) { - val /= total_weight; - } - - float * dst_ptr = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); - *dst_ptr = val; + float * dst_data = (float *)((char *) wdata + (in*ofs0 + iic*ofs1)); // [IH, IW] + dst_data[iih*IW + iiw] = grad; } } } } - } else if (mode == GGML_SCALE_MODE_BILINEAR) { - for (int64_t i3 = 0; i3 < ne3; i3++) { - const int64_t i03 = i3 / sf3; - for (int64_t i2 = ith; i2 < ne2; i2 += nth) { - const int64_t i02 = i2 / sf2; - for (int64_t i1 = 0; i1 < ne1; i1++) { - const float y = ((float)i1 + pixel_offset) / sf1 - pixel_offset; - int64_t y0 = (int64_t)floorf(y); - int64_t y1 = y0 + 1; + } +} - y0 = std::max(int64_t(0), std::min(y0, ne01 - 1)); - y1 = std::max(int64_t(0), std::min(y1, ne01 - 1)); - float dy = y - (float)y0; - dy = std::max(0.0f, std::min(dy, 1.0f)); +// ggml_compute_forward_im2col_3d_f16 +// src0: kernel [OC*IC, KD, KH, KW] +// src1: image [N*IC, ID, IH, IW] +// dst: result [N*OD, OH, OW, IC * KD * KH * KW] +static void ggml_compute_forward_im2col_3d_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { - for (int64_t i0 = 0; i0 < ne0; i0++) { - const float x = ((float)i0 + pixel_offset) / sf0 - pixel_offset; - int64_t x0 = (int64_t)floorf(x); - int64_t x1 = x0 + 1; + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - x0 = std::max(int64_t(0), std::min(x0, ne00 - 1)); - x1 = std::max(int64_t(0), std::min(x1, ne00 - 1)); + GGML_ASSERT(src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F16); - float dx = x - (float)x0; - dx = std::max(0.0f, std::min(dx, 1.0f)); + GGML_TENSOR_BINARY_OP_LOCALS; - // fetch the four surrounding pixel values and interpolate - const float a = *(const float *)((const char *)src0->data + x0*nb00 + y0*nb01 + i02*nb02 + i03*nb03); - const float b = *(const float *)((const char *)src0->data + x1*nb00 + y0*nb01 + i02*nb02 + i03*nb03); - const float c = *(const float *)((const char *)src0->data + x0*nb00 + y1*nb01 + i02*nb02 + i03*nb03); - const float d = *(const float *)((const char *)src0->data + x1*nb00 + y1*nb01 + i02*nb02 + i03*nb03); + const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; + const int32_t s1 = ((const int32_t *)(dst->op_params))[1]; + const int32_t s2 = ((const int32_t *)(dst->op_params))[2]; + const int32_t p0 = ((const int32_t *)(dst->op_params))[3]; + const int32_t p1 = ((const int32_t *)(dst->op_params))[4]; + const int32_t p2 = ((const int32_t *)(dst->op_params))[5]; + const int32_t d0 = ((const int32_t *)(dst->op_params))[6]; + const int32_t d1 = ((const int32_t *)(dst->op_params))[7]; + const int32_t d2 = ((const int32_t *)(dst->op_params))[8]; + const int32_t IC = ((const int32_t *)(dst->op_params))[9]; - const float val = a*(1 - dx)*(1 - dy) + b*dx*(1 - dy) + c*(1 - dx)*dy + d*dx*dy; - float * y_dst = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); - *y_dst = val; - } - } - } - } - } else if (mode == GGML_SCALE_MODE_BICUBIC) { - // https://en.wikipedia.org/wiki/Bicubic_interpolation#Bicubic_convolution_algorithm - const float a = -0.75f; // use alpha = -0.75 (same as PyTorch) - auto weight1 = [a](float x) { return ((a + 2) * x - (a + 3)) * x * x + 1; }; - auto weight2 = [a](float x) { return ((a * x - 5 * a) * x + 8 * a) * x - 4 * a; }; - auto bicubic = [=](float p0, float p1, float p2, float p3, float x) { - const float w0 = weight2(x + 1); - const float w1 = weight1(x + 0); - const float w2 = weight1(1 - x); - const float w3 = weight2(2 - x); - return p0*w0 + p1*w1 + p2*w2 + p3*w3; - }; + const int ith = params->ith; + const int nth = params->nth; - for (int64_t i3 = 0; i3 < ne3; i3++) { - const int64_t i03 = i3 / sf3; - for (int64_t i2 = ith; i2 < ne2; i2 += nth) { - const int64_t i02 = i2 / sf2; - for (int64_t i1 = 0; i1 < ne1; i1++) { - const float y = ((float)i1 + pixel_offset) / sf1 - pixel_offset; - const int64_t y0 = (int64_t)floorf(y); - const float dy = y - (float)y0; + const int64_t N = ne13 / IC; + const int64_t ID = ne12; + const int64_t IH = ne11; + const int64_t IW = ne10; - for (int64_t i0 = 0; i0 < ne0; i0++) { - const float x = ((float)i0 + pixel_offset) / sf0 - pixel_offset; - const int64_t x0 = (int64_t)floorf(x); - const float dx = x - (float)x0; + const int64_t OC = ne03 / IC; + GGML_UNUSED(OC); + const int64_t KD = ne02; + const int64_t KH = ne01; + const int64_t KW = ne00; - auto p = [=](int64_t x_off, int64_t y_off) -> float { - int64_t i00 = std::max(int64_t(0), std::min(x0 + x_off, ne00 - 1)); - int64_t i01 = std::max(int64_t(0), std::min(y0 + y_off, ne01 - 1)); - return *(const float *)((const char *)src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); - }; + const int64_t OD = ne3 / N; + const int64_t OH = ne2; + const int64_t OW = ne1; + const int64_t OH_OW = OH*OW; + const int64_t KD_KH_KW = KD*KH*KW; + const int64_t KH_KW = KH*KW; + const int64_t IC_KD_KH_KW = IC*KD*KH*KW; - const float val = bicubic( - bicubic(p(-1,-1), p(0,-1), p(1,-1), p(2,-1), dx), - bicubic(p(-1, 0), p(0, 0), p(1, 0), p(2, 0), dx), - bicubic(p(-1, 1), p(0, 1), p(1, 1), p(2, 1), dx), - bicubic(p(-1, 2), p(0, 2), p(1, 2), p(2, 2), dx), dy); + GGML_ASSERT(nb10 == sizeof(float)); - float * y_dst = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); - *y_dst = val; - } - } - } - } - } else { - GGML_ABORT("unsupported upscale mode"); - } -} + // im2col: [N*IC, ID, IH, IW] => [N*OD, OH, OW, IC * KD * KH * KW] + { + ggml_fp16_t * const wdata = (ggml_fp16_t *) dst->data; -void ggml_compute_forward_upscale( - const ggml_compute_params * params, - ggml_tensor * dst) { + for (int64_t in = 0; in < N; in++) { + for (int64_t iod = 0; iod < OD; iod++) { + for (int64_t ioh = 0; ioh < OH; ioh++) { + for (int64_t iow = 0; iow < OW; iow++) { + for (int64_t iic = ith; iic < IC; iic += nth) { - const ggml_tensor * src0 = dst->src[0]; + // micro kernel + ggml_fp16_t * dst_data = wdata + (in*OD*OH_OW + iod*OH_OW + ioh*OW + iow)*IC_KD_KH_KW; // [IC, KD, KH, KW] + const float * const src_data = (const float *) ((const char *)src1->data + (in*IC + iic)*nb13); // [ID, IH, IW] - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_upscale_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); + for (int64_t ikd = 0; ikd < KD; ikd++) { + for (int64_t ikh = 0; ikh < KH; ikh++) { + for (int64_t ikw = 0; ikw < KW; ikw++) { + const int64_t iiw = iow*s0 + ikw*d0 - p0; + const int64_t iih = ioh*s1 + ikh*d1 - p1; + const int64_t iid = iod*s2 + ikd*d2 - p2; + + if (iid < 0 || iid >= ID || iih < 0 || iih >= IH || iiw < 0 || iiw >= IW) { + dst_data[iic*KD_KH_KW + ikd * KH_KW + ikh*KW + ikw] = 0; + } else { + const float * const s = (const float *) ((const char *)src_data + iid*nb12 + iih*nb11 + iiw*nb10); // [ID, IH, IW] + dst_data[iic*KD_KH_KW + ikd * KH_KW + ikh*KW + ikw] = GGML_CPU_FP32_TO_FP16(*s); + } + } + } + } + } + } + } } + } } } +// ggml_compute_forward_im2col_3d_f32 +// src0: kernel [OC*IC, KD, KH, KW] +// src1: image [N*IC, ID, IH, IW] +// dst: result [N*OD, OH, OW, IC * KD * KH * KW] +static void ggml_compute_forward_im2col_3d_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; -// ggml_compute_forward_pad + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); -template -static void ggml_compute_forward_pad_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { + GGML_TENSOR_BINARY_OP_LOCALS; - const ggml_tensor * src0 = dst->src[0]; + const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; + const int32_t s1 = ((const int32_t *)(dst->op_params))[1]; + const int32_t s2 = ((const int32_t *)(dst->op_params))[2]; + const int32_t p0 = ((const int32_t *)(dst->op_params))[3]; + const int32_t p1 = ((const int32_t *)(dst->op_params))[4]; + const int32_t p2 = ((const int32_t *)(dst->op_params))[5]; + const int32_t d0 = ((const int32_t *)(dst->op_params))[6]; + const int32_t d1 = ((const int32_t *)(dst->op_params))[7]; + const int32_t d2 = ((const int32_t *)(dst->op_params))[8]; + const int32_t IC = ((const int32_t *)(dst->op_params))[9]; - assert(dst->nb[0] == sizeof(float)); const int ith = params->ith; const int nth = params->nth; - GGML_TENSOR_UNARY_OP_LOCALS + const int64_t N = ne13 / IC; + const int64_t ID = ne12; + const int64_t IH = ne11; + const int64_t IW = ne10; - float * dst_ptr = (float *) dst->data; - const int32_t lp0 = ggml_get_op_params_i32(dst, 0); - const int32_t rp0 = ggml_get_op_params_i32(dst, 1); - const int32_t lp1 = ggml_get_op_params_i32(dst, 2); - const int32_t rp1 = ggml_get_op_params_i32(dst, 3); - const int32_t lp2 = ggml_get_op_params_i32(dst, 4); - const int32_t rp2 = ggml_get_op_params_i32(dst, 5); - const int32_t lp3 = ggml_get_op_params_i32(dst, 6); - const int32_t rp3 = ggml_get_op_params_i32(dst, 7); + const int64_t OC = ne03 / IC; + GGML_UNUSED(OC); + const int64_t KD = ne02; + const int64_t KH = ne01; + const int64_t KW = ne00; - // TODO: optimize + const int64_t OD = ne3 / N; + const int64_t OH = ne2; + const int64_t OW = ne1; - for (int64_t i2 = 0; i2 < ne2; ++i2) { - for (int64_t i1 = ith; i1 < ne1; i1 += nth) { - for (int64_t i0 = 0; i0 < ne0; ++i0) { - for (int64_t i3 = 0; i3 < ne3; ++i3) { - // circular means wrap around on a torus, so x and y loop around - if constexpr (circular_t) { - const int64_t dst_idx = i3*(ne0*ne1*ne2) + i2*(ne0*ne1) + i1*ne0 + i0; - const int64_t src_i0 = ggml_wrap_around(i0 - lp0, ne00); - const int64_t src_i1 = ggml_wrap_around(i1 - lp1, ne01); - const int64_t src_i2 = ggml_wrap_around(i2 - lp2, ne02); - const int64_t src_i3 = ggml_wrap_around(i3 - lp3, ne03); + const int64_t OH_OW = OH*OW; + const int64_t KD_KH_KW = KD*KH*KW; + const int64_t KH_KW = KH*KW; + const int64_t IC_KD_KH_KW = IC*KD*KH*KW; - const int64_t src_idx = - src_i3*nb03 + - src_i2*nb02 + - src_i1*nb01 + - src_i0*nb00; + GGML_ASSERT(nb10 == sizeof(float)); - const float * src_ptr = (const float *)((char *) src0->data + src_idx); - dst_ptr[dst_idx] = *src_ptr; - } else { - const int64_t dst_idx = i3*(ne0*ne1*ne2) + i2*(ne0*ne1) + i1*ne0 + i0; - if ((i0 >= lp0 && i0 < ne0 - rp0) \ - && (i1 >= lp1 && i1 < ne1 - rp1) \ - && (i2 >= lp2 && i2 < ne2 - rp2) \ - && (i3 >= lp3 && i3 < ne3 - rp3)) { - const int64_t src_idx = (i3 - lp3)*nb03 + (i2 - lp2)*nb02 + (i1 - lp1)*nb01 + (i0 - lp0)*nb00; - const float * src_ptr = (const float *)((char *) src0->data + src_idx); - dst_ptr[dst_idx] = *src_ptr; - } else { - dst_ptr[dst_idx] = 0; + // im2col: [N*IC, ID, IH, IW] => [N*OD, OH, OW, IC * KD * KH * KW] + { + float * const wdata = (float *) dst->data; + + for (int64_t in = 0; in < N; in++) { + for (int64_t iod = 0; iod < OD; iod++) { + for (int64_t ioh = 0; ioh < OH; ioh++) { + for (int64_t iow = 0; iow < OW; iow++) { + for (int64_t iic = ith; iic < IC; iic += nth) { + + // micro kernel + float * dst_data = wdata + (in*OD*OH_OW + iod*OH_OW + ioh*OW + iow)*IC_KD_KH_KW; // [IC, KD, KH, KW] + const float * const src_data = (const float *) ((const char *)src1->data + (in*IC + iic)*nb13); // [ID, IH, IW] + + for (int64_t ikd = 0; ikd < KD; ikd++) { + for (int64_t ikh = 0; ikh < KH; ikh++) { + for (int64_t ikw = 0; ikw < KW; ikw++) { + const int64_t iiw = iow*s0 + ikw*d0 - p0; + const int64_t iih = ioh*s1 + ikh*d1 - p1; + const int64_t iid = iod*s2 + ikd*d2 - p2; + + if (iid < 0 || iid >= ID || iih < 0 || iih >= IH || iiw < 0 || iiw >= IW || iid < 0 || iid >= ID) { + dst_data[iic*KD_KH_KW + ikd * KH_KW + ikh*KW + ikw] = 0; + } else { + const float * const s = (const float *) ((const char *)src_data + iid*nb12 + iih*nb11 + iiw*nb10); // [ID, IH, IW] + dst_data[iic*KD_KH_KW + ikd * KH_KW + ikh*KW + ikw] = *s; + } + } + } + } } } } @@ -7812,19 +8657,17 @@ static void ggml_compute_forward_pad_f32( } -void ggml_compute_forward_pad( - const ggml_compute_params * params, - ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - const bool circular = (bool) ggml_get_op_params_i32(dst, 8); - switch (src0->type) { +void ggml_compute_forward_im2col_3d( + const ggml_compute_params * params, + ggml_tensor * dst) { + switch (dst->type) { + case GGML_TYPE_F16: + { + ggml_compute_forward_im2col_3d_f16(params, dst); + } break; case GGML_TYPE_F32: { - if (circular) { - ggml_compute_forward_pad_f32(params, dst); - } else { - ggml_compute_forward_pad_f32(params, dst); - } + ggml_compute_forward_im2col_3d_f32(params, dst); } break; default: { @@ -7833,330 +8676,452 @@ void ggml_compute_forward_pad( } } -// ggml_compute_forward_pad_reflect_1d +static void ggml_call_mul_mat(ggml_type type, const ggml_compute_params * params, int64_t m, int64_t n, int64_t k, + void * a, void * b, float * c) { + const ggml_type_traits * traits = ggml_get_type_traits(type); + struct ggml_tensor src1 = {}; + src1.type = type; + src1.ne[0] = k; + src1.ne[1] = m; + src1.ne[2] = 1; + src1.ne[3] = 1; + src1.nb[0] = traits->type_size; + src1.nb[1] = k * traits->type_size; + src1.nb[2] = src1.nb[1]; + src1.nb[3] = src1.nb[2]; + src1.data = a; -void ggml_compute_forward_pad_reflect_1d( - const ggml_compute_params * params, - ggml_tensor * dst) { + struct ggml_tensor src0 = {}; + src0.type = type; + src0.ne[0] = k; + src0.ne[1] = n; + src0.ne[2] = 1; + src0.ne[3] = 1; + src0.nb[0] = traits->type_size; + src0.nb[1] = k * traits->type_size; + src0.nb[2] = src0.nb[1]; + src0.nb[3] = src0.nb[2]; + src0.data = b; - const ggml_tensor * src0 = dst->src[0]; + struct ggml_tensor dst = {}; + dst.ne[0] = n; + dst.ne[1] = m; + dst.ne[2] = 1; + dst.ne[3] = 1; + dst.nb[0] = sizeof(float); + dst.nb[1] = n * sizeof(float); + dst.nb[2] = dst.nb[1]; + dst.nb[3] = dst.nb[2]; + dst.data = c; + dst.src[0] = &src0; + dst.src[1] = &src1; - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT( dst->type == GGML_TYPE_F32); + ggml_compute_forward_mul_mat(params, &dst); +} - const int ith = params->ith; - const int nth = params->nth; +static inline int64_t ggml_wrap_around(int64_t coord, int64_t size) { + return (coord + size) % size; // adding size avoids negative number weirdness +} - const int32_t * opts = (const int32_t *) dst->op_params; - const int p0 = opts[0]; - const int p1 = opts[1]; +// ggml_compute_forward_conv_2d - GGML_TENSOR_UNARY_OP_LOCALS - for (int64_t i3 = 0; i3 < ne3; i3++) { - for (int64_t i2 = 0; i2 < ne2; i2++) { - for (int64_t i1 = ith; i1 < ne1; i1 += nth) { - float * left = (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + p0*nb0); - float * right = (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + (ne0-p1-1)*nb0); +static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params, + const ggml_tensor * kernel, // [KW, KH, IC, OC] + const ggml_tensor * src, // [W, H, C, N] + ggml_tensor * dst, // [OW, OH, OC, N] + ggml_type kernel_type) { - ggml_vec_cpy_f32(ne00, left, (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); + GGML_ASSERT(ggml_is_contiguous(kernel)); + GGML_ASSERT(kernel_type == GGML_TYPE_F16 || kernel_type == GGML_TYPE_F32); + GGML_ASSERT(kernel->type == kernel_type); - for (int i0 = 1; i0 <= p0; i0++) { left[-i0] = left[i0]; } - for (int i0 = 1; i0 <= p1; i0++) { right[i0] = right[-i0]; } - } - } - } -} + const ggml_type_traits * traits = ggml_get_type_traits(kernel_type); -// ggml_compute_forward_roll + const int32_t stride_x = dst->op_params[0]; + const int32_t stride_y = dst->op_params[1]; + const int32_t pad_x = dst->op_params[2]; + const int32_t pad_y = dst->op_params[3]; + const int32_t dilation_x = dst->op_params[4]; + const int32_t dilation_y = dst->op_params[5]; -static int64_t ggml_wrap_index(int64_t i, int64_t ne) { - if (i < 0) { - return i + ne; - } else if (i >= ne) { - return i - ne; - } - return i; -} - -static void ggml_compute_forward_roll_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - const float * src_data = (const float *) src0->data; - float * dst_data = (float *) dst->data; - - GGML_TENSOR_UNARY_OP_LOCALS - - const int s0 = ggml_get_op_params_i32(dst, 0); - const int s1 = ggml_get_op_params_i32(dst, 1); - const int s2 = ggml_get_op_params_i32(dst, 2); - const int s3 = ggml_get_op_params_i32(dst, 3); - - const int64_t total = ne1 * ne2 * ne3; - const int64_t per_thread = (total + params->nth) / params->nth; - const int64_t start = params->ith * per_thread; - const int64_t end = std::min(start + per_thread, total); + const int64_t c_in = src->ne[2]; + const int64_t c_out = kernel->ne[3]; + GGML_ASSERT(c_in == kernel->ne[2]); - for (int64_t i = start; i < end; ++i) { - const int64_t i1 = i % ne1; - const int64_t i2 = (i / ne1) % ne2; - const int64_t i3 = i / (ne2 * ne1); - float * dst_row = dst_data + (i3*nb3 + i2*nb2 + i1*nb1) / sizeof(float); + const int64_t src_w = src->ne[0]; + const int64_t src_h = src->ne[1]; + const int64_t knl_w = kernel->ne[0]; + const int64_t knl_h = kernel->ne[1]; + const int64_t dst_w = dst->ne[0]; + const int64_t dst_h = dst->ne[1]; - const int64_t i01 = ggml_wrap_index(i1 - s1, ne01); - const int64_t i02 = ggml_wrap_index(i2 - s2, ne02); - const int64_t i03 = ggml_wrap_index(i3 - s3, ne03); - const float * src_row = src_data + (i03*nb03 + i02*nb02 + i01*nb01) / sizeof(float); + const float * src_data = (float *) src->data; + void * knl_data = kernel->data; + float * dst_data = (float *) dst->data; - const int64_t s = ggml_wrap_index(-s0, ne00); - const int64_t n = ne00 - s; - ggml_vec_cpy_f32(n, dst_row, src_row + s); - ggml_vec_cpy_f32(s, dst_row + n, src_row); - } -} + const int64_t knl_n = knl_w * knl_h * c_in; + const int64_t patch_total = dst->ne[3] * dst_w * dst_h; -void ggml_compute_forward_roll( - const ggml_compute_params * params, - ggml_tensor * dst) { + const int64_t space_per_patch = knl_n * traits->type_size + c_out * sizeof(float); + const int64_t batch_size = params->wsize / space_per_patch; + const int64_t patches_per_batch = batch_size > 8 ? (batch_size / 8) * 8 : batch_size; + const int64_t batch_n = (patch_total + patches_per_batch - 1) / patches_per_batch; - const ggml_tensor * src0 = dst->src[0]; + GGML_ASSERT(patches_per_batch > 0 && batch_size >= 1); - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_roll_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} + void * tmp = params->wdata; -// ggml_compute_forward_arange + for (int64_t batch_i = 0; batch_i < batch_n; ++batch_i) { -static void ggml_compute_forward_arange_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { + const int64_t patch_start_batch = batch_i * patches_per_batch; + const int64_t patch_end_batch = std::min(patch_start_batch + patches_per_batch, + patch_total); + const int64_t patch_n = patch_end_batch - patch_start_batch; - GGML_ASSERT(dst->nb[0] == sizeof(float)); + const int64_t patch_per_thread = (patch_n + params->nth - 1) / params->nth; + const int64_t patch_start = patch_start_batch + params->ith * patch_per_thread; + const int64_t patch_end = std::min(patch_start + patch_per_thread, patch_end_batch); - const int ith = params->ith; - const int nth = params->nth; + //im2col for a patch + for (int64_t p = patch_start; p < patch_end; ++p) { + const int64_t batch_n = p / (dst_w * dst_h); + const int64_t src_x = (p / dst_w) % dst_h; + const int64_t src_y = p % dst_w; - const float start = ggml_get_op_params_f32(dst, 0); - const float stop = ggml_get_op_params_f32(dst, 1); - const float step = ggml_get_op_params_f32(dst, 2); + const float * src_base = (const float *)((const char *)src_data + batch_n * src->nb[3]); + char * dst_row = (char *) tmp + (p % patches_per_batch) * knl_n * traits->type_size; - const int64_t steps = (int64_t) ceilf((stop - start) / step); + for (int64_t ic = 0; ic < c_in; ++ic) { + for (int64_t ky = 0; ky < knl_h; ++ky) { + for (int64_t kx = 0; kx < knl_w; ++kx) { + const int64_t sy = src_x * stride_y + ky * dilation_y - pad_y; + const int64_t sx = src_y * stride_x + kx * dilation_x - pad_x; - GGML_ASSERT(ggml_nelements(dst) == steps); + int64_t dst_idx = ic * (knl_h * knl_w) + ky * knl_w + kx; - for (int64_t i = ith; i < steps; i+= nth) { - float value = start + step * i; - ((float *)dst->data)[i] = value; - } -} + float src_val; + if (sy < 0 || sy >= src_h || sx < 0 || sx >= src_w) { + src_val = 0.0f; + } else { + const float * src_ptr = (const float *)((const char *)src_base + sx * src->nb[0] + sy * src->nb[1] + ic * src->nb[2]); + src_val = *src_ptr; + } -void ggml_compute_forward_arange( - const ggml_compute_params * params, - ggml_tensor * dst) { - switch (dst->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_arange_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); + char * element_ptr = dst_row + dst_idx * traits->type_size; + if (kernel_type == GGML_TYPE_F32) { + *(float *) element_ptr = src_val; + } else if (kernel_type == GGML_TYPE_F16) { + *(ggml_fp16_t *) element_ptr = GGML_CPU_FP32_TO_FP16(src_val); + } + } + } } - } -} + } // patches handled by this thread -static void ggml_compute_forward_timestep_embedding_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_barrier(params->threadpool); - const ggml_tensor * src0 = dst->src[0]; + float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n * traits->type_size); - GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(gemm_output + patch_n * c_out <= (float*)tmp + params->wsize); - const int ith = params->ith; - const int nth = params->nth; + // GEMM: patches[patch_n, knl_n] × kernel[knl_n, c_out] = output[patch_n, c_out] + ggml_call_mul_mat(kernel_type, params, patch_n, c_out, knl_n, tmp, knl_data, gemm_output); - GGML_TENSOR_UNARY_OP_LOCALS + ggml_barrier(params->threadpool); - const int dim = ggml_get_op_params_i32(dst, 0); - const int max_period = ggml_get_op_params_i32(dst, 1); - int half = dim / 2; + //permute back [OC, N, OH, OW] to [N, OC, OH, OW] + const int64_t permute_per_thread = (patch_n + params->nth - 1) / params->nth; + const int64_t permute_start = params->ith * permute_per_thread; + const int64_t permute_end = std::min(permute_start + permute_per_thread, patch_n); - for (int64_t i = 0; i < ne00; i++) { - float * embed_data = (float *)((char *) dst->data + i*nb1); - for (int64_t j = ith; j < half; j += nth) { - float timestep = ((float *)src0->data)[i]; - float freq = (float)expf(-logf(max_period) * j / half); - float arg = timestep * freq; - embed_data[j] = cosf(arg); - embed_data[j + half] = sinf(arg); - } - if (dim % 2 != 0 && ith == 0) { - embed_data[2 * half] = 0.f; + for (int64_t i = permute_start; i < permute_end; ++i) { + const int64_t p = patch_start_batch + i; + const int64_t batch_n = p / (dst_w * dst_h); + const int64_t dst_y = (p / dst_w) % dst_h; + const int64_t dst_x = p % dst_w; + + for (int64_t oc = 0; oc < c_out; ++oc) { + const float value = gemm_output[i * c_out + oc]; + float * dst_ptr = (float *)((char *)dst_data + dst_x * dst->nb[0] + dst_y * dst->nb[1] + oc * dst->nb[2] + batch_n * dst->nb[3]); + *dst_ptr = value; + } } } } -void ggml_compute_forward_timestep_embedding( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_conv_2d( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_timestep_embedding_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } + ggml_compute_forward_conv_2d_impl(params, src0, src1, dst, src0->type); } -// ggml_compute_forward_argsort +// ggml_compute_forward_conv_3d -template -struct cmp_argsort { - const float * data; - bool operator()(int32_t a, int32_t b) const { - if constexpr (order == GGML_SORT_ORDER_ASC) { - return data[a] < data[b]; - } else { - return data[a] > data[b]; - } - } -}; +static void ggml_compute_forward_conv_3d_impl(const ggml_compute_params * params, + const ggml_tensor * kernel, + const ggml_tensor * src, + ggml_tensor * dst, + ggml_type kernel_type) { -static void ggml_compute_forward_argsort_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { + GGML_ASSERT(ggml_is_contiguous(kernel)); + GGML_ASSERT(kernel_type == GGML_TYPE_F16 || kernel_type == GGML_TYPE_F32); + GGML_ASSERT(kernel->type == kernel_type); - const ggml_tensor * src0 = dst->src[0]; + const ggml_type_traits * traits = ggml_get_type_traits(kernel_type); - GGML_TENSOR_UNARY_OP_LOCALS + const int32_t s0 = dst->op_params[0]; + const int32_t s1 = dst->op_params[1]; + const int32_t s2 = dst->op_params[2]; + const int32_t p0 = dst->op_params[3]; + const int32_t p1 = dst->op_params[4]; + const int32_t p2 = dst->op_params[5]; + const int32_t d0 = dst->op_params[6]; + const int32_t d1 = dst->op_params[7]; + const int32_t d2 = dst->op_params[8]; + const int32_t c = dst->op_params[9]; + const int32_t n = dst->op_params[10]; + const int32_t oc = dst->op_params[11]; - GGML_ASSERT(nb0 == sizeof(float)); + const int64_t src_w = src->ne[0]; + const int64_t src_h = src->ne[1]; + const int64_t src_d = src->ne[2]; + const int64_t knl_w = kernel->ne[0]; + const int64_t knl_h = kernel->ne[1]; + const int64_t knl_d = kernel->ne[2]; + const int64_t dst_w = dst->ne[0]; + const int64_t dst_h = dst->ne[1]; + const int64_t dst_d = dst->ne[2]; - const int ith = params->ith; - const int nth = params->nth; + const float * src_data = (float *) src->data; + void * knl_data = kernel->data; + float * dst_data = (float *) dst->data; - const int64_t nr = ggml_nrows(src0); + const int64_t knl_n_per_channel = knl_w * knl_h * knl_d; + const int64_t knl_n_total = knl_n_per_channel * c; + const int64_t patch_total = n * dst_w * dst_h * dst_d; - ggml_sort_order order = (ggml_sort_order) ggml_get_op_params_i32(dst, 0); + const int64_t space_per_patch = knl_n_total * traits->type_size + oc * sizeof(float); + const int64_t batch_size = params->wsize / space_per_patch; + const int64_t patches_per_batch = batch_size > 8 ? (batch_size / 8) * 8 : batch_size; + const int64_t batch_n = (patch_total + patches_per_batch - 1) / patches_per_batch; - for (int64_t i = ith; i < nr; i += nth) { - const float * src_data = (float *)((char *) src0->data + i*nb01); + GGML_ASSERT(patches_per_batch > 0 && batch_size >= 1); - int32_t * dst_data = (int32_t *)((char *) dst->data + i*nb1); + void * tmp = params->wdata; - for (int64_t j = 0; j < ne0; j++) { - dst_data[j] = j; - } + for (int64_t batch_i = 0; batch_i < batch_n; ++batch_i) { + const int64_t patch_start_batch = batch_i * patches_per_batch; + const int64_t patch_end_batch = std::min(patch_start_batch + patches_per_batch, patch_total); + const int64_t patch_n_in_batch = patch_end_batch - patch_start_batch; - switch (order) { - case GGML_SORT_ORDER_ASC: - std::sort(dst_data, dst_data + ne0, cmp_argsort{src_data}); - break; + const int64_t patch_per_thread = (patch_n_in_batch + params->nth - 1) / params->nth; + const int64_t patch_start = patch_start_batch + params->ith * patch_per_thread; + const int64_t patch_end = std::min(patch_start + patch_per_thread, patch_end_batch); - case GGML_SORT_ORDER_DESC: - std::sort(dst_data, dst_data + ne0, cmp_argsort{src_data}); - break; + for (int64_t p = patch_start; p < patch_end; ++p) { + const int64_t p_in_batch = p % (dst_w * dst_h * dst_d); + const int64_t p_in_depth = p_in_batch % (dst_w * dst_h); + const int64_t batch_idx = p / (dst_w * dst_h * dst_d); + const int64_t dst_z = p_in_batch / (dst_w * dst_h); + const int64_t dst_y = p_in_depth / dst_w; + const int64_t dst_x = p_in_depth % dst_w; - default: - GGML_ABORT("invalid sort order"); + char * dst_row = (char *) tmp + (p % patches_per_batch) * knl_n_total * traits->type_size; + + for (int64_t ic = 0; ic < c; ++ic) { + for (int64_t kz = 0; kz < knl_d; ++kz) { + for (int64_t ky = 0; ky < knl_h; ++ky) { + for (int64_t kx = 0; kx < knl_w; ++kx) { + const int64_t sz = dst_z * s2 + kz * d2 - p2; + const int64_t sy = dst_y * s1 + ky * d1 - p1; + const int64_t sx = dst_x * s0 + kx * d0 - p0; + + int64_t dst_idx = ic * knl_n_per_channel + kz * (knl_h * knl_w) + ky * knl_w + kx; + + float src_val; + if (sz < 0 || sz >= src_d || sy < 0 || sy >= src_h || sx < 0 || sx >= src_w) { + src_val = 0.0f; + } else { + const int64_t cn_idx = batch_idx * c + ic; + const float * src_ptr = (const float *)((const char *)src_data + sx*src->nb[0] + sy*src->nb[1] + sz*src->nb[2] + cn_idx*src->nb[3]); + src_val = *src_ptr; + } + + char * element_ptr = dst_row + dst_idx * traits->type_size; + if (kernel_type == GGML_TYPE_F32) { + *(float *)element_ptr = src_val; + } else if (kernel_type == GGML_TYPE_F16) { + *(ggml_fp16_t *)element_ptr = GGML_CPU_FP32_TO_FP16(src_val); + } + } + } + } + } } - } -} -void ggml_compute_forward_argsort( - const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_barrier(params->threadpool); - const ggml_tensor * src0 = dst->src[0]; + float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n_total * traits->type_size); + ggml_call_mul_mat(kernel_type, params, patch_n_in_batch, oc, knl_n_total, tmp, knl_data, gemm_output); - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_argsort_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); + ggml_barrier(params->threadpool); + + const int64_t permute_per_thread = (patch_n_in_batch + params->nth - 1) / params->nth; + const int64_t permute_start = params->ith * permute_per_thread; + const int64_t permute_end = std::min(permute_start + permute_per_thread, patch_n_in_batch); + + for (int64_t i = permute_start; i < permute_end; ++i) { + const int64_t p = patch_start_batch + i; + const int64_t p_in_batch = p % (dst_w * dst_h * dst_d); + const int64_t p_in_depth = p_in_batch % (dst_w * dst_h); + const int64_t batch_idx = p / (dst_w * dst_h * dst_d); + const int64_t dst_z = p_in_batch / (dst_w * dst_h); + const int64_t dst_y = p_in_depth / dst_w; + const int64_t dst_x = p_in_depth % dst_w; + + for (int64_t ioc = 0; ioc < oc; ++ioc) { + const float value = gemm_output[i * oc + ioc]; + const int64_t ocn_idx = batch_idx * oc + ioc; + float * dst_ptr = (float *)((char *)dst_data + dst_x*dst->nb[0] + dst_y*dst->nb[1] + dst_z*dst->nb[2] + ocn_idx*dst->nb[3]); + *dst_ptr = value; } + } } } -// ggml_compute_forward_top_k - -struct cmp_top_k { - const float * data; - bool operator()(int32_t a, int32_t b) const { - return data[a] > data[b]; - } -}; +void ggml_compute_forward_conv_3d( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + ggml_compute_forward_conv_3d_impl(params, src0, src1, dst, src0->type); +} -static void ggml_compute_forward_top_k_f32( +template +static void ggml_compute_forward_conv_transpose_2d_impl( const ggml_compute_params * params, - ggml_tensor * dst) { + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - GGML_TENSOR_UNARY_OP_LOCALS + GGML_ASSERT(src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); - GGML_ASSERT(nb0 == sizeof(float)); + GGML_TENSOR_BINARY_OP_LOCALS const int ith = params->ith; const int nth = params->nth; - const int64_t nr = ggml_nrows(src0); + const int nk = ne00*ne01*ne02*ne03; - const int top_k = ne0; + GGML_ASSERT(nb00 == ggml_type_size(src0->type)); + GGML_ASSERT(nb10 == sizeof(float)); - int32_t * tmp = (int32_t *) params->wdata + (ne00 + CACHE_LINE_SIZE_F32) * ith; + if (ith == 0) { + memset(params->wdata, 0, params->wsize); - for (int64_t i = ith; i < nr; i += nth) { - const float * src_data = (float *)((char *) src0->data + i*nb01); + // permute kernel data (src0) from (Kw x Kh x Cout x Cin) to (Cin x Kw x Kh x Cout) + { + kernel_t * const wdata = (kernel_t *) params->wdata + 0; - for (int64_t j = 0; j < ne00; j++) { - tmp[j] = j; + for (int64_t i03 = 0; i03 < ne03; i03++) { + for (int64_t i02 = 0; i02 < ne02; i02++) { + const kernel_t * const src = (kernel_t *)((char *) src0->data + i03*nb03 + i02*nb02); + kernel_t * dst_data = wdata + i02*ne01*ne00*ne03; + for (int64_t i01 = 0; i01 < ne01; i01++) { + for (int64_t i00 = 0; i00 < ne00; i00++) { + dst_data[i01*ne00*ne03 + i00*ne03 + i03] = src[i01 * ne00 + i00]; + } + } + } + } } - std::partial_sort(tmp, tmp + top_k, tmp + ne00, cmp_top_k{src_data}); + // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh) + { + kernel_t * const wdata = (kernel_t *) params->wdata + nk; + for (int i12 = 0; i12 < ne12; i12++) { + for (int i11 = 0; i11 < ne11; i11++) { + const float * const src = (float *)((char *) src1->data + i12*nb12 + i11*nb11); + kernel_t * dst_data = wdata + i11*ne10*ne12; + for (int i10 = 0; i10 < ne10; i10++) { + if constexpr (std::is_same_v) { + dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]); + } else { + dst_data[i10*ne12 + i12] = src[i10]; + } + } + } + } + } - int32_t * dst_data = (int32_t *)((char *) dst->data + i*nb1); + memset(dst->data, 0, ggml_nbytes(dst)); + } + ggml_barrier(params->threadpool); - std::copy(tmp, tmp + top_k, dst_data); + const int32_t stride = ggml_get_op_params_i32(dst, 0); - // emphasize that the order is not important - if (top_k > 1) { - std::swap(dst_data[0], dst_data[1]); + // total patches in dst + const int np = ne2; + + // patches per thread + const int dp = (np + nth - 1)/nth; + + // patch range for this thread + const int ip0 = dp*ith; + const int ip1 = MIN(ip0 + dp, np); + + kernel_t * const wdata = (kernel_t *) params->wdata + 0; + kernel_t * const wdata_src = wdata + nk; + + for (int i2 = ip0; i2 < ip1; i2++) { // Cout + float * dst_data = (float *)((char *) dst->data + i2*nb2); + kernel_t * wdata_kernel = wdata + i2*ne01*ne00*ne03; + for (int i11 = 0; i11 < ne11; i11++) { + for (int i10 = 0; i10 < ne10; i10++) { + const int i1n = i11*ne10*ne12 + i10*ne12; + for (int i01 = 0; i01 < ne01; i01++) { + for (int i00 = 0; i00 < ne00; i00++) { + float v = 0; + if constexpr (std::is_same_v) { + ggml_vec_dot_f16(ne03, &v, 0, + wdata_src + i1n, 0, + wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1); + } else { + ggml_vec_dot_f32(ne03, &v, 0, + wdata_src + i1n, 0, + wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1); + } + dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v; + } + } + } } } } -void ggml_compute_forward_top_k( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_conv_transpose_2d( + const ggml_compute_params * params, + ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { + case GGML_TYPE_F16: + { + ggml_compute_forward_conv_transpose_2d_impl(params, dst); + } break; case GGML_TYPE_F32: { - ggml_compute_forward_top_k_f32(params, dst); + ggml_compute_forward_conv_transpose_2d_impl(params, dst); } break; default: { @@ -8165,744 +9130,876 @@ void ggml_compute_forward_top_k( } } -static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( +// ggml_compute_forward_conv_2d_dw + +struct ggml_conv_2d_dw_params { + int64_t channels; + int64_t batch; + int64_t src_w; + int64_t src_h; + int64_t dst_w; + int64_t dst_h; + int64_t knl_w; + int64_t knl_h; + int stride_x; + int stride_y; + int pad_x; + int pad_y; + int dilation_x; + int dilation_y; +}; + +static void ggml_compute_forward_conv_2d_dw_cwhn( const ggml_compute_params * params, + const ggml_tensor * src, + const ggml_tensor * kernel, ggml_tensor * dst, - int ir0, int ir1, - int64_t ic_start, int64_t ic_end, - float * partials, int64_t partial_stride) { + const ggml_conv_2d_dw_params & p) { - const bool write_partials = (partials != nullptr); - const ggml_tensor * q = dst->src[0]; - const ggml_tensor * k = dst->src[1]; - const ggml_tensor * v = dst->src[2]; - const ggml_tensor * mask = dst->src[3]; - const ggml_tensor * sinks = dst->src[4]; + const int64_t c = p.channels; + const float * knl_data = (const float *)kernel->data; - GGML_TENSOR_LOCALS(int64_t, neq, q, ne) - GGML_TENSOR_LOCALS(size_t, nbq, q, nb) - GGML_TENSOR_LOCALS(int64_t, nek, k, ne) - GGML_TENSOR_LOCALS(size_t, nbk, k, nb) - GGML_TENSOR_LOCALS(int64_t, nev, v, ne) - GGML_TENSOR_LOCALS(size_t, nbv, v, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + const int64_t rows_total = p.dst_h * p.batch; + const int64_t rows_per_thread = (rows_total + params->nth - 1) / params->nth; + const int64_t row_start = params->ith * rows_per_thread; + const int64_t row_end = MIN(row_start + rows_per_thread, rows_total); - const int64_t DK = nek0; - const int64_t DV = nev0; - const int64_t N = neq1; +#ifdef GGML_SIMD + #if defined(__ARM_FEATURE_SVE) + const int64_t pkg_size = svcntw(); + #else + const int64_t pkg_size = GGML_F32_EPR; + #endif + const int64_t pkg_count = c / pkg_size; + const int64_t c_pkg_end = pkg_count * pkg_size; +#else + const int64_t c_pkg_end = 0; +#endif - GGML_ASSERT(ne0 == DV); - GGML_ASSERT(ne2 == N); + for (int64_t row = row_start; row < row_end; ++row) { + const int64_t dst_y = row % p.dst_h; + const float * src_data = (const float *)src->data + (row / p.dst_h) * p.src_w * p.src_h * c; + for (int64_t dst_x = 0; dst_x < p.dst_w; ++dst_x) { + float * dst_data = (float *)dst->data + (row * p.dst_w + dst_x) * c; + const int64_t src_y_base = dst_y * p.stride_y - p.pad_y; + const int64_t src_x_base = dst_x * p.stride_x - p.pad_x; - // input tensor rows must be contiguous - GGML_ASSERT(nbq0 == ggml_type_size(q->type)); - GGML_ASSERT(nbk0 == ggml_type_size(k->type)); - GGML_ASSERT(nbv0 == ggml_type_size(v->type)); +#ifdef GGML_SIMD + // Vectorized loop + for (int64_t c_i = 0; c_i < c_pkg_end; c_i += pkg_size) { + GGML_F32_VEC sum = GGML_F32_VEC_ZERO; + for (int64_t knl_y = 0; knl_y < p.knl_h; ++knl_y) { + const int64_t src_y = src_y_base + knl_y * p.dilation_y; + if (src_y < 0 || src_y >= p.src_h) { + continue; + } + for (int64_t knl_x = 0; knl_x < p.knl_w; ++knl_x) { + const int64_t src_x = src_x_base + knl_x * p.dilation_x; + if (src_x < 0 || src_x >= p.src_w) { + continue; + } + GGML_F32_VEC k = GGML_F32_VEC_LOAD(knl_data + (knl_y * p.knl_w + knl_x) * c + c_i); + GGML_F32_VEC s = GGML_F32_VEC_LOAD(src_data + (src_y * p.src_w + src_x) * c + c_i); + sum = GGML_F32_VEC_FMA(sum, k, s); + } + } + GGML_F32_VEC_STORE(dst_data + c_i, sum); + } +#endif + // Scalar loop + for (int64_t c_i = c_pkg_end; c_i < c; ++c_i) { + float sum = 0.0f; + for (int64_t knl_y = 0; knl_y < p.knl_h; ++knl_y) { + const int64_t src_y = src_y_base + knl_y * p.dilation_y; + if (src_y < 0 || src_y >= p.src_h) { + continue; + } + for (int64_t knl_x = 0; knl_x < p.knl_w; ++knl_x) { + const int64_t src_x = src_x_base + knl_x * p.dilation_x; + if (src_x < 0 || src_x >= p.src_w) { + continue; + } + sum += knl_data[(knl_y * p.knl_w + knl_x) * c + c_i] + * src_data[(src_y * p.src_w + src_x) * c + c_i]; + } + } + dst_data[c_i] = sum; + } + } + } +} - GGML_ASSERT(neq0 == DK); - GGML_ASSERT(nek0 == DK); - GGML_ASSERT(nev0 == DV); +static void ggml_compute_forward_conv_2d_dw_whcn( + const ggml_compute_params * params, + const ggml_tensor * src, + const ggml_tensor * kernel, + ggml_tensor * dst, + const ggml_conv_2d_dw_params & p) { - GGML_ASSERT(neq1 == N); + const int64_t n = p.channels * p.batch; + const int64_t per_thread = (n + params->nth - 1) / params->nth; + const int64_t start = params->ith * per_thread; + const int64_t end = MIN(start + per_thread, n); - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); + for (int64_t i = start; i < end; ++i) { + const float * knl_data = (const float *)kernel->data + (i % p.channels) * p.knl_w * p.knl_h; + const float * src_data = (const float *)src->data + i * p.src_w * p.src_h; + float * dst_data = (float *)dst->data + i * p.dst_w * p.dst_h; - // broadcast factors - const int64_t rk2 = neq2/nek2; - const int64_t rk3 = neq3/nek3; + for (int64_t dst_y = 0; dst_y < p.dst_h; ++dst_y) { + for (int64_t dst_x = 0; dst_x < p.dst_w; ++dst_x) { - const int64_t rv2 = neq2/nev2; - const int64_t rv3 = neq3/nev3; + float sum = 0.0f; + for (int64_t knl_y = 0; knl_y < p.knl_h; ++knl_y) { + const int64_t src_y = dst_y * p.stride_y + knl_y * p.dilation_y - p.pad_y; + if (src_y < 0 || src_y >= p.src_h) { + continue; + } + for (int64_t knl_x = 0; knl_x < p.knl_w; ++knl_x) { + const int64_t src_x = dst_x * p.stride_x + knl_x * p.dilation_x - p.pad_x; + if (src_x < 0 || src_x >= p.src_w) { + continue; + } + sum += knl_data[knl_y * p.knl_w + knl_x] + * src_data[src_y * p.src_w + src_x]; + } + } + dst_data[dst_y * p.dst_w + dst_x] = sum; + } + } + } +} - // parallelize by q rows using ggml_vec_dot_f32 +void ggml_compute_forward_conv_2d_dw( + const ggml_compute_params * params, + ggml_tensor * dst) { - float scale = 1.0f; - float max_bias = 0.0f; - float logit_softcap = 0.0f; + const ggml_tensor * kernel = dst->src[0]; + const ggml_tensor * src = dst->src[1]; + ggml_conv_2d_dw_params p; + p.channels = src->ne[2]; + p.batch = src->ne[3]; + p.src_w = src->ne[0]; + p.src_h = src->ne[1]; + p.dst_w = dst->ne[0]; + p.dst_h = dst->ne[1]; + p.knl_w = kernel->ne[0]; + p.knl_h = kernel->ne[1]; + p.stride_x = dst->op_params[0]; + p.stride_y = dst->op_params[1]; + p.pad_x = dst->op_params[2]; + p.pad_y = dst->op_params[3]; + p.dilation_x = dst->op_params[4]; + p.dilation_y = dst->op_params[5]; - memcpy(&scale, (float *) dst->op_params + 0, sizeof(float)); - memcpy(&max_bias, (float *) dst->op_params + 1, sizeof(float)); - memcpy(&logit_softcap, (float *) dst->op_params + 2, sizeof(float)); + GGML_ASSERT(kernel->ne[3] == p.channels); + GGML_ASSERT(dst->ne[3] == p.batch); - if (logit_softcap != 0) { - scale /= logit_softcap; + if (ggml_is_contiguous(src)) { + ggml_compute_forward_conv_2d_dw_whcn(params, src, kernel, dst, p); + } else if (ggml_is_contiguous_channels(src)) { + // kernel should also have channels most contiguous in memory + GGML_ASSERT(kernel->nb[0] >= kernel->nb[2] && kernel->nb[1] >= kernel->nb[0]); + ggml_compute_forward_conv_2d_dw_cwhn(params, src, kernel, dst, p); + } else { + GGML_ABORT("non-contiguous memory layout not supported"); } +} - const uint32_t n_head = neq2; - const uint32_t n_head_log2 = 1u << (uint32_t) floor(log2(n_head)); +// ggml_compute_forward_pool_1d_ksp +static void ggml_compute_forward_pool_1d_ksp( + const ggml_compute_params * params, + const ggml_op_pool op, + const int k, + const int s, + const int p, + ggml_tensor * dst) { - const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); - const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + const ggml_tensor * src = dst->src[0]; - ggml_type const k_vec_dot_type = ggml_get_type_traits_cpu(k->type)->vec_dot_type; - ggml_from_float_t const q_to_vec_dot = ggml_get_type_traits_cpu(k_vec_dot_type)->from_float; - ggml_vec_dot_t const kq_vec_dot = ggml_get_type_traits_cpu(k->type)->vec_dot; - ggml_to_float_t const v_to_float = ggml_get_type_traits(v->type)->to_float; + assert(src->type == GGML_TYPE_F32 || src->type == GGML_TYPE_F16); - GGML_ASSERT(( q_to_vec_dot) && "fattn: unsupported K-type"); - GGML_ASSERT((v->type == GGML_TYPE_F32 || v_to_float ) && "fattn: unsupported V-type"); + if (params->ith != 0) { + return; + } - int ith = params->ith; + const int64_t IW = src->ne[0]; + const int64_t OW = dst->ne[0]; - for (int ir = ir0; ir < ir1; ++ir) { - // q indices - const int iq3 = ir/(neq2*neq1); - const int iq2 = (ir - iq3*neq2*neq1)/neq1; - const int iq1 = (ir - iq3*neq2*neq1 - iq2*neq1); + const int64_t nr = ggml_nrows(src); - const uint32_t h = iq2; // head index - const float slope = (max_bias > 0.0f) ? h < n_head_log2 ? powf(m0, h + 1) : powf(m1, 2*(h - n_head_log2) + 1) : 1.0f; + for (int64_t ir = 0; ir < nr; ++ir) { + const char * srow_bytes = (const char *) src->data + ir * src->nb[1]; + float * drow = (float *) (( char *) dst->data + ir * dst->nb[1]); - float S = 0.0f; // sum - float M = -INFINITY; // maximum KQ value + for (int64_t ow = 0; ow < OW; ++ow) { + float res = 0; + switch (op) { + case GGML_OP_POOL_AVG: res = 0.0f; break; + case GGML_OP_POOL_MAX: res = -FLT_MAX; break; + case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error"); + } - float * VKQ32 = (float *) params->wdata + ith*(1*DK + 2*DV + CACHE_LINE_SIZE_F32); // FP32 VKQ accumulator - float * V32 = (VKQ32 + 1*DV); // (temporary) FP32 V buffer - ggml_fp16_t * VKQ16 = (ggml_fp16_t *) (VKQ32 + 1*DV); // (temporary) FP16 VKQ accumulator - ggml_fp16_t * Q_q = (ggml_fp16_t *) (VKQ32 + 2*DV); // (temporary) buffer for Q converted to quantized/FP16 + int count = 0; + const int base = (int) ow * s - p; - if (v->type == GGML_TYPE_F16) { - memset(VKQ16, 0, DV*sizeof(ggml_fp16_t)); - } else { - memset(VKQ32, 0, DV*sizeof(float)); - } + for (int ki = 0; ki < k; ++ki) { + const int j = base + ki; + if (j < 0 || j >= (int) IW) { + continue; + } - const ggml_fp16_t * mp = mask ? (ggml_fp16_t *)((char *) mask->data + iq1*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]) : NULL; + float v; + if (src->type == GGML_TYPE_F32) { + v = ((const float *) srow_bytes)[j]; + } else { + v = GGML_CPU_FP16_TO_FP32(((const ggml_fp16_t *) srow_bytes)[j]); + } - // k indices - const int ik3 = iq3 / rk3; - const int ik2 = iq2 / rk2; - - // v indices - const int iv3 = iq3 / rv3; - const int iv2 = iq2 / rv2; - - const float * pq = (const float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3)); - q_to_vec_dot(pq, Q_q, DK); + switch (op) { + case GGML_OP_POOL_AVG: res += v; break; + case GGML_OP_POOL_MAX: res = std::max(v, res); break; + case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error"); + } - // online softmax / attention - // loop over n_kv and n_head_kv - // ref: https://arxiv.org/pdf/2112.05682.pdf + ++count; + } - for (int64_t ic = ic_start; ic < ic_end; ++ic) { - const float mv = mp ? slope*GGML_CPU_FP16_TO_FP32(mp[ic]) : 0.0f; - if (mv == -INFINITY) { - continue; + switch (op) { + case GGML_OP_POOL_AVG: res = (count > 0) ? (res / count) : 0.0f; break; + case GGML_OP_POOL_MAX: break; + case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error"); } - float s; // KQ value + drow[ow] = res; + } + } +} - const char * k_data = (const char *) k->data + ( ic*nbk1 + ik2*nbk2 + ik3*nbk3); - kq_vec_dot(DK, &s, 0, k_data, 0, Q_q, 0, 1); +// ggml_compute_forward_pool_1d - s = s*scale; // scale KQ value +void ggml_compute_forward_pool_1d( + const ggml_compute_params * params, + ggml_tensor * dst) { - if (logit_softcap != 0.0f) { - s = logit_softcap*tanhf(s); - } + const int32_t * opts = (const int32_t *)dst->op_params; + ggml_op_pool op = static_cast(opts[0]); + const int k0 = opts[1]; + const int s0 = opts[2]; + const int p0 = opts[3]; - s += mv; // apply mask + ggml_compute_forward_pool_1d_ksp(params, op, k0, s0, p0, dst); +} - const float Mold = M; +// ggml_compute_forward_pool_2d - float ms = 1.0f; // upon new higher max val, scale VKQ and KQ sum with this value - float vs = 1.0f; // post-softmax KQ value, expf(s - M) +void ggml_compute_forward_pool_2d( + const ggml_compute_params * params, + ggml_tensor * dst) { - const char * v_data = ((const char *) v->data + (ic*nbv1 + iv2*nbv2 + iv3*nbv3)); + const ggml_tensor * src = dst->src[0]; - if (v->type == GGML_TYPE_F16) { - if (s > M) { - // s is new maximum, ms < 1.0f, vs == expf(s - s) == 1.0f - M = s; - ms = expf(Mold - M); + assert(src->type == GGML_TYPE_F32 || src->type == GGML_TYPE_F16); - // V = V*expf(Mold - M) - ggml_vec_scale_f16(DV, VKQ16, ms); - } else { - // no new maximum, ms == 1.0f, vs != 1.0f - vs = expf(s - M); - } + if (params->ith != 0) { + return; + } - // V += v*expf(s - M) - ggml_vec_mad_f16(DV, VKQ16, (const ggml_fp16_t *) v_data, vs); - } else { - if (s > M) { - // s is new maximum, ms < 1.0f, vs == expf(s - s) == 1.0f - M = s; - ms = expf(Mold - M); + const int32_t * opts = (const int32_t *)dst->op_params; - // V = V*expf(Mold - M) - ggml_vec_scale_f32(DV, VKQ32, ms); - } else { - // no new maximum, ms == 1.0f, vs != 1.0f - vs = expf(s - M); - } + ggml_op_pool op = static_cast(opts[0]); + const int k0 = opts[1]; + const int k1 = opts[2]; + const int s0 = opts[3]; + const int s1 = opts[4]; + const int p0 = opts[5]; + const int p1 = opts[6]; + const char * cdata = (const char*)src->data; + const char * const data_end = cdata + ggml_nbytes(src); - // V += v*expf(s - M) - if (v_to_float) { - v_to_float(v_data, V32, DV); - ggml_vec_mad_f32(DV, VKQ32, V32, vs); - } else { - // V is F32 - ggml_vec_mad_f32(DV, VKQ32, (const float *) v_data, vs); - } - } + const int64_t px = dst->ne[0]; + const int64_t py = dst->ne[1]; + const int64_t pa = px * py; - S = S*ms + vs; // scale and increment sum with partial sum - } + float * dplane = (float *)dst->data; - if (v->type == GGML_TYPE_F16) { - for (int64_t d = 0; d < DV; ++d) { - VKQ32[d] = GGML_CPU_FP16_TO_FP32(VKQ16[d]); - } - } + const int ka = k0 * k1; + const int offset0 = -p0; + const int offset1 = -p1; - // sinks - apply only on the first kv-chunk - if (sinks && ic_start == 0) { - const float s = ((float *)((char *) sinks->data))[h]; + while (cdata < data_end) { + for (int oy = 0; oy < py; ++oy) { + float * const drow = dplane + oy * px; + float * const out = drow; - float ms = 1.0f; - float vs = 1.0f; + for (int ox = 0; ox < px; ++ox) { + float res = 0; + switch (op) { + case GGML_OP_POOL_AVG: res = 0; break; + case GGML_OP_POOL_MAX: res = -FLT_MAX; break; + case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error"); + } - if (s > M) { - ms = expf(M - s); - M = s; - ggml_vec_scale_f32(DV, VKQ32, ms); - } else { - vs = expf(s - M); - } + const int ix = offset0 + ox * s0; + const int iy = offset1 + oy * s1; - S = S*ms + vs; - } + for (int ky = 0; ky < k1; ++ky) { + if (iy + ky < 0 || iy + ky >= src->ne[1]) { + continue; + } - if (write_partials) { - // Write M, S, VKQ to partials for later reduction - // partials layout: [M, S, VKQ[DV]] per query head - float * partial = partials + ir * partial_stride; - partial[0] = M; - partial[1] = S; - memcpy(partial + 2, VKQ32, DV * sizeof(float)); - } else { - // V /= S - const float S_inv = S == 0.0f ? 0.0f : 1.0f/S; - ggml_vec_scale_f32(DV, VKQ32, S_inv); + const void * srow = (const void *)(cdata + src->nb[1] * (iy + ky)); + for (int kx = 0; kx < k0; ++kx) { + int j = ix + kx; + if (j < 0 || j >= src->ne[0]) { + continue; + } - // dst indices - const int i1 = iq1; - const int i2 = iq2; - const int i3 = iq3; + const float srow_j = (src->type == GGML_TYPE_F32) ? ((const float*)srow)[j] : GGML_CPU_FP16_TO_FP32(((const ggml_fp16_t*)srow)[j]); + switch (op) { + case GGML_OP_POOL_AVG: res += srow_j; break; + case GGML_OP_POOL_MAX: res = std::max(srow_j, res); break; + case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error"); + } + } + } + switch (op) { + case GGML_OP_POOL_AVG: res /= ka; break; + case GGML_OP_POOL_MAX: break; + case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error"); + } - // permute(0, 2, 1, 3) - memcpy((char *) dst->data + (i3*ne2*ne1 + i2 + i1*ne1)*nb1, VKQ32, nb1); + out[ox] = res; + } } + + cdata += src->nb[2]; + dplane += pa; } } -static void ggml_compute_forward_flash_attn_ext_tiled( +// ggml_compute_forward_pool_2d_back + +void ggml_compute_forward_pool_2d_back( const ggml_compute_params * params, - ggml_tensor * dst, - int ir0, int ir1) { - const ggml_tensor * q = dst->src[0]; - const ggml_tensor * k = dst->src[1]; - const ggml_tensor * v = dst->src[2]; - const ggml_tensor * mask = dst->src[3]; - const ggml_tensor * sinks = dst->src[4]; + ggml_tensor * dst) { - GGML_TENSOR_LOCALS(int64_t, neq, q, ne) - GGML_TENSOR_LOCALS(size_t, nbq, q, nb) - GGML_TENSOR_LOCALS(int64_t, nek, k, ne) - GGML_TENSOR_LOCALS(size_t, nbk, k, nb) - GGML_TENSOR_LOCALS(int64_t, nev, v, ne) - GGML_TENSOR_LOCALS(size_t, nbv, v, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + const ggml_tensor * src = dst->src[0]; + const ggml_tensor * dstf = dst->src[1]; // forward tensor of dst - const int64_t DK = nek0; - const int64_t DV = nev0; - const int64_t N = neq1; + assert(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16); - GGML_ASSERT(ne0 == DV); - GGML_ASSERT(ne2 == N); + if (params->ith != 0) { + return; + } - // input tensor rows must be contiguous - GGML_ASSERT(nbq0 == ggml_type_size(q->type)); - GGML_ASSERT(nbk0 == ggml_type_size(k->type)); - GGML_ASSERT(nbv0 == ggml_type_size(v->type)); + const int32_t * opts = (const int32_t *)dst->op_params; + ggml_op_pool op = static_cast(opts[0]); + const int k0 = opts[1]; + const int k1 = opts[2]; + const int s0 = opts[3]; + const int s1 = opts[4]; + const int p0 = opts[5]; + const int p1 = opts[6]; - GGML_ASSERT(neq0 == DK); - GGML_ASSERT(nek0 == DK); - GGML_ASSERT(nev0 == DV); + char * cdata = (char *) dst->data; + const char * cdataf = (const char *) dstf->data; + const char * const data_end = cdata + ggml_nbytes(dst); - GGML_ASSERT(neq1 == N); + GGML_ASSERT(params->ith == 0); + memset(cdata, 0, ggml_nbytes(dst)); - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); + const int64_t px = src->ne[0]; + const int64_t py = src->ne[1]; + const int64_t pa = px * py; - GGML_ASSERT(k->type == v->type); - const ggml_type kv_type = k->type; + const float * splane = (const float *) src->data; + const int ka = k0 * k1; + const int offset0 = -p0; + const int offset1 = -p1; - // broadcast factors - const int64_t rk2 = neq2/nek2; - const int64_t rk3 = neq3/nek3; + while (cdata < data_end) { + for (int oy = 0; oy < py; ++oy) { + const float * const srow = splane + oy * px; + for (int ox = 0; ox < px; ++ox) { + const float grad0 = srow[ox]; - const int64_t rv2 = neq2/nev2; - const int64_t rv3 = neq3/nev3; + const int ix = offset0 + ox * s0; + const int iy = offset1 + oy * s1; - float scale = 1.0f; - float max_bias = 0.0f; - float logit_softcap = 0.0f; + if (op == GGML_OP_POOL_MAX) { + float maxval = -FLT_MAX; + int kxmax = -1; + int kymax = -1; - memcpy(&scale, (float *) dst->op_params + 0, sizeof(float)); - memcpy(&max_bias, (float *) dst->op_params + 1, sizeof(float)); - memcpy(&logit_softcap, (float *) dst->op_params + 2, sizeof(float)); + for (int ky = 0; ky < k1; ++ky) { + if (iy + ky < 0 || iy + ky >= dst->ne[1]) { + continue; + } + const void * drowf = (const void *)(cdataf + dst->nb[1] * (iy + ky)); + for (int kx = 0; kx < k0; ++kx) { + int j = ix + kx; + if (j < 0 || j >= dst->ne[0]) { + continue; + } - if (logit_softcap != 0) { - scale /= logit_softcap; - } + const float val = dst->type == GGML_TYPE_F32 ? + ((const float *) drowf)[j] : GGML_CPU_FP16_TO_FP32(((const ggml_fp16_t *) drowf)[j]); + if (val <= maxval) { + continue; + } - const uint32_t n_head = neq2; - const uint32_t n_head_log2 = 1u << (uint32_t) floor(log2(n_head)); + maxval = val; + kxmax = kx; + kymax = ky; + } + } - const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); - const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + if (kxmax == -1 || kymax == -1) { + continue; + } - int ith = params->ith; + void * drow = (void *)(cdata + dst->nb[1] * (iy + kymax)); + const int j = ix + kxmax; + if (dst->type == GGML_TYPE_F32) { + ((float *) drow)[j] += grad0; + } else { + ((ggml_fp16_t *) drow)[j] = GGML_CPU_FP32_TO_FP16(grad0 + GGML_CPU_FP16_TO_FP32(((const ggml_fp16_t *) drow)[j])); + } + } else if (op == GGML_OP_POOL_AVG) { + const float grad = grad0 / ka; - static constexpr int Q_TILE_SZ = ggml_fa_tile_config::Q; - static constexpr int KV_TILE_SZ = ggml_fa_tile_config::KV; + for (int ky = 0; ky < k1; ++ky) { + if (iy + ky < 0 || iy + ky >= dst->ne[1]) { + continue; + } + void * drow = (void *)(cdata + dst->nb[1] * (iy + ky)); + for (int kx = 0; kx < k0; ++kx) { + int j = ix + kx; + if (j < 0 || j >= dst->ne[0]) { + continue; + } - int ir = ir0; - while (ir < ir1) { - // q indices for the start of this tile - const int iq3 = ir/(neq2*neq1); - const int iq2 = (ir - iq3*neq2*neq1)/neq1; - const int iq1 = (ir - iq3*neq2*neq1 - iq2*neq1); + if (dst->type == GGML_TYPE_F32) { + ((float *) drow)[j] += grad; + } else { + ((ggml_fp16_t *) drow)[j] += GGML_CPU_FP32_TO_FP16(grad); + } + } + } + } else { + GGML_ASSERT(false); + } + } + } - // Number of valid rows in this tile: - // - limited by tile size (Q_TILE_SZ) - // - limited by chunk boundary (ir1 - ir) - // - limited by head boundary (neq1 - iq1) to avoid crossing into next head - const int tile_rows = MIN(Q_TILE_SZ, MIN((int)(ir1 - ir), (int)(neq1 - iq1))); - GGML_ASSERT(tile_rows > 0); + cdata += dst->nb[2]; + cdataf += dst->nb[2]; + splane += pa; + } +} - const uint32_t h = iq2; // head index - const float slope = (max_bias > 0.0f) ? h < n_head_log2 ? powf(m0, h + 1) : powf(m1, 2*(h - n_head_log2) + 1) : 1.0f; +// ggml_compute_forward_upscale - float S[Q_TILE_SZ]; - float M[Q_TILE_SZ]; +static void ggml_compute_forward_upscale_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - for (int i = 0 ; i < Q_TILE_SZ; ++i) { - S[i] = 0.; - M[i] = -INFINITY; - } + const ggml_tensor * src0 = dst->src[0]; - // Per-thread scratch layout: - // Q_q: Q_TILE_SZ * DK (converted Q tile — F32 for GEMM, KV type for scalar) - // KQ: Q_TILE_SZ * KV_TILE_SZ (attention scores in float) - // mask: Q_TILE_SZ * KV_TILE_SZ (mask in float) - // VKQ32: Q_TILE_SZ * DV (FP32 output accumulator) - // V32: KV_TILE_SZ * DV (F32 buffer for V tile) - // K_f32: KV_TILE_SZ * DK (F32 buffer for K tile — GEMM path) - float * base = (float *) params->wdata + ith*(Q_TILE_SZ*DK + 2*Q_TILE_SZ*KV_TILE_SZ + Q_TILE_SZ*DV + KV_TILE_SZ*DV + KV_TILE_SZ*DK + CACHE_LINE_SIZE_F32); + GGML_ASSERT(src0->type == GGML_TYPE_F32); - void * Q_q = base; - float * KQ = (float *)((char *)base + Q_TILE_SZ * DK * sizeof(float)); - float * mask32 = KQ + Q_TILE_SZ * KV_TILE_SZ; - float * VKQ32 = mask32 + Q_TILE_SZ * KV_TILE_SZ; - float * V32 = VKQ32 + Q_TILE_SZ * DV; - float * K_f32 = V32 + KV_TILE_SZ * DV; + const int ith = params->ith; + const int nth = params->nth; - memset(VKQ32, 0, Q_TILE_SZ * DV * sizeof(float)); - memset(mask32, 0, Q_TILE_SZ * KV_TILE_SZ * sizeof(float)); + GGML_TENSOR_UNARY_OP_LOCALS - // k indices - const int ik3 = iq3 / rk3; - const int ik2 = iq2 / rk2; + float sf0 = (float)ne0/src0->ne[0]; + float sf1 = (float)ne1/src0->ne[1]; + float sf2 = (float)ne2/src0->ne[2]; + float sf3 = (float)ne3/src0->ne[3]; + float pixel_offset = 0.5f; - // v indices - const int iv3 = iq3 / rv3; - const int iv2 = iq2 / rv2; + const int32_t mode_flags = ggml_get_op_params_i32(dst, 0); + const ggml_scale_mode mode = (ggml_scale_mode) (mode_flags & 0xFF); - { - float * Q_f32 = (float *)Q_q; - for (int tq = 0; tq < tile_rows; tq++) { - const float * pq = (const float *) ((char *) q->data + ((iq1 + tq)*nbq1 + iq2*nbq2 + iq3*nbq3)); - memcpy(Q_f32 + tq * DK, pq, DK * sizeof(float)); - } - for (int tq = tile_rows; tq < Q_TILE_SZ; tq++) { - memset(Q_f32 + tq * DK, 0, DK * sizeof(float)); - } - } + if (mode_flags & GGML_SCALE_FLAG_ALIGN_CORNERS) { + pixel_offset = 0.0f; + sf0 = ne0 > 1 && ne00 > 1 ? (float)(ne0 - 1) / (ne00 - 1) : sf0; + sf1 = ne1 > 1 && ne01 > 1 ? (float)(ne1 - 1) / (ne01 - 1) : sf1; + } - memset(K_f32, 0, DK * KV_TILE_SZ * sizeof(float)); - memset(V32, 0, KV_TILE_SZ * DV * sizeof(float)); + if (mode == GGML_SCALE_MODE_NEAREST) { + for (int64_t i3 = 0; i3 < ne3; i3++) { + const int64_t i03 = i3 / sf3; + for (int64_t i2 = ith; i2 < ne2; i2 += nth) { + const int64_t i02 = i2 / sf2; + for (int64_t i1 = 0; i1 < ne1; i1++) { + const int64_t i01 = i1 / sf1; + for (int64_t i0 = 0; i0 < ne0; i0++) { + const int64_t i00 = i0 / sf0; - for (int64_t ic = 0; ic < nek1; ic += KV_TILE_SZ) { - const int kv_tile = (int)std::min((int64_t)KV_TILE_SZ, nek1 - ic); + const float * x = (float *)((char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); + float * y = (float *)((char *) dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); - // skip the tile entirely if all the masks are -inf - if (mask) { - bool can_skip = true; - for (int tq = 0; tq < tile_rows; tq++) { - const ggml_fp16_t * mp_row = (const ggml_fp16_t *)((const char *) mask->data + (iq1 + tq)*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]); - for (int tk = 0; tk < kv_tile; tk++) { - mask32[tq * KV_TILE_SZ + tk] = slope * GGML_CPU_FP16_TO_FP32(mp_row[ic + tk]); - if (mask32[tq * KV_TILE_SZ + tk] != -INFINITY) { - can_skip = false; - } - } - // Pad remaining mask entries with -inf - for (int tk = kv_tile; tk < KV_TILE_SZ; tk++) { - mask32[tq * KV_TILE_SZ + tk] = -INFINITY; + *y = *x; } } - - if (can_skip) { - continue; - } } + } + } else if (mode == GGML_SCALE_MODE_BILINEAR && (mode_flags & GGML_SCALE_FLAG_ANTIALIAS)) { + // Similar to F.interpolate(..., mode="bilinear", align_corners=False, antialias=True) + // https://github.com/pytorch/pytorch/blob/8871ff29b743948d1225389d5b7068f37b22750b/aten/src/ATen/native/cpu/UpSampleKernel.cpp + auto triangle_filter = [](float x) -> float { + return std::max(1.0f - fabsf(x), 0.0f); + }; - // Pack K tile transposed: K_f32[dk][kv] so KV_TILE is contiguous (SIMD dim) - // Zero-pad the last tile so the GEMM always operates on KV_TILE_SZ columns - for (int tk = 0; tk < kv_tile; tk++) { - const char * k_data = (const char *)k->data + (ic + tk)*nbk1 + ik2*nbk2 + ik3*nbk3; - if (kv_type == GGML_TYPE_F16) { - const ggml_fp16_t * k_f16 = (const ggml_fp16_t *)k_data; - for (int64_t dk = 0; dk < DK; dk++) { - K_f32[dk * KV_TILE_SZ + tk] = GGML_CPU_FP16_TO_FP32(k_f16[dk]); - } - } else { - const float * k_f32_src = (const float *)k_data; - for (int64_t dk = 0; dk < DK; dk++) { - K_f32[dk * KV_TILE_SZ + tk] = k_f32_src[dk]; - } - } - } - memset(KQ, 0, Q_TILE_SZ * KV_TILE_SZ * sizeof(float)); - simd_gemm(KQ, (const float *)Q_q, K_f32, Q_TILE_SZ, DK, KV_TILE_SZ); - ggml_vec_scale_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, scale); + // support and invscale, minimum 1 pixel for bilinear + const float support1 = std::max(1.0f, 1.0f / sf1); + const float invscale1 = 1.0f / support1; + const float support0 = std::max(1.0f, 1.0f / sf0); + const float invscale0 = 1.0f / support0; - // Set padded KQ entries to -inf so softmax gives them zero weight - if (kv_tile < KV_TILE_SZ) { - for (int tq = 0; tq < Q_TILE_SZ; tq++) { - for (int tk = kv_tile; tk < KV_TILE_SZ; tk++) { - KQ[tq * KV_TILE_SZ + tk] = -INFINITY; - } - } - } + for (int64_t i3 = 0; i3 < ne3; i3++) { + const int64_t i03 = i3 / sf3; + for (int64_t i2 = ith; i2 < ne2; i2 += nth) { + const int64_t i02 = i2 / sf2; + for (int64_t i1 = 0; i1 < ne1; i1++) { + const float y = ((float) i1 + pixel_offset) / sf1; + for (int64_t i0 = 0; i0 < ne0; i0++) { + const float x = ((float) i0 + pixel_offset) / sf0; - if (logit_softcap != 0.0f) { - ggml_vec_tanh_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, KQ); - ggml_vec_scale_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, logit_softcap); - } + // the range of source pixels that contribute + const int64_t x_min = std::max(x - support0 + pixel_offset, 0); + const int64_t x_max = std::min(x + support0 + pixel_offset, ne00); + const int64_t y_min = std::max(y - support1 + pixel_offset, 0); + const int64_t y_max = std::min(y + support1 + pixel_offset, ne01); - if (mask) { - ggml_vec_add_f32(tile_rows * KV_TILE_SZ, KQ, KQ, mask32); - } + // bilinear filter with antialiasing + float val = 0.0f; + float total_weight = 0.0f; - bool skip[Q_TILE_SZ] = {}; + for (int64_t sy = y_min; sy < y_max; sy++) { + const float weight_y = triangle_filter((sy - y + pixel_offset) * invscale1); - for (int tq = 0; tq < Q_TILE_SZ; tq++) { - float * kq_row = KQ + tq * KV_TILE_SZ; + for (int64_t sx = x_min; sx < x_max; sx++) { + const float weight_x = triangle_filter((sx - x + pixel_offset) * invscale0); + const float weight = weight_x * weight_y; - float tile_max; - ggml_vec_max_f32(KV_TILE_SZ, &tile_max, kq_row); + if (weight <= 0.0f) { + continue; + } - if (tile_max == -INFINITY) { - skip[tq] = true; - continue; - } + const float pixel = *(const float *)((const char *)src0->data + sx*nb00 + sy*nb01 + i02*nb02 + i03*nb03); + val += pixel * weight; + total_weight += weight; + } + } - const float Mold = M[tq]; - const float Mnew = fmaxf(Mold, tile_max); + if (total_weight > 0.0f) { + val /= total_weight; + } - if (Mnew > Mold) { - const float ms = expf(Mold - Mnew); - ggml_vec_scale_f32(DV, VKQ32 + tq * DV, ms); - S[tq] *= ms; + float * dst_ptr = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); + *dst_ptr = val; + } } - M[tq] = Mnew; + } + } + } else if (mode == GGML_SCALE_MODE_BILINEAR) { + for (int64_t i3 = 0; i3 < ne3; i3++) { + const int64_t i03 = i3 / sf3; + for (int64_t i2 = ith; i2 < ne2; i2 += nth) { + const int64_t i02 = i2 / sf2; + for (int64_t i1 = 0; i1 < ne1; i1++) { + const float y = ((float)i1 + pixel_offset) / sf1 - pixel_offset; + int64_t y0 = (int64_t)floorf(y); + int64_t y1 = y0 + 1; + y0 = std::max(int64_t(0), std::min(y0, ne01 - 1)); + y1 = std::max(int64_t(0), std::min(y1, ne01 - 1)); - S[tq] += ggml_vec_soft_max_f32(KV_TILE_SZ, kq_row, kq_row, Mnew); - } + float dy = y - (float)y0; + dy = std::max(0.0f, std::min(dy, 1.0f)); - // V accumulation: VKQ32 += softmax(KQ) * V - // Pack V tile to contiguous F32, zero-padded - for (int tk = 0; tk < kv_tile; tk++) { - const char * v_data = (const char *)v->data + (ic + tk)*nbv1 + iv2*nbv2 + iv3*nbv3; - if (kv_type == GGML_TYPE_F16) { - ggml_fp16_to_fp32_row((const ggml_fp16_t *)v_data, V32 + tk * DV, DV); - } else { - memcpy(V32 + tk * DV, v_data, DV * sizeof(float)); - } - } - for (int tq = 0; tq < Q_TILE_SZ; tq++) { - if (skip[tq]) { - memset(KQ + tq * KV_TILE_SZ, 0, KV_TILE_SZ * sizeof(float)); - } - } - simd_gemm(VKQ32, KQ, V32, Q_TILE_SZ, KV_TILE_SZ, DV); - } + for (int64_t i0 = 0; i0 < ne0; i0++) { + const float x = ((float)i0 + pixel_offset) / sf0 - pixel_offset; + int64_t x0 = (int64_t)floorf(x); + int64_t x1 = x0 + 1; - // sinks (apply only to valid rows in the tile) - if (sinks) { - const float s = ((float *)((char *) sinks->data))[h]; + x0 = std::max(int64_t(0), std::min(x0, ne00 - 1)); + x1 = std::max(int64_t(0), std::min(x1, ne00 - 1)); - for (int tq = 0; tq < tile_rows; tq++) { - float ms = 1.0f; - float vs = 1.0f; + float dx = x - (float)x0; + dx = std::max(0.0f, std::min(dx, 1.0f)); - if (s > M[tq]) { - ms = expf(M[tq] - s); - ggml_vec_scale_f32(DV, VKQ32 + tq * DV, ms); - } else { - vs = expf(s - M[tq]); - } + // fetch the four surrounding pixel values and interpolate + const float a = *(const float *)((const char *)src0->data + x0*nb00 + y0*nb01 + i02*nb02 + i03*nb03); + const float b = *(const float *)((const char *)src0->data + x1*nb00 + y0*nb01 + i02*nb02 + i03*nb03); + const float c = *(const float *)((const char *)src0->data + x0*nb00 + y1*nb01 + i02*nb02 + i03*nb03); + const float d = *(const float *)((const char *)src0->data + x1*nb00 + y1*nb01 + i02*nb02 + i03*nb03); - S[tq] = S[tq] * ms + vs; + const float val = a*(1 - dx)*(1 - dy) + b*dx*(1 - dy) + c*(1 - dx)*dy + d*dx*dy; + + float * y_dst = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); + *y_dst = val; + } + } } } + } else if (mode == GGML_SCALE_MODE_BICUBIC) { + // https://en.wikipedia.org/wiki/Bicubic_interpolation#Bicubic_convolution_algorithm + const float a = -0.75f; // use alpha = -0.75 (same as PyTorch) + auto weight1 = [a](float x) { return ((a + 2) * x - (a + 3)) * x * x + 1; }; + auto weight2 = [a](float x) { return ((a * x - 5 * a) * x + 8 * a) * x - 4 * a; }; + auto bicubic = [=](float p0, float p1, float p2, float p3, float x) { + const float w0 = weight2(x + 1); + const float w1 = weight1(x + 0); + const float w2 = weight1(1 - x); + const float w3 = weight2(2 - x); + return p0*w0 + p1*w1 + p2*w2 + p3*w3; + }; - for (int tq = 0; tq < tile_rows; tq++) { - // V /= S - const float S_inv = S[tq] == 0.0f ? 0.0f : 1.0f / S[tq]; - ggml_vec_scale_f32(DV, VKQ32 + tq * DV, S_inv); + for (int64_t i3 = 0; i3 < ne3; i3++) { + const int64_t i03 = i3 / sf3; + for (int64_t i2 = ith; i2 < ne2; i2 += nth) { + const int64_t i02 = i2 / sf2; + for (int64_t i1 = 0; i1 < ne1; i1++) { + const float y = ((float)i1 + pixel_offset) / sf1 - pixel_offset; + const int64_t y0 = (int64_t)floorf(y); + const float dy = y - (float)y0; - // dst indices - const int i1 = iq1 + tq; - const int i2 = iq2; - const int i3 = iq3; + for (int64_t i0 = 0; i0 < ne0; i0++) { + const float x = ((float)i0 + pixel_offset) / sf0 - pixel_offset; + const int64_t x0 = (int64_t)floorf(x); + const float dx = x - (float)x0; - // permute(0, 2, 1, 3) - memcpy((char *) dst->data + (i3*ne2*ne1 + i2 + i1*ne1)*nb1, VKQ32 + tq * DV, nb1); - } + auto p = [=](int64_t x_off, int64_t y_off) -> float { + int64_t i00 = std::max(int64_t(0), std::min(x0 + x_off, ne00 - 1)); + int64_t i01 = std::max(int64_t(0), std::min(y0 + y_off, ne01 - 1)); + return *(const float *)((const char *)src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); + }; - ir += tile_rows; + const float val = bicubic( + bicubic(p(-1,-1), p(0,-1), p(1,-1), p(2,-1), dx), + bicubic(p(-1, 0), p(0, 0), p(1, 0), p(2, 0), dx), + bicubic(p(-1, 1), p(0, 1), p(1, 1), p(2, 1), dx), + bicubic(p(-1, 2), p(0, 2), p(1, 2), p(2, 2), dx), dy); + + float * y_dst = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); + *y_dst = val; + } + } + } + } + } else { + GGML_ABORT("unsupported upscale mode"); } } -// Reduction function: combines partial results across KV chunks -// Partials layout in wdata: [n_q_heads][n_chunks][2 + DV] -static void ggml_flash_attn_ext_reduce_partials( - const ggml_compute_params * params, - ggml_tensor * dst, - const int64_t n_chunks, - const int64_t chunk_size) { +void ggml_compute_forward_upscale( + const ggml_compute_params * params, + ggml_tensor * dst) { - const ggml_tensor * q = dst->src[0]; - const ggml_tensor * k = dst->src[1]; - const ggml_tensor * v = dst->src[2]; + const ggml_tensor * src0 = dst->src[0]; - const int64_t DK = k->ne[0]; - const int64_t DV = v->ne[0]; - const int64_t nek1 = k->ne[1]; - const int64_t n_q_heads = q->ne[2]; + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_upscale_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - const int ith = params->ith; - const int nth = params->nth; - const int64_t wdata_per_thread = DK + 2*DV + CACHE_LINE_SIZE_F32; - float * thread_wdata = (float *) params->wdata + ith * wdata_per_thread; +// ggml_compute_forward_pad - const int64_t partials_offset = nth * (DK + 2*DV + CACHE_LINE_SIZE_F32); - const int64_t partial_size = 2 + DV; - const float * partials_base = (const float *) params->wdata + partials_offset; +template +static void ggml_compute_forward_pad_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - // Output layout - const int64_t ne1 = dst->ne[1]; - const int64_t ne2 = dst->ne[2]; - const size_t nb1 = dst->nb[1]; + const ggml_tensor * src0 = dst->src[0]; - // Each thread reduces a subset of query heads - for (int64_t q_head = ith; q_head < n_q_heads; q_head += nth) { - float M_final = -INFINITY; - float S_final = 0.0f; - float * VKQ_final = thread_wdata; - memset(VKQ_final, 0, DV * sizeof(float)); + assert(dst->nb[0] == sizeof(float)); - // Combine partials from all chunks - for (int64_t chunk_idx = 0; chunk_idx < n_chunks; ++chunk_idx) { - const int64_t ic_start = chunk_idx * chunk_size; - if (ic_start >= nek1) continue; + const int ith = params->ith; + const int nth = params->nth; - const float * partial = partials_base + (q_head * n_chunks + chunk_idx) * partial_size; - const float M_chunk = partial[0]; - const float S_chunk = partial[1]; - const float * VKQ_chunk = partial + 2; + GGML_TENSOR_UNARY_OP_LOCALS - if (S_chunk == 0.0f) continue; + float * dst_ptr = (float *) dst->data; + const int32_t lp0 = ggml_get_op_params_i32(dst, 0); + const int32_t rp0 = ggml_get_op_params_i32(dst, 1); + const int32_t lp1 = ggml_get_op_params_i32(dst, 2); + const int32_t rp1 = ggml_get_op_params_i32(dst, 3); + const int32_t lp2 = ggml_get_op_params_i32(dst, 4); + const int32_t rp2 = ggml_get_op_params_i32(dst, 5); + const int32_t lp3 = ggml_get_op_params_i32(dst, 6); + const int32_t rp3 = ggml_get_op_params_i32(dst, 7); - const float M_new = fmaxf(M_final, M_chunk); - const float scale_old = expf(M_final - M_new); - const float scale_new = expf(M_chunk - M_new); + // TODO: optimize - for (int64_t d = 0; d < DV; ++d) { - VKQ_final[d] = VKQ_final[d] * scale_old + VKQ_chunk[d] * scale_new; - } - S_final = S_final * scale_old + S_chunk * scale_new; - M_final = M_new; - } + for (int64_t i2 = 0; i2 < ne2; ++i2) { + for (int64_t i1 = ith; i1 < ne1; i1 += nth) { + for (int64_t i0 = 0; i0 < ne0; ++i0) { + for (int64_t i3 = 0; i3 < ne3; ++i3) { + // circular means wrap around on a torus, so x and y loop around + if constexpr (circular_t) { + const int64_t dst_idx = i3*(ne0*ne1*ne2) + i2*(ne0*ne1) + i1*ne0 + i0; + const int64_t src_i0 = ggml_wrap_around(i0 - lp0, ne00); + const int64_t src_i1 = ggml_wrap_around(i1 - lp1, ne01); + const int64_t src_i2 = ggml_wrap_around(i2 - lp2, ne02); + const int64_t src_i3 = ggml_wrap_around(i3 - lp3, ne03); - // Normalize and write to output - if (S_final != 0.0f) { - const float S_inv = 1.0f / S_final; - ggml_vec_scale_f32(DV, VKQ_final, S_inv); + const int64_t src_idx = + src_i3*nb03 + + src_i2*nb02 + + src_i1*nb01 + + src_i0*nb00; + + const float * src_ptr = (const float *)((char *) src0->data + src_idx); + dst_ptr[dst_idx] = *src_ptr; + } else { + const int64_t dst_idx = i3*(ne0*ne1*ne2) + i2*(ne0*ne1) + i1*ne0 + i0; + if ((i0 >= lp0 && i0 < ne0 - rp0) \ + && (i1 >= lp1 && i1 < ne1 - rp1) \ + && (i2 >= lp2 && i2 < ne2 - rp2) \ + && (i3 >= lp3 && i3 < ne3 - rp3)) { + const int64_t src_idx = (i3 - lp3)*nb03 + (i2 - lp2)*nb02 + (i1 - lp1)*nb01 + (i0 - lp0)*nb00; + const float * src_ptr = (const float *)((char *) src0->data + src_idx); + dst_ptr[dst_idx] = *src_ptr; + } else { + dst_ptr[dst_idx] = 0; + } + } + } + } } - // iq1=0, iq3=0 for decode - memcpy((char *) dst->data + (0*ne2*ne1 + q_head + 0*ne1)*nb1, VKQ_final, nb1); } } -static void ggml_compute_forward_flash_attn_ext_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * q = dst->src[0]; - const ggml_tensor * k = dst->src[1]; - const ggml_tensor * v = dst->src[2]; - - GGML_TENSOR_LOCALS(int64_t, neq, q, ne) - GGML_TENSOR_LOCALS(size_t, nbq, q, nb) - GGML_TENSOR_LOCALS(int64_t, nek, k, ne) - GGML_TENSOR_LOCALS(size_t, nbk, k, nb) - GGML_TENSOR_LOCALS(int64_t, nev, v, ne) - GGML_TENSOR_LOCALS(size_t, nbv, v, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) - - const int64_t DK = nek0; - const int64_t DV = nev0; - const int64_t N = neq1; +void ggml_compute_forward_pad( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const bool circular = (bool) ggml_get_op_params_i32(dst, 8); + switch (src0->type) { + case GGML_TYPE_F32: + { + if (circular) { + ggml_compute_forward_pad_f32(params, dst); + } else { + ggml_compute_forward_pad_f32(params, dst); + } + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - GGML_ASSERT(ne0 == DV); - GGML_ASSERT(ne2 == N); - - // input tensor rows must be contiguous - GGML_ASSERT(nbq0 == ggml_type_size(q->type)); - GGML_ASSERT(nbk0 == ggml_type_size(k->type)); - GGML_ASSERT(nbv0 == ggml_type_size(v->type)); +// ggml_compute_forward_pad_reflect_1d - GGML_ASSERT(neq0 == DK); - GGML_ASSERT(nek0 == DK); - GGML_ASSERT(nev0 == DV); +void ggml_compute_forward_pad_reflect_1d( + const ggml_compute_params * params, + ggml_tensor * dst) { - GGML_ASSERT(neq1 == N); + const ggml_tensor * src0 = dst->src[0]; - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); const int ith = params->ith; const int nth = params->nth; - // When use_ref is set, force the vec-only reference implementation (no tiling, no KV-chunking) - const bool use_ref = params->use_ref; - - const bool kv_is_f32_or_f16 = (k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_F16); - const bool use_split_kv_path = !use_ref && (neq1 == 1 && neq3 == 1) && kv_is_f32_or_f16 && (k->type == v->type) && q->type == GGML_TYPE_F32 && nek1 >= 512; - - if (use_split_kv_path) { - const int64_t chunk_size = (nek1 + nth - 1) / nth; + const int32_t * opts = (const int32_t *) dst->op_params; + const int p0 = opts[0]; + const int p1 = opts[1]; - // Partials buffer layout: [q_head][kv_chunk][M, S, VKQ] - const int64_t partial_size = 2 + DV; - float * partials_base = (float *) params->wdata + nth * (DK + 2*DV + CACHE_LINE_SIZE_F32); + GGML_TENSOR_UNARY_OP_LOCALS - const int64_t ic_start = ith * chunk_size; - const int64_t ic_end = std::min(ic_start + chunk_size, nek1); + for (int64_t i3 = 0; i3 < ne3; i3++) { + for (int64_t i2 = 0; i2 < ne2; i2++) { + for (int64_t i1 = ith; i1 < ne1; i1 += nth) { + float * left = (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + p0*nb0); + float * right = (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + (ne0-p1-1)*nb0); - const int64_t partial_stride = nth * partial_size; - float * chunk_partials = partials_base + ith * partial_size; + ggml_vec_cpy_f32(ne00, left, (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01)); - if (ic_start < nek1) { - for (int64_t q_head = 0; q_head < neq2; q_head++) { - ggml_compute_forward_flash_attn_ext_f16_one_chunk( - params, dst, q_head, q_head + 1, ic_start, ic_end, - chunk_partials, partial_stride); - } - } else { - for (int64_t q_head = 0; q_head < neq2; q_head++) { - float * q_partials = chunk_partials + q_head * partial_stride; - q_partials[0] = -INFINITY; // M - q_partials[1] = 0.0f; // S + for (int i0 = 1; i0 <= p0; i0++) { left[-i0] = left[i0]; } + for (int i0 = 1; i0 <= p1; i0++) { right[i0] = right[-i0]; } } } + } +} - ggml_barrier(params->threadpool); - ggml_flash_attn_ext_reduce_partials(params, dst, nth, chunk_size); - } else { - - // total rows in q - const int64_t nr = neq1*neq2*neq3; - - // disable for NUMA - const bool disable_chunking = ggml_is_numa(); +// ggml_compute_forward_roll - // 4x chunks per thread - int nth_scaled = nth * 4; - int64_t chunk_size = (nr + nth_scaled - 1) / nth_scaled; - int64_t nchunk = (nr + chunk_size - 1) / chunk_size; +static int64_t ggml_wrap_index(int64_t i, int64_t ne) { + if (i < 0) { + return i + ne; + } else if (i >= ne) { + return i - ne; + } + return i; +} - if (nth == 1 || nchunk < nth || disable_chunking) { - nchunk = nth; - } +static void ggml_compute_forward_roll_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - if (ith == 0) { - ggml_threadpool_chunk_set(params->threadpool, nth); - } + const ggml_tensor * src0 = dst->src[0]; + const float * src_data = (const float *) src0->data; + float * dst_data = (float *) dst->data; - ggml_barrier(params->threadpool); + GGML_TENSOR_UNARY_OP_LOCALS - const int64_t dr = (nr + nchunk - 1) / nchunk; + const int s0 = ggml_get_op_params_i32(dst, 0); + const int s1 = ggml_get_op_params_i32(dst, 1); + const int s2 = ggml_get_op_params_i32(dst, 2); + const int s3 = ggml_get_op_params_i32(dst, 3); - static constexpr int64_t Q_TILE_SZ = ggml_fa_tile_config::Q; - bool use_tiled = !use_ref && - (q->type == GGML_TYPE_F32 && - kv_is_f32_or_f16 && - k->type == v->type && - neq1 >= Q_TILE_SZ); -#ifdef GGML_SIMD - use_tiled &= (DV % GGML_F32_EPR == 0); -#endif - int current_chunk = ith; + const int64_t total = ne1 * ne2 * ne3; + const int64_t per_thread = (total + params->nth) / params->nth; + const int64_t start = params->ith * per_thread; + const int64_t end = std::min(start + per_thread, total); - while (current_chunk < nchunk) { - const int64_t ir0 = dr * current_chunk; - const int64_t ir1 = MIN(ir0 + dr, nr); + for (int64_t i = start; i < end; ++i) { + const int64_t i1 = i % ne1; + const int64_t i2 = (i / ne1) % ne2; + const int64_t i3 = i / (ne2 * ne1); + float * dst_row = dst_data + (i3*nb3 + i2*nb2 + i1*nb1) / sizeof(float); - if (use_tiled) { - ggml_compute_forward_flash_attn_ext_tiled(params, dst, ir0, ir1); - } else { - ggml_compute_forward_flash_attn_ext_f16_one_chunk(params, dst, ir0, ir1, 0, nek1, nullptr, 0); - } + const int64_t i01 = ggml_wrap_index(i1 - s1, ne01); + const int64_t i02 = ggml_wrap_index(i2 - s2, ne02); + const int64_t i03 = ggml_wrap_index(i3 - s3, ne03); + const float * src_row = src_data + (i03*nb03 + i02*nb02 + i01*nb01) / sizeof(float); - current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1); - } + const int64_t s = ggml_wrap_index(-s0, ne00); + const int64_t n = ne00 - s; + ggml_vec_cpy_f32(n, dst_row, src_row + s); + ggml_vec_cpy_f32(s, dst_row + n, src_row); } } -void ggml_compute_forward_flash_attn_ext( +void ggml_compute_forward_roll( const ggml_compute_params * params, ggml_tensor * dst) { - switch (dst->op_params[3]) { - case GGML_PREC_DEFAULT: - case GGML_PREC_F32: + + const ggml_tensor * src0 = dst->src[0]; + + switch (src0->type) { + case GGML_TYPE_F32: { - // uses F32 accumulators - ggml_compute_forward_flash_attn_ext_f16(params, dst); + ggml_compute_forward_roll_f32(params, dst); } break; default: { @@ -8911,334 +10008,1925 @@ void ggml_compute_forward_flash_attn_ext( } } -// ggml_compute_forward_flash_attn_back - -static void ggml_compute_forward_flash_attn_back_f32( - const ggml_compute_params * params, - const bool masked, - ggml_tensor * dst) { +// ggml_compute_forward_arange - const ggml_tensor * q = dst->src[0]; - const ggml_tensor * k = dst->src[1]; - const ggml_tensor * v = dst->src[2]; - const ggml_tensor * d = dst->src[3]; +static void ggml_compute_forward_arange_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - GGML_TENSOR_LOCALS(int64_t, neq, q, ne) - GGML_TENSOR_LOCALS(size_t, nbq, q, nb) - GGML_TENSOR_LOCALS(int64_t, nek, k, ne) - GGML_TENSOR_LOCALS(size_t, nbk, k, nb) - GGML_TENSOR_LOCALS(int64_t, nev, v, ne) - GGML_TENSOR_LOCALS(size_t, nbv, v, nb) - GGML_TENSOR_LOCALS(int64_t, ned, d, ne) - GGML_TENSOR_LOCALS(size_t, nbd, d, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + GGML_ASSERT(dst->nb[0] == sizeof(float)); const int ith = params->ith; const int nth = params->nth; - const int64_t D = neq0; - const int64_t N = neq1; - const int64_t P = nek1 - N; - const int64_t M = P + N; + const float start = ggml_get_op_params_f32(dst, 0); + const float stop = ggml_get_op_params_f32(dst, 1); + const float step = ggml_get_op_params_f32(dst, 2); - const int Mup = ggml_up(M, GGML_SOFT_MAX_UNROLL); - const int mxDM = MAX(D, Mup); + const int64_t steps = (int64_t) ceilf((stop - start) / step); - // GGML_ASSERT(ne0 == D); - // GGML_ASSERT(ne1 == N); - GGML_ASSERT(P >= 0); + GGML_ASSERT(ggml_nelements(dst) == steps); - GGML_ASSERT(nbq0 == sizeof(float)); - GGML_ASSERT(nbk0 == sizeof(float)); - GGML_ASSERT(nbv0 == sizeof(float)); + for (int64_t i = ith; i < steps; i+= nth) { + float value = start + step * i; + ((float *)dst->data)[i] = value; + } +} - GGML_ASSERT(neq0 == D); - GGML_ASSERT(nek0 == D); - GGML_ASSERT(nev1 == D); - GGML_ASSERT(ned0 == D); +void ggml_compute_forward_arange( + const ggml_compute_params * params, + ggml_tensor * dst) { + switch (dst->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_arange_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - GGML_ASSERT(neq1 == N); - GGML_ASSERT(nek1 == N + P); - GGML_ASSERT(nev1 == D); - GGML_ASSERT(ned1 == N); +static void ggml_compute_forward_timestep_embedding_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); + const ggml_tensor * src0 = dst->src[0]; - if (ith == 0) { - memset(dst->data, 0, nb0*ne0*ne1*ne2*ne3); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + + const int ith = params->ith; + const int nth = params->nth; + + GGML_TENSOR_UNARY_OP_LOCALS + + const int dim = ggml_get_op_params_i32(dst, 0); + const int max_period = ggml_get_op_params_i32(dst, 1); + + int half = dim / 2; + + for (int64_t i = 0; i < ne00; i++) { + float * embed_data = (float *)((char *) dst->data + i*nb1); + for (int64_t j = ith; j < half; j += nth) { + float timestep = ((float *)src0->data)[i]; + float freq = (float)expf(-logf(max_period) * j / half); + float arg = timestep * freq; + embed_data[j] = cosf(arg); + embed_data[j + half] = sinf(arg); + } + if (dim % 2 != 0 && ith == 0) { + embed_data[2 * half] = 0.f; + } } - ggml_barrier(params->threadpool); +} - const int64_t elem_q = ggml_nelements(q); - const int64_t elem_k = ggml_nelements(k); +void ggml_compute_forward_timestep_embedding( + const ggml_compute_params * params, + ggml_tensor * dst) { - ggml_type result_type = dst->type; - GGML_ASSERT(ggml_blck_size(result_type) == 1); - const size_t tsize = ggml_type_size(result_type); + const ggml_tensor * src0 = dst->src[0]; - const size_t offs_q = 0; - const size_t offs_k = offs_q + GGML_PAD(elem_q * tsize, GGML_MEM_ALIGN); - const size_t offs_v = offs_k + GGML_PAD(elem_k * tsize, GGML_MEM_ALIGN); + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_timestep_embedding_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - void * grad_q = (char *) dst->data; - void * grad_k = (char *) dst->data + offs_k; - void * grad_v = (char *) dst->data + offs_v; +// ggml_compute_forward_argsort - const size_t nbgq1 = nb0*neq0; - const size_t nbgq2 = nb0*neq0*neq1; - const size_t nbgq3 = nb0*neq0*neq1*neq2; +template +struct cmp_argsort { + const float * data; + bool operator()(int32_t a, int32_t b) const { + if constexpr (order == GGML_SORT_ORDER_ASC) { + return data[a] < data[b]; + } else { + return data[a] > data[b]; + } + } +}; - const size_t nbgk1 = nb0*nek0; - const size_t nbgk2 = nb0*nek0*nek1; - const size_t nbgk3 = nb0*nek0*nek1*neq2; +static void ggml_compute_forward_argsort_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - const size_t nbgv1 = nb0*nev0; - const size_t nbgv2 = nb0*nev0*nev1; - const size_t nbgv3 = nb0*nev0*nev1*neq2; + const ggml_tensor * src0 = dst->src[0]; - // parallelize by k rows using ggml_vec_dot_f32 + GGML_TENSOR_UNARY_OP_LOCALS - // total rows in k - const int nr = nek2*nek3; + GGML_ASSERT(nb0 == sizeof(float)); - // rows per thread - const int dr = (nr + nth - 1)/nth; + const int ith = params->ith; + const int nth = params->nth; - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); + const int64_t nr = ggml_nrows(src0); - const float scale = 1.0f/sqrtf(D); + ggml_sort_order order = (ggml_sort_order) ggml_get_op_params_i32(dst, 0); - //printf("P=%d N=%d D=%d ir0=%d ir1=%d scale = %f\n", P, N, D, ir0, ir1, scale); + for (int64_t i = ith; i < nr; i += nth) { + const float * src_data = (float *)((char *) src0->data + i*nb01); - // how often k2 (and v2) is repeated in q2 - int nrep = neq2/nek2; + int32_t * dst_data = (int32_t *)((char *) dst->data + i*nb1); - for (int ir = ir0; ir < ir1; ++ir) { - // q indices - const int ik3 = ir/(nek2); - const int ik2 = ir - ik3*nek2; + for (int64_t j = 0; j < ne0; j++) { + dst_data[j] = j; + } - const int iq3 = ik3; - const int id3 = ik3; - const int iv3 = ik3; - const int iv2 = ik2; + switch (order) { + case GGML_SORT_ORDER_ASC: + std::sort(dst_data, dst_data + ne0, cmp_argsort{src_data}); + break; - for (int irep = 0; irep < nrep; ++irep) { - const int iq2 = ik2 + irep*nek2; - const int id2 = iq2; + case GGML_SORT_ORDER_DESC: + std::sort(dst_data, dst_data + ne0, cmp_argsort{src_data}); + break; - // (ik2 + irep*nek2) % nek2 == ik2 - for (int iq1 = 0; iq1 < neq1; ++iq1) { - const int id1 = iq1; + default: + GGML_ABORT("invalid sort order"); + } + } +} - // not sure about CACHE_LINE_SIZE_F32.. - // - maybe it must not be multiplied by 2 and excluded from .. in SM 1*(..) offset? - float * S = (float *) params->wdata + ith*2*(mxDM + CACHE_LINE_SIZE_F32) + 0*(mxDM+CACHE_LINE_SIZE_F32); - float * SM = (float *) params->wdata + ith*2*(mxDM + CACHE_LINE_SIZE_F32) + 1*(mxDM+CACHE_LINE_SIZE_F32); +void ggml_compute_forward_argsort( + const ggml_compute_params * params, + ggml_tensor * dst) { - for (int i = M; i < Mup; ++i) { - S[i] = -INFINITY; - } + const ggml_tensor * src0 = dst->src[0]; - const int64_t masked_begin = masked ? (P + iq1 + 1) : M; - for (int64_t ic = 0; ic < masked_begin; ++ic) { - // k indices - const int ik1 = ic; + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_argsort_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - // S indices - const int i1 = ik1; +// ggml_compute_forward_top_k - ggml_vec_dot_f32(neq0, - S + i1, 0, - (float *) ((char *) k->data + (ik1*nbk1 + ik2*nbk2 + ik3*nbk3)), 0, - (float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3)), 0, 1); - } +struct cmp_top_k { + const float * data; + bool operator()(int32_t a, int32_t b) const { + return data[a] > data[b]; + } +}; - // scale - ggml_vec_scale_f32(masked_begin, S, scale); +static void ggml_compute_forward_top_k_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - for (int64_t i = masked_begin; i < M; i++) { - S[i] = -INFINITY; - } + const ggml_tensor * src0 = dst->src[0]; - // softmax - // exclude known -INF S[..] values from max and loop - // dont forget to set their SM values to zero - { - float max = -INFINITY; - ggml_vec_max_f32(masked_begin, &max, S); + GGML_TENSOR_UNARY_OP_LOCALS - ggml_float sum = 0.0; - { -#ifdef GGML_SOFT_MAX_ACCELERATE - max = -max; - vDSP_vsadd(SM, 1, &max, SM, 1, Mup); - vvexpf(SM, SM, &Mup); - ggml_vec_sum_f32(Mup, &sum, SM); -#else - sum = ggml_vec_soft_max_f32(Mup, SM, S, max); -#endif - } + GGML_ASSERT(nb0 == sizeof(float)); - assert(sum > 0.0); + const int ith = params->ith; + const int nth = params->nth; - sum = 1.0/sum; - ggml_vec_scale_f32(masked_begin, SM, sum); + const int64_t nr = ggml_nrows(src0); + + const int top_k = ne0; + + int32_t * tmp = (int32_t *) params->wdata + (ne00 + CACHE_LINE_SIZE_F32) * ith; + + for (int64_t i = ith; i < nr; i += nth) { + const float * src_data = (float *)((char *) src0->data + i*nb01); + + for (int64_t j = 0; j < ne00; j++) { + tmp[j] = j; + } + + std::partial_sort(tmp, tmp + top_k, tmp + ne00, cmp_top_k{src_data}); + + int32_t * dst_data = (int32_t *)((char *) dst->data + i*nb1); + + std::copy(tmp, tmp + top_k, dst_data); + + // emphasize that the order is not important + if (top_k > 1) { + std::swap(dst_data[0], dst_data[1]); + } + } +} + +void ggml_compute_forward_top_k( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_top_k_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( + const ggml_compute_params * params, + ggml_tensor * dst, + int ir0, int ir1, + int64_t ic_start, int64_t ic_end, + float * partials, int64_t partial_stride) { + + const bool write_partials = (partials != nullptr); + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const ggml_tensor * sinks = dst->src[4]; + + GGML_TENSOR_LOCALS(int64_t, neq, q, ne) + GGML_TENSOR_LOCALS(size_t, nbq, q, nb) + GGML_TENSOR_LOCALS(int64_t, nek, k, ne) + GGML_TENSOR_LOCALS(size_t, nbk, k, nb) + GGML_TENSOR_LOCALS(int64_t, nev, v, ne) + GGML_TENSOR_LOCALS(size_t, nbv, v, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + + const int64_t DK = nek0; + const int64_t DV = nev0; + const int64_t N = neq1; + + GGML_ASSERT(ne0 == DV); + GGML_ASSERT(ne2 == N); + + // input tensor rows must be contiguous + GGML_ASSERT(nbq0 == ggml_type_size(q->type)); + GGML_ASSERT(nbk0 == ggml_type_size(k->type)); + GGML_ASSERT(nbv0 == ggml_type_size(v->type)); + + GGML_ASSERT(neq0 == DK); + GGML_ASSERT(nek0 == DK); + GGML_ASSERT(nev0 == DV); + + GGML_ASSERT(neq1 == N); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + // broadcast factors + const int64_t rk2 = neq2/nek2; + const int64_t rk3 = neq3/nek3; + + const int64_t rv2 = neq2/nev2; + const int64_t rv3 = neq3/nev3; + + // parallelize by q rows using ggml_vec_dot_f32 + + float scale = 1.0f; + float max_bias = 0.0f; + float logit_softcap = 0.0f; + + memcpy(&scale, (float *) dst->op_params + 0, sizeof(float)); + memcpy(&max_bias, (float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (float *) dst->op_params + 2, sizeof(float)); + + if (logit_softcap != 0) { + scale /= logit_softcap; + } + + const uint32_t n_head = neq2; + const uint32_t n_head_log2 = 1u << (uint32_t) floor(log2(n_head)); + + const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); + const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + + ggml_type const k_vec_dot_type = ggml_get_type_traits_cpu(k->type)->vec_dot_type; + ggml_from_float_t const q_to_vec_dot = ggml_get_type_traits_cpu(k_vec_dot_type)->from_float; + ggml_vec_dot_t const kq_vec_dot = ggml_get_type_traits_cpu(k->type)->vec_dot; + ggml_to_float_t const v_to_float = ggml_get_type_traits(v->type)->to_float; + + GGML_ASSERT(( q_to_vec_dot) && "fattn: unsupported K-type"); + GGML_ASSERT((v->type == GGML_TYPE_F32 || v_to_float ) && "fattn: unsupported V-type"); + + int ith = params->ith; + + for (int ir = ir0; ir < ir1; ++ir) { + // q indices + const int iq3 = ir/(neq2*neq1); + const int iq2 = (ir - iq3*neq2*neq1)/neq1; + const int iq1 = (ir - iq3*neq2*neq1 - iq2*neq1); + + const uint32_t h = iq2; // head index + const float slope = (max_bias > 0.0f) ? h < n_head_log2 ? powf(m0, h + 1) : powf(m1, 2*(h - n_head_log2) + 1) : 1.0f; + + float S = 0.0f; // sum + float M = -INFINITY; // maximum KQ value + + float * VKQ32 = (float *) params->wdata + ith*(1*DK + 2*DV + CACHE_LINE_SIZE_F32); // FP32 VKQ accumulator + float * V32 = (VKQ32 + 1*DV); // (temporary) FP32 V buffer + ggml_fp16_t * VKQ16 = (ggml_fp16_t *) (VKQ32 + 1*DV); // (temporary) FP16 VKQ accumulator + ggml_fp16_t * Q_q = (ggml_fp16_t *) (VKQ32 + 2*DV); // (temporary) buffer for Q converted to quantized/FP16 + + if (v->type == GGML_TYPE_F16) { + memset(VKQ16, 0, DV*sizeof(ggml_fp16_t)); + } else { + memset(VKQ32, 0, DV*sizeof(float)); + } + + const ggml_fp16_t * mp = mask ? (ggml_fp16_t *)((char *) mask->data + iq1*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]) : NULL; + + // k indices + const int ik3 = iq3 / rk3; + const int ik2 = iq2 / rk2; + + // v indices + const int iv3 = iq3 / rv3; + const int iv2 = iq2 / rv2; + + const float * pq = (const float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3)); + q_to_vec_dot(pq, Q_q, DK); + + // online softmax / attention + // loop over n_kv and n_head_kv + // ref: https://arxiv.org/pdf/2112.05682.pdf + + for (int64_t ic = ic_start; ic < ic_end; ++ic) { + const float mv = mp ? slope*GGML_CPU_FP16_TO_FP32(mp[ic]) : 0.0f; + if (mv == -INFINITY) { + continue; + } + + float s; // KQ value + + const char * k_data = (const char *) k->data + ( ic*nbk1 + ik2*nbk2 + ik3*nbk3); + kq_vec_dot(DK, &s, 0, k_data, 0, Q_q, 0, 1); + + s = s*scale; // scale KQ value + + if (logit_softcap != 0.0f) { + s = logit_softcap*tanhf(s); + } + + s += mv; // apply mask + + const float Mold = M; + + float ms = 1.0f; // upon new higher max val, scale VKQ and KQ sum with this value + float vs = 1.0f; // post-softmax KQ value, expf(s - M) + + const char * v_data = ((const char *) v->data + (ic*nbv1 + iv2*nbv2 + iv3*nbv3)); + + if (v->type == GGML_TYPE_F16) { + if (s > M) { + // s is new maximum, ms < 1.0f, vs == expf(s - s) == 1.0f + M = s; + ms = expf(Mold - M); + + // V = V*expf(Mold - M) + ggml_vec_scale_f16(DV, VKQ16, ms); + } else { + // no new maximum, ms == 1.0f, vs != 1.0f + vs = expf(s - M); + } + + // V += v*expf(s - M) + ggml_vec_mad_f16(DV, VKQ16, (const ggml_fp16_t *) v_data, vs); + } else { + if (s > M) { + // s is new maximum, ms < 1.0f, vs == expf(s - s) == 1.0f + M = s; + ms = expf(Mold - M); + + // V = V*expf(Mold - M) + ggml_vec_scale_f32(DV, VKQ32, ms); + } else { + // no new maximum, ms == 1.0f, vs != 1.0f + vs = expf(s - M); + } + + // V += v*expf(s - M) + if (v_to_float) { + v_to_float(v_data, V32, DV); + ggml_vec_mad_f32(DV, VKQ32, V32, vs); + } else { + // V is F32 + ggml_vec_mad_f32(DV, VKQ32, (const float *) v_data, vs); + } + } + + S = S*ms + vs; // scale and increment sum with partial sum + } + + if (v->type == GGML_TYPE_F16) { + for (int64_t d = 0; d < DV; ++d) { + VKQ32[d] = GGML_CPU_FP16_TO_FP32(VKQ16[d]); + } + } + + // sinks - apply only on the first kv-chunk + if (sinks && ic_start == 0) { + const float s = ((float *)((char *) sinks->data))[h]; + + float ms = 1.0f; + float vs = 1.0f; + + if (s > M) { + ms = expf(M - s); + M = s; + ggml_vec_scale_f32(DV, VKQ32, ms); + } else { + vs = expf(s - M); + } + + S = S*ms + vs; + } + + if (write_partials) { + // Write M, S, VKQ to partials for later reduction + // partials layout: [M, S, VKQ[DV]] per query head + float * partial = partials + ir * partial_stride; + partial[0] = M; + partial[1] = S; + memcpy(partial + 2, VKQ32, DV * sizeof(float)); + } else { + // V /= S + const float S_inv = S == 0.0f ? 0.0f : 1.0f/S; + ggml_vec_scale_f32(DV, VKQ32, S_inv); + + // dst indices + const int i1 = iq1; + const int i2 = iq2; + const int i3 = iq3; + + // permute(0, 2, 1, 3) + memcpy((char *) dst->data + (i3*ne2*ne1 + i2 + i1*ne1)*nb1, VKQ32, nb1); + } + } +} + +static void ggml_compute_forward_flash_attn_ext_tiled( + const ggml_compute_params * params, + ggml_tensor * dst, + int ir0, int ir1) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const ggml_tensor * sinks = dst->src[4]; + + GGML_TENSOR_LOCALS(int64_t, neq, q, ne) + GGML_TENSOR_LOCALS(size_t, nbq, q, nb) + GGML_TENSOR_LOCALS(int64_t, nek, k, ne) + GGML_TENSOR_LOCALS(size_t, nbk, k, nb) + GGML_TENSOR_LOCALS(int64_t, nev, v, ne) + GGML_TENSOR_LOCALS(size_t, nbv, v, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + + const int64_t DK = nek0; + const int64_t DV = nev0; + const int64_t N = neq1; + + GGML_ASSERT(ne0 == DV); + GGML_ASSERT(ne2 == N); + + // input tensor rows must be contiguous + GGML_ASSERT(nbq0 == ggml_type_size(q->type)); + GGML_ASSERT(nbk0 == ggml_type_size(k->type)); + GGML_ASSERT(nbv0 == ggml_type_size(v->type)); + + GGML_ASSERT(neq0 == DK); + GGML_ASSERT(nek0 == DK); + GGML_ASSERT(nev0 == DV); + + GGML_ASSERT(neq1 == N); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + GGML_ASSERT(k->type == v->type); + const ggml_type kv_type = k->type; + + + // broadcast factors + const int64_t rk2 = neq2/nek2; + const int64_t rk3 = neq3/nek3; + + const int64_t rv2 = neq2/nev2; + const int64_t rv3 = neq3/nev3; + + float scale = 1.0f; + float max_bias = 0.0f; + float logit_softcap = 0.0f; + + memcpy(&scale, (float *) dst->op_params + 0, sizeof(float)); + memcpy(&max_bias, (float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (float *) dst->op_params + 2, sizeof(float)); + + if (logit_softcap != 0) { + scale /= logit_softcap; + } + + const uint32_t n_head = neq2; + const uint32_t n_head_log2 = 1u << (uint32_t) floor(log2(n_head)); + + const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); + const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + + int ith = params->ith; + + static constexpr int Q_TILE_SZ = ggml_fa_tile_config::Q; + static constexpr int KV_TILE_SZ = ggml_fa_tile_config::KV; + + int ir = ir0; + while (ir < ir1) { + // q indices for the start of this tile + const int iq3 = ir/(neq2*neq1); + const int iq2 = (ir - iq3*neq2*neq1)/neq1; + const int iq1 = (ir - iq3*neq2*neq1 - iq2*neq1); + + // Number of valid rows in this tile: + // - limited by tile size (Q_TILE_SZ) + // - limited by chunk boundary (ir1 - ir) + // - limited by head boundary (neq1 - iq1) to avoid crossing into next head + const int tile_rows = MIN(Q_TILE_SZ, MIN((int)(ir1 - ir), (int)(neq1 - iq1))); + GGML_ASSERT(tile_rows > 0); + + const uint32_t h = iq2; // head index + const float slope = (max_bias > 0.0f) ? h < n_head_log2 ? powf(m0, h + 1) : powf(m1, 2*(h - n_head_log2) + 1) : 1.0f; + + float S[Q_TILE_SZ]; + float M[Q_TILE_SZ]; + + for (int i = 0 ; i < Q_TILE_SZ; ++i) { + S[i] = 0.; + M[i] = -INFINITY; + } + + // Per-thread scratch layout: + // Q_q: Q_TILE_SZ * DK (converted Q tile — F32 for GEMM, KV type for scalar) + // KQ: Q_TILE_SZ * KV_TILE_SZ (attention scores in float) + // mask: Q_TILE_SZ * KV_TILE_SZ (mask in float) + // VKQ32: Q_TILE_SZ * DV (FP32 output accumulator) + // V32: KV_TILE_SZ * DV (F32 buffer for V tile) + // K_f32: KV_TILE_SZ * DK (F32 buffer for K tile — GEMM path) + float * base = (float *) params->wdata + ith*(Q_TILE_SZ*DK + 2*Q_TILE_SZ*KV_TILE_SZ + Q_TILE_SZ*DV + KV_TILE_SZ*DV + KV_TILE_SZ*DK + CACHE_LINE_SIZE_F32); + + void * Q_q = base; + float * KQ = (float *)((char *)base + Q_TILE_SZ * DK * sizeof(float)); + float * mask32 = KQ + Q_TILE_SZ * KV_TILE_SZ; + float * VKQ32 = mask32 + Q_TILE_SZ * KV_TILE_SZ; + float * V32 = VKQ32 + Q_TILE_SZ * DV; + float * K_f32 = V32 + KV_TILE_SZ * DV; + + memset(VKQ32, 0, Q_TILE_SZ * DV * sizeof(float)); + memset(mask32, 0, Q_TILE_SZ * KV_TILE_SZ * sizeof(float)); + + // k indices + const int ik3 = iq3 / rk3; + const int ik2 = iq2 / rk2; + + // v indices + const int iv3 = iq3 / rv3; + const int iv2 = iq2 / rv2; + + { + float * Q_f32 = (float *)Q_q; + for (int tq = 0; tq < tile_rows; tq++) { + const float * pq = (const float *) ((char *) q->data + ((iq1 + tq)*nbq1 + iq2*nbq2 + iq3*nbq3)); + memcpy(Q_f32 + tq * DK, pq, DK * sizeof(float)); + } + for (int tq = tile_rows; tq < Q_TILE_SZ; tq++) { + memset(Q_f32 + tq * DK, 0, DK * sizeof(float)); + } + } + + memset(K_f32, 0, DK * KV_TILE_SZ * sizeof(float)); + memset(V32, 0, KV_TILE_SZ * DV * sizeof(float)); + + for (int64_t ic = 0; ic < nek1; ic += KV_TILE_SZ) { + const int kv_tile = (int)std::min((int64_t)KV_TILE_SZ, nek1 - ic); + + // skip the tile entirely if all the masks are -inf + if (mask) { + bool can_skip = true; + for (int tq = 0; tq < tile_rows; tq++) { + const ggml_fp16_t * mp_row = (const ggml_fp16_t *)((const char *) mask->data + (iq1 + tq)*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]); + for (int tk = 0; tk < kv_tile; tk++) { + mask32[tq * KV_TILE_SZ + tk] = slope * GGML_CPU_FP16_TO_FP32(mp_row[ic + tk]); + if (mask32[tq * KV_TILE_SZ + tk] != -INFINITY) { + can_skip = false; + } + } + // Pad remaining mask entries with -inf + for (int tk = kv_tile; tk < KV_TILE_SZ; tk++) { + mask32[tq * KV_TILE_SZ + tk] = -INFINITY; + } + } + + if (can_skip) { + continue; + } + } + + // Pack K tile transposed: K_f32[dk][kv] so KV_TILE is contiguous (SIMD dim) + // Zero-pad the last tile so the GEMM always operates on KV_TILE_SZ columns + for (int tk = 0; tk < kv_tile; tk++) { + const char * k_data = (const char *)k->data + (ic + tk)*nbk1 + ik2*nbk2 + ik3*nbk3; + if (kv_type == GGML_TYPE_F16) { + const ggml_fp16_t * k_f16 = (const ggml_fp16_t *)k_data; + for (int64_t dk = 0; dk < DK; dk++) { + K_f32[dk * KV_TILE_SZ + tk] = GGML_CPU_FP16_TO_FP32(k_f16[dk]); + } + } else { + const float * k_f32_src = (const float *)k_data; + for (int64_t dk = 0; dk < DK; dk++) { + K_f32[dk * KV_TILE_SZ + tk] = k_f32_src[dk]; + } + } + } + memset(KQ, 0, Q_TILE_SZ * KV_TILE_SZ * sizeof(float)); + simd_gemm(KQ, (const float *)Q_q, K_f32, Q_TILE_SZ, DK, KV_TILE_SZ); + ggml_vec_scale_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, scale); + + // Set padded KQ entries to -inf so softmax gives them zero weight + if (kv_tile < KV_TILE_SZ) { + for (int tq = 0; tq < Q_TILE_SZ; tq++) { + for (int tk = kv_tile; tk < KV_TILE_SZ; tk++) { + KQ[tq * KV_TILE_SZ + tk] = -INFINITY; + } + } + } + + if (logit_softcap != 0.0f) { + ggml_vec_tanh_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, KQ); + ggml_vec_scale_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, logit_softcap); + } + + if (mask) { + ggml_vec_add_f32(tile_rows * KV_TILE_SZ, KQ, KQ, mask32); + } + + bool skip[Q_TILE_SZ] = {}; + + for (int tq = 0; tq < Q_TILE_SZ; tq++) { + float * kq_row = KQ + tq * KV_TILE_SZ; + + float tile_max; + ggml_vec_max_f32(KV_TILE_SZ, &tile_max, kq_row); + + if (tile_max == -INFINITY) { + skip[tq] = true; + continue; + } + + const float Mold = M[tq]; + const float Mnew = fmaxf(Mold, tile_max); + + if (Mnew > Mold) { + const float ms = expf(Mold - Mnew); + ggml_vec_scale_f32(DV, VKQ32 + tq * DV, ms); + S[tq] *= ms; + } + M[tq] = Mnew; + + + S[tq] += ggml_vec_soft_max_f32(KV_TILE_SZ, kq_row, kq_row, Mnew); + } + + // V accumulation: VKQ32 += softmax(KQ) * V + // Pack V tile to contiguous F32, zero-padded + for (int tk = 0; tk < kv_tile; tk++) { + const char * v_data = (const char *)v->data + (ic + tk)*nbv1 + iv2*nbv2 + iv3*nbv3; + if (kv_type == GGML_TYPE_F16) { + ggml_fp16_to_fp32_row((const ggml_fp16_t *)v_data, V32 + tk * DV, DV); + } else { + memcpy(V32 + tk * DV, v_data, DV * sizeof(float)); + } + } + for (int tq = 0; tq < Q_TILE_SZ; tq++) { + if (skip[tq]) { + memset(KQ + tq * KV_TILE_SZ, 0, KV_TILE_SZ * sizeof(float)); + } + } + simd_gemm(VKQ32, KQ, V32, Q_TILE_SZ, KV_TILE_SZ, DV); + } + + // sinks (apply only to valid rows in the tile) + if (sinks) { + const float s = ((float *)((char *) sinks->data))[h]; + + for (int tq = 0; tq < tile_rows; tq++) { + float ms = 1.0f; + float vs = 1.0f; + + if (s > M[tq]) { + ms = expf(M[tq] - s); + ggml_vec_scale_f32(DV, VKQ32 + tq * DV, ms); + } else { + vs = expf(s - M[tq]); + } + + S[tq] = S[tq] * ms + vs; + } + } + + for (int tq = 0; tq < tile_rows; tq++) { + // V /= S + const float S_inv = S[tq] == 0.0f ? 0.0f : 1.0f / S[tq]; + ggml_vec_scale_f32(DV, VKQ32 + tq * DV, S_inv); + + // dst indices + const int i1 = iq1 + tq; + const int i2 = iq2; + const int i3 = iq3; + + // permute(0, 2, 1, 3) + memcpy((char *) dst->data + (i3*ne2*ne1 + i2 + i1*ne1)*nb1, VKQ32 + tq * DV, nb1); + } + + ir += tile_rows; + } +} + +// Reduction function: combines partial results across KV chunks +// Partials layout in wdata: [n_q_heads][n_chunks][2 + DV] +static void ggml_flash_attn_ext_reduce_partials( + const ggml_compute_params * params, + ggml_tensor * dst, + const int64_t n_chunks, + const int64_t chunk_size) { + + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + + const int64_t DK = k->ne[0]; + const int64_t DV = v->ne[0]; + const int64_t nek1 = k->ne[1]; + const int64_t n_q_heads = q->ne[2]; + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t wdata_per_thread = DK + 2*DV + CACHE_LINE_SIZE_F32; + float * thread_wdata = (float *) params->wdata + ith * wdata_per_thread; + + const int64_t partials_offset = nth * (DK + 2*DV + CACHE_LINE_SIZE_F32); + const int64_t partial_size = 2 + DV; + const float * partials_base = (const float *) params->wdata + partials_offset; + + // Output layout + const int64_t ne1 = dst->ne[1]; + const int64_t ne2 = dst->ne[2]; + const size_t nb1 = dst->nb[1]; + + // Each thread reduces a subset of query heads + for (int64_t q_head = ith; q_head < n_q_heads; q_head += nth) { + float M_final = -INFINITY; + float S_final = 0.0f; + float * VKQ_final = thread_wdata; + memset(VKQ_final, 0, DV * sizeof(float)); + + // Combine partials from all chunks + for (int64_t chunk_idx = 0; chunk_idx < n_chunks; ++chunk_idx) { + const int64_t ic_start = chunk_idx * chunk_size; + if (ic_start >= nek1) continue; + + const float * partial = partials_base + (q_head * n_chunks + chunk_idx) * partial_size; + const float M_chunk = partial[0]; + const float S_chunk = partial[1]; + const float * VKQ_chunk = partial + 2; + + if (S_chunk == 0.0f) continue; + + const float M_new = fmaxf(M_final, M_chunk); + const float scale_old = expf(M_final - M_new); + const float scale_new = expf(M_chunk - M_new); + + for (int64_t d = 0; d < DV; ++d) { + VKQ_final[d] = VKQ_final[d] * scale_old + VKQ_chunk[d] * scale_new; + } + S_final = S_final * scale_old + S_chunk * scale_new; + M_final = M_new; + } + + // Normalize and write to output + if (S_final != 0.0f) { + const float S_inv = 1.0f / S_final; + ggml_vec_scale_f32(DV, VKQ_final, S_inv); + } + // iq1=0, iq3=0 for decode + memcpy((char *) dst->data + (0*ne2*ne1 + q_head + 0*ne1)*nb1, VKQ_final, nb1); + } +} + +static void ggml_compute_forward_flash_attn_ext_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + + GGML_TENSOR_LOCALS(int64_t, neq, q, ne) + GGML_TENSOR_LOCALS(size_t, nbq, q, nb) + GGML_TENSOR_LOCALS(int64_t, nek, k, ne) + GGML_TENSOR_LOCALS(size_t, nbk, k, nb) + GGML_TENSOR_LOCALS(int64_t, nev, v, ne) + GGML_TENSOR_LOCALS(size_t, nbv, v, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + + const int64_t DK = nek0; + const int64_t DV = nev0; + const int64_t N = neq1; + + + GGML_ASSERT(ne0 == DV); + GGML_ASSERT(ne2 == N); + + // input tensor rows must be contiguous + GGML_ASSERT(nbq0 == ggml_type_size(q->type)); + GGML_ASSERT(nbk0 == ggml_type_size(k->type)); + GGML_ASSERT(nbv0 == ggml_type_size(v->type)); + + GGML_ASSERT(neq0 == DK); + GGML_ASSERT(nek0 == DK); + GGML_ASSERT(nev0 == DV); + + GGML_ASSERT(neq1 == N); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + const int ith = params->ith; + const int nth = params->nth; + + // When use_ref is set, force the vec-only reference implementation (no tiling, no KV-chunking) + const bool use_ref = params->use_ref; + + const bool kv_is_f32_or_f16 = (k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_F16); + const bool use_split_kv_path = !use_ref && (neq1 == 1 && neq3 == 1) && kv_is_f32_or_f16 && (k->type == v->type) && q->type == GGML_TYPE_F32 && nek1 >= 512; + + if (use_split_kv_path) { + const int64_t chunk_size = (nek1 + nth - 1) / nth; + + // Partials buffer layout: [q_head][kv_chunk][M, S, VKQ] + const int64_t partial_size = 2 + DV; + float * partials_base = (float *) params->wdata + nth * (DK + 2*DV + CACHE_LINE_SIZE_F32); + + const int64_t ic_start = ith * chunk_size; + const int64_t ic_end = std::min(ic_start + chunk_size, nek1); + + const int64_t partial_stride = nth * partial_size; + float * chunk_partials = partials_base + ith * partial_size; + + if (ic_start < nek1) { + for (int64_t q_head = 0; q_head < neq2; q_head++) { + ggml_compute_forward_flash_attn_ext_f16_one_chunk( + params, dst, q_head, q_head + 1, ic_start, ic_end, + chunk_partials, partial_stride); + } + } else { + for (int64_t q_head = 0; q_head < neq2; q_head++) { + float * q_partials = chunk_partials + q_head * partial_stride; + q_partials[0] = -INFINITY; // M + q_partials[1] = 0.0f; // S + } + } + + ggml_barrier(params->threadpool); + ggml_flash_attn_ext_reduce_partials(params, dst, nth, chunk_size); + } else { + + // total rows in q + const int64_t nr = neq1*neq2*neq3; + + // disable for NUMA + const bool disable_chunking = ggml_is_numa(); + + // 4x chunks per thread + int nth_scaled = nth * 4; + int64_t chunk_size = (nr + nth_scaled - 1) / nth_scaled; + int64_t nchunk = (nr + chunk_size - 1) / chunk_size; + + if (nth == 1 || nchunk < nth || disable_chunking) { + nchunk = nth; + } + + if (ith == 0) { + ggml_threadpool_chunk_set(params->threadpool, nth); + } + + ggml_barrier(params->threadpool); + + const int64_t dr = (nr + nchunk - 1) / nchunk; + + static constexpr int64_t Q_TILE_SZ = ggml_fa_tile_config::Q; + bool use_tiled = !use_ref && + (q->type == GGML_TYPE_F32 && + kv_is_f32_or_f16 && + k->type == v->type && + neq1 >= Q_TILE_SZ); +#ifdef GGML_SIMD + use_tiled &= (DV % GGML_F32_EPR == 0); +#endif + int current_chunk = ith; + + while (current_chunk < nchunk) { + const int64_t ir0 = dr * current_chunk; + const int64_t ir1 = MIN(ir0 + dr, nr); + + if (use_tiled) { + ggml_compute_forward_flash_attn_ext_tiled(params, dst, ir0, ir1); + } else { + ggml_compute_forward_flash_attn_ext_f16_one_chunk(params, dst, ir0, ir1, 0, nek1, nullptr, 0); + } + + current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1); + } + } +} + +void ggml_compute_forward_flash_attn_ext( + const ggml_compute_params * params, + ggml_tensor * dst) { + switch (dst->op_params[3]) { + case GGML_PREC_DEFAULT: + case GGML_PREC_F32: + { + // uses F32 accumulators + ggml_compute_forward_flash_attn_ext_f16(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +// ggml_compute_forward_flash_attn_back + +static void ggml_compute_forward_flash_attn_back_f32( + const ggml_compute_params * params, + const bool masked, + ggml_tensor * dst) { + + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * d = dst->src[3]; + + GGML_TENSOR_LOCALS(int64_t, neq, q, ne) + GGML_TENSOR_LOCALS(size_t, nbq, q, nb) + GGML_TENSOR_LOCALS(int64_t, nek, k, ne) + GGML_TENSOR_LOCALS(size_t, nbk, k, nb) + GGML_TENSOR_LOCALS(int64_t, nev, v, ne) + GGML_TENSOR_LOCALS(size_t, nbv, v, nb) + GGML_TENSOR_LOCALS(int64_t, ned, d, ne) + GGML_TENSOR_LOCALS(size_t, nbd, d, nb) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t D = neq0; + const int64_t N = neq1; + const int64_t P = nek1 - N; + const int64_t M = P + N; + + const int Mup = ggml_up(M, GGML_SOFT_MAX_UNROLL); + const int mxDM = MAX(D, Mup); + + // GGML_ASSERT(ne0 == D); + // GGML_ASSERT(ne1 == N); + GGML_ASSERT(P >= 0); + + GGML_ASSERT(nbq0 == sizeof(float)); + GGML_ASSERT(nbk0 == sizeof(float)); + GGML_ASSERT(nbv0 == sizeof(float)); + + GGML_ASSERT(neq0 == D); + GGML_ASSERT(nek0 == D); + GGML_ASSERT(nev1 == D); + GGML_ASSERT(ned0 == D); + + GGML_ASSERT(neq1 == N); + GGML_ASSERT(nek1 == N + P); + GGML_ASSERT(nev1 == D); + GGML_ASSERT(ned1 == N); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + if (ith == 0) { + memset(dst->data, 0, nb0*ne0*ne1*ne2*ne3); + } + ggml_barrier(params->threadpool); + + const int64_t elem_q = ggml_nelements(q); + const int64_t elem_k = ggml_nelements(k); + + ggml_type result_type = dst->type; + GGML_ASSERT(ggml_blck_size(result_type) == 1); + const size_t tsize = ggml_type_size(result_type); + + const size_t offs_q = 0; + const size_t offs_k = offs_q + GGML_PAD(elem_q * tsize, GGML_MEM_ALIGN); + const size_t offs_v = offs_k + GGML_PAD(elem_k * tsize, GGML_MEM_ALIGN); + + void * grad_q = (char *) dst->data; + void * grad_k = (char *) dst->data + offs_k; + void * grad_v = (char *) dst->data + offs_v; + + const size_t nbgq1 = nb0*neq0; + const size_t nbgq2 = nb0*neq0*neq1; + const size_t nbgq3 = nb0*neq0*neq1*neq2; + + const size_t nbgk1 = nb0*nek0; + const size_t nbgk2 = nb0*nek0*nek1; + const size_t nbgk3 = nb0*nek0*nek1*neq2; + + const size_t nbgv1 = nb0*nev0; + const size_t nbgv2 = nb0*nev0*nev1; + const size_t nbgv3 = nb0*nev0*nev1*neq2; + + // parallelize by k rows using ggml_vec_dot_f32 + + // total rows in k + const int nr = nek2*nek3; + + // rows per thread + const int dr = (nr + nth - 1)/nth; + + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + + const float scale = 1.0f/sqrtf(D); + + //printf("P=%d N=%d D=%d ir0=%d ir1=%d scale = %f\n", P, N, D, ir0, ir1, scale); + + // how often k2 (and v2) is repeated in q2 + int nrep = neq2/nek2; + + for (int ir = ir0; ir < ir1; ++ir) { + // q indices + const int ik3 = ir/(nek2); + const int ik2 = ir - ik3*nek2; + + const int iq3 = ik3; + const int id3 = ik3; + const int iv3 = ik3; + const int iv2 = ik2; + + for (int irep = 0; irep < nrep; ++irep) { + const int iq2 = ik2 + irep*nek2; + const int id2 = iq2; + + // (ik2 + irep*nek2) % nek2 == ik2 + for (int iq1 = 0; iq1 < neq1; ++iq1) { + const int id1 = iq1; + + // not sure about CACHE_LINE_SIZE_F32.. + // - maybe it must not be multiplied by 2 and excluded from .. in SM 1*(..) offset? + float * S = (float *) params->wdata + ith*2*(mxDM + CACHE_LINE_SIZE_F32) + 0*(mxDM+CACHE_LINE_SIZE_F32); + float * SM = (float *) params->wdata + ith*2*(mxDM + CACHE_LINE_SIZE_F32) + 1*(mxDM+CACHE_LINE_SIZE_F32); + + for (int i = M; i < Mup; ++i) { + S[i] = -INFINITY; + } + + const int64_t masked_begin = masked ? (P + iq1 + 1) : M; + for (int64_t ic = 0; ic < masked_begin; ++ic) { + // k indices + const int ik1 = ic; + + // S indices + const int i1 = ik1; + + ggml_vec_dot_f32(neq0, + S + i1, 0, + (float *) ((char *) k->data + (ik1*nbk1 + ik2*nbk2 + ik3*nbk3)), 0, + (float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3)), 0, 1); + } + + // scale + ggml_vec_scale_f32(masked_begin, S, scale); + + for (int64_t i = masked_begin; i < M; i++) { + S[i] = -INFINITY; + } + + // softmax + // exclude known -INF S[..] values from max and loop + // dont forget to set their SM values to zero + { + float max = -INFINITY; + ggml_vec_max_f32(masked_begin, &max, S); + + ggml_float sum = 0.0; + { +#ifdef GGML_SOFT_MAX_ACCELERATE + max = -max; + vDSP_vsadd(SM, 1, &max, SM, 1, Mup); + vvexpf(SM, SM, &Mup); + ggml_vec_sum_f32(Mup, &sum, SM); +#else + sum = ggml_vec_soft_max_f32(Mup, SM, S, max); +#endif + } + + assert(sum > 0.0); + + sum = 1.0/sum; + ggml_vec_scale_f32(masked_begin, SM, sum); + + } + + // step-by-step explanation + { + // forward-process shape grads from backward process + // parallel_for ik2,ik3: + // for irep: + // iq2 = ik2 + irep*nek2 + // k[:D,:M,:,:] [D,M,:,:] grad[k][:D,:M,ik2,ik3] += grad[kcur] + // q[:D,:N,:,:] [D,N,:,:] grad[q][:D,iq1,iq2,iq3] += grad[qcur] + // v[:M,:D,:,:] [M,D,:,:] grad[v][:M,:D,iv2,iv3] += grad[vcur] + // for iq1: + // kcur = k[:D,:M,ik2,ik3] [D,M,1,1] grad[kcur] = grad[S1].T @ qcur + // qcur = q[:D,iq1,iq2,iq3] [D,1,1,1] grad[qcur] = grad[S1] @ kcur + // vcur = v[:M,:D,iv2,iv3] [M,D,1,1] grad[vcur] = grad[S5].T @ S4 + // S0 = -Inf [D,1,1,1] + // ~S1[i] = dot(kcur[:D,i], qcur) + // S1 = qcur @ kcur.T [M,1,1,1] grad[S1] = grad[S2] * scale + // S2 = S1 * scale [M,1,1,1] grad[S2] = diag_mask_zero(grad[S3], P) + // S3 = diag_mask_inf(S2, P) [M,1,1,1] grad[S3] = S4 * (grad[S4] - dot(S4, grad[S4])) + // S4 = softmax(S3) [M,1,1,1] grad[S4] = grad[S5] @ vcur + // ~S5[i] = dot(vcur[:,i], S4) + // S5 = S4 @ vcur.T [D,1,1,1] grad[S5] = d[:D,id1,id2,id3] + // ~dst[i,iq1,iq2,iq3] = S5[i] ^ + // dst[:D,iq1,iq2,iq3] = S5 | grad[dst[:D,iq1,iq2,iq3]] = d[:D,id1,id2,id3] + // dst backward-/ grad[dst] = d + // + // output gradients with their dependencies: + // + // grad[kcur] = grad[S1].T @ qcur + // grad[S1] = diag_mask_zero(grad[S3], P) * scale + // grad[S3] = S4 * (grad[S4] - dot(S4, grad[S4])) + // grad[S4] = grad[S5] @ vcur + // grad[S4] = d[:D,id1,id2,id3] @ vcur + // grad[qcur] = grad[S1] @ kcur + // grad[vcur] = grad[S5].T @ S4 + // grad[vcur] = d[:D,id1,id2,id3].T @ S4 + // + // in post-order: + // + // S1 = qcur @ kcur.T + // S2 = S1 * scale + // S3 = diag_mask_inf(S2, P) + // S4 = softmax(S3) + // grad[S4] = d[:D,id1,id2,id3] @ vcur + // grad[S3] = S4 * (grad[S4] - dot(S4, grad[S4])) + // grad[S1] = diag_mask_zero(grad[S3], P) * scale + // grad[qcur] = grad[S1] @ kcur + // grad[kcur] = grad[S1].T @ qcur + // grad[vcur] = d[:D,id1,id2,id3].T @ S4 + // + // using less variables (SM=S4): + // + // S = diag_mask_inf(qcur @ kcur.T * scale, P) + // SM = softmax(S) + // S = d[:D,iq1,iq2,iq3] @ vcur + // dot_SM_gradSM = dot(SM, S) + // S = SM * (S - dot(SM, S)) + // S = diag_mask_zero(S, P) * scale + // + // grad[q][:D,iq1,iq2,iq3] += S @ kcur + // grad[k][:D,:M,ik2,ik3] += S.T @ qcur + // grad[v][:M,:D,iv2,iv3] += d[:D,id1,id2,id3].T @ SM + } + + // S = gradSM = d[:D,id1,id2,id3] @ vcur[:,:,iv2,iv3] + // S = d[:D,id1,id2,id3] @ vcur[:,:,iv2,iv3] + // for ic: + // S[:M] += vcur[:M,ic,iv2,iv3] * d[ic,id1,id2,id3] + // exclude known future zero S[..] values from operation + ggml_vec_set_f32(masked_begin, S, 0); + for (int64_t ic = 0; ic < D; ++ic) { + ggml_vec_mad_f32(masked_begin, + S, + (float *) ((char *) v->data + ( ic*nbv1 + iv2*nbv2 + iv3*nbv3)), + *(float *) ((char *) d->data + (ic*nbd0 + id1*nbd1 + id2*nbd2 + id3*nbd3))); + } + + // S = SM * (S - dot(SM, S)) + float dot_SM_gradSM = 0; + ggml_vec_dot_f32 (masked_begin, &dot_SM_gradSM, 0, SM, 0, S, 0, 1); + ggml_vec_acc1_f32(M, S, -dot_SM_gradSM); + ggml_vec_mul_f32 (masked_begin, S, S, SM); + + // S = diag_mask_zero(S, P) * scale + // already done by above ggml_vec_set_f32 + + // exclude known zero S[..] values from operation + ggml_vec_scale_f32(masked_begin, S, scale); + + // S shape [M,1] + // SM shape [M,1] + // kcur shape [D,M] + // qcur shape [D,1] + // vcur shape [M,D] + + // grad[q][:D,iq1,iq2,iq3] += S @ kcur + // grad[q][:D,iq1,iq2,iq3] += shape[M,1] @ shape[D,M] + // for ic: + // grad[q][:D,iq1,iq2,iq3] += S[ic] * kcur[:D,ic,ik2,ik3] + // exclude known zero S[..] values from loop + for (int64_t ic = 0; ic < masked_begin; ++ic) { + ggml_vec_mad_f32(D, + (float *) ((char *) grad_q + (iq1*nbgq1 + iq2*nbgq2 + iq3*nbgq3)), + (float *) ((char *) k->data + (ic*nbk1 + ik2*nbk2 + ik3*nbk3)), + S[ic]); + } + + // grad[k][:D,:M,iq2,iq3] += S.T @ qcur + // for ic: + // grad[k][:D,ic,iq2,iq3] += S.T[0,ic] * qcur[:D,0] + // grad[k][:D,ic,iq2,iq3] += S[ic] * qcur[:D,0] + // exclude known zero S[..] values from loop + for (int64_t ic = 0; ic < masked_begin; ++ic) { + ggml_vec_mad_f32(D, + (float *) ((char *) grad_k + (ic*nbgk1 + ik2*nbgk2 + ik3*nbgk3)), + (float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3)), + S[ic]); + } + + // grad[v][:M,:D,iv2,iv3] += d[:D,id1,id2,id3].T @ SM + // for ic: + // grad[v][:M,ic,iv2,iv3] += d[:D,id1,id2,id3].T[0,ic] * SM[:M] + // grad[v][:M,ic,iv2,iv3] += d[ic,id1,id2,id3] * SM[:M] + // exclude known zero SM[..] values from mad + for (int64_t ic = 0; ic < D; ++ic) { + ggml_vec_mad_f32(masked_begin, + (float *) ((char *) grad_v + ( ic*nbgv1 + iv2*nbgv2 + iv3*nbgv3)), + SM, + *(float *) ((char *) d->data + (ic*nbd0 + id1*nbd1 + id2*nbd2 + id3*nbd3))); + } + } + } + } +} + +void ggml_compute_forward_flash_attn_back( + const ggml_compute_params * params, + const bool masked, + ggml_tensor * dst) { + + const ggml_tensor * q = dst->src[0]; + + switch (q->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_flash_attn_back_f32(params, masked, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +// ggml_compute_forward_ssm_conv + +static void ggml_compute_forward_ssm_conv_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; // conv_x + const ggml_tensor * src1 = dst->src[1]; // conv1d.weight + + const int ith = params->ith; + const int nth = params->nth; + + const int nc = src1->ne[0]; // d_conv + const int ncs = src0->ne[0]; // d_conv - 1 + n_t + const int nr = src0->ne[1]; // d_inner + const int n_t = dst->ne[1]; // tokens per sequence + const int n_s = dst->ne[2]; // number of sequences in the batch + + GGML_ASSERT( dst->ne[0] == nr); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(src1->nb[0] == sizeof(float)); + GGML_ASSERT(src0->nb[1] == src0->ne[0]*sizeof(float)); + + // rows per thread + const int dr = (nr + nth - 1)/nth; + + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + const int ir = ir1 - ir0; + + for (int i3 = 0; i3 < n_s; ++i3) { + for (int i2 = 0; i2 < n_t; ++i2) { + // {d_conv - 1 + n_t, d_inner, n_seqs} + // sliding window + const float * s = (const float *) ((const char *) src0->data + ir0*(src0->nb[1]) + i2*(src0->nb[0]) + i3*(src0->nb[2])); // {d_conv, d_inner, n_s} + const float * c = (const float *) ((const char *) src1->data + ir0*(src1->nb[1])); // {d_conv, d_inner} + float * x = (float *) ((char *) dst->data + ir0*(dst->nb[0]) + i2*(dst->nb[1]) + i3*(dst->nb[2])); // {d_inner, n_t, n_s} + + // TODO: transpose the output for smaller strides for big batches? + // d_inner + for (int i1 = 0; i1 < ir; ++i1) { + // rowwise dot product + // NOTE: not using ggml_vec_dot_f32, because its sum is in double precision + float sumf = 0.0f; + + // d_conv + for (int i0 = 0; i0 < nc; ++i0) { + sumf += s[i0 + i1*ncs] * c[i0 + i1*nc]; + } + x[i1] = sumf; + } + } + } +} + +void ggml_compute_forward_ssm_conv( + const ggml_compute_params * params, + ggml_tensor * dst) { + switch (dst->src[0]->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_ssm_conv_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +// ggml_compute_forward_ssm_scan + +static void ggml_compute_forward_ssm_scan_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; // s {d_state, dim, n_head, n_seqs+} + const ggml_tensor * src1 = dst->src[1]; // x {dim, n_head, n_seq_tokens, n_seqs} + const ggml_tensor * src2 = dst->src[2]; // dt {n_head, n_seq_tokens, n_seqs} + const ggml_tensor * src3 = dst->src[3]; // A {d_state, n_head} or {1, n_head} + const ggml_tensor * src4 = dst->src[4]; // B {d_state, n_group, n_seq_tokens, n_seqs} + const ggml_tensor * src5 = dst->src[5]; // C {d_state, n_group, n_seq_tokens, n_seqs} + const ggml_tensor * src6 = dst->src[6]; // ids {n_seqs} + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t nc = src0->ne[0]; // d_state + const int64_t nr = src0->ne[1]; // dim + const int64_t nh = src1->ne[1]; // n_head + const int64_t ng = src4->ne[1]; + const int64_t nt = src1->ne[2]; // number of tokens per sequence + const int64_t ns = src1->ne[3]; // number of sequences in the batch + + // can't use ggml_nbytes because src1 is not necessarily contiguous + const int64_t s_off = ggml_nelements(src1) * ggml_element_size(src1); + + GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*ns == ggml_nelements(dst)); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(src1->nb[0] == sizeof(float)); + GGML_ASSERT(src2->nb[0] == sizeof(float)); + GGML_ASSERT(src3->nb[0] == sizeof(float)); + GGML_ASSERT(src4->nb[0] == sizeof(float)); + GGML_ASSERT(src5->nb[0] == sizeof(float)); + GGML_ASSERT(src6->nb[0] == sizeof(int32_t)); + GGML_ASSERT(nh % ng == 0); + + // heads per thread + const int dh = (nh + nth - 1)/nth; + + // head range for this thread + const int ih0 = dh*ith; + const int ih1 = MIN(ih0 + dh, nh); + + const int32_t * ids = (const int32_t *) src6->data; + + for (int i3 = 0; i3 < ns; ++i3) { + const float * s0 = (const float *) ((const char *) src0->data + ids[i3]*(src0->nb[3])); // {d_state, dim, nh, ns} + float * s = ( float *) (( char *) dst->data + i3*(src0->nb[3]) + s_off); // {d_state, dim, nh, ns} + + for (int i2 = 0; i2 < nt; ++i2) { + const float * x = (const float *) ((const char *) src1->data + i2*(src1->nb[2]) + i3*(src1->nb[3])); // {dim, nh, nt, ns} + const float * dt = (const float *) ((const char *) src2->data + i2*(src2->nb[1]) + i3*(src2->nb[2])); // {nh, nt, ns} + const float * A = (const float *) ((const char *) src3->data); // {d_state, nh} or {1, nh} + const float * B = (const float *) ((const char *) src4->data + i2*(src4->nb[2]) + i3*(src4->nb[3])); // {d_state, ng, nt, ns} + const float * C = (const float *) ((const char *) src5->data + i2*(src5->nb[2]) + i3*(src5->nb[3])); // {d_state, ng, nt, ns} + float * y = ( float *) (( char *) dst->data + i2*(nh*nr*sizeof(float)) + i3*(nt*nh*nr*sizeof(float))); // {dim, nh, nt, ns} + + if (src3->ne[0] == 1) { + // Mamba-2 has a scalar decay factor per head; dA can be outside the state-wise loop + + // n_head + for (int h = ih0; h < ih1; ++h) { + // ref: https://github.com/state-spaces/mamba/blob/62db608da60f6fc790b8ed9f4b3225e95ca15fde/mamba_ssm/ops/triton/softplus.py#L16 + const float dt_soft_plus = ggml_compute_softplus_f32(dt[h]); + const float dA = expf(dt_soft_plus * A[h]); + const int g = h / (nh / ng); // repeat_interleave + + // dim + for (int i1 = 0; i1 < nr; ++i1) { + const int ii = i1 + h*nr; + const float x_dt = x[ii] * dt_soft_plus; + float sumf = 0.0f; +#if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + const int ggml_f32_epr = svcntw(); + const int ggml_f32_step = 1 * ggml_f32_epr; + + const int np = (nc & ~(ggml_f32_step - 1)); + + GGML_F32_VEC sum = GGML_F32_VEC_ZERO; + + GGML_F32_VEC adA = GGML_F32_VEC_SET1(dA); + GGML_F32_VEC axdt = GGML_F32_VEC_SET1(x_dt); + + for (int i = 0; i < np; i += ggml_f32_step) { + // TODO: maybe unroll more? + for (int j = 0; j < 1; j++) { + GGML_F32_VEC t0 = GGML_F32_VEC_LOAD(s0 + i + j*ggml_f32_epr + ii*nc); + GGML_F32_VEC t1 = GGML_F32_VEC_LOAD(B + i + j*ggml_f32_epr + g*nc); + GGML_F32_VEC t2 = GGML_F32_VEC_LOAD(C + i + j*ggml_f32_epr + g*nc); + + t0 = GGML_F32_VEC_MUL(t0, adA); + t1 = GGML_F32_VEC_MUL(t1, axdt); + + t0 = GGML_F32_VEC_ADD(t0, t1); + + sum = GGML_F32_VEC_FMA(sum, t0, t2); + + GGML_F32_VEC_STORE(s + i + j*ggml_f32_epr + ii*nc, t0); + } + } + + sumf = GGML_F32xt_REDUCE_ONE(sum); + #elif defined(__riscv_v_intrinsic) + // todo: RVV implementation + const int np = 0; + #else + const int np = (nc & ~(GGML_F32_STEP - 1)); + + GGML_F32_VEC sum[GGML_F32_ARR] = { GGML_F32_VEC_ZERO }; + + GGML_F32_VEC adA = GGML_F32_VEC_SET1(dA); + GGML_F32_VEC axdt = GGML_F32_VEC_SET1(x_dt); + + GGML_F32_VEC ax[GGML_F32_ARR]; + GGML_F32_VEC ay[GGML_F32_ARR]; + GGML_F32_VEC az[GGML_F32_ARR]; + + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ax[j] = GGML_F32_VEC_LOAD(s0 + i + j*GGML_F32_EPR + ii*nc); + ay[j] = GGML_F32_VEC_LOAD(B + i + j*GGML_F32_EPR + g*nc); + az[j] = GGML_F32_VEC_LOAD(C + i + j*GGML_F32_EPR + g*nc); + + ax[j] = GGML_F32_VEC_MUL(ax[j], adA); + ay[j] = GGML_F32_VEC_MUL(ay[j], axdt); + + ax[j] = GGML_F32_VEC_ADD(ax[j], ay[j]); + + sum[j] = GGML_F32_VEC_FMA(sum[j], ax[j], az[j]); + + GGML_F32_VEC_STORE(s + i + j*GGML_F32_EPR + ii*nc, ax[j]); + } + } + + // reduce sum0..sum3 to sum0 + GGML_F32_VEC_REDUCE(sumf, sum); + #endif +#else + const int np = 0; +#endif + // d_state + for (int i0 = np; i0 < nc; ++i0) { + const int i = i0 + ii*nc; + const int ig = i0 + g*nc; + // state = prev_state * dA + dB * x + const float state = (s0[i] * dA) + (B[ig] * x_dt); + // y = rowwise_dotprod(state, C) + sumf += state * C[ig]; + s[i] = state; + } + y[ii] = sumf; + } + } + } else { + // Mamba-1 has an element-wise decay factor for the states + + // n_head + for (int h = ih0; h < ih1; ++h) { + // ref: https://github.com/state-spaces/mamba/blob/62db608da60f6fc790b8ed9f4b3225e95ca15fde/mamba_ssm/ops/triton/softplus.py#L16 + const float dt_soft_plus = ggml_compute_softplus_f32(dt[h]); + const int g = h / (nh / ng); // repeat_interleave + + // dim + for (int i1 = 0; i1 < nr; ++i1) { + const int ii = i1 + h*nr; + const float x_dt = x[ii] * dt_soft_plus; +#if defined(__ARM_FEATURE_SVE) + svfloat32_t vx_dt = GGML_F32_VEC_SET1(x_dt); + svfloat32_t vdt_soft_plus = GGML_F32_VEC_SET1(dt_soft_plus); + svfloat32_t r1_vector = GGML_F32_VEC_ZERO; + + // d_state + // TODO: what happens when (d_state % svcntw()) != 0? + for (int64_t k = 0; k < nc; k += svcntw()) { + svfloat32_t vA = GGML_F32_VEC_LOAD(&A[h*nc + k]); + svfloat32_t vB = GGML_F32_VEC_LOAD(&B[k + g*nc]); + svfloat32_t vC = GGML_F32_VEC_LOAD(&C[k + g*nc]); + svfloat32_t vs0 = GGML_F32_VEC_LOAD(&s0[ii*nc + k]); + + svfloat32_t t1 = GGML_F32_VEC_MUL(vdt_soft_plus, vA); + t1 = exp_ps_sve(svptrue_b32(), t1); + svfloat32_t t2 = GGML_F32_VEC_MUL(vx_dt, vB); + + vs0 = GGML_F32_VEC_FMA(t2, vs0, t1); + r1_vector = GGML_F32_VEC_ADD(GGML_F32_VEC_MUL(vs0, vC), r1_vector); + + GGML_F32_VEC_STORE(&s[ii*nc + k], vs0); + } + y[ii] = GGML_F32xt_REDUCE_ONE(r1_vector); +#else + float sumf = 0.0f; + // NOTE: can't really use GGML_SIMD here because d_state is usually 16 + // and also because expf is used within the loop. + // d_state + for (int i0 = 0; i0 < nc; ++i0) { + const int i = i0 + ii*nc; + const int ig = i0 + g*nc; + // state = prev_state * dA + dB * x + const float state = (s0[i] * expf(dt_soft_plus * A[i0 + h*nc])) + (B[ig] * x_dt); + // y = rowwise_dotprod(state, C) + sumf += state * C[ig]; + s[i] = state; + } + y[ii] = sumf; +#endif + } + } + } + // use the output as the source when it's not the first token-wise iteration + s0 = s; + } + } +} + +void ggml_compute_forward_ssm_scan( + const ggml_compute_params * params, + ggml_tensor * dst) { + switch (dst->src[0]->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_ssm_scan_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +// ggml_compute_forward_win_part + +static void ggml_compute_forward_win_part_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + GGML_UNUSED(params); + + const ggml_tensor * src0 = dst->src[0]; + + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + + const int32_t nep0 = ((const int32_t *)(dst->op_params))[0]; + const int32_t nep1 = ((const int32_t *)(dst->op_params))[1]; + const int32_t w = ((const int32_t *)(dst->op_params))[2]; + + assert(ne00 == ne0); + assert(ne3 == nep0*nep1); + + // TODO: optimize / multi-thread + for (int py = 0; py < nep1; ++py) { + for (int px = 0; px < nep0; ++px) { + const int64_t i3 = py*nep0 + px; + for (int64_t i2 = 0; i2 < ne2; ++i2) { + for (int64_t i1 = 0; i1 < ne1; ++i1) { + for (int64_t i0 = 0; i0 < ne0; ++i0) { + const int64_t i02 = py*w + i2; + const int64_t i01 = px*w + i1; + const int64_t i00 = i0; + + const int64_t i = i3*ne2*ne1*ne0 + i2*ne1*ne0 + i1*ne0 + i0; + const int64_t j = i02*ne01*ne00 + i01*ne00 + i00; + + if (py*w + i2 >= ne02 || px*w + i1 >= ne01) { + ((float *) dst->data)[i] = 0.0f; + } else { + ((float *) dst->data)[i] = ((float *) src0->data)[j]; + } + } + } + } + } + } +} + +void ggml_compute_forward_win_part( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_win_part_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +// ggml_compute_forward_win_unpart + +static void ggml_compute_forward_win_unpart_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + GGML_UNUSED(params); + + const ggml_tensor * src0 = dst->src[0]; + + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + + const int32_t w = ((const int32_t *)(dst->op_params))[0]; + + // padding + const int px = (w - ne1%w)%w; + //const int py = (w - ne2%w)%w; + + const int npx = (px + ne1)/w; + //const int npy = (py + ne2)/w; + + assert(ne0 == ne00); + + // TODO: optimize / multi-thread + for (int64_t i2 = 0; i2 < ne2; ++i2) { + for (int64_t i1 = 0; i1 < ne1; ++i1) { + for (int64_t i0 = 0; i0 < ne0; ++i0) { + const int ip2 = i2/w; + const int ip1 = i1/w; + + const int64_t i02 = i2%w; + const int64_t i01 = i1%w; + const int64_t i00 = i0; + + const int64_t i = (ip2*npx + ip1)*ne02*ne01*ne00 + i02*ne01*ne00 + i01*ne00 + i00; + const int64_t j = i2*ne1*ne0 + i1*ne0 + i0; + + ((float *) dst->data)[j] = ((float *) src0->data)[i]; + } + } + } +} + +void ggml_compute_forward_win_unpart( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_win_unpart_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +//ggml_compute_forward_unary + +void ggml_compute_forward_unary( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_unary_op op = ggml_get_unary_op(dst); + + switch (op) { + case GGML_UNARY_OP_ABS: + { + ggml_compute_forward_abs(params, dst); + } break; + case GGML_UNARY_OP_SGN: + { + ggml_compute_forward_sgn(params, dst); + } break; + case GGML_UNARY_OP_NEG: + { + ggml_compute_forward_neg(params, dst); + } break; + case GGML_UNARY_OP_STEP: + { + ggml_compute_forward_step(params, dst); + } break; + case GGML_UNARY_OP_TANH: + { + ggml_compute_forward_tanh(params, dst); + } break; + case GGML_UNARY_OP_ELU: + { + ggml_compute_forward_elu(params, dst); + } break; + case GGML_UNARY_OP_RELU: + { + ggml_compute_forward_relu(params, dst); + } break; + case GGML_UNARY_OP_SIGMOID: + { + ggml_compute_forward_sigmoid(params, dst); + } break; + case GGML_UNARY_OP_GELU: + { + ggml_compute_forward_gelu(params, dst); + } break; + case GGML_UNARY_OP_GELU_ERF: + { + ggml_compute_forward_gelu_erf(params, dst); + } break; + case GGML_UNARY_OP_GELU_QUICK: + { + ggml_compute_forward_gelu_quick(params, dst); + } break; + case GGML_UNARY_OP_SILU: + { + ggml_compute_forward_silu(params, dst); + } break; + case GGML_UNARY_OP_HARDSWISH: + { + ggml_compute_forward_hardswish(params, dst); + } break; + case GGML_UNARY_OP_HARDSIGMOID: + { + ggml_compute_forward_hardsigmoid(params, dst); + } break; + case GGML_UNARY_OP_EXP: + { + ggml_compute_forward_exp(params, dst); + } break; + case GGML_UNARY_OP_FLOOR: + { + ggml_compute_forward_floor(params, dst); + } break; + case GGML_UNARY_OP_CEIL: + { + ggml_compute_forward_ceil(params, dst); + } break; + case GGML_UNARY_OP_ROUND: + { + ggml_compute_forward_round(params, dst); + } break; + case GGML_UNARY_OP_TRUNC: + { + ggml_compute_forward_trunc(params, dst); + } break; + case GGML_UNARY_OP_XIELU: + { + ggml_compute_forward_xielu(params, dst); + } break; + case GGML_UNARY_OP_EXPM1: + { + ggml_compute_forward_expm1(params, dst); + } break; + case GGML_UNARY_OP_SOFTPLUS: + { + ggml_compute_forward_softplus(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +//ggml_compute_forward_glu + +void ggml_compute_forward_glu( + const ggml_compute_params * params, + ggml_tensor * dst) { - } + const ggml_glu_op op = ggml_get_glu_op(dst); - // step-by-step explanation - { - // forward-process shape grads from backward process - // parallel_for ik2,ik3: - // for irep: - // iq2 = ik2 + irep*nek2 - // k[:D,:M,:,:] [D,M,:,:] grad[k][:D,:M,ik2,ik3] += grad[kcur] - // q[:D,:N,:,:] [D,N,:,:] grad[q][:D,iq1,iq2,iq3] += grad[qcur] - // v[:M,:D,:,:] [M,D,:,:] grad[v][:M,:D,iv2,iv3] += grad[vcur] - // for iq1: - // kcur = k[:D,:M,ik2,ik3] [D,M,1,1] grad[kcur] = grad[S1].T @ qcur - // qcur = q[:D,iq1,iq2,iq3] [D,1,1,1] grad[qcur] = grad[S1] @ kcur - // vcur = v[:M,:D,iv2,iv3] [M,D,1,1] grad[vcur] = grad[S5].T @ S4 - // S0 = -Inf [D,1,1,1] - // ~S1[i] = dot(kcur[:D,i], qcur) - // S1 = qcur @ kcur.T [M,1,1,1] grad[S1] = grad[S2] * scale - // S2 = S1 * scale [M,1,1,1] grad[S2] = diag_mask_zero(grad[S3], P) - // S3 = diag_mask_inf(S2, P) [M,1,1,1] grad[S3] = S4 * (grad[S4] - dot(S4, grad[S4])) - // S4 = softmax(S3) [M,1,1,1] grad[S4] = grad[S5] @ vcur - // ~S5[i] = dot(vcur[:,i], S4) - // S5 = S4 @ vcur.T [D,1,1,1] grad[S5] = d[:D,id1,id2,id3] - // ~dst[i,iq1,iq2,iq3] = S5[i] ^ - // dst[:D,iq1,iq2,iq3] = S5 | grad[dst[:D,iq1,iq2,iq3]] = d[:D,id1,id2,id3] - // dst backward-/ grad[dst] = d - // - // output gradients with their dependencies: - // - // grad[kcur] = grad[S1].T @ qcur - // grad[S1] = diag_mask_zero(grad[S3], P) * scale - // grad[S3] = S4 * (grad[S4] - dot(S4, grad[S4])) - // grad[S4] = grad[S5] @ vcur - // grad[S4] = d[:D,id1,id2,id3] @ vcur - // grad[qcur] = grad[S1] @ kcur - // grad[vcur] = grad[S5].T @ S4 - // grad[vcur] = d[:D,id1,id2,id3].T @ S4 - // - // in post-order: - // - // S1 = qcur @ kcur.T - // S2 = S1 * scale - // S3 = diag_mask_inf(S2, P) - // S4 = softmax(S3) - // grad[S4] = d[:D,id1,id2,id3] @ vcur - // grad[S3] = S4 * (grad[S4] - dot(S4, grad[S4])) - // grad[S1] = diag_mask_zero(grad[S3], P) * scale - // grad[qcur] = grad[S1] @ kcur - // grad[kcur] = grad[S1].T @ qcur - // grad[vcur] = d[:D,id1,id2,id3].T @ S4 - // - // using less variables (SM=S4): - // - // S = diag_mask_inf(qcur @ kcur.T * scale, P) - // SM = softmax(S) - // S = d[:D,iq1,iq2,iq3] @ vcur - // dot_SM_gradSM = dot(SM, S) - // S = SM * (S - dot(SM, S)) - // S = diag_mask_zero(S, P) * scale - // - // grad[q][:D,iq1,iq2,iq3] += S @ kcur - // grad[k][:D,:M,ik2,ik3] += S.T @ qcur - // grad[v][:M,:D,iv2,iv3] += d[:D,id1,id2,id3].T @ SM - } + switch (op) { + case GGML_GLU_OP_REGLU: + { + ggml_compute_forward_reglu(params, dst); + } break; + case GGML_GLU_OP_GEGLU: + { + ggml_compute_forward_geglu(params, dst); + } break; + case GGML_GLU_OP_SWIGLU: + { + ggml_compute_forward_swiglu(params, dst); + } break; + case GGML_GLU_OP_SWIGLU_OAI: + { + ggml_compute_forward_swiglu_oai(params, dst); + } break; + case GGML_GLU_OP_GEGLU_ERF: + { + ggml_compute_forward_geglu_erf(params, dst); + } break; + case GGML_GLU_OP_GEGLU_QUICK: + { + ggml_compute_forward_geglu_quick(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - // S = gradSM = d[:D,id1,id2,id3] @ vcur[:,:,iv2,iv3] - // S = d[:D,id1,id2,id3] @ vcur[:,:,iv2,iv3] - // for ic: - // S[:M] += vcur[:M,ic,iv2,iv3] * d[ic,id1,id2,id3] - // exclude known future zero S[..] values from operation - ggml_vec_set_f32(masked_begin, S, 0); - for (int64_t ic = 0; ic < D; ++ic) { - ggml_vec_mad_f32(masked_begin, - S, - (float *) ((char *) v->data + ( ic*nbv1 + iv2*nbv2 + iv3*nbv3)), - *(float *) ((char *) d->data + (ic*nbd0 + id1*nbd1 + id2*nbd2 + id3*nbd3))); - } +// ggml_compute_forward_get_rel_pos - // S = SM * (S - dot(SM, S)) - float dot_SM_gradSM = 0; - ggml_vec_dot_f32 (masked_begin, &dot_SM_gradSM, 0, SM, 0, S, 0, 1); - ggml_vec_acc1_f32(M, S, -dot_SM_gradSM); - ggml_vec_mul_f32 (masked_begin, S, S, SM); +static void ggml_compute_forward_get_rel_pos_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { + GGML_UNUSED(params); - // S = diag_mask_zero(S, P) * scale - // already done by above ggml_vec_set_f32 + const ggml_tensor * src0 = dst->src[0]; - // exclude known zero S[..] values from operation - ggml_vec_scale_f32(masked_begin, S, scale); + // ref: https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/image_encoder.py#L292-L322 - // S shape [M,1] - // SM shape [M,1] - // kcur shape [D,M] - // qcur shape [D,1] - // vcur shape [M,D] + GGML_TENSOR_UNARY_OP_LOCALS - // grad[q][:D,iq1,iq2,iq3] += S @ kcur - // grad[q][:D,iq1,iq2,iq3] += shape[M,1] @ shape[D,M] - // for ic: - // grad[q][:D,iq1,iq2,iq3] += S[ic] * kcur[:D,ic,ik2,ik3] - // exclude known zero S[..] values from loop - for (int64_t ic = 0; ic < masked_begin; ++ic) { - ggml_vec_mad_f32(D, - (float *) ((char *) grad_q + (iq1*nbgq1 + iq2*nbgq2 + iq3*nbgq3)), - (float *) ((char *) k->data + (ic*nbk1 + ik2*nbk2 + ik3*nbk3)), - S[ic]); - } + const int64_t w = ne1; - // grad[k][:D,:M,iq2,iq3] += S.T @ qcur - // for ic: - // grad[k][:D,ic,iq2,iq3] += S.T[0,ic] * qcur[:D,0] - // grad[k][:D,ic,iq2,iq3] += S[ic] * qcur[:D,0] - // exclude known zero S[..] values from loop - for (int64_t ic = 0; ic < masked_begin; ++ic) { - ggml_vec_mad_f32(D, - (float *) ((char *) grad_k + (ic*nbgk1 + ik2*nbgk2 + ik3*nbgk3)), - (float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3)), - S[ic]); - } + ggml_fp16_t * src0_data = (ggml_fp16_t *) src0->data; + ggml_fp16_t * dst_data = (ggml_fp16_t *) dst->data; - // grad[v][:M,:D,iv2,iv3] += d[:D,id1,id2,id3].T @ SM - // for ic: - // grad[v][:M,ic,iv2,iv3] += d[:D,id1,id2,id3].T[0,ic] * SM[:M] - // grad[v][:M,ic,iv2,iv3] += d[ic,id1,id2,id3] * SM[:M] - // exclude known zero SM[..] values from mad - for (int64_t ic = 0; ic < D; ++ic) { - ggml_vec_mad_f32(masked_begin, - (float *) ((char *) grad_v + ( ic*nbgv1 + iv2*nbgv2 + iv3*nbgv3)), - SM, - *(float *) ((char *) d->data + (ic*nbd0 + id1*nbd1 + id2*nbd2 + id3*nbd3))); - } + for (int64_t i2 = 0; i2 < ne2; ++i2) { + for (int64_t i1 = 0; i1 < ne1; ++i1) { + const int64_t pos = (w - i1 - 1) + i2; + for (int64_t i0 = 0; i0 < ne0; ++i0) { + dst_data[i2*ne1*ne0 + i1*ne0 + i0] = src0_data[pos*ne00 + i0]; } } } } -void ggml_compute_forward_flash_attn_back( +void ggml_compute_forward_get_rel_pos( const ggml_compute_params * params, - const bool masked, ggml_tensor * dst) { - const ggml_tensor * q = dst->src[0]; + const ggml_tensor * src0 = dst->src[0]; - switch (q->type) { - case GGML_TYPE_F32: + switch (src0->type) { + case GGML_TYPE_F16: + case GGML_TYPE_BF16: { - ggml_compute_forward_flash_attn_back_f32(params, masked, dst); + ggml_compute_forward_get_rel_pos_f16(params, dst); } break; default: { @@ -9247,68 +11935,79 @@ void ggml_compute_forward_flash_attn_back( } } -// ggml_compute_forward_ssm_conv +// ggml_compute_forward_add_rel_pos -static void ggml_compute_forward_ssm_conv_f32( +static void ggml_compute_forward_add_rel_pos_f32( const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; // conv_x - const ggml_tensor * src1 = dst->src[1]; // conv1d.weight - const int ith = params->ith; - const int nth = params->nth; + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + const ggml_tensor * src2 = dst->src[2]; - const int nc = src1->ne[0]; // d_conv - const int ncs = src0->ne[0]; // d_conv - 1 + n_t - const int nr = src0->ne[1]; // d_inner - const int n_t = dst->ne[1]; // tokens per sequence - const int n_s = dst->ne[2]; // number of sequences in the batch + const bool inplace = (bool) ((int32_t *) dst->op_params)[0]; + if (!inplace) { + if (params->ith == 0) { + memcpy((char *) dst->data, (char *) src0->data, ggml_nbytes(dst)); + } + ggml_barrier(params->threadpool); + } + // ref: https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/image_encoder.py#L357-L359 - GGML_ASSERT( dst->ne[0] == nr); - GGML_ASSERT(src0->nb[0] == sizeof(float)); - GGML_ASSERT(src1->nb[0] == sizeof(float)); - GGML_ASSERT(src0->nb[1] == src0->ne[0]*sizeof(float)); + float * src1_data = (float *) src1->data; + float * src2_data = (float *) src2->data; + float * dst_data = (float *) dst->data; - // rows per thread - const int dr = (nr + nth - 1)/nth; + const int64_t ne10 = src1->ne[0]; + const int64_t ne11 = src1->ne[1]; + const int64_t ne12 = src1->ne[2]; + const int64_t ne13 = src1->ne[3]; - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); - const int ir = ir1 - ir0; + const int ith = params->ith; + const int nth = params->nth; - for (int i3 = 0; i3 < n_s; ++i3) { - for (int i2 = 0; i2 < n_t; ++i2) { - // {d_conv - 1 + n_t, d_inner, n_seqs} - // sliding window - const float * s = (const float *) ((const char *) src0->data + ir0*(src0->nb[1]) + i2*(src0->nb[0]) + i3*(src0->nb[2])); // {d_conv, d_inner, n_s} - const float * c = (const float *) ((const char *) src1->data + ir0*(src1->nb[1])); // {d_conv, d_inner} - float * x = (float *) ((char *) dst->data + ir0*(dst->nb[0]) + i2*(dst->nb[1]) + i3*(dst->nb[2])); // {d_inner, n_t, n_s} + // total patches in dst + const int np = ne13; - // TODO: transpose the output for smaller strides for big batches? - // d_inner - for (int i1 = 0; i1 < ir; ++i1) { - // rowwise dot product - // NOTE: not using ggml_vec_dot_f32, because its sum is in double precision - float sumf = 0.0f; + // patches per thread + const int dp = (np + nth - 1)/nth; + + // patch range for this thread + const int ip0 = dp*ith; + const int ip1 = MIN(ip0 + dp, np); + + for (int64_t i13 = ip0; i13 < ip1; ++i13) { + for (int64_t i12 = 0; i12 < ne12; ++i12) { + for (int64_t i11 = 0; i11 < ne11; ++i11) { + const int64_t jp1 = i13*ne12*ne11*ne10 + i12*ne11*ne10 + i11*ne10; + for (int64_t i10 = 0; i10 < ne10; ++i10) { + const int64_t jp0 = jp1 + i10; + const float src1_e = src1_data[jp0]; + const float src2_e = src2_data[jp0]; - // d_conv - for (int i0 = 0; i0 < nc; ++i0) { - sumf += s[i0 + i1*ncs] * c[i0 + i1*nc]; + const int64_t jdh = jp0 * ne10; + const int64_t jdw = jdh - (ne10 - 1) * i10; + + for (int64_t j = 0; j < ne10; ++j) { + dst_data[jdh + j ] += src2_e; + dst_data[jdw + j*ne10] += src1_e; + } } - x[i1] = sumf; } } } } -void ggml_compute_forward_ssm_conv( +void ggml_compute_forward_add_rel_pos( const ggml_compute_params * params, ggml_tensor * dst) { - switch (dst->src[0]->type) { + + const ggml_tensor * src0 = dst->src[0]; + + switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_ssm_conv_f32(params, dst); + ggml_compute_forward_add_rel_pos_f32(params, dst); } break; default: { @@ -9317,282 +12016,402 @@ void ggml_compute_forward_ssm_conv( } } -// ggml_compute_forward_ssm_scan +// ggml_compute_forward_rwkv_wkv6 -static void ggml_compute_forward_ssm_scan_f32( +static void ggml_compute_forward_rwkv_wkv6_f32( const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; // s {d_state, dim, n_head, n_seqs+} - const ggml_tensor * src1 = dst->src[1]; // x {dim, n_head, n_seq_tokens, n_seqs} - const ggml_tensor * src2 = dst->src[2]; // dt {n_head, n_seq_tokens, n_seqs} - const ggml_tensor * src3 = dst->src[3]; // A {d_state, n_head} or {1, n_head} - const ggml_tensor * src4 = dst->src[4]; // B {d_state, n_group, n_seq_tokens, n_seqs} - const ggml_tensor * src5 = dst->src[5]; // C {d_state, n_group, n_seq_tokens, n_seqs} - const ggml_tensor * src6 = dst->src[6]; // ids {n_seqs} + const int64_t T = dst->src[1]->ne[2]; + const int64_t C = dst->ne[0]; + const int64_t HEADS = dst->src[1]->ne[1]; + const int64_t n_seqs = dst->src[5]->ne[1]; + const int64_t head_size = C / HEADS; + + float * dst_data = (float *) dst->data; + float * state = ((float *) dst->data) + C * T; const int ith = params->ith; const int nth = params->nth; - const int64_t nc = src0->ne[0]; // d_state - const int64_t nr = src0->ne[1]; // dim - const int64_t nh = src1->ne[1]; // n_head - const int64_t ng = src4->ne[1]; - const int64_t nt = src1->ne[2]; // number of tokens per sequence - const int64_t ns = src1->ne[3]; // number of sequences in the batch + const int h_start = (HEADS * (ith )) / nth; + const int h_end = ((HEADS * (ith + 1)) / nth < HEADS) ? + (HEADS * (ith + 1)) / nth : HEADS; - // can't use ggml_nbytes because src1 is not necessarily contiguous - const int64_t s_off = ggml_nelements(src1) * ggml_element_size(src1); + float * k = (float *) dst->src[0]->data; + float * v = (float *) dst->src[1]->data; + float * r = (float *) dst->src[2]->data; + float * time_faaaa = (float *) dst->src[3]->data; + float * time_decay = (float *) dst->src[4]->data; - GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*ns == ggml_nelements(dst)); - GGML_ASSERT(src0->nb[0] == sizeof(float)); - GGML_ASSERT(src1->nb[0] == sizeof(float)); - GGML_ASSERT(src2->nb[0] == sizeof(float)); - GGML_ASSERT(src3->nb[0] == sizeof(float)); - GGML_ASSERT(src4->nb[0] == sizeof(float)); - GGML_ASSERT(src5->nb[0] == sizeof(float)); - GGML_ASSERT(src6->nb[0] == sizeof(int32_t)); - GGML_ASSERT(nh % ng == 0); + size_t t_stride = HEADS * head_size; // Same to C - // heads per thread - const int dh = (nh + nth - 1)/nth; + size_t h_stride = C / HEADS; + GGML_ASSERT(C % HEADS == 0); // C must be divisible by HEADS + size_t h_stride_2d = head_size * head_size; - // head range for this thread - const int ih0 = dh*ith; - const int ih1 = MIN(ih0 + dh, nh); + if (ith == 0) { + memset(dst_data, 0, T * C * sizeof(float)); + } + ggml_barrier(params->threadpool); - const int32_t * ids = (const int32_t *) src6->data; - for (int i3 = 0; i3 < ns; ++i3) { - const float * s0 = (const float *) ((const char *) src0->data + ids[i3]*(src0->nb[3])); // {d_state, dim, nh, ns} - float * s = ( float *) (( char *) dst->data + i3*(src0->nb[3]) + s_off); // {d_state, dim, nh, ns} + #if defined(__AVX__) && !defined(__AVX512F__) + #define GGML_F32X GGML_F32x8 + #define GGML_F32X_SET1 GGML_F32x8_SET1 + #define GGML_F32X_LOAD GGML_F32x8_LOAD + #define GGML_F32X_STORE GGML_F32x8_STORE + #define GGML_F32X_MUL GGML_F32x8_MUL + #define GGML_F32X_FMA GGML_F32x8_FMA + #define WKV_VECTOR_SIZE 8 + #elif defined(__AVX512F__) + #define GGML_F32X GGML_F32x16 + #define GGML_F32X_SET1 GGML_F32x16_SET1 + #define GGML_F32X_LOAD GGML_F32x16_LOAD + #define GGML_F32X_STORE GGML_F32x16_STORE + #define GGML_F32X_MUL GGML_F32x16_MUL + #define GGML_F32X_FMA GGML_F32x16_FMA + #define WKV_VECTOR_SIZE 16 + #elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + #define GGML_F32X GGML_F32xt + #define GGML_F32X_SET1 GGML_F32xt_SET1 + #define GGML_F32X_LOAD GGML_F32xt_LOAD + #define GGML_F32X_STORE GGML_F32xt_STORE + #define GGML_F32X_MUL GGML_F32xt_MUL + #define GGML_F32X_FMA GGML_F32xt_FMA + #define WKV_VECTOR_SIZE 8 + #elif defined(__ARM_NEON) && defined(__aarch64__) + #define GGML_F32X GGML_F32x4 + #define GGML_F32X_SET1 GGML_F32x4_SET1 + #define GGML_F32X_LOAD GGML_F32x4_LOAD + #define GGML_F32X_STORE GGML_F32x4_STORE + #define GGML_F32X_MUL GGML_F32x4_MUL + #define GGML_F32X_FMA GGML_F32x4_FMA + #define WKV_VECTOR_SIZE 4 + #endif - for (int i2 = 0; i2 < nt; ++i2) { - const float * x = (const float *) ((const char *) src1->data + i2*(src1->nb[2]) + i3*(src1->nb[3])); // {dim, nh, nt, ns} - const float * dt = (const float *) ((const char *) src2->data + i2*(src2->nb[1]) + i3*(src2->nb[2])); // {nh, nt, ns} - const float * A = (const float *) ((const char *) src3->data); // {d_state, nh} or {1, nh} - const float * B = (const float *) ((const char *) src4->data + i2*(src4->nb[2]) + i3*(src4->nb[3])); // {d_state, ng, nt, ns} - const float * C = (const float *) ((const char *) src5->data + i2*(src5->nb[2]) + i3*(src5->nb[3])); // {d_state, ng, nt, ns} - float * y = ( float *) (( char *) dst->data + i2*(nh*nr*sizeof(float)) + i3*(nt*nh*nr*sizeof(float))); // {dim, nh, nt, ns} + #ifdef WKV_VECTOR_SIZE + int wkv_vector_size; + #if defined(__ARM_FEATURE_SVE) + wkv_vector_size = svcntw(); + #else + wkv_vector_size = WKV_VECTOR_SIZE; + #endif + const int64_t vec_count = head_size / wkv_vector_size; - if (src3->ne[0] == 1) { - // Mamba-2 has a scalar decay factor per head; dA can be outside the state-wise loop + for (int64_t t = 0; t < T; t++) { + size_t t_offset = t * t_stride; + size_t state_offset = head_size * C * (t / (T / n_seqs)); + float * state_cur = state + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[5]->data + state_offset; - // n_head - for (int h = ih0; h < ih1; ++h) { - // ref: https://github.com/state-spaces/mamba/blob/62db608da60f6fc790b8ed9f4b3225e95ca15fde/mamba_ssm/ops/triton/softplus.py#L16 - const float dt_soft_plus = ggml_compute_softplus_f32(dt[h]); - const float dA = expf(dt_soft_plus * A[h]); - const int g = h / (nh / ng); // repeat_interleave + for (int64_t h = h_start; h < h_end; h++) { + size_t h_offset = h * h_stride; + size_t t_h_offset = t_offset + h_offset; + size_t h_2d_offset = h * h_stride_2d; - // dim - for (int i1 = 0; i1 < nr; ++i1) { - const int ii = i1 + h*nr; - const float x_dt = x[ii] * dt_soft_plus; - float sumf = 0.0f; -#if defined(GGML_SIMD) - #if defined(__ARM_FEATURE_SVE) - const int ggml_f32_epr = svcntw(); - const int ggml_f32_step = 1 * ggml_f32_epr; + for (int64_t i = 0; i < head_size; i++) { + size_t t_h_i_offset = t_h_offset + i; + size_t h_i_offset = h_offset + i; + size_t h_2d_i_offset = h_2d_offset + i * h_stride; + + float k_val = k[t_h_i_offset]; + float r_val = r[t_h_i_offset]; + float time_faaaa_val = time_faaaa[h_i_offset]; + float time_decay_val = time_decay[t_h_i_offset]; + + // Broadcast scalar values to vectors + GGML_F32X k_vec = GGML_F32X_SET1(k_val); + GGML_F32X r_vec = GGML_F32X_SET1(r_val); + GGML_F32X time_faaaa_vec = GGML_F32X_SET1(time_faaaa_val); + GGML_F32X time_decay_vec = GGML_F32X_SET1(time_decay_val); + + for (int64_t j = 0; j < vec_count; j++) { + size_t base_j = j * wkv_vector_size; + size_t t_h_j_offset = t_h_offset + base_j; + size_t h_2d_i_j_offset = h_2d_i_offset + base_j; + + // Load x elements at once + GGML_F32X v_vec = GGML_F32X_LOAD(&v[t_h_j_offset]); + GGML_F32X prev_state_vec = GGML_F32X_LOAD(&state_prev[h_2d_i_j_offset]); + GGML_F32X dst_vec = GGML_F32X_LOAD(&dst_data[t_h_j_offset]); + + // Compute kv = v * k + GGML_F32X kv_vec = GGML_F32X_MUL(v_vec, k_vec); + + // Compute temp = kv * time_faaaa + prev_state + GGML_F32X temp_vec = GGML_F32X_FMA(prev_state_vec, kv_vec, time_faaaa_vec); + + // Update dst: dst += temp * r + dst_vec = GGML_F32X_FMA(dst_vec, temp_vec, r_vec); + GGML_F32X_STORE(&dst_data[t_h_j_offset], dst_vec); + + // Update state: state = prev_state * time_decay + kv + GGML_F32X new_state_vec = GGML_F32X_FMA(kv_vec, prev_state_vec, time_decay_vec); + GGML_F32X_STORE(&state_cur[h_2d_i_j_offset], new_state_vec); + } + + // Handle remaining elements, this will not be used. + for (int64_t j = vec_count * wkv_vector_size; j < head_size; j++) { + size_t t_h_j_offset = t_h_offset + j; + size_t h_2d_i_j_offset = h_2d_i_offset + j; + float v_val = v[t_h_j_offset]; + float kv_val = v_val * k_val; + float prev_state_val = state_prev[h_2d_i_j_offset]; + float temp_val = kv_val * time_faaaa_val + prev_state_val; + dst_data[t_h_j_offset] += temp_val * r_val; + state_cur[h_2d_i_j_offset] = prev_state_val * time_decay_val + kv_val; + } + } + } + } + + #else + // basically fused operations: + // dst = r @ (time_faaaa * (k @ v) + state), + // state = time_decay * state + (k @ v), + // recursive through each token + for (int64_t t = 0; t < T; t++) { + size_t t_offset = t * t_stride; + size_t state_offset = head_size * C * (t / (T / n_seqs)); + float * state_cur = state + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[5]->data + state_offset; + + for (int64_t h = h_start; h < h_end; h++) { + size_t h_offset = h * h_stride; + size_t t_h_offset = t_offset + h_offset; + size_t h_2d_offset = h * h_stride_2d; + + for (int64_t i = 0; i < head_size; i++) { + size_t t_h_i_offset = t_h_offset + i; + size_t h_i_offset = h_offset + i; + size_t h_2d_i_offset = h_2d_offset + i * h_stride; + + float k_val = k[t_h_i_offset]; + float r_val = r[t_h_i_offset]; + float time_faaaa_val = time_faaaa[h_i_offset]; + // RWKV v6: different time_decay for each token. + float time_decay_val = time_decay[t_h_i_offset]; + + for (int64_t j = 0; j < head_size; j++) { + size_t t_h_j_offset = t_h_offset + j; + size_t h_2d_i_j_offset = h_2d_i_offset + j; + + float v_val = v[t_h_j_offset]; + float kv_val = v_val * k_val; + float prev_state_val = state_prev[h_2d_i_j_offset]; + float temp_val = kv_val * time_faaaa_val + prev_state_val; + dst_data[t_h_j_offset] += temp_val * r_val; + state_cur[h_2d_i_j_offset] = prev_state_val * time_decay_val + kv_val; + } + } + } + } + #endif +} + + +void ggml_compute_forward_rwkv_wkv6( + const ggml_compute_params * params, + ggml_tensor * dst) { - const int np = (nc & ~(ggml_f32_step - 1)); + const ggml_tensor * src0 = dst->src[0]; - GGML_F32_VEC sum = GGML_F32_VEC_ZERO; + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_rwkv_wkv6_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - GGML_F32_VEC adA = GGML_F32_VEC_SET1(dA); - GGML_F32_VEC axdt = GGML_F32_VEC_SET1(x_dt); +// ggml_compute_forward_gla - for (int i = 0; i < np; i += ggml_f32_step) { - // TODO: maybe unroll more? - for (int j = 0; j < 1; j++) { - GGML_F32_VEC t0 = GGML_F32_VEC_LOAD(s0 + i + j*ggml_f32_epr + ii*nc); - GGML_F32_VEC t1 = GGML_F32_VEC_LOAD(B + i + j*ggml_f32_epr + g*nc); - GGML_F32_VEC t2 = GGML_F32_VEC_LOAD(C + i + j*ggml_f32_epr + g*nc); +static void ggml_compute_forward_gla_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + const int64_t T = dst->src[1]->ne[2]; + const int64_t C = dst->ne[0]; + const int64_t HEADS = dst->src[1]->ne[1]; + const int64_t n_seqs = dst->src[4]->ne[1]; + const int64_t head_size = C / HEADS; + const float scale = ggml_get_op_params_f32(dst, 0); - t0 = GGML_F32_VEC_MUL(t0, adA); - t1 = GGML_F32_VEC_MUL(t1, axdt); + float * dst_data = (float *) dst->data; + float * state = ((float *) dst->data) + C * T; - t0 = GGML_F32_VEC_ADD(t0, t1); + const int ith = params->ith; + const int nth = params->nth; - sum = GGML_F32_VEC_FMA(sum, t0, t2); + const int h_start = (HEADS * (ith )) / nth; + const int h_end = ((HEADS * (ith + 1)) / nth < HEADS) ? + (HEADS * (ith + 1)) / nth : HEADS; - GGML_F32_VEC_STORE(s + i + j*ggml_f32_epr + ii*nc, t0); - } - } + float * k = (float *) dst->src[0]->data; + float * v = (float *) dst->src[1]->data; + float * q = (float *) dst->src[2]->data; + float * g = (float *) dst->src[3]->data; - sumf = GGML_F32xt_REDUCE_ONE(sum); - #elif defined(__riscv_v_intrinsic) - // todo: RVV implementation - const int np = 0; - #else - const int np = (nc & ~(GGML_F32_STEP - 1)); + size_t t_stride = HEADS * head_size; // Same to C - GGML_F32_VEC sum[GGML_F32_ARR] = { GGML_F32_VEC_ZERO }; + size_t h_stride = C / HEADS; + GGML_ASSERT(C % HEADS == 0); // C must be divisible by HEADS + size_t h_stride_2d = head_size * head_size; - GGML_F32_VEC adA = GGML_F32_VEC_SET1(dA); - GGML_F32_VEC axdt = GGML_F32_VEC_SET1(x_dt); + if (ith == 0) { + memset(dst_data, 0, T * C * sizeof(float)); + } + ggml_barrier(params->threadpool); - GGML_F32_VEC ax[GGML_F32_ARR]; - GGML_F32_VEC ay[GGML_F32_ARR]; - GGML_F32_VEC az[GGML_F32_ARR]; - for (int i = 0; i < np; i += GGML_F32_STEP) { - for (int j = 0; j < GGML_F32_ARR; j++) { - ax[j] = GGML_F32_VEC_LOAD(s0 + i + j*GGML_F32_EPR + ii*nc); - ay[j] = GGML_F32_VEC_LOAD(B + i + j*GGML_F32_EPR + g*nc); - az[j] = GGML_F32_VEC_LOAD(C + i + j*GGML_F32_EPR + g*nc); + #if defined(__AVX__) && !defined(__AVX512F__) + #define GGML_F32X GGML_F32x8 + #define GGML_F32X_SET1 GGML_F32x8_SET1 + #define GGML_F32X_LOAD GGML_F32x8_LOAD + #define GGML_F32X_STORE GGML_F32x8_STORE + #define GGML_F32X_MUL GGML_F32x8_MUL + #define GGML_F32X_FMA GGML_F32x8_FMA + #define GLA_VECTOR_SIZE 8 + #elif defined(__AVX512F__) + #define GGML_F32X GGML_F32x16 + #define GGML_F32X_SET1 GGML_F32x16_SET1 + #define GGML_F32X_LOAD GGML_F32x16_LOAD + #define GGML_F32X_STORE GGML_F32x16_STORE + #define GGML_F32X_MUL GGML_F32x16_MUL + #define GGML_F32X_FMA GGML_F32x16_FMA + #define GLA_VECTOR_SIZE 16 + #elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + #define GGML_F32X GGML_F32xt + #define GGML_F32X_SET1 GGML_F32xt_SET1 + #define GGML_F32X_LOAD GGML_F32xt_LOAD + #define GGML_F32X_STORE GGML_F32xt_STORE + #define GGML_F32X_MUL GGML_F32xt_MUL + #define GGML_F32X_FMA GGML_F32xt_FMA + #define GLA_VECTOR_SIZE 8 + #elif defined(__ARM_NEON) && defined(__aarch64__) + #define GGML_F32X GGML_F32x4 + #define GGML_F32X_SET1 GGML_F32x4_SET1 + #define GGML_F32X_LOAD GGML_F32x4_LOAD + #define GGML_F32X_STORE GGML_F32x4_STORE + #define GGML_F32X_MUL GGML_F32x4_MUL + #define GGML_F32X_FMA GGML_F32x4_FMA + #define GLA_VECTOR_SIZE 4 + #endif - ax[j] = GGML_F32_VEC_MUL(ax[j], adA); - ay[j] = GGML_F32_VEC_MUL(ay[j], axdt); + #ifdef GLA_VECTOR_SIZE + int gla_vector_size; + #if defined(__ARM_FEATURE_SVE) + gla_vector_size = svcntw(); + #else + gla_vector_size = GLA_VECTOR_SIZE; + #endif + const int64_t vec_count = head_size / gla_vector_size; - ax[j] = GGML_F32_VEC_ADD(ax[j], ay[j]); + for (int64_t t = 0; t < T; t++) { + size_t t_offset = t * t_stride; + size_t state_offset = head_size * C * (t / (T / n_seqs)); + float * state_cur = state + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[4]->data + state_offset; - sum[j] = GGML_F32_VEC_FMA(sum[j], ax[j], az[j]); + for (int64_t h = h_start; h < h_end; h++) { + size_t h_offset = h * h_stride; + size_t t_h_offset = t_offset + h_offset; + size_t h_2d_offset = h * h_stride_2d; - GGML_F32_VEC_STORE(s + i + j*GGML_F32_EPR + ii*nc, ax[j]); - } - } + for (int64_t i = 0; i < head_size; i++) { + size_t t_h_i_offset = t_h_offset + i; + size_t h_2d_i_offset = h_2d_offset + i * h_stride; - // reduce sum0..sum3 to sum0 - GGML_F32_VEC_REDUCE(sumf, sum); - #endif -#else - const int np = 0; -#endif - // d_state - for (int i0 = np; i0 < nc; ++i0) { - const int i = i0 + ii*nc; - const int ig = i0 + g*nc; - // state = prev_state * dA + dB * x - const float state = (s0[i] * dA) + (B[ig] * x_dt); - // y = rowwise_dotprod(state, C) - sumf += state * C[ig]; - s[i] = state; - } - y[ii] = sumf; - } - } - } else { - // Mamba-1 has an element-wise decay factor for the states + float k_val = k[t_h_i_offset]; + float q_val = q[t_h_i_offset] * scale; + float g_val = g[t_h_i_offset]; - // n_head - for (int h = ih0; h < ih1; ++h) { - // ref: https://github.com/state-spaces/mamba/blob/62db608da60f6fc790b8ed9f4b3225e95ca15fde/mamba_ssm/ops/triton/softplus.py#L16 - const float dt_soft_plus = ggml_compute_softplus_f32(dt[h]); - const int g = h / (nh / ng); // repeat_interleave + // Broadcast scalar values to vectors + GGML_F32X k_vec = GGML_F32X_SET1(k_val); + GGML_F32X q_vec = GGML_F32X_SET1(q_val); + GGML_F32X g_vec = GGML_F32X_SET1(g_val); - // dim - for (int i1 = 0; i1 < nr; ++i1) { - const int ii = i1 + h*nr; - const float x_dt = x[ii] * dt_soft_plus; -#if defined(__ARM_FEATURE_SVE) - svfloat32_t vx_dt = GGML_F32_VEC_SET1(x_dt); - svfloat32_t vdt_soft_plus = GGML_F32_VEC_SET1(dt_soft_plus); - svfloat32_t r1_vector = GGML_F32_VEC_ZERO; + for (int64_t j = 0; j < vec_count; j++) { + size_t base_j = j * gla_vector_size; + size_t t_h_j_offset = t_h_offset + base_j; + size_t h_2d_i_j_offset = h_2d_i_offset + base_j; - // d_state - // TODO: what happens when (d_state % svcntw()) != 0? - for (int64_t k = 0; k < nc; k += svcntw()) { - svfloat32_t vA = GGML_F32_VEC_LOAD(&A[h*nc + k]); - svfloat32_t vB = GGML_F32_VEC_LOAD(&B[k + g*nc]); - svfloat32_t vC = GGML_F32_VEC_LOAD(&C[k + g*nc]); - svfloat32_t vs0 = GGML_F32_VEC_LOAD(&s0[ii*nc + k]); + // Load x elements at once + GGML_F32X v_vec = GGML_F32X_LOAD(&v[t_h_j_offset]); + GGML_F32X prev_state_vec = GGML_F32X_LOAD(&state_prev[h_2d_i_j_offset]); + GGML_F32X dst_vec = GGML_F32X_LOAD(&dst_data[t_h_j_offset]); - svfloat32_t t1 = GGML_F32_VEC_MUL(vdt_soft_plus, vA); - t1 = exp_ps_sve(svptrue_b32(), t1); - svfloat32_t t2 = GGML_F32_VEC_MUL(vx_dt, vB); + // Compute kv = v * k + GGML_F32X kv_vec = GGML_F32X_MUL(v_vec, k_vec); - vs0 = GGML_F32_VEC_FMA(t2, vs0, t1); - r1_vector = GGML_F32_VEC_ADD(GGML_F32_VEC_MUL(vs0, vC), r1_vector); + // Compute temp = prev_state * g + kv + GGML_F32X temp_vec = GGML_F32X_FMA(kv_vec, prev_state_vec, g_vec); - GGML_F32_VEC_STORE(&s[ii*nc + k], vs0); - } - y[ii] = GGML_F32xt_REDUCE_ONE(r1_vector); -#else - float sumf = 0.0f; - // NOTE: can't really use GGML_SIMD here because d_state is usually 16 - // and also because expf is used within the loop. - // d_state - for (int i0 = 0; i0 < nc; ++i0) { - const int i = i0 + ii*nc; - const int ig = i0 + g*nc; - // state = prev_state * dA + dB * x - const float state = (s0[i] * expf(dt_soft_plus * A[i0 + h*nc])) + (B[ig] * x_dt); - // y = rowwise_dotprod(state, C) - sumf += state * C[ig]; - s[i] = state; - } - y[ii] = sumf; -#endif + // Update dst: dst += temp * q + dst_vec = GGML_F32X_FMA(dst_vec, temp_vec, q_vec); + GGML_F32X_STORE(&dst_data[t_h_j_offset], dst_vec); + + // Update state + GGML_F32X_STORE(&state_cur[h_2d_i_j_offset], temp_vec); } - } - } - // use the output as the source when it's not the first token-wise iteration - s0 = s; - } - } -} -void ggml_compute_forward_ssm_scan( - const ggml_compute_params * params, - ggml_tensor * dst) { - switch (dst->src[0]->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_ssm_scan_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); + // Handle remaining elements, this will not be used. + for (int64_t j = vec_count * gla_vector_size; j < head_size; j++) { + size_t t_h_j_offset = t_h_offset + j; + size_t h_2d_i_j_offset = h_2d_i_offset + j; + float v_val = v[t_h_j_offset]; + float kv_val = v_val * k_val; + float prev_state_val = state_prev[h_2d_i_j_offset]; + float temp_val = kv_val + prev_state_val * g_val; + dst_data[t_h_j_offset] += temp_val * q_val; + state_cur[h_2d_i_j_offset] = temp_val; + } + } } - } -} - -// ggml_compute_forward_win_part - -static void ggml_compute_forward_win_part_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { - GGML_UNUSED(params); - - const ggml_tensor * src0 = dst->src[0]; + } - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) + #else + for (int64_t t = 0; t < T; t++) { + size_t t_offset = t * t_stride; + size_t state_offset = head_size * C * (t / (T / n_seqs)); + float * state_cur = state + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[4]->data + state_offset; - const int32_t nep0 = ((const int32_t *)(dst->op_params))[0]; - const int32_t nep1 = ((const int32_t *)(dst->op_params))[1]; - const int32_t w = ((const int32_t *)(dst->op_params))[2]; + for (int64_t h = h_start; h < h_end; h++) { + size_t h_offset = h * h_stride; + size_t t_h_offset = t_offset + h_offset; + size_t h_2d_offset = h * h_stride_2d; - assert(ne00 == ne0); - assert(ne3 == nep0*nep1); + for (int64_t i = 0; i < head_size; i++) { + size_t t_h_i_offset = t_h_offset + i; + size_t h_2d_i_offset = h_2d_offset + i * h_stride; - // TODO: optimize / multi-thread - for (int py = 0; py < nep1; ++py) { - for (int px = 0; px < nep0; ++px) { - const int64_t i3 = py*nep0 + px; - for (int64_t i2 = 0; i2 < ne2; ++i2) { - for (int64_t i1 = 0; i1 < ne1; ++i1) { - for (int64_t i0 = 0; i0 < ne0; ++i0) { - const int64_t i02 = py*w + i2; - const int64_t i01 = px*w + i1; - const int64_t i00 = i0; + float k_val = k[t_h_i_offset]; + float q_val = q[t_h_i_offset] * scale; + float g_val = g[t_h_i_offset]; - const int64_t i = i3*ne2*ne1*ne0 + i2*ne1*ne0 + i1*ne0 + i0; - const int64_t j = i02*ne01*ne00 + i01*ne00 + i00; + for (int64_t j = 0; j < head_size; j++) { + size_t t_h_j_offset = t_h_offset + j; + size_t h_2d_i_j_offset = h_2d_i_offset + j; - if (py*w + i2 >= ne02 || px*w + i1 >= ne01) { - ((float *) dst->data)[i] = 0.0f; - } else { - ((float *) dst->data)[i] = ((float *) src0->data)[j]; - } + float v_val = v[t_h_j_offset]; + float kv_val = v_val * k_val; + float prev_state_val = state_prev[h_2d_i_j_offset]; + float temp_val = prev_state_val * g_val + kv_val; + dst_data[t_h_j_offset] += temp_val * q_val; + state_cur[h_2d_i_j_offset] = temp_val; } } } } - } + #endif } -void ggml_compute_forward_win_part( + +void ggml_compute_forward_gla( const ggml_compute_params * params, ggml_tensor * dst) { @@ -9601,7 +12420,7 @@ void ggml_compute_forward_win_part( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_win_part_f32(params, dst); + ggml_compute_forward_gla_f32(params, dst); } break; default: { @@ -9610,331 +12429,257 @@ void ggml_compute_forward_win_part( } } -// ggml_compute_forward_win_unpart - -static void ggml_compute_forward_win_unpart_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { - GGML_UNUSED(params); - - const ggml_tensor * src0 = dst->src[0]; - - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - - const int32_t w = ((const int32_t *)(dst->op_params))[0]; - - // padding - const int px = (w - ne1%w)%w; - //const int py = (w - ne2%w)%w; +static void ggml_compute_forward_solve_tri_f32(const struct ggml_compute_params * params, struct ggml_tensor * dst) { + const struct ggml_tensor * src0 = dst->src[0]; // A (lower triangular) + const struct ggml_tensor * src1 = dst->src[1]; // B (RHS) - const int npx = (px + ne1)/w; - //const int npy = (py + ne2)/w; + GGML_TENSOR_BINARY_OP_LOCALS; - assert(ne0 == ne00); + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); - // TODO: optimize / multi-thread - for (int64_t i2 = 0; i2 < ne2; ++i2) { - for (int64_t i1 = 0; i1 < ne1; ++i1) { - for (int64_t i0 = 0; i0 < ne0; ++i0) { - const int ip2 = i2/w; - const int ip1 = i1/w; + GGML_ASSERT(ne00 == ne01); // A must be square + GGML_ASSERT(ne0 == ne10); // solution cols == B cols + GGML_ASSERT(ne1 == ne11); // solution rows == B rows - const int64_t i02 = i2%w; - const int64_t i01 = i1%w; - const int64_t i00 = i0; + GGML_ASSERT(ne02 == ne12 && ne12 == ne2); + GGML_ASSERT(ne03 == ne13 && ne13 == ne3); - const int64_t i = (ip2*npx + ip1)*ne02*ne01*ne00 + i02*ne01*ne00 + i01*ne00 + i00; - const int64_t j = i2*ne1*ne0 + i1*ne0 + i0; + const int ith = params->ith; + const int nth = params->nth; - ((float *) dst->data)[j] = ((float *) src0->data)[i]; - } - } - } -} + const int64_t k = ne10; // number of RHS columns + const int64_t n = ne11; // A is n×n + const int64_t nr = ne02 * ne03 * k; // we're parallelizing on columns here, so seq x token x column will be the unit -void ggml_compute_forward_win_unpart( - const ggml_compute_params * params, - ggml_tensor * dst) { + // chunks per thread + const int64_t dr = (nr + nth - 1)/nth; - const ggml_tensor * src0 = dst->src[0]; + // chunk range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_win_unpart_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} + const float * A = (const float *) src0->data; // [n, n, B1, B2] + const float * B = (const float *) src1->data; // [n, k, B1, B2] + float * X = ( float *) dst->data; // [n, k, B1, B2] -//ggml_compute_forward_unary + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t i03 = ir/(ne02*k); + const int64_t i02 = (ir - i03*ne02*k)/k; + const int64_t i01 = (ir - i03*ne02*k - i02*k); -void ggml_compute_forward_unary( - const ggml_compute_params * params, - ggml_tensor * dst) { + const float * A_batch = A + i02 * nb02 / sizeof(float) + i03 * nb03 / sizeof(float); + const float * B_batch = B + i02 * nb12 / sizeof(float) + i03 * nb13 / sizeof(float); - const ggml_unary_op op = ggml_get_unary_op(dst); + float * X_batch = X + i02 * nb2 / sizeof(float) + i03 * nb3 / sizeof(float); - switch (op) { - case GGML_UNARY_OP_ABS: - { - ggml_compute_forward_abs(params, dst); - } break; - case GGML_UNARY_OP_SGN: - { - ggml_compute_forward_sgn(params, dst); - } break; - case GGML_UNARY_OP_NEG: - { - ggml_compute_forward_neg(params, dst); - } break; - case GGML_UNARY_OP_STEP: - { - ggml_compute_forward_step(params, dst); - } break; - case GGML_UNARY_OP_TANH: - { - ggml_compute_forward_tanh(params, dst); - } break; - case GGML_UNARY_OP_ELU: - { - ggml_compute_forward_elu(params, dst); - } break; - case GGML_UNARY_OP_RELU: - { - ggml_compute_forward_relu(params, dst); - } break; - case GGML_UNARY_OP_SIGMOID: - { - ggml_compute_forward_sigmoid(params, dst); - } break; - case GGML_UNARY_OP_GELU: - { - ggml_compute_forward_gelu(params, dst); - } break; - case GGML_UNARY_OP_GELU_ERF: - { - ggml_compute_forward_gelu_erf(params, dst); - } break; - case GGML_UNARY_OP_GELU_QUICK: - { - ggml_compute_forward_gelu_quick(params, dst); - } break; - case GGML_UNARY_OP_SILU: - { - ggml_compute_forward_silu(params, dst); - } break; - case GGML_UNARY_OP_HARDSWISH: - { - ggml_compute_forward_hardswish(params, dst); - } break; - case GGML_UNARY_OP_HARDSIGMOID: - { - ggml_compute_forward_hardsigmoid(params, dst); - } break; - case GGML_UNARY_OP_EXP: - { - ggml_compute_forward_exp(params, dst); - } break; - case GGML_UNARY_OP_FLOOR: - { - ggml_compute_forward_floor(params, dst); - } break; - case GGML_UNARY_OP_CEIL: - { - ggml_compute_forward_ceil(params, dst); - } break; - case GGML_UNARY_OP_ROUND: - { - ggml_compute_forward_round(params, dst); - } break; - case GGML_UNARY_OP_TRUNC: - { - ggml_compute_forward_trunc(params, dst); - } break; - case GGML_UNARY_OP_XIELU: - { - ggml_compute_forward_xielu(params, dst); - } break; - case GGML_UNARY_OP_EXPM1: - { - ggml_compute_forward_expm1(params, dst); - } break; - case GGML_UNARY_OP_SOFTPLUS: - { - ggml_compute_forward_softplus(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); + for (int64_t i00 = 0; i00 < n; ++i00) { + float sum = 0.0f; + for (int64_t t = 0; t < i00; ++t) { + sum += A_batch[i00 * n + t] * X_batch[t * k + i01]; } - } -} -//ggml_compute_forward_glu + const float diag = A_batch[i00 * n + i00]; + assert(diag != 0.0f && "Zero diagonal in triangular matrix"); -void ggml_compute_forward_glu( - const ggml_compute_params * params, - ggml_tensor * dst) { + X_batch[i00 * k + i01] = (B_batch[i00 * k + i01] - sum) / diag; + } + } +} - const ggml_glu_op op = ggml_get_glu_op(dst); +void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - switch (op) { - case GGML_GLU_OP_REGLU: - { - ggml_compute_forward_reglu(params, dst); - } break; - case GGML_GLU_OP_GEGLU: - { - ggml_compute_forward_geglu(params, dst); - } break; - case GGML_GLU_OP_SWIGLU: - { - ggml_compute_forward_swiglu(params, dst); - } break; - case GGML_GLU_OP_SWIGLU_OAI: - { - ggml_compute_forward_swiglu_oai(params, dst); - } break; - case GGML_GLU_OP_GEGLU_ERF: - { - ggml_compute_forward_geglu_erf(params, dst); - } break; - case GGML_GLU_OP_GEGLU_QUICK: - { - ggml_compute_forward_geglu_quick(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32) { + ggml_compute_forward_solve_tri_f32(params, dst); + } else { + GGML_ABORT("fatal error"); } } -// ggml_compute_forward_get_rel_pos +// ggml_compute_forward_gated_delta_net +static void ggml_compute_forward_gated_delta_net_one_chunk( + const ggml_compute_params * params, + ggml_tensor * dst, + int64_t ir0, + int64_t ir1) { -static void ggml_compute_forward_get_rel_pos_f16( - const ggml_compute_params * params, - ggml_tensor * dst) { - GGML_UNUSED(params); + ggml_tensor * src_q = dst->src[0]; + ggml_tensor * src_k = dst->src[1]; + ggml_tensor * src_v = dst->src[2]; + ggml_tensor * src_g = dst->src[3]; + ggml_tensor * src_beta = dst->src[4]; + ggml_tensor * src_state = dst->src[5]; - const ggml_tensor * src0 = dst->src[0]; + const int64_t S_v = src_v->ne[0]; + const int64_t H = src_v->ne[1]; + const int64_t n_tokens = src_v->ne[2]; + const int64_t n_seqs = src_v->ne[3]; - // ref: https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/image_encoder.py#L292-L322 + GGML_ASSERT(ggml_is_contiguous_rows(src_q)); + GGML_ASSERT(ggml_is_contiguous_rows(src_k)); + GGML_ASSERT(ggml_is_contiguous_rows(src_v)); + GGML_ASSERT(ggml_is_contiguous(src_g)); + GGML_ASSERT(ggml_is_contiguous(src_beta)); + GGML_ASSERT(ggml_is_contiguous(src_state)); - GGML_TENSOR_UNARY_OP_LOCALS + GGML_ASSERT(src_g->ne[0] == 1 || src_g->ne[0] == S_v); + GGML_ASSERT(src_beta->ne[0] == 1); - const int64_t w = ne1; + GGML_TENSOR_LOCALS(int64_t, neq, src_q, ne); + GGML_TENSOR_LOCALS(size_t, nbq, src_q, nb); + GGML_TENSOR_LOCALS(int64_t, nek, src_k, ne); + GGML_TENSOR_LOCALS(size_t, nbk, src_k, nb); + GGML_TENSOR_LOCALS(int64_t, nev, src_v, ne); + GGML_TENSOR_LOCALS(size_t, nbv, src_v, nb); + GGML_TENSOR_LOCALS(int64_t, neg, src_g, ne); + GGML_TENSOR_LOCALS(size_t, nbg, src_g, nb); + GGML_TENSOR_LOCALS(size_t, nbb, src_beta, nb); - ggml_fp16_t * src0_data = (ggml_fp16_t *) src0->data; - ggml_fp16_t * dst_data = (ggml_fp16_t *) dst->data; + const bool kda = (neg0 == S_v); - for (int64_t i2 = 0; i2 < ne2; ++i2) { - for (int64_t i1 = 0; i1 < ne1; ++i1) { - const int64_t pos = (w - i1 - 1) + i2; - for (int64_t i0 = 0; i0 < ne0; ++i0) { - dst_data[i2*ne1*ne0 + i1*ne0 + i0] = src0_data[pos*ne00 + i0]; + // scratch layout per thread: [delta(S_v)] + const int64_t scratch_per_thread = S_v; + const int ith = params->ith; + + float * delta = (float *)params->wdata + ith * scratch_per_thread + CACHE_LINE_SIZE_F32; + + // output layout: [attn_scores | new_states] + // attn_scores: S_v * H * n_tokens * n_seqs floats + // new_states: S_v * S_v * H * n_seqs floats + const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs; + float * attn_out_base = (float *)dst->data; + float * state_out_base = (float *)dst->data + attn_score_elems; + + const float * state_in_base = (const float *)src_state->data; + + //const int64_t rq1 = nev1 / neq1; + //const int64_t rk1 = nev1 / nek1; + const int64_t rq3 = nev3 / neq3; + const int64_t rk3 = nev3 / nek3; + + const float scale = 1.0f / sqrtf((float) S_v); + + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t iv1 = ir % H; // head_index + const int64_t iv3 = ir / H; // sequence + + const int64_t iq1 = iv1 % neq1; + const int64_t ik1 = iv1 % nek1; + + const int64_t iq3 = iv3 / rq3; + const int64_t ik3 = iv3 / rk3; + + float * s_out = state_out_base + (iv3 * H + iv1) * S_v * S_v; + + // copy input state into output buffer and operate in-place + const float * s_in = state_in_base + (iv3 * H + iv1) * S_v * S_v; + memcpy(s_out, s_in, S_v * S_v * sizeof(float)); + + // attn output pointer for first token of this (head, seq) + float * attn_data = attn_out_base + (iv3 * n_tokens * H + iv1) * S_v; + + for (int64_t t = 0; t < n_tokens; t++) { + const float * q_d = (const float *)((const char *)src_q->data + iq3 * nbq3 + t * nbq2 + iq1 * nbq1); + const float * k_d = (const float *)((const char *)src_k->data + ik3 * nbk3 + t * nbk2 + ik1 * nbk1); + const float * v_d = (const float *)((const char *)src_v->data + iv3 * nbv3 + t * nbv2 + iv1 * nbv1); + + const float beta_val = *(const float *)((const char *)src_beta->data + iv3 * nbb3 + t * nbb2 + iv1 * nbb1); + const float * g_d = (const float *)((const char *)src_g->data + iv3 * nbg3 + t * nbg2 + iv1 * nbg1); + + // state is stored transposed: s_out[j*S_v + i] = S[i][j] + // so row j of s_out = column j of S (contiguous access) + + if (kda) { + // precompute exp(g) into delta scratch (reused below) + for (int64_t i = 0; i < S_v; ++i) { + delta[i] = expf(g_d[i]); + } + // S[i][:] *= exp(g[i]) => for each row j of M: M[j][i] *= exp(g[i]) + for (int64_t j = 0; j < S_v; ++j) { + ggml_vec_mul_f32(S_v, &s_out[j * S_v], &s_out[j * S_v], delta); + } + } else { + ggml_vec_scale_f32(S_v * S_v, s_out, expf(g_d[0])); } - } - } -} -void ggml_compute_forward_get_rel_pos( - const ggml_compute_params * params, - ggml_tensor * dst) { + // delta[j] = sum_i S[i][j] * k[i] = dot(row j of M, k) + for (int64_t j = 0; j < S_v; ++j) { + float sum = 0.0f; + ggml_vec_dot_f32(S_v, &sum, 0, &s_out[j * S_v], 0, k_d, 0, 1); + delta[j] = (v_d[j] - sum) * beta_val; + } - const ggml_tensor * src0 = dst->src[0]; + // outer product: S[i][j] += k[i] * delta[j] => M[j][i] += delta[j] * k[i] + for (int64_t j = 0; j < S_v; ++j) { + ggml_vec_mad_f32(S_v, &s_out[j * S_v], k_d, delta[j]); + } - switch (src0->type) { - case GGML_TYPE_F16: - case GGML_TYPE_BF16: - { - ggml_compute_forward_get_rel_pos_f16(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); + // attn_out[j] = sum_i S[i][j] * q[i] = dot(row j of M, q) + for (int64_t j = 0; j < S_v; ++j) { + float sum = 0.0f; + ggml_vec_dot_f32(S_v, &sum, 0, &s_out[j * S_v], 0, q_d, 0, 1); + attn_data[j] = sum * scale; } + + attn_data += S_v * H; // advance to next token + } } } -// ggml_compute_forward_add_rel_pos -static void ggml_compute_forward_add_rel_pos_f32( +static void ggml_compute_forward_gated_delta_net_f32( const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - const ggml_tensor * src2 = dst->src[2]; + ggml_tensor * V = dst->src[2]; + int64_t nr = V->ne[1] * V->ne[3]; - const bool inplace = (bool) ((int32_t *) dst->op_params)[0]; - if (!inplace) { - if (params->ith == 0) { - memcpy((char *) dst->data, (char *) src0->data, ggml_nbytes(dst)); - } - ggml_barrier(params->threadpool); - } - // ref: https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/image_encoder.py#L357-L359 + // disable for NUMA + const bool disable_chunking = ggml_is_numa(); - float * src1_data = (float *) src1->data; - float * src2_data = (float *) src2->data; - float * dst_data = (float *) dst->data; + int nth = params->nth; + int ith = params->ith; - const int64_t ne10 = src1->ne[0]; - const int64_t ne11 = src1->ne[1]; - const int64_t ne12 = src1->ne[2]; - const int64_t ne13 = src1->ne[3]; + // 4x chunks per thread + int nth_scaled = nth * 4; + int64_t chunk_size = (nr + nth_scaled - 1) / nth_scaled; + int64_t nchunk = (nr + chunk_size - 1) / chunk_size; - const int ith = params->ith; - const int nth = params->nth; + if (nth == 1 || nchunk < nth || disable_chunking) { + nchunk = nth; + } - // total patches in dst - const int np = ne13; + if (ith == 0) { + ggml_threadpool_chunk_set(params->threadpool, nth); + } - // patches per thread - const int dp = (np + nth - 1)/nth; + ggml_barrier(params->threadpool); - // patch range for this thread - const int ip0 = dp*ith; - const int ip1 = MIN(ip0 + dp, np); + const int64_t dr = (nr + nchunk - 1) / nchunk; - for (int64_t i13 = ip0; i13 < ip1; ++i13) { - for (int64_t i12 = 0; i12 < ne12; ++i12) { - for (int64_t i11 = 0; i11 < ne11; ++i11) { - const int64_t jp1 = i13*ne12*ne11*ne10 + i12*ne11*ne10 + i11*ne10; - for (int64_t i10 = 0; i10 < ne10; ++i10) { - const int64_t jp0 = jp1 + i10; - const float src1_e = src1_data[jp0]; - const float src2_e = src2_data[jp0]; + int current_chunk = ith; - const int64_t jdh = jp0 * ne10; - const int64_t jdw = jdh - (ne10 - 1) * i10; + while (current_chunk < nchunk) { + const int64_t ir0 = dr * current_chunk; + const int64_t ir1 = MIN(ir0 + dr, nr); - for (int64_t j = 0; j < ne10; ++j) { - dst_data[jdh + j ] += src2_e; - dst_data[jdw + j*ne10] += src1_e; - } - } - } - } + ggml_compute_forward_gated_delta_net_one_chunk(params, dst, ir0, ir1); + current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1); } } -void ggml_compute_forward_add_rel_pos( +void ggml_compute_forward_gated_delta_net( const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_add_rel_pos_f32(params, dst); + ggml_compute_forward_gated_delta_net_f32(params, dst); } break; default: { @@ -9943,15 +12688,15 @@ void ggml_compute_forward_add_rel_pos( } } -// ggml_compute_forward_rwkv_wkv6 +// ggml_compute_forward_rwkv_wkv7 -static void ggml_compute_forward_rwkv_wkv6_f32( +static void ggml_compute_forward_rwkv_wkv7_f32( const ggml_compute_params * params, ggml_tensor * dst) { const int64_t T = dst->src[1]->ne[2]; const int64_t C = dst->ne[0]; const int64_t HEADS = dst->src[1]->ne[1]; - const int64_t n_seqs = dst->src[5]->ne[1]; + const int64_t n_seqs = dst->src[6]->ne[1]; const int64_t head_size = C / HEADS; float * dst_data = (float *) dst->data; @@ -9964,172 +12709,177 @@ static void ggml_compute_forward_rwkv_wkv6_f32( const int h_end = ((HEADS * (ith + 1)) / nth < HEADS) ? (HEADS * (ith + 1)) / nth : HEADS; - float * k = (float *) dst->src[0]->data; - float * v = (float *) dst->src[1]->data; - float * r = (float *) dst->src[2]->data; - float * time_faaaa = (float *) dst->src[3]->data; - float * time_decay = (float *) dst->src[4]->data; + float * r = (float *) dst->src[0]->data; + float * w = (float *) dst->src[1]->data; + float * k = (float *) dst->src[2]->data; + float * v = (float *) dst->src[3]->data; + float * a = (float *) dst->src[4]->data; + float * b = (float *) dst->src[5]->data; - size_t t_stride = HEADS * head_size; // Same to C + int64_t t_stride = HEADS * head_size; // Same to C - size_t h_stride = C / HEADS; + int64_t h_stride = C / HEADS; GGML_ASSERT(C % HEADS == 0); // C must be divisible by HEADS - size_t h_stride_2d = head_size * head_size; + int64_t h_stride_2d = head_size * head_size; - if (ith == 0) { - memset(dst_data, 0, T * C * sizeof(float)); - } - ggml_barrier(params->threadpool); + #if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) || defined(__riscv_v_intrinsic) + // scalar Route to scalar implementation //TODO: Write SVE code and RVV code + for (int64_t t = 0; t < T; t++) { + int64_t t_offset = t * t_stride; + int64_t state_offset = head_size * C * (t / (T / n_seqs)); + float * state_cur = state + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[6]->data + state_offset; + for (int64_t h = h_start; h < h_end; h++) { + int64_t h_offset = h * h_stride; + int64_t t_h_offset = t_offset + h_offset; + int64_t h_2d_offset = h * h_stride_2d; - #if defined(__AVX__) && !defined(__AVX512F__) - #define GGML_F32X GGML_F32x8 - #define GGML_F32X_SET1 GGML_F32x8_SET1 - #define GGML_F32X_LOAD GGML_F32x8_LOAD - #define GGML_F32X_STORE GGML_F32x8_STORE - #define GGML_F32X_MUL GGML_F32x8_MUL - #define GGML_F32X_FMA GGML_F32x8_FMA - #define WKV_VECTOR_SIZE 8 - #elif defined(__AVX512F__) - #define GGML_F32X GGML_F32x16 - #define GGML_F32X_SET1 GGML_F32x16_SET1 - #define GGML_F32X_LOAD GGML_F32x16_LOAD - #define GGML_F32X_STORE GGML_F32x16_STORE - #define GGML_F32X_MUL GGML_F32x16_MUL - #define GGML_F32X_FMA GGML_F32x16_FMA - #define WKV_VECTOR_SIZE 16 - #elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) - #define GGML_F32X GGML_F32xt - #define GGML_F32X_SET1 GGML_F32xt_SET1 - #define GGML_F32X_LOAD GGML_F32xt_LOAD - #define GGML_F32X_STORE GGML_F32xt_STORE - #define GGML_F32X_MUL GGML_F32xt_MUL - #define GGML_F32X_FMA GGML_F32xt_FMA - #define WKV_VECTOR_SIZE 8 - #elif defined(__ARM_NEON) && defined(__aarch64__) - #define GGML_F32X GGML_F32x4 - #define GGML_F32X_SET1 GGML_F32x4_SET1 - #define GGML_F32X_LOAD GGML_F32x4_LOAD - #define GGML_F32X_STORE GGML_F32x4_STORE - #define GGML_F32X_MUL GGML_F32x4_MUL - #define GGML_F32X_FMA GGML_F32x4_FMA - #define WKV_VECTOR_SIZE 4 - #endif + for (int64_t i = 0; i < head_size; i++) { + int64_t t_h_i_offset = t_h_offset + i; + int64_t h_2d_i_offset = h_2d_offset + i * h_stride; - #ifdef WKV_VECTOR_SIZE - int wkv_vector_size; - #if defined(__ARM_FEATURE_SVE) - wkv_vector_size = svcntw(); + float v_val = v[t_h_i_offset]; + + float sa = 0, result = 0; + for (int64_t j = 0; j < head_size; j++) { + sa += a[t_h_offset + j] * state_prev[h_2d_i_offset + j]; + } + + for (int64_t j = 0; j < head_size; j++) { + int64_t t_h_j_offset = t_h_offset + j; + int64_t h_2d_i_j_offset = h_2d_i_offset + j; + + float r_val = r[t_h_j_offset]; + float w_val = w[t_h_j_offset]; + float k_val = k[t_h_j_offset]; + float b_val = b[t_h_j_offset]; + float kv_val = v_val * k_val; + float prev_state_val = state_prev[h_2d_i_j_offset]; + state_cur[h_2d_i_j_offset] = prev_state_val * w_val + kv_val + sa * b_val; + result += state_cur[h_2d_i_j_offset] * r_val; + } + dst_data[t_h_i_offset] = result; + } + } + } #else - wkv_vector_size = WKV_VECTOR_SIZE; - #endif - const int64_t vec_count = head_size / wkv_vector_size; + for (int64_t t = 0; t < T; t++) { + int64_t t_offset = t * t_stride; + int64_t state_offset = head_size * C * (t / (T / n_seqs)); + float * state_cur = state + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[6]->data + state_offset; - for (int64_t t = 0; t < T; t++) { - size_t t_offset = t * t_stride; - size_t state_offset = head_size * C * (t / (T / n_seqs)); - float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[5]->data + state_offset; + for (int64_t h = h_start; h < h_end; h++) { + int64_t h_offset = h * h_stride; + int64_t t_h_offset = t_offset + h_offset; + int64_t h_2d_offset = h * h_stride_2d; - for (int64_t h = h_start; h < h_end; h++) { - size_t h_offset = h * h_stride; - size_t t_h_offset = t_offset + h_offset; - size_t h_2d_offset = h * h_stride_2d; + for (int64_t ii = 0; ii < head_size; ii++) { + int64_t t_h_i_offset = t_h_offset + ii; + int64_t h_2d_i_offset = h_2d_offset + ii * h_stride; - for (int64_t i = 0; i < head_size; i++) { - size_t t_h_i_offset = t_h_offset + i; - size_t h_i_offset = h_offset + i; - size_t h_2d_i_offset = h_2d_offset + i * h_stride; + GGML_F32_VEC v_vec = GGML_F32_VEC_SET1(v[t_h_i_offset]); - float k_val = k[t_h_i_offset]; - float r_val = r[t_h_i_offset]; - float time_faaaa_val = time_faaaa[h_i_offset]; - float time_decay_val = time_decay[t_h_i_offset]; + float sa = 0; + { + GGML_F32_VEC sum[GGML_F32_ARR] = { GGML_F32_VEC_ZERO }; + GGML_F32_VEC ax[GGML_F32_ARR]; + GGML_F32_VEC ay[GGML_F32_ARR]; + for (int64_t j = 0; j < head_size; j += GGML_F32_STEP) { + for (int64_t kk = 0; kk < GGML_F32_ARR; kk++) { + ax[kk] = GGML_F32_VEC_LOAD(&a[t_h_offset + j + kk * GGML_F32_EPR]); + ay[kk] = GGML_F32_VEC_LOAD(&state_prev[h_2d_i_offset + j + kk * GGML_F32_EPR]); + sum[kk] = GGML_F32_VEC_FMA(sum[kk], ax[kk], ay[kk]); + } + } + GGML_F32_VEC_REDUCE(sa, sum); + } - // Broadcast scalar values to vectors - GGML_F32X k_vec = GGML_F32X_SET1(k_val); - GGML_F32X r_vec = GGML_F32X_SET1(r_val); - GGML_F32X time_faaaa_vec = GGML_F32X_SET1(time_faaaa_val); - GGML_F32X time_decay_vec = GGML_F32X_SET1(time_decay_val); + GGML_F32_VEC sa_vec = GGML_F32_VEC_SET1(sa); - for (int64_t j = 0; j < vec_count; j++) { - size_t base_j = j * wkv_vector_size; - size_t t_h_j_offset = t_h_offset + base_j; - size_t h_2d_i_j_offset = h_2d_i_offset + base_j; + int64_t j = 0; + GGML_F32_VEC result_vec[GGML_F32_ARR] = { GGML_F32_VEC_ZERO }; + for (; j < head_size; j += GGML_F32_STEP) { + for (int64_t kk = 0; kk < GGML_F32_ARR; kk++) { + int64_t t_h_j_offset = t_h_offset + j + kk * GGML_F32_EPR; + int64_t h_2d_i_j_offset = h_2d_i_offset + j + kk * GGML_F32_EPR; - // Load x elements at once - GGML_F32X v_vec = GGML_F32X_LOAD(&v[t_h_j_offset]); - GGML_F32X prev_state_vec = GGML_F32X_LOAD(&state_prev[h_2d_i_j_offset]); - GGML_F32X dst_vec = GGML_F32X_LOAD(&dst_data[t_h_j_offset]); + GGML_F32_VEC r_vec = GGML_F32_VEC_LOAD(&r[t_h_j_offset]); + GGML_F32_VEC w_vec = GGML_F32_VEC_LOAD(&w[t_h_j_offset]); + GGML_F32_VEC k_vec = GGML_F32_VEC_LOAD(&k[t_h_j_offset]); + GGML_F32_VEC b_vec = GGML_F32_VEC_LOAD(&b[t_h_j_offset]); - // Compute kv = v * k - GGML_F32X kv_vec = GGML_F32X_MUL(v_vec, k_vec); + k_vec = GGML_F32_VEC_MUL(v_vec, k_vec); - // Compute temp = kv * time_faaaa + prev_state - GGML_F32X temp_vec = GGML_F32X_FMA(prev_state_vec, kv_vec, time_faaaa_vec); + GGML_F32_VEC state_vec = GGML_F32_VEC_LOAD(&state_prev[h_2d_i_j_offset]); + // kv + s * decay + sa * b + state_vec = GGML_F32_VEC_FMA(k_vec, state_vec, w_vec); + state_vec = GGML_F32_VEC_FMA(state_vec, sa_vec, b_vec); + GGML_F32_VEC_STORE(&state_cur[h_2d_i_j_offset], state_vec); - // Update dst: dst += temp * r - dst_vec = GGML_F32X_FMA(dst_vec, temp_vec, r_vec); - GGML_F32X_STORE(&dst_data[t_h_j_offset], dst_vec); + result_vec[kk] = GGML_F32_VEC_FMA(result_vec[kk], state_vec, r_vec); + } + } + GGML_F32_VEC_REDUCE(dst_data[t_h_i_offset], result_vec); - // Update state: state = prev_state * time_decay + kv - GGML_F32X new_state_vec = GGML_F32X_FMA(kv_vec, prev_state_vec, time_decay_vec); - GGML_F32X_STORE(&state_cur[h_2d_i_j_offset], new_state_vec); - } + // There shouldn't be left-overs though. + for (; j < head_size; j++) { + int64_t t_h_j_offset = t_h_offset + j; + int64_t h_2d_i_j_offset = h_2d_i_offset + j; - // Handle remaining elements, this will not be used. - for (int64_t j = vec_count * wkv_vector_size; j < head_size; j++) { - size_t t_h_j_offset = t_h_offset + j; - size_t h_2d_i_j_offset = h_2d_i_offset + j; - float v_val = v[t_h_j_offset]; - float kv_val = v_val * k_val; - float prev_state_val = state_prev[h_2d_i_j_offset]; - float temp_val = kv_val * time_faaaa_val + prev_state_val; - dst_data[t_h_j_offset] += temp_val * r_val; - state_cur[h_2d_i_j_offset] = prev_state_val * time_decay_val + kv_val; + float r_val = r[t_h_j_offset]; + float w_val = w[t_h_j_offset]; + float k_val = k[t_h_j_offset]; + float b_val = b[t_h_j_offset]; + float kv_val = v[t_h_i_offset] * k_val; + + float prev_state_val = state_prev[h_2d_i_j_offset]; + state_cur[h_2d_i_j_offset] = prev_state_val * w_val + kv_val + sa * b_val; + dst_data[t_h_i_offset] += state_cur[h_2d_i_j_offset] * r_val; + } } } } - } - + #endif #else - // basically fused operations: - // dst = r @ (time_faaaa * (k @ v) + state), - // state = time_decay * state + (k @ v), - // recursive through each token for (int64_t t = 0; t < T; t++) { - size_t t_offset = t * t_stride; - size_t state_offset = head_size * C * (t / (T / n_seqs)); + int64_t t_offset = t * t_stride; + int64_t state_offset = head_size * C * (t / (T / n_seqs)); float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[5]->data + state_offset; + float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[6]->data + state_offset; for (int64_t h = h_start; h < h_end; h++) { - size_t h_offset = h * h_stride; - size_t t_h_offset = t_offset + h_offset; - size_t h_2d_offset = h * h_stride_2d; + int64_t h_offset = h * h_stride; + int64_t t_h_offset = t_offset + h_offset; + int64_t h_2d_offset = h * h_stride_2d; for (int64_t i = 0; i < head_size; i++) { - size_t t_h_i_offset = t_h_offset + i; - size_t h_i_offset = h_offset + i; - size_t h_2d_i_offset = h_2d_offset + i * h_stride; + int64_t t_h_i_offset = t_h_offset + i; + int64_t h_2d_i_offset = h_2d_offset + i * h_stride; - float k_val = k[t_h_i_offset]; - float r_val = r[t_h_i_offset]; - float time_faaaa_val = time_faaaa[h_i_offset]; - // RWKV v6: different time_decay for each token. - float time_decay_val = time_decay[t_h_i_offset]; + float v_val = v[t_h_i_offset]; + float sa = 0, result = 0; for (int64_t j = 0; j < head_size; j++) { - size_t t_h_j_offset = t_h_offset + j; - size_t h_2d_i_j_offset = h_2d_i_offset + j; + sa += a[t_h_offset + j] * state_prev[h_2d_i_offset + j]; + } - float v_val = v[t_h_j_offset]; + for (int64_t j = 0; j < head_size; j++) { + int64_t t_h_j_offset = t_h_offset + j; + int64_t h_2d_i_j_offset = h_2d_i_offset + j; + + float r_val = r[t_h_j_offset]; + float w_val = w[t_h_j_offset]; + float k_val = k[t_h_j_offset]; + float b_val = b[t_h_j_offset]; float kv_val = v_val * k_val; float prev_state_val = state_prev[h_2d_i_j_offset]; - float temp_val = kv_val * time_faaaa_val + prev_state_val; - dst_data[t_h_j_offset] += temp_val * r_val; - state_cur[h_2d_i_j_offset] = prev_state_val * time_decay_val + kv_val; + state_cur[h_2d_i_j_offset] = prev_state_val * w_val + kv_val + sa * b_val; + result += state_cur[h_2d_i_j_offset] * r_val; } + dst_data[t_h_i_offset] = result; } } } @@ -10137,7 +12887,7 @@ static void ggml_compute_forward_rwkv_wkv6_f32( } -void ggml_compute_forward_rwkv_wkv6( +void ggml_compute_forward_rwkv_wkv7( const ggml_compute_params * params, ggml_tensor * dst) { @@ -10146,7 +12896,7 @@ void ggml_compute_forward_rwkv_wkv6( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_rwkv_wkv6_f32(params, dst); + ggml_compute_forward_rwkv_wkv7_f32(params, dst); } break; default: { @@ -10155,190 +12905,142 @@ void ggml_compute_forward_rwkv_wkv6( } } -// ggml_compute_forward_gla +// ggml_compute_forward_map_custom1 -static void ggml_compute_forward_gla_f32( +void ggml_compute_forward_map_custom1( const ggml_compute_params * params, - ggml_tensor * dst) { - const int64_t T = dst->src[1]->ne[2]; - const int64_t C = dst->ne[0]; - const int64_t HEADS = dst->src[1]->ne[1]; - const int64_t n_seqs = dst->src[4]->ne[1]; - const int64_t head_size = C / HEADS; - const float scale = ggml_get_op_params_f32(dst, 0); + ggml_tensor * dst) { - float * dst_data = (float *) dst->data; - float * state = ((float *) dst->data) + C * T; + const ggml_tensor * a = dst->src[0]; - const int ith = params->ith; - const int nth = params->nth; + struct ggml_map_custom1_op_params p; + memcpy(&p, dst->op_params, sizeof(p)); - const int h_start = (HEADS * (ith )) / nth; - const int h_end = ((HEADS * (ith + 1)) / nth < HEADS) ? - (HEADS * (ith + 1)) / nth : HEADS; + p.fun(dst, a, params->ith, params->nth, p.userdata); +} - float * k = (float *) dst->src[0]->data; - float * v = (float *) dst->src[1]->data; - float * q = (float *) dst->src[2]->data; - float * g = (float *) dst->src[3]->data; +// ggml_compute_forward_map_custom2 - size_t t_stride = HEADS * head_size; // Same to C +void ggml_compute_forward_map_custom2( + const ggml_compute_params * params, + ggml_tensor * dst) { - size_t h_stride = C / HEADS; - GGML_ASSERT(C % HEADS == 0); // C must be divisible by HEADS - size_t h_stride_2d = head_size * head_size; + const ggml_tensor * a = dst->src[0]; + const ggml_tensor * b = dst->src[1]; - if (ith == 0) { - memset(dst_data, 0, T * C * sizeof(float)); - } - ggml_barrier(params->threadpool); + struct ggml_map_custom2_op_params p; + memcpy(&p, dst->op_params, sizeof(p)); + p.fun(dst, a, b, params->ith, params->nth, p.userdata); +} - #if defined(__AVX__) && !defined(__AVX512F__) - #define GGML_F32X GGML_F32x8 - #define GGML_F32X_SET1 GGML_F32x8_SET1 - #define GGML_F32X_LOAD GGML_F32x8_LOAD - #define GGML_F32X_STORE GGML_F32x8_STORE - #define GGML_F32X_MUL GGML_F32x8_MUL - #define GGML_F32X_FMA GGML_F32x8_FMA - #define GLA_VECTOR_SIZE 8 - #elif defined(__AVX512F__) - #define GGML_F32X GGML_F32x16 - #define GGML_F32X_SET1 GGML_F32x16_SET1 - #define GGML_F32X_LOAD GGML_F32x16_LOAD - #define GGML_F32X_STORE GGML_F32x16_STORE - #define GGML_F32X_MUL GGML_F32x16_MUL - #define GGML_F32X_FMA GGML_F32x16_FMA - #define GLA_VECTOR_SIZE 16 - #elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) - #define GGML_F32X GGML_F32xt - #define GGML_F32X_SET1 GGML_F32xt_SET1 - #define GGML_F32X_LOAD GGML_F32xt_LOAD - #define GGML_F32X_STORE GGML_F32xt_STORE - #define GGML_F32X_MUL GGML_F32xt_MUL - #define GGML_F32X_FMA GGML_F32xt_FMA - #define GLA_VECTOR_SIZE 8 - #elif defined(__ARM_NEON) && defined(__aarch64__) - #define GGML_F32X GGML_F32x4 - #define GGML_F32X_SET1 GGML_F32x4_SET1 - #define GGML_F32X_LOAD GGML_F32x4_LOAD - #define GGML_F32X_STORE GGML_F32x4_STORE - #define GGML_F32X_MUL GGML_F32x4_MUL - #define GGML_F32X_FMA GGML_F32x4_FMA - #define GLA_VECTOR_SIZE 4 - #endif +// ggml_compute_forward_map_custom3 - #ifdef GLA_VECTOR_SIZE - int gla_vector_size; - #if defined(__ARM_FEATURE_SVE) - gla_vector_size = svcntw(); - #else - gla_vector_size = GLA_VECTOR_SIZE; - #endif - const int64_t vec_count = head_size / gla_vector_size; +void ggml_compute_forward_map_custom3( + const ggml_compute_params * params, + ggml_tensor * dst) { - for (int64_t t = 0; t < T; t++) { - size_t t_offset = t * t_stride; - size_t state_offset = head_size * C * (t / (T / n_seqs)); - float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[4]->data + state_offset; + const ggml_tensor * a = dst->src[0]; + const ggml_tensor * b = dst->src[1]; + const ggml_tensor * c = dst->src[2]; - for (int64_t h = h_start; h < h_end; h++) { - size_t h_offset = h * h_stride; - size_t t_h_offset = t_offset + h_offset; - size_t h_2d_offset = h * h_stride_2d; + struct ggml_map_custom3_op_params p; + memcpy(&p, dst->op_params, sizeof(p)); - for (int64_t i = 0; i < head_size; i++) { - size_t t_h_i_offset = t_h_offset + i; - size_t h_2d_i_offset = h_2d_offset + i * h_stride; + p.fun(dst, a, b, c, params->ith, params->nth, p.userdata); +} - float k_val = k[t_h_i_offset]; - float q_val = q[t_h_i_offset] * scale; - float g_val = g[t_h_i_offset]; +// ggml_compute_forward_custom - // Broadcast scalar values to vectors - GGML_F32X k_vec = GGML_F32X_SET1(k_val); - GGML_F32X q_vec = GGML_F32X_SET1(q_val); - GGML_F32X g_vec = GGML_F32X_SET1(g_val); +void ggml_compute_forward_custom( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { - for (int64_t j = 0; j < vec_count; j++) { - size_t base_j = j * gla_vector_size; - size_t t_h_j_offset = t_h_offset + base_j; - size_t h_2d_i_j_offset = h_2d_i_offset + base_j; + struct ggml_custom_op_params p; + memcpy(&p, dst->op_params, sizeof(p)); - // Load x elements at once - GGML_F32X v_vec = GGML_F32X_LOAD(&v[t_h_j_offset]); - GGML_F32X prev_state_vec = GGML_F32X_LOAD(&state_prev[h_2d_i_j_offset]); - GGML_F32X dst_vec = GGML_F32X_LOAD(&dst_data[t_h_j_offset]); + p.fun(dst, params->ith, params->nth, p.userdata); +} - // Compute kv = v * k - GGML_F32X kv_vec = GGML_F32X_MUL(v_vec, k_vec); +// ggml_compute_forward_cross_entropy_loss - // Compute temp = prev_state * g + kv - GGML_F32X temp_vec = GGML_F32X_FMA(kv_vec, prev_state_vec, g_vec); +static void ggml_compute_forward_cross_entropy_loss_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - // Update dst: dst += temp * q - dst_vec = GGML_F32X_FMA(dst_vec, temp_vec, q_vec); - GGML_F32X_STORE(&dst_data[t_h_j_offset], dst_vec); + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - // Update state - GGML_F32X_STORE(&state_cur[h_2d_i_j_offset], temp_vec); - } + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(src0->nb[0] == ggml_type_size(src0->type)); + GGML_ASSERT(src1->nb[0] == ggml_type_size(src1->type)); + GGML_ASSERT(ggml_are_same_shape(src0, src1)); + GGML_ASSERT(ggml_is_scalar(dst)); + GGML_ASSERT(dst->type == GGML_TYPE_F32); - // Handle remaining elements, this will not be used. - for (int64_t j = vec_count * gla_vector_size; j < head_size; j++) { - size_t t_h_j_offset = t_h_offset + j; - size_t h_2d_i_j_offset = h_2d_i_offset + j; - float v_val = v[t_h_j_offset]; - float kv_val = v_val * k_val; - float prev_state_val = state_prev[h_2d_i_j_offset]; - float temp_val = kv_val + prev_state_val * g_val; - dst_data[t_h_j_offset] += temp_val * q_val; - state_cur[h_2d_i_j_offset] = temp_val; - } - } - } - } + // TODO: handle transposed/permuted matrices + const int64_t nc = src0->ne[0]; + const int64_t nr = ggml_nrows(src0); - #else - for (int64_t t = 0; t < T; t++) { - size_t t_offset = t * t_stride; - size_t state_offset = head_size * C * (t / (T / n_seqs)); - float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[4]->data + state_offset; + const int ith = params->ith; + const int nth = params->nth; - for (int64_t h = h_start; h < h_end; h++) { - size_t h_offset = h * h_stride; - size_t t_h_offset = t_offset + h_offset; - size_t h_2d_offset = h * h_stride_2d; + float * sums = (float *) params->wdata; + float * st = ((float *) params->wdata) + nth + ith*nc; + float sum_thread = 0.0f; - for (int64_t i = 0; i < head_size; i++) { - size_t t_h_i_offset = t_h_offset + i; - size_t h_2d_i_offset = h_2d_offset + i * h_stride; + GGML_ASSERT(params->wsize >= sizeof(float) * (nth + nth * nc)); - float k_val = k[t_h_i_offset]; - float q_val = q[t_h_i_offset] * scale; - float g_val = g[t_h_i_offset]; + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; - for (int64_t j = 0; j < head_size; j++) { - size_t t_h_j_offset = t_h_offset + j; - size_t h_2d_i_j_offset = h_2d_i_offset + j; + // row range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); + + for (int64_t i1 = ir0; i1 < ir1; ++i1) { + const float * s0 = (const float *)((const char *) src0->data + i1*src0->nb[1]); + const float * s1 = (const float *)((const char *) src1->data + i1*src1->nb[1]); + +#ifndef NDEBUG + for (int64_t i = 0; i < nc; ++i) { + //printf("p[%d] = %f\n", i, p[i]); + assert(!isnan(s0[i])); + assert(!isnan(s1[i])); + } +#endif // NDEBUG + + float max = -INFINITY; + ggml_vec_max_f32(nc, &max, s0); + const ggml_float sum_softmax = ggml_vec_log_soft_max_f32(nc, st, s0, max); + assert(sum_softmax >= 0.0); + + ggml_vec_add1_f32(nc, st, st, -sum_softmax); + ggml_vec_mul_f32(nc, st, st, s1); + + float sum_st = 0.0f; + ggml_vec_sum_f32(nc, &sum_st, st); + sum_thread += sum_st; - float v_val = v[t_h_j_offset]; - float kv_val = v_val * k_val; - float prev_state_val = state_prev[h_2d_i_j_offset]; - float temp_val = prev_state_val * g_val + kv_val; - dst_data[t_h_j_offset] += temp_val * q_val; - state_cur[h_2d_i_j_offset] = temp_val; - } - } - } +#ifndef NDEBUG + for (int64_t i = 0; i < nc; ++i) { + assert(!isnan(st[i])); + assert(!isinf(st[i])); } - #endif -} +#endif // NDEBUG + } + sums[ith] = sum_thread; + ggml_barrier(params->threadpool); + if (ith == 0) { + float * dp = (float *) dst->data; + ggml_vec_sum_f32(nth, dp, sums); + dp[0] *= -1.0f / (float) nr; + } +} -void ggml_compute_forward_gla( +void ggml_compute_forward_cross_entropy_loss( const ggml_compute_params * params, ggml_tensor * dst) { @@ -10347,7 +13049,7 @@ void ggml_compute_forward_gla( switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_gla_f32(params, dst); + ggml_compute_forward_cross_entropy_loss_f32(params, dst); } break; default: { @@ -10356,859 +13058,1391 @@ void ggml_compute_forward_gla( } } -static void ggml_compute_forward_solve_tri_f32(const struct ggml_compute_params * params, struct ggml_tensor * dst) { - const struct ggml_tensor * src0 = dst->src[0]; // A (lower triangular) - const struct ggml_tensor * src1 = dst->src[1]; // B (RHS) - - GGML_TENSOR_BINARY_OP_LOCALS; +// ggml_compute_forward_cross_entropy_loss_back - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); +static void ggml_compute_forward_cross_entropy_loss_back_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - GGML_ASSERT(ne00 == ne01); // A must be square - GGML_ASSERT(ne0 == ne10); // solution cols == B cols - GGML_ASSERT(ne1 == ne11); // solution rows == B rows + const ggml_tensor * grad = dst->src[0]; // gradient of forward pass output + const ggml_tensor * src0f = dst->src[1]; // src0 of forward pass + const ggml_tensor * src1f = dst->src[2]; // src1 of forward pass - GGML_ASSERT(ne02 == ne12 && ne12 == ne2); - GGML_ASSERT(ne03 == ne13 && ne13 == ne3); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_is_contiguous(src0f)); + GGML_ASSERT(ggml_is_contiguous(src1f)); + GGML_ASSERT(ggml_is_contiguous(grad)); + GGML_ASSERT(ggml_are_same_shape(src0f, src1f) && ggml_are_same_shape(src0f, dst)); - const int ith = params->ith; - const int nth = params->nth; + const int64_t ith = params->ith; + const int64_t nth = params->nth; - const int64_t k = ne10; // number of RHS columns - const int64_t n = ne11; // A is n×n - const int64_t nr = ne02 * ne03 * k; // we're parallelizing on columns here, so seq x token x column will be the unit + // TODO: handle transposed/permuted matrices + const int64_t nc = src0f->ne[0]; + const int64_t nr = ggml_nrows(src0f); - // chunks per thread + // rows per thread const int64_t dr = (nr + nth - 1)/nth; - // chunk range for this thread + // row range for this thread const int64_t ir0 = dr*ith; const int64_t ir1 = MIN(ir0 + dr, nr); - const float * A = (const float *) src0->data; // [n, n, B1, B2] - const float * B = (const float *) src1->data; // [n, k, B1, B2] - float * X = ( float *) dst->data; // [n, k, B1, B2] - - for (int64_t ir = ir0; ir < ir1; ++ir) { - const int64_t i03 = ir/(ne02*k); - const int64_t i02 = (ir - i03*ne02*k)/k; - const int64_t i01 = (ir - i03*ne02*k - i02*k); + const float d_by_nr = ((const float *) grad->data)[0] / (float) nr; - const float * A_batch = A + i02 * nb02 / sizeof(float) + i03 * nb03 / sizeof(float); - const float * B_batch = B + i02 * nb12 / sizeof(float) + i03 * nb13 / sizeof(float); + for (int64_t i1 = ir0; i1 < ir1; i1++) { + float * ds0 = (float *)((char *) dst->data + i1*dst->nb[1]); + const float * s0 = (const float *)((const char *) src0f->data + i1*src0f->nb[1]); + const float * s1 = (const float *)((const char *) src1f->data + i1*src1f->nb[1]); - float * X_batch = X + i02 * nb2 / sizeof(float) + i03 * nb3 / sizeof(float); +#ifndef NDEBUG + for (int64_t i = 0; i < nc; ++i) { + //printf("p[%d] = %f\n", i, p[i]); + assert(!isnan(s0[i])); + assert(!isnan(s1[i])); + } +#endif // NDEBUG - for (int64_t i00 = 0; i00 < n; ++i00) { - float sum = 0.0f; - for (int64_t t = 0; t < i00; ++t) { - sum += A_batch[i00 * n + t] * X_batch[t * k + i01]; - } + // soft_max + float max = -INFINITY; + ggml_vec_max_f32(nc, &max, s0); + const ggml_float sum = ggml_vec_soft_max_f32(nc, ds0, s0, max); + assert(sum > 0.0); + ggml_vec_scale_f32(nc, ds0, 1.0/sum); - const float diag = A_batch[i00 * n + i00]; - assert(diag != 0.0f && "Zero diagonal in triangular matrix"); + // grad(src0f) = (softmax(src0f) - src1f) * grad(cross_entropy_loss(src0f, src1f)) / nr + ggml_vec_sub_f32(nc, ds0, ds0, s1); + ggml_vec_scale_f32(nc, ds0, d_by_nr); - X_batch[i00 * k + i01] = (B_batch[i00 * k + i01] - sum) / diag; +#ifndef NDEBUG + for (int64_t i = 0; i < nc; ++i) { + assert(!isnan(ds0[i])); + assert(!isinf(ds0[i])); } +#endif // NDEBUG } } -void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst) { +void ggml_compute_forward_cross_entropy_loss_back( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32) { - ggml_compute_forward_solve_tri_f32(params, dst); - } else { - GGML_ABORT("fatal error"); + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_cross_entropy_loss_back_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } } } -// ggml_compute_forward_gated_delta_net -static void ggml_compute_forward_gated_delta_net_one_chunk( - const ggml_compute_params * params, - ggml_tensor * dst, - int64_t ir0, - int64_t ir1) { - - ggml_tensor * src_q = dst->src[0]; - ggml_tensor * src_k = dst->src[1]; - ggml_tensor * src_v = dst->src[2]; - ggml_tensor * src_g = dst->src[3]; - ggml_tensor * src_beta = dst->src[4]; - ggml_tensor * src_state = dst->src[5]; - - const int64_t S_v = src_v->ne[0]; - const int64_t H = src_v->ne[1]; - const int64_t n_tokens = src_v->ne[2]; - const int64_t n_seqs = src_v->ne[3]; - - GGML_ASSERT(ggml_is_contiguous_rows(src_q)); - GGML_ASSERT(ggml_is_contiguous_rows(src_k)); - GGML_ASSERT(ggml_is_contiguous_rows(src_v)); - GGML_ASSERT(ggml_is_contiguous(src_g)); - GGML_ASSERT(ggml_is_contiguous(src_beta)); - GGML_ASSERT(ggml_is_contiguous(src_state)); - - GGML_ASSERT(src_g->ne[0] == 1 || src_g->ne[0] == S_v); - GGML_ASSERT(src_beta->ne[0] == 1); +static void ggml_compute_forward_opt_step_adamw_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { - GGML_TENSOR_LOCALS(int64_t, neq, src_q, ne); - GGML_TENSOR_LOCALS(size_t, nbq, src_q, nb); - GGML_TENSOR_LOCALS(int64_t, nek, src_k, ne); - GGML_TENSOR_LOCALS(size_t, nbk, src_k, nb); - GGML_TENSOR_LOCALS(int64_t, nev, src_v, ne); - GGML_TENSOR_LOCALS(size_t, nbv, src_v, nb); - GGML_TENSOR_LOCALS(int64_t, neg, src_g, ne); - GGML_TENSOR_LOCALS(size_t, nbg, src_g, nb); - GGML_TENSOR_LOCALS(size_t, nbb, src_beta, nb); + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src0_grad = dst->src[1]; + const ggml_tensor * src0_grad_m = dst->src[2]; + const ggml_tensor * src0_grad_v = dst->src[3]; + const ggml_tensor * adamw_params = dst->src[4]; - const bool kda = (neg0 == S_v); + GGML_ASSERT(ggml_are_same_shape(src0, src0_grad)); + GGML_ASSERT(ggml_are_same_shape(src0, src0_grad_m)); + GGML_ASSERT(ggml_are_same_shape(src0, src0_grad_v)); + GGML_ASSERT(ggml_nelements(adamw_params) == 7); - // scratch layout per thread: [delta(S_v)] - const int64_t scratch_per_thread = S_v; const int ith = params->ith; + const int nth = params->nth; - float * delta = (float *)params->wdata + ith * scratch_per_thread + CACHE_LINE_SIZE_F32; - - // output layout: [attn_scores | new_states] - // attn_scores: S_v * H * n_tokens * n_seqs floats - // new_states: S_v * S_v * H * n_seqs floats - const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs; - float * attn_out_base = (float *)dst->data; - float * state_out_base = (float *)dst->data + attn_score_elems; - - const float * state_in_base = (const float *)src_state->data; - - //const int64_t rq1 = nev1 / neq1; - //const int64_t rk1 = nev1 / nek1; - const int64_t rq3 = nev3 / neq3; - const int64_t rk3 = nev3 / nek3; - - const float scale = 1.0f / sqrtf((float) S_v); - - for (int64_t ir = ir0; ir < ir1; ++ir) { - const int64_t iv1 = ir % H; // head_index - const int64_t iv3 = ir / H; // sequence - - const int64_t iq1 = iv1 % neq1; - const int64_t ik1 = iv1 % nek1; - - const int64_t iq3 = iv3 / rq3; - const int64_t ik3 = iv3 / rk3; - - float * s_out = state_out_base + (iv3 * H + iv1) * S_v * S_v; - - // copy input state into output buffer and operate in-place - const float * s_in = state_in_base + (iv3 * H + iv1) * S_v * S_v; - memcpy(s_out, s_in, S_v * S_v * sizeof(float)); + const int nr = ggml_nrows(src0); - // attn output pointer for first token of this (head, seq) - float * attn_data = attn_out_base + (iv3 * n_tokens * H + iv1) * S_v; + GGML_TENSOR_UNARY_OP_LOCALS + GGML_ASSERT(nb00 == sizeof(float)); - for (int64_t t = 0; t < n_tokens; t++) { - const float * q_d = (const float *)((const char *)src_q->data + iq3 * nbq3 + t * nbq2 + iq1 * nbq1); - const float * k_d = (const float *)((const char *)src_k->data + ik3 * nbk3 + t * nbk2 + ik1 * nbk1); - const float * v_d = (const float *)((const char *)src_v->data + iv3 * nbv3 + t * nbv2 + iv1 * nbv1); + // rows per thread + const int dr = (nr + nth - 1)/nth; - const float beta_val = *(const float *)((const char *)src_beta->data + iv3 * nbb3 + t * nbb2 + iv1 * nbb1); - const float * g_d = (const float *)((const char *)src_g->data + iv3 * nbg3 + t * nbg2 + iv1 * nbg1); + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); - // state is stored transposed: s_out[j*S_v + i] = S[i][j] - // so row j of s_out = column j of S (contiguous access) + const float * adamw_params_ptr = ggml_get_data_f32(adamw_params); - if (kda) { - // precompute exp(g) into delta scratch (reused below) - for (int64_t i = 0; i < S_v; ++i) { - delta[i] = expf(g_d[i]); - } - // S[i][:] *= exp(g[i]) => for each row j of M: M[j][i] *= exp(g[i]) - for (int64_t j = 0; j < S_v; ++j) { - ggml_vec_mul_f32(S_v, &s_out[j * S_v], &s_out[j * S_v], delta); - } - } else { - ggml_vec_scale_f32(S_v * S_v, s_out, expf(g_d[0])); - } + const float alpha = adamw_params_ptr[0]; + const float beta1 = adamw_params_ptr[1]; + const float beta2 = adamw_params_ptr[2]; + const float eps = adamw_params_ptr[3]; + const float wd = adamw_params_ptr[4]; + const float beta1h = adamw_params_ptr[5]; + const float beta2h = adamw_params_ptr[6]; + const float keep = 1.f - alpha * wd; + for (int ir = ir0; ir < ir1; ++ir) { + const int64_t i03 = ir/(ne02*ne01); + const int64_t i02 = (ir - i03*ne02*ne01)/ne01; + const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); - // delta[j] = sum_i S[i][j] * k[i] = dot(row j of M, k) - for (int64_t j = 0; j < S_v; ++j) { - float sum = 0.0f; - ggml_vec_dot_f32(S_v, &sum, 0, &s_out[j * S_v], 0, k_d, 0, 1); - delta[j] = (v_d[j] - sum) * beta_val; - } + const size_t offset = i03*nb03 + i02*nb02 + i01*nb01; - // outer product: S[i][j] += k[i] * delta[j] => M[j][i] += delta[j] * k[i] - for (int64_t j = 0; j < S_v; ++j) { - ggml_vec_mad_f32(S_v, &s_out[j * S_v], k_d, delta[j]); - } + float * w = (float *) ((char *) src0->data + offset); // weight + const float * g = (const float *) ((const char *) src0_grad->data + offset); // grad + float * m = (float *) ((char *) src0_grad_m->data + offset); + float * v = (float *) ((char *) src0_grad_v->data + offset); - // attn_out[j] = sum_i S[i][j] * q[i] = dot(row j of M, q) - for (int64_t j = 0; j < S_v; ++j) { - float sum = 0.0f; - ggml_vec_dot_f32(S_v, &sum, 0, &s_out[j * S_v], 0, q_d, 0, 1); - attn_data[j] = sum * scale; - } + for (int i00 = 0; i00 < ne00; ++i00) { + m[i00] = m[i00]*beta1 + g[i00]*(1.0f - beta1); + v[i00] = v[i00]*beta2 + g[i00]*g[i00]*(1.0f - beta2); - attn_data += S_v * H; // advance to next token + const float mh = m[i00]*beta1h; + const float vh = sqrtf(v[i00]*beta2h) + eps; + + // The weight decay is applied independently of the Adam momenta m and v. + // This is NOT equivalent to l2 regularization that adds w[i00]*w[i00] to the loss. + // See: https://arxiv.org/pdf/1711.05101v3.pdf + w[i00] = w[i00] * keep - alpha * mh / vh; } } } - -static void ggml_compute_forward_gated_delta_net_f32( +void ggml_compute_forward_opt_step_adamw( const ggml_compute_params * params, ggml_tensor * dst) { - ggml_tensor * V = dst->src[2]; - int64_t nr = V->ne[1] * V->ne[3]; + const ggml_tensor * src0 = dst->src[0]; - // disable for NUMA - const bool disable_chunking = ggml_is_numa(); + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_opt_step_adamw_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} - int nth = params->nth; - int ith = params->ith; +static void ggml_compute_forward_opt_step_sgd_f32(const ggml_compute_params * params, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src0_grad = dst->src[1]; + const ggml_tensor * sgd_params = dst->src[2]; - // 4x chunks per thread - int nth_scaled = nth * 4; - int64_t chunk_size = (nr + nth_scaled - 1) / nth_scaled; - int64_t nchunk = (nr + chunk_size - 1) / chunk_size; + GGML_ASSERT(ggml_are_same_shape(src0, src0_grad)); + GGML_ASSERT(ggml_nelements(sgd_params) == 2); - if (nth == 1 || nchunk < nth || disable_chunking) { - nchunk = nth; - } + const int ith = params->ith; + const int nth = params->nth; - if (ith == 0) { - ggml_threadpool_chunk_set(params->threadpool, nth); - } + const int nr = ggml_nrows(src0); - ggml_barrier(params->threadpool); + GGML_TENSOR_UNARY_OP_LOCALS + GGML_ASSERT(nb00 == sizeof(float)); - const int64_t dr = (nr + nchunk - 1) / nchunk; + // rows per thread + const int dr = (nr + nth - 1) / nth; - int current_chunk = ith; + // row range for this thread + const int ir0 = dr * ith; + const int ir1 = MIN(ir0 + dr, nr); - while (current_chunk < nchunk) { - const int64_t ir0 = dr * current_chunk; - const int64_t ir1 = MIN(ir0 + dr, nr); + // using adamw param subset we care about - alpha, wd - could have a separate struct + const float * sgd_params_ptr = ggml_get_data_f32(sgd_params); + const float alpha = sgd_params_ptr[0]; + const float keep = 1.f - alpha * sgd_params_ptr[1]; - ggml_compute_forward_gated_delta_net_one_chunk(params, dst, ir0, ir1); - current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1); + for (int ir = ir0; ir < ir1; ++ir) { + const int64_t i03 = ir / (ne02 * ne01); + const int64_t i02 = (ir - i03 * ne02 * ne01) / ne01; + const int64_t i01 = (ir - i03 * ne02 * ne01 - i02 * ne01); + + const size_t offset = i03 * nb03 + i02 * nb02 + i01 * nb01; + + float * w = (float *) ((char *) src0->data + offset); // weight + const float * g = (const float *) ((const char *) src0_grad->data + offset); // grad + + for (int i00 = 0; i00 < ne00; ++i00) { + w[i00] = w[i00] * keep - alpha * g[i00]; + } } } -void ggml_compute_forward_gated_delta_net( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_opt_step_sgd(const ggml_compute_params * params, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; switch (src0->type) { case GGML_TYPE_F32: { - ggml_compute_forward_gated_delta_net_f32(params, dst); - } break; + ggml_compute_forward_opt_step_sgd_f32(params, dst); + } + break; default: { - GGML_ABORT("fatal error"); + GGML_ABORT("fatal error - sgd is F32 only"); } } } -// ggml_compute_forward_rwkv_wkv7 -static void ggml_compute_forward_rwkv_wkv7_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { - const int64_t T = dst->src[1]->ne[2]; - const int64_t C = dst->ne[0]; - const int64_t HEADS = dst->src[1]->ne[1]; - const int64_t n_seqs = dst->src[6]->ne[1]; - const int64_t head_size = C / HEADS; +// ---- binary-ops implementation ---- - float * dst_data = (float *) dst->data; - float * state = ((float *) dst->data) + C * T; - const int ith = params->ith; - const int nth = params->nth; +extern "C" { +void ggml_compute_forward_add_non_quantized(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sub(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_mul(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_div(const struct ggml_compute_params * params, struct ggml_tensor * dst); +} - const int h_start = (HEADS * (ith )) / nth; - const int h_end = ((HEADS * (ith + 1)) / nth < HEADS) ? - (HEADS * (ith + 1)) / nth : HEADS; +#if defined(GGML_USE_ACCELERATE) +#include - float * r = (float *) dst->src[0]->data; - float * w = (float *) dst->src[1]->data; - float * k = (float *) dst->src[2]->data; - float * v = (float *) dst->src[3]->data; - float * a = (float *) dst->src[4]->data; - float * b = (float *) dst->src[5]->data; +using vDSP_fn_t = void (*)(const float *, vDSP_Stride, const float *, vDSP_Stride, float *, vDSP_Stride, vDSP_Length); +#endif - int64_t t_stride = HEADS * head_size; // Same to C +static inline float op_add(float a, float b) { + return a + b; +} - int64_t h_stride = C / HEADS; - GGML_ASSERT(C % HEADS == 0); // C must be divisible by HEADS - int64_t h_stride_2d = head_size * head_size; +static inline float op_sub(float a, float b) { + return a - b; +} - #if defined(GGML_SIMD) - #if defined(__ARM_FEATURE_SVE) || defined(__riscv_v_intrinsic) - // scalar Route to scalar implementation //TODO: Write SVE code and RVV code - for (int64_t t = 0; t < T; t++) { - int64_t t_offset = t * t_stride; - int64_t state_offset = head_size * C * (t / (T / n_seqs)); - float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[6]->data + state_offset; +static inline float op_mul(float a, float b) { + return a * b; +} - for (int64_t h = h_start; h < h_end; h++) { - int64_t h_offset = h * h_stride; - int64_t t_h_offset = t_offset + h_offset; - int64_t h_2d_offset = h * h_stride_2d; +static inline float op_div(float a, float b) { + return a / b; +} - for (int64_t i = 0; i < head_size; i++) { - int64_t t_h_i_offset = t_h_offset + i; - int64_t h_2d_i_offset = h_2d_offset + i * h_stride; +template +static inline void vec_binary_op_contiguous(const int64_t n, dst_t * z, const src0_t * x, const src1_t * y) { + constexpr auto src0_to_f32 = type_conversion_table::to_f32; + constexpr auto src1_to_f32 = type_conversion_table::to_f32; + constexpr auto f32_to_dst = type_conversion_table::from_f32; - float v_val = v[t_h_i_offset]; + for (int i = 0; i < n; i++) { + z[i] = f32_to_dst(op(src0_to_f32(x[i]), src1_to_f32(y[i]))); + } +} - float sa = 0, result = 0; - for (int64_t j = 0; j < head_size; j++) { - sa += a[t_h_offset + j] * state_prev[h_2d_i_offset + j]; - } +template +static inline void vec_binary_op_non_contiguous(const int64_t n, const int64_t ne10, const int64_t nb10, dst_t * z, const src0_t * x, const src1_t * y) { + constexpr auto src0_to_f32 = type_conversion_table::to_f32; + constexpr auto src1_to_f32 = type_conversion_table::to_f32; + constexpr auto f32_to_dst = type_conversion_table::from_f32; - for (int64_t j = 0; j < head_size; j++) { - int64_t t_h_j_offset = t_h_offset + j; - int64_t h_2d_i_j_offset = h_2d_i_offset + j; + for (int i = 0; i < n; i++) { + int i10 = i % ne10; + const src1_t * y_ptr = (const src1_t *)((const char *)y + i10*nb10); + z[i] = f32_to_dst(op(src0_to_f32(x[i]), src1_to_f32(*y_ptr))); + } +} - float r_val = r[t_h_j_offset]; - float w_val = w[t_h_j_offset]; - float k_val = k[t_h_j_offset]; - float b_val = b[t_h_j_offset]; - float kv_val = v_val * k_val; - float prev_state_val = state_prev[h_2d_i_j_offset]; - state_cur[h_2d_i_j_offset] = prev_state_val * w_val + kv_val + sa * b_val; - result += state_cur[h_2d_i_j_offset] * r_val; - } - dst_data[t_h_i_offset] = result; +template +static void apply_binary_op(const ggml_compute_params * params, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(ggml_can_repeat(src1, src0) && ggml_are_same_shape(src0, dst)); + + GGML_TENSOR_BINARY_OP_LOCALS + + GGML_ASSERT( nb0 == sizeof(dst_t)); + GGML_ASSERT(nb00 == sizeof(src0_t)); + + const auto [ir0, ir1] = get_thread_range(params, src0); + const bool is_src1_contiguous_rows = ggml_is_contiguous_rows(src1); + +#ifdef GGML_USE_ACCELERATE + vDSP_fn_t vDSP_op = nullptr; + // TODO - avoid the f32-only check using type 'trait' lookup tables and row-based src-to-float conversion functions + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + if (op == op_add) { + vDSP_op = vDSP_vadd; + } else if (op == op_sub) { + vDSP_op = vDSP_vsub; + } else if (op == op_mul) { + vDSP_op = vDSP_vmul; + } else if (op == op_div) { + vDSP_op = vDSP_vdiv; + } + } +#endif + + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t i03 = ir/(ne02*ne01); + const int64_t i02 = (ir - i03*ne02*ne01)/ne01; + const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); + + const int64_t i13 = i03 % ne13; + const int64_t i12 = i02 % ne12; + const int64_t i11 = i01 % ne11; + + dst_t * dst_ptr = (dst_t *) ((char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1 ); + const src0_t * src0_ptr = (const src0_t *) ((const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01); + const src1_t * src1_ptr = (const src1_t *) ((const char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11); + + if (is_src1_contiguous_rows) { + // src1 is broadcastable across src0 and dst in i1, i2, i3 + const int64_t nr0 = ne00 / ne10; + + for (int64_t r = 0; r < nr0; ++r) { +#ifdef GGML_USE_ACCELERATE + if constexpr (std::is_same_v && std::is_same_v && std::is_same_v) { + if (vDSP_op != nullptr) { + vDSP_op(src1_ptr, 1, src0_ptr + r*ne10, 1, dst_ptr + r*ne10, 1, ne10); + continue; } } +#endif + vec_binary_op_contiguous(ne10, dst_ptr + r*ne10, src0_ptr + r*ne10, src1_ptr); } - #else - for (int64_t t = 0; t < T; t++) { - int64_t t_offset = t * t_stride; - int64_t state_offset = head_size * C * (t / (T / n_seqs)); - float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[6]->data + state_offset; - - for (int64_t h = h_start; h < h_end; h++) { - int64_t h_offset = h * h_stride; - int64_t t_h_offset = t_offset + h_offset; - int64_t h_2d_offset = h * h_stride_2d; - - for (int64_t ii = 0; ii < head_size; ii++) { - int64_t t_h_i_offset = t_h_offset + ii; - int64_t h_2d_i_offset = h_2d_offset + ii * h_stride; + } else { + vec_binary_op_non_contiguous(ne0, ne10, nb10, dst_ptr, src0_ptr, src1_ptr); + } + } +} - GGML_F32_VEC v_vec = GGML_F32_VEC_SET1(v[t_h_i_offset]); +// TODO: Use the 'traits' lookup table (for type conversion fns), instead of a mass of 'if' conditions with long templates +template +static void binary_op(const ggml_compute_params * params, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; - float sa = 0; - { - GGML_F32_VEC sum[GGML_F32_ARR] = { GGML_F32_VEC_ZERO }; - GGML_F32_VEC ax[GGML_F32_ARR]; - GGML_F32_VEC ay[GGML_F32_ARR]; - for (int64_t j = 0; j < head_size; j += GGML_F32_STEP) { - for (int64_t kk = 0; kk < GGML_F32_ARR; kk++) { - ax[kk] = GGML_F32_VEC_LOAD(&a[t_h_offset + j + kk * GGML_F32_EPR]); - ay[kk] = GGML_F32_VEC_LOAD(&state_prev[h_2d_i_offset + j + kk * GGML_F32_EPR]); - sum[kk] = GGML_F32_VEC_FMA(sum[kk], ax[kk], ay[kk]); - } - } - GGML_F32_VEC_REDUCE(sa, sum); - } + /* */ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { // all f32 + apply_binary_op(params, dst); + } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { // all f16 + apply_binary_op(params, dst); + } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) { // all bf16 + apply_binary_op(params, dst); + } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_BF16) { + apply_binary_op(params, dst); + } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + apply_binary_op(params, dst); + } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) { + apply_binary_op(params, dst); + } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + apply_binary_op(params, dst); + } else { + GGML_ABORT("%s: unsupported types: dst: %s, src0: %s, src1: %s\n", __func__, + ggml_type_name(dst->type), ggml_type_name(src0->type), ggml_type_name(src1->type)); + } +} - GGML_F32_VEC sa_vec = GGML_F32_VEC_SET1(sa); +void ggml_compute_forward_add_non_quantized(const ggml_compute_params * params, ggml_tensor * dst) { + binary_op(params, dst); +} - int64_t j = 0; - GGML_F32_VEC result_vec[GGML_F32_ARR] = { GGML_F32_VEC_ZERO }; - for (; j < head_size; j += GGML_F32_STEP) { - for (int64_t kk = 0; kk < GGML_F32_ARR; kk++) { - int64_t t_h_j_offset = t_h_offset + j + kk * GGML_F32_EPR; - int64_t h_2d_i_j_offset = h_2d_i_offset + j + kk * GGML_F32_EPR; +void ggml_compute_forward_sub(const ggml_compute_params * params, ggml_tensor * dst) { + binary_op(params, dst); +} - GGML_F32_VEC r_vec = GGML_F32_VEC_LOAD(&r[t_h_j_offset]); - GGML_F32_VEC w_vec = GGML_F32_VEC_LOAD(&w[t_h_j_offset]); - GGML_F32_VEC k_vec = GGML_F32_VEC_LOAD(&k[t_h_j_offset]); - GGML_F32_VEC b_vec = GGML_F32_VEC_LOAD(&b[t_h_j_offset]); +void ggml_compute_forward_mul(const ggml_compute_params * params, ggml_tensor * dst) { + binary_op(params, dst); +} - k_vec = GGML_F32_VEC_MUL(v_vec, k_vec); +void ggml_compute_forward_div(const ggml_compute_params * params, ggml_tensor * dst) { + binary_op(params, dst); +} - GGML_F32_VEC state_vec = GGML_F32_VEC_LOAD(&state_prev[h_2d_i_j_offset]); - // kv + s * decay + sa * b - state_vec = GGML_F32_VEC_FMA(k_vec, state_vec, w_vec); - state_vec = GGML_F32_VEC_FMA(state_vec, sa_vec, b_vec); - GGML_F32_VEC_STORE(&state_cur[h_2d_i_j_offset], state_vec); - result_vec[kk] = GGML_F32_VEC_FMA(result_vec[kk], state_vec, r_vec); - } - } - GGML_F32_VEC_REDUCE(dst_data[t_h_i_offset], result_vec); +// ---- unary-ops implementation ---- - // There shouldn't be left-overs though. - for (; j < head_size; j++) { - int64_t t_h_j_offset = t_h_offset + j; - int64_t h_2d_i_j_offset = h_2d_i_offset + j; - float r_val = r[t_h_j_offset]; - float w_val = w[t_h_j_offset]; - float k_val = k[t_h_j_offset]; - float b_val = b[t_h_j_offset]; - float kv_val = v[t_h_i_offset] * k_val; +extern "C" { +void ggml_compute_forward_abs(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sgn(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_neg(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_step(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_tanh(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_elu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_relu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sigmoid(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_hardsigmoid(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_exp(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_hardswish(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sqr(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sqrt(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sin(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_cos(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_log(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_expm1(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_softplus(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_floor(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_ceil(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_round(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_trunc(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_xielu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +} - float prev_state_val = state_prev[h_2d_i_j_offset]; - state_cur[h_2d_i_j_offset] = prev_state_val * w_val + kv_val + sa * b_val; - dst_data[t_h_i_offset] += state_cur[h_2d_i_j_offset] * r_val; - } - } - } - } - #endif - #else - for (int64_t t = 0; t < T; t++) { - int64_t t_offset = t * t_stride; - int64_t state_offset = head_size * C * (t / (T / n_seqs)); - float * state_cur = state + state_offset; - float * state_prev = t % (T / n_seqs) ? state_cur : (float*)dst->src[6]->data + state_offset; +static inline float op_abs(float x) { + return fabsf(x); +} - for (int64_t h = h_start; h < h_end; h++) { - int64_t h_offset = h * h_stride; - int64_t t_h_offset = t_offset + h_offset; - int64_t h_2d_offset = h * h_stride_2d; +static inline float op_sgn(float x) { + return (x > 0.f) ? 1.f : ((x < 0.f) ? -1.f : 0.f); +} - for (int64_t i = 0; i < head_size; i++) { - int64_t t_h_i_offset = t_h_offset + i; - int64_t h_2d_i_offset = h_2d_offset + i * h_stride; +static inline float op_neg(float x) { + return -x; +} - float v_val = v[t_h_i_offset]; +static inline float op_step(float x) { + return (x > 0.f) ? 1.f : 0.f; +} - float sa = 0, result = 0; - for (int64_t j = 0; j < head_size; j++) { - sa += a[t_h_offset + j] * state_prev[h_2d_i_offset + j]; - } +static inline float op_tanh(float x) { + return tanhf(x); +} - for (int64_t j = 0; j < head_size; j++) { - int64_t t_h_j_offset = t_h_offset + j; - int64_t h_2d_i_j_offset = h_2d_i_offset + j; +static inline float op_elu(float x) { + return (x > 0.f) ? x : expm1f(x); +} - float r_val = r[t_h_j_offset]; - float w_val = w[t_h_j_offset]; - float k_val = k[t_h_j_offset]; - float b_val = b[t_h_j_offset]; - float kv_val = v_val * k_val; - float prev_state_val = state_prev[h_2d_i_j_offset]; - state_cur[h_2d_i_j_offset] = prev_state_val * w_val + kv_val + sa * b_val; - result += state_cur[h_2d_i_j_offset] * r_val; - } - dst_data[t_h_i_offset] = result; - } - } - } - #endif +static inline float op_relu(float x) { + return (x > 0.f) ? x : 0.f; } +static inline float op_sigmoid(float x) { + return 1.f / (1.f + expf(-x)); +} -void ggml_compute_forward_rwkv_wkv7( - const ggml_compute_params * params, - ggml_tensor * dst) { +static inline float op_hardsigmoid(float x) { + return fminf(1.0f, fmaxf(0.0f, (x + 3.0f) / 6.0f)); +} - const ggml_tensor * src0 = dst->src[0]; +static inline float op_exp(float x) { + return expf(x); +} - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_rwkv_wkv7_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } +static inline float op_hardswish(float x) { + return x * fminf(1.0f, fmaxf(0.0f, (x + 3.0f) / 6.0f)); +} + +static inline float op_sqr(float x) { + return x * x; +} + +static inline float op_sqrt(float x) { + return sqrtf(x); +} + +static inline float op_xielu(float x, float alpha_n, float alpha_p, float beta, float eps) { + if (x > 0.0f) { + return alpha_p * x * x + beta * x; + } else { + const float min_x_eps = fminf(x, eps); + return (expm1f(min_x_eps) - x) * alpha_n + beta * x; } } -// ggml_compute_forward_map_custom1 +static inline float op_sin(float x) { + return sinf(x); +} -void ggml_compute_forward_map_custom1( - const ggml_compute_params * params, - ggml_tensor * dst) { +static inline float op_cos(float x) { + return cosf(x); +} - const ggml_tensor * a = dst->src[0]; +static inline float op_log(float x) { + return logf(x); +} - struct ggml_map_custom1_op_params p; - memcpy(&p, dst->op_params, sizeof(p)); +static inline float op_expm1(float x) { + return expf(x) - 1.0f; +} - p.fun(dst, a, params->ith, params->nth, p.userdata); +static inline float op_softplus(float x) { + return (x > 20.0f) ? x : logf(1.0f + expf(x)); } -// ggml_compute_forward_map_custom2 +static inline float op_floor(float x) { + return floorf(x); +} -void ggml_compute_forward_map_custom2( - const ggml_compute_params * params, - ggml_tensor * dst) { +static inline float op_ceil(float x) { + return ceilf(x); +} - const ggml_tensor * a = dst->src[0]; - const ggml_tensor * b = dst->src[1]; +static inline float op_round(float x) { + return roundf(x); +} - struct ggml_map_custom2_op_params p; - memcpy(&p, dst->op_params, sizeof(p)); +static inline float op_trunc(float x) { + return truncf(x); +} - p.fun(dst, a, b, params->ith, params->nth, p.userdata); +template +static inline void vec_unary_op(int64_t n, dst_t * y, const src0_t * x) { + constexpr auto src0_to_f32 = type_conversion_table::to_f32; + constexpr auto f32_to_dst = type_conversion_table::from_f32; + + for (int i = 0; i < n; i++) { + y[i] = f32_to_dst(op(src0_to_f32(x[i]))); + } } -// ggml_compute_forward_map_custom3 +template +static void apply_unary_op(const ggml_compute_params * params, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; -void ggml_compute_forward_map_custom3( - const ggml_compute_params * params, - ggml_tensor * dst) { + GGML_ASSERT(ggml_is_contiguous_rows(src0) && ggml_is_contiguous_rows(dst) && ggml_are_same_shape(src0, dst)); - const ggml_tensor * a = dst->src[0]; - const ggml_tensor * b = dst->src[1]; - const ggml_tensor * c = dst->src[2]; + GGML_TENSOR_UNARY_OP_LOCALS - struct ggml_map_custom3_op_params p; - memcpy(&p, dst->op_params, sizeof(p)); + GGML_ASSERT( nb0 == sizeof(dst_t)); + GGML_ASSERT(nb00 == sizeof(src0_t)); - p.fun(dst, a, b, c, params->ith, params->nth, p.userdata); + const auto [ir0, ir1] = get_thread_range(params, src0); + + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t i03 = ir/(ne02*ne01); + const int64_t i02 = (ir - i03*ne02*ne01)/ne01; + const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); + + dst_t * dst_ptr = (dst_t *) ((char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1 ); + const src0_t * src0_ptr = (const src0_t *) ((const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01); + + vec_unary_op(ne0, dst_ptr, src0_ptr); + } } -// ggml_compute_forward_custom +// TODO: Use the 'traits' lookup table (for type conversion fns), instead of a mass of 'if' conditions with long templates +template +static void unary_op(const ggml_compute_params * params, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; -void ggml_compute_forward_custom( - const struct ggml_compute_params * params, - struct ggml_tensor * dst) { + /* */ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { // all f32 + apply_unary_op(params, dst); + } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { // all f16 + apply_unary_op(params, dst); + } else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) { // all bf16 + apply_unary_op(params, dst); + } else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_F32) { + apply_unary_op(params, dst); + } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) { + apply_unary_op(params, dst); + } else { + fprintf(stderr, "%s: unsupported types: dst: %s, src0: %s\n", __func__, + ggml_type_name(dst->type), ggml_type_name(src0->type)); + GGML_ABORT("fatal error"); + } +} - struct ggml_custom_op_params p; - memcpy(&p, dst->op_params, sizeof(p)); +template +static void unary_op_params(const ggml_compute_params * params, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; - p.fun(dst, params->ith, params->nth, p.userdata); + /* */ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { // all f32 + apply_unary_op(params, dst); + } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { // all f16 + apply_unary_op(params, dst); + } else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) { // all bf16 + apply_unary_op(params, dst); + } else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_F32) { + apply_unary_op(params, dst); + } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) { + apply_unary_op(params, dst); + } else { + fprintf(stderr, "%s: unsupported types: dst: %s, src0: %s\n", __func__, + ggml_type_name(dst->type), ggml_type_name(src0->type)); + GGML_ABORT("fatal error"); + } } -// ggml_compute_forward_cross_entropy_loss +// Extend vec_unary_op to support functors +template +static inline void vec_unary_op_functor(int64_t n, dst_t * y, const src0_t * x, Op op) { + constexpr auto src0_to_f32 = type_conversion_table::to_f32; + constexpr auto f32_to_dst = type_conversion_table::from_f32; -static void ggml_compute_forward_cross_entropy_loss_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { + for (int i = 0; i < n; i++) { + y[i] = f32_to_dst(op(src0_to_f32(x[i]))); + } +} +// Extend apply_unary_op to support functors +template +static void apply_unary_op_functor(const ggml_compute_params * params, ggml_tensor * dst, Op op) { const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(src0->nb[0] == ggml_type_size(src0->type)); - GGML_ASSERT(src1->nb[0] == ggml_type_size(src1->type)); - GGML_ASSERT(ggml_are_same_shape(src0, src1)); - GGML_ASSERT(ggml_is_scalar(dst)); - GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous_1(src0) && ggml_is_contiguous_1(dst) && ggml_are_same_shape(src0, dst)); - // TODO: handle transposed/permuted matrices - const int64_t nc = src0->ne[0]; - const int64_t nr = ggml_nrows(src0); + GGML_TENSOR_UNARY_OP_LOCALS - const int ith = params->ith; - const int nth = params->nth; + GGML_ASSERT( nb0 == sizeof(dst_t)); + GGML_ASSERT(nb00 == sizeof(src0_t)); - float * sums = (float *) params->wdata; - float * st = ((float *) params->wdata) + nth + ith*nc; - float sum_thread = 0.0f; + const auto [ir0, ir1] = get_thread_range(params, src0); - GGML_ASSERT(params->wsize >= sizeof(float) * (nth + nth * nc)); + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t i03 = ir/(ne02*ne01); + const int64_t i02 = (ir - i03*ne02*ne01)/ne01; + const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); - // rows per thread - const int64_t dr = (nr + nth - 1)/nth; + dst_t * dst_ptr = (dst_t *) ((char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1 ); + const src0_t * src0_ptr = (const src0_t *) ((const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01); - // row range for this thread - const int64_t ir0 = dr*ith; - const int64_t ir1 = MIN(ir0 + dr, nr); + vec_unary_op_functor(ne0, dst_ptr, src0_ptr, op); + } +} - for (int64_t i1 = ir0; i1 < ir1; ++i1) { - const float * s0 = (const float *)((const char *) src0->data + i1*src0->nb[1]); - const float * s1 = (const float *)((const char *) src1->data + i1*src1->nb[1]); +// Generic dispatcher for functors +template +static void unary_op_functor(const ggml_compute_params * params, ggml_tensor * dst, Op op) { + const ggml_tensor * src0 = dst->src[0]; -#ifndef NDEBUG - for (int64_t i = 0; i < nc; ++i) { - //printf("p[%d] = %f\n", i, p[i]); - assert(!isnan(s0[i])); - assert(!isnan(s1[i])); - } -#endif // NDEBUG + /* */ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { // all f32 + apply_unary_op_functor(params, dst, op); + } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { // all f16 + apply_unary_op_functor(params, dst, op); + } else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) { // all bf16 + apply_unary_op_functor(params, dst, op); + } else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_F32) { + apply_unary_op_functor(params, dst, op); + } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) { + apply_unary_op_functor(params, dst, op); + } else { + fprintf(stderr, "%s: unsupported types: dst: %s, src0: %s\n", __func__, + ggml_type_name(dst->type), ggml_type_name(src0->type)); + GGML_ABORT("fatal error"); + } +} - float max = -INFINITY; - ggml_vec_max_f32(nc, &max, s0); - const ggml_float sum_softmax = ggml_vec_log_soft_max_f32(nc, st, s0, max); - assert(sum_softmax >= 0.0); +void ggml_compute_forward_abs(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - ggml_vec_add1_f32(nc, st, st, -sum_softmax); - ggml_vec_mul_f32(nc, st, st, s1); +void ggml_compute_forward_sgn(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - float sum_st = 0.0f; - ggml_vec_sum_f32(nc, &sum_st, st); - sum_thread += sum_st; +void ggml_compute_forward_neg(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} -#ifndef NDEBUG - for (int64_t i = 0; i < nc; ++i) { - assert(!isnan(st[i])); - assert(!isinf(st[i])); - } -#endif // NDEBUG - } - sums[ith] = sum_thread; - ggml_barrier(params->threadpool); +void ggml_compute_forward_step(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - if (ith == 0) { - float * dp = (float *) dst->data; - ggml_vec_sum_f32(nth, dp, sums); - dp[0] *= -1.0f / (float) nr; - } +void ggml_compute_forward_tanh(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); } -void ggml_compute_forward_cross_entropy_loss( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_elu(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - const ggml_tensor * src0 = dst->src[0]; +void ggml_compute_forward_relu(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_cross_entropy_loss_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } +void ggml_compute_forward_sigmoid(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); } -// ggml_compute_forward_cross_entropy_loss_back +void ggml_compute_forward_hardsigmoid(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} -static void ggml_compute_forward_cross_entropy_loss_back_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_exp(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - const ggml_tensor * grad = dst->src[0]; // gradient of forward pass output - const ggml_tensor * src0f = dst->src[1]; // src0 of forward pass - const ggml_tensor * src1f = dst->src[2]; // src1 of forward pass +void ggml_compute_forward_hardswish(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - GGML_ASSERT(ggml_is_contiguous(dst)); - GGML_ASSERT(ggml_is_contiguous(src0f)); - GGML_ASSERT(ggml_is_contiguous(src1f)); - GGML_ASSERT(ggml_is_contiguous(grad)); - GGML_ASSERT(ggml_are_same_shape(src0f, src1f) && ggml_are_same_shape(src0f, dst)); +void ggml_compute_forward_sqr(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - const int64_t ith = params->ith; - const int64_t nth = params->nth; +void ggml_compute_forward_sqrt(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - // TODO: handle transposed/permuted matrices - const int64_t nc = src0f->ne[0]; - const int64_t nr = ggml_nrows(src0f); +void ggml_compute_forward_sin(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - // rows per thread - const int64_t dr = (nr + nth - 1)/nth; +void ggml_compute_forward_cos(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - // row range for this thread - const int64_t ir0 = dr*ith; - const int64_t ir1 = MIN(ir0 + dr, nr); +void ggml_compute_forward_log(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - const float d_by_nr = ((const float *) grad->data)[0] / (float) nr; +void ggml_compute_forward_expm1(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - for (int64_t i1 = ir0; i1 < ir1; i1++) { - float * ds0 = (float *)((char *) dst->data + i1*dst->nb[1]); - const float * s0 = (const float *)((const char *) src0f->data + i1*src0f->nb[1]); - const float * s1 = (const float *)((const char *) src1f->data + i1*src1f->nb[1]); +void ggml_compute_forward_softplus(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} -#ifndef NDEBUG - for (int64_t i = 0; i < nc; ++i) { - //printf("p[%d] = %f\n", i, p[i]); - assert(!isnan(s0[i])); - assert(!isnan(s1[i])); - } -#endif // NDEBUG +void ggml_compute_forward_floor(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - // soft_max - float max = -INFINITY; - ggml_vec_max_f32(nc, &max, s0); - const ggml_float sum = ggml_vec_soft_max_f32(nc, ds0, s0, max); - assert(sum > 0.0); - ggml_vec_scale_f32(nc, ds0, 1.0/sum); +void ggml_compute_forward_ceil(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} - // grad(src0f) = (softmax(src0f) - src1f) * grad(cross_entropy_loss(src0f, src1f)) / nr - ggml_vec_sub_f32(nc, ds0, ds0, s1); - ggml_vec_scale_f32(nc, ds0, d_by_nr); +void ggml_compute_forward_round(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); +} -#ifndef NDEBUG - for (int64_t i = 0; i < nc; ++i) { - assert(!isnan(ds0[i])); - assert(!isinf(ds0[i])); - } -#endif // NDEBUG - } +void ggml_compute_forward_trunc(const ggml_compute_params * params, ggml_tensor * dst) { + unary_op(params, dst); } -void ggml_compute_forward_cross_entropy_loss_back( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_xielu(const ggml_compute_params * params, ggml_tensor * dst) { + const float alpha_n = ggml_get_op_params_f32(dst, 1); + const float alpha_p = ggml_get_op_params_f32(dst, 2); + const float beta = ggml_get_op_params_f32(dst, 3); + const float eps = ggml_get_op_params_f32(dst, 4); - const ggml_tensor * src0 = dst->src[0]; + const auto xielu_op_params = [alpha_n, alpha_p, beta, eps](float f) { + return op_xielu(f, alpha_n, alpha_p, beta, eps); + }; - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_cross_entropy_loss_back_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } + unary_op_functor(params, dst, xielu_op_params); } -static void ggml_compute_forward_opt_step_adamw_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src0_grad = dst->src[1]; - const ggml_tensor * src0_grad_m = dst->src[2]; - const ggml_tensor * src0_grad_v = dst->src[3]; - const ggml_tensor * adamw_params = dst->src[4]; +// ---- vec implementation ---- - GGML_ASSERT(ggml_are_same_shape(src0, src0_grad)); - GGML_ASSERT(ggml_are_same_shape(src0, src0_grad_m)); - GGML_ASSERT(ggml_are_same_shape(src0, src0_grad_v)); - GGML_ASSERT(ggml_nelements(adamw_params) == 7); - const int ith = params->ith; - const int nth = params->nth; +#include - const int nr = ggml_nrows(src0); +// precomputed gelu table for f16 (128 KB) +ggml_fp16_t ggml_table_gelu_f16[1 << 16]; - GGML_TENSOR_UNARY_OP_LOCALS - GGML_ASSERT(nb00 == sizeof(float)); +// precomputed quick gelu table for f16 (128 KB) +ggml_fp16_t ggml_table_gelu_quick_f16[1 << 16]; - // rows per thread - const int dr = (nr + nth - 1)/nth; +void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc) { + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); +#if defined(GGML_SIMD) + float sumf = 0.0f; - const float * adamw_params_ptr = ggml_get_data_f32(adamw_params); + #if defined(__ARM_FEATURE_SVE) + const int sve_register_length = ggml_cpu_get_sve_cnt() * 8; + const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16 + const int ggml_f32_step = 8 * ggml_f32_epr; // choose 8 SVE registers + + const int np = (n & ~(ggml_f32_step - 1)); + svfloat32_t sum1 = svdup_n_f32(0.0f); + svfloat32_t sum2 = svdup_n_f32(0.0f); + svfloat32_t sum3 = svdup_n_f32(0.0f); + svfloat32_t sum4 = svdup_n_f32(0.0f); + svfloat32_t sum5 = svdup_n_f32(0.0f); + svfloat32_t sum6 = svdup_n_f32(0.0f); + svfloat32_t sum7 = svdup_n_f32(0.0f); + svfloat32_t sum8 = svdup_n_f32(0.0f); + svfloat32_t ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8; + svfloat32_t ay1,ay2,ay3,ay4,ay5,ay6,ay7,ay8; + for (int i = 0; i < np; i += ggml_f32_step) { + ax1 = GGML_F32_VEC_LOAD(x + i); + ay1 = GGML_F32_VEC_LOAD(y + i); + sum1 = GGML_F32_VEC_FMA(sum1, ax1, ay1); + + ax2 = GGML_F32_VEC_LOAD(x + i + 1*ggml_f32_epr); + ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr); + sum2 = GGML_F32_VEC_FMA(sum2, ax2, ay2); + + ax3 = GGML_F32_VEC_LOAD(x + i + 2*ggml_f32_epr); + ay3 = GGML_F32_VEC_LOAD(y + i + 2*ggml_f32_epr); + sum3 = GGML_F32_VEC_FMA(sum3, ax3, ay3); + + ax4 = GGML_F32_VEC_LOAD(x + i + 3*ggml_f32_epr); + ay4 = GGML_F32_VEC_LOAD(y + i + 3*ggml_f32_epr); + sum4 = GGML_F32_VEC_FMA(sum4, ax4, ay4); + + ax5 = GGML_F32_VEC_LOAD(x + i + 4*ggml_f32_epr); + ay5 = GGML_F32_VEC_LOAD(y + i + 4*ggml_f32_epr); + sum5 = GGML_F32_VEC_FMA(sum5, ax5, ay5); + + ax6 = GGML_F32_VEC_LOAD(x + i + 5*ggml_f32_epr); + ay6 = GGML_F32_VEC_LOAD(y + i + 5*ggml_f32_epr); + sum6 = GGML_F32_VEC_FMA(sum6, ax6, ay6); + + ax7 = GGML_F32_VEC_LOAD(x + i + 6*ggml_f32_epr); + ay7 = GGML_F32_VEC_LOAD(y + i + 6*ggml_f32_epr); + sum7 = GGML_F32_VEC_FMA(sum7, ax7, ay7); + + ax8 = GGML_F32_VEC_LOAD(x + i + 7*ggml_f32_epr); + ay8 = GGML_F32_VEC_LOAD(y + i + 7*ggml_f32_epr); + sum8 = GGML_F32_VEC_FMA(sum8, ax8, ay8); + } + // leftovers + // Since 8 unrolls are done in above loop, leftovers lie in range [0, ggml_f32_step] which is handled in below loop + const int np2 = (n & ~(ggml_f32_epr - 1)); + for (int i = np; i < np2; i += ggml_f32_epr) { + ax1 = GGML_F32_VEC_LOAD(x + i); + ay1 = GGML_F32_VEC_LOAD(y + i); + sum1 = GGML_F32_VEC_FMA(sum1, ax1, ay1); + } + // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only + if (np2 < n) { + svbool_t pg = svwhilelt_b32(np2, n); + ax1 = svld1_f32(pg, x + np2); + ay1 = svld1_f32(pg, y + np2); + sum1 = svmad_f32_m(pg, ax1, ay1, sum1); + } + // reduce sum1,sum2 to sum1 + GGML_F32_VEC_REDUCE(sumf, sum1, sum2, sum3, sum4, sum5, sum6, sum7, sum8); + #elif defined(__riscv_v_intrinsic) + int vl = __riscv_vsetvlmax_e32m8(); + vfloat32m1_t vs = __riscv_vfmv_v_f_f32m1(0.0f, 1); + vfloat32m8_t vsum; + vfloat32m8_t ax; + vfloat32m8_t ay; + vsum = __riscv_vfmv_v_f_f32m8_tu(vsum, 0.0f, vl); + for (int i = 0; i < n; i += vl) { + vl = __riscv_vsetvl_e32m8(n - i); + ax = __riscv_vle32_v_f32m8_tu(ax, &x[i], vl); + ay = __riscv_vle32_v_f32m8_tu(ay, &y[i], vl); + vsum = __riscv_vfmacc_vv_f32m8_tu(vsum, ax, ay, vl); + } + vl = __riscv_vsetvlmax_e32m8(); + vs = __riscv_vfredusum_vs_f32m8_f32m1(vsum, vs, vl); + sumf += __riscv_vfmv_f_s_f32m1_f32(vs); + #else + const int np = (n & ~(GGML_F32_STEP - 1)); - const float alpha = adamw_params_ptr[0]; - const float beta1 = adamw_params_ptr[1]; - const float beta2 = adamw_params_ptr[2]; - const float eps = adamw_params_ptr[3]; - const float wd = adamw_params_ptr[4]; - const float beta1h = adamw_params_ptr[5]; - const float beta2h = adamw_params_ptr[6]; - const float keep = 1.f - alpha * wd; - for (int ir = ir0; ir < ir1; ++ir) { - const int64_t i03 = ir/(ne02*ne01); - const int64_t i02 = (ir - i03*ne02*ne01)/ne01; - const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); + GGML_F32_VEC sum[GGML_F32_ARR] = { GGML_F32_VEC_ZERO }; - const size_t offset = i03*nb03 + i02*nb02 + i01*nb01; + GGML_F32_VEC ax[GGML_F32_ARR]; + GGML_F32_VEC ay[GGML_F32_ARR]; - float * w = (float *) ((char *) src0->data + offset); // weight - const float * g = (const float *) ((const char *) src0_grad->data + offset); // grad - float * m = (float *) ((char *) src0_grad_m->data + offset); - float * v = (float *) ((char *) src0_grad_v->data + offset); + for (int i = 0; i < np; i += GGML_F32_STEP) { + for (int j = 0; j < GGML_F32_ARR; j++) { + ax[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR); + ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); - for (int i00 = 0; i00 < ne00; ++i00) { - m[i00] = m[i00]*beta1 + g[i00]*(1.0f - beta1); - v[i00] = v[i00]*beta2 + g[i00]*g[i00]*(1.0f - beta2); + sum[j] = GGML_F32_VEC_FMA(sum[j], ax[j], ay[j]); + } + } - const float mh = m[i00]*beta1h; - const float vh = sqrtf(v[i00]*beta2h) + eps; + // reduce sum0..sum3 to sum0 + GGML_F32_VEC_REDUCE(sumf, sum); - // The weight decay is applied independently of the Adam momenta m and v. - // This is NOT equivalent to l2 regularization that adds w[i00]*w[i00] to the loss. - // See: https://arxiv.org/pdf/1711.05101v3.pdf - w[i00] = w[i00] * keep - alpha * mh / vh; + // leftovers + for (int i = np; i < n; ++i) { + sumf += x[i]*y[i]; } + #endif +#else + // scalar + ggml_float sumf = 0.0; + for (int i = 0; i < n; ++i) { + sumf += (ggml_float)(x[i]*y[i]); } -} - -void ggml_compute_forward_opt_step_adamw( - const ggml_compute_params * params, - ggml_tensor * dst) { +#endif - const ggml_tensor * src0 = dst->src[0]; + *s = sumf; +} + +void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t * GGML_RESTRICT x, size_t bx, ggml_bf16_t * GGML_RESTRICT y, size_t by, int nrc) { + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + int i = 0; + ggml_float sumf = 0; + +#if defined(__AVX512BF16__) + __m512 c1 = _mm512_setzero_ps(); + __m512 c2 = _mm512_setzero_ps(); + for (; i + 64 <= n; i += 64) { + c1 = _mm512_dpbf16_ps(c1, m512bh(_mm512_loadu_si512((x + i))), + m512bh(_mm512_loadu_si512((y + i)))); + c2 = _mm512_dpbf16_ps(c2, m512bh(_mm512_loadu_si512((x + i + 32))), + m512bh(_mm512_loadu_si512((y + i + 32)))); + } + sumf += (ggml_float)_mm512_reduce_add_ps(c1); + sumf += (ggml_float)_mm512_reduce_add_ps(c2); + +#elif defined(__AVX512F__) +#define LOAD(p) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)(p))), 16)) + __m512 c1 = _mm512_setzero_ps(); + __m512 c2 = _mm512_setzero_ps(); + for (; i + 32 <= n; i += 32) { + c1 = _mm512_add_ps(_mm512_mul_ps(LOAD(x + i), LOAD(y + i)), c1); + c2 = _mm512_add_ps(_mm512_mul_ps(LOAD(x + i + 16), LOAD(y + i + 16)), c2); + } + sumf += (ggml_float)_mm512_reduce_add_ps(c1); + sumf += (ggml_float)_mm512_reduce_add_ps(c2); + +#undef LOAD +#elif defined(__AVX2__) || defined(__AVX__) +#if defined(__AVX2__) +#define LOAD(p) _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(_mm_loadu_si128((const __m128i *)(p))), 16)) +#else +#define LOAD(p) _mm256_castsi256_ps(_mm256_insertf128_si256(_mm256_castsi128_si256(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadu_si128((const __m128i *)(p))), 16)), (_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_bsrli_si128(_mm_loadu_si128((const __m128i *)(p)), 8)), 16)), 1)) +#endif + __m256 c1 = _mm256_setzero_ps(); + __m256 c2 = _mm256_setzero_ps(); + __m256 c3 = _mm256_setzero_ps(); + __m256 c4 = _mm256_setzero_ps(); + for (; i + 32 <= n; i += 32) { + c1 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i), LOAD(y + i)), c1); + c2 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i + 8), LOAD(y + i + 8)), c2); + c3 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i + 16), LOAD(y + i + 16)), c3); + c4 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i + 24), LOAD(y + i + 24)), c4); + } + __m128 g; + c1 = _mm256_add_ps(_mm256_add_ps(c1, c3), + _mm256_add_ps(c2, c4)); + g = _mm_add_ps(_mm256_extractf128_ps(c1, 1), + _mm256_castps256_ps128(c1)); + g = _mm_add_ps(g, _mm_movehl_ps(g, g)); + g = _mm_add_ss(g, _mm_movehdup_ps(g)); + sumf += (ggml_float)_mm_cvtss_f32(g); + +#undef LOAD +#elif defined(__riscv_v_intrinsic) && defined(__riscv_zvfbfwma) + size_t vl = __riscv_vsetvlmax_e32m4(); + + // initialize accumulators to all zeroes + vfloat32m4_t vsum0 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + vfloat32m4_t vsum1 = __riscv_vfmv_v_f_f32m4(0.0f, vl); + + // calculate step size + const size_t epr = __riscv_vsetvlmax_e16m2(); + const size_t step = epr * 2; + const int np = (n & ~(step - 1)); + + // unroll by 2 + for (; i < np; i += step) { + vbfloat16m2_t ax0 = __riscv_vle16_v_bf16m2((const __bf16 *)&x[i], epr); + vbfloat16m2_t ay0 = __riscv_vle16_v_bf16m2((const __bf16 *)&y[i], epr); + vsum0 = __riscv_vfwmaccbf16_vv_f32m4(vsum0, ax0, ay0, epr); + __asm__ __volatile__ ("" ::: "memory"); + + vbfloat16m2_t ax1 = __riscv_vle16_v_bf16m2((const __bf16 *)&x[i + epr], epr); + vbfloat16m2_t ay1 = __riscv_vle16_v_bf16m2((const __bf16 *)&y[i + epr], epr); + vsum1 = __riscv_vfwmaccbf16_vv_f32m4(vsum1, ax1, ay1, epr); + __asm__ __volatile__ ("" ::: "memory"); + } + + // accumulate in 1 register + vsum0 = __riscv_vfadd_vv_f32m4(vsum0, vsum1, vl); + + // leftovers + for (i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m2(n - i); + vbfloat16m2_t ax0 = __riscv_vle16_v_bf16m2((const __bf16 *)&x[i], vl); + vbfloat16m2_t ay0 = __riscv_vle16_v_bf16m2((const __bf16 *)&y[i], vl); + vsum0 = __riscv_vfwmaccbf16_vv_f32m4(vsum0, ax0, ay0, vl); + } + + // reduce + vl = __riscv_vsetvlmax_e32m4(); + vfloat32m1_t redsum = __riscv_vfredusum_vs_f32m4_f32m1(vsum0, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl); + sumf += __riscv_vfmv_f_s_f32m1_f32(redsum); + +#elif defined(__POWER9_VECTOR__) || defined(__VXE__) || defined(__VXE2__) + const int np = (n & ~(GGML_BF16_STEP - 1)); + if (np > 0) { + GGML_F32_VEC sum[4] = {GGML_F32_VEC_ZERO}; + for (; i < np; i += GGML_BF16_STEP) { + GGML_BF16_VEC vx0 = GGML_BF16_VEC_LOAD(x + i); + GGML_BF16_VEC vx1 = GGML_BF16_VEC_LOAD(x + i + 8); + GGML_BF16_VEC vy0 = GGML_BF16_VEC_LOAD(y + i); + GGML_BF16_VEC vy1 = GGML_BF16_VEC_LOAD(y + i + 8); + GGML_BF16_FMA_LO(sum[0], vx0, vy0); + GGML_BF16_FMA_HI(sum[1], vx0, vy0); + GGML_BF16_FMA_LO(sum[2], vx1, vy1); + GGML_BF16_FMA_HI(sum[3], vx1, vy1); + } + GGML_F32x4_REDUCE_4(sumf, sum[0], sum[1], sum[2], sum[3]); + } +#endif - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_opt_step_adamw_f32(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } + for (; i < n; ++i) { + sumf += (ggml_float)(GGML_BF16_TO_FP32(x[i]) * + GGML_BF16_TO_FP32(y[i])); } + *s = sumf; } -static void ggml_compute_forward_opt_step_sgd_f32(const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src0_grad = dst->src[1]; - const ggml_tensor * sgd_params = dst->src[2]; +void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * GGML_RESTRICT x, size_t bx, ggml_fp16_t * GGML_RESTRICT y, size_t by, int nrc) { + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); - GGML_ASSERT(ggml_are_same_shape(src0, src0_grad)); - GGML_ASSERT(ggml_nelements(sgd_params) == 2); + ggml_float sumf = 0.0; - const int ith = params->ith; - const int nth = params->nth; - const int nr = ggml_nrows(src0); +#if defined(GGML_SIMD) + #if defined(__ARM_FEATURE_SVE) + const int sve_register_length = svcntb() * 8; //get vector length + const int ggml_f16_epr = sve_register_length / 16; // running when 16 + const int ggml_f16_step = 8 * ggml_f16_epr; // choose 8 SVE registers - GGML_TENSOR_UNARY_OP_LOCALS - GGML_ASSERT(nb00 == sizeof(float)); + const int np= (n & ~(ggml_f16_step - 1)); + svfloat16_t sum1 = svdup_n_f16(0.0f); + svfloat16_t sum2 = svdup_n_f16(0.0f); + svfloat16_t sum3 = svdup_n_f16(0.0f); + svfloat16_t sum4 = svdup_n_f16(0.0f); - // rows per thread - const int dr = (nr + nth - 1) / nth; + svfloat16_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; + svfloat16_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; + for (int i = 0; i < np; i += ggml_f16_step) { + ax1 = GGML_F16x_VEC_LOAD(x + i + 0 * ggml_f16_epr, 0); + ay1 = GGML_F16x_VEC_LOAD(y + i + 0 * ggml_f16_epr, 0); + sum1 = GGML_F16x_VEC_FMA(sum1, ax1, ay1); - // row range for this thread - const int ir0 = dr * ith; - const int ir1 = MIN(ir0 + dr, nr); + ax2 = GGML_F16x_VEC_LOAD(x + i + 1 * ggml_f16_epr, 1); + ay2 = GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 1); + sum2 = GGML_F16x_VEC_FMA(sum2, ax2, ay2); - // using adamw param subset we care about - alpha, wd - could have a separate struct - const float * sgd_params_ptr = ggml_get_data_f32(sgd_params); - const float alpha = sgd_params_ptr[0]; - const float keep = 1.f - alpha * sgd_params_ptr[1]; + ax3 = GGML_F16x_VEC_LOAD(x + i + 2 * ggml_f16_epr, 2); + ay3 = GGML_F16x_VEC_LOAD(y + i + 2 * ggml_f16_epr, 2); + sum3 = GGML_F16x_VEC_FMA(sum3, ax3, ay3); - for (int ir = ir0; ir < ir1; ++ir) { - const int64_t i03 = ir / (ne02 * ne01); - const int64_t i02 = (ir - i03 * ne02 * ne01) / ne01; - const int64_t i01 = (ir - i03 * ne02 * ne01 - i02 * ne01); + ax4 = GGML_F16x_VEC_LOAD(x + i + 3 * ggml_f16_epr, 3); + ay4 = GGML_F16x_VEC_LOAD(y + i + 3 * ggml_f16_epr, 3); + sum4 = GGML_F16x_VEC_FMA(sum4, ax4, ay4); - const size_t offset = i03 * nb03 + i02 * nb02 + i01 * nb01; + ax5 = GGML_F16x_VEC_LOAD(x + i + 4 * ggml_f16_epr, 4); + ay5 = GGML_F16x_VEC_LOAD(y + i + 4 * ggml_f16_epr, 4); + sum1 = GGML_F16x_VEC_FMA(sum1, ax5, ay5); - float * w = (float *) ((char *) src0->data + offset); // weight - const float * g = (const float *) ((const char *) src0_grad->data + offset); // grad + ax6 = GGML_F16x_VEC_LOAD(x + i + 5 * ggml_f16_epr, 5); + ay6 = GGML_F16x_VEC_LOAD(y + i + 5 * ggml_f16_epr, 5); + sum2 = GGML_F16x_VEC_FMA(sum2, ax6, ay6); - for (int i00 = 0; i00 < ne00; ++i00) { - w[i00] = w[i00] * keep - alpha * g[i00]; + ax7 = GGML_F16x_VEC_LOAD(x + i + 6 * ggml_f16_epr, 6); + ay7 = GGML_F16x_VEC_LOAD(y + i + 6 * ggml_f16_epr, 6); + sum3 = GGML_F16x_VEC_FMA(sum3, ax7, ay7); + + ax8 = GGML_F16x_VEC_LOAD(x + i + 7 * ggml_f16_epr, 7); + ay8 = GGML_F16x_VEC_LOAD(y + i + 7 * ggml_f16_epr, 7); + sum4 = GGML_F16x_VEC_FMA(sum4, ax8, ay8); } - } -} -void ggml_compute_forward_opt_step_sgd(const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; + const int np2 = (n & ~(ggml_f16_epr - 1)); // round down to multiple of 8 + for (int k = np; k < np2; k += ggml_f16_epr) { + svfloat16_t rx = GGML_F16x_VEC_LOAD(x + k, 0); + svfloat16_t ry = GGML_F16x_VEC_LOAD(y + k, 0); + sum1 = GGML_F16x_VEC_FMA(sum1, rx, ry); + } - switch (src0->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_opt_step_sgd_f32(params, dst); + if (np2 < n) { + svbool_t pg = svwhilelt_b16(np2, n); + svfloat16_t hx = svld1_f16(pg, (const __fp16 *)(x + np2)); + svfloat16_t hy = svld1_f16(pg, (const __fp16 *)(y + np2)); + + sum1 = svmad_f16_x(pg, hx, hy, sum1); + } + GGML_F16x_VEC_REDUCE(sumf, sum1, sum2, sum3, sum4); + #elif defined(__riscv_v_intrinsic) + #if defined(__riscv_zvfh) + int vl = __riscv_vsetvlmax_e32m2(); + vfloat32m1_t vs = __riscv_vfmv_v_f_f32m1(0.0f, 1); + vfloat32m2_t vsum; + vfloat16m1_t ax; + vfloat16m1_t ay; + vsum = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vmv_v_x_u32m2(0, vl)); + for (int i = 0; i < n; i += vl) { + vl = __riscv_vsetvl_e16m1(n - i); + ax = __riscv_vle16_v_f16m1_tu(ax, (const _Float16 *)&x[i], vl); + ay = __riscv_vle16_v_f16m1_tu(ay, (const _Float16 *)&y[i], vl); + vsum = __riscv_vfwmacc_vv_f32m2_tu(vsum, ax, ay, vl); + } + vl = __riscv_vsetvlmax_e32m1(); + vfloat32m1_t ac0 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(vsum, 0), __riscv_vget_v_f32m2_f32m1(vsum, 1), vl); + vs = __riscv_vfredusum_vs_f32m1_f32m1(ac0, vs, vl); + sumf += __riscv_vfmv_f_s_f32m1_f32(vs); + #else + for (int i = 0; i < n; ++i) { + sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*GGML_CPU_FP16_TO_FP32(y[i])); } - break; - default: - { - GGML_ABORT("fatal error - sgd is F32 only"); + #endif // __riscv_zvfh + #else + const int np = (n & ~(GGML_F16_STEP - 1)); + + GGML_F16_VEC sum[GGML_F16_ARR] = { GGML_F16_VEC_ZERO }; + + GGML_F16_VEC ax[GGML_F16_ARR]; + GGML_F16_VEC ay[GGML_F16_ARR]; + + for (int i = 0; i < np; i += GGML_F16_STEP) { + for (int j = 0; j < GGML_F16_ARR; j++) { + ax[j] = GGML_F16_VEC_LOAD(x + i + j*GGML_F16_EPR, j); + ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); + + sum[j] = GGML_F16_VEC_FMA(sum[j], ax[j], ay[j]); } + } + + // reduce sum0..sum3 to sum0 + GGML_F16_VEC_REDUCE(sumf, sum); + + // leftovers + for (int i = np; i < n; ++i) { + sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*GGML_CPU_FP16_TO_FP32(y[i])); + } + // if you hit this, you are likely running outside the FP range + assert(!isnan(sumf) && !isinf(sumf)); + #endif +#else + for (int i = 0; i < n; ++i) { + sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*GGML_CPU_FP16_TO_FP32(y[i])); + } +#endif // GGML_SIMD + + *s = sumf; +} + +void ggml_vec_silu_f32(const int n, float * y, const float * x) { + int i = 0; +#if defined(__AVX512F__) && defined(__AVX512DQ__) + for (; i + 15 < n; i += 16) { + _mm512_storeu_ps(y + i, ggml_v_silu(_mm512_loadu_ps(x + i))); + } +#elif defined(__AVX2__) && defined(__FMA__) + for (; i + 7 < n; i += 8) { + _mm256_storeu_ps(y + i, ggml_v_silu(_mm256_loadu_ps(x + i))); + } +#elif defined(__SSE2__) + for (; i + 3 < n; i += 4) { + _mm_storeu_ps(y + i, ggml_v_silu(_mm_loadu_ps(x + i))); + } +#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + const int vlen = svcntw(); + for (; i < n; i += vlen) { + const svbool_t pg = svwhilelt_b32_s32(i, n); + svst1_f32(pg, y + i, ggml_v_silu(pg, svld1_f32(pg, x + i))); + } +#elif defined(__ARM_NEON) && defined(__aarch64__) + for (; i + 3 < n; i += 4) { + vst1q_f32(y + i, ggml_v_silu(vld1q_f32(x + i))); + } +#elif defined(__riscv_v_intrinsic) + for (int vl; i < n; i += vl) { + vl = __riscv_vsetvl_e32m2(n - i); + vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl); + vfloat32m2_t vy = ggml_v_silu_m2(vx, vl); + __riscv_vse32_v_f32m2(&y[i], vy, vl); + } +#endif + for (; i < n; ++i) { + y[i] = ggml_silu_f32(x[i]); + } +} + +void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float * g) { + int i = 0; +#if defined(__AVX512F__) && defined(__AVX512DQ__) + for (; i + 15 < n; i += 16) { + _mm512_storeu_ps(y + i, _mm512_mul_ps(ggml_v_silu(_mm512_loadu_ps(x + i)), _mm512_loadu_ps(g + i))); + } +#elif defined(__AVX2__) && defined(__FMA__) + for (; i + 7 < n; i += 8) { + _mm256_storeu_ps(y + i, _mm256_mul_ps(ggml_v_silu(_mm256_loadu_ps(x + i)), _mm256_loadu_ps(g + i))); + } +#elif defined(__SSE2__) + for (; i + 3 < n; i += 4) { + _mm_storeu_ps(y + i, _mm_mul_ps(ggml_v_silu(_mm_loadu_ps(x + i)), _mm_loadu_ps(g + i))); + } +#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + const int vlen = svcntw(); + for (; i < n; i += vlen) { + const svbool_t pg = svwhilelt_b32_s32(i, n); + svst1_f32(pg, y + i, svmul_f32_x(pg, ggml_v_silu(pg, svld1_f32(pg, x + i)), svld1_f32(pg, g + i))); + } +#elif defined(__ARM_NEON) && defined(__aarch64__) + for (; i + 3 < n; i += 4) { + vst1q_f32(y + i, vmulq_f32(ggml_v_silu(vld1q_f32(x + i)), vld1q_f32(g + i))); + } +#elif defined(__riscv_v_intrinsic) + for (int vl; i < n; i += vl) { + vl = __riscv_vsetvl_e32m2(n - i); + vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl); + vfloat32m2_t vg = __riscv_vle32_v_f32m2(&g[i], vl); + vfloat32m2_t vy = __riscv_vfmul_vv_f32m2(ggml_v_silu_m2(vx, vl), vg, vl); + __riscv_vse32_v_f32m2(&y[i], vy, vl); + } +#endif + for (; i < n; ++i) { + y[i] = ggml_silu_f32(x[i]) * g[i]; + } +} + +ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean) { + int i = 0; + ggml_float sum = 0; +// TODO: optimize to process the remaining elements in groups using the smaller vector sizes from AVX2 and SSE +// ref: https://github.com/ggml-org/llama.cpp/pull/15953#pullrequestreview-3310928344 +#if defined(__AVX512F__) && defined(__AVX512DQ__) + for (; i + 15 < n; i += 16) { + __m512 val = _mm512_sub_ps(_mm512_loadu_ps(x + i), + _mm512_set1_ps(mean)); + _mm512_storeu_ps(y + i, val); + sum += (ggml_float)_mm512_reduce_add_ps(_mm512_mul_ps(val, val)); + } +#elif defined(__AVX2__) && defined(__FMA__) + for (; i + 7 < n; i += 8) { + __m256 val = _mm256_sub_ps(_mm256_loadu_ps(x + i), + _mm256_set1_ps(mean)); + _mm256_storeu_ps(y + i, val); + val = _mm256_mul_ps(val,val); + __m128 val2 = _mm_add_ps(_mm256_extractf128_ps(val, 1), + _mm256_castps256_ps128(val)); + val2 = _mm_add_ps(val2, _mm_movehl_ps(val2, val2)); + val2 = _mm_add_ss(val2, _mm_movehdup_ps(val2)); + sum += (ggml_float)_mm_cvtss_f32(val2); + } +#elif defined(__SSE2__) + for (; i + 3 < n; i += 4) { + __m128 val = _mm_sub_ps(_mm_loadu_ps(x + i), + _mm_set1_ps(mean)); + _mm_storeu_ps(y + i, val); + val = _mm_mul_ps(val, val); +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) + val = _mm_add_ps(val, _mm_movehl_ps(val, val)); + val = _mm_add_ss(val, _mm_movehdup_ps(val)); +#else + __m128 tmp = _mm_shuffle_ps(val, val, _MM_SHUFFLE(2, 3, 0, 1)); + val = _mm_add_ps(val, tmp); + tmp = _mm_movehl_ps(tmp, val); + val = _mm_add_ss(val, tmp); +#endif // __AVX__ || __AVX2__ || __AVX512F__ + sum += (ggml_float)_mm_cvtss_f32(val); + } +#elif defined(__ARM_NEON) && defined(__aarch64__) + for (; i + 3 < n; i += 4) { + float32x4_t val = vsubq_f32(vld1q_f32(x + i), + vdupq_n_f32(mean)); + vst1q_f32(y + i, val); + val = vmulq_f32(val, val); + sum += (ggml_float)vaddvq_f32(val); + } +#elif defined(__VXE__) || defined(__VXE2__) + for (; i + 3 < n; i += 4) { + float32x4_t val = vec_sub(vec_xl(0, x + i), vec_splats(mean)); + vec_xst(val, 0, y + i); + val = vec_mul(val, val); + sum += (ggml_float)vec_hsum_f32x4(val); + } +#elif defined(__riscv_v_intrinsic) + vfloat64m1_t vsum = __riscv_vfmv_v_f_f64m1(0, 1); + for (int vl; i < n; i += vl) { + vl = __riscv_vsetvl_e32m2(n - i); + vfloat32m2_t val = __riscv_vfsub_vf_f32m2(__riscv_vle32_v_f32m2(&x[i], vl), mean, vl); + __riscv_vse32_v_f32m2(&y[i], val, vl); + val = __riscv_vfmul_vv_f32m2(val, val, vl); + vsum = __riscv_vfwredusum_vs_f32m2_f64m1(val, vsum, vl); + } + sum = (ggml_float)__riscv_vfmv_f_s_f64m1_f64(vsum); +#endif + for (; i < n; ++i) { + float val = x[i] - mean; + y[i] = val; + val *= val; + sum += (ggml_float)val; + } + return sum/n; +} + +ggml_float ggml_vec_soft_max_f32(const int n, float * y, const float * x, float max) { + int i = 0; + ggml_float sum = 0; +#if defined(__AVX512F__) && defined(__AVX512DQ__) + for (; i + 15 < n; i += 16) { + __m512 val = ggml_v_expf(_mm512_sub_ps(_mm512_loadu_ps(x + i), + _mm512_set1_ps(max))); + _mm512_storeu_ps(y + i, val); + sum += (ggml_float)_mm512_reduce_add_ps(val); + } +#elif defined(__AVX2__) && defined(__FMA__) + for (; i + 7 < n; i += 8) { + __m256 val = ggml_v_expf(_mm256_sub_ps(_mm256_loadu_ps(x + i), + _mm256_set1_ps(max))); + _mm256_storeu_ps(y + i, val); + __m128 val2 = _mm_add_ps(_mm256_extractf128_ps(val, 1), + _mm256_castps256_ps128(val)); + val2 = _mm_add_ps(val2, _mm_movehl_ps(val2, val2)); + val2 = _mm_add_ss(val2, _mm_movehdup_ps(val2)); + sum += (ggml_float)_mm_cvtss_f32(val2); + } +#elif defined(__SSE2__) + for (; i + 3 < n; i += 4) { + __m128 val = ggml_v_expf(_mm_sub_ps(_mm_loadu_ps(x + i), + _mm_set1_ps(max))); + _mm_storeu_ps(y + i, val); +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) + val = _mm_add_ps(val, _mm_movehl_ps(val, val)); + val = _mm_add_ss(val, _mm_movehdup_ps(val)); +#else + __m128 tmp = _mm_shuffle_ps(val, val, _MM_SHUFFLE(2, 3, 0, 1)); + val = _mm_add_ps(val, tmp); + tmp = _mm_movehl_ps(tmp, val); + val = _mm_add_ss(val, tmp); +#endif + sum += (ggml_float)_mm_cvtss_f32(val); + } +#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) + const int vlen = svcntw(); + for (; i < n; i += vlen) { + const svbool_t pg = svwhilelt_b32_s32(i, n); + svfloat32_t val = ggml_v_expf(pg, svsub_f32_x(pg, svld1_f32(pg, x + i), + svdup_n_f32_x(pg, max))); + svst1_f32(pg, y + i, val); + sum += (ggml_float)svaddv_f32(pg, val); + } +#elif defined(__ARM_NEON) && defined(__aarch64__) + for (; i + 3 < n; i += 4) { + float32x4_t val = ggml_v_expf(vsubq_f32(vld1q_f32(x + i), + vdupq_n_f32(max))); + vst1q_f32(y + i, val); + sum += (ggml_float)vaddvq_f32(val); + } +#elif defined(__riscv_v_intrinsic) + vfloat64m1_t vsum = __riscv_vfmv_v_f_f64m1(0, 1); + for (int avl; i < n; i += avl) { + avl = __riscv_vsetvl_e32m2(n - i); + vfloat32m2_t val = ggml_v_expf_m2(__riscv_vfsub_vf_f32m2(__riscv_vle32_v_f32m2(&x[i], avl), max, avl), avl); + __riscv_vse32_v_f32m2(&y[i], val, avl); + vsum = __riscv_vfwredusum_vs_f32m2_f64m1(val, vsum, avl); + } + return (ggml_float)__riscv_vfmv_f_s_f64m1_f64(vsum); +#endif + for (; i < n; ++i) { + float val = expf(x[i] - max); + sum += (ggml_float)val; + y[i] = val; + } + return sum; +} + +ggml_float ggml_vec_log_soft_max_f32(const int n, float * y, const float * x, float max) { + // log(soft_max) = log(soft_max_i / soft_max_sum) = log(soft_max_i) - log(soft_max_sum) = (logit_i - max) - log(soft_max_i) + + int i = 0; + ggml_float sum = 0; + for (; i < n; ++i) { + float val = x[i] - max; + y[i] = val; + sum += (ggml_float)expf(val); } + return sum = (ggml_float)logf(sum); } diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h deleted file mode 100644 index 3fa1443abc48..000000000000 --- a/ggml/src/ggml-cpu/ops.h +++ /dev/null @@ -1,117 +0,0 @@ -#pragma once - -#include "ggml.h" - -// -// cache line -// - -#if defined(__cpp_lib_hardware_interference_size) -#define CACHE_LINE_SIZE std::hardware_destructive_interference_size -#else -#if defined(__POWER9_VECTOR__) -#define CACHE_LINE_SIZE 128 -#elif defined(__VXE__) || defined(__VXE2__) -#define CACHE_LINE_SIZE 256 -#else -#define CACHE_LINE_SIZE 64 -#endif -#endif - -static const size_t CACHE_LINE_SIZE_F32 = CACHE_LINE_SIZE/sizeof(float); - -// Work buffer size for im2col operations in CONV2D -#define GGML_IM2COL_WORK_SIZE (16 * 1024 * 1024) - -#ifdef __cplusplus -extern "C" { -#endif - -void ggml_compute_forward_dup(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_add(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_add_id(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_add1(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_acc(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_sum(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_sum_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_cumsum(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_mean(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_argmax(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_count_equal(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_repeat(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_repeat_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_concat(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_silu_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_rms_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_rms_norm_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_group_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_l2_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_out_prod(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_scale(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_set(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_cpy(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_cont(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_get_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_get_rows_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_set_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_diag(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_diag_mask_inf(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_diag_mask_zero(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_soft_max(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_soft_max_ext_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_rope(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_rope_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_clamp(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_conv_transpose_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_im2col(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_im2col_back_f32(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_im2col_3d(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_conv_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_conv_3d(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_conv_transpose_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_conv_2d_dw(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_pool_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_pool_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_pool_2d_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_upscale(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_pad(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_pad_reflect_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_roll(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_arange(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_timestep_embedding(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_argsort(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_top_k(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_leaky_relu(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_fill(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_flash_attn_ext(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_flash_attn_back( - const struct ggml_compute_params * params, - const bool masked, - struct ggml_tensor * dst); -void ggml_compute_forward_ssm_conv(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_ssm_scan(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_win_part(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_win_unpart(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_unary(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_glu(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_get_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_add_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_rwkv_wkv6(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_rwkv_wkv7(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_custom(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_cross_entropy_loss(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_cross_entropy_loss_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_opt_step_adamw(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_opt_step_sgd(const struct ggml_compute_params * params, struct ggml_tensor * dst); -#ifdef __cplusplus -} -#endif diff --git a/ggml/src/ggml-cpu/quants.c b/ggml/src/ggml-cpu/quants.c deleted file mode 100644 index e5f9a4083f9c..000000000000 --- a/ggml/src/ggml-cpu/quants.c +++ /dev/null @@ -1,1288 +0,0 @@ -#define GGML_COMMON_IMPL_C -#include "ggml-common.h" - -#include "ggml-cpu-impl.h" -#include "simd-mappings.h" -#include "ggml-quants.h" -#include "quants.h" - -#include "arch-fallback.h" - -#include -#include -#include -#include // for qsort -#include // for GGML_ASSERT - -#define GROUP_MAX_EPS 1e-15f -#define GROUP_MAX_EPS_IQ3_XXS 1e-8f -#define GROUP_MAX_EPS_IQ2_S 1e-8f -#define GROUP_MAX_EPS_IQ1_M 1e-7f -#define GROUP_MAX_EPS_IQ1_S 1e-12f - -#define UNUSED GGML_UNUSED - -void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { - quantize_row_q1_0_ref(x, y, k); -} - -void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { - quantize_row_q4_0_ref(x, y, k); -} - -void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { - quantize_row_q4_1_ref(x, y, k); -} - -void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { - quantize_row_q5_0_ref(x, y, k); -} - -void quantize_row_q5_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { - quantize_row_q5_1_ref(x, y, k); -} - -void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { - quantize_row_q8_0_ref(x, y, k); -} - -void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { - quantize_row_q8_1_ref(x, y, k); -} - -void quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { - quantize_row_mxfp4_ref(x, y, k); -} - -void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { - quantize_row_nvfp4_ref(x, y, k); -} - -// -// 2-6 bit quantization in super-blocks -// - -//========================- 2-bit (de)-quantization - -void quantize_row_q2_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - quantize_row_q2_K_ref(x, vy, k); -} - -//========================= 3-bit (de)-quantization - -void quantize_row_q3_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - quantize_row_q3_K_ref(x, vy, k); -} - -// ====================== 4-bit (de)-quantization - -void quantize_row_q4_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - assert(k % QK_K == 0); - block_q4_K * GGML_RESTRICT y = vy; - quantize_row_q4_K_ref(x, y, k); -} - -// ====================== 5-bit (de)-quantization - -void quantize_row_q5_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - assert(k % QK_K == 0); - block_q5_K * GGML_RESTRICT y = vy; - quantize_row_q5_K_ref(x, y, k); -} - -// ====================== 6-bit (de)-quantization - -void quantize_row_q6_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - assert(k % QK_K == 0); - block_q6_K * GGML_RESTRICT y = vy; - quantize_row_q6_K_ref(x, y, k); -} - -// ====================== Ternary (de)-quantization (BitNet b1.58 and TriLMs) - -void quantize_row_tq1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - assert(k % QK_K == 0); - block_tq1_0 * GGML_RESTRICT y = vy; - quantize_row_tq1_0_ref(x, y, k); -} - -void quantize_row_tq2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - assert(k % QK_K == 0); - block_tq2_0 * GGML_RESTRICT y = vy; - quantize_row_tq2_0_ref(x, y, k); -} - -//===================================== Q8_K ============================================== - -void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { - quantize_row_q8_K_ref(x, y, k); -} - -//===================================== Dot products ================================= - -void ggml_vec_dot_q1_0_q8_0_generic(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) { - const int qk = QK1_0; - const int nb = n / qk; - - assert(n % qk == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_q1_0 * GGML_RESTRICT x = vx; - const block_q8_0 * GGML_RESTRICT y = vy; - - float sumf = 0.0; - - for (int i = 0; i < nb; i++) { - const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d); - - float sumi = 0.0f; - - for (int k = 0; k < 4; k++) { - const block_q8_0 * GGML_RESTRICT yb = &y[i * 4 + k]; - const float d1 = GGML_CPU_FP16_TO_FP32(yb->d); - int sumi_block = 0; - - const uint8_t * GGML_RESTRICT bits = &x[i].qs[k * 4]; - const int8_t * GGML_RESTRICT qy = yb->qs; - - for (int b = 0; b < 4; ++b, qy += 8) { - const unsigned mask = bits[b]; - sumi_block += ((mask & 0x01) ? qy[0] : -qy[0]) - + ((mask & 0x02) ? qy[1] : -qy[1]) - + ((mask & 0x04) ? qy[2] : -qy[2]) - + ((mask & 0x08) ? qy[3] : -qy[3]) - + ((mask & 0x10) ? qy[4] : -qy[4]) - + ((mask & 0x20) ? qy[5] : -qy[5]) - + ((mask & 0x40) ? qy[6] : -qy[6]) - + ((mask & 0x80) ? qy[7] : -qy[7]); - } - - sumi += d1 * sumi_block; - } - - sumf += d0 * sumi; - } - - *s = sumf; -} - - -void ggml_vec_dot_q4_0_q8_0_generic(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) { - const int qk = QK8_0; - const int nb = n / qk; - - assert(n % qk == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_q4_0 * GGML_RESTRICT x = vx; - const block_q8_0 * GGML_RESTRICT y = vy; - - int ib = 0; - float sumf = 0; - - for (; ib < nb; ++ib) { - int sumi0 = 0; - int sumi1 = 0; - - for (int j = 0; j < qk/2; ++j) { - const int v0 = (x[ib].qs[j] & 0x0F) - 8; - const int v1 = (x[ib].qs[j] >> 4) - 8; - - sumi0 += (v0 * y[ib].qs[j]); - sumi1 += (v1 * y[ib].qs[j + qk/2]); - } - - int sumi = sumi0 + sumi1; - sumf += sumi*GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d); - } - - *s = sumf; -} - -// TODO: add WASM SIMD -void ggml_vec_dot_q4_1_q8_1_generic(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) { - const int qk = QK8_1; - const int nb = n / qk; - - assert(n % qk == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_q4_1 * GGML_RESTRICT x = vx; - const block_q8_1 * GGML_RESTRICT y = vy; - - int ib = 0; - float sumf = 0; - - for (; ib < nb; ++ib) { - int sumi0 = 0; - int sumi1 = 0; - - for (int j = 0; j < qk/2; ++j) { - const int v0 = (x[ib].qs[j] & 0x0F); - const int v1 = (x[ib].qs[j] >> 4); - - sumi0 += (v0 * y[ib].qs[j]); - sumi1 += (v1 * y[ib].qs[j + qk/2]); - } - - int sumi = sumi0 + sumi1; - sumf += (GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d))*sumi + GGML_CPU_FP16_TO_FP32(x[ib].m)*GGML_CPU_FP16_TO_FP32(y[ib].s); - } - - *s = sumf; -} - -void ggml_vec_dot_mxfp4_q8_0_generic(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) { - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - assert(n % QK_MXFP4 == 0); - static_assert(QK_MXFP4 == QK8_0, "QK_MXFP4 and QK8_0 must be the same"); - - const block_mxfp4 * GGML_RESTRICT x = vx; - const block_q8_0 * GGML_RESTRICT y = vy; - - const int nb = n / QK_MXFP4; - - int ib = 0; - float sumf = 0; - - for (; ib < nb; ++ib) { - const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_E8M0_TO_FP32_HALF(x[ib].e); - - int sumi1 = 0; - int sumi2 = 0; - for (int j = 0; j < QK_MXFP4/2; ++j) { - sumi1 += y[ib].qs[j + 0] * kvalues_mxfp4[x[ib].qs[j] & 0xf]; - sumi2 += y[ib].qs[j + QK_MXFP4/2] * kvalues_mxfp4[x[ib].qs[j] >> 4]; - } - sumf += d * (sumi1 + sumi2); - } - *s = sumf; -} - -// NVFP4: super-block of 64 elements = 4 sub-blocks of 16 = 2 q8_0 blocks -void ggml_vec_dot_nvfp4_q8_0_generic(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) { - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - assert(n % QK_NVFP4 == 0); - - const block_nvfp4 * GGML_RESTRICT x = vx; - const block_q8_0 * GGML_RESTRICT y = vy; - - const int nb = n / QK_NVFP4; - - float sumf = 0; - - for (int ib = 0; ib < nb; ++ib) { - for (int s_idx = 0; s_idx < 4; ++s_idx) { - const float d = ggml_ue4m3_to_fp32(x[ib].d[s_idx]); - const int q8_block = s_idx / 2; - const int q8_off = (s_idx % 2) * QK_NVFP4_SUB; - const float dy = GGML_CPU_FP16_TO_FP32(y[2*ib + q8_block].d); - - int sumi_lo = 0, sumi_hi = 0; - for (int j = 0; j < QK_NVFP4_SUB/2; ++j) { - const uint8_t qv = x[ib].qs[s_idx*(QK_NVFP4_SUB/2) + j]; - sumi_lo += y[2*ib + q8_block].qs[q8_off + j + 0] * kvalues_mxfp4[qv & 0xf]; - sumi_hi += y[2*ib + q8_block].qs[q8_off + j + QK_NVFP4_SUB/2] * kvalues_mxfp4[qv >> 4]; - } - - sumf += dy * d * (sumi_lo + sumi_hi); - } - } - *s = sumf; -} - -void ggml_vec_dot_q5_0_q8_0_generic(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) { - const int qk = QK8_0; - const int nb = n / qk; - - int ib = 0; - float sumf = 0; - - assert(n % qk == 0); - assert(qk == QK5_0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_q5_0 * GGML_RESTRICT x = vx; - const block_q8_0 * GGML_RESTRICT y = vy; - - for (; ib < nb; ++ib) { - uint32_t qh; - memcpy(&qh, x[ib].qh, sizeof(qh)); - - int sumi0 = 0; - int sumi1 = 0; - - for (int j = 0; j < qk/2; ++j) { - const uint8_t xh_0 = ((qh & (1u << (j + 0 ))) >> (j + 0 )) << 4; - const uint8_t xh_1 = ((qh & (1u << (j + 16))) >> (j + 12)); - - const int32_t x0 = (int8_t)(((x[ib].qs[j] & 0x0F) | xh_0) - 16); - const int32_t x1 = (int8_t)(((x[ib].qs[j] >> 4) | xh_1) - 16); - - sumi0 += (x0 * y[ib].qs[j]); - sumi1 += (x1 * y[ib].qs[j + qk/2]); - } - - int sumi = sumi0 + sumi1; - sumf += (GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d)) * sumi; - } - - *s = sumf; -} - -void ggml_vec_dot_q5_1_q8_1_generic(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) { - const int qk = QK8_1; - const int nb = n / qk; - - int ib = 0; - float sumf = 0; - - assert(n % qk == 0); - assert(qk == QK5_1); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_q5_1 * GGML_RESTRICT x = vx; - const block_q8_1 * GGML_RESTRICT y = vy; - - for (; ib < nb; ++ib) { - uint32_t qh; - memcpy(&qh, x[ib].qh, sizeof(qh)); - - int sumi0 = 0; - int sumi1 = 0; - - for (int j = 0; j < qk/2; ++j) { - const uint8_t xh_0 = ((qh >> (j + 0)) << 4) & 0x10; - const uint8_t xh_1 = ((qh >> (j + 12)) ) & 0x10; - - const int32_t x0 = (x[ib].qs[j] & 0xF) | xh_0; - const int32_t x1 = (x[ib].qs[j] >> 4) | xh_1; - - sumi0 += (x0 * y[ib].qs[j]); - sumi1 += (x1 * y[ib].qs[j + qk/2]); - } - - int sumi = sumi0 + sumi1; - sumf += (GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d))*sumi + GGML_CPU_FP16_TO_FP32(x[ib].m)*GGML_CPU_FP16_TO_FP32(y[ib].s); - } - - *s = sumf; -} - -void ggml_vec_dot_q8_0_q8_0_generic(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) { - const int qk = QK8_0; - const int nb = n / qk; - - assert(n % qk == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_q8_0 * GGML_RESTRICT x = vx; - const block_q8_0 * GGML_RESTRICT y = vy; - - int ib = 0; - float sumf = 0; - - for (; ib < nb; ++ib) { - int sumi = 0; - - for (int j = 0; j < qk; j++) { - sumi += x[ib].qs[j]*y[ib].qs[j]; - } - - sumf += sumi*(GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d)); - } - - *s = sumf; -} - -void ggml_vec_dot_tq1_0_q8_K_generic(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) { - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_tq1_0 * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - const uint8_t pow3[6] = {1, 3, 9, 27, 81, 243}; - - float sumf = 0.0f; - - for (int i = 0; i < nb; ++i) { - int sum = 0; - - for (size_t j = 0; j < sizeof(x->qs) - sizeof(x->qs) % 32; j += 32) { - for (size_t l = 0; l < 5; ++l) { - for (size_t m = 0; m < 32; ++m) { - uint8_t q = x[i].qs[j + m] * pow3[l]; - uint16_t xi = ((uint16_t) q * 3) >> 8; - sum += (xi - 1) * y[i].qs[j*5 + l*32 + m]; - } - } - } - for (size_t j = sizeof(x->qs) - sizeof(x->qs) % 32; j < sizeof(x->qs); j += 16) { - for (size_t l = 0; l < 5; ++l) { - for (size_t m = 0; m < 16; ++m) { - uint8_t q = x[i].qs[j + m] * pow3[l]; - uint16_t xi = ((uint16_t) q * 3) >> 8; - sum += (xi - 1) * y[i].qs[j*5 + l*16 + m]; - } - } - } - - for (size_t l = 0; l < 4; ++l) { - for (size_t j = 0; j < sizeof(x->qh); ++j) { - uint8_t q = x[i].qh[j] * pow3[l]; - uint16_t xi = ((uint16_t) q * 3) >> 8; - sum += (xi - 1) * y[i].qs[sizeof(x->qs)*5 + l*sizeof(x->qh) + j]; - } - } - - sumf += (float) sum * (GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d); - } - - *s = sumf; -} - -void ggml_vec_dot_tq2_0_q8_K_generic(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) { - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_tq2_0 * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - float sumf = 0.0f; - - for (int i = 0; i < nb; ++i) { - int32_t sumi = 0; - - for (size_t j = 0; j < sizeof(x->qs); j += 32) { - for (size_t l = 0; l < 4; ++l) { - for (size_t k = 0; k < 32; ++k) { - sumi += y[i].qs[j*4 + l*32 + k] * (((x[i].qs[j + k] >> (l*2)) & 3) - 1); - } - } - } - - const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); - - sumf += (float) sumi * d; - } - - *s = sumf; -} - -void ggml_vec_dot_q2_K_q8_K_generic(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) { - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_q2_K * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - float sumf = 0; - - for (int i = 0; i < nb; ++i) { - - const uint8_t * q2 = x[i].qs; - const int8_t * q8 = y[i].qs; - const uint8_t * sc = x[i].scales; - - int summs = 0; - for (int j = 0; j < 16; ++j) { - summs += y[i].bsums[j] * (sc[j] >> 4); - } - - const float dall = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); - const float dmin = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); - - int isum = 0; - int is = 0; - int d; - for (int k = 0; k < QK_K/128; ++k) { - int shift = 0; - for (int j = 0; j < 4; ++j) { - d = sc[is++] & 0xF; - int isuml = 0; - for (int l = 0; l < 16; ++l) isuml += q8[l] * ((q2[l] >> shift) & 3); - isum += d * isuml; - d = sc[is++] & 0xF; - isuml = 0; - for (int l = 16; l < 32; ++l) isuml += q8[l] * ((q2[l] >> shift) & 3); - isum += d * isuml; - shift += 2; - q8 += 32; - } - q2 += 32; - } - sumf += dall * isum - dmin * summs; - } - *s = sumf; -} - -void ggml_vec_dot_q3_K_q8_K_generic(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) { - assert(n % QK_K == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const uint32_t kmask1 = 0x03030303; - const uint32_t kmask2 = 0x0f0f0f0f; - - const block_q3_K * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - // scalar version - // This function is written like this so the compiler can manage to vectorize most of it - // Using -Ofast, GCC and clang manage to produce code that is within a factor of 2 or so from the - // manually vectorized version above. Every other version I tried would run at least 4 times slower. - // The ideal situation would be if we could just write the code once, and the compiler would - // automatically produce the best possible set of machine instructions, instead of us having to manually - // write vectorized versions for AVX, ARM_NEON, etc. - - int8_t aux8[QK_K]; - int16_t aux16[8]; - float sums [8]; - int32_t aux32[8]; - memset(sums, 0, 8*sizeof(float)); - - uint32_t auxs[4]; - const int8_t * scales = (const int8_t*)auxs; - - float sumf = 0; - for (int i = 0; i < nb; ++i) { - const uint8_t * GGML_RESTRICT q3 = x[i].qs; - const uint8_t * GGML_RESTRICT hm = x[i].hmask; - const int8_t * GGML_RESTRICT q8 = y[i].qs; - memset(aux32, 0, 8*sizeof(int32_t)); - int8_t * GGML_RESTRICT a = aux8; - uint8_t m = 1; - for (int j = 0; j < QK_K; j += 128) { - for (int l = 0; l < 32; ++l) a[l] = q3[l] & 3; - for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4); - a += 32; m <<= 1; - for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 2) & 3; - for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4); - a += 32; m <<= 1; - for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 4) & 3; - for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4); - a += 32; m <<= 1; - for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 6) & 3; - for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4); - a += 32; m <<= 1; - q3 += 32; - } - a = aux8; - - memcpy(auxs, x[i].scales, 12); - uint32_t tmp = auxs[2]; - auxs[2] = ((auxs[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); - auxs[3] = ((auxs[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); - auxs[0] = (auxs[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); - auxs[1] = (auxs[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); - for (int j = 0; j < QK_K/16; ++j) { - for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; - for (int l = 0; l < 8; ++l) aux32[l] += (scales[j] - 32) * aux16[l]; - q8 += 8; a += 8; - for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; - for (int l = 0; l < 8; ++l) aux32[l] += (scales[j] - 32) * aux16[l]; - q8 += 8; a += 8; - } - const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; - for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l]; - } - for (int l = 0; l < 8; ++l) sumf += sums[l]; - *s = sumf; -} - -void ggml_vec_dot_q4_K_q8_K_generic(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) { - assert(n % QK_K == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_q4_K * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - static const uint32_t kmask1 = 0x3f3f3f3f; - static const uint32_t kmask2 = 0x0f0f0f0f; - static const uint32_t kmask3 = 0x03030303; - - uint32_t utmp[4]; - - const uint8_t * scales = (const uint8_t*)&utmp[0]; - const uint8_t * mins = (const uint8_t*)&utmp[2]; - - int8_t aux8[QK_K]; - int16_t aux16[8]; - float sums [8]; - int32_t aux32[8]; - memset(sums, 0, 8*sizeof(float)); - - float sumf = 0; - for (int i = 0; i < nb; ++i) { - const uint8_t * GGML_RESTRICT q4 = x[i].qs; - const int8_t * GGML_RESTRICT q8 = y[i].qs; - memset(aux32, 0, 8*sizeof(int32_t)); - int8_t * GGML_RESTRICT a = aux8; - for (int j = 0; j < QK_K/64; ++j) { - for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] & 0xF); - a += 32; - for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] >> 4); - a += 32; q4 += 32; - } - memcpy(utmp, x[i].scales, 12); - utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); - const uint32_t uaux = utmp[1] & kmask1; - utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); - utmp[2] = uaux; - utmp[0] &= kmask1; - - int sumi = 0; - for (int j = 0; j < QK_K/16; ++j) sumi += y[i].bsums[j] * mins[j/2]; - a = aux8; - int is = 0; - for (int j = 0; j < QK_K/32; ++j) { - int32_t scale = scales[is++]; - for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; - for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; - q8 += 8; a += 8; - for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; - for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; - q8 += 8; a += 8; - for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; - for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; - q8 += 8; a += 8; - for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; - for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; - q8 += 8; a += 8; - } - const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; - for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l]; - const float dmin = GGML_CPU_FP16_TO_FP32(x[i].dmin) * y[i].d; - sumf -= dmin * sumi; - } - for (int l = 0; l < 8; ++l) sumf += sums[l]; - *s = sumf; -} - -void ggml_vec_dot_q5_K_q8_K_generic(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) { - assert(n % QK_K == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_q5_K * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - static const uint32_t kmask1 = 0x3f3f3f3f; - static const uint32_t kmask2 = 0x0f0f0f0f; - static const uint32_t kmask3 = 0x03030303; - - uint32_t utmp[4]; - - const uint8_t * scales = (const uint8_t*)&utmp[0]; - const uint8_t * mins = (const uint8_t*)&utmp[2]; - - int8_t aux8[QK_K]; - int16_t aux16[8]; - float sums [8]; - int32_t aux32[8]; - memset(sums, 0, 8*sizeof(float)); - - float sumf = 0; - for (int i = 0; i < nb; ++i) { - const uint8_t * GGML_RESTRICT q4 = x[i].qs; - const uint8_t * GGML_RESTRICT hm = x[i].qh; - const int8_t * GGML_RESTRICT q8 = y[i].qs; - memset(aux32, 0, 8*sizeof(int32_t)); - int8_t * GGML_RESTRICT a = aux8; - uint8_t m = 1; - for (int j = 0; j < QK_K/64; ++j) { - for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] & 0xF); - for (int l = 0; l < 32; ++l) a[l] += (hm[l] & m ? 16 : 0); - a += 32; m <<= 1; - for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] >> 4); - for (int l = 0; l < 32; ++l) a[l] += (hm[l] & m ? 16 : 0); - a += 32; m <<= 1; - q4 += 32; - } - memcpy(utmp, x[i].scales, 12); - utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); - const uint32_t uaux = utmp[1] & kmask1; - utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); - utmp[2] = uaux; - utmp[0] &= kmask1; - - int sumi = 0; - for (int j = 0; j < QK_K/16; ++j) sumi += y[i].bsums[j] * mins[j/2]; - a = aux8; - int is = 0; - for (int j = 0; j < QK_K/32; ++j) { - int32_t scale = scales[is++]; - for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; - for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; - q8 += 8; a += 8; - for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; - for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; - q8 += 8; a += 8; - for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; - for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; - q8 += 8; a += 8; - for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; - for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; - q8 += 8; a += 8; - } - const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; - for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l]; - const float dmin = GGML_CPU_FP16_TO_FP32(x[i].dmin) * y[i].d; - sumf -= dmin * sumi; - } - for (int l = 0; l < 8; ++l) sumf += sums[l]; - *s = sumf; -} - -void ggml_vec_dot_q6_K_q8_K_generic(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) { - assert(n % QK_K == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_q6_K * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - int8_t aux8[QK_K]; - int16_t aux16[8]; - float sums [8]; - int32_t aux32[8]; - memset(sums, 0, 8*sizeof(float)); - - float sumf = 0; - for (int i = 0; i < nb; ++i) { - const uint8_t * GGML_RESTRICT q4 = x[i].ql; - const uint8_t * GGML_RESTRICT qh = x[i].qh; - const int8_t * GGML_RESTRICT q8 = y[i].qs; - memset(aux32, 0, 8*sizeof(int32_t)); - int8_t * GGML_RESTRICT a = aux8; - for (int j = 0; j < QK_K; j += 128) { - for (int l = 0; l < 32; ++l) { - a[l + 0] = (int8_t)((q4[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; - a[l + 32] = (int8_t)((q4[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; - a[l + 64] = (int8_t)((q4[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; - a[l + 96] = (int8_t)((q4[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; - } - a += 128; - q4 += 64; - qh += 32; - } - a = aux8; - int is = 0; - for (int j = 0; j < QK_K/16; ++j) { - int scale = x[i].scales[is++]; - for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; - for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; - q8 += 8; a += 8; - for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; - for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; - q8 += 8; a += 8; - } - const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; - for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l]; - } - for (int l = 0; l < 8; ++l) sumf += sums[l]; - *s = sumf; -} - -void ggml_vec_dot_iq2_xxs_q8_K_generic(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) { - assert(n % QK_K == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_iq2_xxs * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - uint32_t aux32[2]; - const uint8_t * aux8 = (const uint8_t *)aux32; - - float sumf = 0.f; - for (int i = 0; i < nb; ++i) { - const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; - const uint16_t * GGML_RESTRICT q2 = x[i].qs; - const int8_t * GGML_RESTRICT q8 = y[i].qs; - int32_t bsum = 0; - for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { - memcpy(aux32, q2, 2*sizeof(uint32_t)); - q2 += 4; - const uint32_t ls = 2*(aux32[1] >> 28) + 1; - int32_t sumi = 0; - for (int l = 0; l < 4; ++l) { - const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[l]); - const uint8_t signs = ksigns_iq2xs[(aux32[1] >> 7*l) & 127]; - for (int j = 0; j < 8; ++j) { - sumi += grid[j] * q8[j] * (signs & kmask_iq2xs[j] ? -1 : 1); - } - q8 += 8; - } - bsum += sumi * ls; - } - sumf += d * bsum; - } - *s = 0.125f * sumf; -} - -void ggml_vec_dot_iq2_xs_q8_K_generic(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) { - assert(n % QK_K == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_iq2_xs * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - float sumf = 0.f; - for (int i = 0; i < nb; ++i) { - const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; - const uint16_t * GGML_RESTRICT q2 = x[i].qs; - const uint8_t * GGML_RESTRICT sc = x[i].scales; - const int8_t * GGML_RESTRICT q8 = y[i].qs; - int32_t bsum = 0; - for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { - const uint16_t ls1 = 2*(sc[ib32] & 0xf) + 1; - const uint16_t ls2 = 2*(sc[ib32] >> 4) + 1; - int32_t sumi = 0; - for (int l = 0; l < 2; ++l) { - const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[l] & 511)); - const uint8_t signs = ksigns_iq2xs[q2[l] >> 9]; - for (int j = 0; j < 8; ++j) { - sumi += grid[j] * q8[j] * (signs & kmask_iq2xs[j] ? -1 : 1); - } - q8 += 8; - } - bsum += sumi * ls1; - sumi = 0; - for (int l = 2; l < 4; ++l) { - const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[l] & 511)); - const uint8_t signs = ksigns_iq2xs[q2[l] >> 9]; - for (int j = 0; j < 8; ++j) { - sumi += grid[j] * q8[j] * (signs & kmask_iq2xs[j] ? -1 : 1); - } - q8 += 8; - } - bsum += sumi * ls2; - q2 += 4; - } - sumf += d * bsum; - } - *s = 0.125f * sumf; -} - -void ggml_vec_dot_iq2_s_q8_K_generic(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) { - assert(n % QK_K == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_iq2_s * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - float sumf = 0; - for (int i = 0; i < nb; i++) { - - const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; - const int8_t * q8 = y[i].qs; - const uint8_t * qs = x[i].qs; - const uint8_t * qh = x[i].qh; - const uint8_t * signs = qs + QK_K/8; - - int bsum = 0; - for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { - int ls1 = 1 + 2*(x[i].scales[ib32] & 0xf); - int ls2 = 1 + 2*(x[i].scales[ib32] >> 4); - int sumi1 = 0, sumi2 = 0; - for (int l = 0; l < 2; ++l) { - const uint8_t * grid = (const uint8_t *)(iq2s_grid + (qs[l] | (qh[ib32] << (8-2*l) & 0x300))); - for (int j = 0; j < 8; ++j) { - sumi1 += q8[j] * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1 : 1); - } - q8 += 8; - } - for (int l = 2; l < 4; ++l) { - const uint8_t * grid = (const uint8_t *)(iq2s_grid + (qs[l] | (qh[ib32] << (8-2*l) & 0x300))); - for (int j = 0; j < 8; ++j) { - sumi2 += q8[j] * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1 : 1); - } - q8 += 8; - } - bsum += ls1 * sumi1 + ls2 * sumi2; - qs += 4; - signs += 4; - } - - sumf += d * bsum; - } - - *s = 0.125f * sumf; -} - -void ggml_vec_dot_iq3_xxs_q8_K_generic(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) { - assert(n % QK_K == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_iq3_xxs * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - uint32_t aux32; - - float sumf = 0.f; - for (int i = 0; i < nb; ++i) { - const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; - const uint8_t * GGML_RESTRICT q3 = x[i].qs; - const uint8_t * GGML_RESTRICT gas = x[i].qs + QK_K/4; - const int8_t * GGML_RESTRICT q8 = y[i].qs; - int32_t bsum = 0; - for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { - memcpy(&aux32, gas, sizeof(uint32_t)); gas += sizeof(uint32_t); - const uint32_t ls = 2*(aux32 >> 28) + 1; - int32_t sumi = 0; - for (int l = 0; l < 4; ++l) { - const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*l+0]); - const uint8_t * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*l+1]); - const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*l) & 127]; - for (int j = 0; j < 4; ++j) { - sumi += grid1[j] * q8[j+0] * (signs & kmask_iq2xs[j+0] ? -1 : 1); - sumi += grid2[j] * q8[j+4] * (signs & kmask_iq2xs[j+4] ? -1 : 1); - } - q8 += 8; - } - q3 += 8; - bsum += sumi * ls; - } - sumf += d * bsum; - } - *s = 0.25f * sumf; -} - -void ggml_vec_dot_iq3_s_q8_K_generic(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) { - assert(n % QK_K == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_iq3_s * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - float sumf = 0.f; - for (int i = 0; i < nb; ++i) { - const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; - const uint8_t * GGML_RESTRICT qs = x[i].qs; - const uint8_t * GGML_RESTRICT qh = x[i].qh; - const uint8_t * GGML_RESTRICT signs = x[i].signs; - const int8_t * GGML_RESTRICT q8 = y[i].qs; - int32_t bsum = 0; - for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { - const uint32_t ls1 = 2*(x[i].scales[ib32/2] & 0xf) + 1; - const uint32_t ls2 = 2*(x[i].scales[ib32/2] >> 4) + 1; - int32_t sumi = 0; - for (int l = 0; l < 4; ++l) { - const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[ib32+0] << (8-2*l)) & 256))); - const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[ib32+0] << (7-2*l)) & 256))); - for (int j = 0; j < 4; ++j) { - sumi += grid1[j] * q8[j+0] * (signs[l] & kmask_iq2xs[j+0] ? -1 : 1); - sumi += grid2[j] * q8[j+4] * (signs[l] & kmask_iq2xs[j+4] ? -1 : 1); - } - q8 += 8; - } - qs += 8; - signs += 4; - bsum += sumi * ls1; - sumi = 0; - for (int l = 0; l < 4; ++l) { - const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[ib32+1] << (8-2*l)) & 256))); - const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[ib32+1] << (7-2*l)) & 256))); - for (int j = 0; j < 4; ++j) { - sumi += grid1[j] * q8[j+0] * (signs[l] & kmask_iq2xs[j+0] ? -1 : 1); - sumi += grid2[j] * q8[j+4] * (signs[l] & kmask_iq2xs[j+4] ? -1 : 1); - } - q8 += 8; - } - qs += 8; - signs += 4; - bsum += sumi * ls2; - } - sumf += d * bsum; - } - *s = sumf; -} - -void ggml_vec_dot_iq1_s_q8_K_generic(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) { - assert(n % QK_K == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_iq1_s * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - float sumf = 0; - for (int i = 0; i < nb; i++) { - - const int8_t * q8 = y[i].qs; - const uint8_t * qs = x[i].qs; - const uint16_t * qh = x[i].qh; - - int sumi = 0, sumi1 = 0; - for (int ib = 0; ib < QK_K/32; ++ib) { - const int ls = 2*((qh[ib] >> 12) & 7) + 1; - const int delta = qh[ib] & 0x8000 ? -1 : 1; - int lsum = 0; - for (int l = 0; l < 4; ++l) { - const int8_t * grid = (const int8_t *)(iq1s_grid + (qs[l] | (((qh[ib] >> 3*l) & 7) << 8))); - for (int j = 0; j < 8; ++j) { - lsum += q8[j] * grid[j]; - } - q8 += 8; - } - sumi += ls * lsum; - sumi1 += ls * delta * (y[i].bsums[2*ib+0] + y[i].bsums[2*ib+1]); - qs += 4; - } - - sumf += GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d * (sumi + IQ1S_DELTA * sumi1); - } - - *s = sumf; -} - -void ggml_vec_dot_iq1_m_q8_K_generic(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) { - assert(n % QK_K == 0); - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - - const block_iq1_m * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - iq1m_scale_t scale; - - int sum1[2], sum2[2], delta[4]; - - float sumf = 0; - for (int i = 0; i < nb; i++) { - - const int8_t * q8 = y[i].qs; - const uint8_t * qs = x[i].qs; - const uint8_t * qh = x[i].qh; - const uint16_t * sc = (const uint16_t *)x[i].scales; - - scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); - - int sumi1 = 0, sumi2 = 0; - for (int ib = 0; ib < QK_K/32; ++ib) { - delta[0] = qh[0] & 0x08 ? -1 : 1; - delta[1] = qh[0] & 0x80 ? -1 : 1; - delta[2] = qh[1] & 0x08 ? -1 : 1; - delta[3] = qh[1] & 0x80 ? -1 : 1; - sum1[0] = sum1[1] = sum2[0] = sum2[1] = 0; - for (int l = 0; l < 4; ++l) { - const int8_t * grid = (const int8_t *)(iq1s_grid + (qs[l] | (((uint16_t)qh[l/2] << (8 - 4*(l%2))) & 0x700))); - int lsum1 = 0, lsum2 = 0; - for (int j = 0; j < 8; ++j) { - lsum1 += q8[j] * grid[j]; - lsum2 += q8[j]; - } - q8 += 8; - sum1[l/2] += lsum1; - sum2[l/2] += lsum2*delta[l]; - } - - const int ls1 = 2*((sc[ib/2] >> (6*(ib%2)+0)) & 0x7) + 1; - const int ls2 = 2*((sc[ib/2] >> (6*(ib%2)+3)) & 0x7) + 1; - - sumi1 += sum1[0] * ls1 + sum1[1] * ls2; - sumi2 += sum2[0] * ls1 + sum2[1] * ls2; - qs += 4; - qh += 2; - } - - sumf += GGML_CPU_FP16_TO_FP32(scale.f16) * y[i].d * (sumi1 + IQ1M_DELTA * sumi2); - } - - *s = sumf; -} - -void ggml_vec_dot_iq4_nl_q8_0_generic(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) { - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - assert(n % QK4_NL == 0); - static_assert(QK4_NL == QK8_0, "QK4_NL and QK8_0 must be the same"); - - const block_iq4_nl * GGML_RESTRICT x = vx; - const block_q8_0 * GGML_RESTRICT y = vy; - - const int nb = n / QK4_NL; - - int ib = 0; - float sumf = 0; - - for (; ib < nb; ++ib) { - const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_FP16_TO_FP32(x[ib].d); - int sumi1 = 0, sumi2 = 0; - for (int j = 0; j < QK4_NL/2; ++j) { - sumi1 += y[ib].qs[j+ 0] * kvalues_iq4nl[x[ib].qs[j] & 0xf]; - sumi2 += y[ib].qs[j+QK4_NL/2] * kvalues_iq4nl[x[ib].qs[j] >> 4]; - } - sumf += d * (sumi1 + sumi2); - } - *s = sumf; -} - -void ggml_vec_dot_iq4_xs_q8_K_generic(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) { - assert(nrc == 1); - UNUSED(nrc); - UNUSED(bx); - UNUSED(by); - UNUSED(bs); - assert(n % QK_K == 0); - - const block_iq4_xs * GGML_RESTRICT x = vx; - const block_q8_K * GGML_RESTRICT y = vy; - - const int nb = n / QK_K; - - float sumf = 0; - for (int ibl = 0; ibl < nb; ++ibl) { - const float d4d8 = GGML_CPU_FP16_TO_FP32(x[ibl].d) * y[ibl].d; - uint16_t h = x[ibl].scales_h; - const uint8_t * qs = x[ibl].qs; - const int8_t * q8 = y[ibl].qs; - for (int ib = 0; ib < QK_K/32; ib += 2) { - const uint8_t ls1 = (x[ibl].scales_l[ib/2] & 0xf) | ((h << 4) & 0x30); - const uint8_t ls2 = (x[ibl].scales_l[ib/2] >> 4) | ((h << 2) & 0x30); - h >>= 4; - const float d1 = d4d8*(ls1 - 32); - const float d2 = d4d8*(ls2 - 32); - int sumi1 = 0, sumi2 = 0; - for (int j = 0; j < 16; ++j) { - sumi1 += q8[j+ 0] * kvalues_iq4nl[qs[j] & 0xf]; - sumi2 += q8[j+16] * kvalues_iq4nl[qs[j] >> 4]; - } - sumf += d1 * (sumi1 + sumi2); - qs += 16; - q8 += 32; - sumi1 = sumi2 = 0; - for (int j = 0; j < 16; ++j) { - sumi1 += q8[j+ 0] * kvalues_iq4nl[qs[j] & 0xf]; - sumi2 += q8[j+16] * kvalues_iq4nl[qs[j] >> 4]; - } - sumf += d2 * (sumi1 + sumi2); - qs += 16; - q8 += 32; - } - } - *s = sumf; -} - -// ============================ 4-bit non-linear quants - -void quantize_row_iq4_nl(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { - assert(k % QK4_NL == 0); - quantize_row_iq4_nl_ref(x, y, k); -} - -void quantize_row_iq4_xs(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - quantize_iq4_xs(x, y, 1, k, NULL); -} diff --git a/ggml/src/ggml-cpu/arch/x86/quants.c b/ggml/src/ggml-cpu/quants.c.inc similarity index 71% rename from ggml/src/ggml-cpu/arch/x86/quants.c rename to ggml/src/ggml-cpu/quants.c.inc index 94b19b82bbc2..e738f76dc0ec 100644 --- a/ggml/src/ggml-cpu/arch/x86/quants.c +++ b/ggml/src/ggml-cpu/quants.c.inc @@ -1,16 +1,464 @@ #define GGML_COMMON_IMPL_C -#include "ggml-common.h" -#include "ggml-quants.h" -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "simd-mappings.h" +#include "ggml-common-defs.inc" + +#include "ggml-cpu-impl-defs.inc" +#include "simd-mappings-defs.inc" +#include "ggml-quants-defs.inc" +#define GGML_COMMON_DECL_C +#include "ggml-common-defs.inc" + +#include "ggml.h.inc" + +// GGML CPU internal header + +#ifdef __cplusplus +extern "C" { +#endif + +// Quantization +void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_q2_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q3_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q6_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_tq1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_tq2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void quantize_row_iq4_nl (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_iq4_xs (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +// Dot product +void ggml_vec_dot_q1_0_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 ggml_vec_dot_q4_0_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 ggml_vec_dot_q4_1_q8_1(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 ggml_vec_dot_q5_0_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 ggml_vec_dot_q5_1_q8_1(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 ggml_vec_dot_q8_0_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 ggml_vec_dot_mxfp4_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 ggml_vec_dot_nvfp4_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 ggml_vec_dot_q2_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 ggml_vec_dot_q3_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 ggml_vec_dot_q4_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 ggml_vec_dot_q5_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 ggml_vec_dot_q6_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 ggml_vec_dot_tq1_0_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 ggml_vec_dot_tq2_0_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 ggml_vec_dot_iq2_xxs_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 ggml_vec_dot_iq2_xs_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 ggml_vec_dot_iq2_s_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 ggml_vec_dot_iq3_xxs_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 ggml_vec_dot_iq1_s_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 ggml_vec_dot_iq1_m_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 ggml_vec_dot_iq4_nl_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 ggml_vec_dot_iq4_xs_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 ggml_vec_dot_iq3_s_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); + +// Generic implementation +void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void ggml_vec_dot_q1_0_q8_0_generic(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 ggml_vec_dot_q4_0_q8_0_generic(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 ggml_vec_dot_q4_1_q8_1_generic(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 ggml_vec_dot_q5_0_q8_0_generic(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 ggml_vec_dot_q5_1_q8_1_generic(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 ggml_vec_dot_q8_0_q8_0_generic(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 ggml_vec_dot_mxfp4_q8_0_generic(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 ggml_vec_dot_nvfp4_q8_0_generic(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 ggml_vec_dot_tq1_0_q8_K_generic(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 ggml_vec_dot_tq2_0_q8_K_generic(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 ggml_vec_dot_q2_K_q8_K_generic(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 ggml_vec_dot_q3_K_q8_K_generic(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 ggml_vec_dot_q4_K_q8_K_generic(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 ggml_vec_dot_q5_K_q8_K_generic(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 ggml_vec_dot_q6_K_q8_K_generic(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 ggml_vec_dot_iq2_xxs_q8_K_generic(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 ggml_vec_dot_iq2_xs_q8_K_generic(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 ggml_vec_dot_iq2_s_q8_K_generic(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 ggml_vec_dot_iq3_xxs_q8_K_generic(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 ggml_vec_dot_iq3_s_q8_K_generic(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 ggml_vec_dot_iq1_s_q8_K_generic(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 ggml_vec_dot_iq1_m_q8_K_generic(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 ggml_vec_dot_iq4_nl_q8_0_generic(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 ggml_vec_dot_iq4_xs_q8_K_generic(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); + +#ifdef __cplusplus +} +#endif + + + +// Rename `_generic` functions if no native implementation is available. +// This effectively selects the generic implementation. + +#if defined(GGML_CPU_GENERIC) +// quants.c +#define quantize_row_q8_0_generic quantize_row_q8_0 +#define quantize_row_q8_1_generic quantize_row_q8_1 +#define quantize_row_q8_K_generic quantize_row_q8_K +#define ggml_vec_dot_q4_0_q8_0_generic ggml_vec_dot_q4_0_q8_0 +#define ggml_vec_dot_q4_1_q8_1_generic ggml_vec_dot_q4_1_q8_1 +#define ggml_vec_dot_q5_0_q8_0_generic ggml_vec_dot_q5_0_q8_0 +#define ggml_vec_dot_q5_1_q8_1_generic ggml_vec_dot_q5_1_q8_1 +#define ggml_vec_dot_q8_0_q8_0_generic ggml_vec_dot_q8_0_q8_0 +#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K +#define ggml_vec_dot_q3_K_q8_K_generic ggml_vec_dot_q3_K_q8_K +#define ggml_vec_dot_q4_K_q8_K_generic ggml_vec_dot_q4_K_q8_K +#define ggml_vec_dot_q5_K_q8_K_generic ggml_vec_dot_q5_K_q8_K +#define ggml_vec_dot_q6_K_q8_K_generic ggml_vec_dot_q6_K_q8_K +#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K +#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K +#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K +#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K +#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K +#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +#define ggml_vec_dot_iq4_nl_q8_0_generic ggml_vec_dot_iq4_nl_q8_0 +#define ggml_vec_dot_iq4_xs_q8_K_generic ggml_vec_dot_iq4_xs_q8_K +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__aarch64__) || defined(__arm__) || defined(_M_ARM) || defined(_M_ARM64) +// repack.cpp +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64) +// quants.c +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__POWERPC__) || defined(__powerpc__) +// ref: https://github.com/ggml-org/llama.cpp/pull/14146#issuecomment-2972561679 +// quants.c +#define quantize_row_q8_K_generic quantize_row_q8_K +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__loongarch64) +// quants.c +#define quantize_row_q8_K_generic quantize_row_q8_K +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__riscv) +// quants.c +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +// repack.cpp +#define ggml_quantize_mat_q8_0_4x1_generic ggml_quantize_mat_q8_0_4x1 +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_K_4x1_generic ggml_quantize_mat_q8_K_4x1 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__s390x__) +// quants.c +#define quantize_row_q8_K_generic quantize_row_q8_K +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K +#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K +#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K +#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K +#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K +#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K +#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__wasm__) +// quants.c +#define ggml_vec_dot_q4_1_q8_1_generic ggml_vec_dot_q4_1_q8_1 +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K +#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K +#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K +#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K +#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K +#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +#define ggml_vec_dot_iq4_nl_q8_0_generic ggml_vec_dot_iq4_nl_q8_0 +#define ggml_vec_dot_iq4_xs_q8_K_generic ggml_vec_dot_iq4_xs_q8_K +#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#endif -#include "../../quants.h" -#include "../../ggml-cpu-impl.h" -#include #include #include +#include #include // for qsort #include // for GGML_ASSERT @@ -22,6 +470,1274 @@ #define UNUSED GGML_UNUSED +void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q1_0_ref(x, y, k); +} + +void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q4_0_ref(x, y, k); +} + +void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q4_1_ref(x, y, k); +} + +void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q5_0_ref(x, y, k); +} + +void quantize_row_q5_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q5_1_ref(x, y, k); +} + +void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q8_0_ref(x, y, k); +} + +void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q8_1_ref(x, y, k); +} + +void quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_mxfp4_ref(x, y, k); +} + +void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_nvfp4_ref(x, y, k); +} + +// +// 2-6 bit quantization in super-blocks +// + +//========================- 2-bit (de)-quantization + +void quantize_row_q2_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + quantize_row_q2_K_ref(x, vy, k); +} + +//========================= 3-bit (de)-quantization + +void quantize_row_q3_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + quantize_row_q3_K_ref(x, vy, k); +} + +// ====================== 4-bit (de)-quantization + +void quantize_row_q4_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_q4_K * GGML_RESTRICT y = vy; + quantize_row_q4_K_ref(x, y, k); +} + +// ====================== 5-bit (de)-quantization + +void quantize_row_q5_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_q5_K * GGML_RESTRICT y = vy; + quantize_row_q5_K_ref(x, y, k); +} + +// ====================== 6-bit (de)-quantization + +void quantize_row_q6_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_q6_K * GGML_RESTRICT y = vy; + quantize_row_q6_K_ref(x, y, k); +} + +// ====================== Ternary (de)-quantization (BitNet b1.58 and TriLMs) + +void quantize_row_tq1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_tq1_0 * GGML_RESTRICT y = vy; + quantize_row_tq1_0_ref(x, y, k); +} + +void quantize_row_tq2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_tq2_0 * GGML_RESTRICT y = vy; + quantize_row_tq2_0_ref(x, y, k); +} + +//===================================== Q8_K ============================================== + +void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q8_K_ref(x, y, k); +} + +//===================================== Dot products ================================= + +void ggml_vec_dot_q1_0_q8_0_generic(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) { + const int qk = QK1_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q1_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + float sumf = 0.0; + + for (int i = 0; i < nb; i++) { + const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d); + + float sumi = 0.0f; + + for (int k = 0; k < 4; k++) { + const block_q8_0 * GGML_RESTRICT yb = &y[i * 4 + k]; + const float d1 = GGML_CPU_FP16_TO_FP32(yb->d); + int sumi_block = 0; + + const uint8_t * GGML_RESTRICT bits = &x[i].qs[k * 4]; + const int8_t * GGML_RESTRICT qy = yb->qs; + + for (int b = 0; b < 4; ++b, qy += 8) { + const unsigned mask = bits[b]; + sumi_block += ((mask & 0x01) ? qy[0] : -qy[0]) + + ((mask & 0x02) ? qy[1] : -qy[1]) + + ((mask & 0x04) ? qy[2] : -qy[2]) + + ((mask & 0x08) ? qy[3] : -qy[3]) + + ((mask & 0x10) ? qy[4] : -qy[4]) + + ((mask & 0x20) ? qy[5] : -qy[5]) + + ((mask & 0x40) ? qy[6] : -qy[6]) + + ((mask & 0x80) ? qy[7] : -qy[7]); + } + + sumi += d1 * sumi_block; + } + + sumf += d0 * sumi; + } + + *s = sumf; +} + + +void ggml_vec_dot_q4_0_q8_0_generic(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) { + const int qk = QK8_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q4_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + int ib = 0; + float sumf = 0; + + for (; ib < nb; ++ib) { + int sumi0 = 0; + int sumi1 = 0; + + for (int j = 0; j < qk/2; ++j) { + const int v0 = (x[ib].qs[j] & 0x0F) - 8; + const int v1 = (x[ib].qs[j] >> 4) - 8; + + sumi0 += (v0 * y[ib].qs[j]); + sumi1 += (v1 * y[ib].qs[j + qk/2]); + } + + int sumi = sumi0 + sumi1; + sumf += sumi*GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d); + } + + *s = sumf; +} + +// TODO: add WASM SIMD +void ggml_vec_dot_q4_1_q8_1_generic(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) { + const int qk = QK8_1; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q4_1 * GGML_RESTRICT x = vx; + const block_q8_1 * GGML_RESTRICT y = vy; + + int ib = 0; + float sumf = 0; + + for (; ib < nb; ++ib) { + int sumi0 = 0; + int sumi1 = 0; + + for (int j = 0; j < qk/2; ++j) { + const int v0 = (x[ib].qs[j] & 0x0F); + const int v1 = (x[ib].qs[j] >> 4); + + sumi0 += (v0 * y[ib].qs[j]); + sumi1 += (v1 * y[ib].qs[j + qk/2]); + } + + int sumi = sumi0 + sumi1; + sumf += (GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d))*sumi + GGML_CPU_FP16_TO_FP32(x[ib].m)*GGML_CPU_FP16_TO_FP32(y[ib].s); + } + + *s = sumf; +} + +void ggml_vec_dot_mxfp4_q8_0_generic(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) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK_MXFP4 == 0); + static_assert(QK_MXFP4 == QK8_0, "QK_MXFP4 and QK8_0 must be the same"); + + const block_mxfp4 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + const int nb = n / QK_MXFP4; + + int ib = 0; + float sumf = 0; + + for (; ib < nb; ++ib) { + const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_E8M0_TO_FP32_HALF(x[ib].e); + + int sumi1 = 0; + int sumi2 = 0; + for (int j = 0; j < QK_MXFP4/2; ++j) { + sumi1 += y[ib].qs[j + 0] * kvalues_mxfp4[x[ib].qs[j] & 0xf]; + sumi2 += y[ib].qs[j + QK_MXFP4/2] * kvalues_mxfp4[x[ib].qs[j] >> 4]; + } + sumf += d * (sumi1 + sumi2); + } + *s = sumf; +} + +// NVFP4: super-block of 64 elements = 4 sub-blocks of 16 = 2 q8_0 blocks +void ggml_vec_dot_nvfp4_q8_0_generic(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) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK_NVFP4 == 0); + + const block_nvfp4 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + const int nb = n / QK_NVFP4; + + float sumf = 0; + + for (int ib = 0; ib < nb; ++ib) { + for (int s_idx = 0; s_idx < 4; ++s_idx) { + const float d = ggml_ue4m3_to_fp32(x[ib].d[s_idx]); + const int q8_block = s_idx / 2; + const int q8_off = (s_idx % 2) * QK_NVFP4_SUB; + const float dy = GGML_CPU_FP16_TO_FP32(y[2*ib + q8_block].d); + + int sumi_lo = 0, sumi_hi = 0; + for (int j = 0; j < QK_NVFP4_SUB/2; ++j) { + const uint8_t qv = x[ib].qs[s_idx*(QK_NVFP4_SUB/2) + j]; + sumi_lo += y[2*ib + q8_block].qs[q8_off + j + 0] * kvalues_mxfp4[qv & 0xf]; + sumi_hi += y[2*ib + q8_block].qs[q8_off + j + QK_NVFP4_SUB/2] * kvalues_mxfp4[qv >> 4]; + } + + sumf += dy * d * (sumi_lo + sumi_hi); + } + } + *s = sumf; +} + +void ggml_vec_dot_q5_0_q8_0_generic(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) { + const int qk = QK8_0; + const int nb = n / qk; + + int ib = 0; + float sumf = 0; + + assert(n % qk == 0); + assert(qk == QK5_0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q5_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + for (; ib < nb; ++ib) { + uint32_t qh; + memcpy(&qh, x[ib].qh, sizeof(qh)); + + int sumi0 = 0; + int sumi1 = 0; + + for (int j = 0; j < qk/2; ++j) { + const uint8_t xh_0 = ((qh & (1u << (j + 0 ))) >> (j + 0 )) << 4; + const uint8_t xh_1 = ((qh & (1u << (j + 16))) >> (j + 12)); + + const int32_t x0 = (int8_t)(((x[ib].qs[j] & 0x0F) | xh_0) - 16); + const int32_t x1 = (int8_t)(((x[ib].qs[j] >> 4) | xh_1) - 16); + + sumi0 += (x0 * y[ib].qs[j]); + sumi1 += (x1 * y[ib].qs[j + qk/2]); + } + + int sumi = sumi0 + sumi1; + sumf += (GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d)) * sumi; + } + + *s = sumf; +} + +void ggml_vec_dot_q5_1_q8_1_generic(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) { + const int qk = QK8_1; + const int nb = n / qk; + + int ib = 0; + float sumf = 0; + + assert(n % qk == 0); + assert(qk == QK5_1); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q5_1 * GGML_RESTRICT x = vx; + const block_q8_1 * GGML_RESTRICT y = vy; + + for (; ib < nb; ++ib) { + uint32_t qh; + memcpy(&qh, x[ib].qh, sizeof(qh)); + + int sumi0 = 0; + int sumi1 = 0; + + for (int j = 0; j < qk/2; ++j) { + const uint8_t xh_0 = ((qh >> (j + 0)) << 4) & 0x10; + const uint8_t xh_1 = ((qh >> (j + 12)) ) & 0x10; + + const int32_t x0 = (x[ib].qs[j] & 0xF) | xh_0; + const int32_t x1 = (x[ib].qs[j] >> 4) | xh_1; + + sumi0 += (x0 * y[ib].qs[j]); + sumi1 += (x1 * y[ib].qs[j + qk/2]); + } + + int sumi = sumi0 + sumi1; + sumf += (GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d))*sumi + GGML_CPU_FP16_TO_FP32(x[ib].m)*GGML_CPU_FP16_TO_FP32(y[ib].s); + } + + *s = sumf; +} + +void ggml_vec_dot_q8_0_q8_0_generic(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) { + const int qk = QK8_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q8_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + int ib = 0; + float sumf = 0; + + for (; ib < nb; ++ib) { + int sumi = 0; + + for (int j = 0; j < qk; j++) { + sumi += x[ib].qs[j]*y[ib].qs[j]; + } + + sumf += sumi*(GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d)); + } + + *s = sumf; +} + +void ggml_vec_dot_tq1_0_q8_K_generic(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) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_tq1_0 * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + const uint8_t pow3[6] = {1, 3, 9, 27, 81, 243}; + + float sumf = 0.0f; + + for (int i = 0; i < nb; ++i) { + int sum = 0; + + for (size_t j = 0; j < sizeof(x->qs) - sizeof(x->qs) % 32; j += 32) { + for (size_t l = 0; l < 5; ++l) { + for (size_t m = 0; m < 32; ++m) { + uint8_t q = x[i].qs[j + m] * pow3[l]; + uint16_t xi = ((uint16_t) q * 3) >> 8; + sum += (xi - 1) * y[i].qs[j*5 + l*32 + m]; + } + } + } + for (size_t j = sizeof(x->qs) - sizeof(x->qs) % 32; j < sizeof(x->qs); j += 16) { + for (size_t l = 0; l < 5; ++l) { + for (size_t m = 0; m < 16; ++m) { + uint8_t q = x[i].qs[j + m] * pow3[l]; + uint16_t xi = ((uint16_t) q * 3) >> 8; + sum += (xi - 1) * y[i].qs[j*5 + l*16 + m]; + } + } + } + + for (size_t l = 0; l < 4; ++l) { + for (size_t j = 0; j < sizeof(x->qh); ++j) { + uint8_t q = x[i].qh[j] * pow3[l]; + uint16_t xi = ((uint16_t) q * 3) >> 8; + sum += (xi - 1) * y[i].qs[sizeof(x->qs)*5 + l*sizeof(x->qh) + j]; + } + } + + sumf += (float) sum * (GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d); + } + + *s = sumf; +} + +void ggml_vec_dot_tq2_0_q8_K_generic(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) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_tq2_0 * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + float sumf = 0.0f; + + for (int i = 0; i < nb; ++i) { + int32_t sumi = 0; + + for (size_t j = 0; j < sizeof(x->qs); j += 32) { + for (size_t l = 0; l < 4; ++l) { + for (size_t k = 0; k < 32; ++k) { + sumi += y[i].qs[j*4 + l*32 + k] * (((x[i].qs[j + k] >> (l*2)) & 3) - 1); + } + } + } + + const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + + sumf += (float) sumi * d; + } + + *s = sumf; +} + +void ggml_vec_dot_q2_K_q8_K_generic(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) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q2_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + float sumf = 0; + + for (int i = 0; i < nb; ++i) { + + const uint8_t * q2 = x[i].qs; + const int8_t * q8 = y[i].qs; + const uint8_t * sc = x[i].scales; + + int summs = 0; + for (int j = 0; j < 16; ++j) { + summs += y[i].bsums[j] * (sc[j] >> 4); + } + + const float dall = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); + const float dmin = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); + + int isum = 0; + int is = 0; + int d; + for (int k = 0; k < QK_K/128; ++k) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + d = sc[is++] & 0xF; + int isuml = 0; + for (int l = 0; l < 16; ++l) isuml += q8[l] * ((q2[l] >> shift) & 3); + isum += d * isuml; + d = sc[is++] & 0xF; + isuml = 0; + for (int l = 16; l < 32; ++l) isuml += q8[l] * ((q2[l] >> shift) & 3); + isum += d * isuml; + shift += 2; + q8 += 32; + } + q2 += 32; + } + sumf += dall * isum - dmin * summs; + } + *s = sumf; +} + +void ggml_vec_dot_q3_K_q8_K_generic(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) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const uint32_t kmask1 = 0x03030303; + const uint32_t kmask2 = 0x0f0f0f0f; + + const block_q3_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + // scalar version + // This function is written like this so the compiler can manage to vectorize most of it + // Using -Ofast, GCC and clang manage to produce code that is within a factor of 2 or so from the + // manually vectorized version above. Every other version I tried would run at least 4 times slower. + // The ideal situation would be if we could just write the code once, and the compiler would + // automatically produce the best possible set of machine instructions, instead of us having to manually + // write vectorized versions for AVX, ARM_NEON, etc. + + int8_t aux8[QK_K]; + int16_t aux16[8]; + float sums [8]; + int32_t aux32[8]; + memset(sums, 0, 8*sizeof(float)); + + uint32_t auxs[4]; + const int8_t * scales = (const int8_t*)auxs; + + float sumf = 0; + for (int i = 0; i < nb; ++i) { + const uint8_t * GGML_RESTRICT q3 = x[i].qs; + const uint8_t * GGML_RESTRICT hm = x[i].hmask; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + memset(aux32, 0, 8*sizeof(int32_t)); + int8_t * GGML_RESTRICT a = aux8; + uint8_t m = 1; + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) a[l] = q3[l] & 3; + for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4); + a += 32; m <<= 1; + for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 2) & 3; + for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4); + a += 32; m <<= 1; + for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 4) & 3; + for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4); + a += 32; m <<= 1; + for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 6) & 3; + for (int l = 0; l < 32; ++l) a[l] -= (hm[l] & m ? 0 : 4); + a += 32; m <<= 1; + q3 += 32; + } + a = aux8; + + memcpy(auxs, x[i].scales, 12); + uint32_t tmp = auxs[2]; + auxs[2] = ((auxs[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + auxs[3] = ((auxs[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + auxs[0] = (auxs[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + auxs[1] = (auxs[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + for (int j = 0; j < QK_K/16; ++j) { + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += (scales[j] - 32) * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += (scales[j] - 32) * aux16[l]; + q8 += 8; a += 8; + } + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l]; + } + for (int l = 0; l < 8; ++l) sumf += sums[l]; + *s = sumf; +} + +void ggml_vec_dot_q4_K_q8_K_generic(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) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q4_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + uint32_t utmp[4]; + + const uint8_t * scales = (const uint8_t*)&utmp[0]; + const uint8_t * mins = (const uint8_t*)&utmp[2]; + + int8_t aux8[QK_K]; + int16_t aux16[8]; + float sums [8]; + int32_t aux32[8]; + memset(sums, 0, 8*sizeof(float)); + + float sumf = 0; + for (int i = 0; i < nb; ++i) { + const uint8_t * GGML_RESTRICT q4 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + memset(aux32, 0, 8*sizeof(int32_t)); + int8_t * GGML_RESTRICT a = aux8; + for (int j = 0; j < QK_K/64; ++j) { + for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] & 0xF); + a += 32; + for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] >> 4); + a += 32; q4 += 32; + } + memcpy(utmp, x[i].scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + + int sumi = 0; + for (int j = 0; j < QK_K/16; ++j) sumi += y[i].bsums[j] * mins[j/2]; + a = aux8; + int is = 0; + for (int j = 0; j < QK_K/32; ++j) { + int32_t scale = scales[is++]; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + } + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l]; + const float dmin = GGML_CPU_FP16_TO_FP32(x[i].dmin) * y[i].d; + sumf -= dmin * sumi; + } + for (int l = 0; l < 8; ++l) sumf += sums[l]; + *s = sumf; +} + +void ggml_vec_dot_q5_K_q8_K_generic(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) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q5_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + uint32_t utmp[4]; + + const uint8_t * scales = (const uint8_t*)&utmp[0]; + const uint8_t * mins = (const uint8_t*)&utmp[2]; + + int8_t aux8[QK_K]; + int16_t aux16[8]; + float sums [8]; + int32_t aux32[8]; + memset(sums, 0, 8*sizeof(float)); + + float sumf = 0; + for (int i = 0; i < nb; ++i) { + const uint8_t * GGML_RESTRICT q4 = x[i].qs; + const uint8_t * GGML_RESTRICT hm = x[i].qh; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + memset(aux32, 0, 8*sizeof(int32_t)); + int8_t * GGML_RESTRICT a = aux8; + uint8_t m = 1; + for (int j = 0; j < QK_K/64; ++j) { + for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] & 0xF); + for (int l = 0; l < 32; ++l) a[l] += (hm[l] & m ? 16 : 0); + a += 32; m <<= 1; + for (int l = 0; l < 32; ++l) a[l] = (int8_t)(q4[l] >> 4); + for (int l = 0; l < 32; ++l) a[l] += (hm[l] & m ? 16 : 0); + a += 32; m <<= 1; + q4 += 32; + } + memcpy(utmp, x[i].scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + + int sumi = 0; + for (int j = 0; j < QK_K/16; ++j) sumi += y[i].bsums[j] * mins[j/2]; + a = aux8; + int is = 0; + for (int j = 0; j < QK_K/32; ++j) { + int32_t scale = scales[is++]; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + } + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l]; + const float dmin = GGML_CPU_FP16_TO_FP32(x[i].dmin) * y[i].d; + sumf -= dmin * sumi; + } + for (int l = 0; l < 8; ++l) sumf += sums[l]; + *s = sumf; +} + +void ggml_vec_dot_q6_K_q8_K_generic(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) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q6_K * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + int8_t aux8[QK_K]; + int16_t aux16[8]; + float sums [8]; + int32_t aux32[8]; + memset(sums, 0, 8*sizeof(float)); + + float sumf = 0; + for (int i = 0; i < nb; ++i) { + const uint8_t * GGML_RESTRICT q4 = x[i].ql; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + memset(aux32, 0, 8*sizeof(int32_t)); + int8_t * GGML_RESTRICT a = aux8; + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + a[l + 0] = (int8_t)((q4[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; + a[l + 32] = (int8_t)((q4[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; + a[l + 64] = (int8_t)((q4[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + a[l + 96] = (int8_t)((q4[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + } + a += 128; + q4 += 64; + qh += 32; + } + a = aux8; + int is = 0; + for (int j = 0; j < QK_K/16; ++j) { + int scale = x[i].scales[is++]; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux16[l] = q8[l] * a[l]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * aux16[l]; + q8 += 8; a += 8; + } + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + for (int l = 0; l < 8; ++l) sums[l] += d * aux32[l]; + } + for (int l = 0; l < 8; ++l) sumf += sums[l]; + *s = sumf; +} + +void ggml_vec_dot_iq2_xxs_q8_K_generic(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) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq2_xxs * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + uint32_t aux32[2]; + const uint8_t * aux8 = (const uint8_t *)aux32; + + float sumf = 0.f; + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint16_t * GGML_RESTRICT q2 = x[i].qs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + memcpy(aux32, q2, 2*sizeof(uint32_t)); + q2 += 4; + const uint32_t ls = 2*(aux32[1] >> 28) + 1; + int32_t sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[l]); + const uint8_t signs = ksigns_iq2xs[(aux32[1] >> 7*l) & 127]; + for (int j = 0; j < 8; ++j) { + sumi += grid[j] * q8[j] * (signs & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + } + bsum += sumi * ls; + } + sumf += d * bsum; + } + *s = 0.125f * sumf; +} + +void ggml_vec_dot_iq2_xs_q8_K_generic(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) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq2_xs * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + float sumf = 0.f; + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint16_t * GGML_RESTRICT q2 = x[i].qs; + const uint8_t * GGML_RESTRICT sc = x[i].scales; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + const uint16_t ls1 = 2*(sc[ib32] & 0xf) + 1; + const uint16_t ls2 = 2*(sc[ib32] >> 4) + 1; + int32_t sumi = 0; + for (int l = 0; l < 2; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[l] & 511)); + const uint8_t signs = ksigns_iq2xs[q2[l] >> 9]; + for (int j = 0; j < 8; ++j) { + sumi += grid[j] * q8[j] * (signs & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + } + bsum += sumi * ls1; + sumi = 0; + for (int l = 2; l < 4; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[l] & 511)); + const uint8_t signs = ksigns_iq2xs[q2[l] >> 9]; + for (int j = 0; j < 8; ++j) { + sumi += grid[j] * q8[j] * (signs & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + } + bsum += sumi * ls2; + q2 += 4; + } + sumf += d * bsum; + } + *s = 0.125f * sumf; +} + +void ggml_vec_dot_iq2_s_q8_K_generic(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) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq2_s * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + float sumf = 0; + for (int i = 0; i < nb; i++) { + + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const int8_t * q8 = y[i].qs; + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + const uint8_t * signs = qs + QK_K/8; + + int bsum = 0; + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + int ls1 = 1 + 2*(x[i].scales[ib32] & 0xf); + int ls2 = 1 + 2*(x[i].scales[ib32] >> 4); + int sumi1 = 0, sumi2 = 0; + for (int l = 0; l < 2; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2s_grid + (qs[l] | (qh[ib32] << (8-2*l) & 0x300))); + for (int j = 0; j < 8; ++j) { + sumi1 += q8[j] * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + } + for (int l = 2; l < 4; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2s_grid + (qs[l] | (qh[ib32] << (8-2*l) & 0x300))); + for (int j = 0; j < 8; ++j) { + sumi2 += q8[j] * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1 : 1); + } + q8 += 8; + } + bsum += ls1 * sumi1 + ls2 * sumi2; + qs += 4; + signs += 4; + } + + sumf += d * bsum; + } + + *s = 0.125f * sumf; +} + +void ggml_vec_dot_iq3_xxs_q8_K_generic(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) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq3_xxs * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + uint32_t aux32; + + float sumf = 0.f; + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint8_t * GGML_RESTRICT q3 = x[i].qs; + const uint8_t * GGML_RESTRICT gas = x[i].qs + QK_K/4; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + memcpy(&aux32, gas, sizeof(uint32_t)); gas += sizeof(uint32_t); + const uint32_t ls = 2*(aux32 >> 28) + 1; + int32_t sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*l+0]); + const uint8_t * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*l+1]); + const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*l) & 127]; + for (int j = 0; j < 4; ++j) { + sumi += grid1[j] * q8[j+0] * (signs & kmask_iq2xs[j+0] ? -1 : 1); + sumi += grid2[j] * q8[j+4] * (signs & kmask_iq2xs[j+4] ? -1 : 1); + } + q8 += 8; + } + q3 += 8; + bsum += sumi * ls; + } + sumf += d * bsum; + } + *s = 0.25f * sumf; +} + +void ggml_vec_dot_iq3_s_q8_K_generic(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) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq3_s * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + float sumf = 0.f; + for (int i = 0; i < nb; ++i) { + const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d; + const uint8_t * GGML_RESTRICT qs = x[i].qs; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const uint8_t * GGML_RESTRICT signs = x[i].signs; + const int8_t * GGML_RESTRICT q8 = y[i].qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { + const uint32_t ls1 = 2*(x[i].scales[ib32/2] & 0xf) + 1; + const uint32_t ls2 = 2*(x[i].scales[ib32/2] >> 4) + 1; + int32_t sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[ib32+0] << (8-2*l)) & 256))); + const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[ib32+0] << (7-2*l)) & 256))); + for (int j = 0; j < 4; ++j) { + sumi += grid1[j] * q8[j+0] * (signs[l] & kmask_iq2xs[j+0] ? -1 : 1); + sumi += grid2[j] * q8[j+4] * (signs[l] & kmask_iq2xs[j+4] ? -1 : 1); + } + q8 += 8; + } + qs += 8; + signs += 4; + bsum += sumi * ls1; + sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[ib32+1] << (8-2*l)) & 256))); + const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[ib32+1] << (7-2*l)) & 256))); + for (int j = 0; j < 4; ++j) { + sumi += grid1[j] * q8[j+0] * (signs[l] & kmask_iq2xs[j+0] ? -1 : 1); + sumi += grid2[j] * q8[j+4] * (signs[l] & kmask_iq2xs[j+4] ? -1 : 1); + } + q8 += 8; + } + qs += 8; + signs += 4; + bsum += sumi * ls2; + } + sumf += d * bsum; + } + *s = sumf; +} + +void ggml_vec_dot_iq1_s_q8_K_generic(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) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq1_s * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + float sumf = 0; + for (int i = 0; i < nb; i++) { + + const int8_t * q8 = y[i].qs; + const uint8_t * qs = x[i].qs; + const uint16_t * qh = x[i].qh; + + int sumi = 0, sumi1 = 0; + for (int ib = 0; ib < QK_K/32; ++ib) { + const int ls = 2*((qh[ib] >> 12) & 7) + 1; + const int delta = qh[ib] & 0x8000 ? -1 : 1; + int lsum = 0; + for (int l = 0; l < 4; ++l) { + const int8_t * grid = (const int8_t *)(iq1s_grid + (qs[l] | (((qh[ib] >> 3*l) & 7) << 8))); + for (int j = 0; j < 8; ++j) { + lsum += q8[j] * grid[j]; + } + q8 += 8; + } + sumi += ls * lsum; + sumi1 += ls * delta * (y[i].bsums[2*ib+0] + y[i].bsums[2*ib+1]); + qs += 4; + } + + sumf += GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d * (sumi + IQ1S_DELTA * sumi1); + } + + *s = sumf; +} + +void ggml_vec_dot_iq1_m_q8_K_generic(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) { + assert(n % QK_K == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_iq1_m * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + iq1m_scale_t scale; + + int sum1[2], sum2[2], delta[4]; + + float sumf = 0; + for (int i = 0; i < nb; i++) { + + const int8_t * q8 = y[i].qs; + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + const uint16_t * sc = (const uint16_t *)x[i].scales; + + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); + + int sumi1 = 0, sumi2 = 0; + for (int ib = 0; ib < QK_K/32; ++ib) { + delta[0] = qh[0] & 0x08 ? -1 : 1; + delta[1] = qh[0] & 0x80 ? -1 : 1; + delta[2] = qh[1] & 0x08 ? -1 : 1; + delta[3] = qh[1] & 0x80 ? -1 : 1; + sum1[0] = sum1[1] = sum2[0] = sum2[1] = 0; + for (int l = 0; l < 4; ++l) { + const int8_t * grid = (const int8_t *)(iq1s_grid + (qs[l] | (((uint16_t)qh[l/2] << (8 - 4*(l%2))) & 0x700))); + int lsum1 = 0, lsum2 = 0; + for (int j = 0; j < 8; ++j) { + lsum1 += q8[j] * grid[j]; + lsum2 += q8[j]; + } + q8 += 8; + sum1[l/2] += lsum1; + sum2[l/2] += lsum2*delta[l]; + } + + const int ls1 = 2*((sc[ib/2] >> (6*(ib%2)+0)) & 0x7) + 1; + const int ls2 = 2*((sc[ib/2] >> (6*(ib%2)+3)) & 0x7) + 1; + + sumi1 += sum1[0] * ls1 + sum1[1] * ls2; + sumi2 += sum2[0] * ls1 + sum2[1] * ls2; + qs += 4; + qh += 2; + } + + sumf += GGML_CPU_FP16_TO_FP32(scale.f16) * y[i].d * (sumi1 + IQ1M_DELTA * sumi2); + } + + *s = sumf; +} + +void ggml_vec_dot_iq4_nl_q8_0_generic(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) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK4_NL == 0); + static_assert(QK4_NL == QK8_0, "QK4_NL and QK8_0 must be the same"); + + const block_iq4_nl * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + const int nb = n / QK4_NL; + + int ib = 0; + float sumf = 0; + + for (; ib < nb; ++ib) { + const float d = GGML_CPU_FP16_TO_FP32(y[ib].d)*GGML_CPU_FP16_TO_FP32(x[ib].d); + int sumi1 = 0, sumi2 = 0; + for (int j = 0; j < QK4_NL/2; ++j) { + sumi1 += y[ib].qs[j+ 0] * kvalues_iq4nl[x[ib].qs[j] & 0xf]; + sumi2 += y[ib].qs[j+QK4_NL/2] * kvalues_iq4nl[x[ib].qs[j] >> 4]; + } + sumf += d * (sumi1 + sumi2); + } + *s = sumf; +} + +void ggml_vec_dot_iq4_xs_q8_K_generic(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) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK_K == 0); + + const block_iq4_xs * GGML_RESTRICT x = vx; + const block_q8_K * GGML_RESTRICT y = vy; + + const int nb = n / QK_K; + + float sumf = 0; + for (int ibl = 0; ibl < nb; ++ibl) { + const float d4d8 = GGML_CPU_FP16_TO_FP32(x[ibl].d) * y[ibl].d; + uint16_t h = x[ibl].scales_h; + const uint8_t * qs = x[ibl].qs; + const int8_t * q8 = y[ibl].qs; + for (int ib = 0; ib < QK_K/32; ib += 2) { + const uint8_t ls1 = (x[ibl].scales_l[ib/2] & 0xf) | ((h << 4) & 0x30); + const uint8_t ls2 = (x[ibl].scales_l[ib/2] >> 4) | ((h << 2) & 0x30); + h >>= 4; + const float d1 = d4d8*(ls1 - 32); + const float d2 = d4d8*(ls2 - 32); + int sumi1 = 0, sumi2 = 0; + for (int j = 0; j < 16; ++j) { + sumi1 += q8[j+ 0] * kvalues_iq4nl[qs[j] & 0xf]; + sumi2 += q8[j+16] * kvalues_iq4nl[qs[j] >> 4]; + } + sumf += d1 * (sumi1 + sumi2); + qs += 16; + q8 += 32; + sumi1 = sumi2 = 0; + for (int j = 0; j < 16; ++j) { + sumi1 += q8[j+ 0] * kvalues_iq4nl[qs[j] & 0xf]; + sumi2 += q8[j+16] * kvalues_iq4nl[qs[j] >> 4]; + } + sumf += d2 * (sumi1 + sumi2); + qs += 16; + q8 += 32; + } + } + *s = sumf; +} + +// ============================ 4-bit non-linear quants + +void quantize_row_iq4_nl(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + assert(k % QK4_NL == 0); + quantize_row_iq4_nl_ref(x, y, k); +} + +void quantize_row_iq4_xs(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq4_xs(x, y, 1, k, NULL); +} + + +/* ---- x86 quantization implementation ---- */ + // some compilers don't provide _mm256_set_m128i, e.g. gcc 7 #define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1) diff --git a/ggml/src/ggml-cpu/quants.h b/ggml/src/ggml-cpu/quants.h deleted file mode 100644 index d4bc87a1c052..000000000000 --- a/ggml/src/ggml-cpu/quants.h +++ /dev/null @@ -1,103 +0,0 @@ -#pragma once - -#define GGML_COMMON_DECL_C -#include "ggml-common.h" - -#include "ggml.h" - -// GGML CPU internal header - -#ifdef __cplusplus -extern "C" { -#endif - -// Quantization -void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_q5_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_q8_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); - -void quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); - -void quantize_row_q2_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_q3_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_q4_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_q5_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_q6_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_q8_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); - -void quantize_row_tq1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_tq2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); - -void quantize_row_iq4_nl (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void quantize_row_iq4_xs (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); - -// Dot product -void ggml_vec_dot_q1_0_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 ggml_vec_dot_q4_0_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 ggml_vec_dot_q4_1_q8_1(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 ggml_vec_dot_q5_0_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 ggml_vec_dot_q5_1_q8_1(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 ggml_vec_dot_q8_0_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 ggml_vec_dot_mxfp4_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 ggml_vec_dot_nvfp4_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 ggml_vec_dot_q2_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 ggml_vec_dot_q3_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 ggml_vec_dot_q4_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 ggml_vec_dot_q5_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 ggml_vec_dot_q6_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 ggml_vec_dot_tq1_0_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 ggml_vec_dot_tq2_0_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 ggml_vec_dot_iq2_xxs_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 ggml_vec_dot_iq2_xs_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 ggml_vec_dot_iq2_s_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 ggml_vec_dot_iq3_xxs_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 ggml_vec_dot_iq1_s_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 ggml_vec_dot_iq1_m_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 ggml_vec_dot_iq4_nl_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 ggml_vec_dot_iq4_xs_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 ggml_vec_dot_iq3_s_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); - -// Generic implementation -void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); -void ggml_vec_dot_q1_0_q8_0_generic(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 ggml_vec_dot_q4_0_q8_0_generic(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 ggml_vec_dot_q4_1_q8_1_generic(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 ggml_vec_dot_q5_0_q8_0_generic(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 ggml_vec_dot_q5_1_q8_1_generic(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 ggml_vec_dot_q8_0_q8_0_generic(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 ggml_vec_dot_mxfp4_q8_0_generic(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 ggml_vec_dot_nvfp4_q8_0_generic(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 ggml_vec_dot_tq1_0_q8_K_generic(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 ggml_vec_dot_tq2_0_q8_K_generic(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 ggml_vec_dot_q2_K_q8_K_generic(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 ggml_vec_dot_q3_K_q8_K_generic(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 ggml_vec_dot_q4_K_q8_K_generic(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 ggml_vec_dot_q5_K_q8_K_generic(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 ggml_vec_dot_q6_K_q8_K_generic(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 ggml_vec_dot_iq2_xxs_q8_K_generic(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 ggml_vec_dot_iq2_xs_q8_K_generic(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 ggml_vec_dot_iq2_s_q8_K_generic(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 ggml_vec_dot_iq3_xxs_q8_K_generic(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 ggml_vec_dot_iq3_s_q8_K_generic(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 ggml_vec_dot_iq1_s_q8_K_generic(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 ggml_vec_dot_iq1_m_q8_K_generic(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 ggml_vec_dot_iq4_nl_q8_0_generic(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 ggml_vec_dot_iq4_xs_q8_K_generic(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); - -#ifdef __cplusplus -} -#endif diff --git a/ggml/src/ggml-cpu/repack.cpp b/ggml/src/ggml-cpu/repack.cpp deleted file mode 100644 index f18758f16bb6..000000000000 --- a/ggml/src/ggml-cpu/repack.cpp +++ /dev/null @@ -1,4836 +0,0 @@ -#define GGML_COMMON_IMPL_CPP -#define GGML_COMMON_DECL_CPP -#include "ggml-common.h" -#include "ggml-backend-impl.h" - -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "ggml-cpu-impl.h" -#include "simd-mappings.h" -#include "traits.h" - -#include "arch-fallback.h" - -#include -#include -#include -#include // for GGML_ASSERT - -#include "repack.h" - -#if defined(__GNUC__) -#pragma GCC diagnostic ignored "-Woverlength-strings" -#endif - -#define UNUSED GGML_UNUSED - -static inline int nearest_int(float fval) { - assert(fabsf(fval) <= 4194303.f); - float val = fval + 12582912.f; - int i; memcpy(&i, &val, sizeof(int)); - return (i & 0x007fffff) - 0x00400000; -} - -// Functions to create the interleaved data layout formats - -// interleave 4 block_q4_0s in blocks of blck_size_interleave -// returns an interleaved block_q4_0x4 -// in the interleaved block_q4_0x4, place deltas for 4 block_q4_0 blocks -// first, then interleave quants from 4 block_q4_0s in blocks of blck_size_interleave -// -// - in : an array of block_q4_0 pointers -// - blck_size_interleave : the block_q4_0 quants bytes are interleaved in blocks of -// blck_size_interleave bytes -// - xor_mask : the mask to convert the nibbles in block_q4_0 quants bytes -// from bias offset form to pure sign form (this saves subtract -// operations durin unpacking) -// - -extern "C" { - -#if defined __riscv_zvfh -void ggml_quantize_mat_q8_0_4x1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - assert(QK8_0 == 32); - assert(k % QK8_0 == 0); - const int nb = k / QK8_0; - - block_q8_0x4 * GGML_RESTRICT y = (block_q8_0x4 *) vy; - - // scalar - const int blck_size_interleave = 1; - float srcv[4][QK8_0]; - float id[4]; - - for (int i = 0; i < nb; i++) { - for (int row_iter = 0; row_iter < 4; row_iter++) { - float amax = 0.0f; // absolute max - - for (int j = 0; j < QK8_0; j++) { - srcv[row_iter][j] = x[row_iter * k + i * QK8_0 + j]; - amax = MAX(amax, fabsf(srcv[row_iter][j])); - } - - const float d = amax / ((1 << 7) - 1); - id[row_iter] = d ? 1.0f / d : 0.0f; - - y[i].d[row_iter] = GGML_CPU_FP32_TO_FP16(d); - } - - for (int j = 0; j < QK8_0 * 4; j++) { - int src_offset = (j / (4 * blck_size_interleave)) * blck_size_interleave; - int src_id = (j % (4 * blck_size_interleave)) / blck_size_interleave; - src_offset += (j % blck_size_interleave); - - float x0 = srcv[src_id][src_offset] * id[src_id]; - y[i].qs[j] = roundf(x0); - } - } -} - -void ggml_quantize_mat_q8_K_4x1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - assert(QK_K == 256); - assert(k % QK_K == 0); - const int nb = k / QK_K; - - block_q8_Kx4 * GGML_RESTRICT y = (block_q8_Kx4 *) vy; - - const int blck_size_interleave = 1; - float srcv[4][QK_K]; - float iscale[4]; - - for (int i = 0; i < nb; i++) { - for (int row_iter = 0; row_iter < 4; row_iter++) { - float amax = 0.0f; // absolute max - float max = 0; - - for (int j = 0; j < QK_K; j++) { - srcv[row_iter][j] = x[row_iter * k + i * QK_K + j]; - // Update the maximum value of the corresponding super block - if(amax < fabsf(srcv[row_iter][j])) { - amax = fabsf(srcv[row_iter][j]); - max = srcv[row_iter][j]; - } - } - - iscale[row_iter] = amax ? -127.f/max : 0; - y[i].d[row_iter] = amax ? 1/iscale[row_iter] : 0; - } - - for (int j = 0; j < QK_K / 4; j++) { - y[i].bsums[j] = 0; - } - for (int j = 0; j < QK_K * 4; j++) { - int src_id = j % 4; - int src_offset = j / 4; - int index = ((j >> 6) << 2) + (j & 3); - - float x0 = srcv[src_id][src_offset] * iscale[src_id]; - y[i].qs[j] = nearest_int(x0); - y[i].bsums[index] += y[i].qs[j]; - } - } -} -#endif - -void ggml_quantize_mat_q8_0_4x4_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - assert(QK8_0 == 32); - assert(k % QK8_0 == 0); - const int nb = k / QK8_0; - - block_q8_0x4 * GGML_RESTRICT y = (block_q8_0x4 *) vy; - - // scalar - const int blck_size_interleave = 4; - float srcv[4][QK8_0]; - float id[4]; - - for (int i = 0; i < nb; i++) { - for (int row_iter = 0; row_iter < 4; row_iter++) { - float amax = 0.0f; // absolute max - - for (int j = 0; j < QK8_0; j++) { - srcv[row_iter][j] = x[row_iter * k + i * QK8_0 + j]; - amax = MAX(amax, fabsf(srcv[row_iter][j])); - } - - const float d = amax / ((1 << 7) - 1); - id[row_iter] = d ? 1.0f / d : 0.0f; - - y[i].d[row_iter] = GGML_CPU_FP32_TO_FP16(d); - } - - for (int j = 0; j < QK8_0 * 4; j++) { - int src_offset = (j / (4 * blck_size_interleave)) * blck_size_interleave; - int src_id = (j % (4 * blck_size_interleave)) / blck_size_interleave; - src_offset += (j % blck_size_interleave); - - float x0 = srcv[src_id][src_offset] * id[src_id]; - y[i].qs[j] = roundf(x0); - } - } -} - -void ggml_quantize_mat_q8_0_4x8_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - assert(QK8_0 == 32); - assert(k % QK8_0 == 0); - const int nb = k / QK8_0; - - block_q8_0x4 * GGML_RESTRICT y = (block_q8_0x4 *) vy; - - // scalar - const int blck_size_interleave = 8; - float srcv[4][QK8_0]; - float id[4]; - - for (int i = 0; i < nb; i++) { - for (int row_iter = 0; row_iter < 4; row_iter++) { - float amax = 0.0f; // absolute max - - for (int j = 0; j < QK8_0; j++) { - srcv[row_iter][j] = x[row_iter * k + i * QK8_0 + j]; - amax = MAX(amax, fabsf(srcv[row_iter][j])); - } - - const float d = amax / ((1 << 7) - 1); - id[row_iter] = d ? 1.0f / d : 0.0f; - - y[i].d[row_iter] = GGML_CPU_FP32_TO_FP16(d); - } - - for (int j = 0; j < QK8_0 * 4; j++) { - int src_offset = (j / (4 * blck_size_interleave)) * blck_size_interleave; - int src_id = (j % (4 * blck_size_interleave)) / blck_size_interleave; - src_offset += (j % blck_size_interleave); - - float x0 = srcv[src_id][src_offset] * id[src_id]; - y[i].qs[j] = roundf(x0); - } - } -} - -void ggml_quantize_mat_q8_K_4x4_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - assert(QK_K == 256); - assert(k % QK_K == 0); - const int nb = k / QK_K; - - block_q8_Kx4 * GGML_RESTRICT y = (block_q8_Kx4 *) vy; - - // scalar - const int blck_size_interleave = 4; - float srcv[4][QK_K]; - float iscale[4]; - - for (int i = 0; i < nb; i++) { - for (int row_iter = 0; row_iter < 4; row_iter++) { - float amax = 0.0f; // absolute max - float max = 0; - - for (int j = 0; j < QK_K; j++) { - srcv[row_iter][j] = x[row_iter * k + i * QK_K + j]; - // Update the maximum value of the corresponding super block - if(amax < fabsf(srcv[row_iter][j])) { - amax = fabsf(srcv[row_iter][j]); - max = srcv[row_iter][j]; - } - } - - iscale[row_iter] = amax ? -127.f/max : 0; - - y[i].d[row_iter] = amax ? 1/iscale[row_iter] : 0; - } - - for (int j = 0; j < QK_K / 4; j++) { - y[i].bsums[j] = 0; - } - - // Quants values are interleaved in sequence of four bytes from corresponding super blocks - // Bsums values are interleaved in sequence of four bsums from each super block taken for interleaving - // i.e first four bsums from the first super block, followed by first four bsums from second super block and so on - for (int j = 0; j < QK_K * 4; j++) { - int src_offset = (j / (4 * blck_size_interleave)) * blck_size_interleave; - int src_id = (j % (4 * blck_size_interleave)) / blck_size_interleave; - src_offset += (j % blck_size_interleave); - int index = (((j & 15) >> 2) << 2) + ((j >> 8) << 4) + ((j >> 6) & 3); - - float x0 = srcv[src_id][src_offset] * iscale[src_id]; - y[i].qs[j] = nearest_int(x0); - y[i].bsums[index] += y[i].qs[j]; - } - } -} - -void ggml_quantize_mat_q8_K_4x8_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { - assert(QK_K == 256); - assert(k % QK_K == 0); - const int nb = k / QK_K; - - block_q8_Kx4 * GGML_RESTRICT y = (block_q8_Kx4 *) vy; - - // scalar - const int blck_size_interleave = 8; - float srcv[4][QK_K]; - float iscale[4]; - - for (int i = 0; i < nb; i++) { - for (int row_iter = 0; row_iter < 4; row_iter++) { - float amax = 0.0f; // absolute max - float max = 0; - - for (int j = 0; j < QK_K; j++) { - srcv[row_iter][j] = x[row_iter * k + i * QK_K + j]; - // Update the maximum value of the corresponding super block - if(amax < fabsf(srcv[row_iter][j])) { - amax = fabsf(srcv[row_iter][j]); - max = srcv[row_iter][j]; - } - } - - iscale[row_iter] = amax ? -127.f/max : 0; - - y[i].d[row_iter] = amax ? 1/iscale[row_iter] : 0; - } - - for (int j = 0; j < QK_K / 4; j++) { - y[i].bsums[j] = 0; - } - - // Quants values are interleaved in sequence of eight bytes from corresponding super blocks - // Bsums values are interleaved in sequence of four bsums from each super block taken for interleaving - // i.e first four bsums from the first super block, followed by first four bsums from second super block and so on - for (int j = 0; j < QK_K * 4; j++) { - int src_offset = (j / (4 * blck_size_interleave)) * blck_size_interleave; - int src_id = (j % (4 * blck_size_interleave)) / blck_size_interleave; - src_offset += (j % blck_size_interleave); - int index = (((j & 31) >> 3) << 2) + ((j >> 8) << 4) + ((j >> 6) & 3); - - float x0 = srcv[src_id][src_offset] * iscale[src_id]; - y[i].qs[j] = nearest_int(x0); - y[i].bsums[index] += y[i].qs[j]; - } - } -} - -} // extern "C" - -template -void ggml_quantize_mat_t(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row); - -template <> void ggml_quantize_mat_t<4, GGML_TYPE_Q8_0>(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row) { - assert(nrow == 4); - UNUSED(nrow); - ggml_quantize_mat_q8_0_4x4(x, vy, n_per_row); -} - -template <> void ggml_quantize_mat_t<8, GGML_TYPE_Q8_0>(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row) { - assert(nrow == 4); - UNUSED(nrow); - ggml_quantize_mat_q8_0_4x8(x, vy, n_per_row); -} - -template <> void ggml_quantize_mat_t<4, GGML_TYPE_Q8_K>(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row) { - assert(nrow == 4); - UNUSED(nrow); - ggml_quantize_mat_q8_K_4x4(x, vy, n_per_row); -} - -template <> void ggml_quantize_mat_t<8, GGML_TYPE_Q8_K>(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row) { - assert(nrow == 4); - UNUSED(nrow); - ggml_quantize_mat_q8_K_4x8(x, vy, n_per_row); -} - -#if defined __riscv_zvfh -template <> void ggml_quantize_mat_t<1, GGML_TYPE_Q8_0>(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row) { - assert(nrow == 4); - UNUSED(nrow); - ggml_quantize_mat_q8_0_4x1(x, vy, n_per_row); -} - -template <> void ggml_quantize_mat_t<1, GGML_TYPE_Q8_K>(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row) { - assert(nrow == 4); - UNUSED(nrow); - ggml_quantize_mat_q8_K_4x1(x, vy, n_per_row); -} -#endif - -template -static void ggml_gemv_q6_K_NxM_q8_K_generic_impl(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * GGML_RESTRICT vy, - int nr, - int nc) { - constexpr int blocklen = M; - constexpr int ncols_interleaved = N; - const int qk = QK_K; - const int nb = n / qk; - const int blocks_per_half = 64 / blocklen; - - assert(n % qk == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - float sumf[8]; - - const block_q8_K * a_ptr = (const block_q8_K *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q6_Kx8 * b_ptr = (const block_q6_Kx8 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) { - sumf[j] = 0.0f; - } - - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - const int base_l = (k / blocks_per_half) * 128 + (k % blocks_per_half) * blocklen; - const int base_h = base_l + 64; - - const int scale_idx_l = base_l / 16; - const int scale_idx_h = base_h / 16; - - const int qh_shift_l = ((base_l % 128) / 32) * 2; - const int qh_shift_h = ((base_h % 128) / 32) * 2; - - const int qh_half_l = (base_l / 128) * 32; - const int qh_half_h = (base_h / 128) * 32; - - for (int j = 0; j < ncols_interleaved; j++) { - const int8_t scale_l = b_ptr[l].scales[scale_idx_l * ncols_interleaved + j]; - const int8_t scale_h = b_ptr[l].scales[scale_idx_h * ncols_interleaved + j]; - - int sumi_l = 0; - int sumi_h = 0; - - for (int i = 0; i < blocklen; i++) { - const int ql_pos = k * ncols_interleaved * blocklen + j * blocklen + i; - const int l_4 = b_ptr[l].ql[ql_pos] & 0xF; - const int hi_4 = (b_ptr[l].ql[ql_pos] >> 4) & 0xF; - - const int qh_idx_l = qh_half_l + ((base_l + i) % 32); - const int qh_chunk_l = qh_idx_l / blocklen; - const int qh_pos_l = qh_idx_l % blocklen; - const int qh_offset_l = qh_chunk_l * (blocklen * ncols_interleaved) + j * blocklen + qh_pos_l; - const int hi_2_l = (b_ptr[l].qh[qh_offset_l] >> qh_shift_l) & 0x3; - - const int qh_idx_h = qh_half_h + ((base_h + i) % 32); - const int qh_chunk_h = qh_idx_h / blocklen; - const int qh_pos_h = qh_idx_h % blocklen; - const int qh_offset_h = qh_chunk_h * (blocklen * ncols_interleaved) + j * blocklen + qh_pos_h; - const int hi_2_h = (b_ptr[l].qh[qh_offset_h] >> qh_shift_h) & 0x3; - - const int q_l = ((hi_2_l << 4) | l_4) - 32; - const int q_h = ((hi_2_h << 4) | hi_4) - 32; - - const int8_t a_l = a_ptr[l].qs[base_l + i]; - const int8_t a_h = a_ptr[l].qs[base_h + i]; - - sumi_l += q_l * a_l; - sumi_h += q_h * a_h; - } - - sumf[j] += - (sumi_l * scale_l + sumi_h * scale_h) * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d; - } - } - } - - for (int j = 0; j < ncols_interleaved; j++) { - s[x * ncols_interleaved + j] = sumf[j]; - } - } -} - -template -static void ggml_gemm_q6_K_NxM_q8_K_generic_impl(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * GGML_RESTRICT vy, - int nr, - int nc) { - constexpr int blocklen = M; - constexpr int ncols_interleaved = N; - const int qk = QK_K; - const int nb = n / qk; - const int blocks_per_half = 64 / blocklen; - const int q8_half_stride = 512; - const int q8_low_high_step = 256; - - assert(n % qk == 0); - assert(nr % 4 == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - - float sumf[4][8]; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_Kx4 * a_ptr = (const block_q8_Kx4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q6_Kx8 * b_ptr = (const block_q6_Kx8 *) vx + (x * nb); - - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumf[m][j] = 0.0f; - } - } - - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - const int base_l = (k / blocks_per_half) * 128 + (k % blocks_per_half) * blocklen; - const int base_h = base_l + 64; - - const int scale_idx_l = base_l / 16; - const int scale_idx_h = base_h / 16; - - const int qh_shift_l = ((base_l % 128) / 32) * 2; - const int qh_shift_h = ((base_h % 128) / 32) * 2; - - const int qh_half_l = (base_l / 128) * 32; - const int qh_half_h = (base_h / 128) * 32; - - const int q8_base = (k / blocks_per_half) * q8_half_stride + (k % blocks_per_half) * (blocklen * 4); - - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - const int8_t scale_l = b_ptr[l].scales[scale_idx_l * ncols_interleaved + j]; - const int8_t scale_h = b_ptr[l].scales[scale_idx_h * ncols_interleaved + j]; - - int sumi_l = 0; - int sumi_h = 0; - - for (int i = 0; i < blocklen; i++) { - const int ql_pos = k * ncols_interleaved * blocklen + j * blocklen + i; - const int l_4 = b_ptr[l].ql[ql_pos] & 0xF; - const int hi_4 = (b_ptr[l].ql[ql_pos] >> 4) & 0xF; - - const int qh_idx_l = qh_half_l + ((base_l + i) % 32); - const int qh_chunk_l = qh_idx_l / blocklen; - const int qh_pos_l = qh_idx_l % blocklen; - const int qh_offset_l = - qh_chunk_l * (blocklen * ncols_interleaved) + j * blocklen + qh_pos_l; - const int hi_2_l = (b_ptr[l].qh[qh_offset_l] >> qh_shift_l) & 0x3; - - const int qh_idx_h = qh_half_h + ((base_h + i) % 32); - const int qh_chunk_h = qh_idx_h / blocklen; - const int qh_pos_h = qh_idx_h % blocklen; - const int qh_offset_h = - qh_chunk_h * (blocklen * ncols_interleaved) + j * blocklen + qh_pos_h; - const int hi_2_h = (b_ptr[l].qh[qh_offset_h] >> qh_shift_h) & 0x3; - - const int q_l = ((hi_2_l << 4) | l_4) - 32; - const int q_h = ((hi_2_h << 4) | hi_4) - 32; - - const int8_t q8_l = a_ptr[l].qs[q8_base + m * blocklen + i]; - const int8_t q8_h = a_ptr[l].qs[q8_base + m * blocklen + i + q8_low_high_step]; - - sumi_l += q_l * q8_l; - sumi_h += q_h * q8_h; - } - - sumf[m][j] += (sumi_l * scale_l + sumi_h * scale_h) * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * - a_ptr[l].d[m]; - } - } - } - } - - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } - } -} - -template -static void ggml_gemv_q5_K_NxM_q8_K_generic_impl(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * GGML_RESTRICT vy, - int nr, - int nc) { - constexpr int blocklen = M; - constexpr int ncols_interleaved = N; - const int qk = QK_K; - const int nb = n / qk; - static const uint32_t kmask1 = 0x3f3f3f3f; - static const uint32_t kmask2 = 0x0f0f0f0f; - static const uint32_t kmask3 = 0x03030303; - - assert(n % qk == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - float sumf[ncols_interleaved]; - float sum_minf[ncols_interleaved]; - uint32_t utmp[32]; - int sumi1; - int sumi2; - int sumi; - - const block_q8_K * a_ptr = (const block_q8_K *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q5_Kx8 * b_ptr = (const block_q5_Kx8 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) { - sumf[j] = 0.0; - sum_minf[j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int sb = 0; sb < 8; sb++) { - memcpy(utmp + sb * 4, b_ptr[l].scales + sb * K_SCALE_SIZE, K_SCALE_SIZE); - utmp[sb * 4 + 3] = ((utmp[sb * 4 + 2] >> 4) & kmask2) | (((utmp[sb * 4 + 1] >> 6) & kmask3) << 4); - const uint32_t uaux_0 = utmp[sb * 4 + 1] & kmask1; - utmp[sb * 4 + 1] = (utmp[sb * 4 + 2] & kmask2) | (((utmp[sb * 4 + 0] >> 6) & kmask3) << 4); - utmp[sb * 4 + 2] = uaux_0; - utmp[sb * 4 + 0] &= kmask1; - } - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - constexpr int scale_stride = 32; - uint8_t * scales_0 = (uint8_t *) utmp + (k / (32 / blocklen)) * scale_stride; - uint8_t * scales_1 = (uint8_t *) utmp + (k / (32 / blocklen)) * scale_stride + 16; - - const int qh_shift = (k / (32 / blocklen)) * 2; - for (int j = 0; j < ncols_interleaved; j++) { - sumi1 = 0; - sumi2 = 0; - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int b_qs_offset = k * ncols_interleaved * blocklen + j * blocklen + i; - - const int qh_idx = (k * blocklen + i) % 32; - const int qh_chunk = qh_idx / blocklen; - const int qh_pos = qh_idx % blocklen; - const int b_qh_offset = qh_chunk * (blocklen * ncols_interleaved) + j * blocklen + qh_pos; - - const uint8_t qh_val = b_ptr[l].qh[b_qh_offset]; - const uint8_t h0 = (qh_val >> qh_shift) & 1; - const uint8_t h1 = (qh_val >> (qh_shift + 1)) & 1; - - const int v0 = (int8_t) ((b_ptr[l].qs[b_qs_offset] & 0xF) | (h0 << 4)); - const int v1 = (int8_t) ((b_ptr[l].qs[b_qs_offset] >> 4) | (h1 << 4)); - - const int q8_offset = (k / (32 / blocklen)) * 64 + (k % (32 / blocklen)) * blocklen + i; - - sumi1 = (v0 * a_ptr[l].qs[q8_offset]); - sumi2 = (v1 * a_ptr[l].qs[q8_offset + 32]); - sumi1 = sumi1 * scales_0[j]; - sumi2 = sumi2 * scales_1[j]; - sumi += sumi1 + sumi2; - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d; - } - } - for (int sb = 0; sb < 8; sb++) { - uint8_t * mins = (uint8_t *) utmp + 8 + sb * 16; - for (int j = 0; j < ncols_interleaved; j++) { - sum_minf[j] += mins[j] * (a_ptr[l].bsums[sb * 2] + a_ptr[l].bsums[sb * 2 + 1]) * - GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d; - } - } - } - for (int j = 0; j < ncols_interleaved; j++) { - s[x * ncols_interleaved + j] = sumf[j] - sum_minf[j]; - } - } -} - -template -static void ggml_gemm_q5_K_NxM_q8_K_generic_impl(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * GGML_RESTRICT vy, - int nr, - int nc) { - constexpr int blocklen = M; - constexpr int ncols_interleaved = N; - const int qk = QK_K; - const int nb = n / qk; - static const uint32_t kmask1 = 0x3f3f3f3f; - static const uint32_t kmask2 = 0x0f0f0f0f; - static const uint32_t kmask3 = 0x03030303; - - assert(n % qk == 0); - assert(nr % 4 == 0); - assert(nc % ncols_interleaved == 0); - - float sumf[4][ncols_interleaved]; - float sum_minf[4][ncols_interleaved]; - uint32_t utmp[32]; - int sumi1; - int sumi2; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_Kx4 * a_ptr = (const block_q8_Kx4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q5_Kx8 * b_ptr = (const block_q5_Kx8 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumf[m][j] = 0.0; - sum_minf[m][j] = 0.0; - } - } - for (int l = 0; l < nb; l++) { - for (int sb = 0; sb < 8; sb++) { - memcpy(utmp + sb * 4, b_ptr[l].scales + sb * K_SCALE_SIZE, K_SCALE_SIZE); - utmp[sb * 4 + 3] = ((utmp[sb * 4 + 2] >> 4) & kmask2) | (((utmp[sb * 4 + 1] >> 6) & kmask3) << 4); - const uint32_t uaux_0 = utmp[sb * 4 + 1] & kmask1; - utmp[sb * 4 + 1] = (utmp[sb * 4 + 2] & kmask2) | (((utmp[sb * 4 + 0] >> 6) & kmask3) << 4); - utmp[sb * 4 + 2] = uaux_0; - utmp[sb * 4 + 0] &= kmask1; - } - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - constexpr int scale_stride = 32; - uint8_t * scales_0 = (uint8_t *) utmp + (k / (32 / blocklen)) * scale_stride; - uint8_t * scales_1 = (uint8_t *) utmp + (k / (32 / blocklen)) * scale_stride + 16; - - const int qh_shift = (k / (32 / blocklen)) * 2; - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi1 = 0; - sumi2 = 0; - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int b_qs_offset = k * ncols_interleaved * blocklen + j * blocklen + i; - - const int qh_idx = (k * blocklen + i) % 32; - const int qh_chunk = qh_idx / blocklen; - const int qh_pos = qh_idx % blocklen; - const int b_qh_offset = - qh_chunk * (blocklen * ncols_interleaved) + j * blocklen + qh_pos; - - const uint8_t qh_val = b_ptr[l].qh[b_qh_offset]; - const uint8_t h0 = (qh_val >> qh_shift) & 1; - const uint8_t h1 = (qh_val >> (qh_shift + 1)) & 1; - - const int v0 = (int8_t) ((b_ptr[l].qs[b_qs_offset] & 0xF) | (h0 << 4)); - const int v1 = (int8_t) ((b_ptr[l].qs[b_qs_offset] >> 4) | (h1 << 4)); - - const int q8_offset = (k / (32 / blocklen)) * 256 + - (k % (32 / blocklen)) * 4 * blocklen + m * blocklen + i; - - sumi1 = (v0 * a_ptr[l].qs[q8_offset]); - sumi2 = (v1 * a_ptr[l].qs[q8_offset + 128]); - sumi1 = sumi1 * scales_0[j]; - sumi2 = sumi2 * scales_1[j]; - sumi += sumi1 + sumi2; - } - sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d[m]; - } - } - } - for (int sb = 0; sb < 8; sb++) { - uint8_t * mins = (uint8_t *) utmp + 8 + sb * 16; - for (int m = 0; m < 4; m++) { - const int16_t * bsums = a_ptr[l].bsums + (sb * 8) + (m * 4) - ((sb % 2) * 6); - for (int j = 0; j < ncols_interleaved; j++) { - sum_minf[m][j] += mins[j] * (bsums[0] + bsums[1]) * - GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d[m]; - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j] - sum_minf[m][j]; - } - } - } - } -} - -extern "C" { - -void ggml_gemv_q4_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 4; - const int blocklen = 4; - - assert(nr == 1); - assert(n % qk == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - float sumf[4]; - int sumi; - - const block_q8_0 * a_ptr = (const block_q8_0 *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_0x4 * b_ptr = (const block_q4_0x4 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); - const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); - sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])) >> 4; - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); - } - } - } - for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; - } -} - -void ggml_gemv_q4_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 4; - const int blocklen = 8; - - assert (n % qk == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - float sumf[4]; - int sumi; - - const block_q8_0 * a_ptr = (const block_q8_0 *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_0x4 * b_ptr = (const block_q4_0x4 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); - const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); - sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])) >> 4; - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); - } - } - } - for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; - } -} - -void ggml_gemv_q4_0_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 8; - const int blocklen = 8; - - assert (n % qk == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - float sumf[8]; - int sumi; - - const block_q8_0 * a_ptr = (const block_q8_0 *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_0x8 * b_ptr = (const block_q4_0x8 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); - const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); - sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])) >> 4; - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); - } - } - } - for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; - } -} - -void ggml_gemv_q4_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK_K; - const int nb = n / qk; - const int ncols_interleaved = 8; - const int blocklen = 4; - static const uint32_t kmask1 = 0x3f3f3f3f; - static const uint32_t kmask2 = 0x0f0f0f0f; - static const uint32_t kmask3 = 0x03030303; - - assert (n % qk == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - float sumf[8]; - float sum_minf[8]; - uint32_t utmp[32]; - int sumi1; - int sumi2; - int sumi; - - const block_q8_K * a_ptr = (const block_q8_K *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_Kx8 * b_ptr = (const block_q4_Kx8 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) { - sumf[j] = 0.0; - sum_minf[j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int sb = 0; sb < 8; sb++) { - memcpy(utmp + sb * 4, b_ptr[l].scales + sb * 12, 12); - utmp[sb * 4 + 3] = ((utmp[sb * 4 + 2] >> 4) & kmask2) | (((utmp[sb * 4 + 1] >> 6) & kmask3) << 4); - const uint32_t uaux_0 = utmp[sb * 4 + 1] & kmask1; - utmp[sb * 4 + 1] = (utmp[sb * 4 + 2] & kmask2) | (((utmp[sb * 4 + 0] >> 6) & kmask3) << 4); - utmp[sb * 4 + 2] = uaux_0; - utmp[sb * 4 + 0] &= kmask1; - } - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - uint8_t * scales_0 = (uint8_t *) utmp + (k / 8) * 32; - uint8_t * scales_1 = (uint8_t *) utmp + (k / 8) * 32 + 16; - for (int j = 0; j < ncols_interleaved; j++) { - sumi1 = 0; - sumi2 = 0; - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF); - const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4); - sumi1 = (v0 * a_ptr[l].qs[(k / 8) * 64 + (k % 8) * blocklen + i]); - sumi2 = (v1 * a_ptr[l].qs[(k / 8) * 64 + (k % 8) * blocklen + i + 32]); - sumi1 = sumi1 * scales_0[j]; - sumi2 = sumi2 * scales_1[j]; - sumi += sumi1 + sumi2; - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d; - } - } - for (int sb = 0; sb < 8; sb++) { - uint8_t * mins = (uint8_t *) utmp + 8 + sb * 16; - for (int j = 0; j < ncols_interleaved; j++) { - sum_minf[j] += mins[j] * (a_ptr[l].bsums[sb * 2] + a_ptr[l].bsums[sb * 2 + 1]) * GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d; - } - } - } - for (int j = 0; j < ncols_interleaved; j++) { - s[x * ncols_interleaved + j] = sumf[j] - sum_minf[j]; - } - } -} - -void ggml_gemv_q4_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK_K; - const int nb = n / qk; - const int ncols_interleaved = 8; - const int blocklen = 8; - static const uint32_t kmask1 = 0x3f3f3f3f; - static const uint32_t kmask2 = 0x0f0f0f0f; - static const uint32_t kmask3 = 0x03030303; - - assert (n % qk == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - float sumf[8]; - float sum_minf[8]; - uint32_t utmp[32]; - int sumi1; - int sumi2; - int sumi; - - const block_q8_K * a_ptr = (const block_q8_K *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_Kx8 * b_ptr = (const block_q4_Kx8 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) { - sumf[j] = 0.0; - sum_minf[j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int sb = 0; sb < 8; sb++) { - memcpy(utmp + sb * 4, b_ptr[l].scales + sb * 12, 12); - utmp[sb * 4 + 3] = ((utmp[sb * 4 + 2] >> 4) & kmask2) | (((utmp[sb * 4 + 1] >> 6) & kmask3) << 4); - const uint32_t uaux_0 = utmp[sb * 4 + 1] & kmask1; - utmp[sb * 4 + 1] = (utmp[sb * 4 + 2] & kmask2) | (((utmp[sb * 4 + 0] >> 6) & kmask3) << 4); - utmp[sb * 4 + 2] = uaux_0; - utmp[sb * 4 + 0] &= kmask1; - } - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - uint8_t *scales_0 = (uint8_t*) utmp + (k / 4) * 32; - uint8_t *scales_1 = (uint8_t*) utmp + (k / 4) * 32 + 16; - for (int j = 0; j < ncols_interleaved; j++) { - sumi1 = 0; - sumi2 = 0; - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF); - const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4); - sumi1 = (v0 * a_ptr[l].qs[(k >> 2) * 64 + (k % 4) * blocklen + i]); - sumi2 = (v1 * a_ptr[l].qs[(k >> 2) * 64 + (k % 4) * blocklen + i + 32]); - sumi1 = sumi1 * scales_0[j]; - sumi2 = sumi2 * scales_1[j]; - sumi += sumi1 + sumi2; - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d; - } - } - for (int sb = 0; sb < 8; sb++) { - uint8_t *mins = (uint8_t*) utmp + 8 + sb * 16; - for (int j = 0; j < ncols_interleaved; j++) { - sum_minf[j] += mins[j] * (a_ptr[l].bsums[sb * 2] + a_ptr[l].bsums[sb * 2 + 1]) * GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d; - } - } - } - for (int j = 0; j < ncols_interleaved; j++) { - s[x * ncols_interleaved + j] = sumf[j] - sum_minf[j]; - } - } -} - -void ggml_gemv_q2_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK_K; - const int nb = n / qk; - const int ncols_interleaved = 8; - const int blocklen = 8; - - assert (n % qk == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - float sumf[8]; - float sum_minf[8]; - int sumi1,sumi2,sumi3,sumi4; - int sumi; - - const block_q8_K * a_ptr = (const block_q8_K *)vy; - for(int x = 0; x < nc / ncols_interleaved; x++) { - const block_q2_Kx8 * b_ptr = (const block_q2_Kx8 *) vx + (x * nb); - for (int j = 0; j < ncols_interleaved; j++) { - sumf[j] = 0.0; - sum_minf[j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (4 * blocklen)); k++) { - const uint8_t *scales_0 = b_ptr[l].scales + (k / 4) * 64 ; - const uint8_t *scales_1 = b_ptr[l].scales + (k / 4) * 64 + 16; - const uint8_t *scales_2 = b_ptr[l].scales + (k / 4) * 64 + 32; - const uint8_t *scales_3 = b_ptr[l].scales + (k / 4) * 64 + 48; - for (int j = 0; j < ncols_interleaved; j++) { - sumi1 = 0; - sumi2 = 0; - sumi3 = 0; - sumi4 = 0; - sumi = 0; - int offset = ((k / 2) % 2) + j * 2; - for (int i = 0; i < blocklen; ++i){ - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 3); - const int v1 = (int8_t) ((b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 2 ) & 3); - const int v2 = (int8_t) ((b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4 ) & 3); - const int v3 = (int8_t) ((b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 6 ) & 3); - sumi1 = (v0 * a_ptr[l].qs[(k >> 2) * 128 + (k % 4) * blocklen + i]); - sumi2 = (v1 * a_ptr[l].qs[(k >> 2) * 128 + (k % 4) * blocklen + i + 32]); - sumi3 = (v2 * a_ptr[l].qs[(k >> 2) * 128 + (k % 4) * blocklen + i + 64]); - sumi4 = (v3 * a_ptr[l].qs[(k >> 2) * 128 + (k % 4) * blocklen + i + 96]); - - sumi1 = sumi1 * (scales_0[offset] & 0xF); - sumi2 = sumi2 * (scales_1[offset] & 0xF); - sumi3 = sumi3 * (scales_2[offset] & 0xF); - sumi4 = sumi4 * (scales_3[offset] & 0xF); - sumi += sumi1 + sumi2 + sumi3 + sumi4; - } - sumf[j] += sumi * GGML_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d; - } - } - for(int sb = 0; sb < 8; sb++) { - const uint8_t *mins = b_ptr[l].scales + sb * 16; - for(int j = 0; j < ncols_interleaved; j++){ - sum_minf[j] += ((mins[j * 2] >> 4) * a_ptr[l].bsums[sb * 2] + (mins[(j * 2)+ 1] >> 4) * a_ptr[l].bsums[sb * 2 + 1]) * GGML_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d; - } - } - } - for (int j = 0; j < ncols_interleaved; j++) { - s[x * ncols_interleaved + j] = sumf[j] - sum_minf[j]; - } - } -} - -void ggml_gemv_q5_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - ggml_gemv_q5_K_NxM_q8_K_generic_impl<4, 8>(n, s, bs, vx, vy, nr, nc); -} - -void ggml_gemv_q5_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - ggml_gemv_q5_K_NxM_q8_K_generic_impl<8, 8>(n, s, bs, vx, vy, nr, nc); -} - - -void ggml_gemv_q6_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - ggml_gemv_q6_K_NxM_q8_K_generic_impl<4, 8>(n, s, bs, vx, vy, nr, nc); -} - -void ggml_gemv_q6_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - ggml_gemv_q6_K_NxM_q8_K_generic_impl<8, 8>(n, s, bs, vx, vy, nr, nc); -} - -void ggml_gemv_iq4_nl_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 4; - const int blocklen = 4; - - assert(nr == 1); - assert(n % qk == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - float sumf[4]; - int sumi; - - const block_q8_0 * a_ptr = (const block_q8_0 *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_iq4_nlx4 * b_ptr = (const block_iq4_nlx4 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; - const int v1 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; - sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])); - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); - } - } - } - for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; - } -} - -void ggml_gemv_iq4_nl_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 8; - const int blocklen = 8; - - assert(nr == 1); - assert(n % qk == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - float sumf[8]; - int sumi; - - const block_q8_0 * a_ptr = (const block_q8_0 *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_iq4_nlx8 * b_ptr = (const block_iq4_nlx8 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; - const int v1 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; - sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])); - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); - } - } - } - for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; - } -} - -void ggml_gemv_mxfp4_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 4; - const int blocklen = 4; - - assert(nr == 1); - assert(n % qk == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - float sumf[4]; - int sumi; - - const block_q8_0 * a_ptr = (const block_q8_0 *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_mxfp4x4 * b_ptr = (const block_mxfp4x4 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; - const int v1 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; - sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])); - } - sumf[j] += sumi * GGML_CPU_E8M0_TO_FP32_HALF(b_ptr[l].e[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); - } - } - } - for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; - } -} - -void ggml_gemv_mxfp4_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 8; - const int blocklen = 8; - - assert(nr == 1); - assert(n % qk == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - float sumf[8]; - int sumi; - - const block_q8_0 * a_ptr = (const block_q8_0 *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_mxfp4x8 * b_ptr = (const block_mxfp4x8 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; - const int v1 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; - sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])); - } - sumf[j] += sumi * GGML_CPU_E8M0_TO_FP32_HALF(b_ptr[l].e[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); - } - } - } - for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; - } -} - -void ggml_gemv_q8_0_4x4_q8_0_generic(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * GGML_RESTRICT vy, - int nr, - int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 4; - const int blocklen = 4; - - assert(nr == 1); - assert(n % qk == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - float sumf[4]; - int sumi; - - const block_q8_0 * a_ptr = (const block_q8_0 *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q8_0x4 * b_ptr = (const block_q8_0x4 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) { - sumf[j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / blocklen); k++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i]; - sumi += v0 * a_ptr[l].qs[k * blocklen + i]; - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); - } - } - } - for (int j = 0; j < ncols_interleaved; j++) { - s[x * ncols_interleaved + j] = sumf[j]; - } - } -} - -void ggml_gemv_q8_0_4x8_q8_0_generic(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * GGML_RESTRICT vy, - int nr, - int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 4; - const int blocklen = 8; - - assert(nr == 1); - assert(n % qk == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - float sumf[4]; - int sumi; - - const block_q8_0 * a_ptr = (const block_q8_0 *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q8_0x4 * b_ptr = (const block_q8_0x4 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) { - sumf[j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / blocklen); k++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i]; - sumi += v0 * a_ptr[l].qs[k * blocklen + i]; - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); - } - } - } - for (int j = 0; j < ncols_interleaved; j++) { - s[x * ncols_interleaved + j] = sumf[j]; - } - } -} - -// Only enable these for RISC-V. -#if defined __riscv_zvfh -void ggml_gemv_q4_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 16; - const int blocklen = 1; - - assert (n % qk == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - float sumf[16]; - int sumi; - - const block_q8_0 * a_ptr = (const block_q8_0 *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_0x16 * b_ptr = (const block_q4_0x16 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); - const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); - sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])) >> 4; - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); - } - } - } - for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; - } -} - -void ggml_gemv_q4_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK_K; - const int nb = n / qk; - const int ncols_interleaved = 16; - const int blocklen = 1; - assert (n % qk == 0); - assert (nc % ncols_interleaved == 0); - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - float sumf[16]; - float sum_minf[16]; - uint8_t scales[128]; - uint8_t mins[128]; - int sumi1; - int sumi2; - int sumi; - const block_q8_K * a_ptr = (const block_q8_K *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_Kx16 * b_ptr = (const block_q4_Kx16 *) vx + (x * nb); - for (int j = 0; j < ncols_interleaved; j++) { - sumf[j] = 0.0f; - sum_minf[j] = 0.0f; - } - for (int l = 0; l < nb; l++) { - for (int i = 0; i < 128; i++) { - scales[i] = b_ptr[l].scales[i] & 0x0F; - mins[i] = b_ptr[l].scales[i] >> 4; - } - for (int i = 0; i < 64; i++) { - scales[i] |= (b_ptr[l].scales[128 + i] & 0x03) << 4; - mins[i] |= (b_ptr[l].scales[128 + i] & 0x0C) << 2; - scales[i + 64] |= (b_ptr[l].scales[128 + i] & 0x30); - mins[i + 64] |= (b_ptr[l].scales[128 + i] & 0xC0) >> 2; - } - for (int sb = 0; sb < 8; sb++) { - uint8_t *min = &mins[sb * 16]; - for (int j = 0; j < ncols_interleaved; j++) { - sum_minf[j] += min[j] * (a_ptr[l].bsums[sb * 2] + a_ptr[l].bsums[sb * 2 + 1]) * GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d; - } - } - for (int sb = 0; sb < 8; sb += 2) { - uint8_t *scales_0 = &scales[sb * 16]; - uint8_t *scales_1 = &scales[(sb + 1) * 16]; - for (int i = 0; i < QK4_0; i++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi1 = 0; - sumi2 = 0; - sumi = 0; - const int v0 = (int8_t) (b_ptr[l].qs[sb * 256 + i * 16 + j] & 0xF); - const int v1 = (int8_t) (b_ptr[l].qs[sb * 256 + i * 16 + j] >> 4); - sumi1 = (v0 * a_ptr[l].qs[sb * 32 + i]); - sumi2 = (v1 * a_ptr[l].qs[sb * 32 + 32 + i]); - sumi1 = sumi1 * scales_0[j]; - sumi2 = sumi2 * scales_1[j]; - sumi += sumi1 + sumi2; - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d; - } - } - } - } - for (int j = 0; j < ncols_interleaved; j++) { - s[x * ncols_interleaved + j] = sumf[j] - sum_minf[j]; - } - } -} - -void ggml_gemv_iq4_nl_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 16; - const int blocklen = 1; - - assert(nr == 1); - assert(n % qk == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - float sumf[16]; - int sumi; - - const block_q8_0 * a_ptr = (const block_q8_0 *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_iq4_nlx16 * b_ptr = (const block_iq4_nlx16 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; - const int v1 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; - sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])); - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); - } - } - } - for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; - } -} - -void ggml_gemv_q8_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 16; - const int blocklen = 1; - - assert(nr == 1); - assert(n % qk == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - float sumf[16]; - int sumi; - - const block_q8_0 * a_ptr = (const block_q8_0 *) vy; - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q8_0x16 * b_ptr = (const block_q8_0x16 *) vx + (x * nb); - - for (int j = 0; j < ncols_interleaved; j++) { - sumf[j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / blocklen); k++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i]; - sumi += v0 * a_ptr[l].qs[k * blocklen + i]; - } - sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); - } - } - } - for (int j = 0; j < ncols_interleaved; j++) { - s[x * ncols_interleaved + j] = sumf[j]; - } - } -} - -void ggml_gemv_q2_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - assert(n % QK_K == 0); - assert(nr == 1); - assert(nc % 16 == 0); - - UNUSED(bs); - UNUSED(nr); - - const int nb = n / QK_K; - const block_q2_Kx16 * x = (const block_q2_Kx16 *)vx; - const block_q8_K * y = (const block_q8_K *)vy; - - // Layout: Even-Low(0,2,4,6), Odd-Low(1,3,5,7), Even-High(8...), Odd-High(9...) - const int sb_perm[16] = { - 0, 4, 1, 5, 2, 6, 3, 7, // 0-7 - 8, 12, 9, 13, 10, 14, 11, 15 // 8-15 - }; - - for (int col_tile = 0; col_tile < nc; col_tile += 16) { - const block_q2_Kx16 * x_ptr = x + (col_tile / 16) * nb; - const block_q8_K * y_ptr = y; - - float sumf[16] = {0}; - - // Loop over K-blocks - for (int k_block = 0; k_block < nb; ++k_block) { - int32_t isum[16] = {0}; - int32_t summs[16] = {0}; - - const uint8_t * qs_rhs = x_ptr[k_block].qs; - const uint8_t * sc_rhs = x_ptr[k_block].scales; - const int8_t * qs_lhs = y_ptr[k_block].qs; - const int16_t * bs_lhs = y_ptr[k_block].bsums; - - // Iterate over sub-blocks 0..15 - for (int sb = 0; sb < 16; ++sb) { - // Correction Term - int16_t bsum = bs_lhs[sb]; - int scale_offset = sb_perm[sb] * 16; - - for (int col = 0; col < 16; ++col) { - uint8_t sc_val = sc_rhs[scale_offset + col]; - summs[col] += bsum * (sc_val >> 4); // Min is high 4 bits - } - - // Main Dot Product - // Calculate base offsets for Q2 unpacking based on SB - int byte_base; - if (sb < 8) byte_base = (sb % 2 == 0) ? 0 : 16; - else byte_base = (sb % 2 == 0) ? 32 : 48; - - int shift = ((sb / 2) % 4) * 2; - - for (int col = 0; col < 16; ++col) { - uint8_t sc_val = sc_rhs[scale_offset + col]; - int32_t d_sb = sc_val & 0xF; // Scale is low 4 bits - - // Process 16 elements (l=0..15) - for (int l = 0; l < 16; ++l) { - // Q2: Interleaved by column. Byte `l` contains 4 k-values. - int qs_idx = (byte_base + l) * 16 + col; - uint8_t q2_val = (qs_rhs[qs_idx] >> shift) & 3; - - // Q8: Linear access - int k = sb * 16 + l; - int8_t q8_val = qs_lhs[k]; - - isum[col] += q8_val * q2_val * d_sb; - } - } - } - - // Finalize K-Block - for (int col = 0; col < 16; ++col) { - float d_lhs = y_ptr[k_block].d; - float d_rhs = GGML_FP16_TO_FP32(x_ptr[k_block].d[col]); - float dm_rhs = GGML_FP16_TO_FP32(x_ptr[k_block].dmin[col]); - - float d_all = d_lhs * d_rhs; - float d_min = d_lhs * dm_rhs; - - sumf[col] += (isum[col] * d_all) - (summs[col] * d_min); - } - } - - for (int col = 0; col < 16; ++col) { - s[col_tile + col] = sumf[col]; - } - } -} -#endif - -void ggml_gemm_q4_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 4; - const int blocklen = 4; - - assert (n % qk == 0); - assert (nr % 4 == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - { - float sumf[4][4]; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_0x4 * b_ptr = (const block_q4_0x4 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); - const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); - sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + - (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])) >> 4; - } - sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } - } -} - -void ggml_gemm_q4_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 4; - const int blocklen = 8; - - assert (n % qk == 0); - assert (nr % 4 == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - float sumf[4][4]; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_0x4 * b_ptr = (const block_q4_0x4 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); - const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); - sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + - (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])) >> 4; - } - sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } -} - -void ggml_gemm_q4_0_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 8; - const int blocklen = 8; - - assert (n % qk == 0); - assert (nr % 4 == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - float sumf[4][8]; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_0x8 * b_ptr = (const block_q4_0x8 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); - const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); - sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + - (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])) >> 4; - } - sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } -} - -void ggml_gemm_q4_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK_K; - const int nb = n / qk; - const int ncols_interleaved = 8; - const int blocklen = 4; - static const uint32_t kmask1 = 0x3f3f3f3f; - static const uint32_t kmask2 = 0x0f0f0f0f; - static const uint32_t kmask3 = 0x03030303; - - assert (n % qk == 0); - assert (nr % 4 == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - float sumf[4][8]; - float sum_minf[4][8]; - uint32_t utmp[32]; - int sumi1; - int sumi2; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_Kx4 * a_ptr = (const block_q8_Kx4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_Kx8 * b_ptr = (const block_q4_Kx8 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumf[m][j] = 0.0; - sum_minf[m][j] = 0.0; - } - } - for (int l = 0; l < nb; l++) { - for (int sb = 0; sb < 8; sb++) { - memcpy(utmp + sb * 4, b_ptr[l].scales + sb * 12, 12); - utmp[sb * 4 + 3] = ((utmp[sb * 4 + 2] >> 4) & kmask2) | (((utmp[sb * 4 + 1] >> 6) & kmask3) << 4); - const uint32_t uaux_0 = utmp[sb * 4 + 1] & kmask1; - utmp[sb * 4 + 1] = (utmp[sb * 4 + 2] & kmask2) | (((utmp[sb * 4 + 0] >> 6) & kmask3) << 4); - utmp[sb * 4 + 2] = uaux_0; - utmp[sb * 4 + 0] &= kmask1; - } - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - uint8_t * scales_0 = (uint8_t *) utmp + (k / 8) * 32; - uint8_t * scales_1 = (uint8_t *) utmp + (k / 8) * 32 + 16; - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi1 = 0; - sumi2 = 0; - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF); - const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4); - sumi1 = (v0 * a_ptr[l].qs[(k / 8) * 256 + (k % 8) * 4 * blocklen + m * blocklen + i]); - sumi2 = (v1 * a_ptr[l].qs[(k / 8) * 256 + (k % 8) * 4 * blocklen + m * blocklen + i + 128]); - sumi1 = sumi1 * scales_0[j]; - sumi2 = sumi2 * scales_1[j]; - sumi += sumi1 + sumi2; - } - sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d[m]; - } - } - } - for (int sb = 0; sb < 8; sb++) { - uint8_t * mins = (uint8_t *) utmp + 8 + sb * 16; - for(int m = 0; m < 4; m++) { - const int16_t * bsums = a_ptr[l].bsums + (sb * 8) + (m * 4) - ((sb % 2) * 6); - for(int j = 0; j < ncols_interleaved; j++) { - sum_minf[m][j] += mins[j] * (bsums[0] + bsums[1]) * GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d[m]; - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j] - sum_minf[m][j]; - } - } - } - } -} - -void ggml_gemm_q4_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK_K; - const int nb = n / qk; - const int ncols_interleaved = 8; - const int blocklen = 8; - static const uint32_t kmask1 = 0x3f3f3f3f; - static const uint32_t kmask2 = 0x0f0f0f0f; - static const uint32_t kmask3 = 0x03030303; - - assert (n % qk == 0); - assert (nr % 4 == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(bs); - - float sumf[4][8]; - float sum_minf[4][8]; - uint32_t utmp[32]; - int sumi1; - int sumi2; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_Kx4 * a_ptr = (const block_q8_Kx4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_Kx8 * b_ptr = (const block_q4_Kx8 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumf[m][j] = 0.0; - sum_minf[m][j] = 0.0; - } - } - for (int l = 0; l < nb; l++) { - for (int sb = 0; sb < 8; sb++) { - memcpy(utmp + sb * 4, b_ptr[l].scales + sb * 12, 12); - utmp[sb * 4 + 3] = ((utmp[sb * 4 + 2] >> 4) & kmask2) | (((utmp[sb * 4 + 1] >> 6) & kmask3) << 4); - const uint32_t uaux_0 = utmp[sb * 4 + 1] & kmask1; - utmp[sb * 4 + 1] = (utmp[sb * 4 + 2] & kmask2) | (((utmp[sb * 4 + 0] >> 6) & kmask3) << 4); - utmp[sb * 4 + 2] = uaux_0; - utmp[sb * 4 + 0] &= kmask1; - } - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - uint8_t *scales_0 = (uint8_t*) utmp + (k / 4) * 32; - uint8_t *scales_1 = (uint8_t*) utmp + (k / 4) * 32 + 16; - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi1 = 0; - sumi2 = 0; - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF); - const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4); - sumi1 = (v0 * a_ptr[l].qs[(k >> 2) * 256 + (k % 4) * 4 * blocklen + m * blocklen + i]); - sumi2 = (v1 * a_ptr[l].qs[(k >> 2) * 256 + (k % 4) * 4 * blocklen + m * blocklen + i + 128]); - sumi1 = sumi1 * scales_0[j]; - sumi2 = sumi2 * scales_1[j]; - sumi += sumi1 + sumi2; - } - sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d[m]; - } - } - } - for (int sb = 0; sb < 8; sb++) { - uint8_t *mins = (uint8_t*) utmp + 8 + sb * 16; - for(int m = 0; m < 4; m++) { - const int16_t *bsums = a_ptr[l].bsums + (sb * 8) + (m * 4) - ((sb % 2) * 6); - for(int j = 0; j < ncols_interleaved; j++) { - sum_minf[m][j] += mins[j] * (bsums[0] + bsums[1]) * GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d[m]; - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j] - sum_minf[m][j]; - } - } - } - } -} - -void ggml_gemm_q2_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK_K; - const int nb = n / qk; - const int ncols_interleaved = 8; - const int blocklen = 8; - - assert (n % qk == 0); - assert (nr % 4 == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - float sumf[4][8]; - float sum_minf[4][8]; - int sumi1, sumi2, sumi3, sumi4; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_Kx4 * a_ptr = (const block_q8_Kx4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q2_Kx8 * b_ptr = (const block_q2_Kx8 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumf[m][j] = 0.0; - sum_minf[m][j] = 0.0; - } - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (4 * blocklen)); k++) { - - const uint8_t *scales_0 = b_ptr[l].scales + (k / 4) * 64 ; - const uint8_t *scales_1 = b_ptr[l].scales + (k / 4) * 64 + 16; - const uint8_t *scales_2 = b_ptr[l].scales + (k / 4) * 64 + 32; - const uint8_t *scales_3 = b_ptr[l].scales + (k / 4) * 64 + 48; - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi1 = 0; - sumi2 = 0; - sumi3 = 0; - sumi4 = 0; - sumi = 0; - int offset = ((k / 2) % 2) + j * 2; - for (int i = 0; i < blocklen; ++i){ - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 3); - const int v1 = (int8_t) ((b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 2 ) & 3); - const int v2 = (int8_t) ((b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4 ) & 3); - const int v3 = (int8_t) ((b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 6 ) & 3); - sumi1 = (v0 * a_ptr[l].qs[(k >> 2) * 512 + (k % 4) * 4 * blocklen + m * blocklen + i]); - sumi2 = (v1 * a_ptr[l].qs[(k >> 2) * 512 + (k % 4) * 4 * blocklen + m * blocklen + i + 128]); - sumi3 = (v2 * a_ptr[l].qs[(k >> 2) * 512 + (k % 4) * 4 * blocklen + m * blocklen + i + 256]); - sumi4 = (v3 * a_ptr[l].qs[(k >> 2) * 512 + (k % 4) * 4 * blocklen + m * blocklen + i + 384]); - sumi1 = sumi1 * (scales_0[offset] & 0xF); - sumi2 = sumi2 * (scales_1[offset] & 0xF); - sumi3 = sumi3 * (scales_2[offset] & 0xF); - sumi4 = sumi4 * (scales_3[offset] & 0xF); - sumi += sumi1 + sumi2 + sumi3 + sumi4; - } - sumf[m][j] += sumi * GGML_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d[m]; - } - } - } - for(int sb = 0; sb < 8; sb++) { - const uint8_t *mins = b_ptr[l].scales + sb * 16; - for(int m = 0; m < 4; m++) { - const int16_t *bsums = a_ptr[l].bsums + (sb * 8) + (m * 4) - ((sb % 2) * 6); - for(int j = 0; j < ncols_interleaved; j++) { - int mins_prod = ((mins[j * 2] >> 4) * bsums[0] + (mins[(j * 2)+ 1] >> 4) * bsums[1]); - sum_minf[m][j] += (mins_prod) * GGML_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d[m]; - } - } - } - } - - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j] - sum_minf[m][j]; - } - } - } - } -} - -void ggml_gemm_q5_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - ggml_gemm_q5_K_NxM_q8_K_generic_impl<4, 8>(n, s, bs, vx, vy, nr, nc); -} - -void ggml_gemm_q5_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - ggml_gemm_q5_K_NxM_q8_K_generic_impl<8, 8>(n, s, bs, vx, vy, nr, nc); -} - -void ggml_gemm_q6_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - ggml_gemm_q6_K_NxM_q8_K_generic_impl<4, 8>(n, s, bs, vx, vy, nr, nc); -} - -void ggml_gemm_q6_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - ggml_gemm_q6_K_NxM_q8_K_generic_impl<8, 8>(n, s, bs, vx, vy, nr, nc); -} - -void ggml_gemm_iq4_nl_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 4; - const int blocklen = 4; - - assert (n % qk == 0); - assert (nr % 4 == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - { - float sumf[4][4]; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_iq4_nlx4 * b_ptr = (const block_iq4_nlx4 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; - const int v1 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; - sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + - (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])); - } - sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } - } -} - -void ggml_gemm_iq4_nl_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 8; - const int blocklen = 8; - - assert(n % qk == 0); - assert(nr % 4 == 0); - assert(nc % ncols_interleaved == 0); - - float sumf[4][8]; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_iq4_nlx8 * b_ptr = (const block_iq4_nlx8 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; - const int v1 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; - sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + - (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])); - } - sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } -} - -void ggml_gemm_mxfp4_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 4; - const int blocklen = 4; - - assert(n % qk == 0); - assert(nr % 4 == 0); - assert(nc % ncols_interleaved == 0); - - float sumf[4][4]; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_mxfp4x4 * b_ptr = (const block_mxfp4x4 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; - const int v1 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; - sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + - (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])); - } - sumf[m][j] += sumi * GGML_CPU_E8M0_TO_FP32_HALF(b_ptr[l].e[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } -} - -void ggml_gemm_mxfp4_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 8; - const int blocklen = 8; - - assert(n % qk == 0); - assert(nr % 4 == 0); - assert(nc % ncols_interleaved == 0); - - float sumf[4][8]; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_mxfp4x8 * b_ptr = (const block_mxfp4x8 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; - const int v1 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; - sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + - (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])); - } - sumf[m][j] += sumi * GGML_CPU_E8M0_TO_FP32_HALF(b_ptr[l].e[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } -} - -void ggml_gemm_q8_0_4x4_q8_0_generic(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * GGML_RESTRICT vy, - int nr, - int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 4; - const int blocklen = 4; - - assert(n % qk == 0); - assert(nr % 4 == 0); - assert(nc % ncols_interleaved == 0); - - float sumf[4][4]; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q8_0x4 * b_ptr = (const block_q8_0x4 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumf[m][j] = 0.0; - } - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / blocklen); k++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i]; - sumi += v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]; - } - sumf[m][j] += - sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } - } -} - - - -void ggml_gemm_q8_0_4x8_q8_0_generic(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * GGML_RESTRICT vy, - int nr, - int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 4; - const int blocklen = 8; - - assert(n % qk == 0); - assert(nr % 4 == 0); - assert(nc % ncols_interleaved == 0); - - float sumf[4][4]; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q8_0x4 * b_ptr = (const block_q8_0x4 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumf[m][j] = 0.0; - } - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / blocklen); k++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i]; - sumi += v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]; - } - sumf[m][j] += - sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } - } -} - -// Only enable these for RISC-V. -#if defined __riscv_zvfh -void ggml_gemm_q4_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 16; - const int blocklen = 1; - - assert (n % qk == 0); - assert (nr % 4 == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - float sumf[4][16]; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_0x16 * b_ptr = (const block_q4_0x16 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); - const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); - sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + - (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])) >> 4; - } - sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } -} - -void ggml_gemm_q4_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK_K; - const int nb = n / qk; - const int ncols_interleaved = 16; - const int blocklen = 1; - - assert (n % qk == 0); - assert (nr % 4 == 0); - assert (nc % ncols_interleaved == 0); - - UNUSED(s); - UNUSED(bs); - UNUSED(vx); - UNUSED(vy); - UNUSED(nr); - UNUSED(nc); - UNUSED(nb); - UNUSED(ncols_interleaved); - UNUSED(blocklen); - - float sumf[4][16]; - float sum_minf[4][16]; - uint8_t scales[128]; - uint8_t mins[128]; - int sumi1; - int sumi2; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_Kx4 * a_ptr = (const block_q8_Kx4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q4_Kx16 * b_ptr = (const block_q4_Kx16 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumf[m][j] = 0.0; - sum_minf[m][j] = 0.0; - } - } - for (int l = 0; l < nb; l++) { - for (int i = 0; i < 128; i++) { - scales[i] = b_ptr[l].scales[i] & 0x0F; - mins[i] = b_ptr[l].scales[i] >> 4; - } - for (int i = 0; i < 64; i++) { - scales[i] |= (b_ptr[l].scales[128 + i] & 0x03) << 4; - mins[i] |= (b_ptr[l].scales[128 + i] & 0x0C) << 2; - scales[i + 64] |= (b_ptr[l].scales[128 + i] & 0x30); - mins[i + 64] |= (b_ptr[l].scales[128 + i] & 0xC0) >> 2; - } - - for (int sb = 0; sb < 8; sb++) { - uint8_t *min = &mins[sb * 16]; - for(int m = 0; m < 4; m++) { - const int16_t bsums = a_ptr[l].bsums[sb * 8 + m] + a_ptr[l].bsums[sb * 8 + m + 4]; - for(int j = 0; j < ncols_interleaved; j++) { - sum_minf[m][j] += min[j] * bsums * GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d[m]; - } - } - } - - for (int sb = 0; sb < 8; sb += 2) { - uint8_t *scales_0 = &scales[sb * 16]; - uint8_t *scales_1 = &scales[(sb + 1) * 16]; - - for (int i = 0; i < QK4_0; i++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi1 = 0; - sumi2 = 0; - sumi = 0; - - const int v0 = (int8_t) (b_ptr[l].qs[sb * 256 + i * 16 + j] & 0xF); - const int v1 = (int8_t) (b_ptr[l].qs[sb * 256 + i * 16 + j] >> 4); - sumi1 = (v0 * a_ptr[l].qs[sb * 4 * 32 + i * 4 + m]); - sumi2 = (v1 * a_ptr[l].qs[sb * 4 * 32 + 32 * 4 + i * 4 + m]); - sumi1 = sumi1 * scales_0[j]; - sumi2 = sumi2 * scales_1[j]; - sumi += sumi1 + sumi2; - - sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d[m]; - } - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j] - sum_minf[m][j]; - } - } - } - } -} - -void ggml_gemm_iq4_nl_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 16; - const int blocklen = 1; - - assert(n % qk == 0); - assert(nr % 4 == 0); - assert(nc % ncols_interleaved == 0); - - float sumf[4][16]; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_iq4_nlx16 * b_ptr = (const block_iq4_nlx16 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / (2 * blocklen)); k++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; - const int v1 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; - sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + - (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + (qk / 2) * 4])); - } - sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } -} - -void ggml_gemm_q8_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - const int qk = QK8_0; - const int nb = n / qk; - const int ncols_interleaved = 16; - const int blocklen = 1; - - assert(n % qk == 0); - assert(nr % 4 == 0); - assert(nc % ncols_interleaved == 0); - - float sumf[4][16]; - int sumi; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_q8_0x16 * b_ptr = (const block_q8_0x16 *) vx + (x * nb); - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumf[m][j] = 0.0; - } - } - for (int l = 0; l < nb; l++) { - for (int k = 0; k < (qk / blocklen); k++) { - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumi = 0; - for (int i = 0; i < blocklen; ++i) { - const int v0 = b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i]; - sumi += v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]; - } - sumf[m][j] += - sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); - } - } - } - } - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } - } -} - - -void ggml_gemm_q2_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { - assert(n % QK_K == 0); - assert(nr % 4 == 0); - assert(nc % 16 == 0); - const int nb = n / QK_K; - const block_q2_Kx16 * x = (const block_q2_Kx16 *)vx; - const block_q8_Kx4 * y = (const block_q8_Kx4 *)vy; - - const int sb_perm[16] = { - 0, 4, 1, 5, 2, 6, 3, 7, - 8, 12, 9, 13, 10, 14, 11, 15 - }; - - // Iterate Rows in tiles of 4 - for (int row_tile = 0; row_tile < nr; row_tile += 4) { - // Iterate Columns in tiles of 16 - for (int col_tile = 0; col_tile < nc; col_tile += 16) { - - const block_q2_Kx16 * x_ptr = x + (col_tile / 16) * nb; - const block_q8_Kx4 * y_ptr = y + (row_tile / 4) * nb; - - float sumf[4][16]; - memset(sumf, 0, sizeof(sumf)); - - for (int k_block = 0; k_block < nb; ++k_block) { - int32_t isum[4][16]; - int32_t summs[4][16]; - memset(isum, 0, sizeof(isum)); - memset(summs, 0, sizeof(summs)); - - const uint8_t * qs_rhs = x_ptr[k_block].qs; - const uint8_t * sc_rhs = x_ptr[k_block].scales; - const int8_t * qs_lhs = y_ptr[k_block].qs; - const int16_t * bs_lhs = y_ptr[k_block].bsums; - - for (int sb = 0; sb < 16; ++sb) { - int scale_offset = sb_perm[sb] * 16; - - int byte_base; - if (sb < 8) byte_base = (sb % 2 == 0) ? 0 : 16; - else byte_base = (sb % 2 == 0) ? 32 : 48; - int shift = ((sb / 2) % 4) * 2; - - for (int col = 0; col < 16; ++col) { - uint8_t sc_val = sc_rhs[scale_offset + col]; - int32_t d_sb = sc_val & 0xF; - int32_t m_sb = sc_val >> 4; - - // Correction Term - for (int r = 0; r < 4; ++r) { - int bsum_idx = (sb / 4) * 16 + r * 4 + (sb % 4); - summs[r][col] += bs_lhs[bsum_idx] * m_sb; - } - - // Main Dot Product - for (int l = 0; l < 16; ++l) { - int qs_idx = (byte_base + l) * 16 + col; - uint8_t q2_val = (qs_rhs[qs_idx] >> shift) & 3; - - // Calculate Q8 index for this specific k and row - int k = sb * 16 + l; - int q8_idx = (k / 4) * 16 + (k % 4); - - for (int r = 0; r < 4; ++r) { - // Add r*4 to jump to the correct row within the 4x4 chunk - int8_t q8_val = qs_lhs[q8_idx + r * 4]; - isum[r][col] += q8_val * q2_val * d_sb; - } - } - } - } - - // Finalize K-Block - for (int col = 0; col < 16; ++col) { - float d_rhs = GGML_FP16_TO_FP32(x_ptr[k_block].d[col]); - float dm_rhs = GGML_FP16_TO_FP32(x_ptr[k_block].dmin[col]); - - for (int r = 0; r < 4; ++r) { - float d_lhs = y_ptr[k_block].d[r]; - float d_all = d_lhs * d_rhs; - float d_min = d_lhs * dm_rhs; - sumf[r][col] += (isum[r][col] * d_all) - (summs[r][col] * d_min); - } - } - } - - for (int r = 0; r < 4; ++r) { - for (int col = 0; col < 16; ++col) { - s[(row_tile + r) * bs + (col_tile + col)] = sumf[r][col]; - } - } - } - } -} -#endif - -} // extern "C" - -static block_q8_0x4 make_block_q8_0x4(block_q8_0 * in, unsigned int blck_size_interleave) { - block_q8_0x4 out; - - for (int i = 0; i < 4; i++) { - out.d[i] = in[i].d; - } - - const int end = QK8_0 * 4 / blck_size_interleave; - for (int i = 0; i < end; ++i) { - int src_id = i % 4; - int src_offset = (i / 4) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], blck_size_interleave); - } - return out; -} - -static block_q4_0x4 make_block_q4_0x4(block_q4_0 * in, unsigned int blck_size_interleave) { - block_q4_0x4 out; - - for (int i = 0; i < 4; i++) { - out.d[i] = in[i].d; - } - - const int end = QK4_0 * 2 / blck_size_interleave; - - if (blck_size_interleave == 8) { - const uint64_t xor_mask = 0x8888888888888888ULL; - for (int i = 0; i < end; ++i) { - int src_id = i % 4; - int src_offset = (i / 4) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - uint64_t elems; - // Using memcpy to avoid unaligned memory accesses - memcpy(&elems, &in[src_id].qs[src_offset], sizeof(uint64_t)); - elems ^= xor_mask; - memcpy(&out.qs[dst_offset], &elems, sizeof(uint64_t)); - } - } else if (blck_size_interleave == 4) { - const uint32_t xor_mask = 0x88888888; - for (int i = 0; i < end; ++i) { - int src_id = i % 4; - int src_offset = (i / 4) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - uint32_t elems; - memcpy(&elems, &in[src_id].qs[src_offset], sizeof(uint32_t)); - elems ^= xor_mask; - memcpy(&out.qs[dst_offset], &elems, sizeof(uint32_t)); - } - } else { - GGML_ASSERT(false); - } - - return out; -} - -// interleave 8 block_q4_0s in blocks of blck_size_interleave -// returns an interleaved block_q4_0x8 -// in the interleaved block_q4_0x8, place deltas for 8 block_q4_0 blocks -// first, then interleave quants from 8 block_q4_0s in blocks of blck_size_interleave -static block_q4_0x8 make_block_q4_0x8(block_q4_0 * in, unsigned int blck_size_interleave) { - block_q4_0x8 out; - - for (int i = 0; i < 8; i++) { - out.d[i] = in[i].d; - } - - const int end = QK4_0 * 4 / blck_size_interleave; - const uint64_t xor_mask = 0x8888888888888888ULL; - - for (int i = 0; i < end; ++i) { - int src_id = i % 8; - int src_offset = (i / 8) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - uint64_t elems; - memcpy(&elems, &in[src_id].qs[src_offset], sizeof(uint64_t)); - elems ^= xor_mask; - memcpy(&out.qs[dst_offset], &elems, sizeof(uint64_t)); - } - - return out; -} - -static block_q4_0x16 make_block_q4_0x16(block_q4_0 * in, unsigned int blck_size_interleave) { - block_q4_0x16 out; - - for (int i = 0; i < 16; i++) { - out.d[i] = in[i].d; - } - - const int end = QK4_0 * 8 / blck_size_interleave; - - if (blck_size_interleave == 1) { - const uint8_t xor_mask = 0x88; - for (int i = 0; i < end; ++i) { - int src_id = i % 16; - int src_offset = i / 16; - int dst_offset = i; - - out.qs[dst_offset] = in[src_id].qs[src_offset] ^ xor_mask; - } - } else { - GGML_ASSERT(false); - } - - return out; -} - -static block_q4_Kx8 make_block_q4_Kx8(block_q4_K * in, unsigned int blck_size_interleave) { - block_q4_Kx8 out; - //Delta(scale) and dmin values of the eight Q4_K structures are copied onto the output interleaved structure - for (int i = 0; i < 8; i++) { - out.d[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d; - } - - for (int i = 0; i < 8; i++) { - out.dmin[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.dmin; - } - - const int end = QK_K * 4 / blck_size_interleave; - - // Interleave Q4_K quants by taking 8 bytes at a time - for (int i = 0; i < end; ++i) { - int src_id = i % 8; - int src_offset = (i / 8) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - // buffer large enough for the max interleave block size (8 bytes) - uint64_t elems; - memcpy(&elems, &in[src_id].qs[src_offset], blck_size_interleave); - memcpy(&out.qs[dst_offset], &elems, blck_size_interleave); - } - - // The below logic is designed so as to unpack and rearrange scales and mins values in Q4_K - // Currently the Q4_K structure has 8 scales and 8 mins packed in 12 bytes ( 6 bits for each value) - // The output Q4_Kx8 structure has 96 bytes - // Every 12 byte is packed such that it contains scales and mins for corresponding sub blocks from Q4_K structure - // For eg - First 12 bytes contains 8 scales and 8 mins - each of first sub block from different Q4_K structures - uint8_t s[8], m[8]; - - for (int i = 0; i < 4; i++) { - for (int j = 0; j < 8; j++) { - s[j] = in[j].scales[i] & 63; - m[j] = in[j].scales[i + 4] & 63; - } - - out.scales[i * 12] = (s[0] & 63) + ((s[4] & 48) << 2); - out.scales[i * 12 + 1] = (s[1] & 63) + ((s[5] & 48) << 2); - out.scales[i * 12 + 2] = (s[2] & 63) + ((s[6] & 48) << 2); - out.scales[i * 12 + 3] = (s[3] & 63) + ((s[7] & 48) << 2); - out.scales[i * 12 + 4] = (m[0] & 63) + ((m[4] & 48) << 2); - out.scales[i * 12 + 5] = (m[1] & 63) + ((m[5] & 48) << 2); - out.scales[i * 12 + 6] = (m[2] & 63) + ((m[6] & 48) << 2); - out.scales[i * 12 + 7] = (m[3] & 63) + ((m[7] & 48) << 2); - out.scales[i * 12 + 8] = (s[4] & 15) + ((m[4] & 15) << 4); - out.scales[i * 12 + 9] = (s[5] & 15) + ((m[5] & 15) << 4); - out.scales[i * 12 + 10] = (s[6] & 15) + ((m[6] & 15) << 4); - out.scales[i * 12 + 11] = (s[7] & 15) + ((m[7] & 15) << 4); - - } - - for (int i = 0; i < 4; i++) { - for (int j = 0; j < 8; j++) { - s[j] = ((in[j].scales[i] & 192) >> 2) | (in[j].scales[i+8] & 15); - m[j] = ((in[j].scales[i + 4] & 192) >> 2) | ((in[j].scales[i+8] & 240) >> 4); - } - - out.scales[i * 12 + 48] = (s[0] & 63) + ((s[4] & 48) << 2); - out.scales[i * 12 + 49] = (s[1] & 63) + ((s[5] & 48) << 2); - out.scales[i * 12 + 50] = (s[2] & 63) + ((s[6] & 48) << 2); - out.scales[i * 12 + 51] = (s[3] & 63) + ((s[7] & 48) << 2); - out.scales[i * 12 + 52] = (m[0] & 63) + ((m[4] & 48) << 2); - out.scales[i * 12 + 53] = (m[1] & 63) + ((m[5] & 48) << 2); - out.scales[i * 12 + 54] = (m[2] & 63) + ((m[6] & 48) << 2); - out.scales[i * 12 + 55] = (m[3] & 63) + ((m[7] & 48) << 2); - out.scales[i * 12 + 56] = (s[4] & 15) + ((m[4] & 15) << 4); - out.scales[i * 12 + 57] = (s[5] & 15) + ((m[5] & 15) << 4); - out.scales[i * 12 + 58] = (s[6] & 15) + ((m[6] & 15) << 4); - out.scales[i * 12 + 59] = (s[7] & 15) + ((m[7] & 15) << 4); - - } - - return out; -} - -static block_q4_Kx16 make_block_q4_Kx16(block_q4_K * in, unsigned int blck_size_interleave) { - block_q4_Kx16 out; - //Delta(scale) and dmin values of the 16 Q4_K structures are copied onto the output interleaved structure - for (int i = 0; i < 16; i++) { - out.d[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d; - } - - for (int i = 0; i < 16; i++) { - out.dmin[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.dmin; - } - - const int end = QK_K * 8 / blck_size_interleave; - - if (blck_size_interleave == 1) { - for (int i = 0; i < end; ++i) { - int src_id = i % 16; - int src_offset = i / 16; - int dst_offset = i; - - out.qs[dst_offset] = in[src_id].qs[src_offset]; - } - - // RVV repacking. - // - // Extract sums and mins for all 8 sub-blocks for each block of Q4_K. - uint8_t s[128], m[128]; - for (int i = 0; i < 4; i++) { - for (int j = 0; j < 16; j++) { - s[i * 16 + j] = in[j].scales[i] & 63; - m[i * 16 + j] = in[j].scales[i + 4] & 63; - } - } - for (int i = 0; i < 4; i++) { - for (int j = 0; j < 16; j++) { - s[64 + i * 16 + j] = ((in[j].scales[i] & 192) >> 2) | (in[j].scales[i+8] & 15); - m[64 + i * 16 + j] = ((in[j].scales[i + 4] & 192) >> 2) | ((in[j].scales[i+8] & 240) >> 4); - } - } - - for (int i = 0; i < 128; i++) { - out.scales[i] = (s[i] & 15) | ((m[i] & 15) << 4); - } - for (int i = 0; i < 64; i++) { - out.scales[128 + i] = ((s[i] & 48) >> 4) | ((m[i] & 48) >> 2) | (s[64 + i] & 48) | ((m[64 + i] & 48) << 2); - } - } else { - GGML_ASSERT(false); - } - - return out; -} - -static block_q2_Kx8 make_block_q2_Kx8(block_q2_K * in, unsigned int blck_size_interleave) { - block_q2_Kx8 out; - - // Delta(scale) and dmin values of the eight Q2_K structures are copied onto the output interleaved structure - for (int i = 0; i < 8; i++) { - out.d[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d; - } - - for (int i = 0; i < 8; i++) { - out.dmin[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.dmin; - } - - const int end = QK_K * 2 / blck_size_interleave; - - // Interleave Q2_K quants by taking 8 bytes at a time - for (int i = 0; i < end; ++i) { - int src_id = i % 8; - int src_offset = (i / 8) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - uint64_t elems; - memcpy(&elems, &in[src_id].qs[src_offset], sizeof(uint64_t)); - memcpy(&out.qs[dst_offset], &elems, sizeof(uint64_t)); - } - - // The below logic is designed so as to unpack and rearrange scales and mins values in Q2_K - // Currently the Q2_K structure has 16 scales and 16 mins packed in 16 bytes ( 4 bits for each value) - // The output Q2_Kx8 structure has 128 bytes for storing scales and mins - // Every 16 byte is packed such that it contains scales and mins for corresponding sub blocks from Q2_K structure - // For eg - First 16 bytes contains 16 scales and 16 mins - each of first and second sub blocks from different Q2_K structures - - for (int i = 0; i < 128; i++) { - // Index for selecting which q2k super block - int src1 = (i % 16) / 2; - // Index for selecting scale - int src2 = ((i / 16) * 2) + (i % 2); - - out.scales[i] = in[src1].scales[src2]; - } - return out; -} - -static block_q5_Kx8 make_block_q5_Kx8(block_q5_K * in, unsigned int blck_size_interleave) { - block_q5_Kx8 out; - //Delta(scale) and dmin values of the eight Q5_K structures are copied onto the output interleaved structure - for (int i = 0; i < 8; i++) { - out.d[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d; - } - - for (int i = 0; i < 8; i++) { - out.dmin[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.dmin; - } - - const int end = QK_K * 4 / blck_size_interleave; - - // Interleave Q5_K quants by taking blck_size_interleave bytes at a time - for (int i = 0; i < end; ++i) { - int src_id = i % 8; - int src_offset = (i / 8) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], blck_size_interleave); - } - - // Repeat for high bits with the same chunk size, since - // the high bits are interleaved in Q5_K and the index is - // qh_idx = (qs_idx % 32); - // qh_val = qh[qh_idx] >> (qs_idx / 32); - for (int i = 0; i < end / 4; ++i) { - int src_id = i % 8; - int src_offset = (i / 8) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - memcpy(&out.qh[dst_offset], &in[src_id].qh[src_offset], blck_size_interleave); - } - - // The below logic is copied over from Q4_K - // The point is to unpack all the scales and mins for each sub block every time we load 12 bytes. - // Currently the Q5_K structure has 8 scales and 8 mins packed in 12 bytes ( 6 bits for each value) - // The output Q5_Kx8 structure has 96 bytes - // Every 12 byte is packed such that it contains scales and mins for corresponding sub blocks from Q5_K structure - // For eg - First 12 bytes contains 8 scales and 8 mins - each of first sub block from different Q5_K structures - uint8_t s[8], m[8]; - - for (int i = 0; i < 4; i++) { - for (int j = 0; j < 8; j++) { - s[j] = in[j].scales[i] & 63; - m[j] = in[j].scales[i + 4] & 63; - } - - out.scales[i * 12] = (s[0] & 63) + ((s[4] & 48) << 2); - out.scales[i * 12 + 1] = (s[1] & 63) + ((s[5] & 48) << 2); - out.scales[i * 12 + 2] = (s[2] & 63) + ((s[6] & 48) << 2); - out.scales[i * 12 + 3] = (s[3] & 63) + ((s[7] & 48) << 2); - out.scales[i * 12 + 4] = (m[0] & 63) + ((m[4] & 48) << 2); - out.scales[i * 12 + 5] = (m[1] & 63) + ((m[5] & 48) << 2); - out.scales[i * 12 + 6] = (m[2] & 63) + ((m[6] & 48) << 2); - out.scales[i * 12 + 7] = (m[3] & 63) + ((m[7] & 48) << 2); - out.scales[i * 12 + 8] = (s[4] & 15) + ((m[4] & 15) << 4); - out.scales[i * 12 + 9] = (s[5] & 15) + ((m[5] & 15) << 4); - out.scales[i * 12 + 10] = (s[6] & 15) + ((m[6] & 15) << 4); - out.scales[i * 12 + 11] = (s[7] & 15) + ((m[7] & 15) << 4); - } - - for (int i = 0; i < 4; i++) { - for (int j = 0; j < 8; j++) { - s[j] = ((in[j].scales[i] & 192) >> 2) | (in[j].scales[i + 8] & 15); - m[j] = ((in[j].scales[i + 4] & 192) >> 2) | ((in[j].scales[i + 8] & 240) >> 4); - } - - out.scales[i * 12 + 48] = (s[0] & 63) + ((s[4] & 48) << 2); - out.scales[i * 12 + 49] = (s[1] & 63) + ((s[5] & 48) << 2); - out.scales[i * 12 + 50] = (s[2] & 63) + ((s[6] & 48) << 2); - out.scales[i * 12 + 51] = (s[3] & 63) + ((s[7] & 48) << 2); - out.scales[i * 12 + 52] = (m[0] & 63) + ((m[4] & 48) << 2); - out.scales[i * 12 + 53] = (m[1] & 63) + ((m[5] & 48) << 2); - out.scales[i * 12 + 54] = (m[2] & 63) + ((m[6] & 48) << 2); - out.scales[i * 12 + 55] = (m[3] & 63) + ((m[7] & 48) << 2); - out.scales[i * 12 + 56] = (s[4] & 15) + ((m[4] & 15) << 4); - out.scales[i * 12 + 57] = (s[5] & 15) + ((m[5] & 15) << 4); - out.scales[i * 12 + 58] = (s[6] & 15) + ((m[6] & 15) << 4); - out.scales[i * 12 + 59] = (s[7] & 15) + ((m[7] & 15) << 4); - } - - return out; -} - -static block_q6_Kx8 make_block_q6_Kx8(block_q6_K * in, unsigned int blck_size_interleave) { - block_q6_Kx8 out; - constexpr int n_blocks = 8; // Kx8 - for (int i = 0; i < n_blocks; i++) { - out.d[i] = in[i].d; - } - - const int end_ls = QK_K * 4 / blck_size_interleave; - // Interleave Q6_K quants by taking blck_size_interleave bytes at a time - for (int i = 0; i < end_ls; ++i) { - int src_id = i % n_blocks; - int src_offset = (i / n_blocks) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - uint64_t elem_ls; - memcpy(&elem_ls, &in[src_id].ql[src_offset], blck_size_interleave); - memcpy(&out.ql[dst_offset], &elem_ls, blck_size_interleave); - } - - // Interleave high bits using same chunk size as low bits - const int end_hs = end_ls / 2; - for (int i = 0; i < end_hs; ++i) { - int src_id = i % n_blocks; - int src_offset = (i / n_blocks) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - uint64_t elem_hs; - memcpy(&elem_hs, &in[src_id].qh[src_offset], blck_size_interleave); - memcpy(&out.qh[dst_offset], &elem_hs, blck_size_interleave); - } - - // The below logic is designed so as to unpack and rearrange scales in Q6_K - // The output Q6_Kx8 structure interleaves the 8 bit scales in the same fashion as the quants - // Q6_K structure has an 8-bit scale per 16 elements -> 16 scales - // scales: [0 bl0 0 bl1 ... 0 bl7][1 bl0 ... 1 bl7] ... [15 bl0 ... 15 bl7] (bl = block) - constexpr int n_scales = QK_K / 16; - - for (int i = 0; i < n_blocks; i++) { - for (int j = 0; j < n_scales; j++) { - out.scales[j * n_blocks + i] = in[i].scales[j]; - } - } - - return out; -} - -static block_q2_Kx16 make_block_q2_Kx16(const block_q2_K * in, unsigned int blck_size_interleave) { - block_q2_Kx16 out; - constexpr int N_COLS = 16; - - // 1. Copy Super-Scales (d) and Super-Mins (dmin) - for (int i = 0; i < N_COLS; i++) { - out.d[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d; - out.dmin[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.dmin; - } - - // 2. Interleave Q2_K Data - const int bytes_per_col = 64; - const int total_bytes = N_COLS * bytes_per_col; - const int end = total_bytes / blck_size_interleave; - - for (int i = 0; i < end; ++i) { - int src_col_id = i % N_COLS; - int src_offset = (i / N_COLS) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - memcpy(&out.qs[dst_offset], &in[src_col_id].qs[src_offset], blck_size_interleave); - } - - // 3. Repack Scales into the Optimized "Sequential-Parallel" Layout - int out_idx = 0; - - // Arrays define the sub-block order for each group - const int even_low_sbs[] = {0, 2, 4, 6}; - const int odd_low_sbs[] = {1, 3, 5, 7}; - const int even_high_sbs[] = {8, 10, 12, 14}; - const int odd_high_sbs[] = {9, 11, 13, 15}; - - // Pack Group 1: Even-Low - for (int sb : even_low_sbs) { - for (int col = 0; col < N_COLS; col++) { - out.scales[out_idx++] = in[col].scales[sb]; - } - } - - // Pack Group 2: Odd-Low - for (int sb : odd_low_sbs) { - for (int col = 0; col < N_COLS; col++) { - out.scales[out_idx++] = in[col].scales[sb]; - } - } - - // Pack Group 3: Even-High - for (int sb : even_high_sbs) { - for (int col = 0; col < N_COLS; col++) { - out.scales[out_idx++] = in[col].scales[sb]; - } - } - - // Pack Group 4: Odd-High - for (int sb : odd_high_sbs) { - for (int col = 0; col < N_COLS; col++) { - out.scales[out_idx++] = in[col].scales[sb]; - } - } - - return out; -} - -static int repack_q4_0_to_q4_0_4_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_Q4_0); - GGML_ASSERT(interleave_block == 4 || interleave_block == 8); - constexpr int nrows_interleaved = 4; - - block_q4_0x4 * dst = (block_q4_0x4 *)t->data; - const block_q4_0 * src = (const block_q4_0 *)data; - block_q4_0 dst_tmp[4]; - int nrow = ggml_nrows(t); - int nblocks = t->ne[0] / QK4_0; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_0)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_q4_0x4(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; - - GGML_UNUSED(data_size); -} - -static int repack_q4_K_to_q4_K_8_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_Q4_K); - GGML_ASSERT(interleave_block == 8 || interleave_block == 4); - constexpr int nrows_interleaved = 8; - - block_q4_Kx8 * dst = (block_q4_Kx8*)t->data; - const block_q4_K * src = (const block_q4_K*) data; - block_q4_K dst_tmp[8]; - int nrow = ggml_nrows(t); - int nblocks = t->ne[0] / QK_K; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_K)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++ ) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_q4_Kx8(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; - - GGML_UNUSED(data_size); -} - -static int repack_q4_K_to_q4_K_16_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_Q4_K); - constexpr int nrows_interleaved = 16; - - block_q4_Kx16 * dst = (block_q4_Kx16*)t->data; - const block_q4_K * src = (const block_q4_K*) data; - block_q4_K dst_tmp[16]; - int nrow = ggml_nrows(t); - int nblocks = t->ne[0] / QK_K; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_K)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++ ) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_q4_Kx16(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; - - GGML_UNUSED(data_size); -} - -static int repack_q2_K_to_q2_K_8_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_Q2_K); - GGML_ASSERT(interleave_block == 8); - constexpr int nrows_interleaved = 8; - - block_q2_Kx8 * dst = (block_q2_Kx8*)t->data; - const block_q2_K * src = (const block_q2_K*) data; - block_q2_K dst_tmp[8]; - int nrow = ggml_nrows(t); - int nblocks = t->ne[0] / QK_K; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q2_K)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_q2_Kx8(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; - - GGML_UNUSED(data_size); -} - -static int repack_q2_K_to_q2_K_16_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_Q2_K); - constexpr int nrows_interleaved = 16; - - block_q2_Kx16 * dst = (block_q2_Kx16*)t->data; - const block_q2_K * src = (const block_q2_K*) data; - - block_q2_K dst_tmp[nrows_interleaved]; - - int nrow = ggml_nrows(t); - int nblocks = t->ne[0] / QK_K; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q2_K)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - // This loop gathers 16 separate blocks (one from each column) - // that correspond to the same K-dimension chunk. - for (int i = 0; i < nrows_interleaved; i++ ) { - dst_tmp[i] = src[x + i * nblocks]; - } - - *dst++ = make_block_q2_Kx16(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; - - GGML_UNUSED(data_size); -} - -static int repack_q4_0_to_q4_0_16_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_Q4_0); - constexpr int nrows_interleaved = 16; - - block_q4_0x16 * dst = (block_q4_0x16*)t->data; - const block_q4_0 * src = (const block_q4_0*) data; - block_q4_0 dst_tmp[16]; - int nrow = ggml_nrows(t); - int nblocks = t->ne[0] / QK4_0; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_0)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++ ) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_q4_0x16(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; - - GGML_UNUSED(data_size); -} - -static int repack_q5_K_to_q5_K_8_bl(struct ggml_tensor * t, - int interleave_block, - const void * GGML_RESTRICT data, - size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_Q5_K); - GGML_ASSERT(interleave_block == 4 || interleave_block == 8); - constexpr int nrows_interleaved = 8; - - block_q5_Kx8 * dst = (block_q5_Kx8 *) t->data; - const block_q5_K * src = (const block_q5_K *) data; - block_q5_K dst_tmp[8]; - int nrow = ggml_nrows(t); - int nblocks = t->ne[0] / QK_K; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q5_K)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_q5_Kx8(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; -} - -static int repack_q6_K_to_q6_K_8_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_Q6_K); - GGML_ASSERT(interleave_block == 4 || interleave_block == 8); - constexpr int nrows_interleaved = 8; - - block_q6_Kx8 * dst = (block_q6_Kx8 *)t->data; - const block_q6_K * src = (const block_q6_K *) data; - block_q6_K dst_tmp[8]; - int nrow = ggml_nrows(t); - int nblocks = t->ne[0] / QK_K; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q6_K)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_q6_Kx8(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; -} - -static int repack_q4_0_to_q4_0_8_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_Q4_0); - GGML_ASSERT(interleave_block == 8); - constexpr int nrows_interleaved = 8; - - block_q4_0x8 * dst = (block_q4_0x8*)t->data; - const block_q4_0 * src = (const block_q4_0*) data; - block_q4_0 dst_tmp[8]; - int nrow = ggml_nrows(t); - int nblocks = t->ne[0] / QK4_0; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_0)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++ ) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_q4_0x8(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; - - GGML_UNUSED(data_size); -} - -static int repack_q8_0_to_q8_0_4_bl(struct ggml_tensor * t, - int interleave_block, - const void * GGML_RESTRICT data, - size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_Q8_0); - GGML_ASSERT(interleave_block == 4 || interleave_block == 8); - constexpr int nrows_interleaved = 4; - - block_q8_0x4 * dst = (block_q8_0x4 *) t->data; - const block_q8_0 * src = (const block_q8_0 *) data; - block_q8_0 dst_tmp[4]; - int nrow = ggml_nrows(t); - int nblocks = t->ne[0] / QK8_0; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q8_0)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_q8_0x4(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; -} - -static block_q8_0x16 make_block_q8_0x16(block_q8_0 * in, unsigned int blck_size_interleave) { - block_q8_0x16 out; - - for (int i = 0; i < 16; i++) { - out.d[i] = in[i].d; - } - - const int end = QK8_0 * 16 / blck_size_interleave; - - if (blck_size_interleave == 1) { - for (int i = 0; i < end; ++i) { - int src_id = i % 16; - int src_offset = i / 16; - int dst_offset = i; - out.qs[dst_offset] = in[src_id].qs[src_offset]; - } - } else { - GGML_ASSERT(false); - } - - return out; -} - -static int repack_q8_0_to_q8_0_16_bl(struct ggml_tensor * t, - int interleave_block, - const void * GGML_RESTRICT data, - size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_Q8_0); - constexpr int nrows_interleaved = 16; - - block_q8_0x16 * dst = (block_q8_0x16 *) t->data; - const block_q8_0 * src = (const block_q8_0 *) data; - block_q8_0 dst_tmp[16]; - int nrow = ggml_nrows(t); - int nblocks = t->ne[0] / QK8_0; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q8_0)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_q8_0x16(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; -} - -static block_iq4_nlx4 make_block_iq4_nlx4(block_iq4_nl * in, unsigned int blck_size_interleave) { - block_iq4_nlx4 out; - - for (int i = 0; i < 4; i++) { - out.d[i] = in[i].d; - } - - const int end = QK4_NL * 2 / blck_size_interleave; - - // TODO: this branch seems wrong - //if (blck_size_interleave == 8) { - // for (int i = 0; i < end; ++i) { - // int src_id = i % 4; - // int src_offset = (i / 4) * blck_size_interleave; - // int dst_offset = i * blck_size_interleave; - - // // Using memcpy to avoid unaligned memory accesses - // memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], sizeof(uint64_t)); - // } - //} else - if (blck_size_interleave == 4) { - for (int i = 0; i < end; ++i) { - int src_id = i % 4; - int src_offset = (i / 4) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], sizeof(uint32_t)); - } - } else { - GGML_ASSERT(false); - } - - return out; -} - -static int repack_iq4_nl_to_iq4_nl_4_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_IQ4_NL); - GGML_ASSERT(interleave_block == 4); - - const block_iq4_nl * src = (const block_iq4_nl *)data; - block_iq4_nlx4 * dst = ( block_iq4_nlx4 *)t->data; - - block_iq4_nl dst_tmp[4]; - - int nrow = ggml_nrows(t); - int nrows_interleaved = 4; - int nblocks = t->ne[0] / QK4_NL; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_iq4_nl)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_iq4_nlx4(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; - - GGML_UNUSED(data_size); -} - -static block_iq4_nlx8 make_block_iq4_nlx8(block_iq4_nl * in, unsigned int blck_size_interleave) { - block_iq4_nlx8 out; - - for (int i = 0; i < 8; i++) { - out.d[i] = in[i].d; - } - - const int end = QK4_NL * 4 / blck_size_interleave; - - if (blck_size_interleave == 8) { - for (int i = 0; i < end; ++i) { - int src_id = i % 8; - int src_offset = (i / 8) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], sizeof(uint64_t)); - } - } else { - GGML_ASSERT(false); - } - - return out; -} - -static int repack_iq4_nl_to_iq4_nl_8_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_IQ4_NL); - GGML_ASSERT(interleave_block == 8); - - const block_iq4_nl * src = (const block_iq4_nl *)data; - block_iq4_nlx8 * dst = ( block_iq4_nlx8 *)t->data; - - block_iq4_nl dst_tmp[8]; - - int nrow = ggml_nrows(t); - int nrows_interleaved = 8; - int nblocks = t->ne[0] / QK4_NL; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_iq4_nl)); - - if (t->ne[1] % nrows_interleaved != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_iq4_nlx8(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; - - GGML_UNUSED(data_size); -} - -static block_iq4_nlx16 make_block_iq4_nlx16(block_iq4_nl * in, unsigned int blck_size_interleave) { - block_iq4_nlx16 out; - - for (int i = 0; i < 16; i++) { - out.d[i] = in[i].d; - } - - const int end = QK4_NL * 8 / blck_size_interleave; - - if (blck_size_interleave == 1) { - for (int i = 0; i < end; ++i) { - int src_id = i % 16; - int src_offset = i / 16; - int dst_offset = i; - - out.qs[dst_offset] = in[src_id].qs[src_offset]; - } - } else { - GGML_ASSERT(false); - } - - return out; -} - -static int repack_iq4_nl_to_iq4_nl_16_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_IQ4_NL); - GGML_ASSERT(interleave_block == 1); - - const block_iq4_nl * src = (const block_iq4_nl *)data; - block_iq4_nlx16 * dst = ( block_iq4_nlx16 *)t->data; - - block_iq4_nl dst_tmp[16]; - - int nrow = ggml_nrows(t); - int nrows_interleaved = 16; - int nblocks = t->ne[0] / QK4_NL; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_iq4_nl)); - - if (t->ne[1] % nrows_interleaved != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_iq4_nlx16(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; - - GGML_UNUSED(data_size); -} - -static block_mxfp4x4 make_block_mxfp4x4(block_mxfp4 * in, unsigned int blck_size_interleave) { - block_mxfp4x4 out; - - for (int i = 0; i < 4; i++) { - out.e[i] = in[i].e; - } - - const int end = QK_MXFP4 * 2 / blck_size_interleave; - - if (blck_size_interleave == 4) { - for (int i = 0; i < end; ++i) { - int src_id = i % 4; - int src_offset = (i / 4) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], sizeof(uint32_t)); - } - } else { - GGML_ASSERT(false); - } - - return out; -} - -static int repack_mxfp4_to_mxfp4_4_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_MXFP4); - GGML_ASSERT(interleave_block == 4); - - const block_mxfp4 * src = (const block_mxfp4 *)data; - block_mxfp4x4 * dst = ( block_mxfp4x4 *)t->data; - - block_mxfp4 dst_tmp[4]; - - int nrow = ggml_nrows(t); - int nrows_interleaved = 4; - int nblocks = t->ne[0] / QK_MXFP4; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_mxfp4)); - - if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_mxfp4x4(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; - - GGML_UNUSED(data_size); -} - -static block_mxfp4x8 make_block_mxfp4x8(block_mxfp4 * in, unsigned int blck_size_interleave) { - block_mxfp4x8 out; - - for (int i = 0; i < 8; i++) { - out.e[i] = in[i].e; - } - - const int end = QK_MXFP4 * 4 / blck_size_interleave; - - if (blck_size_interleave == 8) { - for (int i = 0; i < end; ++i) { - int src_id = i % 8; - int src_offset = (i / 8) * blck_size_interleave; - int dst_offset = i * blck_size_interleave; - - memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], sizeof(uint64_t)); - } - } else { - GGML_ASSERT(false); - } - - return out; -} - -static int repack_mxfp4_to_mxfp4_8_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { - GGML_ASSERT(t->type == GGML_TYPE_MXFP4); - GGML_ASSERT(interleave_block == 8); - - const block_mxfp4 * src = (const block_mxfp4 *)data; - block_mxfp4x8 * dst = ( block_mxfp4x8 *)t->data; - - block_mxfp4 dst_tmp[8]; - - int nrow = ggml_nrows(t); - int nrows_interleaved = 8; - int nblocks = t->ne[0] / QK_MXFP4; - - GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_mxfp4)); - - if (t->ne[1] % nrows_interleaved != 0) { - return -1; - } - - for (int b = 0; b < nrow; b += nrows_interleaved) { - for (int64_t x = 0; x < nblocks; x++) { - for (int i = 0; i < nrows_interleaved; i++) { - dst_tmp[i] = src[x + i * nblocks]; - } - *dst++ = make_block_mxfp4x8(dst_tmp, interleave_block); - } - src += nrows_interleaved * nblocks; - } - return 0; - - GGML_UNUSED(data_size); -} - -namespace ggml::cpu::repack { -// repack -template -int repack(struct ggml_tensor *, const void *, size_t); - -// TODO: generalise. -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q4_0_to_q4_0_4_bl(t, 4, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q4_0_to_q4_0_4_bl(t, 8, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q4_0_to_q4_0_8_bl(t, 8, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q4_K_to_q4_K_8_bl(t, 8, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q4_K_to_q4_K_8_bl(t, 4, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q2_K_to_q2_K_8_bl(t, 8, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q5_K_to_q5_K_8_bl(t, 4, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q5_K_to_q5_K_8_bl(t, 8, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q6_K_to_q6_K_8_bl(t, 4, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q6_K_to_q6_K_8_bl(t, 8, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_iq4_nl_to_iq4_nl_4_bl(t, 4, data, data_size); -} - -// TODO: needs to be revisited -//template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { -// return repack_iq4_nl_to_iq4_nl_4_bl(t, 8, data, data_size); -//} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_iq4_nl_to_iq4_nl_8_bl(t, 8, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_mxfp4_to_mxfp4_4_bl(t, 4, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_mxfp4_to_mxfp4_8_bl(t, 8, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q8_0_to_q8_0_4_bl(t, 4, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q8_0_to_q8_0_4_bl(t, 8, data, data_size); -} - -#if defined __riscv_zvfh -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q4_0_to_q4_0_16_bl(t, 1, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q4_K_to_q4_K_16_bl(t, 1, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_iq4_nl_to_iq4_nl_16_bl(t, 1, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q8_0_to_q8_0_16_bl(t, 1, data, data_size); -} - -template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { - return repack_q2_K_to_q2_K_16_bl(t, 1, data, data_size); -} -#endif - -// gemv -template -void gemv(int, float *, size_t, const void *, const void *, int, int); - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q4_0_4x4_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q4_0_4x8_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q4_0_8x8_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> -void gemv(int n, - float * s, - size_t bs, - const void * vx, - const void * vy, - int nr, - int nc) { - ggml_gemv_q2_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q4_K_8x4_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q4_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q5_K_8x4_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q5_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q6_K_8x4_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q6_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_iq4_nl_4x4_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_iq4_nl_8x8_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_mxfp4_4x4_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_mxfp4_8x8_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q8_0_4x4_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q8_0_4x8_q8_0(n, s, bs, vx, vy, nr, nc); -} - -#if defined __riscv_zvfh -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q4_0_16x1_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q4_K_16x1_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_iq4_nl_16x1_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q8_0_16x1_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemv_q2_K_16x1_q8_K(n, s, bs, vx, vy, nr, nc); -} -#endif - -// gemm -template -void gemm(int, float *, size_t, const void *, const void *, int, int); - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q4_0_4x4_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q4_0_4x8_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> -void gemm(int n, - float * s, - size_t bs, - const void * vx, - const void * vy, - int nr, - int nc) { - ggml_gemm_q4_0_8x8_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q2_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q4_K_8x4_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q4_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q5_K_8x4_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q5_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q6_K_8x4_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q6_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_iq4_nl_4x4_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_iq4_nl_8x8_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_mxfp4_4x4_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_mxfp4_8x8_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q8_0_4x4_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q8_0_4x8_q8_0(n, s, bs, vx, vy, nr, nc); -} - -#if defined __riscv_zvfh -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q4_0_16x1_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q4_K_16x1_q8_K(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_iq4_nl_16x1_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q8_0_16x1_q8_0(n, s, bs, vx, vy, nr, nc); -} - -template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { - ggml_gemm_q2_K_16x1_q8_K(n, s, bs, vx, vy, nr, nc); -} -#endif - -class tensor_traits_base : public ggml::cpu::tensor_traits { - public: - virtual int repack(struct ggml_tensor * t, const void * data, size_t data_size) = 0; -}; - -template class tensor_traits : public tensor_traits_base { - - bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override { - // not realy a GGML_TYPE_Q8_0 but same size. - switch (op->op) { - case GGML_OP_MUL_MAT: - { - size = ggml_row_size(PARAM_TYPE, ggml_nelements(op->src[1])); - return true; - } - case GGML_OP_MUL_MAT_ID: - { - size = ggml_row_size(PARAM_TYPE, ggml_nelements(op->src[1])); - size = GGML_PAD(size, sizeof(int64_t)); // + padding for next block. - - const int64_t ne02 = op->src[0]->ne[2]; // n_as, n_expert - const int64_t ne12 = op->src[1]->ne[2]; // n_tokens - - const size_t sizeof_mmid_row_mapping = sizeof(int64_t); - - size += sizeof_mmid_row_mapping*ne02*(ne12 + 1); - - return true; - } - default: - // GGML_ABORT("fatal error"); - break; - } - return false; - } - - bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) override { - switch (op->op) { - case GGML_OP_MUL_MAT: - forward_mul_mat(params, op); - return true; - case GGML_OP_MUL_MAT_ID: - forward_mul_mat_id(params, op); - return true; - default: - // GGML_ABORT("fatal error"); - break; - } - return false; - } - - void forward_mul_mat_one_chunk(ggml_compute_params * params, - ggml_tensor * op, - int64_t src0_start, - int64_t src0_end, - int64_t src1_start, - int64_t src1_end) { - const ggml_tensor * src0 = op->src[0]; - const ggml_tensor * src1 = op->src[1]; - ggml_tensor * dst = op; - - GGML_TENSOR_BINARY_OP_LOCALS - - const size_t src1_col_stride = ggml_row_size(PARAM_TYPE, ne10); - - GGML_ASSERT(ne03 == 1 && ne13 == 1); - GGML_ASSERT(ne12 % ne02 == 0); - const int64_t r2 = ne12 / ne02; - - const int64_t i12 = src1_start / ne1; - const int64_t i11 = src1_start - i12 * ne1; - - // Determine batch index - const int64_t i02 = i12 / r2; - - const int64_t i1 = i11; - const int64_t i2 = i12; - - const char * src0_ptr = (const char *) src0->data + i02 * nb02; - const char * src1_ptr = (const char *) params->wdata + (i11 + i12 * ne11) * src1_col_stride; - char * dst_ptr = ((char *) dst->data + (i1 * nb1 + i2 * nb2)); - - const int64_t nrows = src1_end - src1_start; - const int64_t ncols = src0_end - src0_start; - - GGML_ASSERT(src1_ptr + src1_col_stride * nrows <= (const char *) params->wdata + params->wsize); - - // If there are more than three rows in src1, use gemm; otherwise, use gemv. - if (nrows > 3) { - gemm(ne00, (float *) (dst_ptr) + src0_start, nb1 / nb0, - src0_ptr + src0_start * nb01, src1_ptr, - nrows - (nrows % 4), ncols); - } - for (int iter = nrows - (nrows % 4); iter < nrows; iter++) { - gemv(ne00, (float *) (dst_ptr + (iter * nb1)) + src0_start, - ne01, src0_ptr + src0_start * nb01, - src1_ptr + (src1_col_stride * iter), 1 /* nrows */, ncols); - } - } - - void forward_mul_mat(ggml_compute_params * params, ggml_tensor * op) { - const ggml_tensor * src0 = op->src[0]; - const ggml_tensor * src1 = op->src[1]; - ggml_tensor * dst = op; - - GGML_TENSOR_BINARY_OP_LOCALS - - const int ith = params->ith; - const int nth = params->nth; - - GGML_ASSERT(ne0 == ne01); - GGML_ASSERT(ne1 == ne11); - GGML_ASSERT(ne2 == ne12); - GGML_ASSERT(ne3 == ne13); - - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); - - // TODO: General batched mul mat for 4D tensors - // Currently only supports 3D tensors - GGML_ASSERT(ne03 == 1); - GGML_ASSERT(ne13 == 1); - GGML_ASSERT(ne3 == 1); - - GGML_ASSERT(src1->type == GGML_TYPE_F32); - - GGML_ASSERT(ggml_n_dims(op->src[0]) == 2); - // GGML_ASSERT(ggml_n_dims(op->src[1]) == 2); - - char * wdata = static_cast(params->wdata); - const size_t nbw1 = ggml_row_size(PARAM_TYPE, ne10); - const size_t nbw2 = nbw1 * ne11; - - assert(params->wsize >= nbw2 * ne12); - - const ggml_from_float_t from_float = ggml_get_type_traits_cpu(PARAM_TYPE)->from_float; - - // INFO: Quantization is done in planes to avoid extra complexity in chunking. - // Flattening dimensions not multiple of INTER_SIZE would require extra handling depending on how - // the planes are broadcast. - for (int64_t i12 = 0; i12 < ne12; i12++) { - char * data_ptr = (char *) src1->data + i12 * nb12; - char * wdata_ptr = wdata + i12 * nbw2; - - for (int64_t i11 = ith * 4; i11 < ne11 - ne11 % 4; i11 += nth * 4) { - ggml_quantize_mat_t((float *) (data_ptr + i11 * nb11), - (void *) (wdata_ptr + i11 * nbw1), 4, ne10); - } - - const int64_t i11_processed = ne11 - ne11 % 4; - for (int64_t i11 = i11_processed + ith; i11 < ne11; i11 += nth) { - from_float((float *) (data_ptr + i11 * nb11), (void *) (wdata_ptr + i11 * nbw1), ne10); - } - } - - // disable for NUMA - const bool disable_chunking = ggml_is_numa(); - - // 4x chunks per thread - const int64_t nr0 = ggml_nrows(op->src[0]); - - int nth_scaled = nth * 4; - int64_t chunk_size0 = (nr0 + nth_scaled - 1) / nth_scaled; - int64_t nchunk0 = (nr0 + chunk_size0 - 1) / chunk_size0; - - // src1 is chunked only by full planes. - // When we flatten we need to address dimensions not multiple of the q8 INTER_SIZE - // to route them thorugh GEMV. - // nchunk1 = ne12 also avoids messing the chunking for models with no 3d tensors - // to avoid affecting their performance - int64_t nchunk1 = ne12; - - // Ensure minimum chunk size to avoid alignment issues with high thread counts - // Minimum chunk size should be at least NB_COLS to prevent overlapping chunks after alignment - const int64_t min_chunk_size = NB_COLS; - if (nchunk0 > 0 && (nr0 / nchunk0) < min_chunk_size && nr0 >= min_chunk_size) { - nchunk0 = (nr0 + min_chunk_size - 1) / min_chunk_size; - } - - int64_t dr0 = (nr0 + nchunk0 - 1) / nchunk0; - // Only increase nchunk0 to nth if it won't make chunks too small - if (nth == 1 || ((nchunk0 < nth || disable_chunking) && (nr0 + nth - 1) / nth >= min_chunk_size)) { - nchunk0 = nth; - dr0 = (nr0 + nchunk0 - 1) / nchunk0; - } - - // Ensure nchunk doesn't exceed the number of rows divided by minimum chunk size - // This prevents creating too many tiny chunks that could overlap after alignment - const int64_t max_nchunk = (nr0 + min_chunk_size - 1) / min_chunk_size; - nchunk0 = MIN(nchunk0, max_nchunk); - - if (ith == 0) { - // Every thread starts at ith, so the first unprocessed chunk is nth. This save a bit of coordination right at the start. - ggml_threadpool_chunk_set(params->threadpool, nth); - } - - ggml_barrier(params->threadpool); - - // The first chunk comes from our thread_id, the rest will get auto-assigned. - int current_chunk = ith; - - while (current_chunk < nchunk0 * nchunk1) { - const int64_t ith0 = current_chunk % nchunk0; - const int64_t ith1 = current_chunk / nchunk0; - - int64_t src0_start = dr0 * ith0; - int64_t src0_end = MIN(src0_start + dr0, nr0); - - // full-plane range for src1 - int64_t src1_start = ith1 * ne11; - int64_t src1_end = (ith1 + 1) * ne11; - - // Align boundaries to NB_COLS - round up to ensure all data is included - // The chunk size limiting above ensures chunks are large enough to prevent overlaps - src0_start = (src0_start % NB_COLS) ? src0_start + NB_COLS - (src0_start % NB_COLS) : src0_start; - src0_end = (src0_end % NB_COLS) ? src0_end + NB_COLS - (src0_end % NB_COLS) : src0_end; - src0_end = MIN(src0_end, ne01); - - // Make sure current plane is the last one before exiting - if (src0_start >= src0_end) { - current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1); - continue; - } - - forward_mul_mat_one_chunk(params, dst, src0_start, src0_end, src1_start, src1_end); - - current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1); - } - } - - void forward_mul_mat_id(ggml_compute_params * params, ggml_tensor * op) { - const ggml_tensor * src0 = op->src[0]; - const ggml_tensor * src1 = op->src[1]; - const ggml_tensor * ids = op->src[2]; - ggml_tensor * dst = op; - - GGML_TENSOR_BINARY_OP_LOCALS - - const int ith = params->ith; - const int nth = params->nth; - - const ggml_from_float_t from_float = ggml_get_type_traits_cpu(PARAM_TYPE)->from_float; - - // we don't support permuted src0 or src1 - GGML_ASSERT(nb00 == ggml_type_size(src0->type)); - GGML_ASSERT(nb10 == ggml_type_size(src1->type)); - - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); - - GGML_ASSERT(ne03 == 1); - GGML_ASSERT(ne13 == 1); - GGML_ASSERT(ne3 == 1); - - GGML_ASSERT(src1->type == GGML_TYPE_F32); - - // row groups - const int n_ids = ids->ne[0]; // n_expert_used - const int n_as = ne02; // n_expert - - const size_t nbw1 = ggml_row_size(PARAM_TYPE, ne10); - const size_t nbw2 = nbw1*ne11; - const size_t nbw3 = nbw2*ne12; - - struct mmid_row_mapping { - int32_t i1; - int32_t i2; - }; - - GGML_ASSERT(params->wsize >= - (GGML_PAD(nbw3, sizeof(int64_t)) + - n_as*(ne12 + 1)*sizeof(mmid_row_mapping)) - ); - - auto * wdata = (char *)params->wdata; - auto * wdata_src1_end = (char *)wdata + GGML_PAD(nbw3, sizeof(int64_t)); - - // total of [n_as][ne12 + 1] elements of type mmid_row_mapping (2*int32_t = int64_t) - auto * matrix_row_counts = (int64_t *) (wdata_src1_end); // [n_as] - struct mmid_row_mapping * matrix_rows = (struct mmid_row_mapping *) (matrix_row_counts + n_as); // [n_as][ne12] - - // src1: float32 => param type - for (int64_t i12 = 0; i12 < ne12; ++i12) { - for (int64_t i11 = ith; i11 < ne11; i11 += nth) { - from_float((float *)((char *) src1->data + i12 * nb12 + i11 * nb11), - (void *) (wdata + i12 * nbw2 + i11 * nbw1), - ne10); - } - } - -#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * ne12 + (i1)] - - if (ith == 0) { - // initialize matrix_row_counts - memset(matrix_row_counts, 0, n_as * sizeof(int64_t)); - - // group rows by src0 matrix - for (int32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { - for (int32_t id = 0; id < n_ids; ++id) { - const int32_t i02 = - *(const int32_t *) ((const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]); - - GGML_ASSERT(i02 >= 0 && i02 < n_as); - - MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = { id, iid1 }; - matrix_row_counts[i02] += 1; - } - } - } - - ggml_barrier(params->threadpool); - - // compute each matrix multiplication in sequence - for (int cur_a = 0; cur_a < n_as; ++cur_a) { - const int64_t cne1 = matrix_row_counts[cur_a]; - - if (cne1 == 0) { - continue; - } - - const auto * src0_cur = (const char *) src0->data + cur_a*nb02; - - //const int64_t nr0 = ne01; // src0 rows - const int64_t nr1 = cne1; // src1 rows - - int64_t src0_cur_start = (ith * ne01) / nth; - int64_t src0_cur_end = ((ith + 1) * ne01) / nth; - - // Align boundaries to NB_COLS - round up to ensure all data is included - src0_cur_start = (src0_cur_start % NB_COLS) ? src0_cur_start + NB_COLS - (src0_cur_start % NB_COLS) : src0_cur_start; - src0_cur_end = (src0_cur_end % NB_COLS) ? src0_cur_end + NB_COLS - (src0_cur_end % NB_COLS) : src0_cur_end; - if (src0_cur_end > ne01) { - src0_cur_end = ne01; - } - - if (src0_cur_start >= src0_cur_end) { - return; - } - - for (int ir1 = 0; ir1 < nr1; ir1++) { - struct mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, ir1); - - const int id = row_mapping.i1; // selected expert index - - const int64_t i11 = id % ne11; - const int64_t i12 = row_mapping.i2; // row index in src1 - - const int64_t i1 = id; // selected expert index - const int64_t i2 = i12; // row - - const auto * src1_col = (const char *) wdata + (i11 * nbw1 + i12 * nbw2); - - gemv( - ne00, (float *) ((char *) dst->data + (i1 * nb1 + i2 * nb2)) + src0_cur_start, ne01, - src0_cur + src0_cur_start * nb01, src1_col, 1, src0_cur_end - src0_cur_start); - } - } -#undef MMID_MATRIX_ROW - } - - int repack(struct ggml_tensor * t, const void * data, size_t data_size) override { - GGML_LOG_DEBUG("%s: repack tensor %s with %s_%dx%d\n", __func__, t->name, ggml_type_name(t->type), - (int) NB_COLS, (int) INTER_SIZE); - return ggml::cpu::repack::repack(t, data, data_size); - } -}; - -} // namespace ggml::cpu::repack - -static const ggml::cpu::tensor_traits * ggml_repack_get_optimal_repack_type(const struct ggml_tensor * cur) { - // instance for Q4 - static const ggml::cpu::repack::tensor_traits q4_0_4x4_q8_0; - static const ggml::cpu::repack::tensor_traits q4_0_4x8_q8_0; - static const ggml::cpu::repack::tensor_traits q4_0_8x8_q8_0; - - // instance for Q4_K - static const ggml::cpu::repack::tensor_traits q4_K_8x4_q8_K; - static const ggml::cpu::repack::tensor_traits q4_K_8x8_q8_K; - - // instance for Q5_K - static const ggml::cpu::repack::tensor_traits q5_K_8x4_q8_K; - static const ggml::cpu::repack::tensor_traits q5_K_8x8_q8_K; - - // instance for Q6_K - static const ggml::cpu::repack::tensor_traits q6_K_8x4_q8_K; - static const ggml::cpu::repack::tensor_traits q6_K_8x8_q8_K; - - // instance for Q2 - static const ggml::cpu::repack::tensor_traits q2_K_8x8_q8_K; - - // instance for IQ4 - static const ggml::cpu::repack::tensor_traits iq4_nl_4x4_q8_0; - static const ggml::cpu::repack::tensor_traits iq4_nl_8x8_q8_0; - - // instance for MXFP4 - static const ggml::cpu::repack::tensor_traits mxfp4_4x4_q8_0; - static const ggml::cpu::repack::tensor_traits mxfp4_8x8_q8_0; - - // instance for Q8_0 - static const ggml::cpu::repack::tensor_traits q8_0_4x4_q8_0; - static const ggml::cpu::repack::tensor_traits q8_0_4x8_q8_0; - - // instances for RISC-V - // - // These implement outer-product style matrix multiplication kernels with - // an interleave of 1. -#if defined __riscv_zvfh - static const ggml::cpu::repack::tensor_traits q4_0_16x1_q8_0; - static const ggml::cpu::repack::tensor_traits q4_K_16x1_q8_K; - static const ggml::cpu::repack::tensor_traits iq4_nl_16x1_q8_0; - static const ggml::cpu::repack::tensor_traits q8_0_16x1_q8_0; - static const ggml::cpu::repack::tensor_traits q2_K_16x1_q8_K; -#endif - - if (cur->type == GGML_TYPE_Q4_0) { - if (ggml_cpu_has_avx2() || (ggml_cpu_has_sve() && ggml_cpu_has_matmul_int8() && ggml_cpu_get_sve_cnt() == QK8_0)) { - if (cur->ne[1] % 8 == 0) { - return &q4_0_8x8_q8_0; - } - } - if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) { - if (cur->ne[1] % 4 == 0) { - return &q4_0_4x8_q8_0; - } - } - if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { - if (cur->ne[1] % 4 == 0) { - return &q4_0_4x4_q8_0; - } - } - if (ggml_cpu_has_riscv_v()) { - #if defined __riscv_zvfh - switch (__riscv_vlenb() * 8) { - case 128: { break; } // TODO - case 256: { if (cur->ne[1] % 16 == 0) { return &q4_0_16x1_q8_0; } break; } - case 512: { break; } // TODO - case 1024: { break; } // TODO - default: { return nullptr; } - } - #endif - } - } else if (cur->type == GGML_TYPE_Q4_K) { - if (ggml_cpu_has_avx2()) { - if (cur->ne[1] % 8 == 0) { - return &q4_K_8x8_q8_K; - } - } - if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) { - if (cur->ne[1] % 8 == 0) { - return &q4_K_8x8_q8_K; - } - } - if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { - if (cur->ne[1] % 8 == 0) { - return &q4_K_8x4_q8_K; - } - } - if (ggml_cpu_has_riscv_v()) { - #if defined __riscv_zvfh - switch (__riscv_vlenb() * 8) { - case 128: { break; } // TODO - case 256: { if (cur->ne[1] % 16 == 0) { return &q4_K_16x1_q8_K; } break; } - case 512: { break; } // TODO - case 1024: { break; } // TODO - default: { return nullptr; } - } - #endif - } - } else if (cur->type == GGML_TYPE_Q2_K) { - if (ggml_cpu_has_avx512()) { - if (cur->ne[1] % 8 == 0) { - return &q2_K_8x8_q8_K; - } - } - if (ggml_cpu_has_riscv_v()) { - #if defined __riscv_zvfh - switch (__riscv_vlenb() * 8) { - case 128: { break; } // TODO - case 256: { if (cur->ne[1] % 16 == 0) { return &q2_K_16x1_q8_K; } break; } - case 512: { break; } // TODO - case 1024: { break; } // TODO - default: { return nullptr; } - } - #endif - } - } else if (cur->type == GGML_TYPE_Q5_K) { - if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) { - if (cur->ne[1] % 8 == 0) { - return &q5_K_8x8_q8_K; - } - } - if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { - if (cur->ne[1] % 8 == 0) { - return &q5_K_8x4_q8_K; - } - } - } else if (cur->type == GGML_TYPE_Q6_K) { - if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) { - if (cur->ne[1] % 8 == 0) { - return &q6_K_8x8_q8_K; - } - } - if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { - if (cur->ne[1] % 8 == 0) { - return &q6_K_8x4_q8_K; - } - } - } else if (cur->type == GGML_TYPE_IQ4_NL) { - if (ggml_cpu_has_avx2()) { - if (cur->ne[1] % 8 == 0) { - return &iq4_nl_8x8_q8_0; - } - } - if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { - if (cur->ne[1] % 4 == 0) { - return &iq4_nl_4x4_q8_0; - } - } - if (ggml_cpu_has_riscv_v()) { - #if defined __riscv_zvfh - switch (__riscv_vlenb() * 8) { - case 128: { break; } // TODO - case 256: { if (cur->ne[1] % 16 == 0) { return &iq4_nl_16x1_q8_0; } break; } - case 512: { break; } // TODO - case 1024: { break; } // TODO - default: { return nullptr; } - } - #endif - } - } else if (cur->type == GGML_TYPE_MXFP4) { - if (ggml_cpu_has_avx2()) { - if (cur->ne[1] % 8 == 0) { - return &mxfp4_8x8_q8_0; - } - } - if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { - if (cur->ne[1] % 4 == 0) { - return &mxfp4_4x4_q8_0; - } - } - } else if (cur->type == GGML_TYPE_Q8_0) { - if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) { - if (cur->ne[1] % 4 == 0) { - return &q8_0_4x8_q8_0; - } - } - if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { - if (cur->ne[1] % 4 == 0) { - return &q8_0_4x4_q8_0; - } - } - if (ggml_cpu_has_riscv_v()) { - #if defined __riscv_zvfh - switch (__riscv_vlenb() * 8) { - case 128: { break; } // TODO - case 256: { if (cur->ne[1] % 16 == 0) { return &q8_0_16x1_q8_0; } break; } - case 512: { break; } // TODO - case 1024: { break; } // TODO - default: { return nullptr; } - } - #endif - } - } - - return nullptr; -} - -static enum ggml_status ggml_backend_cpu_repack_buffer_init_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor) { - tensor->extra = (void *) const_cast(ggml_repack_get_optimal_repack_type(tensor)); - - GGML_UNUSED(buffer); - return GGML_STATUS_SUCCESS; -} - -static void ggml_backend_cpu_repack_buffer_set_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, - const void * data, size_t offset, size_t size) { - GGML_ASSERT(offset == 0); - GGML_ASSERT(size == ggml_nbytes(tensor)); - - auto tensor_traits = (ggml::cpu::repack::tensor_traits_base *) tensor->extra; - auto OK = tensor_traits->repack(tensor, data, size); - - GGML_ASSERT(OK == 0); - GGML_UNUSED(buffer); -} - -static const char * ggml_backend_cpu_repack_buffer_type_get_name(ggml_backend_buffer_type_t buft) { - return "CPU_REPACK"; - - GGML_UNUSED(buft); -} - -static ggml_backend_buffer_t ggml_backend_cpu_repack_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { - ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); - - if (buffer == nullptr) { - return nullptr; - } - - buffer->buft = buft; - buffer->iface.init_tensor = ggml_backend_cpu_repack_buffer_init_tensor; - buffer->iface.set_tensor = ggml_backend_cpu_repack_buffer_set_tensor; - buffer->iface.get_tensor = nullptr; - buffer->iface.cpy_tensor = nullptr; - return buffer; -} - -static size_t ggml_backend_cpu_repack_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { - return TENSOR_ALIGNMENT; - - GGML_UNUSED(buft); -} - -namespace ggml::cpu::repack { -class extra_buffer_type : ggml::cpu::extra_buffer_type { - bool supports_op(ggml_backend_dev_t, const struct ggml_tensor * op) override { - if ( op->op == GGML_OP_MUL_MAT && - op->src[0]->buffer && - (ggml_n_dims(op->src[0]) == 2) && - op->src[0]->buffer->buft == ggml_backend_cpu_repack_buffer_type() && - ggml_repack_get_optimal_repack_type(op->src[0]) - ) { - if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) { - return false; - } - if (op->src[1]->type == GGML_TYPE_F32) { - return true; - } - //if (op->src[1]->type == GGML_TYPE_Q8_0) { - // return true; - //} - // may be possible if Q8_0 packed... - } else if (op->op == GGML_OP_MUL_MAT_ID - && op->src[0]->buffer - && (ggml_n_dims(op->src[0]) == 3) - && op->src[0]->buffer->buft == ggml_backend_cpu_repack_buffer_type() - && ggml_repack_get_optimal_repack_type(op->src[0]) - ) { - if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) { - return false; - } - if (op->src[1]->type == GGML_TYPE_F32) { - return true; - } - //if (op->src[1]->type == GGML_TYPE_Q8_0) { - // return true; - //} - } - return false; - } - - ggml::cpu::tensor_traits * get_tensor_traits(const struct ggml_tensor * op) override { - if (op->op == GGML_OP_MUL_MAT || op->op == GGML_OP_MUL_MAT_ID) { - if (op->src[0]->buffer && op->src[0]->buffer->buft == ggml_backend_cpu_repack_buffer_type()) { - return (ggml::cpu::tensor_traits *) op->src[0]->extra; - } - } - return nullptr; - } -}; -} // namespace ggml::cpu::repack - -ggml_backend_buffer_type_t ggml_backend_cpu_repack_buffer_type(void) { - static struct ggml_backend_buffer_type ggml_backend_cpu_buffer_type_repack = { - /* .iface = */ { - /* .get_name = */ ggml_backend_cpu_repack_buffer_type_get_name, - /* .alloc_buffer = */ ggml_backend_cpu_repack_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_cpu_repack_buffer_type_get_alignment, - /* .get_max_size = */ nullptr, // defaults to SIZE_MAX - /* .get_alloc_size = */ nullptr, // defaults to ggml_nbytes - /* .is_host = */ nullptr, - }, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0), - /* .context = */ new ggml::cpu::repack::extra_buffer_type(), - }; - - return &ggml_backend_cpu_buffer_type_repack; -} diff --git a/ggml/src/ggml-cpu/arch/x86/repack.cpp b/ggml/src/ggml-cpu/repack.cpp.inc similarity index 73% rename from ggml/src/ggml-cpu/arch/x86/repack.cpp rename to ggml/src/ggml-cpu/repack.cpp.inc index af1cebad131d..29b6cf554a4e 100644 --- a/ggml/src/ggml-cpu/arch/x86/repack.cpp +++ b/ggml/src/ggml-cpu/repack.cpp.inc @@ -1,13 +1,5568 @@ #define GGML_COMMON_IMPL_CPP #define GGML_COMMON_DECL_CPP -#include "ggml-common.h" -#include "ggml-backend-impl.h" +#include "ggml-common-defs.inc" +#include "ggml-backend-impl-defs.inc" + +#include "ggml-impl-defs.inc" +#include "ggml-cpu.h.inc" +#include "ggml-cpu-impl-defs.inc" +#include "simd-mappings-defs.inc" +#include "ggml.h.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml-cpu-impl-defs.inc" +#include "ggml-impl-defs.inc" +#include "simd-mappings-defs.inc" + +#ifdef __cplusplus +#include +extern "C" { +#endif + +// return true if op part of extra "accelerator" +bool ggml_cpu_extra_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op); +bool ggml_cpu_extra_work_size(int n_threads, const struct ggml_tensor * op, size_t * size); + +#ifdef __cplusplus +} + +namespace ggml::cpu { +// register in tensor->extra +class tensor_traits { + public: + virtual ~tensor_traits(); + virtual bool work_size(int n_threads, const struct ggml_tensor * op, size_t & size) = 0; + virtual bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) = 0; +}; + +class extra_buffer_type { + public: + virtual ~extra_buffer_type(); + virtual bool supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) = 0; + virtual tensor_traits * get_tensor_traits(const struct ggml_tensor * op) = 0; +}; +} // namespace ggml::cpu + +// implemented in ggml-cpu.cpp. +std::vector & ggml_backend_cpu_get_extra_buffer_types(); + +#endif + +#define GGML_FA_TILE_Q 64 +#define GGML_FA_TILE_KV 64 + +#ifdef __cplusplus + +#include + +// convenience functions/macros for use in template calls +// note: these won't be required after the 'traits' lookup table is used. +static inline ggml_fp16_t f32_to_f16(float x) { + return GGML_CPU_FP32_TO_FP16(x); +} + +static inline float f16_to_f32(ggml_fp16_t x) { + return GGML_CPU_FP16_TO_FP32(x); +} + +static inline ggml_bf16_t f32_to_bf16(float x) { + return GGML_FP32_TO_BF16(x); +} + +static inline float bf16_to_f32(ggml_bf16_t x) { + return GGML_BF16_TO_FP32(x); +} + +static inline float i32_to_f32(int32_t x) { + return x; +} + +static inline int32_t f32_to_i32(float x) { + return x; +} + +static inline float f32_to_f32(float x) { + return x; +} + +// TODO - merge this into the traits table, after using row-based conversions +template +struct type_conversion_table; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_fp16_t) = f16_to_f32; + static constexpr ggml_fp16_t (*from_f32)(float) = f32_to_f16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(float) = f32_to_f32; + static constexpr float (*from_f32)(float) = f32_to_f32; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_bf16_t) = bf16_to_f32; + static constexpr ggml_bf16_t (*from_f32)(float) = f32_to_bf16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(int32_t) = i32_to_f32; + static constexpr int32_t (*from_f32)(float) = f32_to_i32; +}; + +static std::pair get_thread_range(const struct ggml_compute_params * params, const struct ggml_tensor * src0) { + const int64_t ith = params->ith; + const int64_t nth = params->nth; + + const int64_t nr = ggml_nrows(src0); + + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; + + // row range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); + + return {ir0, ir1}; +} + +struct ggml_fa_tile_config { + static constexpr size_t Q = GGML_FA_TILE_Q; + static constexpr size_t KV = GGML_FA_TILE_KV; +}; + +#endif + + + +// Rename `_generic` functions if no native implementation is available. +// This effectively selects the generic implementation. + +#if defined(GGML_CPU_GENERIC) +// quants.c +#define quantize_row_q8_0_generic quantize_row_q8_0 +#define quantize_row_q8_1_generic quantize_row_q8_1 +#define quantize_row_q8_K_generic quantize_row_q8_K +#define ggml_vec_dot_q4_0_q8_0_generic ggml_vec_dot_q4_0_q8_0 +#define ggml_vec_dot_q4_1_q8_1_generic ggml_vec_dot_q4_1_q8_1 +#define ggml_vec_dot_q5_0_q8_0_generic ggml_vec_dot_q5_0_q8_0 +#define ggml_vec_dot_q5_1_q8_1_generic ggml_vec_dot_q5_1_q8_1 +#define ggml_vec_dot_q8_0_q8_0_generic ggml_vec_dot_q8_0_q8_0 +#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K +#define ggml_vec_dot_q3_K_q8_K_generic ggml_vec_dot_q3_K_q8_K +#define ggml_vec_dot_q4_K_q8_K_generic ggml_vec_dot_q4_K_q8_K +#define ggml_vec_dot_q5_K_q8_K_generic ggml_vec_dot_q5_K_q8_K +#define ggml_vec_dot_q6_K_q8_K_generic ggml_vec_dot_q6_K_q8_K +#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K +#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K +#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K +#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K +#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K +#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +#define ggml_vec_dot_iq4_nl_q8_0_generic ggml_vec_dot_iq4_nl_q8_0 +#define ggml_vec_dot_iq4_xs_q8_K_generic ggml_vec_dot_iq4_xs_q8_K +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__aarch64__) || defined(__arm__) || defined(_M_ARM) || defined(_M_ARM64) +// repack.cpp +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64) +// quants.c +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__POWERPC__) || defined(__powerpc__) +// ref: https://github.com/ggml-org/llama.cpp/pull/14146#issuecomment-2972561679 +// quants.c +#define quantize_row_q8_K_generic quantize_row_q8_K +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__loongarch64) +// quants.c +#define quantize_row_q8_K_generic quantize_row_q8_K +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__riscv) +// quants.c +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +// repack.cpp +#define ggml_quantize_mat_q8_0_4x1_generic ggml_quantize_mat_q8_0_4x1 +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_K_4x1_generic ggml_quantize_mat_q8_K_4x1 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__s390x__) +// quants.c +#define quantize_row_q8_K_generic quantize_row_q8_K +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K +#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K +#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K +#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K +#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K +#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K +#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#elif defined(__wasm__) +// quants.c +#define ggml_vec_dot_q4_1_q8_1_generic ggml_vec_dot_q4_1_q8_1 +#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K +#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K +#define ggml_vec_dot_iq2_xxs_q8_K_generic ggml_vec_dot_iq2_xxs_q8_K +#define ggml_vec_dot_iq2_xs_q8_K_generic ggml_vec_dot_iq2_xs_q8_K +#define ggml_vec_dot_iq2_s_q8_K_generic ggml_vec_dot_iq2_s_q8_K +#define ggml_vec_dot_iq3_xxs_q8_K_generic ggml_vec_dot_iq3_xxs_q8_K +#define ggml_vec_dot_iq3_s_q8_K_generic ggml_vec_dot_iq3_s_q8_K +#define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K +#define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K +#define ggml_vec_dot_iq4_nl_q8_0_generic ggml_vec_dot_iq4_nl_q8_0 +#define ggml_vec_dot_iq4_xs_q8_K_generic ggml_vec_dot_iq4_xs_q8_K +#define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 +#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +// repack.cpp +#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 +#define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 +#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 +#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8 +#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0 +#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0 +#define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0 +#define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K +#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K +#define ggml_gemv_q4_K_8x8_q8_K_generic ggml_gemv_q4_K_8x8_q8_K +#define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K +#define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K +#define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K +#define ggml_gemv_q6_K_8x8_q8_K_generic ggml_gemv_q6_K_8x8_q8_K +#define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0 +#define ggml_gemv_iq4_nl_8x8_q8_0_generic ggml_gemv_iq4_nl_8x8_q8_0 +#define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0 +#define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0 +#define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0 +#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0 +#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0 +#define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0 +#define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0 +#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K +#define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K +#define ggml_gemm_q4_K_8x8_q8_K_generic ggml_gemm_q4_K_8x8_q8_K +#define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K +#define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K +#define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K +#define ggml_gemm_q6_K_8x8_q8_K_generic ggml_gemm_q6_K_8x8_q8_K +#define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0 +#define ggml_gemm_iq4_nl_8x8_q8_0_generic ggml_gemm_iq4_nl_8x8_q8_0 +#define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0 +#define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0 +#define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0 +#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0 +#endif + + +#include +#include +#include +#include // for GGML_ASSERT + +#define GGML_COMMON_DECL_CPP +#include "ggml-common-defs.inc" + +#include "ggml.h.inc" + +// GGML internal header + +ggml_backend_buffer_type_t ggml_backend_cpu_repack_buffer_type(void); + +template constexpr int QK_0() { + if constexpr (K == 4) { + return QK4_0; + } + if constexpr (K == 8) { + return QK8_0; + } + return -1; +} + +template struct block { + ggml_half d[N]; // deltas for N qK_0 blocks + int8_t qs[(QK_0() * N * K) / 8]; // quants for N qK_0 blocks +}; + +// control size +static_assert(sizeof(block<4, 4>) == 4 * sizeof(ggml_half) + QK8_0 * 2, "wrong block<4,4> size/padding"); +static_assert(sizeof(block<4, 8>) == 8 * sizeof(ggml_half) + QK8_0 * 4, "wrong block<4,8> size/padding"); +static_assert(sizeof(block<4, 16>) == 16 * sizeof(ggml_half) + QK8_0 * 8, "wrong block<4,16> size/padding"); +static_assert(sizeof(block<8, 4>) == 4 * sizeof(ggml_half) + QK8_0 * 4, "wrong block<8,4> size/padding"); +static_assert(sizeof(block<8, 8>) == 8 * sizeof(ggml_half) + QK8_0 * 8, "wrong block<8,8> size/padding"); +static_assert(sizeof(block<8, 16>) == 16 * sizeof(ggml_half) + QK8_0 * 16, "wrong block<8,16> size/padding"); + +using block_q4_0x4 = block<4, 4>; +using block_q4_0x8 = block<4, 8>; +using block_q4_0x16 = block<4, 16>; +using block_q8_0x4 = block<8, 4>; +using block_q8_0x8 = block<8, 8>; +using block_q8_0x16 = block<8, 16>; + +struct block_q4_Kx8 { + ggml_half d[8]; // super-block scale for quantized scales + ggml_half dmin[8]; // super-block scale for quantized mins + uint8_t scales[96]; // scales and mins, quantized with 6 bits + uint8_t qs[1024]; // 4--bit quants +}; + +static_assert(sizeof(block_q4_Kx8) == sizeof(ggml_half) * 16 + K_SCALE_SIZE * 8 + QK_K * 4, "wrong q4_K block size/padding"); +struct block_q4_Kx16 { + ggml_half d[16]; // super-block scale for quantized scales + ggml_half dmin[16]; // super-block scale for quantized mins + uint8_t scales[192]; // scales and mins, quantized with 6 bits + uint8_t qs[2048]; // 4--bit quants +}; + +static_assert(sizeof(block_q4_Kx16) == sizeof(ggml_half) * 32 + K_SCALE_SIZE * 16 + QK_K * 8, "wrong q4_K block size/padding"); +struct block_q2_Kx8 { + ggml_half d[8]; // super-block scale for quantized scales + ggml_half dmin[8]; // super-block scale for quantized mins + uint8_t scales[128]; // scales and mins, quantized with 4 bits + uint8_t qs[512]; // 2--bit quants +}; + +static_assert(sizeof(block_q2_Kx8) == sizeof(ggml_half) * 16 + QK_K/2 + QK_K * 2, "wrong q2_K block size/padding"); +struct block_q2_Kx16 { + ggml_half d[16]; // Super-block scale for quantized scales + ggml_half dmin[16]; // Super-block scale for quantized mins + uint8_t scales[256]; // Sub-block scales (16 cols * 16 sub-blocks) + uint8_t qs[1024]; // Data (16 cols * 64 bytes per block) +}; +static_assert(sizeof(block_q2_Kx16) == sizeof(ggml_half) * 32 + QK_K + QK_K * 4, "wrong q2_K block size/padding"); + +struct block_q5_Kx8 { + ggml_half d[8]; // super-block scale for quantized scales + ggml_half dmin[8]; // super-block scale for quantized mins + uint8_t scales[96]; // scales and mins, quantized with 6 bits + uint8_t qh[QK_K * 8 / 8]; // high bits of 5-bit quants + uint8_t qs[QK_K * 8 / 2]; // low bits of 5-bit quants (in groups of 4) +}; + +static_assert(sizeof(block_q5_Kx8) == sizeof(ggml_half) * 16 + K_SCALE_SIZE * 8 + QK_K * 5, + "wrong q5_K block size/padding"); + +struct block_q6_Kx8 { + ggml_half d[8]; + int8_t scales[QK_K / 16 * 8]; + uint8_t ql[QK_K / 2 * 8]; // low bits of 6-bit quants (groups of 2) + uint8_t qh[QK_K / 4 * 8]; // high bits of 6-bit quants (groups of 4) +}; + +static_assert(sizeof(block_q6_Kx8) == sizeof(ggml_half) * 8 + QK_K / 16 * 8 + 3 * QK_K / 4 * 8, + "wrong q6_K block size/padding"); + +struct block_q8_Kx4 { + float d[4]; // delta + int8_t qs[QK_K * 4]; // quants + int16_t bsums[QK_K / 4]; // sum of quants in groups of 16 +}; + +static_assert(sizeof(block_q8_Kx4) == sizeof(float) * 4 + QK_K * 4 + (QK_K / 4) * sizeof(int16_t), "wrong q8_K block size/padding"); + +struct block_iq4_nlx4 { + ggml_half d[4]; // deltas for 4 iq4_nl blocks + uint8_t qs[QK4_NL * 2]; // nibbles / quants for 4 iq4_nl blocks +}; + +static_assert(sizeof(block_iq4_nlx4) == 4 * sizeof(ggml_half) + QK4_NL * 2, "wrong iq4_nlx4 block size/padding"); + +struct block_iq4_nlx8 { + ggml_half d[8]; // deltas for 8 iq4_nl blocks + uint8_t qs[QK4_NL * 4]; // nibbles / quants for 8 iq4_nl blocks +}; + +static_assert(sizeof(block_iq4_nlx8) == 8 * sizeof(ggml_half) + QK4_NL * 4, "wrong iq4_nlx8 block size/padding"); + +struct block_iq4_nlx16 { + ggml_half d[16]; // deltas for 16 iq4_nl blocks + uint8_t qs[QK4_NL * 8]; // nibbles / quants for 16 iq4_nl blocks +}; + +static_assert(sizeof(block_iq4_nlx16) == 16 * sizeof(ggml_half) + QK4_NL * 8, "wrong iq4_nlx16 block size/padding"); +struct block_mxfp4x4 { + uint8_t e[4]; + uint8_t qs[QK_MXFP4 * 2]; +}; +static_assert(sizeof(block_mxfp4x4) == 4 + QK_MXFP4 * 2, "wrong mxfp4x4 block size/padding"); + +struct block_mxfp4x8 { + uint8_t e[8]; + uint8_t qs[QK_MXFP4 * 4]; +}; +static_assert(sizeof(block_mxfp4x8) == 8 + QK_MXFP4 * 4, "wrong mxfp4x8 block size/padding"); + +#if defined(__cplusplus) +extern "C" { +#endif + +void ggml_quantize_mat_q8_0_4x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_0_4x8(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_K_4x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_K_4x8(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_gemv_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_0_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q2_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q5_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q5_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q6_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q6_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_iq4_nl_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_iq4_nl_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_mxfp4_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_mxfp4_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q8_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q8_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q2_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q5_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q5_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q6_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q6_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_iq4_nl_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_iq4_nl_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_mxfp4_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_mxfp4_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q8_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q8_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +#if defined __riscv_zvfh +void ggml_quantize_mat_q8_0_4x1(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_K_4x1(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_gemv_q4_0_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_K_16x1_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_iq4_nl_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q8_0_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q2_K_16x1_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_K_16x1_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_iq4_nl_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q8_0_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q2_K_16x1_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +#endif + +// Native implementations +void ggml_quantize_mat_q8_0_4x4_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_0_4x8_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_K_4x4_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_K_4x8_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_gemv_q4_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_0_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q2_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q5_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q5_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q6_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q6_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_iq4_nl_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_iq4_nl_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_mxfp4_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_mxfp4_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q8_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q8_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q2_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q5_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q5_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q6_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q6_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_iq4_nl_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_iq4_nl_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_mxfp4_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_mxfp4_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q8_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q8_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +#if defined __riscv_zvfh +void ggml_quantize_mat_q8_0_4x1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_quantize_mat_q8_K_4x1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void ggml_gemv_q4_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q4_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q8_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_q2_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemv_iq4_nl_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q4_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q8_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_q2_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +void ggml_gemm_iq4_nl_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); +#endif + +#if defined(__cplusplus) +} // extern "C" +#endif + + +#if defined(__GNUC__) +#pragma GCC diagnostic ignored "-Woverlength-strings" +#endif + +#define UNUSED GGML_UNUSED + +static inline int nearest_int(float fval) { + assert(fabsf(fval) <= 4194303.f); + float val = fval + 12582912.f; + int i; memcpy(&i, &val, sizeof(int)); + return (i & 0x007fffff) - 0x00400000; +} + +// Functions to create the interleaved data layout formats + +// interleave 4 block_q4_0s in blocks of blck_size_interleave +// returns an interleaved block_q4_0x4 +// in the interleaved block_q4_0x4, place deltas for 4 block_q4_0 blocks +// first, then interleave quants from 4 block_q4_0s in blocks of blck_size_interleave +// +// - in : an array of block_q4_0 pointers +// - blck_size_interleave : the block_q4_0 quants bytes are interleaved in blocks of +// blck_size_interleave bytes +// - xor_mask : the mask to convert the nibbles in block_q4_0 quants bytes +// from bias offset form to pure sign form (this saves subtract +// operations durin unpacking) +// + +extern "C" { + +#if defined __riscv_zvfh +void ggml_quantize_mat_q8_0_4x1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(QK8_0 == 32); + assert(k % QK8_0 == 0); + const int nb = k / QK8_0; + + block_q8_0x4 * GGML_RESTRICT y = (block_q8_0x4 *) vy; + + // scalar + const int blck_size_interleave = 1; + float srcv[4][QK8_0]; + float id[4]; + + for (int i = 0; i < nb; i++) { + for (int row_iter = 0; row_iter < 4; row_iter++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < QK8_0; j++) { + srcv[row_iter][j] = x[row_iter * k + i * QK8_0 + j]; + amax = MAX(amax, fabsf(srcv[row_iter][j])); + } + + const float d = amax / ((1 << 7) - 1); + id[row_iter] = d ? 1.0f / d : 0.0f; + + y[i].d[row_iter] = GGML_CPU_FP32_TO_FP16(d); + } + + for (int j = 0; j < QK8_0 * 4; j++) { + int src_offset = (j / (4 * blck_size_interleave)) * blck_size_interleave; + int src_id = (j % (4 * blck_size_interleave)) / blck_size_interleave; + src_offset += (j % blck_size_interleave); + + float x0 = srcv[src_id][src_offset] * id[src_id]; + y[i].qs[j] = roundf(x0); + } + } +} + +void ggml_quantize_mat_q8_K_4x1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(QK_K == 256); + assert(k % QK_K == 0); + const int nb = k / QK_K; + + block_q8_Kx4 * GGML_RESTRICT y = (block_q8_Kx4 *) vy; + + const int blck_size_interleave = 1; + float srcv[4][QK_K]; + float iscale[4]; + + for (int i = 0; i < nb; i++) { + for (int row_iter = 0; row_iter < 4; row_iter++) { + float amax = 0.0f; // absolute max + float max = 0; + + for (int j = 0; j < QK_K; j++) { + srcv[row_iter][j] = x[row_iter * k + i * QK_K + j]; + // Update the maximum value of the corresponding super block + if(amax < fabsf(srcv[row_iter][j])) { + amax = fabsf(srcv[row_iter][j]); + max = srcv[row_iter][j]; + } + } + + iscale[row_iter] = amax ? -127.f/max : 0; + y[i].d[row_iter] = amax ? 1/iscale[row_iter] : 0; + } + + for (int j = 0; j < QK_K / 4; j++) { + y[i].bsums[j] = 0; + } + for (int j = 0; j < QK_K * 4; j++) { + int src_id = j % 4; + int src_offset = j / 4; + int index = ((j >> 6) << 2) + (j & 3); + + float x0 = srcv[src_id][src_offset] * iscale[src_id]; + y[i].qs[j] = nearest_int(x0); + y[i].bsums[index] += y[i].qs[j]; + } + } +} +#endif + +void ggml_quantize_mat_q8_0_4x4_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(QK8_0 == 32); + assert(k % QK8_0 == 0); + const int nb = k / QK8_0; + + block_q8_0x4 * GGML_RESTRICT y = (block_q8_0x4 *) vy; + + // scalar + const int blck_size_interleave = 4; + float srcv[4][QK8_0]; + float id[4]; + + for (int i = 0; i < nb; i++) { + for (int row_iter = 0; row_iter < 4; row_iter++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < QK8_0; j++) { + srcv[row_iter][j] = x[row_iter * k + i * QK8_0 + j]; + amax = MAX(amax, fabsf(srcv[row_iter][j])); + } + + const float d = amax / ((1 << 7) - 1); + id[row_iter] = d ? 1.0f / d : 0.0f; + + y[i].d[row_iter] = GGML_CPU_FP32_TO_FP16(d); + } + + for (int j = 0; j < QK8_0 * 4; j++) { + int src_offset = (j / (4 * blck_size_interleave)) * blck_size_interleave; + int src_id = (j % (4 * blck_size_interleave)) / blck_size_interleave; + src_offset += (j % blck_size_interleave); + + float x0 = srcv[src_id][src_offset] * id[src_id]; + y[i].qs[j] = roundf(x0); + } + } +} + +void ggml_quantize_mat_q8_0_4x8_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(QK8_0 == 32); + assert(k % QK8_0 == 0); + const int nb = k / QK8_0; + + block_q8_0x4 * GGML_RESTRICT y = (block_q8_0x4 *) vy; + + // scalar + const int blck_size_interleave = 8; + float srcv[4][QK8_0]; + float id[4]; + + for (int i = 0; i < nb; i++) { + for (int row_iter = 0; row_iter < 4; row_iter++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < QK8_0; j++) { + srcv[row_iter][j] = x[row_iter * k + i * QK8_0 + j]; + amax = MAX(amax, fabsf(srcv[row_iter][j])); + } + + const float d = amax / ((1 << 7) - 1); + id[row_iter] = d ? 1.0f / d : 0.0f; + + y[i].d[row_iter] = GGML_CPU_FP32_TO_FP16(d); + } + + for (int j = 0; j < QK8_0 * 4; j++) { + int src_offset = (j / (4 * blck_size_interleave)) * blck_size_interleave; + int src_id = (j % (4 * blck_size_interleave)) / blck_size_interleave; + src_offset += (j % blck_size_interleave); + + float x0 = srcv[src_id][src_offset] * id[src_id]; + y[i].qs[j] = roundf(x0); + } + } +} + +void ggml_quantize_mat_q8_K_4x4_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(QK_K == 256); + assert(k % QK_K == 0); + const int nb = k / QK_K; + + block_q8_Kx4 * GGML_RESTRICT y = (block_q8_Kx4 *) vy; + + // scalar + const int blck_size_interleave = 4; + float srcv[4][QK_K]; + float iscale[4]; + + for (int i = 0; i < nb; i++) { + for (int row_iter = 0; row_iter < 4; row_iter++) { + float amax = 0.0f; // absolute max + float max = 0; + + for (int j = 0; j < QK_K; j++) { + srcv[row_iter][j] = x[row_iter * k + i * QK_K + j]; + // Update the maximum value of the corresponding super block + if(amax < fabsf(srcv[row_iter][j])) { + amax = fabsf(srcv[row_iter][j]); + max = srcv[row_iter][j]; + } + } + + iscale[row_iter] = amax ? -127.f/max : 0; + + y[i].d[row_iter] = amax ? 1/iscale[row_iter] : 0; + } + + for (int j = 0; j < QK_K / 4; j++) { + y[i].bsums[j] = 0; + } + + // Quants values are interleaved in sequence of four bytes from corresponding super blocks + // Bsums values are interleaved in sequence of four bsums from each super block taken for interleaving + // i.e first four bsums from the first super block, followed by first four bsums from second super block and so on + for (int j = 0; j < QK_K * 4; j++) { + int src_offset = (j / (4 * blck_size_interleave)) * blck_size_interleave; + int src_id = (j % (4 * blck_size_interleave)) / blck_size_interleave; + src_offset += (j % blck_size_interleave); + int index = (((j & 15) >> 2) << 2) + ((j >> 8) << 4) + ((j >> 6) & 3); + + float x0 = srcv[src_id][src_offset] * iscale[src_id]; + y[i].qs[j] = nearest_int(x0); + y[i].bsums[index] += y[i].qs[j]; + } + } +} + +void ggml_quantize_mat_q8_K_4x8_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(QK_K == 256); + assert(k % QK_K == 0); + const int nb = k / QK_K; + + block_q8_Kx4 * GGML_RESTRICT y = (block_q8_Kx4 *) vy; + + // scalar + const int blck_size_interleave = 8; + float srcv[4][QK_K]; + float iscale[4]; + + for (int i = 0; i < nb; i++) { + for (int row_iter = 0; row_iter < 4; row_iter++) { + float amax = 0.0f; // absolute max + float max = 0; + + for (int j = 0; j < QK_K; j++) { + srcv[row_iter][j] = x[row_iter * k + i * QK_K + j]; + // Update the maximum value of the corresponding super block + if(amax < fabsf(srcv[row_iter][j])) { + amax = fabsf(srcv[row_iter][j]); + max = srcv[row_iter][j]; + } + } + + iscale[row_iter] = amax ? -127.f/max : 0; + + y[i].d[row_iter] = amax ? 1/iscale[row_iter] : 0; + } + + for (int j = 0; j < QK_K / 4; j++) { + y[i].bsums[j] = 0; + } + + // Quants values are interleaved in sequence of eight bytes from corresponding super blocks + // Bsums values are interleaved in sequence of four bsums from each super block taken for interleaving + // i.e first four bsums from the first super block, followed by first four bsums from second super block and so on + for (int j = 0; j < QK_K * 4; j++) { + int src_offset = (j / (4 * blck_size_interleave)) * blck_size_interleave; + int src_id = (j % (4 * blck_size_interleave)) / blck_size_interleave; + src_offset += (j % blck_size_interleave); + int index = (((j & 31) >> 3) << 2) + ((j >> 8) << 4) + ((j >> 6) & 3); + + float x0 = srcv[src_id][src_offset] * iscale[src_id]; + y[i].qs[j] = nearest_int(x0); + y[i].bsums[index] += y[i].qs[j]; + } + } +} + +} // extern "C" + +template +void ggml_quantize_mat_t(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row); + +template <> void ggml_quantize_mat_t<4, GGML_TYPE_Q8_0>(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row) { + assert(nrow == 4); + UNUSED(nrow); + ggml_quantize_mat_q8_0_4x4(x, vy, n_per_row); +} + +template <> void ggml_quantize_mat_t<8, GGML_TYPE_Q8_0>(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row) { + assert(nrow == 4); + UNUSED(nrow); + ggml_quantize_mat_q8_0_4x8(x, vy, n_per_row); +} + +template <> void ggml_quantize_mat_t<4, GGML_TYPE_Q8_K>(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row) { + assert(nrow == 4); + UNUSED(nrow); + ggml_quantize_mat_q8_K_4x4(x, vy, n_per_row); +} + +template <> void ggml_quantize_mat_t<8, GGML_TYPE_Q8_K>(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row) { + assert(nrow == 4); + UNUSED(nrow); + ggml_quantize_mat_q8_K_4x8(x, vy, n_per_row); +} + +#if defined __riscv_zvfh +template <> void ggml_quantize_mat_t<1, GGML_TYPE_Q8_0>(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row) { + assert(nrow == 4); + UNUSED(nrow); + ggml_quantize_mat_q8_0_4x1(x, vy, n_per_row); +} + +template <> void ggml_quantize_mat_t<1, GGML_TYPE_Q8_K>(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t nrow, int64_t n_per_row) { + assert(nrow == 4); + UNUSED(nrow); + ggml_quantize_mat_q8_K_4x1(x, vy, n_per_row); +} +#endif + +template +static void ggml_gemv_q6_K_NxM_q8_K_generic_impl(int n, + float * GGML_RESTRICT s, + size_t bs, + const void * GGML_RESTRICT vx, + const void * GGML_RESTRICT vy, + int nr, + int nc) { + constexpr int blocklen = M; + constexpr int ncols_interleaved = N; + const int qk = QK_K; + const int nb = n / qk; + const int blocks_per_half = 64 / blocklen; + + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[8]; + + const block_q8_K * a_ptr = (const block_q8_K *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q6_Kx8 * b_ptr = (const block_q6_Kx8 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) { + sumf[j] = 0.0f; + } + + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + const int base_l = (k / blocks_per_half) * 128 + (k % blocks_per_half) * blocklen; + const int base_h = base_l + 64; + + const int scale_idx_l = base_l / 16; + const int scale_idx_h = base_h / 16; + + const int qh_shift_l = ((base_l % 128) / 32) * 2; + const int qh_shift_h = ((base_h % 128) / 32) * 2; + + const int qh_half_l = (base_l / 128) * 32; + const int qh_half_h = (base_h / 128) * 32; + + for (int j = 0; j < ncols_interleaved; j++) { + const int8_t scale_l = b_ptr[l].scales[scale_idx_l * ncols_interleaved + j]; + const int8_t scale_h = b_ptr[l].scales[scale_idx_h * ncols_interleaved + j]; + + int sumi_l = 0; + int sumi_h = 0; + + for (int i = 0; i < blocklen; i++) { + const int ql_pos = k * ncols_interleaved * blocklen + j * blocklen + i; + const int l_4 = b_ptr[l].ql[ql_pos] & 0xF; + const int hi_4 = (b_ptr[l].ql[ql_pos] >> 4) & 0xF; + + const int qh_idx_l = qh_half_l + ((base_l + i) % 32); + const int qh_chunk_l = qh_idx_l / blocklen; + const int qh_pos_l = qh_idx_l % blocklen; + const int qh_offset_l = qh_chunk_l * (blocklen * ncols_interleaved) + j * blocklen + qh_pos_l; + const int hi_2_l = (b_ptr[l].qh[qh_offset_l] >> qh_shift_l) & 0x3; + + const int qh_idx_h = qh_half_h + ((base_h + i) % 32); + const int qh_chunk_h = qh_idx_h / blocklen; + const int qh_pos_h = qh_idx_h % blocklen; + const int qh_offset_h = qh_chunk_h * (blocklen * ncols_interleaved) + j * blocklen + qh_pos_h; + const int hi_2_h = (b_ptr[l].qh[qh_offset_h] >> qh_shift_h) & 0x3; + + const int q_l = ((hi_2_l << 4) | l_4) - 32; + const int q_h = ((hi_2_h << 4) | hi_4) - 32; + + const int8_t a_l = a_ptr[l].qs[base_l + i]; + const int8_t a_h = a_ptr[l].qs[base_h + i]; + + sumi_l += q_l * a_l; + sumi_h += q_h * a_h; + } + + sumf[j] += + (sumi_l * scale_l + sumi_h * scale_h) * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d; + } + } + } + + for (int j = 0; j < ncols_interleaved; j++) { + s[x * ncols_interleaved + j] = sumf[j]; + } + } +} + +template +static void ggml_gemm_q6_K_NxM_q8_K_generic_impl(int n, + float * GGML_RESTRICT s, + size_t bs, + const void * GGML_RESTRICT vx, + const void * GGML_RESTRICT vy, + int nr, + int nc) { + constexpr int blocklen = M; + constexpr int ncols_interleaved = N; + const int qk = QK_K; + const int nb = n / qk; + const int blocks_per_half = 64 / blocklen; + const int q8_half_stride = 512; + const int q8_low_high_step = 256; + + assert(n % qk == 0); + assert(nr % 4 == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(bs); + + float sumf[4][8]; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_Kx4 * a_ptr = (const block_q8_Kx4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q6_Kx8 * b_ptr = (const block_q6_Kx8 *) vx + (x * nb); + + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumf[m][j] = 0.0f; + } + } + + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + const int base_l = (k / blocks_per_half) * 128 + (k % blocks_per_half) * blocklen; + const int base_h = base_l + 64; + + const int scale_idx_l = base_l / 16; + const int scale_idx_h = base_h / 16; + + const int qh_shift_l = ((base_l % 128) / 32) * 2; + const int qh_shift_h = ((base_h % 128) / 32) * 2; + + const int qh_half_l = (base_l / 128) * 32; + const int qh_half_h = (base_h / 128) * 32; + + const int q8_base = (k / blocks_per_half) * q8_half_stride + (k % blocks_per_half) * (blocklen * 4); + + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + const int8_t scale_l = b_ptr[l].scales[scale_idx_l * ncols_interleaved + j]; + const int8_t scale_h = b_ptr[l].scales[scale_idx_h * ncols_interleaved + j]; + + int sumi_l = 0; + int sumi_h = 0; + + for (int i = 0; i < blocklen; i++) { + const int ql_pos = k * ncols_interleaved * blocklen + j * blocklen + i; + const int l_4 = b_ptr[l].ql[ql_pos] & 0xF; + const int hi_4 = (b_ptr[l].ql[ql_pos] >> 4) & 0xF; + + const int qh_idx_l = qh_half_l + ((base_l + i) % 32); + const int qh_chunk_l = qh_idx_l / blocklen; + const int qh_pos_l = qh_idx_l % blocklen; + const int qh_offset_l = + qh_chunk_l * (blocklen * ncols_interleaved) + j * blocklen + qh_pos_l; + const int hi_2_l = (b_ptr[l].qh[qh_offset_l] >> qh_shift_l) & 0x3; + + const int qh_idx_h = qh_half_h + ((base_h + i) % 32); + const int qh_chunk_h = qh_idx_h / blocklen; + const int qh_pos_h = qh_idx_h % blocklen; + const int qh_offset_h = + qh_chunk_h * (blocklen * ncols_interleaved) + j * blocklen + qh_pos_h; + const int hi_2_h = (b_ptr[l].qh[qh_offset_h] >> qh_shift_h) & 0x3; + + const int q_l = ((hi_2_l << 4) | l_4) - 32; + const int q_h = ((hi_2_h << 4) | hi_4) - 32; + + const int8_t q8_l = a_ptr[l].qs[q8_base + m * blocklen + i]; + const int8_t q8_h = a_ptr[l].qs[q8_base + m * blocklen + i + q8_low_high_step]; + + sumi_l += q_l * q8_l; + sumi_h += q_h * q8_h; + } + + sumf[m][j] += (sumi_l * scale_l + sumi_h * scale_h) * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * + a_ptr[l].d[m]; + } + } + } + } + + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } + } +} + +template +static void ggml_gemv_q5_K_NxM_q8_K_generic_impl(int n, + float * GGML_RESTRICT s, + size_t bs, + const void * GGML_RESTRICT vx, + const void * GGML_RESTRICT vy, + int nr, + int nc) { + constexpr int blocklen = M; + constexpr int ncols_interleaved = N; + const int qk = QK_K; + const int nb = n / qk; + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[ncols_interleaved]; + float sum_minf[ncols_interleaved]; + uint32_t utmp[32]; + int sumi1; + int sumi2; + int sumi; + + const block_q8_K * a_ptr = (const block_q8_K *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q5_Kx8 * b_ptr = (const block_q5_Kx8 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) { + sumf[j] = 0.0; + sum_minf[j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int sb = 0; sb < 8; sb++) { + memcpy(utmp + sb * 4, b_ptr[l].scales + sb * K_SCALE_SIZE, K_SCALE_SIZE); + utmp[sb * 4 + 3] = ((utmp[sb * 4 + 2] >> 4) & kmask2) | (((utmp[sb * 4 + 1] >> 6) & kmask3) << 4); + const uint32_t uaux_0 = utmp[sb * 4 + 1] & kmask1; + utmp[sb * 4 + 1] = (utmp[sb * 4 + 2] & kmask2) | (((utmp[sb * 4 + 0] >> 6) & kmask3) << 4); + utmp[sb * 4 + 2] = uaux_0; + utmp[sb * 4 + 0] &= kmask1; + } + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + constexpr int scale_stride = 32; + uint8_t * scales_0 = (uint8_t *) utmp + (k / (32 / blocklen)) * scale_stride; + uint8_t * scales_1 = (uint8_t *) utmp + (k / (32 / blocklen)) * scale_stride + 16; + + const int qh_shift = (k / (32 / blocklen)) * 2; + for (int j = 0; j < ncols_interleaved; j++) { + sumi1 = 0; + sumi2 = 0; + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int b_qs_offset = k * ncols_interleaved * blocklen + j * blocklen + i; + + const int qh_idx = (k * blocklen + i) % 32; + const int qh_chunk = qh_idx / blocklen; + const int qh_pos = qh_idx % blocklen; + const int b_qh_offset = qh_chunk * (blocklen * ncols_interleaved) + j * blocklen + qh_pos; + + const uint8_t qh_val = b_ptr[l].qh[b_qh_offset]; + const uint8_t h0 = (qh_val >> qh_shift) & 1; + const uint8_t h1 = (qh_val >> (qh_shift + 1)) & 1; + + const int v0 = (int8_t) ((b_ptr[l].qs[b_qs_offset] & 0xF) | (h0 << 4)); + const int v1 = (int8_t) ((b_ptr[l].qs[b_qs_offset] >> 4) | (h1 << 4)); + + const int q8_offset = (k / (32 / blocklen)) * 64 + (k % (32 / blocklen)) * blocklen + i; + + sumi1 = (v0 * a_ptr[l].qs[q8_offset]); + sumi2 = (v1 * a_ptr[l].qs[q8_offset + 32]); + sumi1 = sumi1 * scales_0[j]; + sumi2 = sumi2 * scales_1[j]; + sumi += sumi1 + sumi2; + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d; + } + } + for (int sb = 0; sb < 8; sb++) { + uint8_t * mins = (uint8_t *) utmp + 8 + sb * 16; + for (int j = 0; j < ncols_interleaved; j++) { + sum_minf[j] += mins[j] * (a_ptr[l].bsums[sb * 2] + a_ptr[l].bsums[sb * 2 + 1]) * + GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d; + } + } + } + for (int j = 0; j < ncols_interleaved; j++) { + s[x * ncols_interleaved + j] = sumf[j] - sum_minf[j]; + } + } +} + +template +static void ggml_gemm_q5_K_NxM_q8_K_generic_impl(int n, + float * GGML_RESTRICT s, + size_t bs, + const void * GGML_RESTRICT vx, + const void * GGML_RESTRICT vy, + int nr, + int nc) { + constexpr int blocklen = M; + constexpr int ncols_interleaved = N; + const int qk = QK_K; + const int nb = n / qk; + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + assert(n % qk == 0); + assert(nr % 4 == 0); + assert(nc % ncols_interleaved == 0); + + float sumf[4][ncols_interleaved]; + float sum_minf[4][ncols_interleaved]; + uint32_t utmp[32]; + int sumi1; + int sumi2; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_Kx4 * a_ptr = (const block_q8_Kx4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q5_Kx8 * b_ptr = (const block_q5_Kx8 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumf[m][j] = 0.0; + sum_minf[m][j] = 0.0; + } + } + for (int l = 0; l < nb; l++) { + for (int sb = 0; sb < 8; sb++) { + memcpy(utmp + sb * 4, b_ptr[l].scales + sb * K_SCALE_SIZE, K_SCALE_SIZE); + utmp[sb * 4 + 3] = ((utmp[sb * 4 + 2] >> 4) & kmask2) | (((utmp[sb * 4 + 1] >> 6) & kmask3) << 4); + const uint32_t uaux_0 = utmp[sb * 4 + 1] & kmask1; + utmp[sb * 4 + 1] = (utmp[sb * 4 + 2] & kmask2) | (((utmp[sb * 4 + 0] >> 6) & kmask3) << 4); + utmp[sb * 4 + 2] = uaux_0; + utmp[sb * 4 + 0] &= kmask1; + } + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + constexpr int scale_stride = 32; + uint8_t * scales_0 = (uint8_t *) utmp + (k / (32 / blocklen)) * scale_stride; + uint8_t * scales_1 = (uint8_t *) utmp + (k / (32 / blocklen)) * scale_stride + 16; + + const int qh_shift = (k / (32 / blocklen)) * 2; + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi1 = 0; + sumi2 = 0; + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int b_qs_offset = k * ncols_interleaved * blocklen + j * blocklen + i; + + const int qh_idx = (k * blocklen + i) % 32; + const int qh_chunk = qh_idx / blocklen; + const int qh_pos = qh_idx % blocklen; + const int b_qh_offset = + qh_chunk * (blocklen * ncols_interleaved) + j * blocklen + qh_pos; + + const uint8_t qh_val = b_ptr[l].qh[b_qh_offset]; + const uint8_t h0 = (qh_val >> qh_shift) & 1; + const uint8_t h1 = (qh_val >> (qh_shift + 1)) & 1; + + const int v0 = (int8_t) ((b_ptr[l].qs[b_qs_offset] & 0xF) | (h0 << 4)); + const int v1 = (int8_t) ((b_ptr[l].qs[b_qs_offset] >> 4) | (h1 << 4)); + + const int q8_offset = (k / (32 / blocklen)) * 256 + + (k % (32 / blocklen)) * 4 * blocklen + m * blocklen + i; + + sumi1 = (v0 * a_ptr[l].qs[q8_offset]); + sumi2 = (v1 * a_ptr[l].qs[q8_offset + 128]); + sumi1 = sumi1 * scales_0[j]; + sumi2 = sumi2 * scales_1[j]; + sumi += sumi1 + sumi2; + } + sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d[m]; + } + } + } + for (int sb = 0; sb < 8; sb++) { + uint8_t * mins = (uint8_t *) utmp + 8 + sb * 16; + for (int m = 0; m < 4; m++) { + const int16_t * bsums = a_ptr[l].bsums + (sb * 8) + (m * 4) - ((sb % 2) * 6); + for (int j = 0; j < ncols_interleaved; j++) { + sum_minf[m][j] += mins[j] * (bsums[0] + bsums[1]) * + GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d[m]; + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j] - sum_minf[m][j]; + } + } + } + } +} + +extern "C" { + +void ggml_gemv_q4_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + const int blocklen = 4; + + assert(nr == 1); + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + float sumf[4]; + int sumi; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_0x4 * b_ptr = (const block_q4_0x4 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); + const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); + sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])) >> 4; + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); + } + } + } + for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; + } +} + +void ggml_gemv_q4_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + const int blocklen = 8; + + assert (n % qk == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + float sumf[4]; + int sumi; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_0x4 * b_ptr = (const block_q4_0x4 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); + const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); + sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])) >> 4; + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); + } + } + } + for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; + } +} + +void ggml_gemv_q4_0_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 8; + const int blocklen = 8; + + assert (n % qk == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + float sumf[8]; + int sumi; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_0x8 * b_ptr = (const block_q4_0x8 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); + const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); + sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])) >> 4; + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); + } + } + } + for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; + } +} + +void ggml_gemv_q4_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK_K; + const int nb = n / qk; + const int ncols_interleaved = 8; + const int blocklen = 4; + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + assert (n % qk == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[8]; + float sum_minf[8]; + uint32_t utmp[32]; + int sumi1; + int sumi2; + int sumi; + + const block_q8_K * a_ptr = (const block_q8_K *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_Kx8 * b_ptr = (const block_q4_Kx8 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) { + sumf[j] = 0.0; + sum_minf[j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int sb = 0; sb < 8; sb++) { + memcpy(utmp + sb * 4, b_ptr[l].scales + sb * 12, 12); + utmp[sb * 4 + 3] = ((utmp[sb * 4 + 2] >> 4) & kmask2) | (((utmp[sb * 4 + 1] >> 6) & kmask3) << 4); + const uint32_t uaux_0 = utmp[sb * 4 + 1] & kmask1; + utmp[sb * 4 + 1] = (utmp[sb * 4 + 2] & kmask2) | (((utmp[sb * 4 + 0] >> 6) & kmask3) << 4); + utmp[sb * 4 + 2] = uaux_0; + utmp[sb * 4 + 0] &= kmask1; + } + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + uint8_t * scales_0 = (uint8_t *) utmp + (k / 8) * 32; + uint8_t * scales_1 = (uint8_t *) utmp + (k / 8) * 32 + 16; + for (int j = 0; j < ncols_interleaved; j++) { + sumi1 = 0; + sumi2 = 0; + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF); + const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4); + sumi1 = (v0 * a_ptr[l].qs[(k / 8) * 64 + (k % 8) * blocklen + i]); + sumi2 = (v1 * a_ptr[l].qs[(k / 8) * 64 + (k % 8) * blocklen + i + 32]); + sumi1 = sumi1 * scales_0[j]; + sumi2 = sumi2 * scales_1[j]; + sumi += sumi1 + sumi2; + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d; + } + } + for (int sb = 0; sb < 8; sb++) { + uint8_t * mins = (uint8_t *) utmp + 8 + sb * 16; + for (int j = 0; j < ncols_interleaved; j++) { + sum_minf[j] += mins[j] * (a_ptr[l].bsums[sb * 2] + a_ptr[l].bsums[sb * 2 + 1]) * GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d; + } + } + } + for (int j = 0; j < ncols_interleaved; j++) { + s[x * ncols_interleaved + j] = sumf[j] - sum_minf[j]; + } + } +} + +void ggml_gemv_q4_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK_K; + const int nb = n / qk; + const int ncols_interleaved = 8; + const int blocklen = 8; + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + assert (n % qk == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[8]; + float sum_minf[8]; + uint32_t utmp[32]; + int sumi1; + int sumi2; + int sumi; + + const block_q8_K * a_ptr = (const block_q8_K *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_Kx8 * b_ptr = (const block_q4_Kx8 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) { + sumf[j] = 0.0; + sum_minf[j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int sb = 0; sb < 8; sb++) { + memcpy(utmp + sb * 4, b_ptr[l].scales + sb * 12, 12); + utmp[sb * 4 + 3] = ((utmp[sb * 4 + 2] >> 4) & kmask2) | (((utmp[sb * 4 + 1] >> 6) & kmask3) << 4); + const uint32_t uaux_0 = utmp[sb * 4 + 1] & kmask1; + utmp[sb * 4 + 1] = (utmp[sb * 4 + 2] & kmask2) | (((utmp[sb * 4 + 0] >> 6) & kmask3) << 4); + utmp[sb * 4 + 2] = uaux_0; + utmp[sb * 4 + 0] &= kmask1; + } + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + uint8_t *scales_0 = (uint8_t*) utmp + (k / 4) * 32; + uint8_t *scales_1 = (uint8_t*) utmp + (k / 4) * 32 + 16; + for (int j = 0; j < ncols_interleaved; j++) { + sumi1 = 0; + sumi2 = 0; + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF); + const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4); + sumi1 = (v0 * a_ptr[l].qs[(k >> 2) * 64 + (k % 4) * blocklen + i]); + sumi2 = (v1 * a_ptr[l].qs[(k >> 2) * 64 + (k % 4) * blocklen + i + 32]); + sumi1 = sumi1 * scales_0[j]; + sumi2 = sumi2 * scales_1[j]; + sumi += sumi1 + sumi2; + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d; + } + } + for (int sb = 0; sb < 8; sb++) { + uint8_t *mins = (uint8_t*) utmp + 8 + sb * 16; + for (int j = 0; j < ncols_interleaved; j++) { + sum_minf[j] += mins[j] * (a_ptr[l].bsums[sb * 2] + a_ptr[l].bsums[sb * 2 + 1]) * GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d; + } + } + } + for (int j = 0; j < ncols_interleaved; j++) { + s[x * ncols_interleaved + j] = sumf[j] - sum_minf[j]; + } + } +} + +void ggml_gemv_q2_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK_K; + const int nb = n / qk; + const int ncols_interleaved = 8; + const int blocklen = 8; + + assert (n % qk == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + float sumf[8]; + float sum_minf[8]; + int sumi1,sumi2,sumi3,sumi4; + int sumi; + + const block_q8_K * a_ptr = (const block_q8_K *)vy; + for(int x = 0; x < nc / ncols_interleaved; x++) { + const block_q2_Kx8 * b_ptr = (const block_q2_Kx8 *) vx + (x * nb); + for (int j = 0; j < ncols_interleaved; j++) { + sumf[j] = 0.0; + sum_minf[j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (4 * blocklen)); k++) { + const uint8_t *scales_0 = b_ptr[l].scales + (k / 4) * 64 ; + const uint8_t *scales_1 = b_ptr[l].scales + (k / 4) * 64 + 16; + const uint8_t *scales_2 = b_ptr[l].scales + (k / 4) * 64 + 32; + const uint8_t *scales_3 = b_ptr[l].scales + (k / 4) * 64 + 48; + for (int j = 0; j < ncols_interleaved; j++) { + sumi1 = 0; + sumi2 = 0; + sumi3 = 0; + sumi4 = 0; + sumi = 0; + int offset = ((k / 2) % 2) + j * 2; + for (int i = 0; i < blocklen; ++i){ + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 3); + const int v1 = (int8_t) ((b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 2 ) & 3); + const int v2 = (int8_t) ((b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4 ) & 3); + const int v3 = (int8_t) ((b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 6 ) & 3); + sumi1 = (v0 * a_ptr[l].qs[(k >> 2) * 128 + (k % 4) * blocklen + i]); + sumi2 = (v1 * a_ptr[l].qs[(k >> 2) * 128 + (k % 4) * blocklen + i + 32]); + sumi3 = (v2 * a_ptr[l].qs[(k >> 2) * 128 + (k % 4) * blocklen + i + 64]); + sumi4 = (v3 * a_ptr[l].qs[(k >> 2) * 128 + (k % 4) * blocklen + i + 96]); + + sumi1 = sumi1 * (scales_0[offset] & 0xF); + sumi2 = sumi2 * (scales_1[offset] & 0xF); + sumi3 = sumi3 * (scales_2[offset] & 0xF); + sumi4 = sumi4 * (scales_3[offset] & 0xF); + sumi += sumi1 + sumi2 + sumi3 + sumi4; + } + sumf[j] += sumi * GGML_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d; + } + } + for(int sb = 0; sb < 8; sb++) { + const uint8_t *mins = b_ptr[l].scales + sb * 16; + for(int j = 0; j < ncols_interleaved; j++){ + sum_minf[j] += ((mins[j * 2] >> 4) * a_ptr[l].bsums[sb * 2] + (mins[(j * 2)+ 1] >> 4) * a_ptr[l].bsums[sb * 2 + 1]) * GGML_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d; + } + } + } + for (int j = 0; j < ncols_interleaved; j++) { + s[x * ncols_interleaved + j] = sumf[j] - sum_minf[j]; + } + } +} + +void ggml_gemv_q5_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + ggml_gemv_q5_K_NxM_q8_K_generic_impl<4, 8>(n, s, bs, vx, vy, nr, nc); +} + +void ggml_gemv_q5_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + ggml_gemv_q5_K_NxM_q8_K_generic_impl<8, 8>(n, s, bs, vx, vy, nr, nc); +} + + +void ggml_gemv_q6_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + ggml_gemv_q6_K_NxM_q8_K_generic_impl<4, 8>(n, s, bs, vx, vy, nr, nc); +} + +void ggml_gemv_q6_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + ggml_gemv_q6_K_NxM_q8_K_generic_impl<8, 8>(n, s, bs, vx, vy, nr, nc); +} + +void ggml_gemv_iq4_nl_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + const int blocklen = 4; + + assert(nr == 1); + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[4]; + int sumi; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_iq4_nlx4 * b_ptr = (const block_iq4_nlx4 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; + const int v1 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; + sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])); + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); + } + } + } + for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; + } +} + +void ggml_gemv_iq4_nl_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 8; + const int blocklen = 8; + + assert(nr == 1); + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[8]; + int sumi; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_iq4_nlx8 * b_ptr = (const block_iq4_nlx8 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; + const int v1 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; + sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])); + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); + } + } + } + for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; + } +} + +void ggml_gemv_mxfp4_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + const int blocklen = 4; + + assert(nr == 1); + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[4]; + int sumi; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_mxfp4x4 * b_ptr = (const block_mxfp4x4 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; + const int v1 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; + sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])); + } + sumf[j] += sumi * GGML_CPU_E8M0_TO_FP32_HALF(b_ptr[l].e[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); + } + } + } + for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; + } +} + +void ggml_gemv_mxfp4_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 8; + const int blocklen = 8; + + assert(nr == 1); + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[8]; + int sumi; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_mxfp4x8 * b_ptr = (const block_mxfp4x8 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; + const int v1 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; + sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])); + } + sumf[j] += sumi * GGML_CPU_E8M0_TO_FP32_HALF(b_ptr[l].e[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); + } + } + } + for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; + } +} + +void ggml_gemv_q8_0_4x4_q8_0_generic(int n, + float * GGML_RESTRICT s, + size_t bs, + const void * GGML_RESTRICT vx, + const void * GGML_RESTRICT vy, + int nr, + int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + const int blocklen = 4; + + assert(nr == 1); + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[4]; + int sumi; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q8_0x4 * b_ptr = (const block_q8_0x4 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) { + sumf[j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / blocklen); k++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i]; + sumi += v0 * a_ptr[l].qs[k * blocklen + i]; + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); + } + } + } + for (int j = 0; j < ncols_interleaved; j++) { + s[x * ncols_interleaved + j] = sumf[j]; + } + } +} + +void ggml_gemv_q8_0_4x8_q8_0_generic(int n, + float * GGML_RESTRICT s, + size_t bs, + const void * GGML_RESTRICT vx, + const void * GGML_RESTRICT vy, + int nr, + int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + const int blocklen = 8; + + assert(nr == 1); + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[4]; + int sumi; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q8_0x4 * b_ptr = (const block_q8_0x4 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) { + sumf[j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / blocklen); k++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i]; + sumi += v0 * a_ptr[l].qs[k * blocklen + i]; + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); + } + } + } + for (int j = 0; j < ncols_interleaved; j++) { + s[x * ncols_interleaved + j] = sumf[j]; + } + } +} + +// Only enable these for RISC-V. +#if defined __riscv_zvfh +void ggml_gemv_q4_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 16; + const int blocklen = 1; + + assert (n % qk == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + float sumf[16]; + int sumi; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_0x16 * b_ptr = (const block_q4_0x16 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); + const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); + sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])) >> 4; + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); + } + } + } + for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; + } +} + +void ggml_gemv_q4_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK_K; + const int nb = n / qk; + const int ncols_interleaved = 16; + const int blocklen = 1; + assert (n % qk == 0); + assert (nc % ncols_interleaved == 0); + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + float sumf[16]; + float sum_minf[16]; + uint8_t scales[128]; + uint8_t mins[128]; + int sumi1; + int sumi2; + int sumi; + const block_q8_K * a_ptr = (const block_q8_K *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_Kx16 * b_ptr = (const block_q4_Kx16 *) vx + (x * nb); + for (int j = 0; j < ncols_interleaved; j++) { + sumf[j] = 0.0f; + sum_minf[j] = 0.0f; + } + for (int l = 0; l < nb; l++) { + for (int i = 0; i < 128; i++) { + scales[i] = b_ptr[l].scales[i] & 0x0F; + mins[i] = b_ptr[l].scales[i] >> 4; + } + for (int i = 0; i < 64; i++) { + scales[i] |= (b_ptr[l].scales[128 + i] & 0x03) << 4; + mins[i] |= (b_ptr[l].scales[128 + i] & 0x0C) << 2; + scales[i + 64] |= (b_ptr[l].scales[128 + i] & 0x30); + mins[i + 64] |= (b_ptr[l].scales[128 + i] & 0xC0) >> 2; + } + for (int sb = 0; sb < 8; sb++) { + uint8_t *min = &mins[sb * 16]; + for (int j = 0; j < ncols_interleaved; j++) { + sum_minf[j] += min[j] * (a_ptr[l].bsums[sb * 2] + a_ptr[l].bsums[sb * 2 + 1]) * GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d; + } + } + for (int sb = 0; sb < 8; sb += 2) { + uint8_t *scales_0 = &scales[sb * 16]; + uint8_t *scales_1 = &scales[(sb + 1) * 16]; + for (int i = 0; i < QK4_0; i++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi1 = 0; + sumi2 = 0; + sumi = 0; + const int v0 = (int8_t) (b_ptr[l].qs[sb * 256 + i * 16 + j] & 0xF); + const int v1 = (int8_t) (b_ptr[l].qs[sb * 256 + i * 16 + j] >> 4); + sumi1 = (v0 * a_ptr[l].qs[sb * 32 + i]); + sumi2 = (v1 * a_ptr[l].qs[sb * 32 + 32 + i]); + sumi1 = sumi1 * scales_0[j]; + sumi2 = sumi2 * scales_1[j]; + sumi += sumi1 + sumi2; + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d; + } + } + } + } + for (int j = 0; j < ncols_interleaved; j++) { + s[x * ncols_interleaved + j] = sumf[j] - sum_minf[j]; + } + } +} + +void ggml_gemv_iq4_nl_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 16; + const int blocklen = 1; + + assert(nr == 1); + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[16]; + int sumi; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_iq4_nlx16 * b_ptr = (const block_iq4_nlx16 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) sumf[j] = 0.0; + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; + const int v1 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; + sumi += ((v0 * a_ptr[l].qs[k * blocklen + i]) + (v1 * a_ptr[l].qs[k * blocklen + i + qk / 2])); + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); + } + } + } + for (int j = 0; j < ncols_interleaved; j++) s[x * ncols_interleaved + j] = sumf[j]; + } +} + +void ggml_gemv_q8_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 16; + const int blocklen = 1; + + assert(nr == 1); + assert(n % qk == 0); + assert(nc % ncols_interleaved == 0); + + UNUSED(bs); + UNUSED(nr); + + float sumf[16]; + int sumi; + + const block_q8_0 * a_ptr = (const block_q8_0 *) vy; + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q8_0x16 * b_ptr = (const block_q8_0x16 *) vx + (x * nb); + + for (int j = 0; j < ncols_interleaved; j++) { + sumf[j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / blocklen); k++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i]; + sumi += v0 * a_ptr[l].qs[k * blocklen + i]; + } + sumf[j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d); + } + } + } + for (int j = 0; j < ncols_interleaved; j++) { + s[x * ncols_interleaved + j] = sumf[j]; + } + } +} + +void ggml_gemv_q2_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + assert(n % QK_K == 0); + assert(nr == 1); + assert(nc % 16 == 0); + + UNUSED(bs); + UNUSED(nr); + + const int nb = n / QK_K; + const block_q2_Kx16 * x = (const block_q2_Kx16 *)vx; + const block_q8_K * y = (const block_q8_K *)vy; + + // Layout: Even-Low(0,2,4,6), Odd-Low(1,3,5,7), Even-High(8...), Odd-High(9...) + const int sb_perm[16] = { + 0, 4, 1, 5, 2, 6, 3, 7, // 0-7 + 8, 12, 9, 13, 10, 14, 11, 15 // 8-15 + }; + + for (int col_tile = 0; col_tile < nc; col_tile += 16) { + const block_q2_Kx16 * x_ptr = x + (col_tile / 16) * nb; + const block_q8_K * y_ptr = y; + + float sumf[16] = {0}; + + // Loop over K-blocks + for (int k_block = 0; k_block < nb; ++k_block) { + int32_t isum[16] = {0}; + int32_t summs[16] = {0}; + + const uint8_t * qs_rhs = x_ptr[k_block].qs; + const uint8_t * sc_rhs = x_ptr[k_block].scales; + const int8_t * qs_lhs = y_ptr[k_block].qs; + const int16_t * bs_lhs = y_ptr[k_block].bsums; + + // Iterate over sub-blocks 0..15 + for (int sb = 0; sb < 16; ++sb) { + // Correction Term + int16_t bsum = bs_lhs[sb]; + int scale_offset = sb_perm[sb] * 16; + + for (int col = 0; col < 16; ++col) { + uint8_t sc_val = sc_rhs[scale_offset + col]; + summs[col] += bsum * (sc_val >> 4); // Min is high 4 bits + } + + // Main Dot Product + // Calculate base offsets for Q2 unpacking based on SB + int byte_base; + if (sb < 8) byte_base = (sb % 2 == 0) ? 0 : 16; + else byte_base = (sb % 2 == 0) ? 32 : 48; + + int shift = ((sb / 2) % 4) * 2; + + for (int col = 0; col < 16; ++col) { + uint8_t sc_val = sc_rhs[scale_offset + col]; + int32_t d_sb = sc_val & 0xF; // Scale is low 4 bits + + // Process 16 elements (l=0..15) + for (int l = 0; l < 16; ++l) { + // Q2: Interleaved by column. Byte `l` contains 4 k-values. + int qs_idx = (byte_base + l) * 16 + col; + uint8_t q2_val = (qs_rhs[qs_idx] >> shift) & 3; + + // Q8: Linear access + int k = sb * 16 + l; + int8_t q8_val = qs_lhs[k]; + + isum[col] += q8_val * q2_val * d_sb; + } + } + } + + // Finalize K-Block + for (int col = 0; col < 16; ++col) { + float d_lhs = y_ptr[k_block].d; + float d_rhs = GGML_FP16_TO_FP32(x_ptr[k_block].d[col]); + float dm_rhs = GGML_FP16_TO_FP32(x_ptr[k_block].dmin[col]); + + float d_all = d_lhs * d_rhs; + float d_min = d_lhs * dm_rhs; + + sumf[col] += (isum[col] * d_all) - (summs[col] * d_min); + } + } + + for (int col = 0; col < 16; ++col) { + s[col_tile + col] = sumf[col]; + } + } +} +#endif + +void ggml_gemm_q4_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + const int blocklen = 4; + + assert (n % qk == 0); + assert (nr % 4 == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + { + float sumf[4][4]; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_0x4 * b_ptr = (const block_q4_0x4 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); + const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); + sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + + (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])) >> 4; + } + sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } + } +} + +void ggml_gemm_q4_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + const int blocklen = 8; + + assert (n % qk == 0); + assert (nr % 4 == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + float sumf[4][4]; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_0x4 * b_ptr = (const block_q4_0x4 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); + const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); + sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + + (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])) >> 4; + } + sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } +} + +void ggml_gemm_q4_0_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 8; + const int blocklen = 8; + + assert (n % qk == 0); + assert (nr % 4 == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + float sumf[4][8]; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_0x8 * b_ptr = (const block_q4_0x8 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); + const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); + sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + + (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])) >> 4; + } + sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } +} + +void ggml_gemm_q4_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK_K; + const int nb = n / qk; + const int ncols_interleaved = 8; + const int blocklen = 4; + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + assert (n % qk == 0); + assert (nr % 4 == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + float sumf[4][8]; + float sum_minf[4][8]; + uint32_t utmp[32]; + int sumi1; + int sumi2; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_Kx4 * a_ptr = (const block_q8_Kx4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_Kx8 * b_ptr = (const block_q4_Kx8 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumf[m][j] = 0.0; + sum_minf[m][j] = 0.0; + } + } + for (int l = 0; l < nb; l++) { + for (int sb = 0; sb < 8; sb++) { + memcpy(utmp + sb * 4, b_ptr[l].scales + sb * 12, 12); + utmp[sb * 4 + 3] = ((utmp[sb * 4 + 2] >> 4) & kmask2) | (((utmp[sb * 4 + 1] >> 6) & kmask3) << 4); + const uint32_t uaux_0 = utmp[sb * 4 + 1] & kmask1; + utmp[sb * 4 + 1] = (utmp[sb * 4 + 2] & kmask2) | (((utmp[sb * 4 + 0] >> 6) & kmask3) << 4); + utmp[sb * 4 + 2] = uaux_0; + utmp[sb * 4 + 0] &= kmask1; + } + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + uint8_t * scales_0 = (uint8_t *) utmp + (k / 8) * 32; + uint8_t * scales_1 = (uint8_t *) utmp + (k / 8) * 32 + 16; + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi1 = 0; + sumi2 = 0; + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF); + const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4); + sumi1 = (v0 * a_ptr[l].qs[(k / 8) * 256 + (k % 8) * 4 * blocklen + m * blocklen + i]); + sumi2 = (v1 * a_ptr[l].qs[(k / 8) * 256 + (k % 8) * 4 * blocklen + m * blocklen + i + 128]); + sumi1 = sumi1 * scales_0[j]; + sumi2 = sumi2 * scales_1[j]; + sumi += sumi1 + sumi2; + } + sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d[m]; + } + } + } + for (int sb = 0; sb < 8; sb++) { + uint8_t * mins = (uint8_t *) utmp + 8 + sb * 16; + for(int m = 0; m < 4; m++) { + const int16_t * bsums = a_ptr[l].bsums + (sb * 8) + (m * 4) - ((sb % 2) * 6); + for(int j = 0; j < ncols_interleaved; j++) { + sum_minf[m][j] += mins[j] * (bsums[0] + bsums[1]) * GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d[m]; + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j] - sum_minf[m][j]; + } + } + } + } +} + +void ggml_gemm_q4_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK_K; + const int nb = n / qk; + const int ncols_interleaved = 8; + const int blocklen = 8; + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + assert (n % qk == 0); + assert (nr % 4 == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(bs); + + float sumf[4][8]; + float sum_minf[4][8]; + uint32_t utmp[32]; + int sumi1; + int sumi2; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_Kx4 * a_ptr = (const block_q8_Kx4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_Kx8 * b_ptr = (const block_q4_Kx8 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumf[m][j] = 0.0; + sum_minf[m][j] = 0.0; + } + } + for (int l = 0; l < nb; l++) { + for (int sb = 0; sb < 8; sb++) { + memcpy(utmp + sb * 4, b_ptr[l].scales + sb * 12, 12); + utmp[sb * 4 + 3] = ((utmp[sb * 4 + 2] >> 4) & kmask2) | (((utmp[sb * 4 + 1] >> 6) & kmask3) << 4); + const uint32_t uaux_0 = utmp[sb * 4 + 1] & kmask1; + utmp[sb * 4 + 1] = (utmp[sb * 4 + 2] & kmask2) | (((utmp[sb * 4 + 0] >> 6) & kmask3) << 4); + utmp[sb * 4 + 2] = uaux_0; + utmp[sb * 4 + 0] &= kmask1; + } + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + uint8_t *scales_0 = (uint8_t*) utmp + (k / 4) * 32; + uint8_t *scales_1 = (uint8_t*) utmp + (k / 4) * 32 + 16; + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi1 = 0; + sumi2 = 0; + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF); + const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4); + sumi1 = (v0 * a_ptr[l].qs[(k >> 2) * 256 + (k % 4) * 4 * blocklen + m * blocklen + i]); + sumi2 = (v1 * a_ptr[l].qs[(k >> 2) * 256 + (k % 4) * 4 * blocklen + m * blocklen + i + 128]); + sumi1 = sumi1 * scales_0[j]; + sumi2 = sumi2 * scales_1[j]; + sumi += sumi1 + sumi2; + } + sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d[m]; + } + } + } + for (int sb = 0; sb < 8; sb++) { + uint8_t *mins = (uint8_t*) utmp + 8 + sb * 16; + for(int m = 0; m < 4; m++) { + const int16_t *bsums = a_ptr[l].bsums + (sb * 8) + (m * 4) - ((sb % 2) * 6); + for(int j = 0; j < ncols_interleaved; j++) { + sum_minf[m][j] += mins[j] * (bsums[0] + bsums[1]) * GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d[m]; + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j] - sum_minf[m][j]; + } + } + } + } +} + +void ggml_gemm_q2_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK_K; + const int nb = n / qk; + const int ncols_interleaved = 8; + const int blocklen = 8; + + assert (n % qk == 0); + assert (nr % 4 == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + float sumf[4][8]; + float sum_minf[4][8]; + int sumi1, sumi2, sumi3, sumi4; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_Kx4 * a_ptr = (const block_q8_Kx4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q2_Kx8 * b_ptr = (const block_q2_Kx8 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumf[m][j] = 0.0; + sum_minf[m][j] = 0.0; + } + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (4 * blocklen)); k++) { + + const uint8_t *scales_0 = b_ptr[l].scales + (k / 4) * 64 ; + const uint8_t *scales_1 = b_ptr[l].scales + (k / 4) * 64 + 16; + const uint8_t *scales_2 = b_ptr[l].scales + (k / 4) * 64 + 32; + const uint8_t *scales_3 = b_ptr[l].scales + (k / 4) * 64 + 48; + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi1 = 0; + sumi2 = 0; + sumi3 = 0; + sumi4 = 0; + sumi = 0; + int offset = ((k / 2) % 2) + j * 2; + for (int i = 0; i < blocklen; ++i){ + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 3); + const int v1 = (int8_t) ((b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 2 ) & 3); + const int v2 = (int8_t) ((b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4 ) & 3); + const int v3 = (int8_t) ((b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 6 ) & 3); + sumi1 = (v0 * a_ptr[l].qs[(k >> 2) * 512 + (k % 4) * 4 * blocklen + m * blocklen + i]); + sumi2 = (v1 * a_ptr[l].qs[(k >> 2) * 512 + (k % 4) * 4 * blocklen + m * blocklen + i + 128]); + sumi3 = (v2 * a_ptr[l].qs[(k >> 2) * 512 + (k % 4) * 4 * blocklen + m * blocklen + i + 256]); + sumi4 = (v3 * a_ptr[l].qs[(k >> 2) * 512 + (k % 4) * 4 * blocklen + m * blocklen + i + 384]); + sumi1 = sumi1 * (scales_0[offset] & 0xF); + sumi2 = sumi2 * (scales_1[offset] & 0xF); + sumi3 = sumi3 * (scales_2[offset] & 0xF); + sumi4 = sumi4 * (scales_3[offset] & 0xF); + sumi += sumi1 + sumi2 + sumi3 + sumi4; + } + sumf[m][j] += sumi * GGML_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d[m]; + } + } + } + for(int sb = 0; sb < 8; sb++) { + const uint8_t *mins = b_ptr[l].scales + sb * 16; + for(int m = 0; m < 4; m++) { + const int16_t *bsums = a_ptr[l].bsums + (sb * 8) + (m * 4) - ((sb % 2) * 6); + for(int j = 0; j < ncols_interleaved; j++) { + int mins_prod = ((mins[j * 2] >> 4) * bsums[0] + (mins[(j * 2)+ 1] >> 4) * bsums[1]); + sum_minf[m][j] += (mins_prod) * GGML_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d[m]; + } + } + } + } + + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j] - sum_minf[m][j]; + } + } + } + } +} + +void ggml_gemm_q5_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + ggml_gemm_q5_K_NxM_q8_K_generic_impl<4, 8>(n, s, bs, vx, vy, nr, nc); +} + +void ggml_gemm_q5_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + ggml_gemm_q5_K_NxM_q8_K_generic_impl<8, 8>(n, s, bs, vx, vy, nr, nc); +} + +void ggml_gemm_q6_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + ggml_gemm_q6_K_NxM_q8_K_generic_impl<4, 8>(n, s, bs, vx, vy, nr, nc); +} + +void ggml_gemm_q6_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + ggml_gemm_q6_K_NxM_q8_K_generic_impl<8, 8>(n, s, bs, vx, vy, nr, nc); +} + +void ggml_gemm_iq4_nl_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + const int blocklen = 4; + + assert (n % qk == 0); + assert (nr % 4 == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + { + float sumf[4][4]; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_iq4_nlx4 * b_ptr = (const block_iq4_nlx4 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; + const int v1 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; + sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + + (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])); + } + sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } + } +} + +void ggml_gemm_iq4_nl_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 8; + const int blocklen = 8; + + assert(n % qk == 0); + assert(nr % 4 == 0); + assert(nc % ncols_interleaved == 0); + + float sumf[4][8]; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_iq4_nlx8 * b_ptr = (const block_iq4_nlx8 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; + const int v1 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; + sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + + (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])); + } + sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } +} + +void ggml_gemm_mxfp4_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + const int blocklen = 4; + + assert(n % qk == 0); + assert(nr % 4 == 0); + assert(nc % ncols_interleaved == 0); + + float sumf[4][4]; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_mxfp4x4 * b_ptr = (const block_mxfp4x4 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; + const int v1 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; + sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + + (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])); + } + sumf[m][j] += sumi * GGML_CPU_E8M0_TO_FP32_HALF(b_ptr[l].e[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } +} + +void ggml_gemm_mxfp4_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 8; + const int blocklen = 8; + + assert(n % qk == 0); + assert(nr % 4 == 0); + assert(nc % ncols_interleaved == 0); + + float sumf[4][8]; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_mxfp4x8 * b_ptr = (const block_mxfp4x8 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; + const int v1 = kvalues_mxfp4[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; + sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + + (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])); + } + sumf[m][j] += sumi * GGML_CPU_E8M0_TO_FP32_HALF(b_ptr[l].e[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } +} + +void ggml_gemm_q8_0_4x4_q8_0_generic(int n, + float * GGML_RESTRICT s, + size_t bs, + const void * GGML_RESTRICT vx, + const void * GGML_RESTRICT vy, + int nr, + int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + const int blocklen = 4; + + assert(n % qk == 0); + assert(nr % 4 == 0); + assert(nc % ncols_interleaved == 0); + + float sumf[4][4]; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q8_0x4 * b_ptr = (const block_q8_0x4 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumf[m][j] = 0.0; + } + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / blocklen); k++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i]; + sumi += v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]; + } + sumf[m][j] += + sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } + } +} + + + +void ggml_gemm_q8_0_4x8_q8_0_generic(int n, + float * GGML_RESTRICT s, + size_t bs, + const void * GGML_RESTRICT vx, + const void * GGML_RESTRICT vy, + int nr, + int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 4; + const int blocklen = 8; + + assert(n % qk == 0); + assert(nr % 4 == 0); + assert(nc % ncols_interleaved == 0); + + float sumf[4][4]; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q8_0x4 * b_ptr = (const block_q8_0x4 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumf[m][j] = 0.0; + } + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / blocklen); k++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i]; + sumi += v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]; + } + sumf[m][j] += + sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } + } +} + +// Only enable these for RISC-V. +#if defined __riscv_zvfh +void ggml_gemm_q4_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 16; + const int blocklen = 1; + + assert (n % qk == 0); + assert (nr % 4 == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + float sumf[4][16]; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_0x16 * b_ptr = (const block_q4_0x16 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] << 4); + const int v1 = (int8_t) (b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0xF0); + sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + + (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + qk / 2 * 4])) >> 4; + } + sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } +} + +void ggml_gemm_q4_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK_K; + const int nb = n / qk; + const int ncols_interleaved = 16; + const int blocklen = 1; + + assert (n % qk == 0); + assert (nr % 4 == 0); + assert (nc % ncols_interleaved == 0); + + UNUSED(s); + UNUSED(bs); + UNUSED(vx); + UNUSED(vy); + UNUSED(nr); + UNUSED(nc); + UNUSED(nb); + UNUSED(ncols_interleaved); + UNUSED(blocklen); + + float sumf[4][16]; + float sum_minf[4][16]; + uint8_t scales[128]; + uint8_t mins[128]; + int sumi1; + int sumi2; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_Kx4 * a_ptr = (const block_q8_Kx4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q4_Kx16 * b_ptr = (const block_q4_Kx16 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumf[m][j] = 0.0; + sum_minf[m][j] = 0.0; + } + } + for (int l = 0; l < nb; l++) { + for (int i = 0; i < 128; i++) { + scales[i] = b_ptr[l].scales[i] & 0x0F; + mins[i] = b_ptr[l].scales[i] >> 4; + } + for (int i = 0; i < 64; i++) { + scales[i] |= (b_ptr[l].scales[128 + i] & 0x03) << 4; + mins[i] |= (b_ptr[l].scales[128 + i] & 0x0C) << 2; + scales[i + 64] |= (b_ptr[l].scales[128 + i] & 0x30); + mins[i + 64] |= (b_ptr[l].scales[128 + i] & 0xC0) >> 2; + } + + for (int sb = 0; sb < 8; sb++) { + uint8_t *min = &mins[sb * 16]; + for(int m = 0; m < 4; m++) { + const int16_t bsums = a_ptr[l].bsums[sb * 8 + m] + a_ptr[l].bsums[sb * 8 + m + 4]; + for(int j = 0; j < ncols_interleaved; j++) { + sum_minf[m][j] += min[j] * bsums * GGML_CPU_FP16_TO_FP32(b_ptr[l].dmin[j]) * a_ptr[l].d[m]; + } + } + } + + for (int sb = 0; sb < 8; sb += 2) { + uint8_t *scales_0 = &scales[sb * 16]; + uint8_t *scales_1 = &scales[(sb + 1) * 16]; + + for (int i = 0; i < QK4_0; i++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi1 = 0; + sumi2 = 0; + sumi = 0; + + const int v0 = (int8_t) (b_ptr[l].qs[sb * 256 + i * 16 + j] & 0xF); + const int v1 = (int8_t) (b_ptr[l].qs[sb * 256 + i * 16 + j] >> 4); + sumi1 = (v0 * a_ptr[l].qs[sb * 4 * 32 + i * 4 + m]); + sumi2 = (v1 * a_ptr[l].qs[sb * 4 * 32 + 32 * 4 + i * 4 + m]); + sumi1 = sumi1 * scales_0[j]; + sumi2 = sumi2 * scales_1[j]; + sumi += sumi1 + sumi2; + + sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * a_ptr[l].d[m]; + } + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j] - sum_minf[m][j]; + } + } + } + } +} + +void ggml_gemm_iq4_nl_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 16; + const int blocklen = 1; + + assert(n % qk == 0); + assert(nr % 4 == 0); + assert(nc % ncols_interleaved == 0); + + float sumf[4][16]; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_iq4_nlx16 * b_ptr = (const block_iq4_nlx16 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) sumf[m][j] = 0.0; + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / (2 * blocklen)); k++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] & 0x0F]; + const int v1 = kvalues_iq4nl[b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i] >> 4]; + sumi += ((v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]) + + (v1 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i + (qk / 2) * 4])); + } + sumf[m][j] += sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } +} + +void ggml_gemm_q8_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + const int qk = QK8_0; + const int nb = n / qk; + const int ncols_interleaved = 16; + const int blocklen = 1; + + assert(n % qk == 0); + assert(nr % 4 == 0); + assert(nc % ncols_interleaved == 0); + + float sumf[4][16]; + int sumi; + + for (int y = 0; y < nr / 4; y++) { + const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb); + for (int x = 0; x < nc / ncols_interleaved; x++) { + const block_q8_0x16 * b_ptr = (const block_q8_0x16 *) vx + (x * nb); + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumf[m][j] = 0.0; + } + } + for (int l = 0; l < nb; l++) { + for (int k = 0; k < (qk / blocklen); k++) { + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + sumi = 0; + for (int i = 0; i < blocklen; ++i) { + const int v0 = b_ptr[l].qs[k * ncols_interleaved * blocklen + j * blocklen + i]; + sumi += v0 * a_ptr[l].qs[k * 4 * blocklen + m * blocklen + i]; + } + sumf[m][j] += + sumi * GGML_CPU_FP16_TO_FP32(b_ptr[l].d[j]) * GGML_CPU_FP16_TO_FP32(a_ptr[l].d[m]); + } + } + } + } + for (int m = 0; m < 4; m++) { + for (int j = 0; j < ncols_interleaved; j++) { + s[(y * 4 + m) * bs + x * ncols_interleaved + j] = sumf[m][j]; + } + } + } + } +} + + +void ggml_gemm_q2_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { + assert(n % QK_K == 0); + assert(nr % 4 == 0); + assert(nc % 16 == 0); + const int nb = n / QK_K; + const block_q2_Kx16 * x = (const block_q2_Kx16 *)vx; + const block_q8_Kx4 * y = (const block_q8_Kx4 *)vy; + + const int sb_perm[16] = { + 0, 4, 1, 5, 2, 6, 3, 7, + 8, 12, 9, 13, 10, 14, 11, 15 + }; + + // Iterate Rows in tiles of 4 + for (int row_tile = 0; row_tile < nr; row_tile += 4) { + // Iterate Columns in tiles of 16 + for (int col_tile = 0; col_tile < nc; col_tile += 16) { + + const block_q2_Kx16 * x_ptr = x + (col_tile / 16) * nb; + const block_q8_Kx4 * y_ptr = y + (row_tile / 4) * nb; + + float sumf[4][16]; + memset(sumf, 0, sizeof(sumf)); + + for (int k_block = 0; k_block < nb; ++k_block) { + int32_t isum[4][16]; + int32_t summs[4][16]; + memset(isum, 0, sizeof(isum)); + memset(summs, 0, sizeof(summs)); + + const uint8_t * qs_rhs = x_ptr[k_block].qs; + const uint8_t * sc_rhs = x_ptr[k_block].scales; + const int8_t * qs_lhs = y_ptr[k_block].qs; + const int16_t * bs_lhs = y_ptr[k_block].bsums; + + for (int sb = 0; sb < 16; ++sb) { + int scale_offset = sb_perm[sb] * 16; + + int byte_base; + if (sb < 8) byte_base = (sb % 2 == 0) ? 0 : 16; + else byte_base = (sb % 2 == 0) ? 32 : 48; + int shift = ((sb / 2) % 4) * 2; + + for (int col = 0; col < 16; ++col) { + uint8_t sc_val = sc_rhs[scale_offset + col]; + int32_t d_sb = sc_val & 0xF; + int32_t m_sb = sc_val >> 4; + + // Correction Term + for (int r = 0; r < 4; ++r) { + int bsum_idx = (sb / 4) * 16 + r * 4 + (sb % 4); + summs[r][col] += bs_lhs[bsum_idx] * m_sb; + } + + // Main Dot Product + for (int l = 0; l < 16; ++l) { + int qs_idx = (byte_base + l) * 16 + col; + uint8_t q2_val = (qs_rhs[qs_idx] >> shift) & 3; + + // Calculate Q8 index for this specific k and row + int k = sb * 16 + l; + int q8_idx = (k / 4) * 16 + (k % 4); + + for (int r = 0; r < 4; ++r) { + // Add r*4 to jump to the correct row within the 4x4 chunk + int8_t q8_val = qs_lhs[q8_idx + r * 4]; + isum[r][col] += q8_val * q2_val * d_sb; + } + } + } + } + + // Finalize K-Block + for (int col = 0; col < 16; ++col) { + float d_rhs = GGML_FP16_TO_FP32(x_ptr[k_block].d[col]); + float dm_rhs = GGML_FP16_TO_FP32(x_ptr[k_block].dmin[col]); + + for (int r = 0; r < 4; ++r) { + float d_lhs = y_ptr[k_block].d[r]; + float d_all = d_lhs * d_rhs; + float d_min = d_lhs * dm_rhs; + sumf[r][col] += (isum[r][col] * d_all) - (summs[r][col] * d_min); + } + } + } + + for (int r = 0; r < 4; ++r) { + for (int col = 0; col < 16; ++col) { + s[(row_tile + r) * bs + (col_tile + col)] = sumf[r][col]; + } + } + } + } +} +#endif + +} // extern "C" + +static block_q8_0x4 make_block_q8_0x4(block_q8_0 * in, unsigned int blck_size_interleave) { + block_q8_0x4 out; + + for (int i = 0; i < 4; i++) { + out.d[i] = in[i].d; + } + + const int end = QK8_0 * 4 / blck_size_interleave; + for (int i = 0; i < end; ++i) { + int src_id = i % 4; + int src_offset = (i / 4) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], blck_size_interleave); + } + return out; +} + +static block_q4_0x4 make_block_q4_0x4(block_q4_0 * in, unsigned int blck_size_interleave) { + block_q4_0x4 out; + + for (int i = 0; i < 4; i++) { + out.d[i] = in[i].d; + } + + const int end = QK4_0 * 2 / blck_size_interleave; + + if (blck_size_interleave == 8) { + const uint64_t xor_mask = 0x8888888888888888ULL; + for (int i = 0; i < end; ++i) { + int src_id = i % 4; + int src_offset = (i / 4) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + uint64_t elems; + // Using memcpy to avoid unaligned memory accesses + memcpy(&elems, &in[src_id].qs[src_offset], sizeof(uint64_t)); + elems ^= xor_mask; + memcpy(&out.qs[dst_offset], &elems, sizeof(uint64_t)); + } + } else if (blck_size_interleave == 4) { + const uint32_t xor_mask = 0x88888888; + for (int i = 0; i < end; ++i) { + int src_id = i % 4; + int src_offset = (i / 4) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + uint32_t elems; + memcpy(&elems, &in[src_id].qs[src_offset], sizeof(uint32_t)); + elems ^= xor_mask; + memcpy(&out.qs[dst_offset], &elems, sizeof(uint32_t)); + } + } else { + GGML_ASSERT(false); + } + + return out; +} + +// interleave 8 block_q4_0s in blocks of blck_size_interleave +// returns an interleaved block_q4_0x8 +// in the interleaved block_q4_0x8, place deltas for 8 block_q4_0 blocks +// first, then interleave quants from 8 block_q4_0s in blocks of blck_size_interleave +static block_q4_0x8 make_block_q4_0x8(block_q4_0 * in, unsigned int blck_size_interleave) { + block_q4_0x8 out; + + for (int i = 0; i < 8; i++) { + out.d[i] = in[i].d; + } + + const int end = QK4_0 * 4 / blck_size_interleave; + const uint64_t xor_mask = 0x8888888888888888ULL; + + for (int i = 0; i < end; ++i) { + int src_id = i % 8; + int src_offset = (i / 8) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + uint64_t elems; + memcpy(&elems, &in[src_id].qs[src_offset], sizeof(uint64_t)); + elems ^= xor_mask; + memcpy(&out.qs[dst_offset], &elems, sizeof(uint64_t)); + } + + return out; +} + +static block_q4_0x16 make_block_q4_0x16(block_q4_0 * in, unsigned int blck_size_interleave) { + block_q4_0x16 out; + + for (int i = 0; i < 16; i++) { + out.d[i] = in[i].d; + } + + const int end = QK4_0 * 8 / blck_size_interleave; + + if (blck_size_interleave == 1) { + const uint8_t xor_mask = 0x88; + for (int i = 0; i < end; ++i) { + int src_id = i % 16; + int src_offset = i / 16; + int dst_offset = i; + + out.qs[dst_offset] = in[src_id].qs[src_offset] ^ xor_mask; + } + } else { + GGML_ASSERT(false); + } + + return out; +} + +static block_q4_Kx8 make_block_q4_Kx8(block_q4_K * in, unsigned int blck_size_interleave) { + block_q4_Kx8 out; + //Delta(scale) and dmin values of the eight Q4_K structures are copied onto the output interleaved structure + for (int i = 0; i < 8; i++) { + out.d[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d; + } + + for (int i = 0; i < 8; i++) { + out.dmin[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.dmin; + } + + const int end = QK_K * 4 / blck_size_interleave; + + // Interleave Q4_K quants by taking 8 bytes at a time + for (int i = 0; i < end; ++i) { + int src_id = i % 8; + int src_offset = (i / 8) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + // buffer large enough for the max interleave block size (8 bytes) + uint64_t elems; + memcpy(&elems, &in[src_id].qs[src_offset], blck_size_interleave); + memcpy(&out.qs[dst_offset], &elems, blck_size_interleave); + } + + // The below logic is designed so as to unpack and rearrange scales and mins values in Q4_K + // Currently the Q4_K structure has 8 scales and 8 mins packed in 12 bytes ( 6 bits for each value) + // The output Q4_Kx8 structure has 96 bytes + // Every 12 byte is packed such that it contains scales and mins for corresponding sub blocks from Q4_K structure + // For eg - First 12 bytes contains 8 scales and 8 mins - each of first sub block from different Q4_K structures + uint8_t s[8], m[8]; + + for (int i = 0; i < 4; i++) { + for (int j = 0; j < 8; j++) { + s[j] = in[j].scales[i] & 63; + m[j] = in[j].scales[i + 4] & 63; + } + + out.scales[i * 12] = (s[0] & 63) + ((s[4] & 48) << 2); + out.scales[i * 12 + 1] = (s[1] & 63) + ((s[5] & 48) << 2); + out.scales[i * 12 + 2] = (s[2] & 63) + ((s[6] & 48) << 2); + out.scales[i * 12 + 3] = (s[3] & 63) + ((s[7] & 48) << 2); + out.scales[i * 12 + 4] = (m[0] & 63) + ((m[4] & 48) << 2); + out.scales[i * 12 + 5] = (m[1] & 63) + ((m[5] & 48) << 2); + out.scales[i * 12 + 6] = (m[2] & 63) + ((m[6] & 48) << 2); + out.scales[i * 12 + 7] = (m[3] & 63) + ((m[7] & 48) << 2); + out.scales[i * 12 + 8] = (s[4] & 15) + ((m[4] & 15) << 4); + out.scales[i * 12 + 9] = (s[5] & 15) + ((m[5] & 15) << 4); + out.scales[i * 12 + 10] = (s[6] & 15) + ((m[6] & 15) << 4); + out.scales[i * 12 + 11] = (s[7] & 15) + ((m[7] & 15) << 4); + + } + + for (int i = 0; i < 4; i++) { + for (int j = 0; j < 8; j++) { + s[j] = ((in[j].scales[i] & 192) >> 2) | (in[j].scales[i+8] & 15); + m[j] = ((in[j].scales[i + 4] & 192) >> 2) | ((in[j].scales[i+8] & 240) >> 4); + } + + out.scales[i * 12 + 48] = (s[0] & 63) + ((s[4] & 48) << 2); + out.scales[i * 12 + 49] = (s[1] & 63) + ((s[5] & 48) << 2); + out.scales[i * 12 + 50] = (s[2] & 63) + ((s[6] & 48) << 2); + out.scales[i * 12 + 51] = (s[3] & 63) + ((s[7] & 48) << 2); + out.scales[i * 12 + 52] = (m[0] & 63) + ((m[4] & 48) << 2); + out.scales[i * 12 + 53] = (m[1] & 63) + ((m[5] & 48) << 2); + out.scales[i * 12 + 54] = (m[2] & 63) + ((m[6] & 48) << 2); + out.scales[i * 12 + 55] = (m[3] & 63) + ((m[7] & 48) << 2); + out.scales[i * 12 + 56] = (s[4] & 15) + ((m[4] & 15) << 4); + out.scales[i * 12 + 57] = (s[5] & 15) + ((m[5] & 15) << 4); + out.scales[i * 12 + 58] = (s[6] & 15) + ((m[6] & 15) << 4); + out.scales[i * 12 + 59] = (s[7] & 15) + ((m[7] & 15) << 4); + + } + + return out; +} + +static block_q4_Kx16 make_block_q4_Kx16(block_q4_K * in, unsigned int blck_size_interleave) { + block_q4_Kx16 out; + //Delta(scale) and dmin values of the 16 Q4_K structures are copied onto the output interleaved structure + for (int i = 0; i < 16; i++) { + out.d[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d; + } + + for (int i = 0; i < 16; i++) { + out.dmin[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.dmin; + } + + const int end = QK_K * 8 / blck_size_interleave; + + if (blck_size_interleave == 1) { + for (int i = 0; i < end; ++i) { + int src_id = i % 16; + int src_offset = i / 16; + int dst_offset = i; + + out.qs[dst_offset] = in[src_id].qs[src_offset]; + } + + // RVV repacking. + // + // Extract sums and mins for all 8 sub-blocks for each block of Q4_K. + uint8_t s[128], m[128]; + for (int i = 0; i < 4; i++) { + for (int j = 0; j < 16; j++) { + s[i * 16 + j] = in[j].scales[i] & 63; + m[i * 16 + j] = in[j].scales[i + 4] & 63; + } + } + for (int i = 0; i < 4; i++) { + for (int j = 0; j < 16; j++) { + s[64 + i * 16 + j] = ((in[j].scales[i] & 192) >> 2) | (in[j].scales[i+8] & 15); + m[64 + i * 16 + j] = ((in[j].scales[i + 4] & 192) >> 2) | ((in[j].scales[i+8] & 240) >> 4); + } + } + + for (int i = 0; i < 128; i++) { + out.scales[i] = (s[i] & 15) | ((m[i] & 15) << 4); + } + for (int i = 0; i < 64; i++) { + out.scales[128 + i] = ((s[i] & 48) >> 4) | ((m[i] & 48) >> 2) | (s[64 + i] & 48) | ((m[64 + i] & 48) << 2); + } + } else { + GGML_ASSERT(false); + } + + return out; +} + +static block_q2_Kx8 make_block_q2_Kx8(block_q2_K * in, unsigned int blck_size_interleave) { + block_q2_Kx8 out; + + // Delta(scale) and dmin values of the eight Q2_K structures are copied onto the output interleaved structure + for (int i = 0; i < 8; i++) { + out.d[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d; + } + + for (int i = 0; i < 8; i++) { + out.dmin[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.dmin; + } + + const int end = QK_K * 2 / blck_size_interleave; + + // Interleave Q2_K quants by taking 8 bytes at a time + for (int i = 0; i < end; ++i) { + int src_id = i % 8; + int src_offset = (i / 8) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + uint64_t elems; + memcpy(&elems, &in[src_id].qs[src_offset], sizeof(uint64_t)); + memcpy(&out.qs[dst_offset], &elems, sizeof(uint64_t)); + } + + // The below logic is designed so as to unpack and rearrange scales and mins values in Q2_K + // Currently the Q2_K structure has 16 scales and 16 mins packed in 16 bytes ( 4 bits for each value) + // The output Q2_Kx8 structure has 128 bytes for storing scales and mins + // Every 16 byte is packed such that it contains scales and mins for corresponding sub blocks from Q2_K structure + // For eg - First 16 bytes contains 16 scales and 16 mins - each of first and second sub blocks from different Q2_K structures + + for (int i = 0; i < 128; i++) { + // Index for selecting which q2k super block + int src1 = (i % 16) / 2; + // Index for selecting scale + int src2 = ((i / 16) * 2) + (i % 2); + + out.scales[i] = in[src1].scales[src2]; + } + return out; +} + +static block_q5_Kx8 make_block_q5_Kx8(block_q5_K * in, unsigned int blck_size_interleave) { + block_q5_Kx8 out; + //Delta(scale) and dmin values of the eight Q5_K structures are copied onto the output interleaved structure + for (int i = 0; i < 8; i++) { + out.d[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d; + } + + for (int i = 0; i < 8; i++) { + out.dmin[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.dmin; + } + + const int end = QK_K * 4 / blck_size_interleave; + + // Interleave Q5_K quants by taking blck_size_interleave bytes at a time + for (int i = 0; i < end; ++i) { + int src_id = i % 8; + int src_offset = (i / 8) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], blck_size_interleave); + } + + // Repeat for high bits with the same chunk size, since + // the high bits are interleaved in Q5_K and the index is + // qh_idx = (qs_idx % 32); + // qh_val = qh[qh_idx] >> (qs_idx / 32); + for (int i = 0; i < end / 4; ++i) { + int src_id = i % 8; + int src_offset = (i / 8) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + memcpy(&out.qh[dst_offset], &in[src_id].qh[src_offset], blck_size_interleave); + } + + // The below logic is copied over from Q4_K + // The point is to unpack all the scales and mins for each sub block every time we load 12 bytes. + // Currently the Q5_K structure has 8 scales and 8 mins packed in 12 bytes ( 6 bits for each value) + // The output Q5_Kx8 structure has 96 bytes + // Every 12 byte is packed such that it contains scales and mins for corresponding sub blocks from Q5_K structure + // For eg - First 12 bytes contains 8 scales and 8 mins - each of first sub block from different Q5_K structures + uint8_t s[8], m[8]; + + for (int i = 0; i < 4; i++) { + for (int j = 0; j < 8; j++) { + s[j] = in[j].scales[i] & 63; + m[j] = in[j].scales[i + 4] & 63; + } + + out.scales[i * 12] = (s[0] & 63) + ((s[4] & 48) << 2); + out.scales[i * 12 + 1] = (s[1] & 63) + ((s[5] & 48) << 2); + out.scales[i * 12 + 2] = (s[2] & 63) + ((s[6] & 48) << 2); + out.scales[i * 12 + 3] = (s[3] & 63) + ((s[7] & 48) << 2); + out.scales[i * 12 + 4] = (m[0] & 63) + ((m[4] & 48) << 2); + out.scales[i * 12 + 5] = (m[1] & 63) + ((m[5] & 48) << 2); + out.scales[i * 12 + 6] = (m[2] & 63) + ((m[6] & 48) << 2); + out.scales[i * 12 + 7] = (m[3] & 63) + ((m[7] & 48) << 2); + out.scales[i * 12 + 8] = (s[4] & 15) + ((m[4] & 15) << 4); + out.scales[i * 12 + 9] = (s[5] & 15) + ((m[5] & 15) << 4); + out.scales[i * 12 + 10] = (s[6] & 15) + ((m[6] & 15) << 4); + out.scales[i * 12 + 11] = (s[7] & 15) + ((m[7] & 15) << 4); + } + + for (int i = 0; i < 4; i++) { + for (int j = 0; j < 8; j++) { + s[j] = ((in[j].scales[i] & 192) >> 2) | (in[j].scales[i + 8] & 15); + m[j] = ((in[j].scales[i + 4] & 192) >> 2) | ((in[j].scales[i + 8] & 240) >> 4); + } + + out.scales[i * 12 + 48] = (s[0] & 63) + ((s[4] & 48) << 2); + out.scales[i * 12 + 49] = (s[1] & 63) + ((s[5] & 48) << 2); + out.scales[i * 12 + 50] = (s[2] & 63) + ((s[6] & 48) << 2); + out.scales[i * 12 + 51] = (s[3] & 63) + ((s[7] & 48) << 2); + out.scales[i * 12 + 52] = (m[0] & 63) + ((m[4] & 48) << 2); + out.scales[i * 12 + 53] = (m[1] & 63) + ((m[5] & 48) << 2); + out.scales[i * 12 + 54] = (m[2] & 63) + ((m[6] & 48) << 2); + out.scales[i * 12 + 55] = (m[3] & 63) + ((m[7] & 48) << 2); + out.scales[i * 12 + 56] = (s[4] & 15) + ((m[4] & 15) << 4); + out.scales[i * 12 + 57] = (s[5] & 15) + ((m[5] & 15) << 4); + out.scales[i * 12 + 58] = (s[6] & 15) + ((m[6] & 15) << 4); + out.scales[i * 12 + 59] = (s[7] & 15) + ((m[7] & 15) << 4); + } + + return out; +} + +static block_q6_Kx8 make_block_q6_Kx8(block_q6_K * in, unsigned int blck_size_interleave) { + block_q6_Kx8 out; + constexpr int n_blocks = 8; // Kx8 + for (int i = 0; i < n_blocks; i++) { + out.d[i] = in[i].d; + } + + const int end_ls = QK_K * 4 / blck_size_interleave; + // Interleave Q6_K quants by taking blck_size_interleave bytes at a time + for (int i = 0; i < end_ls; ++i) { + int src_id = i % n_blocks; + int src_offset = (i / n_blocks) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + uint64_t elem_ls; + memcpy(&elem_ls, &in[src_id].ql[src_offset], blck_size_interleave); + memcpy(&out.ql[dst_offset], &elem_ls, blck_size_interleave); + } + + // Interleave high bits using same chunk size as low bits + const int end_hs = end_ls / 2; + for (int i = 0; i < end_hs; ++i) { + int src_id = i % n_blocks; + int src_offset = (i / n_blocks) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + uint64_t elem_hs; + memcpy(&elem_hs, &in[src_id].qh[src_offset], blck_size_interleave); + memcpy(&out.qh[dst_offset], &elem_hs, blck_size_interleave); + } + + // The below logic is designed so as to unpack and rearrange scales in Q6_K + // The output Q6_Kx8 structure interleaves the 8 bit scales in the same fashion as the quants + // Q6_K structure has an 8-bit scale per 16 elements -> 16 scales + // scales: [0 bl0 0 bl1 ... 0 bl7][1 bl0 ... 1 bl7] ... [15 bl0 ... 15 bl7] (bl = block) + constexpr int n_scales = QK_K / 16; + + for (int i = 0; i < n_blocks; i++) { + for (int j = 0; j < n_scales; j++) { + out.scales[j * n_blocks + i] = in[i].scales[j]; + } + } + + return out; +} + +static block_q2_Kx16 make_block_q2_Kx16(const block_q2_K * in, unsigned int blck_size_interleave) { + block_q2_Kx16 out; + constexpr int N_COLS = 16; + + // 1. Copy Super-Scales (d) and Super-Mins (dmin) + for (int i = 0; i < N_COLS; i++) { + out.d[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.d; + out.dmin[i] = in[i].GGML_COMMON_AGGR_U.GGML_COMMON_AGGR_S.dmin; + } + + // 2. Interleave Q2_K Data + const int bytes_per_col = 64; + const int total_bytes = N_COLS * bytes_per_col; + const int end = total_bytes / blck_size_interleave; + + for (int i = 0; i < end; ++i) { + int src_col_id = i % N_COLS; + int src_offset = (i / N_COLS) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + memcpy(&out.qs[dst_offset], &in[src_col_id].qs[src_offset], blck_size_interleave); + } + + // 3. Repack Scales into the Optimized "Sequential-Parallel" Layout + int out_idx = 0; + + // Arrays define the sub-block order for each group + const int even_low_sbs[] = {0, 2, 4, 6}; + const int odd_low_sbs[] = {1, 3, 5, 7}; + const int even_high_sbs[] = {8, 10, 12, 14}; + const int odd_high_sbs[] = {9, 11, 13, 15}; + + // Pack Group 1: Even-Low + for (int sb : even_low_sbs) { + for (int col = 0; col < N_COLS; col++) { + out.scales[out_idx++] = in[col].scales[sb]; + } + } + + // Pack Group 2: Odd-Low + for (int sb : odd_low_sbs) { + for (int col = 0; col < N_COLS; col++) { + out.scales[out_idx++] = in[col].scales[sb]; + } + } + + // Pack Group 3: Even-High + for (int sb : even_high_sbs) { + for (int col = 0; col < N_COLS; col++) { + out.scales[out_idx++] = in[col].scales[sb]; + } + } + + // Pack Group 4: Odd-High + for (int sb : odd_high_sbs) { + for (int col = 0; col < N_COLS; col++) { + out.scales[out_idx++] = in[col].scales[sb]; + } + } + + return out; +} + +static int repack_q4_0_to_q4_0_4_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_Q4_0); + GGML_ASSERT(interleave_block == 4 || interleave_block == 8); + constexpr int nrows_interleaved = 4; + + block_q4_0x4 * dst = (block_q4_0x4 *)t->data; + const block_q4_0 * src = (const block_q4_0 *)data; + block_q4_0 dst_tmp[4]; + int nrow = ggml_nrows(t); + int nblocks = t->ne[0] / QK4_0; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_0)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_q4_0x4(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; + + GGML_UNUSED(data_size); +} + +static int repack_q4_K_to_q4_K_8_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_Q4_K); + GGML_ASSERT(interleave_block == 8 || interleave_block == 4); + constexpr int nrows_interleaved = 8; + + block_q4_Kx8 * dst = (block_q4_Kx8*)t->data; + const block_q4_K * src = (const block_q4_K*) data; + block_q4_K dst_tmp[8]; + int nrow = ggml_nrows(t); + int nblocks = t->ne[0] / QK_K; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_K)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++ ) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_q4_Kx8(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; + + GGML_UNUSED(data_size); +} + +static int repack_q4_K_to_q4_K_16_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_Q4_K); + constexpr int nrows_interleaved = 16; + + block_q4_Kx16 * dst = (block_q4_Kx16*)t->data; + const block_q4_K * src = (const block_q4_K*) data; + block_q4_K dst_tmp[16]; + int nrow = ggml_nrows(t); + int nblocks = t->ne[0] / QK_K; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_K)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++ ) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_q4_Kx16(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; + + GGML_UNUSED(data_size); +} + +static int repack_q2_K_to_q2_K_8_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_Q2_K); + GGML_ASSERT(interleave_block == 8); + constexpr int nrows_interleaved = 8; + + block_q2_Kx8 * dst = (block_q2_Kx8*)t->data; + const block_q2_K * src = (const block_q2_K*) data; + block_q2_K dst_tmp[8]; + int nrow = ggml_nrows(t); + int nblocks = t->ne[0] / QK_K; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q2_K)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_q2_Kx8(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; + + GGML_UNUSED(data_size); +} + +static int repack_q2_K_to_q2_K_16_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_Q2_K); + constexpr int nrows_interleaved = 16; + + block_q2_Kx16 * dst = (block_q2_Kx16*)t->data; + const block_q2_K * src = (const block_q2_K*) data; + + block_q2_K dst_tmp[nrows_interleaved]; + + int nrow = ggml_nrows(t); + int nblocks = t->ne[0] / QK_K; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q2_K)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + // This loop gathers 16 separate blocks (one from each column) + // that correspond to the same K-dimension chunk. + for (int i = 0; i < nrows_interleaved; i++ ) { + dst_tmp[i] = src[x + i * nblocks]; + } + + *dst++ = make_block_q2_Kx16(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; + + GGML_UNUSED(data_size); +} + +static int repack_q4_0_to_q4_0_16_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_Q4_0); + constexpr int nrows_interleaved = 16; + + block_q4_0x16 * dst = (block_q4_0x16*)t->data; + const block_q4_0 * src = (const block_q4_0*) data; + block_q4_0 dst_tmp[16]; + int nrow = ggml_nrows(t); + int nblocks = t->ne[0] / QK4_0; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_0)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++ ) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_q4_0x16(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; + + GGML_UNUSED(data_size); +} + +static int repack_q5_K_to_q5_K_8_bl(struct ggml_tensor * t, + int interleave_block, + const void * GGML_RESTRICT data, + size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_Q5_K); + GGML_ASSERT(interleave_block == 4 || interleave_block == 8); + constexpr int nrows_interleaved = 8; + + block_q5_Kx8 * dst = (block_q5_Kx8 *) t->data; + const block_q5_K * src = (const block_q5_K *) data; + block_q5_K dst_tmp[8]; + int nrow = ggml_nrows(t); + int nblocks = t->ne[0] / QK_K; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q5_K)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_q5_Kx8(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; +} + +static int repack_q6_K_to_q6_K_8_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_Q6_K); + GGML_ASSERT(interleave_block == 4 || interleave_block == 8); + constexpr int nrows_interleaved = 8; + + block_q6_Kx8 * dst = (block_q6_Kx8 *)t->data; + const block_q6_K * src = (const block_q6_K *) data; + block_q6_K dst_tmp[8]; + int nrow = ggml_nrows(t); + int nblocks = t->ne[0] / QK_K; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q6_K)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_q6_Kx8(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; +} + +static int repack_q4_0_to_q4_0_8_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_Q4_0); + GGML_ASSERT(interleave_block == 8); + constexpr int nrows_interleaved = 8; + + block_q4_0x8 * dst = (block_q4_0x8*)t->data; + const block_q4_0 * src = (const block_q4_0*) data; + block_q4_0 dst_tmp[8]; + int nrow = ggml_nrows(t); + int nblocks = t->ne[0] / QK4_0; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q4_0)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++ ) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_q4_0x8(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; + + GGML_UNUSED(data_size); +} + +static int repack_q8_0_to_q8_0_4_bl(struct ggml_tensor * t, + int interleave_block, + const void * GGML_RESTRICT data, + size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_Q8_0); + GGML_ASSERT(interleave_block == 4 || interleave_block == 8); + constexpr int nrows_interleaved = 4; + + block_q8_0x4 * dst = (block_q8_0x4 *) t->data; + const block_q8_0 * src = (const block_q8_0 *) data; + block_q8_0 dst_tmp[4]; + int nrow = ggml_nrows(t); + int nblocks = t->ne[0] / QK8_0; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q8_0)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_q8_0x4(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; +} + +static block_q8_0x16 make_block_q8_0x16(block_q8_0 * in, unsigned int blck_size_interleave) { + block_q8_0x16 out; + + for (int i = 0; i < 16; i++) { + out.d[i] = in[i].d; + } + + const int end = QK8_0 * 16 / blck_size_interleave; + + if (blck_size_interleave == 1) { + for (int i = 0; i < end; ++i) { + int src_id = i % 16; + int src_offset = i / 16; + int dst_offset = i; + out.qs[dst_offset] = in[src_id].qs[src_offset]; + } + } else { + GGML_ASSERT(false); + } + + return out; +} + +static int repack_q8_0_to_q8_0_16_bl(struct ggml_tensor * t, + int interleave_block, + const void * GGML_RESTRICT data, + size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_Q8_0); + constexpr int nrows_interleaved = 16; + + block_q8_0x16 * dst = (block_q8_0x16 *) t->data; + const block_q8_0 * src = (const block_q8_0 *) data; + block_q8_0 dst_tmp[16]; + int nrow = ggml_nrows(t); + int nblocks = t->ne[0] / QK8_0; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q8_0)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_q8_0x16(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; +} + +static block_iq4_nlx4 make_block_iq4_nlx4(block_iq4_nl * in, unsigned int blck_size_interleave) { + block_iq4_nlx4 out; + + for (int i = 0; i < 4; i++) { + out.d[i] = in[i].d; + } + + const int end = QK4_NL * 2 / blck_size_interleave; + + // TODO: this branch seems wrong + //if (blck_size_interleave == 8) { + // for (int i = 0; i < end; ++i) { + // int src_id = i % 4; + // int src_offset = (i / 4) * blck_size_interleave; + // int dst_offset = i * blck_size_interleave; + + // // Using memcpy to avoid unaligned memory accesses + // memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], sizeof(uint64_t)); + // } + //} else + if (blck_size_interleave == 4) { + for (int i = 0; i < end; ++i) { + int src_id = i % 4; + int src_offset = (i / 4) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], sizeof(uint32_t)); + } + } else { + GGML_ASSERT(false); + } + + return out; +} + +static int repack_iq4_nl_to_iq4_nl_4_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_IQ4_NL); + GGML_ASSERT(interleave_block == 4); + + const block_iq4_nl * src = (const block_iq4_nl *)data; + block_iq4_nlx4 * dst = ( block_iq4_nlx4 *)t->data; + + block_iq4_nl dst_tmp[4]; + + int nrow = ggml_nrows(t); + int nrows_interleaved = 4; + int nblocks = t->ne[0] / QK4_NL; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_iq4_nl)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_iq4_nlx4(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; + + GGML_UNUSED(data_size); +} + +static block_iq4_nlx8 make_block_iq4_nlx8(block_iq4_nl * in, unsigned int blck_size_interleave) { + block_iq4_nlx8 out; + + for (int i = 0; i < 8; i++) { + out.d[i] = in[i].d; + } + + const int end = QK4_NL * 4 / blck_size_interleave; + + if (blck_size_interleave == 8) { + for (int i = 0; i < end; ++i) { + int src_id = i % 8; + int src_offset = (i / 8) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], sizeof(uint64_t)); + } + } else { + GGML_ASSERT(false); + } + + return out; +} + +static int repack_iq4_nl_to_iq4_nl_8_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_IQ4_NL); + GGML_ASSERT(interleave_block == 8); + + const block_iq4_nl * src = (const block_iq4_nl *)data; + block_iq4_nlx8 * dst = ( block_iq4_nlx8 *)t->data; + + block_iq4_nl dst_tmp[8]; + + int nrow = ggml_nrows(t); + int nrows_interleaved = 8; + int nblocks = t->ne[0] / QK4_NL; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_iq4_nl)); + + if (t->ne[1] % nrows_interleaved != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_iq4_nlx8(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; + + GGML_UNUSED(data_size); +} + +static block_iq4_nlx16 make_block_iq4_nlx16(block_iq4_nl * in, unsigned int blck_size_interleave) { + block_iq4_nlx16 out; + + for (int i = 0; i < 16; i++) { + out.d[i] = in[i].d; + } + + const int end = QK4_NL * 8 / blck_size_interleave; + + if (blck_size_interleave == 1) { + for (int i = 0; i < end; ++i) { + int src_id = i % 16; + int src_offset = i / 16; + int dst_offset = i; + + out.qs[dst_offset] = in[src_id].qs[src_offset]; + } + } else { + GGML_ASSERT(false); + } + + return out; +} + +static int repack_iq4_nl_to_iq4_nl_16_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_IQ4_NL); + GGML_ASSERT(interleave_block == 1); + + const block_iq4_nl * src = (const block_iq4_nl *)data; + block_iq4_nlx16 * dst = ( block_iq4_nlx16 *)t->data; + + block_iq4_nl dst_tmp[16]; + + int nrow = ggml_nrows(t); + int nrows_interleaved = 16; + int nblocks = t->ne[0] / QK4_NL; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_iq4_nl)); + + if (t->ne[1] % nrows_interleaved != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_iq4_nlx16(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; + + GGML_UNUSED(data_size); +} + +static block_mxfp4x4 make_block_mxfp4x4(block_mxfp4 * in, unsigned int blck_size_interleave) { + block_mxfp4x4 out; + + for (int i = 0; i < 4; i++) { + out.e[i] = in[i].e; + } + + const int end = QK_MXFP4 * 2 / blck_size_interleave; + + if (blck_size_interleave == 4) { + for (int i = 0; i < end; ++i) { + int src_id = i % 4; + int src_offset = (i / 4) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], sizeof(uint32_t)); + } + } else { + GGML_ASSERT(false); + } + + return out; +} + +static int repack_mxfp4_to_mxfp4_4_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_MXFP4); + GGML_ASSERT(interleave_block == 4); + + const block_mxfp4 * src = (const block_mxfp4 *)data; + block_mxfp4x4 * dst = ( block_mxfp4x4 *)t->data; + + block_mxfp4 dst_tmp[4]; + + int nrow = ggml_nrows(t); + int nrows_interleaved = 4; + int nblocks = t->ne[0] / QK_MXFP4; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_mxfp4)); + + if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % 8 != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_mxfp4x4(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; + + GGML_UNUSED(data_size); +} + +static block_mxfp4x8 make_block_mxfp4x8(block_mxfp4 * in, unsigned int blck_size_interleave) { + block_mxfp4x8 out; + + for (int i = 0; i < 8; i++) { + out.e[i] = in[i].e; + } + + const int end = QK_MXFP4 * 4 / blck_size_interleave; + + if (blck_size_interleave == 8) { + for (int i = 0; i < end; ++i) { + int src_id = i % 8; + int src_offset = (i / 8) * blck_size_interleave; + int dst_offset = i * blck_size_interleave; + + memcpy(&out.qs[dst_offset], &in[src_id].qs[src_offset], sizeof(uint64_t)); + } + } else { + GGML_ASSERT(false); + } + + return out; +} + +static int repack_mxfp4_to_mxfp4_8_bl(struct ggml_tensor * t, int interleave_block, const void * GGML_RESTRICT data, size_t data_size) { + GGML_ASSERT(t->type == GGML_TYPE_MXFP4); + GGML_ASSERT(interleave_block == 8); + + const block_mxfp4 * src = (const block_mxfp4 *)data; + block_mxfp4x8 * dst = ( block_mxfp4x8 *)t->data; + + block_mxfp4 dst_tmp[8]; + + int nrow = ggml_nrows(t); + int nrows_interleaved = 8; + int nblocks = t->ne[0] / QK_MXFP4; + + GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_mxfp4)); + + if (t->ne[1] % nrows_interleaved != 0) { + return -1; + } + + for (int b = 0; b < nrow; b += nrows_interleaved) { + for (int64_t x = 0; x < nblocks; x++) { + for (int i = 0; i < nrows_interleaved; i++) { + dst_tmp[i] = src[x + i * nblocks]; + } + *dst++ = make_block_mxfp4x8(dst_tmp, interleave_block); + } + src += nrows_interleaved * nblocks; + } + return 0; + + GGML_UNUSED(data_size); +} + +namespace ggml::cpu::repack { +// repack +template +int repack(struct ggml_tensor *, const void *, size_t); + +// TODO: generalise. +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q4_0_to_q4_0_4_bl(t, 4, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q4_0_to_q4_0_4_bl(t, 8, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q4_0_to_q4_0_8_bl(t, 8, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q4_K_to_q4_K_8_bl(t, 8, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q4_K_to_q4_K_8_bl(t, 4, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q2_K_to_q2_K_8_bl(t, 8, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q5_K_to_q5_K_8_bl(t, 4, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q5_K_to_q5_K_8_bl(t, 8, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q6_K_to_q6_K_8_bl(t, 4, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q6_K_to_q6_K_8_bl(t, 8, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_iq4_nl_to_iq4_nl_4_bl(t, 4, data, data_size); +} + +// TODO: needs to be revisited +//template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { +// return repack_iq4_nl_to_iq4_nl_4_bl(t, 8, data, data_size); +//} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_iq4_nl_to_iq4_nl_8_bl(t, 8, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_mxfp4_to_mxfp4_4_bl(t, 4, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_mxfp4_to_mxfp4_8_bl(t, 8, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q8_0_to_q8_0_4_bl(t, 4, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q8_0_to_q8_0_4_bl(t, 8, data, data_size); +} + +#if defined __riscv_zvfh +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q4_0_to_q4_0_16_bl(t, 1, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q4_K_to_q4_K_16_bl(t, 1, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_iq4_nl_to_iq4_nl_16_bl(t, 1, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q8_0_to_q8_0_16_bl(t, 1, data, data_size); +} + +template <> int repack(struct ggml_tensor * t, const void * data, size_t data_size) { + return repack_q2_K_to_q2_K_16_bl(t, 1, data, data_size); +} +#endif + +// gemv +template +void gemv(int, float *, size_t, const void *, const void *, int, int); + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q4_0_4x4_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q4_0_4x8_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q4_0_8x8_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> +void gemv(int n, + float * s, + size_t bs, + const void * vx, + const void * vy, + int nr, + int nc) { + ggml_gemv_q2_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q4_K_8x4_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q4_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q5_K_8x4_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q5_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q6_K_8x4_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q6_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_iq4_nl_4x4_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_iq4_nl_8x8_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_mxfp4_4x4_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_mxfp4_8x8_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q8_0_4x4_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q8_0_4x8_q8_0(n, s, bs, vx, vy, nr, nc); +} + +#if defined __riscv_zvfh +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q4_0_16x1_q8_0(n, s, bs, vx, vy, nr, nc); +} -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "ggml-cpu-impl.h" -#include "simd-mappings.h" -#include "traits.h" +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q4_K_16x1_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_iq4_nl_16x1_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q8_0_16x1_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemv(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemv_q2_K_16x1_q8_K(n, s, bs, vx, vy, nr, nc); +} +#endif + +// gemm +template +void gemm(int, float *, size_t, const void *, const void *, int, int); + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q4_0_4x4_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q4_0_4x8_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> +void gemm(int n, + float * s, + size_t bs, + const void * vx, + const void * vy, + int nr, + int nc) { + ggml_gemm_q4_0_8x8_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q2_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q4_K_8x4_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q4_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q5_K_8x4_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q5_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q6_K_8x4_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q6_K_8x8_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_iq4_nl_4x4_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_iq4_nl_8x8_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_mxfp4_4x4_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_mxfp4_8x8_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q8_0_4x4_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q8_0_4x8_q8_0(n, s, bs, vx, vy, nr, nc); +} + +#if defined __riscv_zvfh +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q4_0_16x1_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q4_K_16x1_q8_K(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_iq4_nl_16x1_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q8_0_16x1_q8_0(n, s, bs, vx, vy, nr, nc); +} + +template <> void gemm(int n, float * s, size_t bs, const void * vx, const void * vy, int nr, int nc) { + ggml_gemm_q2_K_16x1_q8_K(n, s, bs, vx, vy, nr, nc); +} +#endif + +class tensor_traits_base : public ggml::cpu::tensor_traits { + public: + virtual int repack(struct ggml_tensor * t, const void * data, size_t data_size) = 0; +}; + +template class tensor_traits : public tensor_traits_base { + + bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override { + // not realy a GGML_TYPE_Q8_0 but same size. + switch (op->op) { + case GGML_OP_MUL_MAT: + { + size = ggml_row_size(PARAM_TYPE, ggml_nelements(op->src[1])); + return true; + } + case GGML_OP_MUL_MAT_ID: + { + size = ggml_row_size(PARAM_TYPE, ggml_nelements(op->src[1])); + size = GGML_PAD(size, sizeof(int64_t)); // + padding for next block. + + const int64_t ne02 = op->src[0]->ne[2]; // n_as, n_expert + const int64_t ne12 = op->src[1]->ne[2]; // n_tokens + + const size_t sizeof_mmid_row_mapping = sizeof(int64_t); + + size += sizeof_mmid_row_mapping*ne02*(ne12 + 1); + + return true; + } + default: + // GGML_ABORT("fatal error"); + break; + } + return false; + } + + bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) override { + switch (op->op) { + case GGML_OP_MUL_MAT: + forward_mul_mat(params, op); + return true; + case GGML_OP_MUL_MAT_ID: + forward_mul_mat_id(params, op); + return true; + default: + // GGML_ABORT("fatal error"); + break; + } + return false; + } + + void forward_mul_mat_one_chunk(ggml_compute_params * params, + ggml_tensor * op, + int64_t src0_start, + int64_t src0_end, + int64_t src1_start, + int64_t src1_end) { + const ggml_tensor * src0 = op->src[0]; + const ggml_tensor * src1 = op->src[1]; + ggml_tensor * dst = op; + + GGML_TENSOR_BINARY_OP_LOCALS + + const size_t src1_col_stride = ggml_row_size(PARAM_TYPE, ne10); + + GGML_ASSERT(ne03 == 1 && ne13 == 1); + GGML_ASSERT(ne12 % ne02 == 0); + const int64_t r2 = ne12 / ne02; + + const int64_t i12 = src1_start / ne1; + const int64_t i11 = src1_start - i12 * ne1; + + // Determine batch index + const int64_t i02 = i12 / r2; + + const int64_t i1 = i11; + const int64_t i2 = i12; + + const char * src0_ptr = (const char *) src0->data + i02 * nb02; + const char * src1_ptr = (const char *) params->wdata + (i11 + i12 * ne11) * src1_col_stride; + char * dst_ptr = ((char *) dst->data + (i1 * nb1 + i2 * nb2)); + + const int64_t nrows = src1_end - src1_start; + const int64_t ncols = src0_end - src0_start; + + GGML_ASSERT(src1_ptr + src1_col_stride * nrows <= (const char *) params->wdata + params->wsize); + + // If there are more than three rows in src1, use gemm; otherwise, use gemv. + if (nrows > 3) { + gemm(ne00, (float *) (dst_ptr) + src0_start, nb1 / nb0, + src0_ptr + src0_start * nb01, src1_ptr, + nrows - (nrows % 4), ncols); + } + for (int iter = nrows - (nrows % 4); iter < nrows; iter++) { + gemv(ne00, (float *) (dst_ptr + (iter * nb1)) + src0_start, + ne01, src0_ptr + src0_start * nb01, + src1_ptr + (src1_col_stride * iter), 1 /* nrows */, ncols); + } + } + + void forward_mul_mat(ggml_compute_params * params, ggml_tensor * op) { + const ggml_tensor * src0 = op->src[0]; + const ggml_tensor * src1 = op->src[1]; + ggml_tensor * dst = op; + + GGML_TENSOR_BINARY_OP_LOCALS + + const int ith = params->ith; + const int nth = params->nth; + + GGML_ASSERT(ne0 == ne01); + GGML_ASSERT(ne1 == ne11); + GGML_ASSERT(ne2 == ne12); + GGML_ASSERT(ne3 == ne13); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + // TODO: General batched mul mat for 4D tensors + // Currently only supports 3D tensors + GGML_ASSERT(ne03 == 1); + GGML_ASSERT(ne13 == 1); + GGML_ASSERT(ne3 == 1); + + GGML_ASSERT(src1->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_n_dims(op->src[0]) == 2); + // GGML_ASSERT(ggml_n_dims(op->src[1]) == 2); + + char * wdata = static_cast(params->wdata); + const size_t nbw1 = ggml_row_size(PARAM_TYPE, ne10); + const size_t nbw2 = nbw1 * ne11; + + assert(params->wsize >= nbw2 * ne12); + + const ggml_from_float_t from_float = ggml_get_type_traits_cpu(PARAM_TYPE)->from_float; + + // INFO: Quantization is done in planes to avoid extra complexity in chunking. + // Flattening dimensions not multiple of INTER_SIZE would require extra handling depending on how + // the planes are broadcast. + for (int64_t i12 = 0; i12 < ne12; i12++) { + char * data_ptr = (char *) src1->data + i12 * nb12; + char * wdata_ptr = wdata + i12 * nbw2; + + for (int64_t i11 = ith * 4; i11 < ne11 - ne11 % 4; i11 += nth * 4) { + ggml_quantize_mat_t((float *) (data_ptr + i11 * nb11), + (void *) (wdata_ptr + i11 * nbw1), 4, ne10); + } + + const int64_t i11_processed = ne11 - ne11 % 4; + for (int64_t i11 = i11_processed + ith; i11 < ne11; i11 += nth) { + from_float((float *) (data_ptr + i11 * nb11), (void *) (wdata_ptr + i11 * nbw1), ne10); + } + } + + // disable for NUMA + const bool disable_chunking = ggml_is_numa(); + + // 4x chunks per thread + const int64_t nr0 = ggml_nrows(op->src[0]); + + int nth_scaled = nth * 4; + int64_t chunk_size0 = (nr0 + nth_scaled - 1) / nth_scaled; + int64_t nchunk0 = (nr0 + chunk_size0 - 1) / chunk_size0; + + // src1 is chunked only by full planes. + // When we flatten we need to address dimensions not multiple of the q8 INTER_SIZE + // to route them thorugh GEMV. + // nchunk1 = ne12 also avoids messing the chunking for models with no 3d tensors + // to avoid affecting their performance + int64_t nchunk1 = ne12; + + // Ensure minimum chunk size to avoid alignment issues with high thread counts + // Minimum chunk size should be at least NB_COLS to prevent overlapping chunks after alignment + const int64_t min_chunk_size = NB_COLS; + if (nchunk0 > 0 && (nr0 / nchunk0) < min_chunk_size && nr0 >= min_chunk_size) { + nchunk0 = (nr0 + min_chunk_size - 1) / min_chunk_size; + } + + int64_t dr0 = (nr0 + nchunk0 - 1) / nchunk0; + // Only increase nchunk0 to nth if it won't make chunks too small + if (nth == 1 || ((nchunk0 < nth || disable_chunking) && (nr0 + nth - 1) / nth >= min_chunk_size)) { + nchunk0 = nth; + dr0 = (nr0 + nchunk0 - 1) / nchunk0; + } + + // Ensure nchunk doesn't exceed the number of rows divided by minimum chunk size + // This prevents creating too many tiny chunks that could overlap after alignment + const int64_t max_nchunk = (nr0 + min_chunk_size - 1) / min_chunk_size; + nchunk0 = MIN(nchunk0, max_nchunk); + + if (ith == 0) { + // Every thread starts at ith, so the first unprocessed chunk is nth. This save a bit of coordination right at the start. + ggml_threadpool_chunk_set(params->threadpool, nth); + } + + ggml_barrier(params->threadpool); + + // The first chunk comes from our thread_id, the rest will get auto-assigned. + int current_chunk = ith; + + while (current_chunk < nchunk0 * nchunk1) { + const int64_t ith0 = current_chunk % nchunk0; + const int64_t ith1 = current_chunk / nchunk0; + + int64_t src0_start = dr0 * ith0; + int64_t src0_end = MIN(src0_start + dr0, nr0); + + // full-plane range for src1 + int64_t src1_start = ith1 * ne11; + int64_t src1_end = (ith1 + 1) * ne11; + + // Align boundaries to NB_COLS - round up to ensure all data is included + // The chunk size limiting above ensures chunks are large enough to prevent overlaps + src0_start = (src0_start % NB_COLS) ? src0_start + NB_COLS - (src0_start % NB_COLS) : src0_start; + src0_end = (src0_end % NB_COLS) ? src0_end + NB_COLS - (src0_end % NB_COLS) : src0_end; + src0_end = MIN(src0_end, ne01); + + // Make sure current plane is the last one before exiting + if (src0_start >= src0_end) { + current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1); + continue; + } + + forward_mul_mat_one_chunk(params, dst, src0_start, src0_end, src1_start, src1_end); + + current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1); + } + } + + void forward_mul_mat_id(ggml_compute_params * params, ggml_tensor * op) { + const ggml_tensor * src0 = op->src[0]; + const ggml_tensor * src1 = op->src[1]; + const ggml_tensor * ids = op->src[2]; + ggml_tensor * dst = op; + + GGML_TENSOR_BINARY_OP_LOCALS + + const int ith = params->ith; + const int nth = params->nth; + + const ggml_from_float_t from_float = ggml_get_type_traits_cpu(PARAM_TYPE)->from_float; + + // we don't support permuted src0 or src1 + GGML_ASSERT(nb00 == ggml_type_size(src0->type)); + GGML_ASSERT(nb10 == ggml_type_size(src1->type)); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + GGML_ASSERT(ne03 == 1); + GGML_ASSERT(ne13 == 1); + GGML_ASSERT(ne3 == 1); + + GGML_ASSERT(src1->type == GGML_TYPE_F32); + + // row groups + const int n_ids = ids->ne[0]; // n_expert_used + const int n_as = ne02; // n_expert + + const size_t nbw1 = ggml_row_size(PARAM_TYPE, ne10); + const size_t nbw2 = nbw1*ne11; + const size_t nbw3 = nbw2*ne12; + + struct mmid_row_mapping { + int32_t i1; + int32_t i2; + }; + + GGML_ASSERT(params->wsize >= + (GGML_PAD(nbw3, sizeof(int64_t)) + + n_as*(ne12 + 1)*sizeof(mmid_row_mapping)) + ); + + auto * wdata = (char *)params->wdata; + auto * wdata_src1_end = (char *)wdata + GGML_PAD(nbw3, sizeof(int64_t)); + + // total of [n_as][ne12 + 1] elements of type mmid_row_mapping (2*int32_t = int64_t) + auto * matrix_row_counts = (int64_t *) (wdata_src1_end); // [n_as] + struct mmid_row_mapping * matrix_rows = (struct mmid_row_mapping *) (matrix_row_counts + n_as); // [n_as][ne12] + + // src1: float32 => param type + for (int64_t i12 = 0; i12 < ne12; ++i12) { + for (int64_t i11 = ith; i11 < ne11; i11 += nth) { + from_float((float *)((char *) src1->data + i12 * nb12 + i11 * nb11), + (void *) (wdata + i12 * nbw2 + i11 * nbw1), + ne10); + } + } + +#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * ne12 + (i1)] + + if (ith == 0) { + // initialize matrix_row_counts + memset(matrix_row_counts, 0, n_as * sizeof(int64_t)); + + // group rows by src0 matrix + for (int32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { + for (int32_t id = 0; id < n_ids; ++id) { + const int32_t i02 = + *(const int32_t *) ((const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]); + + GGML_ASSERT(i02 >= 0 && i02 < n_as); + + MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = { id, iid1 }; + matrix_row_counts[i02] += 1; + } + } + } + + ggml_barrier(params->threadpool); + + // compute each matrix multiplication in sequence + for (int cur_a = 0; cur_a < n_as; ++cur_a) { + const int64_t cne1 = matrix_row_counts[cur_a]; + + if (cne1 == 0) { + continue; + } + + const auto * src0_cur = (const char *) src0->data + cur_a*nb02; + + //const int64_t nr0 = ne01; // src0 rows + const int64_t nr1 = cne1; // src1 rows + + int64_t src0_cur_start = (ith * ne01) / nth; + int64_t src0_cur_end = ((ith + 1) * ne01) / nth; + + // Align boundaries to NB_COLS - round up to ensure all data is included + src0_cur_start = (src0_cur_start % NB_COLS) ? src0_cur_start + NB_COLS - (src0_cur_start % NB_COLS) : src0_cur_start; + src0_cur_end = (src0_cur_end % NB_COLS) ? src0_cur_end + NB_COLS - (src0_cur_end % NB_COLS) : src0_cur_end; + if (src0_cur_end > ne01) { + src0_cur_end = ne01; + } + + if (src0_cur_start >= src0_cur_end) { + return; + } + + for (int ir1 = 0; ir1 < nr1; ir1++) { + struct mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, ir1); + + const int id = row_mapping.i1; // selected expert index + + const int64_t i11 = id % ne11; + const int64_t i12 = row_mapping.i2; // row index in src1 + + const int64_t i1 = id; // selected expert index + const int64_t i2 = i12; // row + + const auto * src1_col = (const char *) wdata + (i11 * nbw1 + i12 * nbw2); + + gemv( + ne00, (float *) ((char *) dst->data + (i1 * nb1 + i2 * nb2)) + src0_cur_start, ne01, + src0_cur + src0_cur_start * nb01, src1_col, 1, src0_cur_end - src0_cur_start); + } + } +#undef MMID_MATRIX_ROW + } + + int repack(struct ggml_tensor * t, const void * data, size_t data_size) override { + GGML_LOG_DEBUG("%s: repack tensor %s with %s_%dx%d\n", __func__, t->name, ggml_type_name(t->type), + (int) NB_COLS, (int) INTER_SIZE); + return ggml::cpu::repack::repack(t, data, data_size); + } +}; + +} // namespace ggml::cpu::repack + +static const ggml::cpu::tensor_traits * ggml_repack_get_optimal_repack_type(const struct ggml_tensor * cur) { + // instance for Q4 + static const ggml::cpu::repack::tensor_traits q4_0_4x4_q8_0; + static const ggml::cpu::repack::tensor_traits q4_0_4x8_q8_0; + static const ggml::cpu::repack::tensor_traits q4_0_8x8_q8_0; + + // instance for Q4_K + static const ggml::cpu::repack::tensor_traits q4_K_8x4_q8_K; + static const ggml::cpu::repack::tensor_traits q4_K_8x8_q8_K; + + // instance for Q5_K + static const ggml::cpu::repack::tensor_traits q5_K_8x4_q8_K; + static const ggml::cpu::repack::tensor_traits q5_K_8x8_q8_K; + + // instance for Q6_K + static const ggml::cpu::repack::tensor_traits q6_K_8x4_q8_K; + static const ggml::cpu::repack::tensor_traits q6_K_8x8_q8_K; + + // instance for Q2 + static const ggml::cpu::repack::tensor_traits q2_K_8x8_q8_K; + + // instance for IQ4 + static const ggml::cpu::repack::tensor_traits iq4_nl_4x4_q8_0; + static const ggml::cpu::repack::tensor_traits iq4_nl_8x8_q8_0; + + // instance for MXFP4 + static const ggml::cpu::repack::tensor_traits mxfp4_4x4_q8_0; + static const ggml::cpu::repack::tensor_traits mxfp4_8x8_q8_0; + + // instance for Q8_0 + static const ggml::cpu::repack::tensor_traits q8_0_4x4_q8_0; + static const ggml::cpu::repack::tensor_traits q8_0_4x8_q8_0; + + // instances for RISC-V + // + // These implement outer-product style matrix multiplication kernels with + // an interleave of 1. +#if defined __riscv_zvfh + static const ggml::cpu::repack::tensor_traits q4_0_16x1_q8_0; + static const ggml::cpu::repack::tensor_traits q4_K_16x1_q8_K; + static const ggml::cpu::repack::tensor_traits iq4_nl_16x1_q8_0; + static const ggml::cpu::repack::tensor_traits q8_0_16x1_q8_0; + static const ggml::cpu::repack::tensor_traits q2_K_16x1_q8_K; +#endif + + if (cur->type == GGML_TYPE_Q4_0) { + if (ggml_cpu_has_avx2() || (ggml_cpu_has_sve() && ggml_cpu_has_matmul_int8() && ggml_cpu_get_sve_cnt() == QK8_0)) { + if (cur->ne[1] % 8 == 0) { + return &q4_0_8x8_q8_0; + } + } + if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) { + if (cur->ne[1] % 4 == 0) { + return &q4_0_4x8_q8_0; + } + } + if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { + if (cur->ne[1] % 4 == 0) { + return &q4_0_4x4_q8_0; + } + } + if (ggml_cpu_has_riscv_v()) { + #if defined __riscv_zvfh + switch (__riscv_vlenb() * 8) { + case 128: { break; } // TODO + case 256: { if (cur->ne[1] % 16 == 0) { return &q4_0_16x1_q8_0; } break; } + case 512: { break; } // TODO + case 1024: { break; } // TODO + default: { return nullptr; } + } + #endif + } + } else if (cur->type == GGML_TYPE_Q4_K) { + if (ggml_cpu_has_avx2()) { + if (cur->ne[1] % 8 == 0) { + return &q4_K_8x8_q8_K; + } + } + if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) { + if (cur->ne[1] % 8 == 0) { + return &q4_K_8x8_q8_K; + } + } + if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { + if (cur->ne[1] % 8 == 0) { + return &q4_K_8x4_q8_K; + } + } + if (ggml_cpu_has_riscv_v()) { + #if defined __riscv_zvfh + switch (__riscv_vlenb() * 8) { + case 128: { break; } // TODO + case 256: { if (cur->ne[1] % 16 == 0) { return &q4_K_16x1_q8_K; } break; } + case 512: { break; } // TODO + case 1024: { break; } // TODO + default: { return nullptr; } + } + #endif + } + } else if (cur->type == GGML_TYPE_Q2_K) { + if (ggml_cpu_has_avx512()) { + if (cur->ne[1] % 8 == 0) { + return &q2_K_8x8_q8_K; + } + } + if (ggml_cpu_has_riscv_v()) { + #if defined __riscv_zvfh + switch (__riscv_vlenb() * 8) { + case 128: { break; } // TODO + case 256: { if (cur->ne[1] % 16 == 0) { return &q2_K_16x1_q8_K; } break; } + case 512: { break; } // TODO + case 1024: { break; } // TODO + default: { return nullptr; } + } + #endif + } + } else if (cur->type == GGML_TYPE_Q5_K) { + if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) { + if (cur->ne[1] % 8 == 0) { + return &q5_K_8x8_q8_K; + } + } + if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { + if (cur->ne[1] % 8 == 0) { + return &q5_K_8x4_q8_K; + } + } + } else if (cur->type == GGML_TYPE_Q6_K) { + if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) { + if (cur->ne[1] % 8 == 0) { + return &q6_K_8x8_q8_K; + } + } + if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { + if (cur->ne[1] % 8 == 0) { + return &q6_K_8x4_q8_K; + } + } + } else if (cur->type == GGML_TYPE_IQ4_NL) { + if (ggml_cpu_has_avx2()) { + if (cur->ne[1] % 8 == 0) { + return &iq4_nl_8x8_q8_0; + } + } + if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { + if (cur->ne[1] % 4 == 0) { + return &iq4_nl_4x4_q8_0; + } + } + if (ggml_cpu_has_riscv_v()) { + #if defined __riscv_zvfh + switch (__riscv_vlenb() * 8) { + case 128: { break; } // TODO + case 256: { if (cur->ne[1] % 16 == 0) { return &iq4_nl_16x1_q8_0; } break; } + case 512: { break; } // TODO + case 1024: { break; } // TODO + default: { return nullptr; } + } + #endif + } + } else if (cur->type == GGML_TYPE_MXFP4) { + if (ggml_cpu_has_avx2()) { + if (cur->ne[1] % 8 == 0) { + return &mxfp4_8x8_q8_0; + } + } + if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { + if (cur->ne[1] % 4 == 0) { + return &mxfp4_4x4_q8_0; + } + } + } else if (cur->type == GGML_TYPE_Q8_0) { + if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) { + if (cur->ne[1] % 4 == 0) { + return &q8_0_4x8_q8_0; + } + } + if (ggml_cpu_has_neon() && ggml_cpu_has_dotprod()) { + if (cur->ne[1] % 4 == 0) { + return &q8_0_4x4_q8_0; + } + } + if (ggml_cpu_has_riscv_v()) { + #if defined __riscv_zvfh + switch (__riscv_vlenb() * 8) { + case 128: { break; } // TODO + case 256: { if (cur->ne[1] % 16 == 0) { return &q8_0_16x1_q8_0; } break; } + case 512: { break; } // TODO + case 1024: { break; } // TODO + default: { return nullptr; } + } + #endif + } + } + + return nullptr; +} + +static enum ggml_status ggml_backend_cpu_repack_buffer_init_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor) { + tensor->extra = (void *) const_cast(ggml_repack_get_optimal_repack_type(tensor)); + + GGML_UNUSED(buffer); + return GGML_STATUS_SUCCESS; +} + +static void ggml_backend_cpu_repack_buffer_set_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, + const void * data, size_t offset, size_t size) { + GGML_ASSERT(offset == 0); + GGML_ASSERT(size == ggml_nbytes(tensor)); + + auto tensor_traits = (ggml::cpu::repack::tensor_traits_base *) tensor->extra; + auto OK = tensor_traits->repack(tensor, data, size); + + GGML_ASSERT(OK == 0); + GGML_UNUSED(buffer); +} + +static const char * ggml_backend_cpu_repack_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + return "CPU_REPACK"; + + GGML_UNUSED(buft); +} + +static ggml_backend_buffer_t ggml_backend_cpu_repack_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); + + if (buffer == nullptr) { + return nullptr; + } + + buffer->buft = buft; + buffer->iface.init_tensor = ggml_backend_cpu_repack_buffer_init_tensor; + buffer->iface.set_tensor = ggml_backend_cpu_repack_buffer_set_tensor; + buffer->iface.get_tensor = nullptr; + buffer->iface.cpy_tensor = nullptr; + return buffer; +} + +static size_t ggml_backend_cpu_repack_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + return TENSOR_ALIGNMENT; + + GGML_UNUSED(buft); +} + +namespace ggml::cpu::repack { +class extra_buffer_type : ggml::cpu::extra_buffer_type { + bool supports_op(ggml_backend_dev_t, const struct ggml_tensor * op) override { + if ( op->op == GGML_OP_MUL_MAT && + op->src[0]->buffer && + (ggml_n_dims(op->src[0]) == 2) && + op->src[0]->buffer->buft == ggml_backend_cpu_repack_buffer_type() && + ggml_repack_get_optimal_repack_type(op->src[0]) + ) { + if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) { + return false; + } + if (op->src[1]->type == GGML_TYPE_F32) { + return true; + } + //if (op->src[1]->type == GGML_TYPE_Q8_0) { + // return true; + //} + // may be possible if Q8_0 packed... + } else if (op->op == GGML_OP_MUL_MAT_ID + && op->src[0]->buffer + && (ggml_n_dims(op->src[0]) == 3) + && op->src[0]->buffer->buft == ggml_backend_cpu_repack_buffer_type() + && ggml_repack_get_optimal_repack_type(op->src[0]) + ) { + if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) { + return false; + } + if (op->src[1]->type == GGML_TYPE_F32) { + return true; + } + //if (op->src[1]->type == GGML_TYPE_Q8_0) { + // return true; + //} + } + return false; + } + + ggml::cpu::tensor_traits * get_tensor_traits(const struct ggml_tensor * op) override { + if (op->op == GGML_OP_MUL_MAT || op->op == GGML_OP_MUL_MAT_ID) { + if (op->src[0]->buffer && op->src[0]->buffer->buft == ggml_backend_cpu_repack_buffer_type()) { + return (ggml::cpu::tensor_traits *) op->src[0]->extra; + } + } + return nullptr; + } +}; +} // namespace ggml::cpu::repack + +ggml_backend_buffer_type_t ggml_backend_cpu_repack_buffer_type(void) { + static struct ggml_backend_buffer_type ggml_backend_cpu_buffer_type_repack = { + /* .iface = */ { + /* .get_name = */ ggml_backend_cpu_repack_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_cpu_repack_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_cpu_repack_buffer_type_get_alignment, + /* .get_max_size = */ nullptr, // defaults to SIZE_MAX + /* .get_alloc_size = */ nullptr, // defaults to ggml_nbytes + /* .is_host = */ nullptr, + }, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0), + /* .context = */ new ggml::cpu::repack::extra_buffer_type(), + }; + + return &ggml_backend_cpu_buffer_type_repack; +} + + +// ---- x86 repack implementation ---- + +#define GGML_COMMON_IMPL_CPP +#define GGML_COMMON_DECL_CPP +#include "ggml-common-defs.inc" +#include "ggml-backend-impl-defs.inc" + +#include "ggml-impl-defs.inc" +#include "ggml-cpu.h.inc" +#include "ggml-cpu-impl-defs.inc" +#include "simd-mappings-defs.inc" #include #include @@ -16,7 +5571,6 @@ #include // for GGML_ASSERT #define GGML_CPU_CLANG_WORKAROUND -#include "../../repack.h" #if defined(__GNUC__) #pragma GCC diagnostic ignored "-Woverlength-strings" @@ -105,7 +5659,7 @@ static inline __m256 __avx_rearranged_f32cx8_load(ggml_fp16_t *x, __m128i arrang #endif #endif -static inline int nearest_int(float fval) { +static inline int x86_nearest_int(float fval) { assert(fabsf(fval) <= 4194303.f); float val = fval + 12582912.f; int i; memcpy(&i, &val, sizeof(int)); @@ -160,7 +5714,7 @@ static inline __m256i mul_sum_us8_pairs_acc_int32x8(const __m256i acc, const __m #endif } -// Integer variant of the function defined in ggml-quants.c +// Integer variant of the function defined in the ggml quantization implementation // multiply int8_t, add results pairwise twice and return as 256 bit int vector, then add the accumulator static inline __m256i mul_sum_i8_pairs_acc_int32x8(const __m256i acc, const __m256i x, const __m256i y) { #if defined(__AVXVNNIINT8__) @@ -6405,3 +11959,6 @@ void ggml_gemm_q2_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo #endif } + + +#undef UNUSED diff --git a/ggml/src/ggml-cpu/repack.h b/ggml/src/ggml-cpu/repack.h deleted file mode 100644 index cb21edf62394..000000000000 --- a/ggml/src/ggml-cpu/repack.h +++ /dev/null @@ -1,245 +0,0 @@ -#pragma once - -#define GGML_COMMON_DECL_CPP -#include "ggml-common.h" - -#include "traits.h" -#include "ggml.h" - -// GGML internal header - -ggml_backend_buffer_type_t ggml_backend_cpu_repack_buffer_type(void); - -template constexpr int QK_0() { - if constexpr (K == 4) { - return QK4_0; - } - if constexpr (K == 8) { - return QK8_0; - } - return -1; -} - -template struct block { - ggml_half d[N]; // deltas for N qK_0 blocks - int8_t qs[(QK_0() * N * K) / 8]; // quants for N qK_0 blocks -}; - -// control size -static_assert(sizeof(block<4, 4>) == 4 * sizeof(ggml_half) + QK8_0 * 2, "wrong block<4,4> size/padding"); -static_assert(sizeof(block<4, 8>) == 8 * sizeof(ggml_half) + QK8_0 * 4, "wrong block<4,8> size/padding"); -static_assert(sizeof(block<4, 16>) == 16 * sizeof(ggml_half) + QK8_0 * 8, "wrong block<4,16> size/padding"); -static_assert(sizeof(block<8, 4>) == 4 * sizeof(ggml_half) + QK8_0 * 4, "wrong block<8,4> size/padding"); -static_assert(sizeof(block<8, 8>) == 8 * sizeof(ggml_half) + QK8_0 * 8, "wrong block<8,8> size/padding"); -static_assert(sizeof(block<8, 16>) == 16 * sizeof(ggml_half) + QK8_0 * 16, "wrong block<8,16> size/padding"); - -using block_q4_0x4 = block<4, 4>; -using block_q4_0x8 = block<4, 8>; -using block_q4_0x16 = block<4, 16>; -using block_q8_0x4 = block<8, 4>; -using block_q8_0x8 = block<8, 8>; -using block_q8_0x16 = block<8, 16>; - -struct block_q4_Kx8 { - ggml_half d[8]; // super-block scale for quantized scales - ggml_half dmin[8]; // super-block scale for quantized mins - uint8_t scales[96]; // scales and mins, quantized with 6 bits - uint8_t qs[1024]; // 4--bit quants -}; - -static_assert(sizeof(block_q4_Kx8) == sizeof(ggml_half) * 16 + K_SCALE_SIZE * 8 + QK_K * 4, "wrong q4_K block size/padding"); -struct block_q4_Kx16 { - ggml_half d[16]; // super-block scale for quantized scales - ggml_half dmin[16]; // super-block scale for quantized mins - uint8_t scales[192]; // scales and mins, quantized with 6 bits - uint8_t qs[2048]; // 4--bit quants -}; - -static_assert(sizeof(block_q4_Kx16) == sizeof(ggml_half) * 32 + K_SCALE_SIZE * 16 + QK_K * 8, "wrong q4_K block size/padding"); -struct block_q2_Kx8 { - ggml_half d[8]; // super-block scale for quantized scales - ggml_half dmin[8]; // super-block scale for quantized mins - uint8_t scales[128]; // scales and mins, quantized with 4 bits - uint8_t qs[512]; // 2--bit quants -}; - -static_assert(sizeof(block_q2_Kx8) == sizeof(ggml_half) * 16 + QK_K/2 + QK_K * 2, "wrong q2_K block size/padding"); -struct block_q2_Kx16 { - ggml_half d[16]; // Super-block scale for quantized scales - ggml_half dmin[16]; // Super-block scale for quantized mins - uint8_t scales[256]; // Sub-block scales (16 cols * 16 sub-blocks) - uint8_t qs[1024]; // Data (16 cols * 64 bytes per block) -}; -static_assert(sizeof(block_q2_Kx16) == sizeof(ggml_half) * 32 + QK_K + QK_K * 4, "wrong q2_K block size/padding"); - -struct block_q5_Kx8 { - ggml_half d[8]; // super-block scale for quantized scales - ggml_half dmin[8]; // super-block scale for quantized mins - uint8_t scales[96]; // scales and mins, quantized with 6 bits - uint8_t qh[QK_K * 8 / 8]; // high bits of 5-bit quants - uint8_t qs[QK_K * 8 / 2]; // low bits of 5-bit quants (in groups of 4) -}; - -static_assert(sizeof(block_q5_Kx8) == sizeof(ggml_half) * 16 + K_SCALE_SIZE * 8 + QK_K * 5, - "wrong q5_K block size/padding"); - -struct block_q6_Kx8 { - ggml_half d[8]; - int8_t scales[QK_K / 16 * 8]; - uint8_t ql[QK_K / 2 * 8]; // low bits of 6-bit quants (groups of 2) - uint8_t qh[QK_K / 4 * 8]; // high bits of 6-bit quants (groups of 4) -}; - -static_assert(sizeof(block_q6_Kx8) == sizeof(ggml_half) * 8 + QK_K / 16 * 8 + 3 * QK_K / 4 * 8, - "wrong q6_K block size/padding"); - -struct block_q8_Kx4 { - float d[4]; // delta - int8_t qs[QK_K * 4]; // quants - int16_t bsums[QK_K / 4]; // sum of quants in groups of 16 -}; - -static_assert(sizeof(block_q8_Kx4) == sizeof(float) * 4 + QK_K * 4 + (QK_K / 4) * sizeof(int16_t), "wrong q8_K block size/padding"); - -struct block_iq4_nlx4 { - ggml_half d[4]; // deltas for 4 iq4_nl blocks - uint8_t qs[QK4_NL * 2]; // nibbles / quants for 4 iq4_nl blocks -}; - -static_assert(sizeof(block_iq4_nlx4) == 4 * sizeof(ggml_half) + QK4_NL * 2, "wrong iq4_nlx4 block size/padding"); - -struct block_iq4_nlx8 { - ggml_half d[8]; // deltas for 8 iq4_nl blocks - uint8_t qs[QK4_NL * 4]; // nibbles / quants for 8 iq4_nl blocks -}; - -static_assert(sizeof(block_iq4_nlx8) == 8 * sizeof(ggml_half) + QK4_NL * 4, "wrong iq4_nlx8 block size/padding"); - -struct block_iq4_nlx16 { - ggml_half d[16]; // deltas for 16 iq4_nl blocks - uint8_t qs[QK4_NL * 8]; // nibbles / quants for 16 iq4_nl blocks -}; - -static_assert(sizeof(block_iq4_nlx16) == 16 * sizeof(ggml_half) + QK4_NL * 8, "wrong iq4_nlx16 block size/padding"); -struct block_mxfp4x4 { - uint8_t e[4]; - uint8_t qs[QK_MXFP4 * 2]; -}; -static_assert(sizeof(block_mxfp4x4) == 4 + QK_MXFP4 * 2, "wrong mxfp4x4 block size/padding"); - -struct block_mxfp4x8 { - uint8_t e[8]; - uint8_t qs[QK_MXFP4 * 4]; -}; -static_assert(sizeof(block_mxfp4x8) == 8 + QK_MXFP4 * 4, "wrong mxfp4x8 block size/padding"); - -#if defined(__cplusplus) -extern "C" { -#endif - -void ggml_quantize_mat_q8_0_4x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void ggml_quantize_mat_q8_0_4x8(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void ggml_quantize_mat_q8_K_4x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void ggml_quantize_mat_q8_K_4x8(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void ggml_gemv_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q4_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q4_0_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q2_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q4_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q5_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q5_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q6_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q6_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_iq4_nl_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_iq4_nl_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_mxfp4_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_mxfp4_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q8_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q8_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_0_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q2_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q5_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q5_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q6_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q6_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_iq4_nl_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_iq4_nl_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_mxfp4_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_mxfp4_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q8_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q8_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -#if defined __riscv_zvfh -void ggml_quantize_mat_q8_0_4x1(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void ggml_quantize_mat_q8_K_4x1(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void ggml_gemv_q4_0_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q4_K_16x1_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_iq4_nl_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q8_0_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q2_K_16x1_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_0_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_K_16x1_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_iq4_nl_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q8_0_16x1_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q2_K_16x1_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -#endif - -// Native implementations -void ggml_quantize_mat_q8_0_4x4_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void ggml_quantize_mat_q8_0_4x8_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void ggml_quantize_mat_q8_K_4x4_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void ggml_quantize_mat_q8_K_4x8_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void ggml_gemv_q4_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q4_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q4_0_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q2_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q4_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q4_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q5_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q5_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q6_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q6_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_iq4_nl_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_iq4_nl_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_mxfp4_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_mxfp4_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q8_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q8_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_0_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q2_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q5_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q5_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q6_K_8x4_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q6_K_8x8_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_iq4_nl_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_iq4_nl_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_mxfp4_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_mxfp4_8x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q8_0_4x4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q8_0_4x8_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -#if defined __riscv_zvfh -void ggml_quantize_mat_q8_0_4x1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void ggml_quantize_mat_q8_K_4x1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); -void ggml_gemv_q4_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q4_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q8_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_q2_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemv_iq4_nl_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q4_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q8_0_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_q2_K_16x1_q8_K_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -void ggml_gemm_iq4_nl_16x1_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc); -#endif - -#if defined(__cplusplus) -} // extern "C" -#endif diff --git a/ggml/src/ggml-cpu/simd-gemm.h b/ggml/src/ggml-cpu/simd-gemm.h deleted file mode 100644 index 4119d04f895a..000000000000 --- a/ggml/src/ggml-cpu/simd-gemm.h +++ /dev/null @@ -1,226 +0,0 @@ -#pragma once - -// Computes C[M x N] += A[M x K] * B[K x N] - -#include "simd-mappings.h" - -// TODO: add support for sizeless vector types -#if defined(GGML_SIMD) && !defined(__ARM_FEATURE_SVE) && !defined(__riscv_v_intrinsic) - -// TODO: untested on avx512 -// These are in units of GGML_F32_EPR -#if defined(__AVX512F__) || defined (__ARM_NEON__) - static constexpr int GEMM_RM = 4; - static constexpr int GEMM_RN = 4; // 16+4+1 = 25/32 -#elif defined(__AVX2__) || defined(__AVX__) - static constexpr int GEMM_RM = 6; - static constexpr int GEMM_RN = 2; // 12+2+1 = 15/16 -#else - static constexpr int GEMM_RM = 2; - static constexpr int GEMM_RN = 2; -#endif - -template -static inline void simd_gemm_ukernel( - float * GGML_RESTRICT C, - const float * GGML_RESTRICT A, - const float * GGML_RESTRICT B, - int K, int N) -{ - static constexpr int KN = GGML_F32_EPR; - - GGML_F32_VEC acc[RM][RN]; - for (int64_t i = 0; i < RM; i++) { - for (int r = 0; r < RN; r++) { - acc[i][r] = GGML_F32_VEC_LOAD(C + i * N + r * KN); - } - } - - for (int64_t kk = 0; kk < K; kk++) { - GGML_F32_VEC Bv[RN]; - for (int r = 0; r < RN; r++) { - Bv[r] = GGML_F32_VEC_LOAD(B + kk * N + r * KN); - } - for (int64_t i = 0; i < RM; i++) { - GGML_F32_VEC p = GGML_F32_VEC_SET1(A[i * K + kk]); - for (int r = 0; r < RN; r++) { - acc[i][r] = GGML_F32_VEC_FMA(acc[i][r], Bv[r], p); - } - } - } - - for (int64_t i = 0; i < RM; i++) { - for (int r = 0; r < RN; r++) { - GGML_F32_VEC_STORE(C + i * N + r * KN, acc[i][r]); - } - } -} - -// C[M x N] += A[M x K] * B[K x N] -static void simd_gemm( - float * GGML_RESTRICT C, - const float * GGML_RESTRICT A, - const float * GGML_RESTRICT B, - int M, int K, int N) -{ - static constexpr int KN = GGML_F32_EPR; - - int64_t ii = 0; - for (; ii + GEMM_RM <= M; ii += GEMM_RM) { - int64_t jj = 0; - for (; jj + GEMM_RN * KN <= N; jj += GEMM_RN * KN) { - simd_gemm_ukernel(C + jj, A, B + jj, K, N); - } - for (; jj + KN <= N; jj += KN) { - simd_gemm_ukernel(C + jj, A, B + jj, K, N); - } - for (; jj < N; jj++) { - for (int64_t i = 0; i < GEMM_RM; i++) { - float a = C[i * N + jj]; - for (int64_t kk = 0; kk < K; kk++) { - a += A[i + kk] * B[kk * N + jj]; - } - C[i * N + jj] = a; - } - } - - A += GEMM_RM * K; - C += GEMM_RM * N; - } - - // Tail rows: one at a time - for (; ii < M; ii++) { - int64_t jj = 0; - for (; jj + GEMM_RN * KN <= N; jj += GEMM_RN * KN) { - simd_gemm_ukernel<1, GEMM_RN>(C + jj, A, B + jj, K, N); - } - for (; jj + KN <= N; jj += KN) { - simd_gemm_ukernel<1, 1>(C + jj, A, B + jj, K, N); - } - for (; jj < N; jj++) { - float a = C[jj]; - for (int64_t kk = 0; kk < K; kk++) { - a += A[kk] * B[kk * N + jj]; - } - C[jj] = a; - } - - A += K; - C += N; - } -} -#elif defined(GGML_SIMD) && defined(__riscv_v_intrinsic) -// RM accumulators + 1 B vector = RM + 1 <= 8 => RM <= 7 -// Microkernel: C[RM x vl] += A[RM x K] * B[K x N] -template -static inline void rvv_simd_gemm_ukernel( - float * GGML_RESTRICT C, - const float * GGML_RESTRICT A, - const float * GGML_RESTRICT B, - int K, int N, size_t vl) -{ - static_assert(RM >= 1 && RM <= 7, "RM must be 1..7 for LMUL=4"); - - vfloat32m4_t acc_0 = __riscv_vle32_v_f32m4(C + 0 * N, vl); - vfloat32m4_t acc_1, acc_2, acc_3, acc_4, acc_5, acc_6; - if constexpr (RM > 1) acc_1 = __riscv_vle32_v_f32m4(C + 1 * N, vl); - if constexpr (RM > 2) acc_2 = __riscv_vle32_v_f32m4(C + 2 * N, vl); - if constexpr (RM > 3) acc_3 = __riscv_vle32_v_f32m4(C + 3 * N, vl); - if constexpr (RM > 4) acc_4 = __riscv_vle32_v_f32m4(C + 4 * N, vl); - if constexpr (RM > 5) acc_5 = __riscv_vle32_v_f32m4(C + 5 * N, vl); - if constexpr (RM > 6) acc_6 = __riscv_vle32_v_f32m4(C + 6 * N, vl); - - for (int kk = 0; kk < K; kk++) { - vfloat32m4_t b_0 = __riscv_vle32_v_f32m4(B + kk * N, vl); - - acc_0 = __riscv_vfmacc_vf_f32m4(acc_0, A[0 * K + kk], b_0, vl); - if constexpr (RM > 1) acc_1 = __riscv_vfmacc_vf_f32m4(acc_1, A[1 * K + kk], b_0, vl); - if constexpr (RM > 2) acc_2 = __riscv_vfmacc_vf_f32m4(acc_2, A[2 * K + kk], b_0, vl); - if constexpr (RM > 3) acc_3 = __riscv_vfmacc_vf_f32m4(acc_3, A[3 * K + kk], b_0, vl); - if constexpr (RM > 4) acc_4 = __riscv_vfmacc_vf_f32m4(acc_4, A[4 * K + kk], b_0, vl); - if constexpr (RM > 5) acc_5 = __riscv_vfmacc_vf_f32m4(acc_5, A[5 * K + kk], b_0, vl); - if constexpr (RM > 6) acc_6 = __riscv_vfmacc_vf_f32m4(acc_6, A[6 * K + kk], b_0, vl); - } - - __riscv_vse32_v_f32m4(C + 0 * N, acc_0, vl); - if constexpr (RM > 1) __riscv_vse32_v_f32m4(C + 1 * N, acc_1, vl); - if constexpr (RM > 2) __riscv_vse32_v_f32m4(C + 2 * N, acc_2, vl); - if constexpr (RM > 3) __riscv_vse32_v_f32m4(C + 3 * N, acc_3, vl); - if constexpr (RM > 4) __riscv_vse32_v_f32m4(C + 4 * N, acc_4, vl); - if constexpr (RM > 5) __riscv_vse32_v_f32m4(C + 5 * N, acc_5, vl); - if constexpr (RM > 6) __riscv_vse32_v_f32m4(C + 6 * N, acc_6, vl); -} - -template -static inline void rvv_simd_gemm_dispatch_tail( - float * GGML_RESTRICT C, - const float * GGML_RESTRICT A, - const float * GGML_RESTRICT B, - int K, int N, int KN, int remaining_rows) -{ - if constexpr (RM > 0) { - if (remaining_rows == RM) { - int64_t jj = 0; - for (; jj + KN <= N; jj += KN) { - rvv_simd_gemm_ukernel(C + jj, A, B + jj, K, N, KN); - } - if (jj < N) { - rvv_simd_gemm_ukernel(C + jj, A, B + jj, K, N, N - jj); - } - } else { - rvv_simd_gemm_dispatch_tail(C, A, B, K, N, KN, remaining_rows); - } - } -} - -static constexpr int GEMM_RM = 7; - -// C[M x N] += A[M x K] * B[K x N] -static void simd_gemm( - float * GGML_RESTRICT C, - const float * GGML_RESTRICT A, - const float * GGML_RESTRICT B, - int M, int K, int N) -{ - const int KN = (int)__riscv_vlenb(); - int64_t ii = 0; - for (; ii + GEMM_RM <= M; ii += GEMM_RM) { - int64_t jj = 0; - for (; jj + KN <= N; jj += KN) { - rvv_simd_gemm_ukernel(C + jj, A, B + jj, K, N, KN); - } - if (jj < N) { - rvv_simd_gemm_ukernel(C + jj, A, B + jj, K, N, N - jj); - } - A += GEMM_RM * K; - C += GEMM_RM * N; - } - - int remaining_rows = M - ii; - rvv_simd_gemm_dispatch_tail(C, A, B, K, N, KN, remaining_rows); -} - -#if defined(__GNUC__) && !defined(__clang__) -#pragma GCC diagnostic pop -#endif - -#else // scalar path - -static void simd_gemm( - float * GGML_RESTRICT C, - const float * GGML_RESTRICT A, - const float * GGML_RESTRICT B, - int M, int K, int N) -{ - for (int64_t i = 0; i < M; i++) { - for (int64_t j = 0; j < N; j++) { - float sum = C[i * N + j]; - for (int64_t kk = 0; kk < K; kk++) { - sum += A[i * K + kk] * B[kk * N + j]; - } - C[i * N + j] = sum; - } - } -} - -#endif // GGML_SIMD diff --git a/ggml/src/ggml-cpu/simd-mappings.h b/ggml/src/ggml-cpu/simd-mappings-defs.inc similarity index 99% rename from ggml/src/ggml-cpu/simd-mappings.h rename to ggml/src/ggml-cpu/simd-mappings-defs.inc index 0deda9309856..bdb40f0214c9 100644 --- a/ggml/src/ggml-cpu/simd-mappings.h +++ b/ggml/src/ggml-cpu/simd-mappings-defs.inc @@ -1,6 +1,6 @@ #pragma once -#include "ggml-cpu-impl.h" +#include "ggml-cpu-impl-defs.inc" #ifdef __ARM_FEATURE_SVE #include diff --git a/ggml/src/ggml-cpu/spacemit/ime.cpp b/ggml/src/ggml-cpu/spacemit/ime.cpp.inc similarity index 90% rename from ggml/src/ggml-cpu/spacemit/ime.cpp rename to ggml/src/ggml-cpu/spacemit/ime.cpp.inc index 91fe1925eaa3..4016ad08eb6d 100644 --- a/ggml/src/ggml-cpu/spacemit/ime.cpp +++ b/ggml/src/ggml-cpu/spacemit/ime.cpp.inc @@ -1,13 +1,137 @@ #define GGML_COMMON_IMPL_CPP #define GGML_COMMON_DECL_CPP -#include "ime.h" +#include "ggml-backend-impl-defs.inc" +#include "ggml-common-defs.inc" +#include "ggml-cpu.h.inc" +#include "ggml.h.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml-cpu-impl-defs.inc" +#include "ggml-impl-defs.inc" +#include "simd-mappings-defs.inc" + +#ifdef __cplusplus +#include +extern "C" { +#endif + +// return true if op part of extra "accelerator" +bool ggml_cpu_extra_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op); +bool ggml_cpu_extra_work_size(int n_threads, const struct ggml_tensor * op, size_t * size); + +#ifdef __cplusplus +} + +namespace ggml::cpu { +// register in tensor->extra +class tensor_traits { + public: + virtual ~tensor_traits(); + virtual bool work_size(int n_threads, const struct ggml_tensor * op, size_t & size) = 0; + virtual bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) = 0; +}; + +class extra_buffer_type { + public: + virtual ~extra_buffer_type(); + virtual bool supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) = 0; + virtual tensor_traits * get_tensor_traits(const struct ggml_tensor * op) = 0; +}; +} // namespace ggml::cpu + +// implemented in ggml-cpu.cpp. +std::vector & ggml_backend_cpu_get_extra_buffer_types(); + +#endif + +#define GGML_FA_TILE_Q 64 +#define GGML_FA_TILE_KV 64 + +#ifdef __cplusplus + +#include + +// convenience functions/macros for use in template calls +// note: these won't be required after the 'traits' lookup table is used. +static inline ggml_fp16_t f32_to_f16(float x) { + return GGML_CPU_FP32_TO_FP16(x); +} + +static inline float f16_to_f32(ggml_fp16_t x) { + return GGML_CPU_FP16_TO_FP32(x); +} + +static inline ggml_bf16_t f32_to_bf16(float x) { + return GGML_FP32_TO_BF16(x); +} + +static inline float bf16_to_f32(ggml_bf16_t x) { + return GGML_BF16_TO_FP32(x); +} + +static inline float i32_to_f32(int32_t x) { + return x; +} + +static inline int32_t f32_to_i32(float x) { + return x; +} + +static inline float f32_to_f32(float x) { + return x; +} + +// TODO - merge this into the traits table, after using row-based conversions +template +struct type_conversion_table; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_fp16_t) = f16_to_f32; + static constexpr ggml_fp16_t (*from_f32)(float) = f32_to_f16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(float) = f32_to_f32; + static constexpr float (*from_f32)(float) = f32_to_f32; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(ggml_bf16_t) = bf16_to_f32; + static constexpr ggml_bf16_t (*from_f32)(float) = f32_to_bf16; +}; + +template <> +struct type_conversion_table { + static constexpr float (*to_f32)(int32_t) = i32_to_f32; + static constexpr int32_t (*from_f32)(float) = f32_to_i32; +}; + +static std::pair get_thread_range(const struct ggml_compute_params * params, const struct ggml_tensor * src0) { + const int64_t ith = params->ith; + const int64_t nth = params->nth; + + const int64_t nr = ggml_nrows(src0); + + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; + + // row range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); + + return {ir0, ir1}; +} + +struct ggml_fa_tile_config { + static constexpr size_t Q = GGML_FA_TILE_Q; + static constexpr size_t KV = GGML_FA_TILE_KV; +}; + +#endif -#include "ggml-backend-impl.h" -#include "ggml-common.h" -#include "ggml-cpu.h" -#include "ime_kernels.h" -#include "traits.h" #include #include @@ -16,6 +140,29 @@ #include #include +extern "C" ggml_backend_buffer_type_t ggml_backend_cpu_riscv64_spacemit_buffer_type(void); + +namespace sqnbitgemm_spacemit_ime { +namespace ime1 { +size_t gemm_kernel_i8i4(size_t blk_len, + const std::byte * quant_a_ptr, + const std::byte * quant_b_data, + const float * quant_b_scale, + const std::byte * quant_b_zp, + float * c_ptr, + size_t count_m, + size_t count_n, + size_t count_k, + size_t block_count_k, + size_t ldc, + const float * bias, + const size_t scale_stride); + +void quantize_a_row_i8(size_t blk_len, const float * a_ptr, size_t count_k, std::byte * quant_a_ptr); +void quantize_a_4row_i8(size_t blk_len, const float * a_ptr, size_t count_k, std::byte * quant_a_ptr); +} // namespace ime1 +} // namespace sqnbitgemm_spacemit_ime + // clang-format off #if defined(__riscv) diff --git a/ggml/src/ggml-cpu/spacemit/ime.h b/ggml/src/ggml-cpu/spacemit/ime.h deleted file mode 100644 index 800d91acdaef..000000000000 --- a/ggml/src/ggml-cpu/spacemit/ime.h +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "ggml-alloc.h" - -#ifdef __cplusplus -extern "C" { -#endif - -ggml_backend_buffer_type_t ggml_backend_cpu_riscv64_spacemit_buffer_type(void); - -#ifdef __cplusplus -} -#endif diff --git a/ggml/src/ggml-cpu/spacemit/ime1_kernels.cpp b/ggml/src/ggml-cpu/spacemit/ime1_kernels.cpp.inc similarity index 99% rename from ggml/src/ggml-cpu/spacemit/ime1_kernels.cpp rename to ggml/src/ggml-cpu/spacemit/ime1_kernels.cpp.inc index cbbb6cd91607..470a3aeb3fc8 100644 --- a/ggml/src/ggml-cpu/spacemit/ime1_kernels.cpp +++ b/ggml/src/ggml-cpu/spacemit/ime1_kernels.cpp.inc @@ -1,9 +1,29 @@ -#include "ggml.h" -#include "ime_kernels.h" +#include "ggml.h.inc" #include #include +namespace sqnbitgemm_spacemit_ime { +namespace ime1 { +size_t gemm_kernel_i8i4(size_t blk_len, + const std::byte * quant_a_ptr, + const std::byte * quant_b_data, + const float * quant_b_scale, + const std::byte * quant_b_zp, + float * c_ptr, + size_t count_m, + size_t count_n, + size_t count_k, + size_t block_count_k, + size_t ldc, + const float * bias, + const size_t scale_stride); + +void quantize_a_row_i8(size_t blk_len, const float * a_ptr, size_t count_k, std::byte * quant_a_ptr); +void quantize_a_4row_i8(size_t blk_len, const float * a_ptr, size_t count_k, std::byte * quant_a_ptr); +} // namespace ime1 +} // namespace sqnbitgemm_spacemit_ime + // clang-format off #if defined(__GNUC__) #pragma GCC diagnostic ignored "-Woverlength-strings" diff --git a/ggml/src/ggml-cpu/spacemit/ime_kernels.h b/ggml/src/ggml-cpu/spacemit/ime_kernels.h deleted file mode 100644 index 757063415053..000000000000 --- a/ggml/src/ggml-cpu/spacemit/ime_kernels.h +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -namespace sqnbitgemm_spacemit_ime { -namespace ime1 { -size_t gemm_kernel_i8i4(size_t blk_len, - const std::byte * quant_a_ptr, - const std::byte * quant_b_data, - const float * quant_b_scale, - const std::byte * quant_b_zp, - float * c_ptr, - size_t count_m, - size_t count_n, - size_t count_k, - size_t block_count_k, - size_t ldc, - const float * bias, - const size_t scale_stride); - -void quantize_a_row_i8(size_t blk_len, const float * a_ptr, size_t count_k, std::byte * quant_a_ptr); - -void quantize_a_4row_i8(size_t blk_len, const float * a_ptr, size_t count_k, std::byte * quant_a_ptr); - -} // namespace ime1 -} // namespace sqnbitgemm_spacemit_ime diff --git a/ggml/src/ggml-cpu/traits.cpp b/ggml/src/ggml-cpu/traits.cpp deleted file mode 100644 index 4f32f10255aa..000000000000 --- a/ggml/src/ggml-cpu/traits.cpp +++ /dev/null @@ -1,36 +0,0 @@ -#include "traits.h" - -#include "ggml-backend-impl.h" -#include "ggml-backend.h" - -namespace ggml::cpu { -tensor_traits::~tensor_traits() {} - -extra_buffer_type::~extra_buffer_type() {} -} // namespace ggml::cpu - -bool ggml_cpu_extra_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) { - for (auto extra : ggml_backend_cpu_get_extra_buffer_types()) { - if (extra && extra->context) { - auto buf_extra = (ggml::cpu::extra_buffer_type *) extra->context; - auto tensor_traits = buf_extra->get_tensor_traits(op); - if (tensor_traits && tensor_traits->compute_forward(params, op)) { - return true; - } - } - } - return false; -} - -bool ggml_cpu_extra_work_size(int n_threads, const struct ggml_tensor * op, size_t * size) { - for (auto extra : ggml_backend_cpu_get_extra_buffer_types()) { - if (extra && extra->context) { - auto buf_extra = (ggml::cpu::extra_buffer_type *) extra->context; - auto tensor_traits = buf_extra->get_tensor_traits(op); - if (tensor_traits && tensor_traits->work_size(n_threads, op, *size)) { - return true; - } - } - } - return false; -} diff --git a/ggml/src/ggml-cpu/traits.h b/ggml/src/ggml-cpu/traits.h deleted file mode 100644 index f4e0990ddfc9..000000000000 --- a/ggml/src/ggml-cpu/traits.h +++ /dev/null @@ -1,38 +0,0 @@ -#pragma once -#include "ggml-backend-impl.h" -#include "ggml-cpu-impl.h" -#include "ggml.h" - -#ifdef __cplusplus -# include -extern "C" { -#endif - -// return true if op part of extra "accelerator" -bool ggml_cpu_extra_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op); -bool ggml_cpu_extra_work_size(int n_threads, const struct ggml_tensor * op, size_t * size); - -#ifdef __cplusplus -} - -namespace ggml::cpu { -// register in tensor->extra -class tensor_traits { - public: - virtual ~tensor_traits(); - virtual bool work_size(int n_threads, const struct ggml_tensor * op, size_t & size) = 0; - virtual bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) = 0; -}; - -class extra_buffer_type { - public: - virtual ~extra_buffer_type(); - virtual bool supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) = 0; - virtual tensor_traits * get_tensor_traits(const struct ggml_tensor * op) = 0; -}; -} // namespace ggml::cpu - -// implemented in ggml-cpu.cpp. -std::vector & ggml_backend_cpu_get_extra_buffer_types(); - -#endif diff --git a/ggml/src/ggml-cpu/unary-ops.cpp b/ggml/src/ggml-cpu/unary-ops.cpp deleted file mode 100644 index 1d8344436f02..000000000000 --- a/ggml/src/ggml-cpu/unary-ops.cpp +++ /dev/null @@ -1,337 +0,0 @@ -#include "unary-ops.h" - -static inline float op_abs(float x) { - return fabsf(x); -} - -static inline float op_sgn(float x) { - return (x > 0.f) ? 1.f : ((x < 0.f) ? -1.f : 0.f); -} - -static inline float op_neg(float x) { - return -x; -} - -static inline float op_step(float x) { - return (x > 0.f) ? 1.f : 0.f; -} - -static inline float op_tanh(float x) { - return tanhf(x); -} - -static inline float op_elu(float x) { - return (x > 0.f) ? x : expm1f(x); -} - -static inline float op_relu(float x) { - return (x > 0.f) ? x : 0.f; -} - -static inline float op_sigmoid(float x) { - return 1.f / (1.f + expf(-x)); -} - -static inline float op_hardsigmoid(float x) { - return fminf(1.0f, fmaxf(0.0f, (x + 3.0f) / 6.0f)); -} - -static inline float op_exp(float x) { - return expf(x); -} - -static inline float op_hardswish(float x) { - return x * fminf(1.0f, fmaxf(0.0f, (x + 3.0f) / 6.0f)); -} - -static inline float op_sqr(float x) { - return x * x; -} - -static inline float op_sqrt(float x) { - return sqrtf(x); -} - -static inline float op_xielu(float x, float alpha_n, float alpha_p, float beta, float eps) { - if (x > 0.0f) { - return alpha_p * x * x + beta * x; - } else { - const float min_x_eps = fminf(x, eps); - return (expm1f(min_x_eps) - x) * alpha_n + beta * x; - } -} - -static inline float op_sin(float x) { - return sinf(x); -} - -static inline float op_cos(float x) { - return cosf(x); -} - -static inline float op_log(float x) { - return logf(x); -} - -static inline float op_expm1(float x) { - return expf(x) - 1.0f; -} - -static inline float op_softplus(float x) { - return (x > 20.0f) ? x : logf(1.0f + expf(x)); -} - -static inline float op_floor(float x) { - return floorf(x); -} - -static inline float op_ceil(float x) { - return ceilf(x); -} - -static inline float op_round(float x) { - return roundf(x); -} - -static inline float op_trunc(float x) { - return truncf(x); -} - -template -static inline void vec_unary_op(int64_t n, dst_t * y, const src0_t * x) { - constexpr auto src0_to_f32 = type_conversion_table::to_f32; - constexpr auto f32_to_dst = type_conversion_table::from_f32; - - for (int i = 0; i < n; i++) { - y[i] = f32_to_dst(op(src0_to_f32(x[i]))); - } -} - -template -static void apply_unary_op(const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - - GGML_ASSERT(ggml_is_contiguous_rows(src0) && ggml_is_contiguous_rows(dst) && ggml_are_same_shape(src0, dst)); - - GGML_TENSOR_UNARY_OP_LOCALS - - GGML_ASSERT( nb0 == sizeof(dst_t)); - GGML_ASSERT(nb00 == sizeof(src0_t)); - - const auto [ir0, ir1] = get_thread_range(params, src0); - - for (int64_t ir = ir0; ir < ir1; ++ir) { - const int64_t i03 = ir/(ne02*ne01); - const int64_t i02 = (ir - i03*ne02*ne01)/ne01; - const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); - - dst_t * dst_ptr = (dst_t *) ((char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1 ); - const src0_t * src0_ptr = (const src0_t *) ((const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01); - - vec_unary_op(ne0, dst_ptr, src0_ptr); - } -} - -// TODO: Use the 'traits' lookup table (for type conversion fns), instead of a mass of 'if' conditions with long templates -template -static void unary_op(const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - - /* */ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { // all f32 - apply_unary_op(params, dst); - } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { // all f16 - apply_unary_op(params, dst); - } else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) { // all bf16 - apply_unary_op(params, dst); - } else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_F32) { - apply_unary_op(params, dst); - } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) { - apply_unary_op(params, dst); - } else { - fprintf(stderr, "%s: unsupported types: dst: %s, src0: %s\n", __func__, - ggml_type_name(dst->type), ggml_type_name(src0->type)); - GGML_ABORT("fatal error"); - } -} - -template -static void unary_op_params(const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - - /* */ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { // all f32 - apply_unary_op(params, dst); - } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { // all f16 - apply_unary_op(params, dst); - } else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) { // all bf16 - apply_unary_op(params, dst); - } else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_F32) { - apply_unary_op(params, dst); - } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) { - apply_unary_op(params, dst); - } else { - fprintf(stderr, "%s: unsupported types: dst: %s, src0: %s\n", __func__, - ggml_type_name(dst->type), ggml_type_name(src0->type)); - GGML_ABORT("fatal error"); - } -} - -// Extend vec_unary_op to support functors -template -static inline void vec_unary_op_functor(int64_t n, dst_t * y, const src0_t * x, Op op) { - constexpr auto src0_to_f32 = type_conversion_table::to_f32; - constexpr auto f32_to_dst = type_conversion_table::from_f32; - - for (int i = 0; i < n; i++) { - y[i] = f32_to_dst(op(src0_to_f32(x[i]))); - } -} - -// Extend apply_unary_op to support functors -template -static void apply_unary_op_functor(const ggml_compute_params * params, ggml_tensor * dst, Op op) { - const ggml_tensor * src0 = dst->src[0]; - - GGML_ASSERT(ggml_is_contiguous_1(src0) && ggml_is_contiguous_1(dst) && ggml_are_same_shape(src0, dst)); - - GGML_TENSOR_UNARY_OP_LOCALS - - GGML_ASSERT( nb0 == sizeof(dst_t)); - GGML_ASSERT(nb00 == sizeof(src0_t)); - - const auto [ir0, ir1] = get_thread_range(params, src0); - - for (int64_t ir = ir0; ir < ir1; ++ir) { - const int64_t i03 = ir/(ne02*ne01); - const int64_t i02 = (ir - i03*ne02*ne01)/ne01; - const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01); - - dst_t * dst_ptr = (dst_t *) ((char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1 ); - const src0_t * src0_ptr = (const src0_t *) ((const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01); - - vec_unary_op_functor(ne0, dst_ptr, src0_ptr, op); - } -} - -// Generic dispatcher for functors -template -static void unary_op_functor(const ggml_compute_params * params, ggml_tensor * dst, Op op) { - const ggml_tensor * src0 = dst->src[0]; - - /* */ if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { // all f32 - apply_unary_op_functor(params, dst, op); - } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { // all f16 - apply_unary_op_functor(params, dst, op); - } else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) { // all bf16 - apply_unary_op_functor(params, dst, op); - } else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_F32) { - apply_unary_op_functor(params, dst, op); - } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) { - apply_unary_op_functor(params, dst, op); - } else { - fprintf(stderr, "%s: unsupported types: dst: %s, src0: %s\n", __func__, - ggml_type_name(dst->type), ggml_type_name(src0->type)); - GGML_ABORT("fatal error"); - } -} - -void ggml_compute_forward_abs(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_sgn(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_neg(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_step(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_tanh(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_elu(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_relu(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_sigmoid(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_hardsigmoid(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_exp(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_hardswish(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_sqr(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_sqrt(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_sin(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_cos(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_log(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_expm1(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_softplus(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_floor(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_ceil(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_round(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_trunc(const ggml_compute_params * params, ggml_tensor * dst) { - unary_op(params, dst); -} - -void ggml_compute_forward_xielu(const ggml_compute_params * params, ggml_tensor * dst) { - const float alpha_n = ggml_get_op_params_f32(dst, 1); - const float alpha_p = ggml_get_op_params_f32(dst, 2); - const float beta = ggml_get_op_params_f32(dst, 3); - const float eps = ggml_get_op_params_f32(dst, 4); - - const auto xielu_op_params = [alpha_n, alpha_p, beta, eps](float f) { - return op_xielu(f, alpha_n, alpha_p, beta, eps); - }; - - unary_op_functor(params, dst, xielu_op_params); -} - diff --git a/ggml/src/ggml-cpu/unary-ops.h b/ggml/src/ggml-cpu/unary-ops.h deleted file mode 100644 index bcad5a3af1a9..000000000000 --- a/ggml/src/ggml-cpu/unary-ops.h +++ /dev/null @@ -1,35 +0,0 @@ -#pragma once - -#include "common.h" - -#ifdef __cplusplus -extern "C" { -#endif - -void ggml_compute_forward_abs(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_sgn(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_neg(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_step(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_tanh(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_elu(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_relu(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_sigmoid(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_hardsigmoid(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_exp(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_hardswish(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_sqr(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_sqrt(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_sin(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_cos(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_log(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_expm1(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_softplus(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_floor(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_ceil(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_round(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_trunc(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_xielu(const struct ggml_compute_params * params, struct ggml_tensor * dst); - -#ifdef __cplusplus -} -#endif diff --git a/ggml/src/ggml-cpu/vec.cpp b/ggml/src/ggml-cpu/vec.cpp deleted file mode 100644 index d0e4001338a2..000000000000 --- a/ggml/src/ggml-cpu/vec.cpp +++ /dev/null @@ -1,629 +0,0 @@ -#include "vec.h" - -#include - -// precomputed gelu table for f16 (128 KB) -ggml_fp16_t ggml_table_gelu_f16[1 << 16]; - -// precomputed quick gelu table for f16 (128 KB) -ggml_fp16_t ggml_table_gelu_quick_f16[1 << 16]; - -void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc) { - assert(nrc == 1); - GGML_UNUSED(nrc); - GGML_UNUSED(bx); - GGML_UNUSED(by); - GGML_UNUSED(bs); - -#if defined(GGML_SIMD) - float sumf = 0.0f; - - #if defined(__ARM_FEATURE_SVE) - const int sve_register_length = ggml_cpu_get_sve_cnt() * 8; - const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16 - const int ggml_f32_step = 8 * ggml_f32_epr; // choose 8 SVE registers - - const int np = (n & ~(ggml_f32_step - 1)); - svfloat32_t sum1 = svdup_n_f32(0.0f); - svfloat32_t sum2 = svdup_n_f32(0.0f); - svfloat32_t sum3 = svdup_n_f32(0.0f); - svfloat32_t sum4 = svdup_n_f32(0.0f); - svfloat32_t sum5 = svdup_n_f32(0.0f); - svfloat32_t sum6 = svdup_n_f32(0.0f); - svfloat32_t sum7 = svdup_n_f32(0.0f); - svfloat32_t sum8 = svdup_n_f32(0.0f); - svfloat32_t ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8; - svfloat32_t ay1,ay2,ay3,ay4,ay5,ay6,ay7,ay8; - for (int i = 0; i < np; i += ggml_f32_step) { - ax1 = GGML_F32_VEC_LOAD(x + i); - ay1 = GGML_F32_VEC_LOAD(y + i); - sum1 = GGML_F32_VEC_FMA(sum1, ax1, ay1); - - ax2 = GGML_F32_VEC_LOAD(x + i + 1*ggml_f32_epr); - ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr); - sum2 = GGML_F32_VEC_FMA(sum2, ax2, ay2); - - ax3 = GGML_F32_VEC_LOAD(x + i + 2*ggml_f32_epr); - ay3 = GGML_F32_VEC_LOAD(y + i + 2*ggml_f32_epr); - sum3 = GGML_F32_VEC_FMA(sum3, ax3, ay3); - - ax4 = GGML_F32_VEC_LOAD(x + i + 3*ggml_f32_epr); - ay4 = GGML_F32_VEC_LOAD(y + i + 3*ggml_f32_epr); - sum4 = GGML_F32_VEC_FMA(sum4, ax4, ay4); - - ax5 = GGML_F32_VEC_LOAD(x + i + 4*ggml_f32_epr); - ay5 = GGML_F32_VEC_LOAD(y + i + 4*ggml_f32_epr); - sum5 = GGML_F32_VEC_FMA(sum5, ax5, ay5); - - ax6 = GGML_F32_VEC_LOAD(x + i + 5*ggml_f32_epr); - ay6 = GGML_F32_VEC_LOAD(y + i + 5*ggml_f32_epr); - sum6 = GGML_F32_VEC_FMA(sum6, ax6, ay6); - - ax7 = GGML_F32_VEC_LOAD(x + i + 6*ggml_f32_epr); - ay7 = GGML_F32_VEC_LOAD(y + i + 6*ggml_f32_epr); - sum7 = GGML_F32_VEC_FMA(sum7, ax7, ay7); - - ax8 = GGML_F32_VEC_LOAD(x + i + 7*ggml_f32_epr); - ay8 = GGML_F32_VEC_LOAD(y + i + 7*ggml_f32_epr); - sum8 = GGML_F32_VEC_FMA(sum8, ax8, ay8); - } - // leftovers - // Since 8 unrolls are done in above loop, leftovers lie in range [0, ggml_f32_step] which is handled in below loop - const int np2 = (n & ~(ggml_f32_epr - 1)); - for (int i = np; i < np2; i += ggml_f32_epr) { - ax1 = GGML_F32_VEC_LOAD(x + i); - ay1 = GGML_F32_VEC_LOAD(y + i); - sum1 = GGML_F32_VEC_FMA(sum1, ax1, ay1); - } - // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only - if (np2 < n) { - svbool_t pg = svwhilelt_b32(np2, n); - ax1 = svld1_f32(pg, x + np2); - ay1 = svld1_f32(pg, y + np2); - sum1 = svmad_f32_m(pg, ax1, ay1, sum1); - } - // reduce sum1,sum2 to sum1 - GGML_F32_VEC_REDUCE(sumf, sum1, sum2, sum3, sum4, sum5, sum6, sum7, sum8); - #elif defined(__riscv_v_intrinsic) - int vl = __riscv_vsetvlmax_e32m8(); - vfloat32m1_t vs = __riscv_vfmv_v_f_f32m1(0.0f, 1); - vfloat32m8_t vsum; - vfloat32m8_t ax; - vfloat32m8_t ay; - vsum = __riscv_vfmv_v_f_f32m8_tu(vsum, 0.0f, vl); - for (int i = 0; i < n; i += vl) { - vl = __riscv_vsetvl_e32m8(n - i); - ax = __riscv_vle32_v_f32m8_tu(ax, &x[i], vl); - ay = __riscv_vle32_v_f32m8_tu(ay, &y[i], vl); - vsum = __riscv_vfmacc_vv_f32m8_tu(vsum, ax, ay, vl); - } - vl = __riscv_vsetvlmax_e32m8(); - vs = __riscv_vfredusum_vs_f32m8_f32m1(vsum, vs, vl); - sumf += __riscv_vfmv_f_s_f32m1_f32(vs); - #else - const int np = (n & ~(GGML_F32_STEP - 1)); - - GGML_F32_VEC sum[GGML_F32_ARR] = { GGML_F32_VEC_ZERO }; - - GGML_F32_VEC ax[GGML_F32_ARR]; - GGML_F32_VEC ay[GGML_F32_ARR]; - - for (int i = 0; i < np; i += GGML_F32_STEP) { - for (int j = 0; j < GGML_F32_ARR; j++) { - ax[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR); - ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); - - sum[j] = GGML_F32_VEC_FMA(sum[j], ax[j], ay[j]); - } - } - - // reduce sum0..sum3 to sum0 - GGML_F32_VEC_REDUCE(sumf, sum); - - // leftovers - for (int i = np; i < n; ++i) { - sumf += x[i]*y[i]; - } - #endif -#else - // scalar - ggml_float sumf = 0.0; - for (int i = 0; i < n; ++i) { - sumf += (ggml_float)(x[i]*y[i]); - } -#endif - - *s = sumf; -} - -void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t * GGML_RESTRICT x, size_t bx, ggml_bf16_t * GGML_RESTRICT y, size_t by, int nrc) { - assert(nrc == 1); - GGML_UNUSED(nrc); - GGML_UNUSED(bx); - GGML_UNUSED(by); - GGML_UNUSED(bs); - int i = 0; - ggml_float sumf = 0; - -#if defined(__AVX512BF16__) - __m512 c1 = _mm512_setzero_ps(); - __m512 c2 = _mm512_setzero_ps(); - for (; i + 64 <= n; i += 64) { - c1 = _mm512_dpbf16_ps(c1, m512bh(_mm512_loadu_si512((x + i))), - m512bh(_mm512_loadu_si512((y + i)))); - c2 = _mm512_dpbf16_ps(c2, m512bh(_mm512_loadu_si512((x + i + 32))), - m512bh(_mm512_loadu_si512((y + i + 32)))); - } - sumf += (ggml_float)_mm512_reduce_add_ps(c1); - sumf += (ggml_float)_mm512_reduce_add_ps(c2); - -#elif defined(__AVX512F__) -#define LOAD(p) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)(p))), 16)) - __m512 c1 = _mm512_setzero_ps(); - __m512 c2 = _mm512_setzero_ps(); - for (; i + 32 <= n; i += 32) { - c1 = _mm512_add_ps(_mm512_mul_ps(LOAD(x + i), LOAD(y + i)), c1); - c2 = _mm512_add_ps(_mm512_mul_ps(LOAD(x + i + 16), LOAD(y + i + 16)), c2); - } - sumf += (ggml_float)_mm512_reduce_add_ps(c1); - sumf += (ggml_float)_mm512_reduce_add_ps(c2); - -#undef LOAD -#elif defined(__AVX2__) || defined(__AVX__) -#if defined(__AVX2__) -#define LOAD(p) _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(_mm_loadu_si128((const __m128i *)(p))), 16)) -#else -#define LOAD(p) _mm256_castsi256_ps(_mm256_insertf128_si256(_mm256_castsi128_si256(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadu_si128((const __m128i *)(p))), 16)), (_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_bsrli_si128(_mm_loadu_si128((const __m128i *)(p)), 8)), 16)), 1)) -#endif - __m256 c1 = _mm256_setzero_ps(); - __m256 c2 = _mm256_setzero_ps(); - __m256 c3 = _mm256_setzero_ps(); - __m256 c4 = _mm256_setzero_ps(); - for (; i + 32 <= n; i += 32) { - c1 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i), LOAD(y + i)), c1); - c2 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i + 8), LOAD(y + i + 8)), c2); - c3 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i + 16), LOAD(y + i + 16)), c3); - c4 = _mm256_add_ps(_mm256_mul_ps(LOAD(x + i + 24), LOAD(y + i + 24)), c4); - } - __m128 g; - c1 = _mm256_add_ps(_mm256_add_ps(c1, c3), - _mm256_add_ps(c2, c4)); - g = _mm_add_ps(_mm256_extractf128_ps(c1, 1), - _mm256_castps256_ps128(c1)); - g = _mm_add_ps(g, _mm_movehl_ps(g, g)); - g = _mm_add_ss(g, _mm_movehdup_ps(g)); - sumf += (ggml_float)_mm_cvtss_f32(g); - -#undef LOAD -#elif defined(__riscv_v_intrinsic) && defined(__riscv_zvfbfwma) - size_t vl = __riscv_vsetvlmax_e32m4(); - - // initialize accumulators to all zeroes - vfloat32m4_t vsum0 = __riscv_vfmv_v_f_f32m4(0.0f, vl); - vfloat32m4_t vsum1 = __riscv_vfmv_v_f_f32m4(0.0f, vl); - - // calculate step size - const size_t epr = __riscv_vsetvlmax_e16m2(); - const size_t step = epr * 2; - const int np = (n & ~(step - 1)); - - // unroll by 2 - for (; i < np; i += step) { - vbfloat16m2_t ax0 = __riscv_vle16_v_bf16m2((const __bf16 *)&x[i], epr); - vbfloat16m2_t ay0 = __riscv_vle16_v_bf16m2((const __bf16 *)&y[i], epr); - vsum0 = __riscv_vfwmaccbf16_vv_f32m4(vsum0, ax0, ay0, epr); - __asm__ __volatile__ ("" ::: "memory"); - - vbfloat16m2_t ax1 = __riscv_vle16_v_bf16m2((const __bf16 *)&x[i + epr], epr); - vbfloat16m2_t ay1 = __riscv_vle16_v_bf16m2((const __bf16 *)&y[i + epr], epr); - vsum1 = __riscv_vfwmaccbf16_vv_f32m4(vsum1, ax1, ay1, epr); - __asm__ __volatile__ ("" ::: "memory"); - } - - // accumulate in 1 register - vsum0 = __riscv_vfadd_vv_f32m4(vsum0, vsum1, vl); - - // leftovers - for (i = np; i < n; i += vl) { - vl = __riscv_vsetvl_e16m2(n - i); - vbfloat16m2_t ax0 = __riscv_vle16_v_bf16m2((const __bf16 *)&x[i], vl); - vbfloat16m2_t ay0 = __riscv_vle16_v_bf16m2((const __bf16 *)&y[i], vl); - vsum0 = __riscv_vfwmaccbf16_vv_f32m4(vsum0, ax0, ay0, vl); - } - - // reduce - vl = __riscv_vsetvlmax_e32m4(); - vfloat32m1_t redsum = __riscv_vfredusum_vs_f32m4_f32m1(vsum0, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl); - sumf += __riscv_vfmv_f_s_f32m1_f32(redsum); - -#elif defined(__POWER9_VECTOR__) || defined(__VXE__) || defined(__VXE2__) - const int np = (n & ~(GGML_BF16_STEP - 1)); - if (np > 0) { - GGML_F32_VEC sum[4] = {GGML_F32_VEC_ZERO}; - for (; i < np; i += GGML_BF16_STEP) { - GGML_BF16_VEC vx0 = GGML_BF16_VEC_LOAD(x + i); - GGML_BF16_VEC vx1 = GGML_BF16_VEC_LOAD(x + i + 8); - GGML_BF16_VEC vy0 = GGML_BF16_VEC_LOAD(y + i); - GGML_BF16_VEC vy1 = GGML_BF16_VEC_LOAD(y + i + 8); - GGML_BF16_FMA_LO(sum[0], vx0, vy0); - GGML_BF16_FMA_HI(sum[1], vx0, vy0); - GGML_BF16_FMA_LO(sum[2], vx1, vy1); - GGML_BF16_FMA_HI(sum[3], vx1, vy1); - } - GGML_F32x4_REDUCE_4(sumf, sum[0], sum[1], sum[2], sum[3]); - } -#endif - - for (; i < n; ++i) { - sumf += (ggml_float)(GGML_BF16_TO_FP32(x[i]) * - GGML_BF16_TO_FP32(y[i])); - } - *s = sumf; -} - -void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * GGML_RESTRICT x, size_t bx, ggml_fp16_t * GGML_RESTRICT y, size_t by, int nrc) { - assert(nrc == 1); - GGML_UNUSED(nrc); - GGML_UNUSED(bx); - GGML_UNUSED(by); - GGML_UNUSED(bs); - - ggml_float sumf = 0.0; - - -#if defined(GGML_SIMD) - #if defined(__ARM_FEATURE_SVE) - const int sve_register_length = svcntb() * 8; //get vector length - const int ggml_f16_epr = sve_register_length / 16; // running when 16 - const int ggml_f16_step = 8 * ggml_f16_epr; // choose 8 SVE registers - - const int np= (n & ~(ggml_f16_step - 1)); - svfloat16_t sum1 = svdup_n_f16(0.0f); - svfloat16_t sum2 = svdup_n_f16(0.0f); - svfloat16_t sum3 = svdup_n_f16(0.0f); - svfloat16_t sum4 = svdup_n_f16(0.0f); - - svfloat16_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; - svfloat16_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; - for (int i = 0; i < np; i += ggml_f16_step) { - ax1 = GGML_F16x_VEC_LOAD(x + i + 0 * ggml_f16_epr, 0); - ay1 = GGML_F16x_VEC_LOAD(y + i + 0 * ggml_f16_epr, 0); - sum1 = GGML_F16x_VEC_FMA(sum1, ax1, ay1); - - ax2 = GGML_F16x_VEC_LOAD(x + i + 1 * ggml_f16_epr, 1); - ay2 = GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 1); - sum2 = GGML_F16x_VEC_FMA(sum2, ax2, ay2); - - ax3 = GGML_F16x_VEC_LOAD(x + i + 2 * ggml_f16_epr, 2); - ay3 = GGML_F16x_VEC_LOAD(y + i + 2 * ggml_f16_epr, 2); - sum3 = GGML_F16x_VEC_FMA(sum3, ax3, ay3); - - ax4 = GGML_F16x_VEC_LOAD(x + i + 3 * ggml_f16_epr, 3); - ay4 = GGML_F16x_VEC_LOAD(y + i + 3 * ggml_f16_epr, 3); - sum4 = GGML_F16x_VEC_FMA(sum4, ax4, ay4); - - ax5 = GGML_F16x_VEC_LOAD(x + i + 4 * ggml_f16_epr, 4); - ay5 = GGML_F16x_VEC_LOAD(y + i + 4 * ggml_f16_epr, 4); - sum1 = GGML_F16x_VEC_FMA(sum1, ax5, ay5); - - ax6 = GGML_F16x_VEC_LOAD(x + i + 5 * ggml_f16_epr, 5); - ay6 = GGML_F16x_VEC_LOAD(y + i + 5 * ggml_f16_epr, 5); - sum2 = GGML_F16x_VEC_FMA(sum2, ax6, ay6); - - ax7 = GGML_F16x_VEC_LOAD(x + i + 6 * ggml_f16_epr, 6); - ay7 = GGML_F16x_VEC_LOAD(y + i + 6 * ggml_f16_epr, 6); - sum3 = GGML_F16x_VEC_FMA(sum3, ax7, ay7); - - ax8 = GGML_F16x_VEC_LOAD(x + i + 7 * ggml_f16_epr, 7); - ay8 = GGML_F16x_VEC_LOAD(y + i + 7 * ggml_f16_epr, 7); - sum4 = GGML_F16x_VEC_FMA(sum4, ax8, ay8); - } - - const int np2 = (n & ~(ggml_f16_epr - 1)); // round down to multiple of 8 - for (int k = np; k < np2; k += ggml_f16_epr) { - svfloat16_t rx = GGML_F16x_VEC_LOAD(x + k, 0); - svfloat16_t ry = GGML_F16x_VEC_LOAD(y + k, 0); - sum1 = GGML_F16x_VEC_FMA(sum1, rx, ry); - } - - if (np2 < n) { - svbool_t pg = svwhilelt_b16(np2, n); - svfloat16_t hx = svld1_f16(pg, (const __fp16 *)(x + np2)); - svfloat16_t hy = svld1_f16(pg, (const __fp16 *)(y + np2)); - - sum1 = svmad_f16_x(pg, hx, hy, sum1); - } - GGML_F16x_VEC_REDUCE(sumf, sum1, sum2, sum3, sum4); - #elif defined(__riscv_v_intrinsic) - #if defined(__riscv_zvfh) - int vl = __riscv_vsetvlmax_e32m2(); - vfloat32m1_t vs = __riscv_vfmv_v_f_f32m1(0.0f, 1); - vfloat32m2_t vsum; - vfloat16m1_t ax; - vfloat16m1_t ay; - vsum = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vmv_v_x_u32m2(0, vl)); - for (int i = 0; i < n; i += vl) { - vl = __riscv_vsetvl_e16m1(n - i); - ax = __riscv_vle16_v_f16m1_tu(ax, (const _Float16 *)&x[i], vl); - ay = __riscv_vle16_v_f16m1_tu(ay, (const _Float16 *)&y[i], vl); - vsum = __riscv_vfwmacc_vv_f32m2_tu(vsum, ax, ay, vl); - } - vl = __riscv_vsetvlmax_e32m1(); - vfloat32m1_t ac0 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(vsum, 0), __riscv_vget_v_f32m2_f32m1(vsum, 1), vl); - vs = __riscv_vfredusum_vs_f32m1_f32m1(ac0, vs, vl); - sumf += __riscv_vfmv_f_s_f32m1_f32(vs); - #else - for (int i = 0; i < n; ++i) { - sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*GGML_CPU_FP16_TO_FP32(y[i])); - } - #endif // __riscv_zvfh - #else - const int np = (n & ~(GGML_F16_STEP - 1)); - - GGML_F16_VEC sum[GGML_F16_ARR] = { GGML_F16_VEC_ZERO }; - - GGML_F16_VEC ax[GGML_F16_ARR]; - GGML_F16_VEC ay[GGML_F16_ARR]; - - for (int i = 0; i < np; i += GGML_F16_STEP) { - for (int j = 0; j < GGML_F16_ARR; j++) { - ax[j] = GGML_F16_VEC_LOAD(x + i + j*GGML_F16_EPR, j); - ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); - - sum[j] = GGML_F16_VEC_FMA(sum[j], ax[j], ay[j]); - } - } - - // reduce sum0..sum3 to sum0 - GGML_F16_VEC_REDUCE(sumf, sum); - - // leftovers - for (int i = np; i < n; ++i) { - sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*GGML_CPU_FP16_TO_FP32(y[i])); - } - // if you hit this, you are likely running outside the FP range - assert(!isnan(sumf) && !isinf(sumf)); - #endif -#else - for (int i = 0; i < n; ++i) { - sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*GGML_CPU_FP16_TO_FP32(y[i])); - } -#endif // GGML_SIMD - - *s = sumf; -} - -void ggml_vec_silu_f32(const int n, float * y, const float * x) { - int i = 0; -#if defined(__AVX512F__) && defined(__AVX512DQ__) - for (; i + 15 < n; i += 16) { - _mm512_storeu_ps(y + i, ggml_v_silu(_mm512_loadu_ps(x + i))); - } -#elif defined(__AVX2__) && defined(__FMA__) - for (; i + 7 < n; i += 8) { - _mm256_storeu_ps(y + i, ggml_v_silu(_mm256_loadu_ps(x + i))); - } -#elif defined(__SSE2__) - for (; i + 3 < n; i += 4) { - _mm_storeu_ps(y + i, ggml_v_silu(_mm_loadu_ps(x + i))); - } -#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) - const int vlen = svcntw(); - for (; i < n; i += vlen) { - const svbool_t pg = svwhilelt_b32_s32(i, n); - svst1_f32(pg, y + i, ggml_v_silu(pg, svld1_f32(pg, x + i))); - } -#elif defined(__ARM_NEON) && defined(__aarch64__) - for (; i + 3 < n; i += 4) { - vst1q_f32(y + i, ggml_v_silu(vld1q_f32(x + i))); - } -#elif defined(__riscv_v_intrinsic) - for (int vl; i < n; i += vl) { - vl = __riscv_vsetvl_e32m2(n - i); - vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl); - vfloat32m2_t vy = ggml_v_silu_m2(vx, vl); - __riscv_vse32_v_f32m2(&y[i], vy, vl); - } -#endif - for (; i < n; ++i) { - y[i] = ggml_silu_f32(x[i]); - } -} - -void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float * g) { - int i = 0; -#if defined(__AVX512F__) && defined(__AVX512DQ__) - for (; i + 15 < n; i += 16) { - _mm512_storeu_ps(y + i, _mm512_mul_ps(ggml_v_silu(_mm512_loadu_ps(x + i)), _mm512_loadu_ps(g + i))); - } -#elif defined(__AVX2__) && defined(__FMA__) - for (; i + 7 < n; i += 8) { - _mm256_storeu_ps(y + i, _mm256_mul_ps(ggml_v_silu(_mm256_loadu_ps(x + i)), _mm256_loadu_ps(g + i))); - } -#elif defined(__SSE2__) - for (; i + 3 < n; i += 4) { - _mm_storeu_ps(y + i, _mm_mul_ps(ggml_v_silu(_mm_loadu_ps(x + i)), _mm_loadu_ps(g + i))); - } -#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) - const int vlen = svcntw(); - for (; i < n; i += vlen) { - const svbool_t pg = svwhilelt_b32_s32(i, n); - svst1_f32(pg, y + i, svmul_f32_x(pg, ggml_v_silu(pg, svld1_f32(pg, x + i)), svld1_f32(pg, g + i))); - } -#elif defined(__ARM_NEON) && defined(__aarch64__) - for (; i + 3 < n; i += 4) { - vst1q_f32(y + i, vmulq_f32(ggml_v_silu(vld1q_f32(x + i)), vld1q_f32(g + i))); - } -#elif defined(__riscv_v_intrinsic) - for (int vl; i < n; i += vl) { - vl = __riscv_vsetvl_e32m2(n - i); - vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl); - vfloat32m2_t vg = __riscv_vle32_v_f32m2(&g[i], vl); - vfloat32m2_t vy = __riscv_vfmul_vv_f32m2(ggml_v_silu_m2(vx, vl), vg, vl); - __riscv_vse32_v_f32m2(&y[i], vy, vl); - } -#endif - for (; i < n; ++i) { - y[i] = ggml_silu_f32(x[i]) * g[i]; - } -} - -ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean) { - int i = 0; - ggml_float sum = 0; -// TODO: optimize to process the remaining elements in groups using the smaller vector sizes from AVX2 and SSE -// ref: https://github.com/ggml-org/llama.cpp/pull/15953#pullrequestreview-3310928344 -#if defined(__AVX512F__) && defined(__AVX512DQ__) - for (; i + 15 < n; i += 16) { - __m512 val = _mm512_sub_ps(_mm512_loadu_ps(x + i), - _mm512_set1_ps(mean)); - _mm512_storeu_ps(y + i, val); - sum += (ggml_float)_mm512_reduce_add_ps(_mm512_mul_ps(val, val)); - } -#elif defined(__AVX2__) && defined(__FMA__) - for (; i + 7 < n; i += 8) { - __m256 val = _mm256_sub_ps(_mm256_loadu_ps(x + i), - _mm256_set1_ps(mean)); - _mm256_storeu_ps(y + i, val); - val = _mm256_mul_ps(val,val); - __m128 val2 = _mm_add_ps(_mm256_extractf128_ps(val, 1), - _mm256_castps256_ps128(val)); - val2 = _mm_add_ps(val2, _mm_movehl_ps(val2, val2)); - val2 = _mm_add_ss(val2, _mm_movehdup_ps(val2)); - sum += (ggml_float)_mm_cvtss_f32(val2); - } -#elif defined(__SSE2__) - for (; i + 3 < n; i += 4) { - __m128 val = _mm_sub_ps(_mm_loadu_ps(x + i), - _mm_set1_ps(mean)); - _mm_storeu_ps(y + i, val); - val = _mm_mul_ps(val, val); -#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) - val = _mm_add_ps(val, _mm_movehl_ps(val, val)); - val = _mm_add_ss(val, _mm_movehdup_ps(val)); -#else - __m128 tmp = _mm_shuffle_ps(val, val, _MM_SHUFFLE(2, 3, 0, 1)); - val = _mm_add_ps(val, tmp); - tmp = _mm_movehl_ps(tmp, val); - val = _mm_add_ss(val, tmp); -#endif // __AVX__ || __AVX2__ || __AVX512F__ - sum += (ggml_float)_mm_cvtss_f32(val); - } -#elif defined(__ARM_NEON) && defined(__aarch64__) - for (; i + 3 < n; i += 4) { - float32x4_t val = vsubq_f32(vld1q_f32(x + i), - vdupq_n_f32(mean)); - vst1q_f32(y + i, val); - val = vmulq_f32(val, val); - sum += (ggml_float)vaddvq_f32(val); - } -#elif defined(__VXE__) || defined(__VXE2__) - for (; i + 3 < n; i += 4) { - float32x4_t val = vec_sub(vec_xl(0, x + i), vec_splats(mean)); - vec_xst(val, 0, y + i); - val = vec_mul(val, val); - sum += (ggml_float)vec_hsum_f32x4(val); - } -#elif defined(__riscv_v_intrinsic) - vfloat64m1_t vsum = __riscv_vfmv_v_f_f64m1(0, 1); - for (int vl; i < n; i += vl) { - vl = __riscv_vsetvl_e32m2(n - i); - vfloat32m2_t val = __riscv_vfsub_vf_f32m2(__riscv_vle32_v_f32m2(&x[i], vl), mean, vl); - __riscv_vse32_v_f32m2(&y[i], val, vl); - val = __riscv_vfmul_vv_f32m2(val, val, vl); - vsum = __riscv_vfwredusum_vs_f32m2_f64m1(val, vsum, vl); - } - sum = (ggml_float)__riscv_vfmv_f_s_f64m1_f64(vsum); -#endif - for (; i < n; ++i) { - float val = x[i] - mean; - y[i] = val; - val *= val; - sum += (ggml_float)val; - } - return sum/n; -} - -ggml_float ggml_vec_soft_max_f32(const int n, float * y, const float * x, float max) { - int i = 0; - ggml_float sum = 0; -#if defined(__AVX512F__) && defined(__AVX512DQ__) - for (; i + 15 < n; i += 16) { - __m512 val = ggml_v_expf(_mm512_sub_ps(_mm512_loadu_ps(x + i), - _mm512_set1_ps(max))); - _mm512_storeu_ps(y + i, val); - sum += (ggml_float)_mm512_reduce_add_ps(val); - } -#elif defined(__AVX2__) && defined(__FMA__) - for (; i + 7 < n; i += 8) { - __m256 val = ggml_v_expf(_mm256_sub_ps(_mm256_loadu_ps(x + i), - _mm256_set1_ps(max))); - _mm256_storeu_ps(y + i, val); - __m128 val2 = _mm_add_ps(_mm256_extractf128_ps(val, 1), - _mm256_castps256_ps128(val)); - val2 = _mm_add_ps(val2, _mm_movehl_ps(val2, val2)); - val2 = _mm_add_ss(val2, _mm_movehdup_ps(val2)); - sum += (ggml_float)_mm_cvtss_f32(val2); - } -#elif defined(__SSE2__) - for (; i + 3 < n; i += 4) { - __m128 val = ggml_v_expf(_mm_sub_ps(_mm_loadu_ps(x + i), - _mm_set1_ps(max))); - _mm_storeu_ps(y + i, val); -#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) - val = _mm_add_ps(val, _mm_movehl_ps(val, val)); - val = _mm_add_ss(val, _mm_movehdup_ps(val)); -#else - __m128 tmp = _mm_shuffle_ps(val, val, _MM_SHUFFLE(2, 3, 0, 1)); - val = _mm_add_ps(val, tmp); - tmp = _mm_movehl_ps(tmp, val); - val = _mm_add_ss(val, tmp); -#endif - sum += (ggml_float)_mm_cvtss_f32(val); - } -#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__) - const int vlen = svcntw(); - for (; i < n; i += vlen) { - const svbool_t pg = svwhilelt_b32_s32(i, n); - svfloat32_t val = ggml_v_expf(pg, svsub_f32_x(pg, svld1_f32(pg, x + i), - svdup_n_f32_x(pg, max))); - svst1_f32(pg, y + i, val); - sum += (ggml_float)svaddv_f32(pg, val); - } -#elif defined(__ARM_NEON) && defined(__aarch64__) - for (; i + 3 < n; i += 4) { - float32x4_t val = ggml_v_expf(vsubq_f32(vld1q_f32(x + i), - vdupq_n_f32(max))); - vst1q_f32(y + i, val); - sum += (ggml_float)vaddvq_f32(val); - } -#elif defined(__riscv_v_intrinsic) - vfloat64m1_t vsum = __riscv_vfmv_v_f_f64m1(0, 1); - for (int avl; i < n; i += avl) { - avl = __riscv_vsetvl_e32m2(n - i); - vfloat32m2_t val = ggml_v_expf_m2(__riscv_vfsub_vf_f32m2(__riscv_vle32_v_f32m2(&x[i], avl), max, avl), avl); - __riscv_vse32_v_f32m2(&y[i], val, avl); - vsum = __riscv_vfwredusum_vs_f32m2_f64m1(val, vsum, avl); - } - return (ggml_float)__riscv_vfmv_f_s_f64m1_f64(vsum); -#endif - for (; i < n; ++i) { - float val = expf(x[i] - max); - sum += (ggml_float)val; - y[i] = val; - } - return sum; -} - -ggml_float ggml_vec_log_soft_max_f32(const int n, float * y, const float * x, float max) { - // log(soft_max) = log(soft_max_i / soft_max_sum) = log(soft_max_i) - log(soft_max_sum) = (logit_i - max) - log(soft_max_i) - - int i = 0; - ggml_float sum = 0; - for (; i < n; ++i) { - float val = x[i] - max; - y[i] = val; - sum += (ggml_float)expf(val); - } - return sum = (ggml_float)logf(sum); -} diff --git a/ggml/src/ggml-cpu/vec.h b/ggml/src/ggml-cpu/vec.h deleted file mode 100644 index bcd68da9aa9e..000000000000 --- a/ggml/src/ggml-cpu/vec.h +++ /dev/null @@ -1,1588 +0,0 @@ -// Vectorized functions for fundamental operations - -#pragma once - -#include "ggml-impl.h" -#include "simd-mappings.h" -#include "ggml.h" -#include "ggml-cpu.h" - -#if defined(GGML_USE_ACCELERATE) -#include -#endif - -// floating point type used to accumulate sums -typedef double ggml_float; - -#define GGML_GELU_FP16 -#define GGML_GELU_QUICK_FP16 - -#define GGML_SOFT_MAX_UNROLL 4 -#define GGML_VEC_DOT_UNROLL 2 -#define GGML_VEC_MAD_UNROLL 32 - -#ifdef __cplusplus -extern "C" { -#endif - -// -// global data -// - -// precomputed gelu table for f16 (128 KB) -extern ggml_fp16_t ggml_table_gelu_f16[1 << 16]; - -// precomputed quick gelu table for f16 (128 KB) -extern ggml_fp16_t ggml_table_gelu_quick_f16[1 << 16]; - -// -// fundamental operations -// - -void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc); -void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t * GGML_RESTRICT x, size_t bx, ggml_bf16_t * GGML_RESTRICT y, size_t by, int nrc); -void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * GGML_RESTRICT x, size_t bx, ggml_fp16_t * GGML_RESTRICT y, size_t by, int nrc); - -void ggml_vec_silu_f32(const int n, float * y, const float * x); -ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean); //it will also center y ( y = y - mean ) -ggml_float ggml_vec_soft_max_f32(const int n, float * y, const float * x, float max); -ggml_float ggml_vec_log_soft_max_f32(const int n, float * y, const float * x, float max); - -inline static void ggml_vec_set_i8(const int n, int8_t * x, const int8_t v) { for (int i = 0; i < n; ++i) x[i] = v; } -inline static void ggml_vec_set_i16(const int n, int16_t * x, const int16_t v) { for (int i = 0; i < n; ++i) x[i] = v; } - -inline static void ggml_vec_set_i32(const int n, int32_t * x, const int32_t v) { for (int i = 0; i < n; ++i) x[i] = v; } -inline static void ggml_vec_cpy_i32(const int n, int32_t * y, const int32_t * x) { for (int i = 0; i < n; ++i) y[i] = x[i]; } - -inline static void ggml_vec_set_f16(const int n, ggml_fp16_t * x, const ggml_fp16_t v) { for (int i = 0; i < n; ++i) x[i] = v; } -inline static void ggml_vec_set_bf16(const int n, ggml_bf16_t * x, const ggml_bf16_t v) { for (int i = 0; i < n; ++i) x[i] = v; } - -inline static void ggml_vec_add_f32 (const int n, float * z, const float * x, const float * y) { - int i = 0; -#if defined(__AVX2__) - for (; i + 7 < n; i += 8) { - __m256 vx = _mm256_loadu_ps(x + i); - __m256 vy = _mm256_loadu_ps(y + i); - __m256 vz = _mm256_add_ps(vx, vy); - _mm256_storeu_ps(z + i, vz); - } -#endif - for (; i < n; ++i) { - z[i] = x[i] + y[i]; - } -} - -inline static void ggml_vec_add_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { - for (int i = 0; i < n; ++i) { - z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) + GGML_CPU_FP16_TO_FP32(y[i])); - } -} -inline static void ggml_vec_add1_f32(const int n, float * z, const float * x, const float v) { for (int i = 0; i < n; ++i) z[i] = x[i] + v; } -inline static void ggml_vec_acc_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] += x[i]; } -inline static void ggml_vec_acc1_f32(const int n, float * y, const float v) { for (int i = 0; i < n; ++i) y[i] += v; } -inline static void ggml_vec_sub_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i] = x[i] - y[i]; } -inline static void ggml_vec_sub_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { - for (int i = 0; i < n; ++i) { - z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) - GGML_CPU_FP16_TO_FP32(y[i])); - } -} -inline static void ggml_vec_set_f32 (const int n, float * x, const float v) { for (int i = 0; i < n; ++i) x[i] = v; } -inline static void ggml_vec_cpy_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i]; } -inline static void ggml_vec_neg_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = -x[i]; } -inline static void ggml_vec_neg_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(-GGML_CPU_FP16_TO_FP32(x[i])); - } -} - -inline static void ggml_vec_mul_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i] = x[i]*y[i]; } -inline static void ggml_vec_mul_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { - for (int i = 0; i < n; ++i) { - z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) * GGML_CPU_FP16_TO_FP32(y[i])); - } -} -inline static void ggml_vec_div_f32 (const int n, float * z, const float * x, const float * y) { for (int i = 0; i < n; ++i) z[i] = x[i]/y[i]; } -inline static void ggml_vec_div_f16 (const int n, ggml_fp16_t * z, const ggml_fp16_t * x, const ggml_fp16_t * y) { - for (int i = 0; i < n; ++i) { - z[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(x[i]) / GGML_CPU_FP16_TO_FP32(y[i])); - } -} - -// compute GGML_VEC_DOT_UNROLL dot products at once -// xs - x row stride in bytes -inline static void ggml_vec_dot_f16_unroll(const int n, const int xs, float * GGML_RESTRICT s, void * GGML_RESTRICT xv, ggml_fp16_t * GGML_RESTRICT y) { - ggml_float sumf[GGML_VEC_DOT_UNROLL] = { 0.0 }; - - ggml_fp16_t * GGML_RESTRICT x[GGML_VEC_DOT_UNROLL]; - - for (int i = 0; i < GGML_VEC_DOT_UNROLL; ++i) { - x[i] = (ggml_fp16_t *) ((char *) xv + i*xs); - } - -#if defined(GGML_SIMD) - #if defined(__ARM_FEATURE_SVE) - - const int sve_register_length = svcntb() * 8; - const int ggml_f16_epr = sve_register_length / 16; // running when 16 - const int ggml_f16_step = 8 * ggml_f16_epr; // choose 8 SVE registers - - int np = (n & ~(ggml_f16_step - 1)); - - svfloat16_t sum_00 = svdup_n_f16(0.0f); - svfloat16_t sum_01 = svdup_n_f16(0.0f); - svfloat16_t sum_02 = svdup_n_f16(0.0f); - svfloat16_t sum_03 = svdup_n_f16(0.0f); - - svfloat16_t sum_10 = svdup_n_f16(0.0f); - svfloat16_t sum_11 = svdup_n_f16(0.0f); - svfloat16_t sum_12 = svdup_n_f16(0.0f); - svfloat16_t sum_13 = svdup_n_f16(0.0f); - - svfloat16_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; - svfloat16_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; - - for (int i = 0; i < np; i += ggml_f16_step) { - ay1 = GGML_F16x_VEC_LOAD(y + i + 0 * ggml_f16_epr, 0); // 8 elements - - ax1 = GGML_F16x_VEC_LOAD(x[0] + i + 0*ggml_f16_epr, 0); // 8 elements - sum_00 = GGML_F16x_VEC_FMA(sum_00, ax1, ay1); // sum_00 = sum_00+ax1*ay1 - ax1 = GGML_F16x_VEC_LOAD(x[1] + i + 0*ggml_f16_epr, 0); // 8 elements - sum_10 = GGML_F16x_VEC_FMA(sum_10, ax1, ay1); - - ay2 = GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 1); // next 8 elements - - ax2 = GGML_F16x_VEC_LOAD(x[0] + i + 1*ggml_f16_epr, 1); // next 8 elements - sum_01 = GGML_F16x_VEC_FMA(sum_01, ax2, ay2); - ax2 = GGML_F16x_VEC_LOAD(x[1] + i + 1*ggml_f16_epr, 1); - sum_11 = GGML_F16x_VEC_FMA(sum_11, ax2, ay2); - - ay3 = GGML_F16x_VEC_LOAD(y + i + 2 * ggml_f16_epr, 2); - - ax3 = GGML_F16x_VEC_LOAD(x[0] + i + 2*ggml_f16_epr, 2); - sum_02 = GGML_F16x_VEC_FMA(sum_02, ax3, ay3); - ax3 = GGML_F16x_VEC_LOAD(x[1] + i + 2*ggml_f16_epr, 2); - sum_12 = GGML_F16x_VEC_FMA(sum_12, ax3, ay3); - - ay4 = GGML_F16x_VEC_LOAD(y + i + 3 * ggml_f16_epr, 3); - - ax4 = GGML_F16x_VEC_LOAD(x[0] + i + 3*ggml_f16_epr, 3); - sum_03 = GGML_F16x_VEC_FMA(sum_03, ax4, ay4); - ax4 = GGML_F16x_VEC_LOAD(x[1] + i + 3*ggml_f16_epr, 3); - sum_13 = GGML_F16x_VEC_FMA(sum_13, ax4, ay4); - - ay5 = GGML_F16x_VEC_LOAD(y + i + 4 * ggml_f16_epr, 4); - - ax5 = GGML_F16x_VEC_LOAD(x[0] + i + 4*ggml_f16_epr, 4); - - sum_00 = GGML_F16x_VEC_FMA(sum_00, ax5, ay5); - ax5 = GGML_F16x_VEC_LOAD(x[1] + i + 4*ggml_f16_epr, 4); - sum_10 = GGML_F16x_VEC_FMA(sum_10, ax5, ay5); - - ay6 = GGML_F16x_VEC_LOAD(y + i + 5 * ggml_f16_epr, 5); - - ax6 = GGML_F16x_VEC_LOAD(x[0] + i + 5*ggml_f16_epr, 5); - - sum_01 = GGML_F16x_VEC_FMA(sum_01, ax6, ay6); - ax6 = GGML_F16x_VEC_LOAD(x[1] + i + 5*ggml_f16_epr, 5); - sum_11 = GGML_F16x_VEC_FMA(sum_11, ax6, ay6); - - ay7 = GGML_F16x_VEC_LOAD(y + i + 6 * ggml_f16_epr, 6); - - ax7 = GGML_F16x_VEC_LOAD(x[0] + i + 6*ggml_f16_epr, 6); - - sum_02 = GGML_F16x_VEC_FMA(sum_02, ax7, ay7); - ax7 = GGML_F16x_VEC_LOAD(x[1] + i + 6*ggml_f16_epr, 6); - sum_12 = GGML_F16x_VEC_FMA(sum_12, ax7, ay7); - - ay8 = GGML_F16x_VEC_LOAD(y + i + 7 * ggml_f16_epr, 7); - - ax8 = GGML_F16x_VEC_LOAD(x[0] + i + 7*ggml_f16_epr, 7); - - sum_03 = GGML_F16x_VEC_FMA(sum_03, ax8, ay8); - ax8 = GGML_F16x_VEC_LOAD(x[1] + i + 7*ggml_f16_epr, 7); - sum_13 = GGML_F16x_VEC_FMA(sum_13, ax8, ay8); - } - - const int np2 = (n & ~(ggml_f16_epr - 1)); - for (int k = np; k < np2; k += ggml_f16_epr) { - svfloat16_t ry = GGML_F16x_VEC_LOAD(y + k, 0); - - svfloat16_t rx = GGML_F16x_VEC_LOAD(x[0] + k, 0); - sum_00 = GGML_F16x_VEC_FMA(sum_00, rx, ry); - rx = GGML_F16x_VEC_LOAD(x[1] + k, 0); - sum_10 = GGML_F16x_VEC_FMA(sum_10, rx, ry); - } - - if (np2 < n) { - svbool_t pg = svwhilelt_b16(np2, n); - svfloat16_t hx_0 = svld1_f16(pg, (const __fp16 *)(x[0] + np2)); - svfloat16_t hx_1 = svld1_f16(pg, (const __fp16 *)(x[1] + np2)); - svfloat16_t hy = svld1_f16(pg, (const __fp16 *)(y + np2)); - - sum_00 = svmad_f16_x(pg, hx_0, hy, sum_00); - sum_10 = svmad_f16_x(pg, hx_1, hy, sum_10); - } - GGML_F16x_VEC_REDUCE(sumf[0], sum_00, sum_01, sum_02, sum_03); - GGML_F16x_VEC_REDUCE(sumf[1], sum_10, sum_11, sum_12, sum_13); - np = n; - #elif defined(__riscv_v_intrinsic) - #if defined(__riscv_zvfh) - size_t vl = __riscv_vsetvlmax_e32m4(); - - // initialize accumulators to all zeroes - vfloat32m4_t vsum0_0 = __riscv_vfmv_v_f_f32m4(0.0f, vl); - vfloat32m4_t vsum0_1 = __riscv_vfmv_v_f_f32m4(0.0f, vl); - vfloat32m4_t vsum1_0 = __riscv_vfmv_v_f_f32m4(0.0f, vl); - vfloat32m4_t vsum1_1 = __riscv_vfmv_v_f_f32m4(0.0f, vl); - - // calculate step size - const size_t epr = __riscv_vsetvlmax_e16m2(); - const size_t step = epr * 2; - int np = (n & ~(step - 1)); - - // unroll by 2 along the row dimension - for (int i = 0; i < np; i += step) { - vfloat16m2_t ay0 = __riscv_vle16_v_f16m2((const _Float16 *)(y + i), epr); - vfloat16m2_t ax0_0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i), epr); - vfloat16m2_t ax1_0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i), epr); - vsum0_0 = __riscv_vfwmacc_vv_f32m4(vsum0_0, ax0_0, ay0, epr); - vsum1_0 = __riscv_vfwmacc_vv_f32m4(vsum1_0, ax1_0, ay0, epr); - - vfloat16m2_t ay1 = __riscv_vle16_v_f16m2((const _Float16 *)(y + i + epr), epr); - vfloat16m2_t ax0_1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i + epr), epr); - vfloat16m2_t ax1_1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i + epr), epr); - vsum0_1 = __riscv_vfwmacc_vv_f32m4(vsum0_1, ax0_1, ay1, epr); - vsum1_1 = __riscv_vfwmacc_vv_f32m4(vsum1_1, ax1_1, ay1, epr); - } - - vfloat32m4_t vsum0 = __riscv_vfadd_vv_f32m4(vsum0_0, vsum0_1, vl); - vfloat32m4_t vsum1 = __riscv_vfadd_vv_f32m4(vsum1_0, vsum1_1, vl); - - // leftovers - for (int i = np; i < n; i += vl) { - vl = __riscv_vsetvl_e16m2(n - i); - vfloat16m2_t ay = __riscv_vle16_v_f16m2((const _Float16 *)(y + i), vl); - vfloat16m2_t ax0 = __riscv_vle16_v_f16m2((const _Float16 *)(x[0] + i), vl); - vfloat16m2_t ax1 = __riscv_vle16_v_f16m2((const _Float16 *)(x[1] + i), vl); - - vsum0 = __riscv_vfwmacc_vv_f32m4(vsum0, ax0, ay, vl); - vsum1 = __riscv_vfwmacc_vv_f32m4(vsum1, ax1, ay, vl); - } - - // reduce - vl = __riscv_vsetvlmax_e32m2(); - vfloat32m2_t acc0_0 = __riscv_vfadd_vv_f32m2(__riscv_vget_v_f32m4_f32m2(vsum0, 0), - __riscv_vget_v_f32m4_f32m2(vsum0, 1), vl); - vl = __riscv_vsetvlmax_e32m1(); - vfloat32m1_t acc0_1 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(acc0_0, 0), - __riscv_vget_v_f32m2_f32m1(acc0_0, 1), vl); - vfloat32m1_t redsum0 = __riscv_vfredusum_vs_f32m1_f32m1( - acc0_1, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl); - - vl = __riscv_vsetvlmax_e32m2(); - vfloat32m2_t acc1_0 = __riscv_vfadd_vv_f32m2(__riscv_vget_v_f32m4_f32m2(vsum1, 0), - __riscv_vget_v_f32m4_f32m2(vsum1, 1), vl); - vl = __riscv_vsetvlmax_e32m1(); - vfloat32m1_t acc1_1 = __riscv_vfadd_vv_f32m1(__riscv_vget_v_f32m2_f32m1(acc1_0, 0), - __riscv_vget_v_f32m2_f32m1(acc1_0, 1), vl); - vfloat32m1_t redsum1 = __riscv_vfredusum_vs_f32m1_f32m1( - acc1_1, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl); - sumf[0] = __riscv_vfmv_f_s_f32m1_f32(redsum0); - sumf[1] = __riscv_vfmv_f_s_f32m1_f32(redsum1); - np = n; - #else - const int np = 0; - #endif - #else - const int np = (n & ~(GGML_F16_STEP - 1)); - - GGML_F16_VEC sum[GGML_VEC_DOT_UNROLL][GGML_F16_ARR] = { { GGML_F16_VEC_ZERO } }; - - GGML_F16_VEC ax[GGML_F16_ARR]; - GGML_F16_VEC ay[GGML_F16_ARR]; - - for (int i = 0; i < np; i += GGML_F16_STEP) { - for (int j = 0; j < GGML_F16_ARR; j++) { - ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); - - for (int k = 0; k < GGML_VEC_DOT_UNROLL; ++k) { - ax[j] = GGML_F16_VEC_LOAD(x[k] + i + j*GGML_F16_EPR, j); - - sum[k][j] = GGML_F16_VEC_FMA(sum[k][j], ax[j], ay[j]); - } - } - } - - // reduce sum0..sum3 to sum0 - for (int k = 0; k < GGML_VEC_DOT_UNROLL; ++k) { - GGML_F16_VEC_REDUCE(sumf[k], sum[k]); - } - #endif -#else - // scalar path - const int np = 0; -#endif - // scalar and leftovers - for (int i = np; i < n; ++i) { - for (int j = 0; j < GGML_VEC_DOT_UNROLL; ++j) { - sumf[j] += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[j][i])*GGML_CPU_FP16_TO_FP32(y[i])); - } - } - - for (int i = 0; i < GGML_VEC_DOT_UNROLL; ++i) { - s[i] = (float)sumf[i]; - } -} - -inline static void ggml_vec_mad_f32(const int n, float * GGML_RESTRICT y, const float * GGML_RESTRICT x, const float v) { -#if defined(GGML_SIMD) - #if defined(__ARM_FEATURE_SVE) - - const int sve_register_length = ggml_cpu_get_sve_cnt() * 8; - const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16 - const int ggml_f32_step = 8 * ggml_f32_epr; // choose 8 SVE registers - GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); - - const int np = (n & ~(ggml_f32_step - 1)); - svfloat32_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; - svfloat32_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; - for (int i = 0; i < np; i += ggml_f32_step) { - - ax1 = GGML_F32_VEC_LOAD(x + i); - ay1 = GGML_F32_VEC_LOAD(y + i); - ay1 = GGML_F32_VEC_FMA(ay1, ax1, vx); - - GGML_F32_VEC_STORE(y + i, ay1); - - ax2 = GGML_F32_VEC_LOAD(x + i + 1*ggml_f32_epr); - ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr); - ay2 = GGML_F32_VEC_FMA(ay2, ax2, vx); - - GGML_F32_VEC_STORE(y + i + 1*ggml_f32_epr, ay2); - - ax3 = GGML_F32_VEC_LOAD(x + i + 2*ggml_f32_epr); - ay3 = GGML_F32_VEC_LOAD(y + i + 2*ggml_f32_epr); - ay3 = GGML_F32_VEC_FMA(ay3, ax3, vx); - - GGML_F32_VEC_STORE(y + i + 2*ggml_f32_epr, ay3); - - ax4 = GGML_F32_VEC_LOAD(x + i + 3*ggml_f32_epr); - ay4 = GGML_F32_VEC_LOAD(y + i + 3*ggml_f32_epr); - ay4 = GGML_F32_VEC_FMA(ay4, ax4, vx); - - GGML_F32_VEC_STORE(y + i + 3*ggml_f32_epr, ay4); - - ax5 = GGML_F32_VEC_LOAD(x + i + 4*ggml_f32_epr); - ay5 = GGML_F32_VEC_LOAD(y + i + 4*ggml_f32_epr); - ay5 = GGML_F32_VEC_FMA(ay5, ax5, vx); - - GGML_F32_VEC_STORE(y + i + 4*ggml_f32_epr, ay5); - - ax6 = GGML_F32_VEC_LOAD(x + i + 5*ggml_f32_epr); - ay6 = GGML_F32_VEC_LOAD(y + i + 5*ggml_f32_epr); - ay6 = GGML_F32_VEC_FMA(ay6, ax6, vx); - - GGML_F32_VEC_STORE(y + i + 5*ggml_f32_epr, ay6); - - ax7 = GGML_F32_VEC_LOAD(x + i + 6*ggml_f32_epr); - ay7 = GGML_F32_VEC_LOAD(y + i + 6*ggml_f32_epr); - ay7 = GGML_F32_VEC_FMA(ay7, ax7, vx); - - GGML_F32_VEC_STORE(y + i + 6*ggml_f32_epr, ay7); - - ax8 = GGML_F32_VEC_LOAD(x + i + 7*ggml_f32_epr); - ay8 = GGML_F32_VEC_LOAD(y + i + 7*ggml_f32_epr); - ay8 = GGML_F32_VEC_FMA(ay8, ax8, vx); - - GGML_F32_VEC_STORE(y + i + 7*ggml_f32_epr, ay8); - } - // leftovers - // Since 8 unrolls are done in above loop, leftovers lie in range [0, ggml_f32_step] which is handled in below loop - const int np2 = (n & ~(ggml_f32_epr - 1)); - for (int i = np; i < np2; i += ggml_f32_epr) { - ax1 = GGML_F32_VEC_LOAD(x + i); - ay1 = GGML_F32_VEC_LOAD(y + i); - ay1 = GGML_F32_VEC_FMA(ay1, ax1, vx); - - GGML_F32_VEC_STORE(y + i, ay1); - } - // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only - if (np2 < n) { - svbool_t pg =svwhilelt_b32(np2, n); - ax1 = svld1_f32(pg, x + np2); - ay1 = svld1_f32(pg, y + np2); - ay1 = svmad_f32_m(pg, ax1, vx, ay1); - - svst1_f32(pg, y + np2, ay1); - } - #elif defined(__riscv_v_intrinsic) - for (int i = 0, avl; i < n; i += avl) { - avl = __riscv_vsetvl_e32m8(n - i); - vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[i], avl); - vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl); - vfloat32m8_t ny = __riscv_vfmadd_vf_f32m8(ax, v, ay, avl); - __riscv_vse32_v_f32m8(&y[i], ny, avl); - } - #else - const int np = (n & ~(GGML_F32_STEP - 1)); - - GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); - - GGML_F32_VEC ax[GGML_F32_ARR]; - GGML_F32_VEC ay[GGML_F32_ARR]; - - for (int i = 0; i < np; i += GGML_F32_STEP) { - for (int j = 0; j < GGML_F32_ARR; j++) { - ax[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR); - ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); - ay[j] = GGML_F32_VEC_FMA(ay[j], ax[j], vx); - - GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); - } - } - - // leftovers - for (int i = np; i < n; ++i) { - y[i] += x[i]*v; - } - #endif -#else - // scalar - for (int i = 0; i < n; ++i) { - y[i] += x[i]*v; - } -#endif -} - -inline static void ggml_vec_mad_f16(const int n, ggml_fp16_t * GGML_RESTRICT y, const ggml_fp16_t * GGML_RESTRICT x, const float v) { -#if defined(GGML_SIMD) && defined(__ARM_FEATURE_SVE) - const int sve_register_length = svcntb() * 8; - const int ggml_f16_epr = sve_register_length / 16; - const int ggml_f16_step = 8 * ggml_f16_epr; - - GGML_F16x_VEC vx = GGML_F16x_VEC_SET1(v); - - int np = (n & ~(ggml_f16_step - 1)); - - svfloat16_t ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8; - svfloat16_t ay1, ay2, ay3, ay4, ay5, ay6, ay7, ay8; - for (int i = 0; i < np; i += ggml_f16_step) { - ax1 = GGML_F16x_VEC_LOAD(x + i + 0 * ggml_f16_epr, 0); - ay1 = GGML_F16x_VEC_LOAD(y + i + 0 * ggml_f16_epr, 0); - ay1 = GGML_F16x_VEC_FMA(ay1, ax1, vx); - - GGML_F16x_VEC_STORE(y + i + 0 * ggml_f16_epr, ay1, 0); - - ax2 = GGML_F16x_VEC_LOAD(x + i + 1 * ggml_f16_epr, 1); - ay2 = GGML_F16x_VEC_LOAD(y + i + 1 * ggml_f16_epr, 1); - ay2 = GGML_F16x_VEC_FMA(ay2, ax2, vx); - - GGML_F16x_VEC_STORE(y + i + 1 * ggml_f16_epr, ay2, 1); - - ax3 = GGML_F16x_VEC_LOAD(x + i + 2 * ggml_f16_epr, 2); - ay3 = GGML_F16x_VEC_LOAD(y + i + 2 * ggml_f16_epr, 2); - ay3 = GGML_F16x_VEC_FMA(ay3, ax3, vx); - - GGML_F16x_VEC_STORE(y + i + 2 * ggml_f16_epr, ay3, 2); - - ax4 = GGML_F16x_VEC_LOAD(x + i + 3 * ggml_f16_epr, 3); - ay4 = GGML_F16x_VEC_LOAD(y + i + 3 * ggml_f16_epr, 3); - ay4 = GGML_F16x_VEC_FMA(ay4, ax4, vx); - - GGML_F16x_VEC_STORE(y + i + 3 * ggml_f16_epr, ay4, 3); - - ax5 = GGML_F16x_VEC_LOAD(x + i + 4 * ggml_f16_epr, 4); - ay5 = GGML_F16x_VEC_LOAD(y + i + 4 * ggml_f16_epr, 4); - ay5 = GGML_F16x_VEC_FMA(ay5, ax5, vx); - - GGML_F16x_VEC_STORE(y + i + 4 * ggml_f16_epr, ay5, 4); - - ax6 = GGML_F16x_VEC_LOAD(x + i + 5 * ggml_f16_epr, 5); - ay6 = GGML_F16x_VEC_LOAD(y + i + 5 * ggml_f16_epr, 5); - ay6 = GGML_F16x_VEC_FMA(ay6, ax6, vx); - - GGML_F16x_VEC_STORE(y + i + 5 * ggml_f16_epr, ay6, 5); - - ax7 = GGML_F16x_VEC_LOAD(x + i + 6 * ggml_f16_epr, 6); - ay7 = GGML_F16x_VEC_LOAD(y + i + 6 * ggml_f16_epr, 6); - ay7 = GGML_F16x_VEC_FMA(ay7, ax7, vx); - - GGML_F16x_VEC_STORE(y + i + 6 * ggml_f16_epr, ay7, 6); - - ax8 = GGML_F16x_VEC_LOAD(x + i + 7 * ggml_f16_epr, 7); - ay8 = GGML_F16x_VEC_LOAD(y + i + 7 * ggml_f16_epr, 7); - ay8 = GGML_F16x_VEC_FMA(ay8, ax8, vx); - - GGML_F16x_VEC_STORE(y + i + 7 * ggml_f16_epr, ay8, 7); - } - const int np2 = (n & ~(ggml_f16_epr - 1)); - for (int k = np; k < np2; k += ggml_f16_epr) { - svfloat16_t rx = GGML_F16x_VEC_LOAD(x + k, 0); - svfloat16_t ry = GGML_F16x_VEC_LOAD(y + k, 0); - ry = GGML_F16x_VEC_FMA(ry, rx, vx); - - GGML_F16x_VEC_STORE(y + k, ry, 0); - } - - if (np2 < n) { - svbool_t pg = svwhilelt_b16(np2, n); - svfloat16_t hx = svld1_f16(pg, (const __fp16 *)(x + np2)); - svfloat16_t hy = svld1_f16(pg, (const __fp16 *)(y + np2)); - hy = svmad_f16_x(pg, hx, vx, hy); - svst1_f16(pg, (__fp16 *)(y + np2), hy); - } - np = n; -#elif defined(__riscv_v_intrinsic) // implies __riscv_v_intrinsic - #if defined (__riscv_zvfh) - const ggml_fp16_t s = GGML_CPU_FP32_TO_FP16(v); - const _Float16 scale = *(const _Float16*)(&s); - - // calculate step size - const int epr = __riscv_vsetvlmax_e16m4(); - const int step = epr * 2; - int np = (n & ~(step - 1)); - - // unroll by 2 - for (int i = 0; i < np; i += step) { - vfloat16m4_t ax0 = __riscv_vle16_v_f16m4((const _Float16*)x + i, epr); - vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, epr); - ay0 = __riscv_vfmacc_vf_f16m4(ay0, scale, ax0, epr); - __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, epr); - __asm__ __volatile__ ("" ::: "memory"); - - vfloat16m4_t ax1 = __riscv_vle16_v_f16m4((const _Float16*)x + i + epr, epr); - vfloat16m4_t ay1 = __riscv_vle16_v_f16m4((const _Float16*)y + i + epr, epr); - ay1 = __riscv_vfmacc_vf_f16m4(ay1, scale, ax1, epr); - __riscv_vse16_v_f16m4((_Float16*)y + i + epr, ay1, epr); - __asm__ __volatile__ ("" ::: "memory"); - } - - // leftovers - int vl; - for (int i = np; i < n; i += vl) { - vl = __riscv_vsetvl_e16m4(n - i); - vfloat16m4_t ax0 = __riscv_vle16_v_f16m4((const _Float16*)x + i, vl); - vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, vl); - ay0 = __riscv_vfmacc_vf_f16m4(ay0, scale, ax0, vl); - __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, vl); - } - np = n; - #else - // fall to scalar path - const int np = 0; - #endif -#elif defined(GGML_SIMD) - const int np = (n & ~(GGML_F16_STEP - 1)); - - GGML_F16_VEC vx = GGML_F16_VEC_SET1(v); - - GGML_F16_VEC ax[GGML_F16_ARR]; - GGML_F16_VEC ay[GGML_F16_ARR]; - - for (int i = 0; i < np; i += GGML_F16_STEP) { - for (int j = 0; j < GGML_F16_ARR; j++) { - ax[j] = GGML_F16_VEC_LOAD(x + i + j*GGML_F16_EPR, j); - ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); - ay[j] = GGML_F16_VEC_FMA(ay[j], ax[j], vx); - - GGML_F16_VEC_STORE(y + i + j*GGML_F16_EPR, ay, j); - } - } -#else - // scalar path - const int np = 0; -#endif - - // scalar and leftovers - for (int i = np; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(y[i]) + GGML_CPU_FP16_TO_FP32(x[i])*v); - } -} - -// xs and vs are byte strides of x and v -inline static void ggml_vec_mad_f32_unroll(const int n, const int xs, const int vs, float * GGML_RESTRICT y, const float * GGML_RESTRICT xv, const float * GGML_RESTRICT vv) { - - const float * GGML_RESTRICT x[GGML_VEC_MAD_UNROLL]; - const float * GGML_RESTRICT v[GGML_VEC_MAD_UNROLL]; - - for (int i = 0; i < GGML_VEC_MAD_UNROLL; ++i) { - x[i] = (const float *) ((const char *) xv + i*xs); - v[i] = (const float *) ((const char *) vv + i*vs); - } - -#if defined(GGML_SIMD) - #if defined(__ARM_FEATURE_SVE) - // scalar Route to scalar implementation //TODO: Write SVE code - for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { - for (int i = 0; i < n; ++i) { - y[i] += x[k][i]*v[k][0]; - } - } - #elif defined(__riscv_v_intrinsic) - for (int i = 0, avl; i < n; i += avl) { - avl = __riscv_vsetvl_e32m8(n - i); - vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl); - for (int k = 0; k < GGML_VEC_MAD_UNROLL; k++) { - vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[k][i], avl); - ay = __riscv_vfmadd_vf_f32m8(ax, v[k][0], ay, avl); - } - __riscv_vse32_v_f32m8(&y[i], ay, avl); - } - #else - const int np = (n & ~(GGML_F32_STEP - 1)); - - GGML_F32_VEC vx[GGML_VEC_MAD_UNROLL]; - - for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { - vx[k] = GGML_F32_VEC_SET1(v[k][0]); - } - - GGML_F32_VEC ax[GGML_VEC_MAD_UNROLL][GGML_F32_ARR]; - GGML_F32_VEC ay[GGML_F32_ARR]; - - for (int i = 0; i < np; i += GGML_F32_STEP) { - for (int j = 0; j < GGML_F32_ARR; j++) { - ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); - - for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { - ax[k][j] = GGML_F32_VEC_LOAD(x[k] + i + j*GGML_F32_EPR); - ay[j] = GGML_F32_VEC_FMA(ay[j], ax[k][j], vx[k]); - } - - GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); - } - } - - // leftovers - for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { - for (int i = np; i < n; ++i) { - y[i] += x[k][i]*v[k][0]; - } - } - #endif -#else - // scalar - for (int k = 0; k < GGML_VEC_MAD_UNROLL; ++k) { - for (int i = 0; i < n; ++i) { - y[i] += x[k][i]*v[k][0]; - } - } -#endif -} - -inline static void ggml_vec_mad1_f32(const int n, float * y, const float * x, const float s, const float b) { -#if defined(GGML_USE_ACCELERATE) - vDSP_vsmsa(x, 1, &s, &b, y, 1, n); -#elif defined(GGML_SIMD) - #if defined(__ARM_FEATURE_SVE) - // scalar ; TODO: Write SVE code - for (int i = 0; i < n; ++i) { - y[i] = x[i]*s + b; - } - #elif defined(__riscv_v_intrinsic) - for (int i = 0, avl; i < n; i += avl) { - avl = __riscv_vsetvl_e32m8(n - i); - vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[i], avl); - vfloat32m8_t vb = __riscv_vfmv_v_f_f32m8(b, avl); - vfloat32m8_t ny = __riscv_vfmadd_vf_f32m8(ax, s, vb, avl); - __riscv_vse32_v_f32m8(&y[i], ny, avl); - } - #else - const int np = (n & ~(GGML_F32_STEP - 1)); - - GGML_F32_VEC vs = GGML_F32_VEC_SET1(s); - GGML_F32_VEC vb = GGML_F32_VEC_SET1(b); - - GGML_F32_VEC ay[GGML_F32_ARR]; - - for (int i = 0; i < np; i += GGML_F32_STEP) { - for (int j = 0; j < GGML_F32_ARR; j++) { - ay[j] = GGML_F32_VEC_LOAD(x + i + j*GGML_F32_EPR); - ay[j] = GGML_F32_VEC_FMA(vb, ay[j], vs); - - GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); - } - } - - // leftovers - for (int i = np; i < n; ++i) { - y[i] = x[i]*s + b; - } - #endif -#else - // scalar - for (int i = 0; i < n; ++i) { - y[i] = x[i]*s + b; - } -#endif -} - -//inline static void ggml_vec_scale_f32(const int n, float * y, const float v) { for (int i = 0; i < n; ++i) y[i] *= v; } -inline static void ggml_vec_scale_f32(const int n, float * y, const float v) { -#if defined(GGML_USE_ACCELERATE) - vDSP_vsmul(y, 1, &v, y, 1, n); -#elif defined(GGML_SIMD) - #if defined(__ARM_FEATURE_SVE) - const int sve_register_length = ggml_cpu_get_sve_cnt() * 8; - const int ggml_f32_epr = sve_register_length / 32;//8;//svcntw(); // SVE128:4, SVE256:8, SVE512:16 - const int ggml_f32_step = 2 * ggml_f32_epr; - - GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); - const int np = (n & ~(ggml_f32_step - 1)); - svfloat32_t ay1; - svfloat32_t ay2; - for (int i = 0; i < np; i += ggml_f32_step) { - ay1 = GGML_F32_VEC_LOAD(y + i); - ay1 = GGML_F32_VEC_MUL(ay1, vx); - GGML_F32_VEC_STORE(y + i, ay1); - - ay2 = GGML_F32_VEC_LOAD(y + i + 1*ggml_f32_epr); - ay2 = GGML_F32_VEC_MUL(ay2, vx); - GGML_F32_VEC_STORE(y + i + 1*ggml_f32_epr, ay2); - } - // leftovers - // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only - for (int i = np; i < n; i += ggml_f32_epr) { - svbool_t pg = svwhilelt_b32(i, n); - ay1 = svld1_f32(pg, y + i); - ay1 = svmul_f32_m(pg, ay1, vx); - svst1_f32(pg, y + i, ay1); - } - #elif defined(__riscv_v_intrinsic) - for (int i = 0, avl; i < n; i += avl) { - avl = __riscv_vsetvl_e32m8(n - i); - vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl); - vfloat32m8_t ny = __riscv_vfmul_vf_f32m8(ay, v, avl); - __riscv_vse32_v_f32m8(&y[i], ny, avl); - } - #else - const int np = (n & ~(GGML_F32_STEP - 1)); - - GGML_F32_VEC vx = GGML_F32_VEC_SET1(v); - - GGML_F32_VEC ay[GGML_F32_ARR]; - - for (int i = 0; i < np; i += GGML_F32_STEP) { - for (int j = 0; j < GGML_F32_ARR; j++) { - ay[j] = GGML_F32_VEC_LOAD(y + i + j*GGML_F32_EPR); - ay[j] = GGML_F32_VEC_MUL(ay[j], vx); - - GGML_F32_VEC_STORE(y + i + j*GGML_F32_EPR, ay[j]); - } - } - - // leftovers - for (int i = np; i < n; ++i) { - y[i] *= v; - } - #endif -#else - // scalar - for (int i = 0; i < n; ++i) { - y[i] *= v; - } -#endif -} - -inline static void ggml_vec_scale_f16(const int n, ggml_fp16_t * y, const float v) { -#if defined(GGML_SIMD) && defined(__ARM_FEATURE_SVE) - const int sve_register_length = svcntb() * 8; - const int ggml_f16_epr = sve_register_length / 16; - const int ggml_f16_step = 2 * ggml_f16_epr; - - GGML_F16x_VEC vx = GGML_F16x_VEC_SET1(v); - int np = (n & ~(ggml_f16_step - 1)); - svfloat16_t ay1, ay2; - - for (int i = 0; i < np; i += ggml_f16_step) { - ay1 = GGML_F16x_VEC_LOAD(y + i + 0*ggml_f16_epr, 0); - ay1 = GGML_F16x_VEC_MUL(ay1, vx); - GGML_F16x_VEC_STORE(y + i + 0*ggml_f16_epr, ay1, 0); - - ay2 = GGML_F16x_VEC_LOAD(y + i + 1*ggml_f16_epr, 1); - ay2 = GGML_F16x_VEC_MUL(ay2, vx); - GGML_F16x_VEC_STORE(y + i + 1*ggml_f16_epr, ay2, 1); - } - // leftovers - // maximum number of leftover elements will be less that ggmlF_16x_epr. Apply predicated svmad on available elements only - if (np < n) { - svbool_t pg = svwhilelt_b16(np, n); - svfloat16_t hy = svld1_f16(pg, (__fp16 *)(y + np)); - svfloat16_t out = svmul_f16_m(pg, hy, vx); - svst1_f16(pg, (__fp16 *)(y + np), out); - } - np = n; -#elif defined(__riscv_v_intrinsic) - #if defined(__riscv_zvfh) - const ggml_fp16_t s = GGML_CPU_FP32_TO_FP16(v); - const _Float16 scale = *(const _Float16*)(&s); - - // calculate step size - const int epr = __riscv_vsetvlmax_e16m4(); - const int step = epr * 2; - int np = (n & ~(step - 1)); - - // unroll by 2 - for (int i = 0; i < np; i += step) { - vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, epr); - ay0 = __riscv_vfmul_vf_f16m4(ay0, scale, epr); - __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, epr); - __asm__ __volatile__ ("" ::: "memory"); - - vfloat16m4_t ay1 = __riscv_vle16_v_f16m4((const _Float16*)y + i + epr, epr); - ay1 = __riscv_vfmul_vf_f16m4(ay1, scale, epr); - __riscv_vse16_v_f16m4((_Float16*)y + i + epr, ay1, epr); - __asm__ __volatile__ ("" ::: "memory"); - } - - // leftovers - int vl; - for (int i = np; i < n; i += vl) { - vl = __riscv_vsetvl_e16m4(n - i); - vfloat16m4_t ay0 = __riscv_vle16_v_f16m4((const _Float16*)y + i, vl); - ay0 = __riscv_vfmul_vf_f16m4(ay0, scale, vl); - __riscv_vse16_v_f16m4((_Float16*)y + i, ay0, vl); - } - np = n; - #else - // fall to scalar path - const int np = 0; - #endif -#elif defined(GGML_SIMD) - const int np = (n & ~(GGML_F16_STEP - 1)); - - GGML_F16_VEC vx = GGML_F16_VEC_SET1(v); - - GGML_F16_VEC ay[GGML_F16_ARR]; - - for (int i = 0; i < np; i += GGML_F16_STEP) { - for (int j = 0; j < GGML_F16_ARR; j++) { - ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j); - ay[j] = GGML_F16_VEC_MUL(ay[j], vx); - - GGML_F16_VEC_STORE(y + i + j*GGML_F16_EPR, ay, j); - } - } -#else - // scalar path - const int np = 0; -#endif - // scalar and leftovers - for (int i = np; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(y[i])*v); - } -} - -inline static void ggml_vec_norm_f32 (const int n, float * s, const float * x) { ggml_vec_dot_f32(n, s, 0, x, 0, x, 0, 1); *s = sqrtf(*s); } -inline static void ggml_vec_sqr_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i]*x[i]; } -inline static void ggml_vec_sqr_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - float v = GGML_CPU_FP16_TO_FP32(x[i]); - y[i] = GGML_CPU_FP32_TO_FP16(v*v); - } -} -inline static void ggml_vec_sqrt_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = sqrtf(x[i]); } -inline static void ggml_vec_sqrt_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(sqrtf(GGML_CPU_FP16_TO_FP32(x[i]))); - } -} -inline static void ggml_vec_log_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = logf(x[i]); } -inline static void ggml_vec_log_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(logf(GGML_CPU_FP16_TO_FP32(x[i]))); - } -} -inline static void ggml_vec_sin_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = sinf(x[i]); } -inline static void ggml_vec_sin_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(sinf(GGML_CPU_FP16_TO_FP32(x[i]))); - } -} -inline static void ggml_vec_cos_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = cosf(x[i]); } -inline static void ggml_vec_cos_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(cosf(GGML_CPU_FP16_TO_FP32(x[i]))); - } -} -inline static void ggml_vec_abs_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = fabsf(x[i]); } -inline static void ggml_vec_abs_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(fabsf(GGML_CPU_FP16_TO_FP32(x[i]))); - } -} -inline static void ggml_vec_sgn_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? 1.f : ((x[i] < 0.f) ? -1.f : 0.f); } -inline static void ggml_vec_sgn_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - float v = GGML_CPU_FP16_TO_FP32(x[i]); - y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? 1.f : ((v < 0.f) ? -1.f : 0.f)); - } -} -inline static void ggml_vec_step_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? 1.f : 0.f; } -inline static void ggml_vec_step_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16((GGML_CPU_FP16_TO_FP32(x[i]) > 0.f) ? 1.f : 0.f); - } -} -inline static void ggml_vec_tanh_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = tanhf(x[i]); } -inline static void ggml_vec_tanh_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(tanhf(GGML_CPU_FP16_TO_FP32(x[i]))); - } -} -inline static void ggml_vec_elu_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? x[i] : expm1f(x[i]); } -inline static void ggml_vec_elu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - const float v = GGML_CPU_FP16_TO_FP32(x[i]); - y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v : expm1f(v)); - } -} -inline static void ggml_vec_relu_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = (x[i] > 0.f) ? x[i] : 0.f; } -inline static void ggml_vec_relu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - float v = GGML_CPU_FP16_TO_FP32(x[i]); - y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v : 0.f); - } -} -inline static void ggml_vec_leaky_relu_f32 (const int n, float * y, const float * x, const float ns) { for (int i = 0; i < n; ++i) y[i] = ((x[i] > 0.f) ? x[i] : 0.f) + ns * ((x[i] < 0.0f) ? x[i] : 0.f); } -inline static void ggml_vec_leaky_relu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const float ns) { - for (int i = 0; i < n; ++i) { - float v = GGML_CPU_FP16_TO_FP32(x[i]); - y[i] = GGML_CPU_FP32_TO_FP16(((v > 0.f) ? v : 0.f) + ns * ((v < 0.0f) ? v : 0.f)); - } -} -inline static void ggml_vec_sigmoid_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = 1.f / (1.f + expf(-x[i])); } -inline static void ggml_vec_sigmoid_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(1.f / (1.f + expf(-GGML_CPU_FP16_TO_FP32(x[i])))); - } -} -// TODO: optimize performance -inline static void ggml_vec_hardswish_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = x[i] * fminf(1.0f, fmaxf(0.0f, (x[i] + 3.0f) / 6.0f)); } -inline static void ggml_vec_hardswish_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - float v = GGML_CPU_FP16_TO_FP32(x[i]); - y[i] = GGML_CPU_FP32_TO_FP16(v * fminf(1.0f, fmaxf(0.0f, (v + 3.0f) / 6.0f))); - } -} -inline static void ggml_vec_hardsigmoid_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = fminf(1.0f, fmaxf(0.0f, (x[i] + 3.0f) / 6.0f)); } -inline static void ggml_vec_hardsigmoid_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(fminf(1.0f, fmaxf(0.0f, (GGML_CPU_FP16_TO_FP32(x[i]) + 3.0f) / 6.0f))); - } -} -inline static void ggml_vec_exp_f32 (const int n, float * y, const float * x) { for (int i = 0; i < n; ++i) y[i] = expf(x[i]); } -inline static void ggml_vec_exp_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(expf(GGML_CPU_FP16_TO_FP32(x[i]))); - } -} - -static const float GELU_COEF_A = 0.044715f; -static const float GELU_QUICK_COEF = -1.702f; -static const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; -static const float SQRT_2_INV = 0.70710678118654752440084436210484f; - -inline static float ggml_gelu_f32(float x) { - return 0.5f*x*(1.0f + tanhf(SQRT_2_OVER_PI*x*(1.0f + GELU_COEF_A*x*x))); -} - -inline static void ggml_vec_gelu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - const uint16_t * i16 = (const uint16_t *) x; - for (int i = 0; i < n; ++i) { - y[i] = ggml_table_gelu_f16[i16[i]]; - } -} - -inline static void ggml_vec_gelu_erf_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - float xi = GGML_CPU_FP16_TO_FP32(x[i]); - float res = 0.5f*xi*(1.0f + erff(xi*SQRT_2_INV)); - y[i] = GGML_CPU_FP32_TO_FP16(res); - } -} - -#ifdef GGML_GELU_FP16 -inline static void ggml_vec_gelu_f32(const int n, float * y, const float * x) { - uint16_t t; - for (int i = 0; i < n; ++i) { - if (x[i] <= -10.0f) { - y[i] = 0.0f; - } else if (x[i] >= 10.0f) { - y[i] = x[i]; - } else { - ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); - memcpy(&t, &fp16, sizeof(uint16_t)); - y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[t]); - } - } -} -#else -inline static void ggml_vec_gelu_f32(const int n, float * y, const float * x) { - for (int i = 0; i < n; ++i) { - y[i] = ggml_gelu_f32(x[i]); - } -} -#endif - -inline static void ggml_vec_gelu_erf_f32(const int n, float * y, const float * x) { - for (int i = 0; i < n; ++i) { - float xi = x[i]; - y[i] = 0.5f*xi*(1.0f + erff(xi*SQRT_2_INV)); - } -} - -inline static float ggml_gelu_quick_f32(float x) { - return x*(1.0f/(1.0f+expf(GELU_QUICK_COEF*x))); -} - -inline static void ggml_vec_gelu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - const uint16_t * i16 = (const uint16_t *) x; - for (int i = 0; i < n; ++i) { - y[i] = ggml_table_gelu_quick_f16[i16[i]]; - } -} - -#ifdef GGML_GELU_QUICK_FP16 -inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) { - uint16_t t; - for (int i = 0; i < n; ++i) { - ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); - memcpy(&t, &fp16, sizeof(uint16_t)); - y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[t]); - } -} -#else -inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) { - for (int i = 0; i < n; ++i) { - y[i] = ggml_gelu_quick_f32(x[i]); - } -} -#endif - -// Sigmoid Linear Unit (SiLU) function -inline static float ggml_silu_f32(float x) { - return x/(1.0f + expf(-x)); -} -inline static ggml_fp16_t ggml_silu_f16(ggml_fp16_t x) { - float v = GGML_CPU_FP16_TO_FP32(x); - return GGML_CPU_FP32_TO_FP16(v/(1.0f + expf(-v))); -} - -#if __FINITE_MATH_ONLY__ -#error "some routines in ggml.c require non-finite math arithmetics -- pass -fno-finite-math-only to the compiler to fix" -#error "ref: https://github.com/ggml-org/llama.cpp/pull/7154#issuecomment-2143844461" -#endif - -/* Below function was borrowed from the GitHub repository: -https://github.com/openvinotoolkit/openvino/blob/master/src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/common.hpp */ -#if defined(__ARM_FEATURE_SVE) && defined(__aarch64__) - inline static svfloat32_t exp_ps_sve(svbool_t pg, svfloat32_t src) { - // Constants - const svfloat32_t log2_e = svdup_n_f32(1.4426950409f); - const svfloat32_t ln2 = svdup_n_f32(0.6931473921f); - const svfloat32_t half_ln2_sq = svdup_n_f32(0.2413862043f); - const svuint32_t not_mask17 = svdup_n_u32(~((1u << 17) - 1)); - const svfloat32_t one = svdup_n_f32(1.0f); - const svfloat32_t inactive1 = svdup_n_f32(0.0f); - const svint32_t inactive2 = svdup_n_s32(0); - - // Algorithm starts here - svfloat32_t t0 = svmul_f32_m(pg, src, log2_e); // y = x * log2(e) - svfloat32_t t1 = svrintm_f32_m(inactive1, pg, t0); // rount to int (float) - svint32_t t2 = svcvt_s32_f32_m(inactive2, pg, t1); // n - - t1 = svsub_f32_m(pg, t0, t1); // a = y - floor(y) - t1 = svadd_f32_m(pg, t1, one); // b = a + 1 - - svuint32_t t3 = svlsr_n_u32_m(pg, svreinterpret_u32_f32(t1), 17); // v = b >> 17 (u32) - svfloat32_t t4 = svexpa_f32(t3); // c = fexpa(v) - t4 = svscale_f32_m(pg, t4, t2); // fexpa(v) * 2^(n) - - // and_(t2.d, t1.d, not_mask17.d) - svfloat32_t t5 = svreinterpret_f32_u32(svand_u32_m(pg, svreinterpret_u32_f32(t1), not_mask17)); - t5 = svsub_f32_m(pg, t1, t5); // z - t0 = svmla_f32_m(pg, ln2, t5, half_ln2_sq); // ln2 + half_ln2_sq * z - t0 = svmla_f32_m(pg, one, t5, t0); // 1 + (ln2 * z) + (half_ln2_sq * z * z) - t0 = svmul_f32_m(pg, t0, t4); // Final result - - return t0; - } -#endif - -#if defined(__ARM_FEATURE_SVE) && defined(__aarch64__) - -inline static svfloat32_t ggml_v_expf(svbool_t pg, svfloat32_t x) { - const svfloat32_t r = svdup_n_f32_x(pg, 0x1.8p23f); - const svfloat32_t z = svmla_n_f32_x(pg, r, x, 0x1.715476p+0f); - const svfloat32_t n = svsub_f32_x(pg, z, r); - const svfloat32_t b = svmls_n_f32_x(pg, svmls_n_f32_x(pg, x, n, 0x1.62e4p-1f), n, 0x1.7f7d1cp-20f); - const svuint32_t e = svlsl_n_u32_x(pg, svreinterpret_u32_f32(z), 23); - const svfloat32_t k = svreinterpret_f32_u32(svadd_u32_x(pg, e, svreinterpret_u32_f32(svdup_n_f32_x(pg, 1)))); - const svbool_t c = svacgt_n_f32(pg, n, 126); - const svfloat32_t u = svmul_f32_x(pg, b, b); - const svfloat32_t j = svmla_f32_x(pg, - svmul_n_f32_x(pg, b, 0x1.ffffecp-1f), - svmla_f32_x(pg, svmla_f32_x(pg, svdup_n_f32_x(pg, 0x1.fffdb6p-2f), svdup_n_f32_x(pg, 0x1.555e66p-3f), b), - svmla_f32_x(pg, svdup_n_f32_x(pg, 0x1.573e2ep-5f), svdup_n_f32_x(pg, 0x1.0e4020p-7f), b), u), u); - const svuint32_t d = svdup_n_u32_z(svcmple_n_f32(pg, n, 0.0), 0x82000000); - const svfloat32_t s1 = svreinterpret_f32_u32(svadd_n_u32_x(pg, d, 0x7f000000)); - const svfloat32_t s2 = svreinterpret_f32_u32(svsub_u32_x(pg, e, d)); - return svsel_f32(svacgt_f32(pg, n, svdup_n_f32_x(pg, 192)), svmul_f32_x(pg, s1, s1), - svsel_f32(c, svmul_f32_x(pg, svmla_f32_x(pg, s2, s2, j), s1), svmla_f32_x(pg, k, k, j))); -} - -// computes silu x/(1+exp(-x)) in single precision vector -inline static svfloat32_t ggml_v_silu(svbool_t pg, svfloat32_t x) { - const svfloat32_t one = svdup_n_f32_x(pg, 1.0f); - const svfloat32_t zero = svdup_n_f32_x(pg, 0.0f); - const svfloat32_t neg_x = svsub_f32_x(pg, zero, x); - const svfloat32_t exp_neg_x = ggml_v_expf(pg, neg_x); - const svfloat32_t one_plus_exp_neg_x = svadd_f32_x(pg, one, exp_neg_x); - return svdiv_f32_x(pg, x, one_plus_exp_neg_x); -} - -#elif defined(__ARM_NEON) && defined(__aarch64__) - -// 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 -inline static float32x4_t ggml_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))); -} - -// computes silu x/(1+exp(-x)) in single precision vector -inline static float32x4_t ggml_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 = ggml_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); -} - -#elif defined(__AVX512F__) && defined(__AVX512DQ__) - -// 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 -inline static __m512 ggml_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); -} - -// computes silu x/(1+exp(-x)) in single precision vector -inline static __m512 ggml_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 = ggml_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); -} - -#elif 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 -inline static __m256 ggml_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))))); -} - -// computes silu x/(1+exp(-x)) in single precision vector -inline static __m256 ggml_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 = ggml_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); -} - -#elif defined(__SSE2__) // __AVX2__ / __ARM_NEON - -#if defined(__FMA__) -#define MADD128(x, y, z) _mm_fmadd_ps(x, y, z) -#define NMADD128(x, y, z) _mm_fnmadd_ps(x, y, z) -#else -#define MADD128(x, y, z) _mm_add_ps(_mm_mul_ps(x, y), z) -#define NMADD128(x, y, z) _mm_sub_ps(z, _mm_mul_ps(x, y)) -#endif - -// 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 -inline static __m128 ggml_v_expf(__m128 x) { - const __m128 r = _mm_set1_ps(0x1.8p23f); - const __m128 z = MADD128(x, _mm_set1_ps(0x1.715476p+0f), r); - const __m128 n = _mm_sub_ps(z, r); - const __m128 b = - NMADD128(n, _mm_set1_ps(0x1.7f7d1cp-20f), NMADD128(n, _mm_set1_ps(0x1.62e4p-1f), x)); - const __m128i e = _mm_slli_epi32(_mm_castps_si128(z), 23); - const __m128 k = _mm_castsi128_ps(_mm_add_epi32(e, _mm_castps_si128(_mm_set1_ps(1)))); - const __m128i c = - _mm_castps_si128(_mm_cmpgt_ps(_mm_andnot_ps(_mm_set1_ps(-0.f), n), _mm_set1_ps(126))); - const __m128 u = _mm_mul_ps(b, b); - const __m128 j = - MADD128(MADD128(MADD128(_mm_set1_ps(0x1.0e4020p-7f), b, _mm_set1_ps(0x1.573e2ep-5f)), u, - MADD128(_mm_set1_ps(0x1.555e66p-3f), b, _mm_set1_ps(0x1.fffdb6p-2f))), - u, _mm_mul_ps(_mm_set1_ps(0x1.ffffecp-1f), b)); - if (!_mm_movemask_epi8(c)) - return MADD128(j, k, k); - const __m128i g = _mm_and_si128(_mm_castps_si128(_mm_cmple_ps(n, _mm_setzero_ps())), - _mm_set1_epi32(0x82000000u)); - const __m128 s1 = _mm_castsi128_ps(_mm_add_epi32(g, _mm_set1_epi32(0x7f000000u))); - const __m128 s2 = _mm_castsi128_ps(_mm_sub_epi32(e, g)); - const __m128i d = - _mm_castps_si128(_mm_cmpgt_ps(_mm_andnot_ps(_mm_set1_ps(-0.f), n), _mm_set1_ps(192))); - return _mm_or_ps( - _mm_and_ps(_mm_castsi128_ps(d), _mm_mul_ps(s1, s1)), - _mm_andnot_ps(_mm_castsi128_ps(d), - _mm_or_ps(_mm_and_ps(_mm_castsi128_ps(c), _mm_mul_ps(MADD128(s2, j, s2), s1)), - _mm_andnot_ps(_mm_castsi128_ps(c), MADD128(k, j, k))))); -} - -// computes silu x/(1+exp(-x)) in single precision vector -inline static __m128 ggml_v_silu(__m128 x) { - const __m128 one = _mm_set1_ps(1); - const __m128 zero = _mm_setzero_ps(); - const __m128 neg_x = _mm_sub_ps(zero, x); - const __m128 exp_neg_x = ggml_v_expf(neg_x); - const __m128 one_plus_exp_neg_x = _mm_add_ps(one, exp_neg_x); - return _mm_div_ps(x, one_plus_exp_neg_x); -} - -#elif defined(__riscv_v_intrinsic) - -// 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 -inline static vfloat32m2_t ggml_v_expf_m2(vfloat32m2_t x, int vl) { - const vfloat32m2_t r = __riscv_vfmv_v_f_f32m2(0x1.8p23f, vl); -#ifdef __riscv_xtheadvector - // workaround for compiler bug (gcc 14.3.0: Error: unrecognized opcode `th.vmv1r.v v2,v4') - vfloat32m2_t z = __riscv_vfadd_vf_f32m2(r, 0.0f, vl); - z = __riscv_vfmacc_vf_f32m2(z, 0x1.715476p+0f, x, vl); -#else - const vfloat32m2_t z = __riscv_vfmacc_vf_f32m2(r, 0x1.715476p+0f, x, vl); -#endif - const vfloat32m2_t n = __riscv_vfsub_vv_f32m2(z, r, vl); - const vfloat32m2_t b = __riscv_vfnmsac_vf_f32m2(__riscv_vfnmsac_vf_f32m2(x, 0x1.62e4p-1f, n, vl), - 0x1.7f7d1cp-20f, n, vl); - const vuint32m2_t e = __riscv_vsll_vx_u32m2(__riscv_vreinterpret_v_f32m2_u32m2(z), 23, vl); - const vfloat32m2_t k = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vadd_vx_u32m2(e, 0x3f800000, vl)); // 1.0f - const vbool16_t c = __riscv_vmfgt_vf_f32m2_b16(__riscv_vfabs_v_f32m2(n, vl), 126.0f, vl); - const vfloat32m2_t u = __riscv_vfmul_vv_f32m2(b, b, vl); - const vfloat32m2_t j = __riscv_vfmacc_vv_f32m2( - __riscv_vfmul_vf_f32m2(b, 0x1.ffffecp-1f, vl), - __riscv_vfmacc_vv_f32m2( - __riscv_vfmacc_vf_f32m2(__riscv_vfmv_v_f_f32m2(0x1.fffdb6p-2f, vl), 0x1.555e66p-3f, b, vl), - __riscv_vfmacc_vf_f32m2(__riscv_vfmv_v_f_f32m2(0x1.573e2ep-5f, vl), 0x1.0e4020p-7f, b, vl), - u, vl), u, vl); - if (!__riscv_vcpop_m_b16(c, vl)) - return __riscv_vfmacc_vv_f32m2(k, j, k, vl); - const vbool16_t dm = __riscv_vmfle_vf_f32m2_b16(n, 0.0f, vl); - const vuint32m2_t d = __riscv_vmerge_vxm_u32m2(__riscv_vmv_v_x_u32m2(0, vl), 0x82000000, dm, vl); - const vfloat32m2_t s1 = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vadd_vx_u32m2(d, 0x7f000000, vl)); - const vfloat32m2_t s2 = __riscv_vreinterpret_v_u32m2_f32m2(__riscv_vsub_vv_u32m2(e, d, vl)); - const vfloat32m2_t r1 = __riscv_vmerge_vvm_f32m2( - __riscv_vfmacc_vv_f32m2(k, k, j, vl), - __riscv_vfmul_vv_f32m2(__riscv_vfmacc_vv_f32m2(s2, s2, j, vl), s1, vl), - c, vl); - return __riscv_vmerge_vvm_f32m2( - r1, __riscv_vfmul_vv_f32m2(s1, s1, vl), - __riscv_vmfgt_vf_f32m2_b16(__riscv_vfabs_v_f32m2(n, vl), 192.0f, vl), - vl); -} - -// computes silu x/(1+exp(-x)) in single precision vector -inline static vfloat32m2_t ggml_v_silu_m2(vfloat32m2_t x, int vl) { - const vfloat32m2_t neg_x = __riscv_vfneg_v_f32m2(x, vl); - const vfloat32m2_t exp_neg_x = ggml_v_expf_m2(neg_x, vl); - const vfloat32m2_t one_plus_exp_neg_x = __riscv_vfadd_vf_f32m2(exp_neg_x, 1.0f, vl); - return __riscv_vfdiv_vv_f32m2(x, one_plus_exp_neg_x, vl); -} - -#endif // __ARM_NEON / __AVX2__ / __SSE2__ / __riscv_v_intrinsic - -inline static void ggml_vec_silu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) { - for (int i = 0; i < n; ++i) { - y[i] = ggml_silu_f16(x[i]); - } -} - -inline static float ggml_silu_backward_f32(float x, float dy) { - const float s = 1.0f/(1.0f + expf(-x)); - return dy*s*(1.0f + x*(1.0f - s)); -} - -inline static ggml_fp16_t ggml_silu_backward_f16(ggml_fp16_t x, ggml_fp16_t dy) { - const float v = GGML_CPU_FP16_TO_FP32(x); - const float s = 1.0f/(1.0f + expf(-v)); - return GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(dy)*s*(1.0f + v*(1.0f - s))); -} - -inline static void ggml_vec_silu_backward_f32(const int n, float * dx, const float * x, const float * dy) { - for (int i = 0; i < n; ++i) { - dx[i] = ggml_silu_backward_f32(x[i], dy[i]); - } -} - -inline static void ggml_vec_silu_backward_f16(const int n, ggml_fp16_t * dx, const ggml_fp16_t * x, const ggml_fp16_t * dy) { - for (int i = 0; i < n; ++i) { - dx[i] = ggml_silu_backward_f16(x[i], dy[i]); - } -} - -inline static void ggml_vec_reglu_f32 (const int n, float * y, const float * x, const float * g) { - for (int i = 0; i < n; ++i) { - y[i] = (x[i] > 0.f) ? x[i] * g[i] : 0.f; - } -} - -inline static void ggml_vec_reglu_f16 (const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { - for (int i = 0; i < n; ++i) { - float v = GGML_CPU_FP16_TO_FP32(x[i]); - y[i] = GGML_CPU_FP32_TO_FP16((v > 0.f) ? v * GGML_CPU_FP16_TO_FP32(g[i]) : 0.f); - } -} - -#ifdef GGML_GELU_FP16 -inline static void ggml_vec_geglu_f32(const int n, float * y, const float * x, const float * g) { - uint16_t t; - for (int i = 0; i < n; ++i) { - if (x[i] <= -10.0f) { - y[i] = 0.0f; - } else if (x[i] >= 10.0f) { - y[i] = x[i] * g[i]; - } else { - ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); - memcpy(&t, &fp16, sizeof(uint16_t)); - y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[t]) * g[i]; - } - } -} -#else -inline static void ggml_vec_geglu_f32(const int n, float * y, const float * x, const float * g) { - for (int i = 0; i < n; ++i) { - y[i] = ggml_gelu_f32(x[i]) * g[i]; - } -} -#endif - -inline static void ggml_vec_geglu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { - const uint16_t * i16 = (const uint16_t *) x; - for (int i = 0; i < n; ++i) { - float v = GGML_CPU_FP16_TO_FP32(g[i]); - y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(ggml_table_gelu_f16[i16[i]]) * v); - } -} - -void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float * g); - -inline static void ggml_vec_swiglu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { - for (int i = 0; i < n; ++i) { - float xi = GGML_CPU_FP16_TO_FP32(x[i]); - float gi = GGML_CPU_FP16_TO_FP32(g[i]); - y[i] = GGML_CPU_FP32_TO_FP16((xi/(1.0f + expf(-xi))) * gi); - } -} - -inline static void ggml_vec_geglu_erf_f32(const int n, float * y, const float * x, const float * g) { - for (int i = 0; i < n; ++i) { - float xi = x[i]; - y[i] = 0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * g[i]; - } -} - -inline static void ggml_vec_geglu_erf_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { - for (int i = 0; i < n; ++i) { - float xi = GGML_CPU_FP16_TO_FP32(x[i]); - float gi = GGML_CPU_FP16_TO_FP32(g[i]); - y[i] = GGML_CPU_FP32_TO_FP16(0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * gi); - } -} - -#ifdef GGML_GELU_QUICK_FP16 -inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) { - uint16_t t; - for (int i = 0; i < n; ++i) { - ggml_fp16_t fp16 = GGML_CPU_FP32_TO_FP16(x[i]); - memcpy(&t, &fp16, sizeof(uint16_t)); - y[i] = GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[t]) * g[i]; - } -} -#else -inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) { - for (int i = 0; i < n; ++i) { - y[i] = ggml_gelu_quick_f32(x[i]) * g[i]; - } -} -#endif - -inline static void ggml_vec_geglu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) { - const uint16_t * i16 = (const uint16_t *) x; - for (int i = 0; i < n; ++i) { - float v = GGML_CPU_FP16_TO_FP32(g[i]); - y[i] = GGML_CPU_FP32_TO_FP16(GGML_CPU_FP16_TO_FP32(ggml_table_gelu_quick_f16[i16[i]]) * v); - } -} - -inline static void ggml_vec_sum_f32(const int n, float * s, const float * x) { -#ifndef GGML_USE_ACCELERATE - ggml_float sum = 0.0; - for (int i = 0; i < n; ++i) { - sum += (ggml_float)x[i]; - } - *s = (float)sum; -#else - vDSP_sve(x, 1, s, n); -#endif -} - -inline static void ggml_vec_cumsum_f32(const int n, float * y, const float * x) { - for (int i = 0; i < n; ++i) { - if (i == 0) { - y[i] = x[i]; - } else { - y[i] = y[i - 1] + x[i]; - } - } -} - -inline static void ggml_vec_sum_f32_ggf(const int n, ggml_float * s, const float * x) { - ggml_float sum = 0.0; - for (int i = 0; i < n; ++i) { - sum += (ggml_float)x[i]; - } - *s = sum; -} - -inline static void ggml_vec_sum_f16_ggf(const int n, float * s, const ggml_fp16_t * x) { - float sum = 0.0f; - for (int i = 0; i < n; ++i) { - sum += GGML_CPU_FP16_TO_FP32(x[i]); - } - *s = sum; -} - -inline static void ggml_vec_sum_bf16_ggf(const int n, float * s, const ggml_bf16_t * x) { - float sum = 0.0f; - for (int i = 0; i < n; ++i) { - sum += GGML_BF16_TO_FP32(x[i]); - } - *s = sum; -} - -inline static void ggml_vec_max_f32(const int n, float * s, const float * x) { -#ifndef GGML_USE_ACCELERATE - float max = -INFINITY; - for (int i = 0; i < n; ++i) { - max = MAX(max, x[i]); - } - *s = max; -#else - vDSP_maxv(x, 1, s, n); -#endif -} - -inline static void ggml_vec_norm_inv_f32(const int n, float * s, const float * x) { - ggml_vec_norm_f32(n, s, x); - *s = 1.f/(*s); -} - -inline static void ggml_vec_argmax_f32(const int n, int * s, const float * x) { - float max = -INFINITY; - int idx = 0; - for (int i = 0; i < n; ++i) { - max = MAX(max, x[i]); - if (max == x[i]) { idx = i; } - } - *s = idx; -} - -#ifdef __cplusplus -} -#endif diff --git a/ggml/src/ggml-cuda/CMakeLists.txt b/ggml/src/ggml-cuda/CMakeLists.txt index b54d4a6b1079..07332e5ae3a2 100644 --- a/ggml/src/ggml-cuda/CMakeLists.txt +++ b/ggml/src/ggml-cuda/CMakeLists.txt @@ -98,31 +98,33 @@ if (CUDAToolkit_FOUND) endif() message(STATUS "Using CMAKE_CUDA_ARCHITECTURES=${CMAKE_CUDA_ARCHITECTURES} CMAKE_CUDA_ARCHITECTURES_NATIVE=${CMAKE_CUDA_ARCHITECTURES_NATIVE}") - file(GLOB GGML_HEADERS_CUDA "*.cuh") - list(APPEND GGML_HEADERS_CUDA "../../include/ggml-cuda.h") + file(GLOB GGML_HEADERS_CUDA "*.cuh.inc") + list(APPEND GGML_HEADERS_CUDA "../../include/ggml-cuda.h.inc") - file(GLOB GGML_SOURCES_CUDA "*.cu") - file(GLOB SRCS "template-instances/fattn-tile*.cu") + file(GLOB GGML_SOURCES_CUDA "*.cu.inc") + file(GLOB SRCS "template-instances/fattn-tile*.cu.inc") list(APPEND GGML_SOURCES_CUDA ${SRCS}) - file(GLOB SRCS "template-instances/fattn-mma*.cu") + file(GLOB SRCS "template-instances/fattn-mma*.cu.inc") list(APPEND GGML_SOURCES_CUDA ${SRCS}) - file(GLOB SRCS "template-instances/mmq*.cu") + file(GLOB SRCS "template-instances/mmq*.cu.inc") list(APPEND GGML_SOURCES_CUDA ${SRCS}) - file(GLOB SRCS "template-instances/mmf*.cu") + file(GLOB SRCS "template-instances/mmf*.cu.inc") list(APPEND GGML_SOURCES_CUDA ${SRCS}) if (GGML_CUDA_FA_ALL_QUANTS) - file(GLOB SRCS "template-instances/fattn-vec*.cu") + file(GLOB SRCS "template-instances/fattn-vec*.cu.inc") list(APPEND GGML_SOURCES_CUDA ${SRCS}) add_compile_definitions(GGML_CUDA_FA_ALL_QUANTS) else() list(APPEND GGML_SOURCES_CUDA - template-instances/fattn-vec-instance-f16-f16.cu - template-instances/fattn-vec-instance-q4_0-q4_0.cu - template-instances/fattn-vec-instance-q8_0-q8_0.cu - template-instances/fattn-vec-instance-bf16-bf16.cu) + template-instances/fattn-vec-instance-f16-f16.cu.inc + template-instances/fattn-vec-instance-q4_0-q4_0.cu.inc + template-instances/fattn-vec-instance-q8_0-q8_0.cu.inc + template-instances/fattn-vec-instance-bf16-bf16.cu.inc) endif() + set_source_files_properties(${GGML_SOURCES_CUDA} PROPERTIES LANGUAGE CUDA) + ggml_add_backend_library(ggml-cuda ${GGML_HEADERS_CUDA} ${GGML_SOURCES_CUDA} diff --git a/ggml/src/ggml-cuda/acc.cu b/ggml/src/ggml-cuda/acc.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/acc.cu rename to ggml/src/ggml-cuda/acc.cu.inc index e084607c029a..7652ae3cf3e8 100644 --- a/ggml/src/ggml-cuda/acc.cu +++ b/ggml/src/ggml-cuda/acc.cu.inc @@ -1,4 +1,4 @@ -#include "acc.cuh" +#include "acc.cuh.inc" static __global__ void acc_f32(const float * x, const float * y, float * dst, const int64_t ne, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t ne13, diff --git a/ggml/src/ggml-cuda/acc.cuh b/ggml/src/ggml-cuda/acc.cuh.inc similarity index 80% rename from ggml/src/ggml-cuda/acc.cuh rename to ggml/src/ggml-cuda/acc.cuh.inc index 1168ea1b2e87..7ecb06103816 100644 --- a/ggml/src/ggml-cuda/acc.cuh +++ b/ggml/src/ggml-cuda/acc.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_ACC_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/add-id.cu b/ggml/src/ggml-cuda/add-id.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/add-id.cu rename to ggml/src/ggml-cuda/add-id.cu.inc index 8d9cf692b4b5..26758aa911c9 100644 --- a/ggml/src/ggml-cuda/add-id.cu +++ b/ggml/src/ggml-cuda/add-id.cu.inc @@ -1,4 +1,4 @@ -#include "add-id.cuh" +#include "add-id.cuh.inc" static __global__ void add_id_kernel( const float * src0, const float * src1, const int32_t * src2, float * dst, diff --git a/ggml/src/ggml-cuda/add-id.cuh b/ggml/src/ggml-cuda/add-id.cuh.inc similarity index 75% rename from ggml/src/ggml-cuda/add-id.cuh rename to ggml/src/ggml-cuda/add-id.cuh.inc index 30b1721ac324..b0e72a8561b0 100644 --- a/ggml/src/ggml-cuda/add-id.cuh +++ b/ggml/src/ggml-cuda/add-id.cuh.inc @@ -1,3 +1,3 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_op_add_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/arange.cu b/ggml/src/ggml-cuda/arange.cu.inc similarity index 97% rename from ggml/src/ggml-cuda/arange.cu rename to ggml/src/ggml-cuda/arange.cu.inc index b5e495a24622..d7b4427ea07a 100644 --- a/ggml/src/ggml-cuda/arange.cu +++ b/ggml/src/ggml-cuda/arange.cu.inc @@ -1,4 +1,4 @@ -#include "arange.cuh" +#include "arange.cuh.inc" static __global__ void arange_f32(float * dst, const int ne0, const float start, const float step) { // blockIDx.x: idx of ne0 / BLOCK_SIZE diff --git a/ggml/src/ggml-cuda/arange.cuh b/ggml/src/ggml-cuda/arange.cuh.inc similarity index 81% rename from ggml/src/ggml-cuda/arange.cuh rename to ggml/src/ggml-cuda/arange.cuh.inc index 41e74fdfc203..c96b2474b814 100644 --- a/ggml/src/ggml-cuda/arange.cuh +++ b/ggml/src/ggml-cuda/arange.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_ARANGE_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/argmax.cu b/ggml/src/ggml-cuda/argmax.cu.inc similarity index 97% rename from ggml/src/ggml-cuda/argmax.cu rename to ggml/src/ggml-cuda/argmax.cu.inc index 51967c667cfd..e14dff5378ad 100644 --- a/ggml/src/ggml-cuda/argmax.cu +++ b/ggml/src/ggml-cuda/argmax.cu.inc @@ -1,9 +1,9 @@ #include #include -#include "argmax.cuh" -#include "common.cuh" -#include "sum.cuh" +#include "argmax.cuh.inc" +#include "common.cuh.inc" +#include "sum.cuh.inc" static __global__ void argmax_f32(const float * __restrict__ x, int32_t * __restrict__ dst, const int64_t ncols) { const int64_t row = blockIdx.x; diff --git a/ggml/src/ggml-cuda/argmax.cuh b/ggml/src/ggml-cuda/argmax.cuh.inc similarity index 74% rename from ggml/src/ggml-cuda/argmax.cuh rename to ggml/src/ggml-cuda/argmax.cuh.inc index 5b7223adc6ba..53c7b7d763fd 100644 --- a/ggml/src/ggml-cuda/argmax.cuh +++ b/ggml/src/ggml-cuda/argmax.cuh.inc @@ -1,3 +1,3 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_argmax(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/argsort.cu b/ggml/src/ggml-cuda/argsort.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/argsort.cu rename to ggml/src/ggml-cuda/argsort.cu.inc index 0f3f017b534f..e126cb479fe9 100644 --- a/ggml/src/ggml-cuda/argsort.cu +++ b/ggml/src/ggml-cuda/argsort.cu.inc @@ -1,4 +1,4 @@ -#include "argsort.cuh" +#include "argsort.cuh.inc" #ifdef GGML_CUDA_USE_CUB # include diff --git a/ggml/src/ggml-cuda/argsort.cuh b/ggml/src/ggml-cuda/argsort.cuh.inc similarity index 97% rename from ggml/src/ggml-cuda/argsort.cuh rename to ggml/src/ggml-cuda/argsort.cuh.inc index 22b7306f2020..a004315b0418 100644 --- a/ggml/src/ggml-cuda/argsort.cuh +++ b/ggml/src/ggml-cuda/argsort.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/binbcast.cu b/ggml/src/ggml-cuda/binbcast.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/binbcast.cu rename to ggml/src/ggml-cuda/binbcast.cu.inc index adb4d5f0cb90..4385e05ffe7f 100644 --- a/ggml/src/ggml-cuda/binbcast.cu +++ b/ggml/src/ggml-cuda/binbcast.cu.inc @@ -1,4 +1,4 @@ -#include "binbcast.cuh" +#include "binbcast.cuh.inc" #include #include diff --git a/ggml/src/ggml-cuda/binbcast.cuh b/ggml/src/ggml-cuda/binbcast.cuh.inc similarity index 96% rename from ggml/src/ggml-cuda/binbcast.cuh rename to ggml/src/ggml-cuda/binbcast.cuh.inc index 12624785b444..0ea2ad6efd89 100644 --- a/ggml/src/ggml-cuda/binbcast.cuh +++ b/ggml/src/ggml-cuda/binbcast.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_op_repeat(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/clamp.cu b/ggml/src/ggml-cuda/clamp.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/clamp.cu rename to ggml/src/ggml-cuda/clamp.cu.inc index fe415e7f78dd..ac7b17ff72ad 100644 --- a/ggml/src/ggml-cuda/clamp.cu +++ b/ggml/src/ggml-cuda/clamp.cu.inc @@ -1,4 +1,4 @@ -#include "clamp.cuh" +#include "clamp.cuh.inc" static __device__ __forceinline__ float op_clamp(float x, float min, float max) { return fminf(fmaxf(x, min), max); diff --git a/ggml/src/ggml-cuda/clamp.cuh b/ggml/src/ggml-cuda/clamp.cuh.inc similarity index 81% rename from ggml/src/ggml-cuda/clamp.cuh rename to ggml/src/ggml-cuda/clamp.cuh.inc index 7f9559dd17eb..38d2ce53192c 100644 --- a/ggml/src/ggml-cuda/clamp.cuh +++ b/ggml/src/ggml-cuda/clamp.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_CLAMP_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh.inc similarity index 99% rename from ggml/src/ggml-cuda/common.cuh rename to ggml/src/ggml-cuda/common.cuh.inc index 10817505d9f0..63b8076efe62 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh.inc @@ -1,8 +1,8 @@ #pragma once -#include "ggml.h" -#include "ggml-impl.h" -#include "ggml-cuda.h" +#include "ggml.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-cuda.h.inc" #include #include @@ -18,7 +18,7 @@ #define GGML_COMMON_IMPL_MUSA #endif #endif -#include "ggml-common.h" +#include "ggml-common-defs.inc" #include #include @@ -30,11 +30,11 @@ #include #if defined(GGML_USE_HIP) -#include "vendors/hip.h" +#include "vendors/hip.h.inc" #elif defined(GGML_USE_MUSA) -#include "vendors/musa.h" +#include "vendors/musa.h.inc" #else -#include "vendors/cuda.h" +#include "vendors/cuda.h.inc" #endif // defined(GGML_USE_HIP) #define STRINGIZE_IMPL(...) #__VA_ARGS__ @@ -367,7 +367,7 @@ static __device__ void no_device_code( file_name, line, function_name, arch); GGML_UNUSED(arch_list); #else - printf("%s:%d: ERROR: CUDA kernel %s has no device code compatible with CUDA arch %d. ggml-cuda.cu was compiled for: %s\n", + printf("%s:%d: ERROR: CUDA kernel %s has no device code compatible with CUDA arch %d. ggml-cuda.cu.inc was compiled for: %s\n", file_name, line, function_name, arch, arch_list); #endif // defined(GGML_USE_HIP) __trap(); diff --git a/ggml/src/ggml-cuda/concat.cu b/ggml/src/ggml-cuda/concat.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/concat.cu rename to ggml/src/ggml-cuda/concat.cu.inc index 102f944f924a..0e53424635b0 100644 --- a/ggml/src/ggml-cuda/concat.cu +++ b/ggml/src/ggml-cuda/concat.cu.inc @@ -1,4 +1,4 @@ -#include "concat.cuh" +#include "concat.cuh.inc" // contiguous kernels template diff --git a/ggml/src/ggml-cuda/concat.cuh b/ggml/src/ggml-cuda/concat.cuh.inc similarity index 81% rename from ggml/src/ggml-cuda/concat.cuh rename to ggml/src/ggml-cuda/concat.cuh.inc index aa506a05f2cc..150ffd9914ef 100644 --- a/ggml/src/ggml-cuda/concat.cuh +++ b/ggml/src/ggml-cuda/concat.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_CONCAT_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/conv-transpose-1d.cu b/ggml/src/ggml-cuda/conv-transpose-1d.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/conv-transpose-1d.cu rename to ggml/src/ggml-cuda/conv-transpose-1d.cu.inc index 8418ba667318..a2155cbda7ff 100644 --- a/ggml/src/ggml-cuda/conv-transpose-1d.cu +++ b/ggml/src/ggml-cuda/conv-transpose-1d.cu.inc @@ -1,4 +1,4 @@ -#include "conv-transpose-1d.cuh" +#include "conv-transpose-1d.cuh.inc" static __global__ void conv_transpose_1d_kernel( const int s0, const int p0, const int d0, const int output_size, diff --git a/ggml/src/ggml-cuda/conv-transpose-1d.cuh b/ggml/src/ggml-cuda/conv-transpose-1d.cuh.inc similarity index 83% rename from ggml/src/ggml-cuda/conv-transpose-1d.cuh rename to ggml/src/ggml-cuda/conv-transpose-1d.cuh.inc index 6c2cf666b68d..53fe0dbd0dc5 100644 --- a/ggml/src/ggml-cuda/conv-transpose-1d.cuh +++ b/ggml/src/ggml-cuda/conv-transpose-1d.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_CONV_TRANPOSE_1D_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/conv2d-dw.cu b/ggml/src/ggml-cuda/conv2d-dw.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/conv2d-dw.cu rename to ggml/src/ggml-cuda/conv2d-dw.cu.inc index 7583233b1b7c..836088943549 100644 --- a/ggml/src/ggml-cuda/conv2d-dw.cu +++ b/ggml/src/ggml-cuda/conv2d-dw.cu.inc @@ -1,4 +1,4 @@ -#include "conv2d-dw.cuh" +#include "conv2d-dw.cuh.inc" struct conv_params { int in_w, in_h; diff --git a/ggml/src/ggml-cuda/conv2d-dw.cuh b/ggml/src/ggml-cuda/conv2d-dw.cuh.inc similarity index 83% rename from ggml/src/ggml-cuda/conv2d-dw.cuh rename to ggml/src/ggml-cuda/conv2d-dw.cuh.inc index b5d5a69d345c..b3fa1c720036 100644 --- a/ggml/src/ggml-cuda/conv2d-dw.cuh +++ b/ggml/src/ggml-cuda/conv2d-dw.cuh.inc @@ -1,5 +1,5 @@ #pragma once -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_CONV2D_DW_BLOCK_SIZE 256 void ggml_cuda_op_conv2d_dw(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/conv2d-transpose.cu b/ggml/src/ggml-cuda/conv2d-transpose.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/conv2d-transpose.cu rename to ggml/src/ggml-cuda/conv2d-transpose.cu.inc index 6cbd6f879e6f..ff1fe13030ad 100644 --- a/ggml/src/ggml-cuda/conv2d-transpose.cu +++ b/ggml/src/ggml-cuda/conv2d-transpose.cu.inc @@ -1,5 +1,5 @@ -#include "conv2d-transpose.cuh" -#include "convert.cuh" +#include "conv2d-transpose.cuh.inc" +#include "convert.cuh.inc" template static __global__ void conv2d_transpose_kernel(const float * __restrict__ input, diff --git a/ggml/src/ggml-cuda/conv2d-transpose.cuh b/ggml/src/ggml-cuda/conv2d-transpose.cuh.inc similarity index 83% rename from ggml/src/ggml-cuda/conv2d-transpose.cuh rename to ggml/src/ggml-cuda/conv2d-transpose.cuh.inc index 72889c5f0fa8..94f5395a6404 100644 --- a/ggml/src/ggml-cuda/conv2d-transpose.cuh +++ b/ggml/src/ggml-cuda/conv2d-transpose.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_CONV2D_TRANSPOSE_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/conv2d.cu b/ggml/src/ggml-cuda/conv2d.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/conv2d.cu rename to ggml/src/ggml-cuda/conv2d.cu.inc index 142dd66903aa..1992adde79b2 100644 --- a/ggml/src/ggml-cuda/conv2d.cu +++ b/ggml/src/ggml-cuda/conv2d.cu.inc @@ -1,5 +1,5 @@ -#include "conv2d.cuh" -#include "convert.cuh" +#include "conv2d.cuh.inc" +#include "convert.cuh.inc" struct conv_params { const int64_t IW, IH; diff --git a/ggml/src/ggml-cuda/conv2d.cuh b/ggml/src/ggml-cuda/conv2d.cuh.inc similarity index 83% rename from ggml/src/ggml-cuda/conv2d.cuh rename to ggml/src/ggml-cuda/conv2d.cuh.inc index ce4802c7ed79..25b8efe3ea55 100644 --- a/ggml/src/ggml-cuda/conv2d.cuh +++ b/ggml/src/ggml-cuda/conv2d.cuh.inc @@ -1,5 +1,5 @@ #pragma once -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_CONV2D_BLOCK_SIZE 256 void ggml_cuda_op_conv2d(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/convert.cu b/ggml/src/ggml-cuda/convert.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/convert.cu rename to ggml/src/ggml-cuda/convert.cu.inc index 61630a35a29b..3116f7469bd4 100644 --- a/ggml/src/ggml-cuda/convert.cu +++ b/ggml/src/ggml-cuda/convert.cu.inc @@ -1,5 +1,5 @@ -#include "convert.cuh" -#include "dequantize.cuh" +#include "convert.cuh.inc" +#include "dequantize.cuh.inc" #include diff --git a/ggml/src/ggml-cuda/convert.cuh b/ggml/src/ggml-cuda/convert.cuh.inc similarity index 98% rename from ggml/src/ggml-cuda/convert.cuh rename to ggml/src/ggml-cuda/convert.cuh.inc index f5d37c7b9987..5f28702d5143 100644 --- a/ggml/src/ggml-cuda/convert.cuh +++ b/ggml/src/ggml-cuda/convert.cuh.inc @@ -1,5 +1,5 @@ #pragma once -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_DEQUANTIZE_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/count-equal.cu b/ggml/src/ggml-cuda/count-equal.cu.inc similarity index 97% rename from ggml/src/ggml-cuda/count-equal.cu rename to ggml/src/ggml-cuda/count-equal.cu.inc index 08898115daed..0da06f355cff 100644 --- a/ggml/src/ggml-cuda/count-equal.cu +++ b/ggml/src/ggml-cuda/count-equal.cu.inc @@ -1,5 +1,5 @@ -#include "common.cuh" -#include "count-equal.cuh" +#include "common.cuh.inc" +#include "count-equal.cuh.inc" #include diff --git a/ggml/src/ggml-cuda/count-equal.cuh b/ggml/src/ggml-cuda/count-equal.cuh.inc similarity index 82% rename from ggml/src/ggml-cuda/count-equal.cuh rename to ggml/src/ggml-cuda/count-equal.cuh.inc index 8467da79e0c7..1b369a67ec5b 100644 --- a/ggml/src/ggml-cuda/count-equal.cuh +++ b/ggml/src/ggml-cuda/count-equal.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_COUNT_EQUAL_CHUNK_SIZE 128 diff --git a/ggml/src/ggml-cuda/cp-async.cuh b/ggml/src/ggml-cuda/cp-async.cuh.inc similarity index 98% rename from ggml/src/ggml-cuda/cp-async.cuh rename to ggml/src/ggml-cuda/cp-async.cuh.inc index 63d0c482ff72..f3f3cb3b4109 100644 --- a/ggml/src/ggml-cuda/cp-async.cuh +++ b/ggml/src/ggml-cuda/cp-async.cuh.inc @@ -1,6 +1,6 @@ // Simplified API for asynchronous data loading. -#include "common.cuh" +#include "common.cuh.inc" static __device__ __forceinline__ unsigned int ggml_cuda_cvta_generic_to_shared(void * generic_ptr) { diff --git a/ggml/src/ggml-cuda/cpy-utils.cuh b/ggml/src/ggml-cuda/cpy-utils.cuh.inc similarity index 98% rename from ggml/src/ggml-cuda/cpy-utils.cuh rename to ggml/src/ggml-cuda/cpy-utils.cuh.inc index 7697c292dd6f..c72624f20997 100644 --- a/ggml/src/ggml-cuda/cpy-utils.cuh +++ b/ggml/src/ggml-cuda/cpy-utils.cuh.inc @@ -1,7 +1,7 @@ #pragma once -#include "ggml-common.h" -#include "convert.cuh" +#include "ggml-common-defs.inc" +#include "convert.cuh.inc" static __device__ __forceinline__ int best_index_int8(int n, const int8_t * val, float x) { if (x <= val[0]) return 0; @@ -186,7 +186,7 @@ static __device__ void quantize_f32_iq4_nl_block(const float * __restrict__ x, b y->d = sumq2 > 0 ? sumqx/sumq2 : d; } -// Wrapper functions for cpy.cu compatibility +// Wrapper functions for cpy.cu.inc compatibility static __device__ void cpy_blck_f32_q4_0(const char * cxi, char * cdsti) { quantize_f32_q4_0_block((const float *)cxi, (block_q4_0 *)cdsti); } diff --git a/ggml/src/ggml-cuda/cpy.cu b/ggml/src/ggml-cuda/cpy.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/cpy.cu rename to ggml/src/ggml-cuda/cpy.cu.inc index d208acf2d5f0..2783298f0c4f 100644 --- a/ggml/src/ggml-cuda/cpy.cu +++ b/ggml/src/ggml-cuda/cpy.cu.inc @@ -1,8 +1,8 @@ -#include "cpy.cuh" -#include "dequantize.cuh" -#include "cpy-utils.cuh" +#include "cpy.cuh.inc" +#include "dequantize.cuh.inc" +#include "cpy-utils.cuh.inc" #if defined(GGML_USE_MUSA) && defined(GGML_MUSA_MUDNN_COPY) -#include "ggml-musa/mudnn.cuh" +#include "ggml-musa/mudnn.cuh.inc" #endif // GGML_USE_MUSA && GGML_MUSA_MUDNN_COPY typedef void (*cpy_kernel_t)(const char * cx, char * cdst); diff --git a/ggml/src/ggml-cuda/cpy.cuh b/ggml/src/ggml-cuda/cpy.cuh.inc similarity index 88% rename from ggml/src/ggml-cuda/cpy.cuh rename to ggml/src/ggml-cuda/cpy.cuh.inc index a7a87d8fcfb7..74c689eb110d 100644 --- a/ggml/src/ggml-cuda/cpy.cuh +++ b/ggml/src/ggml-cuda/cpy.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_CPY_BLOCK_SIZE 64 diff --git a/ggml/src/ggml-cuda/cross-entropy-loss.cu b/ggml/src/ggml-cuda/cross-entropy-loss.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/cross-entropy-loss.cu rename to ggml/src/ggml-cuda/cross-entropy-loss.cu.inc index 0c8b0819724e..6ff589e135d0 100644 --- a/ggml/src/ggml-cuda/cross-entropy-loss.cu +++ b/ggml/src/ggml-cuda/cross-entropy-loss.cu.inc @@ -1,6 +1,6 @@ -#include "common.cuh" -#include "cross-entropy-loss.cuh" -#include "sum.cuh" +#include "common.cuh.inc" +#include "cross-entropy-loss.cuh.inc" +#include "sum.cuh.inc" #include #include diff --git a/ggml/src/ggml-cuda/cross-entropy-loss.cuh b/ggml/src/ggml-cuda/cross-entropy-loss.cuh.inc similarity index 89% rename from ggml/src/ggml-cuda/cross-entropy-loss.cuh rename to ggml/src/ggml-cuda/cross-entropy-loss.cuh.inc index 9ec7152ff451..ac231d332197 100644 --- a/ggml/src/ggml-cuda/cross-entropy-loss.cuh +++ b/ggml/src/ggml-cuda/cross-entropy-loss.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_CROSS_ENTROPY_LOSS_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/cumsum.cu b/ggml/src/ggml-cuda/cumsum.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/cumsum.cu rename to ggml/src/ggml-cuda/cumsum.cu.inc index def9c32955f2..dbc15e07e2bc 100644 --- a/ggml/src/ggml-cuda/cumsum.cu +++ b/ggml/src/ggml-cuda/cumsum.cu.inc @@ -1,8 +1,8 @@ #include -#include "cumsum.cuh" -#include "convert.cuh" -#include "ggml-cuda/common.cuh" -#include "ggml.h" +#include "cumsum.cuh.inc" +#include "convert.cuh.inc" +#include "ggml-cuda/common.cuh.inc" +#include "ggml.h.inc" #ifdef GGML_CUDA_USE_CUB # include diff --git a/ggml/src/ggml-cuda/cumsum.cuh b/ggml/src/ggml-cuda/cumsum.cuh.inc similarity index 81% rename from ggml/src/ggml-cuda/cumsum.cuh rename to ggml/src/ggml-cuda/cumsum.cuh.inc index 782d1d92e9bb..3a0bd85afc45 100644 --- a/ggml/src/ggml-cuda/cumsum.cuh +++ b/ggml/src/ggml-cuda/cumsum.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_CUMSUM_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/dequantize.cuh b/ggml/src/ggml-cuda/dequantize.cuh.inc similarity index 99% rename from ggml/src/ggml-cuda/dequantize.cuh rename to ggml/src/ggml-cuda/dequantize.cuh.inc index 9ae1342fc0ef..6bc91d462c6d 100644 --- a/ggml/src/ggml-cuda/dequantize.cuh +++ b/ggml/src/ggml-cuda/dequantize.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" static __device__ __forceinline__ void dequantize_q1_0(const void * vx, const int64_t ib, const int iqs, float2 & v){ const block_q1_0 * x = (const block_q1_0 *) vx; diff --git a/ggml/src/ggml-cuda/diag.cu b/ggml/src/ggml-cuda/diag.cu.inc similarity index 97% rename from ggml/src/ggml-cuda/diag.cu rename to ggml/src/ggml-cuda/diag.cu.inc index 5cea210517f6..c5ce10a37bde 100644 --- a/ggml/src/ggml-cuda/diag.cu +++ b/ggml/src/ggml-cuda/diag.cu.inc @@ -1,6 +1,6 @@ -#include "convert.cuh" -#include "diag.cuh" -#include "ggml.h" +#include "convert.cuh.inc" +#include "diag.cuh.inc" +#include "ggml.h.inc" template static __global__ void diag_kernel(T * __restrict__ dst, diff --git a/ggml/src/ggml-cuda/diag.cuh b/ggml/src/ggml-cuda/diag.cuh.inc similarity index 81% rename from ggml/src/ggml-cuda/diag.cuh rename to ggml/src/ggml-cuda/diag.cuh.inc index 7d73e6a8eb44..bce89bda8cf7 100644 --- a/ggml/src/ggml-cuda/diag.cuh +++ b/ggml/src/ggml-cuda/diag.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_DIAG_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/diagmask.cu b/ggml/src/ggml-cuda/diagmask.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/diagmask.cu rename to ggml/src/ggml-cuda/diagmask.cu.inc index 4b713ba22eb5..94b08976878e 100644 --- a/ggml/src/ggml-cuda/diagmask.cu +++ b/ggml/src/ggml-cuda/diagmask.cu.inc @@ -1,4 +1,4 @@ -#include "diagmask.cuh" +#include "diagmask.cuh.inc" static __global__ void diag_mask_inf_f32(const float * x, float * dst, const int ncols, const int rows_per_channel, const int n_past) { const int col = blockDim.y*blockIdx.y + threadIdx.y; diff --git a/ggml/src/ggml-cuda/diagmask.cuh b/ggml/src/ggml-cuda/diagmask.cuh.inc similarity index 83% rename from ggml/src/ggml-cuda/diagmask.cuh rename to ggml/src/ggml-cuda/diagmask.cuh.inc index 6cdbef17e345..9f9e9e240f4a 100644 --- a/ggml/src/ggml-cuda/diagmask.cuh +++ b/ggml/src/ggml-cuda/diagmask.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_DIAG_MASK_INF_BLOCK_SIZE 32 diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh.inc similarity index 99% rename from ggml/src/ggml-cuda/fattn-common.cuh rename to ggml/src/ggml-cuda/fattn-common.cuh.inc index beeb52389464..34244b3558fe 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh.inc @@ -1,8 +1,8 @@ #pragma once -#include "common.cuh" -#include "convert.cuh" -#include "vecdotq.cuh" +#include "common.cuh.inc" +#include "convert.cuh.inc" +#include "vecdotq.cuh.inc" #include diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh.inc similarity index 99% rename from ggml/src/ggml-cuda/fattn-mma-f16.cuh rename to ggml/src/ggml-cuda/fattn-mma-f16.cuh.inc index 3f01e858de79..1049c3dc134d 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh.inc @@ -1,7 +1,7 @@ -#include "common.cuh" -#include "cp-async.cuh" -#include "mma.cuh" -#include "fattn-common.cuh" +#include "common.cuh.inc" +#include "cp-async.cuh.inc" +#include "mma.cuh.inc" +#include "fattn-common.cuh.inc" using namespace ggml_cuda_mma; @@ -1729,7 +1729,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml const int warp_size_host = ggml_cuda_info().devices[ctx.device].warp_size; const int nwarps = nthreads / warp_size_host; - constexpr bool V_is_K_view = DKQ == 576; // Guaranteed by the kernel selection logic in fattn.cu + constexpr bool V_is_K_view = DKQ == 576; // Guaranteed by the kernel selection logic in fattn.cu.inc const size_t nbytes_shared_KV_1stage = nbatch_fa * std::max(nbatch_K2 + 4, nbatch_V2 + 4) * sizeof(half2); const size_t nbytes_shared_KV_2stage = nbatch_fa * (nbatch_K2 + 4 + nbatch_V2 + 4) * sizeof(half2); diff --git a/ggml/src/ggml-cuda/fattn-tile.cu b/ggml/src/ggml-cuda/fattn-tile.cu.inc similarity index 95% rename from ggml/src/ggml-cuda/fattn-tile.cu rename to ggml/src/ggml-cuda/fattn-tile.cu.inc index d60634cc0e96..79d938096ced 100644 --- a/ggml/src/ggml-cuda/fattn-tile.cu +++ b/ggml/src/ggml-cuda/fattn-tile.cu.inc @@ -1,6 +1,6 @@ -#include "common.cuh" -#include "fattn-tile.cuh" -#include "fattn-wmma-f16.cuh" +#include "common.cuh.inc" +#include "fattn-tile.cuh.inc" +#include "fattn-wmma-f16.cuh.inc" void ggml_cuda_flash_attn_ext_tile(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * K = dst->src[1]; diff --git a/ggml/src/ggml-cuda/fattn-tile.cuh b/ggml/src/ggml-cuda/fattn-tile.cuh.inc similarity index 99% rename from ggml/src/ggml-cuda/fattn-tile.cuh rename to ggml/src/ggml-cuda/fattn-tile.cuh.inc index 585f2c228532..3b10b38189a5 100644 --- a/ggml/src/ggml-cuda/fattn-tile.cuh +++ b/ggml/src/ggml-cuda/fattn-tile.cuh.inc @@ -1,6 +1,6 @@ -#include "common.cuh" -#include "fattn-common.cuh" -#include "fattn-wmma-f16.cuh" +#include "common.cuh.inc" +#include "fattn-common.cuh.inc" +#include "fattn-wmma-f16.cuh.inc" // nbatch_fa == number of KQ rows to process per iteration // nbatch_K == number of K columns to load in parallel for KQ calculation diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh.inc similarity index 99% rename from ggml/src/ggml-cuda/fattn-vec.cuh rename to ggml/src/ggml-cuda/fattn-vec.cuh.inc index f0bd42a57612..b0e254602e45 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh.inc @@ -1,5 +1,5 @@ -#include "common.cuh" -#include "fattn-common.cuh" +#include "common.cuh.inc" +#include "fattn-common.cuh.inc" static int ggml_cuda_fattn_vec_get_nthreads_host(const int cc) { return 128; diff --git a/ggml/src/ggml-cuda/fattn-wmma-f16.cu b/ggml/src/ggml-cuda/fattn-wmma-f16.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/fattn-wmma-f16.cu rename to ggml/src/ggml-cuda/fattn-wmma-f16.cu.inc index f19defbff939..5856049582d7 100644 --- a/ggml/src/ggml-cuda/fattn-wmma-f16.cu +++ b/ggml/src/ggml-cuda/fattn-wmma-f16.cu.inc @@ -2,9 +2,9 @@ // It is still needed for Volta since the memory layout of NVIDIA tensor cores changed with Turing. // Long-term the WMMA code should be replaced with a dedicated Volta implementation. -#include "common.cuh" -#include "fattn-common.cuh" -#include "fattn-wmma-f16.cuh" +#include "common.cuh.inc" +#include "fattn-common.cuh.inc" +#include "fattn-wmma-f16.cuh.inc" #ifdef GGML_USE_WMMA_FATTN #if !defined(GGML_USE_HIP) diff --git a/ggml/src/ggml-cuda/fattn-wmma-f16.cuh b/ggml/src/ggml-cuda/fattn-wmma-f16.cuh.inc similarity index 98% rename from ggml/src/ggml-cuda/fattn-wmma-f16.cuh rename to ggml/src/ggml-cuda/fattn-wmma-f16.cuh.inc index aaf711a618cb..6b55b73ab510 100644 --- a/ggml/src/ggml-cuda/fattn-wmma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-wmma-f16.cuh.inc @@ -1,6 +1,6 @@ #pragma once -#include "common.cuh" +#include "common.cuh.inc" #if defined(GGML_USE_MUSA) #define GGML_USE_WMMA_FATTN diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/fattn.cu rename to ggml/src/ggml-cuda/fattn.cu.inc index 8256591b21d6..18b00ff8e449 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu.inc @@ -1,10 +1,10 @@ -#include "common.cuh" -#include "fattn-common.cuh" -#include "fattn-mma-f16.cuh" -#include "fattn-tile.cuh" -#include "fattn-vec.cuh" -#include "fattn-wmma-f16.cuh" -#include "fattn.cuh" +#include "common.cuh.inc" +#include "fattn-common.cuh.inc" +#include "fattn-mma-f16.cuh.inc" +#include "fattn-tile.cuh.inc" +#include "fattn-vec.cuh.inc" +#include "fattn-wmma-f16.cuh.inc" +#include "fattn.cuh.inc" template static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { diff --git a/ggml/src/ggml-cuda/fattn.cuh b/ggml/src/ggml-cuda/fattn.cuh.inc similarity index 86% rename from ggml/src/ggml-cuda/fattn.cuh rename to ggml/src/ggml-cuda/fattn.cuh.inc index 78705d59951c..a0e78858deb1 100644 --- a/ggml/src/ggml-cuda/fattn.cuh +++ b/ggml/src/ggml-cuda/fattn.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/fill.cu b/ggml/src/ggml-cuda/fill.cu.inc similarity index 95% rename from ggml/src/ggml-cuda/fill.cu rename to ggml/src/ggml-cuda/fill.cu.inc index 739062c4057a..6ba4b5e23556 100644 --- a/ggml/src/ggml-cuda/fill.cu +++ b/ggml/src/ggml-cuda/fill.cu.inc @@ -1,5 +1,5 @@ -#include "fill.cuh" -#include "convert.cuh" +#include "fill.cuh.inc" +#include "convert.cuh.inc" #define CUDA_FILL_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/fill.cuh b/ggml/src/ggml-cuda/fill.cuh.inc similarity index 74% rename from ggml/src/ggml-cuda/fill.cuh rename to ggml/src/ggml-cuda/fill.cuh.inc index 8443c83620de..8215c4db6e0f 100644 --- a/ggml/src/ggml-cuda/fill.cuh +++ b/ggml/src/ggml-cuda/fill.cuh.inc @@ -1,3 +1,3 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_op_fill(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/gated_delta_net.cu rename to ggml/src/ggml-cuda/gated_delta_net.cu.inc index 6b44bec73174..435650981aae 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu.inc @@ -1,4 +1,4 @@ -#include "gated_delta_net.cuh" +#include "gated_delta_net.cuh.inc" template __global__ void __launch_bounds__((ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v) * 4, 2) diff --git a/ggml/src/ggml-cuda/gated_delta_net.cuh b/ggml/src/ggml-cuda/gated_delta_net.cuh.inc similarity index 64% rename from ggml/src/ggml-cuda/gated_delta_net.cuh rename to ggml/src/ggml-cuda/gated_delta_net.cuh.inc index 7375e81c0c36..e065c6eab8f1 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cuh +++ b/ggml/src/ggml-cuda/gated_delta_net.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" -#include "ggml.h" +#include "common.cuh.inc" +#include "ggml.h.inc" void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/getrows.cu b/ggml/src/ggml-cuda/getrows.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/getrows.cu rename to ggml/src/ggml-cuda/getrows.cu.inc index e99cba63d344..3b71bd6c7412 100644 --- a/ggml/src/ggml-cuda/getrows.cu +++ b/ggml/src/ggml-cuda/getrows.cu.inc @@ -1,6 +1,6 @@ -#include "getrows.cuh" -#include "dequantize.cuh" -#include "convert.cuh" +#include "getrows.cuh.inc" +#include "dequantize.cuh.inc" +#include "convert.cuh.inc" template static __global__ void k_get_rows( diff --git a/ggml/src/ggml-cuda/getrows.cuh b/ggml/src/ggml-cuda/getrows.cuh.inc similarity index 95% rename from ggml/src/ggml-cuda/getrows.cuh rename to ggml/src/ggml-cuda/getrows.cuh.inc index 3c5bea5f48c1..095bee4da5be 100644 --- a/ggml/src/ggml-cuda/getrows.cuh +++ b/ggml/src/ggml-cuda/getrows.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_GET_ROWS_BLOCK_SIZE 256 #define CUDA_GET_ROWS_BACK_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/ggml-cuda.cu rename to ggml/src/ggml-cuda/ggml-cuda.cu.inc index fbe0fa06242c..381c1d35dee6 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu.inc @@ -1,67 +1,67 @@ -#include "ggml-cuda.h" -#include "ggml-impl.h" -#include "ggml-backend-impl.h" - -#include "ggml-cuda/common.cuh" -#include "ggml-cuda/acc.cuh" -#include "ggml-cuda/add-id.cuh" -#include "ggml-cuda/arange.cuh" -#include "ggml-cuda/argmax.cuh" -#include "ggml-cuda/argsort.cuh" -#include "ggml-cuda/binbcast.cuh" -#include "ggml-cuda/clamp.cuh" -#include "ggml-cuda/concat.cuh" -#include "ggml-cuda/conv-transpose-1d.cuh" -#include "ggml-cuda/conv2d.cuh" -#include "ggml-cuda/conv2d-dw.cuh" -#include "ggml-cuda/conv2d-transpose.cuh" -#include "ggml-cuda/convert.cuh" -#include "ggml-cuda/count-equal.cuh" -#include "ggml-cuda/cpy.cuh" -#include "ggml-cuda/cross-entropy-loss.cuh" -#include "ggml-cuda/cumsum.cuh" -#include "ggml-cuda/diagmask.cuh" -#include "ggml-cuda/diag.cuh" -#include "ggml-cuda/fattn.cuh" -#include "ggml-cuda/getrows.cuh" -#include "ggml-cuda/im2col.cuh" -#include "ggml-cuda/mmf.cuh" -#include "ggml-cuda/mmq.cuh" -#include "ggml-cuda/mmvf.cuh" -#include "ggml-cuda/mmvq.cuh" -#include "ggml-cuda/norm.cuh" -#include "ggml-cuda/opt-step-adamw.cuh" -#include "ggml-cuda/opt-step-sgd.cuh" -#include "ggml-cuda/out-prod.cuh" -#include "ggml-cuda/pad.cuh" -#include "ggml-cuda/pool2d.cuh" -#include "ggml-cuda/quantize.cuh" -#include "ggml-cuda/rope.cuh" -#include "ggml-cuda/roll.cuh" -#include "ggml-cuda/scale.cuh" -#include "ggml-cuda/softcap.cuh" -#include "ggml-cuda/softmax.cuh" -#include "ggml-cuda/ssm-conv.cuh" -#include "ggml-cuda/ssm-scan.cuh" -#include "ggml-cuda/sum.cuh" -#include "ggml-cuda/sumrows.cuh" -#include "ggml-cuda/top-k.cuh" -#include "ggml-cuda/mean.cuh" -#include "ggml-cuda/tsembd.cuh" -#include "ggml-cuda/topk-moe.cuh" -#include "ggml-cuda/unary.cuh" -#include "ggml-cuda/upscale.cuh" -#include "ggml-cuda/wkv.cuh" -#include "ggml-cuda/gla.cuh" -#include "ggml-cuda/gated_delta_net.cuh" -#include "ggml-cuda/set.cuh" -#include "ggml-cuda/set-rows.cuh" -#include "ggml-cuda/pad_reflect_1d.cuh" -#include "ggml-cuda/solve_tri.cuh" -#include "ggml-cuda/tri.cuh" -#include "ggml-cuda/cumsum.cuh" -#include "ggml-cuda/fill.cuh" -#include "ggml.h" +#include "ggml-cuda.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-backend-impl-defs.inc" + +#include "ggml-cuda/common.cuh.inc" +#include "ggml-cuda/acc.cuh.inc" +#include "ggml-cuda/add-id.cuh.inc" +#include "ggml-cuda/arange.cuh.inc" +#include "ggml-cuda/argmax.cuh.inc" +#include "ggml-cuda/argsort.cuh.inc" +#include "ggml-cuda/binbcast.cuh.inc" +#include "ggml-cuda/clamp.cuh.inc" +#include "ggml-cuda/concat.cuh.inc" +#include "ggml-cuda/conv-transpose-1d.cuh.inc" +#include "ggml-cuda/conv2d.cuh.inc" +#include "ggml-cuda/conv2d-dw.cuh.inc" +#include "ggml-cuda/conv2d-transpose.cuh.inc" +#include "ggml-cuda/convert.cuh.inc" +#include "ggml-cuda/count-equal.cuh.inc" +#include "ggml-cuda/cpy.cuh.inc" +#include "ggml-cuda/cross-entropy-loss.cuh.inc" +#include "ggml-cuda/cumsum.cuh.inc" +#include "ggml-cuda/diagmask.cuh.inc" +#include "ggml-cuda/diag.cuh.inc" +#include "ggml-cuda/fattn.cuh.inc" +#include "ggml-cuda/getrows.cuh.inc" +#include "ggml-cuda/im2col.cuh.inc" +#include "ggml-cuda/mmf.cuh.inc" +#include "ggml-cuda/mmq.cuh.inc" +#include "ggml-cuda/mmvf.cuh.inc" +#include "ggml-cuda/mmvq.cuh.inc" +#include "ggml-cuda/norm.cuh.inc" +#include "ggml-cuda/opt-step-adamw.cuh.inc" +#include "ggml-cuda/opt-step-sgd.cuh.inc" +#include "ggml-cuda/out-prod.cuh.inc" +#include "ggml-cuda/pad.cuh.inc" +#include "ggml-cuda/pool2d.cuh.inc" +#include "ggml-cuda/quantize.cuh.inc" +#include "ggml-cuda/rope.cuh.inc" +#include "ggml-cuda/roll.cuh.inc" +#include "ggml-cuda/scale.cuh.inc" +#include "ggml-cuda/softcap.cuh.inc" +#include "ggml-cuda/softmax.cuh.inc" +#include "ggml-cuda/ssm-conv.cuh.inc" +#include "ggml-cuda/ssm-scan.cuh.inc" +#include "ggml-cuda/sum.cuh.inc" +#include "ggml-cuda/sumrows.cuh.inc" +#include "ggml-cuda/top-k.cuh.inc" +#include "ggml-cuda/mean.cuh.inc" +#include "ggml-cuda/tsembd.cuh.inc" +#include "ggml-cuda/topk-moe.cuh.inc" +#include "ggml-cuda/unary.cuh.inc" +#include "ggml-cuda/upscale.cuh.inc" +#include "ggml-cuda/wkv.cuh.inc" +#include "ggml-cuda/gla.cuh.inc" +#include "ggml-cuda/gated_delta_net.cuh.inc" +#include "ggml-cuda/set.cuh.inc" +#include "ggml-cuda/set-rows.cuh.inc" +#include "ggml-cuda/pad_reflect_1d.cuh.inc" +#include "ggml-cuda/solve_tri.cuh.inc" +#include "ggml-cuda/tri.cuh.inc" +#include "ggml-cuda/cumsum.cuh.inc" +#include "ggml-cuda/fill.cuh.inc" +#include "ggml.h.inc" #include #include diff --git a/ggml/src/ggml-cuda/gla.cu b/ggml/src/ggml-cuda/gla.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/gla.cu rename to ggml/src/ggml-cuda/gla.cu.inc index f7d615a8282f..b8dc5493ba2b 100644 --- a/ggml/src/ggml-cuda/gla.cu +++ b/ggml/src/ggml-cuda/gla.cu.inc @@ -1,5 +1,5 @@ -#include "common.cuh" -#include "gla.cuh" +#include "common.cuh.inc" +#include "gla.cuh.inc" template static __global__ void gated_linear_attn_f32(const int B, const int T, const int C, const int H, const float scale, diff --git a/ggml/src/ggml-cuda/gla.cuh b/ggml/src/ggml-cuda/gla.cuh.inc similarity index 77% rename from ggml/src/ggml-cuda/gla.cuh rename to ggml/src/ggml-cuda/gla.cuh.inc index 2c82ad7dd722..81042a671a59 100644 --- a/ggml/src/ggml-cuda/gla.cuh +++ b/ggml/src/ggml-cuda/gla.cuh.inc @@ -1,3 +1,3 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_op_gated_linear_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/im2col.cu b/ggml/src/ggml-cuda/im2col.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/im2col.cu rename to ggml/src/ggml-cuda/im2col.cu.inc index 56dc0545742e..9d9790239fb8 100644 --- a/ggml/src/ggml-cuda/im2col.cu +++ b/ggml/src/ggml-cuda/im2col.cu.inc @@ -1,4 +1,4 @@ -#include "im2col.cuh" +#include "im2col.cuh.inc" #define MAX_GRIDDIM_Z 65535 diff --git a/ggml/src/ggml-cuda/im2col.cuh b/ggml/src/ggml-cuda/im2col.cuh.inc similarity index 88% rename from ggml/src/ggml-cuda/im2col.cuh rename to ggml/src/ggml-cuda/im2col.cuh.inc index 2da1223d6345..f77a284b3c40 100644 --- a/ggml/src/ggml-cuda/im2col.cuh +++ b/ggml/src/ggml-cuda/im2col.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_IM2COL_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/mean.cu b/ggml/src/ggml-cuda/mean.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/mean.cu rename to ggml/src/ggml-cuda/mean.cu.inc index 49af5389957e..677b6d422cf9 100644 --- a/ggml/src/ggml-cuda/mean.cu +++ b/ggml/src/ggml-cuda/mean.cu.inc @@ -1,5 +1,5 @@ -#include "mean.cuh" -#include "reduce_rows.cuh" +#include "mean.cuh.inc" +#include "reduce_rows.cuh.inc" #ifdef GGML_CUDA_USE_CUB #include diff --git a/ggml/src/ggml-cuda/mean.cuh b/ggml/src/ggml-cuda/mean.cuh.inc similarity index 74% rename from ggml/src/ggml-cuda/mean.cuh rename to ggml/src/ggml-cuda/mean.cuh.inc index 2b9b10433438..881d6db0943f 100644 --- a/ggml/src/ggml-cuda/mean.cuh +++ b/ggml/src/ggml-cuda/mean.cuh.inc @@ -1,3 +1,3 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_op_mean(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/mma.cuh b/ggml/src/ggml-cuda/mma.cuh.inc similarity index 99% rename from ggml/src/ggml-cuda/mma.cuh rename to ggml/src/ggml-cuda/mma.cuh.inc index 79bb2934c5f9..c0bd23e77f60 100644 --- a/ggml/src/ggml-cuda/mma.cuh +++ b/ggml/src/ggml-cuda/mma.cuh.inc @@ -16,7 +16,7 @@ // As described in the PTX documentation, all pointers for load_ldmatrix must be to shared memory and aligned to 16 bytes. // The API in this file also assumes that the pointers for load_generic are aligned to 16 bytes, unaligned pointers are considered undefined behavior. -#include "common.cuh" +#include "common.cuh.inc" // On Volta each warp is doing 4 8x8 mma operations in parallel. // The basic memory layout for a 32x8 output tile is to stack 4 input tiles in I direction and to mirror the B tile. diff --git a/ggml/src/ggml-cuda/mmf.cu b/ggml/src/ggml-cuda/mmf.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/mmf.cu rename to ggml/src/ggml-cuda/mmf.cu.inc index aad4c34aa668..bdc902afe4a0 100644 --- a/ggml/src/ggml-cuda/mmf.cu +++ b/ggml/src/ggml-cuda/mmf.cu.inc @@ -1,6 +1,6 @@ -#include "ggml.h" -#include "mmf.cuh" -#include "mmid.cuh" +#include "ggml.h.inc" +#include "mmf.cuh.inc" +#include "mmid.cuh.inc" static __forceinline__ int mmf_get_rows_per_block(const int cc) { if (GGML_CUDA_CC_IS_CDNA(cc)) { diff --git a/ggml/src/ggml-cuda/mmf.cuh b/ggml/src/ggml-cuda/mmf.cuh.inc similarity index 99% rename from ggml/src/ggml-cuda/mmf.cuh rename to ggml/src/ggml-cuda/mmf.cuh.inc index c2a8d54c95ac..2cc55bd6c4bb 100644 --- a/ggml/src/ggml-cuda/mmf.cuh +++ b/ggml/src/ggml-cuda/mmf.cuh.inc @@ -1,8 +1,8 @@ #pragma once -#include "mma.cuh" -#include "common.cuh" -#include "convert.cuh" +#include "mma.cuh.inc" +#include "common.cuh.inc" +#include "convert.cuh.inc" using namespace ggml_cuda_mma; diff --git a/ggml/src/ggml-cuda/mmid.cu b/ggml/src/ggml-cuda/mmid.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/mmid.cu rename to ggml/src/ggml-cuda/mmid.cu.inc index 3c61e4595a7b..28c98ad3cb50 100644 --- a/ggml/src/ggml-cuda/mmid.cu +++ b/ggml/src/ggml-cuda/mmid.cu.inc @@ -1,5 +1,5 @@ -#include "common.cuh" -#include "mmid.cuh" +#include "common.cuh.inc" +#include "mmid.cuh.inc" // To reduce shared memory use, store "it" and "iex_used" with 22/10 bits each. struct mm_ids_helper_store { diff --git a/ggml/src/ggml-cuda/mmid.cuh b/ggml/src/ggml-cuda/mmid.cuh.inc similarity index 100% rename from ggml/src/ggml-cuda/mmid.cuh rename to ggml/src/ggml-cuda/mmid.cuh.inc diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/mmq.cu rename to ggml/src/ggml-cuda/mmq.cu.inc index e1add5e03316..e3c5a70b81fd 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu.inc @@ -1,7 +1,7 @@ -#include "common.cuh" -#include "mmq.cuh" -#include "quantize.cuh" -#include "mmid.cuh" +#include "common.cuh.inc" +#include "mmq.cuh.inc" +#include "quantize.cuh.inc" +#include "mmid.cuh.inc" static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { switch (args.type_x) { diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh.inc similarity index 99% rename from ggml/src/ggml-cuda/mmq.cuh rename to ggml/src/ggml-cuda/mmq.cuh.inc index edf546d8f1e2..16a39cc5078b 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh.inc @@ -1,8 +1,8 @@ #pragma once -#include "common.cuh" -#include "vecdotq.cuh" -#include "mma.cuh" +#include "common.cuh.inc" +#include "vecdotq.cuh.inc" +#include "mma.cuh.inc" #include #include diff --git a/ggml/src/ggml-cuda/mmvf.cu b/ggml/src/ggml-cuda/mmvf.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/mmvf.cu rename to ggml/src/ggml-cuda/mmvf.cu.inc index d91472024296..09fc3417c560 100644 --- a/ggml/src/ggml-cuda/mmvf.cu +++ b/ggml/src/ggml-cuda/mmvf.cu.inc @@ -1,8 +1,8 @@ -#include "ggml.h" -#include "common.cuh" -#include "unary.cuh" -#include "mmvf.cuh" -#include "convert.cuh" +#include "ggml.h.inc" +#include "common.cuh.inc" +#include "unary.cuh.inc" +#include "mmvf.cuh.inc" +#include "convert.cuh.inc" template static __global__ void mul_mat_vec_f( diff --git a/ggml/src/ggml-cuda/mmvf.cuh b/ggml/src/ggml-cuda/mmvf.cuh.inc similarity index 96% rename from ggml/src/ggml-cuda/mmvf.cuh rename to ggml/src/ggml-cuda/mmvf.cuh.inc index a50f7c021801..a073e639f2af 100644 --- a/ggml/src/ggml-cuda/mmvf.cuh +++ b/ggml/src/ggml-cuda/mmvf.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define MMVF_MAX_BATCH_SIZE 8 // Max. batch size for which to use MMVF kernels. diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/mmvq.cu rename to ggml/src/ggml-cuda/mmvq.cu.inc index da48f313a38b..d17b4c86d68b 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu.inc @@ -1,7 +1,7 @@ -#include "mmvq.cuh" -#include "quantize.cuh" -#include "unary.cuh" -#include "vecdotq.cuh" +#include "mmvq.cuh.inc" +#include "quantize.cuh.inc" +#include "unary.cuh.inc" +#include "vecdotq.cuh.inc" #include diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh.inc similarity index 97% rename from ggml/src/ggml-cuda/mmvq.cuh rename to ggml/src/ggml-cuda/mmvq.cuh.inc index 6bf0a8e8677d..58cc1f922224 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/ggml/src/ggml-cuda/mmvq.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define MMVQ_MAX_BATCH_SIZE 8 // Max. batch size for which to use MMVQ kernels. diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/norm.cu rename to ggml/src/ggml-cuda/norm.cu.inc index ef98f675aa71..4739fe77cfaf 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu.inc @@ -1,4 +1,4 @@ -#include "norm.cuh" +#include "norm.cuh.inc" #include template diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh.inc similarity index 96% rename from ggml/src/ggml-cuda/norm.cuh rename to ggml/src/ggml-cuda/norm.cuh.inc index a74f6376720a..85c9df7545e4 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_op_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/opt-step-adamw.cu b/ggml/src/ggml-cuda/opt-step-adamw.cu.inc similarity index 97% rename from ggml/src/ggml-cuda/opt-step-adamw.cu rename to ggml/src/ggml-cuda/opt-step-adamw.cu.inc index 35154f299665..8f04a18f0a27 100644 --- a/ggml/src/ggml-cuda/opt-step-adamw.cu +++ b/ggml/src/ggml-cuda/opt-step-adamw.cu.inc @@ -1,5 +1,5 @@ -#include "ggml-impl.h" -#include "opt-step-adamw.cuh" +#include "ggml-impl-defs.inc" +#include "opt-step-adamw.cuh.inc" #include diff --git a/ggml/src/ggml-cuda/opt-step-adamw.cuh b/ggml/src/ggml-cuda/opt-step-adamw.cuh.inc similarity index 83% rename from ggml/src/ggml-cuda/opt-step-adamw.cuh rename to ggml/src/ggml-cuda/opt-step-adamw.cuh.inc index 58d6f6e5dfc5..d44ba211bf9e 100644 --- a/ggml/src/ggml-cuda/opt-step-adamw.cuh +++ b/ggml/src/ggml-cuda/opt-step-adamw.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_OPT_STEP_ADAMW_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/opt-step-sgd.cu b/ggml/src/ggml-cuda/opt-step-sgd.cu.inc similarity index 96% rename from ggml/src/ggml-cuda/opt-step-sgd.cu rename to ggml/src/ggml-cuda/opt-step-sgd.cu.inc index 460b16de447a..411176c9fedb 100644 --- a/ggml/src/ggml-cuda/opt-step-sgd.cu +++ b/ggml/src/ggml-cuda/opt-step-sgd.cu.inc @@ -1,5 +1,5 @@ -#include "ggml-impl.h" -#include "opt-step-sgd.cuh" +#include "ggml-impl-defs.inc" +#include "opt-step-sgd.cuh.inc" #include diff --git a/ggml/src/ggml-cuda/opt-step-sgd.cuh b/ggml/src/ggml-cuda/opt-step-sgd.cuh.inc similarity index 82% rename from ggml/src/ggml-cuda/opt-step-sgd.cuh rename to ggml/src/ggml-cuda/opt-step-sgd.cuh.inc index f97ab7d9bede..777ca7bcd6c3 100644 --- a/ggml/src/ggml-cuda/opt-step-sgd.cuh +++ b/ggml/src/ggml-cuda/opt-step-sgd.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_OPT_STEP_SGD_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/out-prod.cu b/ggml/src/ggml-cuda/out-prod.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/out-prod.cu rename to ggml/src/ggml-cuda/out-prod.cu.inc index c9b2b699c6a5..11f9aa1f8b41 100644 --- a/ggml/src/ggml-cuda/out-prod.cu +++ b/ggml/src/ggml-cuda/out-prod.cu.inc @@ -1,4 +1,4 @@ -#include "out-prod.cuh" +#include "out-prod.cuh.inc" #include diff --git a/ggml/src/ggml-cuda/out-prod.cuh b/ggml/src/ggml-cuda/out-prod.cuh.inc similarity index 75% rename from ggml/src/ggml-cuda/out-prod.cuh rename to ggml/src/ggml-cuda/out-prod.cuh.inc index a0046f5f8f48..e7b94dd15457 100644 --- a/ggml/src/ggml-cuda/out-prod.cuh +++ b/ggml/src/ggml-cuda/out-prod.cuh.inc @@ -1,3 +1,3 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_out_prod(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/pad.cu b/ggml/src/ggml-cuda/pad.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/pad.cu rename to ggml/src/ggml-cuda/pad.cu.inc index 31cd00f77816..1437fb42e2dc 100644 --- a/ggml/src/ggml-cuda/pad.cu +++ b/ggml/src/ggml-cuda/pad.cu.inc @@ -1,4 +1,4 @@ -#include "pad.cuh" +#include "pad.cuh.inc" #include diff --git a/ggml/src/ggml-cuda/pad.cuh b/ggml/src/ggml-cuda/pad.cuh.inc similarity index 80% rename from ggml/src/ggml-cuda/pad.cuh rename to ggml/src/ggml-cuda/pad.cuh.inc index 8fd386b008c1..a54d14fb393e 100644 --- a/ggml/src/ggml-cuda/pad.cuh +++ b/ggml/src/ggml-cuda/pad.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_PAD_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/pad_reflect_1d.cu b/ggml/src/ggml-cuda/pad_reflect_1d.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/pad_reflect_1d.cu rename to ggml/src/ggml-cuda/pad_reflect_1d.cu.inc index 32993eb59130..8a124322fdcd 100644 --- a/ggml/src/ggml-cuda/pad_reflect_1d.cu +++ b/ggml/src/ggml-cuda/pad_reflect_1d.cu.inc @@ -1,4 +1,4 @@ -#include "pad_reflect_1d.cuh" +#include "pad_reflect_1d.cuh.inc" static __global__ __launch_bounds__(CUDA_PAD_REFLECT_1D_BLOCK_SIZE, 1) void pad_reflect_1d_kernel_f32( @@ -28,7 +28,7 @@ static __global__ __launch_bounds__(CUDA_PAD_REFLECT_1D_BLOCK_SIZE, 1) void const int64_t i1 = tile1; const int64_t i0 = threadIdx.x + tile0 * blockDim.x; - // ne01.z is original value of unpacked ne01 (see init_fastdiv_values in common.cuh) + // ne01.z is original value of unpacked ne01 (see init_fastdiv_values in common.cuh.inc) if (i0 >= ne0 || i1 >= ne01.z || i2 >= ne02 || i3 >= ne03) { return; } diff --git a/ggml/src/ggml-cuda/pad_reflect_1d.cuh b/ggml/src/ggml-cuda/pad_reflect_1d.cuh.inc similarity index 83% rename from ggml/src/ggml-cuda/pad_reflect_1d.cuh rename to ggml/src/ggml-cuda/pad_reflect_1d.cuh.inc index 15f2ed1737b1..7d0737c3e36a 100644 --- a/ggml/src/ggml-cuda/pad_reflect_1d.cuh +++ b/ggml/src/ggml-cuda/pad_reflect_1d.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_PAD_REFLECT_1D_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/pool2d.cu b/ggml/src/ggml-cuda/pool2d.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/pool2d.cu rename to ggml/src/ggml-cuda/pool2d.cu.inc index c6d51e4d655a..9f2437bec381 100644 --- a/ggml/src/ggml-cuda/pool2d.cu +++ b/ggml/src/ggml-cuda/pool2d.cu.inc @@ -1,4 +1,4 @@ -#include "pool2d.cuh" +#include "pool2d.cuh.inc" template static __global__ void pool2d_nchw_kernel( diff --git a/ggml/src/ggml-cuda/pool2d.cuh b/ggml/src/ggml-cuda/pool2d.cuh.inc similarity index 81% rename from ggml/src/ggml-cuda/pool2d.cuh rename to ggml/src/ggml-cuda/pool2d.cuh.inc index 7841292bcc27..27aba6d10bf3 100644 --- a/ggml/src/ggml-cuda/pool2d.cuh +++ b/ggml/src/ggml-cuda/pool2d.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_POOL2D_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/quantize.cu b/ggml/src/ggml-cuda/quantize.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/quantize.cu rename to ggml/src/ggml-cuda/quantize.cu.inc index 52f664719ae1..33e3d0cdbec6 100644 --- a/ggml/src/ggml-cuda/quantize.cu +++ b/ggml/src/ggml-cuda/quantize.cu.inc @@ -1,4 +1,4 @@ -#include "quantize.cuh" +#include "quantize.cuh.inc" #include __launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1) diff --git a/ggml/src/ggml-cuda/quantize.cuh b/ggml/src/ggml-cuda/quantize.cuh.inc similarity index 97% rename from ggml/src/ggml-cuda/quantize.cuh rename to ggml/src/ggml-cuda/quantize.cuh.inc index 768a3ae6de6c..f784f8fa5da1 100644 --- a/ggml/src/ggml-cuda/quantize.cuh +++ b/ggml/src/ggml-cuda/quantize.cuh.inc @@ -1,7 +1,7 @@ #pragma once -#include "common.cuh" -#include "mmq.cuh" +#include "common.cuh.inc" +#include "mmq.cuh.inc" #include diff --git a/ggml/src/ggml-cuda/reduce_rows.cuh b/ggml/src/ggml-cuda/reduce_rows.cuh.inc similarity index 97% rename from ggml/src/ggml-cuda/reduce_rows.cuh rename to ggml/src/ggml-cuda/reduce_rows.cuh.inc index de240fd44135..fd0d6e527ed7 100644 --- a/ggml/src/ggml-cuda/reduce_rows.cuh +++ b/ggml/src/ggml-cuda/reduce_rows.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" // Row reduction kernel template - compute sum (norm=false) or mean (norm=true) template diff --git a/ggml/src/ggml-cuda/roll.cu b/ggml/src/ggml-cuda/roll.cu.inc similarity index 97% rename from ggml/src/ggml-cuda/roll.cu rename to ggml/src/ggml-cuda/roll.cu.inc index a339dfc1ae0b..4f16491f0161 100644 --- a/ggml/src/ggml-cuda/roll.cu +++ b/ggml/src/ggml-cuda/roll.cu.inc @@ -1,5 +1,5 @@ -#include "ggml-cuda/common.cuh" -#include "roll.cuh" +#include "ggml-cuda/common.cuh.inc" +#include "roll.cuh.inc" static __forceinline__ __device__ int64_t wrap_index(const int64_t idx, const int64_t ne) { if (idx < 0) { diff --git a/ggml/src/ggml-cuda/roll.cuh b/ggml/src/ggml-cuda/roll.cuh.inc similarity index 81% rename from ggml/src/ggml-cuda/roll.cuh rename to ggml/src/ggml-cuda/roll.cuh.inc index 322d55436e25..4b44ac99d1c6 100644 --- a/ggml/src/ggml-cuda/roll.cuh +++ b/ggml/src/ggml-cuda/roll.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_ROLL_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/rope.cu b/ggml/src/ggml-cuda/rope.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/rope.cu rename to ggml/src/ggml-cuda/rope.cu.inc index 45a49a5dc2a3..5bcba45895b6 100644 --- a/ggml/src/ggml-cuda/rope.cu +++ b/ggml/src/ggml-cuda/rope.cu.inc @@ -1,7 +1,7 @@ -#include "convert.cuh" -#include "ggml-cuda/common.cuh" -#include "ggml.h" -#include "rope.cuh" +#include "convert.cuh.inc" +#include "ggml-cuda/common.cuh.inc" +#include "ggml.h.inc" +#include "rope.cuh.inc" struct rope_corr_dims { float v[2]; diff --git a/ggml/src/ggml-cuda/rope.cuh b/ggml/src/ggml-cuda/rope.cuh.inc similarity index 92% rename from ggml/src/ggml-cuda/rope.cuh rename to ggml/src/ggml-cuda/rope.cuh.inc index 72af086cd1b4..0317d3245cb6 100644 --- a/ggml/src/ggml-cuda/rope.cuh +++ b/ggml/src/ggml-cuda/rope.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_ROPE_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/scale.cu b/ggml/src/ggml-cuda/scale.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/scale.cu rename to ggml/src/ggml-cuda/scale.cu.inc index 0ddeff6a1755..8e80462bc14b 100644 --- a/ggml/src/ggml-cuda/scale.cu +++ b/ggml/src/ggml-cuda/scale.cu.inc @@ -1,4 +1,4 @@ -#include "scale.cuh" +#include "scale.cuh.inc" #define MAX_GRIDDIM_X 0x7FFFFFFF diff --git a/ggml/src/ggml-cuda/scale.cuh b/ggml/src/ggml-cuda/scale.cuh.inc similarity index 81% rename from ggml/src/ggml-cuda/scale.cuh rename to ggml/src/ggml-cuda/scale.cuh.inc index 8ff75c8298b0..ccf97e9f3bb4 100644 --- a/ggml/src/ggml-cuda/scale.cuh +++ b/ggml/src/ggml-cuda/scale.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_SCALE_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/set-rows.cu b/ggml/src/ggml-cuda/set-rows.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/set-rows.cu rename to ggml/src/ggml-cuda/set-rows.cu.inc index 631de7e8fa51..f2dd68f3faee 100644 --- a/ggml/src/ggml-cuda/set-rows.cu +++ b/ggml/src/ggml-cuda/set-rows.cu.inc @@ -1,5 +1,5 @@ -#include "set-rows.cuh" -#include "cpy-utils.cuh" +#include "set-rows.cuh.inc" +#include "cpy-utils.cuh.inc" typedef void (*set_rows_kernel_t)(const char * src, char * dst); diff --git a/ggml/src/ggml-cuda/set-rows.cuh b/ggml/src/ggml-cuda/set-rows.cuh.inc similarity index 83% rename from ggml/src/ggml-cuda/set-rows.cuh rename to ggml/src/ggml-cuda/set-rows.cuh.inc index c140c0873c8a..34ee02b6fea9 100644 --- a/ggml/src/ggml-cuda/set-rows.cuh +++ b/ggml/src/ggml-cuda/set-rows.cuh.inc @@ -1,6 +1,6 @@ #pragma once -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_SET_ROWS_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/set.cu b/ggml/src/ggml-cuda/set.cu.inc similarity index 96% rename from ggml/src/ggml-cuda/set.cu rename to ggml/src/ggml-cuda/set.cu.inc index 04bfe07ba033..38fbf1cb7a4b 100644 --- a/ggml/src/ggml-cuda/set.cu +++ b/ggml/src/ggml-cuda/set.cu.inc @@ -1,5 +1,5 @@ -#include "set.cuh" -#include "cpy.cuh" +#include "set.cuh.inc" +#include "cpy.cuh.inc" void ggml_cuda_op_set(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; diff --git a/ggml/src/ggml-cuda/set.cuh b/ggml/src/ggml-cuda/set.cuh.inc similarity index 82% rename from ggml/src/ggml-cuda/set.cuh rename to ggml/src/ggml-cuda/set.cuh.inc index dd09529f3e42..3b10ea2ea58c 100644 --- a/ggml/src/ggml-cuda/set.cuh +++ b/ggml/src/ggml-cuda/set.cuh.inc @@ -1,6 +1,6 @@ #pragma once -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_SET_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/softcap.cu b/ggml/src/ggml-cuda/softcap.cu.inc similarity index 97% rename from ggml/src/ggml-cuda/softcap.cu rename to ggml/src/ggml-cuda/softcap.cu.inc index 40dfe45d65cf..ad85a984bcad 100644 --- a/ggml/src/ggml-cuda/softcap.cu +++ b/ggml/src/ggml-cuda/softcap.cu.inc @@ -1,4 +1,4 @@ -#include "softcap.cuh" +#include "softcap.cuh.inc" static __global__ void softcap_f32(const float * x, float * dst, const float scale, const float softcap, const int k) { const int i = blockDim.x*blockIdx.x + threadIdx.x; diff --git a/ggml/src/ggml-cuda/softcap.cuh b/ggml/src/ggml-cuda/softcap.cuh.inc similarity index 83% rename from ggml/src/ggml-cuda/softcap.cuh rename to ggml/src/ggml-cuda/softcap.cuh.inc index 6d34fb2bee41..a6ceb4f56564 100644 --- a/ggml/src/ggml-cuda/softcap.cuh +++ b/ggml/src/ggml-cuda/softcap.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_SOFTCAP_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/softmax.cu b/ggml/src/ggml-cuda/softmax.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/softmax.cu rename to ggml/src/ggml-cuda/softmax.cu.inc index 285c0e9543a2..715c78103e75 100644 --- a/ggml/src/ggml-cuda/softmax.cu +++ b/ggml/src/ggml-cuda/softmax.cu.inc @@ -1,6 +1,6 @@ -#include "common.cuh" -#include "ggml.h" -#include "softmax.cuh" +#include "common.cuh.inc" +#include "ggml.h.inc" +#include "softmax.cuh.inc" #ifdef GGML_USE_HIP #include diff --git a/ggml/src/ggml-cuda/softmax.cuh b/ggml/src/ggml-cuda/softmax.cuh.inc similarity index 88% rename from ggml/src/ggml-cuda/softmax.cuh rename to ggml/src/ggml-cuda/softmax.cuh.inc index 93dfee835f6f..a2e4d30cdcd7 100644 --- a/ggml/src/ggml-cuda/softmax.cuh +++ b/ggml/src/ggml-cuda/softmax.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_SOFT_MAX_BLOCK_SIZE 1024 diff --git a/ggml/src/ggml-cuda/solve_tri.cu b/ggml/src/ggml-cuda/solve_tri.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/solve_tri.cu rename to ggml/src/ggml-cuda/solve_tri.cu.inc index 07ca33f513b0..b060e7ff44fe 100644 --- a/ggml/src/ggml-cuda/solve_tri.cu +++ b/ggml/src/ggml-cuda/solve_tri.cu.inc @@ -1,6 +1,6 @@ -#include "common.cuh" -#include "ggml.h" -#include "solve_tri.cuh" +#include "common.cuh.inc" +#include "ggml.h.inc" +#include "solve_tri.cuh.inc" #define MAX_N_FAST 64 #define MAX_K_FAST 32 @@ -72,7 +72,7 @@ static void solve_tri_f32_cublas(ggml_backend_cuda_context & ctx, CUBLAS_CHECK(cublasStrsmBatched(ctx.cublas_handle(id), CUBLAS_SIDE_RIGHT, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT, k, n, &alpha, A_ptrs_dev, n, X_ptrs_dev, k, total_batches)); - // revert to standard mode from common.cuh + // revert to standard mode from common.cuh.inc CUBLAS_CHECK(cublasSetMathMode(ctx.cublas_handle(id), CUBLAS_TF32_TENSOR_OP_MATH)); GGML_UNUSED_VARS(s12, s13); diff --git a/ggml/src/ggml-cuda/solve_tri.cuh b/ggml/src/ggml-cuda/solve_tri.cuh.inc similarity index 75% rename from ggml/src/ggml-cuda/solve_tri.cuh rename to ggml/src/ggml-cuda/solve_tri.cuh.inc index 639992396a35..ad51ca6db91d 100644 --- a/ggml/src/ggml-cuda/solve_tri.cuh +++ b/ggml/src/ggml-cuda/solve_tri.cuh.inc @@ -1,3 +1,3 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_op_solve_tri(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ssm-conv.cu b/ggml/src/ggml-cuda/ssm-conv.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/ssm-conv.cu rename to ggml/src/ggml-cuda/ssm-conv.cu.inc index 4841389fbc88..ec3881fe3fb1 100644 --- a/ggml/src/ggml-cuda/ssm-conv.cu +++ b/ggml/src/ggml-cuda/ssm-conv.cu.inc @@ -1,5 +1,5 @@ -#include "ssm-conv.cuh" -#include "unary.cuh" +#include "ssm-conv.cuh.inc" +#include "unary.cuh.inc" template static __global__ void ssm_conv_f32(const float * __restrict__ src0, const float * __restrict__ src1, diff --git a/ggml/src/ggml-cuda/ssm-conv.cuh b/ggml/src/ggml-cuda/ssm-conv.cuh.inc similarity index 85% rename from ggml/src/ggml-cuda/ssm-conv.cuh rename to ggml/src/ggml-cuda/ssm-conv.cuh.inc index 8514ca84920f..3e96c495ca0c 100644 --- a/ggml/src/ggml-cuda/ssm-conv.cuh +++ b/ggml/src/ggml-cuda/ssm-conv.cuh.inc @@ -1,3 +1,3 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_op_ssm_conv(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * bias_add_node = nullptr, ggml_tensor * silu_dst = nullptr); diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/ssm-scan.cu rename to ggml/src/ggml-cuda/ssm-scan.cu.inc index c1d4e2bc8dfd..74849e57faf1 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu.inc @@ -7,7 +7,7 @@ using namespace cub; #endif // USE_CUB -#include "ssm-scan.cuh" +#include "ssm-scan.cuh.inc" // We would like to keep pragma unroll for cases where L_template is not 0, // so we suppress the clang transformation warning. diff --git a/ggml/src/ggml-cuda/ssm-scan.cuh b/ggml/src/ggml-cuda/ssm-scan.cuh.inc similarity index 75% rename from ggml/src/ggml-cuda/ssm-scan.cuh rename to ggml/src/ggml-cuda/ssm-scan.cuh.inc index ee078f5ebb8c..95c0080f5c5a 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cuh +++ b/ggml/src/ggml-cuda/ssm-scan.cuh.inc @@ -1,3 +1,3 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/sum.cu b/ggml/src/ggml-cuda/sum.cu.inc similarity index 96% rename from ggml/src/ggml-cuda/sum.cu rename to ggml/src/ggml-cuda/sum.cu.inc index c56257b44066..cbca43910152 100644 --- a/ggml/src/ggml-cuda/sum.cu +++ b/ggml/src/ggml-cuda/sum.cu.inc @@ -1,5 +1,5 @@ -#include "sum.cuh" -#include "sumrows.cuh" +#include "sum.cuh.inc" +#include "sumrows.cuh.inc" #ifdef GGML_CUDA_USE_CUB #include diff --git a/ggml/src/ggml-cuda/sum.cuh b/ggml/src/ggml-cuda/sum.cuh.inc similarity index 87% rename from ggml/src/ggml-cuda/sum.cuh rename to ggml/src/ggml-cuda/sum.cuh.inc index 8cadc3736f07..6ebf5397a3d6 100644 --- a/ggml/src/ggml-cuda/sum.cuh +++ b/ggml/src/ggml-cuda/sum.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" void sum_f32_cuda(ggml_cuda_pool & pool, const float * x, float * dst, const int64_t ne, cudaStream_t stream); diff --git a/ggml/src/ggml-cuda/sumrows.cu b/ggml/src/ggml-cuda/sumrows.cu.inc similarity index 96% rename from ggml/src/ggml-cuda/sumrows.cu rename to ggml/src/ggml-cuda/sumrows.cu.inc index 4025771aadb9..28d2f745e9ce 100644 --- a/ggml/src/ggml-cuda/sumrows.cu +++ b/ggml/src/ggml-cuda/sumrows.cu.inc @@ -1,5 +1,5 @@ -#include "reduce_rows.cuh" -#include "sumrows.cuh" +#include "reduce_rows.cuh.inc" +#include "sumrows.cuh.inc" void sum_rows_f32_cuda(const float * x, float * dst, const int ncols, const int nrows, cudaStream_t stream) { const int id = ggml_cuda_get_device(); diff --git a/ggml/src/ggml-cuda/sumrows.cuh b/ggml/src/ggml-cuda/sumrows.cuh.inc similarity index 87% rename from ggml/src/ggml-cuda/sumrows.cuh rename to ggml/src/ggml-cuda/sumrows.cuh.inc index 3431c599b1b8..5ddcfe40eb05 100644 --- a/ggml/src/ggml-cuda/sumrows.cuh +++ b/ggml/src/ggml-cuda/sumrows.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" void sum_rows_f32_cuda(const float * x, float * dst, const int ncols, const int nrows, cudaStream_t stream); void ggml_cuda_op_sum_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_16.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_16.cu.inc similarity index 77% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_16.cu.inc index fb26abeb0dab..3cedc51946ca 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_16.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(576, 512, 1, 16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_32.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_32.cu.inc similarity index 82% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_32.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_32.cu.inc index 8fc3b17976e7..aa716b9ea364 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_32.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_32.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(320, 256, 1, 32); DECL_FATTN_MMA_F16_CASE(576, 512, 1, 32); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu.inc similarity index 91% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu.inc index 22d383173f36..0f90edda1089 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 1, 8); DECL_FATTN_MMA_F16_CASE(80, 80, 1, 8); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_1.cu.inc similarity index 90% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_1.cu.inc index 9d3cfd8edf74..e00bc54182f2 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 16, 1); DECL_FATTN_MMA_F16_CASE(80, 80, 16, 1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_2.cu.inc similarity index 90% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_2.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_2.cu.inc index 2e1883af40ed..ef30c4bd4f2c 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_2.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_2.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 16, 2); DECL_FATTN_MMA_F16_CASE(80, 80, 16, 2); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_4.cu.inc similarity index 92% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_4.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_4.cu.inc index d2415bfa957f..39c5b0ca54f9 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_4.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_4.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 16, 4); DECL_FATTN_MMA_F16_CASE(80, 80, 16, 4); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_16.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_16.cu.inc similarity index 77% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_16.cu.inc index f011a208cd27..307309c935d1 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_16.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(576, 512, 2, 16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_32.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_32.cu.inc similarity index 82% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_32.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_32.cu.inc index abd2b21ce044..de2d276546f1 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_32.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_32.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(320, 256, 2, 32); DECL_FATTN_MMA_F16_CASE(576, 512, 2, 32); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_4.cu.inc similarity index 91% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_4.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_4.cu.inc index 8eec1d74e293..20ec67c4f869 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_4.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_4.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 2, 4); DECL_FATTN_MMA_F16_CASE(80, 80, 2, 4); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu.inc similarity index 91% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu.inc index 84b674cd05a6..2e4796289c91 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 2, 8); DECL_FATTN_MMA_F16_CASE(80, 80, 2, 8); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_1.cu.inc similarity index 90% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_1.cu.inc index 0543532ea347..2abf540605f9 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 32, 1); DECL_FATTN_MMA_F16_CASE(80, 80, 32, 1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_2.cu.inc similarity index 90% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_2.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_2.cu.inc index 407b6cf4c702..8acb36786014 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_2.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_2.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 32, 2); DECL_FATTN_MMA_F16_CASE(80, 80, 32, 2); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_16.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_16.cu.inc similarity index 77% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_16.cu.inc index f5fd0e2369cf..49f1dfa5801a 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_16.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(576, 512, 4, 16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_2.cu.inc similarity index 90% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_2.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_2.cu.inc index 5e46685024b8..a046fb20d091 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_2.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_2.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 4, 2); DECL_FATTN_MMA_F16_CASE(80, 80, 4, 2); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_4.cu.inc similarity index 91% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_4.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_4.cu.inc index 3475dfea08a1..d63f6e5be4dc 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_4.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_4.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 4, 4); DECL_FATTN_MMA_F16_CASE(80, 80, 4, 4); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu.inc similarity index 91% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu.inc index 5906398db912..3616f5d933ab 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 4, 8); DECL_FATTN_MMA_F16_CASE(80, 80, 4, 8); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_64-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_64-ncols2_1.cu.inc similarity index 90% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_64-ncols2_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_64-ncols2_1.cu.inc index 0d7a9c728537..0397ebe3f63e 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_64-ncols2_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_64-ncols2_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 64, 1); DECL_FATTN_MMA_F16_CASE(80, 80, 64, 1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_1.cu.inc similarity index 90% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_1.cu.inc index 9d5a9976f0ed..1238cb81463c 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 8, 1); DECL_FATTN_MMA_F16_CASE(80, 80, 8, 1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_2.cu.inc similarity index 90% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_2.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_2.cu.inc index a6e6f093dcb2..c73e0bc10e61 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_2.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_2.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 8, 2); DECL_FATTN_MMA_F16_CASE(80, 80, 8, 2); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_4.cu.inc similarity index 91% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_4.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_4.cu.inc index 684cd25ce0de..66228815fff4 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_4.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_4.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 8, 4); DECL_FATTN_MMA_F16_CASE(80, 80, 8, 4); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu.inc similarity index 91% rename from ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu rename to ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu.inc index 4bc60d62f910..b37a9d651335 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-mma-f16.cuh" +#include "../fattn-mma-f16.cuh.inc" DECL_FATTN_MMA_F16_CASE(64, 64, 8, 8); DECL_FATTN_MMA_F16_CASE(80, 80, 8, 8); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq112-dv112.cu b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq112-dv112.cu deleted file mode 100644 index a8b15ad72a91..000000000000 --- a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq112-dv112.cu +++ /dev/null @@ -1,5 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-tile.cuh" - -DECL_FATTN_TILE_CASE(112, 112); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq112-dv112.cpp b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq112-dv112.cu.inc similarity index 78% rename from ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq112-dv112.cpp rename to ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq112-dv112.cu.inc index 5c06d42fdbdd..e8541dc15d78 100644 --- a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq112-dv112.cpp +++ b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq112-dv112.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-tile.hpp" +#include "../fattn-tile.cuh.inc" DECL_FATTN_TILE_CASE(112, 112); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq128-dv128.cu b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq128-dv128.cu.inc similarity index 78% rename from ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq128-dv128.cu rename to ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq128-dv128.cu.inc index 1da18105508a..8f10f47ffe32 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq128-dv128.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq128-dv128.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-tile.cuh" +#include "../fattn-tile.cuh.inc" DECL_FATTN_TILE_CASE(128, 128); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq256-dv256.cu b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq256-dv256.cu deleted file mode 100644 index bc65c723eca9..000000000000 --- a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq256-dv256.cu +++ /dev/null @@ -1,5 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-tile.cuh" - -DECL_FATTN_TILE_CASE(256, 256); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq256-dv256.cpp b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq256-dv256.cu.inc similarity index 78% rename from ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq256-dv256.cpp rename to ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq256-dv256.cu.inc index b574fe9308dc..5cc267c9c6ad 100644 --- a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq256-dv256.cpp +++ b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq256-dv256.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-tile.hpp" +#include "../fattn-tile.cuh.inc" DECL_FATTN_TILE_CASE(256, 256); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq320-dv256.cu b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq320-dv256.cu.inc similarity index 78% rename from ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq320-dv256.cu rename to ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq320-dv256.cu.inc index c91f508079d5..e9536d021e9b 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq320-dv256.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq320-dv256.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-tile.cuh" +#include "../fattn-tile.cuh.inc" DECL_FATTN_TILE_CASE(320, 256); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq40-dv40.cu b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq40-dv40.cu.inc similarity index 77% rename from ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq40-dv40.cu rename to ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq40-dv40.cu.inc index 10b330fa6c03..88e6cb4413e1 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq40-dv40.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq40-dv40.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-tile.cuh" +#include "../fattn-tile.cuh.inc" DECL_FATTN_TILE_CASE(40, 40); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq512-dv512.cu b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq512-dv512.cu.inc similarity index 78% rename from ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq512-dv512.cu rename to ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq512-dv512.cu.inc index 7c61d8d2ecdb..de30c7fce345 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq512-dv512.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq512-dv512.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-tile.cuh" +#include "../fattn-tile.cuh.inc" DECL_FATTN_TILE_CASE(512, 512); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq576-dv512.cu b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq576-dv512.cu.inc similarity index 78% rename from ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq576-dv512.cu rename to ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq576-dv512.cu.inc index 254b7d2e1dc2..96e84ed585ef 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq576-dv512.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq576-dv512.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-tile.cuh" +#include "../fattn-tile.cuh.inc" DECL_FATTN_TILE_CASE(576, 512); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq64-dv64.cu b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq64-dv64.cu.inc similarity index 77% rename from ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq64-dv64.cu rename to ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq64-dv64.cu.inc index 5caffac0467d..15662212a6d9 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq64-dv64.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq64-dv64.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-tile.cuh" +#include "../fattn-tile.cuh.inc" DECL_FATTN_TILE_CASE(64, 64); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq72-dv72.cu b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq72-dv72.cu.inc similarity index 77% rename from ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq72-dv72.cu rename to ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq72-dv72.cu.inc index 8f9d5315f2ac..80a8fd54cee3 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq72-dv72.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq72-dv72.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-tile.cuh" +#include "../fattn-tile.cuh.inc" DECL_FATTN_TILE_CASE(72, 72); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq80-dv80.cu b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq80-dv80.cu.inc similarity index 77% rename from ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq80-dv80.cu rename to ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq80-dv80.cu.inc index 90abb3b18626..010e0248cc12 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq80-dv80.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq80-dv80.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-tile.cuh" +#include "../fattn-tile.cuh.inc" DECL_FATTN_TILE_CASE(80, 80); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq96-dv96.cu b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq96-dv96.cu.inc similarity index 77% rename from ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq96-dv96.cu rename to ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq96-dv96.cu.inc index 7292c0aab8f9..f5974861dea6 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq96-dv96.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-tile-instance-dkq96-dv96.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-tile.cuh" +#include "../fattn-tile.cuh.inc" DECL_FATTN_TILE_CASE(96, 96); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu.inc index 3a2fa99b05ba..87fc31de4216 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_BF16, GGML_TYPE_BF16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_BF16, GGML_TYPE_BF16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-f16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-f16.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-f16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-f16.cu.inc index 60f0f6f79527..f4b2a7ceace8 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-f16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-f16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_BF16, GGML_TYPE_F16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_BF16, GGML_TYPE_F16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q4_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q4_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q4_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q4_0.cu.inc index 489e05f08c39..33805dfbdd2d 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q4_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q4_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_BF16, GGML_TYPE_Q4_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_BF16, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q4_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q4_1.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q4_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q4_1.cu.inc index 6fa3c26d3091..3b50aa5cfffd 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q4_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q4_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_BF16, GGML_TYPE_Q4_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_BF16, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q5_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q5_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q5_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q5_0.cu.inc index 421027fb29d0..c38e287816c1 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q5_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q5_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_BF16, GGML_TYPE_Q5_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_BF16, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q5_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q5_1.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q5_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q5_1.cu.inc index abbc94348023..9315202fc18c 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q5_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q5_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_BF16, GGML_TYPE_Q5_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_BF16, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q8_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q8_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q8_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q8_0.cu.inc index d641f859d81b..706061d2ffa1 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q8_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-bf16-q8_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_BF16, GGML_TYPE_Q8_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_BF16, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-bf16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-bf16.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-bf16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-bf16.cu.inc index d1071dc24381..54215a2bed0e 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-bf16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-bf16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_BF16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_BF16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu.inc index c357abd80d3c..2b03e0edd2a7 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_F16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_F16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q4_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q4_0.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q4_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q4_0.cu.inc index 4b148656f929..609c746c50ec 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q4_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q4_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q4_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q4_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q4_1.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q4_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q4_1.cu.inc index ef7715758c91..3a6faf15a024 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q4_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q4_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q4_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q5_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q5_0.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q5_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q5_0.cu.inc index 9ae11cc5423c..e8715f21fa8a 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q5_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q5_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q5_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q5_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q5_1.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q5_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q5_1.cu.inc index 10ed48affa47..633fa2e0dce4 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q5_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q5_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q5_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q8_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q8_0.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q8_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q8_0.cu.inc index 4fcc3f337764..bbfbe5bb768f 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q8_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-f16-q8_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q8_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-bf16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-bf16.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-bf16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-bf16.cu.inc index 8afda3142389..2ba53cea83c3 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-bf16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-bf16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_BF16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_BF16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-f16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-f16.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-f16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-f16.cu.inc index 7ca50531fb24..a194926811c4 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-f16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-f16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_F16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_F16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu.inc index 6ef1a48fdb02..f19bf72f61da 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_1.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_1.cu.inc index 4c0532ca7ebb..7576589870f7 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q5_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q5_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q5_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q5_0.cu.inc index ed3d7bad3953..ee85d0272afc 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q5_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q5_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_Q5_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q5_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q5_1.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q5_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q5_1.cu.inc index 687f25406813..a1080d6dfb85 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q5_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q5_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_Q5_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q8_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q8_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q8_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q8_0.cu.inc index 41107c45f464..167f98f7c18f 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q8_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_0-q8_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-bf16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-bf16.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-bf16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-bf16.cu.inc index 506864ac18dc..f607fe2946f5 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-bf16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-bf16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_BF16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_BF16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-f16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-f16.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-f16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-f16.cu.inc index d523ce01cc58..dbf42223c85a 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-f16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-f16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_F16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_F16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q4_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q4_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q4_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q4_0.cu.inc index 8b9ed358eca2..98ab1a6ff8cc 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q4_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q4_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q4_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q4_1.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q4_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q4_1.cu.inc index 0553e464c49d..2586d8b69058 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q4_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q4_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_Q4_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q5_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q5_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q5_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q5_0.cu.inc index 8390eaf1c88b..9f4555a3c493 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q5_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q5_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q5_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q5_1.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q5_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q5_1.cu.inc index f61e19d6a390..276a1e74e389 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q5_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q5_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_Q5_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q8_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q8_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q8_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q8_0.cu.inc index 86a188269c7c..6a78a9ee2f49 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q8_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q4_1-q8_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_Q8_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-bf16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-bf16.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-bf16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-bf16.cu.inc index 0bbda8371e62..0461e877037a 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-bf16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-bf16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_BF16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_BF16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-f16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-f16.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-f16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-f16.cu.inc index 1d7af474b484..ccc299d01042 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-f16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-f16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_F16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_F16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q4_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q4_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q4_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q4_0.cu.inc index 837224d36095..dcdc8a64308b 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q4_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q4_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_Q4_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q4_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q4_1.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q4_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q4_1.cu.inc index 0dd7dd693f16..8ff81e200ea9 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q4_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q4_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q5_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q5_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q5_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q5_0.cu.inc index 41b859f45d72..c51f5d8b3e17 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q5_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q5_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_Q5_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q5_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q5_1.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q5_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q5_1.cu.inc index d2e5ffd0ac58..75b80f6b60fe 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q5_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q5_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q8_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q8_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q8_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q8_0.cu.inc index 81ff740b5852..164729e911ae 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q8_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_0-q8_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-bf16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-bf16.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-bf16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-bf16.cu.inc index 79be24daf9e0..9eb827566af8 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-bf16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-bf16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_BF16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_BF16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-f16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-f16.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-f16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-f16.cu.inc index a38dae19221e..40aba5da859f 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-f16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-f16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_F16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_F16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q4_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q4_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q4_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q4_0.cu.inc index 2304571e2404..923aa99904b5 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q4_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q4_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_Q4_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q4_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q4_1.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q4_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q4_1.cu.inc index 84b83e5544ca..da32d11f21bd 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q4_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q4_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_Q4_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q5_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q5_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q5_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q5_0.cu.inc index 39f80e218d36..3a3617ae9b65 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q5_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q5_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q5_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q5_1.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q5_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q5_1.cu.inc index cf4e66112b65..16c320c3b1f4 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q5_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q5_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_Q5_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q8_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q8_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q8_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q8_0.cu.inc index 65654182e552..ecc3fb3d9daf 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q8_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q5_1-q8_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-bf16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-bf16.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-bf16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-bf16.cu.inc index 45636e5e70c6..07e592e3f007 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-bf16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-bf16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_BF16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_BF16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-f16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-f16.cu.inc similarity index 88% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-f16.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-f16.cu.inc index a1bc3f5a6aa3..70fc38ddea88 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-f16.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-f16.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_F16); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_F16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q4_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q4_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q4_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q4_0.cu.inc index 4b76a9be232f..3f5a058a25b5 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q4_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q4_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q4_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q4_1.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q4_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q4_1.cu.inc index 77d04125f7b4..45410dff65cf 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q4_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q4_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_Q4_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q5_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q5_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q5_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q5_0.cu.inc index 6e170fe36f2c..27f898466e24 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q5_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q5_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_Q5_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q5_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q5_1.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q5_1.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q5_1.cu.inc index b617cd73b567..d42240540434 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q5_1.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q5_1.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu.inc similarity index 89% rename from ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu rename to ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu.inc index a5b768b111b8..a6ddd63c3a9b 100644 --- a/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu.inc @@ -1,6 +1,6 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../fattn-vec.cuh" +#include "../fattn-vec.cuh.inc" DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0); DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_1.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_1.cu.inc similarity index 79% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_1.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_1.cu.inc index f594d5d51d29..1a2a0578aef2 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_1.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_1.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(1); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_10.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_10.cu.inc similarity index 80% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_10.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_10.cu.inc index 9cc67725421c..52420676e38e 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_10.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_10.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(10); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_11.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_11.cu.inc similarity index 80% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_11.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_11.cu.inc index 317f487d7a79..526ae0a721b2 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_11.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_11.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(11); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_12.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_12.cu.inc similarity index 80% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_12.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_12.cu.inc index dc0033227c0e..0fde70ad97c8 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_12.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_12.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(12); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_13.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_13.cu.inc similarity index 80% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_13.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_13.cu.inc index 078210175306..c687fc6873b2 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_13.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_13.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(13); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_14.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_14.cu.inc similarity index 80% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_14.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_14.cu.inc index a23ad6ae262d..8dcc9bb85b29 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_14.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_14.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(14); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_15.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_15.cu.inc similarity index 80% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_15.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_15.cu.inc index 0fe3f7821eed..3eb42cfdee70 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_15.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_15.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(15); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_16.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_16.cu.inc similarity index 80% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_16.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_16.cu.inc index 544086375e88..3c6f374738e5 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_16.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_16.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(16); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_2.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_2.cu.inc similarity index 79% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_2.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_2.cu.inc index 3b901797cfb7..09d7ed5b16ac 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_2.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_2.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(2); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_3.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_3.cu.inc similarity index 79% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_3.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_3.cu.inc index 56e940bba08b..6f73736a87a7 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_3.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_3.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(3); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_4.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_4.cu.inc similarity index 79% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_4.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_4.cu.inc index a7665d49d0b7..b1bef141b9fa 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_4.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_4.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(4); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_5.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_5.cu.inc similarity index 79% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_5.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_5.cu.inc index 3a1dff2587a1..45e4de6ae3fc 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_5.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_5.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(5); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_6.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_6.cu.inc similarity index 79% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_6.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_6.cu.inc index 400fb7c66310..abab5dbd179f 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_6.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_6.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(6); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_7.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_7.cu.inc similarity index 79% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_7.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_7.cu.inc index 954a1c7e032f..6b3b48807c3f 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_7.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_7.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(7); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_8.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_8.cu.inc similarity index 79% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_8.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_8.cu.inc index f1bd09c9458e..545dc4bd597e 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_8.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_8.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(8); diff --git a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_9.cu b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_9.cu.inc similarity index 79% rename from ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_9.cu rename to ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_9.cu.inc index 1255ac2af661..82b650e08642 100644 --- a/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_9.cu +++ b/ggml/src/ggml-cuda/template-instances/mmf-instance-ncols_9.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmf.cuh" +#include "../mmf.cuh.inc" DECL_MMF_CASE(9); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq1_s.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq1_s.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-iq1_s.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-iq1_s.cu.inc index 84ec85029473..fe5a1cddc382 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq1_s.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq1_s.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_IQ1_S); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_s.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_s.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_s.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_s.cu.inc index 583c4e5a5150..1546c3fe049f 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_s.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_s.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_IQ2_S); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu.inc index edaf1560defd..582f2e545b7b 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_xs.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_IQ2_XS); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu.inc similarity index 82% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu.inc index 233d9342c994..82013c644a50 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq2_xxs.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_IQ2_XXS); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq3_s.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq3_s.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-iq3_s.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-iq3_s.cu.inc index 6092dc713634..8b35c15b4eb5 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq3_s.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq3_s.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_IQ3_S); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu.inc similarity index 82% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu.inc index 1d5bd201fb5a..7290e8a79eec 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq3_xxs.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_IQ3_XXS); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu.inc index eb02fab002ec..ab3f5d9680db 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_nl.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_IQ4_NL); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu.inc index 1eb3b7430726..fb6c87aa05ba 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-iq4_xs.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu.inc index c14624c52cad..f596eb5677a6 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-mxfp4.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_MXFP4); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu.inc index 2cb140d35a3c..5e79d3b9ecf1 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-nvfp4.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_NVFP4); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-q1_0.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-q1_0.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-q1_0.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-q1_0.cu.inc index f0686b0d0d85..85db79c842c2 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-q1_0.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-q1_0.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_Q1_0); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_k.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_k.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-q2_k.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-q2_k.cu.inc index 6415369dc1d9..d6ed7767f374 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_k.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_k.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_Q2_K); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_k.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_k.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-q3_k.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-q3_k.cu.inc index ffb6213af83e..8c387055ccb3 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_k.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_k.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_Q3_K); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0.cu.inc index 0c0b0c8a8ed2..45ba235b7628 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_1.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_1.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-q4_1.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-q4_1.cu.inc index ee67f6942a8f..5e30dc5b708e 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_1.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_1.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu.inc index 9eeb3cd7f3cc..3543bdcc3638 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_k.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_Q4_K); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_0.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_0.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-q5_0.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-q5_0.cu.inc index cc57fb9753c9..079cf7f730ef 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_0.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_0.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_1.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_1.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-q5_1.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-q5_1.cu.inc index 721ac790c44f..6ec4222d1cdb 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_1.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_1.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu.inc index a2e90ffd5d0a..31f6775bde8c 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-q5_k.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_Q5_K); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu.inc index 470938fef8a0..434d500b9373 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-q6_k.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_Q6_K); diff --git a/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu b/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu.inc similarity index 81% rename from ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu rename to ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu.inc index 974477bbb73a..a30d20a748c1 100644 --- a/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu +++ b/ggml/src/ggml-cuda/template-instances/mmq-instance-q8_0.cu.inc @@ -1,5 +1,5 @@ // This file has been autogenerated by generate_cu_files.py, do not edit manually. -#include "../mmq.cuh" +#include "../mmq.cuh.inc" DECL_MMQ_CASE(GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/top-k.cu rename to ggml/src/ggml-cuda/top-k.cu.inc index 59ce36fb1c95..730759b9d1ab 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu.inc @@ -1,5 +1,5 @@ -#include "argsort.cuh" -#include "top-k.cuh" +#include "argsort.cuh.inc" +#include "top-k.cuh.inc" #ifdef GGML_CUDA_USE_CUB # include diff --git a/ggml/src/ggml-cuda/top-k.cuh b/ggml/src/ggml-cuda/top-k.cuh.inc similarity index 75% rename from ggml/src/ggml-cuda/top-k.cuh rename to ggml/src/ggml-cuda/top-k.cuh.inc index f4d8f61e5b3f..480891aed2f4 100644 --- a/ggml/src/ggml-cuda/top-k.cuh +++ b/ggml/src/ggml-cuda/top-k.cuh.inc @@ -1,3 +1,3 @@ -#include "common.cuh" +#include "common.cuh.inc" void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/topk-moe.cu b/ggml/src/ggml-cuda/topk-moe.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/topk-moe.cu rename to ggml/src/ggml-cuda/topk-moe.cu.inc index 3020e5c74336..31566613d703 100644 --- a/ggml/src/ggml-cuda/topk-moe.cu +++ b/ggml/src/ggml-cuda/topk-moe.cu.inc @@ -1,6 +1,6 @@ -#include "ggml-cuda/common.cuh" -#include "ggml.h" -#include "topk-moe.cuh" +#include "ggml-cuda/common.cuh.inc" +#include "ggml.h.inc" +#include "topk-moe.cuh.inc" #include #include diff --git a/ggml/src/ggml-cuda/topk-moe.cuh b/ggml/src/ggml-cuda/topk-moe.cuh.inc similarity index 95% rename from ggml/src/ggml-cuda/topk-moe.cuh rename to ggml/src/ggml-cuda/topk-moe.cuh.inc index 243dc2f1c41b..82bc5fec0564 100644 --- a/ggml/src/ggml-cuda/topk-moe.cuh +++ b/ggml/src/ggml-cuda/topk-moe.cuh.inc @@ -1,5 +1,5 @@ -#include "common.cuh" -#include "ggml.h" +#include "common.cuh.inc" +#include "ggml.h.inc" #include diff --git a/ggml/src/ggml-cuda/tri.cu b/ggml/src/ggml-cuda/tri.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/tri.cu rename to ggml/src/ggml-cuda/tri.cu.inc index 44156b63e705..bd3bcc4bec3e 100644 --- a/ggml/src/ggml-cuda/tri.cu +++ b/ggml/src/ggml-cuda/tri.cu.inc @@ -1,7 +1,7 @@ -#include "common.cuh" -#include "convert.cuh" -#include "tri.cuh" -#include "ggml.h" +#include "common.cuh.inc" +#include "convert.cuh.inc" +#include "tri.cuh.inc" +#include "ggml.h.inc" template static __global__ void tri_kernel( diff --git a/ggml/src/ggml-cuda/tri.cuh b/ggml/src/ggml-cuda/tri.cuh.inc similarity index 80% rename from ggml/src/ggml-cuda/tri.cuh rename to ggml/src/ggml-cuda/tri.cuh.inc index a4cc66750d3b..328e59ae3264 100644 --- a/ggml/src/ggml-cuda/tri.cuh +++ b/ggml/src/ggml-cuda/tri.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_TRI_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/tsembd.cu b/ggml/src/ggml-cuda/tsembd.cu.inc similarity index 98% rename from ggml/src/ggml-cuda/tsembd.cu rename to ggml/src/ggml-cuda/tsembd.cu.inc index b91a26fc80e6..107ab99a2d4a 100644 --- a/ggml/src/ggml-cuda/tsembd.cu +++ b/ggml/src/ggml-cuda/tsembd.cu.inc @@ -1,4 +1,4 @@ -#include "tsembd.cuh" +#include "tsembd.cuh.inc" static __global__ void timestep_embedding_f32(const float * timesteps, float * dst, const int nb1, const int dim, const int max_period) { // blockIDx.y: idx of timesteps->ne[0] diff --git a/ggml/src/ggml-cuda/tsembd.cuh b/ggml/src/ggml-cuda/tsembd.cuh.inc similarity index 84% rename from ggml/src/ggml-cuda/tsembd.cuh rename to ggml/src/ggml-cuda/tsembd.cuh.inc index 84340e3d7d2c..91b7a222075f 100644 --- a/ggml/src/ggml-cuda/tsembd.cuh +++ b/ggml/src/ggml-cuda/tsembd.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_TIMESTEP_EMBEDDING_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/unary.cu rename to ggml/src/ggml-cuda/unary.cu.inc index 2aeba26f4140..2e07ffdfe48e 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu.inc @@ -1,5 +1,5 @@ -#include "unary.cuh" -#include "convert.cuh" +#include "unary.cuh.inc" +#include "convert.cuh.inc" static __device__ __forceinline__ float op_abs(float x) { return fabsf(x); diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh.inc similarity index 99% rename from ggml/src/ggml-cuda/unary.cuh rename to ggml/src/ggml-cuda/unary.cuh.inc index 81ed873ecc30..84d4b3bc985f 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh.inc @@ -1,5 +1,5 @@ #pragma once -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_NEG_BLOCK_SIZE 256 #define CUDA_STEP_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/upscale.cu b/ggml/src/ggml-cuda/upscale.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/upscale.cu rename to ggml/src/ggml-cuda/upscale.cu.inc index 6bdf3cd996bf..23418f37ffd6 100644 --- a/ggml/src/ggml-cuda/upscale.cu +++ b/ggml/src/ggml-cuda/upscale.cu.inc @@ -1,4 +1,4 @@ -#include "upscale.cuh" +#include "upscale.cuh.inc" static __global__ void upscale_f32(const float * x, float * dst, const int nb00, const int nb01, const int nb02, const int nb03, diff --git a/ggml/src/ggml-cuda/upscale.cuh b/ggml/src/ggml-cuda/upscale.cuh.inc similarity index 81% rename from ggml/src/ggml-cuda/upscale.cuh rename to ggml/src/ggml-cuda/upscale.cuh.inc index d4d7652308e9..faee6fac59e4 100644 --- a/ggml/src/ggml-cuda/upscale.cuh +++ b/ggml/src/ggml-cuda/upscale.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_UPSCALE_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh.inc similarity index 99% rename from ggml/src/ggml-cuda/vecdotq.cuh rename to ggml/src/ggml-cuda/vecdotq.cuh.inc index d1741cc8d7ba..e695fb918e81 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh.inc @@ -1,6 +1,6 @@ #pragma once -#include "common.cuh" +#include "common.cuh.inc" #include diff --git a/ggml/src/ggml-cuda/vendors/cuda.h b/ggml/src/ggml-cuda/vendors/cuda.h.inc similarity index 97% rename from ggml/src/ggml-cuda/vendors/cuda.h rename to ggml/src/ggml-cuda/vendors/cuda.h.inc index 323c98019347..0745ea45f0ed 100644 --- a/ggml/src/ggml-cuda/vendors/cuda.h +++ b/ggml/src/ggml-cuda/vendors/cuda.h.inc @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include #include #include diff --git a/ggml/src/ggml-cuda/vendors/hip.h b/ggml/src/ggml-cuda/vendors/hip.h.inc similarity index 100% rename from ggml/src/ggml-cuda/vendors/hip.h rename to ggml/src/ggml-cuda/vendors/hip.h.inc diff --git a/ggml/src/ggml-cuda/vendors/musa.h b/ggml/src/ggml-cuda/vendors/musa.h.inc similarity index 99% rename from ggml/src/ggml-cuda/vendors/musa.h rename to ggml/src/ggml-cuda/vendors/musa.h.inc index 8aa056e91743..fd14ed64f24f 100644 --- a/ggml/src/ggml-cuda/vendors/musa.h +++ b/ggml/src/ggml-cuda/vendors/musa.h.inc @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include #include #include diff --git a/ggml/src/ggml-cuda/wkv.cu b/ggml/src/ggml-cuda/wkv.cu.inc similarity index 99% rename from ggml/src/ggml-cuda/wkv.cu rename to ggml/src/ggml-cuda/wkv.cu.inc index d2fced705e09..e8240b826b10 100644 --- a/ggml/src/ggml-cuda/wkv.cu +++ b/ggml/src/ggml-cuda/wkv.cu.inc @@ -1,5 +1,5 @@ -#include "common.cuh" -#include "wkv.cuh" +#include "common.cuh.inc" +#include "wkv.cuh.inc" template static __global__ void rwkv_wkv_f32(const int B, const int T, const int C, const int H, const float * k, const float * v, const float * r, const float * tf, const float * td, const float * s, float * dst) { diff --git a/ggml/src/ggml-cuda/wkv.cuh b/ggml/src/ggml-cuda/wkv.cuh.inc similarity index 88% rename from ggml/src/ggml-cuda/wkv.cuh rename to ggml/src/ggml-cuda/wkv.cuh.inc index 9623dd7f8c7a..cee696775bf2 100644 --- a/ggml/src/ggml-cuda/wkv.cuh +++ b/ggml/src/ggml-cuda/wkv.cuh.inc @@ -1,4 +1,4 @@ -#include "common.cuh" +#include "common.cuh.inc" #define CUDA_WKV_BLOCK_SIZE 64 diff --git a/ggml/src/ggml-hexagon/CMakeLists.txt b/ggml/src/ggml-hexagon/CMakeLists.txt index b82bae0c1035..253776b9f8b2 100644 --- a/ggml/src/ggml-hexagon/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/CMakeLists.txt @@ -49,12 +49,11 @@ else() endif() set(TARGET_NAME ggml-hexagon) +set_source_files_properties(ggml-hexagon.cpp.inc htp-drv.cpp.inc PROPERTIES LANGUAGE CXX) ggml_add_backend_library(${TARGET_NAME} - ggml-hexagon.cpp - htp-drv.cpp - htp-drv.h - libdl.h - ../../include/ggml-hexagon.h) + ggml-hexagon.cpp.inc + htp-drv.cpp.inc + ../../include/ggml-hexagon.h.inc) target_link_libraries(${TARGET_NAME} PRIVATE htp_iface) target_include_directories(${TARGET_NAME} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/htp ${CMAKE_CURRENT_BINARY_DIR}) diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp.inc similarity index 94% rename from ggml/src/ggml-hexagon/ggml-hexagon.cpp rename to ggml/src/ggml-hexagon/ggml-hexagon.cpp.inc index df4ed1014642..3e966430c916 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp.inc @@ -34,15 +34,283 @@ #include #define GGML_COMMON_IMPL_CPP -#include "ggml-backend-impl.h" -#include "ggml-common.h" -#include "ggml-hexagon.h" -#include "ggml-impl.h" -#include "ggml-quants.h" -#include "op-desc.h" -#include "htp-ops.h" +#include "ggml-backend-impl-defs.inc" +#include "ggml-common-defs.inc" +#include "ggml-hexagon.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-quants-defs.inc" +#define GGML_COMMON_IMPL_CPP +#include "ggml-backend-impl-defs.inc" +#include "ggml-common-defs.inc" + +#include +#include + +struct op_desc { + char strides[64 * GGML_MAX_SRC]; + char dims[64 * GGML_MAX_SRC]; + char types[16 * GGML_MAX_SRC]; + char buffs[64 * GGML_MAX_SRC]; + char names[64 * GGML_MAX_SRC]; + + int format_tensor_dims(char * str, const struct ggml_tensor * t) { + if (t->ne[2] == 1 && t->ne[3] == 1) { + return sprintf(str, "%d:%d", (int) t->ne[0], (int) t->ne[1]); + } else { + return sprintf(str, "%d:%d:%d:%d", (int) t->ne[0], (int) t->ne[1], (int) t->ne[2], (int) t->ne[3]); + } + } + + void format_op_dims(char * str, const struct ggml_tensor * t) { + char * p = str; + + // append src0 and src1 (if any) + if (t->src[0]) { + p += format_tensor_dims(p, t->src[0]); + + for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) { + p += sprintf(p, " x "); + p += format_tensor_dims(p, t->src[i]); + } + + p += sprintf(p, " -> "); + } + + // format self dims separately for better visual alignment + char self[64]; + format_tensor_dims(self, t); + + p += sprintf(p, "%s", self); + } + + int format_tensor_strides(char * str, const struct ggml_tensor * t) { + const char * c = ggml_is_contiguous(t) ? "" : "!"; + + if (t->ne[2] == 1 && t->ne[3] == 1) { + return sprintf(str, "%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], c); + } else { + return sprintf(str, "%zu:%zu:%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], (size_t) t->nb[2], (size_t) t->nb[3], c); + } + } + + void format_op_strides(char * str, const struct ggml_tensor * t) { + char * p = str; + + // append src0 and src1 (if any) + if (t->src[0]) { + p += format_tensor_strides(p, t->src[0]); + + for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) { + p += sprintf(p, " x "); + p += format_tensor_strides(p, t->src[i]); + } + + p += sprintf(p, " -> "); + } + + // format self dims separately for better visual alignment + char self[64]; + format_tensor_strides(self, t); + + p += sprintf(p, "%s", self); + } + + void format_op_types(char * str, const struct ggml_tensor * t) { + char * p = str; + + // append src0 and src1 (if any) + if (t->src[0]) { + p += sprintf(p, "%s", ggml_type_name(t->src[0]->type)); + + for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) { + p += sprintf(p, " x "); + p += sprintf(p, "%s", ggml_type_name(t->src[i]->type)); + } + + p += sprintf(p, " -> "); + } + + p += sprintf(p, "%s", ggml_type_name(t->type)); + } + + const char * tensor_buff_name(const struct ggml_tensor * t) { + if (t->buffer) { + return ggml_backend_buffer_name(t->buffer); + } + return "NONE"; + } + + void format_op_buffs(char * str, const struct ggml_tensor * t) { + char * p = str; + + // append src0 and src1 (if any) + if (t->src[0]) { + p += sprintf(p, "%s", tensor_buff_name(t->src[0])); + + for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) { + p += sprintf(p, " x "); + p += sprintf(p, "%s", tensor_buff_name(t->src[i])); + } + + p += sprintf(p, " -> "); + } + + p += sprintf(p, "%s", tensor_buff_name(t)); + } + + void format_op_names(char * str, const struct ggml_tensor * t) { + char * p = str; + + // append src0 and src1 (if any) + if (t->src[0]) { + p += sprintf(p, "%s", t->src[0]->name); + + for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) { + p += sprintf(p, " x "); + p += sprintf(p, "%s", t->src[i]->name); + } + + p += sprintf(p, " -> "); + } + + p += sprintf(p, "%s", t->name); + } + + void format(const ggml_tensor * op) { + format_op_dims(dims, op); + format_op_strides(strides, op); + format_op_types(types, op); + format_op_buffs(buffs, op); + format_op_names(names, op); + } + + op_desc() {} + op_desc(const ggml_tensor * op) { format(op); } +}; + + +#include "htp-ops-defs.inc" #include "htp_iface.h" -#include "htp-drv.h" +#ifdef __cplusplus +extern "C" { +#endif + +#ifdef _WIN32 +# pragma clang diagnostic ignored "-Wignored-attributes" +#endif + +#include +#include +#include +#include + +#if defined(_WIN32) && !defined(__MINGW32__) +# ifdef GGML_BACKEND_BUILD +# define HTPDRV_API __declspec(dllexport) extern +# else +# define HTPDRV_API __declspec(dllimport) extern +# endif +#else +# define HTPDRV_API __attribute__ ((visibility ("default"))) extern +#endif + +/* Offset to differentiate HLOS and Hexagon error codes. + Stores the value of AEE_EOFFSET for Hexagon. */ +#ifndef DSP_OFFSET +# define DSP_OFFSET 0x80000400 +#endif + +/* Errno for connection reset by peer. */ +#ifndef ECONNRESET +# ifdef __hexagon__ +# define ECONNRESET 104 +# endif +#endif + +/* Abstraction of different OS specific sleep APIs. + SLEEP accepts input in seconds. */ +#ifndef SLEEP +# ifdef __hexagon__ +# define SLEEP(x) \ + { /* Do nothing for simulator. */ \ + } +# else +# ifdef _WIN32 +# define SLEEP(x) Sleep(1000 * x) /* Sleep accepts input in milliseconds. */ +# else +# define SLEEP(x) sleep(x) /* sleep accepts input in seconds. */ +# endif +# endif +#endif + +/* Include windows specific header files. */ +#ifdef _WIN32 +# include +# include +# define _CRT_SECURE_NO_WARNINGS 1 +# define _WINSOCK_DEPRECATED_NO_WARNINGS 1 +#endif + +/* Includes and defines for all HLOS except windows */ +#if !defined(__hexagon__) && !defined(_WIN32) +# include "unistd.h" + +# include +#endif + +/* Includes and defines for Hexagon and all HLOS except Windows. */ +#if !defined(_WIN32) +/* Weak reference to remote symbol for compilation. */ +# pragma weak remote_session_control +# pragma weak remote_handle_control +# pragma weak remote_handle64_control +# pragma weak fastrpc_mmap +# pragma weak fastrpc_munmap +# pragma weak rpcmem_alloc2 +#endif + +#if !defined(_WIN32) +# pragma weak remote_system_request +#endif + +#ifdef _WIN32 +# define DSPQUEUE_TIMEOUT DSPQUEUE_TIMEOUT_NONE +#else +# define DSPQUEUE_TIMEOUT 1000000 +#endif + +/** + * htpdrv_init API: driver interface entry point + * + * @return Return AEE error codes as defined in Hexagon SDK. + */ +HTPDRV_API int htpdrv_init(void); + +/** + * get_domain API: get domain struct from domain value. + * + * @param[in] domain value of a domain + * @return Returns domain struct of the domain if it is supported or else + * returns NULL. + * + */ +HTPDRV_API domain * get_domain(int domain_id); + +/** + * get_hex_arch_ver API: query the Hexagon processor architecture version information + * + * @param[in] domain_id value of a domain + * @param[out] Arch version (73, 75, ...) + * @return 0 if query is successful. + * non-zero if error, return value points to the error. + * + */ +HTPDRV_API int get_hex_arch_ver(int domain, int * arch); + +#ifdef __cplusplus +} +#endif + using intvec = std::vector; using uintvec = std::vector; diff --git a/ggml/src/ggml-hexagon/htp-drv.cpp b/ggml/src/ggml-hexagon/htp-drv.cpp.inc similarity index 78% rename from ggml/src/ggml-hexagon/htp-drv.cpp rename to ggml/src/ggml-hexagon/htp-drv.cpp.inc index 4c376b5fc918..16f690a469f1 100644 --- a/ggml/src/ggml-hexagon/htp-drv.cpp +++ b/ggml/src/ggml-hexagon/htp-drv.cpp.inc @@ -19,9 +19,205 @@ # include # include #endif -#include "ggml-impl.h" -#include "htp-drv.h" -#include "libdl.h" +#include "ggml-impl-defs.inc" +#ifdef __cplusplus +extern "C" { +#endif + +#ifdef _WIN32 +# pragma clang diagnostic ignored "-Wignored-attributes" +#endif + +#include +#include +#include +#include + +#if defined(_WIN32) && !defined(__MINGW32__) +# ifdef GGML_BACKEND_BUILD +# define HTPDRV_API __declspec(dllexport) extern +# else +# define HTPDRV_API __declspec(dllimport) extern +# endif +#else +# define HTPDRV_API __attribute__ ((visibility ("default"))) extern +#endif + +/* Offset to differentiate HLOS and Hexagon error codes. + Stores the value of AEE_EOFFSET for Hexagon. */ +#ifndef DSP_OFFSET +# define DSP_OFFSET 0x80000400 +#endif + +/* Errno for connection reset by peer. */ +#ifndef ECONNRESET +# ifdef __hexagon__ +# define ECONNRESET 104 +# endif +#endif + +/* Abstraction of different OS specific sleep APIs. + SLEEP accepts input in seconds. */ +#ifndef SLEEP +# ifdef __hexagon__ +# define SLEEP(x) \ + { /* Do nothing for simulator. */ \ + } +# else +# ifdef _WIN32 +# define SLEEP(x) Sleep(1000 * x) /* Sleep accepts input in milliseconds. */ +# else +# define SLEEP(x) sleep(x) /* sleep accepts input in seconds. */ +# endif +# endif +#endif + +/* Include windows specific header files. */ +#ifdef _WIN32 +# include +# include +# define _CRT_SECURE_NO_WARNINGS 1 +# define _WINSOCK_DEPRECATED_NO_WARNINGS 1 +#endif + +/* Includes and defines for all HLOS except windows */ +#if !defined(__hexagon__) && !defined(_WIN32) +# include "unistd.h" + +# include +#endif + +/* Includes and defines for Hexagon and all HLOS except Windows. */ +#if !defined(_WIN32) +/* Weak reference to remote symbol for compilation. */ +# pragma weak remote_session_control +# pragma weak remote_handle_control +# pragma weak remote_handle64_control +# pragma weak fastrpc_mmap +# pragma weak fastrpc_munmap +# pragma weak rpcmem_alloc2 +#endif + +#if !defined(_WIN32) +# pragma weak remote_system_request +#endif + +#ifdef _WIN32 +# define DSPQUEUE_TIMEOUT DSPQUEUE_TIMEOUT_NONE +#else +# define DSPQUEUE_TIMEOUT 1000000 +#endif + +/** + * htpdrv_init API: driver interface entry point + * + * @return Return AEE error codes as defined in Hexagon SDK. + */ +HTPDRV_API int htpdrv_init(void); + +/** + * get_domain API: get domain struct from domain value. + * + * @param[in] domain value of a domain + * @return Returns domain struct of the domain if it is supported or else + * returns NULL. + * + */ +HTPDRV_API domain * get_domain(int domain_id); + +/** + * get_hex_arch_ver API: query the Hexagon processor architecture version information + * + * @param[in] domain_id value of a domain + * @param[out] Arch version (73, 75, ...) + * @return 0 if query is successful. + * non-zero if error, return value points to the error. + * + */ +HTPDRV_API int get_hex_arch_ver(int domain, int * arch); + +#ifdef __cplusplus +} +#endif + +#ifdef _WIN32 +# define WIN32_LEAN_AND_MEAN +# ifndef NOMINMAX +# define NOMINMAX +# endif +# include +# include +#else +# include +# include +#endif +#include + +namespace fs = std::filesystem; + +#ifdef _WIN32 + +using dl_handle = std::remove_pointer_t; + +struct dl_handle_deleter { + void operator()(HMODULE handle) { + FreeLibrary(handle); + } +}; + +static inline dl_handle * dl_load_library(const fs::path & path) { + // suppress error dialogs for missing DLLs + DWORD old_mode = SetErrorMode(SEM_FAILCRITICALERRORS); + SetErrorMode(old_mode | SEM_FAILCRITICALERRORS); + + HMODULE handle = LoadLibraryW(path.wstring().c_str()); + + SetErrorMode(old_mode); + + return handle; +} + +static inline void * dl_get_sym(dl_handle * handle, const char * name) { + DWORD old_mode = SetErrorMode(SEM_FAILCRITICALERRORS); + SetErrorMode(old_mode | SEM_FAILCRITICALERRORS); + + void * p = (void *) GetProcAddress(handle, name); + + SetErrorMode(old_mode); + + return p; +} + +static inline const char * dl_error() { + return ""; +} + +#else + +using dl_handle = void; + +struct dl_handle_deleter { + void operator()(void * handle) { + dlclose(handle); + } +}; + +static inline dl_handle * dl_load_library(const fs::path & path) { + dl_handle * handle = dlopen(path.string().c_str(), RTLD_NOW | RTLD_LOCAL); + return handle; +} + +static inline void * dl_get_sym(dl_handle * handle, const char * name) { + return dlsym(handle, name); +} + +static inline const char * dl_error() { + const char *rslt = dlerror(); + return rslt != nullptr ? rslt : ""; +} + +#endif + #include diff --git a/ggml/src/ggml-hexagon/htp-drv.h b/ggml/src/ggml-hexagon/htp-drv.h deleted file mode 100644 index 6eba7ba17d8d..000000000000 --- a/ggml/src/ggml-hexagon/htp-drv.h +++ /dev/null @@ -1,121 +0,0 @@ -#pragma once - -#ifdef __cplusplus -extern "C" { -#endif - -#ifdef _WIN32 -# pragma clang diagnostic ignored "-Wignored-attributes" -#endif - -#include -#include -#include -#include - -#if defined(_WIN32) && !defined(__MINGW32__) -# ifdef GGML_BACKEND_BUILD -# define HTPDRV_API __declspec(dllexport) extern -# else -# define HTPDRV_API __declspec(dllimport) extern -# endif -#else -# define HTPDRV_API __attribute__ ((visibility ("default"))) extern -#endif - -/* Offset to differentiate HLOS and Hexagon error codes. - Stores the value of AEE_EOFFSET for Hexagon. */ -#ifndef DSP_OFFSET -# define DSP_OFFSET 0x80000400 -#endif - -/* Errno for connection reset by peer. */ -#ifndef ECONNRESET -# ifdef __hexagon__ -# define ECONNRESET 104 -# endif -#endif - -/* Abstraction of different OS specific sleep APIs. - SLEEP accepts input in seconds. */ -#ifndef SLEEP -# ifdef __hexagon__ -# define SLEEP(x) \ - { /* Do nothing for simulator. */ \ - } -# else -# ifdef _WIN32 -# define SLEEP(x) Sleep(1000 * x) /* Sleep accepts input in milliseconds. */ -# else -# define SLEEP(x) sleep(x) /* sleep accepts input in seconds. */ -# endif -# endif -#endif - -/* Include windows specific header files. */ -#ifdef _WIN32 -# include -# include -# define _CRT_SECURE_NO_WARNINGS 1 -# define _WINSOCK_DEPRECATED_NO_WARNINGS 1 -#endif - -/* Includes and defines for all HLOS except windows */ -#if !defined(__hexagon__) && !defined(_WIN32) -# include "unistd.h" - -# include -#endif - -/* Includes and defines for Hexagon and all HLOS except Windows. */ -#if !defined(_WIN32) -/* Weak reference to remote symbol for compilation. */ -# pragma weak remote_session_control -# pragma weak remote_handle_control -# pragma weak remote_handle64_control -# pragma weak fastrpc_mmap -# pragma weak fastrpc_munmap -# pragma weak rpcmem_alloc2 -#endif - -#if !defined(_WIN32) -# pragma weak remote_system_request -#endif - -#ifdef _WIN32 -# define DSPQUEUE_TIMEOUT DSPQUEUE_TIMEOUT_NONE -#else -# define DSPQUEUE_TIMEOUT 1000000 -#endif - -/** - * htpdrv_init API: driver interface entry point - * - * @return Return AEE error codes as defined in Hexagon SDK. - */ -HTPDRV_API int htpdrv_init(void); - -/** - * get_domain API: get domain struct from domain value. - * - * @param[in] domain value of a domain - * @return Returns domain struct of the domain if it is supported or else - * returns NULL. - * - */ -HTPDRV_API domain * get_domain(int domain_id); - -/** - * get_hex_arch_ver API: query the Hexagon processor architecture version information - * - * @param[in] domain_id value of a domain - * @param[out] Arch version (73, 75, ...) - * @return 0 if query is successful. - * non-zero if error, return value points to the error. - * - */ -HTPDRV_API int get_hex_arch_ver(int domain, int * arch); - -#ifdef __cplusplus -} -#endif diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index 7c9e4cda5f1a..58125cafaec5 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -13,32 +13,33 @@ include_directories( ${CMAKE_CURRENT_BINARY_DIR}) set(HTP_LIB ggml-htp-${DSP_VERSION}) - -add_library(${HTP_LIB} SHARED - main.c +set(HTP_SOURCES + main.inc htp_iface_skel.c - worker-pool.c - hex-dma.c - matmul-ops.c - binary-ops.c - unary-ops.c - sum-rows-ops.c - softmax-ops.c - act-ops.c - rope-ops.c - flash-attn-ops.c - set-rows-ops.c - get-rows-ops.c - cpy-ops.c - repeat-ops.c - argsort-ops.c - ssm-conv.c - cumsum-ops.c - fill-ops.c - diag-ops.c - solve-tri-ops.c + matmul-ops.inc + binary-ops.inc + unary-ops.inc + sum-rows-ops.inc + softmax-ops.inc + act-ops.inc + rope-ops.inc + flash-attn-ops.inc + set-rows-ops.inc + get-rows-ops.inc + cpy-ops.inc + repeat-ops.inc + argsort-ops.inc + ssm-conv.inc + cumsum-ops.inc + fill-ops.inc + diag-ops.inc + solve-tri-ops.inc ) +set_source_files_properties(${HTP_SOURCES} PROPERTIES LANGUAGE C) + +add_library(${HTP_LIB} SHARED ${HTP_SOURCES}) + target_compile_definitions(${HTP_LIB} PRIVATE $,HTP_DEBUG=1,NDEBUG=1> $,FARF_HIGH=1,> @@ -54,16 +55,17 @@ set(HTP_HMX_VERSIONS v73 v75 v79 v81) list(FIND HTP_HMX_VERSIONS ${DSP_VERSION} _hmx_idx) if (_hmx_idx GREATER_EQUAL 0) - target_sources(${HTP_LIB} PRIVATE - hmx-queue.c - hmx-matmul-ops.c - hmx-flash-attn-ops.c + set(HTP_HMX_SOURCES + hmx-matmul-ops.inc + hmx-flash-attn-ops.inc ) + set_source_files_properties(${HTP_HMX_SOURCES} PROPERTIES LANGUAGE C) + target_sources(${HTP_LIB} PRIVATE ${HTP_HMX_SOURCES}) # -mhmx enables HMX instruction set (needed by files that include hmx-utils.h) set_source_files_properties( - hmx-matmul-ops.c - hmx-flash-attn-ops.c + hmx-matmul-ops.inc + hmx-flash-attn-ops.inc PROPERTIES COMPILE_OPTIONS "-mhmx" ) diff --git a/ggml/src/ggml-hexagon/htp/act-ops.c b/ggml/src/ggml-hexagon/htp/act-ops.inc similarity index 99% rename from ggml/src/ggml-hexagon/htp/act-ops.c rename to ggml/src/ggml-hexagon/htp/act-ops.inc index 6416d2dfbc38..ff079bcf05bd 100644 --- a/ggml/src/ggml-hexagon/htp/act-ops.c +++ b/ggml/src/ggml-hexagon/htp/act-ops.inc @@ -8,14 +8,14 @@ #include #include -#include "hex-dma.h" -#include "hvx-utils.h" +#include "hex-dma-defs.inc" +#include "hvx-utils-defs.inc" #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" #define htp_act_preamble \ const struct htp_tensor * src0 = actx->octx->src[0]; \ diff --git a/ggml/src/ggml-hexagon/htp/argsort-ops.c b/ggml/src/ggml-hexagon/htp/argsort-ops.inc similarity index 97% rename from ggml/src/ggml-hexagon/htp/argsort-ops.c rename to ggml/src/ggml-hexagon/htp/argsort-ops.inc index bdd0623615d5..029924297b75 100644 --- a/ggml/src/ggml-hexagon/htp/argsort-ops.c +++ b/ggml/src/ggml-hexagon/htp/argsort-ops.inc @@ -5,15 +5,15 @@ #include #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "ggml.h" +#include "ggml-common-defs.inc" +#include "ggml.h.inc" -#include "hvx-utils.h" -#include "hex-dma.h" +#include "hvx-utils-defs.inc" +#include "hex-dma-defs.inc" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" #ifndef MIN #define MIN(a, b) ((a) < (b) ? (a) : (b)) diff --git a/ggml/src/ggml-hexagon/htp/binary-ops.c b/ggml/src/ggml-hexagon/htp/binary-ops.inc similarity index 99% rename from ggml/src/ggml-hexagon/htp/binary-ops.c rename to ggml/src/ggml-hexagon/htp/binary-ops.inc index 52013ad0fec5..ee97de67fcd6 100644 --- a/ggml/src/ggml-hexagon/htp/binary-ops.c +++ b/ggml/src/ggml-hexagon/htp/binary-ops.inc @@ -8,14 +8,14 @@ #include #include -#include "hex-dma.h" -#include "hvx-utils.h" +#include "hex-dma-defs.inc" +#include "hvx-utils-defs.inc" #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" #ifndef MIN #define MIN(a, b) ((a) < (b) ? (a) : (b)) diff --git a/ggml/src/ggml-hexagon/htp/cpy-ops.c b/ggml/src/ggml-hexagon/htp/cpy-ops.inc similarity index 98% rename from ggml/src/ggml-hexagon/htp/cpy-ops.c rename to ggml/src/ggml-hexagon/htp/cpy-ops.inc index e5b9d350fd77..8179337c653a 100644 --- a/ggml/src/ggml-hexagon/htp/cpy-ops.c +++ b/ggml/src/ggml-hexagon/htp/cpy-ops.inc @@ -9,11 +9,11 @@ #include #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" -#include "hvx-utils.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" +#include "hvx-utils-defs.inc" struct htp_copy_context { struct htp_ops_context * octx; diff --git a/ggml/src/ggml-hexagon/htp/cumsum-ops.c b/ggml/src/ggml-hexagon/htp/cumsum-ops.inc similarity index 98% rename from ggml/src/ggml-hexagon/htp/cumsum-ops.c rename to ggml/src/ggml-hexagon/htp/cumsum-ops.inc index 2ced19712362..6ddb8e20d08e 100644 --- a/ggml/src/ggml-hexagon/htp/cumsum-ops.c +++ b/ggml/src/ggml-hexagon/htp/cumsum-ops.inc @@ -6,12 +6,12 @@ #include #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "hvx-types.h" -#include "hvx-utils.h" -#include "hex-dma.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "hvx-types-defs.inc" +#include "hvx-utils-defs.inc" +#include "hex-dma-defs.inc" #define htp_cumsum_tensors_preamble \ const struct htp_tensor * restrict src0 = octx->src[0]; \ diff --git a/ggml/src/ggml-hexagon/htp/diag-ops.c b/ggml/src/ggml-hexagon/htp/diag-ops.c deleted file mode 100644 index 9b3194d90846..000000000000 --- a/ggml/src/ggml-hexagon/htp/diag-ops.c +++ /dev/null @@ -1,216 +0,0 @@ -#pragma clang diagnostic ignored "-Wunused-but-set-variable" - -#include -#include - -#define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "hvx-types.h" -#include "hex-utils.h" -#include "hvx-copy.h" -#include "hex-dma.h" - -#define htp_diag_tensors_preamble \ - const struct htp_tensor * restrict src0 = octx->src[0]; \ - const struct htp_tensor * restrict dst = octx->dst; \ - \ - const uint32_t ne02 = src0->ne[2]; \ - \ - const uint32_t ne0 = dst->ne[0]; \ - const uint32_t ne1 = dst->ne[1]; \ - \ - const uint32_t nb02 = src0->nb[2]; \ - const uint32_t nb03 = src0->nb[3]; \ - \ - const uint32_t nb1 = dst->nb[1]; \ - const uint32_t nb2 = dst->nb[2]; \ - const uint32_t nb3 = dst->nb[3]; - -struct htp_diag_context { - struct htp_ops_context * octx; - size_t src_batch_size; - size_t dst_row_size; - size_t src_batch_size_aligned; - size_t dst_row_size_aligned; - uint32_t batches_per_thread; - uint32_t total_batches; -}; - -#define htp_diag_preamble \ - struct htp_diag_context * dctx = (struct htp_diag_context *) data; \ - struct htp_ops_context * octx = dctx->octx; \ - htp_diag_tensors_preamble; - -static inline void hvx_diag_row_f32(const float * restrict src, float * restrict dst, - uint32_t row_idx, uint32_t n) { - hvx_splat_f32_a((uint8_t *) dst, 0.0f, n); - dst[row_idx] = src[row_idx]; -} - -// --------------------------------------------------------------------------- -// Per thread worker: DMA src fetch, compute in VTCM, DMA dst writeback -// --------------------------------------------------------------------------- - -static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data) { - htp_diag_preamble; - dma_queue * dma_queue = octx->ctx->dma[ith]; - - uint64_t t1, t2; - t1 = HAP_perf_get_qtimer_count(); - - const uint32_t ib0 = dctx->batches_per_thread * ith; - const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->total_batches); - - if (ib0 >= ib1) { - return; - } - - const size_t src_batch_size = dctx->src_batch_size; - const size_t dst_row_size = dctx->dst_row_size; - const size_t src_batch_size_aligned = dctx->src_batch_size_aligned; - const size_t dst_row_size_aligned = dctx->dst_row_size_aligned; - - const uint8_t * src_data = (const uint8_t *) src0->data; - uint8_t * dst_data = (uint8_t *) dst->data; - - // 1 src buffer + 1 dst row buffer per thread in VTCM - uint8_t * src_spad = octx->src0_spad.data + (ith * src_batch_size_aligned); - uint8_t * dst_spad = octx->dst_spad.data + (ith * dst_row_size_aligned); - - for (uint32_t ib = ib0; ib < ib1; ib++) { - const uint32_t i3 = ib / ne02; - const uint32_t i2 = ib % ne02; - - const uint8_t * src_batch = src_data + i3 * nb03 + i2 * nb02; - - // Fetch source vector into VTCM - dma_queue_push_ddr_to_vtcm(dma_queue, - dma_make_ptr(src_spad, src_batch), - src_batch_size_aligned, src_batch_size, 1); - dma_queue_flush(dma_queue); - - const float * src_spad_f32 = (const float *) src_spad; - float * dst_spad_f32 = (float *) dst_spad; - - for (uint32_t i1 = 0; i1 < ne1; i1++) { - // Compute row in VTCM - hvx_diag_row_f32(src_spad_f32, dst_spad_f32, i1, ne0); - - // Write completed row back to DDR - uint8_t * dst_row = dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1; - dma_queue_push_vtcm_to_ddr(dma_queue, - dma_make_ptr(dst_row, dst_spad), - dst_row_size, dst_row_size_aligned, 1); - dma_queue_flush(dma_queue); - } - } - - t2 = HAP_perf_get_qtimer_count(); - - FARF(HIGH, "diag-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n", - ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ib0, ib1, - dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], - (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1)); -} - -// --------------------------------------------------------------------------- -// Per thread worker: Direct HVX (no DMA) -// --------------------------------------------------------------------------- - -static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) { - htp_diag_preamble; - - uint64_t t1, t2; - t1 = HAP_perf_get_qtimer_count(); - - const uint8_t * src_data = (const uint8_t *) src0->data; - uint8_t * dst_data = (uint8_t *) dst->data; - - const uint32_t ib0 = dctx->batches_per_thread * ith; - const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->total_batches); - - for (uint32_t ib = ib0; ib < ib1; ib++) { - const uint32_t i3 = ib / ne02; - const uint32_t i2 = ib % ne02; - - const float * restrict src_batch = (const float *)(src_data + i3 * nb03 + i2 * nb02); - - for (uint32_t i1 = 0; i1 < ne1; i1++) { - float * restrict dst_row = (float *)(dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1); - hvx_diag_row_f32(src_batch, dst_row, i1, ne0); - } - } - - t2 = HAP_perf_get_qtimer_count(); - - FARF(HIGH, "diag-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n", - ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ib0, ib1, - dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], - (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1)); -} - -int op_diag_f32(struct htp_ops_context * octx) { - const struct htp_tensor * src0 = octx->src[0]; - const struct htp_tensor * dst = octx->dst; - - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; - } - - const uint32_t total_batches = src0->ne[2] * src0->ne[3]; - const uint32_t n_threads = MIN(octx->n_threads, total_batches); - - const size_t src_batch_size = src0->ne[0] * sizeof(float); - const size_t dst_row_size = dst->ne[0] * sizeof(float); - const size_t src_batch_size_aligned = hex_round_up(src_batch_size, VLEN); - const size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN); - - // 1 src buffer + 1 dst row buffer per thread - const size_t spad_per_thread = src_batch_size_aligned + dst_row_size_aligned; - - octx->src0_spad.size_per_thread = src_batch_size_aligned; - octx->dst_spad.size_per_thread = dst_row_size_aligned; - - octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread; - octx->dst_spad.size = n_threads * octx->dst_spad.size_per_thread; - - octx->src0_spad.data = octx->ctx->vtcm_base; octx->src0_spad.src = NULL; - octx->dst_spad.data = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src = NULL; - - struct htp_diag_context dctx = { - .octx = octx, - .src_batch_size = src_batch_size, - .dst_row_size = dst_row_size, - .src_batch_size_aligned = src_batch_size_aligned, - .dst_row_size_aligned = dst_row_size_aligned, - .batches_per_thread = (total_batches + n_threads - 1) / n_threads, - .total_batches = total_batches, - }; - - if (octx->ctx->vtcm_size < spad_per_thread * n_threads) { - worker_pool_run_func(octx->ctx->worker_pool, diag_thread_f32, &dctx, n_threads); - } else { - worker_pool_run_func(octx->ctx->worker_pool, diag_thread_f32_dma, &dctx, n_threads); - } - - return HTP_STATUS_OK; -} - -int op_diag(struct htp_ops_context * octx) { - const struct htp_tensor * dst = octx->dst; - - int err = HTP_STATUS_OK; - - switch (dst->type) { - case HTP_TYPE_F32: - err = op_diag_f32(octx); - break; - default: - err = HTP_STATUS_NO_SUPPORT; - break; - } - - return err; -} diff --git a/ggml/src/ggml-hexagon/htp/hvx-copy.h b/ggml/src/ggml-hexagon/htp/diag-ops.inc similarity index 63% rename from ggml/src/ggml-hexagon/htp/hvx-copy.h rename to ggml/src/ggml-hexagon/htp/diag-ops.inc index a3e33c3b3af0..e454091c0ade 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-copy.h +++ b/ggml/src/ggml-hexagon/htp/diag-ops.inc @@ -1,11 +1,19 @@ -#ifndef HVX_COPY_H -#define HVX_COPY_H +#pragma clang diagnostic ignored "-Wunused-but-set-variable" +#include +#include + +#define GGML_COMMON_DECL_C +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "hvx-types-defs.inc" +#include "hex-utils-defs.inc" #include #include #include -#include "hvx-base.h" +#include "hvx-base-defs.inc" #define hvx_splat_pragma(x) _Pragma(#x) #define hvx_splat_loop_body(dst_type, vec_store, unroll_cnt) \ @@ -259,4 +267,207 @@ static inline void hvx_copy_f32_f16_uu(uint8_t * restrict dst, const uint8_t * r hvx_copy_f32_f16_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); } -#endif // HVX_COPY_H +#include "hex-dma-defs.inc" + +#define htp_diag_tensors_preamble \ + const struct htp_tensor * restrict src0 = octx->src[0]; \ + const struct htp_tensor * restrict dst = octx->dst; \ + \ + const uint32_t ne02 = src0->ne[2]; \ + \ + const uint32_t ne0 = dst->ne[0]; \ + const uint32_t ne1 = dst->ne[1]; \ + \ + const uint32_t nb02 = src0->nb[2]; \ + const uint32_t nb03 = src0->nb[3]; \ + \ + const uint32_t nb1 = dst->nb[1]; \ + const uint32_t nb2 = dst->nb[2]; \ + const uint32_t nb3 = dst->nb[3]; + +struct htp_diag_context { + struct htp_ops_context * octx; + size_t src_batch_size; + size_t dst_row_size; + size_t src_batch_size_aligned; + size_t dst_row_size_aligned; + uint32_t batches_per_thread; + uint32_t total_batches; +}; + +#define htp_diag_preamble \ + struct htp_diag_context * dctx = (struct htp_diag_context *) data; \ + struct htp_ops_context * octx = dctx->octx; \ + htp_diag_tensors_preamble; + +static inline void hvx_diag_row_f32(const float * restrict src, float * restrict dst, + uint32_t row_idx, uint32_t n) { + hvx_splat_f32_a((uint8_t *) dst, 0.0f, n); + dst[row_idx] = src[row_idx]; +} + +// --------------------------------------------------------------------------- +// Per thread worker: DMA src fetch, compute in VTCM, DMA dst writeback +// --------------------------------------------------------------------------- + +static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data) { + htp_diag_preamble; + dma_queue * dma_queue = octx->ctx->dma[ith]; + + uint64_t t1, t2; + t1 = HAP_perf_get_qtimer_count(); + + const uint32_t ib0 = dctx->batches_per_thread * ith; + const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->total_batches); + + if (ib0 >= ib1) { + return; + } + + const size_t src_batch_size = dctx->src_batch_size; + const size_t dst_row_size = dctx->dst_row_size; + const size_t src_batch_size_aligned = dctx->src_batch_size_aligned; + const size_t dst_row_size_aligned = dctx->dst_row_size_aligned; + + const uint8_t * src_data = (const uint8_t *) src0->data; + uint8_t * dst_data = (uint8_t *) dst->data; + + // 1 src buffer + 1 dst row buffer per thread in VTCM + uint8_t * src_spad = octx->src0_spad.data + (ith * src_batch_size_aligned); + uint8_t * dst_spad = octx->dst_spad.data + (ith * dst_row_size_aligned); + + for (uint32_t ib = ib0; ib < ib1; ib++) { + const uint32_t i3 = ib / ne02; + const uint32_t i2 = ib % ne02; + + const uint8_t * src_batch = src_data + i3 * nb03 + i2 * nb02; + + // Fetch source vector into VTCM + dma_queue_push_ddr_to_vtcm(dma_queue, + dma_make_ptr(src_spad, src_batch), + src_batch_size_aligned, src_batch_size, 1); + dma_queue_flush(dma_queue); + + const float * src_spad_f32 = (const float *) src_spad; + float * dst_spad_f32 = (float *) dst_spad; + + for (uint32_t i1 = 0; i1 < ne1; i1++) { + // Compute row in VTCM + hvx_diag_row_f32(src_spad_f32, dst_spad_f32, i1, ne0); + + // Write completed row back to DDR + uint8_t * dst_row = dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1; + dma_queue_push_vtcm_to_ddr(dma_queue, + dma_make_ptr(dst_row, dst_spad), + dst_row_size, dst_row_size_aligned, 1); + dma_queue_flush(dma_queue); + } + } + + t2 = HAP_perf_get_qtimer_count(); + + FARF(HIGH, "diag-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n", + ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ib0, ib1, + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], + (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1)); +} + +// --------------------------------------------------------------------------- +// Per thread worker: Direct HVX (no DMA) +// --------------------------------------------------------------------------- + +static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) { + htp_diag_preamble; + + uint64_t t1, t2; + t1 = HAP_perf_get_qtimer_count(); + + const uint8_t * src_data = (const uint8_t *) src0->data; + uint8_t * dst_data = (uint8_t *) dst->data; + + const uint32_t ib0 = dctx->batches_per_thread * ith; + const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->total_batches); + + for (uint32_t ib = ib0; ib < ib1; ib++) { + const uint32_t i3 = ib / ne02; + const uint32_t i2 = ib % ne02; + + const float * restrict src_batch = (const float *)(src_data + i3 * nb03 + i2 * nb02); + + for (uint32_t i1 = 0; i1 < ne1; i1++) { + float * restrict dst_row = (float *)(dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1); + hvx_diag_row_f32(src_batch, dst_row, i1, ne0); + } + } + + t2 = HAP_perf_get_qtimer_count(); + + FARF(HIGH, "diag-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n", + ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ib0, ib1, + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], + (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1)); +} + +int op_diag_f32(struct htp_ops_context * octx) { + const struct htp_tensor * src0 = octx->src[0]; + const struct htp_tensor * dst = octx->dst; + + if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { + return HTP_STATUS_OK; + } + + const uint32_t total_batches = src0->ne[2] * src0->ne[3]; + const uint32_t n_threads = MIN(octx->n_threads, total_batches); + + const size_t src_batch_size = src0->ne[0] * sizeof(float); + const size_t dst_row_size = dst->ne[0] * sizeof(float); + const size_t src_batch_size_aligned = hex_round_up(src_batch_size, VLEN); + const size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN); + + // 1 src buffer + 1 dst row buffer per thread + const size_t spad_per_thread = src_batch_size_aligned + dst_row_size_aligned; + + octx->src0_spad.size_per_thread = src_batch_size_aligned; + octx->dst_spad.size_per_thread = dst_row_size_aligned; + + octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread; + octx->dst_spad.size = n_threads * octx->dst_spad.size_per_thread; + + octx->src0_spad.data = octx->ctx->vtcm_base; octx->src0_spad.src = NULL; + octx->dst_spad.data = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src = NULL; + + struct htp_diag_context dctx = { + .octx = octx, + .src_batch_size = src_batch_size, + .dst_row_size = dst_row_size, + .src_batch_size_aligned = src_batch_size_aligned, + .dst_row_size_aligned = dst_row_size_aligned, + .batches_per_thread = (total_batches + n_threads - 1) / n_threads, + .total_batches = total_batches, + }; + + if (octx->ctx->vtcm_size < spad_per_thread * n_threads) { + worker_pool_run_func(octx->ctx->worker_pool, diag_thread_f32, &dctx, n_threads); + } else { + worker_pool_run_func(octx->ctx->worker_pool, diag_thread_f32_dma, &dctx, n_threads); + } + + return HTP_STATUS_OK; +} + +int op_diag(struct htp_ops_context * octx) { + const struct htp_tensor * dst = octx->dst; + + int err = HTP_STATUS_OK; + + switch (dst->type) { + case HTP_TYPE_F32: + err = op_diag_f32(octx); + break; + default: + err = HTP_STATUS_NO_SUPPORT; + break; + } + + return err; +} diff --git a/ggml/src/ggml-hexagon/htp/fill-ops.c b/ggml/src/ggml-hexagon/htp/fill-ops.inc similarity index 97% rename from ggml/src/ggml-hexagon/htp/fill-ops.c rename to ggml/src/ggml-hexagon/htp/fill-ops.inc index 3ccfbe74ee45..31d9997c49dc 100644 --- a/ggml/src/ggml-hexagon/htp/fill-ops.c +++ b/ggml/src/ggml-hexagon/htp/fill-ops.inc @@ -7,13 +7,12 @@ #include -#include "hvx-copy.h" -#include "hvx-utils.h" +#include "hvx-utils-defs.inc" #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" // ggml op_params layout for FILL: // op_params[0] (as float) - the scalar fill value diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.inc similarity index 83% rename from ggml/src/ggml-hexagon/htp/flash-attn-ops.c rename to ggml/src/ggml-hexagon/htp/flash-attn-ops.inc index d95df6ac9d5d..7e245386e472 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.inc @@ -8,16 +8,211 @@ #include #include -#include "hex-dma.h" -#include "hvx-utils.h" -#include "hvx-dump.h" +#include "hex-dma-defs.inc" +#include "hvx-utils-defs.inc" +#include + +#include +#include + +#include "hex-utils-defs.inc" +#include "hvx-types-defs.inc" + +static void hvx_vec_dump_f16_n(char * pref, HVX_Vector v, uint32_t n) { + HVX_VectorAlias u = { .v = v }; + + const uint32_t n0 = n / 16; + const uint32_t n1 = n % 16; + int i = 0; + for (; i < n0; i++) { + hex_dump_f16_line(pref, u.fp16 + (16 * i), 16); + } + if (n1) { + hex_dump_f16_line(pref, u.fp16 + (16 * i), n1); + } +} + +static void hvx_vec_dump_f16(char * pref, HVX_Vector v) { + hvx_vec_dump_f16_n(pref, v, 64); +} + +static void hvx_vec_dump_f32_n(char * pref, HVX_Vector v, uint32_t n) { + HVX_VectorAlias u = { .v = v }; + + const uint32_t n0 = n / 16; + const uint32_t n1 = n % 16; + int i = 0; + for (; i < n0; i++) { + hex_dump_f32_line(pref, u.fp32 + (16 * i), 16); + } + if (n1) { + hex_dump_f32_line(pref, u.fp32 + (16 * i), n1); + } +} + +static void hvx_vec_dump_f32_hmt(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + float d[32]; + } u = { .v = v }; + + FARF(HIGH, "%s: %.6f %.6f %.6f %.6f ... %.6f %.6f %.6f %.6f ... %.6f %.6f %.6f %.6f\n", pref, u.d[0], u.d[1], + u.d[2], u.d[3], u.d[12], u.d[13], u.d[14], u.d[15], u.d[28], u.d[29], u.d[30], u.d[31]); +} + +static void hvx_vec_dump_f32(char * pref, HVX_Vector v) { + hvx_vec_dump_f32_n(pref, v, 32); +} + +static void hvx_vec_dump_int32(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int32_t d[32]; + } u = { .v = v }; + + for (int i = 0; i < 32 / 16; i++) { + hex_dump_int32_line(pref, u.d + (16 * i), 16); + } +} + +static void hvx_vec_dump_int32_hmt(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int32_t d[32]; + } u = { .v = v }; + + FARF(HIGH, "%s: %d %d %d %d ... %d %d %d %d ... %d %d %d %d\n", pref, u.d[0], u.d[1], u.d[2], u.d[3], u.d[12], + u.d[13], u.d[14], u.d[15], u.d[28], u.d[29], u.d[30], u.d[31]); +} + +static void hvx_vec_dump_int8_hmt(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int8_t d[128]; + } u = { .v = v }; + + FARF(HIGH, "%s: %d %d %d %d ... %d %d %d %d ... %d %d %d %d\n", pref, u.d[0], u.d[1], u.d[2], u.d[3], u.d[60], + u.d[61], u.d[62], u.d[63], u.d[124], u.d[125], u.d[126], u.d[127]); +} + +static void hvx_vec_dump_int8(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int8_t d[128]; + } u = { .v = v }; + + for (int i = 0; i < 128 / 16; i++) { + hex_dump_int8_line(pref, u.d + (16 * i), 16); + } +} + +static void hvx_vec_dump_uint8(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + uint8_t d[128]; + } u = { .v = v }; + + for (int i = 0; i < 128 / 16; i++) { + hex_dump_uint8_line(pref, u.d + (16 * i), 16); + } +} + +static bool hvx_vec_eq(HVX_Vector v0, HVX_Vector v1, size_t n) { + typedef union { + HVX_Vector v; + int8_t d[128]; + } U; + + U u0 = { .v = v0 }; + U u1 = { .v = v1 }; + + for (int i = 0; i < n; i++) { + if (u0.d[i] != u1.d[i]) { + return false; + } + } + + return true; +} + #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" -#include "hmx-ops.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" +// HMX operation entry-point declarations. +// Ported from htp-ops-lib/include/dsp/ops.h (renamed, benchmark kernels removed). (https://github.com/haozixu/htp-ops-lib) + +#ifndef HMX_OPS_H +#define HMX_OPS_H + +#include +#include + +#include "htp-ops-defs.inc" + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct { + float *dst; + const float *activation; + const __fp16 *permuted_weight; + int m; + int k; + int n; + int act_stride; + int weight_stride; + int dst_stride; + int ne02; + int ne03; + int ne12; + int ne13; + size_t src0_nb2; + size_t src0_nb3; + size_t src1_nb2; + size_t src1_nb3; + size_t dst_nb2; + size_t dst_nb3; +} hmx_matmul_w16a32_batched_params_t; + +// HMX matrix multiplication — tile-permuted FP16 weights, FP32 activation/output +// act_stride: activation row stride in elements (= k for contiguous, or +// nb[1]/sizeof(float) for permuted tensors like attention Q). +// weight_stride: weight row stride in elements (= k for compact weights, or +// nb[1]/sizeof(__fp16) for permuted KV-cache views used by QK). +int hmx_mat_mul_permuted_w16a32(struct htp_context *ctx, + float *restrict dst, + const float *activation, + const __fp16 *permuted_weight, + int m, int k, int n, + int act_stride, + int weight_stride); + +// Batched F16 wrapper over hmx_mat_mul_permuted_w16a32. +// Batch semantics match ggml_mul_mat(): src0 broadcasts to src1 in dims 2/3. +int hmx_mat_mul_permuted_w16a32_batched(struct htp_context *ctx, + const hmx_matmul_w16a32_batched_params_t *params); + +// HMX matrix multiplication — tile-permuted quantised weights (Q4_0/Q8_0/IQ4_NL) +int hmx_mat_mul_permuted_qk_0_d16a32(struct htp_context *ctx, + float *restrict dst, + const float *activation, + const uint8_t *permuted_weight, + int m, int k, int n, + int weight_type); + +// HMX flash attention +int hmx_flash_attn_ext(struct htp_ops_context * octx); + +#ifdef __cplusplus +} +#endif + +#endif // HMX_OPS_H + // Must be multiple of 32 #define FLASH_ATTN_BLOCK_SIZE (32 * 2) diff --git a/ggml/src/ggml-hexagon/htp/get-rows-ops.c b/ggml/src/ggml-hexagon/htp/get-rows-ops.inc similarity index 96% rename from ggml/src/ggml-hexagon/htp/get-rows-ops.c rename to ggml/src/ggml-hexagon/htp/get-rows-ops.inc index 5a1dc933860e..e38e895c1928 100644 --- a/ggml/src/ggml-hexagon/htp/get-rows-ops.c +++ b/ggml/src/ggml-hexagon/htp/get-rows-ops.inc @@ -9,11 +9,11 @@ #include #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" -#include "hvx-utils.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" +#include "hvx-utils-defs.inc" struct get_rows_context { struct htp_ops_context * octx; diff --git a/ggml/src/ggml-hexagon/htp/hex-dma.h b/ggml/src/ggml-hexagon/htp/hex-dma-defs.inc similarity index 100% rename from ggml/src/ggml-hexagon/htp/hex-dma.h rename to ggml/src/ggml-hexagon/htp/hex-dma-defs.inc diff --git a/ggml/src/ggml-hexagon/htp/hex-dma.c b/ggml/src/ggml-hexagon/htp/hex-dma.inc similarity index 98% rename from ggml/src/ggml-hexagon/htp/hex-dma.c rename to ggml/src/ggml-hexagon/htp/hex-dma.inc index b66e2d2603ce..fac59197de2f 100644 --- a/ggml/src/ggml-hexagon/htp/hex-dma.c +++ b/ggml/src/ggml-hexagon/htp/hex-dma.inc @@ -1,4 +1,4 @@ -#include "hex-dma.h" +#include "hex-dma-defs.inc" #include #include diff --git a/ggml/src/ggml-hexagon/htp/hex-dump.h b/ggml/src/ggml-hexagon/htp/hex-dump.h deleted file mode 100644 index 19d173c22328..000000000000 --- a/ggml/src/ggml-hexagon/htp/hex-dump.h +++ /dev/null @@ -1,86 +0,0 @@ -#ifndef HEX_DUMP_H -#define HEX_DUMP_H - -#include - -static inline void hex_dump_int8_line(char * pref, const int8_t * x, int n) { - char str[1024], *p = str, *p_end = str + sizeof(str); - p += snprintf(p, p_end - p, "%s: ", pref); - for (int i = 0; i < n && p < p_end; i++) { - p += snprintf(p, p_end - p, "%d, ", x[i]); - } - FARF(HIGH, "%s\n", str); -} - -static inline void hex_dump_uint8_line(char * pref, const uint8_t * x, uint32_t n) { - char str[1024], *p = str, *p_end = str + sizeof(str); - p += snprintf(p, p_end - p, "%s: ", pref); - for (int i = 0; i < n && p < p_end; i++) { - p += snprintf(p, p_end - p, "%d, ", x[i]); - } - FARF(HIGH, "%s\n", str); -} - -static inline void hex_dump_uint32_line(char * pref, const uint32_t * x, uint32_t n) { - char str[1024], *p = str, *p_end = str + sizeof(str); - p += snprintf(p, p_end - p, "%s: ", pref); - for (int i = 0; i < n; i++) { - p += snprintf(p, p_end - p, "%u, ", (unsigned int) x[i]); - } - FARF(HIGH, "%s\n", str); -} - -static inline void hex_dump_int32_line(char * pref, const int32_t * x, uint32_t n) { - char str[1024], *p = str, *p_end = str + sizeof(str); - p += snprintf(p, p_end - p, "%s: ", pref); - for (int i = 0; i < n; i++) { - p += snprintf(p, p_end - p, "%d, ", (int) x[i]); - } - FARF(HIGH, "%s\n", str); -} - -static inline void hex_dump_f16_line(char * pref, const __fp16 * x, uint32_t n) { - char str[1024], *p = str, *p_end = str + sizeof(str); - p += snprintf(p, p_end - p, "%s: ", pref); - for (int i = 0; i < n; i++) { - p += snprintf(p, p_end - p, "%.6f, ", (float) x[i]); - } - FARF(HIGH, "%s\n", str); -} - -static inline void hex_dump_f32_line(char * pref, const float * x, uint32_t n) { - char str[1024], *p = str, *p_end = str + sizeof(str); - p += snprintf(p, p_end - p, "%s: ", pref); - for (int i = 0; i < n; i++) { - p += snprintf(p, p_end - p, "%.6f, ", x[i]); - } - FARF(HIGH, "%s\n", str); -} - -static inline void hex_dump_f32(char * pref, const float * x, uint32_t n) { - uint32_t n0 = n / 16; - uint32_t n1 = n % 16; - - uint32_t i = 0; - for (; i < n0; i++) { - hex_dump_f32_line(pref, x + (16 * i), 16); - } - if (n1) { - hex_dump_f32_line(pref, x + (16 * i), n1); - } -} - -static inline void hex_dump_f16(char * pref, const __fp16 * x, uint32_t n) { - uint32_t n0 = n / 16; - uint32_t n1 = n % 16; - - uint32_t i = 0; - for (; i < n0; i++) { - hex_dump_f16_line(pref, x + (16 * i), 16); - } - if (n1) { - hex_dump_f16_line(pref, x + (16 * i), n1); - } -} - -#endif /* HEX_DUMP_H */ diff --git a/ggml/src/ggml-hexagon/htp/hex-fastdiv.h b/ggml/src/ggml-hexagon/htp/hex-fastdiv.h deleted file mode 100644 index b7b5867593fc..000000000000 --- a/ggml/src/ggml-hexagon/htp/hex-fastdiv.h +++ /dev/null @@ -1,37 +0,0 @@ -#ifndef HEX_FASTDIV_H -#define HEX_FASTDIV_H - -// See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1. -// Precompute mp (m' in the paper) and L such that division -// can be computed using a multiply (high 32b of 64b result) -// and a shift: -// -// n/d = (mulhi(n, mp) + n) >> L; -struct fastdiv_values { - uint32_t mp; - uint32_t l; -}; - -static inline struct fastdiv_values init_fastdiv_values(uint32_t d) { - struct fastdiv_values result = { 0, 0 }; - // compute L = ceil(log2(d)); - while (result.l < 32 && ((uint32_t) 1 << result.l) < d) { - ++(result.l); - } - - result.mp = (uint32_t) (((uint64_t) 1 << 32) * (((uint64_t) 1 << result.l) - d) / d + 1); - return result; -} - -static inline uint32_t fastdiv(uint32_t n, const struct fastdiv_values * vals) { - // Compute high 32 bits of n * mp - const uint32_t hi = (uint32_t) (((uint64_t) n * vals->mp) >> 32); // mulhi(n, mp) - // add n, apply bit shift - return (hi + n) >> vals->l; -} - -static inline uint32_t fastmodulo(uint32_t n, uint32_t d, const struct fastdiv_values * vals) { - return n - fastdiv(n, vals) * d; -} - -#endif /* HEX_FASTDIV_H */ diff --git a/ggml/src/ggml-hexagon/htp/hex-utils.h b/ggml/src/ggml-hexagon/htp/hex-utils-defs.inc similarity index 51% rename from ggml/src/ggml-hexagon/htp/hex-utils.h rename to ggml/src/ggml-hexagon/htp/hex-utils-defs.inc index 6239ceff4b41..7817968a8399 100644 --- a/ggml/src/ggml-hexagon/htp/hex-utils.h +++ b/ggml/src/ggml-hexagon/htp/hex-utils-defs.inc @@ -9,8 +9,123 @@ #include "hexagon_types.h" #include "hexagon_protos.h" -#include "hex-fastdiv.h" -#include "hex-dump.h" +// See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1. +// Precompute mp (m' in the paper) and L such that division +// can be computed using a multiply (high 32b of 64b result) +// and a shift: +// +// n/d = (mulhi(n, mp) + n) >> L; +struct fastdiv_values { + uint32_t mp; + uint32_t l; +}; + +static inline struct fastdiv_values init_fastdiv_values(uint32_t d) { + struct fastdiv_values result = { 0, 0 }; + // compute L = ceil(log2(d)); + while (result.l < 32 && ((uint32_t) 1 << result.l) < d) { + ++(result.l); + } + + result.mp = (uint32_t) (((uint64_t) 1 << 32) * (((uint64_t) 1 << result.l) - d) / d + 1); + return result; +} + +static inline uint32_t fastdiv(uint32_t n, const struct fastdiv_values * vals) { + // Compute high 32 bits of n * mp + const uint32_t hi = (uint32_t) (((uint64_t) n * vals->mp) >> 32); // mulhi(n, mp) + // add n, apply bit shift + return (hi + n) >> vals->l; +} + +static inline uint32_t fastmodulo(uint32_t n, uint32_t d, const struct fastdiv_values * vals) { + return n - fastdiv(n, vals) * d; +} + + +#include + +static inline void hex_dump_int8_line(char * pref, const int8_t * x, int n) { + char str[1024], *p = str, *p_end = str + sizeof(str); + p += snprintf(p, p_end - p, "%s: ", pref); + for (int i = 0; i < n && p < p_end; i++) { + p += snprintf(p, p_end - p, "%d, ", x[i]); + } + FARF(HIGH, "%s\n", str); +} + +static inline void hex_dump_uint8_line(char * pref, const uint8_t * x, uint32_t n) { + char str[1024], *p = str, *p_end = str + sizeof(str); + p += snprintf(p, p_end - p, "%s: ", pref); + for (int i = 0; i < n && p < p_end; i++) { + p += snprintf(p, p_end - p, "%d, ", x[i]); + } + FARF(HIGH, "%s\n", str); +} + +static inline void hex_dump_uint32_line(char * pref, const uint32_t * x, uint32_t n) { + char str[1024], *p = str, *p_end = str + sizeof(str); + p += snprintf(p, p_end - p, "%s: ", pref); + for (int i = 0; i < n; i++) { + p += snprintf(p, p_end - p, "%u, ", (unsigned int) x[i]); + } + FARF(HIGH, "%s\n", str); +} + +static inline void hex_dump_int32_line(char * pref, const int32_t * x, uint32_t n) { + char str[1024], *p = str, *p_end = str + sizeof(str); + p += snprintf(p, p_end - p, "%s: ", pref); + for (int i = 0; i < n; i++) { + p += snprintf(p, p_end - p, "%d, ", (int) x[i]); + } + FARF(HIGH, "%s\n", str); +} + +static inline void hex_dump_f16_line(char * pref, const __fp16 * x, uint32_t n) { + char str[1024], *p = str, *p_end = str + sizeof(str); + p += snprintf(p, p_end - p, "%s: ", pref); + for (int i = 0; i < n; i++) { + p += snprintf(p, p_end - p, "%.6f, ", (float) x[i]); + } + FARF(HIGH, "%s\n", str); +} + +static inline void hex_dump_f32_line(char * pref, const float * x, uint32_t n) { + char str[1024], *p = str, *p_end = str + sizeof(str); + p += snprintf(p, p_end - p, "%s: ", pref); + for (int i = 0; i < n; i++) { + p += snprintf(p, p_end - p, "%.6f, ", x[i]); + } + FARF(HIGH, "%s\n", str); +} + +static inline void hex_dump_f32(char * pref, const float * x, uint32_t n) { + uint32_t n0 = n / 16; + uint32_t n1 = n % 16; + + uint32_t i = 0; + for (; i < n0; i++) { + hex_dump_f32_line(pref, x + (16 * i), 16); + } + if (n1) { + hex_dump_f32_line(pref, x + (16 * i), n1); + } +} + +static inline void hex_dump_f16(char * pref, const __fp16 * x, uint32_t n) { + uint32_t n0 = n / 16; + uint32_t n1 = n % 16; + + uint32_t i = 0; + for (; i < n0; i++) { + hex_dump_f16_line(pref, x + (16 * i), 16); + } + if (n1) { + hex_dump_f16_line(pref, x + (16 * i), n1); + } +} + + #ifndef MAX #define MAX(a, b) ((a) > (b) ? (a) : (b)) diff --git a/ggml/src/ggml-hexagon/htp/hmx-flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/hmx-flash-attn-ops.inc similarity index 77% rename from ggml/src/ggml-hexagon/htp/hmx-flash-attn-ops.c rename to ggml/src/ggml-hexagon/htp/hmx-flash-attn-ops.inc index 8a6d7c14edfb..2f5bd76fe1bc 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/hmx-flash-attn-ops.inc @@ -15,18 +15,673 @@ #include #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "hex-dma.h" -#include "hmx-profile.h" -#include "hmx-queue.h" -#include "hmx-utils.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "hvx-dump.h" -#include "hvx-reduce.h" -#include "hvx-utils.h" -#include "vtcm-utils.h" -#include "worker-pool.h" +#include "ggml-common-defs.inc" +#include "hex-dma-defs.inc" +// Conditional fine-grained profiling macros for HMX operations. +// +// Define ENABLE_PROFILE_TIMERS (via compiler flag or before including this +// header) to instrument sub-operation latencies with HAP qtimer. When the +// macro is not defined the TIMER_* helpers expand to nothing so there is zero +// overhead. +// +// Usage: +// TIMER_DEFINE(my_phase); // declare accumulator variable +// TIMER_START(my_phase); // snapshot start time +// ... work ... +// TIMER_STOP(my_phase); // accumulate elapsed ticks +// FARF(ALWAYS, "my_phase: %lld us", TIMER_US(my_phase)); + +#include + +// #define ENABLE_PROFILE_TIMERS + +#if defined(ENABLE_PROFILE_TIMERS) +# define TIMER_DEFINE(name) int64_t name##_ticks = 0 +# define TIMER_START(name) int64_t name##_t0 = HAP_perf_get_qtimer_count() +# define TIMER_STOP(name) name##_ticks += HAP_perf_get_qtimer_count() - name##_t0 +# define TIMER_US(name) HAP_perf_qtimer_count_to_us(name##_ticks) +#else +# define TIMER_DEFINE(name) +# define TIMER_START(name) +# define TIMER_STOP(name) +# define TIMER_US(name) 0LL +#endif + + +#include "hmx-queue-defs.inc" +// HMX tile-level inline helpers (FP16 32x32 tile operations). +// Ported from htp-ops-lib/include/dsp/hmx_utils.h. (https://github.com/haozixu/htp-ops-lib) + +#include "hvx-base-defs.inc" + +#include +#include +#include + +#define HMX_FP16_TILE_N_ROWS 32 +#define HMX_FP16_TILE_N_COLS 32 +#define HMX_FP16_TILE_N_ELMS 1024 +#define HMX_FP16_TILE_SIZE 2048 + +// Initialise aligned 256-byte area with scale vector + zero padding. +static inline void hmx_init_column_scales(void *out_scales, HVX_Vector v_scale) { + volatile HVX_Vector *pv = (HVX_Vector *) out_scales; + pv[0] = v_scale; + pv[1] = Q6_V_vzero(); +} + +// --- Shared scatter offsets and interleave helper --- + +// vscatter offsets for fused dequant+transpose: write K-values directly to [K][N] tile. +// word[i] = i*128 maps K-row-pair i to byte offset i*128. +// Column offset (n*4) is added at runtime. Entries 0..15 cover one tile (region 2047); +// entries 16..31 cover the next adjacent tile (region 4095) — pick region size at the +// call site to scatter into one tile (masked) or two contiguous tiles (unmasked). +static const int32_t hmx_transpose_scatter_offsets[32] __attribute__((aligned(VLEN))) = { + 0 * 128, 1 * 128, 2 * 128, 3 * 128, 4 * 128, 5 * 128, 6 * 128, 7 * 128, 8 * 128, 9 * 128, 10 * 128, + 11 * 128, 12 * 128, 13 * 128, 14 * 128, 15 * 128, 16 * 128, 17 * 128, 18 * 128, 19 * 128, 20 * 128, 21 * 128, + 22 * 128, 23 * 128, 24 * 128, 25 * 128, 26 * 128, 27 * 128, 28 * 128, 29 * 128, 30 * 128, 31 * 128, +}; + +// Scatter row-major FP16 data (in VTCM scratch) into transposed [K][N] tiles. +// vtcm_src: [n_cols][src_stride] row-major fp16 (only first k elements per row are used) +// vtcm_dst: [n_col_tiles][n_k_tiles][HMX_FP16_TILE_N_ELMS] tile-major interleaved fp16 +// Processes rows [start_row, end_row) for multi-thread slicing. +// Full range: start_row=0, end_row=n_cols. +static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst, + const __fp16 * restrict vtcm_src, + int n_cols, + int k, + int src_stride, + int start_row, + int end_row) { + assert(k % HMX_FP16_TILE_N_COLS == 0); + + const int n_k_tiles = k / HMX_FP16_TILE_N_COLS; + const HVX_Vector v_scat_base = hvx_vmem(hmx_transpose_scatter_offsets); + const HVX_Vector v_scat_step = Q6_V_vsplat_R(4); + const HVX_VectorPred q_mask64 = Q6_Q_vsetq_R(64); + // Each hvx_vmemu load brings 64 fp16 = 128 bytes covering 2 adjacent K-tiles. + // When n_k_tiles is even, scatter into 2 K-tiles per call (region 4095, no mask) + // using the upper half of hmx_transpose_scatter_offsets. Tail one K-tile (when + // n_k_tiles is odd) falls back to single-tile masked scatter. + const bool pair_scatter = (n_k_tiles & 1) == 0; + const size_t pair_region = (size_t) (2 * HMX_FP16_TILE_SIZE - 1); + const size_t single_region = (size_t) (HMX_FP16_TILE_SIZE - 1); + __builtin_assume(k > 0); + __builtin_assume(end_row > start_row); + + if (pair_scatter) { + // Step c by 64 fp16 (two K-tiles per scatter), advance dst by 2 tiles per iter. + const int c_step = 2 * HMX_FP16_TILE_N_COLS; + const size_t c_byte_step = (size_t) c_step * sizeof(__fp16); + const size_t dst_step = 2 * (size_t) HMX_FP16_TILE_N_ELMS; + const int n_c_iters = k / c_step; + + for (int r = start_row; r < end_row; r += 2) { + const int ct = r / HMX_FP16_TILE_N_ROWS; + const int local_r = r % HMX_FP16_TILE_N_ROWS; + const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_cols; + const HVX_Vector v_off0 = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4)); + const HVX_Vector v_off1 = Q6_Vw_vadd_VwVw(v_off0, v_scat_step); + + __fp16 * tile_base = vtcm_dst + (size_t) ct * n_k_tiles * HMX_FP16_TILE_N_ELMS; + const uint8_t * p0 = (const uint8_t *) (vtcm_src + r * src_stride); + const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL; + + if (p1) { + for (int i = 0; i < n_c_iters; ++i) { + HVX_Vector v0 = hvx_vmemu(p0); + p0 += c_byte_step; + HVX_Vector v1 = hvx_vmemu(p1); + p1 += c_byte_step; + Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0); + Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, v1); + tile_base += dst_step; + } + } else { + const HVX_Vector vzero = Q6_V_vzero(); + for (int i = 0; i < n_c_iters; ++i) { + HVX_Vector v0 = hvx_vmemu(p0); + p0 += c_byte_step; + Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0); + Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, vzero); + tile_base += dst_step; + } + } + } + } else { + // Fallback: scatter one K-tile per call (region 2047, masked). + const int c_step = HMX_FP16_TILE_N_COLS; + const size_t c_byte_step = (size_t) c_step * sizeof(__fp16); + const size_t dst_step = (size_t) HMX_FP16_TILE_N_ELMS; + const int n_c_iters = k / c_step; + + for (int r = start_row; r < end_row; r += 2) { + const int ct = r / HMX_FP16_TILE_N_ROWS; + const int local_r = r % HMX_FP16_TILE_N_ROWS; + const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_cols; + const HVX_Vector v_off0 = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4)); + const HVX_Vector v_off1 = Q6_Vw_vadd_VwVw(v_off0, v_scat_step); + + __fp16 * tile_base = vtcm_dst + (size_t) ct * n_k_tiles * HMX_FP16_TILE_N_ELMS; + const uint8_t * p0 = (const uint8_t *) (vtcm_src + r * src_stride); + const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL; + + if (p1) { + for (int i = 0; i < n_c_iters; ++i) { + HVX_Vector v0 = hvx_vmemu(p0); + p0 += c_byte_step; + HVX_Vector v1 = hvx_vmemu(p1); + p1 += c_byte_step; + Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0); + Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, v1); + tile_base += dst_step; + } + } else { + const HVX_Vector vzero = Q6_V_vzero(); + for (int i = 0; i < n_c_iters; ++i) { + HVX_Vector v0 = hvx_vmemu(p0); + p0 += c_byte_step; + Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0); + Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, vzero); + tile_base += dst_step; + } + } + } + } +} + +// Interleave row-major FP16 data into column-major tile format. +// Input: [n_rows, head_dim] row-major. Output: tile[dim_tile][row_tile]. +// Processes rows [start_row, end_row) for multi-thread slicing. +// Full range: start_row=0, end_row=n_rows. +static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out, + const __fp16 * restrict src, + int n_rows, + int head_dim, + int src_stride, + int n_row_tiles, + int start_row, + int end_row) { + __builtin_assume(head_dim > 0); + const size_t tile_stride_elms = (size_t) n_row_tiles * HMX_FP16_TILE_N_ELMS; + + for (int r = start_row; r < end_row; r += 2) { + const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_rows; + + const HVX_Vector * pv_in0 = (const HVX_Vector *) (src + r * src_stride); + const HVX_Vector * pv_in1 = next_row_valid ? (const HVX_Vector *) (src + (r + 1) * src_stride) : NULL; + + // Row-pair invariants hoisted out of the c loop. + const int r0 = r / HMX_FP16_TILE_N_ROWS; + const int r1_half = (r % HMX_FP16_TILE_N_ROWS) / 2; + + // tb0 starts at tile (c0=0, r0); tb1 at the adjacent dim-tile (c0=1, r0). + // Each c step (+= 64) advances both by 2 dim-tiles worth of fp16. + __fp16 * tb0 = tiles_out + (size_t) r0 * HMX_FP16_TILE_N_ELMS; + __fp16 * tb1 = tb0 + tile_stride_elms; + const size_t tb_step = 2 * tile_stride_elms; + + if (pv_in1) { + for (int c = 0; c < head_dim; c += 64) { + HVX_Vector v0 = *pv_in0++; + HVX_Vector v1 = *pv_in1++; + HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2); + ((HVX_Vector *) tb0)[r1_half] = Q6_V_lo_W(vp); + ((HVX_Vector *) tb1)[r1_half] = Q6_V_hi_W(vp); + tb0 += tb_step; + tb1 += tb_step; + } + } else { + const HVX_Vector vzero = Q6_V_vzero(); + for (int c = 0; c < head_dim; c += 64) { + HVX_Vector v0 = *pv_in0++; + HVX_VectorPair vp = Q6_W_vshuff_VVR(vzero, v0, -2); + ((HVX_Vector *) tb0)[r1_half] = Q6_V_lo_W(vp); + ((HVX_Vector *) tb1)[r1_half] = Q6_V_hi_W(vp); + tb0 += tb_step; + tb1 += tb_step; + } + } + } +} + + +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include + +#include +#include + +#include "hex-utils-defs.inc" +#include "hvx-types-defs.inc" + +static void hvx_vec_dump_f16_n(char * pref, HVX_Vector v, uint32_t n) { + HVX_VectorAlias u = { .v = v }; + + const uint32_t n0 = n / 16; + const uint32_t n1 = n % 16; + int i = 0; + for (; i < n0; i++) { + hex_dump_f16_line(pref, u.fp16 + (16 * i), 16); + } + if (n1) { + hex_dump_f16_line(pref, u.fp16 + (16 * i), n1); + } +} + +static void hvx_vec_dump_f16(char * pref, HVX_Vector v) { + hvx_vec_dump_f16_n(pref, v, 64); +} + +static void hvx_vec_dump_f32_n(char * pref, HVX_Vector v, uint32_t n) { + HVX_VectorAlias u = { .v = v }; + + const uint32_t n0 = n / 16; + const uint32_t n1 = n % 16; + int i = 0; + for (; i < n0; i++) { + hex_dump_f32_line(pref, u.fp32 + (16 * i), 16); + } + if (n1) { + hex_dump_f32_line(pref, u.fp32 + (16 * i), n1); + } +} + +static void hvx_vec_dump_f32_hmt(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + float d[32]; + } u = { .v = v }; + + FARF(HIGH, "%s: %.6f %.6f %.6f %.6f ... %.6f %.6f %.6f %.6f ... %.6f %.6f %.6f %.6f\n", pref, u.d[0], u.d[1], + u.d[2], u.d[3], u.d[12], u.d[13], u.d[14], u.d[15], u.d[28], u.d[29], u.d[30], u.d[31]); +} + +static void hvx_vec_dump_f32(char * pref, HVX_Vector v) { + hvx_vec_dump_f32_n(pref, v, 32); +} + +static void hvx_vec_dump_int32(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int32_t d[32]; + } u = { .v = v }; + + for (int i = 0; i < 32 / 16; i++) { + hex_dump_int32_line(pref, u.d + (16 * i), 16); + } +} + +static void hvx_vec_dump_int32_hmt(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int32_t d[32]; + } u = { .v = v }; + + FARF(HIGH, "%s: %d %d %d %d ... %d %d %d %d ... %d %d %d %d\n", pref, u.d[0], u.d[1], u.d[2], u.d[3], u.d[12], + u.d[13], u.d[14], u.d[15], u.d[28], u.d[29], u.d[30], u.d[31]); +} + +static void hvx_vec_dump_int8_hmt(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int8_t d[128]; + } u = { .v = v }; + + FARF(HIGH, "%s: %d %d %d %d ... %d %d %d %d ... %d %d %d %d\n", pref, u.d[0], u.d[1], u.d[2], u.d[3], u.d[60], + u.d[61], u.d[62], u.d[63], u.d[124], u.d[125], u.d[126], u.d[127]); +} + +static void hvx_vec_dump_int8(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int8_t d[128]; + } u = { .v = v }; + + for (int i = 0; i < 128 / 16; i++) { + hex_dump_int8_line(pref, u.d + (16 * i), 16); + } +} + +static void hvx_vec_dump_uint8(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + uint8_t d[128]; + } u = { .v = v }; + + for (int i = 0; i < 128 / 16; i++) { + hex_dump_uint8_line(pref, u.d + (16 * i), 16); + } +} + +static bool hvx_vec_eq(HVX_Vector v0, HVX_Vector v1, size_t n) { + typedef union { + HVX_Vector v; + int8_t d[128]; + } U; + + U u0 = { .v = v0 }; + U u1 = { .v = v1 }; + + for (int i = 0; i < n; i++) { + if (u0.d[i] != u1.d[i]) { + return false; + } + } + + return true; +} + +#include +#include +#include +#include + +#include "hex-utils-defs.inc" +#include "hvx-base-defs.inc" +#include "hvx-types-defs.inc" + +static inline HVX_Vector hvx_vec_reduce_sum_n_i32(HVX_Vector in, unsigned int n) { + unsigned int total = n * 4; // total vec nbytes + unsigned int width = 4; // int32 + + HVX_Vector sum = in, sum_t; + while (width < total) { + sum_t = Q6_V_vror_VR(sum, width); // rotate right + sum = Q6_Vw_vadd_VwVw(sum_t, sum); // elementwise sum + width = width << 1; + } + return sum; +} + +static inline HVX_Vector hvx_vec_reduce_sum_i32(HVX_Vector in) { + return hvx_vec_reduce_sum_n_i32(in, 32); +} + +static inline HVX_Vector hvx_vec_reduce_sum_n_qf32(HVX_Vector in, unsigned int n) { + unsigned int total = n * 4; // total vec nbytes + unsigned int width = 4; // fp32 nbytes + + HVX_Vector sum = in, sum_t; + while (width < total) { + sum_t = Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum), width); // rotate right + sum = Q6_Vqf32_vadd_Vqf32Vsf(sum, sum_t); // elementwise sum + width = width << 1; + } + return sum; +} + +static inline HVX_Vector hvx_vec_reduce_sum_qf32(HVX_Vector in) { + return hvx_vec_reduce_sum_n_qf32(in, 32); +} + +#if __HVX_ARCH__ > 75 + +static inline HVX_Vector hvx_vec_reduce_sum_f32x4(HVX_Vector_x4 in) { + HVX_VectorPair sum_p01 = Q6_W_vshuff_VVR(in.v[1], in.v[0], 4); + HVX_VectorPair sum_p23 = Q6_W_vshuff_VVR(in.v[3], in.v[2], 4); + HVX_Vector sum_sf01 = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p01), Q6_V_hi_W(sum_p01)); + HVX_Vector sum_sf23 = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p23), Q6_V_hi_W(sum_p23)); + + HVX_VectorPair sum_p0123 = Q6_W_vshuff_VVR(sum_sf23, sum_sf01, 8); + HVX_Vector sum_sf = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p0123), Q6_V_hi_W(sum_p0123)); + + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 2)); + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 4)); + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 8)); + return sum_sf; +} + +static inline HVX_Vector hvx_vec_reduce_sum_f32x2(HVX_Vector in0, HVX_Vector in1) { + HVX_VectorPair sump = Q6_W_vshuff_VVR(in1, in0, 4); + HVX_Vector sum_sf = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sump), Q6_V_hi_W(sump)); + + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 2)); + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 4)); + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 8)); + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 16)); + return sum_sf; +} + +static inline HVX_Vector hvx_vec_reduce_sum_n_f32(HVX_Vector in, unsigned int n) { + unsigned int total = n * 4; // total vec nbytes + unsigned int width = 4; // fp32 nbytes + + HVX_Vector sum = in, sum_t; + while (width < total) { + sum_t = Q6_V_vror_VR(sum, width); // rotate right + sum = Q6_Vsf_vadd_VsfVsf(sum, sum_t); // elementwise sum + width = width << 1; + } + return sum; +} + +#else + +static inline HVX_Vector hvx_vec_reduce_sum_f32x4(HVX_Vector_x4 in) { + HVX_VectorPair sum_p01 = Q6_W_vshuff_VVR(in.v[1], in.v[0], 4); + HVX_VectorPair sum_p23 = Q6_W_vshuff_VVR(in.v[3], in.v[2], 4); + HVX_Vector sum_qf01 = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p01), Q6_V_hi_W(sum_p01)); + HVX_Vector sum_qf23 = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p23), Q6_V_hi_W(sum_p23)); + + HVX_VectorPair sum_p0123 = Q6_W_vshuff_VVR(Q6_Vsf_equals_Vqf32(sum_qf23), Q6_Vsf_equals_Vqf32(sum_qf01), 8); + HVX_Vector sum_qf = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p0123), Q6_V_hi_W(sum_p0123)); + + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 2)); + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 4)); + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 8)); + return Q6_Vsf_equals_Vqf32(sum_qf); +} + +static inline HVX_Vector hvx_vec_reduce_sum_f32x2(HVX_Vector in0, HVX_Vector in1) { + HVX_VectorPair sump = Q6_W_vshuff_VVR(in1, in0, 4); + HVX_Vector sum_qf = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sump), Q6_V_hi_W(sump)); + + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 2)); + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 4)); + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 8)); + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 16)); + return Q6_Vsf_equals_Vqf32(sum_qf); +} + +static inline HVX_Vector hvx_vec_reduce_sum_n_f32(HVX_Vector in, unsigned int n) { + unsigned int total = n * 4; // total vec nbytes + unsigned int width = 4; // fp32 nbytes + + HVX_Vector sum = in, sum_t; + while (width < total) { + sum_t = Q6_V_vror_VR(sum, width); // rotate right + sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(sum, sum_t)); // elementwise sum + width = width << 1; + } + return sum; +} + +#endif + +static inline HVX_Vector hvx_vec_reduce_sum_f32(HVX_Vector in) { + return hvx_vec_reduce_sum_n_f32(in, 32); +} + +static inline HVX_Vector hvx_vec_reduce_max_f16(HVX_Vector in) { + unsigned total = 128; // total vec nbytes + unsigned width = 2; // fp16 nbytes + + HVX_Vector _max = in, _max_t; + while (width < total) { + _max_t = Q6_V_vror_VR(_max, width); // rotate right + _max = Q6_Vhf_vmax_VhfVhf(_max_t, _max); // elementwise max + width = width << 1; + } + + return _max; +} + +static inline HVX_Vector hvx_vec_reduce_max2_f16(HVX_Vector in, HVX_Vector _max) { + unsigned total = 128; // total vec nbytes + unsigned width = 2; // fp32 nbytes + + HVX_Vector _max_t; + + _max = Q6_Vhf_vmax_VhfVhf(in, _max); + while (width < total) { + _max_t = Q6_V_vror_VR(_max, width); // rotate right + _max = Q6_Vhf_vmax_VhfVhf(_max_t, _max); // elementwise max + width = width << 1; + } + + return _max; +} + +static inline HVX_Vector hvx_vec_reduce_max_f32(HVX_Vector in) { + unsigned total = 128; // total vec nbytes + unsigned width = 4; // fp32 nbytes + + HVX_Vector _max = in, _max_t; + while (width < total) { + _max_t = Q6_V_vror_VR(_max, width); // rotate right + _max = Q6_Vsf_vmax_VsfVsf(_max_t, _max); // elementwise max + width = width << 1; + } + + return _max; +} + +static inline HVX_Vector hvx_vec_reduce_max2_f32(HVX_Vector in, HVX_Vector _max) { + unsigned total = 128; // total vec nbytes + unsigned width = 4; // fp32 nbytes + + HVX_Vector _max_t; + + _max = Q6_Vsf_vmax_VsfVsf(in, _max); + while (width < total) { + _max_t = Q6_V_vror_VR(_max, width); // rotate right + _max = Q6_Vsf_vmax_VsfVsf(_max_t, _max); // elementwise max + width = width << 1; + } + + return _max; +} + +#define hvx_reduce_loop_body(src_type, init_vec, pad_vec, vec_op, reduce_op, scalar_reduce) \ + do { \ + src_type * restrict vsrc = (src_type *) src; \ + HVX_Vector acc = init_vec; \ + \ + const uint32_t elem_size = sizeof(float); \ + const uint32_t epv = 128 / elem_size; \ + const uint32_t nvec = num_elems / epv; \ + const uint32_t nloe = num_elems % epv; \ + \ + uint32_t i = 0; \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + acc = vec_op(acc, vsrc[i]); \ + } \ + if (nloe) { \ + const float * srcf = (const float *) src + i * epv; \ + HVX_Vector in = *(HVX_UVector *) srcf; \ + HVX_Vector temp = Q6_V_valign_VVR(in, pad_vec, nloe * elem_size); \ + acc = vec_op(acc, temp); \ + } \ + HVX_Vector v = reduce_op(acc); \ + return scalar_reduce(v); \ + } while(0) + +#define HVX_REDUCE_MAX_OP(acc, val) Q6_Vsf_vmax_VsfVsf(acc, val) +#define HVX_REDUCE_SUM_OP(acc, val) Q6_Vqf32_vadd_VsfVsf(Q6_Vsf_equals_Vqf32(acc), val) +#define HVX_SUM_SQ_OP(acc, val) Q6_Vqf32_vadd_Vqf32Vqf32(acc, Q6_Vqf32_vmpy_VsfVsf(val, val)) +#define HVX_REDUCE_MAX_SCALAR(v) hvx_vec_get_f32(v) +#define HVX_REDUCE_SUM_SCALAR(v) hvx_vec_get_f32(Q6_Vsf_equals_Vqf32(v)) + +// Max variants + +static inline float hvx_reduce_max_f32_a(const uint8_t * restrict src, const int num_elems) { + HVX_Vector init_vec = hvx_vec_splat_f32(((const float *) src)[0]); + assert((unsigned long) src % 128 == 0); + hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_REDUCE_MAX_OP, hvx_vec_reduce_max_f32, HVX_REDUCE_MAX_SCALAR); +} + +static inline float hvx_reduce_max_f32_u(const uint8_t * restrict src, const int num_elems) { + HVX_Vector init_vec = hvx_vec_splat_f32(((const float *) src)[0]); + hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_REDUCE_MAX_OP, hvx_vec_reduce_max_f32, HVX_REDUCE_MAX_SCALAR); +} + +static inline float hvx_reduce_max_f32(const uint8_t * restrict src, const int num_elems) { + if (hex_is_aligned((void *) src, 128)) { + return hvx_reduce_max_f32_a(src, num_elems); + } else { + return hvx_reduce_max_f32_u(src, num_elems); + } +} + +// Sum variants + +static inline float hvx_reduce_sum_f32_a(const uint8_t * restrict src, const int num_elems) { + HVX_Vector init_vec = Q6_V_vsplat_R(0); + assert((unsigned long) src % 128 == 0); + hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_REDUCE_SUM_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR); +} + +static inline float hvx_reduce_sum_f32_u(const uint8_t * restrict src, const int num_elems) { + HVX_Vector init_vec = Q6_V_vsplat_R(0); + hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_REDUCE_SUM_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR); +} + +static inline float hvx_reduce_sum_f32(const uint8_t * restrict src, const int num_elems) { + if (hex_is_aligned((void *) src, 128)) { + return hvx_reduce_sum_f32_a(src, num_elems); + } else { + return hvx_reduce_sum_f32_u(src, num_elems); + } +} + +// Sum of squares variants + +static inline float hvx_sum_of_squares_f32_a(const uint8_t * restrict src, const int num_elems) { + HVX_Vector init_vec = Q6_V_vsplat_R(0); + assert((uintptr_t) src % 128 == 0); + hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_SUM_SQ_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR); +} + +static inline float hvx_sum_of_squares_f32_u(const uint8_t * restrict src, const int num_elems) { + HVX_Vector init_vec = Q6_V_vsplat_R(0); + hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_SUM_SQ_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR); +} + +static inline float hvx_sum_of_squares_f32(const uint8_t * restrict src, const int num_elems) { + if (hex_is_aligned((void *) src, 128)) { + return hvx_sum_of_squares_f32_a(src, num_elems); + } else { + return hvx_sum_of_squares_f32_u(src, num_elems); + } +} + +#undef hvx_reduce_loop_body +#undef HVX_REDUCE_MAX_OP +#undef HVX_REDUCE_SUM_OP +#undef HVX_REDUCE_MAX_SCALAR +#undef HVX_REDUCE_SUM_SCALAR +#undef HVX_SUM_SQ_OP + +#include "hvx-utils-defs.inc" +#include "hex-utils-defs.inc" + +#include +#include +#include + +static inline uint8_t *vtcm_seq_alloc(uint8_t **vtcm_ptr, size_t size) { + uint8_t *p = *vtcm_ptr; + *vtcm_ptr += size; + return p; +} + + +#include "worker-pool-defs.inc" // ============================================================================ // Constants diff --git a/ggml/src/ggml-hexagon/htp/hmx-matmul-ops.c b/ggml/src/ggml-hexagon/htp/hmx-matmul-ops.inc similarity index 83% rename from ggml/src/ggml-hexagon/htp/hmx-matmul-ops.c rename to ggml/src/ggml-hexagon/htp/hmx-matmul-ops.inc index 2666a78a96ad..d1cda6d38c36 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/hmx-matmul-ops.inc @@ -13,22 +13,457 @@ #include #define GGML_COMMON_DECL_C -#include "ggml-common.h" +#include "ggml-common-defs.inc" -#include "hex-dma.h" -#include "worker-pool.h" +#include "hex-dma-defs.inc" +#include "worker-pool-defs.inc" -#include "hvx-utils.h" -#include "hvx-dump.h" -#include "htp-ctx.h" -#include "htp-ops.h" +#include "hvx-utils-defs.inc" +#include + +#include +#include + +#include "hex-utils-defs.inc" +#include "hvx-types-defs.inc" + +static void hvx_vec_dump_f16_n(char * pref, HVX_Vector v, uint32_t n) { + HVX_VectorAlias u = { .v = v }; + + const uint32_t n0 = n / 16; + const uint32_t n1 = n % 16; + int i = 0; + for (; i < n0; i++) { + hex_dump_f16_line(pref, u.fp16 + (16 * i), 16); + } + if (n1) { + hex_dump_f16_line(pref, u.fp16 + (16 * i), n1); + } +} + +static void hvx_vec_dump_f16(char * pref, HVX_Vector v) { + hvx_vec_dump_f16_n(pref, v, 64); +} + +static void hvx_vec_dump_f32_n(char * pref, HVX_Vector v, uint32_t n) { + HVX_VectorAlias u = { .v = v }; + + const uint32_t n0 = n / 16; + const uint32_t n1 = n % 16; + int i = 0; + for (; i < n0; i++) { + hex_dump_f32_line(pref, u.fp32 + (16 * i), 16); + } + if (n1) { + hex_dump_f32_line(pref, u.fp32 + (16 * i), n1); + } +} + +static void hvx_vec_dump_f32_hmt(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + float d[32]; + } u = { .v = v }; + + FARF(HIGH, "%s: %.6f %.6f %.6f %.6f ... %.6f %.6f %.6f %.6f ... %.6f %.6f %.6f %.6f\n", pref, u.d[0], u.d[1], + u.d[2], u.d[3], u.d[12], u.d[13], u.d[14], u.d[15], u.d[28], u.d[29], u.d[30], u.d[31]); +} + +static void hvx_vec_dump_f32(char * pref, HVX_Vector v) { + hvx_vec_dump_f32_n(pref, v, 32); +} + +static void hvx_vec_dump_int32(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int32_t d[32]; + } u = { .v = v }; + + for (int i = 0; i < 32 / 16; i++) { + hex_dump_int32_line(pref, u.d + (16 * i), 16); + } +} + +static void hvx_vec_dump_int32_hmt(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int32_t d[32]; + } u = { .v = v }; + + FARF(HIGH, "%s: %d %d %d %d ... %d %d %d %d ... %d %d %d %d\n", pref, u.d[0], u.d[1], u.d[2], u.d[3], u.d[12], + u.d[13], u.d[14], u.d[15], u.d[28], u.d[29], u.d[30], u.d[31]); +} + +static void hvx_vec_dump_int8_hmt(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int8_t d[128]; + } u = { .v = v }; + + FARF(HIGH, "%s: %d %d %d %d ... %d %d %d %d ... %d %d %d %d\n", pref, u.d[0], u.d[1], u.d[2], u.d[3], u.d[60], + u.d[61], u.d[62], u.d[63], u.d[124], u.d[125], u.d[126], u.d[127]); +} + +static void hvx_vec_dump_int8(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int8_t d[128]; + } u = { .v = v }; + + for (int i = 0; i < 128 / 16; i++) { + hex_dump_int8_line(pref, u.d + (16 * i), 16); + } +} + +static void hvx_vec_dump_uint8(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + uint8_t d[128]; + } u = { .v = v }; + + for (int i = 0; i < 128 / 16; i++) { + hex_dump_uint8_line(pref, u.d + (16 * i), 16); + } +} + +static bool hvx_vec_eq(HVX_Vector v0, HVX_Vector v1, size_t n) { + typedef union { + HVX_Vector v; + int8_t d[128]; + } U; + + U u0 = { .v = v0 }; + U u1 = { .v = v1 }; + + for (int i = 0; i < n; i++) { + if (u0.d[i] != u1.d[i]) { + return false; + } + } + + return true; +} + +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" + +// HMX operation entry-point declarations. +// Ported from htp-ops-lib/include/dsp/ops.h (renamed, benchmark kernels removed). (https://github.com/haozixu/htp-ops-lib) + +#ifndef HMX_OPS_H +#define HMX_OPS_H + +#include +#include + +#include "htp-ops-defs.inc" + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct { + float *dst; + const float *activation; + const __fp16 *permuted_weight; + int m; + int k; + int n; + int act_stride; + int weight_stride; + int dst_stride; + int ne02; + int ne03; + int ne12; + int ne13; + size_t src0_nb2; + size_t src0_nb3; + size_t src1_nb2; + size_t src1_nb3; + size_t dst_nb2; + size_t dst_nb3; +} hmx_matmul_w16a32_batched_params_t; + +// HMX matrix multiplication — tile-permuted FP16 weights, FP32 activation/output +// act_stride: activation row stride in elements (= k for contiguous, or +// nb[1]/sizeof(float) for permuted tensors like attention Q). +// weight_stride: weight row stride in elements (= k for compact weights, or +// nb[1]/sizeof(__fp16) for permuted KV-cache views used by QK). +int hmx_mat_mul_permuted_w16a32(struct htp_context *ctx, + float *restrict dst, + const float *activation, + const __fp16 *permuted_weight, + int m, int k, int n, + int act_stride, + int weight_stride); + +// Batched F16 wrapper over hmx_mat_mul_permuted_w16a32. +// Batch semantics match ggml_mul_mat(): src0 broadcasts to src1 in dims 2/3. +int hmx_mat_mul_permuted_w16a32_batched(struct htp_context *ctx, + const hmx_matmul_w16a32_batched_params_t *params); + +// HMX matrix multiplication — tile-permuted quantised weights (Q4_0/Q8_0/IQ4_NL) +int hmx_mat_mul_permuted_qk_0_d16a32(struct htp_context *ctx, + float *restrict dst, + const float *activation, + const uint8_t *permuted_weight, + int m, int k, int n, + int weight_type); + +// HMX flash attention +int hmx_flash_attn_ext(struct htp_ops_context * octx); + +#ifdef __cplusplus +} +#endif + +#endif // HMX_OPS_H + +// HMX tile-level inline helpers (FP16 32x32 tile operations). +// Ported from htp-ops-lib/include/dsp/hmx_utils.h. (https://github.com/haozixu/htp-ops-lib) + +#include "hvx-base-defs.inc" + +#include +#include +#include + +#define HMX_FP16_TILE_N_ROWS 32 +#define HMX_FP16_TILE_N_COLS 32 +#define HMX_FP16_TILE_N_ELMS 1024 +#define HMX_FP16_TILE_SIZE 2048 + +// Initialise aligned 256-byte area with scale vector + zero padding. +static inline void hmx_init_column_scales(void *out_scales, HVX_Vector v_scale) { + volatile HVX_Vector *pv = (HVX_Vector *) out_scales; + pv[0] = v_scale; + pv[1] = Q6_V_vzero(); +} + +// --- Shared scatter offsets and interleave helper --- + +// vscatter offsets for fused dequant+transpose: write K-values directly to [K][N] tile. +// word[i] = i*128 maps K-row-pair i to byte offset i*128. +// Column offset (n*4) is added at runtime. Entries 0..15 cover one tile (region 2047); +// entries 16..31 cover the next adjacent tile (region 4095) — pick region size at the +// call site to scatter into one tile (masked) or two contiguous tiles (unmasked). +static const int32_t hmx_transpose_scatter_offsets[32] __attribute__((aligned(VLEN))) = { + 0 * 128, 1 * 128, 2 * 128, 3 * 128, 4 * 128, 5 * 128, 6 * 128, 7 * 128, 8 * 128, 9 * 128, 10 * 128, + 11 * 128, 12 * 128, 13 * 128, 14 * 128, 15 * 128, 16 * 128, 17 * 128, 18 * 128, 19 * 128, 20 * 128, 21 * 128, + 22 * 128, 23 * 128, 24 * 128, 25 * 128, 26 * 128, 27 * 128, 28 * 128, 29 * 128, 30 * 128, 31 * 128, +}; + +// Scatter row-major FP16 data (in VTCM scratch) into transposed [K][N] tiles. +// vtcm_src: [n_cols][src_stride] row-major fp16 (only first k elements per row are used) +// vtcm_dst: [n_col_tiles][n_k_tiles][HMX_FP16_TILE_N_ELMS] tile-major interleaved fp16 +// Processes rows [start_row, end_row) for multi-thread slicing. +// Full range: start_row=0, end_row=n_cols. +static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst, + const __fp16 * restrict vtcm_src, + int n_cols, + int k, + int src_stride, + int start_row, + int end_row) { + assert(k % HMX_FP16_TILE_N_COLS == 0); + + const int n_k_tiles = k / HMX_FP16_TILE_N_COLS; + const HVX_Vector v_scat_base = hvx_vmem(hmx_transpose_scatter_offsets); + const HVX_Vector v_scat_step = Q6_V_vsplat_R(4); + const HVX_VectorPred q_mask64 = Q6_Q_vsetq_R(64); + // Each hvx_vmemu load brings 64 fp16 = 128 bytes covering 2 adjacent K-tiles. + // When n_k_tiles is even, scatter into 2 K-tiles per call (region 4095, no mask) + // using the upper half of hmx_transpose_scatter_offsets. Tail one K-tile (when + // n_k_tiles is odd) falls back to single-tile masked scatter. + const bool pair_scatter = (n_k_tiles & 1) == 0; + const size_t pair_region = (size_t) (2 * HMX_FP16_TILE_SIZE - 1); + const size_t single_region = (size_t) (HMX_FP16_TILE_SIZE - 1); + __builtin_assume(k > 0); + __builtin_assume(end_row > start_row); + + if (pair_scatter) { + // Step c by 64 fp16 (two K-tiles per scatter), advance dst by 2 tiles per iter. + const int c_step = 2 * HMX_FP16_TILE_N_COLS; + const size_t c_byte_step = (size_t) c_step * sizeof(__fp16); + const size_t dst_step = 2 * (size_t) HMX_FP16_TILE_N_ELMS; + const int n_c_iters = k / c_step; + + for (int r = start_row; r < end_row; r += 2) { + const int ct = r / HMX_FP16_TILE_N_ROWS; + const int local_r = r % HMX_FP16_TILE_N_ROWS; + const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_cols; + const HVX_Vector v_off0 = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4)); + const HVX_Vector v_off1 = Q6_Vw_vadd_VwVw(v_off0, v_scat_step); + + __fp16 * tile_base = vtcm_dst + (size_t) ct * n_k_tiles * HMX_FP16_TILE_N_ELMS; + const uint8_t * p0 = (const uint8_t *) (vtcm_src + r * src_stride); + const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL; + + if (p1) { + for (int i = 0; i < n_c_iters; ++i) { + HVX_Vector v0 = hvx_vmemu(p0); + p0 += c_byte_step; + HVX_Vector v1 = hvx_vmemu(p1); + p1 += c_byte_step; + Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0); + Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, v1); + tile_base += dst_step; + } + } else { + const HVX_Vector vzero = Q6_V_vzero(); + for (int i = 0; i < n_c_iters; ++i) { + HVX_Vector v0 = hvx_vmemu(p0); + p0 += c_byte_step; + Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0); + Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, vzero); + tile_base += dst_step; + } + } + } + } else { + // Fallback: scatter one K-tile per call (region 2047, masked). + const int c_step = HMX_FP16_TILE_N_COLS; + const size_t c_byte_step = (size_t) c_step * sizeof(__fp16); + const size_t dst_step = (size_t) HMX_FP16_TILE_N_ELMS; + const int n_c_iters = k / c_step; + + for (int r = start_row; r < end_row; r += 2) { + const int ct = r / HMX_FP16_TILE_N_ROWS; + const int local_r = r % HMX_FP16_TILE_N_ROWS; + const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_cols; + const HVX_Vector v_off0 = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4)); + const HVX_Vector v_off1 = Q6_Vw_vadd_VwVw(v_off0, v_scat_step); + + __fp16 * tile_base = vtcm_dst + (size_t) ct * n_k_tiles * HMX_FP16_TILE_N_ELMS; + const uint8_t * p0 = (const uint8_t *) (vtcm_src + r * src_stride); + const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL; + + if (p1) { + for (int i = 0; i < n_c_iters; ++i) { + HVX_Vector v0 = hvx_vmemu(p0); + p0 += c_byte_step; + HVX_Vector v1 = hvx_vmemu(p1); + p1 += c_byte_step; + Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0); + Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, v1); + tile_base += dst_step; + } + } else { + const HVX_Vector vzero = Q6_V_vzero(); + for (int i = 0; i < n_c_iters; ++i) { + HVX_Vector v0 = hvx_vmemu(p0); + p0 += c_byte_step; + Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0); + Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, vzero); + tile_base += dst_step; + } + } + } + } +} + +// Interleave row-major FP16 data into column-major tile format. +// Input: [n_rows, head_dim] row-major. Output: tile[dim_tile][row_tile]. +// Processes rows [start_row, end_row) for multi-thread slicing. +// Full range: start_row=0, end_row=n_rows. +static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out, + const __fp16 * restrict src, + int n_rows, + int head_dim, + int src_stride, + int n_row_tiles, + int start_row, + int end_row) { + __builtin_assume(head_dim > 0); + const size_t tile_stride_elms = (size_t) n_row_tiles * HMX_FP16_TILE_N_ELMS; + + for (int r = start_row; r < end_row; r += 2) { + const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_rows; + + const HVX_Vector * pv_in0 = (const HVX_Vector *) (src + r * src_stride); + const HVX_Vector * pv_in1 = next_row_valid ? (const HVX_Vector *) (src + (r + 1) * src_stride) : NULL; + + // Row-pair invariants hoisted out of the c loop. + const int r0 = r / HMX_FP16_TILE_N_ROWS; + const int r1_half = (r % HMX_FP16_TILE_N_ROWS) / 2; + + // tb0 starts at tile (c0=0, r0); tb1 at the adjacent dim-tile (c0=1, r0). + // Each c step (+= 64) advances both by 2 dim-tiles worth of fp16. + __fp16 * tb0 = tiles_out + (size_t) r0 * HMX_FP16_TILE_N_ELMS; + __fp16 * tb1 = tb0 + tile_stride_elms; + const size_t tb_step = 2 * tile_stride_elms; + + if (pv_in1) { + for (int c = 0; c < head_dim; c += 64) { + HVX_Vector v0 = *pv_in0++; + HVX_Vector v1 = *pv_in1++; + HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2); + ((HVX_Vector *) tb0)[r1_half] = Q6_V_lo_W(vp); + ((HVX_Vector *) tb1)[r1_half] = Q6_V_hi_W(vp); + tb0 += tb_step; + tb1 += tb_step; + } + } else { + const HVX_Vector vzero = Q6_V_vzero(); + for (int c = 0; c < head_dim; c += 64) { + HVX_Vector v0 = *pv_in0++; + HVX_VectorPair vp = Q6_W_vshuff_VVR(vzero, v0, -2); + ((HVX_Vector *) tb0)[r1_half] = Q6_V_lo_W(vp); + ((HVX_Vector *) tb1)[r1_half] = Q6_V_hi_W(vp); + tb0 += tb_step; + tb1 += tb_step; + } + } + } +} + + +#include "hmx-queue-defs.inc" +// Conditional fine-grained profiling macros for HMX operations. +// +// Define ENABLE_PROFILE_TIMERS (via compiler flag or before including this +// header) to instrument sub-operation latencies with HAP qtimer. When the +// macro is not defined the TIMER_* helpers expand to nothing so there is zero +// overhead. +// +// Usage: +// TIMER_DEFINE(my_phase); // declare accumulator variable +// TIMER_START(my_phase); // snapshot start time +// ... work ... +// TIMER_STOP(my_phase); // accumulate elapsed ticks +// FARF(ALWAYS, "my_phase: %lld us", TIMER_US(my_phase)); + +#include + +// #define ENABLE_PROFILE_TIMERS + +#if defined(ENABLE_PROFILE_TIMERS) +# define TIMER_DEFINE(name) int64_t name##_ticks = 0 +# define TIMER_START(name) int64_t name##_t0 = HAP_perf_get_qtimer_count() +# define TIMER_STOP(name) name##_ticks += HAP_perf_get_qtimer_count() - name##_t0 +# define TIMER_US(name) HAP_perf_qtimer_count_to_us(name##_ticks) +#else +# define TIMER_DEFINE(name) +# define TIMER_START(name) +# define TIMER_STOP(name) +# define TIMER_US(name) 0LL +#endif + + + +#include "hex-utils-defs.inc" + +#include +#include +#include + +static inline uint8_t *vtcm_seq_alloc(uint8_t **vtcm_ptr, size_t size) { + uint8_t *p = *vtcm_ptr; + *vtcm_ptr += size; + return p; +} -#include "hmx-ops.h" -#include "hmx-utils.h" -#include "hmx-queue.h" -#include "hmx-profile.h" -#include "vtcm-utils.h" static const __fp16 q4_0_to_fp16_lut[64] __attribute__((aligned(VLEN))) = { -8, 0, -7, 0, -6, 0, -5, 0, -4, 0, -3, 0, -2, 0, -1, 0, 0, 0, 1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0, 7, 0, diff --git a/ggml/src/ggml-hexagon/htp/hmx-ops.h b/ggml/src/ggml-hexagon/htp/hmx-ops.h deleted file mode 100644 index 1c78ffadd1cd..000000000000 --- a/ggml/src/ggml-hexagon/htp/hmx-ops.h +++ /dev/null @@ -1,71 +0,0 @@ -// HMX operation entry-point declarations. -// Ported from htp-ops-lib/include/dsp/ops.h (renamed, benchmark kernels removed). (https://github.com/haozixu/htp-ops-lib) - -#ifndef HMX_OPS_H -#define HMX_OPS_H - -#include -#include - -#include "htp-ops.h" - -#ifdef __cplusplus -extern "C" { -#endif - -typedef struct { - float *dst; - const float *activation; - const __fp16 *permuted_weight; - int m; - int k; - int n; - int act_stride; - int weight_stride; - int dst_stride; - int ne02; - int ne03; - int ne12; - int ne13; - size_t src0_nb2; - size_t src0_nb3; - size_t src1_nb2; - size_t src1_nb3; - size_t dst_nb2; - size_t dst_nb3; -} hmx_matmul_w16a32_batched_params_t; - -// HMX matrix multiplication — tile-permuted FP16 weights, FP32 activation/output -// act_stride: activation row stride in elements (= k for contiguous, or -// nb[1]/sizeof(float) for permuted tensors like attention Q). -// weight_stride: weight row stride in elements (= k for compact weights, or -// nb[1]/sizeof(__fp16) for permuted KV-cache views used by QK). -int hmx_mat_mul_permuted_w16a32(struct htp_context *ctx, - float *restrict dst, - const float *activation, - const __fp16 *permuted_weight, - int m, int k, int n, - int act_stride, - int weight_stride); - -// Batched F16 wrapper over hmx_mat_mul_permuted_w16a32. -// Batch semantics match ggml_mul_mat(): src0 broadcasts to src1 in dims 2/3. -int hmx_mat_mul_permuted_w16a32_batched(struct htp_context *ctx, - const hmx_matmul_w16a32_batched_params_t *params); - -// HMX matrix multiplication — tile-permuted quantised weights (Q4_0/Q8_0/IQ4_NL) -int hmx_mat_mul_permuted_qk_0_d16a32(struct htp_context *ctx, - float *restrict dst, - const float *activation, - const uint8_t *permuted_weight, - int m, int k, int n, - int weight_type); - -// HMX flash attention -int hmx_flash_attn_ext(struct htp_ops_context * octx); - -#ifdef __cplusplus -} -#endif - -#endif // HMX_OPS_H diff --git a/ggml/src/ggml-hexagon/htp/hmx-profile.h b/ggml/src/ggml-hexagon/htp/hmx-profile.h deleted file mode 100644 index 01eece720c57..000000000000 --- a/ggml/src/ggml-hexagon/htp/hmx-profile.h +++ /dev/null @@ -1,34 +0,0 @@ -// Conditional fine-grained profiling macros for HMX operations. -// -// Define ENABLE_PROFILE_TIMERS (via compiler flag or before including this -// header) to instrument sub-operation latencies with HAP qtimer. When the -// macro is not defined the TIMER_* helpers expand to nothing so there is zero -// overhead. -// -// Usage: -// TIMER_DEFINE(my_phase); // declare accumulator variable -// TIMER_START(my_phase); // snapshot start time -// ... work ... -// TIMER_STOP(my_phase); // accumulate elapsed ticks -// FARF(ALWAYS, "my_phase: %lld us", TIMER_US(my_phase)); - -#ifndef HMX_PROFILE_H -#define HMX_PROFILE_H - -#include - -// #define ENABLE_PROFILE_TIMERS - -#if defined(ENABLE_PROFILE_TIMERS) -# define TIMER_DEFINE(name) int64_t name##_ticks = 0 -# define TIMER_START(name) int64_t name##_t0 = HAP_perf_get_qtimer_count() -# define TIMER_STOP(name) name##_ticks += HAP_perf_get_qtimer_count() - name##_t0 -# define TIMER_US(name) HAP_perf_qtimer_count_to_us(name##_ticks) -#else -# define TIMER_DEFINE(name) -# define TIMER_START(name) -# define TIMER_STOP(name) -# define TIMER_US(name) 0LL -#endif - -#endif // HMX_PROFILE_H diff --git a/ggml/src/ggml-hexagon/htp/hmx-queue.h b/ggml/src/ggml-hexagon/htp/hmx-queue-defs.inc similarity index 99% rename from ggml/src/ggml-hexagon/htp/hmx-queue.h rename to ggml/src/ggml-hexagon/htp/hmx-queue-defs.inc index 0d48c280f526..81b91ed69b55 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-queue.h +++ b/ggml/src/ggml-hexagon/htp/hmx-queue-defs.inc @@ -10,7 +10,7 @@ #include #include -#include "hex-utils.h" +#include "hex-utils-defs.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/src/ggml-hexagon/htp/hmx-queue.c b/ggml/src/ggml-hexagon/htp/hmx-queue.inc similarity index 99% rename from ggml/src/ggml-hexagon/htp/hmx-queue.c rename to ggml/src/ggml-hexagon/htp/hmx-queue.inc index 5b1d83a0cbf0..744a70ba7114 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-queue.c +++ b/ggml/src/ggml-hexagon/htp/hmx-queue.inc @@ -9,7 +9,7 @@ #include -#include "hmx-queue.h" +#include "hmx-queue-defs.inc" #define QURT_LOWEST_PRIO (254) diff --git a/ggml/src/ggml-hexagon/htp/hmx-utils.h b/ggml/src/ggml-hexagon/htp/hmx-utils.h deleted file mode 100644 index 68f174d6937c..000000000000 --- a/ggml/src/ggml-hexagon/htp/hmx-utils.h +++ /dev/null @@ -1,202 +0,0 @@ -// HMX tile-level inline helpers (FP16 32x32 tile operations). -// Ported from htp-ops-lib/include/dsp/hmx_utils.h. (https://github.com/haozixu/htp-ops-lib) - -#ifndef HMX_UTILS_H -#define HMX_UTILS_H - -#include "hvx-base.h" - -#include -#include -#include - -#define HMX_FP16_TILE_N_ROWS 32 -#define HMX_FP16_TILE_N_COLS 32 -#define HMX_FP16_TILE_N_ELMS 1024 -#define HMX_FP16_TILE_SIZE 2048 - -// Initialise aligned 256-byte area with scale vector + zero padding. -static inline void hmx_init_column_scales(void *out_scales, HVX_Vector v_scale) { - volatile HVX_Vector *pv = (HVX_Vector *) out_scales; - pv[0] = v_scale; - pv[1] = Q6_V_vzero(); -} - -// --- Shared scatter offsets and interleave helper --- - -// vscatter offsets for fused dequant+transpose: write K-values directly to [K][N] tile. -// word[i] = i*128 maps K-row-pair i to byte offset i*128. -// Column offset (n*4) is added at runtime. Entries 0..15 cover one tile (region 2047); -// entries 16..31 cover the next adjacent tile (region 4095) — pick region size at the -// call site to scatter into one tile (masked) or two contiguous tiles (unmasked). -static const int32_t hmx_transpose_scatter_offsets[32] __attribute__((aligned(VLEN))) = { - 0 * 128, 1 * 128, 2 * 128, 3 * 128, 4 * 128, 5 * 128, 6 * 128, 7 * 128, 8 * 128, 9 * 128, 10 * 128, - 11 * 128, 12 * 128, 13 * 128, 14 * 128, 15 * 128, 16 * 128, 17 * 128, 18 * 128, 19 * 128, 20 * 128, 21 * 128, - 22 * 128, 23 * 128, 24 * 128, 25 * 128, 26 * 128, 27 * 128, 28 * 128, 29 * 128, 30 * 128, 31 * 128, -}; - -// Scatter row-major FP16 data (in VTCM scratch) into transposed [K][N] tiles. -// vtcm_src: [n_cols][src_stride] row-major fp16 (only first k elements per row are used) -// vtcm_dst: [n_col_tiles][n_k_tiles][HMX_FP16_TILE_N_ELMS] tile-major interleaved fp16 -// Processes rows [start_row, end_row) for multi-thread slicing. -// Full range: start_row=0, end_row=n_cols. -static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst, - const __fp16 * restrict vtcm_src, - int n_cols, - int k, - int src_stride, - int start_row, - int end_row) { - assert(k % HMX_FP16_TILE_N_COLS == 0); - - const int n_k_tiles = k / HMX_FP16_TILE_N_COLS; - const HVX_Vector v_scat_base = hvx_vmem(hmx_transpose_scatter_offsets); - const HVX_Vector v_scat_step = Q6_V_vsplat_R(4); - const HVX_VectorPred q_mask64 = Q6_Q_vsetq_R(64); - // Each hvx_vmemu load brings 64 fp16 = 128 bytes covering 2 adjacent K-tiles. - // When n_k_tiles is even, scatter into 2 K-tiles per call (region 4095, no mask) - // using the upper half of hmx_transpose_scatter_offsets. Tail one K-tile (when - // n_k_tiles is odd) falls back to single-tile masked scatter. - const bool pair_scatter = (n_k_tiles & 1) == 0; - const size_t pair_region = (size_t) (2 * HMX_FP16_TILE_SIZE - 1); - const size_t single_region = (size_t) (HMX_FP16_TILE_SIZE - 1); - __builtin_assume(k > 0); - __builtin_assume(end_row > start_row); - - if (pair_scatter) { - // Step c by 64 fp16 (two K-tiles per scatter), advance dst by 2 tiles per iter. - const int c_step = 2 * HMX_FP16_TILE_N_COLS; - const size_t c_byte_step = (size_t) c_step * sizeof(__fp16); - const size_t dst_step = 2 * (size_t) HMX_FP16_TILE_N_ELMS; - const int n_c_iters = k / c_step; - - for (int r = start_row; r < end_row; r += 2) { - const int ct = r / HMX_FP16_TILE_N_ROWS; - const int local_r = r % HMX_FP16_TILE_N_ROWS; - const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_cols; - const HVX_Vector v_off0 = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4)); - const HVX_Vector v_off1 = Q6_Vw_vadd_VwVw(v_off0, v_scat_step); - - __fp16 * tile_base = vtcm_dst + (size_t) ct * n_k_tiles * HMX_FP16_TILE_N_ELMS; - const uint8_t * p0 = (const uint8_t *) (vtcm_src + r * src_stride); - const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL; - - if (p1) { - for (int i = 0; i < n_c_iters; ++i) { - HVX_Vector v0 = hvx_vmemu(p0); - p0 += c_byte_step; - HVX_Vector v1 = hvx_vmemu(p1); - p1 += c_byte_step; - Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0); - Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, v1); - tile_base += dst_step; - } - } else { - const HVX_Vector vzero = Q6_V_vzero(); - for (int i = 0; i < n_c_iters; ++i) { - HVX_Vector v0 = hvx_vmemu(p0); - p0 += c_byte_step; - Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0); - Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, vzero); - tile_base += dst_step; - } - } - } - } else { - // Fallback: scatter one K-tile per call (region 2047, masked). - const int c_step = HMX_FP16_TILE_N_COLS; - const size_t c_byte_step = (size_t) c_step * sizeof(__fp16); - const size_t dst_step = (size_t) HMX_FP16_TILE_N_ELMS; - const int n_c_iters = k / c_step; - - for (int r = start_row; r < end_row; r += 2) { - const int ct = r / HMX_FP16_TILE_N_ROWS; - const int local_r = r % HMX_FP16_TILE_N_ROWS; - const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_cols; - const HVX_Vector v_off0 = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4)); - const HVX_Vector v_off1 = Q6_Vw_vadd_VwVw(v_off0, v_scat_step); - - __fp16 * tile_base = vtcm_dst + (size_t) ct * n_k_tiles * HMX_FP16_TILE_N_ELMS; - const uint8_t * p0 = (const uint8_t *) (vtcm_src + r * src_stride); - const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL; - - if (p1) { - for (int i = 0; i < n_c_iters; ++i) { - HVX_Vector v0 = hvx_vmemu(p0); - p0 += c_byte_step; - HVX_Vector v1 = hvx_vmemu(p1); - p1 += c_byte_step; - Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0); - Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, v1); - tile_base += dst_step; - } - } else { - const HVX_Vector vzero = Q6_V_vzero(); - for (int i = 0; i < n_c_iters; ++i) { - HVX_Vector v0 = hvx_vmemu(p0); - p0 += c_byte_step; - Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0); - Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, vzero); - tile_base += dst_step; - } - } - } - } -} - -// Interleave row-major FP16 data into column-major tile format. -// Input: [n_rows, head_dim] row-major. Output: tile[dim_tile][row_tile]. -// Processes rows [start_row, end_row) for multi-thread slicing. -// Full range: start_row=0, end_row=n_rows. -static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out, - const __fp16 * restrict src, - int n_rows, - int head_dim, - int src_stride, - int n_row_tiles, - int start_row, - int end_row) { - __builtin_assume(head_dim > 0); - const size_t tile_stride_elms = (size_t) n_row_tiles * HMX_FP16_TILE_N_ELMS; - - for (int r = start_row; r < end_row; r += 2) { - const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_rows; - - const HVX_Vector * pv_in0 = (const HVX_Vector *) (src + r * src_stride); - const HVX_Vector * pv_in1 = next_row_valid ? (const HVX_Vector *) (src + (r + 1) * src_stride) : NULL; - - // Row-pair invariants hoisted out of the c loop. - const int r0 = r / HMX_FP16_TILE_N_ROWS; - const int r1_half = (r % HMX_FP16_TILE_N_ROWS) / 2; - - // tb0 starts at tile (c0=0, r0); tb1 at the adjacent dim-tile (c0=1, r0). - // Each c step (+= 64) advances both by 2 dim-tiles worth of fp16. - __fp16 * tb0 = tiles_out + (size_t) r0 * HMX_FP16_TILE_N_ELMS; - __fp16 * tb1 = tb0 + tile_stride_elms; - const size_t tb_step = 2 * tile_stride_elms; - - if (pv_in1) { - for (int c = 0; c < head_dim; c += 64) { - HVX_Vector v0 = *pv_in0++; - HVX_Vector v1 = *pv_in1++; - HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2); - ((HVX_Vector *) tb0)[r1_half] = Q6_V_lo_W(vp); - ((HVX_Vector *) tb1)[r1_half] = Q6_V_hi_W(vp); - tb0 += tb_step; - tb1 += tb_step; - } - } else { - const HVX_Vector vzero = Q6_V_vzero(); - for (int c = 0; c < head_dim; c += 64) { - HVX_Vector v0 = *pv_in0++; - HVX_VectorPair vp = Q6_W_vshuff_VVR(vzero, v0, -2); - ((HVX_Vector *) tb0)[r1_half] = Q6_V_lo_W(vp); - ((HVX_Vector *) tb1)[r1_half] = Q6_V_hi_W(vp); - tb0 += tb_step; - tb1 += tb_step; - } - } - } -} - -#endif // HMX_UTILS_H diff --git a/ggml/src/ggml-hexagon/htp/htp-ctx.h b/ggml/src/ggml-hexagon/htp/htp-ctx-defs.inc similarity index 96% rename from ggml/src/ggml-hexagon/htp/htp-ctx.h rename to ggml/src/ggml-hexagon/htp/htp-ctx-defs.inc index e9c563ca887b..c3d8c763ff44 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ctx.h +++ b/ggml/src/ggml-hexagon/htp/htp-ctx-defs.inc @@ -1,10 +1,10 @@ #ifndef HTP_CTX_H #define HTP_CTX_H -#include "hex-dma.h" -#include "hmx-queue.h" -#include "htp-ops.h" -#include "worker-pool.h" +#include "hex-dma-defs.inc" +#include "hmx-queue-defs.inc" +#include "htp-ops-defs.inc" +#include "worker-pool-defs.inc" #include #include diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops-defs.inc similarity index 97% rename from ggml/src/ggml-hexagon/htp/htp-ops.h rename to ggml/src/ggml-hexagon/htp/htp-ops-defs.inc index 66a3150c1a08..61e00ac0e47e 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops-defs.inc @@ -3,7 +3,7 @@ #include -// ggml-common.h must be included prio to this header +// ggml-common-defs.inc must be included prio to this header enum htp_status { HTP_STATUS_OK = 1, @@ -14,7 +14,7 @@ enum htp_status { }; // First set of values must match the ggml_type. -// Duplicated here because we can't include full ggml.h in the htp build. +// Duplicated here because we can't include full ggml.h.inc in the htp build. // We have some static_asserts in the cpp code to ensure things are in sync. enum htp_data_type { HTP_TYPE_F32 = 0, @@ -111,7 +111,7 @@ struct htp_tensor { uint16_t type; // Data type uint16_t bi; // Buffer index uint32_t ne[HTP_OP_MAX_DIMS]; // Number of elements - uint32_t nb[HTP_OP_MAX_DIMS]; // Stride in bytes (see ggml.h ggml_tensor) + uint32_t nb[HTP_OP_MAX_DIMS]; // Stride in bytes (see ggml.h.inc ggml_tensor) }; // Buffer descriptor diff --git a/ggml/src/ggml-hexagon/htp/hvx-arith.h b/ggml/src/ggml-hexagon/htp/hvx-arith.h deleted file mode 100644 index 82e3416970b4..000000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-arith.h +++ /dev/null @@ -1,443 +0,0 @@ -#ifndef HVX_ARITH_H -#define HVX_ARITH_H - -#include -#include -#include -#include - -#include "hvx-base.h" -#include "hex-utils.h" - -// -// Binary operations (add, mul, sub) -// - -#define UNUSED(x) (void)(x) - -#define hvx_arith_loop_body(dst_type, src0_type, src1_type, elem_size, vec_store, vec_op) \ - do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src0_type * restrict vsrc0 = (src0_type *) src0; \ - src1_type * restrict vsrc1 = (src1_type *) src1; \ - \ - const uint32_t epv = 128 / (elem_size); \ - const uint32_t nvec = n / epv; \ - const uint32_t nloe = n % epv; \ - \ - uint32_t i = 0; \ - \ - _Pragma("unroll(4)") \ - for (; i < nvec; i++) { \ - vdst[i] = vec_op(vsrc0[i], vsrc1[i]); \ - } \ - if (nloe) { \ - HVX_Vector v = vec_op(vsrc0[i], vsrc1[i]); \ - vec_store((void *) &vdst[i], nloe * (elem_size), v); \ - } \ - } while(0) - -#if __HVX_ARCH__ < 79 - -#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b)) -#define HVX_OP_SUB_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(a, b)) -#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b)) - -#else - -#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b) -#define HVX_OP_SUB_F32(a, b) Q6_Vsf_vsub_VsfVsf(a, b) -#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b) - -#endif - -#define HVX_OP_ADD_F16(a, b) hvx_vec_add_f16_f16(a, b) -#define HVX_OP_SUB_F16(a, b) hvx_vec_sub_f16_f16(a, b) -#define HVX_OP_MUL_F16(a, b) hvx_vec_mul_f16_f16(a, b) - -// Generic macro to define alignment permutations for an op -#define DEFINE_HVX_BINARY_OP_VARIANTS(OP_NAME, OP_MACRO, ELEM_TYPE) \ -static inline void OP_NAME##_aaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) dst % 128 == 0); \ - assert((uintptr_t) src0 % 128 == 0); \ - assert((uintptr_t) src1 % 128 == 0); \ - hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ -} \ -static inline void OP_NAME##_aau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) dst % 128 == 0); \ - assert((uintptr_t) src0 % 128 == 0); \ - hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ -} \ -static inline void OP_NAME##_aua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) dst % 128 == 0); \ - assert((uintptr_t) src1 % 128 == 0); \ - hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ -} \ -static inline void OP_NAME##_auu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) dst % 128 == 0); \ - hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ -} \ -static inline void OP_NAME##_uaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) src0 % 128 == 0); \ - assert((uintptr_t) src1 % 128 == 0); \ - hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ -} \ -static inline void OP_NAME##_uau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) src0 % 128 == 0); \ - hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ -} \ -static inline void OP_NAME##_uua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) src1 % 128 == 0); \ - hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ -} \ -static inline void OP_NAME##_uuu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ -} \ - -DEFINE_HVX_BINARY_OP_VARIANTS(hvx_add_f32, HVX_OP_ADD_F32, float) -DEFINE_HVX_BINARY_OP_VARIANTS(hvx_sub_f32, HVX_OP_SUB_F32, float) -DEFINE_HVX_BINARY_OP_VARIANTS(hvx_mul_f32, HVX_OP_MUL_F32, float) - -DEFINE_HVX_BINARY_OP_VARIANTS(hvx_add_f16, HVX_OP_ADD_F16, _Float16) -DEFINE_HVX_BINARY_OP_VARIANTS(hvx_sub_f16, HVX_OP_SUB_F16, _Float16) -DEFINE_HVX_BINARY_OP_VARIANTS(hvx_mul_f16, HVX_OP_MUL_F16, _Float16) - -// Dispatcher logic -#define HVX_BINARY_DISPATCHER(OP_NAME) \ -static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, const uint32_t num_elems) { \ - if (hex_is_aligned((void *) dst, 128)) { \ - if (hex_is_aligned((void *) src0, 128)) { \ - if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems); \ - else OP_NAME##_aau(dst, src0, src1, num_elems); \ - } else { \ - if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems); \ - else OP_NAME##_auu(dst, src0, src1, num_elems); \ - } \ - } else { \ - if (hex_is_aligned((void *) src0, 128)) { \ - if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems); \ - else OP_NAME##_uau(dst, src0, src1, num_elems); \ - } else { \ - if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems); \ - else OP_NAME##_uuu(dst, src0, src1, num_elems); \ - } \ - } \ -} - -HVX_BINARY_DISPATCHER(hvx_add_f32) -HVX_BINARY_DISPATCHER(hvx_sub_f32) -HVX_BINARY_DISPATCHER(hvx_mul_f32) - -HVX_BINARY_DISPATCHER(hvx_add_f16) -HVX_BINARY_DISPATCHER(hvx_sub_f16) -HVX_BINARY_DISPATCHER(hvx_mul_f16) - -// Mul-Mul Optimized -static inline void hvx_mul_mul_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, const uint8_t * restrict src2, const uint32_t num_elems) { - assert((unsigned long) dst % 128 == 0); - assert((unsigned long) src0 % 128 == 0); - assert((unsigned long) src1 % 128 == 0); - assert((unsigned long) src2 % 128 == 0); - - HVX_Vector * restrict vdst = (HVX_Vector *) dst; - HVX_Vector * restrict vsrc0 = (HVX_Vector *) src0; - HVX_Vector * restrict vsrc1 = (HVX_Vector *) src1; - HVX_Vector * restrict vsrc2 = (HVX_Vector *) src2; - - const uint32_t elem_size = sizeof(float); - const uint32_t epv = 128 / elem_size; - const uint32_t nvec = num_elems / epv; - const uint32_t nloe = num_elems % epv; - - uint32_t i = 0; - - _Pragma("unroll(4)") - for (; i < nvec; i++) { - HVX_Vector v1 = HVX_OP_MUL_F32(vsrc0[i], vsrc1[i]); - vdst[i] = HVX_OP_MUL(v1, vsrc2[i]); - } - - if (nloe) { - HVX_Vector v1 = HVX_OP_MUL_F32(vsrc0[i], vsrc1[i]); - HVX_Vector v2 = HVX_OP_MUL_F32(v1, vsrc2[i]); - hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v2); - } -} - -// Scalar Operations - -#define hvx_scalar_loop_body(dst_type, src_type, elem_size, vec_store, scalar_op_macro) \ - do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src_type * restrict vsrc = (src_type *) src; \ - \ - const uint32_t epv = 128 / (elem_size); \ - const uint32_t nvec = n / epv; \ - const uint32_t nloe = n % epv; \ - \ - uint32_t i = 0; \ - \ - _Pragma("unroll(4)") \ - for (; i < nvec; i++) { \ - HVX_Vector v = vsrc[i]; \ - vdst[i] = scalar_op_macro(v); \ - } \ - if (nloe) { \ - HVX_Vector v = vsrc[i]; \ - v = scalar_op_macro(v); \ - vec_store((void *) &vdst[i], nloe * (elem_size), v); \ - } \ - } while(0) - -#define HVX_OP_ADD_SCALAR_F32(v) \ - ({ \ - const HVX_VectorPred pred_inf = Q6_Q_vcmp_eq_VwVw(inf, v); \ - HVX_Vector out = HVX_OP_ADD_F32(v, val_vec); \ - Q6_V_vmux_QVV(pred_inf, inf, out); \ - }) - -#define HVX_OP_MUL_SCALAR_F32(v) HVX_OP_MUL_F32(v, val_vec) -#define HVX_OP_SUB_SCALAR_F32(v) HVX_OP_SUB_F32(v, val_vec) - -#define HVX_OP_ADD_SCALAR_F16(v) \ - ({ \ - const HVX_VectorPred pred_inf = Q6_Q_vcmp_eq_VhVh(inf, v); \ - HVX_Vector out = HVX_OP_ADD_F16(v, val_vec); \ - Q6_V_vmux_QVV(pred_inf, inf, out); \ - }) - -#define HVX_OP_MUL_SCALAR_F16(v) HVX_OP_MUL_F16(v, val_vec) -#define HVX_OP_SUB_SCALAR_F16(v) HVX_OP_SUB_F16(v, val_vec) - -// Scalar Variants - -// Generic macro to define alignment permutations for an op -#define DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(OP_NAME, OP_MACRO, SPLAT_MACRO, ELEM_TYPE) \ -static inline void OP_NAME##_aa(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \ - const HVX_Vector val_vec = SPLAT_MACRO(val); \ - const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \ - assert((uintptr_t) dst % 128 == 0); \ - assert((uintptr_t) src % 128 == 0); \ - hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ -} \ -static inline void OP_NAME##_au(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \ - const HVX_Vector val_vec = SPLAT_MACRO(val); \ - const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \ - assert((uintptr_t) dst % 128 == 0); \ - hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ -} \ -static inline void OP_NAME##_ua(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \ - const HVX_Vector val_vec = SPLAT_MACRO(val); \ - const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \ - assert((uintptr_t) src % 128 == 0); \ - hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ -} \ -static inline void OP_NAME##_uu(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \ - const HVX_Vector val_vec = SPLAT_MACRO(val); \ - const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \ - hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ -} \ - -DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_add_scalar_f32, HVX_OP_ADD_SCALAR_F32, hvx_vec_splat_f32, float) -DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_sub_scalar_f32, HVX_OP_SUB_SCALAR_F32, hvx_vec_splat_f32, float) -DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_mul_scalar_f32, HVX_OP_MUL_SCALAR_F32, hvx_vec_splat_f32, float) - -DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_add_scalar_f16, HVX_OP_ADD_SCALAR_F16, hvx_vec_splat_f16, _Float16) -DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_sub_scalar_f16, HVX_OP_SUB_SCALAR_F16, hvx_vec_splat_f16, _Float16) -DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_mul_scalar_f16, HVX_OP_MUL_SCALAR_F16, hvx_vec_splat_f16, _Float16) - -// Dispatcher logic -#define HVX_BINARY_SCALAR_DISPATCHER(OP_NAME, ELEM_TYPE) \ -static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, const uint32_t num_elems) { \ - if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { \ - OP_NAME##_aa(dst, src, val, num_elems); \ - } else if (hex_is_aligned((void *) dst, 128)) { \ - OP_NAME##_au(dst, src, val, num_elems); \ - } else if (hex_is_aligned((void *) src, 128)) { \ - OP_NAME##_ua(dst, src, val, num_elems); \ - } else { \ - OP_NAME##_uu(dst, src, val, num_elems); \ - } \ -} - -HVX_BINARY_SCALAR_DISPATCHER(hvx_add_scalar_f32, float) -HVX_BINARY_SCALAR_DISPATCHER(hvx_sub_scalar_f32, float) -HVX_BINARY_SCALAR_DISPATCHER(hvx_mul_scalar_f32, float) - -HVX_BINARY_SCALAR_DISPATCHER(hvx_add_scalar_f16, _Float16) -HVX_BINARY_SCALAR_DISPATCHER(hvx_sub_scalar_f16, _Float16) -HVX_BINARY_SCALAR_DISPATCHER(hvx_mul_scalar_f16, _Float16) - -// MIN Scalar variants - -#define HVX_OP_MIN_SCALAR(v) Q6_Vsf_vmin_VsfVsf(val_vec, v) - -static inline void hvx_min_scalar_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) { - const HVX_Vector val_vec = hvx_vec_splat_f32(val); - assert((unsigned long) dst % 128 == 0); - assert((unsigned long) src % 128 == 0); - hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(float), hvx_vec_store_a, HVX_OP_MIN_SCALAR); -} - -static inline void hvx_min_scalar_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) { - const HVX_Vector val_vec = hvx_vec_splat_f32(val); - assert((unsigned long) dst % 128 == 0); - hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(float), hvx_vec_store_a, HVX_OP_MIN_SCALAR); -} - -static inline void hvx_min_scalar_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) { - const HVX_Vector val_vec = hvx_vec_splat_f32(val); - assert((unsigned long) src % 128 == 0); - hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(float), hvx_vec_store_u, HVX_OP_MIN_SCALAR); -} - -static inline void hvx_min_scalar_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) { - const HVX_Vector val_vec = hvx_vec_splat_f32(val); - hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(float), hvx_vec_store_u, HVX_OP_MIN_SCALAR); -} - -static inline void hvx_min_scalar_f32(uint8_t * restrict dst, const uint8_t * restrict src, const float val, const int num_elems) { - if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { - hvx_min_scalar_f32_aa(dst, src, val, num_elems); - } else if (hex_is_aligned((void *) dst, 128)) { - hvx_min_scalar_f32_au(dst, src, val, num_elems); - } else if (hex_is_aligned((void *) src, 128)) { - hvx_min_scalar_f32_ua(dst, src, val, num_elems); - } else { - hvx_min_scalar_f32_uu(dst, src, val, num_elems); - } -} - -// CLAMP Scalar variants - -#define HVX_OP_CLAMP_SCALAR(v) \ - ({ \ - HVX_VectorPred pred_cap_right = Q6_Q_vcmp_gt_VsfVsf(v, max_vec); \ - HVX_VectorPred pred_cap_left = Q6_Q_vcmp_gt_VsfVsf(min_vec, v); \ - HVX_Vector tmp = Q6_V_vmux_QVV(pred_cap_right, max_vec, v); \ - Q6_V_vmux_QVV(pred_cap_left, min_vec, tmp); \ - }) - -static inline void hvx_clamp_scalar_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) { - const HVX_Vector min_vec = hvx_vec_splat_f32(min); - const HVX_Vector max_vec = hvx_vec_splat_f32(max); - assert((unsigned long) dst % 128 == 0); - assert((unsigned long) src % 128 == 0); - hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(float), hvx_vec_store_a, HVX_OP_CLAMP_SCALAR); -} - -static inline void hvx_clamp_scalar_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) { - const HVX_Vector min_vec = hvx_vec_splat_f32(min); - const HVX_Vector max_vec = hvx_vec_splat_f32(max); - assert((unsigned long) dst % 128 == 0); - hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(float), hvx_vec_store_a, HVX_OP_CLAMP_SCALAR); -} - -static inline void hvx_clamp_scalar_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) { - const HVX_Vector min_vec = hvx_vec_splat_f32(min); - const HVX_Vector max_vec = hvx_vec_splat_f32(max); - assert((unsigned long) src % 128 == 0); - hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(float), hvx_vec_store_u, HVX_OP_CLAMP_SCALAR); -} - -static inline void hvx_clamp_scalar_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) { - const HVX_Vector min_vec = hvx_vec_splat_f32(min); - const HVX_Vector max_vec = hvx_vec_splat_f32(max); - hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(float), hvx_vec_store_u, HVX_OP_CLAMP_SCALAR); -} - -static inline void hvx_clamp_scalar_f32(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, const int num_elems) { - if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { - hvx_clamp_scalar_f32_aa(dst, src, min, max, num_elems); - } else if (hex_is_aligned((void *) dst, 128)) { - hvx_clamp_scalar_f32_au(dst, src, min, max, num_elems); - } else if (hex_is_aligned((void *) src, 128)) { - hvx_clamp_scalar_f32_ua(dst, src, min, max, num_elems); - } else { - hvx_clamp_scalar_f32_uu(dst, src, min, max, num_elems); - } -} - -// -// Square -// - -#define hvx_sqr_f32_loop_body(dst_type, src_type, vec_store) \ - do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src_type * restrict vsrc = (src_type *) src; \ - \ - const uint32_t elem_size = sizeof(float); \ - const uint32_t epv = 128 / elem_size; \ - const uint32_t nvec = n / epv; \ - const uint32_t nloe = n % epv; \ - \ - uint32_t i = 0; \ - \ - _Pragma("unroll(4)") \ - for (; i < nvec; i++) { \ - vdst[i] = HVX_OP_MUL_F32(vsrc[i], vsrc[i]); \ - } \ - if (nloe) { \ - HVX_Vector v = HVX_OP_MUL_F32(vsrc[i], vsrc[i]); \ - vec_store((void *) &vdst[i], nloe * elem_size, v); \ - } \ - } while(0) - -static inline void hvx_sqr_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - assert((unsigned long) dst % 128 == 0); - assert((unsigned long) src % 128 == 0); - hvx_sqr_f32_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); -} - -static inline void hvx_sqr_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - assert((unsigned long) dst % 128 == 0); - hvx_sqr_f32_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); -} - -static inline void hvx_sqr_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - assert((unsigned long) src % 128 == 0); - hvx_sqr_f32_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); -} - -static inline void hvx_sqr_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - hvx_sqr_f32_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); -} - -static inline void hvx_sqr_f32(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t num_elems) { - if (hex_is_aligned((void *) dst, 128)) { - if (hex_is_aligned((void *) src, 128)) { - hvx_sqr_f32_aa(dst, src, num_elems); - } else { - hvx_sqr_f32_au(dst, src, num_elems); - } - } else { - if (hex_is_aligned((void *) src, 128)) { - hvx_sqr_f32_ua(dst, src, num_elems); - } else { - hvx_sqr_f32_uu(dst, src, num_elems); - } - } -} - -#undef HVX_OP_ADD_F32 -#undef HVX_OP_SUB_F32 -#undef HVX_OP_MUL_F32 -#undef HVX_OP_ADD_F16 -#undef HVX_OP_SUB_F16 -#undef HVX_OP_MUL_F16 -#undef hvx_arith_loop_body -#undef HVX_OP_ADD_SCALAR_F32 -#undef HVX_OP_SUB_SCALAR_F32 -#undef HVX_OP_MUL_SCALAR_F32 -#undef HVX_OP_ADD_SCALAR_F16 -#undef HVX_OP_SUB_SCALAR_F16 -#undef HVX_OP_MUL_SCALAR_F16 -#undef hvx_scalar_loop_body -#undef HVX_OP_MIN_SCALAR -#undef HVX_OP_CLAMP_SCALAR -#undef DEFINE_HVX_BINARY_OP_VARIANTS -#undef HVX_BINARY_DISPATCHER -#undef UNUSED - -#endif // HVX_ARITH_H diff --git a/ggml/src/ggml-hexagon/htp/hvx-base.h b/ggml/src/ggml-hexagon/htp/hvx-base-defs.inc similarity index 99% rename from ggml/src/ggml-hexagon/htp/hvx-base.h rename to ggml/src/ggml-hexagon/htp/hvx-base-defs.inc index f6cb02951d0c..8d5b04470e35 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-base.h +++ b/ggml/src/ggml-hexagon/htp/hvx-base-defs.inc @@ -6,8 +6,8 @@ #include #include -#include "hex-utils.h" -#include "hvx-types.h" +#include "hex-utils-defs.inc" +#include "hvx-types-defs.inc" #define hvx_vmem(A) *((HVX_Vector *)(A)) #define hvx_vmemu(A) *((HVX_UVector *)(A)) diff --git a/ggml/src/ggml-hexagon/htp/hvx-div.h b/ggml/src/ggml-hexagon/htp/hvx-div.h deleted file mode 100644 index 53ee304e749b..000000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-div.h +++ /dev/null @@ -1,291 +0,0 @@ -#ifndef HVX_DIV_H -#define HVX_DIV_H - -#include - -#include -#include -#include -#include -#include - -#include "hvx-base.h" -#include "hex-utils.h" -#include "hvx-inverse.h" -#include "hvx-arith.h" - -#if __HVX_ARCH__ < 79 -#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b)) -#define HVX_OP_MUL_F16(a, b) Q6_Vhf_equals_Wqf32(Q6_Wqf32_vmpy_VhfVhf(a, b)) -#else -#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b) -#define HVX_OP_MUL_F16(a, b) Q6_Vhf_vmpy_VhfVhf(a, b) -#endif - -// Compute div by scaler in f32. Requires first by expanding fp32 to fp16 and converting the result back to fp32. -static inline HVX_Vector hvx_div_mul_f16_const_using_f32(HVX_Vector vec1_hf, HVX_Vector vec2_sf_const, HVX_Vector vec_hf_one_1_0) { -#if __HVX_ARCH__ < 79 - HVX_VectorPair src_to_f32 = Q6_Wqf32_vmpy_VhfVhf(vec1_hf, vec_hf_one_1_0); - HVX_Vector src_to_f32_0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(src_to_f32)); - HVX_Vector src_to_f32_1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(src_to_f32)); -#else - HVX_VectorPair src_to_f32 = Q6_Wsf_vmpy_VhfVhf(vec1_hf, vec_hf_one_1_0); - HVX_Vector src_to_f32_0 = Q6_V_lo_W(src_to_f32); - HVX_Vector src_to_f32_1 = Q6_V_hi_W(src_to_f32); -#endif - - HVX_Vector div_f32_0 = HVX_OP_MUL_F32(src_to_f32_0, vec2_sf_const); - HVX_Vector div_f32_1 = HVX_OP_MUL_F32(src_to_f32_1, vec2_sf_const); - -#if __HVX_ARCH__ < 79 - HVX_Vector res = hvx_vec_f32_to_f16(div_f32_0, div_f32_1); -#else - HVX_Vector res = Q6_Vhf_vcvt_VsfVsf(div_f32_0, div_f32_1); -#endif - return res; -} - -// Variant for =v79 -static inline HVX_Vector hvx_vec_hybrid_div_f16(HVX_Vector vec1, HVX_Vector vec2, HVX_Vector f32_nan_inf_mask, HVX_Vector f16_nan_inf_mask, HVX_Vector vec_hf_one_1_0) { -#if __HVX_ARCH__ < 79 - // For older architectures, use f16 reciprocal to avoid NaN/-inf issues - HVX_Vector vec2_inv = hvx_vec_inverse_f16_guard(vec2, f16_nan_inf_mask); - return HVX_OP_MUL_F16(vec1, vec2_inv); -#else - return hvx_vec_div_f16_using_f32(vec1, vec2, f32_nan_inf_mask, vec_hf_one_1_0); -#endif -} - -#define hvx_div_f16_loop_body(dst_type, src0_type, src1_type, vec_store) \ - do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src0_type * restrict vsrc0 = (src0_type *) src0; \ - src1_type * restrict vsrc1 = (src1_type *) src1; \ - \ - const HVX_Vector f32_nan_inf_mask = Q6_V_vsplat_R(0x7f800000); \ - const HVX_Vector f16_nan_inf_mask = Q6_Vh_vsplat_R(0x7c00); \ - const HVX_Vector hf_one = Q6_Vh_vsplat_R(0x3C00); \ - \ - const uint32_t nvec = n / VLEN_FP16; \ - const uint32_t nloe = n % VLEN_FP16; \ - \ - uint32_t i = 0; \ - \ - _Pragma("unroll(4)") \ - for (; i < nvec; i++) { \ - HVX_Vector res = hvx_vec_hybrid_div_f16(vsrc0[i], vsrc1[i], \ - f32_nan_inf_mask, f16_nan_inf_mask, \ - hf_one); \ - vdst[i] = res; \ - } \ - if (nloe) { \ - HVX_Vector res = hvx_vec_hybrid_div_f16(vsrc0[i], vsrc1[i], \ - f32_nan_inf_mask, f16_nan_inf_mask, \ - hf_one); \ - vec_store((void *) &vdst[i], nloe * SIZEOF_FP16, res); \ - } \ - } while(0) - -#define hvx_div_f32_loop_body(dst_type, src0_type, src1_type, vec_store) \ - do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src0_type * restrict vsrc0 = (src0_type *) src0; \ - src1_type * restrict vsrc1 = (src1_type *) src1; \ - \ - const HVX_Vector nan_inf_mask = Q6_V_vsplat_R(0x7f800000); \ - \ - const uint32_t nvec = n / VLEN_FP32; \ - const uint32_t nloe = n % VLEN_FP32; \ - \ - uint32_t i = 0; \ - \ - _Pragma("unroll(4)") \ - for (; i < nvec; i++) { \ - HVX_Vector inv_src1 = hvx_vec_inverse_f32_guard(vsrc1[i], nan_inf_mask); \ - HVX_Vector res = HVX_OP_MUL_F32(vsrc0[i], inv_src1); \ - vdst[i] = res; \ - } \ - if (nloe) { \ - HVX_Vector inv_src1 = hvx_vec_inverse_f32_guard(vsrc1[i], nan_inf_mask); \ - HVX_Vector res = HVX_OP_MUL_F32(vsrc0[i], inv_src1); \ - vec_store((void *) &vdst[i], nloe * SIZEOF_FP32, res); \ - } \ - } while(0) - -// Generic macro to define alignment permutations for an op -#define DEFINE_HVX_DIV_OP_VARIANTS(OP_NAME, OP_LOOP_BODY) \ -static inline void OP_NAME##_aaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) dst % 128 == 0); \ - assert((uintptr_t) src0 % 128 == 0); \ - assert((uintptr_t) src1 % 128 == 0); \ - OP_LOOP_BODY(HVX_Vector, HVX_Vector, HVX_Vector, hvx_vec_store_a); \ -} \ -static inline void OP_NAME##_aau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) dst % 128 == 0); \ - assert((uintptr_t) src0 % 128 == 0); \ - OP_LOOP_BODY(HVX_Vector, HVX_Vector, HVX_UVector, hvx_vec_store_a); \ -} \ -static inline void OP_NAME##_aua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) dst % 128 == 0); \ - assert((uintptr_t) src1 % 128 == 0); \ - OP_LOOP_BODY(HVX_Vector, HVX_UVector, HVX_Vector, hvx_vec_store_a); \ -} \ -static inline void OP_NAME##_auu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) dst % 128 == 0); \ - OP_LOOP_BODY(HVX_Vector, HVX_UVector, HVX_UVector, hvx_vec_store_a); \ -} \ -static inline void OP_NAME##_uaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) src0 % 128 == 0); \ - assert((uintptr_t) src1 % 128 == 0); \ - OP_LOOP_BODY(HVX_UVector, HVX_Vector, HVX_Vector, hvx_vec_store_u); \ -} \ -static inline void OP_NAME##_uau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) src0 % 128 == 0); \ - OP_LOOP_BODY(HVX_UVector, HVX_Vector, HVX_UVector, hvx_vec_store_u); \ -} \ -static inline void OP_NAME##_uua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - assert((uintptr_t) src1 % 128 == 0); \ - OP_LOOP_BODY(HVX_UVector, HVX_UVector, HVX_Vector, hvx_vec_store_u); \ -} \ -static inline void OP_NAME##_uuu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ - OP_LOOP_BODY(HVX_UVector, HVX_UVector, HVX_UVector, hvx_vec_store_u); \ -} \ - -// Dispatcher logic -#define HVX_DIV_DISPATCHER(OP_NAME) \ -static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, const uint32_t num_elems) { \ - if (hex_is_aligned((void *) dst, 128)) { \ - if (hex_is_aligned((void *) src0, 128)) { \ - if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems); \ - else OP_NAME##_aau(dst, src0, src1, num_elems); \ - } else { \ - if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems); \ - else OP_NAME##_auu(dst, src0, src1, num_elems); \ - } \ - } else { \ - if (hex_is_aligned((void *) src0, 128)) { \ - if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems); \ - else OP_NAME##_uau(dst, src0, src1, num_elems); \ - } else { \ - if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems); \ - else OP_NAME##_uuu(dst, src0, src1, num_elems); \ - } \ - } \ -} - -DEFINE_HVX_DIV_OP_VARIANTS(hvx_div_f32, hvx_div_f32_loop_body) -DEFINE_HVX_DIV_OP_VARIANTS(hvx_div_f16, hvx_div_f16_loop_body) - -HVX_DIV_DISPATCHER(hvx_div_f32) -HVX_DIV_DISPATCHER(hvx_div_f16) - -#undef HVX_OP_MUL_F32 -#undef HVX_OP_MUL_F16 - -#endif // HVX_DIV_H diff --git a/ggml/src/ggml-hexagon/htp/hvx-dump.h b/ggml/src/ggml-hexagon/htp/hvx-dump.h deleted file mode 100644 index 85201fc34532..000000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-dump.h +++ /dev/null @@ -1,129 +0,0 @@ -#ifndef HVX_DUMP_H -#define HVX_DUMP_H - -#include - -#include -#include - -#include "hex-utils.h" -#include "hvx-types.h" - -static void hvx_vec_dump_f16_n(char * pref, HVX_Vector v, uint32_t n) { - HVX_VectorAlias u = { .v = v }; - - const uint32_t n0 = n / 16; - const uint32_t n1 = n % 16; - int i = 0; - for (; i < n0; i++) { - hex_dump_f16_line(pref, u.fp16 + (16 * i), 16); - } - if (n1) { - hex_dump_f16_line(pref, u.fp16 + (16 * i), n1); - } -} - -static void hvx_vec_dump_f16(char * pref, HVX_Vector v) { - hvx_vec_dump_f16_n(pref, v, 64); -} - -static void hvx_vec_dump_f32_n(char * pref, HVX_Vector v, uint32_t n) { - HVX_VectorAlias u = { .v = v }; - - const uint32_t n0 = n / 16; - const uint32_t n1 = n % 16; - int i = 0; - for (; i < n0; i++) { - hex_dump_f32_line(pref, u.fp32 + (16 * i), 16); - } - if (n1) { - hex_dump_f32_line(pref, u.fp32 + (16 * i), n1); - } -} - -static void hvx_vec_dump_f32_hmt(char * pref, HVX_Vector v) { - union { - HVX_Vector v; - float d[32]; - } u = { .v = v }; - - FARF(HIGH, "%s: %.6f %.6f %.6f %.6f ... %.6f %.6f %.6f %.6f ... %.6f %.6f %.6f %.6f\n", pref, u.d[0], u.d[1], - u.d[2], u.d[3], u.d[12], u.d[13], u.d[14], u.d[15], u.d[28], u.d[29], u.d[30], u.d[31]); -} - -static void hvx_vec_dump_f32(char * pref, HVX_Vector v) { - hvx_vec_dump_f32_n(pref, v, 32); -} - -static void hvx_vec_dump_int32(char * pref, HVX_Vector v) { - union { - HVX_Vector v; - int32_t d[32]; - } u = { .v = v }; - - for (int i = 0; i < 32 / 16; i++) { - hex_dump_int32_line(pref, u.d + (16 * i), 16); - } -} - -static void hvx_vec_dump_int32_hmt(char * pref, HVX_Vector v) { - union { - HVX_Vector v; - int32_t d[32]; - } u = { .v = v }; - - FARF(HIGH, "%s: %d %d %d %d ... %d %d %d %d ... %d %d %d %d\n", pref, u.d[0], u.d[1], u.d[2], u.d[3], u.d[12], - u.d[13], u.d[14], u.d[15], u.d[28], u.d[29], u.d[30], u.d[31]); -} - -static void hvx_vec_dump_int8_hmt(char * pref, HVX_Vector v) { - union { - HVX_Vector v; - int8_t d[128]; - } u = { .v = v }; - - FARF(HIGH, "%s: %d %d %d %d ... %d %d %d %d ... %d %d %d %d\n", pref, u.d[0], u.d[1], u.d[2], u.d[3], u.d[60], - u.d[61], u.d[62], u.d[63], u.d[124], u.d[125], u.d[126], u.d[127]); -} - -static void hvx_vec_dump_int8(char * pref, HVX_Vector v) { - union { - HVX_Vector v; - int8_t d[128]; - } u = { .v = v }; - - for (int i = 0; i < 128 / 16; i++) { - hex_dump_int8_line(pref, u.d + (16 * i), 16); - } -} - -static void hvx_vec_dump_uint8(char * pref, HVX_Vector v) { - union { - HVX_Vector v; - uint8_t d[128]; - } u = { .v = v }; - - for (int i = 0; i < 128 / 16; i++) { - hex_dump_uint8_line(pref, u.d + (16 * i), 16); - } -} - -static bool hvx_vec_eq(HVX_Vector v0, HVX_Vector v1, size_t n) { - typedef union { - HVX_Vector v; - int8_t d[128]; - } U; - - U u0 = { .v = v0 }; - U u1 = { .v = v1 }; - - for (int i = 0; i < n; i++) { - if (u0.d[i] != u1.d[i]) { - return false; - } - } - - return true; -} - -#endif /* HVX_DUMP_H */ diff --git a/ggml/src/ggml-hexagon/htp/hvx-exp.h b/ggml/src/ggml-hexagon/htp/hvx-exp.h deleted file mode 100644 index e71ec4909a6f..000000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-exp.h +++ /dev/null @@ -1,216 +0,0 @@ -#ifndef HVX_EXP_H -#define HVX_EXP_H - -#include -#include -#include - -#include "hvx-base.h" -#include "hvx-floor.h" - -#define EXP_COEFF_5 (0x39506967) // 0.000198757 = 1/(7!) -#define EXP_COEFF_4 (0x3AB743CE) // 0.0013982 = 1/(6!) -#define EXP_COEFF_3 (0x3C088908) // 0.00833345 = 1/(5!) -#define EXP_COEFF_2 (0x3D2AA9C1) // 0.416658 = 1/(4!) -#define EXP_COEFF_1 (0x3E2AAAAA) // 0.16666667 = 1/(3!) -#define EXP_COEFF_0 (0x3F000000) // 0.5 = 1/(2!) -#define EXP_LOGN2 (0x3F317218) // ln(2) = 0.6931471805 -#define EXP_LOG2E (0x3FB8AA3B) // log2(e) = 1/ln(2) = 1.4426950408 -#define EXP_ONE (0x3f800000) // 1.0 -#define EXP_RANGE_R (0x42B17218) // ln(FLT_MAX) approx = 88.7228 -#define EXP_RANGE_L (0xC2B00000) // -88.0 (approx log(FLT_MIN)) - -static inline HVX_Vector hvx_vec_exp_f32(HVX_Vector in_vec) { - HVX_Vector z_qf32_v; - HVX_Vector x_v; - HVX_Vector x_qf32_v; - HVX_Vector y_v; - HVX_Vector k_v; - HVX_Vector f_v; - HVX_Vector epsilon_v; - HVX_Vector log2e = Q6_V_vsplat_R(EXP_LOG2E); - HVX_Vector logn2 = Q6_V_vsplat_R(EXP_LOGN2); - HVX_Vector E_const; - HVX_Vector zero_v = Q6_V_vzero(); - - // exp(x) is approximated as follows: - // f = floor(x/ln(2)) = floor(x*log2(e)) - // epsilon = x - f*ln(2) - // exp(x) = exp(epsilon+f*ln(2)) - // = exp(epsilon)*exp(f*ln(2)) - // = exp(epsilon)*2^f - // - // Since epsilon is close to zero, it can be approximated with its Taylor series: - // exp(x) ~= 1+x+x^2/2!+x^3/3!+...+x^n/n!+... - // Preserving the first eight elements, we get: - // exp(x) ~= 1+x+e0*x^2+e1*x^3+e2*x^4+e3*x^5+e4*x^6+e5*x^7 - // = 1+x+(E0+(E1+(E2+(E3+(E4+E5*x)*x)*x)*x)*x)*x^2 - - HVX_Vector temp_v = in_vec; - - // Clamp inputs to (-88.0, 88.0) to avoid overflow/underflow - HVX_VectorPred pred_cap_right = Q6_Q_vcmp_gt_VsfVsf(in_vec, Q6_V_vsplat_R(EXP_RANGE_R)); - HVX_VectorPred pred_cap_left = Q6_Q_vcmp_gt_VsfVsf(Q6_V_vsplat_R(EXP_RANGE_L), in_vec); - - in_vec = Q6_V_vmux_QVV(pred_cap_right, Q6_V_vsplat_R(EXP_RANGE_R), temp_v); - in_vec = Q6_V_vmux_QVV(pred_cap_left, Q6_V_vsplat_R(EXP_RANGE_L), in_vec); - - epsilon_v = Q6_Vqf32_vmpy_VsfVsf(log2e, in_vec); - epsilon_v = Q6_Vsf_equals_Vqf32(epsilon_v); - - // f_v is the floating point result and k_v is the integer result - f_v = hvx_vec_floor_f32(epsilon_v); - k_v = hvx_vec_truncate_f32(f_v); - - x_qf32_v = Q6_Vqf32_vadd_VsfVsf(in_vec, zero_v); - - // x = x - f_v * logn2; - epsilon_v = Q6_Vqf32_vmpy_VsfVsf(f_v, logn2); - x_qf32_v = Q6_Vqf32_vsub_Vqf32Vqf32(x_qf32_v, epsilon_v); - // normalize before every QFloat's vmpy - x_qf32_v = Q6_Vqf32_vadd_Vqf32Vsf(x_qf32_v, zero_v); - - x_v = Q6_Vsf_equals_Vqf32(x_qf32_v); - - // z = x * x; - z_qf32_v = Q6_Vqf32_vmpy_Vqf32Vqf32(x_qf32_v, x_qf32_v); - z_qf32_v = Q6_Vqf32_vadd_Vqf32Vsf(z_qf32_v, zero_v); - - // y = E4 + E5 * x; - E_const = Q6_V_vsplat_R(EXP_COEFF_5); - y_v = Q6_Vqf32_vmpy_VsfVsf(E_const, x_v); - E_const = Q6_V_vsplat_R(EXP_COEFF_4); - y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); - y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); - - // y = E3 + y * x; - E_const = Q6_V_vsplat_R(EXP_COEFF_3); - y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v); - y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); - y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); - - // y = E2 + y * x; - E_const = Q6_V_vsplat_R(EXP_COEFF_2); - y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v); - y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); - y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); - - // y = E1 + y * x; - E_const = Q6_V_vsplat_R(EXP_COEFF_1); - y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v); - y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); - y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); - - // y = E0 + y * x; - E_const = Q6_V_vsplat_R(EXP_COEFF_0); - y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v); - y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); - y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); - - // y = x + y * z; - y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, z_qf32_v); - y_v = Q6_Vqf32_vadd_Vqf32Vqf32(y_v, x_qf32_v); - y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); - - // y = y + 1.0; - y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, Q6_V_vsplat_R(EXP_ONE)); - - // insert exponents - // y = ldexpf(y, k); - // y_v += k_v; // qf32 - // modify exponent - - y_v = Q6_Vsf_equals_Vqf32(y_v); - - // add k_v to the exponent of y_v - HVX_Vector y_v_exponent = Q6_Vw_vasl_VwR(y_v, 1); - - y_v_exponent = Q6_Vuw_vlsr_VuwR(y_v_exponent, IEEE_VSF_MANTLEN + 1); - y_v_exponent = Q6_Vw_vadd_VwVw(k_v, y_v_exponent); - - // exponent cannot be negative; if overflow is detected, result is set to zero - HVX_VectorPred qy_v_negative_exponent = Q6_Q_vcmp_gt_VwVw(zero_v, y_v_exponent); - - y_v = Q6_Vw_vaslacc_VwVwR(y_v, k_v, IEEE_VSF_MANTLEN); - - y_v = Q6_V_vmux_QVV(qy_v_negative_exponent, zero_v, y_v); - - return y_v; -} - -static inline HVX_Vector hvx_vec_exp_f32_guard(HVX_Vector in_vec, HVX_Vector max_exp, HVX_Vector inf) { - const HVX_VectorPred pred0 = Q6_Q_vcmp_gt_VsfVsf(in_vec, max_exp); - - HVX_Vector out = hvx_vec_exp_f32(in_vec); - - return Q6_V_vmux_QVV(pred0, inf, out); -} - -static inline void hvx_exp_f32(uint8_t * restrict dst, const uint8_t * restrict src, const int num_elems, bool negate) { - int left_over = num_elems & (VLEN_FP32 - 1); - int num_elems_whole = num_elems - left_over; - - int unaligned_addr = 0; - int unaligned_loop = 0; - if ((0 == hex_is_aligned((void *) src, VLEN)) || (0 == hex_is_aligned((void *) dst, VLEN))) { - unaligned_addr = 1; - } - // assert((0 == unaligned_addr) || (0 == num_elems_whole)); - if ((1 == unaligned_addr) && (num_elems_whole != 0)) { - unaligned_loop = 1; - } - - HVX_Vector vec_out = Q6_V_vzero(); - - static const float kInf = INFINITY; - static const float kMaxExp = 88.7228f; - - const HVX_Vector max_exp = hvx_vec_splat_f32(kMaxExp); - const HVX_Vector inf = hvx_vec_splat_f32(kInf); - - if (0 == unaligned_loop) { - HVX_Vector * p_vec_in1 = (HVX_Vector *) src; - HVX_Vector * p_vec_out = (HVX_Vector *) dst; - - #pragma unroll(4) - for (int i = 0; i < num_elems_whole; i += VLEN_FP32) { - if (true == negate) { - HVX_Vector neg_vec_in = hvx_vec_neg_f32(*p_vec_in1++); - *p_vec_out++ = hvx_vec_exp_f32_guard(neg_vec_in, max_exp, inf); - } else { - *p_vec_out++ = hvx_vec_exp_f32_guard(*p_vec_in1++, max_exp, inf); - } - } - } else { - #pragma unroll(4) - for (int i = 0; i < num_elems_whole; i += VLEN_FP32) { - HVX_Vector in = *(HVX_UVector *) (src + i * SIZEOF_FP32); - - if (true == negate) { - HVX_Vector neg_vec_in = hvx_vec_neg_f32(in); - *(HVX_UVector *) (dst + i * SIZEOF_FP32) = hvx_vec_exp_f32_guard(neg_vec_in, max_exp, inf); - } else { - *(HVX_UVector *) (dst + i * SIZEOF_FP32) = hvx_vec_exp_f32_guard(in, max_exp, inf); - } - } - } - - if (left_over > 0) { - const float * srcf = (float *) src + num_elems_whole; - float * dstf = (float *) dst + num_elems_whole; - - HVX_Vector in = *(HVX_UVector *) srcf; - - if (true == negate) { - HVX_Vector neg_vec_in = hvx_vec_neg_f32(in); - - vec_out = hvx_vec_exp_f32_guard(neg_vec_in, max_exp, inf); - } else { - vec_out = hvx_vec_exp_f32_guard(in, max_exp, inf); - } - - hvx_vec_store_u((void *) dstf, left_over * SIZEOF_FP32, vec_out); - } -} - -#endif /* HVX_EXP_H */ diff --git a/ggml/src/ggml-hexagon/htp/hvx-floor.h b/ggml/src/ggml-hexagon/htp/hvx-floor.h deleted file mode 100644 index 6a1bfde56757..000000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-floor.h +++ /dev/null @@ -1,100 +0,0 @@ -#ifndef HVX_FLOOR_H -#define HVX_FLOOR_H - -#include -#include - -#include "hvx-base.h" - -#define IEEE_VSF_EXPLEN (8) -#define IEEE_VSF_EXPBIAS (127) -#define IEEE_VSF_EXPMASK (0xFF) -#define IEEE_VSF_MANTLEN (23) -#define IEEE_VSF_MANTMASK (0x7FFFFF) -#define IEEE_VSF_MIMPMASK (0x800000) - -static inline HVX_Vector hvx_vec_truncate_f32(HVX_Vector in_vec) { - HVX_Vector mask_mant_v = Q6_V_vsplat_R(IEEE_VSF_MANTMASK); - HVX_Vector mask_impl_v = Q6_V_vsplat_R(IEEE_VSF_MIMPMASK); - HVX_Vector const_zero_v = Q6_V_vzero(); - - HVX_VectorPred q_negative = Q6_Q_vcmp_gt_VwVw(const_zero_v, in_vec); - - HVX_Vector expval_v = in_vec >> IEEE_VSF_MANTLEN; - expval_v &= IEEE_VSF_EXPMASK; - expval_v -= IEEE_VSF_EXPBIAS; - - // negative exp == fractional value - HVX_VectorPred q_negexp = Q6_Q_vcmp_gt_VwVw(const_zero_v, expval_v); - - HVX_Vector rshift_v = IEEE_VSF_MANTLEN - expval_v; // fractional bits - exp shift - - HVX_Vector mant_v = in_vec & mask_mant_v; // obtain mantissa - HVX_Vector vout = Q6_Vw_vadd_VwVw(mant_v, mask_impl_v); // add implicit 1.0 - - vout = Q6_Vw_vasr_VwVw(vout, rshift_v); // shift to obtain truncated integer - vout = Q6_V_vmux_QVV(q_negexp, const_zero_v, vout); // expval<0 -> 0 - - HVX_Vector neg_vout = -vout; - - vout = Q6_V_vmux_QVV(q_negative, neg_vout, vout); // handle negatives - - return (vout); -} - -static inline HVX_Vector hvx_vec_floor_f32(HVX_Vector in_vec) { - HVX_Vector mask_mant_v = Q6_V_vsplat_R(IEEE_VSF_MANTMASK); - HVX_Vector mask_impl_v = Q6_V_vsplat_R(IEEE_VSF_MIMPMASK); - HVX_Vector const_mnlen_v = Q6_V_vsplat_R(IEEE_VSF_MANTLEN); - HVX_Vector const_zero_v = Q6_V_vzero(); - HVX_Vector const_negone_v = Q6_V_vsplat_R(0xbf800000); // -1 IEEE vsf - - HVX_VectorPred q_negative = Q6_Q_vcmp_gt_VwVw(const_zero_v, in_vec); - - HVX_Vector expval_v = in_vec >> IEEE_VSF_MANTLEN; - expval_v &= IEEE_VSF_EXPMASK; - expval_v -= IEEE_VSF_EXPBIAS; - - HVX_VectorPred q_negexp = Q6_Q_vcmp_gt_VwVw(const_zero_v, expval_v); - HVX_VectorPred q_expltmn = Q6_Q_vcmp_gt_VwVw(const_mnlen_v, expval_v); - HVX_VectorPred q_negexp_pos = Q6_Q_vcmp_gtand_QVwVw(q_negexp, in_vec, const_zero_v); - HVX_VectorPred q_negexp_neg = Q6_Q_vcmp_gtand_QVwVw(q_negexp, const_zero_v, in_vec); - - // if expval < 0 (q_negexp) // <0, floor is 0 - // if vin > 0 - // floor = 0 - // if vin < 0 - // floor = -1 - // if expval < mant_len (q_expltmn) // >0, but fraction may exist - // get sign (q_negative) - // mask >> expval // fraction bits to mask off - // vout = ~(mask) // apply mask to remove fraction - // if (qneg) // negative floor is one less (more, sign bit for neg) - // vout += ((impl_mask) >> expval) - // if (mask && vin) - // vout = vin - // else // already an integer - // ; // no change - - // compute floor - mask_mant_v >>= expval_v; - HVX_Vector neg_addin_v = mask_impl_v >> expval_v; - HVX_Vector vout_neg_addin = Q6_Vw_vadd_VwVw(in_vec, neg_addin_v); - HVX_Vector vout = Q6_V_vmux_QVV(q_negative, vout_neg_addin, in_vec); - - HVX_Vector mask_chk_v = Q6_V_vand_VV(in_vec, mask_mant_v); // chk if bits set - HVX_VectorPred q_integral = Q6_Q_vcmp_eq_VwVw(const_zero_v, mask_chk_v); - - HVX_Vector not_mask_v = Q6_V_vnot_V(mask_mant_v); // frac bits to clear - HVX_Vector vfrfloor_v = Q6_V_vand_VV(vout, not_mask_v); // clear frac bits - - vout = in_vec; - vout = Q6_V_vmux_QVV(q_expltmn, vfrfloor_v, vout); // expval0 -> 0 - vout = Q6_V_vmux_QVV(q_negexp_neg, const_negone_v, vout); // expval<0 x<0 -> -1 - - return vout; -} - -#endif /* HVX_FLOOR_H */ diff --git a/ggml/src/ggml-hexagon/htp/hvx-inverse.h b/ggml/src/ggml-hexagon/htp/hvx-inverse.h deleted file mode 100644 index f2054f45baca..000000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-inverse.h +++ /dev/null @@ -1,210 +0,0 @@ -#ifndef HVX_INVERSE_H -#define HVX_INVERSE_H - -#include - -#include -#include -#include -#include -#include - -#include "hvx-base.h" - -// ==================================================== -// FUNCTION: 1/(x+1) y(0) = 1, y(0.5) = 0.6667, y(1) = 0.5 -// Order:3; continuity: True; Ends forced: True -// Mode: unsigned; Result fractional bits: 14 -// Peak Error: 1.1295e-04 Rms Error: 2.8410e-05 Mean Error: 1.1370e-05 -// 32769 -32706 31252 -10589 -// 32590 -30635 22793 -4493 -// 32066 -27505 16481 -2348 -// 31205 -24054 11849 -1306 - -static inline HVX_Vector hvx_vec_recip_xp1_O3_unsigned(HVX_Vector vx) { - // input is 0..0xffff representing 0.0 .. 1.0 - HVX_Vector p; - p = Q6_Vh_vlut4_VuhPh(vx, 0xFAE6F6D4EE73D6A3ull); - p = Q6_Vh_vmpa_VhVhVuhPuh_sat(p, vx, 0x2E49406159097A14ull); - p = Q6_Vh_vmps_VhVhVuhPuh_sat(p, vx, 0x5DF66B7177AB7FC2ull); - p = Q6_Vh_vmpa_VhVhVuhPuh_sat(p, vx, 0x79E57D427F4E8001ull); - return p; // signed result, 14 fractional bits -} - -// Find reciprocal of fp16. -// (1) first, convert to fp32, multiplying by 1.0; this is done to -// handle denormals. Ignoring sign and zero, result should be at -// least 5.9604645e-08 (32-bit code 0x33800000) and at most 131008 (0x47ffe000) -// (exponent in range [103,143]) -// (2) extract the mantissa into 16-bit unsigned; find reciprocal using a fitted poly -// (3) put this, along with '253-exp' (exp from (1)) together to make an qf32 -// (4) convert that to fp16 -// (5) put sign back in. Also, if the original value (w/o sign) was <0x81, replace -// the result with the max value. -static inline HVX_Vector hvx_vec_inverse_f16(HVX_Vector vals) { - HVX_Vector em_mask = Q6_Vh_vsplat_R(0x7FFF); - HVX_Vector avals = Q6_V_vand_VV(vals, em_mask); - HVX_VectorPred is_neg = Q6_Q_vcmp_gt_VhVh(avals, vals); - // is too small to 1/x ? for 'standard' fp16, this would be 0x101 - HVX_VectorPred is_small = Q6_Q_vcmp_gt_VhVh(Q6_Vh_vsplat_R(0x101), avals); - - HVX_VectorPair to_qf32 = Q6_Wqf32_vmpy_VhfVhf(avals, Q6_Vh_vsplat_R(0x3C00)); // *1.0 - HVX_Vector to_f32_0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(to_qf32)); - HVX_Vector to_f32_1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(to_qf32)); - - // bits 22..13 contain the mantissa now (w/o hidden bit); move to bit 14..5 of a 16-bit vector - HVX_Vector mant_u16 = Q6_Vh_vshuffo_VhVh(Q6_Vw_vasl_VwR(to_f32_1, 9), Q6_Vw_vasl_VwR(to_f32_0, 9)); - // likewise extract the upper 16 from each, containing the exponents in range 103..142 - HVX_Vector exp_u16 = Q6_Vh_vshuffo_VhVh(to_f32_1, to_f32_0); - //Get exponent in IEEE 32-bit representation - exp_u16 = Q6_Vuh_vlsr_VuhR(exp_u16, 7); - - // so, mant_u16 contains an unbiased mantissa in upper 10 bits of each u16 lane - // We can consider it to be x-1.0, with 16 fractional bits, where 'x' is in range [1.0,2.0) - // Use poly to transform to 1/x, with 14 fractional bits - // - HVX_Vector rm = hvx_vec_recip_xp1_O3_unsigned(mant_u16); - - HVX_Vector vcl0 = Q6_Vuh_vcl0_Vuh(rm); //count leading zeros - - // Get mantissa for 16-bit representation - HVX_Vector mant_recip = Q6_V_vand_VV(Q6_Vh_vasr_VhR(Q6_Vh_vasl_VhVh(rm, vcl0), 5), Q6_Vh_vsplat_R(0x03FF)); - - //Compute Reciprocal Exponent - HVX_Vector exp_recip = - Q6_Vh_vsub_VhVh(Q6_Vh_vsub_VhVh(Q6_Vh_vsplat_R(254), exp_u16), Q6_Vh_vsub_VhVh(vcl0, Q6_Vh_vsplat_R(1))); - //Convert it for 16-bit representation - exp_recip = Q6_Vh_vadd_VhVh_sat(Q6_Vh_vsub_VhVh(exp_recip, Q6_Vh_vsplat_R(127)), Q6_Vh_vsplat_R(15)); - exp_recip = Q6_Vh_vasl_VhR(exp_recip, 10); - - //Merge exponent and mantissa for reciprocal - HVX_Vector recip = Q6_V_vor_VV(exp_recip, mant_recip); - // map 'small' inputs to standard largest value 0x7bff - recip = Q6_V_vmux_QVV(is_small, Q6_Vh_vsplat_R(0x7bff), recip); - // add sign back - recip = Q6_V_vandor_VQR(recip, is_neg, 0x80008000); - return recip; -} - -static inline HVX_Vector hvx_vec_inverse_f32(HVX_Vector v_sf) { - HVX_Vector inv_aprox_sf = Q6_V_vsplat_R(0x7EEEEBB3); - HVX_Vector two_sf = hvx_vec_splat_f32(2.0); - - // First approximation - HVX_Vector i_sf = Q6_Vw_vsub_VwVw(inv_aprox_sf, v_sf); - - HVX_Vector r_qf; - - // Refine - r_qf = Q6_Vqf32_vmpy_VsfVsf( - i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v_sf))))); - r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32( - r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v_sf)))); - r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32( - r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v_sf)))); - - return Q6_Vsf_equals_Vqf32(r_qf); -} - -static inline HVX_Vector hvx_vec_inverse_f32_guard(HVX_Vector v_sf, HVX_Vector nan_inf_mask) { - HVX_Vector out = hvx_vec_inverse_f32(v_sf); - - HVX_Vector masked_out = Q6_V_vand_VV(out, nan_inf_mask); - const HVX_VectorPred pred = Q6_Q_vcmp_eq_VwVw(nan_inf_mask, masked_out); - - return Q6_V_vmux_QVV(pred, Q6_V_vzero(), out); -} - -#define hvx_inverse_f32_loop_body(dst_type, src_type, vec_store) \ - do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src_type * restrict vsrc = (src_type *) src; \ - \ - const HVX_Vector nan_inf_mask = Q6_V_vsplat_R(0x7f800000); \ - \ - const uint32_t nvec = n / VLEN_FP32; \ - const uint32_t nloe = n % VLEN_FP32; \ - \ - uint32_t i = 0; \ - \ - _Pragma("unroll(4)") \ - for (; i < nvec; i++) { \ - vdst[i] = hvx_vec_inverse_f32_guard(vsrc[i], nan_inf_mask); \ - } \ - if (nloe) { \ - HVX_Vector v = hvx_vec_inverse_f32_guard(vsrc[i], nan_inf_mask); \ - vec_store((void *) &vdst[i], nloe * SIZEOF_FP32, v); \ - } \ - } while(0) - -static inline HVX_Vector hvx_vec_inverse_f16_guard(HVX_Vector v_sf, HVX_Vector nan_inf_mask) { - HVX_Vector out = hvx_vec_inverse_f16(v_sf); - - HVX_Vector masked_out = Q6_V_vand_VV(out, nan_inf_mask); - const HVX_VectorPred pred = Q6_Q_vcmp_eq_VhVh(nan_inf_mask, masked_out); - - return Q6_V_vmux_QVV(pred, Q6_V_vzero(), out); -} - -#define hvx_inverse_f16_loop_body(dst_type, src_type, vec_store) \ - do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src_type * restrict vsrc = (src_type *) src; \ - \ - const HVX_Vector nan_inf_mask = Q6_Vh_vsplat_R(0x7c00); \ - \ - const uint32_t nvec = n / VLEN_FP16; \ - const uint32_t nloe = n % VLEN_FP16; \ - \ - uint32_t i = 0; \ - \ - _Pragma("unroll(4)") \ - for (; i < nvec; i++) { \ - vdst[i] = hvx_vec_inverse_f16_guard(vsrc[i], nan_inf_mask); \ - } \ - if (nloe) { \ - HVX_Vector v = hvx_vec_inverse_f16_guard(vsrc[i], nan_inf_mask); \ - vec_store((void *) &vdst[i], nloe * SIZEOF_FP16, v); \ - } \ - } while(0) - -// Generic macro to define alignment permutations for an op -#define DEFINE_HVX_INV_OP_VARIANTS(OP_NAME, OP_LOOP_BODY) \ -static inline void OP_NAME##_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \ - assert((uintptr_t) dst % 128 == 0); \ - assert((uintptr_t) src % 128 == 0); \ - OP_LOOP_BODY(HVX_Vector, HVX_Vector, hvx_vec_store_a); \ -} \ -static inline void OP_NAME##_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \ - assert((uintptr_t) dst % 128 == 0); \ - OP_LOOP_BODY(HVX_Vector, HVX_UVector, hvx_vec_store_a); \ -} \ -static inline void OP_NAME##_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \ - assert((uintptr_t) src % 128 == 0); \ - OP_LOOP_BODY(HVX_UVector, HVX_Vector, hvx_vec_store_u); \ -} \ -static inline void OP_NAME##_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \ - OP_LOOP_BODY(HVX_UVector, HVX_UVector, hvx_vec_store_u); \ -} \ - -// Dispatcher logic -#define HVX_INV_DISPATCHER(OP_NAME) \ -static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t num_elems) { \ - if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { \ - OP_NAME##_aa(dst, src, num_elems); \ - } else if (hex_is_aligned((void *) dst, 128)) { \ - OP_NAME##_au(dst, src, num_elems); \ - } else if (hex_is_aligned((void *) src, 128)) { \ - OP_NAME##_ua(dst, src, num_elems); \ - } else { \ - OP_NAME##_uu(dst, src, num_elems); \ - } \ -} - -DEFINE_HVX_INV_OP_VARIANTS(hvx_inverse_f32, hvx_inverse_f32_loop_body) -DEFINE_HVX_INV_OP_VARIANTS(hvx_inverse_f16, hvx_inverse_f16_loop_body) - -HVX_INV_DISPATCHER(hvx_inverse_f32) -HVX_INV_DISPATCHER(hvx_inverse_f16) - -#endif // HVX_INVERSE_H diff --git a/ggml/src/ggml-hexagon/htp/hvx-reduce.h b/ggml/src/ggml-hexagon/htp/hvx-reduce.h deleted file mode 100644 index 3c0073ef6d80..000000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-reduce.h +++ /dev/null @@ -1,296 +0,0 @@ -#ifndef HVX_REDUCE_H -#define HVX_REDUCE_H - -#include -#include -#include -#include - -#include "hex-utils.h" -#include "hvx-base.h" -#include "hvx-types.h" - -static inline HVX_Vector hvx_vec_reduce_sum_n_i32(HVX_Vector in, unsigned int n) { - unsigned int total = n * 4; // total vec nbytes - unsigned int width = 4; // int32 - - HVX_Vector sum = in, sum_t; - while (width < total) { - sum_t = Q6_V_vror_VR(sum, width); // rotate right - sum = Q6_Vw_vadd_VwVw(sum_t, sum); // elementwise sum - width = width << 1; - } - return sum; -} - -static inline HVX_Vector hvx_vec_reduce_sum_i32(HVX_Vector in) { - return hvx_vec_reduce_sum_n_i32(in, 32); -} - -static inline HVX_Vector hvx_vec_reduce_sum_n_qf32(HVX_Vector in, unsigned int n) { - unsigned int total = n * 4; // total vec nbytes - unsigned int width = 4; // fp32 nbytes - - HVX_Vector sum = in, sum_t; - while (width < total) { - sum_t = Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum), width); // rotate right - sum = Q6_Vqf32_vadd_Vqf32Vsf(sum, sum_t); // elementwise sum - width = width << 1; - } - return sum; -} - -static inline HVX_Vector hvx_vec_reduce_sum_qf32(HVX_Vector in) { - return hvx_vec_reduce_sum_n_qf32(in, 32); -} - -#if __HVX_ARCH__ > 75 - -static inline HVX_Vector hvx_vec_reduce_sum_f32x4(HVX_Vector_x4 in) { - HVX_VectorPair sum_p01 = Q6_W_vshuff_VVR(in.v[1], in.v[0], 4); - HVX_VectorPair sum_p23 = Q6_W_vshuff_VVR(in.v[3], in.v[2], 4); - HVX_Vector sum_sf01 = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p01), Q6_V_hi_W(sum_p01)); - HVX_Vector sum_sf23 = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p23), Q6_V_hi_W(sum_p23)); - - HVX_VectorPair sum_p0123 = Q6_W_vshuff_VVR(sum_sf23, sum_sf01, 8); - HVX_Vector sum_sf = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p0123), Q6_V_hi_W(sum_p0123)); - - sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 2)); - sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 4)); - sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 8)); - return sum_sf; -} - -static inline HVX_Vector hvx_vec_reduce_sum_f32x2(HVX_Vector in0, HVX_Vector in1) { - HVX_VectorPair sump = Q6_W_vshuff_VVR(in1, in0, 4); - HVX_Vector sum_sf = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sump), Q6_V_hi_W(sump)); - - sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 2)); - sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 4)); - sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 8)); - sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 16)); - return sum_sf; -} - -static inline HVX_Vector hvx_vec_reduce_sum_n_f32(HVX_Vector in, unsigned int n) { - unsigned int total = n * 4; // total vec nbytes - unsigned int width = 4; // fp32 nbytes - - HVX_Vector sum = in, sum_t; - while (width < total) { - sum_t = Q6_V_vror_VR(sum, width); // rotate right - sum = Q6_Vsf_vadd_VsfVsf(sum, sum_t); // elementwise sum - width = width << 1; - } - return sum; -} - -#else - -static inline HVX_Vector hvx_vec_reduce_sum_f32x4(HVX_Vector_x4 in) { - HVX_VectorPair sum_p01 = Q6_W_vshuff_VVR(in.v[1], in.v[0], 4); - HVX_VectorPair sum_p23 = Q6_W_vshuff_VVR(in.v[3], in.v[2], 4); - HVX_Vector sum_qf01 = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p01), Q6_V_hi_W(sum_p01)); - HVX_Vector sum_qf23 = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p23), Q6_V_hi_W(sum_p23)); - - HVX_VectorPair sum_p0123 = Q6_W_vshuff_VVR(Q6_Vsf_equals_Vqf32(sum_qf23), Q6_Vsf_equals_Vqf32(sum_qf01), 8); - HVX_Vector sum_qf = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p0123), Q6_V_hi_W(sum_p0123)); - - sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 2)); - sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 4)); - sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 8)); - return Q6_Vsf_equals_Vqf32(sum_qf); -} - -static inline HVX_Vector hvx_vec_reduce_sum_f32x2(HVX_Vector in0, HVX_Vector in1) { - HVX_VectorPair sump = Q6_W_vshuff_VVR(in1, in0, 4); - HVX_Vector sum_qf = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sump), Q6_V_hi_W(sump)); - - sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 2)); - sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 4)); - sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 8)); - sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 16)); - return Q6_Vsf_equals_Vqf32(sum_qf); -} - -static inline HVX_Vector hvx_vec_reduce_sum_n_f32(HVX_Vector in, unsigned int n) { - unsigned int total = n * 4; // total vec nbytes - unsigned int width = 4; // fp32 nbytes - - HVX_Vector sum = in, sum_t; - while (width < total) { - sum_t = Q6_V_vror_VR(sum, width); // rotate right - sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(sum, sum_t)); // elementwise sum - width = width << 1; - } - return sum; -} - -#endif - -static inline HVX_Vector hvx_vec_reduce_sum_f32(HVX_Vector in) { - return hvx_vec_reduce_sum_n_f32(in, 32); -} - -static inline HVX_Vector hvx_vec_reduce_max_f16(HVX_Vector in) { - unsigned total = 128; // total vec nbytes - unsigned width = 2; // fp16 nbytes - - HVX_Vector _max = in, _max_t; - while (width < total) { - _max_t = Q6_V_vror_VR(_max, width); // rotate right - _max = Q6_Vhf_vmax_VhfVhf(_max_t, _max); // elementwise max - width = width << 1; - } - - return _max; -} - -static inline HVX_Vector hvx_vec_reduce_max2_f16(HVX_Vector in, HVX_Vector _max) { - unsigned total = 128; // total vec nbytes - unsigned width = 2; // fp32 nbytes - - HVX_Vector _max_t; - - _max = Q6_Vhf_vmax_VhfVhf(in, _max); - while (width < total) { - _max_t = Q6_V_vror_VR(_max, width); // rotate right - _max = Q6_Vhf_vmax_VhfVhf(_max_t, _max); // elementwise max - width = width << 1; - } - - return _max; -} - -static inline HVX_Vector hvx_vec_reduce_max_f32(HVX_Vector in) { - unsigned total = 128; // total vec nbytes - unsigned width = 4; // fp32 nbytes - - HVX_Vector _max = in, _max_t; - while (width < total) { - _max_t = Q6_V_vror_VR(_max, width); // rotate right - _max = Q6_Vsf_vmax_VsfVsf(_max_t, _max); // elementwise max - width = width << 1; - } - - return _max; -} - -static inline HVX_Vector hvx_vec_reduce_max2_f32(HVX_Vector in, HVX_Vector _max) { - unsigned total = 128; // total vec nbytes - unsigned width = 4; // fp32 nbytes - - HVX_Vector _max_t; - - _max = Q6_Vsf_vmax_VsfVsf(in, _max); - while (width < total) { - _max_t = Q6_V_vror_VR(_max, width); // rotate right - _max = Q6_Vsf_vmax_VsfVsf(_max_t, _max); // elementwise max - width = width << 1; - } - - return _max; -} - -#define hvx_reduce_loop_body(src_type, init_vec, pad_vec, vec_op, reduce_op, scalar_reduce) \ - do { \ - src_type * restrict vsrc = (src_type *) src; \ - HVX_Vector acc = init_vec; \ - \ - const uint32_t elem_size = sizeof(float); \ - const uint32_t epv = 128 / elem_size; \ - const uint32_t nvec = num_elems / epv; \ - const uint32_t nloe = num_elems % epv; \ - \ - uint32_t i = 0; \ - _Pragma("unroll(4)") \ - for (; i < nvec; i++) { \ - acc = vec_op(acc, vsrc[i]); \ - } \ - if (nloe) { \ - const float * srcf = (const float *) src + i * epv; \ - HVX_Vector in = *(HVX_UVector *) srcf; \ - HVX_Vector temp = Q6_V_valign_VVR(in, pad_vec, nloe * elem_size); \ - acc = vec_op(acc, temp); \ - } \ - HVX_Vector v = reduce_op(acc); \ - return scalar_reduce(v); \ - } while(0) - -#define HVX_REDUCE_MAX_OP(acc, val) Q6_Vsf_vmax_VsfVsf(acc, val) -#define HVX_REDUCE_SUM_OP(acc, val) Q6_Vqf32_vadd_VsfVsf(Q6_Vsf_equals_Vqf32(acc), val) -#define HVX_SUM_SQ_OP(acc, val) Q6_Vqf32_vadd_Vqf32Vqf32(acc, Q6_Vqf32_vmpy_VsfVsf(val, val)) -#define HVX_REDUCE_MAX_SCALAR(v) hvx_vec_get_f32(v) -#define HVX_REDUCE_SUM_SCALAR(v) hvx_vec_get_f32(Q6_Vsf_equals_Vqf32(v)) - -// Max variants - -static inline float hvx_reduce_max_f32_a(const uint8_t * restrict src, const int num_elems) { - HVX_Vector init_vec = hvx_vec_splat_f32(((const float *) src)[0]); - assert((unsigned long) src % 128 == 0); - hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_REDUCE_MAX_OP, hvx_vec_reduce_max_f32, HVX_REDUCE_MAX_SCALAR); -} - -static inline float hvx_reduce_max_f32_u(const uint8_t * restrict src, const int num_elems) { - HVX_Vector init_vec = hvx_vec_splat_f32(((const float *) src)[0]); - hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_REDUCE_MAX_OP, hvx_vec_reduce_max_f32, HVX_REDUCE_MAX_SCALAR); -} - -static inline float hvx_reduce_max_f32(const uint8_t * restrict src, const int num_elems) { - if (hex_is_aligned((void *) src, 128)) { - return hvx_reduce_max_f32_a(src, num_elems); - } else { - return hvx_reduce_max_f32_u(src, num_elems); - } -} - -// Sum variants - -static inline float hvx_reduce_sum_f32_a(const uint8_t * restrict src, const int num_elems) { - HVX_Vector init_vec = Q6_V_vsplat_R(0); - assert((unsigned long) src % 128 == 0); - hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_REDUCE_SUM_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR); -} - -static inline float hvx_reduce_sum_f32_u(const uint8_t * restrict src, const int num_elems) { - HVX_Vector init_vec = Q6_V_vsplat_R(0); - hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_REDUCE_SUM_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR); -} - -static inline float hvx_reduce_sum_f32(const uint8_t * restrict src, const int num_elems) { - if (hex_is_aligned((void *) src, 128)) { - return hvx_reduce_sum_f32_a(src, num_elems); - } else { - return hvx_reduce_sum_f32_u(src, num_elems); - } -} - -// Sum of squares variants - -static inline float hvx_sum_of_squares_f32_a(const uint8_t * restrict src, const int num_elems) { - HVX_Vector init_vec = Q6_V_vsplat_R(0); - assert((uintptr_t) src % 128 == 0); - hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_SUM_SQ_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR); -} - -static inline float hvx_sum_of_squares_f32_u(const uint8_t * restrict src, const int num_elems) { - HVX_Vector init_vec = Q6_V_vsplat_R(0); - hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_SUM_SQ_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR); -} - -static inline float hvx_sum_of_squares_f32(const uint8_t * restrict src, const int num_elems) { - if (hex_is_aligned((void *) src, 128)) { - return hvx_sum_of_squares_f32_a(src, num_elems); - } else { - return hvx_sum_of_squares_f32_u(src, num_elems); - } -} - -#undef hvx_reduce_loop_body -#undef HVX_REDUCE_MAX_OP -#undef HVX_REDUCE_SUM_OP -#undef HVX_REDUCE_MAX_SCALAR -#undef HVX_REDUCE_SUM_SCALAR -#undef HVX_SUM_SQ_OP - -#endif /* HVX_REDUCE_H */ diff --git a/ggml/src/ggml-hexagon/htp/hvx-scale.h b/ggml/src/ggml-hexagon/htp/hvx-scale.h deleted file mode 100644 index c65c98639dc0..000000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-scale.h +++ /dev/null @@ -1,133 +0,0 @@ -#ifndef HVX_SCALE_H -#define HVX_SCALE_H - -#include -#include -#include - -#include "hvx-base.h" - -#define hvx_scale_f32_loop_body(dst_type, src_type, vec_store) \ - do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src_type * restrict vsrc = (src_type *) src; \ - \ - HVX_Vector vs = hvx_vec_splat_f32(scale); \ - \ - const uint32_t elem_size = sizeof(float); \ - const uint32_t epv = 128 / elem_size; \ - const uint32_t nvec = n / epv; \ - const uint32_t nloe = n % epv; \ - \ - uint32_t i = 0; \ - \ - _Pragma("unroll(4)") \ - for (; i < nvec; ++i) { \ - HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(vsrc[i], vs); \ - vdst[i] = Q6_Vsf_equals_Vqf32(v); \ - } \ - if (nloe) { \ - HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(vsrc[i], vs); \ - vec_store((void *) &vdst[i], nloe * elem_size, Q6_Vsf_equals_Vqf32(v)); \ - } \ - } while(0) - -static inline void hvx_scale_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale) { - assert((size_t) dst % 128 == 0); - assert((size_t) src % 128 == 0); - hvx_scale_f32_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); -} - -static inline void hvx_scale_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale) { - assert((size_t) dst % 128 == 0); - hvx_scale_f32_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); -} - -static inline void hvx_scale_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale) { - assert((size_t) src % 128 == 0); - hvx_scale_f32_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); -} - -static inline void hvx_scale_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale) { - hvx_scale_f32_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); -} - -static inline void hvx_scale_f32(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale) { - if (((size_t) dst & 127) == 0) { - if (((size_t) src & 127) == 0) { - hvx_scale_f32_aa(dst, src, n, scale); - } else { - hvx_scale_f32_au(dst, src, n, scale); - } - } else { - if (((size_t) src & 127) == 0) { - hvx_scale_f32_ua(dst, src, n, scale); - } else { - hvx_scale_f32_uu(dst, src, n, scale); - } - } -} - -#define hvx_scale_offset_f32_loop_body(dst_type, src_type, vec_store) \ - do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src_type * restrict vsrc = (src_type *) src; \ - \ - HVX_Vector vs = hvx_vec_splat_f32(scale); \ - HVX_Vector vo = hvx_vec_splat_f32(offset); \ - \ - const uint32_t elem_size = sizeof(float); \ - const uint32_t epv = 128 / elem_size; \ - const uint32_t nvec = n / epv; \ - const uint32_t nloe = n % epv; \ - \ - uint32_t i = 0; \ - \ - _Pragma("unroll(4)") \ - for (; i < nvec; ++i) { \ - HVX_Vector v = Q6_Vqf32_vadd_Vqf32Vsf(Q6_Vqf32_vmpy_VsfVsf(vsrc[i], vs), vo); \ - vdst[i] = Q6_Vsf_equals_Vqf32(v); \ - } \ - if (nloe) { \ - HVX_Vector v = Q6_Vqf32_vadd_Vqf32Vsf(Q6_Vqf32_vmpy_VsfVsf(vsrc[i], vs), vo); \ - vec_store((void *) &vdst[i], nloe * elem_size, Q6_Vsf_equals_Vqf32(v)); \ - } \ - } while(0) - -static inline void hvx_scale_offset_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) { - assert((size_t) dst % 128 == 0); - assert((size_t) src % 128 == 0); - hvx_scale_offset_f32_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); -} - -static inline void hvx_scale_offset_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) { - assert((size_t) dst % 128 == 0); - hvx_scale_offset_f32_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); -} - -static inline void hvx_scale_offset_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) { - assert((size_t) src % 128 == 0); - hvx_scale_offset_f32_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); -} - -static inline void hvx_scale_offset_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) { - hvx_scale_offset_f32_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); -} - -static inline void hvx_scale_offset_f32(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) { - if (((size_t) dst & 127) == 0) { - if (((size_t) src & 127) == 0) { - hvx_scale_offset_f32_aa(dst, src, n, scale, offset); - } else { - hvx_scale_offset_f32_au(dst, src, n, scale, offset); - } - } else { - if (((size_t) src & 127) == 0) { - hvx_scale_offset_f32_ua(dst, src, n, scale, offset); - } else { - hvx_scale_offset_f32_uu(dst, src, n, scale, offset); - } - } -} - -#endif // HVX_SCALE_H diff --git a/ggml/src/ggml-hexagon/htp/hvx-sigmoid.h b/ggml/src/ggml-hexagon/htp/hvx-sigmoid.h deleted file mode 100644 index 37f3e7b6faec..000000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-sigmoid.h +++ /dev/null @@ -1,142 +0,0 @@ -#ifndef HVX_SIGMOID_H -#define HVX_SIGMOID_H - -#include "hvx-base.h" -#include "hvx-inverse.h" - -#define FAST_SIGMOID_LOG2F (0x3fb8aa3b) // 1.442695022 -#define FAST_SIGMOID_C1 (0x3d009076) // 0.03138777 -#define FAST_SIGMOID_C2 (0x3e8d74bd) // 0.276281267 -#define FAST_SIGMOID_C3 (0x3f000000) // 0.5 - -static inline HVX_Vector hvx_vec_fast_sigmoid_f32(HVX_Vector v) { - v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F)); - v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3)); - - HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v)); - HVX_Vector x = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int)); - HVX_Vector xx = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x); - - HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2)); - v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F)); - - HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1)); - v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx); - v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x); - - HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1)); - HVX_Vector v3_exponent = Q6_Vw_vasl_VwR(v3, 1); - v3_exponent = Q6_Vuw_vlsr_VuwR(v3_exponent, 24); - v3_exponent = Q6_Vw_vadd_VwVw(in_int, v3_exponent); - v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24); - - HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1)); - HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4)); - - HVX_Vector res = hvx_vec_inverse_f32(v5); - res = Q6_Vqf32_vmpy_VsfVsf(v3, res); - - return Q6_Vsf_equals_Vqf32(res); -} - -static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard(HVX_Vector v, - HVX_Vector one, - HVX_Vector max_exp, - HVX_Vector min_exp) { - const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v); - const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp); - - HVX_Vector out = hvx_vec_fast_sigmoid_f32(v); - out = Q6_V_vmux_QVV(pred_max, out, one); - return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero()); -} - -static inline HVX_Vector hvx_vec_tanh_f32(HVX_Vector x) { - // tanh(x) = 2 * sigmoid(2x) - 1 - HVX_Vector two = hvx_vec_splat_f32(2.0f); - HVX_Vector one = hvx_vec_splat_f32(1.0f); - HVX_Vector x2 = Q6_Vqf32_vmpy_VsfVsf(x, two); - - HVX_Vector max_exp = hvx_vec_splat_f32(87.f); - HVX_Vector min_exp = hvx_vec_splat_f32(-87.f); - - HVX_Vector sig2x = hvx_vec_fast_sigmoid_f32_guard(Q6_Vsf_equals_Vqf32(x2), one, max_exp, min_exp); - - HVX_Vector res = Q6_Vqf32_vmpy_VsfVsf(sig2x, two); - res = Q6_Vqf32_vsub_Vqf32Vsf(res, one); - return Q6_Vsf_equals_Vqf32(res); -} - -#define hvx_sigmoid_loop_body(dst_type, src_type, vec_store) \ - do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src_type * restrict vsrc = (src_type *) src; \ - \ - const HVX_Vector one = hvx_vec_splat_f32(1.f); \ - const HVX_Vector max_exp = hvx_vec_splat_f32(87.f); \ - const HVX_Vector min_exp = hvx_vec_splat_f32(-87.f); \ - \ - const uint32_t epv = 128 / sizeof(float); \ - const uint32_t nvec = n / epv; \ - const uint32_t nloe = n % epv; \ - \ - uint32_t i = 0; \ - \ - _Pragma("unroll(4)") \ - for (; i < nvec; i++) { \ - vdst[i] = hvx_vec_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp); \ - } \ - if (nloe) { \ - HVX_Vector tmp = hvx_vec_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp); \ - vec_store((void *) &vdst[i], nloe * sizeof(float), tmp); \ - } \ - } while(0) - -#define hvx_tanh_loop_body(dst_type, src_type, vec_store) \ - do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src_type * restrict vsrc = (src_type *) src; \ - \ - const uint32_t epv = 128 / sizeof(float); \ - const uint32_t nvec = n / epv; \ - const uint32_t nloe = n % epv; \ - \ - uint32_t i = 0; \ - \ - _Pragma("unroll(4)") \ - for (; i < nvec; i++) { \ - vdst[i] = hvx_vec_tanh_f32(vsrc[i]); \ - } \ - if (nloe) { \ - HVX_Vector tmp = hvx_vec_tanh_f32(vsrc[i]); \ - vec_store((void *) &vdst[i], nloe * sizeof(float), tmp); \ - } \ - } while(0) - -static inline void hvx_sigmoid_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - assert((unsigned long) dst % 128 == 0); - assert((unsigned long) src % 128 == 0); - hvx_sigmoid_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); -} - -static inline void hvx_sigmoid_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - assert((unsigned long) dst % 128 == 0); - hvx_sigmoid_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); -} - -static inline void hvx_sigmoid_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - assert((unsigned long) src % 128 == 0); - hvx_sigmoid_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); -} - -static inline void hvx_sigmoid_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - hvx_sigmoid_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); -} - -static inline void hvx_tanh_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - assert((unsigned long) dst % 128 == 0); - assert((unsigned long) src % 128 == 0); - hvx_tanh_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); -} - -#endif /* HVX_SIGMOID_H */ diff --git a/ggml/src/ggml-hexagon/htp/hvx-sqrt.h b/ggml/src/ggml-hexagon/htp/hvx-sqrt.h deleted file mode 100644 index e31a1006d213..000000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-sqrt.h +++ /dev/null @@ -1,126 +0,0 @@ -#ifndef HVX_SQRT_H -#define HVX_SQRT_H - -#include -#include - -#include "hex-utils.h" - -#include "hvx-base.h" - -#define RSQRT_CONST 0x5f3759df // Constant for fast inverse square root calculation -#define RSQRT_ONE_HALF 0x3f000000 // 0.5 -#define RSQRT_THREE_HALVES 0x3fc00000 // 1.5 - -#if __HVX_ARCH__ < 79 -#define HVX_OP_MUL(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b)) -#else -#define HVX_OP_MUL(a, b) Q6_Vsf_vmpy_VsfVsf(a, b) -#endif - -static inline HVX_Vector hvx_vec_rsqrt_f32(HVX_Vector in_vec) { - //Algorithm : - // x2 = input*0.5 - // y = * (long *) &input - // y = 0x5f3759df - (y>>1) - // y = y*(threehalfs - x2*y*y) - - HVX_Vector rsqrtconst = Q6_V_vsplat_R(RSQRT_CONST); - HVX_Vector onehalf = Q6_V_vsplat_R(RSQRT_ONE_HALF); - HVX_Vector threehalfs = Q6_V_vsplat_R(RSQRT_THREE_HALVES); - - HVX_Vector x2, y, ypower2, temp; - - x2 = Q6_Vqf32_vmpy_VsfVsf(in_vec, onehalf); - x2 = Q6_Vqf32_vadd_Vqf32Vsf(x2, Q6_V_vzero()); - - y = Q6_Vw_vasr_VwR(in_vec, 1); - y = Q6_Vw_vsub_VwVw(rsqrtconst, y); - - // 1st iteration - ypower2 = Q6_Vqf32_vmpy_VsfVsf(y, y); - ypower2 = Q6_Vqf32_vadd_Vqf32Vsf(ypower2, Q6_V_vzero()); - temp = Q6_Vqf32_vmpy_Vqf32Vqf32(x2, ypower2); - temp = Q6_Vqf32_vsub_VsfVsf(threehalfs, Q6_Vsf_equals_Vqf32(temp)); - temp = Q6_Vqf32_vmpy_VsfVsf(y, Q6_Vsf_equals_Vqf32(temp)); - - // 2nd iteration - y = Q6_Vqf32_vadd_Vqf32Vsf(temp, Q6_V_vzero()); - ypower2 = Q6_Vqf32_vmpy_Vqf32Vqf32(y, y); - ypower2 = Q6_Vqf32_vadd_Vqf32Vsf(ypower2, Q6_V_vzero()); - temp = Q6_Vqf32_vmpy_Vqf32Vqf32(x2, ypower2); - temp = Q6_Vqf32_vsub_VsfVsf(threehalfs, Q6_Vsf_equals_Vqf32(temp)); - temp = Q6_Vqf32_vmpy_Vqf32Vqf32(y, temp); - - // 3rd iteration - y = Q6_Vqf32_vadd_Vqf32Vsf(temp, Q6_V_vzero()); - ypower2 = Q6_Vqf32_vmpy_Vqf32Vqf32(y, y); - ypower2 = Q6_Vqf32_vadd_Vqf32Vsf(ypower2, Q6_V_vzero()); - temp = Q6_Vqf32_vmpy_Vqf32Vqf32(x2, ypower2); - temp = Q6_Vqf32_vsub_VsfVsf(threehalfs, Q6_Vsf_equals_Vqf32(temp)); - temp = Q6_Vqf32_vmpy_Vqf32Vqf32(y, temp); - - return Q6_Vsf_equals_Vqf32(temp); -} - -// Compute sqrt(x) as x*inv_sqrt(x) -#define hvx_sqrt_f32_loop_body(dst_type, src_type, vec_store) \ - do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src_type * restrict vsrc = (src_type *) src; \ - \ - const uint32_t nvec = n / VLEN_FP32; \ - const uint32_t nloe = n % VLEN_FP32; \ - \ - uint32_t i = 0; \ - \ - _Pragma("unroll(4)") \ - for (; i < nvec; i++) { \ - HVX_Vector inv_sqrt = hvx_vec_rsqrt_f32(vsrc[i]); \ - HVX_Vector sqrt_res = HVX_OP_MUL(inv_sqrt, vsrc[i]); \ - vdst[i] = sqrt_res; \ - } \ - if (nloe) { \ - HVX_Vector inv_sqrt = hvx_vec_rsqrt_f32(vsrc[i]); \ - HVX_Vector sqrt_res = HVX_OP_MUL(inv_sqrt, vsrc[i]); \ - vec_store((void *) &vdst[i], nloe * SIZEOF_FP32, sqrt_res); \ - } \ - } while(0) - -static inline void hvx_sqrt_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - assert((unsigned long) dst % 128 == 0); - assert((unsigned long) src % 128 == 0); - hvx_sqrt_f32_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); -} - -static inline void hvx_sqrt_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - assert((unsigned long) dst % 128 == 0); - hvx_sqrt_f32_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); -} - -static inline void hvx_sqrt_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - assert((unsigned long) src % 128 == 0); - hvx_sqrt_f32_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); -} - -static inline void hvx_sqrt_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { - hvx_sqrt_f32_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); -} - -static inline void hvx_sqrt_f32(uint8_t * restrict dst, const uint8_t * restrict src, const int num_elems) { - if ((unsigned long) dst % 128 == 0) { - if ((unsigned long) src % 128 == 0) { - hvx_sqrt_f32_aa(dst, src, num_elems); - } else { - hvx_sqrt_f32_au(dst, src, num_elems); - } - } else { - if ((unsigned long) src % 128 == 0) { - hvx_sqrt_f32_ua(dst, src, num_elems); - } else { - hvx_sqrt_f32_uu(dst, src, num_elems); - } - } -} - -#endif /* HVX_SQRT_H */ diff --git a/ggml/src/ggml-hexagon/htp/hvx-types.h b/ggml/src/ggml-hexagon/htp/hvx-types-defs.inc similarity index 100% rename from ggml/src/ggml-hexagon/htp/hvx-types.h rename to ggml/src/ggml-hexagon/htp/hvx-types-defs.inc diff --git a/ggml/src/ggml-hexagon/htp/hvx-utils-defs.inc b/ggml/src/ggml-hexagon/htp/hvx-utils-defs.inc new file mode 100644 index 000000000000..8544d4c07af0 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/hvx-utils-defs.inc @@ -0,0 +1,2188 @@ +#ifndef HVX_UTILS_H +#define HVX_UTILS_H + +#include "hex-utils-defs.inc" + +#include "hvx-types-defs.inc" +#include +#include +#include + +#include "hvx-base-defs.inc" + +#define hvx_splat_pragma(x) _Pragma(#x) +#define hvx_splat_loop_body(dst_type, vec_store, unroll_cnt) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + \ + uint32_t nvec = n / (128 / elem_size); \ + uint32_t nloe = n % (128 / elem_size); \ + \ + uint32_t i = 0; \ + \ + hvx_splat_pragma(unroll(unroll_cnt)) \ + for (; i < nvec; i++) { \ + vdst[i] = src; \ + } \ + if (nloe) { \ + vec_store((void *) &vdst[i], nloe * elem_size, src); \ + } \ + } while(0) + +static inline void hvx_splat_a(void * restrict dst, HVX_Vector src, uint32_t n, uint32_t elem_size) { + assert((unsigned long) dst % 128 == 0); + hvx_splat_loop_body(HVX_Vector, hvx_vec_store_a, 4); +} + +static inline void hvx_splat_u(void * restrict dst, HVX_Vector src, uint32_t n, uint32_t elem_size) { + hvx_splat_loop_body(HVX_UVector, hvx_vec_store_u, 4); +} + +static inline void hvx_splat_f32_a(void * restrict dst, float v, uint32_t n) { + hvx_splat_a(dst, hvx_vec_splat_f32(v), n, sizeof(float)); +} + +static inline void hvx_splat_f32_u(void * restrict dst, float v, uint32_t n) { + hvx_splat_u(dst, hvx_vec_splat_f32(v), n, sizeof(float)); +} + +static inline void hvx_splat_f16_a(void * restrict dst, _Float16 v, uint32_t n) { + hvx_splat_u(dst, hvx_vec_splat_f16(v), n, sizeof(__fp16)); +} + +static inline void hvx_splat_f16_u(void * restrict dst, _Float16 v, uint32_t n) { + hvx_splat_u(dst, hvx_vec_splat_f16(v), n, sizeof(__fp16)); +} + +static inline void hvx_splat_u16_a(void * restrict dst, uint16_t v, uint32_t n) { + hvx_splat_a(dst, Q6_Vh_vsplat_R(v), n, sizeof(uint16_t)); +} + +static inline void hvx_splat_u16_u(void * restrict dst, uint16_t v, uint32_t n) { + hvx_splat_u(dst, Q6_Vh_vsplat_R(v), n, sizeof(uint16_t)); +} + +static inline void hvx_splat_u8_a(void * restrict dst, uint8_t v, uint32_t n) { + hvx_splat_a(dst, Q6_Vb_vsplat_R(v), n, 1); +} + +static inline void hvx_splat_u8_u(void * restrict dst, uint8_t v, uint32_t n) { + hvx_splat_u(dst, Q6_Vb_vsplat_R(v), n, 1); +} + +#define hvx_copy_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const uint32_t epv = 128 / elem_size; \ + const uint32_t nvec = n / epv; \ + const uint32_t nloe = n % epv; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { vdst[i] = vsrc[i]; } \ + if (nloe) { \ + vec_store((void *) &vdst[i], nloe * elem_size, vsrc[i]); \ + } \ + } while(0) + +// Generic copy routines +static inline void hvx_copy_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n, uint32_t elem_size) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + hvx_copy_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); +} + +static inline void hvx_copy_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n, uint32_t elem_size) { + assert((unsigned long) dst % 128 == 0); + hvx_copy_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); +} + +static inline void hvx_copy_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n, uint32_t elem_size) { + assert((unsigned long) src % 128 == 0); + hvx_copy_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); +} + +static inline void hvx_copy_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n, uint32_t elem_size) { + hvx_copy_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); +} + +// copy n fp16 elements : source and destination are aligned to HVX Vector (128) +static inline void hvx_copy_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_copy_aa(dst, src, n, sizeof(__fp16)); +} + +// copy n fp16 elements : source is aligned, destination is potentially unaligned +static inline void hvx_copy_f16_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_copy_au(dst, src, n, sizeof(__fp16)); +} + +// copy n fp16 elements : source is aligned, destination is potentially unaligned +static inline void hvx_copy_f16_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_copy_ua(dst, src, n, sizeof(__fp16)); +} + +// copy n fp16 elements : source is aligned, destination is potentially unaligned +static inline void hvx_copy_f16_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_copy_uu(dst, src, n, sizeof(__fp16)); +} + +// copy n fp32 elements : source and destination are aligned to HVX Vector (128) +static inline void hvx_copy_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_copy_aa(dst, src, n, sizeof(float)); +} + +// copy n fp32 elements : source is aligned, destination is unaligned +static inline void hvx_copy_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_copy_ua(dst, src, n, sizeof(float)); +} + +// copy n fp32 elements : source is unaligned, destination is aligned +static inline void hvx_copy_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_copy_au(dst, src, n, sizeof(float)); +} + +// copy n fp32 elements : source is unaligned, destination unaligned +static inline void hvx_copy_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_copy_uu(dst, src, n, sizeof(float)); +} + +//// fp32 -> fp16 + +#define hvx_copy_f16_f32_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const uint32_t elem_size = sizeof(__fp16); \ + const uint32_t epv = 128 / elem_size; \ + const uint32_t nvec = n / epv; \ + const uint32_t nloe = n % epv; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + vdst[i] = hvx_vec_f32_to_f16(vsrc[i*2+0], vsrc[i*2+1]); \ + } \ + if (nloe) { \ + HVX_Vector v = hvx_vec_f32_to_f16(vsrc[i*2+0], vsrc[i*2+1]); \ + vec_store((void *) &vdst[i], nloe * elem_size, v); \ + } \ + } while(0) + +// copy/convert n fp32 elements into n fp16 elements : source is aligned, destination is aligned +static inline void hvx_copy_f16_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + hvx_copy_f16_f32_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); +} + +// copy/convert n fp32 elements into n fp16 elements : source is unaligned, destination is aligned +static inline void hvx_copy_f16_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + hvx_copy_f16_f32_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); +} + +// copy/convert n fp32 elements into n fp16 elements : source is aligned, destination is unaligned +static inline void hvx_copy_f16_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) src % 128 == 0); + hvx_copy_f16_f32_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); +} + +// copy/convert n fp32 elements into n fp16 elements : source is unaligned, destination is unaligned +static inline void hvx_copy_f16_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_copy_f16_f32_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); +} + +//// fp16 -> fp32 + +#define hvx_copy_f32_f16_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const HVX_Vector one = hvx_vec_splat_f16(1.0); \ + \ + const uint32_t elem_size = sizeof(__fp16); \ + const uint32_t epv = 128 / elem_size; \ + const uint32_t nvec = n / epv; \ + uint32_t nloe = n % epv; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (i = 0; i < nvec; ++i) { \ + HVX_VectorPair p = Q6_Wqf32_vmpy_VhfVhf(Q6_Vh_vshuff_Vh(vsrc[i]), one); \ + vdst[i*2] = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(p)); \ + vdst[i*2+1] = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(p)); \ + } \ + \ + if (nloe) { \ + HVX_VectorPair p = Q6_Wqf32_vmpy_VhfVhf(Q6_Vh_vshuff_Vh(vsrc[i]), one); \ + \ + HVX_Vector vd = Q6_V_lo_W(p); \ + i = 2 * i; \ + \ + if (nloe >= 32) { \ + vdst[i] = Q6_Vsf_equals_Vqf32(vd); \ + nloe -= 32; ++i; vd = Q6_V_hi_W(p); \ + } \ + \ + if (nloe) { \ + vd = Q6_Vsf_equals_Vqf32(vd); \ + hvx_vec_store_u(&vdst[i], nloe * sizeof(float), vd); \ + } \ + } \ + } while(0) + +// copy/convert n fp16 elements into n fp32 elements : source is aligned, destination is aligned +static inline void hvx_copy_f32_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + hvx_copy_f32_f16_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); +} + +// copy/convert n fp16 elements into n fp32 elements : source is unaligned, destination is aligned +static inline void hvx_copy_f32_f16_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + hvx_copy_f32_f16_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); +} + +// copy/convert n fp16 elements into n fp32 elements : source is aligned, destination is unaligned +static inline void hvx_copy_f32_f16_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) src % 128 == 0); + hvx_copy_f32_f16_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); +} + +// copy/convert n fp16 elements into n fp32 elements : source is unaligned, destination is unaligned +static inline void hvx_copy_f32_f16_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_copy_f32_f16_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); +} + +#include +#include +#include + +#include "hvx-base-defs.inc" + +#define hvx_scale_f32_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + HVX_Vector vs = hvx_vec_splat_f32(scale); \ + \ + const uint32_t elem_size = sizeof(float); \ + const uint32_t epv = 128 / elem_size; \ + const uint32_t nvec = n / epv; \ + const uint32_t nloe = n % epv; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; ++i) { \ + HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(vsrc[i], vs); \ + vdst[i] = Q6_Vsf_equals_Vqf32(v); \ + } \ + if (nloe) { \ + HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(vsrc[i], vs); \ + vec_store((void *) &vdst[i], nloe * elem_size, Q6_Vsf_equals_Vqf32(v)); \ + } \ + } while(0) + +static inline void hvx_scale_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale) { + assert((size_t) dst % 128 == 0); + assert((size_t) src % 128 == 0); + hvx_scale_f32_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); +} + +static inline void hvx_scale_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale) { + assert((size_t) dst % 128 == 0); + hvx_scale_f32_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); +} + +static inline void hvx_scale_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale) { + assert((size_t) src % 128 == 0); + hvx_scale_f32_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); +} + +static inline void hvx_scale_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale) { + hvx_scale_f32_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); +} + +static inline void hvx_scale_f32(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale) { + if (((size_t) dst & 127) == 0) { + if (((size_t) src & 127) == 0) { + hvx_scale_f32_aa(dst, src, n, scale); + } else { + hvx_scale_f32_au(dst, src, n, scale); + } + } else { + if (((size_t) src & 127) == 0) { + hvx_scale_f32_ua(dst, src, n, scale); + } else { + hvx_scale_f32_uu(dst, src, n, scale); + } + } +} + +#define hvx_scale_offset_f32_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + HVX_Vector vs = hvx_vec_splat_f32(scale); \ + HVX_Vector vo = hvx_vec_splat_f32(offset); \ + \ + const uint32_t elem_size = sizeof(float); \ + const uint32_t epv = 128 / elem_size; \ + const uint32_t nvec = n / epv; \ + const uint32_t nloe = n % epv; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; ++i) { \ + HVX_Vector v = Q6_Vqf32_vadd_Vqf32Vsf(Q6_Vqf32_vmpy_VsfVsf(vsrc[i], vs), vo); \ + vdst[i] = Q6_Vsf_equals_Vqf32(v); \ + } \ + if (nloe) { \ + HVX_Vector v = Q6_Vqf32_vadd_Vqf32Vsf(Q6_Vqf32_vmpy_VsfVsf(vsrc[i], vs), vo); \ + vec_store((void *) &vdst[i], nloe * elem_size, Q6_Vsf_equals_Vqf32(v)); \ + } \ + } while(0) + +static inline void hvx_scale_offset_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) { + assert((size_t) dst % 128 == 0); + assert((size_t) src % 128 == 0); + hvx_scale_offset_f32_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); +} + +static inline void hvx_scale_offset_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) { + assert((size_t) dst % 128 == 0); + hvx_scale_offset_f32_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); +} + +static inline void hvx_scale_offset_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) { + assert((size_t) src % 128 == 0); + hvx_scale_offset_f32_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); +} + +static inline void hvx_scale_offset_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) { + hvx_scale_offset_f32_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); +} + +static inline void hvx_scale_offset_f32(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) { + if (((size_t) dst & 127) == 0) { + if (((size_t) src & 127) == 0) { + hvx_scale_offset_f32_aa(dst, src, n, scale, offset); + } else { + hvx_scale_offset_f32_au(dst, src, n, scale, offset); + } + } else { + if (((size_t) src & 127) == 0) { + hvx_scale_offset_f32_ua(dst, src, n, scale, offset); + } else { + hvx_scale_offset_f32_uu(dst, src, n, scale, offset); + } + } +} + + +#include +#include +#include + +#include "hvx-base-defs.inc" +#include +#include + +#include "hvx-base-defs.inc" + +#define IEEE_VSF_EXPLEN (8) +#define IEEE_VSF_EXPBIAS (127) +#define IEEE_VSF_EXPMASK (0xFF) +#define IEEE_VSF_MANTLEN (23) +#define IEEE_VSF_MANTMASK (0x7FFFFF) +#define IEEE_VSF_MIMPMASK (0x800000) + +static inline HVX_Vector hvx_vec_truncate_f32(HVX_Vector in_vec) { + HVX_Vector mask_mant_v = Q6_V_vsplat_R(IEEE_VSF_MANTMASK); + HVX_Vector mask_impl_v = Q6_V_vsplat_R(IEEE_VSF_MIMPMASK); + HVX_Vector const_zero_v = Q6_V_vzero(); + + HVX_VectorPred q_negative = Q6_Q_vcmp_gt_VwVw(const_zero_v, in_vec); + + HVX_Vector expval_v = in_vec >> IEEE_VSF_MANTLEN; + expval_v &= IEEE_VSF_EXPMASK; + expval_v -= IEEE_VSF_EXPBIAS; + + // negative exp == fractional value + HVX_VectorPred q_negexp = Q6_Q_vcmp_gt_VwVw(const_zero_v, expval_v); + + HVX_Vector rshift_v = IEEE_VSF_MANTLEN - expval_v; // fractional bits - exp shift + + HVX_Vector mant_v = in_vec & mask_mant_v; // obtain mantissa + HVX_Vector vout = Q6_Vw_vadd_VwVw(mant_v, mask_impl_v); // add implicit 1.0 + + vout = Q6_Vw_vasr_VwVw(vout, rshift_v); // shift to obtain truncated integer + vout = Q6_V_vmux_QVV(q_negexp, const_zero_v, vout); // expval<0 -> 0 + + HVX_Vector neg_vout = -vout; + + vout = Q6_V_vmux_QVV(q_negative, neg_vout, vout); // handle negatives + + return (vout); +} + +static inline HVX_Vector hvx_vec_floor_f32(HVX_Vector in_vec) { + HVX_Vector mask_mant_v = Q6_V_vsplat_R(IEEE_VSF_MANTMASK); + HVX_Vector mask_impl_v = Q6_V_vsplat_R(IEEE_VSF_MIMPMASK); + HVX_Vector const_mnlen_v = Q6_V_vsplat_R(IEEE_VSF_MANTLEN); + HVX_Vector const_zero_v = Q6_V_vzero(); + HVX_Vector const_negone_v = Q6_V_vsplat_R(0xbf800000); // -1 IEEE vsf + + HVX_VectorPred q_negative = Q6_Q_vcmp_gt_VwVw(const_zero_v, in_vec); + + HVX_Vector expval_v = in_vec >> IEEE_VSF_MANTLEN; + expval_v &= IEEE_VSF_EXPMASK; + expval_v -= IEEE_VSF_EXPBIAS; + + HVX_VectorPred q_negexp = Q6_Q_vcmp_gt_VwVw(const_zero_v, expval_v); + HVX_VectorPred q_expltmn = Q6_Q_vcmp_gt_VwVw(const_mnlen_v, expval_v); + HVX_VectorPred q_negexp_pos = Q6_Q_vcmp_gtand_QVwVw(q_negexp, in_vec, const_zero_v); + HVX_VectorPred q_negexp_neg = Q6_Q_vcmp_gtand_QVwVw(q_negexp, const_zero_v, in_vec); + + // if expval < 0 (q_negexp) // <0, floor is 0 + // if vin > 0 + // floor = 0 + // if vin < 0 + // floor = -1 + // if expval < mant_len (q_expltmn) // >0, but fraction may exist + // get sign (q_negative) + // mask >> expval // fraction bits to mask off + // vout = ~(mask) // apply mask to remove fraction + // if (qneg) // negative floor is one less (more, sign bit for neg) + // vout += ((impl_mask) >> expval) + // if (mask && vin) + // vout = vin + // else // already an integer + // ; // no change + + // compute floor + mask_mant_v >>= expval_v; + HVX_Vector neg_addin_v = mask_impl_v >> expval_v; + HVX_Vector vout_neg_addin = Q6_Vw_vadd_VwVw(in_vec, neg_addin_v); + HVX_Vector vout = Q6_V_vmux_QVV(q_negative, vout_neg_addin, in_vec); + + HVX_Vector mask_chk_v = Q6_V_vand_VV(in_vec, mask_mant_v); // chk if bits set + HVX_VectorPred q_integral = Q6_Q_vcmp_eq_VwVw(const_zero_v, mask_chk_v); + + HVX_Vector not_mask_v = Q6_V_vnot_V(mask_mant_v); // frac bits to clear + HVX_Vector vfrfloor_v = Q6_V_vand_VV(vout, not_mask_v); // clear frac bits + + vout = in_vec; + vout = Q6_V_vmux_QVV(q_expltmn, vfrfloor_v, vout); // expval0 -> 0 + vout = Q6_V_vmux_QVV(q_negexp_neg, const_negone_v, vout); // expval<0 x<0 -> -1 + + return vout; +} + + + +#define EXP_COEFF_5 (0x39506967) // 0.000198757 = 1/(7!) +#define EXP_COEFF_4 (0x3AB743CE) // 0.0013982 = 1/(6!) +#define EXP_COEFF_3 (0x3C088908) // 0.00833345 = 1/(5!) +#define EXP_COEFF_2 (0x3D2AA9C1) // 0.416658 = 1/(4!) +#define EXP_COEFF_1 (0x3E2AAAAA) // 0.16666667 = 1/(3!) +#define EXP_COEFF_0 (0x3F000000) // 0.5 = 1/(2!) +#define EXP_LOGN2 (0x3F317218) // ln(2) = 0.6931471805 +#define EXP_LOG2E (0x3FB8AA3B) // log2(e) = 1/ln(2) = 1.4426950408 +#define EXP_ONE (0x3f800000) // 1.0 +#define EXP_RANGE_R (0x42B17218) // ln(FLT_MAX) approx = 88.7228 +#define EXP_RANGE_L (0xC2B00000) // -88.0 (approx log(FLT_MIN)) + +static inline HVX_Vector hvx_vec_exp_f32(HVX_Vector in_vec) { + HVX_Vector z_qf32_v; + HVX_Vector x_v; + HVX_Vector x_qf32_v; + HVX_Vector y_v; + HVX_Vector k_v; + HVX_Vector f_v; + HVX_Vector epsilon_v; + HVX_Vector log2e = Q6_V_vsplat_R(EXP_LOG2E); + HVX_Vector logn2 = Q6_V_vsplat_R(EXP_LOGN2); + HVX_Vector E_const; + HVX_Vector zero_v = Q6_V_vzero(); + + // exp(x) is approximated as follows: + // f = floor(x/ln(2)) = floor(x*log2(e)) + // epsilon = x - f*ln(2) + // exp(x) = exp(epsilon+f*ln(2)) + // = exp(epsilon)*exp(f*ln(2)) + // = exp(epsilon)*2^f + // + // Since epsilon is close to zero, it can be approximated with its Taylor series: + // exp(x) ~= 1+x+x^2/2!+x^3/3!+...+x^n/n!+... + // Preserving the first eight elements, we get: + // exp(x) ~= 1+x+e0*x^2+e1*x^3+e2*x^4+e3*x^5+e4*x^6+e5*x^7 + // = 1+x+(E0+(E1+(E2+(E3+(E4+E5*x)*x)*x)*x)*x)*x^2 + + HVX_Vector temp_v = in_vec; + + // Clamp inputs to (-88.0, 88.0) to avoid overflow/underflow + HVX_VectorPred pred_cap_right = Q6_Q_vcmp_gt_VsfVsf(in_vec, Q6_V_vsplat_R(EXP_RANGE_R)); + HVX_VectorPred pred_cap_left = Q6_Q_vcmp_gt_VsfVsf(Q6_V_vsplat_R(EXP_RANGE_L), in_vec); + + in_vec = Q6_V_vmux_QVV(pred_cap_right, Q6_V_vsplat_R(EXP_RANGE_R), temp_v); + in_vec = Q6_V_vmux_QVV(pred_cap_left, Q6_V_vsplat_R(EXP_RANGE_L), in_vec); + + epsilon_v = Q6_Vqf32_vmpy_VsfVsf(log2e, in_vec); + epsilon_v = Q6_Vsf_equals_Vqf32(epsilon_v); + + // f_v is the floating point result and k_v is the integer result + f_v = hvx_vec_floor_f32(epsilon_v); + k_v = hvx_vec_truncate_f32(f_v); + + x_qf32_v = Q6_Vqf32_vadd_VsfVsf(in_vec, zero_v); + + // x = x - f_v * logn2; + epsilon_v = Q6_Vqf32_vmpy_VsfVsf(f_v, logn2); + x_qf32_v = Q6_Vqf32_vsub_Vqf32Vqf32(x_qf32_v, epsilon_v); + // normalize before every QFloat's vmpy + x_qf32_v = Q6_Vqf32_vadd_Vqf32Vsf(x_qf32_v, zero_v); + + x_v = Q6_Vsf_equals_Vqf32(x_qf32_v); + + // z = x * x; + z_qf32_v = Q6_Vqf32_vmpy_Vqf32Vqf32(x_qf32_v, x_qf32_v); + z_qf32_v = Q6_Vqf32_vadd_Vqf32Vsf(z_qf32_v, zero_v); + + // y = E4 + E5 * x; + E_const = Q6_V_vsplat_R(EXP_COEFF_5); + y_v = Q6_Vqf32_vmpy_VsfVsf(E_const, x_v); + E_const = Q6_V_vsplat_R(EXP_COEFF_4); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); + + // y = E3 + y * x; + E_const = Q6_V_vsplat_R(EXP_COEFF_3); + y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); + + // y = E2 + y * x; + E_const = Q6_V_vsplat_R(EXP_COEFF_2); + y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); + + // y = E1 + y * x; + E_const = Q6_V_vsplat_R(EXP_COEFF_1); + y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); + + // y = E0 + y * x; + E_const = Q6_V_vsplat_R(EXP_COEFF_0); + y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); + + // y = x + y * z; + y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, z_qf32_v); + y_v = Q6_Vqf32_vadd_Vqf32Vqf32(y_v, x_qf32_v); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); + + // y = y + 1.0; + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, Q6_V_vsplat_R(EXP_ONE)); + + // insert exponents + // y = ldexpf(y, k); + // y_v += k_v; // qf32 + // modify exponent + + y_v = Q6_Vsf_equals_Vqf32(y_v); + + // add k_v to the exponent of y_v + HVX_Vector y_v_exponent = Q6_Vw_vasl_VwR(y_v, 1); + + y_v_exponent = Q6_Vuw_vlsr_VuwR(y_v_exponent, IEEE_VSF_MANTLEN + 1); + y_v_exponent = Q6_Vw_vadd_VwVw(k_v, y_v_exponent); + + // exponent cannot be negative; if overflow is detected, result is set to zero + HVX_VectorPred qy_v_negative_exponent = Q6_Q_vcmp_gt_VwVw(zero_v, y_v_exponent); + + y_v = Q6_Vw_vaslacc_VwVwR(y_v, k_v, IEEE_VSF_MANTLEN); + + y_v = Q6_V_vmux_QVV(qy_v_negative_exponent, zero_v, y_v); + + return y_v; +} + +static inline HVX_Vector hvx_vec_exp_f32_guard(HVX_Vector in_vec, HVX_Vector max_exp, HVX_Vector inf) { + const HVX_VectorPred pred0 = Q6_Q_vcmp_gt_VsfVsf(in_vec, max_exp); + + HVX_Vector out = hvx_vec_exp_f32(in_vec); + + return Q6_V_vmux_QVV(pred0, inf, out); +} + +static inline void hvx_exp_f32(uint8_t * restrict dst, const uint8_t * restrict src, const int num_elems, bool negate) { + int left_over = num_elems & (VLEN_FP32 - 1); + int num_elems_whole = num_elems - left_over; + + int unaligned_addr = 0; + int unaligned_loop = 0; + if ((0 == hex_is_aligned((void *) src, VLEN)) || (0 == hex_is_aligned((void *) dst, VLEN))) { + unaligned_addr = 1; + } + // assert((0 == unaligned_addr) || (0 == num_elems_whole)); + if ((1 == unaligned_addr) && (num_elems_whole != 0)) { + unaligned_loop = 1; + } + + HVX_Vector vec_out = Q6_V_vzero(); + + static const float kInf = INFINITY; + static const float kMaxExp = 88.7228f; + + const HVX_Vector max_exp = hvx_vec_splat_f32(kMaxExp); + const HVX_Vector inf = hvx_vec_splat_f32(kInf); + + if (0 == unaligned_loop) { + HVX_Vector * p_vec_in1 = (HVX_Vector *) src; + HVX_Vector * p_vec_out = (HVX_Vector *) dst; + + #pragma unroll(4) + for (int i = 0; i < num_elems_whole; i += VLEN_FP32) { + if (true == negate) { + HVX_Vector neg_vec_in = hvx_vec_neg_f32(*p_vec_in1++); + *p_vec_out++ = hvx_vec_exp_f32_guard(neg_vec_in, max_exp, inf); + } else { + *p_vec_out++ = hvx_vec_exp_f32_guard(*p_vec_in1++, max_exp, inf); + } + } + } else { + #pragma unroll(4) + for (int i = 0; i < num_elems_whole; i += VLEN_FP32) { + HVX_Vector in = *(HVX_UVector *) (src + i * SIZEOF_FP32); + + if (true == negate) { + HVX_Vector neg_vec_in = hvx_vec_neg_f32(in); + *(HVX_UVector *) (dst + i * SIZEOF_FP32) = hvx_vec_exp_f32_guard(neg_vec_in, max_exp, inf); + } else { + *(HVX_UVector *) (dst + i * SIZEOF_FP32) = hvx_vec_exp_f32_guard(in, max_exp, inf); + } + } + } + + if (left_over > 0) { + const float * srcf = (float *) src + num_elems_whole; + float * dstf = (float *) dst + num_elems_whole; + + HVX_Vector in = *(HVX_UVector *) srcf; + + if (true == negate) { + HVX_Vector neg_vec_in = hvx_vec_neg_f32(in); + + vec_out = hvx_vec_exp_f32_guard(neg_vec_in, max_exp, inf); + } else { + vec_out = hvx_vec_exp_f32_guard(in, max_exp, inf); + } + + hvx_vec_store_u((void *) dstf, left_over * SIZEOF_FP32, vec_out); + } +} + +#include + +#include +#include +#include +#include +#include + +#include "hvx-base-defs.inc" + +// ==================================================== +// FUNCTION: 1/(x+1) y(0) = 1, y(0.5) = 0.6667, y(1) = 0.5 +// Order:3; continuity: True; Ends forced: True +// Mode: unsigned; Result fractional bits: 14 +// Peak Error: 1.1295e-04 Rms Error: 2.8410e-05 Mean Error: 1.1370e-05 +// 32769 -32706 31252 -10589 +// 32590 -30635 22793 -4493 +// 32066 -27505 16481 -2348 +// 31205 -24054 11849 -1306 + +static inline HVX_Vector hvx_vec_recip_xp1_O3_unsigned(HVX_Vector vx) { + // input is 0..0xffff representing 0.0 .. 1.0 + HVX_Vector p; + p = Q6_Vh_vlut4_VuhPh(vx, 0xFAE6F6D4EE73D6A3ull); + p = Q6_Vh_vmpa_VhVhVuhPuh_sat(p, vx, 0x2E49406159097A14ull); + p = Q6_Vh_vmps_VhVhVuhPuh_sat(p, vx, 0x5DF66B7177AB7FC2ull); + p = Q6_Vh_vmpa_VhVhVuhPuh_sat(p, vx, 0x79E57D427F4E8001ull); + return p; // signed result, 14 fractional bits +} + +// Find reciprocal of fp16. +// (1) first, convert to fp32, multiplying by 1.0; this is done to +// handle denormals. Ignoring sign and zero, result should be at +// least 5.9604645e-08 (32-bit code 0x33800000) and at most 131008 (0x47ffe000) +// (exponent in range [103,143]) +// (2) extract the mantissa into 16-bit unsigned; find reciprocal using a fitted poly +// (3) put this, along with '253-exp' (exp from (1)) together to make an qf32 +// (4) convert that to fp16 +// (5) put sign back in. Also, if the original value (w/o sign) was <0x81, replace +// the result with the max value. +static inline HVX_Vector hvx_vec_inverse_f16(HVX_Vector vals) { + HVX_Vector em_mask = Q6_Vh_vsplat_R(0x7FFF); + HVX_Vector avals = Q6_V_vand_VV(vals, em_mask); + HVX_VectorPred is_neg = Q6_Q_vcmp_gt_VhVh(avals, vals); + // is too small to 1/x ? for 'standard' fp16, this would be 0x101 + HVX_VectorPred is_small = Q6_Q_vcmp_gt_VhVh(Q6_Vh_vsplat_R(0x101), avals); + + HVX_VectorPair to_qf32 = Q6_Wqf32_vmpy_VhfVhf(avals, Q6_Vh_vsplat_R(0x3C00)); // *1.0 + HVX_Vector to_f32_0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(to_qf32)); + HVX_Vector to_f32_1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(to_qf32)); + + // bits 22..13 contain the mantissa now (w/o hidden bit); move to bit 14..5 of a 16-bit vector + HVX_Vector mant_u16 = Q6_Vh_vshuffo_VhVh(Q6_Vw_vasl_VwR(to_f32_1, 9), Q6_Vw_vasl_VwR(to_f32_0, 9)); + // likewise extract the upper 16 from each, containing the exponents in range 103..142 + HVX_Vector exp_u16 = Q6_Vh_vshuffo_VhVh(to_f32_1, to_f32_0); + //Get exponent in IEEE 32-bit representation + exp_u16 = Q6_Vuh_vlsr_VuhR(exp_u16, 7); + + // so, mant_u16 contains an unbiased mantissa in upper 10 bits of each u16 lane + // We can consider it to be x-1.0, with 16 fractional bits, where 'x' is in range [1.0,2.0) + // Use poly to transform to 1/x, with 14 fractional bits + // + HVX_Vector rm = hvx_vec_recip_xp1_O3_unsigned(mant_u16); + + HVX_Vector vcl0 = Q6_Vuh_vcl0_Vuh(rm); //count leading zeros + + // Get mantissa for 16-bit representation + HVX_Vector mant_recip = Q6_V_vand_VV(Q6_Vh_vasr_VhR(Q6_Vh_vasl_VhVh(rm, vcl0), 5), Q6_Vh_vsplat_R(0x03FF)); + + //Compute Reciprocal Exponent + HVX_Vector exp_recip = + Q6_Vh_vsub_VhVh(Q6_Vh_vsub_VhVh(Q6_Vh_vsplat_R(254), exp_u16), Q6_Vh_vsub_VhVh(vcl0, Q6_Vh_vsplat_R(1))); + //Convert it for 16-bit representation + exp_recip = Q6_Vh_vadd_VhVh_sat(Q6_Vh_vsub_VhVh(exp_recip, Q6_Vh_vsplat_R(127)), Q6_Vh_vsplat_R(15)); + exp_recip = Q6_Vh_vasl_VhR(exp_recip, 10); + + //Merge exponent and mantissa for reciprocal + HVX_Vector recip = Q6_V_vor_VV(exp_recip, mant_recip); + // map 'small' inputs to standard largest value 0x7bff + recip = Q6_V_vmux_QVV(is_small, Q6_Vh_vsplat_R(0x7bff), recip); + // add sign back + recip = Q6_V_vandor_VQR(recip, is_neg, 0x80008000); + return recip; +} + +static inline HVX_Vector hvx_vec_inverse_f32(HVX_Vector v_sf) { + HVX_Vector inv_aprox_sf = Q6_V_vsplat_R(0x7EEEEBB3); + HVX_Vector two_sf = hvx_vec_splat_f32(2.0); + + // First approximation + HVX_Vector i_sf = Q6_Vw_vsub_VwVw(inv_aprox_sf, v_sf); + + HVX_Vector r_qf; + + // Refine + r_qf = Q6_Vqf32_vmpy_VsfVsf( + i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v_sf))))); + r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32( + r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v_sf)))); + r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32( + r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v_sf)))); + + return Q6_Vsf_equals_Vqf32(r_qf); +} + +static inline HVX_Vector hvx_vec_inverse_f32_guard(HVX_Vector v_sf, HVX_Vector nan_inf_mask) { + HVX_Vector out = hvx_vec_inverse_f32(v_sf); + + HVX_Vector masked_out = Q6_V_vand_VV(out, nan_inf_mask); + const HVX_VectorPred pred = Q6_Q_vcmp_eq_VwVw(nan_inf_mask, masked_out); + + return Q6_V_vmux_QVV(pred, Q6_V_vzero(), out); +} + +#define hvx_inverse_f32_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const HVX_Vector nan_inf_mask = Q6_V_vsplat_R(0x7f800000); \ + \ + const uint32_t nvec = n / VLEN_FP32; \ + const uint32_t nloe = n % VLEN_FP32; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + vdst[i] = hvx_vec_inverse_f32_guard(vsrc[i], nan_inf_mask); \ + } \ + if (nloe) { \ + HVX_Vector v = hvx_vec_inverse_f32_guard(vsrc[i], nan_inf_mask); \ + vec_store((void *) &vdst[i], nloe * SIZEOF_FP32, v); \ + } \ + } while(0) + +static inline HVX_Vector hvx_vec_inverse_f16_guard(HVX_Vector v_sf, HVX_Vector nan_inf_mask) { + HVX_Vector out = hvx_vec_inverse_f16(v_sf); + + HVX_Vector masked_out = Q6_V_vand_VV(out, nan_inf_mask); + const HVX_VectorPred pred = Q6_Q_vcmp_eq_VhVh(nan_inf_mask, masked_out); + + return Q6_V_vmux_QVV(pred, Q6_V_vzero(), out); +} + +#define hvx_inverse_f16_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const HVX_Vector nan_inf_mask = Q6_Vh_vsplat_R(0x7c00); \ + \ + const uint32_t nvec = n / VLEN_FP16; \ + const uint32_t nloe = n % VLEN_FP16; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + vdst[i] = hvx_vec_inverse_f16_guard(vsrc[i], nan_inf_mask); \ + } \ + if (nloe) { \ + HVX_Vector v = hvx_vec_inverse_f16_guard(vsrc[i], nan_inf_mask); \ + vec_store((void *) &vdst[i], nloe * SIZEOF_FP16, v); \ + } \ + } while(0) + +// Generic macro to define alignment permutations for an op +#define DEFINE_HVX_INV_OP_VARIANTS(OP_NAME, OP_LOOP_BODY) \ +static inline void OP_NAME##_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \ + assert((uintptr_t) dst % 128 == 0); \ + assert((uintptr_t) src % 128 == 0); \ + OP_LOOP_BODY(HVX_Vector, HVX_Vector, hvx_vec_store_a); \ +} \ +static inline void OP_NAME##_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \ + assert((uintptr_t) dst % 128 == 0); \ + OP_LOOP_BODY(HVX_Vector, HVX_UVector, hvx_vec_store_a); \ +} \ +static inline void OP_NAME##_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \ + assert((uintptr_t) src % 128 == 0); \ + OP_LOOP_BODY(HVX_UVector, HVX_Vector, hvx_vec_store_u); \ +} \ +static inline void OP_NAME##_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \ + OP_LOOP_BODY(HVX_UVector, HVX_UVector, hvx_vec_store_u); \ +} \ + +// Dispatcher logic +#define HVX_INV_DISPATCHER(OP_NAME) \ +static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t num_elems) { \ + if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { \ + OP_NAME##_aa(dst, src, num_elems); \ + } else if (hex_is_aligned((void *) dst, 128)) { \ + OP_NAME##_au(dst, src, num_elems); \ + } else if (hex_is_aligned((void *) src, 128)) { \ + OP_NAME##_ua(dst, src, num_elems); \ + } else { \ + OP_NAME##_uu(dst, src, num_elems); \ + } \ +} + +DEFINE_HVX_INV_OP_VARIANTS(hvx_inverse_f32, hvx_inverse_f32_loop_body) +DEFINE_HVX_INV_OP_VARIANTS(hvx_inverse_f16, hvx_inverse_f16_loop_body) + +HVX_INV_DISPATCHER(hvx_inverse_f32) +HVX_INV_DISPATCHER(hvx_inverse_f16) + +#include +#include +#include +#include + +#include "hex-utils-defs.inc" +#include "hvx-base-defs.inc" +#include "hvx-types-defs.inc" + +static inline HVX_Vector hvx_vec_reduce_sum_n_i32(HVX_Vector in, unsigned int n) { + unsigned int total = n * 4; // total vec nbytes + unsigned int width = 4; // int32 + + HVX_Vector sum = in, sum_t; + while (width < total) { + sum_t = Q6_V_vror_VR(sum, width); // rotate right + sum = Q6_Vw_vadd_VwVw(sum_t, sum); // elementwise sum + width = width << 1; + } + return sum; +} + +static inline HVX_Vector hvx_vec_reduce_sum_i32(HVX_Vector in) { + return hvx_vec_reduce_sum_n_i32(in, 32); +} + +static inline HVX_Vector hvx_vec_reduce_sum_n_qf32(HVX_Vector in, unsigned int n) { + unsigned int total = n * 4; // total vec nbytes + unsigned int width = 4; // fp32 nbytes + + HVX_Vector sum = in, sum_t; + while (width < total) { + sum_t = Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum), width); // rotate right + sum = Q6_Vqf32_vadd_Vqf32Vsf(sum, sum_t); // elementwise sum + width = width << 1; + } + return sum; +} + +static inline HVX_Vector hvx_vec_reduce_sum_qf32(HVX_Vector in) { + return hvx_vec_reduce_sum_n_qf32(in, 32); +} + +#if __HVX_ARCH__ > 75 + +static inline HVX_Vector hvx_vec_reduce_sum_f32x4(HVX_Vector_x4 in) { + HVX_VectorPair sum_p01 = Q6_W_vshuff_VVR(in.v[1], in.v[0], 4); + HVX_VectorPair sum_p23 = Q6_W_vshuff_VVR(in.v[3], in.v[2], 4); + HVX_Vector sum_sf01 = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p01), Q6_V_hi_W(sum_p01)); + HVX_Vector sum_sf23 = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p23), Q6_V_hi_W(sum_p23)); + + HVX_VectorPair sum_p0123 = Q6_W_vshuff_VVR(sum_sf23, sum_sf01, 8); + HVX_Vector sum_sf = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p0123), Q6_V_hi_W(sum_p0123)); + + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 2)); + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 4)); + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 8)); + return sum_sf; +} + +static inline HVX_Vector hvx_vec_reduce_sum_f32x2(HVX_Vector in0, HVX_Vector in1) { + HVX_VectorPair sump = Q6_W_vshuff_VVR(in1, in0, 4); + HVX_Vector sum_sf = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sump), Q6_V_hi_W(sump)); + + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 2)); + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 4)); + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 8)); + sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 16)); + return sum_sf; +} + +static inline HVX_Vector hvx_vec_reduce_sum_n_f32(HVX_Vector in, unsigned int n) { + unsigned int total = n * 4; // total vec nbytes + unsigned int width = 4; // fp32 nbytes + + HVX_Vector sum = in, sum_t; + while (width < total) { + sum_t = Q6_V_vror_VR(sum, width); // rotate right + sum = Q6_Vsf_vadd_VsfVsf(sum, sum_t); // elementwise sum + width = width << 1; + } + return sum; +} + +#else + +static inline HVX_Vector hvx_vec_reduce_sum_f32x4(HVX_Vector_x4 in) { + HVX_VectorPair sum_p01 = Q6_W_vshuff_VVR(in.v[1], in.v[0], 4); + HVX_VectorPair sum_p23 = Q6_W_vshuff_VVR(in.v[3], in.v[2], 4); + HVX_Vector sum_qf01 = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p01), Q6_V_hi_W(sum_p01)); + HVX_Vector sum_qf23 = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p23), Q6_V_hi_W(sum_p23)); + + HVX_VectorPair sum_p0123 = Q6_W_vshuff_VVR(Q6_Vsf_equals_Vqf32(sum_qf23), Q6_Vsf_equals_Vqf32(sum_qf01), 8); + HVX_Vector sum_qf = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p0123), Q6_V_hi_W(sum_p0123)); + + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 2)); + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 4)); + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 8)); + return Q6_Vsf_equals_Vqf32(sum_qf); +} + +static inline HVX_Vector hvx_vec_reduce_sum_f32x2(HVX_Vector in0, HVX_Vector in1) { + HVX_VectorPair sump = Q6_W_vshuff_VVR(in1, in0, 4); + HVX_Vector sum_qf = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sump), Q6_V_hi_W(sump)); + + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 2)); + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 4)); + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 8)); + sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 16)); + return Q6_Vsf_equals_Vqf32(sum_qf); +} + +static inline HVX_Vector hvx_vec_reduce_sum_n_f32(HVX_Vector in, unsigned int n) { + unsigned int total = n * 4; // total vec nbytes + unsigned int width = 4; // fp32 nbytes + + HVX_Vector sum = in, sum_t; + while (width < total) { + sum_t = Q6_V_vror_VR(sum, width); // rotate right + sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(sum, sum_t)); // elementwise sum + width = width << 1; + } + return sum; +} + +#endif + +static inline HVX_Vector hvx_vec_reduce_sum_f32(HVX_Vector in) { + return hvx_vec_reduce_sum_n_f32(in, 32); +} + +static inline HVX_Vector hvx_vec_reduce_max_f16(HVX_Vector in) { + unsigned total = 128; // total vec nbytes + unsigned width = 2; // fp16 nbytes + + HVX_Vector _max = in, _max_t; + while (width < total) { + _max_t = Q6_V_vror_VR(_max, width); // rotate right + _max = Q6_Vhf_vmax_VhfVhf(_max_t, _max); // elementwise max + width = width << 1; + } + + return _max; +} + +static inline HVX_Vector hvx_vec_reduce_max2_f16(HVX_Vector in, HVX_Vector _max) { + unsigned total = 128; // total vec nbytes + unsigned width = 2; // fp32 nbytes + + HVX_Vector _max_t; + + _max = Q6_Vhf_vmax_VhfVhf(in, _max); + while (width < total) { + _max_t = Q6_V_vror_VR(_max, width); // rotate right + _max = Q6_Vhf_vmax_VhfVhf(_max_t, _max); // elementwise max + width = width << 1; + } + + return _max; +} + +static inline HVX_Vector hvx_vec_reduce_max_f32(HVX_Vector in) { + unsigned total = 128; // total vec nbytes + unsigned width = 4; // fp32 nbytes + + HVX_Vector _max = in, _max_t; + while (width < total) { + _max_t = Q6_V_vror_VR(_max, width); // rotate right + _max = Q6_Vsf_vmax_VsfVsf(_max_t, _max); // elementwise max + width = width << 1; + } + + return _max; +} + +static inline HVX_Vector hvx_vec_reduce_max2_f32(HVX_Vector in, HVX_Vector _max) { + unsigned total = 128; // total vec nbytes + unsigned width = 4; // fp32 nbytes + + HVX_Vector _max_t; + + _max = Q6_Vsf_vmax_VsfVsf(in, _max); + while (width < total) { + _max_t = Q6_V_vror_VR(_max, width); // rotate right + _max = Q6_Vsf_vmax_VsfVsf(_max_t, _max); // elementwise max + width = width << 1; + } + + return _max; +} + +#define hvx_reduce_loop_body(src_type, init_vec, pad_vec, vec_op, reduce_op, scalar_reduce) \ + do { \ + src_type * restrict vsrc = (src_type *) src; \ + HVX_Vector acc = init_vec; \ + \ + const uint32_t elem_size = sizeof(float); \ + const uint32_t epv = 128 / elem_size; \ + const uint32_t nvec = num_elems / epv; \ + const uint32_t nloe = num_elems % epv; \ + \ + uint32_t i = 0; \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + acc = vec_op(acc, vsrc[i]); \ + } \ + if (nloe) { \ + const float * srcf = (const float *) src + i * epv; \ + HVX_Vector in = *(HVX_UVector *) srcf; \ + HVX_Vector temp = Q6_V_valign_VVR(in, pad_vec, nloe * elem_size); \ + acc = vec_op(acc, temp); \ + } \ + HVX_Vector v = reduce_op(acc); \ + return scalar_reduce(v); \ + } while(0) + +#define HVX_REDUCE_MAX_OP(acc, val) Q6_Vsf_vmax_VsfVsf(acc, val) +#define HVX_REDUCE_SUM_OP(acc, val) Q6_Vqf32_vadd_VsfVsf(Q6_Vsf_equals_Vqf32(acc), val) +#define HVX_SUM_SQ_OP(acc, val) Q6_Vqf32_vadd_Vqf32Vqf32(acc, Q6_Vqf32_vmpy_VsfVsf(val, val)) +#define HVX_REDUCE_MAX_SCALAR(v) hvx_vec_get_f32(v) +#define HVX_REDUCE_SUM_SCALAR(v) hvx_vec_get_f32(Q6_Vsf_equals_Vqf32(v)) + +// Max variants + +static inline float hvx_reduce_max_f32_a(const uint8_t * restrict src, const int num_elems) { + HVX_Vector init_vec = hvx_vec_splat_f32(((const float *) src)[0]); + assert((unsigned long) src % 128 == 0); + hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_REDUCE_MAX_OP, hvx_vec_reduce_max_f32, HVX_REDUCE_MAX_SCALAR); +} + +static inline float hvx_reduce_max_f32_u(const uint8_t * restrict src, const int num_elems) { + HVX_Vector init_vec = hvx_vec_splat_f32(((const float *) src)[0]); + hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_REDUCE_MAX_OP, hvx_vec_reduce_max_f32, HVX_REDUCE_MAX_SCALAR); +} + +static inline float hvx_reduce_max_f32(const uint8_t * restrict src, const int num_elems) { + if (hex_is_aligned((void *) src, 128)) { + return hvx_reduce_max_f32_a(src, num_elems); + } else { + return hvx_reduce_max_f32_u(src, num_elems); + } +} + +// Sum variants + +static inline float hvx_reduce_sum_f32_a(const uint8_t * restrict src, const int num_elems) { + HVX_Vector init_vec = Q6_V_vsplat_R(0); + assert((unsigned long) src % 128 == 0); + hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_REDUCE_SUM_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR); +} + +static inline float hvx_reduce_sum_f32_u(const uint8_t * restrict src, const int num_elems) { + HVX_Vector init_vec = Q6_V_vsplat_R(0); + hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_REDUCE_SUM_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR); +} + +static inline float hvx_reduce_sum_f32(const uint8_t * restrict src, const int num_elems) { + if (hex_is_aligned((void *) src, 128)) { + return hvx_reduce_sum_f32_a(src, num_elems); + } else { + return hvx_reduce_sum_f32_u(src, num_elems); + } +} + +// Sum of squares variants + +static inline float hvx_sum_of_squares_f32_a(const uint8_t * restrict src, const int num_elems) { + HVX_Vector init_vec = Q6_V_vsplat_R(0); + assert((uintptr_t) src % 128 == 0); + hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_SUM_SQ_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR); +} + +static inline float hvx_sum_of_squares_f32_u(const uint8_t * restrict src, const int num_elems) { + HVX_Vector init_vec = Q6_V_vsplat_R(0); + hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_SUM_SQ_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR); +} + +static inline float hvx_sum_of_squares_f32(const uint8_t * restrict src, const int num_elems) { + if (hex_is_aligned((void *) src, 128)) { + return hvx_sum_of_squares_f32_a(src, num_elems); + } else { + return hvx_sum_of_squares_f32_u(src, num_elems); + } +} + +#undef hvx_reduce_loop_body +#undef HVX_REDUCE_MAX_OP +#undef HVX_REDUCE_SUM_OP +#undef HVX_REDUCE_MAX_SCALAR +#undef HVX_REDUCE_SUM_SCALAR +#undef HVX_SUM_SQ_OP + +#include "hvx-base-defs.inc" + +#define FAST_SIGMOID_LOG2F (0x3fb8aa3b) // 1.442695022 +#define FAST_SIGMOID_C1 (0x3d009076) // 0.03138777 +#define FAST_SIGMOID_C2 (0x3e8d74bd) // 0.276281267 +#define FAST_SIGMOID_C3 (0x3f000000) // 0.5 + +static inline HVX_Vector hvx_vec_fast_sigmoid_f32(HVX_Vector v) { + v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F)); + v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3)); + + HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v)); + HVX_Vector x = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int)); + HVX_Vector xx = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x); + + HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2)); + v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F)); + + HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1)); + v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx); + v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x); + + HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1)); + HVX_Vector v3_exponent = Q6_Vw_vasl_VwR(v3, 1); + v3_exponent = Q6_Vuw_vlsr_VuwR(v3_exponent, 24); + v3_exponent = Q6_Vw_vadd_VwVw(in_int, v3_exponent); + v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24); + + HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1)); + HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4)); + + HVX_Vector res = hvx_vec_inverse_f32(v5); + res = Q6_Vqf32_vmpy_VsfVsf(v3, res); + + return Q6_Vsf_equals_Vqf32(res); +} + +static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard(HVX_Vector v, + HVX_Vector one, + HVX_Vector max_exp, + HVX_Vector min_exp) { + const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v); + const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp); + + HVX_Vector out = hvx_vec_fast_sigmoid_f32(v); + out = Q6_V_vmux_QVV(pred_max, out, one); + return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero()); +} + +static inline HVX_Vector hvx_vec_tanh_f32(HVX_Vector x) { + // tanh(x) = 2 * sigmoid(2x) - 1 + HVX_Vector two = hvx_vec_splat_f32(2.0f); + HVX_Vector one = hvx_vec_splat_f32(1.0f); + HVX_Vector x2 = Q6_Vqf32_vmpy_VsfVsf(x, two); + + HVX_Vector max_exp = hvx_vec_splat_f32(87.f); + HVX_Vector min_exp = hvx_vec_splat_f32(-87.f); + + HVX_Vector sig2x = hvx_vec_fast_sigmoid_f32_guard(Q6_Vsf_equals_Vqf32(x2), one, max_exp, min_exp); + + HVX_Vector res = Q6_Vqf32_vmpy_VsfVsf(sig2x, two); + res = Q6_Vqf32_vsub_Vqf32Vsf(res, one); + return Q6_Vsf_equals_Vqf32(res); +} + +#define hvx_sigmoid_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const HVX_Vector one = hvx_vec_splat_f32(1.f); \ + const HVX_Vector max_exp = hvx_vec_splat_f32(87.f); \ + const HVX_Vector min_exp = hvx_vec_splat_f32(-87.f); \ + \ + const uint32_t epv = 128 / sizeof(float); \ + const uint32_t nvec = n / epv; \ + const uint32_t nloe = n % epv; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + vdst[i] = hvx_vec_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp); \ + } \ + if (nloe) { \ + HVX_Vector tmp = hvx_vec_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp); \ + vec_store((void *) &vdst[i], nloe * sizeof(float), tmp); \ + } \ + } while(0) + +#define hvx_tanh_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const uint32_t epv = 128 / sizeof(float); \ + const uint32_t nvec = n / epv; \ + const uint32_t nloe = n % epv; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + vdst[i] = hvx_vec_tanh_f32(vsrc[i]); \ + } \ + if (nloe) { \ + HVX_Vector tmp = hvx_vec_tanh_f32(vsrc[i]); \ + vec_store((void *) &vdst[i], nloe * sizeof(float), tmp); \ + } \ + } while(0) + +static inline void hvx_sigmoid_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + hvx_sigmoid_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); +} + +static inline void hvx_sigmoid_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + hvx_sigmoid_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); +} + +static inline void hvx_sigmoid_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) src % 128 == 0); + hvx_sigmoid_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); +} + +static inline void hvx_sigmoid_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_sigmoid_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); +} + +static inline void hvx_tanh_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + hvx_tanh_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); +} + +#include +#include + +#include "hex-utils-defs.inc" + +#include "hvx-base-defs.inc" + +#define RSQRT_CONST 0x5f3759df // Constant for fast inverse square root calculation +#define RSQRT_ONE_HALF 0x3f000000 // 0.5 +#define RSQRT_THREE_HALVES 0x3fc00000 // 1.5 + +#if __HVX_ARCH__ < 79 +#define HVX_OP_MUL(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b)) +#else +#define HVX_OP_MUL(a, b) Q6_Vsf_vmpy_VsfVsf(a, b) +#endif + +static inline HVX_Vector hvx_vec_rsqrt_f32(HVX_Vector in_vec) { + //Algorithm : + // x2 = input*0.5 + // y = * (long *) &input + // y = 0x5f3759df - (y>>1) + // y = y*(threehalfs - x2*y*y) + + HVX_Vector rsqrtconst = Q6_V_vsplat_R(RSQRT_CONST); + HVX_Vector onehalf = Q6_V_vsplat_R(RSQRT_ONE_HALF); + HVX_Vector threehalfs = Q6_V_vsplat_R(RSQRT_THREE_HALVES); + + HVX_Vector x2, y, ypower2, temp; + + x2 = Q6_Vqf32_vmpy_VsfVsf(in_vec, onehalf); + x2 = Q6_Vqf32_vadd_Vqf32Vsf(x2, Q6_V_vzero()); + + y = Q6_Vw_vasr_VwR(in_vec, 1); + y = Q6_Vw_vsub_VwVw(rsqrtconst, y); + + // 1st iteration + ypower2 = Q6_Vqf32_vmpy_VsfVsf(y, y); + ypower2 = Q6_Vqf32_vadd_Vqf32Vsf(ypower2, Q6_V_vzero()); + temp = Q6_Vqf32_vmpy_Vqf32Vqf32(x2, ypower2); + temp = Q6_Vqf32_vsub_VsfVsf(threehalfs, Q6_Vsf_equals_Vqf32(temp)); + temp = Q6_Vqf32_vmpy_VsfVsf(y, Q6_Vsf_equals_Vqf32(temp)); + + // 2nd iteration + y = Q6_Vqf32_vadd_Vqf32Vsf(temp, Q6_V_vzero()); + ypower2 = Q6_Vqf32_vmpy_Vqf32Vqf32(y, y); + ypower2 = Q6_Vqf32_vadd_Vqf32Vsf(ypower2, Q6_V_vzero()); + temp = Q6_Vqf32_vmpy_Vqf32Vqf32(x2, ypower2); + temp = Q6_Vqf32_vsub_VsfVsf(threehalfs, Q6_Vsf_equals_Vqf32(temp)); + temp = Q6_Vqf32_vmpy_Vqf32Vqf32(y, temp); + + // 3rd iteration + y = Q6_Vqf32_vadd_Vqf32Vsf(temp, Q6_V_vzero()); + ypower2 = Q6_Vqf32_vmpy_Vqf32Vqf32(y, y); + ypower2 = Q6_Vqf32_vadd_Vqf32Vsf(ypower2, Q6_V_vzero()); + temp = Q6_Vqf32_vmpy_Vqf32Vqf32(x2, ypower2); + temp = Q6_Vqf32_vsub_VsfVsf(threehalfs, Q6_Vsf_equals_Vqf32(temp)); + temp = Q6_Vqf32_vmpy_Vqf32Vqf32(y, temp); + + return Q6_Vsf_equals_Vqf32(temp); +} + +// Compute sqrt(x) as x*inv_sqrt(x) +#define hvx_sqrt_f32_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const uint32_t nvec = n / VLEN_FP32; \ + const uint32_t nloe = n % VLEN_FP32; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + HVX_Vector inv_sqrt = hvx_vec_rsqrt_f32(vsrc[i]); \ + HVX_Vector sqrt_res = HVX_OP_MUL(inv_sqrt, vsrc[i]); \ + vdst[i] = sqrt_res; \ + } \ + if (nloe) { \ + HVX_Vector inv_sqrt = hvx_vec_rsqrt_f32(vsrc[i]); \ + HVX_Vector sqrt_res = HVX_OP_MUL(inv_sqrt, vsrc[i]); \ + vec_store((void *) &vdst[i], nloe * SIZEOF_FP32, sqrt_res); \ + } \ + } while(0) + +static inline void hvx_sqrt_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + hvx_sqrt_f32_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); +} + +static inline void hvx_sqrt_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + hvx_sqrt_f32_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); +} + +static inline void hvx_sqrt_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) src % 128 == 0); + hvx_sqrt_f32_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); +} + +static inline void hvx_sqrt_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_sqrt_f32_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); +} + +static inline void hvx_sqrt_f32(uint8_t * restrict dst, const uint8_t * restrict src, const int num_elems) { + if ((unsigned long) dst % 128 == 0) { + if ((unsigned long) src % 128 == 0) { + hvx_sqrt_f32_aa(dst, src, num_elems); + } else { + hvx_sqrt_f32_au(dst, src, num_elems); + } + } else { + if ((unsigned long) src % 128 == 0) { + hvx_sqrt_f32_ua(dst, src, num_elems); + } else { + hvx_sqrt_f32_uu(dst, src, num_elems); + } + } +} + + +#include +#include +#include +#include + +#include "hvx-base-defs.inc" +#include "hex-utils-defs.inc" + +// +// Binary operations (add, mul, sub) +// + +#define UNUSED(x) (void)(x) + +#define hvx_arith_loop_body(dst_type, src0_type, src1_type, elem_size, vec_store, vec_op) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src0_type * restrict vsrc0 = (src0_type *) src0; \ + src1_type * restrict vsrc1 = (src1_type *) src1; \ + \ + const uint32_t epv = 128 / (elem_size); \ + const uint32_t nvec = n / epv; \ + const uint32_t nloe = n % epv; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + vdst[i] = vec_op(vsrc0[i], vsrc1[i]); \ + } \ + if (nloe) { \ + HVX_Vector v = vec_op(vsrc0[i], vsrc1[i]); \ + vec_store((void *) &vdst[i], nloe * (elem_size), v); \ + } \ + } while(0) + +#if __HVX_ARCH__ < 79 + +#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b)) +#define HVX_OP_SUB_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(a, b)) +#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b)) + +#else + +#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b) +#define HVX_OP_SUB_F32(a, b) Q6_Vsf_vsub_VsfVsf(a, b) +#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b) + +#endif + +#define HVX_OP_ADD_F16(a, b) hvx_vec_add_f16_f16(a, b) +#define HVX_OP_SUB_F16(a, b) hvx_vec_sub_f16_f16(a, b) +#define HVX_OP_MUL_F16(a, b) hvx_vec_mul_f16_f16(a, b) + +// Generic macro to define alignment permutations for an op +#define DEFINE_HVX_BINARY_OP_VARIANTS(OP_NAME, OP_MACRO, ELEM_TYPE) \ +static inline void OP_NAME##_aaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) dst % 128 == 0); \ + assert((uintptr_t) src0 % 128 == 0); \ + assert((uintptr_t) src1 % 128 == 0); \ + hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ +} \ +static inline void OP_NAME##_aau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) dst % 128 == 0); \ + assert((uintptr_t) src0 % 128 == 0); \ + hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ +} \ +static inline void OP_NAME##_aua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) dst % 128 == 0); \ + assert((uintptr_t) src1 % 128 == 0); \ + hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ +} \ +static inline void OP_NAME##_auu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) dst % 128 == 0); \ + hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ +} \ +static inline void OP_NAME##_uaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) src0 % 128 == 0); \ + assert((uintptr_t) src1 % 128 == 0); \ + hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ +} \ +static inline void OP_NAME##_uau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) src0 % 128 == 0); \ + hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ +} \ +static inline void OP_NAME##_uua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) src1 % 128 == 0); \ + hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ +} \ +static inline void OP_NAME##_uuu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ +} \ + +DEFINE_HVX_BINARY_OP_VARIANTS(hvx_add_f32, HVX_OP_ADD_F32, float) +DEFINE_HVX_BINARY_OP_VARIANTS(hvx_sub_f32, HVX_OP_SUB_F32, float) +DEFINE_HVX_BINARY_OP_VARIANTS(hvx_mul_f32, HVX_OP_MUL_F32, float) + +DEFINE_HVX_BINARY_OP_VARIANTS(hvx_add_f16, HVX_OP_ADD_F16, _Float16) +DEFINE_HVX_BINARY_OP_VARIANTS(hvx_sub_f16, HVX_OP_SUB_F16, _Float16) +DEFINE_HVX_BINARY_OP_VARIANTS(hvx_mul_f16, HVX_OP_MUL_F16, _Float16) + +// Dispatcher logic +#define HVX_BINARY_DISPATCHER(OP_NAME) \ +static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, const uint32_t num_elems) { \ + if (hex_is_aligned((void *) dst, 128)) { \ + if (hex_is_aligned((void *) src0, 128)) { \ + if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems); \ + else OP_NAME##_aau(dst, src0, src1, num_elems); \ + } else { \ + if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems); \ + else OP_NAME##_auu(dst, src0, src1, num_elems); \ + } \ + } else { \ + if (hex_is_aligned((void *) src0, 128)) { \ + if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems); \ + else OP_NAME##_uau(dst, src0, src1, num_elems); \ + } else { \ + if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems); \ + else OP_NAME##_uuu(dst, src0, src1, num_elems); \ + } \ + } \ +} + +HVX_BINARY_DISPATCHER(hvx_add_f32) +HVX_BINARY_DISPATCHER(hvx_sub_f32) +HVX_BINARY_DISPATCHER(hvx_mul_f32) + +HVX_BINARY_DISPATCHER(hvx_add_f16) +HVX_BINARY_DISPATCHER(hvx_sub_f16) +HVX_BINARY_DISPATCHER(hvx_mul_f16) + +// Mul-Mul Optimized +static inline void hvx_mul_mul_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, const uint8_t * restrict src2, const uint32_t num_elems) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src0 % 128 == 0); + assert((unsigned long) src1 % 128 == 0); + assert((unsigned long) src2 % 128 == 0); + + HVX_Vector * restrict vdst = (HVX_Vector *) dst; + HVX_Vector * restrict vsrc0 = (HVX_Vector *) src0; + HVX_Vector * restrict vsrc1 = (HVX_Vector *) src1; + HVX_Vector * restrict vsrc2 = (HVX_Vector *) src2; + + const uint32_t elem_size = sizeof(float); + const uint32_t epv = 128 / elem_size; + const uint32_t nvec = num_elems / epv; + const uint32_t nloe = num_elems % epv; + + uint32_t i = 0; + + _Pragma("unroll(4)") + for (; i < nvec; i++) { + HVX_Vector v1 = HVX_OP_MUL_F32(vsrc0[i], vsrc1[i]); + vdst[i] = HVX_OP_MUL(v1, vsrc2[i]); + } + + if (nloe) { + HVX_Vector v1 = HVX_OP_MUL_F32(vsrc0[i], vsrc1[i]); + HVX_Vector v2 = HVX_OP_MUL_F32(v1, vsrc2[i]); + hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v2); + } +} + +// Scalar Operations + +#define hvx_scalar_loop_body(dst_type, src_type, elem_size, vec_store, scalar_op_macro) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const uint32_t epv = 128 / (elem_size); \ + const uint32_t nvec = n / epv; \ + const uint32_t nloe = n % epv; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + HVX_Vector v = vsrc[i]; \ + vdst[i] = scalar_op_macro(v); \ + } \ + if (nloe) { \ + HVX_Vector v = vsrc[i]; \ + v = scalar_op_macro(v); \ + vec_store((void *) &vdst[i], nloe * (elem_size), v); \ + } \ + } while(0) + +#define HVX_OP_ADD_SCALAR_F32(v) \ + ({ \ + const HVX_VectorPred pred_inf = Q6_Q_vcmp_eq_VwVw(inf, v); \ + HVX_Vector out = HVX_OP_ADD_F32(v, val_vec); \ + Q6_V_vmux_QVV(pred_inf, inf, out); \ + }) + +#define HVX_OP_MUL_SCALAR_F32(v) HVX_OP_MUL_F32(v, val_vec) +#define HVX_OP_SUB_SCALAR_F32(v) HVX_OP_SUB_F32(v, val_vec) + +#define HVX_OP_ADD_SCALAR_F16(v) \ + ({ \ + const HVX_VectorPred pred_inf = Q6_Q_vcmp_eq_VhVh(inf, v); \ + HVX_Vector out = HVX_OP_ADD_F16(v, val_vec); \ + Q6_V_vmux_QVV(pred_inf, inf, out); \ + }) + +#define HVX_OP_MUL_SCALAR_F16(v) HVX_OP_MUL_F16(v, val_vec) +#define HVX_OP_SUB_SCALAR_F16(v) HVX_OP_SUB_F16(v, val_vec) + +// Scalar Variants + +// Generic macro to define alignment permutations for an op +#define DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(OP_NAME, OP_MACRO, SPLAT_MACRO, ELEM_TYPE) \ +static inline void OP_NAME##_aa(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \ + const HVX_Vector val_vec = SPLAT_MACRO(val); \ + const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \ + assert((uintptr_t) dst % 128 == 0); \ + assert((uintptr_t) src % 128 == 0); \ + hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ +} \ +static inline void OP_NAME##_au(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \ + const HVX_Vector val_vec = SPLAT_MACRO(val); \ + const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \ + assert((uintptr_t) dst % 128 == 0); \ + hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ +} \ +static inline void OP_NAME##_ua(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \ + const HVX_Vector val_vec = SPLAT_MACRO(val); \ + const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \ + assert((uintptr_t) src % 128 == 0); \ + hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ +} \ +static inline void OP_NAME##_uu(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \ + const HVX_Vector val_vec = SPLAT_MACRO(val); \ + const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \ + hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ +} \ + +DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_add_scalar_f32, HVX_OP_ADD_SCALAR_F32, hvx_vec_splat_f32, float) +DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_sub_scalar_f32, HVX_OP_SUB_SCALAR_F32, hvx_vec_splat_f32, float) +DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_mul_scalar_f32, HVX_OP_MUL_SCALAR_F32, hvx_vec_splat_f32, float) + +DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_add_scalar_f16, HVX_OP_ADD_SCALAR_F16, hvx_vec_splat_f16, _Float16) +DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_sub_scalar_f16, HVX_OP_SUB_SCALAR_F16, hvx_vec_splat_f16, _Float16) +DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_mul_scalar_f16, HVX_OP_MUL_SCALAR_F16, hvx_vec_splat_f16, _Float16) + +// Dispatcher logic +#define HVX_BINARY_SCALAR_DISPATCHER(OP_NAME, ELEM_TYPE) \ +static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, const uint32_t num_elems) { \ + if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { \ + OP_NAME##_aa(dst, src, val, num_elems); \ + } else if (hex_is_aligned((void *) dst, 128)) { \ + OP_NAME##_au(dst, src, val, num_elems); \ + } else if (hex_is_aligned((void *) src, 128)) { \ + OP_NAME##_ua(dst, src, val, num_elems); \ + } else { \ + OP_NAME##_uu(dst, src, val, num_elems); \ + } \ +} + +HVX_BINARY_SCALAR_DISPATCHER(hvx_add_scalar_f32, float) +HVX_BINARY_SCALAR_DISPATCHER(hvx_sub_scalar_f32, float) +HVX_BINARY_SCALAR_DISPATCHER(hvx_mul_scalar_f32, float) + +HVX_BINARY_SCALAR_DISPATCHER(hvx_add_scalar_f16, _Float16) +HVX_BINARY_SCALAR_DISPATCHER(hvx_sub_scalar_f16, _Float16) +HVX_BINARY_SCALAR_DISPATCHER(hvx_mul_scalar_f16, _Float16) + +// MIN Scalar variants + +#define HVX_OP_MIN_SCALAR(v) Q6_Vsf_vmin_VsfVsf(val_vec, v) + +static inline void hvx_min_scalar_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) { + const HVX_Vector val_vec = hvx_vec_splat_f32(val); + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(float), hvx_vec_store_a, HVX_OP_MIN_SCALAR); +} + +static inline void hvx_min_scalar_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) { + const HVX_Vector val_vec = hvx_vec_splat_f32(val); + assert((unsigned long) dst % 128 == 0); + hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(float), hvx_vec_store_a, HVX_OP_MIN_SCALAR); +} + +static inline void hvx_min_scalar_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) { + const HVX_Vector val_vec = hvx_vec_splat_f32(val); + assert((unsigned long) src % 128 == 0); + hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(float), hvx_vec_store_u, HVX_OP_MIN_SCALAR); +} + +static inline void hvx_min_scalar_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) { + const HVX_Vector val_vec = hvx_vec_splat_f32(val); + hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(float), hvx_vec_store_u, HVX_OP_MIN_SCALAR); +} + +static inline void hvx_min_scalar_f32(uint8_t * restrict dst, const uint8_t * restrict src, const float val, const int num_elems) { + if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { + hvx_min_scalar_f32_aa(dst, src, val, num_elems); + } else if (hex_is_aligned((void *) dst, 128)) { + hvx_min_scalar_f32_au(dst, src, val, num_elems); + } else if (hex_is_aligned((void *) src, 128)) { + hvx_min_scalar_f32_ua(dst, src, val, num_elems); + } else { + hvx_min_scalar_f32_uu(dst, src, val, num_elems); + } +} + +// CLAMP Scalar variants + +#define HVX_OP_CLAMP_SCALAR(v) \ + ({ \ + HVX_VectorPred pred_cap_right = Q6_Q_vcmp_gt_VsfVsf(v, max_vec); \ + HVX_VectorPred pred_cap_left = Q6_Q_vcmp_gt_VsfVsf(min_vec, v); \ + HVX_Vector tmp = Q6_V_vmux_QVV(pred_cap_right, max_vec, v); \ + Q6_V_vmux_QVV(pred_cap_left, min_vec, tmp); \ + }) + +static inline void hvx_clamp_scalar_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) { + const HVX_Vector min_vec = hvx_vec_splat_f32(min); + const HVX_Vector max_vec = hvx_vec_splat_f32(max); + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(float), hvx_vec_store_a, HVX_OP_CLAMP_SCALAR); +} + +static inline void hvx_clamp_scalar_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) { + const HVX_Vector min_vec = hvx_vec_splat_f32(min); + const HVX_Vector max_vec = hvx_vec_splat_f32(max); + assert((unsigned long) dst % 128 == 0); + hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(float), hvx_vec_store_a, HVX_OP_CLAMP_SCALAR); +} + +static inline void hvx_clamp_scalar_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) { + const HVX_Vector min_vec = hvx_vec_splat_f32(min); + const HVX_Vector max_vec = hvx_vec_splat_f32(max); + assert((unsigned long) src % 128 == 0); + hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(float), hvx_vec_store_u, HVX_OP_CLAMP_SCALAR); +} + +static inline void hvx_clamp_scalar_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) { + const HVX_Vector min_vec = hvx_vec_splat_f32(min); + const HVX_Vector max_vec = hvx_vec_splat_f32(max); + hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(float), hvx_vec_store_u, HVX_OP_CLAMP_SCALAR); +} + +static inline void hvx_clamp_scalar_f32(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, const int num_elems) { + if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { + hvx_clamp_scalar_f32_aa(dst, src, min, max, num_elems); + } else if (hex_is_aligned((void *) dst, 128)) { + hvx_clamp_scalar_f32_au(dst, src, min, max, num_elems); + } else if (hex_is_aligned((void *) src, 128)) { + hvx_clamp_scalar_f32_ua(dst, src, min, max, num_elems); + } else { + hvx_clamp_scalar_f32_uu(dst, src, min, max, num_elems); + } +} + +// +// Square +// + +#define hvx_sqr_f32_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const uint32_t elem_size = sizeof(float); \ + const uint32_t epv = 128 / elem_size; \ + const uint32_t nvec = n / epv; \ + const uint32_t nloe = n % epv; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + vdst[i] = HVX_OP_MUL_F32(vsrc[i], vsrc[i]); \ + } \ + if (nloe) { \ + HVX_Vector v = HVX_OP_MUL_F32(vsrc[i], vsrc[i]); \ + vec_store((void *) &vdst[i], nloe * elem_size, v); \ + } \ + } while(0) + +static inline void hvx_sqr_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + hvx_sqr_f32_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); +} + +static inline void hvx_sqr_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + hvx_sqr_f32_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); +} + +static inline void hvx_sqr_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) src % 128 == 0); + hvx_sqr_f32_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); +} + +static inline void hvx_sqr_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_sqr_f32_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); +} + +static inline void hvx_sqr_f32(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t num_elems) { + if (hex_is_aligned((void *) dst, 128)) { + if (hex_is_aligned((void *) src, 128)) { + hvx_sqr_f32_aa(dst, src, num_elems); + } else { + hvx_sqr_f32_au(dst, src, num_elems); + } + } else { + if (hex_is_aligned((void *) src, 128)) { + hvx_sqr_f32_ua(dst, src, num_elems); + } else { + hvx_sqr_f32_uu(dst, src, num_elems); + } + } +} + +#undef HVX_OP_ADD_F32 +#undef HVX_OP_SUB_F32 +#undef HVX_OP_MUL_F32 +#undef HVX_OP_ADD_F16 +#undef HVX_OP_SUB_F16 +#undef HVX_OP_MUL_F16 +#undef hvx_arith_loop_body +#undef HVX_OP_ADD_SCALAR_F32 +#undef HVX_OP_SUB_SCALAR_F32 +#undef HVX_OP_MUL_SCALAR_F32 +#undef HVX_OP_ADD_SCALAR_F16 +#undef HVX_OP_SUB_SCALAR_F16 +#undef HVX_OP_MUL_SCALAR_F16 +#undef hvx_scalar_loop_body +#undef HVX_OP_MIN_SCALAR +#undef HVX_OP_CLAMP_SCALAR +#undef DEFINE_HVX_BINARY_OP_VARIANTS +#undef HVX_BINARY_DISPATCHER +#undef UNUSED + +#include + +#include +#include +#include +#include +#include + +#include "hvx-base-defs.inc" +#include "hex-utils-defs.inc" + +#if __HVX_ARCH__ < 79 +#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b)) +#define HVX_OP_MUL_F16(a, b) Q6_Vhf_equals_Wqf32(Q6_Wqf32_vmpy_VhfVhf(a, b)) +#else +#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b) +#define HVX_OP_MUL_F16(a, b) Q6_Vhf_vmpy_VhfVhf(a, b) +#endif + +// Compute div by scaler in f32. Requires first by expanding fp32 to fp16 and converting the result back to fp32. +static inline HVX_Vector hvx_div_mul_f16_const_using_f32(HVX_Vector vec1_hf, HVX_Vector vec2_sf_const, HVX_Vector vec_hf_one_1_0) { +#if __HVX_ARCH__ < 79 + HVX_VectorPair src_to_f32 = Q6_Wqf32_vmpy_VhfVhf(vec1_hf, vec_hf_one_1_0); + HVX_Vector src_to_f32_0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(src_to_f32)); + HVX_Vector src_to_f32_1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(src_to_f32)); +#else + HVX_VectorPair src_to_f32 = Q6_Wsf_vmpy_VhfVhf(vec1_hf, vec_hf_one_1_0); + HVX_Vector src_to_f32_0 = Q6_V_lo_W(src_to_f32); + HVX_Vector src_to_f32_1 = Q6_V_hi_W(src_to_f32); +#endif + + HVX_Vector div_f32_0 = HVX_OP_MUL_F32(src_to_f32_0, vec2_sf_const); + HVX_Vector div_f32_1 = HVX_OP_MUL_F32(src_to_f32_1, vec2_sf_const); + +#if __HVX_ARCH__ < 79 + HVX_Vector res = hvx_vec_f32_to_f16(div_f32_0, div_f32_1); +#else + HVX_Vector res = Q6_Vhf_vcvt_VsfVsf(div_f32_0, div_f32_1); +#endif + return res; +} + +// Variant for =v79 +static inline HVX_Vector hvx_vec_hybrid_div_f16(HVX_Vector vec1, HVX_Vector vec2, HVX_Vector f32_nan_inf_mask, HVX_Vector f16_nan_inf_mask, HVX_Vector vec_hf_one_1_0) { +#if __HVX_ARCH__ < 79 + // For older architectures, use f16 reciprocal to avoid NaN/-inf issues + HVX_Vector vec2_inv = hvx_vec_inverse_f16_guard(vec2, f16_nan_inf_mask); + return HVX_OP_MUL_F16(vec1, vec2_inv); +#else + return hvx_vec_div_f16_using_f32(vec1, vec2, f32_nan_inf_mask, vec_hf_one_1_0); +#endif +} + +#define hvx_div_f16_loop_body(dst_type, src0_type, src1_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src0_type * restrict vsrc0 = (src0_type *) src0; \ + src1_type * restrict vsrc1 = (src1_type *) src1; \ + \ + const HVX_Vector f32_nan_inf_mask = Q6_V_vsplat_R(0x7f800000); \ + const HVX_Vector f16_nan_inf_mask = Q6_Vh_vsplat_R(0x7c00); \ + const HVX_Vector hf_one = Q6_Vh_vsplat_R(0x3C00); \ + \ + const uint32_t nvec = n / VLEN_FP16; \ + const uint32_t nloe = n % VLEN_FP16; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + HVX_Vector res = hvx_vec_hybrid_div_f16(vsrc0[i], vsrc1[i], \ + f32_nan_inf_mask, f16_nan_inf_mask, \ + hf_one); \ + vdst[i] = res; \ + } \ + if (nloe) { \ + HVX_Vector res = hvx_vec_hybrid_div_f16(vsrc0[i], vsrc1[i], \ + f32_nan_inf_mask, f16_nan_inf_mask, \ + hf_one); \ + vec_store((void *) &vdst[i], nloe * SIZEOF_FP16, res); \ + } \ + } while(0) + +#define hvx_div_f32_loop_body(dst_type, src0_type, src1_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src0_type * restrict vsrc0 = (src0_type *) src0; \ + src1_type * restrict vsrc1 = (src1_type *) src1; \ + \ + const HVX_Vector nan_inf_mask = Q6_V_vsplat_R(0x7f800000); \ + \ + const uint32_t nvec = n / VLEN_FP32; \ + const uint32_t nloe = n % VLEN_FP32; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + HVX_Vector inv_src1 = hvx_vec_inverse_f32_guard(vsrc1[i], nan_inf_mask); \ + HVX_Vector res = HVX_OP_MUL_F32(vsrc0[i], inv_src1); \ + vdst[i] = res; \ + } \ + if (nloe) { \ + HVX_Vector inv_src1 = hvx_vec_inverse_f32_guard(vsrc1[i], nan_inf_mask); \ + HVX_Vector res = HVX_OP_MUL_F32(vsrc0[i], inv_src1); \ + vec_store((void *) &vdst[i], nloe * SIZEOF_FP32, res); \ + } \ + } while(0) + +// Generic macro to define alignment permutations for an op +#define DEFINE_HVX_DIV_OP_VARIANTS(OP_NAME, OP_LOOP_BODY) \ +static inline void OP_NAME##_aaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) dst % 128 == 0); \ + assert((uintptr_t) src0 % 128 == 0); \ + assert((uintptr_t) src1 % 128 == 0); \ + OP_LOOP_BODY(HVX_Vector, HVX_Vector, HVX_Vector, hvx_vec_store_a); \ +} \ +static inline void OP_NAME##_aau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) dst % 128 == 0); \ + assert((uintptr_t) src0 % 128 == 0); \ + OP_LOOP_BODY(HVX_Vector, HVX_Vector, HVX_UVector, hvx_vec_store_a); \ +} \ +static inline void OP_NAME##_aua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) dst % 128 == 0); \ + assert((uintptr_t) src1 % 128 == 0); \ + OP_LOOP_BODY(HVX_Vector, HVX_UVector, HVX_Vector, hvx_vec_store_a); \ +} \ +static inline void OP_NAME##_auu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) dst % 128 == 0); \ + OP_LOOP_BODY(HVX_Vector, HVX_UVector, HVX_UVector, hvx_vec_store_a); \ +} \ +static inline void OP_NAME##_uaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) src0 % 128 == 0); \ + assert((uintptr_t) src1 % 128 == 0); \ + OP_LOOP_BODY(HVX_UVector, HVX_Vector, HVX_Vector, hvx_vec_store_u); \ +} \ +static inline void OP_NAME##_uau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) src0 % 128 == 0); \ + OP_LOOP_BODY(HVX_UVector, HVX_Vector, HVX_UVector, hvx_vec_store_u); \ +} \ +static inline void OP_NAME##_uua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + assert((uintptr_t) src1 % 128 == 0); \ + OP_LOOP_BODY(HVX_UVector, HVX_UVector, HVX_Vector, hvx_vec_store_u); \ +} \ +static inline void OP_NAME##_uuu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ + OP_LOOP_BODY(HVX_UVector, HVX_UVector, HVX_UVector, hvx_vec_store_u); \ +} \ + +// Dispatcher logic +#define HVX_DIV_DISPATCHER(OP_NAME) \ +static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, const uint32_t num_elems) { \ + if (hex_is_aligned((void *) dst, 128)) { \ + if (hex_is_aligned((void *) src0, 128)) { \ + if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems); \ + else OP_NAME##_aau(dst, src0, src1, num_elems); \ + } else { \ + if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems); \ + else OP_NAME##_auu(dst, src0, src1, num_elems); \ + } \ + } else { \ + if (hex_is_aligned((void *) src0, 128)) { \ + if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems); \ + else OP_NAME##_uau(dst, src0, src1, num_elems); \ + } else { \ + if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems); \ + else OP_NAME##_uuu(dst, src0, src1, num_elems); \ + } \ + } \ +} + +DEFINE_HVX_DIV_OP_VARIANTS(hvx_div_f32, hvx_div_f32_loop_body) +DEFINE_HVX_DIV_OP_VARIANTS(hvx_div_f16, hvx_div_f16_loop_body) + +HVX_DIV_DISPATCHER(hvx_div_f32) +HVX_DIV_DISPATCHER(hvx_div_f16) + +#undef HVX_OP_MUL_F32 +#undef HVX_OP_MUL_F16 + + +#include "hvx-base-defs.inc" + +#endif /* HVX_UTILS_H */ diff --git a/ggml/src/ggml-hexagon/htp/hvx-utils.h b/ggml/src/ggml-hexagon/htp/hvx-utils.h deleted file mode 100644 index a518ad37331d..000000000000 --- a/ggml/src/ggml-hexagon/htp/hvx-utils.h +++ /dev/null @@ -1,18 +0,0 @@ -#ifndef HVX_UTILS_H -#define HVX_UTILS_H - -#include "hex-utils.h" - -#include "hvx-types.h" -#include "hvx-copy.h" -#include "hvx-scale.h" -#include "hvx-exp.h" -#include "hvx-inverse.h" -#include "hvx-reduce.h" -#include "hvx-sigmoid.h" -#include "hvx-sqrt.h" -#include "hvx-arith.h" -#include "hvx-div.h" -#include "hvx-base.h" - -#endif /* HVX_UTILS_H */ diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.inc similarity index 98% rename from ggml/src/ggml-hexagon/htp/main.c rename to ggml/src/ggml-hexagon/htp/main.inc index 49c1a15b344c..aa65c5992370 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.inc @@ -18,17 +18,17 @@ #include #include -#include "hex-utils.h" -#include "hex-dma.h" -#include "hmx-queue.h" +#include "hex-utils-defs.inc" +#include "hex-dma-defs.inc" +#include "hmx-queue-defs.inc" #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" #include "htp_iface.h" -#include "worker-pool.h" +#include "worker-pool-defs.inc" AEEResult htp_iface_open(const char * uri, remote_handle64 * handle) { struct htp_context * ctx; @@ -873,3 +873,10 @@ static void htp_packet_callback(dspqueue_t queue, int error, void * context) { vtcm_release(ctx); } + +#include "hex-dma.inc" +#include "worker-pool.inc" + +#ifdef HTP_HAS_HMX +#include "hmx-queue.inc" +#endif diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.inc similarity index 96% rename from ggml/src/ggml-hexagon/htp/matmul-ops.c rename to ggml/src/ggml-hexagon/htp/matmul-ops.inc index a0c265132c84..40731b1e3ed4 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.inc @@ -9,16 +9,211 @@ #include #include -#include "hex-dma.h" -#include "hvx-utils.h" -#include "hvx-dump.h" +#include "hex-dma-defs.inc" +#include "hvx-utils-defs.inc" +#include + +#include +#include + +#include "hex-utils-defs.inc" +#include "hvx-types-defs.inc" + +static void hvx_vec_dump_f16_n(char * pref, HVX_Vector v, uint32_t n) { + HVX_VectorAlias u = { .v = v }; + + const uint32_t n0 = n / 16; + const uint32_t n1 = n % 16; + int i = 0; + for (; i < n0; i++) { + hex_dump_f16_line(pref, u.fp16 + (16 * i), 16); + } + if (n1) { + hex_dump_f16_line(pref, u.fp16 + (16 * i), n1); + } +} + +static void hvx_vec_dump_f16(char * pref, HVX_Vector v) { + hvx_vec_dump_f16_n(pref, v, 64); +} + +static void hvx_vec_dump_f32_n(char * pref, HVX_Vector v, uint32_t n) { + HVX_VectorAlias u = { .v = v }; + + const uint32_t n0 = n / 16; + const uint32_t n1 = n % 16; + int i = 0; + for (; i < n0; i++) { + hex_dump_f32_line(pref, u.fp32 + (16 * i), 16); + } + if (n1) { + hex_dump_f32_line(pref, u.fp32 + (16 * i), n1); + } +} + +static void hvx_vec_dump_f32_hmt(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + float d[32]; + } u = { .v = v }; + + FARF(HIGH, "%s: %.6f %.6f %.6f %.6f ... %.6f %.6f %.6f %.6f ... %.6f %.6f %.6f %.6f\n", pref, u.d[0], u.d[1], + u.d[2], u.d[3], u.d[12], u.d[13], u.d[14], u.d[15], u.d[28], u.d[29], u.d[30], u.d[31]); +} + +static void hvx_vec_dump_f32(char * pref, HVX_Vector v) { + hvx_vec_dump_f32_n(pref, v, 32); +} + +static void hvx_vec_dump_int32(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int32_t d[32]; + } u = { .v = v }; + + for (int i = 0; i < 32 / 16; i++) { + hex_dump_int32_line(pref, u.d + (16 * i), 16); + } +} + +static void hvx_vec_dump_int32_hmt(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int32_t d[32]; + } u = { .v = v }; + + FARF(HIGH, "%s: %d %d %d %d ... %d %d %d %d ... %d %d %d %d\n", pref, u.d[0], u.d[1], u.d[2], u.d[3], u.d[12], + u.d[13], u.d[14], u.d[15], u.d[28], u.d[29], u.d[30], u.d[31]); +} + +static void hvx_vec_dump_int8_hmt(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int8_t d[128]; + } u = { .v = v }; + + FARF(HIGH, "%s: %d %d %d %d ... %d %d %d %d ... %d %d %d %d\n", pref, u.d[0], u.d[1], u.d[2], u.d[3], u.d[60], + u.d[61], u.d[62], u.d[63], u.d[124], u.d[125], u.d[126], u.d[127]); +} + +static void hvx_vec_dump_int8(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + int8_t d[128]; + } u = { .v = v }; + + for (int i = 0; i < 128 / 16; i++) { + hex_dump_int8_line(pref, u.d + (16 * i), 16); + } +} + +static void hvx_vec_dump_uint8(char * pref, HVX_Vector v) { + union { + HVX_Vector v; + uint8_t d[128]; + } u = { .v = v }; + + for (int i = 0; i < 128 / 16; i++) { + hex_dump_uint8_line(pref, u.d + (16 * i), 16); + } +} + +static bool hvx_vec_eq(HVX_Vector v0, HVX_Vector v1, size_t n) { + typedef union { + HVX_Vector v; + int8_t d[128]; + } U; + + U u0 = { .v = v0 }; + U u1 = { .v = v1 }; + + for (int i = 0; i < n; i++) { + if (u0.d[i] != u1.d[i]) { + return false; + } + } + + return true; +} + #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" -#include "hmx-ops.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" +// HMX operation entry-point declarations. +// Ported from htp-ops-lib/include/dsp/ops.h (renamed, benchmark kernels removed). (https://github.com/haozixu/htp-ops-lib) + +#ifndef HMX_OPS_H +#define HMX_OPS_H + +#include +#include + +#include "htp-ops-defs.inc" + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct { + float *dst; + const float *activation; + const __fp16 *permuted_weight; + int m; + int k; + int n; + int act_stride; + int weight_stride; + int dst_stride; + int ne02; + int ne03; + int ne12; + int ne13; + size_t src0_nb2; + size_t src0_nb3; + size_t src1_nb2; + size_t src1_nb3; + size_t dst_nb2; + size_t dst_nb3; +} hmx_matmul_w16a32_batched_params_t; + +// HMX matrix multiplication — tile-permuted FP16 weights, FP32 activation/output +// act_stride: activation row stride in elements (= k for contiguous, or +// nb[1]/sizeof(float) for permuted tensors like attention Q). +// weight_stride: weight row stride in elements (= k for compact weights, or +// nb[1]/sizeof(__fp16) for permuted KV-cache views used by QK). +int hmx_mat_mul_permuted_w16a32(struct htp_context *ctx, + float *restrict dst, + const float *activation, + const __fp16 *permuted_weight, + int m, int k, int n, + int act_stride, + int weight_stride); + +// Batched F16 wrapper over hmx_mat_mul_permuted_w16a32. +// Batch semantics match ggml_mul_mat(): src0 broadcasts to src1 in dims 2/3. +int hmx_mat_mul_permuted_w16a32_batched(struct htp_context *ctx, + const hmx_matmul_w16a32_batched_params_t *params); + +// HMX matrix multiplication — tile-permuted quantised weights (Q4_0/Q8_0/IQ4_NL) +int hmx_mat_mul_permuted_qk_0_d16a32(struct htp_context *ctx, + float *restrict dst, + const float *activation, + const uint8_t *permuted_weight, + int m, int k, int n, + int weight_type); + +// HMX flash attention +int hmx_flash_attn_ext(struct htp_ops_context * octx); + +#ifdef __cplusplus +} +#endif + +#endif // HMX_OPS_H + #define MM_SPAD_SRC0_NROWS 16 #define MM_SPAD_SRC1_NROWS 16 diff --git a/ggml/src/ggml-hexagon/htp/repeat-ops.c b/ggml/src/ggml-hexagon/htp/repeat-ops.inc similarity index 97% rename from ggml/src/ggml-hexagon/htp/repeat-ops.c rename to ggml/src/ggml-hexagon/htp/repeat-ops.inc index a6f2f0ed5f3f..0cbee33f0ec3 100644 --- a/ggml/src/ggml-hexagon/htp/repeat-ops.c +++ b/ggml/src/ggml-hexagon/htp/repeat-ops.inc @@ -7,13 +7,13 @@ #include -#include "hvx-utils.h" +#include "hvx-utils-defs.inc" #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" struct htp_repeat_context { struct htp_ops_context * octx; diff --git a/ggml/src/ggml-hexagon/htp/rope-ops.c b/ggml/src/ggml-hexagon/htp/rope-ops.inc similarity index 94% rename from ggml/src/ggml-hexagon/htp/rope-ops.c rename to ggml/src/ggml-hexagon/htp/rope-ops.inc index 1d8b0796bc9e..6d21818345cb 100644 --- a/ggml/src/ggml-hexagon/htp/rope-ops.c +++ b/ggml/src/ggml-hexagon/htp/rope-ops.inc @@ -8,17 +8,50 @@ #include #include -#include "hex-dma.h" -#include "hvx-utils.h" -#include "hex-fastdiv.h" +#include "hex-dma-defs.inc" +#include "hvx-utils-defs.inc" +// See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1. +// Precompute mp (m' in the paper) and L such that division +// can be computed using a multiply (high 32b of 64b result) +// and a shift: +// +// n/d = (mulhi(n, mp) + n) >> L; +struct fastdiv_values { + uint32_t mp; + uint32_t l; +}; + +static inline struct fastdiv_values init_fastdiv_values(uint32_t d) { + struct fastdiv_values result = { 0, 0 }; + // compute L = ceil(log2(d)); + while (result.l < 32 && ((uint32_t) 1 << result.l) < d) { + ++(result.l); + } + + result.mp = (uint32_t) (((uint64_t) 1 << 32) * (((uint64_t) 1 << result.l) - d) / d + 1); + return result; +} + +static inline uint32_t fastdiv(uint32_t n, const struct fastdiv_values * vals) { + // Compute high 32 bits of n * mp + const uint32_t hi = (uint32_t) (((uint64_t) n * vals->mp) >> 32); // mulhi(n, mp) + // add n, apply bit shift + return (hi + n) >> vals->l; +} + +static inline uint32_t fastmodulo(uint32_t n, uint32_t d, const struct fastdiv_values * vals) { + return n - fastdiv(n, vals) * d; +} + + #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" -// Redefined the types GGML_ROPE_TYPE_NORMAL & GGML_ROPE_TYPE_NEOX as we can't include ggml.h +// Redefined the types GGML_ROPE_TYPE_NORMAL & GGML_ROPE_TYPE_NEOX as we can't include ggml.h.inc #define HTP_ROPE_TYPE_NORMAL 0 #define HTP_ROPE_TYPE_NEOX 2 diff --git a/ggml/src/ggml-hexagon/htp/set-rows-ops.c b/ggml/src/ggml-hexagon/htp/set-rows-ops.inc similarity index 97% rename from ggml/src/ggml-hexagon/htp/set-rows-ops.c rename to ggml/src/ggml-hexagon/htp/set-rows-ops.inc index 0def7b408bfe..d16904af217d 100644 --- a/ggml/src/ggml-hexagon/htp/set-rows-ops.c +++ b/ggml/src/ggml-hexagon/htp/set-rows-ops.inc @@ -8,14 +8,14 @@ #include #include -#include "hex-dma.h" -#include "hvx-utils.h" +#include "hex-dma-defs.inc" +#include "hvx-utils-defs.inc" #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" #define set_rows_preamble \ const uint32_t ne00 = octx->src[0]->ne[0]; \ diff --git a/ggml/src/ggml-hexagon/htp/softmax-ops.c b/ggml/src/ggml-hexagon/htp/softmax-ops.inc similarity index 92% rename from ggml/src/ggml-hexagon/htp/softmax-ops.c rename to ggml/src/ggml-hexagon/htp/softmax-ops.inc index d78bcc0eb24e..cb7b459d233d 100644 --- a/ggml/src/ggml-hexagon/htp/softmax-ops.c +++ b/ggml/src/ggml-hexagon/htp/softmax-ops.inc @@ -8,15 +8,48 @@ #include #include -#include "hex-dma.h" -#include "hvx-utils.h" -#include "hex-fastdiv.h" +#include "hex-dma-defs.inc" +#include "hvx-utils-defs.inc" +// See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1. +// Precompute mp (m' in the paper) and L such that division +// can be computed using a multiply (high 32b of 64b result) +// and a shift: +// +// n/d = (mulhi(n, mp) + n) >> L; +struct fastdiv_values { + uint32_t mp; + uint32_t l; +}; + +static inline struct fastdiv_values init_fastdiv_values(uint32_t d) { + struct fastdiv_values result = { 0, 0 }; + // compute L = ceil(log2(d)); + while (result.l < 32 && ((uint32_t) 1 << result.l) < d) { + ++(result.l); + } + + result.mp = (uint32_t) (((uint64_t) 1 << 32) * (((uint64_t) 1 << result.l) - d) / d + 1); + return result; +} + +static inline uint32_t fastdiv(uint32_t n, const struct fastdiv_values * vals) { + // Compute high 32 bits of n * mp + const uint32_t hi = (uint32_t) (((uint64_t) n * vals->mp) >> 32); // mulhi(n, mp) + // add n, apply bit shift + return (hi + n) >> vals->l; +} + +static inline uint32_t fastmodulo(uint32_t n, uint32_t d, const struct fastdiv_values * vals) { + return n - fastdiv(n, vals) * d; +} + + #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" #define htp_softmax_preamble3 \ const uint32_t ne00 = src0->ne[0]; \ diff --git a/ggml/src/ggml-hexagon/htp/solve-tri-ops.c b/ggml/src/ggml-hexagon/htp/solve-tri-ops.inc similarity index 98% rename from ggml/src/ggml-hexagon/htp/solve-tri-ops.c rename to ggml/src/ggml-hexagon/htp/solve-tri-ops.inc index ae8e1a50495f..c51bc69a34e2 100644 --- a/ggml/src/ggml-hexagon/htp/solve-tri-ops.c +++ b/ggml/src/ggml-hexagon/htp/solve-tri-ops.inc @@ -5,11 +5,11 @@ #include #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "hvx-types.h" -#include "hvx-utils.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "hvx-types-defs.inc" +#include "hvx-utils-defs.inc" struct htp_solve_tri_context { struct htp_ops_context * octx; diff --git a/ggml/src/ggml-hexagon/htp/ssm-conv.c b/ggml/src/ggml-hexagon/htp/ssm-conv.inc similarity index 98% rename from ggml/src/ggml-hexagon/htp/ssm-conv.c rename to ggml/src/ggml-hexagon/htp/ssm-conv.inc index a28fd03e9781..1f3b94d30c80 100644 --- a/ggml/src/ggml-hexagon/htp/ssm-conv.c +++ b/ggml/src/ggml-hexagon/htp/ssm-conv.inc @@ -13,12 +13,12 @@ #include #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "hex-dma.h" -#include "htp-ops.h" -#include "htp-ops.h" -#include "hvx-utils.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "hex-dma-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" +#include "hvx-utils-defs.inc" #define htp_ssm_conv_tensors_preamble \ const struct htp_tensor * restrict src0 = octx->src[0]; \ diff --git a/ggml/src/ggml-hexagon/htp/sum-rows-ops.c b/ggml/src/ggml-hexagon/htp/sum-rows-ops.inc similarity index 95% rename from ggml/src/ggml-hexagon/htp/sum-rows-ops.c rename to ggml/src/ggml-hexagon/htp/sum-rows-ops.inc index 874c41ab2ac7..c81f299572ff 100644 --- a/ggml/src/ggml-hexagon/htp/sum-rows-ops.c +++ b/ggml/src/ggml-hexagon/htp/sum-rows-ops.inc @@ -8,14 +8,14 @@ #include #include -#include "hex-dma.h" -#include "hvx-utils.h" +#include "hex-dma-defs.inc" +#include "hvx-utils-defs.inc" #define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" #define sum_rows_preamble \ const struct htp_tensor *src0 = octx->src[0]; \ diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.c b/ggml/src/ggml-hexagon/htp/unary-ops.c deleted file mode 100644 index 819cdc49bd9b..000000000000 --- a/ggml/src/ggml-hexagon/htp/unary-ops.c +++ /dev/null @@ -1,556 +0,0 @@ -#pragma clang diagnostic ignored "-Wunused-variable" -#pragma clang diagnostic ignored "-Wunused-function" -#pragma clang diagnostic ignored "-Wunused-but-set-variable" - -#include -#include - -#include -#include - -#include "hex-dma.h" -#include "hvx-exp.h" -#include "hvx-sigmoid.h" -#include "hvx-utils.h" - -#define GGML_COMMON_DECL_C -#include "ggml-common.h" -#include "htp-ctx.h" -#include "htp-ops.h" -#include "htp-ops.h" - -struct htp_unary_context { - struct htp_ops_context * octx; - - // Precomputed values - const uint8_t * data_src0; - uint8_t * data_dst; - - size_t src0_data_row_size; // actual data bytes per row - size_t dst_data_row_size; // actual data bytes per row - - size_t src0_row_size_aligned; - size_t dst_row_size_aligned; - - size_t src0_spad_half_size; - size_t dst_spad_half_size; - - uint32_t block; - uint32_t src0_nrows; - uint32_t src0_nrows_per_thread; - uint32_t nc; -}; - -// Convert flat row index to DDR byte offset using the tensor's actual strides. -// ir = i1 + ne1*(i2 + ne2*i3) => offset = i1*nb1 + i2*nb2 + i3*nb3 -static inline size_t unary_row_offset(uint32_t ir, - uint32_t ne1, uint32_t ne2, - size_t nb1, size_t nb2, size_t nb3) { - const uint32_t i1 = ir % ne1; - const uint32_t i2 = (ir / ne1) % ne2; - const uint32_t i3 = ir / (ne1 * ne2); - return i1 * nb1 + i2 * nb2 + i3 * nb3; -} -// Safe DMA block size from row `ir`: clamp to the tighter dim-1 slice -// boundary of src and dst so the nb1 stride stays valid for all rows. -static inline uint32_t unary_block_size(uint32_t ir, - uint32_t end_row, - uint32_t block, - bool src_contig, - bool dst_contig, - uint32_t src_ne1, - uint32_t dst_ne1) { - uint32_t limit = MIN(block, end_row - ir); - - if (!src_contig) { - const uint32_t src_slice_end = (ir / src_ne1 + 1) * src_ne1; - limit = MIN(limit, src_slice_end - ir); - } - - if (!dst_contig) { - const uint32_t dst_slice_end = (ir / dst_ne1 + 1) * dst_ne1; - limit = MIN(limit, dst_slice_end - ir); - } - - return limit; -} - -#define htp_unary_preamble \ - const uint32_t ne00 = src->ne[0]; \ - const uint32_t ne01 = src->ne[1]; \ - const uint32_t ne02 = src->ne[2]; \ - const uint32_t ne03 = src->ne[3]; \ - \ - const uint32_t ne0 = dst->ne[0]; \ - const uint32_t ne1 = dst->ne[1]; \ - const uint32_t ne2 = dst->ne[2]; \ - const uint32_t ne3 = dst->ne[3]; \ - \ - const uint32_t nb00 = src->nb[0]; \ - const uint32_t nb01 = src->nb[1]; \ - const uint32_t nb02 = src->nb[2]; \ - const uint32_t nb03 = src->nb[3]; \ - \ - const uint32_t nb0 = dst->nb[0]; \ - const uint32_t nb1 = dst->nb[1]; \ - const uint32_t nb2 = dst->nb[2]; \ - const uint32_t nb3 = dst->nb[3]; - -static void hvx_fast_rms_norm_f32(const uint8_t * restrict src, - uint8_t * restrict dst, - uint8_t * restrict pad, - const int num_elems, - float epsilon) { - (void)pad; - - const HVX_Vector * restrict v_src = (HVX_Vector *) src; - HVX_Vector * restrict v_dst = (HVX_Vector *) dst; - - const int nvec = num_elems / VLEN_FP32; // number of full vectors - const int nloe = num_elems % VLEN_FP32; // leftover elements - - // Compute sum of squares for full vectors - HVX_Vector sum_v = Q6_V_vsplat_R(0x00000000); - HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon); - - #pragma unroll(4) - for (int i = 0; i < nvec; i++) { - HVX_Vector v1 = v_src[i]; - HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1); - sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2); - } - - // Handle tail elements using vectorized ops with masking - if (nloe > 0) { - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); - HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]); - HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1); - sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2); - } - - // Reduce HVX sum - sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v)); - - HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems); - HVX_Vector denom_v = hvx_vec_inverse_f32(t_v); - HVX_Vector mean_v = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v); - HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v); - - // Scale full vectors - HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v)); - - #pragma unroll(4) - for (int i = 0; i < nvec; i++) { - HVX_Vector v1 = v_src[i]; - HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v); - v_dst[i] = Q6_Vsf_equals_Vqf32(v2); - } - - // Handle tail elements using vectorized ops with masking - if (nloe > 0) { - - HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); - HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]); - HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v); - HVX_Vector result = Q6_Vsf_equals_Vqf32(v2); - - // Store with masking to avoid overwriting memory beyond the tensor - hvx_vec_store_a(&v_dst[nvec], nloe * 4, result); - } -} - -static void scale_f32(const float * restrict src, - float * restrict dst, - uint8_t * restrict spad, - const uint32_t num_rows, - const uint32_t row_elems, - const size_t row_size, - int32_t * op_params) { - float scale = 0.f; - float bias = 0.f; - memcpy(&scale, &op_params[0], sizeof(float)); - memcpy(&bias, &op_params[1], sizeof(float)); - - for (uint32_t ir = 0; ir < num_rows; ir++) { - const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); - uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); - - hvx_scale_offset_f32_aa((uint8_t *) dst_local, (const uint8_t *) src_local, row_elems, scale, bias); - } -} - -static void rms_norm_f32(const float * restrict src, - float * restrict dst, - uint8_t * restrict spad, - const uint32_t num_rows, - const uint32_t row_elems, - const size_t row_size, - int32_t * op_params) { - float epsilon = 0.f; - memcpy(&epsilon, op_params, sizeof(float)); - - for (uint32_t ir = 0; ir < num_rows; ir++) { - const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); - uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); - - hvx_fast_rms_norm_f32((const uint8_t *) src_local, (uint8_t *) dst_local, spad, row_elems, epsilon); - } -} - -static void sqr_f32(const float * restrict src, - float * restrict dst, - uint8_t * restrict spad, - const uint32_t num_rows, - const uint32_t row_elems, - const size_t row_size, - int32_t * op_params) { - - for (uint32_t ir = 0; ir < num_rows; ir++) { - const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); - uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); - - hvx_sqr_f32_aa((uint8_t *) dst_local, (const uint8_t *) src_local, row_elems); - } -} - -static void sqrt_f32(const float * restrict src, - float * restrict dst, - uint8_t * restrict spad, - const uint32_t num_rows, - const uint32_t row_elems, - const size_t row_size, - int32_t * op_params) { - - for (uint32_t ir = 0; ir < num_rows; ir++) { - const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); - uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); - - hvx_sqrt_f32_aa((uint8_t *) dst_local, (const uint8_t *) src_local, row_elems); - } -} - -static void neg_f32(const float * restrict src, - float * restrict dst, - uint8_t * restrict spad, - const uint32_t num_rows, - const uint32_t row_elems, - const size_t row_size, - int32_t * op_params) { - - for (uint32_t ir = 0; ir < num_rows; ir++) { - const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); - uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); - - hvx_scale_f32_aa(dst_local, src_local, row_elems, -1.0f); - } -} - -static void exp_f32(const float * restrict src, - float * restrict dst, - uint8_t * restrict spad, - const uint32_t num_rows, - const uint32_t row_elems, - const size_t row_size, - int32_t * op_params) { - - for (uint32_t ir = 0; ir < num_rows; ir++) { - const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); - uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); - - hvx_exp_f32(dst_local, src_local, row_elems, false); - } -} - -static void sigmoid_f32(const float * restrict src, - float * restrict dst, - uint8_t * restrict spad, - const uint32_t num_rows, - const uint32_t row_elems, - const size_t row_size, - int32_t * op_params) { - - for (uint32_t ir = 0; ir < num_rows; ir++) { - const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); - uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); - - hvx_sigmoid_f32_aa(dst_local, src_local, row_elems); - } -} - -static void softplus_f32(const float * restrict src, - float * restrict dst, - uint8_t * restrict spad, - const uint32_t num_rows, - const uint32_t row_elems, - const size_t row_size, - int32_t * op_params) { - // softplus(x) = log(1 + exp(x)) - // Match CPU reference: ggml_compute_softplus_f32() in ggml-impl.h - for (uint32_t ir = 0; ir < num_rows; ir++) { - const float * restrict src_f = (const float *)((const uint8_t *)src + (ir * row_size)); - float * restrict dst_f = (float *)((uint8_t *)dst + (ir * row_size)); - - for (uint32_t i = 0; i < row_elems; i++) { - float x = src_f[i]; - // For x > 20: softplus(x) ≈ x (avoids exp overflow) - dst_f[i] = (x > 20.0f) ? x : logf(1.0f + expf(x)); - } - } -} - -static void unary_job_f32_per_thread(unsigned int nth, unsigned int ith, void * data) { - const struct htp_unary_context * uctx = (const struct htp_unary_context *) data; - struct htp_ops_context * octx = uctx->octx; - const struct htp_tensor * src = octx->src[0]; - const struct htp_tensor * dst = octx->dst; - - htp_unary_preamble; - - int htp_op = octx->op; - int32_t * op_params = octx->op_params; - uint32_t src0_nrows_per_thread = uctx->src0_nrows_per_thread; - - const size_t src0_data_row_size = uctx->src0_data_row_size; - const size_t dst_data_row_size = uctx->dst_data_row_size; - - const size_t src0_row_size_aligned = uctx->src0_row_size_aligned; - const size_t dst_row_size_aligned = uctx->dst_row_size_aligned; - - const uint32_t src0_nrows = uctx->src0_nrows; - const uint32_t src0_start_row = src0_nrows_per_thread * ith; - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); - - // no work for this thread - if (src0_start_row >= src0_end_row) { - return; - } - - uint64_t t1, t2; - t1 = HAP_perf_get_qtimer_count(); - - const uint8_t * restrict data_src = uctx->data_src0; - uint8_t * restrict data_dst = uctx->data_dst; - - uint8_t * src0_spad_data = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread); - uint8_t * dst_spad_data = octx->dst_spad.data + (ith * octx->dst_spad.size_per_thread); - - size_t src0_spad_half_size = uctx->src0_spad_half_size; - size_t dst_spad_half_size = uctx->dst_spad_half_size; - - // Non-contiguous tensors have gaps at dim-2/3 boundaries that a single-stride - // 2D DMA descriptor cannot span. Clamp BLOCK to ne1 (one dim-1 slice) so every - // transfer stays within a nb1-uniform region. Skipped for contiguous tensors. - const bool src0_contig = (nb02 == (size_t)ne01 * nb01) && - (nb03 == (size_t)ne02 * nb02); - const bool dst_contig = (nb2 == (size_t)ne1 * nb1) && - (nb3 == (size_t)ne2 * nb2); - const uint32_t src0_max_block = src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); - const uint32_t dst_max_block = dst_contig ? uctx->block : MIN((uint32_t)uctx->block, ne1); - const uint32_t BLOCK = MIN(src0_max_block, dst_max_block); - if (BLOCK == 0) { - FARF(ERROR, "unary-f32 : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", - octx->src0_spad.size_per_thread, src0_row_size_aligned); - return; - } - - dma_queue * dma_queue = octx->ctx->dma[ith]; - - for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; spad_idx++) { - const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, ne1); - - // Dummy DMA transation for sequencing (interleaving dst,src,dst,...) - dma_queue_push(dma_queue, - dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)), - nb1, dst_row_size_aligned, dst_data_row_size, 0); - - const size_t src0_off = unary_row_offset(ir, ne01, ne02, nb01, nb02, nb03); - dma_queue_push(dma_queue, - dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src + src0_off), - src0_row_size_aligned, nb01, src0_data_row_size, block_size); - ir += block_size; - } - - for (uint32_t ir = src0_start_row; ir < src0_end_row; ) { - const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, ne1); - - float * dst_spad = (float *) dma_queue_pop(dma_queue).src; - float * src0_spad = (float *) dma_queue_pop(dma_queue).dst; - - // Process block in VTCM - switch (htp_op) { - case HTP_OP_RMS_NORM: - rms_norm_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); - break; - case HTP_OP_SCALE: - scale_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); - break; - case HTP_OP_SQR: - sqr_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); - break; - case HTP_OP_SQRT: - sqrt_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); - break; - case HTP_OP_UNARY_NEG: - neg_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); - break; - case HTP_OP_UNARY_EXP: - exp_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); - break; - case HTP_OP_UNARY_SIGMOID: - sigmoid_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); - break; - case HTP_OP_UNARY_SOFTPLUS: - softplus_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); - break; - default: - break; - } - - const size_t dst_off = unary_row_offset(ir, ne1, ne2, nb1, nb2, nb3); - dma_queue_push(dma_queue, - dma_make_ptr(data_dst + dst_off, dst_spad), - nb1, dst_row_size_aligned, dst_data_row_size, block_size); - - // prefetch N+2 loop iteration if any - const uint32_t next_ir = ir + block_size; - if (next_ir < src0_end_row) { - const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, ne1); - const uint32_t pref_ir = next_ir + next_block_size; - if (pref_ir < src0_end_row) { - const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, ne1); - const size_t src0_pref_off = unary_row_offset(pref_ir, ne01, ne02, nb01, nb02, nb03); - dma_queue_push(dma_queue, - dma_make_ptr(src0_spad, data_src + src0_pref_off), - src0_row_size_aligned, nb01, src0_data_row_size, pref_block_size); - } - } - ir += block_size; - } - - dma_queue_flush(dma_queue); - - t2 = HAP_perf_get_qtimer_count(); - - FARF(HIGH, "unary-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n", ith, nth, src->ne[0], - src->ne[1], src->ne[2], src->ne[3], src0_start_row, src0_end_row, dst->ne[0], dst->ne[1], dst->ne[2], - dst->ne[3], (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1)); -} - -static int execute_op_unary_f32(struct htp_ops_context * octx) { - int err = HTP_STATUS_OK; - - const struct htp_tensor * src0 = octx->src[0]; - const struct htp_tensor * dst = octx->dst; - - const char * op_type = NULL; - - switch (octx->op) { - case HTP_OP_RMS_NORM: - op_type = "rmsnorm-f32"; - break; - case HTP_OP_SCALE: - op_type = "scale-f32"; - break; - case HTP_OP_SQR: - op_type = "sqr-f32"; - break; - case HTP_OP_SQRT: - op_type = "sqrt-f32"; - break; - case HTP_OP_UNARY_NEG: - op_type = "neg-f32"; - break; - case HTP_OP_UNARY_EXP: - op_type = "exp-f32"; - break; - case HTP_OP_UNARY_SIGMOID: - op_type = "sigmoid-f32"; - break; - case HTP_OP_UNARY_SOFTPLUS: - op_type = "softplus-f32"; - break; - - default: - FARF(ERROR, "Unsupported unary Op %u\n", octx->op); - return HTP_STATUS_NO_SUPPORT; - } - - const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3]; - const uint32_t n_threads = MIN(octx->n_threads, src0_nrows); - - const size_t src0_data_row_size = src0->ne[0] * sizeof(float); - const size_t dst_data_row_size = dst->ne[0] * sizeof(float); - - const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, VLEN); - const size_t dst_row_size_aligned = hex_round_up(dst_data_row_size, VLEN); - - // VTCM scratchpads for all tensors - // N rows per thread, padded to HVX vector size - // Double buffering requires 2x size per buffer - - size_t spad_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned); - size_t vtcm_row_per_thread = (octx->ctx->vtcm_size)/ (n_threads * spad_size_per_row); - - // Make sure the reserved vtcm size is sufficient - if (vtcm_row_per_thread == 0) { - FARF(ERROR, "unary-%s : current VTCM reservation %zu is too small, needed %zu\n", op_type, octx->ctx->vtcm_size, - spad_size_per_row * n_threads); - return HTP_STATUS_VTCM_TOO_SMALL; - } - - octx->src0_spad.size_per_thread = src0_row_size_aligned * vtcm_row_per_thread * 2; - octx->dst_spad.size_per_thread = dst_row_size_aligned * vtcm_row_per_thread * 2; - - octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread; - octx->dst_spad.size = n_threads * octx->dst_spad.size_per_thread; - - octx->src0_spad.data = octx->ctx->vtcm_base; - octx->dst_spad.data = octx->src0_spad.data + octx->src0_spad.size; - - FARF(HIGH, "%s: (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-spad-size %u src1-spad-size %u dst-spad-size %u\n", op_type, - src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], - octx->src0_spad.size, octx->src1_spad.size, octx->dst_spad.size); - - if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) { - struct htp_unary_context uctx = { - .octx = octx, - .src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads, - .src0_nrows = src0_nrows, - - .data_src0 = (const uint8_t *)src0->data, - .data_dst = (uint8_t *)dst->data, - - .src0_data_row_size = src0_data_row_size, - .dst_data_row_size = dst_data_row_size, - - .src0_row_size_aligned = src0_row_size_aligned, - .dst_row_size_aligned = dst_row_size_aligned, - - .src0_spad_half_size = octx->src0_spad.size_per_thread / 2, - .dst_spad_half_size = octx->dst_spad.size_per_thread / 2, - - .block = (octx->src0_spad.size_per_thread / 2) / src0_row_size_aligned, - .nc = src0->ne[0], - }; - - worker_pool_run_func(octx->ctx->worker_pool, unary_job_f32_per_thread, &uctx, n_threads); - } - - return err; -} - -int op_unary(struct htp_ops_context * octx) { - int err = HTP_STATUS_OK; - - switch (octx->src[0]->type) { - case HTP_TYPE_F32: - err = execute_op_unary_f32(octx); - break; - - default: - err = HTP_STATUS_NO_SUPPORT; - break; - } - - return err; -} diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.inc b/ggml/src/ggml-hexagon/htp/unary-ops.inc new file mode 100644 index 000000000000..3507f51379bf --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/unary-ops.inc @@ -0,0 +1,1205 @@ +#pragma clang diagnostic ignored "-Wunused-variable" +#pragma clang diagnostic ignored "-Wunused-function" +#pragma clang diagnostic ignored "-Wunused-but-set-variable" + +#include +#include + +#include +#include + +#include "hex-dma-defs.inc" +#include +#include +#include + +#include "hvx-base-defs.inc" +#include +#include + +#include "hvx-base-defs.inc" + +#define IEEE_VSF_EXPLEN (8) +#define IEEE_VSF_EXPBIAS (127) +#define IEEE_VSF_EXPMASK (0xFF) +#define IEEE_VSF_MANTLEN (23) +#define IEEE_VSF_MANTMASK (0x7FFFFF) +#define IEEE_VSF_MIMPMASK (0x800000) + +static inline HVX_Vector hvx_vec_truncate_f32(HVX_Vector in_vec) { + HVX_Vector mask_mant_v = Q6_V_vsplat_R(IEEE_VSF_MANTMASK); + HVX_Vector mask_impl_v = Q6_V_vsplat_R(IEEE_VSF_MIMPMASK); + HVX_Vector const_zero_v = Q6_V_vzero(); + + HVX_VectorPred q_negative = Q6_Q_vcmp_gt_VwVw(const_zero_v, in_vec); + + HVX_Vector expval_v = in_vec >> IEEE_VSF_MANTLEN; + expval_v &= IEEE_VSF_EXPMASK; + expval_v -= IEEE_VSF_EXPBIAS; + + // negative exp == fractional value + HVX_VectorPred q_negexp = Q6_Q_vcmp_gt_VwVw(const_zero_v, expval_v); + + HVX_Vector rshift_v = IEEE_VSF_MANTLEN - expval_v; // fractional bits - exp shift + + HVX_Vector mant_v = in_vec & mask_mant_v; // obtain mantissa + HVX_Vector vout = Q6_Vw_vadd_VwVw(mant_v, mask_impl_v); // add implicit 1.0 + + vout = Q6_Vw_vasr_VwVw(vout, rshift_v); // shift to obtain truncated integer + vout = Q6_V_vmux_QVV(q_negexp, const_zero_v, vout); // expval<0 -> 0 + + HVX_Vector neg_vout = -vout; + + vout = Q6_V_vmux_QVV(q_negative, neg_vout, vout); // handle negatives + + return (vout); +} + +static inline HVX_Vector hvx_vec_floor_f32(HVX_Vector in_vec) { + HVX_Vector mask_mant_v = Q6_V_vsplat_R(IEEE_VSF_MANTMASK); + HVX_Vector mask_impl_v = Q6_V_vsplat_R(IEEE_VSF_MIMPMASK); + HVX_Vector const_mnlen_v = Q6_V_vsplat_R(IEEE_VSF_MANTLEN); + HVX_Vector const_zero_v = Q6_V_vzero(); + HVX_Vector const_negone_v = Q6_V_vsplat_R(0xbf800000); // -1 IEEE vsf + + HVX_VectorPred q_negative = Q6_Q_vcmp_gt_VwVw(const_zero_v, in_vec); + + HVX_Vector expval_v = in_vec >> IEEE_VSF_MANTLEN; + expval_v &= IEEE_VSF_EXPMASK; + expval_v -= IEEE_VSF_EXPBIAS; + + HVX_VectorPred q_negexp = Q6_Q_vcmp_gt_VwVw(const_zero_v, expval_v); + HVX_VectorPred q_expltmn = Q6_Q_vcmp_gt_VwVw(const_mnlen_v, expval_v); + HVX_VectorPred q_negexp_pos = Q6_Q_vcmp_gtand_QVwVw(q_negexp, in_vec, const_zero_v); + HVX_VectorPred q_negexp_neg = Q6_Q_vcmp_gtand_QVwVw(q_negexp, const_zero_v, in_vec); + + // if expval < 0 (q_negexp) // <0, floor is 0 + // if vin > 0 + // floor = 0 + // if vin < 0 + // floor = -1 + // if expval < mant_len (q_expltmn) // >0, but fraction may exist + // get sign (q_negative) + // mask >> expval // fraction bits to mask off + // vout = ~(mask) // apply mask to remove fraction + // if (qneg) // negative floor is one less (more, sign bit for neg) + // vout += ((impl_mask) >> expval) + // if (mask && vin) + // vout = vin + // else // already an integer + // ; // no change + + // compute floor + mask_mant_v >>= expval_v; + HVX_Vector neg_addin_v = mask_impl_v >> expval_v; + HVX_Vector vout_neg_addin = Q6_Vw_vadd_VwVw(in_vec, neg_addin_v); + HVX_Vector vout = Q6_V_vmux_QVV(q_negative, vout_neg_addin, in_vec); + + HVX_Vector mask_chk_v = Q6_V_vand_VV(in_vec, mask_mant_v); // chk if bits set + HVX_VectorPred q_integral = Q6_Q_vcmp_eq_VwVw(const_zero_v, mask_chk_v); + + HVX_Vector not_mask_v = Q6_V_vnot_V(mask_mant_v); // frac bits to clear + HVX_Vector vfrfloor_v = Q6_V_vand_VV(vout, not_mask_v); // clear frac bits + + vout = in_vec; + vout = Q6_V_vmux_QVV(q_expltmn, vfrfloor_v, vout); // expval0 -> 0 + vout = Q6_V_vmux_QVV(q_negexp_neg, const_negone_v, vout); // expval<0 x<0 -> -1 + + return vout; +} + + + +#define EXP_COEFF_5 (0x39506967) // 0.000198757 = 1/(7!) +#define EXP_COEFF_4 (0x3AB743CE) // 0.0013982 = 1/(6!) +#define EXP_COEFF_3 (0x3C088908) // 0.00833345 = 1/(5!) +#define EXP_COEFF_2 (0x3D2AA9C1) // 0.416658 = 1/(4!) +#define EXP_COEFF_1 (0x3E2AAAAA) // 0.16666667 = 1/(3!) +#define EXP_COEFF_0 (0x3F000000) // 0.5 = 1/(2!) +#define EXP_LOGN2 (0x3F317218) // ln(2) = 0.6931471805 +#define EXP_LOG2E (0x3FB8AA3B) // log2(e) = 1/ln(2) = 1.4426950408 +#define EXP_ONE (0x3f800000) // 1.0 +#define EXP_RANGE_R (0x42B17218) // ln(FLT_MAX) approx = 88.7228 +#define EXP_RANGE_L (0xC2B00000) // -88.0 (approx log(FLT_MIN)) + +static inline HVX_Vector hvx_vec_exp_f32(HVX_Vector in_vec) { + HVX_Vector z_qf32_v; + HVX_Vector x_v; + HVX_Vector x_qf32_v; + HVX_Vector y_v; + HVX_Vector k_v; + HVX_Vector f_v; + HVX_Vector epsilon_v; + HVX_Vector log2e = Q6_V_vsplat_R(EXP_LOG2E); + HVX_Vector logn2 = Q6_V_vsplat_R(EXP_LOGN2); + HVX_Vector E_const; + HVX_Vector zero_v = Q6_V_vzero(); + + // exp(x) is approximated as follows: + // f = floor(x/ln(2)) = floor(x*log2(e)) + // epsilon = x - f*ln(2) + // exp(x) = exp(epsilon+f*ln(2)) + // = exp(epsilon)*exp(f*ln(2)) + // = exp(epsilon)*2^f + // + // Since epsilon is close to zero, it can be approximated with its Taylor series: + // exp(x) ~= 1+x+x^2/2!+x^3/3!+...+x^n/n!+... + // Preserving the first eight elements, we get: + // exp(x) ~= 1+x+e0*x^2+e1*x^3+e2*x^4+e3*x^5+e4*x^6+e5*x^7 + // = 1+x+(E0+(E1+(E2+(E3+(E4+E5*x)*x)*x)*x)*x)*x^2 + + HVX_Vector temp_v = in_vec; + + // Clamp inputs to (-88.0, 88.0) to avoid overflow/underflow + HVX_VectorPred pred_cap_right = Q6_Q_vcmp_gt_VsfVsf(in_vec, Q6_V_vsplat_R(EXP_RANGE_R)); + HVX_VectorPred pred_cap_left = Q6_Q_vcmp_gt_VsfVsf(Q6_V_vsplat_R(EXP_RANGE_L), in_vec); + + in_vec = Q6_V_vmux_QVV(pred_cap_right, Q6_V_vsplat_R(EXP_RANGE_R), temp_v); + in_vec = Q6_V_vmux_QVV(pred_cap_left, Q6_V_vsplat_R(EXP_RANGE_L), in_vec); + + epsilon_v = Q6_Vqf32_vmpy_VsfVsf(log2e, in_vec); + epsilon_v = Q6_Vsf_equals_Vqf32(epsilon_v); + + // f_v is the floating point result and k_v is the integer result + f_v = hvx_vec_floor_f32(epsilon_v); + k_v = hvx_vec_truncate_f32(f_v); + + x_qf32_v = Q6_Vqf32_vadd_VsfVsf(in_vec, zero_v); + + // x = x - f_v * logn2; + epsilon_v = Q6_Vqf32_vmpy_VsfVsf(f_v, logn2); + x_qf32_v = Q6_Vqf32_vsub_Vqf32Vqf32(x_qf32_v, epsilon_v); + // normalize before every QFloat's vmpy + x_qf32_v = Q6_Vqf32_vadd_Vqf32Vsf(x_qf32_v, zero_v); + + x_v = Q6_Vsf_equals_Vqf32(x_qf32_v); + + // z = x * x; + z_qf32_v = Q6_Vqf32_vmpy_Vqf32Vqf32(x_qf32_v, x_qf32_v); + z_qf32_v = Q6_Vqf32_vadd_Vqf32Vsf(z_qf32_v, zero_v); + + // y = E4 + E5 * x; + E_const = Q6_V_vsplat_R(EXP_COEFF_5); + y_v = Q6_Vqf32_vmpy_VsfVsf(E_const, x_v); + E_const = Q6_V_vsplat_R(EXP_COEFF_4); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); + + // y = E3 + y * x; + E_const = Q6_V_vsplat_R(EXP_COEFF_3); + y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); + + // y = E2 + y * x; + E_const = Q6_V_vsplat_R(EXP_COEFF_2); + y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); + + // y = E1 + y * x; + E_const = Q6_V_vsplat_R(EXP_COEFF_1); + y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); + + // y = E0 + y * x; + E_const = Q6_V_vsplat_R(EXP_COEFF_0); + y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); + + // y = x + y * z; + y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, z_qf32_v); + y_v = Q6_Vqf32_vadd_Vqf32Vqf32(y_v, x_qf32_v); + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v); + + // y = y + 1.0; + y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, Q6_V_vsplat_R(EXP_ONE)); + + // insert exponents + // y = ldexpf(y, k); + // y_v += k_v; // qf32 + // modify exponent + + y_v = Q6_Vsf_equals_Vqf32(y_v); + + // add k_v to the exponent of y_v + HVX_Vector y_v_exponent = Q6_Vw_vasl_VwR(y_v, 1); + + y_v_exponent = Q6_Vuw_vlsr_VuwR(y_v_exponent, IEEE_VSF_MANTLEN + 1); + y_v_exponent = Q6_Vw_vadd_VwVw(k_v, y_v_exponent); + + // exponent cannot be negative; if overflow is detected, result is set to zero + HVX_VectorPred qy_v_negative_exponent = Q6_Q_vcmp_gt_VwVw(zero_v, y_v_exponent); + + y_v = Q6_Vw_vaslacc_VwVwR(y_v, k_v, IEEE_VSF_MANTLEN); + + y_v = Q6_V_vmux_QVV(qy_v_negative_exponent, zero_v, y_v); + + return y_v; +} + +static inline HVX_Vector hvx_vec_exp_f32_guard(HVX_Vector in_vec, HVX_Vector max_exp, HVX_Vector inf) { + const HVX_VectorPred pred0 = Q6_Q_vcmp_gt_VsfVsf(in_vec, max_exp); + + HVX_Vector out = hvx_vec_exp_f32(in_vec); + + return Q6_V_vmux_QVV(pred0, inf, out); +} + +static inline void hvx_exp_f32(uint8_t * restrict dst, const uint8_t * restrict src, const int num_elems, bool negate) { + int left_over = num_elems & (VLEN_FP32 - 1); + int num_elems_whole = num_elems - left_over; + + int unaligned_addr = 0; + int unaligned_loop = 0; + if ((0 == hex_is_aligned((void *) src, VLEN)) || (0 == hex_is_aligned((void *) dst, VLEN))) { + unaligned_addr = 1; + } + // assert((0 == unaligned_addr) || (0 == num_elems_whole)); + if ((1 == unaligned_addr) && (num_elems_whole != 0)) { + unaligned_loop = 1; + } + + HVX_Vector vec_out = Q6_V_vzero(); + + static const float kInf = INFINITY; + static const float kMaxExp = 88.7228f; + + const HVX_Vector max_exp = hvx_vec_splat_f32(kMaxExp); + const HVX_Vector inf = hvx_vec_splat_f32(kInf); + + if (0 == unaligned_loop) { + HVX_Vector * p_vec_in1 = (HVX_Vector *) src; + HVX_Vector * p_vec_out = (HVX_Vector *) dst; + + #pragma unroll(4) + for (int i = 0; i < num_elems_whole; i += VLEN_FP32) { + if (true == negate) { + HVX_Vector neg_vec_in = hvx_vec_neg_f32(*p_vec_in1++); + *p_vec_out++ = hvx_vec_exp_f32_guard(neg_vec_in, max_exp, inf); + } else { + *p_vec_out++ = hvx_vec_exp_f32_guard(*p_vec_in1++, max_exp, inf); + } + } + } else { + #pragma unroll(4) + for (int i = 0; i < num_elems_whole; i += VLEN_FP32) { + HVX_Vector in = *(HVX_UVector *) (src + i * SIZEOF_FP32); + + if (true == negate) { + HVX_Vector neg_vec_in = hvx_vec_neg_f32(in); + *(HVX_UVector *) (dst + i * SIZEOF_FP32) = hvx_vec_exp_f32_guard(neg_vec_in, max_exp, inf); + } else { + *(HVX_UVector *) (dst + i * SIZEOF_FP32) = hvx_vec_exp_f32_guard(in, max_exp, inf); + } + } + } + + if (left_over > 0) { + const float * srcf = (float *) src + num_elems_whole; + float * dstf = (float *) dst + num_elems_whole; + + HVX_Vector in = *(HVX_UVector *) srcf; + + if (true == negate) { + HVX_Vector neg_vec_in = hvx_vec_neg_f32(in); + + vec_out = hvx_vec_exp_f32_guard(neg_vec_in, max_exp, inf); + } else { + vec_out = hvx_vec_exp_f32_guard(in, max_exp, inf); + } + + hvx_vec_store_u((void *) dstf, left_over * SIZEOF_FP32, vec_out); + } +} + +#include + +#include +#include +#include +#include +#include + +#include "hvx-base-defs.inc" + +// ==================================================== +// FUNCTION: 1/(x+1) y(0) = 1, y(0.5) = 0.6667, y(1) = 0.5 +// Order:3; continuity: True; Ends forced: True +// Mode: unsigned; Result fractional bits: 14 +// Peak Error: 1.1295e-04 Rms Error: 2.8410e-05 Mean Error: 1.1370e-05 +// 32769 -32706 31252 -10589 +// 32590 -30635 22793 -4493 +// 32066 -27505 16481 -2348 +// 31205 -24054 11849 -1306 + +static inline HVX_Vector hvx_vec_recip_xp1_O3_unsigned(HVX_Vector vx) { + // input is 0..0xffff representing 0.0 .. 1.0 + HVX_Vector p; + p = Q6_Vh_vlut4_VuhPh(vx, 0xFAE6F6D4EE73D6A3ull); + p = Q6_Vh_vmpa_VhVhVuhPuh_sat(p, vx, 0x2E49406159097A14ull); + p = Q6_Vh_vmps_VhVhVuhPuh_sat(p, vx, 0x5DF66B7177AB7FC2ull); + p = Q6_Vh_vmpa_VhVhVuhPuh_sat(p, vx, 0x79E57D427F4E8001ull); + return p; // signed result, 14 fractional bits +} + +// Find reciprocal of fp16. +// (1) first, convert to fp32, multiplying by 1.0; this is done to +// handle denormals. Ignoring sign and zero, result should be at +// least 5.9604645e-08 (32-bit code 0x33800000) and at most 131008 (0x47ffe000) +// (exponent in range [103,143]) +// (2) extract the mantissa into 16-bit unsigned; find reciprocal using a fitted poly +// (3) put this, along with '253-exp' (exp from (1)) together to make an qf32 +// (4) convert that to fp16 +// (5) put sign back in. Also, if the original value (w/o sign) was <0x81, replace +// the result with the max value. +static inline HVX_Vector hvx_vec_inverse_f16(HVX_Vector vals) { + HVX_Vector em_mask = Q6_Vh_vsplat_R(0x7FFF); + HVX_Vector avals = Q6_V_vand_VV(vals, em_mask); + HVX_VectorPred is_neg = Q6_Q_vcmp_gt_VhVh(avals, vals); + // is too small to 1/x ? for 'standard' fp16, this would be 0x101 + HVX_VectorPred is_small = Q6_Q_vcmp_gt_VhVh(Q6_Vh_vsplat_R(0x101), avals); + + HVX_VectorPair to_qf32 = Q6_Wqf32_vmpy_VhfVhf(avals, Q6_Vh_vsplat_R(0x3C00)); // *1.0 + HVX_Vector to_f32_0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(to_qf32)); + HVX_Vector to_f32_1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(to_qf32)); + + // bits 22..13 contain the mantissa now (w/o hidden bit); move to bit 14..5 of a 16-bit vector + HVX_Vector mant_u16 = Q6_Vh_vshuffo_VhVh(Q6_Vw_vasl_VwR(to_f32_1, 9), Q6_Vw_vasl_VwR(to_f32_0, 9)); + // likewise extract the upper 16 from each, containing the exponents in range 103..142 + HVX_Vector exp_u16 = Q6_Vh_vshuffo_VhVh(to_f32_1, to_f32_0); + //Get exponent in IEEE 32-bit representation + exp_u16 = Q6_Vuh_vlsr_VuhR(exp_u16, 7); + + // so, mant_u16 contains an unbiased mantissa in upper 10 bits of each u16 lane + // We can consider it to be x-1.0, with 16 fractional bits, where 'x' is in range [1.0,2.0) + // Use poly to transform to 1/x, with 14 fractional bits + // + HVX_Vector rm = hvx_vec_recip_xp1_O3_unsigned(mant_u16); + + HVX_Vector vcl0 = Q6_Vuh_vcl0_Vuh(rm); //count leading zeros + + // Get mantissa for 16-bit representation + HVX_Vector mant_recip = Q6_V_vand_VV(Q6_Vh_vasr_VhR(Q6_Vh_vasl_VhVh(rm, vcl0), 5), Q6_Vh_vsplat_R(0x03FF)); + + //Compute Reciprocal Exponent + HVX_Vector exp_recip = + Q6_Vh_vsub_VhVh(Q6_Vh_vsub_VhVh(Q6_Vh_vsplat_R(254), exp_u16), Q6_Vh_vsub_VhVh(vcl0, Q6_Vh_vsplat_R(1))); + //Convert it for 16-bit representation + exp_recip = Q6_Vh_vadd_VhVh_sat(Q6_Vh_vsub_VhVh(exp_recip, Q6_Vh_vsplat_R(127)), Q6_Vh_vsplat_R(15)); + exp_recip = Q6_Vh_vasl_VhR(exp_recip, 10); + + //Merge exponent and mantissa for reciprocal + HVX_Vector recip = Q6_V_vor_VV(exp_recip, mant_recip); + // map 'small' inputs to standard largest value 0x7bff + recip = Q6_V_vmux_QVV(is_small, Q6_Vh_vsplat_R(0x7bff), recip); + // add sign back + recip = Q6_V_vandor_VQR(recip, is_neg, 0x80008000); + return recip; +} + +static inline HVX_Vector hvx_vec_inverse_f32(HVX_Vector v_sf) { + HVX_Vector inv_aprox_sf = Q6_V_vsplat_R(0x7EEEEBB3); + HVX_Vector two_sf = hvx_vec_splat_f32(2.0); + + // First approximation + HVX_Vector i_sf = Q6_Vw_vsub_VwVw(inv_aprox_sf, v_sf); + + HVX_Vector r_qf; + + // Refine + r_qf = Q6_Vqf32_vmpy_VsfVsf( + i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v_sf))))); + r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32( + r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v_sf)))); + r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32( + r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v_sf)))); + + return Q6_Vsf_equals_Vqf32(r_qf); +} + +static inline HVX_Vector hvx_vec_inverse_f32_guard(HVX_Vector v_sf, HVX_Vector nan_inf_mask) { + HVX_Vector out = hvx_vec_inverse_f32(v_sf); + + HVX_Vector masked_out = Q6_V_vand_VV(out, nan_inf_mask); + const HVX_VectorPred pred = Q6_Q_vcmp_eq_VwVw(nan_inf_mask, masked_out); + + return Q6_V_vmux_QVV(pred, Q6_V_vzero(), out); +} + +#define hvx_inverse_f32_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const HVX_Vector nan_inf_mask = Q6_V_vsplat_R(0x7f800000); \ + \ + const uint32_t nvec = n / VLEN_FP32; \ + const uint32_t nloe = n % VLEN_FP32; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + vdst[i] = hvx_vec_inverse_f32_guard(vsrc[i], nan_inf_mask); \ + } \ + if (nloe) { \ + HVX_Vector v = hvx_vec_inverse_f32_guard(vsrc[i], nan_inf_mask); \ + vec_store((void *) &vdst[i], nloe * SIZEOF_FP32, v); \ + } \ + } while(0) + +static inline HVX_Vector hvx_vec_inverse_f16_guard(HVX_Vector v_sf, HVX_Vector nan_inf_mask) { + HVX_Vector out = hvx_vec_inverse_f16(v_sf); + + HVX_Vector masked_out = Q6_V_vand_VV(out, nan_inf_mask); + const HVX_VectorPred pred = Q6_Q_vcmp_eq_VhVh(nan_inf_mask, masked_out); + + return Q6_V_vmux_QVV(pred, Q6_V_vzero(), out); +} + +#define hvx_inverse_f16_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const HVX_Vector nan_inf_mask = Q6_Vh_vsplat_R(0x7c00); \ + \ + const uint32_t nvec = n / VLEN_FP16; \ + const uint32_t nloe = n % VLEN_FP16; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + vdst[i] = hvx_vec_inverse_f16_guard(vsrc[i], nan_inf_mask); \ + } \ + if (nloe) { \ + HVX_Vector v = hvx_vec_inverse_f16_guard(vsrc[i], nan_inf_mask); \ + vec_store((void *) &vdst[i], nloe * SIZEOF_FP16, v); \ + } \ + } while(0) + +// Generic macro to define alignment permutations for an op +#define DEFINE_HVX_INV_OP_VARIANTS(OP_NAME, OP_LOOP_BODY) \ +static inline void OP_NAME##_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \ + assert((uintptr_t) dst % 128 == 0); \ + assert((uintptr_t) src % 128 == 0); \ + OP_LOOP_BODY(HVX_Vector, HVX_Vector, hvx_vec_store_a); \ +} \ +static inline void OP_NAME##_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \ + assert((uintptr_t) dst % 128 == 0); \ + OP_LOOP_BODY(HVX_Vector, HVX_UVector, hvx_vec_store_a); \ +} \ +static inline void OP_NAME##_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \ + assert((uintptr_t) src % 128 == 0); \ + OP_LOOP_BODY(HVX_UVector, HVX_Vector, hvx_vec_store_u); \ +} \ +static inline void OP_NAME##_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \ + OP_LOOP_BODY(HVX_UVector, HVX_UVector, hvx_vec_store_u); \ +} \ + +// Dispatcher logic +#define HVX_INV_DISPATCHER(OP_NAME) \ +static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t num_elems) { \ + if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { \ + OP_NAME##_aa(dst, src, num_elems); \ + } else if (hex_is_aligned((void *) dst, 128)) { \ + OP_NAME##_au(dst, src, num_elems); \ + } else if (hex_is_aligned((void *) src, 128)) { \ + OP_NAME##_ua(dst, src, num_elems); \ + } else { \ + OP_NAME##_uu(dst, src, num_elems); \ + } \ +} + +DEFINE_HVX_INV_OP_VARIANTS(hvx_inverse_f32, hvx_inverse_f32_loop_body) +DEFINE_HVX_INV_OP_VARIANTS(hvx_inverse_f16, hvx_inverse_f16_loop_body) + +HVX_INV_DISPATCHER(hvx_inverse_f32) +HVX_INV_DISPATCHER(hvx_inverse_f16) + +#include "hvx-base-defs.inc" + +#define FAST_SIGMOID_LOG2F (0x3fb8aa3b) // 1.442695022 +#define FAST_SIGMOID_C1 (0x3d009076) // 0.03138777 +#define FAST_SIGMOID_C2 (0x3e8d74bd) // 0.276281267 +#define FAST_SIGMOID_C3 (0x3f000000) // 0.5 + +static inline HVX_Vector hvx_vec_fast_sigmoid_f32(HVX_Vector v) { + v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F)); + v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3)); + + HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v)); + HVX_Vector x = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int)); + HVX_Vector xx = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x); + + HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2)); + v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F)); + + HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1)); + v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx); + v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x); + + HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1)); + HVX_Vector v3_exponent = Q6_Vw_vasl_VwR(v3, 1); + v3_exponent = Q6_Vuw_vlsr_VuwR(v3_exponent, 24); + v3_exponent = Q6_Vw_vadd_VwVw(in_int, v3_exponent); + v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24); + + HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1)); + HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4)); + + HVX_Vector res = hvx_vec_inverse_f32(v5); + res = Q6_Vqf32_vmpy_VsfVsf(v3, res); + + return Q6_Vsf_equals_Vqf32(res); +} + +static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard(HVX_Vector v, + HVX_Vector one, + HVX_Vector max_exp, + HVX_Vector min_exp) { + const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v); + const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp); + + HVX_Vector out = hvx_vec_fast_sigmoid_f32(v); + out = Q6_V_vmux_QVV(pred_max, out, one); + return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero()); +} + +static inline HVX_Vector hvx_vec_tanh_f32(HVX_Vector x) { + // tanh(x) = 2 * sigmoid(2x) - 1 + HVX_Vector two = hvx_vec_splat_f32(2.0f); + HVX_Vector one = hvx_vec_splat_f32(1.0f); + HVX_Vector x2 = Q6_Vqf32_vmpy_VsfVsf(x, two); + + HVX_Vector max_exp = hvx_vec_splat_f32(87.f); + HVX_Vector min_exp = hvx_vec_splat_f32(-87.f); + + HVX_Vector sig2x = hvx_vec_fast_sigmoid_f32_guard(Q6_Vsf_equals_Vqf32(x2), one, max_exp, min_exp); + + HVX_Vector res = Q6_Vqf32_vmpy_VsfVsf(sig2x, two); + res = Q6_Vqf32_vsub_Vqf32Vsf(res, one); + return Q6_Vsf_equals_Vqf32(res); +} + +#define hvx_sigmoid_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const HVX_Vector one = hvx_vec_splat_f32(1.f); \ + const HVX_Vector max_exp = hvx_vec_splat_f32(87.f); \ + const HVX_Vector min_exp = hvx_vec_splat_f32(-87.f); \ + \ + const uint32_t epv = 128 / sizeof(float); \ + const uint32_t nvec = n / epv; \ + const uint32_t nloe = n % epv; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + vdst[i] = hvx_vec_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp); \ + } \ + if (nloe) { \ + HVX_Vector tmp = hvx_vec_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp); \ + vec_store((void *) &vdst[i], nloe * sizeof(float), tmp); \ + } \ + } while(0) + +#define hvx_tanh_loop_body(dst_type, src_type, vec_store) \ + do { \ + dst_type * restrict vdst = (dst_type *) dst; \ + src_type * restrict vsrc = (src_type *) src; \ + \ + const uint32_t epv = 128 / sizeof(float); \ + const uint32_t nvec = n / epv; \ + const uint32_t nloe = n % epv; \ + \ + uint32_t i = 0; \ + \ + _Pragma("unroll(4)") \ + for (; i < nvec; i++) { \ + vdst[i] = hvx_vec_tanh_f32(vsrc[i]); \ + } \ + if (nloe) { \ + HVX_Vector tmp = hvx_vec_tanh_f32(vsrc[i]); \ + vec_store((void *) &vdst[i], nloe * sizeof(float), tmp); \ + } \ + } while(0) + +static inline void hvx_sigmoid_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + hvx_sigmoid_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); +} + +static inline void hvx_sigmoid_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + hvx_sigmoid_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a); +} + +static inline void hvx_sigmoid_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) src % 128 == 0); + hvx_sigmoid_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u); +} + +static inline void hvx_sigmoid_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + hvx_sigmoid_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u); +} + +static inline void hvx_tanh_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + hvx_tanh_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a); +} + +#include "hvx-utils-defs.inc" + +#define GGML_COMMON_DECL_C +#include "ggml-common-defs.inc" +#include "htp-ctx-defs.inc" +#include "htp-ops-defs.inc" +#include "htp-ops-defs.inc" + +struct htp_unary_context { + struct htp_ops_context * octx; + + // Precomputed values + const uint8_t * data_src0; + uint8_t * data_dst; + + size_t src0_data_row_size; // actual data bytes per row + size_t dst_data_row_size; // actual data bytes per row + + size_t src0_row_size_aligned; + size_t dst_row_size_aligned; + + size_t src0_spad_half_size; + size_t dst_spad_half_size; + + uint32_t block; + uint32_t src0_nrows; + uint32_t src0_nrows_per_thread; + uint32_t nc; +}; + +// Convert flat row index to DDR byte offset using the tensor's actual strides. +// ir = i1 + ne1*(i2 + ne2*i3) => offset = i1*nb1 + i2*nb2 + i3*nb3 +static inline size_t unary_row_offset(uint32_t ir, + uint32_t ne1, uint32_t ne2, + size_t nb1, size_t nb2, size_t nb3) { + const uint32_t i1 = ir % ne1; + const uint32_t i2 = (ir / ne1) % ne2; + const uint32_t i3 = ir / (ne1 * ne2); + return i1 * nb1 + i2 * nb2 + i3 * nb3; +} +// Safe DMA block size from row `ir`: clamp to the tighter dim-1 slice +// boundary of src and dst so the nb1 stride stays valid for all rows. +static inline uint32_t unary_block_size(uint32_t ir, + uint32_t end_row, + uint32_t block, + bool src_contig, + bool dst_contig, + uint32_t src_ne1, + uint32_t dst_ne1) { + uint32_t limit = MIN(block, end_row - ir); + + if (!src_contig) { + const uint32_t src_slice_end = (ir / src_ne1 + 1) * src_ne1; + limit = MIN(limit, src_slice_end - ir); + } + + if (!dst_contig) { + const uint32_t dst_slice_end = (ir / dst_ne1 + 1) * dst_ne1; + limit = MIN(limit, dst_slice_end - ir); + } + + return limit; +} + +#define htp_unary_preamble \ + const uint32_t ne00 = src->ne[0]; \ + const uint32_t ne01 = src->ne[1]; \ + const uint32_t ne02 = src->ne[2]; \ + const uint32_t ne03 = src->ne[3]; \ + \ + const uint32_t ne0 = dst->ne[0]; \ + const uint32_t ne1 = dst->ne[1]; \ + const uint32_t ne2 = dst->ne[2]; \ + const uint32_t ne3 = dst->ne[3]; \ + \ + const uint32_t nb00 = src->nb[0]; \ + const uint32_t nb01 = src->nb[1]; \ + const uint32_t nb02 = src->nb[2]; \ + const uint32_t nb03 = src->nb[3]; \ + \ + const uint32_t nb0 = dst->nb[0]; \ + const uint32_t nb1 = dst->nb[1]; \ + const uint32_t nb2 = dst->nb[2]; \ + const uint32_t nb3 = dst->nb[3]; + +static void hvx_fast_rms_norm_f32(const uint8_t * restrict src, + uint8_t * restrict dst, + uint8_t * restrict pad, + const int num_elems, + float epsilon) { + (void)pad; + + const HVX_Vector * restrict v_src = (HVX_Vector *) src; + HVX_Vector * restrict v_dst = (HVX_Vector *) dst; + + const int nvec = num_elems / VLEN_FP32; // number of full vectors + const int nloe = num_elems % VLEN_FP32; // leftover elements + + // Compute sum of squares for full vectors + HVX_Vector sum_v = Q6_V_vsplat_R(0x00000000); + HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon); + + #pragma unroll(4) + for (int i = 0; i < nvec; i++) { + HVX_Vector v1 = v_src[i]; + HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1); + sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2); + } + + // Handle tail elements using vectorized ops with masking + if (nloe > 0) { + HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); + HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]); + HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1); + sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2); + } + + // Reduce HVX sum + sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v)); + + HVX_Vector t_v = hvx_vec_splat_f32((float) num_elems); + HVX_Vector denom_v = hvx_vec_inverse_f32(t_v); + HVX_Vector mean_v = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v); + HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v); + + // Scale full vectors + HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v)); + + #pragma unroll(4) + for (int i = 0; i < nvec; i++) { + HVX_Vector v1 = v_src[i]; + HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v); + v_dst[i] = Q6_Vsf_equals_Vqf32(v2); + } + + // Handle tail elements using vectorized ops with masking + if (nloe > 0) { + + HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4); + HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]); + HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v); + HVX_Vector result = Q6_Vsf_equals_Vqf32(v2); + + // Store with masking to avoid overwriting memory beyond the tensor + hvx_vec_store_a(&v_dst[nvec], nloe * 4, result); + } +} + +static void scale_f32(const float * restrict src, + float * restrict dst, + uint8_t * restrict spad, + const uint32_t num_rows, + const uint32_t row_elems, + const size_t row_size, + int32_t * op_params) { + float scale = 0.f; + float bias = 0.f; + memcpy(&scale, &op_params[0], sizeof(float)); + memcpy(&bias, &op_params[1], sizeof(float)); + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); + + hvx_scale_offset_f32_aa((uint8_t *) dst_local, (const uint8_t *) src_local, row_elems, scale, bias); + } +} + +static void rms_norm_f32(const float * restrict src, + float * restrict dst, + uint8_t * restrict spad, + const uint32_t num_rows, + const uint32_t row_elems, + const size_t row_size, + int32_t * op_params) { + float epsilon = 0.f; + memcpy(&epsilon, op_params, sizeof(float)); + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); + + hvx_fast_rms_norm_f32((const uint8_t *) src_local, (uint8_t *) dst_local, spad, row_elems, epsilon); + } +} + +static void sqr_f32(const float * restrict src, + float * restrict dst, + uint8_t * restrict spad, + const uint32_t num_rows, + const uint32_t row_elems, + const size_t row_size, + int32_t * op_params) { + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); + + hvx_sqr_f32_aa((uint8_t *) dst_local, (const uint8_t *) src_local, row_elems); + } +} + +static void sqrt_f32(const float * restrict src, + float * restrict dst, + uint8_t * restrict spad, + const uint32_t num_rows, + const uint32_t row_elems, + const size_t row_size, + int32_t * op_params) { + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); + + hvx_sqrt_f32_aa((uint8_t *) dst_local, (const uint8_t *) src_local, row_elems); + } +} + +static void neg_f32(const float * restrict src, + float * restrict dst, + uint8_t * restrict spad, + const uint32_t num_rows, + const uint32_t row_elems, + const size_t row_size, + int32_t * op_params) { + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); + + hvx_scale_f32_aa(dst_local, src_local, row_elems, -1.0f); + } +} + +static void exp_f32(const float * restrict src, + float * restrict dst, + uint8_t * restrict spad, + const uint32_t num_rows, + const uint32_t row_elems, + const size_t row_size, + int32_t * op_params) { + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); + + hvx_exp_f32(dst_local, src_local, row_elems, false); + } +} + +static void sigmoid_f32(const float * restrict src, + float * restrict dst, + uint8_t * restrict spad, + const uint32_t num_rows, + const uint32_t row_elems, + const size_t row_size, + int32_t * op_params) { + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * row_size); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * row_size); + + hvx_sigmoid_f32_aa(dst_local, src_local, row_elems); + } +} + +static void softplus_f32(const float * restrict src, + float * restrict dst, + uint8_t * restrict spad, + const uint32_t num_rows, + const uint32_t row_elems, + const size_t row_size, + int32_t * op_params) { + // softplus(x) = log(1 + exp(x)) + // Match CPU reference: ggml_compute_softplus_f32() in ggml-impl-defs.inc + for (uint32_t ir = 0; ir < num_rows; ir++) { + const float * restrict src_f = (const float *)((const uint8_t *)src + (ir * row_size)); + float * restrict dst_f = (float *)((uint8_t *)dst + (ir * row_size)); + + for (uint32_t i = 0; i < row_elems; i++) { + float x = src_f[i]; + // For x > 20: softplus(x) ≈ x (avoids exp overflow) + dst_f[i] = (x > 20.0f) ? x : logf(1.0f + expf(x)); + } + } +} + +static void unary_job_f32_per_thread(unsigned int nth, unsigned int ith, void * data) { + const struct htp_unary_context * uctx = (const struct htp_unary_context *) data; + struct htp_ops_context * octx = uctx->octx; + const struct htp_tensor * src = octx->src[0]; + const struct htp_tensor * dst = octx->dst; + + htp_unary_preamble; + + int htp_op = octx->op; + int32_t * op_params = octx->op_params; + uint32_t src0_nrows_per_thread = uctx->src0_nrows_per_thread; + + const size_t src0_data_row_size = uctx->src0_data_row_size; + const size_t dst_data_row_size = uctx->dst_data_row_size; + + const size_t src0_row_size_aligned = uctx->src0_row_size_aligned; + const size_t dst_row_size_aligned = uctx->dst_row_size_aligned; + + const uint32_t src0_nrows = uctx->src0_nrows; + const uint32_t src0_start_row = src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); + + // no work for this thread + if (src0_start_row >= src0_end_row) { + return; + } + + uint64_t t1, t2; + t1 = HAP_perf_get_qtimer_count(); + + const uint8_t * restrict data_src = uctx->data_src0; + uint8_t * restrict data_dst = uctx->data_dst; + + uint8_t * src0_spad_data = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread); + uint8_t * dst_spad_data = octx->dst_spad.data + (ith * octx->dst_spad.size_per_thread); + + size_t src0_spad_half_size = uctx->src0_spad_half_size; + size_t dst_spad_half_size = uctx->dst_spad_half_size; + + // Non-contiguous tensors have gaps at dim-2/3 boundaries that a single-stride + // 2D DMA descriptor cannot span. Clamp BLOCK to ne1 (one dim-1 slice) so every + // transfer stays within a nb1-uniform region. Skipped for contiguous tensors. + const bool src0_contig = (nb02 == (size_t)ne01 * nb01) && + (nb03 == (size_t)ne02 * nb02); + const bool dst_contig = (nb2 == (size_t)ne1 * nb1) && + (nb3 == (size_t)ne2 * nb2); + const uint32_t src0_max_block = src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); + const uint32_t dst_max_block = dst_contig ? uctx->block : MIN((uint32_t)uctx->block, ne1); + const uint32_t BLOCK = MIN(src0_max_block, dst_max_block); + if (BLOCK == 0) { + FARF(ERROR, "unary-f32 : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", + octx->src0_spad.size_per_thread, src0_row_size_aligned); + return; + } + + dma_queue * dma_queue = octx->ctx->dma[ith]; + + for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; spad_idx++) { + const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, ne1); + + // Dummy DMA transation for sequencing (interleaving dst,src,dst,...) + dma_queue_push(dma_queue, + dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)), + nb1, dst_row_size_aligned, dst_data_row_size, 0); + + const size_t src0_off = unary_row_offset(ir, ne01, ne02, nb01, nb02, nb03); + dma_queue_push(dma_queue, + dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src + src0_off), + src0_row_size_aligned, nb01, src0_data_row_size, block_size); + ir += block_size; + } + + for (uint32_t ir = src0_start_row; ir < src0_end_row; ) { + const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, ne1); + + float * dst_spad = (float *) dma_queue_pop(dma_queue).src; + float * src0_spad = (float *) dma_queue_pop(dma_queue).dst; + + // Process block in VTCM + switch (htp_op) { + case HTP_OP_RMS_NORM: + rms_norm_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); + break; + case HTP_OP_SCALE: + scale_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); + break; + case HTP_OP_SQR: + sqr_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); + break; + case HTP_OP_SQRT: + sqrt_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); + break; + case HTP_OP_UNARY_NEG: + neg_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); + break; + case HTP_OP_UNARY_EXP: + exp_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); + break; + case HTP_OP_UNARY_SIGMOID: + sigmoid_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); + break; + case HTP_OP_UNARY_SOFTPLUS: + softplus_f32(src0_spad, dst_spad, NULL, block_size, ne0, src0_row_size_aligned, op_params); + break; + default: + break; + } + + const size_t dst_off = unary_row_offset(ir, ne1, ne2, nb1, nb2, nb3); + dma_queue_push(dma_queue, + dma_make_ptr(data_dst + dst_off, dst_spad), + nb1, dst_row_size_aligned, dst_data_row_size, block_size); + + // prefetch N+2 loop iteration if any + const uint32_t next_ir = ir + block_size; + if (next_ir < src0_end_row) { + const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, ne1); + const uint32_t pref_ir = next_ir + next_block_size; + if (pref_ir < src0_end_row) { + const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, ne1); + const size_t src0_pref_off = unary_row_offset(pref_ir, ne01, ne02, nb01, nb02, nb03); + dma_queue_push(dma_queue, + dma_make_ptr(src0_spad, data_src + src0_pref_off), + src0_row_size_aligned, nb01, src0_data_row_size, pref_block_size); + } + } + ir += block_size; + } + + dma_queue_flush(dma_queue); + + t2 = HAP_perf_get_qtimer_count(); + + FARF(HIGH, "unary-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n", ith, nth, src->ne[0], + src->ne[1], src->ne[2], src->ne[3], src0_start_row, src0_end_row, dst->ne[0], dst->ne[1], dst->ne[2], + dst->ne[3], (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1)); +} + +static int execute_op_unary_f32(struct htp_ops_context * octx) { + int err = HTP_STATUS_OK; + + const struct htp_tensor * src0 = octx->src[0]; + const struct htp_tensor * dst = octx->dst; + + const char * op_type = NULL; + + switch (octx->op) { + case HTP_OP_RMS_NORM: + op_type = "rmsnorm-f32"; + break; + case HTP_OP_SCALE: + op_type = "scale-f32"; + break; + case HTP_OP_SQR: + op_type = "sqr-f32"; + break; + case HTP_OP_SQRT: + op_type = "sqrt-f32"; + break; + case HTP_OP_UNARY_NEG: + op_type = "neg-f32"; + break; + case HTP_OP_UNARY_EXP: + op_type = "exp-f32"; + break; + case HTP_OP_UNARY_SIGMOID: + op_type = "sigmoid-f32"; + break; + case HTP_OP_UNARY_SOFTPLUS: + op_type = "softplus-f32"; + break; + + default: + FARF(ERROR, "Unsupported unary Op %u\n", octx->op); + return HTP_STATUS_NO_SUPPORT; + } + + const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3]; + const uint32_t n_threads = MIN(octx->n_threads, src0_nrows); + + const size_t src0_data_row_size = src0->ne[0] * sizeof(float); + const size_t dst_data_row_size = dst->ne[0] * sizeof(float); + + const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, VLEN); + const size_t dst_row_size_aligned = hex_round_up(dst_data_row_size, VLEN); + + // VTCM scratchpads for all tensors + // N rows per thread, padded to HVX vector size + // Double buffering requires 2x size per buffer + + size_t spad_size_per_row = 2 * (src0_row_size_aligned + dst_row_size_aligned); + size_t vtcm_row_per_thread = (octx->ctx->vtcm_size)/ (n_threads * spad_size_per_row); + + // Make sure the reserved vtcm size is sufficient + if (vtcm_row_per_thread == 0) { + FARF(ERROR, "unary-%s : current VTCM reservation %zu is too small, needed %zu\n", op_type, octx->ctx->vtcm_size, + spad_size_per_row * n_threads); + return HTP_STATUS_VTCM_TOO_SMALL; + } + + octx->src0_spad.size_per_thread = src0_row_size_aligned * vtcm_row_per_thread * 2; + octx->dst_spad.size_per_thread = dst_row_size_aligned * vtcm_row_per_thread * 2; + + octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread; + octx->dst_spad.size = n_threads * octx->dst_spad.size_per_thread; + + octx->src0_spad.data = octx->ctx->vtcm_base; + octx->dst_spad.data = octx->src0_spad.data + octx->src0_spad.size; + + FARF(HIGH, "%s: (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-spad-size %u src1-spad-size %u dst-spad-size %u\n", op_type, + src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], + octx->src0_spad.size, octx->src1_spad.size, octx->dst_spad.size); + + if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) { + struct htp_unary_context uctx = { + .octx = octx, + .src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads, + .src0_nrows = src0_nrows, + + .data_src0 = (const uint8_t *)src0->data, + .data_dst = (uint8_t *)dst->data, + + .src0_data_row_size = src0_data_row_size, + .dst_data_row_size = dst_data_row_size, + + .src0_row_size_aligned = src0_row_size_aligned, + .dst_row_size_aligned = dst_row_size_aligned, + + .src0_spad_half_size = octx->src0_spad.size_per_thread / 2, + .dst_spad_half_size = octx->dst_spad.size_per_thread / 2, + + .block = (octx->src0_spad.size_per_thread / 2) / src0_row_size_aligned, + .nc = src0->ne[0], + }; + + worker_pool_run_func(octx->ctx->worker_pool, unary_job_f32_per_thread, &uctx, n_threads); + } + + return err; +} + +int op_unary(struct htp_ops_context * octx) { + int err = HTP_STATUS_OK; + + switch (octx->src[0]->type) { + case HTP_TYPE_F32: + err = execute_op_unary_f32(octx); + break; + + default: + err = HTP_STATUS_NO_SUPPORT; + break; + } + + return err; +} diff --git a/ggml/src/ggml-hexagon/htp/vtcm-utils.h b/ggml/src/ggml-hexagon/htp/vtcm-utils.h deleted file mode 100644 index b129fb74e315..000000000000 --- a/ggml/src/ggml-hexagon/htp/vtcm-utils.h +++ /dev/null @@ -1,16 +0,0 @@ -#ifndef VTCM_UTILS_H -#define VTCM_UTILS_H - -#include "hex-utils.h" - -#include -#include -#include - -static inline uint8_t *vtcm_seq_alloc(uint8_t **vtcm_ptr, size_t size) { - uint8_t *p = *vtcm_ptr; - *vtcm_ptr += size; - return p; -} - -#endif // VTCM_UTILS_H diff --git a/ggml/src/ggml-hexagon/htp/worker-pool.h b/ggml/src/ggml-hexagon/htp/worker-pool-defs.inc similarity index 100% rename from ggml/src/ggml-hexagon/htp/worker-pool.h rename to ggml/src/ggml-hexagon/htp/worker-pool-defs.inc diff --git a/ggml/src/ggml-hexagon/htp/worker-pool.c b/ggml/src/ggml-hexagon/htp/worker-pool.inc similarity index 99% rename from ggml/src/ggml-hexagon/htp/worker-pool.c rename to ggml/src/ggml-hexagon/htp/worker-pool.inc index 172e28908ebb..7633a0b5cbdc 100644 --- a/ggml/src/ggml-hexagon/htp/worker-pool.c +++ b/ggml/src/ggml-hexagon/htp/worker-pool.inc @@ -1,4 +1,4 @@ -#include "worker-pool.h" +#include "worker-pool-defs.inc" #include #include diff --git a/ggml/src/ggml-hexagon/libdl.h b/ggml/src/ggml-hexagon/libdl.h deleted file mode 100644 index 8ca5016f0396..000000000000 --- a/ggml/src/ggml-hexagon/libdl.h +++ /dev/null @@ -1,79 +0,0 @@ -#pragma once - -#ifdef _WIN32 -# define WIN32_LEAN_AND_MEAN -# ifndef NOMINMAX -# define NOMINMAX -# endif -# include -# include -#else -# include -# include -#endif -#include - -namespace fs = std::filesystem; - -#ifdef _WIN32 - -using dl_handle = std::remove_pointer_t; - -struct dl_handle_deleter { - void operator()(HMODULE handle) { - FreeLibrary(handle); - } -}; - -static inline dl_handle * dl_load_library(const fs::path & path) { - // suppress error dialogs for missing DLLs - DWORD old_mode = SetErrorMode(SEM_FAILCRITICALERRORS); - SetErrorMode(old_mode | SEM_FAILCRITICALERRORS); - - HMODULE handle = LoadLibraryW(path.wstring().c_str()); - - SetErrorMode(old_mode); - - return handle; -} - -static inline void * dl_get_sym(dl_handle * handle, const char * name) { - DWORD old_mode = SetErrorMode(SEM_FAILCRITICALERRORS); - SetErrorMode(old_mode | SEM_FAILCRITICALERRORS); - - void * p = (void *) GetProcAddress(handle, name); - - SetErrorMode(old_mode); - - return p; -} - -static inline const char * dl_error() { - return ""; -} - -#else - -using dl_handle = void; - -struct dl_handle_deleter { - void operator()(void * handle) { - dlclose(handle); - } -}; - -static inline dl_handle * dl_load_library(const fs::path & path) { - dl_handle * handle = dlopen(path.string().c_str(), RTLD_NOW | RTLD_LOCAL); - return handle; -} - -static inline void * dl_get_sym(dl_handle * handle, const char * name) { - return dlsym(handle, name); -} - -static inline const char * dl_error() { - const char *rslt = dlerror(); - return rslt != nullptr ? rslt : ""; -} - -#endif diff --git a/ggml/src/ggml-hexagon/op-desc.h b/ggml/src/ggml-hexagon/op-desc.h deleted file mode 100644 index a1e8ddd8b975..000000000000 --- a/ggml/src/ggml-hexagon/op-desc.h +++ /dev/null @@ -1,153 +0,0 @@ -#ifndef OP_DESC_H -#define OP_DESC_H - -#define GGML_COMMON_IMPL_CPP -#include "ggml-backend-impl.h" -#include "ggml-common.h" - -#include -#include - -struct op_desc { - char strides[64 * GGML_MAX_SRC]; - char dims[64 * GGML_MAX_SRC]; - char types[16 * GGML_MAX_SRC]; - char buffs[64 * GGML_MAX_SRC]; - char names[64 * GGML_MAX_SRC]; - - int format_tensor_dims(char * str, const struct ggml_tensor * t) { - if (t->ne[2] == 1 && t->ne[3] == 1) { - return sprintf(str, "%d:%d", (int) t->ne[0], (int) t->ne[1]); - } else { - return sprintf(str, "%d:%d:%d:%d", (int) t->ne[0], (int) t->ne[1], (int) t->ne[2], (int) t->ne[3]); - } - } - - void format_op_dims(char * str, const struct ggml_tensor * t) { - char * p = str; - - // append src0 and src1 (if any) - if (t->src[0]) { - p += format_tensor_dims(p, t->src[0]); - - for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) { - p += sprintf(p, " x "); - p += format_tensor_dims(p, t->src[i]); - } - - p += sprintf(p, " -> "); - } - - // format self dims separately for better visual alignment - char self[64]; - format_tensor_dims(self, t); - - p += sprintf(p, "%s", self); - } - - int format_tensor_strides(char * str, const struct ggml_tensor * t) { - const char * c = ggml_is_contiguous(t) ? "" : "!"; - - if (t->ne[2] == 1 && t->ne[3] == 1) { - return sprintf(str, "%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], c); - } else { - return sprintf(str, "%zu:%zu:%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], (size_t) t->nb[2], (size_t) t->nb[3], c); - } - } - - void format_op_strides(char * str, const struct ggml_tensor * t) { - char * p = str; - - // append src0 and src1 (if any) - if (t->src[0]) { - p += format_tensor_strides(p, t->src[0]); - - for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) { - p += sprintf(p, " x "); - p += format_tensor_strides(p, t->src[i]); - } - - p += sprintf(p, " -> "); - } - - // format self dims separately for better visual alignment - char self[64]; - format_tensor_strides(self, t); - - p += sprintf(p, "%s", self); - } - - void format_op_types(char * str, const struct ggml_tensor * t) { - char * p = str; - - // append src0 and src1 (if any) - if (t->src[0]) { - p += sprintf(p, "%s", ggml_type_name(t->src[0]->type)); - - for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) { - p += sprintf(p, " x "); - p += sprintf(p, "%s", ggml_type_name(t->src[i]->type)); - } - - p += sprintf(p, " -> "); - } - - p += sprintf(p, "%s", ggml_type_name(t->type)); - } - - const char * tensor_buff_name(const struct ggml_tensor * t) { - if (t->buffer) { - return ggml_backend_buffer_name(t->buffer); - } - return "NONE"; - } - - void format_op_buffs(char * str, const struct ggml_tensor * t) { - char * p = str; - - // append src0 and src1 (if any) - if (t->src[0]) { - p += sprintf(p, "%s", tensor_buff_name(t->src[0])); - - for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) { - p += sprintf(p, " x "); - p += sprintf(p, "%s", tensor_buff_name(t->src[i])); - } - - p += sprintf(p, " -> "); - } - - p += sprintf(p, "%s", tensor_buff_name(t)); - } - - void format_op_names(char * str, const struct ggml_tensor * t) { - char * p = str; - - // append src0 and src1 (if any) - if (t->src[0]) { - p += sprintf(p, "%s", t->src[0]->name); - - for (int i = 1; i < GGML_MAX_SRC && t->src[i]; i++) { - p += sprintf(p, " x "); - p += sprintf(p, "%s", t->src[i]->name); - } - - p += sprintf(p, " -> "); - } - - p += sprintf(p, "%s", t->name); - } - - void format(const ggml_tensor * op) { - format_op_dims(dims, op); - format_op_strides(strides, op); - format_op_types(types, op); - format_op_buffs(buffs, op); - format_op_names(names, op); - } - - op_desc() {} - op_desc(const ggml_tensor * op) { format(op); } -}; - -#endif // OP_DESC_H diff --git a/ggml/src/ggml-hip/CMakeLists.txt b/ggml/src/ggml-hip/CMakeLists.txt index a7d4e0ea2b53..ab7e889f942a 100644 --- a/ggml/src/ggml-hip/CMakeLists.txt +++ b/ggml/src/ggml-hip/CMakeLists.txt @@ -57,29 +57,29 @@ endif() message(STATUS "HIP and hipBLAS found") -file(GLOB GGML_HEADERS_ROCM "../ggml-cuda/*.cuh") -list(APPEND GGML_HEADERS_ROCM "../../include/ggml-cuda.h") +file(GLOB GGML_HEADERS_ROCM "../ggml-cuda/*.cuh.inc") +list(APPEND GGML_HEADERS_ROCM "../../include/ggml-cuda.h.inc") -file(GLOB GGML_SOURCES_ROCM "../ggml-cuda/*.cu") -file(GLOB SRCS "../ggml-cuda/template-instances/fattn-tile*.cu") +file(GLOB GGML_SOURCES_ROCM "../ggml-cuda/*.cu.inc") +file(GLOB SRCS "../ggml-cuda/template-instances/fattn-tile*.cu.inc") list(APPEND GGML_SOURCES_ROCM ${SRCS}) -file(GLOB SRCS "../ggml-cuda/template-instances/fattn-mma*.cu") +file(GLOB SRCS "../ggml-cuda/template-instances/fattn-mma*.cu.inc") list(APPEND GGML_SOURCES_ROCM ${SRCS}) -file(GLOB SRCS "../ggml-cuda/template-instances/mmq*.cu") +file(GLOB SRCS "../ggml-cuda/template-instances/mmq*.cu.inc") list(APPEND GGML_SOURCES_ROCM ${SRCS}) -file(GLOB SRCS "../ggml-cuda/template-instances/mmf*.cu") +file(GLOB SRCS "../ggml-cuda/template-instances/mmf*.cu.inc") list(APPEND GGML_SOURCES_ROCM ${SRCS}) if (GGML_CUDA_FA_ALL_QUANTS) - file(GLOB SRCS "../ggml-cuda/template-instances/fattn-vec*.cu") + file(GLOB SRCS "../ggml-cuda/template-instances/fattn-vec*.cu.inc") list(APPEND GGML_SOURCES_ROCM ${SRCS}) add_compile_definitions(GGML_CUDA_FA_ALL_QUANTS) else() list(APPEND GGML_SOURCES_ROCM - ../ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu - ../ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu - ../ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu - ../ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu) + ../ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu.inc + ../ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu.inc + ../ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu.inc + ../ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu.inc) endif() ggml_add_backend_library(ggml-hip diff --git a/ggml/src/ggml-impl.h b/ggml/src/ggml-impl-defs.inc similarity index 97% rename from ggml/src/ggml-impl.h rename to ggml/src/ggml-impl-defs.inc index 62b76abbcec9..4f88a7d6ade4 100644 --- a/ggml/src/ggml-impl.h +++ b/ggml/src/ggml-impl-defs.inc @@ -2,8 +2,8 @@ // GGML internal header -#include "ggml.h" -#include "gguf.h" +#include "ggml.h.inc" +#include "gguf.h.inc" #include #include @@ -31,6 +31,14 @@ extern "C" { void ggml_print_backtrace(void); uint64_t ggml_graph_next_uid(void); +bool ggml_op_is_empty_rust(int op); +size_t ggml_bitset_size_rust(size_t n); +size_t ggml_aligned_offset_rust(const void * buffer, size_t offset, size_t alignment); +int ggml_get_node_buffer_id_rust(const int * node_buffer_ids, int index); +bool ggml_buffer_address_less_rust(int a_chunk, size_t a_offset, int b_chunk, size_t b_offset); +bool ggml_isinf_fp16_rust(uint16_t value); +bool ggml_isnan_fp16_rust(uint16_t value); +bool ggml_is_invalid_e8m0_rust(uint8_t value); #ifndef MIN # define MIN(a, b) ((a) < (b) ? (a) : (b)) @@ -71,7 +79,7 @@ static inline int ggml_up(int n, int m) { return (n + m - 1) & ~(m - 1); } -// TODO: move to ggml.h? (won't be able to inline) +// TODO: move to ggml.h.inc? (won't be able to inline) static bool ggml_are_same_layout(const struct ggml_tensor * a, const struct ggml_tensor * b) { if (a->type != b->type) { return false; @@ -87,18 +95,7 @@ static bool ggml_are_same_layout(const struct ggml_tensor * a, const struct ggml return true; } -static bool ggml_op_is_empty(enum ggml_op op) { - switch (op) { - case GGML_OP_NONE: - case GGML_OP_RESHAPE: - case GGML_OP_TRANSPOSE: - case GGML_OP_VIEW: - case GGML_OP_PERMUTE: - return true; - default: - return false; - } -} +#define ggml_op_is_empty(op) ggml_op_is_empty_rust((int) (op)) static inline bool ggml_impl_is_view(const struct ggml_tensor * t) { return t->view_src != NULL; @@ -202,9 +199,7 @@ static_assert(sizeof(ggml_bitset_t) == 4, "bitset_t constants must be updated"); #define BITSET_SHR 5 // log2(sizeof(ggml_bitset_t)*8) #define BITSET_MASK (sizeof(ggml_bitset_t)*8 - 1) -static size_t ggml_bitset_size(size_t n) { - return (n + BITSET_MASK) >> BITSET_SHR; -} +#define ggml_bitset_size(n) ggml_bitset_size_rust(n) static inline bool ggml_bitset_get(const ggml_bitset_t * bitset, size_t i) { return !!(bitset[i >> BITSET_SHR] & (1u << (i & BITSET_MASK))); @@ -351,7 +346,7 @@ struct ggml_cgraph { // if you need the gradients, get them from the original graph struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph, int i0, int i1); -// ggml-alloc.c: true if the operation can reuse memory from its sources +// allocator implementation: true if the operation can reuse memory from its sources GGML_API bool ggml_op_can_inplace(enum ggml_op op); diff --git a/ggml/src/ggml-metal/CMakeLists.txt b/ggml/src/ggml-metal/CMakeLists.txt index 42054d841aa3..a7cc6d6650db 100644 --- a/ggml/src/ggml-metal/CMakeLists.txt +++ b/ggml/src/ggml-metal/CMakeLists.txt @@ -4,13 +4,18 @@ find_library(METALKIT_FRAMEWORK MetalKit REQUIRED) message(STATUS "Metal framework found") +set(GGML_METAL_SOURCES + ggml-metal.inc + ggml-metal-device.inc + ggml-metal-common.inc + ggml-metal-ops.inc +) +set_source_files_properties(${GGML_METAL_SOURCES} PROPERTIES LANGUAGE CXX) + ggml_add_backend_library(ggml-metal - ggml-metal.cpp + ${GGML_METAL_SOURCES} ggml-metal-device.m - ggml-metal-device.cpp - ggml-metal-common.cpp ggml-metal-context.m - ggml-metal-ops.cpp ) target_link_libraries(ggml-metal PRIVATE @@ -23,18 +28,18 @@ if (GGML_METAL_NDEBUG) add_compile_definitions(GGML_METAL_NDEBUG) endif() -set(METALLIB_COMMON "${CMAKE_CURRENT_SOURCE_DIR}/../ggml-common.h") +set(METALLIB_COMMON "${CMAKE_CURRENT_SOURCE_DIR}/../ggml-common-defs.inc") if (GGML_METAL_EMBED_LIBRARY) enable_language(ASM) add_compile_definitions(GGML_METAL_EMBED_LIBRARY) set(METALLIB_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal.metal") - set(METALLIB_IMPL "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal-impl.h") + set(METALLIB_IMPL "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal-impl-defs.inc") file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/autogenerated") - # merge ggml-common.h and ggml-metal.metal into a single file + # merge ggml-common-defs.inc and ggml-metal.metal into a single file set(METALLIB_EMBED_ASM "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.s") set(METALLIB_SOURCE_EMBED "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.metal") set(METALLIB_SOURCE_EMBED_TMP "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.metal.tmp") @@ -42,15 +47,15 @@ if (GGML_METAL_EMBED_LIBRARY) add_custom_command( OUTPUT "${METALLIB_EMBED_ASM}" COMMAND echo "Embedding Metal library" - COMMAND sed -e "/__embed_ggml-common.h__/r ${METALLIB_COMMON}" -e "/__embed_ggml-common.h__/d" < "${METALLIB_SOURCE}" > "${METALLIB_SOURCE_EMBED_TMP}" - COMMAND sed -e "/\#include \"ggml-metal-impl.h\"/r ${METALLIB_IMPL}" -e "/\#include \"ggml-metal-impl.h\"/d" < "${METALLIB_SOURCE_EMBED_TMP}" > "${METALLIB_SOURCE_EMBED}" + COMMAND sed -e "/__embed_ggml-common-defs.inc__/r ${METALLIB_COMMON}" -e "/__embed_ggml-common-defs.inc__/d" < "${METALLIB_SOURCE}" > "${METALLIB_SOURCE_EMBED_TMP}" + COMMAND sed -e "/\#include \"ggml-metal-impl-defs.inc\"/r ${METALLIB_IMPL}" -e "/\#include \"ggml-metal-impl-defs.inc\"/d" < "${METALLIB_SOURCE_EMBED_TMP}" > "${METALLIB_SOURCE_EMBED}" COMMAND echo ".section __DATA,__ggml_metallib" > "${METALLIB_EMBED_ASM}" COMMAND echo ".globl _ggml_metallib_start" >> "${METALLIB_EMBED_ASM}" COMMAND echo "_ggml_metallib_start:" >> "${METALLIB_EMBED_ASM}" COMMAND echo .incbin "\"${METALLIB_SOURCE_EMBED}\"" >> "${METALLIB_EMBED_ASM}" COMMAND echo ".globl _ggml_metallib_end" >> "${METALLIB_EMBED_ASM}" COMMAND echo "_ggml_metallib_end:" >> "${METALLIB_EMBED_ASM}" - DEPENDS ../ggml-common.h ggml-metal.metal ggml-metal-impl.h + DEPENDS ../ggml-common-defs.inc ggml-metal.metal ggml-metal-impl-defs.inc COMMENT "Generate assembly for embedded Metal library" VERBATIM ) @@ -58,9 +63,9 @@ if (GGML_METAL_EMBED_LIBRARY) target_sources(ggml-metal PRIVATE "${METALLIB_EMBED_ASM}") else() # copy metal files to bin directory - configure_file(../ggml-common.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h COPYONLY) + configure_file(../ggml-common-defs.inc ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common-defs.inc COPYONLY) configure_file(ggml-metal.metal ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal COPYONLY) - configure_file(ggml-metal-impl.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal-impl.h COPYONLY) + configure_file(ggml-metal-impl-defs.inc ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal-impl-defs.inc COPYONLY) if (GGML_METAL_SHADER_DEBUG) # custom command to do the following: @@ -94,7 +99,7 @@ else() OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib COMMAND xcrun -sdk macosx metal ${XC_FLAGS} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal -o - | xcrun -sdk macosx metallib - -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib - COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h + COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common-defs.inc COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal DEPENDS ggml-metal.metal ${METALLIB_COMMON} COMMENT "Compiling Metal kernels" diff --git a/ggml/src/ggml-metal/ggml-metal-common.h b/ggml/src/ggml-metal/ggml-metal-common-defs.inc similarity index 100% rename from ggml/src/ggml-metal/ggml-metal-common.h rename to ggml/src/ggml-metal/ggml-metal-common-defs.inc diff --git a/ggml/src/ggml-metal/ggml-metal-common.cpp b/ggml/src/ggml-metal/ggml-metal-common.inc similarity index 99% rename from ggml/src/ggml-metal/ggml-metal-common.cpp rename to ggml/src/ggml-metal/ggml-metal-common.inc index 2eb9820bff91..546cef4d620b 100644 --- a/ggml/src/ggml-metal/ggml-metal-common.cpp +++ b/ggml/src/ggml-metal/ggml-metal-common.inc @@ -1,7 +1,7 @@ -#include "ggml-metal-common.h" +#include "ggml-metal-common-defs.inc" -#include "ggml-impl.h" -#include "ggml-backend-impl.h" +#include "ggml-impl-defs.inc" +#include "ggml-backend-impl-defs.inc" #include diff --git a/ggml/src/ggml-metal/ggml-metal-context.h b/ggml/src/ggml-metal/ggml-metal-context-defs.inc similarity index 97% rename from ggml/src/ggml-metal/ggml-metal-context.h rename to ggml/src/ggml-metal/ggml-metal-context-defs.inc index abf4b06ed2ab..83396c4f0736 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.h +++ b/ggml/src/ggml-metal/ggml-metal-context-defs.inc @@ -1,6 +1,6 @@ #pragma once -#include "ggml-metal-device.h" +#include "ggml-metal-device-defs.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/src/ggml-metal/ggml-metal-context.m b/ggml/src/ggml-metal/ggml-metal-context.m index 32d97cd5d0af..616b35afec6b 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.m +++ b/ggml/src/ggml-metal/ggml-metal-context.m @@ -1,11 +1,11 @@ -#import "ggml-metal-context.h" +#import "ggml-metal-context-defs.inc" -#import "ggml-impl.h" -#import "ggml-backend-impl.h" +#import "ggml-impl-defs.inc" +#import "ggml-backend-impl-defs.inc" -#import "ggml-metal-impl.h" -#import "ggml-metal-common.h" -#import "ggml-metal-ops.h" +#import "ggml-metal-impl-defs.inc" +#import "ggml-metal-common-defs.inc" +#import "ggml-metal-ops-defs.inc" #import diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device-defs.inc similarity index 99% rename from ggml/src/ggml-metal/ggml-metal-device.h rename to ggml/src/ggml-metal/ggml-metal-device-defs.inc index a6c1dab55158..db7d75c66256 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device-defs.inc @@ -1,6 +1,6 @@ #pragma once -#include "ggml.h" +#include "ggml.h.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.inc similarity index 99% rename from ggml/src/ggml-metal/ggml-metal-device.cpp rename to ggml/src/ggml-metal/ggml-metal-device.inc index d211bf79f143..2a1ddc20623a 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.inc @@ -1,8 +1,8 @@ -#include "ggml-metal-device.h" +#include "ggml-metal-device-defs.inc" -#include "ggml-metal-impl.h" +#include "ggml-metal-impl-defs.inc" -#include "ggml-impl.h" +#include "ggml-impl-defs.inc" #include #include diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index fe90aafe7bc2..95d3255d8ff0 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1,6 +1,6 @@ -#import "ggml-metal-device.h" +#import "ggml-metal-device-defs.inc" -#import "ggml-impl.h" +#import "ggml-impl-defs.inc" #include diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl-defs.inc similarity index 100% rename from ggml/src/ggml-metal/ggml-metal-impl.h rename to ggml/src/ggml-metal/ggml-metal-impl-defs.inc diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops-defs.inc similarity index 99% rename from ggml/src/ggml-metal/ggml-metal-ops.h rename to ggml/src/ggml-metal/ggml-metal-ops-defs.inc index 36c61071b4fa..574301a0ebaa 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops-defs.inc @@ -1,6 +1,6 @@ #pragma once -#include "ggml-metal-device.h" +#include "ggml-metal-device-defs.inc" #ifdef __cplusplus extern "C" { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.inc similarity index 99% rename from ggml/src/ggml-metal/ggml-metal-ops.cpp rename to ggml/src/ggml-metal/ggml-metal-ops.inc index 5fa162c875ca..95a579607d79 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.inc @@ -1,12 +1,12 @@ -#include "ggml-metal-ops.h" +#include "ggml-metal-ops-defs.inc" -#include "ggml.h" -#include "ggml-impl.h" -#include "ggml-backend-impl.h" +#include "ggml.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-backend-impl-defs.inc" -#include "ggml-metal-impl.h" -#include "ggml-metal-common.h" -#include "ggml-metal-device.h" +#include "ggml-metal-impl-defs.inc" +#include "ggml-metal-common-defs.inc" +#include "ggml-metal-device-defs.inc" #include #include diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.inc similarity index 99% rename from ggml/src/ggml-metal/ggml-metal.cpp rename to ggml/src/ggml-metal/ggml-metal.inc index cc329d675941..7ecb341c553d 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.inc @@ -1,11 +1,11 @@ -#include "ggml-metal.h" +#include "ggml-metal.h.inc" -#include "ggml-impl.h" -#include "ggml-backend-impl.h" +#include "ggml-impl-defs.inc" +#include "ggml-backend-impl-defs.inc" -#include "ggml-metal-device.h" -#include "ggml-metal-context.h" -#include "ggml-metal-ops.h" +#include "ggml-metal-device-defs.inc" +#include "ggml-metal-context-defs.inc" +#include "ggml-metal-ops-defs.inc" #include #include diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index c372eaedeae4..a02af8580cb0 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -1,11 +1,11 @@ #define GGML_COMMON_DECL_METAL #define GGML_COMMON_IMPL_METAL #if defined(GGML_METAL_EMBED_LIBRARY) -__embed_ggml-common.h__ +__embed_ggml-common-defs.inc__ #else -#include "ggml-common.h" +#include "ggml-common-defs.inc" #endif -#include "ggml-metal-impl.h" +#include "ggml-metal-impl-defs.inc" #include @@ -2073,7 +2073,7 @@ template [[host_name("kernel_soft_max_f32")]] kernel kernel_soft_max_t kerne template [[host_name("kernel_soft_max_f16_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4; template [[host_name("kernel_soft_max_f32_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4; -// ref: ggml.c:ggml_compute_forward_ssm_conv_f32 +// ref: ggml.c.inc:ggml_compute_forward_ssm_conv_f32 kernel void kernel_ssm_conv_f32_f32( constant ggml_metal_kargs_ssm_conv & args, device const void * src0, @@ -2230,7 +2230,7 @@ kernel void kernel_ssm_conv_f32_f32_batched_4( x[0] = sumf; } -// ref: ggml.c:ggml_compute_forward_ssm_scan_f32, Mamba-2 part +// ref: ggml.c.inc:ggml_compute_forward_ssm_scan_f32, Mamba-2 part // Optimized version: reduces redundant memory loads by having one thread load shared values kernel void kernel_ssm_scan_f32( constant ggml_metal_kargs_ssm_scan & args, diff --git a/ggml/src/ggml-musa/CMakeLists.txt b/ggml/src/ggml-musa/CMakeLists.txt index cc53c812ce5f..8b55ce56a278 100644 --- a/ggml/src/ggml-musa/CMakeLists.txt +++ b/ggml/src/ggml-musa/CMakeLists.txt @@ -25,34 +25,34 @@ if (MUSAToolkit_FOUND) endif() message(STATUS "Using MUSA architectures: ${MUSA_ARCHITECTURES}") - file(GLOB GGML_HEADERS_MUSA "../ggml-cuda/*.cuh") - list(APPEND GGML_HEADERS_MUSA "../../include/ggml-cuda.h") - list(APPEND GGML_HEADERS_MUSA "../ggml-musa/mudnn.cuh") + file(GLOB GGML_HEADERS_MUSA "../ggml-cuda/*.cuh.inc") + list(APPEND GGML_HEADERS_MUSA "../../include/ggml-cuda.h.inc") + list(APPEND GGML_HEADERS_MUSA "../ggml-musa/mudnn.cuh.inc") - file(GLOB GGML_SOURCES_MUSA "../ggml-cuda/*.cu") - file(GLOB SRCS "../ggml-cuda/template-instances/fattn-tile*.cu") + file(GLOB GGML_SOURCES_MUSA "../ggml-cuda/*.cu.inc") + file(GLOB SRCS "../ggml-cuda/template-instances/fattn-tile*.cu.inc") list(APPEND GGML_SOURCES_MUSA ${SRCS}) - file(GLOB SRCS "../ggml-cuda/template-instances/fattn-mma*.cu") + file(GLOB SRCS "../ggml-cuda/template-instances/fattn-mma*.cu.inc") list(APPEND GGML_SOURCES_MUSA ${SRCS}) - file(GLOB SRCS "../ggml-cuda/template-instances/mmq*.cu") + file(GLOB SRCS "../ggml-cuda/template-instances/mmq*.cu.inc") list(APPEND GGML_SOURCES_MUSA ${SRCS}) if (GGML_MUSA_MUDNN_COPY) - file(GLOB SRCS "../ggml-musa/*.cu") + file(GLOB SRCS "../ggml-musa/*.cu.inc") list(APPEND GGML_SOURCES_MUSA ${SRCS}) add_compile_definitions(GGML_MUSA_MUDNN_COPY) endif() if (GGML_CUDA_FA_ALL_QUANTS) - file(GLOB SRCS "../ggml-cuda/template-instances/fattn-vec*.cu") + file(GLOB SRCS "../ggml-cuda/template-instances/fattn-vec*.cu.inc") list(APPEND GGML_SOURCES_MUSA ${SRCS}) add_compile_definitions(GGML_CUDA_FA_ALL_QUANTS) else() list(APPEND GGML_SOURCES_MUSA - ../ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu - ../ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu - ../ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu - ../ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu) + ../ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu.inc + ../ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu.inc + ../ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu.inc + ../ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu.inc) endif() set_source_files_properties(${GGML_SOURCES_MUSA} PROPERTIES LANGUAGE CXX) diff --git a/ggml/src/ggml-musa/mudnn.cu b/ggml/src/ggml-musa/mudnn.cu.inc similarity index 99% rename from ggml/src/ggml-musa/mudnn.cu rename to ggml/src/ggml-musa/mudnn.cu.inc index 020c1702c45c..4a1da1bb126f 100644 --- a/ggml/src/ggml-musa/mudnn.cu +++ b/ggml/src/ggml-musa/mudnn.cu.inc @@ -1,7 +1,7 @@ #include #include -#include "mudnn.cuh" +#include "mudnn.cuh.inc" namespace mudnn = musa::dnn; diff --git a/ggml/src/ggml-musa/mudnn.cuh b/ggml/src/ggml-musa/mudnn.cuh.inc similarity index 82% rename from ggml/src/ggml-musa/mudnn.cuh rename to ggml/src/ggml-musa/mudnn.cuh.inc index c30128561e81..91ce8e3fec09 100644 --- a/ggml/src/ggml-musa/mudnn.cuh +++ b/ggml/src/ggml-musa/mudnn.cuh.inc @@ -1,7 +1,7 @@ #pragma once -#include "ggml-cuda/common.cuh" -#include "ggml.h" +#include "ggml-cuda/common.cuh.inc" +#include "ggml.h.inc" // Asynchronously copies data from src tensor to dst tensor using the provided context. // Returns a musaError_t indicating success or failure. diff --git a/ggml/src/ggml-opencl/CMakeLists.txt b/ggml/src/ggml-opencl/CMakeLists.txt index 35d425a431fa..26dc75d6b425 100644 --- a/ggml/src/ggml-opencl/CMakeLists.txt +++ b/ggml/src/ggml-opencl/CMakeLists.txt @@ -3,9 +3,10 @@ find_package(Python3 REQUIRED) set(TARGET_NAME ggml-opencl) +set_source_files_properties(ggml-opencl.cpp.inc PROPERTIES LANGUAGE CXX) ggml_add_backend_library(${TARGET_NAME} - ggml-opencl.cpp - ../../include/ggml-opencl.h) + ggml-opencl.cpp.inc + ../../include/ggml-opencl.h.inc) target_link_libraries(${TARGET_NAME} PRIVATE ${OpenCL_LIBRARIES}) target_include_directories(${TARGET_NAME} PRIVATE ${OpenCL_INCLUDE_DIRS}) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp.inc similarity index 99% rename from ggml/src/ggml-opencl/ggml-opencl.cpp rename to ggml/src/ggml-opencl/ggml-opencl.cpp.inc index 74948c27e4e8..e79789419949 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp.inc @@ -7,11 +7,11 @@ #pragma GCC diagnostic ignored "-Wgnu-anonymous-struct" #endif -#include "ggml-opencl.h" -#include "ggml-backend.h" -#include "ggml-impl.h" -#include "ggml-backend-impl.h" -#include "ggml.h" +#include "ggml-opencl.h.inc" +#include "ggml-backend.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml.h.inc" #include diff --git a/ggml/src/ggml-openvino/CMakeLists.txt b/ggml/src/ggml-openvino/CMakeLists.txt index 175b585661d3..48c9d75b7297 100644 --- a/ggml/src/ggml-openvino/CMakeLists.txt +++ b/ggml/src/ggml-openvino/CMakeLists.txt @@ -3,8 +3,9 @@ find_package(OpenCL REQUIRED) include("${OpenVINO_DIR}/../3rdparty/tbb/lib/cmake/TBB/TBBConfig.cmake") -file(GLOB_RECURSE GGML_HEADERS_OPENVINO "*.h" "*.hpp") -file(GLOB_RECURSE GGML_SOURCES_OPENVINO "*.cpp") +set(GGML_HEADERS_OPENVINO) +set(GGML_SOURCES_OPENVINO ggml-openvino.cpp.inc) +set_source_files_properties(${GGML_SOURCES_OPENVINO} PROPERTIES LANGUAGE CXX) ggml_add_backend_library(ggml-openvino ${GGML_SOURCES_OPENVINO} diff --git a/ggml/src/ggml-openvino/ggml-decoder.cpp b/ggml/src/ggml-openvino/ggml-decoder.cpp deleted file mode 100644 index 5095e7998493..000000000000 --- a/ggml/src/ggml-openvino/ggml-decoder.cpp +++ /dev/null @@ -1,985 +0,0 @@ -#include "ggml-decoder.h" - -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "ggml-openvino-extra.h" -#include "ggml-openvino.h" -#include "ggml-quants.h" - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -GgmlOvDecoder::GgmlOvDecoder(ggml_cgraph * cgraph, - ModelParams & model_params, - ComputeParams & compute_params, - std::map> & model_weights, - bool is_static, - bool is_stateful, - bool is_prefill, - int prefill_chunk_size) : - m_is_static(is_static), - m_is_stateful(is_stateful), - m_is_prefill(is_prefill), - m_naive(false), - m_prefill_chunk_size(prefill_chunk_size), - m_cgraph(cgraph), - m_model_weights(model_weights), - m_model_params(model_params), - m_compute_params(compute_params) { - if (auto * env = getenv("GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS"); env && std::string(env) != "0") { -#ifdef _WIN32 - _putenv_s("GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS", ""); -#else - unsetenv("GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS"); -#endif - print_tensor_address_map(cgraph); - } - - validate_cgraph(); - - set_input_output(); - compute_model_inputs(); - compute_model_outputs(); - - for (int node_n = 0; node_n < cgraph->n_nodes; node_n++) { - m_node_info_list[node_n].node_op_case = compute_op_case(m_node_info_list[node_n].node); - m_node_info_list[node_n].node_op_type = compute_op_type(m_node_info_list[node_n].node); - } - - add_extra_inputs(); -} - -void GgmlOvDecoder::update_io(ggml_cgraph * cgraph) { - m_cgraph = cgraph; - m_model_inputs.clear(); - m_model_outputs.clear(); - m_node_info_list.clear(); - set_input_output(); - compute_model_inputs(); - compute_model_outputs(); -} - -GgmlOvDecoder::GgmlOvDecoder(ggml_cgraph * cgraph, std::map> & model_weights) { - m_cgraph = cgraph; - m_model_weights = model_weights; - m_naive = true; - set_input_output(); - compute_model_inputs(); - compute_model_outputs(); - for (int node_n = 0; node_n < cgraph->n_nodes; node_n++) { - m_node_info_list[node_n].node_op_case = compute_op_case(m_node_info_list[node_n].node); - m_node_info_list[node_n].node_op_type = compute_op_type(m_node_info_list[node_n].node); - } -} - -void GgmlOvDecoder::set_input_output() { - for (int node_n = 0; node_n < m_cgraph->n_nodes; node_n++) { - auto node = m_cgraph->nodes[node_n]; - - NodeInfo current_node_info; - auto node_name = std::string(node->name); - auto node_output_name = node_name; - auto * node_output = node; - if (node->op == GGML_OP_SET_ROWS) { - // SET_ROWS updates the tensor in place. For later ov op that uses the - // the view_src of SET_ROWS, we need to make sure they get the updated tensor - // by putting the view_src name in the tensor_map in - // /src/frontends/ggml/src/translate_session.cpp - node_output_name = std::string(node->view_src->name); - node_output = node->view_src; - } - - current_node_info.node = node; - current_node_info.node_name = node_name; - current_node_info.node_output = node_output; - current_node_info.node_output_name = node_output_name; - current_node_info.node_op_case = 0; - current_node_info.data_addr = node->data; - - for (int i = 0; i < GGML_MAX_SRC; i++) { - auto * src = node->src[i]; - if (src == nullptr) { - continue; - } - auto src_name = std::string(src->name); - if (src->flags & GGML_TENSOR_FLAG_INPUT) { - src_name = get_graph_input_ov_name(src, node); - } - current_node_info.node_inputs[src_name] = src; - current_node_info.node_inputs_names.push_back(src_name); - } - - m_node_info_list.push_back(current_node_info); - } -} - -int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const { - int op_case = 0; - switch (node->op) { - case GGML_OP_RESHAPE: { - auto * src = node->src[0]; - if (src->op == GGML_OP_RESHAPE && src->src[0]->ne[0] == node->ne[0] && src->src[0]->ne[1] == node->ne[1]) { - op_case = 4; - } else if (node->ne[0] * node->ne[1] == src->ne[0]) { - op_case = 1; - } else if (src->ne[0] * src->ne[1] == node->ne[0]) { - op_case = 2; - if (src->ne[2] * src->ne[3] == node->ne[1]) { - op_case = 5; - } - } else if (src->ne[0] * src->ne[1] == node->ne[1]) { - op_case = 3; - } else if (src->ne[1] * src->ne[2] == node->ne[1]) { - op_case = 6; - } - break; - } - case GGML_OP_CONT: { - if (node->src[0]->op == GGML_OP_PERMUTE) { - op_case = 1; - } else if (node->src[0]->op == GGML_OP_TRANSPOSE) { - op_case = 2; - } else if (node->src[0]->op == GGML_OP_VIEW) { - op_case = 3; - } - break; - } - case GGML_OP_PERMUTE: { - if (node->src[0]->op != GGML_OP_VIEW) { - op_case = 1; - } else if (node->src[0]->src[0]->op == GGML_OP_NONE) { - // kv cache tensor - std::string src_name(node->view_src->name); - int layer = extract_layer_from_name(src_name); - if (!is_swa_layer(layer)) { - op_case = 2; - } else { - op_case = 3; - } - } else { - // rope'ed query tensor - op_case = 4; - } - break; - } - case GGML_OP_MUL_MAT: { - if (node->src[0]->op == GGML_OP_CONT && node->src[0]->src[0]->op == GGML_OP_TRANSPOSE) { - op_case = 2; - } else if (node->src[0]->op == GGML_OP_VIEW && node->src[1]->op == GGML_OP_VIEW) { - op_case = 3; - } - break; - } - case GGML_OP_GET_ROWS: { - if (node->src[1]->op == GGML_OP_VIEW) { - op_case = 2; - } - break; - } - case GGML_OP_ROPE: { - const int mode = node->op_params[2]; - switch (mode) { - case GGML_ROPE_TYPE_NEOX: { - op_case = 0x00010000; - break; - } - case GGML_ROPE_TYPE_IMROPE: { - op_case = 0x00020000; - break; - } - default: - op_case = 0x00000000; - break; - } - if (node->src[0]->op == GGML_OP_VIEW) { - op_case = (op_case | 0x00000002); - } - break; - } - case GGML_OP_VIEW: { - if (node->src[0]->op == GGML_OP_VIEW) { - auto * src = node->src[0]; - if (ggml_nelements(node) != ggml_nelements(src)) { - throw std::runtime_error("Unsupported VIEW case"); - } - op_case = 2; - } - { - auto * src = node->src[0]; - if ((ggml_nelements(node) != ggml_nelements(src)) && m_naive) { - // Compare each dimension of node and src, if only one dimension differs then op_case=3 - int diff_count = 0; - for (int i = 0; i < GGML_MAX_DIMS; i++) { - if (node->ne[i] != src->ne[i]) { - diff_count++; - } - } - if (diff_count == 1) { - op_case = 3; - } - } - } - break; - } - default: - break; - } - return op_case; -} - -int extract_layer_from_name(const std::string & name) { - size_t pos1 = name.find("_l"); - assert(pos1 != std::string::npos); - pos1 += 2; - size_t pos2 = name.find(' ', pos1); - if (pos2 == std::string::npos) { - pos2 = name.length(); - } - std::string layer_str = name.substr(pos1, pos2 - pos1); - int layer = std::stoi(layer_str); - return layer; -} - -std::pair GgmlOvDecoder::compute_llm_params(ggml_cgraph * cgraph, bool is_static) { - ModelParams model_params; - ComputeParams compute_params; - for (int i = 0; i < cgraph->n_nodes; i++) { - auto * node = cgraph->nodes[i]; - std::string name = std::string(node->name); - if (node->op == GGML_OP_FLASH_ATTN_EXT) { - model_params.n_heads = node->src[0]->ne[2]; - model_params.n_heads_kv = node->src[1]->ne[2]; - model_params.head_size = node->src[0]->ne[0]; - compute_params.input_len = node->src[0]->ne[1]; - - auto * cache_k_perm = node->src[1]; - if (cache_k_perm->op == GGML_OP_CPY) { - cache_k_perm = cache_k_perm->src[0]; - } - assert(cache_k_perm->op == GGML_OP_PERMUTE); - auto * cache_k_view = cache_k_perm->src[0]; - assert(cache_k_view->op == GGML_OP_VIEW); - - auto * cache_k = cache_k_view->src[0]; - int layer = extract_layer_from_name(cache_k->name); - auto * mask = node->src[3]; - std::string mask_name(mask->name); - - model_params.kv_buffer_ctx_id = ggml_backend_openvino_buffer_get_ctx_id(cache_k->buffer); - if (mask_name.find("swa") != std::string::npos) { - model_params.swa_layers.push_back(layer); - model_params.ctx_per_seq_swa = cache_k->ne[1]; - } else { - model_params.ctx_per_seq = cache_k->ne[1]; - model_params.n_seq = cache_k->ne[2]; - } - - compute_params.n_seq_active = mask->ne[3]; - auto seq_size = cache_k->ne[0] * cache_k->ne[1] * ggml_type_size(cache_k->type); - size_t offset; - memcpy(&offset, cache_k_view->op_params, sizeof(size_t)); - compute_params.seq_active_start = offset / seq_size; - compute_params.token_len_per_seq = node->ne[2]; - - if (mask_name.find("swa") != std::string::npos) { - compute_params.attention_size_swa = mask->ne[0]; - } else { - compute_params.attention_size = mask->ne[0]; - } - if (is_static) { - compute_params.attention_size = model_params.ctx_per_seq; - compute_params.attention_size_swa = model_params.ctx_per_seq_swa; - compute_params.token_len_per_seq = 1; - } - break; - } - if (node->op == GGML_OP_ROPE) { - memcpy(model_params.rope_params, node->op_params, sizeof(int32_t) * 15); - } - } - auto * output_tensor = cgraph->nodes[cgraph->n_nodes - 1]; - compute_params.output_len = output_tensor->ne[1]; - // for NPU, output_len is always 1 except for llama-perplexity - if (is_static && compute_params.output_len == 0) { - compute_params.output_len = 1; - } - model_params.ctx = model_params.ctx_per_seq * model_params.n_seq; - model_params.ctx_swa = model_params.ctx_per_seq_swa * model_params.n_seq; - return {model_params, compute_params}; -} - -void GgmlOvDecoder::validate_cgraph() const { - if (m_model_params.n_seq > 1 && m_is_static == true) { - throw std::runtime_error("n_seq > 1 is not supported on NPU. Try setting -np 1."); - } -} - -ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * op, const ggml_tensor * input) const { - if (m_naive) { - return input!= nullptr ? ov::PartialShape{get_shape(input)} : ov::PartialShape{get_shape(op)}; - } - auto name = std::string(input->name); - ov::PartialShape input_shape; - - if (is_inp_tok(input, op) || is_inp_pos(input, op)) { - // tokens or positions - int len = m_is_static ? (m_is_prefill ? m_prefill_chunk_size : 1) : -1; - input_shape = ov::PartialShape{1, 1, 1, len}; - - } else if (is_output_idx(input, op)) { - // output index - input_shape = ov::PartialShape{1, 1, 1, m_is_static ? m_compute_params.output_len : -1}; - - } else if (is_inp_mask(input, op)) { - // mask - if (m_is_static) { - input_shape = ov::PartialShape{1, 1, m_is_prefill ? m_prefill_chunk_size : 1, m_model_params.ctx}; - } else if (m_is_stateful) { - input_shape = ov::PartialShape{1, 1, -1, -1}; - } else { - input_shape = ov::PartialShape{-1, 1, -1, -1}; - } - - } else if (is_kvcache(input, op)) { - // kvcache - input_shape = ov::PartialShape{get_shape(input)}; - if (!m_is_static) { - // do not fix ctx size to make llama-bench work across test params - input_shape[2] = -1; - } - if (is_stateful()) { - // Convert stateless KV cache layout [1, 1, seq, n_heads_kv * head_size] - // to stateful layout [1, seq, n_heads_kv, head_size]. - assert(input_shape.size() == 4 && input_shape[0] == 1 && input_shape[1] == 1 && - input_shape[2].is_dynamic() && - input_shape[3] == (m_model_params.n_heads_kv * m_model_params.head_size)); - input_shape = {input_shape[0], ov::Dimension::dynamic(), m_model_params.n_heads_kv, - m_model_params.head_size}; - } - - } else if (is_kv_idx(input, op)) { - // kv update index - int len = m_is_static ? (m_is_prefill ? m_prefill_chunk_size : 1) : -1; - input_shape = ov::PartialShape{1, 1, 1, len}; - - } else { - input_shape = ov::PartialShape{get_shape(input)}; - } - return input_shape; -} - -void GgmlOvDecoder::add_extra_inputs() { - // Extra inputs: - // 1. `attention_size`, used in FLASH_ATTN where the shape of the matmul's are 256 aligned, - // see llama_kv_cache_unified::get_n_kv and llama_kv_cache_unified::get_padding. - // 2. `n_seq_active` and `seq_active_start`, used in FLASH_ATTN_EXT to indicate the active sequences in the batch - - auto create_1d_input = [this](const std::string & name, int64_t value) { - if (m_is_static) { - auto constant = - std::make_shared(ov::element::i64, ov::Shape{1}, std::vector{value}); - constant->set_friendly_name(name); - m_model_extra_inputs[name] = constant; - } else { - auto param_node = std::make_shared(ov::element::i64, ov::Shape{1}); - param_node->set_friendly_name(name); - param_node->output(0).get_tensor().set_names({name}); - m_model_extra_inputs[name] = param_node; - - auto tensor = std::make_shared(ov::element::i64, ov::Shape{1}); - *tensor->data() = value; - m_model_extra_input_values[name] = tensor; - } - }; - - create_1d_input("attention_size", m_compute_params.attention_size); - if (m_compute_params.attention_size_swa != -1) { - create_1d_input("attention_size_swa", m_compute_params.attention_size_swa); - } - create_1d_input("n_seq_active", m_compute_params.n_seq_active); - create_1d_input("seq_active_start", m_compute_params.seq_active_start); - create_1d_input("seq_active_end", m_compute_params.seq_active_start + m_compute_params.n_seq_active); - create_1d_input("token_len_per_seq", m_compute_params.token_len_per_seq); - // create_1d_input("token_len", m_token_len_per_seq * m_n_seq_active); -} - -bool GgmlOvDecoder::node_is_used_as_src(const int node_idx) { - ggml_tensor * node = m_cgraph->nodes[node_idx]; - for (int i = node_idx; i < m_cgraph->n_nodes; i++) { - ggml_tensor * other_node = m_cgraph->nodes[i]; - for (int j = 0; j < GGML_MAX_SRC; j++) { - if (other_node->src[j] == node) { - return true; - } - } - } - return false; -} - -void GgmlOvDecoder::compute_model_inputs() { - m_model_inputs.clear(); - m_inputs.clear(); - for (int i = 0; i < m_cgraph->n_nodes; i++) { - ggml_tensor * node = m_cgraph->nodes[i]; - // the node op is NONE means this node maybe as input of later nodes, we should add it to model inputs for this node. - if (node->op == GGML_OP_NONE && node_is_used_as_src(i)) { - std::string node_name(node->name); - if (m_model_weights.find(node_name) == m_model_weights.end()) { - m_inputs[node_name] = node; - auto param_node = - std::make_shared(get_ov_type(node), get_graph_input_shape(node, nullptr)); - param_node->set_friendly_name(node_name); - param_node->output(0).get_tensor().set_names({node_name}); - m_model_inputs[node_name] = param_node; - } - continue; - } - for (int i = 0; i < GGML_MAX_SRC; i++) { - auto * src = node->src[i]; - if (src == nullptr) { - continue; - } - std::string src_name = std::string(src->name); - if (src->flags & GGML_TENSOR_FLAG_INPUT) { - src_name = get_graph_input_ov_name(src, node); - } - if (m_model_weights.find(src_name) != m_model_weights.end()) { - continue; - } - - bool is_intermediate_node = false; - for (const auto & node_info : m_node_info_list) { - if (node_info.node == src) { - is_intermediate_node = true; - break; - } - } - if (is_intermediate_node) { - continue; - } - if (m_model_inputs.find(src_name) != m_model_inputs.end()) { - continue; - } - - m_inputs[src_name] = src; - - ggml_backend_buffer * buffer = src->buffer; - // GGML_BACKEND_BUFFER_USAGE_ANY are kv caches - if (buffer->usage == GGML_BACKEND_BUFFER_USAGE_ANY) { - if (auto it = std::find(m_model_params.kv_names.begin(), m_model_params.kv_names.end(), src_name); - it == m_model_params.kv_names.end()) { - m_model_params.kv_names.push_back(src_name); - } - } - ov::PartialShape param_shape = get_graph_input_shape(node, src); - auto param_node = std::make_shared(get_ov_type(src), param_shape); - param_node->set_friendly_name(src_name); - param_node->output(0).get_tensor().set_names({src_name}); - m_model_inputs[src_name] = param_node; - } - } -} - -void GgmlOvDecoder::compute_model_outputs() { - m_model_outputs.clear(); - m_model_output_names.clear(); - for (int node_n = 0; node_n < m_cgraph->n_nodes; node_n++) { - auto * cur_node = m_cgraph->nodes[node_n]; - // if the node op is NONE means this node is not used at all, we can skip it directly without adding to model outputs. - if (cur_node->op == GGML_OP_NONE) { - continue; - } - auto cur_node_use_count = m_cgraph->use_counts[ggml_hash_find(&m_cgraph->visited_hash_set, cur_node)]; - if (cur_node_use_count == 0) { - // The output of SET_ROWS is the view_src tensor, which is updated in place. We should use the view_src name as the output name to make sure it can be correctly matched with the later ops that use the view_src. - if (cur_node != nullptr && cur_node->op == GGML_OP_SET_ROWS) { - cur_node = cur_node->view_src; - } - } else { - int input_use_count = 0; - for (int i = 0; i < m_cgraph->n_nodes; i++) { - ggml_tensor * node = m_cgraph->nodes[i]; - for (int j = 0; j < GGML_MAX_SRC; j++) { - if (node->src[j] != NULL && node->src[j] == cur_node) { - input_use_count++; - } - } - } - if (input_use_count == cur_node_use_count) { - cur_node = nullptr; - } - } - if (cur_node != nullptr) { - std::string node_output_name(cur_node->name); - m_model_outputs[node_output_name] = cur_node; - m_model_output_names.push_back(node_output_name); - } - } -} - -const ggml_tensor * GgmlOvDecoder::get_tensor_used_op(const ggml_tensor * tensor) const { - if (tensor == nullptr) { - return nullptr; - } - for (int i = 0; i < m_cgraph->n_nodes; i++) { - const auto * node = m_cgraph->nodes[i]; - for (int j = 0; j < GGML_MAX_SRC; j++) { - if (node->src[j] == tensor) { - return node; - } - } - } - return nullptr; -} - -const ggml_tensor * GgmlOvDecoder::get_tensor_from_name(const std::string & name) const { - for (int i = 0; i < m_cgraph->n_nodes; i++) { - const auto * node = m_cgraph->nodes[i]; - for (int j = 0; j < GGML_MAX_SRC; j++) { - const auto * src = node->src[j]; - if (src == nullptr) { - break; - } - if (std::string(src->name) == name) { - return src; - } - } - } - return nullptr; -} - -std::map GgmlOvDecoder::get_kv_param_res_names() const { - std::map kv_param_res_names; - for (const auto & name : m_model_params.kv_names) { - kv_param_res_names[name] = name; - } - return kv_param_res_names; -} - -std::map> GgmlOvDecoder::create_weight_nodes(ggml_cgraph * cgraph, bool naive) { - std::map> model_weights; - auto * nodes = cgraph->nodes; - auto n_nodes = cgraph->n_nodes; - for (int node_i = 0; node_i < n_nodes; node_i++) { - auto * node = nodes[node_i]; - for (int i = 0; i < GGML_MAX_SRC; i++) { - auto * src = node->src[i]; - if (src == nullptr) { - continue; - } - - std::string src_name(src->name); - if (is_rope_freqs_weight(src, node)) { - src_name = "rope_freqs.weight"; - } - if (!src->view_src) { - ggml_backend_buffer * buffer = src->buffer; - if (buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS || ggml_is_quantized(src->type)) { - if (model_weights.find(src_name) == model_weights.end()) { - auto weight_node = create_weight_node(src, naive); - weight_node->set_friendly_name(src_name); - model_weights[src_name] = weight_node; - } - } - } - } - } - return model_weights; -} - -std::shared_ptr GgmlOvDecoder::create_weight_node(ggml_tensor * tensor, bool naive) { - const bool is_ov_buffer = ggml_backend_buffer_is_openvino(tensor->buffer); - - // Check if we have a pre-built constant from the OpenVINO backend buffer - // This is set during ggml_backend_openvino_buffer_set_tensor - if (tensor->extra) { - OPENVINO_ASSERT(is_ov_buffer, "Unsupported weight tensor: " + std::string(tensor->name) + - " Possibly this is a cpu backend repacked quantized weights"); - // Cast to our extra base type and check the type - auto * extra_base = static_cast(tensor->extra); - - if (extra_base->type == ggml_openvino_extra_base::Type::WEIGHT) { - // F16/F32/BF16 weight with shared-memory constant - auto * weight_extra = static_cast(tensor->extra); - if (weight_extra->weight_node) { - // GGML_LOG_DEBUG("%s: using pre-built weight node for %s\n", __func__, tensor->name); - return weight_extra->weight_node; - } - } else if (extra_base->type == ggml_openvino_extra_base::Type::QUANTIZED_WEIGHT) { - // Quantized weight with pre-extracted data - auto * quant_extra = static_cast(tensor->extra); - if (quant_extra->weight_node) { - // GGML_LOG_DEBUG("%s: using pre-extracted quantized weight node for %s\n", __func__, tensor->name); - return quant_extra->weight_node; - } - } - } - - // There are three cases where we need to create a new weight node: - // 1. weights are in openvino_host_buffer. Weight loading to host buffer will not trigger backend_buffer_set_tensor - // 2. weights are in cpu/cpu_mapped buffer. On token_embd.weight goes to case 1 or 2, depending on whether mmap or direct_io is used - // 3. test-backend-ops. buffers in test-backend-ops does not set USAGE_WEIGHT so backend_buffer_set_tensor will not create weight node - - // GGML_LOG_DEBUG("%s: creating new weight node for %s\n", __func__, tensor->name); - static const std::set weight_types = {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, - GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, - GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K}; - if (weight_types.find(tensor->type) == weight_types.end()) { - throw std::runtime_error("Unexpected weight tensor type: " + std::string(tensor->name) + " with type " + - ggml_type_name(tensor->type)); - } - - OvWeight ov_weight; - if (ggml_is_quantized(tensor->type)) { - auto use_bias = naive; - if (is_ov_buffer) { - // For quantized weights, copy raw data to a temp buffer first because - // process_weight_tensor reads from data and writes extracted results - // (weights/scales/zp) to output_base_ptr — they would overlap if both - // point to tensor->data. - size_t raw_size = ggml_nbytes(tensor); - std::vector tmp(raw_size); - memcpy(tmp.data(), tensor->data, raw_size); - ov_weight = process_weight_tensor(tensor, tmp.data(), tensor->data, use_bias); - } else { - ov_weight = process_weight_tensor(tensor, tensor->data, nullptr, use_bias); - } - } else { - // For non-quantized weights (F16/F32/BF16), data is already in tensor->data. - // process_weight_tensor will create an ov::Tensor wrapping tensor->data directly. - ov_weight = process_weight_tensor(tensor, tensor->data, tensor->data); - } - - ov_weight.weight_node->set_friendly_name(tensor->name); - if (!is_ov_buffer) { - return ov_weight.weight_node; - } - - ggml_openvino_extra_base * extra; - if (ov_weight.is_quantized()) { - extra = new ggml_openvino_quantized_weight_extra(std::move(ov_weight.weights), std::move(ov_weight.scales), - std::move(ov_weight.zp), ov_weight.weight_node); - } else { - extra = new ggml_openvino_weight_extra(std::move(ov_weight.weights), ov_weight.weight_node); - } - ggml_openvino_buffer_register_extra(tensor, extra); - - return ov_weight.weight_node; -} - -void GgmlOvDecoder::dump_cgraph(const ggml_cgraph * cgraph, std::string & filename) { - std::ofstream file(filename); - if (!file.is_open()) { - std::cerr << "Failed to open file" << std::endl; - return; - } - - file << "=== GRAPH ===\n"; - - // clang-format off - file << "n_nodes = " << cgraph->n_nodes << "\n"; - file << " " << std::setw(3) << "nodes" - << std::setw(15) << "shape" - << std::setw(20) << "op" - << std::setw(20) << "name" - << std::setw(3) << " " - << std::setw(62) << "stride" - << std::setw(20) << "buffer_type" - << "\n"; - for (int i = 0; i < cgraph->n_nodes; i++) { - ggml_tensor * node = cgraph->nodes[i]; - - // Get buffer type name - const char * buf_name = "none"; - ggml_backend_buffer_t buf = node->view_src ? node->view_src->buffer : node->buffer; - if (buf) { - buf_name = ggml_backend_buffer_name(buf); - } - - file << " - " << std::setw(3) << i << ": [ " - << std::setw(5) << node->ne[0] << ", " - << std::setw(5) << node->ne[1] << ", " - << std::setw(5) << node->ne[2] << ", " - << std::setw(5) << node->ne[3] << "] " - << std::left << std::setw(20) << ggml_op_name(node->op) << std::right << " " - << std::left << std::setw(45) << node->name << std::right - << std::setw(2) << "[ " - << std::setw(0) << node->nb[0] << ", " - << std::setw(5) << node->nb[1] << ", " - << std::setw(5) << node->nb[2] << ", " - << std::setw(5) << node->nb[3] << "] " - << std::right << std::setw(15) << buf_name << std::right - << "\n"; - - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (auto* src = node->src[i]) { - // Get buffer type name for source - const char * src_buf_name = "none"; - ggml_backend_buffer_t src_buf = src->view_src ? src->view_src->buffer : src->buffer; - if (src_buf) { - src_buf_name = ggml_backend_buffer_name(src_buf); - } - - file << std::setw(10) << " [ " - << std::setw(5) << src->ne[0] << ", " - << std::setw(5) << src->ne[1] << ", " - << std::setw(5) << src->ne[2] << ", " - << std::setw(5) << src->ne[3] << "] " - << std::setw(12) - << i << ": " << std::left << std::setw(12) << ggml_op_name(src->op) << std::right; - file << std::left << std::setw(30) << src->name << std::right - << std::setw(16) << "[ " - << std::setw(0) << src->nb[0] << ", " - << std::setw(5) << src->nb[1] << ", " - << std::setw(5) << src->nb[2] << ", " - << std::setw(5) << src->nb[3] << "] " - << std::right << std::setw(15) << src_buf_name << std::right - << "\n"; - } - } - } - - file << "n_leafs = " << cgraph->n_leafs << "\n"; - for (int i = 0; i < cgraph->n_leafs; i++) { - ggml_tensor * node = cgraph->leafs[i]; - - // Get buffer type name for leaf - const char * leaf_buf_name = "none"; - ggml_backend_buffer_t leaf_buf = node->view_src ? node->view_src->buffer : node->buffer; - if (leaf_buf) { - leaf_buf_name = ggml_backend_buffer_name(leaf_buf); - } - - file << " - " << std::setw(3) << i << ": [ " - << std::setw(5) << node->ne[0] << ", " - << std::setw(5) << node->ne[1] << "] " - << std::setw(8) << ggml_op_name(node->op) << " " - << std::setw(16) << ggml_get_name(node) - << std::setw(20) << leaf_buf_name << "\n"; - } - // clang-format on - file << "========================================\n"; - - file.close(); -} - -void print_tensor_address_map(const ggml_cgraph * cgraph) { - std::map> address_map; - for (int node_n = 0; node_n < cgraph->n_nodes; node_n++) { - auto * node = cgraph->nodes[node_n]; - if (node->data) { - auto it = address_map.find(node->data); - if (it == address_map.end()) { - address_map[node->data] = std::vector(); - } - address_map[node->data].push_back(node->name); - } - } - for (const auto & pair : address_map) { - std::cout << "Address: " << pair.first << std::endl; - for (const auto & name : pair.second) { - std::cout << name << " ; "; - } - std::cout << std::endl << std::endl; - } -} - -ov::Shape GgmlOvDecoder::get_shape(const ggml_tensor * tensor) { - std::vector shape; - for (int i = GGML_MAX_DIMS - 1; i >= 0; --i) { - shape.push_back(static_cast(tensor->ne[i])); - } - return shape; -} - -std::vector GgmlOvDecoder::get_stride(const ggml_tensor * tensor) { - std::vector stride; - for (int i = GGML_MAX_DIMS - 1; i >= 0; --i) { - stride.push_back(static_cast(tensor->nb[i])); - } - return stride; -} - -ov::element::Type GgmlOvDecoder::get_ov_type(const ggml_tensor * tensor) { - switch (tensor->type) { - case GGML_TYPE_F64: - return ov::element::f64; - case GGML_TYPE_F32: - return ov::element::f32; - case GGML_TYPE_F16: - return ov::element::f16; - case GGML_TYPE_BF16: - return ov::element::bf16; - case GGML_TYPE_I8: - return ov::element::i8; - case GGML_TYPE_I16: - return ov::element::i16; - case GGML_TYPE_I32: - return ov::element::i32; - case GGML_TYPE_I64: - return ov::element::i64; - default: - return ov::element::dynamic; - } -} - -ov::PartialShape GgmlOvDecoder::get_input_shape(int node_idx, const std::string & name) const { - return ov::PartialShape(get_shape(m_node_info_list[node_idx].node_inputs.at(name))); -} - -std::vector GgmlOvDecoder::get_input_stride(int node_idx, const std::string & name) const { - return get_stride(m_node_info_list[node_idx].node_inputs.at(name)); -} - -ov::element::Type GgmlOvDecoder::get_input_type(int node_idx, const std::string & name) const { - return get_ov_type(m_node_info_list[node_idx].node_inputs.at(name)); -} - -size_t GgmlOvDecoder::get_input_size() const { - return m_model_inputs.size(); -} - -size_t GgmlOvDecoder::get_input_size(int node_idx) const { - return m_node_info_list[node_idx].node_inputs_names.size(); -} - -std::vector GgmlOvDecoder::get_input_names(int node_idx) const { - return m_node_info_list[node_idx].node_inputs_names; -} - -ov::PartialShape GgmlOvDecoder::get_output_shape(int node_idx) const { - auto * ggml_tensor = m_node_info_list[node_idx].node_output; - return ov::PartialShape(get_shape(ggml_tensor)); -} - -ov::element::Type GgmlOvDecoder::get_output_type(const int node_idx) const { - return get_ov_type(m_node_info_list[node_idx].node); -} - -std::vector GgmlOvDecoder::get_output_names(int node_idx) const { - return {m_node_info_list[node_idx].node_output_name}; -} - -const std::string & GgmlOvDecoder::get_op_name() const { - static const std::string unknown_name = "UNKNOWN_OP_NAME"; - return unknown_name; -} - -const std::string & GgmlOvDecoder::get_op_name(int node_idx) const { - return m_node_info_list[node_idx].node_name; -} - -int32_t * GgmlOvDecoder::get_input_op_params(int node_idx, const std::string & name) const { - return m_node_info_list[node_idx].node_inputs.at(name)->op_params; -} - -int32_t * GgmlOvDecoder::get_output_op_params(int node_idx) const { - return m_node_info_list[node_idx].node->op_params; -} - -void GgmlOvDecoder::visit_subgraph(std::function, int node_idx)> node_visitor) const { - for (int node_idx = 0; node_idx < m_cgraph->n_nodes; node_idx++) { - if (m_cgraph->nodes[node_idx]->op == GGML_OP_NONE) { - continue; - } - node_visitor(std::make_shared(*this), node_idx); - } -} - -std::string GgmlOvDecoder::compute_op_type(const ggml_tensor * node) { - static const std::map ops = { - {GGML_OP_NONE, "GGML_OP_NONE" }, - {GGML_OP_ACC, "GGML_OP_ACC" }, - {GGML_OP_ADD, "GGML_OP_ADD" }, - {GGML_OP_ADD1, "GGML_OP_ADD1" }, - {GGML_OP_CONT, "GGML_OP_CONT" }, - {GGML_OP_DIV, "GGML_OP_DIV" }, - {GGML_OP_DUP, "GGML_OP_DUP" }, - {GGML_OP_GET_ROWS, "GGML_OP_GET_ROWS" }, - {GGML_OP_MUL, "GGML_OP_MUL" }, - {GGML_OP_MUL_MAT, "GGML_OP_MUL_MAT" }, - {GGML_OP_PERMUTE, "GGML_OP_PERMUTE" }, - {GGML_OP_RESHAPE, "GGML_OP_RESHAPE" }, - {GGML_OP_RMS_NORM, "GGML_OP_RMS_NORM" }, - {GGML_OP_ROPE, "GGML_OP_ROPE" }, - {GGML_OP_SCALE, "GGML_OP_SCALE" }, - {GGML_OP_SOFT_MAX, "GGML_OP_SOFT_MAX" }, - {GGML_OP_SUB, "GGML_OP_SUB" }, - {GGML_OP_TRANSPOSE, "GGML_OP_TRANSPOSE" }, - {GGML_OP_VIEW, "GGML_OP_VIEW" }, - {GGML_OP_SET_ROWS, "GGML_OP_SET_ROWS" }, - {GGML_OP_CPY, "GGML_OP_CPY" }, - {GGML_OP_FLASH_ATTN_EXT, "GGML_OP_FLASH_ATTN_EXT"}, - }; - static const std::map unary_ops = { - {GGML_UNARY_OP_ABS, "GGML_UNARY_OP_ABS" }, - {GGML_UNARY_OP_SGN, "GGML_UNARY_OP_SGN" }, - {GGML_UNARY_OP_NEG, "GGML_UNARY_OP_NEG" }, - {GGML_UNARY_OP_STEP, "GGML_UNARY_OP_STEP" }, - {GGML_UNARY_OP_TANH, "GGML_UNARY_OP_TANH" }, - {GGML_UNARY_OP_ELU, "GGML_UNARY_OP_ELU" }, - {GGML_UNARY_OP_RELU, "GGML_UNARY_OP_RELU" }, - {GGML_UNARY_OP_SIGMOID, "GGML_UNARY_OP_SIGMOID" }, - {GGML_UNARY_OP_GELU, "GGML_UNARY_OP_GELU" }, - {GGML_UNARY_OP_GELU_QUICK, "GGML_UNARY_OP_GELU_QUICK" }, - {GGML_UNARY_OP_SILU, "GGML_UNARY_OP_SILU" }, - {GGML_UNARY_OP_HARDSWISH, "GGML_UNARY_OP_HARDSWISH" }, - {GGML_UNARY_OP_HARDSIGMOID, "GGML_UNARY_OP_HARDSIGMOID"}, - {GGML_UNARY_OP_EXP, "GGML_UNARY_OP_EXP" }, - {GGML_UNARY_OP_COUNT, "GGML_UNARY_OP_COUNT" } - }; - static const std::map glu_ops = { - {GGML_GLU_OP_SWIGLU, "GGML_GLU_OP_SWIGLU"}, - {GGML_GLU_OP_GEGLU, "GGML_GLU_OP_GEGLU" }, - {GGML_GLU_OP_REGLU, "GGML_GLU_OP_REGLU" } - }; - - switch (node->op) { - case GGML_OP_UNARY: - return unary_ops.at(ggml_get_unary_op(node)); - case GGML_OP_GLU: - return glu_ops.at(ggml_get_glu_op(node)); - default: - return ops.at(node->op); - } - static const std::string unknown_op = "UNKNOWN_GGML_OP"; - return unknown_op; -} - -const std::string & GgmlOvDecoder::get_op_type(int node_idx) const { - return m_node_info_list[node_idx].node_op_type; -} - -const std::string & GgmlOvDecoder::get_op_type() const { - static const std::string unknown_op = "UNKNOWN_GGML_OP"; - return unknown_op; -} diff --git a/ggml/src/ggml-openvino/ggml-decoder.h b/ggml/src/ggml-openvino/ggml-decoder.h deleted file mode 100644 index 3ae25ddda320..000000000000 --- a/ggml/src/ggml-openvino/ggml-decoder.h +++ /dev/null @@ -1,294 +0,0 @@ -#pragma once - -#include "ggml-quants.h" -#include "ggml.h" -#include "openvino/decoder.h" - -#include -#include -#include -#include -#include -#include -#include - -struct ModelParams { - int ctx = -1; - int ctx_swa = -1; - int ctx_per_seq = -1; - int ctx_per_seq_swa = -1; - int n_seq = 1; - int n_heads = -1; - int n_heads_kv = -1; - int head_size = -1; - int32_t rope_params[15]; - std::vector swa_layers; - - std::vector kv_names; - size_t kv_buffer_ctx_id = 0; - - bool same_rope_params(const ModelParams & other) const { - return memcmp(rope_params, other.rope_params, sizeof(int32_t) * 15) == 0; - } - - bool can_reuse_dynamically(const ModelParams & other) const { return same_rope_params(other); } - - bool can_reuse_statically(const ModelParams & other) const { return same_rope_params(other) && ctx == other.ctx; } - - bool kv_buffer_changed(const ModelParams & other) const { return kv_buffer_ctx_id != other.kv_buffer_ctx_id; } -}; - -struct ComputeParams { - int n_seq_active = 1; - int seq_active_start = 0; - int attention_size = -1; - int attention_size_swa = -1; - int input_len = -1; - int token_len_per_seq = -1; - int past_kv_len = -1; - int output_len = 1; -}; - -class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder { -public: - struct NodeInfo { - ggml_tensor * node; - std::string node_name; - std::string node_op_type; - std::map node_inputs; - std::vector node_inputs_names; - ggml_tensor * node_output; - std::string node_output_name; - int node_op_case = 0; - void * data_addr; - }; - // Graph decoder - GgmlOvDecoder(ggml_cgraph * cgraph, - ModelParams & model_params, - ComputeParams & compute_params, - std::map> & model_weights, - bool is_static, - bool is_stateful = false, - bool is_prefill = false, - int prefill_chunk_size = 256); - - // Naive graph decoder - GgmlOvDecoder(ggml_cgraph * cgraph, std::map> & model_weights); - - virtual ov::Any get_attribute(const std::string & name) const override { - return nullptr; - GGML_UNUSED(name); - } - - virtual ov::PartialShape get_input_shape(int node_idx, const std::string & name) const override; - - virtual std::vector get_input_stride(int node_idx, const std::string & name) const override; - - virtual ov::element::Type get_input_type(int node_idx, const std::string & name) const override; - - virtual size_t get_input_size() const override; - - virtual size_t get_input_size(int node_idx) const override; - - virtual void get_input_node(size_t input_port_idx, - std::string & producer_name, - std::string & producer_output_port_name, - size_t & producer_output_port_index) const override { - GGML_UNUSED(input_port_idx); - GGML_UNUSED(producer_name); - GGML_UNUSED(producer_output_port_name); - GGML_UNUSED(producer_output_port_index); - } - - virtual std::vector get_input_names(int node_idx) const override; - - virtual ov::PartialShape get_output_shape(int node_idx) const override; - - virtual ov::element::Type get_output_type(int node_idx) const override; - - virtual int32_t * get_input_op_params(int node_idx, const std::string & name) const override; - - virtual int32_t * get_output_op_params(int node_idx) const override; - - virtual std::vector get_output_names(int node_idx) const override; - - virtual const std::string & get_op_type() const override; - - virtual const std::string & get_op_type(int node_idx) const override; - - virtual const std::string & get_op_name() const override; - - virtual const std::string & get_op_name(int node_idx) const override; - - virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const override; - - ggml_tensor * get_input_ggml_tensor(const std::string & name) const { return m_inputs.at(name); } - - virtual int get_op_case(int node_idx) const override { return m_node_info_list[node_idx].node_op_case; } - - virtual const std::map> & get_model_inputs() const override { - return m_model_inputs; - } - - virtual const std::map> & get_model_extra_inputs() const override { - return m_model_extra_inputs; - } - - virtual const std::map> & get_model_extra_input_values() const { - return m_model_extra_input_values; - } - - virtual const std::map> & get_model_weights() const override { - return m_model_weights; - } - - virtual std::vector get_model_output_names() const override { - return m_model_output_names; - } - - const std::map & get_model_outputs() const { return m_model_outputs; } - - virtual int get_ctx_size() const { return m_model_params.ctx; } - - virtual int get_ctx_swa_size() const { return m_model_params.ctx_swa; } - - virtual int get_ctx_per_seq() const { return m_model_params.ctx_per_seq; } - - virtual int get_ctx_per_seq_swa() const { return m_model_params.ctx_per_seq_swa; } - - virtual int get_n_seq() const { return m_model_params.n_seq; } - - virtual int is_swa_layer(int layer) const override { - return std::find(m_model_params.swa_layers.begin(), m_model_params.swa_layers.end(), layer) != - m_model_params.swa_layers.end(); - } - - int get_past_kv_len() const { return m_compute_params.past_kv_len; } - - int get_input_len() const { return m_compute_params.input_len; } - - virtual int32_t * get_rope_params() const override { return const_cast(m_model_params.rope_params); } - - virtual std::map get_kv_param_res_names() const override; - - virtual bool is_static() const override { return m_is_static; } - - virtual bool is_stateful() const override { return m_is_stateful; } - - ov::PartialShape get_graph_input_shape(const ggml_tensor * op, const ggml_tensor * input) const; - - static void dump_cgraph(const ggml_cgraph * cgraph, std::string & filename); - - static std::shared_ptr create_weight_node(ggml_tensor * tensor, bool naive = false); - - static std::map> create_weight_nodes(ggml_cgraph * cgraph, - bool naive = false); - - const ggml_tensor * get_tensor_used_op(const ggml_tensor * tensor) const; - - const ggml_tensor * get_tensor_from_name(const std::string & name) const; - - void clear_model_weights() { m_model_weights.clear(); } - - static std::pair compute_llm_params(ggml_cgraph * cgraph, bool is_static); - - ModelParams get_model_params() const { return m_model_params; } - - ComputeParams get_compute_params() const { return m_compute_params; } - - void set_model_params(const ModelParams & model_params) { m_model_params = model_params; } - - void set_compute_params(const ComputeParams & compute_params) { m_compute_params = compute_params; } - - bool m_is_static = false; - bool m_is_stateful = false; - bool m_is_prefill = false; - bool m_naive = false; - int m_prefill_chunk_size = 0; - - static ov::Shape get_shape(const ggml_tensor * tensor); - static std::vector get_stride(const ggml_tensor * tensor); - static ov::element::Type get_ov_type(const ggml_tensor * tensor); - static std::string compute_op_type(const ggml_tensor * node); - void add_extra_inputs(); - - void update_io(ggml_cgraph * cgraph); - - inline static bool is_inp_tok(const ggml_tensor * tensor, const ggml_tensor * op) { - return op->op == GGML_OP_GET_ROWS && tensor == op->src[1] && op->src[0]->op == GGML_OP_NONE; - } - - inline static bool is_inp_pos(const ggml_tensor * tensor, const ggml_tensor * op) { - return op->op == GGML_OP_ROPE && tensor == op->src[1]; - } - - inline static bool is_inp_emb(const ggml_tensor * tensor, const ggml_tensor * op) { - return tensor->op == GGML_OP_GET_ROWS && op->op == GGML_OP_RMS_NORM; - } - - inline static bool is_inp_mask(const ggml_tensor * tensor, const ggml_tensor * op) { - return op->op == GGML_OP_CPY || (op->op == GGML_OP_FLASH_ATTN_EXT && tensor == op->src[3]); - } - - inline static bool is_rope_freqs_weight(const ggml_tensor * tensor, const ggml_tensor * op) { - return op->op == GGML_OP_ROPE && tensor == op->src[2]; - } - - inline static bool is_kvcache(const ggml_tensor * tensor, const ggml_tensor * op) { - return op->op == GGML_OP_SET_ROWS && op->src[2] == tensor; - } - - inline static bool is_kv_idx(const ggml_tensor * tensor, const ggml_tensor * op) { - return op->op == GGML_OP_SET_ROWS && op->src[1] == tensor; - } - - inline static bool is_output_idx(const ggml_tensor * tensor, const ggml_tensor * op) { - return op->op == GGML_OP_GET_ROWS && tensor == op->src[1] && op->src[0]->op != GGML_OP_NONE; - } - - static std::string get_graph_input_ov_name(const ggml_tensor * tensor, const ggml_tensor * op) { - if (is_inp_tok(tensor, op)) { - return "inp_tokens"; - } - if (is_inp_pos(tensor, op)) { - return "inp_pos"; - } - if (is_inp_emb(tensor, op)) { - return "embd"; - } - if (is_output_idx(tensor, op)) { - return "inp_out_ids"; - } - if (is_inp_mask(tensor, op)) { - return std::string(tensor->name).find("swa") == std::string::npos ? "self_kq_mask" : "self_kq_mask_swa"; - } - return tensor->name; - } - -private: - void set_input_output(); - int compute_op_case(const ggml_tensor * node) const; - bool node_is_used_as_src(const int node_idx); - void compute_model_inputs(); - void compute_model_outputs(); - - void validate_cgraph() const; - - ggml_cgraph * m_cgraph = nullptr; - std::map m_inputs; - - std::map> m_model_inputs; - std::map> m_model_extra_inputs; - std::map> m_model_extra_input_values; - std::map> m_model_weights; - std::map m_model_outputs; - std::vector m_model_output_names; - std::vector m_node_info_list; - - ModelParams m_model_params; - ComputeParams m_compute_params; -}; - -void print_tensor_address_map(const ggml_cgraph * cgraph); - -int extract_layer_from_name(const std::string & name); diff --git a/ggml/src/ggml-openvino/ggml-openvino-extra.cpp b/ggml/src/ggml-openvino/ggml-openvino-extra.cpp deleted file mode 100644 index 4140136aca25..000000000000 --- a/ggml/src/ggml-openvino/ggml-openvino-extra.cpp +++ /dev/null @@ -1,380 +0,0 @@ -#include "ggml-openvino-extra.h" - -#include "ggml-impl.h" -#include "ggml.h" - -#include -#include -#include -#include -#include - -ov::Core & ov_singleton_core() { - static ov::Core core; - return core; -} - -// ===================================================== -// Device Configuration Implementations -// ===================================================== - -void ggml_openvino_device_config::init() { - if (initialized) { - return; - } - device_name = getenv("GGML_OPENVINO_DEVICE") ? getenv("GGML_OPENVINO_DEVICE") : "CPU"; - auto available_devices = ov_singleton_core().get_available_devices(); - if (std::find(available_devices.begin(), available_devices.end(), device_name) == available_devices.end()) { - GGML_LOG_WARN("GGML OpenVINO Backend: device %s is not available, fallback to CPU\n", device_name.c_str()); - device_name = "CPU"; - } - is_npu = (device_name == "NPU"); - - auto * cache_dir = getenv("GGML_OPENVINO_CACHE_DIR"); - if (device_name == "NPU") { - compile_config = { - {"NPU_COMPILER_DYNAMIC_QUANTIZATION", "YES" }, - {"NPU_USE_NPUW", "YES" }, - {"NPUW_DEVICES", "NPU" }, - {"NPUW_FOLD", "YES" }, - {"NPUW_WEIGHTS_BANK", "shared"}, - {"NPUW_FUNCALL_FOR_ALL", "YES" }, - {"NPUW_FUNCALL_ASYNC", "YES" }, - {"NPUW_DQ", "YES" }, - {"NPUW_DQ_FULL", "NO" }, - }; - if (cache_dir && strlen(cache_dir) > 0) { - compile_config["NPUW_CACHE_DIR"] = cache_dir; - compile_config.insert(ov::cache_mode(ov::CacheMode::OPTIMIZE_SIZE)); - } - } else if (cache_dir && strlen(cache_dir) > 0) { - compile_config.insert(ov::cache_dir(cache_dir)); - compile_config.insert(ov::cache_mode(ov::CacheMode::OPTIMIZE_SIZE)); - } - - // Initialize remote context with queue sharing for GPU - if (device_name == "GPU") { - // Create OpenCL context and queue - cl_int err; - cl_platform_id platform; - err = clGetPlatformIDs(1, &platform, nullptr); - if (err != CL_SUCCESS) { - GGML_LOG_ERROR("Failed to get OpenCL platform: %d\n", err); - return; - } - - cl_device_id cl_device; - err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_GPU, 1, &cl_device, nullptr); - if (err != CL_SUCCESS) { - GGML_LOG_ERROR("Failed to get OpenCL device: %d\n", err); - return; - } - - cl_context cl_ctx = clCreateContext(nullptr, 1, &cl_device, nullptr, nullptr, &err); - if (err != CL_SUCCESS) { - GGML_LOG_ERROR("Failed to create OpenCL context: %d\n", err); - return; - } - - cl_queue = clCreateCommandQueueWithProperties(cl_ctx, cl_device, nullptr, &err); - if (err != CL_SUCCESS) { - GGML_LOG_ERROR("Failed to create OpenCL command queue: %d\n", err); - clReleaseContext(cl_ctx); - return; - } - - // Create OpenVINO remote context with queue sharing - remote_context = ov::intel_gpu::ocl::ClContext(ov_singleton_core(), cl_queue); - - // Release the context (queue keeps a reference) - clReleaseContext(cl_ctx); - } else if (device_name == "NPU") { - // remote tensor is not used for NPU yet - // remote_context = ov_singleton_core().get_default_context(device_name); - } - - initialized = true; -} - -ggml_openvino_device_config::~ggml_openvino_device_config() { - if (cl_queue != nullptr) { - clReleaseCommandQueue(cl_queue); - cl_queue = nullptr; - } -} - -// Get the global device config singleton -ggml_openvino_device_config & ggml_openvino_get_device_config() { - static ggml_openvino_device_config config; - return config; -} - -// Initialize device config (call during backend init) -void ggml_openvino_init_device_config() { - ggml_openvino_get_device_config().init(); -} - -// Get the device name -const std::string & ggml_openvino_get_device_name() { - return ggml_openvino_get_device_config().device_name; -} - -// Check if running on NPU -bool ggml_openvino_is_npu() { - return ggml_openvino_get_device_config().is_npu; -} - -// Get the remote context for the current device (returns empty optional for CPU) -std::optional ggml_openvino_get_remote_context() { - return ggml_openvino_get_device_config().remote_context; -} - -// Get the compile config for the current device -const ov::AnyMap & ggml_openvino_get_compile_config() { - return ggml_openvino_get_device_config().compile_config; -} - -// Get the OpenCL command queue for GPU operations -cl_command_queue ggml_openvino_get_cl_queue() { - return ggml_openvino_get_device_config().cl_queue; -} - -// Get the clEnqueueMemFillINTEL function pointer (lazy load) -clEnqueueMemFillINTEL_fn ggml_openvino_get_clEnqueueMemFillINTEL() { - static clEnqueueMemFillINTEL_fn fn = nullptr; - static bool loaded = false; - if (!loaded) { - loaded = true; - cl_platform_id platform; - if (clGetPlatformIDs(1, &platform, nullptr) == CL_SUCCESS) { - fn = (clEnqueueMemFillINTEL_fn) clGetExtensionFunctionAddressForPlatform(platform, "clEnqueueMemFillINTEL"); - } - } - return fn; -} - -// Get the clEnqueueMemcpyINTEL function pointer (lazy load) -clEnqueueMemcpyINTEL_fn ggml_openvino_get_clEnqueueMemcpyINTEL() { - static clEnqueueMemcpyINTEL_fn fn = nullptr; - static bool loaded = false; - if (!loaded) { - loaded = true; - cl_platform_id platform; - if (clGetPlatformIDs(1, &platform, nullptr) == CL_SUCCESS) { - fn = (clEnqueueMemcpyINTEL_fn) clGetExtensionFunctionAddressForPlatform(platform, "clEnqueueMemcpyINTEL"); - } - } - return fn; -} - -// Get requantization type for a tensor type (returns nullopt if no requant needed) -std::optional ggml_openvino_get_requant_type(const ggml_tensor * tensor, bool no_requant) { - if (no_requant) { - return std::nullopt; - } - if (strncmp(tensor->name, "token_embd.weight", 17) == 0) { - return ((ggml_openvino_is_npu() && tensor->type == GGML_TYPE_Q6_K) ? ExtraQuantType::F16 : ExtraQuantType::Q8_0_C); - } - if (strncmp(tensor->name, "output.weight", 13) == 0) { - return ExtraQuantType::Q8_0_C; - } - if (ggml_openvino_is_npu()) { - return ExtraQuantType::Q4_0_128; - } - switch (tensor->type) { - case GGML_TYPE_Q6_K: - case GGML_TYPE_Q5_K: - return ExtraQuantType::Q8_0_C; - default: - return std::nullopt; - } -} - -// ===================================================== -// Extracted Layout Calculation -// ===================================================== - -ggml_openvino_extracted_layout ggml_openvino_get_extracted_layout(const ggml_tensor * tensor, bool use_bias) { - ggml_openvino_extracted_layout layout = {}; - layout.is_symmetric = false; - - if (!ggml_is_quantized(tensor->type)) { - return layout; - } - - // Only handle 2D weight tensors - if (tensor->ne[2] != 1 || tensor->ne[3] != 1) { - return layout; - } - - int64_t n_elements = ggml_nelements(tensor); - const size_t alignment = 64; // Good for SIMD - - // Check if requantization is needed (NPU-specific) - auto requant_type = ggml_openvino_get_requant_type(tensor, use_bias); - if (requant_type.has_value()) { - layout.is_requant = true; - layout.requant_type = requant_type; - - // Special case: requant to F16 - just store F16 weights, no scales/zp - if (requant_type.value() == ExtraQuantType::F16) { - layout.weights_size = n_elements * sizeof(uint16_t); // F16 = 2 bytes - layout.total_size = layout.weights_size; - layout.weights_offset = 0; - // No scales/zp for F16 - return layout; - } - - // Requant to different quantized format (e.g., Q4_0_128) - switch (requant_type.value()) { - case ExtraQuantType::Q4_0_128: - layout.is_u4 = true; - layout.weights_per_block = 128; - layout.is_symmetric = true; - break; - case ExtraQuantType::Q4_0_C: - layout.is_u4 = true; - layout.weights_per_block = tensor->ne[0]; - layout.is_symmetric = true; - break; - case ExtraQuantType::Q8_0_32: - layout.is_u4 = false; - layout.weights_per_block = 32; - layout.is_symmetric = true; - break; - case ExtraQuantType::Q8_0_C: - layout.is_u4 = false; - layout.weights_per_block = tensor->ne[0]; - layout.is_symmetric = true; - break; - case ExtraQuantType::Q8_1_C: - layout.is_u4 = false; - layout.weights_per_block = tensor->ne[0]; - break; - default: - layout.weights_per_block = -1; - GGML_ABORT("Code of re-quantizing to channel-wise is not updated"); - break; - } - - if (layout.is_requant) { - // Calculate sizes for requantized format - layout.weights_size = layout.is_u4 ? (n_elements / 2) : n_elements; - int64_t n_blocks = n_elements / layout.weights_per_block; - layout.scales_size = n_blocks * sizeof(uint16_t); - // For symmetric quantization, no zp needed (weights stored as signed) - if (layout.is_symmetric) { - layout.zp_size = 0; - } else { - layout.zp_size = layout.is_u4 ? ((n_blocks + 1) / 2) : n_blocks; - } - - layout.weights_offset = 0; - layout.scales_offset = ((layout.weights_size + alignment - 1) / alignment) * alignment; - layout.zp_offset = layout.scales_offset + ((layout.scales_size + alignment - 1) / alignment) * alignment; - layout.total_size = layout.zp_offset + layout.zp_size; - layout.total_size = std::max(layout.total_size, ggml_nbytes(tensor)); - return layout; - } - } - - // Normal extraction (no requant) - determine format based on tensor type - layout.is_u4 = false; - layout.weights_per_block = 32; - layout.is_symmetric = false; - - switch (tensor->type) { - case GGML_TYPE_Q4_0: - layout.is_u4 = true; - layout.is_symmetric = true; - break; - - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q4_K: - layout.is_u4 = true; - break; - - case GGML_TYPE_Q8_0: - layout.is_symmetric = true; - break; - - case GGML_TYPE_Q6_K: - layout.weights_per_block = 16; - layout.is_symmetric = true; - break; - - case GGML_TYPE_Q5_K: - break; - - default: - // Unsupported quantization type - return layout; - } - - // Calculate sizes - // Weights: U4 = n_elements/2 bytes, U8 = n_elements bytes - layout.weights_size = layout.is_u4 ? (n_elements / 2) : n_elements; - - // Scales: F16 per block - int64_t n_blocks = n_elements / layout.weights_per_block; - layout.scales_size = n_blocks * sizeof(uint16_t); // F16 = 2 bytes - // For symmetric quantization, no zp needed (weights stored as signed) - if (layout.is_symmetric) { - layout.zp_size = 0; - } else { - layout.zp_size = layout.is_u4 ? ((n_blocks + 1) / 2) : n_blocks; - } - - // Layout in buffer: [weights | scales | zp] with alignment - layout.weights_offset = 0; - layout.scales_offset = ((layout.weights_size + alignment - 1) / alignment) * alignment; - layout.zp_offset = layout.scales_offset + ((layout.scales_size + alignment - 1) / alignment) * alignment; - layout.total_size = layout.zp_offset + layout.zp_size; - layout.total_size = std::max(layout.total_size, ggml_nbytes(tensor)); - - return layout; -} - -ggml_openvino_tensor_extra * ggml_openvino_create_tensor_extra(const ggml_tensor * tensor, bool is_remote) { - ov::Shape shape; - for (int i = GGML_MAX_DIMS - 1; i >= 0; --i) { - shape.push_back(static_cast(tensor->ne[i])); - } - - ov::element::Type element_type; - switch (tensor->type) { - case GGML_TYPE_F32: - element_type = ov::element::f32; - break; - case GGML_TYPE_F16: - element_type = ov::element::f16; - break; - case GGML_TYPE_BF16: - element_type = ov::element::bf16; - break; - case GGML_TYPE_I32: - element_type = ov::element::i32; - break; - case GGML_TYPE_I64: - element_type = ov::element::i64; - break; - default: - // GGML_LOG_WARN("%s: unsupported tensor type for ov::Tensor: %s\n", __func__, ggml_type_name(tensor->type)); - return nullptr; - } - - const auto & device_name = ggml_openvino_get_device_name(); - auto remote_context = ggml_openvino_get_remote_context(); - - std::shared_ptr ov_tensor; - if (is_remote) { - GGML_ASSERT(device_name == "GPU"); - auto gpu_context = remote_context->as(); - auto usm_tensor = gpu_context.create_tensor(element_type, shape, tensor->data); - ov_tensor = std::make_shared(std::move(usm_tensor)); - } else { - ov_tensor = std::make_shared(element_type, shape, tensor->data); - } - - return new ggml_openvino_tensor_extra(ov_tensor); -} diff --git a/ggml/src/ggml-openvino/ggml-openvino-extra.h b/ggml/src/ggml-openvino/ggml-openvino-extra.h deleted file mode 100644 index cd0baf4a681b..000000000000 --- a/ggml/src/ggml-openvino/ggml-openvino-extra.h +++ /dev/null @@ -1,182 +0,0 @@ -#pragma once - -#include "ggml.h" -#include "openvino/runtime/core.hpp" - -#define CL_TARGET_OPENCL_VERSION 300 -#include - -#include -#include -#include -#include -#include -#include -#include - -// ExtraQuantType enum - defines requantization target formats -enum class ExtraQuantType { F16, Q4_0_C, Q8_1_C, Q4_0_128, Q8_0_C, Q8_0_32 }; - -ov::Core & ov_singleton_core(); - -// Get the remote context for the current device (returns empty optional for CPU) -std::optional ggml_openvino_get_remote_context(); - -// Get the compile config for the current device -const ov::AnyMap & ggml_openvino_get_compile_config(); - -// Get the OpenCL command queue for GPU operations (returns nullptr for CPU/NPU) -cl_command_queue ggml_openvino_get_cl_queue(); - -// Intel USM extension function type -typedef cl_int(CL_API_CALL * clEnqueueMemFillINTEL_fn)(cl_command_queue queue, - void * dst_ptr, - const void * pattern, - size_t pattern_size, - size_t size, - cl_uint num_events_in_wait_list, - const cl_event * event_wait_list, - cl_event * event); - -typedef cl_int(CL_API_CALL * clEnqueueMemcpyINTEL_fn)(cl_command_queue queue, - cl_bool blocking, - void * dst_ptr, - const void * src_ptr, - size_t size, - cl_uint num_events_in_wait_list, - const cl_event * event_wait_list, - cl_event * event); - -// Get the clEnqueueMemFillINTEL function pointer (returns nullptr if not available) -clEnqueueMemFillINTEL_fn ggml_openvino_get_clEnqueueMemFillINTEL(); - -// Get the clEnqueueMemcpyINTEL function pointer (returns nullptr if not available) -clEnqueueMemcpyINTEL_fn ggml_openvino_get_clEnqueueMemcpyINTEL(); - -// ===================================================== -// Global Device Configuration (singleton) -// ===================================================== -// Initialized once during backend init from GGML_OPENVINO_DEVICE env var - -struct ggml_openvino_device_config { - std::string device_name = "CPU"; - bool is_npu = false; - bool initialized = false; - std::optional remote_context; - ov::AnyMap compile_config; - cl_command_queue cl_queue = nullptr; - - void init(); - ~ggml_openvino_device_config(); -}; - -// Get the global device config singleton -ggml_openvino_device_config & ggml_openvino_get_device_config(); - -// Initialize device config (call during backend init) -void ggml_openvino_init_device_config(); - -// Get the device name -const std::string & ggml_openvino_get_device_name(); - -// Check if running on NPU -bool ggml_openvino_is_npu(); - -// Get requantization type for a tensor type (returns nullopt if no requant needed) -std::optional ggml_openvino_get_requant_type(const ggml_tensor * tensor, bool no_requant = false); - -// ===================================================== -// OpenVINO Tensor Extra Types -// ===================================================== -// These types are stored in tensor->extra by the OpenVINO backend buffer. -// They allow: -// 1. Pre-built ov::Constant nodes for weights (avoiding memcpy during graph construction) -// 2. ov::Tensor wrappers for KV cache / compute tensors (for direct use with infer_request) - -// Base class for OpenVINO tensor extra data -struct ggml_openvino_extra_base { - enum class Type { WEIGHT, QUANTIZED_WEIGHT, TENSOR }; - Type type; - virtual ~ggml_openvino_extra_base() = default; -protected: - explicit ggml_openvino_extra_base(Type t) : type(t) {} -}; - -// Extra data for F16/F32/BF16 weight tensors - stores the pre-built weight node -struct ggml_openvino_weight_extra : public ggml_openvino_extra_base { - ov::Tensor weights; // The underlying weight data tensor - std::shared_ptr weight_node; // Pre-built OpenVINO weight node - - ggml_openvino_weight_extra(ov::Tensor w, std::shared_ptr n) : - ggml_openvino_extra_base(Type::WEIGHT), - weights(std::move(w)), - weight_node(std::move(n)) {} -}; - -// Extra data for quantized weight tensors - stores extracted weights/scales/zp and weight node -struct ggml_openvino_quantized_weight_extra : public ggml_openvino_extra_base { - ov::Tensor weights; // U4 or U8 extracted weights - ov::Tensor scales; // F16 scales - ov::Tensor zp; // U4 or U8 zero points (same type as weights) - std::shared_ptr weight_node; // Pre-built OpenVINO weight subgraph - - ggml_openvino_quantized_weight_extra(ov::Tensor w, ov::Tensor s, ov::Tensor z, std::shared_ptr n) : - ggml_openvino_extra_base(Type::QUANTIZED_WEIGHT), - weights(std::move(w)), - scales(std::move(s)), - zp(std::move(z)), - weight_node(std::move(n)) {} -}; - -// Extra data for KV cache / compute tensors - stores ov::Tensor for infer_request -struct ggml_openvino_tensor_extra : public ggml_openvino_extra_base { - std::shared_ptr tensor; // For direct use with infer_request - - explicit ggml_openvino_tensor_extra(std::shared_ptr t) - : ggml_openvino_extra_base(Type::TENSOR), tensor(std::move(t)) {} -}; - -// ===================================================== -// Extracted Size Calculation for Quantized Tensors -// ===================================================== -// For quantized tensors, we need extra space to store extracted weights, scales, and zero points. -// Returns the total size needed in the buffer for extracted data. - -struct ggml_openvino_extracted_layout { - size_t total_size = 0; // Total bytes needed - size_t weights_offset = 0; // Offset to weights in buffer - size_t weights_size = 0; // Size of weights in bytes - size_t scales_offset = 0; // Offset to scales in buffer - size_t scales_size = 0; // Size of scales in bytes - size_t zp_offset = 0; // Offset to zero points in buffer - size_t zp_size = 0; // Size of zero points in bytes (U4 or U8) - bool is_u4; // true for U4 weights, false for U8 - int64_t weights_per_block; // weights per scale/zp block - bool is_symmetric; // true for symmetric quantization - - // Requantization info - bool is_requant = false; // true if this tensor needs requantization - std::optional requant_type; // target requant type if is_requant -}; - -// Calculate the buffer layout for extracted quantized data -ggml_openvino_extracted_layout ggml_openvino_get_extracted_layout(const ggml_tensor * tensor, bool use_bias = false); - -ggml_openvino_tensor_extra * ggml_openvino_create_tensor_extra(const ggml_tensor * tensor, bool is_remote); - -// Register an extra with the tensor's OpenVINO buffer context for proper lifetime management. -// This sets tensor->extra and tracks the extra in the buffer context for cleanup. -void ggml_openvino_buffer_register_extra(ggml_tensor * tensor, ggml_openvino_extra_base * extra); - -// ===================================================== -// OpenVINO Backend Context and Interface -// ===================================================== -struct ggml_backend_openvino_context { - int device = 0; - std::string name = "OpenVINO"; - std::string description = "OpenVINO Backend Context"; - - std::shared_ptr runtime_context = nullptr; - - ggml_backend_openvino_context() = default; -}; diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp deleted file mode 100644 index 4f3ebf2536b0..000000000000 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ /dev/null @@ -1,1132 +0,0 @@ -#include "ggml-openvino.h" - -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "ggml-impl.h" -#include "ggml-openvino-extra.h" -#include "ggml-openvino/utils.h" -#include "ggml-quants.h" -#include "ggml.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#if defined(_WIN32) -# define WIN32_LEAN_AND_MEAN -# ifndef NOMINMAX -# define NOMINMAX -# endif -# include -#else -# include -#endif - -// ===================================================== -// OpenVINO Buffer Implementation using ov::Tensor -// ===================================================== -// -// Design: This implementation uses a hybrid approach: -// 1. For weight tensors: Store a pre-built ov::op::v0::Constant in tensor->extra -// - This avoids the memcpy during graph construction -// - For quantized weights, the constant is already converted to OpenVINO format -// 2. For KV cache / compute tensors: Store an ov::Tensor in tensor->extra -// - This can be directly passed to infer_request -// - Future: can be changed to ov::RemoteTensor for GPU/NPU -// -// This design is similar to: -// - CUDA split buffer: tensor->extra stores device pointers -// - CPU repack buffer: tensor->extra stores tensor_traits with repacked data -// ===================================================== - -// Buffer context that manages per-tensor allocations (no contiguous buffer for weights) -struct ggml_backend_openvino_buffer_context { - int device; - std::string name; - size_t id; - - // For non-weight buffers (KV cache, compute), we still use contiguous allocation - void * data; - size_t size; - bool is_remote; - - // Wrapping of the buffer - std::shared_ptr ov_buffer; - - // Track all extras for cleanup - std::map tensor_extras; - - // Used for re-allocation on device for kvcache - void * data_prev; - - ggml_backend_openvino_buffer_context(int device, size_t size, bool is_remote = false) : - device(device), - name(std::string(GGML_OPENVINO_NAME) + std::to_string(device)), - id([]() { - static std::atomic next_id{1}; - return next_id.fetch_add(1); - }()), - data(nullptr), - size(size), - is_remote(is_remote) { - if (size == 0) { - return; - } - - const auto & device_name = ggml_openvino_get_device_name(); - - if (is_remote) { - GGML_ASSERT(device_name == "GPU"); - auto remote_context = ggml_openvino_get_remote_context(); - auto gpu_context = remote_context->as(); - ov::intel_gpu::ocl::USMTensor usm_tensor = - gpu_context.create_usm_device_tensor(ov::element::u8, ov::Shape{size}); - data = usm_tensor.get(); - ov_buffer = std::make_shared(std::move(usm_tensor)); - } else { - data = ggml_aligned_malloc(size); - GGML_ASSERT(data); - memset(data, 0, size); - ov_buffer = std::make_shared(ov::element::u8, ov::Shape{size}, data); - } - - if (data == nullptr) { - GGML_LOG_ERROR("%s: failed to allocate %zu bytes\n", __func__, size); - return; - } - - if (reinterpret_cast(data) % TENSOR_ALIGNMENT != 0) { - GGML_LOG_ERROR("%s: %s buffer is not aligned to %d bytes\n", __func__, device_name.c_str(), - TENSOR_ALIGNMENT); - GGML_ABORT("fatal error"); - } - } - - ~ggml_backend_openvino_buffer_context() { - // Clean up all tensor extras - // GGML_LOG_DEBUG("Deleting OpenVINO buffer context #%zu for device %d, size %zu MB\n", id, device, - // size / 1024 / 1024); - for (auto & pair : tensor_extras) { - delete pair.second; - } - tensor_extras.clear(); - if (!is_remote && data != nullptr) { - ggml_aligned_free(data, size); - } - } -}; - -// Buffer type context (per-device) -struct ggml_backend_openvino_buffer_type_context { - int device; - std::string name; -}; - -// Buffer interface functions -static void ggml_backend_openvino_buffer_free_buffer(ggml_backend_buffer_t buffer) { - ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; - delete ctx; -} - -static void * ggml_backend_openvino_buffer_get_base(ggml_backend_buffer_t buffer) { - ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; - return ctx->data; -} - -static bool is_stateful_enabled() { - static const auto * stateful = getenv("GGML_OPENVINO_STATEFUL_EXECUTION"); - return stateful && *stateful != '\0' && strcmp(stateful, "0") != 0; -} - -static enum ggml_status ggml_backend_openvino_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { - // GGML_LOG_DEBUG("%s: buffer usage=%d, tensor name=%s\n", __func__, buffer->usage, tensor->name); - ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; - - // Put kvcache on device memory for GPU (NPU memory is too small even for kvcache) - if (strncmp(tensor->name, "cache_", 6) == 0 && !ctx->is_remote && ggml_openvino_get_device_name() == "GPU" && - !is_stateful_enabled()) { - GGML_ASSERT(ctx->tensor_extras.empty()); - auto device = ctx->device; - auto size = ctx->size; - auto * data_prev = ctx->data; - delete ctx; - ctx = new ggml_backend_openvino_buffer_context(device, size, true); - buffer->context = ctx; - tensor->data = (char *) ctx->data + ((char *) tensor->data - (char *) data_prev); - } - - // Views share the extra from view_src - if (tensor->view_src != nullptr) { - GGML_ASSERT(tensor->view_src->buffer->buft == buffer->buft); - if (tensor->view_src->extra != nullptr) { - tensor->extra = tensor->view_src->extra; - } - return GGML_STATUS_SUCCESS; - } - - ctx = (ggml_backend_openvino_buffer_context *) buffer->context; - - if (tensor->data != nullptr && !ggml_is_quantized(tensor->type)) { - ggml_openvino_tensor_extra * extra = ggml_openvino_create_tensor_extra(tensor, ctx->is_remote); - if (extra != nullptr) { - auto it = ctx->tensor_extras.find(tensor); - if (it != ctx->tensor_extras.end()) { - delete it->second; - } - ctx->tensor_extras[tensor] = extra; - tensor->extra = extra; - } - } - - return GGML_STATUS_SUCCESS; -} - -static void ggml_backend_openvino_buffer_memset_tensor(ggml_backend_buffer_t buffer, - ggml_tensor * tensor, - uint8_t value, - size_t offset, - size_t size) { - // GGML_LOG_DEBUG("%s: buffer usage=%d, tensor name=%s\n", __func__, buffer->usage, tensor->name); - GGML_ASSERT(tensor != nullptr && tensor->data != nullptr); - ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; - - if (ctx->is_remote) { - // For remote (device) buffers, use OpenCL USM memfill - cl_command_queue queue = ggml_openvino_get_cl_queue(); - auto mem_fill_fn = ggml_openvino_get_clEnqueueMemFillINTEL(); - if (queue != nullptr && mem_fill_fn != nullptr) { - uint8_t pattern = value; - cl_int err = mem_fill_fn(queue, (char *) tensor->data + offset, &pattern, sizeof(pattern), size, 0, nullptr, - nullptr); - if (err != CL_SUCCESS) { - GGML_LOG_ERROR("%s: clEnqueueMemFillINTEL failed with error %d\n", __func__, err); - } - clFinish(queue); - } else { - GGML_LOG_ERROR("%s: no OpenCL queue or clEnqueueMemFillINTEL not available for GPU buffer\n", __func__); - } - } else { - memset((char *) tensor->data + offset, value, size); - } -} - -static void ggml_backend_openvino_buffer_set_tensor(ggml_backend_buffer_t buffer, - ggml_tensor * tensor, - const void * data, - size_t offset, - size_t size) { - // GGML_LOG_DEBUG("%s: buffer usage=%d, tensor name=%s\n", __func__, buffer->usage, tensor->name); - GGML_ASSERT(tensor != nullptr && tensor->data != nullptr); - ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; - - // Check if this is a weight buffer (usage is set BEFORE set_tensor is called, except in test-backend-ops) - bool is_weight_buffer = (buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - // Full tensor set: offset=0, full size, not a view - bool is_full_tensor_set = (offset == 0 && size == ggml_nbytes(tensor) && tensor->view_src == nullptr); - // 2D tensor (typical weight shape) - bool is_2d = (tensor->ne[2] == 1 && tensor->ne[3] == 1); - - if (is_weight_buffer && is_full_tensor_set && is_2d) { - try { - auto result = process_weight_tensor(tensor, data, tensor->data); - result.weight_node->set_friendly_name(tensor->name); - - // const auto & layout = result.layout; - ggml_openvino_extra_base * extra; - - // Quantized path with extracted weight/scale/zp tensors - if (result.is_quantized()) { - extra = new ggml_openvino_quantized_weight_extra(std::move(result.weights), std::move(result.scales), - std::move(result.zp), result.weight_node); - - // if (layout.is_requant) { - // GGML_LOG_DEBUG("%s: requantized %s to %s (u%d, block_size=%ld)\n", __func__, tensor->name, - // extra_quant_type_name(layout.requant_type.value()), layout.is_u4 ? 4 : 8, - // layout.weights_per_block); - // } else { - // int64_t n_blocks = ggml_nelements(tensor) / layout.weights_per_block; - // GGML_LOG_DEBUG("%s: extracted quantized weight node for %s (u%d, %zu weights, %ld blocks)\n", - // __func__, tensor->name, layout.is_u4 ? 4 : 8, layout.weights_size, n_blocks); - // } - } else { - // F16/F32/BF16 weight or F16-requant - extra = new ggml_openvino_weight_extra(std::move(result.weights), result.weight_node); - - // if (layout.total_size > 0) { - // GGML_LOG_DEBUG("%s: requantized %s to F16\n", __func__, tensor->name); - // } else { - // GGML_LOG_DEBUG("%s: created shared-memory weight node for %s\n", __func__, tensor->name); - // } - } - - ctx->tensor_extras[tensor] = extra; - tensor->extra = extra; - - } catch (const std::exception & e) { - GGML_LOG_ERROR("%s: failed to process weight tensor for %s: %s\n", __func__, tensor->name, e.what()); - memcpy((char *) tensor->data + offset, data, size); - } - } else { - // Non-weight tensor (KV cache, activations, etc.) - copy data. test-backend-ops also goes here - if (ctx->is_remote) { - cl_command_queue queue = ggml_openvino_get_cl_queue(); - auto mem_cpy_fn = ggml_openvino_get_clEnqueueMemcpyINTEL(); - if (queue != nullptr && mem_cpy_fn != nullptr) { - cl_int err = - mem_cpy_fn(queue, CL_TRUE, (char *) tensor->data + offset, data, size, 0, nullptr, nullptr); - if (err != CL_SUCCESS) { - GGML_LOG_ERROR("%s: clEnqueueMemcpyINTEL failed with error %d\n", __func__, err); - } - } else { - GGML_LOG_ERROR("%s: no OpenCL queue or clEnqueueMemcpyINTEL not available for GPU buffer\n", __func__); - } - } else { - memcpy((char *) tensor->data + offset, data, size); - } - - ggml_openvino_tensor_extra * extra = ggml_openvino_create_tensor_extra(tensor, ctx->is_remote); - if (extra == nullptr) { - // GGML_LOG_ERROR("%s: failed to create tensor extra for %s\n", __func__, tensor->name); - return; - } - - auto it = ctx->tensor_extras.find(tensor); - if (it != ctx->tensor_extras.end()) { - delete it->second; - } - ctx->tensor_extras[tensor] = extra; - tensor->extra = extra; - } -} - -static void ggml_backend_openvino_buffer_get_tensor(ggml_backend_buffer_t buffer, - const ggml_tensor * tensor, - void * data, - size_t offset, - size_t size) { - // GGML_LOG_DEBUG("%s: buffer usage=%d, tensor name=%s\n", __func__, buffer->usage, tensor->name); - GGML_ASSERT(tensor != nullptr && tensor->data != nullptr); - ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; - - if (ctx->is_remote) { - // For remote (device) buffers, use OpenCL USM memcpy (device-to-host) - cl_command_queue queue = ggml_openvino_get_cl_queue(); - auto mem_cpy_fn = ggml_openvino_get_clEnqueueMemcpyINTEL(); - if (queue != nullptr && mem_cpy_fn != nullptr) { - cl_int err = - mem_cpy_fn(queue, CL_TRUE, data, (const char *) tensor->data + offset, size, 0, nullptr, nullptr); - if (err != CL_SUCCESS) { - GGML_LOG_ERROR("%s: clEnqueueMemcpyINTEL failed with error %d\n", __func__, err); - } - } else { - GGML_LOG_ERROR("%s: no OpenCL queue or clEnqueueMemcpyINTEL not available for GPU buffer\n", __func__); - } - } else { - memcpy(data, (const char *) tensor->data + offset, size); - } -} - -static bool ggml_backend_openvino_buffer_cpy_tensor(ggml_backend_buffer_t buffer, - const ggml_tensor * src, - ggml_tensor * dst) { - // GGML_LOG_DEBUG("%s: src tensor name=%s, dst tensor name=%s\n", __func__, src->name, dst->name); - GGML_ASSERT(src != nullptr && dst != nullptr); - ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; - - if (ctx->is_remote) { - // For remote (device) buffers, use OpenCL USM memcpy - cl_command_queue queue = ggml_openvino_get_cl_queue(); - auto mem_cpy_fn = ggml_openvino_get_clEnqueueMemcpyINTEL(); - if (queue == nullptr || mem_cpy_fn == nullptr) { - GGML_LOG_ERROR("%s: no OpenCL queue or clEnqueueMemcpyINTEL not available for GPU buffer\n", __func__); - return false; - } - // Can copy from host to device - if (ggml_backend_buffer_is_host(src->buffer)) { - cl_int err = mem_cpy_fn(queue, CL_TRUE, dst->data, src->data, ggml_nbytes(src), 0, nullptr, nullptr); - if (err != CL_SUCCESS) { - GGML_LOG_ERROR("%s: clEnqueueMemcpyINTEL (host-to-device) failed with error %d\n", __func__, err); - return false; - } - return true; - } - // Can also copy from device to device if both are OpenVINO remote buffers - if (ggml_backend_buffer_is_openvino(src->buffer)) { - ggml_backend_openvino_buffer_context * src_ctx = - (ggml_backend_openvino_buffer_context *) src->buffer->context; - if (src_ctx->is_remote) { - cl_int err = - mem_cpy_fn(queue, CL_TRUE, dst->data, src->data, ggml_nbytes(src), 0, nullptr, nullptr); - if (err != CL_SUCCESS) { - GGML_LOG_ERROR("%s: clEnqueueMemcpyINTEL (device-to-device) failed with error %d\n", __func__, - err); - return false; - } - return true; - } - } - return false; - } - - // Host buffer - can copy from any host buffer - if (ggml_backend_buffer_is_host(src->buffer)) { - memcpy(dst->data, src->data, ggml_nbytes(src)); - return true; - } - return false; -} - -static void ggml_backend_openvino_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { - ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; - GGML_ASSERT(ctx->data != nullptr); - if (ctx->is_remote) { - cl_command_queue queue = ggml_openvino_get_cl_queue(); - auto mem_fill_fn = ggml_openvino_get_clEnqueueMemFillINTEL(); - if (queue != nullptr && mem_fill_fn != nullptr) { - uint8_t pattern = value; - cl_int err = mem_fill_fn(queue, ctx->data, &pattern, sizeof(pattern), ctx->size, 0, nullptr, nullptr); - if (err != CL_SUCCESS) { - GGML_LOG_WARN("%s: clEnqueueMemFillINTEL failed with error %d\n", __func__, err); - } - clFinish(queue); - } else { - GGML_LOG_WARN("%s: no OpenCL queue or clEnqueueMemFillINTEL not available for GPU buffer clear\n", - __func__); - } - } else { - memset(ctx->data, value, ctx->size); - } -} - -static const ggml_backend_buffer_i ggml_backend_openvino_buffer_interface = { - /* .free_buffer = */ ggml_backend_openvino_buffer_free_buffer, - /* .get_base = */ ggml_backend_openvino_buffer_get_base, - /* .init_tensor = */ ggml_backend_openvino_buffer_init_tensor, - /* .memset_tensor = */ ggml_backend_openvino_buffer_memset_tensor, - /* .set_tensor = */ ggml_backend_openvino_buffer_set_tensor, - /* .get_tensor = */ ggml_backend_openvino_buffer_get_tensor, - /* .set_tensor_2d = */ NULL, - /* .get_tensor_2d = */ NULL, - /* .cpy_tensor = */ ggml_backend_openvino_buffer_cpy_tensor, - /* .clear = */ ggml_backend_openvino_buffer_clear, - /* .reset = */ NULL, -}; - -// Buffer type interface functions -static const char * ggml_backend_openvino_buffer_type_get_name(ggml_backend_buffer_type_t buft) { - ggml_backend_openvino_buffer_type_context * ctx = (ggml_backend_openvino_buffer_type_context *) buft->context; - return ctx->name.c_str(); -} - -static ggml_backend_buffer_t ggml_backend_openvino_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, - size_t size) { - ggml_backend_openvino_buffer_type_context * buft_ctx = (ggml_backend_openvino_buffer_type_context *) buft->context; - - // Create buffer context with contiguous memory allocation - ggml_backend_openvino_buffer_context * ctx = new ggml_backend_openvino_buffer_context(buft_ctx->device, size); - - if (ctx->data == nullptr && size > 0) { - GGML_LOG_ERROR("%s: failed to allocate buffer of size %zu\n", __func__, size); - delete ctx; - return nullptr; - } - - return ggml_backend_buffer_init(buft, ggml_backend_openvino_buffer_interface, ctx, size); -} - -static size_t ggml_backend_openvino_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { - GGML_UNUSED(buft); - return TENSOR_ALIGNMENT; -} - -static size_t ggml_backend_openvino_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { - GGML_UNUSED(buft); - return SIZE_MAX; -} - -static size_t ggml_backend_openvino_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, - const ggml_tensor * tensor) { - GGML_UNUSED(buft); - - // For quantized 2D tensors (weights), we need extra space for extracted data - if (ggml_is_quantized(tensor->type) && tensor->ne[2] == 1 && tensor->ne[3] == 1) { - ggml_openvino_extracted_layout layout = ggml_openvino_get_extracted_layout(tensor); - if (layout.total_size > 0) { - // GGML_LOG_DEBUG("%s: tensor %s needs %zu bytes (original %zu, extracted: weights=%zu scales=%zu zp=%zu)\n", - // __func__, tensor->name, layout.total_size, ggml_nbytes(tensor), layout.weights_size, - // layout.scales_size, layout.zp_size); - return layout.total_size; - } - } - - return ggml_nbytes(tensor); -} - -static const ggml_backend_buffer_type_i ggml_backend_openvino_buffer_type_interface = { - /* .get_name = */ ggml_backend_openvino_buffer_type_get_name, - /* .alloc_buffer = */ ggml_backend_openvino_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_openvino_buffer_type_get_alignment, - /* .get_max_size = */ ggml_backend_openvino_buffer_type_get_max_size, - /* .get_alloc_size = */ ggml_backend_openvino_buffer_type_get_alloc_size, - /* .is_host = */ nullptr, -}; - -// Get buffer type for a specific device -GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_openvino_buffer_type(int device) { - GGML_ASSERT(device >= 0 && device < ggml_backend_openvino_get_device_count()); - - static std::mutex mutex; - std::lock_guard lock(mutex); - - static std::vector buffer_types; - static std::vector buffer_type_contexts; - - if (buffer_types.empty()) { - int device_count = ggml_backend_openvino_get_device_count(); - buffer_types.resize(device_count); - buffer_type_contexts.resize(device_count); - - for (int i = 0; i < device_count; i++) { - buffer_type_contexts[i].device = i; - buffer_type_contexts[i].name = std::string(GGML_OPENVINO_NAME) + std::to_string(i); - - buffer_types[i] = ggml_backend_buffer_type{ - /* .iface = */ ggml_backend_openvino_buffer_type_interface, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_openvino_reg(), i), - /* .context = */ &buffer_type_contexts[i], - }; - } - } - - return &buffer_types[device]; -} - -// ===================================================== -// OpenVINO Host Buffer Implementation -// ===================================================== - -static const char * ggml_backend_openvino_host_buffer_type_get_name(ggml_backend_buffer_type_t buft) { - ggml_backend_openvino_buffer_type_context * ctx = (ggml_backend_openvino_buffer_type_context *) buft->context; - static std::string name; - name = ctx->name + "_HOST"; - return name.c_str(); -} - -static bool ggml_backend_openvino_host_buffer_type_is_host(ggml_backend_buffer_type_t buft) { - GGML_UNUSED(buft); - return true; -} - -static const ggml_backend_buffer_type_i ggml_backend_openvino_host_buffer_type_interface = { - /* .get_name = */ ggml_backend_openvino_host_buffer_type_get_name, - /* .alloc_buffer = */ ggml_backend_openvino_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_openvino_buffer_type_get_alignment, - /* .get_max_size = */ ggml_backend_openvino_buffer_type_get_max_size, - /* .get_alloc_size = */ ggml_backend_openvino_buffer_type_get_alloc_size, - /* .is_host = */ ggml_backend_openvino_host_buffer_type_is_host, -}; - -GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_openvino_host_buffer_type(int device) { - GGML_ASSERT(device >= 0 && device < ggml_backend_openvino_get_device_count()); - - static std::mutex mutex; - std::lock_guard lock(mutex); - - static std::vector buffer_types; - static std::vector buffer_type_contexts; - - if (buffer_types.empty()) { - int device_count = ggml_backend_openvino_get_device_count(); - buffer_types.resize(device_count); - buffer_type_contexts.resize(device_count); - - for (int i = 0; i < device_count; i++) { - buffer_type_contexts[i].device = i; - buffer_type_contexts[i].name = std::string(GGML_OPENVINO_NAME) + std::to_string(i); - - buffer_types[i] = ggml_backend_buffer_type{ - /* .iface = */ ggml_backend_openvino_host_buffer_type_interface, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_openvino_reg(), i), - /* .context = */ &buffer_type_contexts[i], - }; - } - } - - return &buffer_types[device]; -} - -bool ggml_backend_buffer_is_openvino(ggml_backend_buffer_t buffer) { - return buffer->iface.free_buffer == ggml_backend_openvino_buffer_free_buffer; -} - -size_t ggml_backend_openvino_buffer_get_ctx_id(ggml_backend_buffer_t buffer) { - if (!ggml_backend_buffer_is_openvino(buffer)) { - return 0; - } - ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; - return ctx->id; -} - -void ggml_openvino_buffer_register_extra(ggml_tensor * tensor, ggml_openvino_extra_base * extra) { - GGML_ASSERT(tensor != nullptr); - GGML_ASSERT(tensor->buffer != nullptr); - GGML_ASSERT(ggml_backend_buffer_is_openvino(tensor->buffer)); - - auto * ctx = static_cast(tensor->buffer->context); - - auto it = ctx->tensor_extras.find(tensor); - if (it != ctx->tensor_extras.end()) { - delete it->second; - } - - ctx->tensor_extras[tensor] = extra; - tensor->extra = extra; -} - -bool ggml_backend_buft_is_openvino(ggml_backend_buffer_type_t buft) { - return buft->iface.get_name == ggml_backend_openvino_buffer_type_get_name; -} - -bool ggml_backend_buft_is_openvino_host(ggml_backend_buffer_type_t buft) { - return buft->iface.get_name == ggml_backend_openvino_host_buffer_type_get_name; -} - -static void ggml_backend_openvino_free(ggml_backend_t backend) { - ggml_backend_openvino_context * ctx = (ggml_backend_openvino_context *) backend->context; - - if (ctx->runtime_context) { - auto r_ctx = std::static_pointer_cast(ctx->runtime_context); - if (--r_ctx->backend_count == 0) { - r_ctx->clear_caches(); - } - } - - delete ctx; - delete backend; -} - -static const char * ggml_backend_openvino_get_name(ggml_backend_t backend) { - return GGML_OPENVINO_NAME; - GGML_UNUSED(backend); -} - -static enum ggml_status ggml_backend_openvino_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { - return ov_graph_compute(cgraph, backend); - GGML_UNUSED(backend); -} - -static const ggml_backend_i ggml_backend_openvino_interface = { - /* .get_name = */ ggml_backend_openvino_get_name, - /* .free = */ ggml_backend_openvino_free, - /* .set_tensor_async = */ NULL, - /* .get_tensor_async = */ NULL, - /* .set_tensor_2d_async = */ NULL, - /* .get_tensor_2d_async = */ NULL, - /* .cpy_tensor_async = */ NULL, - /* .synchronize = */ NULL, - /* .graph_plan_create = */ NULL, - /* .graph_plan_free = */ NULL, - /* .graph_plan_update = */ NULL, - /* .graph_plan_compute = */ NULL, - /* .graph_compute = */ ggml_backend_openvino_graph_compute, - /* .event_record = */ NULL, - /* .event_wait = */ NULL, - /* .graph_optimize = */ NULL, -}; - -int ggml_backend_openvino_get_device_count() { - return 1; -} - -static ggml_guid_t ggml_backend_openvino_guid(void) { - static ggml_guid guid = {0x12, 0xa8, 0xae, 0xf4, 0xc0, 0x1e, 0x61, 0x97, - 0x8f, 0xeb, 0x33, 0x04, 0xa1, 0x33, 0x51, 0x2d}; - return &guid; -} - -static std::shared_ptr get_ov_runtime_context_ptr() { - static std::shared_ptr r_ctx = [] { - auto ctx = std::make_shared(); - ctx->device = ggml_openvino_get_device_name(); - ctx->stateful = is_stateful_enabled() && !ggml_openvino_is_npu(); - return ctx; - }(); - return r_ctx; -} - -// backend API -GGML_BACKEND_API ggml_backend_t ggml_backend_openvino_init(int device) { - if (device < 0 || device >= ggml_backend_openvino_get_device_count()) { - GGML_LOG_ERROR("%s: invalid device %d\n", __func__, device); - return nullptr; - } - - ggml_backend_openvino_context * ctx = new ggml_backend_openvino_context; - if (ctx == nullptr) { - GGML_LOG_ERROR("%s: failed to allocate context\n", __func__); - return nullptr; - } - - ctx->runtime_context = get_ov_runtime_context_ptr(); - if (ctx->runtime_context == nullptr) { - GGML_LOG_ERROR("%s: failed to allocate runtime context\n", __func__); - delete ctx; - return nullptr; - } - - std::shared_ptr r_ctx = std::static_pointer_cast(ctx->runtime_context); - r_ctx->backend_count++; - - ggml_backend_t openvino_backend = new ggml_backend{ - /* .guid = */ ggml_backend_openvino_guid(), - /* .interface = */ ggml_backend_openvino_interface, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_openvino_reg(), device), - /* .context = */ ctx, - }; - - return openvino_backend; -} - -GGML_BACKEND_API bool ggml_backend_is_openvino(ggml_backend_t backend) { - return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_openvino_guid()); -} - -struct ggml_backend_openvino_device_context { - int device; - std::string name; - std::string description; -}; - -static const char * ggml_backend_openvino_device_get_name(ggml_backend_dev_t dev) { - ggml_backend_openvino_device_context * ctx = (ggml_backend_openvino_device_context *) dev->context; - return ctx->name.c_str(); -} - -static const char * ggml_backend_openvino_device_get_description(ggml_backend_dev_t dev) { - ggml_backend_openvino_device_context * ctx = (ggml_backend_openvino_device_context *) dev->context; - return ctx->description.c_str(); -} - -static void ggml_backend_openvino_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { -#ifdef _WIN32 - MEMORYSTATUSEX status; - status.dwLength = sizeof(status); - GlobalMemoryStatusEx(&status); - *total = status.ullTotalPhys; - *free = status.ullAvailPhys; -#else - long pages = sysconf(_SC_PHYS_PAGES); - long page_size = sysconf(_SC_PAGE_SIZE); - *total = pages * page_size; - - // "free" system memory is ill-defined, for practical purposes assume that all of it is free: - *free = *total; -#endif // _WIN32 - - GGML_UNUSED(dev); -} - -static enum ggml_backend_dev_type ggml_backend_openvino_device_get_type(ggml_backend_dev_t dev) { - GGML_UNUSED(dev); - return GGML_BACKEND_DEVICE_TYPE_GPU; -} - -static void ggml_backend_openvino_device_get_props(ggml_backend_dev_t dev, ggml_backend_dev_props * props) { - props->name = ggml_backend_openvino_device_get_name(dev); - props->description = ggml_backend_openvino_device_get_description(dev); - props->type = ggml_backend_openvino_device_get_type(dev); - ggml_backend_openvino_device_get_memory(dev, &props->memory_free, &props->memory_total); - - props->caps = { - /* .async = */ false, - /* .host_buffer = */ false, - /* .buffer_from_host_ptr = */ false, - /* .events = */ false, - }; -} - -static ggml_backend_t ggml_backend_openvino_device_init(ggml_backend_dev_t dev, const char * params) { - GGML_UNUSED(params); - ggml_backend_openvino_device_context * ctx = (ggml_backend_openvino_device_context *) dev->context; - return ggml_backend_openvino_init(ctx->device); -} - -static ggml_backend_buffer_type_t ggml_backend_openvino_device_get_buffer_type(ggml_backend_dev_t dev) { - ggml_backend_openvino_device_context * ctx = (ggml_backend_openvino_device_context *) dev->context; - return ggml_backend_openvino_buffer_type(ctx->device); -} - -static ggml_backend_buffer_type_t ggml_backend_openvino_device_get_host_buffer_type(ggml_backend_dev_t dev) { - ggml_backend_openvino_device_context * ctx = (ggml_backend_openvino_device_context *) dev->context; - return ggml_backend_openvino_host_buffer_type(ctx->device); -} - -static bool has_view_op_input(const ggml_tensor * op) { - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (op->src[i] == nullptr) { - break; - } - if (op->src[i]->op == GGML_OP_VIEW) { - return true; - } - } - return false; -} - -static bool is_supported_flash_attn_pattern(const ggml_tensor * op) { - // pattern of q,k,v should be q->op==PERMUTE, q->src[0]->op==VIEW, q->src[0]->src[0]->view_src==nullptr - for (int i = 0; i < 3; i++) { - const ggml_tensor * src = op->src[i]; - if (src->op != GGML_OP_PERMUTE || src->src[0] == nullptr || src->src[0]->op != GGML_OP_VIEW || - src->src[0]->src[0] == nullptr || src->src[0]->src[0]->view_src != nullptr) { - return false; - } - } - return true; -} - -static bool is_op_unsupported_case(const ggml_tensor * op) { - switch (op->op) { - case GGML_OP_GET_ROWS: - case GGML_OP_SET_ROWS: { - if (op->ne[3] != 1) { - return true; - } - break; - } - case GGML_OP_ADD: - case GGML_OP_MUL: { - if (op->src[1]->op == GGML_OP_PERMUTE) { - return true; - } - for (int i = 0; i < 4; i++) { - if (op->src[0]->ne[i] != op->src[1]->ne[i] && (op->src[0]->ne[i] != 1 && op->src[1]->ne[i] != 1)) { - return true; - } - } - break; - } - case GGML_OP_SOFT_MAX: { - if (op->src[2] != nullptr) { - // GGML_LOG_WARN("OpenVINO backend does not support SOFT_MAX with sinks\n"); - return true; - } - float scale = 1.0f; - float max_bias = 0.0f; - const auto * op_params = op->op_params; - memcpy(&scale, (const float *) op_params + 0, sizeof(float)); - memcpy(&max_bias, (const float *) op_params + 1, sizeof(float)); - if (max_bias > 0) { - // GGML_LOG_WARN("OpenVINO backend does not support SOFT_MAX with max_bias > 0\n"); - return true; - } - break; - } - case GGML_OP_FLASH_ATTN_EXT: { - if (op->src[4] != nullptr) { - // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with sinks\n"); - return true; - } - if (!is_supported_flash_attn_pattern(op)) { - return true; - } - float scale = 1.0f; - float max_bias = 0.0f; - float logit_softcap = 0.0f; - const auto * op_params = op->op_params; - memcpy(&scale, (const float *) op_params + 0, sizeof(float)); - memcpy(&max_bias, (const float *) op_params + 1, sizeof(float)); - memcpy(&logit_softcap, (const float *) op_params + 2, sizeof(float)); - if (max_bias > 0) { - // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with max_bias > 0\n"); - return true; - } - if (logit_softcap != 0) { - // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with logit_softcap != 0\n"); - return true; - } - break; - } - case GGML_OP_PERMUTE: { - if (op->type == GGML_TYPE_BF16) { - // err msg: [GPU] Could not find a suitable kernel for transpose - // GGML_LOG_WARN("OpenVINO backend does not support PERMUTE with BF16 type\n"); - return true; - } - break; - } - case GGML_OP_CPY: { - if (op->src[1] != op) { - // GGML_LOG_WARN("OpenVINO backend only supports CPY that is a cast\n"); - return true; - } - break; - } - case GGML_OP_MUL_MAT: { - if (op->src[0]->type == GGML_TYPE_F16 && op->src[1]->type == GGML_TYPE_F16) { - // Has accuracy issue, try enabling this and see `test-backend-ops -o "MUL_MAT"` - // GGML_LOG_WARN("OpenVINO backend does not support MUL_MAT with two F16 tensors\n"); - return true; - } - if (op->src[0]->ne[3] != op->src[1]->ne[3] && op->src[0]->ne[3] != 1 && op->src[1]->ne[3] != 1) { - return true; - } - if (op->src[0]->op == GGML_OP_PERMUTE || op->src[1]->op == GGML_OP_PERMUTE) { - return true; - } - if (ggml_is_quantized(op->src[0]->type) && op->src[0]->ne[1] == 1) { - // MUL_MAT(type_a=q4_0,type_b=f32,m=1,n=2048,k=8192,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1) - // triggers a bug in ov matmul_shape_inference.hpp - return true; - } - if (op->src[0]->op == GGML_OP_VIEW && op->src[1]->op == GGML_OP_VIEW) { - return true; - } - break; - } - case GGML_OP_ROPE: { - const int32_t * op_params = op->op_params; - const int n_dims = op_params[1]; - const int mode = op_params[2]; - if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX && mode != GGML_ROPE_TYPE_IMROPE) { - // GGML_LOG_WARN("OpenVINO backend does not support ROPE with mode %d\n", mode); - return true; - } - if (n_dims != 0.0f && n_dims != op->src[0]->ne[0]) { - // GGML_LOG_WARN("OpenVINO backend does not support ROPE with n_dims %d != src[0]->ne[0] %ld\n", n_dims, - // op->src[0]->ne[0]); - return true; - } - if (op->type != GGML_TYPE_F32) { - // GGML_LOG_WARN("OpenVINO backend does not support ROPE with type %s\n", ggml_type_name(op->type)); - return true; - } - if (op->src[0]->op == GGML_OP_VIEW) { - if (op->src[0]->view_src->ne[1] != op->src[0]->ne[2]) { - // GGML_LOG_WARN( - // "OpenVINO backend does not support ROPE with src[0]->view_src->ne[1] %ld != src[0]->ne[2] " - // "%ld\n", - // op->src[0]->view_src->ne[1], op->src[0]->ne[2]); - return true; - } - } - if (mode == GGML_ROPE_TYPE_IMROPE && - (op->src[2] != 0 || ((const float *) op_params)[6] != 1 || ((const float *) op_params)[7] != 0 || - ((const float *) op_params)[8] != 1)) { - // GGML_LOG_WARN("OpenVINO backend does not support IMROPE with freq_factors, freq_scale, ext_factor, and attn_factor\n"); - return true; - } - break; - } - default: - break; - } - if (op->op == GGML_OP_GET_ROWS) { - if (op->ne[0] == 256 && (op->src[0]->type == GGML_TYPE_Q4_K || op->src[0]->type == GGML_TYPE_Q5_K)) { - // ERR = 0.000000306 > 0.000000100 GET_ROWS(type=q4_K,n=256,m=5,r=4,be1=1,be2=1,v=0) - // ERR = 0.000000197 > 0.000000100 GET_ROWS(type=q5_K,n=256,m=5,r=4,be1=1,be2=1,v=0) - return true; - } - } - return false; -} - -static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { - GGML_ASSERT(dev->reg != nullptr); - - static std::set supported_types{GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_I64, - GGML_TYPE_I32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_K, - GGML_TYPE_Q5_K, GGML_TYPE_Q8_0, GGML_TYPE_Q6_K}; - - static const std::set supported_ops{GGML_OP_NONE, GGML_OP_ADD, GGML_OP_MUL, GGML_OP_MUL_MAT, GGML_OP_VIEW, - /*GGML_OP_CONT,*/ GGML_OP_RESHAPE, GGML_OP_PERMUTE, GGML_OP_TRANSPOSE, - GGML_OP_GET_ROWS, GGML_OP_ROPE, GGML_OP_RMS_NORM, GGML_OP_SCALE, - // softmax is not updated due to replaced by flash_attn_ext - // GGML_OP_SOFT_MAX, - GGML_OP_SET_ROWS, GGML_OP_FLASH_ATTN_EXT, GGML_OP_CPY}; - static const std::set supported_unary_ops{ - GGML_UNARY_OP_GELU, - GGML_UNARY_OP_SILU, - }; - static const std::set supported_glu_ops{ - GGML_GLU_OP_SWIGLU, - GGML_GLU_OP_GEGLU, - }; - - switch (op->op) { - case GGML_OP_UNARY: { - auto supported = supported_unary_ops.find(ggml_get_unary_op(op)) != supported_unary_ops.end(); - if (!supported) { - // GGML_LOG_WARN("OpenVINO backend does not support unary op %s\n", ggml_unary_op_name(ggml_get_unary_op(op))); - return false; - } - if (has_view_op_input(op)) { - // GGML_LOG_WARN("OpenVINO backend does not support unary op %s with view input\n", - // ggml_unary_op_name(ggml_get_unary_op(op))); - return false; - } - break; - } - case GGML_OP_GLU: { - auto supported = supported_glu_ops.find(ggml_get_glu_op(op)) != supported_glu_ops.end(); - if (!supported) { - // GGML_LOG_WARN("OpenVINO backend does not support GLU op %s\n", ggml_glu_op_name(ggml_get_glu_op(op))); - return false; - } - if (has_view_op_input(op)) { - // GGML_LOG_WARN("OpenVINO backend does not support unary op %s with view input\n", - // ggml_glu_op_name(ggml_get_glu_op(op))); - return false; - } - if (op->src[1] == nullptr && op->src[0]->ne[0] % 2 != 0) { - // triggers bug in ov gpu - return false; - } - break; - } - default: { - auto supported = supported_ops.find(op->op) != supported_ops.end(); - if (!supported) { - // GGML_LOG_WARN("OpenVINO backend does not support op %s\n", ggml_op_name(op->op)); - return false; - } - static std::set ops_not_support_view_input{ - GGML_OP_GET_ROWS, - GGML_OP_RMS_NORM, - }; - if (ops_not_support_view_input.find(op->op) != ops_not_support_view_input.end() && has_view_op_input(op)) { - // GGML_LOG_WARN("OpenVINO backend does not support op %s with view input\n", ggml_op_name(op->op)); - return false; - } - } - } - - if (supported_types.find(op->type) == supported_types.end()) { - // GGML_LOG_WARN("OpenVINO backend does not support tensor type %s\n", ggml_type_name(op->type)); - return false; - } - for (int i = 0; i < GGML_MAX_SRC; i++) { - auto * src = op->src[i]; - if (src == nullptr) { - break; - } - if (supported_types.find(src->type) == supported_types.end()) { - // GGML_LOG_WARN("OpenVINO backend does not support tensor type %s\n", ggml_type_name(src->type)); - return false; - } - if (ggml_is_quantized(src->type) && src->ne[2] != 1) { - // GGML_LOG_WARN("OpenVINO backend does not support 3D quantized tensors\n"); - return false; - } - } - - if (is_op_unsupported_case(op)) { - return false; - } - return true; -} - -static bool ggml_backend_openvino_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { - return ggml_backend_buft_is_openvino(buft) || ggml_backend_buft_is_host(buft); - GGML_UNUSED(dev); -} - -static const struct ggml_backend_device_i ggml_backend_openvino_device_interface = { - /* .get_name = */ ggml_backend_openvino_device_get_name, - /* .get_description = */ ggml_backend_openvino_device_get_description, - /* .get_memory = */ ggml_backend_openvino_device_get_memory, - /* .get_type = */ ggml_backend_openvino_device_get_type, - /* .get_props = */ ggml_backend_openvino_device_get_props, - /* .init_backend = */ ggml_backend_openvino_device_init, - /* .get_buffer_type = */ ggml_backend_openvino_device_get_buffer_type, - /* .get_host_buffer_type = */ ggml_backend_openvino_device_get_host_buffer_type, - /* .buffer_from_host_ptr = */ NULL, - /* .supports_op = */ ggml_backend_openvino_device_supports_op, - /* .supports_buft = */ ggml_backend_openvino_device_supports_buft, - /* .offload_op = */ NULL, - /* .event_new = */ NULL, - /* .event_free = */ NULL, - /* .event_synchronize = */ NULL, -}; - -struct ggml_backend_openvino_reg_context { - std::vector devices; -}; - -static const char * ggml_backend_openvino_reg_get_name(ggml_backend_reg_t reg) { - return GGML_OPENVINO_NAME; - GGML_UNUSED(reg); -} - -static size_t ggml_backend_openvino_reg_get_device_count(ggml_backend_reg_t reg) { - GGML_UNUSED(reg); - return (size_t) ggml_backend_openvino_get_device_count(); -} - -static ggml_backend_dev_t ggml_backend_openvino_reg_get_device(ggml_backend_reg_t reg, size_t index) { - ggml_backend_openvino_reg_context * ctx = (ggml_backend_openvino_reg_context *) reg->context; - GGML_ASSERT(index < ctx->devices.size()); - return ctx->devices[index]; -} - -static const struct ggml_backend_reg_i ggml_backend_openvino_reg_interface = { - /* .get_name = */ ggml_backend_openvino_reg_get_name, - /* .get_device_count = */ ggml_backend_openvino_reg_get_device_count, - /* .get_device = */ ggml_backend_openvino_reg_get_device, - /* .get_proc_address = */ NULL, -}; - -static void ggml_openvino_init() { - // Initialize device config singleton from env var - ggml_openvino_init_device_config(); - GGML_LOG_INFO("OpenVINO: using device %s\n", ggml_openvino_get_device_name().c_str()); -} - -GGML_BACKEND_API ggml_backend_reg_t ggml_backend_openvino_reg(void) { - static ggml_backend_reg reg; - - static bool initialized = false; - { - static std::mutex mutex; - std::lock_guard lock(mutex); - if (!initialized) { - ggml_openvino_init(); - - ggml_backend_openvino_reg_context * ctx = new ggml_backend_openvino_reg_context; - - for (int i = 0; i < ggml_backend_openvino_get_device_count(); i++) { - ggml_backend_openvino_device_context * dev_ctx = new ggml_backend_openvino_device_context; - dev_ctx->device = i; - dev_ctx->name = GGML_OPENVINO_NAME + std::to_string(i); - - dev_ctx->description = ov::get_openvino_version().description; - - ggml_backend_dev_t dev = - new ggml_backend_device{/* .interface = */ ggml_backend_openvino_device_interface, - /* .reg = */ ®, - /* .context = */ dev_ctx}; - ctx->devices.push_back(dev); - } - - reg = ggml_backend_reg{/* .api_version = */ GGML_BACKEND_API_VERSION, - /* .iface = */ ggml_backend_openvino_reg_interface, - /* .context = */ ctx}; - } - - initialized = true; - } - - return ® -} diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp.inc b/ggml/src/ggml-openvino/ggml-openvino.cpp.inc new file mode 100644 index 000000000000..e61d22f60700 --- /dev/null +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp.inc @@ -0,0 +1,12186 @@ +/* Inlined openvino/utils.h */ +#include +#include +#include +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +std::string getCurrentTime(); + +void dump_ov_model(std::shared_ptr model); + +void num_inputs_check(const NodeContext& context, size_t min_inputs, size_t max_inputs); + +int non_cont_dim(std::vector ne, std::vector nb); + +template +std::vector argsort_descend(const std::vector& v) { + std::vector idx(v.size()); + std::iota(idx.begin(), idx.end(), 0); + std::sort(idx.begin(), idx.end(), [&v](int i1, int i2) { + return v[i1] > v[i2]; + }); + return idx; +} + +template +std::vector sorted_descend(std::vector v) { + std::sort(v.begin(), v.end(), [](T a, T b) { + return a > b; + }); + return v; +} + +template +bool is_permuted(const std::vector& strides) { + for (size_t i = 0; i < strides.size() - 1; ++i) { + if (strides[i] < strides[i + 1]) { + return true; + } + } + return false; +} + +template +std::vector permute(const std::vector& x, const std::vector& perm) { + std::vector result; + result.reserve(perm.size()); + for (int i : perm) { + result.push_back(x[i]); + } + return result; +} + +std::shared_ptr get_dimensions(const std::shared_ptr& shape, + const std::vector& dims); +std::shared_ptr get_dimensions(const std::shared_ptr& node, const std::vector& dims); + +OutputVector rename_outputs_with_suffix(const OutputVector& outputs, const std::string& suffix); + +std::pair, ov::Output> make_sin_cos(int32_t* rope_params, + std::shared_ptr inp_pos, + std::shared_ptr rope_freqs_weight = nullptr, + bool imrope = false, + bool stateful = false); + +ov::Output process_view_input(const NodeContext& context, int input_index, int slice_len = 0); + +namespace op { +template +OutputVector translate_1to1_match_2_inputs(const NodeContext& context) { + num_inputs_check(context, 2, 2); + auto res = std::make_shared(context.get_input(0), context.get_input(1)); + return rename_outputs_with_suffix({res}, context.get_name()); +} +} // namespace op + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +#include "ggml-openvino.h.inc" + +#include "ggml-backend-impl-defs.inc" +#include "ggml-backend.h.inc" +#include "ggml-impl-defs.inc" +/* Inlined ggml-openvino-extra.h */ +#include "ggml.h.inc" +#include "openvino/runtime/core.hpp" + +#define CL_TARGET_OPENCL_VERSION 300 +#include + +#include +#include +#include +#include +#include +#include +#include + +// ExtraQuantType enum - defines requantization target formats +enum class ExtraQuantType { F16, Q4_0_C, Q8_1_C, Q4_0_128, Q8_0_C, Q8_0_32 }; + +ov::Core & ov_singleton_core(); + +// Get the remote context for the current device (returns empty optional for CPU) +std::optional ggml_openvino_get_remote_context(); + +// Get the compile config for the current device +const ov::AnyMap & ggml_openvino_get_compile_config(); + +// Get the OpenCL command queue for GPU operations (returns nullptr for CPU/NPU) +cl_command_queue ggml_openvino_get_cl_queue(); + +// Intel USM extension function type +typedef cl_int(CL_API_CALL * clEnqueueMemFillINTEL_fn)(cl_command_queue queue, + void * dst_ptr, + const void * pattern, + size_t pattern_size, + size_t size, + cl_uint num_events_in_wait_list, + const cl_event * event_wait_list, + cl_event * event); + +typedef cl_int(CL_API_CALL * clEnqueueMemcpyINTEL_fn)(cl_command_queue queue, + cl_bool blocking, + void * dst_ptr, + const void * src_ptr, + size_t size, + cl_uint num_events_in_wait_list, + const cl_event * event_wait_list, + cl_event * event); + +// Get the clEnqueueMemFillINTEL function pointer (returns nullptr if not available) +clEnqueueMemFillINTEL_fn ggml_openvino_get_clEnqueueMemFillINTEL(); + +// Get the clEnqueueMemcpyINTEL function pointer (returns nullptr if not available) +clEnqueueMemcpyINTEL_fn ggml_openvino_get_clEnqueueMemcpyINTEL(); + +// ===================================================== +// Global Device Configuration (singleton) +// ===================================================== +// Initialized once during backend init from GGML_OPENVINO_DEVICE env var + +struct ggml_openvino_device_config { + std::string device_name = "CPU"; + bool is_npu = false; + bool initialized = false; + std::optional remote_context; + ov::AnyMap compile_config; + cl_command_queue cl_queue = nullptr; + + void init(); + ~ggml_openvino_device_config(); +}; + +// Get the global device config singleton +ggml_openvino_device_config & ggml_openvino_get_device_config(); + +// Initialize device config (call during backend init) +void ggml_openvino_init_device_config(); + +// Get the device name +const std::string & ggml_openvino_get_device_name(); + +// Check if running on NPU +bool ggml_openvino_is_npu(); + +// Get requantization type for a tensor type (returns nullopt if no requant needed) +std::optional ggml_openvino_get_requant_type(const ggml_tensor * tensor, bool no_requant = false); + +// ===================================================== +// OpenVINO Tensor Extra Types +// ===================================================== +// These types are stored in tensor->extra by the OpenVINO backend buffer. +// They allow: +// 1. Pre-built ov::Constant nodes for weights (avoiding memcpy during graph construction) +// 2. ov::Tensor wrappers for KV cache / compute tensors (for direct use with infer_request) + +// Base class for OpenVINO tensor extra data +struct ggml_openvino_extra_base { + enum class Type { WEIGHT, QUANTIZED_WEIGHT, TENSOR }; + Type type; + virtual ~ggml_openvino_extra_base() = default; +protected: + explicit ggml_openvino_extra_base(Type t) : type(t) {} +}; + +// Extra data for F16/F32/BF16 weight tensors - stores the pre-built weight node +struct ggml_openvino_weight_extra : public ggml_openvino_extra_base { + ov::Tensor weights; // The underlying weight data tensor + std::shared_ptr weight_node; // Pre-built OpenVINO weight node + + ggml_openvino_weight_extra(ov::Tensor w, std::shared_ptr n) : + ggml_openvino_extra_base(Type::WEIGHT), + weights(std::move(w)), + weight_node(std::move(n)) {} +}; + +// Extra data for quantized weight tensors - stores extracted weights/scales/zp and weight node +struct ggml_openvino_quantized_weight_extra : public ggml_openvino_extra_base { + ov::Tensor weights; // U4 or U8 extracted weights + ov::Tensor scales; // F16 scales + ov::Tensor zp; // U4 or U8 zero points (same type as weights) + std::shared_ptr weight_node; // Pre-built OpenVINO weight subgraph + + ggml_openvino_quantized_weight_extra(ov::Tensor w, ov::Tensor s, ov::Tensor z, std::shared_ptr n) : + ggml_openvino_extra_base(Type::QUANTIZED_WEIGHT), + weights(std::move(w)), + scales(std::move(s)), + zp(std::move(z)), + weight_node(std::move(n)) {} +}; + +// Extra data for KV cache / compute tensors - stores ov::Tensor for infer_request +struct ggml_openvino_tensor_extra : public ggml_openvino_extra_base { + std::shared_ptr tensor; // For direct use with infer_request + + explicit ggml_openvino_tensor_extra(std::shared_ptr t) + : ggml_openvino_extra_base(Type::TENSOR), tensor(std::move(t)) {} +}; + +// ===================================================== +// Extracted Size Calculation for Quantized Tensors +// ===================================================== +// For quantized tensors, we need extra space to store extracted weights, scales, and zero points. +// Returns the total size needed in the buffer for extracted data. + +struct ggml_openvino_extracted_layout { + size_t total_size = 0; // Total bytes needed + size_t weights_offset = 0; // Offset to weights in buffer + size_t weights_size = 0; // Size of weights in bytes + size_t scales_offset = 0; // Offset to scales in buffer + size_t scales_size = 0; // Size of scales in bytes + size_t zp_offset = 0; // Offset to zero points in buffer + size_t zp_size = 0; // Size of zero points in bytes (U4 or U8) + bool is_u4; // true for U4 weights, false for U8 + int64_t weights_per_block; // weights per scale/zp block + bool is_symmetric; // true for symmetric quantization + + // Requantization info + bool is_requant = false; // true if this tensor needs requantization + std::optional requant_type; // target requant type if is_requant +}; + +// Calculate the buffer layout for extracted quantized data +ggml_openvino_extracted_layout ggml_openvino_get_extracted_layout(const ggml_tensor * tensor, bool use_bias = false); + +ggml_openvino_tensor_extra * ggml_openvino_create_tensor_extra(const ggml_tensor * tensor, bool is_remote); + +// Register an extra with the tensor's OpenVINO buffer context for proper lifetime management. +// This sets tensor->extra and tracks the extra in the buffer context for cleanup. +void ggml_openvino_buffer_register_extra(ggml_tensor * tensor, ggml_openvino_extra_base * extra); + +// ===================================================== +// OpenVINO Backend Context and Interface +// ===================================================== +struct ggml_backend_openvino_context { + int device = 0; + std::string name = "OpenVINO"; + std::string description = "OpenVINO Backend Context"; + + std::shared_ptr runtime_context = nullptr; + + ggml_backend_openvino_context() = default; +}; + +#include "ggml-openvino/utils.h" +/* Inlined ggml-quants-defs.inc */ +#include "ggml.h.inc" + +#include +#include +#include + +void unpack_32_4(const uint8_t* data, uint8_t* dst); + +void extract_q4_0_data(const ggml_tensor * tensor, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr); + +void extract_q4_1_data(const ggml_tensor * tensor, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr, + bool use_bias = false); + +void extract_q8_0_data(const ggml_tensor * tensor, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr); + +void unpack_256_4(const uint8_t* data, uint8_t* dst); + +void extract_q4_k_data(const ggml_tensor * tensor, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr, + bool use_bias = false); + +void extract_q5_k_data(const ggml_tensor * tensor, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr, + bool use_bias = false); + +void extract_q6_k_data(const ggml_tensor * tensor, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr); + +static constexpr size_t GGML_QUANTIZATION_GROUP_SIZE = 32; + +ov::Output make_int8_weights(ov::Tensor & weight, + ov::Tensor & scales, + ov::Tensor & zp, + size_t group_size = GGML_QUANTIZATION_GROUP_SIZE, + bool use_bias = false); + +ov::Output make_int4_weights(ov::Tensor & weight, + ov::Tensor & scales, + ov::Tensor & zp, + size_t group_size = GGML_QUANTIZATION_GROUP_SIZE, + bool use_bias = false); + +// Extract quantized weights from tensor and create weight subgraph +// If weights/scales/zp are provided (non-empty), uses them as output buffers +// Otherwise allocates new ov::Tensors internally +// Returns the weight node (make_int4_weights or make_int8_weights result) +std::shared_ptr extract_quantized_weights( + const ggml_tensor * tensor, + const void * data, // Source data pointer (may differ from tensor->data) + ov::Tensor & weights, + ov::Tensor & scales, + ov::Tensor & zp, + bool use_bias = false); // Use fp bias instead of quantized zero_point (for test-backend-ops) + +// Requantize weights from tensor to target format, writing to provided buffers +// For F16 target, only weights buffer is used (scales/zp ignored) +// Returns the weight node +std::shared_ptr requantize_to_buffers(const ggml_tensor * tensor, + const void * data, // Source data pointer + ExtraQuantType requant_type, + int64_t block_size, + ov::Tensor & weights, + ov::Tensor & scales, + ov::Tensor & zp); + +inline const char * extra_quant_type_name(ExtraQuantType t) { + switch (t) { + case ExtraQuantType::F16: + return "F16"; + case ExtraQuantType::Q4_0_C: + return "Q4_0_C"; + case ExtraQuantType::Q4_0_128: + return "Q4_0_128"; + case ExtraQuantType::Q8_0_C: + return "Q8_0_C"; + case ExtraQuantType::Q8_0_32: + return "Q8_0_32"; + case ExtraQuantType::Q8_1_C: + return "Q8_1_C"; + default: + return "unknown"; + } +} + +// Result from process_weight_tensor containing the weight node and tensors. +// For quantized weights, also contains the extracted layout and scale/zp tensors. +struct OvWeight { + std::shared_ptr weight_node; + ggml_openvino_extracted_layout layout; // Only meaningful for quantized (layout.total_size > 0) + ov::Tensor weights; + ov::Tensor scales; + ov::Tensor zp; + + bool is_quantized() const { return layout.scales_size > 0; } +}; + +// Process weight tensor and create an OpenVINO weight node +// Handles F16/F32/BF16 and quantized weights, with optional requantization +// If output_base_ptr is nullptr, allocates internal buffers (for decoder use) +// If output_base_ptr is provided, uses pre-allocated buffers at specified offsets (for backend buffer use) +// Returns OvWeight with the weight node and optional quantized tensors +OvWeight process_weight_tensor( + const ggml_tensor * tensor, + const void * data, // Source data pointer (may differ from tensor->data) + void * output_base_ptr = nullptr, // Base pointer for output buffers (or nullptr for internal allocation) + bool use_bias = false); // Use fp bias instead of quantized zero_point, only used in test-backend-ops + +void quantize_q4_0(const float * x, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr, + int64_t k, + int64_t qk); +void quantize_q8_1(const float * x, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr, + int64_t k, + int64_t qk); +void quantize_q8_0(const float * x, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr, + int64_t k, + int64_t qk); + +namespace ov { +namespace op { +namespace util { +// From /src/common/transformations/include/transformations/utils/utils.hpp +bool get_single_value(const std::shared_ptr& const_node, + float& value, + bool check_value_range = true); +} // namespace util +} // namespace op +} // namespace ov + +#include "ggml.h.inc" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +# define WIN32_LEAN_AND_MEAN +# ifndef NOMINMAX +# define NOMINMAX +# endif +# include +#else +# include +#endif + +// ===================================================== +// OpenVINO Buffer Implementation using ov::Tensor +// ===================================================== +// +// Design: This implementation uses a hybrid approach: +// 1. For weight tensors: Store a pre-built ov::op::v0::Constant in tensor->extra +// - This avoids the memcpy during graph construction +// - For quantized weights, the constant is already converted to OpenVINO format +// 2. For KV cache / compute tensors: Store an ov::Tensor in tensor->extra +// - This can be directly passed to infer_request +// - Future: can be changed to ov::RemoteTensor for GPU/NPU +// +// This design is similar to: +// - CUDA split buffer: tensor->extra stores device pointers +// - CPU repack buffer: tensor->extra stores tensor_traits with repacked data +// ===================================================== + +// Buffer context that manages per-tensor allocations (no contiguous buffer for weights) +struct ggml_backend_openvino_buffer_context { + int device; + std::string name; + size_t id; + + // For non-weight buffers (KV cache, compute), we still use contiguous allocation + void * data; + size_t size; + bool is_remote; + + // Wrapping of the buffer + std::shared_ptr ov_buffer; + + // Track all extras for cleanup + std::map tensor_extras; + + // Used for re-allocation on device for kvcache + void * data_prev; + + ggml_backend_openvino_buffer_context(int device, size_t size, bool is_remote = false) : + device(device), + name(std::string(GGML_OPENVINO_NAME) + std::to_string(device)), + id([]() { + static std::atomic next_id{1}; + return next_id.fetch_add(1); + }()), + data(nullptr), + size(size), + is_remote(is_remote) { + if (size == 0) { + return; + } + + const auto & device_name = ggml_openvino_get_device_name(); + + if (is_remote) { + GGML_ASSERT(device_name == "GPU"); + auto remote_context = ggml_openvino_get_remote_context(); + auto gpu_context = remote_context->as(); + ov::intel_gpu::ocl::USMTensor usm_tensor = + gpu_context.create_usm_device_tensor(ov::element::u8, ov::Shape{size}); + data = usm_tensor.get(); + ov_buffer = std::make_shared(std::move(usm_tensor)); + } else { + data = ggml_aligned_malloc(size); + GGML_ASSERT(data); + memset(data, 0, size); + ov_buffer = std::make_shared(ov::element::u8, ov::Shape{size}, data); + } + + if (data == nullptr) { + GGML_LOG_ERROR("%s: failed to allocate %zu bytes\n", __func__, size); + return; + } + + if (reinterpret_cast(data) % TENSOR_ALIGNMENT != 0) { + GGML_LOG_ERROR("%s: %s buffer is not aligned to %d bytes\n", __func__, device_name.c_str(), + TENSOR_ALIGNMENT); + GGML_ABORT("fatal error"); + } + } + + ~ggml_backend_openvino_buffer_context() { + // Clean up all tensor extras + // GGML_LOG_DEBUG("Deleting OpenVINO buffer context #%zu for device %d, size %zu MB\n", id, device, + // size / 1024 / 1024); + for (auto & pair : tensor_extras) { + delete pair.second; + } + tensor_extras.clear(); + if (!is_remote && data != nullptr) { + ggml_aligned_free(data, size); + } + } +}; + +// Buffer type context (per-device) +struct ggml_backend_openvino_buffer_type_context { + int device; + std::string name; +}; + +// Buffer interface functions +static void ggml_backend_openvino_buffer_free_buffer(ggml_backend_buffer_t buffer) { + ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; + delete ctx; +} + +static void * ggml_backend_openvino_buffer_get_base(ggml_backend_buffer_t buffer) { + ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; + return ctx->data; +} + +static bool is_stateful_enabled() { + static const auto * stateful = getenv("GGML_OPENVINO_STATEFUL_EXECUTION"); + return stateful && *stateful != '\0' && strcmp(stateful, "0") != 0; +} + +static enum ggml_status ggml_backend_openvino_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + // GGML_LOG_DEBUG("%s: buffer usage=%d, tensor name=%s\n", __func__, buffer->usage, tensor->name); + ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; + + // Put kvcache on device memory for GPU (NPU memory is too small even for kvcache) + if (strncmp(tensor->name, "cache_", 6) == 0 && !ctx->is_remote && ggml_openvino_get_device_name() == "GPU" && + !is_stateful_enabled()) { + GGML_ASSERT(ctx->tensor_extras.empty()); + auto device = ctx->device; + auto size = ctx->size; + auto * data_prev = ctx->data; + delete ctx; + ctx = new ggml_backend_openvino_buffer_context(device, size, true); + buffer->context = ctx; + tensor->data = (char *) ctx->data + ((char *) tensor->data - (char *) data_prev); + } + + // Views share the extra from view_src + if (tensor->view_src != nullptr) { + GGML_ASSERT(tensor->view_src->buffer->buft == buffer->buft); + if (tensor->view_src->extra != nullptr) { + tensor->extra = tensor->view_src->extra; + } + return GGML_STATUS_SUCCESS; + } + + ctx = (ggml_backend_openvino_buffer_context *) buffer->context; + + if (tensor->data != nullptr && !ggml_is_quantized(tensor->type)) { + ggml_openvino_tensor_extra * extra = ggml_openvino_create_tensor_extra(tensor, ctx->is_remote); + if (extra != nullptr) { + auto it = ctx->tensor_extras.find(tensor); + if (it != ctx->tensor_extras.end()) { + delete it->second; + } + ctx->tensor_extras[tensor] = extra; + tensor->extra = extra; + } + } + + return GGML_STATUS_SUCCESS; +} + +static void ggml_backend_openvino_buffer_memset_tensor(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + uint8_t value, + size_t offset, + size_t size) { + // GGML_LOG_DEBUG("%s: buffer usage=%d, tensor name=%s\n", __func__, buffer->usage, tensor->name); + GGML_ASSERT(tensor != nullptr && tensor->data != nullptr); + ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; + + if (ctx->is_remote) { + // For remote (device) buffers, use OpenCL USM memfill + cl_command_queue queue = ggml_openvino_get_cl_queue(); + auto mem_fill_fn = ggml_openvino_get_clEnqueueMemFillINTEL(); + if (queue != nullptr && mem_fill_fn != nullptr) { + uint8_t pattern = value; + cl_int err = mem_fill_fn(queue, (char *) tensor->data + offset, &pattern, sizeof(pattern), size, 0, nullptr, + nullptr); + if (err != CL_SUCCESS) { + GGML_LOG_ERROR("%s: clEnqueueMemFillINTEL failed with error %d\n", __func__, err); + } + clFinish(queue); + } else { + GGML_LOG_ERROR("%s: no OpenCL queue or clEnqueueMemFillINTEL not available for GPU buffer\n", __func__); + } + } else { + memset((char *) tensor->data + offset, value, size); + } +} + +static void ggml_backend_openvino_buffer_set_tensor(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size) { + // GGML_LOG_DEBUG("%s: buffer usage=%d, tensor name=%s\n", __func__, buffer->usage, tensor->name); + GGML_ASSERT(tensor != nullptr && tensor->data != nullptr); + ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; + + // Check if this is a weight buffer (usage is set BEFORE set_tensor is called, except in test-backend-ops) + bool is_weight_buffer = (buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + // Full tensor set: offset=0, full size, not a view + bool is_full_tensor_set = (offset == 0 && size == ggml_nbytes(tensor) && tensor->view_src == nullptr); + // 2D tensor (typical weight shape) + bool is_2d = (tensor->ne[2] == 1 && tensor->ne[3] == 1); + + if (is_weight_buffer && is_full_tensor_set && is_2d) { + try { + auto result = process_weight_tensor(tensor, data, tensor->data); + result.weight_node->set_friendly_name(tensor->name); + + // const auto & layout = result.layout; + ggml_openvino_extra_base * extra; + + // Quantized path with extracted weight/scale/zp tensors + if (result.is_quantized()) { + extra = new ggml_openvino_quantized_weight_extra(std::move(result.weights), std::move(result.scales), + std::move(result.zp), result.weight_node); + + // if (layout.is_requant) { + // GGML_LOG_DEBUG("%s: requantized %s to %s (u%d, block_size=%ld)\n", __func__, tensor->name, + // extra_quant_type_name(layout.requant_type.value()), layout.is_u4 ? 4 : 8, + // layout.weights_per_block); + // } else { + // int64_t n_blocks = ggml_nelements(tensor) / layout.weights_per_block; + // GGML_LOG_DEBUG("%s: extracted quantized weight node for %s (u%d, %zu weights, %ld blocks)\n", + // __func__, tensor->name, layout.is_u4 ? 4 : 8, layout.weights_size, n_blocks); + // } + } else { + // F16/F32/BF16 weight or F16-requant + extra = new ggml_openvino_weight_extra(std::move(result.weights), result.weight_node); + + // if (layout.total_size > 0) { + // GGML_LOG_DEBUG("%s: requantized %s to F16\n", __func__, tensor->name); + // } else { + // GGML_LOG_DEBUG("%s: created shared-memory weight node for %s\n", __func__, tensor->name); + // } + } + + ctx->tensor_extras[tensor] = extra; + tensor->extra = extra; + + } catch (const std::exception & e) { + GGML_LOG_ERROR("%s: failed to process weight tensor for %s: %s\n", __func__, tensor->name, e.what()); + memcpy((char *) tensor->data + offset, data, size); + } + } else { + // Non-weight tensor (KV cache, activations, etc.) - copy data. test-backend-ops also goes here + if (ctx->is_remote) { + cl_command_queue queue = ggml_openvino_get_cl_queue(); + auto mem_cpy_fn = ggml_openvino_get_clEnqueueMemcpyINTEL(); + if (queue != nullptr && mem_cpy_fn != nullptr) { + cl_int err = + mem_cpy_fn(queue, CL_TRUE, (char *) tensor->data + offset, data, size, 0, nullptr, nullptr); + if (err != CL_SUCCESS) { + GGML_LOG_ERROR("%s: clEnqueueMemcpyINTEL failed with error %d\n", __func__, err); + } + } else { + GGML_LOG_ERROR("%s: no OpenCL queue or clEnqueueMemcpyINTEL not available for GPU buffer\n", __func__); + } + } else { + memcpy((char *) tensor->data + offset, data, size); + } + + ggml_openvino_tensor_extra * extra = ggml_openvino_create_tensor_extra(tensor, ctx->is_remote); + if (extra == nullptr) { + // GGML_LOG_ERROR("%s: failed to create tensor extra for %s\n", __func__, tensor->name); + return; + } + + auto it = ctx->tensor_extras.find(tensor); + if (it != ctx->tensor_extras.end()) { + delete it->second; + } + ctx->tensor_extras[tensor] = extra; + tensor->extra = extra; + } +} + +static void ggml_backend_openvino_buffer_get_tensor(ggml_backend_buffer_t buffer, + const ggml_tensor * tensor, + void * data, + size_t offset, + size_t size) { + // GGML_LOG_DEBUG("%s: buffer usage=%d, tensor name=%s\n", __func__, buffer->usage, tensor->name); + GGML_ASSERT(tensor != nullptr && tensor->data != nullptr); + ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; + + if (ctx->is_remote) { + // For remote (device) buffers, use OpenCL USM memcpy (device-to-host) + cl_command_queue queue = ggml_openvino_get_cl_queue(); + auto mem_cpy_fn = ggml_openvino_get_clEnqueueMemcpyINTEL(); + if (queue != nullptr && mem_cpy_fn != nullptr) { + cl_int err = + mem_cpy_fn(queue, CL_TRUE, data, (const char *) tensor->data + offset, size, 0, nullptr, nullptr); + if (err != CL_SUCCESS) { + GGML_LOG_ERROR("%s: clEnqueueMemcpyINTEL failed with error %d\n", __func__, err); + } + } else { + GGML_LOG_ERROR("%s: no OpenCL queue or clEnqueueMemcpyINTEL not available for GPU buffer\n", __func__); + } + } else { + memcpy(data, (const char *) tensor->data + offset, size); + } +} + +static bool ggml_backend_openvino_buffer_cpy_tensor(ggml_backend_buffer_t buffer, + const ggml_tensor * src, + ggml_tensor * dst) { + // GGML_LOG_DEBUG("%s: src tensor name=%s, dst tensor name=%s\n", __func__, src->name, dst->name); + GGML_ASSERT(src != nullptr && dst != nullptr); + ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; + + if (ctx->is_remote) { + // For remote (device) buffers, use OpenCL USM memcpy + cl_command_queue queue = ggml_openvino_get_cl_queue(); + auto mem_cpy_fn = ggml_openvino_get_clEnqueueMemcpyINTEL(); + if (queue == nullptr || mem_cpy_fn == nullptr) { + GGML_LOG_ERROR("%s: no OpenCL queue or clEnqueueMemcpyINTEL not available for GPU buffer\n", __func__); + return false; + } + // Can copy from host to device + if (ggml_backend_buffer_is_host(src->buffer)) { + cl_int err = mem_cpy_fn(queue, CL_TRUE, dst->data, src->data, ggml_nbytes(src), 0, nullptr, nullptr); + if (err != CL_SUCCESS) { + GGML_LOG_ERROR("%s: clEnqueueMemcpyINTEL (host-to-device) failed with error %d\n", __func__, err); + return false; + } + return true; + } + // Can also copy from device to device if both are OpenVINO remote buffers + if (ggml_backend_buffer_is_openvino(src->buffer)) { + ggml_backend_openvino_buffer_context * src_ctx = + (ggml_backend_openvino_buffer_context *) src->buffer->context; + if (src_ctx->is_remote) { + cl_int err = + mem_cpy_fn(queue, CL_TRUE, dst->data, src->data, ggml_nbytes(src), 0, nullptr, nullptr); + if (err != CL_SUCCESS) { + GGML_LOG_ERROR("%s: clEnqueueMemcpyINTEL (device-to-device) failed with error %d\n", __func__, + err); + return false; + } + return true; + } + } + return false; + } + + // Host buffer - can copy from any host buffer + if (ggml_backend_buffer_is_host(src->buffer)) { + memcpy(dst->data, src->data, ggml_nbytes(src)); + return true; + } + return false; +} + +static void ggml_backend_openvino_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; + GGML_ASSERT(ctx->data != nullptr); + if (ctx->is_remote) { + cl_command_queue queue = ggml_openvino_get_cl_queue(); + auto mem_fill_fn = ggml_openvino_get_clEnqueueMemFillINTEL(); + if (queue != nullptr && mem_fill_fn != nullptr) { + uint8_t pattern = value; + cl_int err = mem_fill_fn(queue, ctx->data, &pattern, sizeof(pattern), ctx->size, 0, nullptr, nullptr); + if (err != CL_SUCCESS) { + GGML_LOG_WARN("%s: clEnqueueMemFillINTEL failed with error %d\n", __func__, err); + } + clFinish(queue); + } else { + GGML_LOG_WARN("%s: no OpenCL queue or clEnqueueMemFillINTEL not available for GPU buffer clear\n", + __func__); + } + } else { + memset(ctx->data, value, ctx->size); + } +} + +static const ggml_backend_buffer_i ggml_backend_openvino_buffer_interface = { + /* .free_buffer = */ ggml_backend_openvino_buffer_free_buffer, + /* .get_base = */ ggml_backend_openvino_buffer_get_base, + /* .init_tensor = */ ggml_backend_openvino_buffer_init_tensor, + /* .memset_tensor = */ ggml_backend_openvino_buffer_memset_tensor, + /* .set_tensor = */ ggml_backend_openvino_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_openvino_buffer_get_tensor, + /* .set_tensor_2d = */ NULL, + /* .get_tensor_2d = */ NULL, + /* .cpy_tensor = */ ggml_backend_openvino_buffer_cpy_tensor, + /* .clear = */ ggml_backend_openvino_buffer_clear, + /* .reset = */ NULL, +}; + +// Buffer type interface functions +static const char * ggml_backend_openvino_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + ggml_backend_openvino_buffer_type_context * ctx = (ggml_backend_openvino_buffer_type_context *) buft->context; + return ctx->name.c_str(); +} + +static ggml_backend_buffer_t ggml_backend_openvino_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, + size_t size) { + ggml_backend_openvino_buffer_type_context * buft_ctx = (ggml_backend_openvino_buffer_type_context *) buft->context; + + // Create buffer context with contiguous memory allocation + ggml_backend_openvino_buffer_context * ctx = new ggml_backend_openvino_buffer_context(buft_ctx->device, size); + + if (ctx->data == nullptr && size > 0) { + GGML_LOG_ERROR("%s: failed to allocate buffer of size %zu\n", __func__, size); + delete ctx; + return nullptr; + } + + return ggml_backend_buffer_init(buft, ggml_backend_openvino_buffer_interface, ctx, size); +} + +static size_t ggml_backend_openvino_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + GGML_UNUSED(buft); + return TENSOR_ALIGNMENT; +} + +static size_t ggml_backend_openvino_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { + GGML_UNUSED(buft); + return SIZE_MAX; +} + +static size_t ggml_backend_openvino_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, + const ggml_tensor * tensor) { + GGML_UNUSED(buft); + + // For quantized 2D tensors (weights), we need extra space for extracted data + if (ggml_is_quantized(tensor->type) && tensor->ne[2] == 1 && tensor->ne[3] == 1) { + ggml_openvino_extracted_layout layout = ggml_openvino_get_extracted_layout(tensor); + if (layout.total_size > 0) { + // GGML_LOG_DEBUG("%s: tensor %s needs %zu bytes (original %zu, extracted: weights=%zu scales=%zu zp=%zu)\n", + // __func__, tensor->name, layout.total_size, ggml_nbytes(tensor), layout.weights_size, + // layout.scales_size, layout.zp_size); + return layout.total_size; + } + } + + return ggml_nbytes(tensor); +} + +static const ggml_backend_buffer_type_i ggml_backend_openvino_buffer_type_interface = { + /* .get_name = */ ggml_backend_openvino_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_openvino_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_openvino_buffer_type_get_alignment, + /* .get_max_size = */ ggml_backend_openvino_buffer_type_get_max_size, + /* .get_alloc_size = */ ggml_backend_openvino_buffer_type_get_alloc_size, + /* .is_host = */ nullptr, +}; + +// Get buffer type for a specific device +GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_openvino_buffer_type(int device) { + GGML_ASSERT(device >= 0 && device < ggml_backend_openvino_get_device_count()); + + static std::mutex mutex; + std::lock_guard lock(mutex); + + static std::vector buffer_types; + static std::vector buffer_type_contexts; + + if (buffer_types.empty()) { + int device_count = ggml_backend_openvino_get_device_count(); + buffer_types.resize(device_count); + buffer_type_contexts.resize(device_count); + + for (int i = 0; i < device_count; i++) { + buffer_type_contexts[i].device = i; + buffer_type_contexts[i].name = std::string(GGML_OPENVINO_NAME) + std::to_string(i); + + buffer_types[i] = ggml_backend_buffer_type{ + /* .iface = */ ggml_backend_openvino_buffer_type_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_openvino_reg(), i), + /* .context = */ &buffer_type_contexts[i], + }; + } + } + + return &buffer_types[device]; +} + +// ===================================================== +// OpenVINO Host Buffer Implementation +// ===================================================== + +static const char * ggml_backend_openvino_host_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + ggml_backend_openvino_buffer_type_context * ctx = (ggml_backend_openvino_buffer_type_context *) buft->context; + static std::string name; + name = ctx->name + "_HOST"; + return name.c_str(); +} + +static bool ggml_backend_openvino_host_buffer_type_is_host(ggml_backend_buffer_type_t buft) { + GGML_UNUSED(buft); + return true; +} + +static const ggml_backend_buffer_type_i ggml_backend_openvino_host_buffer_type_interface = { + /* .get_name = */ ggml_backend_openvino_host_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_openvino_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_openvino_buffer_type_get_alignment, + /* .get_max_size = */ ggml_backend_openvino_buffer_type_get_max_size, + /* .get_alloc_size = */ ggml_backend_openvino_buffer_type_get_alloc_size, + /* .is_host = */ ggml_backend_openvino_host_buffer_type_is_host, +}; + +GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_openvino_host_buffer_type(int device) { + GGML_ASSERT(device >= 0 && device < ggml_backend_openvino_get_device_count()); + + static std::mutex mutex; + std::lock_guard lock(mutex); + + static std::vector buffer_types; + static std::vector buffer_type_contexts; + + if (buffer_types.empty()) { + int device_count = ggml_backend_openvino_get_device_count(); + buffer_types.resize(device_count); + buffer_type_contexts.resize(device_count); + + for (int i = 0; i < device_count; i++) { + buffer_type_contexts[i].device = i; + buffer_type_contexts[i].name = std::string(GGML_OPENVINO_NAME) + std::to_string(i); + + buffer_types[i] = ggml_backend_buffer_type{ + /* .iface = */ ggml_backend_openvino_host_buffer_type_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_openvino_reg(), i), + /* .context = */ &buffer_type_contexts[i], + }; + } + } + + return &buffer_types[device]; +} + +bool ggml_backend_buffer_is_openvino(ggml_backend_buffer_t buffer) { + return buffer->iface.free_buffer == ggml_backend_openvino_buffer_free_buffer; +} + +size_t ggml_backend_openvino_buffer_get_ctx_id(ggml_backend_buffer_t buffer) { + if (!ggml_backend_buffer_is_openvino(buffer)) { + return 0; + } + ggml_backend_openvino_buffer_context * ctx = (ggml_backend_openvino_buffer_context *) buffer->context; + return ctx->id; +} + +void ggml_openvino_buffer_register_extra(ggml_tensor * tensor, ggml_openvino_extra_base * extra) { + GGML_ASSERT(tensor != nullptr); + GGML_ASSERT(tensor->buffer != nullptr); + GGML_ASSERT(ggml_backend_buffer_is_openvino(tensor->buffer)); + + auto * ctx = static_cast(tensor->buffer->context); + + auto it = ctx->tensor_extras.find(tensor); + if (it != ctx->tensor_extras.end()) { + delete it->second; + } + + ctx->tensor_extras[tensor] = extra; + tensor->extra = extra; +} + +bool ggml_backend_buft_is_openvino(ggml_backend_buffer_type_t buft) { + return buft->iface.get_name == ggml_backend_openvino_buffer_type_get_name; +} + +bool ggml_backend_buft_is_openvino_host(ggml_backend_buffer_type_t buft) { + return buft->iface.get_name == ggml_backend_openvino_host_buffer_type_get_name; +} + +static void ggml_backend_openvino_free(ggml_backend_t backend) { + ggml_backend_openvino_context * ctx = (ggml_backend_openvino_context *) backend->context; + + if (ctx->runtime_context) { + auto r_ctx = std::static_pointer_cast(ctx->runtime_context); + if (--r_ctx->backend_count == 0) { + r_ctx->clear_caches(); + } + } + + delete ctx; + delete backend; +} + +static const char * ggml_backend_openvino_get_name(ggml_backend_t backend) { + return GGML_OPENVINO_NAME; + GGML_UNUSED(backend); +} + +static enum ggml_status ggml_backend_openvino_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { + return ov_graph_compute(cgraph, backend); + GGML_UNUSED(backend); +} + +static const ggml_backend_i ggml_backend_openvino_interface = { + /* .get_name = */ ggml_backend_openvino_get_name, + /* .free = */ ggml_backend_openvino_free, + /* .set_tensor_async = */ NULL, + /* .get_tensor_async = */ NULL, + /* .set_tensor_2d_async = */ NULL, + /* .get_tensor_2d_async = */ NULL, + /* .cpy_tensor_async = */ NULL, + /* .synchronize = */ NULL, + /* .graph_plan_create = */ NULL, + /* .graph_plan_free = */ NULL, + /* .graph_plan_update = */ NULL, + /* .graph_plan_compute = */ NULL, + /* .graph_compute = */ ggml_backend_openvino_graph_compute, + /* .event_record = */ NULL, + /* .event_wait = */ NULL, + /* .graph_optimize = */ NULL, +}; + +int ggml_backend_openvino_get_device_count() { + return 1; +} + +static ggml_guid_t ggml_backend_openvino_guid(void) { + static ggml_guid guid = {0x12, 0xa8, 0xae, 0xf4, 0xc0, 0x1e, 0x61, 0x97, + 0x8f, 0xeb, 0x33, 0x04, 0xa1, 0x33, 0x51, 0x2d}; + return &guid; +} + +static std::shared_ptr get_ov_runtime_context_ptr() { + static std::shared_ptr r_ctx = [] { + auto ctx = std::make_shared(); + ctx->device = ggml_openvino_get_device_name(); + ctx->stateful = is_stateful_enabled() && !ggml_openvino_is_npu(); + return ctx; + }(); + return r_ctx; +} + +// backend API +GGML_BACKEND_API ggml_backend_t ggml_backend_openvino_init(int device) { + if (device < 0 || device >= ggml_backend_openvino_get_device_count()) { + GGML_LOG_ERROR("%s: invalid device %d\n", __func__, device); + return nullptr; + } + + ggml_backend_openvino_context * ctx = new ggml_backend_openvino_context; + if (ctx == nullptr) { + GGML_LOG_ERROR("%s: failed to allocate context\n", __func__); + return nullptr; + } + + ctx->runtime_context = get_ov_runtime_context_ptr(); + if (ctx->runtime_context == nullptr) { + GGML_LOG_ERROR("%s: failed to allocate runtime context\n", __func__); + delete ctx; + return nullptr; + } + + std::shared_ptr r_ctx = std::static_pointer_cast(ctx->runtime_context); + r_ctx->backend_count++; + + ggml_backend_t openvino_backend = new ggml_backend{ + /* .guid = */ ggml_backend_openvino_guid(), + /* .interface = */ ggml_backend_openvino_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_openvino_reg(), device), + /* .context = */ ctx, + }; + + return openvino_backend; +} + +GGML_BACKEND_API bool ggml_backend_is_openvino(ggml_backend_t backend) { + return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_openvino_guid()); +} + +struct ggml_backend_openvino_device_context { + int device; + std::string name; + std::string description; +}; + +static const char * ggml_backend_openvino_device_get_name(ggml_backend_dev_t dev) { + ggml_backend_openvino_device_context * ctx = (ggml_backend_openvino_device_context *) dev->context; + return ctx->name.c_str(); +} + +static const char * ggml_backend_openvino_device_get_description(ggml_backend_dev_t dev) { + ggml_backend_openvino_device_context * ctx = (ggml_backend_openvino_device_context *) dev->context; + return ctx->description.c_str(); +} + +static void ggml_backend_openvino_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { +#ifdef _WIN32 + MEMORYSTATUSEX status; + status.dwLength = sizeof(status); + GlobalMemoryStatusEx(&status); + *total = status.ullTotalPhys; + *free = status.ullAvailPhys; +#else + long pages = sysconf(_SC_PHYS_PAGES); + long page_size = sysconf(_SC_PAGE_SIZE); + *total = pages * page_size; + + // "free" system memory is ill-defined, for practical purposes assume that all of it is free: + *free = *total; +#endif // _WIN32 + + GGML_UNUSED(dev); +} + +static enum ggml_backend_dev_type ggml_backend_openvino_device_get_type(ggml_backend_dev_t dev) { + GGML_UNUSED(dev); + return GGML_BACKEND_DEVICE_TYPE_GPU; +} + +static void ggml_backend_openvino_device_get_props(ggml_backend_dev_t dev, ggml_backend_dev_props * props) { + props->name = ggml_backend_openvino_device_get_name(dev); + props->description = ggml_backend_openvino_device_get_description(dev); + props->type = ggml_backend_openvino_device_get_type(dev); + ggml_backend_openvino_device_get_memory(dev, &props->memory_free, &props->memory_total); + + props->caps = { + /* .async = */ false, + /* .host_buffer = */ false, + /* .buffer_from_host_ptr = */ false, + /* .events = */ false, + }; +} + +static ggml_backend_t ggml_backend_openvino_device_init(ggml_backend_dev_t dev, const char * params) { + GGML_UNUSED(params); + ggml_backend_openvino_device_context * ctx = (ggml_backend_openvino_device_context *) dev->context; + return ggml_backend_openvino_init(ctx->device); +} + +static ggml_backend_buffer_type_t ggml_backend_openvino_device_get_buffer_type(ggml_backend_dev_t dev) { + ggml_backend_openvino_device_context * ctx = (ggml_backend_openvino_device_context *) dev->context; + return ggml_backend_openvino_buffer_type(ctx->device); +} + +static ggml_backend_buffer_type_t ggml_backend_openvino_device_get_host_buffer_type(ggml_backend_dev_t dev) { + ggml_backend_openvino_device_context * ctx = (ggml_backend_openvino_device_context *) dev->context; + return ggml_backend_openvino_host_buffer_type(ctx->device); +} + +static bool has_view_op_input(const ggml_tensor * op) { + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (op->src[i] == nullptr) { + break; + } + if (op->src[i]->op == GGML_OP_VIEW) { + return true; + } + } + return false; +} + +static bool is_supported_flash_attn_pattern(const ggml_tensor * op) { + // pattern of q,k,v should be q->op==PERMUTE, q->src[0]->op==VIEW, q->src[0]->src[0]->view_src==nullptr + for (int i = 0; i < 3; i++) { + const ggml_tensor * src = op->src[i]; + if (src->op != GGML_OP_PERMUTE || src->src[0] == nullptr || src->src[0]->op != GGML_OP_VIEW || + src->src[0]->src[0] == nullptr || src->src[0]->src[0]->view_src != nullptr) { + return false; + } + } + return true; +} + +static bool is_op_unsupported_case(const ggml_tensor * op) { + switch (op->op) { + case GGML_OP_GET_ROWS: + case GGML_OP_SET_ROWS: { + if (op->ne[3] != 1) { + return true; + } + break; + } + case GGML_OP_ADD: + case GGML_OP_MUL: { + if (op->src[1]->op == GGML_OP_PERMUTE) { + return true; + } + for (int i = 0; i < 4; i++) { + if (op->src[0]->ne[i] != op->src[1]->ne[i] && (op->src[0]->ne[i] != 1 && op->src[1]->ne[i] != 1)) { + return true; + } + } + break; + } + case GGML_OP_SOFT_MAX: { + if (op->src[2] != nullptr) { + // GGML_LOG_WARN("OpenVINO backend does not support SOFT_MAX with sinks\n"); + return true; + } + float scale = 1.0f; + float max_bias = 0.0f; + const auto * op_params = op->op_params; + memcpy(&scale, (const float *) op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) op_params + 1, sizeof(float)); + if (max_bias > 0) { + // GGML_LOG_WARN("OpenVINO backend does not support SOFT_MAX with max_bias > 0\n"); + return true; + } + break; + } + case GGML_OP_FLASH_ATTN_EXT: { + if (op->src[4] != nullptr) { + // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with sinks\n"); + return true; + } + if (!is_supported_flash_attn_pattern(op)) { + return true; + } + float scale = 1.0f; + float max_bias = 0.0f; + float logit_softcap = 0.0f; + const auto * op_params = op->op_params; + memcpy(&scale, (const float *) op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) op_params + 2, sizeof(float)); + if (max_bias > 0) { + // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with max_bias > 0\n"); + return true; + } + if (logit_softcap != 0) { + // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with logit_softcap != 0\n"); + return true; + } + break; + } + case GGML_OP_PERMUTE: { + if (op->type == GGML_TYPE_BF16) { + // err msg: [GPU] Could not find a suitable kernel for transpose + // GGML_LOG_WARN("OpenVINO backend does not support PERMUTE with BF16 type\n"); + return true; + } + break; + } + case GGML_OP_CPY: { + if (op->src[1] != op) { + // GGML_LOG_WARN("OpenVINO backend only supports CPY that is a cast\n"); + return true; + } + break; + } + case GGML_OP_MUL_MAT: { + if (op->src[0]->type == GGML_TYPE_F16 && op->src[1]->type == GGML_TYPE_F16) { + // Has accuracy issue, try enabling this and see `test-backend-ops -o "MUL_MAT"` + // GGML_LOG_WARN("OpenVINO backend does not support MUL_MAT with two F16 tensors\n"); + return true; + } + if (op->src[0]->ne[3] != op->src[1]->ne[3] && op->src[0]->ne[3] != 1 && op->src[1]->ne[3] != 1) { + return true; + } + if (op->src[0]->op == GGML_OP_PERMUTE || op->src[1]->op == GGML_OP_PERMUTE) { + return true; + } + if (ggml_is_quantized(op->src[0]->type) && op->src[0]->ne[1] == 1) { + // MUL_MAT(type_a=q4_0,type_b=f32,m=1,n=2048,k=8192,bs=[1,1],nr=[1,1],per=[0,1,2,3],k_v=0,o=1) + // triggers a bug in ov matmul_shape_inference.hpp + return true; + } + if (op->src[0]->op == GGML_OP_VIEW && op->src[1]->op == GGML_OP_VIEW) { + return true; + } + break; + } + case GGML_OP_ROPE: { + const int32_t * op_params = op->op_params; + const int n_dims = op_params[1]; + const int mode = op_params[2]; + if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX && mode != GGML_ROPE_TYPE_IMROPE) { + // GGML_LOG_WARN("OpenVINO backend does not support ROPE with mode %d\n", mode); + return true; + } + if (n_dims != 0.0f && n_dims != op->src[0]->ne[0]) { + // GGML_LOG_WARN("OpenVINO backend does not support ROPE with n_dims %d != src[0]->ne[0] %ld\n", n_dims, + // op->src[0]->ne[0]); + return true; + } + if (op->type != GGML_TYPE_F32) { + // GGML_LOG_WARN("OpenVINO backend does not support ROPE with type %s\n", ggml_type_name(op->type)); + return true; + } + if (op->src[0]->op == GGML_OP_VIEW) { + if (op->src[0]->view_src->ne[1] != op->src[0]->ne[2]) { + // GGML_LOG_WARN( + // "OpenVINO backend does not support ROPE with src[0]->view_src->ne[1] %ld != src[0]->ne[2] " + // "%ld\n", + // op->src[0]->view_src->ne[1], op->src[0]->ne[2]); + return true; + } + } + if (mode == GGML_ROPE_TYPE_IMROPE && + (op->src[2] != 0 || ((const float *) op_params)[6] != 1 || ((const float *) op_params)[7] != 0 || + ((const float *) op_params)[8] != 1)) { + // GGML_LOG_WARN("OpenVINO backend does not support IMROPE with freq_factors, freq_scale, ext_factor, and attn_factor\n"); + return true; + } + break; + } + default: + break; + } + if (op->op == GGML_OP_GET_ROWS) { + if (op->ne[0] == 256 && (op->src[0]->type == GGML_TYPE_Q4_K || op->src[0]->type == GGML_TYPE_Q5_K)) { + // ERR = 0.000000306 > 0.000000100 GET_ROWS(type=q4_K,n=256,m=5,r=4,be1=1,be2=1,v=0) + // ERR = 0.000000197 > 0.000000100 GET_ROWS(type=q5_K,n=256,m=5,r=4,be1=1,be2=1,v=0) + return true; + } + } + return false; +} + +static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + GGML_ASSERT(dev->reg != nullptr); + + static std::set supported_types{GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_I64, + GGML_TYPE_I32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_K, + GGML_TYPE_Q5_K, GGML_TYPE_Q8_0, GGML_TYPE_Q6_K}; + + static const std::set supported_ops{GGML_OP_NONE, GGML_OP_ADD, GGML_OP_MUL, GGML_OP_MUL_MAT, GGML_OP_VIEW, + /*GGML_OP_CONT,*/ GGML_OP_RESHAPE, GGML_OP_PERMUTE, GGML_OP_TRANSPOSE, + GGML_OP_GET_ROWS, GGML_OP_ROPE, GGML_OP_RMS_NORM, GGML_OP_SCALE, + // softmax is not updated due to replaced by flash_attn_ext + // GGML_OP_SOFT_MAX, + GGML_OP_SET_ROWS, GGML_OP_FLASH_ATTN_EXT, GGML_OP_CPY}; + static const std::set supported_unary_ops{ + GGML_UNARY_OP_GELU, + GGML_UNARY_OP_SILU, + }; + static const std::set supported_glu_ops{ + GGML_GLU_OP_SWIGLU, + GGML_GLU_OP_GEGLU, + }; + + switch (op->op) { + case GGML_OP_UNARY: { + auto supported = supported_unary_ops.find(ggml_get_unary_op(op)) != supported_unary_ops.end(); + if (!supported) { + // GGML_LOG_WARN("OpenVINO backend does not support unary op %s\n", ggml_unary_op_name(ggml_get_unary_op(op))); + return false; + } + if (has_view_op_input(op)) { + // GGML_LOG_WARN("OpenVINO backend does not support unary op %s with view input\n", + // ggml_unary_op_name(ggml_get_unary_op(op))); + return false; + } + break; + } + case GGML_OP_GLU: { + auto supported = supported_glu_ops.find(ggml_get_glu_op(op)) != supported_glu_ops.end(); + if (!supported) { + // GGML_LOG_WARN("OpenVINO backend does not support GLU op %s\n", ggml_glu_op_name(ggml_get_glu_op(op))); + return false; + } + if (has_view_op_input(op)) { + // GGML_LOG_WARN("OpenVINO backend does not support unary op %s with view input\n", + // ggml_glu_op_name(ggml_get_glu_op(op))); + return false; + } + if (op->src[1] == nullptr && op->src[0]->ne[0] % 2 != 0) { + // triggers bug in ov gpu + return false; + } + break; + } + default: { + auto supported = supported_ops.find(op->op) != supported_ops.end(); + if (!supported) { + // GGML_LOG_WARN("OpenVINO backend does not support op %s\n", ggml_op_name(op->op)); + return false; + } + static std::set ops_not_support_view_input{ + GGML_OP_GET_ROWS, + GGML_OP_RMS_NORM, + }; + if (ops_not_support_view_input.find(op->op) != ops_not_support_view_input.end() && has_view_op_input(op)) { + // GGML_LOG_WARN("OpenVINO backend does not support op %s with view input\n", ggml_op_name(op->op)); + return false; + } + } + } + + if (supported_types.find(op->type) == supported_types.end()) { + // GGML_LOG_WARN("OpenVINO backend does not support tensor type %s\n", ggml_type_name(op->type)); + return false; + } + for (int i = 0; i < GGML_MAX_SRC; i++) { + auto * src = op->src[i]; + if (src == nullptr) { + break; + } + if (supported_types.find(src->type) == supported_types.end()) { + // GGML_LOG_WARN("OpenVINO backend does not support tensor type %s\n", ggml_type_name(src->type)); + return false; + } + if (ggml_is_quantized(src->type) && src->ne[2] != 1) { + // GGML_LOG_WARN("OpenVINO backend does not support 3D quantized tensors\n"); + return false; + } + } + + if (is_op_unsupported_case(op)) { + return false; + } + return true; +} + +static bool ggml_backend_openvino_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { + return ggml_backend_buft_is_openvino(buft) || ggml_backend_buft_is_host(buft); + GGML_UNUSED(dev); +} + +static const struct ggml_backend_device_i ggml_backend_openvino_device_interface = { + /* .get_name = */ ggml_backend_openvino_device_get_name, + /* .get_description = */ ggml_backend_openvino_device_get_description, + /* .get_memory = */ ggml_backend_openvino_device_get_memory, + /* .get_type = */ ggml_backend_openvino_device_get_type, + /* .get_props = */ ggml_backend_openvino_device_get_props, + /* .init_backend = */ ggml_backend_openvino_device_init, + /* .get_buffer_type = */ ggml_backend_openvino_device_get_buffer_type, + /* .get_host_buffer_type = */ ggml_backend_openvino_device_get_host_buffer_type, + /* .buffer_from_host_ptr = */ NULL, + /* .supports_op = */ ggml_backend_openvino_device_supports_op, + /* .supports_buft = */ ggml_backend_openvino_device_supports_buft, + /* .offload_op = */ NULL, + /* .event_new = */ NULL, + /* .event_free = */ NULL, + /* .event_synchronize = */ NULL, +}; + +struct ggml_backend_openvino_reg_context { + std::vector devices; +}; + +static const char * ggml_backend_openvino_reg_get_name(ggml_backend_reg_t reg) { + return GGML_OPENVINO_NAME; + GGML_UNUSED(reg); +} + +static size_t ggml_backend_openvino_reg_get_device_count(ggml_backend_reg_t reg) { + GGML_UNUSED(reg); + return (size_t) ggml_backend_openvino_get_device_count(); +} + +static ggml_backend_dev_t ggml_backend_openvino_reg_get_device(ggml_backend_reg_t reg, size_t index) { + ggml_backend_openvino_reg_context * ctx = (ggml_backend_openvino_reg_context *) reg->context; + GGML_ASSERT(index < ctx->devices.size()); + return ctx->devices[index]; +} + +static const struct ggml_backend_reg_i ggml_backend_openvino_reg_interface = { + /* .get_name = */ ggml_backend_openvino_reg_get_name, + /* .get_device_count = */ ggml_backend_openvino_reg_get_device_count, + /* .get_device = */ ggml_backend_openvino_reg_get_device, + /* .get_proc_address = */ NULL, +}; + +static void ggml_openvino_init() { + // Initialize device config singleton from env var + ggml_openvino_init_device_config(); + GGML_LOG_INFO("OpenVINO: using device %s\n", ggml_openvino_get_device_name().c_str()); +} + +GGML_BACKEND_API ggml_backend_reg_t ggml_backend_openvino_reg(void) { + static ggml_backend_reg reg; + + static bool initialized = false; + { + static std::mutex mutex; + std::lock_guard lock(mutex); + if (!initialized) { + ggml_openvino_init(); + + ggml_backend_openvino_reg_context * ctx = new ggml_backend_openvino_reg_context; + + for (int i = 0; i < ggml_backend_openvino_get_device_count(); i++) { + ggml_backend_openvino_device_context * dev_ctx = new ggml_backend_openvino_device_context; + dev_ctx->device = i; + dev_ctx->name = GGML_OPENVINO_NAME + std::to_string(i); + + dev_ctx->description = ov::get_openvino_version().description; + + ggml_backend_dev_t dev = + new ggml_backend_device{/* .interface = */ ggml_backend_openvino_device_interface, + /* .reg = */ ®, + /* .context = */ dev_ctx}; + ctx->devices.push_back(dev); + } + + reg = ggml_backend_reg{/* .api_version = */ GGML_BACKEND_API_VERSION, + /* .iface = */ ggml_backend_openvino_reg_interface, + /* .context = */ ctx}; + } + + initialized = true; + } + + return ® +} + + +/* Inlined ggml-openvino-extra.cpp */ + +#include "ggml-impl-defs.inc" +#include "ggml.h.inc" + +#include +#include +#include +#include +#include + +ov::Core & ov_singleton_core() { + static ov::Core core; + return core; +} + +// ===================================================== +// Device Configuration Implementations +// ===================================================== + +void ggml_openvino_device_config::init() { + if (initialized) { + return; + } + device_name = getenv("GGML_OPENVINO_DEVICE") ? getenv("GGML_OPENVINO_DEVICE") : "CPU"; + auto available_devices = ov_singleton_core().get_available_devices(); + if (std::find(available_devices.begin(), available_devices.end(), device_name) == available_devices.end()) { + GGML_LOG_WARN("GGML OpenVINO Backend: device %s is not available, fallback to CPU\n", device_name.c_str()); + device_name = "CPU"; + } + is_npu = (device_name == "NPU"); + + auto * cache_dir = getenv("GGML_OPENVINO_CACHE_DIR"); + if (device_name == "NPU") { + compile_config = { + {"NPU_COMPILER_DYNAMIC_QUANTIZATION", "YES" }, + {"NPU_USE_NPUW", "YES" }, + {"NPUW_DEVICES", "NPU" }, + {"NPUW_FOLD", "YES" }, + {"NPUW_WEIGHTS_BANK", "shared"}, + {"NPUW_FUNCALL_FOR_ALL", "YES" }, + {"NPUW_FUNCALL_ASYNC", "YES" }, + {"NPUW_DQ", "YES" }, + {"NPUW_DQ_FULL", "NO" }, + }; + if (cache_dir && strlen(cache_dir) > 0) { + compile_config["NPUW_CACHE_DIR"] = cache_dir; + compile_config.insert(ov::cache_mode(ov::CacheMode::OPTIMIZE_SIZE)); + } + } else if (cache_dir && strlen(cache_dir) > 0) { + compile_config.insert(ov::cache_dir(cache_dir)); + compile_config.insert(ov::cache_mode(ov::CacheMode::OPTIMIZE_SIZE)); + } + + // Initialize remote context with queue sharing for GPU + if (device_name == "GPU") { + // Create OpenCL context and queue + cl_int err; + cl_platform_id platform; + err = clGetPlatformIDs(1, &platform, nullptr); + if (err != CL_SUCCESS) { + GGML_LOG_ERROR("Failed to get OpenCL platform: %d\n", err); + return; + } + + cl_device_id cl_device; + err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_GPU, 1, &cl_device, nullptr); + if (err != CL_SUCCESS) { + GGML_LOG_ERROR("Failed to get OpenCL device: %d\n", err); + return; + } + + cl_context cl_ctx = clCreateContext(nullptr, 1, &cl_device, nullptr, nullptr, &err); + if (err != CL_SUCCESS) { + GGML_LOG_ERROR("Failed to create OpenCL context: %d\n", err); + return; + } + + cl_queue = clCreateCommandQueueWithProperties(cl_ctx, cl_device, nullptr, &err); + if (err != CL_SUCCESS) { + GGML_LOG_ERROR("Failed to create OpenCL command queue: %d\n", err); + clReleaseContext(cl_ctx); + return; + } + + // Create OpenVINO remote context with queue sharing + remote_context = ov::intel_gpu::ocl::ClContext(ov_singleton_core(), cl_queue); + + // Release the context (queue keeps a reference) + clReleaseContext(cl_ctx); + } else if (device_name == "NPU") { + // remote tensor is not used for NPU yet + // remote_context = ov_singleton_core().get_default_context(device_name); + } + + initialized = true; +} + +ggml_openvino_device_config::~ggml_openvino_device_config() { + if (cl_queue != nullptr) { + clReleaseCommandQueue(cl_queue); + cl_queue = nullptr; + } +} + +// Get the global device config singleton +ggml_openvino_device_config & ggml_openvino_get_device_config() { + static ggml_openvino_device_config config; + return config; +} + +// Initialize device config (call during backend init) +void ggml_openvino_init_device_config() { + ggml_openvino_get_device_config().init(); +} + +// Get the device name +const std::string & ggml_openvino_get_device_name() { + return ggml_openvino_get_device_config().device_name; +} + +// Check if running on NPU +bool ggml_openvino_is_npu() { + return ggml_openvino_get_device_config().is_npu; +} + +// Get the remote context for the current device (returns empty optional for CPU) +std::optional ggml_openvino_get_remote_context() { + return ggml_openvino_get_device_config().remote_context; +} + +// Get the compile config for the current device +const ov::AnyMap & ggml_openvino_get_compile_config() { + return ggml_openvino_get_device_config().compile_config; +} + +// Get the OpenCL command queue for GPU operations +cl_command_queue ggml_openvino_get_cl_queue() { + return ggml_openvino_get_device_config().cl_queue; +} + +// Get the clEnqueueMemFillINTEL function pointer (lazy load) +clEnqueueMemFillINTEL_fn ggml_openvino_get_clEnqueueMemFillINTEL() { + static clEnqueueMemFillINTEL_fn fn = nullptr; + static bool loaded = false; + if (!loaded) { + loaded = true; + cl_platform_id platform; + if (clGetPlatformIDs(1, &platform, nullptr) == CL_SUCCESS) { + fn = (clEnqueueMemFillINTEL_fn) clGetExtensionFunctionAddressForPlatform(platform, "clEnqueueMemFillINTEL"); + } + } + return fn; +} + +// Get the clEnqueueMemcpyINTEL function pointer (lazy load) +clEnqueueMemcpyINTEL_fn ggml_openvino_get_clEnqueueMemcpyINTEL() { + static clEnqueueMemcpyINTEL_fn fn = nullptr; + static bool loaded = false; + if (!loaded) { + loaded = true; + cl_platform_id platform; + if (clGetPlatformIDs(1, &platform, nullptr) == CL_SUCCESS) { + fn = (clEnqueueMemcpyINTEL_fn) clGetExtensionFunctionAddressForPlatform(platform, "clEnqueueMemcpyINTEL"); + } + } + return fn; +} + +// Get requantization type for a tensor type (returns nullopt if no requant needed) +std::optional ggml_openvino_get_requant_type(const ggml_tensor * tensor, bool no_requant) { + if (no_requant) { + return std::nullopt; + } + if (strncmp(tensor->name, "token_embd.weight", 17) == 0) { + return ((ggml_openvino_is_npu() && tensor->type == GGML_TYPE_Q6_K) ? ExtraQuantType::F16 : ExtraQuantType::Q8_0_C); + } + if (strncmp(tensor->name, "output.weight", 13) == 0) { + return ExtraQuantType::Q8_0_C; + } + if (ggml_openvino_is_npu()) { + return ExtraQuantType::Q4_0_128; + } + switch (tensor->type) { + case GGML_TYPE_Q6_K: + case GGML_TYPE_Q5_K: + return ExtraQuantType::Q8_0_C; + default: + return std::nullopt; + } +} + +// ===================================================== +// Extracted Layout Calculation +// ===================================================== + +ggml_openvino_extracted_layout ggml_openvino_get_extracted_layout(const ggml_tensor * tensor, bool use_bias) { + ggml_openvino_extracted_layout layout = {}; + layout.is_symmetric = false; + + if (!ggml_is_quantized(tensor->type)) { + return layout; + } + + // Only handle 2D weight tensors + if (tensor->ne[2] != 1 || tensor->ne[3] != 1) { + return layout; + } + + int64_t n_elements = ggml_nelements(tensor); + const size_t alignment = 64; // Good for SIMD + + // Check if requantization is needed (NPU-specific) + auto requant_type = ggml_openvino_get_requant_type(tensor, use_bias); + if (requant_type.has_value()) { + layout.is_requant = true; + layout.requant_type = requant_type; + + // Special case: requant to F16 - just store F16 weights, no scales/zp + if (requant_type.value() == ExtraQuantType::F16) { + layout.weights_size = n_elements * sizeof(uint16_t); // F16 = 2 bytes + layout.total_size = layout.weights_size; + layout.weights_offset = 0; + // No scales/zp for F16 + return layout; + } + + // Requant to different quantized format (e.g., Q4_0_128) + switch (requant_type.value()) { + case ExtraQuantType::Q4_0_128: + layout.is_u4 = true; + layout.weights_per_block = 128; + layout.is_symmetric = true; + break; + case ExtraQuantType::Q4_0_C: + layout.is_u4 = true; + layout.weights_per_block = tensor->ne[0]; + layout.is_symmetric = true; + break; + case ExtraQuantType::Q8_0_32: + layout.is_u4 = false; + layout.weights_per_block = 32; + layout.is_symmetric = true; + break; + case ExtraQuantType::Q8_0_C: + layout.is_u4 = false; + layout.weights_per_block = tensor->ne[0]; + layout.is_symmetric = true; + break; + case ExtraQuantType::Q8_1_C: + layout.is_u4 = false; + layout.weights_per_block = tensor->ne[0]; + break; + default: + layout.weights_per_block = -1; + GGML_ABORT("Code of re-quantizing to channel-wise is not updated"); + break; + } + + if (layout.is_requant) { + // Calculate sizes for requantized format + layout.weights_size = layout.is_u4 ? (n_elements / 2) : n_elements; + int64_t n_blocks = n_elements / layout.weights_per_block; + layout.scales_size = n_blocks * sizeof(uint16_t); + // For symmetric quantization, no zp needed (weights stored as signed) + if (layout.is_symmetric) { + layout.zp_size = 0; + } else { + layout.zp_size = layout.is_u4 ? ((n_blocks + 1) / 2) : n_blocks; + } + + layout.weights_offset = 0; + layout.scales_offset = ((layout.weights_size + alignment - 1) / alignment) * alignment; + layout.zp_offset = layout.scales_offset + ((layout.scales_size + alignment - 1) / alignment) * alignment; + layout.total_size = layout.zp_offset + layout.zp_size; + layout.total_size = std::max(layout.total_size, ggml_nbytes(tensor)); + return layout; + } + } + + // Normal extraction (no requant) - determine format based on tensor type + layout.is_u4 = false; + layout.weights_per_block = 32; + layout.is_symmetric = false; + + switch (tensor->type) { + case GGML_TYPE_Q4_0: + layout.is_u4 = true; + layout.is_symmetric = true; + break; + + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q4_K: + layout.is_u4 = true; + break; + + case GGML_TYPE_Q8_0: + layout.is_symmetric = true; + break; + + case GGML_TYPE_Q6_K: + layout.weights_per_block = 16; + layout.is_symmetric = true; + break; + + case GGML_TYPE_Q5_K: + break; + + default: + // Unsupported quantization type + return layout; + } + + // Calculate sizes + // Weights: U4 = n_elements/2 bytes, U8 = n_elements bytes + layout.weights_size = layout.is_u4 ? (n_elements / 2) : n_elements; + + // Scales: F16 per block + int64_t n_blocks = n_elements / layout.weights_per_block; + layout.scales_size = n_blocks * sizeof(uint16_t); // F16 = 2 bytes + // For symmetric quantization, no zp needed (weights stored as signed) + if (layout.is_symmetric) { + layout.zp_size = 0; + } else { + layout.zp_size = layout.is_u4 ? ((n_blocks + 1) / 2) : n_blocks; + } + + // Layout in buffer: [weights | scales | zp] with alignment + layout.weights_offset = 0; + layout.scales_offset = ((layout.weights_size + alignment - 1) / alignment) * alignment; + layout.zp_offset = layout.scales_offset + ((layout.scales_size + alignment - 1) / alignment) * alignment; + layout.total_size = layout.zp_offset + layout.zp_size; + layout.total_size = std::max(layout.total_size, ggml_nbytes(tensor)); + + return layout; +} + +ggml_openvino_tensor_extra * ggml_openvino_create_tensor_extra(const ggml_tensor * tensor, bool is_remote) { + ov::Shape shape; + for (int i = GGML_MAX_DIMS - 1; i >= 0; --i) { + shape.push_back(static_cast(tensor->ne[i])); + } + + ov::element::Type element_type; + switch (tensor->type) { + case GGML_TYPE_F32: + element_type = ov::element::f32; + break; + case GGML_TYPE_F16: + element_type = ov::element::f16; + break; + case GGML_TYPE_BF16: + element_type = ov::element::bf16; + break; + case GGML_TYPE_I32: + element_type = ov::element::i32; + break; + case GGML_TYPE_I64: + element_type = ov::element::i64; + break; + default: + // GGML_LOG_WARN("%s: unsupported tensor type for ov::Tensor: %s\n", __func__, ggml_type_name(tensor->type)); + return nullptr; + } + + const auto & device_name = ggml_openvino_get_device_name(); + auto remote_context = ggml_openvino_get_remote_context(); + + std::shared_ptr ov_tensor; + if (is_remote) { + GGML_ASSERT(device_name == "GPU"); + auto gpu_context = remote_context->as(); + auto usm_tensor = gpu_context.create_tensor(element_type, shape, tensor->data); + ov_tensor = std::make_shared(std::move(usm_tensor)); + } else { + ov_tensor = std::make_shared(element_type, shape, tensor->data); + } + + return new ggml_openvino_tensor_extra(ov_tensor); +} + + +/* Inlined ggml-quants.cpp */ + +#include "ggml-common-defs.inc" +#include "ggml-impl-defs.inc" +#include "ggml.h.inc" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +void unpack_32_4(const uint8_t * data, uint8_t * dst) { + std::fill_n(dst, 16, 0); + for (int j = 0; j < 16; ++j) { + uint8_t x = (data[j] & 0x0F); + uint8_t y = (data[j] >> 4); + if (j % 2 != 0) { + x <<= 4; + y <<= 4; + } + dst[j / 2] |= x; + dst[8 + j / 2] |= y; // Last 16 weights are in the higher bits + } +} + +// Extracts (weight, scales, zp) from Q4_0 tensors. +// Data layout is: |16 bit scale|32 x 4bit weights|. +// When zp_arr is empty (symmetric), weights are stored as signed i4 (value - 8). +void extract_q4_0_data(const ggml_tensor * tensor, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr) { + const uint64_t bytes_per_block = 18; // 2 bytes scale, 32x0.5 byte weights + + auto * data = static_cast(tensor->data); + auto * weights = static_cast(weights_arr.data()); + auto * scales = scales_arr.data::value_type>(); + + bool is_symmetric = (weights_arr.get_element_type() == ov::element::i4); // Signed i4 path + + if (!is_symmetric) { + auto * zp = static_cast(zp_arr.data()); + ov::parallel_for(scales_arr.get_size(), [&](size_t i) { + scales[i] = ov::float16::from_bits(*((uint16_t *) (data + i * bytes_per_block))); + // Pack two 4-bit zero points per byte + if (i % 2 == 0) { + zp[i / 2] = 8; // Lower nibble + } else { + zp[i / 2] |= (8 << 4); // Upper nibble + } + unpack_32_4(data + i * bytes_per_block + 2, weights + i * 16); + }); + } else { + // Symmetric: unpack as u4 then convert to i4 by subtracting 8 (XOR each nibble) + ov::parallel_for(scales_arr.get_size(), [&](size_t i) { + scales[i] = ov::float16::from_bits(*((uint16_t *) (data + i * bytes_per_block))); + unpack_32_4(data + i * bytes_per_block + 2, weights + i * 16); + // Convert u4 to i4: subtract 8 from each nibble. XOR 0x88 flips each nibble by 8. + for (int j = 0; j < 16; ++j) { + weights[i * 16 + j] ^= 0x88; + } + }); + } +} + +// Extracts (weight, scales, zp) from Q4_1 tensors. +// Data layout is: |16 bit scale|16 bit min|32 x 4bit weights|. +void extract_q4_1_data(const ggml_tensor * tensor, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr, + bool use_bias) { + const uint64_t bytes_per_block = 20; // 2 bytes scale, 2 bytes min, 32x0.5 byte weights + + auto * data = static_cast(tensor->data); + auto * weights = static_cast(weights_arr.data()); + auto * scales = scales_arr.data::value_type>(); + + if (use_bias) { + // Store bias (min) directly as f16 instead of computing u4 zero points + auto * bias = zp_arr.data::value_type>(); + ov::parallel_for(scales_arr.get_size(), [&](size_t i) { + float scale = static_cast(ov::float16::from_bits(*((uint16_t *) (data + i * bytes_per_block)))); + float min = static_cast(ov::float16::from_bits(*((uint16_t *) (data + i * bytes_per_block + 2)))); + scales[i] = ov::float16(scale); + bias[i] = ov::float16(min); // bias = min, dequant: w*s + bias + unpack_32_4(data + i * bytes_per_block + 4, weights + i * 16); + }); + } else { + auto * zp = static_cast(zp_arr.data()); + ov::parallel_for(scales_arr.get_size(), [&](size_t i) { + float scale = static_cast(ov::float16::from_bits(*((uint16_t *) (data + i * bytes_per_block)))); + float min = static_cast(ov::float16::from_bits(*((uint16_t *) (data + i * bytes_per_block + 2)))); + scales[i] = ov::float16(scale); + // zp = -min / scale (bias = min, so zp = -bias/scale) + uint8_t zp_val = (scale != 0.0f) ? (uint8_t) std::round(-min / scale) : 0; + // Pack two 4-bit zero points per byte + if (i % 2 == 0) { + zp[i / 2] = zp_val & 0x0F; // Lower nibble + } else { + zp[i / 2] |= (zp_val << 4); // Upper nibble + } + unpack_32_4(data + i * bytes_per_block + 4, weights + i * 16); + }); + } +} + +// Extracts (weight, scales, zp) from Q8_0 tensors. +// Data layout is: |16 bit scale|32 x 8bit weights|. +// When zp_arr is empty (symmetric), weights are stored as signed i8 directly. +void extract_q8_0_data(const ggml_tensor * tensor, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr) { + const uint64_t weights_per_block = 32; + const uint64_t bytes_per_block = 34; // 2 bytes scale, 32x1 byte weights + + auto * data = static_cast(tensor->data); + auto * weights = static_cast(weights_arr.data()); + auto * scales = scales_arr.data::value_type>(); + + bool is_symmetric = (weights_arr.get_element_type() == ov::element::i8); // Signed i8 path + + if (!is_symmetric) { + auto * zp = static_cast(zp_arr.data()); + ov::parallel_for(scales_arr.get_size(), [&](size_t i) { + uint8_t * block_data = data + i * bytes_per_block; + scales[i] = ov::float16::from_bits(*(uint16_t *) block_data); + zp[i] = 128; + for (size_t j = 0; j < weights_per_block; ++j) { + uint8_t x = block_data[j + 2]; + x ^= 1 << 7; // Convert int8 to uint8 by flipping sign bit + weights[i * weights_per_block + j] = x; + } + }); + } else { + // Symmetric: store original int8 values directly (no unsigned bias) + ov::parallel_for(scales_arr.get_size(), [&](size_t i) { + uint8_t * block_data = data + i * bytes_per_block; + scales[i] = ov::float16::from_bits(*(uint16_t *) block_data); + // Copy int8 weights as-is (the tensor element type is i8) + memcpy(weights + i * weights_per_block, block_data + 2, weights_per_block); + }); + } +} + +void unpack_256_4(const uint8_t * data, uint8_t * dst) { + // Initialize the output array with zeros + std::fill_n(dst, 128, 0); + + for (size_t i = 0; i < 4; ++i) { + for (int j = 0; j < 32; ++j) { + uint8_t x = (data[i * 32 + j] & 0x0F); + uint8_t y = (data[i * 32 + j] >> 4); + if (j % 2 != 0) { + x <<= 4; + y <<= 4; + } + dst[i * 32 + j / 2] |= x; + dst[i * 32 + 16 + j / 2] |= y; // Last 16 weights are in the higher bits + } + } +} + +void extract_q4_k_data(const ggml_tensor * tensor, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr, + bool use_bias) { + const uint64_t bytes_per_block = 2 + 2 + 12 + 128; + const uint64_t n_super_block = tensor->nb[3] / bytes_per_block; + + auto * data = static_cast(tensor->data); + auto * weights = static_cast(weights_arr.data()); + auto * scales = scales_arr.data::value_type>(); + + // For bias path, zp_arr holds f16 bias values; for zp path, it holds packed u4 zero points + auto * zp_u4 = use_bias ? nullptr : static_cast(zp_arr.data()); + auto * bias_f16 = use_bias ? zp_arr.data::value_type>() : nullptr; + + ov::parallel_for(n_super_block, [&](size_t i) { + uint8_t * block_data = data + i * bytes_per_block; + + // Extract scale factors and offsets + float scale_scales = static_cast(ov::float16::from_bits(*((uint16_t *) block_data))); + float scale_mins = static_cast(ov::float16::from_bits(*((uint16_t *) block_data + 1))); + + // Extract qs1 and qs2 + uint8_t * qs1 = block_data + 4; + + // Calculate scales + float scale_vals[8]; + scale_vals[0] = scale_scales * static_cast((*(qs1) & 0b111111)); + scale_vals[1] = scale_scales * static_cast((*(qs1 + 1) & 0b111111)); + scale_vals[2] = scale_scales * static_cast((*(qs1 + 2) & 0b111111)); + scale_vals[3] = scale_scales * static_cast((*(qs1 + 3) & 0b111111)); + scale_vals[4] = scale_scales * static_cast((*(qs1 + 8) & 0b00001111) | ((*(qs1) >> 6) << 4)); + scale_vals[5] = scale_scales * static_cast((*(qs1 + 9) & 0b00001111) | ((*(qs1 + 1) >> 6) << 4)); + scale_vals[6] = scale_scales * static_cast((*(qs1 + 10) & 0b00001111) | ((*(qs1 + 2) >> 6) << 4)); + scale_vals[7] = scale_scales * static_cast((*(qs1 + 11) & 0b00001111) | ((*(qs1 + 3) >> 6) << 4)); + + // Calculate min values (bias = -min) + float min_vals[8]; + min_vals[0] = scale_mins * static_cast((*(qs1 + 4) & 0b111111)); + min_vals[1] = scale_mins * static_cast((*(qs1 + 5) & 0b111111)); + min_vals[2] = scale_mins * static_cast((*(qs1 + 6) & 0b111111)); + min_vals[3] = scale_mins * static_cast((*(qs1 + 7) & 0b111111)); + min_vals[4] = scale_mins * static_cast((*(qs1 + 8) >> 4) | ((*(qs1 + 4) >> 6) << 4)); + min_vals[5] = scale_mins * static_cast((*(qs1 + 9) >> 4) | ((*(qs1 + 5) >> 6) << 4)); + min_vals[6] = scale_mins * static_cast((*(qs1 + 10) >> 4) | ((*(qs1 + 6) >> 6) << 4)); + min_vals[7] = scale_mins * static_cast((*(qs1 + 11) >> 4) | ((*(qs1 + 7) >> 6) << 4)); + + // Store scales and compute zero points or bias + for (int j = 0; j < 8; j++) { + scales[i * 8 + j] = ov::float16(scale_vals[j]); + if (use_bias) { + // Store bias = -min directly as f16, dequant: w*s + bias + bias_f16[i * 8 + j] = ov::float16(-min_vals[j]); + } else { + // zp = min / scale (since bias = -min and zp = -bias/scale) + uint8_t zp_val = (scale_vals[j] != 0.0f) ? (uint8_t) std::round(min_vals[j] / scale_vals[j]) : 0; + // Pack two 4-bit zero points per byte + size_t idx = i * 8 + j; + if (idx % 2 == 0) { + zp_u4[idx / 2] = zp_val & 0x0F; + } else { + zp_u4[idx / 2] |= (zp_val << 4); + } + } + } + unpack_256_4(block_data + 16, weights + i * 128); + }); +} + +void extract_q6_k_data(const ggml_tensor * tensor, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr) { + const uint64_t bytes_per_block = 128 + 64 + 16 + 2; + const uint64_t n_super_block = tensor->nb[3] / bytes_per_block; + + auto * data = static_cast(tensor->data); + auto * weights = static_cast(weights_arr.data()); + auto * scales = scales_arr.data::value_type>(); + + bool is_symmetric = (weights_arr.get_element_type() == ov::element::i8); // Signed i8 path + + if (!is_symmetric) { + auto * zp = static_cast(zp_arr.data()); + ov::parallel_for(n_super_block, [&](size_t i) { + uint8_t * block_data = data + i * bytes_per_block; + float scale_factor = static_cast(ov::float16::from_bits(*((uint16_t *) block_data + 104))); + for (size_t j = 0; j < 16; j++) { + scales[j + i * 16] = + ov::float16(scale_factor * static_cast(*((int8_t *) (block_data + 128 + 64 + j)))); + zp[j + i * 16] = 32; + } + uint8_t * ql = block_data; + uint8_t * qh = block_data + 128; + for (int64_t j = 0; j < 32; ++j) { + weights[i * 256 + j] = (ql[j] & 0xF) | (((qh[j] >> 0) & 3) << 4); + weights[i * 256 + j + 32] = (ql[32 + j] & 0xF) | (((qh[j] >> 2) & 3) << 4); + weights[i * 256 + j + 64] = (ql[j] >> 4) | (((qh[j] >> 4) & 3) << 4); + weights[i * 256 + j + 96] = (ql[32 + j] >> 4) | (((qh[j] >> 6) & 3) << 4); + weights[i * 256 + j + 128] = (ql[64 + j] & 0xF) | (((qh[32 + j] >> 0) & 3) << 4); + weights[i * 256 + j + 160] = (ql[96 + j] & 0xF) | (((qh[32 + j] >> 2) & 3) << 4); + weights[i * 256 + j + 192] = (ql[64 + j] >> 4) | (((qh[32 + j] >> 4) & 3) << 4); + weights[i * 256 + j + 224] = (ql[96 + j] >> 4) | (((qh[32 + j] >> 6) & 3) << 4); + } + }); + } else { + // Symmetric: subtract 32 from each weight to store as signed i8 + ov::parallel_for(n_super_block, [&](size_t i) { + uint8_t * block_data = data + i * bytes_per_block; + float scale_factor = static_cast(ov::float16::from_bits(*((uint16_t *) block_data + 104))); + for (size_t j = 0; j < 16; j++) { + scales[j + i * 16] = + ov::float16(scale_factor * static_cast(*((int8_t *) (block_data + 128 + 64 + j)))); + } + uint8_t * ql = block_data; + uint8_t * qh = block_data + 128; + auto * signed_weights = reinterpret_cast(weights); + for (int64_t j = 0; j < 32; ++j) { + signed_weights[i * 256 + j] = static_cast((ql[j] & 0xF) | (((qh[j] >> 0) & 3) << 4)) - 32; + signed_weights[i * 256 + j + 32] = + static_cast((ql[32 + j] & 0xF) | (((qh[j] >> 2) & 3) << 4)) - 32; + signed_weights[i * 256 + j + 64] = static_cast((ql[j] >> 4) | (((qh[j] >> 4) & 3) << 4)) - 32; + signed_weights[i * 256 + j + 96] = + static_cast((ql[32 + j] >> 4) | (((qh[j] >> 6) & 3) << 4)) - 32; + signed_weights[i * 256 + j + 128] = + static_cast((ql[64 + j] & 0xF) | (((qh[32 + j] >> 0) & 3) << 4)) - 32; + signed_weights[i * 256 + j + 160] = + static_cast((ql[96 + j] & 0xF) | (((qh[32 + j] >> 2) & 3) << 4)) - 32; + signed_weights[i * 256 + j + 192] = + static_cast((ql[64 + j] >> 4) | (((qh[32 + j] >> 4) & 3) << 4)) - 32; + signed_weights[i * 256 + j + 224] = + static_cast((ql[96 + j] >> 4) | (((qh[32 + j] >> 6) & 3) << 4)) - 32; + } + }); + } +} + +static 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); + } +} + +void extract_q5_k_data(const ggml_tensor * tensor, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr, + bool use_bias) { + const uint64_t bytes_per_block = 4 + 12 + 32 + 128; + const uint64_t n_super_block = tensor->nb[3] / bytes_per_block; + + auto * data = static_cast(tensor->data); + auto * weights = static_cast(weights_arr.data()); + auto * scales = scales_arr.data::value_type>(); + + // For bias path, zp_arr holds f16 bias values; for zp path, it holds u8 zero points + auto * zp_u8 = use_bias ? nullptr : static_cast(zp_arr.data()); + auto * bias_f16 = use_bias ? zp_arr.data::value_type>() : nullptr; + + ov::parallel_for(n_super_block, [&](size_t i) { + uint8_t * block_data = data + i * bytes_per_block; + + const float d = static_cast(ov::float16::from_bits(*((uint16_t *) block_data))); + const float min_factor = static_cast(ov::float16::from_bits(*((uint16_t *) block_data + 1))); + + const uint8_t * scales_data = block_data + 4; // 12 bytes of scales + const uint8_t * qh = block_data + 4 + 12; // 32 bytes of high bits + const uint8_t * ql = block_data + 4 + 12 + 32; // 128 bytes of low bits + + int is = 0; + uint8_t u1 = 1; + uint8_t u2 = 2; + + // Process 2 blocks in one iteration + for (int j = 0; j < 256; j += 64) { // 256 = QK_K, so 4 iterations of 64 + uint8_t sc; + uint8_t m; + + // Get scale and min for first 32 elements + get_scale_min_k4(is + 0, scales_data, &sc, &m); + const float d1 = d * sc; + const float m1 = min_factor * m; + + // Get scale and min for second 32 elements + get_scale_min_k4(is + 1, scales_data, &sc, &m); + const float d2 = d * sc; + const float m2 = min_factor * m; + + scales[i * 8 + is] = ov::float16(d1); + scales[i * 8 + is + 1] = ov::float16(d2); + if (use_bias) { + // Store bias = -min directly as f16, dequant: w*s + bias + bias_f16[i * 8 + is] = ov::float16(-m1); + bias_f16[i * 8 + is + 1] = ov::float16(-m2); + } else { + // zp = min / scale (since bias = -min and zp = -bias/scale) + zp_u8[i * 8 + is] = (d1 != 0.0f) ? (uint8_t) std::round(m1 / d1) : 0; + zp_u8[i * 8 + is + 1] = (d2 != 0.0f) ? (uint8_t) std::round(m2 / d2) : 0; + } + + // Extract weights for first 32 elements (matching deq formula exactly) + for (int l = 0; l < 32; ++l) { + weights[i * 256 + j + l] = (ql[l] & 0xF) + ((qh[l] & u1) ? 16 : 0); + } + + // Extract weights for second 32 elements + for (int l = 0; l < 32; ++l) { + weights[i * 256 + j + l + 32] = (ql[l] >> 4) + ((qh[l] & u2) ? 16 : 0); + } + + ql += 32; + is += 2; + u1 <<= 2; + u2 <<= 2; + } + }); +} + +// TODO Reorder for make_intX_weights + +ov::Output make_int8_weights(ov::Tensor & weight, + ov::Tensor & scales, + ov::Tensor & zp, + size_t group_size, + bool use_bias) { + ov::Shape orig_shape = weight.get_shape(); + bool is_signed = (weight.get_element_type() == ov::element::i8); // Symmetric: signed weights, no ZP + + // Expand dimensions for scales and zp/bias + auto scale_shape = scales.get_shape(); + + ov::Shape packed_shape = {orig_shape[0], orig_shape[1] / group_size, group_size}; + + if (packed_shape[1] == 1) { + // Requantized channel-wise case + packed_shape.erase(packed_shape.begin() + 1); + } else { + scale_shape.push_back(1); + scales.set_shape(scale_shape); + if (!is_signed && zp.get_size() > 0) { + auto zp_shape = zp.get_shape(); + zp_shape.push_back(1); + zp.set_shape(zp_shape); + } + } + + auto scales_f16 = std::make_shared(scales); + + ov::Output result; + if (is_signed) { + // Signed path: q * s (no zero point subtraction needed) + auto weights_node = std::make_shared(ov::element::i8, packed_shape, + static_cast(weight.data()), nullptr); + weights_node->get_rt_info()["__gguf_tensor_holder"] = weight; + auto weights_f16 = std::make_shared(weights_node, ov::element::f16); + result = std::make_shared(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY); + } else { + // Unsigned path + auto weights_node = std::make_shared(ov::element::u8, packed_shape, + static_cast(weight.data()), nullptr); + weights_node->get_rt_info()["__gguf_tensor_holder"] = weight; + auto weights_f16 = std::make_shared(weights_node, ov::element::f16); + + if (use_bias && zp.get_size() > 0) { + // Bias path: w * s + b (zp tensor holds f16 bias values) + auto bias_f16 = std::make_shared(zp); + auto w_s = + std::make_shared(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY); + result = std::make_shared(w_s, bias_f16, ov::op::AutoBroadcastType::NUMPY); + } else { + // Zero point path: (w - zp) * s + auto zero_point = std::make_shared(zp); + float zp_value; + if (ov::op::util::get_single_value(zero_point, zp_value)) { + zero_point = ov::op::v0::Constant::create(zero_point->get_element_type(), {}, {zp_value}); + } + auto zero_point_f16 = std::make_shared(zero_point, ov::element::f16); + auto w_zp = + std::make_shared(weights_f16, zero_point_f16, ov::op::AutoBroadcastType::NUMPY); + result = std::make_shared(w_zp, scales_f16, ov::op::AutoBroadcastType::NUMPY); + } + } + + if (packed_shape.size() != 2) { + // If not requantized channel-wise case, reshape back to original shape + auto final_shape = + std::make_shared(ov::element::i64, ov::Shape{orig_shape.size()}, orig_shape); + result = std::make_shared(result, final_shape, false); + } + + return std::make_shared(result, ov::element::f32); +} + +ov::Output make_int4_weights(ov::Tensor & weight, + ov::Tensor & scales, + ov::Tensor & zp, + size_t group_size, + bool use_bias) { + ov::Shape orig_weight_shape = weight.get_shape(); + bool is_signed = (weight.get_element_type() == ov::element::i4); // Symmetric: signed weights, no ZP + + // Expand dimensions for scales and zp/bias + ov::Shape scale_shape = scales.get_shape(); + + // Create INT4 weight tensor + ov::Shape packed_shape = {orig_weight_shape[0], orig_weight_shape[1] / group_size, group_size}; + + if (packed_shape[1] == 1) { + // Requantized channel-wise case + packed_shape.erase(packed_shape.begin() + 1); + } else { + scale_shape.push_back(1); + scales.set_shape(scale_shape); + if (!is_signed && zp.get_size() > 0) { + auto zp_shape = zp.get_shape(); + zp_shape.push_back(1); + zp.set_shape(zp_shape); + } + } + + auto scales_f16 = std::make_shared(scales); + + ov::Output result; + if (is_signed) { + // Signed path: q * s (no zero point subtraction needed) + auto weights_node = std::make_shared(ov::element::i4, packed_shape, + static_cast(weight.data()), nullptr); + weights_node->get_rt_info()["__gguf_tensor_holder"] = weight; + auto weights_f16 = std::make_shared(weights_node, ov::element::f16); + result = std::make_shared(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY); + } else { + // Unsigned path + auto weights_node = std::make_shared(ov::element::u4, packed_shape, + static_cast(weight.data()), nullptr); + weights_node->get_rt_info()["__gguf_tensor_holder"] = weight; + auto weights_f16 = std::make_shared(weights_node, ov::element::f16); + + if (use_bias && zp.get_size() > 0) { + // Bias path: w * s + b (zp tensor holds f16 bias values) + auto bias_f16 = std::make_shared(zp); + auto w_s = + std::make_shared(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY); + result = std::make_shared(w_s, bias_f16, ov::op::AutoBroadcastType::NUMPY); + } else { + // Zero point path: (w - zp) * s + auto zero_points_node = std::make_shared(zp); + float zp_value; + if (ov::op::util::get_single_value(zero_points_node, zp_value)) { + zero_points_node = ov::op::v0::Constant::create(zero_points_node->get_element_type(), {}, {zp_value}); + } + auto zero_points_f16 = std::make_shared(zero_points_node, ov::element::f16); + auto w_zp = + std::make_shared(weights_f16, zero_points_f16, ov::op::AutoBroadcastType::NUMPY); + result = std::make_shared(w_zp, scales_f16, ov::op::AutoBroadcastType::NUMPY); + } + } + + if (packed_shape.size() != 2) { + // If not requantized channel-wise case, reshape back to original shape + auto final_shape = std::make_shared(ov::element::i64, ov::Shape{orig_weight_shape.size()}, + orig_weight_shape); + result = std::make_shared(result, final_shape, false); + } + + return std::make_shared(result, ov::element::f32); +} + +// Extract quantized weights from tensor and create weight subgraph +std::shared_ptr extract_quantized_weights(const ggml_tensor * tensor, + const void * data, + ov::Tensor & weights, + ov::Tensor & scales, + ov::Tensor & zp, + bool use_bias) { + // Create a temporary tensor for extraction functions that read from tensor->data + ggml_tensor temp_tensor = *tensor; + temp_tensor.data = const_cast(data); + + // Determine block size based on tensor type + int64_t weights_per_block; + bool is_u4; + switch (tensor->type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q4_K: + is_u4 = true; + weights_per_block = 32; + break; + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q5_K: + is_u4 = false; + weights_per_block = 32; + break; + case GGML_TYPE_Q6_K: + is_u4 = false; + weights_per_block = 16; + break; + default: + throw std::runtime_error("Unsupported quantized type for extraction: " + + std::string(ggml_type_name(tensor->type))); + } + + // Extract quantized data + switch (tensor->type) { + case GGML_TYPE_Q4_0: + extract_q4_0_data(&temp_tensor, weights, scales, zp); + break; + case GGML_TYPE_Q4_1: + extract_q4_1_data(&temp_tensor, weights, scales, zp, use_bias); + break; + case GGML_TYPE_Q4_K: + extract_q4_k_data(&temp_tensor, weights, scales, zp, use_bias); + break; + case GGML_TYPE_Q8_0: + extract_q8_0_data(&temp_tensor, weights, scales, zp); + break; + case GGML_TYPE_Q6_K: + extract_q6_k_data(&temp_tensor, weights, scales, zp); + break; + case GGML_TYPE_Q5_K: + extract_q5_k_data(&temp_tensor, weights, scales, zp, use_bias); + break; + default: + throw std::runtime_error("Unsupported quantized type: " + std::string(ggml_type_name(tensor->type))); + } + + // Create the OpenVINO weight subgraph + ov::Output weight_node; + if (is_u4) { + weight_node = make_int4_weights(weights, scales, zp, weights_per_block, use_bias); + } else { + weight_node = make_int8_weights(weights, scales, zp, weights_per_block, use_bias); + } + + auto result = weight_node.get_node_shared_ptr(); + result->set_friendly_name(tensor->name); + return result; +} + +// Requantize weights to target format, writing to provided buffers +std::shared_ptr requantize_to_buffers(const ggml_tensor * tensor, + const void * data, + ExtraQuantType requant_type, + int64_t block_size, + ov::Tensor & weights, + ov::Tensor & scales, + ov::Tensor & zp) { + int64_t n_elements = ggml_nelements(tensor); + + // First dequantize to F32 + std::vector weights_f32(n_elements); + ggml_get_type_traits(tensor->type)->to_float(data, weights_f32.data(), n_elements); + + // Handle F16 case - just convert and create constant + if (requant_type == ExtraQuantType::F16) { + ggml_get_type_traits(GGML_TYPE_F16)->from_float_ref(weights_f32.data(), weights.data(), n_elements); + auto result = std::make_shared(weights); + result->set_friendly_name(tensor->name); + return result; + } + + // Requantize to target quantized format + bool is_u4 = (requant_type == ExtraQuantType::Q4_0_C || requant_type == ExtraQuantType::Q4_0_128); + + if (is_u4) { + quantize_q4_0(weights_f32.data(), weights, scales, zp, n_elements, block_size); + } else if (requant_type == ExtraQuantType::Q8_1_C) { + quantize_q8_1(weights_f32.data(), weights, scales, zp, n_elements, block_size); + } else { + quantize_q8_0(weights_f32.data(), weights, scales, zp, n_elements, block_size); + } + + // Create the OpenVINO weight subgraph + ov::Output weight_node; + if (is_u4) { + weight_node = make_int4_weights(weights, scales, zp, block_size); + } else { + weight_node = make_int8_weights(weights, scales, zp, block_size); + } + + auto result = weight_node.get_node_shared_ptr(); + result->set_friendly_name(tensor->name); + return result; +} + +OvWeight process_weight_tensor(const ggml_tensor * tensor, const void * data, void * output_base_ptr, bool use_bias) { + GGML_ASSERT(tensor != nullptr); + GGML_ASSERT(data != nullptr); + + OvWeight result; + + // Get 2D shape for weights [rows, cols] + ov::Shape node_shape = {static_cast(tensor->ne[1]), static_cast(tensor->ne[0])}; + + // Handle F16/F32/BF16 weights + if (tensor->type == GGML_TYPE_F32 || tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_BF16) { + ov::element::Type element_type; + switch (tensor->type) { + case GGML_TYPE_F32: + element_type = ov::element::f32; + break; + case GGML_TYPE_F16: + element_type = ov::element::f16; + break; + case GGML_TYPE_BF16: + element_type = ov::element::bf16; + break; + default: + OPENVINO_THROW("Unexpected tensor type in F16/F32/BF16 path"); + } + + if (output_base_ptr && output_base_ptr != data) { + // Using external buffer - copy data and create shared-memory constant + size_t tensor_bytes = ggml_nbytes(tensor); + memcpy(output_base_ptr, data, tensor_bytes); + result.weights = ov::Tensor(element_type, node_shape, output_base_ptr); + } else { + result.weights = ov::Tensor(element_type, node_shape, data); + } + result.weight_node = std::make_shared(result.weights); + return result; + } + + // Handle quantized weights + if (!ggml_is_quantized(tensor->type)) { + OPENVINO_THROW("Unsupported weight tensor type: ", ggml_type_name(tensor->type)); + } + + result.layout = ggml_openvino_get_extracted_layout(tensor, use_bias); + const auto & layout = result.layout; + if (layout.total_size == 0) { + OPENVINO_THROW("Unsupported quantized type: ", ggml_type_name(tensor->type)); + } + + if (use_bias) { + OPENVINO_ASSERT(!layout.is_requant, + "use_bias is only used for test-backend-ops, which should not have requantization"); + // bias node will be created on the fly and not use backend buffer + output_base_ptr = nullptr; + } + + // F16 requant path - no separate scales/zp needed in result + if (layout.is_requant && layout.requant_type.has_value() && layout.requant_type.value() == ExtraQuantType::F16) { + if (output_base_ptr) { + result.weights = ov::Tensor(ov::element::f16, node_shape, + static_cast(output_base_ptr) + layout.weights_offset); + } else { + result.weights = ov::Tensor(ov::element::f16, node_shape); + } + ov::Tensor dummy_scales, dummy_zp; // Not used for F16 + result.weight_node = + requantize_to_buffers(tensor, data, ExtraQuantType::F16, 0, result.weights, dummy_scales, dummy_zp); + return result; + } + + // Quantized path (normal extraction or quantized requant) + // Create weight/scale/zp tensors - shared between both paths + // For symmetric quantization, use signed types (i4/i8) and no ZP tensor + ov::element::Type weight_type = layout.is_symmetric ? (layout.is_u4 ? ov::element::i4 : ov::element::i8) : + (layout.is_u4 ? ov::element::u4 : ov::element::u8); + ov::Shape scale_shape = {node_shape[0], node_shape[1] / layout.weights_per_block}; + + if (output_base_ptr) { + uint8_t * buf_base = static_cast(output_base_ptr); + result.weights = ov::Tensor(weight_type, node_shape, buf_base + layout.weights_offset); + result.scales = ov::Tensor(ov::element::f16, scale_shape, buf_base + layout.scales_offset); + if (!layout.is_symmetric) { + ov::element::Type zp_type = layout.is_u4 ? ov::element::u4 : ov::element::u8; + result.zp = ov::Tensor(zp_type, scale_shape, buf_base + layout.zp_offset); + } + // else: result.zp remains default-constructed (empty) for symmetric + } else { + result.weights = ov::Tensor(weight_type, node_shape); + result.scales = ov::Tensor(ov::element::f16, scale_shape); + if (!layout.is_symmetric) { + if (use_bias) { + result.zp = ov::Tensor(ov::element::f16, scale_shape); + } else { + ov::element::Type zp_type = layout.is_u4 ? ov::element::u4 : ov::element::u8; + result.zp = ov::Tensor(zp_type, scale_shape); + } + } + // else: result.zp remains default-constructed (empty) for symmetric + } + + if (layout.is_requant && layout.requant_type.has_value()) { + result.weight_node = requantize_to_buffers(tensor, data, layout.requant_type.value(), layout.weights_per_block, + result.weights, result.scales, result.zp); + } else { + result.weight_node = + extract_quantized_weights(tensor, data, result.weights, result.scales, result.zp, use_bias); + } + + return result; +} + +void quantize_q4_0(const float * x, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr, + int64_t k, + int64_t qk) { + assert(k % qk == 0); + const int nb = k / qk; + + auto * weights = static_cast(weights_arr.data()); + auto * scales = scales_arr.data::value_type>(); + bool is_symmetric = (weights_arr.get_element_type() == ov::element::i4); // Signed i4 path + + if (!is_symmetric) { + auto * zp = static_cast(zp_arr.data()); + for (int i = 0; i < nb; i++) { + float amax = 0.0f; + float max = 0.0f; + for (int j = 0; j < qk; j++) { + const float v = x[i * qk + j]; + if (amax < fabsf(v)) { + amax = fabsf(v); + max = v; + } + } + const float d = max / -8; + if (d == 0) { + scales[i] = ov::float16(1.0f); + if (i % 2 == 0) { + zp[i / 2] = 8; + } else { + zp[i / 2] |= (8 << 4); + } + memset(weights + i * qk / 2, 8 | (8 << 4), qk / 2); + continue; + } + const float id = 1.0f / d; + scales[i] = ov::float16(d); + if (i % 2 == 0) { + zp[i / 2] = 8; + } else { + zp[i / 2] |= (8 << 4); + } + for (int j = 0; j < qk / 2; ++j) { + const float x0 = x[i * qk + 2 * j] * id; + const float x1 = x[i * qk + 2 * j + 1] * id; + const uint8_t xi0 = MIN(15, (int8_t) (x0 + 8.5f)); + const uint8_t xi1 = MIN(15, (int8_t) (x1 + 8.5f)); + weights[i * qk / 2 + j] = xi0 | (xi1 << 4); + } + } + } else { + // Symmetric: produce signed i4 values in [-8, 7] + for (int i = 0; i < nb; i++) { + float amax = 0.0f; + float max = 0.0f; + for (int j = 0; j < qk; j++) { + const float v = x[i * qk + j]; + if (amax < fabsf(v)) { + amax = fabsf(v); + max = v; + } + } + const float d = max / -8; + if (d == 0) { + scales[i] = ov::float16(1.0f); + // i4 value 0 packed: 0x00 + memset(weights + i * qk / 2, 0, qk / 2); + continue; + } + const float id = 1.0f / d; + scales[i] = ov::float16(d); + for (int j = 0; j < qk / 2; ++j) { + const float x0 = x[i * qk + 2 * j] * id; + const float x1 = x[i * qk + 2 * j + 1] * id; + // Signed i4: range [-8, 7]. Quantize as round(x*id), then pack as 4-bit two's complement. + int8_t si0 = (int8_t) std::max(-8, std::min(7, (int) roundf(x0))); + int8_t si1 = (int8_t) std::max(-8, std::min(7, (int) roundf(x1))); + weights[i * qk / 2 + j] = (si0 & 0x0F) | ((si1 & 0x0F) << 4); + } + } + } +} + +void quantize_q8_0(const float * x, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr, + int64_t k, + int64_t qk) { + assert(k % qk == 0); + const int nb = k / qk; + + auto * weights = static_cast(weights_arr.data()); + auto * scales = scales_arr.data::value_type>(); + bool is_symmetric = (weights_arr.get_element_type() == ov::element::i8); // Signed i8 path + + if (!is_symmetric) { + auto * zp = static_cast(zp_arr.data()); + for (int i = 0; i < nb; i++) { + float amax = 0.0f; + for (int j = 0; j < qk; j++) { + const float v = x[i * qk + j]; + amax = std::max(amax, fabsf(v)); + } + const float d = amax / 127.0f; + const float id = d ? 1.0f / d : 0.0f; + scales[i] = ov::float16(d); + zp[i] = 128; + for (int j = 0; j < qk; ++j) { + const float x0 = x[i * qk + j] * id; + const int8_t xi0 = roundf(x0); + weights[i * qk + j] = (uint8_t) (xi0 + 128); + } + } + } else { + // Symmetric: store signed int8 values directly + auto * signed_weights = reinterpret_cast(weights); + for (int i = 0; i < nb; i++) { + float amax = 0.0f; + for (int j = 0; j < qk; j++) { + const float v = x[i * qk + j]; + amax = std::max(amax, fabsf(v)); + } + const float d = amax / 127.0f; + const float id = d ? 1.0f / d : 0.0f; + scales[i] = ov::float16(d); + for (int j = 0; j < qk; ++j) { + const float x0 = x[i * qk + j] * id; + signed_weights[i * qk + j] = (int8_t) roundf(x0); + } + } + } +} + +void quantize_q8_1(const float * x, + ov::Tensor & weights_arr, + ov::Tensor & scales_arr, + ov::Tensor & zp_arr, + int64_t k, + int64_t qk) { + assert(k % qk == 0); + const int nb = k / qk; + + auto * weights = static_cast(weights_arr.data()); + auto * scales = scales_arr.data::value_type>(); + auto * zp = static_cast(zp_arr.data()); + for (int i = 0; i < nb; i++) { + float min = std::numeric_limits::max(); + float max = std::numeric_limits::lowest(); + + for (int j = 0; j < qk; j++) { + const float v = x[i * qk + j]; + min = std::min(v, min); + max = std::max(v, max); + } + + const float d = (max - min) / ((1 << 8) - 1); + const float id = d ? 1.0f / d : 0.0f; + scales[i] = ov::float16(d); + // zp = -min / scale (Q8_1 is asymmetric) + zp[i] = (d != 0.0f) ? (uint8_t) std::round(-min / d) : 0; + + for (int j = 0; j < qk; ++j) { + const float x0 = (x[i * qk + j] - min) * id; + const uint8_t xi0 = roundf(x0); + weights[i * qk + j] = xi0; + } + } +} + + +/* Inlined utils.cpp */ +/* Inlined utils.h */ +#include "ggml-backend-impl-defs.inc" +#include "ggml.h.inc" +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +#include +#include +#include +#include +#include +#include +#include + +struct ModelParams { + int ctx = -1; + int ctx_swa = -1; + int ctx_per_seq = -1; + int ctx_per_seq_swa = -1; + int n_seq = 1; + int n_heads = -1; + int n_heads_kv = -1; + int head_size = -1; + int32_t rope_params[15]; + std::vector swa_layers; + + std::vector kv_names; + size_t kv_buffer_ctx_id = 0; + + bool same_rope_params(const ModelParams & other) const { + return memcmp(rope_params, other.rope_params, sizeof(int32_t) * 15) == 0; + } + + bool can_reuse_dynamically(const ModelParams & other) const { return same_rope_params(other); } + + bool can_reuse_statically(const ModelParams & other) const { return same_rope_params(other) && ctx == other.ctx; } + + bool kv_buffer_changed(const ModelParams & other) const { return kv_buffer_ctx_id != other.kv_buffer_ctx_id; } +}; + +struct ComputeParams { + int n_seq_active = 1; + int seq_active_start = 0; + int attention_size = -1; + int attention_size_swa = -1; + int input_len = -1; + int token_len_per_seq = -1; + int past_kv_len = -1; + int output_len = 1; +}; + +class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder { +public: + struct NodeInfo { + ggml_tensor * node; + std::string node_name; + std::string node_op_type; + std::map node_inputs; + std::vector node_inputs_names; + ggml_tensor * node_output; + std::string node_output_name; + int node_op_case = 0; + void * data_addr; + }; + // Graph decoder + GgmlOvDecoder(ggml_cgraph * cgraph, + ModelParams & model_params, + ComputeParams & compute_params, + std::map> & model_weights, + bool is_static, + bool is_stateful = false, + bool is_prefill = false, + int prefill_chunk_size = 256); + + // Naive graph decoder + GgmlOvDecoder(ggml_cgraph * cgraph, std::map> & model_weights); + + virtual ov::Any get_attribute(const std::string & name) const override { + return nullptr; + GGML_UNUSED(name); + } + + virtual ov::PartialShape get_input_shape(int node_idx, const std::string & name) const override; + + virtual std::vector get_input_stride(int node_idx, const std::string & name) const override; + + virtual ov::element::Type get_input_type(int node_idx, const std::string & name) const override; + + virtual size_t get_input_size() const override; + + virtual size_t get_input_size(int node_idx) const override; + + virtual void get_input_node(size_t input_port_idx, + std::string & producer_name, + std::string & producer_output_port_name, + size_t & producer_output_port_index) const override { + GGML_UNUSED(input_port_idx); + GGML_UNUSED(producer_name); + GGML_UNUSED(producer_output_port_name); + GGML_UNUSED(producer_output_port_index); + } + + virtual std::vector get_input_names(int node_idx) const override; + + virtual ov::PartialShape get_output_shape(int node_idx) const override; + + virtual ov::element::Type get_output_type(int node_idx) const override; + + virtual int32_t * get_input_op_params(int node_idx, const std::string & name) const override; + + virtual int32_t * get_output_op_params(int node_idx) const override; + + virtual std::vector get_output_names(int node_idx) const override; + + virtual const std::string & get_op_type() const override; + + virtual const std::string & get_op_type(int node_idx) const override; + + virtual const std::string & get_op_name() const override; + + virtual const std::string & get_op_name(int node_idx) const override; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const override; + + ggml_tensor * get_input_ggml_tensor(const std::string & name) const { return m_inputs.at(name); } + + virtual int get_op_case(int node_idx) const override { return m_node_info_list[node_idx].node_op_case; } + + virtual const std::map> & get_model_inputs() const override { + return m_model_inputs; + } + + virtual const std::map> & get_model_extra_inputs() const override { + return m_model_extra_inputs; + } + + virtual const std::map> & get_model_extra_input_values() const { + return m_model_extra_input_values; + } + + virtual const std::map> & get_model_weights() const override { + return m_model_weights; + } + + virtual std::vector get_model_output_names() const override { + return m_model_output_names; + } + + const std::map & get_model_outputs() const { return m_model_outputs; } + + virtual int get_ctx_size() const { return m_model_params.ctx; } + + virtual int get_ctx_swa_size() const { return m_model_params.ctx_swa; } + + virtual int get_ctx_per_seq() const { return m_model_params.ctx_per_seq; } + + virtual int get_ctx_per_seq_swa() const { return m_model_params.ctx_per_seq_swa; } + + virtual int get_n_seq() const { return m_model_params.n_seq; } + + virtual int is_swa_layer(int layer) const override { + return std::find(m_model_params.swa_layers.begin(), m_model_params.swa_layers.end(), layer) != + m_model_params.swa_layers.end(); + } + + int get_past_kv_len() const { return m_compute_params.past_kv_len; } + + int get_input_len() const { return m_compute_params.input_len; } + + virtual int32_t * get_rope_params() const override { return const_cast(m_model_params.rope_params); } + + virtual std::map get_kv_param_res_names() const override; + + virtual bool is_static() const override { return m_is_static; } + + virtual bool is_stateful() const override { return m_is_stateful; } + + ov::PartialShape get_graph_input_shape(const ggml_tensor * op, const ggml_tensor * input) const; + + static void dump_cgraph(const ggml_cgraph * cgraph, std::string & filename); + + static std::shared_ptr create_weight_node(ggml_tensor * tensor, bool naive = false); + + static std::map> create_weight_nodes(ggml_cgraph * cgraph, + bool naive = false); + + const ggml_tensor * get_tensor_used_op(const ggml_tensor * tensor) const; + + const ggml_tensor * get_tensor_from_name(const std::string & name) const; + + void clear_model_weights() { m_model_weights.clear(); } + + static std::pair compute_llm_params(ggml_cgraph * cgraph, bool is_static); + + ModelParams get_model_params() const { return m_model_params; } + + ComputeParams get_compute_params() const { return m_compute_params; } + + void set_model_params(const ModelParams & model_params) { m_model_params = model_params; } + + void set_compute_params(const ComputeParams & compute_params) { m_compute_params = compute_params; } + + bool m_is_static = false; + bool m_is_stateful = false; + bool m_is_prefill = false; + bool m_naive = false; + int m_prefill_chunk_size = 0; + + static ov::Shape get_shape(const ggml_tensor * tensor); + static std::vector get_stride(const ggml_tensor * tensor); + static ov::element::Type get_ov_type(const ggml_tensor * tensor); + static std::string compute_op_type(const ggml_tensor * node); + void add_extra_inputs(); + + void update_io(ggml_cgraph * cgraph); + + inline static bool is_inp_tok(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_GET_ROWS && tensor == op->src[1] && op->src[0]->op == GGML_OP_NONE; + } + + inline static bool is_inp_pos(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_ROPE && tensor == op->src[1]; + } + + inline static bool is_inp_emb(const ggml_tensor * tensor, const ggml_tensor * op) { + return tensor->op == GGML_OP_GET_ROWS && op->op == GGML_OP_RMS_NORM; + } + + inline static bool is_inp_mask(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_CPY || (op->op == GGML_OP_FLASH_ATTN_EXT && tensor == op->src[3]); + } + + inline static bool is_rope_freqs_weight(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_ROPE && tensor == op->src[2]; + } + + inline static bool is_kvcache(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_SET_ROWS && op->src[2] == tensor; + } + + inline static bool is_kv_idx(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_SET_ROWS && op->src[1] == tensor; + } + + inline static bool is_output_idx(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_GET_ROWS && tensor == op->src[1] && op->src[0]->op != GGML_OP_NONE; + } + + static std::string get_graph_input_ov_name(const ggml_tensor * tensor, const ggml_tensor * op) { + if (is_inp_tok(tensor, op)) { + return "inp_tokens"; + } + if (is_inp_pos(tensor, op)) { + return "inp_pos"; + } + if (is_inp_emb(tensor, op)) { + return "embd"; + } + if (is_output_idx(tensor, op)) { + return "inp_out_ids"; + } + if (is_inp_mask(tensor, op)) { + return std::string(tensor->name).find("swa") == std::string::npos ? "self_kq_mask" : "self_kq_mask_swa"; + } + return tensor->name; + } + +private: + void set_input_output(); + int compute_op_case(const ggml_tensor * node) const; + bool node_is_used_as_src(const int node_idx); + void compute_model_inputs(); + void compute_model_outputs(); + + void validate_cgraph() const; + + ggml_cgraph * m_cgraph = nullptr; + std::map m_inputs; + + std::map> m_model_inputs; + std::map> m_model_extra_inputs; + std::map> m_model_extra_input_values; + std::map> m_model_weights; + std::map m_model_outputs; + std::vector m_model_output_names; + std::vector m_node_info_list; + + ModelParams m_model_params; + ComputeParams m_compute_params; +}; + +void print_tensor_address_map(const ggml_cgraph * cgraph); + +int extract_layer_from_name(const std::string & name); + +#include "ggml-impl-defs.inc" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct graph_key { + int n_nodes; + std::string first_node_name; + std::string last_node_name; + + graph_key(const ggml_cgraph * cgraph) : n_nodes(cgraph->n_nodes) { + if (n_nodes > 0) { + first_node_name = cgraph->nodes[0]->name; + last_node_name = cgraph->nodes[n_nodes - 1]->name; + } + } + + bool operator==(const graph_key & other) const { + return n_nodes == other.n_nodes && first_node_name == other.first_node_name && + last_node_name == other.last_node_name; + } +}; + +struct graph_key_hash { + size_t operator()(const graph_key & key) const { + size_t h = std::hash{}(key.n_nodes); + if (key.n_nodes > 0) { + h ^= std::hash{}(key.first_node_name) + 0x9e3779b9 + (h << 6) + (h >> 2); + h ^= std::hash{}(key.last_node_name) + 0x9e3779b9 + (h << 6) + (h >> 2); + } + return h; + } +}; + +struct decoder_runtime_ctx { + decoder_runtime_ctx(std::shared_ptr mutex) : mutex(std::move(mutex)) {} + std::shared_ptr mutex; + std::shared_ptr ptr; +}; + +struct ov_runtime_context { + mutable std::mutex ctx_mutex; + std::string device; + bool stateful; + std::unordered_map, graph_key_hash> decoder_cache; + std::unordered_map, graph_key_hash> infer_request_cache; + std::unordered_map, graph_key_hash> infer_request_cache_prefill; + std::unordered_map, graph_key_hash> ov_input_names_cache; + std::unordered_map, graph_key_hash> ov_output_names_cache; + //TODO: Stateful is only supported for single request at a time. + // Simultanous stateful inference request support to be added. + size_t stateful_kv_size; + std::map kv_state_input_name_map; + std::atomic backend_count; + + ov_runtime_context() : + device("CPU"), + stateful(false), + stateful_kv_size(0), + backend_count(0) {} + + void clear_caches() { + std::lock_guard lock(ctx_mutex); + decoder_cache.clear(); + infer_request_cache.clear(); + infer_request_cache_prefill.clear(); + ov_input_names_cache.clear(); + ov_output_names_cache.clear(); + } +}; + +enum ggml_status ov_graph_compute(struct ggml_cgraph * cgraph, ggml_backend_t backend); + +enum ggml_status ov_graph_compute_dynamic(struct ggml_cgraph * cgraph, std::shared_ptr r_ctx); +enum ggml_status ov_graph_compute_static(struct ggml_cgraph * cgraph, std::shared_ptr r_ctx); + +size_t checksum(const void * data, size_t size); + +void print_input_tensor_info(const std::string & name, const ov::Tensor & tensor); + +void print_output_tensor_info(const std::string & name, const ov::Tensor & tensor, const void * output_dst); + +template +std::vector pad_input(const T * data, + size_t rows, + size_t cols, + size_t padded_rows, + size_t padded_cols, + T pad_value) { + std::vector padded(padded_rows * padded_cols, pad_value); + + for (size_t i = 0; i < std::min(rows, padded_rows); ++i) { + for (size_t j = 0; j < std::min(cols, padded_cols); ++j) { + padded[i * padded_cols + j] = data[i * cols + j]; + } + } + + return padded; +} + +template +std::vector pad_input(const ggml_tensor * tensor, size_t padded_rows, size_t padded_cols, T pad_value) { + return pad_input(reinterpret_cast(tensor->data), + static_cast(tensor->ne[1]), // rows + static_cast(tensor->ne[0]), // cols + padded_rows, padded_cols, pad_value); +} + +void set_zero_diagonal(std::vector & matrix, size_t rows, size_t cols); + +const ggml_tensor * get_inp_pos_tensor(struct ggml_cgraph * cgraph); + +bool get_is_prefill(const ggml_tensor * inp_pos); + +ov::Tensor get_ov_input_tensor(std::shared_ptr ggml_decoder, const std::string & param_name); +ov::Tensor get_ov_input_tensor_static_decode(std::shared_ptr ggml_decoder, + const std::string & param_name); +ov::Tensor get_ov_input_tensor_static_prefill(std::shared_ptr ggml_decoder, + const std::string & param_name, + int chunk_index); + +ov::Tensor create_ov_output_tensor(std::shared_ptr ggml_decoder, + std::shared_ptr infer_request, + int output_index, + const ggml_tensor * ggml_tensor); + +bool is_naive(struct ggml_cgraph * cgraph); + +enum ggml_status naive_compute(struct ggml_cgraph * cgraph, + ov::Core & core, + const std::string & device, + const ov::AnyMap & config); + + +#include "ggml-impl-defs.inc" +#include "ggml.h.inc" +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class FrontEnd { +public: + using Ptr = std::shared_ptr; + FrontEnd(); + + static std::shared_ptr convert(const InputModel::Ptr& model, bool naive = false); +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Suppress deprecation warning for ov::Tensor::data() +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" + +enum ggml_status ov_graph_compute(ggml_cgraph * cgraph, ggml_backend_t backend) { + ggml_backend_openvino_context * ctx = (ggml_backend_openvino_context *) backend->context; + try { + if (getenv("GGML_OPENVINO_DUMP_CGRAPH")) { + std::string filename = "cgraph_ov.txt"; + GgmlOvDecoder::dump_cgraph(cgraph, filename); + } + + const auto is_static = ggml_openvino_is_npu(); + + GGML_ASSERT(ctx->runtime_context != nullptr); + std::shared_ptr r_ctx = std::static_pointer_cast(ctx->runtime_context); + + return is_static ? ov_graph_compute_static(cgraph, r_ctx) : ov_graph_compute_dynamic(cgraph, r_ctx); + } catch (const ov::Exception & e) { + GGML_LOG_ERROR("GGML OpenVINO backend ov::Exception: %s\n", e.what()); + return GGML_STATUS_FAILED; + } catch (const std::exception & e) { + GGML_LOG_ERROR("GGML OpenVINO backend std::exception: %s\n", e.what()); + return GGML_STATUS_FAILED; + } catch (...) { + GGML_LOG_ERROR("GGML OpenVINO backend unknown exception\n"); + return GGML_STATUS_FAILED; + } +} + +ov::Tensor create_ov_output_tensor(std::shared_ptr ggml_decoder, + std::shared_ptr infer_request, + int output_index, + const ggml_tensor * ggml_tensor) { + auto output_type = ggml_decoder->get_ov_type(ggml_tensor); + ov::Shape output_shape; + if (ggml_decoder->is_static()) { + output_shape = infer_request->get_output_tensor(output_index).get_shape(); + } else { + output_shape = ggml_decoder->get_shape(ggml_tensor); + } + + ov::Tensor output_tensor(output_type, output_shape, ggml_tensor->data); + return output_tensor; +} + +enum ggml_status ov_graph_compute_dynamic(ggml_cgraph * cgraph, std::shared_ptr r_ctx) { + auto & core = ov_singleton_core(); + const auto & config = ggml_openvino_get_compile_config(); + const auto & device = r_ctx->device; + const auto & stateful = r_ctx->stateful; + static auto is_static = false; + + if (is_naive(cgraph)) { + return naive_compute(cgraph, core, device, config); + } + + auto start_time = ggml_time_us(); + + std::shared_ptr ggml_decoder; + std::shared_ptr infer_request; + ModelParams m_params; + ComputeParams c_params; + std::tie(m_params, c_params) = GgmlOvDecoder::compute_llm_params(cgraph, is_static); + + graph_key key(cgraph); + bool cache_hit; + + int64_t decoder_end_time; + int64_t conversion_end_time; + int64_t compile_end_time; + int64_t infer_end_time; + + { + std::shared_ptr entry; + ModelParams old_m_params; + + { + std::lock_guard map_lock(r_ctx->ctx_mutex); + auto it = r_ctx->decoder_cache.find(key); + cache_hit = it != r_ctx->decoder_cache.end(); + if (cache_hit) { + entry = it->second; + } else { + auto mutex = std::make_shared(); + entry = std::make_shared(mutex); + r_ctx->decoder_cache[key] = entry; + } + } + + std::lock_guard lock(*(entry->mutex)); + + if (cache_hit) { + ggml_decoder = entry->ptr; + old_m_params = ggml_decoder->get_model_params(); + cache_hit = old_m_params.can_reuse_dynamically(m_params); + } + + if (cache_hit) { + std::map> model_weights; + ggml_decoder->set_compute_params(c_params); + ggml_decoder->set_model_params(m_params); + if (old_m_params.kv_buffer_changed(m_params)) { + ggml_decoder->update_io(cgraph); + } + ggml_decoder->add_extra_inputs(); + { + std::lock_guard map_lock(r_ctx->ctx_mutex); + infer_request = r_ctx->infer_request_cache.at(key); + } + + if (stateful) { + const auto * inp_pos = get_inp_pos_tensor(cgraph); + int32_t * pos_data = (int32_t *) inp_pos->data; + auto pos_shape = ggml_decoder->get_shape(inp_pos); + if (pos_data[0] == 0) { + infer_request->reset_state(); + r_ctx->stateful_kv_size = pos_shape[3]; + } else if (r_ctx->stateful_kv_size == static_cast(pos_data[0])) { + r_ctx->stateful_kv_size += pos_shape[3]; + } else { + auto states = infer_request->query_state(); + for (auto state : states) { + auto state_tensor = state.get_state(); + auto state_tensor_shape = state_tensor.get_shape(); + if (static_cast(pos_data[0]) > r_ctx->stateful_kv_size) { + std::string state_name; + try { + state_name = r_ctx->kv_state_input_name_map.at(state.get_name()); + } catch (...) { + GGML_LOG_ERROR("GGML OpenVINO backend stateful inference failed: no input found for the state\n"); + return GGML_STATUS_FAILED; + } + auto kv_tensor = get_ov_input_tensor(ggml_decoder, state_name); + kv_tensor.set_shape({state_tensor_shape[0], kv_tensor.get_shape()[2], + state_tensor_shape[2], state_tensor_shape[3]}); + state_tensor = kv_tensor; + state_tensor_shape = state_tensor.get_shape(); + } + ov::Coordinate begin = {0, 0, 0, 0}; + ov::Coordinate end = {state_tensor_shape[0], static_cast(pos_data[0]), + state_tensor_shape[2], state_tensor_shape[3]}; + ov::Tensor new_state_tensor(state_tensor, begin, end); + state.set_state(new_state_tensor); + } + r_ctx->stateful_kv_size = pos_data[0] + 1; + } + } + + decoder_end_time = ggml_time_us(); + conversion_end_time = decoder_end_time; + compile_end_time = decoder_end_time; + } else { + { + std::lock_guard map_lock(r_ctx->ctx_mutex); + r_ctx->infer_request_cache.erase(key); + } + + std::shared_ptr model; + auto model_weights = GgmlOvDecoder::create_weight_nodes(cgraph); + + ggml_decoder = std::make_shared(cgraph, m_params, c_params, model_weights, is_static, stateful); + decoder_end_time = ggml_time_us(); + + auto input_model = std::make_shared(ggml_decoder); + model = ov::frontend::ggml::FrontEnd::convert(input_model); + ggml_decoder->clear_model_weights(); + conversion_end_time = ggml_time_us(); + + if (getenv("GGML_OPENVINO_DUMP_IR")) { + char timestamped_filename[64]; + auto timestamp = (long long) ggml_time_us(); + snprintf(timestamped_filename, sizeof(timestamped_filename), "model_%lld.xml", timestamp); + ov::serialize(model, timestamped_filename); + } + + ov::CompiledModel compiled_model; + auto remote_context = ggml_openvino_get_remote_context(); + if (remote_context.has_value()) { + compiled_model = core.compile_model(model, remote_context.value(), config); + } else { + compiled_model = core.compile_model(model, device, config); + } + compile_end_time = ggml_time_us(); + infer_request = std::make_shared(compiled_model.create_infer_request()); + entry->ptr = ggml_decoder; + + std::vector ov_input_names; + std::vector ov_output_names; + for (const auto & ov_param : model->get_parameters()) { + ov_input_names.push_back(ov_param->get_friendly_name()); + } + for (const auto & ov_output : model->get_results()) { + ov_output_names.push_back(ov_output->get_friendly_name()); + } + + { + std::lock_guard map_lock(r_ctx->ctx_mutex); + r_ctx->infer_request_cache[key] = infer_request; + r_ctx->ov_input_names_cache[key] = std::move(ov_input_names); + r_ctx->ov_output_names_cache[key] = std::move(ov_output_names); + } + + if (stateful) { + const auto * inp_pos = get_inp_pos_tensor(cgraph); + auto pos_shape = ggml_decoder->get_shape(inp_pos); + r_ctx->stateful_kv_size = pos_shape[3]; + const auto kv_param_res_names = ggml_decoder->get_kv_param_res_names(); + for (const auto& pair : kv_param_res_names) { + r_ctx->kv_state_input_name_map[pair.first+pair.second] = pair.first; + } + } + } + + std::vector ov_input_names; + std::vector ov_output_names; + { + std::lock_guard map_lock(r_ctx->ctx_mutex); + ov_input_names = r_ctx->ov_input_names_cache[key]; + ov_output_names = r_ctx->ov_output_names_cache[key]; + } + + for (size_t i = 0; i < ov_input_names.size(); i++) { + auto param_name = ov_input_names[i]; + auto input_tensor = get_ov_input_tensor(ggml_decoder, param_name); + infer_request->set_input_tensor(i, input_tensor); + + if (getenv("GGML_OPENVINO_DEBUG_INPUT")) { + print_input_tensor_info(param_name, input_tensor); + } + } + + for (size_t i = 0; i < ov_output_names.size(); i++) { + auto * ggml_tensor = ggml_decoder->get_model_outputs().at(ov_output_names[i]); + auto output_tensor = create_ov_output_tensor(ggml_decoder, infer_request, i, ggml_tensor); + infer_request->set_output_tensor(i, output_tensor); + } + + infer_request->infer(); + infer_end_time = ggml_time_us(); + + if (getenv("GGML_OPENVINO_DEBUG_OUTPUT")) { + for (size_t i = 0; i < ov_output_names.size(); i++) { + const auto output_tensor = infer_request->get_output_tensor(i); + print_output_tensor_info(ov_output_names[i], output_tensor, output_tensor.data()); + } + } + + if (getenv("GGML_OPENVINO_PROFILING")) { + GGML_LOG_INFO("\nGGML OpenVINO Backend: \n"); + GGML_LOG_INFO(" - Graph decoder time: %ld ms \n", (decoder_end_time - start_time) / 1000); + if (!cache_hit) { + GGML_LOG_INFO(" - Graph conversion time: %ld ms \n", (conversion_end_time - decoder_end_time) / 1000); + GGML_LOG_INFO(" - Graph compile time: %ld ms \n", (compile_end_time - conversion_end_time) / 1000); + } + GGML_LOG_INFO(" - Graph inference time: %ld ms \n", (infer_end_time - compile_end_time) / 1000); + } + } + + return GGML_STATUS_SUCCESS; +} + +enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptr r_ctx) { + auto & core = ov_singleton_core(); + + auto get_prefill_chunk_size = [] { + const char * chunk_size_str = getenv("GGML_OPENVINO_PREFILL_CHUNK_SIZE"); + if (chunk_size_str && atoi(chunk_size_str) > 0) { + return atoi(chunk_size_str); + } + return 256; + }; + + static std::string device = "NPU"; + static auto is_static = true; + static auto stateful = false; + static auto prefill_chunk_size = get_prefill_chunk_size(); + const auto & config = ggml_openvino_get_compile_config(); + + if (is_naive(cgraph)) { + return naive_compute(cgraph, core, device, config); + } + + auto start_time = ggml_time_us(); + + std::shared_ptr ggml_decoder; + std::shared_ptr infer_request; + ModelParams m_params; + ComputeParams c_params; + std::tie(m_params, c_params) = GgmlOvDecoder::compute_llm_params(cgraph, is_static); + + const auto * inp_pos = get_inp_pos_tensor(cgraph); + const auto is_prefill = get_is_prefill(inp_pos); + graph_key key(cgraph); + bool cache_hit; + + int64_t decoder_end_time; + int64_t conversion_end_time; + int64_t compile_end_time; + int64_t infer_end_time; + + std::shared_ptr entry; + ModelParams old_m_params; + + { + std::lock_guard map_lock(r_ctx->ctx_mutex); + auto it = r_ctx->decoder_cache.find(key); + cache_hit = it != r_ctx->decoder_cache.end(); + if (cache_hit) { + entry = it->second; + } else { + auto mutex = std::make_shared(); + entry = std::make_shared(mutex); + r_ctx->decoder_cache[key] = entry; + } + } + + std::lock_guard lock(*(entry->mutex)); + + if (cache_hit) { + ggml_decoder = entry->ptr; + old_m_params = ggml_decoder->get_model_params(); + cache_hit = old_m_params.can_reuse_statically(m_params); + } + + if (cache_hit) { + std::map> model_weights; + ggml_decoder->m_is_prefill = is_prefill; + ggml_decoder->set_model_params(m_params); + ggml_decoder->set_compute_params(c_params); + if (old_m_params.kv_buffer_changed(m_params)) { + ggml_decoder->update_io(cgraph); + } + ggml_decoder->add_extra_inputs(); + { + std::lock_guard map_lock(r_ctx->ctx_mutex); + infer_request = + is_prefill ? r_ctx->infer_request_cache_prefill.at(key) : r_ctx->infer_request_cache.at(key); + } + + decoder_end_time = ggml_time_us(); + conversion_end_time = decoder_end_time; + compile_end_time = decoder_end_time; + } else { + { + std::lock_guard map_lock(r_ctx->ctx_mutex); + r_ctx->infer_request_cache.erase(key); + r_ctx->infer_request_cache_prefill.erase(key); + } + + std::shared_ptr model; + auto model_weights = GgmlOvDecoder::create_weight_nodes(cgraph); + + auto ggml_decoder_prefill = std::make_shared(cgraph, m_params, c_params, model_weights, + is_static, stateful, true, prefill_chunk_size); + auto ggml_decoder_decode = std::make_shared(cgraph, m_params, c_params, model_weights, is_static, + stateful, false, prefill_chunk_size); + decoder_end_time = ggml_time_us(); + + auto input_model_prefill = std::make_shared(ggml_decoder_prefill); + auto input_model_decode = std::make_shared(ggml_decoder_decode); + + auto model_prefill = ov::frontend::ggml::FrontEnd::convert(input_model_prefill); + ggml_decoder_prefill->clear_model_weights(); + auto model_decode = ov::frontend::ggml::FrontEnd::convert(input_model_decode); + ggml_decoder_decode->clear_model_weights(); + conversion_end_time = ggml_time_us(); + + if (getenv("GGML_OPENVINO_DUMP_IR")) { + char timestamped_filename[64]; + auto timestamp = (long long) ggml_time_us(); + snprintf(timestamped_filename, sizeof(timestamped_filename), "model_prefill_%lld.xml", timestamp); + ov::serialize(model_prefill, timestamped_filename); + snprintf(timestamped_filename, sizeof(timestamped_filename), "model_decode_%lld.xml", timestamp); + ov::serialize(model_decode, timestamped_filename); + } + + ov::CompiledModel compiled_model_prefill; + ov::CompiledModel compiled_model_decode; + auto remote_context = ggml_openvino_get_remote_context(); + if (remote_context.has_value()) { + compiled_model_prefill = core.compile_model(model_prefill, remote_context.value(), config); + compiled_model_decode = core.compile_model(model_decode, remote_context.value(), config); + } else { + compiled_model_prefill = core.compile_model(model_prefill, device, config); + compiled_model_decode = core.compile_model(model_decode, device, config); + } + + auto infer_request_prefill = std::make_shared(compiled_model_prefill.create_infer_request()); + auto infer_request_decode = std::make_shared(compiled_model_decode.create_infer_request()); + compile_end_time = ggml_time_us(); + + model = is_prefill ? model_prefill : model_decode; + ggml_decoder = is_prefill ? ggml_decoder_prefill : ggml_decoder_decode; + infer_request = is_prefill ? infer_request_prefill : infer_request_decode; + entry->ptr = ggml_decoder; + + std::vector ov_input_names; + std::vector ov_output_names; + for (const auto & ov_param : model->get_parameters()) { + ov_input_names.push_back(ov_param->get_friendly_name()); + } + for (const auto & ov_output : model->get_results()) { + ov_output_names.push_back(ov_output->get_friendly_name()); + } + + { + std::lock_guard map_lock(r_ctx->ctx_mutex); + r_ctx->infer_request_cache_prefill[key] = infer_request_prefill; + r_ctx->infer_request_cache[key] = infer_request_decode; + r_ctx->ov_input_names_cache[key] = std::move(ov_input_names); + r_ctx->ov_output_names_cache[key] = std::move(ov_output_names); + } + } + + std::vector ov_input_names_local; + std::vector ov_output_names_local; + { + std::lock_guard map_lock(r_ctx->ctx_mutex); + ov_input_names_local = r_ctx->ov_input_names_cache[key]; + ov_output_names_local = r_ctx->ov_output_names_cache[key]; + } + + if (is_prefill) { + auto inp_len = inp_pos->ne[0]; + for (int chunk_index = 0; chunk_index * prefill_chunk_size < inp_len; chunk_index++) { + for (size_t i = 0; i < ov_input_names_local.size(); i++) { + auto param_name = ov_input_names_local[i]; + auto input_tensor = get_ov_input_tensor_static_prefill(ggml_decoder, param_name, chunk_index); + infer_request->set_input_tensor(i, input_tensor); + + if (getenv("GGML_OPENVINO_DEBUG_INPUT")) { + const auto input_tensor = infer_request->get_input_tensor(i); + print_input_tensor_info(param_name, input_tensor); + } + } + + for (size_t i = 0; i < ov_output_names_local.size(); i++) { + auto * ggml_tensor = ggml_decoder->get_model_outputs().at(ov_output_names_local[i]); + auto output_tensor = create_ov_output_tensor(ggml_decoder, infer_request, i, ggml_tensor); + infer_request->set_output_tensor(i, output_tensor); + } + + infer_request->infer(); + + if (getenv("GGML_OPENVINO_DEBUG_OUTPUT")) { + for (size_t i = 0; i < ov_output_names_local.size(); i++) { + const auto output_tensor = infer_request->get_output_tensor(i); + print_output_tensor_info(ov_output_names_local[i], output_tensor, output_tensor.data()); + } + } + } + infer_end_time = ggml_time_us(); + } else { + for (size_t i = 0; i < ov_input_names_local.size(); i++) { + auto param_name = ov_input_names_local[i]; + auto input_tensor = get_ov_input_tensor_static_decode(ggml_decoder, param_name); + infer_request->set_input_tensor(i, input_tensor); + + if (getenv("GGML_OPENVINO_DEBUG_INPUT")) { + const auto input_tensor = infer_request->get_input_tensor(i); + print_input_tensor_info(param_name, input_tensor); + } + } + + for (size_t i = 0; i < ov_output_names_local.size(); i++) { + auto * ggml_tensor = ggml_decoder->get_model_outputs().at(ov_output_names_local[i]); + auto output_tensor = create_ov_output_tensor(ggml_decoder, infer_request, i, ggml_tensor); + infer_request->set_output_tensor(i, output_tensor); + } + + infer_request->infer(); + infer_end_time = ggml_time_us(); + + if (getenv("GGML_OPENVINO_DEBUG_OUTPUT")) { + for (size_t i = 0; i < ov_output_names_local.size(); i++) { + const auto output_tensor = infer_request->get_output_tensor(i); + print_output_tensor_info(ov_output_names_local[i], output_tensor, output_tensor.data()); + } + } + } + + if (getenv("GGML_OPENVINO_PROFILING")) { + GGML_LOG_INFO("\nGGML OpenVINO Backend: \n"); + GGML_LOG_INFO(" - Graph decoder time: %ld ms \n", (decoder_end_time - start_time) / 1000); + if (!cache_hit) { + GGML_LOG_INFO(" - Graph conversion time: %ld ms \n", (conversion_end_time - decoder_end_time) / 1000); + GGML_LOG_INFO(" - Graph compile time: %ld ms \n", (compile_end_time - conversion_end_time) / 1000); + } + GGML_LOG_INFO(" - Graph inference time: %ld ms \n", (infer_end_time - compile_end_time) / 1000); + } + + return GGML_STATUS_SUCCESS; +} + +bool is_naive(ggml_cgraph * cgraph) { + constexpr int naive_graph_size_threshold = 20; + int count = 0; + for (int i = 0; i < cgraph->n_nodes; i++) { + if (cgraph->nodes[i]->op != GGML_OP_NONE) { + count++; + } + } + return count < naive_graph_size_threshold; +} + +enum ggml_status naive_compute(ggml_cgraph * cgraph, + ov::Core & core, + const std::string & device, + const ov::AnyMap & config) { + if (cgraph->n_nodes == 1 && (cgraph->nodes[0]->op == GGML_OP_NONE || cgraph->nodes[0]->op == GGML_OP_VIEW)) { + return GGML_STATUS_SUCCESS; + } + + bool naive = true; + auto model_weights = GgmlOvDecoder::create_weight_nodes(cgraph, naive); + auto decoder = std::make_shared(cgraph, model_weights); + auto input_model = std::make_shared(decoder); + auto model = ov::frontend::ggml::FrontEnd::convert(input_model, naive); + if (getenv("GGML_OPENVINO_DUMP_IR")) { + ov::serialize(model, "IR_naive.xml"); + } + + std::shared_ptr infer_request; + auto remote_context = ggml_openvino_get_remote_context(); + if (cgraph->nodes[0]->op == GGML_OP_MUL_MAT) { + // TODO ACCURACY hint triggers a bug in GPU plugin/driver on Lunar Lake. Remove once CVS-182166 is resolved + core.set_property(device, ov::hint::execution_mode(ov::hint::ExecutionMode::PERFORMANCE)); + } else { + core.set_property(device, ov::hint::execution_mode(ov::hint::ExecutionMode::ACCURACY)); + } + if (remote_context.has_value()) { + infer_request = std::make_shared( + core.compile_model(model, remote_context.value(), config).create_infer_request()); + } else { + infer_request = + std::make_shared(core.compile_model(model, device, config).create_infer_request()); + } + + auto ov_params = model->get_parameters(); + for (size_t i = 0; i < ov_params.size(); i++) { + auto param_name = ov_params[i]->get_friendly_name(); + auto input_tensor = get_ov_input_tensor(decoder, param_name); + infer_request->set_input_tensor(i, input_tensor); + } + + auto ov_results = model->get_results(); + for (size_t i = 0; i < ov_results.size(); i++) { + auto * ggml_tensor = decoder->get_model_outputs().at(ov_results[i]->get_friendly_name()); + auto output_tensor = create_ov_output_tensor(decoder, infer_request, i, ggml_tensor); + infer_request->set_output_tensor(i, output_tensor); + } + + infer_request->infer(); + return GGML_STATUS_SUCCESS; +} + +namespace { +ov::Tensor convert_ggml_input_to_ov(std::shared_ptr ggml_decoder, const std::string & name) { + const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(name); + + if (ggml_tensor->extra != nullptr) { + // GGML_LOG_DEBUG("Using ggml_tensor->extra as ov::Tensor for input: %s\n", name.c_str()); + auto * extra_base = static_cast(ggml_tensor->extra); + if (extra_base->type != ggml_openvino_extra_base::Type::TENSOR) { + throw std::runtime_error("ggml tensor extra is not of type TENSOR for input: " + name); + } + auto * tensor_extra = static_cast(extra_base); + return *tensor_extra->tensor; + } + + // GGML_LOG_DEBUG("Converting ggml tensor to ov::Tensor for input: %s\n", name.c_str()); + auto * input_data = ggml_tensor->data; + ov::Shape input_shape; + if (ggml_tensor->op == GGML_OP_VIEW) { + // This case is added to make test-backend-ops work + input_shape = ggml_decoder->get_shape(ggml_tensor->view_src); + } else { + input_shape = ggml_decoder->get_shape(ggml_tensor); + } + auto input_tensor = ov::Tensor(ggml_decoder->get_ov_type(ggml_tensor), input_shape, input_data); + return input_tensor; +} +} // namespace + +ov::Tensor get_ov_input_tensor(std::shared_ptr ggml_decoder, const std::string & param_name) { + ov::Tensor input_tensor; + if (ggml_decoder->get_model_extra_inputs().find(param_name) != ggml_decoder->get_model_extra_inputs().end()) { + input_tensor = *ggml_decoder->get_model_extra_input_values().at(param_name); + } else { + input_tensor = convert_ggml_input_to_ov(ggml_decoder, param_name); + } + return input_tensor; +} + +ov::Tensor get_ov_input_tensor_static_decode(std::shared_ptr ggml_decoder, + const std::string & param_name) { + // NPU decoding stage + const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name); + const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor); + + if (GgmlOvDecoder::is_inp_tok(ggml_tensor, op) || GgmlOvDecoder::is_inp_pos(ggml_tensor, op) || + GgmlOvDecoder::is_kv_idx(ggml_tensor, op)) { + assert(ggml_tensor->ne[0] == 1); + ov::Shape input_shape = {1, 1, 1, 1}; + ov::Tensor input_tensor(ggml_decoder->get_ov_type(ggml_tensor), input_shape); + if (ggml_tensor->type == GGML_TYPE_I32) { + *input_tensor.data() = *((int32_t *) ggml_tensor->data); + } else if (ggml_tensor->type == GGML_TYPE_I64) { + *input_tensor.data() = *((int64_t *) ggml_tensor->data); + } else { + throw std::runtime_error("Unexpected tensor type for " + param_name); + } + return input_tensor; + } + + if (GgmlOvDecoder::is_output_idx(ggml_tensor, op)) { + ov::Shape input_shape = {1, 1, 1, 1}; + ov::Tensor input_tensor(ggml_decoder->get_ov_type(ggml_tensor), input_shape); + int32_t inp_out_id = *((int32_t *) ggml_tensor->data); + assert(ggml_tensor->ne[0] == 1); + assert(inp_out_id == 0); + *input_tensor.data() = inp_out_id; + return input_tensor; + } + + if (GgmlOvDecoder::is_inp_mask(ggml_tensor, op)) { + size_t context_size = ggml_decoder->get_ctx_size(); + std::vector padded_data = pad_input(ggml_tensor, 1, context_size, -INFINITY); + ov::Tensor input_tensor(ov::element::f32, ov::Shape{1, 1, 1, context_size}); + auto * data_ptr = input_tensor.data(); + std::copy(padded_data.begin(), padded_data.begin() + context_size, data_ptr); + return input_tensor; + } + + return get_ov_input_tensor(ggml_decoder, param_name); +} + +ov::Tensor get_ov_input_tensor_static_prefill(std::shared_ptr ggml_decoder, + const std::string & param_name, + int chunk_index) { + // NPU prompt processing stage + const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name); + const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor); + + const size_t input_len = ggml_decoder->get_input_len(); + const size_t chunk_size = ggml_decoder->m_prefill_chunk_size; + const size_t chunk_valid_size = std::min(chunk_size, input_len - chunk_index * chunk_size); + const size_t chunk_pad_size = chunk_size - chunk_valid_size; + + if (GgmlOvDecoder::is_inp_tok(ggml_tensor, op) || GgmlOvDecoder::is_inp_pos(ggml_tensor, op) || + GgmlOvDecoder::is_kv_idx(ggml_tensor, op)) { + ov::Shape input_shape = {1, 1, 1, chunk_size}; + ov::Tensor input_tensor(ggml_decoder->get_ov_type(ggml_tensor), input_shape); + // copy the chunk_index-th chunk from ggml_tensor + size_t element_size = ggml_type_size(ggml_tensor->type); + void * input_data = (char *) ggml_tensor->data + chunk_index * chunk_size * element_size; + std::memcpy(input_tensor.data(), input_data, chunk_valid_size * element_size); + // pad the rest with last_value + 1, so that kv's of padded positions are inserted + // to the next row after the valids row in the kvcache + if (chunk_pad_size > 0) { + if (ggml_tensor->type == GGML_TYPE_I32) { + int32_t last_value = + *((int32_t *) ggml_tensor->data + (chunk_index * chunk_size + chunk_valid_size - 1)); + int32_t * output_data = input_tensor.data(); + std::fill(output_data + chunk_valid_size, output_data + chunk_size, last_value + 1); + } else if (ggml_tensor->type == GGML_TYPE_I64) { + int64_t last_value = + *((int64_t *) ggml_tensor->data + (chunk_index * chunk_size + chunk_valid_size - 1)); + int64_t * output_data = input_tensor.data(); + std::fill(output_data + chunk_valid_size, output_data + chunk_size, last_value + 1); + } else { + throw std::runtime_error("Unexpected tensor type for " + param_name); + } + } + return input_tensor; + } + + if (GgmlOvDecoder::is_output_idx(ggml_tensor, op)) { + size_t output_len = ggml_decoder->get_compute_params().output_len; + ov::Shape input_shape = {1, 1, 1, output_len}; + ov::Tensor input_tensor(ggml_decoder->get_ov_type(ggml_tensor), input_shape); + if (ggml_tensor->ne[0] == 0) { + *input_tensor.data() = 0; + } else { + auto * data_addr = input_tensor.data(); + for (size_t i = 0; i < output_len; i++) { + data_addr[i] = ((int32_t *) ggml_tensor->data)[i] % chunk_size; + } + } + return input_tensor; + } + + if (GgmlOvDecoder::is_inp_mask(ggml_tensor, op)) { + size_t cols = ggml_tensor->ne[0]; + size_t rows = ggml_tensor->ne[1]; + float * ggml_data = (float *) ggml_tensor->data + chunk_index * chunk_size * cols; + size_t chunk_valid_rows = std::min(chunk_size, rows - chunk_index * chunk_size); + size_t context_size = ggml_decoder->get_ctx_size(); + std::vector padded_data = + pad_input(ggml_data, chunk_valid_rows, cols, chunk_size, context_size, -INFINITY); + set_zero_diagonal(padded_data, chunk_size, context_size); + ov::Tensor input_tensor(ov::element::f32, ov::Shape{1, 1, chunk_size, context_size}); + auto * data_ptr = input_tensor.data(); + std::copy(padded_data.begin(), padded_data.begin() + chunk_size * context_size, data_ptr); + return input_tensor; + } + + return get_ov_input_tensor(ggml_decoder, param_name); +} + +size_t checksum(const void * data, size_t size) { + const uint8_t * bytes = static_cast(data); + size_t sum = 0; + for (size_t i = 0; i < size; ++i) { + sum += (uint8_t) i; + sum += bytes[i]; + } + return sum; +} + +void print_input_tensor_info(const std::string & name, const ov::Tensor & tensor) { + std::cout << "Input name: " << name << ", Input shape: " << tensor.get_shape() << ", Address: " << tensor.data() + << std::endl; + switch (tensor.get_element_type()) { + case ov::element::f32: { + if (name.find("self_kq_mask") == std::string::npos) { + std::cout << *(tensor.data()) << std::endl; + } else { + size_t rows = tensor.get_shape()[2]; + size_t cols = tensor.get_shape()[3]; + auto * data = tensor.data(); + for (size_t i = 0; i < rows; ++i) { + for (size_t j = 0; j < cols; ++j) { + float val = data[i * cols + j]; + if (std::isinf(val) && val < 0) { + std::cout << std::setw(5) << "-inf"; + } else { + std::cout << std::setw(5) << val; + } + } + std::cout << std::endl; + } + } + + break; + } + case ov::element::f16: + std::cout << *(tensor.data()) << std::endl; + break; + case ov::element::i32: + for (size_t i = 0; i < tensor.get_size(); ++i) { + std::cout << tensor.data()[i] << " "; + } + std::cout << std::endl; + break; + case ov::element::i64: + for (size_t i = 0; i < tensor.get_size(); ++i) { + std::cout << tensor.data()[i] << " "; + } + std::cout << std::endl; + break; + default: + break; + } +} + +void print_output_tensor_info(const std::string & name, const ov::Tensor & tensor, const void * output_dst) { + std::cout << "Output name: " << name << ", Output shape: " << tensor.get_shape() << ", Address: " << output_dst + << std::endl; + + auto print_float_stats = [](const std::string & type_name, size_t size, auto get_value) { + if (size == 0) { + return; + } + + float first = get_value(0); + float min = first; + float max = first; + double sum = first; + + for (size_t i = 1; i < size; ++i) { + float v = get_value(i); + if (v < min) { + min = v; + } + if (v > max) { + max = v; + } + sum += v; + } + double mean = sum / size; + + std::cout << std::right << std::setw(6) << type_name << std::right << std::setw(12) << "First" << std::setw(12) + << "Min" << std::setw(12) << "Max" << std::setw(12) << "Mean" << std::endl; + std::cout << std::right << std::setw(6) << "" << std::right << std::setw(12) << first << std::setw(12) << min + << std::setw(12) << max << std::setw(12) << mean << std::endl; + }; + + switch (tensor.get_element_type()) { + case ov::element::f32: { + const float * data = tensor.data(); + size_t size = tensor.get_size(); + print_float_stats("[f32]", size, [data](size_t i) { return data[i]; }); + break; + } + case ov::element::f16: { + const ov::float16 * data = tensor.data(); + size_t size = tensor.get_size(); + print_float_stats("[f16]", size, [data](size_t i) { return static_cast(data[i]); }); + break; + } + default: + break; + } +} + +void set_zero_diagonal(std::vector & matrix, size_t rows, size_t cols) { + for (size_t i = 0; i < rows; ++i) { + size_t diag_col = std::min(i, cols - 1); + matrix[i * cols + diag_col] = 0.0f; + } +} + +const ggml_tensor * get_inp_pos_tensor(ggml_cgraph * cgraph) { + for (int i = 0; i < cgraph->n_nodes; ++i) { + auto * op = cgraph->nodes[i]; + for (int j = 0; j < GGML_MAX_SRC; ++j) { + auto * src = op->src[j]; + if (src == nullptr) { + break; + } + if (GgmlOvDecoder::is_inp_pos(src, op)) { + return src; + } + } + } + GGML_LOG_ERROR("get_inp_pos_tensor: inp_pos not found in cgraph"); + throw std::runtime_error("get_inp_pos_tensor: inp_pos not found in cgraph"); +} + +bool get_is_prefill(const ggml_tensor * inp_pos) { + return inp_pos->ne[0] > 1; +} + +#pragma GCC diagnostic pop + + +/* Inlined ggml-decoder.cpp */ +#include "ggml.h.inc" +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +#include +#include +#include +#include +#include +#include +#include + +struct ModelParams { + int ctx = -1; + int ctx_swa = -1; + int ctx_per_seq = -1; + int ctx_per_seq_swa = -1; + int n_seq = 1; + int n_heads = -1; + int n_heads_kv = -1; + int head_size = -1; + int32_t rope_params[15]; + std::vector swa_layers; + + std::vector kv_names; + size_t kv_buffer_ctx_id = 0; + + bool same_rope_params(const ModelParams & other) const { + return memcmp(rope_params, other.rope_params, sizeof(int32_t) * 15) == 0; + } + + bool can_reuse_dynamically(const ModelParams & other) const { return same_rope_params(other); } + + bool can_reuse_statically(const ModelParams & other) const { return same_rope_params(other) && ctx == other.ctx; } + + bool kv_buffer_changed(const ModelParams & other) const { return kv_buffer_ctx_id != other.kv_buffer_ctx_id; } +}; + +struct ComputeParams { + int n_seq_active = 1; + int seq_active_start = 0; + int attention_size = -1; + int attention_size_swa = -1; + int input_len = -1; + int token_len_per_seq = -1; + int past_kv_len = -1; + int output_len = 1; +}; + +class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder { +public: + struct NodeInfo { + ggml_tensor * node; + std::string node_name; + std::string node_op_type; + std::map node_inputs; + std::vector node_inputs_names; + ggml_tensor * node_output; + std::string node_output_name; + int node_op_case = 0; + void * data_addr; + }; + // Graph decoder + GgmlOvDecoder(ggml_cgraph * cgraph, + ModelParams & model_params, + ComputeParams & compute_params, + std::map> & model_weights, + bool is_static, + bool is_stateful = false, + bool is_prefill = false, + int prefill_chunk_size = 256); + + // Naive graph decoder + GgmlOvDecoder(ggml_cgraph * cgraph, std::map> & model_weights); + + virtual ov::Any get_attribute(const std::string & name) const override { + return nullptr; + GGML_UNUSED(name); + } + + virtual ov::PartialShape get_input_shape(int node_idx, const std::string & name) const override; + + virtual std::vector get_input_stride(int node_idx, const std::string & name) const override; + + virtual ov::element::Type get_input_type(int node_idx, const std::string & name) const override; + + virtual size_t get_input_size() const override; + + virtual size_t get_input_size(int node_idx) const override; + + virtual void get_input_node(size_t input_port_idx, + std::string & producer_name, + std::string & producer_output_port_name, + size_t & producer_output_port_index) const override { + GGML_UNUSED(input_port_idx); + GGML_UNUSED(producer_name); + GGML_UNUSED(producer_output_port_name); + GGML_UNUSED(producer_output_port_index); + } + + virtual std::vector get_input_names(int node_idx) const override; + + virtual ov::PartialShape get_output_shape(int node_idx) const override; + + virtual ov::element::Type get_output_type(int node_idx) const override; + + virtual int32_t * get_input_op_params(int node_idx, const std::string & name) const override; + + virtual int32_t * get_output_op_params(int node_idx) const override; + + virtual std::vector get_output_names(int node_idx) const override; + + virtual const std::string & get_op_type() const override; + + virtual const std::string & get_op_type(int node_idx) const override; + + virtual const std::string & get_op_name() const override; + + virtual const std::string & get_op_name(int node_idx) const override; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const override; + + ggml_tensor * get_input_ggml_tensor(const std::string & name) const { return m_inputs.at(name); } + + virtual int get_op_case(int node_idx) const override { return m_node_info_list[node_idx].node_op_case; } + + virtual const std::map> & get_model_inputs() const override { + return m_model_inputs; + } + + virtual const std::map> & get_model_extra_inputs() const override { + return m_model_extra_inputs; + } + + virtual const std::map> & get_model_extra_input_values() const { + return m_model_extra_input_values; + } + + virtual const std::map> & get_model_weights() const override { + return m_model_weights; + } + + virtual std::vector get_model_output_names() const override { + return m_model_output_names; + } + + const std::map & get_model_outputs() const { return m_model_outputs; } + + virtual int get_ctx_size() const { return m_model_params.ctx; } + + virtual int get_ctx_swa_size() const { return m_model_params.ctx_swa; } + + virtual int get_ctx_per_seq() const { return m_model_params.ctx_per_seq; } + + virtual int get_ctx_per_seq_swa() const { return m_model_params.ctx_per_seq_swa; } + + virtual int get_n_seq() const { return m_model_params.n_seq; } + + virtual int is_swa_layer(int layer) const override { + return std::find(m_model_params.swa_layers.begin(), m_model_params.swa_layers.end(), layer) != + m_model_params.swa_layers.end(); + } + + int get_past_kv_len() const { return m_compute_params.past_kv_len; } + + int get_input_len() const { return m_compute_params.input_len; } + + virtual int32_t * get_rope_params() const override { return const_cast(m_model_params.rope_params); } + + virtual std::map get_kv_param_res_names() const override; + + virtual bool is_static() const override { return m_is_static; } + + virtual bool is_stateful() const override { return m_is_stateful; } + + ov::PartialShape get_graph_input_shape(const ggml_tensor * op, const ggml_tensor * input) const; + + static void dump_cgraph(const ggml_cgraph * cgraph, std::string & filename); + + static std::shared_ptr create_weight_node(ggml_tensor * tensor, bool naive = false); + + static std::map> create_weight_nodes(ggml_cgraph * cgraph, + bool naive = false); + + const ggml_tensor * get_tensor_used_op(const ggml_tensor * tensor) const; + + const ggml_tensor * get_tensor_from_name(const std::string & name) const; + + void clear_model_weights() { m_model_weights.clear(); } + + static std::pair compute_llm_params(ggml_cgraph * cgraph, bool is_static); + + ModelParams get_model_params() const { return m_model_params; } + + ComputeParams get_compute_params() const { return m_compute_params; } + + void set_model_params(const ModelParams & model_params) { m_model_params = model_params; } + + void set_compute_params(const ComputeParams & compute_params) { m_compute_params = compute_params; } + + bool m_is_static = false; + bool m_is_stateful = false; + bool m_is_prefill = false; + bool m_naive = false; + int m_prefill_chunk_size = 0; + + static ov::Shape get_shape(const ggml_tensor * tensor); + static std::vector get_stride(const ggml_tensor * tensor); + static ov::element::Type get_ov_type(const ggml_tensor * tensor); + static std::string compute_op_type(const ggml_tensor * node); + void add_extra_inputs(); + + void update_io(ggml_cgraph * cgraph); + + inline static bool is_inp_tok(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_GET_ROWS && tensor == op->src[1] && op->src[0]->op == GGML_OP_NONE; + } + + inline static bool is_inp_pos(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_ROPE && tensor == op->src[1]; + } + + inline static bool is_inp_emb(const ggml_tensor * tensor, const ggml_tensor * op) { + return tensor->op == GGML_OP_GET_ROWS && op->op == GGML_OP_RMS_NORM; + } + + inline static bool is_inp_mask(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_CPY || (op->op == GGML_OP_FLASH_ATTN_EXT && tensor == op->src[3]); + } + + inline static bool is_rope_freqs_weight(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_ROPE && tensor == op->src[2]; + } + + inline static bool is_kvcache(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_SET_ROWS && op->src[2] == tensor; + } + + inline static bool is_kv_idx(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_SET_ROWS && op->src[1] == tensor; + } + + inline static bool is_output_idx(const ggml_tensor * tensor, const ggml_tensor * op) { + return op->op == GGML_OP_GET_ROWS && tensor == op->src[1] && op->src[0]->op != GGML_OP_NONE; + } + + static std::string get_graph_input_ov_name(const ggml_tensor * tensor, const ggml_tensor * op) { + if (is_inp_tok(tensor, op)) { + return "inp_tokens"; + } + if (is_inp_pos(tensor, op)) { + return "inp_pos"; + } + if (is_inp_emb(tensor, op)) { + return "embd"; + } + if (is_output_idx(tensor, op)) { + return "inp_out_ids"; + } + if (is_inp_mask(tensor, op)) { + return std::string(tensor->name).find("swa") == std::string::npos ? "self_kq_mask" : "self_kq_mask_swa"; + } + return tensor->name; + } + +private: + void set_input_output(); + int compute_op_case(const ggml_tensor * node) const; + bool node_is_used_as_src(const int node_idx); + void compute_model_inputs(); + void compute_model_outputs(); + + void validate_cgraph() const; + + ggml_cgraph * m_cgraph = nullptr; + std::map m_inputs; + + std::map> m_model_inputs; + std::map> m_model_extra_inputs; + std::map> m_model_extra_input_values; + std::map> m_model_weights; + std::map m_model_outputs; + std::vector m_model_output_names; + std::vector m_node_info_list; + + ModelParams m_model_params; + ComputeParams m_compute_params; +}; + +void print_tensor_address_map(const ggml_cgraph * cgraph); + +int extract_layer_from_name(const std::string & name); + + +#include "ggml-backend-impl-defs.inc" +#include "ggml-backend.h.inc" +#include "ggml-openvino.h.inc" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +GgmlOvDecoder::GgmlOvDecoder(ggml_cgraph * cgraph, + ModelParams & model_params, + ComputeParams & compute_params, + std::map> & model_weights, + bool is_static, + bool is_stateful, + bool is_prefill, + int prefill_chunk_size) : + m_is_static(is_static), + m_is_stateful(is_stateful), + m_is_prefill(is_prefill), + m_naive(false), + m_prefill_chunk_size(prefill_chunk_size), + m_cgraph(cgraph), + m_model_weights(model_weights), + m_model_params(model_params), + m_compute_params(compute_params) { + if (auto * env = getenv("GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS"); env && std::string(env) != "0") { +#ifdef _WIN32 + _putenv_s("GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS", ""); +#else + unsetenv("GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS"); +#endif + print_tensor_address_map(cgraph); + } + + validate_cgraph(); + + set_input_output(); + compute_model_inputs(); + compute_model_outputs(); + + for (int node_n = 0; node_n < cgraph->n_nodes; node_n++) { + m_node_info_list[node_n].node_op_case = compute_op_case(m_node_info_list[node_n].node); + m_node_info_list[node_n].node_op_type = compute_op_type(m_node_info_list[node_n].node); + } + + add_extra_inputs(); +} + +void GgmlOvDecoder::update_io(ggml_cgraph * cgraph) { + m_cgraph = cgraph; + m_model_inputs.clear(); + m_model_outputs.clear(); + m_node_info_list.clear(); + set_input_output(); + compute_model_inputs(); + compute_model_outputs(); +} + +GgmlOvDecoder::GgmlOvDecoder(ggml_cgraph * cgraph, std::map> & model_weights) { + m_cgraph = cgraph; + m_model_weights = model_weights; + m_naive = true; + set_input_output(); + compute_model_inputs(); + compute_model_outputs(); + for (int node_n = 0; node_n < cgraph->n_nodes; node_n++) { + m_node_info_list[node_n].node_op_case = compute_op_case(m_node_info_list[node_n].node); + m_node_info_list[node_n].node_op_type = compute_op_type(m_node_info_list[node_n].node); + } +} + +void GgmlOvDecoder::set_input_output() { + for (int node_n = 0; node_n < m_cgraph->n_nodes; node_n++) { + auto node = m_cgraph->nodes[node_n]; + + NodeInfo current_node_info; + auto node_name = std::string(node->name); + auto node_output_name = node_name; + auto * node_output = node; + if (node->op == GGML_OP_SET_ROWS) { + // SET_ROWS updates the tensor in place. For later ov op that uses the + // the view_src of SET_ROWS, we need to make sure they get the updated tensor + // by putting the view_src name in the tensor_map in + // /src/frontends/ggml/src/translate_session.cpp + node_output_name = std::string(node->view_src->name); + node_output = node->view_src; + } + + current_node_info.node = node; + current_node_info.node_name = node_name; + current_node_info.node_output = node_output; + current_node_info.node_output_name = node_output_name; + current_node_info.node_op_case = 0; + current_node_info.data_addr = node->data; + + for (int i = 0; i < GGML_MAX_SRC; i++) { + auto * src = node->src[i]; + if (src == nullptr) { + continue; + } + auto src_name = std::string(src->name); + if (src->flags & GGML_TENSOR_FLAG_INPUT) { + src_name = get_graph_input_ov_name(src, node); + } + current_node_info.node_inputs[src_name] = src; + current_node_info.node_inputs_names.push_back(src_name); + } + + m_node_info_list.push_back(current_node_info); + } +} + +int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const { + int op_case = 0; + switch (node->op) { + case GGML_OP_RESHAPE: { + auto * src = node->src[0]; + if (src->op == GGML_OP_RESHAPE && src->src[0]->ne[0] == node->ne[0] && src->src[0]->ne[1] == node->ne[1]) { + op_case = 4; + } else if (node->ne[0] * node->ne[1] == src->ne[0]) { + op_case = 1; + } else if (src->ne[0] * src->ne[1] == node->ne[0]) { + op_case = 2; + if (src->ne[2] * src->ne[3] == node->ne[1]) { + op_case = 5; + } + } else if (src->ne[0] * src->ne[1] == node->ne[1]) { + op_case = 3; + } else if (src->ne[1] * src->ne[2] == node->ne[1]) { + op_case = 6; + } + break; + } + case GGML_OP_CONT: { + if (node->src[0]->op == GGML_OP_PERMUTE) { + op_case = 1; + } else if (node->src[0]->op == GGML_OP_TRANSPOSE) { + op_case = 2; + } else if (node->src[0]->op == GGML_OP_VIEW) { + op_case = 3; + } + break; + } + case GGML_OP_PERMUTE: { + if (node->src[0]->op != GGML_OP_VIEW) { + op_case = 1; + } else if (node->src[0]->src[0]->op == GGML_OP_NONE) { + // kv cache tensor + std::string src_name(node->view_src->name); + int layer = extract_layer_from_name(src_name); + if (!is_swa_layer(layer)) { + op_case = 2; + } else { + op_case = 3; + } + } else { + // rope'ed query tensor + op_case = 4; + } + break; + } + case GGML_OP_MUL_MAT: { + if (node->src[0]->op == GGML_OP_CONT && node->src[0]->src[0]->op == GGML_OP_TRANSPOSE) { + op_case = 2; + } else if (node->src[0]->op == GGML_OP_VIEW && node->src[1]->op == GGML_OP_VIEW) { + op_case = 3; + } + break; + } + case GGML_OP_GET_ROWS: { + if (node->src[1]->op == GGML_OP_VIEW) { + op_case = 2; + } + break; + } + case GGML_OP_ROPE: { + const int mode = node->op_params[2]; + switch (mode) { + case GGML_ROPE_TYPE_NEOX: { + op_case = 0x00010000; + break; + } + case GGML_ROPE_TYPE_IMROPE: { + op_case = 0x00020000; + break; + } + default: + op_case = 0x00000000; + break; + } + if (node->src[0]->op == GGML_OP_VIEW) { + op_case = (op_case | 0x00000002); + } + break; + } + case GGML_OP_VIEW: { + if (node->src[0]->op == GGML_OP_VIEW) { + auto * src = node->src[0]; + if (ggml_nelements(node) != ggml_nelements(src)) { + throw std::runtime_error("Unsupported VIEW case"); + } + op_case = 2; + } + { + auto * src = node->src[0]; + if ((ggml_nelements(node) != ggml_nelements(src)) && m_naive) { + // Compare each dimension of node and src, if only one dimension differs then op_case=3 + int diff_count = 0; + for (int i = 0; i < GGML_MAX_DIMS; i++) { + if (node->ne[i] != src->ne[i]) { + diff_count++; + } + } + if (diff_count == 1) { + op_case = 3; + } + } + } + break; + } + default: + break; + } + return op_case; +} + +int extract_layer_from_name(const std::string & name) { + size_t pos1 = name.find("_l"); + assert(pos1 != std::string::npos); + pos1 += 2; + size_t pos2 = name.find(' ', pos1); + if (pos2 == std::string::npos) { + pos2 = name.length(); + } + std::string layer_str = name.substr(pos1, pos2 - pos1); + int layer = std::stoi(layer_str); + return layer; +} + +std::pair GgmlOvDecoder::compute_llm_params(ggml_cgraph * cgraph, bool is_static) { + ModelParams model_params; + ComputeParams compute_params; + for (int i = 0; i < cgraph->n_nodes; i++) { + auto * node = cgraph->nodes[i]; + std::string name = std::string(node->name); + if (node->op == GGML_OP_FLASH_ATTN_EXT) { + model_params.n_heads = node->src[0]->ne[2]; + model_params.n_heads_kv = node->src[1]->ne[2]; + model_params.head_size = node->src[0]->ne[0]; + compute_params.input_len = node->src[0]->ne[1]; + + auto * cache_k_perm = node->src[1]; + if (cache_k_perm->op == GGML_OP_CPY) { + cache_k_perm = cache_k_perm->src[0]; + } + assert(cache_k_perm->op == GGML_OP_PERMUTE); + auto * cache_k_view = cache_k_perm->src[0]; + assert(cache_k_view->op == GGML_OP_VIEW); + + auto * cache_k = cache_k_view->src[0]; + int layer = extract_layer_from_name(cache_k->name); + auto * mask = node->src[3]; + std::string mask_name(mask->name); + + model_params.kv_buffer_ctx_id = ggml_backend_openvino_buffer_get_ctx_id(cache_k->buffer); + if (mask_name.find("swa") != std::string::npos) { + model_params.swa_layers.push_back(layer); + model_params.ctx_per_seq_swa = cache_k->ne[1]; + } else { + model_params.ctx_per_seq = cache_k->ne[1]; + model_params.n_seq = cache_k->ne[2]; + } + + compute_params.n_seq_active = mask->ne[3]; + auto seq_size = cache_k->ne[0] * cache_k->ne[1] * ggml_type_size(cache_k->type); + size_t offset; + memcpy(&offset, cache_k_view->op_params, sizeof(size_t)); + compute_params.seq_active_start = offset / seq_size; + compute_params.token_len_per_seq = node->ne[2]; + + if (mask_name.find("swa") != std::string::npos) { + compute_params.attention_size_swa = mask->ne[0]; + } else { + compute_params.attention_size = mask->ne[0]; + } + if (is_static) { + compute_params.attention_size = model_params.ctx_per_seq; + compute_params.attention_size_swa = model_params.ctx_per_seq_swa; + compute_params.token_len_per_seq = 1; + } + break; + } + if (node->op == GGML_OP_ROPE) { + memcpy(model_params.rope_params, node->op_params, sizeof(int32_t) * 15); + } + } + auto * output_tensor = cgraph->nodes[cgraph->n_nodes - 1]; + compute_params.output_len = output_tensor->ne[1]; + // for NPU, output_len is always 1 except for llama-perplexity + if (is_static && compute_params.output_len == 0) { + compute_params.output_len = 1; + } + model_params.ctx = model_params.ctx_per_seq * model_params.n_seq; + model_params.ctx_swa = model_params.ctx_per_seq_swa * model_params.n_seq; + return {model_params, compute_params}; +} + +void GgmlOvDecoder::validate_cgraph() const { + if (m_model_params.n_seq > 1 && m_is_static == true) { + throw std::runtime_error("n_seq > 1 is not supported on NPU. Try setting -np 1."); + } +} + +ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * op, const ggml_tensor * input) const { + if (m_naive) { + return input!= nullptr ? ov::PartialShape{get_shape(input)} : ov::PartialShape{get_shape(op)}; + } + auto name = std::string(input->name); + ov::PartialShape input_shape; + + if (is_inp_tok(input, op) || is_inp_pos(input, op)) { + // tokens or positions + int len = m_is_static ? (m_is_prefill ? m_prefill_chunk_size : 1) : -1; + input_shape = ov::PartialShape{1, 1, 1, len}; + + } else if (is_output_idx(input, op)) { + // output index + input_shape = ov::PartialShape{1, 1, 1, m_is_static ? m_compute_params.output_len : -1}; + + } else if (is_inp_mask(input, op)) { + // mask + if (m_is_static) { + input_shape = ov::PartialShape{1, 1, m_is_prefill ? m_prefill_chunk_size : 1, m_model_params.ctx}; + } else if (m_is_stateful) { + input_shape = ov::PartialShape{1, 1, -1, -1}; + } else { + input_shape = ov::PartialShape{-1, 1, -1, -1}; + } + + } else if (is_kvcache(input, op)) { + // kvcache + input_shape = ov::PartialShape{get_shape(input)}; + if (!m_is_static) { + // do not fix ctx size to make llama-bench work across test params + input_shape[2] = -1; + } + if (is_stateful()) { + // Convert stateless KV cache layout [1, 1, seq, n_heads_kv * head_size] + // to stateful layout [1, seq, n_heads_kv, head_size]. + assert(input_shape.size() == 4 && input_shape[0] == 1 && input_shape[1] == 1 && + input_shape[2].is_dynamic() && + input_shape[3] == (m_model_params.n_heads_kv * m_model_params.head_size)); + input_shape = {input_shape[0], ov::Dimension::dynamic(), m_model_params.n_heads_kv, + m_model_params.head_size}; + } + + } else if (is_kv_idx(input, op)) { + // kv update index + int len = m_is_static ? (m_is_prefill ? m_prefill_chunk_size : 1) : -1; + input_shape = ov::PartialShape{1, 1, 1, len}; + + } else { + input_shape = ov::PartialShape{get_shape(input)}; + } + return input_shape; +} + +void GgmlOvDecoder::add_extra_inputs() { + // Extra inputs: + // 1. `attention_size`, used in FLASH_ATTN where the shape of the matmul's are 256 aligned, + // see llama_kv_cache_unified::get_n_kv and llama_kv_cache_unified::get_padding. + // 2. `n_seq_active` and `seq_active_start`, used in FLASH_ATTN_EXT to indicate the active sequences in the batch + + auto create_1d_input = [this](const std::string & name, int64_t value) { + if (m_is_static) { + auto constant = + std::make_shared(ov::element::i64, ov::Shape{1}, std::vector{value}); + constant->set_friendly_name(name); + m_model_extra_inputs[name] = constant; + } else { + auto param_node = std::make_shared(ov::element::i64, ov::Shape{1}); + param_node->set_friendly_name(name); + param_node->output(0).get_tensor().set_names({name}); + m_model_extra_inputs[name] = param_node; + + auto tensor = std::make_shared(ov::element::i64, ov::Shape{1}); + *tensor->data() = value; + m_model_extra_input_values[name] = tensor; + } + }; + + create_1d_input("attention_size", m_compute_params.attention_size); + if (m_compute_params.attention_size_swa != -1) { + create_1d_input("attention_size_swa", m_compute_params.attention_size_swa); + } + create_1d_input("n_seq_active", m_compute_params.n_seq_active); + create_1d_input("seq_active_start", m_compute_params.seq_active_start); + create_1d_input("seq_active_end", m_compute_params.seq_active_start + m_compute_params.n_seq_active); + create_1d_input("token_len_per_seq", m_compute_params.token_len_per_seq); + // create_1d_input("token_len", m_token_len_per_seq * m_n_seq_active); +} + +bool GgmlOvDecoder::node_is_used_as_src(const int node_idx) { + ggml_tensor * node = m_cgraph->nodes[node_idx]; + for (int i = node_idx; i < m_cgraph->n_nodes; i++) { + ggml_tensor * other_node = m_cgraph->nodes[i]; + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (other_node->src[j] == node) { + return true; + } + } + } + return false; +} + +void GgmlOvDecoder::compute_model_inputs() { + m_model_inputs.clear(); + m_inputs.clear(); + for (int i = 0; i < m_cgraph->n_nodes; i++) { + ggml_tensor * node = m_cgraph->nodes[i]; + // the node op is NONE means this node maybe as input of later nodes, we should add it to model inputs for this node. + if (node->op == GGML_OP_NONE && node_is_used_as_src(i)) { + std::string node_name(node->name); + if (m_model_weights.find(node_name) == m_model_weights.end()) { + m_inputs[node_name] = node; + auto param_node = + std::make_shared(get_ov_type(node), get_graph_input_shape(node, nullptr)); + param_node->set_friendly_name(node_name); + param_node->output(0).get_tensor().set_names({node_name}); + m_model_inputs[node_name] = param_node; + } + continue; + } + for (int i = 0; i < GGML_MAX_SRC; i++) { + auto * src = node->src[i]; + if (src == nullptr) { + continue; + } + std::string src_name = std::string(src->name); + if (src->flags & GGML_TENSOR_FLAG_INPUT) { + src_name = get_graph_input_ov_name(src, node); + } + if (m_model_weights.find(src_name) != m_model_weights.end()) { + continue; + } + + bool is_intermediate_node = false; + for (const auto & node_info : m_node_info_list) { + if (node_info.node == src) { + is_intermediate_node = true; + break; + } + } + if (is_intermediate_node) { + continue; + } + if (m_model_inputs.find(src_name) != m_model_inputs.end()) { + continue; + } + + m_inputs[src_name] = src; + + ggml_backend_buffer * buffer = src->buffer; + // GGML_BACKEND_BUFFER_USAGE_ANY are kv caches + if (buffer->usage == GGML_BACKEND_BUFFER_USAGE_ANY) { + if (auto it = std::find(m_model_params.kv_names.begin(), m_model_params.kv_names.end(), src_name); + it == m_model_params.kv_names.end()) { + m_model_params.kv_names.push_back(src_name); + } + } + ov::PartialShape param_shape = get_graph_input_shape(node, src); + auto param_node = std::make_shared(get_ov_type(src), param_shape); + param_node->set_friendly_name(src_name); + param_node->output(0).get_tensor().set_names({src_name}); + m_model_inputs[src_name] = param_node; + } + } +} + +void GgmlOvDecoder::compute_model_outputs() { + m_model_outputs.clear(); + m_model_output_names.clear(); + for (int node_n = 0; node_n < m_cgraph->n_nodes; node_n++) { + auto * cur_node = m_cgraph->nodes[node_n]; + // if the node op is NONE means this node is not used at all, we can skip it directly without adding to model outputs. + if (cur_node->op == GGML_OP_NONE) { + continue; + } + auto cur_node_use_count = m_cgraph->use_counts[ggml_hash_find(&m_cgraph->visited_hash_set, cur_node)]; + if (cur_node_use_count == 0) { + // The output of SET_ROWS is the view_src tensor, which is updated in place. We should use the view_src name as the output name to make sure it can be correctly matched with the later ops that use the view_src. + if (cur_node != nullptr && cur_node->op == GGML_OP_SET_ROWS) { + cur_node = cur_node->view_src; + } + } else { + int input_use_count = 0; + for (int i = 0; i < m_cgraph->n_nodes; i++) { + ggml_tensor * node = m_cgraph->nodes[i]; + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (node->src[j] != NULL && node->src[j] == cur_node) { + input_use_count++; + } + } + } + if (input_use_count == cur_node_use_count) { + cur_node = nullptr; + } + } + if (cur_node != nullptr) { + std::string node_output_name(cur_node->name); + m_model_outputs[node_output_name] = cur_node; + m_model_output_names.push_back(node_output_name); + } + } +} + +const ggml_tensor * GgmlOvDecoder::get_tensor_used_op(const ggml_tensor * tensor) const { + if (tensor == nullptr) { + return nullptr; + } + for (int i = 0; i < m_cgraph->n_nodes; i++) { + const auto * node = m_cgraph->nodes[i]; + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (node->src[j] == tensor) { + return node; + } + } + } + return nullptr; +} + +const ggml_tensor * GgmlOvDecoder::get_tensor_from_name(const std::string & name) const { + for (int i = 0; i < m_cgraph->n_nodes; i++) { + const auto * node = m_cgraph->nodes[i]; + for (int j = 0; j < GGML_MAX_SRC; j++) { + const auto * src = node->src[j]; + if (src == nullptr) { + break; + } + if (std::string(src->name) == name) { + return src; + } + } + } + return nullptr; +} + +std::map GgmlOvDecoder::get_kv_param_res_names() const { + std::map kv_param_res_names; + for (const auto & name : m_model_params.kv_names) { + kv_param_res_names[name] = name; + } + return kv_param_res_names; +} + +std::map> GgmlOvDecoder::create_weight_nodes(ggml_cgraph * cgraph, bool naive) { + std::map> model_weights; + auto * nodes = cgraph->nodes; + auto n_nodes = cgraph->n_nodes; + for (int node_i = 0; node_i < n_nodes; node_i++) { + auto * node = nodes[node_i]; + for (int i = 0; i < GGML_MAX_SRC; i++) { + auto * src = node->src[i]; + if (src == nullptr) { + continue; + } + + std::string src_name(src->name); + if (is_rope_freqs_weight(src, node)) { + src_name = "rope_freqs.weight"; + } + if (!src->view_src) { + ggml_backend_buffer * buffer = src->buffer; + if (buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS || ggml_is_quantized(src->type)) { + if (model_weights.find(src_name) == model_weights.end()) { + auto weight_node = create_weight_node(src, naive); + weight_node->set_friendly_name(src_name); + model_weights[src_name] = weight_node; + } + } + } + } + } + return model_weights; +} + +std::shared_ptr GgmlOvDecoder::create_weight_node(ggml_tensor * tensor, bool naive) { + const bool is_ov_buffer = ggml_backend_buffer_is_openvino(tensor->buffer); + + // Check if we have a pre-built constant from the OpenVINO backend buffer + // This is set during ggml_backend_openvino_buffer_set_tensor + if (tensor->extra) { + OPENVINO_ASSERT(is_ov_buffer, "Unsupported weight tensor: " + std::string(tensor->name) + + " Possibly this is a cpu backend repacked quantized weights"); + // Cast to our extra base type and check the type + auto * extra_base = static_cast(tensor->extra); + + if (extra_base->type == ggml_openvino_extra_base::Type::WEIGHT) { + // F16/F32/BF16 weight with shared-memory constant + auto * weight_extra = static_cast(tensor->extra); + if (weight_extra->weight_node) { + // GGML_LOG_DEBUG("%s: using pre-built weight node for %s\n", __func__, tensor->name); + return weight_extra->weight_node; + } + } else if (extra_base->type == ggml_openvino_extra_base::Type::QUANTIZED_WEIGHT) { + // Quantized weight with pre-extracted data + auto * quant_extra = static_cast(tensor->extra); + if (quant_extra->weight_node) { + // GGML_LOG_DEBUG("%s: using pre-extracted quantized weight node for %s\n", __func__, tensor->name); + return quant_extra->weight_node; + } + } + } + + // There are three cases where we need to create a new weight node: + // 1. weights are in openvino_host_buffer. Weight loading to host buffer will not trigger backend_buffer_set_tensor + // 2. weights are in cpu/cpu_mapped buffer. On token_embd.weight goes to case 1 or 2, depending on whether mmap or direct_io is used + // 3. test-backend-ops. buffers in test-backend-ops does not set USAGE_WEIGHT so backend_buffer_set_tensor will not create weight node + + // GGML_LOG_DEBUG("%s: creating new weight node for %s\n", __func__, tensor->name); + static const std::set weight_types = {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, + GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, + GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K}; + if (weight_types.find(tensor->type) == weight_types.end()) { + throw std::runtime_error("Unexpected weight tensor type: " + std::string(tensor->name) + " with type " + + ggml_type_name(tensor->type)); + } + + OvWeight ov_weight; + if (ggml_is_quantized(tensor->type)) { + auto use_bias = naive; + if (is_ov_buffer) { + // For quantized weights, copy raw data to a temp buffer first because + // process_weight_tensor reads from data and writes extracted results + // (weights/scales/zp) to output_base_ptr — they would overlap if both + // point to tensor->data. + size_t raw_size = ggml_nbytes(tensor); + std::vector tmp(raw_size); + memcpy(tmp.data(), tensor->data, raw_size); + ov_weight = process_weight_tensor(tensor, tmp.data(), tensor->data, use_bias); + } else { + ov_weight = process_weight_tensor(tensor, tensor->data, nullptr, use_bias); + } + } else { + // For non-quantized weights (F16/F32/BF16), data is already in tensor->data. + // process_weight_tensor will create an ov::Tensor wrapping tensor->data directly. + ov_weight = process_weight_tensor(tensor, tensor->data, tensor->data); + } + + ov_weight.weight_node->set_friendly_name(tensor->name); + if (!is_ov_buffer) { + return ov_weight.weight_node; + } + + ggml_openvino_extra_base * extra; + if (ov_weight.is_quantized()) { + extra = new ggml_openvino_quantized_weight_extra(std::move(ov_weight.weights), std::move(ov_weight.scales), + std::move(ov_weight.zp), ov_weight.weight_node); + } else { + extra = new ggml_openvino_weight_extra(std::move(ov_weight.weights), ov_weight.weight_node); + } + ggml_openvino_buffer_register_extra(tensor, extra); + + return ov_weight.weight_node; +} + +void GgmlOvDecoder::dump_cgraph(const ggml_cgraph * cgraph, std::string & filename) { + std::ofstream file(filename); + if (!file.is_open()) { + std::cerr << "Failed to open file" << std::endl; + return; + } + + file << "=== GRAPH ===\n"; + + // clang-format off + file << "n_nodes = " << cgraph->n_nodes << "\n"; + file << " " << std::setw(3) << "nodes" + << std::setw(15) << "shape" + << std::setw(20) << "op" + << std::setw(20) << "name" + << std::setw(3) << " " + << std::setw(62) << "stride" + << std::setw(20) << "buffer_type" + << "\n"; + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_tensor * node = cgraph->nodes[i]; + + // Get buffer type name + const char * buf_name = "none"; + ggml_backend_buffer_t buf = node->view_src ? node->view_src->buffer : node->buffer; + if (buf) { + buf_name = ggml_backend_buffer_name(buf); + } + + file << " - " << std::setw(3) << i << ": [ " + << std::setw(5) << node->ne[0] << ", " + << std::setw(5) << node->ne[1] << ", " + << std::setw(5) << node->ne[2] << ", " + << std::setw(5) << node->ne[3] << "] " + << std::left << std::setw(20) << ggml_op_name(node->op) << std::right << " " + << std::left << std::setw(45) << node->name << std::right + << std::setw(2) << "[ " + << std::setw(0) << node->nb[0] << ", " + << std::setw(5) << node->nb[1] << ", " + << std::setw(5) << node->nb[2] << ", " + << std::setw(5) << node->nb[3] << "] " + << std::right << std::setw(15) << buf_name << std::right + << "\n"; + + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (auto* src = node->src[i]) { + // Get buffer type name for source + const char * src_buf_name = "none"; + ggml_backend_buffer_t src_buf = src->view_src ? src->view_src->buffer : src->buffer; + if (src_buf) { + src_buf_name = ggml_backend_buffer_name(src_buf); + } + + file << std::setw(10) << " [ " + << std::setw(5) << src->ne[0] << ", " + << std::setw(5) << src->ne[1] << ", " + << std::setw(5) << src->ne[2] << ", " + << std::setw(5) << src->ne[3] << "] " + << std::setw(12) + << i << ": " << std::left << std::setw(12) << ggml_op_name(src->op) << std::right; + file << std::left << std::setw(30) << src->name << std::right + << std::setw(16) << "[ " + << std::setw(0) << src->nb[0] << ", " + << std::setw(5) << src->nb[1] << ", " + << std::setw(5) << src->nb[2] << ", " + << std::setw(5) << src->nb[3] << "] " + << std::right << std::setw(15) << src_buf_name << std::right + << "\n"; + } + } + } + + file << "n_leafs = " << cgraph->n_leafs << "\n"; + for (int i = 0; i < cgraph->n_leafs; i++) { + ggml_tensor * node = cgraph->leafs[i]; + + // Get buffer type name for leaf + const char * leaf_buf_name = "none"; + ggml_backend_buffer_t leaf_buf = node->view_src ? node->view_src->buffer : node->buffer; + if (leaf_buf) { + leaf_buf_name = ggml_backend_buffer_name(leaf_buf); + } + + file << " - " << std::setw(3) << i << ": [ " + << std::setw(5) << node->ne[0] << ", " + << std::setw(5) << node->ne[1] << "] " + << std::setw(8) << ggml_op_name(node->op) << " " + << std::setw(16) << ggml_get_name(node) + << std::setw(20) << leaf_buf_name << "\n"; + } + // clang-format on + file << "========================================\n"; + + file.close(); +} + +void print_tensor_address_map(const ggml_cgraph * cgraph) { + std::map> address_map; + for (int node_n = 0; node_n < cgraph->n_nodes; node_n++) { + auto * node = cgraph->nodes[node_n]; + if (node->data) { + auto it = address_map.find(node->data); + if (it == address_map.end()) { + address_map[node->data] = std::vector(); + } + address_map[node->data].push_back(node->name); + } + } + for (const auto & pair : address_map) { + std::cout << "Address: " << pair.first << std::endl; + for (const auto & name : pair.second) { + std::cout << name << " ; "; + } + std::cout << std::endl << std::endl; + } +} + +ov::Shape GgmlOvDecoder::get_shape(const ggml_tensor * tensor) { + std::vector shape; + for (int i = GGML_MAX_DIMS - 1; i >= 0; --i) { + shape.push_back(static_cast(tensor->ne[i])); + } + return shape; +} + +std::vector GgmlOvDecoder::get_stride(const ggml_tensor * tensor) { + std::vector stride; + for (int i = GGML_MAX_DIMS - 1; i >= 0; --i) { + stride.push_back(static_cast(tensor->nb[i])); + } + return stride; +} + +ov::element::Type GgmlOvDecoder::get_ov_type(const ggml_tensor * tensor) { + switch (tensor->type) { + case GGML_TYPE_F64: + return ov::element::f64; + case GGML_TYPE_F32: + return ov::element::f32; + case GGML_TYPE_F16: + return ov::element::f16; + case GGML_TYPE_BF16: + return ov::element::bf16; + case GGML_TYPE_I8: + return ov::element::i8; + case GGML_TYPE_I16: + return ov::element::i16; + case GGML_TYPE_I32: + return ov::element::i32; + case GGML_TYPE_I64: + return ov::element::i64; + default: + return ov::element::dynamic; + } +} + +ov::PartialShape GgmlOvDecoder::get_input_shape(int node_idx, const std::string & name) const { + return ov::PartialShape(get_shape(m_node_info_list[node_idx].node_inputs.at(name))); +} + +std::vector GgmlOvDecoder::get_input_stride(int node_idx, const std::string & name) const { + return get_stride(m_node_info_list[node_idx].node_inputs.at(name)); +} + +ov::element::Type GgmlOvDecoder::get_input_type(int node_idx, const std::string & name) const { + return get_ov_type(m_node_info_list[node_idx].node_inputs.at(name)); +} + +size_t GgmlOvDecoder::get_input_size() const { + return m_model_inputs.size(); +} + +size_t GgmlOvDecoder::get_input_size(int node_idx) const { + return m_node_info_list[node_idx].node_inputs_names.size(); +} + +std::vector GgmlOvDecoder::get_input_names(int node_idx) const { + return m_node_info_list[node_idx].node_inputs_names; +} + +ov::PartialShape GgmlOvDecoder::get_output_shape(int node_idx) const { + auto * ggml_tensor = m_node_info_list[node_idx].node_output; + return ov::PartialShape(get_shape(ggml_tensor)); +} + +ov::element::Type GgmlOvDecoder::get_output_type(const int node_idx) const { + return get_ov_type(m_node_info_list[node_idx].node); +} + +std::vector GgmlOvDecoder::get_output_names(int node_idx) const { + return {m_node_info_list[node_idx].node_output_name}; +} + +const std::string & GgmlOvDecoder::get_op_name() const { + static const std::string unknown_name = "UNKNOWN_OP_NAME"; + return unknown_name; +} + +const std::string & GgmlOvDecoder::get_op_name(int node_idx) const { + return m_node_info_list[node_idx].node_name; +} + +int32_t * GgmlOvDecoder::get_input_op_params(int node_idx, const std::string & name) const { + return m_node_info_list[node_idx].node_inputs.at(name)->op_params; +} + +int32_t * GgmlOvDecoder::get_output_op_params(int node_idx) const { + return m_node_info_list[node_idx].node->op_params; +} + +void GgmlOvDecoder::visit_subgraph(std::function, int node_idx)> node_visitor) const { + for (int node_idx = 0; node_idx < m_cgraph->n_nodes; node_idx++) { + if (m_cgraph->nodes[node_idx]->op == GGML_OP_NONE) { + continue; + } + node_visitor(std::make_shared(*this), node_idx); + } +} + +std::string GgmlOvDecoder::compute_op_type(const ggml_tensor * node) { + static const std::map ops = { + {GGML_OP_NONE, "GGML_OP_NONE" }, + {GGML_OP_ACC, "GGML_OP_ACC" }, + {GGML_OP_ADD, "GGML_OP_ADD" }, + {GGML_OP_ADD1, "GGML_OP_ADD1" }, + {GGML_OP_CONT, "GGML_OP_CONT" }, + {GGML_OP_DIV, "GGML_OP_DIV" }, + {GGML_OP_DUP, "GGML_OP_DUP" }, + {GGML_OP_GET_ROWS, "GGML_OP_GET_ROWS" }, + {GGML_OP_MUL, "GGML_OP_MUL" }, + {GGML_OP_MUL_MAT, "GGML_OP_MUL_MAT" }, + {GGML_OP_PERMUTE, "GGML_OP_PERMUTE" }, + {GGML_OP_RESHAPE, "GGML_OP_RESHAPE" }, + {GGML_OP_RMS_NORM, "GGML_OP_RMS_NORM" }, + {GGML_OP_ROPE, "GGML_OP_ROPE" }, + {GGML_OP_SCALE, "GGML_OP_SCALE" }, + {GGML_OP_SOFT_MAX, "GGML_OP_SOFT_MAX" }, + {GGML_OP_SUB, "GGML_OP_SUB" }, + {GGML_OP_TRANSPOSE, "GGML_OP_TRANSPOSE" }, + {GGML_OP_VIEW, "GGML_OP_VIEW" }, + {GGML_OP_SET_ROWS, "GGML_OP_SET_ROWS" }, + {GGML_OP_CPY, "GGML_OP_CPY" }, + {GGML_OP_FLASH_ATTN_EXT, "GGML_OP_FLASH_ATTN_EXT"}, + }; + static const std::map unary_ops = { + {GGML_UNARY_OP_ABS, "GGML_UNARY_OP_ABS" }, + {GGML_UNARY_OP_SGN, "GGML_UNARY_OP_SGN" }, + {GGML_UNARY_OP_NEG, "GGML_UNARY_OP_NEG" }, + {GGML_UNARY_OP_STEP, "GGML_UNARY_OP_STEP" }, + {GGML_UNARY_OP_TANH, "GGML_UNARY_OP_TANH" }, + {GGML_UNARY_OP_ELU, "GGML_UNARY_OP_ELU" }, + {GGML_UNARY_OP_RELU, "GGML_UNARY_OP_RELU" }, + {GGML_UNARY_OP_SIGMOID, "GGML_UNARY_OP_SIGMOID" }, + {GGML_UNARY_OP_GELU, "GGML_UNARY_OP_GELU" }, + {GGML_UNARY_OP_GELU_QUICK, "GGML_UNARY_OP_GELU_QUICK" }, + {GGML_UNARY_OP_SILU, "GGML_UNARY_OP_SILU" }, + {GGML_UNARY_OP_HARDSWISH, "GGML_UNARY_OP_HARDSWISH" }, + {GGML_UNARY_OP_HARDSIGMOID, "GGML_UNARY_OP_HARDSIGMOID"}, + {GGML_UNARY_OP_EXP, "GGML_UNARY_OP_EXP" }, + {GGML_UNARY_OP_COUNT, "GGML_UNARY_OP_COUNT" } + }; + static const std::map glu_ops = { + {GGML_GLU_OP_SWIGLU, "GGML_GLU_OP_SWIGLU"}, + {GGML_GLU_OP_GEGLU, "GGML_GLU_OP_GEGLU" }, + {GGML_GLU_OP_REGLU, "GGML_GLU_OP_REGLU" } + }; + + switch (node->op) { + case GGML_OP_UNARY: + return unary_ops.at(ggml_get_unary_op(node)); + case GGML_OP_GLU: + return glu_ops.at(ggml_get_glu_op(node)); + default: + return ops.at(node->op); + } + static const std::string unknown_op = "UNKNOWN_GGML_OP"; + return unknown_op; +} + +const std::string & GgmlOvDecoder::get_op_type(int node_idx) const { + return m_node_info_list[node_idx].node_op_type; +} + +const std::string & GgmlOvDecoder::get_op_type() const { + static const std::string unknown_op = "UNKNOWN_GGML_OP"; + return unknown_op; +} + + +/* Inlined openvino/frontend.cpp */ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class FrontEnd { +public: + using Ptr = std::shared_ptr; + FrontEnd(); + + static std::shared_ptr convert(const InputModel::Ptr& model, bool naive = false); +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class FrontEnd; +class GgmlDecoder; +using ov::frontend::ggml::GgmlDecoder; + +class InputModel : public ov::frontend::InputModel { + friend class ::ov::frontend::ggml::FrontEnd; + +public: + explicit InputModel(const std::shared_ptr& gdecoder); + + const std::shared_ptr& get_model_decoder() const; + +private: + std::shared_ptr m_decoder; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +namespace op { + +#define GGML_OP_CONVERTER(op) OutputVector op(const NodeContext& context) + +GGML_OP_CONVERTER(translate_add); +GGML_OP_CONVERTER(translate_cont); +GGML_OP_CONVERTER(translate_get_rows); +GGML_OP_CONVERTER(translate_mul); +GGML_OP_CONVERTER(translate_mulmat); +GGML_OP_CONVERTER(translate_permute); +GGML_OP_CONVERTER(translate_reshape); +GGML_OP_CONVERTER(translate_rms_norm); +GGML_OP_CONVERTER(translate_rope); +GGML_OP_CONVERTER(translate_scale); +GGML_OP_CONVERTER(translate_unary_silu); +GGML_OP_CONVERTER(translate_unary_gelu); +GGML_OP_CONVERTER(translate_soft_max); +GGML_OP_CONVERTER(translate_transpose); +GGML_OP_CONVERTER(translate_view); +GGML_OP_CONVERTER(translate_glu_swiglu); +GGML_OP_CONVERTER(translate_glu_geglu); +GGML_OP_CONVERTER(translate_set_rows); +GGML_OP_CONVERTER(translate_cpy); +GGML_OP_CONVERTER(translate_flash_attn_ext); + +} // namespace op + +std::unordered_map get_supported_ops(); + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession { +public: + TranslateSession(const frontend::InputModel::Ptr& input_model, + const std::unordered_map& translator_map, bool naive = false); + + std::shared_ptr get_converted_model(); + std::shared_ptr translate_graph(const frontend::InputModel::Ptr& input_model); + +private: + std::shared_ptr apply_transformations(std::shared_ptr model); + const frontend::InputModel::Ptr m_input_model; + const std::unordered_map& m_translator_map; + std::shared_ptr m_ov_model; + bool m_naive; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +FrontEnd::FrontEnd() {} + +std::shared_ptr FrontEnd::convert(const InputModel::Ptr & model, bool naive) { + auto ggml_model = std::dynamic_pointer_cast(model); + FRONT_END_GENERAL_CHECK(ggml_model, "Invalid input model"); + std::shared_ptr converted_model; + const auto & supported_ops = get_supported_ops(); + { + TranslateSession translate_session(model, supported_ops, naive); + converted_model = translate_session.get_converted_model(); + } + return converted_model; +} + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined openvino/input_model.cpp */ +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class FrontEnd; +class GgmlDecoder; +using ov::frontend::ggml::GgmlDecoder; + +class InputModel : public ov::frontend::InputModel { + friend class ::ov::frontend::ggml::FrontEnd; + +public: + explicit InputModel(const std::shared_ptr& gdecoder); + + const std::shared_ptr& get_model_decoder() const; + +private: + std::shared_ptr m_decoder; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + + +namespace ov { +namespace frontend { +namespace ggml { + +InputModel::InputModel(const std::shared_ptr & gdecoder) : m_decoder(gdecoder) {} + +const std::shared_ptr & InputModel::get_model_decoder() const { + return m_decoder; +} + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined openvino/op_table.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +namespace op { + +#define GGML_OP_CONVERTER(op) OutputVector op(const NodeContext& context) + +GGML_OP_CONVERTER(translate_add); +GGML_OP_CONVERTER(translate_cont); +GGML_OP_CONVERTER(translate_get_rows); +GGML_OP_CONVERTER(translate_mul); +GGML_OP_CONVERTER(translate_mulmat); +GGML_OP_CONVERTER(translate_permute); +GGML_OP_CONVERTER(translate_reshape); +GGML_OP_CONVERTER(translate_rms_norm); +GGML_OP_CONVERTER(translate_rope); +GGML_OP_CONVERTER(translate_scale); +GGML_OP_CONVERTER(translate_unary_silu); +GGML_OP_CONVERTER(translate_unary_gelu); +GGML_OP_CONVERTER(translate_soft_max); +GGML_OP_CONVERTER(translate_transpose); +GGML_OP_CONVERTER(translate_view); +GGML_OP_CONVERTER(translate_glu_swiglu); +GGML_OP_CONVERTER(translate_glu_geglu); +GGML_OP_CONVERTER(translate_set_rows); +GGML_OP_CONVERTER(translate_cpy); +GGML_OP_CONVERTER(translate_flash_attn_ext); + +} // namespace op + +std::unordered_map get_supported_ops(); + +} // namespace ggml +} // namespace frontend +} // namespace ov + + + +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +std::unordered_map get_supported_ops() { + using namespace ov::op; + return { + {"GGML_OP_ADD", op::translate_1to1_match_2_inputs }, + {"GGML_OP_ADD1", op::translate_1to1_match_2_inputs }, + {"GGML_OP_CONT", op::translate_cont }, + {"GGML_OP_DIV", op::translate_1to1_match_2_inputs }, + {"GGML_OP_GET_ROWS", op::translate_get_rows }, + {"GGML_OP_MUL", op::translate_1to1_match_2_inputs}, + {"GGML_OP_MUL_MAT", op::translate_mulmat }, + {"GGML_OP_PERMUTE", op::translate_permute }, + {"GGML_OP_RESHAPE", op::translate_reshape }, + {"GGML_OP_RMS_NORM", op::translate_rms_norm }, + {"GGML_OP_ROPE", op::translate_rope }, + {"GGML_OP_SCALE", op::translate_scale }, + {"GGML_OP_SOFT_MAX", op::translate_soft_max }, + {"GGML_OP_SUB", op::translate_1to1_match_2_inputs}, + {"GGML_OP_TRANSPOSE", op::translate_transpose }, + {"GGML_UNARY_OP_GELU", op::translate_unary_gelu }, + {"GGML_UNARY_OP_SILU", op::translate_unary_silu }, + {"GGML_OP_VIEW", op::translate_view }, + {"GGML_GLU_OP_SWIGLU", op::translate_glu_swiglu }, + {"GGML_GLU_OP_GEGLU", op::translate_glu_geglu }, + {"GGML_OP_SET_ROWS", op::translate_set_rows }, + {"GGML_OP_CPY", op::translate_cpy }, + {"GGML_OP_FLASH_ATTN_EXT", op::translate_flash_attn_ext }, + }; +} + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/cont.cpp */ + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_cont(const NodeContext & context) { + num_inputs_check(context, 1, 1); + + int op_case = context.get_op_case(); + FRONT_END_CHECK_IMPLEMENTED(op_case == 1 || op_case == 2 || op_case == 3, "Unsupported CONT case"); + + auto src_shape = context.get_input_shape(0).to_shape(); + auto dst_shape = context.get_output_shape().to_shape(); + ov::Output res; + + if (op_case == 1) { + // The input comes from a PERMUTE + throw std::runtime_error("Code of this case might be outdated"); + dst_shape[1] = -1; + res = std::make_shared( + context.get_input(0), ov::op::v0::Constant::create(ov::element::i64, {dst_shape.size()}, dst_shape), false); + } else if (op_case == 2) { + // The input comes from a TRANSPOSE + return {context.get_input(0)}; + } else { + // The input comes from a VIEW + res = process_view_input(context, 0); + } + + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/cpy.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_cpy(const NodeContext & context) { + auto res = std::make_shared(context.get_input(0), context.get_output_type()); + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/flash_attn_ext.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_flash_attn_ext(const NodeContext & context) { + num_inputs_check(context, 4, 4); + auto q_f32 = context.get_input(0); + auto k = context.get_input(1); + auto v = context.get_input(2); + auto mask = context.get_input(3); + + float * params = reinterpret_cast(context.get_output_op_params()); + float scale = params[0]; + // float max_bias = params[1]; + // float logit_softcap = params[2]; + + auto q = std::make_shared(q_f32, ov::element::f16); + auto scale_node = std::make_shared(ov::element::f16, ov::Shape{}, std::vector{scale}); + + ov::Output mask_sliced, res; + std::string mask_name = "KQ_mask_sliced"; + if (context.get_input_names()[3].find("swa") != std::string::npos) { + mask_name = "KQ_mask_swa_sliced"; + } + if (context.has_input(mask_name)) { + mask_sliced = context.get_input(mask_name); + } else { + auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); + auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto two = ov::op::v0::Constant::create(ov::element::i64, {1}, {2}); + auto token_len = get_dimensions(q, {2}); + mask_sliced = std::make_shared(mask, zero, token_len, one, two); + } + + if (mask_sliced.get_element_type() != ov::element::f16) { + mask_sliced = std::make_shared(mask_sliced, ov::element::f16); + } + + auto tile_kv = [&](int64_t num_heads, int64_t num_heads_kv, int64_t head_size, ov::Output kv) { + int64_t factor = num_heads / num_heads_kv; + if (factor > 1 && num_heads_kv > 1) { + ov::Output kv_broadcast_shape, kv_unsqueezed, new_kv_shape; + auto unsqueeze_axes = ov::op::v0::Constant::create(ov::element::i64, Shape{}, {2}); + kv_unsqueezed = std::make_shared(kv, unsqueeze_axes); + + kv_broadcast_shape = ov::op::v0::Constant::create( + ov::element::i64, {5}, {(int64_t) 1, (int64_t) 1, factor, (int64_t) 1, (int64_t) 1}); + new_kv_shape = + ov::op::v0::Constant::create(ov::element::i64, {4}, {(int64_t) 0, num_heads, (int64_t) -1, head_size}); + + kv = std::make_shared(kv_unsqueezed, kv_broadcast_shape, + ov::op::BroadcastType::BIDIRECTIONAL); + kv = std::make_shared(kv, new_kv_shape, true); + } + return kv; + }; + + auto q_shape = context.get_input_shape(0).to_shape(); + auto k_shape = context.get_input_shape(1).to_shape(); + k = tile_kv(q_shape[1], k_shape[1], q_shape[3], k); + v = tile_kv(q_shape[1], k_shape[1], q_shape[3], v); + + auto sdpa = std::make_shared(q, k, v, mask_sliced, scale_node, false); + res = std::make_shared(sdpa, + ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3})); + res = std::make_shared(res, ov::element::f32); + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/get_rows.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_get_rows(const NodeContext & context) { + num_inputs_check(context, 2, 2); + + int op_case = context.get_op_case(); + + Output res; + auto data = context.get_input(0); + auto indices = context.get_input(1); + + if (op_case == 2) { + // The input comes from a VIEW + indices = process_view_input(context, 1); + } + + // data[1,b,x,y] ind[1,1,b,x'] test-backend-ops case + // data[x,y] ind[1,1,1,x'] normal case + indices = + std::make_shared(indices, ov::op::v0::Constant::create(ov::element::i64, {2}, {0, 1})); + if (data.get_partial_shape().rank() == 4) { + if (!(data.get_partial_shape()[1].is_dynamic()) && data.get_partial_shape()[1].get_length() == 1) { + // Work-around for a bug in ov cpu plugin for test-backend-ops + data = std::make_shared(data, + ov::op::v0::Constant::create(ov::element::i64, {2}, {0, 1})); + auto axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {0}); + res = std::make_shared(data, indices, axis); + } else { + auto axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {1}); + data = + std::make_shared(data, ov::op::v0::Constant::create(ov::element::i64, {1}, {0})); + res = std::make_shared(data, indices, axis, 1); + } + } else if (context.is_stateful() && data.get_partial_shape().rank() == 3) { + auto axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {1}); + res = std::make_shared(data, indices, axis, 1); + } else { + auto axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {0}); + res = std::make_shared(data, indices, axis); + } + + if (res.get_element_type() != context.get_output_type()) { + res = std::make_shared(res, context.get_output_type()); + } + if (!(context.is_stateful())) { + res = std::make_shared(res, ov::op::v0::Constant::create(ov::element::i64, {1}, {0})); + } + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/glu_geglu.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_glu_geglu(const NodeContext & context) { + num_inputs_check(context, 1, 2); + + ov::Output src0; + ov::Output src1; + if (context.get_input_size() == 2) { + src0 = context.get_input(0); + src1 = context.get_input(1); + } else { + // GGML splits along ne[0] (OV last axis) using floor division: nc = ne[0] / 2. + // Both halves are nc elements; if the dimension is odd, the last element is dropped. + // Use Slice instead of Split to handle odd dimensions correctly. + auto combined = context.get_input(0); + auto combined_shape = combined.get_partial_shape(); + int64_t last_dim_val = combined_shape[combined_shape.rank().get_length() - 1].get_length(); + int64_t nc = last_dim_val / 2; + + auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}); + auto step = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto start0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); + auto stop0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc}); + auto start1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc}); + auto stop1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {2 * nc}); + + src0 = std::make_shared(combined, start0, stop0, step, axis); + src1 = std::make_shared(combined, start1, stop1, step, axis); + } + + int32_t * params = context.get_output_op_params(); + const int32_t swapped = params[1]; + if (swapped) { + std::swap(src0, src1); + } + + auto gelu = std::make_shared(src0); + auto res = std::make_shared(gelu, src1); + + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/glu_swiglu.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_glu_swiglu(const NodeContext & context) { + num_inputs_check(context, 1, 2); + + ov::Output src0; + ov::Output src1; + if (context.get_input_size() == 2) { + src0 = context.get_input(0); + src1 = context.get_input(1); + } else { + // GGML splits along ne[0] (OV last axis) using floor division: nc = ne[0] / 2. + // Both halves are nc elements; if the dimension is odd, the last element is dropped. + // Use Slice instead of Split to handle odd dimensions correctly. + auto combined = context.get_input(0); + auto combined_shape = combined.get_partial_shape(); + int64_t last_dim_val = combined_shape[combined_shape.rank().get_length() - 1].get_length(); + int64_t nc = last_dim_val / 2; + + auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}); + auto step = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto start0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); + auto stop0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc}); + auto start1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc}); + auto stop1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {2 * nc}); + + src0 = std::make_shared(combined, start0, stop0, step, axis); + src1 = std::make_shared(combined, start1, stop1, step, axis); + } + + int32_t * params = context.get_output_op_params(); + const int32_t swapped = params[1]; + if (swapped) { + std::swap(src0, src1); + } + + auto sigmoid = std::make_shared(src0); + auto silu = std::make_shared(src0, sigmoid); + auto res = std::make_shared(silu, src1); + + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/mulmat.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_mulmat(const NodeContext & context) { + num_inputs_check(context, 2, 2); + + int op_case = context.get_op_case(); + + ov::Output res; + ov::Output B = context.get_input(0); + ov::Output A = context.get_input(1); + + bool transpose_b = true; + if (op_case == 2) { + B = B.get_node_shared_ptr()->input_value(0); + transpose_b = false; + } else if (op_case == 3) { + B = process_view_input(context, 0); + A = process_view_input(context, 1); + } + if (A.get_element_type() != B.get_element_type()) { + B = std::make_shared(context.get_input(0), context.get_input_type(1)); + } + + auto B_shape = context.get_input_shape(0).to_shape(); + auto A_shape = context.get_input_shape(1).to_shape(); + int64_t A_batch = A_shape[1]; + int64_t B_batch = B_shape[1]; + + auto A_batch_larger = A_batch > B_batch; + auto batch_large = A_batch_larger ? A_batch : B_batch; + auto batch_small = A_batch_larger ? B_batch : A_batch; + + Output Z = A_batch_larger ? B : A; + int64_t factor = batch_large / batch_small; + if (factor > 1 && batch_small > 1) { + auto batch_large_node = ov::op::v0::Constant::create(ov::element::i64, {1}, std::vector{batch_large}); + auto batch_small_node = ov::op::v0::Constant::create(ov::element::i64, {1}, std::vector{batch_small}); + auto factor_node = ov::op::v0::Constant::create(ov::element::i64, {1}, std::vector{factor}); + + auto unsqueeze_axes = ov::op::v0::Constant::create(ov::element::i64, Shape{}, {2}); + auto Z_unsqueezed = std::make_shared(Z, unsqueeze_axes); + + auto broadcast_shape = ov::op::v0::Constant::create( + ov::element::i64, {5}, {(int64_t) 1, (int64_t) 1, factor, (int64_t) 1, (int64_t) 1}); + auto new_Z_shape = ov::op::v0::Constant::create(ov::element::i64, {4}, + {(int64_t) 0, batch_large, (int64_t) -1, (int64_t) A_shape[3]}); + + auto Z_broadcasted = std::make_shared(Z_unsqueezed, broadcast_shape, + ov::op::BroadcastType::BIDIRECTIONAL); + Z = std::make_shared(Z_broadcasted, new_Z_shape, true); + } + if (A_batch_larger) { + B = Z; + } else { + A = Z; + } + + res = std::make_shared(A, B, false, transpose_b); + + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/permute.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_permute(const NodeContext & context) { + num_inputs_check(context, 1, 1); + + int op_case = context.get_op_case(); + FRONT_END_CHECK_IMPLEMENTED(op_case == 1 || op_case == 2 || op_case == 3 || op_case == 4, + "Unsupported PERMUTE case"); + + ov::Output res; + auto src = context.get_input(0); + auto perm = ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3}); + + if (op_case == 1 || context.is_stateful()) { + res = std::make_shared(src, perm); + } else if (op_case == 4) { + auto output_shape = context.get_output_shape().to_shape(); + auto n_heads = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[1]}); + auto head_size = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[3]}); + auto n_seq_active = context.has_input("n_seq_active") ? + context.get_input("n_seq_active") : + ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[0]}); + auto neg_one = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}); + + auto new_shape = + std::make_shared(ov::OutputVector{n_seq_active, neg_one, n_heads, head_size}, 0); + + // // Alternative + // auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); + // auto new_shape = std::make_shared(ov::OutputVector{n_seq_active, neg_one, zero, zero}, 0); + + auto reshaped = std::make_shared(src, new_shape, true); + res = std::make_shared(reshaped, perm); + } else { + auto cache_shape = src.get_partial_shape(); + auto output_shape = context.get_output_shape().to_shape(); + int64_t head_size = output_shape[3]; + int64_t n_heads = output_shape[1]; + int64_t ctx_per_seq = cache_shape[2].is_static() ? cache_shape[2].get_length() : -1; + int64_t n_seq = cache_shape[1].get_length(); + + Output attention_size; + if (!context.has_input("attention_size")) { + attention_size = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[2]}); + } else if (op_case == 2) { + attention_size = context.get_input("attention_size"); + } else { + attention_size = context.get_input("attention_size_swa"); + } + + Output seq_active_start; + Output seq_active_end; + if (context.has_input("seq_active_start")) { + seq_active_start = context.get_input("seq_active_start"); + seq_active_end = context.get_input("seq_active_end"); + } else { + int64_t n_seq_active = output_shape[0]; + size_t offset = *((size_t *) context.get_input_op_params(0)); + int64_t seq_active_start_val = offset / context.get_input_stride(0)[0]; + int64_t seq_active_end_val = seq_active_start_val + n_seq_active; + seq_active_start = ov::op::v0::Constant::create(ov::element::i64, {1}, {seq_active_start_val}); + seq_active_end = ov::op::v0::Constant::create(ov::element::i64, {1}, {seq_active_end_val}); + } + + // 1. reshape to [n_seq, ctx_per_seq, n_heads, head_size] + // 2. slice out the active sequences + // 3. slice out the attention part in each sequence + // 4. permute + auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); + auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + + auto src_reshaped = std::make_shared( + src, ov::op::v0::Constant::create(ov::element::i64, {4}, {n_seq, ctx_per_seq, n_heads, head_size}), false); + auto slice1 = std::make_shared(src_reshaped, seq_active_start, seq_active_end, one, zero); + auto slice2 = std::make_shared(slice1, zero, attention_size, one, one); + res = std::make_shared(slice2, perm); + } + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/reshape.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_reshape(const NodeContext & context) { + num_inputs_check(context, 1, 1); + if (context.get_input_shape(0) == context.get_output_shape()) { + return {context.get_input(0)}; + } + + int op_case = context.get_op_case(); + FRONT_END_CHECK_IMPLEMENTED( + op_case == 1 || op_case == 2 || op_case == 3 || op_case == 4 || op_case == 5 || op_case == 6, + "Unsupported RESHAPE case"); + + auto output_shape = context.get_output_shape().to_shape(); + std::shared_ptr new_shape_node; + if (op_case == 1) { + if (context.is_stateful()) { + new_shape_node = ov::op::v0::Constant::create( + ov::element::i64, {3}, + std::vector{-1, (int64_t) output_shape[2], (int64_t) output_shape[3]}); + } else { + new_shape_node = ov::op::v0::Constant::create( + ov::element::i64, {4}, + std::vector{(int64_t) output_shape[0], -1, (int64_t) output_shape[2], (int64_t) output_shape[3]}); + } + } else if (op_case == 2) { + new_shape_node = ov::op::v0::Constant::create( + ov::element::i64, {4}, + std::vector{(int64_t) output_shape[0], (int64_t) output_shape[1], -1, (int64_t) output_shape[3]}); + + } else if (op_case == 3) { + throw std::runtime_error("might be outdated RESHAPE case"); + new_shape_node = ov::op::v0::Constant::create( + ov::element::i64, {4}, std::vector{(int64_t) output_shape[0], (int64_t) output_shape[1], -1, 1}); + + } else if (op_case == 4) { + return {context.get_input(0).get_node_shared_ptr()->input_value(0)}; + + } else if (op_case == 5) { + if (context.is_stateful()) { + std::vector shape_vec = {1, -1, (int64_t) context.get_output_shape().to_shape()[3]}; + new_shape_node = ov::op::v0::Constant::create(ov::element::i64, {3}, shape_vec); + } else { + std::vector shape_vec = {1, 1, -1, (int64_t) context.get_output_shape().to_shape()[3]}; + new_shape_node = ov::op::v0::Constant::create(ov::element::i64, {4}, shape_vec); + } + + // // Alternative + // auto token_len = context.get_input("token_len"); + // auto emb_size = + // ov::op::v0::Constant::create(ov::element::i64, {1}, {(int64_t) context.get_output_shape().to_shape()[3]}); + // auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + // new_shape_node = std::make_shared(ov::OutputVector{one, one, token_len, emb_size}, 0); + + } else if (op_case == 6) { + new_shape_node = ov::op::v0::Constant::create(ov::element::i64, {4}, context.get_output_shape().to_shape()); + } + auto res = std::make_shared(context.get_input(0), new_shape_node, false); + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/rms_norm.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_rms_norm(const NodeContext & context) { + num_inputs_check(context, 1, 1); + + auto input_node = context.get_input(0); + auto square = std::make_shared( + input_node, ov::op::v0::Constant::create(ov::element::f32, ov::Shape{1}, {2.0f})); + + auto mean = std::make_shared( + square, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {-1}), true); + + float eps; + memcpy(&eps, context.get_output_op_params(), sizeof(float)); + + auto rms = std::make_shared( + std::make_shared(mean, ov::op::v0::Constant::create(ov::element::f32, ov::Shape{1}, {eps}))); + + auto reciprocal = + std::make_shared(ov::op::v0::Constant::create(ov::element::f32, ov::Shape{1}, {1.0f}), rms); + + auto res = std::make_shared(input_node, reciprocal); + + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/rope.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_rope(const NodeContext & context) { + num_inputs_check(context, 2, 3); + + int op_case = context.get_op_case(); + + ov::Output res; + + auto data_node = context.get_input(0).get_node_shared_ptr(); + auto output_shape = context.get_output_shape().to_shape(); + int32_t * op_params = context.get_output_op_params(); + const int mode = (op_case & 0xFFFF0000) >> 16; + op_case = (op_case & 0x0000FFFF); + + constexpr int TYPE_NORMAL = 0; + constexpr int TYPE_NEOX = 1; + constexpr int TYPE_IMROPE = 2; + + Output cos_theta_node; + Output sin_theta_node; + if (context.has_input("rope_cos")) { + cos_theta_node = context.get_input("rope_cos"); + sin_theta_node = context.get_input("rope_sin"); + } else { + auto inp_pos = context.get_input(1).get_node_shared_ptr(); + std::shared_ptr rope_freqs_weight; + if (context.get_input_size() == 3) { + rope_freqs_weight = context.get_input(2).get_node_shared_ptr(); + } + auto sin_cos = make_sin_cos(op_params, inp_pos, rope_freqs_weight, mode == TYPE_IMROPE); + sin_theta_node = sin_cos.first; + cos_theta_node = sin_cos.second; + } + + if (op_case == 2) { + // The input comes from a VIEW + int slice_len = output_shape[2] * output_shape[3]; + data_node = process_view_input(context, 0, slice_len).get_node_shared_ptr(); + if (context.is_stateful()) { + auto data_shape = ov::op::v0::Constant::create( + ov::element::i64, {3}, std::vector{-1, (int64_t) output_shape[2], (int64_t) output_shape[3]}); + data_node = std::make_shared(data_node, data_shape, false); + } else { + auto data_shape = ov::op::v0::Constant::create( + ov::element::i64, {4}, std::vector{1, -1, (int64_t) output_shape[2], (int64_t) output_shape[3]}); + data_node = std::make_shared(data_node, data_shape, false); + } + } + + if (mode == TYPE_NORMAL) { + auto neg_one = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}); + auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); + auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto two = ov::op::v0::Constant::create(ov::element::i64, {1}, {2}); + auto end = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[3]}); + Output even_slice; + Output odd_slice; + int32_t unsqueeze_dim = context.is_stateful() ? 3 : 4; + even_slice = std::make_shared(data_node, zero, end, two, neg_one); + odd_slice = std::make_shared(data_node, one, end, two, neg_one); + + Output first_half = + std::make_shared(std::make_shared(even_slice, cos_theta_node), + std::make_shared(odd_slice, sin_theta_node)); + Output second_half = + std::make_shared(std::make_shared(even_slice, sin_theta_node), + std::make_shared(odd_slice, cos_theta_node)); + + first_half = std::make_shared(first_half, + ov::op::v0::Constant::create(ov::element::i64, {1}, {unsqueeze_dim})); + second_half = std::make_shared(second_half, + ov::op::v0::Constant::create(ov::element::i64, {1}, {unsqueeze_dim})); + auto stack = std::make_shared(OutputVector{first_half, second_half}, unsqueeze_dim); + + auto data_shape = ov::op::v0::Constant::create( + ov::element::i64, {4}, std::vector{1, -1, (int64_t) output_shape[2], (int64_t) output_shape[3]}); + res = std::make_shared(stack, data_shape, false); + } else if (mode == TYPE_NEOX) { + auto data_split = std::make_shared( + data_node, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {-1}), 2); + Output slice_data_node_0 = data_split->outputs()[0]; + Output slice_data_node_1 = data_split->outputs()[1]; + + auto first_half_node = std::make_shared( + std::make_shared(slice_data_node_0, cos_theta_node), + std::make_shared(slice_data_node_1, sin_theta_node)); + + auto second_half_node = std::make_shared( + std::make_shared(slice_data_node_0, sin_theta_node), + std::make_shared(slice_data_node_1, cos_theta_node)); + + res = std::make_shared(ov::OutputVector{first_half_node, second_half_node}, -1); + } else if (mode == TYPE_IMROPE) { + int64_t n_dims = data_node->get_shape()[3]; + auto cos_sin_shape = std::make_shared(ov::element::i64, ov::Shape{4}, std::vector{1,-1,1,(n_dims >> 1)}); + auto cos_reshaped = std::make_shared(cos_theta_node, cos_sin_shape, true); + auto sin_reshaped = std::make_shared(sin_theta_node, cos_sin_shape, true); + + auto split_axis = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {3}); + auto split_a = std::make_shared(data_node, split_axis, 2); + auto x0 = split_a->output(0); + auto x1 = split_a->output(1); + auto mul_a = std::make_shared(x0, cos_reshaped); + auto mul_b = std::make_shared(x1, sin_reshaped); + auto sub = std::make_shared(mul_a, mul_b); + + auto mul_c = std::make_shared(x0, sin_reshaped); + auto mul_d = std::make_shared(x1, cos_reshaped); + auto add = std::make_shared(mul_c, mul_d); + + res = std::make_shared(ov::OutputVector{sub, add}, 3); + } + + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/scale.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_scale(const NodeContext & context) { + num_inputs_check(context, 1, 1); + + float scale; + float bias; + memcpy(&scale, (float *) context.get_output_op_params() + 0, sizeof(float)); + memcpy(&bias, (float *) context.get_output_op_params() + 1, sizeof(float)); + + auto scale_node = std::make_shared(ov::element::f32, ov::Shape{}, std::vector{scale}); + auto scaled = std::make_shared(context.get_input(0), scale_node); + + std::shared_ptr res; + if (bias != 0.0f) { + auto bias_node = + std::make_shared(ov::element::f32, ov::Shape{}, std::vector{bias}); + res = std::make_shared(scaled, bias_node); + } else { + res = scaled; + } + + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/set_rows.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_set_rows(const NodeContext & context) { + num_inputs_check(context, 3, 3); + + auto data = context.get_input(0); + auto indices = context.get_input(1); + auto dst = context.get_input(2); + + data = std::make_shared(data, context.get_output_type()); + + auto dst_shape = context.get_output_shape().to_shape(); + + auto ind_squeezed = + std::make_shared(indices, ov::op::v0::Constant::create(ov::element::i64, {3}, {0, 1, 2})); + auto data_reshaped = std::make_shared( + data, + ov::op::v0::Constant::create(ov::element::i64, {4}, + {(int64_t) 1, (int64_t) 1, (int64_t) -1, (int64_t) dst_shape[3]}), + false); + auto axes = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {2}); + + Output res; + if (context.is_stateful()) { + int concat_axis = 1; + int64_t dim2 = dst.get_partial_shape()[2].get_length(); + int64_t dim3 = dst.get_partial_shape()[3].get_length(); + data = std::make_shared( + data, ov::op::v0::Constant::create(ov::element::i64, {4}, {(int64_t) 1, (int64_t) -1, dim2, dim3}), false); + res = std::make_shared(OutputVector{dst, data}, concat_axis); + } else { + res = std::make_shared(dst, ind_squeezed, data_reshaped, axes); + } + + if (auto dst_reshape = std::dynamic_pointer_cast(dst.get_node_shared_ptr())) { + // Fix the case of multiple sequences, reshape back to original shape [1, n_seq, ctx_per_seq, emb] + // ctx_per_seq is not fixed due to llama-bench compatibility + auto dst_shape_partial = dst_reshape->get_input_partial_shape(0); + std::vector dst_shape = {dst_shape_partial[0].get_length(), dst_shape_partial[1].get_length(), + dst_shape_partial[2].is_static() ? dst_shape_partial[2].get_length() : -1, + dst_shape_partial[3].get_length()}; + res = std::make_shared(res, ov::op::v0::Constant::create(ov::element::i64, {4}, dst_shape), + false); + } + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/softmax.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_soft_max(const NodeContext & context) { + // TODO code is outdated + num_inputs_check(context, 1, 2); + + auto input_node = context.get_input(0).get_node_shared_ptr(); + ov::Output res; + + float scale = 1.0f; + float max_bias = 0.0f; + auto * op_params = context.get_output_op_params(); + memcpy(&scale, (float *) op_params + 0, sizeof(float)); + memcpy(&max_bias, (float *) op_params + 1, sizeof(float)); + auto src0_shape = context.get_input_shape(0).get_shape(); + const uint32_t h = src0_shape[2]; + const uint32_t n_head = src0_shape[0]; + const uint32_t n_head_log2 = 1u << (uint32_t) floor(log2(n_head)); + + const float m0 = powf(2.0f, -(max_bias) / n_head_log2); + const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + const float slope = + (max_bias > 0.0f) ? h < n_head_log2 ? powf(m0, h + 1) : powf(m1, 2 * (h - n_head_log2) + 1) : 1.0f; + + auto scale_node = std::make_shared(ov::element::f32, ov::Shape{}, std::vector{scale}); + auto scaled_input = std::make_shared(input_node, scale_node); + + if (context.get_input_size() < 2) { + res = std::make_shared(scaled_input, 2); + return rename_outputs_with_suffix({res}, context.get_name()); + } + + ov::Output mask_node_sliced; + if (context.has_input("KQ_mask_sliced")) { + mask_node_sliced = context.get_input("KQ_mask_sliced"); + } else { + auto token_len = get_dimensions(input_node, {1}); + auto mask_node = context.get_input(1); + auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); + auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + mask_node_sliced = std::make_shared(mask_node, zero, token_len, one, one); + } + + if (mask_node_sliced.get_element_type() != context.get_output_type()) { + mask_node_sliced = std::make_shared(mask_node_sliced, context.get_output_type()); + } + + Output slope_mask; + if (slope != 1.0f) { + auto slope_node = + std::make_shared(ov::element::f32, ov::Shape{}, std::vector{slope}); + slope_mask = std::make_shared(mask_node_sliced, slope_node); + throw std::runtime_error("Slope != 1.0f in softmax has not been tested, verify it before use."); + } + slope_mask = mask_node_sliced; + + auto input_slope_mask_node = std::make_shared(scaled_input, slope_mask); + + res = std::make_shared(input_slope_mask_node, 2); + + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/transpose.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_transpose(const NodeContext & context) { + num_inputs_check(context, 1, 1); + + auto res = std::make_shared( + context.get_input(0), ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 1, 3, 2})); + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/unary_gelu.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_unary_gelu(const NodeContext & context) { + num_inputs_check(context, 1, 1); + + auto input = context.get_input(0); + auto res = std::make_shared(input); + + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/unary_silu.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" + +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_unary_silu(const NodeContext & context) { + num_inputs_check(context, 1, 1); + + auto input = context.get_input(0); + auto sigmoid = std::make_shared(input); + auto res = std::make_shared(input, sigmoid); + + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined op/view.cpp */ +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include "../utils.h" +#include +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_view(const NodeContext & context) { + num_inputs_check(context, 1, 1); + + if (context.get_op_case() == 2) { + auto dst_shape = context.get_output_shape().to_shape(); + return rename_outputs_with_suffix({process_view_input(context, 0, dst_shape[2] * dst_shape[3])}, + context.get_name()); + } + // op_case 3 + if (context.get_op_case() == 3) { + auto input = context.get_input(0); + auto input_ov_shape = input.get_partial_shape(); + + auto input_llama_shape = context.get_input_shape(0).to_shape(); + + // if the input ov shape size is different from the input llama shape size, it means the input is already reshaped and we need to reshape it back to the original shape before slicing + if (input_ov_shape.size() != input_llama_shape.size()) { + input = std::make_shared(input, ov::op::v0::Constant::create(ov::element::i64, {input_llama_shape.size()}, input_llama_shape), false); + } + + auto dst_shape = context.get_output_shape().to_shape(); + + // find the index of dst_shape that is different from input shape, and use that index to slice the input + int slice_dim = -1; + for (size_t i = 0; i < dst_shape.size(); ++i) { + if (dst_shape[i] != input_llama_shape[i]) { + slice_dim = i; + break; + } + } + + auto begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); + auto end = ov::op::v0::Constant::create(ov::element::i64, {1}, {dst_shape[slice_dim]}); + auto stride = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto axes = ov::op::v0::Constant::create(ov::element::i64, {1}, {slice_dim}); + auto sliced = std::make_shared(input, begin, end, stride, axes); + return {sliced}; + } + return {context.get_input(0)}; +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined pass/fuse_to_sdpa.cpp */ +#include "openvino/pass/matcher_pass.hpp" + +namespace ov { +namespace frontend { +namespace ggml { +namespace pass { + +class FuseToSDPA : public ov::pass::MatcherPass { +public: + OPENVINO_MATCHER_PASS_RTTI("ov::frontend::ggml::pass::FuseToSDPA") + FuseToSDPA(); +}; + +} // namespace pass +} // namespace ggml +} // namespace frontend +} // namespace ov + + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace pass { + +FuseToSDPA::FuseToSDPA() { + // Not maintained since FLASH_ATTN_EXT has replaced this pattern + const auto m_k = ov::pass::pattern::any_input(); + const auto m_q = ov::pass::pattern::any_input(); + const auto m_qk = ov::pass::pattern::wrap_type({m_q, m_k}); + const auto m_qk_f32 = ov::pass::pattern::wrap_type({m_qk}); + const auto m_scale = ov::pass::pattern::any_input(); + const auto m_scaled_qk = ov::pass::pattern::wrap_type({m_qk_f32, m_scale}); + const auto m_mask = ov::pass::pattern::any_input(); + const auto m_masked_qk = ov::pass::pattern::wrap_type({m_scaled_qk, m_mask}); + const auto m_softmax_qk = ov::pass::pattern::wrap_type({m_masked_qk}); + const auto m_softmax_qk_f16 = ov::pass::pattern::wrap_type({m_softmax_qk}); + const auto m_v = ov::pass::pattern::any_input(); + const auto m_qkv = ov::pass::pattern::wrap_type({m_softmax_qk_f16, m_v}); + + const auto callback = [=](ov::pass::pattern::Matcher & m) { + auto & pattern_to_output = m.get_pattern_value_map(); + auto k = pattern_to_output[m_k]; + auto q = pattern_to_output[m_q]; + auto v = pattern_to_output[m_v]; + auto mask = pattern_to_output[m_mask]; + auto scale = pattern_to_output[m_scale]; + + auto mask_f16 = register_new_node(mask, ov::element::f16); + auto scale_f16 = register_new_node(scale, ov::element::f16); + auto sdpa = std::make_shared(q, k, v, mask_f16, scale_f16, false); + + ov::replace_node(m.get_match_root(), sdpa); + ov::copy_runtime_info(m.get_matched_nodes(), sdpa); + + return true; + }; + register_matcher(std::make_shared(m_qkv, "ov::frontend::ggml::pass::FuseToSDPA"), + callback); +} + +} // namespace pass +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined pass/squeeze_matmul.cpp */ +#include "openvino/pass/matcher_pass.hpp" + +namespace ov { +namespace frontend { +namespace ggml { +namespace pass { + +class SqueezeMatmul : public ov::pass::MatcherPass { +public: + OPENVINO_MATCHER_PASS_RTTI("ov::frontend::ggml::pass::SqueezeMatmul") + SqueezeMatmul(); +}; + +} // namespace pass +} // namespace ggml +} // namespace frontend +} // namespace ov + + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace opp = ov::pass::pattern; + +namespace ov { +namespace frontend { +namespace ggml { +namespace pass { + +// For quantized models, NPUW expects the activation to be 3d in DQ(DynamicQuantization) opt, e.g. DQMatMulGQ2i +SqueezeMatmul::SqueezeMatmul() { + auto m_act = opp::any_input(); + auto m_wei = opp::any_input(); + auto m_matmul = opp::wrap_type({m_act, m_wei}); + + const auto callback = [=](ov::pass::pattern::Matcher & m) { + const auto & pattern_map = m.get_pattern_value_map(); + auto matmul_node = + std::dynamic_pointer_cast(pattern_map.at(m_matmul).get_node_shared_ptr()); + auto act = pattern_map.at(m_act); + auto wei = pattern_map.at(m_wei); + auto act_shape = act.get_partial_shape(); + auto wei_shape = wei.get_partial_shape(); + if (act_shape.rank().is_dynamic() || wei_shape.rank().is_dynamic()) { + return false; + } + if (act_shape.rank().get_length() == 4 && wei_shape.rank().get_length() == 2) { + auto axis = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {0}); + auto squeezed_act = std::make_shared(act, axis); + auto new_matmul = std::make_shared(squeezed_act, wei, matmul_node->get_transpose_a(), + matmul_node->get_transpose_b()); + auto unsqueezed_output = std::make_shared(new_matmul, axis); + unsqueezed_output->set_friendly_name(matmul_node->get_friendly_name()); + ov::copy_runtime_info(matmul_node, {squeezed_act, new_matmul, unsqueezed_output}); + ov::replace_node(matmul_node, unsqueezed_output); + return true; + } + return false; + }; + + register_matcher(std::make_shared(m_matmul, "ov::frontend::ggml::pass::SqueezeMatmul"), + callback); +} + +} // namespace pass +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined openvino/translate_session.cpp */ +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class FrontEnd; +class GgmlDecoder; +using ov::frontend::ggml::GgmlDecoder; + +class InputModel : public ov::frontend::InputModel { + friend class ::ov::frontend::ggml::FrontEnd; + +public: + explicit InputModel(const std::shared_ptr& gdecoder); + + const std::shared_ptr& get_model_decoder() const; + +private: + std::shared_ptr m_decoder; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +class GgmlDecoder : public DecoderBase { +public: + virtual ov::Any get_attribute(const std::string& name) const = 0; + + virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; + + virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; + + virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; + + virtual size_t get_input_size() const = 0; + + virtual size_t get_input_size(int node_idx) const = 0; + + virtual void get_input_node(size_t input_port_idx, + std::string& producer_name, + std::string& producer_output_port_name, + size_t& producer_output_port_index) const = 0; + + virtual std::vector get_input_names(int node_idx) const = 0; + + virtual PartialShape get_output_shape(int node_idx) const = 0; + + virtual element::Type get_output_type(const int node_idx) const = 0; + + virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; + + virtual int32_t * get_output_op_params(int node_idx) const = 0; + + virtual std::vector get_output_names(int node_idx) const = 0; + + virtual const std::string& get_op_type() const = 0; + + virtual const std::string& get_op_type(int node_idx) const = 0; + + virtual const std::string& get_op_name() const = 0; + + virtual const std::string& get_op_name(int node_idx) const = 0; + + virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; + + virtual int get_op_case(int node_idx) const = 0; + + virtual const std::map>& get_model_inputs() const = 0; + virtual const std::map>& get_model_extra_inputs() const = 0; + virtual const std::map>& get_model_weights() const = 0; + virtual std::vector get_model_output_names() const = 0; + + virtual int32_t* get_rope_params() const = 0; + + virtual std::map get_kv_param_res_names() const = 0; + + virtual bool is_static() const = 0; + + virtual bool is_stateful() const = 0; + + virtual int is_swa_layer(int layer) const = 0; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession; + +typedef std::map> TensorMap; + +class NodeContext : public frontend::NodeContext { +public: + NodeContext(const std::shared_ptr& decoder, + std::shared_ptr& tensor_map, + int node_idx, + TranslateSession* translate_session = nullptr) + : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), + m_decoder(decoder), + m_tensor_map(tensor_map), + m_node_idx(node_idx), + m_translate_session(translate_session) { + m_input_names = decoder->get_input_names(m_node_idx); + m_output_names = decoder->get_output_names(m_node_idx); + } + + TranslateSession* get_translate_session() const { + return m_translate_session; + } + + const std::vector& get_input_names() const { return m_input_names; } + + size_t get_input_size() const override { + return m_decoder->get_input_size(m_node_idx); + } + + ov::element::Type get_input_type(size_t index) const { + return m_decoder->get_input_type(m_node_idx, m_input_names[index]); + } + + PartialShape get_input_shape(size_t input_index) const { + return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); + } + + std::vector get_input_stride(size_t index) const { + return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); + } + + std::string get_output_name() const { return m_output_names[0]; } + + PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } + + int32_t* get_input_op_params(size_t index) const { + return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); + } + + int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } + + ov::element::Type get_output_type() const { + return m_decoder->get_output_type(m_node_idx); + } + + Output get_input(int idx) const override { + return m_tensor_map->at(m_input_names[idx]); + } + + Output get_input(const std::string& name) const override { + if (m_tensor_map->find(name) == m_tensor_map->end()) { + throw std::runtime_error("'" + name + "' not found in tensor map."); + } + return m_tensor_map->at(name); + } + + bool has_input(const std::string& name) const { + return m_tensor_map->find(name) != m_tensor_map->end(); + } + + const std::string& get_name() const override { + return m_decoder->get_op_name(m_node_idx); + } + + ov::Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); + } + + int get_op_case() const { + return m_decoder->get_op_case(m_node_idx); + } + + bool is_static() const { return m_decoder->is_static(); } + + bool is_stateful() const { return m_decoder->is_stateful(); } + +private: + std::shared_ptr m_decoder; + std::shared_ptr& m_tensor_map; + int m_node_idx; + TranslateSession* m_translate_session; + std::vector m_input_names; + std::vector m_output_names; +}; + +using CreatorFunction = std::function; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +namespace ov { +namespace frontend { +namespace ggml { + +class TranslateSession { +public: + TranslateSession(const frontend::InputModel::Ptr& input_model, + const std::unordered_map& translator_map, bool naive = false); + + std::shared_ptr get_converted_model(); + std::shared_ptr translate_graph(const frontend::InputModel::Ptr& input_model); + +private: + std::shared_ptr apply_transformations(std::shared_ptr model); + const frontend::InputModel::Ptr m_input_model; + const std::unordered_map& m_translator_map; + std::shared_ptr m_ov_model; + bool m_naive; +}; + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +#include "ggml-openvino/openvino/utils.h" +#include "openvino/pass/matcher_pass.hpp" +#include "openvino/core/visibility.hpp" + +#ifdef OPENVINO_STATIC_LIBRARY +# define TRANSFORMATIONS_API +#else +# ifdef IMPLEMENT_OPENVINO_API +# define TRANSFORMATIONS_API OPENVINO_CORE_EXPORTS +# else +# define TRANSFORMATIONS_API OPENVINO_CORE_IMPORTS +# endif // IMPLEMENT_OPENVINO_API +#endif // OPENVINO_STATIC_LIBRARY + +namespace ov { +namespace pass { + +class TRANSFORMATIONS_API MarkCompressedFloatConstants; + +} // namespace pass +} // namespace ov + +class ov::pass::MarkCompressedFloatConstants : public MatcherPass { +public: + OPENVINO_MATCHER_PASS_RTTI("MarkCompressedFloatConstants") + MarkCompressedFloatConstants(); +}; + +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include +#include +#include + +namespace ov { + +/** + * @brief Holds weightless caching attributes of a single constant. + * + * WeightlessCacheAttribute class represents runtime info attribute that holds + * the values of original size of the constant in bytes and the binary offset of the + * constant's data in the weights file used by the weightless caching mechanism. It's + * not copyable in case the data was changed (the original node was replaced by a new + * one produced during the tranformation pipeline) - in that case weightless caching + * can't be used for that constant. + */ +class OPENVINO_API WeightlessCacheAttribute : public RuntimeAttribute { +public: + OPENVINO_RTTI("WeightlessCacheAttribute", "0", RuntimeAttribute) + + WeightlessCacheAttribute() = delete; + + WeightlessCacheAttribute(size_t original_size, size_t bin_offset, ov::element::Type original_dtype) + : original_size(original_size), + bin_offset(bin_offset), + original_dtype(original_dtype) {} + + bool is_copyable() const override; + + size_t original_size; + size_t bin_offset; + ov::element::Type original_dtype; +}; + +} // namespace ov + + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +using namespace ov::op; + +namespace { + +ov::pass::MakeStateful::ParamResPairs get_kv_param_res_pairs( + const std::shared_ptr & model, + const std::map & kv_param_res_names) { + ov::pass::MakeStateful::ParamResPairs pairs; + const auto & params = model->get_parameters(); + const auto & results = model->get_results(); + + for (const auto & param_res : kv_param_res_names) { + const auto & param_name = param_res.first; + const auto & res_name = param_res.second; + + auto param_it = std::find_if(params.begin(), params.end(), [&](const std::shared_ptr & node) { + return node->get_friendly_name() == param_name; + }); + + OPENVINO_ASSERT(param_it != params.end(), "The tensor name ", param_name, + " is not associated with any of " + "Parameters in the network."); + + auto res_it = std::find_if(results.begin(), results.end(), [&](const std::shared_ptr & node) { + return node->get_friendly_name() == res_name; + }); + + OPENVINO_ASSERT(res_it != results.end(), "The tensor name ", res_name, + " is not associated with any of " + "Results in the network."); + + std::shared_ptr param = *param_it; + std::shared_ptr res = *res_it; + pairs.emplace_back(param, res); + } + return pairs; +} + +void add_sliced_mask(TensorMap & tensor_map, GgmlDecoder & ggml_model_decoder) { + + auto create_sliced_mask = [&](const std::string & mask_name, const std::string & sliced_name, bool is_static) { + if ((tensor_map.find(mask_name) != tensor_map.end()) && + (tensor_map.find("token_len_per_seq") != tensor_map.end())) { + auto token_len_per_seq = tensor_map.at("token_len_per_seq").get_node_shared_ptr(); + auto mask = tensor_map.at(mask_name).get_node_shared_ptr(); + std::shared_ptr mask_sliced; + if (is_static) { + mask_sliced = mask; + } else if (ggml_model_decoder.is_stateful()) { + auto zero_2d = ov::op::v0::Constant::create(ov::element::i64, {2}, {0,0}); + auto one_2d = ov::op::v0::Constant::create(ov::element::i64, {2}, {1,1}); + auto zero_1d = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); + auto three_1d = ov::op::v0::Constant::create(ov::element::i64, {1}, {3}); + auto neg_one_1d = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}); + auto axes = ov::op::v0::Constant::create(ov::element::i64, {2}, {-2,-1}); + auto inp_pos = tensor_map.at("inp_pos").get_node_shared_ptr(); + auto gather_inp_pos = std::make_shared(inp_pos, neg_one_1d, three_1d); + auto reshaped_inp_pos = std::make_shared(gather_inp_pos, ov::op::v0::Constant::create(ov::element::i64, {1}, {1}), false); + auto inp_pos_incremented = std::make_shared(reshaped_inp_pos, ov::op::v0::Constant::create(ov::element::i32, ov::Shape{1}, {1})); + auto stop = std::make_shared(ov::OutputVector{token_len_per_seq, std::make_shared(inp_pos_incremented, token_len_per_seq)}, 0); + mask_sliced = + std::make_shared(mask, zero_2d, stop, one_2d, axes); + mask_sliced = std::make_shared(mask_sliced, ov::element::f16); + mask_sliced->set_friendly_name(sliced_name); + } else { + auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); + auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto two = ov::op::v0::Constant::create(ov::element::i64, {1}, {2}); + mask_sliced = std::make_shared(mask, zero, token_len_per_seq, one, two); + mask_sliced = std::make_shared(mask_sliced, ov::element::f16); + mask_sliced->set_friendly_name(sliced_name); + } + tensor_map.insert({sliced_name, mask_sliced->output(0)}); + } + }; + + create_sliced_mask("self_kq_mask", "KQ_mask_sliced", ggml_model_decoder.is_static()); + create_sliced_mask("self_kq_mask_swa", "KQ_mask_swa_sliced", ggml_model_decoder.is_static()); +} + +void add_rope_sin_cos(TensorMap & tensor_map, GgmlDecoder & ggml_model_decoder) { + int32_t * rope_params = ggml_model_decoder.get_rope_params(); + if (tensor_map.find("inp_pos") == tensor_map.end() || rope_params == nullptr) { + return; + } + auto inp_pos = tensor_map.at("inp_pos").get_node_shared_ptr(); + std::shared_ptr rope_freqs_weight; + if (tensor_map.find("rope_freqs.weight") != tensor_map.end()) { + rope_freqs_weight = tensor_map.at("rope_freqs.weight").get_node_shared_ptr(); + } + + auto sin_cos = make_sin_cos(rope_params, inp_pos, rope_freqs_weight); + auto sin_theta = sin_cos.first; + auto cos_theta = sin_cos.second; + + cos_theta.get_node_shared_ptr()->set_friendly_name("rope_cos"); + sin_theta.get_node_shared_ptr()->set_friendly_name("rope_sin"); + tensor_map.insert({"rope_cos", cos_theta}); + tensor_map.insert({"rope_sin", sin_theta}); +} + +// Create common patterns +void preprocess(TensorMap & tensor_map, GgmlDecoder & ggml_model_decoder) { + add_sliced_mask(tensor_map, ggml_model_decoder); + add_rope_sin_cos(tensor_map, ggml_model_decoder); +} + +} // namespace + +TranslateSession::TranslateSession(const frontend::InputModel::Ptr & input_model, + const std::unordered_map & translator_map, + bool naive) : + m_input_model(input_model), + m_translator_map(translator_map), + m_ov_model(nullptr), + m_naive(naive) {} + +std::shared_ptr TranslateSession::get_converted_model() { + if (m_ov_model) { + return m_ov_model; + } + m_ov_model = translate_graph(m_input_model); + return m_ov_model; +} + +std::shared_ptr TranslateSession::translate_graph(const frontend::InputModel::Ptr & input_model) { + ov::ParameterVector params; + ov::ResultVector results; + auto tensor_map = std::make_shared(); + std::shared_ptr resulting_model; + + const auto & ggml_model = std::dynamic_pointer_cast(input_model); + std::shared_ptr ggml_model_decoder = ggml_model->get_model_decoder(); + + for (const auto & it : ggml_model_decoder->get_model_inputs()) { + params.push_back(std::dynamic_pointer_cast(it.second)); + (*tensor_map)[it.first] = it.second; + } + + for (const auto & it : ggml_model_decoder->get_model_extra_inputs()) { + if (std::dynamic_pointer_cast(it.second)) { + params.push_back(std::dynamic_pointer_cast(it.second)); + } + (*tensor_map)[it.first] = it.second; + } + + for (const auto & it : ggml_model_decoder->get_model_weights()) { + (*tensor_map)[it.first] = it.second; + } + + auto node_visitor = [&](std::shared_ptr decoder, int node_idx) { + auto operation_type = decoder->get_op_type(node_idx); + if (operation_type == "GGML_OP_NONE") { + return; + } + + ov::OutputVector converted_outputs; + auto it = m_translator_map.find(operation_type); + FRONT_END_OP_CONVERSION_CHECK(it != m_translator_map.end(), "Translation for operation type ", operation_type, + " is not implemented."); + NodeContext node_context(decoder, tensor_map, node_idx, this); + converted_outputs = it->second(node_context); + + const auto & node_output_names = decoder->get_output_names(node_idx); + FRONT_END_OP_CONVERSION_CHECK(node_output_names.size() == converted_outputs.size(), "Number of ", + operation_type, " outputs greater than number of converted outputs, which are ", + node_output_names.size(), " and ", converted_outputs.size(), " respectively."); + + for (size_t i = 0; i < node_output_names.size(); ++i) { + auto output_name = node_output_names[i]; + if (i < converted_outputs.size() && converted_outputs[i].get_node_shared_ptr() != nullptr) { + (*tensor_map)[output_name] = converted_outputs[i]; + } + } + }; + + if (!m_naive) { + preprocess(*tensor_map, *ggml_model_decoder); + } + ggml_model_decoder->visit_subgraph(node_visitor); + + for (const auto & name : ggml_model_decoder->get_model_output_names()) { + FRONT_END_GENERAL_CHECK(tensor_map->find(name) != tensor_map->end(), + "Output name not found in tensor map: ", name); + auto result = std::make_shared(tensor_map->at(name)); + result->set_friendly_name(name); + results.push_back(result); + } + + ov::ParameterVector used_params; + for (const auto & param : params) { + if (!param->output(0).get_target_inputs().empty()) { + used_params.push_back(param); + } + } + // if (auto diff = params.size() - used_params.size()) { + // GGML_LOG_INFO("%zu parameters are not used in the model.", diff); + // } + resulting_model = std::make_shared(results, used_params); + + apply_transformations(resulting_model); + + // Set WeightlessCacheAttribute on large constants to avoid unnecessary memory copies + // in the NPUW plugin. Without this attribute, NPUW's LazyTensor constructor + // (lazy_tensor.cpp, op::Const::Const) will memcpy every constant "in case export + // occurs", doubling memory usage per compile_model call. + // + // The bin_offset field serves as a unique key (not a real file offset) — this is + // the same convention the GPU plugin uses for non-IR models (see + // Plugin::set_weightless_cache_attributes in intel_gpu/src/plugin/plugin.cpp). + // Each constant must have a distinct bin_offset, otherwise GPU's weightless cache + // import will map multiple constants to the same data. + // + // Small constants (< 16 elements) are excluded since they may be introduced by + // optimization patterns and the overhead is negligible. + size_t offset = 0; + for (auto & node : resulting_model->get_ordered_ops()) { + if (auto cnst = ov::as_type_ptr(node); + cnst && cnst->get_byte_size() / cnst->get_element_type().size() >= 16) { + auto & rt_info = cnst->get_rt_info(); + if (rt_info.find(ov::WeightlessCacheAttribute::get_type_info_static()) == rt_info.end()) { + rt_info[ov::WeightlessCacheAttribute::get_type_info_static()] = + ov::WeightlessCacheAttribute(cnst->get_byte_size(), offset++, cnst->get_element_type()); + } + } + } + return resulting_model; +} + +std::shared_ptr TranslateSession::apply_transformations(std::shared_ptr model) { + auto ggml_model_decoder = std::dynamic_pointer_cast(m_input_model)->get_model_decoder(); + { + ov::pass::Manager manager; + manager.set_per_pass_validation(true); + manager.register_pass(); + + if (ggml_model_decoder->is_stateful()) { + const auto kv_param_res_names = ggml_model_decoder->get_kv_param_res_names(); + const auto kv_param_res_pairs = get_kv_param_res_pairs(model, kv_param_res_names); + manager.register_pass(kv_param_res_pairs); + } + + if (ggml_model_decoder->is_static()) { + manager.register_pass(); + } + manager.run_passes(model); + if (ggml_model_decoder->is_stateful()) { + auto output_names = ggml_model_decoder->get_model_output_names(); + std::map model_output_indexes; + for (size_t i=0; iget_output_size(); i++) { + auto output_friendly_name = model->output(i).get_node_shared_ptr()->get_friendly_name(); + auto output_id = model_output_indexes[output_friendly_name]; + auto model_output_shape = model->output(i).get_partial_shape(); + auto decoder_output_shape = ggml_model_decoder->get_output_shape(output_id); + if (model_output_shape.rank().is_static() && decoder_output_shape.rank().is_static() + && model_output_shape.rank().get_length() + 1 == decoder_output_shape.rank().get_length() + && decoder_output_shape[0].is_static() && decoder_output_shape[0].get_length() == 1) { + ppp.output(i).postprocess().custom([](const ov::Output& node) { + auto axes = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{1}, {0}); + return std::make_shared(node, axes); + }); + } + } + model = ppp.build(); + } + } + return model; +} + +} // namespace ggml +} // namespace frontend +} // namespace ov + + +/* Inlined openvino/utils.cpp */ + +#include "ggml-impl-defs.inc" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { + +std::string getCurrentTime() { + std::time_t now = std::time(nullptr); + char buf[100]; + std::strftime(buf, sizeof(buf), "%Y-%m-%d %H:%M:%S", std::localtime(&now)); + return buf; +} + +void num_inputs_check(const NodeContext & context, size_t min_inputs, size_t max_inputs) { + auto input_size = context.get_input_size(); + FRONT_END_OP_CONVERSION_CHECK(input_size >= min_inputs, "Got less inputs than expected"); + FRONT_END_OP_CONVERSION_CHECK(input_size <= max_inputs, "Got more inputs than expected"); +} + +int non_cont_dim(std::vector ne, std::vector nb) { + int dim = nb.size() - 1; + size_t bytes = nb[dim]; + for (int i = dim; i > 0; i--) { + bytes *= ne[i]; + if (bytes != nb[i - 1]) { + return i; + } + } + return 0; +} + +std::shared_ptr get_dimensions(const std::shared_ptr & shape, + const std::vector & dims) { + using namespace ov::op; + const auto zero = v0::Constant::create(ov::element::i32, ov::Shape{}, {0}); + const auto dims_const = v0::Constant::create(ov::element::i32, ov::Shape{dims.size()}, dims); + return std::make_shared(shape, dims_const, zero); +} + +std::shared_ptr get_dimensions(const std::shared_ptr & node, const std::vector & dims) { + return get_dimensions(std::make_shared(node), dims); +} + +OutputVector rename_outputs_with_suffix(const OutputVector & outputs, const std::string & suffix) { + for (const auto & output : outputs) { + auto node = output.get_node_shared_ptr(); + std::string name = node->get_friendly_name(); + name += "_"; + name += suffix; + node->set_friendly_name(name); + // std::cout << name << " " << output.get_partial_shape() << std::endl; + } + return outputs; +} + +namespace { +ov::Output rope_yarn_ramp_mix(int n_dims, const float corr_dims[2], float ext_factor) { + int half_n_dims = n_dims / 2; + std::vector dim_ids_vec(half_n_dims); + std::iota(dim_ids_vec.begin(), dim_ids_vec.end(), 0); + auto dim_ids = ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1, (size_t) half_n_dims}, dim_ids_vec); + auto corr_low = ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1, 1}, {corr_dims[0]}); + auto corr_high = ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1, 1}, {corr_dims[1]}); + auto denom = std::make_shared( + std::make_shared(corr_high, corr_low), + ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1, 1}, {0.001f})); + auto ramp_y = + std::make_shared(std::make_shared(dim_ids, corr_low), denom); + auto ramp_clamped = std::make_shared(ramp_y, 0.0f, 1.0f); + // rope_yarn_ramp returns (1 - clamp(y)), so invert before scaling + auto one = ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1, 1}, {1.0f}); + auto ramp_inverted = std::make_shared(one, ramp_clamped); + auto ext_factor_node = ov::op::v0::Constant::create(ov::element::f32, Shape{}, {ext_factor}); + auto ramp_mix = std::make_shared(ramp_inverted, ext_factor_node); + return ramp_mix; +} + +float ggml_rope_yarn_corr_dim(int n_dims, int n_ctx_orig, float n_rot, float base) { +#ifndef M_PI +# define M_PI 3.14159265358979323846 +#endif + return n_dims * logf(n_ctx_orig / (n_rot * 2 * (float) M_PI)) / (2 * logf(base)); +} + +void ggml_rope_yarn_corr_dims(int n_dims, + int n_ctx_orig, + float freq_base, + float beta_fast, + float beta_slow, + float dims[2]) { + float start = floorf(ggml_rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_fast, freq_base)); + float end = ceilf(ggml_rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_slow, freq_base)); + dims[0] = std::max(0.0f, start); + dims[1] = std::min(static_cast(n_dims - 1), end); +} +} // namespace + +std::pair, ov::Output> make_sin_cos(int32_t * rope_params, + std::shared_ptr inp_pos, + std::shared_ptr rope_freqs_weight, + bool imrope, + bool stateful) { + if (stateful) { + inp_pos = std::make_shared(inp_pos, ov::op::v0::Constant::create(ov::element::i64, {1}, {0})); + inp_pos = std::make_shared(inp_pos, ov::element::f32); + auto pos_perm = + std::make_shared(ov::element::i64, ov::Shape{3}, std::vector{2, 1, 0}); + inp_pos = std::make_shared(inp_pos, pos_perm); + } else if (imrope) { + inp_pos = std::make_shared(inp_pos, ov::element::f32); + auto pos_shape = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{5}, {0, 0, 0, 4, -1}); + inp_pos = std::make_shared(inp_pos, pos_shape, true); + auto pos_transpose_shape = + std::make_shared(ov::element::i64, ov::Shape{5}, std::vector{0, 1, 2, 4, 3}); + inp_pos = std::make_shared(inp_pos, pos_transpose_shape); + } else { + inp_pos = std::make_shared(inp_pos, ov::element::f32); + auto pos_perm = + std::make_shared(ov::element::i64, ov::Shape{4}, std::vector{0, 3, 1, 2}); + inp_pos = std::make_shared(inp_pos, pos_perm); + } + + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + const int n_dims = rope_params[1]; + const size_t n_dims_half = n_dims >> 1; + const int n_ctx_orig = rope_params[4]; + memcpy(&freq_base, rope_params + 5, sizeof(float)); + memcpy(&freq_scale, rope_params + 6, sizeof(float)); + memcpy(&ext_factor, rope_params + 7, sizeof(float)); + memcpy(&attn_factor, rope_params + 8, sizeof(float)); + memcpy(&beta_fast, rope_params + 9, sizeof(float)); + memcpy(&beta_slow, rope_params + 10, sizeof(float)); + + const float theta_scale = powf(freq_base, -2.0f / n_dims); + + std::vector factor(n_dims_half); + + Output freq_factors; + + Output theta; + float mscale = attn_factor; + if (imrope) { + std::vector gather_indices(n_dims_half); + for (size_t j = 0; j < n_dims_half; j++) { + gather_indices[j] = j % 3; + factor[j] = std::pow(theta_scale, j); + } + auto gather_indices_const = + std::make_shared(ov::element::i64, ov::Shape{n_dims_half}, gather_indices); + auto gather_axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {4}); + inp_pos = std::make_shared(inp_pos, gather_indices_const, gather_axis); + auto factor_const = std::make_shared(ov::element::f32, ov::Shape{n_dims_half}, factor); + theta = std::make_shared(inp_pos, factor_const); + } else { + float corr_dims[2]; + ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); + factor[0] = 1.0f; + for (size_t i = 1; i < factor.size(); i++) { + factor[i] = theta_scale * factor[i - 1]; + } + if (stateful) { + freq_factors = + std::make_shared(ov::element::f32, ov::Shape{1, 1, factor.size()}, factor); + } else { + freq_factors = + std::make_shared(ov::element::f32, ov::Shape{1, 1, 1, factor.size()}, factor); + } + if (rope_freqs_weight) { + freq_factors = std::make_shared(freq_factors, rope_freqs_weight); + } + + auto theta_extrap = std::make_shared(freq_factors, inp_pos); + auto theta_interp = std::make_shared( + theta_extrap, ov::op::v0::Constant::create(ov::element::f32, {1}, {freq_scale})); + + if (ext_factor == 0.0f) { + theta = theta_interp; + } else { + auto ramp_mix = rope_yarn_ramp_mix(n_dims, corr_dims, ext_factor); + Output one; + if (stateful) { + one = ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1}, {1.0f}); + } else { + one = ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1, 1}, {1.0f}); + } + auto one_minus_ramp = std::make_shared(one, ramp_mix); + + theta = std::make_shared(std::make_shared(theta_interp, one_minus_ramp), + std::make_shared(theta_extrap, ramp_mix)); + mscale *= (1.0f + 0.1f * std::log(1.0f / freq_scale)); + } + } + + Output cos_theta = std::make_shared(theta); + Output sin_theta = std::make_shared(theta); + + if (!imrope) { + auto mscale_node = ov::op::v0::Constant::create(ov::element::f32, Shape{}, {mscale}); + + cos_theta = std::make_shared(cos_theta, mscale_node); + sin_theta = std::make_shared(sin_theta, mscale_node); + } + + return std::make_pair(sin_theta, cos_theta); +} + +ov::Output process_view_input(const NodeContext & context, int input_index, int slice_len) { + // Only works for VIEW operations that slice at the lowest dimension + // If the VIEW also reshape the result, `slice_len` should be provided + auto input = context.get_input(input_index); + auto * op_params = (size_t *) context.get_input_op_params(input_index); + auto src1_stride = context.get_input_stride(input_index); + + int64_t split_addr = op_params[0] / src1_stride[3]; + if (slice_len == 0) { + slice_len = context.get_input_shape(input_index)[3].get_length(); + } + int64_t slice_end = split_addr + slice_len; + + auto begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {split_addr}); + auto end = ov::op::v0::Constant::create(ov::element::i64, {1}, {slice_end}); + auto stride = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto axes = ov::op::v0::Constant::create(ov::element::i64, {1}, {context.is_stateful() ? 2 : 3}); + auto sliced = std::make_shared(input, begin, end, stride, axes); + return sliced; +} + +} // namespace ggml +} // namespace frontend +} // namespace ov diff --git a/ggml/src/ggml-openvino/ggml-quants.cpp b/ggml/src/ggml-openvino/ggml-quants.cpp deleted file mode 100644 index 57d66df4f017..000000000000 --- a/ggml/src/ggml-openvino/ggml-quants.cpp +++ /dev/null @@ -1,956 +0,0 @@ -#include "ggml-quants.h" - -#include "ggml-common.h" -#include "ggml-impl.h" -#include "ggml.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -void unpack_32_4(const uint8_t * data, uint8_t * dst) { - std::fill_n(dst, 16, 0); - for (int j = 0; j < 16; ++j) { - uint8_t x = (data[j] & 0x0F); - uint8_t y = (data[j] >> 4); - if (j % 2 != 0) { - x <<= 4; - y <<= 4; - } - dst[j / 2] |= x; - dst[8 + j / 2] |= y; // Last 16 weights are in the higher bits - } -} - -// Extracts (weight, scales, zp) from Q4_0 tensors. -// Data layout is: |16 bit scale|32 x 4bit weights|. -// When zp_arr is empty (symmetric), weights are stored as signed i4 (value - 8). -void extract_q4_0_data(const ggml_tensor * tensor, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr) { - const uint64_t bytes_per_block = 18; // 2 bytes scale, 32x0.5 byte weights - - auto * data = static_cast(tensor->data); - auto * weights = static_cast(weights_arr.data()); - auto * scales = scales_arr.data::value_type>(); - - bool is_symmetric = (weights_arr.get_element_type() == ov::element::i4); // Signed i4 path - - if (!is_symmetric) { - auto * zp = static_cast(zp_arr.data()); - ov::parallel_for(scales_arr.get_size(), [&](size_t i) { - scales[i] = ov::float16::from_bits(*((uint16_t *) (data + i * bytes_per_block))); - // Pack two 4-bit zero points per byte - if (i % 2 == 0) { - zp[i / 2] = 8; // Lower nibble - } else { - zp[i / 2] |= (8 << 4); // Upper nibble - } - unpack_32_4(data + i * bytes_per_block + 2, weights + i * 16); - }); - } else { - // Symmetric: unpack as u4 then convert to i4 by subtracting 8 (XOR each nibble) - ov::parallel_for(scales_arr.get_size(), [&](size_t i) { - scales[i] = ov::float16::from_bits(*((uint16_t *) (data + i * bytes_per_block))); - unpack_32_4(data + i * bytes_per_block + 2, weights + i * 16); - // Convert u4 to i4: subtract 8 from each nibble. XOR 0x88 flips each nibble by 8. - for (int j = 0; j < 16; ++j) { - weights[i * 16 + j] ^= 0x88; - } - }); - } -} - -// Extracts (weight, scales, zp) from Q4_1 tensors. -// Data layout is: |16 bit scale|16 bit min|32 x 4bit weights|. -void extract_q4_1_data(const ggml_tensor * tensor, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr, - bool use_bias) { - const uint64_t bytes_per_block = 20; // 2 bytes scale, 2 bytes min, 32x0.5 byte weights - - auto * data = static_cast(tensor->data); - auto * weights = static_cast(weights_arr.data()); - auto * scales = scales_arr.data::value_type>(); - - if (use_bias) { - // Store bias (min) directly as f16 instead of computing u4 zero points - auto * bias = zp_arr.data::value_type>(); - ov::parallel_for(scales_arr.get_size(), [&](size_t i) { - float scale = static_cast(ov::float16::from_bits(*((uint16_t *) (data + i * bytes_per_block)))); - float min = static_cast(ov::float16::from_bits(*((uint16_t *) (data + i * bytes_per_block + 2)))); - scales[i] = ov::float16(scale); - bias[i] = ov::float16(min); // bias = min, dequant: w*s + bias - unpack_32_4(data + i * bytes_per_block + 4, weights + i * 16); - }); - } else { - auto * zp = static_cast(zp_arr.data()); - ov::parallel_for(scales_arr.get_size(), [&](size_t i) { - float scale = static_cast(ov::float16::from_bits(*((uint16_t *) (data + i * bytes_per_block)))); - float min = static_cast(ov::float16::from_bits(*((uint16_t *) (data + i * bytes_per_block + 2)))); - scales[i] = ov::float16(scale); - // zp = -min / scale (bias = min, so zp = -bias/scale) - uint8_t zp_val = (scale != 0.0f) ? (uint8_t) std::round(-min / scale) : 0; - // Pack two 4-bit zero points per byte - if (i % 2 == 0) { - zp[i / 2] = zp_val & 0x0F; // Lower nibble - } else { - zp[i / 2] |= (zp_val << 4); // Upper nibble - } - unpack_32_4(data + i * bytes_per_block + 4, weights + i * 16); - }); - } -} - -// Extracts (weight, scales, zp) from Q8_0 tensors. -// Data layout is: |16 bit scale|32 x 8bit weights|. -// When zp_arr is empty (symmetric), weights are stored as signed i8 directly. -void extract_q8_0_data(const ggml_tensor * tensor, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr) { - const uint64_t weights_per_block = 32; - const uint64_t bytes_per_block = 34; // 2 bytes scale, 32x1 byte weights - - auto * data = static_cast(tensor->data); - auto * weights = static_cast(weights_arr.data()); - auto * scales = scales_arr.data::value_type>(); - - bool is_symmetric = (weights_arr.get_element_type() == ov::element::i8); // Signed i8 path - - if (!is_symmetric) { - auto * zp = static_cast(zp_arr.data()); - ov::parallel_for(scales_arr.get_size(), [&](size_t i) { - uint8_t * block_data = data + i * bytes_per_block; - scales[i] = ov::float16::from_bits(*(uint16_t *) block_data); - zp[i] = 128; - for (size_t j = 0; j < weights_per_block; ++j) { - uint8_t x = block_data[j + 2]; - x ^= 1 << 7; // Convert int8 to uint8 by flipping sign bit - weights[i * weights_per_block + j] = x; - } - }); - } else { - // Symmetric: store original int8 values directly (no unsigned bias) - ov::parallel_for(scales_arr.get_size(), [&](size_t i) { - uint8_t * block_data = data + i * bytes_per_block; - scales[i] = ov::float16::from_bits(*(uint16_t *) block_data); - // Copy int8 weights as-is (the tensor element type is i8) - memcpy(weights + i * weights_per_block, block_data + 2, weights_per_block); - }); - } -} - -void unpack_256_4(const uint8_t * data, uint8_t * dst) { - // Initialize the output array with zeros - std::fill_n(dst, 128, 0); - - for (size_t i = 0; i < 4; ++i) { - for (int j = 0; j < 32; ++j) { - uint8_t x = (data[i * 32 + j] & 0x0F); - uint8_t y = (data[i * 32 + j] >> 4); - if (j % 2 != 0) { - x <<= 4; - y <<= 4; - } - dst[i * 32 + j / 2] |= x; - dst[i * 32 + 16 + j / 2] |= y; // Last 16 weights are in the higher bits - } - } -} - -void extract_q4_k_data(const ggml_tensor * tensor, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr, - bool use_bias) { - const uint64_t bytes_per_block = 2 + 2 + 12 + 128; - const uint64_t n_super_block = tensor->nb[3] / bytes_per_block; - - auto * data = static_cast(tensor->data); - auto * weights = static_cast(weights_arr.data()); - auto * scales = scales_arr.data::value_type>(); - - // For bias path, zp_arr holds f16 bias values; for zp path, it holds packed u4 zero points - auto * zp_u4 = use_bias ? nullptr : static_cast(zp_arr.data()); - auto * bias_f16 = use_bias ? zp_arr.data::value_type>() : nullptr; - - ov::parallel_for(n_super_block, [&](size_t i) { - uint8_t * block_data = data + i * bytes_per_block; - - // Extract scale factors and offsets - float scale_scales = static_cast(ov::float16::from_bits(*((uint16_t *) block_data))); - float scale_mins = static_cast(ov::float16::from_bits(*((uint16_t *) block_data + 1))); - - // Extract qs1 and qs2 - uint8_t * qs1 = block_data + 4; - - // Calculate scales - float scale_vals[8]; - scale_vals[0] = scale_scales * static_cast((*(qs1) & 0b111111)); - scale_vals[1] = scale_scales * static_cast((*(qs1 + 1) & 0b111111)); - scale_vals[2] = scale_scales * static_cast((*(qs1 + 2) & 0b111111)); - scale_vals[3] = scale_scales * static_cast((*(qs1 + 3) & 0b111111)); - scale_vals[4] = scale_scales * static_cast((*(qs1 + 8) & 0b00001111) | ((*(qs1) >> 6) << 4)); - scale_vals[5] = scale_scales * static_cast((*(qs1 + 9) & 0b00001111) | ((*(qs1 + 1) >> 6) << 4)); - scale_vals[6] = scale_scales * static_cast((*(qs1 + 10) & 0b00001111) | ((*(qs1 + 2) >> 6) << 4)); - scale_vals[7] = scale_scales * static_cast((*(qs1 + 11) & 0b00001111) | ((*(qs1 + 3) >> 6) << 4)); - - // Calculate min values (bias = -min) - float min_vals[8]; - min_vals[0] = scale_mins * static_cast((*(qs1 + 4) & 0b111111)); - min_vals[1] = scale_mins * static_cast((*(qs1 + 5) & 0b111111)); - min_vals[2] = scale_mins * static_cast((*(qs1 + 6) & 0b111111)); - min_vals[3] = scale_mins * static_cast((*(qs1 + 7) & 0b111111)); - min_vals[4] = scale_mins * static_cast((*(qs1 + 8) >> 4) | ((*(qs1 + 4) >> 6) << 4)); - min_vals[5] = scale_mins * static_cast((*(qs1 + 9) >> 4) | ((*(qs1 + 5) >> 6) << 4)); - min_vals[6] = scale_mins * static_cast((*(qs1 + 10) >> 4) | ((*(qs1 + 6) >> 6) << 4)); - min_vals[7] = scale_mins * static_cast((*(qs1 + 11) >> 4) | ((*(qs1 + 7) >> 6) << 4)); - - // Store scales and compute zero points or bias - for (int j = 0; j < 8; j++) { - scales[i * 8 + j] = ov::float16(scale_vals[j]); - if (use_bias) { - // Store bias = -min directly as f16, dequant: w*s + bias - bias_f16[i * 8 + j] = ov::float16(-min_vals[j]); - } else { - // zp = min / scale (since bias = -min and zp = -bias/scale) - uint8_t zp_val = (scale_vals[j] != 0.0f) ? (uint8_t) std::round(min_vals[j] / scale_vals[j]) : 0; - // Pack two 4-bit zero points per byte - size_t idx = i * 8 + j; - if (idx % 2 == 0) { - zp_u4[idx / 2] = zp_val & 0x0F; - } else { - zp_u4[idx / 2] |= (zp_val << 4); - } - } - } - unpack_256_4(block_data + 16, weights + i * 128); - }); -} - -void extract_q6_k_data(const ggml_tensor * tensor, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr) { - const uint64_t bytes_per_block = 128 + 64 + 16 + 2; - const uint64_t n_super_block = tensor->nb[3] / bytes_per_block; - - auto * data = static_cast(tensor->data); - auto * weights = static_cast(weights_arr.data()); - auto * scales = scales_arr.data::value_type>(); - - bool is_symmetric = (weights_arr.get_element_type() == ov::element::i8); // Signed i8 path - - if (!is_symmetric) { - auto * zp = static_cast(zp_arr.data()); - ov::parallel_for(n_super_block, [&](size_t i) { - uint8_t * block_data = data + i * bytes_per_block; - float scale_factor = static_cast(ov::float16::from_bits(*((uint16_t *) block_data + 104))); - for (size_t j = 0; j < 16; j++) { - scales[j + i * 16] = - ov::float16(scale_factor * static_cast(*((int8_t *) (block_data + 128 + 64 + j)))); - zp[j + i * 16] = 32; - } - uint8_t * ql = block_data; - uint8_t * qh = block_data + 128; - for (int64_t j = 0; j < 32; ++j) { - weights[i * 256 + j] = (ql[j] & 0xF) | (((qh[j] >> 0) & 3) << 4); - weights[i * 256 + j + 32] = (ql[32 + j] & 0xF) | (((qh[j] >> 2) & 3) << 4); - weights[i * 256 + j + 64] = (ql[j] >> 4) | (((qh[j] >> 4) & 3) << 4); - weights[i * 256 + j + 96] = (ql[32 + j] >> 4) | (((qh[j] >> 6) & 3) << 4); - weights[i * 256 + j + 128] = (ql[64 + j] & 0xF) | (((qh[32 + j] >> 0) & 3) << 4); - weights[i * 256 + j + 160] = (ql[96 + j] & 0xF) | (((qh[32 + j] >> 2) & 3) << 4); - weights[i * 256 + j + 192] = (ql[64 + j] >> 4) | (((qh[32 + j] >> 4) & 3) << 4); - weights[i * 256 + j + 224] = (ql[96 + j] >> 4) | (((qh[32 + j] >> 6) & 3) << 4); - } - }); - } else { - // Symmetric: subtract 32 from each weight to store as signed i8 - ov::parallel_for(n_super_block, [&](size_t i) { - uint8_t * block_data = data + i * bytes_per_block; - float scale_factor = static_cast(ov::float16::from_bits(*((uint16_t *) block_data + 104))); - for (size_t j = 0; j < 16; j++) { - scales[j + i * 16] = - ov::float16(scale_factor * static_cast(*((int8_t *) (block_data + 128 + 64 + j)))); - } - uint8_t * ql = block_data; - uint8_t * qh = block_data + 128; - auto * signed_weights = reinterpret_cast(weights); - for (int64_t j = 0; j < 32; ++j) { - signed_weights[i * 256 + j] = static_cast((ql[j] & 0xF) | (((qh[j] >> 0) & 3) << 4)) - 32; - signed_weights[i * 256 + j + 32] = - static_cast((ql[32 + j] & 0xF) | (((qh[j] >> 2) & 3) << 4)) - 32; - signed_weights[i * 256 + j + 64] = static_cast((ql[j] >> 4) | (((qh[j] >> 4) & 3) << 4)) - 32; - signed_weights[i * 256 + j + 96] = - static_cast((ql[32 + j] >> 4) | (((qh[j] >> 6) & 3) << 4)) - 32; - signed_weights[i * 256 + j + 128] = - static_cast((ql[64 + j] & 0xF) | (((qh[32 + j] >> 0) & 3) << 4)) - 32; - signed_weights[i * 256 + j + 160] = - static_cast((ql[96 + j] & 0xF) | (((qh[32 + j] >> 2) & 3) << 4)) - 32; - signed_weights[i * 256 + j + 192] = - static_cast((ql[64 + j] >> 4) | (((qh[32 + j] >> 4) & 3) << 4)) - 32; - signed_weights[i * 256 + j + 224] = - static_cast((ql[96 + j] >> 4) | (((qh[32 + j] >> 6) & 3) << 4)) - 32; - } - }); - } -} - -static 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); - } -} - -void extract_q5_k_data(const ggml_tensor * tensor, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr, - bool use_bias) { - const uint64_t bytes_per_block = 4 + 12 + 32 + 128; - const uint64_t n_super_block = tensor->nb[3] / bytes_per_block; - - auto * data = static_cast(tensor->data); - auto * weights = static_cast(weights_arr.data()); - auto * scales = scales_arr.data::value_type>(); - - // For bias path, zp_arr holds f16 bias values; for zp path, it holds u8 zero points - auto * zp_u8 = use_bias ? nullptr : static_cast(zp_arr.data()); - auto * bias_f16 = use_bias ? zp_arr.data::value_type>() : nullptr; - - ov::parallel_for(n_super_block, [&](size_t i) { - uint8_t * block_data = data + i * bytes_per_block; - - const float d = static_cast(ov::float16::from_bits(*((uint16_t *) block_data))); - const float min_factor = static_cast(ov::float16::from_bits(*((uint16_t *) block_data + 1))); - - const uint8_t * scales_data = block_data + 4; // 12 bytes of scales - const uint8_t * qh = block_data + 4 + 12; // 32 bytes of high bits - const uint8_t * ql = block_data + 4 + 12 + 32; // 128 bytes of low bits - - int is = 0; - uint8_t u1 = 1; - uint8_t u2 = 2; - - // Process 2 blocks in one iteration - for (int j = 0; j < 256; j += 64) { // 256 = QK_K, so 4 iterations of 64 - uint8_t sc; - uint8_t m; - - // Get scale and min for first 32 elements - get_scale_min_k4(is + 0, scales_data, &sc, &m); - const float d1 = d * sc; - const float m1 = min_factor * m; - - // Get scale and min for second 32 elements - get_scale_min_k4(is + 1, scales_data, &sc, &m); - const float d2 = d * sc; - const float m2 = min_factor * m; - - scales[i * 8 + is] = ov::float16(d1); - scales[i * 8 + is + 1] = ov::float16(d2); - if (use_bias) { - // Store bias = -min directly as f16, dequant: w*s + bias - bias_f16[i * 8 + is] = ov::float16(-m1); - bias_f16[i * 8 + is + 1] = ov::float16(-m2); - } else { - // zp = min / scale (since bias = -min and zp = -bias/scale) - zp_u8[i * 8 + is] = (d1 != 0.0f) ? (uint8_t) std::round(m1 / d1) : 0; - zp_u8[i * 8 + is + 1] = (d2 != 0.0f) ? (uint8_t) std::round(m2 / d2) : 0; - } - - // Extract weights for first 32 elements (matching deq formula exactly) - for (int l = 0; l < 32; ++l) { - weights[i * 256 + j + l] = (ql[l] & 0xF) + ((qh[l] & u1) ? 16 : 0); - } - - // Extract weights for second 32 elements - for (int l = 0; l < 32; ++l) { - weights[i * 256 + j + l + 32] = (ql[l] >> 4) + ((qh[l] & u2) ? 16 : 0); - } - - ql += 32; - is += 2; - u1 <<= 2; - u2 <<= 2; - } - }); -} - -// TODO Reorder for make_intX_weights - -ov::Output make_int8_weights(ov::Tensor & weight, - ov::Tensor & scales, - ov::Tensor & zp, - size_t group_size, - bool use_bias) { - ov::Shape orig_shape = weight.get_shape(); - bool is_signed = (weight.get_element_type() == ov::element::i8); // Symmetric: signed weights, no ZP - - // Expand dimensions for scales and zp/bias - auto scale_shape = scales.get_shape(); - - ov::Shape packed_shape = {orig_shape[0], orig_shape[1] / group_size, group_size}; - - if (packed_shape[1] == 1) { - // Requantized channel-wise case - packed_shape.erase(packed_shape.begin() + 1); - } else { - scale_shape.push_back(1); - scales.set_shape(scale_shape); - if (!is_signed && zp.get_size() > 0) { - auto zp_shape = zp.get_shape(); - zp_shape.push_back(1); - zp.set_shape(zp_shape); - } - } - - auto scales_f16 = std::make_shared(scales); - - ov::Output result; - if (is_signed) { - // Signed path: q * s (no zero point subtraction needed) - auto weights_node = std::make_shared(ov::element::i8, packed_shape, - static_cast(weight.data()), nullptr); - weights_node->get_rt_info()["__gguf_tensor_holder"] = weight; - auto weights_f16 = std::make_shared(weights_node, ov::element::f16); - result = std::make_shared(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY); - } else { - // Unsigned path - auto weights_node = std::make_shared(ov::element::u8, packed_shape, - static_cast(weight.data()), nullptr); - weights_node->get_rt_info()["__gguf_tensor_holder"] = weight; - auto weights_f16 = std::make_shared(weights_node, ov::element::f16); - - if (use_bias && zp.get_size() > 0) { - // Bias path: w * s + b (zp tensor holds f16 bias values) - auto bias_f16 = std::make_shared(zp); - auto w_s = - std::make_shared(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY); - result = std::make_shared(w_s, bias_f16, ov::op::AutoBroadcastType::NUMPY); - } else { - // Zero point path: (w - zp) * s - auto zero_point = std::make_shared(zp); - float zp_value; - if (ov::op::util::get_single_value(zero_point, zp_value)) { - zero_point = ov::op::v0::Constant::create(zero_point->get_element_type(), {}, {zp_value}); - } - auto zero_point_f16 = std::make_shared(zero_point, ov::element::f16); - auto w_zp = - std::make_shared(weights_f16, zero_point_f16, ov::op::AutoBroadcastType::NUMPY); - result = std::make_shared(w_zp, scales_f16, ov::op::AutoBroadcastType::NUMPY); - } - } - - if (packed_shape.size() != 2) { - // If not requantized channel-wise case, reshape back to original shape - auto final_shape = - std::make_shared(ov::element::i64, ov::Shape{orig_shape.size()}, orig_shape); - result = std::make_shared(result, final_shape, false); - } - - return std::make_shared(result, ov::element::f32); -} - -ov::Output make_int4_weights(ov::Tensor & weight, - ov::Tensor & scales, - ov::Tensor & zp, - size_t group_size, - bool use_bias) { - ov::Shape orig_weight_shape = weight.get_shape(); - bool is_signed = (weight.get_element_type() == ov::element::i4); // Symmetric: signed weights, no ZP - - // Expand dimensions for scales and zp/bias - ov::Shape scale_shape = scales.get_shape(); - - // Create INT4 weight tensor - ov::Shape packed_shape = {orig_weight_shape[0], orig_weight_shape[1] / group_size, group_size}; - - if (packed_shape[1] == 1) { - // Requantized channel-wise case - packed_shape.erase(packed_shape.begin() + 1); - } else { - scale_shape.push_back(1); - scales.set_shape(scale_shape); - if (!is_signed && zp.get_size() > 0) { - auto zp_shape = zp.get_shape(); - zp_shape.push_back(1); - zp.set_shape(zp_shape); - } - } - - auto scales_f16 = std::make_shared(scales); - - ov::Output result; - if (is_signed) { - // Signed path: q * s (no zero point subtraction needed) - auto weights_node = std::make_shared(ov::element::i4, packed_shape, - static_cast(weight.data()), nullptr); - weights_node->get_rt_info()["__gguf_tensor_holder"] = weight; - auto weights_f16 = std::make_shared(weights_node, ov::element::f16); - result = std::make_shared(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY); - } else { - // Unsigned path - auto weights_node = std::make_shared(ov::element::u4, packed_shape, - static_cast(weight.data()), nullptr); - weights_node->get_rt_info()["__gguf_tensor_holder"] = weight; - auto weights_f16 = std::make_shared(weights_node, ov::element::f16); - - if (use_bias && zp.get_size() > 0) { - // Bias path: w * s + b (zp tensor holds f16 bias values) - auto bias_f16 = std::make_shared(zp); - auto w_s = - std::make_shared(weights_f16, scales_f16, ov::op::AutoBroadcastType::NUMPY); - result = std::make_shared(w_s, bias_f16, ov::op::AutoBroadcastType::NUMPY); - } else { - // Zero point path: (w - zp) * s - auto zero_points_node = std::make_shared(zp); - float zp_value; - if (ov::op::util::get_single_value(zero_points_node, zp_value)) { - zero_points_node = ov::op::v0::Constant::create(zero_points_node->get_element_type(), {}, {zp_value}); - } - auto zero_points_f16 = std::make_shared(zero_points_node, ov::element::f16); - auto w_zp = - std::make_shared(weights_f16, zero_points_f16, ov::op::AutoBroadcastType::NUMPY); - result = std::make_shared(w_zp, scales_f16, ov::op::AutoBroadcastType::NUMPY); - } - } - - if (packed_shape.size() != 2) { - // If not requantized channel-wise case, reshape back to original shape - auto final_shape = std::make_shared(ov::element::i64, ov::Shape{orig_weight_shape.size()}, - orig_weight_shape); - result = std::make_shared(result, final_shape, false); - } - - return std::make_shared(result, ov::element::f32); -} - -// Extract quantized weights from tensor and create weight subgraph -std::shared_ptr extract_quantized_weights(const ggml_tensor * tensor, - const void * data, - ov::Tensor & weights, - ov::Tensor & scales, - ov::Tensor & zp, - bool use_bias) { - // Create a temporary tensor for extraction functions that read from tensor->data - ggml_tensor temp_tensor = *tensor; - temp_tensor.data = const_cast(data); - - // Determine block size based on tensor type - int64_t weights_per_block; - bool is_u4; - switch (tensor->type) { - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q4_K: - is_u4 = true; - weights_per_block = 32; - break; - case GGML_TYPE_Q8_0: - case GGML_TYPE_Q5_K: - is_u4 = false; - weights_per_block = 32; - break; - case GGML_TYPE_Q6_K: - is_u4 = false; - weights_per_block = 16; - break; - default: - throw std::runtime_error("Unsupported quantized type for extraction: " + - std::string(ggml_type_name(tensor->type))); - } - - // Extract quantized data - switch (tensor->type) { - case GGML_TYPE_Q4_0: - extract_q4_0_data(&temp_tensor, weights, scales, zp); - break; - case GGML_TYPE_Q4_1: - extract_q4_1_data(&temp_tensor, weights, scales, zp, use_bias); - break; - case GGML_TYPE_Q4_K: - extract_q4_k_data(&temp_tensor, weights, scales, zp, use_bias); - break; - case GGML_TYPE_Q8_0: - extract_q8_0_data(&temp_tensor, weights, scales, zp); - break; - case GGML_TYPE_Q6_K: - extract_q6_k_data(&temp_tensor, weights, scales, zp); - break; - case GGML_TYPE_Q5_K: - extract_q5_k_data(&temp_tensor, weights, scales, zp, use_bias); - break; - default: - throw std::runtime_error("Unsupported quantized type: " + std::string(ggml_type_name(tensor->type))); - } - - // Create the OpenVINO weight subgraph - ov::Output weight_node; - if (is_u4) { - weight_node = make_int4_weights(weights, scales, zp, weights_per_block, use_bias); - } else { - weight_node = make_int8_weights(weights, scales, zp, weights_per_block, use_bias); - } - - auto result = weight_node.get_node_shared_ptr(); - result->set_friendly_name(tensor->name); - return result; -} - -// Requantize weights to target format, writing to provided buffers -std::shared_ptr requantize_to_buffers(const ggml_tensor * tensor, - const void * data, - ExtraQuantType requant_type, - int64_t block_size, - ov::Tensor & weights, - ov::Tensor & scales, - ov::Tensor & zp) { - int64_t n_elements = ggml_nelements(tensor); - - // First dequantize to F32 - std::vector weights_f32(n_elements); - ggml_get_type_traits(tensor->type)->to_float(data, weights_f32.data(), n_elements); - - // Handle F16 case - just convert and create constant - if (requant_type == ExtraQuantType::F16) { - ggml_get_type_traits(GGML_TYPE_F16)->from_float_ref(weights_f32.data(), weights.data(), n_elements); - auto result = std::make_shared(weights); - result->set_friendly_name(tensor->name); - return result; - } - - // Requantize to target quantized format - bool is_u4 = (requant_type == ExtraQuantType::Q4_0_C || requant_type == ExtraQuantType::Q4_0_128); - - if (is_u4) { - quantize_q4_0(weights_f32.data(), weights, scales, zp, n_elements, block_size); - } else if (requant_type == ExtraQuantType::Q8_1_C) { - quantize_q8_1(weights_f32.data(), weights, scales, zp, n_elements, block_size); - } else { - quantize_q8_0(weights_f32.data(), weights, scales, zp, n_elements, block_size); - } - - // Create the OpenVINO weight subgraph - ov::Output weight_node; - if (is_u4) { - weight_node = make_int4_weights(weights, scales, zp, block_size); - } else { - weight_node = make_int8_weights(weights, scales, zp, block_size); - } - - auto result = weight_node.get_node_shared_ptr(); - result->set_friendly_name(tensor->name); - return result; -} - -OvWeight process_weight_tensor(const ggml_tensor * tensor, const void * data, void * output_base_ptr, bool use_bias) { - GGML_ASSERT(tensor != nullptr); - GGML_ASSERT(data != nullptr); - - OvWeight result; - - // Get 2D shape for weights [rows, cols] - ov::Shape node_shape = {static_cast(tensor->ne[1]), static_cast(tensor->ne[0])}; - - // Handle F16/F32/BF16 weights - if (tensor->type == GGML_TYPE_F32 || tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_BF16) { - ov::element::Type element_type; - switch (tensor->type) { - case GGML_TYPE_F32: - element_type = ov::element::f32; - break; - case GGML_TYPE_F16: - element_type = ov::element::f16; - break; - case GGML_TYPE_BF16: - element_type = ov::element::bf16; - break; - default: - OPENVINO_THROW("Unexpected tensor type in F16/F32/BF16 path"); - } - - if (output_base_ptr && output_base_ptr != data) { - // Using external buffer - copy data and create shared-memory constant - size_t tensor_bytes = ggml_nbytes(tensor); - memcpy(output_base_ptr, data, tensor_bytes); - result.weights = ov::Tensor(element_type, node_shape, output_base_ptr); - } else { - result.weights = ov::Tensor(element_type, node_shape, data); - } - result.weight_node = std::make_shared(result.weights); - return result; - } - - // Handle quantized weights - if (!ggml_is_quantized(tensor->type)) { - OPENVINO_THROW("Unsupported weight tensor type: ", ggml_type_name(tensor->type)); - } - - result.layout = ggml_openvino_get_extracted_layout(tensor, use_bias); - const auto & layout = result.layout; - if (layout.total_size == 0) { - OPENVINO_THROW("Unsupported quantized type: ", ggml_type_name(tensor->type)); - } - - if (use_bias) { - OPENVINO_ASSERT(!layout.is_requant, - "use_bias is only used for test-backend-ops, which should not have requantization"); - // bias node will be created on the fly and not use backend buffer - output_base_ptr = nullptr; - } - - // F16 requant path - no separate scales/zp needed in result - if (layout.is_requant && layout.requant_type.has_value() && layout.requant_type.value() == ExtraQuantType::F16) { - if (output_base_ptr) { - result.weights = ov::Tensor(ov::element::f16, node_shape, - static_cast(output_base_ptr) + layout.weights_offset); - } else { - result.weights = ov::Tensor(ov::element::f16, node_shape); - } - ov::Tensor dummy_scales, dummy_zp; // Not used for F16 - result.weight_node = - requantize_to_buffers(tensor, data, ExtraQuantType::F16, 0, result.weights, dummy_scales, dummy_zp); - return result; - } - - // Quantized path (normal extraction or quantized requant) - // Create weight/scale/zp tensors - shared between both paths - // For symmetric quantization, use signed types (i4/i8) and no ZP tensor - ov::element::Type weight_type = layout.is_symmetric ? (layout.is_u4 ? ov::element::i4 : ov::element::i8) : - (layout.is_u4 ? ov::element::u4 : ov::element::u8); - ov::Shape scale_shape = {node_shape[0], node_shape[1] / layout.weights_per_block}; - - if (output_base_ptr) { - uint8_t * buf_base = static_cast(output_base_ptr); - result.weights = ov::Tensor(weight_type, node_shape, buf_base + layout.weights_offset); - result.scales = ov::Tensor(ov::element::f16, scale_shape, buf_base + layout.scales_offset); - if (!layout.is_symmetric) { - ov::element::Type zp_type = layout.is_u4 ? ov::element::u4 : ov::element::u8; - result.zp = ov::Tensor(zp_type, scale_shape, buf_base + layout.zp_offset); - } - // else: result.zp remains default-constructed (empty) for symmetric - } else { - result.weights = ov::Tensor(weight_type, node_shape); - result.scales = ov::Tensor(ov::element::f16, scale_shape); - if (!layout.is_symmetric) { - if (use_bias) { - result.zp = ov::Tensor(ov::element::f16, scale_shape); - } else { - ov::element::Type zp_type = layout.is_u4 ? ov::element::u4 : ov::element::u8; - result.zp = ov::Tensor(zp_type, scale_shape); - } - } - // else: result.zp remains default-constructed (empty) for symmetric - } - - if (layout.is_requant && layout.requant_type.has_value()) { - result.weight_node = requantize_to_buffers(tensor, data, layout.requant_type.value(), layout.weights_per_block, - result.weights, result.scales, result.zp); - } else { - result.weight_node = - extract_quantized_weights(tensor, data, result.weights, result.scales, result.zp, use_bias); - } - - return result; -} - -void quantize_q4_0(const float * x, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr, - int64_t k, - int64_t qk) { - assert(k % qk == 0); - const int nb = k / qk; - - auto * weights = static_cast(weights_arr.data()); - auto * scales = scales_arr.data::value_type>(); - bool is_symmetric = (weights_arr.get_element_type() == ov::element::i4); // Signed i4 path - - if (!is_symmetric) { - auto * zp = static_cast(zp_arr.data()); - for (int i = 0; i < nb; i++) { - float amax = 0.0f; - float max = 0.0f; - for (int j = 0; j < qk; j++) { - const float v = x[i * qk + j]; - if (amax < fabsf(v)) { - amax = fabsf(v); - max = v; - } - } - const float d = max / -8; - if (d == 0) { - scales[i] = ov::float16(1.0f); - if (i % 2 == 0) { - zp[i / 2] = 8; - } else { - zp[i / 2] |= (8 << 4); - } - memset(weights + i * qk / 2, 8 | (8 << 4), qk / 2); - continue; - } - const float id = 1.0f / d; - scales[i] = ov::float16(d); - if (i % 2 == 0) { - zp[i / 2] = 8; - } else { - zp[i / 2] |= (8 << 4); - } - for (int j = 0; j < qk / 2; ++j) { - const float x0 = x[i * qk + 2 * j] * id; - const float x1 = x[i * qk + 2 * j + 1] * id; - const uint8_t xi0 = MIN(15, (int8_t) (x0 + 8.5f)); - const uint8_t xi1 = MIN(15, (int8_t) (x1 + 8.5f)); - weights[i * qk / 2 + j] = xi0 | (xi1 << 4); - } - } - } else { - // Symmetric: produce signed i4 values in [-8, 7] - for (int i = 0; i < nb; i++) { - float amax = 0.0f; - float max = 0.0f; - for (int j = 0; j < qk; j++) { - const float v = x[i * qk + j]; - if (amax < fabsf(v)) { - amax = fabsf(v); - max = v; - } - } - const float d = max / -8; - if (d == 0) { - scales[i] = ov::float16(1.0f); - // i4 value 0 packed: 0x00 - memset(weights + i * qk / 2, 0, qk / 2); - continue; - } - const float id = 1.0f / d; - scales[i] = ov::float16(d); - for (int j = 0; j < qk / 2; ++j) { - const float x0 = x[i * qk + 2 * j] * id; - const float x1 = x[i * qk + 2 * j + 1] * id; - // Signed i4: range [-8, 7]. Quantize as round(x*id), then pack as 4-bit two's complement. - int8_t si0 = (int8_t) std::max(-8, std::min(7, (int) roundf(x0))); - int8_t si1 = (int8_t) std::max(-8, std::min(7, (int) roundf(x1))); - weights[i * qk / 2 + j] = (si0 & 0x0F) | ((si1 & 0x0F) << 4); - } - } - } -} - -void quantize_q8_0(const float * x, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr, - int64_t k, - int64_t qk) { - assert(k % qk == 0); - const int nb = k / qk; - - auto * weights = static_cast(weights_arr.data()); - auto * scales = scales_arr.data::value_type>(); - bool is_symmetric = (weights_arr.get_element_type() == ov::element::i8); // Signed i8 path - - if (!is_symmetric) { - auto * zp = static_cast(zp_arr.data()); - for (int i = 0; i < nb; i++) { - float amax = 0.0f; - for (int j = 0; j < qk; j++) { - const float v = x[i * qk + j]; - amax = std::max(amax, fabsf(v)); - } - const float d = amax / 127.0f; - const float id = d ? 1.0f / d : 0.0f; - scales[i] = ov::float16(d); - zp[i] = 128; - for (int j = 0; j < qk; ++j) { - const float x0 = x[i * qk + j] * id; - const int8_t xi0 = roundf(x0); - weights[i * qk + j] = (uint8_t) (xi0 + 128); - } - } - } else { - // Symmetric: store signed int8 values directly - auto * signed_weights = reinterpret_cast(weights); - for (int i = 0; i < nb; i++) { - float amax = 0.0f; - for (int j = 0; j < qk; j++) { - const float v = x[i * qk + j]; - amax = std::max(amax, fabsf(v)); - } - const float d = amax / 127.0f; - const float id = d ? 1.0f / d : 0.0f; - scales[i] = ov::float16(d); - for (int j = 0; j < qk; ++j) { - const float x0 = x[i * qk + j] * id; - signed_weights[i * qk + j] = (int8_t) roundf(x0); - } - } - } -} - -void quantize_q8_1(const float * x, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr, - int64_t k, - int64_t qk) { - assert(k % qk == 0); - const int nb = k / qk; - - auto * weights = static_cast(weights_arr.data()); - auto * scales = scales_arr.data::value_type>(); - auto * zp = static_cast(zp_arr.data()); - for (int i = 0; i < nb; i++) { - float min = std::numeric_limits::max(); - float max = std::numeric_limits::lowest(); - - for (int j = 0; j < qk; j++) { - const float v = x[i * qk + j]; - min = std::min(v, min); - max = std::max(v, max); - } - - const float d = (max - min) / ((1 << 8) - 1); - const float id = d ? 1.0f / d : 0.0f; - scales[i] = ov::float16(d); - // zp = -min / scale (Q8_1 is asymmetric) - zp[i] = (d != 0.0f) ? (uint8_t) std::round(-min / d) : 0; - - for (int j = 0; j < qk; ++j) { - const float x0 = (x[i * qk + j] - min) * id; - const uint8_t xi0 = roundf(x0); - weights[i * qk + j] = xi0; - } - } -} diff --git a/ggml/src/ggml-openvino/ggml-quants.h b/ggml/src/ggml-openvino/ggml-quants.h deleted file mode 100644 index e4a02297cae4..000000000000 --- a/ggml/src/ggml-openvino/ggml-quants.h +++ /dev/null @@ -1,153 +0,0 @@ -#pragma once -#include "ggml-openvino-extra.h" // For ExtraQuantType -#include "ggml.h" - -#include -#include -#include - -void unpack_32_4(const uint8_t* data, uint8_t* dst); - -void extract_q4_0_data(const ggml_tensor * tensor, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr); - -void extract_q4_1_data(const ggml_tensor * tensor, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr, - bool use_bias = false); - -void extract_q8_0_data(const ggml_tensor * tensor, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr); - -void unpack_256_4(const uint8_t* data, uint8_t* dst); - -void extract_q4_k_data(const ggml_tensor * tensor, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr, - bool use_bias = false); - -void extract_q5_k_data(const ggml_tensor * tensor, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr, - bool use_bias = false); - -void extract_q6_k_data(const ggml_tensor * tensor, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr); - -static constexpr size_t GGML_QUANTIZATION_GROUP_SIZE = 32; - -ov::Output make_int8_weights(ov::Tensor & weight, - ov::Tensor & scales, - ov::Tensor & zp, - size_t group_size = GGML_QUANTIZATION_GROUP_SIZE, - bool use_bias = false); - -ov::Output make_int4_weights(ov::Tensor & weight, - ov::Tensor & scales, - ov::Tensor & zp, - size_t group_size = GGML_QUANTIZATION_GROUP_SIZE, - bool use_bias = false); - -// Extract quantized weights from tensor and create weight subgraph -// If weights/scales/zp are provided (non-empty), uses them as output buffers -// Otherwise allocates new ov::Tensors internally -// Returns the weight node (make_int4_weights or make_int8_weights result) -std::shared_ptr extract_quantized_weights( - const ggml_tensor * tensor, - const void * data, // Source data pointer (may differ from tensor->data) - ov::Tensor & weights, - ov::Tensor & scales, - ov::Tensor & zp, - bool use_bias = false); // Use fp bias instead of quantized zero_point (for test-backend-ops) - -// Requantize weights from tensor to target format, writing to provided buffers -// For F16 target, only weights buffer is used (scales/zp ignored) -// Returns the weight node -std::shared_ptr requantize_to_buffers(const ggml_tensor * tensor, - const void * data, // Source data pointer - ExtraQuantType requant_type, - int64_t block_size, - ov::Tensor & weights, - ov::Tensor & scales, - ov::Tensor & zp); - -inline const char * extra_quant_type_name(ExtraQuantType t) { - switch (t) { - case ExtraQuantType::F16: - return "F16"; - case ExtraQuantType::Q4_0_C: - return "Q4_0_C"; - case ExtraQuantType::Q4_0_128: - return "Q4_0_128"; - case ExtraQuantType::Q8_0_C: - return "Q8_0_C"; - case ExtraQuantType::Q8_0_32: - return "Q8_0_32"; - case ExtraQuantType::Q8_1_C: - return "Q8_1_C"; - default: - return "unknown"; - } -} - -// Result from process_weight_tensor containing the weight node and tensors. -// For quantized weights, also contains the extracted layout and scale/zp tensors. -struct OvWeight { - std::shared_ptr weight_node; - ggml_openvino_extracted_layout layout; // Only meaningful for quantized (layout.total_size > 0) - ov::Tensor weights; - ov::Tensor scales; - ov::Tensor zp; - - bool is_quantized() const { return layout.scales_size > 0; } -}; - -// Process weight tensor and create an OpenVINO weight node -// Handles F16/F32/BF16 and quantized weights, with optional requantization -// If output_base_ptr is nullptr, allocates internal buffers (for decoder use) -// If output_base_ptr is provided, uses pre-allocated buffers at specified offsets (for backend buffer use) -// Returns OvWeight with the weight node and optional quantized tensors -OvWeight process_weight_tensor( - const ggml_tensor * tensor, - const void * data, // Source data pointer (may differ from tensor->data) - void * output_base_ptr = nullptr, // Base pointer for output buffers (or nullptr for internal allocation) - bool use_bias = false); // Use fp bias instead of quantized zero_point, only used in test-backend-ops - -void quantize_q4_0(const float * x, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr, - int64_t k, - int64_t qk); -void quantize_q8_1(const float * x, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr, - int64_t k, - int64_t qk); -void quantize_q8_0(const float * x, - ov::Tensor & weights_arr, - ov::Tensor & scales_arr, - ov::Tensor & zp_arr, - int64_t k, - int64_t qk); - -namespace ov { -namespace op { -namespace util { -// From /src/common/transformations/include/transformations/utils/utils.hpp -bool get_single_value(const std::shared_ptr& const_node, - float& value, - bool check_value_range = true); -} // namespace util -} // namespace op -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/decoder.h b/ggml/src/ggml-openvino/openvino/decoder.h deleted file mode 100644 index 3b8da2be5d2b..000000000000 --- a/ggml/src/ggml-openvino/openvino/decoder.h +++ /dev/null @@ -1,74 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { - -class GgmlDecoder : public DecoderBase { -public: - virtual ov::Any get_attribute(const std::string& name) const = 0; - - virtual PartialShape get_input_shape(int node_idx, const std::string& name) const = 0; - - virtual std::vector get_input_stride(int node_idx, const std::string& name) const = 0; - - virtual element::Type get_input_type(int node_idx, const std::string& name) const = 0; - - virtual size_t get_input_size() const = 0; - - virtual size_t get_input_size(int node_idx) const = 0; - - virtual void get_input_node(size_t input_port_idx, - std::string& producer_name, - std::string& producer_output_port_name, - size_t& producer_output_port_index) const = 0; - - virtual std::vector get_input_names(int node_idx) const = 0; - - virtual PartialShape get_output_shape(int node_idx) const = 0; - - virtual element::Type get_output_type(const int node_idx) const = 0; - - virtual int32_t* get_input_op_params(int node_idx, const std::string& name) const = 0; - - virtual int32_t * get_output_op_params(int node_idx) const = 0; - - virtual std::vector get_output_names(int node_idx) const = 0; - - virtual const std::string& get_op_type() const = 0; - - virtual const std::string& get_op_type(int node_idx) const = 0; - - virtual const std::string& get_op_name() const = 0; - - virtual const std::string& get_op_name(int node_idx) const = 0; - - virtual void visit_subgraph(std::function, int node_idx)> node_visitor) const = 0; - - virtual int get_op_case(int node_idx) const = 0; - - virtual const std::map>& get_model_inputs() const = 0; - virtual const std::map>& get_model_extra_inputs() const = 0; - virtual const std::map>& get_model_weights() const = 0; - virtual std::vector get_model_output_names() const = 0; - - virtual int32_t* get_rope_params() const = 0; - - virtual std::map get_kv_param_res_names() const = 0; - - virtual bool is_static() const = 0; - - virtual bool is_stateful() const = 0; - - virtual int is_swa_layer(int layer) const = 0; -}; - -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/frontend.cpp b/ggml/src/ggml-openvino/openvino/frontend.cpp deleted file mode 100644 index c2ba14e66e6e..000000000000 --- a/ggml/src/ggml-openvino/openvino/frontend.cpp +++ /dev/null @@ -1,27 +0,0 @@ -#include "frontend.h" - -#include "input_model.h" -#include "op_table.h" -#include "translate_session.h" - -namespace ov { -namespace frontend { -namespace ggml { - -FrontEnd::FrontEnd() {} - -std::shared_ptr FrontEnd::convert(const InputModel::Ptr & model, bool naive) { - auto ggml_model = std::dynamic_pointer_cast(model); - FRONT_END_GENERAL_CHECK(ggml_model, "Invalid input model"); - std::shared_ptr converted_model; - const auto & supported_ops = get_supported_ops(); - { - TranslateSession translate_session(model, supported_ops, naive); - converted_model = translate_session.get_converted_model(); - } - return converted_model; -} - -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/frontend.h b/ggml/src/ggml-openvino/openvino/frontend.h deleted file mode 100644 index f1c6f0c3e3ce..000000000000 --- a/ggml/src/ggml-openvino/openvino/frontend.h +++ /dev/null @@ -1,23 +0,0 @@ -// Copyright (C) 2018-2024 Intel Corporation -// SPDX-License-Identifier: Apache-2.0 -// - -#pragma once - -#include - -namespace ov { -namespace frontend { -namespace ggml { - -class FrontEnd { -public: - using Ptr = std::shared_ptr; - FrontEnd(); - - static std::shared_ptr convert(const InputModel::Ptr& model, bool naive = false); -}; - -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/input_model.cpp b/ggml/src/ggml-openvino/openvino/input_model.cpp deleted file mode 100644 index 39b004c93176..000000000000 --- a/ggml/src/ggml-openvino/openvino/input_model.cpp +++ /dev/null @@ -1,17 +0,0 @@ -#include "input_model.h" - -#include "decoder.h" - -namespace ov { -namespace frontend { -namespace ggml { - -InputModel::InputModel(const std::shared_ptr & gdecoder) : m_decoder(gdecoder) {} - -const std::shared_ptr & InputModel::get_model_decoder() const { - return m_decoder; -} - -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/input_model.h b/ggml/src/ggml-openvino/openvino/input_model.h deleted file mode 100644 index ce8434426c90..000000000000 --- a/ggml/src/ggml-openvino/openvino/input_model.h +++ /dev/null @@ -1,29 +0,0 @@ -#pragma once - -#include - -#include "decoder.h" - -namespace ov { -namespace frontend { -namespace ggml { - -class FrontEnd; -class GgmlDecoder; -using ov::frontend::ggml::GgmlDecoder; - -class InputModel : public ov::frontend::InputModel { - friend class ::ov::frontend::ggml::FrontEnd; - -public: - explicit InputModel(const std::shared_ptr& gdecoder); - - const std::shared_ptr& get_model_decoder() const; - -private: - std::shared_ptr m_decoder; -}; - -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/node_context.h b/ggml/src/ggml-openvino/openvino/node_context.h deleted file mode 100644 index aa484128a952..000000000000 --- a/ggml/src/ggml-openvino/openvino/node_context.h +++ /dev/null @@ -1,112 +0,0 @@ -#pragma once - -#include -#include -#include - -#include "decoder.h" - -namespace ov { -namespace frontend { -namespace ggml { - -class TranslateSession; - -typedef std::map> TensorMap; - -class NodeContext : public frontend::NodeContext { -public: - NodeContext(const std::shared_ptr& decoder, - std::shared_ptr& tensor_map, - int node_idx, - TranslateSession* translate_session = nullptr) - : ov::frontend::NodeContext(decoder->get_op_type(node_idx)), - m_decoder(decoder), - m_tensor_map(tensor_map), - m_node_idx(node_idx), - m_translate_session(translate_session) { - m_input_names = decoder->get_input_names(m_node_idx); - m_output_names = decoder->get_output_names(m_node_idx); - } - - TranslateSession* get_translate_session() const { - return m_translate_session; - } - - const std::vector& get_input_names() const { return m_input_names; } - - size_t get_input_size() const override { - return m_decoder->get_input_size(m_node_idx); - } - - ov::element::Type get_input_type(size_t index) const { - return m_decoder->get_input_type(m_node_idx, m_input_names[index]); - } - - PartialShape get_input_shape(size_t input_index) const { - return m_decoder->get_input_shape(m_node_idx, m_input_names[input_index]); - } - - std::vector get_input_stride(size_t index) const { - return m_decoder->get_input_stride(m_node_idx, m_input_names[index]); - } - - std::string get_output_name() const { return m_output_names[0]; } - - PartialShape get_output_shape() const { return m_decoder->get_output_shape(m_node_idx); } - - int32_t* get_input_op_params(size_t index) const { - return m_decoder->get_input_op_params(m_node_idx, m_input_names[index]); - } - - int32_t * get_output_op_params() const { return m_decoder->get_output_op_params(m_node_idx); } - - ov::element::Type get_output_type() const { - return m_decoder->get_output_type(m_node_idx); - } - - Output get_input(int idx) const override { - return m_tensor_map->at(m_input_names[idx]); - } - - Output get_input(const std::string& name) const override { - if (m_tensor_map->find(name) == m_tensor_map->end()) { - throw std::runtime_error("'" + name + "' not found in tensor map."); - } - return m_tensor_map->at(name); - } - - bool has_input(const std::string& name) const { - return m_tensor_map->find(name) != m_tensor_map->end(); - } - - const std::string& get_name() const override { - return m_decoder->get_op_name(m_node_idx); - } - - ov::Any get_attribute_as_any(const std::string& name) const override { - return m_decoder->get_attribute(name); - } - - int get_op_case() const { - return m_decoder->get_op_case(m_node_idx); - } - - bool is_static() const { return m_decoder->is_static(); } - - bool is_stateful() const { return m_decoder->is_stateful(); } - -private: - std::shared_ptr m_decoder; - std::shared_ptr& m_tensor_map; - int m_node_idx; - TranslateSession* m_translate_session; - std::vector m_input_names; - std::vector m_output_names; -}; - -using CreatorFunction = std::function; - -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/cont.cpp b/ggml/src/ggml-openvino/openvino/op/cont.cpp deleted file mode 100644 index 6160dd744446..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/cont.cpp +++ /dev/null @@ -1,48 +0,0 @@ - -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_cont(const NodeContext & context) { - num_inputs_check(context, 1, 1); - - int op_case = context.get_op_case(); - FRONT_END_CHECK_IMPLEMENTED(op_case == 1 || op_case == 2 || op_case == 3, "Unsupported CONT case"); - - auto src_shape = context.get_input_shape(0).to_shape(); - auto dst_shape = context.get_output_shape().to_shape(); - ov::Output res; - - if (op_case == 1) { - // The input comes from a PERMUTE - throw std::runtime_error("Code of this case might be outdated"); - dst_shape[1] = -1; - res = std::make_shared( - context.get_input(0), ov::op::v0::Constant::create(ov::element::i64, {dst_shape.size()}, dst_shape), false); - } else if (op_case == 2) { - // The input comes from a TRANSPOSE - return {context.get_input(0)}; - } else { - // The input comes from a VIEW - res = process_view_input(context, 0); - } - - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/cpy.cpp b/ggml/src/ggml-openvino/openvino/op/cpy.cpp deleted file mode 100644 index 831117208be4..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/cpy.cpp +++ /dev/null @@ -1,21 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_cpy(const NodeContext & context) { - auto res = std::make_shared(context.get_input(0), context.get_output_type()); - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp b/ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp deleted file mode 100644 index 42602a730a4f..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp +++ /dev/null @@ -1,90 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_flash_attn_ext(const NodeContext & context) { - num_inputs_check(context, 4, 4); - auto q_f32 = context.get_input(0); - auto k = context.get_input(1); - auto v = context.get_input(2); - auto mask = context.get_input(3); - - float * params = reinterpret_cast(context.get_output_op_params()); - float scale = params[0]; - // float max_bias = params[1]; - // float logit_softcap = params[2]; - - auto q = std::make_shared(q_f32, ov::element::f16); - auto scale_node = std::make_shared(ov::element::f16, ov::Shape{}, std::vector{scale}); - - ov::Output mask_sliced, res; - std::string mask_name = "KQ_mask_sliced"; - if (context.get_input_names()[3].find("swa") != std::string::npos) { - mask_name = "KQ_mask_swa_sliced"; - } - if (context.has_input(mask_name)) { - mask_sliced = context.get_input(mask_name); - } else { - auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); - auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); - auto two = ov::op::v0::Constant::create(ov::element::i64, {1}, {2}); - auto token_len = get_dimensions(q, {2}); - mask_sliced = std::make_shared(mask, zero, token_len, one, two); - } - - if (mask_sliced.get_element_type() != ov::element::f16) { - mask_sliced = std::make_shared(mask_sliced, ov::element::f16); - } - - auto tile_kv = [&](int64_t num_heads, int64_t num_heads_kv, int64_t head_size, ov::Output kv) { - int64_t factor = num_heads / num_heads_kv; - if (factor > 1 && num_heads_kv > 1) { - ov::Output kv_broadcast_shape, kv_unsqueezed, new_kv_shape; - auto unsqueeze_axes = ov::op::v0::Constant::create(ov::element::i64, Shape{}, {2}); - kv_unsqueezed = std::make_shared(kv, unsqueeze_axes); - - kv_broadcast_shape = ov::op::v0::Constant::create( - ov::element::i64, {5}, {(int64_t) 1, (int64_t) 1, factor, (int64_t) 1, (int64_t) 1}); - new_kv_shape = - ov::op::v0::Constant::create(ov::element::i64, {4}, {(int64_t) 0, num_heads, (int64_t) -1, head_size}); - - kv = std::make_shared(kv_unsqueezed, kv_broadcast_shape, - ov::op::BroadcastType::BIDIRECTIONAL); - kv = std::make_shared(kv, new_kv_shape, true); - } - return kv; - }; - - auto q_shape = context.get_input_shape(0).to_shape(); - auto k_shape = context.get_input_shape(1).to_shape(); - k = tile_kv(q_shape[1], k_shape[1], q_shape[3], k); - v = tile_kv(q_shape[1], k_shape[1], q_shape[3], v); - - auto sdpa = std::make_shared(q, k, v, mask_sliced, scale_node, false); - res = std::make_shared(sdpa, - ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3})); - res = std::make_shared(res, ov::element::f32); - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/get_rows.cpp b/ggml/src/ggml-openvino/openvino/op/get_rows.cpp deleted file mode 100644 index 49f51b7ca3fc..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/get_rows.cpp +++ /dev/null @@ -1,69 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_get_rows(const NodeContext & context) { - num_inputs_check(context, 2, 2); - - int op_case = context.get_op_case(); - - Output res; - auto data = context.get_input(0); - auto indices = context.get_input(1); - - if (op_case == 2) { - // The input comes from a VIEW - indices = process_view_input(context, 1); - } - - // data[1,b,x,y] ind[1,1,b,x'] test-backend-ops case - // data[x,y] ind[1,1,1,x'] normal case - indices = - std::make_shared(indices, ov::op::v0::Constant::create(ov::element::i64, {2}, {0, 1})); - if (data.get_partial_shape().rank() == 4) { - if (!(data.get_partial_shape()[1].is_dynamic()) && data.get_partial_shape()[1].get_length() == 1) { - // Work-around for a bug in ov cpu plugin for test-backend-ops - data = std::make_shared(data, - ov::op::v0::Constant::create(ov::element::i64, {2}, {0, 1})); - auto axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {0}); - res = std::make_shared(data, indices, axis); - } else { - auto axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {1}); - data = - std::make_shared(data, ov::op::v0::Constant::create(ov::element::i64, {1}, {0})); - res = std::make_shared(data, indices, axis, 1); - } - } else if (context.is_stateful() && data.get_partial_shape().rank() == 3) { - auto axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {1}); - res = std::make_shared(data, indices, axis, 1); - } else { - auto axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {0}); - res = std::make_shared(data, indices, axis); - } - - if (res.get_element_type() != context.get_output_type()) { - res = std::make_shared(res, context.get_output_type()); - } - if (!(context.is_stateful())) { - res = std::make_shared(res, ov::op::v0::Constant::create(ov::element::i64, {1}, {0})); - } - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/glu_geglu.cpp b/ggml/src/ggml-openvino/openvino/op/glu_geglu.cpp deleted file mode 100644 index d9fa4c24367c..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/glu_geglu.cpp +++ /dev/null @@ -1,61 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_glu_geglu(const NodeContext & context) { - num_inputs_check(context, 1, 2); - - ov::Output src0; - ov::Output src1; - if (context.get_input_size() == 2) { - src0 = context.get_input(0); - src1 = context.get_input(1); - } else { - // GGML splits along ne[0] (OV last axis) using floor division: nc = ne[0] / 2. - // Both halves are nc elements; if the dimension is odd, the last element is dropped. - // Use Slice instead of Split to handle odd dimensions correctly. - auto combined = context.get_input(0); - auto combined_shape = combined.get_partial_shape(); - int64_t last_dim_val = combined_shape[combined_shape.rank().get_length() - 1].get_length(); - int64_t nc = last_dim_val / 2; - - auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}); - auto step = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); - auto start0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); - auto stop0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc}); - auto start1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc}); - auto stop1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {2 * nc}); - - src0 = std::make_shared(combined, start0, stop0, step, axis); - src1 = std::make_shared(combined, start1, stop1, step, axis); - } - - int32_t * params = context.get_output_op_params(); - const int32_t swapped = params[1]; - if (swapped) { - std::swap(src0, src1); - } - - auto gelu = std::make_shared(src0); - auto res = std::make_shared(gelu, src1); - - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/glu_swiglu.cpp b/ggml/src/ggml-openvino/openvino/op/glu_swiglu.cpp deleted file mode 100644 index 00ed7951a03d..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/glu_swiglu.cpp +++ /dev/null @@ -1,62 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_glu_swiglu(const NodeContext & context) { - num_inputs_check(context, 1, 2); - - ov::Output src0; - ov::Output src1; - if (context.get_input_size() == 2) { - src0 = context.get_input(0); - src1 = context.get_input(1); - } else { - // GGML splits along ne[0] (OV last axis) using floor division: nc = ne[0] / 2. - // Both halves are nc elements; if the dimension is odd, the last element is dropped. - // Use Slice instead of Split to handle odd dimensions correctly. - auto combined = context.get_input(0); - auto combined_shape = combined.get_partial_shape(); - int64_t last_dim_val = combined_shape[combined_shape.rank().get_length() - 1].get_length(); - int64_t nc = last_dim_val / 2; - - auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}); - auto step = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); - auto start0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); - auto stop0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc}); - auto start1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc}); - auto stop1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {2 * nc}); - - src0 = std::make_shared(combined, start0, stop0, step, axis); - src1 = std::make_shared(combined, start1, stop1, step, axis); - } - - int32_t * params = context.get_output_op_params(); - const int32_t swapped = params[1]; - if (swapped) { - std::swap(src0, src1); - } - - auto sigmoid = std::make_shared(src0); - auto silu = std::make_shared(src0, sigmoid); - auto res = std::make_shared(silu, src1); - - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/mulmat.cpp b/ggml/src/ggml-openvino/openvino/op/mulmat.cpp deleted file mode 100644 index 38edec85ddf7..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/mulmat.cpp +++ /dev/null @@ -1,90 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_mulmat(const NodeContext & context) { - num_inputs_check(context, 2, 2); - - int op_case = context.get_op_case(); - - ov::Output res; - ov::Output B = context.get_input(0); - ov::Output A = context.get_input(1); - - bool transpose_b = true; - if (op_case == 2) { - B = B.get_node_shared_ptr()->input_value(0); - transpose_b = false; - } else if (op_case == 3) { - B = process_view_input(context, 0); - A = process_view_input(context, 1); - } - if (A.get_element_type() != B.get_element_type()) { - B = std::make_shared(context.get_input(0), context.get_input_type(1)); - } - - auto B_shape = context.get_input_shape(0).to_shape(); - auto A_shape = context.get_input_shape(1).to_shape(); - int64_t A_batch = A_shape[1]; - int64_t B_batch = B_shape[1]; - - auto A_batch_larger = A_batch > B_batch; - auto batch_large = A_batch_larger ? A_batch : B_batch; - auto batch_small = A_batch_larger ? B_batch : A_batch; - - Output Z = A_batch_larger ? B : A; - int64_t factor = batch_large / batch_small; - if (factor > 1 && batch_small > 1) { - auto batch_large_node = ov::op::v0::Constant::create(ov::element::i64, {1}, std::vector{batch_large}); - auto batch_small_node = ov::op::v0::Constant::create(ov::element::i64, {1}, std::vector{batch_small}); - auto factor_node = ov::op::v0::Constant::create(ov::element::i64, {1}, std::vector{factor}); - - auto unsqueeze_axes = ov::op::v0::Constant::create(ov::element::i64, Shape{}, {2}); - auto Z_unsqueezed = std::make_shared(Z, unsqueeze_axes); - - auto broadcast_shape = ov::op::v0::Constant::create( - ov::element::i64, {5}, {(int64_t) 1, (int64_t) 1, factor, (int64_t) 1, (int64_t) 1}); - auto new_Z_shape = ov::op::v0::Constant::create(ov::element::i64, {4}, - {(int64_t) 0, batch_large, (int64_t) -1, (int64_t) A_shape[3]}); - - auto Z_broadcasted = std::make_shared(Z_unsqueezed, broadcast_shape, - ov::op::BroadcastType::BIDIRECTIONAL); - Z = std::make_shared(Z_broadcasted, new_Z_shape, true); - } - if (A_batch_larger) { - B = Z; - } else { - A = Z; - } - - res = std::make_shared(A, B, false, transpose_b); - - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/permute.cpp b/ggml/src/ggml-openvino/openvino/op/permute.cpp deleted file mode 100644 index 4c800f9ee4f6..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/permute.cpp +++ /dev/null @@ -1,102 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_permute(const NodeContext & context) { - num_inputs_check(context, 1, 1); - - int op_case = context.get_op_case(); - FRONT_END_CHECK_IMPLEMENTED(op_case == 1 || op_case == 2 || op_case == 3 || op_case == 4, - "Unsupported PERMUTE case"); - - ov::Output res; - auto src = context.get_input(0); - auto perm = ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3}); - - if (op_case == 1 || context.is_stateful()) { - res = std::make_shared(src, perm); - } else if (op_case == 4) { - auto output_shape = context.get_output_shape().to_shape(); - auto n_heads = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[1]}); - auto head_size = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[3]}); - auto n_seq_active = context.has_input("n_seq_active") ? - context.get_input("n_seq_active") : - ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[0]}); - auto neg_one = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}); - - auto new_shape = - std::make_shared(ov::OutputVector{n_seq_active, neg_one, n_heads, head_size}, 0); - - // // Alternative - // auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); - // auto new_shape = std::make_shared(ov::OutputVector{n_seq_active, neg_one, zero, zero}, 0); - - auto reshaped = std::make_shared(src, new_shape, true); - res = std::make_shared(reshaped, perm); - } else { - auto cache_shape = src.get_partial_shape(); - auto output_shape = context.get_output_shape().to_shape(); - int64_t head_size = output_shape[3]; - int64_t n_heads = output_shape[1]; - int64_t ctx_per_seq = cache_shape[2].is_static() ? cache_shape[2].get_length() : -1; - int64_t n_seq = cache_shape[1].get_length(); - - Output attention_size; - if (!context.has_input("attention_size")) { - attention_size = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[2]}); - } else if (op_case == 2) { - attention_size = context.get_input("attention_size"); - } else { - attention_size = context.get_input("attention_size_swa"); - } - - Output seq_active_start; - Output seq_active_end; - if (context.has_input("seq_active_start")) { - seq_active_start = context.get_input("seq_active_start"); - seq_active_end = context.get_input("seq_active_end"); - } else { - int64_t n_seq_active = output_shape[0]; - size_t offset = *((size_t *) context.get_input_op_params(0)); - int64_t seq_active_start_val = offset / context.get_input_stride(0)[0]; - int64_t seq_active_end_val = seq_active_start_val + n_seq_active; - seq_active_start = ov::op::v0::Constant::create(ov::element::i64, {1}, {seq_active_start_val}); - seq_active_end = ov::op::v0::Constant::create(ov::element::i64, {1}, {seq_active_end_val}); - } - - // 1. reshape to [n_seq, ctx_per_seq, n_heads, head_size] - // 2. slice out the active sequences - // 3. slice out the attention part in each sequence - // 4. permute - auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); - auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); - - auto src_reshaped = std::make_shared( - src, ov::op::v0::Constant::create(ov::element::i64, {4}, {n_seq, ctx_per_seq, n_heads, head_size}), false); - auto slice1 = std::make_shared(src_reshaped, seq_active_start, seq_active_end, one, zero); - auto slice2 = std::make_shared(slice1, zero, attention_size, one, one); - res = std::make_shared(slice2, perm); - } - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/reshape.cpp b/ggml/src/ggml-openvino/openvino/op/reshape.cpp deleted file mode 100644 index efd9a5a860ab..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/reshape.cpp +++ /dev/null @@ -1,83 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_reshape(const NodeContext & context) { - num_inputs_check(context, 1, 1); - if (context.get_input_shape(0) == context.get_output_shape()) { - return {context.get_input(0)}; - } - - int op_case = context.get_op_case(); - FRONT_END_CHECK_IMPLEMENTED( - op_case == 1 || op_case == 2 || op_case == 3 || op_case == 4 || op_case == 5 || op_case == 6, - "Unsupported RESHAPE case"); - - auto output_shape = context.get_output_shape().to_shape(); - std::shared_ptr new_shape_node; - if (op_case == 1) { - if (context.is_stateful()) { - new_shape_node = ov::op::v0::Constant::create( - ov::element::i64, {3}, - std::vector{-1, (int64_t) output_shape[2], (int64_t) output_shape[3]}); - } else { - new_shape_node = ov::op::v0::Constant::create( - ov::element::i64, {4}, - std::vector{(int64_t) output_shape[0], -1, (int64_t) output_shape[2], (int64_t) output_shape[3]}); - } - } else if (op_case == 2) { - new_shape_node = ov::op::v0::Constant::create( - ov::element::i64, {4}, - std::vector{(int64_t) output_shape[0], (int64_t) output_shape[1], -1, (int64_t) output_shape[3]}); - - } else if (op_case == 3) { - throw std::runtime_error("might be outdated RESHAPE case"); - new_shape_node = ov::op::v0::Constant::create( - ov::element::i64, {4}, std::vector{(int64_t) output_shape[0], (int64_t) output_shape[1], -1, 1}); - - } else if (op_case == 4) { - return {context.get_input(0).get_node_shared_ptr()->input_value(0)}; - - } else if (op_case == 5) { - if (context.is_stateful()) { - std::vector shape_vec = {1, -1, (int64_t) context.get_output_shape().to_shape()[3]}; - new_shape_node = ov::op::v0::Constant::create(ov::element::i64, {3}, shape_vec); - } else { - std::vector shape_vec = {1, 1, -1, (int64_t) context.get_output_shape().to_shape()[3]}; - new_shape_node = ov::op::v0::Constant::create(ov::element::i64, {4}, shape_vec); - } - - // // Alternative - // auto token_len = context.get_input("token_len"); - // auto emb_size = - // ov::op::v0::Constant::create(ov::element::i64, {1}, {(int64_t) context.get_output_shape().to_shape()[3]}); - // auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); - // new_shape_node = std::make_shared(ov::OutputVector{one, one, token_len, emb_size}, 0); - - } else if (op_case == 6) { - new_shape_node = ov::op::v0::Constant::create(ov::element::i64, {4}, context.get_output_shape().to_shape()); - } - auto res = std::make_shared(context.get_input(0), new_shape_node, false); - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/rms_norm.cpp b/ggml/src/ggml-openvino/openvino/op/rms_norm.cpp deleted file mode 100644 index 72cf92283e9e..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/rms_norm.cpp +++ /dev/null @@ -1,46 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_rms_norm(const NodeContext & context) { - num_inputs_check(context, 1, 1); - - auto input_node = context.get_input(0); - auto square = std::make_shared( - input_node, ov::op::v0::Constant::create(ov::element::f32, ov::Shape{1}, {2.0f})); - - auto mean = std::make_shared( - square, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {-1}), true); - - float eps; - memcpy(&eps, context.get_output_op_params(), sizeof(float)); - - auto rms = std::make_shared( - std::make_shared(mean, ov::op::v0::Constant::create(ov::element::f32, ov::Shape{1}, {eps}))); - - auto reciprocal = - std::make_shared(ov::op::v0::Constant::create(ov::element::f32, ov::Shape{1}, {1.0f}), rms); - - auto res = std::make_shared(input_node, reciprocal); - - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/rope.cpp b/ggml/src/ggml-openvino/openvino/op/rope.cpp deleted file mode 100644 index a8db9b38930f..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/rope.cpp +++ /dev/null @@ -1,149 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_rope(const NodeContext & context) { - num_inputs_check(context, 2, 3); - - int op_case = context.get_op_case(); - - ov::Output res; - - auto data_node = context.get_input(0).get_node_shared_ptr(); - auto output_shape = context.get_output_shape().to_shape(); - int32_t * op_params = context.get_output_op_params(); - const int mode = (op_case & 0xFFFF0000) >> 16; - op_case = (op_case & 0x0000FFFF); - - constexpr int TYPE_NORMAL = 0; - constexpr int TYPE_NEOX = 1; - constexpr int TYPE_IMROPE = 2; - - Output cos_theta_node; - Output sin_theta_node; - if (context.has_input("rope_cos")) { - cos_theta_node = context.get_input("rope_cos"); - sin_theta_node = context.get_input("rope_sin"); - } else { - auto inp_pos = context.get_input(1).get_node_shared_ptr(); - std::shared_ptr rope_freqs_weight; - if (context.get_input_size() == 3) { - rope_freqs_weight = context.get_input(2).get_node_shared_ptr(); - } - auto sin_cos = make_sin_cos(op_params, inp_pos, rope_freqs_weight, mode == TYPE_IMROPE); - sin_theta_node = sin_cos.first; - cos_theta_node = sin_cos.second; - } - - if (op_case == 2) { - // The input comes from a VIEW - int slice_len = output_shape[2] * output_shape[3]; - data_node = process_view_input(context, 0, slice_len).get_node_shared_ptr(); - if (context.is_stateful()) { - auto data_shape = ov::op::v0::Constant::create( - ov::element::i64, {3}, std::vector{-1, (int64_t) output_shape[2], (int64_t) output_shape[3]}); - data_node = std::make_shared(data_node, data_shape, false); - } else { - auto data_shape = ov::op::v0::Constant::create( - ov::element::i64, {4}, std::vector{1, -1, (int64_t) output_shape[2], (int64_t) output_shape[3]}); - data_node = std::make_shared(data_node, data_shape, false); - } - } - - if (mode == TYPE_NORMAL) { - auto neg_one = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}); - auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); - auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); - auto two = ov::op::v0::Constant::create(ov::element::i64, {1}, {2}); - auto end = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[3]}); - Output even_slice; - Output odd_slice; - int32_t unsqueeze_dim = context.is_stateful() ? 3 : 4; - even_slice = std::make_shared(data_node, zero, end, two, neg_one); - odd_slice = std::make_shared(data_node, one, end, two, neg_one); - - Output first_half = - std::make_shared(std::make_shared(even_slice, cos_theta_node), - std::make_shared(odd_slice, sin_theta_node)); - Output second_half = - std::make_shared(std::make_shared(even_slice, sin_theta_node), - std::make_shared(odd_slice, cos_theta_node)); - - first_half = std::make_shared(first_half, - ov::op::v0::Constant::create(ov::element::i64, {1}, {unsqueeze_dim})); - second_half = std::make_shared(second_half, - ov::op::v0::Constant::create(ov::element::i64, {1}, {unsqueeze_dim})); - auto stack = std::make_shared(OutputVector{first_half, second_half}, unsqueeze_dim); - - auto data_shape = ov::op::v0::Constant::create( - ov::element::i64, {4}, std::vector{1, -1, (int64_t) output_shape[2], (int64_t) output_shape[3]}); - res = std::make_shared(stack, data_shape, false); - } else if (mode == TYPE_NEOX) { - auto data_split = std::make_shared( - data_node, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {-1}), 2); - Output slice_data_node_0 = data_split->outputs()[0]; - Output slice_data_node_1 = data_split->outputs()[1]; - - auto first_half_node = std::make_shared( - std::make_shared(slice_data_node_0, cos_theta_node), - std::make_shared(slice_data_node_1, sin_theta_node)); - - auto second_half_node = std::make_shared( - std::make_shared(slice_data_node_0, sin_theta_node), - std::make_shared(slice_data_node_1, cos_theta_node)); - - res = std::make_shared(ov::OutputVector{first_half_node, second_half_node}, -1); - } else if (mode == TYPE_IMROPE) { - int64_t n_dims = data_node->get_shape()[3]; - auto cos_sin_shape = std::make_shared(ov::element::i64, ov::Shape{4}, std::vector{1,-1,1,(n_dims >> 1)}); - auto cos_reshaped = std::make_shared(cos_theta_node, cos_sin_shape, true); - auto sin_reshaped = std::make_shared(sin_theta_node, cos_sin_shape, true); - - auto split_axis = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {3}); - auto split_a = std::make_shared(data_node, split_axis, 2); - auto x0 = split_a->output(0); - auto x1 = split_a->output(1); - auto mul_a = std::make_shared(x0, cos_reshaped); - auto mul_b = std::make_shared(x1, sin_reshaped); - auto sub = std::make_shared(mul_a, mul_b); - - auto mul_c = std::make_shared(x0, sin_reshaped); - auto mul_d = std::make_shared(x1, cos_reshaped); - auto add = std::make_shared(mul_c, mul_d); - - res = std::make_shared(ov::OutputVector{sub, add}, 3); - } - - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/scale.cpp b/ggml/src/ggml-openvino/openvino/op/scale.cpp deleted file mode 100644 index 0f3d800c1990..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/scale.cpp +++ /dev/null @@ -1,41 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_scale(const NodeContext & context) { - num_inputs_check(context, 1, 1); - - float scale; - float bias; - memcpy(&scale, (float *) context.get_output_op_params() + 0, sizeof(float)); - memcpy(&bias, (float *) context.get_output_op_params() + 1, sizeof(float)); - - auto scale_node = std::make_shared(ov::element::f32, ov::Shape{}, std::vector{scale}); - auto scaled = std::make_shared(context.get_input(0), scale_node); - - std::shared_ptr res; - if (bias != 0.0f) { - auto bias_node = - std::make_shared(ov::element::f32, ov::Shape{}, std::vector{bias}); - res = std::make_shared(scaled, bias_node); - } else { - res = scaled; - } - - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/set_rows.cpp b/ggml/src/ggml-openvino/openvino/op/set_rows.cpp deleted file mode 100644 index 136e4265b429..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/set_rows.cpp +++ /dev/null @@ -1,76 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_set_rows(const NodeContext & context) { - num_inputs_check(context, 3, 3); - - auto data = context.get_input(0); - auto indices = context.get_input(1); - auto dst = context.get_input(2); - - data = std::make_shared(data, context.get_output_type()); - - auto dst_shape = context.get_output_shape().to_shape(); - - auto ind_squeezed = - std::make_shared(indices, ov::op::v0::Constant::create(ov::element::i64, {3}, {0, 1, 2})); - auto data_reshaped = std::make_shared( - data, - ov::op::v0::Constant::create(ov::element::i64, {4}, - {(int64_t) 1, (int64_t) 1, (int64_t) -1, (int64_t) dst_shape[3]}), - false); - auto axes = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {2}); - - Output res; - if (context.is_stateful()) { - int concat_axis = 1; - int64_t dim2 = dst.get_partial_shape()[2].get_length(); - int64_t dim3 = dst.get_partial_shape()[3].get_length(); - data = std::make_shared( - data, ov::op::v0::Constant::create(ov::element::i64, {4}, {(int64_t) 1, (int64_t) -1, dim2, dim3}), false); - res = std::make_shared(OutputVector{dst, data}, concat_axis); - } else { - res = std::make_shared(dst, ind_squeezed, data_reshaped, axes); - } - - if (auto dst_reshape = std::dynamic_pointer_cast(dst.get_node_shared_ptr())) { - // Fix the case of multiple sequences, reshape back to original shape [1, n_seq, ctx_per_seq, emb] - // ctx_per_seq is not fixed due to llama-bench compatibility - auto dst_shape_partial = dst_reshape->get_input_partial_shape(0); - std::vector dst_shape = {dst_shape_partial[0].get_length(), dst_shape_partial[1].get_length(), - dst_shape_partial[2].is_static() ? dst_shape_partial[2].get_length() : -1, - dst_shape_partial[3].get_length()}; - res = std::make_shared(res, ov::op::v0::Constant::create(ov::element::i64, {4}, dst_shape), - false); - } - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/softmax.cpp b/ggml/src/ggml-openvino/openvino/op/softmax.cpp deleted file mode 100644 index 9f6330862be4..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/softmax.cpp +++ /dev/null @@ -1,89 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_soft_max(const NodeContext & context) { - // TODO code is outdated - num_inputs_check(context, 1, 2); - - auto input_node = context.get_input(0).get_node_shared_ptr(); - ov::Output res; - - float scale = 1.0f; - float max_bias = 0.0f; - auto * op_params = context.get_output_op_params(); - memcpy(&scale, (float *) op_params + 0, sizeof(float)); - memcpy(&max_bias, (float *) op_params + 1, sizeof(float)); - auto src0_shape = context.get_input_shape(0).get_shape(); - const uint32_t h = src0_shape[2]; - const uint32_t n_head = src0_shape[0]; - const uint32_t n_head_log2 = 1u << (uint32_t) floor(log2(n_head)); - - const float m0 = powf(2.0f, -(max_bias) / n_head_log2); - const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); - const float slope = - (max_bias > 0.0f) ? h < n_head_log2 ? powf(m0, h + 1) : powf(m1, 2 * (h - n_head_log2) + 1) : 1.0f; - - auto scale_node = std::make_shared(ov::element::f32, ov::Shape{}, std::vector{scale}); - auto scaled_input = std::make_shared(input_node, scale_node); - - if (context.get_input_size() < 2) { - res = std::make_shared(scaled_input, 2); - return rename_outputs_with_suffix({res}, context.get_name()); - } - - ov::Output mask_node_sliced; - if (context.has_input("KQ_mask_sliced")) { - mask_node_sliced = context.get_input("KQ_mask_sliced"); - } else { - auto token_len = get_dimensions(input_node, {1}); - auto mask_node = context.get_input(1); - auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); - auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); - mask_node_sliced = std::make_shared(mask_node, zero, token_len, one, one); - } - - if (mask_node_sliced.get_element_type() != context.get_output_type()) { - mask_node_sliced = std::make_shared(mask_node_sliced, context.get_output_type()); - } - - Output slope_mask; - if (slope != 1.0f) { - auto slope_node = - std::make_shared(ov::element::f32, ov::Shape{}, std::vector{slope}); - slope_mask = std::make_shared(mask_node_sliced, slope_node); - throw std::runtime_error("Slope != 1.0f in softmax has not been tested, verify it before use."); - } - slope_mask = mask_node_sliced; - - auto input_slope_mask_node = std::make_shared(scaled_input, slope_mask); - - res = std::make_shared(input_slope_mask_node, 2); - - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/transpose.cpp b/ggml/src/ggml-openvino/openvino/op/transpose.cpp deleted file mode 100644 index 8e62e83c0d78..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/transpose.cpp +++ /dev/null @@ -1,23 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_transpose(const NodeContext & context) { - num_inputs_check(context, 1, 1); - - auto res = std::make_shared( - context.get_input(0), ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 1, 3, 2})); - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/unary_gelu.cpp b/ggml/src/ggml-openvino/openvino/op/unary_gelu.cpp deleted file mode 100644 index d1e9efc33a55..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/unary_gelu.cpp +++ /dev/null @@ -1,25 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_unary_gelu(const NodeContext & context) { - num_inputs_check(context, 1, 1); - - auto input = context.get_input(0); - auto res = std::make_shared(input); - - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/unary_silu.cpp b/ggml/src/ggml-openvino/openvino/op/unary_silu.cpp deleted file mode 100644 index 037e0b94df1f..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/unary_silu.cpp +++ /dev/null @@ -1,27 +0,0 @@ -#include "../node_context.h" -#include "../op_table.h" -#include "../utils.h" - -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_unary_silu(const NodeContext & context) { - num_inputs_check(context, 1, 1); - - auto input = context.get_input(0); - auto sigmoid = std::make_shared(input); - auto res = std::make_shared(input, sigmoid); - - return rename_outputs_with_suffix({res}, context.get_name()); -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/view.cpp b/ggml/src/ggml-openvino/openvino/op/view.cpp deleted file mode 100644 index 8528d2523367..000000000000 --- a/ggml/src/ggml-openvino/openvino/op/view.cpp +++ /dev/null @@ -1,53 +0,0 @@ -#include "../op_table.h" -#include "../utils.h" -#include -namespace ov { -namespace frontend { -namespace ggml { -namespace op { - -OutputVector translate_view(const NodeContext & context) { - num_inputs_check(context, 1, 1); - - if (context.get_op_case() == 2) { - auto dst_shape = context.get_output_shape().to_shape(); - return rename_outputs_with_suffix({process_view_input(context, 0, dst_shape[2] * dst_shape[3])}, - context.get_name()); - } - // op_case 3 - if (context.get_op_case() == 3) { - auto input = context.get_input(0); - auto input_ov_shape = input.get_partial_shape(); - - auto input_llama_shape = context.get_input_shape(0).to_shape(); - - // if the input ov shape size is different from the input llama shape size, it means the input is already reshaped and we need to reshape it back to the original shape before slicing - if (input_ov_shape.size() != input_llama_shape.size()) { - input = std::make_shared(input, ov::op::v0::Constant::create(ov::element::i64, {input_llama_shape.size()}, input_llama_shape), false); - } - - auto dst_shape = context.get_output_shape().to_shape(); - - // find the index of dst_shape that is different from input shape, and use that index to slice the input - int slice_dim = -1; - for (size_t i = 0; i < dst_shape.size(); ++i) { - if (dst_shape[i] != input_llama_shape[i]) { - slice_dim = i; - break; - } - } - - auto begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); - auto end = ov::op::v0::Constant::create(ov::element::i64, {1}, {dst_shape[slice_dim]}); - auto stride = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); - auto axes = ov::op::v0::Constant::create(ov::element::i64, {1}, {slice_dim}); - auto sliced = std::make_shared(input, begin, end, stride, axes); - return {sliced}; - } - return {context.get_input(0)}; -} - -} // namespace op -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op_table.cpp b/ggml/src/ggml-openvino/openvino/op_table.cpp deleted file mode 100644 index 1385539279cb..000000000000 --- a/ggml/src/ggml-openvino/openvino/op_table.cpp +++ /dev/null @@ -1,47 +0,0 @@ -#include "op_table.h" - -#include "utils.h" - -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { - -std::unordered_map get_supported_ops() { - using namespace ov::op; - return { - {"GGML_OP_ADD", op::translate_1to1_match_2_inputs }, - {"GGML_OP_ADD1", op::translate_1to1_match_2_inputs }, - {"GGML_OP_CONT", op::translate_cont }, - {"GGML_OP_DIV", op::translate_1to1_match_2_inputs }, - {"GGML_OP_GET_ROWS", op::translate_get_rows }, - {"GGML_OP_MUL", op::translate_1to1_match_2_inputs}, - {"GGML_OP_MUL_MAT", op::translate_mulmat }, - {"GGML_OP_PERMUTE", op::translate_permute }, - {"GGML_OP_RESHAPE", op::translate_reshape }, - {"GGML_OP_RMS_NORM", op::translate_rms_norm }, - {"GGML_OP_ROPE", op::translate_rope }, - {"GGML_OP_SCALE", op::translate_scale }, - {"GGML_OP_SOFT_MAX", op::translate_soft_max }, - {"GGML_OP_SUB", op::translate_1to1_match_2_inputs}, - {"GGML_OP_TRANSPOSE", op::translate_transpose }, - {"GGML_UNARY_OP_GELU", op::translate_unary_gelu }, - {"GGML_UNARY_OP_SILU", op::translate_unary_silu }, - {"GGML_OP_VIEW", op::translate_view }, - {"GGML_GLU_OP_SWIGLU", op::translate_glu_swiglu }, - {"GGML_GLU_OP_GEGLU", op::translate_glu_geglu }, - {"GGML_OP_SET_ROWS", op::translate_set_rows }, - {"GGML_OP_CPY", op::translate_cpy }, - {"GGML_OP_FLASH_ATTN_EXT", op::translate_flash_attn_ext }, - }; -} - -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op_table.h b/ggml/src/ggml-openvino/openvino/op_table.h deleted file mode 100644 index f546796d2ee0..000000000000 --- a/ggml/src/ggml-openvino/openvino/op_table.h +++ /dev/null @@ -1,40 +0,0 @@ -#pragma once - -#include "node_context.h" - -namespace ov { -namespace frontend { -namespace ggml { - -namespace op { - -#define GGML_OP_CONVERTER(op) OutputVector op(const NodeContext& context) - -GGML_OP_CONVERTER(translate_add); -GGML_OP_CONVERTER(translate_cont); -GGML_OP_CONVERTER(translate_get_rows); -GGML_OP_CONVERTER(translate_mul); -GGML_OP_CONVERTER(translate_mulmat); -GGML_OP_CONVERTER(translate_permute); -GGML_OP_CONVERTER(translate_reshape); -GGML_OP_CONVERTER(translate_rms_norm); -GGML_OP_CONVERTER(translate_rope); -GGML_OP_CONVERTER(translate_scale); -GGML_OP_CONVERTER(translate_unary_silu); -GGML_OP_CONVERTER(translate_unary_gelu); -GGML_OP_CONVERTER(translate_soft_max); -GGML_OP_CONVERTER(translate_transpose); -GGML_OP_CONVERTER(translate_view); -GGML_OP_CONVERTER(translate_glu_swiglu); -GGML_OP_CONVERTER(translate_glu_geglu); -GGML_OP_CONVERTER(translate_set_rows); -GGML_OP_CONVERTER(translate_cpy); -GGML_OP_CONVERTER(translate_flash_attn_ext); - -} // namespace op - -std::unordered_map get_supported_ops(); - -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/pass/fuse_to_sdpa.cpp b/ggml/src/ggml-openvino/openvino/pass/fuse_to_sdpa.cpp deleted file mode 100644 index 0671542ee383..000000000000 --- a/ggml/src/ggml-openvino/openvino/pass/fuse_to_sdpa.cpp +++ /dev/null @@ -1,60 +0,0 @@ -#include "fuse_to_sdpa.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { -namespace pass { - -FuseToSDPA::FuseToSDPA() { - // Not maintained since FLASH_ATTN_EXT has replaced this pattern - const auto m_k = ov::pass::pattern::any_input(); - const auto m_q = ov::pass::pattern::any_input(); - const auto m_qk = ov::pass::pattern::wrap_type({m_q, m_k}); - const auto m_qk_f32 = ov::pass::pattern::wrap_type({m_qk}); - const auto m_scale = ov::pass::pattern::any_input(); - const auto m_scaled_qk = ov::pass::pattern::wrap_type({m_qk_f32, m_scale}); - const auto m_mask = ov::pass::pattern::any_input(); - const auto m_masked_qk = ov::pass::pattern::wrap_type({m_scaled_qk, m_mask}); - const auto m_softmax_qk = ov::pass::pattern::wrap_type({m_masked_qk}); - const auto m_softmax_qk_f16 = ov::pass::pattern::wrap_type({m_softmax_qk}); - const auto m_v = ov::pass::pattern::any_input(); - const auto m_qkv = ov::pass::pattern::wrap_type({m_softmax_qk_f16, m_v}); - - const auto callback = [=](ov::pass::pattern::Matcher & m) { - auto & pattern_to_output = m.get_pattern_value_map(); - auto k = pattern_to_output[m_k]; - auto q = pattern_to_output[m_q]; - auto v = pattern_to_output[m_v]; - auto mask = pattern_to_output[m_mask]; - auto scale = pattern_to_output[m_scale]; - - auto mask_f16 = register_new_node(mask, ov::element::f16); - auto scale_f16 = register_new_node(scale, ov::element::f16); - auto sdpa = std::make_shared(q, k, v, mask_f16, scale_f16, false); - - ov::replace_node(m.get_match_root(), sdpa); - ov::copy_runtime_info(m.get_matched_nodes(), sdpa); - - return true; - }; - register_matcher(std::make_shared(m_qkv, "ov::frontend::ggml::pass::FuseToSDPA"), - callback); -} - -} // namespace pass -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/pass/fuse_to_sdpa.h b/ggml/src/ggml-openvino/openvino/pass/fuse_to_sdpa.h deleted file mode 100644 index 8b5164d23293..000000000000 --- a/ggml/src/ggml-openvino/openvino/pass/fuse_to_sdpa.h +++ /dev/null @@ -1,17 +0,0 @@ -#include "openvino/pass/matcher_pass.hpp" - -namespace ov { -namespace frontend { -namespace ggml { -namespace pass { - -class FuseToSDPA : public ov::pass::MatcherPass { -public: - OPENVINO_MATCHER_PASS_RTTI("ov::frontend::ggml::pass::FuseToSDPA") - FuseToSDPA(); -}; - -} // namespace pass -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/pass/mark_decompression_convert_constant_folding.h b/ggml/src/ggml-openvino/openvino/pass/mark_decompression_convert_constant_folding.h deleted file mode 100644 index b95385611e88..000000000000 --- a/ggml/src/ggml-openvino/openvino/pass/mark_decompression_convert_constant_folding.h +++ /dev/null @@ -1,29 +0,0 @@ -#pragma once - -#include "mark_decompression_convert_constant_folding.h" -#include "openvino/pass/matcher_pass.hpp" -#include "openvino/core/visibility.hpp" - -#ifdef OPENVINO_STATIC_LIBRARY -# define TRANSFORMATIONS_API -#else -# ifdef IMPLEMENT_OPENVINO_API -# define TRANSFORMATIONS_API OPENVINO_CORE_EXPORTS -# else -# define TRANSFORMATIONS_API OPENVINO_CORE_IMPORTS -# endif // IMPLEMENT_OPENVINO_API -#endif // OPENVINO_STATIC_LIBRARY - -namespace ov { -namespace pass { - -class TRANSFORMATIONS_API MarkCompressedFloatConstants; - -} // namespace pass -} // namespace ov - -class ov::pass::MarkCompressedFloatConstants : public MatcherPass { -public: - OPENVINO_MATCHER_PASS_RTTI("MarkCompressedFloatConstants") - MarkCompressedFloatConstants(); -}; diff --git a/ggml/src/ggml-openvino/openvino/pass/squeeze_matmul.cpp b/ggml/src/ggml-openvino/openvino/pass/squeeze_matmul.cpp deleted file mode 100644 index 20a3a374934b..000000000000 --- a/ggml/src/ggml-openvino/openvino/pass/squeeze_matmul.cpp +++ /dev/null @@ -1,58 +0,0 @@ -#include "squeeze_matmul.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace opp = ov::pass::pattern; - -namespace ov { -namespace frontend { -namespace ggml { -namespace pass { - -// For quantized models, NPUW expects the activation to be 3d in DQ(DynamicQuantization) opt, e.g. DQMatMulGQ2i -SqueezeMatmul::SqueezeMatmul() { - auto m_act = opp::any_input(); - auto m_wei = opp::any_input(); - auto m_matmul = opp::wrap_type({m_act, m_wei}); - - const auto callback = [=](ov::pass::pattern::Matcher & m) { - const auto & pattern_map = m.get_pattern_value_map(); - auto matmul_node = - std::dynamic_pointer_cast(pattern_map.at(m_matmul).get_node_shared_ptr()); - auto act = pattern_map.at(m_act); - auto wei = pattern_map.at(m_wei); - auto act_shape = act.get_partial_shape(); - auto wei_shape = wei.get_partial_shape(); - if (act_shape.rank().is_dynamic() || wei_shape.rank().is_dynamic()) { - return false; - } - if (act_shape.rank().get_length() == 4 && wei_shape.rank().get_length() == 2) { - auto axis = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {0}); - auto squeezed_act = std::make_shared(act, axis); - auto new_matmul = std::make_shared(squeezed_act, wei, matmul_node->get_transpose_a(), - matmul_node->get_transpose_b()); - auto unsqueezed_output = std::make_shared(new_matmul, axis); - unsqueezed_output->set_friendly_name(matmul_node->get_friendly_name()); - ov::copy_runtime_info(matmul_node, {squeezed_act, new_matmul, unsqueezed_output}); - ov::replace_node(matmul_node, unsqueezed_output); - return true; - } - return false; - }; - - register_matcher(std::make_shared(m_matmul, "ov::frontend::ggml::pass::SqueezeMatmul"), - callback); -} - -} // namespace pass -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/pass/squeeze_matmul.h b/ggml/src/ggml-openvino/openvino/pass/squeeze_matmul.h deleted file mode 100644 index f8fbc69d546e..000000000000 --- a/ggml/src/ggml-openvino/openvino/pass/squeeze_matmul.h +++ /dev/null @@ -1,17 +0,0 @@ -#include "openvino/pass/matcher_pass.hpp" - -namespace ov { -namespace frontend { -namespace ggml { -namespace pass { - -class SqueezeMatmul : public ov::pass::MatcherPass { -public: - OPENVINO_MATCHER_PASS_RTTI("ov::frontend::ggml::pass::SqueezeMatmul") - SqueezeMatmul(); -}; - -} // namespace pass -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/rt_info/weightless_caching_attributes.hpp b/ggml/src/ggml-openvino/openvino/rt_info/weightless_caching_attributes.hpp deleted file mode 100644 index f051891c481e..000000000000 --- a/ggml/src/ggml-openvino/openvino/rt_info/weightless_caching_attributes.hpp +++ /dev/null @@ -1,41 +0,0 @@ -// Copyright (C) 2018-2026 Intel Corporation -// SPDX-License-Identifier: Apache-2.0 -// - -#pragma once - -#include -#include -#include - -namespace ov { - -/** - * @brief Holds weightless caching attributes of a single constant. - * - * WeightlessCacheAttribute class represents runtime info attribute that holds - * the values of original size of the constant in bytes and the binary offset of the - * constant's data in the weights file used by the weightless caching mechanism. It's - * not copyable in case the data was changed (the original node was replaced by a new - * one produced during the tranformation pipeline) - in that case weightless caching - * can't be used for that constant. - */ -class OPENVINO_API WeightlessCacheAttribute : public RuntimeAttribute { -public: - OPENVINO_RTTI("WeightlessCacheAttribute", "0", RuntimeAttribute) - - WeightlessCacheAttribute() = delete; - - WeightlessCacheAttribute(size_t original_size, size_t bin_offset, ov::element::Type original_dtype) - : original_size(original_size), - bin_offset(bin_offset), - original_dtype(original_dtype) {} - - bool is_copyable() const override; - - size_t original_size; - size_t bin_offset; - ov::element::Type original_dtype; -}; - -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/translate_session.cpp b/ggml/src/ggml-openvino/openvino/translate_session.cpp deleted file mode 100644 index 0f68a1f50623..000000000000 --- a/ggml/src/ggml-openvino/openvino/translate_session.cpp +++ /dev/null @@ -1,317 +0,0 @@ -#include "translate_session.h" - -#include "ggml-openvino/openvino/node_context.h" -#include "ggml-openvino/openvino/utils.h" -#include "input_model.h" -#include "pass/mark_decompression_convert_constant_folding.h" -#include "pass/squeeze_matmul.h" -#include "rt_info/weightless_caching_attributes.hpp" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { - -using namespace ov::op; - -namespace { - -ov::pass::MakeStateful::ParamResPairs get_kv_param_res_pairs( - const std::shared_ptr & model, - const std::map & kv_param_res_names) { - ov::pass::MakeStateful::ParamResPairs pairs; - const auto & params = model->get_parameters(); - const auto & results = model->get_results(); - - for (const auto & param_res : kv_param_res_names) { - const auto & param_name = param_res.first; - const auto & res_name = param_res.second; - - auto param_it = std::find_if(params.begin(), params.end(), [&](const std::shared_ptr & node) { - return node->get_friendly_name() == param_name; - }); - - OPENVINO_ASSERT(param_it != params.end(), "The tensor name ", param_name, - " is not associated with any of " - "Parameters in the network."); - - auto res_it = std::find_if(results.begin(), results.end(), [&](const std::shared_ptr & node) { - return node->get_friendly_name() == res_name; - }); - - OPENVINO_ASSERT(res_it != results.end(), "The tensor name ", res_name, - " is not associated with any of " - "Results in the network."); - - std::shared_ptr param = *param_it; - std::shared_ptr res = *res_it; - pairs.emplace_back(param, res); - } - return pairs; -} - -void add_sliced_mask(TensorMap & tensor_map, GgmlDecoder & ggml_model_decoder) { - - auto create_sliced_mask = [&](const std::string & mask_name, const std::string & sliced_name, bool is_static) { - if ((tensor_map.find(mask_name) != tensor_map.end()) && - (tensor_map.find("token_len_per_seq") != tensor_map.end())) { - auto token_len_per_seq = tensor_map.at("token_len_per_seq").get_node_shared_ptr(); - auto mask = tensor_map.at(mask_name).get_node_shared_ptr(); - std::shared_ptr mask_sliced; - if (is_static) { - mask_sliced = mask; - } else if (ggml_model_decoder.is_stateful()) { - auto zero_2d = ov::op::v0::Constant::create(ov::element::i64, {2}, {0,0}); - auto one_2d = ov::op::v0::Constant::create(ov::element::i64, {2}, {1,1}); - auto zero_1d = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); - auto three_1d = ov::op::v0::Constant::create(ov::element::i64, {1}, {3}); - auto neg_one_1d = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}); - auto axes = ov::op::v0::Constant::create(ov::element::i64, {2}, {-2,-1}); - auto inp_pos = tensor_map.at("inp_pos").get_node_shared_ptr(); - auto gather_inp_pos = std::make_shared(inp_pos, neg_one_1d, three_1d); - auto reshaped_inp_pos = std::make_shared(gather_inp_pos, ov::op::v0::Constant::create(ov::element::i64, {1}, {1}), false); - auto inp_pos_incremented = std::make_shared(reshaped_inp_pos, ov::op::v0::Constant::create(ov::element::i32, ov::Shape{1}, {1})); - auto stop = std::make_shared(ov::OutputVector{token_len_per_seq, std::make_shared(inp_pos_incremented, token_len_per_seq)}, 0); - mask_sliced = - std::make_shared(mask, zero_2d, stop, one_2d, axes); - mask_sliced = std::make_shared(mask_sliced, ov::element::f16); - mask_sliced->set_friendly_name(sliced_name); - } else { - auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); - auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); - auto two = ov::op::v0::Constant::create(ov::element::i64, {1}, {2}); - mask_sliced = std::make_shared(mask, zero, token_len_per_seq, one, two); - mask_sliced = std::make_shared(mask_sliced, ov::element::f16); - mask_sliced->set_friendly_name(sliced_name); - } - tensor_map.insert({sliced_name, mask_sliced->output(0)}); - } - }; - - create_sliced_mask("self_kq_mask", "KQ_mask_sliced", ggml_model_decoder.is_static()); - create_sliced_mask("self_kq_mask_swa", "KQ_mask_swa_sliced", ggml_model_decoder.is_static()); -} - -void add_rope_sin_cos(TensorMap & tensor_map, GgmlDecoder & ggml_model_decoder) { - int32_t * rope_params = ggml_model_decoder.get_rope_params(); - if (tensor_map.find("inp_pos") == tensor_map.end() || rope_params == nullptr) { - return; - } - auto inp_pos = tensor_map.at("inp_pos").get_node_shared_ptr(); - std::shared_ptr rope_freqs_weight; - if (tensor_map.find("rope_freqs.weight") != tensor_map.end()) { - rope_freqs_weight = tensor_map.at("rope_freqs.weight").get_node_shared_ptr(); - } - - auto sin_cos = make_sin_cos(rope_params, inp_pos, rope_freqs_weight); - auto sin_theta = sin_cos.first; - auto cos_theta = sin_cos.second; - - cos_theta.get_node_shared_ptr()->set_friendly_name("rope_cos"); - sin_theta.get_node_shared_ptr()->set_friendly_name("rope_sin"); - tensor_map.insert({"rope_cos", cos_theta}); - tensor_map.insert({"rope_sin", sin_theta}); -} - -// Create common patterns -void preprocess(TensorMap & tensor_map, GgmlDecoder & ggml_model_decoder) { - add_sliced_mask(tensor_map, ggml_model_decoder); - add_rope_sin_cos(tensor_map, ggml_model_decoder); -} - -} // namespace - -TranslateSession::TranslateSession(const frontend::InputModel::Ptr & input_model, - const std::unordered_map & translator_map, - bool naive) : - m_input_model(input_model), - m_translator_map(translator_map), - m_ov_model(nullptr), - m_naive(naive) {} - -std::shared_ptr TranslateSession::get_converted_model() { - if (m_ov_model) { - return m_ov_model; - } - m_ov_model = translate_graph(m_input_model); - return m_ov_model; -} - -std::shared_ptr TranslateSession::translate_graph(const frontend::InputModel::Ptr & input_model) { - ov::ParameterVector params; - ov::ResultVector results; - auto tensor_map = std::make_shared(); - std::shared_ptr resulting_model; - - const auto & ggml_model = std::dynamic_pointer_cast(input_model); - std::shared_ptr ggml_model_decoder = ggml_model->get_model_decoder(); - - for (const auto & it : ggml_model_decoder->get_model_inputs()) { - params.push_back(std::dynamic_pointer_cast(it.second)); - (*tensor_map)[it.first] = it.second; - } - - for (const auto & it : ggml_model_decoder->get_model_extra_inputs()) { - if (std::dynamic_pointer_cast(it.second)) { - params.push_back(std::dynamic_pointer_cast(it.second)); - } - (*tensor_map)[it.first] = it.second; - } - - for (const auto & it : ggml_model_decoder->get_model_weights()) { - (*tensor_map)[it.first] = it.second; - } - - auto node_visitor = [&](std::shared_ptr decoder, int node_idx) { - auto operation_type = decoder->get_op_type(node_idx); - if (operation_type == "GGML_OP_NONE") { - return; - } - - ov::OutputVector converted_outputs; - auto it = m_translator_map.find(operation_type); - FRONT_END_OP_CONVERSION_CHECK(it != m_translator_map.end(), "Translation for operation type ", operation_type, - " is not implemented."); - NodeContext node_context(decoder, tensor_map, node_idx, this); - converted_outputs = it->second(node_context); - - const auto & node_output_names = decoder->get_output_names(node_idx); - FRONT_END_OP_CONVERSION_CHECK(node_output_names.size() == converted_outputs.size(), "Number of ", - operation_type, " outputs greater than number of converted outputs, which are ", - node_output_names.size(), " and ", converted_outputs.size(), " respectively."); - - for (size_t i = 0; i < node_output_names.size(); ++i) { - auto output_name = node_output_names[i]; - if (i < converted_outputs.size() && converted_outputs[i].get_node_shared_ptr() != nullptr) { - (*tensor_map)[output_name] = converted_outputs[i]; - } - } - }; - - if (!m_naive) { - preprocess(*tensor_map, *ggml_model_decoder); - } - ggml_model_decoder->visit_subgraph(node_visitor); - - for (const auto & name : ggml_model_decoder->get_model_output_names()) { - FRONT_END_GENERAL_CHECK(tensor_map->find(name) != tensor_map->end(), - "Output name not found in tensor map: ", name); - auto result = std::make_shared(tensor_map->at(name)); - result->set_friendly_name(name); - results.push_back(result); - } - - ov::ParameterVector used_params; - for (const auto & param : params) { - if (!param->output(0).get_target_inputs().empty()) { - used_params.push_back(param); - } - } - // if (auto diff = params.size() - used_params.size()) { - // GGML_LOG_INFO("%zu parameters are not used in the model.", diff); - // } - resulting_model = std::make_shared(results, used_params); - - apply_transformations(resulting_model); - - // Set WeightlessCacheAttribute on large constants to avoid unnecessary memory copies - // in the NPUW plugin. Without this attribute, NPUW's LazyTensor constructor - // (lazy_tensor.cpp, op::Const::Const) will memcpy every constant "in case export - // occurs", doubling memory usage per compile_model call. - // - // The bin_offset field serves as a unique key (not a real file offset) — this is - // the same convention the GPU plugin uses for non-IR models (see - // Plugin::set_weightless_cache_attributes in intel_gpu/src/plugin/plugin.cpp). - // Each constant must have a distinct bin_offset, otherwise GPU's weightless cache - // import will map multiple constants to the same data. - // - // Small constants (< 16 elements) are excluded since they may be introduced by - // optimization patterns and the overhead is negligible. - size_t offset = 0; - for (auto & node : resulting_model->get_ordered_ops()) { - if (auto cnst = ov::as_type_ptr(node); - cnst && cnst->get_byte_size() / cnst->get_element_type().size() >= 16) { - auto & rt_info = cnst->get_rt_info(); - if (rt_info.find(ov::WeightlessCacheAttribute::get_type_info_static()) == rt_info.end()) { - rt_info[ov::WeightlessCacheAttribute::get_type_info_static()] = - ov::WeightlessCacheAttribute(cnst->get_byte_size(), offset++, cnst->get_element_type()); - } - } - } - return resulting_model; -} - -std::shared_ptr TranslateSession::apply_transformations(std::shared_ptr model) { - auto ggml_model_decoder = std::dynamic_pointer_cast(m_input_model)->get_model_decoder(); - { - ov::pass::Manager manager; - manager.set_per_pass_validation(true); - manager.register_pass(); - - if (ggml_model_decoder->is_stateful()) { - const auto kv_param_res_names = ggml_model_decoder->get_kv_param_res_names(); - const auto kv_param_res_pairs = get_kv_param_res_pairs(model, kv_param_res_names); - manager.register_pass(kv_param_res_pairs); - } - - if (ggml_model_decoder->is_static()) { - manager.register_pass(); - } - manager.run_passes(model); - if (ggml_model_decoder->is_stateful()) { - auto output_names = ggml_model_decoder->get_model_output_names(); - std::map model_output_indexes; - for (size_t i=0; iget_output_size(); i++) { - auto output_friendly_name = model->output(i).get_node_shared_ptr()->get_friendly_name(); - auto output_id = model_output_indexes[output_friendly_name]; - auto model_output_shape = model->output(i).get_partial_shape(); - auto decoder_output_shape = ggml_model_decoder->get_output_shape(output_id); - if (model_output_shape.rank().is_static() && decoder_output_shape.rank().is_static() - && model_output_shape.rank().get_length() + 1 == decoder_output_shape.rank().get_length() - && decoder_output_shape[0].is_static() && decoder_output_shape[0].get_length() == 1) { - ppp.output(i).postprocess().custom([](const ov::Output& node) { - auto axes = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{1}, {0}); - return std::make_shared(node, axes); - }); - } - } - model = ppp.build(); - } - } - return model; -} - -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/translate_session.h b/ggml/src/ggml-openvino/openvino/translate_session.h deleted file mode 100644 index 56a14ae7c07d..000000000000 --- a/ggml/src/ggml-openvino/openvino/translate_session.h +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -#include "input_model.h" -#include "node_context.h" - -namespace ov { -namespace frontend { -namespace ggml { - -class TranslateSession { -public: - TranslateSession(const frontend::InputModel::Ptr& input_model, - const std::unordered_map& translator_map, bool naive = false); - - std::shared_ptr get_converted_model(); - std::shared_ptr translate_graph(const frontend::InputModel::Ptr& input_model); - -private: - std::shared_ptr apply_transformations(std::shared_ptr model); - const frontend::InputModel::Ptr m_input_model; - const std::unordered_map& m_translator_map; - std::shared_ptr m_ov_model; - bool m_naive; -}; - -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/utils.cpp b/ggml/src/ggml-openvino/openvino/utils.cpp deleted file mode 100644 index 0baaf88e17a7..000000000000 --- a/ggml/src/ggml-openvino/openvino/utils.cpp +++ /dev/null @@ -1,257 +0,0 @@ -#include "utils.h" - -#include "ggml-impl.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace ov { -namespace frontend { -namespace ggml { - -std::string getCurrentTime() { - std::time_t now = std::time(nullptr); - char buf[100]; - std::strftime(buf, sizeof(buf), "%Y-%m-%d %H:%M:%S", std::localtime(&now)); - return buf; -} - -void num_inputs_check(const NodeContext & context, size_t min_inputs, size_t max_inputs) { - auto input_size = context.get_input_size(); - FRONT_END_OP_CONVERSION_CHECK(input_size >= min_inputs, "Got less inputs than expected"); - FRONT_END_OP_CONVERSION_CHECK(input_size <= max_inputs, "Got more inputs than expected"); -} - -int non_cont_dim(std::vector ne, std::vector nb) { - int dim = nb.size() - 1; - size_t bytes = nb[dim]; - for (int i = dim; i > 0; i--) { - bytes *= ne[i]; - if (bytes != nb[i - 1]) { - return i; - } - } - return 0; -} - -std::shared_ptr get_dimensions(const std::shared_ptr & shape, - const std::vector & dims) { - using namespace ov::op; - const auto zero = v0::Constant::create(ov::element::i32, ov::Shape{}, {0}); - const auto dims_const = v0::Constant::create(ov::element::i32, ov::Shape{dims.size()}, dims); - return std::make_shared(shape, dims_const, zero); -} - -std::shared_ptr get_dimensions(const std::shared_ptr & node, const std::vector & dims) { - return get_dimensions(std::make_shared(node), dims); -} - -OutputVector rename_outputs_with_suffix(const OutputVector & outputs, const std::string & suffix) { - for (const auto & output : outputs) { - auto node = output.get_node_shared_ptr(); - std::string name = node->get_friendly_name(); - name += "_"; - name += suffix; - node->set_friendly_name(name); - // std::cout << name << " " << output.get_partial_shape() << std::endl; - } - return outputs; -} - -namespace { -ov::Output rope_yarn_ramp_mix(int n_dims, const float corr_dims[2], float ext_factor) { - int half_n_dims = n_dims / 2; - std::vector dim_ids_vec(half_n_dims); - std::iota(dim_ids_vec.begin(), dim_ids_vec.end(), 0); - auto dim_ids = ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1, (size_t) half_n_dims}, dim_ids_vec); - auto corr_low = ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1, 1}, {corr_dims[0]}); - auto corr_high = ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1, 1}, {corr_dims[1]}); - auto denom = std::make_shared( - std::make_shared(corr_high, corr_low), - ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1, 1}, {0.001f})); - auto ramp_y = - std::make_shared(std::make_shared(dim_ids, corr_low), denom); - auto ramp_clamped = std::make_shared(ramp_y, 0.0f, 1.0f); - // rope_yarn_ramp returns (1 - clamp(y)), so invert before scaling - auto one = ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1, 1}, {1.0f}); - auto ramp_inverted = std::make_shared(one, ramp_clamped); - auto ext_factor_node = ov::op::v0::Constant::create(ov::element::f32, Shape{}, {ext_factor}); - auto ramp_mix = std::make_shared(ramp_inverted, ext_factor_node); - return ramp_mix; -} - -float ggml_rope_yarn_corr_dim(int n_dims, int n_ctx_orig, float n_rot, float base) { -#ifndef M_PI -# define M_PI 3.14159265358979323846 -#endif - return n_dims * logf(n_ctx_orig / (n_rot * 2 * (float) M_PI)) / (2 * logf(base)); -} - -void ggml_rope_yarn_corr_dims(int n_dims, - int n_ctx_orig, - float freq_base, - float beta_fast, - float beta_slow, - float dims[2]) { - float start = floorf(ggml_rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_fast, freq_base)); - float end = ceilf(ggml_rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_slow, freq_base)); - dims[0] = std::max(0.0f, start); - dims[1] = std::min(static_cast(n_dims - 1), end); -} -} // namespace - -std::pair, ov::Output> make_sin_cos(int32_t * rope_params, - std::shared_ptr inp_pos, - std::shared_ptr rope_freqs_weight, - bool imrope, - bool stateful) { - if (stateful) { - inp_pos = std::make_shared(inp_pos, ov::op::v0::Constant::create(ov::element::i64, {1}, {0})); - inp_pos = std::make_shared(inp_pos, ov::element::f32); - auto pos_perm = - std::make_shared(ov::element::i64, ov::Shape{3}, std::vector{2, 1, 0}); - inp_pos = std::make_shared(inp_pos, pos_perm); - } else if (imrope) { - inp_pos = std::make_shared(inp_pos, ov::element::f32); - auto pos_shape = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{5}, {0, 0, 0, 4, -1}); - inp_pos = std::make_shared(inp_pos, pos_shape, true); - auto pos_transpose_shape = - std::make_shared(ov::element::i64, ov::Shape{5}, std::vector{0, 1, 2, 4, 3}); - inp_pos = std::make_shared(inp_pos, pos_transpose_shape); - } else { - inp_pos = std::make_shared(inp_pos, ov::element::f32); - auto pos_perm = - std::make_shared(ov::element::i64, ov::Shape{4}, std::vector{0, 3, 1, 2}); - inp_pos = std::make_shared(inp_pos, pos_perm); - } - - float freq_base; - float freq_scale; - float ext_factor; - float attn_factor; - float beta_fast; - float beta_slow; - const int n_dims = rope_params[1]; - const size_t n_dims_half = n_dims >> 1; - const int n_ctx_orig = rope_params[4]; - memcpy(&freq_base, rope_params + 5, sizeof(float)); - memcpy(&freq_scale, rope_params + 6, sizeof(float)); - memcpy(&ext_factor, rope_params + 7, sizeof(float)); - memcpy(&attn_factor, rope_params + 8, sizeof(float)); - memcpy(&beta_fast, rope_params + 9, sizeof(float)); - memcpy(&beta_slow, rope_params + 10, sizeof(float)); - - const float theta_scale = powf(freq_base, -2.0f / n_dims); - - std::vector factor(n_dims_half); - - Output freq_factors; - - Output theta; - float mscale = attn_factor; - if (imrope) { - std::vector gather_indices(n_dims_half); - for (size_t j = 0; j < n_dims_half; j++) { - gather_indices[j] = j % 3; - factor[j] = std::pow(theta_scale, j); - } - auto gather_indices_const = - std::make_shared(ov::element::i64, ov::Shape{n_dims_half}, gather_indices); - auto gather_axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {4}); - inp_pos = std::make_shared(inp_pos, gather_indices_const, gather_axis); - auto factor_const = std::make_shared(ov::element::f32, ov::Shape{n_dims_half}, factor); - theta = std::make_shared(inp_pos, factor_const); - } else { - float corr_dims[2]; - ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); - factor[0] = 1.0f; - for (size_t i = 1; i < factor.size(); i++) { - factor[i] = theta_scale * factor[i - 1]; - } - if (stateful) { - freq_factors = - std::make_shared(ov::element::f32, ov::Shape{1, 1, factor.size()}, factor); - } else { - freq_factors = - std::make_shared(ov::element::f32, ov::Shape{1, 1, 1, factor.size()}, factor); - } - if (rope_freqs_weight) { - freq_factors = std::make_shared(freq_factors, rope_freqs_weight); - } - - auto theta_extrap = std::make_shared(freq_factors, inp_pos); - auto theta_interp = std::make_shared( - theta_extrap, ov::op::v0::Constant::create(ov::element::f32, {1}, {freq_scale})); - - if (ext_factor == 0.0f) { - theta = theta_interp; - } else { - auto ramp_mix = rope_yarn_ramp_mix(n_dims, corr_dims, ext_factor); - Output one; - if (stateful) { - one = ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1}, {1.0f}); - } else { - one = ov::op::v0::Constant::create(ov::element::f32, Shape{1, 1, 1, 1}, {1.0f}); - } - auto one_minus_ramp = std::make_shared(one, ramp_mix); - - theta = std::make_shared(std::make_shared(theta_interp, one_minus_ramp), - std::make_shared(theta_extrap, ramp_mix)); - mscale *= (1.0f + 0.1f * std::log(1.0f / freq_scale)); - } - } - - Output cos_theta = std::make_shared(theta); - Output sin_theta = std::make_shared(theta); - - if (!imrope) { - auto mscale_node = ov::op::v0::Constant::create(ov::element::f32, Shape{}, {mscale}); - - cos_theta = std::make_shared(cos_theta, mscale_node); - sin_theta = std::make_shared(sin_theta, mscale_node); - } - - return std::make_pair(sin_theta, cos_theta); -} - -ov::Output process_view_input(const NodeContext & context, int input_index, int slice_len) { - // Only works for VIEW operations that slice at the lowest dimension - // If the VIEW also reshape the result, `slice_len` should be provided - auto input = context.get_input(input_index); - auto * op_params = (size_t *) context.get_input_op_params(input_index); - auto src1_stride = context.get_input_stride(input_index); - - int64_t split_addr = op_params[0] / src1_stride[3]; - if (slice_len == 0) { - slice_len = context.get_input_shape(input_index)[3].get_length(); - } - int64_t slice_end = split_addr + slice_len; - - auto begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {split_addr}); - auto end = ov::op::v0::Constant::create(ov::element::i64, {1}, {slice_end}); - auto stride = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); - auto axes = ov::op::v0::Constant::create(ov::element::i64, {1}, {context.is_stateful() ? 2 : 3}); - auto sliced = std::make_shared(input, begin, end, stride, axes); - return sliced; -} - -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/utils.h b/ggml/src/ggml-openvino/openvino/utils.h deleted file mode 100644 index 767dd4c53ea5..000000000000 --- a/ggml/src/ggml-openvino/openvino/utils.h +++ /dev/null @@ -1,86 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include - -#include "node_context.h" - -namespace ov { -namespace frontend { -namespace ggml { - -std::string getCurrentTime(); - -void dump_ov_model(std::shared_ptr model); - -void num_inputs_check(const NodeContext& context, size_t min_inputs, size_t max_inputs); - -int non_cont_dim(std::vector ne, std::vector nb); - -template -std::vector argsort_descend(const std::vector& v) { - std::vector idx(v.size()); - std::iota(idx.begin(), idx.end(), 0); - std::sort(idx.begin(), idx.end(), [&v](int i1, int i2) { - return v[i1] > v[i2]; - }); - return idx; -} - -template -std::vector sorted_descend(std::vector v) { - std::sort(v.begin(), v.end(), [](T a, T b) { - return a > b; - }); - return v; -} - -template -bool is_permuted(const std::vector& strides) { - for (size_t i = 0; i < strides.size() - 1; ++i) { - if (strides[i] < strides[i + 1]) { - return true; - } - } - return false; -} - -template -std::vector permute(const std::vector& x, const std::vector& perm) { - std::vector result; - result.reserve(perm.size()); - for (int i : perm) { - result.push_back(x[i]); - } - return result; -} - -std::shared_ptr get_dimensions(const std::shared_ptr& shape, - const std::vector& dims); -std::shared_ptr get_dimensions(const std::shared_ptr& node, const std::vector& dims); - -OutputVector rename_outputs_with_suffix(const OutputVector& outputs, const std::string& suffix); - -std::pair, ov::Output> make_sin_cos(int32_t* rope_params, - std::shared_ptr inp_pos, - std::shared_ptr rope_freqs_weight = nullptr, - bool imrope = false, - bool stateful = false); - -ov::Output process_view_input(const NodeContext& context, int input_index, int slice_len = 0); - -namespace op { -template -OutputVector translate_1to1_match_2_inputs(const NodeContext& context) { - num_inputs_check(context, 2, 2); - auto res = std::make_shared(context.get_input(0), context.get_input(1)); - return rename_outputs_with_suffix({res}, context.get_name()); -} -} // namespace op - -} // namespace ggml -} // namespace frontend -} // namespace ov diff --git a/ggml/src/ggml-openvino/utils.cpp b/ggml/src/ggml-openvino/utils.cpp deleted file mode 100644 index 998ef7c9eb4f..000000000000 --- a/ggml/src/ggml-openvino/utils.cpp +++ /dev/null @@ -1,880 +0,0 @@ -#include "utils.h" - -#include "ggml-impl.h" -#include "ggml-openvino-extra.h" -#include "ggml-openvino/ggml-decoder.h" -#include "ggml.h" -#include "openvino/frontend.h" -#include "openvino/input_model.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -// Suppress deprecation warning for ov::Tensor::data() -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wdeprecated-declarations" - -enum ggml_status ov_graph_compute(ggml_cgraph * cgraph, ggml_backend_t backend) { - ggml_backend_openvino_context * ctx = (ggml_backend_openvino_context *) backend->context; - try { - if (getenv("GGML_OPENVINO_DUMP_CGRAPH")) { - std::string filename = "cgraph_ov.txt"; - GgmlOvDecoder::dump_cgraph(cgraph, filename); - } - - const auto is_static = ggml_openvino_is_npu(); - - GGML_ASSERT(ctx->runtime_context != nullptr); - std::shared_ptr r_ctx = std::static_pointer_cast(ctx->runtime_context); - - return is_static ? ov_graph_compute_static(cgraph, r_ctx) : ov_graph_compute_dynamic(cgraph, r_ctx); - } catch (const ov::Exception & e) { - GGML_LOG_ERROR("GGML OpenVINO backend ov::Exception: %s\n", e.what()); - return GGML_STATUS_FAILED; - } catch (const std::exception & e) { - GGML_LOG_ERROR("GGML OpenVINO backend std::exception: %s\n", e.what()); - return GGML_STATUS_FAILED; - } catch (...) { - GGML_LOG_ERROR("GGML OpenVINO backend unknown exception\n"); - return GGML_STATUS_FAILED; - } -} - -ov::Tensor create_ov_output_tensor(std::shared_ptr ggml_decoder, - std::shared_ptr infer_request, - int output_index, - const ggml_tensor * ggml_tensor) { - auto output_type = ggml_decoder->get_ov_type(ggml_tensor); - ov::Shape output_shape; - if (ggml_decoder->is_static()) { - output_shape = infer_request->get_output_tensor(output_index).get_shape(); - } else { - output_shape = ggml_decoder->get_shape(ggml_tensor); - } - - ov::Tensor output_tensor(output_type, output_shape, ggml_tensor->data); - return output_tensor; -} - -enum ggml_status ov_graph_compute_dynamic(ggml_cgraph * cgraph, std::shared_ptr r_ctx) { - auto & core = ov_singleton_core(); - const auto & config = ggml_openvino_get_compile_config(); - const auto & device = r_ctx->device; - const auto & stateful = r_ctx->stateful; - static auto is_static = false; - - if (is_naive(cgraph)) { - return naive_compute(cgraph, core, device, config); - } - - auto start_time = ggml_time_us(); - - std::shared_ptr ggml_decoder; - std::shared_ptr infer_request; - ModelParams m_params; - ComputeParams c_params; - std::tie(m_params, c_params) = GgmlOvDecoder::compute_llm_params(cgraph, is_static); - - graph_key key(cgraph); - bool cache_hit; - - int64_t decoder_end_time; - int64_t conversion_end_time; - int64_t compile_end_time; - int64_t infer_end_time; - - { - std::shared_ptr entry; - ModelParams old_m_params; - - { - std::lock_guard map_lock(r_ctx->ctx_mutex); - auto it = r_ctx->decoder_cache.find(key); - cache_hit = it != r_ctx->decoder_cache.end(); - if (cache_hit) { - entry = it->second; - } else { - auto mutex = std::make_shared(); - entry = std::make_shared(mutex); - r_ctx->decoder_cache[key] = entry; - } - } - - std::lock_guard lock(*(entry->mutex)); - - if (cache_hit) { - ggml_decoder = entry->ptr; - old_m_params = ggml_decoder->get_model_params(); - cache_hit = old_m_params.can_reuse_dynamically(m_params); - } - - if (cache_hit) { - std::map> model_weights; - ggml_decoder->set_compute_params(c_params); - ggml_decoder->set_model_params(m_params); - if (old_m_params.kv_buffer_changed(m_params)) { - ggml_decoder->update_io(cgraph); - } - ggml_decoder->add_extra_inputs(); - { - std::lock_guard map_lock(r_ctx->ctx_mutex); - infer_request = r_ctx->infer_request_cache.at(key); - } - - if (stateful) { - const auto * inp_pos = get_inp_pos_tensor(cgraph); - int32_t * pos_data = (int32_t *) inp_pos->data; - auto pos_shape = ggml_decoder->get_shape(inp_pos); - if (pos_data[0] == 0) { - infer_request->reset_state(); - r_ctx->stateful_kv_size = pos_shape[3]; - } else if (r_ctx->stateful_kv_size == static_cast(pos_data[0])) { - r_ctx->stateful_kv_size += pos_shape[3]; - } else { - auto states = infer_request->query_state(); - for (auto state : states) { - auto state_tensor = state.get_state(); - auto state_tensor_shape = state_tensor.get_shape(); - if (static_cast(pos_data[0]) > r_ctx->stateful_kv_size) { - std::string state_name; - try { - state_name = r_ctx->kv_state_input_name_map.at(state.get_name()); - } catch (...) { - GGML_LOG_ERROR("GGML OpenVINO backend stateful inference failed: no input found for the state\n"); - return GGML_STATUS_FAILED; - } - auto kv_tensor = get_ov_input_tensor(ggml_decoder, state_name); - kv_tensor.set_shape({state_tensor_shape[0], kv_tensor.get_shape()[2], - state_tensor_shape[2], state_tensor_shape[3]}); - state_tensor = kv_tensor; - state_tensor_shape = state_tensor.get_shape(); - } - ov::Coordinate begin = {0, 0, 0, 0}; - ov::Coordinate end = {state_tensor_shape[0], static_cast(pos_data[0]), - state_tensor_shape[2], state_tensor_shape[3]}; - ov::Tensor new_state_tensor(state_tensor, begin, end); - state.set_state(new_state_tensor); - } - r_ctx->stateful_kv_size = pos_data[0] + 1; - } - } - - decoder_end_time = ggml_time_us(); - conversion_end_time = decoder_end_time; - compile_end_time = decoder_end_time; - } else { - { - std::lock_guard map_lock(r_ctx->ctx_mutex); - r_ctx->infer_request_cache.erase(key); - } - - std::shared_ptr model; - auto model_weights = GgmlOvDecoder::create_weight_nodes(cgraph); - - ggml_decoder = std::make_shared(cgraph, m_params, c_params, model_weights, is_static, stateful); - decoder_end_time = ggml_time_us(); - - auto input_model = std::make_shared(ggml_decoder); - model = ov::frontend::ggml::FrontEnd::convert(input_model); - ggml_decoder->clear_model_weights(); - conversion_end_time = ggml_time_us(); - - if (getenv("GGML_OPENVINO_DUMP_IR")) { - char timestamped_filename[64]; - auto timestamp = (long long) ggml_time_us(); - snprintf(timestamped_filename, sizeof(timestamped_filename), "model_%lld.xml", timestamp); - ov::serialize(model, timestamped_filename); - } - - ov::CompiledModel compiled_model; - auto remote_context = ggml_openvino_get_remote_context(); - if (remote_context.has_value()) { - compiled_model = core.compile_model(model, remote_context.value(), config); - } else { - compiled_model = core.compile_model(model, device, config); - } - compile_end_time = ggml_time_us(); - infer_request = std::make_shared(compiled_model.create_infer_request()); - entry->ptr = ggml_decoder; - - std::vector ov_input_names; - std::vector ov_output_names; - for (const auto & ov_param : model->get_parameters()) { - ov_input_names.push_back(ov_param->get_friendly_name()); - } - for (const auto & ov_output : model->get_results()) { - ov_output_names.push_back(ov_output->get_friendly_name()); - } - - { - std::lock_guard map_lock(r_ctx->ctx_mutex); - r_ctx->infer_request_cache[key] = infer_request; - r_ctx->ov_input_names_cache[key] = std::move(ov_input_names); - r_ctx->ov_output_names_cache[key] = std::move(ov_output_names); - } - - if (stateful) { - const auto * inp_pos = get_inp_pos_tensor(cgraph); - auto pos_shape = ggml_decoder->get_shape(inp_pos); - r_ctx->stateful_kv_size = pos_shape[3]; - const auto kv_param_res_names = ggml_decoder->get_kv_param_res_names(); - for (const auto& pair : kv_param_res_names) { - r_ctx->kv_state_input_name_map[pair.first+pair.second] = pair.first; - } - } - } - - std::vector ov_input_names; - std::vector ov_output_names; - { - std::lock_guard map_lock(r_ctx->ctx_mutex); - ov_input_names = r_ctx->ov_input_names_cache[key]; - ov_output_names = r_ctx->ov_output_names_cache[key]; - } - - for (size_t i = 0; i < ov_input_names.size(); i++) { - auto param_name = ov_input_names[i]; - auto input_tensor = get_ov_input_tensor(ggml_decoder, param_name); - infer_request->set_input_tensor(i, input_tensor); - - if (getenv("GGML_OPENVINO_DEBUG_INPUT")) { - print_input_tensor_info(param_name, input_tensor); - } - } - - for (size_t i = 0; i < ov_output_names.size(); i++) { - auto * ggml_tensor = ggml_decoder->get_model_outputs().at(ov_output_names[i]); - auto output_tensor = create_ov_output_tensor(ggml_decoder, infer_request, i, ggml_tensor); - infer_request->set_output_tensor(i, output_tensor); - } - - infer_request->infer(); - infer_end_time = ggml_time_us(); - - if (getenv("GGML_OPENVINO_DEBUG_OUTPUT")) { - for (size_t i = 0; i < ov_output_names.size(); i++) { - const auto output_tensor = infer_request->get_output_tensor(i); - print_output_tensor_info(ov_output_names[i], output_tensor, output_tensor.data()); - } - } - - if (getenv("GGML_OPENVINO_PROFILING")) { - GGML_LOG_INFO("\nGGML OpenVINO Backend: \n"); - GGML_LOG_INFO(" - Graph decoder time: %ld ms \n", (decoder_end_time - start_time) / 1000); - if (!cache_hit) { - GGML_LOG_INFO(" - Graph conversion time: %ld ms \n", (conversion_end_time - decoder_end_time) / 1000); - GGML_LOG_INFO(" - Graph compile time: %ld ms \n", (compile_end_time - conversion_end_time) / 1000); - } - GGML_LOG_INFO(" - Graph inference time: %ld ms \n", (infer_end_time - compile_end_time) / 1000); - } - } - - return GGML_STATUS_SUCCESS; -} - -enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptr r_ctx) { - auto & core = ov_singleton_core(); - - auto get_prefill_chunk_size = [] { - const char * chunk_size_str = getenv("GGML_OPENVINO_PREFILL_CHUNK_SIZE"); - if (chunk_size_str && atoi(chunk_size_str) > 0) { - return atoi(chunk_size_str); - } - return 256; - }; - - static std::string device = "NPU"; - static auto is_static = true; - static auto stateful = false; - static auto prefill_chunk_size = get_prefill_chunk_size(); - const auto & config = ggml_openvino_get_compile_config(); - - if (is_naive(cgraph)) { - return naive_compute(cgraph, core, device, config); - } - - auto start_time = ggml_time_us(); - - std::shared_ptr ggml_decoder; - std::shared_ptr infer_request; - ModelParams m_params; - ComputeParams c_params; - std::tie(m_params, c_params) = GgmlOvDecoder::compute_llm_params(cgraph, is_static); - - const auto * inp_pos = get_inp_pos_tensor(cgraph); - const auto is_prefill = get_is_prefill(inp_pos); - graph_key key(cgraph); - bool cache_hit; - - int64_t decoder_end_time; - int64_t conversion_end_time; - int64_t compile_end_time; - int64_t infer_end_time; - - std::shared_ptr entry; - ModelParams old_m_params; - - { - std::lock_guard map_lock(r_ctx->ctx_mutex); - auto it = r_ctx->decoder_cache.find(key); - cache_hit = it != r_ctx->decoder_cache.end(); - if (cache_hit) { - entry = it->second; - } else { - auto mutex = std::make_shared(); - entry = std::make_shared(mutex); - r_ctx->decoder_cache[key] = entry; - } - } - - std::lock_guard lock(*(entry->mutex)); - - if (cache_hit) { - ggml_decoder = entry->ptr; - old_m_params = ggml_decoder->get_model_params(); - cache_hit = old_m_params.can_reuse_statically(m_params); - } - - if (cache_hit) { - std::map> model_weights; - ggml_decoder->m_is_prefill = is_prefill; - ggml_decoder->set_model_params(m_params); - ggml_decoder->set_compute_params(c_params); - if (old_m_params.kv_buffer_changed(m_params)) { - ggml_decoder->update_io(cgraph); - } - ggml_decoder->add_extra_inputs(); - { - std::lock_guard map_lock(r_ctx->ctx_mutex); - infer_request = - is_prefill ? r_ctx->infer_request_cache_prefill.at(key) : r_ctx->infer_request_cache.at(key); - } - - decoder_end_time = ggml_time_us(); - conversion_end_time = decoder_end_time; - compile_end_time = decoder_end_time; - } else { - { - std::lock_guard map_lock(r_ctx->ctx_mutex); - r_ctx->infer_request_cache.erase(key); - r_ctx->infer_request_cache_prefill.erase(key); - } - - std::shared_ptr model; - auto model_weights = GgmlOvDecoder::create_weight_nodes(cgraph); - - auto ggml_decoder_prefill = std::make_shared(cgraph, m_params, c_params, model_weights, - is_static, stateful, true, prefill_chunk_size); - auto ggml_decoder_decode = std::make_shared(cgraph, m_params, c_params, model_weights, is_static, - stateful, false, prefill_chunk_size); - decoder_end_time = ggml_time_us(); - - auto input_model_prefill = std::make_shared(ggml_decoder_prefill); - auto input_model_decode = std::make_shared(ggml_decoder_decode); - - auto model_prefill = ov::frontend::ggml::FrontEnd::convert(input_model_prefill); - ggml_decoder_prefill->clear_model_weights(); - auto model_decode = ov::frontend::ggml::FrontEnd::convert(input_model_decode); - ggml_decoder_decode->clear_model_weights(); - conversion_end_time = ggml_time_us(); - - if (getenv("GGML_OPENVINO_DUMP_IR")) { - char timestamped_filename[64]; - auto timestamp = (long long) ggml_time_us(); - snprintf(timestamped_filename, sizeof(timestamped_filename), "model_prefill_%lld.xml", timestamp); - ov::serialize(model_prefill, timestamped_filename); - snprintf(timestamped_filename, sizeof(timestamped_filename), "model_decode_%lld.xml", timestamp); - ov::serialize(model_decode, timestamped_filename); - } - - ov::CompiledModel compiled_model_prefill; - ov::CompiledModel compiled_model_decode; - auto remote_context = ggml_openvino_get_remote_context(); - if (remote_context.has_value()) { - compiled_model_prefill = core.compile_model(model_prefill, remote_context.value(), config); - compiled_model_decode = core.compile_model(model_decode, remote_context.value(), config); - } else { - compiled_model_prefill = core.compile_model(model_prefill, device, config); - compiled_model_decode = core.compile_model(model_decode, device, config); - } - - auto infer_request_prefill = std::make_shared(compiled_model_prefill.create_infer_request()); - auto infer_request_decode = std::make_shared(compiled_model_decode.create_infer_request()); - compile_end_time = ggml_time_us(); - - model = is_prefill ? model_prefill : model_decode; - ggml_decoder = is_prefill ? ggml_decoder_prefill : ggml_decoder_decode; - infer_request = is_prefill ? infer_request_prefill : infer_request_decode; - entry->ptr = ggml_decoder; - - std::vector ov_input_names; - std::vector ov_output_names; - for (const auto & ov_param : model->get_parameters()) { - ov_input_names.push_back(ov_param->get_friendly_name()); - } - for (const auto & ov_output : model->get_results()) { - ov_output_names.push_back(ov_output->get_friendly_name()); - } - - { - std::lock_guard map_lock(r_ctx->ctx_mutex); - r_ctx->infer_request_cache_prefill[key] = infer_request_prefill; - r_ctx->infer_request_cache[key] = infer_request_decode; - r_ctx->ov_input_names_cache[key] = std::move(ov_input_names); - r_ctx->ov_output_names_cache[key] = std::move(ov_output_names); - } - } - - std::vector ov_input_names_local; - std::vector ov_output_names_local; - { - std::lock_guard map_lock(r_ctx->ctx_mutex); - ov_input_names_local = r_ctx->ov_input_names_cache[key]; - ov_output_names_local = r_ctx->ov_output_names_cache[key]; - } - - if (is_prefill) { - auto inp_len = inp_pos->ne[0]; - for (int chunk_index = 0; chunk_index * prefill_chunk_size < inp_len; chunk_index++) { - for (size_t i = 0; i < ov_input_names_local.size(); i++) { - auto param_name = ov_input_names_local[i]; - auto input_tensor = get_ov_input_tensor_static_prefill(ggml_decoder, param_name, chunk_index); - infer_request->set_input_tensor(i, input_tensor); - - if (getenv("GGML_OPENVINO_DEBUG_INPUT")) { - const auto input_tensor = infer_request->get_input_tensor(i); - print_input_tensor_info(param_name, input_tensor); - } - } - - for (size_t i = 0; i < ov_output_names_local.size(); i++) { - auto * ggml_tensor = ggml_decoder->get_model_outputs().at(ov_output_names_local[i]); - auto output_tensor = create_ov_output_tensor(ggml_decoder, infer_request, i, ggml_tensor); - infer_request->set_output_tensor(i, output_tensor); - } - - infer_request->infer(); - - if (getenv("GGML_OPENVINO_DEBUG_OUTPUT")) { - for (size_t i = 0; i < ov_output_names_local.size(); i++) { - const auto output_tensor = infer_request->get_output_tensor(i); - print_output_tensor_info(ov_output_names_local[i], output_tensor, output_tensor.data()); - } - } - } - infer_end_time = ggml_time_us(); - } else { - for (size_t i = 0; i < ov_input_names_local.size(); i++) { - auto param_name = ov_input_names_local[i]; - auto input_tensor = get_ov_input_tensor_static_decode(ggml_decoder, param_name); - infer_request->set_input_tensor(i, input_tensor); - - if (getenv("GGML_OPENVINO_DEBUG_INPUT")) { - const auto input_tensor = infer_request->get_input_tensor(i); - print_input_tensor_info(param_name, input_tensor); - } - } - - for (size_t i = 0; i < ov_output_names_local.size(); i++) { - auto * ggml_tensor = ggml_decoder->get_model_outputs().at(ov_output_names_local[i]); - auto output_tensor = create_ov_output_tensor(ggml_decoder, infer_request, i, ggml_tensor); - infer_request->set_output_tensor(i, output_tensor); - } - - infer_request->infer(); - infer_end_time = ggml_time_us(); - - if (getenv("GGML_OPENVINO_DEBUG_OUTPUT")) { - for (size_t i = 0; i < ov_output_names_local.size(); i++) { - const auto output_tensor = infer_request->get_output_tensor(i); - print_output_tensor_info(ov_output_names_local[i], output_tensor, output_tensor.data()); - } - } - } - - if (getenv("GGML_OPENVINO_PROFILING")) { - GGML_LOG_INFO("\nGGML OpenVINO Backend: \n"); - GGML_LOG_INFO(" - Graph decoder time: %ld ms \n", (decoder_end_time - start_time) / 1000); - if (!cache_hit) { - GGML_LOG_INFO(" - Graph conversion time: %ld ms \n", (conversion_end_time - decoder_end_time) / 1000); - GGML_LOG_INFO(" - Graph compile time: %ld ms \n", (compile_end_time - conversion_end_time) / 1000); - } - GGML_LOG_INFO(" - Graph inference time: %ld ms \n", (infer_end_time - compile_end_time) / 1000); - } - - return GGML_STATUS_SUCCESS; -} - -bool is_naive(ggml_cgraph * cgraph) { - constexpr int naive_graph_size_threshold = 20; - int count = 0; - for (int i = 0; i < cgraph->n_nodes; i++) { - if (cgraph->nodes[i]->op != GGML_OP_NONE) { - count++; - } - } - return count < naive_graph_size_threshold; -} - -enum ggml_status naive_compute(ggml_cgraph * cgraph, - ov::Core & core, - const std::string & device, - const ov::AnyMap & config) { - if (cgraph->n_nodes == 1 && (cgraph->nodes[0]->op == GGML_OP_NONE || cgraph->nodes[0]->op == GGML_OP_VIEW)) { - return GGML_STATUS_SUCCESS; - } - - bool naive = true; - auto model_weights = GgmlOvDecoder::create_weight_nodes(cgraph, naive); - auto decoder = std::make_shared(cgraph, model_weights); - auto input_model = std::make_shared(decoder); - auto model = ov::frontend::ggml::FrontEnd::convert(input_model, naive); - if (getenv("GGML_OPENVINO_DUMP_IR")) { - ov::serialize(model, "IR_naive.xml"); - } - - std::shared_ptr infer_request; - auto remote_context = ggml_openvino_get_remote_context(); - if (cgraph->nodes[0]->op == GGML_OP_MUL_MAT) { - // TODO ACCURACY hint triggers a bug in GPU plugin/driver on Lunar Lake. Remove once CVS-182166 is resolved - core.set_property(device, ov::hint::execution_mode(ov::hint::ExecutionMode::PERFORMANCE)); - } else { - core.set_property(device, ov::hint::execution_mode(ov::hint::ExecutionMode::ACCURACY)); - } - if (remote_context.has_value()) { - infer_request = std::make_shared( - core.compile_model(model, remote_context.value(), config).create_infer_request()); - } else { - infer_request = - std::make_shared(core.compile_model(model, device, config).create_infer_request()); - } - - auto ov_params = model->get_parameters(); - for (size_t i = 0; i < ov_params.size(); i++) { - auto param_name = ov_params[i]->get_friendly_name(); - auto input_tensor = get_ov_input_tensor(decoder, param_name); - infer_request->set_input_tensor(i, input_tensor); - } - - auto ov_results = model->get_results(); - for (size_t i = 0; i < ov_results.size(); i++) { - auto * ggml_tensor = decoder->get_model_outputs().at(ov_results[i]->get_friendly_name()); - auto output_tensor = create_ov_output_tensor(decoder, infer_request, i, ggml_tensor); - infer_request->set_output_tensor(i, output_tensor); - } - - infer_request->infer(); - return GGML_STATUS_SUCCESS; -} - -namespace { -ov::Tensor convert_ggml_input_to_ov(std::shared_ptr ggml_decoder, const std::string & name) { - const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(name); - - if (ggml_tensor->extra != nullptr) { - // GGML_LOG_DEBUG("Using ggml_tensor->extra as ov::Tensor for input: %s\n", name.c_str()); - auto * extra_base = static_cast(ggml_tensor->extra); - if (extra_base->type != ggml_openvino_extra_base::Type::TENSOR) { - throw std::runtime_error("ggml tensor extra is not of type TENSOR for input: " + name); - } - auto * tensor_extra = static_cast(extra_base); - return *tensor_extra->tensor; - } - - // GGML_LOG_DEBUG("Converting ggml tensor to ov::Tensor for input: %s\n", name.c_str()); - auto * input_data = ggml_tensor->data; - ov::Shape input_shape; - if (ggml_tensor->op == GGML_OP_VIEW) { - // This case is added to make test-backend-ops work - input_shape = ggml_decoder->get_shape(ggml_tensor->view_src); - } else { - input_shape = ggml_decoder->get_shape(ggml_tensor); - } - auto input_tensor = ov::Tensor(ggml_decoder->get_ov_type(ggml_tensor), input_shape, input_data); - return input_tensor; -} -} // namespace - -ov::Tensor get_ov_input_tensor(std::shared_ptr ggml_decoder, const std::string & param_name) { - ov::Tensor input_tensor; - if (ggml_decoder->get_model_extra_inputs().find(param_name) != ggml_decoder->get_model_extra_inputs().end()) { - input_tensor = *ggml_decoder->get_model_extra_input_values().at(param_name); - } else { - input_tensor = convert_ggml_input_to_ov(ggml_decoder, param_name); - } - return input_tensor; -} - -ov::Tensor get_ov_input_tensor_static_decode(std::shared_ptr ggml_decoder, - const std::string & param_name) { - // NPU decoding stage - const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name); - const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor); - - if (GgmlOvDecoder::is_inp_tok(ggml_tensor, op) || GgmlOvDecoder::is_inp_pos(ggml_tensor, op) || - GgmlOvDecoder::is_kv_idx(ggml_tensor, op)) { - assert(ggml_tensor->ne[0] == 1); - ov::Shape input_shape = {1, 1, 1, 1}; - ov::Tensor input_tensor(ggml_decoder->get_ov_type(ggml_tensor), input_shape); - if (ggml_tensor->type == GGML_TYPE_I32) { - *input_tensor.data() = *((int32_t *) ggml_tensor->data); - } else if (ggml_tensor->type == GGML_TYPE_I64) { - *input_tensor.data() = *((int64_t *) ggml_tensor->data); - } else { - throw std::runtime_error("Unexpected tensor type for " + param_name); - } - return input_tensor; - } - - if (GgmlOvDecoder::is_output_idx(ggml_tensor, op)) { - ov::Shape input_shape = {1, 1, 1, 1}; - ov::Tensor input_tensor(ggml_decoder->get_ov_type(ggml_tensor), input_shape); - int32_t inp_out_id = *((int32_t *) ggml_tensor->data); - assert(ggml_tensor->ne[0] == 1); - assert(inp_out_id == 0); - *input_tensor.data() = inp_out_id; - return input_tensor; - } - - if (GgmlOvDecoder::is_inp_mask(ggml_tensor, op)) { - size_t context_size = ggml_decoder->get_ctx_size(); - std::vector padded_data = pad_input(ggml_tensor, 1, context_size, -INFINITY); - ov::Tensor input_tensor(ov::element::f32, ov::Shape{1, 1, 1, context_size}); - auto * data_ptr = input_tensor.data(); - std::copy(padded_data.begin(), padded_data.begin() + context_size, data_ptr); - return input_tensor; - } - - return get_ov_input_tensor(ggml_decoder, param_name); -} - -ov::Tensor get_ov_input_tensor_static_prefill(std::shared_ptr ggml_decoder, - const std::string & param_name, - int chunk_index) { - // NPU prompt processing stage - const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name); - const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor); - - const size_t input_len = ggml_decoder->get_input_len(); - const size_t chunk_size = ggml_decoder->m_prefill_chunk_size; - const size_t chunk_valid_size = std::min(chunk_size, input_len - chunk_index * chunk_size); - const size_t chunk_pad_size = chunk_size - chunk_valid_size; - - if (GgmlOvDecoder::is_inp_tok(ggml_tensor, op) || GgmlOvDecoder::is_inp_pos(ggml_tensor, op) || - GgmlOvDecoder::is_kv_idx(ggml_tensor, op)) { - ov::Shape input_shape = {1, 1, 1, chunk_size}; - ov::Tensor input_tensor(ggml_decoder->get_ov_type(ggml_tensor), input_shape); - // copy the chunk_index-th chunk from ggml_tensor - size_t element_size = ggml_type_size(ggml_tensor->type); - void * input_data = (char *) ggml_tensor->data + chunk_index * chunk_size * element_size; - std::memcpy(input_tensor.data(), input_data, chunk_valid_size * element_size); - // pad the rest with last_value + 1, so that kv's of padded positions are inserted - // to the next row after the valids row in the kvcache - if (chunk_pad_size > 0) { - if (ggml_tensor->type == GGML_TYPE_I32) { - int32_t last_value = - *((int32_t *) ggml_tensor->data + (chunk_index * chunk_size + chunk_valid_size - 1)); - int32_t * output_data = input_tensor.data(); - std::fill(output_data + chunk_valid_size, output_data + chunk_size, last_value + 1); - } else if (ggml_tensor->type == GGML_TYPE_I64) { - int64_t last_value = - *((int64_t *) ggml_tensor->data + (chunk_index * chunk_size + chunk_valid_size - 1)); - int64_t * output_data = input_tensor.data(); - std::fill(output_data + chunk_valid_size, output_data + chunk_size, last_value + 1); - } else { - throw std::runtime_error("Unexpected tensor type for " + param_name); - } - } - return input_tensor; - } - - if (GgmlOvDecoder::is_output_idx(ggml_tensor, op)) { - size_t output_len = ggml_decoder->get_compute_params().output_len; - ov::Shape input_shape = {1, 1, 1, output_len}; - ov::Tensor input_tensor(ggml_decoder->get_ov_type(ggml_tensor), input_shape); - if (ggml_tensor->ne[0] == 0) { - *input_tensor.data() = 0; - } else { - auto * data_addr = input_tensor.data(); - for (size_t i = 0; i < output_len; i++) { - data_addr[i] = ((int32_t *) ggml_tensor->data)[i] % chunk_size; - } - } - return input_tensor; - } - - if (GgmlOvDecoder::is_inp_mask(ggml_tensor, op)) { - size_t cols = ggml_tensor->ne[0]; - size_t rows = ggml_tensor->ne[1]; - float * ggml_data = (float *) ggml_tensor->data + chunk_index * chunk_size * cols; - size_t chunk_valid_rows = std::min(chunk_size, rows - chunk_index * chunk_size); - size_t context_size = ggml_decoder->get_ctx_size(); - std::vector padded_data = - pad_input(ggml_data, chunk_valid_rows, cols, chunk_size, context_size, -INFINITY); - set_zero_diagonal(padded_data, chunk_size, context_size); - ov::Tensor input_tensor(ov::element::f32, ov::Shape{1, 1, chunk_size, context_size}); - auto * data_ptr = input_tensor.data(); - std::copy(padded_data.begin(), padded_data.begin() + chunk_size * context_size, data_ptr); - return input_tensor; - } - - return get_ov_input_tensor(ggml_decoder, param_name); -} - -size_t checksum(const void * data, size_t size) { - const uint8_t * bytes = static_cast(data); - size_t sum = 0; - for (size_t i = 0; i < size; ++i) { - sum += (uint8_t) i; - sum += bytes[i]; - } - return sum; -} - -void print_input_tensor_info(const std::string & name, const ov::Tensor & tensor) { - std::cout << "Input name: " << name << ", Input shape: " << tensor.get_shape() << ", Address: " << tensor.data() - << std::endl; - switch (tensor.get_element_type()) { - case ov::element::f32: { - if (name.find("self_kq_mask") == std::string::npos) { - std::cout << *(tensor.data()) << std::endl; - } else { - size_t rows = tensor.get_shape()[2]; - size_t cols = tensor.get_shape()[3]; - auto * data = tensor.data(); - for (size_t i = 0; i < rows; ++i) { - for (size_t j = 0; j < cols; ++j) { - float val = data[i * cols + j]; - if (std::isinf(val) && val < 0) { - std::cout << std::setw(5) << "-inf"; - } else { - std::cout << std::setw(5) << val; - } - } - std::cout << std::endl; - } - } - - break; - } - case ov::element::f16: - std::cout << *(tensor.data()) << std::endl; - break; - case ov::element::i32: - for (size_t i = 0; i < tensor.get_size(); ++i) { - std::cout << tensor.data()[i] << " "; - } - std::cout << std::endl; - break; - case ov::element::i64: - for (size_t i = 0; i < tensor.get_size(); ++i) { - std::cout << tensor.data()[i] << " "; - } - std::cout << std::endl; - break; - default: - break; - } -} - -void print_output_tensor_info(const std::string & name, const ov::Tensor & tensor, const void * output_dst) { - std::cout << "Output name: " << name << ", Output shape: " << tensor.get_shape() << ", Address: " << output_dst - << std::endl; - - auto print_float_stats = [](const std::string & type_name, size_t size, auto get_value) { - if (size == 0) { - return; - } - - float first = get_value(0); - float min = first; - float max = first; - double sum = first; - - for (size_t i = 1; i < size; ++i) { - float v = get_value(i); - if (v < min) { - min = v; - } - if (v > max) { - max = v; - } - sum += v; - } - double mean = sum / size; - - std::cout << std::right << std::setw(6) << type_name << std::right << std::setw(12) << "First" << std::setw(12) - << "Min" << std::setw(12) << "Max" << std::setw(12) << "Mean" << std::endl; - std::cout << std::right << std::setw(6) << "" << std::right << std::setw(12) << first << std::setw(12) << min - << std::setw(12) << max << std::setw(12) << mean << std::endl; - }; - - switch (tensor.get_element_type()) { - case ov::element::f32: { - const float * data = tensor.data(); - size_t size = tensor.get_size(); - print_float_stats("[f32]", size, [data](size_t i) { return data[i]; }); - break; - } - case ov::element::f16: { - const ov::float16 * data = tensor.data(); - size_t size = tensor.get_size(); - print_float_stats("[f16]", size, [data](size_t i) { return static_cast(data[i]); }); - break; - } - default: - break; - } -} - -void set_zero_diagonal(std::vector & matrix, size_t rows, size_t cols) { - for (size_t i = 0; i < rows; ++i) { - size_t diag_col = std::min(i, cols - 1); - matrix[i * cols + diag_col] = 0.0f; - } -} - -const ggml_tensor * get_inp_pos_tensor(ggml_cgraph * cgraph) { - for (int i = 0; i < cgraph->n_nodes; ++i) { - auto * op = cgraph->nodes[i]; - for (int j = 0; j < GGML_MAX_SRC; ++j) { - auto * src = op->src[j]; - if (src == nullptr) { - break; - } - if (GgmlOvDecoder::is_inp_pos(src, op)) { - return src; - } - } - } - GGML_LOG_ERROR("get_inp_pos_tensor: inp_pos not found in cgraph"); - throw std::runtime_error("get_inp_pos_tensor: inp_pos not found in cgraph"); -} - -bool get_is_prefill(const ggml_tensor * inp_pos) { - return inp_pos->ne[0] > 1; -} - -#pragma GCC diagnostic pop diff --git a/ggml/src/ggml-openvino/utils.h b/ggml/src/ggml-openvino/utils.h deleted file mode 100644 index 2c72e33c352f..000000000000 --- a/ggml/src/ggml-openvino/utils.h +++ /dev/null @@ -1,143 +0,0 @@ -#include "ggml-backend-impl.h" -#include "ggml-decoder.h" -#include "ggml-impl.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -struct graph_key { - int n_nodes; - std::string first_node_name; - std::string last_node_name; - - graph_key(const ggml_cgraph * cgraph) : n_nodes(cgraph->n_nodes) { - if (n_nodes > 0) { - first_node_name = cgraph->nodes[0]->name; - last_node_name = cgraph->nodes[n_nodes - 1]->name; - } - } - - bool operator==(const graph_key & other) const { - return n_nodes == other.n_nodes && first_node_name == other.first_node_name && - last_node_name == other.last_node_name; - } -}; - -struct graph_key_hash { - size_t operator()(const graph_key & key) const { - size_t h = std::hash{}(key.n_nodes); - if (key.n_nodes > 0) { - h ^= std::hash{}(key.first_node_name) + 0x9e3779b9 + (h << 6) + (h >> 2); - h ^= std::hash{}(key.last_node_name) + 0x9e3779b9 + (h << 6) + (h >> 2); - } - return h; - } -}; - -struct decoder_runtime_ctx { - decoder_runtime_ctx(std::shared_ptr mutex) : mutex(std::move(mutex)) {} - std::shared_ptr mutex; - std::shared_ptr ptr; -}; - -struct ov_runtime_context { - mutable std::mutex ctx_mutex; - std::string device; - bool stateful; - std::unordered_map, graph_key_hash> decoder_cache; - std::unordered_map, graph_key_hash> infer_request_cache; - std::unordered_map, graph_key_hash> infer_request_cache_prefill; - std::unordered_map, graph_key_hash> ov_input_names_cache; - std::unordered_map, graph_key_hash> ov_output_names_cache; - //TODO: Stateful is only supported for single request at a time. - // Simultanous stateful inference request support to be added. - size_t stateful_kv_size; - std::map kv_state_input_name_map; - std::atomic backend_count; - - ov_runtime_context() : - device("CPU"), - stateful(false), - stateful_kv_size(0), - backend_count(0) {} - - void clear_caches() { - std::lock_guard lock(ctx_mutex); - decoder_cache.clear(); - infer_request_cache.clear(); - infer_request_cache_prefill.clear(); - ov_input_names_cache.clear(); - ov_output_names_cache.clear(); - } -}; - -enum ggml_status ov_graph_compute(struct ggml_cgraph * cgraph, ggml_backend_t backend); - -enum ggml_status ov_graph_compute_dynamic(struct ggml_cgraph * cgraph, std::shared_ptr r_ctx); -enum ggml_status ov_graph_compute_static(struct ggml_cgraph * cgraph, std::shared_ptr r_ctx); - -size_t checksum(const void * data, size_t size); - -void print_input_tensor_info(const std::string & name, const ov::Tensor & tensor); - -void print_output_tensor_info(const std::string & name, const ov::Tensor & tensor, const void * output_dst); - -template -std::vector pad_input(const T * data, - size_t rows, - size_t cols, - size_t padded_rows, - size_t padded_cols, - T pad_value) { - std::vector padded(padded_rows * padded_cols, pad_value); - - for (size_t i = 0; i < std::min(rows, padded_rows); ++i) { - for (size_t j = 0; j < std::min(cols, padded_cols); ++j) { - padded[i * padded_cols + j] = data[i * cols + j]; - } - } - - return padded; -} - -template -std::vector pad_input(const ggml_tensor * tensor, size_t padded_rows, size_t padded_cols, T pad_value) { - return pad_input(reinterpret_cast(tensor->data), - static_cast(tensor->ne[1]), // rows - static_cast(tensor->ne[0]), // cols - padded_rows, padded_cols, pad_value); -} - -void set_zero_diagonal(std::vector & matrix, size_t rows, size_t cols); - -const ggml_tensor * get_inp_pos_tensor(struct ggml_cgraph * cgraph); - -bool get_is_prefill(const ggml_tensor * inp_pos); - -ov::Tensor get_ov_input_tensor(std::shared_ptr ggml_decoder, const std::string & param_name); -ov::Tensor get_ov_input_tensor_static_decode(std::shared_ptr ggml_decoder, - const std::string & param_name); -ov::Tensor get_ov_input_tensor_static_prefill(std::shared_ptr ggml_decoder, - const std::string & param_name, - int chunk_index); - -ov::Tensor create_ov_output_tensor(std::shared_ptr ggml_decoder, - std::shared_ptr infer_request, - int output_index, - const ggml_tensor * ggml_tensor); - -bool is_naive(struct ggml_cgraph * cgraph); - -enum ggml_status naive_compute(struct ggml_cgraph * cgraph, - ov::Core & core, - const std::string & device, - const ov::AnyMap & config); diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp deleted file mode 100644 index 53903defa8f4..000000000000 --- a/ggml/src/ggml-opt.cpp +++ /dev/null @@ -1,1094 +0,0 @@ -#include "ggml-opt.h" - -#include "ggml.h" -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "ggml-impl.h" - -#include -#include -#include -#include -#include -#include -#include - -struct ggml_opt_dataset { - struct ggml_context * ctx = nullptr; - ggml_backend_buffer_t buf = nullptr; - struct ggml_tensor * data = nullptr; - struct ggml_tensor * labels = nullptr; - - int64_t ndata = -1; - int64_t ndata_shard = -1; - size_t nbs_data = -1; - size_t nbs_labels = -1; - - std::vector permutation; -}; - -struct ggml_opt_context { - ggml_backend_sched_t backend_sched = nullptr; - ggml_cgraph * allocated_graph = nullptr; - ggml_cgraph * allocated_graph_copy = nullptr; - struct ggml_context * ctx_static = nullptr; - struct ggml_context * ctx_cpu = nullptr; - struct ggml_context * ctx_compute = nullptr; - struct ggml_context * ctx_copy = nullptr; - ggml_backend_buffer_t buf_static = nullptr; - ggml_backend_buffer_t buf_cpu = nullptr; - std::mt19937 rng; - enum ggml_opt_loss_type loss_type; - enum ggml_opt_build_type build_type; - enum ggml_opt_build_type build_type_alloc; - - struct ggml_tensor * inputs = nullptr; - struct ggml_tensor * outputs = nullptr; - struct ggml_tensor * labels = nullptr; - - struct ggml_tensor * loss = nullptr; - struct ggml_tensor * pred = nullptr; - struct ggml_tensor * ncorrect = nullptr; - - struct ggml_cgraph * gf = nullptr; - struct ggml_cgraph * gb_grad = nullptr; - struct ggml_cgraph * gb_opt = nullptr; - bool static_graphs = false; - bool eval_ready = false; - std::vector grad_accs; - std::vector grad_m; - std::vector grad_v; - - int64_t iter = 1; - int32_t opt_period = 1; - int32_t opt_i = 0; - bool loss_per_datapoint = false; - - ggml_opt_get_optimizer_params get_opt_pars = nullptr; - void * get_opt_pars_ud = nullptr; - struct ggml_tensor * opt_step_params = nullptr; // Stores output of get_opt_pars. - - enum ggml_opt_optimizer_type optimizer = GGML_OPT_OPTIMIZER_TYPE_ADAMW; -}; - -struct ggml_opt_result { - int64_t ndata = 0; - std::vector loss; - std::vector pred; - int64_t ncorrect = 0; - - int64_t opt_period = -1; - bool loss_per_datapoint = false; -}; - -// ====== Dataset ====== - -ggml_opt_dataset_t ggml_opt_dataset_init( - enum ggml_type type_data, - enum ggml_type type_label, - int64_t ne_datapoint, - int64_t ne_label, - int64_t ndata, - int64_t ndata_shard) { - GGML_ASSERT(ne_datapoint > 0); - GGML_ASSERT(ne_label >= 0); - GGML_ASSERT(ndata > 0); - GGML_ASSERT(ndata_shard > 0); - - ggml_opt_dataset_t result = new ggml_opt_dataset; - result->ndata = ndata; - result->ndata_shard = ndata_shard; - - { - struct ggml_init_params params = { - /*.mem_size =*/ 2*ggml_tensor_overhead(), - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - result->ctx = ggml_init(params); - } - - result->data = ggml_new_tensor_2d(result->ctx, type_data, ne_datapoint, ndata); - result->nbs_data = ggml_nbytes(result->data) * ndata_shard/ndata; - - if (ne_label > 0) { - result->labels = ggml_new_tensor_2d(result->ctx, type_label, ne_label, ndata); - result->nbs_labels = ggml_nbytes(result->labels) * ndata_shard/ndata; - } else { - result->labels = nullptr; - result->nbs_labels = 0; - } - - result->buf = ggml_backend_alloc_ctx_tensors_from_buft(result->ctx, ggml_backend_cpu_buffer_type()); - - const int64_t nshards = ndata/ndata_shard; - result->permutation.resize(nshards); - for (int64_t i = 0; i < nshards; ++i) { - result->permutation[i] = i; - } - return result; -} - -void ggml_opt_dataset_free(ggml_opt_dataset_t dataset) { - ggml_backend_buffer_free(dataset->buf); - ggml_free(dataset->ctx); - delete dataset; -} - -int64_t ggml_opt_dataset_ndata(ggml_opt_dataset_t dataset) { - return dataset->ndata; -} - -struct ggml_tensor * ggml_opt_dataset_data(ggml_opt_dataset_t dataset) { - return dataset->data; -} - -struct ggml_tensor * ggml_opt_dataset_labels(ggml_opt_dataset_t dataset) { - return dataset->labels; -} - -void ggml_opt_dataset_shuffle(ggml_opt_context_t opt_ctx, ggml_opt_dataset_t dataset, int64_t idata) { - GGML_ASSERT(idata <= dataset->ndata); - - if (idata < 0) { - std::shuffle(dataset->permutation.begin(), dataset->permutation.end(), opt_ctx->rng); - return; - } - - GGML_ASSERT(idata % dataset->ndata_shard == 0); - const int64_t ishard_max = idata / dataset->ndata_shard; - std::shuffle(dataset->permutation.begin(), dataset->permutation.begin() + ishard_max, opt_ctx->rng); -} - -void ggml_opt_dataset_get_batch(ggml_opt_dataset_t dataset, struct ggml_tensor * data_batch, struct ggml_tensor * labels_batch, int64_t ibatch) { - GGML_ASSERT( data_batch && ggml_is_contiguous(data_batch)); - GGML_ASSERT(!labels_batch || ggml_is_contiguous(labels_batch)); - GGML_ASSERT((labels_batch == nullptr) == (dataset->labels == nullptr)); - GGML_ASSERT( data_batch->type == dataset->data->type); - GGML_ASSERT(!labels_batch || labels_batch->type == dataset->labels->type); - - const size_t nb_data_batch = ggml_nbytes(data_batch); - GGML_ASSERT(nb_data_batch % dataset->nbs_data == 0); - const int64_t shards_per_batch = nb_data_batch / dataset->nbs_data; - - if (labels_batch) { - const size_t nb_labels_batch = ggml_nbytes(labels_batch); - GGML_ASSERT(nb_labels_batch == shards_per_batch*dataset->nbs_labels); - } - - GGML_ASSERT((ibatch + 1)*shards_per_batch <= int64_t(dataset->permutation.size())); - - for (int64_t ishard_batch = 0; ishard_batch < shards_per_batch; ++ishard_batch) { - const int64_t ishard = dataset->permutation[ibatch*shards_per_batch + ishard_batch]; - - const char * ptr_data = (const char *) dataset->data->data + ishard*dataset->nbs_data; - ggml_backend_tensor_set(data_batch, ptr_data, ishard_batch*dataset->nbs_data, dataset->nbs_data); - - if (!labels_batch) { - continue; - } - - const char * ptr_labels = (const char *) dataset->labels->data + ishard*dataset->nbs_labels; - ggml_backend_tensor_set(labels_batch, ptr_labels, ishard_batch*dataset->nbs_labels, dataset->nbs_labels); - } -} - -void ggml_opt_dataset_get_batch_host(ggml_opt_dataset_t dataset, void * data_batch, size_t nb_data_batch, void * labels_batch, int64_t ibatch) { - GGML_ASSERT((labels_batch == nullptr) == (dataset->labels == nullptr)); - GGML_ASSERT(nb_data_batch % dataset->nbs_data == 0); - - const int64_t shards_per_batch = nb_data_batch / dataset->nbs_data; - - GGML_ASSERT((ibatch + 1)*shards_per_batch <= int64_t(dataset->permutation.size())); - - for (int64_t ishard_batch = 0; ishard_batch < shards_per_batch; ++ishard_batch) { - const int64_t ishard = dataset->permutation[ibatch*shards_per_batch + ishard_batch]; - - const char * ptr_data = (const char *) dataset->data->data + ishard *dataset->nbs_data; - char * ptr_data_batch = (char *) data_batch + ishard_batch*dataset->nbs_data; - memcpy(ptr_data_batch, ptr_data, dataset->nbs_data); - - if (!labels_batch) { - continue; - } - - const char * ptr_labels = (const char *) dataset->labels->data + ishard *dataset->nbs_labels; - char * ptr_labels_batch = (char *) labels_batch + ishard_batch*dataset->nbs_labels; - memcpy(ptr_labels_batch, ptr_labels, dataset->nbs_labels); - } -} - -// ====== Model / Context ====== - -struct ggml_opt_optimizer_params ggml_opt_get_default_optimizer_params(void * userdata) { - GGML_UNUSED(userdata); - - ggml_opt_optimizer_params result; - - result.adamw.alpha = 0.001f; - result.adamw.beta1 = 0.9f; - result.adamw.beta2 = 0.999f; - result.adamw.eps = 1e-8f; - result.adamw.wd = 0.0f; - - result.sgd.alpha = 1e-3f; - result.sgd.wd = 0.0f; - - return result; -} - - -struct ggml_opt_optimizer_params ggml_opt_get_constant_optimizer_params(void * userdata) { - return *((struct ggml_opt_optimizer_params *) userdata); -} - -struct ggml_opt_params ggml_opt_default_params( - ggml_backend_sched_t backend_sched, - enum ggml_opt_loss_type loss_type) { - return { - /*backend_sched =*/ backend_sched, - /*ctx_compute =*/ nullptr, - /*inputs =*/ nullptr, - /*logits =*/ nullptr, - /*loss_type =*/ loss_type, - /*build_type =*/ GGML_OPT_BUILD_TYPE_OPT, - /*opt_period =*/ 1, - /*get_opt_pars =*/ ggml_opt_get_default_optimizer_params, - /*get_opt_pars_ud =*/ nullptr, - /*optimizer =*/ GGML_OPT_OPTIMIZER_TYPE_ADAMW, - }; -} - -static ggml_tensor * map_tensor(std::map & tensor_map, ggml_context * ctx, ggml_tensor * tensor) { - if (!tensor) { - return nullptr; - } - - if (tensor_map.find(tensor) != tensor_map.end()) { - return tensor_map[tensor]; - } - - ggml_tensor * new_tensor = ggml_dup_tensor(ctx, tensor); - tensor_map[tensor] = new_tensor; - - new_tensor->op = tensor->op; - for (int i = 0; i < GGML_MAX_DIMS; i++) { - new_tensor->nb[i] = tensor->nb[i]; - } - new_tensor->flags = tensor->flags; - memcpy(new_tensor->op_params, tensor->op_params, sizeof(tensor->op_params)); - strcpy(new_tensor->name, tensor->name); - new_tensor->data = tensor->data; - new_tensor->buffer = tensor->buffer; - new_tensor->extra = tensor->extra; - new_tensor->view_offs = tensor->view_offs; - new_tensor->view_src = map_tensor(tensor_map, ctx, tensor->view_src); - for (int i = 0; i < GGML_MAX_SRC; i++) { - new_tensor->src[i] = map_tensor(tensor_map, ctx, tensor->src[i]); - } - - return new_tensor; -} - -static ggml_cgraph * dup_graph(ggml_context * ctx, ggml_cgraph * src) { - std::map tensor_map; - - ggml_cgraph * dst = ggml_new_graph_custom(ctx, src->size, /*grads =*/ true); - - for (int i = 0; i < src->n_leafs; i++) { - ggml_build_forward_expand(dst, map_tensor(tensor_map, ctx, src->leafs[i])); - } - GGML_ASSERT(dst->n_leafs == src->n_leafs); - for (int i = 0; i < src->n_nodes; i++) { - ggml_build_forward_expand(dst, map_tensor(tensor_map, ctx, src->nodes[i])); - } - GGML_ASSERT(dst->n_nodes == src->n_nodes); - for (int i = 0; i < src->n_nodes; ++i) { - const size_t igrad_src = ggml_hash_find(&src->visited_hash_set, src->nodes[i]); - const size_t igrad_dst = ggml_hash_find(&dst->visited_hash_set, dst->nodes[i]); - - GGML_ASSERT(igrad_src != GGML_HASHSET_FULL); - GGML_ASSERT(ggml_bitset_get(src->visited_hash_set.used, igrad_src)); - GGML_ASSERT(igrad_dst != GGML_HASHSET_FULL); - GGML_ASSERT(ggml_bitset_get(dst->visited_hash_set.used, igrad_dst)); - - dst->grads[igrad_dst] = src->grads[igrad_src]; - dst->grad_accs[igrad_dst] = src->grad_accs[igrad_src]; - } - - return dst; -} - -static void ggml_opt_build(ggml_opt_context_t opt_ctx) { - GGML_ASSERT(opt_ctx->ctx_compute && "no compute context set, either use static graphs or set one with ggml_opt_prepare_alloc"); - GGML_ASSERT((!opt_ctx->static_graphs || opt_ctx->inputs->data) && "when using static graphs the inputs must be allocated statically"); - - const enum ggml_opt_optimizer_type optimizer = opt_ctx->optimizer; - - const bool accumulate = opt_ctx->build_type_alloc >= GGML_OPT_BUILD_TYPE_GRAD && - !(opt_ctx->static_graphs && opt_ctx->build_type_alloc == GGML_OPT_BUILD_TYPE_OPT && opt_ctx->opt_period == 1); - - const bool need_momenta = opt_ctx->build_type_alloc == GGML_OPT_BUILD_TYPE_OPT && - opt_ctx->optimizer == GGML_OPT_OPTIMIZER_TYPE_ADAMW; - - ggml_set_input(opt_ctx->inputs); - ggml_set_output(opt_ctx->outputs); - - int n_param = 0; - for (int i = 0; i < opt_ctx->gf->n_nodes; ++i) { - const struct ggml_tensor * node = opt_ctx->gf->nodes[i]; - if (node->flags & GGML_TENSOR_FLAG_PARAM) { - n_param++; - } - GGML_ASSERT(!(node->flags & GGML_TENSOR_FLAG_LOSS) && "support for extra loss terms not implemented"); - } - - if (!opt_ctx->ctx_static) { - // The static context is used for: - // - gradients (1 per loss, 1 tensor per param if using gradient accumulation) - // - optimizer momenta (2 tensors per param) - // - labels (if using static graphs) - // - loss (if using static graphs, up to 5 tensors) - // - pred (if using static graphs) - // - ncorrect (if using static graphs, 2 tensors). - constexpr size_t n_loss = 1; - const size_t tensors_per_param = (accumulate ? 1 : 0) + (need_momenta ? 2 : 0); - const size_t tensors_const = opt_ctx->static_graphs ? 9 : 0; - const size_t size_meta = (n_loss + tensors_per_param*n_param + tensors_const) * ggml_tensor_overhead(); - struct ggml_init_params params = { - /*.mem_size =*/ size_meta, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - opt_ctx->ctx_static = ggml_init(params); - } - GGML_ASSERT(opt_ctx->build_type <= opt_ctx->build_type_alloc); - - { - // The cpu context is allocated statically if using static graphs, dynamically otherwise. - // It is used for: - // - optimizer parameters (1 shared for all optimizer invocations) - const size_t size_meta = 1 * ggml_tensor_overhead(); - struct ggml_init_params params = { - /*.mem_size =*/ size_meta, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - ggml_free(opt_ctx->ctx_cpu); - opt_ctx->ctx_cpu = ggml_init(params); - - ggml_backend_buffer_free(opt_ctx->buf_cpu); - opt_ctx->buf_cpu = nullptr; - } - - struct ggml_context * ctx_results = opt_ctx->static_graphs ? opt_ctx->ctx_static : opt_ctx->ctx_compute; - - switch (opt_ctx->loss_type) { - case GGML_OPT_LOSS_TYPE_MEAN: { - opt_ctx->loss = ggml_sum(ctx_results, opt_ctx->outputs); - ggml_set_name(opt_ctx->loss, "loss_sum"); - const float scale = 1.0f / (opt_ctx->opt_period * ggml_nelements(opt_ctx->outputs)); - opt_ctx->loss = ggml_scale(ctx_results, opt_ctx->loss, scale); - ggml_set_name(opt_ctx->loss, "loss_mean"); - opt_ctx->loss_per_datapoint = true; - break; - } - case GGML_OPT_LOSS_TYPE_SUM: { - opt_ctx->loss = ggml_sum(ctx_results, opt_ctx->outputs); - ggml_set_name(opt_ctx->loss, "loss_sum"); - opt_ctx->loss_per_datapoint = false; - break; - } - case GGML_OPT_LOSS_TYPE_CROSS_ENTROPY: { - opt_ctx->labels = ggml_dup_tensor(ctx_results, opt_ctx->outputs); - ggml_set_input(opt_ctx->labels); - ggml_set_name(opt_ctx->labels, "labels"); - opt_ctx->loss = ggml_cross_entropy_loss(ctx_results, opt_ctx->outputs, opt_ctx->labels); - ggml_set_name(opt_ctx->loss, "loss_cross_entropy"); - if (opt_ctx->opt_period > 1) { - opt_ctx->loss = ggml_scale(ctx_results, opt_ctx->loss, 1.0f / opt_ctx->opt_period); - ggml_set_name(opt_ctx->loss, "loss_cross_entropy_scaled"); - } - opt_ctx->loss_per_datapoint = true; - break; - } - case GGML_OPT_LOSS_TYPE_MEAN_SQUARED_ERROR: { - opt_ctx->labels = ggml_dup_tensor(ctx_results, opt_ctx->outputs); - ggml_set_input(opt_ctx->labels); - ggml_set_name(opt_ctx->labels, "labels"); - opt_ctx->loss = ggml_sub(ctx_results, opt_ctx->outputs, opt_ctx->labels); - ggml_set_name(opt_ctx->loss, "loss_error"); - opt_ctx->loss = ggml_sqr(ctx_results, opt_ctx->loss); - ggml_set_name(opt_ctx->loss, "loss_squared_error"); - opt_ctx->loss = ggml_sum(ctx_results, opt_ctx->loss); - ggml_set_name(opt_ctx->loss, "loss_sum_squared_error"); - const float scale = 1.0f / (opt_ctx->opt_period * ggml_nelements(opt_ctx->outputs)); - opt_ctx->loss = ggml_scale(ctx_results, opt_ctx->loss, scale); - ggml_set_name(opt_ctx->loss, "loss_mean_squared_error"); - opt_ctx->loss_per_datapoint = true; - break; - } - } - ggml_set_output(opt_ctx->loss); - ggml_set_loss(opt_ctx->loss); - ggml_build_forward_expand(opt_ctx->gf, opt_ctx->loss); - - if (opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY) { - opt_ctx->pred = ggml_argmax(ctx_results, opt_ctx->outputs); - ggml_set_name(opt_ctx->pred, "pred"); - ggml_set_output(opt_ctx->pred); - ggml_build_forward_expand(opt_ctx->gf, opt_ctx->pred); - - opt_ctx->ncorrect = ggml_count_equal(ctx_results, opt_ctx->pred, ggml_argmax(ctx_results, opt_ctx->labels)); - ggml_set_name(opt_ctx->ncorrect, "ncorrect"); - ggml_set_output(opt_ctx->ncorrect); - ggml_build_forward_expand(opt_ctx->gf, opt_ctx->ncorrect); - } - - if (opt_ctx->buf_static) { - if (opt_ctx->build_type == GGML_OPT_BUILD_TYPE_FORWARD) { - return; - } - } else if (opt_ctx->build_type_alloc == GGML_OPT_BUILD_TYPE_FORWARD) { - opt_ctx->buf_static = ggml_backend_alloc_ctx_tensors( - opt_ctx->ctx_static, ggml_backend_sched_get_backend(opt_ctx->backend_sched, 0)); - return; - } - - if (opt_ctx->grad_accs.empty()) { - GGML_ASSERT(opt_ctx->build_type_alloc >= GGML_OPT_BUILD_TYPE_GRAD); - - const int n_nodes = opt_ctx->gf->n_nodes; - opt_ctx->grad_accs.resize(n_nodes); - for (int i = 0; i < n_nodes; ++i) { - ggml_tensor * node = opt_ctx->gf->nodes[i]; - if ((accumulate && (node->flags & GGML_TENSOR_FLAG_PARAM)) || (node->flags & GGML_TENSOR_FLAG_LOSS)) { - opt_ctx->grad_accs[i] = ggml_new_tensor(opt_ctx->ctx_static, GGML_TYPE_F32, GGML_MAX_DIMS, node->ne); - } else { - opt_ctx->grad_accs[i] = nullptr; - } - } - - if (need_momenta && opt_ctx->build_type_alloc >= GGML_OPT_BUILD_TYPE_OPT) { - opt_ctx->grad_m.resize(n_nodes); - opt_ctx->grad_v.resize(n_nodes); - for (int i = 0; i < n_nodes; ++i) { - ggml_tensor * node = opt_ctx->gf->nodes[i]; - if (node->flags & GGML_TENSOR_FLAG_PARAM) { - opt_ctx->grad_m[i] = ggml_new_tensor(opt_ctx->ctx_static, GGML_TYPE_F32, GGML_MAX_DIMS, node->ne); - opt_ctx->grad_v[i] = ggml_new_tensor(opt_ctx->ctx_static, GGML_TYPE_F32, GGML_MAX_DIMS, node->ne); - } else { - opt_ctx->grad_m[i] = nullptr; - opt_ctx->grad_v[i] = nullptr; - } - } - } - } - - // gb_grad == graph backward gradients, forward pass, then backward pass to calculate gradients. - opt_ctx->gb_grad = ggml_graph_dup(opt_ctx->ctx_compute, opt_ctx->gf, /*force_grads =*/ true); - ggml_build_backward_expand(opt_ctx->ctx_compute, opt_ctx->gb_grad, opt_ctx->grad_accs.data()); - - if (opt_ctx->buf_static) { - if (opt_ctx->build_type == GGML_OPT_BUILD_TYPE_GRAD) { - return; - } - } else if (opt_ctx->build_type_alloc == GGML_OPT_BUILD_TYPE_GRAD) { - opt_ctx->buf_static = ggml_backend_alloc_ctx_tensors(opt_ctx->ctx_static, ggml_backend_sched_get_backend(opt_ctx->backend_sched, 0)); - ggml_graph_reset(opt_ctx->gb_grad); - } - - GGML_ASSERT(opt_ctx->build_type_alloc == GGML_OPT_BUILD_TYPE_OPT); - - // gb_opt == graph backward optimize, forward pass, then backward pass to calculate gradients, then optimizer step. - opt_ctx->gb_opt = ggml_graph_dup(opt_ctx->ctx_compute, opt_ctx->gb_grad, /*force_grads =*/ true); - - opt_ctx->opt_step_params = ggml_new_tensor_1d(opt_ctx->ctx_cpu, GGML_TYPE_F32, need_momenta ? 7 : 2); - ggml_tensor * adamw_params = opt_ctx->opt_step_params; - ggml_set_input(adamw_params); - const char * optimizer_name = ggml_opt_optimizer_name(opt_ctx->optimizer); - ggml_format_name(adamw_params, "%s_params", optimizer_name); - for (int i = opt_ctx->gf->n_nodes-1; i >= 0; --i) { - struct ggml_tensor * node = opt_ctx->gb_opt->nodes[i]; - struct ggml_tensor * grad = ggml_graph_get_grad(opt_ctx->gb_opt, node); - - if (grad && (node->flags & GGML_TENSOR_FLAG_PARAM)) { - struct ggml_tensor * m = nullptr; - struct ggml_tensor * v = nullptr; - if (need_momenta) { - m = opt_ctx->grad_m[i]; - v = opt_ctx->grad_v[i]; - ggml_format_name(m, "AdamW m for %s", node->name); - ggml_format_name(v, "AdamW v for %s", node->name); - } - struct ggml_tensor * opt_step; - switch (optimizer) { - case GGML_OPT_OPTIMIZER_TYPE_ADAMW: - opt_step = ggml_opt_step_adamw(opt_ctx->ctx_compute, node, grad, m, v, adamw_params); - break; - case GGML_OPT_OPTIMIZER_TYPE_SGD: - opt_step = ggml_opt_step_sgd(opt_ctx->ctx_compute, node, grad, adamw_params); - break; - default: - GGML_ABORT("fatal error"); - } - ggml_format_name(opt_step, "%s step for %s", optimizer_name, node->name); - ggml_build_forward_expand(opt_ctx->gb_opt, opt_step); - } - } - - if (!opt_ctx->buf_static) { - opt_ctx->buf_static = ggml_backend_alloc_ctx_tensors( - opt_ctx->ctx_static, ggml_backend_sched_get_backend(opt_ctx->backend_sched, 0)); - ggml_graph_reset(opt_ctx->gb_opt); - } - - opt_ctx->buf_cpu = ggml_backend_alloc_ctx_tensors_from_buft(opt_ctx->ctx_cpu, ggml_backend_cpu_buffer_type()); -} - -ggml_opt_context_t ggml_opt_init(struct ggml_opt_params params) { - ggml_opt_context_t result = new struct ggml_opt_context; - result->backend_sched = params.backend_sched; - result->ctx_compute = params.ctx_compute; - result->loss_type = params.loss_type; - result->build_type = params.build_type; - result->build_type_alloc = params.build_type; - result->inputs = params.inputs; - result->outputs = params.outputs; - result->opt_period = params.opt_period; - result->get_opt_pars = params.get_opt_pars; - result->get_opt_pars_ud = params.get_opt_pars_ud; - result->optimizer = params.optimizer; - - GGML_ASSERT(result->opt_period >= 1); - - result->static_graphs = result->ctx_compute; - - if (!result->static_graphs) { - GGML_ASSERT(!result->inputs); - GGML_ASSERT(!result->outputs); - return result; - } - - GGML_ASSERT(result->inputs); - GGML_ASSERT(result->outputs); - - result->gf = ggml_new_graph_custom(result->ctx_compute, GGML_DEFAULT_GRAPH_SIZE, /*grads =*/ true); // Forward pass. - ggml_build_forward_expand(result->gf, result->outputs); - - ggml_opt_build(result); - - return result; -} - -void ggml_opt_free(ggml_opt_context_t opt_ctx) { - if (opt_ctx == nullptr) { - return; - } - ggml_backend_buffer_free(opt_ctx->buf_static); - ggml_backend_buffer_free(opt_ctx->buf_cpu); - ggml_free(opt_ctx->ctx_static); - ggml_free(opt_ctx->ctx_cpu); - ggml_free(opt_ctx->ctx_copy); - delete opt_ctx; -} - -void ggml_opt_reset(ggml_opt_context_t opt_ctx, bool optimizer) { - if (optimizer) { - ggml_graph_reset(opt_ctx->gb_opt); - opt_ctx->iter = 1; - } else { - ggml_graph_reset(opt_ctx->gb_grad); - } -} - -bool ggml_opt_static_graphs(ggml_opt_context_t opt_ctx) { - return opt_ctx->static_graphs; -} - -struct ggml_tensor * ggml_opt_inputs(ggml_opt_context_t opt_ctx) { - return opt_ctx->inputs; -} - -struct ggml_tensor * ggml_opt_outputs(ggml_opt_context_t opt_ctx) { - return opt_ctx->outputs; -} - -struct ggml_tensor * ggml_opt_labels(ggml_opt_context_t opt_ctx) { - return opt_ctx->labels; -} - -struct ggml_tensor * ggml_opt_loss(ggml_opt_context_t opt_ctx) { - return opt_ctx->loss; -} - -struct ggml_tensor * ggml_opt_pred(ggml_opt_context_t opt_ctx) { - return opt_ctx->pred; -} - -struct ggml_tensor * ggml_opt_ncorrect(ggml_opt_context_t opt_ctx) { - return opt_ctx->ncorrect; -} - -struct ggml_tensor * ggml_opt_grad_acc(ggml_opt_context_t opt_ctx, struct ggml_tensor * node) { - return ggml_graph_get_grad_acc(opt_ctx->gb_opt, node); -} - -// ====== Optimization Result ====== - -ggml_opt_result_t ggml_opt_result_init() { - return new ggml_opt_result; -} - -void ggml_opt_result_free(ggml_opt_result_t result) { - delete result; -} - -void ggml_opt_result_reset(ggml_opt_result_t result) { - result->ndata = 0; - result->loss.clear(); - result->pred.clear(); - result->ncorrect = 0; -} - -void ggml_opt_result_ndata(ggml_opt_result_t result, int64_t * ndata) { - *ndata = result->ndata; -} - -void ggml_opt_result_loss(ggml_opt_result_t result, double * loss, double * unc) { - const int64_t nbatches = result->loss.size(); // Number of physical batches. - - if (nbatches == 0) { - *loss = 0.0; - *unc = NAN; - return; - } - - double sum = 0.0; - double sum_squared = 0.0; - - for (const float & loss : result->loss) { - // If the loss is per datapoint it was scaled by 1.0f/opt_period for each physical batch. - const float loss_scaled = result->loss_per_datapoint ? loss*result->opt_period : loss; - sum += loss_scaled; - sum_squared += loss_scaled*loss_scaled; - } - - const double mean = sum/nbatches; - *loss = result->loss_per_datapoint ? mean : sum; - - if (!unc) { - return; - } - - if (nbatches < 2) { - *unc = NAN; - return; - } - - const double var_sum = sum_squared/nbatches - mean*mean; // variance without Bessel's correction, i.e. nbatches/(nbatches-1) - *unc = result->loss_per_datapoint ? sqrt(var_sum / (nbatches - 1)) : sqrt(var_sum * nbatches/(nbatches - 1)); -} - -void ggml_opt_result_pred(ggml_opt_result_t result, int32_t * pred) { - for (size_t i = 0; i < result->pred.size(); ++i) { - pred[i] = result->pred[i]; - } -} - -void ggml_opt_result_accuracy(ggml_opt_result_t result, double * accuracy, double * unc) { - *accuracy = result->ncorrect >= 0 ? double(result->ncorrect) / double(result->ndata) : NAN; - - if (!unc) { - return; - } - - *unc = result->ncorrect >= 0 && result->ndata >= 2 ? - sqrt((*accuracy) * (1.0 - (*accuracy)) / double(result->ndata - 1)) : NAN; -} - -// ====== Computation ====== - -void ggml_opt_prepare_alloc( - ggml_opt_context_t opt_ctx, - struct ggml_context * ctx_compute, - struct ggml_cgraph * gf, - struct ggml_tensor * inputs, - struct ggml_tensor * outputs) { - GGML_ASSERT(!opt_ctx->static_graphs); - opt_ctx->ctx_compute = ctx_compute; - opt_ctx->gf = gf; - opt_ctx->inputs = inputs; - opt_ctx->outputs = outputs; -} - -void ggml_opt_alloc(ggml_opt_context_t opt_ctx, bool backward) { - GGML_ASSERT(!opt_ctx->eval_ready); - if (opt_ctx->build_type == GGML_OPT_BUILD_TYPE_OPT && opt_ctx->opt_period > 1 && opt_ctx->opt_i == 0) { - ggml_graph_reset(opt_ctx->gb_grad); - } - if (backward) { - const int32_t opt_i_next = (opt_ctx->opt_i + 1) % opt_ctx->opt_period; - opt_ctx->build_type = opt_i_next == 0 ? GGML_OPT_BUILD_TYPE_OPT : GGML_OPT_BUILD_TYPE_GRAD; - } else { - opt_ctx->build_type = GGML_OPT_BUILD_TYPE_FORWARD; - } - - if (!opt_ctx->static_graphs) { - ggml_opt_build(opt_ctx); - } - - struct ggml_cgraph * graph = nullptr; - switch (opt_ctx->build_type) { - case GGML_OPT_BUILD_TYPE_FORWARD: { - graph = opt_ctx->gf; - } break; - case GGML_OPT_BUILD_TYPE_GRAD: { - graph = opt_ctx->gb_grad; - } break; - case GGML_OPT_BUILD_TYPE_OPT: { - graph = opt_ctx->gb_opt; - } break; - } - GGML_ASSERT(graph); - - if (opt_ctx->allocated_graph == graph) { - opt_ctx->eval_ready = true; - return; - } - - ggml_backend_sched_reset(opt_ctx->backend_sched); // clear allocation of previous graph - - if (opt_ctx->static_graphs) { - ggml_init_params params = { - /*.mem_size =*/ graph->size*ggml_tensor_overhead() + ggml_graph_overhead_custom(graph->size, graph->grads), - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - ggml_free(opt_ctx->ctx_copy); - opt_ctx->ctx_copy = ggml_init(params); - - opt_ctx->allocated_graph_copy = dup_graph(opt_ctx->ctx_copy, graph); - } else { - opt_ctx->allocated_graph_copy = graph; - } - - ggml_backend_sched_alloc_graph(opt_ctx->backend_sched, opt_ctx->allocated_graph_copy); - opt_ctx->allocated_graph = graph; - - opt_ctx->eval_ready = true; -} - -void ggml_opt_eval(ggml_opt_context_t opt_ctx, ggml_opt_result_t result) { - GGML_ASSERT(opt_ctx->eval_ready); - if (opt_ctx->allocated_graph == opt_ctx->gb_opt) { - const ggml_opt_optimizer_params & opt_pars = opt_ctx->get_opt_pars(opt_ctx->get_opt_pars_ud); - - switch (opt_ctx->optimizer) { - case GGML_OPT_OPTIMIZER_TYPE_ADAMW: { - GGML_ASSERT(opt_pars.adamw.alpha > 0.0f); - GGML_ASSERT(opt_pars.adamw.beta1 >= 0.0f); - GGML_ASSERT(opt_pars.adamw.beta1 <= 1.0f); - GGML_ASSERT(opt_pars.adamw.beta2 >= 0.0f); - GGML_ASSERT(opt_pars.adamw.beta2 <= 1.0f); - GGML_ASSERT(opt_pars.adamw.eps >= 0.0f); - GGML_ASSERT(opt_pars.adamw.wd >= 0.0f); - GGML_ASSERT(opt_pars.adamw.wd <= 1.0f); - - // beta1, beta2 after applying warmup - const float beta1h = 1.0f / (1.0f - powf(opt_pars.adamw.beta1, opt_ctx->iter)); - const float beta2h = 1.0f / (1.0f - powf(opt_pars.adamw.beta2, opt_ctx->iter)); - - float * adamw_par_data = ggml_get_data_f32(opt_ctx->opt_step_params); - adamw_par_data[0] = opt_pars.adamw.alpha; - adamw_par_data[1] = opt_pars.adamw.beta1; - adamw_par_data[2] = opt_pars.adamw.beta2; - adamw_par_data[3] = opt_pars.adamw.eps; - adamw_par_data[4] = opt_pars.adamw.wd; - adamw_par_data[5] = beta1h; - adamw_par_data[6] = beta2h; - } break; - case GGML_OPT_OPTIMIZER_TYPE_SGD: { - GGML_ASSERT(opt_pars.sgd.alpha > 0.0f); - GGML_ASSERT(opt_pars.sgd.wd >= 0.0f); - GGML_ASSERT(opt_pars.sgd.wd <= 1.0f); - float * sgd = ggml_get_data_f32(opt_ctx->opt_step_params); - sgd[0] = opt_pars.sgd.alpha; - sgd[1] = opt_pars.sgd.wd; - } break; - default: - GGML_ABORT("fatal error"); - } - } - - ggml_backend_sched_graph_compute(opt_ctx->backend_sched, opt_ctx->allocated_graph_copy); - opt_ctx->iter += opt_ctx->allocated_graph == opt_ctx->gb_opt; - opt_ctx->opt_i = (opt_ctx->opt_i + 1) % opt_ctx->opt_period; - - if (!opt_ctx->static_graphs) { - opt_ctx->gf = nullptr; - opt_ctx->gb_grad = nullptr; - opt_ctx->gb_opt = nullptr; - opt_ctx->allocated_graph = nullptr; - opt_ctx->allocated_graph_copy = nullptr; - } - - opt_ctx->eval_ready = false; - - if (!result) { - return; - } - - if (result->ndata == 0) { - result->loss_per_datapoint = opt_ctx->loss_per_datapoint; - result->opt_period = opt_ctx->opt_period; - } else { - GGML_ASSERT(result->loss_per_datapoint == opt_ctx->loss_per_datapoint); - GGML_ASSERT(result->opt_period == opt_ctx->opt_period); - } - - const int64_t ndata = opt_ctx->outputs->ne[1]; - GGML_ASSERT(result->ndata == ndata*int64_t(result->loss.size()) && "varying batch size not supported"); - result->ndata += ndata; - - GGML_ASSERT(ggml_is_scalar(opt_ctx->loss)); - GGML_ASSERT(opt_ctx->loss->type == GGML_TYPE_F32); - float loss; - ggml_backend_tensor_get(opt_ctx->loss, &loss, 0, ggml_nbytes(opt_ctx->loss)); - result->loss.push_back(loss); - - if (opt_ctx->pred) { - GGML_ASSERT(opt_ctx->pred->type == GGML_TYPE_I32); - std::vector pred(ndata); - ggml_backend_tensor_get(opt_ctx->pred, pred.data(), 0, ggml_nbytes(opt_ctx->pred)); - result->pred.insert(result->pred.end(), pred.begin(), pred.end()); - } - - if (!opt_ctx->ncorrect || result->ncorrect < 0) { - result->ncorrect = -1; - return; - } - - GGML_ASSERT(ggml_is_scalar(opt_ctx->ncorrect)); - GGML_ASSERT(opt_ctx->ncorrect->type == GGML_TYPE_I64); - int64_t ncorrect; - ggml_backend_tensor_get(opt_ctx->ncorrect, &ncorrect, 0, ggml_nbytes(opt_ctx->ncorrect)); - result->ncorrect += ncorrect; -} - -// ====== High-Level Functions ====== - -void ggml_opt_epoch( - ggml_opt_context_t opt_ctx, - ggml_opt_dataset_t dataset, - ggml_opt_result_t result_train, - ggml_opt_result_t result_eval, - int64_t idata_split, - ggml_opt_epoch_callback callback_train, - ggml_opt_epoch_callback callback_eval) { - GGML_ASSERT(ggml_opt_static_graphs(opt_ctx) && "ggml_opt_epoch requires static graphs"); - struct ggml_tensor * inputs = ggml_opt_inputs(opt_ctx); - struct ggml_tensor * labels = ggml_opt_labels(opt_ctx); - struct ggml_tensor * data = ggml_opt_dataset_data(dataset); - GGML_ASSERT(data->ne[0] == inputs->ne[0]); - - const int64_t ndata = data->ne[1]; - const int64_t ndata_batch = inputs->ne[1]; - - GGML_ASSERT(data->ne[1] % inputs->ne[1] == 0); - const int64_t nbatches = ndata/ndata_batch; - - idata_split = idata_split < 0 ? ndata : idata_split; - GGML_ASSERT(idata_split % ndata_batch == 0); - const int64_t ibatch_split = idata_split / ndata_batch; - - int64_t ibatch = 0; - int64_t t_loop_start = ggml_time_us(); - for (; ibatch < ibatch_split; ++ibatch) { - ggml_opt_alloc(opt_ctx, /*backward =*/ true); - ggml_opt_dataset_get_batch(dataset, inputs, labels, ibatch); - ggml_opt_eval(opt_ctx, result_train); - if (callback_train) { - callback_train(true, opt_ctx, dataset, result_train, ibatch+1, ibatch_split, t_loop_start); - } - } - t_loop_start = ggml_time_us(); - for (; ibatch < nbatches; ++ibatch) { - ggml_opt_alloc(opt_ctx, /*backward =*/ false); - ggml_opt_dataset_get_batch(dataset, inputs, labels, ibatch); - ggml_opt_eval(opt_ctx, result_eval); - if (callback_eval) { - callback_eval(false, opt_ctx, dataset, result_eval, ibatch+1-ibatch_split, nbatches-ibatch_split, t_loop_start); - } - } -} - -void ggml_opt_epoch_callback_progress_bar( - bool train, - ggml_opt_context_t opt_ctx, - ggml_opt_dataset_t dataset, - ggml_opt_result_t result, - int64_t ibatch, - int64_t ibatch_max, - int64_t t_start_us) { - fprintf(stderr, "%s[", train ? "train: " : "val: "); - - // The progress bar consists of partially filled blocks, unicode has 8 separate fill levels. - constexpr int64_t bar_length = 8; - const int64_t ibatch8 = 8 * ibatch; - for (int64_t j = 0; j < bar_length; ++j) { - if (ibatch_max * (8*j + 8) / bar_length < ibatch8) { - fprintf(stderr, "\u2588"); // full block - } else if (ibatch_max * (8*j + 7) / bar_length < ibatch8) { - fprintf(stderr, "\u2589"); // 7/8 filled - } else if (ibatch_max * (8*j + 6) / bar_length < ibatch8) { - fprintf(stderr, "\u258A"); // 6/8 filled - } else if (ibatch_max * (8*j + 5) / bar_length < ibatch8) { - fprintf(stderr, "\u258B"); // 5/8 filled - } else if (ibatch_max * (8*j + 4) / bar_length < ibatch8) { - fprintf(stderr, "\u258C"); // 4/8 filled - } else if (ibatch_max * (8*j + 3) / bar_length < ibatch8) { - fprintf(stderr, "\u258D"); // 3/8 filled - } else if (ibatch_max * (8*j + 2) / bar_length < ibatch8) { - fprintf(stderr, "\u258E"); // 2/8 filled - } else if (ibatch_max * (8*j + 1) / bar_length < ibatch8) { - fprintf(stderr, "\u258F"); // 1/8 filled - } else { - fprintf(stderr, " "); - } - } - - const int64_t batch_size = ggml_opt_inputs(opt_ctx)->ne[1]; - const int64_t idata = ibatch*batch_size; - const int64_t idata_max = ibatch_max*batch_size; - - double loss; - double loss_unc; - ggml_opt_result_loss(result, &loss, &loss_unc); - - double accuracy; - double accuracy_unc; - ggml_opt_result_accuracy(result, &accuracy, &accuracy_unc); - - const int64_t t_ibatch_us = ggml_time_us() - t_start_us; - int64_t t_ibatch_s = t_ibatch_us / 1000000; - const int64_t t_ibatch_h = t_ibatch_s / 3600; - t_ibatch_s -= t_ibatch_h * 3600; - const int64_t t_ibatch_m = t_ibatch_s / 60; - t_ibatch_s -= t_ibatch_m * 60; - - const int64_t t_eta_us = t_ibatch_us * (ibatch_max - ibatch)/ibatch; - int64_t t_eta_s = t_eta_us / 1000000; - const int64_t t_eta_h = t_eta_s / 3600; - t_eta_s -= t_eta_h * 3600; - const int64_t t_eta_m = t_eta_s / 60; - t_eta_s -= t_eta_m * 60; - - fprintf(stderr, "] data=%07" PRId64 "/%07" PRId64 " loss=%.5lf±%.5lf acc=%.2lf±%.2lf%% " - "t=%02" PRId64 ":%02" PRId64 ":%02" PRId64 " ETA=%02" PRId64 ":%02" PRId64 ":%02" PRId64 " \r", - idata, idata_max, loss, loss_unc, 100.0*accuracy, 100.0*accuracy_unc, - t_ibatch_h, t_ibatch_m, t_ibatch_s, t_eta_h, t_eta_m, t_eta_s); - if (ibatch == ibatch_max) { - fprintf(stderr, "\n"); - } - fflush(stderr); - - GGML_UNUSED(dataset); -} - -void ggml_opt_fit( - ggml_backend_sched_t backend_sched, - ggml_context * ctx_compute, - ggml_tensor * inputs, - ggml_tensor * outputs, - ggml_opt_dataset_t dataset, - enum ggml_opt_loss_type loss_type, - enum ggml_opt_optimizer_type optimizer, - ggml_opt_get_optimizer_params get_opt_pars, - int64_t nepoch, - int64_t nbatch_logical, - float val_split, - bool silent) { - ggml_time_init(); - const int64_t t_start_us = ggml_time_us(); - - const int64_t ndata = ggml_opt_dataset_data(dataset)->ne[1]; - const int64_t nbatch_physical = inputs->ne[1]; - GGML_ASSERT(ndata % nbatch_logical == 0); - GGML_ASSERT(nbatch_logical % nbatch_physical == 0); - - const int64_t opt_period = nbatch_logical / nbatch_physical; - const int64_t nbatches_logical = ndata / nbatch_logical; - - GGML_ASSERT(val_split >= 0.0f); - GGML_ASSERT(val_split < 1.0f); - const int64_t ibatch_split = int64_t(((1.0f - val_split) * nbatches_logical)) * opt_period; // train <-> val split index (physical) - const int64_t idata_split = ibatch_split * nbatch_physical; - - int64_t epoch = 1; - - ggml_opt_params params = ggml_opt_default_params(backend_sched, loss_type); - params.ctx_compute = ctx_compute; - params.inputs = inputs; - params.outputs = outputs; - params.opt_period = opt_period; - params.get_opt_pars = get_opt_pars; - params.get_opt_pars_ud = &epoch; - params.optimizer = optimizer; - ggml_opt_context_t opt_ctx = ggml_opt_init(params); - - // Shuffling the data is generally useful but there is only a point if not all data is used in a single batch. - if (nbatch_logical < ndata) { - ggml_opt_dataset_shuffle(opt_ctx, dataset, -1); // Shuffle all data (train + validation). - } - - ggml_opt_result_t result_train = ggml_opt_result_init(); - ggml_opt_result_t result_val = ggml_opt_result_init(); - - ggml_opt_epoch_callback epoch_callback = silent ? nullptr : ggml_opt_epoch_callback_progress_bar; - - for (; epoch <= nepoch; ++epoch) { - if (nbatch_logical < idata_split) { - ggml_opt_dataset_shuffle(opt_ctx, dataset, idata_split); - } - - ggml_opt_result_reset(result_train); - ggml_opt_result_reset(result_val); - - if (!silent) { - fprintf(stderr, "%s: epoch %04" PRId64 "/%04" PRId64 ":\n", __func__, epoch, nepoch); - } - ggml_opt_epoch(opt_ctx, dataset, result_train, result_val, idata_split, epoch_callback, epoch_callback); - if (!silent) { - fprintf(stderr, "\n"); - } - } - - if (!silent) { - int64_t t_total_s = (ggml_time_us() - t_start_us) / 1000000; - const int64_t t_total_h = t_total_s / 3600; - t_total_s -= t_total_h * 3600; - const int64_t t_total_m = t_total_s / 60; - t_total_s -= t_total_m * 60; - fprintf(stderr, "%s: training took %02" PRId64 ":%02" PRId64 ":%02" PRId64 "\n", __func__, t_total_h, t_total_m, t_total_s); - } - - ggml_opt_free(opt_ctx); - ggml_opt_result_free(result_train); - ggml_opt_result_free(result_val); -} - -enum ggml_opt_optimizer_type ggml_opt_context_optimizer_type(ggml_opt_context_t c) { - return c->optimizer; -} - -GGML_API const char * ggml_opt_optimizer_name(enum ggml_opt_optimizer_type o) { - switch (o) { - case GGML_OPT_OPTIMIZER_TYPE_ADAMW: - return "adamw"; - case GGML_OPT_OPTIMIZER_TYPE_SGD: - return "sgd"; - default: - return "undefined"; - }; -} diff --git a/ggml/src/ggml-quants.h b/ggml/src/ggml-quants-defs.inc similarity index 99% rename from ggml/src/ggml-quants.h rename to ggml/src/ggml-quants-defs.inc index d56c86da8909..7e825c550515 100644 --- a/ggml/src/ggml-quants.h +++ b/ggml/src/ggml-quants-defs.inc @@ -1,9 +1,9 @@ #pragma once #define GGML_COMMON_DECL_C -#include "ggml-common.h" +#include "ggml-common-defs.inc" -#include "ggml.h" +#include "ggml.h.inc" // GGML internal header diff --git a/ggml/src/ggml-quants.c b/ggml/src/ggml-quants.c deleted file mode 100644 index 15443aa554a4..000000000000 --- a/ggml/src/ggml-quants.c +++ /dev/null @@ -1,5491 +0,0 @@ -#define GGML_COMMON_IMPL_C -#include "ggml-common.h" - -#include "ggml-quants.h" -#include "ggml-impl.h" -#include "ggml-cpu/ggml-cpu-impl.h" -#include "ggml-cpu.h" - -#include -#include -#include -#include -#include // for qsort -#include // for GGML_ASSERT - -#define GROUP_MAX_EPS 1e-15f -#define GROUP_MAX_EPS_IQ3_XXS 1e-8f -#define GROUP_MAX_EPS_IQ2_S 1e-8f -#define GROUP_MAX_EPS_IQ1_M 1e-7f -#define GROUP_MAX_EPS_IQ1_S 1e-12f - -#define UNUSED GGML_UNUSED - -static inline int best_index_int8(int n, const int8_t * 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; -} - -// reference implementation for deterministic creation of model files -void quantize_row_q1_0_ref(const float * GGML_RESTRICT x, block_q1_0 * GGML_RESTRICT y, int64_t k) { - static const int qk = QK1_0; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - float sum_abs = 0.0f; - for (int j = 0; j < qk; j++) { - sum_abs += fabsf(x[i*qk + j]); - } - const float d = sum_abs / qk; - - y[i].d = GGML_FP32_TO_FP16(d); - - // Clear all bits first - for (int j = 0; j < qk / 8; ++j) { - y[i].qs[j] = 0; - } - - // Just store sign of each weight directly (no normalization) - for (int j = 0; j < qk; ++j) { - const int bit_index = j; - const int byte_index = bit_index / 8; - const int bit_offset = bit_index % 8; - - if (x[i*qk + j] >= 0.0f) { - y[i].qs[byte_index] |= (1 << bit_offset); - } - } - } -} - -// reference implementation for deterministic creation of model files -void quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t k) { - static const int qk = QK4_0; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - float amax = 0.0f; // absolute max - float max = 0.0f; - - for (int j = 0; j < qk; j++) { - const float v = x[i*qk + j]; - if (amax < fabsf(v)) { - amax = fabsf(v); - max = v; - } - } - - const float d = max / -8; - const float id = d ? 1.0f/d : 0.0f; - - y[i].d = GGML_FP32_TO_FP16(d); - - for (int j = 0; j < qk/2; ++j) { - const float x0 = x[i*qk + 0 + j]*id; - const float x1 = x[i*qk + qk/2 + j]*id; - - const uint8_t xi0 = MIN(15, (int8_t)(x0 + 8.5f)); - const uint8_t xi1 = MIN(15, (int8_t)(x1 + 8.5f)); - - y[i].qs[j] = xi0; - y[i].qs[j] |= xi1 << 4; - } - } -} - -void quantize_row_q4_1_ref(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t k) { - const int qk = QK4_1; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - float min = FLT_MAX; - float max = -FLT_MAX; - - for (int j = 0; j < qk; j++) { - const float v = x[i*qk + j]; - - if (v < min) min = v; - if (v > max) max = v; - } - - const float d = (max - min) / ((1 << 4) - 1); - const float id = d ? 1.0f/d : 0.0f; - - y[i].d = GGML_FP32_TO_FP16(d); - y[i].m = GGML_FP32_TO_FP16(min); - - for (int j = 0; j < qk/2; ++j) { - const float x0 = (x[i*qk + 0 + j] - min)*id; - const float x1 = (x[i*qk + qk/2 + j] - min)*id; - - const uint8_t xi0 = MIN(15, (int8_t)(x0 + 0.5f)); - const uint8_t xi1 = MIN(15, (int8_t)(x1 + 0.5f)); - - y[i].qs[j] = xi0; - y[i].qs[j] |= xi1 << 4; - } - } -} - -void quantize_row_q5_0_ref(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t k) { - static const int qk = QK5_0; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - float amax = 0.0f; // absolute max - float max = 0.0f; - - for (int j = 0; j < qk; j++) { - const float v = x[i*qk + j]; - if (amax < fabsf(v)) { - amax = fabsf(v); - max = v; - } - } - - const float d = max / -16; - const float id = d ? 1.0f/d : 0.0f; - - y[i].d = GGML_FP32_TO_FP16(d); - - uint32_t qh = 0; - - for (int j = 0; j < qk/2; ++j) { - const float x0 = x[i*qk + 0 + j]*id; - const float x1 = x[i*qk + qk/2 + j]*id; - - const uint8_t xi0 = MIN(31, (int8_t)(x0 + 16.5f)); - const uint8_t xi1 = MIN(31, (int8_t)(x1 + 16.5f)); - - y[i].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4); - - // get the 5-th bit and store it in qh at the right position - qh |= ((xi0 & 0x10u) >> 4) << (j + 0); - qh |= ((xi1 & 0x10u) >> 4) << (j + qk/2); - } - - memcpy(&y[i].qh, &qh, sizeof(qh)); - } -} - -void quantize_row_q5_1_ref(const float * GGML_RESTRICT x, block_q5_1 * GGML_RESTRICT y, int64_t k) { - const int qk = QK5_1; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - float min = FLT_MAX; - float max = -FLT_MAX; - - for (int j = 0; j < qk; j++) { - const float v = x[i*qk + j]; - - if (v < min) min = v; - if (v > max) max = v; - } - - const float d = (max - min) / ((1 << 5) - 1); - const float id = d ? 1.0f/d : 0.0f; - - y[i].d = GGML_FP32_TO_FP16(d); - y[i].m = GGML_FP32_TO_FP16(min); - - uint32_t qh = 0; - - for (int j = 0; j < qk/2; ++j) { - const float x0 = (x[i*qk + 0 + j] - min)*id; - const float x1 = (x[i*qk + qk/2 + j] - min)*id; - - const uint8_t xi0 = (uint8_t)(x0 + 0.5f); - const uint8_t xi1 = (uint8_t)(x1 + 0.5f); - - y[i].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4); - - // get the 5-th bit and store it in qh at the right position - qh |= ((xi0 & 0x10u) >> 4) << (j + 0); - qh |= ((xi1 & 0x10u) >> 4) << (j + qk/2); - } - - memcpy(&y[i].qh, &qh, sizeof(y[i].qh)); - } -} - -// reference implementation for deterministic creation of model files -void quantize_row_q8_0_ref(const float * GGML_RESTRICT x, block_q8_0 * GGML_RESTRICT y, int64_t k) { - assert(k % QK8_0 == 0); - const int nb = k / QK8_0; - - for (int i = 0; i < nb; i++) { - float amax = 0.0f; // absolute max - - for (int j = 0; j < QK8_0; j++) { - const float v = x[i*QK8_0 + j]; - amax = MAX(amax, fabsf(v)); - } - - const float d = amax / ((1 << 7) - 1); - const float id = d ? 1.0f/d : 0.0f; - - y[i].d = GGML_FP32_TO_FP16(d); - - for (int j = 0; j < QK8_0; ++j) { - const float x0 = x[i*QK8_0 + j]*id; - - y[i].qs[j] = roundf(x0); - } - } -} - -// reference implementation for deterministic creation of model files -void quantize_row_q8_1_ref(const float * GGML_RESTRICT x, block_q8_1 * GGML_RESTRICT y, int64_t k) { - assert(QK8_1 == 32); - assert(k % QK8_1 == 0); - const int nb = k / QK8_1; - - for (int i = 0; i < nb; i++) { - float amax = 0.0f; // absolute max - - for (int j = 0; j < QK8_1; j++) { - const float v = x[i*QK8_1 + j]; - amax = MAX(amax, fabsf(v)); - } - - const float d = amax / ((1 << 7) - 1); - const float id = d ? 1.0f/d : 0.0f; - - y[i].d = GGML_FP32_TO_FP16(d); - - int sum = 0; - - for (int j = 0; j < QK8_1/2; ++j) { - const float v0 = x[i*QK8_1 + j]*id; - const float v1 = x[i*QK8_1 + QK8_1/2 + j]*id; - - y[i].qs[ j] = roundf(v0); - y[i].qs[QK8_1/2 + j] = roundf(v1); - - sum += y[i].qs[ j]; - sum += y[i].qs[QK8_1/2 + j]; - } - - y[i].s = GGML_FP32_TO_FP16(sum*d); - } -} - -static inline int best_index_mxfp4(float x, float e) { - int best_index = 0; - float best_err = fabsf(kvalues_mxfp4[0]*e - x); - for (int i = 1; i < 16; i++) { - float err = fabsf(kvalues_mxfp4[i]*e - x); - if (err < best_err) { - best_index = i; - best_err = err; - } - } - return best_index; -} - -void quantize_row_mxfp4_ref(const float * GGML_RESTRICT x, block_mxfp4 * GGML_RESTRICT y, int64_t k) { - static const int qk = QK_MXFP4; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - float amax = 0.0f; // absolute max - - for (int j = 0; j < qk; j++) { - const float v = x[i*qk + j]; - - if (amax < fabsf(v)) { - amax = fabsf(v); - } - } - - const uint8_t e = amax > 0.0f ? (uint8_t) (floorf(log2f(amax)) - 2 + 127) : 0; - - const float d = GGML_E8M0_TO_FP32_HALF(e); - - y[i].e = e; - - for (int j = 0; j < qk/2; ++j) { - const uint8_t x0 = best_index_mxfp4(x[i*qk + 0 + j], d); - const uint8_t x1 = best_index_mxfp4(x[i*qk + qk/2 + j], d); - - y[i].qs[j] = x0; - y[i].qs[j] |= x1 << 4; - } - } -} - -void quantize_row_nvfp4_ref(const float * GGML_RESTRICT x, block_nvfp4 * GGML_RESTRICT y, int64_t k) { - static const int qk = QK_NVFP4; - static const int qk_sub = QK_NVFP4_SUB; - static const int n_sub = QK_NVFP4 / QK_NVFP4_SUB; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - for (int s = 0; s < n_sub; s++) { - const float * xb = x + i*qk + s*qk_sub; - - float amax = 0.0f; - for (int j = 0; j < qk_sub; j++) { - if (amax < fabsf(xb[j])) { - amax = fabsf(xb[j]); - } - } - - // UE4M3 scale: amax / 6.0 maps the max E2M1 value (6.0) to amax - const uint8_t ue = ggml_fp32_to_ue4m3(amax / 6.0f); - y[i].d[s] = ue; - const float d = ggml_ue4m3_to_fp32(ue); - - for (int j = 0; j < qk_sub/2; ++j) { - const uint8_t x0 = best_index_mxfp4(xb[0 + j], d); - const uint8_t x1 = best_index_mxfp4(xb[qk_sub/2 + j], d); - - y[i].qs[s*(qk_sub/2) + j] = x0 | (x1 << 4); - } - } - } -} - -void dequantize_row_q1_0(const block_q1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - static const int qk = QK1_0; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - const float d = GGML_FP16_TO_FP32(x[i].d); - const float neg_d = -d; - - for (int j = 0; j < qk; ++j) { - const int byte_index = j / 8; - const int bit_offset = j % 8; - const uint8_t bit = (x[i].qs[byte_index] >> bit_offset) & 1; - y[i*qk + j] = bit ? d : neg_d; - } - } -} - -void dequantize_row_q4_0(const block_q4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - static const int qk = QK4_0; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - const float d = GGML_FP16_TO_FP32(x[i].d); - - for (int j = 0; j < qk/2; ++j) { - const int x0 = (x[i].qs[j] & 0x0F) - 8; - const int x1 = (x[i].qs[j] >> 4) - 8; - - y[i*qk + j + 0 ] = x0*d; - y[i*qk + j + qk/2] = x1*d; - } - } -} - -void dequantize_row_q4_1(const block_q4_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - static const int qk = QK4_1; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - const float d = GGML_FP16_TO_FP32(x[i].d); - const float m = GGML_FP16_TO_FP32(x[i].m); - - for (int j = 0; j < qk/2; ++j) { - const int x0 = (x[i].qs[j] & 0x0F); - const int x1 = (x[i].qs[j] >> 4); - - y[i*qk + j + 0 ] = x0*d + m; - y[i*qk + j + qk/2] = x1*d + m; - } - } -} - -void dequantize_row_q5_0(const block_q5_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - static const int qk = QK5_0; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - const float d = GGML_FP16_TO_FP32(x[i].d); - - uint32_t qh; - memcpy(&qh, x[i].qh, sizeof(qh)); - - for (int j = 0; j < qk/2; ++j) { - const uint8_t xh_0 = ((qh >> (j + 0)) << 4) & 0x10; - const uint8_t xh_1 = ((qh >> (j + 12)) ) & 0x10; - - const int32_t x0 = ((x[i].qs[j] & 0x0F) | xh_0) - 16; - const int32_t x1 = ((x[i].qs[j] >> 4) | xh_1) - 16; - - y[i*qk + j + 0 ] = x0*d; - y[i*qk + j + qk/2] = x1*d; - } - } -} - -void dequantize_row_q5_1(const block_q5_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - static const int qk = QK5_1; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - const float d = GGML_FP16_TO_FP32(x[i].d); - const float m = GGML_FP16_TO_FP32(x[i].m); - - uint32_t qh; - memcpy(&qh, x[i].qh, sizeof(qh)); - - for (int j = 0; j < qk/2; ++j) { - const uint8_t xh_0 = ((qh >> (j + 0)) << 4) & 0x10; - const uint8_t xh_1 = ((qh >> (j + 12)) ) & 0x10; - - const int x0 = (x[i].qs[j] & 0x0F) | xh_0; - const int x1 = (x[i].qs[j] >> 4) | xh_1; - - y[i*qk + j + 0 ] = x0*d + m; - y[i*qk + j + qk/2] = x1*d + m; - } - } -} - -void dequantize_row_q8_0(const block_q8_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - static const int qk = QK8_0; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - const float d = GGML_FP16_TO_FP32(x[i].d); - - for (int j = 0; j < qk; ++j) { - y[i*qk + j] = x[i].qs[j]*d; - } - } -} - -void dequantize_row_mxfp4(const block_mxfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - static const int qk = QK_MXFP4; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - const float d = GGML_E8M0_TO_FP32_HALF(x[i].e); - - for (int j = 0; j < qk/2; ++j) { - const int8_t x0 = kvalues_mxfp4[x[i].qs[j] & 0x0F]; - const int8_t x1 = kvalues_mxfp4[x[i].qs[j] >> 4]; - - y[i*qk + j + 0 ] = x0*d; - y[i*qk + j + qk/2] = x1*d; - } - } -} - -void dequantize_row_nvfp4(const block_nvfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - static const int qk = QK_NVFP4; - static const int qk_sub = QK_NVFP4_SUB; - static const int n_sub = QK_NVFP4 / QK_NVFP4_SUB; - - assert(k % qk == 0); - - const int nb = k / qk; - - for (int i = 0; i < nb; i++) { - for (int s = 0; s < n_sub; s++) { - const float d = ggml_ue4m3_to_fp32(x[i].d[s]); - float * yb = y + i*qk + s*qk_sub; - - for (int j = 0; j < qk_sub/2; ++j) { - const int8_t v0 = kvalues_mxfp4[x[i].qs[s*(qk_sub/2) + j] & 0x0F]; - const int8_t v1 = kvalues_mxfp4[x[i].qs[s*(qk_sub/2) + j] >> 4]; - - yb[j + 0 ] = v0*d; - yb[j + qk_sub/2] = v1*d; - } - } - } -} - -// -// 2-6 bit quantization in super-blocks -// - -// -// ===================== Helper functions -// -static inline int nearest_int(float fval) { - assert(fabsf(fval) <= 4194303.f); - float val = fval + 12582912.f; - int i; memcpy(&i, &val, sizeof(int)); - return (i & 0x007fffff) - 0x00400000; -} - -static float make_qx_quants(int n, int nmax, const float * GGML_RESTRICT x, int8_t * GGML_RESTRICT L, int rmse_type, - const float * GGML_RESTRICT 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 < GROUP_MAX_EPS) { // all zero - for (int i = 0; i < n; ++i) { - L[i] = 0; - } - return 0.f; - } - float iscale = -nmax / max; - if (rmse_type == 0) { - for (int i = 0; i < n; ++i) { - int l = nearest_int(iscale * x[i]); - L[i] = nmax + MAX(-nmax, MIN(nmax-1, l)); - } - return 1/iscale; - } - bool return_early = false; - if (rmse_type < 0) { - rmse_type = -rmse_type; - return_early = true; - } - float sumlx = 0; - float suml2 = 0; -#ifdef HAVE_BUGGY_APPLE_LINKER - // use 'volatile' to prevent unroll and work around a bug in Apple ld64 1015.7 - for (volatile int i = 0; i < n; ++i) { -#else - for (int i = 0; i < n; ++i) { -#endif - int l = nearest_int(iscale * x[i]); - l = MAX(-nmax, MIN(nmax-1, l)); - L[i] = l + nmax; - float w = qw ? qw[i] : rmse_type == 1 ? x[i] * x[i] : rmse_type == 2 ? 1 : rmse_type == 3 ? fabsf(x[i]) : sqrtf(fabsf(x[i])); - sumlx += w*x[i]*l; - suml2 += w*l*l; - } - float scale = suml2 ? sumlx/suml2 : 0.0f; - if (return_early) return suml2 > 0 ? 0.5f*(scale + 1/iscale) : 1/iscale; - 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 = MAX(-nmax, MIN(nmax-1, l)); - float w = qw ? qw[i] : rmse_type == 1 ? x[i] * x[i] : rmse_type == 2 ? 1 : rmse_type == 3 ? fabsf(x[i]) : sqrtf(fabsf(x[i])); - sumlx += w*x[i]*l; - suml2 += w*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 + MAX(-nmax, MIN(nmax-1, l)); - } - scale = sumlx/suml2; best = scale*sumlx; - } - } - return scale; -} - -static float make_q3_quants(int n, int nmax, const float * GGML_RESTRICT x, int8_t * GGML_RESTRICT L, bool do_rmse) { - 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 < GROUP_MAX_EPS) { // all zero - for (int i = 0; i < n; ++i) { L[i] = 0; } - return 0.f; - } - float iscale = -nmax / max; - if (do_rmse) { - float sumlx = 0; - float suml2 = 0; - for (int i = 0; i < n; ++i) { - int l = nearest_int(iscale * x[i]); - l = MAX(-nmax, MIN(nmax-1, l)); - L[i] = l; - float w = x[i]*x[i]; - sumlx += w*x[i]*l; - suml2 += w*l*l; - } - for (int itry = 0; itry < 5; ++itry) { - int n_changed = 0; - for (int i = 0; i < n; ++i) { - float w = x[i]*x[i]; - float slx = sumlx - w*x[i]*L[i]; - if (slx > 0) { - float sl2 = suml2 - w*L[i]*L[i]; - int new_l = nearest_int(x[i] * sl2 / slx); - new_l = MAX(-nmax, MIN(nmax-1, new_l)); - if (new_l != L[i]) { - slx += w*x[i]*new_l; - sl2 += w*new_l*new_l; - if (sl2 > 0 && slx*slx*suml2 > sumlx*sumlx*sl2) { - L[i] = new_l; sumlx = slx; suml2 = sl2; - ++n_changed; - } - } - } - } - if (!n_changed) { - break; - } - } - for (int i = 0; i < n; ++i) { - L[i] += nmax; - } - return suml2 > 0.0f ? sumlx / suml2 : 0.0f; - } - for (int i = 0; i < n; ++i) { - int l = nearest_int(iscale * x[i]); - l = MAX(-nmax, MIN(nmax-1, l)); - L[i] = l + nmax; - } - return 1/iscale; -} - -static float make_qkx1_quants(int n, int nmax, const float * GGML_RESTRICT x, uint8_t * GGML_RESTRICT L, float * GGML_RESTRICT the_min, - int ntry, float alpha) { - float min = x[0]; - float max = x[0]; - for (int i = 1; i < n; ++i) { - if (x[i] < min) min = x[i]; - if (x[i] > max) max = x[i]; - } - if (max == min) { - for (int i = 0; i < n; ++i) L[i] = 0; - *the_min = 0; - return 0.f; - } - if (min > 0) min = 0; - float iscale = nmax/(max - min); - float scale = 1/iscale; - for (int itry = 0; itry < ntry; ++itry) { - float sumlx = 0; int suml2 = 0; - bool did_change = false; - for (int i = 0; i < n; ++i) { - int l = nearest_int(iscale*(x[i] - min)); - l = MAX(0, MIN(nmax, l)); - if (l != L[i]) { - L[i] = l; - did_change = true; - } - sumlx += (x[i] - min)*l; - suml2 += l*l; - } - scale = sumlx/suml2; - float sum = 0; - for (int i = 0; i < n; ++i) { - sum += x[i] - scale*L[i]; - } - min = alpha*min + (1 - alpha)*sum/n; - if (min > 0) min = 0; - iscale = 1/scale; - if (!did_change) break; - } - *the_min = -min; - return scale; -} - -static float make_qkx2_quants(int n, int nmax, const float * GGML_RESTRICT x, const float * GGML_RESTRICT weights, - uint8_t * GGML_RESTRICT L, float * GGML_RESTRICT the_min, uint8_t * GGML_RESTRICT Laux, - float rmin, float rdelta, int nstep, bool use_mad) { - float min = x[0]; - float max = x[0]; - float sum_w = weights[0]; - float sum_x = sum_w * x[0]; -#ifdef HAVE_BUGGY_APPLE_LINKER - // use 'volatile' to prevent unroll and work around a bug in Apple ld64 1015.7 - for (volatile int i = 1; i < n; ++i) { -#else - for (int i = 1; i < n; ++i) { -#endif - if (x[i] < min) min = x[i]; - if (x[i] > max) max = x[i]; - float w = weights[i]; - sum_w += w; - sum_x += w * x[i]; - } - if (min > 0) min = 0; - if (max == min) { - for (int i = 0; i < n; ++i) L[i] = 0; - *the_min = -min; - return 0.f; - } - float iscale = nmax/(max - min); - float scale = 1/iscale; - float best_error = 0; - for (int i = 0; i < n; ++i) { - int l = nearest_int(iscale*(x[i] - min)); - L[i] = MAX(0, MIN(nmax, l)); - float diff = scale * L[i] + min - x[i]; - diff = use_mad ? fabsf(diff) : diff * diff; - float w = weights[i]; - best_error += w * diff; - } - if (nstep < 1) { - *the_min = -min; - return scale; - } - for (int is = 0; is <= nstep; ++is) { - iscale = (rmin + rdelta*is + nmax)/(max - min); - float sum_l = 0, sum_l2 = 0, sum_xl = 0; - for (int i = 0; i < n; ++i) { - int l = nearest_int(iscale*(x[i] - min)); - l = MAX(0, MIN(nmax, l)); - Laux[i] = l; - float w = weights[i]; - sum_l += w*l; - sum_l2 += w*l*l; - sum_xl += w*l*x[i]; - } - float D = sum_w * sum_l2 - sum_l * sum_l; - if (D > 0) { - float this_scale = (sum_w * sum_xl - sum_x * sum_l)/D; - float this_min = (sum_l2 * sum_x - sum_l * sum_xl)/D; - if (this_min > 0) { - this_min = 0; - this_scale = sum_xl / sum_l2; - } - float cur_error = 0; - for (int i = 0; i < n; ++i) { - float diff = this_scale * Laux[i] + this_min - x[i]; - diff = use_mad ? fabsf(diff) : diff * diff; - float w = weights[i]; - cur_error += w * diff; - } - if (cur_error < best_error) { - for (int i = 0; i < n; ++i) { - L[i] = Laux[i]; - } - best_error = cur_error; - scale = this_scale; - min = this_min; - } - } - } - *the_min = -min; - return scale; -} - -static inline void get_scale_min_k4(int j, const uint8_t * GGML_RESTRICT q, uint8_t * GGML_RESTRICT d, uint8_t * GGML_RESTRICT 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); - } -} - -//========================- 2-bit (de)-quantization - -void quantize_row_q2_K_ref(const float * GGML_RESTRICT x, block_q2_K * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int nb = k / QK_K; - - uint8_t L[QK_K]; - uint8_t Laux[16]; - float weights[16]; - float mins[QK_K/16]; - float scales[QK_K/16]; - - const float q4scale = 15.f; - - for (int i = 0; i < nb; i++) { - float max_scale = 0; // as we are deducting the min, scales are always positive - float max_min = 0; - for (int j = 0; j < QK_K/16; ++j) { - for (int l = 0; l < 16; ++l) weights[l] = fabsf(x[16*j + l]); - scales[j] = make_qkx2_quants(16, 3, x + 16*j, weights, L + 16*j, &mins[j], Laux, -0.5f, 0.1f, 15, true); - float scale = scales[j]; - if (scale > max_scale) { - max_scale = scale; - } - float min = mins[j]; - if (min > max_min) { - max_min = min; - } - } - - if (max_scale > 0) { - float iscale = q4scale/max_scale; - for (int j = 0; j < QK_K/16; ++j) { - int l = nearest_int(iscale*scales[j]); - y[i].scales[j] = l; - } - y[i].d = GGML_FP32_TO_FP16(max_scale/q4scale); - } else { - for (int j = 0; j < QK_K/16; ++j) y[i].scales[j] = 0; - y[i].d = GGML_FP32_TO_FP16(0.f); - } - if (max_min > 0) { - float iscale = q4scale/max_min; - for (int j = 0; j < QK_K/16; ++j) { - int l = nearest_int(iscale*mins[j]); - y[i].scales[j] |= (l << 4); - } - y[i].dmin = GGML_FP32_TO_FP16(max_min/q4scale); - } else { - y[i].dmin = GGML_FP32_TO_FP16(0.f); - } - for (int j = 0; j < QK_K/16; ++j) { - const float d = GGML_FP16_TO_FP32(y[i].d) * (y[i].scales[j] & 0xF); - if (!d) continue; - const float dm = GGML_FP16_TO_FP32(y[i].dmin) * (y[i].scales[j] >> 4); - for (int ii = 0; ii < 16; ++ii) { - int l = nearest_int((x[16*j + ii] + dm)/d); - l = MAX(0, MIN(3, l)); - L[16*j + ii] = l; - } - } - - for (int j = 0; j < QK_K; j += 128) { - for (int l = 0; l < 32; ++l) { - y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6); - } - } - - x += QK_K; - } -} - -void dequantize_row_q2_K(const block_q2_K * GGML_RESTRICT x, float * GGML_RESTRICT 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 float min = GGML_FP16_TO_FP32(x[i].dmin); - - const uint8_t * q = x[i].qs; - - int is = 0; - float dl, ml; - for (int n = 0; n < QK_K; n += 128) { - int shift = 0; - for (int j = 0; j < 4; ++j) { - - uint8_t sc = x[i].scales[is++]; - dl = d * (sc & 0xF); ml = min * (sc >> 4); - for (int l = 0; l < 16; ++l) *y++ = dl * ((int8_t)((q[l] >> shift) & 3)) - ml; - - sc = x[i].scales[is++]; - dl = d * (sc & 0xF); ml = min * (sc >> 4); - for (int l = 0; l < 16; ++l) *y++ = dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml; - - shift += 2; - } - q += 32; - } - } -} - -static float make_qkx3_quants(int n, int nmax, const float * GGML_RESTRICT x, const float * GGML_RESTRICT weights, - uint8_t * GGML_RESTRICT L, float * GGML_RESTRICT the_min, uint8_t * GGML_RESTRICT Laux, - float rmin, float rdelta, int nstep, bool use_mad) { - float min = x[0]; - float max = x[0]; - float sum_w = weights ? weights[0] : x[0]*x[0]; - float sum_x = sum_w * x[0]; -#ifdef HAVE_BUGGY_APPLE_LINKER - // use 'volatile' to prevent unroll and work around a bug in Apple ld64 1015.7 - for (volatile int i = 1; i < n; ++i) { -#else - for (int i = 1; i < n; ++i) { -#endif - if (x[i] < min) min = x[i]; - if (x[i] > max) max = x[i]; - float w = weights ? weights[i] : x[i]*x[i]; - sum_w += w; - sum_x += w * x[i]; - } - if (min > 0) { - min = 0; - } - if (max <= min) { - memset(L, 0, n); - *the_min = -min; - return 0.f; - } - float iscale = nmax/(max - min); - float scale = 1/iscale; - float best_mad = 0; - for (int i = 0; i < n; ++i) { - int l = nearest_int(iscale*(x[i] - min)); - L[i] = MAX(0, MIN(nmax, l)); - float diff = scale * L[i] + min - x[i]; - diff = use_mad ? fabsf(diff) : diff*diff; - float w = weights ? weights[i] : x[i]*x[i]; - best_mad += w * diff; - } - if (nstep < 1) { - *the_min = -min; - return scale; - } - for (int is = 0; is <= nstep; ++is) { - iscale = (rmin + rdelta*is + nmax)/(max - min); - float sum_l = 0, sum_l2 = 0, sum_xl = 0; - for (int i = 0; i < n; ++i) { - int l = nearest_int(iscale*(x[i] - min)); - l = MAX(0, MIN(nmax, l)); - Laux[i] = l; - float w = weights ? weights[i] : x[i]*x[i]; - sum_l += w*l; - sum_l2 += w*l*l; - sum_xl += w*l*x[i]; - } - float D = sum_w * sum_l2 - sum_l * sum_l; - if (D > 0) { - float this_scale = (sum_w * sum_xl - sum_x * sum_l)/D; - float this_min = (sum_l2 * sum_x - sum_l * sum_xl)/D; - if (this_min > 0) { - this_min = 0; - this_scale = sum_xl / sum_l2; - } - float mad = 0; - for (int i = 0; i < n; ++i) { - float diff = this_scale * Laux[i] + this_min - x[i]; - diff = use_mad ? fabsf(diff) : diff*diff; - float w = weights ? weights[i] : x[i]*x[i]; - mad += w * diff; - } - if (mad < best_mad) { - for (int i = 0; i < n; ++i) { - L[i] = Laux[i]; - } - best_mad = mad; - scale = this_scale; - min = this_min; - } - } - } - *the_min = -min; - return scale; -} - -static float make_qp_quants(int n, int nmax, const float * GGML_RESTRICT x, uint8_t * GGML_RESTRICT L, const float * quant_weights) { - float max = 0; - for (int i = 0; i < n; ++i) { - max = MAX(max, x[i]); - } - if (max < GROUP_MAX_EPS) { // all zero - for (int i = 0; i < n; ++i) { L[i] = 0; } - return 0.f; - } - float iscale = nmax / max; - for (int i = 0; i < n; ++i) { - L[i] = nearest_int(iscale * x[i]); - } - float scale = 1/iscale; - float best_mse = 0; - for (int i = 0; i < n; ++i) { - float diff = x[i] - scale*L[i]; - float w = quant_weights[i]; - best_mse += w*diff*diff; - } - for (int is = -4; is <= 4; ++is) { - if (is == 0) continue; - float iscale_is = (0.1f*is + nmax)/max; - float scale_is = 1/iscale_is; - float mse = 0; - for (int i = 0; i < n; ++i) { - int l = nearest_int(iscale_is*x[i]); - l = MIN(nmax, l); - float diff = x[i] - scale_is*l; - float w = quant_weights[i]; - mse += w*diff*diff; - } - if (mse < best_mse) { - best_mse = mse; - iscale = iscale_is; - } - } - float sumlx = 0; - float suml2 = 0; - for (int i = 0; i < n; ++i) { - int l = nearest_int(iscale * x[i]); - l = MIN(nmax, l); - L[i] = l; - float w = quant_weights[i]; - sumlx += w*x[i]*l; - suml2 += w*l*l; - } - for (int itry = 0; itry < 5; ++itry) { - int n_changed = 0; - for (int i = 0; i < n; ++i) { - float w = quant_weights[i]; - float slx = sumlx - w*x[i]*L[i]; - float sl2 = suml2 - w*L[i]*L[i]; - if (slx > 0 && sl2 > 0) { - int new_l = nearest_int(x[i] * sl2 / slx); - new_l = MIN(nmax, new_l); - if (new_l != L[i]) { - slx += w*x[i]*new_l; - sl2 += w*new_l*new_l; - if (slx*slx*suml2 > sumlx*sumlx*sl2) { - L[i] = new_l; sumlx = slx; suml2 = sl2; - ++n_changed; - } - } - } - } - if (!n_changed) { - break; - } - } - return suml2 > 0.0f ? sumlx / suml2 : 0.0f; -} - -static void quantize_row_q2_K_impl(const float * GGML_RESTRICT x, block_q2_K * GGML_RESTRICT y, int k, const float * GGML_RESTRICT quant_weights) { - GGML_ASSERT(quant_weights); - assert(k % QK_K == 0); - const int nb = k / QK_K; - const bool requantize = true; - - uint8_t L[QK_K]; - uint8_t Laux[16]; - float mins[QK_K/16]; - float scales[QK_K/16]; - float sw[QK_K/16]; - float weight[16]; - uint8_t Ls[QK_K/16], Lm[QK_K/16]; - - for (int i = 0; i < nb; i++) { - memset(sw, 0, QK_K/16*sizeof(float)); - float sumx2 = 0; - for (int j = 0; j < QK_K; ++j) sumx2 += x[j]*x[j]; - float sigma2 = sumx2/QK_K; - for (int j = 0; j < QK_K/16; ++j) { - const float * GGML_RESTRICT qw = quant_weights + QK_K * i + 16*j; - for (int l = 0; l < 16; ++l) weight[l] = qw[l] * sqrtf(sigma2 + x[16*j + l]*x[16*j + l]); - for (int l = 0; l < QK_K/16; ++l) sw[j] += weight[l]; - scales[j] = make_qkx3_quants(16, 3, x + 16*j, weight, L + 16*j, &mins[j], Laux, -0.9f, 0.05f, 36, false); - } - - float dm, mm; - dm = make_qp_quants(QK_K/16, 15, scales, Ls, sw); - mm = make_qp_quants(QK_K/16, 15, mins, Lm, sw); - - y[i].d = GGML_FP32_TO_FP16(dm); - y[i].dmin = GGML_FP32_TO_FP16(mm); - dm = GGML_FP16_TO_FP32(y[i].d); - mm = GGML_FP16_TO_FP32(y[i].dmin); - - for (int j = 0; j < QK_K/16; ++j) { - y[i].scales[j] = Ls[j] | (Lm[j] << 4); - } - - if (requantize) { - for (int j = 0; j < QK_K/16; ++j) { - const float d = dm * (y[i].scales[j] & 0xF); - if (!d) continue; - const float m = mm * (y[i].scales[j] >> 4); - for (int ii = 0; ii < 16; ++ii) { - int l = nearest_int((x[16*j + ii] + m)/d); - l = MAX(0, MIN(3, l)); - L[16*j + ii] = l; - } - } - } - - for (int j = 0; j < QK_K; j += 128) { - for (int l = 0; l < 32; ++l) { - y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6); - } - } - - x += QK_K; - } -} - -size_t quantize_q2_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - size_t row_size = ggml_row_size(GGML_TYPE_Q2_K, n_per_row); - if (!quant_weights) { - quantize_row_q2_K_ref(src, dst, (int64_t)nrow*n_per_row); - } - else { - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_q2_K_impl(src, (block_q2_K*)qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += row_size; - } - } - return nrow * row_size; -} - -//========================= 3-bit (de)-quantization - -void quantize_row_q3_K_ref(const float * GGML_RESTRICT x, block_q3_K * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int nb = k / QK_K; - - int8_t L[QK_K]; - float scales[QK_K / 16]; - - for (int i = 0; i < nb; i++) { - - float max_scale = 0; - float amax = 0; - for (int j = 0; j < QK_K/16; ++j) { - scales[j] = make_q3_quants(16, 4, x + 16*j, L + 16*j, true); - float scale = fabsf(scales[j]); - if (scale > amax) { - amax = scale; max_scale = scales[j]; - } - } - - memset(y[i].scales, 0, 12); - if (max_scale) { - float iscale = -32.f/max_scale; - for (int j = 0; j < QK_K/16; ++j) { - int8_t l = nearest_int(iscale*scales[j]); - l = MAX(-32, MIN(31, l)) + 32; - if (j < 8) { - y[i].scales[j] = l & 0xF; - } else { - y[i].scales[j-8] |= ((l & 0xF) << 4); - } - l >>= 4; - y[i].scales[j%4 + 8] |= (l << (2*(j/4))); - } - y[i].d = GGML_FP32_TO_FP16(1/iscale); - } else { - y[i].d = GGML_FP32_TO_FP16(0.f); - } - - int8_t sc; - for (int j = 0; j < QK_K/16; ++j) { - sc = j < 8 ? y[i].scales[j] & 0xF : y[i].scales[j-8] >> 4; - sc = (sc | (((y[i].scales[8 + j%4] >> (2*(j/4))) & 3) << 4)) - 32; - float d = GGML_FP16_TO_FP32(y[i].d) * sc; - if (!d) { - continue; - } - for (int ii = 0; ii < 16; ++ii) { - int l = nearest_int(x[16*j + ii]/d); - l = MAX(-4, MIN(3, l)); - L[16*j + ii] = l + 4; - } - } - - memset(y[i].hmask, 0, QK_K/8); - // We put the high-bit for the 1st 8 quants into bit 0, the next 8 into bit 1, etc. - int m = 0; - uint8_t hm = 1; - for (int j = 0; j < QK_K; ++j) { - if (L[j] > 3) { - y[i].hmask[m] |= hm; - L[j] -= 4; - } - if (++m == QK_K/8) { - m = 0; hm <<= 1; - } - } - for (int j = 0; j < QK_K; j += 128) { - for (int l = 0; l < 32; ++l) { - y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6); - } - } - - x += QK_K; - } -} - -void dequantize_row_q3_K(const block_q3_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int nb = k / QK_K; - - const uint32_t kmask1 = 0x03030303; - const uint32_t kmask2 = 0x0f0f0f0f; - - uint32_t aux[4]; - const int8_t * scales = (const int8_t*)aux; - - for (int i = 0; i < nb; i++) { - - const float d_all = GGML_FP16_TO_FP32(x[i].d); - - const uint8_t * GGML_RESTRICT q = x[i].qs; - const uint8_t * GGML_RESTRICT hm = x[i].hmask; - uint8_t m = 1; - - memcpy(aux, x[i].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); - - int is = 0; - float dl; - for (int n = 0; n < QK_K; n += 128) { - int shift = 0; - for (int j = 0; j < 4; ++j) { - - dl = d_all * (scales[is++] - 32); - for (int l = 0; l < 16; ++l) { - *y++ = dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4)); - } - - dl = d_all * (scales[is++] - 32); - for (int l = 0; l < 16; ++l) { - *y++ = dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4)); - } - - shift += 2; - m <<= 1; - } - q += 32; - } - - } -} - -static void quantize_row_q3_K_impl(const float * GGML_RESTRICT x, block_q3_K * GGML_RESTRICT y, int64_t n_per_row, const float * GGML_RESTRICT quant_weights) { - assert(n_per_row % QK_K == 0); - const int nb = n_per_row / QK_K; - - int8_t L[QK_K]; - float scales[QK_K / 16]; - float weight[16]; - float sw[QK_K / 16]; - int8_t Ls[QK_K / 16]; - - for (int i = 0; i < nb; i++) { - - float sumx2 = 0; - for (int j = 0; j < QK_K; ++j) sumx2 += x[j]*x[j]; - float sigma2 = 2*sumx2/QK_K; - - for (int j = 0; j < QK_K/16; ++j) { - if (quant_weights) { - const float * qw = quant_weights + QK_K * i + 16*j; - for (int l = 0; l < 16; ++l) weight[l] = qw[l] * sqrtf(sigma2 + x[16*j+l]*x[16*j+l]); - } else { - for (int l = 0; l < 16; ++l) weight[l] = x[16*j+l]*x[16*j+l]; - } - float sumw = 0; - for (int l = 0; l < 16; ++l) sumw += weight[l]; - sw[j] = sumw; - - scales[j] = make_qx_quants(16, 4, x + 16*j, L + 16*j, 1, weight); - - } - - memset(y[i].scales, 0, 12); - - float d_block = make_qx_quants(QK_K/16, 32, scales, Ls, 1, sw); - for (int j = 0; j < QK_K/16; ++j) { - int l = Ls[j]; - if (j < 8) { - y[i].scales[j] = l & 0xF; - } else { - y[i].scales[j-8] |= ((l & 0xF) << 4); - } - l >>= 4; - y[i].scales[j%4 + 8] |= (l << (2*(j/4))); - } - y[i].d = GGML_FP32_TO_FP16(d_block); - - int8_t sc; - for (int j = 0; j < QK_K/16; ++j) { - sc = j < 8 ? y[i].scales[j] & 0xF : y[i].scales[j-8] >> 4; - sc = (sc | (((y[i].scales[8 + j%4] >> (2*(j/4))) & 3) << 4)) - 32; - float d = GGML_FP16_TO_FP32(y[i].d) * sc; - if (!d) { - continue; - } - for (int ii = 0; ii < 16; ++ii) { - int l = nearest_int(x[16*j + ii]/d); - l = MAX(-4, MIN(3, l)); - L[16*j + ii] = l + 4; - } - } - - memset(y[i].hmask, 0, QK_K/8); - // We put the high-bit for the 1st 8 quants into bit 0, the next 8 into bit 1, etc. - int m = 0; - uint8_t hm = 1; - for (int j = 0; j < QK_K; ++j) { - if (L[j] > 3) { - y[i].hmask[m] |= hm; - L[j] -= 4; - } - if (++m == QK_K/8) { - m = 0; hm <<= 1; - } - } - for (int j = 0; j < QK_K; j += 128) { - for (int l = 0; l < 32; ++l) { - y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6); - } - } - - x += QK_K; - } -} - -size_t quantize_q3_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - size_t row_size = ggml_row_size(GGML_TYPE_Q3_K, n_per_row); - if (!quant_weights) { - quantize_row_q3_K_ref(src, dst, (int64_t)nrow*n_per_row); - } - else { - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_q3_K_impl(src, (block_q3_K*)qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += row_size; - } - } - return nrow * row_size; -} - -// ====================== 4-bit (de)-quantization - -void quantize_row_q4_K_ref(const float * GGML_RESTRICT x, block_q4_K * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int nb = k / QK_K; - - uint8_t L[QK_K]; - uint8_t Laux[32]; - float weights[32]; - float mins[QK_K/32]; - float scales[QK_K/32]; - - for (int i = 0; i < nb; i++) { - float max_scale = 0; // as we are deducting the min, scales are always positive - float max_min = 0; - for (int j = 0; j < QK_K/32; ++j) { - //scales[j] = make_qkx1_quants(32, 15, x + 32*j, L + 32*j, &mins[j], 9, 0.5f); - float sum_x2 = 0; - for (int l = 0; l < 32; ++l) sum_x2 += x[32*j + l] * x[32*j + l]; - float av_x = sqrtf(sum_x2/32); - for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]); - scales[j] = make_qkx2_quants(32, 15, x + 32*j, weights, L + 32*j, &mins[j], Laux, -1.f, 0.1f, 20, false); - float scale = scales[j]; - if (scale > max_scale) { - max_scale = scale; - } - float min = mins[j]; - if (min > max_min) { - max_min = min; - } - } - - float inv_scale = max_scale > 0 ? 63.f/max_scale : 0.f; - float inv_min = max_min > 0 ? 63.f/max_min : 0.f; - for (int j = 0; j < QK_K/32; ++j) { - uint8_t ls = nearest_int(inv_scale*scales[j]); - uint8_t lm = nearest_int(inv_min*mins[j]); - ls = MIN(63, ls); - lm = MIN(63, lm); - if (j < 4) { - y[i].scales[j] = ls; - y[i].scales[j+4] = lm; - } else { - y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4); - y[i].scales[j-4] |= ((ls >> 4) << 6); - y[i].scales[j-0] |= ((lm >> 4) << 6); - } - } - y[i].d = GGML_FP32_TO_FP16(max_scale/63.f); - y[i].dmin = GGML_FP32_TO_FP16(max_min/63.f); - - uint8_t sc, m; - for (int j = 0; j < QK_K/32; ++j) { - get_scale_min_k4(j, y[i].scales, &sc, &m); - const float d = GGML_FP16_TO_FP32(y[i].d) * sc; - if (!d) continue; - const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m; - for (int ii = 0; ii < 32; ++ii) { - int l = nearest_int((x[32*j + ii] + dm)/d); - l = MAX(0, MIN(15, l)); - L[32*j + ii] = l; - } - } - - uint8_t * q = y[i].qs; - for (int j = 0; j < QK_K; j += 64) { - for (int l = 0; l < 32; ++l) q[l] = L[j + l] | (L[j + l + 32] << 4); - q += 32; - } - - x += QK_K; - } -} - -void dequantize_row_q4_K(const block_q4_K * GGML_RESTRICT x, float * GGML_RESTRICT 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 * q = x[i].qs; - - const float d = GGML_FP16_TO_FP32(x[i].d); - const float min = GGML_FP16_TO_FP32(x[i].dmin); - - int is = 0; - uint8_t sc, m; - for (int j = 0; j < QK_K; j += 64) { - get_scale_min_k4(is + 0, x[i].scales, &sc, &m); - const float d1 = d * sc; const float m1 = min * m; - get_scale_min_k4(is + 1, x[i].scales, &sc, &m); - const float d2 = d * sc; const float m2 = min * m; - for (int l = 0; l < 32; ++l) *y++ = d1 * (q[l] & 0xF) - m1; - for (int l = 0; l < 32; ++l) *y++ = d2 * (q[l] >> 4) - m2; - q += 32; is += 2; - } - } -} - -static void quantize_row_q4_K_impl(const float * GGML_RESTRICT x, block_q4_K * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { - assert(n_per_row % QK_K == 0); - const int64_t nb = n_per_row / QK_K; - - uint8_t L[QK_K]; - uint8_t Laux[32]; - uint8_t Ls[QK_K/32]; - uint8_t Lm[QK_K/32]; - float weights[32]; - float sw[QK_K/32]; - float mins[QK_K/32]; - float scales[QK_K/32]; - - for (int i = 0; i < nb; i++) { - - float sum_x2 = 0; - for (int l = 0; l < QK_K; ++l) sum_x2 += x[l] * x[l]; - float sigma2 = 2*sum_x2/QK_K; - float av_x = sqrtf(sigma2); - - for (int j = 0; j < QK_K/32; ++j) { - if (quant_weights) { - const float * qw = quant_weights + QK_K*i + 32*j; - for (int l = 0; l < 32; ++l) weights[l] = qw[l] * sqrtf(sigma2 + x[32*j + l]*x[32*j + l]); - } else { - for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]); - } - float sumw = 0; - for (int l = 0; l < 32; ++l) sumw += weights[l]; - sw[j] = sumw; - scales[j] = make_qkx3_quants(32, 15, x + 32*j, weights, L + 32*j, &mins[j], Laux, -0.9f, 0.05f, 36, false); - } - - float d_block = make_qp_quants(QK_K/32, 63, scales, Ls, sw); - float m_block = make_qp_quants(QK_K/32, 63, mins, Lm, sw); - for (int j = 0; j < QK_K/32; ++j) { - uint8_t ls = Ls[j]; - uint8_t lm = Lm[j]; - if (j < 4) { - y[i].scales[j] = ls; - y[i].scales[j+4] = lm; - } else { - y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4); - y[i].scales[j-4] |= ((ls >> 4) << 6); - y[i].scales[j-0] |= ((lm >> 4) << 6); - } - } - y[i].d = GGML_FP32_TO_FP16(d_block); - y[i].dmin = GGML_FP32_TO_FP16(m_block); - - uint8_t sc, m; - for (int j = 0; j < QK_K/32; ++j) { - get_scale_min_k4(j, y[i].scales, &sc, &m); - const float d = GGML_FP16_TO_FP32(y[i].d) * sc; - if (!d) continue; - const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m; - for (int ii = 0; ii < 32; ++ii) { - int l = nearest_int((x[32*j + ii] + dm)/d); - l = MAX(0, MIN(15, l)); - L[32*j + ii] = l; - } - } - uint8_t * q = y[i].qs; - for (int j = 0; j < QK_K; j += 64) { - for (int l = 0; l < 32; ++l) q[l] = L[j + l] | (L[j + l + 32] << 4); - q += 32; - } - - x += QK_K; - - } -} - -size_t quantize_q4_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - size_t row_size = ggml_row_size(GGML_TYPE_Q4_K, n_per_row); - if (!quant_weights) { - quantize_row_q4_K_ref(src, dst, (int64_t)nrow*n_per_row); - } - else { - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_q4_K_impl(src, (block_q4_K*)qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += row_size; - } - } - return nrow * row_size; -} - -// ====================== 5-bit (de)-quantization - -void quantize_row_q5_K_ref(const float * GGML_RESTRICT x, block_q5_K * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - uint8_t L[QK_K]; - float mins[QK_K/32]; - float scales[QK_K/32]; - float weights[32]; - uint8_t Laux[32]; - - for (int i = 0; i < nb; i++) { - float max_scale = 0; // as we are deducting the min, scales are always positive - float max_min = 0; - for (int j = 0; j < QK_K/32; ++j) { - //scales[j] = make_qkx1_quants(32, 31, x + 32*j, L + 32*j, &mins[j], 9, 0.5f); - float sum_x2 = 0; - for (int l = 0; l < 32; ++l) sum_x2 += x[32*j + l] * x[32*j + l]; - float av_x = sqrtf(sum_x2/32); - for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]); - scales[j] = make_qkx2_quants(32, 31, x + 32*j, weights, L + 32*j, &mins[j], Laux, -0.5f, 0.1f, 15, false); - float scale = scales[j]; - if (scale > max_scale) { - max_scale = scale; - } - float min = mins[j]; - if (min > max_min) { - max_min = min; - } - } - - float inv_scale = max_scale > 0 ? 63.f/max_scale : 0.f; - float inv_min = max_min > 0 ? 63.f/max_min : 0.f; - for (int j = 0; j < QK_K/32; ++j) { - uint8_t ls = nearest_int(inv_scale*scales[j]); - uint8_t lm = nearest_int(inv_min*mins[j]); - ls = MIN(63, ls); - lm = MIN(63, lm); - if (j < 4) { - y[i].scales[j] = ls; - y[i].scales[j+4] = lm; - } else { - y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4); - y[i].scales[j-4] |= ((ls >> 4) << 6); - y[i].scales[j-0] |= ((lm >> 4) << 6); - } - } - y[i].d = GGML_FP32_TO_FP16(max_scale/63.f); - y[i].dmin = GGML_FP32_TO_FP16(max_min/63.f); - - uint8_t sc, m; - for (int j = 0; j < QK_K/32; ++j) { - get_scale_min_k4(j, y[i].scales, &sc, &m); - const float d = GGML_FP16_TO_FP32(y[i].d) * sc; - if (!d) continue; - const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m; - for (int ii = 0; ii < 32; ++ii) { - int l = nearest_int((x[32*j + ii] + dm)/d); - l = MAX(0, MIN(31, l)); - L[32*j + ii] = l; - } - } - - uint8_t * GGML_RESTRICT qh = y[i].qh; - uint8_t * GGML_RESTRICT ql = y[i].qs; - memset(qh, 0, QK_K/8); - - uint8_t m1 = 1, m2 = 2; - for (int n = 0; n < QK_K; n += 64) { - for (int j = 0; j < 32; ++j) { - int l1 = L[n + j]; - if (l1 > 15) { - l1 -= 16; qh[j] |= m1; - } - int l2 = L[n + j + 32]; - if (l2 > 15) { - l2 -= 16; qh[j] |= m2; - } - ql[j] = l1 | (l2 << 4); - } - m1 <<= 2; m2 <<= 2; - ql += 32; - } - - x += QK_K; - } -} - -void dequantize_row_q5_K(const block_q5_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - for (int i = 0; i < nb; i++) { - const uint8_t * ql = x[i].qs; - const uint8_t * qh = x[i].qh; - - const float d = GGML_FP16_TO_FP32(x[i].d); - const float min = GGML_FP16_TO_FP32(x[i].dmin); - - int is = 0; - uint8_t sc, m; - uint8_t u1 = 1, u2 = 2; - for (int j = 0; j < QK_K; j += 64) { - get_scale_min_k4(is + 0, x[i].scales, &sc, &m); - const float d1 = d * sc; const float m1 = min * m; - get_scale_min_k4(is + 1, x[i].scales, &sc, &m); - const float d2 = d * sc; const float m2 = min * m; - for (int l = 0; l < 32; ++l) *y++ = d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1; - for (int l = 0; l < 32; ++l) *y++ = d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2; - ql += 32; is += 2; - u1 <<= 2; u2 <<= 2; - } - } -} - -static void quantize_row_q5_K_impl(const float * GGML_RESTRICT x, block_q5_K * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { - assert(n_per_row % QK_K == 0); - const int64_t nb = n_per_row / QK_K; - - uint8_t L[QK_K]; - uint8_t Laux[32]; - uint8_t Ls[QK_K/32]; - uint8_t Lm[QK_K/32]; - float mins[QK_K/32]; - float scales[QK_K/32]; - float sw[QK_K/32]; - float weights[32]; - - for (int i = 0; i < nb; i++) { - - float sum_x2 = 0; - for (int l = 0; l < QK_K; ++l) sum_x2 += x[l] * x[l]; - float sigma2 = 2*sum_x2/QK_K; - float av_x = sqrtf(sigma2); - - for (int j = 0; j < QK_K/32; ++j) { - if (quant_weights) { - const float * qw = quant_weights + QK_K*i + 32*j; - for (int l = 0; l < 32; ++l) weights[l] = qw[l] * sqrtf(sigma2 + x[32*j + l]*x[32*j + l]); - } else { - for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]); - } - float sumw = 0; - for (int l = 0; l < 32; ++l) sumw += weights[l]; - sw[j] = sumw; - - scales[j] = make_qkx3_quants(32, 31, x + 32*j, weights, L + 32*j, &mins[j], Laux, -0.9f, 0.05f, 36, false); - } - - float d_block = make_qp_quants(QK_K/32, 63, scales, Ls, sw); - float m_block = make_qp_quants(QK_K/32, 63, mins, Lm, sw); - - for (int j = 0; j < QK_K/32; ++j) { - uint8_t ls = Ls[j]; - uint8_t lm = Lm[j]; - ls = MIN(63, ls); - lm = MIN(63, lm); - if (j < 4) { - y[i].scales[j] = ls; - y[i].scales[j+4] = lm; - } else { - y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4); - y[i].scales[j-4] |= ((ls >> 4) << 6); - y[i].scales[j-0] |= ((lm >> 4) << 6); - } - } - y[i].d = GGML_FP32_TO_FP16(d_block); - y[i].dmin = GGML_FP32_TO_FP16(m_block); - - uint8_t sc, m; - for (int j = 0; j < QK_K/32; ++j) { - get_scale_min_k4(j, y[i].scales, &sc, &m); - const float d = GGML_FP16_TO_FP32(y[i].d) * sc; - if (!d) continue; - const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m; - for (int ii = 0; ii < 32; ++ii) { - int l = nearest_int((x[32*j + ii] + dm)/d); - l = MAX(0, MIN(31, l)); - L[32*j + ii] = l; - } - } - - uint8_t * GGML_RESTRICT qh = y[i].qh; - uint8_t * GGML_RESTRICT ql = y[i].qs; - memset(qh, 0, QK_K/8); - - uint8_t m1 = 1, m2 = 2; - for (int n = 0; n < QK_K; n += 64) { - for (int j = 0; j < 32; ++j) { - int l1 = L[n + j]; - if (l1 > 15) { - l1 -= 16; qh[j] |= m1; - } - int l2 = L[n + j + 32]; - if (l2 > 15) { - l2 -= 16; qh[j] |= m2; - } - ql[j] = l1 | (l2 << 4); - } - m1 <<= 2; m2 <<= 2; - ql += 32; - } - - x += QK_K; - - } -} - -size_t quantize_q5_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - size_t row_size = ggml_row_size(GGML_TYPE_Q5_K, n_per_row); - if (!quant_weights) { - quantize_row_q5_K_ref(src, dst, (int64_t)nrow*n_per_row); - } - else { - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_q5_K_impl(src, (block_q5_K*)qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += row_size; - } - } - return nrow * row_size; -} - -// ====================== 6-bit (de)-quantization - -void quantize_row_q6_K_ref(const float * GGML_RESTRICT x, block_q6_K * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - int8_t L[QK_K]; - float scales[QK_K/16]; - - for (int i = 0; i < nb; i++) { - - float max_scale = 0; - float max_abs_scale = 0; - - for (int ib = 0; ib < QK_K/16; ++ib) { - - const float scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, NULL); - scales[ib] = scale; - - const float abs_scale = fabsf(scale); - if (abs_scale > max_abs_scale) { - max_abs_scale = abs_scale; - max_scale = scale; - } - - } - - if (max_abs_scale < GROUP_MAX_EPS) { - memset(&y[i], 0, sizeof(block_q6_K)); - y[i].d = GGML_FP32_TO_FP16(0.f); - x += QK_K; - continue; - } - - float iscale = -128.f/max_scale; - y[i].d = GGML_FP32_TO_FP16(1/iscale); - for (int ib = 0; ib < QK_K/16; ++ib) { - y[i].scales[ib] = MIN(127, nearest_int(iscale*scales[ib])); - } - - for (int j = 0; j < QK_K/16; ++j) { - float d = GGML_FP16_TO_FP32(y[i].d) * y[i].scales[j]; - if (!d) { - continue; - } - for (int ii = 0; ii < 16; ++ii) { - int l = nearest_int(x[16*j + ii]/d); - l = MAX(-32, MIN(31, l)); - L[16*j + ii] = l + 32; - } - } - - uint8_t * GGML_RESTRICT ql = y[i].ql; - uint8_t * GGML_RESTRICT qh = y[i].qh; - for (int j = 0; j < QK_K; j += 128) { - for (int l = 0; l < 32; ++l) { - const uint8_t q1 = L[j + l + 0] & 0xF; - const uint8_t q2 = L[j + l + 32] & 0xF; - const uint8_t q3 = L[j + l + 64] & 0xF; - const uint8_t q4 = L[j + l + 96] & 0xF; - ql[l+ 0] = q1 | (q3 << 4); - ql[l+32] = q2 | (q4 << 4); - qh[l] = (L[j + l] >> 4) | ((L[j + l + 32] >> 4) << 2) | ((L[j + l + 64] >> 4) << 4) | ((L[j + l + 96] >> 4) << 6); - } - ql += 64; - qh += 32; - } - - x += QK_K; - } -} - -void dequantize_row_q6_K(const block_q6_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - for (int i = 0; i < nb; i++) { - const float d = GGML_FP16_TO_FP32(x[i].d); - - const uint8_t * GGML_RESTRICT ql = x[i].ql; - const uint8_t * GGML_RESTRICT qh = x[i].qh; - const int8_t * GGML_RESTRICT sc = x[i].scales; - - for (int n = 0; n < QK_K; n += 128) { - for (int l = 0; l < 32; ++l) { - int is = l/16; - const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; - const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; - const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; - const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; - y[l + 0] = d * sc[is + 0] * q1; - y[l + 32] = d * sc[is + 2] * q2; - y[l + 64] = d * sc[is + 4] * q3; - y[l + 96] = d * sc[is + 6] * q4; - } - y += 128; - ql += 64; - qh += 32; - sc += 8; - } - } -} - -static void quantize_row_q6_K_impl(const float * GGML_RESTRICT x, block_q6_K * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { - assert(n_per_row % QK_K == 0); - const int64_t nb = n_per_row / QK_K; - - int8_t L[QK_K]; - float scales[QK_K/16]; - //float weights[16]; - - for (int i = 0; i < nb; i++) { - - //float sum_x2 = 0; - //for (int j = 0; j < QK_K; ++j) sum_x2 += x[j]*x[j]; - //float sigma2 = sum_x2/QK_K; - - float max_scale = 0; - float max_abs_scale = 0; - - for (int ib = 0; ib < QK_K/16; ++ib) { - - float scale; - if (quant_weights) { - const float * qw = quant_weights + QK_K*i + 16*ib; - //for (int j = 0; j < 16; ++j) weights[j] = qw[j] * sqrtf(sigma2 + x[16*ib + j]*x[16*ib + j]); - //scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, weights); - scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, qw); - } else { - scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, NULL); - } - scales[ib] = scale; - - const float abs_scale = fabsf(scale); - if (abs_scale > max_abs_scale) { - max_abs_scale = abs_scale; - max_scale = scale; - } - - } - - if (max_abs_scale < GROUP_MAX_EPS) { - memset(&y[i], 0, sizeof(block_q6_K)); - y[i].d = GGML_FP32_TO_FP16(0.f); - x += QK_K; - continue; - } - - float iscale = -128.f/max_scale; - y[i].d = GGML_FP32_TO_FP16(1/iscale); - for (int ib = 0; ib < QK_K/16; ++ib) { - y[i].scales[ib] = MIN(127, nearest_int(iscale*scales[ib])); - } - - for (int j = 0; j < QK_K/16; ++j) { - float d = GGML_FP16_TO_FP32(y[i].d) * y[i].scales[j]; - if (!d) { - continue; - } - for (int ii = 0; ii < 16; ++ii) { - int l = nearest_int(x[16*j + ii]/d); - l = MAX(-32, MIN(31, l)); - L[16*j + ii] = l + 32; - } - } - - uint8_t * GGML_RESTRICT ql = y[i].ql; - uint8_t * GGML_RESTRICT qh = y[i].qh; - for (int j = 0; j < QK_K; j += 128) { - for (int l = 0; l < 32; ++l) { - const uint8_t q1 = L[j + l + 0] & 0xF; - const uint8_t q2 = L[j + l + 32] & 0xF; - const uint8_t q3 = L[j + l + 64] & 0xF; - const uint8_t q4 = L[j + l + 96] & 0xF; - ql[l+ 0] = q1 | (q3 << 4); - ql[l+32] = q2 | (q4 << 4); - qh[l] = (L[j + l] >> 4) | ((L[j + l + 32] >> 4) << 2) | ((L[j + l + 64] >> 4) << 4) | ((L[j + l + 96] >> 4) << 6); - } - ql += 64; - qh += 32; - } - - x += QK_K; - - } -} - -size_t quantize_q6_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - size_t row_size = ggml_row_size(GGML_TYPE_Q6_K, n_per_row); - if (!quant_weights) { - quantize_row_q6_K_ref(src, dst, (int64_t)nrow*n_per_row); - } - else { - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_q6_K_impl(src, (block_q6_K*)qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += row_size; - } - } - return nrow * row_size; -} - -static void quantize_row_q4_0_impl(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { - static_assert(QK4_0 == 32, "QK4_0 must be 32"); - - if (!quant_weights) { - quantize_row_q4_0_ref(x, y, n_per_row); - return; - } - - float weight[QK4_0]; - int8_t L[QK4_0]; - - float sum_x2 = 0; - for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j]; - float sigma2 = sum_x2/n_per_row; - - const int64_t nb = n_per_row/QK4_0; - for (int ib = 0; ib < nb; ++ib) { - const float * xb = x + QK4_0 * ib; - const float * qw = quant_weights + QK4_0 * ib; - for (int j = 0; j < QK4_0; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); - float d = make_qx_quants(QK4_0, 8, xb, L, 1, weight); - y[ib].d = GGML_FP32_TO_FP16(d); - for (int j = 0; j < 16; ++j) { - y[ib].qs[j] = L[j] | (L[j+16] << 4); - } - } -} - -size_t quantize_q1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - if (!quant_weights) { - quantize_row_q1_0_ref(src, dst, (int64_t)nrow*n_per_row); - return nrow * ggml_row_size(GGML_TYPE_Q1_0, n_per_row); - } - size_t row_size = ggml_row_size(GGML_TYPE_Q1_0, n_per_row); - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_q1_0_ref(src, (block_q1_0*)qrow, n_per_row); - src += n_per_row; - qrow += row_size; - } - return nrow * row_size; -} - - -size_t quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - if (!quant_weights) { - quantize_row_q4_0_ref(src, dst, (int64_t)nrow*n_per_row); - return nrow * ggml_row_size(GGML_TYPE_Q4_0, n_per_row); - } - size_t row_size = ggml_row_size(GGML_TYPE_Q4_0, n_per_row); - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_q4_0_impl(src, (block_q4_0*)qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += row_size; - } - return nrow * row_size; -} - -static void quantize_row_q4_1_impl(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { - static_assert(QK4_1 == 32, "QK4_1 must be 32"); - - if (!quant_weights) { - quantize_row_q4_1_ref(x, y, n_per_row); - return; - } - - float weight[QK4_1]; - uint8_t L[QK4_1], Laux[QK4_1]; - - float sum_x2 = 0; - for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j]; - float sigma2 = sum_x2/n_per_row; - - const int64_t nb = n_per_row/QK4_1; - for (int ib = 0; ib < nb; ++ib) { - const float * xb = x + QK4_1 * ib; - const float * qw = quant_weights + QK4_1 * ib; - for (int j = 0; j < QK4_1; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); - float min; - float d = make_qkx3_quants(QK4_1, 15, xb, weight, L, &min, Laux, -0.9f, 0.05f, 36, false); - y[ib].d = GGML_FP32_TO_FP16(d); - y[ib].m = GGML_FP32_TO_FP16(-min); - for (int j = 0; j < 16; ++j) { - y[ib].qs[j] = L[j] | (L[j+16] << 4); - } - } -} - -size_t quantize_q4_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - if (!quant_weights) { - quantize_row_q4_1_ref(src, dst, (int64_t)nrow*n_per_row); - return nrow * ggml_row_size(GGML_TYPE_Q4_1, n_per_row); - } - size_t row_size = ggml_row_size(GGML_TYPE_Q4_1, n_per_row); - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_q4_1_impl(src, (block_q4_1*)qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += row_size; - } - return nrow * row_size; -} - -static void quantize_row_q5_0_impl(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { - static_assert(QK5_0 == 32, "QK5_0 must be 32"); - - if (!quant_weights) { - quantize_row_q5_0_ref(x, y, n_per_row); - return; - } - - float weight[QK5_0]; - int8_t L[QK5_0]; - - float sum_x2 = 0; - for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j]; - float sigma2 = sum_x2/n_per_row; - - const int64_t nb = n_per_row/QK5_0; - for (int ib = 0; ib < nb; ++ib) { - const float * xb = x + QK5_0 * ib; - const float * qw = quant_weights + QK5_0 * ib; - for (int j = 0; j < QK5_0; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); - float d = make_qx_quants(QK5_0, 16, xb, L, 1, weight); - y[ib].d = GGML_FP32_TO_FP16(d); - - uint32_t qh = 0; - - for (int j = 0; j < 16; ++j) { - const uint8_t xi0 = L[j]; - const uint8_t xi1 = L[j+16]; - y[ib].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4); - - // get the 5-th bit and store it in qh at the right position - qh |= ((xi0 & 0x10u) >> 4) << (j + 0); - qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_0/2); - } - - memcpy(&y[ib].qh, &qh, sizeof(qh)); - } -} - -size_t quantize_q5_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - if (!quant_weights) { - quantize_row_q5_0_ref(src, dst, (int64_t)nrow*n_per_row); - return nrow * ggml_row_size(GGML_TYPE_Q5_0, n_per_row); - } - size_t row_size = ggml_row_size(GGML_TYPE_Q5_0, n_per_row); - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_q5_0_impl(src, (block_q5_0*)qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += row_size; - } - return nrow * row_size; -} - -static void quantize_row_q5_1_impl(const float * GGML_RESTRICT x, block_q5_1 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { - static_assert(QK5_1 == 32, "QK5_1 must be 32"); - - if (!quant_weights) { - quantize_row_q5_1_ref(x, y, n_per_row); - return; - } - - float weight[QK5_1]; - uint8_t L[QK5_1], Laux[QK5_1]; - - float sum_x2 = 0; - for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j]; - float sigma2 = sum_x2/n_per_row; - - const int64_t nb = n_per_row/QK5_1; - for (int ib = 0; ib < nb; ++ib) { - const float * xb = x + QK5_1 * ib; - const float * qw = quant_weights + QK5_1 * ib; - for (int j = 0; j < QK5_1; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); - float min; - float d = make_qkx3_quants(QK5_1, 31, xb, weight, L, &min, Laux, -0.9f, 0.05f, 36, false); - y[ib].d = GGML_FP32_TO_FP16(d); - y[ib].m = GGML_FP32_TO_FP16(-min); - - uint32_t qh = 0; - for (int j = 0; j < 16; ++j) { - const uint8_t xi0 = L[j]; - const uint8_t xi1 = L[j+16]; - y[ib].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4); - // get the 5-th bit and store it in qh at the right position - qh |= ((xi0 & 0x10u) >> 4) << (j + 0); - qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_0/2); - } - memcpy(&y[ib].qh, &qh, sizeof(qh)); - } -} - -size_t quantize_q5_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - if (!quant_weights) { - quantize_row_q5_1_ref(src, dst, (int64_t)nrow*n_per_row); - return nrow * ggml_row_size(GGML_TYPE_Q5_1, n_per_row); - } - size_t row_size = ggml_row_size(GGML_TYPE_Q5_1, n_per_row); - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_q5_1_impl(src, (block_q5_1*)qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += row_size; - } - return nrow * row_size; -} - -size_t quantize_q8_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - (void)quant_weights; // not used - const size_t row_size = ggml_row_size(GGML_TYPE_Q8_0, n_per_row); - quantize_row_q8_0_ref(src, dst, (int64_t)nrow*n_per_row); - return nrow * row_size; -} - -size_t quantize_mxfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - GGML_UNUSED(quant_weights); - quantize_row_mxfp4_ref(src, dst, (int64_t)nrow*n_per_row); - return nrow * ggml_row_size(GGML_TYPE_MXFP4, n_per_row); -} - -size_t quantize_nvfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - GGML_UNUSED(quant_weights); - quantize_row_nvfp4_ref(src, dst, (int64_t)nrow*n_per_row); - return nrow * ggml_row_size(GGML_TYPE_NVFP4, n_per_row); -} - -// ====================== Ternary (de)-quantization (BitNet b1.58 and TriLMs) - -void quantize_row_tq1_0_ref(const float * GGML_RESTRICT x, block_tq1_0 * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - for (int64_t i = 0; i < nb; i++) { - float amax = 0.0f; // absolute max - - for (int j = 0; j < QK_K; j++) { - const float v = x[j]; - amax = MAX(amax, fabsf(v)); - } - - const float d = amax; - const float id = d ? 1.0f/d : 0.0f; - - y[i].d = GGML_FP32_TO_FP16(d); - - // 5 elements per byte, along 32 bytes - for (size_t j = 0; j < sizeof(y->qs) - sizeof(y->qs) % 32; j += 32) { - for (size_t m = 0; m < 32; ++m) { - uint8_t q = 0; - for (size_t n = 0; n < 5; ++n) { - int xi = lroundf(x[m + n*32] * id) + 1; // -1, 0, 1 -> 0, 1, 2 - q *= 3; - q += xi; - } - // ceiling division (243 == pow(3, 5)) - q = ((uint16_t)q * 256 + (243 - 1)) / 243; - y[i].qs[j + m] = q; - } - x += 5*32; - } - // along 16 bytes - for (size_t j = sizeof(y->qs) - sizeof(y->qs) % 32; j < sizeof(y->qs); j += 16) { - for (size_t m = 0; m < 16; ++m) { - uint8_t q = 0; - for (size_t n = 0; n < 5; ++n) { - int xi = lroundf(x[m + n*16] * id) + 1; // -1, 0, 1 -> 0, 1, 2 - q *= 3; - q += xi; - } - // ceiling division (243 == pow(3, 5)) - q = ((uint16_t)q * 256 + (243 - 1)) / 243; - y[i].qs[j + m] = q; - } - x += 5*16; - } - // 4 elements per byte - for (size_t j = 0; j < sizeof(y->qh); ++j) { - uint8_t q = 0; - for (size_t m = 0; m < 4; ++m) { - // -1, 0, 1 -> 0, 1, 2 - int xi = lroundf(x[j + m*sizeof(y->qh)] * id) + 1; - q *= 3; - q += xi; - } - // shift the first value to the most significant trit - q *= 3; - // ceiling division (243 == pow(3, 5)) - q = ((uint16_t)q * 256 + (243 - 1)) / 243; - y[i].qh[j] = q; - } - x += 4*sizeof(y->qh); - } -} - -void quantize_row_tq2_0_ref(const float * GGML_RESTRICT x, block_tq2_0 * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - for (int64_t i = 0; i < nb; i++) { - float amax = 0.0f; // absolute max - - for (int j = 0; j < QK_K; j++) { - const float v = x[j]; - amax = MAX(amax, fabsf(v)); - } - - const float d = amax; - const float id = d ? 1.0f/d : 0.0f; - - y[i].d = GGML_FP32_TO_FP16(d); - - for (size_t j = 0; j < sizeof(y->qs); j += 32) { - for (size_t m = 0; m < 32; ++m) { - uint8_t q = 0; - for (size_t n = 0; n < 4; ++n) { - // -1, 0, 1 -> 0, 1, 2 - int xi = lroundf(x[m + n*32] * id) + 1; - q += (xi & 3) << (2*n); - } - y[i].qs[j + m] = q; - } - x += 4*32; - } - } -} - -size_t quantize_tq1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - (void)quant_weights; // not used - const size_t row_size = ggml_row_size(GGML_TYPE_TQ1_0, n_per_row); - quantize_row_tq1_0_ref(src, dst, (int64_t)nrow*n_per_row); - return nrow * row_size; -} - -size_t quantize_tq2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - (void)quant_weights; // not used - const size_t row_size = ggml_row_size(GGML_TYPE_TQ2_0, n_per_row); - quantize_row_tq2_0_ref(src, dst, (int64_t)nrow*n_per_row); - return nrow * row_size; -} - -void dequantize_row_tq1_0(const block_tq1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - const uint8_t pow3[6] = {1, 3, 9, 27, 81, 243}; - - for (int64_t i = 0; i < nb; ++i) { - - const float d = GGML_FP16_TO_FP32(x[i].d); - - for (size_t j = 0; j < sizeof(x->qs) - sizeof(x->qs) % 32; j += 32) { - for (size_t n = 0; n < 5; ++n) { - for (size_t m = 0; m < 32; ++m) { - uint8_t q = x[i].qs[j + m] * pow3[n]; - int16_t xi = ((uint16_t) q * 3) >> 8; - *y++ = (float) (xi - 1) * d; - } - } - } - for (size_t j = sizeof(x->qs) - sizeof(x->qs) % 32; j < sizeof(x->qs); j += 16) { - for (size_t n = 0; n < 5; ++n) { - for (size_t m = 0; m < 16; ++m) { - uint8_t q = x[i].qs[j + m] * pow3[n]; - int16_t xi = ((uint16_t) q * 3) >> 8; - *y++ = (float) (xi - 1) * d; - } - } - } - - for (size_t n = 0; n < 4; ++n) { - for (size_t j = 0; j < sizeof(x->qh); ++j) { - uint8_t q = x[i].qh[j] * pow3[n]; - int16_t xi = ((uint16_t) q * 3) >> 8; - *y++ = (float) (xi - 1) * d; - } - } - } -} - -void dequantize_row_tq2_0(const block_tq2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - for (int64_t i = 0; i < nb; ++i) { - - const float d = GGML_FP16_TO_FP32(x[i].d); - - for (size_t j = 0; j < sizeof(x->qs); j += 32) { - for (size_t l = 0; l < 4; ++l) { - for (size_t m = 0; m < 32; ++m) { - int8_t q = (x[i].qs[j + m] >> (l*2)) & 3; - *y++ = (float) (q - 1) * d; - } - } - } - } -} - -// ====================== "True" 2-bit (de)-quantization - -void dequantize_row_iq2_xxs(const block_iq2_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - uint32_t aux32[2]; - const uint8_t * aux8 = (const uint8_t *)aux32; - - for (int i = 0; i < nb; i++) { - - const float d = GGML_FP16_TO_FP32(x[i].d); - - for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { - memcpy(aux32, x[i].qs + 4*ib32, 2*sizeof(uint32_t)); - const float db = d * (0.5f + (aux32[1] >> 28)) * 0.25f; - for (int l = 0; l < 4; ++l) { - const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[l]); - const uint8_t signs = ksigns_iq2xs[(aux32[1] >> 7*l) & 127]; - for (int j = 0; j < 8; ++j) { - y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); - } - y += 8; - } - } - } -} - -// ====================== 2.3125 bpw (de)-quantization - -void dequantize_row_iq2_xs(const block_iq2_xs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - float db[2]; - - for (int i = 0; i < nb; i++) { - - const float d = GGML_FP16_TO_FP32(x[i].d); - - for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { - db[0] = d * (0.5f + (x[i].scales[ib32] & 0xf)) * 0.25f; - db[1] = d * (0.5f + (x[i].scales[ib32] >> 4)) * 0.25f; - for (int l = 0; l < 4; ++l) { - const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (x[i].qs[4*ib32 + l] & 511)); - const uint8_t signs = ksigns_iq2xs[x[i].qs[4*ib32 + l] >> 9]; - for (int j = 0; j < 8; ++j) { - y[j] = db[l/2] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); - } - y += 8; - } - } - } -} - -// ====================== 2.5625 bpw (de)-quantization - -void dequantize_row_iq2_s(const block_iq2_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - float db[2]; - - 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 * signs = qs + QK_K/8; - - for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { - db[0] = d * (0.5f + (x[i].scales[ib32] & 0xf)) * 0.25f; - db[1] = d * (0.5f + (x[i].scales[ib32] >> 4)) * 0.25f; - for (int l = 0; l < 4; ++l) { - const float dl = db[l/2]; - const uint8_t * grid = (const uint8_t *)(iq2s_grid + (qs[l] | (qh[ib32] << (8-2*l) & 0x300))); - for (int j = 0; j < 8; ++j) { - y[j] = dl * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1.f : 1.f); - } - y += 8; - } - qs += 4; - signs += 4; - } - } -} - -// ====================== 3.0625 bpw (de)-quantization - -void dequantize_row_iq3_xxs(const block_iq3_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - uint32_t aux32; - - 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 * scales_and_signs = qs + QK_K/4; - - for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { - memcpy(&aux32, scales_and_signs + 4*ib32, sizeof(uint32_t)); - const float db = d * (0.5f + (aux32 >> 28)) * 0.5f; - for (int l = 0; l < 4; ++l) { - const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*l) & 127]; - const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + qs[2*l+0]); - const uint8_t * grid2 = (const uint8_t *)(iq3xxs_grid + qs[2*l+1]); - for (int j = 0; j < 4; ++j) { - y[j+0] = db * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f); - y[j+4] = db * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f); - } - y += 8; - } - qs += 8; - } - } -} - -// ====================== 3.3125 bpw (de)-quantization - -void dequantize_row_iq3_s(const block_iq3_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t 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 * signs = x[i].signs; - - for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { - const float db1 = d * (1 + 2*(x[i].scales[ib32/2] & 0xf)); - const float db2 = d * (1 + 2*(x[i].scales[ib32/2] >> 4)); - for (int l = 0; l < 4; ++l) { - const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[0] << (8-2*l)) & 256))); - const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[0] << (7-2*l)) & 256))); - for (int j = 0; j < 4; ++j) { - y[j+0] = db1 * grid1[j] * (signs[l] & kmask_iq2xs[j+0] ? -1.f : 1.f); - y[j+4] = db1 * grid2[j] * (signs[l] & kmask_iq2xs[j+4] ? -1.f : 1.f); - } - y += 8; - } - qs += 8; - signs += 4; - for (int l = 0; l < 4; ++l) { - const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[1] << (8-2*l)) & 256))); - const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[1] << (7-2*l)) & 256))); - for (int j = 0; j < 4; ++j) { - y[j+0] = db2 * grid1[j] * (signs[l] & kmask_iq2xs[j+0] ? -1.f : 1.f); - y[j+4] = db2 * grid2[j] * (signs[l] & kmask_iq2xs[j+4] ? -1.f : 1.f); - } - y += 8; - } - qh += 2; - qs += 8; - signs += 4; - } - } -} - -// ====================== 1.5625 bpw (de)-quantization - -void dequantize_row_iq1_s(const block_iq1_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t 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 uint16_t * qh = x[i].qh; - - for (int ib = 0; ib < QK_K/32; ++ib) { - const float dl = d * (2*((qh[ib] >> 12) & 7) + 1); - const float delta = qh[ib] & 0x8000 ? -IQ1S_DELTA : IQ1S_DELTA; - for (int l = 0; l < 4; ++l) { - const int8_t * grid = (const int8_t *)(iq1s_grid + (qs[l] | (((qh[ib] >> 3*l) & 7) << 8))); - for (int j = 0; j < 8; ++j) { - y[j] = dl * (grid[j] + delta); - } - y += 8; - } - qs += 4; - } - } -} - -void dequantize_row_iq1_m(const block_iq1_m * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - float delta[4]; - uint16_t idx[4]; - - iq1m_scale_t scale; - - for (int i = 0; i < nb; i++) { - - const uint16_t * sc = (const uint16_t *)x[i].scales; - scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); - const float d = GGML_FP16_TO_FP32(scale.f16); - - const uint8_t * qs = x[i].qs; - const uint8_t * qh = x[i].qh; - - for (int ib = 0; ib < QK_K/32; ++ib) { - const float dl1 = d * (2*((sc[ib/2] >> (6*(ib%2)+0)) & 0x7) + 1); - const float dl2 = d * (2*((sc[ib/2] >> (6*(ib%2)+3)) & 0x7) + 1); - - idx[0] = qs[0] | ((qh[0] << 8) & 0x700); - idx[1] = qs[1] | ((qh[0] << 4) & 0x700); - idx[2] = qs[2] | ((qh[1] << 8) & 0x700); - idx[3] = qs[3] | ((qh[1] << 4) & 0x700); - delta[0] = qh[0] & 0x08 ? -IQ1S_DELTA : IQ1S_DELTA; - delta[1] = qh[0] & 0x80 ? -IQ1S_DELTA : IQ1S_DELTA; - delta[2] = qh[1] & 0x08 ? -IQ1S_DELTA : IQ1S_DELTA; - delta[3] = qh[1] & 0x80 ? -IQ1S_DELTA : IQ1S_DELTA; - for (int l = 0; l < 2; ++l) { - const int8_t * grid = (const int8_t *)(iq1s_grid + idx[l]); - for (int j = 0; j < 8; ++j) { - y[j] = dl1 * (grid[j] + delta[l]); - } - y += 8; - } - for (int l = 2; l < 4; ++l) { - const int8_t * grid = (const int8_t *)(iq1s_grid + idx[l]); - for (int j = 0; j < 8; ++j) { - y[j] = dl2 * (grid[j] + delta[l]); - } - y += 8; - } - qs += 4; - qh += 2; - } - } -} - -void dequantize_row_iq4_nl(const block_iq4_nl * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK4_NL == 0); - const int64_t nb = k / QK4_NL; - - for (int i = 0; i < nb; i++) { - - const uint8_t * qs = x[i].qs; - - const float d = GGML_FP16_TO_FP32(x[i].d); - for (int j = 0; j < QK4_NL/2; ++j) { - y[j+ 0] = d * kvalues_iq4nl[qs[j] & 0xf]; - y[j+QK4_NL/2] = d * kvalues_iq4nl[qs[j] >> 4]; - } - y += QK4_NL; - qs += QK4_NL/2; - } -} - -void dequantize_row_iq4_xs(const block_iq4_xs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t 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); - - for (int ib = 0; ib < QK_K/32; ++ib) { - const int ls = ((x[i].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[i].scales_h >> 2*ib) & 3) << 4); - const float dl = d * (ls - 32); - for (int j = 0; j < 16; ++j) { - y[j+ 0] = dl * kvalues_iq4nl[qs[j] & 0xf]; - y[j+16] = dl * kvalues_iq4nl[qs[j] >> 4]; - } - y += 32; - qs += 16; - } - } -} - -//===================================== Q8_K ============================================== - -void quantize_row_q8_K_ref(const float * GGML_RESTRICT x, block_q8_K * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - 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); - } - 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; - } - y[i].d = 1/iscale; - x += QK_K; - } -} - -void dequantize_row_q8_K(const block_q8_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - const int64_t nb = k / QK_K; - - for (int i = 0; i < nb; i++) { - for (int j = 0; j < QK_K; ++j) { - *y++ = x[i].d * x[i].qs[j]; - } - } -} - -// ================================ IQ2 quantization ============================================= - -typedef struct { - uint64_t * grid; - int * map; - uint16_t * neighbours; -} iq2_entry_t; - -static iq2_entry_t iq2_data[4] = { - {NULL, NULL, NULL}, - {NULL, NULL, NULL}, - {NULL, NULL, NULL}, - {NULL, NULL, NULL}, -}; - -static inline int iq2_data_index(enum ggml_type type) { - GGML_ASSERT(type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M || type == GGML_TYPE_IQ2_S); - return type == GGML_TYPE_IQ2_XXS ? 0 : - type == GGML_TYPE_IQ2_XS ? 1 : - type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? 2 : 3; -} - -static inline int iq2_grid_size(enum ggml_type type) { - GGML_ASSERT(type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M || type == GGML_TYPE_IQ2_S); - return type == GGML_TYPE_IQ2_XXS ? 256 : - type == GGML_TYPE_IQ2_XS ? 512 : - type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? NGRID_IQ1S : 1024; -} - -static int iq2_compare_func(const void * left, const void * right) { - const int * l = (const int *)left; - const int * r = (const int *)right; - return l[0] < r[0] ? -1 : l[0] > r[0] ? 1 : l[1] < r[1] ? -1 : l[1] > r[1] ? 1 : 0; -} - -void iq2xs_init_impl(enum ggml_type type) { - const int gindex = iq2_data_index(type); - const int grid_size = iq2_grid_size(type); - if (iq2_data[gindex].grid) { - return; - } - static const uint16_t kgrid_2bit_256[256] = { - 0, 2, 5, 8, 10, 17, 20, 32, 34, 40, 42, 65, 68, 80, 88, 97, - 100, 128, 130, 138, 162, 257, 260, 272, 277, 320, 388, 408, 512, 514, 546, 642, - 1025, 1028, 1040, 1057, 1060, 1088, 1090, 1096, 1120, 1153, 1156, 1168, 1188, 1280, 1282, 1288, - 1312, 1350, 1385, 1408, 1425, 1545, 1552, 1600, 1668, 1700, 2048, 2053, 2056, 2068, 2088, 2113, - 2116, 2128, 2130, 2184, 2308, 2368, 2562, 2580, 4097, 4100, 4112, 4129, 4160, 4192, 4228, 4240, - 4245, 4352, 4360, 4384, 4432, 4442, 4480, 4644, 4677, 5120, 5128, 5152, 5157, 5193, 5248, 5400, - 5474, 5632, 5654, 6145, 6148, 6160, 6208, 6273, 6400, 6405, 6560, 6737, 8192, 8194, 8202, 8260, - 8289, 8320, 8322, 8489, 8520, 8704, 8706, 9217, 9220, 9232, 9280, 9302, 9472, 9537, 9572, 9872, - 10248, 10272, 10388, 10820, 16385, 16388, 16400, 16408, 16417, 16420, 16448, 16456, 16470, 16480, 16513, 16516, - 16528, 16640, 16672, 16737, 16768, 16773, 16897, 16912, 16968, 16982, 17000, 17408, 17416, 17440, 17536, 17561, - 17682, 17700, 17920, 18433, 18436, 18448, 18496, 18501, 18688, 18776, 18785, 18818, 19013, 19088, 20480, 20488, - 20497, 20505, 20512, 20608, 20616, 20740, 20802, 20900, 21137, 21648, 21650, 21770, 22017, 22100, 22528, 22545, - 22553, 22628, 22848, 23048, 24580, 24592, 24640, 24680, 24832, 24917, 25112, 25184, 25600, 25605, 25872, 25874, - 25988, 26690, 32768, 32770, 32778, 32833, 32898, 33028, 33048, 33088, 33297, 33793, 33796, 33808, 33813, 33856, - 33888, 34048, 34118, 34196, 34313, 34368, 34400, 34818, 35076, 35345, 36868, 36880, 36900, 36928, 37025, 37142, - 37248, 37445, 37888, 37922, 37956, 38225, 39041, 39200, 40962, 41040, 41093, 41225, 41472, 42008, 43088, 43268, - }; - static const uint16_t kgrid_2bit_512[512] = { - 0, 2, 5, 8, 10, 17, 20, 22, 25, 32, 34, 37, 40, 65, 68, 70, - 73, 80, 82, 85, 88, 97, 100, 128, 130, 133, 136, 145, 148, 153, 160, 257, - 260, 262, 265, 272, 274, 277, 280, 282, 289, 292, 320, 322, 325, 328, 337, 340, - 352, 360, 385, 388, 400, 512, 514, 517, 520, 529, 532, 544, 577, 580, 592, 597, - 640, 650, 1025, 1028, 1030, 1033, 1040, 1042, 1045, 1048, 1057, 1060, 1088, 1090, 1093, 1096, - 1105, 1108, 1110, 1120, 1153, 1156, 1168, 1280, 1282, 1285, 1288, 1297, 1300, 1312, 1345, 1348, - 1360, 1377, 1408, 1537, 1540, 1552, 1574, 1600, 1602, 1668, 2048, 2050, 2053, 2056, 2058, 2065, - 2068, 2080, 2085, 2113, 2116, 2128, 2136, 2176, 2208, 2218, 2305, 2308, 2320, 2368, 2433, 2441, - 2560, 2592, 2600, 2710, 2720, 4097, 4100, 4102, 4105, 4112, 4114, 4117, 4120, 4129, 4132, 4160, - 4162, 4165, 4168, 4177, 4180, 4192, 4202, 4225, 4228, 4240, 4352, 4354, 4357, 4360, 4369, 4372, - 4384, 4417, 4420, 4432, 4480, 4500, 4502, 4609, 4612, 4614, 4624, 4672, 4704, 5120, 5122, 5125, - 5128, 5137, 5140, 5152, 5185, 5188, 5193, 5200, 5220, 5248, 5377, 5380, 5392, 5440, 5632, 5652, - 5705, 6145, 6148, 6160, 6162, 6208, 6228, 6278, 6400, 6405, 6502, 6737, 6825, 8192, 8194, 8197, - 8200, 8202, 8209, 8212, 8224, 8257, 8260, 8272, 8320, 8352, 8449, 8452, 8464, 8512, 8520, 8549, - 8704, 8738, 8832, 8872, 9217, 9220, 9232, 9257, 9280, 9472, 9537, 9554, 9625, 9729, 9754, 9894, - 10240, 10248, 10250, 10272, 10325, 10376, 10402, 10600, 10640, 10760, 10784, 10882, 10888, 10890, 16385, 16388, - 16390, 16393, 16400, 16402, 16405, 16408, 16417, 16420, 16448, 16450, 16453, 16456, 16458, 16465, 16468, 16480, - 16485, 16513, 16516, 16528, 16640, 16642, 16645, 16648, 16657, 16660, 16672, 16705, 16708, 16720, 16768, 16773, - 16802, 16897, 16900, 16912, 16914, 16937, 16960, 17408, 17410, 17413, 17416, 17425, 17428, 17433, 17440, 17473, - 17476, 17488, 17536, 17556, 17665, 17668, 17680, 17700, 17728, 17818, 17920, 17930, 17988, 18000, 18433, 18436, - 18448, 18496, 18501, 18516, 18530, 18688, 18705, 18756, 18768, 18793, 18948, 20480, 20482, 20485, 20488, 20497, - 20500, 20512, 20520, 20545, 20548, 20560, 20608, 20737, 20740, 20752, 20757, 20800, 20802, 20992, 21060, 21162, - 21505, 21508, 21520, 21537, 21568, 21600, 21633, 21665, 21760, 21768, 21888, 21896, 22049, 22120, 22177, 22528, - 22548, 22593, 22608, 22681, 22810, 22848, 22850, 23173, 24577, 24580, 24592, 24640, 24660, 24674, 24710, 24745, - 24832, 25124, 25162, 25234, 25600, 25622, 25872, 25920, 25925, 26020, 26625, 26730, 26917, 27142, 27220, 27234, - 32768, 32770, 32773, 32776, 32785, 32788, 32800, 32810, 32833, 32836, 32848, 32896, 32898, 32936, 32938, 33025, - 33028, 33030, 33040, 33088, 33105, 33113, 33280, 33312, 33408, 33410, 33440, 33448, 33793, 33796, 33808, 33810, - 33813, 33856, 33888, 33929, 34048, 34116, 34213, 34328, 34410, 34816, 34824, 34853, 34906, 34944, 34946, 34984, - 35078, 35362, 35456, 35464, 35478, 35496, 36865, 36868, 36880, 36928, 36950, 36996, 37120, 37154, 37220, 37462, - 37513, 37888, 37893, 37956, 37968, 37976, 38185, 38288, 38290, 38465, 38993, 39078, 39241, 39445, 39520, 40960, - 40962, 40968, 40970, 40992, 41002, 41120, 41297, 41305, 41382, 41472, 41474, 41480, 41514, 41600, 41632, 42048, - 42133, 42597, 42648, 43018, 43040, 43042, 43048, 43168, 43176, 43268, 43396, 43398, 43560, 43562, 43665, 43690, - }; - static const uint16_t kgrid_1bit_2048[NGRID_IQ1S] = { - 0, 2, 5, 8, 10, 17, 21, 32, 34, 40, 42, 69, 81, 84, 86, 101, - 128, 130, 136, 138, 149, 160, 162, 168, 170, 260, 261, 273, 276, 278, 281, 282, - 293, 321, 326, 329, 338, 341, 346, 353, 356, 358, 360, 389, 401, 404, 406, 421, - 512, 514, 520, 522, 533, 544, 546, 552, 554, 581, 593, 601, 612, 617, 640, 642, - 648, 650, 657, 661, 665, 672, 674, 680, 682, 1041, 1044, 1046, 1061, 1089, 1097, 1109, - 1114, 1124, 1125, 1169, 1177, 1189, 1281, 1284, 1285, 1286, 1301, 1304, 1306, 1321, 1344, 1349, - 1354, 1360, 1361, 1364, 1365, 1366, 1369, 1376, 1378, 1381, 1384, 1386, 1409, 1425, 1429, 1432, - 1434, 1441, 1444, 1445, 1446, 1449, 1556, 1561, 1601, 1604, 1616, 1618, 1621, 1624, 1632, 1633, - 1638, 1641, 1669, 1681, 1684, 1689, 2048, 2050, 2056, 2058, 2069, 2080, 2082, 2088, 2090, 2117, - 2129, 2134, 2149, 2176, 2178, 2184, 2186, 2197, 2208, 2210, 2216, 2218, 2309, 2321, 2324, 2329, - 2340, 2341, 2369, 2384, 2385, 2389, 2401, 2404, 2409, 2449, 2452, 2454, 2457, 2469, 2560, 2562, - 2568, 2570, 2581, 2592, 2594, 2600, 2602, 2629, 2641, 2649, 2657, 2661, 2688, 2690, 2693, 2696, - 2698, 2709, 2720, 2722, 2728, 2730, 4112, 4113, 4116, 4121, 4132, 4133, 4161, 4164, 4176, 4181, - 4184, 4193, 4196, 4197, 4201, 4241, 4244, 4246, 4257, 4261, 4353, 4356, 4358, 4361, 4368, 4370, - 4373, 4376, 4385, 4388, 4393, 4421, 4426, 4432, 4433, 4434, 4436, 4437, 4438, 4441, 4448, 4453, - 4484, 4498, 4501, 4513, 4516, 4625, 4628, 4630, 4645, 4672, 4678, 4681, 4690, 4693, 4696, 4698, - 4708, 4710, 4741, 4753, 4756, 4758, 4773, 5121, 5126, 5129, 5140, 5141, 5144, 5145, 5153, 5158, - 5185, 5189, 5190, 5192, 5194, 5201, 5204, 5205, 5206, 5209, 5218, 5221, 5224, 5252, 5257, 5264, - 5268, 5269, 5272, 5273, 5274, 5281, 5284, 5285, 5289, 5378, 5381, 5386, 5393, 5396, 5397, 5398, - 5401, 5408, 5410, 5413, 5416, 5418, 5441, 5444, 5445, 5446, 5457, 5458, 5460, 5461, 5462, 5465, - 5466, 5473, 5476, 5477, 5478, 5481, 5504, 5506, 5508, 5509, 5512, 5514, 5520, 5521, 5524, 5525, - 5526, 5529, 5530, 5536, 5538, 5541, 5633, 5636, 5637, 5638, 5653, 5654, 5656, 5658, 5665, 5670, - 5696, 5698, 5700, 5701, 5704, 5706, 5713, 5717, 5718, 5720, 5721, 5729, 5732, 5733, 5736, 5737, - 5738, 5766, 5770, 5778, 5781, 5796, 5801, 6161, 6166, 6181, 6209, 6212, 6214, 6217, 6224, 6229, - 6232, 6234, 6240, 6241, 6244, 6246, 6249, 6277, 6289, 6292, 6309, 6416, 6418, 6421, 6426, 6433, - 6437, 6466, 6468, 6469, 6472, 6481, 6484, 6485, 6486, 6489, 6490, 6496, 6501, 6506, 6537, 6545, - 6546, 6549, 6552, 6561, 6566, 6569, 6665, 6678, 6692, 6694, 6724, 6726, 6729, 6736, 6738, 6741, - 6744, 6753, 6758, 6761, 6789, 6801, 6806, 6810, 8192, 8194, 8200, 8202, 8213, 8224, 8226, 8229, - 8232, 8234, 8261, 8273, 8281, 8289, 8293, 8320, 8322, 8328, 8330, 8341, 8352, 8354, 8357, 8360, - 8362, 8453, 8465, 8468, 8473, 8485, 8514, 8516, 8521, 8533, 8536, 8538, 8545, 8548, 8549, 8550, - 8581, 8592, 8598, 8601, 8613, 8705, 8712, 8714, 8721, 8725, 8736, 8738, 8744, 8746, 8773, 8785, - 8790, 8793, 8805, 8833, 8840, 8842, 8849, 8853, 8864, 8866, 8872, 8874, 9221, 9236, 9238, 9241, - 9253, 9284, 9285, 9286, 9289, 9298, 9301, 9304, 9306, 9318, 9349, 9361, 9364, 9369, 9377, 9381, - 9481, 9493, 9505, 9513, 9536, 9541, 9544, 9553, 9556, 9557, 9561, 9570, 9573, 9576, 9609, 9616, - 9620, 9621, 9624, 9626, 9633, 9636, 9638, 9641, 9733, 9744, 9746, 9753, 9765, 9793, 9801, 9813, - 9824, 9825, 9833, 9860, 9862, 9872, 9882, 10240, 10242, 10248, 10250, 10261, 10272, 10274, 10280, 10282, - 10309, 10321, 10324, 10341, 10368, 10370, 10376, 10378, 10400, 10402, 10408, 10410, 10505, 10513, 10516, 10521, - 10533, 10566, 10569, 10578, 10581, 10593, 10596, 10598, 10601, 10629, 10640, 10646, 10649, 10660, 10661, 10752, - 10754, 10760, 10762, 10784, 10786, 10792, 10794, 10821, 10833, 10838, 10841, 10853, 10880, 10882, 10888, 10890, - 10901, 10912, 10914, 10920, 10922, 16389, 16401, 16406, 16421, 16457, 16466, 16469, 16472, 16474, 16481, 16484, - 16486, 16532, 16537, 16545, 16550, 16640, 16641, 16644, 16646, 16649, 16658, 16661, 16662, 16664, 16666, 16673, - 16678, 16681, 16709, 16712, 16714, 16721, 16724, 16725, 16726, 16729, 16730, 16741, 16744, 16746, 16769, 16772, - 16774, 16784, 16786, 16789, 16800, 16801, 16802, 16901, 16913, 16916, 16918, 16933, 16961, 16978, 16981, 16986, - 16996, 17001, 17033, 17044, 17061, 17409, 17429, 17433, 17449, 17477, 17480, 17482, 17489, 17492, 17493, 17494, - 17505, 17506, 17509, 17512, 17514, 17537, 17542, 17545, 17552, 17554, 17557, 17568, 17569, 17577, 17665, 17666, - 17669, 17674, 17681, 17684, 17685, 17686, 17689, 17696, 17701, 17706, 17729, 17732, 17733, 17734, 17737, 17744, - 17745, 17748, 17749, 17750, 17752, 17753, 17761, 17764, 17765, 17766, 17769, 17794, 17796, 17797, 17800, 17809, - 17812, 17813, 17814, 17817, 17818, 17829, 17832, 17834, 17921, 17925, 17929, 17940, 17941, 17944, 17946, 17953, - 17956, 17961, 17984, 17986, 17989, 17992, 18000, 18001, 18002, 18005, 18006, 18009, 18018, 18021, 18024, 18049, - 18053, 18058, 18068, 18069, 18081, 18084, 18086, 18437, 18449, 18453, 18458, 18469, 18498, 18505, 18512, 18517, - 18520, 18529, 18532, 18534, 18537, 18565, 18577, 18580, 18582, 18585, 18597, 18689, 18693, 18694, 18698, 18704, - 18708, 18709, 18712, 18721, 18724, 18726, 18752, 18757, 18762, 18769, 18770, 18772, 18773, 18774, 18777, 18784, - 18786, 18789, 18790, 18794, 18822, 18825, 18834, 18837, 18838, 18840, 18849, 18852, 18854, 18857, 18966, 19012, - 19014, 19017, 19029, 19032, 19034, 19044, 19049, 19092, 19109, 20481, 20484, 20485, 20486, 20489, 20498, 20501, - 20506, 20513, 20516, 20521, 20544, 20549, 20552, 20561, 20564, 20565, 20566, 20569, 20581, 20584, 20614, 20617, - 20629, 20632, 20640, 20641, 20646, 20649, 20741, 20744, 20745, 20746, 20753, 20756, 20757, 20758, 20760, 20761, - 20768, 20773, 20774, 20776, 20778, 20801, 20804, 20805, 20806, 20809, 20816, 20817, 20818, 20820, 20821, 20822, - 20824, 20825, 20826, 20833, 20836, 20837, 20838, 20841, 20866, 20869, 20881, 20884, 20885, 20886, 20889, 20896, - 20901, 20906, 20993, 20998, 21010, 21013, 21018, 21025, 21028, 21058, 21061, 21066, 21073, 21076, 21077, 21078, - 21081, 21090, 21093, 21125, 21136, 21138, 21141, 21145, 21146, 21156, 21508, 21509, 21521, 21524, 21525, 21526, - 21528, 21529, 21537, 21541, 21544, 21546, 21569, 21572, 21573, 21574, 21577, 21578, 21584, 21585, 21588, 21589, - 21590, 21592, 21593, 21594, 21601, 21602, 21604, 21605, 21606, 21609, 21632, 21640, 21642, 21649, 21652, 21653, - 21654, 21657, 21665, 21668, 21669, 21674, 21761, 21762, 21764, 21765, 21766, 21769, 21776, 21777, 21778, 21780, - 21781, 21782, 21785, 21786, 21793, 21796, 21797, 21798, 21801, 21824, 21825, 21826, 21828, 21829, 21830, 21832, - 21833, 21840, 21841, 21842, 21844, 21845, 21846, 21848, 21849, 21850, 21856, 21857, 21860, 21861, 21862, 21864, - 21865, 21866, 21889, 21892, 21893, 21897, 21898, 21904, 21905, 21908, 21909, 21910, 21912, 21913, 21921, 21924, - 21925, 21926, 21929, 22016, 22017, 22018, 22020, 22022, 22024, 22025, 22033, 22036, 22037, 22040, 22041, 22048, - 22049, 22050, 22052, 22053, 22054, 22056, 22057, 22081, 22085, 22086, 22088, 22089, 22090, 22096, 22097, 22098, - 22100, 22101, 22102, 22104, 22105, 22106, 22113, 22116, 22117, 22121, 22146, 22149, 22150, 22152, 22153, 22154, - 22161, 22165, 22170, 22178, 22181, 22182, 22184, 22185, 22532, 22533, 22534, 22537, 22544, 22549, 22552, 22561, - 22570, 22597, 22600, 22602, 22609, 22612, 22613, 22614, 22616, 22617, 22624, 22626, 22628, 22629, 22658, 22665, - 22672, 22674, 22677, 22680, 22689, 22697, 22785, 22786, 22789, 22794, 22801, 22804, 22805, 22806, 22809, 22821, - 22849, 22852, 22853, 22854, 22857, 22864, 22865, 22866, 22868, 22869, 22870, 22872, 22873, 22874, 22881, 22884, - 22885, 22886, 22889, 22913, 22917, 22921, 22929, 22932, 22933, 22934, 22936, 22937, 22949, 23044, 23048, 23061, - 23066, 23072, 23077, 23078, 23081, 23109, 23112, 23113, 23121, 23125, 23126, 23128, 23129, 23138, 23141, 23144, - 23146, 23169, 23178, 23186, 23189, 23190, 23192, 23194, 23201, 24581, 24596, 24598, 24601, 24613, 24644, 24656, - 24661, 24662, 24664, 24666, 24673, 24676, 24678, 24681, 24705, 24726, 24741, 24833, 24836, 24838, 24841, 24850, - 24853, 24865, 24866, 24870, 24873, 24901, 24905, 24913, 24917, 24918, 24921, 24933, 24934, 24938, 24964, 24970, - 24978, 24981, 24993, 24998, 25001, 25105, 25110, 25113, 25152, 25153, 25158, 25173, 25174, 25176, 25184, 25221, - 25233, 25238, 25253, 25617, 25618, 25621, 25622, 25626, 25633, 25638, 25641, 25664, 25666, 25669, 25672, 25674, - 25681, 25684, 25685, 25686, 25689, 25690, 25696, 25698, 25701, 25732, 25733, 25737, 25744, 25746, 25748, 25749, - 25750, 25752, 25754, 25761, 25764, 25769, 25861, 25864, 25866, 25873, 25877, 25878, 25881, 25924, 25925, 25926, - 25929, 25936, 25937, 25940, 25941, 25942, 25945, 25953, 25956, 25957, 25958, 25961, 25990, 25993, 25994, 26001, - 26005, 26006, 26009, 26010, 26018, 26021, 26022, 26024, 26114, 26121, 26133, 26144, 26150, 26152, 26153, 26176, - 26181, 26184, 26186, 26193, 26196, 26197, 26198, 26200, 26202, 26208, 26213, 26216, 26240, 26242, 26245, 26250, - 26260, 26262, 26264, 26265, 26272, 26276, 26278, 26282, 26646, 26649, 26661, 26689, 26706, 26709, 26714, 26721, - 26729, 26757, 26769, 26776, 26790, 26881, 26884, 26896, 26901, 26913, 26916, 26918, 26921, 26944, 26945, 26949, - 26950, 26952, 26961, 26964, 26965, 26966, 26969, 26976, 26981, 26986, 27010, 27012, 27018, 27029, 27041, 27044, - 27045, 27049, 27153, 27158, 27160, 27201, 27204, 27209, 27216, 27221, 27224, 27226, 27236, 27237, 27241, 27270, - 27284, 27288, 27290, 27302, 32768, 32770, 32776, 32778, 32800, 32802, 32808, 32810, 32837, 32848, 32849, 32852, - 32854, 32857, 32869, 32896, 32898, 32904, 32906, 32917, 32928, 32930, 32936, 32938, 33029, 33041, 33044, 33046, - 33049, 33061, 33089, 33092, 33097, 33104, 33106, 33109, 33110, 33112, 33113, 33124, 33126, 33129, 33157, 33161, - 33172, 33174, 33177, 33189, 33280, 33282, 33288, 33290, 33301, 33312, 33314, 33320, 33322, 33361, 33364, 33369, - 33381, 33408, 33410, 33416, 33418, 33429, 33440, 33442, 33448, 33450, 33812, 33817, 33857, 33860, 33873, 33877, - 33882, 33889, 33892, 33897, 33940, 33945, 34049, 34057, 34066, 34069, 34074, 34086, 34089, 34112, 34113, 34117, - 34120, 34129, 34132, 34133, 34134, 34137, 34138, 34149, 34150, 34152, 34154, 34177, 34180, 34182, 34185, 34192, - 34194, 34197, 34200, 34214, 34321, 34326, 34329, 34341, 34369, 34372, 34377, 34378, 34384, 34389, 34393, 34394, - 34401, 34406, 34410, 34437, 34449, 34458, 34468, 34816, 34818, 34824, 34826, 34837, 34848, 34850, 34856, 34858, - 34881, 34885, 34897, 34900, 34905, 34917, 34921, 34944, 34946, 34952, 34954, 34965, 34976, 34978, 34984, 34986, - 35077, 35078, 35089, 35092, 35094, 35109, 35137, 35140, 35142, 35145, 35152, 35154, 35157, 35162, 35169, 35172, - 35205, 35222, 35225, 35237, 35328, 35330, 35336, 35338, 35349, 35360, 35362, 35368, 35370, 35397, 35409, 35412, - 35414, 35456, 35458, 35464, 35466, 35477, 35488, 35490, 35496, 35498, 36869, 36881, 36886, 36888, 36889, 36901, - 36929, 36934, 36937, 36949, 36952, 36954, 36969, 36970, 36997, 37009, 37012, 37014, 37017, 37029, 37121, 37124, - 37126, 37129, 37136, 37141, 37144, 37146, 37153, 37156, 37158, 37161, 37184, 37189, 37200, 37201, 37204, 37205, - 37206, 37209, 37218, 37221, 37252, 37254, 37266, 37269, 37272, 37281, 37284, 37286, 37289, 37381, 37393, 37396, - 37401, 37413, 37444, 37446, 37449, 37456, 37458, 37461, 37464, 37478, 37481, 37509, 37524, 37526, 37545, 37889, - 37892, 37894, 37904, 37909, 37912, 37926, 37952, 37962, 37969, 37972, 37973, 37974, 37976, 37977, 37984, 37985, - 37986, 37989, 38020, 38022, 38034, 38036, 38037, 38040, 38049, 38057, 38144, 38149, 38152, 38154, 38160, 38161, - 38164, 38165, 38166, 38169, 38177, 38181, 38185, 38186, 38209, 38212, 38213, 38214, 38217, 38224, 38225, 38226, - 38228, 38229, 38230, 38232, 38233, 38234, 38241, 38244, 38245, 38246, 38249, 38273, 38277, 38280, 38289, 38290, - 38292, 38293, 38294, 38297, 38298, 38304, 38306, 38309, 38312, 38314, 38401, 38404, 38416, 38421, 38425, 38432, - 38438, 38441, 38469, 38472, 38473, 38481, 38482, 38485, 38486, 38489, 38501, 38504, 38530, 38532, 38537, 38538, - 38546, 38548, 38549, 38564, 38566, 38569, 38917, 38934, 38937, 38949, 38977, 38982, 38992, 38994, 38997, 38998, - 39002, 39012, 39013, 39045, 39057, 39062, 39065, 39077, 39172, 39174, 39177, 39184, 39186, 39189, 39192, 39194, - 39200, 39201, 39204, 39206, 39232, 39234, 39237, 39240, 39242, 39249, 39252, 39253, 39254, 39257, 39266, 39269, - 39270, 39274, 39297, 39300, 39312, 39314, 39317, 39322, 39329, 39334, 39429, 39445, 39461, 39492, 39494, 39497, - 39504, 39509, 39512, 39521, 39557, 39569, 39572, 39573, 39574, 40960, 40962, 40968, 40970, 40981, 40992, 40994, - 41000, 41002, 41029, 41041, 41044, 41046, 41049, 41088, 41090, 41096, 41098, 41109, 41120, 41122, 41128, 41130, - 41221, 41225, 41233, 41236, 41238, 41241, 41242, 41286, 41289, 41297, 41301, 41304, 41306, 41313, 41316, 41349, - 41360, 41362, 41366, 41369, 41474, 41480, 41482, 41488, 41497, 41506, 41512, 41514, 41541, 41553, 41558, 41561, - 41573, 41600, 41602, 41608, 41610, 41621, 41632, 41634, 41640, 41642, 42009, 42021, 42049, 42052, 42064, 42068, - 42069, 42072, 42074, 42081, 42085, 42086, 42088, 42089, 42117, 42246, 42249, 42256, 42258, 42261, 42264, 42278, - 42281, 42306, 42309, 42321, 42324, 42325, 42326, 42329, 42341, 42346, 42369, 42372, 42373, 42374, 42377, 42386, - 42389, 42392, 42501, 42513, 42518, 42522, 42529, 42533, 42564, 42566, 42570, 42578, 42581, 42582, 42584, 42592, - 42594, 42630, 42640, 42645, 42646, 42649, 42657, 42660, 42662, 43008, 43010, 43016, 43018, 43040, 43042, 43048, - 43050, 43089, 43092, 43094, 43097, 43136, 43138, 43144, 43146, 43157, 43168, 43170, 43176, 43178, 43269, 43284, - 43289, 43297, 43301, 43329, 43344, 43349, 43354, 43361, 43366, 43369, 43408, 43414, 43520, 43522, 43528, 43530, - 43552, 43554, 43560, 43562, 43601, 43604, 43606, 43648, 43650, 43656, 43658, 43669, 43680, 43682, 43688, 43690, - }; - static const uint16_t kgrid_2bit_1024[1024] = { - 0, 2, 5, 8, 10, 17, 20, 22, 25, 32, 34, 37, 40, 65, 68, 70, - 73, 80, 82, 85, 88, 97, 100, 102, 105, 128, 130, 133, 136, 145, 148, 160, - 165, 170, 257, 260, 262, 265, 272, 274, 277, 280, 289, 292, 320, 322, 325, 328, - 337, 340, 342, 345, 352, 357, 360, 385, 388, 400, 402, 405, 417, 420, 512, 514, - 517, 520, 529, 532, 544, 554, 577, 580, 582, 585, 592, 597, 640, 645, 650, 660, - 674, 1025, 1028, 1030, 1033, 1040, 1042, 1045, 1048, 1057, 1060, 1062, 1065, 1088, 1090, 1093, - 1096, 1098, 1105, 1108, 1110, 1113, 1120, 1122, 1125, 1153, 1156, 1158, 1161, 1168, 1173, 1176, - 1185, 1188, 1280, 1282, 1285, 1288, 1290, 1297, 1300, 1302, 1305, 1312, 1317, 1320, 1345, 1348, - 1350, 1353, 1360, 1362, 1365, 1368, 1377, 1380, 1408, 1410, 1413, 1416, 1425, 1428, 1440, 1537, - 1540, 1542, 1545, 1552, 1557, 1600, 1605, 1608, 1617, 1620, 1632, 1665, 1668, 1680, 2048, 2050, - 2053, 2056, 2065, 2068, 2070, 2073, 2080, 2085, 2090, 2113, 2116, 2118, 2121, 2128, 2130, 2133, - 2136, 2145, 2148, 2176, 2181, 2196, 2218, 2305, 2308, 2320, 2322, 2325, 2328, 2337, 2368, 2373, - 2376, 2385, 2388, 2400, 2433, 2448, 2560, 2577, 2580, 2594, 2600, 2602, 2640, 2713, 4097, 4100, - 4102, 4105, 4112, 4114, 4117, 4120, 4129, 4132, 4134, 4160, 4162, 4165, 4168, 4177, 4180, 4182, - 4185, 4192, 4194, 4197, 4200, 4225, 4228, 4230, 4240, 4245, 4248, 4257, 4260, 4352, 4354, 4357, - 4360, 4362, 4369, 4372, 4374, 4377, 4384, 4386, 4389, 4392, 4417, 4420, 4422, 4425, 4432, 4434, - 4437, 4440, 4449, 4452, 4480, 4482, 4485, 4488, 4497, 4500, 4609, 4612, 4617, 4624, 4629, 4641, - 4644, 4672, 4677, 4689, 4692, 4737, 4740, 4752, 5120, 5122, 5125, 5128, 5137, 5140, 5142, 5145, - 5152, 5157, 5160, 5185, 5188, 5190, 5193, 5200, 5202, 5205, 5208, 5217, 5220, 5248, 5250, 5253, - 5256, 5265, 5268, 5280, 5377, 5380, 5382, 5385, 5392, 5394, 5397, 5400, 5409, 5412, 5440, 5442, - 5445, 5448, 5457, 5460, 5472, 5505, 5508, 5520, 5632, 5637, 5640, 5649, 5652, 5664, 5697, 5700, - 5712, 5760, 5802, 6145, 6148, 6150, 6153, 6160, 6165, 6168, 6177, 6208, 6210, 6213, 6216, 6225, - 6228, 6240, 6273, 6276, 6400, 6402, 6405, 6408, 6417, 6420, 6432, 6465, 6468, 6480, 6505, 6562, - 6660, 6672, 6720, 6742, 8192, 8194, 8197, 8200, 8209, 8212, 8214, 8217, 8224, 8229, 8234, 8257, - 8260, 8272, 8274, 8277, 8292, 8320, 8330, 8340, 8362, 8449, 8452, 8464, 8466, 8469, 8481, 8512, - 8514, 8517, 8529, 8532, 8544, 8577, 8580, 8592, 8704, 8714, 8738, 8744, 8746, 8772, 8784, 8840, - 8842, 8872, 9217, 9220, 9222, 9225, 9232, 9237, 9240, 9249, 9252, 9280, 9282, 9285, 9288, 9297, - 9300, 9312, 9345, 9348, 9360, 9472, 9477, 9480, 9489, 9492, 9504, 9537, 9540, 9552, 9574, 9600, - 9729, 9732, 9744, 9792, 9817, 10240, 10245, 10257, 10260, 10305, 10308, 10320, 10378, 10410, 10497, 10500, - 10512, 10645, 10762, 10786, 10852, 10888, 10890, 16385, 16388, 16390, 16393, 16400, 16402, 16405, 16408, 16410, - 16417, 16420, 16422, 16448, 16450, 16453, 16456, 16458, 16465, 16468, 16470, 16473, 16480, 16482, 16485, 16513, - 16516, 16528, 16533, 16536, 16545, 16548, 16640, 16642, 16645, 16648, 16657, 16660, 16662, 16665, 16672, 16674, - 16677, 16705, 16708, 16710, 16713, 16720, 16722, 16725, 16728, 16737, 16740, 16768, 16770, 16773, 16776, 16785, - 16788, 16800, 16897, 16900, 16912, 16914, 16917, 16920, 16932, 16960, 16965, 16968, 16977, 16980, 16992, 17025, - 17028, 17408, 17410, 17413, 17416, 17418, 17425, 17428, 17430, 17433, 17440, 17442, 17445, 17448, 17473, 17476, - 17478, 17481, 17488, 17490, 17493, 17496, 17505, 17508, 17536, 17538, 17541, 17544, 17553, 17556, 17568, 17665, - 17668, 17670, 17673, 17680, 17682, 17685, 17688, 17697, 17700, 17728, 17730, 17733, 17736, 17745, 17748, 17760, - 17770, 17793, 17796, 17808, 17920, 17922, 17925, 17928, 17937, 17940, 17952, 17985, 17988, 18000, 18048, 18085, - 18433, 18436, 18441, 18448, 18450, 18453, 18456, 18465, 18468, 18496, 18498, 18501, 18504, 18513, 18516, 18528, - 18564, 18576, 18688, 18690, 18693, 18696, 18705, 18708, 18720, 18753, 18756, 18768, 18816, 18838, 18945, 18948, - 18960, 19008, 20480, 20482, 20485, 20488, 20497, 20500, 20502, 20505, 20512, 20514, 20517, 20520, 20545, 20548, - 20550, 20553, 20560, 20562, 20565, 20568, 20577, 20580, 20608, 20610, 20613, 20616, 20625, 20628, 20737, 20740, - 20742, 20745, 20752, 20754, 20757, 20760, 20769, 20772, 20800, 20802, 20805, 20808, 20817, 20820, 20832, 20865, - 20868, 20880, 20992, 20997, 21000, 21009, 21012, 21024, 21057, 21060, 21072, 21097, 21120, 21505, 21508, 21510, - 21513, 21520, 21522, 21525, 21528, 21537, 21540, 21568, 21570, 21573, 21576, 21585, 21588, 21600, 21633, 21636, - 21648, 21760, 21762, 21765, 21768, 21777, 21780, 21792, 21825, 21828, 21840, 21888, 22017, 22020, 22032, 22054, - 22080, 22528, 22530, 22533, 22536, 22545, 22548, 22560, 22593, 22596, 22608, 22618, 22656, 22785, 22788, 22800, - 22848, 23040, 23065, 23173, 23208, 24577, 24580, 24582, 24592, 24594, 24597, 24600, 24609, 24612, 24640, 24645, - 24648, 24657, 24660, 24672, 24708, 24720, 24832, 24834, 24837, 24840, 24849, 24852, 24864, 24897, 24900, 24912, - 24960, 24985, 25092, 25104, 25152, 25174, 25249, 25600, 25605, 25608, 25617, 25620, 25632, 25665, 25668, 25680, - 25728, 25857, 25860, 25872, 25920, 25930, 25960, 26002, 26112, 26260, 26625, 26628, 26640, 26725, 26776, 26880, - 26922, 27202, 27297, 32768, 32770, 32773, 32776, 32785, 32788, 32793, 32800, 32805, 32833, 32836, 32848, 32850, - 32853, 32856, 32865, 32896, 32901, 32913, 32916, 33025, 33028, 33033, 33040, 33042, 33045, 33048, 33057, 33060, - 33088, 33090, 33093, 33096, 33105, 33108, 33153, 33156, 33168, 33193, 33280, 33285, 33290, 33297, 33300, 33345, - 33348, 33360, 33793, 33796, 33798, 33801, 33808, 33810, 33813, 33816, 33825, 33856, 33858, 33861, 33864, 33873, - 33876, 33888, 33921, 33924, 33936, 34048, 34050, 34053, 34056, 34065, 34068, 34080, 34113, 34116, 34128, 34176, - 34186, 34305, 34308, 34320, 34345, 34368, 34816, 34821, 34833, 34836, 34881, 34884, 34896, 34978, 35073, 35076, - 35136, 35173, 35362, 35416, 35418, 35458, 35490, 36865, 36868, 36873, 36880, 36882, 36885, 36888, 36900, 36928, - 36930, 36933, 36936, 36945, 36948, 36960, 36993, 36996, 37008, 37120, 37125, 37137, 37140, 37185, 37188, 37200, - 37210, 37377, 37380, 37392, 37440, 37542, 37888, 37890, 37893, 37896, 37905, 37908, 37920, 37953, 37956, 37968, - 38016, 38038, 38145, 38148, 38160, 38208, 38296, 38305, 38400, 38470, 38500, 38913, 38916, 38928, 38950, 38976, - 39081, 39168, 39241, 39250, 39568, 40960, 40965, 40970, 40980, 40994, 41002, 41025, 41028, 41040, 41122, 41130, - 41280, 41317, 41474, 41482, 41506, 41512, 41514, 41602, 41608, 41610, 41640, 41985, 41988, 42000, 42048, 42121, - 42148, 42240, 42265, 42577, 43018, 43048, 43170, 43348, 43398, 43528, 43530, 43552, 43554, 43560, 43656, 43690, - }; - - const int kmap_size = 43692; - //const int nwant = type == GGML_TYPE_IQ1_S ? 3 : 2; - const int nwant = type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? 3 : type == GGML_TYPE_IQ2_S ? 1 : 2; - const uint16_t * kgrid = type == GGML_TYPE_IQ2_XXS ? kgrid_2bit_256 : - type == GGML_TYPE_IQ2_XS ? kgrid_2bit_512 : - type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? kgrid_1bit_2048 : kgrid_2bit_1024; - uint64_t * kgrid_q2xs; - int * kmap_q2xs; - uint16_t * kneighbors_q2xs; - - //printf("================================================================= %s(grid_size = %d)\n", __func__, grid_size); - uint64_t * the_grid = (uint64_t *)malloc(grid_size*sizeof(uint64_t)); - for (int k = 0; k < grid_size; ++k) { - int8_t * pos = (int8_t *)(the_grid + k); - for (int i = 0; i < 8; ++i) { - int l = (kgrid[k] >> 2*i) & 0x3; - pos[i] = 2*l + 1; - } - } - kgrid_q2xs = the_grid; - iq2_data[gindex].grid = the_grid; - kmap_q2xs = (int *)malloc(kmap_size*sizeof(int)); - iq2_data[gindex].map = kmap_q2xs; - for (int i = 0; i < kmap_size; ++i) kmap_q2xs[i] = -1; - uint64_t aux64; - uint8_t * aux8 = (uint8_t *)&aux64; - for (int i = 0; i < grid_size; ++i) { - aux64 = kgrid_q2xs[i]; - uint16_t index = 0; - for (int k=0; k<8; ++k) { - uint16_t q = (aux8[k] - 1)/2; - index |= (q << 2*k); - } - kmap_q2xs[index] = i; - } - int8_t pos[8]; - int * dist2 = (int *)malloc(2*grid_size*sizeof(int)); - int num_neighbors = 0, num_not_in_map = 0; - for (int i = 0; i < kmap_size; ++i) { - if (kmap_q2xs[i] >= 0) continue; - ++num_not_in_map; - for (int k = 0; k < 8; ++k) { - int l = (i >> 2*k) & 0x3; - pos[k] = 2*l + 1; - } - for (int j = 0; j < grid_size; ++j) { - const int8_t * pg = (const int8_t *)(kgrid_q2xs + j); - int d2 = 0; - for (int k = 0; k < 8; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]); - dist2[2*j+0] = d2; - dist2[2*j+1] = j; - } - qsort(dist2, grid_size, 2*sizeof(int), iq2_compare_func); - int n = 0; int d2 = dist2[0]; - int nhave = 1; - for (int j = 0; j < grid_size; ++j) { - if (dist2[2*j] > d2) { - if (nhave == nwant) break; - d2 = dist2[2*j]; - ++nhave; - } - ++n; - } - num_neighbors += n; - } - //printf("%s: %d neighbours in total\n", __func__, num_neighbors); - kneighbors_q2xs = (uint16_t *)malloc((num_neighbors + num_not_in_map)*sizeof(uint16_t)); - iq2_data[gindex].neighbours = kneighbors_q2xs; - int counter = 0; - for (int i = 0; i < kmap_size; ++i) { - if (kmap_q2xs[i] >= 0) continue; - for (int k = 0; k < 8; ++k) { - int l = (i >> 2*k) & 0x3; - pos[k] = 2*l + 1; - } - for (int j = 0; j < grid_size; ++j) { - const int8_t * pg = (const int8_t *)(kgrid_q2xs + j); - int d2 = 0; - for (int k = 0; k < 8; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]); - dist2[2*j+0] = d2; - dist2[2*j+1] = j; - } - qsort(dist2, grid_size, 2*sizeof(int), iq2_compare_func); - kmap_q2xs[i] = -(counter + 1); - int d2 = dist2[0]; - uint16_t * start = &kneighbors_q2xs[counter++]; - int n = 0, nhave = 1; - for (int j = 0; j < grid_size; ++j) { - if (dist2[2*j] > d2) { - if (nhave == nwant) break; - d2 = dist2[2*j]; - ++nhave; - } - kneighbors_q2xs[counter++] = dist2[2*j+1]; - ++n; - } - *start = n; - } - free(dist2); -} - -void iq2xs_free_impl(enum ggml_type type) { - GGML_ASSERT(type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M || type == GGML_TYPE_IQ2_S); - const int gindex = iq2_data_index(type); - if (iq2_data[gindex].grid) { - free(iq2_data[gindex].grid); iq2_data[gindex].grid = NULL; - free(iq2_data[gindex].map); iq2_data[gindex].map = NULL; - free(iq2_data[gindex].neighbours); iq2_data[gindex].neighbours = NULL; - } -} - -static int iq2_find_best_neighbour(const uint16_t * GGML_RESTRICT neighbours, const uint64_t * GGML_RESTRICT grid, - const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float scale, int8_t * GGML_RESTRICT L) { - int num_neighbors = neighbours[0]; - GGML_ASSERT(num_neighbors > 0); - float best_d2 = FLT_MAX; - int grid_index = -1; - for (int j = 1; j <= num_neighbors; ++j) { - const int8_t * pg = (const int8_t *)(grid + neighbours[j]); - float d2 = 0; - for (int i = 0; i < 8; ++i) { - float q = pg[i]; - float diff = scale*q - xval[i]; - d2 += weight[i]*diff*diff; - } - if (d2 < best_d2) { - best_d2 = d2; grid_index = neighbours[j]; - } - } - GGML_ASSERT(grid_index >= 0); - const int8_t * pg = (const int8_t *)(grid + grid_index); - for (int i = 0; i < 8; ++i) L[i] = (pg[i] - 1)/2; - return grid_index; -} - -static void quantize_row_iq2_xxs_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights) { - - const int gindex = iq2_data_index(GGML_TYPE_IQ2_XXS); - - const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; - const int * kmap_q2xs = iq2_data[gindex].map; - const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; - - GGML_ASSERT(quant_weights && "missing quantization weights"); - GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(n%QK_K == 0); - - const int kMaxQ = 3; - - const int64_t nbl = n/QK_K; - - block_iq2_xxs * y = vy; - - float scales[QK_K/32]; - float weight[32]; - float xval[32]; - int8_t L[32]; - int8_t Laux[32]; - float waux[32]; - uint8_t block_signs[4]; - uint32_t q2[2*(QK_K/32)]; - - for (int ibl = 0; ibl < nbl; ++ibl) { - - y[ibl].d = GGML_FP32_TO_FP16(0.f); - memset(q2, 0, QK_K/4); - - float max_scale = 0; - - const float * xbl = x + QK_K*ibl; - float sumx2 = 0; - for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; - float sigma2 = sumx2/QK_K; - - for (int ib = 0; ib < QK_K/32; ++ib) { - const float * xb = xbl + 32*ib; - const float * qw = quant_weights + QK_K*ibl + 32*ib; - for (int i = 0; i < 32; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); - for (int i = 0; i < 32; ++i) waux[i] = sqrtf(weight[i]); - for (int k = 0; k < 4; ++k) { - int nflip = 0; - uint8_t s = 0; - for (int i = 0; i < 8; ++i) { - if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; - else { - xval[8*k + i] = -xb[8*k + i]; ++nflip; s |= (1 << i); - } - } - if (nflip%2) { - int imin = 0; float min = weight[8*k+imin]*xb[8*k+imin]*xb[8*k+imin]; - for (int i = 1; i < 8; ++i) { - float ax = weight[8*k+i]*xb[8*k+i]*xb[8*k+i]; - if (ax < min) { - min = ax; imin = i; - } - } - xval[8*k+imin] = -xval[8*k+imin]; - s ^= (1 << imin); - } - block_signs[k] = s & 127; - } - float max = xval[0]; - for (int i = 1; i < 32; ++i) max = MAX(max, xval[i]); - if (max < GROUP_MAX_EPS) { - scales[ib] = 0; - memset(L, 0, 32); - continue; - } - float scale = make_qp_quants(32, kMaxQ+1, xval, (uint8_t*)L, weight); - float eff_max = scale*kMaxQ; - if (eff_max <= 0) { - scales[ib] = 0; - memset(L, 0, 32); - continue; - } - float best = 0; - for (int is = -6; is <= 6; ++is) { - float id = (2*kMaxQ-1+is*0.1f)/eff_max; - float this_scale = 1/id; - for (int k = 0; k < 4; ++k) { - for (int i = 0; i < 8; ++i) { - int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); - Laux[8*k+i] = MAX(0, MIN(kMaxQ-1, l)); - } - uint16_t u = 0; - for (int i = 0; i < 8; ++i) u |= (Laux[8*k+i] << 2*i); - int grid_index = kmap_q2xs[u]; - if (grid_index < 0) { - const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; - grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, this_scale, Laux + 8*k); - } - } - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < 32; ++i) { - float w = weight[i]; - float q = 2*Laux[i] + 1; - sumqx += w*xval[i]*q; - sumq2 += w*q*q; - } - if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { - scale = sumqx/sumq2; best = scale*sumqx; - memcpy(L, Laux, 32); - } - } - if (scale > 0) { - float id = 1/scale; - for (int k = 0; k < 4; ++k) { - uint16_t u = 0; - for (int i = 0; i < 8; ++i) { - int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); - l = MAX(0, MIN(kMaxQ-1, l)); - u |= (l << 2*i); - } - int grid_index = kmap_q2xs[u]; - if (grid_index < 0) { - const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; - grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, scale, L + 8*k); - } - const int8_t * pg = (const int8_t *)(kgrid_q2xs + grid_index); - for (int i = 0; i < 8; ++i) L[8*k+i] = (pg[i] - 1)/2; - } - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < 32; ++i) { - float w = weight[i]; - float q = 2*L[i] + 1; - sumqx += w*xval[i]*q; - sumq2 += w*q*q; - } - if (sumq2 > 0) scale = sumqx/sumq2; - } - if (scale < 0) { - // This should never happen, but just in case, flip scale so that it is positive (we use uint's to encode the scale) - // and correspondingly flip quant signs. - scale = -scale; - for (int k = 0; k < 4; ++k) block_signs[k] = (~block_signs[k]) & 127; - } - for (int k = 0; k < 4; ++k) { - uint16_t u = 0; - for (int i = 0; i < 8; ++i) u |= (L[8*k+i] << 2*i); - int grid_index = kmap_q2xs[u]; - if (grid_index < 0) { - printf("Oops: found point %u not on grid:", u); - for (int i = 0; i < 8; ++i) printf(" %d", L[8*k+i]); - printf("\n"); - GGML_ABORT("fatal error"); - } - q2[2*ib+0] |= ((uint32_t) grid_index << 8*k); - q2[2*ib+1] |= (block_signs[k] << 7*k); - } - GGML_ASSERT(scale >= 0); - scales[ib] = scale; - max_scale = MAX(max_scale, scale); - } - - if (!max_scale) { - memset(y[ibl].qs, 0, QK_K/4); - continue; - } - - float d = max_scale/31; - y[ibl].d = GGML_FP32_TO_FP16(d); - float id = 1/d; - for (int ib = 0; ib < QK_K/32; ++ib) { - int l = nearest_int(0.5f*(id*scales[ib]-1)); - l = MAX(0, MIN(15, l)); - q2[2*ib+1] |= ((uint32_t)l << 28); - } - memcpy(y[ibl].qs, q2, QK_K/4); - } -} - -static void quantize_row_iq2_xs_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights) { - - const int gindex = iq2_data_index(GGML_TYPE_IQ2_XS); - - const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; - const int * kmap_q2xs = iq2_data[gindex].map; - const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; - - GGML_ASSERT(quant_weights && "missing quantization weights"); - GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(n%QK_K == 0); - - const int kMaxQ = 3; - - const int64_t nbl = n/QK_K; - - block_iq2_xs * y = vy; - - float scales[QK_K/16]; - float weight[16]; - float xval[16]; - int8_t L[16]; - int8_t Laux[16]; - float waux[16]; - bool is_on_grid[2]; - bool is_on_grid_aux[2]; - uint8_t block_signs[2]; - uint16_t q2[2*(QK_K/16)]; - - for (int ibl = 0; ibl < nbl; ++ibl) { - - y[ibl].d = GGML_FP32_TO_FP16(0.f); - memset(q2, 0, QK_K/4); - memset(y[ibl].scales, 0, QK_K/32); - - float max_scale = 0; - - const float * xbl = x + QK_K*ibl; - float sumx2 = 0; - for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; - float sigma2 = sumx2/QK_K; - - for (int ib = 0; ib < QK_K/16; ++ib) { - const float * xb = xbl + 16*ib; - const float * qw = quant_weights + QK_K*ibl + 16*ib; - for (int i = 0; i < 16; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); - for (int i = 0; i < 16; ++i) waux[i] = sqrtf(weight[i]); - for (int k = 0; k < 2; ++k) { - int nflip = 0; - uint8_t s = 0; - for (int i = 0; i < 8; ++i) { - if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; - else { - xval[8*k + i] = -xb[8*k + i]; ++nflip; s |= (1 << i); - } - } - if (nflip%2) { - int imin = 0; float min = weight[8*k+imin]*xb[8*k+imin]*xb[8*k+imin]; - for (int i = 1; i < 8; ++i) { - float ax = weight[8*k+i]*xb[8*k+i]*xb[8*k+i]; - if (ax < min) { - min = ax; imin = i; - } - } - xval[8*k+imin] = -xval[8*k+imin]; - s ^= (1 << imin); - } - block_signs[k] = s & 127; - } - float max = xval[0]; - for (int i = 1; i < 16; ++i) max = MAX(max, xval[i]); - memset(L, 0, 16); - if (max < GROUP_MAX_EPS) { - scales[ib] = 0; - continue; - } - float best = 0; - float scale = max/(2*kMaxQ-1); - is_on_grid[0] = is_on_grid[1] = true; - for (int is = -9; is <= 9; ++is) { - float id = (2*kMaxQ-1+is*0.1f)/max; - float this_scale = 1/id; - for (int k = 0; k < 2; ++k) { - for (int i = 0; i < 8; ++i) { - int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); - Laux[8*k+i] = MAX(0, MIN(kMaxQ-1, l)); - } - uint16_t u = 0; - for (int i = 0; i < 8; ++i) u |= (Laux[8*k+i] << 2*i); - int grid_index = kmap_q2xs[u]; - is_on_grid_aux[k] = true; - if (grid_index < 0) { - is_on_grid_aux[k] = false; - const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; - grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, this_scale, Laux + 8*k); - } - } - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < 16; ++i) { - float w = weight[i]; - float q = 2*Laux[i] + 1; - sumqx += w*xval[i]*q; - sumq2 += w*q*q; - } - if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { - scale = sumqx/sumq2; best = scale*sumqx; - for (int i = 0; i < 16; ++i) L[i] = Laux[i]; - for (int k = 0; k < 2; ++k) is_on_grid[k] = is_on_grid_aux[k]; - } - } - int n_not_ongrid = 0; - for (int k = 0; k < 2; ++k) if (!is_on_grid[k]) ++n_not_ongrid; - if (n_not_ongrid > 0 && scale > 0) { - float id = 1/scale; - for (int k = 0; k < 2; ++k) { - if (is_on_grid[k]) continue; - uint16_t u = 0; - for (int i = 0; i < 8; ++i) { - int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); - l = MAX(0, MIN(kMaxQ-1, l)); - u |= (l << 2*i); - L[8*k + i] = l; - } - int grid_index = kmap_q2xs[u]; - if (grid_index < 0) { - const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; - grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, scale, L + 8*k); - } - } - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < 16; ++i) { - float w = weight[i]; - float q = 2*L[i] + 1; - sumqx += w*xval[i]*q; - sumq2 += w*q*q; - } - if (sumq2 > 0) scale = sumqx/sumq2; - } - if (scale < 0) { - scale = -scale; - for (int k = 0; k < 2; ++k) block_signs[k] = (~block_signs[k]) & 127; - } - for (int k = 0; k < 2; ++k) { - uint16_t u = 0; - for (int i = 0; i < 8; ++i) u |= (L[8*k+i] << 2*i); - int grid_index = kmap_q2xs[u]; - if (grid_index < 0) { - printf("Oops: found point %u not on grid:", u); - for (int i = 0; i < 8; ++i) printf(" %d", L[8*k+i]); - printf("\n"); - GGML_ABORT("fatal error"); - } - q2[2*ib+k] = grid_index | (block_signs[k] << 9); - } - GGML_ASSERT(scale >= 0); - scales[ib] = scale; - max_scale = MAX(max_scale, scale); - } - - if (!max_scale) { - memset(y[ibl].qs, 0, QK_K/4); - continue; - } - - float d = max_scale/31; - y[ibl].d = GGML_FP32_TO_FP16(d); - float id = 1/d; - for (int ib = 0; ib < QK_K/16; ++ib) { - int l = nearest_int(0.5f*(id*scales[ib]-1)); - l = MAX(0, MIN(15, l)); - if (ib%2 == 0) y[ibl].scales[ib/2] = l; - else y[ibl].scales[ib/2] |= (l << 4); - } - memcpy(y[ibl].qs, q2, QK_K/4); - - } -} - -size_t quantize_iq2_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - GGML_ASSERT(n_per_row%QK_K == 0); - int64_t nblock = n_per_row/QK_K; - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_iq2_xxs_impl(src, qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += nblock*sizeof(block_iq2_xxs); - } - return nrow * nblock * sizeof(block_iq2_xxs); -} - -size_t quantize_iq2_xs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - GGML_ASSERT(n_per_row%QK_K == 0); - int64_t nblock = n_per_row/QK_K; - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_iq2_xs_impl(src, qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += nblock*sizeof(block_iq2_xs); - } - return nrow * nblock * sizeof(block_iq2_xs); -} - -// -// ============================================= 3-bit using D4 lattice -// - -typedef struct { - uint32_t * grid; - int * map; - uint16_t * neighbours; -} iq3_entry_t; - -static iq3_entry_t iq3_data[2] = { - {NULL, NULL, NULL}, - {NULL, NULL, NULL}, -}; - -static inline int iq3_data_index(int grid_size) { - (void)grid_size; - GGML_ASSERT(grid_size == 256 || grid_size == 512); - return grid_size == 256 ? 0 : 1; -} - -static int iq3_compare_func(const void * left, const void * right) { - const int * l = (const int *)left; - const int * r = (const int *)right; - return l[0] < r[0] ? -1 : l[0] > r[0] ? 1 : l[1] < r[1] ? -1 : l[1] > r[1] ? 1 : 0; -} - -void iq3xs_init_impl(int grid_size) { - const int gindex = iq3_data_index(grid_size); - if (iq3_data[gindex].grid) { - return; - } - static const uint16_t kgrid_256[256] = { - 0, 2, 4, 9, 11, 15, 16, 18, 25, 34, 59, 61, 65, 67, 72, 74, - 81, 85, 88, 90, 97, 108, 120, 128, 130, 132, 137, 144, 146, 153, 155, 159, - 169, 175, 189, 193, 199, 200, 202, 213, 248, 267, 287, 292, 303, 315, 317, 321, - 327, 346, 362, 413, 436, 456, 460, 462, 483, 497, 513, 515, 520, 522, 529, 531, - 536, 538, 540, 551, 552, 576, 578, 585, 592, 594, 641, 643, 648, 650, 657, 664, - 698, 704, 706, 720, 729, 742, 758, 769, 773, 808, 848, 852, 870, 889, 901, 978, - 992, 1024, 1026, 1033, 1035, 1040, 1042, 1046, 1049, 1058, 1089, 1091, 1093, 1096, 1098, 1105, - 1112, 1139, 1143, 1144, 1152, 1154, 1161, 1167, 1168, 1170, 1183, 1184, 1197, 1217, 1224, 1228, - 1272, 1276, 1309, 1323, 1347, 1367, 1377, 1404, 1473, 1475, 1486, 1509, 1537, 1544, 1546, 1553, - 1555, 1576, 1589, 1594, 1600, 1602, 1616, 1625, 1636, 1638, 1665, 1667, 1672, 1685, 1706, 1722, - 1737, 1755, 1816, 1831, 1850, 1856, 1862, 1874, 1901, 1932, 1950, 1971, 2011, 2032, 2052, 2063, - 2077, 2079, 2091, 2095, 2172, 2192, 2207, 2208, 2224, 2230, 2247, 2277, 2308, 2345, 2356, 2389, - 2403, 2424, 2501, 2504, 2506, 2520, 2570, 2593, 2616, 2624, 2630, 2646, 2669, 2700, 2714, 2746, - 2754, 2795, 2824, 2835, 2839, 2874, 2882, 2905, 2984, 3028, 3042, 3092, 3108, 3110, 3124, 3153, - 3185, 3215, 3252, 3288, 3294, 3364, 3397, 3434, 3483, 3523, 3537, 3587, 3589, 3591, 3592, 3610, - 3626, 3670, 3680, 3722, 3749, 3754, 3776, 3789, 3803, 3824, 3857, 3873, 3904, 3906, 3924, 3992, - }; - static const uint16_t kgrid_512[512] = { - 0, 1, 2, 5, 7, 8, 9, 10, 12, 14, 16, 17, 21, 27, 32, 34, - 37, 39, 41, 43, 48, 50, 57, 60, 63, 64, 65, 66, 68, 72, 73, 77, - 80, 83, 87, 89, 93, 100, 113, 117, 122, 128, 129, 133, 135, 136, 139, 142, - 145, 149, 152, 156, 162, 165, 167, 169, 171, 184, 187, 195, 201, 205, 208, 210, - 217, 219, 222, 228, 232, 234, 247, 249, 253, 256, 267, 271, 273, 276, 282, 288, - 291, 297, 312, 322, 324, 336, 338, 342, 347, 353, 357, 359, 374, 379, 390, 393, - 395, 409, 426, 441, 448, 450, 452, 464, 466, 470, 475, 488, 492, 512, 513, 514, - 516, 520, 521, 523, 525, 527, 528, 530, 537, 540, 542, 556, 558, 561, 570, 576, - 577, 579, 582, 584, 588, 593, 600, 603, 609, 616, 618, 632, 638, 640, 650, 653, - 655, 656, 660, 666, 672, 675, 685, 688, 698, 705, 708, 711, 712, 715, 721, 727, - 728, 732, 737, 754, 760, 771, 773, 778, 780, 793, 795, 802, 806, 808, 812, 833, - 840, 843, 849, 856, 858, 873, 912, 916, 919, 932, 934, 961, 963, 968, 970, 977, - 989, 993, 1010, 1016, 1024, 1025, 1027, 1029, 1031, 1032, 1034, 1036, 1038, 1041, 1043, 1047, - 1048, 1050, 1057, 1059, 1061, 1064, 1066, 1079, 1080, 1083, 1085, 1088, 1090, 1096, 1099, 1103, - 1106, 1109, 1113, 1116, 1122, 1129, 1153, 1156, 1159, 1169, 1171, 1176, 1183, 1185, 1195, 1199, - 1209, 1212, 1216, 1218, 1221, 1225, 1234, 1236, 1241, 1243, 1250, 1256, 1270, 1281, 1287, 1296, - 1299, 1306, 1309, 1313, 1338, 1341, 1348, 1353, 1362, 1375, 1376, 1387, 1400, 1408, 1410, 1415, - 1425, 1453, 1457, 1477, 1481, 1494, 1496, 1507, 1512, 1538, 1545, 1547, 1549, 1551, 1554, 1561, - 1563, 1565, 1570, 1572, 1575, 1577, 1587, 1593, 1601, 1603, 1605, 1612, 1617, 1619, 1632, 1648, - 1658, 1662, 1664, 1674, 1680, 1690, 1692, 1704, 1729, 1736, 1740, 1745, 1747, 1751, 1752, 1761, - 1763, 1767, 1773, 1787, 1795, 1801, 1806, 1810, 1817, 1834, 1840, 1844, 1857, 1864, 1866, 1877, - 1882, 1892, 1902, 1915, 1934, 1953, 1985, 1987, 2000, 2002, 2013, 2048, 2052, 2058, 2064, 2068, - 2071, 2074, 2081, 2088, 2104, 2114, 2119, 2121, 2123, 2130, 2136, 2141, 2147, 2153, 2157, 2177, - 2179, 2184, 2189, 2193, 2203, 2208, 2223, 2226, 2232, 2244, 2249, 2251, 2256, 2258, 2265, 2269, - 2304, 2306, 2324, 2335, 2336, 2361, 2373, 2375, 2385, 2418, 2443, 2460, 2480, 2504, 2509, 2520, - 2531, 2537, 2562, 2568, 2572, 2578, 2592, 2596, 2599, 2602, 2614, 2620, 2625, 2627, 2629, 2634, - 2641, 2650, 2682, 2688, 2697, 2707, 2712, 2718, 2731, 2754, 2759, 2760, 2775, 2788, 2793, 2805, - 2811, 2817, 2820, 2832, 2842, 2854, 2890, 2902, 2921, 2923, 2978, 3010, 3012, 3026, 3081, 3083, - 3085, 3097, 3099, 3120, 3136, 3152, 3159, 3188, 3210, 3228, 3234, 3245, 3250, 3256, 3264, 3276, - 3281, 3296, 3349, 3363, 3378, 3392, 3395, 3420, 3440, 3461, 3488, 3529, 3531, 3584, 3588, 3591, - 3600, 3602, 3614, 3616, 3628, 3634, 3650, 3657, 3668, 3683, 3685, 3713, 3716, 3720, 3726, 3729, - 3736, 3753, 3778, 3802, 3805, 3819, 3841, 3845, 3851, 3856, 3880, 3922, 3938, 3970, 3993, 4032, - }; - - const int kmap_size = 4096; - const int nwant = grid_size == 256 ? 2 : 3; - const uint16_t * kgrid = grid_size == 256 ? kgrid_256 : kgrid_512; - uint32_t * kgrid_q3xs; - int * kmap_q3xs; - uint16_t * kneighbors_q3xs; - - //printf("================================================================= %s(grid_size = %d)\n", __func__, grid_size); - uint32_t * the_grid = (uint32_t *)malloc(grid_size*sizeof(uint32_t)); - for (int k = 0; k < grid_size; ++k) { - int8_t * pos = (int8_t *)(the_grid + k); - for (int i = 0; i < 4; ++i) { - int l = (kgrid[k] >> 3*i) & 0x7; - pos[i] = 2*l + 1; - } - } - kgrid_q3xs = the_grid; - iq3_data[gindex].grid = the_grid; - kmap_q3xs = (int *)malloc(kmap_size*sizeof(int)); - iq3_data[gindex].map = kmap_q3xs; - for (int i = 0; i < kmap_size; ++i) kmap_q3xs[i] = -1; - uint32_t aux32; - uint8_t * aux8 = (uint8_t *)&aux32; - for (int i = 0; i < grid_size; ++i) { - aux32 = kgrid_q3xs[i]; - uint16_t index = 0; - for (int k=0; k<4; ++k) { - uint16_t q = (aux8[k] - 1)/2; - index |= (q << 3*k); - } - kmap_q3xs[index] = i; - } - int8_t pos[4]; - int * dist2 = (int *)malloc(2*grid_size*sizeof(int)); - int num_neighbors = 0, num_not_in_map = 0; - for (int i = 0; i < kmap_size; ++i) { - if (kmap_q3xs[i] >= 0) continue; - ++num_not_in_map; - for (int k = 0; k < 4; ++k) { - int l = (i >> 3*k) & 0x7; - pos[k] = 2*l + 1; - } - for (int j = 0; j < grid_size; ++j) { - const int8_t * pg = (const int8_t *)(kgrid_q3xs + j); - int d2 = 0; - for (int k = 0; k < 4; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]); - dist2[2*j+0] = d2; - dist2[2*j+1] = j; - } - qsort(dist2, grid_size, 2*sizeof(int), iq3_compare_func); - int n = 0; int d2 = dist2[0]; - int nhave = 1; - for (int j = 0; j < grid_size; ++j) { - if (dist2[2*j] > d2) { - if (nhave == nwant) break; - d2 = dist2[2*j]; - ++nhave; - } - ++n; - } - num_neighbors += n; - } - //printf("%s: %d neighbours in total\n", __func__, num_neighbors); - kneighbors_q3xs = (uint16_t *)malloc((num_neighbors + num_not_in_map)*sizeof(uint16_t)); - iq3_data[gindex].neighbours = kneighbors_q3xs; - int counter = 0; - for (int i = 0; i < kmap_size; ++i) { - if (kmap_q3xs[i] >= 0) continue; - for (int k = 0; k < 4; ++k) { - int l = (i >> 3*k) & 0x7; - pos[k] = 2*l + 1; - } - for (int j = 0; j < grid_size; ++j) { - const int8_t * pg = (const int8_t *)(kgrid_q3xs + j); - int d2 = 0; - for (int k = 0; k < 4; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]); - dist2[2*j+0] = d2; - dist2[2*j+1] = j; - } - qsort(dist2, grid_size, 2*sizeof(int), iq3_compare_func); - kmap_q3xs[i] = -(counter + 1); - int d2 = dist2[0]; - uint16_t * start = &kneighbors_q3xs[counter++]; - int n = 0, nhave = 1; - for (int j = 0; j < grid_size; ++j) { - if (dist2[2*j] > d2) { - if (nhave == nwant) break; - d2 = dist2[2*j]; - ++nhave; - } - kneighbors_q3xs[counter++] = dist2[2*j+1]; - ++n; - } - *start = n; - } - free(dist2); -} - -void iq3xs_free_impl(int grid_size) { - GGML_ASSERT(grid_size == 256 || grid_size == 512); - const int gindex = iq3_data_index(grid_size); - if (iq3_data[gindex].grid) { - free(iq3_data[gindex].grid); iq3_data[gindex].grid = NULL; - free(iq3_data[gindex].map); iq3_data[gindex].map = NULL; - free(iq3_data[gindex].neighbours); iq3_data[gindex].neighbours = NULL; - } -} - -static int iq3_find_best_neighbour(const uint16_t * GGML_RESTRICT neighbours, const uint32_t * GGML_RESTRICT grid, - const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float scale, int8_t * GGML_RESTRICT L) { - int num_neighbors = neighbours[0]; - GGML_ASSERT(num_neighbors > 0); - float best_d2 = FLT_MAX; - int grid_index = -1; - for (int j = 1; j <= num_neighbors; ++j) { - const int8_t * pg = (const int8_t *)(grid + neighbours[j]); - float d2 = 0; - for (int i = 0; i < 4; ++i) { - float q = pg[i]; - float diff = scale*q - xval[i]; - d2 += weight[i]*diff*diff; - } - if (d2 < best_d2) { - best_d2 = d2; grid_index = neighbours[j]; - } - } - GGML_ASSERT(grid_index >= 0); - const int8_t * pg = (const int8_t *)(grid + grid_index); - for (int i = 0; i < 4; ++i) L[i] = (pg[i] - 1)/2; - return grid_index; -} - -static void quantize_row_iq3_xxs_impl(int grid_size, const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, - const float * GGML_RESTRICT quant_weights) { - - const int gindex = iq3_data_index(grid_size); - - const uint32_t * kgrid_q3xs = iq3_data[gindex].grid; - const int * kmap_q3xs = iq3_data[gindex].map; - const uint16_t * kneighbors_q3xs = iq3_data[gindex].neighbours; - - //GGML_ASSERT(quant_weights && "missing quantization weights"); - GGML_ASSERT(kgrid_q3xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kmap_q3xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kneighbors_q3xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(n%QK_K == 0); - - const int kMaxQ = 8; - - const int64_t nbl = n/QK_K; - - ggml_fp16_t * dh; - uint8_t * qs; - int block_size; - if (grid_size == 256) { - block_iq3_xxs * y = vy; - dh = &y->d; - qs = y->qs; - block_size = sizeof(block_iq3_xxs); - } else { - block_iq3_s * y = vy; - dh = &y->d; - qs = y->qs; - block_size = sizeof(block_iq3_s); - } - int quant_size = block_size - sizeof(ggml_fp16_t); - - float scales[QK_K/32]; - float weight[32]; - float xval[32]; - int8_t L[32]; - int8_t Laux[32]; - float waux[32]; - bool is_on_grid[8]; - bool is_on_grid_aux[8]; - uint8_t block_signs[8]; - uint8_t q3[3*(QK_K/8)+QK_K/32]; - uint32_t * scales_and_signs = (uint32_t *)(q3 + QK_K/4); - uint8_t * qh = q3 + 3*(QK_K/8); - - for (int ibl = 0; ibl < nbl; ++ibl) { - - dh[0] = GGML_FP32_TO_FP16(0.f); - memset(q3, 0, 3*QK_K/8+QK_K/32); - - float max_scale = 0; - - const float * xbl = x + QK_K*ibl; - float sumx2 = 0; - for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; - float sigma2 = 2*sumx2/QK_K; - - for (int ib = 0; ib < QK_K/32; ++ib) { - const float * xb = xbl + 32*ib; - if (quant_weights) { - const float * qw = quant_weights + QK_K*ibl + 32*ib; - for (int i = 0; i < 32; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); - } else { - for (int i = 0; i < 32; ++i) weight[i] = xb[i]*xb[i]; - } - for (int i = 0; i < 32; ++i) waux[i] = sqrtf(weight[i]); - for (int k = 0; k < 4; ++k) { - int nflip = 0; - uint8_t s = 0; - for (int i = 0; i < 8; ++i) { - if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; - else { - xval[8*k + i] = -xb[8*k + i]; ++nflip; s |= (1 << i); - } - } - if (nflip%2) { - int imin = 0; float min = weight[8*k+imin]*xb[8*k+imin]*xb[8*k+imin]; - for (int i = 1; i < 8; ++i) { - float ax = weight[8*k+i]*xb[8*k+i]*xb[8*k+i]; - if (ax < min) { - min = ax; imin = i; - } - } - xval[8*k+imin] = -xval[8*k+imin]; - s ^= (1 << imin); - } - block_signs[k] = s & 127; - } - float max = xval[0]; - for (int i = 1; i < 32; ++i) max = MAX(max, xval[i]); - memset(L, 0, 32); - if (max < GROUP_MAX_EPS_IQ3_XXS) { - scales[ib] = 0; - continue; - } - float best = 0; - float scale = max/(2*kMaxQ-1); - for (int k = 0; k < 8; ++k) is_on_grid[k] = true; - for (int is = -15; is <= 15; ++is) { - float id = (2*kMaxQ-1+is*0.2f)/max; - float this_scale = 1/id; - for (int k = 0; k < 8; ++k) { - for (int i = 0; i < 4; ++i) { - int l = nearest_int(0.5f*(id*xval[4*k+i]-1)); - Laux[4*k+i] = MAX(0, MIN(kMaxQ-1, l)); - } - uint16_t u = 0; - for (int i = 0; i < 4; ++i) u |= (Laux[4*k+i] << 3*i); - int grid_index = kmap_q3xs[u]; - is_on_grid_aux[k] = true; - if (grid_index < 0) { - is_on_grid_aux[k] = false; - const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1; - grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, this_scale, Laux + 4*k); - } - } - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < 32; ++i) { - float w = weight[i]; - float q = 2*Laux[i] + 1; - sumqx += w*xval[i]*q; - sumq2 += w*q*q; - } - if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { - scale = sumqx/sumq2; best = scale*sumqx; - for (int i = 0; i < 32; ++i) L[i] = Laux[i]; - for (int k = 0; k < 8; ++k) is_on_grid[k] = is_on_grid_aux[k]; - } - } - int n_not_ongrid = 0; - for (int k = 0; k < 8; ++k) if (!is_on_grid[k]) ++n_not_ongrid; - if (n_not_ongrid > 0 && scale > 0) { - float id = 1/scale; - for (int k = 0; k < 8; ++k) { - if (is_on_grid[k]) continue; - uint16_t u = 0; - for (int i = 0; i < 4; ++i) { - int l = nearest_int(0.5f*(id*xval[4*k+i]-1)); - l = MAX(0, MIN(kMaxQ-1, l)); - u |= (l << 3*i); - } - int grid_index = kmap_q3xs[u]; - if (grid_index < 0) { - const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1; - grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, scale, L + 4*k); - } - const int8_t * pg = (const int8_t *)(kgrid_q3xs + grid_index); - for (int i = 0; i < 4; ++i) L[4*k+i] = (pg[i] - 1)/2; - } - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < 32; ++i) { - float w = weight[i]; - float q = 2*L[i] + 1; - sumqx += w*xval[i]*q; - sumq2 += w*q*q; - } - if (sumq2 > 0) scale = sumqx/sumq2; - } - if (scale < 0) { - // This should never happen, but just in case, flip scale so that it is positive (we use uint's to encode the scale) - // and correspondingly flip quant signs. - scale = -scale; - for (int k = 0; k < 4; ++k) block_signs[k] = (~block_signs[k]) & 127; - } - for (int k = 0; k < 8; ++k) { - uint16_t u = 0; - for (int i = 0; i < 4; ++i) u |= (L[4*k+i] << 3*i); - int grid_index = kmap_q3xs[u]; - if (grid_index < 0) { - printf("Oops: found point %u not on grid:", u); - for (int i = 0; i < 4; ++i) printf(" %d", L[4*k+i]); - printf("\n"); - GGML_ABORT("fatal error"); - } - if (grid_size == 256) { - q3[8*ib+k] = grid_index; - } else { - q3[8*ib+k] = grid_index & 255; - qh[ib] |= ((grid_index >> 8) << k); - } - - } - scales_and_signs[ib] = block_signs[0] | (block_signs[1] << 7) | (block_signs[2] << 14) | (block_signs[3] << 21); - GGML_ASSERT(scale >= 0); - scales[ib] = scale; - max_scale = MAX(max_scale, scale); - } - - if (!max_scale) { - memset(qs, 0, quant_size); - dh += block_size/sizeof(ggml_fp16_t); - qs += block_size; - continue; - } - - float d = max_scale/31; - dh[0] = GGML_FP32_TO_FP16(d * 1.0125f); // small improvement via this fudge factor - float id = 1/d; - for (int ib = 0; ib < QK_K/32; ++ib) { - int l = nearest_int(0.5f*(id*scales[ib]-1)); - l = MAX(0, MIN(15, l)); - scales_and_signs[ib] |= ((uint32_t)l << 28); - } - memcpy(qs, q3, quant_size); - - dh += block_size/sizeof(ggml_fp16_t); - qs += block_size; - - } -} - -size_t quantize_iq3_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - GGML_ASSERT(n_per_row%QK_K == 0); - int64_t nblock = n_per_row/QK_K; - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_iq3_xxs_impl(256, src, qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += nblock*sizeof(block_iq3_xxs); - } - return nrow * nblock * sizeof(block_iq3_xxs); -} - -void quantize_row_iq3_xxs_ref(const float * GGML_RESTRICT x, block_iq3_xxs * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - quantize_row_iq3_xxs_impl(256, x, y, k, NULL); -} - -static void quantize_row_iq3_s_impl(int block_size, const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int n, - const float * GGML_RESTRICT quant_weights, - float * scales, - float * weight, - float * xval, - int8_t * L, - int8_t * Laux, - float * waux, - bool * is_on_grid, - bool * is_on_grid_aux, - uint8_t * block_signs) { - - const int gindex = iq3_data_index(512); - - const uint32_t * kgrid_q3xs = iq3_data[gindex].grid; - const int * kmap_q3xs = iq3_data[gindex].map; - const uint16_t * kneighbors_q3xs = iq3_data[gindex].neighbours; - - //GGML_ASSERT(quant_weights && "missing quantization weights"); - GGML_ASSERT(kgrid_q3xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kmap_q3xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kneighbors_q3xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(n%QK_K == 0); - - const int kMaxQ = 8; - - const int64_t nbl = n/QK_K; - - block_iq3_s * y = vy; - - const int bs4 = block_size/4; - const int bs8 = block_size/8; - - for (int ibl = 0; ibl < nbl; ++ibl) { - - memset(&y[ibl], 0, sizeof(block_iq3_s)); - y[ibl].d = GGML_FP32_TO_FP16(0.f); - - uint8_t * qs = y[ibl].qs; - uint8_t * qh = y[ibl].qh; - uint8_t * signs = y[ibl].signs; - - float max_scale = 0; - - const float * xbl = x + QK_K*ibl; - float sumx2 = 0; - for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; - float sigma2 = 2*sumx2/QK_K; - - for (int ib = 0; ib < QK_K/block_size; ++ib) { - const float * xb = xbl + block_size*ib; - if (quant_weights) { - const float * qw = quant_weights + QK_K*ibl + block_size*ib; - for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); - } else { - for (int i = 0; i < block_size; ++i) weight[i] = xb[i]*xb[i]; - } - for (int i = 0; i < block_size; ++i) waux[i] = sqrtf(weight[i]); - for (int k = 0; k < bs8; ++k) { - uint8_t s = 0; - for (int i = 0; i < 8; ++i) { - if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; - else { - xval[8*k + i] = -xb[8*k + i]; s |= (1 << i); - } - } - block_signs[k] = s; - } - float max = xval[0]; - for (int i = 1; i < block_size; ++i) max = MAX(max, xval[i]); - memset(L, 0, block_size); - if (!max) { - scales[ib] = 0; - continue; - } - float best = 0; - float scale = max/(2*kMaxQ-1); - for (int k = 0; k < bs4; ++k) is_on_grid[k] = false; - for (int is = -9; is <= 9; ++is) { - float id = (2*kMaxQ-1+is*0.2f)/max; - float this_scale = 1/id; - for (int k = 0; k < bs4; ++k) { - for (int i = 0; i < 4; ++i) { - int l = nearest_int(0.5f*(id*xval[4*k+i]-1)); - Laux[4*k+i] = MAX(0, MIN(kMaxQ-1, l)); - } - uint16_t u = 0; - for (int i = 0; i < 4; ++i) u |= (Laux[4*k+i] << 3*i); - int grid_index = kmap_q3xs[u]; - is_on_grid_aux[k] = true; - if (grid_index < 0) { - is_on_grid_aux[k] = false; - const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1; - grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, this_scale, Laux + 4*k); - } - } - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < block_size; ++i) { - float w = weight[i]; - float q = 2*Laux[i] + 1; - sumqx += w*xval[i]*q; - sumq2 += w*q*q; - } - if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { - scale = sumqx/sumq2; best = scale*sumqx; - for (int i = 0; i < block_size; ++i) L[i] = Laux[i]; - for (int k = 0; k < bs4; ++k) is_on_grid[k] = is_on_grid_aux[k]; - } - } - int n_not_ongrid = 0; - for (int k = 0; k < bs4; ++k) if (!is_on_grid[k]) ++n_not_ongrid; - if (n_not_ongrid > 0 && scale > 0) { - float id = 1/scale; - for (int k = 0; k < bs4; ++k) { - //if (is_on_grid[k]) continue; - uint16_t u = 0; - for (int i = 0; i < 4; ++i) { - int l = nearest_int(0.5f*(id*xval[4*k+i]-1)); - l = MAX(0, MIN(kMaxQ-1, l)); - u |= (l << 3*i); - } - int grid_index = kmap_q3xs[u]; - if (grid_index < 0) { - const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1; - grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, scale, L + 4*k); - } - const int8_t * pg = (const int8_t *)(kgrid_q3xs + grid_index); - for (int i = 0; i < 4; ++i) L[4*k+i] = (pg[i] - 1)/2; - } - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < block_size; ++i) { - float w = weight[i]; - float q = 2*L[i] + 1; - sumqx += w*xval[i]*q; - sumq2 += w*q*q; - } - if (sumq2 > 0) scale = sumqx/sumq2; - } - if (scale < 0) { - // This should never happen, but just in case, flip scale so that it is positive (we use uint's to encode the scale) - // and correspondingly flip quant signs. - scale = -scale; - for (int k = 0; k < bs8; ++k) block_signs[k] = ~block_signs[k]; - } - for (int k = 0; k < bs4; ++k) { - uint16_t u = 0; - for (int i = 0; i < 4; ++i) u |= (L[4*k+i] << 3*i); - int grid_index = kmap_q3xs[u]; - if (grid_index < 0) { - printf("Oops: found point %u not on grid:", u); - for (int i = 0; i < 4; ++i) printf(" %d", L[4*k+i]); - printf("\n"); - GGML_ABORT("fatal error"); - } - qs[k] = grid_index & 255; - qh[(ib*bs4+k)/8] |= ((grid_index >> 8) << ((ib*bs4+k)%8)); - } - qs += bs4; - for (int k = 0; k < bs8; ++k) signs[k] = block_signs[k]; - signs += bs8; - GGML_ASSERT(scale >= 0); - scales[ib] = scale; - max_scale = MAX(max_scale, scale); - } - - if (!max_scale) { - continue; - } - - float d = max_scale/31; - y[ibl].d = GGML_FP32_TO_FP16(d * 1.033f); - float id = 1/d; - for (int ib = 0; ib < QK_K/block_size; ib += 2) { - int l1 = nearest_int(0.5f*(id*scales[ib+0]-1)); - l1 = MAX(0, MIN(15, l1)); - int l2 = nearest_int(0.5f*(id*scales[ib+1]-1)); - l2 = MAX(0, MIN(15, l2)); - y[ibl].scales[ib/2] = l1 | (l2 << 4); - } - - } -} - -#define IQ3S_BLOCK_SIZE 32 -size_t quantize_iq3_s(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - GGML_ASSERT(n_per_row%QK_K == 0); - int64_t nblock = n_per_row/QK_K; - float scales[QK_K/IQ3S_BLOCK_SIZE]; - float weight[IQ3S_BLOCK_SIZE]; - float xval[IQ3S_BLOCK_SIZE]; - int8_t L[IQ3S_BLOCK_SIZE]; - int8_t Laux[IQ3S_BLOCK_SIZE]; - float waux[IQ3S_BLOCK_SIZE]; - bool is_on_grid[IQ3S_BLOCK_SIZE/4]; - bool is_on_grid_aux[IQ3S_BLOCK_SIZE/4]; - uint8_t block_signs[IQ3S_BLOCK_SIZE/8]; - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_iq3_s_impl(IQ3S_BLOCK_SIZE, src, qrow, n_per_row, quant_weights, - scales, weight, xval, L, Laux, waux, is_on_grid, is_on_grid_aux, block_signs); - src += n_per_row; - qrow += nblock*sizeof(block_iq3_s); - } - return nrow * nblock * sizeof(block_iq3_s); -} - -void quantize_row_iq3_s_ref(const float * GGML_RESTRICT x, block_iq3_s * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - quantize_iq3_s(x, y, 1, k, NULL); -} - - -// =================================== 1.5 bpw =================================================== - -static int iq1_find_best_neighbour(const uint16_t * GGML_RESTRICT neighbours, const uint64_t * GGML_RESTRICT grid, - const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float * scale, int8_t * GGML_RESTRICT L, int ngrid) { - int num_neighbors = neighbours[0]; - GGML_ASSERT(num_neighbors > 0); - float best_score = -FLT_MAX; - int grid_index = -1; - for (int j = 1; j <= num_neighbors; ++j) { - const int8_t * pg = (const int8_t *)(grid + neighbours[j]); - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < 8; ++i) { - float q = (pg[i] - 3)/2; - float w = weight[i]; - sumqx += w*q*xval[i]; - sumq2 += w*q*q; - } - if (sumqx > 0 && sumq2 > 0 && sumqx*sumqx > best_score*sumq2) { - *scale = sumqx/sumq2; best_score = *scale * sumqx; - grid_index = neighbours[j]; - } - } - if (grid_index < 0) { - for (int i = 0; i < ngrid; ++i) { - const int8_t * grid_i = (const int8_t *)(grid + i); - float sumqx = 0, sumq2 = 0; - for (int j = 0; j < 8; ++j) { - float w = weight[j]; - float q = (grid_i[j] - 3)/2; - sumqx += w*q*xval[j]; - sumq2 += w*q*q; - } - if (sumqx > 0 && sumq2 > 0 && sumqx*sumqx > best_score*sumq2) { - *scale = sumqx/sumq2; best_score = *scale*sumqx; - grid_index = i; - } - } - } - if (grid_index < 0) { - printf("Oops, did not find grid point\n"); - printf("Have %d neighbours\n", num_neighbors); - for (int j = 1; j <= num_neighbors; ++j) { - const int8_t * pg = (const int8_t *)(grid + neighbours[j]); - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < 8; ++i) { - float q = (pg[i] - 3)/2; - float w = weight[i]; - sumqx += w*q*xval[i]; - sumq2 += w*q*q; - } - printf(" neighbour %d: sumqx = %g sumq2 = %g\n", j, (double)sumqx, (double)sumq2); - } - } - GGML_ASSERT(grid_index >= 0); - //!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! - *scale *= 1.05f; // This is a fudge factor. Don't ask me why it improves the result. - //!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! - const int8_t * pg = (const int8_t *)(grid + grid_index); - for (int i = 0; i < 8; ++i) L[i] = (pg[i] - 1)/2; - return grid_index; -} - -static int iq1_find_best_neighbour2(const uint16_t * GGML_RESTRICT neighbours, const uint64_t * GGML_RESTRICT grid, - const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float scale, const float * GGML_RESTRICT xg, int8_t * GGML_RESTRICT L, int ngrid) { - int num_neighbors = neighbours[0]; - GGML_ASSERT(num_neighbors > 0); - float best_score = FLT_MAX; - int grid_index = -1; - for (int j = 1; j <= num_neighbors; ++j) { - const int8_t * pg = (const int8_t *)(grid + neighbours[j]); - float d2 = 0; - for (int i = 0; i < 8; ++i) { - float q = xg[(pg[i] - 1)/2]; - float w = weight[i]; - float diff = scale*q - xval[i]; - d2 += w*diff*diff; - } - if (d2 < best_score) { - best_score = d2; - grid_index = neighbours[j]; - } - } - if (grid_index < 0) { - for (int i = 0; i < ngrid; ++i) { - const int8_t * grid_i = (const int8_t *)(grid + i); - float d2 = 0; - for (int j = 0; j < 8; ++j) { - float w = weight[j]; - float q = xg[(grid_i[j] - 1)/2]; - float diff = scale*q - xval[i]; - d2 += w*diff*diff; - } - if (d2 < best_score) { - best_score = d2; - grid_index = i; - } - } - } - if (grid_index < 0) { - printf("Oops, did not find grid point\n"); - printf("Have %d neighbours\n", num_neighbors); - for (int j = 1; j <= num_neighbors; ++j) { - const int8_t * pg = (const int8_t *)(grid + neighbours[j]); - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < 8; ++i) { - float q = xg[(pg[i] - 1)/2]; - float w = weight[i]; - sumqx += w*q*xval[i]; - sumq2 += w*q*q; - } - printf(" neighbour %d: sumqx = %g sumq2 = %g\n", j, (double)sumqx, (double)sumq2); - } - } - GGML_ASSERT(grid_index >= 0); - const int8_t * pg = (const int8_t *)(grid + grid_index); - for (int i = 0; i < 8; ++i) L[i] = (pg[i] - 1)/2; - return grid_index; -} - -static int iq1_sort_helper(const void * left, const void * right) { - const float * l = left; - const float * r = right; - return *l < *r ? -1 : *l > *r ? 1 : 0; -} - -#define IQ1S_BLOCK_SIZE 32 -#define IQ1M_BLOCK_SIZE 16 -static void quantize_row_iq1_s_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights, - float * scales, - float * weight, - float * sumx, - float * sumw, - float * pairs, - int8_t * L, - uint16_t * index, - int8_t * shifts) { - - const int gindex = iq2_data_index(GGML_TYPE_IQ1_S); - - const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; - const int * kmap_q2xs = iq2_data[gindex].map; - const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; - - GGML_ASSERT(quant_weights && "missing quantization weights"); - GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(n%QK_K == 0); - - block_iq1_s * y = vy; - - const int64_t nbl = n/QK_K; - - const int block_size = IQ1S_BLOCK_SIZE; - - const float x_p[3] = {-1 + IQ1S_DELTA, IQ1S_DELTA, 1 + IQ1S_DELTA}; - const float x_m[3] = {-1 - IQ1S_DELTA, -IQ1S_DELTA, 1 - IQ1S_DELTA}; - - - int * idx = (int *)(pairs + 1); - - for (int ibl = 0; ibl < nbl; ++ibl) { - - y[ibl].d = GGML_FP32_TO_FP16(0.f); - memset(y[ibl].qs, 0, QK_K/8); - memset(y[ibl].qh, 0, QK_K/16); - - float max_scale = 0; - - const float * xbl = x + QK_K*ibl; - float sumx2 = 0; - for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; - float sigma2 = 2*sumx2/QK_K; - - for (int ib = 0; ib < QK_K/block_size; ++ib) { - const float * xb = xbl + block_size*ib; - const float * qw = quant_weights + QK_K*ibl + block_size*ib; - for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); - float max = fabsf(xb[0]); - for (int i = 1; i < block_size; ++i) max = MAX(max, fabsf(xb[i])); - if (max < GROUP_MAX_EPS_IQ1_S) { - scales[ib] = 0; - shifts[ib] = 1; - memset(L, 1, block_size); - continue; - } - // Here we solve exactly the sum of squared difference (SSD) weighted minimization problem. - // With just 3 allowed quant values (-1, 0, 1), we can search exhaustively for the two - // boundaries that split the weights xb[i] into 3 groups. To do so, we sort the weights - // in ascending order, compute Si = sum[weight[j] xb[j], j = 0...i] and - // Wi = sum[weight[j], j = 0...i], and use these to quckly get get the optimum scale - // for each possible and score for each split. - for (int j = 0; j < block_size; ++j) { - pairs[2*j] = xb[j]; - idx[2*j] = j; - } - qsort(pairs, block_size, 2*sizeof(float), iq1_sort_helper); - { - sumx[0] = sumw[0] = 0; - for (int j = 0; j < block_size; ++j) { - int i = idx[2*j]; - sumx[j+1] = sumx[j] + weight[i]*xb[i]; - sumw[j+1] = sumw[j] + weight[i]; - } - } - float best_score = -FLT_MAX, scale = max; - int besti1 = -1, besti2 = -1, best_shift = 0; - for (int i1 = 0; i1 <= block_size; ++i1) { - for (int i2 = i1; i2 <= block_size; ++i2) { - float sumqx = (sumx[i1] - sumx[0])*x_p[0] + (sumx[i2] - sumx[i1])*x_p[1] + (sumx[block_size] - sumx[i2])*x_p[2]; - float sumq2 = (sumw[i1] - sumw[0])*x_p[0]*x_p[0] + (sumw[i2] - sumw[i1])*x_p[1]*x_p[1] + (sumw[block_size] - sumw[i2])*x_p[2]*x_p[2]; - if (sumq2 > 0 && sumqx*sumqx > best_score*sumq2) { - scale = sumqx/sumq2; best_score = scale*sumqx; - besti1 = i1; besti2 = i2; best_shift = 1; - } - sumqx = (sumx[i1] - sumx[0])*x_m[0] + (sumx[i2] - sumx[i1])*x_m[1] + (sumx[block_size] - sumx[i2])*x_m[2]; - sumq2 = (sumw[i1] - sumw[0])*x_m[0]*x_m[0] + (sumw[i2] - sumw[i1])*x_m[1]*x_m[1] + (sumw[block_size] - sumw[i2])*x_m[2]*x_m[2]; - if (sumq2 > 0 && sumqx*sumqx > best_score*sumq2) { - scale = sumqx/sumq2; best_score = scale*sumqx; - besti1 = i1; besti2 = i2; best_shift = -1; - } - } - } - if (besti1 < 0 || besti2 < 0 || best_shift == 0) { - scales[ib] = 0; - shifts[ib] = 1; - memset(L, 1, block_size); - continue; - } - for (int j = 0; j < besti1; ++j) L[idx[2*j]] = 0; - for (int j = besti1; j < besti2; ++j) L[idx[2*j]] = 1; - for (int j = besti2; j < block_size; ++j) L[idx[2*j]] = 2; - if (scale < 0) { - for (int j = 0; j < block_size; ++j) L[j] = 2 - L[j]; - scale = -scale; best_shift = -best_shift; - } - bool all_on_grid = true; - const float * xx = best_shift == 1 ? x_p : x_m; - for (int k = 0; k < block_size/8; ++k) { - uint16_t u = 0; - for (int j = 0; j < 8; ++j) u |= (L[8*k+j] << 2*j); - int grid_index = kmap_q2xs[u]; - if (grid_index < 0) { - all_on_grid = false; - const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; - grid_index = iq1_find_best_neighbour2(neighbours, kgrid_q2xs, xb + 8*k, weight + 8*k, scale, xx, L + 8*k, NGRID_IQ1S); - GGML_ASSERT(grid_index >= 0); - } - index[k] = grid_index; - } - if (!all_on_grid) { - float sumqx = 0, sumq2 = 0; - for (int k = 0; k < block_size/8; ++k) { - const int8_t * pg = (const int8_t *)(kgrid_q2xs + index[k]); - for (int j = 0; j < 8; ++j) { - float w = weight[8*k + j]; - float q = xx[(pg[j] - 1)/2]; - sumqx += w*q*xb[8*k+j]; - sumq2 += w*q*q; - } - } - if (sumqx > 0 && sumq2 > 0) scale = sumqx/sumq2; - } - uint16_t h = 0; - for (int k = 0; k < block_size/8; ++k) { - y[ibl].qs[(block_size/8)*ib + k] = index[k] & 255; - h |= (index[k] >> 8) << 3*k; - } - y[ibl].qh[ib] = h; - GGML_ASSERT(scale >= 0); - scales[ib] = scale; - shifts[ib] = best_shift; - max_scale = MAX(max_scale, scale); - } - - if (!max_scale) { - continue; - } - - float d = max_scale/15; - y[ibl].d = GGML_FP32_TO_FP16(d*1.125f); // 1.125f is another fudge factor. Don't ask me why it is needed. - float id = 1/d; - for (int ib = 0; ib < QK_K/block_size; ++ib) { - int l = nearest_int(0.5f*(id*scales[ib]-1)); - l = MAX(0, MIN(7, l)); - if (shifts[ib] == -1) l |= 8; - y[ibl].qh[ib] |= (l << 12); - } - } -} - -size_t quantize_iq1_s(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - GGML_ASSERT(n_per_row%QK_K == 0); - float scales[QK_K/IQ1S_BLOCK_SIZE]; - float weight[IQ1S_BLOCK_SIZE]; - int8_t L[IQ1S_BLOCK_SIZE]; - float sumx[IQ1S_BLOCK_SIZE+1]; - float sumw[IQ1S_BLOCK_SIZE+1]; - float pairs[2*IQ1S_BLOCK_SIZE]; - uint16_t index[IQ1S_BLOCK_SIZE/8]; - int8_t shifts[QK_K/IQ1S_BLOCK_SIZE]; - int64_t nblock = n_per_row/QK_K; - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_iq1_s_impl(src, qrow, n_per_row, quant_weights, scales, weight, sumx, sumw, pairs, L, index, shifts); - src += n_per_row; - qrow += nblock*sizeof(block_iq1_s); - } - return nrow * nblock * sizeof(block_iq1_s); -} - -static void quantize_row_iq1_m_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights, - float * scales, - float * weight, - float * pairs, - int8_t * L, - uint16_t * index, - int8_t * shifts) { - - const int gindex = iq2_data_index(GGML_TYPE_IQ1_M); - - const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; - const int * kmap_q2xs = iq2_data[gindex].map; - const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; - - //GGML_ASSERT(quant_weights && "missing quantization weights"); - GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(n%QK_K == 0); - - block_iq1_m * y = vy; - - const int64_t nbl = n/QK_K; - - const int block_size = IQ1M_BLOCK_SIZE; - - const float x_p[3] = {-1 + IQ1M_DELTA, IQ1M_DELTA, 1 + IQ1M_DELTA}; - const float x_m[3] = {-1 - IQ1M_DELTA, -IQ1M_DELTA, 1 - IQ1M_DELTA}; - const uint8_t masks[4] = {0x00, 0x80, 0x08, 0x88}; - - int * idx = (int *)(pairs + 1); - - float sumqx[4], sumq2[4]; - - iq1m_scale_t s; - const float * xx; - - for (int ibl = 0; ibl < nbl; ++ibl) { - memset(y[ibl].qs, 0, QK_K/8); - memset(y[ibl].qh, 0, QK_K/16); - memset(y[ibl].scales, 0, QK_K/32); - - float max_scale = 0; - - const float * xbl = x + QK_K*ibl; - float sumx2 = 0; - for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; - float sigma2 = 2*sumx2/QK_K; - - for (int ib = 0; ib < QK_K/block_size; ++ib) { - const float * xb = xbl + block_size*ib; - if (quant_weights) { - const float * qw = quant_weights + QK_K*ibl + block_size*ib; - for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); - } else { - for (int i = 0; i < block_size; ++i) weight[i] = xb[i]*xb[i]; - } - float max = fabsf(xb[0]); - for (int i = 1; i < block_size; ++i) max = MAX(max, fabsf(xb[i])); - if (max < GROUP_MAX_EPS_IQ1_M) { - scales[ib] = 0; - shifts[ib] = 0; - memset(L, 1, block_size); - continue; - } - // Here we solve exactly the sum of squared difference (SSD) weighted minimization problem. - // With just 3 allowed quant values (-1, 0, 1), we can search exhaustively for the two - // boundaries that split the weights xb[i] into 3 groups. To do so, we sort the weights - // in ascending order, compute Si = sum[weight[j] xb[j], j = 0...i] and - // Wi = sum[weight[j], j = 0...i], and use these to quckly get get the optimum scale - // for each possible and score for each split. - for (int j = 0; j < block_size; ++j) { - pairs[2*j] = xb[j]; - idx[2*j] = j; - } - qsort(pairs, block_size, 2*sizeof(float), iq1_sort_helper); - float best_score = -FLT_MAX, scale = max; - int besti1 = -1, besti2 = -1, best_k = -1; - // 0: +, + - // 1: +, - - // 2: -, + - // 3: -, - - for (int i1 = 0; i1 <= block_size; ++i1) { - for (int i2 = i1; i2 <= block_size; ++i2) { - memset(sumqx, 0, 4*sizeof(float)); - memset(sumq2, 0, 4*sizeof(float)); - for (int j = 0; j < i1; ++j) { - int i = idx[2*j]; - if (i < block_size/2) { - sumqx[0] += weight[i]*x_p[0]*xb[i]; - sumqx[1] += weight[i]*x_p[0]*xb[i]; - sumqx[2] += weight[i]*x_m[0]*xb[i]; - sumqx[3] += weight[i]*x_m[0]*xb[i]; - sumq2[0] += weight[i]*x_p[0]*x_p[0]; - sumq2[1] += weight[i]*x_p[0]*x_p[0]; - sumq2[2] += weight[i]*x_m[0]*x_m[0]; - sumq2[3] += weight[i]*x_m[0]*x_m[0]; - } else { - sumqx[0] += weight[i]*x_p[0]*xb[i]; - sumqx[2] += weight[i]*x_p[0]*xb[i]; - sumqx[1] += weight[i]*x_m[0]*xb[i]; - sumqx[3] += weight[i]*x_m[0]*xb[i]; - sumq2[0] += weight[i]*x_p[0]*x_p[0]; - sumq2[2] += weight[i]*x_p[0]*x_p[0]; - sumq2[1] += weight[i]*x_m[0]*x_m[0]; - sumq2[3] += weight[i]*x_m[0]*x_m[0]; - } - } - for (int j = i1; j < i2; ++j) { - int i = idx[2*j]; - if (i < block_size/2) { - sumqx[0] += weight[i]*x_p[1]*xb[i]; - sumqx[1] += weight[i]*x_p[1]*xb[i]; - sumqx[2] += weight[i]*x_m[1]*xb[i]; - sumqx[3] += weight[i]*x_m[1]*xb[i]; - sumq2[0] += weight[i]*x_p[1]*x_p[1]; - sumq2[1] += weight[i]*x_p[1]*x_p[1]; - sumq2[2] += weight[i]*x_m[1]*x_m[1]; - sumq2[3] += weight[i]*x_m[1]*x_m[1]; - } else { - sumqx[0] += weight[i]*x_p[1]*xb[i]; - sumqx[2] += weight[i]*x_p[1]*xb[i]; - sumqx[1] += weight[i]*x_m[1]*xb[i]; - sumqx[3] += weight[i]*x_m[1]*xb[i]; - sumq2[0] += weight[i]*x_p[1]*x_p[1]; - sumq2[2] += weight[i]*x_p[1]*x_p[1]; - sumq2[1] += weight[i]*x_m[1]*x_m[1]; - sumq2[3] += weight[i]*x_m[1]*x_m[1]; - } - } - for (int j = i2; j < block_size; ++j) { - int i = idx[2*j]; - if (i < block_size/2) { - sumqx[0] += weight[i]*x_p[2]*xb[i]; - sumqx[1] += weight[i]*x_p[2]*xb[i]; - sumqx[2] += weight[i]*x_m[2]*xb[i]; - sumqx[3] += weight[i]*x_m[2]*xb[i]; - sumq2[0] += weight[i]*x_p[2]*x_p[2]; - sumq2[1] += weight[i]*x_p[2]*x_p[2]; - sumq2[2] += weight[i]*x_m[2]*x_m[2]; - sumq2[3] += weight[i]*x_m[2]*x_m[2]; - } else { - sumqx[0] += weight[i]*x_p[2]*xb[i]; - sumqx[2] += weight[i]*x_p[2]*xb[i]; - sumqx[1] += weight[i]*x_m[2]*xb[i]; - sumqx[3] += weight[i]*x_m[2]*xb[i]; - sumq2[0] += weight[i]*x_p[2]*x_p[2]; - sumq2[2] += weight[i]*x_p[2]*x_p[2]; - sumq2[1] += weight[i]*x_m[2]*x_m[2]; - sumq2[3] += weight[i]*x_m[2]*x_m[2]; - } - } - for (int k = 0; k < 4; ++k) { - if (sumq2[k] > 0 && sumqx[k]*sumqx[k] > best_score*sumq2[k]) { - scale = sumqx[k]/sumq2[k]; best_score = scale*sumqx[k]; - besti1 = i1; besti2 = i2; best_k = k; - } - } - } - } - if (besti1 < 0 || besti2 < 0 || best_k < 0) { - scales[ib] = 0; - shifts[ib] = 0; - memset(L, 1, block_size); - continue; - } - for (int j = 0; j < besti1; ++j) L[idx[2*j]] = 0; - for (int j = besti1; j < besti2; ++j) L[idx[2*j]] = 1; - for (int j = besti2; j < block_size; ++j) L[idx[2*j]] = 2; - if (scale < 0) { - for (int j = 0; j < block_size; ++j) L[j] = 2 - L[j]; - scale = -scale; - best_k = best_k == 0 ? 3 : best_k == 1 ? 2 : best_k == 2 ? 1 : 0; - } - bool all_on_grid = true; - for (int k = 0; k < block_size/8; ++k) { - if (k == 0) xx = best_k < 2 ? x_p : x_m; - else xx = best_k%2 == 0 ? x_p : x_m; - uint16_t u = 0; - for (int j = 0; j < 8; ++j) u |= (L[8*k+j] << 2*j); - int grid_index = kmap_q2xs[u]; - if (grid_index < 0) { - all_on_grid = false; - const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; - grid_index = iq1_find_best_neighbour2(neighbours, kgrid_q2xs, xb + 8*k, weight + 8*k, scale, xx, L + 8*k, NGRID_IQ1S); - GGML_ASSERT(grid_index >= 0); - } - index[k] = grid_index; - } - if (!all_on_grid) { - float sumqx_f = 0, sumq2_f = 0; - for (int k = 0; k < block_size/8; ++k) { - if (k == 0) xx = best_k < 2 ? x_p : x_m; - else xx = best_k%2 == 0 ? x_p : x_m; - const int8_t * pg = (const int8_t *)(kgrid_q2xs + index[k]); - for (int j = 0; j < 8; ++j) { - float w = weight[8*k + j]; - float q = xx[(pg[j] - 1)/2]; - sumqx_f += w*q*xb[8*k+j]; - sumq2_f += w*q*q; - } - } - if (sumqx_f > 0 && sumq2_f > 0) scale = sumqx_f/sumq2_f; - } - y[ibl].qs[2*ib + 0] = index[0] & 255; - y[ibl].qs[2*ib + 1] = index[1] & 255; - y[ibl].qh[ib] = (index[0] >> 8) | ((index[1] >> 8) << 4); - GGML_ASSERT(scale >= 0); - scales[ib] = scale; - shifts[ib] = best_k; - max_scale = MAX(max_scale, scale); - } - - if (!max_scale) { - continue; - } - - uint16_t * sc = (uint16_t *)y[ibl].scales; - float d = max_scale/15; - float id = 1/d; - float sumqx_f = 0, sumq2_f = 0; - for (int ib = 0; ib < QK_K/block_size; ++ib) { - int l = nearest_int(0.5f*(id*scales[ib+0]-1)); - l = MAX(0, MIN(7, l)); - sc[ib/4] |= (l << 3*(ib%4)); - y[ibl].qh[ib] |= masks[shifts[ib]]; - const float * xb = xbl + block_size*ib; - if (quant_weights) { - const float * qw = quant_weights + QK_K*ibl + block_size*ib; - for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); - } else { - for (int i = 0; i < block_size; ++i) weight[i] = xb[i]*xb[i]; - } - for (int k = 0; k < block_size/8; ++k) { - if (k == 0) xx = shifts[ib] < 2 ? x_p : x_m; - else xx = shifts[ib]%2 == 0 ? x_p : x_m; - const int8_t * pg = (const int8_t *)(kgrid_q2xs + y[ibl].qs[2*ib+k] + ((y[ibl].qh[ib] << (8 - 4*k)) & 0x700)); - for (int j = 0; j < 8; ++j) { - float w = weight[8*k + j]; - float q = xx[(pg[j] - 1)/2]*(2*l+1); - sumqx_f += w*q*xb[8*k+j]; - sumq2_f += w*q*q; - } - } - } - if (sumq2_f > 0) d = sumqx_f/sumq2_f; - s.f16 = GGML_FP32_TO_FP16(d*1.1125f); // 1.1125f is another fudge factor. Don't ask me why it is needed. - sc[0] |= ((s.u16 & 0x000f) << 12); - sc[1] |= ((s.u16 & 0x00f0) << 8); - sc[2] |= ((s.u16 & 0x0f00) << 4); - sc[3] |= ((s.u16 & 0xf000) << 0); - } -} - -size_t quantize_iq1_m(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - GGML_ASSERT(n_per_row%QK_K == 0); - float scales[QK_K/IQ1M_BLOCK_SIZE]; - float weight[IQ1M_BLOCK_SIZE]; - int8_t L[IQ1M_BLOCK_SIZE]; - float pairs[2*IQ1M_BLOCK_SIZE]; - uint16_t index[IQ1M_BLOCK_SIZE/8]; - int8_t shifts[QK_K/IQ1M_BLOCK_SIZE]; - int64_t nblock = n_per_row/QK_K; - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_iq1_m_impl(src, qrow, n_per_row, quant_weights, scales, weight, pairs, L, index, shifts); - src += n_per_row; - qrow += nblock*sizeof(block_iq1_m); - } - return nrow * nblock * sizeof(block_iq1_m); -} - -// ============================ 4-bit non-linear quants - -static void quantize_row_iq4_nl_impl(const int super_block_size, const int block_size, const float * GGML_RESTRICT x, - ggml_fp16_t * dh, uint8_t * q4, uint16_t * scales_h, uint8_t * scales_l, - float * scales, float * weight, uint8_t * L, - const int8_t * values, - const float * quant_weights, - const int ntry) { - - float sigma2 = 0; - for (int j = 0; j < super_block_size; ++j) sigma2 += x[j]*x[j]; - sigma2 *= 2.f/super_block_size; - - memset(q4, 0, super_block_size/2); - dh[0] = GGML_FP32_TO_FP16(0.f); - - float max_scale = 0, amax_scale = 0; - for (int ib = 0; ib < super_block_size/block_size; ++ib) { - const float * xb = x + ib*block_size; - uint8_t * Lb = L + 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 < GROUP_MAX_EPS) { - scales[ib] = 0; - continue; - } - float d = ntry > 0 ? -max/values[0] : max/values[0]; - float id = 1/d; - float sumqx = 0, sumq2 = 0; - for (int j = 0; j < block_size; ++j) { - float al = id*xb[j]; - int l = best_index_int8(16, values, al); - Lb[j] = l; - float q = values[l]; - float w = weight[j]; - sumqx += w*q*xb[j]; - sumq2 += w*q*q; - } - d = sumq2 > 0 ? sumqx/sumq2 : 0.f; - float best = d*sumqx; - for (int itry = -ntry; itry <= ntry; ++itry) { - id = (itry + values[0])/max; - sumqx = sumq2 = 0; - for (int j = 0; j < block_size; ++j) { - float al = id*xb[j]; - int l = best_index_int8(16, values, al); - float q = values[l]; - float w = weight[j]; - sumqx += w*q*xb[j]; - sumq2 += w*q*q; - } - if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { - d = sumqx/sumq2; best = d * sumqx; - } - } - scales[ib] = d; - float abs_d = fabsf(d); - if (abs_d > amax_scale) { - amax_scale = abs_d; max_scale = d; - } - } - - if (super_block_size/block_size > 1) { - int nb = super_block_size/block_size; - memset(scales_h, 0, ((nb+7)/8)*sizeof(uint16_t)); - float d = -max_scale/32; - dh[0] = GGML_FP32_TO_FP16(d); - float id = d ? 1/d : 0.f; - for (int ib = 0; ib < super_block_size/block_size; ++ib) { - 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; - for (int j = 0; j < block_size; ++j) { - Lb[j] = best_index_int8(16, values, idl*xb[j]); - } - l += 32; - uint8_t l_l = l & 0xf; - uint8_t l_h = l >> 4; - if (ib%2 == 0) scales_l[ib/2] = l_l; - else scales_l[ib/2] |= (l_l << 4); - scales_h[ib/8] |= (l_h << 2*(ib%8)); - } - } else { - dh[0] = GGML_FP32_TO_FP16(scales[0]); - if (ntry > 0) { - float id = scales[0] ? 1/scales[0] : 0; - for (int j = 0; j < super_block_size; ++j) { - L[j] = best_index_int8(16, values, id*x[j]); - } - } - } - - for (int i = 0; i < super_block_size/32; ++i) { - for (int j = 0; j < 16; ++j) { - q4[16*i + j] = L[32*i + j] | (L[32*i + 16 + j] << 4); - } - } -} - -size_t quantize_iq4_nl(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - GGML_ASSERT(n_per_row%QK4_NL == 0); - int64_t nblock = n_per_row/QK4_NL; - char * qrow = (char *)dst; - uint8_t L[QK4_NL]; - float weight[QK4_NL]; - uint16_t unused_h; - uint8_t * unused_l = NULL; - float scale; - for (int64_t row = 0; row < nrow; ++row) { - block_iq4_nl * iq4 = (block_iq4_nl *)qrow; - for (int ibl = 0; ibl < nblock; ++ibl) { - const float * qw = quant_weights ? quant_weights + QK4_NL*ibl : NULL; - quantize_row_iq4_nl_impl(QK4_NL, 32, src + QK4_NL*ibl, &iq4[ibl].d, iq4[ibl].qs, &unused_h, unused_l, - &scale, weight, L, kvalues_iq4nl, qw, 7); - } - src += n_per_row; - qrow += nblock*sizeof(block_iq4_nl); - } - return nrow * nblock * sizeof(block_iq4_nl); -} - -//void quantize_row_iq4_nl_ref(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { -void quantize_row_iq4_nl_ref(const float * GGML_RESTRICT x, block_iq4_nl * GGML_RESTRICT y, int64_t k) { - GGML_ASSERT(k%QK4_NL == 0); - int64_t nblock = k/QK4_NL; - uint8_t L[QK4_NL]; - float weight[QK4_NL]; - uint16_t unused_h; - uint8_t * unused_l = NULL; - float scale; - block_iq4_nl * iq4 = y; - for (int ibl = 0; ibl < nblock; ++ibl) { - quantize_row_iq4_nl_impl(QK4_NL, 32, x + QK4_NL*ibl, &iq4[ibl].d, iq4[ibl].qs, &unused_h, unused_l, - &scale, weight, L, kvalues_iq4nl, NULL, -1); - } -} - -size_t quantize_iq4_xs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - GGML_ASSERT(n_per_row%QK_K == 0); - int64_t nblock = n_per_row/QK_K; - char * qrow = (char *)dst; - uint8_t L[QK_K]; - float weight[32]; - float scales[QK_K/32]; - for (int64_t row = 0; row < nrow; ++row) { - block_iq4_xs * iq4 = (block_iq4_xs *)qrow; - for (int ibl = 0; ibl < nblock; ++ibl) { - const float * qw = quant_weights ? quant_weights + QK_K*ibl : NULL; - quantize_row_iq4_nl_impl(QK_K, 32, src + QK_K*ibl, &iq4[ibl].d, iq4[ibl].qs, &iq4[ibl].scales_h, iq4[ibl].scales_l, - scales, weight, L, kvalues_iq4nl, qw, 7); - } - src += n_per_row; - qrow += nblock*sizeof(block_iq4_xs); - } - return nrow * nblock * sizeof(block_iq4_xs); -} - -void quantize_row_iq4_xs_ref(const float * GGML_RESTRICT x, block_iq4_xs * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - quantize_iq4_xs(x, y, 1, k, NULL); -} - -// =============================== 2.5625 bpw - -static void quantize_row_iq2_s_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights) { - - const int gindex = iq2_data_index(GGML_TYPE_IQ2_S); - - const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; - const int * kmap_q2xs = iq2_data[gindex].map; - const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; - - GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); - GGML_ASSERT(n%QK_K == 0); - - const int kMaxQ = 3; - - const int64_t nbl = n/QK_K; - - block_iq2_s * y = vy; - - float scales[QK_K/16]; - float weight[16]; - float xval[16]; - int8_t L[16]; - int8_t Laux[16]; - float waux[16]; - bool is_on_grid[2]; - bool is_on_grid_aux[2]; - uint8_t block_signs[2]; - - for (int ibl = 0; ibl < nbl; ++ibl) { - - memset(&y[ibl], 0, sizeof(block_iq2_s)); - y[ibl].d = GGML_FP32_TO_FP16(0.f); - - float max_scale = 0; - - const float * xbl = x + QK_K*ibl; - float sumx2 = 0; - for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; - float sigma2 = 2*sumx2/QK_K; - - for (int ib = 0; ib < QK_K/16; ++ib) { - const float * xb = xbl + 16*ib; - if (quant_weights) { - const float * qw = quant_weights + QK_K*ibl + 16*ib; - for (int i = 0; i < 16; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); - } else { - for (int i = 0; i < 16; ++i) weight[i] = 0.25f*sigma2 + xb[i]*xb[i]; - } - for (int i = 0; i < 16; ++i) waux[i] = sqrtf(weight[i]); - for (int k = 0; k < 2; ++k) { - uint8_t s = 0; - for (int i = 0; i < 8; ++i) { - if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; - else { - xval[8*k + i] = -xb[8*k + i]; s |= (1 << i); - } - } - block_signs[k] = s; - } - float max = xval[0]; - for (int i = 1; i < 16; ++i) max = MAX(max, xval[i]); - memset(L, 0, 16); - if (max < GROUP_MAX_EPS_IQ2_S) { - scales[ib] = 0; - continue; - } - float best = 0; - float scale = max/(2*kMaxQ-1); - is_on_grid[0] = is_on_grid[1] = true; - for (int is = -9; is <= 9; ++is) { - float id = (2*kMaxQ-1+is*0.1f)/max; - float this_scale = 1/id; - for (int k = 0; k < 2; ++k) { - for (int i = 0; i < 8; ++i) { - int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); - Laux[8*k+i] = MAX(0, MIN(kMaxQ-1, l)); - } - uint16_t u = 0; - for (int i = 0; i < 8; ++i) u |= (Laux[8*k+i] << 2*i); - int grid_index = kmap_q2xs[u]; - is_on_grid_aux[k] = true; - if (grid_index < 0) { - is_on_grid_aux[k] = false; - const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; - grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, this_scale, Laux + 8*k); - } - } - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < 16; ++i) { - float w = weight[i]; - float q = 2*Laux[i] + 1; - sumqx += w*xval[i]*q; - sumq2 += w*q*q; - } - if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { - scale = sumqx/sumq2; best = scale*sumqx; - for (int i = 0; i < 16; ++i) L[i] = Laux[i]; - for (int k = 0; k < 2; ++k) is_on_grid[k] = is_on_grid_aux[k]; - } - } - int n_not_ongrid = 0; - for (int k = 0; k < 2; ++k) if (!is_on_grid[k]) ++n_not_ongrid; - if (n_not_ongrid > 0 && scale > 0) { - float id = 1/scale; - for (int k = 0; k < 2; ++k) { - if (is_on_grid[k]) continue; - uint16_t u = 0; - for (int i = 0; i < 8; ++i) { - int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); - l = MAX(0, MIN(kMaxQ-1, l)); - u |= (l << 2*i); - L[8*k + i] = l; - } - int grid_index = kmap_q2xs[u]; - if (grid_index < 0) { - const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; - grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, scale, L + 8*k); - } - } - float sumqx = 0, sumq2 = 0; - for (int i = 0; i < 16; ++i) { - float w = weight[i]; - float q = 2*L[i] + 1; - sumqx += w*xval[i]*q; - sumq2 += w*q*q; - } - if (sumq2 > 0) scale = sumqx/sumq2; - } - if (scale < 0) { - scale = -scale; - for (int k = 0; k < 2; ++k) block_signs[k] = ~block_signs[k]; - } - for (int k = 0; k < 2; ++k) { - uint16_t u = 0; - for (int i = 0; i < 8; ++i) u |= (L[8*k+i] << 2*i); - int grid_index = kmap_q2xs[u]; - if (grid_index < 0) { - printf("Oops: found point %u not on grid:", u); - for (int i = 0; i < 8; ++i) printf(" %d", L[8*k+i]); - printf("\n"); - GGML_ABORT("fatal error"); - } - const int i8 = 2*ib + k; - y[ibl].qs[i8] = grid_index & 255; - y[ibl].qh[i8/4] |= ((grid_index >> 8) << 2*(i8%4)); - y[ibl].qs[QK_K/8 + i8] = block_signs[k]; - } - GGML_ASSERT(scale >= 0); - scales[ib] = scale; - max_scale = MAX(max_scale, scale); - } - - if (!max_scale) { - continue; - } - - float d = max_scale/31; - y[ibl].d = GGML_FP32_TO_FP16(d * 0.9875f); - float id = 1/d; - for (int ib = 0; ib < QK_K/16; ++ib) { - int l = nearest_int(0.5f*(id*scales[ib]-1)); - l = MAX(0, MIN(15, l)); - if (ib%2 == 0) y[ibl].scales[ib/2] = l; - else y[ibl].scales[ib/2] |= (l << 4); - } - } -} - -size_t quantize_iq2_s(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { - GGML_ASSERT(n_per_row%QK_K == 0); - int64_t nblock = n_per_row/QK_K; - char * qrow = (char *)dst; - for (int64_t row = 0; row < nrow; ++row) { - quantize_row_iq2_s_impl(src, qrow, n_per_row, quant_weights); - src += n_per_row; - qrow += nblock*sizeof(block_iq2_s); - } - return nrow * nblock * sizeof(block_iq2_s); -} - -void quantize_row_iq2_s_ref(const float * GGML_RESTRICT x, block_iq2_s * GGML_RESTRICT y, int64_t k) { - assert(k % QK_K == 0); - quantize_iq2_s(x, y, 1, k, NULL); -} - -// =============================== data validation - -static bool validate_float(float f, size_t i) { - if (isinf(f)) { - fprintf(stderr, "ggml_validate_row_data: found inf value at block %zu\n", i); - return false; - } - - if (isnan(f)) { - fprintf(stderr, "ggml_validate_row_data: found nan value at block %zu\n", i); - return false; - } - - return true; -} - -static bool isinf_fp16(ggml_fp16_t f) { - return (f & 0x7c00) == 0x7c00 && (f & 0x03ff) == 0; -} - -static bool isnan_fp16(ggml_fp16_t f) { - return (f & 0x7c00) == 0x7c00 && (f & 0x03ff) != 0; -} - -static bool validate_fp16(ggml_fp16_t f, size_t i) { - if (isinf_fp16(f)) { - fprintf(stderr, "ggml_validate_row_data: found inf value at block %zu\n", i); - return false; - } - - if (isnan_fp16(f)) { - fprintf(stderr, "ggml_validate_row_data: found nan value at block %zu\n", i); - return false; - } - - return true; -} - -static bool validate_e_e8m0(uint8_t e, size_t i) { - if (e == 0xff) { - fprintf(stderr, "ggml_validate_row_data: found invalid e value %d at block %zu\n", e, i); - return false; - } - - return true; -} - -#define VALIDATE_ROW_DATA_D_F16_IMPL(type, data, nb) \ - const type * q = (const type *) (data); \ - for (size_t i = 0; i < (nb); ++i) { \ - if (!validate_fp16(q[i].d, i)) { \ - return false; \ - } \ - } - -#define VALIDATE_ROW_DATA_DM_F16_IMPL(type, data, nb, d, m) \ - const type * q = (const type *) (data); \ - for (size_t i = 0; i < (nb); ++i) { \ - if (!validate_fp16(q[i].d, i) || !validate_fp16(q[i].m, i)) { \ - return false; \ - } \ - } - -#define VALIDATE_ROW_DATA_E_E8M0_IMPL(type, data, nb) \ - const type * q = (const type *) (data); \ - for (size_t i = 0; i < (nb); ++i) { \ - if (!validate_e_e8m0(q[i].e, i)) { \ - return false; \ - } \ - } - -#define VALIDATE_ROW_DATA_DVEC_F16_IMPL(type, data, nb, nr) \ - const type * q = (const type *) (data); \ - for (size_t i = 0; i < (nb); ++i) { \ - for (size_t j = 0; j < (nr); ++j) { \ - if (!validate_fp16(q[i].d[j], i)) { \ - return false; \ - } \ - } \ - } - -bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbytes) { - if (type < 0 || type >= GGML_TYPE_COUNT) { - fprintf(stderr, "%s: invalid type %d\n", __func__, type); - return false; - } - - if (nbytes % ggml_type_size(type) != 0) { - fprintf(stderr, "%s: invalid size %zu for type %s (type size = %zu)\n", __func__, nbytes, ggml_type_name(type), ggml_type_size(type)); - return false; - } - - const size_t nb = nbytes/ggml_type_size(type); - - switch (type) { - case GGML_TYPE_BF16: - { - int nans = 0; - int infs = 0; - const unsigned short * f = (const unsigned short *) data; - for (size_t i = 0; i < nb; ++i) { - nans += (f[i] & 0x7fff) > 0x7f80; - infs += (f[i] & 0x7fff) == 0x7f80; - } - if (nans) { - fprintf(stderr, "%s: found %d NaNs in row of %zu BF16 values\n", __func__, nans, nb); - return false; - } - if (infs) { - fprintf(stderr, "%s: found %d infinities in row of %zu BF16 values\n", __func__, infs, nb); - return false; - } - } break; - case GGML_TYPE_F16: - { - const ggml_fp16_t * f = (const ggml_fp16_t *) data; - size_t i = 0; -#if defined(__AVX2__) - for (; i + 15 < nb; i += 16) { - __m256i v = _mm256_loadu_si256((const __m256i *)(f + i)); - __m256i vexp = _mm256_and_si256(v, _mm256_set1_epi16(0x7c00)); - __m256i cmp = _mm256_cmpeq_epi16(vexp, _mm256_set1_epi16(0x7c00)); - int mask = _mm256_movemask_epi8(cmp); - if (mask) { - for (size_t j = 0; j < 16; ++j) { - if (!validate_fp16(f[i + j], i + j)) { - return false; - } - } - GGML_UNREACHABLE(); - } - } -#elif defined(__ARM_NEON) - for (; i + 7 < nb; i += 8) { - uint16x8_t v = vld1q_u16(f + i); - uint16x8_t vexp = vandq_u16(v, vdupq_n_u16(0x7c00)); - uint16x8_t cmp = vceqq_u16(vexp, vdupq_n_u16(0x7c00)); - uint64_t mask = vget_lane_u64(vreinterpret_u64_u8(vshrn_n_u16(cmp, 4)), 0); - if (mask) { - for (size_t j = 0; j < 8; ++j) { - if (!validate_fp16(f[i + j], i + j)) { - return false; - } - } - GGML_UNREACHABLE(); - } - } -#endif - for (; i < nb; ++i) { - if (!validate_fp16(f[i], i)) { - return false; - } - } - } break; - case GGML_TYPE_F32: - { - const float * f = (const float *) data; - size_t i = 0; -#if defined(__AVX2__) - for (; i + 7 < nb; i += 8) { - __m256i v = _mm256_loadu_si256((const __m256i *)(f + i)); - __m256i vexp = _mm256_and_si256(v, _mm256_set1_epi32(0x7f800000)); - __m256i cmp = _mm256_cmpeq_epi32(vexp, _mm256_set1_epi32(0x7f800000)); - int mask = _mm256_movemask_epi8(cmp); - if (mask) { - for (size_t j = 0; j < 8; ++j) { - if (!validate_float(f[i + j], i + j)) { - return false; - } - } - GGML_UNREACHABLE(); - } - } -#elif defined(__ARM_NEON) - for (; i + 3 < nb; i += 4) { - uint32x4_t v = vld1q_u32((const uint32_t *)f + i); - uint32x4_t vexp = vandq_u32(v, vdupq_n_u32(0x7f800000)); - uint32x4_t cmp = vceqq_u32(vexp, vdupq_n_u32(0x7f800000)); - uint64_t mask = vget_lane_u64(vreinterpret_u64_u16(vshrn_n_u32(cmp, 8)), 0); - if (mask) { - for (size_t j = 0; j < 4; ++j) { - if (!validate_float(f[i + j], i + j)) { - return false; - } - } - GGML_UNREACHABLE(); - } - } -#endif - for (; i < nb; ++i) { - if (!validate_float(f[i], i)) { - return false; - } - } - } break; - case GGML_TYPE_F64: - { - const double * f = (const double *) data; - for (size_t i = 0; i < nb; ++i) { - if (!validate_float(f[i], i)) { - return false; - } - } - } break; - case GGML_TYPE_Q1_0: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_q1_0, data, nb); - } break; - case GGML_TYPE_Q4_0: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_q4_0, data, nb); - } break; - case GGML_TYPE_Q4_1: - { - VALIDATE_ROW_DATA_DM_F16_IMPL(block_q4_1, data, nb, d, m); - } break; - case GGML_TYPE_Q5_0: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_q5_0, data, nb); - } break; - case GGML_TYPE_Q5_1: - { - VALIDATE_ROW_DATA_DM_F16_IMPL(block_q5_1, data, nb, d, m); - } break; - case GGML_TYPE_Q8_0: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_q8_0, data, nb); - } break; - case GGML_TYPE_MXFP4: - { - VALIDATE_ROW_DATA_E_E8M0_IMPL(block_mxfp4, data, nb); - } break; - case GGML_TYPE_NVFP4: - { - // UE4M3 scales are uint8_t — all byte values are valid - GGML_UNUSED(data); - GGML_UNUSED(nb); - } break; - case GGML_TYPE_Q2_K: - { - VALIDATE_ROW_DATA_DM_F16_IMPL(block_q2_K, data, nb, d, dmin); - } break; - case GGML_TYPE_Q3_K: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_q3_K, data, nb); - } break; - case GGML_TYPE_Q4_K: - { - VALIDATE_ROW_DATA_DM_F16_IMPL(block_q4_K, data, nb, d, dmin); - } break; - case GGML_TYPE_Q5_K: - { - VALIDATE_ROW_DATA_DM_F16_IMPL(block_q5_K, data, nb, d, dmin); - } break; - case GGML_TYPE_Q6_K: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_q6_K, data, nb); - } break; - case GGML_TYPE_Q8_K: - { - const block_q8_K * q = (const block_q8_K *) data; - for (size_t i = 0; i < nb; ++i) { - if (!validate_float(q[i].d, i)) { - return false; - } - } - } break; - case GGML_TYPE_TQ1_0: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_tq1_0, data, nb); - } break; - case GGML_TYPE_TQ2_0: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_tq2_0, data, nb); - } break; - case GGML_TYPE_IQ1_S: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_iq1_s, data, nb); - } break; - case GGML_TYPE_IQ1_M: - { - const block_iq1_m * q = (const block_iq1_m *) data; - for (size_t i = 0; i < nb; ++i) { - iq1m_scale_t scale; - const uint16_t * sc = (const uint16_t *)q[i].scales; - scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); - if (!validate_fp16(scale.f16, i)) { - return false; - } - } - } break; - case GGML_TYPE_IQ2_XXS: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_iq2_xxs, data, nb); - } break; - case GGML_TYPE_IQ2_XS: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_iq2_xs, data, nb); - } break; - case GGML_TYPE_IQ2_S: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_iq2_s, data, nb); - } break; - case GGML_TYPE_IQ3_XXS: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_iq3_xxs, data, nb); - } break; - - case GGML_TYPE_IQ3_S: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_iq3_s, data, nb); - } break; - case GGML_TYPE_IQ4_XS: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_iq4_xs, data, nb); - } break; - case GGML_TYPE_IQ4_NL: - { - VALIDATE_ROW_DATA_D_F16_IMPL(block_iq4_nl, data, nb); - } break; - - case GGML_TYPE_I8: - case GGML_TYPE_I16: - case GGML_TYPE_I32: - case GGML_TYPE_I64: - // nothing to validate - break; - default: - { - fprintf(stderr, "%s: invalid type %d\n", __func__, type); - return false; - } - } - - return true; -} diff --git a/ggml/src/ggml-rpc/CMakeLists.txt b/ggml/src/ggml-rpc/CMakeLists.txt index 40e11fead63a..18165d28389e 100644 --- a/ggml/src/ggml-rpc/CMakeLists.txt +++ b/ggml/src/ggml-rpc/CMakeLists.txt @@ -1,8 +1,8 @@ message(STATUS "Using RPC backend") +set_source_files_properties(ggml-rpc.cpp.inc PROPERTIES LANGUAGE CXX) ggml_add_backend_library(ggml-rpc - ggml-rpc.cpp - transport.cpp + ggml-rpc.cpp.inc ) if (WIN32) diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp.inc similarity index 77% rename from ggml/src/ggml-rpc/ggml-rpc.cpp rename to ggml/src/ggml-rpc/ggml-rpc.cpp.inc index 7176d2feef9a..f2af5bb4d982 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp.inc @@ -1,8 +1,39 @@ -#include "ggml-rpc.h" -#include "ggml-impl.h" -#include "ggml-backend-impl.h" -#include "ggml-cpp.h" -#include "transport.h" +#include "ggml-rpc.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-backend-impl-defs.inc" +#include +#include +#include + +struct socket_t; +typedef std::shared_ptr socket_ptr; + +static constexpr size_t MAX_CHUNK_SIZE = 1024ull * 1024ull * 1024ull; // 1 GiB +static constexpr size_t RPC_CONN_CAPS_SIZE = 24; + +struct socket_t { + ~socket_t(); + + bool send_data(const void * data, size_t size); + bool recv_data(void * data, size_t size); + + socket_ptr accept(); + + void get_caps(uint8_t * local_caps); + void update_caps(const uint8_t * remote_caps); + + static socket_ptr create_server(const char * host, int port); + static socket_ptr connect(const char * host, int port); + +private: + struct impl; + explicit socket_t(std::unique_ptr p); + std::unique_ptr pimpl; +}; + +bool rpc_transport_init(); +void rpc_transport_shutdown(); + #include #include @@ -26,6 +57,12 @@ static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG"); namespace fs = std::filesystem; +struct ggml_context_deleter { + void operator()(ggml_context * ctx) { ggml_free(ctx); } +}; + +using ggml_context_ptr = std::unique_ptr; + // macro for nicer error messages on server crash #define RPC_STATUS_ASSERT(x) if (!(x)) GGML_ABORT("Remote RPC server crashed or returned malformed response") @@ -1996,3 +2033,683 @@ ggml_backend_reg_t ggml_backend_rpc_add_server(const char * endpoint) { GGML_BACKEND_DL_IMPL(ggml_backend_rpc_reg) + + +/* Inlined transport.cpp */ +#include "ggml-impl-defs.inc" + +#ifdef _WIN32 +# define WIN32_LEAN_AND_MEAN +# ifndef NOMINMAX +# define NOMINMAX +# endif +# include +# include +#else +# include +# include +# include +# include +# include +# include +# include +#endif +#include +#include +#include + +#ifdef GGML_RPC_RDMA +# include +# include +# ifndef _WIN32 +# include +# endif +#endif // GGML_RPC_RDMA + +#ifdef _WIN32 +typedef SOCKET sockfd_t; +using ssize_t = __int64; +#else +typedef int sockfd_t; +#endif + +#ifdef GGML_RPC_RDMA +static constexpr size_t RDMA_CHUNK = 256 * 1024; // 256 KiB per send/recv (fits default 8 MiB memlock) +static constexpr int RDMA_RX_DEPTH = 24; // pre-posted recv ring: 24 × 256 KiB = 6 MiB +static constexpr size_t RDMA_GID_SIZE = 16; // RoCE GID / IB GID is always 16 bytes +using rdma_gid_t = std::array; + +struct rdma_conn { + struct ibv_context * ctx = nullptr; + struct ibv_pd * pd = nullptr; + struct ibv_cq * scq = nullptr; // send completions + struct ibv_cq * rcq = nullptr; // recv completions + struct ibv_qp * qp = nullptr; + + void * tx_buf = nullptr; + struct ibv_mr * tx_mr = nullptr; + + void * rx_buf = nullptr; // RDMA_RX_DEPTH × RDMA_CHUNK contiguous + struct ibv_mr * rx_mr = nullptr; + int rx_head = 0; + + uint32_t max_inline = 0; + + uint8_t * rx_slot(int i) const { + return static_cast(rx_buf) + static_cast(i) * RDMA_CHUNK; + } + + bool post_rx(int i) { + struct ibv_sge sge = {}; + sge.addr = (uintptr_t)rx_slot(i); + sge.length = RDMA_CHUNK; + sge.lkey = rx_mr->lkey; + struct ibv_recv_wr wr = {}, * bad = nullptr; + wr.wr_id = (uint64_t)i; + wr.sg_list = &sge; + wr.num_sge = 1; + return ibv_post_recv(qp, &wr, &bad) == 0; + } + + ~rdma_conn() { + if (tx_mr) ibv_dereg_mr(tx_mr); + if (rx_mr) ibv_dereg_mr(rx_mr); + free(tx_buf); + free(rx_buf); + if (qp) ibv_destroy_qp(qp); + if (scq) ibv_destroy_cq(scq); + if (rcq) ibv_destroy_cq(rcq); + if (pd) ibv_dealloc_pd(pd); + if (ctx) ibv_close_device(ctx); + } +}; + +// Local RDMA parameters captured during the probe phase and later consumed +// by rdma_activate() after the remote side's caps arrive via HELLO. +struct rdma_local_info { + uint32_t qpn = 0; + uint32_t psn = 0; + uint8_t gid[RDMA_GID_SIZE] = {}; + uint8_t ib_port = 0; + int gid_idx = 0; + enum ibv_mtu path_mtu = IBV_MTU_1024; +}; + +struct rdma_caps { + uint32_t qpn; + uint32_t psn; + uint8_t gid[RDMA_GID_SIZE]; +}; + +static_assert(sizeof(rdma_caps) == RPC_CONN_CAPS_SIZE, "rdma_caps must match conn_caps size"); + +#endif // GGML_RPC_RDMA + +struct socket_t::impl { + impl(sockfd_t fd) : use_rdma(false), fd(fd) {} + ~impl(); + bool send_data(const void * data, size_t size); + bool recv_data(void * data, size_t size); + void get_caps(uint8_t * local_caps); + void update_caps(const uint8_t * remote_caps); + +#ifdef GGML_RPC_RDMA + bool tcp_peer_closed(); + std::optional rdma_build_target_gid(); + bool rdma_probe(); + bool rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid); + bool rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc); + bool rdma_send(const void * data, size_t size); + bool rdma_recv(void * data, size_t size); + + std::unique_ptr rdma; + rdma_local_info rdma_local = {}; +#endif // GGML_RPC_RDMA + bool use_rdma; + sockfd_t fd; +}; + +socket_t::impl::~impl() { +#ifdef GGML_RPC_RDMA + rdma.reset(); +#endif // GGML_RPC_RDMA + LOG_DBG("[%s] closing socket %d\n", __func__, this->fd); +#ifdef _WIN32 + if (fd != INVALID_SOCKET) closesocket(this->fd); +#else + if (fd >= 0) close(this->fd); +#endif +} + +#ifdef GGML_RPC_RDMA + +bool socket_t::impl::tcp_peer_closed() { + if (fd < 0) return false; +#ifndef _WIN32 + struct pollfd pfd = { fd, POLLIN | POLLRDHUP, 0 }; + int r = poll(&pfd, 1, 0); + return r > 0 && (pfd.revents & (POLLHUP | POLLERR | POLLRDHUP)); +#else + return false; +#endif +} + +// Build a RoCE GID-shaped 16-byte target from a TCP socket's local address. +// Used to match the socket's local IP against the kernel's GID table so that +// a single memcmp handles IPv4, IPv4-mapped IPv6, and native IPv6 uniformly: +// AF_INET -> ::ffff:a.b.c.d (bytes 10-11 = 0xff, last 4 = IPv4) +// AF_INET6 (IPv4-mapped) -> ::ffff:a.b.c.d (already in GID shape) +// AF_INET6 (native v6) -> the 16-byte IPv6 address as-is +// Returns std::nullopt on unsupported family or getsockname failure. +std::optional socket_t::impl::rdma_build_target_gid() { + sockaddr_storage addr = {}; + socklen_t addr_len = sizeof(addr); + if (getsockname(fd, reinterpret_cast(&addr), &addr_len) != 0) { + return std::nullopt; + } + rdma_gid_t target = {}; + if (addr.ss_family == AF_INET) { + const auto * a = reinterpret_cast(&addr); + target[10] = 0xff; + target[11] = 0xff; + memcpy(&target[12], &a->sin_addr, 4); + return target; + } + if (addr.ss_family == AF_INET6) { + const auto * a = reinterpret_cast(&addr); + memcpy(target.data(), &a->sin6_addr, RDMA_GID_SIZE); + return target; + } + return std::nullopt; +} + +bool socket_t::impl::rdma_probe() { + const char * dev_env = std::getenv("GGML_RDMA_DEV"); + const char * gid_env = std::getenv("GGML_RDMA_GID"); + + auto target_gid = rdma_build_target_gid(); + if (!target_gid) { + return false; + } + + const uint8_t ib_port = 1; + int num_devs = 0; + ibv_device ** devs = ibv_get_device_list(&num_devs); + if (!devs || num_devs == 0) return false; + + ibv_context * ibctx = nullptr; + const char * matched_dev = nullptr; + int gid_idx = gid_env ? atoi(gid_env) : -1; + int gid_version = IBV_GID_TYPE_IB; // 0 = unknown/IB + + for (int d = 0; d < num_devs; d++) { + const char * dn = ibv_get_device_name(devs[d]); + if (dev_env && strcmp(dev_env, dn) != 0) continue; + + ibv_context * ctx = ibv_open_device(devs[d]); + if (!ctx) continue; + + ibv_port_attr pa; + if (ibv_query_port(ctx, ib_port, &pa) != 0) { ibv_close_device(ctx); continue; } + + int found_gid = gid_idx; + int found_version = IBV_GID_TYPE_IB; + if (found_gid < 0) { + // Find a GID on this port whose bytes equal the local TCP address + // (IPv4 or IPv6). Prefer RoCE v2 (UDP/IP, L3-routable) over v1 + // (raw Ethernet, same-L2 only) so silent hangs on L3-routed paths + // are avoided. ibv_query_gid_ex returns gid+type in one call. + int v2_idx = -1; + int v1_idx = -1; + for (int i = 0; i < pa.gid_tbl_len; i++) { + ibv_gid_entry entry = {}; + if (ibv_query_gid_ex(ctx, ib_port, i, &entry, 0) != 0) continue; + if (memcmp(entry.gid.raw, target_gid->data(), RDMA_GID_SIZE) != 0) continue; + if (entry.gid_type == IBV_GID_TYPE_ROCE_V2 && v2_idx < 0) { + v2_idx = i; + } else if (entry.gid_type == IBV_GID_TYPE_ROCE_V1 && v1_idx < 0) { + v1_idx = i; + } + } + if (v2_idx >= 0) { + found_gid = v2_idx; + found_version = IBV_GID_TYPE_ROCE_V2; + } else if (v1_idx >= 0) { + found_gid = v1_idx; + found_version = IBV_GID_TYPE_ROCE_V1; + } + } else { + // Explicit GID index from GGML_RDMA_GID — fetch its type for logging. + ibv_gid_entry entry = {}; + if (ibv_query_gid_ex(ctx, ib_port, found_gid, &entry, 0) == 0) { + found_version = entry.gid_type; + } + } + if (found_gid >= 0) { + ibctx = ctx; + gid_idx = found_gid; + gid_version = found_version; + matched_dev = dn; + rdma_local.path_mtu = pa.active_mtu; + break; + } + ibv_close_device(ctx); + } + ibv_free_device_list(devs); + if (!ibctx) return false; + + rdma_local.ib_port = ib_port; + rdma_local.gid_idx = gid_idx; + + rdma = std::make_unique(); + rdma->ctx = ibctx; + + rdma->pd = ibv_alloc_pd(ibctx); + if (!rdma->pd) return false; + + rdma->scq = ibv_create_cq(ibctx, 16, nullptr, nullptr, 0); + rdma->rcq = ibv_create_cq(ibctx, RDMA_RX_DEPTH + 4, nullptr, nullptr, 0); + if (!rdma->scq || !rdma->rcq) return false; + + ibv_qp_init_attr qia = {}; + qia.send_cq = rdma->scq; + qia.recv_cq = rdma->rcq; + qia.qp_type = IBV_QPT_RC; + qia.cap.max_send_wr = 4; + qia.cap.max_recv_wr = RDMA_RX_DEPTH + 4; + qia.cap.max_send_sge = 1; + qia.cap.max_recv_sge = 1; + qia.cap.max_inline_data = 256; + + rdma->qp = ibv_create_qp(rdma->pd, &qia); + if (!rdma->qp) return false; + rdma->max_inline = qia.cap.max_inline_data; + + rdma->tx_buf = aligned_alloc(4096, RDMA_CHUNK); + rdma->rx_buf = aligned_alloc(4096, static_cast(RDMA_RX_DEPTH) * RDMA_CHUNK); + if (!rdma->tx_buf || !rdma->rx_buf) return false; + + rdma->tx_mr = ibv_reg_mr(rdma->pd, rdma->tx_buf, RDMA_CHUNK, IBV_ACCESS_LOCAL_WRITE); + rdma->rx_mr = ibv_reg_mr(rdma->pd, rdma->rx_buf, static_cast(RDMA_RX_DEPTH) * RDMA_CHUNK, + IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); + if (!rdma->tx_mr || !rdma->rx_mr) return false; + + ibv_gid local_gid; + if (ibv_query_gid(ibctx, ib_port, gid_idx, &local_gid) != 0) return false; + + rdma_local.qpn = rdma->qp->qp_num; + rdma_local.psn = rdma->qp->qp_num & 0xffffff; + memcpy(&rdma_local.gid, &local_gid, RDMA_GID_SIZE); + + const char * ver_str = ""; + if (gid_version == IBV_GID_TYPE_ROCE_V2) { + ver_str = " RoCEv2"; + } else if (gid_version == IBV_GID_TYPE_ROCE_V1) { + ver_str = " RoCEv1"; + } + GGML_LOG_INFO("RDMA probed: dev=%s gid=%d%s qpn=%u inline=%u\n", + matched_dev, gid_idx, ver_str, rdma_local.qpn, rdma->max_inline); + return true; +} + +// Phase 2: Given remote QPN/PSN/GID, transition QP: RESET->INIT->pre-post->RTR->RTS. +// On success, the connection is live and ready for rdma_send/rdma_recv. +bool socket_t::impl::rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid) { + // RESET -> INIT + { + struct ibv_qp_attr a = {}; + a.qp_state = IBV_QPS_INIT; + a.port_num = rdma_local.ib_port; + a.pkey_index = 0; + a.qp_access_flags = IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_LOCAL_WRITE; + if (ibv_modify_qp(rdma->qp, &a, + IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS) != 0) { + return false; + } + } + + for (int i = 0; i < RDMA_RX_DEPTH; i++) { + if (!rdma->post_rx(i)) return false; + } + + // INIT -> RTR + { + struct ibv_qp_attr a = {}; + a.qp_state = IBV_QPS_RTR; + a.path_mtu = rdma_local.path_mtu; + a.dest_qp_num = remote_qpn; + a.rq_psn = remote_psn; + a.max_dest_rd_atomic = 1; + a.min_rnr_timer = 1; + a.ah_attr.is_global = 1; + memcpy(&a.ah_attr.grh.dgid, remote_gid, RDMA_GID_SIZE); + a.ah_attr.grh.hop_limit = 1; + a.ah_attr.grh.sgid_index = rdma_local.gid_idx; + a.ah_attr.dlid = 0; + a.ah_attr.port_num = rdma_local.ib_port; + if (ibv_modify_qp(rdma->qp, &a, + IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN | + IBV_QP_RQ_PSN | IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER) != 0) { + return false; + } + } + + // RTR -> RTS + { + struct ibv_qp_attr a = {}; + a.qp_state = IBV_QPS_RTS; + a.timeout = 14; + a.retry_cnt = 7; + a.rnr_retry = 7; + a.sq_psn = rdma_local.psn; + a.max_rd_atomic = 1; + if (ibv_modify_qp(rdma->qp, &a, + IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT | IBV_QP_RNR_RETRY | + IBV_QP_SQ_PSN | IBV_QP_MAX_QP_RD_ATOMIC) != 0) { + return false; + } + } + + GGML_LOG_INFO("RDMA activated: qpn=%u->%u mtu=%d rx_depth=%d\n", + rdma_local.qpn, remote_qpn, 128 << rdma_local.path_mtu, RDMA_RX_DEPTH); + return true; +} + +bool socket_t::impl::rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc) { + for (uint64_t s = 0; ; s++) { + int n = ibv_poll_cq(cq, 1, wc); + if (n > 0) { + if (wc->status != IBV_WC_SUCCESS) { + GGML_LOG_ERROR("RDMA CQ wc error: status=%d (%s) vendor_err=0x%x\n", + wc->status, ibv_wc_status_str(wc->status), wc->vendor_err); + } + return wc->status == IBV_WC_SUCCESS; + } + if (n < 0) return false; + if ((s & 0xFFFFF) == 0 && s > 0) { + if (tcp_peer_closed()) { + return false; + } + } + } +} + +bool socket_t::impl::rdma_send(const void * data, size_t size) { + rdma_conn * c = rdma.get(); + const uint8_t * src = (const uint8_t *)data; + size_t rem = size; + while (rem > 0) { + size_t chunk = std::min(rem, RDMA_CHUNK); + + struct ibv_sge sge = {}; + struct ibv_send_wr wr = {}, * bad = nullptr; + wr.opcode = IBV_WR_SEND; + wr.sg_list = &sge; + wr.num_sge = 1; + + if (chunk <= c->max_inline) { + sge.addr = (uintptr_t)src; + sge.length = chunk; + wr.send_flags = IBV_SEND_SIGNALED | IBV_SEND_INLINE; + } else { + memcpy(c->tx_buf, src, chunk); + sge.addr = (uintptr_t)c->tx_buf; + sge.length = chunk; + sge.lkey = c->tx_mr->lkey; + wr.send_flags = IBV_SEND_SIGNALED; + } + + if (ibv_post_send(c->qp, &wr, &bad) != 0) return false; + struct ibv_wc wc; + if (!rdma_poll(c->scq, &wc)) return false; + + src += chunk; + rem -= chunk; + } + return true; +} + +bool socket_t::impl::rdma_recv(void * data, size_t size) { + rdma_conn * c = rdma.get(); + uint8_t * dst = (uint8_t *)data; + size_t rem = size; + while (rem > 0) { + struct ibv_wc wc; + if (!rdma_poll(c->rcq, &wc)) return false; + + int slot = (int)wc.wr_id; + size_t got = wc.byte_len; + memcpy(dst, c->rx_slot(slot), got); + + if (!c->post_rx(slot)) return false; + + dst += got; + rem -= got; + } + return true; +} + +#endif // GGML_RPC_RDMA + +bool socket_t::impl::send_data(const void * data, size_t size) { +#ifdef GGML_RPC_RDMA + if (use_rdma) { + return rdma_send(data, size); + } +#endif + size_t bytes_sent = 0; + while (bytes_sent < size) { + size_t size_to_send = std::min(size - bytes_sent, MAX_CHUNK_SIZE); + ssize_t n = send(fd, (const char *)data + bytes_sent, size_to_send, 0); + if (n < 0) { + GGML_LOG_ERROR("send failed (bytes_sent=%zu, size_to_send=%zu)\n", + bytes_sent, size_to_send); + return false; + } + bytes_sent += (size_t)n; + } + return true; +} + +bool socket_t::impl::recv_data(void * data, size_t size) { +#ifdef GGML_RPC_RDMA + if (use_rdma) { + return rdma_recv(data, size); + } +#endif + size_t bytes_recv = 0; + while (bytes_recv < size) { + size_t size_to_recv = std::min(size - bytes_recv, MAX_CHUNK_SIZE); + ssize_t n = recv(fd, (char *)data + bytes_recv, size_to_recv, 0); + if (n < 0) { + GGML_LOG_ERROR("recv failed (bytes_recv=%zu, size_to_recv=%zu)\n", + bytes_recv, size_to_recv); + return false; + } + if (n == 0) { + LOG_DBG("recv returned 0 (peer closed?)\n"); + return false; + } + bytes_recv += (size_t)n; + } + return true; +} + +void socket_t::impl::get_caps(uint8_t * local_caps) { + memset(local_caps, 0, RPC_CONN_CAPS_SIZE); +#ifdef GGML_RPC_RDMA + rdma_local = {}; + if (rdma_probe()) { + rdma_caps rc = {}; + rc.qpn = rdma_local.qpn; + rc.psn = rdma_local.psn; + memcpy(rc.gid, rdma_local.gid, RDMA_GID_SIZE); + memcpy(local_caps, &rc, sizeof(rc)); + } else { + rdma.reset(); + } +#endif // GGML_RPC_RDMA +} + +void socket_t::impl::update_caps(const uint8_t * remote_caps) { +#ifdef GGML_RPC_RDMA + if (!rdma) { + return; + } + rdma_caps rc = {}; + memcpy(&rc, remote_caps, sizeof(rc)); + if (rc.qpn == 0) { + rdma.reset(); + return; + } + if (rdma_activate(rc.qpn, rc.psn, rc.gid)) { + use_rdma = true; + } else { + GGML_LOG_ERROR("RDMA activate failed, staying on TCP\n"); + rdma.reset(); + } +#else + (void)remote_caps; +#endif // GGML_RPC_RDMA +} + + +///////////////////////////////////////////////////////////////////////////// + +socket_t::socket_t(std::unique_ptr p) : pimpl(std::move(p)) {} + +socket_t::~socket_t() = default; + +bool socket_t::send_data(const void * data, size_t size) { + return pimpl->send_data(data, size); +} + +bool socket_t::recv_data(void * data, size_t size) { + return pimpl->recv_data(data, size); +} + +void socket_t::get_caps(uint8_t * local_caps) { + return pimpl->get_caps(local_caps); +} + +void socket_t::update_caps(const uint8_t * remote_caps) { + return pimpl->update_caps(remote_caps); +} + +static bool is_valid_fd(sockfd_t sockfd) { +#ifdef _WIN32 + return sockfd != INVALID_SOCKET; +#else + return sockfd >= 0; +#endif +} + +static bool set_no_delay(sockfd_t sockfd) { + int flag = 1; + // set TCP_NODELAY to disable Nagle's algorithm + int ret = setsockopt(sockfd, IPPROTO_TCP, TCP_NODELAY, (char *)&flag, sizeof(int)); + return ret == 0; +} + +static bool set_reuse_addr(sockfd_t sockfd) { + int flag = 1; + int ret = setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, (char *)&flag, sizeof(int)); + return ret == 0; +} + +socket_ptr socket_t::accept() { + auto client_socket_fd = ::accept(pimpl->fd, NULL, NULL); + if (!is_valid_fd(client_socket_fd)) { + return nullptr; + } + if (!set_no_delay(client_socket_fd)) { + GGML_LOG_ERROR("Failed to set TCP_NODELAY\n"); + return nullptr; + } + return socket_ptr(new socket_t(std::make_unique(client_socket_fd))); +} + +socket_ptr socket_t::create_server(const char * host, int port) { + auto sockfd = socket(AF_INET, SOCK_STREAM, 0); + if (!is_valid_fd(sockfd)) { + return nullptr; + } + if (!set_reuse_addr(sockfd)) { + GGML_LOG_ERROR("Failed to set SO_REUSEADDR\n"); + return nullptr; + } + if (inet_addr(host) == INADDR_NONE) { + GGML_LOG_ERROR("Invalid host address: %s\n", host); + return nullptr; + } + struct sockaddr_in serv_addr; + serv_addr.sin_family = AF_INET; + serv_addr.sin_addr.s_addr = inet_addr(host); + serv_addr.sin_port = htons(port); + + if (bind(sockfd, (struct sockaddr *) &serv_addr, sizeof(serv_addr)) < 0) { + return nullptr; + } + if (listen(sockfd, 1) < 0) { + return nullptr; + } + return socket_ptr(new socket_t(std::make_unique(sockfd))); +} + +socket_ptr socket_t::connect(const char * host, int port) { + auto sockfd = socket(AF_INET, SOCK_STREAM, 0); + if (!is_valid_fd(sockfd)) { + return nullptr; + } + if (!set_no_delay(sockfd)) { + GGML_LOG_ERROR("Failed to set TCP_NODELAY\n"); + return nullptr; + } + struct sockaddr_in addr; + addr.sin_family = AF_INET; + addr.sin_port = htons(port); + struct hostent * server = gethostbyname(host); + if (server == NULL) { + GGML_LOG_ERROR("Cannot resolve host '%s'\n", host); + return nullptr; + } + memcpy(&addr.sin_addr.s_addr, server->h_addr, server->h_length); + if (::connect(sockfd, (struct sockaddr *)&addr, sizeof(addr)) < 0) { + return nullptr; + } + return socket_ptr(new socket_t(std::make_unique(sockfd))); +} + +#ifdef _WIN32 +static std::mutex g_rpc_transport_mu; +static bool g_rpc_transport_wsa_started = false; +#endif + +bool rpc_transport_init() { +#ifdef _WIN32 + std::lock_guard lock(g_rpc_transport_mu); + if (g_rpc_transport_wsa_started) { + return true; + } + WSADATA wsaData; + int res = WSAStartup(MAKEWORD(2, 2), &wsaData); + if (res != 0) { + return false; + } + g_rpc_transport_wsa_started = true; + return true; +#else + return true; +#endif +} + +void rpc_transport_shutdown() { +#ifdef _WIN32 + std::lock_guard lock(g_rpc_transport_mu); + if (!g_rpc_transport_wsa_started) { + return; + } + WSACleanup(); + g_rpc_transport_wsa_started = false; +#endif +} diff --git a/ggml/src/ggml-rpc/transport.cpp b/ggml/src/ggml-rpc/transport.cpp deleted file mode 100644 index a728152421f7..000000000000 --- a/ggml/src/ggml-rpc/transport.cpp +++ /dev/null @@ -1,683 +0,0 @@ -#include "transport.h" -#include "ggml-impl.h" - -#ifdef _WIN32 -# define WIN32_LEAN_AND_MEAN -# ifndef NOMINMAX -# define NOMINMAX -# endif -# include -# include -#else -# include -# include -# include -# include -# include -# include -# include -#endif -#include -#include -#include - -#ifdef GGML_RPC_RDMA -# include -# include -# ifndef _WIN32 -# include -# endif -#endif // GGML_RPC_RDMA - -#ifdef _WIN32 -typedef SOCKET sockfd_t; -using ssize_t = __int64; -#else -typedef int sockfd_t; -#endif - -static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG"); - -#define LOG_DBG(...) \ - do { if (RPC_DEBUG) GGML_LOG_DEBUG(__VA_ARGS__); } while (0) - -#ifdef GGML_RPC_RDMA -static constexpr size_t RDMA_CHUNK = 256 * 1024; // 256 KiB per send/recv (fits default 8 MiB memlock) -static constexpr int RDMA_RX_DEPTH = 24; // pre-posted recv ring: 24 × 256 KiB = 6 MiB -static constexpr size_t RDMA_GID_SIZE = 16; // RoCE GID / IB GID is always 16 bytes -using rdma_gid_t = std::array; - -struct rdma_conn { - struct ibv_context * ctx = nullptr; - struct ibv_pd * pd = nullptr; - struct ibv_cq * scq = nullptr; // send completions - struct ibv_cq * rcq = nullptr; // recv completions - struct ibv_qp * qp = nullptr; - - void * tx_buf = nullptr; - struct ibv_mr * tx_mr = nullptr; - - void * rx_buf = nullptr; // RDMA_RX_DEPTH × RDMA_CHUNK contiguous - struct ibv_mr * rx_mr = nullptr; - int rx_head = 0; - - uint32_t max_inline = 0; - - uint8_t * rx_slot(int i) const { - return static_cast(rx_buf) + static_cast(i) * RDMA_CHUNK; - } - - bool post_rx(int i) { - struct ibv_sge sge = {}; - sge.addr = (uintptr_t)rx_slot(i); - sge.length = RDMA_CHUNK; - sge.lkey = rx_mr->lkey; - struct ibv_recv_wr wr = {}, * bad = nullptr; - wr.wr_id = (uint64_t)i; - wr.sg_list = &sge; - wr.num_sge = 1; - return ibv_post_recv(qp, &wr, &bad) == 0; - } - - ~rdma_conn() { - if (tx_mr) ibv_dereg_mr(tx_mr); - if (rx_mr) ibv_dereg_mr(rx_mr); - free(tx_buf); - free(rx_buf); - if (qp) ibv_destroy_qp(qp); - if (scq) ibv_destroy_cq(scq); - if (rcq) ibv_destroy_cq(rcq); - if (pd) ibv_dealloc_pd(pd); - if (ctx) ibv_close_device(ctx); - } -}; - -// Local RDMA parameters captured during the probe phase and later consumed -// by rdma_activate() after the remote side's caps arrive via HELLO. -struct rdma_local_info { - uint32_t qpn = 0; - uint32_t psn = 0; - uint8_t gid[RDMA_GID_SIZE] = {}; - uint8_t ib_port = 0; - int gid_idx = 0; - enum ibv_mtu path_mtu = IBV_MTU_1024; -}; - -struct rdma_caps { - uint32_t qpn; - uint32_t psn; - uint8_t gid[RDMA_GID_SIZE]; -}; - -static_assert(sizeof(rdma_caps) == RPC_CONN_CAPS_SIZE, "rdma_caps must match conn_caps size"); - -#endif // GGML_RPC_RDMA - -struct socket_t::impl { - impl(sockfd_t fd) : use_rdma(false), fd(fd) {} - ~impl(); - bool send_data(const void * data, size_t size); - bool recv_data(void * data, size_t size); - void get_caps(uint8_t * local_caps); - void update_caps(const uint8_t * remote_caps); - -#ifdef GGML_RPC_RDMA - bool tcp_peer_closed(); - std::optional rdma_build_target_gid(); - bool rdma_probe(); - bool rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid); - bool rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc); - bool rdma_send(const void * data, size_t size); - bool rdma_recv(void * data, size_t size); - - std::unique_ptr rdma; - rdma_local_info rdma_local = {}; -#endif // GGML_RPC_RDMA - bool use_rdma; - sockfd_t fd; -}; - -socket_t::impl::~impl() { -#ifdef GGML_RPC_RDMA - rdma.reset(); -#endif // GGML_RPC_RDMA - LOG_DBG("[%s] closing socket %d\n", __func__, this->fd); -#ifdef _WIN32 - if (fd != INVALID_SOCKET) closesocket(this->fd); -#else - if (fd >= 0) close(this->fd); -#endif -} - -#ifdef GGML_RPC_RDMA - -bool socket_t::impl::tcp_peer_closed() { - if (fd < 0) return false; -#ifndef _WIN32 - struct pollfd pfd = { fd, POLLIN | POLLRDHUP, 0 }; - int r = poll(&pfd, 1, 0); - return r > 0 && (pfd.revents & (POLLHUP | POLLERR | POLLRDHUP)); -#else - return false; -#endif -} - -// Build a RoCE GID-shaped 16-byte target from a TCP socket's local address. -// Used to match the socket's local IP against the kernel's GID table so that -// a single memcmp handles IPv4, IPv4-mapped IPv6, and native IPv6 uniformly: -// AF_INET -> ::ffff:a.b.c.d (bytes 10-11 = 0xff, last 4 = IPv4) -// AF_INET6 (IPv4-mapped) -> ::ffff:a.b.c.d (already in GID shape) -// AF_INET6 (native v6) -> the 16-byte IPv6 address as-is -// Returns std::nullopt on unsupported family or getsockname failure. -std::optional socket_t::impl::rdma_build_target_gid() { - sockaddr_storage addr = {}; - socklen_t addr_len = sizeof(addr); - if (getsockname(fd, reinterpret_cast(&addr), &addr_len) != 0) { - return std::nullopt; - } - rdma_gid_t target = {}; - if (addr.ss_family == AF_INET) { - const auto * a = reinterpret_cast(&addr); - target[10] = 0xff; - target[11] = 0xff; - memcpy(&target[12], &a->sin_addr, 4); - return target; - } - if (addr.ss_family == AF_INET6) { - const auto * a = reinterpret_cast(&addr); - memcpy(target.data(), &a->sin6_addr, RDMA_GID_SIZE); - return target; - } - return std::nullopt; -} - -bool socket_t::impl::rdma_probe() { - const char * dev_env = std::getenv("GGML_RDMA_DEV"); - const char * gid_env = std::getenv("GGML_RDMA_GID"); - - auto target_gid = rdma_build_target_gid(); - if (!target_gid) { - return false; - } - - const uint8_t ib_port = 1; - int num_devs = 0; - ibv_device ** devs = ibv_get_device_list(&num_devs); - if (!devs || num_devs == 0) return false; - - ibv_context * ibctx = nullptr; - const char * matched_dev = nullptr; - int gid_idx = gid_env ? atoi(gid_env) : -1; - int gid_version = IBV_GID_TYPE_IB; // 0 = unknown/IB - - for (int d = 0; d < num_devs; d++) { - const char * dn = ibv_get_device_name(devs[d]); - if (dev_env && strcmp(dev_env, dn) != 0) continue; - - ibv_context * ctx = ibv_open_device(devs[d]); - if (!ctx) continue; - - ibv_port_attr pa; - if (ibv_query_port(ctx, ib_port, &pa) != 0) { ibv_close_device(ctx); continue; } - - int found_gid = gid_idx; - int found_version = IBV_GID_TYPE_IB; - if (found_gid < 0) { - // Find a GID on this port whose bytes equal the local TCP address - // (IPv4 or IPv6). Prefer RoCE v2 (UDP/IP, L3-routable) over v1 - // (raw Ethernet, same-L2 only) so silent hangs on L3-routed paths - // are avoided. ibv_query_gid_ex returns gid+type in one call. - int v2_idx = -1; - int v1_idx = -1; - for (int i = 0; i < pa.gid_tbl_len; i++) { - ibv_gid_entry entry = {}; - if (ibv_query_gid_ex(ctx, ib_port, i, &entry, 0) != 0) continue; - if (memcmp(entry.gid.raw, target_gid->data(), RDMA_GID_SIZE) != 0) continue; - if (entry.gid_type == IBV_GID_TYPE_ROCE_V2 && v2_idx < 0) { - v2_idx = i; - } else if (entry.gid_type == IBV_GID_TYPE_ROCE_V1 && v1_idx < 0) { - v1_idx = i; - } - } - if (v2_idx >= 0) { - found_gid = v2_idx; - found_version = IBV_GID_TYPE_ROCE_V2; - } else if (v1_idx >= 0) { - found_gid = v1_idx; - found_version = IBV_GID_TYPE_ROCE_V1; - } - } else { - // Explicit GID index from GGML_RDMA_GID — fetch its type for logging. - ibv_gid_entry entry = {}; - if (ibv_query_gid_ex(ctx, ib_port, found_gid, &entry, 0) == 0) { - found_version = entry.gid_type; - } - } - if (found_gid >= 0) { - ibctx = ctx; - gid_idx = found_gid; - gid_version = found_version; - matched_dev = dn; - rdma_local.path_mtu = pa.active_mtu; - break; - } - ibv_close_device(ctx); - } - ibv_free_device_list(devs); - if (!ibctx) return false; - - rdma_local.ib_port = ib_port; - rdma_local.gid_idx = gid_idx; - - rdma = std::make_unique(); - rdma->ctx = ibctx; - - rdma->pd = ibv_alloc_pd(ibctx); - if (!rdma->pd) return false; - - rdma->scq = ibv_create_cq(ibctx, 16, nullptr, nullptr, 0); - rdma->rcq = ibv_create_cq(ibctx, RDMA_RX_DEPTH + 4, nullptr, nullptr, 0); - if (!rdma->scq || !rdma->rcq) return false; - - ibv_qp_init_attr qia = {}; - qia.send_cq = rdma->scq; - qia.recv_cq = rdma->rcq; - qia.qp_type = IBV_QPT_RC; - qia.cap.max_send_wr = 4; - qia.cap.max_recv_wr = RDMA_RX_DEPTH + 4; - qia.cap.max_send_sge = 1; - qia.cap.max_recv_sge = 1; - qia.cap.max_inline_data = 256; - - rdma->qp = ibv_create_qp(rdma->pd, &qia); - if (!rdma->qp) return false; - rdma->max_inline = qia.cap.max_inline_data; - - rdma->tx_buf = aligned_alloc(4096, RDMA_CHUNK); - rdma->rx_buf = aligned_alloc(4096, static_cast(RDMA_RX_DEPTH) * RDMA_CHUNK); - if (!rdma->tx_buf || !rdma->rx_buf) return false; - - rdma->tx_mr = ibv_reg_mr(rdma->pd, rdma->tx_buf, RDMA_CHUNK, IBV_ACCESS_LOCAL_WRITE); - rdma->rx_mr = ibv_reg_mr(rdma->pd, rdma->rx_buf, static_cast(RDMA_RX_DEPTH) * RDMA_CHUNK, - IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); - if (!rdma->tx_mr || !rdma->rx_mr) return false; - - ibv_gid local_gid; - if (ibv_query_gid(ibctx, ib_port, gid_idx, &local_gid) != 0) return false; - - rdma_local.qpn = rdma->qp->qp_num; - rdma_local.psn = rdma->qp->qp_num & 0xffffff; - memcpy(&rdma_local.gid, &local_gid, RDMA_GID_SIZE); - - const char * ver_str = ""; - if (gid_version == IBV_GID_TYPE_ROCE_V2) { - ver_str = " RoCEv2"; - } else if (gid_version == IBV_GID_TYPE_ROCE_V1) { - ver_str = " RoCEv1"; - } - GGML_LOG_INFO("RDMA probed: dev=%s gid=%d%s qpn=%u inline=%u\n", - matched_dev, gid_idx, ver_str, rdma_local.qpn, rdma->max_inline); - return true; -} - -// Phase 2: Given remote QPN/PSN/GID, transition QP: RESET->INIT->pre-post->RTR->RTS. -// On success, the connection is live and ready for rdma_send/rdma_recv. -bool socket_t::impl::rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid) { - // RESET -> INIT - { - struct ibv_qp_attr a = {}; - a.qp_state = IBV_QPS_INIT; - a.port_num = rdma_local.ib_port; - a.pkey_index = 0; - a.qp_access_flags = IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_LOCAL_WRITE; - if (ibv_modify_qp(rdma->qp, &a, - IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS) != 0) { - return false; - } - } - - for (int i = 0; i < RDMA_RX_DEPTH; i++) { - if (!rdma->post_rx(i)) return false; - } - - // INIT -> RTR - { - struct ibv_qp_attr a = {}; - a.qp_state = IBV_QPS_RTR; - a.path_mtu = rdma_local.path_mtu; - a.dest_qp_num = remote_qpn; - a.rq_psn = remote_psn; - a.max_dest_rd_atomic = 1; - a.min_rnr_timer = 1; - a.ah_attr.is_global = 1; - memcpy(&a.ah_attr.grh.dgid, remote_gid, RDMA_GID_SIZE); - a.ah_attr.grh.hop_limit = 1; - a.ah_attr.grh.sgid_index = rdma_local.gid_idx; - a.ah_attr.dlid = 0; - a.ah_attr.port_num = rdma_local.ib_port; - if (ibv_modify_qp(rdma->qp, &a, - IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN | - IBV_QP_RQ_PSN | IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER) != 0) { - return false; - } - } - - // RTR -> RTS - { - struct ibv_qp_attr a = {}; - a.qp_state = IBV_QPS_RTS; - a.timeout = 14; - a.retry_cnt = 7; - a.rnr_retry = 7; - a.sq_psn = rdma_local.psn; - a.max_rd_atomic = 1; - if (ibv_modify_qp(rdma->qp, &a, - IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT | IBV_QP_RNR_RETRY | - IBV_QP_SQ_PSN | IBV_QP_MAX_QP_RD_ATOMIC) != 0) { - return false; - } - } - - GGML_LOG_INFO("RDMA activated: qpn=%u->%u mtu=%d rx_depth=%d\n", - rdma_local.qpn, remote_qpn, 128 << rdma_local.path_mtu, RDMA_RX_DEPTH); - return true; -} - -bool socket_t::impl::rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc) { - for (uint64_t s = 0; ; s++) { - int n = ibv_poll_cq(cq, 1, wc); - if (n > 0) { - if (wc->status != IBV_WC_SUCCESS) { - GGML_LOG_ERROR("RDMA CQ wc error: status=%d (%s) vendor_err=0x%x\n", - wc->status, ibv_wc_status_str(wc->status), wc->vendor_err); - } - return wc->status == IBV_WC_SUCCESS; - } - if (n < 0) return false; - if ((s & 0xFFFFF) == 0 && s > 0) { - if (tcp_peer_closed()) { - return false; - } - } - } -} - -bool socket_t::impl::rdma_send(const void * data, size_t size) { - rdma_conn * c = rdma.get(); - const uint8_t * src = (const uint8_t *)data; - size_t rem = size; - while (rem > 0) { - size_t chunk = std::min(rem, RDMA_CHUNK); - - struct ibv_sge sge = {}; - struct ibv_send_wr wr = {}, * bad = nullptr; - wr.opcode = IBV_WR_SEND; - wr.sg_list = &sge; - wr.num_sge = 1; - - if (chunk <= c->max_inline) { - sge.addr = (uintptr_t)src; - sge.length = chunk; - wr.send_flags = IBV_SEND_SIGNALED | IBV_SEND_INLINE; - } else { - memcpy(c->tx_buf, src, chunk); - sge.addr = (uintptr_t)c->tx_buf; - sge.length = chunk; - sge.lkey = c->tx_mr->lkey; - wr.send_flags = IBV_SEND_SIGNALED; - } - - if (ibv_post_send(c->qp, &wr, &bad) != 0) return false; - struct ibv_wc wc; - if (!rdma_poll(c->scq, &wc)) return false; - - src += chunk; - rem -= chunk; - } - return true; -} - -bool socket_t::impl::rdma_recv(void * data, size_t size) { - rdma_conn * c = rdma.get(); - uint8_t * dst = (uint8_t *)data; - size_t rem = size; - while (rem > 0) { - struct ibv_wc wc; - if (!rdma_poll(c->rcq, &wc)) return false; - - int slot = (int)wc.wr_id; - size_t got = wc.byte_len; - memcpy(dst, c->rx_slot(slot), got); - - if (!c->post_rx(slot)) return false; - - dst += got; - rem -= got; - } - return true; -} - -#endif // GGML_RPC_RDMA - -bool socket_t::impl::send_data(const void * data, size_t size) { -#ifdef GGML_RPC_RDMA - if (use_rdma) { - return rdma_send(data, size); - } -#endif - size_t bytes_sent = 0; - while (bytes_sent < size) { - size_t size_to_send = std::min(size - bytes_sent, MAX_CHUNK_SIZE); - ssize_t n = send(fd, (const char *)data + bytes_sent, size_to_send, 0); - if (n < 0) { - GGML_LOG_ERROR("send failed (bytes_sent=%zu, size_to_send=%zu)\n", - bytes_sent, size_to_send); - return false; - } - bytes_sent += (size_t)n; - } - return true; -} - -bool socket_t::impl::recv_data(void * data, size_t size) { -#ifdef GGML_RPC_RDMA - if (use_rdma) { - return rdma_recv(data, size); - } -#endif - size_t bytes_recv = 0; - while (bytes_recv < size) { - size_t size_to_recv = std::min(size - bytes_recv, MAX_CHUNK_SIZE); - ssize_t n = recv(fd, (char *)data + bytes_recv, size_to_recv, 0); - if (n < 0) { - GGML_LOG_ERROR("recv failed (bytes_recv=%zu, size_to_recv=%zu)\n", - bytes_recv, size_to_recv); - return false; - } - if (n == 0) { - LOG_DBG("recv returned 0 (peer closed?)\n"); - return false; - } - bytes_recv += (size_t)n; - } - return true; -} - -void socket_t::impl::get_caps(uint8_t * local_caps) { - memset(local_caps, 0, RPC_CONN_CAPS_SIZE); -#ifdef GGML_RPC_RDMA - rdma_local = {}; - if (rdma_probe()) { - rdma_caps rc = {}; - rc.qpn = rdma_local.qpn; - rc.psn = rdma_local.psn; - memcpy(rc.gid, rdma_local.gid, RDMA_GID_SIZE); - memcpy(local_caps, &rc, sizeof(rc)); - } else { - rdma.reset(); - } -#endif // GGML_RPC_RDMA -} - -void socket_t::impl::update_caps(const uint8_t * remote_caps) { -#ifdef GGML_RPC_RDMA - if (!rdma) { - return; - } - rdma_caps rc = {}; - memcpy(&rc, remote_caps, sizeof(rc)); - if (rc.qpn == 0) { - rdma.reset(); - return; - } - if (rdma_activate(rc.qpn, rc.psn, rc.gid)) { - use_rdma = true; - } else { - GGML_LOG_ERROR("RDMA activate failed, staying on TCP\n"); - rdma.reset(); - } -#else - (void)remote_caps; -#endif // GGML_RPC_RDMA -} - - -///////////////////////////////////////////////////////////////////////////// - -socket_t::socket_t(std::unique_ptr p) : pimpl(std::move(p)) {} - -socket_t::~socket_t() = default; - -bool socket_t::send_data(const void * data, size_t size) { - return pimpl->send_data(data, size); -} - -bool socket_t::recv_data(void * data, size_t size) { - return pimpl->recv_data(data, size); -} - -void socket_t::get_caps(uint8_t * local_caps) { - return pimpl->get_caps(local_caps); -} - -void socket_t::update_caps(const uint8_t * remote_caps) { - return pimpl->update_caps(remote_caps); -} - -static bool is_valid_fd(sockfd_t sockfd) { -#ifdef _WIN32 - return sockfd != INVALID_SOCKET; -#else - return sockfd >= 0; -#endif -} - -static bool set_no_delay(sockfd_t sockfd) { - int flag = 1; - // set TCP_NODELAY to disable Nagle's algorithm - int ret = setsockopt(sockfd, IPPROTO_TCP, TCP_NODELAY, (char *)&flag, sizeof(int)); - return ret == 0; -} - -static bool set_reuse_addr(sockfd_t sockfd) { - int flag = 1; - int ret = setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, (char *)&flag, sizeof(int)); - return ret == 0; -} - -socket_ptr socket_t::accept() { - auto client_socket_fd = ::accept(pimpl->fd, NULL, NULL); - if (!is_valid_fd(client_socket_fd)) { - return nullptr; - } - if (!set_no_delay(client_socket_fd)) { - GGML_LOG_ERROR("Failed to set TCP_NODELAY\n"); - return nullptr; - } - return socket_ptr(new socket_t(std::make_unique(client_socket_fd))); -} - -socket_ptr socket_t::create_server(const char * host, int port) { - auto sockfd = socket(AF_INET, SOCK_STREAM, 0); - if (!is_valid_fd(sockfd)) { - return nullptr; - } - if (!set_reuse_addr(sockfd)) { - GGML_LOG_ERROR("Failed to set SO_REUSEADDR\n"); - return nullptr; - } - if (inet_addr(host) == INADDR_NONE) { - GGML_LOG_ERROR("Invalid host address: %s\n", host); - return nullptr; - } - struct sockaddr_in serv_addr; - serv_addr.sin_family = AF_INET; - serv_addr.sin_addr.s_addr = inet_addr(host); - serv_addr.sin_port = htons(port); - - if (bind(sockfd, (struct sockaddr *) &serv_addr, sizeof(serv_addr)) < 0) { - return nullptr; - } - if (listen(sockfd, 1) < 0) { - return nullptr; - } - return socket_ptr(new socket_t(std::make_unique(sockfd))); -} - -socket_ptr socket_t::connect(const char * host, int port) { - auto sockfd = socket(AF_INET, SOCK_STREAM, 0); - if (!is_valid_fd(sockfd)) { - return nullptr; - } - if (!set_no_delay(sockfd)) { - GGML_LOG_ERROR("Failed to set TCP_NODELAY\n"); - return nullptr; - } - struct sockaddr_in addr; - addr.sin_family = AF_INET; - addr.sin_port = htons(port); - struct hostent * server = gethostbyname(host); - if (server == NULL) { - GGML_LOG_ERROR("Cannot resolve host '%s'\n", host); - return nullptr; - } - memcpy(&addr.sin_addr.s_addr, server->h_addr, server->h_length); - if (::connect(sockfd, (struct sockaddr *)&addr, sizeof(addr)) < 0) { - return nullptr; - } - return socket_ptr(new socket_t(std::make_unique(sockfd))); -} - -#ifdef _WIN32 -static std::mutex g_rpc_transport_mu; -static bool g_rpc_transport_wsa_started = false; -#endif - -bool rpc_transport_init() { -#ifdef _WIN32 - std::lock_guard lock(g_rpc_transport_mu); - if (g_rpc_transport_wsa_started) { - return true; - } - WSADATA wsaData; - int res = WSAStartup(MAKEWORD(2, 2), &wsaData); - if (res != 0) { - return false; - } - g_rpc_transport_wsa_started = true; - return true; -#else - return true; -#endif -} - -void rpc_transport_shutdown() { -#ifdef _WIN32 - std::lock_guard lock(g_rpc_transport_mu); - if (!g_rpc_transport_wsa_started) { - return; - } - WSACleanup(); - g_rpc_transport_wsa_started = false; -#endif -} diff --git a/ggml/src/ggml-rpc/transport.h b/ggml/src/ggml-rpc/transport.h deleted file mode 100644 index 73b85cc530a0..000000000000 --- a/ggml/src/ggml-rpc/transport.h +++ /dev/null @@ -1,34 +0,0 @@ -#pragma once - -#include -#include -#include - -struct socket_t; -typedef std::shared_ptr socket_ptr; - -static constexpr size_t MAX_CHUNK_SIZE = 1024ull * 1024ull * 1024ull; // 1 GiB -static constexpr size_t RPC_CONN_CAPS_SIZE = 24; - -struct socket_t { - ~socket_t(); - - bool send_data(const void * data, size_t size); - bool recv_data(void * data, size_t size); - - socket_ptr accept(); - - void get_caps(uint8_t * local_caps); - void update_caps(const uint8_t * remote_caps); - - static socket_ptr create_server(const char * host, int port); - static socket_ptr connect(const char * host, int port); - -private: - struct impl; - explicit socket_t(std::unique_ptr p); - std::unique_ptr pimpl; -}; - -bool rpc_transport_init(); -void rpc_transport_shutdown(); diff --git a/ggml/src/ggml-sycl/CMakeLists.txt b/ggml/src/ggml-sycl/CMakeLists.txt index 8e589fa238dc..fe7ac0ca56f2 100644 --- a/ggml/src/ggml-sycl/CMakeLists.txt +++ b/ggml/src/ggml-sycl/CMakeLists.txt @@ -18,20 +18,52 @@ endif() message(STATUS "SYCL found") #todo: AOT +set(GGML_SOURCES_SYCL + ggml-sycl.inc + common.inc + add-id.inc + binbcast.inc + concat.inc + conv.inc + convert.inc + count-equal.inc + cpy.inc + dmmv.inc + element_wise.inc + fattn.inc + fattn-tile.inc + gated_delta_net.inc + getrows.inc + gla.inc + im2col.inc + mmq.inc + mmvq.inc + norm.inc + outprod.inc + pad.inc + pad_reflect_1d.inc + repeat_back.inc + roll.inc + rope.inc + set.inc + set_rows.inc + softmax.inc + ssm_conv.inc + sycl_hw.inc + tsembd.inc + upscale.inc + wkv.inc + template-instances/fattn-tile-instances.inc + template-instances/fattn-vec-instances.inc +) + +set_source_files_properties(${GGML_SOURCES_SYCL} PROPERTIES LANGUAGE CXX) + ggml_add_backend_library(ggml-sycl - ggml-sycl.cpp - ../../include/ggml-sycl.h + ${GGML_SOURCES_SYCL} + ../../include/ggml-sycl.h.inc ) -file(GLOB GGML_HEADERS_SYCL "*.hpp") -file(GLOB GGML_SOURCES_SYCL "*.cpp") -file(GLOB SRCS "template-instances/fattn-tile*.cpp") -list(APPEND GGML_SOURCES_SYCL ${SRCS}) -file(GLOB SRCS "template-instances/fattn-vec*.cpp") -list(APPEND GGML_SOURCES_SYCL ${SRCS}) - -target_sources(ggml-sycl PRIVATE ${GGML_HEADERS_SYCL} ${GGML_SOURCES_SYCL}) - if (WIN32) # To generate a Visual Studio solution, using Intel C++ Compiler for ggml-sycl is mandatory if( ${CMAKE_GENERATOR} MATCHES "Visual Studio" AND NOT (${CMAKE_GENERATOR_TOOLSET} MATCHES "Intel C")) @@ -163,4 +195,3 @@ if (GGML_SYCL_DEVICE_ARCH) target_compile_options(ggml-sycl PRIVATE -Xsycl-target-backend --offload-arch=${GGML_SYCL_DEVICE_ARCH}) target_link_options(ggml-sycl PRIVATE -Xsycl-target-backend --offload-arch=${GGML_SYCL_DEVICE_ARCH}) endif() - diff --git a/ggml/src/ggml-sycl/add-id.hpp b/ggml/src/ggml-sycl/add-id.hpp deleted file mode 100644 index e1b09ee8c7c8..000000000000 --- a/ggml/src/ggml-sycl/add-id.hpp +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef GGML_SYCL_ADD_ID_HPP -#define GGML_SYCL_ADD_ID_HPP - -#include "common.hpp" - -void ggml_sycl_add_id(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -#endif // GGML_SYCL_ADD_ID_HPP diff --git a/ggml/src/ggml-sycl/add-id.cpp b/ggml/src/ggml-sycl/add-id.inc similarity index 98% rename from ggml/src/ggml-sycl/add-id.cpp rename to ggml/src/ggml-sycl/add-id.inc index e0adc4fe4230..c6d43f48787a 100644 --- a/ggml/src/ggml-sycl/add-id.cpp +++ b/ggml/src/ggml-sycl/add-id.inc @@ -1,6 +1,5 @@ #include -#include "common.hpp" -#include "add-id.hpp" +#include "common-defs.inc" static void add_id_kernel( const float* src0, diff --git a/ggml/src/ggml-sycl/backend-defs.inc b/ggml/src/ggml-sycl/backend-defs.inc new file mode 100644 index 000000000000..4d3aabcd274e --- /dev/null +++ b/ggml/src/ggml-sycl/backend-defs.inc @@ -0,0 +1,48 @@ +// +// MIT license +// Copyright (C) 2024 Intel Corporation +// SPDX-License-Identifier: MIT +// + +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// + +#ifndef GGML_SYCL_BACKEND_HPP +#define GGML_SYCL_BACKEND_HPP + +#include "binbcast-defs.inc" +#include "common-defs.inc" +#include "concat-defs.inc" +#include "conv-defs.inc" +#include "convert-defs.inc" +#include "count-equal-defs.inc" +#include "cpy-defs.inc" +#include "dequantize-defs.inc" +#include "dmmv-defs.inc" +#include "element_wise-defs.inc" +#include "fattn-defs.inc" +#include "gated_delta_net-defs.inc" +#include "gla-defs.inc" +#include "im2col-defs.inc" +#include "mmq-defs.inc" +#include "mmvq-defs.inc" +#include "norm-defs.inc" +#include "outprod-defs.inc" +#include "pad-defs.inc" +#include "pad_reflect_1d-defs.inc" +#include "quantize-defs.inc" +#include "quants-defs.inc" +#include "roll-defs.inc" +#include "rope-defs.inc" +#include "set_rows-defs.inc" +#include "ssm_conv-defs.inc" +#include "softmax-defs.inc" +#include "tsembd-defs.inc" +#include "upscale-defs.inc" +#include "wkv-defs.inc" + + +#endif // GGML_SYCL_BACKEND_HPP diff --git a/ggml/src/ggml-sycl/backend.hpp b/ggml/src/ggml-sycl/backend.hpp deleted file mode 100644 index a526d8e58bc9..000000000000 --- a/ggml/src/ggml-sycl/backend.hpp +++ /dev/null @@ -1,48 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_BACKEND_HPP -#define GGML_SYCL_BACKEND_HPP - -#include "binbcast.hpp" -#include "common.hpp" -#include "concat.hpp" -#include "conv.hpp" -#include "convert.hpp" -#include "count-equal.hpp" -#include "cpy.hpp" -#include "dequantize.hpp" -#include "dmmv.hpp" -#include "element_wise.hpp" -#include "fattn.hpp" -#include "gated_delta_net.hpp" -#include "gla.hpp" -#include "im2col.hpp" -#include "mmq.hpp" -#include "mmvq.hpp" -#include "norm.hpp" -#include "outprod.hpp" -#include "pad.hpp" -#include "pad_reflect_1d.hpp" -#include "quantize.hpp" -#include "quants.hpp" -#include "roll.hpp" -#include "rope.hpp" -#include "set_rows.hpp" -#include "ssm_conv.hpp" -#include "softmax.hpp" -#include "tsembd.hpp" -#include "upscale.hpp" -#include "wkv.hpp" - - -#endif // GGML_SYCL_BACKEND_HPP diff --git a/ggml/src/ggml-sycl/binbcast.hpp b/ggml/src/ggml-sycl/binbcast.hpp deleted file mode 100644 index 9cce0f053a58..000000000000 --- a/ggml/src/ggml-sycl/binbcast.hpp +++ /dev/null @@ -1,39 +0,0 @@ -#ifndef GGML_SYCL_BINBCAST_HPP -#define GGML_SYCL_BINBCAST_HPP -#include "common.hpp" - - -static __dpct_inline__ float op_repeat(const float a, const float b) { - return b; - GGML_UNUSED(a); -} - -static __dpct_inline__ float op_add(const float a, const float b) { - return a + b; -} - -static __dpct_inline__ float op_sub(const float a, const float b) { - return a - b; -} - -static __dpct_inline__ float op_mul(const float a, const float b) { - return a * b; -} - -static __dpct_inline__ float op_div(const float a, const float b) { - return a / b; -} - -void ggml_sycl_add(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -void ggml_sycl_sub(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -void ggml_sycl_mul(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -void ggml_sycl_div(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -void ggml_sycl_repeat(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - - -#endif //GGML_SYCL_BINBCAST_HPP - diff --git a/ggml/src/ggml-sycl/binbcast.cpp b/ggml/src/ggml-sycl/binbcast.inc similarity index 96% rename from ggml/src/ggml-sycl/binbcast.cpp rename to ggml/src/ggml-sycl/binbcast.inc index 92dd18889f4f..e1892560984d 100644 --- a/ggml/src/ggml-sycl/binbcast.cpp +++ b/ggml/src/ggml-sycl/binbcast.inc @@ -1,10 +1,31 @@ -#include "binbcast.hpp" +#include "common-defs.inc" #include #include #include -#include "ggml.h" +#include "ggml.h.inc" + +static __dpct_inline__ float op_repeat(const float a, const float b) { + return b; + GGML_UNUSED(a); +} + +static __dpct_inline__ float op_add(const float a, const float b) { + return a + b; +} + +static __dpct_inline__ float op_sub(const float a, const float b) { + return a - b; +} + +static __dpct_inline__ float op_mul(const float a, const float b) { + return a * b; +} + +static __dpct_inline__ float op_div(const float a, const float b) { + return a / b; +} template static void k_bin_bcast(const src0_t * src0, const src1_t * src1, dst_t * dst, @@ -343,4 +364,3 @@ void ggml_sycl_repeat(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1); ggml_sycl_op_repeat(ctx, dst); } - diff --git a/ggml/src/ggml-sycl/common.hpp b/ggml/src/ggml-sycl/common-defs.inc similarity index 97% rename from ggml/src/ggml-sycl/common.hpp rename to ggml/src/ggml-sycl/common-defs.inc index 5abf22906518..9d02419b0484 100644 --- a/ggml/src/ggml-sycl/common.hpp +++ b/ggml/src/ggml-sycl/common-defs.inc @@ -18,15 +18,34 @@ #include #include -#include "dpct/helper.hpp" -#include "ggml.h" -#include "ggml-impl.h" -#include "ggml-sycl.h" -#include "presets.hpp" -#include "type.hpp" -#include "sycl_hw.hpp" +#include "dpct/helper-defs.inc" +#include "ggml.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-sycl.h.inc" +#include "presets-defs.inc" +#include "type-defs.inc" namespace syclexp = sycl::ext::oneapi::experimental; +namespace syclex = sycl::ext::oneapi::experimental; +using gpu_arch = sycl::ext::oneapi::experimental::architecture; + +enum sycl_intel_gpu_family { + GPU_FAMILY_UKNOWN = -1, + GPU_FAMILY_IGPU_NON_XE = 0, + GPU_FAMILY_IGPU_XE = 1, + GPU_FAMILY_DGPU_CLIENT_GAME = 2, + GPU_FAMILY_DGPU_CLOUD = 3 +}; + +struct sycl_hw_info { + syclex::architecture arch; + const char* arch_name; + int32_t device_id; + std::string name; + sycl_intel_gpu_family gpu_family; +}; + +sycl_hw_info get_device_hw_info(sycl::device * device_ptr); #if defined(__INTEL_LLVM_COMPILER) && __has_include() #include @@ -48,9 +67,9 @@ namespace syclexp = sycl::ext::oneapi::experimental; /* suppress warning spam */ #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wnested-anon-types" -#include "ggml-common.h" +#include "ggml-common-defs.inc" #pragma clang diagnostic pop -#include "ggml-impl.h" +#include "ggml-impl-defs.inc" void* ggml_sycl_host_malloc(size_t size); void ggml_sycl_host_free(void* ptr); diff --git a/ggml/src/ggml-sycl/common.cpp b/ggml/src/ggml-sycl/common.inc similarity index 96% rename from ggml/src/ggml-sycl/common.cpp rename to ggml/src/ggml-sycl/common.inc index 05fd5ef46c76..97e52642ff78 100644 --- a/ggml/src/ggml-sycl/common.cpp +++ b/ggml/src/ggml-sycl/common.inc @@ -10,10 +10,10 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // -#include "common.hpp" +#include "common-defs.inc" -#include "ggml-backend-impl.h" -#include "ggml-impl.h" +#include "ggml-backend-impl-defs.inc" +#include "ggml-impl-defs.inc" int get_current_device_id() { return dpct::dev_mgr::instance().current_device_id(); diff --git a/ggml/src/ggml-sycl/concat.hpp b/ggml/src/ggml-sycl/concat.hpp deleted file mode 100644 index e5cb7314c9f3..000000000000 --- a/ggml/src/ggml-sycl/concat.hpp +++ /dev/null @@ -1,20 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_CONCAT_HPP -#define GGML_SYCL_CONCAT_HPP - -#include "common.hpp" - -void ggml_sycl_op_concat(ggml_backend_sycl_context & ctx, ggml_tensor *dst); - -#endif // GGML_SYCL_CONCAT_HPP diff --git a/ggml/src/ggml-sycl/concat.cpp b/ggml/src/ggml-sycl/concat.inc similarity index 99% rename from ggml/src/ggml-sycl/concat.cpp rename to ggml/src/ggml-sycl/concat.inc index d16215bc91cc..9c813e842d20 100644 --- a/ggml/src/ggml-sycl/concat.cpp +++ b/ggml/src/ggml-sycl/concat.inc @@ -10,7 +10,7 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // -#include "concat.hpp" +#include "common-defs.inc" static inline size_t elem_size(ggml_type t) { return ggml_type_size(t) / ggml_blck_size(t); diff --git a/ggml/src/ggml-sycl/conv.hpp b/ggml/src/ggml-sycl/conv.hpp deleted file mode 100644 index f9e60dc75802..000000000000 --- a/ggml/src/ggml-sycl/conv.hpp +++ /dev/null @@ -1,20 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_CONV_HPP -#define GGML_SYCL_CONV_HPP - -#include "common.hpp" - -void ggml_sycl_op_conv_transpose_1d(ggml_backend_sycl_context & ctx, ggml_tensor *dst); - -#endif // GGML_SYCL_CONV_HPP diff --git a/ggml/src/ggml-sycl/conv.cpp b/ggml/src/ggml-sycl/conv.inc similarity index 99% rename from ggml/src/ggml-sycl/conv.cpp rename to ggml/src/ggml-sycl/conv.inc index 475bd34a25d5..89f25b02cfbe 100644 --- a/ggml/src/ggml-sycl/conv.cpp +++ b/ggml/src/ggml-sycl/conv.inc @@ -10,7 +10,7 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // -#include "conv.hpp" +#include "common-defs.inc" static void conv_transpose_1d_kernel( const int s0, const int output_size, @@ -98,4 +98,3 @@ void ggml_sycl_op_conv_transpose_1d(ggml_backend_sycl_context & ctx, ggml_tensor src1->ne[0], dst->ne[0], src0_d, src1_d, dst_d, stream); } - diff --git a/ggml/src/ggml-sycl/convert.hpp b/ggml/src/ggml-sycl/convert-defs.inc similarity index 98% rename from ggml/src/ggml-sycl/convert.hpp rename to ggml/src/ggml-sycl/convert-defs.inc index 8de79d10ff69..0246b5d72961 100644 --- a/ggml/src/ggml-sycl/convert.hpp +++ b/ggml/src/ggml-sycl/convert-defs.inc @@ -13,7 +13,7 @@ #ifndef GGML_SYCL_CONVERT_HPP #define GGML_SYCL_CONVERT_HPP -#include "common.hpp" +#include "common-defs.inc" template using to_t_sycl_t = void (*)(const void * __restrict__ x, T * __restrict__ y, int64_t k, dpct::queue_ptr stream); diff --git a/ggml/src/ggml-sycl/convert.cpp b/ggml/src/ggml-sycl/convert.inc similarity index 99% rename from ggml/src/ggml-sycl/convert.cpp rename to ggml/src/ggml-sycl/convert.inc index 67b9c06f3e44..7f255fb95efa 100644 --- a/ggml/src/ggml-sycl/convert.cpp +++ b/ggml/src/ggml-sycl/convert.inc @@ -1,6 +1,6 @@ -#include "convert.hpp" -#include "dequantize.hpp" -#include "presets.hpp" +#include "convert-defs.inc" +#include "dequantize-defs.inc" +#include "presets-defs.inc" template static void dequantize_block(const void * __restrict__ vx, dst_t * __restrict__ y, const int64_t k, diff --git a/ggml/src/ggml-sycl/count-equal.hpp b/ggml/src/ggml-sycl/count-equal.hpp deleted file mode 100644 index f7f4fcbd0bad..000000000000 --- a/ggml/src/ggml-sycl/count-equal.hpp +++ /dev/null @@ -1,9 +0,0 @@ -#ifndef GGML_SYCL_COUNT_EQUAL_HPP -#define GGML_SYCL_COUNT_EQUAL_HPP -#include "common.hpp" - -#define SYCL_COUNT_EQUAL_CHUNK_SIZE 128 - -void ggml_sycl_count_equal(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -#endif //GGML_SYCL_COUNT_EQUAL_HPP diff --git a/ggml/src/ggml-sycl/count-equal.cpp b/ggml/src/ggml-sycl/count-equal.inc similarity index 97% rename from ggml/src/ggml-sycl/count-equal.cpp rename to ggml/src/ggml-sycl/count-equal.inc index 4580354cd9d0..c477f3bbb4a5 100644 --- a/ggml/src/ggml-sycl/count-equal.cpp +++ b/ggml/src/ggml-sycl/count-equal.inc @@ -1,7 +1,9 @@ -#include "count-equal.hpp" +#include "common-defs.inc" #include +#define SYCL_COUNT_EQUAL_CHUNK_SIZE 128 + template static void count_equal(const T *__restrict__ x, const T *__restrict__ y, int64_t *__restrict__ dst, const int64_t dk, diff --git a/ggml/src/ggml-sycl/cpy.hpp b/ggml/src/ggml-sycl/cpy-defs.inc similarity index 99% rename from ggml/src/ggml-sycl/cpy.hpp rename to ggml/src/ggml-sycl/cpy-defs.inc index 3c331f1ef27b..745f63b39246 100644 --- a/ggml/src/ggml-sycl/cpy.hpp +++ b/ggml/src/ggml-sycl/cpy-defs.inc @@ -1,7 +1,7 @@ #ifndef GGML_SYCL_CPY_HPP #define GGML_SYCL_CPY_HPP -#include "common.hpp" +#include "common-defs.inc" #include typedef void (*cpy_kernel_t)(const char * cx, char * cdst); diff --git a/ggml/src/ggml-sycl/cpy.cpp b/ggml/src/ggml-sycl/cpy.inc similarity index 99% rename from ggml/src/ggml-sycl/cpy.cpp rename to ggml/src/ggml-sycl/cpy.inc index 96709554cf69..57344a2d1946 100644 --- a/ggml/src/ggml-sycl/cpy.cpp +++ b/ggml/src/ggml-sycl/cpy.inc @@ -1,11 +1,11 @@ -#include "cpy.hpp" +#include "cpy-defs.inc" #include -#include "dequantize.hpp" -#include "ggml-sycl/common.hpp" -#include "ggml-sycl/presets.hpp" -#include "ggml.h" +#include "dequantize-defs.inc" +#include "ggml-sycl/common-defs.inc" +#include "ggml-sycl/presets-defs.inc" +#include "ggml.h.inc" static void cpy_1_f32_f32(const char * cxi, char * cdsti) { diff --git a/ggml/src/ggml-sycl/dequantize.hpp b/ggml/src/ggml-sycl/dequantize-defs.inc similarity index 99% rename from ggml/src/ggml-sycl/dequantize.hpp rename to ggml/src/ggml-sycl/dequantize-defs.inc index 19fa88680d69..a9556460cf19 100644 --- a/ggml/src/ggml-sycl/dequantize.hpp +++ b/ggml/src/ggml-sycl/dequantize-defs.inc @@ -13,8 +13,8 @@ #ifndef GGML_SYCL_DEQUANTIZE_HPP #define GGML_SYCL_DEQUANTIZE_HPP -#include "common.hpp" -#include "convert.hpp" +#include "common-defs.inc" +#include "convert-defs.inc" typedef void (*dequantize_kernel_t)(const void * vx, const int64_t ib, const int iqs, dfloat2 & v); typedef void (*dequantize_kernel_t_reorder)(const void *d, const int64_t ib, const void *qs, diff --git a/ggml/src/ggml-sycl/dmmv.hpp b/ggml/src/ggml-sycl/dmmv.hpp deleted file mode 100644 index bd8373564153..000000000000 --- a/ggml/src/ggml-sycl/dmmv.hpp +++ /dev/null @@ -1,27 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_DMMV_HPP -#define GGML_SYCL_DMMV_HPP - -#include "common.hpp" - - -void ggml_sycl_op_dequantize_mul_mat_vec( - ggml_backend_sycl_context & ctx, - const ggml_tensor *src0, const ggml_tensor *src1, ggml_tensor *dst, - const char *src0_dd_i, const float *src1_ddf_i, const char *src1_ddq_i, - float *dst_dd_i, const int64_t row_low, const int64_t row_high, - const int64_t src1_ncols, const int64_t src1_padded_row_size, - const dpct::queue_ptr &stream); - -#endif // GGML_SYCL_DMMV_HPP diff --git a/ggml/src/ggml-sycl/dmmv.cpp b/ggml/src/ggml-sycl/dmmv.inc similarity index 99% rename from ggml/src/ggml-sycl/dmmv.cpp rename to ggml/src/ggml-sycl/dmmv.inc index 5577bf73b28a..202c19ade2c2 100644 --- a/ggml/src/ggml-sycl/dmmv.cpp +++ b/ggml/src/ggml-sycl/dmmv.inc @@ -1,7 +1,7 @@ -#include "convert.hpp" -#include "dmmv.hpp" -#include "dequantize.hpp" -#include "presets.hpp" +#include "convert-defs.inc" +#include "common-defs.inc" +#include "dequantize-defs.inc" +#include "presets-defs.inc" static void convert_f16(const void * vx, const int64_t ib, const int iqs, dfloat2 & v){ const sycl::half *x = (const sycl::half *)vx; diff --git a/ggml/src/ggml-sycl/dpct/helper.hpp b/ggml/src/ggml-sycl/dpct/helper-defs.inc similarity index 99% rename from ggml/src/ggml-sycl/dpct/helper.hpp rename to ggml/src/ggml-sycl/dpct/helper-defs.inc index 791d3cac52e1..9b963e7a56b4 100644 --- a/ggml/src/ggml-sycl/dpct/helper.hpp +++ b/ggml/src/ggml-sycl/dpct/helper-defs.inc @@ -19,7 +19,7 @@ #include -#include "ggml.h" +#include "ggml.h.inc" #if defined(__linux__) #include diff --git a/ggml/src/ggml-sycl/element_wise.hpp b/ggml/src/ggml-sycl/element_wise-defs.inc similarity index 98% rename from ggml/src/ggml-sycl/element_wise.hpp rename to ggml/src/ggml-sycl/element_wise-defs.inc index 997132166ab2..99629344bd3f 100644 --- a/ggml/src/ggml-sycl/element_wise.hpp +++ b/ggml/src/ggml-sycl/element_wise-defs.inc @@ -1,8 +1,8 @@ #ifndef GGML_SYCL_ELEMENTWISE_HPP #define GGML_SYCL_ELEMENTWISE_HPP -#include "common.hpp" -#include "ggml.h" +#include "common-defs.inc" +#include "ggml.h.inc" #include // For std::numeric_limits #define SYCL_GLU_BLOCK_SIZE 256 diff --git a/ggml/src/ggml-sycl/element_wise.cpp b/ggml/src/ggml-sycl/element_wise.inc similarity index 99% rename from ggml/src/ggml-sycl/element_wise.cpp rename to ggml/src/ggml-sycl/element_wise.inc index 249e80c826ea..d646f511251f 100644 --- a/ggml/src/ggml-sycl/element_wise.cpp +++ b/ggml/src/ggml-sycl/element_wise.inc @@ -1,7 +1,7 @@ -#include "common.hpp" -#include "ggml-sycl/presets.hpp" -#include "ggml.h" -#include "element_wise.hpp" +#include "common-defs.inc" +#include "ggml-sycl/presets-defs.inc" +#include "ggml.h.inc" +#include "element_wise-defs.inc" #define SYCL_GLOBAL_ID_LOOP(K, ITEM) \ for (auto i = ITEM.get_global_id(0); i < (size_t)K; i += ITEM.get_global_range(0)) diff --git a/ggml/src/ggml-sycl/fattn-common.hpp b/ggml/src/ggml-sycl/fattn-common-defs.inc similarity index 99% rename from ggml/src/ggml-sycl/fattn-common.hpp rename to ggml/src/ggml-sycl/fattn-common-defs.inc index ed00d03c3b67..3640282a9ae6 100644 --- a/ggml/src/ggml-sycl/fattn-common.hpp +++ b/ggml/src/ggml-sycl/fattn-common-defs.inc @@ -1,12 +1,12 @@ #pragma once #include -#include "dpct/helper.hpp" -#include "common.hpp" -#include "convert.hpp" -#include "vecdotq.hpp" +#include "dpct/helper-defs.inc" +#include "common-defs.inc" +#include "convert-defs.inc" +#include "vecdotq-defs.inc" -#include "ggml.h" +#include "ggml.h.inc" #include #include diff --git a/ggml/src/ggml-sycl/fattn-tile.hpp b/ggml/src/ggml-sycl/fattn-tile-defs.inc similarity index 99% rename from ggml/src/ggml-sycl/fattn-tile.hpp rename to ggml/src/ggml-sycl/fattn-tile-defs.inc index 9ba5296968d5..183897510583 100644 --- a/ggml/src/ggml-sycl/fattn-tile.hpp +++ b/ggml/src/ggml-sycl/fattn-tile-defs.inc @@ -1,8 +1,8 @@ #include #include -#include "dpct/helper.hpp" -#include "common.hpp" -#include "fattn-common.hpp" +#include "dpct/helper-defs.inc" +#include "common-defs.inc" +#include "fattn-common-defs.inc" #include #include diff --git a/ggml/src/ggml-sycl/fattn-tile.cpp b/ggml/src/ggml-sycl/fattn-tile.inc similarity index 94% rename from ggml/src/ggml-sycl/fattn-tile.cpp rename to ggml/src/ggml-sycl/fattn-tile.inc index 9449d75784d0..98f9a7f7a307 100644 --- a/ggml/src/ggml-sycl/fattn-tile.cpp +++ b/ggml/src/ggml-sycl/fattn-tile.inc @@ -1,9 +1,9 @@ #include #include -#include "dpct/helper.hpp" -#include "common.hpp" -#include "fattn-common.hpp" -#include "fattn-tile.hpp" +#include "dpct/helper-defs.inc" +#include "common-defs.inc" +#include "fattn-common-defs.inc" +#include "fattn-tile-defs.inc" #include #include namespace syclex = sycl::ext::oneapi::experimental; diff --git a/ggml/src/ggml-sycl/fattn-vec.hpp b/ggml/src/ggml-sycl/fattn-vec-defs.inc similarity index 99% rename from ggml/src/ggml-sycl/fattn-vec.hpp rename to ggml/src/ggml-sycl/fattn-vec-defs.inc index 8031acfdff88..9f7bfbdc730c 100644 --- a/ggml/src/ggml-sycl/fattn-vec.hpp +++ b/ggml/src/ggml-sycl/fattn-vec-defs.inc @@ -6,10 +6,10 @@ #include #include -#include "dpct/helper.hpp" -#include "common.hpp" -#include "ggml.h" -#include "fattn-common.hpp" +#include "dpct/helper-defs.inc" +#include "common-defs.inc" +#include "ggml.h.inc" +#include "fattn-common-defs.inc" #include #include diff --git a/ggml/src/ggml-sycl/fattn.hpp b/ggml/src/ggml-sycl/fattn.hpp deleted file mode 100644 index f2a8ffc97dee..000000000000 --- a/ggml/src/ggml-sycl/fattn.hpp +++ /dev/null @@ -1,22 +0,0 @@ -// -// MIT license -// Copyright (C) 2025 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_FATTN_HPP -#define GGML_SYCL_FATTN_HPP - -#include "common.hpp" - -void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst); - -#endif // GGML_SYCL_FATTN_HPP diff --git a/ggml/src/ggml-sycl/fattn.cpp b/ggml/src/ggml-sycl/fattn.inc similarity index 97% rename from ggml/src/ggml-sycl/fattn.cpp rename to ggml/src/ggml-sycl/fattn.inc index 7c6e6112fdcd..aedf63016201 100644 --- a/ggml/src/ggml-sycl/fattn.cpp +++ b/ggml/src/ggml-sycl/fattn.inc @@ -12,12 +12,12 @@ #include -#include "dpct/helper.hpp" -#include "common.hpp" -#include "fattn-common.hpp" -#include "fattn-tile.hpp" -#include "fattn-vec.hpp" -#include "fattn.hpp" +#include "dpct/helper-defs.inc" +#include "common-defs.inc" +#include "fattn-common-defs.inc" +#include "fattn-tile-defs.inc" +#include "fattn-vec-defs.inc" +#include "common-defs.inc" #define FATTN_VEC_CASE(D, type_K, type_V) \ diff --git a/ggml/src/ggml-sycl/gated_delta_net.hpp b/ggml/src/ggml-sycl/gated_delta_net.hpp deleted file mode 100644 index a3308ee8763c..000000000000 --- a/ggml/src/ggml-sycl/gated_delta_net.hpp +++ /dev/null @@ -1,8 +0,0 @@ -#pragma once - -#include -#include "dpct/helper.hpp" -#include "common.hpp" -#include "ggml.h" - -void ggml_sycl_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-sycl/gated_delta_net.cpp b/ggml/src/ggml-sycl/gated_delta_net.inc similarity index 99% rename from ggml/src/ggml-sycl/gated_delta_net.cpp rename to ggml/src/ggml-sycl/gated_delta_net.inc index ebc587524bf4..f33e23bbf8cd 100644 --- a/ggml/src/ggml-sycl/gated_delta_net.cpp +++ b/ggml/src/ggml-sycl/gated_delta_net.inc @@ -1,8 +1,7 @@ #include -#include "dpct/helper.hpp" -#include "common.hpp" -#include "ggml.h" -#include "gated_delta_net.hpp" +#include "dpct/helper-defs.inc" +#include "common-defs.inc" +#include "ggml.h.inc" #include diff --git a/ggml/src/ggml-sycl/gemm.hpp b/ggml/src/ggml-sycl/gemm-defs.inc similarity index 99% rename from ggml/src/ggml-sycl/gemm.hpp rename to ggml/src/ggml-sycl/gemm-defs.inc index c202da110beb..a1f05fc98282 100644 --- a/ggml/src/ggml-sycl/gemm.hpp +++ b/ggml/src/ggml-sycl/gemm-defs.inc @@ -13,7 +13,7 @@ #ifndef GGML_SYCL_GEMM_HPP #define GGML_SYCL_GEMM_HPP -#include "ggml-sycl.h" +#include "ggml-sycl.h.inc" #if GGML_SYCL_DNNL diff --git a/ggml/src/ggml-sycl/getrows.hpp b/ggml/src/ggml-sycl/getrows.hpp deleted file mode 100644 index 1c560cd9f894..000000000000 --- a/ggml/src/ggml-sycl/getrows.hpp +++ /dev/null @@ -1,20 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_GETROWS_HPP -#define GGML_SYCL_GETROWS_HPP - -#include "common.hpp" - -void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor *dst); - -#endif // GGML_SYCL_GETROWS_HPP diff --git a/ggml/src/ggml-sycl/getrows.cpp b/ggml/src/ggml-sycl/getrows.inc similarity index 98% rename from ggml/src/ggml-sycl/getrows.cpp rename to ggml/src/ggml-sycl/getrows.inc index 03f8dd907485..ebf42f760ed4 100644 --- a/ggml/src/ggml-sycl/getrows.cpp +++ b/ggml/src/ggml-sycl/getrows.inc @@ -10,10 +10,10 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // -#include "ggml-impl.h" -#include "common.hpp" -#include "dequantize.hpp" -#include "getrows.hpp" +#include "ggml-impl-defs.inc" +#include "common-defs.inc" +#include "dequantize-defs.inc" +#include "common-defs.inc" template diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.inc similarity index 99% rename from ggml/src/ggml-sycl/ggml-sycl.cpp rename to ggml/src/ggml-sycl/ggml-sycl.inc index f06147eeeb8c..1e059b4e5606 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.inc @@ -35,25 +35,17 @@ #endif #include -#include "ggml.h" -#include "ggml-sycl.h" -#include "ggml-impl.h" -#include "ggml-backend-impl.h" - -#include "ggml-sycl/add-id.hpp" -#include "ggml-sycl/backend.hpp" -#include "ggml-sycl/common.hpp" -#include "ggml-sycl/element_wise.hpp" -#include "ggml-sycl/gemm.hpp" -#include "ggml-sycl/getrows.hpp" -#include "ggml-sycl/norm.hpp" -#include "ggml-sycl/presets.hpp" -#include "ggml-sycl/quantize.hpp" -#include "ggml-sycl/repeat_back.hpp" -#include "ggml-sycl/set_rows.hpp" -#include "ggml-sycl/set.hpp" -#include "ggml-sycl/ssm_conv.hpp" -#include "ggml-sycl/sycl_hw.hpp" +#include "ggml.h.inc" +#include "ggml-sycl.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-backend-impl-defs.inc" + +#include "ggml-sycl/backend-defs.inc" +#include "ggml-sycl/common-defs.inc" +#include "ggml-sycl/element_wise-defs.inc" +#include "ggml-sycl/gemm-defs.inc" +#include "ggml-sycl/presets-defs.inc" +#include "ggml-sycl/quantize-defs.inc" static bool g_sycl_loaded = false; diff --git a/ggml/src/ggml-sycl/gla.hpp b/ggml/src/ggml-sycl/gla.hpp deleted file mode 100644 index 607cf3a7f304..000000000000 --- a/ggml/src/ggml-sycl/gla.hpp +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef GGML_SYCL_GLA_HPP -#define GGML_SYCL_GLA_HPP - -#include "common.hpp" - -void ggml_sycl_op_gated_linear_attn(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -#endif // GGML_SYCL_GLA_HPP diff --git a/ggml/src/ggml-sycl/gla.cpp b/ggml/src/ggml-sycl/gla.inc similarity index 99% rename from ggml/src/ggml-sycl/gla.cpp rename to ggml/src/ggml-sycl/gla.inc index 879184fdd311..2183e257ac84 100644 --- a/ggml/src/ggml-sycl/gla.cpp +++ b/ggml/src/ggml-sycl/gla.inc @@ -1,6 +1,6 @@ #include -#include "common.hpp" +#include "common-defs.inc" template static void gated_linear_attn_f32_kernel(const dpct::queue_ptr stream, u_int B, u_int T, u_int C, u_int H, float scale, diff --git a/ggml/src/ggml-sycl/im2col.hpp b/ggml/src/ggml-sycl/im2col.hpp deleted file mode 100644 index dbbb248ddb4f..000000000000 --- a/ggml/src/ggml-sycl/im2col.hpp +++ /dev/null @@ -1,21 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_IM2COL_HPP -#define GGML_SYCL_IM2COL_HPP - -#include "common.hpp" - -void ggml_sycl_op_im2col( - ggml_backend_sycl_context & ctx, ggml_tensor *dst); - -#endif // GGML_SYCL_IM2COL_HPP diff --git a/ggml/src/ggml-sycl/im2col.cpp b/ggml/src/ggml-sycl/im2col.inc similarity index 99% rename from ggml/src/ggml-sycl/im2col.cpp rename to ggml/src/ggml-sycl/im2col.inc index 6d75d34d83f4..aa062b9822b1 100644 --- a/ggml/src/ggml-sycl/im2col.cpp +++ b/ggml/src/ggml-sycl/im2col.inc @@ -10,12 +10,12 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // -#include "im2col.hpp" +#include "common-defs.inc" #include #include // For std::is_same_v -#include "ggml.h" +#include "ggml.h.inc" template static void im2col_kernel(const float * x, T * dst, int64_t batch_offset, int64_t offset_delta, int64_t IC, int64_t IW, diff --git a/ggml/src/ggml-sycl/mmq.hpp b/ggml/src/ggml-sycl/mmq.hpp deleted file mode 100644 index 3f5297aaa537..000000000000 --- a/ggml/src/ggml-sycl/mmq.hpp +++ /dev/null @@ -1,33 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_MMQ_HPP -#define GGML_SYCL_MMQ_HPP - -#include "common.hpp" - -void ggml_sycl_op_mul_mat_q( - ggml_backend_sycl_context & ctx, - const ggml_tensor* src0, - const ggml_tensor* src1, - ggml_tensor* dst, - const char* src0_dd_i, - const float* src1_ddf_i, - const char* src1_ddq_i, - float* dst_dd_i, - const int64_t row_low, - const int64_t row_high, - const int64_t src1_ncols, - const int64_t src1_padded_row_size, - const dpct::queue_ptr& stream); - -#endif // GGML_SYCL_MMQ_HPP diff --git a/ggml/src/ggml-sycl/mmq.cpp b/ggml/src/ggml-sycl/mmq.inc similarity index 99% rename from ggml/src/ggml-sycl/mmq.cpp rename to ggml/src/ggml-sycl/mmq.inc index ffb272aa2837..e01167bae896 100644 --- a/ggml/src/ggml-sycl/mmq.cpp +++ b/ggml/src/ggml-sycl/mmq.inc @@ -10,8 +10,8 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // -#include "mmq.hpp" -#include "vecdotq.hpp" +#include "common-defs.inc" +#include "vecdotq-defs.inc" typedef void (*allocate_tiles_sycl_t)( int** x_ql, diff --git a/ggml/src/ggml-sycl/mmvq.hpp b/ggml/src/ggml-sycl/mmvq.hpp deleted file mode 100644 index d674dc1d61ea..000000000000 --- a/ggml/src/ggml-sycl/mmvq.hpp +++ /dev/null @@ -1,43 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_MMVQ_HPP -#define GGML_SYCL_MMVQ_HPP - -#include "common.hpp" - - -void ggml_sycl_op_mul_mat_vec_q( - ggml_backend_sycl_context & ctx, - const ggml_tensor *src0, const ggml_tensor *src1, ggml_tensor *dst, - const char *src0_dd_i, const float *src1_ddf_i, const char *src1_ddq_i, - float *dst_dd_i, const int64_t row_low, const int64_t row_high, - const int64_t src1_ncols, const int64_t src1_padded_row_size, - const dpct::queue_ptr &stream); - -// Requires standard (non-reorder) block layout for src0. -// Returns false if src0_type isn't handled; caller should fall back. -bool ggml_sycl_mul_mat_vec_q_id( - enum ggml_type src0_type, - const void * vx_base, // start of stacked expert weights - const void * vy, // pre-quantized src1 (Q8_1) - const int32_t * ids_dev, // device-side int32, length n_experts_used - float * dst_base, - int ncols, - int nrows, - int n_experts_used, - size_t expert_weight_stride, // bytes between experts in vx_base - size_t dst_row_stride, // bytes between dst rows - size_t src1_row_stride, // 0 = shared src1, else per-expert stride in bytes - dpct::queue_ptr stream); - -#endif // GGML_SYCL_MMVQ_HPP diff --git a/ggml/src/ggml-sycl/mmvq.cpp b/ggml/src/ggml-sycl/mmvq.inc similarity index 99% rename from ggml/src/ggml-sycl/mmvq.cpp rename to ggml/src/ggml-sycl/mmvq.inc index 8fa2198f35af..d17ec241c91f 100644 --- a/ggml/src/ggml-sycl/mmvq.cpp +++ b/ggml/src/ggml-sycl/mmvq.inc @@ -1,9 +1,9 @@ -#include "mmvq.hpp" +#include "common-defs.inc" -#include "ggml.h" -#include "common.hpp" -#include "quants.hpp" -#include "vecdotq.hpp" +#include "ggml.h.inc" +#include "common-defs.inc" +#include "quants-defs.inc" +#include "vecdotq-defs.inc" template static void mul_mat_vec_q_reorder(const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, diff --git a/ggml/src/ggml-sycl/norm.hpp b/ggml/src/ggml-sycl/norm.hpp deleted file mode 100644 index 8cb885eb2eed..000000000000 --- a/ggml/src/ggml-sycl/norm.hpp +++ /dev/null @@ -1,28 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_NORM_HPP -#define GGML_SYCL_NORM_HPP - -#include "common.hpp" - -void ggml_sycl_op_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst); - -void ggml_sycl_op_rms_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst); - -void ggml_sycl_op_rms_norm_back(ggml_backend_sycl_context& ctx, ggml_tensor* dst); - -void ggml_sycl_op_group_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst); - -void ggml_sycl_op_l2_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst); - -#endif // GGML_SYCL_NORM_HPP diff --git a/ggml/src/ggml-sycl/norm.cpp b/ggml/src/ggml-sycl/norm.inc similarity index 99% rename from ggml/src/ggml-sycl/norm.cpp rename to ggml/src/ggml-sycl/norm.inc index 09fce1280adc..7ca562f9d9c5 100644 --- a/ggml/src/ggml-sycl/norm.cpp +++ b/ggml/src/ggml-sycl/norm.inc @@ -1,6 +1,6 @@ -#include "norm.hpp" -#include "ggml-sycl/common.hpp" -#include "ggml-sycl/presets.hpp" +#include "common-defs.inc" +#include "ggml-sycl/common-defs.inc" +#include "ggml-sycl/presets-defs.inc" static void norm_f32(const float* x, float* dst, const int ncols, const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, const sycl::nd_item<3>& item_ct1, sycl::float2* s_sum, int block_size) { diff --git a/ggml/src/ggml-sycl/outprod.hpp b/ggml/src/ggml-sycl/outprod.hpp deleted file mode 100644 index f50413d3f7a2..000000000000 --- a/ggml/src/ggml-sycl/outprod.hpp +++ /dev/null @@ -1,10 +0,0 @@ -#ifndef GGML_SYCL_OUTPROD_HPP -#define GGML_SYCL_OUTPROD_HPP - -#include "common.hpp" - -void ggml_sycl_op_out_prod(ggml_backend_sycl_context& ctx, ggml_tensor* dst); - - -#endif // GGML_SYCL_OUTPROD_HPP - diff --git a/ggml/src/ggml-sycl/outprod.cpp b/ggml/src/ggml-sycl/outprod.inc similarity index 98% rename from ggml/src/ggml-sycl/outprod.cpp rename to ggml/src/ggml-sycl/outprod.inc index f52b11f0d6ed..191ce8e2c58e 100644 --- a/ggml/src/ggml-sycl/outprod.cpp +++ b/ggml/src/ggml-sycl/outprod.inc @@ -1,4 +1,4 @@ -#include "outprod.hpp" +#include "common-defs.inc" void ggml_sycl_op_out_prod(ggml_backend_sycl_context& ctx, ggml_tensor* dst) { scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2); diff --git a/ggml/src/ggml-sycl/pad.hpp b/ggml/src/ggml-sycl/pad.hpp deleted file mode 100644 index b099e9b73a4a..000000000000 --- a/ggml/src/ggml-sycl/pad.hpp +++ /dev/null @@ -1,24 +0,0 @@ -// -// MIT license -// Copyright (C) 2025 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_PAD_HPP -#define GGML_SYCL_PAD_HPP - -#include "common.hpp" - -#define SYCL_PAD_BLOCK_SIZE 256 - -void ggml_sycl_pad(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -void ggml_sycl_op_pad(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -#endif // GGML_SYCL_PAD_HPP diff --git a/ggml/src/ggml-sycl/pad.cpp b/ggml/src/ggml-sycl/pad.inc similarity index 98% rename from ggml/src/ggml-sycl/pad.cpp rename to ggml/src/ggml-sycl/pad.inc index f989c5e4b8bb..7d6963dda1af 100644 --- a/ggml/src/ggml-sycl/pad.cpp +++ b/ggml/src/ggml-sycl/pad.inc @@ -10,8 +10,8 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // -//#include "common.hpp" -#include "pad.hpp" +//#include "common-defs.inc" +#include "common-defs.inc" static void pad_f32(const float * src, float * dst, const int lp0, const int rp0, const int lp1, const int rp1, diff --git a/ggml/src/ggml-sycl/pad_reflect_1d.hpp b/ggml/src/ggml-sycl/pad_reflect_1d.hpp deleted file mode 100644 index 45aaf9a9111f..000000000000 --- a/ggml/src/ggml-sycl/pad_reflect_1d.hpp +++ /dev/null @@ -1,10 +0,0 @@ -#ifndef GGML_SYCL_PAD_REFLECT_1D_HPP -#define GGML_SYCL_PAD_REFLECT_1D_HPP - -#include "common.hpp" - -#define SYCL_PAD_REFLECT_1D_BLOCK_SIZE 256 - -void ggml_sycl_op_pad_reflect_1d(ggml_backend_sycl_context& ctx, ggml_tensor* dst); - -#endif // GGML_SYCL_PAD_REFLECT_1D_HPP diff --git a/ggml/src/ggml-sycl/pad_reflect_1d.cpp b/ggml/src/ggml-sycl/pad_reflect_1d.inc similarity index 98% rename from ggml/src/ggml-sycl/pad_reflect_1d.cpp rename to ggml/src/ggml-sycl/pad_reflect_1d.inc index 85e993628c65..6670c4756182 100644 --- a/ggml/src/ggml-sycl/pad_reflect_1d.cpp +++ b/ggml/src/ggml-sycl/pad_reflect_1d.inc @@ -1,4 +1,6 @@ -#include "pad_reflect_1d.hpp" +#include "common-defs.inc" + +#define SYCL_PAD_REFLECT_1D_BLOCK_SIZE 256 static void pad_reflect_1d_kernel_f32( const void *__restrict__ src0, void *__restrict__ dst, const int64_t ne0, diff --git a/ggml/src/ggml-sycl/presets.hpp b/ggml/src/ggml-sycl/presets-defs.inc similarity index 100% rename from ggml/src/ggml-sycl/presets.hpp rename to ggml/src/ggml-sycl/presets-defs.inc diff --git a/ggml/src/ggml-sycl/quantize.hpp b/ggml/src/ggml-sycl/quantize-defs.inc similarity index 99% rename from ggml/src/ggml-sycl/quantize.hpp rename to ggml/src/ggml-sycl/quantize-defs.inc index b5c7a54b7924..f0cdc1370339 100644 --- a/ggml/src/ggml-sycl/quantize.hpp +++ b/ggml/src/ggml-sycl/quantize-defs.inc @@ -21,7 +21,7 @@ #include -#include "ggml-sycl/dpct/helper.hpp" +#include "ggml-sycl/dpct/helper-defs.inc" template __dpct_inline__ static void quantize_q8_1_impl(const float * __restrict__ x, diff --git a/ggml/src/ggml-sycl/quants.hpp b/ggml/src/ggml-sycl/quants-defs.inc similarity index 97% rename from ggml/src/ggml-sycl/quants.hpp rename to ggml/src/ggml-sycl/quants-defs.inc index 1f5b62740a8e..5846cf5732ca 100644 --- a/ggml/src/ggml-sycl/quants.hpp +++ b/ggml/src/ggml-sycl/quants-defs.inc @@ -16,8 +16,8 @@ #include -#include "ggml-common.h" -#include "ggml.h" +#include "ggml-common-defs.inc" +#include "ggml.h.inc" namespace ggml_sycl_reordered { @@ -38,7 +38,7 @@ template struct block_q_t; // for quantization types that has low and high bits split, qr is calculated with // using the lower bits, e.g for Q6 quants QR6 is 2 // qi number of 32 bit integers needed to represent all the quants from a block (`qs` field) -// See ggml-common.h to see how these are calculated +// See ggml-common-defs.inc to see how these are calculated template <> struct block_q_t { struct traits { static constexpr uint32_t qk = QK4_0; diff --git a/ggml/src/ggml-sycl/repeat_back.hpp b/ggml/src/ggml-sycl/repeat_back.hpp deleted file mode 100644 index 17a87f3e159b..000000000000 --- a/ggml/src/ggml-sycl/repeat_back.hpp +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef GGML_SYCL_REPEAT_BACK_HPP -#define GGML_SYCL_REPEAT_BACK_HPP - -#include "common.hpp" - -void ggml_sycl_op_repeat_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -#endif // GGML_SYCL_REPEAT_BACK_HPP diff --git a/ggml/src/ggml-sycl/repeat_back.cpp b/ggml/src/ggml-sycl/repeat_back.inc similarity index 98% rename from ggml/src/ggml-sycl/repeat_back.cpp rename to ggml/src/ggml-sycl/repeat_back.inc index 845b48468c1d..a7b096648763 100644 --- a/ggml/src/ggml-sycl/repeat_back.cpp +++ b/ggml/src/ggml-sycl/repeat_back.inc @@ -1,6 +1,4 @@ -#include "repeat_back.hpp" - -#include "common.hpp" +#include "common-defs.inc" #define GGML_ASSERT_TENSOR_FITS_INT(t) \ GGML_ASSERT((t)->ne[0] < INT_MAX && (t)->ne[1] < INT_MAX && (t)->ne[2] < INT_MAX && (t)->ne[3] < INT_MAX) diff --git a/ggml/src/ggml-sycl/roll.hpp b/ggml/src/ggml-sycl/roll.hpp deleted file mode 100644 index 97dc03d64b24..000000000000 --- a/ggml/src/ggml-sycl/roll.hpp +++ /dev/null @@ -1,20 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_ROLL_HPP -#define GGML_SYCL_ROLL_HPP - -#include "common.hpp" - -void ggml_sycl_roll(ggml_backend_sycl_context & ctx, ggml_tensor *dst); - -#endif // GGML_SYCL_ROLL_HPP diff --git a/ggml/src/ggml-sycl/roll.cpp b/ggml/src/ggml-sycl/roll.inc similarity index 98% rename from ggml/src/ggml-sycl/roll.cpp rename to ggml/src/ggml-sycl/roll.inc index 1e05181789c2..8d60b7dbcb56 100644 --- a/ggml/src/ggml-sycl/roll.cpp +++ b/ggml/src/ggml-sycl/roll.inc @@ -1,5 +1,5 @@ -#include "roll.hpp" -#include "common.hpp" +#include "common-defs.inc" +#include "common-defs.inc" using namespace sycl; diff --git a/ggml/src/ggml-sycl/rope.hpp b/ggml/src/ggml-sycl/rope.hpp deleted file mode 100644 index b95a585808be..000000000000 --- a/ggml/src/ggml-sycl/rope.hpp +++ /dev/null @@ -1,26 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_ROPE_HPP -#define GGML_SYCL_ROPE_HPP - -#include "common.hpp" - -#define SYCL_ROPE_BLOCK_SIZE 256 - -void ggml_sycl_rope(ggml_backend_sycl_context & ctx, ggml_tensor *dst); - -void ggml_sycl_rope_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -void ggml_sycl_rope_fused(ggml_backend_sycl_context & ctx, ggml_tensor * dst, ggml_tensor * set_rows); - -#endif // GGML_SYCL_ROPE_HPP diff --git a/ggml/src/ggml-sycl/rope.cpp b/ggml/src/ggml-sycl/rope.inc similarity index 99% rename from ggml/src/ggml-sycl/rope.cpp rename to ggml/src/ggml-sycl/rope.inc index 9d83a1e9fa09..31f89299d19d 100644 --- a/ggml/src/ggml-sycl/rope.cpp +++ b/ggml/src/ggml-sycl/rope.inc @@ -1,7 +1,7 @@ -#include "rope.hpp" -#include "convert.hpp" -#include "ggml-sycl/common.hpp" -#include "ggml.h" +#include "common-defs.inc" +#include "convert-defs.inc" +#include "ggml-sycl/common-defs.inc" +#include "ggml.h.inc" struct rope_corr_dims { float v[2]; diff --git a/ggml/src/ggml-sycl/set.hpp b/ggml/src/ggml-sycl/set.hpp deleted file mode 100644 index 657d7ac9a7b0..000000000000 --- a/ggml/src/ggml-sycl/set.hpp +++ /dev/null @@ -1,5 +0,0 @@ -#pragma once -#include "backend.hpp" -#include "ggml.h" - -void ggml_sycl_op_set(ggml_backend_sycl_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-sycl/set.cpp b/ggml/src/ggml-sycl/set.inc similarity index 97% rename from ggml/src/ggml-sycl/set.cpp rename to ggml/src/ggml-sycl/set.inc index 381326d230ab..06227a3a4727 100644 --- a/ggml/src/ggml-sycl/set.cpp +++ b/ggml/src/ggml-sycl/set.inc @@ -1,7 +1,6 @@ -#include "presets.hpp" -#include "common.hpp" -#include "ggml.h" -#include "set.hpp" +#include "presets-defs.inc" +#include "common-defs.inc" +#include "ggml.h.inc" #include #include using namespace sycl; diff --git a/ggml/src/ggml-sycl/set_rows.hpp b/ggml/src/ggml-sycl/set_rows.hpp deleted file mode 100644 index 27fcc8f90175..000000000000 --- a/ggml/src/ggml-sycl/set_rows.hpp +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef GGML_SYCL_SET_ROWS_HPP -#define GGML_SYCL_SET_ROWS_HPP - -#include "common.hpp" - -void ggml_sycl_op_set_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -#endif // GGML_SYCL_SET_ROWS_HPP diff --git a/ggml/src/ggml-sycl/set_rows.cpp b/ggml/src/ggml-sycl/set_rows.inc similarity index 99% rename from ggml/src/ggml-sycl/set_rows.cpp rename to ggml/src/ggml-sycl/set_rows.inc index 8fb41943525c..9e04de531733 100644 --- a/ggml/src/ggml-sycl/set_rows.cpp +++ b/ggml/src/ggml-sycl/set_rows.inc @@ -1,5 +1,5 @@ -#include "set_rows.hpp" -#include "cpy.hpp" +#include "common-defs.inc" +#include "cpy-defs.inc" namespace utils { template diff --git a/ggml/src/ggml-sycl/softmax.hpp b/ggml/src/ggml-sycl/softmax.hpp deleted file mode 100644 index 23f1e5a9d65e..000000000000 --- a/ggml/src/ggml-sycl/softmax.hpp +++ /dev/null @@ -1,24 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_SOFTMAX_HPP -#define GGML_SYCL_SOFTMAX_HPP - -#include "common.hpp" - -#define SYCL_SOFT_MAX_BLOCK_SIZE 1024 - -void ggml_sycl_op_soft_max(ggml_backend_sycl_context &ctx, ggml_tensor *dst); - -void ggml_sycl_op_soft_max_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -#endif // GGML_SYCL_SOFTMAX_HPP diff --git a/ggml/src/ggml-sycl/softmax.cpp b/ggml/src/ggml-sycl/softmax.inc similarity index 99% rename from ggml/src/ggml-sycl/softmax.cpp rename to ggml/src/ggml-sycl/softmax.inc index fdf9b843e015..be95bcf1bf32 100644 --- a/ggml/src/ggml-sycl/softmax.cpp +++ b/ggml/src/ggml-sycl/softmax.inc @@ -1,4 +1,6 @@ -#include "softmax.hpp" +#include "common-defs.inc" + +#define SYCL_SOFT_MAX_BLOCK_SIZE 1024 #include #include #include diff --git a/ggml/src/ggml-sycl/ssm_conv.hpp b/ggml/src/ggml-sycl/ssm_conv.hpp deleted file mode 100644 index 1a8ad05f0c7f..000000000000 --- a/ggml/src/ggml-sycl/ssm_conv.hpp +++ /dev/null @@ -1,5 +0,0 @@ -#pragma once - -#include "common.hpp" - -void ggml_sycl_ssm_conv(ggml_backend_sycl_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-sycl/ssm_conv.cpp b/ggml/src/ggml-sycl/ssm_conv.inc similarity index 98% rename from ggml/src/ggml-sycl/ssm_conv.cpp rename to ggml/src/ggml-sycl/ssm_conv.inc index eea9a73d67e5..2036348be2c9 100644 --- a/ggml/src/ggml-sycl/ssm_conv.cpp +++ b/ggml/src/ggml-sycl/ssm_conv.inc @@ -1,5 +1,4 @@ -#include "ssm_conv.hpp" -#include "common.hpp" +#include "common-defs.inc" #include diff --git a/ggml/src/ggml-sycl/sycl_hw.hpp b/ggml/src/ggml-sycl/sycl_hw.hpp deleted file mode 100644 index a5d204625723..000000000000 --- a/ggml/src/ggml-sycl/sycl_hw.hpp +++ /dev/null @@ -1,38 +0,0 @@ -#ifndef SYCL_HW_HPP -#define SYCL_HW_HPP - -#include -#include -#include -#include - -#include - -namespace syclex = sycl::ext::oneapi::experimental; -using gpu_arch = sycl::ext::oneapi::experimental::architecture; - -// It's used to mark the GPU computing capacity -// The value must flow the order of performance. -enum sycl_intel_gpu_family { - GPU_FAMILY_UKNOWN = -1, - // iGPU without Xe core, before Meteor Lake iGPU(Xe) - GPU_FAMILY_IGPU_NON_XE = 0, - // iGPU with Xe core, Meteor Lake iGPU or newer. - GPU_FAMILY_IGPU_XE = 1, - // dGPU for gaming in client/data center (DG1/FLex 140 or newer). - GPU_FAMILY_DGPU_CLIENT_GAME = 2, - // dGPU for AI in cloud, PVC or newer. - GPU_FAMILY_DGPU_CLOUD = 3 -}; - -struct sycl_hw_info { - syclex::architecture arch; - const char* arch_name; - int32_t device_id; - std::string name; - sycl_intel_gpu_family gpu_family; -}; - -sycl_hw_info get_device_hw_info(sycl::device *device_ptr); - -#endif // SYCL_HW_HPP diff --git a/ggml/src/ggml-sycl/sycl_hw.cpp b/ggml/src/ggml-sycl/sycl_hw.inc similarity index 99% rename from ggml/src/ggml-sycl/sycl_hw.cpp rename to ggml/src/ggml-sycl/sycl_hw.inc index 03b0c37a3cd6..7281cf649e9f 100644 --- a/ggml/src/ggml-sycl/sycl_hw.cpp +++ b/ggml/src/ggml-sycl/sycl_hw.inc @@ -1,4 +1,4 @@ -#include "sycl_hw.hpp" +#include "common-defs.inc" using namespace std; diff --git a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq128-dv128.cpp b/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq128-dv128.cpp deleted file mode 100644 index f74e1202b838..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq128-dv128.cpp +++ /dev/null @@ -1,5 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-tile.hpp" - -DECL_FATTN_TILE_CASE(128, 128); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq40-dv40.cpp b/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq40-dv40.cpp deleted file mode 100644 index 8c8fb692c43f..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq40-dv40.cpp +++ /dev/null @@ -1,5 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-tile.hpp" - -DECL_FATTN_TILE_CASE(40, 40); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq512-dv512.cpp b/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq512-dv512.cpp deleted file mode 100644 index 9a6a18775666..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq512-dv512.cpp +++ /dev/null @@ -1,6 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-tile.hpp" - -DECL_FATTN_TILE_CASE(512, 512); - diff --git a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq576-dv512.cpp b/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq576-dv512.cpp deleted file mode 100644 index f218552e85f7..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq576-dv512.cpp +++ /dev/null @@ -1,5 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-tile.hpp" - -DECL_FATTN_TILE_CASE(576, 512); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq64-dv64.cpp b/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq64-dv64.cpp deleted file mode 100644 index 99303a53a3c9..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq64-dv64.cpp +++ /dev/null @@ -1,5 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-tile.hpp" - -DECL_FATTN_TILE_CASE(64, 64); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq72-dv72.cpp b/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq72-dv72.cpp deleted file mode 100644 index 50592768afd4..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq72-dv72.cpp +++ /dev/null @@ -1,5 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-tile.hpp" - -DECL_FATTN_TILE_CASE(72, 72); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq80-dv80.cpp b/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq80-dv80.cpp deleted file mode 100644 index 74f1ea5e90c3..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq80-dv80.cpp +++ /dev/null @@ -1,5 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-tile.hpp" - -DECL_FATTN_TILE_CASE(80, 80); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq96-dv96.cpp b/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq96-dv96.cpp deleted file mode 100644 index cefb46dddc78..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-tile-instance-dkq96-dv96.cpp +++ /dev/null @@ -1,5 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-tile.hpp" - -DECL_FATTN_TILE_CASE(96, 96); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-tile-instances.inc b/ggml/src/ggml-sycl/template-instances/fattn-tile-instances.inc new file mode 100644 index 000000000000..c59ab3db2e0b --- /dev/null +++ b/ggml/src/ggml-sycl/template-instances/fattn-tile-instances.inc @@ -0,0 +1,14 @@ +// Consolidated explicit instantiations for SYCL flash-attention tile cases. + +#include "../fattn-tile-defs.inc" + +DECL_FATTN_TILE_CASE( 40, 40); +DECL_FATTN_TILE_CASE( 64, 64); +DECL_FATTN_TILE_CASE( 72, 72); +DECL_FATTN_TILE_CASE( 80, 80); +DECL_FATTN_TILE_CASE( 96, 96); +DECL_FATTN_TILE_CASE(112, 112); +DECL_FATTN_TILE_CASE(128, 128); +DECL_FATTN_TILE_CASE(256, 256); +DECL_FATTN_TILE_CASE(512, 512); +DECL_FATTN_TILE_CASE(576, 512); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-f16.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-f16.cpp deleted file mode 100644 index 43ef94c118cd..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-f16.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_F16, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_F16); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q4_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q4_0.cpp deleted file mode 100644 index 9404061d456d..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q4_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_F16, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q4_1.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q4_1.cpp deleted file mode 100644 index a8bb9f52d0c8..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q4_1.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_F16, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q5_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q5_0.cpp deleted file mode 100644 index 7d61f6ab0afe..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q5_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_F16, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q5_1.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q5_1.cpp deleted file mode 100644 index 753bae09f838..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q5_1.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_F16, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q8_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q8_0.cpp deleted file mode 100644 index 546a93b25707..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-f16-q8_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_F16, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_F16, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_F16, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_F16, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-f16.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-f16.cpp deleted file mode 100644 index 53c8c2f2654b..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-f16.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_0, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q4_0, GGML_TYPE_F16); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q4_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q4_0.cpp deleted file mode 100644 index 5b409c55f217..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q4_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q4_1.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q4_1.cpp deleted file mode 100644 index 8c4ef588d634..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q4_1.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q5_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q5_0.cpp deleted file mode 100644 index 83f0a07552e5..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q5_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_0, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q4_0, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q5_1.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q5_1.cpp deleted file mode 100644 index 9df9b03bba43..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q5_1.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_0, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q4_0, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q8_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q8_0.cpp deleted file mode 100644 index 6980c2a65bb0..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_0-q8_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-f16.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-f16.cpp deleted file mode 100644 index bd61bc1dc2b2..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-f16.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_1, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q4_1, GGML_TYPE_F16); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q4_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q4_0.cpp deleted file mode 100644 index 492e229a58e8..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q4_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q4_1.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q4_1.cpp deleted file mode 100644 index 30f88a2ebd5f..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q4_1.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_1, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q4_1, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q5_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q5_0.cpp deleted file mode 100644 index db76663604ed..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q5_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q5_1.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q5_1.cpp deleted file mode 100644 index 1dbcc8a85a8f..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q5_1.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_1, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q4_1, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q8_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q8_0.cpp deleted file mode 100644 index d30996a62598..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q4_1-q8_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q4_1, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q4_1, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q4_1, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q4_1, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-f16.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-f16.cpp deleted file mode 100644 index bc0f635d922d..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-f16.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_0, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q5_0, GGML_TYPE_F16); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q4_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q4_0.cpp deleted file mode 100644 index 9e0378107cb1..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q4_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_0, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q5_0, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q4_1.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q4_1.cpp deleted file mode 100644 index a8535ac9156d..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q4_1.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q5_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q5_0.cpp deleted file mode 100644 index 43d4fae9a619..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q5_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_0, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q5_0, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q5_1.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q5_1.cpp deleted file mode 100644 index 23335a41640b..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q5_1.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q8_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q8_0.cpp deleted file mode 100644 index 52550a33757b..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_0-q8_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q5_0, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-f16.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-f16.cpp deleted file mode 100644 index 4651f14c0505..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-f16.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_1, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q5_1, GGML_TYPE_F16); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q4_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q4_0.cpp deleted file mode 100644 index 2310fd8792c3..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q4_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_1, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q5_1, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q4_1.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q4_1.cpp deleted file mode 100644 index d2494048bc14..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q4_1.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_1, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q5_1, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q5_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q5_0.cpp deleted file mode 100644 index be3a1fe97f54..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q5_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q5_1.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q5_1.cpp deleted file mode 100644 index be0a89409caf..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q5_1.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_1, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q5_1, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q8_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q8_0.cpp deleted file mode 100644 index 6781efcb0d22..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q5_1-q8_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-f16.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-f16.cpp deleted file mode 100644 index 43a70ae3543f..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-f16.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q8_0, GGML_TYPE_F16); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q8_0, GGML_TYPE_F16); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q4_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q4_0.cpp deleted file mode 100644 index fa7eb8163cac..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q4_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q4_1.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q4_1.cpp deleted file mode 100644 index 79d9cfbee965..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q4_1.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q8_0, GGML_TYPE_Q4_1); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q8_0, GGML_TYPE_Q4_1); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q5_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q5_0.cpp deleted file mode 100644 index 86befd5d3273..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q5_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q8_0, GGML_TYPE_Q5_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q8_0, GGML_TYPE_Q5_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q5_1.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q5_1.cpp deleted file mode 100644 index c2f619b0b166..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q5_1.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q8_0.cpp b/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q8_0.cpp deleted file mode 100644 index 7cf31f8b8a13..000000000000 --- a/ggml/src/ggml-sycl/template-instances/fattn-vec-instance-q8_0-q8_0.cpp +++ /dev/null @@ -1,8 +0,0 @@ -// This file has been autogenerated by generate_cu_files.py, do not edit manually. - -#include "../fattn-vec.hpp" - -DECL_FATTN_VEC_CASE( 64, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(256, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0); -DECL_FATTN_VEC_CASE(512, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-sycl/template-instances/fattn-vec-instances.inc b/ggml/src/ggml-sycl/template-instances/fattn-vec-instances.inc new file mode 100644 index 000000000000..7e035883f314 --- /dev/null +++ b/ggml/src/ggml-sycl/template-instances/fattn-vec-instances.inc @@ -0,0 +1,53 @@ +// Consolidated explicit instantiations for SYCL flash-attention vector cases. + +#include "../fattn-vec-defs.inc" + +#define DECL_FATTN_VEC_CASES(type_k, type_v) \ + DECL_FATTN_VEC_CASE( 64, type_k, type_v); \ + DECL_FATTN_VEC_CASE(128, type_k, type_v); \ + DECL_FATTN_VEC_CASE(256, type_k, type_v); \ + DECL_FATTN_VEC_CASE(512, type_k, type_v) + +DECL_FATTN_VEC_CASES(GGML_TYPE_F16, GGML_TYPE_F16); +DECL_FATTN_VEC_CASES(GGML_TYPE_F16, GGML_TYPE_Q4_0); +DECL_FATTN_VEC_CASES(GGML_TYPE_F16, GGML_TYPE_Q4_1); +DECL_FATTN_VEC_CASES(GGML_TYPE_F16, GGML_TYPE_Q5_0); +DECL_FATTN_VEC_CASES(GGML_TYPE_F16, GGML_TYPE_Q5_1); +DECL_FATTN_VEC_CASES(GGML_TYPE_F16, GGML_TYPE_Q8_0); + +DECL_FATTN_VEC_CASES(GGML_TYPE_Q4_0, GGML_TYPE_F16); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q4_0, GGML_TYPE_Q4_0); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q4_0, GGML_TYPE_Q4_1); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q4_0, GGML_TYPE_Q5_0); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q4_0, GGML_TYPE_Q5_1); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q4_0, GGML_TYPE_Q8_0); + +DECL_FATTN_VEC_CASES(GGML_TYPE_Q4_1, GGML_TYPE_F16); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q4_1, GGML_TYPE_Q4_0); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q4_1, GGML_TYPE_Q4_1); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q4_1, GGML_TYPE_Q5_0); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q4_1, GGML_TYPE_Q5_1); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q4_1, GGML_TYPE_Q8_0); + +DECL_FATTN_VEC_CASES(GGML_TYPE_Q5_0, GGML_TYPE_F16); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q5_0, GGML_TYPE_Q4_0); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q5_0, GGML_TYPE_Q4_1); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q5_0, GGML_TYPE_Q5_0); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q5_0, GGML_TYPE_Q5_1); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q5_0, GGML_TYPE_Q8_0); + +DECL_FATTN_VEC_CASES(GGML_TYPE_Q5_1, GGML_TYPE_F16); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q5_1, GGML_TYPE_Q4_0); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q5_1, GGML_TYPE_Q4_1); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q5_1, GGML_TYPE_Q5_0); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q5_1, GGML_TYPE_Q5_1); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q5_1, GGML_TYPE_Q8_0); + +DECL_FATTN_VEC_CASES(GGML_TYPE_Q8_0, GGML_TYPE_F16); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q8_0, GGML_TYPE_Q4_0); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q8_0, GGML_TYPE_Q4_1); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q8_0, GGML_TYPE_Q5_0); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q8_0, GGML_TYPE_Q5_1); +DECL_FATTN_VEC_CASES(GGML_TYPE_Q8_0, GGML_TYPE_Q8_0); + +#undef DECL_FATTN_VEC_CASES diff --git a/ggml/src/ggml-sycl/tsembd.hpp b/ggml/src/ggml-sycl/tsembd.hpp deleted file mode 100644 index 4c18748bbffc..000000000000 --- a/ggml/src/ggml-sycl/tsembd.hpp +++ /dev/null @@ -1,20 +0,0 @@ -// -// MIT license -// Copyright (C) 2024 Intel Corporation -// SPDX-License-Identifier: MIT -// - -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// - -#ifndef GGML_SYCL_TSEMBD_HPP -#define GGML_SYCL_TSEMBD_HPP - -#include "common.hpp" - -void ggml_sycl_op_timestep_embedding(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -#endif // GGML_SYCL_TSEMBD_HPP diff --git a/ggml/src/ggml-sycl/tsembd.cpp b/ggml/src/ggml-sycl/tsembd.inc similarity index 98% rename from ggml/src/ggml-sycl/tsembd.cpp rename to ggml/src/ggml-sycl/tsembd.inc index f2003794d3f5..3e78dd543446 100644 --- a/ggml/src/ggml-sycl/tsembd.cpp +++ b/ggml/src/ggml-sycl/tsembd.inc @@ -10,7 +10,7 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // -#include "tsembd.hpp" +#include "common-defs.inc" static void timestep_embedding_f32( const float * timesteps, float * dst, const int nb1, diff --git a/ggml/src/ggml-sycl/type.hpp b/ggml/src/ggml-sycl/type-defs.inc similarity index 100% rename from ggml/src/ggml-sycl/type.hpp rename to ggml/src/ggml-sycl/type-defs.inc diff --git a/ggml/src/ggml-sycl/upscale.hpp b/ggml/src/ggml-sycl/upscale.hpp deleted file mode 100644 index c36c1bdc970d..000000000000 --- a/ggml/src/ggml-sycl/upscale.hpp +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include -#include "dpct/helper.hpp" -#include "common.hpp" - -#define SYCL_UPSCALE_BLOCK_SIZE 256 - -void ggml_sycl_upscale(ggml_backend_sycl_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-sycl/upscale.cpp b/ggml/src/ggml-sycl/upscale.inc similarity index 99% rename from ggml/src/ggml-sycl/upscale.cpp rename to ggml/src/ggml-sycl/upscale.inc index e42cb419d838..e18cad765e01 100644 --- a/ggml/src/ggml-sycl/upscale.cpp +++ b/ggml/src/ggml-sycl/upscale.inc @@ -1,4 +1,4 @@ -#include "upscale.hpp" +#include "common-defs.inc" static void upscale_f32(const float * x, float * dst, const int nb00, const int nb01, const int nb02, const int nb03, diff --git a/ggml/src/ggml-sycl/vecdotq.hpp b/ggml/src/ggml-sycl/vecdotq-defs.inc similarity index 99% rename from ggml/src/ggml-sycl/vecdotq.hpp rename to ggml/src/ggml-sycl/vecdotq-defs.inc index 9253168e5ea2..483cd1fa48e3 100644 --- a/ggml/src/ggml-sycl/vecdotq.hpp +++ b/ggml/src/ggml-sycl/vecdotq-defs.inc @@ -13,10 +13,10 @@ #ifndef GGML_SYCL_VECDOTQ_HPP #define GGML_SYCL_VECDOTQ_HPP -#include "dpct/helper.hpp" -#include "ggml.h" -#include "type.hpp" -#include "quants.hpp" +#include "dpct/helper-defs.inc" +#include "ggml.h.inc" +#include "type-defs.inc" +#include "quants-defs.inc" typedef float (*vec_dot_q_sycl_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs); diff --git a/ggml/src/ggml-sycl/wkv.hpp b/ggml/src/ggml-sycl/wkv.hpp deleted file mode 100644 index 9f34a1001fd6..000000000000 --- a/ggml/src/ggml-sycl/wkv.hpp +++ /dev/null @@ -1,10 +0,0 @@ -#ifndef GGML_SYCL_WKV_HPP -#define GGML_SYCL_WKV_HPP - -#include "common.hpp" - -void ggml_sycl_op_rwkv_wkv6(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -void ggml_sycl_op_rwkv_wkv7(ggml_backend_sycl_context & ctx, ggml_tensor * dst); - -#endif // GGML_SYCL_WKV_HPP diff --git a/ggml/src/ggml-sycl/wkv.cpp b/ggml/src/ggml-sycl/wkv.inc similarity index 99% rename from ggml/src/ggml-sycl/wkv.cpp rename to ggml/src/ggml-sycl/wkv.inc index b56e0c2400f4..30b247746e5f 100644 --- a/ggml/src/ggml-sycl/wkv.cpp +++ b/ggml/src/ggml-sycl/wkv.inc @@ -1,5 +1,5 @@ #include -#include "wkv.hpp" +#include "common-defs.inc" constexpr int WKV_BLOCK_SIZE = 64; diff --git a/ggml/src/ggml-threading-rust/Cargo.toml b/ggml/src/ggml-threading-rust/Cargo.toml new file mode 100644 index 000000000000..b26c895673f7 --- /dev/null +++ b/ggml/src/ggml-threading-rust/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "ggml-threading-rust" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["staticlib", "rlib"] diff --git a/ggml/src/ggml-threading-rust/src/lib.rs b/ggml/src/ggml-threading-rust/src/lib.rs new file mode 100644 index 000000000000..a9a06c87984f --- /dev/null +++ b/ggml/src/ggml-threading-rust/src/lib.rs @@ -0,0 +1,150 @@ +use std::sync::{Mutex, MutexGuard}; + +static CRITICAL_SECTION: Mutex<()> = Mutex::new(()); +static mut CRITICAL_SECTION_GUARD: Option> = None; + +#[no_mangle] +pub extern "C" fn ggml_critical_section_start() { + let guard = CRITICAL_SECTION + .lock() + .unwrap_or_else(|err| err.into_inner()); + unsafe { + CRITICAL_SECTION_GUARD = Some(guard); + } +} + +#[no_mangle] +pub extern "C" fn ggml_critical_section_end() { + unsafe { + CRITICAL_SECTION_GUARD = None; + } +} + +#[no_mangle] +pub extern "C" fn ggml_op_is_empty_rust(op: i32) -> bool { + matches!(op, 0 | 36 | 37 | 38 | 39) +} + +#[no_mangle] +pub extern "C" fn ggml_bitset_size_rust(n: usize) -> usize { + (n + 31) >> 5 +} + +#[no_mangle] +pub extern "C" fn ggml_aligned_offset_rust(buffer: *const std::ffi::c_void, offset: usize, alignment: usize) -> usize { + debug_assert!(alignment != 0 && alignment.is_power_of_two()); + let address = buffer as usize + offset; + let align = (alignment - (address % alignment)) % alignment; + offset + align +} + +#[no_mangle] +pub unsafe extern "C" fn ggml_get_node_buffer_id_rust(node_buffer_ids: *const i32, index: i32) -> i32 { + if node_buffer_ids.is_null() { + return 0; + } + unsafe { *node_buffer_ids.add(index as usize) } +} + +#[no_mangle] +pub extern "C" fn ggml_buffer_address_less_rust( + a_chunk: i32, + a_offset: usize, + b_chunk: i32, + b_offset: usize, +) -> bool { + if a_chunk != b_chunk { + return a_chunk < b_chunk; + } + a_offset < b_offset +} + +#[no_mangle] +pub extern "C" fn ggml_isinf_fp16_rust(value: u16) -> bool { + (value & 0x7c00) == 0x7c00 && (value & 0x03ff) == 0 +} + +#[no_mangle] +pub extern "C" fn ggml_isnan_fp16_rust(value: u16) -> bool { + (value & 0x7c00) == 0x7c00 && (value & 0x03ff) != 0 +} + +#[no_mangle] +pub extern "C" fn ggml_is_invalid_e8m0_rust(value: u8) -> bool { + value == 0xff +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn locks_and_unlocks() { + ggml_critical_section_start(); + ggml_critical_section_end(); + } + + #[test] + fn identifies_empty_ops() { + assert!(ggml_op_is_empty_rust(0)); + assert!(ggml_op_is_empty_rust(36)); + assert!(ggml_op_is_empty_rust(37)); + assert!(ggml_op_is_empty_rust(38)); + assert!(ggml_op_is_empty_rust(39)); + assert!(!ggml_op_is_empty_rust(1)); + } + + #[test] + fn computes_bitset_word_count() { + assert_eq!(ggml_bitset_size_rust(0), 0); + assert_eq!(ggml_bitset_size_rust(1), 1); + assert_eq!(ggml_bitset_size_rust(32), 1); + assert_eq!(ggml_bitset_size_rust(33), 2); + } + + #[test] + fn computes_aligned_offsets() { + assert_eq!(ggml_aligned_offset_rust(std::ptr::null(), 0, 32), 0); + assert_eq!(ggml_aligned_offset_rust(std::ptr::null(), 1, 32), 32); + assert_eq!(ggml_aligned_offset_rust(std::ptr::null(), 64, 32), 64); + + let data = [0u8; 64]; + let base = data.as_ptr() as usize; + let expected = (16 - ((base + 3) % 16)) % 16 + 3; + assert_eq!(ggml_aligned_offset_rust(data.as_ptr().cast(), 3, 16), expected); + } + + #[test] + fn reads_node_buffer_ids_with_default() { + let ids = [7, 11, 13]; + unsafe { + assert_eq!(ggml_get_node_buffer_id_rust(std::ptr::null(), 2), 0); + assert_eq!(ggml_get_node_buffer_id_rust(ids.as_ptr(), 0), 7); + assert_eq!(ggml_get_node_buffer_id_rust(ids.as_ptr(), 2), 13); + } + } + + #[test] + fn compares_buffer_addresses() { + assert!(ggml_buffer_address_less_rust(0, 99, 1, 0)); + assert!(ggml_buffer_address_less_rust(1, 4, 1, 5)); + assert!(!ggml_buffer_address_less_rust(2, 0, 1, 99)); + assert!(!ggml_buffer_address_less_rust(1, 5, 1, 5)); + } + + #[test] + fn classifies_fp16_and_e8m0_values() { + assert!(ggml_isinf_fp16_rust(0x7c00)); + assert!(ggml_isinf_fp16_rust(0xfc00)); + assert!(!ggml_isinf_fp16_rust(0x7c01)); + assert!(!ggml_isinf_fp16_rust(0x3c00)); + + assert!(ggml_isnan_fp16_rust(0x7c01)); + assert!(ggml_isnan_fp16_rust(0x7fff)); + assert!(!ggml_isnan_fp16_rust(0x7c00)); + assert!(!ggml_isnan_fp16_rust(0x3c00)); + + assert!(ggml_is_invalid_e8m0_rust(0xff)); + assert!(!ggml_is_invalid_e8m0_rust(0xfe)); + } +} diff --git a/ggml/src/ggml-threading.cpp b/ggml/src/ggml-threading.cpp deleted file mode 100644 index 25a19eedb905..000000000000 --- a/ggml/src/ggml-threading.cpp +++ /dev/null @@ -1,12 +0,0 @@ -#include "ggml-threading.h" -#include - -std::mutex ggml_critical_section_mutex; - -void ggml_critical_section_start() { - ggml_critical_section_mutex.lock(); -} - -void ggml_critical_section_end(void) { - ggml_critical_section_mutex.unlock(); -} diff --git a/ggml/src/ggml-threading.h b/ggml/src/ggml-threading.h deleted file mode 100644 index dec2c8840aa3..000000000000 --- a/ggml/src/ggml-threading.h +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "ggml.h" - -#ifdef __cplusplus -extern "C" { -#endif - -GGML_API void ggml_critical_section_start(void); -GGML_API void ggml_critical_section_end(void); - -#ifdef __cplusplus -} -#endif diff --git a/ggml/src/ggml-virtgpu/CMakeLists.txt b/ggml/src/ggml-virtgpu/CMakeLists.txt index e6b020beb5bf..d3765744a394 100644 --- a/ggml/src/ggml-virtgpu/CMakeLists.txt +++ b/ggml/src/ggml-virtgpu/CMakeLists.txt @@ -29,24 +29,10 @@ if (NOT GGML_VIRTGPU_BACKEND STREQUAL "ONLY") target_compile_definitions(ggml PUBLIC "GGML_USE_VIRTGPU_FRONTEND") endif() + set_source_files_properties(virtgpu.cpp.inc PROPERTIES LANGUAGE CXX) ggml_add_backend_library(ggml-virtgpu - ggml-backend-buffer.cpp - ggml-backend.cpp - ggml-backend-device.cpp - ggml-backend-reg.cpp - ggml-backend-buffer-type.cpp - virtgpu-apir.h - virtgpu-forward.gen.h - virtgpu.cpp - virtgpu-shm.cpp - virtgpu-utils.cpp - virtgpu-forward-device.cpp - virtgpu-forward-buffer-type.cpp - virtgpu-forward-buffer.cpp - virtgpu-forward-backend.cpp - virtgpu-forward-impl.h - apir_cs_ggml-rpc-front.cpp - ../../include/ggml-virtgpu.h) + virtgpu.cpp.inc + ../../include/ggml-virtgpu.h.inc) target_include_directories(ggml-virtgpu PUBLIC /usr/include/libdrm/) diff --git a/ggml/src/ggml-virtgpu/apir_cs_ggml-rpc-front.cpp b/ggml/src/ggml-virtgpu/apir_cs_ggml-rpc-front.cpp deleted file mode 100644 index d2e87330a63d..000000000000 --- a/ggml/src/ggml-virtgpu/apir_cs_ggml-rpc-front.cpp +++ /dev/null @@ -1,87 +0,0 @@ -#include "backend/shared/apir_cs_rpc.h" -#include "ggml-backend-impl.h" -#include "ggml-impl.h" -#include "ggml-remoting.h" - -#include -#include -#include -#include - -apir_rpc_tensor apir_serialize_tensor(const ggml_tensor * tensor) { - apir_rpc_tensor result; - result.id = reinterpret_cast(tensor); - result.type = tensor->type; - if (tensor->buffer) { - ggml_backend_buffer_t buffer = tensor->buffer; - - result.buffer = BUFFER_TO_HOST_HANDLE(buffer); - } else { - result.buffer = 0; - } - for (uint32_t i = 0; i < GGML_MAX_DIMS; i++) { - result.ne[i] = tensor->ne[i]; - result.nb[i] = tensor->nb[i]; - } - result.op = tensor->op; - for (uint32_t i = 0; i < GGML_MAX_OP_PARAMS / sizeof(int32_t); i++) { - result.op_params[i] = tensor->op_params[i]; - } - result.flags = tensor->flags; - for (uint32_t i = 0; i < GGML_MAX_SRC; i++) { - result.src[i] = reinterpret_cast(tensor->src[i]); - } - result.view_src = reinterpret_cast(tensor->view_src); - result.view_offs = tensor->view_offs; - result.data = reinterpret_cast(tensor->data); - if (tensor->data) { - if (!tensor->buffer) { - GGML_ABORT("%s: tensor has data but not buffer", __func__); - } - // tensor->data is serialized as an offset to the buffer base address - result.data -= reinterpret_cast(BUFFER_TO_GGML_CONTEXT(tensor->buffer)->base); - } - snprintf(result.name, GGML_MAX_NAME, "%s", tensor->name); - return result; -} - -void apir_add_tensor(ggml_tensor * tensor, - std::vector & tensors, - std::unordered_set & visited) { - if (tensor == nullptr) { - return; - } - if (visited.find(tensor) != visited.end()) { - return; - } - visited.insert(tensor); - for (int i = 0; i < GGML_MAX_SRC; i++) { - apir_add_tensor(tensor->src[i], tensors, visited); - } - apir_add_tensor(tensor->view_src, tensors, visited); - tensors.push_back(apir_serialize_tensor(tensor)); -} - -void apir_serialize_graph(const ggml_cgraph * cgraph, std::vector & output) { - uint32_t n_nodes = cgraph->n_nodes; - std::vector tensors; - std::unordered_set visited; - for (uint32_t i = 0; i < n_nodes; i++) { - apir_add_tensor(cgraph->nodes[i], tensors, visited); - } - // serialization format: - // | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(apir_rpc_tensor)) | - uint32_t n_tensors = tensors.size(); - int output_size = - sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(apir_rpc_tensor); - output.resize(output_size, 0); - memcpy(output.data(), &n_nodes, sizeof(n_nodes)); - for (uint32_t i = 0; i < n_nodes; i++) { - memcpy(output.data() + sizeof(n_nodes) + i * sizeof(uint64_t), &cgraph->nodes[i], sizeof(uint64_t)); - } - uint32_t * out_ntensors = (uint32_t *) (output.data() + sizeof(n_nodes) + n_nodes * sizeof(uint64_t)); - *out_ntensors = n_tensors; - apir_rpc_tensor * out_tensors = - (apir_rpc_tensor *) (output.data() + sizeof(n_nodes) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t)); - memcpy(out_tensors, tensors.data(), n_tensors * sizeof(apir_rpc_tensor)); -} diff --git a/ggml/src/ggml-virtgpu/backend/CMakeLists.txt b/ggml/src/ggml-virtgpu/backend/CMakeLists.txt index 0b49c403b9a0..392f8525227d 100644 --- a/ggml/src/ggml-virtgpu/backend/CMakeLists.txt +++ b/ggml/src/ggml-virtgpu/backend/CMakeLists.txt @@ -3,19 +3,29 @@ cmake_policy(SET CMP0114 NEW) message(STATUS "Enable the VirtGPU/Virglrenderer backend library") +set_source_files_properties(backend.cpp.inc PROPERTIES LANGUAGE CXX) ggml_add_backend_library(ggml-virtgpu-backend - backend.cpp - backend-dispatched.cpp - backend-dispatched-backend.cpp - backend-dispatched-device.cpp - backend-dispatched-buffer.cpp - backend-dispatched-buffer-type.cpp - shared/api_remoting.h - shared/apir_backend.h - shared/apir_cs.h - apir_cs_ggml-rpc-back.cpp) + backend.cpp.inc) + +set(GGML_VIRTGPU_BACKEND_DISPATCHED_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/backend-dispatched-rust/Cargo.toml) +set(GGML_VIRTGPU_BACKEND_DISPATCHED_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/ggml-virtgpu-backend-dispatched) +set(GGML_VIRTGPU_BACKEND_DISPATCHED_RUST_LIB ${GGML_VIRTGPU_BACKEND_DISPATCHED_RUST_TARGET_DIR}/debug/${CMAKE_STATIC_LIBRARY_PREFIX}ggml_virtgpu_backend_dispatched_rust${CMAKE_STATIC_LIBRARY_SUFFIX}) + +add_custom_command( + OUTPUT ${GGML_VIRTGPU_BACKEND_DISPATCHED_RUST_LIB} + COMMAND ${CMAKE_COMMAND} -E env "RUSTFLAGS=-Crelocation-model=pic -Cpanic=abort" + ${CARGO_EXECUTABLE} build --manifest-path ${GGML_VIRTGPU_BACKEND_DISPATCHED_RUST_MANIFEST} --target-dir ${GGML_VIRTGPU_BACKEND_DISPATCHED_RUST_TARGET_DIR} + DEPENDS + ${GGML_VIRTGPU_BACKEND_DISPATCHED_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/backend-dispatched-rust/src/lib.rs +) +add_custom_target(ggml-virtgpu-backend-dispatched-rust-build DEPENDS ${GGML_VIRTGPU_BACKEND_DISPATCHED_RUST_LIB}) +add_library(ggml-virtgpu-backend-dispatched-rust STATIC IMPORTED GLOBAL) +set_target_properties(ggml-virtgpu-backend-dispatched-rust PROPERTIES IMPORTED_LOCATION ${GGML_VIRTGPU_BACKEND_DISPATCHED_RUST_LIB}) +add_dependencies(ggml-virtgpu-backend-dispatched-rust ggml-virtgpu-backend-dispatched-rust-build) +target_link_libraries(ggml-virtgpu-backend PRIVATE "$") target_compile_options(ggml-virtgpu-backend PRIVATE -std=c++20) -# Add include directory for ggml-backend-impl.h and other core headers +# Add include directory for ggml-backend-impl-defs.inc and other core headers target_include_directories(ggml-virtgpu-backend PRIVATE ../..) diff --git a/ggml/src/ggml-virtgpu/backend/apir_cs_ggml-rpc-back.cpp b/ggml/src/ggml-virtgpu/backend/apir_cs_ggml-rpc-back.cpp deleted file mode 100644 index 60a8a93bfb81..000000000000 --- a/ggml/src/ggml-virtgpu/backend/apir_cs_ggml-rpc-back.cpp +++ /dev/null @@ -1,115 +0,0 @@ -#include "ggml-backend-impl.h" -#include "ggml-impl.h" -#include "shared/apir_cs_rpc.h" - -#include -#include -#include -#include - -std::unordered_set backend_buffers; - -void apir_track_backend_buffer(ggml_backend_buffer_t buffer) { - backend_buffers.insert(buffer); -} - -bool apir_untrack_backend_buffer(ggml_backend_buffer_t buffer) { - auto it = backend_buffers.find(buffer); - if (it == backend_buffers.end()) { - return false; - } - - backend_buffers.erase(it); - return true; -} - -std::unordered_set apir_get_track_backend_buffers() { - return backend_buffers; -} - -ggml_tensor * apir_deserialize_tensor(ggml_context * ctx, const apir_rpc_tensor * tensor) { - ggml_tensor * result = - ggml_new_tensor_4d(ctx, (ggml_type) tensor->type, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); - for (uint32_t i = 0; i < GGML_MAX_DIMS; i++) { - result->nb[i] = tensor->nb[i]; - } - result->buffer = reinterpret_cast(tensor->buffer); - if (result->buffer && backend_buffers.find(result->buffer) == backend_buffers.end()) { - printf("WARNING: HOST BUFFER NOT FOUND | %p\n", (void *) result->buffer); - result->buffer = nullptr; - } - - uint64_t tensor_data = tensor->data; - if (result->buffer) { - // require that the tensor data does not go beyond the buffer end - uint64_t tensor_size = (uint64_t) ggml_nbytes(result); - uint64_t buffer_start = (uint64_t) ggml_backend_buffer_get_base(result->buffer); - uint64_t buffer_size = (uint64_t) ggml_backend_buffer_get_size(result->buffer); - - // tensor->data is serialized as an offset to the buffer base address - tensor_data += buffer_start; - - GGML_ASSERT(tensor_data + tensor_size >= tensor_data); // check for overflow - GGML_ASSERT(tensor_data >= buffer_start && tensor_data + tensor_size <= buffer_start + buffer_size); - } - - result->op = (ggml_op) tensor->op; - for (uint32_t i = 0; i < GGML_MAX_OP_PARAMS / sizeof(int32_t); i++) { - result->op_params[i] = tensor->op_params[i]; - } - result->flags = tensor->flags; - result->data = reinterpret_cast(tensor_data); - ggml_set_name(result, tensor->name); - return result; -} - -ggml_tensor * apir_create_node(uint64_t id, - ggml_context * ctx, - const std::unordered_map & tensor_ptrs, - std::unordered_map & tensor_map) { - if (id == 0) { - return nullptr; - } - if (tensor_map.find(id) != tensor_map.end()) { - return tensor_map[id]; - } - const apir_rpc_tensor * tensor = tensor_ptrs.at(id); - ggml_tensor * result = apir_deserialize_tensor(ctx, tensor); - if (result == nullptr) { - return nullptr; - } - tensor_map[id] = result; - for (int i = 0; i < GGML_MAX_SRC; i++) { - result->src[i] = apir_create_node(tensor->src[i], ctx, tensor_ptrs, tensor_map); - } - result->view_src = apir_create_node(tensor->view_src, ctx, tensor_ptrs, tensor_map); - result->view_offs = tensor->view_offs; - return result; -} - -ggml_cgraph * apir_deserialize_graph(uint32_t n_nodes, - uint32_t n_tensors, - const apir_rpc_tensor * tensors, - const uint64_t * nodes) { - size_t buf_size = ggml_tensor_overhead() * (n_nodes + n_tensors) + ggml_graph_overhead_custom(n_nodes, false); - ggml_init_params params = { - /*.mem_size =*/buf_size, - /*.mem_buffer =*/NULL, - /*.no_alloc =*/true, - }; - ggml_context * ctx = ggml_init(params); - ggml_cgraph * graph = ggml_new_graph_custom(ctx, n_nodes, false); - graph->n_nodes = n_nodes; - std::unordered_map tensor_ptrs; - for (uint32_t i = 0; i < n_tensors; i++) { - tensor_ptrs[tensors[i].id] = &tensors[i]; - } - std::unordered_map tensor_map; - for (uint32_t i = 0; i < n_nodes; i++) { - int64_t id; - memcpy(&id, &nodes[i], sizeof(id)); - graph->nodes[i] = apir_create_node(id, ctx, tensor_ptrs, tensor_map); - } - - return graph; -} diff --git a/ggml/src/ggml-virtgpu/backend/backend-convert.h b/ggml/src/ggml-virtgpu/backend/backend-convert.h deleted file mode 100644 index 1978d21f7efe..000000000000 --- a/ggml/src/ggml-virtgpu/backend/backend-convert.h +++ /dev/null @@ -1,13 +0,0 @@ -#include "shared/apir_backend.h" - -#define BUFFER_TO_HOST_HANDLE(name) ggml_buffer_to_apir_handle(name) - -static inline apir_buffer_host_handle_t ggml_buffer_to_apir_handle(ggml_backend_buffer_t buffer) { - // in the backend, the buffer handle is the buffer pointer - return (apir_buffer_host_handle_t) buffer; -} - -static inline apir_buffer_type_host_handle_t ggml_buffer_type_to_apir_handle(ggml_backend_buffer_type_t buft) { - // in the backend, the buffer handle is the buffer pointer - return (apir_buffer_type_host_handle_t) buft; -} diff --git a/ggml/src/ggml-virtgpu/backend/backend-dispatched-backend.cpp b/ggml/src/ggml-virtgpu/backend/backend-dispatched-backend.cpp deleted file mode 100644 index 03a037f1cbd9..000000000000 --- a/ggml/src/ggml-virtgpu/backend/backend-dispatched-backend.cpp +++ /dev/null @@ -1,102 +0,0 @@ -#include "backend-dispatched.h" -#include "backend-virgl-apir.h" -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "ggml-impl.h" -#include "shared/apir_backend.h" - -#include - -static uint32_t validate_graph_operation(size_t cgraph_size, uint32_t shmem_res_id, const char * operation) { - if (cgraph_size == 0) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Zero-size computation graph\n", operation); - return 1; - } - - // place-holder: validate that the size of shmem_res_id is <= cgraph_size - // need to add another method in the Virgl->APIR callback interface - GGML_UNUSED(shmem_res_id); - - return 0; // Valid -} - -uint32_t backend_backend_graph_compute(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - - static bool async_backend_initialized = false; - static bool async_backend; - - if (!async_backend_initialized) { - ggml_backend_dev_props props; - - dev->iface.get_props(dev, &props); - async_backend = props.caps.async; - async_backend_initialized = true; - } - - uint32_t shmem_res_id; - apir_decode_virtgpu_shmem_res_id(dec, &shmem_res_id); - - const void * shmem_data = ctx->iface->get_shmem_ptr(ctx->ctx_id, shmem_res_id); - if (!shmem_data) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Couldn't get the shmem addr from virgl\n", __func__); - apir_decoder_set_fatal(dec); - return 1; - } - size_t cgraph_size; - apir_decode_size_t(dec, &cgraph_size); - - if (validate_graph_operation(cgraph_size, shmem_res_id, __func__) != 0) { - apir_decoder_set_fatal(dec); - return 1; - } - - apir_decoder secondary_dec = apir_new_decoder((const char *) shmem_data, cgraph_size); - - ggml_cgraph * cgraph = apir_decode_ggml_cgraph(&secondary_dec, cgraph_size); - - if (!cgraph || apir_decoder_get_fatal(&secondary_dec)) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Failed to deserialize computation graph\n", __func__); - return 1; - } - - if (cgraph->n_nodes < 0 || cgraph->n_leafs < 0) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid negative node/leaf count: nodes=%d leafs=%d\n", __func__, - cgraph->n_nodes, cgraph->n_leafs); - return 1; - } - - ggml_status status; -#if APIR_BACKEND_CHECK_SUPPORTS_OP == 1 - for (int idx = 0; idx < cgraph->n_nodes; idx++) { - ggml_tensor * op = ggml_graph_node(cgraph, idx); - if (dev->iface.supports_op(dev, op)) { - continue; - } - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Graph node %d (%s) not supported by the backend\n", __func__, idx, - ggml_op_desc(op)); - - status = GGML_STATUS_ABORTED; - apir_encode_ggml_status(enc, &status); - - return 0; - } -#endif - - // Check if backend is properly initialized - if (!bck) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Backend not initialized (bck is null)\n", __func__); - - return 1; - } - - status = bck->iface.graph_compute(bck, cgraph); - - if (async_backend && bck->iface.synchronize) { - bck->iface.synchronize(bck); - } - - apir_encode_ggml_status(enc, &status); - - return 0; -} diff --git a/ggml/src/ggml-virtgpu/backend/backend-dispatched-buffer-type.cpp b/ggml/src/ggml-virtgpu/backend/backend-dispatched-buffer-type.cpp deleted file mode 100644 index c66dbaa9e8f3..000000000000 --- a/ggml/src/ggml-virtgpu/backend/backend-dispatched-buffer-type.cpp +++ /dev/null @@ -1,105 +0,0 @@ -#include "backend-dispatched.h" -#include "backend-virgl-apir.h" -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "ggml-impl.h" - -#include - -uint32_t backend_buffer_type_get_name(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - ggml_backend_buffer_type_t buft; - buft = apir_decode_ggml_buffer_type(dec); - - const char * string = buft->iface.get_name(buft); - - const size_t string_size = strlen(string) + 1; - apir_encode_array_size(enc, string_size); - apir_encode_char_array(enc, string, string_size); - - return 0; -} - -uint32_t backend_buffer_type_get_alignment(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - ggml_backend_buffer_type_t buft; - buft = apir_decode_ggml_buffer_type(dec); - - size_t value = buft->iface.get_alignment(buft); - apir_encode_size_t(enc, &value); - - return 0; -} - -uint32_t backend_buffer_type_get_max_size(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - ggml_backend_buffer_type_t buft; - buft = apir_decode_ggml_buffer_type(dec); - - size_t value = SIZE_MAX; - if (buft->iface.get_max_size) { - value = buft->iface.get_max_size(buft); - } - - apir_encode_size_t(enc, &value); - - return 0; -} - -/* APIR_COMMAND_TYPE_BUFFER_TYPE_IS_HOST is deprecated. Keeping the handler for backward compatibility. */ -uint32_t backend_buffer_type_is_host(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(dec); - const bool is_host = false; - - apir_encode_bool_t(enc, &is_host); - - return 0; -} - -uint32_t backend_buffer_type_alloc_buffer(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - ggml_backend_buffer_type_t buft; - buft = apir_decode_ggml_buffer_type(dec); - - size_t size; - apir_decode_size_t(dec, &size); - - ggml_backend_buffer_t buffer; - - buffer = buft->iface.alloc_buffer(buft, size); - - apir_encode_ggml_buffer(enc, buffer); - - if (buffer) { - apir_track_backend_buffer(buffer); - } - - return 0; -} - -uint32_t backend_buffer_type_get_alloc_size(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - ggml_backend_buffer_type_t buft; - buft = apir_decode_ggml_buffer_type(dec); - - const ggml_tensor * op = apir_decode_ggml_tensor_inplace(dec); - - // Check for decode error - if (op == nullptr) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Failed to decode tensor\n", __func__); - apir_decoder_set_fatal(dec); - return 1; - } - - size_t value; - if (buft->iface.get_alloc_size) { - value = buft->iface.get_alloc_size(buft, op); - } else { - value = ggml_nbytes(op); // Default fallback - } - - apir_encode_size_t(enc, &value); - - return 0; -} diff --git a/ggml/src/ggml-virtgpu/backend/backend-dispatched-buffer.cpp b/ggml/src/ggml-virtgpu/backend/backend-dispatched-buffer.cpp deleted file mode 100644 index 3ade8d99b4ed..000000000000 --- a/ggml/src/ggml-virtgpu/backend/backend-dispatched-buffer.cpp +++ /dev/null @@ -1,179 +0,0 @@ -#include "backend-dispatched.h" -#include "backend-virgl-apir.h" -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "ggml-impl.h" - -#include - -static uint32_t validate_buffer_operation(size_t offset, size_t size, const char * operation) { - // Only check for critical integer overflow - no arbitrary size limits - if (offset > SIZE_MAX - size) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Integer overflow in offset+size: %zu + %zu\n", operation, offset, size); - return 1; - } - - return 0; // Valid -} - -uint32_t backend_buffer_get_base(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - ggml_backend_buffer_t buffer; - buffer = apir_decode_ggml_buffer(dec); - - if (!buffer || apir_decoder_get_fatal(dec)) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid buffer handle from guest\n", __func__); - return 1; - } - - uintptr_t base = (uintptr_t) buffer->iface.get_base(buffer); - apir_encode_uintptr_t(enc, &base); - - return 0; -} - -uint32_t backend_buffer_set_tensor(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(enc); - - ggml_backend_buffer_t buffer; - buffer = apir_decode_ggml_buffer(dec); - - if (!buffer || apir_decoder_get_fatal(dec)) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid buffer handle from guest\n", __func__); - return 1; - } - - ggml_tensor * tensor; - // safe to remove the const qualifier here - tensor = (ggml_tensor *) (uintptr_t) apir_decode_ggml_tensor(dec); - - uint32_t shmem_res_id; - apir_decode_virtgpu_shmem_res_id(dec, &shmem_res_id); - - size_t offset; - apir_decode_size_t(dec, &offset); - - size_t size; - apir_decode_size_t(dec, &size); - - if (validate_buffer_operation(offset, size, __func__) != 0) { - return 1; - } - - void * shmem_data = ctx->iface->get_shmem_ptr(ctx->ctx_id, shmem_res_id); - - if (!shmem_data) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Couldn't get the shmem addr from virgl\n", __func__); - return 1; - } - - buffer->iface.set_tensor(buffer, tensor, shmem_data, offset, size); - - return 0; -} - -uint32_t backend_buffer_get_tensor(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(enc); - - ggml_backend_buffer_t buffer; - buffer = apir_decode_ggml_buffer(dec); - - if (!buffer || apir_decoder_get_fatal(dec)) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid buffer handle from guest\n", __func__); - return 1; - } - - const ggml_tensor * tensor; - // safe to remove the const qualifier here - tensor = apir_decode_ggml_tensor(dec); - - uint32_t shmem_res_id; - apir_decode_virtgpu_shmem_res_id(dec, &shmem_res_id); - - size_t offset; - apir_decode_size_t(dec, &offset); - - size_t size; - apir_decode_size_t(dec, &size); - - if (validate_buffer_operation(offset, size, __func__) != 0) { - return 1; - } - - void * shmem_data = ctx->iface->get_shmem_ptr(ctx->ctx_id, shmem_res_id); - if (!shmem_data) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Couldn't get the shmem addr from virgl\n", __func__); - return 1; - } - - buffer->iface.get_tensor(buffer, tensor, shmem_data, offset, size); - - return 0; -} - -uint32_t backend_buffer_cpy_tensor(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - - ggml_backend_buffer_t buffer; - buffer = apir_decode_ggml_buffer(dec); - - if (!buffer || apir_decoder_get_fatal(dec)) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid buffer handle from guest\n", __func__); - return 1; - } - - const ggml_tensor * src; - // safe to remove the const qualifier here - src = apir_decode_ggml_tensor(dec); - ggml_tensor * dst = (ggml_tensor *) (uintptr_t) apir_decode_ggml_tensor(dec); - - bool ret = buffer->iface.cpy_tensor(buffer, src, (ggml_tensor *) dst); - - apir_encode_bool_t(enc, &ret); - - return 0; -} - -uint32_t backend_buffer_clear(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(enc); - - ggml_backend_buffer_t buffer; - buffer = apir_decode_ggml_buffer(dec); - - if (!buffer || apir_decoder_get_fatal(dec)) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid buffer handle from guest\n", __func__); - return 1; - } - - uint8_t value; - apir_decode_uint8_t(dec, &value); - - buffer->iface.clear(buffer, value); - - return 0; -} - -uint32_t backend_buffer_free_buffer(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(enc); - - ggml_backend_buffer_t buffer; - buffer = apir_decode_ggml_buffer(dec); - - if (!buffer || apir_decoder_get_fatal(dec)) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid buffer handle from guest\n", __func__); - return 1; - } - - if (!apir_untrack_backend_buffer(buffer)) { - GGML_LOG_WARN(GGML_VIRTGPU_BCK "%s: unknown buffer %p\n", __func__, (void *) buffer); - return 1; - } - - buffer->iface.free_buffer(buffer); - - return 0; -} diff --git a/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp b/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp deleted file mode 100644 index c7acb8b51ce7..000000000000 --- a/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp +++ /dev/null @@ -1,148 +0,0 @@ -#include "backend-dispatched.h" -#include "backend-virgl-apir.h" -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "ggml-impl.h" - -#include - -uint32_t backend_device_get_device_count(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(ctx); - GGML_UNUSED(dec); - - int32_t dev_count = reg->iface.get_device_count(reg); - apir_encode_int32_t(enc, &dev_count); - - return 0; -} - -uint32_t backend_device_get_count(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(ctx); - GGML_UNUSED(dec); - - int32_t dev_count = reg->iface.get_device_count(reg); - apir_encode_int32_t(enc, &dev_count); - - return 0; -} - -uint32_t backend_device_get_name(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(dec); - - const char * string = dev->iface.get_name(dev); - - const size_t string_size = strlen(string) + 1; - apir_encode_array_size(enc, string_size); - apir_encode_char_array(enc, string, string_size); - - return 0; -} - -uint32_t backend_device_get_description(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(dec); - - const char * string = dev->iface.get_description(dev); - - const size_t string_size = strlen(string) + 1; - apir_encode_array_size(enc, string_size); - apir_encode_char_array(enc, string, string_size); - - return 0; -} - -uint32_t backend_device_get_type(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(dec); - - uint32_t type = dev->iface.get_type(dev); - apir_encode_uint32_t(enc, &type); - - return 0; -} - -uint32_t backend_device_get_memory(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(dec); - - size_t free, total; - dev->iface.get_memory(dev, &free, &total); - - apir_encode_size_t(enc, &free); - apir_encode_size_t(enc, &total); - - return 0; -} - -uint32_t backend_device_supports_op(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - - const ggml_tensor * op = apir_decode_ggml_tensor_inplace(dec); - - bool supports_op = dev->iface.supports_op(dev, op); - - apir_encode_bool_t(enc, &supports_op); - - return 0; -} - -uint32_t backend_device_get_buffer_type(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(dec); - - ggml_backend_buffer_type_t bufft = dev->iface.get_buffer_type(dev); - - apir_encode_ggml_buffer_type(enc, bufft); - - return 0; -} - -uint32_t backend_device_get_props(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(dec); - - ggml_backend_dev_props props; - dev->iface.get_props(dev, &props); - - apir_encode_bool_t(enc, &props.caps.async); - apir_encode_bool_t(enc, &props.caps.host_buffer); - apir_encode_bool_t(enc, &props.caps.buffer_from_host_ptr); - apir_encode_bool_t(enc, &props.caps.events); - - return 0; -} - -uint32_t backend_device_buffer_from_ptr(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { - GGML_UNUSED(ctx); - GGML_UNUSED(dec); - - uint32_t shmem_res_id; - apir_decode_virtgpu_shmem_res_id(dec, &shmem_res_id); - - void * shmem_ptr = ctx->iface->get_shmem_ptr(ctx->ctx_id, shmem_res_id); - if (!shmem_ptr) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Couldn't get the shmem addr from virgl\n", __func__); - apir_decoder_set_fatal(dec); - return 1; - } - - size_t size; - apir_decode_size_t(dec, &size); - size_t max_tensor_size; - apir_decode_size_t(dec, &max_tensor_size); - - ggml_backend_buffer_t buffer; - buffer = dev->iface.buffer_from_host_ptr(dev, shmem_ptr, size, max_tensor_size); - - apir_encode_ggml_buffer(enc, buffer); - apir_encode_ggml_buffer_type(enc, buffer->buft); - - if (buffer) { - apir_track_backend_buffer(buffer); - } - - return 0; -} diff --git a/ggml/src/ggml-virtgpu/backend/backend-dispatched-rust/Cargo.toml b/ggml/src/ggml-virtgpu/backend/backend-dispatched-rust/Cargo.toml new file mode 100644 index 000000000000..243a9cff1c97 --- /dev/null +++ b/ggml/src/ggml-virtgpu/backend/backend-dispatched-rust/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "ggml-virtgpu-backend-dispatched-rust" +version = "0.1.0" +edition = "2021" + +[lib] +crate-type = ["staticlib", "rlib"] diff --git a/ggml/src/ggml-virtgpu/backend/backend-dispatched-rust/src/lib.rs b/ggml/src/ggml-virtgpu/backend/backend-dispatched-rust/src/lib.rs new file mode 100644 index 000000000000..8f44bb264814 --- /dev/null +++ b/ggml/src/ggml-virtgpu/backend/backend-dispatched-rust/src/lib.rs @@ -0,0 +1,173 @@ +use std::ffi::{c_char, c_void}; + +const APIR_BACKEND_INITIALIZE_SUCCESS: u32 = 0; +const APIR_BACKEND_INITIALIZE_BACKEND_REG_FAILED: u32 = 6; +const APIR_BACKEND_INITIALIZE_ALREADY_INITED: u32 = 7; +const APIR_BACKEND_INITIALIZE_NO_DEVICE: u32 = 8; +const APIR_BACKEND_INITIALIZE_BACKEND_INIT_FAILED: u32 = 9; + +#[repr(C)] +struct GgmlBackendReg { + api_version: i32, + iface: GgmlBackendRegI, + context: *mut c_void, +} + +#[repr(C)] +struct GgmlBackendRegI { + get_name: Option *const c_char>, + get_device_count: Option usize>, + get_device: Option *mut GgmlBackendDevice>, + get_proc_address: Option *mut c_void>, +} + +#[repr(C)] +struct GgmlBackendDevice { + iface: GgmlBackendDeviceI, + reg: *mut GgmlBackendReg, + context: *mut c_void, +} + +#[repr(C)] +struct GgmlBackendDeviceI { + get_name: Option *const c_char>, + get_description: Option *const c_char>, + get_memory: Option, + get_type: Option i32>, + get_props: Option, + init_backend: Option *mut GgmlBackend>, +} + +#[repr(C)] +struct GgmlBackend { + _private: [u8; 0], +} + +type BackendRegFn = unsafe extern "C" fn() -> *mut GgmlBackendReg; + +#[no_mangle] +pub static mut reg: *mut c_void = std::ptr::null_mut(); +#[no_mangle] +pub static mut dev: *mut c_void = std::ptr::null_mut(); +#[no_mangle] +pub static mut bck: *mut c_void = std::ptr::null_mut(); + +#[no_mangle] +pub static mut timer_start: u64 = 0; +#[no_mangle] +pub static mut timer_total: u64 = 0; +#[no_mangle] +pub static mut timer_count: u64 = 0; + +unsafe fn initialize_from_fn(backend_reg_fn: BackendRegFn) -> u32 { + if !reg.is_null() { + return APIR_BACKEND_INITIALIZE_ALREADY_INITED; + } + + reg = backend_reg_fn().cast(); + if reg.is_null() { + return APIR_BACKEND_INITIALIZE_BACKEND_REG_FAILED; + } + + let reg_ptr = reg.cast::(); + + let Some(get_device_count) = (*reg_ptr).iface.get_device_count else { + return APIR_BACKEND_INITIALIZE_BACKEND_REG_FAILED; + }; + if get_device_count(reg_ptr) == 0 { + return APIR_BACKEND_INITIALIZE_NO_DEVICE; + } + + let Some(get_device) = (*reg_ptr).iface.get_device else { + return APIR_BACKEND_INITIALIZE_NO_DEVICE; + }; + dev = get_device(reg_ptr, 0).cast(); + if dev.is_null() { + return APIR_BACKEND_INITIALIZE_NO_DEVICE; + } + + let dev_ptr = dev.cast::(); + + let Some(init_backend) = (*dev_ptr).iface.init_backend else { + return APIR_BACKEND_INITIALIZE_BACKEND_INIT_FAILED; + }; + bck = init_backend(dev_ptr, std::ptr::null()).cast(); + if bck.is_null() { + return APIR_BACKEND_INITIALIZE_BACKEND_INIT_FAILED; + } + + APIR_BACKEND_INITIALIZE_SUCCESS +} + +#[no_mangle] +pub unsafe extern "C" fn backend_dispatch_initialize(ggml_backend_reg_fct_p: *mut c_void) -> u32 { + if ggml_backend_reg_fct_p.is_null() { + return APIR_BACKEND_INITIALIZE_BACKEND_REG_FAILED; + } + + let backend_reg_fn = std::mem::transmute::<*mut c_void, BackendRegFn>(ggml_backend_reg_fct_p); + initialize_from_fn(backend_reg_fn) +} + +#[cfg(test)] +mod tests { + use super::*; + + unsafe extern "C" fn fake_device_count(_reg: *mut GgmlBackendReg) -> usize { + 1 + } + + unsafe extern "C" fn fake_get_device(_reg: *mut GgmlBackendReg, _index: usize) -> *mut GgmlBackendDevice { + &raw mut FAKE_DEV + } + + unsafe extern "C" fn fake_init_backend( + _dev: *mut GgmlBackendDevice, + _params: *const c_char, + ) -> *mut GgmlBackend { + &raw mut FAKE_BACKEND + } + + static mut FAKE_BACKEND: GgmlBackend = GgmlBackend { _private: [] }; + static mut FAKE_DEV: GgmlBackendDevice = GgmlBackendDevice { + iface: GgmlBackendDeviceI { + get_name: None, + get_description: None, + get_memory: None, + get_type: None, + get_props: None, + init_backend: Some(fake_init_backend), + }, + reg: std::ptr::null_mut(), + context: std::ptr::null_mut(), + }; + static mut FAKE_REG: GgmlBackendReg = GgmlBackendReg { + api_version: 2, + iface: GgmlBackendRegI { + get_name: None, + get_device_count: Some(fake_device_count), + get_device: Some(fake_get_device), + get_proc_address: None, + }, + context: std::ptr::null_mut(), + }; + + unsafe extern "C" fn fake_reg() -> *mut GgmlBackendReg { + &raw mut FAKE_REG + } + + #[test] + fn initializes_backend_once() { + unsafe { + reg = std::ptr::null_mut(); + dev = std::ptr::null_mut(); + bck = std::ptr::null_mut(); + + assert_eq!(initialize_from_fn(fake_reg), APIR_BACKEND_INITIALIZE_SUCCESS); + assert!(!reg.is_null()); + assert!(!dev.is_null()); + assert!(!bck.is_null()); + assert_eq!(initialize_from_fn(fake_reg), APIR_BACKEND_INITIALIZE_ALREADY_INITED); + } + } +} diff --git a/ggml/src/ggml-virtgpu/backend/backend-dispatched.cpp b/ggml/src/ggml-virtgpu/backend/backend-dispatched.cpp deleted file mode 100644 index c80e4aabe1fe..000000000000 --- a/ggml/src/ggml-virtgpu/backend/backend-dispatched.cpp +++ /dev/null @@ -1,51 +0,0 @@ -#include "backend-dispatched.h" - -#include "backend-virgl-apir.h" -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "ggml-impl.h" - -#include - -ggml_backend_reg_t reg = NULL; -ggml_backend_dev_t dev = NULL; -ggml_backend_t bck = NULL; - -uint64_t timer_start = 0; -uint64_t timer_total = 0; -uint64_t timer_count = 0; - -uint32_t backend_dispatch_initialize(void * ggml_backend_reg_fct_p) { - if (reg != NULL) { - GGML_LOG_WARN(GGML_VIRTGPU_BCK "%s: already initialized\n", __func__); - return APIR_BACKEND_INITIALIZE_ALREADY_INITED; - } - ggml_backend_reg_t (*ggml_backend_reg_fct)(void) = (ggml_backend_reg_t (*)()) ggml_backend_reg_fct_p; - - reg = ggml_backend_reg_fct(); - if (reg == NULL) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: backend registration failed\n", __func__); - return APIR_BACKEND_INITIALIZE_BACKEND_REG_FAILED; - } - - size_t device_count = reg->iface.get_device_count(reg); - if (!device_count) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: no device found\n", __func__); - return APIR_BACKEND_INITIALIZE_NO_DEVICE; - } - - dev = reg->iface.get_device(reg, 0); - - if (!dev) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: failed to get device\n", __func__); - return APIR_BACKEND_INITIALIZE_NO_DEVICE; - } - - bck = dev->iface.init_backend(dev, NULL); - if (!bck) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: backend initialization failed\n", __func__); - return APIR_BACKEND_INITIALIZE_BACKEND_INIT_FAILED; - } - - return APIR_BACKEND_INITIALIZE_SUCCESS; -} diff --git a/ggml/src/ggml-virtgpu/backend/backend-dispatched.gen.h b/ggml/src/ggml-virtgpu/backend/backend-dispatched.gen.h deleted file mode 100644 index 3dc334e4ce4f..000000000000 --- a/ggml/src/ggml-virtgpu/backend/backend-dispatched.gen.h +++ /dev/null @@ -1,73 +0,0 @@ -#pragma once - -/* device */ -uint32_t backend_device_get_device_count(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_device_get_count(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_device_get_name(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_device_get_description(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_device_get_type(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_device_get_memory(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_device_supports_op(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_device_get_buffer_type(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_device_get_props(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_device_buffer_from_ptr(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); - -/* buffer-type */ -uint32_t backend_buffer_type_get_name(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_buffer_type_get_alignment(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_buffer_type_get_max_size(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -/* APIR_COMMAND_TYPE_BUFFER_TYPE_IS_HOST is deprecated. Keeping the handler for backward compatibility. */ -uint32_t backend_buffer_type_is_host(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_buffer_type_alloc_buffer(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_buffer_type_get_alloc_size(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); - -/* buffer */ -uint32_t backend_buffer_get_base(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_buffer_set_tensor(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_buffer_get_tensor(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_buffer_cpy_tensor(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_buffer_clear(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); -uint32_t backend_buffer_free_buffer(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); - -/* backend */ -uint32_t backend_backend_graph_compute(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); - -extern "C" { -static const backend_dispatch_t apir_backend_dispatch_table[APIR_BACKEND_DISPATCH_TABLE_COUNT] = { - - /* device */ - - /* APIR_COMMAND_TYPE_DEVICE_GET_DEVICE_COUNT = */ backend_device_get_device_count, - /* APIR_COMMAND_TYPE_DEVICE_GET_COUNT = */ backend_device_get_count, - /* APIR_COMMAND_TYPE_DEVICE_GET_NAME = */ backend_device_get_name, - /* APIR_COMMAND_TYPE_DEVICE_GET_DESCRIPTION = */ backend_device_get_description, - /* APIR_COMMAND_TYPE_DEVICE_GET_TYPE = */ backend_device_get_type, - /* APIR_COMMAND_TYPE_DEVICE_GET_MEMORY = */ backend_device_get_memory, - /* APIR_COMMAND_TYPE_DEVICE_SUPPORTS_OP = */ backend_device_supports_op, - /* APIR_COMMAND_TYPE_DEVICE_GET_BUFFER_TYPE = */ backend_device_get_buffer_type, - /* APIR_COMMAND_TYPE_DEVICE_GET_PROPS = */ backend_device_get_props, - /* APIR_COMMAND_TYPE_DEVICE_BUFFER_FROM_PTR = */ backend_device_buffer_from_ptr, - - /* buffer-type */ - - /* APIR_COMMAND_TYPE_BUFFER_TYPE_GET_NAME = */ backend_buffer_type_get_name, - /* APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALIGNMENT = */ backend_buffer_type_get_alignment, - /* APIR_COMMAND_TYPE_BUFFER_TYPE_GET_MAX_SIZE = */ backend_buffer_type_get_max_size, - /* APIR_COMMAND_TYPE_BUFFER_TYPE_IS_HOST = */ backend_buffer_type_is_host /* DEPRECATED */, - /* APIR_COMMAND_TYPE_BUFFER_TYPE_ALLOC_BUFFER = */ backend_buffer_type_alloc_buffer, - /* APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALLOC_SIZE = */ backend_buffer_type_get_alloc_size, - - /* buffer */ - - /* APIR_COMMAND_TYPE_BUFFER_GET_BASE = */ backend_buffer_get_base, - /* APIR_COMMAND_TYPE_BUFFER_SET_TENSOR = */ backend_buffer_set_tensor, - /* APIR_COMMAND_TYPE_BUFFER_GET_TENSOR = */ backend_buffer_get_tensor, - /* APIR_COMMAND_TYPE_BUFFER_CPY_TENSOR = */ backend_buffer_cpy_tensor, - /* APIR_COMMAND_TYPE_BUFFER_CLEAR = */ backend_buffer_clear, - /* APIR_COMMAND_TYPE_BUFFER_FREE_BUFFER = */ backend_buffer_free_buffer, - - /* backend */ - - /* APIR_COMMAND_TYPE_BACKEND_GRAPH_COMPUTE = */ backend_backend_graph_compute, -}; -} diff --git a/ggml/src/ggml-virtgpu/backend/backend-dispatched.h b/ggml/src/ggml-virtgpu/backend/backend-dispatched.h deleted file mode 100644 index 740ee9e3ffc8..000000000000 --- a/ggml/src/ggml-virtgpu/backend/backend-dispatched.h +++ /dev/null @@ -1,27 +0,0 @@ -#pragma once - -// clang-format off -#include -#include - -#include - -#include "backend-convert.h" -#include "backend-virgl-apir.h" -#include "shared/apir_backend.h" -#include "shared/apir_cs.h" -#include "shared/apir_cs_ggml.h" -// clang-format on - -#define GGML_VIRTGPU_BCK "ggml-virtgpu-backend: " - -struct virgl_apir_context { - uint32_t ctx_id; - virgl_apir_callbacks * iface; -}; - -typedef uint32_t (*backend_dispatch_t)(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); - -#include "backend-dispatched.gen.h" - -uint32_t backend_dispatch_initialize(void * ggml_backend_reg_fct_p); diff --git a/ggml/src/ggml-virtgpu/backend/backend-virgl-apir.h b/ggml/src/ggml-virtgpu/backend/backend-virgl-apir.h deleted file mode 100644 index c65a01cdf9b4..000000000000 --- a/ggml/src/ggml-virtgpu/backend/backend-virgl-apir.h +++ /dev/null @@ -1,32 +0,0 @@ -#pragma once - -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "ggml-impl.h" -#include "shared/api_remoting.h" - -#include -#include -#include - -extern ggml_backend_reg_t reg; -extern ggml_backend_dev_t dev; -extern ggml_backend_t bck; - -struct virgl_apir_callbacks { - const char * (*get_config)(uint32_t virgl_ctx_id, const char * key); - void * (*get_shmem_ptr)(uint32_t virgl_ctx_id, uint32_t res_id); -}; - -extern "C" { -ApirLoadLibraryReturnCode apir_backend_initialize(uint32_t virgl_ctx_id, struct virgl_apir_callbacks * virgl_cbs); -void apir_backend_deinit(uint32_t virgl_ctx_id); -uint32_t apir_backend_dispatcher(uint32_t virgl_ctx_id, - virgl_apir_callbacks * virgl_cbs, - uint32_t cmd_type, - char * dec_cur, - const char * dec_end, - char * enc_cur, - const char * enc_end, - char ** enc_cur_after); -} diff --git a/ggml/src/ggml-virtgpu/backend/backend.cpp b/ggml/src/ggml-virtgpu/backend/backend.cpp deleted file mode 100644 index 535a05f3e69b..000000000000 --- a/ggml/src/ggml-virtgpu/backend/backend.cpp +++ /dev/null @@ -1,144 +0,0 @@ -#include "backend-dispatched.h" -#include "backend-virgl-apir.h" -#include "shared/api_remoting.h" -#include "shared/apir_backend.h" -#include "shared/apir_cs.h" - -#include -#include - -#include - -#define APIR_LLAMA_CPP_GGML_LIBRARY_PATH_ENV "APIR_LLAMA_CPP_GGML_LIBRARY_PATH" -#define APIR_LLAMA_CPP_GGML_LIBRARY_REG_ENV "APIR_LLAMA_CPP_GGML_LIBRARY_REG" -#define APIR_LLAMA_CPP_LOG_TO_FILE_ENV "APIR_LLAMA_CPP_LOG_TO_FILE" - -#define GGML_DEFAULT_BACKEND_REG "ggml_backend_init" - -static void * backend_library_handle = NULL; -static FILE * apir_logfile = NULL; - -static void log_to_file_callback(enum ggml_log_level level, const char * text, void * user_data) { - FILE * logfile = (FILE *) user_data; - fprintf(logfile, "[%d] %s", level, text); - fflush(logfile); -} - -extern "C" { -void apir_backend_deinit(uint32_t virgl_ctx_id) { - GGML_UNUSED(virgl_ctx_id); - - auto buffers = apir_get_track_backend_buffers(); - for (const auto & buffer : buffers) { - apir_untrack_backend_buffer(buffer); - buffer->iface.free_buffer(buffer); - } - - if (backend_library_handle) { - GGML_LOG_INFO(GGML_VIRTGPU_BCK "The GGML backend library was loaded. Unloading it.\n"); - dlclose(backend_library_handle); - backend_library_handle = NULL; - } - - if (apir_logfile) { - fclose(apir_logfile); - apir_logfile = NULL; - } -} - -#define APIR_GGML_LIBRARY_PATH_KEY "ggml.library.path" -#define APIR_GGML_LIBRARY_REG_KEY "ggml.library.reg" - -ApirLoadLibraryReturnCode apir_backend_initialize(uint32_t virgl_ctx_id, struct virgl_apir_callbacks * virgl_cbs) { - const char * dlsym_error; - - const char * apir_log_to_file = getenv(APIR_LLAMA_CPP_LOG_TO_FILE_ENV); - if (apir_log_to_file) { - apir_logfile = fopen(apir_log_to_file, "w"); - if (apir_logfile) { - ggml_log_set(log_to_file_callback, apir_logfile); - } else { - GGML_LOG_INFO(GGML_VIRTGPU_BCK "Could not open the log file at '%s'\n", apir_log_to_file); - } - } - - const char * library_name = virgl_cbs->get_config(virgl_ctx_id, APIR_GGML_LIBRARY_PATH_KEY); - const char * virgl_library_reg = virgl_cbs->get_config(virgl_ctx_id, APIR_GGML_LIBRARY_REG_KEY); - const char * library_reg = virgl_library_reg ? virgl_library_reg : GGML_DEFAULT_BACKEND_REG; - - if (!library_name) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: cannot open the GGML library: env var '%s' not defined\n", __func__, - APIR_LLAMA_CPP_GGML_LIBRARY_PATH_ENV); - - return APIR_LOAD_LIBRARY_ENV_VAR_MISSING; - } - - backend_library_handle = dlopen(library_name, RTLD_LAZY); - - if (!backend_library_handle) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: cannot open the GGML library: %s\n", __func__, dlerror()); - - return APIR_LOAD_LIBRARY_CANNOT_OPEN; - } - - if (!library_reg) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: cannot register the GGML library: env var '%s' not defined\n", __func__, - APIR_LLAMA_CPP_GGML_LIBRARY_REG_ENV); - - return APIR_LOAD_LIBRARY_ENV_VAR_MISSING; - } - - void * ggml_backend_reg_fct = dlsym(backend_library_handle, library_reg); - dlsym_error = dlerror(); - if (dlsym_error) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: cannot find the GGML backend registration symbol '%s' (from %s): %s\n", - __func__, library_reg, APIR_LLAMA_CPP_GGML_LIBRARY_REG_ENV, dlsym_error); - - return APIR_LOAD_LIBRARY_SYMBOL_MISSING; - } - - uint32_t ret = backend_dispatch_initialize(ggml_backend_reg_fct); - - return (ApirLoadLibraryReturnCode) (APIR_LOAD_LIBRARY_INIT_BASE_INDEX + ret); -} - -uint32_t apir_backend_dispatcher(uint32_t virgl_ctx_id, - virgl_apir_callbacks * virgl_cbs, - uint32_t cmd_type, - char * dec_cur, - const char * dec_end, - char * enc_cur, - const char * enc_end, - char ** enc_cur_after) { - apir_encoder enc = { - .cur = enc_cur, - .start = enc_cur, - .end = enc_end, - .fatal = false, - }; - - apir_decoder dec = { - .cur = dec_cur, - .end = dec_end, - .fatal = false, - }; - - virgl_apir_context ctx = { - .ctx_id = virgl_ctx_id, - .iface = virgl_cbs, - }; - - if (cmd_type >= APIR_BACKEND_DISPATCH_TABLE_COUNT) { - GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Received an invalid dispatch index (%d >= %d)\n", __func__, cmd_type, - APIR_BACKEND_DISPATCH_TABLE_COUNT); - return APIR_BACKEND_FORWARD_INDEX_INVALID; - } - - backend_dispatch_t forward_fct = apir_backend_dispatch_table[cmd_type]; - uint32_t ret = forward_fct(&enc, &dec, &ctx); - - *enc_cur_after = enc.cur; - - return ret; -} -} diff --git a/ggml/src/ggml-virtgpu/backend/backend.cpp.inc b/ggml/src/ggml-virtgpu/backend/backend.cpp.inc new file mode 100644 index 000000000000..40f306248146 --- /dev/null +++ b/ggml/src/ggml-virtgpu/backend/backend.cpp.inc @@ -0,0 +1,1831 @@ +// clang-format off +#include +#include + +#include + +typedef enum ApirBackendCommandType { + + /* device */ + APIR_COMMAND_TYPE_DEVICE_GET_DEVICE_COUNT = 0, + APIR_COMMAND_TYPE_DEVICE_GET_COUNT = 1, + APIR_COMMAND_TYPE_DEVICE_GET_NAME = 2, + APIR_COMMAND_TYPE_DEVICE_GET_DESCRIPTION = 3, + APIR_COMMAND_TYPE_DEVICE_GET_TYPE = 4, + APIR_COMMAND_TYPE_DEVICE_GET_MEMORY = 5, + APIR_COMMAND_TYPE_DEVICE_SUPPORTS_OP = 6, + APIR_COMMAND_TYPE_DEVICE_GET_BUFFER_TYPE = 7, + APIR_COMMAND_TYPE_DEVICE_GET_PROPS = 8, + APIR_COMMAND_TYPE_DEVICE_BUFFER_FROM_PTR = 9, + + /* buffer-type */ + APIR_COMMAND_TYPE_BUFFER_TYPE_GET_NAME = 10, + APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALIGNMENT = 11, + APIR_COMMAND_TYPE_BUFFER_TYPE_GET_MAX_SIZE = 12, + APIR_COMMAND_TYPE_BUFFER_TYPE_IS_HOST = 13, + APIR_COMMAND_TYPE_BUFFER_TYPE_ALLOC_BUFFER = 14, + APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALLOC_SIZE = 15, + + /* buffer */ + APIR_COMMAND_TYPE_BUFFER_GET_BASE = 16, + APIR_COMMAND_TYPE_BUFFER_SET_TENSOR = 17, + APIR_COMMAND_TYPE_BUFFER_GET_TENSOR = 18, + APIR_COMMAND_TYPE_BUFFER_CPY_TENSOR = 19, + APIR_COMMAND_TYPE_BUFFER_CLEAR = 20, + APIR_COMMAND_TYPE_BUFFER_FREE_BUFFER = 21, + + /* backend */ + APIR_COMMAND_TYPE_BACKEND_GRAPH_COMPUTE = 22, + + // last command_type index + 1 + APIR_BACKEND_DISPATCH_TABLE_COUNT = 23, +} ApirBackendCommandType; + +static inline const char * apir_dispatch_command_name(ApirBackendCommandType type) { + switch (type) { + /* device */ + case APIR_COMMAND_TYPE_DEVICE_GET_DEVICE_COUNT: + return "device_get_device_count"; + case APIR_COMMAND_TYPE_DEVICE_GET_COUNT: + return "device_get_count"; + case APIR_COMMAND_TYPE_DEVICE_GET_NAME: + return "device_get_name"; + case APIR_COMMAND_TYPE_DEVICE_GET_DESCRIPTION: + return "device_get_description"; + case APIR_COMMAND_TYPE_DEVICE_GET_TYPE: + return "device_get_type"; + case APIR_COMMAND_TYPE_DEVICE_GET_MEMORY: + return "device_get_memory"; + case APIR_COMMAND_TYPE_DEVICE_SUPPORTS_OP: + return "device_supports_op"; + case APIR_COMMAND_TYPE_DEVICE_GET_BUFFER_TYPE: + return "device_get_buffer_type"; + case APIR_COMMAND_TYPE_DEVICE_GET_PROPS: + return "device_get_props"; + case APIR_COMMAND_TYPE_DEVICE_BUFFER_FROM_PTR: + return "device_buffer_from_ptr"; + /* buffer-type */ + case APIR_COMMAND_TYPE_BUFFER_TYPE_GET_NAME: + return "buffer_type_get_name"; + case APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALIGNMENT: + return "buffer_type_get_alignment"; + case APIR_COMMAND_TYPE_BUFFER_TYPE_GET_MAX_SIZE: + return "buffer_type_get_max_size"; + case APIR_COMMAND_TYPE_BUFFER_TYPE_IS_HOST: + return "buffer_type_is_host"; + case APIR_COMMAND_TYPE_BUFFER_TYPE_ALLOC_BUFFER: + return "buffer_type_alloc_buffer"; + case APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALLOC_SIZE: + return "buffer_type_get_alloc_size"; + /* buffer */ + case APIR_COMMAND_TYPE_BUFFER_GET_BASE: + return "buffer_get_base"; + case APIR_COMMAND_TYPE_BUFFER_SET_TENSOR: + return "buffer_set_tensor"; + case APIR_COMMAND_TYPE_BUFFER_GET_TENSOR: + return "buffer_get_tensor"; + case APIR_COMMAND_TYPE_BUFFER_CPY_TENSOR: + return "buffer_cpy_tensor"; + case APIR_COMMAND_TYPE_BUFFER_CLEAR: + return "buffer_clear"; + case APIR_COMMAND_TYPE_BUFFER_FREE_BUFFER: + return "buffer_free_buffer"; + /* backend */ + case APIR_COMMAND_TYPE_BACKEND_GRAPH_COMPUTE: + return "backend_graph_compute"; + + default: + return "unknown"; + } +} + + +#include // for uintptr_t +#include // for timespec, clock_gettime + +#define APIR_BACKEND_INITIALIZE_SUCCESS 0 +#define APIR_BACKEND_INITIALIZE_CANNOT_OPEN_BACKEND_LIBRARY 1 +#define APIR_BACKEND_INITIALIZE_CANNOT_OPEN_GGML_LIBRARY 2 +#define APIR_BACKEND_INITIALIZE_MISSING_BACKEND_SYMBOLS 3 +#define APIR_BACKEND_INITIALIZE_MISSING_GGML_SYMBOLS 4 +#define APIR_BACKEND_INITIALIZE_BACKEND_FAILED 5 +#define APIR_BACKEND_INITIALIZE_BACKEND_REG_FAILED 6 +#define APIR_BACKEND_INITIALIZE_ALREADY_INITED 7 +#define APIR_BACKEND_INITIALIZE_NO_DEVICE 8 +#define APIR_BACKEND_INITIALIZE_BACKEND_INIT_FAILED 9 + +// new entries here need to be added to the apir_backend_initialize_error function below + +#define APIR_BACKEND_FORWARD_INDEX_INVALID 6 + +// 0 is fast, 1 avoids the backend to crash if an unsupported tensor is received +#define APIR_BACKEND_CHECK_SUPPORTS_OP 0 + +typedef uintptr_t apir_buffer_type_host_handle_t; +typedef uintptr_t apir_buffer_host_handle_t; + +static const char * apir_backend_initialize_error(int code) { +#define APIR_BACKEND_INITIALIZE_ERROR(code_name) \ + do { \ + if (code == code_name) \ + return #code_name; \ + } while (0) + + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_SUCCESS); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_CANNOT_OPEN_BACKEND_LIBRARY); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_CANNOT_OPEN_GGML_LIBRARY); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_MISSING_BACKEND_SYMBOLS); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_MISSING_GGML_SYMBOLS); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_BACKEND_FAILED); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_BACKEND_REG_FAILED); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_ALREADY_INITED); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_NO_DEVICE); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_BACKEND_INIT_FAILED); + + return "Unknown APIR_BACKEND_INITIALIZE error:/"; + +#undef APIR_BACKEND_INITIALIZE_ERROR +} + + +#define BUFFER_TO_HOST_HANDLE(name) ggml_buffer_to_apir_handle(name) + +static inline apir_buffer_host_handle_t ggml_buffer_to_apir_handle(ggml_backend_buffer_t buffer) { + // in the backend, the buffer handle is the buffer pointer + return (apir_buffer_host_handle_t) buffer; +} + +static inline apir_buffer_type_host_handle_t ggml_buffer_type_to_apir_handle(ggml_backend_buffer_type_t buft) { + // in the backend, the buffer handle is the buffer pointer + return (apir_buffer_type_host_handle_t) buft; +} + +#include "ggml-backend-impl-defs.inc" +#include "ggml-backend.h.inc" +#include "ggml-impl-defs.inc" +/* the rest of this file must match virglrenderer/src/apir-protocol.h */ + +#include + +#include + +#define APIR_PROTOCOL_MAJOR 0 +#define APIR_PROTOCOL_MINOR 1 + +#define APIR_HANDSHAKE_MAGIC 0xab1e + +enum ApirCommandType { + APIR_COMMAND_TYPE_HANDSHAKE = 0, + APIR_COMMAND_TYPE_LOADLIBRARY = 1, + APIR_COMMAND_TYPE_FORWARD = 2, + + APIR_COMMAND_TYPE_LENGTH = 3, +}; + +typedef uint64_t ApirCommandFlags; + +enum ApirLoadLibraryReturnCode { + APIR_LOAD_LIBRARY_SUCCESS = 0, + // these error codes are returned by the Virglrenderer APIR component + APIR_LOAD_LIBRARY_HYPERCALL_INITIALIZATION_ERROR = 1, + APIR_LOAD_LIBRARY_ALREADY_LOADED = 2, + APIR_LOAD_LIBRARY_ENV_VAR_MISSING = 3, + APIR_LOAD_LIBRARY_CANNOT_OPEN = 4, + APIR_LOAD_LIBRARY_SYMBOL_MISSING = 5, + // any value greater than this is an APIR *backend library* initialization return code + APIR_LOAD_LIBRARY_INIT_BASE_INDEX = 6, +}; + +enum ApirForwardReturnCode { + APIR_FORWARD_SUCCESS = 0, + // these error codes are returned by the Virglrenderer APIR component + APIR_FORWARD_NO_DISPATCH_FCT = 1, + APIR_FORWARD_TIMEOUT = 2, + APIR_FORWARD_FAILED_TO_SYNC_STREAMS = 3, + // any value greater than this index an APIR *backend library* forward return code + APIR_FORWARD_BASE_INDEX = 4, +}; + +__attribute__((unused)) static inline const char * apir_command_name(ApirCommandType type) { + switch (type) { + case APIR_COMMAND_TYPE_HANDSHAKE: + return "HandShake"; + case APIR_COMMAND_TYPE_LOADLIBRARY: + return "LoadLibrary"; + case APIR_COMMAND_TYPE_FORWARD: + return "Forward"; + default: + return "unknown"; + } +} + +__attribute__((unused)) static const char * apir_load_library_error(ApirLoadLibraryReturnCode code) { +#define APIR_LOAD_LIBRARY_ERROR(code_name) \ + do { \ + if (code == code_name) \ + return #code_name; \ + } while (0) + + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_SUCCESS); + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_HYPERCALL_INITIALIZATION_ERROR); + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_ALREADY_LOADED); + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_ENV_VAR_MISSING); + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_CANNOT_OPEN); + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_SYMBOL_MISSING); + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_INIT_BASE_INDEX); + + return "Unknown APIR_COMMAND_TYPE_LoadLibrary error"; + +#undef APIR_LOAD_LIBRARY_ERROR +} + +__attribute__((unused)) static const char * apir_forward_error(ApirForwardReturnCode code) { +#define APIR_FORWARD_ERROR(code_name) \ + do { \ + if (code == code_name) \ + return #code_name; \ + } while (0) + + APIR_FORWARD_ERROR(APIR_FORWARD_SUCCESS); + APIR_FORWARD_ERROR(APIR_FORWARD_NO_DISPATCH_FCT); + APIR_FORWARD_ERROR(APIR_FORWARD_TIMEOUT); + APIR_FORWARD_ERROR(APIR_FORWARD_FAILED_TO_SYNC_STREAMS); + APIR_FORWARD_ERROR(APIR_FORWARD_BASE_INDEX); + + return "Unknown APIR_COMMAND_TYPE_FORWARD error"; + +#undef APIR_FORWARD_ERROR +} + + +#include +#include +#include + +extern ggml_backend_reg_t reg; +extern ggml_backend_dev_t dev; +extern ggml_backend_t bck; + +struct virgl_apir_callbacks { + const char * (*get_config)(uint32_t virgl_ctx_id, const char * key); + void * (*get_shmem_ptr)(uint32_t virgl_ctx_id, uint32_t res_id); +}; + +extern "C" { +ApirLoadLibraryReturnCode apir_backend_initialize(uint32_t virgl_ctx_id, struct virgl_apir_callbacks * virgl_cbs); +void apir_backend_deinit(uint32_t virgl_ctx_id); +uint32_t apir_backend_dispatcher(uint32_t virgl_ctx_id, + virgl_apir_callbacks * virgl_cbs, + uint32_t cmd_type, + char * dec_cur, + const char * dec_end, + char * enc_cur, + const char * enc_end, + char ** enc_cur_after); +} + +#include "ggml-impl-defs.inc" + +#include +#include + +#define likely(x) __builtin_expect(!!(x), 1) +#define unlikely(x) __builtin_expect(!!(x), 0) + +struct apir_encoder { + char * cur; + const char * start; + const char * end; + bool fatal; +}; + +struct apir_decoder { + const char * cur; + const char * end; + bool fatal; +}; + +/* + * new encoder and decoder + */ + +static apir_decoder apir_new_decoder(const char * ptr, size_t size) { + apir_decoder dec = { + .cur = ptr, + .end = ptr + size, + .fatal = false, + }; + + return dec; +} + +static apir_encoder apir_new_encoder(char * ptr, size_t size) { + apir_encoder enc = { + .cur = ptr, + .start = ptr, + .end = ptr + size, + .fatal = false, + }; + + return enc; +} + +/* + * fatal flag handling + */ + +static inline void apir_encoder_reset_fatal(apir_encoder * enc) { + enc->fatal = false; +} + +static inline void apir_encoder_set_fatal(apir_encoder * enc) { + enc->fatal = true; +} + +static inline bool apir_encoder_get_fatal(const apir_encoder * enc) { + return enc->fatal; +} + +static inline void apir_decoder_reset_fatal(apir_decoder * dec) { + dec->fatal = false; +} + +static inline void apir_decoder_set_fatal(apir_decoder * dec) { + dec->fatal = true; +} + +static inline bool apir_decoder_get_fatal(const apir_decoder * dec) { + return dec->fatal; +} + +/* + * encode peek + */ + +static inline bool apir_decoder_peek_internal(apir_decoder * dec, size_t size, void * val, size_t val_size) { + assert(val_size <= size); + + if (unlikely(size > (size_t) (dec->end - dec->cur))) { + GGML_LOG_ERROR("%s: reading too much from the decoder ...\n", __func__); + apir_decoder_set_fatal(dec); + memset(val, 0, val_size); + return false; + } + + /* we should not rely on the compiler to optimize away memcpy... */ + memcpy(val, dec->cur, val_size); + return true; +} + +static inline void apir_decoder_peek(apir_decoder * dec, size_t size, void * val, size_t val_size) { + apir_decoder_peek_internal(dec, size, val, val_size); +} + +static inline const void * apir_decoder_use_inplace(apir_decoder * dec, size_t size) { + if (unlikely(size > (size_t) (dec->end - dec->cur))) { + GGML_LOG_ERROR("%s: reading too much from the decoder ...\n", __func__); + apir_decoder_set_fatal(dec); + return NULL; + } + const void * addr = dec->cur; + dec->cur += size; + + return addr; +} + +/* + * read/write + */ + +static inline void apir_decoder_read(apir_decoder * dec, size_t size, void * val, size_t val_size) { + if (apir_decoder_peek_internal(dec, size, val, val_size)) { + dec->cur += size; + } +} + +static inline char * apir_encoder_write(apir_encoder * enc, size_t size, const void * val, size_t val_size) { + assert(val_size <= size); + assert(size <= ((size_t) (enc->end - enc->cur))); + + char * write_addr = enc->cur; + /* we should not rely on the compiler to optimize away memcpy... */ + memcpy(write_addr, val, val_size); + enc->cur += size; + + return write_addr; +} + +/* + * encode/decode + */ + +static inline void apir_decode(apir_decoder * dec, size_t size, void * data, size_t data_size) { + assert(size % 4 == 0); + apir_decoder_read(dec, size, data, data_size); +} + +static inline void apir_encode(apir_encoder * enc, size_t size, const void * data, size_t data_size) { + assert(size % 4 == 0); + apir_encoder_write(enc, size, data, data_size); +} + +/* + * typed encode/decode + */ + +/* uint8_t */ + +static inline void apir_encode_uint8_t(apir_encoder * enc, const uint8_t * val) { + apir_encode(enc, sizeof(int), val, sizeof(*val)); +} + +static inline void apir_decode_uint8_t(apir_decoder * dec, uint8_t * val) { + apir_decode(dec, sizeof(int), val, sizeof(*val)); +} + +/* uint64_t */ + +static inline void apir_encode_uint64_t(apir_encoder * enc, const uint64_t * val) { + apir_encode(enc, 8, val, sizeof(*val)); +} + +static inline void apir_decode_uint64_t(apir_decoder * dec, uint64_t * val) { + apir_decode(dec, 8, val, sizeof(*val)); +} + +static inline void apir_encode_uint64_t_array(apir_encoder * enc, const uint64_t * val, uint32_t count) { + const size_t size = sizeof(*val) * count; + assert(size >= count); + apir_encode(enc, size, val, size); +} + +static inline void apir_decode_uint64_t_array(apir_decoder * dec, uint64_t * val, uint32_t count) { + const size_t size = sizeof(*val) * count; + assert(size >= count); + apir_decode(dec, size, val, size); +} + +static inline const uint64_t * apir_decode_uint64_t_array_inplace(apir_decoder * dec, uint32_t count) { + return (uint64_t *) (uintptr_t) apir_decoder_use_inplace(dec, count * sizeof(uint64_t)); +} + +/* int32_t */ + +static inline void apir_encode_int32_t(apir_encoder * enc, const int32_t * val) { + apir_encode(enc, 4, val, sizeof(*val)); +} + +static inline void apir_decode_int32_t(apir_decoder * dec, int32_t * val) { + apir_decode(dec, 4, val, sizeof(*val)); +} + +static inline void apir_encode_int32_t_array(apir_encoder * enc, const int32_t * val, uint32_t count) { + const size_t size = sizeof(*val) * count; + assert(size >= count); + apir_encode(enc, size, val, size); +} + +static inline void apir_decode_int32_t_array(apir_decoder * dec, int32_t * val, uint32_t count) { + const size_t size = sizeof(*val) * count; + assert(size >= count); + apir_decode(dec, size, val, size); +} + +/* array size (uint64_t) */ + +static inline void apir_encode_array_size(apir_encoder * enc, uint64_t size) { + apir_encode_uint64_t(enc, &size); +} + +static inline uint64_t apir_decode_array_size(apir_decoder * dec, uint64_t expected_size) { + uint64_t size; + apir_decode_uint64_t(dec, &size); + if (size != expected_size) { + GGML_LOG_ERROR("%s: Couldn't decode array from the decoder\n", __func__); + apir_decoder_set_fatal(dec); + size = 0; + } + return size; +} + +static inline uint64_t apir_decode_array_size_unchecked(apir_decoder * dec) { + uint64_t size; + apir_decode_uint64_t(dec, &size); + return size; +} + +/* non-array pointer */ + +static inline bool apir_encode_simple_pointer(apir_encoder * enc, const void * val) { + apir_encode_array_size(enc, val ? 1 : 0); + return val; +} + +static inline bool apir_decode_simple_pointer(apir_decoder * dec) { + return apir_decode_array_size_unchecked(dec); +} + +/* uint32_t */ + +static inline void apir_encode_uint32_t(apir_encoder * enc, const uint32_t * val) { + apir_encode(enc, 4, val, sizeof(*val)); +} + +static inline void apir_decode_uint32_t(apir_decoder * dec, uint32_t * val) { + apir_decode(dec, 4, val, sizeof(*val)); +} + +static inline void apir_encode_uint32_t_array(apir_encoder * enc, const uint32_t * val, uint32_t count) { + const size_t size = sizeof(*val) * count; + assert(size >= count); + apir_encode(enc, size, val, size); +} + +static inline void apir_decode_uint32_t_array(apir_decoder * dec, uint32_t * val, uint32_t count) { + const size_t size = sizeof(*val) * count; + assert(size >= count); + apir_decode(dec, size, val, size); +} + +/* size_t */ + +static inline void apir_encode_size_t(apir_encoder * enc, const size_t * val) { + const uint64_t tmp = *val; + apir_encode_uint64_t(enc, &tmp); +} + +static inline void apir_decode_size_t(apir_decoder * dec, size_t * val) { + uint64_t tmp; + apir_decode_uint64_t(dec, &tmp); + *val = tmp; +} + +static inline void apir_encode_size_t_array(apir_encoder * enc, const size_t * val, uint32_t count) { + if (sizeof(size_t) == sizeof(uint64_t)) { + apir_encode_uint64_t_array(enc, (const uint64_t *) val, count); + } else { + for (uint32_t i = 0; i < count; i++) { + apir_encode_size_t(enc, &val[i]); + } + } +} + +static inline void apir_decode_size_t_array(apir_decoder * dec, size_t * val, uint32_t count) { + if (sizeof(size_t) == sizeof(uint64_t)) { + apir_decode_uint64_t_array(dec, (uint64_t *) val, count); + } else { + for (uint32_t i = 0; i < count; i++) { + apir_decode_size_t(dec, &val[i]); + } + } +} + +/* opaque blob */ + +static inline void apir_encode_blob_array(apir_encoder * enc, const void * val, size_t size) { + apir_encode(enc, (size + 3) & ~3, val, size); +} + +static inline void apir_decode_blob_array(apir_decoder * dec, void * val, size_t size) { + apir_decode(dec, (size + 3) & ~3, val, size); +} + +/* string */ + +static inline void apir_encode_char_array(apir_encoder * enc, const char * val, size_t size) { + assert(size && strlen(val) < size); + apir_encode_blob_array(enc, val, size); +} + +static inline void apir_decode_char_array(apir_decoder * dec, char * val, size_t size) { + apir_decode_blob_array(dec, val, size); + if (size) { + val[size - 1] = '\0'; + } else { + GGML_LOG_ERROR("%s: Couldn't decode the blog array\n", __func__); + apir_decoder_set_fatal(dec); + } +} + +/* (temp) buffer allocation */ + +static inline void * apir_decoder_alloc_array(size_t size, size_t count) { + size_t alloc_size; + if (unlikely(__builtin_mul_overflow(size, count, &alloc_size))) { + GGML_LOG_ERROR("%s: overflow in array allocation of %zu * %zu bytes\n", __func__, size, count); + return NULL; + } + + return malloc(alloc_size); +} + +/* bool */ + +static inline void apir_encode_bool_t(apir_encoder * enc, const bool * val) { + apir_encode(enc, sizeof(int), val, sizeof(bool)); +} + +static inline void apir_decode_bool_t(apir_decoder * dec, bool * val) { + apir_decode(dec, sizeof(int), val, sizeof(bool)); +} + +/* apir_buffer_type_host_handle_t */ + +static inline void apir_encode_apir_buffer_type_host_handle_t(apir_encoder * enc, + const apir_buffer_type_host_handle_t * val) { + apir_encode(enc, sizeof(apir_buffer_type_host_handle_t), val, sizeof(apir_buffer_type_host_handle_t)); +} + +static inline void apir_decode_apir_buffer_type_host_handle_t(apir_decoder * dec, + apir_buffer_type_host_handle_t * val) { + apir_decode(dec, sizeof(apir_buffer_type_host_handle_t), val, sizeof(apir_buffer_type_host_handle_t)); +} + +/* apir_buffer_host_handle_t */ + +static inline void apir_encode_apir_buffer_host_handle_t(apir_encoder * enc, const apir_buffer_host_handle_t * val) { + apir_encode(enc, sizeof(apir_buffer_host_handle_t), val, sizeof(apir_buffer_host_handle_t)); +} + +static inline void apir_decode_apir_buffer_host_handle_t(apir_decoder * dec, apir_buffer_host_handle_t * val) { + apir_decode(dec, sizeof(apir_buffer_host_handle_t), val, sizeof(apir_buffer_host_handle_t)); +} + +/* uintptr_t */ + +static inline void apir_encode_uintptr_t(apir_encoder * enc, const uintptr_t * val) { + apir_encode(enc, sizeof(*val), val, sizeof(*val)); +} + +static inline void apir_decode_uintptr_t(apir_decoder * dec, uintptr_t * val) { + apir_decode(dec, sizeof(*val), val, sizeof(*val)); +} + +// clang-format off +#include "ggml.h.inc" +#include "ggml-backend-impl-defs.inc" + +#include +#include +#include +#include +// clang-format on + +// ggml_tensor is serialized into apir_rpc_tensor +struct apir_rpc_tensor { + uint64_t id; + uint32_t type; + uint64_t buffer; + uint32_t ne[GGML_MAX_DIMS]; + uint32_t nb[GGML_MAX_DIMS]; + uint32_t op; + int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; + int32_t flags; + uint64_t src[GGML_MAX_SRC]; + uint64_t view_src; + uint64_t view_offs; + uint64_t data; + char name[GGML_MAX_NAME]; + + char padding[4]; +}; + +/* frontend */ + +apir_rpc_tensor apir_serialize_tensor(const ggml_tensor * tensor); + +void apir_serialize_graph(const ggml_cgraph * cgraph, std::vector & output); + +/* backend */ + +void apir_track_backend_buffer(ggml_backend_buffer_t buffer); +bool apir_untrack_backend_buffer(ggml_backend_buffer_t buffer); +std::unordered_set apir_get_track_backend_buffers(); + +void apir_add_tensor(ggml_tensor * tensor, + std::vector & tensors, + std::unordered_set & visited); + +ggml_tensor * apir_deserialize_tensor(ggml_context * ctx, const apir_rpc_tensor * tensor); + +ggml_tensor * apir_create_node(uint64_t id, + ggml_context * ctx, + const std::unordered_map & tensor_ptrs, + std::unordered_map & tensor_map); + +ggml_cgraph * apir_deserialize_graph(uint32_t n_nodes, + uint32_t n_tensors, + const apir_rpc_tensor * tensors, + const uint64_t * nodes); + +#include "ggml-impl-defs.inc" + +// ggml_buffer_to_apir_host_handle(ggml_backend_buffer_t buffer); + +static inline void apir_encode_ggml_buffer_host_handle(apir_encoder * enc, const apir_buffer_host_handle_t * handle); + +static inline ggml_backend_buffer_t apir_decode_ggml_buffer(apir_decoder * dec); + +/* apir_rpc_tensor */ + +static inline void apir_encode_rcp_tensor(apir_encoder * enc, const apir_rpc_tensor * apir_rpc_tensor) { + size_t apir_rpc_tensor_size = sizeof(*apir_rpc_tensor); + apir_encode(enc, apir_rpc_tensor_size, apir_rpc_tensor, apir_rpc_tensor_size); +} + +static inline apir_rpc_tensor * apir_decode_apir_rpc_tensor_inplace(apir_decoder * dec) { + size_t apir_rpc_tensor_size = sizeof(apir_rpc_tensor); + + return (apir_rpc_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, apir_rpc_tensor_size); +} + +static inline apir_rpc_tensor * apir_decode_apir_rpc_tensor_array_inplace(apir_decoder * dec, uint32_t n_tensors) { + size_t apir_rpc_tensor_size = sizeof(apir_rpc_tensor) * n_tensors; + + return (apir_rpc_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, apir_rpc_tensor_size); +} + +/* ggml_tensor */ + +static inline void apir_encode_ggml_tensor(apir_encoder * enc, const ggml_tensor * tensor) { + apir_rpc_tensor serialized = apir_serialize_tensor(tensor); + + apir_encode_rcp_tensor(enc, &serialized); +} + +static inline const ggml_tensor * apir_decode_ggml_tensor(apir_decoder * dec) { + const apir_rpc_tensor * apir_rpc_tensor = apir_decode_apir_rpc_tensor_inplace(dec); + + if (!apir_rpc_tensor) { + return NULL; + } + + ggml_init_params params{ + /*.mem_size =*/ggml_tensor_overhead(), + /*.mem_buffer =*/NULL, + /*.no_alloc =*/true, + }; + + ggml_context * ctx = ggml_init(params); + + const ggml_tensor * tensor = apir_deserialize_tensor(ctx, apir_rpc_tensor); + + return tensor; +} + +/* *** ggml_backend_buffer_type_t *** */ + +// ggml_backend_buffer_type_t is a POINTER (to a struct). +// Only the host pointer is shared between the host and guest. +// The guest stores it in `buft->context`. +// The host simply writes the pointer address in the buffer variable. + +static inline void apir_encode_ggml_buffer_type(apir_encoder * enc, ggml_backend_buffer_type_t buft) { + apir_buffer_type_host_handle_t handle = ggml_buffer_type_to_apir_handle(buft); + apir_encoder_write(enc, sizeof(handle), &handle, sizeof(handle)); +} + +static inline ggml_backend_buffer_type_t apir_decode_ggml_buffer_type(apir_decoder * dec) { + apir_buffer_type_host_handle_t handle; + + apir_decoder_read(dec, sizeof(handle), &handle, sizeof(handle)); + + return (ggml_backend_buffer_type_t) handle; +} + +static inline void apir_encode_apir_buffer_type_host_handle(apir_encoder * enc, apir_buffer_type_host_handle_t handle) { + apir_encoder_write(enc, sizeof(handle), &handle, sizeof(handle)); +} + +static inline apir_buffer_type_host_handle_t apir_decode_apir_buffer_type_host_handle(apir_decoder * dec) { + apir_buffer_type_host_handle_t handle; + + apir_decoder_read(dec, sizeof(handle), &handle, sizeof(handle)); + + return handle; +} + +/* *** ggml_backend_type_t *** */ + +// ggml_backend_buffer_t is a POINTER. +// same logic as for ggml_backend_buffer_type_t + +static inline void apir_encode_ggml_buffer(apir_encoder * enc, const ggml_backend_buffer_t buffer) { + apir_buffer_host_handle_t handle = BUFFER_TO_HOST_HANDLE(buffer); + apir_encoder_write(enc, sizeof(handle), &handle, sizeof(handle)); +} + +static inline ggml_backend_buffer_t apir_decode_ggml_buffer(apir_decoder * dec) { + ggml_backend_buffer_t buffer; + size_t buffer_ptr_size = sizeof(buffer); + + apir_decoder_read(dec, buffer_ptr_size, &buffer, buffer_ptr_size); + + // SECURITY: Validate buffer handle against tracked buffers to prevent + // guest VM from providing arbitrary host memory addresses + if (buffer) { + extern std::unordered_set backend_buffers; + if (backend_buffers.find(buffer) == backend_buffers.end()) { + GGML_LOG_WARN("ggml-virtgpu-backend: %s: Invalid buffer handle from guest: %p\n", __func__, + (void *) buffer); + // Set fatal flag to prevent further processing with invalid handle + apir_decoder_set_fatal(dec); + return NULL; + } + } + + return buffer; +} + +/* enum ggml_status */ + +static inline void apir_encode_ggml_status(apir_encoder * enc, const ggml_status * status) { + apir_encoder_write(enc, sizeof(*status), status, sizeof(*status)); +} + +static inline void apir_decode_ggml_status(apir_decoder * dec, ggml_status * status) { + apir_decoder_read(dec, sizeof(*status), status, sizeof(*status)); +} + +/* virtgpu_shmem */ + +static inline void apir_encode_virtgpu_shmem_res_id(apir_encoder * enc, uint32_t shmem_res_id) { + apir_encode_uint32_t(enc, &shmem_res_id); +} + +static inline void apir_decode_virtgpu_shmem_res_id(apir_decoder * dec, uint32_t * shmem_res_id) { + apir_decode_uint32_t(dec, shmem_res_id); +} + +/* ggml_cgraph */ + +static inline size_t apir_serialize_ggml_cgraph(ggml_cgraph * cgraph, std::vector & cgraph_data) { + apir_serialize_graph(cgraph, cgraph_data); + + return cgraph_data.size(); +} + +static inline void apir_encode_cgraph_data(apir_encoder * enc, std::vector & cgraph_data) { + size_t cgraph_size = cgraph_data.size(); + + apir_encode(enc, cgraph_size, cgraph_data.data(), cgraph_size); +} + +static inline ggml_cgraph * apir_decode_ggml_cgraph(apir_decoder * dec, size_t cgraph_size) { + GGML_UNUSED(cgraph_size); + + uint32_t n_nodes; + apir_decode_uint32_t(dec, &n_nodes); + const uint64_t * nodes = apir_decode_uint64_t_array_inplace(dec, n_nodes); + + uint32_t n_tensors; + apir_decode_uint32_t(dec, &n_tensors); + const apir_rpc_tensor * tensors = apir_decode_apir_rpc_tensor_array_inplace(dec, n_tensors); + + return apir_deserialize_graph(n_nodes, n_tensors, tensors, nodes); +} + +static inline void apir_encode_ggml_buffer_handle(apir_encoder * enc, const apir_buffer_host_handle_t * handle) { + apir_encoder_write(enc, sizeof(*handle), &handle, sizeof(*handle)); +} + +static inline void apir_encode_ggml_tensor_inline(apir_encoder * enc, const ggml_tensor * tensor) { + size_t tensor_size = sizeof(*tensor); + + if (tensor->extra) { + GGML_ABORT("%s: Cannot pass tensors with extra", __func__); + } + + if (tensor->src[0] && tensor->buffer) { + static int first = 1; + if (first) { + GGML_LOG_WARN("%s: Cannot pass tensors with src and buffer\n", __func__); + first = 0; + } + } + + apir_encoder_write(enc, tensor_size, tensor, tensor_size); + + // tensor->data is a pointer inside the device buffer. No need to touch it + // tensor->buffer is a pointer to a buffer. Encoding the buffer handle in sequence. + // (could also make a copy of the tensor, and update locally.) + + if (tensor->buffer) { + apir_buffer_host_handle_t buffer_handle = ggml_buffer_to_apir_handle(tensor->buffer); + apir_encode_ggml_buffer_handle(enc, &buffer_handle); + } + + if (tensor->view_src) { + apir_encoder_write(enc, tensor_size, tensor->view_src, tensor_size); + } + + for (int i = 0; tensor->src[i]; i++) { + const ggml_tensor * tensor_src = tensor->src[i]; + apir_encoder_write(enc, tensor_size, tensor_src, tensor_size); + } +} + +static inline const ggml_tensor * apir_decode_ggml_tensor_inplace(apir_decoder * dec) { + // it safe to remove the `const` qualifier here, we *do* want to + // modify the shared memory data to fix the `src` pointers. + ggml_tensor * tensor = (ggml_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, sizeof(ggml_tensor)); + + // tensor->data is a pointer inside the device buffer. No need to touch it + // tensor->buffer is a pointer to a buffer. Decode the buffer handle encoded in sequence. + if (tensor->buffer) { + tensor->buffer = apir_decode_ggml_buffer(dec); + } + + if (tensor->view_src) { + ggml_tensor * tensor_view_src = (ggml_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, sizeof(ggml_tensor)); + tensor->view_src = tensor_view_src; + } + + for (int i = 0; tensor->src[i]; i++) { + ggml_tensor * tensor_src = (ggml_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, sizeof(ggml_tensor)); + tensor->src[i] = tensor_src; // overwrite op->src[i] pointer with the actual location of the src tensor + } + + return tensor; +} + +// clang-format on + +#define GGML_VIRTGPU_BCK "ggml-virtgpu-backend: " + +struct virgl_apir_context { + uint32_t ctx_id; + virgl_apir_callbacks * iface; +}; + +typedef uint32_t (*backend_dispatch_t)(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); + +/* device */ +uint32_t backend_device_get_device_count(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_device_get_count(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_device_get_name(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_device_get_description(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_device_get_type(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_device_get_memory(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_device_supports_op(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_device_get_buffer_type(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_device_get_props(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_device_buffer_from_ptr(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); + +/* buffer-type */ +uint32_t backend_buffer_type_get_name(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_buffer_type_get_alignment(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_buffer_type_get_max_size(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +/* APIR_COMMAND_TYPE_BUFFER_TYPE_IS_HOST is deprecated. Keeping the handler for backward compatibility. */ +uint32_t backend_buffer_type_is_host(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_buffer_type_alloc_buffer(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_buffer_type_get_alloc_size(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); + +/* buffer */ +uint32_t backend_buffer_get_base(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_buffer_set_tensor(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_buffer_get_tensor(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_buffer_cpy_tensor(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_buffer_clear(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); +uint32_t backend_buffer_free_buffer(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); + +/* backend */ +uint32_t backend_backend_graph_compute(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx); + +extern "C" { +static const backend_dispatch_t apir_backend_dispatch_table[APIR_BACKEND_DISPATCH_TABLE_COUNT] = { + + /* device */ + + /* APIR_COMMAND_TYPE_DEVICE_GET_DEVICE_COUNT = */ backend_device_get_device_count, + /* APIR_COMMAND_TYPE_DEVICE_GET_COUNT = */ backend_device_get_count, + /* APIR_COMMAND_TYPE_DEVICE_GET_NAME = */ backend_device_get_name, + /* APIR_COMMAND_TYPE_DEVICE_GET_DESCRIPTION = */ backend_device_get_description, + /* APIR_COMMAND_TYPE_DEVICE_GET_TYPE = */ backend_device_get_type, + /* APIR_COMMAND_TYPE_DEVICE_GET_MEMORY = */ backend_device_get_memory, + /* APIR_COMMAND_TYPE_DEVICE_SUPPORTS_OP = */ backend_device_supports_op, + /* APIR_COMMAND_TYPE_DEVICE_GET_BUFFER_TYPE = */ backend_device_get_buffer_type, + /* APIR_COMMAND_TYPE_DEVICE_GET_PROPS = */ backend_device_get_props, + /* APIR_COMMAND_TYPE_DEVICE_BUFFER_FROM_PTR = */ backend_device_buffer_from_ptr, + + /* buffer-type */ + + /* APIR_COMMAND_TYPE_BUFFER_TYPE_GET_NAME = */ backend_buffer_type_get_name, + /* APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALIGNMENT = */ backend_buffer_type_get_alignment, + /* APIR_COMMAND_TYPE_BUFFER_TYPE_GET_MAX_SIZE = */ backend_buffer_type_get_max_size, + /* APIR_COMMAND_TYPE_BUFFER_TYPE_IS_HOST = */ backend_buffer_type_is_host /* DEPRECATED */, + /* APIR_COMMAND_TYPE_BUFFER_TYPE_ALLOC_BUFFER = */ backend_buffer_type_alloc_buffer, + /* APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALLOC_SIZE = */ backend_buffer_type_get_alloc_size, + + /* buffer */ + + /* APIR_COMMAND_TYPE_BUFFER_GET_BASE = */ backend_buffer_get_base, + /* APIR_COMMAND_TYPE_BUFFER_SET_TENSOR = */ backend_buffer_set_tensor, + /* APIR_COMMAND_TYPE_BUFFER_GET_TENSOR = */ backend_buffer_get_tensor, + /* APIR_COMMAND_TYPE_BUFFER_CPY_TENSOR = */ backend_buffer_cpy_tensor, + /* APIR_COMMAND_TYPE_BUFFER_CLEAR = */ backend_buffer_clear, + /* APIR_COMMAND_TYPE_BUFFER_FREE_BUFFER = */ backend_buffer_free_buffer, + + /* backend */ + + /* APIR_COMMAND_TYPE_BACKEND_GRAPH_COMPUTE = */ backend_backend_graph_compute, +}; +} + + +extern "C" uint32_t backend_dispatch_initialize(void * ggml_backend_reg_fct_p); + + +#include +#include + +#include + +#define APIR_LLAMA_CPP_GGML_LIBRARY_PATH_ENV "APIR_LLAMA_CPP_GGML_LIBRARY_PATH" +#define APIR_LLAMA_CPP_GGML_LIBRARY_REG_ENV "APIR_LLAMA_CPP_GGML_LIBRARY_REG" +#define APIR_LLAMA_CPP_LOG_TO_FILE_ENV "APIR_LLAMA_CPP_LOG_TO_FILE" + +#define GGML_DEFAULT_BACKEND_REG "ggml_backend_init" + +static void * backend_library_handle = NULL; +static FILE * apir_logfile = NULL; + +static void log_to_file_callback(enum ggml_log_level level, const char * text, void * user_data) { + FILE * logfile = (FILE *) user_data; + fprintf(logfile, "[%d] %s", level, text); + fflush(logfile); +} + +extern "C" { +void apir_backend_deinit(uint32_t virgl_ctx_id) { + GGML_UNUSED(virgl_ctx_id); + + auto buffers = apir_get_track_backend_buffers(); + for (const auto & buffer : buffers) { + apir_untrack_backend_buffer(buffer); + buffer->iface.free_buffer(buffer); + } + + if (backend_library_handle) { + GGML_LOG_INFO(GGML_VIRTGPU_BCK "The GGML backend library was loaded. Unloading it.\n"); + dlclose(backend_library_handle); + backend_library_handle = NULL; + } + + if (apir_logfile) { + fclose(apir_logfile); + apir_logfile = NULL; + } +} + +#define APIR_GGML_LIBRARY_PATH_KEY "ggml.library.path" +#define APIR_GGML_LIBRARY_REG_KEY "ggml.library.reg" + +ApirLoadLibraryReturnCode apir_backend_initialize(uint32_t virgl_ctx_id, struct virgl_apir_callbacks * virgl_cbs) { + const char * dlsym_error; + + const char * apir_log_to_file = getenv(APIR_LLAMA_CPP_LOG_TO_FILE_ENV); + if (apir_log_to_file) { + apir_logfile = fopen(apir_log_to_file, "w"); + if (apir_logfile) { + ggml_log_set(log_to_file_callback, apir_logfile); + } else { + GGML_LOG_INFO(GGML_VIRTGPU_BCK "Could not open the log file at '%s'\n", apir_log_to_file); + } + } + + const char * library_name = virgl_cbs->get_config(virgl_ctx_id, APIR_GGML_LIBRARY_PATH_KEY); + const char * virgl_library_reg = virgl_cbs->get_config(virgl_ctx_id, APIR_GGML_LIBRARY_REG_KEY); + const char * library_reg = virgl_library_reg ? virgl_library_reg : GGML_DEFAULT_BACKEND_REG; + + if (!library_name) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: cannot open the GGML library: env var '%s' not defined\n", __func__, + APIR_LLAMA_CPP_GGML_LIBRARY_PATH_ENV); + + return APIR_LOAD_LIBRARY_ENV_VAR_MISSING; + } + + backend_library_handle = dlopen(library_name, RTLD_LAZY); + + if (!backend_library_handle) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: cannot open the GGML library: %s\n", __func__, dlerror()); + + return APIR_LOAD_LIBRARY_CANNOT_OPEN; + } + + if (!library_reg) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: cannot register the GGML library: env var '%s' not defined\n", __func__, + APIR_LLAMA_CPP_GGML_LIBRARY_REG_ENV); + + return APIR_LOAD_LIBRARY_ENV_VAR_MISSING; + } + + void * ggml_backend_reg_fct = dlsym(backend_library_handle, library_reg); + dlsym_error = dlerror(); + if (dlsym_error) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: cannot find the GGML backend registration symbol '%s' (from %s): %s\n", + __func__, library_reg, APIR_LLAMA_CPP_GGML_LIBRARY_REG_ENV, dlsym_error); + + return APIR_LOAD_LIBRARY_SYMBOL_MISSING; + } + + uint32_t ret = backend_dispatch_initialize(ggml_backend_reg_fct); + + return (ApirLoadLibraryReturnCode) (APIR_LOAD_LIBRARY_INIT_BASE_INDEX + ret); +} + +uint32_t apir_backend_dispatcher(uint32_t virgl_ctx_id, + virgl_apir_callbacks * virgl_cbs, + uint32_t cmd_type, + char * dec_cur, + const char * dec_end, + char * enc_cur, + const char * enc_end, + char ** enc_cur_after) { + apir_encoder enc = { + .cur = enc_cur, + .start = enc_cur, + .end = enc_end, + .fatal = false, + }; + + apir_decoder dec = { + .cur = dec_cur, + .end = dec_end, + .fatal = false, + }; + + virgl_apir_context ctx = { + .ctx_id = virgl_ctx_id, + .iface = virgl_cbs, + }; + + if (cmd_type >= APIR_BACKEND_DISPATCH_TABLE_COUNT) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Received an invalid dispatch index (%d >= %d)\n", __func__, cmd_type, + APIR_BACKEND_DISPATCH_TABLE_COUNT); + return APIR_BACKEND_FORWARD_INDEX_INVALID; + } + + backend_dispatch_t forward_fct = apir_backend_dispatch_table[cmd_type]; + uint32_t ret = forward_fct(&enc, &dec, &ctx); + + *enc_cur_after = enc.cur; + + return ret; +} +} + + +/* Inlined backend-dispatched-device.cpp */ +#include "ggml-backend-impl-defs.inc" +#include "ggml-backend.h.inc" +#include "ggml-impl-defs.inc" + +#include + +uint32_t backend_device_get_device_count(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(ctx); + GGML_UNUSED(dec); + + int32_t dev_count = reg->iface.get_device_count(reg); + apir_encode_int32_t(enc, &dev_count); + + return 0; +} + +uint32_t backend_device_get_count(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(ctx); + GGML_UNUSED(dec); + + int32_t dev_count = reg->iface.get_device_count(reg); + apir_encode_int32_t(enc, &dev_count); + + return 0; +} + +uint32_t backend_device_get_name(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(dec); + + const char * string = dev->iface.get_name(dev); + + const size_t string_size = strlen(string) + 1; + apir_encode_array_size(enc, string_size); + apir_encode_char_array(enc, string, string_size); + + return 0; +} + +uint32_t backend_device_get_description(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(dec); + + const char * string = dev->iface.get_description(dev); + + const size_t string_size = strlen(string) + 1; + apir_encode_array_size(enc, string_size); + apir_encode_char_array(enc, string, string_size); + + return 0; +} + +uint32_t backend_device_get_type(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(dec); + + uint32_t type = dev->iface.get_type(dev); + apir_encode_uint32_t(enc, &type); + + return 0; +} + +uint32_t backend_device_get_memory(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(dec); + + size_t free, total; + dev->iface.get_memory(dev, &free, &total); + + apir_encode_size_t(enc, &free); + apir_encode_size_t(enc, &total); + + return 0; +} + +uint32_t backend_device_supports_op(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + + const ggml_tensor * op = apir_decode_ggml_tensor_inplace(dec); + + bool supports_op = dev->iface.supports_op(dev, op); + + apir_encode_bool_t(enc, &supports_op); + + return 0; +} + +uint32_t backend_device_get_buffer_type(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(dec); + + ggml_backend_buffer_type_t bufft = dev->iface.get_buffer_type(dev); + + apir_encode_ggml_buffer_type(enc, bufft); + + return 0; +} + +uint32_t backend_device_get_props(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(dec); + + ggml_backend_dev_props props; + dev->iface.get_props(dev, &props); + + apir_encode_bool_t(enc, &props.caps.async); + apir_encode_bool_t(enc, &props.caps.host_buffer); + apir_encode_bool_t(enc, &props.caps.buffer_from_host_ptr); + apir_encode_bool_t(enc, &props.caps.events); + + return 0; +} + +uint32_t backend_device_buffer_from_ptr(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(dec); + + uint32_t shmem_res_id; + apir_decode_virtgpu_shmem_res_id(dec, &shmem_res_id); + + void * shmem_ptr = ctx->iface->get_shmem_ptr(ctx->ctx_id, shmem_res_id); + if (!shmem_ptr) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Couldn't get the shmem addr from virgl\n", __func__); + apir_decoder_set_fatal(dec); + return 1; + } + + size_t size; + apir_decode_size_t(dec, &size); + size_t max_tensor_size; + apir_decode_size_t(dec, &max_tensor_size); + + ggml_backend_buffer_t buffer; + buffer = dev->iface.buffer_from_host_ptr(dev, shmem_ptr, size, max_tensor_size); + + apir_encode_ggml_buffer(enc, buffer); + apir_encode_ggml_buffer_type(enc, buffer->buft); + + if (buffer) { + apir_track_backend_buffer(buffer); + } + + return 0; +} + + +/* Inlined backend-dispatched-buffer-type.cpp */ +#include "ggml-backend-impl-defs.inc" +#include "ggml-backend.h.inc" +#include "ggml-impl-defs.inc" + +#include + +uint32_t backend_buffer_type_get_name(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + ggml_backend_buffer_type_t buft; + buft = apir_decode_ggml_buffer_type(dec); + + const char * string = buft->iface.get_name(buft); + + const size_t string_size = strlen(string) + 1; + apir_encode_array_size(enc, string_size); + apir_encode_char_array(enc, string, string_size); + + return 0; +} + +uint32_t backend_buffer_type_get_alignment(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + ggml_backend_buffer_type_t buft; + buft = apir_decode_ggml_buffer_type(dec); + + size_t value = buft->iface.get_alignment(buft); + apir_encode_size_t(enc, &value); + + return 0; +} + +uint32_t backend_buffer_type_get_max_size(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + ggml_backend_buffer_type_t buft; + buft = apir_decode_ggml_buffer_type(dec); + + size_t value = SIZE_MAX; + if (buft->iface.get_max_size) { + value = buft->iface.get_max_size(buft); + } + + apir_encode_size_t(enc, &value); + + return 0; +} + +/* APIR_COMMAND_TYPE_BUFFER_TYPE_IS_HOST is deprecated. Keeping the handler for backward compatibility. */ +uint32_t backend_buffer_type_is_host(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(dec); + const bool is_host = false; + + apir_encode_bool_t(enc, &is_host); + + return 0; +} + +uint32_t backend_buffer_type_alloc_buffer(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + ggml_backend_buffer_type_t buft; + buft = apir_decode_ggml_buffer_type(dec); + + size_t size; + apir_decode_size_t(dec, &size); + + ggml_backend_buffer_t buffer; + + buffer = buft->iface.alloc_buffer(buft, size); + + apir_encode_ggml_buffer(enc, buffer); + + if (buffer) { + apir_track_backend_buffer(buffer); + } + + return 0; +} + +uint32_t backend_buffer_type_get_alloc_size(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + ggml_backend_buffer_type_t buft; + buft = apir_decode_ggml_buffer_type(dec); + + const ggml_tensor * op = apir_decode_ggml_tensor_inplace(dec); + + // Check for decode error + if (op == nullptr) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Failed to decode tensor\n", __func__); + apir_decoder_set_fatal(dec); + return 1; + } + + size_t value; + if (buft->iface.get_alloc_size) { + value = buft->iface.get_alloc_size(buft, op); + } else { + value = ggml_nbytes(op); // Default fallback + } + + apir_encode_size_t(enc, &value); + + return 0; +} + + +/* Inlined backend-dispatched-buffer.cpp */ +#include "ggml-backend-impl-defs.inc" +#include "ggml-backend.h.inc" +#include "ggml-impl-defs.inc" + +#include + +static uint32_t validate_buffer_operation(size_t offset, size_t size, const char * operation) { + // Only check for critical integer overflow - no arbitrary size limits + if (offset > SIZE_MAX - size) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Integer overflow in offset+size: %zu + %zu\n", operation, offset, size); + return 1; + } + + return 0; // Valid +} + +uint32_t backend_buffer_get_base(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + ggml_backend_buffer_t buffer; + buffer = apir_decode_ggml_buffer(dec); + + if (!buffer || apir_decoder_get_fatal(dec)) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid buffer handle from guest\n", __func__); + return 1; + } + + uintptr_t base = (uintptr_t) buffer->iface.get_base(buffer); + apir_encode_uintptr_t(enc, &base); + + return 0; +} + +uint32_t backend_buffer_set_tensor(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(enc); + + ggml_backend_buffer_t buffer; + buffer = apir_decode_ggml_buffer(dec); + + if (!buffer || apir_decoder_get_fatal(dec)) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid buffer handle from guest\n", __func__); + return 1; + } + + ggml_tensor * tensor; + // safe to remove the const qualifier here + tensor = (ggml_tensor *) (uintptr_t) apir_decode_ggml_tensor(dec); + + uint32_t shmem_res_id; + apir_decode_virtgpu_shmem_res_id(dec, &shmem_res_id); + + size_t offset; + apir_decode_size_t(dec, &offset); + + size_t size; + apir_decode_size_t(dec, &size); + + if (validate_buffer_operation(offset, size, __func__) != 0) { + return 1; + } + + void * shmem_data = ctx->iface->get_shmem_ptr(ctx->ctx_id, shmem_res_id); + + if (!shmem_data) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Couldn't get the shmem addr from virgl\n", __func__); + return 1; + } + + buffer->iface.set_tensor(buffer, tensor, shmem_data, offset, size); + + return 0; +} + +uint32_t backend_buffer_get_tensor(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(enc); + + ggml_backend_buffer_t buffer; + buffer = apir_decode_ggml_buffer(dec); + + if (!buffer || apir_decoder_get_fatal(dec)) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid buffer handle from guest\n", __func__); + return 1; + } + + const ggml_tensor * tensor; + // safe to remove the const qualifier here + tensor = apir_decode_ggml_tensor(dec); + + uint32_t shmem_res_id; + apir_decode_virtgpu_shmem_res_id(dec, &shmem_res_id); + + size_t offset; + apir_decode_size_t(dec, &offset); + + size_t size; + apir_decode_size_t(dec, &size); + + if (validate_buffer_operation(offset, size, __func__) != 0) { + return 1; + } + + void * shmem_data = ctx->iface->get_shmem_ptr(ctx->ctx_id, shmem_res_id); + if (!shmem_data) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Couldn't get the shmem addr from virgl\n", __func__); + return 1; + } + + buffer->iface.get_tensor(buffer, tensor, shmem_data, offset, size); + + return 0; +} + +uint32_t backend_buffer_cpy_tensor(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + + ggml_backend_buffer_t buffer; + buffer = apir_decode_ggml_buffer(dec); + + if (!buffer || apir_decoder_get_fatal(dec)) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid buffer handle from guest\n", __func__); + return 1; + } + + const ggml_tensor * src; + // safe to remove the const qualifier here + src = apir_decode_ggml_tensor(dec); + ggml_tensor * dst = (ggml_tensor *) (uintptr_t) apir_decode_ggml_tensor(dec); + + bool ret = buffer->iface.cpy_tensor(buffer, src, (ggml_tensor *) dst); + + apir_encode_bool_t(enc, &ret); + + return 0; +} + +uint32_t backend_buffer_clear(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(enc); + + ggml_backend_buffer_t buffer; + buffer = apir_decode_ggml_buffer(dec); + + if (!buffer || apir_decoder_get_fatal(dec)) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid buffer handle from guest\n", __func__); + return 1; + } + + uint8_t value; + apir_decode_uint8_t(dec, &value); + + buffer->iface.clear(buffer, value); + + return 0; +} + +uint32_t backend_buffer_free_buffer(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + GGML_UNUSED(enc); + + ggml_backend_buffer_t buffer; + buffer = apir_decode_ggml_buffer(dec); + + if (!buffer || apir_decoder_get_fatal(dec)) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid buffer handle from guest\n", __func__); + return 1; + } + + if (!apir_untrack_backend_buffer(buffer)) { + GGML_LOG_WARN(GGML_VIRTGPU_BCK "%s: unknown buffer %p\n", __func__, (void *) buffer); + return 1; + } + + buffer->iface.free_buffer(buffer); + + return 0; +} + + +/* Inlined backend-dispatched-backend.cpp */ +#include "ggml-backend-impl-defs.inc" +#include "ggml-backend.h.inc" +#include "ggml-impl-defs.inc" + +#include + +static uint32_t validate_graph_operation(size_t cgraph_size, uint32_t shmem_res_id, const char * operation) { + if (cgraph_size == 0) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Zero-size computation graph\n", operation); + return 1; + } + + // place-holder: validate that the size of shmem_res_id is <= cgraph_size + // need to add another method in the Virgl->APIR callback interface + GGML_UNUSED(shmem_res_id); + + return 0; // Valid +} + +uint32_t backend_backend_graph_compute(apir_encoder * enc, apir_decoder * dec, virgl_apir_context * ctx) { + GGML_UNUSED(ctx); + + static bool async_backend_initialized = false; + static bool async_backend; + + if (!async_backend_initialized) { + ggml_backend_dev_props props; + + dev->iface.get_props(dev, &props); + async_backend = props.caps.async; + async_backend_initialized = true; + } + + uint32_t shmem_res_id; + apir_decode_virtgpu_shmem_res_id(dec, &shmem_res_id); + + const void * shmem_data = ctx->iface->get_shmem_ptr(ctx->ctx_id, shmem_res_id); + if (!shmem_data) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Couldn't get the shmem addr from virgl\n", __func__); + apir_decoder_set_fatal(dec); + return 1; + } + size_t cgraph_size; + apir_decode_size_t(dec, &cgraph_size); + + if (validate_graph_operation(cgraph_size, shmem_res_id, __func__) != 0) { + apir_decoder_set_fatal(dec); + return 1; + } + + apir_decoder secondary_dec = apir_new_decoder((const char *) shmem_data, cgraph_size); + + ggml_cgraph * cgraph = apir_decode_ggml_cgraph(&secondary_dec, cgraph_size); + + if (!cgraph || apir_decoder_get_fatal(&secondary_dec)) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Failed to deserialize computation graph\n", __func__); + return 1; + } + + if (cgraph->n_nodes < 0 || cgraph->n_leafs < 0) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Invalid negative node/leaf count: nodes=%d leafs=%d\n", __func__, + cgraph->n_nodes, cgraph->n_leafs); + return 1; + } + + ggml_status status; +#if APIR_BACKEND_CHECK_SUPPORTS_OP == 1 + for (int idx = 0; idx < cgraph->n_nodes; idx++) { + ggml_tensor * op = ggml_graph_node(cgraph, idx); + if (dev->iface.supports_op(dev, op)) { + continue; + } + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Graph node %d (%s) not supported by the backend\n", __func__, idx, + ggml_op_desc(op)); + + status = GGML_STATUS_ABORTED; + apir_encode_ggml_status(enc, &status); + + return 0; + } +#endif + + // Check if backend is properly initialized + if (!bck) { + GGML_LOG_ERROR(GGML_VIRTGPU_BCK "%s: Backend not initialized (bck is null)\n", __func__); + + return 1; + } + + status = bck->iface.graph_compute(bck, cgraph); + + if (async_backend && bck->iface.synchronize) { + bck->iface.synchronize(bck); + } + + apir_encode_ggml_status(enc, &status); + + return 0; +} + + +/* Inlined apir_cs_ggml-rpc-back.cpp */ +#include "ggml-backend-impl-defs.inc" +#include "ggml-impl-defs.inc" + + +#include +#include +#include +#include + +std::unordered_set backend_buffers; + +void apir_track_backend_buffer(ggml_backend_buffer_t buffer) { + backend_buffers.insert(buffer); +} + +bool apir_untrack_backend_buffer(ggml_backend_buffer_t buffer) { + auto it = backend_buffers.find(buffer); + if (it == backend_buffers.end()) { + return false; + } + + backend_buffers.erase(it); + return true; +} + +std::unordered_set apir_get_track_backend_buffers() { + return backend_buffers; +} + +ggml_tensor * apir_deserialize_tensor(ggml_context * ctx, const apir_rpc_tensor * tensor) { + ggml_tensor * result = + ggml_new_tensor_4d(ctx, (ggml_type) tensor->type, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); + for (uint32_t i = 0; i < GGML_MAX_DIMS; i++) { + result->nb[i] = tensor->nb[i]; + } + result->buffer = reinterpret_cast(tensor->buffer); + if (result->buffer && backend_buffers.find(result->buffer) == backend_buffers.end()) { + printf("WARNING: HOST BUFFER NOT FOUND | %p\n", (void *) result->buffer); + result->buffer = nullptr; + } + + uint64_t tensor_data = tensor->data; + if (result->buffer) { + // require that the tensor data does not go beyond the buffer end + uint64_t tensor_size = (uint64_t) ggml_nbytes(result); + uint64_t buffer_start = (uint64_t) ggml_backend_buffer_get_base(result->buffer); + uint64_t buffer_size = (uint64_t) ggml_backend_buffer_get_size(result->buffer); + + // tensor->data is serialized as an offset to the buffer base address + tensor_data += buffer_start; + + GGML_ASSERT(tensor_data + tensor_size >= tensor_data); // check for overflow + GGML_ASSERT(tensor_data >= buffer_start && tensor_data + tensor_size <= buffer_start + buffer_size); + } + + result->op = (ggml_op) tensor->op; + for (uint32_t i = 0; i < GGML_MAX_OP_PARAMS / sizeof(int32_t); i++) { + result->op_params[i] = tensor->op_params[i]; + } + result->flags = tensor->flags; + result->data = reinterpret_cast(tensor_data); + ggml_set_name(result, tensor->name); + return result; +} + +ggml_tensor * apir_create_node(uint64_t id, + ggml_context * ctx, + const std::unordered_map & tensor_ptrs, + std::unordered_map & tensor_map) { + if (id == 0) { + return nullptr; + } + if (tensor_map.find(id) != tensor_map.end()) { + return tensor_map[id]; + } + const apir_rpc_tensor * tensor = tensor_ptrs.at(id); + ggml_tensor * result = apir_deserialize_tensor(ctx, tensor); + if (result == nullptr) { + return nullptr; + } + tensor_map[id] = result; + for (int i = 0; i < GGML_MAX_SRC; i++) { + result->src[i] = apir_create_node(tensor->src[i], ctx, tensor_ptrs, tensor_map); + } + result->view_src = apir_create_node(tensor->view_src, ctx, tensor_ptrs, tensor_map); + result->view_offs = tensor->view_offs; + return result; +} + +ggml_cgraph * apir_deserialize_graph(uint32_t n_nodes, + uint32_t n_tensors, + const apir_rpc_tensor * tensors, + const uint64_t * nodes) { + size_t buf_size = ggml_tensor_overhead() * (n_nodes + n_tensors) + ggml_graph_overhead_custom(n_nodes, false); + ggml_init_params params = { + /*.mem_size =*/buf_size, + /*.mem_buffer =*/NULL, + /*.no_alloc =*/true, + }; + ggml_context * ctx = ggml_init(params); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, n_nodes, false); + graph->n_nodes = n_nodes; + std::unordered_map tensor_ptrs; + for (uint32_t i = 0; i < n_tensors; i++) { + tensor_ptrs[tensors[i].id] = &tensors[i]; + } + std::unordered_map tensor_map; + for (uint32_t i = 0; i < n_nodes; i++) { + int64_t id; + memcpy(&id, &nodes[i], sizeof(id)); + graph->nodes[i] = apir_create_node(id, ctx, tensor_ptrs, tensor_map); + } + + return graph; +} diff --git a/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h b/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h deleted file mode 100644 index 6bf97e8a3a24..000000000000 --- a/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h +++ /dev/null @@ -1,95 +0,0 @@ -#pragma once - -/* the rest of this file must match virglrenderer/src/apir-protocol.h */ - -#include - -#include - -#define APIR_PROTOCOL_MAJOR 0 -#define APIR_PROTOCOL_MINOR 1 - -#define APIR_HANDSHAKE_MAGIC 0xab1e - -enum ApirCommandType { - APIR_COMMAND_TYPE_HANDSHAKE = 0, - APIR_COMMAND_TYPE_LOADLIBRARY = 1, - APIR_COMMAND_TYPE_FORWARD = 2, - - APIR_COMMAND_TYPE_LENGTH = 3, -}; - -typedef uint64_t ApirCommandFlags; - -enum ApirLoadLibraryReturnCode { - APIR_LOAD_LIBRARY_SUCCESS = 0, - // these error codes are returned by the Virglrenderer APIR component - APIR_LOAD_LIBRARY_HYPERCALL_INITIALIZATION_ERROR = 1, - APIR_LOAD_LIBRARY_ALREADY_LOADED = 2, - APIR_LOAD_LIBRARY_ENV_VAR_MISSING = 3, - APIR_LOAD_LIBRARY_CANNOT_OPEN = 4, - APIR_LOAD_LIBRARY_SYMBOL_MISSING = 5, - // any value greater than this is an APIR *backend library* initialization return code - APIR_LOAD_LIBRARY_INIT_BASE_INDEX = 6, -}; - -enum ApirForwardReturnCode { - APIR_FORWARD_SUCCESS = 0, - // these error codes are returned by the Virglrenderer APIR component - APIR_FORWARD_NO_DISPATCH_FCT = 1, - APIR_FORWARD_TIMEOUT = 2, - APIR_FORWARD_FAILED_TO_SYNC_STREAMS = 3, - // any value greater than this index an APIR *backend library* forward return code - APIR_FORWARD_BASE_INDEX = 4, -}; - -__attribute__((unused)) static inline const char * apir_command_name(ApirCommandType type) { - switch (type) { - case APIR_COMMAND_TYPE_HANDSHAKE: - return "HandShake"; - case APIR_COMMAND_TYPE_LOADLIBRARY: - return "LoadLibrary"; - case APIR_COMMAND_TYPE_FORWARD: - return "Forward"; - default: - return "unknown"; - } -} - -__attribute__((unused)) static const char * apir_load_library_error(ApirLoadLibraryReturnCode code) { -#define APIR_LOAD_LIBRARY_ERROR(code_name) \ - do { \ - if (code == code_name) \ - return #code_name; \ - } while (0) - - APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_SUCCESS); - APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_HYPERCALL_INITIALIZATION_ERROR); - APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_ALREADY_LOADED); - APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_ENV_VAR_MISSING); - APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_CANNOT_OPEN); - APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_SYMBOL_MISSING); - APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_INIT_BASE_INDEX); - - return "Unknown APIR_COMMAND_TYPE_LoadLibrary error"; - -#undef APIR_LOAD_LIBRARY_ERROR -} - -__attribute__((unused)) static const char * apir_forward_error(ApirForwardReturnCode code) { -#define APIR_FORWARD_ERROR(code_name) \ - do { \ - if (code == code_name) \ - return #code_name; \ - } while (0) - - APIR_FORWARD_ERROR(APIR_FORWARD_SUCCESS); - APIR_FORWARD_ERROR(APIR_FORWARD_NO_DISPATCH_FCT); - APIR_FORWARD_ERROR(APIR_FORWARD_TIMEOUT); - APIR_FORWARD_ERROR(APIR_FORWARD_FAILED_TO_SYNC_STREAMS); - APIR_FORWARD_ERROR(APIR_FORWARD_BASE_INDEX); - - return "Unknown APIR_COMMAND_TYPE_FORWARD error"; - -#undef APIR_FORWARD_ERROR -} diff --git a/ggml/src/ggml-virtgpu/backend/shared/apir_backend.gen.h b/ggml/src/ggml-virtgpu/backend/shared/apir_backend.gen.h deleted file mode 100644 index 520ac9c72990..000000000000 --- a/ggml/src/ggml-virtgpu/backend/shared/apir_backend.gen.h +++ /dev/null @@ -1,94 +0,0 @@ -typedef enum ApirBackendCommandType { - - /* device */ - APIR_COMMAND_TYPE_DEVICE_GET_DEVICE_COUNT = 0, - APIR_COMMAND_TYPE_DEVICE_GET_COUNT = 1, - APIR_COMMAND_TYPE_DEVICE_GET_NAME = 2, - APIR_COMMAND_TYPE_DEVICE_GET_DESCRIPTION = 3, - APIR_COMMAND_TYPE_DEVICE_GET_TYPE = 4, - APIR_COMMAND_TYPE_DEVICE_GET_MEMORY = 5, - APIR_COMMAND_TYPE_DEVICE_SUPPORTS_OP = 6, - APIR_COMMAND_TYPE_DEVICE_GET_BUFFER_TYPE = 7, - APIR_COMMAND_TYPE_DEVICE_GET_PROPS = 8, - APIR_COMMAND_TYPE_DEVICE_BUFFER_FROM_PTR = 9, - - /* buffer-type */ - APIR_COMMAND_TYPE_BUFFER_TYPE_GET_NAME = 10, - APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALIGNMENT = 11, - APIR_COMMAND_TYPE_BUFFER_TYPE_GET_MAX_SIZE = 12, - APIR_COMMAND_TYPE_BUFFER_TYPE_IS_HOST = 13, - APIR_COMMAND_TYPE_BUFFER_TYPE_ALLOC_BUFFER = 14, - APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALLOC_SIZE = 15, - - /* buffer */ - APIR_COMMAND_TYPE_BUFFER_GET_BASE = 16, - APIR_COMMAND_TYPE_BUFFER_SET_TENSOR = 17, - APIR_COMMAND_TYPE_BUFFER_GET_TENSOR = 18, - APIR_COMMAND_TYPE_BUFFER_CPY_TENSOR = 19, - APIR_COMMAND_TYPE_BUFFER_CLEAR = 20, - APIR_COMMAND_TYPE_BUFFER_FREE_BUFFER = 21, - - /* backend */ - APIR_COMMAND_TYPE_BACKEND_GRAPH_COMPUTE = 22, - - // last command_type index + 1 - APIR_BACKEND_DISPATCH_TABLE_COUNT = 23, -} ApirBackendCommandType; - -static inline const char * apir_dispatch_command_name(ApirBackendCommandType type) { - switch (type) { - /* device */ - case APIR_COMMAND_TYPE_DEVICE_GET_DEVICE_COUNT: - return "device_get_device_count"; - case APIR_COMMAND_TYPE_DEVICE_GET_COUNT: - return "device_get_count"; - case APIR_COMMAND_TYPE_DEVICE_GET_NAME: - return "device_get_name"; - case APIR_COMMAND_TYPE_DEVICE_GET_DESCRIPTION: - return "device_get_description"; - case APIR_COMMAND_TYPE_DEVICE_GET_TYPE: - return "device_get_type"; - case APIR_COMMAND_TYPE_DEVICE_GET_MEMORY: - return "device_get_memory"; - case APIR_COMMAND_TYPE_DEVICE_SUPPORTS_OP: - return "device_supports_op"; - case APIR_COMMAND_TYPE_DEVICE_GET_BUFFER_TYPE: - return "device_get_buffer_type"; - case APIR_COMMAND_TYPE_DEVICE_GET_PROPS: - return "device_get_props"; - case APIR_COMMAND_TYPE_DEVICE_BUFFER_FROM_PTR: - return "device_buffer_from_ptr"; - /* buffer-type */ - case APIR_COMMAND_TYPE_BUFFER_TYPE_GET_NAME: - return "buffer_type_get_name"; - case APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALIGNMENT: - return "buffer_type_get_alignment"; - case APIR_COMMAND_TYPE_BUFFER_TYPE_GET_MAX_SIZE: - return "buffer_type_get_max_size"; - case APIR_COMMAND_TYPE_BUFFER_TYPE_IS_HOST: - return "buffer_type_is_host"; - case APIR_COMMAND_TYPE_BUFFER_TYPE_ALLOC_BUFFER: - return "buffer_type_alloc_buffer"; - case APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALLOC_SIZE: - return "buffer_type_get_alloc_size"; - /* buffer */ - case APIR_COMMAND_TYPE_BUFFER_GET_BASE: - return "buffer_get_base"; - case APIR_COMMAND_TYPE_BUFFER_SET_TENSOR: - return "buffer_set_tensor"; - case APIR_COMMAND_TYPE_BUFFER_GET_TENSOR: - return "buffer_get_tensor"; - case APIR_COMMAND_TYPE_BUFFER_CPY_TENSOR: - return "buffer_cpy_tensor"; - case APIR_COMMAND_TYPE_BUFFER_CLEAR: - return "buffer_clear"; - case APIR_COMMAND_TYPE_BUFFER_FREE_BUFFER: - return "buffer_free_buffer"; - /* backend */ - case APIR_COMMAND_TYPE_BACKEND_GRAPH_COMPUTE: - return "backend_graph_compute"; - - default: - return "unknown"; - } -} diff --git a/ggml/src/ggml-virtgpu/backend/shared/apir_backend.h b/ggml/src/ggml-virtgpu/backend/shared/apir_backend.h deleted file mode 100644 index da1e21b5b2fd..000000000000 --- a/ggml/src/ggml-virtgpu/backend/shared/apir_backend.h +++ /dev/null @@ -1,50 +0,0 @@ -#pragma once - -#include "apir_backend.gen.h" - -#include // for uintptr_t -#include // for timespec, clock_gettime - -#define APIR_BACKEND_INITIALIZE_SUCCESS 0 -#define APIR_BACKEND_INITIALIZE_CANNOT_OPEN_BACKEND_LIBRARY 1 -#define APIR_BACKEND_INITIALIZE_CANNOT_OPEN_GGML_LIBRARY 2 -#define APIR_BACKEND_INITIALIZE_MISSING_BACKEND_SYMBOLS 3 -#define APIR_BACKEND_INITIALIZE_MISSING_GGML_SYMBOLS 4 -#define APIR_BACKEND_INITIALIZE_BACKEND_FAILED 5 -#define APIR_BACKEND_INITIALIZE_BACKEND_REG_FAILED 6 -#define APIR_BACKEND_INITIALIZE_ALREADY_INITED 7 -#define APIR_BACKEND_INITIALIZE_NO_DEVICE 8 -#define APIR_BACKEND_INITIALIZE_BACKEND_INIT_FAILED 9 - -// new entries here need to be added to the apir_backend_initialize_error function below - -#define APIR_BACKEND_FORWARD_INDEX_INVALID 6 - -// 0 is fast, 1 avoids the backend to crash if an unsupported tensor is received -#define APIR_BACKEND_CHECK_SUPPORTS_OP 0 - -typedef uintptr_t apir_buffer_type_host_handle_t; -typedef uintptr_t apir_buffer_host_handle_t; - -static const char * apir_backend_initialize_error(int code) { -#define APIR_BACKEND_INITIALIZE_ERROR(code_name) \ - do { \ - if (code == code_name) \ - return #code_name; \ - } while (0) - - APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_SUCCESS); - APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_CANNOT_OPEN_BACKEND_LIBRARY); - APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_CANNOT_OPEN_GGML_LIBRARY); - APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_MISSING_BACKEND_SYMBOLS); - APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_MISSING_GGML_SYMBOLS); - APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_BACKEND_FAILED); - APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_BACKEND_REG_FAILED); - APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_ALREADY_INITED); - APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_NO_DEVICE); - APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_BACKEND_INIT_FAILED); - - return "Unknown APIR_BACKEND_INITIALIZE error:/"; - -#undef APIR_BACKEND_INITIALIZE_ERROR -} diff --git a/ggml/src/ggml-virtgpu/backend/shared/apir_cs.h b/ggml/src/ggml-virtgpu/backend/shared/apir_cs.h deleted file mode 100644 index 64bf2ec9609d..000000000000 --- a/ggml/src/ggml-virtgpu/backend/shared/apir_cs.h +++ /dev/null @@ -1,378 +0,0 @@ -#pragma once - -#include "ggml-impl.h" - -#include -#include - -#define likely(x) __builtin_expect(!!(x), 1) -#define unlikely(x) __builtin_expect(!!(x), 0) - -struct apir_encoder { - char * cur; - const char * start; - const char * end; - bool fatal; -}; - -struct apir_decoder { - const char * cur; - const char * end; - bool fatal; -}; - -/* - * new encoder and decoder - */ - -static apir_decoder apir_new_decoder(const char * ptr, size_t size) { - apir_decoder dec = { - .cur = ptr, - .end = ptr + size, - .fatal = false, - }; - - return dec; -} - -static apir_encoder apir_new_encoder(char * ptr, size_t size) { - apir_encoder enc = { - .cur = ptr, - .start = ptr, - .end = ptr + size, - .fatal = false, - }; - - return enc; -} - -/* - * fatal flag handling - */ - -static inline void apir_encoder_reset_fatal(apir_encoder * enc) { - enc->fatal = false; -} - -static inline void apir_encoder_set_fatal(apir_encoder * enc) { - enc->fatal = true; -} - -static inline bool apir_encoder_get_fatal(const apir_encoder * enc) { - return enc->fatal; -} - -static inline void apir_decoder_reset_fatal(apir_decoder * dec) { - dec->fatal = false; -} - -static inline void apir_decoder_set_fatal(apir_decoder * dec) { - dec->fatal = true; -} - -static inline bool apir_decoder_get_fatal(const apir_decoder * dec) { - return dec->fatal; -} - -/* - * encode peek - */ - -static inline bool apir_decoder_peek_internal(apir_decoder * dec, size_t size, void * val, size_t val_size) { - assert(val_size <= size); - - if (unlikely(size > (size_t) (dec->end - dec->cur))) { - GGML_LOG_ERROR("%s: reading too much from the decoder ...\n", __func__); - apir_decoder_set_fatal(dec); - memset(val, 0, val_size); - return false; - } - - /* we should not rely on the compiler to optimize away memcpy... */ - memcpy(val, dec->cur, val_size); - return true; -} - -static inline void apir_decoder_peek(apir_decoder * dec, size_t size, void * val, size_t val_size) { - apir_decoder_peek_internal(dec, size, val, val_size); -} - -static inline const void * apir_decoder_use_inplace(apir_decoder * dec, size_t size) { - if (unlikely(size > (size_t) (dec->end - dec->cur))) { - GGML_LOG_ERROR("%s: reading too much from the decoder ...\n", __func__); - apir_decoder_set_fatal(dec); - return NULL; - } - const void * addr = dec->cur; - dec->cur += size; - - return addr; -} - -/* - * read/write - */ - -static inline void apir_decoder_read(apir_decoder * dec, size_t size, void * val, size_t val_size) { - if (apir_decoder_peek_internal(dec, size, val, val_size)) { - dec->cur += size; - } -} - -static inline char * apir_encoder_write(apir_encoder * enc, size_t size, const void * val, size_t val_size) { - assert(val_size <= size); - assert(size <= ((size_t) (enc->end - enc->cur))); - - char * write_addr = enc->cur; - /* we should not rely on the compiler to optimize away memcpy... */ - memcpy(write_addr, val, val_size); - enc->cur += size; - - return write_addr; -} - -/* - * encode/decode - */ - -static inline void apir_decode(apir_decoder * dec, size_t size, void * data, size_t data_size) { - assert(size % 4 == 0); - apir_decoder_read(dec, size, data, data_size); -} - -static inline void apir_encode(apir_encoder * enc, size_t size, const void * data, size_t data_size) { - assert(size % 4 == 0); - apir_encoder_write(enc, size, data, data_size); -} - -/* - * typed encode/decode - */ - -/* uint8_t */ - -static inline void apir_encode_uint8_t(apir_encoder * enc, const uint8_t * val) { - apir_encode(enc, sizeof(int), val, sizeof(*val)); -} - -static inline void apir_decode_uint8_t(apir_decoder * dec, uint8_t * val) { - apir_decode(dec, sizeof(int), val, sizeof(*val)); -} - -/* uint64_t */ - -static inline void apir_encode_uint64_t(apir_encoder * enc, const uint64_t * val) { - apir_encode(enc, 8, val, sizeof(*val)); -} - -static inline void apir_decode_uint64_t(apir_decoder * dec, uint64_t * val) { - apir_decode(dec, 8, val, sizeof(*val)); -} - -static inline void apir_encode_uint64_t_array(apir_encoder * enc, const uint64_t * val, uint32_t count) { - const size_t size = sizeof(*val) * count; - assert(size >= count); - apir_encode(enc, size, val, size); -} - -static inline void apir_decode_uint64_t_array(apir_decoder * dec, uint64_t * val, uint32_t count) { - const size_t size = sizeof(*val) * count; - assert(size >= count); - apir_decode(dec, size, val, size); -} - -static inline const uint64_t * apir_decode_uint64_t_array_inplace(apir_decoder * dec, uint32_t count) { - return (uint64_t *) (uintptr_t) apir_decoder_use_inplace(dec, count * sizeof(uint64_t)); -} - -/* int32_t */ - -static inline void apir_encode_int32_t(apir_encoder * enc, const int32_t * val) { - apir_encode(enc, 4, val, sizeof(*val)); -} - -static inline void apir_decode_int32_t(apir_decoder * dec, int32_t * val) { - apir_decode(dec, 4, val, sizeof(*val)); -} - -static inline void apir_encode_int32_t_array(apir_encoder * enc, const int32_t * val, uint32_t count) { - const size_t size = sizeof(*val) * count; - assert(size >= count); - apir_encode(enc, size, val, size); -} - -static inline void apir_decode_int32_t_array(apir_decoder * dec, int32_t * val, uint32_t count) { - const size_t size = sizeof(*val) * count; - assert(size >= count); - apir_decode(dec, size, val, size); -} - -/* array size (uint64_t) */ - -static inline void apir_encode_array_size(apir_encoder * enc, uint64_t size) { - apir_encode_uint64_t(enc, &size); -} - -static inline uint64_t apir_decode_array_size(apir_decoder * dec, uint64_t expected_size) { - uint64_t size; - apir_decode_uint64_t(dec, &size); - if (size != expected_size) { - GGML_LOG_ERROR("%s: Couldn't decode array from the decoder\n", __func__); - apir_decoder_set_fatal(dec); - size = 0; - } - return size; -} - -static inline uint64_t apir_decode_array_size_unchecked(apir_decoder * dec) { - uint64_t size; - apir_decode_uint64_t(dec, &size); - return size; -} - -/* non-array pointer */ - -static inline bool apir_encode_simple_pointer(apir_encoder * enc, const void * val) { - apir_encode_array_size(enc, val ? 1 : 0); - return val; -} - -static inline bool apir_decode_simple_pointer(apir_decoder * dec) { - return apir_decode_array_size_unchecked(dec); -} - -/* uint32_t */ - -static inline void apir_encode_uint32_t(apir_encoder * enc, const uint32_t * val) { - apir_encode(enc, 4, val, sizeof(*val)); -} - -static inline void apir_decode_uint32_t(apir_decoder * dec, uint32_t * val) { - apir_decode(dec, 4, val, sizeof(*val)); -} - -static inline void apir_encode_uint32_t_array(apir_encoder * enc, const uint32_t * val, uint32_t count) { - const size_t size = sizeof(*val) * count; - assert(size >= count); - apir_encode(enc, size, val, size); -} - -static inline void apir_decode_uint32_t_array(apir_decoder * dec, uint32_t * val, uint32_t count) { - const size_t size = sizeof(*val) * count; - assert(size >= count); - apir_decode(dec, size, val, size); -} - -/* size_t */ - -static inline void apir_encode_size_t(apir_encoder * enc, const size_t * val) { - const uint64_t tmp = *val; - apir_encode_uint64_t(enc, &tmp); -} - -static inline void apir_decode_size_t(apir_decoder * dec, size_t * val) { - uint64_t tmp; - apir_decode_uint64_t(dec, &tmp); - *val = tmp; -} - -static inline void apir_encode_size_t_array(apir_encoder * enc, const size_t * val, uint32_t count) { - if (sizeof(size_t) == sizeof(uint64_t)) { - apir_encode_uint64_t_array(enc, (const uint64_t *) val, count); - } else { - for (uint32_t i = 0; i < count; i++) { - apir_encode_size_t(enc, &val[i]); - } - } -} - -static inline void apir_decode_size_t_array(apir_decoder * dec, size_t * val, uint32_t count) { - if (sizeof(size_t) == sizeof(uint64_t)) { - apir_decode_uint64_t_array(dec, (uint64_t *) val, count); - } else { - for (uint32_t i = 0; i < count; i++) { - apir_decode_size_t(dec, &val[i]); - } - } -} - -/* opaque blob */ - -static inline void apir_encode_blob_array(apir_encoder * enc, const void * val, size_t size) { - apir_encode(enc, (size + 3) & ~3, val, size); -} - -static inline void apir_decode_blob_array(apir_decoder * dec, void * val, size_t size) { - apir_decode(dec, (size + 3) & ~3, val, size); -} - -/* string */ - -static inline void apir_encode_char_array(apir_encoder * enc, const char * val, size_t size) { - assert(size && strlen(val) < size); - apir_encode_blob_array(enc, val, size); -} - -static inline void apir_decode_char_array(apir_decoder * dec, char * val, size_t size) { - apir_decode_blob_array(dec, val, size); - if (size) { - val[size - 1] = '\0'; - } else { - GGML_LOG_ERROR("%s: Couldn't decode the blog array\n", __func__); - apir_decoder_set_fatal(dec); - } -} - -/* (temp) buffer allocation */ - -static inline void * apir_decoder_alloc_array(size_t size, size_t count) { - size_t alloc_size; - if (unlikely(__builtin_mul_overflow(size, count, &alloc_size))) { - GGML_LOG_ERROR("%s: overflow in array allocation of %zu * %zu bytes\n", __func__, size, count); - return NULL; - } - - return malloc(alloc_size); -} - -/* bool */ - -static inline void apir_encode_bool_t(apir_encoder * enc, const bool * val) { - apir_encode(enc, sizeof(int), val, sizeof(bool)); -} - -static inline void apir_decode_bool_t(apir_decoder * dec, bool * val) { - apir_decode(dec, sizeof(int), val, sizeof(bool)); -} - -/* apir_buffer_type_host_handle_t */ - -static inline void apir_encode_apir_buffer_type_host_handle_t(apir_encoder * enc, - const apir_buffer_type_host_handle_t * val) { - apir_encode(enc, sizeof(apir_buffer_type_host_handle_t), val, sizeof(apir_buffer_type_host_handle_t)); -} - -static inline void apir_decode_apir_buffer_type_host_handle_t(apir_decoder * dec, - apir_buffer_type_host_handle_t * val) { - apir_decode(dec, sizeof(apir_buffer_type_host_handle_t), val, sizeof(apir_buffer_type_host_handle_t)); -} - -/* apir_buffer_host_handle_t */ - -static inline void apir_encode_apir_buffer_host_handle_t(apir_encoder * enc, const apir_buffer_host_handle_t * val) { - apir_encode(enc, sizeof(apir_buffer_host_handle_t), val, sizeof(apir_buffer_host_handle_t)); -} - -static inline void apir_decode_apir_buffer_host_handle_t(apir_decoder * dec, apir_buffer_host_handle_t * val) { - apir_decode(dec, sizeof(apir_buffer_host_handle_t), val, sizeof(apir_buffer_host_handle_t)); -} - -/* uintptr_t */ - -static inline void apir_encode_uintptr_t(apir_encoder * enc, const uintptr_t * val) { - apir_encode(enc, sizeof(*val), val, sizeof(*val)); -} - -static inline void apir_decode_uintptr_t(apir_decoder * dec, uintptr_t * val) { - apir_decode(dec, sizeof(*val), val, sizeof(*val)); -} diff --git a/ggml/src/ggml-virtgpu/backend/shared/apir_cs_ggml.h b/ggml/src/ggml-virtgpu/backend/shared/apir_cs_ggml.h deleted file mode 100644 index fabe3e401ca2..000000000000 --- a/ggml/src/ggml-virtgpu/backend/shared/apir_cs_ggml.h +++ /dev/null @@ -1,232 +0,0 @@ -#include "apir_cs.h" -#include "apir_cs_rpc.h" -#include "ggml-impl.h" - -// ggml_buffer_to_apir_host_handle(ggml_backend_buffer_t buffer); - -static inline void apir_encode_ggml_buffer_host_handle(apir_encoder * enc, const apir_buffer_host_handle_t * handle); - -static inline ggml_backend_buffer_t apir_decode_ggml_buffer(apir_decoder * dec); - -/* apir_rpc_tensor */ - -static inline void apir_encode_rcp_tensor(apir_encoder * enc, const apir_rpc_tensor * apir_rpc_tensor) { - size_t apir_rpc_tensor_size = sizeof(*apir_rpc_tensor); - apir_encode(enc, apir_rpc_tensor_size, apir_rpc_tensor, apir_rpc_tensor_size); -} - -static inline apir_rpc_tensor * apir_decode_apir_rpc_tensor_inplace(apir_decoder * dec) { - size_t apir_rpc_tensor_size = sizeof(apir_rpc_tensor); - - return (apir_rpc_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, apir_rpc_tensor_size); -} - -static inline apir_rpc_tensor * apir_decode_apir_rpc_tensor_array_inplace(apir_decoder * dec, uint32_t n_tensors) { - size_t apir_rpc_tensor_size = sizeof(apir_rpc_tensor) * n_tensors; - - return (apir_rpc_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, apir_rpc_tensor_size); -} - -/* ggml_tensor */ - -static inline void apir_encode_ggml_tensor(apir_encoder * enc, const ggml_tensor * tensor) { - apir_rpc_tensor serialized = apir_serialize_tensor(tensor); - - apir_encode_rcp_tensor(enc, &serialized); -} - -static inline const ggml_tensor * apir_decode_ggml_tensor(apir_decoder * dec) { - const apir_rpc_tensor * apir_rpc_tensor = apir_decode_apir_rpc_tensor_inplace(dec); - - if (!apir_rpc_tensor) { - return NULL; - } - - ggml_init_params params{ - /*.mem_size =*/ggml_tensor_overhead(), - /*.mem_buffer =*/NULL, - /*.no_alloc =*/true, - }; - - ggml_context * ctx = ggml_init(params); - - const ggml_tensor * tensor = apir_deserialize_tensor(ctx, apir_rpc_tensor); - - return tensor; -} - -/* *** ggml_backend_buffer_type_t *** */ - -// ggml_backend_buffer_type_t is a POINTER (to a struct). -// Only the host pointer is shared between the host and guest. -// The guest stores it in `buft->context`. -// The host simply writes the pointer address in the buffer variable. - -static inline void apir_encode_ggml_buffer_type(apir_encoder * enc, ggml_backend_buffer_type_t buft) { - apir_buffer_type_host_handle_t handle = ggml_buffer_type_to_apir_handle(buft); - apir_encoder_write(enc, sizeof(handle), &handle, sizeof(handle)); -} - -static inline ggml_backend_buffer_type_t apir_decode_ggml_buffer_type(apir_decoder * dec) { - apir_buffer_type_host_handle_t handle; - - apir_decoder_read(dec, sizeof(handle), &handle, sizeof(handle)); - - return (ggml_backend_buffer_type_t) handle; -} - -static inline void apir_encode_apir_buffer_type_host_handle(apir_encoder * enc, apir_buffer_type_host_handle_t handle) { - apir_encoder_write(enc, sizeof(handle), &handle, sizeof(handle)); -} - -static inline apir_buffer_type_host_handle_t apir_decode_apir_buffer_type_host_handle(apir_decoder * dec) { - apir_buffer_type_host_handle_t handle; - - apir_decoder_read(dec, sizeof(handle), &handle, sizeof(handle)); - - return handle; -} - -/* *** ggml_backend_type_t *** */ - -// ggml_backend_buffer_t is a POINTER. -// same logic as for ggml_backend_buffer_type_t - -static inline void apir_encode_ggml_buffer(apir_encoder * enc, const ggml_backend_buffer_t buffer) { - apir_buffer_host_handle_t handle = BUFFER_TO_HOST_HANDLE(buffer); - apir_encoder_write(enc, sizeof(handle), &handle, sizeof(handle)); -} - -static inline ggml_backend_buffer_t apir_decode_ggml_buffer(apir_decoder * dec) { - ggml_backend_buffer_t buffer; - size_t buffer_ptr_size = sizeof(buffer); - - apir_decoder_read(dec, buffer_ptr_size, &buffer, buffer_ptr_size); - - // SECURITY: Validate buffer handle against tracked buffers to prevent - // guest VM from providing arbitrary host memory addresses - if (buffer) { - extern std::unordered_set backend_buffers; - if (backend_buffers.find(buffer) == backend_buffers.end()) { - GGML_LOG_WARN("ggml-virtgpu-backend: %s: Invalid buffer handle from guest: %p\n", __func__, - (void *) buffer); - // Set fatal flag to prevent further processing with invalid handle - apir_decoder_set_fatal(dec); - return NULL; - } - } - - return buffer; -} - -/* enum ggml_status */ - -static inline void apir_encode_ggml_status(apir_encoder * enc, const ggml_status * status) { - apir_encoder_write(enc, sizeof(*status), status, sizeof(*status)); -} - -static inline void apir_decode_ggml_status(apir_decoder * dec, ggml_status * status) { - apir_decoder_read(dec, sizeof(*status), status, sizeof(*status)); -} - -/* virtgpu_shmem */ - -static inline void apir_encode_virtgpu_shmem_res_id(apir_encoder * enc, uint32_t shmem_res_id) { - apir_encode_uint32_t(enc, &shmem_res_id); -} - -static inline void apir_decode_virtgpu_shmem_res_id(apir_decoder * dec, uint32_t * shmem_res_id) { - apir_decode_uint32_t(dec, shmem_res_id); -} - -/* ggml_cgraph */ - -static inline size_t apir_serialize_ggml_cgraph(ggml_cgraph * cgraph, std::vector & cgraph_data) { - apir_serialize_graph(cgraph, cgraph_data); - - return cgraph_data.size(); -} - -static inline void apir_encode_cgraph_data(apir_encoder * enc, std::vector & cgraph_data) { - size_t cgraph_size = cgraph_data.size(); - - apir_encode(enc, cgraph_size, cgraph_data.data(), cgraph_size); -} - -static inline ggml_cgraph * apir_decode_ggml_cgraph(apir_decoder * dec, size_t cgraph_size) { - GGML_UNUSED(cgraph_size); - - uint32_t n_nodes; - apir_decode_uint32_t(dec, &n_nodes); - const uint64_t * nodes = apir_decode_uint64_t_array_inplace(dec, n_nodes); - - uint32_t n_tensors; - apir_decode_uint32_t(dec, &n_tensors); - const apir_rpc_tensor * tensors = apir_decode_apir_rpc_tensor_array_inplace(dec, n_tensors); - - return apir_deserialize_graph(n_nodes, n_tensors, tensors, nodes); -} - -static inline void apir_encode_ggml_buffer_handle(apir_encoder * enc, const apir_buffer_host_handle_t * handle) { - apir_encoder_write(enc, sizeof(*handle), &handle, sizeof(*handle)); -} - -static inline void apir_encode_ggml_tensor_inline(apir_encoder * enc, const ggml_tensor * tensor) { - size_t tensor_size = sizeof(*tensor); - - if (tensor->extra) { - GGML_ABORT("%s: Cannot pass tensors with extra", __func__); - } - - if (tensor->src[0] && tensor->buffer) { - static int first = 1; - if (first) { - GGML_LOG_WARN("%s: Cannot pass tensors with src and buffer\n", __func__); - first = 0; - } - } - - apir_encoder_write(enc, tensor_size, tensor, tensor_size); - - // tensor->data is a pointer inside the device buffer. No need to touch it - // tensor->buffer is a pointer to a buffer. Encoding the buffer handle in sequence. - // (could also make a copy of the tensor, and update locally.) - - if (tensor->buffer) { - apir_buffer_host_handle_t buffer_handle = ggml_buffer_to_apir_handle(tensor->buffer); - apir_encode_ggml_buffer_handle(enc, &buffer_handle); - } - - if (tensor->view_src) { - apir_encoder_write(enc, tensor_size, tensor->view_src, tensor_size); - } - - for (int i = 0; tensor->src[i]; i++) { - const ggml_tensor * tensor_src = tensor->src[i]; - apir_encoder_write(enc, tensor_size, tensor_src, tensor_size); - } -} - -static inline const ggml_tensor * apir_decode_ggml_tensor_inplace(apir_decoder * dec) { - // it safe to remove the `const` qualifier here, we *do* want to - // modify the shared memory data to fix the `src` pointers. - ggml_tensor * tensor = (ggml_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, sizeof(ggml_tensor)); - - // tensor->data is a pointer inside the device buffer. No need to touch it - // tensor->buffer is a pointer to a buffer. Decode the buffer handle encoded in sequence. - if (tensor->buffer) { - tensor->buffer = apir_decode_ggml_buffer(dec); - } - - if (tensor->view_src) { - ggml_tensor * tensor_view_src = (ggml_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, sizeof(ggml_tensor)); - tensor->view_src = tensor_view_src; - } - - for (int i = 0; tensor->src[i]; i++) { - ggml_tensor * tensor_src = (ggml_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, sizeof(ggml_tensor)); - tensor->src[i] = tensor_src; // overwrite op->src[i] pointer with the actual location of the src tensor - } - - return tensor; -} diff --git a/ggml/src/ggml-virtgpu/backend/shared/apir_cs_rpc.h b/ggml/src/ggml-virtgpu/backend/shared/apir_cs_rpc.h deleted file mode 100644 index 4cb2f047d1e5..000000000000 --- a/ggml/src/ggml-virtgpu/backend/shared/apir_cs_rpc.h +++ /dev/null @@ -1,58 +0,0 @@ -#pragma once - -// clang-format off -#include "ggml.h" -#include "ggml-backend-impl.h" - -#include -#include -#include -#include -// clang-format on - -// ggml_tensor is serialized into apir_rpc_tensor -struct apir_rpc_tensor { - uint64_t id; - uint32_t type; - uint64_t buffer; - uint32_t ne[GGML_MAX_DIMS]; - uint32_t nb[GGML_MAX_DIMS]; - uint32_t op; - int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; - int32_t flags; - uint64_t src[GGML_MAX_SRC]; - uint64_t view_src; - uint64_t view_offs; - uint64_t data; - char name[GGML_MAX_NAME]; - - char padding[4]; -}; - -/* frontend */ - -apir_rpc_tensor apir_serialize_tensor(const ggml_tensor * tensor); - -void apir_serialize_graph(const ggml_cgraph * cgraph, std::vector & output); - -/* backend */ - -void apir_track_backend_buffer(ggml_backend_buffer_t buffer); -bool apir_untrack_backend_buffer(ggml_backend_buffer_t buffer); -std::unordered_set apir_get_track_backend_buffers(); - -void apir_add_tensor(ggml_tensor * tensor, - std::vector & tensors, - std::unordered_set & visited); - -ggml_tensor * apir_deserialize_tensor(ggml_context * ctx, const apir_rpc_tensor * tensor); - -ggml_tensor * apir_create_node(uint64_t id, - ggml_context * ctx, - const std::unordered_map & tensor_ptrs, - std::unordered_map & tensor_map); - -ggml_cgraph * apir_deserialize_graph(uint32_t n_nodes, - uint32_t n_tensors, - const apir_rpc_tensor * tensors, - const uint64_t * nodes); diff --git a/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp b/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp deleted file mode 100644 index 8fa20ff43bd5..000000000000 --- a/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp +++ /dev/null @@ -1,81 +0,0 @@ -#include "ggml-remoting.h" - -static ggml_backend_buffer_t ggml_backend_remoting_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, - size_t size) { - virtgpu * gpu = BUFT_TO_GPU(buft); - - ggml_backend_remoting_buffer_context * context = (ggml_backend_remoting_buffer_context *) malloc(sizeof(*context)); - if (!context) { - GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate the buffer context ...", __func__); - } - - context->gpu = gpu; - - bool async__unused, host_buffer__unused, events__unused; - bool buffer_from_host_ptr; - apir_device_get_props(gpu, &async__unused, &host_buffer__unused, &buffer_from_host_ptr, &events__unused); - - if (buffer_from_host_ptr) { - context->apir_context = apir_device_buffer_from_ptr(gpu, size, size); - context->base = context->apir_context.shmem.mmap_ptr; - context->is_from_ptr = true; - } else { - context->apir_context = apir_buffer_type_alloc_buffer(gpu, gpu->cached_buffer_type.host_handle, size); - context->is_from_ptr = false; - context->base = NULL; - } - - ggml_backend_buffer_t buffer = - ggml_backend_buffer_init(buft, ggml_backend_remoting_buffer_interface, (void *) context, size); - - return buffer; -} - -static const char * ggml_backend_remoting_buffer_type_get_name(ggml_backend_buffer_type_t buft) { - virtgpu * gpu = BUFT_TO_GPU(buft); - - // Return the prefixed name that was built once during initialization - return gpu->cached_buffer_type.name; -} - -static size_t ggml_backend_remoting_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { - virtgpu * gpu = BUFT_TO_GPU(buft); - - return gpu->cached_buffer_type.alignment; -} - -static size_t ggml_backend_remoting_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { - virtgpu * gpu = BUFT_TO_GPU(buft); - - return gpu->cached_buffer_type.max_size; -} - -static size_t ggml_backend_remoting_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, - const ggml_tensor * tensor) { - virtgpu * gpu = BUFT_TO_GPU(buft); - - if (tensor->buffer == NULL || !tensor->buffer->context || - !buft->device->iface.supports_buft(buft->device, tensor->buffer->buft)) { - return ggml_nbytes(tensor); - } - - return apir_buffer_type_get_alloc_size(gpu, gpu->cached_buffer_type.host_handle, tensor); -} - -const ggml_backend_buffer_type_i ggml_backend_remoting_buffer_type_interface = { - /* .get_name = */ ggml_backend_remoting_buffer_type_get_name, - /* .alloc_buffer = */ ggml_backend_remoting_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_remoting_buffer_type_get_alignment, - /* .get_max_size = */ ggml_backend_remoting_buffer_type_get_max_size, - /* .get_alloc_size = */ ggml_backend_remoting_buffer_type_get_alloc_size, - /* .is_host = */ NULL, -}; - -const ggml_backend_buffer_type_i ggml_backend_remoting_buffer_from_ptr_type_interface = { - /* .get_name = */ ggml_backend_remoting_buffer_type_get_name, - /* .alloc_buffer = */ NULL, - /* .get_alignment = */ ggml_backend_remoting_buffer_type_get_alignment, - /* .get_max_size = */ ggml_backend_remoting_buffer_type_get_max_size, - /* .get_alloc_size = */ ggml_backend_remoting_buffer_type_get_alloc_size, - /* .is_host = */ NULL, -}; diff --git a/ggml/src/ggml-virtgpu/ggml-backend-buffer.cpp b/ggml/src/ggml-virtgpu/ggml-backend-buffer.cpp deleted file mode 100644 index b6c561cd61ee..000000000000 --- a/ggml/src/ggml-virtgpu/ggml-backend-buffer.cpp +++ /dev/null @@ -1,123 +0,0 @@ -#include "ggml-remoting.h" - -#define BUFFER_TO_GPU(name) ((ggml_backend_remoting_buffer_context *) (name)->context)->gpu - -static void * ggml_backend_remoting_buffer_get_base(ggml_backend_buffer_t buffer) { - ggml_backend_remoting_buffer_context * context = (ggml_backend_remoting_buffer_context *) buffer->context; - if (context->base) { - return context->base; - } - - context->base = apir_buffer_get_base(BUFFER_TO_GPU(buffer), BUFFER_TO_APIR_CONTEXT(buffer)); - - return context->base; -} - -static void ggml_backend_remoting_buffer_set_tensor(ggml_backend_buffer_t buffer, - ggml_tensor * tensor, - const void * data, - size_t offset, - size_t size) { - virtgpu * gpu = BUFFER_TO_GPU(buffer); - - ggml_backend_remoting_buffer_context * context = BUFFER_TO_GGML_CONTEXT(buffer); - if (context->is_from_ptr) { - memcpy((char *) tensor->data + offset, data, size); - } else { - apir_buffer_set_tensor(gpu, BUFFER_TO_APIR_CONTEXT(buffer), tensor, data, offset, size); - } - - return; -} - -static void ggml_backend_remoting_buffer_get_tensor(ggml_backend_buffer_t buffer, - const ggml_tensor * tensor, - void * data, - size_t offset, - size_t size) { - virtgpu * gpu = BUFFER_TO_GPU(buffer); - ggml_backend_remoting_buffer_context * context = BUFFER_TO_GGML_CONTEXT(buffer); - if (context->is_from_ptr) { - memcpy(data, (const char *) tensor->data + offset, size); - } else { - apir_buffer_get_tensor(gpu, BUFFER_TO_APIR_CONTEXT(buffer), tensor, data, offset, size); - } -} - -static void ggml_backend_remoting_buffer_set_tensor_from_ptr(ggml_backend_buffer_t buffer, - ggml_tensor * tensor, - const void * data, - size_t offset, - size_t size) { - UNUSED(buffer); - - memcpy((char *) tensor->data + offset, data, size); - - return; -} - -static void ggml_backend_remoting_buffer_get_tensor_from_ptr(ggml_backend_buffer_t buffer, - const ggml_tensor * tensor, - void * data, - size_t offset, - size_t size) { - UNUSED(buffer); - - memcpy(data, (const char *) tensor->data + offset, size); -} - -static bool ggml_backend_remoting_buffer_cpy_tensor(ggml_backend_buffer_t buffer, - const ggml_tensor * src, - ggml_tensor * dst) { - virtgpu * gpu = BUFFER_TO_GPU(buffer); - - bool ret = apir_buffer_cpy_tensor(gpu, BUFFER_TO_APIR_CONTEXT(buffer), src, dst); - - return ret; -} - -static void ggml_backend_remoting_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { - virtgpu * gpu = BUFFER_TO_GPU(buffer); - - apir_buffer_clear(gpu, BUFFER_TO_APIR_CONTEXT(buffer), value); - - return; -} - -static void ggml_backend_remoting_buffer_free_buffer(ggml_backend_buffer_t buffer) { - virtgpu * gpu = BUFFER_TO_GPU(buffer); - - apir_buffer_free_buffer(gpu, BUFFER_TO_APIR_CONTEXT(buffer)); - - ggml_backend_remoting_buffer_context * context = BUFFER_TO_GGML_CONTEXT(buffer); - free(context); - buffer->context = NULL; -} - -const ggml_backend_buffer_i ggml_backend_remoting_buffer_interface = { - /* .free_buffer = */ ggml_backend_remoting_buffer_free_buffer, - /* .get_base = */ ggml_backend_remoting_buffer_get_base, - /* .init_tensor = */ NULL, - /* .memset_tensor = */ NULL, - /* .set_tensor = */ ggml_backend_remoting_buffer_set_tensor, - /* .get_tensor = */ ggml_backend_remoting_buffer_get_tensor, - /* .set_tensor_2d = */ NULL, - /* .get_tensor_2d = */ NULL, - /* .cpy_tensor = */ ggml_backend_remoting_buffer_cpy_tensor, - /* .clear = */ ggml_backend_remoting_buffer_clear, - /* .reset = */ NULL, -}; - -const ggml_backend_buffer_i ggml_backend_remoting_buffer_from_ptr_interface = { - /* .free_buffer = */ ggml_backend_remoting_buffer_free_buffer, - /* .get_base = */ ggml_backend_remoting_buffer_get_base, - /* .init_tensor = */ NULL, - /* .memset_tensor = */ NULL, - /* .set_tensor = */ ggml_backend_remoting_buffer_set_tensor_from_ptr, - /* .get_tensor = */ ggml_backend_remoting_buffer_get_tensor_from_ptr, - /* .set_tensor_2d = */ NULL, - /* .get_tensor_2d = */ NULL, - /* .cpy_tensor = */ ggml_backend_remoting_buffer_cpy_tensor, - /* .clear = */ ggml_backend_remoting_buffer_clear, - /* .reset = */ NULL, -}; diff --git a/ggml/src/ggml-virtgpu/ggml-backend-device.cpp b/ggml/src/ggml-virtgpu/ggml-backend-device.cpp deleted file mode 100644 index ec8156bb868d..000000000000 --- a/ggml/src/ggml-virtgpu/ggml-backend-device.cpp +++ /dev/null @@ -1,158 +0,0 @@ -#include "ggml-remoting.h" - -static const char * ggml_backend_remoting_device_get_name(ggml_backend_dev_t dev) { - virtgpu * gpu = DEV_TO_GPU(dev); - - // Return the prefixed name that was built once during initialization - return gpu->cached_device_info.name; -} - -static const char * ggml_backend_remoting_device_get_description(ggml_backend_dev_t dev) { - virtgpu * gpu = DEV_TO_GPU(dev); - - // Return the pre-cached description from the virtgpu structure - return gpu->cached_device_info.description; -} - -static enum ggml_backend_dev_type ggml_backend_remoting_device_get_type(ggml_backend_dev_t dev) { - virtgpu * gpu = DEV_TO_GPU(dev); - - return (enum ggml_backend_dev_type) gpu->cached_device_info.type; -} - -static void ggml_backend_remoting_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { - virtgpu * gpu = DEV_TO_GPU(dev); - - *free = gpu->cached_device_info.memory_free; - *total = gpu->cached_device_info.memory_total; -} - -static bool ggml_backend_remoting_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { -#if USE_ALWAYS_TRUE_SUPPORTS_OP == 1 - /* ggml-rpc cheats it like this */ - /* with the current implementation of serialize_tensor, the src/view aren't properly passed */ - UNUSED(dev); - UNUSED(op); - - return true; -#else - virtgpu * gpu = DEV_TO_GPU(dev); - - return apir_device_supports_op(gpu, op); -#endif -} - -static bool ggml_backend_remoting_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { - bool supported = buft->device == dev; - - return supported; -} - -static bool ggml_backend_remoting_device_offload_op(ggml_backend_dev_t dev, const ggml_tensor * op) { - UNUSED(dev); - UNUSED(op); - - return false; -} - -static void ggml_backend_remoting_device_get_props(ggml_backend_dev_t dev, ggml_backend_dev_props * props) { - props->name = ggml_backend_remoting_device_get_name(dev); - props->description = ggml_backend_remoting_device_get_description(dev); - props->type = ggml_backend_remoting_device_get_type(dev); - ggml_backend_remoting_device_get_memory(dev, &props->memory_free, &props->memory_total); - - virtgpu * gpu = DEV_TO_GPU(dev); - apir_device_get_props(gpu, &props->caps.async, &props->caps.host_buffer, &props->caps.buffer_from_host_ptr, - &props->caps.events); - - props->caps.buffer_from_host_ptr = false; - props->caps.async = false; - props->caps.events = false; -} - -ggml_backend_buffer_type_t ggml_backend_remoting_device_get_buffer_type(ggml_backend_dev_t dev) { - virtgpu * gpu = DEV_TO_GPU(dev); - - static std::atomic initialized = false; - static ggml_backend_buffer_type buft; - - if (!initialized) { - static std::mutex mutex; - std::lock_guard lock(mutex); - - if (!initialized) { - buft = { - /* .iface = */ ggml_backend_remoting_buffer_type_interface, - /* .device = */ dev, - /* .context = */ (void *) gpu->cached_buffer_type.host_handle, - }; - initialized = true; - } - } - - return &buft; -} - -static ggml_backend_buffer_type_t ggml_backend_remoting_device_get_buffer_from_ptr_type(ggml_backend_dev_t dev) { - virtgpu * gpu = DEV_TO_GPU(dev); - - static std::atomic initialized = false; - static ggml_backend_buffer_type buft; - - if (!initialized) { - static std::mutex mutex; - std::lock_guard lock(mutex); - - if (!initialized) { - buft = { - /* .iface = */ ggml_backend_remoting_buffer_from_ptr_type_interface, - /* .device = */ dev, - /* .context = */ (void *) gpu->cached_buffer_type.host_handle, - }; - initialized = true; - } - } - - return &buft; -} - -static ggml_backend_buffer_t ggml_backend_remoting_device_buffer_from_ptr(ggml_backend_dev_t dev, - void * ptr, - size_t size, - size_t max_tensor_size) { - virtgpu * gpu = DEV_TO_GPU(dev); - - ggml_backend_remoting_buffer_context * context = (ggml_backend_remoting_buffer_context *) malloc(sizeof(*context)); - if (!context) { - GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate the buffer context ...", __func__); - } - - context->gpu = gpu; - context->apir_context = apir_device_buffer_from_ptr(gpu, size, max_tensor_size); - context->base = ptr; - context->is_from_ptr = true; - - ggml_backend_buffer_t buffer = - ggml_backend_buffer_init(ggml_backend_remoting_device_get_buffer_from_ptr_type(dev), - ggml_backend_remoting_buffer_from_ptr_interface, (void *) context, size); - - return buffer; -} - -const ggml_backend_device_i ggml_backend_remoting_device_interface = { - /* .get_name = */ ggml_backend_remoting_device_get_name, - /* .get_description = */ ggml_backend_remoting_device_get_description, - /* .get_memory = */ ggml_backend_remoting_device_get_memory, - /* .get_type = */ ggml_backend_remoting_device_get_type, - /* .get_props = */ ggml_backend_remoting_device_get_props, - /* .init_backend = */ ggml_backend_remoting_device_init, - /* .get_buffer_type = */ ggml_backend_remoting_device_get_buffer_type, - /* .get_host_buffer_type = */ NULL, - /* .buffer_from_host_ptr = */ ggml_backend_remoting_device_buffer_from_ptr, - /* .supports_op = */ ggml_backend_remoting_device_supports_op, - /* .supports_buft = */ ggml_backend_remoting_device_supports_buft, - /* .offload_op = */ ggml_backend_remoting_device_offload_op, - /* .event_new = */ NULL, - /* .event_free = */ NULL, - /* .event_synchronize = */ NULL, -}; diff --git a/ggml/src/ggml-virtgpu/ggml-backend-reg.cpp b/ggml/src/ggml-virtgpu/ggml-backend-reg.cpp deleted file mode 100644 index a4df5956aa39..000000000000 --- a/ggml/src/ggml-virtgpu/ggml-backend-reg.cpp +++ /dev/null @@ -1,213 +0,0 @@ -#include "ggml-remoting.h" -#include "ggml-virtgpu.h" - -#include -#include - -void ggml_virtgpu_cleanup(virtgpu * gpu); - -static virtgpu * apir_initialize() { - static virtgpu * gpu = NULL; - static std::atomic initialized = false; - - if (initialized) { - // fast track - return gpu; - } - - { - static std::mutex mutex; - std::lock_guard lock(mutex); - - if (initialized) { - // thread safe - return gpu; - } - - gpu = create_virtgpu(); - if (!gpu) { - initialized = true; - return NULL; - } - - // Pre-fetch and cache all device information, it will not change - gpu->cached_device_info.description = apir_device_get_description(gpu); - if (!gpu->cached_device_info.description) { - GGML_ABORT(GGML_VIRTGPU "%s: failed to initialize the virtgpu device description", __func__); - } - gpu->cached_device_info.device_count = apir_device_get_count(gpu); - gpu->cached_device_info.type = apir_device_get_type(gpu); - - { - // Get the remote name and create prefixed version - char * rmt_device_name = apir_device_get_name(gpu); - if (!rmt_device_name) { - GGML_ABORT(GGML_VIRTGPU "%s: failed to get the virtgpu device name", __func__); - } - - size_t device_name_len = strlen(rmt_device_name) + 11; // "[virtgpu] " + null terminator - gpu->cached_device_info.name = (char *) malloc(device_name_len); - if (!gpu->cached_device_info.name) { - free(rmt_device_name); - GGML_ABORT(GGML_VIRTGPU "%s: failed to allocate memory for prefixed device name", __func__); - } - snprintf(gpu->cached_device_info.name, device_name_len, "[virtgpu] %s", rmt_device_name); - free(rmt_device_name); - } - - apir_device_get_memory(gpu, &gpu->cached_device_info.memory_free, &gpu->cached_device_info.memory_total); - - apir_buffer_type_host_handle_t buft_host_handle = apir_device_get_buffer_type(gpu); - gpu->cached_buffer_type.host_handle = buft_host_handle; - { - // Get the remote name and create prefixed version - char * rmt_name = apir_buffer_type_get_name(gpu, buft_host_handle); - if (!rmt_name) { - GGML_ABORT(GGML_VIRTGPU "%s: failed to get the virtgpu buffer type name", __func__); - } - - size_t prefixed_len = strlen(rmt_name) + 11; // "[virtgpu] " + null terminator - gpu->cached_buffer_type.name = (char *) malloc(prefixed_len); - if (!gpu->cached_buffer_type.name) { - free(rmt_name); - GGML_ABORT(GGML_VIRTGPU "%s: failed to allocate memory for prefixed buffer type name", __func__); - } - snprintf(gpu->cached_buffer_type.name, prefixed_len, "[virtgpu] %s", rmt_name); - free(rmt_name); - } - - gpu->cached_buffer_type.alignment = apir_buffer_type_get_alignment(gpu, buft_host_handle); - gpu->cached_buffer_type.max_size = apir_buffer_type_get_max_size(gpu, buft_host_handle); - - initialized = true; - } - - return gpu; -} - -static int ggml_backend_remoting_get_device_count() { - virtgpu * gpu = apir_initialize(); - if (!gpu) { - return 0; - } - - return gpu->cached_device_info.device_count; -} - -static size_t ggml_backend_remoting_reg_get_device_count(ggml_backend_reg_t reg) { - UNUSED(reg); - - return ggml_backend_remoting_get_device_count(); -} - -static std::vector devices; - -ggml_backend_dev_t ggml_backend_remoting_get_device(size_t device) { - GGML_ASSERT(device < devices.size()); - return devices[device]; -} - -static void ggml_backend_remoting_reg_init_devices(ggml_backend_reg_t reg) { - if (devices.size() > 0) { - GGML_LOG_INFO(GGML_VIRTGPU "%s: already initialized\n", __func__); - return; - } - - virtgpu * gpu = apir_initialize(); - if (!gpu) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: apir_initialize failed\n", __func__); - return; - } - - static std::atomic initialized = false; - - if (initialized) { - return; // fast track - } - - { - static std::mutex mutex; - std::lock_guard lock(mutex); - if (!initialized) { - for (int i = 0; i < ggml_backend_remoting_get_device_count(); i++) { - ggml_backend_remoting_device_context * ctx = new ggml_backend_remoting_device_context; - char desc[256] = "ggml-virtgpu API Remoting device"; - - ctx->device = i; - ctx->name = GGML_VIRTGPU_NAME + std::to_string(i); - ctx->description = desc; - ctx->gpu = gpu; - - ggml_backend_dev_t dev = new ggml_backend_device{ - /* .iface = */ ggml_backend_remoting_device_interface, - /* .reg = */ reg, - /* .context = */ ctx, - }; - devices.push_back(dev); - } - initialized = true; - } - } -} - -static ggml_backend_dev_t ggml_backend_remoting_reg_get_device(ggml_backend_reg_t reg, size_t device) { - UNUSED(reg); - - return ggml_backend_remoting_get_device(device); -} - -static const char * ggml_backend_remoting_reg_get_name(ggml_backend_reg_t reg) { - UNUSED(reg); - - return GGML_VIRTGPU_NAME; -} - -static const ggml_backend_reg_i ggml_backend_remoting_reg_i = { - /* .get_name = */ ggml_backend_remoting_reg_get_name, - /* .get_device_count = */ ggml_backend_remoting_reg_get_device_count, - /* .get_device = */ ggml_backend_remoting_reg_get_device, - /* .get_proc_address = */ NULL, -}; - -ggml_backend_reg_t ggml_backend_virtgpu_reg() { - virtgpu * gpu = apir_initialize(); - if (!gpu) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: virtgpu_apir_initialize failed\n", __func__); - } - - static ggml_backend_reg reg = { - /* .api_version = */ GGML_BACKEND_API_VERSION, - /* .iface = */ ggml_backend_remoting_reg_i, - /* .context = */ gpu, - }; - - static bool initialized = false; - if (initialized) { - return ® - } - initialized = true; - - ggml_backend_remoting_reg_init_devices(®); - - return ® -} - -// public function, not exposed in the GGML interface at the moment -void ggml_virtgpu_cleanup(virtgpu * gpu) { - if (gpu->cached_device_info.name) { - free(gpu->cached_device_info.name); - gpu->cached_device_info.name = NULL; - } - if (gpu->cached_device_info.description) { - free(gpu->cached_device_info.description); - gpu->cached_device_info.description = NULL; - } - if (gpu->cached_buffer_type.name) { - free(gpu->cached_buffer_type.name); - gpu->cached_buffer_type.name = NULL; - } - - mtx_destroy(&gpu->data_shmem_mutex); -} - -GGML_BACKEND_DL_IMPL(ggml_backend_virtgpu_reg) diff --git a/ggml/src/ggml-virtgpu/ggml-backend.cpp b/ggml/src/ggml-virtgpu/ggml-backend.cpp deleted file mode 100644 index 12756c9282f7..000000000000 --- a/ggml/src/ggml-virtgpu/ggml-backend.cpp +++ /dev/null @@ -1,71 +0,0 @@ -#include "../../include/ggml-virtgpu.h" -#include "ggml-remoting.h" - -static const char * ggml_backend_remoting_get_name(ggml_backend_t backend) { - UNUSED(backend); - - return "API Remoting backend"; -} - -static void ggml_backend_remoting_free(ggml_backend_t backend) { - delete backend; -} - -static ggml_status ggml_backend_remoting_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { - virtgpu * gpu = DEV_TO_GPU(backend->device); - - return apir_backend_graph_compute(gpu, cgraph); -} - -static void ggml_backend_remoting_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { - virtgpu * gpu = DEV_TO_GPU(backend->device); -#if true - UNUSED(gpu); - UNUSED(cgraph); -#else - // not working yet - - apir_backend_graph_optimize(gpu, cgraph); -#endif -} - -static ggml_backend_i ggml_backend_remoting_interface = { - /* .get_name = */ ggml_backend_remoting_get_name, - /* .free = */ ggml_backend_remoting_free, - /* .set_tensor_async = */ NULL, // ggml_backend_remoting_set_tensor_async, - /* .get_tensor_async = */ NULL, // ggml_backend_remoting_get_tensor_async, - /* .set_tensor_2d_async = */ NULL, - /* .get_tensor_2d_async = */ NULL, - /* .cpy_tensor_async = */ NULL, // ggml_backend_remoting_cpy_tensor_async, - /* .synchronize = */ NULL, // ggml_backend_remoting_synchronize, - /* .graph_plan_create = */ NULL, - /* .graph_plan_free = */ NULL, - /* .graph_plan_update = */ NULL, - /* .graph_plan_compute = */ NULL, - /* .graph_compute = */ ggml_backend_remoting_graph_compute, - /* .event_record = */ NULL, - /* .event_wait = */ NULL, - /* .graph_optimize = */ ggml_backend_remoting_graph_optimize, -}; - -static ggml_guid_t ggml_backend_remoting_guid() { - static ggml_guid guid = { 0xb8, 0xf7, 0x4f, 0x86, 0x14, 0x03, 0x86, 0x02, - 0x91, 0xc8, 0xdd, 0xe9, 0x02, 0x3f, 0xc0, 0x2b }; - - return &guid; -} - -ggml_backend_t ggml_backend_remoting_device_init(ggml_backend_dev_t dev, const char * params) { - UNUSED(params); - - ggml_backend_remoting_device_context * ctx = (ggml_backend_remoting_device_context *) dev->context; - - ggml_backend_t remoting_backend = new ggml_backend{ - /* .guid = */ ggml_backend_remoting_guid(), - /* .interface = */ ggml_backend_remoting_interface, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_virtgpu_reg(), ctx->device), - /* .context = */ ctx, - }; - - return remoting_backend; -} diff --git a/ggml/src/ggml-virtgpu/ggml-remoting.h b/ggml/src/ggml-virtgpu/ggml-remoting.h deleted file mode 100644 index 4f70326bee2c..000000000000 --- a/ggml/src/ggml-virtgpu/ggml-remoting.h +++ /dev/null @@ -1,71 +0,0 @@ -#pragma once - -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "ggml-impl.h" -#include "virtgpu.h" - -#include -#include - -#define GGML_VIRTGPU_NAME "ggml-virtgpu" -#define GGML_VIRTGPU "ggml-virtgpu: " - -// USE_ALWAYS_TRUE_SUPPORTS_OP: 1 is fast, 0 avoid micro-benchmark crashes - -#define USE_ALWAYS_TRUE_SUPPORTS_OP 1 -#define USE_METAL_GUEST_SUPPORTS_OP 0 - -#define DEV_TO_GPU(name) ((ggml_backend_remoting_device_context *) (name)->context)->gpu - -#define BUFFER_TO_GGML_CONTEXT(name) ((ggml_backend_remoting_buffer_context *) (name)->context) - -#define BUFFER_TO_APIR_CONTEXT(name) &((ggml_backend_remoting_buffer_context *) (name)->context)->apir_context - -#define BUFFER_TO_HOST_HANDLE(name) ((ggml_backend_remoting_buffer_context *) (name)->context)->apir_context.host_handle - -#define GET_DEVICE_CONTEXT() (ggml_backend_remoting_device_context *) ggml_backend_remoting_get_device(0)->context - -#define BUFT_TO_GPU(name) ((ggml_backend_remoting_device_context *) (name)->device->context)->gpu - -struct ggml_backend_remoting_device_context { - size_t device; - std::string name; - std::string description; - - std::vector> shared_memory; - - virtgpu * gpu; -}; - -struct ggml_backend_remoting_buffer_context { - apir_buffer_context_t apir_context; - - virtgpu * gpu; - - void * base; - - bool is_from_ptr; -}; - -extern const ggml_backend_buffer_type_i ggml_backend_remoting_buffer_type_interface; -extern const ggml_backend_device_i ggml_backend_remoting_device_interface; -extern const ggml_backend_buffer_i ggml_backend_remoting_buffer_interface; -extern const ggml_backend_buffer_type_i ggml_backend_remoting_buffer_from_ptr_type_interface; -extern const ggml_backend_buffer_i ggml_backend_remoting_buffer_from_ptr_interface; - -ggml_backend_dev_t ggml_backend_remoting_get_device(size_t device); -ggml_backend_t ggml_backend_remoting_device_init(ggml_backend_dev_t dev, const char * params); -ggml_backend_buffer_type_t ggml_backend_remoting_device_get_buffer_type(ggml_backend_dev_t dev); - -static inline apir_buffer_type_host_handle_t ggml_buffer_type_to_apir_handle(ggml_backend_buffer_type_t buft) { - // in the backend, the buffer handle is the buffer pointer - return (apir_buffer_type_host_handle_t) buft->context; -} - -static inline apir_buffer_host_handle_t ggml_buffer_to_apir_handle(ggml_backend_buffer_t buffer) { - if (!buffer->context) { - GGML_ABORT(GGML_VIRTGPU "%s: no context available :/", __func__); - } - return BUFFER_TO_HOST_HANDLE(buffer); -} diff --git a/ggml/src/ggml-virtgpu/include/apir_hw.h b/ggml/src/ggml-virtgpu/include/apir_hw.h deleted file mode 100644 index 7d6ea2265db8..000000000000 --- a/ggml/src/ggml-virtgpu/include/apir_hw.h +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include - -struct virgl_renderer_capset_apir { - uint32_t apir_version; - uint32_t supports_blob_resources; - uint32_t reserved[4]; // For future expansion -}; diff --git a/ggml/src/ggml-virtgpu/regenerate_remoting.py b/ggml/src/ggml-virtgpu/regenerate_remoting.py index dae75fd1c80f..7dce57131d47 100755 --- a/ggml/src/ggml-virtgpu/regenerate_remoting.py +++ b/ggml/src/ggml-virtgpu/regenerate_remoting.py @@ -268,46 +268,6 @@ def regenerate_codebase(self) -> None: base_path = self.config_data.get('base_path', 'ggml/src') frontend_base = Path(base_path) / "ggml-virtgpu" - # Compute final file paths - backend_base = frontend_base / "backend" - apir_backend_path = backend_base / "shared" / "apir_backend.gen.h" - backend_dispatched_path = backend_base / "backend-dispatched.gen.h" - virtgpu_forward_path = frontend_base / "virtgpu-forward.gen.h" - - # Create output directories for each file - apir_backend_path.parent.mkdir(parents=True, exist_ok=True) - backend_dispatched_path.parent.mkdir(parents=True, exist_ok=True) - virtgpu_forward_path.parent.mkdir(parents=True, exist_ok=True) - - # Generate header files - logging.info("📁 Generating header files...") - - apir_backend_content = self.generate_apir_backend_header() - apir_backend_path.write_text(apir_backend_content) - logging.info(f" ✅ {apir_backend_path.resolve()}") - - backend_dispatched_content = self.generate_backend_dispatched_header() - backend_dispatched_path.write_text(backend_dispatched_content) - logging.info(f" ✅ {backend_dispatched_path.resolve()}") - - virtgpu_forward_content = self.generate_virtgpu_forward_header() - virtgpu_forward_path.write_text(virtgpu_forward_content) - logging.info(f" ✅ {virtgpu_forward_path.resolve()}") - - # Format generated files with clang-format - generated_files = [apir_backend_path, backend_dispatched_path, virtgpu_forward_path] - - if not self.clang_format_available: - logging.warning("\n⚠️clang-format not found in PATH. Generated files will not be formatted.\n" - " Install clang-format to enable automatic code formatting.") - else: - logging.info("\n🎨 Formatting files with clang-format...") - for file_path in generated_files: - if self._format_file_with_clang_format(file_path): - logging.info(f" ✅ Formatted {file_path.name}") - else: - logging.warning(f" ❌ Failed to format {file_path.name}") - # Generate summary functions = self.get_enabled_functions() total_functions = len(functions) @@ -316,7 +276,7 @@ def regenerate_codebase(self) -> None: logging.info("=" * 50) logging.info(f" Total functions: {total_functions}") logging.info(f" Function groups: {len(self.functions)}") - logging.info(" Header files: 3") + logging.info(" Header files: 0") logging.info(f" Working directory: {current_dir}") diff --git a/ggml/src/ggml-virtgpu/virtgpu-apir.h b/ggml/src/ggml-virtgpu/virtgpu-apir.h deleted file mode 100644 index 238f960acd23..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu-apir.h +++ /dev/null @@ -1,15 +0,0 @@ -#include "backend/shared/apir_backend.h" -#include "ggml-alloc.h" -#include "ggml-impl.h" -#include "ggml.h" -#include "virtgpu-shm.h" -#include "virtgpu-utils.h" - -struct apir_buffer_context_t { - apir_buffer_host_handle_t host_handle; - - struct virtgpu_shmem shmem; - apir_buffer_type_host_handle_t buft_host_handle; -}; - -#include "virtgpu-forward.gen.h" diff --git a/ggml/src/ggml-virtgpu/virtgpu-forward-backend.cpp b/ggml/src/ggml-virtgpu/virtgpu-forward-backend.cpp deleted file mode 100644 index 4593690c6385..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu-forward-backend.cpp +++ /dev/null @@ -1,58 +0,0 @@ -#include "virtgpu-forward-impl.h" - -static long long current_time_ms() { - timespec ts; - clock_gettime(CLOCK_REALTIME, &ts); // Use CLOCK_MONOTONIC for elapsed time - return (long long) ts.tv_sec * 1000000000LL + ts.tv_nsec; -} - -ggml_status apir_backend_graph_compute(virtgpu * gpu, ggml_cgraph * cgraph) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BACKEND_GRAPH_COMPUTE); - - std::vector cgraph_data; - size_t cgraph_size = apir_serialize_ggml_cgraph(cgraph, cgraph_data); - - virtgpu_shmem temp_shmem; // Local storage for large buffers - virtgpu_shmem * shmem = &temp_shmem; - bool using_shared_shmem = false; - - if (cgraph_size <= gpu->data_shmem.mmap_size) { - // Lock mutex before using shared data_shmem buffer - if (mtx_lock(&gpu->data_shmem_mutex) != thrd_success) { - GGML_ABORT(GGML_VIRTGPU "%s: Failed to lock data_shmem mutex", __func__); - } - using_shared_shmem = true; - shmem = &gpu->data_shmem; - } else if (virtgpu_shmem_create(gpu, cgraph_size, shmem)) { - GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate the guest-host shared buffer", __func__); - } - - apir_encode_virtgpu_shmem_res_id(encoder, shmem->res_id); - - apir_encode_size_t(encoder, &cgraph_size); - - char * shmem_data = (char *) shmem->mmap_ptr; - apir_encoder secondary_enc = apir_new_encoder(shmem_data, cgraph_size); - - apir_encode_cgraph_data(&secondary_enc, cgraph_data); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - ggml_status status = GGML_STATUS_ABORTED; - apir_decode_ggml_status(decoder, &status); - - remote_call_finish(gpu, encoder, decoder); - - // Unlock mutex before cleanup - if (using_shared_shmem) { - mtx_unlock(&gpu->data_shmem_mutex); - } else { - virtgpu_shmem_destroy(gpu, shmem); - } - - return status; -} diff --git a/ggml/src/ggml-virtgpu/virtgpu-forward-buffer-type.cpp b/ggml/src/ggml-virtgpu/virtgpu-forward-buffer-type.cpp deleted file mode 100644 index 38f8ec945e06..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu-forward-buffer-type.cpp +++ /dev/null @@ -1,110 +0,0 @@ -#include "virtgpu-forward-impl.h" - -char * apir_buffer_type_get_name(virtgpu * gpu, apir_buffer_type_host_handle_t host_handle) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_TYPE_GET_NAME); - - apir_encode_apir_buffer_type_host_handle(encoder, host_handle); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - const size_t string_size = apir_decode_array_size_unchecked(decoder); - char * string = (char *) apir_decoder_alloc_array(sizeof(char), string_size); - if (!string) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: Could not allocate the device name buffer\n", __func__); - apir_decoder_set_fatal(decoder); - } - apir_decode_char_array(decoder, string, string_size); - - remote_call_finish(gpu, encoder, decoder); - - return string; -} - -size_t apir_buffer_type_get_alignment(virtgpu * gpu, apir_buffer_type_host_handle_t host_handle) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALIGNMENT); - - apir_encode_apir_buffer_type_host_handle(encoder, host_handle); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - size_t alignment; - apir_decode_size_t(decoder, &alignment); - - remote_call_finish(gpu, encoder, decoder); - - return alignment; -} - -size_t apir_buffer_type_get_max_size(virtgpu * gpu, apir_buffer_type_host_handle_t host_handle) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_TYPE_GET_MAX_SIZE); - - apir_encode_apir_buffer_type_host_handle(encoder, host_handle); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - size_t max_size; - apir_decode_size_t(decoder, &max_size); - - remote_call_finish(gpu, encoder, decoder); - - return max_size; -} - -apir_buffer_context_t apir_buffer_type_alloc_buffer(virtgpu * gpu, - apir_buffer_type_host_handle_t host_handle, - size_t size) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - apir_buffer_context_t buffer_context; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_TYPE_ALLOC_BUFFER); - - apir_encode_apir_buffer_type_host_handle(encoder, host_handle); - - apir_encode_size_t(encoder, &size); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - apir_decode_apir_buffer_host_handle_t(decoder, &buffer_context.host_handle); - - remote_call_finish(gpu, encoder, decoder); - - return buffer_context; -} - -size_t apir_buffer_type_get_alloc_size(virtgpu * gpu, - apir_buffer_type_host_handle_t host_handle, - const ggml_tensor * op) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALLOC_SIZE); - - apir_encode_apir_buffer_type_host_handle(encoder, host_handle); - - apir_encode_ggml_tensor_inline(encoder, op); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - size_t alloc_size; - apir_decode_size_t(decoder, &alloc_size); - - remote_call_finish(gpu, encoder, decoder); - - return alloc_size; -} diff --git a/ggml/src/ggml-virtgpu/virtgpu-forward-buffer.cpp b/ggml/src/ggml-virtgpu/virtgpu-forward-buffer.cpp deleted file mode 100644 index 228284f4a42a..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu-forward-buffer.cpp +++ /dev/null @@ -1,173 +0,0 @@ -#include "virtgpu-forward-impl.h" - -void * apir_buffer_get_base(virtgpu * gpu, apir_buffer_context_t * buffer_context) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_GET_BASE); - - apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - uintptr_t base; - apir_decode_uintptr_t(decoder, &base); - - remote_call_finish(gpu, encoder, decoder); - - return (void *) base; -} - -void apir_buffer_set_tensor(virtgpu * gpu, - apir_buffer_context_t * buffer_context, - ggml_tensor * tensor, - const void * data, - size_t offset, - size_t size) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_SET_TENSOR); - - apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle); - apir_encode_ggml_tensor(encoder, tensor); - - virtgpu_shmem temp_shmem; // Local storage for large buffers - virtgpu_shmem * shmem = &temp_shmem; - bool using_shared_shmem = false; - - if (size <= gpu->data_shmem.mmap_size) { - // Lock mutex before using shared data_shmem buffer - if (mtx_lock(&gpu->data_shmem_mutex) != thrd_success) { - GGML_ABORT(GGML_VIRTGPU "%s: Failed to lock data_shmem mutex", __func__); - } - using_shared_shmem = true; - shmem = &gpu->data_shmem; - - } else if (virtgpu_shmem_create(gpu, size, shmem)) { - GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate the guest-host shared buffer", __func__); - } - - memcpy(shmem->mmap_ptr, data, size); - apir_encode_virtgpu_shmem_res_id(encoder, shmem->res_id); - - apir_encode_size_t(encoder, &offset); - apir_encode_size_t(encoder, &size); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - remote_call_finish(gpu, encoder, decoder); - - // Unlock mutex before cleanup - if (using_shared_shmem) { - mtx_unlock(&gpu->data_shmem_mutex); - } else { - virtgpu_shmem_destroy(gpu, shmem); - } - - return; -} - -void apir_buffer_get_tensor(virtgpu * gpu, - apir_buffer_context_t * buffer_context, - const ggml_tensor * tensor, - void * data, - size_t offset, - size_t size) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_GET_TENSOR); - - apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle); - apir_encode_ggml_tensor(encoder, tensor); - - virtgpu_shmem temp_shmem; // Local storage for large buffers - virtgpu_shmem * shmem = &temp_shmem; - bool using_shared_shmem = false; - - if (size <= gpu->data_shmem.mmap_size) { - // Lock mutex before using shared data_shmem buffer - if (mtx_lock(&gpu->data_shmem_mutex) != thrd_success) { - GGML_ABORT(GGML_VIRTGPU "%s: Failed to lock data_shmem mutex", __func__); - } - using_shared_shmem = true; - shmem = &gpu->data_shmem; - - } else if (virtgpu_shmem_create(gpu, size, shmem)) { - GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate the guest-host shared buffer", __func__); - } - - apir_encode_virtgpu_shmem_res_id(encoder, shmem->res_id); - apir_encode_size_t(encoder, &offset); - apir_encode_size_t(encoder, &size); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - memcpy(data, shmem->mmap_ptr, size); - - remote_call_finish(gpu, encoder, decoder); - - // Unlock mutex before cleanup - if (using_shared_shmem) { - mtx_unlock(&gpu->data_shmem_mutex); - } else { - virtgpu_shmem_destroy(gpu, shmem); - } -} - -bool apir_buffer_cpy_tensor(virtgpu * gpu, - apir_buffer_context_t * buffer_context, - const ggml_tensor * src, - const ggml_tensor * dst) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_CPY_TENSOR); - - apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle); - apir_encode_ggml_tensor(encoder, src); - apir_encode_ggml_tensor(encoder, dst); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - bool ret_val; - apir_decode_bool_t(decoder, &ret_val); - - remote_call_finish(gpu, encoder, decoder); - - return ret_val; -} - -void apir_buffer_clear(virtgpu * gpu, apir_buffer_context_t * buffer_context, uint8_t value) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_CLEAR); - - apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle); - apir_encode_uint8_t(encoder, &value); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - remote_call_finish(gpu, encoder, decoder); -} - -void apir_buffer_free_buffer(virtgpu * gpu, apir_buffer_context_t * buffer_context) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_FREE_BUFFER); - - apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - remote_call_finish(gpu, encoder, decoder); -} diff --git a/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp b/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp deleted file mode 100644 index 9f513c138dd2..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp +++ /dev/null @@ -1,192 +0,0 @@ -#include "virtgpu-forward-impl.h" -#include "virtgpu-shm.h" - -int apir_device_get_count(virtgpu * gpu) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_COUNT); - REMOTE_CALL(gpu, encoder, decoder, ret); - - int32_t dev_count = -1; - apir_decode_int32_t(decoder, &dev_count); - - remote_call_finish(gpu, encoder, decoder); - - return dev_count; -} - -char * apir_device_get_name(virtgpu * gpu) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_NAME); - REMOTE_CALL(gpu, encoder, decoder, ret); - - const size_t string_size = apir_decode_array_size_unchecked(decoder); - char * string = (char *) apir_decoder_alloc_array(sizeof(char), string_size); - if (!string) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: Could not allocate the device name buffer\n", __func__); - return NULL; - } - apir_decode_char_array(decoder, string, string_size); - - remote_call_finish(gpu, encoder, decoder); - - return string; -} - -char * apir_device_get_description(virtgpu * gpu) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_DESCRIPTION); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - const size_t string_size = apir_decode_array_size_unchecked(decoder); - char * string = (char *) apir_decoder_alloc_array(sizeof(char), string_size); - if (!string) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: Could not allocate the device description buffer\n", __func__); - - return NULL; - } - apir_decode_char_array(decoder, string, string_size); - - remote_call_finish(gpu, encoder, decoder); - - return string; -} - -uint32_t apir_device_get_type(virtgpu * gpu) { - static uint32_t dev_type = 255; - if (dev_type != 255) { - return dev_type; - } - - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_TYPE); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - apir_decode_uint32_t(decoder, &dev_type); - - remote_call_finish(gpu, encoder, decoder); - - return dev_type; -} - -void apir_device_get_memory(virtgpu * gpu, size_t * free, size_t * total) { - static size_t dev_free = 0; - static size_t dev_total = 0; - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_MEMORY); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - apir_decode_size_t(decoder, &dev_free); - apir_decode_size_t(decoder, &dev_total); - - *free = dev_free; - *total = dev_total; - - remote_call_finish(gpu, encoder, decoder); - - return; -} - -bool apir_device_supports_op(virtgpu * gpu, const ggml_tensor * op) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_SUPPORTS_OP); - - apir_encode_ggml_tensor_inline(encoder, op); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - bool supports_op; - apir_decode_bool_t(decoder, &supports_op); - - remote_call_finish(gpu, encoder, decoder); - - return supports_op; -} - -apir_buffer_type_host_handle_t apir_device_get_buffer_type(virtgpu * gpu) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_BUFFER_TYPE); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - apir_buffer_type_host_handle_t buft_handle; - apir_decode_apir_buffer_type_host_handle_t(decoder, &buft_handle); - - remote_call_finish(gpu, encoder, decoder); - - return buft_handle; -} - -void apir_device_get_props(virtgpu * gpu, - bool * async, - bool * host_buffer, - bool * buffer_from_host_ptr, - bool * events) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_PROPS); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - apir_decode_bool_t(decoder, async); - apir_decode_bool_t(decoder, host_buffer); - apir_decode_bool_t(decoder, buffer_from_host_ptr); - apir_decode_bool_t(decoder, events); - - remote_call_finish(gpu, encoder, decoder); - - return; -} - -apir_buffer_context_t apir_device_buffer_from_ptr(virtgpu * gpu, size_t size, size_t max_tensor_size) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirForwardReturnCode ret; - - apir_buffer_context_t buffer_context; - - REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_BUFFER_FROM_PTR); - - if (virtgpu_shmem_create(gpu, size, &buffer_context.shmem)) { - GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate %ldb of guest-host shared buffer", __func__, size); - } - - apir_encode_virtgpu_shmem_res_id(encoder, buffer_context.shmem.res_id); - - apir_encode_size_t(encoder, &size); - apir_encode_size_t(encoder, &max_tensor_size); - - REMOTE_CALL(gpu, encoder, decoder, ret); - - apir_decode_apir_buffer_host_handle_t(decoder, &buffer_context.host_handle); - buffer_context.buft_host_handle = apir_decode_apir_buffer_type_host_handle(decoder); - - remote_call_finish(gpu, encoder, decoder); - - return buffer_context; -} diff --git a/ggml/src/ggml-virtgpu/virtgpu-forward-impl.h b/ggml/src/ggml-virtgpu/virtgpu-forward-impl.h deleted file mode 100644 index 4d0b6e05c740..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu-forward-impl.h +++ /dev/null @@ -1,36 +0,0 @@ -#pragma once - -// clang-format off -#include "virtgpu.h" -#include "ggml-remoting.h" -#include "backend/shared/apir_backend.h" -#include "backend/shared/apir_cs_ggml.h" -#include "ggml-backend-impl.h" -// clang-format on - -#define REMOTE_CALL_PREPARE(gpu_dev_name, encoder_name, apir_command_type__) \ - int32_t REMOTE_CALL_PREPARE_forward_flag = (int32_t) apir_command_type__; \ - const char * REMOTE_CALL_PREPARE_command_name = apir_dispatch_command_name(apir_command_type__); \ - do { \ - encoder_name = remote_call_prepare(gpu_dev_name, APIR_COMMAND_TYPE_FORWARD, REMOTE_CALL_PREPARE_forward_flag); \ - if (!encoder_name) { \ - GGML_ABORT(GGML_VIRTGPU "%s: failed to prepare the remote call encoder", __func__); \ - } \ - } while (0) - -#define REMOTE_CALL(gpu_dev_name, encoder_name, decoder_name, ret_name) \ - do { \ - ret_name = (ApirForwardReturnCode) remote_call(gpu_dev_name, encoder_name, &decoder_name, 0, NULL); \ - if (!decoder_name) { \ - GGML_ABORT(GGML_VIRTGPU "%s: failed to kick the remote call", __func__); \ - } \ - if (ret_name < APIR_FORWARD_BASE_INDEX) { \ - GGML_ABORT(GGML_VIRTGPU "%s: failed to forward the API call: %s: code %d", __func__, \ - apir_forward_error(ret_name), ret_name); \ - } \ - ret_name = (ApirForwardReturnCode) (ret_name - APIR_FORWARD_BASE_INDEX); \ - if (ret_name != 0) { \ - GGML_ABORT(GGML_VIRTGPU "backend function '%s' failed (return code: %d)", \ - REMOTE_CALL_PREPARE_command_name, ret_name); \ - } \ - } while (0) diff --git a/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h b/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h deleted file mode 100644 index 44b0ad1ffa1d..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h +++ /dev/null @@ -1,53 +0,0 @@ -#pragma once - -/* device */ -void apir_device_get_device_count(struct virtgpu * gpu); -int apir_device_get_count(struct virtgpu * gpu); -char * apir_device_get_name(struct virtgpu * gpu); -char * apir_device_get_description(struct virtgpu * gpu); -uint32_t apir_device_get_type(struct virtgpu * gpu); -void apir_device_get_memory(struct virtgpu * gpu, size_t * free, size_t * total); -bool apir_device_supports_op(struct virtgpu * gpu, const ggml_tensor * op); -apir_buffer_type_host_handle_t apir_device_get_buffer_type(struct virtgpu * gpu); -void apir_device_get_props(struct virtgpu * gpu, - bool * async, - bool * host_buffer, - bool * buffer_from_host_ptr, - bool * events); -apir_buffer_context_t apir_device_buffer_from_ptr(struct virtgpu * gpu, size_t size, size_t max_tensor_size); - -/* buffer-type */ -char * apir_buffer_type_get_name(struct virtgpu * gpu, apir_buffer_type_host_handle_t host_handle); -size_t apir_buffer_type_get_alignment(struct virtgpu * gpu, apir_buffer_type_host_handle_t host_handle); -size_t apir_buffer_type_get_max_size(struct virtgpu * gpu, apir_buffer_type_host_handle_t host_handle); -/* apir_buffer_type_is_host is deprecated. */ -apir_buffer_context_t apir_buffer_type_alloc_buffer(struct virtgpu * gpu, - apir_buffer_type_host_handle_t host_handle, - size_t size); -size_t apir_buffer_type_get_alloc_size(struct virtgpu * gpu, - apir_buffer_type_host_handle_t host_handle, - const ggml_tensor * op); - -/* buffer */ -void * apir_buffer_get_base(struct virtgpu * gpu, apir_buffer_context_t * buffer_context); -void apir_buffer_set_tensor(struct virtgpu * gpu, - apir_buffer_context_t * buffer_context, - ggml_tensor * tensor, - const void * data, - size_t offset, - size_t size); -void apir_buffer_get_tensor(struct virtgpu * gpu, - apir_buffer_context_t * buffer_context, - const ggml_tensor * tensor, - void * data, - size_t offset, - size_t size); -bool apir_buffer_cpy_tensor(struct virtgpu * gpu, - apir_buffer_context_t * buffer_context, - const ggml_tensor * src, - const ggml_tensor * dst); -void apir_buffer_clear(struct virtgpu * gpu, apir_buffer_context_t * buffer_context, uint8_t value); -void apir_buffer_free_buffer(struct virtgpu * gpu, apir_buffer_context_t * buffer_context); - -/* backend */ -ggml_status apir_backend_graph_compute(struct virtgpu * gpu, ggml_cgraph * cgraph); diff --git a/ggml/src/ggml-virtgpu/virtgpu-shm.cpp b/ggml/src/ggml-virtgpu/virtgpu-shm.cpp deleted file mode 100644 index 7f2c2322d919..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu-shm.cpp +++ /dev/null @@ -1,99 +0,0 @@ -#include "virtgpu-shm.h" - -#include "virtgpu.h" -#include "ggml-remoting.h" - -#include - -static uint32_t virtgpu_ioctl_resource_create_blob(virtgpu * gpu, - uint32_t blob_mem, - uint32_t blob_flags, - size_t blob_size, - uint64_t blob_id, - uint32_t * res_id) { -#ifdef SIMULATE_BO_SIZE_FIX - blob_size = align64(blob_size, 4096); -#endif - - drm_virtgpu_resource_create_blob args = { - .blob_mem = blob_mem, - .blob_flags = blob_flags, - .bo_handle = 0, - .res_handle = 0, - .size = blob_size, - .pad = 0, - .cmd_size = 0, - .cmd = 0, - .blob_id = blob_id, - }; - - if (virtgpu_ioctl(gpu, DRM_IOCTL_VIRTGPU_RESOURCE_CREATE_BLOB, &args)) { - return 0; - } - - *res_id = args.res_handle; - return args.bo_handle; -} - -static void virtgpu_ioctl_gem_close(virtgpu * gpu, uint32_t gem_handle) { - drm_gem_close args = { - .handle = gem_handle, - .pad = 0, - }; - - const int ret = virtgpu_ioctl(gpu, DRM_IOCTL_GEM_CLOSE, &args); - assert(!ret); -#ifdef NDEBUG - UNUSED(ret); -#endif -} - -static void * virtgpu_ioctl_map(virtgpu * gpu, uint32_t gem_handle, size_t size) { - drm_virtgpu_map args = { - .offset = 0, - .handle = gem_handle, - .pad = 0, - }; - - if (virtgpu_ioctl(gpu, DRM_IOCTL_VIRTGPU_MAP, &args)) { - return NULL; - } - - void * ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, gpu->fd, args.offset); - if (ptr == MAP_FAILED) { - return NULL; - } - - return ptr; -} - -void virtgpu_shmem_destroy(virtgpu * gpu, virtgpu_shmem * shmem) { - munmap(shmem->mmap_ptr, shmem->mmap_size); - virtgpu_ioctl_gem_close(gpu, shmem->gem_handle); -} - -int virtgpu_shmem_create(virtgpu * gpu, size_t size, virtgpu_shmem * shmem) { - size = align64(size, 16384); - - uint32_t res_id; - uint32_t gem_handle = virtgpu_ioctl_resource_create_blob(gpu, VIRTGPU_BLOB_MEM_HOST3D, - VIRTGPU_BLOB_FLAG_USE_MAPPABLE, size, 0, &res_id); - - if (!gem_handle) { - return 1; - } - - void * ptr = virtgpu_ioctl_map(gpu, gem_handle, size); - if (!ptr) { - virtgpu_ioctl_gem_close(gpu, gem_handle); - GGML_LOG_ERROR(GGML_VIRTGPU "%s: virtgpu_ioctl_map failed\n", __func__); - return 1; - } - - shmem->res_id = res_id; - shmem->mmap_size = size; - shmem->mmap_ptr = ptr; - shmem->gem_handle = gem_handle; - - return 0; -} diff --git a/ggml/src/ggml-virtgpu/virtgpu-shm.h b/ggml/src/ggml-virtgpu/virtgpu-shm.h deleted file mode 100644 index 606860a0946a..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu-shm.h +++ /dev/null @@ -1,23 +0,0 @@ -#pragma once - -#include "virtgpu-utils.h" - -#include - -#include -#include -#include -#include - -struct virtgpu; - -struct virtgpu_shmem { - uint32_t res_id; - size_t mmap_size; - void * mmap_ptr; - - uint32_t gem_handle; -}; - -int virtgpu_shmem_create(virtgpu * gpu, size_t size, virtgpu_shmem * shmem); -void virtgpu_shmem_destroy(virtgpu * gpu, virtgpu_shmem * shmem); diff --git a/ggml/src/ggml-virtgpu/virtgpu-utils.cpp b/ggml/src/ggml-virtgpu/virtgpu-utils.cpp deleted file mode 100644 index 8a2805e99023..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu-utils.cpp +++ /dev/null @@ -1,179 +0,0 @@ -#include "virtgpu-utils.h" - -#include -#include - -#include - -#define NODE_ALLOC_ALIGN 64 -#define NODE_PTR_MASK (~((uintptr_t) NODE_ALLOC_ALIGN - 1)) -#define NODE_LEVEL_MASK ((uintptr_t) NODE_ALLOC_ALIGN - 1) -#define NULL_NODE 0 - -#define os_malloc_aligned(_size, _align) _aligned_malloc(_size, _align) -#define os_free_aligned(_ptr) free(_ptr) -#define p_atomic_cmpxchg(v, old, _new) __sync_val_compare_and_swap((v), (old), (_new)) - -static inline uint64_t util_logbase2_64(uint64_t n) { -#if defined(HAVE___BUILTIN_CLZLL) - return ((sizeof(uint64_t) * 8 - 1) - __builtin_clzll(n | 1)); -#else - uint64_t pos = 0ull; - if (n >= 1ull << 32) { - n >>= 32; - pos += 32; - } - if (n >= 1ull << 16) { - n >>= 16; - pos += 16; - } - if (n >= 1ull << 8) { - n >>= 8; - pos += 8; - } - if (n >= 1ull << 4) { - n >>= 4; - pos += 4; - } - if (n >= 1ull << 2) { - n >>= 2; - pos += 2; - } - if (n >= 1ull << 1) { - pos += 1; - } - return pos; -#endif -} - -void util_sparse_array_init(util_sparse_array * arr, size_t elem_size, size_t node_size) { - memset(arr, 0, sizeof(*arr)); - arr->elem_size = elem_size; - arr->node_size_log2 = util_logbase2_64(node_size); - assert(node_size >= 2 && node_size == (1ull << arr->node_size_log2)); -} - -static inline void * os_malloc_aligned(size_t size, size_t alignment) { - void * ptr; - alignment = (alignment + sizeof(void *) - 1) & ~(sizeof(void *) - 1); - if (posix_memalign(&ptr, alignment, size) != 0) { - return NULL; - } - return ptr; -} - -static inline void * _util_sparse_array_node_data(uintptr_t handle) { - return (void *) (handle & NODE_PTR_MASK); -} - -static inline unsigned _util_sparse_array_node_level(uintptr_t handle) { - return handle & NODE_LEVEL_MASK; -} - -static inline void _util_sparse_array_node_finish(util_sparse_array * arr, uintptr_t node) { - if (_util_sparse_array_node_level(node) > 0) { - uintptr_t * children = (uintptr_t *) _util_sparse_array_node_data(node); - size_t node_size = 1ull << arr->node_size_log2; - for (size_t i = 0; i < node_size; i++) { - if (children[i]) { - _util_sparse_array_node_finish(arr, children[i]); - } - } - } - - os_free_aligned(_util_sparse_array_node_data(node)); -} - -static inline uintptr_t _util_sparse_array_node(void * data, unsigned level) { - assert(data != NULL); - assert(((uintptr_t) data & NODE_LEVEL_MASK) == 0); - assert((level & NODE_PTR_MASK) == 0); - return (uintptr_t) data | level; -} - -inline uintptr_t _util_sparse_array_node_alloc(util_sparse_array * arr, unsigned level) { - size_t size; - if (level == 0) { - size = arr->elem_size << arr->node_size_log2; - } else { - size = sizeof(uintptr_t) << arr->node_size_log2; - } - - void * data = os_malloc_aligned(size, NODE_ALLOC_ALIGN); - memset(data, 0, size); - - return _util_sparse_array_node(data, level); -} - -static inline uintptr_t _util_sparse_array_set_or_free_node(uintptr_t * node_ptr, uintptr_t cmp_node, uintptr_t node) { - uintptr_t prev_node = p_atomic_cmpxchg(node_ptr, cmp_node, node); - - if (prev_node != cmp_node) { - /* We lost the race. Free this one and return the one that was already - * allocated. - */ - os_free_aligned(_util_sparse_array_node_data(node)); - return prev_node; - } else { - return node; - } -} - -void * util_sparse_array_get(util_sparse_array * arr, uint64_t idx) { - const unsigned node_size_log2 = arr->node_size_log2; - uintptr_t root = p_atomic_read(&arr->root); - if (unlikely(!root)) { - unsigned root_level = 0; - uint64_t idx_iter = idx >> node_size_log2; - while (idx_iter) { - idx_iter >>= node_size_log2; - root_level++; - } - uintptr_t new_root = _util_sparse_array_node_alloc(arr, root_level); - root = _util_sparse_array_set_or_free_node(&arr->root, NULL_NODE, new_root); - } - - while (1) { - unsigned root_level = _util_sparse_array_node_level(root); - uint64_t root_idx = idx >> (root_level * node_size_log2); - if (likely(root_idx < (1ull << node_size_log2))) { - break; - } - - /* In this case, we have a root but its level is low enough that the - * requested index is out-of-bounds. - */ - uintptr_t new_root = _util_sparse_array_node_alloc(arr, root_level + 1); - - uintptr_t * new_root_children = (uintptr_t *) _util_sparse_array_node_data(new_root); - new_root_children[0] = root; - - /* We only add one at a time instead of the whole tree because it's - * easier to ensure correctness of both the tree building and the - * clean-up path. Because we're only adding one node we never have to - * worry about trying to free multiple things without freeing the old - * things. - */ - root = _util_sparse_array_set_or_free_node(&arr->root, root, new_root); - } - - void * node_data = _util_sparse_array_node_data(root); - unsigned node_level = _util_sparse_array_node_level(root); - while (node_level > 0) { - uint64_t child_idx = (idx >> (node_level * node_size_log2)) & ((1ull << node_size_log2) - 1); - - uintptr_t * children = (uintptr_t *) node_data; - uintptr_t child = p_atomic_read(&children[child_idx]); - - if (unlikely(!child)) { - child = _util_sparse_array_node_alloc(arr, node_level - 1); - child = _util_sparse_array_set_or_free_node(&children[child_idx], NULL_NODE, child); - } - - node_data = _util_sparse_array_node_data(child); - node_level = _util_sparse_array_node_level(child); - } - - uint64_t elem_idx = idx & ((1ull << node_size_log2) - 1); - return (void *) ((char *) node_data + (elem_idx * arr->elem_size)); -} diff --git a/ggml/src/ggml-virtgpu/virtgpu-utils.h b/ggml/src/ggml-virtgpu/virtgpu-utils.h deleted file mode 100644 index a0036b4e2bc8..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu-utils.h +++ /dev/null @@ -1,86 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#define unlikely(x) __builtin_expect(!!(x), 0) -#define likely(x) __builtin_expect(!!(x), 1) - -#ifndef UNUSED -# define UNUSED(x) (void) (x) -#endif - -/** Checks is a value is a power of two. Does not handle zero. */ -#define IS_POT(v) (((v) & ((v) - 1)) == 0) - -/** Checks is a value is a power of two. Zero handled. */ -#define IS_POT_NONZERO(v) ((v) != 0 && IS_POT(v)) - -/** Align a value to a power of two */ -#define ALIGN_POT(x, pot_align) (((x) + (pot_align) - 1) & ~((pot_align) - 1)) - -#define p_atomic_read(_v) __atomic_load_n((_v), __ATOMIC_ACQUIRE) - -static inline bool util_is_power_of_two_nonzero64(uint64_t v) { - return IS_POT_NONZERO(v); -} - -static inline uint64_t align64(uint64_t value, uint64_t alignment) { - assert(util_is_power_of_two_nonzero64(alignment)); - return ALIGN_POT(value, alignment); -} - -struct list_head { - list_head * prev; - list_head * next; -}; - -struct util_sparse_array { - size_t elem_size; - unsigned node_size_log2; - - uintptr_t root; -}; - -void * util_sparse_array_get(util_sparse_array * arr, uint64_t idx); -void util_sparse_array_init(util_sparse_array * arr, size_t elem_size, size_t node_size); - -inline void os_time_sleep(int64_t usecs) { - timespec time; - time.tv_sec = usecs / 1000000; - time.tv_nsec = (usecs % 1000000) * 1000; - while (clock_nanosleep(CLOCK_MONOTONIC, 0, &time, &time) == EINTR) - ; -} - -struct timer_data { - long long start; - long long total; - long long count; -}; - -static inline void start_timer(timer_data * timer) { - timespec ts; - clock_gettime(CLOCK_MONOTONIC, &ts); - timer->start = (long long) ts.tv_sec * 1000000000LL + ts.tv_nsec; -} - -// returns the duration in ns -static inline long long stop_timer(timer_data * timer) { - timespec ts; - clock_gettime(CLOCK_MONOTONIC, &ts); - long long timer_end = (long long) ts.tv_sec * 1000000000LL + ts.tv_nsec; - - long long duration = (timer_end - timer->start); - timer->total += duration; - timer->count += 1; - - return duration; -} diff --git a/ggml/src/ggml-virtgpu/virtgpu.cpp b/ggml/src/ggml-virtgpu/virtgpu.cpp deleted file mode 100644 index e3ae1cc75e00..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu.cpp +++ /dev/null @@ -1,545 +0,0 @@ -#include "virtgpu.h" -#include "ggml-remoting.h" - -#include -#include - -#include -#include -#include - -static virt_gpu_result_t virtgpu_open_device(virtgpu * gpu, const drmDevicePtr dev); -static virt_gpu_result_t virtgpu_open(virtgpu * gpu); - -static virt_gpu_result_t virtgpu_init_capset(virtgpu * gpu); -static virt_gpu_result_t virtgpu_init_context(virtgpu * gpu); - -static int virtgpu_ioctl_context_init(virtgpu * gpu, virgl_renderer_capset capset_id); -static int virtgpu_ioctl_get_caps(virtgpu * gpu, - virgl_renderer_capset id, - uint32_t version, - void * capset, - size_t capset_size); -static uint64_t virtgpu_ioctl_getparam(virtgpu * gpu, uint64_t param); -static void virtgpu_init_renderer_info(virtgpu * gpu); - -static void log_call_duration(long long call_duration_ns, const char * name); - -const uint64_t APIR_HANDSHAKE_MAX_WAIT_MS = 2 * 1000; // 2s -const uint64_t APIR_LOADLIBRARY_MAX_WAIT_MS = 60 * 1000; // 60s - -static int virtgpu_handshake(virtgpu * gpu) { - apir_encoder * encoder; - apir_decoder * decoder; - - encoder = remote_call_prepare(gpu, APIR_COMMAND_TYPE_HANDSHAKE, 0); - if (!encoder) { - GGML_ABORT(GGML_VIRTGPU "%s: failed to prepare the remote call encoder", __func__); - return 1; - } - - /* write handshake props */ - - uint32_t guest_major = APIR_PROTOCOL_MAJOR; - uint32_t guest_minor = APIR_PROTOCOL_MINOR; - apir_encode_uint32_t(encoder, &guest_major); - apir_encode_uint32_t(encoder, &guest_minor); - - /* *** */ - - uint32_t ret_magic; - long long call_duration_ns; - ret_magic = remote_call(gpu, encoder, &decoder, APIR_HANDSHAKE_MAX_WAIT_MS, &call_duration_ns); - log_call_duration(call_duration_ns, "API Remoting handshake"); - - if (!decoder) { - GGML_ABORT(GGML_VIRTGPU - "%s: failed to initiate the communication with the virglrenderer library. " - "Most likely, the wrong virglrenderer library was loaded in the hypervisor.", - __func__); - return 1; - } - - /* read handshake return values */ - - uint32_t host_major; - uint32_t host_minor; - - if (ret_magic != APIR_HANDSHAKE_MAGIC) { - GGML_ABORT(GGML_VIRTGPU "%s: handshake with the virglrenderer failed (code=%d | %s)", __func__, ret_magic, - apir_backend_initialize_error(ret_magic)); - } else { - apir_decode_uint32_t(decoder, &host_major); - apir_decode_uint32_t(decoder, &host_minor); - } - - remote_call_finish(gpu, encoder, decoder); - - if (ret_magic != APIR_HANDSHAKE_MAGIC) { - return 1; - } - - GGML_LOG_INFO(GGML_VIRTGPU "%s: Guest is running with %u.%u\n", __func__, guest_major, guest_minor); - GGML_LOG_INFO(GGML_VIRTGPU "%s: Host is running with %u.%u\n", __func__, host_major, host_minor); - - if (guest_major != host_major) { - GGML_LOG_ERROR(GGML_VIRTGPU "Host major (%d) and guest major (%d) version differ\n", host_major, guest_major); - } else if (guest_minor != host_minor) { - GGML_LOG_WARN(GGML_VIRTGPU "Host minor (%d) and guest minor (%d) version differ\n", host_minor, guest_minor); - } - - return 0; -} - -static ApirLoadLibraryReturnCode virtgpu_load_library(virtgpu * gpu) { - apir_encoder * encoder; - apir_decoder * decoder; - ApirLoadLibraryReturnCode ret; - - encoder = remote_call_prepare(gpu, APIR_COMMAND_TYPE_LOADLIBRARY, 0); - if (!encoder) { - GGML_ABORT(GGML_VIRTGPU "%s: hypercall error: failed to prepare the API Remoting command encoder", __func__); - return APIR_LOAD_LIBRARY_HYPERCALL_INITIALIZATION_ERROR; - } - - long long call_duration_ns; - - ret = (ApirLoadLibraryReturnCode) remote_call(gpu, encoder, &decoder, APIR_LOADLIBRARY_MAX_WAIT_MS, - &call_duration_ns); - log_call_duration(call_duration_ns, "API Remoting LoadLibrary"); - - if (!decoder) { - GGML_ABORT(GGML_VIRTGPU "%s: hypercall error: failed to trigger the API Remoting hypercall.\n", __func__); - return APIR_LOAD_LIBRARY_HYPERCALL_INITIALIZATION_ERROR; - } - - remote_call_finish(gpu, encoder, decoder); - - if (ret == APIR_LOAD_LIBRARY_SUCCESS) { - GGML_LOG_INFO(GGML_VIRTGPU "The API Remoting backend was successfully loaded and initialized\n"); - - return ret; - } - - // something wrong happened, find out what. - if (ret < APIR_LOAD_LIBRARY_INIT_BASE_INDEX) { - if (ret == APIR_LOAD_LIBRARY_ENV_VAR_MISSING) { - GGML_ABORT(GGML_VIRTGPU - "%s: virglrenderer could not open the API Remoting backend library, " - "some environment variables are missing. " - "Make sure virglrenderer is correctly configured by the hypervisor. (%s)", - __func__, apir_load_library_error(ret)); - } else if (ret == APIR_LOAD_LIBRARY_CANNOT_OPEN) { - GGML_ABORT(GGML_VIRTGPU - "%s: virglrenderer could not open the API Remoting backend library. " - "Make sure virglrenderer is correctly configured by the hypervisor. (%s)", - __func__, apir_load_library_error(ret)); - } else if (ret == APIR_LOAD_LIBRARY_ENV_VAR_MISSING) { - GGML_ABORT(GGML_VIRTGPU - "%s: could not load the backend library, some symbols are missing. " - "Make sure virglrenderer is correctly configured by the hypervisor. (%s) ", - __func__, apir_load_library_error(ret)); - } else { - GGML_ABORT(GGML_VIRTGPU "%s: virglrenderer could not load the API Remoting backend library. (%s - code %d)", - __func__, apir_load_library_error(ret), ret); - } - return ret; - } - - GGML_LOG_INFO(GGML_VIRTGPU "%s: virglrenderer successfully loaded the API Remoting backend library.\n", __func__); - - ApirLoadLibraryReturnCode apir_ret = (ApirLoadLibraryReturnCode) (ret - APIR_LOAD_LIBRARY_INIT_BASE_INDEX); - - if (apir_ret == APIR_LOAD_LIBRARY_CANNOT_OPEN) { - GGML_ABORT(GGML_VIRTGPU - "%s: the API Remoting backend library couldn't load the GGML backend library. " - "Make sure virglrenderer is correctly configured by the hypervisor. (%s)", - __func__, apir_load_library_error(apir_ret)); - } else if (apir_ret == APIR_LOAD_LIBRARY_SYMBOL_MISSING) { - GGML_ABORT( - GGML_VIRTGPU - "%s: the API Remoting backend library couldn't load the GGML backend library, some symbols are missing. " - "Make sure virglrenderer is correctly configured by the hypervisor. (%s)", - __func__, apir_load_library_error(apir_ret)); - } else if (apir_ret < APIR_LOAD_LIBRARY_INIT_BASE_INDEX) { - GGML_ABORT(GGML_VIRTGPU - "%s: the API Remoting backend library couldn't load the GGML backend library: apir code=%d | %s)", - __func__, apir_ret, apir_load_library_error(apir_ret)); - } else { - uint32_t lib_ret = apir_ret - APIR_LOAD_LIBRARY_INIT_BASE_INDEX; - GGML_ABORT(GGML_VIRTGPU - "%s: the API Remoting backend library failed to initialize its backend library: apir code=%d)", - __func__, lib_ret); - } - return ret; -} - -virtgpu * create_virtgpu() { - virtgpu * gpu = new virtgpu(); - - gpu->use_apir_capset = getenv("GGML_REMOTING_USE_APIR_CAPSET") != nullptr; - util_sparse_array_init(&gpu->shmem_array, sizeof(virtgpu_shmem), 1024); - - // Initialize mutex to protect shared data_shmem buffer - if (mtx_init(&gpu->data_shmem_mutex, mtx_plain) != thrd_success) { - delete gpu; - GGML_ABORT(GGML_VIRTGPU "%s: failed to initialize data_shmem mutex", __func__); - return NULL; - } - - if (virtgpu_open(gpu) != APIR_SUCCESS) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: failed to open the virtgpu device\n", __func__); - return NULL; - } - - if (virtgpu_init_capset(gpu) != APIR_SUCCESS) { - if (gpu->use_apir_capset) { - GGML_ABORT(GGML_VIRTGPU - "%s: failed to initialize the virtgpu APIR capset. Make sure that the virglrenderer library " - "supports it.", - __func__); - } else { - GGML_ABORT(GGML_VIRTGPU "%s: failed to initialize the virtgpu Venus capset", __func__); - } - return NULL; - } - - if (virtgpu_init_context(gpu) != APIR_SUCCESS) { - GGML_ABORT(GGML_VIRTGPU "%s: failed to initialize the GPU context", __func__); - return NULL; - } - - if (virtgpu_shmem_create(gpu, SHMEM_REPLY_SIZE, &gpu->reply_shmem)) { - GGML_ABORT(GGML_VIRTGPU "%s: failed to create the shared reply memory pages", __func__); - return NULL; - } - - if (virtgpu_shmem_create(gpu, SHMEM_DATA_SIZE, &gpu->data_shmem)) { - GGML_ABORT(GGML_VIRTGPU "%s: failed to create the shared data memory pages", __func__); - return NULL; - } - - if (virtgpu_handshake(gpu)) { - GGML_ABORT(GGML_VIRTGPU "%s: failed to handshake with the virglrenderer library", __func__); - return NULL; - } - - if (virtgpu_load_library(gpu) != APIR_LOAD_LIBRARY_SUCCESS) { - GGML_ABORT(GGML_VIRTGPU "%s: failed to load the backend library", __func__); - return NULL; - } - - return gpu; -} - -static virt_gpu_result_t virtgpu_open(virtgpu * gpu) { - drmDevicePtr devs[8]; - int count = drmGetDevices2(0, devs, ARRAY_SIZE(devs)); - if (count < 0) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: failed to enumerate DRM devices\n", __func__); - return APIR_ERROR_INITIALIZATION_FAILED; - } - - virt_gpu_result_t result = APIR_ERROR_INITIALIZATION_FAILED; - for (int i = 0; i < count; i++) { - result = virtgpu_open_device(gpu, devs[i]); - if (result == APIR_SUCCESS) { - break; - } - } - - drmFreeDevices(devs, count); - - return result; -} - -static virt_gpu_result_t virtgpu_open_device(virtgpu * gpu, const drmDevicePtr dev) { - const char * node_path = dev->nodes[DRM_NODE_RENDER]; - - int fd = open(node_path, O_RDWR | O_CLOEXEC); - if (fd < 0) { - GGML_ABORT(GGML_VIRTGPU "%s: failed to open %s", __func__, node_path); - return APIR_ERROR_INITIALIZATION_FAILED; - } - - drmVersionPtr version = drmGetVersion(fd); - if (!version || strcmp(version->name, "virtio_gpu") || version->version_major != 0) { - if (version) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: unknown DRM driver %s version %d\n", __func__, version->name, - version->version_major); - } else { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: failed to get DRM driver version\n", __func__); - } - - if (version) { - drmFreeVersion(version); - } - close(fd); - return APIR_ERROR_INITIALIZATION_FAILED; - } - - gpu->fd = fd; - - drmFreeVersion(version); - - GGML_LOG_INFO(GGML_VIRTGPU "using DRM device %s\n", node_path); - - return APIR_SUCCESS; -} - -static virt_gpu_result_t virtgpu_init_context(virtgpu * gpu) { - assert(!gpu->capset.version); - const int ret = virtgpu_ioctl_context_init(gpu, gpu->capset.id); - if (ret) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: failed to initialize context: %s\n", __func__, strerror(errno)); - return APIR_ERROR_INITIALIZATION_FAILED; - } - - return APIR_SUCCESS; -} - -static virt_gpu_result_t virtgpu_init_capset(virtgpu * gpu) { - if (gpu->use_apir_capset) { - GGML_LOG_INFO(GGML_VIRTGPU "Using the APIR capset\n"); - gpu->capset.id = VIRTGPU_DRM_CAPSET_APIR; - } else { - GGML_LOG_INFO(GGML_VIRTGPU "Using the Venus capset\n"); - gpu->capset.id = VIRTGPU_DRM_CAPSET_VENUS; - } - gpu->capset.version = 0; - - int ret = - virtgpu_ioctl_get_caps(gpu, gpu->capset.id, gpu->capset.version, &gpu->capset.data, sizeof(gpu->capset.data)); - - if (ret) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: failed to get APIR v%d capset: %s\n", __func__, gpu->capset.version, - strerror(errno)); - return APIR_ERROR_INITIALIZATION_FAILED; - } - - assert(gpu->capset.data.supports_blob_resources); - - return APIR_SUCCESS; -} - -static int virtgpu_ioctl_context_init(virtgpu * gpu, virgl_renderer_capset capset_id) { - drm_virtgpu_context_set_param ctx_set_params[3] = { - { - .param = VIRTGPU_CONTEXT_PARAM_CAPSET_ID, - .value = capset_id, - }, - { - .param = VIRTGPU_CONTEXT_PARAM_NUM_RINGS, - .value = 1, - }, - { - .param = VIRTGPU_CONTEXT_PARAM_POLL_RINGS_MASK, - .value = 0, /* don't generate drm_events on fence signaling */ - }, - }; - - drm_virtgpu_context_init args = { - .num_params = ARRAY_SIZE(ctx_set_params), - .pad = 0, - .ctx_set_params = (uintptr_t) &ctx_set_params, - }; - - return virtgpu_ioctl(gpu, DRM_IOCTL_VIRTGPU_CONTEXT_INIT, &args); -} - -static int virtgpu_ioctl_get_caps(virtgpu * gpu, - virgl_renderer_capset id, - uint32_t version, - void * capset, - size_t capset_size) { - drm_virtgpu_get_caps args = { - .cap_set_id = id, - .cap_set_ver = version, - .addr = (uintptr_t) capset, - .size = (__u32) capset_size, - .pad = 0, - }; - - return virtgpu_ioctl(gpu, DRM_IOCTL_VIRTGPU_GET_CAPS, &args); -} - -static uint64_t virtgpu_ioctl_getparam(virtgpu * gpu, uint64_t param) { - /* val must be zeroed because kernel only writes the lower 32 bits */ - uint64_t val = 0; - drm_virtgpu_getparam args = { - .param = param, - .value = (uintptr_t) &val, - }; - - const int ret = virtgpu_ioctl(gpu, DRM_IOCTL_VIRTGPU_GETPARAM, &args); - return ret ? 0 : val; -} - -apir_encoder * remote_call_prepare(virtgpu * gpu, ApirCommandType apir_cmd_type, int32_t cmd_flags) { - /* - * Prepare the command encoder and its buffer - */ - - thread_local char encoder_buffer[4096]; - - thread_local apir_encoder enc; - enc = { - .cur = encoder_buffer, - .start = encoder_buffer, - .end = encoder_buffer + sizeof(encoder_buffer), - .fatal = false, - }; - - /* - * Fill the command encoder with the common args: - * - cmd_type (int32_t) - * - cmd_flags (int32_t) - * - reply res id (uint32_t) - */ - - int32_t cmd_type = apir_cmd_type; - - // for testing during the hypervisor transition - if (!gpu->use_apir_capset) { - cmd_type += VENUS_COMMAND_TYPE_LENGTH; - } - apir_encode_int32_t(&enc, &cmd_type); - apir_encode_int32_t(&enc, &cmd_flags); - - uint32_t reply_res_id = gpu->reply_shmem.res_id; - apir_encode_uint32_t(&enc, &reply_res_id); - - return &enc; -} - -void remote_call_finish(virtgpu * gpu, apir_encoder * enc, apir_decoder * dec) { - UNUSED(gpu); - - if (!enc) { - GGML_ABORT(GGML_VIRTGPU "%s: Invalid (null) encoder", __func__); - } - - if (!dec) { - GGML_ABORT(GGML_VIRTGPU "%s: Invalid (null) decoder", __func__); - } - - if (apir_encoder_get_fatal(enc)) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: Failed to encode the output parameters.", __func__); - } - - if (apir_decoder_get_fatal(dec)) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: Failed to decode the input parameters.", __func__); - } -} - -uint32_t remote_call(virtgpu * gpu, - apir_encoder * encoder, - apir_decoder ** decoder, - float max_wait_ms, - long long * call_duration_ns) { - /* - * Prepare the reply notification pointer - */ - - volatile std::atomic_uint * atomic_reply_notif = (volatile std::atomic_uint *) gpu->reply_shmem.mmap_ptr; - *atomic_reply_notif = 0; - - /* - * Trigger the execbuf ioctl - */ - - drm_virtgpu_execbuffer args = { - .flags = VIRTGPU_EXECBUF_RING_IDX, - .size = (uint32_t) (encoder->cur - encoder->start), - .command = (uintptr_t) encoder->start, - - .bo_handles = 0, - .num_bo_handles = 0, - - .fence_fd = 0, - .ring_idx = 0, - .syncobj_stride = 0, - .num_in_syncobjs = 0, - .num_out_syncobjs = 0, - .in_syncobjs = 0, - .out_syncobjs = 0, - }; - - *decoder = NULL; - - int ret = drmIoctl(gpu->fd, DRM_IOCTL_VIRTGPU_EXECBUFFER, &args); - - if (ret != 0) { - GGML_ABORT(GGML_VIRTGPU "%s: the virtgpu EXECBUFFER ioctl failed (%d)", __func__, ret); - } - - /* - * Wait for the response notification - */ - timer_data wait_host_reply_timer = { 0, 0, 0 }; - - start_timer(&wait_host_reply_timer); - - timespec ts_start, ts_end; - clock_gettime(CLOCK_MONOTONIC, &ts_start); - long long start_time = (long long) ts_start.tv_sec * 1000000000LL + ts_start.tv_nsec; - - bool timedout = false; - uint32_t notif_value = 0; - while (true) { - notif_value = std::atomic_load_explicit(atomic_reply_notif, std::memory_order_acquire); - - if (notif_value != 0) { - break; - } - - int64_t base_sleep_us = 15; - - os_time_sleep(base_sleep_us); - - if (max_wait_ms) { - clock_gettime(CLOCK_MONOTONIC, &ts_end); - long long end_time = (long long) ts_end.tv_sec * 1000000000LL + ts_end.tv_nsec; - float duration_ms = (end_time - start_time) / 1000000; - - if (duration_ms > max_wait_ms) { - timedout = true; - break; - } - } - } - - if (call_duration_ns) { - *call_duration_ns = stop_timer(&wait_host_reply_timer); - } - - if (max_wait_ms && timedout) { - GGML_LOG_ERROR(GGML_VIRTGPU "%s: timed out waiting for the host answer...\n", __func__); - return APIR_FORWARD_TIMEOUT; - } - - /* - * Prepare the decoder - */ - static apir_decoder response_dec; - response_dec.cur = (char *) gpu->reply_shmem.mmap_ptr + sizeof(*atomic_reply_notif); - response_dec.end = (char *) gpu->reply_shmem.mmap_ptr + gpu->reply_shmem.mmap_size; - *decoder = &response_dec; - - // extract the actual return value from the notif flag - uint32_t returned_value = notif_value - 1; - return returned_value; -} - -static void log_call_duration(long long call_duration_ns, const char * name) { - double call_duration_ms = (double) call_duration_ns / 1e6; // 1 millisecond = 1e6 nanoseconds - double call_duration_s = (double) call_duration_ns / 1e9; // 1 second = 1e9 nanoseconds - - if (call_duration_s > 1) { - GGML_LOG_INFO(GGML_VIRTGPU "waited %.2fs for the %s host reply...\n", call_duration_s, name); - } else if (call_duration_ms > 1) { - GGML_LOG_INFO(GGML_VIRTGPU "waited %.2fms for the %s host reply...\n", call_duration_ms, name); - } else { - GGML_LOG_INFO(GGML_VIRTGPU "waited %lldns for the %s host reply...\n", call_duration_ns, name); - } -} diff --git a/ggml/src/ggml-virtgpu/virtgpu.cpp.inc b/ggml/src/ggml-virtgpu/virtgpu.cpp.inc new file mode 100644 index 000000000000..7685f9835f75 --- /dev/null +++ b/ggml/src/ggml-virtgpu/virtgpu.cpp.inc @@ -0,0 +1,3376 @@ +// clang-format off +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define unlikely(x) __builtin_expect(!!(x), 0) +#define likely(x) __builtin_expect(!!(x), 1) + +#ifndef UNUSED +# define UNUSED(x) (void) (x) +#endif + +/** Checks is a value is a power of two. Does not handle zero. */ +#define IS_POT(v) (((v) & ((v) - 1)) == 0) + +/** Checks is a value is a power of two. Zero handled. */ +#define IS_POT_NONZERO(v) ((v) != 0 && IS_POT(v)) + +/** Align a value to a power of two */ +#define ALIGN_POT(x, pot_align) (((x) + (pot_align) - 1) & ~((pot_align) - 1)) + +#define p_atomic_read(_v) __atomic_load_n((_v), __ATOMIC_ACQUIRE) + +static inline bool util_is_power_of_two_nonzero64(uint64_t v) { + return IS_POT_NONZERO(v); +} + +static inline uint64_t align64(uint64_t value, uint64_t alignment) { + assert(util_is_power_of_two_nonzero64(alignment)); + return ALIGN_POT(value, alignment); +} + +struct list_head { + list_head * prev; + list_head * next; +}; + +struct util_sparse_array { + size_t elem_size; + unsigned node_size_log2; + + uintptr_t root; +}; + +void * util_sparse_array_get(util_sparse_array * arr, uint64_t idx); +void util_sparse_array_init(util_sparse_array * arr, size_t elem_size, size_t node_size); + +inline void os_time_sleep(int64_t usecs) { + timespec time; + time.tv_sec = usecs / 1000000; + time.tv_nsec = (usecs % 1000000) * 1000; + while (clock_nanosleep(CLOCK_MONOTONIC, 0, &time, &time) == EINTR) + ; +} + +struct timer_data { + long long start; + long long total; + long long count; +}; + +static inline void start_timer(timer_data * timer) { + timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + timer->start = (long long) ts.tv_sec * 1000000000LL + ts.tv_nsec; +} + +// returns the duration in ns +static inline long long stop_timer(timer_data * timer) { + timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + long long timer_end = (long long) ts.tv_sec * 1000000000LL + ts.tv_nsec; + + long long duration = (timer_end - timer->start); + timer->total += duration; + timer->count += 1; + + return duration; +} + +struct virtgpu; + +struct virtgpu_shmem { + uint32_t res_id; + size_t mmap_size; + void * mmap_ptr; + + uint32_t gem_handle; +}; + +int virtgpu_shmem_create(virtgpu * gpu, size_t size, virtgpu_shmem * shmem); +void virtgpu_shmem_destroy(virtgpu * gpu, virtgpu_shmem * shmem); + +typedef enum ApirBackendCommandType { + + /* device */ + APIR_COMMAND_TYPE_DEVICE_GET_DEVICE_COUNT = 0, + APIR_COMMAND_TYPE_DEVICE_GET_COUNT = 1, + APIR_COMMAND_TYPE_DEVICE_GET_NAME = 2, + APIR_COMMAND_TYPE_DEVICE_GET_DESCRIPTION = 3, + APIR_COMMAND_TYPE_DEVICE_GET_TYPE = 4, + APIR_COMMAND_TYPE_DEVICE_GET_MEMORY = 5, + APIR_COMMAND_TYPE_DEVICE_SUPPORTS_OP = 6, + APIR_COMMAND_TYPE_DEVICE_GET_BUFFER_TYPE = 7, + APIR_COMMAND_TYPE_DEVICE_GET_PROPS = 8, + APIR_COMMAND_TYPE_DEVICE_BUFFER_FROM_PTR = 9, + + /* buffer-type */ + APIR_COMMAND_TYPE_BUFFER_TYPE_GET_NAME = 10, + APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALIGNMENT = 11, + APIR_COMMAND_TYPE_BUFFER_TYPE_GET_MAX_SIZE = 12, + APIR_COMMAND_TYPE_BUFFER_TYPE_IS_HOST = 13, + APIR_COMMAND_TYPE_BUFFER_TYPE_ALLOC_BUFFER = 14, + APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALLOC_SIZE = 15, + + /* buffer */ + APIR_COMMAND_TYPE_BUFFER_GET_BASE = 16, + APIR_COMMAND_TYPE_BUFFER_SET_TENSOR = 17, + APIR_COMMAND_TYPE_BUFFER_GET_TENSOR = 18, + APIR_COMMAND_TYPE_BUFFER_CPY_TENSOR = 19, + APIR_COMMAND_TYPE_BUFFER_CLEAR = 20, + APIR_COMMAND_TYPE_BUFFER_FREE_BUFFER = 21, + + /* backend */ + APIR_COMMAND_TYPE_BACKEND_GRAPH_COMPUTE = 22, + + // last command_type index + 1 + APIR_BACKEND_DISPATCH_TABLE_COUNT = 23, +} ApirBackendCommandType; + +static inline const char * apir_dispatch_command_name(ApirBackendCommandType type) { + switch (type) { + /* device */ + case APIR_COMMAND_TYPE_DEVICE_GET_DEVICE_COUNT: + return "device_get_device_count"; + case APIR_COMMAND_TYPE_DEVICE_GET_COUNT: + return "device_get_count"; + case APIR_COMMAND_TYPE_DEVICE_GET_NAME: + return "device_get_name"; + case APIR_COMMAND_TYPE_DEVICE_GET_DESCRIPTION: + return "device_get_description"; + case APIR_COMMAND_TYPE_DEVICE_GET_TYPE: + return "device_get_type"; + case APIR_COMMAND_TYPE_DEVICE_GET_MEMORY: + return "device_get_memory"; + case APIR_COMMAND_TYPE_DEVICE_SUPPORTS_OP: + return "device_supports_op"; + case APIR_COMMAND_TYPE_DEVICE_GET_BUFFER_TYPE: + return "device_get_buffer_type"; + case APIR_COMMAND_TYPE_DEVICE_GET_PROPS: + return "device_get_props"; + case APIR_COMMAND_TYPE_DEVICE_BUFFER_FROM_PTR: + return "device_buffer_from_ptr"; + /* buffer-type */ + case APIR_COMMAND_TYPE_BUFFER_TYPE_GET_NAME: + return "buffer_type_get_name"; + case APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALIGNMENT: + return "buffer_type_get_alignment"; + case APIR_COMMAND_TYPE_BUFFER_TYPE_GET_MAX_SIZE: + return "buffer_type_get_max_size"; + case APIR_COMMAND_TYPE_BUFFER_TYPE_IS_HOST: + return "buffer_type_is_host"; + case APIR_COMMAND_TYPE_BUFFER_TYPE_ALLOC_BUFFER: + return "buffer_type_alloc_buffer"; + case APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALLOC_SIZE: + return "buffer_type_get_alloc_size"; + /* buffer */ + case APIR_COMMAND_TYPE_BUFFER_GET_BASE: + return "buffer_get_base"; + case APIR_COMMAND_TYPE_BUFFER_SET_TENSOR: + return "buffer_set_tensor"; + case APIR_COMMAND_TYPE_BUFFER_GET_TENSOR: + return "buffer_get_tensor"; + case APIR_COMMAND_TYPE_BUFFER_CPY_TENSOR: + return "buffer_cpy_tensor"; + case APIR_COMMAND_TYPE_BUFFER_CLEAR: + return "buffer_clear"; + case APIR_COMMAND_TYPE_BUFFER_FREE_BUFFER: + return "buffer_free_buffer"; + /* backend */ + case APIR_COMMAND_TYPE_BACKEND_GRAPH_COMPUTE: + return "backend_graph_compute"; + + default: + return "unknown"; + } +} + + +#include // for uintptr_t +#include // for timespec, clock_gettime + +#define APIR_BACKEND_INITIALIZE_SUCCESS 0 +#define APIR_BACKEND_INITIALIZE_CANNOT_OPEN_BACKEND_LIBRARY 1 +#define APIR_BACKEND_INITIALIZE_CANNOT_OPEN_GGML_LIBRARY 2 +#define APIR_BACKEND_INITIALIZE_MISSING_BACKEND_SYMBOLS 3 +#define APIR_BACKEND_INITIALIZE_MISSING_GGML_SYMBOLS 4 +#define APIR_BACKEND_INITIALIZE_BACKEND_FAILED 5 +#define APIR_BACKEND_INITIALIZE_BACKEND_REG_FAILED 6 +#define APIR_BACKEND_INITIALIZE_ALREADY_INITED 7 +#define APIR_BACKEND_INITIALIZE_NO_DEVICE 8 +#define APIR_BACKEND_INITIALIZE_BACKEND_INIT_FAILED 9 + +// new entries here need to be added to the apir_backend_initialize_error function below + +#define APIR_BACKEND_FORWARD_INDEX_INVALID 6 + +// 0 is fast, 1 avoids the backend to crash if an unsupported tensor is received +#define APIR_BACKEND_CHECK_SUPPORTS_OP 0 + +typedef uintptr_t apir_buffer_type_host_handle_t; +typedef uintptr_t apir_buffer_host_handle_t; + +static const char * apir_backend_initialize_error(int code) { +#define APIR_BACKEND_INITIALIZE_ERROR(code_name) \ + do { \ + if (code == code_name) \ + return #code_name; \ + } while (0) + + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_SUCCESS); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_CANNOT_OPEN_BACKEND_LIBRARY); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_CANNOT_OPEN_GGML_LIBRARY); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_MISSING_BACKEND_SYMBOLS); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_MISSING_GGML_SYMBOLS); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_BACKEND_FAILED); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_BACKEND_REG_FAILED); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_ALREADY_INITED); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_NO_DEVICE); + APIR_BACKEND_INITIALIZE_ERROR(APIR_BACKEND_INITIALIZE_BACKEND_INIT_FAILED); + + return "Unknown APIR_BACKEND_INITIALIZE error:/"; + +#undef APIR_BACKEND_INITIALIZE_ERROR +} + +#include "ggml-alloc.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml.h.inc" + +struct apir_buffer_context_t { + apir_buffer_host_handle_t host_handle; + + struct virtgpu_shmem shmem; + apir_buffer_type_host_handle_t buft_host_handle; +}; + +/* device */ +void apir_device_get_device_count(struct virtgpu * gpu); +int apir_device_get_count(struct virtgpu * gpu); +char * apir_device_get_name(struct virtgpu * gpu); +char * apir_device_get_description(struct virtgpu * gpu); +uint32_t apir_device_get_type(struct virtgpu * gpu); +void apir_device_get_memory(struct virtgpu * gpu, size_t * free, size_t * total); +bool apir_device_supports_op(struct virtgpu * gpu, const ggml_tensor * op); +apir_buffer_type_host_handle_t apir_device_get_buffer_type(struct virtgpu * gpu); +void apir_device_get_props(struct virtgpu * gpu, + bool * async, + bool * host_buffer, + bool * buffer_from_host_ptr, + bool * events); +apir_buffer_context_t apir_device_buffer_from_ptr(struct virtgpu * gpu, size_t size, size_t max_tensor_size); + +/* buffer-type */ +char * apir_buffer_type_get_name(struct virtgpu * gpu, apir_buffer_type_host_handle_t host_handle); +size_t apir_buffer_type_get_alignment(struct virtgpu * gpu, apir_buffer_type_host_handle_t host_handle); +size_t apir_buffer_type_get_max_size(struct virtgpu * gpu, apir_buffer_type_host_handle_t host_handle); +/* apir_buffer_type_is_host is deprecated. */ +apir_buffer_context_t apir_buffer_type_alloc_buffer(struct virtgpu * gpu, + apir_buffer_type_host_handle_t host_handle, + size_t size); +size_t apir_buffer_type_get_alloc_size(struct virtgpu * gpu, + apir_buffer_type_host_handle_t host_handle, + const ggml_tensor * op); + +/* buffer */ +void * apir_buffer_get_base(struct virtgpu * gpu, apir_buffer_context_t * buffer_context); +void apir_buffer_set_tensor(struct virtgpu * gpu, + apir_buffer_context_t * buffer_context, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size); +void apir_buffer_get_tensor(struct virtgpu * gpu, + apir_buffer_context_t * buffer_context, + const ggml_tensor * tensor, + void * data, + size_t offset, + size_t size); +bool apir_buffer_cpy_tensor(struct virtgpu * gpu, + apir_buffer_context_t * buffer_context, + const ggml_tensor * src, + const ggml_tensor * dst); +void apir_buffer_clear(struct virtgpu * gpu, apir_buffer_context_t * buffer_context, uint8_t value); +void apir_buffer_free_buffer(struct virtgpu * gpu, apir_buffer_context_t * buffer_context); + +/* backend */ +ggml_status apir_backend_graph_compute(struct virtgpu * gpu, ggml_cgraph * cgraph); + + + +/* the rest of this file must match virglrenderer/src/apir-protocol.h */ + +#include + +#include + +#define APIR_PROTOCOL_MAJOR 0 +#define APIR_PROTOCOL_MINOR 1 + +#define APIR_HANDSHAKE_MAGIC 0xab1e + +enum ApirCommandType { + APIR_COMMAND_TYPE_HANDSHAKE = 0, + APIR_COMMAND_TYPE_LOADLIBRARY = 1, + APIR_COMMAND_TYPE_FORWARD = 2, + + APIR_COMMAND_TYPE_LENGTH = 3, +}; + +typedef uint64_t ApirCommandFlags; + +enum ApirLoadLibraryReturnCode { + APIR_LOAD_LIBRARY_SUCCESS = 0, + // these error codes are returned by the Virglrenderer APIR component + APIR_LOAD_LIBRARY_HYPERCALL_INITIALIZATION_ERROR = 1, + APIR_LOAD_LIBRARY_ALREADY_LOADED = 2, + APIR_LOAD_LIBRARY_ENV_VAR_MISSING = 3, + APIR_LOAD_LIBRARY_CANNOT_OPEN = 4, + APIR_LOAD_LIBRARY_SYMBOL_MISSING = 5, + // any value greater than this is an APIR *backend library* initialization return code + APIR_LOAD_LIBRARY_INIT_BASE_INDEX = 6, +}; + +enum ApirForwardReturnCode { + APIR_FORWARD_SUCCESS = 0, + // these error codes are returned by the Virglrenderer APIR component + APIR_FORWARD_NO_DISPATCH_FCT = 1, + APIR_FORWARD_TIMEOUT = 2, + APIR_FORWARD_FAILED_TO_SYNC_STREAMS = 3, + // any value greater than this index an APIR *backend library* forward return code + APIR_FORWARD_BASE_INDEX = 4, +}; + +__attribute__((unused)) static inline const char * apir_command_name(ApirCommandType type) { + switch (type) { + case APIR_COMMAND_TYPE_HANDSHAKE: + return "HandShake"; + case APIR_COMMAND_TYPE_LOADLIBRARY: + return "LoadLibrary"; + case APIR_COMMAND_TYPE_FORWARD: + return "Forward"; + default: + return "unknown"; + } +} + +__attribute__((unused)) static const char * apir_load_library_error(ApirLoadLibraryReturnCode code) { +#define APIR_LOAD_LIBRARY_ERROR(code_name) \ + do { \ + if (code == code_name) \ + return #code_name; \ + } while (0) + + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_SUCCESS); + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_HYPERCALL_INITIALIZATION_ERROR); + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_ALREADY_LOADED); + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_ENV_VAR_MISSING); + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_CANNOT_OPEN); + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_SYMBOL_MISSING); + APIR_LOAD_LIBRARY_ERROR(APIR_LOAD_LIBRARY_INIT_BASE_INDEX); + + return "Unknown APIR_COMMAND_TYPE_LoadLibrary error"; + +#undef APIR_LOAD_LIBRARY_ERROR +} + +__attribute__((unused)) static const char * apir_forward_error(ApirForwardReturnCode code) { +#define APIR_FORWARD_ERROR(code_name) \ + do { \ + if (code == code_name) \ + return #code_name; \ + } while (0) + + APIR_FORWARD_ERROR(APIR_FORWARD_SUCCESS); + APIR_FORWARD_ERROR(APIR_FORWARD_NO_DISPATCH_FCT); + APIR_FORWARD_ERROR(APIR_FORWARD_TIMEOUT); + APIR_FORWARD_ERROR(APIR_FORWARD_FAILED_TO_SYNC_STREAMS); + APIR_FORWARD_ERROR(APIR_FORWARD_BASE_INDEX); + + return "Unknown APIR_COMMAND_TYPE_FORWARD error"; + +#undef APIR_FORWARD_ERROR +} + +#include "ggml-impl-defs.inc" + +#include +#include + +#define likely(x) __builtin_expect(!!(x), 1) +#define unlikely(x) __builtin_expect(!!(x), 0) + +struct apir_encoder { + char * cur; + const char * start; + const char * end; + bool fatal; +}; + +struct apir_decoder { + const char * cur; + const char * end; + bool fatal; +}; + +/* + * new encoder and decoder + */ + +static apir_decoder apir_new_decoder(const char * ptr, size_t size) { + apir_decoder dec = { + .cur = ptr, + .end = ptr + size, + .fatal = false, + }; + + return dec; +} + +static apir_encoder apir_new_encoder(char * ptr, size_t size) { + apir_encoder enc = { + .cur = ptr, + .start = ptr, + .end = ptr + size, + .fatal = false, + }; + + return enc; +} + +/* + * fatal flag handling + */ + +static inline void apir_encoder_reset_fatal(apir_encoder * enc) { + enc->fatal = false; +} + +static inline void apir_encoder_set_fatal(apir_encoder * enc) { + enc->fatal = true; +} + +static inline bool apir_encoder_get_fatal(const apir_encoder * enc) { + return enc->fatal; +} + +static inline void apir_decoder_reset_fatal(apir_decoder * dec) { + dec->fatal = false; +} + +static inline void apir_decoder_set_fatal(apir_decoder * dec) { + dec->fatal = true; +} + +static inline bool apir_decoder_get_fatal(const apir_decoder * dec) { + return dec->fatal; +} + +/* + * encode peek + */ + +static inline bool apir_decoder_peek_internal(apir_decoder * dec, size_t size, void * val, size_t val_size) { + assert(val_size <= size); + + if (unlikely(size > (size_t) (dec->end - dec->cur))) { + GGML_LOG_ERROR("%s: reading too much from the decoder ...\n", __func__); + apir_decoder_set_fatal(dec); + memset(val, 0, val_size); + return false; + } + + /* we should not rely on the compiler to optimize away memcpy... */ + memcpy(val, dec->cur, val_size); + return true; +} + +static inline void apir_decoder_peek(apir_decoder * dec, size_t size, void * val, size_t val_size) { + apir_decoder_peek_internal(dec, size, val, val_size); +} + +static inline const void * apir_decoder_use_inplace(apir_decoder * dec, size_t size) { + if (unlikely(size > (size_t) (dec->end - dec->cur))) { + GGML_LOG_ERROR("%s: reading too much from the decoder ...\n", __func__); + apir_decoder_set_fatal(dec); + return NULL; + } + const void * addr = dec->cur; + dec->cur += size; + + return addr; +} + +/* + * read/write + */ + +static inline void apir_decoder_read(apir_decoder * dec, size_t size, void * val, size_t val_size) { + if (apir_decoder_peek_internal(dec, size, val, val_size)) { + dec->cur += size; + } +} + +static inline char * apir_encoder_write(apir_encoder * enc, size_t size, const void * val, size_t val_size) { + assert(val_size <= size); + assert(size <= ((size_t) (enc->end - enc->cur))); + + char * write_addr = enc->cur; + /* we should not rely on the compiler to optimize away memcpy... */ + memcpy(write_addr, val, val_size); + enc->cur += size; + + return write_addr; +} + +/* + * encode/decode + */ + +static inline void apir_decode(apir_decoder * dec, size_t size, void * data, size_t data_size) { + assert(size % 4 == 0); + apir_decoder_read(dec, size, data, data_size); +} + +static inline void apir_encode(apir_encoder * enc, size_t size, const void * data, size_t data_size) { + assert(size % 4 == 0); + apir_encoder_write(enc, size, data, data_size); +} + +/* + * typed encode/decode + */ + +/* uint8_t */ + +static inline void apir_encode_uint8_t(apir_encoder * enc, const uint8_t * val) { + apir_encode(enc, sizeof(int), val, sizeof(*val)); +} + +static inline void apir_decode_uint8_t(apir_decoder * dec, uint8_t * val) { + apir_decode(dec, sizeof(int), val, sizeof(*val)); +} + +/* uint64_t */ + +static inline void apir_encode_uint64_t(apir_encoder * enc, const uint64_t * val) { + apir_encode(enc, 8, val, sizeof(*val)); +} + +static inline void apir_decode_uint64_t(apir_decoder * dec, uint64_t * val) { + apir_decode(dec, 8, val, sizeof(*val)); +} + +static inline void apir_encode_uint64_t_array(apir_encoder * enc, const uint64_t * val, uint32_t count) { + const size_t size = sizeof(*val) * count; + assert(size >= count); + apir_encode(enc, size, val, size); +} + +static inline void apir_decode_uint64_t_array(apir_decoder * dec, uint64_t * val, uint32_t count) { + const size_t size = sizeof(*val) * count; + assert(size >= count); + apir_decode(dec, size, val, size); +} + +static inline const uint64_t * apir_decode_uint64_t_array_inplace(apir_decoder * dec, uint32_t count) { + return (uint64_t *) (uintptr_t) apir_decoder_use_inplace(dec, count * sizeof(uint64_t)); +} + +/* int32_t */ + +static inline void apir_encode_int32_t(apir_encoder * enc, const int32_t * val) { + apir_encode(enc, 4, val, sizeof(*val)); +} + +static inline void apir_decode_int32_t(apir_decoder * dec, int32_t * val) { + apir_decode(dec, 4, val, sizeof(*val)); +} + +static inline void apir_encode_int32_t_array(apir_encoder * enc, const int32_t * val, uint32_t count) { + const size_t size = sizeof(*val) * count; + assert(size >= count); + apir_encode(enc, size, val, size); +} + +static inline void apir_decode_int32_t_array(apir_decoder * dec, int32_t * val, uint32_t count) { + const size_t size = sizeof(*val) * count; + assert(size >= count); + apir_decode(dec, size, val, size); +} + +/* array size (uint64_t) */ + +static inline void apir_encode_array_size(apir_encoder * enc, uint64_t size) { + apir_encode_uint64_t(enc, &size); +} + +static inline uint64_t apir_decode_array_size(apir_decoder * dec, uint64_t expected_size) { + uint64_t size; + apir_decode_uint64_t(dec, &size); + if (size != expected_size) { + GGML_LOG_ERROR("%s: Couldn't decode array from the decoder\n", __func__); + apir_decoder_set_fatal(dec); + size = 0; + } + return size; +} + +static inline uint64_t apir_decode_array_size_unchecked(apir_decoder * dec) { + uint64_t size; + apir_decode_uint64_t(dec, &size); + return size; +} + +/* non-array pointer */ + +static inline bool apir_encode_simple_pointer(apir_encoder * enc, const void * val) { + apir_encode_array_size(enc, val ? 1 : 0); + return val; +} + +static inline bool apir_decode_simple_pointer(apir_decoder * dec) { + return apir_decode_array_size_unchecked(dec); +} + +/* uint32_t */ + +static inline void apir_encode_uint32_t(apir_encoder * enc, const uint32_t * val) { + apir_encode(enc, 4, val, sizeof(*val)); +} + +static inline void apir_decode_uint32_t(apir_decoder * dec, uint32_t * val) { + apir_decode(dec, 4, val, sizeof(*val)); +} + +static inline void apir_encode_uint32_t_array(apir_encoder * enc, const uint32_t * val, uint32_t count) { + const size_t size = sizeof(*val) * count; + assert(size >= count); + apir_encode(enc, size, val, size); +} + +static inline void apir_decode_uint32_t_array(apir_decoder * dec, uint32_t * val, uint32_t count) { + const size_t size = sizeof(*val) * count; + assert(size >= count); + apir_decode(dec, size, val, size); +} + +/* size_t */ + +static inline void apir_encode_size_t(apir_encoder * enc, const size_t * val) { + const uint64_t tmp = *val; + apir_encode_uint64_t(enc, &tmp); +} + +static inline void apir_decode_size_t(apir_decoder * dec, size_t * val) { + uint64_t tmp; + apir_decode_uint64_t(dec, &tmp); + *val = tmp; +} + +static inline void apir_encode_size_t_array(apir_encoder * enc, const size_t * val, uint32_t count) { + if (sizeof(size_t) == sizeof(uint64_t)) { + apir_encode_uint64_t_array(enc, (const uint64_t *) val, count); + } else { + for (uint32_t i = 0; i < count; i++) { + apir_encode_size_t(enc, &val[i]); + } + } +} + +static inline void apir_decode_size_t_array(apir_decoder * dec, size_t * val, uint32_t count) { + if (sizeof(size_t) == sizeof(uint64_t)) { + apir_decode_uint64_t_array(dec, (uint64_t *) val, count); + } else { + for (uint32_t i = 0; i < count; i++) { + apir_decode_size_t(dec, &val[i]); + } + } +} + +/* opaque blob */ + +static inline void apir_encode_blob_array(apir_encoder * enc, const void * val, size_t size) { + apir_encode(enc, (size + 3) & ~3, val, size); +} + +static inline void apir_decode_blob_array(apir_decoder * dec, void * val, size_t size) { + apir_decode(dec, (size + 3) & ~3, val, size); +} + +/* string */ + +static inline void apir_encode_char_array(apir_encoder * enc, const char * val, size_t size) { + assert(size && strlen(val) < size); + apir_encode_blob_array(enc, val, size); +} + +static inline void apir_decode_char_array(apir_decoder * dec, char * val, size_t size) { + apir_decode_blob_array(dec, val, size); + if (size) { + val[size - 1] = '\0'; + } else { + GGML_LOG_ERROR("%s: Couldn't decode the blog array\n", __func__); + apir_decoder_set_fatal(dec); + } +} + +/* (temp) buffer allocation */ + +static inline void * apir_decoder_alloc_array(size_t size, size_t count) { + size_t alloc_size; + if (unlikely(__builtin_mul_overflow(size, count, &alloc_size))) { + GGML_LOG_ERROR("%s: overflow in array allocation of %zu * %zu bytes\n", __func__, size, count); + return NULL; + } + + return malloc(alloc_size); +} + +/* bool */ + +static inline void apir_encode_bool_t(apir_encoder * enc, const bool * val) { + apir_encode(enc, sizeof(int), val, sizeof(bool)); +} + +static inline void apir_decode_bool_t(apir_decoder * dec, bool * val) { + apir_decode(dec, sizeof(int), val, sizeof(bool)); +} + +/* apir_buffer_type_host_handle_t */ + +static inline void apir_encode_apir_buffer_type_host_handle_t(apir_encoder * enc, + const apir_buffer_type_host_handle_t * val) { + apir_encode(enc, sizeof(apir_buffer_type_host_handle_t), val, sizeof(apir_buffer_type_host_handle_t)); +} + +static inline void apir_decode_apir_buffer_type_host_handle_t(apir_decoder * dec, + apir_buffer_type_host_handle_t * val) { + apir_decode(dec, sizeof(apir_buffer_type_host_handle_t), val, sizeof(apir_buffer_type_host_handle_t)); +} + +/* apir_buffer_host_handle_t */ + +static inline void apir_encode_apir_buffer_host_handle_t(apir_encoder * enc, const apir_buffer_host_handle_t * val) { + apir_encode(enc, sizeof(apir_buffer_host_handle_t), val, sizeof(apir_buffer_host_handle_t)); +} + +static inline void apir_decode_apir_buffer_host_handle_t(apir_decoder * dec, apir_buffer_host_handle_t * val) { + apir_decode(dec, sizeof(apir_buffer_host_handle_t), val, sizeof(apir_buffer_host_handle_t)); +} + +/* uintptr_t */ + +static inline void apir_encode_uintptr_t(apir_encoder * enc, const uintptr_t * val) { + apir_encode(enc, sizeof(*val), val, sizeof(*val)); +} + +static inline void apir_decode_uintptr_t(apir_decoder * dec, uintptr_t * val) { + apir_decode(dec, sizeof(*val), val, sizeof(*val)); +} + + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#define VIRGL_RENDERER_UNSTABLE_APIS 1 +#include + +struct virgl_renderer_capset_apir { + uint32_t apir_version; + uint32_t supports_blob_resources; + uint32_t reserved[4]; // For future expansion +}; + +#include +#include "venus_hw.h" +// clang-format on + +#ifndef VIRTGPU_DRM_CAPSET_APIR +// Will be defined include/drm/virtgpu_drm.h when +// https://gitlab.freedesktop.org/virgl/virglrenderer/-/merge_requests/1590/diffs +// is merged +# define VIRTGPU_DRM_CAPSET_APIR 10 +#endif + +// Mesa/Virlgrenderer Venus internal. Only necessary during the +// Venus->APIR transition in Virglrenderer +#define VENUS_COMMAND_TYPE_LENGTH 331 + +#ifndef VIRTGPU_DRM_CAPSET_VENUS // only available with Linux >= v6.16 +# define VIRTGPU_DRM_CAPSET_VENUS 4 +#endif + +typedef uint32_t virgl_renderer_capset; + +/* from src/virtio/vulkan/vn_renderer_virtgpu.c */ +#define VIRTGPU_PCI_VENDOR_ID 0x1af4 +#define VIRTGPU_PCI_DEVICE_ID 0x1050 +#define VIRTGPU_BLOB_MEM_GUEST_VRAM 0x0004 +#define VIRTGPU_PARAM_GUEST_VRAM 9 + +#define SHMEM_DATA_SIZE 0x1830000 // 24MiB +#define SHMEM_REPLY_SIZE 0x4000 + +#define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0])) + +enum virt_gpu_result_t { + APIR_SUCCESS = 0, + APIR_ERROR_INITIALIZATION_FAILED = -1, +}; + +#define PRINTFLIKE(f, a) __attribute__((format(__printf__, f, a))) + +struct virtgpu { + bool use_apir_capset; + + int fd; + + struct { + virgl_renderer_capset id; + uint32_t version; + virgl_renderer_capset_apir data; + } capset; + + util_sparse_array shmem_array; + + /* APIR communication pages */ + virtgpu_shmem reply_shmem; + virtgpu_shmem data_shmem; + + /* Mutex to protect shared data_shmem buffer from concurrent access */ + mtx_t data_shmem_mutex; + + /* Cached device information to prevent memory leaks and race conditions */ + struct { + char * description; + char * name; + int32_t device_count; + uint32_t type; + size_t memory_free; + size_t memory_total; + } cached_device_info; + + /* Cached buffer type information to prevent memory leaks and race conditions */ + struct { + apir_buffer_type_host_handle_t host_handle; + char * name; + size_t alignment; + size_t max_size; + } cached_buffer_type; +}; + +static inline int virtgpu_ioctl(virtgpu * gpu, unsigned long request, void * args) { + return drmIoctl(gpu->fd, request, args); +} + +virtgpu * create_virtgpu(); + +#include "ggml-backend-impl-defs.inc" +#include "ggml-backend.h.inc" + +#define GGML_VIRTGPU_NAME "ggml-virtgpu" +#define GGML_VIRTGPU "ggml-virtgpu: " + +// USE_ALWAYS_TRUE_SUPPORTS_OP: 1 is fast, 0 avoid micro-benchmark crashes + +#define USE_ALWAYS_TRUE_SUPPORTS_OP 1 +#define USE_METAL_GUEST_SUPPORTS_OP 0 + +#define DEV_TO_GPU(name) ((ggml_backend_remoting_device_context *) (name)->context)->gpu + +#define BUFFER_TO_GGML_CONTEXT(name) ((ggml_backend_remoting_buffer_context *) (name)->context) + +#define BUFFER_TO_APIR_CONTEXT(name) &((ggml_backend_remoting_buffer_context *) (name)->context)->apir_context + +#define BUFFER_TO_HOST_HANDLE(name) ((ggml_backend_remoting_buffer_context *) (name)->context)->apir_context.host_handle + +#define GET_DEVICE_CONTEXT() (ggml_backend_remoting_device_context *) ggml_backend_remoting_get_device(0)->context + +#define BUFT_TO_GPU(name) ((ggml_backend_remoting_device_context *) (name)->device->context)->gpu + +struct ggml_backend_remoting_device_context { + size_t device; + std::string name; + std::string description; + + std::vector> shared_memory; + + virtgpu * gpu; +}; + +struct ggml_backend_remoting_buffer_context { + apir_buffer_context_t apir_context; + + virtgpu * gpu; + + void * base; + + bool is_from_ptr; +}; + +extern const ggml_backend_buffer_type_i ggml_backend_remoting_buffer_type_interface; +extern const ggml_backend_device_i ggml_backend_remoting_device_interface; +extern const ggml_backend_buffer_i ggml_backend_remoting_buffer_interface; +extern const ggml_backend_buffer_type_i ggml_backend_remoting_buffer_from_ptr_type_interface; +extern const ggml_backend_buffer_i ggml_backend_remoting_buffer_from_ptr_interface; + +ggml_backend_dev_t ggml_backend_remoting_get_device(size_t device); +ggml_backend_t ggml_backend_remoting_device_init(ggml_backend_dev_t dev, const char * params); +ggml_backend_buffer_type_t ggml_backend_remoting_device_get_buffer_type(ggml_backend_dev_t dev); + +static inline apir_buffer_type_host_handle_t ggml_buffer_type_to_apir_handle(ggml_backend_buffer_type_t buft) { + // in the backend, the buffer handle is the buffer pointer + return (apir_buffer_type_host_handle_t) buft->context; +} + +static inline apir_buffer_host_handle_t ggml_buffer_to_apir_handle(ggml_backend_buffer_t buffer) { + if (!buffer->context) { + GGML_ABORT(GGML_VIRTGPU "%s: no context available :/", __func__); + } + return BUFFER_TO_HOST_HANDLE(buffer); +} + +apir_encoder * remote_call_prepare(virtgpu * gpu, ApirCommandType apir_cmd_type, int32_t cmd_flags); + +uint32_t remote_call(virtgpu * gpu, + apir_encoder * enc, + apir_decoder ** dec, + float max_wait_ms, + long long * call_duration_ns); + +void remote_call_finish(virtgpu * gpu, apir_encoder * enc, apir_decoder * dec); + +struct apir_rpc_tensor { + uint64_t id; + uint32_t type; + uint64_t buffer; + uint32_t ne[GGML_MAX_DIMS]; + uint32_t nb[GGML_MAX_DIMS]; + uint32_t op; + int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; + int32_t flags; + uint64_t src[GGML_MAX_SRC]; + uint64_t view_src; + uint64_t view_offs; + uint64_t data; + char name[GGML_MAX_NAME]; + + char padding[4]; +}; + +/* frontend */ + +apir_rpc_tensor apir_serialize_tensor(const ggml_tensor * tensor); + +void apir_serialize_graph(const ggml_cgraph * cgraph, std::vector & output); + +/* backend */ + +void apir_track_backend_buffer(ggml_backend_buffer_t buffer); +bool apir_untrack_backend_buffer(ggml_backend_buffer_t buffer); +std::unordered_set apir_get_track_backend_buffers(); + +void apir_add_tensor(ggml_tensor * tensor, + std::vector & tensors, + std::unordered_set & visited); + +ggml_tensor * apir_deserialize_tensor(ggml_context * ctx, const apir_rpc_tensor * tensor); + +ggml_tensor * apir_create_node(uint64_t id, + ggml_context * ctx, + const std::unordered_map & tensor_ptrs, + std::unordered_map & tensor_map); + +ggml_cgraph * apir_deserialize_graph(uint32_t n_nodes, + uint32_t n_tensors, + const apir_rpc_tensor * tensors, + const uint64_t * nodes); + +#include "ggml-impl-defs.inc" + +// ggml_buffer_to_apir_host_handle(ggml_backend_buffer_t buffer); + +static inline void apir_encode_ggml_buffer_host_handle(apir_encoder * enc, const apir_buffer_host_handle_t * handle); + +static inline ggml_backend_buffer_t apir_decode_ggml_buffer(apir_decoder * dec); + +/* apir_rpc_tensor */ + +static inline void apir_encode_rcp_tensor(apir_encoder * enc, const apir_rpc_tensor * apir_rpc_tensor) { + size_t apir_rpc_tensor_size = sizeof(*apir_rpc_tensor); + apir_encode(enc, apir_rpc_tensor_size, apir_rpc_tensor, apir_rpc_tensor_size); +} + +static inline apir_rpc_tensor * apir_decode_apir_rpc_tensor_inplace(apir_decoder * dec) { + size_t apir_rpc_tensor_size = sizeof(apir_rpc_tensor); + + return (apir_rpc_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, apir_rpc_tensor_size); +} + +static inline apir_rpc_tensor * apir_decode_apir_rpc_tensor_array_inplace(apir_decoder * dec, uint32_t n_tensors) { + size_t apir_rpc_tensor_size = sizeof(apir_rpc_tensor) * n_tensors; + + return (apir_rpc_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, apir_rpc_tensor_size); +} + +/* ggml_tensor */ + +static inline void apir_encode_ggml_tensor(apir_encoder * enc, const ggml_tensor * tensor) { + apir_rpc_tensor serialized = apir_serialize_tensor(tensor); + + apir_encode_rcp_tensor(enc, &serialized); +} + +static inline const ggml_tensor * apir_decode_ggml_tensor(apir_decoder * dec) { + const apir_rpc_tensor * apir_rpc_tensor = apir_decode_apir_rpc_tensor_inplace(dec); + + if (!apir_rpc_tensor) { + return NULL; + } + + ggml_init_params params{ + /*.mem_size =*/ggml_tensor_overhead(), + /*.mem_buffer =*/NULL, + /*.no_alloc =*/true, + }; + + ggml_context * ctx = ggml_init(params); + + const ggml_tensor * tensor = apir_deserialize_tensor(ctx, apir_rpc_tensor); + + return tensor; +} + +/* *** ggml_backend_buffer_type_t *** */ + +// ggml_backend_buffer_type_t is a POINTER (to a struct). +// Only the host pointer is shared between the host and guest. +// The guest stores it in `buft->context`. +// The host simply writes the pointer address in the buffer variable. + +static inline void apir_encode_ggml_buffer_type(apir_encoder * enc, ggml_backend_buffer_type_t buft) { + apir_buffer_type_host_handle_t handle = ggml_buffer_type_to_apir_handle(buft); + apir_encoder_write(enc, sizeof(handle), &handle, sizeof(handle)); +} + +static inline ggml_backend_buffer_type_t apir_decode_ggml_buffer_type(apir_decoder * dec) { + apir_buffer_type_host_handle_t handle; + + apir_decoder_read(dec, sizeof(handle), &handle, sizeof(handle)); + + return (ggml_backend_buffer_type_t) handle; +} + +static inline void apir_encode_apir_buffer_type_host_handle(apir_encoder * enc, apir_buffer_type_host_handle_t handle) { + apir_encoder_write(enc, sizeof(handle), &handle, sizeof(handle)); +} + +static inline apir_buffer_type_host_handle_t apir_decode_apir_buffer_type_host_handle(apir_decoder * dec) { + apir_buffer_type_host_handle_t handle; + + apir_decoder_read(dec, sizeof(handle), &handle, sizeof(handle)); + + return handle; +} + +/* *** ggml_backend_type_t *** */ + +// ggml_backend_buffer_t is a POINTER. +// same logic as for ggml_backend_buffer_type_t + +static inline void apir_encode_ggml_buffer(apir_encoder * enc, const ggml_backend_buffer_t buffer) { + apir_buffer_host_handle_t handle = BUFFER_TO_HOST_HANDLE(buffer); + apir_encoder_write(enc, sizeof(handle), &handle, sizeof(handle)); +} + +static inline ggml_backend_buffer_t apir_decode_ggml_buffer(apir_decoder * dec) { + ggml_backend_buffer_t buffer; + size_t buffer_ptr_size = sizeof(buffer); + + apir_decoder_read(dec, buffer_ptr_size, &buffer, buffer_ptr_size); + + // SECURITY: Validate buffer handle against tracked buffers to prevent + // guest VM from providing arbitrary host memory addresses + if (buffer) { + extern std::unordered_set backend_buffers; + if (backend_buffers.find(buffer) == backend_buffers.end()) { + GGML_LOG_WARN("ggml-virtgpu-backend: %s: Invalid buffer handle from guest: %p\n", __func__, + (void *) buffer); + // Set fatal flag to prevent further processing with invalid handle + apir_decoder_set_fatal(dec); + return NULL; + } + } + + return buffer; +} + +/* enum ggml_status */ + +static inline void apir_encode_ggml_status(apir_encoder * enc, const ggml_status * status) { + apir_encoder_write(enc, sizeof(*status), status, sizeof(*status)); +} + +static inline void apir_decode_ggml_status(apir_decoder * dec, ggml_status * status) { + apir_decoder_read(dec, sizeof(*status), status, sizeof(*status)); +} + +/* virtgpu_shmem */ + +static inline void apir_encode_virtgpu_shmem_res_id(apir_encoder * enc, uint32_t shmem_res_id) { + apir_encode_uint32_t(enc, &shmem_res_id); +} + +static inline void apir_decode_virtgpu_shmem_res_id(apir_decoder * dec, uint32_t * shmem_res_id) { + apir_decode_uint32_t(dec, shmem_res_id); +} + +/* ggml_cgraph */ + +static inline size_t apir_serialize_ggml_cgraph(ggml_cgraph * cgraph, std::vector & cgraph_data) { + apir_serialize_graph(cgraph, cgraph_data); + + return cgraph_data.size(); +} + +static inline void apir_encode_cgraph_data(apir_encoder * enc, std::vector & cgraph_data) { + size_t cgraph_size = cgraph_data.size(); + + apir_encode(enc, cgraph_size, cgraph_data.data(), cgraph_size); +} + +static inline ggml_cgraph * apir_decode_ggml_cgraph(apir_decoder * dec, size_t cgraph_size) { + GGML_UNUSED(cgraph_size); + + uint32_t n_nodes; + apir_decode_uint32_t(dec, &n_nodes); + const uint64_t * nodes = apir_decode_uint64_t_array_inplace(dec, n_nodes); + + uint32_t n_tensors; + apir_decode_uint32_t(dec, &n_tensors); + const apir_rpc_tensor * tensors = apir_decode_apir_rpc_tensor_array_inplace(dec, n_tensors); + + return apir_deserialize_graph(n_nodes, n_tensors, tensors, nodes); +} + +static inline void apir_encode_ggml_buffer_handle(apir_encoder * enc, const apir_buffer_host_handle_t * handle) { + apir_encoder_write(enc, sizeof(*handle), &handle, sizeof(*handle)); +} + +static inline void apir_encode_ggml_tensor_inline(apir_encoder * enc, const ggml_tensor * tensor) { + size_t tensor_size = sizeof(*tensor); + + if (tensor->extra) { + GGML_ABORT("%s: Cannot pass tensors with extra", __func__); + } + + if (tensor->src[0] && tensor->buffer) { + static int first = 1; + if (first) { + GGML_LOG_WARN("%s: Cannot pass tensors with src and buffer\n", __func__); + first = 0; + } + } + + apir_encoder_write(enc, tensor_size, tensor, tensor_size); + + // tensor->data is a pointer inside the device buffer. No need to touch it + // tensor->buffer is a pointer to a buffer. Encoding the buffer handle in sequence. + // (could also make a copy of the tensor, and update locally.) + + if (tensor->buffer) { + apir_buffer_host_handle_t buffer_handle = ggml_buffer_to_apir_handle(tensor->buffer); + apir_encode_ggml_buffer_handle(enc, &buffer_handle); + } + + if (tensor->view_src) { + apir_encoder_write(enc, tensor_size, tensor->view_src, tensor_size); + } + + for (int i = 0; tensor->src[i]; i++) { + const ggml_tensor * tensor_src = tensor->src[i]; + apir_encoder_write(enc, tensor_size, tensor_src, tensor_size); + } +} + +static inline const ggml_tensor * apir_decode_ggml_tensor_inplace(apir_decoder * dec) { + // it safe to remove the `const` qualifier here, we *do* want to + // modify the shared memory data to fix the `src` pointers. + ggml_tensor * tensor = (ggml_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, sizeof(ggml_tensor)); + + // tensor->data is a pointer inside the device buffer. No need to touch it + // tensor->buffer is a pointer to a buffer. Decode the buffer handle encoded in sequence. + if (tensor->buffer) { + tensor->buffer = apir_decode_ggml_buffer(dec); + } + + if (tensor->view_src) { + ggml_tensor * tensor_view_src = (ggml_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, sizeof(ggml_tensor)); + tensor->view_src = tensor_view_src; + } + + for (int i = 0; tensor->src[i]; i++) { + ggml_tensor * tensor_src = (ggml_tensor *) (uintptr_t) apir_decoder_use_inplace(dec, sizeof(ggml_tensor)); + tensor->src[i] = tensor_src; // overwrite op->src[i] pointer with the actual location of the src tensor + } + + return tensor; +} + + +#define REMOTE_CALL_PREPARE(gpu_dev_name, encoder_name, apir_command_type__) \ + int32_t REMOTE_CALL_PREPARE_forward_flag = (int32_t) apir_command_type__; \ + const char * REMOTE_CALL_PREPARE_command_name = apir_dispatch_command_name(apir_command_type__); \ + do { \ + encoder_name = remote_call_prepare(gpu_dev_name, APIR_COMMAND_TYPE_FORWARD, REMOTE_CALL_PREPARE_forward_flag); \ + if (!encoder_name) { \ + GGML_ABORT(GGML_VIRTGPU "%s: failed to prepare the remote call encoder", __func__); \ + } \ + } while (0) + +#define REMOTE_CALL(gpu_dev_name, encoder_name, decoder_name, ret_name) \ + do { \ + ret_name = (ApirForwardReturnCode) remote_call(gpu_dev_name, encoder_name, &decoder_name, 0, NULL); \ + if (!decoder_name) { \ + GGML_ABORT(GGML_VIRTGPU "%s: failed to kick the remote call", __func__); \ + } \ + if (ret_name < APIR_FORWARD_BASE_INDEX) { \ + GGML_ABORT(GGML_VIRTGPU "%s: failed to forward the API call: %s: code %d", __func__, \ + apir_forward_error(ret_name), ret_name); \ + } \ + ret_name = (ApirForwardReturnCode) (ret_name - APIR_FORWARD_BASE_INDEX); \ + if (ret_name != 0) { \ + GGML_ABORT(GGML_VIRTGPU "backend function '%s' failed (return code: %d)", \ + REMOTE_CALL_PREPARE_command_name, ret_name); \ + } \ + } while (0) + + +#include +#include + +#include +#include +#include + +static virt_gpu_result_t virtgpu_open_device(virtgpu * gpu, const drmDevicePtr dev); +static virt_gpu_result_t virtgpu_open(virtgpu * gpu); + +static virt_gpu_result_t virtgpu_init_capset(virtgpu * gpu); +static virt_gpu_result_t virtgpu_init_context(virtgpu * gpu); + +static int virtgpu_ioctl_context_init(virtgpu * gpu, virgl_renderer_capset capset_id); +static int virtgpu_ioctl_get_caps(virtgpu * gpu, + virgl_renderer_capset id, + uint32_t version, + void * capset, + size_t capset_size); +static uint64_t virtgpu_ioctl_getparam(virtgpu * gpu, uint64_t param); +static void virtgpu_init_renderer_info(virtgpu * gpu); + +static void log_call_duration(long long call_duration_ns, const char * name); + +const uint64_t APIR_HANDSHAKE_MAX_WAIT_MS = 2 * 1000; // 2s +const uint64_t APIR_LOADLIBRARY_MAX_WAIT_MS = 60 * 1000; // 60s + +static int virtgpu_handshake(virtgpu * gpu) { + apir_encoder * encoder; + apir_decoder * decoder; + + encoder = remote_call_prepare(gpu, APIR_COMMAND_TYPE_HANDSHAKE, 0); + if (!encoder) { + GGML_ABORT(GGML_VIRTGPU "%s: failed to prepare the remote call encoder", __func__); + return 1; + } + + /* write handshake props */ + + uint32_t guest_major = APIR_PROTOCOL_MAJOR; + uint32_t guest_minor = APIR_PROTOCOL_MINOR; + apir_encode_uint32_t(encoder, &guest_major); + apir_encode_uint32_t(encoder, &guest_minor); + + /* *** */ + + uint32_t ret_magic; + long long call_duration_ns; + ret_magic = remote_call(gpu, encoder, &decoder, APIR_HANDSHAKE_MAX_WAIT_MS, &call_duration_ns); + log_call_duration(call_duration_ns, "API Remoting handshake"); + + if (!decoder) { + GGML_ABORT(GGML_VIRTGPU + "%s: failed to initiate the communication with the virglrenderer library. " + "Most likely, the wrong virglrenderer library was loaded in the hypervisor.", + __func__); + return 1; + } + + /* read handshake return values */ + + uint32_t host_major; + uint32_t host_minor; + + if (ret_magic != APIR_HANDSHAKE_MAGIC) { + GGML_ABORT(GGML_VIRTGPU "%s: handshake with the virglrenderer failed (code=%d | %s)", __func__, ret_magic, + apir_backend_initialize_error(ret_magic)); + } else { + apir_decode_uint32_t(decoder, &host_major); + apir_decode_uint32_t(decoder, &host_minor); + } + + remote_call_finish(gpu, encoder, decoder); + + if (ret_magic != APIR_HANDSHAKE_MAGIC) { + return 1; + } + + GGML_LOG_INFO(GGML_VIRTGPU "%s: Guest is running with %u.%u\n", __func__, guest_major, guest_minor); + GGML_LOG_INFO(GGML_VIRTGPU "%s: Host is running with %u.%u\n", __func__, host_major, host_minor); + + if (guest_major != host_major) { + GGML_LOG_ERROR(GGML_VIRTGPU "Host major (%d) and guest major (%d) version differ\n", host_major, guest_major); + } else if (guest_minor != host_minor) { + GGML_LOG_WARN(GGML_VIRTGPU "Host minor (%d) and guest minor (%d) version differ\n", host_minor, guest_minor); + } + + return 0; +} + +static ApirLoadLibraryReturnCode virtgpu_load_library(virtgpu * gpu) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirLoadLibraryReturnCode ret; + + encoder = remote_call_prepare(gpu, APIR_COMMAND_TYPE_LOADLIBRARY, 0); + if (!encoder) { + GGML_ABORT(GGML_VIRTGPU "%s: hypercall error: failed to prepare the API Remoting command encoder", __func__); + return APIR_LOAD_LIBRARY_HYPERCALL_INITIALIZATION_ERROR; + } + + long long call_duration_ns; + + ret = (ApirLoadLibraryReturnCode) remote_call(gpu, encoder, &decoder, APIR_LOADLIBRARY_MAX_WAIT_MS, + &call_duration_ns); + log_call_duration(call_duration_ns, "API Remoting LoadLibrary"); + + if (!decoder) { + GGML_ABORT(GGML_VIRTGPU "%s: hypercall error: failed to trigger the API Remoting hypercall.\n", __func__); + return APIR_LOAD_LIBRARY_HYPERCALL_INITIALIZATION_ERROR; + } + + remote_call_finish(gpu, encoder, decoder); + + if (ret == APIR_LOAD_LIBRARY_SUCCESS) { + GGML_LOG_INFO(GGML_VIRTGPU "The API Remoting backend was successfully loaded and initialized\n"); + + return ret; + } + + // something wrong happened, find out what. + if (ret < APIR_LOAD_LIBRARY_INIT_BASE_INDEX) { + if (ret == APIR_LOAD_LIBRARY_ENV_VAR_MISSING) { + GGML_ABORT(GGML_VIRTGPU + "%s: virglrenderer could not open the API Remoting backend library, " + "some environment variables are missing. " + "Make sure virglrenderer is correctly configured by the hypervisor. (%s)", + __func__, apir_load_library_error(ret)); + } else if (ret == APIR_LOAD_LIBRARY_CANNOT_OPEN) { + GGML_ABORT(GGML_VIRTGPU + "%s: virglrenderer could not open the API Remoting backend library. " + "Make sure virglrenderer is correctly configured by the hypervisor. (%s)", + __func__, apir_load_library_error(ret)); + } else if (ret == APIR_LOAD_LIBRARY_ENV_VAR_MISSING) { + GGML_ABORT(GGML_VIRTGPU + "%s: could not load the backend library, some symbols are missing. " + "Make sure virglrenderer is correctly configured by the hypervisor. (%s) ", + __func__, apir_load_library_error(ret)); + } else { + GGML_ABORT(GGML_VIRTGPU "%s: virglrenderer could not load the API Remoting backend library. (%s - code %d)", + __func__, apir_load_library_error(ret), ret); + } + return ret; + } + + GGML_LOG_INFO(GGML_VIRTGPU "%s: virglrenderer successfully loaded the API Remoting backend library.\n", __func__); + + ApirLoadLibraryReturnCode apir_ret = (ApirLoadLibraryReturnCode) (ret - APIR_LOAD_LIBRARY_INIT_BASE_INDEX); + + if (apir_ret == APIR_LOAD_LIBRARY_CANNOT_OPEN) { + GGML_ABORT(GGML_VIRTGPU + "%s: the API Remoting backend library couldn't load the GGML backend library. " + "Make sure virglrenderer is correctly configured by the hypervisor. (%s)", + __func__, apir_load_library_error(apir_ret)); + } else if (apir_ret == APIR_LOAD_LIBRARY_SYMBOL_MISSING) { + GGML_ABORT( + GGML_VIRTGPU + "%s: the API Remoting backend library couldn't load the GGML backend library, some symbols are missing. " + "Make sure virglrenderer is correctly configured by the hypervisor. (%s)", + __func__, apir_load_library_error(apir_ret)); + } else if (apir_ret < APIR_LOAD_LIBRARY_INIT_BASE_INDEX) { + GGML_ABORT(GGML_VIRTGPU + "%s: the API Remoting backend library couldn't load the GGML backend library: apir code=%d | %s)", + __func__, apir_ret, apir_load_library_error(apir_ret)); + } else { + uint32_t lib_ret = apir_ret - APIR_LOAD_LIBRARY_INIT_BASE_INDEX; + GGML_ABORT(GGML_VIRTGPU + "%s: the API Remoting backend library failed to initialize its backend library: apir code=%d)", + __func__, lib_ret); + } + return ret; +} + +virtgpu * create_virtgpu() { + virtgpu * gpu = new virtgpu(); + + gpu->use_apir_capset = getenv("GGML_REMOTING_USE_APIR_CAPSET") != nullptr; + util_sparse_array_init(&gpu->shmem_array, sizeof(virtgpu_shmem), 1024); + + // Initialize mutex to protect shared data_shmem buffer + if (mtx_init(&gpu->data_shmem_mutex, mtx_plain) != thrd_success) { + delete gpu; + GGML_ABORT(GGML_VIRTGPU "%s: failed to initialize data_shmem mutex", __func__); + return NULL; + } + + if (virtgpu_open(gpu) != APIR_SUCCESS) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: failed to open the virtgpu device\n", __func__); + return NULL; + } + + if (virtgpu_init_capset(gpu) != APIR_SUCCESS) { + if (gpu->use_apir_capset) { + GGML_ABORT(GGML_VIRTGPU + "%s: failed to initialize the virtgpu APIR capset. Make sure that the virglrenderer library " + "supports it.", + __func__); + } else { + GGML_ABORT(GGML_VIRTGPU "%s: failed to initialize the virtgpu Venus capset", __func__); + } + return NULL; + } + + if (virtgpu_init_context(gpu) != APIR_SUCCESS) { + GGML_ABORT(GGML_VIRTGPU "%s: failed to initialize the GPU context", __func__); + return NULL; + } + + if (virtgpu_shmem_create(gpu, SHMEM_REPLY_SIZE, &gpu->reply_shmem)) { + GGML_ABORT(GGML_VIRTGPU "%s: failed to create the shared reply memory pages", __func__); + return NULL; + } + + if (virtgpu_shmem_create(gpu, SHMEM_DATA_SIZE, &gpu->data_shmem)) { + GGML_ABORT(GGML_VIRTGPU "%s: failed to create the shared data memory pages", __func__); + return NULL; + } + + if (virtgpu_handshake(gpu)) { + GGML_ABORT(GGML_VIRTGPU "%s: failed to handshake with the virglrenderer library", __func__); + return NULL; + } + + if (virtgpu_load_library(gpu) != APIR_LOAD_LIBRARY_SUCCESS) { + GGML_ABORT(GGML_VIRTGPU "%s: failed to load the backend library", __func__); + return NULL; + } + + return gpu; +} + +static virt_gpu_result_t virtgpu_open(virtgpu * gpu) { + drmDevicePtr devs[8]; + int count = drmGetDevices2(0, devs, ARRAY_SIZE(devs)); + if (count < 0) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: failed to enumerate DRM devices\n", __func__); + return APIR_ERROR_INITIALIZATION_FAILED; + } + + virt_gpu_result_t result = APIR_ERROR_INITIALIZATION_FAILED; + for (int i = 0; i < count; i++) { + result = virtgpu_open_device(gpu, devs[i]); + if (result == APIR_SUCCESS) { + break; + } + } + + drmFreeDevices(devs, count); + + return result; +} + +static virt_gpu_result_t virtgpu_open_device(virtgpu * gpu, const drmDevicePtr dev) { + const char * node_path = dev->nodes[DRM_NODE_RENDER]; + + int fd = open(node_path, O_RDWR | O_CLOEXEC); + if (fd < 0) { + GGML_ABORT(GGML_VIRTGPU "%s: failed to open %s", __func__, node_path); + return APIR_ERROR_INITIALIZATION_FAILED; + } + + drmVersionPtr version = drmGetVersion(fd); + if (!version || strcmp(version->name, "virtio_gpu") || version->version_major != 0) { + if (version) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: unknown DRM driver %s version %d\n", __func__, version->name, + version->version_major); + } else { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: failed to get DRM driver version\n", __func__); + } + + if (version) { + drmFreeVersion(version); + } + close(fd); + return APIR_ERROR_INITIALIZATION_FAILED; + } + + gpu->fd = fd; + + drmFreeVersion(version); + + GGML_LOG_INFO(GGML_VIRTGPU "using DRM device %s\n", node_path); + + return APIR_SUCCESS; +} + +static virt_gpu_result_t virtgpu_init_context(virtgpu * gpu) { + assert(!gpu->capset.version); + const int ret = virtgpu_ioctl_context_init(gpu, gpu->capset.id); + if (ret) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: failed to initialize context: %s\n", __func__, strerror(errno)); + return APIR_ERROR_INITIALIZATION_FAILED; + } + + return APIR_SUCCESS; +} + +static virt_gpu_result_t virtgpu_init_capset(virtgpu * gpu) { + if (gpu->use_apir_capset) { + GGML_LOG_INFO(GGML_VIRTGPU "Using the APIR capset\n"); + gpu->capset.id = VIRTGPU_DRM_CAPSET_APIR; + } else { + GGML_LOG_INFO(GGML_VIRTGPU "Using the Venus capset\n"); + gpu->capset.id = VIRTGPU_DRM_CAPSET_VENUS; + } + gpu->capset.version = 0; + + int ret = + virtgpu_ioctl_get_caps(gpu, gpu->capset.id, gpu->capset.version, &gpu->capset.data, sizeof(gpu->capset.data)); + + if (ret) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: failed to get APIR v%d capset: %s\n", __func__, gpu->capset.version, + strerror(errno)); + return APIR_ERROR_INITIALIZATION_FAILED; + } + + assert(gpu->capset.data.supports_blob_resources); + + return APIR_SUCCESS; +} + +static int virtgpu_ioctl_context_init(virtgpu * gpu, virgl_renderer_capset capset_id) { + drm_virtgpu_context_set_param ctx_set_params[3] = { + { + .param = VIRTGPU_CONTEXT_PARAM_CAPSET_ID, + .value = capset_id, + }, + { + .param = VIRTGPU_CONTEXT_PARAM_NUM_RINGS, + .value = 1, + }, + { + .param = VIRTGPU_CONTEXT_PARAM_POLL_RINGS_MASK, + .value = 0, /* don't generate drm_events on fence signaling */ + }, + }; + + drm_virtgpu_context_init args = { + .num_params = ARRAY_SIZE(ctx_set_params), + .pad = 0, + .ctx_set_params = (uintptr_t) &ctx_set_params, + }; + + return virtgpu_ioctl(gpu, DRM_IOCTL_VIRTGPU_CONTEXT_INIT, &args); +} + +static int virtgpu_ioctl_get_caps(virtgpu * gpu, + virgl_renderer_capset id, + uint32_t version, + void * capset, + size_t capset_size) { + drm_virtgpu_get_caps args = { + .cap_set_id = id, + .cap_set_ver = version, + .addr = (uintptr_t) capset, + .size = (__u32) capset_size, + .pad = 0, + }; + + return virtgpu_ioctl(gpu, DRM_IOCTL_VIRTGPU_GET_CAPS, &args); +} + +static uint64_t virtgpu_ioctl_getparam(virtgpu * gpu, uint64_t param) { + /* val must be zeroed because kernel only writes the lower 32 bits */ + uint64_t val = 0; + drm_virtgpu_getparam args = { + .param = param, + .value = (uintptr_t) &val, + }; + + const int ret = virtgpu_ioctl(gpu, DRM_IOCTL_VIRTGPU_GETPARAM, &args); + return ret ? 0 : val; +} + +apir_encoder * remote_call_prepare(virtgpu * gpu, ApirCommandType apir_cmd_type, int32_t cmd_flags) { + /* + * Prepare the command encoder and its buffer + */ + + thread_local char encoder_buffer[4096]; + + thread_local apir_encoder enc; + enc = { + .cur = encoder_buffer, + .start = encoder_buffer, + .end = encoder_buffer + sizeof(encoder_buffer), + .fatal = false, + }; + + /* + * Fill the command encoder with the common args: + * - cmd_type (int32_t) + * - cmd_flags (int32_t) + * - reply res id (uint32_t) + */ + + int32_t cmd_type = apir_cmd_type; + + // for testing during the hypervisor transition + if (!gpu->use_apir_capset) { + cmd_type += VENUS_COMMAND_TYPE_LENGTH; + } + apir_encode_int32_t(&enc, &cmd_type); + apir_encode_int32_t(&enc, &cmd_flags); + + uint32_t reply_res_id = gpu->reply_shmem.res_id; + apir_encode_uint32_t(&enc, &reply_res_id); + + return &enc; +} + +void remote_call_finish(virtgpu * gpu, apir_encoder * enc, apir_decoder * dec) { + UNUSED(gpu); + + if (!enc) { + GGML_ABORT(GGML_VIRTGPU "%s: Invalid (null) encoder", __func__); + } + + if (!dec) { + GGML_ABORT(GGML_VIRTGPU "%s: Invalid (null) decoder", __func__); + } + + if (apir_encoder_get_fatal(enc)) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: Failed to encode the output parameters.", __func__); + } + + if (apir_decoder_get_fatal(dec)) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: Failed to decode the input parameters.", __func__); + } +} + +uint32_t remote_call(virtgpu * gpu, + apir_encoder * encoder, + apir_decoder ** decoder, + float max_wait_ms, + long long * call_duration_ns) { + /* + * Prepare the reply notification pointer + */ + + volatile std::atomic_uint * atomic_reply_notif = (volatile std::atomic_uint *) gpu->reply_shmem.mmap_ptr; + *atomic_reply_notif = 0; + + /* + * Trigger the execbuf ioctl + */ + + drm_virtgpu_execbuffer args = { + .flags = VIRTGPU_EXECBUF_RING_IDX, + .size = (uint32_t) (encoder->cur - encoder->start), + .command = (uintptr_t) encoder->start, + + .bo_handles = 0, + .num_bo_handles = 0, + + .fence_fd = 0, + .ring_idx = 0, + .syncobj_stride = 0, + .num_in_syncobjs = 0, + .num_out_syncobjs = 0, + .in_syncobjs = 0, + .out_syncobjs = 0, + }; + + *decoder = NULL; + + int ret = drmIoctl(gpu->fd, DRM_IOCTL_VIRTGPU_EXECBUFFER, &args); + + if (ret != 0) { + GGML_ABORT(GGML_VIRTGPU "%s: the virtgpu EXECBUFFER ioctl failed (%d)", __func__, ret); + } + + /* + * Wait for the response notification + */ + timer_data wait_host_reply_timer = { 0, 0, 0 }; + + start_timer(&wait_host_reply_timer); + + timespec ts_start, ts_end; + clock_gettime(CLOCK_MONOTONIC, &ts_start); + long long start_time = (long long) ts_start.tv_sec * 1000000000LL + ts_start.tv_nsec; + + bool timedout = false; + uint32_t notif_value = 0; + while (true) { + notif_value = std::atomic_load_explicit(atomic_reply_notif, std::memory_order_acquire); + + if (notif_value != 0) { + break; + } + + int64_t base_sleep_us = 15; + + os_time_sleep(base_sleep_us); + + if (max_wait_ms) { + clock_gettime(CLOCK_MONOTONIC, &ts_end); + long long end_time = (long long) ts_end.tv_sec * 1000000000LL + ts_end.tv_nsec; + float duration_ms = (end_time - start_time) / 1000000; + + if (duration_ms > max_wait_ms) { + timedout = true; + break; + } + } + } + + if (call_duration_ns) { + *call_duration_ns = stop_timer(&wait_host_reply_timer); + } + + if (max_wait_ms && timedout) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: timed out waiting for the host answer...\n", __func__); + return APIR_FORWARD_TIMEOUT; + } + + /* + * Prepare the decoder + */ + static apir_decoder response_dec; + response_dec.cur = (char *) gpu->reply_shmem.mmap_ptr + sizeof(*atomic_reply_notif); + response_dec.end = (char *) gpu->reply_shmem.mmap_ptr + gpu->reply_shmem.mmap_size; + *decoder = &response_dec; + + // extract the actual return value from the notif flag + uint32_t returned_value = notif_value - 1; + return returned_value; +} + +static void log_call_duration(long long call_duration_ns, const char * name) { + double call_duration_ms = (double) call_duration_ns / 1e6; // 1 millisecond = 1e6 nanoseconds + double call_duration_s = (double) call_duration_ns / 1e9; // 1 second = 1e9 nanoseconds + + if (call_duration_s > 1) { + GGML_LOG_INFO(GGML_VIRTGPU "waited %.2fs for the %s host reply...\n", call_duration_s, name); + } else if (call_duration_ms > 1) { + GGML_LOG_INFO(GGML_VIRTGPU "waited %.2fms for the %s host reply...\n", call_duration_ms, name); + } else { + GGML_LOG_INFO(GGML_VIRTGPU "waited %lldns for the %s host reply...\n", call_duration_ns, name); + } +} + + +/* Inlined virtgpu-shm.cpp */ + +#include +#include + +static uint32_t virtgpu_ioctl_resource_create_blob(virtgpu * gpu, + uint32_t blob_mem, + uint32_t blob_flags, + size_t blob_size, + uint64_t blob_id, + uint32_t * res_id) { +#ifdef SIMULATE_BO_SIZE_FIX + blob_size = align64(blob_size, 4096); +#endif + + drm_virtgpu_resource_create_blob args = { + .blob_mem = blob_mem, + .blob_flags = blob_flags, + .bo_handle = 0, + .res_handle = 0, + .size = blob_size, + .pad = 0, + .cmd_size = 0, + .cmd = 0, + .blob_id = blob_id, + }; + + if (virtgpu_ioctl(gpu, DRM_IOCTL_VIRTGPU_RESOURCE_CREATE_BLOB, &args)) { + return 0; + } + + *res_id = args.res_handle; + return args.bo_handle; +} + +static void virtgpu_ioctl_gem_close(virtgpu * gpu, uint32_t gem_handle) { + drm_gem_close args = { + .handle = gem_handle, + .pad = 0, + }; + + const int ret = virtgpu_ioctl(gpu, DRM_IOCTL_GEM_CLOSE, &args); + assert(!ret); +#ifdef NDEBUG + UNUSED(ret); +#endif +} + +static void * virtgpu_ioctl_map(virtgpu * gpu, uint32_t gem_handle, size_t size) { + drm_virtgpu_map args = { + .offset = 0, + .handle = gem_handle, + .pad = 0, + }; + + if (virtgpu_ioctl(gpu, DRM_IOCTL_VIRTGPU_MAP, &args)) { + return NULL; + } + + void * ptr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, gpu->fd, args.offset); + if (ptr == MAP_FAILED) { + return NULL; + } + + return ptr; +} + +void virtgpu_shmem_destroy(virtgpu * gpu, virtgpu_shmem * shmem) { + munmap(shmem->mmap_ptr, shmem->mmap_size); + virtgpu_ioctl_gem_close(gpu, shmem->gem_handle); +} + +int virtgpu_shmem_create(virtgpu * gpu, size_t size, virtgpu_shmem * shmem) { + size = align64(size, 16384); + + uint32_t res_id; + uint32_t gem_handle = virtgpu_ioctl_resource_create_blob(gpu, VIRTGPU_BLOB_MEM_HOST3D, + VIRTGPU_BLOB_FLAG_USE_MAPPABLE, size, 0, &res_id); + + if (!gem_handle) { + return 1; + } + + void * ptr = virtgpu_ioctl_map(gpu, gem_handle, size); + if (!ptr) { + virtgpu_ioctl_gem_close(gpu, gem_handle); + GGML_LOG_ERROR(GGML_VIRTGPU "%s: virtgpu_ioctl_map failed\n", __func__); + return 1; + } + + shmem->res_id = res_id; + shmem->mmap_size = size; + shmem->mmap_ptr = ptr; + shmem->gem_handle = gem_handle; + + return 0; +} + + +/* Inlined virtgpu-utils.cpp */ + +#include +#include + +#include + +#define NODE_ALLOC_ALIGN 64 +#define NODE_PTR_MASK (~((uintptr_t) NODE_ALLOC_ALIGN - 1)) +#define NODE_LEVEL_MASK ((uintptr_t) NODE_ALLOC_ALIGN - 1) +#define NULL_NODE 0 + +#define os_malloc_aligned(_size, _align) _aligned_malloc(_size, _align) +#define os_free_aligned(_ptr) free(_ptr) +#define p_atomic_cmpxchg(v, old, _new) __sync_val_compare_and_swap((v), (old), (_new)) + +static inline uint64_t util_logbase2_64(uint64_t n) { +#if defined(HAVE___BUILTIN_CLZLL) + return ((sizeof(uint64_t) * 8 - 1) - __builtin_clzll(n | 1)); +#else + uint64_t pos = 0ull; + if (n >= 1ull << 32) { + n >>= 32; + pos += 32; + } + if (n >= 1ull << 16) { + n >>= 16; + pos += 16; + } + if (n >= 1ull << 8) { + n >>= 8; + pos += 8; + } + if (n >= 1ull << 4) { + n >>= 4; + pos += 4; + } + if (n >= 1ull << 2) { + n >>= 2; + pos += 2; + } + if (n >= 1ull << 1) { + pos += 1; + } + return pos; +#endif +} + +void util_sparse_array_init(util_sparse_array * arr, size_t elem_size, size_t node_size) { + memset(arr, 0, sizeof(*arr)); + arr->elem_size = elem_size; + arr->node_size_log2 = util_logbase2_64(node_size); + assert(node_size >= 2 && node_size == (1ull << arr->node_size_log2)); +} + +static inline void * os_malloc_aligned(size_t size, size_t alignment) { + void * ptr; + alignment = (alignment + sizeof(void *) - 1) & ~(sizeof(void *) - 1); + if (posix_memalign(&ptr, alignment, size) != 0) { + return NULL; + } + return ptr; +} + +static inline void * _util_sparse_array_node_data(uintptr_t handle) { + return (void *) (handle & NODE_PTR_MASK); +} + +static inline unsigned _util_sparse_array_node_level(uintptr_t handle) { + return handle & NODE_LEVEL_MASK; +} + +static inline void _util_sparse_array_node_finish(util_sparse_array * arr, uintptr_t node) { + if (_util_sparse_array_node_level(node) > 0) { + uintptr_t * children = (uintptr_t *) _util_sparse_array_node_data(node); + size_t node_size = 1ull << arr->node_size_log2; + for (size_t i = 0; i < node_size; i++) { + if (children[i]) { + _util_sparse_array_node_finish(arr, children[i]); + } + } + } + + os_free_aligned(_util_sparse_array_node_data(node)); +} + +static inline uintptr_t _util_sparse_array_node(void * data, unsigned level) { + assert(data != NULL); + assert(((uintptr_t) data & NODE_LEVEL_MASK) == 0); + assert((level & NODE_PTR_MASK) == 0); + return (uintptr_t) data | level; +} + +inline uintptr_t _util_sparse_array_node_alloc(util_sparse_array * arr, unsigned level) { + size_t size; + if (level == 0) { + size = arr->elem_size << arr->node_size_log2; + } else { + size = sizeof(uintptr_t) << arr->node_size_log2; + } + + void * data = os_malloc_aligned(size, NODE_ALLOC_ALIGN); + memset(data, 0, size); + + return _util_sparse_array_node(data, level); +} + +static inline uintptr_t _util_sparse_array_set_or_free_node(uintptr_t * node_ptr, uintptr_t cmp_node, uintptr_t node) { + uintptr_t prev_node = p_atomic_cmpxchg(node_ptr, cmp_node, node); + + if (prev_node != cmp_node) { + /* We lost the race. Free this one and return the one that was already + * allocated. + */ + os_free_aligned(_util_sparse_array_node_data(node)); + return prev_node; + } else { + return node; + } +} + +void * util_sparse_array_get(util_sparse_array * arr, uint64_t idx) { + const unsigned node_size_log2 = arr->node_size_log2; + uintptr_t root = p_atomic_read(&arr->root); + if (unlikely(!root)) { + unsigned root_level = 0; + uint64_t idx_iter = idx >> node_size_log2; + while (idx_iter) { + idx_iter >>= node_size_log2; + root_level++; + } + uintptr_t new_root = _util_sparse_array_node_alloc(arr, root_level); + root = _util_sparse_array_set_or_free_node(&arr->root, NULL_NODE, new_root); + } + + while (1) { + unsigned root_level = _util_sparse_array_node_level(root); + uint64_t root_idx = idx >> (root_level * node_size_log2); + if (likely(root_idx < (1ull << node_size_log2))) { + break; + } + + /* In this case, we have a root but its level is low enough that the + * requested index is out-of-bounds. + */ + uintptr_t new_root = _util_sparse_array_node_alloc(arr, root_level + 1); + + uintptr_t * new_root_children = (uintptr_t *) _util_sparse_array_node_data(new_root); + new_root_children[0] = root; + + /* We only add one at a time instead of the whole tree because it's + * easier to ensure correctness of both the tree building and the + * clean-up path. Because we're only adding one node we never have to + * worry about trying to free multiple things without freeing the old + * things. + */ + root = _util_sparse_array_set_or_free_node(&arr->root, root, new_root); + } + + void * node_data = _util_sparse_array_node_data(root); + unsigned node_level = _util_sparse_array_node_level(root); + while (node_level > 0) { + uint64_t child_idx = (idx >> (node_level * node_size_log2)) & ((1ull << node_size_log2) - 1); + + uintptr_t * children = (uintptr_t *) node_data; + uintptr_t child = p_atomic_read(&children[child_idx]); + + if (unlikely(!child)) { + child = _util_sparse_array_node_alloc(arr, node_level - 1); + child = _util_sparse_array_set_or_free_node(&children[child_idx], NULL_NODE, child); + } + + node_data = _util_sparse_array_node_data(child); + node_level = _util_sparse_array_node_level(child); + } + + uint64_t elem_idx = idx & ((1ull << node_size_log2) - 1); + return (void *) ((char *) node_data + (elem_idx * arr->elem_size)); +} + + +/* Inlined virtgpu-forward-device.cpp */ + +int apir_device_get_count(virtgpu * gpu) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_COUNT); + REMOTE_CALL(gpu, encoder, decoder, ret); + + int32_t dev_count = -1; + apir_decode_int32_t(decoder, &dev_count); + + remote_call_finish(gpu, encoder, decoder); + + return dev_count; +} + +char * apir_device_get_name(virtgpu * gpu) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_NAME); + REMOTE_CALL(gpu, encoder, decoder, ret); + + const size_t string_size = apir_decode_array_size_unchecked(decoder); + char * string = (char *) apir_decoder_alloc_array(sizeof(char), string_size); + if (!string) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: Could not allocate the device name buffer\n", __func__); + return NULL; + } + apir_decode_char_array(decoder, string, string_size); + + remote_call_finish(gpu, encoder, decoder); + + return string; +} + +char * apir_device_get_description(virtgpu * gpu) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_DESCRIPTION); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + const size_t string_size = apir_decode_array_size_unchecked(decoder); + char * string = (char *) apir_decoder_alloc_array(sizeof(char), string_size); + if (!string) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: Could not allocate the device description buffer\n", __func__); + + return NULL; + } + apir_decode_char_array(decoder, string, string_size); + + remote_call_finish(gpu, encoder, decoder); + + return string; +} + +uint32_t apir_device_get_type(virtgpu * gpu) { + static uint32_t dev_type = 255; + if (dev_type != 255) { + return dev_type; + } + + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_TYPE); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + apir_decode_uint32_t(decoder, &dev_type); + + remote_call_finish(gpu, encoder, decoder); + + return dev_type; +} + +void apir_device_get_memory(virtgpu * gpu, size_t * free, size_t * total) { + static size_t dev_free = 0; + static size_t dev_total = 0; + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_MEMORY); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + apir_decode_size_t(decoder, &dev_free); + apir_decode_size_t(decoder, &dev_total); + + *free = dev_free; + *total = dev_total; + + remote_call_finish(gpu, encoder, decoder); + + return; +} + +bool apir_device_supports_op(virtgpu * gpu, const ggml_tensor * op) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_SUPPORTS_OP); + + apir_encode_ggml_tensor_inline(encoder, op); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + bool supports_op; + apir_decode_bool_t(decoder, &supports_op); + + remote_call_finish(gpu, encoder, decoder); + + return supports_op; +} + +apir_buffer_type_host_handle_t apir_device_get_buffer_type(virtgpu * gpu) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_BUFFER_TYPE); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + apir_buffer_type_host_handle_t buft_handle; + apir_decode_apir_buffer_type_host_handle_t(decoder, &buft_handle); + + remote_call_finish(gpu, encoder, decoder); + + return buft_handle; +} + +void apir_device_get_props(virtgpu * gpu, + bool * async, + bool * host_buffer, + bool * buffer_from_host_ptr, + bool * events) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_GET_PROPS); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + apir_decode_bool_t(decoder, async); + apir_decode_bool_t(decoder, host_buffer); + apir_decode_bool_t(decoder, buffer_from_host_ptr); + apir_decode_bool_t(decoder, events); + + remote_call_finish(gpu, encoder, decoder); + + return; +} + +apir_buffer_context_t apir_device_buffer_from_ptr(virtgpu * gpu, size_t size, size_t max_tensor_size) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + apir_buffer_context_t buffer_context; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_DEVICE_BUFFER_FROM_PTR); + + if (virtgpu_shmem_create(gpu, size, &buffer_context.shmem)) { + GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate %ldb of guest-host shared buffer", __func__, size); + } + + apir_encode_virtgpu_shmem_res_id(encoder, buffer_context.shmem.res_id); + + apir_encode_size_t(encoder, &size); + apir_encode_size_t(encoder, &max_tensor_size); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + apir_decode_apir_buffer_host_handle_t(decoder, &buffer_context.host_handle); + buffer_context.buft_host_handle = apir_decode_apir_buffer_type_host_handle(decoder); + + remote_call_finish(gpu, encoder, decoder); + + return buffer_context; +} + + +/* Inlined virtgpu-forward-buffer-type.cpp */ + +char * apir_buffer_type_get_name(virtgpu * gpu, apir_buffer_type_host_handle_t host_handle) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_TYPE_GET_NAME); + + apir_encode_apir_buffer_type_host_handle(encoder, host_handle); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + const size_t string_size = apir_decode_array_size_unchecked(decoder); + char * string = (char *) apir_decoder_alloc_array(sizeof(char), string_size); + if (!string) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: Could not allocate the device name buffer\n", __func__); + apir_decoder_set_fatal(decoder); + } + apir_decode_char_array(decoder, string, string_size); + + remote_call_finish(gpu, encoder, decoder); + + return string; +} + +size_t apir_buffer_type_get_alignment(virtgpu * gpu, apir_buffer_type_host_handle_t host_handle) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALIGNMENT); + + apir_encode_apir_buffer_type_host_handle(encoder, host_handle); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + size_t alignment; + apir_decode_size_t(decoder, &alignment); + + remote_call_finish(gpu, encoder, decoder); + + return alignment; +} + +size_t apir_buffer_type_get_max_size(virtgpu * gpu, apir_buffer_type_host_handle_t host_handle) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_TYPE_GET_MAX_SIZE); + + apir_encode_apir_buffer_type_host_handle(encoder, host_handle); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + size_t max_size; + apir_decode_size_t(decoder, &max_size); + + remote_call_finish(gpu, encoder, decoder); + + return max_size; +} + +apir_buffer_context_t apir_buffer_type_alloc_buffer(virtgpu * gpu, + apir_buffer_type_host_handle_t host_handle, + size_t size) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + apir_buffer_context_t buffer_context; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_TYPE_ALLOC_BUFFER); + + apir_encode_apir_buffer_type_host_handle(encoder, host_handle); + + apir_encode_size_t(encoder, &size); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + apir_decode_apir_buffer_host_handle_t(decoder, &buffer_context.host_handle); + + remote_call_finish(gpu, encoder, decoder); + + return buffer_context; +} + +size_t apir_buffer_type_get_alloc_size(virtgpu * gpu, + apir_buffer_type_host_handle_t host_handle, + const ggml_tensor * op) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_TYPE_GET_ALLOC_SIZE); + + apir_encode_apir_buffer_type_host_handle(encoder, host_handle); + + apir_encode_ggml_tensor_inline(encoder, op); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + size_t alloc_size; + apir_decode_size_t(decoder, &alloc_size); + + remote_call_finish(gpu, encoder, decoder); + + return alloc_size; +} + + +/* Inlined virtgpu-forward-buffer.cpp */ + +void * apir_buffer_get_base(virtgpu * gpu, apir_buffer_context_t * buffer_context) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_GET_BASE); + + apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + uintptr_t base; + apir_decode_uintptr_t(decoder, &base); + + remote_call_finish(gpu, encoder, decoder); + + return (void *) base; +} + +void apir_buffer_set_tensor(virtgpu * gpu, + apir_buffer_context_t * buffer_context, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_SET_TENSOR); + + apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle); + apir_encode_ggml_tensor(encoder, tensor); + + virtgpu_shmem temp_shmem; // Local storage for large buffers + virtgpu_shmem * shmem = &temp_shmem; + bool using_shared_shmem = false; + + if (size <= gpu->data_shmem.mmap_size) { + // Lock mutex before using shared data_shmem buffer + if (mtx_lock(&gpu->data_shmem_mutex) != thrd_success) { + GGML_ABORT(GGML_VIRTGPU "%s: Failed to lock data_shmem mutex", __func__); + } + using_shared_shmem = true; + shmem = &gpu->data_shmem; + + } else if (virtgpu_shmem_create(gpu, size, shmem)) { + GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate the guest-host shared buffer", __func__); + } + + memcpy(shmem->mmap_ptr, data, size); + apir_encode_virtgpu_shmem_res_id(encoder, shmem->res_id); + + apir_encode_size_t(encoder, &offset); + apir_encode_size_t(encoder, &size); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + remote_call_finish(gpu, encoder, decoder); + + // Unlock mutex before cleanup + if (using_shared_shmem) { + mtx_unlock(&gpu->data_shmem_mutex); + } else { + virtgpu_shmem_destroy(gpu, shmem); + } + + return; +} + +void apir_buffer_get_tensor(virtgpu * gpu, + apir_buffer_context_t * buffer_context, + const ggml_tensor * tensor, + void * data, + size_t offset, + size_t size) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_GET_TENSOR); + + apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle); + apir_encode_ggml_tensor(encoder, tensor); + + virtgpu_shmem temp_shmem; // Local storage for large buffers + virtgpu_shmem * shmem = &temp_shmem; + bool using_shared_shmem = false; + + if (size <= gpu->data_shmem.mmap_size) { + // Lock mutex before using shared data_shmem buffer + if (mtx_lock(&gpu->data_shmem_mutex) != thrd_success) { + GGML_ABORT(GGML_VIRTGPU "%s: Failed to lock data_shmem mutex", __func__); + } + using_shared_shmem = true; + shmem = &gpu->data_shmem; + + } else if (virtgpu_shmem_create(gpu, size, shmem)) { + GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate the guest-host shared buffer", __func__); + } + + apir_encode_virtgpu_shmem_res_id(encoder, shmem->res_id); + apir_encode_size_t(encoder, &offset); + apir_encode_size_t(encoder, &size); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + memcpy(data, shmem->mmap_ptr, size); + + remote_call_finish(gpu, encoder, decoder); + + // Unlock mutex before cleanup + if (using_shared_shmem) { + mtx_unlock(&gpu->data_shmem_mutex); + } else { + virtgpu_shmem_destroy(gpu, shmem); + } +} + +bool apir_buffer_cpy_tensor(virtgpu * gpu, + apir_buffer_context_t * buffer_context, + const ggml_tensor * src, + const ggml_tensor * dst) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_CPY_TENSOR); + + apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle); + apir_encode_ggml_tensor(encoder, src); + apir_encode_ggml_tensor(encoder, dst); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + bool ret_val; + apir_decode_bool_t(decoder, &ret_val); + + remote_call_finish(gpu, encoder, decoder); + + return ret_val; +} + +void apir_buffer_clear(virtgpu * gpu, apir_buffer_context_t * buffer_context, uint8_t value) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_CLEAR); + + apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle); + apir_encode_uint8_t(encoder, &value); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + remote_call_finish(gpu, encoder, decoder); +} + +void apir_buffer_free_buffer(virtgpu * gpu, apir_buffer_context_t * buffer_context) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_FREE_BUFFER); + + apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + remote_call_finish(gpu, encoder, decoder); +} + + +/* Inlined virtgpu-forward-backend.cpp */ + +static long long current_time_ms() { + timespec ts; + clock_gettime(CLOCK_REALTIME, &ts); // Use CLOCK_MONOTONIC for elapsed time + return (long long) ts.tv_sec * 1000000000LL + ts.tv_nsec; +} + +ggml_status apir_backend_graph_compute(virtgpu * gpu, ggml_cgraph * cgraph) { + apir_encoder * encoder; + apir_decoder * decoder; + ApirForwardReturnCode ret; + + REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BACKEND_GRAPH_COMPUTE); + + std::vector cgraph_data; + size_t cgraph_size = apir_serialize_ggml_cgraph(cgraph, cgraph_data); + + virtgpu_shmem temp_shmem; // Local storage for large buffers + virtgpu_shmem * shmem = &temp_shmem; + bool using_shared_shmem = false; + + if (cgraph_size <= gpu->data_shmem.mmap_size) { + // Lock mutex before using shared data_shmem buffer + if (mtx_lock(&gpu->data_shmem_mutex) != thrd_success) { + GGML_ABORT(GGML_VIRTGPU "%s: Failed to lock data_shmem mutex", __func__); + } + using_shared_shmem = true; + shmem = &gpu->data_shmem; + } else if (virtgpu_shmem_create(gpu, cgraph_size, shmem)) { + GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate the guest-host shared buffer", __func__); + } + + apir_encode_virtgpu_shmem_res_id(encoder, shmem->res_id); + + apir_encode_size_t(encoder, &cgraph_size); + + char * shmem_data = (char *) shmem->mmap_ptr; + apir_encoder secondary_enc = apir_new_encoder(shmem_data, cgraph_size); + + apir_encode_cgraph_data(&secondary_enc, cgraph_data); + + REMOTE_CALL(gpu, encoder, decoder, ret); + + ggml_status status = GGML_STATUS_ABORTED; + apir_decode_ggml_status(decoder, &status); + + remote_call_finish(gpu, encoder, decoder); + + // Unlock mutex before cleanup + if (using_shared_shmem) { + mtx_unlock(&gpu->data_shmem_mutex); + } else { + virtgpu_shmem_destroy(gpu, shmem); + } + + return status; +} + + +/* Inlined ggml-backend-buffer.cpp */ + +#define BUFFER_TO_GPU(name) ((ggml_backend_remoting_buffer_context *) (name)->context)->gpu + +static void * ggml_backend_remoting_buffer_get_base(ggml_backend_buffer_t buffer) { + ggml_backend_remoting_buffer_context * context = (ggml_backend_remoting_buffer_context *) buffer->context; + if (context->base) { + return context->base; + } + + context->base = apir_buffer_get_base(BUFFER_TO_GPU(buffer), BUFFER_TO_APIR_CONTEXT(buffer)); + + return context->base; +} + +static void ggml_backend_remoting_buffer_set_tensor(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size) { + virtgpu * gpu = BUFFER_TO_GPU(buffer); + + ggml_backend_remoting_buffer_context * context = BUFFER_TO_GGML_CONTEXT(buffer); + if (context->is_from_ptr) { + memcpy((char *) tensor->data + offset, data, size); + } else { + apir_buffer_set_tensor(gpu, BUFFER_TO_APIR_CONTEXT(buffer), tensor, data, offset, size); + } + + return; +} + +static void ggml_backend_remoting_buffer_get_tensor(ggml_backend_buffer_t buffer, + const ggml_tensor * tensor, + void * data, + size_t offset, + size_t size) { + virtgpu * gpu = BUFFER_TO_GPU(buffer); + ggml_backend_remoting_buffer_context * context = BUFFER_TO_GGML_CONTEXT(buffer); + if (context->is_from_ptr) { + memcpy(data, (const char *) tensor->data + offset, size); + } else { + apir_buffer_get_tensor(gpu, BUFFER_TO_APIR_CONTEXT(buffer), tensor, data, offset, size); + } +} + +static void ggml_backend_remoting_buffer_set_tensor_from_ptr(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size) { + UNUSED(buffer); + + memcpy((char *) tensor->data + offset, data, size); + + return; +} + +static void ggml_backend_remoting_buffer_get_tensor_from_ptr(ggml_backend_buffer_t buffer, + const ggml_tensor * tensor, + void * data, + size_t offset, + size_t size) { + UNUSED(buffer); + + memcpy(data, (const char *) tensor->data + offset, size); +} + +static bool ggml_backend_remoting_buffer_cpy_tensor(ggml_backend_buffer_t buffer, + const ggml_tensor * src, + ggml_tensor * dst) { + virtgpu * gpu = BUFFER_TO_GPU(buffer); + + bool ret = apir_buffer_cpy_tensor(gpu, BUFFER_TO_APIR_CONTEXT(buffer), src, dst); + + return ret; +} + +static void ggml_backend_remoting_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + virtgpu * gpu = BUFFER_TO_GPU(buffer); + + apir_buffer_clear(gpu, BUFFER_TO_APIR_CONTEXT(buffer), value); + + return; +} + +static void ggml_backend_remoting_buffer_free_buffer(ggml_backend_buffer_t buffer) { + virtgpu * gpu = BUFFER_TO_GPU(buffer); + + apir_buffer_free_buffer(gpu, BUFFER_TO_APIR_CONTEXT(buffer)); + + ggml_backend_remoting_buffer_context * context = BUFFER_TO_GGML_CONTEXT(buffer); + free(context); + buffer->context = NULL; +} + +const ggml_backend_buffer_i ggml_backend_remoting_buffer_interface = { + /* .free_buffer = */ ggml_backend_remoting_buffer_free_buffer, + /* .get_base = */ ggml_backend_remoting_buffer_get_base, + /* .init_tensor = */ NULL, + /* .memset_tensor = */ NULL, + /* .set_tensor = */ ggml_backend_remoting_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_remoting_buffer_get_tensor, + /* .set_tensor_2d = */ NULL, + /* .get_tensor_2d = */ NULL, + /* .cpy_tensor = */ ggml_backend_remoting_buffer_cpy_tensor, + /* .clear = */ ggml_backend_remoting_buffer_clear, + /* .reset = */ NULL, +}; + +const ggml_backend_buffer_i ggml_backend_remoting_buffer_from_ptr_interface = { + /* .free_buffer = */ ggml_backend_remoting_buffer_free_buffer, + /* .get_base = */ ggml_backend_remoting_buffer_get_base, + /* .init_tensor = */ NULL, + /* .memset_tensor = */ NULL, + /* .set_tensor = */ ggml_backend_remoting_buffer_set_tensor_from_ptr, + /* .get_tensor = */ ggml_backend_remoting_buffer_get_tensor_from_ptr, + /* .set_tensor_2d = */ NULL, + /* .get_tensor_2d = */ NULL, + /* .cpy_tensor = */ ggml_backend_remoting_buffer_cpy_tensor, + /* .clear = */ ggml_backend_remoting_buffer_clear, + /* .reset = */ NULL, +}; + + +/* Inlined ggml-backend-buffer-type.cpp */ + +static ggml_backend_buffer_t ggml_backend_remoting_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, + size_t size) { + virtgpu * gpu = BUFT_TO_GPU(buft); + + ggml_backend_remoting_buffer_context * context = (ggml_backend_remoting_buffer_context *) malloc(sizeof(*context)); + if (!context) { + GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate the buffer context ...", __func__); + } + + context->gpu = gpu; + + bool async__unused, host_buffer__unused, events__unused; + bool buffer_from_host_ptr; + apir_device_get_props(gpu, &async__unused, &host_buffer__unused, &buffer_from_host_ptr, &events__unused); + + if (buffer_from_host_ptr) { + context->apir_context = apir_device_buffer_from_ptr(gpu, size, size); + context->base = context->apir_context.shmem.mmap_ptr; + context->is_from_ptr = true; + } else { + context->apir_context = apir_buffer_type_alloc_buffer(gpu, gpu->cached_buffer_type.host_handle, size); + context->is_from_ptr = false; + context->base = NULL; + } + + ggml_backend_buffer_t buffer = + ggml_backend_buffer_init(buft, ggml_backend_remoting_buffer_interface, (void *) context, size); + + return buffer; +} + +static const char * ggml_backend_remoting_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + virtgpu * gpu = BUFT_TO_GPU(buft); + + // Return the prefixed name that was built once during initialization + return gpu->cached_buffer_type.name; +} + +static size_t ggml_backend_remoting_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + virtgpu * gpu = BUFT_TO_GPU(buft); + + return gpu->cached_buffer_type.alignment; +} + +static size_t ggml_backend_remoting_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { + virtgpu * gpu = BUFT_TO_GPU(buft); + + return gpu->cached_buffer_type.max_size; +} + +static size_t ggml_backend_remoting_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, + const ggml_tensor * tensor) { + virtgpu * gpu = BUFT_TO_GPU(buft); + + if (tensor->buffer == NULL || !tensor->buffer->context || + !buft->device->iface.supports_buft(buft->device, tensor->buffer->buft)) { + return ggml_nbytes(tensor); + } + + return apir_buffer_type_get_alloc_size(gpu, gpu->cached_buffer_type.host_handle, tensor); +} + +const ggml_backend_buffer_type_i ggml_backend_remoting_buffer_type_interface = { + /* .get_name = */ ggml_backend_remoting_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_remoting_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_remoting_buffer_type_get_alignment, + /* .get_max_size = */ ggml_backend_remoting_buffer_type_get_max_size, + /* .get_alloc_size = */ ggml_backend_remoting_buffer_type_get_alloc_size, + /* .is_host = */ NULL, +}; + +const ggml_backend_buffer_type_i ggml_backend_remoting_buffer_from_ptr_type_interface = { + /* .get_name = */ ggml_backend_remoting_buffer_type_get_name, + /* .alloc_buffer = */ NULL, + /* .get_alignment = */ ggml_backend_remoting_buffer_type_get_alignment, + /* .get_max_size = */ ggml_backend_remoting_buffer_type_get_max_size, + /* .get_alloc_size = */ ggml_backend_remoting_buffer_type_get_alloc_size, + /* .is_host = */ NULL, +}; + + +/* Inlined ggml-backend-device.cpp */ + +static const char * ggml_backend_remoting_device_get_name(ggml_backend_dev_t dev) { + virtgpu * gpu = DEV_TO_GPU(dev); + + // Return the prefixed name that was built once during initialization + return gpu->cached_device_info.name; +} + +static const char * ggml_backend_remoting_device_get_description(ggml_backend_dev_t dev) { + virtgpu * gpu = DEV_TO_GPU(dev); + + // Return the pre-cached description from the virtgpu structure + return gpu->cached_device_info.description; +} + +static enum ggml_backend_dev_type ggml_backend_remoting_device_get_type(ggml_backend_dev_t dev) { + virtgpu * gpu = DEV_TO_GPU(dev); + + return (enum ggml_backend_dev_type) gpu->cached_device_info.type; +} + +static void ggml_backend_remoting_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { + virtgpu * gpu = DEV_TO_GPU(dev); + + *free = gpu->cached_device_info.memory_free; + *total = gpu->cached_device_info.memory_total; +} + +static bool ggml_backend_remoting_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { +#if USE_ALWAYS_TRUE_SUPPORTS_OP == 1 + /* ggml-rpc cheats it like this */ + /* with the current implementation of serialize_tensor, the src/view aren't properly passed */ + UNUSED(dev); + UNUSED(op); + + return true; +#else + virtgpu * gpu = DEV_TO_GPU(dev); + + return apir_device_supports_op(gpu, op); +#endif +} + +static bool ggml_backend_remoting_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { + bool supported = buft->device == dev; + + return supported; +} + +static bool ggml_backend_remoting_device_offload_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + UNUSED(dev); + UNUSED(op); + + return false; +} + +static void ggml_backend_remoting_device_get_props(ggml_backend_dev_t dev, ggml_backend_dev_props * props) { + props->name = ggml_backend_remoting_device_get_name(dev); + props->description = ggml_backend_remoting_device_get_description(dev); + props->type = ggml_backend_remoting_device_get_type(dev); + ggml_backend_remoting_device_get_memory(dev, &props->memory_free, &props->memory_total); + + virtgpu * gpu = DEV_TO_GPU(dev); + apir_device_get_props(gpu, &props->caps.async, &props->caps.host_buffer, &props->caps.buffer_from_host_ptr, + &props->caps.events); + + props->caps.buffer_from_host_ptr = false; + props->caps.async = false; + props->caps.events = false; +} + +ggml_backend_buffer_type_t ggml_backend_remoting_device_get_buffer_type(ggml_backend_dev_t dev) { + virtgpu * gpu = DEV_TO_GPU(dev); + + static std::atomic initialized = false; + static ggml_backend_buffer_type buft; + + if (!initialized) { + static std::mutex mutex; + std::lock_guard lock(mutex); + + if (!initialized) { + buft = { + /* .iface = */ ggml_backend_remoting_buffer_type_interface, + /* .device = */ dev, + /* .context = */ (void *) gpu->cached_buffer_type.host_handle, + }; + initialized = true; + } + } + + return &buft; +} + +static ggml_backend_buffer_type_t ggml_backend_remoting_device_get_buffer_from_ptr_type(ggml_backend_dev_t dev) { + virtgpu * gpu = DEV_TO_GPU(dev); + + static std::atomic initialized = false; + static ggml_backend_buffer_type buft; + + if (!initialized) { + static std::mutex mutex; + std::lock_guard lock(mutex); + + if (!initialized) { + buft = { + /* .iface = */ ggml_backend_remoting_buffer_from_ptr_type_interface, + /* .device = */ dev, + /* .context = */ (void *) gpu->cached_buffer_type.host_handle, + }; + initialized = true; + } + } + + return &buft; +} + +static ggml_backend_buffer_t ggml_backend_remoting_device_buffer_from_ptr(ggml_backend_dev_t dev, + void * ptr, + size_t size, + size_t max_tensor_size) { + virtgpu * gpu = DEV_TO_GPU(dev); + + ggml_backend_remoting_buffer_context * context = (ggml_backend_remoting_buffer_context *) malloc(sizeof(*context)); + if (!context) { + GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate the buffer context ...", __func__); + } + + context->gpu = gpu; + context->apir_context = apir_device_buffer_from_ptr(gpu, size, max_tensor_size); + context->base = ptr; + context->is_from_ptr = true; + + ggml_backend_buffer_t buffer = + ggml_backend_buffer_init(ggml_backend_remoting_device_get_buffer_from_ptr_type(dev), + ggml_backend_remoting_buffer_from_ptr_interface, (void *) context, size); + + return buffer; +} + +const ggml_backend_device_i ggml_backend_remoting_device_interface = { + /* .get_name = */ ggml_backend_remoting_device_get_name, + /* .get_description = */ ggml_backend_remoting_device_get_description, + /* .get_memory = */ ggml_backend_remoting_device_get_memory, + /* .get_type = */ ggml_backend_remoting_device_get_type, + /* .get_props = */ ggml_backend_remoting_device_get_props, + /* .init_backend = */ ggml_backend_remoting_device_init, + /* .get_buffer_type = */ ggml_backend_remoting_device_get_buffer_type, + /* .get_host_buffer_type = */ NULL, + /* .buffer_from_host_ptr = */ ggml_backend_remoting_device_buffer_from_ptr, + /* .supports_op = */ ggml_backend_remoting_device_supports_op, + /* .supports_buft = */ ggml_backend_remoting_device_supports_buft, + /* .offload_op = */ ggml_backend_remoting_device_offload_op, + /* .event_new = */ NULL, + /* .event_free = */ NULL, + /* .event_synchronize = */ NULL, +}; + + +/* Inlined ggml-backend-reg.cpp.inc */ +#include "ggml-virtgpu.h.inc" + +#include +#include + +void ggml_virtgpu_cleanup(virtgpu * gpu); + +static virtgpu * apir_initialize() { + static virtgpu * gpu = NULL; + static std::atomic initialized = false; + + if (initialized) { + // fast track + return gpu; + } + + { + static std::mutex mutex; + std::lock_guard lock(mutex); + + if (initialized) { + // thread safe + return gpu; + } + + gpu = create_virtgpu(); + if (!gpu) { + initialized = true; + return NULL; + } + + // Pre-fetch and cache all device information, it will not change + gpu->cached_device_info.description = apir_device_get_description(gpu); + if (!gpu->cached_device_info.description) { + GGML_ABORT(GGML_VIRTGPU "%s: failed to initialize the virtgpu device description", __func__); + } + gpu->cached_device_info.device_count = apir_device_get_count(gpu); + gpu->cached_device_info.type = apir_device_get_type(gpu); + + { + // Get the remote name and create prefixed version + char * rmt_device_name = apir_device_get_name(gpu); + if (!rmt_device_name) { + GGML_ABORT(GGML_VIRTGPU "%s: failed to get the virtgpu device name", __func__); + } + + size_t device_name_len = strlen(rmt_device_name) + 11; // "[virtgpu] " + null terminator + gpu->cached_device_info.name = (char *) malloc(device_name_len); + if (!gpu->cached_device_info.name) { + free(rmt_device_name); + GGML_ABORT(GGML_VIRTGPU "%s: failed to allocate memory for prefixed device name", __func__); + } + snprintf(gpu->cached_device_info.name, device_name_len, "[virtgpu] %s", rmt_device_name); + free(rmt_device_name); + } + + apir_device_get_memory(gpu, &gpu->cached_device_info.memory_free, &gpu->cached_device_info.memory_total); + + apir_buffer_type_host_handle_t buft_host_handle = apir_device_get_buffer_type(gpu); + gpu->cached_buffer_type.host_handle = buft_host_handle; + { + // Get the remote name and create prefixed version + char * rmt_name = apir_buffer_type_get_name(gpu, buft_host_handle); + if (!rmt_name) { + GGML_ABORT(GGML_VIRTGPU "%s: failed to get the virtgpu buffer type name", __func__); + } + + size_t prefixed_len = strlen(rmt_name) + 11; // "[virtgpu] " + null terminator + gpu->cached_buffer_type.name = (char *) malloc(prefixed_len); + if (!gpu->cached_buffer_type.name) { + free(rmt_name); + GGML_ABORT(GGML_VIRTGPU "%s: failed to allocate memory for prefixed buffer type name", __func__); + } + snprintf(gpu->cached_buffer_type.name, prefixed_len, "[virtgpu] %s", rmt_name); + free(rmt_name); + } + + gpu->cached_buffer_type.alignment = apir_buffer_type_get_alignment(gpu, buft_host_handle); + gpu->cached_buffer_type.max_size = apir_buffer_type_get_max_size(gpu, buft_host_handle); + + initialized = true; + } + + return gpu; +} + +static int ggml_backend_remoting_get_device_count() { + virtgpu * gpu = apir_initialize(); + if (!gpu) { + return 0; + } + + return gpu->cached_device_info.device_count; +} + +static size_t ggml_backend_remoting_reg_get_device_count(ggml_backend_reg_t reg) { + UNUSED(reg); + + return ggml_backend_remoting_get_device_count(); +} + +static std::vector devices; + +ggml_backend_dev_t ggml_backend_remoting_get_device(size_t device) { + GGML_ASSERT(device < devices.size()); + return devices[device]; +} + +static void ggml_backend_remoting_reg_init_devices(ggml_backend_reg_t reg) { + if (devices.size() > 0) { + GGML_LOG_INFO(GGML_VIRTGPU "%s: already initialized\n", __func__); + return; + } + + virtgpu * gpu = apir_initialize(); + if (!gpu) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: apir_initialize failed\n", __func__); + return; + } + + static std::atomic initialized = false; + + if (initialized) { + return; // fast track + } + + { + static std::mutex mutex; + std::lock_guard lock(mutex); + if (!initialized) { + for (int i = 0; i < ggml_backend_remoting_get_device_count(); i++) { + ggml_backend_remoting_device_context * ctx = new ggml_backend_remoting_device_context; + char desc[256] = "ggml-virtgpu API Remoting device"; + + ctx->device = i; + ctx->name = GGML_VIRTGPU_NAME + std::to_string(i); + ctx->description = desc; + ctx->gpu = gpu; + + ggml_backend_dev_t dev = new ggml_backend_device{ + /* .iface = */ ggml_backend_remoting_device_interface, + /* .reg = */ reg, + /* .context = */ ctx, + }; + devices.push_back(dev); + } + initialized = true; + } + } +} + +static ggml_backend_dev_t ggml_backend_remoting_reg_get_device(ggml_backend_reg_t reg, size_t device) { + UNUSED(reg); + + return ggml_backend_remoting_get_device(device); +} + +static const char * ggml_backend_remoting_reg_get_name(ggml_backend_reg_t reg) { + UNUSED(reg); + + return GGML_VIRTGPU_NAME; +} + +static const ggml_backend_reg_i ggml_backend_remoting_reg_i = { + /* .get_name = */ ggml_backend_remoting_reg_get_name, + /* .get_device_count = */ ggml_backend_remoting_reg_get_device_count, + /* .get_device = */ ggml_backend_remoting_reg_get_device, + /* .get_proc_address = */ NULL, +}; + +ggml_backend_reg_t ggml_backend_virtgpu_reg() { + virtgpu * gpu = apir_initialize(); + if (!gpu) { + GGML_LOG_ERROR(GGML_VIRTGPU "%s: virtgpu_apir_initialize failed\n", __func__); + } + + static ggml_backend_reg reg = { + /* .api_version = */ GGML_BACKEND_API_VERSION, + /* .iface = */ ggml_backend_remoting_reg_i, + /* .context = */ gpu, + }; + + static bool initialized = false; + if (initialized) { + return ® + } + initialized = true; + + ggml_backend_remoting_reg_init_devices(®); + + return ® +} + +// public function, not exposed in the GGML interface at the moment +void ggml_virtgpu_cleanup(virtgpu * gpu) { + if (gpu->cached_device_info.name) { + free(gpu->cached_device_info.name); + gpu->cached_device_info.name = NULL; + } + if (gpu->cached_device_info.description) { + free(gpu->cached_device_info.description); + gpu->cached_device_info.description = NULL; + } + if (gpu->cached_buffer_type.name) { + free(gpu->cached_buffer_type.name); + gpu->cached_buffer_type.name = NULL; + } + + mtx_destroy(&gpu->data_shmem_mutex); +} + +GGML_BACKEND_DL_IMPL(ggml_backend_virtgpu_reg) + + +/* Inlined ggml-backend.cpp.inc */ +#include "../../include/ggml-virtgpu.h.inc" + +static const char * ggml_backend_remoting_get_name(ggml_backend_t backend) { + UNUSED(backend); + + return "API Remoting backend"; +} + +static void ggml_backend_remoting_free(ggml_backend_t backend) { + delete backend; +} + +static ggml_status ggml_backend_remoting_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { + virtgpu * gpu = DEV_TO_GPU(backend->device); + + return apir_backend_graph_compute(gpu, cgraph); +} + +static void ggml_backend_remoting_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { + virtgpu * gpu = DEV_TO_GPU(backend->device); +#if true + UNUSED(gpu); + UNUSED(cgraph); +#else + // not working yet + + apir_backend_graph_optimize(gpu, cgraph); +#endif +} + +static ggml_backend_i ggml_backend_remoting_interface = { + /* .get_name = */ ggml_backend_remoting_get_name, + /* .free = */ ggml_backend_remoting_free, + /* .set_tensor_async = */ NULL, // ggml_backend_remoting_set_tensor_async, + /* .get_tensor_async = */ NULL, // ggml_backend_remoting_get_tensor_async, + /* .set_tensor_2d_async = */ NULL, + /* .get_tensor_2d_async = */ NULL, + /* .cpy_tensor_async = */ NULL, // ggml_backend_remoting_cpy_tensor_async, + /* .synchronize = */ NULL, // ggml_backend_remoting_synchronize, + /* .graph_plan_create = */ NULL, + /* .graph_plan_free = */ NULL, + /* .graph_plan_update = */ NULL, + /* .graph_plan_compute = */ NULL, + /* .graph_compute = */ ggml_backend_remoting_graph_compute, + /* .event_record = */ NULL, + /* .event_wait = */ NULL, + /* .graph_optimize = */ ggml_backend_remoting_graph_optimize, +}; + +static ggml_guid_t ggml_backend_remoting_guid() { + static ggml_guid guid = { 0xb8, 0xf7, 0x4f, 0x86, 0x14, 0x03, 0x86, 0x02, + 0x91, 0xc8, 0xdd, 0xe9, 0x02, 0x3f, 0xc0, 0x2b }; + + return &guid; +} + +ggml_backend_t ggml_backend_remoting_device_init(ggml_backend_dev_t dev, const char * params) { + UNUSED(params); + + ggml_backend_remoting_device_context * ctx = (ggml_backend_remoting_device_context *) dev->context; + + ggml_backend_t remoting_backend = new ggml_backend{ + /* .guid = */ ggml_backend_remoting_guid(), + /* .interface = */ ggml_backend_remoting_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_virtgpu_reg(), ctx->device), + /* .context = */ ctx, + }; + + return remoting_backend; +} + + +/* Inlined apir_cs_ggml-rpc-front.cpp */ +#include "ggml-backend-impl-defs.inc" +#include "ggml-impl-defs.inc" + +#include +#include +#include +#include + +apir_rpc_tensor apir_serialize_tensor(const ggml_tensor * tensor) { + apir_rpc_tensor result; + result.id = reinterpret_cast(tensor); + result.type = tensor->type; + if (tensor->buffer) { + ggml_backend_buffer_t buffer = tensor->buffer; + + result.buffer = BUFFER_TO_HOST_HANDLE(buffer); + } else { + result.buffer = 0; + } + for (uint32_t i = 0; i < GGML_MAX_DIMS; i++) { + result.ne[i] = tensor->ne[i]; + result.nb[i] = tensor->nb[i]; + } + result.op = tensor->op; + for (uint32_t i = 0; i < GGML_MAX_OP_PARAMS / sizeof(int32_t); i++) { + result.op_params[i] = tensor->op_params[i]; + } + result.flags = tensor->flags; + for (uint32_t i = 0; i < GGML_MAX_SRC; i++) { + result.src[i] = reinterpret_cast(tensor->src[i]); + } + result.view_src = reinterpret_cast(tensor->view_src); + result.view_offs = tensor->view_offs; + result.data = reinterpret_cast(tensor->data); + if (tensor->data) { + if (!tensor->buffer) { + GGML_ABORT("%s: tensor has data but not buffer", __func__); + } + // tensor->data is serialized as an offset to the buffer base address + result.data -= reinterpret_cast(BUFFER_TO_GGML_CONTEXT(tensor->buffer)->base); + } + snprintf(result.name, GGML_MAX_NAME, "%s", tensor->name); + return result; +} + +void apir_add_tensor(ggml_tensor * tensor, + std::vector & tensors, + std::unordered_set & visited) { + if (tensor == nullptr) { + return; + } + if (visited.find(tensor) != visited.end()) { + return; + } + visited.insert(tensor); + for (int i = 0; i < GGML_MAX_SRC; i++) { + apir_add_tensor(tensor->src[i], tensors, visited); + } + apir_add_tensor(tensor->view_src, tensors, visited); + tensors.push_back(apir_serialize_tensor(tensor)); +} + +void apir_serialize_graph(const ggml_cgraph * cgraph, std::vector & output) { + uint32_t n_nodes = cgraph->n_nodes; + std::vector tensors; + std::unordered_set visited; + for (uint32_t i = 0; i < n_nodes; i++) { + apir_add_tensor(cgraph->nodes[i], tensors, visited); + } + // serialization format: + // | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(apir_rpc_tensor)) | + uint32_t n_tensors = tensors.size(); + int output_size = + sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(apir_rpc_tensor); + output.resize(output_size, 0); + memcpy(output.data(), &n_nodes, sizeof(n_nodes)); + for (uint32_t i = 0; i < n_nodes; i++) { + memcpy(output.data() + sizeof(n_nodes) + i * sizeof(uint64_t), &cgraph->nodes[i], sizeof(uint64_t)); + } + uint32_t * out_ntensors = (uint32_t *) (output.data() + sizeof(n_nodes) + n_nodes * sizeof(uint64_t)); + *out_ntensors = n_tensors; + apir_rpc_tensor * out_tensors = + (apir_rpc_tensor *) (output.data() + sizeof(n_nodes) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t)); + memcpy(out_tensors, tensors.data(), n_tensors * sizeof(apir_rpc_tensor)); +} diff --git a/ggml/src/ggml-virtgpu/virtgpu.h b/ggml/src/ggml-virtgpu/virtgpu.h deleted file mode 100644 index 6b8de583893e..000000000000 --- a/ggml/src/ggml-virtgpu/virtgpu.h +++ /dev/null @@ -1,115 +0,0 @@ -#pragma once - -// clang-format off -#include "virtgpu-utils.h" -#include "virtgpu-shm.h" -#include "virtgpu-apir.h" - -#include "backend/shared/api_remoting.h" -#include "backend/shared/apir_cs.h" - -#include -#include -#include -#include -#include -#include -#include - -#include - -#define VIRGL_RENDERER_UNSTABLE_APIS 1 -#include "apir_hw.h" -#include -#include "venus_hw.h" -// clang-format on - -#ifndef VIRTGPU_DRM_CAPSET_APIR -// Will be defined include/drm/virtgpu_drm.h when -// https://gitlab.freedesktop.org/virgl/virglrenderer/-/merge_requests/1590/diffs -// is merged -# define VIRTGPU_DRM_CAPSET_APIR 10 -#endif - -// Mesa/Virlgrenderer Venus internal. Only necessary during the -// Venus->APIR transition in Virglrenderer -#define VENUS_COMMAND_TYPE_LENGTH 331 - -#ifndef VIRTGPU_DRM_CAPSET_VENUS // only available with Linux >= v6.16 -# define VIRTGPU_DRM_CAPSET_VENUS 4 -#endif - -typedef uint32_t virgl_renderer_capset; - -/* from src/virtio/vulkan/vn_renderer_virtgpu.c */ -#define VIRTGPU_PCI_VENDOR_ID 0x1af4 -#define VIRTGPU_PCI_DEVICE_ID 0x1050 -#define VIRTGPU_BLOB_MEM_GUEST_VRAM 0x0004 -#define VIRTGPU_PARAM_GUEST_VRAM 9 - -#define SHMEM_DATA_SIZE 0x1830000 // 24MiB -#define SHMEM_REPLY_SIZE 0x4000 - -#define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0])) - -enum virt_gpu_result_t { - APIR_SUCCESS = 0, - APIR_ERROR_INITIALIZATION_FAILED = -1, -}; - -#define PRINTFLIKE(f, a) __attribute__((format(__printf__, f, a))) - -struct virtgpu { - bool use_apir_capset; - - int fd; - - struct { - virgl_renderer_capset id; - uint32_t version; - virgl_renderer_capset_apir data; - } capset; - - util_sparse_array shmem_array; - - /* APIR communication pages */ - virtgpu_shmem reply_shmem; - virtgpu_shmem data_shmem; - - /* Mutex to protect shared data_shmem buffer from concurrent access */ - mtx_t data_shmem_mutex; - - /* Cached device information to prevent memory leaks and race conditions */ - struct { - char * description; - char * name; - int32_t device_count; - uint32_t type; - size_t memory_free; - size_t memory_total; - } cached_device_info; - - /* Cached buffer type information to prevent memory leaks and race conditions */ - struct { - apir_buffer_type_host_handle_t host_handle; - char * name; - size_t alignment; - size_t max_size; - } cached_buffer_type; -}; - -static inline int virtgpu_ioctl(virtgpu * gpu, unsigned long request, void * args) { - return drmIoctl(gpu->fd, request, args); -} - -virtgpu * create_virtgpu(); - -apir_encoder * remote_call_prepare(virtgpu * gpu, ApirCommandType apir_cmd_type, int32_t cmd_flags); - -uint32_t remote_call(virtgpu * gpu, - apir_encoder * enc, - apir_decoder ** dec, - float max_wait_ms, - long long * call_duration_ns); - -void remote_call_finish(virtgpu * gpu, apir_encoder * enc, apir_decoder * dec); diff --git a/ggml/src/ggml-vulkan/CMakeLists.txt b/ggml/src/ggml-vulkan/CMakeLists.txt index 715a263a6d09..183d32b9de45 100644 --- a/ggml/src/ggml-vulkan/CMakeLists.txt +++ b/ggml/src/ggml-vulkan/CMakeLists.txt @@ -54,9 +54,10 @@ endfunction() if (Vulkan_FOUND) message(STATUS "Vulkan found") + set_source_files_properties(ggml-vulkan.cpp.inc PROPERTIES LANGUAGE CXX) ggml_add_backend_library(ggml-vulkan - ggml-vulkan.cpp - ../../include/ggml-vulkan.h + ggml-vulkan.cpp.inc + ../../include/ggml-vulkan.h.inc ) set(VULKAN_SHADER_GEN_CMAKE_ARGS "") diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp.inc similarity index 99% rename from ggml/src/ggml-vulkan/ggml-vulkan.cpp rename to ggml/src/ggml-vulkan/ggml-vulkan.cpp.inc index c2f1883328f0..0d897b74d3c3 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp.inc @@ -1,8 +1,8 @@ -#include "ggml-vulkan.h" +#include "ggml-vulkan.h.inc" #include #if defined(GGML_VULKAN_RUN_TESTS) || defined(GGML_VULKAN_CHECK_RESULTS) #include -#include "ggml-cpu.h" +#include "ggml-cpu.h.inc" #endif // See https://github.com/KhronosGroup/Vulkan-Hpp?tab=readme-ov-file#extensions--per-device-function-pointers- @@ -76,8 +76,8 @@ DispatchLoaderDynamic & ggml_vk_default_dispatcher(); #define YIELD() #endif -#include "ggml-impl.h" -#include "ggml-backend-impl.h" +#include "ggml-impl-defs.inc" +#include "ggml-backend-impl-defs.inc" #include "ggml-vulkan-shaders.hpp" @@ -15946,7 +15946,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32; case GGML_OP_RWKV_WKV6: case GGML_OP_RWKV_WKV7: - return true; // all inputs are contiguous, see ggml.c + return true; // all inputs are contiguous, see ggml.c.inc case GGML_OP_GATED_DELTA_NET: { const uint32_t S_v = op->src[2]->ne[0]; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt b/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt index e1f613fb4f68..678ab00acdb6 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt +++ b/ggml/src/ggml-vulkan/vulkan-shaders/CMakeLists.txt @@ -25,7 +25,8 @@ if (GGML_VULKAN_SHADER_DEBUG_INFO) endif() set(TARGET vulkan-shaders-gen) -add_executable(${TARGET} vulkan-shaders-gen.cpp) +set_source_files_properties(vulkan-shaders-gen.cpp.inc PROPERTIES LANGUAGE CXX) +add_executable(${TARGET} vulkan-shaders-gen.cpp.inc) install(TARGETS ${TARGET} RUNTIME) target_compile_features(${TARGET} PRIVATE cxx_std_17) target_link_libraries(vulkan-shaders-gen PUBLIC Threads::Threads) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp b/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp index f7d12a8dda6c..4795c3abb38a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/upscale.comp @@ -17,7 +17,7 @@ layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; -// from ggml.h: enum ggml_scale_mode, enum ggml_scale_flag +// from ggml.h.inc: enum ggml_scale_mode, enum ggml_scale_flag #define NEAREST 0 #define BILINEAR 1 #define BICUBIC 2 diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp.inc similarity index 100% rename from ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp rename to ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp.inc diff --git a/ggml/src/ggml-webgpu/CMakeLists.txt b/ggml/src/ggml-webgpu/CMakeLists.txt index 3ccce58aa39e..1d2930ede005 100644 --- a/ggml/src/ggml-webgpu/CMakeLists.txt +++ b/ggml/src/ggml-webgpu/CMakeLists.txt @@ -28,10 +28,11 @@ add_custom_command( add_custom_target(generate_shaders DEPENDS ${SHADER_HEADER}) +set_source_files_properties(ggml-webgpu.cpp.inc PROPERTIES LANGUAGE CXX) ggml_add_backend_library(ggml-webgpu - ggml-webgpu.cpp + ggml-webgpu.cpp.inc ${SHADER_HEADER} - ../../include/ggml-webgpu.h + ../../include/ggml-webgpu.h.inc ) add_dependencies(ggml-webgpu generate_shaders) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp deleted file mode 100644 index cff93b8d1705..000000000000 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ /dev/null @@ -1,3068 +0,0 @@ -#ifndef GGML_WEBGPU_SHADER_LIB_HPP -#define GGML_WEBGPU_SHADER_LIB_HPP - -#include "ggml-impl.h" -#include "ggml-wgsl-shaders.hpp" -#include "ggml.h" -#include "pre_wgsl.hpp" - -#include - -#include -#include -#include -#include -#include - -#define GGML_WEBGPU_F16_SIZE_BYTES 2 -#define GGML_WEBGPU_F32_SIZE_BYTES 4 -#define GGML_WEBGPU_I32_SIZE_BYTES 4 -#define GGML_WEBGPU_FLASH_ATTN_PREFERRED_KV_SG_TILES 8u -#define GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE 128u -// Matches GGML_PAD(..., 256) in src/llama-context.cpp for KV cache sizing. -#define GGML_WEBGPU_KV_SEQ_PAD 256u - -#define GGML_WEBGPU_ARGSORT_MERGE_MAX_WG_SIZE 512u - -// Matrix multiplication parameters - -// Register tiling parameters -#define WEBGPU_MUL_MAT_TILE_M 4 -#define WEBGPU_MUL_MAT_TILE_N 4 -#define WEBGPU_MUL_MAT_WG_SIZE_M 8 -#define WEBGPU_MUL_MAT_WG_SIZE_N 8 -#define WEBGPU_MUL_MAT_REG_TILE_K_FLOAT 8 -#define WEBGPU_MUL_MAT_REG_TILE_K_QUANT 32 - -// Subgroup matrix parameters -// The number of subgroups in the M dimension -#define WEBGPU_MUL_MAT_SUBGROUP_M 2 -// The number of subgroups in the N dimension -#define WEBGPU_MUL_MAT_SUBGROUP_N 4 -// The number of subgroup matrices each subgroup accumulates over -#define WEBGPU_MUL_MAT_SUBGROUP_MATRIX_M 4 -#define WEBGPU_MUL_MAT_SUBGROUP_MATRIX_N 2 -#define WEBGPU_MUL_MAT_SUBGROUP_TILE_K_FLOAT 32 -#define WEBGPU_MUL_MAT_SUBGROUP_TILE_K_QUANT 32 - -// Matrix-vector multiplication parameters -#define WEBGPU_MUL_MAT_VEC_WG_SIZE 256 - -#define WEBGPU_MUL_MAT_VEC_FLOAT_OUTPUTS_PER_WG 4 -#define WEBGPU_MUL_MAT_VEC_LEGACY_Q_OUTPUTS_PER_WG 4 -#define WEBGPU_MUL_MAT_VEC_K_Q_OUTPUTS_PER_WG 4 - -// default size for legacy matrix multiplication -#define WEBGPU_MUL_MAT_WG_SIZE 256 - -// Same hash combine function as in boost -template inline void ggml_webgpu_hash_combine(size_t & seed, const T & value) { - seed ^= std::hash{}(value) + 0x9e3779b9 + (seed << 6) + (seed >> 2); -} - -// Calculates base address of a tensor ignoring the fake base pointer -inline uintptr_t ggml_webgpu_tensor_addr(const ggml_tensor * tensor) { - const ggml_tensor * base_tensor = tensor->view_src ? tensor->view_src : tensor; - return (uintptr_t) base_tensor->data + tensor->view_offs; -} - -inline bool ggml_webgpu_tensor_equal(const ggml_tensor * a, const ggml_tensor * b) { - return a->buffer == b->buffer && ggml_webgpu_tensor_addr(a) == ggml_webgpu_tensor_addr(b); -} - -inline bool ggml_webgpu_tensor_overlap(const ggml_tensor * a, const ggml_tensor * b) { - return a->buffer == b->buffer && ggml_webgpu_tensor_addr(a) < ggml_webgpu_tensor_addr(b) + ggml_nbytes(b) && - ggml_webgpu_tensor_addr(b) < ggml_webgpu_tensor_addr(a) + ggml_nbytes(a); -} - -struct ggml_webgpu_shader_lib_context { - ggml_tensor * src0; - ggml_tensor * src1; - ggml_tensor * src2; - ggml_tensor * src3; - ggml_tensor * src4; - ggml_tensor * src5; - ggml_tensor * dst; - - uint32_t max_wg_size; - size_t wg_mem_limit_bytes = 0; - bool supports_subgroups = false; - bool supports_subgroup_matrix = false; - uint32_t sg_mat_m = 0; - uint32_t sg_mat_n = 0; - uint32_t sg_mat_k = 0; - uint32_t max_subgroup_size = 0; -}; - -struct webgpu_pipeline { - wgpu::ComputePipeline pipeline; - std::string name; - std::shared_ptr context = nullptr; -}; - -struct ggml_webgpu_generic_shader_decisions { - uint32_t wg_size = 0; - bool inplace = false; -}; - -struct ggml_webgpu_binary_shader_decisions { - uint32_t wg_size = 0; - bool inplace = false; - bool overlap = false; - bool src_overlap = false; -}; - -struct ggml_webgpu_processed_shader { - std::string wgsl; - std::string variant; - std::shared_ptr decisions; -}; - -struct ggml_webgpu_ssm_conv_shader_decisions { - uint32_t block_size; - uint32_t tokens_per_wg; -}; - -struct ggml_webgpu_ssm_scan_pipeline_key { - int type; - int d_state; - bool xbc_overlap; - - bool operator==(const ggml_webgpu_ssm_scan_pipeline_key & other) const { - return type == other.type && d_state == other.d_state && xbc_overlap == other.xbc_overlap; - } -}; - -struct ggml_webgpu_ssm_scan_pipeline_key_hash { - size_t operator()(const ggml_webgpu_ssm_scan_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.type); - ggml_webgpu_hash_combine(seed, key.d_state); - ggml_webgpu_hash_combine(seed, key.xbc_overlap); - return seed; - } -}; - -struct ggml_webgpu_ssm_scan_shader_decisions { - uint32_t wg_size; - uint32_t tokens_per_tile; - bool xbc_overlap = false; -}; - -/** Argsort **/ - -struct ggml_webgpu_argsort_shader_lib_context { - uint32_t max_wg_size; - size_t wg_mem_limit_bytes; - int32_t order; -}; - -/** Set Rows **/ - -struct ggml_webgpu_set_rows_pipeline_key { - int dst_type; - int vec4; - int i64_idx; - - bool operator==(const ggml_webgpu_set_rows_pipeline_key & other) const { - return dst_type == other.dst_type && vec4 == other.vec4 && i64_idx == other.i64_idx; - } -}; - -struct ggml_webgpu_set_rows_pipeline_key_hash { - size_t operator()(const ggml_webgpu_set_rows_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.dst_type); - ggml_webgpu_hash_combine(seed, key.vec4); - ggml_webgpu_hash_combine(seed, key.i64_idx); - return seed; - } -}; - -struct ggml_webgpu_set_rows_shader_decisions { - bool vec4; - bool i64_idx; - uint32_t wg_size; -}; - -/** Set **/ - -struct ggml_webgpu_set_pipeline_key { - ggml_type type; - bool inplace; - - bool operator==(const ggml_webgpu_set_pipeline_key & other) const { - return type == other.type && inplace == other.inplace; - } -}; - -struct ggml_webgpu_set_pipeline_key_hash { - size_t operator()(const ggml_webgpu_set_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.type); - ggml_webgpu_hash_combine(seed, key.inplace); - return seed; - } -}; - -/** Get Rows **/ - -struct ggml_webgpu_get_rows_pipeline_key { - ggml_type src_type; - int vectorized; - - bool operator==(const ggml_webgpu_get_rows_pipeline_key & other) const { - return src_type == other.src_type && vectorized == other.vectorized; - } -}; - -struct ggml_webgpu_get_rows_pipeline_key_hash { - size_t operator()(const ggml_webgpu_get_rows_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.src_type); - ggml_webgpu_hash_combine(seed, key.vectorized); - return seed; - } -}; - -/** Row Norm **/ - -struct ggml_webgpu_row_norm_pipeline_key { - ggml_op op; - bool inplace; - - bool operator==(const ggml_webgpu_row_norm_pipeline_key & other) const { - return op == other.op && inplace == other.inplace; - } -}; - -struct ggml_webgpu_row_norm_pipeline_key_hash { - size_t operator()(const ggml_webgpu_row_norm_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.op); - ggml_webgpu_hash_combine(seed, key.inplace); - return seed; - } -}; - -/** RMS_NORM + MUL **/ - -struct ggml_webgpu_rms_norm_mul_pipeline_key { - bool inplace; // rn_src == dst - bool overlap; // mul_src == dst - bool src_overlap; // rn_src == mul_src - - bool operator==(const ggml_webgpu_rms_norm_mul_pipeline_key & other) const { - return inplace == other.inplace && overlap == other.overlap && src_overlap == other.src_overlap; - } -}; - -struct ggml_webgpu_rms_norm_mul_pipeline_key_hash { - size_t operator()(const ggml_webgpu_rms_norm_mul_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.inplace); - ggml_webgpu_hash_combine(seed, key.overlap); - ggml_webgpu_hash_combine(seed, key.src_overlap); - return seed; - } -}; - -struct ggml_webgpu_rms_norm_mul_shader_decisions { - uint32_t wg_size = 0; - bool inplace = false; - bool overlap = false; - bool src_overlap = false; -}; - -/** Pad **/ -struct ggml_webgpu_pad_pipeline_key { - bool circular; - - bool operator==(const ggml_webgpu_pad_pipeline_key & other) const { return circular == other.circular; } -}; - -struct ggml_webgpu_pad_pipeline_key_hash { - size_t operator()(const ggml_webgpu_pad_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.circular); - return seed; - } -}; - -/** Solve Tri **/ -struct ggml_webgpu_solve_tri_pipeline_key { - int type; - int n; - int k; - - bool operator==(const ggml_webgpu_solve_tri_pipeline_key & other) const { - return type == other.type && n == other.n && k == other.k; - } -}; - -struct ggml_webgpu_solve_tri_pipeline_key_hash { - size_t operator()(const ggml_webgpu_solve_tri_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.type); - ggml_webgpu_hash_combine(seed, key.n); - ggml_webgpu_hash_combine(seed, key.k); - return seed; - } -}; - -/** SSM Conv **/ -struct ggml_webgpu_ssm_conv_pipeline_key { - int type; - int vectorized; - - bool operator==(const ggml_webgpu_ssm_conv_pipeline_key & other) const { - return type == other.type && vectorized == other.vectorized; - } -}; - -/** CONV 2D */ -struct ggml_webgpu_conv2d_pipeline_key { - ggml_type weight_type; - ggml_type input_type; - ggml_type output_type; - - bool operator==(const ggml_webgpu_conv2d_pipeline_key & other) const { - return weight_type == other.weight_type && input_type == other.input_type && output_type == other.output_type; - } -}; - -struct ggml_webgpu_conv2d_pipeline_key_hash { - size_t operator()(const ggml_webgpu_conv2d_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.weight_type); - ggml_webgpu_hash_combine(seed, key.input_type); - ggml_webgpu_hash_combine(seed, key.output_type); - return seed; - } -}; - -/** Im2Col **/ -struct ggml_webgpu_im2col_pipeline_key { - ggml_type input_type; - ggml_type output_type; - - bool operator==(const ggml_webgpu_im2col_pipeline_key & other) const { - return input_type == other.input_type && output_type == other.output_type; - } -}; - -struct ggml_webgpu_im2col_pipeline_key_hash { - size_t operator()(const ggml_webgpu_im2col_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.input_type); - ggml_webgpu_hash_combine(seed, key.output_type); - return seed; - } -}; - -/** Gated Delta Net **/ -struct ggml_webgpu_gated_delta_net_pipeline_key { - int type; - int s_v; - int kda; - - bool operator==(const ggml_webgpu_gated_delta_net_pipeline_key & other) const { - return type == other.type && s_v == other.s_v && kda == other.kda; - } -}; - -struct ggml_webgpu_gated_delta_net_pipeline_key_hash { - size_t operator()(const ggml_webgpu_gated_delta_net_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.type); - ggml_webgpu_hash_combine(seed, key.s_v); - ggml_webgpu_hash_combine(seed, key.kda); - return seed; - } -}; - -struct ggml_webgpu_ssm_conv_pipeline_key_hash { - size_t operator()(const ggml_webgpu_ssm_conv_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.type); - ggml_webgpu_hash_combine(seed, key.vectorized); - return seed; - } -}; - -/** Scale **/ - -struct ggml_webgpu_scale_pipeline_key { - int inplace; - - bool operator==(const ggml_webgpu_scale_pipeline_key & other) const { return inplace == other.inplace; } -}; - -struct ggml_webgpu_scale_pipeline_key_hash { - size_t operator()(const ggml_webgpu_scale_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.inplace); - return seed; - } -}; - -/** Upscale **/ - -struct ggml_webgpu_upscale_pipeline_key { - ggml_type input_type; - ggml_type output_type; - uint32_t base_mode; - bool antialias; - - bool operator==(const ggml_webgpu_upscale_pipeline_key & other) const { - return input_type == other.input_type && output_type == other.output_type && base_mode == other.base_mode && - antialias == other.antialias; - } -}; - -struct ggml_webgpu_upscale_pipeline_key_hash { - size_t operator()(const ggml_webgpu_upscale_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.input_type); - ggml_webgpu_hash_combine(seed, key.output_type); - ggml_webgpu_hash_combine(seed, key.base_mode); - ggml_webgpu_hash_combine(seed, key.antialias); - return seed; - } -}; - -/** Concat **/ - -struct ggml_webgpu_concat_pipeline_key { - int type; - - bool operator==(const ggml_webgpu_concat_pipeline_key & other) const { return type == other.type; } -}; - -struct ggml_webgpu_concat_pipeline_key_hash { - size_t operator()(const ggml_webgpu_concat_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.type); - return seed; - } -}; - -/** Repeat **/ - -struct ggml_webgpu_repeat_pipeline_key { - int type; - - bool operator==(const ggml_webgpu_repeat_pipeline_key & other) const { return type == other.type; } -}; - -struct ggml_webgpu_repeat_pipeline_key_hash { - size_t operator()(const ggml_webgpu_repeat_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.type); - return seed; - } -}; - -/** Binary **/ - -struct ggml_webgpu_binary_pipeline_key { - int type; - int op; - bool inplace; - bool overlap; - bool src_overlap; - - bool operator==(const ggml_webgpu_binary_pipeline_key & other) const { - return type == other.type && op == other.op && inplace == other.inplace && overlap == other.overlap && - src_overlap == other.src_overlap; - } -}; - -struct ggml_webgpu_binary_pipeline_key_hash { - size_t operator()(const ggml_webgpu_binary_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.type); - ggml_webgpu_hash_combine(seed, key.op); - ggml_webgpu_hash_combine(seed, key.inplace); - ggml_webgpu_hash_combine(seed, key.overlap); - ggml_webgpu_hash_combine(seed, key.src_overlap); - return seed; - } -}; - -/** Unary **/ - -struct ggml_webgpu_unary_pipeline_key { - int type; - int op; - bool is_unary; // many unary operators fall under the GGML_OP_UNARY umbrella - bool inplace; - ggml_tri_type ttype; // only used for GGML_OP_TRI - - bool operator==(const ggml_webgpu_unary_pipeline_key & other) const { - return type == other.type && op == other.op && is_unary == other.is_unary && inplace == other.inplace && - ttype == other.ttype; - } -}; - -struct ggml_webgpu_unary_pipeline_key_hash { - size_t operator()(const ggml_webgpu_unary_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.type); - ggml_webgpu_hash_combine(seed, key.op); - ggml_webgpu_hash_combine(seed, key.is_unary); - ggml_webgpu_hash_combine(seed, key.inplace); - ggml_webgpu_hash_combine(seed, key.ttype); - return seed; - } -}; - -/** FlashAttention */ - -enum ggml_webgpu_flash_attn_path : uint32_t { - GGML_WEBGPU_FLASH_ATTN_PATH_NONE = 0u, - GGML_WEBGPU_FLASH_ATTN_PATH_SUBGROUP_MATRIX = 1u, - GGML_WEBGPU_FLASH_ATTN_PATH_TILE = 2u, - GGML_WEBGPU_FLASH_ATTN_PATH_VEC = 3u, -}; - -struct ggml_webgpu_flash_attn_pipeline_key { - ggml_type kv_type; - uint32_t head_dim_qk; - uint32_t head_dim_v; - bool kv_direct; - bool kv_overlap; - bool has_mask; - bool has_sinks; - bool uses_logit_softcap; - uint32_t path; - - bool operator==(const ggml_webgpu_flash_attn_pipeline_key & other) const { - return kv_type == other.kv_type && head_dim_qk == other.head_dim_qk && head_dim_v == other.head_dim_v && - kv_direct == other.kv_direct && kv_overlap == other.kv_overlap && has_mask == other.has_mask && - has_sinks == other.has_sinks && uses_logit_softcap == other.uses_logit_softcap && path == other.path; - } -}; - -struct ggml_webgpu_flash_attn_pipeline_key_hash { - size_t operator()(const ggml_webgpu_flash_attn_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.kv_type); - ggml_webgpu_hash_combine(seed, key.head_dim_qk); - ggml_webgpu_hash_combine(seed, key.head_dim_v); - ggml_webgpu_hash_combine(seed, key.kv_direct); - ggml_webgpu_hash_combine(seed, key.kv_overlap); - ggml_webgpu_hash_combine(seed, key.has_mask); - ggml_webgpu_hash_combine(seed, key.has_sinks); - ggml_webgpu_hash_combine(seed, key.uses_logit_softcap); - ggml_webgpu_hash_combine(seed, key.path); - return seed; - } -}; - -struct ggml_webgpu_flash_attn_decisions { - uint32_t path = GGML_WEBGPU_FLASH_ATTN_PATH_NONE; - uint32_t q_tile = 0; - uint32_t kv_tile = 0; - uint32_t wg_size = 0; - bool kv_direct = false; - bool kv_overlap = false; -}; - -inline constexpr uint32_t GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH = 4u; -inline constexpr uint32_t GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE = 4u; - -inline uint32_t ggml_webgpu_flash_attn_pick_vec_ne(const ggml_webgpu_flash_attn_pipeline_key & key) { - if (key.path != GGML_WEBGPU_FLASH_ATTN_PATH_VEC || key.kv_type != GGML_TYPE_F16 || - key.head_dim_qk != key.head_dim_v) { - return 1u; - } - - switch (key.head_dim_qk) { - case 64: - case 192: - case 576: - return 2u; - case 96: - return 4u; - default: - return 1u; - } -} - -inline ggml_webgpu_flash_attn_pipeline_key ggml_webgpu_flash_attn_make_pipeline_key( - const ggml_webgpu_shader_lib_context & context, - uint32_t path) { - const bool has_mask = context.src3 != nullptr; - const bool has_sinks = context.src4 != nullptr; - bool kv_direct = false; - if (path != GGML_WEBGPU_FLASH_ATTN_PATH_TILE) { - uint32_t kv_direct_align = GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH; - if (path == GGML_WEBGPU_FLASH_ATTN_PATH_SUBGROUP_MATRIX) { - kv_direct_align = context.sg_mat_k; - } - kv_direct = (context.src1->type == GGML_TYPE_F16) && - (context.src0->ne[0] % std::max(1u, kv_direct_align) == 0) && - (context.src1->ne[1] % GGML_WEBGPU_KV_SEQ_PAD == 0); - } - - ggml_webgpu_flash_attn_pipeline_key key = {}; - key.kv_type = context.src1->type; - key.head_dim_qk = (uint32_t) context.src0->ne[0]; - key.head_dim_v = (uint32_t) context.src2->ne[0]; - key.kv_direct = kv_direct; - key.kv_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src2); - key.has_mask = has_mask; - key.has_sinks = has_sinks; - key.uses_logit_softcap = ggml_get_op_params_f32(context.dst, 2) != 0.0f; - key.path = path; - return key; -} - -struct ggml_webgpu_flash_attn_vec_reduce_pipeline_key { - uint32_t head_dim_v; - uint32_t wg_size; -}; - -struct ggml_webgpu_flash_attn_vec_reduce_pipeline_key_hash { - size_t operator()(const ggml_webgpu_flash_attn_vec_reduce_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.head_dim_v); - ggml_webgpu_hash_combine(seed, key.wg_size); - return seed; - } -}; - -inline bool operator==(const ggml_webgpu_flash_attn_vec_reduce_pipeline_key & lhs, - const ggml_webgpu_flash_attn_vec_reduce_pipeline_key & rhs) { - return lhs.head_dim_v == rhs.head_dim_v && lhs.wg_size == rhs.wg_size; -} - -struct ggml_webgpu_flash_attn_blk_pipeline_key { - uint32_t kv_tile; - - bool operator==(const ggml_webgpu_flash_attn_blk_pipeline_key & other) const { return kv_tile == other.kv_tile; } -}; - -struct ggml_webgpu_flash_attn_blk_pipeline_key_hash { - size_t operator()(const ggml_webgpu_flash_attn_blk_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.kv_tile); - return seed; - } -}; - -// This is exposed because it's necessary in supports_op -inline size_t ggml_webgpu_flash_attn_wg_mem_bytes(uint32_t q_tile, - uint32_t kv_tile, - uint32_t head_dim_qk, - uint32_t head_dim_v, - bool has_mask, - bool kv_direct) { - const uint32_t max_head_dim = std::max(head_dim_qk, head_dim_v); - size_t f16_elems = 0; - size_t f32_elems = 0; - f16_elems += q_tile * head_dim_qk; // q_shmem - if (!kv_direct) { - f16_elems += kv_tile * max_head_dim; // kv_shmem - } - f16_elems += q_tile * head_dim_v; // o_shmem - if (has_mask) { - f16_elems += q_tile * kv_tile; // mask_shmem - } - f16_elems += q_tile * kv_tile; // inter_shmem - f32_elems += q_tile; // row_max_shmem - f32_elems += q_tile; // exp_sum_shmem - return f16_elems * GGML_WEBGPU_F16_SIZE_BYTES + f32_elems * GGML_WEBGPU_F32_SIZE_BYTES; -} - -inline uint32_t ggml_webgpu_flash_attn_max_kv_tile(const ggml_webgpu_shader_lib_context & context, - const ggml_webgpu_flash_attn_pipeline_key & key) { - const size_t limit_bytes = context.wg_mem_limit_bytes; - uint32_t q_tile = context.sg_mat_m; - uint32_t kv_granularity = context.sg_mat_n; - if (key.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE) { - q_tile = GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE; - kv_granularity = std::max(1u, context.max_subgroup_size); - } else if (key.path == GGML_WEBGPU_FLASH_ATTN_PATH_VEC) { - q_tile = 1u; - kv_granularity = 8u; - } - const size_t base_q_bytes = (key.head_dim_qk + key.head_dim_v) * q_tile * GGML_WEBGPU_F16_SIZE_BYTES + - 2 * q_tile * GGML_WEBGPU_F32_SIZE_BYTES; - size_t bytes_per_kv = 0; - if (!key.kv_direct) { - bytes_per_kv += std::max(key.head_dim_qk, key.head_dim_v); - } - if (key.has_mask) { - bytes_per_kv += q_tile; - } - bytes_per_kv += q_tile; - bytes_per_kv *= GGML_WEBGPU_F16_SIZE_BYTES; - const uint32_t max_kv_tile = (limit_bytes - base_q_bytes) / bytes_per_kv; - return (max_kv_tile / kv_granularity) * kv_granularity; -} - -inline ggml_webgpu_flash_attn_decisions ggml_webgpu_flash_attn_get_decisions( - const ggml_webgpu_shader_lib_context & context, - size_t storage_offset_alignment) { - ggml_webgpu_flash_attn_decisions decisions = {}; - const size_t alignment = std::max(1u, storage_offset_alignment); - const auto * K = context.src1; - const auto * V = context.src2; - GGML_ASSERT(K != nullptr); - GGML_ASSERT(V != nullptr); - - const auto flash_attn_tensor_offset = [](const ggml_tensor * tensor) -> size_t { - constexpr uintptr_t ptr_base_addr = 0x1000u; - const ggml_tensor * base = tensor->view_src != nullptr ? tensor->view_src : tensor; - return reinterpret_cast(base->data) - ptr_base_addr + tensor->view_offs; - }; - - const uint32_t k_offset_elems = - (uint32_t) ((flash_attn_tensor_offset(K) & (alignment - 1)) / ggml_type_size(K->type)); - const uint32_t v_offset_elems = - (uint32_t) ((flash_attn_tensor_offset(V) & (alignment - 1)) / ggml_type_size(V->type)); - const bool f16_vec4_aligned = (k_offset_elems % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0u) && - (v_offset_elems % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0u); - const bool kv_vec_type_supported = - K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q4_0 || K->type == GGML_TYPE_Q8_0; - const bool use_vec = context.supports_subgroups && (context.src0->ne[1] < 20) && (context.src0->ne[0] % 32 == 0) && - (context.src2->ne[0] % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0) && - kv_vec_type_supported && (K->type != GGML_TYPE_F16 || f16_vec4_aligned) && - (context.src2->type == K->type); - const bool use_tile = context.supports_subgroups && !context.supports_subgroup_matrix && K->type == GGML_TYPE_F16 && - V->type == GGML_TYPE_F16 && f16_vec4_aligned && - (context.src0->ne[0] % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0) && - (context.src2->ne[0] % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0) && !use_vec; - - decisions.path = use_vec ? GGML_WEBGPU_FLASH_ATTN_PATH_VEC : - use_tile ? GGML_WEBGPU_FLASH_ATTN_PATH_TILE : - context.supports_subgroup_matrix ? GGML_WEBGPU_FLASH_ATTN_PATH_SUBGROUP_MATRIX : - GGML_WEBGPU_FLASH_ATTN_PATH_NONE; - - if (decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_NONE) { - return decisions; - } - - const ggml_webgpu_flash_attn_pipeline_key key = ggml_webgpu_flash_attn_make_pipeline_key(context, decisions.path); - decisions.kv_direct = key.kv_direct; - const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile(context, key); - // invalidate if even the smallest kv_tile doesn't fit in shared memory - if (max_kv_tile == 0) { - decisions.path = GGML_WEBGPU_FLASH_ATTN_PATH_NONE; - return decisions; - } - - if (decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_VEC) { - decisions.q_tile = 1u; - decisions.kv_tile = std::max(8u, std::min(32u, max_kv_tile)); - decisions.kv_tile = (decisions.kv_tile / 8u) * 8u; - decisions.wg_size = std::max(1u, std::min(32u, context.max_subgroup_size)); - if (decisions.kv_direct) { - decisions.kv_tile = std::min(decisions.kv_tile, GGML_WEBGPU_KV_SEQ_PAD); - while (GGML_WEBGPU_KV_SEQ_PAD % decisions.kv_tile != 0) { - decisions.kv_tile -= 8u; - } - } - return decisions; - } - - decisions.q_tile = - decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE ? GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE : context.sg_mat_m; - decisions.kv_tile = decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE ? - std::min(64u, max_kv_tile) : - std::min(max_kv_tile, context.sg_mat_n * GGML_WEBGPU_FLASH_ATTN_PREFERRED_KV_SG_TILES); - decisions.wg_size = decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE ? - GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE : - std::max(context.max_subgroup_size, GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE); - - if (decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE) { - const uint32_t tile_kv_granularity = std::max(1u, context.max_subgroup_size); - decisions.kv_tile = - std::max(tile_kv_granularity, (decisions.kv_tile / tile_kv_granularity) * tile_kv_granularity); - } - - if (decisions.kv_direct) { - GGML_ASSERT(decisions.kv_tile <= GGML_WEBGPU_KV_SEQ_PAD); - while (GGML_WEBGPU_KV_SEQ_PAD % decisions.kv_tile != 0) { - decisions.kv_tile -= decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE ? - std::max(1u, context.max_subgroup_size) : - context.sg_mat_n; - } - } - return decisions; -} - -/** Matrix Multiplication **/ - -struct ggml_webgpu_legacy_mul_mat_pipeline_key { - ggml_type src0_type; - ggml_type src1_type; - - bool operator==(const ggml_webgpu_legacy_mul_mat_pipeline_key & other) const { - return src0_type == other.src0_type && src1_type == other.src1_type; - } -}; - -struct ggml_webgpu_legacy_mul_mat_pipeline_key_hash { - size_t operator()(const ggml_webgpu_legacy_mul_mat_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.src0_type); - ggml_webgpu_hash_combine(seed, key.src1_type); - return seed; - } -}; - -struct ggml_webgpu_mul_mat_vec_pipeline_key { - ggml_type src0_type; - ggml_type src1_type; - int vectorized; - - bool operator==(const ggml_webgpu_mul_mat_vec_pipeline_key & other) const { - return src0_type == other.src0_type && src1_type == other.src1_type && vectorized == other.vectorized; - } -}; - -struct ggml_webgpu_mul_mat_vec_pipeline_key_hash { - size_t operator()(const ggml_webgpu_mul_mat_vec_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.src0_type); - ggml_webgpu_hash_combine(seed, key.src1_type); - ggml_webgpu_hash_combine(seed, key.vectorized); - return seed; - } -}; - -struct ggml_webgpu_mul_mat_vec_shader_decisions { - uint32_t wg_size; - uint32_t outputs_per_wg; - uint32_t vec_size; -}; - -struct ggml_webgpu_mul_mat_pipeline_key { - ggml_type src0_type; - ggml_type src1_type; - int vectorized; - int use_subgroup_matrix; - - bool operator==(const ggml_webgpu_mul_mat_pipeline_key & other) const { - return src0_type == other.src0_type && src1_type == other.src1_type && vectorized == other.vectorized && - use_subgroup_matrix == other.use_subgroup_matrix; - } -}; - -struct ggml_webgpu_mul_mat_pipeline_key_hash { - size_t operator()(const ggml_webgpu_mul_mat_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.src0_type); - ggml_webgpu_hash_combine(seed, key.src1_type); - ggml_webgpu_hash_combine(seed, key.vectorized); - ggml_webgpu_hash_combine(seed, key.use_subgroup_matrix); - return seed; - } -}; - -struct ggml_webgpu_mul_mat_shader_decisions { - uint32_t tile_k; - uint32_t wg_size_m; - uint32_t wg_size_n; - uint32_t wg_size; - uint32_t outputs_per_wg; - int use_subgroup_matrix; - - uint32_t tile_m; - uint32_t tile_n; - - // Subgroup matrix parameters - uint32_t subgroup_m; - uint32_t subgroup_n; - uint32_t subgroup_matrix_m; - uint32_t subgroup_matrix_n; - - uint32_t mul_mat_wg_size; -}; - -/** MUL_MAT_ID **/ - -struct ggml_webgpu_mul_mat_id_pipeline_key { - ggml_type src0_type; - ggml_type src1_type; - uint32_t n_experts; - int vectorized; - - bool operator==(const ggml_webgpu_mul_mat_id_pipeline_key & other) const { - return src0_type == other.src0_type && src1_type == other.src1_type && n_experts == other.n_experts && - vectorized == other.vectorized; - } -}; - -struct ggml_webgpu_mul_mat_id_pipeline_key_hash { - size_t operator()(const ggml_webgpu_mul_mat_id_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.src0_type); - ggml_webgpu_hash_combine(seed, key.src1_type); - ggml_webgpu_hash_combine(seed, key.n_experts); - ggml_webgpu_hash_combine(seed, key.vectorized); - return seed; - } -}; - -/** Cpy **/ - -struct ggml_webgpu_cpy_pipeline_key { - ggml_type src_type; - ggml_type dst_type; - - bool operator==(const ggml_webgpu_cpy_pipeline_key & other) const { - return src_type == other.src_type && dst_type == other.dst_type; - } -}; - -struct ggml_webgpu_cpy_pipeline_key_hash { - size_t operator()(const ggml_webgpu_cpy_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.src_type); - ggml_webgpu_hash_combine(seed, key.dst_type); - return seed; - } -}; - -/** Glu **/ - -struct ggml_webgpu_glu_pipeline_key { - ggml_glu_op glu_op; - ggml_type type; - bool split; - - bool operator==(const ggml_webgpu_glu_pipeline_key & other) const { - return glu_op == other.glu_op && type == other.type && split == other.split; - } -}; - -struct ggml_webgpu_glu_pipeline_key_hash { - size_t operator()(const ggml_webgpu_glu_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.glu_op); - ggml_webgpu_hash_combine(seed, key.type); - ggml_webgpu_hash_combine(seed, key.split); - return seed; - } -}; - -/** Rope **/ - -struct ggml_webgpu_rope_pipeline_key { - ggml_type type; - bool inplace; - bool has_ff; - - bool operator==(const ggml_webgpu_rope_pipeline_key & other) const { - return type == other.type && inplace == other.inplace && has_ff == other.has_ff; - } -}; - -struct ggml_webgpu_rope_pipeline_key_hash { - size_t operator()(const ggml_webgpu_rope_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.type); - ggml_webgpu_hash_combine(seed, key.inplace); - ggml_webgpu_hash_combine(seed, key.has_ff); - return seed; - } -}; - -/** SoftMax **/ - -struct ggml_webgpu_soft_max_pipeline_key { - ggml_type mask_type; - bool has_mask; - bool has_sink; - bool inplace; - - bool operator==(const ggml_webgpu_soft_max_pipeline_key & other) const { - return mask_type == other.mask_type && has_mask == other.has_mask && has_sink == other.has_sink && - inplace == other.inplace; - } -}; - -struct ggml_webgpu_soft_max_pipeline_key_hash { - size_t operator()(const ggml_webgpu_soft_max_pipeline_key & key) const { - size_t seed = 0; - ggml_webgpu_hash_combine(seed, key.mask_type); - ggml_webgpu_hash_combine(seed, key.has_mask); - ggml_webgpu_hash_combine(seed, key.has_sink); - ggml_webgpu_hash_combine(seed, key.inplace); - return seed; - } -}; - -class ggml_webgpu_shader_lib { - wgpu::Device device; - pre_wgsl::Preprocessor preprocessor; - - std::unordered_map sum_rows_pipelines; // key is fixed, no variants yet - std::unordered_map argmax_pipelines; // key is vec4 - std::unordered_map argsort_pipelines; // key is order - std::unordered_map argsort_merge_pipelines; // key is order - std::unordered_map cumsum_pipelines; // key is fixed, no variants yet - std::unordered_map - row_norm_pipelines; // op/inplace - - std::unordered_map - get_rows_pipelines; // src_type, vectorized - std::unordered_map - unary_pipelines; // type/op/inplace - std::unordered_map - scale_pipelines; // inplace - std::unordered_map - solve_tri_pipelines; // type - std::unordered_map - ssm_conv_pipelines; // type/vectorized - std::unordered_map - ssm_scan_pipelines; // type/d_state - std::unordered_map - gated_delta_net_pipelines; // type/S_v/kda - std::unordered_map - pad_pipelines; // circular/non-circular - std::unordered_map - binary_pipelines; // type/op/inplace/overlap - std::unordered_map - concat_pipelines; // type - std::unordered_map - repeat_pipelines; // type - std::unordered_map - flash_attn_pipelines; - std::unordered_map - flash_attn_vec_reduce_pipelines; - std::unordered_map - flash_attn_blk_pipelines; - std::unordered_map - mul_mat_legacy_pipelines; // legacy mul_mat (non-subgroup/non-regtile/non-vec) - std::unordered_map - mul_mat_vec_pipelines; // fast mat-vec (n==1) - std::unordered_map - mul_mat_fast_pipelines; // fast mat-mat (reg-tile or subgroup) - std::unordered_map mul_mat_id_gather_pipelines; // key is fixed - std::unordered_map - mul_mat_id_pipelines; // src0_type/src1_type - std::unordered_map - mul_mat_id_vec_pipelines; // src0_type/src1_type - - std::unordered_map - set_rows_pipelines; - std::unordered_map set_pipelines; - std::unordered_map cpy_pipelines; - std::unordered_map glu_pipelines; - std::unordered_map - rope_pipelines; - std::unordered_map - soft_max_pipelines; - std::unordered_map - conv2d_pipelines; - std::unordered_map - im2col_pipelines; - - std::unordered_map - rms_norm_mul_pipelines; - std::unordered_map - upscale_pipelines; - - public: - ggml_webgpu_shader_lib(wgpu::Device device) { this->device = device; } - - webgpu_pipeline get_sum_rows_pipeline(const ggml_webgpu_shader_lib_context & context) { - auto it = sum_rows_pipelines.find(1); - if (it != sum_rows_pipelines.end()) { - return it->second; - } - std::vector defines; - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_sum_rows, defines); - sum_rows_pipelines[1] = ggml_webgpu_create_pipeline(device, processed, "sum_rows"); - return sum_rows_pipelines[1]; - } - - webgpu_pipeline get_row_norm_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_row_norm_pipeline_key key = {}; - key.op = context.dst->op; - key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); - - auto it = row_norm_pipelines.find(key); - if (it != row_norm_pipelines.end()) { - return it->second; - } - std::vector defines; - std::string variant; - - switch (key.op) { - case GGML_OP_RMS_NORM: - defines.push_back("RMS_NORM"); - variant = "rms_norm"; - break; - case GGML_OP_L2_NORM: - defines.push_back("L2_NORM"); - variant = "l2_norm"; - break; - default: - GGML_ABORT("Unsupported op for row_norm shader"); - } - - if (key.inplace) { - defines.push_back("INPLACE"); - variant += "_inplace"; - } - - const uint32_t row_norm_wg_size = 128u; - uint32_t wg_size = std::min(context.max_wg_size, row_norm_wg_size); - defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); - auto processed = preprocessor.preprocess(wgsl_row_norm, defines); - auto decisions = std::make_shared(); - decisions->wg_size = wg_size; - decisions->inplace = key.inplace; - row_norm_pipelines[key] = ggml_webgpu_create_pipeline(device, processed, variant); - row_norm_pipelines[key].context = decisions; - return row_norm_pipelines[key]; - } - - webgpu_pipeline get_argmax_pipeline(const ggml_webgpu_shader_lib_context & context) { - bool vec4 = context.src0->ne[0] % 4 == 0; - - auto it = argmax_pipelines.find(vec4); - if (it != argmax_pipelines.end()) { - return it->second; - } - std::string variant = "argmax"; - std::vector defines; - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - if (vec4) { - defines.push_back("VEC4"); - variant += "_vec4"; - } - - auto processed = preprocessor.preprocess(wgsl_argmax, defines); - argmax_pipelines[vec4] = ggml_webgpu_create_pipeline(device, processed, variant); - return argmax_pipelines.at(vec4); - } - - webgpu_pipeline get_set_rows_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_set_rows_pipeline_key key = {}; - key.dst_type = context.dst->type; - key.vec4 = context.src0->ne[0] % 4 == 0; - key.i64_idx = context.src1->type == GGML_TYPE_I64; - - auto it = set_rows_pipelines.find(key); - if (it != set_rows_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "set_rows"; - - switch (context.dst->type) { - case GGML_TYPE_F32: - defines.push_back("DST_F32"); - variant += "_dstf32"; - break; - case GGML_TYPE_F16: - defines.push_back("DST_F16"); - variant += "_dstf16"; - break; - default: - GGML_ABORT("Unsupported dst type for set_rows shader"); - } - - if (key.vec4) { - defines.push_back("VEC4"); - variant += "_vec4"; - } - if (key.i64_idx) { - defines.push_back("I64_IDX"); - variant += "_i64idx"; - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_set_rows, defines); - auto decisions = std::make_shared(); - decisions->vec4 = key.vec4; - decisions->i64_idx = key.i64_idx; - decisions->wg_size = context.max_wg_size; - set_rows_pipelines[key] = ggml_webgpu_create_pipeline(device, processed, variant); - set_rows_pipelines[key].context = decisions; - return set_rows_pipelines[key]; - } - - webgpu_pipeline get_set_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_set_pipeline_key key = {}; - key.type = context.dst->type; - key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); - - auto it = set_pipelines.find(key); - if (it != set_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "set"; - - switch (key.type) { - case GGML_TYPE_F32: - defines.push_back("TYPE_F32"); - variant += "_f32"; - break; - case GGML_TYPE_I32: - defines.push_back("TYPE_I32"); - variant += "_i32"; - break; - default: - GGML_ABORT("Unsupported type for set shader"); - } - - if (key.inplace) { - defines.push_back("INPLACE"); - variant += "_inplace"; - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_set, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - decisions->inplace = key.inplace; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - set_pipelines[key] = pipeline; - return set_pipelines[key]; - } - - webgpu_pipeline get_cumsum_pipeline(const ggml_webgpu_shader_lib_context & context) { - auto it = cumsum_pipelines.find(1); - if (it != cumsum_pipelines.end()) { - return it->second; - } - - std::vector defines; - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_cumsum, defines); - cumsum_pipelines[1] = ggml_webgpu_create_pipeline(device, processed, "cumsum"); - return cumsum_pipelines[1]; - } - - webgpu_pipeline get_argsort_pipeline(const ggml_webgpu_shader_lib_context & context) { - bool is_top_k = context.dst->op == GGML_OP_TOP_K; - // ascending order is 0, descending order is 1 - const int32_t order = - is_top_k ? (int32_t) GGML_SORT_ORDER_DESC : (int32_t) ggml_get_op_params_i32(context.dst, 0); - - auto it = argsort_pipelines.find(order); - if (it != argsort_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "argsort"; - defines.push_back(std::string("ORDER=") + std::to_string(order)); - variant += std::string("_order") + std::to_string(order); - uint32_t wg_size = 1; - while (wg_size * 2 <= context.max_wg_size && - wg_size * GGML_WEBGPU_I32_SIZE_BYTES <= context.wg_mem_limit_bytes / 2) { - wg_size *= 2; - } - defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); - auto processed = preprocessor.preprocess(wgsl_argsort, defines); - auto decisions = std::make_shared(); - decisions->wg_size = wg_size; - argsort_pipelines[order] = ggml_webgpu_create_pipeline(device, processed, variant); - argsort_pipelines[order].context = decisions; - return argsort_pipelines[order]; - } - - webgpu_pipeline get_argsort_merge_pipeline(const ggml_webgpu_shader_lib_context & context) { - bool is_top_k = context.dst->op == GGML_OP_TOP_K; - // ascending order is 0, descending order is 1 - const int32_t order = - is_top_k ? (int32_t) GGML_SORT_ORDER_DESC : (int32_t) ggml_get_op_params_i32(context.dst, 0); - - auto it = argsort_merge_pipelines.find(order); - if (it != argsort_merge_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "argsort_merge"; - defines.push_back(std::string("ORDER=") + std::to_string(order)); - variant += std::string("_order") + std::to_string(order); - uint32_t wg_size = std::min(GGML_WEBGPU_ARGSORT_MERGE_MAX_WG_SIZE, context.max_wg_size); - defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); - - auto processed = preprocessor.preprocess(wgsl_argsort_merge, defines); - argsort_merge_pipelines[order] = ggml_webgpu_create_pipeline(device, processed, variant); - return argsort_merge_pipelines[order]; - } - - webgpu_pipeline get_get_rows_pipeline(const ggml_webgpu_shader_lib_context & context) { - const bool vectorized = context.src0->type == GGML_TYPE_F32 && context.dst->ne[0] % 4 == 0; - ggml_webgpu_get_rows_pipeline_key key = {}; - key.src_type = context.src0->type; - key.vectorized = (int) vectorized; - - auto it = get_rows_pipelines.find(key); - if (it != get_rows_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "get_rows"; - - const struct ggml_type_traits * type_traits = ggml_get_type_traits(key.src_type); - const char * type_str = type_traits->type_name; - - switch (key.src_type) { - case GGML_TYPE_F32: - defines.push_back("FLOAT_PARALLEL"); - if (key.vectorized) { - defines.push_back("F32_VEC"); - defines.push_back("SRC_TYPE=vec4"); - defines.push_back("DST_TYPE=vec4"); - defines.push_back("BLOCK_SIZE=4u"); - } else { - defines.push_back("F32"); - defines.push_back("SRC_TYPE=f32"); - defines.push_back("DST_TYPE=f32"); - defines.push_back("BLOCK_SIZE=1u"); - } - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("FLOAT_PARALLEL"); - defines.push_back("F16"); - defines.push_back("SRC_TYPE=f16"); - defines.push_back("DST_TYPE=f32"); - defines.push_back("BLOCK_SIZE=1u"); - variant += "_f16"; - break; - case GGML_TYPE_I32: - defines.push_back("FLOAT_PARALLEL"); - defines.push_back("I32"); - defines.push_back("SRC_TYPE=i32"); - defines.push_back("DST_TYPE=i32"); - defines.push_back("BLOCK_SIZE=1u"); - variant += "_i32"; - break; - default: - { - std::string type_upper = type_str; - std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper); - - switch (key.src_type) { - case GGML_TYPE_Q1_0: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q8_0: - case GGML_TYPE_Q3_K: - case GGML_TYPE_Q6_K: - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ2_S: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ4_NL: - { - // Quantized types using u32 buffers for portability. - defines.push_back("SRC_TYPE=u32"); - defines.push_back("U32_DEQUANT_HELPERS"); - break; - } - default: - { - defines.push_back(std::string("SRC_TYPE=") + type_str); - } - } - - defines.push_back("BYTE_HELPERS"); - defines.push_back(type_upper + "_T"); - defines.push_back(type_upper); - defines.push_back(type_upper + "_SCALE_MIN"); - defines.push_back(type_upper + "_TABLES"); - defines.push_back(type_upper + "_GRID"); - - variant += "_"; - variant += type_str; - - defines.push_back("DST_TYPE=f32"); - - if (key.src_type == GGML_TYPE_Q1_0) { - defines.push_back("BLOCK_SIZE=128u"); - } else if ((key.src_type >= GGML_TYPE_Q4_0 && key.src_type <= GGML_TYPE_Q8_1) || - key.src_type == GGML_TYPE_IQ4_NL) { - defines.push_back("BLOCK_SIZE=32u"); - } else if (key.src_type >= GGML_TYPE_Q2_K) { - defines.push_back("BLOCK_SIZE=256u"); - } else { - defines.push_back("BLOCK_SIZE=1u"); - } - break; - } - } - - if (key.vectorized) { - variant += "_vec"; - } - - defines.push_back("WG_SIZE=" + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_get_rows, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - get_rows_pipelines[key] = pipeline; - return get_rows_pipelines[key]; - } - - webgpu_pipeline get_scale_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_scale_pipeline_key key = {}; - key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); - - auto it = scale_pipelines.find(key); - if (it != scale_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "scale"; - - if (key.inplace) { - defines.push_back("INPLACE"); - variant += "_inplace"; - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_scale, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - decisions->inplace = key.inplace; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - scale_pipelines[key] = pipeline; - return scale_pipelines[key]; - } - - webgpu_pipeline get_solve_tri_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_solve_tri_pipeline_key key = {}; - key.type = context.dst->type; - key.n = (int) context.src0->ne[0]; - key.k = (int) context.src1->ne[0]; - - auto it = solve_tri_pipelines.find(key); - if (it != solve_tri_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "solve_tri"; - - switch (key.type) { - case GGML_TYPE_F32: - variant += "_f32"; - break; - default: - GGML_ABORT("Unsupported type for solve_tri shader"); - } - - const uint32_t wg_size = std::min((uint32_t) key.n, context.max_wg_size); - const uint32_t k_tile = wg_size; - const uint32_t bytes_per_row = ((uint32_t) key.n + wg_size) * GGML_WEBGPU_F32_SIZE_BYTES; - const uint32_t batch_n = (uint32_t) (context.wg_mem_limit_bytes / bytes_per_row); - - defines.push_back(std::string("N=") + std::to_string(key.n)); - defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); - defines.push_back(std::string("K_TILE=") + std::to_string(k_tile)); - defines.push_back(std::string("BATCH_N=") + std::to_string(batch_n)); - - auto processed = preprocessor.preprocess(wgsl_solve_tri, defines); - auto decisions = std::make_shared(); - decisions->wg_size = wg_size; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - solve_tri_pipelines[key] = pipeline; - return solve_tri_pipelines[key]; - } - - webgpu_pipeline get_ssm_conv_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_ssm_conv_pipeline_key key = {}; - key.type = context.dst->type; - key.vectorized = context.src1->ne[0] == 4; - - auto it = ssm_conv_pipelines.find(key); - if (it != ssm_conv_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "ssm_conv"; - - switch (key.type) { - case GGML_TYPE_F32: - variant += "_f32"; - break; - default: - GGML_ABORT("Unsupported type for ssm_conv shader"); - } - - if (key.vectorized) { - defines.push_back("VECTORIZED"); - variant += "_vec4"; - } - - constexpr uint32_t block_size = 32u; - constexpr uint32_t tokens_per_wg = 8u; - - defines.push_back("BLOCK_SIZE=" + std::to_string(block_size) + "u"); - defines.push_back("TOKENS_PER_WG=" + std::to_string(tokens_per_wg) + "u"); - - auto processed = preprocessor.preprocess(wgsl_ssm_conv, defines); - auto decisions = std::make_shared(); - decisions->block_size = block_size; - decisions->tokens_per_wg = tokens_per_wg; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - ssm_conv_pipelines[key] = pipeline; - return ssm_conv_pipelines[key]; - } - - webgpu_pipeline get_ssm_scan_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_ssm_scan_pipeline_key key = {}; - key.type = context.dst->type; - key.d_state = (int) context.src0->ne[0]; - key.xbc_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src4) && - ggml_webgpu_tensor_overlap(context.src1, context.src5); - - auto it = ssm_scan_pipelines.find(key); - if (it != ssm_scan_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "ssm_scan"; - - switch (key.type) { - case GGML_TYPE_F32: - variant += "_f32"; - break; - default: - GGML_ABORT("Unsupported type for ssm_scan shader"); - } - - const uint32_t wg_size = (uint32_t) key.d_state; - - constexpr uint32_t tokens_per_tile = 4u; - - defines.push_back("WG_SIZE=" + std::to_string(wg_size) + "u"); - defines.push_back("TOKENS_PER_TILE=" + std::to_string(tokens_per_tile) + "u"); - - if (context.supports_subgroups) { - defines.push_back("USE_SUBGROUP_REDUCTION"); - variant += "_sg_reduce"; - } else { - variant += "_wg_reduce"; - } - - if (key.xbc_overlap) { - defines.push_back("XBC_OVERLAP"); - } - - variant += "_d" + std::to_string(key.d_state); - - auto processed = preprocessor.preprocess(wgsl_ssm_scan, defines); - auto decisions = std::make_shared(); - decisions->wg_size = wg_size; - decisions->tokens_per_tile = tokens_per_tile; - decisions->xbc_overlap = key.xbc_overlap; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - ssm_scan_pipelines[key] = pipeline; - return ssm_scan_pipelines[key]; - } - - webgpu_pipeline get_gated_delta_net_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_gated_delta_net_pipeline_key key = {}; - key.type = context.dst->type; - key.s_v = (int) context.src2->ne[0]; - key.kda = context.src3->ne[0] == context.src2->ne[0]; - - auto it = gated_delta_net_pipelines.find(key); - if (it != gated_delta_net_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "gated_delta_net"; - - switch (key.type) { - case GGML_TYPE_F32: - variant += "_f32"; - break; - default: - GGML_ABORT("Unsupported type for gated_delta_net shader"); - } - - if (key.kda) { - defines.push_back("KDA"); - variant += "_kda"; - } - - defines.push_back("S_V=" + std::to_string(key.s_v) + "u"); - defines.push_back("WG_SIZE=" + std::to_string(key.s_v) + "u"); - - auto processed = preprocessor.preprocess(wgsl_gated_delta_net, defines); - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - gated_delta_net_pipelines[key] = pipeline; - return gated_delta_net_pipelines[key]; - } - - webgpu_pipeline get_pad_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_pad_pipeline_key key = {}; - key.circular = ggml_get_op_params_i32(context.dst, 8) != 0; - - auto it = pad_pipelines.find(key); - if (it != pad_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "pad"; - - if (key.circular) { - defines.push_back("CIRCULAR"); - variant += "_circular"; - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_pad, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - pad_pipelines[key] = pipeline; - return pad_pipelines[key]; - } - - webgpu_pipeline get_mul_mat_vec_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_mul_mat_vec_pipeline_key key = {}; - key.src0_type = context.src0->type; - key.src1_type = context.src1->type; - key.vectorized = (context.src0->ne[0] % 4 == 0 && - (context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ? - 1 : - 0; - - auto it = mul_mat_vec_pipelines.find(key); - if (it != mul_mat_vec_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "mul_mat_vec"; - const char * shader_src = wgsl_mul_mat_vec; - - // src0 type (matrix row) - switch (context.src0->type) { - case GGML_TYPE_F32: - defines.push_back("SRC0_INNER_TYPE=f32"); - defines.push_back("MUL_ACC_FLOAT"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("SRC0_INNER_TYPE=f16"); - defines.push_back("MUL_ACC_FLOAT"); - variant += "_f16"; - break; - default: - { - // Quantized types: use helpers but accumulate in f16 - const struct ggml_type_traits * src0_traits = ggml_get_type_traits(context.src0->type); - std::string src0_name = src0_traits->type_name; - std::string type_upper = src0_name; - variant += "_" + src0_name; - std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper); - - defines.push_back("BYTE_HELPERS"); - defines.push_back("MUL_ACC_" + type_upper); - defines.push_back("U32_DEQUANT_HELPERS"); - defines.push_back("SRC0_INNER_TYPE=u32"); - switch (context.src0->type) { - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ2_S: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - defines.push_back(type_upper + "_GRID"); - break; - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ3_XXS: - defines.push_back(type_upper + "_GRID"); - defines.push_back(type_upper + "_TABLES"); - break; - default: - break; - } - break; - } - } - - // src1 type (vector) - switch (context.src1->type) { - case GGML_TYPE_F32: - defines.push_back("SRC1_INNER_TYPE=f32"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("SRC1_INNER_TYPE=f16"); - variant += "_f16"; - break; - default: - GGML_ABORT("Unsupported src1 type for mul_mat_vec shader"); - } - - // VEC/SCALAR controls - defines.push_back(key.vectorized ? "VEC" : "SCALAR"); - - uint32_t wg_size = WEBGPU_MUL_MAT_VEC_WG_SIZE; - uint32_t outputs_per_wg = WEBGPU_MUL_MAT_VEC_FLOAT_OUTPUTS_PER_WG; - - if (key.src0_type == GGML_TYPE_Q1_0) { - outputs_per_wg = WEBGPU_MUL_MAT_VEC_LEGACY_Q_OUTPUTS_PER_WG; - } else if (key.src0_type >= GGML_TYPE_Q2_K) { - outputs_per_wg = WEBGPU_MUL_MAT_VEC_K_Q_OUTPUTS_PER_WG; - } else if (key.src0_type >= GGML_TYPE_Q4_0) { - outputs_per_wg = WEBGPU_MUL_MAT_VEC_LEGACY_Q_OUTPUTS_PER_WG; - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); - defines.push_back(std::string("OUTPUTS_PER_WG=") + std::to_string(outputs_per_wg)); - defines.push_back(context.supports_subgroups ? "USE_SUBGROUP_REDUCTION" : "USE_WORKGROUP_REDUCTION"); - variant += context.supports_subgroups ? "_sg_reduce" : "_wg_reduce"; - if (key.vectorized) { - variant += "_vectorized"; - } - - auto processed = preprocessor.preprocess(shader_src, defines); - auto decisions = std::make_shared(); - decisions->wg_size = wg_size; - decisions->outputs_per_wg = outputs_per_wg; - decisions->vec_size = key.vectorized ? 4 : 1; - - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - mul_mat_vec_pipelines[key] = pipeline; - return mul_mat_vec_pipelines[key]; - } - - webgpu_pipeline get_mul_mat_fast_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_mul_mat_pipeline_key key = {}; - key.src0_type = context.src0->type; - key.src1_type = context.src1->type; - key.vectorized = (context.src0->ne[0] % 4 == 0 && context.dst->ne[0] % 4 == 0 && - (context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ? - 1 : - 0; - key.use_subgroup_matrix = context.supports_subgroup_matrix; - - auto it = mul_mat_fast_pipelines.find(key); - if (it != mul_mat_fast_pipelines.end()) { - return it->second; - } - - const char * shader_src = key.use_subgroup_matrix ? wgsl_mul_mat_subgroup_matrix : wgsl_mul_mat_reg_tile; - std::vector defines; - std::string variant = key.use_subgroup_matrix ? "mul_mat_subgroup_matrix" : "mul_mat_reg_tile"; - - // src1 type - switch (context.src1->type) { - case GGML_TYPE_F32: - defines.push_back("SRC1_INNER_TYPE=f32"); - break; - case GGML_TYPE_F16: - defines.push_back("SRC1_INNER_TYPE=f16"); - break; - default: - GGML_ABORT("Unsupported src1 type for mul_mat fast shader"); - } - - // src0 type - const struct ggml_type_traits * src0_traits = ggml_get_type_traits(context.src0->type); - const char * src0_name = src0_traits->type_name; - - switch (context.src0->type) { - case GGML_TYPE_F32: - defines.push_back("SRC0_INNER_TYPE=f32"); - defines.push_back("FLOAT"); - defines.push_back("MUL_ACC_FLOAT"); - defines.push_back("INIT_SRC0_SHMEM_FLOAT"); - defines.push_back("INIT_SRC1_SHMEM_FLOAT"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("SRC0_INNER_TYPE=f16"); - defines.push_back("FLOAT"); - defines.push_back("MUL_ACC_FLOAT"); - defines.push_back("INIT_SRC0_SHMEM_FLOAT"); - defines.push_back("INIT_SRC1_SHMEM_FLOAT"); - variant += "_f16"; - break; - default: - { - std::string type_upper = src0_name; - std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper); - - defines.push_back("BYTE_HELPERS"); - defines.push_back("MUL_ACC_" + type_upper); - defines.push_back("INIT_SRC0_SHMEM_" + type_upper); - defines.push_back("INIT_SRC1_SHMEM_FLOAT"); - defines.push_back("U32_DEQUANT_HELPERS"); - defines.push_back("SRC0_INNER_TYPE=u32"); - - switch (context.src0->type) { - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - defines.push_back(type_upper + "_GRID"); - break; - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ2_S: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ3_S: - defines.push_back(type_upper + "_GRID"); - defines.push_back(type_upper + "_TABLES"); - break; - default: - break; - } - - variant += std::string("_") + src0_name; - break; - } - } - - // VEC/SCALAR controls - defines.push_back(key.vectorized ? "VEC" : "SCALAR"); - - const bool is_quant = ggml_is_quantized(context.src0->type); - - uint32_t tile_k; - if (key.use_subgroup_matrix) { - tile_k = is_quant ? WEBGPU_MUL_MAT_SUBGROUP_TILE_K_QUANT : WEBGPU_MUL_MAT_SUBGROUP_TILE_K_FLOAT; - } else { - tile_k = is_quant ? WEBGPU_MUL_MAT_REG_TILE_K_QUANT : WEBGPU_MUL_MAT_REG_TILE_K_FLOAT; - } - - // Tiles - defines.push_back("TILE_M=" + std::to_string(WEBGPU_MUL_MAT_TILE_M) + "u"); - defines.push_back("TILE_N=" + std::to_string(WEBGPU_MUL_MAT_TILE_N) + "u"); - - // Subgroup matrix specifics - if (key.use_subgroup_matrix) { - defines.push_back("TILE_K=" + std::to_string(tile_k) + "u"); - defines.push_back("MAX_SUBGROUP_SIZE=" + std::to_string(context.max_subgroup_size) + "u"); - defines.push_back("SUBGROUP_M=" + std::to_string(WEBGPU_MUL_MAT_SUBGROUP_M) + "u"); - defines.push_back("SUBGROUP_N=" + std::to_string(WEBGPU_MUL_MAT_SUBGROUP_N) + "u"); - defines.push_back("SUBGROUP_MATRIX_M=" + std::to_string(WEBGPU_MUL_MAT_SUBGROUP_MATRIX_M) + "u"); - defines.push_back("SUBGROUP_MATRIX_N=" + std::to_string(WEBGPU_MUL_MAT_SUBGROUP_MATRIX_N) + "u"); - defines.push_back("SUBGROUP_MATRIX_M_SIZE=" + std::to_string(context.sg_mat_m) + "u"); - defines.push_back("SUBGROUP_MATRIX_N_SIZE=" + std::to_string(context.sg_mat_n) + "u"); - defines.push_back("SUBGROUP_MATRIX_K_SIZE=" + std::to_string(context.sg_mat_k) + "u"); - } - - // variant suffix for src1 type - variant += std::string("_") + (context.src1->type == GGML_TYPE_F32 ? "f32" : "f16"); - if (key.vectorized) { - variant += "_vectorized"; - } - - if (!key.use_subgroup_matrix) { - defines.push_back("WORKGROUP_SIZE_M=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_M) + "u"); - defines.push_back("WORKGROUP_SIZE_N=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_N) + "u"); - defines.push_back("TILE_K=" + std::to_string(tile_k) + "u"); - } - - auto processed = preprocessor.preprocess(shader_src, defines); - - auto decisions = std::make_shared(); - decisions->tile_k = tile_k; - decisions->tile_m = WEBGPU_MUL_MAT_TILE_M; - decisions->tile_n = WEBGPU_MUL_MAT_TILE_N; - decisions->use_subgroup_matrix = key.use_subgroup_matrix; - if (key.use_subgroup_matrix) { - decisions->subgroup_m = WEBGPU_MUL_MAT_SUBGROUP_M; - decisions->subgroup_n = WEBGPU_MUL_MAT_SUBGROUP_N; - decisions->subgroup_matrix_m = WEBGPU_MUL_MAT_SUBGROUP_MATRIX_M; - decisions->subgroup_matrix_n = WEBGPU_MUL_MAT_SUBGROUP_MATRIX_N; - decisions->wg_size = context.max_subgroup_size; - } else { - decisions->wg_size_m = WEBGPU_MUL_MAT_WG_SIZE_M; - decisions->wg_size_n = WEBGPU_MUL_MAT_WG_SIZE_N; - decisions->wg_size = WEBGPU_MUL_MAT_WG_SIZE_M * WEBGPU_MUL_MAT_WG_SIZE_N; - decisions->mul_mat_wg_size = WEBGPU_MUL_MAT_WG_SIZE; - } - - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - mul_mat_fast_pipelines[key] = pipeline; - return mul_mat_fast_pipelines[key]; - } - - webgpu_pipeline get_mul_mat_legacy_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_legacy_mul_mat_pipeline_key key = {}; - key.src0_type = context.src0->type; - key.src1_type = context.src1->type; - - auto it = mul_mat_legacy_pipelines.find(key); - if (it != mul_mat_legacy_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "mul_mat"; - - switch (context.src1->type) { - case GGML_TYPE_F32: - defines.push_back("SRC1_TYPE=f32"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("SRC1_TYPE=f16"); - variant += "_f16"; - break; - default: - GGML_ABORT("Unsupported src1 type for mul_mat legacy shader"); - } - - const struct ggml_type_traits * src0_traits = ggml_get_type_traits(context.src0->type); - const char * src0_name = src0_traits->type_name; - - switch (context.src0->type) { - case GGML_TYPE_F32: - defines.push_back("SRC0_TYPE=f32"); - defines.push_back("FLOAT"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("SRC0_TYPE=f16"); - defines.push_back("FLOAT"); - variant += "_f16"; - break; - default: - { - std::string type_upper = src0_name; - std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper); - - switch (context.src0->type) { - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q8_0: - case GGML_TYPE_Q3_K: - case GGML_TYPE_Q6_K: - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ2_S: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ4_NL: - { - // Quantized types using u32 buffers for portability. - defines.push_back("SRC0_TYPE=u32"); - defines.push_back("U32_DEQUANT_HELPERS"); - break; - } - default: - { - defines.push_back(std::string("SRC0_TYPE=") + src0_name); - } - } - - defines.push_back("BYTE_HELPERS"); - defines.push_back(type_upper + "_T"); - defines.push_back(type_upper); - defines.push_back(type_upper + "_SCALE_MIN"); - defines.push_back(type_upper + "_TABLES"); - defines.push_back(type_upper + "_GRID"); - - variant += std::string("_") + src0_name; - break; - } - } - - auto processed = preprocessor.preprocess(wgsl_mul_mat, defines); - - auto decisions = std::make_shared(); - decisions->wg_size = WEBGPU_MUL_MAT_WG_SIZE; - - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - mul_mat_legacy_pipelines[key] = pipeline; - return mul_mat_legacy_pipelines[key]; - } - - webgpu_pipeline get_mul_mat_id_gather_pipeline(const ggml_webgpu_shader_lib_context & context) { - auto it = mul_mat_id_gather_pipelines.find(1); - if (it != mul_mat_id_gather_pipelines.end()) { - return it->second; - } - std::vector defines; - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_mul_mat_id_gather, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, "mul_mat_id_gather"); - pipeline.context = decisions; - mul_mat_id_gather_pipelines[1] = pipeline; - return pipeline; - } - - webgpu_pipeline get_mul_mat_id_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_mul_mat_id_pipeline_key key = {}; - key.src0_type = context.src0->type; - key.src1_type = context.src1->type; - key.n_experts = context.src0->ne[2]; - key.vectorized = (context.src0->ne[0] % 4 == 0 && context.src0->ne[1] % 4 == 0 && - (context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ? - 1 : - 0; - - auto it = mul_mat_id_pipelines.find(key); - if (it != mul_mat_id_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "mul_mat_id"; - defines.push_back("MUL_MAT_ID"); - - // src1 type - switch (context.src1->type) { - case GGML_TYPE_F32: - defines.push_back("SRC1_INNER_TYPE=f32"); - break; - case GGML_TYPE_F16: - defines.push_back("SRC1_INNER_TYPE=f16"); - break; - default: - GGML_ABORT("Unsupported src1 type for mul_mat fast shader"); - } - - // src0 type - const struct ggml_type_traits * src0_traits = ggml_get_type_traits(context.src0->type); - const char * src0_name = src0_traits->type_name; - - switch (context.src0->type) { - case GGML_TYPE_F32: - defines.push_back("SRC0_INNER_TYPE=f32"); - defines.push_back("INIT_SRC0_SHMEM_FLOAT"); - defines.push_back("INIT_SRC1_SHMEM_FLOAT"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("SRC0_INNER_TYPE=f16"); - defines.push_back("INIT_SRC0_SHMEM_FLOAT"); - defines.push_back("INIT_SRC1_SHMEM_FLOAT"); - variant += "_f16"; - break; - default: - { - std::string type_upper = src0_name; - std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper); - - defines.push_back("BYTE_HELPERS"); - defines.push_back("INIT_SRC0_SHMEM_" + type_upper); - defines.push_back("INIT_SRC1_SHMEM_FLOAT"); - defines.push_back("U32_DEQUANT_HELPERS"); - defines.push_back("SRC0_INNER_TYPE=u32"); - - switch (context.src0->type) { - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - defines.push_back(type_upper + "_GRID"); - break; - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ2_S: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ3_S: - defines.push_back(type_upper + "_GRID"); - defines.push_back(type_upper + "_TABLES"); - break; - default: - break; - } - - variant += std::string("_") + src0_name; - break; - } - } - - // VEC/SCALAR controls - defines.push_back(key.vectorized ? "VEC" : "SCALAR"); - - // mul_mat_id is register-tile only. - const uint32_t tile_k = - ggml_is_quantized(context.src0->type) ? WEBGPU_MUL_MAT_REG_TILE_K_QUANT : WEBGPU_MUL_MAT_REG_TILE_K_FLOAT; - - // Tiles - defines.push_back("TILE_M=" + std::to_string(WEBGPU_MUL_MAT_TILE_M) + "u"); - defines.push_back("TILE_N=" + std::to_string(WEBGPU_MUL_MAT_TILE_N) + "u"); - defines.push_back("TILE_K=" + std::to_string(tile_k) + "u"); - - defines.push_back("WORKGROUP_SIZE_M=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_M) + "u"); - defines.push_back("WORKGROUP_SIZE_N=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_N) + "u"); - - // variant suffix for src1 type - variant += std::string("_") + (context.src1->type == GGML_TYPE_F32 ? "f32" : "f16"); - if (key.vectorized) { - variant += "_vectorized"; - } - - auto processed = preprocessor.preprocess(wgsl_mul_mat_id, defines); - - auto decisions = std::make_shared(); - decisions->tile_k = tile_k; - decisions->tile_m = WEBGPU_MUL_MAT_TILE_M; - decisions->tile_n = WEBGPU_MUL_MAT_TILE_N; - decisions->wg_size_m = WEBGPU_MUL_MAT_WG_SIZE_M; - decisions->wg_size_n = WEBGPU_MUL_MAT_WG_SIZE_N; - decisions->wg_size = WEBGPU_MUL_MAT_WG_SIZE_M * WEBGPU_MUL_MAT_WG_SIZE_N; - - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - mul_mat_id_pipelines[key] = pipeline; - return mul_mat_id_pipelines[key]; - } - - webgpu_pipeline get_mul_mat_id_vec_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_mul_mat_id_pipeline_key key = {}; - key.src0_type = context.src0->type; - key.src1_type = context.src1->type; - key.n_experts = context.src0->ne[2]; - key.vectorized = (context.src0->ne[0] % 4 == 0 && - (context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ? - 1 : - 0; - - auto it = mul_mat_id_vec_pipelines.find(key); - if (it != mul_mat_id_vec_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "mul_mat_id_vec"; - const char * shader_src = wgsl_mul_mat_id_vec; - - // src1 type - switch (context.src1->type) { - case GGML_TYPE_F32: - defines.push_back("SRC1_INNER_TYPE=f32"); - break; - case GGML_TYPE_F16: - defines.push_back("SRC1_INNER_TYPE=f16"); - break; - default: - GGML_ABORT("Unsupported src1 type for mul_mat fast shader"); - } - - // src0 type - switch (context.src0->type) { - case GGML_TYPE_F32: - defines.push_back("SRC0_INNER_TYPE=f32"); - defines.push_back("MUL_ACC_FLOAT"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("SRC0_INNER_TYPE=f16"); - defines.push_back("MUL_ACC_FLOAT"); - variant += "_f16"; - break; - default: - { - // Quantized types: use helpers but accumulate in f16 - const struct ggml_type_traits * src0_traits = ggml_get_type_traits(context.src0->type); - std::string src0_name = src0_traits->type_name; - std::string type_upper = src0_name; - variant += "_" + src0_name; - std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper); - - defines.push_back("BYTE_HELPERS"); - defines.push_back("MUL_ACC_" + type_upper); - defines.push_back("U32_DEQUANT_HELPERS"); - defines.push_back("SRC0_INNER_TYPE=u32"); - switch (context.src0->type) { - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ2_S: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - defines.push_back(type_upper + "_GRID"); - break; - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ3_XXS: - defines.push_back(type_upper + "_GRID"); - defines.push_back(type_upper + "_TABLES"); - break; - default: - break; - } - break; - } - } - - // VEC/SCALAR controls - defines.push_back(key.vectorized ? "VEC" : "SCALAR"); - - uint32_t wg_size = WEBGPU_MUL_MAT_VEC_WG_SIZE; - uint32_t outputs_per_wg = WEBGPU_MUL_MAT_VEC_FLOAT_OUTPUTS_PER_WG; - - if (key.src0_type == GGML_TYPE_Q1_0) { - outputs_per_wg = WEBGPU_MUL_MAT_VEC_LEGACY_Q_OUTPUTS_PER_WG; - } else if (key.src0_type >= GGML_TYPE_Q2_K) { - outputs_per_wg = WEBGPU_MUL_MAT_VEC_K_Q_OUTPUTS_PER_WG; - } else if (key.src0_type >= GGML_TYPE_Q4_0) { - outputs_per_wg = WEBGPU_MUL_MAT_VEC_LEGACY_Q_OUTPUTS_PER_WG; - } - - // variant suffix for src1 type - variant += std::string("_") + (context.src1->type == GGML_TYPE_F32 ? "f32" : "f16"); - - defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); - defines.push_back(std::string("OUTPUTS_PER_WG=") + std::to_string(outputs_per_wg)); - defines.push_back(context.supports_subgroups ? "USE_SUBGROUP_REDUCTION" : "USE_WORKGROUP_REDUCTION"); - variant += context.supports_subgroups ? "_sg_reduce" : "_wg_reduce"; - if (key.vectorized) { - variant += "_vectorized"; - } - - defines.push_back(std::string("N_EXPERTS=") + std::to_string(key.n_experts)); - - auto processed = preprocessor.preprocess(shader_src, defines); - - auto decisions = std::make_shared(); - decisions->wg_size = wg_size; - decisions->outputs_per_wg = outputs_per_wg; - - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - mul_mat_id_vec_pipelines[key] = pipeline; - return mul_mat_id_vec_pipelines[key]; - } - - webgpu_pipeline get_unary_pipeline(const ggml_webgpu_shader_lib_context & context) { - const bool is_unary = context.dst->op == GGML_OP_UNARY; - const int op = is_unary ? (int) ggml_get_unary_op(context.dst) : context.dst->op; - ggml_webgpu_unary_pipeline_key key = {}; - key.type = context.dst->type; - key.op = op; - key.is_unary = is_unary; - key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst) || context.dst->op == GGML_OP_FILL; - key.ttype = (ggml_tri_type) ggml_get_op_params_i32(context.dst, 0); - - auto it = unary_pipelines.find(key); - if (it != unary_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = - key.is_unary ? ggml_unary_op_name((ggml_unary_op) key.op) : ggml_op_name((ggml_op) key.op); - defines.push_back(variant); - - switch (key.type) { - case GGML_TYPE_F32: - defines.push_back("TYPE_F32"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("TYPE_F16"); - variant += "_f16"; - break; - default: - GGML_ABORT("Unsupported type for unary shader"); - } - - if (key.inplace) { - defines.push_back("INPLACE"); - variant += "_inplace"; - } - - if (op == GGML_OP_TRI) { - switch (key.ttype) { - case GGML_TRI_TYPE_LOWER: - defines.push_back("TRI_TYPE_LOWER"); - variant += "_tri_type_lower"; - break; - case GGML_TRI_TYPE_LOWER_DIAG: - defines.push_back("TRI_TYPE_LOWER_DIAG"); - variant += "_tri_type_lower_diag"; - break; - case GGML_TRI_TYPE_UPPER: - defines.push_back("TRI_TYPE_UPPER"); - variant += "_tri_type_upper"; - break; - case GGML_TRI_TYPE_UPPER_DIAG: - defines.push_back("TRI_TYPE_UPPER_DIAG"); - variant += "_tri_upper_diag"; - break; - default: - GGML_ABORT("Unsupported ggml_tri_type for unary shader"); - } - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_unary, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - decisions->inplace = key.inplace; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - unary_pipelines[key] = pipeline; - return unary_pipelines[key]; - } - - webgpu_pipeline get_rms_norm_mul_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_rms_norm_mul_pipeline_key key = {}; - key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); - key.overlap = ggml_webgpu_tensor_equal(context.src1, context.dst); - key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1); - - auto it = rms_norm_mul_pipelines.find(key); - if (it != rms_norm_mul_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string op_name = "RMS_NORM_MUL"; - std::string variant = op_name; - - if (key.inplace) { - defines.push_back("INPLACE"); - variant += "_inplace"; - } else if (key.overlap) { - defines.push_back("OVERLAP"); - variant += "_overlap"; - } else if (key.src_overlap) { - defines.push_back("SRC_OVERLAP"); - variant += "_src_overlap"; - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_rms_norm_mul, defines); - auto pipeline_decisions = std::make_shared(); - pipeline_decisions->wg_size = context.max_wg_size; - pipeline_decisions->inplace = key.inplace; - pipeline_decisions->overlap = key.overlap; - pipeline_decisions->src_overlap = key.src_overlap; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = pipeline_decisions; - rms_norm_mul_pipelines[key] = pipeline; - return rms_norm_mul_pipelines[key]; - } - - webgpu_pipeline get_binary_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_binary_pipeline_key key = {}; - key.type = context.dst->type; - key.op = context.dst->op; - key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); - key.overlap = ggml_webgpu_tensor_equal(context.src1, context.dst); - key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1); - - auto it = binary_pipelines.find(key); - if (it != binary_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string op_name = ggml_op_name((ggml_op) key.op); - std::string variant = op_name; - - defines.push_back(std::string("OP_") + op_name); - - switch (key.type) { - case GGML_TYPE_F32: - defines.push_back("TYPE_F32"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("TYPE_F16"); - variant += "_f16"; - break; - default: - GGML_ABORT("Unsupported type for binary shader"); - } - - if (key.inplace) { - defines.push_back("INPLACE"); - variant += "_inplace"; - } else if (key.overlap) { - defines.push_back("OVERLAP"); - variant += "_overlap"; - } else if (key.src_overlap) { - defines.push_back("SRC_OVERLAP"); - variant += "_src_overlap"; - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_binary, defines); - auto pipeline_decisions = std::make_shared(); - pipeline_decisions->wg_size = context.max_wg_size; - pipeline_decisions->inplace = key.inplace; - pipeline_decisions->overlap = key.overlap; - pipeline_decisions->src_overlap = key.src_overlap; - - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = pipeline_decisions; - binary_pipelines[key] = pipeline; - return binary_pipelines[key]; - } - - webgpu_pipeline get_concat_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_concat_pipeline_key key = {}; - key.type = context.dst->type; - - auto it = concat_pipelines.find(key); - if (it != concat_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "concat"; - - switch (key.type) { - case GGML_TYPE_F32: - defines.push_back("TYPE_F32"); - variant += "_f32"; - break; - case GGML_TYPE_I32: - defines.push_back("TYPE_I32"); - variant += "_i32"; - break; - default: - GGML_ABORT("Unsupported type for concat shader"); - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_concat, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - concat_pipelines[key] = pipeline; - return concat_pipelines[key]; - } - - webgpu_pipeline get_repeat_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_repeat_pipeline_key key = {}; - key.type = context.dst->type; - - auto it = repeat_pipelines.find(key); - if (it != repeat_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "repeat"; - - switch (key.type) { - case GGML_TYPE_F32: - defines.push_back("TYPE_F32"); - variant += "_f32"; - break; - case GGML_TYPE_I32: - defines.push_back("TYPE_I32"); - variant += "_i32"; - break; - case GGML_TYPE_I16: - defines.push_back("TYPE_I16"); - variant += "_i16"; - break; - default: - GGML_ABORT("Unsupported type for repeat shader"); - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_repeat, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - repeat_pipelines[key] = pipeline; - return repeat_pipelines[key]; - } - - webgpu_pipeline get_flash_attn_pipeline(const ggml_webgpu_shader_lib_context & context, - size_t storage_offset_alignment) { - const ggml_webgpu_flash_attn_decisions decisions = - ggml_webgpu_flash_attn_get_decisions(context, storage_offset_alignment); - GGML_ASSERT(decisions.path != GGML_WEBGPU_FLASH_ATTN_PATH_NONE); - ggml_webgpu_flash_attn_pipeline_key key = ggml_webgpu_flash_attn_make_pipeline_key(context, decisions.path); - auto it = flash_attn_pipelines.find(key); - if (it != flash_attn_pipelines.end()) { - return it->second; - } - std::vector defines; - std::string variant = decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_VEC ? "flash_attn_vec" : - decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE ? "flash_attn_tile" : - "flash_attn"; - - switch (key.kv_type) { - case GGML_TYPE_F32: - defines.push_back("KV_F32"); - break; - case GGML_TYPE_F16: - defines.push_back("KV_F16"); - break; - case GGML_TYPE_Q4_0: - defines.push_back("KV_Q4_0"); - break; - case GGML_TYPE_Q8_0: - defines.push_back("KV_Q8_0"); - break; - default: - GGML_ABORT("Unsupported KV type for flash attention shader"); - } - variant += std::string("_") + ggml_type_name(key.kv_type); - - if (key.has_mask) { - defines.push_back("MASK"); - if (key.path == GGML_WEBGPU_FLASH_ATTN_PATH_VEC) { - defines.push_back("BLK"); - variant += "_mask_blk"; - } else { - variant += "_mask"; - } - } - if (key.has_sinks) { - defines.push_back("SINKS"); - variant += "_sinks"; - } - if (key.uses_logit_softcap) { - defines.push_back("LOGIT_SOFTCAP"); - variant += "_lgsc"; - } - if (key.kv_direct) { - defines.push_back("KV_DIRECT"); - variant += "_kvdirect"; - } - if (key.kv_overlap) { - defines.push_back("KV_OVERLAP"); - variant += "_kv_overlap"; - } - - defines.push_back(std::string("HEAD_DIM_QK=") + std::to_string(key.head_dim_qk)); - variant += std::string("_hsqk") + std::to_string(key.head_dim_qk); - - defines.push_back(std::string("HEAD_DIM_V=") + std::to_string(key.head_dim_v)); - variant += std::string("_hsv") + std::to_string(key.head_dim_v); - - const char * shader_src = wgsl_flash_attn; - if (key.path == GGML_WEBGPU_FLASH_ATTN_PATH_VEC) { - defines.push_back("KV_GRANULARITY=8"); - defines.push_back(std::string("VEC_NE=") + std::to_string(ggml_webgpu_flash_attn_pick_vec_ne(key)) + "u"); - shader_src = wgsl_flash_attn_vec_split; - } else if (key.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE) { - shader_src = wgsl_flash_attn_tile; - defines.push_back("MAX_SUBGROUP_SIZE=" + std::to_string(context.max_subgroup_size)); - defines.push_back("KV_STAGE_STRIDE=" + std::to_string(std::max(key.head_dim_qk, key.head_dim_v))); - variant += "_tile"; - } else { - defines.push_back(std::string("SG_MAT_M=") + std::to_string(context.sg_mat_m)); - defines.push_back(std::string("SG_MAT_N=") + std::to_string(context.sg_mat_n)); - defines.push_back(std::string("SG_MAT_K=") + std::to_string(context.sg_mat_k)); - } - - auto pipeline_decisions = std::make_shared(decisions); - pipeline_decisions->kv_overlap = key.kv_overlap; - defines.push_back(std::string("Q_TILE=") + std::to_string(decisions.q_tile)); - defines.push_back(std::string("KV_TILE=") + std::to_string(decisions.kv_tile)); - defines.push_back(std::string("WG_SIZE=") + std::to_string(decisions.wg_size)); - - webgpu_pipeline pipeline = - ggml_webgpu_create_pipeline(device, preprocessor.preprocess(shader_src, defines), variant); - pipeline.context = pipeline_decisions; - flash_attn_pipelines[key] = pipeline; - return flash_attn_pipelines[key]; - } - - webgpu_pipeline get_flash_attn_blk_pipeline(const ggml_webgpu_shader_lib_context & context, uint32_t kv_tile) { - ggml_webgpu_flash_attn_blk_pipeline_key key = {}; - key.kv_tile = kv_tile; - auto it = flash_attn_blk_pipelines.find(key); - if (it != flash_attn_blk_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "flash_attn_vec_blk"; - - defines.push_back(std::string("KV_TILE=") + std::to_string(key.kv_tile)); - variant += std::string("_kvt") + std::to_string(key.kv_tile); - - uint32_t wg_size = 1; - while ((wg_size << 1) <= context.max_wg_size) { - wg_size <<= 1; - } - defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); - variant += std::string("_wg") + std::to_string(wg_size); - - webgpu_pipeline pipeline = - ggml_webgpu_create_pipeline(device, preprocessor.preprocess(wgsl_flash_attn_vec_blk, defines), variant); - flash_attn_blk_pipelines[key] = pipeline; - return flash_attn_blk_pipelines[key]; - } - - webgpu_pipeline get_flash_attn_vec_reduce_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_flash_attn_vec_reduce_pipeline_key key = {}; - key.head_dim_v = (uint32_t) context.src2->ne[0]; - key.wg_size = context.max_wg_size; - auto it = flash_attn_vec_reduce_pipelines.find(key); - if (it != flash_attn_vec_reduce_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "flash_attn_vec_reduce"; - - defines.push_back(std::string("HEAD_DIM_V=") + std::to_string(key.head_dim_v)); - variant += std::string("_hsv") + std::to_string(key.head_dim_v); - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - variant += std::string("_wg") + std::to_string(context.max_wg_size); - - webgpu_pipeline pipeline = - ggml_webgpu_create_pipeline(device, preprocessor.preprocess(wgsl_flash_attn_vec_reduce, defines), variant); - flash_attn_vec_reduce_pipelines[key] = pipeline; - return flash_attn_vec_reduce_pipelines[key]; - } - - webgpu_pipeline get_cpy_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_cpy_pipeline_key key = {}; - key.src_type = context.src0->type; - key.dst_type = context.dst->type; - - auto it = cpy_pipelines.find(key); - if (it != cpy_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "cpy"; - - switch (key.src_type) { - case GGML_TYPE_F32: - defines.push_back("SRC_F32"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("SRC_F16"); - variant += "_f16"; - break; - default: - GGML_ABORT("Unsupported src type for cpy shader"); - } - - switch (key.dst_type) { - case GGML_TYPE_F32: - defines.push_back("DST_F32"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("DST_F16"); - variant += "_f16"; - break; - case GGML_TYPE_I32: - defines.push_back("DST_I32"); - variant += "_i32"; - break; - default: - GGML_ABORT("Unsupported dst type for cpy shader"); - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_cpy, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - cpy_pipelines[key] = pipeline; - return cpy_pipelines[key]; - } - - webgpu_pipeline get_glu_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_glu_pipeline_key key = {}; - key.glu_op = ggml_get_glu_op(context.dst); - key.type = context.dst->type; - key.split = (context.src1 != nullptr); - - auto it = glu_pipelines.find(key); - if (it != glu_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "glu"; - - switch (key.glu_op) { - case GGML_GLU_OP_REGLU: - defines.push_back("OP_REGLU"); - variant += "_reglu"; - break; - case GGML_GLU_OP_GEGLU: - defines.push_back("OP_GEGLU"); - variant += "_geglu"; - break; - case GGML_GLU_OP_SWIGLU: - defines.push_back("OP_SWIGLU"); - variant += "_swiglu"; - break; - case GGML_GLU_OP_SWIGLU_OAI: - defines.push_back("OP_SWIGLU_OAI"); - variant += "_swiglu_oai"; - break; - case GGML_GLU_OP_GEGLU_ERF: - defines.push_back("OP_GEGLU_ERF"); - variant += "_geglu_erf"; - break; - case GGML_GLU_OP_GEGLU_QUICK: - defines.push_back("OP_GEGLU_QUICK"); - variant += "_geglu_quick"; - break; - default: - GGML_ABORT("Unsupported GLU op"); - } - switch (key.type) { - case GGML_TYPE_F32: - defines.push_back("TYPE_F32"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("TYPE_F16"); - variant += "_f16"; - break; - default: - GGML_ABORT("Unsupported type for GLU shader"); - } - - if (key.split) { - variant += "_split"; - } else { - defines.push_back("NO_SPLIT"); - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_glu, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - glu_pipelines[key] = pipeline; - return glu_pipelines[key]; - } - - webgpu_pipeline get_rope_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_rope_pipeline_key key = {}; - key.type = context.dst->type; - key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); - key.has_ff = (context.src2 != nullptr); - - auto it = rope_pipelines.find(key); - if (it != rope_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "rope"; - - switch (key.type) { - case GGML_TYPE_F32: - defines.push_back("TYPE_F32"); - variant += "_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("TYPE_F16"); - variant += "_f16"; - break; - default: - GGML_ABORT("Unsupported type for ROPE shader"); - } - - if (key.inplace) { - defines.push_back("INPLACE"); - variant += "_inplace"; - } - - if (key.has_ff) { - defines.push_back("FF_FUNC"); - variant += "_ff"; - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_rope, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - decisions->inplace = key.inplace; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - rope_pipelines[key] = pipeline; - return rope_pipelines[key]; - } - - webgpu_pipeline get_soft_max_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_soft_max_pipeline_key key = {}; - key.mask_type = context.src1 ? context.src1->type : GGML_TYPE_F32; - key.has_mask = (context.src1 != nullptr); - key.has_sink = (context.src2 != nullptr); - key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); - - auto it = soft_max_pipelines.find(key); - if (it != soft_max_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "soft_max"; - - if (key.has_mask) { - defines.push_back("HAS_MASK"); - switch (key.mask_type) { - case GGML_TYPE_F32: - defines.push_back("MASK_F32"); - variant += "_mask_f32"; - break; - case GGML_TYPE_F16: - defines.push_back("MASK_F16"); - variant += "_mask_f16"; - break; - default: - GGML_ABORT("Unsupported type for SOFT_MAX shader"); - } - } - - if (key.has_sink) { - defines.push_back("HAS_SINK"); - variant += "_sink"; - } - - if (key.inplace) { - defines.push_back("INPLACE"); - variant += "_inplace"; - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_soft_max, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - decisions->inplace = key.inplace; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - soft_max_pipelines[key] = pipeline; - return soft_max_pipelines[key]; - } - - webgpu_pipeline get_conv2d_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_conv2d_pipeline_key key = {}; - key.weight_type = context.src0->type; - key.input_type = context.src1->type; - key.output_type = context.dst->type; - - auto it = conv2d_pipelines.find(key); - if (it != conv2d_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "conv_2d"; - - auto push_type_defines = [&](const char * prefix, ggml_type type) { - std::string s_prefix = prefix; - if (type == GGML_TYPE_F32) { - defines.push_back(s_prefix + "_F32"); - } else if (type == GGML_TYPE_F16) { - defines.push_back(s_prefix + "_F16"); - } else { - GGML_ABORT("Unsupported type for CONV_2D shader"); - } - }; - - push_type_defines("WEIGHT", key.weight_type); - push_type_defines("INPUT", key.input_type); - push_type_defines("OUTPUT", key.output_type); - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_conv2d, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - conv2d_pipelines[key] = pipeline; - return conv2d_pipelines[key]; - } - - webgpu_pipeline get_im2col_pipeline(const ggml_webgpu_shader_lib_context & context) { - ggml_webgpu_im2col_pipeline_key key = {}; - key.input_type = context.src1->type; - key.output_type = context.dst->type; - - auto it = im2col_pipelines.find(key); - if (it != im2col_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "im2col"; - - auto push_type_defines = [&](const char * prefix, ggml_type type) { - std::string s_prefix = prefix; - if (type == GGML_TYPE_F32) { - defines.push_back(s_prefix + "_F32"); - } else if (type == GGML_TYPE_F16) { - defines.push_back(s_prefix + "_F16"); - } else { - GGML_ABORT("Unsupported type for IM2COL shader"); - } - }; - - push_type_defines("INPUT", key.input_type); - push_type_defines("OUTPUT", key.output_type); - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_im2col, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - im2col_pipelines[key] = pipeline; - return im2col_pipelines[key]; - } - - webgpu_pipeline get_upscale_pipeline(const ggml_webgpu_shader_lib_context & context) { - const uint32_t mode_flags = (uint32_t) ggml_get_op_params_i32(context.dst, 0); - const uint32_t base_mode = mode_flags & 0xFFu; - const bool antialias = (mode_flags & GGML_SCALE_FLAG_ANTIALIAS) != 0u; - - ggml_webgpu_upscale_pipeline_key key = {}; - key.input_type = context.src0->type; - key.output_type = context.dst->type; - key.base_mode = base_mode; - key.antialias = antialias; - - auto it = upscale_pipelines.find(key); - if (it != upscale_pipelines.end()) { - return it->second; - } - - std::vector defines; - std::string variant = "upscale"; - - if (key.input_type == GGML_TYPE_F16) { - defines.push_back("SRC_F16"); - variant += "_src_f16"; - } else { - variant += "_src_f32"; - } - - if (key.output_type == GGML_TYPE_F16) { - defines.push_back("DST_F16"); - variant += "_dst_f16"; - } else { - variant += "_dst_f32"; - } - - switch (base_mode) { - case GGML_SCALE_MODE_NEAREST: - defines.push_back("NEAREST"); - variant += "_nearest"; - break; - case GGML_SCALE_MODE_BILINEAR: - defines.push_back("BILINEAR"); - variant += "_bilinear"; - break; - case GGML_SCALE_MODE_BICUBIC: - defines.push_back("BICUBIC"); - variant += "_bicubic"; - break; - default: - GGML_ABORT("Unsupported upscale mode"); - } - - if (antialias) { - defines.push_back("ANTIALIAS"); - variant += "_aa"; - } - - defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); - - auto processed = preprocessor.preprocess(wgsl_upscale, defines); - auto decisions = std::make_shared(); - decisions->wg_size = context.max_wg_size; - webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); - pipeline.context = decisions; - upscale_pipelines[key] = pipeline; - return upscale_pipelines[key]; - } - - private: - static webgpu_pipeline ggml_webgpu_create_pipeline(wgpu::Device & device, - std::string shader_code, - std::string label) { - wgpu::ShaderSourceWGSL shader_source; - shader_source.code = shader_code.c_str(); - - wgpu::ShaderModuleDescriptor shader_desc; - shader_desc.nextInChain = &shader_source; - - wgpu::ShaderModule shader_module = device.CreateShaderModule(&shader_desc); - - wgpu::ComputePipelineDescriptor pipeline_desc; - pipeline_desc.label = label.c_str(); - pipeline_desc.compute.module = shader_module; - pipeline_desc.compute.entryPoint = "main"; // Entry point in the WGSL code - pipeline_desc.layout = nullptr; // nullptr means auto layout - return { device.CreateComputePipeline(&pipeline_desc), label }; - } -}; - -#endif // GGML_WEBGPU_SHADER_LIB_HPP diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp.inc similarity index 58% rename from ggml/src/ggml-webgpu/ggml-webgpu.cpp rename to ggml/src/ggml-webgpu/ggml-webgpu.cpp.inc index cab0aead1987..b653afa84e0c 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp.inc @@ -3,12 +3,3850 @@ Note: Use ClangFormat to format this file. */ -#include "ggml-webgpu.h" +#include "ggml-webgpu.h.inc" + +#include "ggml-backend-impl-defs.inc" +#include "ggml-impl-defs.inc" +#include "ggml-impl-defs.inc" +#include "ggml-wgsl-shaders.hpp" +#include "ggml.h.inc" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace pre_wgsl { + +//============================================================== +// Options +//============================================================== +struct Options { + std::string include_path = "."; + std::vector macros; +}; + +//============================================================== +// Utility: trim +//============================================================== +static std::string trim(const std::string & s) { + size_t a = 0; + while (a < s.size() && std::isspace((unsigned char) s[a])) { + a++; + } + size_t b = s.size(); + while (b > a && std::isspace((unsigned char) s[b - 1])) { + b--; + } + return s.substr(a, b - a); +} + +static std::string trim_value(std::istream & is) { + std::string str; + std::getline(is, str); + return trim(str); +} + +static bool isIdentChar(char c) { + return std::isalnum(static_cast(c)) || c == '_'; +} + +static std::string expandMacrosRecursiveInternal(const std::string & line, + const std::unordered_map & macros, + std::unordered_set & visiting); + +static std::string expandMacroValue(const std::string & name, + const std::unordered_map & macros, + std::unordered_set & visiting) { + if (visiting.count(name)) { + throw std::runtime_error("Recursive macro: " + name); + } + visiting.insert(name); + + auto it = macros.find(name); + if (it == macros.end()) { + visiting.erase(name); + return name; + } + + const std::string & value = it->second; + if (value.empty()) { + visiting.erase(name); + return ""; + } + + std::string expanded = expandMacrosRecursiveInternal(value, macros, visiting); + visiting.erase(name); + return expanded; +} + +static std::string expandMacrosRecursiveInternal(const std::string & line, + const std::unordered_map & macros, + std::unordered_set & visiting) { + std::string result; + result.reserve(line.size()); + + size_t i = 0; + while (i < line.size()) { + if (isIdentChar(line[i])) { + size_t start = i; + while (i < line.size() && isIdentChar(line[i])) { + i++; + } + std::string token = line.substr(start, i - start); + + auto it = macros.find(token); + if (it != macros.end()) { + result += expandMacroValue(token, macros, visiting); + } else { + result += token; + } + } else { + result += line[i]; + i++; + } + } + + return result; +} + +static std::string expandMacrosRecursive(const std::string & line, + const std::unordered_map & macros) { + std::unordered_set visiting; + return expandMacrosRecursiveInternal(line, macros, visiting); +} + +//============================================================== +// Tokenizer for expressions in #if/#elif +//============================================================== +class ExprLexer { + public: + enum Kind { END, IDENT, NUMBER, OP, LPAREN, RPAREN }; + + struct Tok { + Kind kind; + std::string text; + }; + + explicit ExprLexer(std::string_view sv) : src(sv), pos(0) {} + + Tok next() { + skipWS(); + if (pos >= src.size()) { + return { END, "" }; + } + + char c = src[pos]; + + // number + if (std::isdigit((unsigned char) c)) { + size_t start = pos; + while (pos < src.size() && std::isdigit((unsigned char) src[pos])) { + pos++; + } + return { NUMBER, std::string(src.substr(start, pos - start)) }; + } + + // identifier + if (std::isalpha((unsigned char) c) || c == '_') { + size_t start = pos; + while (pos < src.size() && (std::isalnum((unsigned char) src[pos]) || src[pos] == '_')) { + pos++; + } + return { IDENT, std::string(src.substr(start, pos - start)) }; + } + + if (c == '(') { + pos++; + return { LPAREN, "(" }; + } + if (c == ')') { + pos++; + return { RPAREN, ")" }; + } + + // multi-char operators + static const char * two_ops[] = { "==", "!=", "<=", ">=", "&&", "||", "<<", ">>" }; + for (auto op : two_ops) { + if (src.substr(pos, 2) == op) { + pos += 2; + return { OP, std::string(op) }; + } + } + + // single-char operators + if (std::string("+-*/%<>!").find(c) != std::string::npos) { + pos++; + return { OP, std::string(1, c) }; + } + + // unexpected + pos++; + return { END, "" }; + } + + private: + std::string_view src; + size_t pos; + + void skipWS() { + while (pos < src.size() && std::isspace((unsigned char) src[pos])) { + pos++; + } + } +}; + +//============================================================== +// Expression Parser (recursive descent) +//============================================================== +class ExprParser { + public: + ExprParser(std::string_view expr, + const std::unordered_map & macros, + std::unordered_set & visiting) : + lex(expr), + macros(macros), + visiting(visiting) { + advance(); + } + + int parse() { return parseLogicalOr(); } + + private: + ExprLexer lex; + ExprLexer::Tok tok; + const std::unordered_map & macros; + std::unordered_set & visiting; + + void advance() { tok = lex.next(); } + + bool acceptOp(const std::string & s) { + if (tok.kind == ExprLexer::OP && tok.text == s) { + advance(); + return true; + } + return false; + } + + bool acceptKind(ExprLexer::Kind k) { + if (tok.kind == k) { + advance(); + return true; + } + return false; + } + + int parseLogicalOr() { + int v = parseLogicalAnd(); + while (acceptOp("||")) { + int rhs = parseLogicalAnd(); + v = (v || rhs); + } + return v; + } + + int parseLogicalAnd() { + int v = parseEquality(); + while (acceptOp("&&")) { + int rhs = parseEquality(); + v = (v && rhs); + } + return v; + } + + int parseEquality() { + int v = parseRelational(); + for (;;) { + if (acceptOp("==")) { + int rhs = parseRelational(); + v = (v == rhs); + } else if (acceptOp("!=")) { + int rhs = parseRelational(); + v = (v != rhs); + } else { + break; + } + } + return v; + } + + int parseRelational() { + int v = parseShift(); + for (;;) { + if (acceptOp("<")) { + int rhs = parseShift(); + v = (v < rhs); + } else if (acceptOp(">")) { + int rhs = parseShift(); + v = (v > rhs); + } else if (acceptOp("<=")) { + int rhs = parseShift(); + v = (v <= rhs); + } else if (acceptOp(">=")) { + int rhs = parseShift(); + v = (v >= rhs); + } else { + break; + } + } + return v; + } + + int parseShift() { + int v = parseAdd(); + for (;;) { + if (acceptOp("<<")) { + int rhs = parseAdd(); + v = (v << rhs); + } else if (acceptOp(">>")) { + int rhs = parseAdd(); + v = (v >> rhs); + } else { + break; + } + } + return v; + } + + int parseAdd() { + int v = parseMult(); + for (;;) { + if (acceptOp("+")) { + int rhs = parseMult(); + v = (v + rhs); + } else if (acceptOp("-")) { + int rhs = parseMult(); + v = (v - rhs); + } else { + break; + } + } + return v; + } + + int parseMult() { + int v = parseUnary(); + for (;;) { + if (acceptOp("*")) { + int rhs = parseUnary(); + v = (v * rhs); + } else if (acceptOp("/")) { + int rhs = parseUnary(); + v = (rhs == 0 ? 0 : v / rhs); + } else if (acceptOp("%")) { + int rhs = parseUnary(); + v = (rhs == 0 ? 0 : v % rhs); + } else { + break; + } + } + return v; + } + + int parseUnary() { + if (acceptOp("!")) { + return !parseUnary(); + } + if (acceptOp("-")) { + return -parseUnary(); + } + if (acceptOp("+")) { + return +parseUnary(); + } + return parsePrimary(); + } + + int parsePrimary() { + // '(' expr ')' + if (acceptKind(ExprLexer::LPAREN)) { + int v = parse(); + if (!acceptKind(ExprLexer::RPAREN)) { + throw std::runtime_error("missing ')'"); + } + return v; + } + + // number + if (tok.kind == ExprLexer::NUMBER) { + int v = std::stoi(tok.text); + advance(); + return v; + } + + // defined(identifier) + if (tok.kind == ExprLexer::IDENT && tok.text == "defined") { + advance(); + if (acceptKind(ExprLexer::LPAREN)) { + if (tok.kind != ExprLexer::IDENT) { + throw std::runtime_error("expected identifier in defined()"); + } + std::string name = tok.text; + advance(); + if (!acceptKind(ExprLexer::RPAREN)) { + throw std::runtime_error("missing ) in defined()"); + } + return macros.count(name) ? 1 : 0; + } else { + // defined NAME + if (tok.kind != ExprLexer::IDENT) { + throw std::runtime_error("expected identifier in defined NAME"); + } + std::string name = tok.text; + advance(); + return macros.count(name) ? 1 : 0; + } + } + + // identifier -> treat as integer, if defined use its value else 0 + if (tok.kind == ExprLexer::IDENT) { + std::string name = tok.text; + advance(); + auto it = macros.find(name); + if (it == macros.end()) { + return 0; + } + if (it->second.empty()) { + return 1; + } + return evalMacroExpression(name, it->second); + } + + // unexpected + return 0; + } + + int evalMacroExpression(const std::string & name, const std::string & value) { + if (visiting.count(name)) { + throw std::runtime_error("Recursive macro: " + name); + } + + visiting.insert(name); + ExprParser ep(value, macros, visiting); + int v = ep.parse(); + visiting.erase(name); + return v; + } +}; + +//============================================================== +// Preprocessor +//============================================================== +class Preprocessor { + public: + explicit Preprocessor(Options opts = {}) : opts_(std::move(opts)) { + // Treat empty include path as current directory + if (opts_.include_path.empty()) { + opts_.include_path = "."; + } + parseMacroDefinitions(opts_.macros); + } + + std::string preprocess_file(const std::string & filename, const std::vector & additional_macros = {}) { + std::unordered_map macros; + std::unordered_set predefined; + std::unordered_set include_stack; + buildMacros(additional_macros, macros, predefined); + + std::string result = processFile(filename, macros, predefined, include_stack, DirectiveMode::All); + return result; + } + + std::string preprocess(const std::string & contents, const std::vector & additional_macros = {}) { + std::unordered_map macros; + std::unordered_set predefined; + std::unordered_set include_stack; + buildMacros(additional_macros, macros, predefined); + + std::string result = processString(contents, macros, predefined, include_stack, DirectiveMode::All); + return result; + } + + std::string preprocess_includes_file(const std::string & filename) { + std::unordered_map macros; + std::unordered_set predefined; + std::unordered_set include_stack; + std::string result = processFile(filename, macros, predefined, include_stack, DirectiveMode::IncludesOnly); + return result; + } + + std::string preprocess_includes(const std::string & contents) { + std::unordered_map macros; + std::unordered_set predefined; + std::unordered_set include_stack; + std::string result = processString(contents, macros, predefined, include_stack, DirectiveMode::IncludesOnly); + return result; + } + + private: + Options opts_; + std::unordered_map global_macros; + + enum class DirectiveMode { All, IncludesOnly }; + + struct Cond { + bool parent_active; + bool active; + bool taken; + }; + + //---------------------------------------------------------- + // Parse macro definitions into global_macros + //---------------------------------------------------------- + void parseMacroDefinitions(const std::vector & macro_defs) { + for (const auto & def : macro_defs) { + size_t eq_pos = def.find('='); + if (eq_pos != std::string::npos) { + // Format: NAME=VALUE + std::string name = trim(def.substr(0, eq_pos)); + std::string value = trim(def.substr(eq_pos + 1)); + global_macros[name] = value; + } else { + // Format: NAME + std::string name = trim(def); + global_macros[name] = ""; + } + } + } + + //---------------------------------------------------------- + // Build combined macro map and predefined set for a preprocessing operation + //---------------------------------------------------------- + void buildMacros(const std::vector & additional_macros, + std::unordered_map & macros, + std::unordered_set & predefined) { + macros = global_macros; + predefined.clear(); + + for (const auto & [name, value] : global_macros) { + predefined.insert(name); + } + + for (const auto & def : additional_macros) { + size_t eq_pos = def.find('='); + std::string name, value; + if (eq_pos != std::string::npos) { + name = trim(def.substr(0, eq_pos)); + value = trim(def.substr(eq_pos + 1)); + } else { + name = trim(def); + value = ""; + } + + // Add to macros map (will override global if same name) + macros[name] = value; + predefined.insert(name); + } + } + + //---------------------------------------------------------- + // Helpers + //---------------------------------------------------------- + std::string loadFile(const std::string & fname) { + std::ifstream f(fname); + if (!f.is_open()) { + throw std::runtime_error("Could not open file: " + fname); + } + std::stringstream ss; + ss << f.rdbuf(); + return ss.str(); + } + + bool condActive(const std::vector & cond) const { + if (cond.empty()) { + return true; + } + return cond.back().active; + } + + //---------------------------------------------------------- + // Process a file + //---------------------------------------------------------- + std::string processFile(const std::string & name, + std::unordered_map & macros, + const std::unordered_set & predefined_macros, + std::unordered_set & include_stack, + DirectiveMode mode) { + if (include_stack.count(name)) { + throw std::runtime_error("Recursive include: " + name); + } + + include_stack.insert(name); + std::string shader_code = loadFile(name); + std::string out = processString(shader_code, macros, predefined_macros, include_stack, mode); + include_stack.erase(name); + return out; + } + + std::string processIncludeFile(const std::string & fname, + std::unordered_map & macros, + const std::unordered_set & predefined_macros, + std::unordered_set & include_stack, + DirectiveMode mode) { + std::string full_path = opts_.include_path + "/" + fname; + return processFile(full_path, macros, predefined_macros, include_stack, mode); + } + + //---------------------------------------------------------- + // Process text + //---------------------------------------------------------- + std::string processString(const std::string & shader_code, + std::unordered_map & macros, + const std::unordered_set & predefined_macros, + std::unordered_set & include_stack, + DirectiveMode mode) { + std::vector cond; // Conditional stack for this shader + std::stringstream out; + std::istringstream in(shader_code); + std::string line; + + while (std::getline(in, line)) { + std::string t = trim(line); + + if (!t.empty() && t[0] == '#') { + bool handled = handleDirective(t, out, macros, predefined_macros, cond, include_stack, mode); + if (mode == DirectiveMode::IncludesOnly && !handled) { + out << line << "\n"; + } + } else { + if (mode == DirectiveMode::IncludesOnly) { + out << line << "\n"; + } else if (condActive(cond)) { + // Expand macros in the line before outputting + std::string expanded = expandMacrosRecursive(line, macros); + out << expanded << "\n"; + } + } + } + + if (mode == DirectiveMode::All && !cond.empty()) { + throw std::runtime_error("Unclosed #if directive"); + } + + return out.str(); + } + + //---------------------------------------------------------- + // Directive handler + //---------------------------------------------------------- + bool handleDirective(const std::string & t, + std::stringstream & out, + std::unordered_map & macros, + const std::unordered_set & predefined_macros, + std::vector & cond, + std::unordered_set & include_stack, + DirectiveMode mode) { + // split into tokens + std::string body = t.substr(1); + std::istringstream iss(body); + std::string cmd; + iss >> cmd; + + if (cmd == "include") { + if (mode == DirectiveMode::All && !condActive(cond)) { + return true; + } + std::string file; + iss >> file; + if (file.size() >= 2 && file.front() == '"' && file.back() == '"') { + file = file.substr(1, file.size() - 2); + } + out << processIncludeFile(file, macros, predefined_macros, include_stack, mode); + return true; + } + + if (mode == DirectiveMode::IncludesOnly) { + return false; + } + + if (cmd == "define") { + if (!condActive(cond)) { + return true; + } + std::string name; + iss >> name; + // Don't override predefined macros from options + if (predefined_macros.count(name)) { + return true; + } + std::string value = trim_value(iss); + macros[name] = value; + return true; + } + + if (cmd == "undef") { + if (!condActive(cond)) { + return true; + } + std::string name; + iss >> name; + // Don't undef predefined macros from options + if (predefined_macros.count(name)) { + return true; + } + macros.erase(name); + return true; + } + + if (cmd == "ifdef") { + std::string name; + iss >> name; + bool p = condActive(cond); + bool v = macros.count(name); + cond.push_back({ p, p && v, p && v }); + return true; + } + + if (cmd == "ifndef") { + std::string name; + iss >> name; + bool p = condActive(cond); + bool v = !macros.count(name); + cond.push_back({ p, p && v, p && v }); + return true; + } + + if (cmd == "if") { + std::string expr = trim_value(iss); + bool p = condActive(cond); + bool v = false; + if (p) { + std::unordered_set visiting; + ExprParser ep(expr, macros, visiting); + v = ep.parse() != 0; + } + cond.push_back({ p, p && v, p && v }); + return true; + } + + if (cmd == "elif") { + std::string expr = trim_value(iss); + + if (cond.empty()) { + throw std::runtime_error("#elif without #if"); + } + + Cond & c = cond.back(); + if (!c.parent_active) { + c.active = false; + return true; + } + + if (c.taken) { + c.active = false; + return true; + } + + std::unordered_set visiting; + ExprParser ep(expr, macros, visiting); + bool v = ep.parse() != 0; + c.active = v; + if (v) { + c.taken = true; + } + return true; + } + + if (cmd == "else") { + if (cond.empty()) { + throw std::runtime_error("#else without #if"); + } + + Cond & c = cond.back(); + if (!c.parent_active) { + c.active = false; + return true; + } + if (c.taken) { + c.active = false; + } else { + c.active = true; + c.taken = true; + } + return true; + } + + if (cmd == "endif") { + if (cond.empty()) { + throw std::runtime_error("#endif without #if"); + } + cond.pop_back(); + return true; + } + + // Unknown directive + throw std::runtime_error("Unknown directive: #" + cmd); + } +}; + +} // namespace pre_wgsl + + + +#include + +#include +#include +#include +#include +#include + +#define GGML_WEBGPU_F16_SIZE_BYTES 2 +#define GGML_WEBGPU_F32_SIZE_BYTES 4 +#define GGML_WEBGPU_I32_SIZE_BYTES 4 +#define GGML_WEBGPU_FLASH_ATTN_PREFERRED_KV_SG_TILES 8u +#define GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE 128u +// Matches GGML_PAD(..., 256) in src/llama-context.cpp for KV cache sizing. +#define GGML_WEBGPU_KV_SEQ_PAD 256u + +#define GGML_WEBGPU_ARGSORT_MERGE_MAX_WG_SIZE 512u + +// Matrix multiplication parameters + +// Register tiling parameters +#define WEBGPU_MUL_MAT_TILE_M 4 +#define WEBGPU_MUL_MAT_TILE_N 4 +#define WEBGPU_MUL_MAT_WG_SIZE_M 8 +#define WEBGPU_MUL_MAT_WG_SIZE_N 8 +#define WEBGPU_MUL_MAT_REG_TILE_K_FLOAT 8 +#define WEBGPU_MUL_MAT_REG_TILE_K_QUANT 32 + +// Subgroup matrix parameters +// The number of subgroups in the M dimension +#define WEBGPU_MUL_MAT_SUBGROUP_M 2 +// The number of subgroups in the N dimension +#define WEBGPU_MUL_MAT_SUBGROUP_N 4 +// The number of subgroup matrices each subgroup accumulates over +#define WEBGPU_MUL_MAT_SUBGROUP_MATRIX_M 4 +#define WEBGPU_MUL_MAT_SUBGROUP_MATRIX_N 2 +#define WEBGPU_MUL_MAT_SUBGROUP_TILE_K_FLOAT 32 +#define WEBGPU_MUL_MAT_SUBGROUP_TILE_K_QUANT 32 + +// Matrix-vector multiplication parameters +#define WEBGPU_MUL_MAT_VEC_WG_SIZE 256 + +#define WEBGPU_MUL_MAT_VEC_FLOAT_OUTPUTS_PER_WG 4 +#define WEBGPU_MUL_MAT_VEC_LEGACY_Q_OUTPUTS_PER_WG 4 +#define WEBGPU_MUL_MAT_VEC_K_Q_OUTPUTS_PER_WG 4 + +// default size for legacy matrix multiplication +#define WEBGPU_MUL_MAT_WG_SIZE 256 + +// Same hash combine function as in boost +template inline void ggml_webgpu_hash_combine(size_t & seed, const T & value) { + seed ^= std::hash{}(value) + 0x9e3779b9 + (seed << 6) + (seed >> 2); +} + +// Calculates base address of a tensor ignoring the fake base pointer +inline uintptr_t ggml_webgpu_tensor_addr(const ggml_tensor * tensor) { + const ggml_tensor * base_tensor = tensor->view_src ? tensor->view_src : tensor; + return (uintptr_t) base_tensor->data + tensor->view_offs; +} + +inline bool ggml_webgpu_tensor_equal(const ggml_tensor * a, const ggml_tensor * b) { + return a->buffer == b->buffer && ggml_webgpu_tensor_addr(a) == ggml_webgpu_tensor_addr(b); +} + +inline bool ggml_webgpu_tensor_overlap(const ggml_tensor * a, const ggml_tensor * b) { + return a->buffer == b->buffer && ggml_webgpu_tensor_addr(a) < ggml_webgpu_tensor_addr(b) + ggml_nbytes(b) && + ggml_webgpu_tensor_addr(b) < ggml_webgpu_tensor_addr(a) + ggml_nbytes(a); +} + +struct ggml_webgpu_shader_lib_context { + ggml_tensor * src0; + ggml_tensor * src1; + ggml_tensor * src2; + ggml_tensor * src3; + ggml_tensor * src4; + ggml_tensor * src5; + ggml_tensor * dst; + + uint32_t max_wg_size; + size_t wg_mem_limit_bytes = 0; + bool supports_subgroups = false; + bool supports_subgroup_matrix = false; + uint32_t sg_mat_m = 0; + uint32_t sg_mat_n = 0; + uint32_t sg_mat_k = 0; + uint32_t max_subgroup_size = 0; +}; + +struct webgpu_pipeline { + wgpu::ComputePipeline pipeline; + std::string name; + std::shared_ptr context = nullptr; +}; + +struct ggml_webgpu_generic_shader_decisions { + uint32_t wg_size = 0; + bool inplace = false; +}; + +struct ggml_webgpu_binary_shader_decisions { + uint32_t wg_size = 0; + bool inplace = false; + bool overlap = false; + bool src_overlap = false; +}; + +struct ggml_webgpu_processed_shader { + std::string wgsl; + std::string variant; + std::shared_ptr decisions; +}; + +struct ggml_webgpu_ssm_conv_shader_decisions { + uint32_t block_size; + uint32_t tokens_per_wg; +}; + +struct ggml_webgpu_ssm_scan_pipeline_key { + int type; + int d_state; + bool xbc_overlap; + + bool operator==(const ggml_webgpu_ssm_scan_pipeline_key & other) const { + return type == other.type && d_state == other.d_state && xbc_overlap == other.xbc_overlap; + } +}; + +struct ggml_webgpu_ssm_scan_pipeline_key_hash { + size_t operator()(const ggml_webgpu_ssm_scan_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.type); + ggml_webgpu_hash_combine(seed, key.d_state); + ggml_webgpu_hash_combine(seed, key.xbc_overlap); + return seed; + } +}; + +struct ggml_webgpu_ssm_scan_shader_decisions { + uint32_t wg_size; + uint32_t tokens_per_tile; + bool xbc_overlap = false; +}; + +/** Argsort **/ + +struct ggml_webgpu_argsort_shader_lib_context { + uint32_t max_wg_size; + size_t wg_mem_limit_bytes; + int32_t order; +}; + +/** Set Rows **/ + +struct ggml_webgpu_set_rows_pipeline_key { + int dst_type; + int vec4; + int i64_idx; + + bool operator==(const ggml_webgpu_set_rows_pipeline_key & other) const { + return dst_type == other.dst_type && vec4 == other.vec4 && i64_idx == other.i64_idx; + } +}; + +struct ggml_webgpu_set_rows_pipeline_key_hash { + size_t operator()(const ggml_webgpu_set_rows_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.dst_type); + ggml_webgpu_hash_combine(seed, key.vec4); + ggml_webgpu_hash_combine(seed, key.i64_idx); + return seed; + } +}; + +struct ggml_webgpu_set_rows_shader_decisions { + bool vec4; + bool i64_idx; + uint32_t wg_size; +}; + +/** Set **/ + +struct ggml_webgpu_set_pipeline_key { + ggml_type type; + bool inplace; + + bool operator==(const ggml_webgpu_set_pipeline_key & other) const { + return type == other.type && inplace == other.inplace; + } +}; + +struct ggml_webgpu_set_pipeline_key_hash { + size_t operator()(const ggml_webgpu_set_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.type); + ggml_webgpu_hash_combine(seed, key.inplace); + return seed; + } +}; + +/** Get Rows **/ + +struct ggml_webgpu_get_rows_pipeline_key { + ggml_type src_type; + int vectorized; + + bool operator==(const ggml_webgpu_get_rows_pipeline_key & other) const { + return src_type == other.src_type && vectorized == other.vectorized; + } +}; + +struct ggml_webgpu_get_rows_pipeline_key_hash { + size_t operator()(const ggml_webgpu_get_rows_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.src_type); + ggml_webgpu_hash_combine(seed, key.vectorized); + return seed; + } +}; + +/** Row Norm **/ + +struct ggml_webgpu_row_norm_pipeline_key { + ggml_op op; + bool inplace; + + bool operator==(const ggml_webgpu_row_norm_pipeline_key & other) const { + return op == other.op && inplace == other.inplace; + } +}; + +struct ggml_webgpu_row_norm_pipeline_key_hash { + size_t operator()(const ggml_webgpu_row_norm_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.op); + ggml_webgpu_hash_combine(seed, key.inplace); + return seed; + } +}; + +/** RMS_NORM + MUL **/ + +struct ggml_webgpu_rms_norm_mul_pipeline_key { + bool inplace; // rn_src == dst + bool overlap; // mul_src == dst + bool src_overlap; // rn_src == mul_src + + bool operator==(const ggml_webgpu_rms_norm_mul_pipeline_key & other) const { + return inplace == other.inplace && overlap == other.overlap && src_overlap == other.src_overlap; + } +}; + +struct ggml_webgpu_rms_norm_mul_pipeline_key_hash { + size_t operator()(const ggml_webgpu_rms_norm_mul_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.inplace); + ggml_webgpu_hash_combine(seed, key.overlap); + ggml_webgpu_hash_combine(seed, key.src_overlap); + return seed; + } +}; + +struct ggml_webgpu_rms_norm_mul_shader_decisions { + uint32_t wg_size = 0; + bool inplace = false; + bool overlap = false; + bool src_overlap = false; +}; + +/** Pad **/ +struct ggml_webgpu_pad_pipeline_key { + bool circular; + + bool operator==(const ggml_webgpu_pad_pipeline_key & other) const { return circular == other.circular; } +}; + +struct ggml_webgpu_pad_pipeline_key_hash { + size_t operator()(const ggml_webgpu_pad_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.circular); + return seed; + } +}; + +/** Solve Tri **/ +struct ggml_webgpu_solve_tri_pipeline_key { + int type; + int n; + int k; + + bool operator==(const ggml_webgpu_solve_tri_pipeline_key & other) const { + return type == other.type && n == other.n && k == other.k; + } +}; + +struct ggml_webgpu_solve_tri_pipeline_key_hash { + size_t operator()(const ggml_webgpu_solve_tri_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.type); + ggml_webgpu_hash_combine(seed, key.n); + ggml_webgpu_hash_combine(seed, key.k); + return seed; + } +}; + +/** SSM Conv **/ +struct ggml_webgpu_ssm_conv_pipeline_key { + int type; + int vectorized; + + bool operator==(const ggml_webgpu_ssm_conv_pipeline_key & other) const { + return type == other.type && vectorized == other.vectorized; + } +}; + +/** CONV 2D */ +struct ggml_webgpu_conv2d_pipeline_key { + ggml_type weight_type; + ggml_type input_type; + ggml_type output_type; + + bool operator==(const ggml_webgpu_conv2d_pipeline_key & other) const { + return weight_type == other.weight_type && input_type == other.input_type && output_type == other.output_type; + } +}; + +struct ggml_webgpu_conv2d_pipeline_key_hash { + size_t operator()(const ggml_webgpu_conv2d_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.weight_type); + ggml_webgpu_hash_combine(seed, key.input_type); + ggml_webgpu_hash_combine(seed, key.output_type); + return seed; + } +}; + +/** Im2Col **/ +struct ggml_webgpu_im2col_pipeline_key { + ggml_type input_type; + ggml_type output_type; + + bool operator==(const ggml_webgpu_im2col_pipeline_key & other) const { + return input_type == other.input_type && output_type == other.output_type; + } +}; + +struct ggml_webgpu_im2col_pipeline_key_hash { + size_t operator()(const ggml_webgpu_im2col_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.input_type); + ggml_webgpu_hash_combine(seed, key.output_type); + return seed; + } +}; + +/** Gated Delta Net **/ +struct ggml_webgpu_gated_delta_net_pipeline_key { + int type; + int s_v; + int kda; + + bool operator==(const ggml_webgpu_gated_delta_net_pipeline_key & other) const { + return type == other.type && s_v == other.s_v && kda == other.kda; + } +}; + +struct ggml_webgpu_gated_delta_net_pipeline_key_hash { + size_t operator()(const ggml_webgpu_gated_delta_net_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.type); + ggml_webgpu_hash_combine(seed, key.s_v); + ggml_webgpu_hash_combine(seed, key.kda); + return seed; + } +}; + +struct ggml_webgpu_ssm_conv_pipeline_key_hash { + size_t operator()(const ggml_webgpu_ssm_conv_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.type); + ggml_webgpu_hash_combine(seed, key.vectorized); + return seed; + } +}; + +/** Scale **/ + +struct ggml_webgpu_scale_pipeline_key { + int inplace; + + bool operator==(const ggml_webgpu_scale_pipeline_key & other) const { return inplace == other.inplace; } +}; + +struct ggml_webgpu_scale_pipeline_key_hash { + size_t operator()(const ggml_webgpu_scale_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.inplace); + return seed; + } +}; + +/** Upscale **/ + +struct ggml_webgpu_upscale_pipeline_key { + ggml_type input_type; + ggml_type output_type; + uint32_t base_mode; + bool antialias; + + bool operator==(const ggml_webgpu_upscale_pipeline_key & other) const { + return input_type == other.input_type && output_type == other.output_type && base_mode == other.base_mode && + antialias == other.antialias; + } +}; + +struct ggml_webgpu_upscale_pipeline_key_hash { + size_t operator()(const ggml_webgpu_upscale_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.input_type); + ggml_webgpu_hash_combine(seed, key.output_type); + ggml_webgpu_hash_combine(seed, key.base_mode); + ggml_webgpu_hash_combine(seed, key.antialias); + return seed; + } +}; + +/** Concat **/ + +struct ggml_webgpu_concat_pipeline_key { + int type; + + bool operator==(const ggml_webgpu_concat_pipeline_key & other) const { return type == other.type; } +}; + +struct ggml_webgpu_concat_pipeline_key_hash { + size_t operator()(const ggml_webgpu_concat_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.type); + return seed; + } +}; + +/** Repeat **/ + +struct ggml_webgpu_repeat_pipeline_key { + int type; + + bool operator==(const ggml_webgpu_repeat_pipeline_key & other) const { return type == other.type; } +}; + +struct ggml_webgpu_repeat_pipeline_key_hash { + size_t operator()(const ggml_webgpu_repeat_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.type); + return seed; + } +}; + +/** Binary **/ + +struct ggml_webgpu_binary_pipeline_key { + int type; + int op; + bool inplace; + bool overlap; + bool src_overlap; + + bool operator==(const ggml_webgpu_binary_pipeline_key & other) const { + return type == other.type && op == other.op && inplace == other.inplace && overlap == other.overlap && + src_overlap == other.src_overlap; + } +}; + +struct ggml_webgpu_binary_pipeline_key_hash { + size_t operator()(const ggml_webgpu_binary_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.type); + ggml_webgpu_hash_combine(seed, key.op); + ggml_webgpu_hash_combine(seed, key.inplace); + ggml_webgpu_hash_combine(seed, key.overlap); + ggml_webgpu_hash_combine(seed, key.src_overlap); + return seed; + } +}; + +/** Unary **/ + +struct ggml_webgpu_unary_pipeline_key { + int type; + int op; + bool is_unary; // many unary operators fall under the GGML_OP_UNARY umbrella + bool inplace; + ggml_tri_type ttype; // only used for GGML_OP_TRI + + bool operator==(const ggml_webgpu_unary_pipeline_key & other) const { + return type == other.type && op == other.op && is_unary == other.is_unary && inplace == other.inplace && + ttype == other.ttype; + } +}; + +struct ggml_webgpu_unary_pipeline_key_hash { + size_t operator()(const ggml_webgpu_unary_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.type); + ggml_webgpu_hash_combine(seed, key.op); + ggml_webgpu_hash_combine(seed, key.is_unary); + ggml_webgpu_hash_combine(seed, key.inplace); + ggml_webgpu_hash_combine(seed, key.ttype); + return seed; + } +}; + +/** FlashAttention */ + +enum ggml_webgpu_flash_attn_path : uint32_t { + GGML_WEBGPU_FLASH_ATTN_PATH_NONE = 0u, + GGML_WEBGPU_FLASH_ATTN_PATH_SUBGROUP_MATRIX = 1u, + GGML_WEBGPU_FLASH_ATTN_PATH_TILE = 2u, + GGML_WEBGPU_FLASH_ATTN_PATH_VEC = 3u, +}; + +struct ggml_webgpu_flash_attn_pipeline_key { + ggml_type kv_type; + uint32_t head_dim_qk; + uint32_t head_dim_v; + bool kv_direct; + bool kv_overlap; + bool has_mask; + bool has_sinks; + bool uses_logit_softcap; + uint32_t path; + + bool operator==(const ggml_webgpu_flash_attn_pipeline_key & other) const { + return kv_type == other.kv_type && head_dim_qk == other.head_dim_qk && head_dim_v == other.head_dim_v && + kv_direct == other.kv_direct && kv_overlap == other.kv_overlap && has_mask == other.has_mask && + has_sinks == other.has_sinks && uses_logit_softcap == other.uses_logit_softcap && path == other.path; + } +}; + +struct ggml_webgpu_flash_attn_pipeline_key_hash { + size_t operator()(const ggml_webgpu_flash_attn_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.kv_type); + ggml_webgpu_hash_combine(seed, key.head_dim_qk); + ggml_webgpu_hash_combine(seed, key.head_dim_v); + ggml_webgpu_hash_combine(seed, key.kv_direct); + ggml_webgpu_hash_combine(seed, key.kv_overlap); + ggml_webgpu_hash_combine(seed, key.has_mask); + ggml_webgpu_hash_combine(seed, key.has_sinks); + ggml_webgpu_hash_combine(seed, key.uses_logit_softcap); + ggml_webgpu_hash_combine(seed, key.path); + return seed; + } +}; + +struct ggml_webgpu_flash_attn_decisions { + uint32_t path = GGML_WEBGPU_FLASH_ATTN_PATH_NONE; + uint32_t q_tile = 0; + uint32_t kv_tile = 0; + uint32_t wg_size = 0; + bool kv_direct = false; + bool kv_overlap = false; +}; + +inline constexpr uint32_t GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH = 4u; +inline constexpr uint32_t GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE = 4u; + +inline uint32_t ggml_webgpu_flash_attn_pick_vec_ne(const ggml_webgpu_flash_attn_pipeline_key & key) { + if (key.path != GGML_WEBGPU_FLASH_ATTN_PATH_VEC || key.kv_type != GGML_TYPE_F16 || + key.head_dim_qk != key.head_dim_v) { + return 1u; + } + + switch (key.head_dim_qk) { + case 64: + case 192: + case 576: + return 2u; + case 96: + return 4u; + default: + return 1u; + } +} + +inline ggml_webgpu_flash_attn_pipeline_key ggml_webgpu_flash_attn_make_pipeline_key( + const ggml_webgpu_shader_lib_context & context, + uint32_t path) { + const bool has_mask = context.src3 != nullptr; + const bool has_sinks = context.src4 != nullptr; + bool kv_direct = false; + if (path != GGML_WEBGPU_FLASH_ATTN_PATH_TILE) { + uint32_t kv_direct_align = GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH; + if (path == GGML_WEBGPU_FLASH_ATTN_PATH_SUBGROUP_MATRIX) { + kv_direct_align = context.sg_mat_k; + } + kv_direct = (context.src1->type == GGML_TYPE_F16) && + (context.src0->ne[0] % std::max(1u, kv_direct_align) == 0) && + (context.src1->ne[1] % GGML_WEBGPU_KV_SEQ_PAD == 0); + } + + ggml_webgpu_flash_attn_pipeline_key key = {}; + key.kv_type = context.src1->type; + key.head_dim_qk = (uint32_t) context.src0->ne[0]; + key.head_dim_v = (uint32_t) context.src2->ne[0]; + key.kv_direct = kv_direct; + key.kv_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src2); + key.has_mask = has_mask; + key.has_sinks = has_sinks; + key.uses_logit_softcap = ggml_get_op_params_f32(context.dst, 2) != 0.0f; + key.path = path; + return key; +} + +struct ggml_webgpu_flash_attn_vec_reduce_pipeline_key { + uint32_t head_dim_v; + uint32_t wg_size; +}; + +struct ggml_webgpu_flash_attn_vec_reduce_pipeline_key_hash { + size_t operator()(const ggml_webgpu_flash_attn_vec_reduce_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.head_dim_v); + ggml_webgpu_hash_combine(seed, key.wg_size); + return seed; + } +}; + +inline bool operator==(const ggml_webgpu_flash_attn_vec_reduce_pipeline_key & lhs, + const ggml_webgpu_flash_attn_vec_reduce_pipeline_key & rhs) { + return lhs.head_dim_v == rhs.head_dim_v && lhs.wg_size == rhs.wg_size; +} + +struct ggml_webgpu_flash_attn_blk_pipeline_key { + uint32_t kv_tile; + + bool operator==(const ggml_webgpu_flash_attn_blk_pipeline_key & other) const { return kv_tile == other.kv_tile; } +}; + +struct ggml_webgpu_flash_attn_blk_pipeline_key_hash { + size_t operator()(const ggml_webgpu_flash_attn_blk_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.kv_tile); + return seed; + } +}; + +// This is exposed because it's necessary in supports_op +inline size_t ggml_webgpu_flash_attn_wg_mem_bytes(uint32_t q_tile, + uint32_t kv_tile, + uint32_t head_dim_qk, + uint32_t head_dim_v, + bool has_mask, + bool kv_direct) { + const uint32_t max_head_dim = std::max(head_dim_qk, head_dim_v); + size_t f16_elems = 0; + size_t f32_elems = 0; + f16_elems += q_tile * head_dim_qk; // q_shmem + if (!kv_direct) { + f16_elems += kv_tile * max_head_dim; // kv_shmem + } + f16_elems += q_tile * head_dim_v; // o_shmem + if (has_mask) { + f16_elems += q_tile * kv_tile; // mask_shmem + } + f16_elems += q_tile * kv_tile; // inter_shmem + f32_elems += q_tile; // row_max_shmem + f32_elems += q_tile; // exp_sum_shmem + return f16_elems * GGML_WEBGPU_F16_SIZE_BYTES + f32_elems * GGML_WEBGPU_F32_SIZE_BYTES; +} + +inline uint32_t ggml_webgpu_flash_attn_max_kv_tile(const ggml_webgpu_shader_lib_context & context, + const ggml_webgpu_flash_attn_pipeline_key & key) { + const size_t limit_bytes = context.wg_mem_limit_bytes; + uint32_t q_tile = context.sg_mat_m; + uint32_t kv_granularity = context.sg_mat_n; + if (key.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE) { + q_tile = GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE; + kv_granularity = std::max(1u, context.max_subgroup_size); + } else if (key.path == GGML_WEBGPU_FLASH_ATTN_PATH_VEC) { + q_tile = 1u; + kv_granularity = 8u; + } + const size_t base_q_bytes = (key.head_dim_qk + key.head_dim_v) * q_tile * GGML_WEBGPU_F16_SIZE_BYTES + + 2 * q_tile * GGML_WEBGPU_F32_SIZE_BYTES; + size_t bytes_per_kv = 0; + if (!key.kv_direct) { + bytes_per_kv += std::max(key.head_dim_qk, key.head_dim_v); + } + if (key.has_mask) { + bytes_per_kv += q_tile; + } + bytes_per_kv += q_tile; + bytes_per_kv *= GGML_WEBGPU_F16_SIZE_BYTES; + const uint32_t max_kv_tile = (limit_bytes - base_q_bytes) / bytes_per_kv; + return (max_kv_tile / kv_granularity) * kv_granularity; +} + +inline ggml_webgpu_flash_attn_decisions ggml_webgpu_flash_attn_get_decisions( + const ggml_webgpu_shader_lib_context & context, + size_t storage_offset_alignment) { + ggml_webgpu_flash_attn_decisions decisions = {}; + const size_t alignment = std::max(1u, storage_offset_alignment); + const auto * K = context.src1; + const auto * V = context.src2; + GGML_ASSERT(K != nullptr); + GGML_ASSERT(V != nullptr); + + const auto flash_attn_tensor_offset = [](const ggml_tensor * tensor) -> size_t { + constexpr uintptr_t ptr_base_addr = 0x1000u; + const ggml_tensor * base = tensor->view_src != nullptr ? tensor->view_src : tensor; + return reinterpret_cast(base->data) - ptr_base_addr + tensor->view_offs; + }; + + const uint32_t k_offset_elems = + (uint32_t) ((flash_attn_tensor_offset(K) & (alignment - 1)) / ggml_type_size(K->type)); + const uint32_t v_offset_elems = + (uint32_t) ((flash_attn_tensor_offset(V) & (alignment - 1)) / ggml_type_size(V->type)); + const bool f16_vec4_aligned = (k_offset_elems % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0u) && + (v_offset_elems % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0u); + const bool kv_vec_type_supported = + K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q4_0 || K->type == GGML_TYPE_Q8_0; + const bool use_vec = context.supports_subgroups && (context.src0->ne[1] < 20) && (context.src0->ne[0] % 32 == 0) && + (context.src2->ne[0] % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0) && + kv_vec_type_supported && (K->type != GGML_TYPE_F16 || f16_vec4_aligned) && + (context.src2->type == K->type); + const bool use_tile = context.supports_subgroups && !context.supports_subgroup_matrix && K->type == GGML_TYPE_F16 && + V->type == GGML_TYPE_F16 && f16_vec4_aligned && + (context.src0->ne[0] % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0) && + (context.src2->ne[0] % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0) && !use_vec; + + decisions.path = use_vec ? GGML_WEBGPU_FLASH_ATTN_PATH_VEC : + use_tile ? GGML_WEBGPU_FLASH_ATTN_PATH_TILE : + context.supports_subgroup_matrix ? GGML_WEBGPU_FLASH_ATTN_PATH_SUBGROUP_MATRIX : + GGML_WEBGPU_FLASH_ATTN_PATH_NONE; + + if (decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_NONE) { + return decisions; + } + + const ggml_webgpu_flash_attn_pipeline_key key = ggml_webgpu_flash_attn_make_pipeline_key(context, decisions.path); + decisions.kv_direct = key.kv_direct; + const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile(context, key); + // invalidate if even the smallest kv_tile doesn't fit in shared memory + if (max_kv_tile == 0) { + decisions.path = GGML_WEBGPU_FLASH_ATTN_PATH_NONE; + return decisions; + } + + if (decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_VEC) { + decisions.q_tile = 1u; + decisions.kv_tile = std::max(8u, std::min(32u, max_kv_tile)); + decisions.kv_tile = (decisions.kv_tile / 8u) * 8u; + decisions.wg_size = std::max(1u, std::min(32u, context.max_subgroup_size)); + if (decisions.kv_direct) { + decisions.kv_tile = std::min(decisions.kv_tile, GGML_WEBGPU_KV_SEQ_PAD); + while (GGML_WEBGPU_KV_SEQ_PAD % decisions.kv_tile != 0) { + decisions.kv_tile -= 8u; + } + } + return decisions; + } + + decisions.q_tile = + decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE ? GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE : context.sg_mat_m; + decisions.kv_tile = decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE ? + std::min(64u, max_kv_tile) : + std::min(max_kv_tile, context.sg_mat_n * GGML_WEBGPU_FLASH_ATTN_PREFERRED_KV_SG_TILES); + decisions.wg_size = decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE ? + GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE : + std::max(context.max_subgroup_size, GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE); + + if (decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE) { + const uint32_t tile_kv_granularity = std::max(1u, context.max_subgroup_size); + decisions.kv_tile = + std::max(tile_kv_granularity, (decisions.kv_tile / tile_kv_granularity) * tile_kv_granularity); + } + + if (decisions.kv_direct) { + GGML_ASSERT(decisions.kv_tile <= GGML_WEBGPU_KV_SEQ_PAD); + while (GGML_WEBGPU_KV_SEQ_PAD % decisions.kv_tile != 0) { + decisions.kv_tile -= decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE ? + std::max(1u, context.max_subgroup_size) : + context.sg_mat_n; + } + } + return decisions; +} + +/** Matrix Multiplication **/ + +struct ggml_webgpu_legacy_mul_mat_pipeline_key { + ggml_type src0_type; + ggml_type src1_type; + + bool operator==(const ggml_webgpu_legacy_mul_mat_pipeline_key & other) const { + return src0_type == other.src0_type && src1_type == other.src1_type; + } +}; + +struct ggml_webgpu_legacy_mul_mat_pipeline_key_hash { + size_t operator()(const ggml_webgpu_legacy_mul_mat_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.src0_type); + ggml_webgpu_hash_combine(seed, key.src1_type); + return seed; + } +}; + +struct ggml_webgpu_mul_mat_vec_pipeline_key { + ggml_type src0_type; + ggml_type src1_type; + int vectorized; + + bool operator==(const ggml_webgpu_mul_mat_vec_pipeline_key & other) const { + return src0_type == other.src0_type && src1_type == other.src1_type && vectorized == other.vectorized; + } +}; + +struct ggml_webgpu_mul_mat_vec_pipeline_key_hash { + size_t operator()(const ggml_webgpu_mul_mat_vec_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.src0_type); + ggml_webgpu_hash_combine(seed, key.src1_type); + ggml_webgpu_hash_combine(seed, key.vectorized); + return seed; + } +}; + +struct ggml_webgpu_mul_mat_vec_shader_decisions { + uint32_t wg_size; + uint32_t outputs_per_wg; + uint32_t vec_size; +}; + +struct ggml_webgpu_mul_mat_pipeline_key { + ggml_type src0_type; + ggml_type src1_type; + int vectorized; + int use_subgroup_matrix; + + bool operator==(const ggml_webgpu_mul_mat_pipeline_key & other) const { + return src0_type == other.src0_type && src1_type == other.src1_type && vectorized == other.vectorized && + use_subgroup_matrix == other.use_subgroup_matrix; + } +}; + +struct ggml_webgpu_mul_mat_pipeline_key_hash { + size_t operator()(const ggml_webgpu_mul_mat_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.src0_type); + ggml_webgpu_hash_combine(seed, key.src1_type); + ggml_webgpu_hash_combine(seed, key.vectorized); + ggml_webgpu_hash_combine(seed, key.use_subgroup_matrix); + return seed; + } +}; + +struct ggml_webgpu_mul_mat_shader_decisions { + uint32_t tile_k; + uint32_t wg_size_m; + uint32_t wg_size_n; + uint32_t wg_size; + uint32_t outputs_per_wg; + int use_subgroup_matrix; + + uint32_t tile_m; + uint32_t tile_n; + + // Subgroup matrix parameters + uint32_t subgroup_m; + uint32_t subgroup_n; + uint32_t subgroup_matrix_m; + uint32_t subgroup_matrix_n; + + uint32_t mul_mat_wg_size; +}; + +/** MUL_MAT_ID **/ + +struct ggml_webgpu_mul_mat_id_pipeline_key { + ggml_type src0_type; + ggml_type src1_type; + uint32_t n_experts; + int vectorized; + + bool operator==(const ggml_webgpu_mul_mat_id_pipeline_key & other) const { + return src0_type == other.src0_type && src1_type == other.src1_type && n_experts == other.n_experts && + vectorized == other.vectorized; + } +}; + +struct ggml_webgpu_mul_mat_id_pipeline_key_hash { + size_t operator()(const ggml_webgpu_mul_mat_id_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.src0_type); + ggml_webgpu_hash_combine(seed, key.src1_type); + ggml_webgpu_hash_combine(seed, key.n_experts); + ggml_webgpu_hash_combine(seed, key.vectorized); + return seed; + } +}; + +/** Cpy **/ + +struct ggml_webgpu_cpy_pipeline_key { + ggml_type src_type; + ggml_type dst_type; + + bool operator==(const ggml_webgpu_cpy_pipeline_key & other) const { + return src_type == other.src_type && dst_type == other.dst_type; + } +}; + +struct ggml_webgpu_cpy_pipeline_key_hash { + size_t operator()(const ggml_webgpu_cpy_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.src_type); + ggml_webgpu_hash_combine(seed, key.dst_type); + return seed; + } +}; + +/** Glu **/ + +struct ggml_webgpu_glu_pipeline_key { + ggml_glu_op glu_op; + ggml_type type; + bool split; + + bool operator==(const ggml_webgpu_glu_pipeline_key & other) const { + return glu_op == other.glu_op && type == other.type && split == other.split; + } +}; + +struct ggml_webgpu_glu_pipeline_key_hash { + size_t operator()(const ggml_webgpu_glu_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.glu_op); + ggml_webgpu_hash_combine(seed, key.type); + ggml_webgpu_hash_combine(seed, key.split); + return seed; + } +}; + +/** Rope **/ + +struct ggml_webgpu_rope_pipeline_key { + ggml_type type; + bool inplace; + bool has_ff; + + bool operator==(const ggml_webgpu_rope_pipeline_key & other) const { + return type == other.type && inplace == other.inplace && has_ff == other.has_ff; + } +}; + +struct ggml_webgpu_rope_pipeline_key_hash { + size_t operator()(const ggml_webgpu_rope_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.type); + ggml_webgpu_hash_combine(seed, key.inplace); + ggml_webgpu_hash_combine(seed, key.has_ff); + return seed; + } +}; + +/** SoftMax **/ + +struct ggml_webgpu_soft_max_pipeline_key { + ggml_type mask_type; + bool has_mask; + bool has_sink; + bool inplace; + + bool operator==(const ggml_webgpu_soft_max_pipeline_key & other) const { + return mask_type == other.mask_type && has_mask == other.has_mask && has_sink == other.has_sink && + inplace == other.inplace; + } +}; + +struct ggml_webgpu_soft_max_pipeline_key_hash { + size_t operator()(const ggml_webgpu_soft_max_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.mask_type); + ggml_webgpu_hash_combine(seed, key.has_mask); + ggml_webgpu_hash_combine(seed, key.has_sink); + ggml_webgpu_hash_combine(seed, key.inplace); + return seed; + } +}; + +class ggml_webgpu_shader_lib { + wgpu::Device device; + pre_wgsl::Preprocessor preprocessor; + + std::unordered_map sum_rows_pipelines; // key is fixed, no variants yet + std::unordered_map argmax_pipelines; // key is vec4 + std::unordered_map argsort_pipelines; // key is order + std::unordered_map argsort_merge_pipelines; // key is order + std::unordered_map cumsum_pipelines; // key is fixed, no variants yet + std::unordered_map + row_norm_pipelines; // op/inplace + + std::unordered_map + get_rows_pipelines; // src_type, vectorized + std::unordered_map + unary_pipelines; // type/op/inplace + std::unordered_map + scale_pipelines; // inplace + std::unordered_map + solve_tri_pipelines; // type + std::unordered_map + ssm_conv_pipelines; // type/vectorized + std::unordered_map + ssm_scan_pipelines; // type/d_state + std::unordered_map + gated_delta_net_pipelines; // type/S_v/kda + std::unordered_map + pad_pipelines; // circular/non-circular + std::unordered_map + binary_pipelines; // type/op/inplace/overlap + std::unordered_map + concat_pipelines; // type + std::unordered_map + repeat_pipelines; // type + std::unordered_map + flash_attn_pipelines; + std::unordered_map + flash_attn_vec_reduce_pipelines; + std::unordered_map + flash_attn_blk_pipelines; + std::unordered_map + mul_mat_legacy_pipelines; // legacy mul_mat (non-subgroup/non-regtile/non-vec) + std::unordered_map + mul_mat_vec_pipelines; // fast mat-vec (n==1) + std::unordered_map + mul_mat_fast_pipelines; // fast mat-mat (reg-tile or subgroup) + std::unordered_map mul_mat_id_gather_pipelines; // key is fixed + std::unordered_map + mul_mat_id_pipelines; // src0_type/src1_type + std::unordered_map + mul_mat_id_vec_pipelines; // src0_type/src1_type + + std::unordered_map + set_rows_pipelines; + std::unordered_map set_pipelines; + std::unordered_map cpy_pipelines; + std::unordered_map glu_pipelines; + std::unordered_map + rope_pipelines; + std::unordered_map + soft_max_pipelines; + std::unordered_map + conv2d_pipelines; + std::unordered_map + im2col_pipelines; + + std::unordered_map + rms_norm_mul_pipelines; + std::unordered_map + upscale_pipelines; + + public: + ggml_webgpu_shader_lib(wgpu::Device device) { this->device = device; } + + webgpu_pipeline get_sum_rows_pipeline(const ggml_webgpu_shader_lib_context & context) { + auto it = sum_rows_pipelines.find(1); + if (it != sum_rows_pipelines.end()) { + return it->second; + } + std::vector defines; + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_sum_rows, defines); + sum_rows_pipelines[1] = ggml_webgpu_create_pipeline(device, processed, "sum_rows"); + return sum_rows_pipelines[1]; + } + + webgpu_pipeline get_row_norm_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_row_norm_pipeline_key key = {}; + key.op = context.dst->op; + key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); + + auto it = row_norm_pipelines.find(key); + if (it != row_norm_pipelines.end()) { + return it->second; + } + std::vector defines; + std::string variant; + + switch (key.op) { + case GGML_OP_RMS_NORM: + defines.push_back("RMS_NORM"); + variant = "rms_norm"; + break; + case GGML_OP_L2_NORM: + defines.push_back("L2_NORM"); + variant = "l2_norm"; + break; + default: + GGML_ABORT("Unsupported op for row_norm shader"); + } + + if (key.inplace) { + defines.push_back("INPLACE"); + variant += "_inplace"; + } + + const uint32_t row_norm_wg_size = 128u; + uint32_t wg_size = std::min(context.max_wg_size, row_norm_wg_size); + defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); + auto processed = preprocessor.preprocess(wgsl_row_norm, defines); + auto decisions = std::make_shared(); + decisions->wg_size = wg_size; + decisions->inplace = key.inplace; + row_norm_pipelines[key] = ggml_webgpu_create_pipeline(device, processed, variant); + row_norm_pipelines[key].context = decisions; + return row_norm_pipelines[key]; + } + + webgpu_pipeline get_argmax_pipeline(const ggml_webgpu_shader_lib_context & context) { + bool vec4 = context.src0->ne[0] % 4 == 0; + + auto it = argmax_pipelines.find(vec4); + if (it != argmax_pipelines.end()) { + return it->second; + } + std::string variant = "argmax"; + std::vector defines; + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + if (vec4) { + defines.push_back("VEC4"); + variant += "_vec4"; + } + + auto processed = preprocessor.preprocess(wgsl_argmax, defines); + argmax_pipelines[vec4] = ggml_webgpu_create_pipeline(device, processed, variant); + return argmax_pipelines.at(vec4); + } + + webgpu_pipeline get_set_rows_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_set_rows_pipeline_key key = {}; + key.dst_type = context.dst->type; + key.vec4 = context.src0->ne[0] % 4 == 0; + key.i64_idx = context.src1->type == GGML_TYPE_I64; + + auto it = set_rows_pipelines.find(key); + if (it != set_rows_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "set_rows"; + + switch (context.dst->type) { + case GGML_TYPE_F32: + defines.push_back("DST_F32"); + variant += "_dstf32"; + break; + case GGML_TYPE_F16: + defines.push_back("DST_F16"); + variant += "_dstf16"; + break; + default: + GGML_ABORT("Unsupported dst type for set_rows shader"); + } + + if (key.vec4) { + defines.push_back("VEC4"); + variant += "_vec4"; + } + if (key.i64_idx) { + defines.push_back("I64_IDX"); + variant += "_i64idx"; + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_set_rows, defines); + auto decisions = std::make_shared(); + decisions->vec4 = key.vec4; + decisions->i64_idx = key.i64_idx; + decisions->wg_size = context.max_wg_size; + set_rows_pipelines[key] = ggml_webgpu_create_pipeline(device, processed, variant); + set_rows_pipelines[key].context = decisions; + return set_rows_pipelines[key]; + } + + webgpu_pipeline get_set_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_set_pipeline_key key = {}; + key.type = context.dst->type; + key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); + + auto it = set_pipelines.find(key); + if (it != set_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "set"; + + switch (key.type) { + case GGML_TYPE_F32: + defines.push_back("TYPE_F32"); + variant += "_f32"; + break; + case GGML_TYPE_I32: + defines.push_back("TYPE_I32"); + variant += "_i32"; + break; + default: + GGML_ABORT("Unsupported type for set shader"); + } + + if (key.inplace) { + defines.push_back("INPLACE"); + variant += "_inplace"; + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_set, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + decisions->inplace = key.inplace; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + set_pipelines[key] = pipeline; + return set_pipelines[key]; + } + + webgpu_pipeline get_cumsum_pipeline(const ggml_webgpu_shader_lib_context & context) { + auto it = cumsum_pipelines.find(1); + if (it != cumsum_pipelines.end()) { + return it->second; + } + + std::vector defines; + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_cumsum, defines); + cumsum_pipelines[1] = ggml_webgpu_create_pipeline(device, processed, "cumsum"); + return cumsum_pipelines[1]; + } + + webgpu_pipeline get_argsort_pipeline(const ggml_webgpu_shader_lib_context & context) { + bool is_top_k = context.dst->op == GGML_OP_TOP_K; + // ascending order is 0, descending order is 1 + const int32_t order = + is_top_k ? (int32_t) GGML_SORT_ORDER_DESC : (int32_t) ggml_get_op_params_i32(context.dst, 0); + + auto it = argsort_pipelines.find(order); + if (it != argsort_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "argsort"; + defines.push_back(std::string("ORDER=") + std::to_string(order)); + variant += std::string("_order") + std::to_string(order); + uint32_t wg_size = 1; + while (wg_size * 2 <= context.max_wg_size && + wg_size * GGML_WEBGPU_I32_SIZE_BYTES <= context.wg_mem_limit_bytes / 2) { + wg_size *= 2; + } + defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); + auto processed = preprocessor.preprocess(wgsl_argsort, defines); + auto decisions = std::make_shared(); + decisions->wg_size = wg_size; + argsort_pipelines[order] = ggml_webgpu_create_pipeline(device, processed, variant); + argsort_pipelines[order].context = decisions; + return argsort_pipelines[order]; + } + + webgpu_pipeline get_argsort_merge_pipeline(const ggml_webgpu_shader_lib_context & context) { + bool is_top_k = context.dst->op == GGML_OP_TOP_K; + // ascending order is 0, descending order is 1 + const int32_t order = + is_top_k ? (int32_t) GGML_SORT_ORDER_DESC : (int32_t) ggml_get_op_params_i32(context.dst, 0); + + auto it = argsort_merge_pipelines.find(order); + if (it != argsort_merge_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "argsort_merge"; + defines.push_back(std::string("ORDER=") + std::to_string(order)); + variant += std::string("_order") + std::to_string(order); + uint32_t wg_size = std::min(GGML_WEBGPU_ARGSORT_MERGE_MAX_WG_SIZE, context.max_wg_size); + defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); + + auto processed = preprocessor.preprocess(wgsl_argsort_merge, defines); + argsort_merge_pipelines[order] = ggml_webgpu_create_pipeline(device, processed, variant); + return argsort_merge_pipelines[order]; + } + + webgpu_pipeline get_get_rows_pipeline(const ggml_webgpu_shader_lib_context & context) { + const bool vectorized = context.src0->type == GGML_TYPE_F32 && context.dst->ne[0] % 4 == 0; + ggml_webgpu_get_rows_pipeline_key key = {}; + key.src_type = context.src0->type; + key.vectorized = (int) vectorized; + + auto it = get_rows_pipelines.find(key); + if (it != get_rows_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "get_rows"; + + const struct ggml_type_traits * type_traits = ggml_get_type_traits(key.src_type); + const char * type_str = type_traits->type_name; + + switch (key.src_type) { + case GGML_TYPE_F32: + defines.push_back("FLOAT_PARALLEL"); + if (key.vectorized) { + defines.push_back("F32_VEC"); + defines.push_back("SRC_TYPE=vec4"); + defines.push_back("DST_TYPE=vec4"); + defines.push_back("BLOCK_SIZE=4u"); + } else { + defines.push_back("F32"); + defines.push_back("SRC_TYPE=f32"); + defines.push_back("DST_TYPE=f32"); + defines.push_back("BLOCK_SIZE=1u"); + } + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("FLOAT_PARALLEL"); + defines.push_back("F16"); + defines.push_back("SRC_TYPE=f16"); + defines.push_back("DST_TYPE=f32"); + defines.push_back("BLOCK_SIZE=1u"); + variant += "_f16"; + break; + case GGML_TYPE_I32: + defines.push_back("FLOAT_PARALLEL"); + defines.push_back("I32"); + defines.push_back("SRC_TYPE=i32"); + defines.push_back("DST_TYPE=i32"); + defines.push_back("BLOCK_SIZE=1u"); + variant += "_i32"; + break; + default: + { + std::string type_upper = type_str; + std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper); + + switch (key.src_type) { + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ4_NL: + { + // Quantized types using u32 buffers for portability. + defines.push_back("SRC_TYPE=u32"); + defines.push_back("U32_DEQUANT_HELPERS"); + break; + } + default: + { + defines.push_back(std::string("SRC_TYPE=") + type_str); + } + } + + defines.push_back("BYTE_HELPERS"); + defines.push_back(type_upper + "_T"); + defines.push_back(type_upper); + defines.push_back(type_upper + "_SCALE_MIN"); + defines.push_back(type_upper + "_TABLES"); + defines.push_back(type_upper + "_GRID"); + + variant += "_"; + variant += type_str; + + defines.push_back("DST_TYPE=f32"); + + if (key.src_type == GGML_TYPE_Q1_0) { + defines.push_back("BLOCK_SIZE=128u"); + } else if ((key.src_type >= GGML_TYPE_Q4_0 && key.src_type <= GGML_TYPE_Q8_1) || + key.src_type == GGML_TYPE_IQ4_NL) { + defines.push_back("BLOCK_SIZE=32u"); + } else if (key.src_type >= GGML_TYPE_Q2_K) { + defines.push_back("BLOCK_SIZE=256u"); + } else { + defines.push_back("BLOCK_SIZE=1u"); + } + break; + } + } + + if (key.vectorized) { + variant += "_vec"; + } + + defines.push_back("WG_SIZE=" + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_get_rows, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + get_rows_pipelines[key] = pipeline; + return get_rows_pipelines[key]; + } + + webgpu_pipeline get_scale_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_scale_pipeline_key key = {}; + key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); + + auto it = scale_pipelines.find(key); + if (it != scale_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "scale"; + + if (key.inplace) { + defines.push_back("INPLACE"); + variant += "_inplace"; + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_scale, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + decisions->inplace = key.inplace; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + scale_pipelines[key] = pipeline; + return scale_pipelines[key]; + } + + webgpu_pipeline get_solve_tri_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_solve_tri_pipeline_key key = {}; + key.type = context.dst->type; + key.n = (int) context.src0->ne[0]; + key.k = (int) context.src1->ne[0]; + + auto it = solve_tri_pipelines.find(key); + if (it != solve_tri_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "solve_tri"; + + switch (key.type) { + case GGML_TYPE_F32: + variant += "_f32"; + break; + default: + GGML_ABORT("Unsupported type for solve_tri shader"); + } + + const uint32_t wg_size = std::min((uint32_t) key.n, context.max_wg_size); + const uint32_t k_tile = wg_size; + const uint32_t bytes_per_row = ((uint32_t) key.n + wg_size) * GGML_WEBGPU_F32_SIZE_BYTES; + const uint32_t batch_n = (uint32_t) (context.wg_mem_limit_bytes / bytes_per_row); + + defines.push_back(std::string("N=") + std::to_string(key.n)); + defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); + defines.push_back(std::string("K_TILE=") + std::to_string(k_tile)); + defines.push_back(std::string("BATCH_N=") + std::to_string(batch_n)); + + auto processed = preprocessor.preprocess(wgsl_solve_tri, defines); + auto decisions = std::make_shared(); + decisions->wg_size = wg_size; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + solve_tri_pipelines[key] = pipeline; + return solve_tri_pipelines[key]; + } + + webgpu_pipeline get_ssm_conv_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_ssm_conv_pipeline_key key = {}; + key.type = context.dst->type; + key.vectorized = context.src1->ne[0] == 4; + + auto it = ssm_conv_pipelines.find(key); + if (it != ssm_conv_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "ssm_conv"; + + switch (key.type) { + case GGML_TYPE_F32: + variant += "_f32"; + break; + default: + GGML_ABORT("Unsupported type for ssm_conv shader"); + } + + if (key.vectorized) { + defines.push_back("VECTORIZED"); + variant += "_vec4"; + } + + constexpr uint32_t block_size = 32u; + constexpr uint32_t tokens_per_wg = 8u; + + defines.push_back("BLOCK_SIZE=" + std::to_string(block_size) + "u"); + defines.push_back("TOKENS_PER_WG=" + std::to_string(tokens_per_wg) + "u"); + + auto processed = preprocessor.preprocess(wgsl_ssm_conv, defines); + auto decisions = std::make_shared(); + decisions->block_size = block_size; + decisions->tokens_per_wg = tokens_per_wg; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + ssm_conv_pipelines[key] = pipeline; + return ssm_conv_pipelines[key]; + } + + webgpu_pipeline get_ssm_scan_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_ssm_scan_pipeline_key key = {}; + key.type = context.dst->type; + key.d_state = (int) context.src0->ne[0]; + key.xbc_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src4) && + ggml_webgpu_tensor_overlap(context.src1, context.src5); + + auto it = ssm_scan_pipelines.find(key); + if (it != ssm_scan_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "ssm_scan"; + + switch (key.type) { + case GGML_TYPE_F32: + variant += "_f32"; + break; + default: + GGML_ABORT("Unsupported type for ssm_scan shader"); + } + + const uint32_t wg_size = (uint32_t) key.d_state; + + constexpr uint32_t tokens_per_tile = 4u; + + defines.push_back("WG_SIZE=" + std::to_string(wg_size) + "u"); + defines.push_back("TOKENS_PER_TILE=" + std::to_string(tokens_per_tile) + "u"); + + if (context.supports_subgroups) { + defines.push_back("USE_SUBGROUP_REDUCTION"); + variant += "_sg_reduce"; + } else { + variant += "_wg_reduce"; + } + + if (key.xbc_overlap) { + defines.push_back("XBC_OVERLAP"); + } + + variant += "_d" + std::to_string(key.d_state); + + auto processed = preprocessor.preprocess(wgsl_ssm_scan, defines); + auto decisions = std::make_shared(); + decisions->wg_size = wg_size; + decisions->tokens_per_tile = tokens_per_tile; + decisions->xbc_overlap = key.xbc_overlap; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + ssm_scan_pipelines[key] = pipeline; + return ssm_scan_pipelines[key]; + } + + webgpu_pipeline get_gated_delta_net_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_gated_delta_net_pipeline_key key = {}; + key.type = context.dst->type; + key.s_v = (int) context.src2->ne[0]; + key.kda = context.src3->ne[0] == context.src2->ne[0]; + + auto it = gated_delta_net_pipelines.find(key); + if (it != gated_delta_net_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "gated_delta_net"; + + switch (key.type) { + case GGML_TYPE_F32: + variant += "_f32"; + break; + default: + GGML_ABORT("Unsupported type for gated_delta_net shader"); + } + + if (key.kda) { + defines.push_back("KDA"); + variant += "_kda"; + } + + defines.push_back("S_V=" + std::to_string(key.s_v) + "u"); + defines.push_back("WG_SIZE=" + std::to_string(key.s_v) + "u"); + + auto processed = preprocessor.preprocess(wgsl_gated_delta_net, defines); + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + gated_delta_net_pipelines[key] = pipeline; + return gated_delta_net_pipelines[key]; + } + + webgpu_pipeline get_pad_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_pad_pipeline_key key = {}; + key.circular = ggml_get_op_params_i32(context.dst, 8) != 0; + + auto it = pad_pipelines.find(key); + if (it != pad_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "pad"; + + if (key.circular) { + defines.push_back("CIRCULAR"); + variant += "_circular"; + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_pad, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + pad_pipelines[key] = pipeline; + return pad_pipelines[key]; + } + + webgpu_pipeline get_mul_mat_vec_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_mul_mat_vec_pipeline_key key = {}; + key.src0_type = context.src0->type; + key.src1_type = context.src1->type; + key.vectorized = (context.src0->ne[0] % 4 == 0 && + (context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ? + 1 : + 0; + + auto it = mul_mat_vec_pipelines.find(key); + if (it != mul_mat_vec_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "mul_mat_vec"; + const char * shader_src = wgsl_mul_mat_vec; + + // src0 type (matrix row) + switch (context.src0->type) { + case GGML_TYPE_F32: + defines.push_back("SRC0_INNER_TYPE=f32"); + defines.push_back("MUL_ACC_FLOAT"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("SRC0_INNER_TYPE=f16"); + defines.push_back("MUL_ACC_FLOAT"); + variant += "_f16"; + break; + default: + { + // Quantized types: use helpers but accumulate in f16 + const struct ggml_type_traits * src0_traits = ggml_get_type_traits(context.src0->type); + std::string src0_name = src0_traits->type_name; + std::string type_upper = src0_name; + variant += "_" + src0_name; + std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper); + + defines.push_back("BYTE_HELPERS"); + defines.push_back("MUL_ACC_" + type_upper); + defines.push_back("U32_DEQUANT_HELPERS"); + defines.push_back("SRC0_INNER_TYPE=u32"); + switch (context.src0->type) { + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + defines.push_back(type_upper + "_GRID"); + break; + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + defines.push_back(type_upper + "_GRID"); + defines.push_back(type_upper + "_TABLES"); + break; + default: + break; + } + break; + } + } + + // src1 type (vector) + switch (context.src1->type) { + case GGML_TYPE_F32: + defines.push_back("SRC1_INNER_TYPE=f32"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("SRC1_INNER_TYPE=f16"); + variant += "_f16"; + break; + default: + GGML_ABORT("Unsupported src1 type for mul_mat_vec shader"); + } + + // VEC/SCALAR controls + defines.push_back(key.vectorized ? "VEC" : "SCALAR"); + + uint32_t wg_size = WEBGPU_MUL_MAT_VEC_WG_SIZE; + uint32_t outputs_per_wg = WEBGPU_MUL_MAT_VEC_FLOAT_OUTPUTS_PER_WG; + + if (key.src0_type == GGML_TYPE_Q1_0) { + outputs_per_wg = WEBGPU_MUL_MAT_VEC_LEGACY_Q_OUTPUTS_PER_WG; + } else if (key.src0_type >= GGML_TYPE_Q2_K) { + outputs_per_wg = WEBGPU_MUL_MAT_VEC_K_Q_OUTPUTS_PER_WG; + } else if (key.src0_type >= GGML_TYPE_Q4_0) { + outputs_per_wg = WEBGPU_MUL_MAT_VEC_LEGACY_Q_OUTPUTS_PER_WG; + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); + defines.push_back(std::string("OUTPUTS_PER_WG=") + std::to_string(outputs_per_wg)); + defines.push_back(context.supports_subgroups ? "USE_SUBGROUP_REDUCTION" : "USE_WORKGROUP_REDUCTION"); + variant += context.supports_subgroups ? "_sg_reduce" : "_wg_reduce"; + if (key.vectorized) { + variant += "_vectorized"; + } + + auto processed = preprocessor.preprocess(shader_src, defines); + auto decisions = std::make_shared(); + decisions->wg_size = wg_size; + decisions->outputs_per_wg = outputs_per_wg; + decisions->vec_size = key.vectorized ? 4 : 1; + + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + mul_mat_vec_pipelines[key] = pipeline; + return mul_mat_vec_pipelines[key]; + } + + webgpu_pipeline get_mul_mat_fast_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_mul_mat_pipeline_key key = {}; + key.src0_type = context.src0->type; + key.src1_type = context.src1->type; + key.vectorized = (context.src0->ne[0] % 4 == 0 && context.dst->ne[0] % 4 == 0 && + (context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ? + 1 : + 0; + key.use_subgroup_matrix = context.supports_subgroup_matrix; + + auto it = mul_mat_fast_pipelines.find(key); + if (it != mul_mat_fast_pipelines.end()) { + return it->second; + } + + const char * shader_src = key.use_subgroup_matrix ? wgsl_mul_mat_subgroup_matrix : wgsl_mul_mat_reg_tile; + std::vector defines; + std::string variant = key.use_subgroup_matrix ? "mul_mat_subgroup_matrix" : "mul_mat_reg_tile"; + + // src1 type + switch (context.src1->type) { + case GGML_TYPE_F32: + defines.push_back("SRC1_INNER_TYPE=f32"); + break; + case GGML_TYPE_F16: + defines.push_back("SRC1_INNER_TYPE=f16"); + break; + default: + GGML_ABORT("Unsupported src1 type for mul_mat fast shader"); + } + + // src0 type + const struct ggml_type_traits * src0_traits = ggml_get_type_traits(context.src0->type); + const char * src0_name = src0_traits->type_name; + + switch (context.src0->type) { + case GGML_TYPE_F32: + defines.push_back("SRC0_INNER_TYPE=f32"); + defines.push_back("FLOAT"); + defines.push_back("MUL_ACC_FLOAT"); + defines.push_back("INIT_SRC0_SHMEM_FLOAT"); + defines.push_back("INIT_SRC1_SHMEM_FLOAT"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("SRC0_INNER_TYPE=f16"); + defines.push_back("FLOAT"); + defines.push_back("MUL_ACC_FLOAT"); + defines.push_back("INIT_SRC0_SHMEM_FLOAT"); + defines.push_back("INIT_SRC1_SHMEM_FLOAT"); + variant += "_f16"; + break; + default: + { + std::string type_upper = src0_name; + std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper); + + defines.push_back("BYTE_HELPERS"); + defines.push_back("MUL_ACC_" + type_upper); + defines.push_back("INIT_SRC0_SHMEM_" + type_upper); + defines.push_back("INIT_SRC1_SHMEM_FLOAT"); + defines.push_back("U32_DEQUANT_HELPERS"); + defines.push_back("SRC0_INNER_TYPE=u32"); + + switch (context.src0->type) { + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + defines.push_back(type_upper + "_GRID"); + break; + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ3_S: + defines.push_back(type_upper + "_GRID"); + defines.push_back(type_upper + "_TABLES"); + break; + default: + break; + } + + variant += std::string("_") + src0_name; + break; + } + } + + // VEC/SCALAR controls + defines.push_back(key.vectorized ? "VEC" : "SCALAR"); + + const bool is_quant = ggml_is_quantized(context.src0->type); + + uint32_t tile_k; + if (key.use_subgroup_matrix) { + tile_k = is_quant ? WEBGPU_MUL_MAT_SUBGROUP_TILE_K_QUANT : WEBGPU_MUL_MAT_SUBGROUP_TILE_K_FLOAT; + } else { + tile_k = is_quant ? WEBGPU_MUL_MAT_REG_TILE_K_QUANT : WEBGPU_MUL_MAT_REG_TILE_K_FLOAT; + } + + // Tiles + defines.push_back("TILE_M=" + std::to_string(WEBGPU_MUL_MAT_TILE_M) + "u"); + defines.push_back("TILE_N=" + std::to_string(WEBGPU_MUL_MAT_TILE_N) + "u"); + + // Subgroup matrix specifics + if (key.use_subgroup_matrix) { + defines.push_back("TILE_K=" + std::to_string(tile_k) + "u"); + defines.push_back("MAX_SUBGROUP_SIZE=" + std::to_string(context.max_subgroup_size) + "u"); + defines.push_back("SUBGROUP_M=" + std::to_string(WEBGPU_MUL_MAT_SUBGROUP_M) + "u"); + defines.push_back("SUBGROUP_N=" + std::to_string(WEBGPU_MUL_MAT_SUBGROUP_N) + "u"); + defines.push_back("SUBGROUP_MATRIX_M=" + std::to_string(WEBGPU_MUL_MAT_SUBGROUP_MATRIX_M) + "u"); + defines.push_back("SUBGROUP_MATRIX_N=" + std::to_string(WEBGPU_MUL_MAT_SUBGROUP_MATRIX_N) + "u"); + defines.push_back("SUBGROUP_MATRIX_M_SIZE=" + std::to_string(context.sg_mat_m) + "u"); + defines.push_back("SUBGROUP_MATRIX_N_SIZE=" + std::to_string(context.sg_mat_n) + "u"); + defines.push_back("SUBGROUP_MATRIX_K_SIZE=" + std::to_string(context.sg_mat_k) + "u"); + } + + // variant suffix for src1 type + variant += std::string("_") + (context.src1->type == GGML_TYPE_F32 ? "f32" : "f16"); + if (key.vectorized) { + variant += "_vectorized"; + } + + if (!key.use_subgroup_matrix) { + defines.push_back("WORKGROUP_SIZE_M=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_M) + "u"); + defines.push_back("WORKGROUP_SIZE_N=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_N) + "u"); + defines.push_back("TILE_K=" + std::to_string(tile_k) + "u"); + } + + auto processed = preprocessor.preprocess(shader_src, defines); + + auto decisions = std::make_shared(); + decisions->tile_k = tile_k; + decisions->tile_m = WEBGPU_MUL_MAT_TILE_M; + decisions->tile_n = WEBGPU_MUL_MAT_TILE_N; + decisions->use_subgroup_matrix = key.use_subgroup_matrix; + if (key.use_subgroup_matrix) { + decisions->subgroup_m = WEBGPU_MUL_MAT_SUBGROUP_M; + decisions->subgroup_n = WEBGPU_MUL_MAT_SUBGROUP_N; + decisions->subgroup_matrix_m = WEBGPU_MUL_MAT_SUBGROUP_MATRIX_M; + decisions->subgroup_matrix_n = WEBGPU_MUL_MAT_SUBGROUP_MATRIX_N; + decisions->wg_size = context.max_subgroup_size; + } else { + decisions->wg_size_m = WEBGPU_MUL_MAT_WG_SIZE_M; + decisions->wg_size_n = WEBGPU_MUL_MAT_WG_SIZE_N; + decisions->wg_size = WEBGPU_MUL_MAT_WG_SIZE_M * WEBGPU_MUL_MAT_WG_SIZE_N; + decisions->mul_mat_wg_size = WEBGPU_MUL_MAT_WG_SIZE; + } + + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + mul_mat_fast_pipelines[key] = pipeline; + return mul_mat_fast_pipelines[key]; + } + + webgpu_pipeline get_mul_mat_legacy_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_legacy_mul_mat_pipeline_key key = {}; + key.src0_type = context.src0->type; + key.src1_type = context.src1->type; + + auto it = mul_mat_legacy_pipelines.find(key); + if (it != mul_mat_legacy_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "mul_mat"; + + switch (context.src1->type) { + case GGML_TYPE_F32: + defines.push_back("SRC1_TYPE=f32"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("SRC1_TYPE=f16"); + variant += "_f16"; + break; + default: + GGML_ABORT("Unsupported src1 type for mul_mat legacy shader"); + } + + const struct ggml_type_traits * src0_traits = ggml_get_type_traits(context.src0->type); + const char * src0_name = src0_traits->type_name; + + switch (context.src0->type) { + case GGML_TYPE_F32: + defines.push_back("SRC0_TYPE=f32"); + defines.push_back("FLOAT"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("SRC0_TYPE=f16"); + defines.push_back("FLOAT"); + variant += "_f16"; + break; + default: + { + std::string type_upper = src0_name; + std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper); + + switch (context.src0->type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ4_NL: + { + // Quantized types using u32 buffers for portability. + defines.push_back("SRC0_TYPE=u32"); + defines.push_back("U32_DEQUANT_HELPERS"); + break; + } + default: + { + defines.push_back(std::string("SRC0_TYPE=") + src0_name); + } + } + + defines.push_back("BYTE_HELPERS"); + defines.push_back(type_upper + "_T"); + defines.push_back(type_upper); + defines.push_back(type_upper + "_SCALE_MIN"); + defines.push_back(type_upper + "_TABLES"); + defines.push_back(type_upper + "_GRID"); + + variant += std::string("_") + src0_name; + break; + } + } + + auto processed = preprocessor.preprocess(wgsl_mul_mat, defines); + + auto decisions = std::make_shared(); + decisions->wg_size = WEBGPU_MUL_MAT_WG_SIZE; + + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + mul_mat_legacy_pipelines[key] = pipeline; + return mul_mat_legacy_pipelines[key]; + } + + webgpu_pipeline get_mul_mat_id_gather_pipeline(const ggml_webgpu_shader_lib_context & context) { + auto it = mul_mat_id_gather_pipelines.find(1); + if (it != mul_mat_id_gather_pipelines.end()) { + return it->second; + } + std::vector defines; + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_mul_mat_id_gather, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, "mul_mat_id_gather"); + pipeline.context = decisions; + mul_mat_id_gather_pipelines[1] = pipeline; + return pipeline; + } + + webgpu_pipeline get_mul_mat_id_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_mul_mat_id_pipeline_key key = {}; + key.src0_type = context.src0->type; + key.src1_type = context.src1->type; + key.n_experts = context.src0->ne[2]; + key.vectorized = (context.src0->ne[0] % 4 == 0 && context.src0->ne[1] % 4 == 0 && + (context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ? + 1 : + 0; + + auto it = mul_mat_id_pipelines.find(key); + if (it != mul_mat_id_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "mul_mat_id"; + defines.push_back("MUL_MAT_ID"); + + // src1 type + switch (context.src1->type) { + case GGML_TYPE_F32: + defines.push_back("SRC1_INNER_TYPE=f32"); + break; + case GGML_TYPE_F16: + defines.push_back("SRC1_INNER_TYPE=f16"); + break; + default: + GGML_ABORT("Unsupported src1 type for mul_mat fast shader"); + } + + // src0 type + const struct ggml_type_traits * src0_traits = ggml_get_type_traits(context.src0->type); + const char * src0_name = src0_traits->type_name; + + switch (context.src0->type) { + case GGML_TYPE_F32: + defines.push_back("SRC0_INNER_TYPE=f32"); + defines.push_back("INIT_SRC0_SHMEM_FLOAT"); + defines.push_back("INIT_SRC1_SHMEM_FLOAT"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("SRC0_INNER_TYPE=f16"); + defines.push_back("INIT_SRC0_SHMEM_FLOAT"); + defines.push_back("INIT_SRC1_SHMEM_FLOAT"); + variant += "_f16"; + break; + default: + { + std::string type_upper = src0_name; + std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper); + + defines.push_back("BYTE_HELPERS"); + defines.push_back("INIT_SRC0_SHMEM_" + type_upper); + defines.push_back("INIT_SRC1_SHMEM_FLOAT"); + defines.push_back("U32_DEQUANT_HELPERS"); + defines.push_back("SRC0_INNER_TYPE=u32"); + + switch (context.src0->type) { + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + defines.push_back(type_upper + "_GRID"); + break; + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ3_S: + defines.push_back(type_upper + "_GRID"); + defines.push_back(type_upper + "_TABLES"); + break; + default: + break; + } + + variant += std::string("_") + src0_name; + break; + } + } + + // VEC/SCALAR controls + defines.push_back(key.vectorized ? "VEC" : "SCALAR"); + + // mul_mat_id is register-tile only. + const uint32_t tile_k = + ggml_is_quantized(context.src0->type) ? WEBGPU_MUL_MAT_REG_TILE_K_QUANT : WEBGPU_MUL_MAT_REG_TILE_K_FLOAT; + + // Tiles + defines.push_back("TILE_M=" + std::to_string(WEBGPU_MUL_MAT_TILE_M) + "u"); + defines.push_back("TILE_N=" + std::to_string(WEBGPU_MUL_MAT_TILE_N) + "u"); + defines.push_back("TILE_K=" + std::to_string(tile_k) + "u"); + + defines.push_back("WORKGROUP_SIZE_M=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_M) + "u"); + defines.push_back("WORKGROUP_SIZE_N=" + std::to_string(WEBGPU_MUL_MAT_WG_SIZE_N) + "u"); + + // variant suffix for src1 type + variant += std::string("_") + (context.src1->type == GGML_TYPE_F32 ? "f32" : "f16"); + if (key.vectorized) { + variant += "_vectorized"; + } + + auto processed = preprocessor.preprocess(wgsl_mul_mat_id, defines); + + auto decisions = std::make_shared(); + decisions->tile_k = tile_k; + decisions->tile_m = WEBGPU_MUL_MAT_TILE_M; + decisions->tile_n = WEBGPU_MUL_MAT_TILE_N; + decisions->wg_size_m = WEBGPU_MUL_MAT_WG_SIZE_M; + decisions->wg_size_n = WEBGPU_MUL_MAT_WG_SIZE_N; + decisions->wg_size = WEBGPU_MUL_MAT_WG_SIZE_M * WEBGPU_MUL_MAT_WG_SIZE_N; + + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + mul_mat_id_pipelines[key] = pipeline; + return mul_mat_id_pipelines[key]; + } + + webgpu_pipeline get_mul_mat_id_vec_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_mul_mat_id_pipeline_key key = {}; + key.src0_type = context.src0->type; + key.src1_type = context.src1->type; + key.n_experts = context.src0->ne[2]; + key.vectorized = (context.src0->ne[0] % 4 == 0 && + (context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ? + 1 : + 0; + + auto it = mul_mat_id_vec_pipelines.find(key); + if (it != mul_mat_id_vec_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "mul_mat_id_vec"; + const char * shader_src = wgsl_mul_mat_id_vec; + + // src1 type + switch (context.src1->type) { + case GGML_TYPE_F32: + defines.push_back("SRC1_INNER_TYPE=f32"); + break; + case GGML_TYPE_F16: + defines.push_back("SRC1_INNER_TYPE=f16"); + break; + default: + GGML_ABORT("Unsupported src1 type for mul_mat fast shader"); + } + + // src0 type + switch (context.src0->type) { + case GGML_TYPE_F32: + defines.push_back("SRC0_INNER_TYPE=f32"); + defines.push_back("MUL_ACC_FLOAT"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("SRC0_INNER_TYPE=f16"); + defines.push_back("MUL_ACC_FLOAT"); + variant += "_f16"; + break; + default: + { + // Quantized types: use helpers but accumulate in f16 + const struct ggml_type_traits * src0_traits = ggml_get_type_traits(context.src0->type); + std::string src0_name = src0_traits->type_name; + std::string type_upper = src0_name; + variant += "_" + src0_name; + std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper); + + defines.push_back("BYTE_HELPERS"); + defines.push_back("MUL_ACC_" + type_upper); + defines.push_back("U32_DEQUANT_HELPERS"); + defines.push_back("SRC0_INNER_TYPE=u32"); + switch (context.src0->type) { + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + defines.push_back(type_upper + "_GRID"); + break; + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ3_XXS: + defines.push_back(type_upper + "_GRID"); + defines.push_back(type_upper + "_TABLES"); + break; + default: + break; + } + break; + } + } + + // VEC/SCALAR controls + defines.push_back(key.vectorized ? "VEC" : "SCALAR"); + + uint32_t wg_size = WEBGPU_MUL_MAT_VEC_WG_SIZE; + uint32_t outputs_per_wg = WEBGPU_MUL_MAT_VEC_FLOAT_OUTPUTS_PER_WG; + + if (key.src0_type == GGML_TYPE_Q1_0) { + outputs_per_wg = WEBGPU_MUL_MAT_VEC_LEGACY_Q_OUTPUTS_PER_WG; + } else if (key.src0_type >= GGML_TYPE_Q2_K) { + outputs_per_wg = WEBGPU_MUL_MAT_VEC_K_Q_OUTPUTS_PER_WG; + } else if (key.src0_type >= GGML_TYPE_Q4_0) { + outputs_per_wg = WEBGPU_MUL_MAT_VEC_LEGACY_Q_OUTPUTS_PER_WG; + } + + // variant suffix for src1 type + variant += std::string("_") + (context.src1->type == GGML_TYPE_F32 ? "f32" : "f16"); + + defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); + defines.push_back(std::string("OUTPUTS_PER_WG=") + std::to_string(outputs_per_wg)); + defines.push_back(context.supports_subgroups ? "USE_SUBGROUP_REDUCTION" : "USE_WORKGROUP_REDUCTION"); + variant += context.supports_subgroups ? "_sg_reduce" : "_wg_reduce"; + if (key.vectorized) { + variant += "_vectorized"; + } + + defines.push_back(std::string("N_EXPERTS=") + std::to_string(key.n_experts)); + + auto processed = preprocessor.preprocess(shader_src, defines); + + auto decisions = std::make_shared(); + decisions->wg_size = wg_size; + decisions->outputs_per_wg = outputs_per_wg; + + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + mul_mat_id_vec_pipelines[key] = pipeline; + return mul_mat_id_vec_pipelines[key]; + } + + webgpu_pipeline get_unary_pipeline(const ggml_webgpu_shader_lib_context & context) { + const bool is_unary = context.dst->op == GGML_OP_UNARY; + const int op = is_unary ? (int) ggml_get_unary_op(context.dst) : context.dst->op; + ggml_webgpu_unary_pipeline_key key = {}; + key.type = context.dst->type; + key.op = op; + key.is_unary = is_unary; + key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst) || context.dst->op == GGML_OP_FILL; + key.ttype = (ggml_tri_type) ggml_get_op_params_i32(context.dst, 0); + + auto it = unary_pipelines.find(key); + if (it != unary_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = + key.is_unary ? ggml_unary_op_name((ggml_unary_op) key.op) : ggml_op_name((ggml_op) key.op); + defines.push_back(variant); + + switch (key.type) { + case GGML_TYPE_F32: + defines.push_back("TYPE_F32"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("TYPE_F16"); + variant += "_f16"; + break; + default: + GGML_ABORT("Unsupported type for unary shader"); + } + + if (key.inplace) { + defines.push_back("INPLACE"); + variant += "_inplace"; + } + + if (op == GGML_OP_TRI) { + switch (key.ttype) { + case GGML_TRI_TYPE_LOWER: + defines.push_back("TRI_TYPE_LOWER"); + variant += "_tri_type_lower"; + break; + case GGML_TRI_TYPE_LOWER_DIAG: + defines.push_back("TRI_TYPE_LOWER_DIAG"); + variant += "_tri_type_lower_diag"; + break; + case GGML_TRI_TYPE_UPPER: + defines.push_back("TRI_TYPE_UPPER"); + variant += "_tri_type_upper"; + break; + case GGML_TRI_TYPE_UPPER_DIAG: + defines.push_back("TRI_TYPE_UPPER_DIAG"); + variant += "_tri_upper_diag"; + break; + default: + GGML_ABORT("Unsupported ggml_tri_type for unary shader"); + } + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_unary, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + decisions->inplace = key.inplace; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + unary_pipelines[key] = pipeline; + return unary_pipelines[key]; + } + + webgpu_pipeline get_rms_norm_mul_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_rms_norm_mul_pipeline_key key = {}; + key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); + key.overlap = ggml_webgpu_tensor_equal(context.src1, context.dst); + key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1); + + auto it = rms_norm_mul_pipelines.find(key); + if (it != rms_norm_mul_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string op_name = "RMS_NORM_MUL"; + std::string variant = op_name; + + if (key.inplace) { + defines.push_back("INPLACE"); + variant += "_inplace"; + } else if (key.overlap) { + defines.push_back("OVERLAP"); + variant += "_overlap"; + } else if (key.src_overlap) { + defines.push_back("SRC_OVERLAP"); + variant += "_src_overlap"; + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_rms_norm_mul, defines); + auto pipeline_decisions = std::make_shared(); + pipeline_decisions->wg_size = context.max_wg_size; + pipeline_decisions->inplace = key.inplace; + pipeline_decisions->overlap = key.overlap; + pipeline_decisions->src_overlap = key.src_overlap; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = pipeline_decisions; + rms_norm_mul_pipelines[key] = pipeline; + return rms_norm_mul_pipelines[key]; + } + + webgpu_pipeline get_binary_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_binary_pipeline_key key = {}; + key.type = context.dst->type; + key.op = context.dst->op; + key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); + key.overlap = ggml_webgpu_tensor_equal(context.src1, context.dst); + key.src_overlap = ggml_webgpu_tensor_overlap(context.src0, context.src1); + + auto it = binary_pipelines.find(key); + if (it != binary_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string op_name = ggml_op_name((ggml_op) key.op); + std::string variant = op_name; + + defines.push_back(std::string("OP_") + op_name); + + switch (key.type) { + case GGML_TYPE_F32: + defines.push_back("TYPE_F32"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("TYPE_F16"); + variant += "_f16"; + break; + default: + GGML_ABORT("Unsupported type for binary shader"); + } + + if (key.inplace) { + defines.push_back("INPLACE"); + variant += "_inplace"; + } else if (key.overlap) { + defines.push_back("OVERLAP"); + variant += "_overlap"; + } else if (key.src_overlap) { + defines.push_back("SRC_OVERLAP"); + variant += "_src_overlap"; + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_binary, defines); + auto pipeline_decisions = std::make_shared(); + pipeline_decisions->wg_size = context.max_wg_size; + pipeline_decisions->inplace = key.inplace; + pipeline_decisions->overlap = key.overlap; + pipeline_decisions->src_overlap = key.src_overlap; + + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = pipeline_decisions; + binary_pipelines[key] = pipeline; + return binary_pipelines[key]; + } + + webgpu_pipeline get_concat_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_concat_pipeline_key key = {}; + key.type = context.dst->type; + + auto it = concat_pipelines.find(key); + if (it != concat_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "concat"; + + switch (key.type) { + case GGML_TYPE_F32: + defines.push_back("TYPE_F32"); + variant += "_f32"; + break; + case GGML_TYPE_I32: + defines.push_back("TYPE_I32"); + variant += "_i32"; + break; + default: + GGML_ABORT("Unsupported type for concat shader"); + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_concat, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + concat_pipelines[key] = pipeline; + return concat_pipelines[key]; + } + + webgpu_pipeline get_repeat_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_repeat_pipeline_key key = {}; + key.type = context.dst->type; + + auto it = repeat_pipelines.find(key); + if (it != repeat_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "repeat"; + + switch (key.type) { + case GGML_TYPE_F32: + defines.push_back("TYPE_F32"); + variant += "_f32"; + break; + case GGML_TYPE_I32: + defines.push_back("TYPE_I32"); + variant += "_i32"; + break; + case GGML_TYPE_I16: + defines.push_back("TYPE_I16"); + variant += "_i16"; + break; + default: + GGML_ABORT("Unsupported type for repeat shader"); + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_repeat, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + repeat_pipelines[key] = pipeline; + return repeat_pipelines[key]; + } + + webgpu_pipeline get_flash_attn_pipeline(const ggml_webgpu_shader_lib_context & context, + size_t storage_offset_alignment) { + const ggml_webgpu_flash_attn_decisions decisions = + ggml_webgpu_flash_attn_get_decisions(context, storage_offset_alignment); + GGML_ASSERT(decisions.path != GGML_WEBGPU_FLASH_ATTN_PATH_NONE); + ggml_webgpu_flash_attn_pipeline_key key = ggml_webgpu_flash_attn_make_pipeline_key(context, decisions.path); + auto it = flash_attn_pipelines.find(key); + if (it != flash_attn_pipelines.end()) { + return it->second; + } + std::vector defines; + std::string variant = decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_VEC ? "flash_attn_vec" : + decisions.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE ? "flash_attn_tile" : + "flash_attn"; + + switch (key.kv_type) { + case GGML_TYPE_F32: + defines.push_back("KV_F32"); + break; + case GGML_TYPE_F16: + defines.push_back("KV_F16"); + break; + case GGML_TYPE_Q4_0: + defines.push_back("KV_Q4_0"); + break; + case GGML_TYPE_Q8_0: + defines.push_back("KV_Q8_0"); + break; + default: + GGML_ABORT("Unsupported KV type for flash attention shader"); + } + variant += std::string("_") + ggml_type_name(key.kv_type); + + if (key.has_mask) { + defines.push_back("MASK"); + if (key.path == GGML_WEBGPU_FLASH_ATTN_PATH_VEC) { + defines.push_back("BLK"); + variant += "_mask_blk"; + } else { + variant += "_mask"; + } + } + if (key.has_sinks) { + defines.push_back("SINKS"); + variant += "_sinks"; + } + if (key.uses_logit_softcap) { + defines.push_back("LOGIT_SOFTCAP"); + variant += "_lgsc"; + } + if (key.kv_direct) { + defines.push_back("KV_DIRECT"); + variant += "_kvdirect"; + } + if (key.kv_overlap) { + defines.push_back("KV_OVERLAP"); + variant += "_kv_overlap"; + } + + defines.push_back(std::string("HEAD_DIM_QK=") + std::to_string(key.head_dim_qk)); + variant += std::string("_hsqk") + std::to_string(key.head_dim_qk); + + defines.push_back(std::string("HEAD_DIM_V=") + std::to_string(key.head_dim_v)); + variant += std::string("_hsv") + std::to_string(key.head_dim_v); + + const char * shader_src = wgsl_flash_attn; + if (key.path == GGML_WEBGPU_FLASH_ATTN_PATH_VEC) { + defines.push_back("KV_GRANULARITY=8"); + defines.push_back(std::string("VEC_NE=") + std::to_string(ggml_webgpu_flash_attn_pick_vec_ne(key)) + "u"); + shader_src = wgsl_flash_attn_vec_split; + } else if (key.path == GGML_WEBGPU_FLASH_ATTN_PATH_TILE) { + shader_src = wgsl_flash_attn_tile; + defines.push_back("MAX_SUBGROUP_SIZE=" + std::to_string(context.max_subgroup_size)); + defines.push_back("KV_STAGE_STRIDE=" + std::to_string(std::max(key.head_dim_qk, key.head_dim_v))); + variant += "_tile"; + } else { + defines.push_back(std::string("SG_MAT_M=") + std::to_string(context.sg_mat_m)); + defines.push_back(std::string("SG_MAT_N=") + std::to_string(context.sg_mat_n)); + defines.push_back(std::string("SG_MAT_K=") + std::to_string(context.sg_mat_k)); + } + + auto pipeline_decisions = std::make_shared(decisions); + pipeline_decisions->kv_overlap = key.kv_overlap; + defines.push_back(std::string("Q_TILE=") + std::to_string(decisions.q_tile)); + defines.push_back(std::string("KV_TILE=") + std::to_string(decisions.kv_tile)); + defines.push_back(std::string("WG_SIZE=") + std::to_string(decisions.wg_size)); + + webgpu_pipeline pipeline = + ggml_webgpu_create_pipeline(device, preprocessor.preprocess(shader_src, defines), variant); + pipeline.context = pipeline_decisions; + flash_attn_pipelines[key] = pipeline; + return flash_attn_pipelines[key]; + } + + webgpu_pipeline get_flash_attn_blk_pipeline(const ggml_webgpu_shader_lib_context & context, uint32_t kv_tile) { + ggml_webgpu_flash_attn_blk_pipeline_key key = {}; + key.kv_tile = kv_tile; + auto it = flash_attn_blk_pipelines.find(key); + if (it != flash_attn_blk_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "flash_attn_vec_blk"; + + defines.push_back(std::string("KV_TILE=") + std::to_string(key.kv_tile)); + variant += std::string("_kvt") + std::to_string(key.kv_tile); + + uint32_t wg_size = 1; + while ((wg_size << 1) <= context.max_wg_size) { + wg_size <<= 1; + } + defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size)); + variant += std::string("_wg") + std::to_string(wg_size); + + webgpu_pipeline pipeline = + ggml_webgpu_create_pipeline(device, preprocessor.preprocess(wgsl_flash_attn_vec_blk, defines), variant); + flash_attn_blk_pipelines[key] = pipeline; + return flash_attn_blk_pipelines[key]; + } + + webgpu_pipeline get_flash_attn_vec_reduce_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_flash_attn_vec_reduce_pipeline_key key = {}; + key.head_dim_v = (uint32_t) context.src2->ne[0]; + key.wg_size = context.max_wg_size; + auto it = flash_attn_vec_reduce_pipelines.find(key); + if (it != flash_attn_vec_reduce_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "flash_attn_vec_reduce"; + + defines.push_back(std::string("HEAD_DIM_V=") + std::to_string(key.head_dim_v)); + variant += std::string("_hsv") + std::to_string(key.head_dim_v); + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + variant += std::string("_wg") + std::to_string(context.max_wg_size); + + webgpu_pipeline pipeline = + ggml_webgpu_create_pipeline(device, preprocessor.preprocess(wgsl_flash_attn_vec_reduce, defines), variant); + flash_attn_vec_reduce_pipelines[key] = pipeline; + return flash_attn_vec_reduce_pipelines[key]; + } + + webgpu_pipeline get_cpy_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_cpy_pipeline_key key = {}; + key.src_type = context.src0->type; + key.dst_type = context.dst->type; + + auto it = cpy_pipelines.find(key); + if (it != cpy_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "cpy"; + + switch (key.src_type) { + case GGML_TYPE_F32: + defines.push_back("SRC_F32"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("SRC_F16"); + variant += "_f16"; + break; + default: + GGML_ABORT("Unsupported src type for cpy shader"); + } + + switch (key.dst_type) { + case GGML_TYPE_F32: + defines.push_back("DST_F32"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("DST_F16"); + variant += "_f16"; + break; + case GGML_TYPE_I32: + defines.push_back("DST_I32"); + variant += "_i32"; + break; + default: + GGML_ABORT("Unsupported dst type for cpy shader"); + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_cpy, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + cpy_pipelines[key] = pipeline; + return cpy_pipelines[key]; + } + + webgpu_pipeline get_glu_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_glu_pipeline_key key = {}; + key.glu_op = ggml_get_glu_op(context.dst); + key.type = context.dst->type; + key.split = (context.src1 != nullptr); + + auto it = glu_pipelines.find(key); + if (it != glu_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "glu"; + + switch (key.glu_op) { + case GGML_GLU_OP_REGLU: + defines.push_back("OP_REGLU"); + variant += "_reglu"; + break; + case GGML_GLU_OP_GEGLU: + defines.push_back("OP_GEGLU"); + variant += "_geglu"; + break; + case GGML_GLU_OP_SWIGLU: + defines.push_back("OP_SWIGLU"); + variant += "_swiglu"; + break; + case GGML_GLU_OP_SWIGLU_OAI: + defines.push_back("OP_SWIGLU_OAI"); + variant += "_swiglu_oai"; + break; + case GGML_GLU_OP_GEGLU_ERF: + defines.push_back("OP_GEGLU_ERF"); + variant += "_geglu_erf"; + break; + case GGML_GLU_OP_GEGLU_QUICK: + defines.push_back("OP_GEGLU_QUICK"); + variant += "_geglu_quick"; + break; + default: + GGML_ABORT("Unsupported GLU op"); + } + switch (key.type) { + case GGML_TYPE_F32: + defines.push_back("TYPE_F32"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("TYPE_F16"); + variant += "_f16"; + break; + default: + GGML_ABORT("Unsupported type for GLU shader"); + } + + if (key.split) { + variant += "_split"; + } else { + defines.push_back("NO_SPLIT"); + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_glu, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + glu_pipelines[key] = pipeline; + return glu_pipelines[key]; + } + + webgpu_pipeline get_rope_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_rope_pipeline_key key = {}; + key.type = context.dst->type; + key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); + key.has_ff = (context.src2 != nullptr); + + auto it = rope_pipelines.find(key); + if (it != rope_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "rope"; + + switch (key.type) { + case GGML_TYPE_F32: + defines.push_back("TYPE_F32"); + variant += "_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("TYPE_F16"); + variant += "_f16"; + break; + default: + GGML_ABORT("Unsupported type for ROPE shader"); + } + + if (key.inplace) { + defines.push_back("INPLACE"); + variant += "_inplace"; + } + + if (key.has_ff) { + defines.push_back("FF_FUNC"); + variant += "_ff"; + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_rope, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + decisions->inplace = key.inplace; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + rope_pipelines[key] = pipeline; + return rope_pipelines[key]; + } + + webgpu_pipeline get_soft_max_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_soft_max_pipeline_key key = {}; + key.mask_type = context.src1 ? context.src1->type : GGML_TYPE_F32; + key.has_mask = (context.src1 != nullptr); + key.has_sink = (context.src2 != nullptr); + key.inplace = ggml_webgpu_tensor_equal(context.src0, context.dst); + + auto it = soft_max_pipelines.find(key); + if (it != soft_max_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "soft_max"; + + if (key.has_mask) { + defines.push_back("HAS_MASK"); + switch (key.mask_type) { + case GGML_TYPE_F32: + defines.push_back("MASK_F32"); + variant += "_mask_f32"; + break; + case GGML_TYPE_F16: + defines.push_back("MASK_F16"); + variant += "_mask_f16"; + break; + default: + GGML_ABORT("Unsupported type for SOFT_MAX shader"); + } + } + + if (key.has_sink) { + defines.push_back("HAS_SINK"); + variant += "_sink"; + } + + if (key.inplace) { + defines.push_back("INPLACE"); + variant += "_inplace"; + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_soft_max, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + decisions->inplace = key.inplace; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + soft_max_pipelines[key] = pipeline; + return soft_max_pipelines[key]; + } + + webgpu_pipeline get_conv2d_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_conv2d_pipeline_key key = {}; + key.weight_type = context.src0->type; + key.input_type = context.src1->type; + key.output_type = context.dst->type; + + auto it = conv2d_pipelines.find(key); + if (it != conv2d_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "conv_2d"; + + auto push_type_defines = [&](const char * prefix, ggml_type type) { + std::string s_prefix = prefix; + if (type == GGML_TYPE_F32) { + defines.push_back(s_prefix + "_F32"); + } else if (type == GGML_TYPE_F16) { + defines.push_back(s_prefix + "_F16"); + } else { + GGML_ABORT("Unsupported type for CONV_2D shader"); + } + }; + + push_type_defines("WEIGHT", key.weight_type); + push_type_defines("INPUT", key.input_type); + push_type_defines("OUTPUT", key.output_type); + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_conv2d, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + conv2d_pipelines[key] = pipeline; + return conv2d_pipelines[key]; + } + + webgpu_pipeline get_im2col_pipeline(const ggml_webgpu_shader_lib_context & context) { + ggml_webgpu_im2col_pipeline_key key = {}; + key.input_type = context.src1->type; + key.output_type = context.dst->type; + + auto it = im2col_pipelines.find(key); + if (it != im2col_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "im2col"; + + auto push_type_defines = [&](const char * prefix, ggml_type type) { + std::string s_prefix = prefix; + if (type == GGML_TYPE_F32) { + defines.push_back(s_prefix + "_F32"); + } else if (type == GGML_TYPE_F16) { + defines.push_back(s_prefix + "_F16"); + } else { + GGML_ABORT("Unsupported type for IM2COL shader"); + } + }; + + push_type_defines("INPUT", key.input_type); + push_type_defines("OUTPUT", key.output_type); + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_im2col, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + im2col_pipelines[key] = pipeline; + return im2col_pipelines[key]; + } + + webgpu_pipeline get_upscale_pipeline(const ggml_webgpu_shader_lib_context & context) { + const uint32_t mode_flags = (uint32_t) ggml_get_op_params_i32(context.dst, 0); + const uint32_t base_mode = mode_flags & 0xFFu; + const bool antialias = (mode_flags & GGML_SCALE_FLAG_ANTIALIAS) != 0u; + + ggml_webgpu_upscale_pipeline_key key = {}; + key.input_type = context.src0->type; + key.output_type = context.dst->type; + key.base_mode = base_mode; + key.antialias = antialias; + + auto it = upscale_pipelines.find(key); + if (it != upscale_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = "upscale"; + + if (key.input_type == GGML_TYPE_F16) { + defines.push_back("SRC_F16"); + variant += "_src_f16"; + } else { + variant += "_src_f32"; + } + + if (key.output_type == GGML_TYPE_F16) { + defines.push_back("DST_F16"); + variant += "_dst_f16"; + } else { + variant += "_dst_f32"; + } + + switch (base_mode) { + case GGML_SCALE_MODE_NEAREST: + defines.push_back("NEAREST"); + variant += "_nearest"; + break; + case GGML_SCALE_MODE_BILINEAR: + defines.push_back("BILINEAR"); + variant += "_bilinear"; + break; + case GGML_SCALE_MODE_BICUBIC: + defines.push_back("BICUBIC"); + variant += "_bicubic"; + break; + default: + GGML_ABORT("Unsupported upscale mode"); + } + + if (antialias) { + defines.push_back("ANTIALIAS"); + variant += "_aa"; + } + + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_upscale, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + upscale_pipelines[key] = pipeline; + return upscale_pipelines[key]; + } + + private: + static webgpu_pipeline ggml_webgpu_create_pipeline(wgpu::Device & device, + std::string shader_code, + std::string label) { + wgpu::ShaderSourceWGSL shader_source; + shader_source.code = shader_code.c_str(); + + wgpu::ShaderModuleDescriptor shader_desc; + shader_desc.nextInChain = &shader_source; + + wgpu::ShaderModule shader_module = device.CreateShaderModule(&shader_desc); + + wgpu::ComputePipelineDescriptor pipeline_desc; + pipeline_desc.label = label.c_str(); + pipeline_desc.compute.module = shader_module; + pipeline_desc.compute.entryPoint = "main"; // Entry point in the WGSL code + pipeline_desc.layout = nullptr; // nullptr means auto layout + return { device.CreateComputePipeline(&pipeline_desc), label }; + } +}; + -#include "ggml-backend-impl.h" -#include "ggml-impl.h" -#include "ggml-webgpu-shader-lib.hpp" -#include "ggml.h" +#include "ggml.h.inc" #ifdef __EMSCRIPTEN__ # include diff --git a/ggml/src/ggml-webgpu/pre_wgsl.hpp b/ggml/src/ggml-webgpu/pre_wgsl.hpp deleted file mode 100644 index 4d4359463cac..000000000000 --- a/ggml/src/ggml-webgpu/pre_wgsl.hpp +++ /dev/null @@ -1,778 +0,0 @@ -#ifndef PRE_WGSL_HPP -#define PRE_WGSL_HPP - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace pre_wgsl { - -//============================================================== -// Options -//============================================================== -struct Options { - std::string include_path = "."; - std::vector macros; -}; - -//============================================================== -// Utility: trim -//============================================================== -static std::string trim(const std::string & s) { - size_t a = 0; - while (a < s.size() && std::isspace((unsigned char) s[a])) { - a++; - } - size_t b = s.size(); - while (b > a && std::isspace((unsigned char) s[b - 1])) { - b--; - } - return s.substr(a, b - a); -} - -static std::string trim_value(std::istream & is) { - std::string str; - std::getline(is, str); - return trim(str); -} - -static bool isIdentChar(char c) { - return std::isalnum(static_cast(c)) || c == '_'; -} - -static std::string expandMacrosRecursiveInternal(const std::string & line, - const std::unordered_map & macros, - std::unordered_set & visiting); - -static std::string expandMacroValue(const std::string & name, - const std::unordered_map & macros, - std::unordered_set & visiting) { - if (visiting.count(name)) { - throw std::runtime_error("Recursive macro: " + name); - } - visiting.insert(name); - - auto it = macros.find(name); - if (it == macros.end()) { - visiting.erase(name); - return name; - } - - const std::string & value = it->second; - if (value.empty()) { - visiting.erase(name); - return ""; - } - - std::string expanded = expandMacrosRecursiveInternal(value, macros, visiting); - visiting.erase(name); - return expanded; -} - -static std::string expandMacrosRecursiveInternal(const std::string & line, - const std::unordered_map & macros, - std::unordered_set & visiting) { - std::string result; - result.reserve(line.size()); - - size_t i = 0; - while (i < line.size()) { - if (isIdentChar(line[i])) { - size_t start = i; - while (i < line.size() && isIdentChar(line[i])) { - i++; - } - std::string token = line.substr(start, i - start); - - auto it = macros.find(token); - if (it != macros.end()) { - result += expandMacroValue(token, macros, visiting); - } else { - result += token; - } - } else { - result += line[i]; - i++; - } - } - - return result; -} - -static std::string expandMacrosRecursive(const std::string & line, - const std::unordered_map & macros) { - std::unordered_set visiting; - return expandMacrosRecursiveInternal(line, macros, visiting); -} - -//============================================================== -// Tokenizer for expressions in #if/#elif -//============================================================== -class ExprLexer { - public: - enum Kind { END, IDENT, NUMBER, OP, LPAREN, RPAREN }; - - struct Tok { - Kind kind; - std::string text; - }; - - explicit ExprLexer(std::string_view sv) : src(sv), pos(0) {} - - Tok next() { - skipWS(); - if (pos >= src.size()) { - return { END, "" }; - } - - char c = src[pos]; - - // number - if (std::isdigit((unsigned char) c)) { - size_t start = pos; - while (pos < src.size() && std::isdigit((unsigned char) src[pos])) { - pos++; - } - return { NUMBER, std::string(src.substr(start, pos - start)) }; - } - - // identifier - if (std::isalpha((unsigned char) c) || c == '_') { - size_t start = pos; - while (pos < src.size() && (std::isalnum((unsigned char) src[pos]) || src[pos] == '_')) { - pos++; - } - return { IDENT, std::string(src.substr(start, pos - start)) }; - } - - if (c == '(') { - pos++; - return { LPAREN, "(" }; - } - if (c == ')') { - pos++; - return { RPAREN, ")" }; - } - - // multi-char operators - static const char * two_ops[] = { "==", "!=", "<=", ">=", "&&", "||", "<<", ">>" }; - for (auto op : two_ops) { - if (src.substr(pos, 2) == op) { - pos += 2; - return { OP, std::string(op) }; - } - } - - // single-char operators - if (std::string("+-*/%<>!").find(c) != std::string::npos) { - pos++; - return { OP, std::string(1, c) }; - } - - // unexpected - pos++; - return { END, "" }; - } - - private: - std::string_view src; - size_t pos; - - void skipWS() { - while (pos < src.size() && std::isspace((unsigned char) src[pos])) { - pos++; - } - } -}; - -//============================================================== -// Expression Parser (recursive descent) -//============================================================== -class ExprParser { - public: - ExprParser(std::string_view expr, - const std::unordered_map & macros, - std::unordered_set & visiting) : - lex(expr), - macros(macros), - visiting(visiting) { - advance(); - } - - int parse() { return parseLogicalOr(); } - - private: - ExprLexer lex; - ExprLexer::Tok tok; - const std::unordered_map & macros; - std::unordered_set & visiting; - - void advance() { tok = lex.next(); } - - bool acceptOp(const std::string & s) { - if (tok.kind == ExprLexer::OP && tok.text == s) { - advance(); - return true; - } - return false; - } - - bool acceptKind(ExprLexer::Kind k) { - if (tok.kind == k) { - advance(); - return true; - } - return false; - } - - int parseLogicalOr() { - int v = parseLogicalAnd(); - while (acceptOp("||")) { - int rhs = parseLogicalAnd(); - v = (v || rhs); - } - return v; - } - - int parseLogicalAnd() { - int v = parseEquality(); - while (acceptOp("&&")) { - int rhs = parseEquality(); - v = (v && rhs); - } - return v; - } - - int parseEquality() { - int v = parseRelational(); - for (;;) { - if (acceptOp("==")) { - int rhs = parseRelational(); - v = (v == rhs); - } else if (acceptOp("!=")) { - int rhs = parseRelational(); - v = (v != rhs); - } else { - break; - } - } - return v; - } - - int parseRelational() { - int v = parseShift(); - for (;;) { - if (acceptOp("<")) { - int rhs = parseShift(); - v = (v < rhs); - } else if (acceptOp(">")) { - int rhs = parseShift(); - v = (v > rhs); - } else if (acceptOp("<=")) { - int rhs = parseShift(); - v = (v <= rhs); - } else if (acceptOp(">=")) { - int rhs = parseShift(); - v = (v >= rhs); - } else { - break; - } - } - return v; - } - - int parseShift() { - int v = parseAdd(); - for (;;) { - if (acceptOp("<<")) { - int rhs = parseAdd(); - v = (v << rhs); - } else if (acceptOp(">>")) { - int rhs = parseAdd(); - v = (v >> rhs); - } else { - break; - } - } - return v; - } - - int parseAdd() { - int v = parseMult(); - for (;;) { - if (acceptOp("+")) { - int rhs = parseMult(); - v = (v + rhs); - } else if (acceptOp("-")) { - int rhs = parseMult(); - v = (v - rhs); - } else { - break; - } - } - return v; - } - - int parseMult() { - int v = parseUnary(); - for (;;) { - if (acceptOp("*")) { - int rhs = parseUnary(); - v = (v * rhs); - } else if (acceptOp("/")) { - int rhs = parseUnary(); - v = (rhs == 0 ? 0 : v / rhs); - } else if (acceptOp("%")) { - int rhs = parseUnary(); - v = (rhs == 0 ? 0 : v % rhs); - } else { - break; - } - } - return v; - } - - int parseUnary() { - if (acceptOp("!")) { - return !parseUnary(); - } - if (acceptOp("-")) { - return -parseUnary(); - } - if (acceptOp("+")) { - return +parseUnary(); - } - return parsePrimary(); - } - - int parsePrimary() { - // '(' expr ')' - if (acceptKind(ExprLexer::LPAREN)) { - int v = parse(); - if (!acceptKind(ExprLexer::RPAREN)) { - throw std::runtime_error("missing ')'"); - } - return v; - } - - // number - if (tok.kind == ExprLexer::NUMBER) { - int v = std::stoi(tok.text); - advance(); - return v; - } - - // defined(identifier) - if (tok.kind == ExprLexer::IDENT && tok.text == "defined") { - advance(); - if (acceptKind(ExprLexer::LPAREN)) { - if (tok.kind != ExprLexer::IDENT) { - throw std::runtime_error("expected identifier in defined()"); - } - std::string name = tok.text; - advance(); - if (!acceptKind(ExprLexer::RPAREN)) { - throw std::runtime_error("missing ) in defined()"); - } - return macros.count(name) ? 1 : 0; - } else { - // defined NAME - if (tok.kind != ExprLexer::IDENT) { - throw std::runtime_error("expected identifier in defined NAME"); - } - std::string name = tok.text; - advance(); - return macros.count(name) ? 1 : 0; - } - } - - // identifier -> treat as integer, if defined use its value else 0 - if (tok.kind == ExprLexer::IDENT) { - std::string name = tok.text; - advance(); - auto it = macros.find(name); - if (it == macros.end()) { - return 0; - } - if (it->second.empty()) { - return 1; - } - return evalMacroExpression(name, it->second); - } - - // unexpected - return 0; - } - - int evalMacroExpression(const std::string & name, const std::string & value) { - if (visiting.count(name)) { - throw std::runtime_error("Recursive macro: " + name); - } - - visiting.insert(name); - ExprParser ep(value, macros, visiting); - int v = ep.parse(); - visiting.erase(name); - return v; - } -}; - -//============================================================== -// Preprocessor -//============================================================== -class Preprocessor { - public: - explicit Preprocessor(Options opts = {}) : opts_(std::move(opts)) { - // Treat empty include path as current directory - if (opts_.include_path.empty()) { - opts_.include_path = "."; - } - parseMacroDefinitions(opts_.macros); - } - - std::string preprocess_file(const std::string & filename, const std::vector & additional_macros = {}) { - std::unordered_map macros; - std::unordered_set predefined; - std::unordered_set include_stack; - buildMacros(additional_macros, macros, predefined); - - std::string result = processFile(filename, macros, predefined, include_stack, DirectiveMode::All); - return result; - } - - std::string preprocess(const std::string & contents, const std::vector & additional_macros = {}) { - std::unordered_map macros; - std::unordered_set predefined; - std::unordered_set include_stack; - buildMacros(additional_macros, macros, predefined); - - std::string result = processString(contents, macros, predefined, include_stack, DirectiveMode::All); - return result; - } - - std::string preprocess_includes_file(const std::string & filename) { - std::unordered_map macros; - std::unordered_set predefined; - std::unordered_set include_stack; - std::string result = processFile(filename, macros, predefined, include_stack, DirectiveMode::IncludesOnly); - return result; - } - - std::string preprocess_includes(const std::string & contents) { - std::unordered_map macros; - std::unordered_set predefined; - std::unordered_set include_stack; - std::string result = processString(contents, macros, predefined, include_stack, DirectiveMode::IncludesOnly); - return result; - } - - private: - Options opts_; - std::unordered_map global_macros; - - enum class DirectiveMode { All, IncludesOnly }; - - struct Cond { - bool parent_active; - bool active; - bool taken; - }; - - //---------------------------------------------------------- - // Parse macro definitions into global_macros - //---------------------------------------------------------- - void parseMacroDefinitions(const std::vector & macro_defs) { - for (const auto & def : macro_defs) { - size_t eq_pos = def.find('='); - if (eq_pos != std::string::npos) { - // Format: NAME=VALUE - std::string name = trim(def.substr(0, eq_pos)); - std::string value = trim(def.substr(eq_pos + 1)); - global_macros[name] = value; - } else { - // Format: NAME - std::string name = trim(def); - global_macros[name] = ""; - } - } - } - - //---------------------------------------------------------- - // Build combined macro map and predefined set for a preprocessing operation - //---------------------------------------------------------- - void buildMacros(const std::vector & additional_macros, - std::unordered_map & macros, - std::unordered_set & predefined) { - macros = global_macros; - predefined.clear(); - - for (const auto & [name, value] : global_macros) { - predefined.insert(name); - } - - for (const auto & def : additional_macros) { - size_t eq_pos = def.find('='); - std::string name, value; - if (eq_pos != std::string::npos) { - name = trim(def.substr(0, eq_pos)); - value = trim(def.substr(eq_pos + 1)); - } else { - name = trim(def); - value = ""; - } - - // Add to macros map (will override global if same name) - macros[name] = value; - predefined.insert(name); - } - } - - //---------------------------------------------------------- - // Helpers - //---------------------------------------------------------- - std::string loadFile(const std::string & fname) { - std::ifstream f(fname); - if (!f.is_open()) { - throw std::runtime_error("Could not open file: " + fname); - } - std::stringstream ss; - ss << f.rdbuf(); - return ss.str(); - } - - bool condActive(const std::vector & cond) const { - if (cond.empty()) { - return true; - } - return cond.back().active; - } - - //---------------------------------------------------------- - // Process a file - //---------------------------------------------------------- - std::string processFile(const std::string & name, - std::unordered_map & macros, - const std::unordered_set & predefined_macros, - std::unordered_set & include_stack, - DirectiveMode mode) { - if (include_stack.count(name)) { - throw std::runtime_error("Recursive include: " + name); - } - - include_stack.insert(name); - std::string shader_code = loadFile(name); - std::string out = processString(shader_code, macros, predefined_macros, include_stack, mode); - include_stack.erase(name); - return out; - } - - std::string processIncludeFile(const std::string & fname, - std::unordered_map & macros, - const std::unordered_set & predefined_macros, - std::unordered_set & include_stack, - DirectiveMode mode) { - std::string full_path = opts_.include_path + "/" + fname; - return processFile(full_path, macros, predefined_macros, include_stack, mode); - } - - //---------------------------------------------------------- - // Process text - //---------------------------------------------------------- - std::string processString(const std::string & shader_code, - std::unordered_map & macros, - const std::unordered_set & predefined_macros, - std::unordered_set & include_stack, - DirectiveMode mode) { - std::vector cond; // Conditional stack for this shader - std::stringstream out; - std::istringstream in(shader_code); - std::string line; - - while (std::getline(in, line)) { - std::string t = trim(line); - - if (!t.empty() && t[0] == '#') { - bool handled = handleDirective(t, out, macros, predefined_macros, cond, include_stack, mode); - if (mode == DirectiveMode::IncludesOnly && !handled) { - out << line << "\n"; - } - } else { - if (mode == DirectiveMode::IncludesOnly) { - out << line << "\n"; - } else if (condActive(cond)) { - // Expand macros in the line before outputting - std::string expanded = expandMacrosRecursive(line, macros); - out << expanded << "\n"; - } - } - } - - if (mode == DirectiveMode::All && !cond.empty()) { - throw std::runtime_error("Unclosed #if directive"); - } - - return out.str(); - } - - //---------------------------------------------------------- - // Directive handler - //---------------------------------------------------------- - bool handleDirective(const std::string & t, - std::stringstream & out, - std::unordered_map & macros, - const std::unordered_set & predefined_macros, - std::vector & cond, - std::unordered_set & include_stack, - DirectiveMode mode) { - // split into tokens - std::string body = t.substr(1); - std::istringstream iss(body); - std::string cmd; - iss >> cmd; - - if (cmd == "include") { - if (mode == DirectiveMode::All && !condActive(cond)) { - return true; - } - std::string file; - iss >> file; - if (file.size() >= 2 && file.front() == '"' && file.back() == '"') { - file = file.substr(1, file.size() - 2); - } - out << processIncludeFile(file, macros, predefined_macros, include_stack, mode); - return true; - } - - if (mode == DirectiveMode::IncludesOnly) { - return false; - } - - if (cmd == "define") { - if (!condActive(cond)) { - return true; - } - std::string name; - iss >> name; - // Don't override predefined macros from options - if (predefined_macros.count(name)) { - return true; - } - std::string value = trim_value(iss); - macros[name] = value; - return true; - } - - if (cmd == "undef") { - if (!condActive(cond)) { - return true; - } - std::string name; - iss >> name; - // Don't undef predefined macros from options - if (predefined_macros.count(name)) { - return true; - } - macros.erase(name); - return true; - } - - if (cmd == "ifdef") { - std::string name; - iss >> name; - bool p = condActive(cond); - bool v = macros.count(name); - cond.push_back({ p, p && v, p && v }); - return true; - } - - if (cmd == "ifndef") { - std::string name; - iss >> name; - bool p = condActive(cond); - bool v = !macros.count(name); - cond.push_back({ p, p && v, p && v }); - return true; - } - - if (cmd == "if") { - std::string expr = trim_value(iss); - bool p = condActive(cond); - bool v = false; - if (p) { - std::unordered_set visiting; - ExprParser ep(expr, macros, visiting); - v = ep.parse() != 0; - } - cond.push_back({ p, p && v, p && v }); - return true; - } - - if (cmd == "elif") { - std::string expr = trim_value(iss); - - if (cond.empty()) { - throw std::runtime_error("#elif without #if"); - } - - Cond & c = cond.back(); - if (!c.parent_active) { - c.active = false; - return true; - } - - if (c.taken) { - c.active = false; - return true; - } - - std::unordered_set visiting; - ExprParser ep(expr, macros, visiting); - bool v = ep.parse() != 0; - c.active = v; - if (v) { - c.taken = true; - } - return true; - } - - if (cmd == "else") { - if (cond.empty()) { - throw std::runtime_error("#else without #if"); - } - - Cond & c = cond.back(); - if (!c.parent_active) { - c.active = false; - return true; - } - if (c.taken) { - c.active = false; - } else { - c.active = true; - c.taken = true; - } - return true; - } - - if (cmd == "endif") { - if (cond.empty()) { - throw std::runtime_error("#endif without #if"); - } - cond.pop_back(); - return true; - } - - // Unknown directive - throw std::runtime_error("Unknown directive: #" + cmd); - } -}; - -} // namespace pre_wgsl - -#endif // PRE_WGSL_HPP diff --git a/ggml/src/ggml-zdnn/CMakeLists.txt b/ggml/src/ggml-zdnn/CMakeLists.txt index 0a723ce4de28..8254ecf9b28a 100644 --- a/ggml/src/ggml-zdnn/CMakeLists.txt +++ b/ggml/src/ggml-zdnn/CMakeLists.txt @@ -25,8 +25,9 @@ else() message(FATAL_ERROR "zdnn: library not found, please set ZDNN_ROOT to the proper path if necessary") endif() -file(GLOB GGML_SOURCES_ZDNN "*.c" "*.cpp") -file(GLOB GGML_HEADERS_ZDNN "*.h" "*.hpp") +set(GGML_SOURCES_ZDNN ggml-zdnn.cpp.inc) +set(GGML_HEADERS_ZDNN) +set_source_files_properties(${GGML_SOURCES_ZDNN} PROPERTIES LANGUAGE CXX) ggml_add_backend_library(ggml-zdnn ${GGML_HEADERS_ZDNN} ${GGML_SOURCES_ZDNN}) target_link_libraries(ggml-zdnn PRIVATE ${ZDNN_LIB}) diff --git a/ggml/src/ggml-zdnn/common.hpp b/ggml/src/ggml-zdnn/common.hpp deleted file mode 100644 index 2462ded55b7f..000000000000 --- a/ggml/src/ggml-zdnn/common.hpp +++ /dev/null @@ -1,59 +0,0 @@ -#ifndef GGML_ZDNN_COMMON_HPP -#define GGML_ZDNN_COMMON_HPP - -#include "ggml.h" -#include "ggml-impl.h" - -#include "zdnn.h" - -#include -#include - -#define GGML_ZDNN_NAME "zDNN" -#define GGML_ZDNN_VERSION ZDNN_VERNUM - -#define ZDNN_CHECK(stmt) \ - do { \ - zdnn_status status = (stmt); \ - GGML_ASSERT(status == ZDNN_OK); \ - } while (0); - -struct ggml_backend_zdnn_device_context { - int zdnn_device; - int zdnn_device_ref_count; - - bool has_parmblkformat_0; - bool has_parmblkformat_1; // checks for z17 - - size_t max_size; - - char name[128]; -}; - -struct ggml_backend_zdnn_context { - int device; - ggml_cgraph * gf; -}; - -struct ggml_backend_zdnn_buffer { - void * data; - ggml_backend_zdnn_buffer * extra; // for bias, etc. - size_t size; - - zdnn_tensor_desc pre_tfm_desc; - zdnn_tensor_desc tfm_desc; - zdnn_ztensor ztensor; - - char name[GGML_MAX_NAME]; -}; - -struct ggml_backend_zdnn_buffer_context { - void * all_data; - size_t all_size; - bool owned; - - int n_buffers; - std::vector> buffers; -}; - -#endif // GGML_ZDNN_COMMON_HPP diff --git a/ggml/src/ggml-zdnn/ggml-zdnn.cpp b/ggml/src/ggml-zdnn/ggml-zdnn.cpp.inc similarity index 69% rename from ggml/src/ggml-zdnn/ggml-zdnn.cpp rename to ggml/src/ggml-zdnn/ggml-zdnn.cpp.inc index 639b818d128e..d8f179308638 100644 --- a/ggml/src/ggml-zdnn/ggml-zdnn.cpp +++ b/ggml/src/ggml-zdnn/ggml-zdnn.cpp.inc @@ -1,11 +1,87 @@ -#include "ggml-zdnn.h" -#include "ggml-impl.h" -#include "ggml-backend-impl.h" +#include "ggml-zdnn.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-backend-impl-defs.inc" -#include "ggml-zdnn/common.hpp" -#include "ggml-zdnn/mmf.hpp" -#include "ggml-zdnn/utils.hpp" -#include "ggml.h" +/* Inlined common.h.incpp */ +#include "ggml.h.inc" +#include "ggml-impl-defs.inc" + +#include "zdnn.h" + +#include +#include + +#define GGML_ZDNN_NAME "zDNN" +#define GGML_ZDNN_VERSION ZDNN_VERNUM + +#define ZDNN_CHECK(stmt) \ + do { \ + zdnn_status status = (stmt); \ + GGML_ASSERT(status == ZDNN_OK); \ + } while (0); + +struct ggml_backend_zdnn_device_context { + int zdnn_device; + int zdnn_device_ref_count; + + bool has_parmblkformat_0; + bool has_parmblkformat_1; // checks for z17 + + size_t max_size; + + char name[128]; +}; + +struct ggml_backend_zdnn_context { + int device; + ggml_cgraph * gf; +}; + +struct ggml_backend_zdnn_buffer { + void * data; + ggml_backend_zdnn_buffer * extra; // for bias, etc. + size_t size; + + zdnn_tensor_desc pre_tfm_desc; + zdnn_tensor_desc tfm_desc; + zdnn_ztensor ztensor; + + char name[GGML_MAX_NAME]; +}; + +struct ggml_backend_zdnn_buffer_context { + void * all_data; + size_t all_size; + bool owned; + + int n_buffers; + std::vector> buffers; +}; + + +void ggml_zdnn_mul_mat_f( + const ggml_backend_zdnn_context * ctx, + const ggml_tensor * src0, + const ggml_tensor * src1, + ggml_tensor * dst); + + + +zdnn_data_types ggml_zdnn_type_mapping(ggml_type type); + +void ggml_zdnn_create_tensor(zdnn_tensor_desc & pre_tfm_desc, + zdnn_tensor_desc & tfm_desc, + zdnn_ztensor & ztensor, + const ggml_tensor * src, + const int64_t * ne, + const zdnn_data_layouts layout); + +void ggml_zdnn_load_tensor(zdnn_ztensor & ztensor, void * buffer); + +void ggml_zdnn_init_tensor(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor); + + +#include "ggml.h.inc" #include #include @@ -635,3 +711,189 @@ ggml_backend_reg_t ggml_backend_zdnn_reg(void) { } GGML_BACKEND_DL_IMPL(ggml_backend_zdnn_reg) + + +/* Inlined utils.cpp */ +#include "ggml.h.inc" + +zdnn_data_types ggml_zdnn_type_mapping(ggml_type type); + +void ggml_zdnn_create_tensor(zdnn_tensor_desc & pre_tfm_desc, + zdnn_tensor_desc & tfm_desc, + zdnn_ztensor & ztensor, + const ggml_tensor * src, + const int64_t * ne, + const zdnn_data_layouts layout); + +void ggml_zdnn_load_tensor(zdnn_ztensor & ztensor, void * buffer); + +void ggml_zdnn_init_tensor(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor); + + + +zdnn_data_types ggml_zdnn_type_mapping(ggml_type type) { + switch (type) { + case GGML_TYPE_F32: + return FP32; + case GGML_TYPE_F16: + return FP16; + case GGML_TYPE_BF16: + return BFLOAT; + case GGML_TYPE_Q8_0: + return INT8; + case GGML_TYPE_I8: + return INT8; + case GGML_TYPE_I32: + return INT32; + default: + GGML_ABORT("%s: fatal: unable to determine zTensor data type", + __func__); + break; + } +} + +void ggml_zdnn_create_tensor(zdnn_tensor_desc & pre_tfm_desc, + zdnn_tensor_desc & tfm_desc, + zdnn_ztensor & ztensor, + const ggml_tensor * src, + const int64_t * ne, + const zdnn_data_layouts layout) { + zdnn_init_pre_transformed_desc( + layout, + ggml_zdnn_type_mapping(src->type), + &pre_tfm_desc, + ne[3], ne[2], ne[1], ne[0] + ); + + ZDNN_CHECK(zdnn_generate_transformed_desc(&pre_tfm_desc, &tfm_desc)); + ZDNN_CHECK(zdnn_init_ztensor_with_malloc(&pre_tfm_desc, &tfm_desc, &ztensor)); +} + +void ggml_zdnn_load_tensor(zdnn_ztensor & ztensor, void * buffer) { + ZDNN_CHECK(zdnn_transform_ztensor(&ztensor, buffer)); +} + +void ggml_zdnn_init_tensor(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor) { + switch (tensor->op) { + case GGML_OP_MUL_MAT: + { + zdnn_init_pre_transformed_desc( + ZDNN_2D, + ggml_zdnn_type_mapping(tensor->type), + &buffer->pre_tfm_desc, + tensor->ne[1], tensor->ne[0] + ); + } break; + + default: + { + // For 4D tensors, GGML uses NCHW layout. However, because zDNN + // automatically transforms everything to NHWC, we will use it + // directly to avoid the performance penalty changing the + // layout and reshaping the tensor. + zdnn_init_pre_transformed_desc( + ZDNN_NHWC, + ggml_zdnn_type_mapping(tensor->type), + &buffer->pre_tfm_desc, + tensor->ne[3], tensor->ne[2], tensor->ne[1], tensor->ne[0] + ); + + // TODO: Consider adding a ggml check. + // TODO: If tensor = 4D, use ZDNN_NCHW by default. + // TODO: If tensor = 2D, use ZDNN_NHWC by default. + } break; + } + + ZDNN_CHECK(zdnn_generate_transformed_desc(&buffer->pre_tfm_desc, &buffer->tfm_desc)); + ZDNN_CHECK(zdnn_init_ztensor_with_malloc(&buffer->pre_tfm_desc, &buffer->tfm_desc, &buffer->ztensor)); +} + + +/* Inlined mmf.cpp */ +#include "ggml.h.inc" + +void ggml_zdnn_mul_mat_f( + const ggml_backend_zdnn_context * ctx, + const ggml_tensor * src0, + const ggml_tensor * src1, + ggml_tensor * dst); + + + +void ggml_zdnn_mul_mat_f( + const ggml_backend_zdnn_context * ctx, + const ggml_tensor * src0, + const ggml_tensor * src1, + ggml_tensor * dst) { + GGML_TENSOR_BINARY_OP_LOCALS; + + const enum ggml_type type = src0->type; + + GGML_ASSERT(ne0 == ne01); + GGML_ASSERT(ne1 == ne11); + GGML_ASSERT(ne2 == ne12); + GGML_ASSERT(ne3 == ne13); + + // we don't support permuted src0 or src1 + GGML_ASSERT(nb00 == ggml_type_size(type)); + GGML_ASSERT(nb10 == ggml_type_size(src1->type)); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + const ggml_tensor * weights = src0; + const ggml_tensor * inputs = src1; + ggml_tensor * output = dst; + + ggml_backend_zdnn_buffer * weights_extra = (ggml_backend_zdnn_buffer *)weights->extra; + ggml_backend_zdnn_buffer * inputs_extra = (ggml_backend_zdnn_buffer *)inputs->extra; + ggml_backend_zdnn_buffer * output_extra = (ggml_backend_zdnn_buffer *)output->extra; + ggml_backend_zdnn_buffer * bias_extra = (ggml_backend_zdnn_buffer *)output_extra->extra; + + const int64_t weights_rows = ne01; + const int64_t weights_cols = ne00; + const int64_t inputs_rows = ne11; + const int64_t inputs_cols = ne10; + + assert(inputs_cols == weights_cols); + + const int64_t output_rows = ne1; + const int64_t output_cols = ne0; + + // GGML_LOG_INFO("%s: tensor '%s' tensor dimensions: [%ld, %ld, %ld, %ld] pre_tfm_desc dimensions: [%ld, %ld, %ld, %ld]\n", + // __func__, weights_extra->name, + // weights->ne[3], weights->ne[2], weights->ne[1], weights->ne[0], + // weights_extra->pre_tfm_desc.dim1, + // weights_extra->pre_tfm_desc.dim2, + // weights_extra->pre_tfm_desc.dim3, + // weights_extra->pre_tfm_desc.dim4); + + // GGML_LOG_INFO("%s: tensor '%s' tensor dimensions: [%ld, %ld, %ld, %ld] pre_tfm_desc dimensions: [%ld, %ld, %ld, %ld]\n", + // __func__, inputs_extra->name, + // inputs->ne[3], inputs->ne[2], inputs->ne[1], inputs->ne[0], + // inputs_extra->pre_tfm_desc.dim1, + // inputs_extra->pre_tfm_desc.dim2, + // inputs_extra->pre_tfm_desc.dim3, + // inputs_extra->pre_tfm_desc.dim4); + + GGML_ASSERT(weights_extra->pre_tfm_desc.dim1 == weights->ne[0] && "weights_extra->pre_tfm_desc.dim1 must match weights->ne[0]"); + GGML_ASSERT(weights_extra->pre_tfm_desc.dim2 == weights->ne[1] && "weights_extra->pre_tfm_desc.dim2 must match weights->ne[1]"); + GGML_ASSERT(inputs_extra->pre_tfm_desc.dim1 == inputs->ne[0] && "inputs_extra->pre_tfm_desc.dim1 must match inputs->ne[0]"); + GGML_ASSERT(inputs_extra->pre_tfm_desc.dim2 == inputs->ne[1] && "inputs_extra->pre_tfm_desc.dim2 must match inputs->ne[1]"); + + ZDNN_CHECK(zdnn_matmul_transpose_op(&inputs_extra->ztensor, &weights_extra->ztensor, &bias_extra->ztensor, + false, true, MATMUL_OP_ADDITION, &output_extra->ztensor)); + // TODO: Remove in the future as we are currently DLF16 -> FP32 then in the next op, FP32 -> DLF16 again. Inefficient. + ZDNN_CHECK(zdnn_transform_origtensor(&output_extra->ztensor, output->data)); + + GGML_UNUSED(ctx); + GGML_UNUSED(weights_rows); + GGML_UNUSED(weights_cols); + GGML_UNUSED(inputs_rows); + GGML_UNUSED(inputs_cols); + GGML_UNUSED(output_rows); + GGML_UNUSED(output_cols); +} diff --git a/ggml/src/ggml-zdnn/mmf.cpp b/ggml/src/ggml-zdnn/mmf.cpp deleted file mode 100644 index 3ac9cf3c931e..000000000000 --- a/ggml/src/ggml-zdnn/mmf.cpp +++ /dev/null @@ -1,80 +0,0 @@ -#include "ggml.h" -#include "mmf.hpp" - -void ggml_zdnn_mul_mat_f( - const ggml_backend_zdnn_context * ctx, - const ggml_tensor * src0, - const ggml_tensor * src1, - ggml_tensor * dst) { - GGML_TENSOR_BINARY_OP_LOCALS; - - const enum ggml_type type = src0->type; - - GGML_ASSERT(ne0 == ne01); - GGML_ASSERT(ne1 == ne11); - GGML_ASSERT(ne2 == ne12); - GGML_ASSERT(ne3 == ne13); - - // we don't support permuted src0 or src1 - GGML_ASSERT(nb00 == ggml_type_size(type)); - GGML_ASSERT(nb10 == ggml_type_size(src1->type)); - - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); - - const ggml_tensor * weights = src0; - const ggml_tensor * inputs = src1; - ggml_tensor * output = dst; - - ggml_backend_zdnn_buffer * weights_extra = (ggml_backend_zdnn_buffer *)weights->extra; - ggml_backend_zdnn_buffer * inputs_extra = (ggml_backend_zdnn_buffer *)inputs->extra; - ggml_backend_zdnn_buffer * output_extra = (ggml_backend_zdnn_buffer *)output->extra; - ggml_backend_zdnn_buffer * bias_extra = (ggml_backend_zdnn_buffer *)output_extra->extra; - - const int64_t weights_rows = ne01; - const int64_t weights_cols = ne00; - const int64_t inputs_rows = ne11; - const int64_t inputs_cols = ne10; - - assert(inputs_cols == weights_cols); - - const int64_t output_rows = ne1; - const int64_t output_cols = ne0; - - // GGML_LOG_INFO("%s: tensor '%s' tensor dimensions: [%ld, %ld, %ld, %ld] pre_tfm_desc dimensions: [%ld, %ld, %ld, %ld]\n", - // __func__, weights_extra->name, - // weights->ne[3], weights->ne[2], weights->ne[1], weights->ne[0], - // weights_extra->pre_tfm_desc.dim1, - // weights_extra->pre_tfm_desc.dim2, - // weights_extra->pre_tfm_desc.dim3, - // weights_extra->pre_tfm_desc.dim4); - - // GGML_LOG_INFO("%s: tensor '%s' tensor dimensions: [%ld, %ld, %ld, %ld] pre_tfm_desc dimensions: [%ld, %ld, %ld, %ld]\n", - // __func__, inputs_extra->name, - // inputs->ne[3], inputs->ne[2], inputs->ne[1], inputs->ne[0], - // inputs_extra->pre_tfm_desc.dim1, - // inputs_extra->pre_tfm_desc.dim2, - // inputs_extra->pre_tfm_desc.dim3, - // inputs_extra->pre_tfm_desc.dim4); - - GGML_ASSERT(weights_extra->pre_tfm_desc.dim1 == weights->ne[0] && "weights_extra->pre_tfm_desc.dim1 must match weights->ne[0]"); - GGML_ASSERT(weights_extra->pre_tfm_desc.dim2 == weights->ne[1] && "weights_extra->pre_tfm_desc.dim2 must match weights->ne[1]"); - GGML_ASSERT(inputs_extra->pre_tfm_desc.dim1 == inputs->ne[0] && "inputs_extra->pre_tfm_desc.dim1 must match inputs->ne[0]"); - GGML_ASSERT(inputs_extra->pre_tfm_desc.dim2 == inputs->ne[1] && "inputs_extra->pre_tfm_desc.dim2 must match inputs->ne[1]"); - - ZDNN_CHECK(zdnn_matmul_transpose_op(&inputs_extra->ztensor, &weights_extra->ztensor, &bias_extra->ztensor, - false, true, MATMUL_OP_ADDITION, &output_extra->ztensor)); - // TODO: Remove in the future as we are currently DLF16 -> FP32 then in the next op, FP32 -> DLF16 again. Inefficient. - ZDNN_CHECK(zdnn_transform_origtensor(&output_extra->ztensor, output->data)); - - GGML_UNUSED(ctx); - GGML_UNUSED(weights_rows); - GGML_UNUSED(weights_cols); - GGML_UNUSED(inputs_rows); - GGML_UNUSED(inputs_cols); - GGML_UNUSED(output_rows); - GGML_UNUSED(output_cols); -} diff --git a/ggml/src/ggml-zdnn/mmf.hpp b/ggml/src/ggml-zdnn/mmf.hpp deleted file mode 100644 index a12f1b8f8a0e..000000000000 --- a/ggml/src/ggml-zdnn/mmf.hpp +++ /dev/null @@ -1,12 +0,0 @@ -#ifndef GGML_ZDNN_MMF_HPP -#define GGML_ZDNN_MMF_HPP - -#include "common.hpp" - -void ggml_zdnn_mul_mat_f( - const ggml_backend_zdnn_context * ctx, - const ggml_tensor * src0, - const ggml_tensor * src1, - ggml_tensor * dst); - -#endif // GGML_ZDNN_MMF_HPP diff --git a/ggml/src/ggml-zdnn/utils.cpp b/ggml/src/ggml-zdnn/utils.cpp deleted file mode 100644 index 2977cb0fe3bd..000000000000 --- a/ggml/src/ggml-zdnn/utils.cpp +++ /dev/null @@ -1,79 +0,0 @@ -#include "ggml.h" -#include "utils.hpp" - -zdnn_data_types ggml_zdnn_type_mapping(ggml_type type) { - switch (type) { - case GGML_TYPE_F32: - return FP32; - case GGML_TYPE_F16: - return FP16; - case GGML_TYPE_BF16: - return BFLOAT; - case GGML_TYPE_Q8_0: - return INT8; - case GGML_TYPE_I8: - return INT8; - case GGML_TYPE_I32: - return INT32; - default: - GGML_ABORT("%s: fatal: unable to determine zTensor data type", - __func__); - break; - } -} - -void ggml_zdnn_create_tensor(zdnn_tensor_desc & pre_tfm_desc, - zdnn_tensor_desc & tfm_desc, - zdnn_ztensor & ztensor, - const ggml_tensor * src, - const int64_t * ne, - const zdnn_data_layouts layout) { - zdnn_init_pre_transformed_desc( - layout, - ggml_zdnn_type_mapping(src->type), - &pre_tfm_desc, - ne[3], ne[2], ne[1], ne[0] - ); - - ZDNN_CHECK(zdnn_generate_transformed_desc(&pre_tfm_desc, &tfm_desc)); - ZDNN_CHECK(zdnn_init_ztensor_with_malloc(&pre_tfm_desc, &tfm_desc, &ztensor)); -} - -void ggml_zdnn_load_tensor(zdnn_ztensor & ztensor, void * buffer) { - ZDNN_CHECK(zdnn_transform_ztensor(&ztensor, buffer)); -} - -void ggml_zdnn_init_tensor(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor) { - switch (tensor->op) { - case GGML_OP_MUL_MAT: - { - zdnn_init_pre_transformed_desc( - ZDNN_2D, - ggml_zdnn_type_mapping(tensor->type), - &buffer->pre_tfm_desc, - tensor->ne[1], tensor->ne[0] - ); - } break; - - default: - { - // For 4D tensors, GGML uses NCHW layout. However, because zDNN - // automatically transforms everything to NHWC, we will use it - // directly to avoid the performance penalty changing the - // layout and reshaping the tensor. - zdnn_init_pre_transformed_desc( - ZDNN_NHWC, - ggml_zdnn_type_mapping(tensor->type), - &buffer->pre_tfm_desc, - tensor->ne[3], tensor->ne[2], tensor->ne[1], tensor->ne[0] - ); - - // TODO: Consider adding a ggml check. - // TODO: If tensor = 4D, use ZDNN_NCHW by default. - // TODO: If tensor = 2D, use ZDNN_NHWC by default. - } break; - } - - ZDNN_CHECK(zdnn_generate_transformed_desc(&buffer->pre_tfm_desc, &buffer->tfm_desc)); - ZDNN_CHECK(zdnn_init_ztensor_with_malloc(&buffer->pre_tfm_desc, &buffer->tfm_desc, &buffer->ztensor)); -} diff --git a/ggml/src/ggml-zdnn/utils.hpp b/ggml/src/ggml-zdnn/utils.hpp deleted file mode 100644 index c1e2028edbca..000000000000 --- a/ggml/src/ggml-zdnn/utils.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#ifndef GGML_ZDNN_UTILITIES_HPP -#define GGML_ZDNN_UTILITIES_HPP - -#include "common.hpp" - -zdnn_data_types ggml_zdnn_type_mapping(ggml_type type); - -void ggml_zdnn_create_tensor(zdnn_tensor_desc & pre_tfm_desc, - zdnn_tensor_desc & tfm_desc, - zdnn_ztensor & ztensor, - const ggml_tensor * src, - const int64_t * ne, - const zdnn_data_layouts layout); - -void ggml_zdnn_load_tensor(zdnn_ztensor & ztensor, void * buffer); - -void ggml_zdnn_init_tensor(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor); - -#endif // GGML_ZDNN_UTILITIES_HPP diff --git a/ggml/src/ggml-zendnn/CMakeLists.txt b/ggml/src/ggml-zendnn/CMakeLists.txt index 4f321a252570..b25d5bb7d36f 100644 --- a/ggml/src/ggml-zendnn/CMakeLists.txt +++ b/ggml/src/ggml-zendnn/CMakeLists.txt @@ -1,5 +1,6 @@ +set_source_files_properties(ggml-zendnn.cpp.inc PROPERTIES LANGUAGE CXX) ggml_add_backend_library(ggml-zendnn - ggml-zendnn.cpp) + ggml-zendnn.cpp.inc) if (NOT DEFINED ZENDNN_ROOT OR ZENDNN_ROOT STREQUAL "") set(ZENDNN_ROOT "$ENV{ZENDNN_ROOT}") diff --git a/ggml/src/ggml-zendnn/ggml-zendnn.cpp b/ggml/src/ggml-zendnn/ggml-zendnn.cpp.inc similarity index 99% rename from ggml/src/ggml-zendnn/ggml-zendnn.cpp rename to ggml/src/ggml-zendnn/ggml-zendnn.cpp.inc index 2b82c7c1dbb8..0105aabdf08f 100644 --- a/ggml/src/ggml-zendnn/ggml-zendnn.cpp +++ b/ggml/src/ggml-zendnn/ggml-zendnn.cpp.inc @@ -1,7 +1,7 @@ -#include "ggml-zendnn.h" +#include "ggml-zendnn.h.inc" -#include "ggml-backend-impl.h" -#include "ggml-impl.h" +#include "ggml-backend-impl-defs.inc" +#include "ggml-impl-defs.inc" #include "zendnnl.hpp" #include diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c deleted file mode 100644 index 81343eeb14ce..000000000000 --- a/ggml/src/ggml.c +++ /dev/null @@ -1,7765 +0,0 @@ -#define _CRT_SECURE_NO_DEPRECATE // Disables "unsafe" warnings on Windows -#define _USE_MATH_DEFINES // For M_PI on MSVC - -#include "ggml-backend.h" -#include "ggml-impl.h" -#include "ggml-threading.h" -#include "ggml-cpu.h" -#include "ggml.h" - -// FIXME: required here for quantization functions -#include "ggml-quants.h" - -#ifdef GGML_USE_CPU_HBM -#include -#endif - -#if defined(_MSC_VER) || defined(__MINGW32__) -#include // using malloc.h with MSC/MINGW -#elif !defined(__FreeBSD__) && !defined(__NetBSD__) && !defined(__OpenBSD__) -#include -#endif - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#if defined(__gnu_linux__) -#include -#endif - -#if defined(__APPLE__) -#include -#include -#include -#endif - -#if defined(_WIN32) -#define WIN32_LEAN_AND_MEAN -#ifndef NOMINMAX - #define NOMINMAX -#endif -#include -#endif - -#define UNUSED GGML_UNUSED - -uint64_t ggml_graph_next_uid(void) { -#ifdef _MSC_VER -#if defined(_WIN32) - static volatile LONG counter = 1; - return (uint64_t) InterlockedIncrement(&counter) - 1; -#else - static volatile long long counter = 1; - return (uint64_t) _InterlockedIncrement64(&counter) - 1; -#endif -#else - static uint64_t counter = 1; - return __atomic_fetch_add(&counter, 1, __ATOMIC_RELAXED); -#endif -} - -// Needed for ggml_fp32_to_bf16_row() -#if defined(__AVX512BF16__) -#if defined(_MSC_VER) -#define m512i(p) p -#else -#include -#define m512i(p) (__m512i)(p) -#endif // defined(_MSC_VER) -#endif // defined(__AVX512BF16__) - -#if defined(__linux__) || \ - defined(__FreeBSD__) || defined(__NetBSD__) || defined(__OpenBSD__) || \ - (defined(__APPLE__) && !TARGET_OS_TV && !TARGET_OS_WATCH) - -#include -#include -#include -#include -#if defined(__linux__) -#include -#endif - -#if defined(__ANDROID__) -#include -#include -#include - -struct backtrace_state { - void ** current; - void ** end; -}; - -static _Unwind_Reason_Code unwind_callback(struct _Unwind_Context* context, void* arg) { - struct backtrace_state * state = (struct backtrace_state *)arg; - uintptr_t pc = _Unwind_GetIP(context); - if (pc) { - if (state->current == state->end) { - return _URC_END_OF_STACK; - } else { - *state->current++ = (void*)pc; - } - } - return _URC_NO_REASON; -} - -static void ggml_print_backtrace_symbols(void) { - const int max = 100; - void* buffer[max]; - - struct backtrace_state state = {buffer, buffer + max}; - _Unwind_Backtrace(unwind_callback, &state); - - int count = state.current - buffer; - - for (int idx = 0; idx < count; ++idx) { - const void * addr = buffer[idx]; - const char * symbol = ""; - - Dl_info info; - if (dladdr(addr, &info) && info.dli_sname) { - symbol = info.dli_sname; - } - - fprintf(stderr, "%d: %p %s\n", idx, addr, symbol); - } -} -#elif defined(__linux__) && defined(__GLIBC__) -#include -static void ggml_print_backtrace_symbols(void) { - void * trace[100]; - int nptrs = backtrace(trace, sizeof(trace)/sizeof(trace[0])); - backtrace_symbols_fd(trace, nptrs, STDERR_FILENO); -} -#elif defined(__APPLE__) -#include -static void ggml_print_backtrace_symbols(void) { - void * trace[100]; - int nptrs = backtrace(trace, sizeof(trace)/sizeof(trace[0])); - backtrace_symbols_fd(trace, nptrs, STDERR_FILENO); -} -#else -static void ggml_print_backtrace_symbols(void) { - // platform not supported -} -#endif - -void ggml_print_backtrace(void) { - const char * GGML_NO_BACKTRACE = getenv("GGML_NO_BACKTRACE"); - if (GGML_NO_BACKTRACE) { - return; - } -#if defined(__APPLE__) - // On macOS, fork+debugger attachment is problematic due to: - // 1. libdispatch "poisons" forked child processes - // 2. lldb has issues attaching to parent from forked child - // Use simple backtrace() instead to avoid Terminal.app crashes - const char * GGML_BACKTRACE_LLDB = getenv("GGML_BACKTRACE_LLDB"); - if (!GGML_BACKTRACE_LLDB) { - fprintf(stderr, "WARNING: Using native backtrace. Set GGML_BACKTRACE_LLDB for more info.\n"); - fprintf(stderr, "WARNING: GGML_BACKTRACE_LLDB may cause native MacOS Terminal.app to crash.\n"); - fprintf(stderr, "See: https://github.com/ggml-org/llama.cpp/pull/17869\n"); - ggml_print_backtrace_symbols(); - return; - } -#endif -#if defined(__linux__) - FILE * f = fopen("/proc/self/status", "r"); - size_t size = 0; - char * line = NULL; - ssize_t length = 0; - while ((length = getline(&line, &size, f)) > 0) { - if (!strncmp(line, "TracerPid:", sizeof("TracerPid:") - 1) && - (length != sizeof("TracerPid:\t0\n") - 1 || line[length - 2] != '0')) { - // Already being debugged, and the breakpoint is the later abort() - free(line); - fclose(f); - return; - } - } - free(line); - fclose(f); - int lock[2] = { -1, -1 }; - (void) !pipe(lock); // Don't start gdb until after PR_SET_PTRACER -#endif - const int parent_pid = getpid(); - const int child_pid = fork(); - if (child_pid < 0) { // error -#if defined(__linux__) - close(lock[1]); - close(lock[0]); -#endif - return; - } else if (child_pid == 0) { // child - char attach[32]; - snprintf(attach, sizeof(attach), "attach %d", parent_pid); -#if defined(__linux__) - close(lock[1]); - (void) !read(lock[0], lock, 1); - close(lock[0]); -#endif - // try gdb - execlp("gdb", "gdb", "--batch", - "-ex", "set style enabled on", - "-ex", attach, - "-ex", "bt -frame-info source-and-location", - "-ex", "detach", - "-ex", "quit", - (char *) NULL); - // try lldb - execlp("lldb", "lldb", "--batch", - "-o", "bt", - "-o", "quit", - "-p", &attach[sizeof("attach ") - 1], - (char *) NULL); - // gdb failed, fallback to backtrace_symbols - ggml_print_backtrace_symbols(); - _Exit(0); - } else { // parent -#if defined(__linux__) - prctl(PR_SET_PTRACER, child_pid); - close(lock[1]); - close(lock[0]); -#endif - waitpid(child_pid, NULL, 0); - } -} -#else -void ggml_print_backtrace(void) { - // platform not supported -} -#endif - -static ggml_abort_callback_t g_abort_callback = NULL; - -// Set the abort callback (passing null will restore original abort functionality: printing a message to stdout) -GGML_API ggml_abort_callback_t ggml_set_abort_callback(ggml_abort_callback_t callback) { - ggml_abort_callback_t ret_val = g_abort_callback; - g_abort_callback = callback; - return ret_val; -} - -void ggml_abort(const char * file, int line, const char * fmt, ...) { - fflush(stdout); - - char message[2048]; - int offset = snprintf(message, sizeof(message), "%s:%d: ", file, line); - - va_list args; - va_start(args, fmt); - vsnprintf(message + offset, sizeof(message) - offset, fmt, args); - va_end(args); - - if (g_abort_callback) { - g_abort_callback(message); - } else { - // default: print error and backtrace to stderr - fprintf(stderr, "%s\n", message); - ggml_print_backtrace(); - } - - abort(); -} - -// ggml_print_backtrace is registered with std::set_terminate by ggml.cpp - -// -// logging -// - -struct ggml_logger_state { - ggml_log_callback log_callback; - void * log_callback_user_data; -}; -static struct ggml_logger_state g_logger_state = {ggml_log_callback_default, NULL}; - -static void ggml_log_internal_v(enum ggml_log_level level, const char * format, va_list args) { - if (format == NULL) { - return; - } - va_list args_copy; - va_copy(args_copy, args); - char buffer[128]; - int len = vsnprintf(buffer, 128, format, args); - if (len < 128) { - g_logger_state.log_callback(level, buffer, g_logger_state.log_callback_user_data); - } else { - char * buffer2 = (char *) calloc(len + 1, sizeof(char)); - vsnprintf(buffer2, len + 1, format, args_copy); - buffer2[len] = 0; - g_logger_state.log_callback(level, buffer2, g_logger_state.log_callback_user_data); - free(buffer2); - } - va_end(args_copy); -} - -void ggml_log_internal(enum ggml_log_level level, const char * format, ...) { - va_list args; - va_start(args, format); - ggml_log_internal_v(level, format, args); - va_end(args); -} - -void ggml_log_callback_default(enum ggml_log_level level, const char * text, void * user_data) { - (void) level; - (void) user_data; - fputs(text, stderr); - fflush(stderr); -} - -// -// end of logging block -// - -#ifdef GGML_USE_ACCELERATE -// uncomment to use vDSP for soft max computation -// note: not sure if it is actually faster -//#define GGML_SOFT_MAX_ACCELERATE -#endif - - -void * ggml_aligned_malloc(size_t size) { -#if defined(__s390x__) - const int alignment = 256; -#else - const int alignment = 64; -#endif - -#if defined(_MSC_VER) || defined(__MINGW32__) - return _aligned_malloc(size, alignment); -#else - if (size == 0) { - GGML_LOG_WARN("Behavior may be unexpected when allocating 0 bytes for ggml_aligned_malloc!\n"); - return NULL; - } - void * aligned_memory = NULL; - #ifdef GGML_USE_CPU_HBM - int result = hbw_posix_memalign(&aligned_memory, alignment, size); - #elif TARGET_OS_OSX - GGML_UNUSED(alignment); - kern_return_t alloc_status = vm_allocate((vm_map_t) mach_task_self(), (vm_address_t *) &aligned_memory, size, VM_FLAGS_ANYWHERE); - int result = EFAULT; - switch (alloc_status) { - case KERN_SUCCESS: - result = 0; - break; - case KERN_INVALID_ADDRESS: - result = EINVAL; - break; - case KERN_NO_SPACE: - result = ENOMEM; - break; - default: - result = EFAULT; - break; - } - #else - int result = posix_memalign(&aligned_memory, alignment, size); - #endif - if (result != 0) { - // Handle allocation failure - const char *error_desc = "unknown allocation error"; - switch (result) { - case EINVAL: - error_desc = "invalid alignment value"; - break; - case ENOMEM: - error_desc = "insufficient memory"; - break; - } - GGML_LOG_ERROR("%s: %s (attempted to allocate %6.2f MB)\n", __func__, error_desc, size/(1024.0*1024.0)); - return NULL; - } - return aligned_memory; -#endif -} - -void ggml_aligned_free(void * ptr, size_t size) { - GGML_UNUSED(size); -#if defined(_MSC_VER) || defined(__MINGW32__) - _aligned_free(ptr); -#elif GGML_USE_CPU_HBM - if (ptr != NULL) { - hbw_free(ptr); - } -#elif TARGET_OS_OSX - if (ptr != NULL) { - vm_deallocate((vm_map_t)mach_task_self(), (vm_address_t)ptr, size); - } -#else - free(ptr); -#endif -} - - -inline static void * ggml_malloc(size_t size) { - if (size == 0) { - GGML_LOG_WARN("Behavior may be unexpected when allocating 0 bytes for ggml_malloc!\n"); - return NULL; - } - void * result = malloc(size); - if (result == NULL) { - GGML_LOG_ERROR("%s: failed to allocate %6.2f MB\n", __func__, size/(1024.0*1024.0)); - GGML_ABORT("fatal error"); - } - return result; -} - -// calloc -inline static void * ggml_calloc(size_t num, size_t size) { - if (num == 0 || size == 0) { - GGML_LOG_WARN("Behavior may be unexpected when allocating 0 bytes for ggml_calloc!\n"); - return NULL; - } - void * result = calloc(num, size); - if (result == NULL) { - GGML_LOG_ERROR("%s: failed to allocate %6.2f MB\n", __func__, size/(1024.0*1024.0)); - GGML_ABORT("fatal error"); - } - return result; -} - -#define GGML_MALLOC(size) ggml_malloc(size) -#define GGML_CALLOC(num, size) ggml_calloc(num, size) - -#define GGML_FREE(ptr) free(ptr) - -const char * ggml_status_to_string(enum ggml_status status) { - switch (status) { - case GGML_STATUS_ALLOC_FAILED: return "GGML status: error (failed to allocate memory)"; - case GGML_STATUS_FAILED: return "GGML status: error (operation failed)"; - case GGML_STATUS_SUCCESS: return "GGML status: success"; - case GGML_STATUS_ABORTED: return "GGML status: warning (operation aborted)"; - } - - return "GGML status: unknown"; -} - -float ggml_fp16_to_fp32(ggml_fp16_t x) { -#define ggml_fp16_to_fp32 do_not_use__ggml_fp16_to_fp32__in_ggml - return GGML_FP16_TO_FP32(x); -} - -ggml_fp16_t ggml_fp32_to_fp16(float x) { -#define ggml_fp32_to_fp16 do_not_use__ggml_fp32_to_fp16__in_ggml - return GGML_FP32_TO_FP16(x); -} - -float ggml_bf16_to_fp32(ggml_bf16_t x) { -#define ggml_bf16_to_fp32 do_not_use__ggml_bf16_to_fp32__in_ggml - return GGML_BF16_TO_FP32(x); // it just left shifts -} - -ggml_bf16_t ggml_fp32_to_bf16(float x) { -#define ggml_fp32_to_bf16 do_not_use__ggml_fp32_to_bf16__in_ggml - return GGML_FP32_TO_BF16(x); -} - -void ggml_fp16_to_fp32_row(const ggml_fp16_t * x, float * y, int64_t n) { - for (int64_t i = 0; i < n; i++) { - y[i] = GGML_FP16_TO_FP32(x[i]); - } -} - -void ggml_fp32_to_fp16_row(const float * x, ggml_fp16_t * y, int64_t n) { - int i = 0; - for (; i < n; ++i) { - y[i] = GGML_FP32_TO_FP16(x[i]); - } -} - -void ggml_bf16_to_fp32_row(const ggml_bf16_t * x, float * y, int64_t n) { - int i = 0; - for (; i < n; ++i) { - y[i] = GGML_BF16_TO_FP32(x[i]); - } -} - -void ggml_fp32_to_bf16_row_ref(const float * x, ggml_bf16_t * y, int64_t n) { - for (int i = 0; i < n; i++) { - y[i] = ggml_compute_fp32_to_bf16(x[i]); - } -} - -void ggml_fp32_to_bf16_row(const float * x, ggml_bf16_t * y, int64_t n) { - int i = 0; -#if defined(__AVX512BF16__) - // subnormals are flushed to zero on this platform - for (; i + 32 <= n; i += 32) { - _mm512_storeu_si512( - (__m512i *)(y + i), - m512i(_mm512_cvtne2ps_pbh(_mm512_loadu_ps(x + i + 16), - _mm512_loadu_ps(x + i)))); - } -#endif - for (; i < n; i++) { - y[i] = GGML_FP32_TO_BF16(x[i]); - } -} - -bool ggml_guid_matches(ggml_guid_t guid_a, ggml_guid_t guid_b) { - return memcmp(guid_a, guid_b, sizeof(ggml_guid)) == 0; -} - -const char * ggml_version(void) { - return GGML_VERSION; -} - -const char * ggml_commit(void) { - return GGML_COMMIT; -} - -// -// timing -// - -#if defined(_MSC_VER) || defined(__MINGW32__) -static int64_t timer_freq, timer_start; -void ggml_time_init(void) { - LARGE_INTEGER t; - QueryPerformanceFrequency(&t); - timer_freq = t.QuadPart; - - // The multiplication by 1000 or 1000000 below can cause an overflow if timer_freq - // and the uptime is high enough. - // We subtract the program start time to reduce the likelihood of that happening. - QueryPerformanceCounter(&t); - timer_start = t.QuadPart; -} -int64_t ggml_time_ms(void) { - LARGE_INTEGER t; - QueryPerformanceCounter(&t); - return ((t.QuadPart-timer_start) * 1000) / timer_freq; -} -int64_t ggml_time_us(void) { - LARGE_INTEGER t; - QueryPerformanceCounter(&t); - return ((t.QuadPart-timer_start) * 1000000) / timer_freq; -} -#else -void ggml_time_init(void) {} -int64_t ggml_time_ms(void) { - struct timespec ts; - clock_gettime(CLOCK_MONOTONIC, &ts); - return (int64_t)ts.tv_sec*1000 + (int64_t)ts.tv_nsec/1000000; -} - -int64_t ggml_time_us(void) { - struct timespec ts; - clock_gettime(CLOCK_MONOTONIC, &ts); - return (int64_t)ts.tv_sec*1000000 + (int64_t)ts.tv_nsec/1000; -} -#endif - -int64_t ggml_cycles(void) { - return clock(); -} - -int64_t ggml_cycles_per_ms(void) { - return CLOCKS_PER_SEC/1000; -} - -// -// cross-platform UTF-8 file paths -// - -#ifdef _WIN32 -static wchar_t * ggml_mbstowcs(const char * mbs) { - int wlen = MultiByteToWideChar(CP_UTF8, 0, mbs, -1, NULL, 0); - if (!wlen) { - errno = EINVAL; - return NULL; - } - - wchar_t * wbuf = GGML_MALLOC(wlen * sizeof(wchar_t)); - wlen = MultiByteToWideChar(CP_UTF8, 0, mbs, -1, wbuf, wlen); - if (!wlen) { - GGML_FREE(wbuf); - errno = EINVAL; - return NULL; - } - - return wbuf; -} -#endif - -FILE * ggml_fopen(const char * fname, const char * mode) { -#ifdef _WIN32 - FILE * file = NULL; - - // convert fname (UTF-8) - wchar_t * wfname = ggml_mbstowcs(fname); - if (wfname) { - // convert mode (ANSI) - wchar_t * wmode = GGML_MALLOC((strlen(mode) + 1) * sizeof(wchar_t)); - wchar_t * wmode_p = wmode; - do { - *wmode_p++ = (wchar_t)*mode; - } while (*mode++); - - // open file - file = _wfopen(wfname, wmode); - - GGML_FREE(wfname); - GGML_FREE(wmode); - } - - return file; -#else - return fopen(fname, mode); -#endif - -} - -static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { - [GGML_TYPE_I8] = { - .type_name = "i8", - .blck_size = 1, - .type_size = sizeof(int8_t), - .is_quantized = false, - }, - [GGML_TYPE_I16] = { - .type_name = "i16", - .blck_size = 1, - .type_size = sizeof(int16_t), - .is_quantized = false, - }, - [GGML_TYPE_I32] = { - .type_name = "i32", - .blck_size = 1, - .type_size = sizeof(int32_t), - .is_quantized = false, - }, - [GGML_TYPE_I64] = { - .type_name = "i64", - .blck_size = 1, - .type_size = sizeof(int64_t), - .is_quantized = false, - }, - [GGML_TYPE_F64] = { - .type_name = "f64", - .blck_size = 1, - .type_size = sizeof(double), - .is_quantized = false, - }, - [GGML_TYPE_F32] = { - .type_name = "f32", - .blck_size = 1, - .type_size = sizeof(float), - .is_quantized = false, - }, - [GGML_TYPE_F16] = { - .type_name = "f16", - .blck_size = 1, - .type_size = sizeof(ggml_fp16_t), - .is_quantized = false, - .to_float = (ggml_to_float_t) ggml_fp16_to_fp32_row, - .from_float_ref = (ggml_from_float_t) ggml_fp32_to_fp16_row, - }, - [GGML_TYPE_Q1_0] = { - .type_name = "q1_0", - .blck_size = QK1_0, - .type_size = sizeof(block_q1_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q1_0, - .from_float_ref = (ggml_from_float_t) quantize_row_q1_0_ref, - }, - [GGML_TYPE_Q4_0] = { - .type_name = "q4_0", - .blck_size = QK4_0, - .type_size = sizeof(block_q4_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q4_0, - .from_float_ref = (ggml_from_float_t) quantize_row_q4_0_ref, - }, - [GGML_TYPE_Q4_1] = { - .type_name = "q4_1", - .blck_size = QK4_1, - .type_size = sizeof(block_q4_1), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q4_1, - .from_float_ref = (ggml_from_float_t) quantize_row_q4_1_ref, - }, - [4] = { // GGML_TYPE_Q4_2 - .type_name = "DEPRECATED", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [5] = { // GGML_TYPE_Q4_3 - .type_name = "DEPRECATED", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [GGML_TYPE_Q5_0] = { - .type_name = "q5_0", - .blck_size = QK5_0, - .type_size = sizeof(block_q5_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q5_0, - .from_float_ref = (ggml_from_float_t) quantize_row_q5_0_ref, - }, - [GGML_TYPE_Q5_1] = { - .type_name = "q5_1", - .blck_size = QK5_1, - .type_size = sizeof(block_q5_1), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q5_1, - .from_float_ref = (ggml_from_float_t) quantize_row_q5_1_ref, - }, - [GGML_TYPE_Q8_0] = { - .type_name = "q8_0", - .blck_size = QK8_0, - .type_size = sizeof(block_q8_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q8_0, - .from_float_ref = (ggml_from_float_t) quantize_row_q8_0_ref, - }, - [GGML_TYPE_Q8_1] = { - .type_name = "q8_1", - .blck_size = QK8_1, - .type_size = sizeof(block_q8_1), - .is_quantized = true, - .from_float_ref = (ggml_from_float_t) quantize_row_q8_1_ref, - }, - [GGML_TYPE_MXFP4] = { - .type_name = "mxfp4", - .blck_size = QK_MXFP4, - .type_size = sizeof(block_mxfp4), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_mxfp4, - .from_float_ref = (ggml_from_float_t)quantize_row_mxfp4_ref, - }, - [GGML_TYPE_NVFP4] = { - .type_name = "nvfp4", - .blck_size = QK_NVFP4, - .type_size = sizeof(block_nvfp4), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_nvfp4, - .from_float_ref = (ggml_from_float_t)quantize_row_nvfp4_ref, - }, - [GGML_TYPE_Q2_K] = { - .type_name = "q2_K", - .blck_size = QK_K, - .type_size = sizeof(block_q2_K), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q2_K, - .from_float_ref = (ggml_from_float_t) quantize_row_q2_K_ref, - }, - [GGML_TYPE_Q3_K] = { - .type_name = "q3_K", - .blck_size = QK_K, - .type_size = sizeof(block_q3_K), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q3_K, - .from_float_ref = (ggml_from_float_t) quantize_row_q3_K_ref, - }, - [GGML_TYPE_Q4_K] = { - .type_name = "q4_K", - .blck_size = QK_K, - .type_size = sizeof(block_q4_K), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q4_K, - .from_float_ref = (ggml_from_float_t) quantize_row_q4_K_ref, - }, - [GGML_TYPE_Q5_K] = { - .type_name = "q5_K", - .blck_size = QK_K, - .type_size = sizeof(block_q5_K), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q5_K, - .from_float_ref = (ggml_from_float_t) quantize_row_q5_K_ref, - }, - [GGML_TYPE_Q6_K] = { - .type_name = "q6_K", - .blck_size = QK_K, - .type_size = sizeof(block_q6_K), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q6_K, - .from_float_ref = (ggml_from_float_t) quantize_row_q6_K_ref, - }, - [GGML_TYPE_IQ2_XXS] = { - .type_name = "iq2_xxs", - .blck_size = QK_K, - .type_size = sizeof(block_iq2_xxs), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq2_xxs, - .from_float_ref = NULL, - }, - [GGML_TYPE_IQ2_XS] = { - .type_name = "iq2_xs", - .blck_size = QK_K, - .type_size = sizeof(block_iq2_xs), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq2_xs, - .from_float_ref = NULL, - }, - [GGML_TYPE_IQ3_XXS] = { - .type_name = "iq3_xxs", - .blck_size = QK_K, - .type_size = sizeof(block_iq3_xxs), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq3_xxs, - .from_float_ref = (ggml_from_float_t)quantize_row_iq3_xxs_ref, - }, - [GGML_TYPE_IQ3_S] = { - .type_name = "iq3_s", - .blck_size = QK_K, - .type_size = sizeof(block_iq3_s), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq3_s, - .from_float_ref = (ggml_from_float_t)quantize_row_iq3_s_ref, - }, - [GGML_TYPE_IQ2_S] = { - .type_name = "iq2_s", - .blck_size = QK_K, - .type_size = sizeof(block_iq2_s), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq2_s, - .from_float_ref = (ggml_from_float_t)quantize_row_iq2_s_ref, - }, - [GGML_TYPE_IQ1_S] = { - .type_name = "iq1_s", - .blck_size = QK_K, - .type_size = sizeof(block_iq1_s), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq1_s, - .from_float_ref = NULL, - }, - [GGML_TYPE_IQ1_M] = { - .type_name = "iq1_m", - .blck_size = QK_K, - .type_size = sizeof(block_iq1_m), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq1_m, - .from_float_ref = NULL, - }, - [GGML_TYPE_IQ4_NL] = { - .type_name = "iq4_nl", - .blck_size = QK4_NL, - .type_size = sizeof(block_iq4_nl), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq4_nl, - .from_float_ref = (ggml_from_float_t)quantize_row_iq4_nl_ref, - }, - [GGML_TYPE_IQ4_XS] = { - .type_name = "iq4_xs", - .blck_size = QK_K, - .type_size = sizeof(block_iq4_xs), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq4_xs, - .from_float_ref = (ggml_from_float_t)quantize_row_iq4_xs_ref, - }, - [GGML_TYPE_Q8_K] = { - .type_name = "q8_K", - .blck_size = QK_K, - .type_size = sizeof(block_q8_K), - .is_quantized = true, - }, - [GGML_TYPE_BF16] = { - .type_name = "bf16", - .blck_size = 1, - .type_size = sizeof(ggml_bf16_t), - .is_quantized = false, - .to_float = (ggml_to_float_t) ggml_bf16_to_fp32_row, - .from_float_ref = (ggml_from_float_t) ggml_fp32_to_bf16_row_ref, - }, - [31] = { // GGML_TYPE_Q4_0_4_4 - .type_name = "TYPE_Q4_0_4_4 REMOVED, use Q4_0 with runtime repacking", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [32] = { // GGML_TYPE_Q4_0_4_8 - .type_name = "TYPE_Q4_0_4_8 REMOVED, use Q4_0 with runtime repacking", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [33] = { // GGML_TYPE_Q4_0_8_8 - .type_name = "TYPE_Q4_0_8_8 REMOVED, use Q4_0 with runtime repacking", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [GGML_TYPE_TQ1_0] = { - .type_name = "tq1_0", - .blck_size = QK_K, - .type_size = sizeof(block_tq1_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_tq1_0, - .from_float_ref = (ggml_from_float_t) quantize_row_tq1_0_ref, - }, - [GGML_TYPE_TQ2_0] = { - .type_name = "tq2_0", - .blck_size = QK_K, - .type_size = sizeof(block_tq2_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_tq2_0, - .from_float_ref = (ggml_from_float_t) quantize_row_tq2_0_ref, - }, - [36] = { // GGML_TYPE_IQ4_NL_4_4 - .type_name = "TYPE_IQ4_NL_4_4 REMOVED, use IQ4_NL with runtime repacking", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [37] = { // GGML_TYPE_IQ4_NL_4_8 - .type_name = "TYPE_IQ4_NL_4_8 REMOVED, use IQ4_NL with runtime repacking", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [38] = { // GGML_TYPE_IQ4_NL_8_8 - .type_name = "TYPE_IQ4_NL_8_8 REMOVED, use IQ4_NL with runtime repacking", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, -}; - -const struct ggml_type_traits * ggml_get_type_traits(enum ggml_type type) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - return &type_traits[type]; -} - -// -// ggml object -// - -struct ggml_object { - size_t offs; - size_t size; - - struct ggml_object * next; - - enum ggml_object_type type; - - char padding[4]; -}; - -static const size_t GGML_OBJECT_SIZE = sizeof(struct ggml_object); - -// -// ggml context -// - -struct ggml_context { - size_t mem_size; - void * mem_buffer; - bool mem_buffer_owned; - bool no_alloc; - - int n_objects; - - struct ggml_object * objects_begin; - struct ggml_object * objects_end; -}; - -// -// data types -// - -static const char * GGML_OP_NAME[GGML_OP_COUNT] = { - "NONE", - - "DUP", - "ADD", - "ADD_ID", - "ADD1", - "ACC", - "SUB", - "MUL", - "DIV", - "SQR", - "SQRT", - "LOG", - "SIN", - "COS", - "SUM", - "SUM_ROWS", - "CUMSUM", - "MEAN", - "ARGMAX", - "COUNT_EQUAL", - "REPEAT", - "REPEAT_BACK", - "CONCAT", - "SILU_BACK", - "NORM", - "RMS_NORM", - "RMS_NORM_BACK", - "GROUP_NORM", - "L2_NORM", - - "MUL_MAT", - "MUL_MAT_ID", - "OUT_PROD", - - "SCALE", - "SET", - "CPY", - "CONT", - "RESHAPE", - "VIEW", - "PERMUTE", - "TRANSPOSE", - "GET_ROWS", - "GET_ROWS_BACK", - "SET_ROWS", - "DIAG", - "DIAG_MASK_INF", - "DIAG_MASK_ZERO", - "SOFT_MAX", - "SOFT_MAX_BACK", - "ROPE", - "ROPE_BACK", - "CLAMP", - "CONV_TRANSPOSE_1D", - "IM2COL", - "IM2COL_BACK", - "IM2COL_3D", - "CONV_2D", - "CONV_3D", - "CONV_2D_DW", - "CONV_TRANSPOSE_2D", - "POOL_1D", - "POOL_2D", - "POOL_2D_BACK", - "UPSCALE", - "PAD", - "PAD_REFLECT_1D", - "ROLL", - "ARANGE", - "TIMESTEP_EMBEDDING", - "ARGSORT", - "TOP_K", - "LEAKY_RELU", - "TRI", - "FILL", - - "FLASH_ATTN_EXT", - "FLASH_ATTN_BACK", - "SSM_CONV", - "SSM_SCAN", - "WIN_PART", - "WIN_UNPART", - "GET_REL_POS", - "ADD_REL_POS", - "RWKV_WKV6", - "GATED_LINEAR_ATTN", - "RWKV_WKV7", - "SOLVE_TRI", - "GATED_DELTA_NET", - - "UNARY", - - "MAP_CUSTOM1", - "MAP_CUSTOM2", - "MAP_CUSTOM3", - - "CUSTOM", - - "CROSS_ENTROPY_LOSS", - "CROSS_ENTROPY_LOSS_BACK", - "OPT_STEP_ADAMW", - "OPT_STEP_SGD", - - "GLU", -}; - -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT != 96"); - -static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { - "none", - - "x", - "x+y", - "x[i]+y", - "x+y", - "view(x,nb,offset)+=y->x", - "x-y", - "x*y", - "x/y", - "x^2", - "√x", - "log(x)", - "sin(x)", - "cos(x)", - "Σx", - "Σx_k", - "cumsum(x)", - "Σx/n", - "argmax(x)", - "count_equal(x)", - "repeat(x)", - "repeat_back(x)", - "concat(x, y)", - "silu_back(x)", - "norm(x)", - "rms_norm(x)", - "rms_norm_back(x)", - "group_norm(x)", - "l2_norm(x)", - - "X*Y", - "X[i]*Y", - "X*Y", - - "x*v", - "y-\\>view(x)", - "x-\\>y", - "cont(x)", - "reshape(x)", - "view(x)", - "permute(x)", - "transpose(x)", - "get_rows(x)", - "get_rows_back(x)", - "set_rows(x)", - "diag(x)", - "diag_mask_inf(x)", - "diag_mask_zero(x)", - "soft_max(x)", - "soft_max_back(x)", - "rope(x)", - "rope_back(x)", - "clamp(x)", - "conv_transpose_1d(x)", - "im2col(x)", - "im2col_back(x)", - "im2col_3d(x)", - "conv_2d(x)", - "conv_3d(x)", - "conv_2d_dw(x)", - "conv_transpose_2d(x)", - "pool_1d(x)", - "pool_2d(x)", - "pool_2d_back(x)", - "upscale(x)", - "pad(x)", - "pad_reflect_1d(x)", - "roll(x)", - "arange(start, stop, step)", - "timestep_embedding(timesteps, dim, max_period)", - "argsort(x)", - "top_k(x)", - "leaky_relu(x)", - "tri(x)", - "fill(x, c)", - - "flash_attn_ext(x)", - "flash_attn_back(x)", - "ssm_conv(x)", - "ssm_scan(x)", - "win_part(x)", - "win_unpart(x)", - "get_rel_pos(x)", - "add_rel_pos(x)", - "rwkv_wkv6(k, v, r, tf, td, s)", - "gated_linear_attn(k, v, q, gate, s)", - "rwkv_wkv7(r, w, k, v, a, b, s)", - "A X = B, A triangular, solve X", - "gated_delta_net(q, k, v, g, beta, s)", - - "unary(x)", - - "map_custom(x)", - "map_custom(x,y)", - "map_custom(x,y,z)", - - "custom(x)", - - "cross_entropy_loss(x,y)", - "cross_entropy_loss_back(x,y)", - "adamw(x)", - "sgd(x)", - - "glu(x)", -}; - -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT != 96"); - -static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); - -static const char * GGML_UNARY_OP_NAME[GGML_UNARY_OP_COUNT] = { - "ABS", - "SGN", - "NEG", - "STEP", - "TANH", - "ELU", - "RELU", - "SIGMOID", - "GELU", - "GELU_QUICK", - "SILU", - "HARDSWISH", - "HARDSIGMOID", - "EXP", - "EXPM1", - "SOFTPLUS", - "GELU_ERF", - "XIELU", - "FLOOR", - "CEIL", - "ROUND", - "TRUNC", -}; - -static_assert(GGML_UNARY_OP_COUNT == 22, "GGML_UNARY_OP_COUNT != 22"); - -static const char * GGML_GLU_OP_NAME[GGML_GLU_OP_COUNT] = { - "REGLU", - "GEGLU", - "SWIGLU", - "SWIGLU_OAI", - "GEGLU_ERF", - "GEGLU_QUICK", -}; - -static_assert(GGML_GLU_OP_COUNT == 6, "GGML_GLU_OP_COUNT != 6"); - - -static_assert(sizeof(struct ggml_object)%GGML_MEM_ALIGN == 0, "ggml_object size must be a multiple of GGML_MEM_ALIGN"); -static_assert(sizeof(struct ggml_tensor)%GGML_MEM_ALIGN == 0, "ggml_tensor size must be a multiple of GGML_MEM_ALIGN"); - - -//////////////////////////////////////////////////////////////////////////////// - -void ggml_print_object(const struct ggml_object * obj) { - GGML_LOG_INFO(" - ggml_object: type = %d, offset = %zu, size = %zu, next = %p\n", - obj->type, obj->offs, obj->size, (const void *) obj->next); -} - -void ggml_print_objects(const struct ggml_context * ctx) { - struct ggml_object * obj = ctx->objects_begin; - - GGML_LOG_INFO("%s: objects in context %p:\n", __func__, (const void *) ctx); - - while (obj != NULL) { - ggml_print_object(obj); - obj = obj->next; - } - - GGML_LOG_INFO("%s: --- end ---\n", __func__); -} - -int64_t ggml_nelements(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return tensor->ne[0]*tensor->ne[1]*tensor->ne[2]*tensor->ne[3]; -} - -int64_t ggml_nrows(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return tensor->ne[1]*tensor->ne[2]*tensor->ne[3]; -} - -size_t ggml_nbytes(const struct ggml_tensor * tensor) { - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - if (tensor->ne[i] <= 0) { - return 0; - } - } - - size_t nbytes; - const size_t blck_size = ggml_blck_size(tensor->type); - if (blck_size == 1) { - nbytes = ggml_type_size(tensor->type); - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - nbytes += (tensor->ne[i] - 1)*tensor->nb[i]; - } - } - else { - nbytes = tensor->ne[0]*tensor->nb[0]/blck_size; - for (int i = 1; i < GGML_MAX_DIMS; ++i) { - nbytes += (tensor->ne[i] - 1)*tensor->nb[i]; - } - } - - return nbytes; -} - -size_t ggml_nbytes_pad(const struct ggml_tensor * tensor) { - return GGML_PAD(ggml_nbytes(tensor), GGML_MEM_ALIGN); -} - -int64_t ggml_blck_size(enum ggml_type type) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - return type_traits[type].blck_size; -} - -size_t ggml_type_size(enum ggml_type type) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - return type_traits[type].type_size; -} - -size_t ggml_row_size(enum ggml_type type, int64_t ne) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - assert(ne % ggml_blck_size(type) == 0); - return ggml_type_size(type)*ne/ggml_blck_size(type); -} - -double ggml_type_sizef(enum ggml_type type) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - return ((double)(type_traits[type].type_size))/type_traits[type].blck_size; -} - -const char * ggml_type_name(enum ggml_type type) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - return type_traits[type].type_name; -} - -bool ggml_is_quantized(enum ggml_type type) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - return type_traits[type].is_quantized; -} - -const char * ggml_op_name(enum ggml_op op) { - return GGML_OP_NAME[op]; -} - -const char * ggml_op_symbol(enum ggml_op op) { - return GGML_OP_SYMBOL[op]; -} - -const char * ggml_unary_op_name(enum ggml_unary_op op) { - return GGML_UNARY_OP_NAME[op]; -} - -const char * ggml_glu_op_name(enum ggml_glu_op op) { - return GGML_GLU_OP_NAME[op]; -} - -const char * ggml_op_desc(const struct ggml_tensor * t) { - if (t->op == GGML_OP_UNARY) { - enum ggml_unary_op uop = ggml_get_unary_op(t); - return ggml_unary_op_name(uop); - } - if (t->op == GGML_OP_GLU) { - enum ggml_glu_op gop = ggml_get_glu_op(t); - return ggml_glu_op_name(gop); - } - return ggml_op_name(t->op); -} - -size_t ggml_element_size(const struct ggml_tensor * tensor) { - return ggml_type_size(tensor->type); -} - -bool ggml_is_scalar(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return tensor->ne[0] == 1 && tensor->ne[1] == 1 && tensor->ne[2] == 1 && tensor->ne[3] == 1; -} - -bool ggml_is_vector(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return tensor->ne[1] == 1 && tensor->ne[2] == 1 && tensor->ne[3] == 1; -} - -bool ggml_is_matrix(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return tensor->ne[2] == 1 && tensor->ne[3] == 1; -} - -bool ggml_is_3d(const struct ggml_tensor * tensor) { - return tensor->ne[3] == 1; -} - -int ggml_n_dims(const struct ggml_tensor * tensor) { - for (int i = GGML_MAX_DIMS - 1; i >= 1; --i) { - if (tensor->ne[i] > 1) { - return i + 1; - } - } - return 1; -} - -enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype) { - enum ggml_type wtype = GGML_TYPE_COUNT; - - switch (ftype) { - case GGML_FTYPE_ALL_F32: wtype = GGML_TYPE_F32; break; - case GGML_FTYPE_MOSTLY_F16: wtype = GGML_TYPE_F16; break; - case GGML_FTYPE_MOSTLY_BF16: wtype = GGML_TYPE_BF16; break; - case GGML_FTYPE_MOSTLY_Q4_0: wtype = GGML_TYPE_Q4_0; break; - case GGML_FTYPE_MOSTLY_Q4_1: wtype = GGML_TYPE_Q4_1; break; - case GGML_FTYPE_MOSTLY_Q1_0: wtype = GGML_TYPE_Q1_0; break; - case GGML_FTYPE_MOSTLY_Q5_0: wtype = GGML_TYPE_Q5_0; break; - case GGML_FTYPE_MOSTLY_Q5_1: wtype = GGML_TYPE_Q5_1; break; - case GGML_FTYPE_MOSTLY_Q8_0: wtype = GGML_TYPE_Q8_0; break; - case GGML_FTYPE_MOSTLY_MXFP4: wtype = GGML_TYPE_MXFP4; break; - case GGML_FTYPE_MOSTLY_NVFP4: wtype = GGML_TYPE_NVFP4; break; - case GGML_FTYPE_MOSTLY_Q2_K: wtype = GGML_TYPE_Q2_K; break; - case GGML_FTYPE_MOSTLY_Q3_K: wtype = GGML_TYPE_Q3_K; break; - case GGML_FTYPE_MOSTLY_Q4_K: wtype = GGML_TYPE_Q4_K; break; - case GGML_FTYPE_MOSTLY_Q5_K: wtype = GGML_TYPE_Q5_K; break; - case GGML_FTYPE_MOSTLY_Q6_K: wtype = GGML_TYPE_Q6_K; break; - case GGML_FTYPE_MOSTLY_IQ2_XXS: wtype = GGML_TYPE_IQ2_XXS; break; - case GGML_FTYPE_MOSTLY_IQ2_XS: wtype = GGML_TYPE_IQ2_XS; break; - case GGML_FTYPE_MOSTLY_IQ3_XXS: wtype = GGML_TYPE_IQ3_XXS; break; - case GGML_FTYPE_MOSTLY_IQ1_S: wtype = GGML_TYPE_IQ1_S; break; - case GGML_FTYPE_MOSTLY_IQ1_M: wtype = GGML_TYPE_IQ1_M; break; - case GGML_FTYPE_MOSTLY_IQ4_NL: wtype = GGML_TYPE_IQ4_NL; break; - case GGML_FTYPE_MOSTLY_IQ4_XS: wtype = GGML_TYPE_IQ4_XS; break; - case GGML_FTYPE_MOSTLY_IQ3_S: wtype = GGML_TYPE_IQ3_S; break; - case GGML_FTYPE_MOSTLY_IQ2_S: wtype = GGML_TYPE_IQ2_S; break; - case GGML_FTYPE_UNKNOWN: wtype = GGML_TYPE_COUNT; break; - case GGML_FTYPE_MOSTLY_Q4_1_SOME_F16: wtype = GGML_TYPE_COUNT; break; - } - - GGML_ASSERT(wtype != GGML_TYPE_COUNT); - - return wtype; -} - -size_t ggml_tensor_overhead(void) { - return GGML_OBJECT_SIZE + GGML_TENSOR_SIZE; -} - -bool ggml_is_transposed(const struct ggml_tensor * tensor) { - return tensor->nb[0] > tensor->nb[1]; -} - -static bool ggml_is_contiguous_n(const struct ggml_tensor * tensor, int n) { - size_t next_nb = ggml_type_size(tensor->type); - if (tensor->ne[0] != ggml_blck_size(tensor->type) && tensor->nb[0] != next_nb) { - return false; - } - next_nb *= tensor->ne[0]/ggml_blck_size(tensor->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - if (i > n) { - if (tensor->ne[i] != 1 && tensor->nb[i] != next_nb) { - return false; - } - next_nb *= tensor->ne[i]; - } else { - // this dimension does not need to be contiguous - next_nb = tensor->ne[i]*tensor->nb[i]; - } - } - return true; -} - -bool ggml_is_contiguous(const struct ggml_tensor * tensor) { - return ggml_is_contiguous_0(tensor); -} - -bool ggml_is_contiguous_0(const struct ggml_tensor * tensor) { - return ggml_is_contiguous_n(tensor, 0); -} - -bool ggml_is_contiguous_1(const struct ggml_tensor * tensor) { - return ggml_is_contiguous_n(tensor, 1); -} - -bool ggml_is_contiguous_2(const struct ggml_tensor * tensor) { - return ggml_is_contiguous_n(tensor, 2); -} - -bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor) { - return ggml_nbytes(tensor) == ggml_nelements(tensor) * ggml_type_size(tensor->type)/ggml_blck_size(tensor->type); -} - -bool ggml_is_permuted(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return tensor->nb[0] > tensor->nb[1] || tensor->nb[1] > tensor->nb[2] || tensor->nb[2] > tensor->nb[3]; -} - -bool ggml_is_contiguous_channels(const struct ggml_tensor * tensor) { - return - tensor->nb[0] > tensor->nb[2] && - tensor->nb[1] > tensor->nb[0] && - tensor->nb[2] == ggml_type_size(tensor->type); -} - -bool ggml_is_contiguous_rows(const struct ggml_tensor * tensor) { - return - tensor->ne[0] == ggml_blck_size(tensor->type) || - tensor->nb[0] == ggml_type_size(tensor->type); -} - -static inline bool ggml_is_padded_1d(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return - tensor->nb[0] == ggml_type_size(tensor->type) && - tensor->nb[2] == tensor->nb[1]*tensor->ne[1] && - tensor->nb[3] == tensor->nb[2]*tensor->ne[2]; -} - -bool ggml_is_empty(const struct ggml_tensor * tensor) { - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - if (tensor->ne[i] == 0) { - // empty if any dimension has no elements - return true; - } - } - return false; -} - -bool ggml_are_same_shape(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return - (t0->ne[0] == t1->ne[0]) && - (t0->ne[1] == t1->ne[1]) && - (t0->ne[2] == t1->ne[2]) && - (t0->ne[3] == t1->ne[3]); -} - -bool ggml_are_same_stride(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return - (t0->nb[0] == t1->nb[0]) && - (t0->nb[1] == t1->nb[1]) && - (t0->nb[2] == t1->nb[2]) && - (t0->nb[3] == t1->nb[3]); -} - -bool ggml_is_view(const struct ggml_tensor * t) { - return ggml_impl_is_view(t); -} - -// check if t1 can be represented as a repetition of t0 -bool ggml_can_repeat(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return ggml_is_empty(t0) ? ggml_is_empty(t1) : - (t1->ne[0]%t0->ne[0] == 0) && - (t1->ne[1]%t0->ne[1] == 0) && - (t1->ne[2]%t0->ne[2] == 0) && - (t1->ne[3]%t0->ne[3] == 0); -} - -static inline bool ggml_can_repeat_rows(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return (t0->ne[0] == t1->ne[0]) && ggml_can_repeat(t0, t1); -} - -// assert that pointer is aligned to GGML_MEM_ALIGN -#define GGML_ASSERT_ALIGNED(ptr) \ - GGML_ASSERT(((uintptr_t) (ptr))%GGML_MEM_ALIGN == 0) - -//////////////////////////////////////////////////////////////////////////////// - -struct ggml_context * ggml_init(struct ggml_init_params params) { - static bool is_first_call = true; - - ggml_critical_section_start(); - - if (is_first_call) { - // initialize time system (required on Windows) - ggml_time_init(); - - is_first_call = false; - } - - ggml_critical_section_end(); - - struct ggml_context * ctx = GGML_MALLOC(sizeof(struct ggml_context)); - - // allow to call ggml_init with 0 size - if (params.mem_size == 0) { - params.mem_size = GGML_MEM_ALIGN; - } - - const size_t mem_size = params.mem_buffer ? params.mem_size : GGML_PAD(params.mem_size, GGML_MEM_ALIGN); - - *ctx = (struct ggml_context) { - /*.mem_size =*/ mem_size, - /*.mem_buffer =*/ params.mem_buffer ? params.mem_buffer : ggml_aligned_malloc(mem_size), - /*.mem_buffer_owned =*/ params.mem_buffer ? false : true, - /*.no_alloc =*/ params.no_alloc, - /*.n_objects =*/ 0, - /*.objects_begin =*/ NULL, - /*.objects_end =*/ NULL, - }; - - GGML_ASSERT(ctx->mem_buffer != NULL); - - GGML_ASSERT_ALIGNED(ctx->mem_buffer); - - GGML_PRINT_DEBUG("%s: context initialized\n", __func__); - - return ctx; -} - -void ggml_reset(struct ggml_context * ctx) { - if (ctx == NULL) { - return; - } - - ctx->n_objects = 0; - ctx->objects_begin = NULL; - ctx->objects_end = NULL; -} - -void ggml_free(struct ggml_context * ctx) { - if (ctx == NULL) { - return; - } - - if (ctx->mem_buffer_owned) { - ggml_aligned_free(ctx->mem_buffer, ctx->mem_size); - } - - GGML_FREE(ctx); -} - -size_t ggml_used_mem(const struct ggml_context * ctx) { - return ctx->objects_end == NULL ? 0 : ctx->objects_end->offs + ctx->objects_end->size; -} - -bool ggml_get_no_alloc(struct ggml_context * ctx) { - return ctx->no_alloc; -} - -void ggml_set_no_alloc(struct ggml_context * ctx, bool no_alloc) { - ctx->no_alloc = no_alloc; -} - -void * ggml_get_mem_buffer(const struct ggml_context * ctx) { - return ctx->mem_buffer; -} - -size_t ggml_get_mem_size(const struct ggml_context * ctx) { - return ctx->mem_size; -} - -size_t ggml_get_max_tensor_size(const struct ggml_context * ctx) { - size_t max_size = 0; - - for (struct ggml_tensor * tensor = ggml_get_first_tensor(ctx); tensor != NULL; tensor = ggml_get_next_tensor(ctx, tensor)) { - size_t bytes = ggml_nbytes(tensor); - max_size = MAX(max_size, bytes); - } - - return max_size; -} - -//////////////////////////////////////////////////////////////////////////////// - -static struct ggml_object * ggml_new_object(struct ggml_context * ctx, enum ggml_object_type type, size_t size) { - // always insert objects at the end of the context's memory pool - struct ggml_object * obj_cur = ctx->objects_end; - - const size_t cur_offs = obj_cur == NULL ? 0 : obj_cur->offs; - const size_t cur_size = obj_cur == NULL ? 0 : obj_cur->size; - const size_t cur_end = cur_offs + cur_size; - - // align to GGML_MEM_ALIGN - GGML_ASSERT(size <= SIZE_MAX - (GGML_MEM_ALIGN - 1)); - size_t size_needed = GGML_PAD(size, GGML_MEM_ALIGN); - - char * const mem_buffer = ctx->mem_buffer; - struct ggml_object * const obj_new = (struct ggml_object *)(mem_buffer + cur_end); - - // integer overflow checks - if (cur_end > SIZE_MAX - size_needed) { - GGML_LOG_WARN("%s: overflow detected in cur_end (%zu) + size_needed (%zu)\n", __func__, cur_end, size_needed); - return NULL; - } - if (cur_end + size_needed > SIZE_MAX - GGML_OBJECT_SIZE) { - GGML_LOG_WARN("%s: overflow detected in cur_end (%zu) + size_needed (%zu) + GGML_OBJECT_SIZE (%zu)\n", __func__, - cur_end, size_needed, (size_t) GGML_OBJECT_SIZE); - return NULL; - } - - if (cur_end + size_needed + GGML_OBJECT_SIZE > ctx->mem_size) { - GGML_LOG_WARN("%s: not enough space in the context's memory pool (needed %zu, available %zu)\n", - __func__, cur_end + size_needed + GGML_OBJECT_SIZE, ctx->mem_size); -#ifndef NDEBUG - GGML_ABORT("not enough space in the context's memory pool"); -#endif - return NULL; - } - - *obj_new = (struct ggml_object) { - .offs = cur_end + GGML_OBJECT_SIZE, - .size = size_needed, - .next = NULL, - .type = type, - }; - - GGML_ASSERT_ALIGNED(mem_buffer + obj_new->offs); - - if (obj_cur != NULL) { - obj_cur->next = obj_new; - } else { - // this is the first object in this context - ctx->objects_begin = obj_new; - } - - ctx->objects_end = obj_new; - - //printf("%s: inserted new object at %zu, size = %zu\n", __func__, cur_end, obj_new->size); - - return obj_new; -} - -static struct ggml_tensor * ggml_new_tensor_impl( - struct ggml_context * ctx, - enum ggml_type type, - int n_dims, - const int64_t * ne, - struct ggml_tensor * view_src, - size_t view_offs) { - - GGML_ASSERT(type >= 0 && type < GGML_TYPE_COUNT); - GGML_ASSERT(n_dims >= 1 && n_dims <= GGML_MAX_DIMS); - - // find the base tensor and absolute offset - if (view_src != NULL && view_src->view_src != NULL) { - view_offs += view_src->view_offs; - view_src = view_src->view_src; - } - - size_t data_size = ggml_row_size(type, ne[0]); - for (int i = 1; i < n_dims; i++) { - data_size *= ne[i]; - } - - GGML_ASSERT(view_src == NULL || data_size == 0 || data_size + view_offs <= ggml_nbytes(view_src)); - - void * data = view_src != NULL ? view_src->data : NULL; - if (data != NULL) { - data = (char *) data + view_offs; - } - - size_t obj_alloc_size = 0; - - if (view_src == NULL && !ctx->no_alloc) { - // allocate tensor data in the context's memory pool - obj_alloc_size = data_size; - } - - GGML_ASSERT(GGML_TENSOR_SIZE <= SIZE_MAX - obj_alloc_size); - - struct ggml_object * const obj_new = ggml_new_object(ctx, GGML_OBJECT_TYPE_TENSOR, GGML_TENSOR_SIZE + obj_alloc_size); - GGML_ASSERT(obj_new); - - struct ggml_tensor * const result = (struct ggml_tensor *)((char *)ctx->mem_buffer + obj_new->offs); - - *result = (struct ggml_tensor) { - /*.type =*/ type, - /*.buffer =*/ NULL, - /*.ne =*/ { 1, 1, 1, 1 }, - /*.nb =*/ { 0, 0, 0, 0 }, - /*.op =*/ GGML_OP_NONE, - /*.op_params =*/ { 0 }, - /*.flags =*/ 0, - /*.src =*/ { NULL }, - /*.view_src =*/ view_src, - /*.view_offs =*/ view_offs, - /*.data =*/ obj_alloc_size > 0 ? (void *)(result + 1) : data, - /*.name =*/ { 0 }, - /*.extra =*/ NULL, - /*.padding =*/ { 0 }, - }; - - // TODO: this should not be needed as long as we don't rely on aligned SIMD loads - //GGML_ASSERT_ALIGNED(result->data); - - for (int i = 0; i < n_dims; i++) { - result->ne[i] = ne[i]; - } - - result->nb[0] = ggml_type_size(type); - result->nb[1] = result->nb[0]*(result->ne[0]/ggml_blck_size(type)); - for (int i = 2; i < GGML_MAX_DIMS; i++) { - result->nb[i] = result->nb[i - 1]*result->ne[i - 1]; - } - - ctx->n_objects++; - - return result; -} - -struct ggml_tensor * ggml_new_tensor( - struct ggml_context * ctx, - enum ggml_type type, - int n_dims, - const int64_t * ne) { - return ggml_new_tensor_impl(ctx, type, n_dims, ne, NULL, 0); -} - -struct ggml_tensor * ggml_new_tensor_1d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0) { - return ggml_new_tensor(ctx, type, 1, &ne0); -} - -struct ggml_tensor * ggml_new_tensor_2d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0, - int64_t ne1) { - const int64_t ne[2] = { ne0, ne1 }; - return ggml_new_tensor(ctx, type, 2, ne); -} - -struct ggml_tensor * ggml_new_tensor_3d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0, - int64_t ne1, - int64_t ne2) { - const int64_t ne[3] = { ne0, ne1, ne2 }; - return ggml_new_tensor(ctx, type, 3, ne); -} - -struct ggml_tensor * ggml_new_tensor_4d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3) { - const int64_t ne[4] = { ne0, ne1, ne2, ne3 }; - return ggml_new_tensor(ctx, type, 4, ne); -} - -void * ggml_new_buffer(struct ggml_context * ctx, size_t nbytes) { - struct ggml_object * obj = ggml_new_object(ctx, GGML_OBJECT_TYPE_WORK_BUFFER, nbytes); - - return (uint8_t *)ctx->mem_buffer + obj->offs; -} - -struct ggml_tensor * ggml_dup_tensor(struct ggml_context * ctx, const struct ggml_tensor * src) { - return ggml_new_tensor(ctx, src->type, GGML_MAX_DIMS, src->ne); -} - -void ggml_unravel_index(const struct ggml_tensor * tensor, int64_t i, int64_t * i0, int64_t * i1, int64_t * i2, int64_t * i3) { - const int64_t ne2 = tensor->ne[2]; - const int64_t ne1 = tensor->ne[1]; - const int64_t ne0 = tensor->ne[0]; - - const int64_t i3_ = (i/(ne2*ne1*ne0)); - const int64_t i2_ = (i - i3_*ne2*ne1*ne0)/(ne1*ne0); - const int64_t i1_ = (i - i3_*ne2*ne1*ne0 - i2_*ne1*ne0)/ne0; - const int64_t i0_ = (i - i3_*ne2*ne1*ne0 - i2_*ne1*ne0 - i1_*ne0); - - if (i0) { - * i0 = i0_; - } - if (i1) { - * i1 = i1_; - } - if (i2) { - * i2 = i2_; - } - if (i3) { - * i3 = i3_; - } -} - -void * ggml_get_data(const struct ggml_tensor * tensor) { - return tensor->data; -} - -float * ggml_get_data_f32(const struct ggml_tensor * tensor) { - assert(tensor->type == GGML_TYPE_F32); - return (float *)(tensor->data); -} - -enum ggml_unary_op ggml_get_unary_op(const struct ggml_tensor * tensor) { - GGML_ASSERT(tensor->op == GGML_OP_UNARY); - return (enum ggml_unary_op) ggml_get_op_params_i32(tensor, 0); -} - -enum ggml_glu_op ggml_get_glu_op(const struct ggml_tensor * tensor) { - GGML_ASSERT(tensor->op == GGML_OP_GLU); - return (enum ggml_glu_op) ggml_get_op_params_i32(tensor, 0); -} - -const char * ggml_get_name(const struct ggml_tensor * tensor) { - return tensor->name; -} - -struct ggml_tensor * ggml_set_name(struct ggml_tensor * tensor, const char * name) { - size_t i; - for (i = 0; i < sizeof(tensor->name) - 1 && name[i] != '\0'; i++) { - tensor->name[i] = name[i]; - } - tensor->name[i] = '\0'; - return tensor; -} - -struct ggml_tensor * ggml_format_name(struct ggml_tensor * tensor, const char * fmt, ...) { - va_list args; - va_start(args, fmt); - vsnprintf(tensor->name, sizeof(tensor->name), fmt, args); - va_end(args); - return tensor; -} - -struct ggml_tensor * ggml_view_tensor( - struct ggml_context * ctx, - struct ggml_tensor * src) { - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, src->type, GGML_MAX_DIMS, src->ne, src, 0); - ggml_format_name(result, "%s (view)", src->name); - - for (int i = 0; i < GGML_MAX_DIMS; i++) { - result->nb[i] = src->nb[i]; - } - - return result; -} - -struct ggml_tensor * ggml_get_first_tensor(const struct ggml_context * ctx) { - struct ggml_object * obj = ctx->objects_begin; - - char * const mem_buffer = ctx->mem_buffer; - - while (obj != NULL) { - if (obj->type == GGML_OBJECT_TYPE_TENSOR) { - return (struct ggml_tensor *)(mem_buffer + obj->offs); - } - - obj = obj->next; - } - - return NULL; -} - -struct ggml_tensor * ggml_get_next_tensor(const struct ggml_context * ctx, struct ggml_tensor * tensor) { - struct ggml_object * obj = (struct ggml_object *) ((char *)tensor - GGML_OBJECT_SIZE); - obj = obj->next; - - char * const mem_buffer = ctx->mem_buffer; - - while (obj != NULL) { - if (obj->type == GGML_OBJECT_TYPE_TENSOR) { - return (struct ggml_tensor *)(mem_buffer + obj->offs); - } - - obj = obj->next; - } - - return NULL; -} - -struct ggml_tensor * ggml_get_tensor(struct ggml_context * ctx, const char * name) { - struct ggml_object * obj = ctx->objects_begin; - - char * const mem_buffer = ctx->mem_buffer; - - while (obj != NULL) { - if (obj->type == GGML_OBJECT_TYPE_TENSOR) { - struct ggml_tensor * cur = (struct ggml_tensor *)(mem_buffer + obj->offs); - if (strcmp(cur->name, name) == 0) { - return cur; - } - } - - obj = obj->next; - } - - return NULL; -} - -//////////////////////////////////////////////////////////////////////////////// - -// ggml_dup - -static struct ggml_tensor * ggml_dup_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_DUP; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_dup( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_dup_impl(ctx, a, false); -} - -struct ggml_tensor * ggml_dup_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_dup_impl(ctx, a, true); -} - -// ggml_add - -static struct ggml_tensor * ggml_add_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool inplace) { - GGML_ASSERT(ggml_can_repeat(b, a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_ADD; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_add( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_add_impl(ctx, a, b, false); -} - -struct ggml_tensor * ggml_add_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_add_impl(ctx, a, b, true); -} - -// ggml_add_cast - -static struct ggml_tensor * ggml_add_cast_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - enum ggml_type type) { - // TODO: support less-strict constraint - // GGML_ASSERT(ggml_can_repeat(b, a)); - GGML_ASSERT(ggml_can_repeat_rows(b, a)); - - // currently only supported for quantized input and f16 - GGML_ASSERT(ggml_is_quantized(a->type) || - a->type == GGML_TYPE_F16 || - a->type == GGML_TYPE_BF16); - - struct ggml_tensor * result = ggml_new_tensor(ctx, type, GGML_MAX_DIMS, a->ne); - - result->op = GGML_OP_ADD; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_add_cast( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - enum ggml_type type) { - return ggml_add_cast_impl(ctx, a, b, type); -} - -struct ggml_tensor * ggml_add_id( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * ids) { - - GGML_ASSERT(a->ne[0] == b->ne[0]); - GGML_ASSERT(a->ne[1] == ids->ne[0]); - GGML_ASSERT(a->ne[2] == ids->ne[1]); - GGML_ASSERT(ids->type == GGML_TYPE_I32); - - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_ADD_ID; - result->src[0] = a; - result->src[1] = b; - result->src[2] = ids; - - return result; -} - -// ggml_add1 - -static struct ggml_tensor * ggml_add1_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool inplace) { - GGML_ASSERT(ggml_is_scalar(b)); - GGML_ASSERT(ggml_is_padded_1d(a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_ADD1; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_add1( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_add1_impl(ctx, a, b, false); -} - -struct ggml_tensor * ggml_add1_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_add1_impl(ctx, a, b, true); -} - -// ggml_acc - -static struct ggml_tensor * ggml_acc_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset, - bool inplace) { - GGML_ASSERT(ggml_nelements(b) <= ggml_nelements(a)); - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(a->type == GGML_TYPE_F32); - GGML_ASSERT(b->type == GGML_TYPE_F32); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - int32_t params[] = { nb1, nb2, nb3, offset, inplace ? 1 : 0 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_ACC; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_acc( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset) { - return ggml_acc_impl(ctx, a, b, nb1, nb2, nb3, offset, false); -} - -struct ggml_tensor * ggml_acc_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset) { - return ggml_acc_impl(ctx, a, b, nb1, nb2, nb3, offset, true); -} - -// ggml_sub - -static struct ggml_tensor * ggml_sub_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool inplace) { - GGML_ASSERT(ggml_can_repeat(b, a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_SUB; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_sub( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_sub_impl(ctx, a, b, false); -} - -struct ggml_tensor * ggml_sub_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_sub_impl(ctx, a, b, true); -} - -// ggml_mul - -static struct ggml_tensor * ggml_mul_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool inplace) { - GGML_ASSERT(ggml_can_repeat(b, a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_MUL; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_mul( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_mul_impl(ctx, a, b, false); -} - -struct ggml_tensor * ggml_mul_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_mul_impl(ctx, a, b, true); -} - -// ggml_div - -static struct ggml_tensor * ggml_div_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool inplace) { - GGML_ASSERT(ggml_can_repeat(b, a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_DIV; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_div( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_div_impl(ctx, a, b, false); -} - -struct ggml_tensor * ggml_div_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_div_impl(ctx, a, b, true); -} - -// ggml_sqr - -static struct ggml_tensor * ggml_sqr_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_SQR; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_sqr( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_sqr_impl(ctx, a, false); -} - -struct ggml_tensor * ggml_sqr_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_sqr_impl(ctx, a, true); -} - -// ggml_sqrt - -static struct ggml_tensor * ggml_sqrt_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_SQRT; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_sqrt( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_sqrt_impl(ctx, a, false); -} - -struct ggml_tensor * ggml_sqrt_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_sqrt_impl(ctx, a, true); -} - -// ggml_log - -static struct ggml_tensor * ggml_log_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_LOG; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_log( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_log_impl(ctx, a, false); -} - -struct ggml_tensor * ggml_log_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_log_impl(ctx, a, true); -} - -struct ggml_tensor * ggml_expm1( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_EXPM1); -} - -struct ggml_tensor * ggml_expm1_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_EXPM1); -} - -struct ggml_tensor * ggml_softplus( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_SOFTPLUS); -} - -struct ggml_tensor * ggml_softplus_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SOFTPLUS); -} - -// ggml_sin - -static struct ggml_tensor * ggml_sin_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_SIN; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_sin( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_sin_impl(ctx, a, false); -} - -struct ggml_tensor * ggml_sin_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_sin_impl(ctx, a, true); -} - -// ggml_cos - -static struct ggml_tensor * ggml_cos_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_COS; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_cos( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_cos_impl(ctx, a, false); -} - -struct ggml_tensor * ggml_cos_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_cos_impl(ctx, a, true); -} - -// ggml_sum - -struct ggml_tensor * ggml_sum( - struct ggml_context * ctx, - struct ggml_tensor * a) { - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, a->type, 1); - - result->op = GGML_OP_SUM; - result->src[0] = a; - - return result; -} - -// ggml_sum_rows - -struct ggml_tensor * ggml_sum_rows( - struct ggml_context * ctx, - struct ggml_tensor * a) { - int64_t ne[GGML_MAX_DIMS] = { 1 }; - for (int i = 1; i < GGML_MAX_DIMS; ++i) { - ne[i] = a->ne[i]; - } - - struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne); - - result->op = GGML_OP_SUM_ROWS; - result->src[0] = a; - - return result; -} - -// ggml_cumsum - -struct ggml_tensor * ggml_cumsum( - struct ggml_context * ctx, - struct ggml_tensor * a) { - GGML_ASSERT(a->type == GGML_TYPE_F32); - - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_CUMSUM; - result->src[0] = a; - - return result; -} - -// ggml_mean - -struct ggml_tensor * ggml_mean( - struct ggml_context * ctx, - struct ggml_tensor * a) { - int64_t ne[4] = { 1, a->ne[1], a->ne[2], a->ne[3] }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_MEAN; - result->src[0] = a; - - return result; -} - -// ggml_argmax - -struct ggml_tensor * ggml_argmax( - struct ggml_context * ctx, - struct ggml_tensor * a) { - GGML_ASSERT(ggml_is_matrix(a)); - GGML_ASSERT(a->ne[0] <= INT32_MAX); - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, a->ne[1]); - - result->op = GGML_OP_ARGMAX; - result->src[0] = a; - - return result; -} - -// ggml_count_equal - -struct ggml_tensor * ggml_count_equal( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_are_same_shape(a, b)); - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); - - result->op = GGML_OP_COUNT_EQUAL; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_repeat - -struct ggml_tensor * ggml_repeat( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_can_repeat(a, b)); - - struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, b->ne); - - result->op = GGML_OP_REPEAT; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_repeat_4d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { - const bool can_repeat = ggml_is_empty(a) || ( - (ne0 % a->ne[0] == 0) && - (ne1 % a->ne[1] == 0) && - (ne2 % a->ne[2] == 0) && - (ne3 % a->ne[3] == 0) - ); - GGML_ASSERT(can_repeat); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, ne0, ne1, ne2, ne3); - - result->op = GGML_OP_REPEAT; - result->src[0] = a; - - return result; -} - -// ggml_repeat_back - -struct ggml_tensor * ggml_repeat_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_can_repeat(b, a)); - - struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, b->ne); - - result->op = GGML_OP_REPEAT_BACK; - result->src[0] = a; - - return result; -} - -// ggml_concat - -struct ggml_tensor * ggml_concat( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int dim) { - GGML_ASSERT(dim >= 0 && dim < GGML_MAX_DIMS); - GGML_ASSERT(a->type == b->type); - - int64_t ne[GGML_MAX_DIMS]; - for (int d = 0; d < GGML_MAX_DIMS; ++d) { - if (d == dim) { - ne[d] = a->ne[d] + b->ne[d]; - continue; - } - GGML_ASSERT(a->ne[d] == b->ne[d]); - ne[d] = a->ne[d]; - } - - struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne); - - ggml_set_op_params_i32(result, 0, dim); - - result->op = GGML_OP_CONCAT; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_abs - -struct ggml_tensor * ggml_abs( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_ABS); -} - -struct ggml_tensor * ggml_abs_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_ABS); -} - -// ggml_sgn - -struct ggml_tensor * ggml_sgn( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_SGN); -} - -struct ggml_tensor * ggml_sgn_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SGN); -} - -// ggml_neg - -struct ggml_tensor * ggml_neg( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_NEG); -} - -struct ggml_tensor * ggml_neg_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_NEG); -} - -// ggml_step - -struct ggml_tensor * ggml_step( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_STEP); -} - -struct ggml_tensor * ggml_step_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_STEP); -} - -// ggml_tanh - -struct ggml_tensor * ggml_tanh( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_TANH); -} - -struct ggml_tensor * ggml_tanh_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_TANH); -} - -// ggml_elu - -struct ggml_tensor * ggml_elu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_ELU); -} - -struct ggml_tensor * ggml_elu_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_ELU); -} - -// ggml_relu - -struct ggml_tensor * ggml_relu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_RELU); -} - -struct ggml_tensor * ggml_relu_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_RELU); -} - -// ggml_leaky_relu - -struct ggml_tensor * ggml_leaky_relu( - struct ggml_context * ctx, - struct ggml_tensor * a, - float negative_slope, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params(result, &negative_slope, sizeof(negative_slope)); - - result->op = GGML_OP_LEAKY_RELU; - result->src[0] = a; - - return result; -} - -// ggml_sigmoid - -struct ggml_tensor * ggml_sigmoid( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_SIGMOID); -} - -struct ggml_tensor * ggml_sigmoid_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SIGMOID); -} - -// ggml_gelu - -struct ggml_tensor * ggml_gelu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_GELU); -} - -struct ggml_tensor * ggml_gelu_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_GELU); -} - -// ggml_gelu_erf - -struct ggml_tensor * ggml_gelu_erf( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_GELU_ERF); -} - -struct ggml_tensor * ggml_gelu_erf_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_GELU_ERF); -} - -// ggml_gelu_quick - -struct ggml_tensor * ggml_gelu_quick( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_GELU_QUICK); -} - -struct ggml_tensor * ggml_gelu_quick_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_GELU_QUICK); -} - -// ggml_silu - -struct ggml_tensor * ggml_silu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_SILU); -} - -struct ggml_tensor * ggml_silu_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SILU); -} - -// ggml_xielu - -struct ggml_tensor * ggml_xielu( - struct ggml_context * ctx, - struct ggml_tensor * a, - float alpha_n, - float alpha_p, - float beta, - float eps) { - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - ggml_set_op_params_i32(result, 0, (int32_t) GGML_UNARY_OP_XIELU); - ggml_set_op_params_f32(result, 1, beta + ggml_compute_softplus_f32(alpha_n)); - ggml_set_op_params_f32(result, 2, ggml_compute_softplus_f32(alpha_p)); - ggml_set_op_params_f32(result, 3, beta); - ggml_set_op_params_f32(result, 4, eps); - - result->op = GGML_OP_UNARY; - result->src[0] = a; - - return result; -} - -// ggml_silu_back - -struct ggml_tensor * ggml_silu_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_SILU_BACK; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml hardswish - -struct ggml_tensor * ggml_hardswish( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_HARDSWISH); -} - -// ggml hardsigmoid - -struct ggml_tensor * ggml_hardsigmoid( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_HARDSIGMOID); -} - -// ggml exp - -struct ggml_tensor * ggml_exp( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_EXP); -} - -struct ggml_tensor * ggml_exp_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_EXP); -} - -// ggml_glu - -static struct ggml_tensor * ggml_glu_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - enum ggml_glu_op op, - bool swapped) { - GGML_ASSERT(ggml_is_contiguous_1(a)); - - if (b) { - GGML_ASSERT(ggml_is_contiguous_1(b)); - GGML_ASSERT(ggml_are_same_shape(a, b)); - GGML_ASSERT(a->type == b->type); - } - - int64_t ne[GGML_MAX_DIMS] = { a->ne[0] / 2 }; for (int i = 1; i < GGML_MAX_DIMS; i++) ne[i] = a->ne[i]; - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b ? a->ne : ne, NULL, 0); - - ggml_set_op_params_i32(result, 0, (int32_t) op); - ggml_set_op_params_i32(result, 1, (int32_t) swapped); - - result->op = GGML_OP_GLU; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_floor - -struct ggml_tensor * ggml_floor( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_FLOOR); -} - -struct ggml_tensor * ggml_floor_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_FLOOR); -} - -// ggml_ceil - -struct ggml_tensor * ggml_ceil( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_CEIL); -} - -struct ggml_tensor * ggml_ceil_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_CEIL); -} - -//ggml_round - -struct ggml_tensor * ggml_round( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_ROUND); -} - -struct ggml_tensor * ggml_round_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_ROUND); -} - -//ggml_trunc - -struct ggml_tensor * ggml_trunc( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_TRUNC); -} - -struct ggml_tensor * ggml_trunc_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_TRUNC); -} - -struct ggml_tensor * ggml_glu( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_glu_op op, - bool swapped) { - return ggml_glu_impl(ctx, a, NULL, op, swapped); -} - -struct ggml_tensor * ggml_glu_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - enum ggml_glu_op op) { - return ggml_glu_impl(ctx, a, b, op, false); -} - -// ggml_reglu - -struct ggml_tensor * ggml_reglu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_REGLU, false); -} - -struct ggml_tensor * ggml_reglu_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_REGLU, true); -} - -struct ggml_tensor * ggml_reglu_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_REGLU, false); -} - -// ggml_geglu - -struct ggml_tensor * ggml_geglu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU, false); -} - -struct ggml_tensor * ggml_geglu_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU, true); -} - -struct ggml_tensor * ggml_geglu_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_GEGLU, false); -} - -// ggml_swiglu - -struct ggml_tensor * ggml_swiglu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_SWIGLU, false); -} - -struct ggml_tensor * ggml_swiglu_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_SWIGLU, true); -} - -struct ggml_tensor * ggml_swiglu_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_SWIGLU, false); -} - -// ggml_geglu_erf - -struct ggml_tensor * ggml_geglu_erf( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_ERF, false); -} - -struct ggml_tensor * ggml_geglu_erf_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_ERF, true); -} - -struct ggml_tensor * ggml_geglu_erf_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_GEGLU_ERF, false); -} - -// ggml_geglu_quick - -struct ggml_tensor * ggml_geglu_quick( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_QUICK, false); -} - -struct ggml_tensor * ggml_geglu_quick_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_QUICK, true); -} - -struct ggml_tensor * ggml_geglu_quick_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_GEGLU_QUICK, false); -} - -struct ggml_tensor * ggml_swiglu_oai( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float alpha, - float limit) { - struct ggml_tensor * result = ggml_glu_impl(ctx, a, b, GGML_GLU_OP_SWIGLU_OAI, false); - ggml_set_op_params_f32(result, 2, alpha); - ggml_set_op_params_f32(result, 3, limit); - - return result; -} - -// ggml_norm - -static struct ggml_tensor * ggml_norm_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params(result, &eps, sizeof(eps)); - - result->op = GGML_OP_NORM; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_norm( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps) { - return ggml_norm_impl(ctx, a, eps, false); -} - -struct ggml_tensor * ggml_norm_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps) { - return ggml_norm_impl(ctx, a, eps, true); -} - -// ggml_rms_norm - -static struct ggml_tensor * ggml_rms_norm_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params(result, &eps, sizeof(eps)); - - result->op = GGML_OP_RMS_NORM; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_rms_norm( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps) { - return ggml_rms_norm_impl(ctx, a, eps, false); -} - -struct ggml_tensor * ggml_rms_norm_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps) { - return ggml_rms_norm_impl(ctx, a, eps, true); -} - -// ggml_rms_norm_back - -struct ggml_tensor * ggml_rms_norm_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float eps) { - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - ggml_set_op_params(result, &eps, sizeof(eps)); - - result->op = GGML_OP_RMS_NORM_BACK; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_group_norm - -static struct ggml_tensor * ggml_group_norm_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_groups, - float eps, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params_i32(result, 0, n_groups); - ggml_set_op_params_f32(result, 1, eps); - - result->op = GGML_OP_GROUP_NORM; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_group_norm( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_groups, - float eps) { - return ggml_group_norm_impl(ctx, a, n_groups, eps, false); -} - -struct ggml_tensor * ggml_group_norm_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_groups, - float eps) { - return ggml_group_norm_impl(ctx, a, n_groups, eps, true); -} - -// ggml_l2_norm - -static struct ggml_tensor * ggml_l2_norm_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params_f32(result, 0, eps); - - result->op = GGML_OP_L2_NORM; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_l2_norm( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps) { - return ggml_l2_norm_impl(ctx, a, eps, false); -} - -struct ggml_tensor * ggml_l2_norm_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps) { - return ggml_l2_norm_impl(ctx, a, eps, true); -} - -// ggml_mul_mat - -static inline bool ggml_can_mul_mat(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return (t0->ne[0] == t1->ne[0]) && - (t1->ne[2]%t0->ne[2] == 0) && // verify t0 is broadcastable - (t1->ne[3]%t0->ne[3] == 0); -} - -struct ggml_tensor * ggml_mul_mat( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_can_mul_mat(a, b)); - GGML_ASSERT(!ggml_is_transposed(a)); - - const int64_t ne[4] = { a->ne[1], b->ne[1], b->ne[2], b->ne[3] }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_MUL_MAT; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -void ggml_mul_mat_set_prec( - struct ggml_tensor * a, - enum ggml_prec prec) { - GGML_ASSERT(a->op == GGML_OP_MUL_MAT); - - const int32_t prec_i32 = (int32_t) prec; - - ggml_set_op_params_i32(a, 0, prec_i32); -} - -// ggml_mul_mat_id - -/* - c = ggml_mul_mat_id(ctx, as, b, ids); - - as -> [cols, rows, n_expert] - b -> [cols, n_expert_used, n_tokens] - ids -> [n_expert_used, n_tokens] (i32) - c -> [rows, n_expert_used, n_tokens] - - in b, n_expert_used can be broadcasted to match the n_expert_used of ids - - c ~= as[:,:,i] @ b[:,i%r,t], i = ids[e,t] for all e,t in ids -*/ -struct ggml_tensor * ggml_mul_mat_id( - struct ggml_context * ctx, - struct ggml_tensor * as, - struct ggml_tensor * b, - struct ggml_tensor * ids) { - GGML_ASSERT(!ggml_is_transposed(as)); - GGML_ASSERT(ids->type == GGML_TYPE_I32); - - GGML_ASSERT(as->ne[3] == 1); // as is 3d (one matrix per expert) - GGML_ASSERT(b->ne[3] == 1); // b is 3d - GGML_ASSERT(ids->ne[2] == 1 && ids->ne[3] == 1); // ids is 2d - GGML_ASSERT(ids->ne[1] == b->ne[2]); // must have an expert list per b row - GGML_ASSERT(as->ne[0] == b->ne[0]); // can_mul_mat - GGML_ASSERT(ids->ne[0] % b->ne[1] == 0); // can broadcast - - const int64_t ne[4] = { as->ne[1], ids->ne[0], b->ne[2], 1 }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_MUL_MAT_ID; - result->src[0] = as; - result->src[1] = b; - result->src[2] = ids; - - return result; -} - -// ggml_out_prod - -static inline bool ggml_can_out_prod(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return (t0->ne[1] == t1->ne[1]) && - (t1->ne[2]%t0->ne[2] == 0) && // verify t0 is broadcastable - (t1->ne[3]%t0->ne[3] == 0); -} - -struct ggml_tensor * ggml_out_prod( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_can_out_prod(a, b)); - GGML_ASSERT(!ggml_is_transposed(a)); - - // a is broadcastable to b for ne[2] and ne[3] -> use b->ne[2] and b->ne[3] - const int64_t ne[4] = { a->ne[0], b->ne[0], b->ne[2], b->ne[3] }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_OUT_PROD; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_scale - -static struct ggml_tensor * ggml_scale_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s, - float b, - bool inplace) { - GGML_ASSERT(ggml_is_padded_1d(a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - float params[2] = { s, b }; - ggml_set_op_params(result, ¶ms, sizeof(params)); - - result->op = GGML_OP_SCALE; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_scale( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s) { - return ggml_scale_impl(ctx, a, s, 0.0, false); -} - -struct ggml_tensor * ggml_scale_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s) { - return ggml_scale_impl(ctx, a, s, 0.0, true); -} - -struct ggml_tensor * ggml_scale_bias( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s, - float b) { - return ggml_scale_impl(ctx, a, s, b, false); -} - -struct ggml_tensor * ggml_scale_bias_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s, - float b) { - return ggml_scale_impl(ctx, a, s, b, true); -} - -// ggml_set - -static struct ggml_tensor * ggml_set_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset, - bool inplace) { - GGML_ASSERT(ggml_nelements(a) >= ggml_nelements(b)); - - // make a view of the destination - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - GGML_ASSERT(offset < (size_t)(1 << 30)); - int32_t params[] = { nb1, nb2, nb3, offset, inplace ? 1 : 0 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_SET; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_set( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset) { - return ggml_set_impl(ctx, a, b, nb1, nb2, nb3, offset, false); -} - -struct ggml_tensor * ggml_set_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset) { - return ggml_set_impl(ctx, a, b, nb1, nb2, nb3, offset, true); -} - -struct ggml_tensor * ggml_set_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t offset) { - return ggml_set_impl(ctx, a, b, a->nb[1], a->nb[2], a->nb[3], offset, false); -} - -struct ggml_tensor * ggml_set_1d_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t offset) { - return ggml_set_impl(ctx, a, b, a->nb[1], a->nb[2], a->nb[3], offset, true); -} - -struct ggml_tensor * ggml_set_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t offset) { - return ggml_set_impl(ctx, a, b, nb1, a->nb[2], a->nb[3], offset, false); -} - -struct ggml_tensor * ggml_set_2d_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t offset) { - return ggml_set_impl(ctx, a, b, nb1, a->nb[2], a->nb[3], offset, true); -} - -// ggml_cpy - -static struct ggml_tensor * ggml_cpy_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b)); - - // make a view of the destination - struct ggml_tensor * result = ggml_view_tensor(ctx, b); - if (strlen(b->name) > 0) { - ggml_format_name(result, "%s (copy of %s)", b->name, a->name); - } else { - ggml_format_name(result, "%s (copy)", a->name); - } - - result->op = GGML_OP_CPY; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_cpy( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_cpy_impl(ctx, a, b); -} - -struct ggml_tensor * ggml_cast( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_type type) { - struct ggml_tensor * result = ggml_new_tensor(ctx, type, GGML_MAX_DIMS, a->ne); - ggml_format_name(result, "%s (copy)", a->name); - - result->op = GGML_OP_CPY; - result->src[0] = a; - result->src[1] = result; // note: this self-reference might seem redundant, but it's actually needed by some - // backends for consistency with ggml_cpy_impl() above - - return result; -} - -// ggml_cont - -static struct ggml_tensor * ggml_cont_impl( - struct ggml_context * ctx, - struct ggml_tensor * a) { - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - ggml_format_name(result, "%s (cont)", a->name); - - result->op = GGML_OP_CONT; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_cont( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_cont_impl(ctx, a); -} - -// make contiguous, with new shape -GGML_API struct ggml_tensor * ggml_cont_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0) { - return ggml_cont_4d(ctx, a, ne0, 1, 1, 1); -} - -GGML_API struct ggml_tensor * ggml_cont_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1) { - return ggml_cont_4d(ctx, a, ne0, ne1, 1, 1); -} - -GGML_API struct ggml_tensor * ggml_cont_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2) { - return ggml_cont_4d(ctx, a, ne0, ne1, ne2, 1); -} - -struct ggml_tensor * ggml_cont_4d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3) { - GGML_ASSERT(ggml_nelements(a) == (ne0*ne1*ne2*ne3)); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, ne0, ne1, ne2, ne3); - ggml_format_name(result, "%s (cont)", a->name); - - result->op = GGML_OP_CONT; - result->src[0] = a; - - return result; -} - -// ggml_reshape - -struct ggml_tensor * ggml_reshape( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_is_contiguous(a)); - // as only the shape of b is relevant, and not its memory layout, b is allowed to be non contiguous. - GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b)); - - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b->ne, a, 0); - ggml_format_name(result, "%s (reshaped)", a->name); - - result->op = GGML_OP_RESHAPE; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_reshape_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0) { - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_nelements(a) == ne0); - - const int64_t ne[1] = { ne0 }; - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 1, ne, a, 0); - ggml_format_name(result, "%s (reshaped)", a->name); - - result->op = GGML_OP_RESHAPE; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_reshape_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1) { - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_nelements(a) == ne0*ne1); - - const int64_t ne[2] = { ne0, ne1 }; - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 2, ne, a, 0); - ggml_format_name(result, "%s (reshaped)", a->name); - - result->op = GGML_OP_RESHAPE; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_reshape_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2) { - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_nelements(a) == ne0*ne1*ne2); - - const int64_t ne[3] = { ne0, ne1, ne2 }; - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 3, ne, a, 0); - ggml_format_name(result, "%s (reshaped)", a->name); - - result->op = GGML_OP_RESHAPE; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_reshape_4d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3) { - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_nelements(a) == ne0*ne1*ne2*ne3); - - const int64_t ne[4] = { ne0, ne1, ne2, ne3 }; - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 4, ne, a, 0); - ggml_format_name(result, "%s (reshaped)", a->name); - - result->op = GGML_OP_RESHAPE; - result->src[0] = a; - - return result; -} - -static struct ggml_tensor * ggml_view_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_dims, - const int64_t * ne, - size_t offset) { - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, n_dims, ne, a, offset); - ggml_format_name(result, "%s (view)", a->name); - - ggml_set_op_params(result, &offset, sizeof(offset)); - - result->op = GGML_OP_VIEW; - result->src[0] = a; - - return result; -} - -// ggml_view_1d - -struct ggml_tensor * ggml_view_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - size_t offset) { - struct ggml_tensor * result = ggml_view_impl(ctx, a, 1, &ne0, offset); - - return result; -} - -// ggml_view_2d - -struct ggml_tensor * ggml_view_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - size_t nb1, - size_t offset) { - const int64_t ne[2] = { ne0, ne1 }; - - struct ggml_tensor * result = ggml_view_impl(ctx, a, 2, ne, offset); - - result->nb[1] = nb1; - result->nb[2] = result->nb[1]*ne1; - result->nb[3] = result->nb[2]; - - return result; -} - -// ggml_view_3d - -struct ggml_tensor * ggml_view_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - size_t nb1, - size_t nb2, - size_t offset) { - const int64_t ne[3] = { ne0, ne1, ne2 }; - - struct ggml_tensor * result = ggml_view_impl(ctx, a, 3, ne, offset); - - result->nb[1] = nb1; - result->nb[2] = nb2; - result->nb[3] = result->nb[2]*ne2; - - return result; -} - -// ggml_view_4d - -struct ggml_tensor * ggml_view_4d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset) { - const int64_t ne[4] = { ne0, ne1, ne2, ne3 }; - - struct ggml_tensor * result = ggml_view_impl(ctx, a, 4, ne, offset); - - result->nb[1] = nb1; - result->nb[2] = nb2; - result->nb[3] = nb3; - - return result; -} - -// ggml_permute - -struct ggml_tensor * ggml_permute( - struct ggml_context * ctx, - struct ggml_tensor * a, - int axis0, - int axis1, - int axis2, - int axis3) { - GGML_ASSERT(axis0 >= 0 && axis0 < GGML_MAX_DIMS); - GGML_ASSERT(axis1 >= 0 && axis1 < GGML_MAX_DIMS); - GGML_ASSERT(axis2 >= 0 && axis2 < GGML_MAX_DIMS); - GGML_ASSERT(axis3 >= 0 && axis3 < GGML_MAX_DIMS); - - GGML_ASSERT(axis0 != axis1); - GGML_ASSERT(axis0 != axis2); - GGML_ASSERT(axis0 != axis3); - GGML_ASSERT(axis1 != axis2); - GGML_ASSERT(axis1 != axis3); - GGML_ASSERT(axis2 != axis3); - - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - ggml_format_name(result, "%s (permuted)", a->name); - - int ne[GGML_MAX_DIMS]; - int nb[GGML_MAX_DIMS]; - - ne[axis0] = a->ne[0]; - ne[axis1] = a->ne[1]; - ne[axis2] = a->ne[2]; - ne[axis3] = a->ne[3]; - - nb[axis0] = a->nb[0]; - nb[axis1] = a->nb[1]; - nb[axis2] = a->nb[2]; - nb[axis3] = a->nb[3]; - - result->ne[0] = ne[0]; - result->ne[1] = ne[1]; - result->ne[2] = ne[2]; - result->ne[3] = ne[3]; - - result->nb[0] = nb[0]; - result->nb[1] = nb[1]; - result->nb[2] = nb[2]; - result->nb[3] = nb[3]; - - result->op = GGML_OP_PERMUTE; - result->src[0] = a; - - int32_t params[] = { axis0, axis1, axis2, axis3 }; - ggml_set_op_params(result, params, sizeof(params)); - - return result; -} - -// ggml_transpose - -struct ggml_tensor * ggml_transpose( - struct ggml_context * ctx, - struct ggml_tensor * a) { - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - ggml_format_name(result, "%s (transposed)", a->name); - - result->ne[0] = a->ne[1]; - result->ne[1] = a->ne[0]; - - result->nb[0] = a->nb[1]; - result->nb[1] = a->nb[0]; - - result->op = GGML_OP_TRANSPOSE; - result->src[0] = a; - - return result; -} - -// ggml_get_rows - -struct ggml_tensor * ggml_get_rows( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(a->ne[2] == b->ne[1]); - GGML_ASSERT(a->ne[3] == b->ne[2]); - GGML_ASSERT(b->ne[3] == 1); - GGML_ASSERT(b->type == GGML_TYPE_I32); - - // TODO: implement non F32 return - enum ggml_type type = GGML_TYPE_F32; - if (a->type == GGML_TYPE_I32) { - type = a->type; - } - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, type, a->ne[0], b->ne[0], b->ne[1], b->ne[2]); - - result->op = GGML_OP_GET_ROWS; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_get_rows_back - -struct ggml_tensor * ggml_get_rows_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c) { - GGML_ASSERT(ggml_is_matrix(a) && ggml_is_vector(b) && b->type == GGML_TYPE_I32); - GGML_ASSERT(ggml_is_matrix(c) && (a->ne[0] == c->ne[0])); - - // TODO: implement non F32 return - //struct ggml_tensor * result = ggml_new_tensor_2d(ctx, a->type, a->ne[0], b->ne[0]); - struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, c->ne[0], c->ne[1]); - - result->op = GGML_OP_GET_ROWS_BACK; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_set_rows - -struct ggml_tensor * ggml_set_rows( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c) { - GGML_ASSERT(a->ne[0] == b->ne[0]); - GGML_ASSERT(a->ne[2] == b->ne[2]); - GGML_ASSERT(a->ne[3] == b->ne[3]); - GGML_ASSERT(b->ne[1] == c->ne[0]); - GGML_ASSERT(b->ne[2] % c->ne[1] == 0); - GGML_ASSERT(b->ne[3] % c->ne[2] == 0); - GGML_ASSERT(c->ne[3] == 1); - GGML_ASSERT(b->type == GGML_TYPE_F32); - GGML_ASSERT(c->type == GGML_TYPE_I64 || c->type == GGML_TYPE_I32); - - GGML_ASSERT(ggml_is_contiguous_rows(a)); - GGML_ASSERT(ggml_is_contiguous_rows(b)); - - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - - result->op = GGML_OP_SET_ROWS; - result->src[0] = b; - result->src[1] = c; - result->src[2] = a; // note: order is weird due to legacy reasons (https://github.com/ggml-org/llama.cpp/pull/16063#discussion_r2385795931) - - return result; -} - -// ggml_diag - -struct ggml_tensor * ggml_diag( - struct ggml_context * ctx, - struct ggml_tensor * a) { - GGML_ASSERT(a->ne[1] == 1); - - const int64_t ne[4] = { a->ne[0], a->ne[0], a->ne[2], a->ne[3] }; - struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, 4, ne); - - result->op = GGML_OP_DIAG; - result->src[0] = a; - - return result; -} - -// ggml_diag_mask_inf - -static struct ggml_tensor * ggml_diag_mask_inf_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - int32_t params[] = { n_past }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_DIAG_MASK_INF; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_diag_mask_inf( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past) { - return ggml_diag_mask_inf_impl(ctx, a, n_past, false); -} - -struct ggml_tensor * ggml_diag_mask_inf_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past) { - return ggml_diag_mask_inf_impl(ctx, a, n_past, true); -} - -// ggml_diag_mask_zero - -static struct ggml_tensor * ggml_diag_mask_zero_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - int32_t params[] = { n_past }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_DIAG_MASK_ZERO; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_diag_mask_zero( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past) { - return ggml_diag_mask_zero_impl(ctx, a, n_past, false); -} - -struct ggml_tensor * ggml_diag_mask_zero_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past) { - return ggml_diag_mask_zero_impl(ctx, a, n_past, true); -} - -// ggml_soft_max - -static struct ggml_tensor * ggml_soft_max_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * mask, - float scale, - float max_bias, - bool inplace) { - GGML_ASSERT(ggml_is_contiguous(a)); - - if (mask) { - GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_contiguous(mask)); - GGML_ASSERT(mask->ne[0] == a->ne[0]); - GGML_ASSERT(mask->ne[1] >= a->ne[1]); - GGML_ASSERT(a->ne[2]%mask->ne[2] == 0); - GGML_ASSERT(a->ne[3]%mask->ne[3] == 0); - } - - if (max_bias > 0.0f) { - GGML_ASSERT(mask); - } - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - float params[] = { scale, max_bias }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_SOFT_MAX; - result->src[0] = a; - result->src[1] = mask; - - return result; -} - -struct ggml_tensor * ggml_soft_max( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_soft_max_impl(ctx, a, NULL, 1.0f, 0.0f, false); -} - -struct ggml_tensor * ggml_soft_max_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_soft_max_impl(ctx, a, NULL, 1.0f, 0.0f, true); -} - -struct ggml_tensor * ggml_soft_max_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * mask, - float scale, - float max_bias) { - return ggml_soft_max_impl(ctx, a, mask, scale, max_bias, false); -} - -struct ggml_tensor * ggml_soft_max_ext_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * mask, - float scale, - float max_bias) { - return ggml_soft_max_impl(ctx, a, mask, scale, max_bias, true); -} - -void ggml_soft_max_add_sinks( - struct ggml_tensor * a, - struct ggml_tensor * sinks) { - if (!sinks) { - a->src[2] = NULL; - return; - } - - GGML_ASSERT(a->op == GGML_OP_SOFT_MAX); - GGML_ASSERT(a->src[2] == NULL); - GGML_ASSERT(a->src[0]->ne[2] == sinks->ne[0]); - GGML_ASSERT(sinks->type == GGML_TYPE_F32); - - a->src[2] = sinks; -} - -// ggml_soft_max_ext_back - -static struct ggml_tensor * ggml_soft_max_ext_back_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float scale, - float max_bias, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_SOFT_MAX_BACK; - result->src[0] = a; - result->src[1] = b; - - memcpy((float *) result->op_params + 0, &scale, sizeof(float)); - memcpy((float *) result->op_params + 1, &max_bias, sizeof(float)); - - return result; -} - -struct ggml_tensor * ggml_soft_max_ext_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float scale, - float max_bias) { - return ggml_soft_max_ext_back_impl(ctx, a, b, scale, max_bias, false); -} - -struct ggml_tensor * ggml_soft_max_ext_back_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float scale, - float max_bias) { - return ggml_soft_max_ext_back_impl(ctx, a, b, scale, max_bias, true); -} - -// ggml_rope - -static struct ggml_tensor * ggml_rope_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int sections[GGML_MROPE_SECTIONS], - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - bool inplace) { - GGML_ASSERT((mode & 1) == 0 && "mode & 1 == 1 is no longer supported"); - - GGML_ASSERT(ggml_is_vector(b)); - GGML_ASSERT(b->type == GGML_TYPE_I32); - - bool mrope_used = mode & GGML_ROPE_TYPE_MROPE; - if (mrope_used) { - GGML_ASSERT(a->ne[2] * 4 == b->ne[0]); // mrope expecting 4 position ids per token - } else { - GGML_ASSERT(a->ne[2] == b->ne[0]); - } - - if (c) { - GGML_ASSERT(c->type == GGML_TYPE_F32); - GGML_ASSERT(c->ne[0] >= n_dims / 2); - } - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - int32_t params[15] = { /*n_past*/ 0, n_dims, mode, /*n_ctx*/ 0, n_ctx_orig }; - memcpy(params + 5, &freq_base, sizeof(float)); - memcpy(params + 6, &freq_scale, sizeof(float)); - memcpy(params + 7, &ext_factor, sizeof(float)); - memcpy(params + 8, &attn_factor, sizeof(float)); - memcpy(params + 9, &beta_fast, sizeof(float)); - memcpy(params + 10, &beta_slow, sizeof(float)); - if (mrope_used && sections) { - memcpy(params + 11, sections, sizeof(int32_t) * GGML_MROPE_SECTIONS); - } else { - memset(params + 11, 0, sizeof(int32_t) * GGML_MROPE_SECTIONS); - } - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_ROPE; - result->src[0] = a; - result->src[1] = b; - result->src[2] = c; - - return result; -} - -struct ggml_tensor * ggml_rope( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int n_dims, - int mode) { - return ggml_rope_impl( - ctx, a, b, NULL, n_dims, NULL, mode, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f, false - ); -} - -struct ggml_tensor * ggml_rope_multi( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int sections[GGML_MROPE_SECTIONS], - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ggml_rope_impl( - ctx, a, b, c, n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow, false - ); -} - -struct ggml_tensor * ggml_rope_multi_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int sections[GGML_MROPE_SECTIONS], - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ggml_rope_impl( - ctx, a, b, c, n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow, true - ); -} - -struct ggml_tensor * ggml_rope_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int n_dims, - int mode) { - return ggml_rope_impl( - ctx, a, b, NULL, n_dims, NULL, mode, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f, true - ); -} - -struct ggml_tensor * ggml_rope_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ggml_rope_impl( - ctx, a, b, c, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow, false - ); -} - -struct ggml_tensor * ggml_rope_ext_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ggml_rope_impl( - ctx, a, b, c, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow, true - ); -} - -struct ggml_tensor * ggml_rope_custom( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ggml_rope_impl( - ctx, a, b, NULL, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow, false - ); -} - -struct ggml_tensor * ggml_rope_custom_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ggml_rope_impl( - ctx, a, b, NULL, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow, true - ); -} - -// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get -// `corr_dim(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))` -static float ggml_rope_yarn_corr_dim(int n_dims, int n_ctx_orig, float n_rot, float base) { - return n_dims * logf(n_ctx_orig / (n_rot * 2 * (float)M_PI)) / (2 * logf(base)); -} - -void ggml_rope_yarn_corr_dims( - int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2] -) { - // start and end correction dims - float start = floorf(ggml_rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_fast, freq_base)); - float end = ceilf(ggml_rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_slow, freq_base)); - dims[0] = MAX(0, start); - dims[1] = MIN(n_dims - 1, end); -} - -// ggml_rope_back - -struct ggml_tensor * ggml_rope_ext_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - struct ggml_tensor * result = ggml_rope_ext( - ctx, a, b, c, n_dims, mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - result->op = GGML_OP_ROPE_BACK; - return result; -} - -struct ggml_tensor * ggml_rope_multi_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int sections[4], - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - struct ggml_tensor * result = ggml_rope_multi( - ctx, a, b, c, n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - result->op = GGML_OP_ROPE_BACK; - return result; -} -// ggml_clamp - -struct ggml_tensor * ggml_clamp( - struct ggml_context * ctx, - struct ggml_tensor * a, - float min, - float max) { - // TODO: when implement backward, fix this: - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - - float params[] = { min, max }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_CLAMP; - result->src[0] = a; - - return result; -} - -static int64_t ggml_calc_conv_output_size(int64_t ins, int64_t ks, int s, int p, int d) { - return (ins + 2 * p - d * (ks - 1) - 1) / s + 1; -} - -// im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW] -// a: [OC，IC, KH, KW] -// b: [N, IC, IH, IW] -// result: [N, OH, OW, IC*KH*KW] -struct ggml_tensor * ggml_im2col( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int s1, - int p0, - int p1, - int d0, - int d1, - bool is_2D, - enum ggml_type dst_type) { - if (is_2D) { - GGML_ASSERT(a->ne[2] == b->ne[2]); - } else { - //GGML_ASSERT(b->ne[1] % a->ne[1] == 0); - GGML_ASSERT(b->ne[1] == a->ne[1]); - GGML_ASSERT(b->ne[3] == 1); - } - - const int64_t OH = is_2D ? ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1) : 0; - const int64_t OW = ggml_calc_conv_output_size(b->ne[0], a->ne[0], s0, p0, d0); - - GGML_ASSERT((!is_2D || OH > 0) && "b too small compared to a"); - GGML_ASSERT((OW > 0) && "b too small compared to a"); - - const int64_t ne[4] = { - is_2D ? (a->ne[2] * a->ne[1] * a->ne[0]) : a->ne[1] * a->ne[0], - OW, - is_2D ? OH : b->ne[2], - is_2D ? b->ne[3] : 1, - }; - - struct ggml_tensor * result = ggml_new_tensor(ctx, dst_type, 4, ne); - int32_t params[] = { s0, s1, p0, p1, d0, d1, (is_2D ? 1 : 0) }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_IM2COL; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_im2col_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int64_t * ne, - int s0, - int s1, - int p0, - int p1, - int d0, - int d1, - bool is_2D) { - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - int32_t params[] = { s0, s1, p0, p1, d0, d1, (is_2D ? 1 : 0) }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_IM2COL_BACK; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_conv_1d - -struct ggml_tensor * ggml_conv_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int p0, - int d0) { - struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16); // [N, OL, IC * K] - - struct ggml_tensor * result = - ggml_mul_mat(ctx, - ggml_reshape_2d(ctx, im2col, im2col->ne[0], (im2col->ne[2] * im2col->ne[1])), // [N, OL, IC * K] => [N*OL, IC * K] - ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1]), a->ne[2])); // [OC，IC, K] => [OC, IC * K] - - result = ggml_reshape_3d(ctx, result, im2col->ne[1], a->ne[2], im2col->ne[2]); // [N, OC, OL] - - return result; -} - -// ggml_conv_1d_ph - -struct ggml_tensor* ggml_conv_1d_ph( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s, - int d) { - return ggml_conv_1d(ctx, a, b, s, a->ne[0] / 2, d); -} - -// ggml_conv_1d_dw - -struct ggml_tensor * ggml_conv_1d_dw( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int p0, - int d0) { - struct ggml_tensor * new_b = ggml_reshape_4d(ctx, b, b->ne[0], 1, b->ne[1], b->ne[2]); - - struct ggml_tensor * im2col = ggml_im2col(ctx, a, new_b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16); - - struct ggml_tensor * result = ggml_mul_mat(ctx, im2col, a); - - result = ggml_reshape_3d(ctx, result, result->ne[0], result->ne[2], 1); - - return result; -} - -// ggml_conv_1d_dw_ph - -struct ggml_tensor * ggml_conv_1d_dw_ph( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int d0) { - return ggml_conv_1d_dw(ctx, a, b, s0, a->ne[0] / 2, d0); -} - -// ggml_conv_transpose_1d - -static int64_t ggml_calc_conv_transpose_1d_output_size(int64_t ins, int64_t ks, int s, int p, int d) { - return (ins - 1) * s - 2 * p + d * (ks - 1) + 1; -} - -GGML_API struct ggml_tensor * ggml_conv_transpose_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int p0, - int d0) { - GGML_ASSERT(ggml_is_matrix(b)); - GGML_ASSERT(a->ne[2] == b->ne[1]); - GGML_ASSERT(a->ne[3] == 1); - - GGML_ASSERT(p0 == 0); - GGML_ASSERT(d0 == 1); - - const int64_t ne[4] = { - ggml_calc_conv_transpose_1d_output_size(b->ne[0], a->ne[0], s0, 0 /*p0*/, 1 /*d0*/), - a->ne[1], b->ne[2], 1, - }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - int32_t params[] = { s0, p0, d0 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_CONV_TRANSPOSE_1D; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_conv_2d - -// a: [OC，IC, KH, KW] -// b: [N, IC, IH, IW] -// result: [N, OC, OH, OW] -struct ggml_tensor * ggml_conv_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int s1, - int p0, - int p1, - int d0, - int d1) { - struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, s1, p0, p1, d0, d1, true, a->type); // [N, OH, OW, IC * KH * KW] - - struct ggml_tensor * result = - ggml_mul_mat(ctx, - ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), // [N, OH, OW, IC * KH * KW] => [N*OH*OW, IC * KH * KW] - ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1] * a->ne[2]), a->ne[3])); // [OC，IC, KH, KW] => [OC, IC * KH * KW] - - result = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], im2col->ne[3], a->ne[3]); // [OC, N, OH, OW] - result = ggml_cont(ctx, ggml_permute(ctx, result, 0, 1, 3, 2)); // [N, OC, OH, OW] - - - return result; -} - -// a: [OC*IC, KD, KH, KW] -// b: [N*IC, ID, IH, IW] -// result: [N*OD, OH, OW, IC * KD * KH * KW] -struct ggml_tensor * ggml_im2col_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int64_t IC, - int s0, // stride width - int s1, // stride height - int s2, // stride depth - int p0, // padding width - int p1, // padding height - int p2, // padding depth - int d0, // dilation width - int d1, // dilation height - int d2, // dilation depth - enum ggml_type dst_type) { - const int64_t N = b->ne[3] / IC; - const int64_t ID = b->ne[2]; - const int64_t IH = b->ne[1]; - const int64_t IW = b->ne[0]; - - const int64_t OC = a->ne[3] / IC; - UNUSED(OC); - const int64_t KD = a->ne[2]; - const int64_t KH = a->ne[1]; - const int64_t KW = a->ne[0]; - const int64_t OD = ggml_calc_conv_output_size(ID, KD, s2, p2, d2); - const int64_t OH = ggml_calc_conv_output_size(IH, KH, s1, p1, d1); - const int64_t OW = ggml_calc_conv_output_size(IW, KW, s0, p0, d0); - - GGML_ASSERT((OD > 0) && "b too small compared to a"); - GGML_ASSERT((OH > 0) && "b too small compared to a"); - GGML_ASSERT((OW > 0) && "b too small compared to a"); - - - const int64_t ne[4] = {KW*KH*KD*IC, OW, OH, OD*N}; - - struct ggml_tensor * result = ggml_new_tensor(ctx, dst_type, 4, ne); - int32_t params[] = { s0, s1, s2, p0, p1, p2, d0, d1, d2, (int32_t)IC}; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_IM2COL_3D; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// a: [OC*IC, KD, KH, KW] -// b: [N*IC, ID, IH, IW] -// result: [N*OC, OD, OH, OW] -struct ggml_tensor * ggml_conv_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int64_t IC, - int s0, // stride width - int s1, // stride height - int s2, // stride depth - int p0, // padding width - int p1, // padding height - int p2, // padding depth - int d0, // dilation width - int d1, // dilation height - int d2 // dilation depth - ) { - struct ggml_tensor * im2col = ggml_im2col_3d(ctx, a, b, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, a->type); // [N*OD, OH, OW, IC * KD * KH * KW] - - int64_t OC = a->ne[3] / IC; - int64_t N = b->ne[3] / IC; - struct ggml_tensor * result = - ggml_mul_mat(ctx, - ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), // [N*OD, OH, OW, IC * KD * KH * KW] => [N*OD*OH*OW, IC * KD * KH * KW] - ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1] * a->ne[2] * IC), OC)); // [OC*IC, KD, KH, KW] => [OC, IC * KD * KH * KW] - - int64_t OD = im2col->ne[3] / N; - result = ggml_reshape_4d(ctx, result, im2col->ne[1]*im2col->ne[2], OD, N, OC); // [OC, N*OD*OH*OW] => [OC, N, OD, OH*OW] - result = ggml_cont(ctx, ggml_permute(ctx, result, 0, 1, 3, 2)); // [N, OC, OD, OH*OW] - result = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], OD, OC * N); // [N*OC, OD, OH, OW] - - return result; -} - -// ggml_conv_2d_sk_p0 - -struct ggml_tensor * ggml_conv_2d_sk_p0( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_conv_2d(ctx, a, b, a->ne[0], a->ne[1], 0, 0, 1, 1); -} - -// ggml_conv_2d_s1_ph - -struct ggml_tensor * ggml_conv_2d_s1_ph( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_conv_2d(ctx, a, b, 1, 1, a->ne[0] / 2, a->ne[1] / 2, 1, 1); -} - -// ggml_conv_2d_dw - -struct ggml_tensor * ggml_conv_2d_dw( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int s1, - int p0, - int p1, - int d0, - int d1) { - struct ggml_tensor * new_a = ggml_reshape_4d(ctx, a, a->ne[0], a->ne[1], 1, a->ne[2] * a->ne[3]); - struct ggml_tensor * im2col = ggml_im2col(ctx, new_a, - ggml_reshape_4d(ctx, b, b->ne[0], b->ne[1], 1, b->ne[2] * b->ne[3]), - s0, s1, p0, p1, d0, d1, true, GGML_TYPE_F16); // [N * IC, OH, OW, KH * KW] - struct ggml_tensor * new_b = ggml_reshape_4d(ctx, im2col, im2col->ne[0], im2col->ne[2] * im2col->ne[1], b->ne[2], b->ne[3]); // [N * IC, OH, OW, KH * KW] => [N, IC, OH * OW, KH * KW] - - new_a = ggml_reshape_4d(ctx, new_a, (new_a->ne[0] * new_a->ne[1]), new_a->ne[2], new_a->ne[3], 1); // [OC，1, KH, KW] => [1, OC, 1, KH * KW] - struct ggml_tensor * result = ggml_mul_mat(ctx, new_a, new_b); - result = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], b->ne[2], b->ne[3]); // [N, OC, OH, OW] - - return result; -} - -// ggml_conv_2d_dw_direct - -struct ggml_tensor * ggml_conv_2d_dw_direct( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int stride0, - int stride1, - int pad0, - int pad1, - int dilation0, - int dilation1) { - GGML_ASSERT(a->ne[2] == 1); - GGML_ASSERT(a->ne[3] == b->ne[2]); - int64_t ne[4]; - ne[0] = ggml_calc_conv_output_size(b->ne[0], a->ne[0], stride0, pad0, dilation0); - ne[1] = ggml_calc_conv_output_size(b->ne[1], a->ne[1], stride1, pad1, dilation1); - ne[2] = b->ne[2]; - ne[3] = b->ne[3]; - - struct ggml_tensor * result = ggml_new_tensor(ctx, b->type, 4, ne); - - if (ggml_is_contiguous_channels(b)) { - // Result will be permuted the same way as input (CWHN order) - const int64_t type_size = ggml_type_size(result->type); - GGML_ASSERT(ggml_blck_size(result->type) == 1); - result->nb[0] = result->ne[2] * type_size; - result->nb[1] = result->ne[0] * result->nb[0]; - result->nb[2] = type_size; - } - - int32_t params[] = { stride0, stride1, pad0, pad1, dilation0, dilation1 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_CONV_2D_DW; - result->src[0] = a; - result->src[1] = b; - return result; -} - -// ggml_conv_2d_direct - -struct ggml_tensor * ggml_conv_2d_direct( - struct ggml_context * ctx, - struct ggml_tensor * a, // convolution kernel [KW, KH, IC, OC] - struct ggml_tensor * b, // input data [W, H, C, N] - int s0, // stride dimension 0 - int s1, // stride dimension 1 - int p0, // padding dimension 0 - int p1, // padding dimension 1 - int d0, // dilation dimension 0 - int d1) {// dilation dimension 1 - - GGML_ASSERT(a->ne[2] == b->ne[2]); - //GGML_ASSERT(a->type == b->type); - - int64_t ne[4]; - ne[0] = ggml_calc_conv_output_size(b->ne[0], a->ne[0], s0, p0, d0); - ne[1] = ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1); - ne[2] = a->ne[3]; - ne[3] = b->ne[3]; - - struct ggml_tensor * result = ggml_new_tensor(ctx, b->type, 4, ne); - - ggml_set_op_params_i32(result, 0, s0); - ggml_set_op_params_i32(result, 1, s1); - ggml_set_op_params_i32(result, 2, p0); - ggml_set_op_params_i32(result, 3, p1); - ggml_set_op_params_i32(result, 4, d0); - ggml_set_op_params_i32(result, 5, d1); - - result->op = GGML_OP_CONV_2D; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_conv_3d_direct - -struct ggml_tensor * ggml_conv_3d_direct( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int s1, - int s2, - int p0, - int p1, - int p2, - int d0, - int d1, - int d2, - int c, - int n, - int oc) { - - GGML_ASSERT(a->ne[3] == (int64_t) c * oc); - GGML_ASSERT(b->ne[3] == (int64_t) c * n); - - int64_t ne[4]; - ne[0] = ggml_calc_conv_output_size(b->ne[0], a->ne[0], s0, p0, d0); - ne[1] = ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1); - ne[2] = ggml_calc_conv_output_size(b->ne[2], a->ne[2], s2, p2, d2); - ne[3] = (int64_t) oc * n; - - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - ggml_set_op_params_i32(result, 0, s0); - ggml_set_op_params_i32(result, 1, s1); - ggml_set_op_params_i32(result, 2, s2); - ggml_set_op_params_i32(result, 3, p0); - ggml_set_op_params_i32(result, 4, p1); - ggml_set_op_params_i32(result, 5, p2); - ggml_set_op_params_i32(result, 6, d0); - ggml_set_op_params_i32(result, 7, d1); - ggml_set_op_params_i32(result, 8, d2); - ggml_set_op_params_i32(result, 9, c); - ggml_set_op_params_i32(result, 10, n); - ggml_set_op_params_i32(result, 11, oc); - - result->op = GGML_OP_CONV_3D; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_conv_transpose_2d_p0 - -static int64_t ggml_calc_conv_transpose_output_size(int64_t ins, int64_t ks, int s, int p) { - return (ins - 1) * s - 2 * p + ks; -} - -struct ggml_tensor * ggml_conv_transpose_2d_p0( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int stride) { - GGML_ASSERT(a->ne[3] == b->ne[2]); - - const int64_t ne[4] = { - ggml_calc_conv_transpose_output_size(b->ne[0], a->ne[0], stride, 0 /*p0*/), - ggml_calc_conv_transpose_output_size(b->ne[1], a->ne[1], stride, 0 /*p1*/), - a->ne[2], b->ne[3], - }; - - struct ggml_tensor* result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - ggml_set_op_params_i32(result, 0, stride); - - result->op = GGML_OP_CONV_TRANSPOSE_2D; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_pool_* - -static int64_t ggml_calc_pool_output_size(int64_t ins, int ks, int s, float p) { - return (ins + 2 * p - ks) / s + 1; -} - -// ggml_pool_1d - -struct ggml_tensor * ggml_pool_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_op_pool op, - int k0, - int s0, - int p0) { - const int64_t ne[4] = { - ggml_calc_pool_output_size(a->ne[0], k0, s0, p0), - a->ne[1], - a->ne[2], - a->ne[3], - }; - GGML_ASSERT(ne[0] > 0); - - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - int32_t params[] = { op, k0, s0, p0 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_POOL_1D; - result->src[0] = a; - - return result; -} - -// ggml_pool_2d - -struct ggml_tensor * ggml_pool_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_op_pool op, - int k0, - int k1, - int s0, - int s1, - float p0, - float p1) { - struct ggml_tensor * result; - const int64_t ne[4] = { - ggml_calc_pool_output_size(a->ne[0], k0, s0, p0), - ggml_calc_pool_output_size(a->ne[1], k1, s1, p1), - a->ne[2], - a->ne[3], - }; - GGML_ASSERT(ne[0] > 0); - GGML_ASSERT(ne[1] > 0); - - result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - int32_t params[] = { op, k0, k1, s0, s1, p0, p1 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_POOL_2D; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_pool_2d_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * af, - enum ggml_op_pool op, - int k0, - int k1, - int s0, - int s1, - float p0, - float p1) { - struct ggml_tensor * result; - result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, af->ne); - - int32_t params[] = { op, k0, k1, s0, s1, p0, p1 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_POOL_2D_BACK; - result->src[0] = a; - result->src[1] = af; - - return result; -} - -// ggml_upscale / ggml_interpolate - -static struct ggml_tensor * ggml_interpolate_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3, - uint32_t mode) { - GGML_ASSERT((mode & 0xFF) < GGML_SCALE_MODE_COUNT); - // TODO: implement antialias for modes other than bilinear - GGML_ASSERT(!(mode & GGML_SCALE_FLAG_ANTIALIAS) || (mode & 0xFF) == GGML_SCALE_MODE_BILINEAR); - GGML_ASSERT(a->type == GGML_TYPE_F32); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, ne0, ne1, ne2, ne3); - - ggml_set_op_params_i32(result, 0, (int32_t)mode); - - result->op = GGML_OP_UPSCALE; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_upscale( - struct ggml_context * ctx, - struct ggml_tensor * a, - int scale_factor, - enum ggml_scale_mode mode) { - GGML_ASSERT(scale_factor > 1); - return ggml_interpolate_impl(ctx, a, a->ne[0] * scale_factor, a->ne[1] * scale_factor, a->ne[2], a->ne[3], mode); -} - -struct ggml_tensor * ggml_upscale_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - int ne0, - int ne1, - int ne2, - int ne3, - enum ggml_scale_mode mode) { - return ggml_interpolate_impl(ctx, a, ne0, ne1, ne2, ne3, mode); -} - -struct ggml_tensor * ggml_interpolate( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3, - uint32_t mode) { - return ggml_interpolate_impl(ctx, a, ne0, ne1, ne2, ne3, mode); -} - -// ggml_pad - -struct ggml_tensor * ggml_pad( - struct ggml_context * ctx, - struct ggml_tensor * a, - int p0, - int p1, - int p2, - int p3) { - return ggml_pad_ext(ctx, a, 0, p0, 0, p1, 0, p2, 0, p3); -} - -// ggml_pad_circular - -struct ggml_tensor * ggml_pad_circular( - struct ggml_context * ctx, - struct ggml_tensor * a, - int p0, - int p1, - int p2, - int p3) { - return ggml_pad_ext_circular(ctx, a, 0, p0, 0, p1, 0, p2, 0, p3); -} - -struct ggml_tensor * ggml_pad_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - int lp0, - int rp0, - int lp1, - int rp1, - int lp2, - int rp2, - int lp3, - int rp3 - ) { - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, - a->ne[0] + lp0 + rp0, - a->ne[1] + lp1 + rp1, - a->ne[2] + lp2 + rp2, - a->ne[3] + lp3 + rp3); - - ggml_set_op_params_i32(result, 0, lp0); - ggml_set_op_params_i32(result, 1, rp0); - ggml_set_op_params_i32(result, 2, lp1); - ggml_set_op_params_i32(result, 3, rp1); - ggml_set_op_params_i32(result, 4, lp2); - ggml_set_op_params_i32(result, 5, rp2); - ggml_set_op_params_i32(result, 6, lp3); - ggml_set_op_params_i32(result, 7, rp3); - ggml_set_op_params_i32(result, 8, 0); // not circular by default - - - result->op = GGML_OP_PAD; - result->src[0] = a; - - return result; -} - -// ggml_pad_ext_circular - -struct ggml_tensor * ggml_pad_ext_circular( - struct ggml_context * ctx, - struct ggml_tensor * a, - int lp0, - int rp0, - int lp1, - int rp1, - int lp2, - int rp2, - int lp3, - int rp3 - ) { - struct ggml_tensor * result = ggml_pad_ext(ctx, a, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); - ggml_set_op_params_i32(result, 8, 1); // circular - return result; -} - -// ggml_pad_reflect_1d - -struct ggml_tensor * ggml_pad_reflect_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int p0, - int p1) { - GGML_ASSERT(p0 >= 0); - GGML_ASSERT(p1 >= 0); - - GGML_ASSERT(p0 < a->ne[0]); // padding length on each size must be less than the - GGML_ASSERT(p1 < a->ne[0]); // existing length of the dimension being padded - - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(a->type == GGML_TYPE_F32); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, - a->ne[0] + p0 + p1, - a->ne[1], - a->ne[2], - a->ne[3]); - - int32_t params[] = { p0, p1 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_PAD_REFLECT_1D; - result->src[0] = a; - - return result; -} - -// ggml_roll - -struct ggml_tensor * ggml_roll( - struct ggml_context * ctx, - struct ggml_tensor * a, - int shift0, - int shift1, - int shift2, - int shift3) { - GGML_ASSERT(a->nb[0] == ggml_type_size(a->type)); - GGML_ASSERT(abs(shift0) < a->ne[0]); - GGML_ASSERT(abs(shift1) < a->ne[1]); - GGML_ASSERT(abs(shift2) < a->ne[2]); - GGML_ASSERT(abs(shift3) < a->ne[3]); - - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - ggml_set_op_params_i32(result, 0, shift0); - ggml_set_op_params_i32(result, 1, shift1); - ggml_set_op_params_i32(result, 2, shift2); - ggml_set_op_params_i32(result, 3, shift3); - - result->op = GGML_OP_ROLL; - result->src[0] = a; - - return result; -} - -// ggml_timestep_embedding - -struct ggml_tensor * ggml_timestep_embedding( - struct ggml_context * ctx, - struct ggml_tensor * timesteps, - int dim, - int max_period) { - - struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, timesteps->ne[0]); - - ggml_set_op_params_i32(result, 0, dim); - ggml_set_op_params_i32(result, 1, max_period); - - result->op = GGML_OP_TIMESTEP_EMBEDDING; - result->src[0] = timesteps; - - return result; -} - -// ggml_tri - -struct ggml_tensor * ggml_tri( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_tri_type type) { - GGML_ASSERT(a->type == GGML_TYPE_F32); - - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(a->ne[0] == a->ne[1]); - - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - ggml_set_op_params_i32(result, 0, type); - - result->op = GGML_OP_TRI; - result->src[0] = a; - - return result; -} - -// ggml_fill - -static struct ggml_tensor * ggml_fill_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - float c, - bool inplace) { - GGML_ASSERT(a->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_contiguous(a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params_f32(result, 0, c); - - result->op = GGML_OP_FILL; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_fill( - struct ggml_context * ctx, - struct ggml_tensor * a, - float c) { - return ggml_fill_impl(ctx, a, c, false); -} - -struct ggml_tensor * ggml_fill_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float c) { - return ggml_fill_impl(ctx, a, c, true); -} - -// ggml_argsort - -struct ggml_tensor * ggml_argsort( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_sort_order order) { - GGML_ASSERT(a->ne[0] <= INT32_MAX); - - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_I32, GGML_MAX_DIMS, a->ne); - - ggml_set_op_params_i32(result, 0, (int32_t) order); - - result->op = GGML_OP_ARGSORT; - result->src[0] = a; - - return result; -} - -// ggml_argsort_top_k - -struct ggml_tensor * ggml_argsort_top_k( - struct ggml_context * ctx, - struct ggml_tensor * a, - int k) { - GGML_ASSERT(a->ne[0] >= k); - - struct ggml_tensor * result = ggml_argsort(ctx, a, GGML_SORT_ORDER_DESC); - - result = ggml_view_4d(ctx, result, - k, result->ne[1], result->ne[2], result->ne[3], - result->nb[1], result->nb[2], result->nb[3], - 0); - - return result; -} - -// ggml_top_k - -struct ggml_tensor * ggml_top_k( - struct ggml_context * ctx, - struct ggml_tensor * a, - int k) { - GGML_ASSERT(a->ne[0] >= k); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, k, a->ne[1], a->ne[2], a->ne[3]); - - result->op = GGML_OP_TOP_K; - result->src[0] = a; - - return result; -} - -// ggml_arange - -struct ggml_tensor * ggml_arange( - struct ggml_context * ctx, - float start, - float stop, - float step) { - GGML_ASSERT(stop > start); - - const int64_t steps = (int64_t) ceilf((stop - start) / step); - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, steps); - - ggml_set_op_params_f32(result, 0, start); - ggml_set_op_params_f32(result, 1, stop); - ggml_set_op_params_f32(result, 2, step); - - result->op = GGML_OP_ARANGE; - - return result; -} - -// ggml_flash_attn_ext - -struct ggml_tensor * ggml_flash_attn_ext( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * mask, - float scale, - float max_bias, - float logit_softcap) { - GGML_ASSERT(ggml_can_mul_mat(k, q)); - // TODO: check if vT can be multiplied by (k*qT) - - GGML_ASSERT(q->ne[3] == k->ne[3]); - GGML_ASSERT(q->ne[3] == v->ne[3]); - - if (mask) { - GGML_ASSERT(mask->type == GGML_TYPE_F16); - GGML_ASSERT(ggml_is_contiguous(mask)); - //GGML_ASSERT(ggml_can_repeat_rows(mask, qk)); - - GGML_ASSERT(q->ne[2] % mask->ne[2] == 0); - GGML_ASSERT(q->ne[3] % mask->ne[3] == 0); - } - - if (max_bias > 0.0f) { - GGML_ASSERT(mask); - } - - // permute(0, 2, 1, 3) - int64_t ne[4] = { v->ne[0], q->ne[2], q->ne[1], q->ne[3] }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - float params[] = { scale, max_bias, logit_softcap }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_FLASH_ATTN_EXT; - result->src[0] = q; - result->src[1] = k; - result->src[2] = v; - result->src[3] = mask; - - return result; -} - -void ggml_flash_attn_ext_set_prec( - struct ggml_tensor * a, - enum ggml_prec prec) { - GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); - - const int32_t prec_i32 = (int32_t) prec; - - ggml_set_op_params_i32(a, 3, prec_i32); // scale is on first pos, max_bias on second -} - -enum ggml_prec ggml_flash_attn_ext_get_prec( - const struct ggml_tensor * a) { - GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); - - const int32_t prec_i32 = ggml_get_op_params_i32(a, 3); - - return (enum ggml_prec) prec_i32; -} - -void ggml_flash_attn_ext_add_sinks( - struct ggml_tensor * a, - struct ggml_tensor * sinks) { - if (!sinks) { - a->src[4] = NULL; - return; - } - - GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); - GGML_ASSERT(a->src[4] == NULL); - GGML_ASSERT(a->src[0]->ne[2] == sinks->ne[0]); - GGML_ASSERT(sinks->type == GGML_TYPE_F32); - - a->src[4] = sinks; -} - -// ggml_flash_attn_back - -struct ggml_tensor * ggml_flash_attn_back( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * d, - bool masked) { - GGML_ABORT("TODO: adapt to ggml_flash_attn_ext() changes"); - - GGML_ASSERT(ggml_can_mul_mat(k, q)); - // TODO: check if vT can be multiplied by (k*qT) - - // d shape [D,N,ne2,ne3] - // q shape [D,N,ne2,ne3] - // k shape [D,M,kvne2,ne3] - // v shape [M,D,kvne2,ne3] - - const int64_t D = q->ne[0]; - const int64_t N = q->ne[1]; - const int64_t M = k->ne[1]; - const int64_t ne2 = q->ne[2]; - const int64_t ne3 = q->ne[3]; - const int64_t kvne2 = k->ne[2]; - - GGML_ASSERT(k->ne[0] == D); - GGML_ASSERT(v->ne[0] == M); - GGML_ASSERT(v->ne[1] == D); - GGML_ASSERT(d->ne[0] == D); - GGML_ASSERT(d->ne[1] == N); - GGML_ASSERT(k->ne[2] == kvne2); - GGML_ASSERT(k->ne[3] == ne3); - GGML_ASSERT(v->ne[2] == kvne2); - GGML_ASSERT(v->ne[3] == ne3); - GGML_ASSERT(d->ne[2] == ne2); - GGML_ASSERT(d->ne[3] == ne3); - - GGML_ASSERT(ne2 % kvne2 == 0); - - // store gradients of q, k and v as continuous tensors concatenated in result. - // note: v and gradv are actually transposed, i.e. v->ne[0] != D. - const int64_t elem_q = ggml_nelements(q); - const int64_t elem_k = ggml_nelements(k); - const int64_t elem_v = ggml_nelements(v); - - enum ggml_type result_type = GGML_TYPE_F32; - GGML_ASSERT(ggml_blck_size(result_type) == 1); - const size_t tsize = ggml_type_size(result_type); - - const size_t offs_q = 0; - const size_t offs_k = offs_q + GGML_PAD(elem_q * tsize, GGML_MEM_ALIGN); - const size_t offs_v = offs_k + GGML_PAD(elem_k * tsize, GGML_MEM_ALIGN); - const size_t end = offs_v + GGML_PAD(elem_v * tsize, GGML_MEM_ALIGN); - - const size_t nelements = (end + tsize - 1)/tsize; - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, nelements); - - int32_t masked_i = masked ? 1 : 0; - ggml_set_op_params(result, &masked_i, sizeof(masked_i)); - - result->op = GGML_OP_FLASH_ATTN_BACK; - result->src[0] = q; - result->src[1] = k; - result->src[2] = v; - result->src[3] = d; - - return result; -} - -// ggml_ssm_conv - -struct ggml_tensor * ggml_ssm_conv( - struct ggml_context * ctx, - struct ggml_tensor * sx, - struct ggml_tensor * c) { - GGML_ASSERT(ggml_is_3d(sx)); - GGML_ASSERT(ggml_is_matrix(c)); - - const int64_t d_conv = c->ne[0]; - const int64_t d_inner = c->ne[1]; - const int64_t n_t = sx->ne[0] - d_conv + 1; // tokens per sequence - const int64_t n_s = sx->ne[2]; - - // TODO: maybe support other strides than 1? - GGML_ASSERT(sx->ne[0] == d_conv - 1 + n_t); - GGML_ASSERT(sx->ne[1] == d_inner); - GGML_ASSERT(n_t >= 0); - - struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, d_inner, n_t, n_s); - - result->op = GGML_OP_SSM_CONV; - result->src[0] = sx; - result->src[1] = c; - - return result; -} - -// ggml_ssm_scan - -struct ggml_tensor * ggml_ssm_scan( - struct ggml_context * ctx, - struct ggml_tensor * s, - struct ggml_tensor * x, - struct ggml_tensor * dt, - struct ggml_tensor * A, - struct ggml_tensor * B, - struct ggml_tensor * C, - struct ggml_tensor * ids) { - GGML_ASSERT(ggml_is_contiguous(s)); - GGML_ASSERT(ggml_is_contiguous(dt)); - GGML_ASSERT(ggml_is_contiguous(A)); - GGML_ASSERT(x->nb[0] == ggml_type_size(x->type)); - GGML_ASSERT(B->nb[0] == ggml_type_size(B->type)); - GGML_ASSERT(C->nb[0] == ggml_type_size(C->type)); - GGML_ASSERT(x->nb[1] == x->ne[0]*x->nb[0]); - GGML_ASSERT(B->nb[1] == B->ne[0]*B->nb[0]); - GGML_ASSERT(C->nb[1] == C->ne[0]*C->nb[0]); - GGML_ASSERT(ggml_are_same_shape(B, C)); - GGML_ASSERT(ids->type == GGML_TYPE_I32); - - { - const int64_t d_state = s->ne[0]; - const int64_t head_dim = x->ne[0]; - const int64_t n_head = x->ne[1]; - const int64_t n_seq_tokens = x->ne[2]; - const int64_t n_seqs = x->ne[3]; - - GGML_ASSERT(dt->ne[0] == n_head); - GGML_ASSERT(dt->ne[1] == n_seq_tokens); - GGML_ASSERT(dt->ne[2] == n_seqs); - GGML_ASSERT(ggml_is_3d(dt)); - GGML_ASSERT(s->ne[1] == head_dim); - GGML_ASSERT(s->ne[2] == n_head); - GGML_ASSERT(B->ne[0] == d_state); - GGML_ASSERT(B->ne[2] == n_seq_tokens); - GGML_ASSERT(B->ne[3] == n_seqs); - GGML_ASSERT(ids->ne[0] == n_seqs); - GGML_ASSERT(ggml_is_vector(ids)); - GGML_ASSERT(A->ne[1] == n_head); - GGML_ASSERT(ggml_is_matrix(A)); - - if (A->ne[0] != 1) { - // Mamba-1 has more granular decay factors - GGML_ASSERT(A->ne[0] == d_state); - } - } - - // concatenated y + ssm_states - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]); - - result->op = GGML_OP_SSM_SCAN; - result->src[0] = s; - result->src[1] = x; - result->src[2] = dt; - result->src[3] = A; - result->src[4] = B; - result->src[5] = C; - result->src[6] = ids; - - return result; -} - -// ggml_win_part - -struct ggml_tensor * ggml_win_part( - struct ggml_context * ctx, - struct ggml_tensor * a, - int w) { - GGML_ASSERT(a->ne[3] == 1); - GGML_ASSERT(a->type == GGML_TYPE_F32); - - // padding - const int px = (w - a->ne[1]%w)%w; - const int py = (w - a->ne[2]%w)%w; - - const int npx = (px + a->ne[1])/w; - const int npy = (py + a->ne[2])/w; - const int np = npx*npy; - - const int64_t ne[4] = { a->ne[0], w, w, np, }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - int32_t params[] = { npx, npy, w }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_WIN_PART; - result->src[0] = a; - - return result; -} - -// ggml_win_unpart - -struct ggml_tensor * ggml_win_unpart( - struct ggml_context * ctx, - struct ggml_tensor * a, - int w0, - int h0, - int w) { - GGML_ASSERT(a->type == GGML_TYPE_F32); - - const int64_t ne[4] = { a->ne[0], w0, h0, 1, }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 3, ne); - - int32_t params[] = { w }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_WIN_UNPART; - result->src[0] = a; - - return result; -} - -// ggml_get_rel_pos - -struct ggml_tensor * ggml_get_rel_pos( - struct ggml_context * ctx, - struct ggml_tensor * a, - int qh, - int kh) { - GGML_ASSERT(qh == kh); - GGML_ASSERT(2*MAX(qh, kh) - 1 == a->ne[1]); - - const int64_t ne[4] = { a->ne[0], kh, qh, 1, }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F16, 3, ne); - - result->op = GGML_OP_GET_REL_POS; - result->src[0] = a; - - return result; -} - -// ggml_add_rel_pos - -static struct ggml_tensor * ggml_add_rel_pos_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * pw, - struct ggml_tensor * ph, - bool inplace) { - GGML_ASSERT(ggml_are_same_shape(pw, ph)); - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_is_contiguous(pw)); - GGML_ASSERT(ggml_is_contiguous(ph)); - GGML_ASSERT(ph->type == GGML_TYPE_F32); - GGML_ASSERT(pw->type == GGML_TYPE_F32); - GGML_ASSERT(pw->ne[3] == a->ne[2]); - GGML_ASSERT(pw->ne[0]*pw->ne[0] == a->ne[0]); - GGML_ASSERT(pw->ne[1]*pw->ne[2] == a->ne[1]); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - ggml_set_op_params_i32(result, 0, inplace ? 1 : 0); - - result->op = GGML_OP_ADD_REL_POS; - result->src[0] = a; - result->src[1] = pw; - result->src[2] = ph; - - return result; -} - -struct ggml_tensor * ggml_add_rel_pos( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * pw, - struct ggml_tensor * ph) { - return ggml_add_rel_pos_impl(ctx, a, pw, ph, false); -} - -struct ggml_tensor * ggml_add_rel_pos_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * pw, - struct ggml_tensor * ph) { - return ggml_add_rel_pos_impl(ctx, a, pw, ph, true); -} - -// ggml_rwkv_wkv6 - -struct ggml_tensor * ggml_rwkv_wkv6( - struct ggml_context * ctx, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * r, - struct ggml_tensor * tf, - struct ggml_tensor * td, - struct ggml_tensor * state) { - GGML_ASSERT(ggml_is_contiguous(k)); - GGML_ASSERT(ggml_is_contiguous(v)); - GGML_ASSERT(ggml_is_contiguous(r)); - GGML_ASSERT(ggml_is_contiguous(tf)); - GGML_ASSERT(ggml_is_contiguous(td)); - GGML_ASSERT(ggml_is_contiguous(state)); - - const int64_t S = k->ne[0]; - const int64_t H = k->ne[1]; - const int64_t n_tokens = k->ne[2]; - const int64_t n_seqs = state->ne[1]; - { - GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); - GGML_ASSERT(r->ne[0] == S && r->ne[1] == H && r->ne[2] == n_tokens); - GGML_ASSERT(td->ne[0] == S && td->ne[1] == H && td->ne[2] == n_tokens); - GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); - } - - // concat output and new_state - const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_RWKV_WKV6; - result->src[0] = k; - result->src[1] = v; - result->src[2] = r; - result->src[3] = tf; - result->src[4] = td; - result->src[5] = state; - - return result; -} - -// ggml_gated_linear_attn - -struct ggml_tensor * ggml_gated_linear_attn( - struct ggml_context * ctx, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * q, - struct ggml_tensor * g, - struct ggml_tensor * state, - float scale) { - GGML_ASSERT(ggml_is_contiguous(k)); - GGML_ASSERT(ggml_is_contiguous(v)); - GGML_ASSERT(ggml_is_contiguous(q)); - GGML_ASSERT(ggml_is_contiguous(g)); - GGML_ASSERT(ggml_is_contiguous(state)); - - const int64_t S = k->ne[0]; - const int64_t H = k->ne[1]; - const int64_t n_tokens = k->ne[2]; - const int64_t n_seqs = state->ne[1]; - { - GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); - GGML_ASSERT(q->ne[0] == S && q->ne[1] == H && q->ne[2] == n_tokens); - GGML_ASSERT(g->ne[0] == S && g->ne[1] == H && g->ne[2] == n_tokens); - GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); - } - - // concat output and new_state - const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - ggml_set_op_params_f32(result, 0, scale); - - result->op = GGML_OP_GATED_LINEAR_ATTN; - result->src[0] = k; - result->src[1] = v; - result->src[2] = q; - result->src[3] = g; - result->src[4] = state; - - return result; -} - -// ggml_rwkv_wkv7 - -struct ggml_tensor * ggml_rwkv_wkv7( - struct ggml_context * ctx, - struct ggml_tensor * r, - struct ggml_tensor * w, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * state) { - GGML_ASSERT(ggml_is_contiguous(r)); - GGML_ASSERT(ggml_is_contiguous(w)); - GGML_ASSERT(ggml_is_contiguous(k)); - GGML_ASSERT(ggml_is_contiguous(v)); - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_is_contiguous(b)); - GGML_ASSERT(ggml_is_contiguous(state)); - - const int64_t S = k->ne[0]; - const int64_t H = k->ne[1]; - const int64_t n_tokens = k->ne[2]; - const int64_t n_seqs = state->ne[1]; - { - GGML_ASSERT(w->ne[0] == S && w->ne[1] == H && w->ne[2] == n_tokens); - GGML_ASSERT(k->ne[0] == S && k->ne[1] == H && k->ne[2] == n_tokens); - GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); - GGML_ASSERT(a->ne[0] == S && a->ne[1] == H && a->ne[2] == n_tokens); - GGML_ASSERT(b->ne[0] == S && b->ne[1] == H && b->ne[2] == n_tokens); - GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); - } - - // concat output and new_state - const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_RWKV_WKV7; - result->src[0] = r; - result->src[1] = w; - result->src[2] = k; - result->src[3] = v; - result->src[4] = a; - result->src[5] = b; - result->src[6] = state; - - return result; -} - -// ggml_unary - -static struct ggml_tensor * ggml_unary_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_unary_op op, - bool inplace) { - GGML_ASSERT(ggml_is_contiguous_rows(a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params_i32(result, 0, (int32_t) op); - - result->op = GGML_OP_UNARY; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_unary( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_unary_op op) { - return ggml_unary_impl(ctx, a, op, false); -} - -struct ggml_tensor * ggml_unary_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_unary_op op) { - return ggml_unary_impl(ctx, a, op, true); -} - -// ggml_map_custom1 - -static struct ggml_tensor * ggml_map_custom1_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - const ggml_custom1_op_t fun, - int n_tasks, - void * userdata, - bool inplace) { - GGML_ASSERT(n_tasks == GGML_N_TASKS_MAX || n_tasks > 0); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - struct ggml_map_custom1_op_params params = { - /*.fun =*/ fun, - /*.n_tasks =*/ n_tasks, - /*.userdata =*/ userdata - }; - ggml_set_op_params(result, ¶ms, sizeof(params)); - - result->op = GGML_OP_MAP_CUSTOM1; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_map_custom1( - struct ggml_context * ctx, - struct ggml_tensor * a, - const ggml_custom1_op_t fun, - int n_tasks, - void * userdata) { - return ggml_map_custom1_impl(ctx, a, fun, n_tasks, userdata, false); -} - -struct ggml_tensor * ggml_map_custom1_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - const ggml_custom1_op_t fun, - int n_tasks, - void * userdata) { - return ggml_map_custom1_impl(ctx, a, fun, n_tasks, userdata, true); -} - -// ggml_map_custom2 - -static struct ggml_tensor * ggml_map_custom2_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - const ggml_custom2_op_t fun, - int n_tasks, - void * userdata, - bool inplace) { - GGML_ASSERT(n_tasks == GGML_N_TASKS_MAX || n_tasks > 0); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - struct ggml_map_custom2_op_params params = { - /*.fun =*/ fun, - /*.n_tasks =*/ n_tasks, - /*.userdata =*/ userdata - }; - ggml_set_op_params(result, ¶ms, sizeof(params)); - - result->op = GGML_OP_MAP_CUSTOM2; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_map_custom2( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - const ggml_custom2_op_t fun, - int n_tasks, - void * userdata) { - return ggml_map_custom2_impl(ctx, a, b, fun, n_tasks, userdata, false); -} - -struct ggml_tensor * ggml_map_custom2_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - const ggml_custom2_op_t fun, - int n_tasks, - void * userdata) { - return ggml_map_custom2_impl(ctx, a, b, fun, n_tasks, userdata, true); -} - -// ggml_map_custom3 - -static struct ggml_tensor * ggml_map_custom3_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - const ggml_custom3_op_t fun, - int n_tasks, - void * userdata, - bool inplace) { - GGML_ASSERT(n_tasks == GGML_N_TASKS_MAX || n_tasks > 0); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - struct ggml_map_custom3_op_params params = { - /*.fun =*/ fun, - /*.n_tasks =*/ n_tasks, - /*.userdata =*/ userdata - }; - ggml_set_op_params(result, ¶ms, sizeof(params)); - - result->op = GGML_OP_MAP_CUSTOM3; - result->src[0] = a; - result->src[1] = b; - result->src[2] = c; - - return result; -} - -struct ggml_tensor * ggml_map_custom3( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - const ggml_custom3_op_t fun, - int n_tasks, - void * userdata) { - return ggml_map_custom3_impl(ctx, a, b, c, fun, n_tasks, userdata, false); -} - -struct ggml_tensor * ggml_map_custom3_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - const ggml_custom3_op_t fun, - int n_tasks, - void * userdata) { - return ggml_map_custom3_impl(ctx, a, b, c, fun, n_tasks, userdata, true); -} - -struct ggml_tensor * ggml_custom_4d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3, - struct ggml_tensor ** args, - int n_args, - ggml_custom_op_t fun, - int n_tasks, - void * userdata) { - - GGML_ASSERT(n_args < GGML_MAX_SRC); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, type, ne0, ne1, ne2, ne3); - - struct ggml_custom_op_params params = { - /*.fun =*/ fun, - /*.n_tasks =*/ n_tasks, - /*.userdata =*/ userdata - }; - ggml_set_op_params(result, ¶ms, sizeof(params)); - - result->op = GGML_OP_CUSTOM; - for (int i = 0; i < n_args; i++) { - result->src[i] = args[i]; - } - - return result; -} - -struct ggml_tensor * ggml_custom_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor ** args, - int n_args, - ggml_custom_op_t fun, - int n_tasks, - void * userdata) { - - GGML_ASSERT(n_args < GGML_MAX_SRC - 1); - - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - - struct ggml_custom_op_params params = { - /*.fun =*/ fun, - /*.n_tasks =*/ n_tasks, - /*.userdata =*/ userdata - }; - ggml_set_op_params(result, ¶ms, sizeof(params)); - - result->op = GGML_OP_CUSTOM; - result->src[0] = a; - for (int i = 0; i < n_args; i++) { - result->src[i + 1] = args[i]; - } - - return result; -} -// ggml_cross_entropy_loss - -struct ggml_tensor * ggml_cross_entropy_loss( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_are_same_shape(a, b)); - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, a->type, 1); - - result->op = GGML_OP_CROSS_ENTROPY_LOSS; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_cross_entropy_loss_back - -struct ggml_tensor * ggml_cross_entropy_loss_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c) { - GGML_ASSERT(ggml_is_scalar(a)); - GGML_ASSERT(ggml_are_same_shape(b, c)); - - struct ggml_tensor * result = ggml_dup_tensor(ctx, b); - - result->op = GGML_OP_CROSS_ENTROPY_LOSS_BACK; - result->src[0] = a; - result->src[1] = b; - result->src[2] = c; - - return result; -} - -// opt_step_adamw - -struct ggml_tensor * ggml_opt_step_adamw( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * grad, - struct ggml_tensor * m, - struct ggml_tensor * v, - struct ggml_tensor * adamw_params) { - GGML_ASSERT(a->flags & GGML_TENSOR_FLAG_PARAM); - GGML_ASSERT(ggml_are_same_shape(a, grad)); - GGML_ASSERT(ggml_are_same_shape(a, m)); - GGML_ASSERT(ggml_are_same_shape(a, v)); - GGML_ASSERT(adamw_params->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_nelements(adamw_params) == 7); - - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - - result->op = GGML_OP_OPT_STEP_ADAMW; - result->src[0] = a; - result->src[1] = grad; - result->src[2] = m; - result->src[3] = v; - result->src[4] = adamw_params; - - return result; -} - -// opt_step_sgd - -struct ggml_tensor * ggml_opt_step_sgd( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * grad, - struct ggml_tensor * params) { - GGML_ASSERT(a->flags & GGML_TENSOR_FLAG_PARAM); - GGML_ASSERT(ggml_are_same_shape(a, grad)); - GGML_ASSERT(params->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_nelements(params) == 2); - - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - - result->op = GGML_OP_OPT_STEP_SGD; - result->src[0] = a; - result->src[1] = grad; - result->src[2] = params; - - return result; -} - -// solve_tri - -struct ggml_tensor * ggml_solve_tri( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool left, - bool lower, - bool uni) { - GGML_ASSERT(a->type == GGML_TYPE_F32); - GGML_ASSERT(b->type == GGML_TYPE_F32); - - // A must be square and lower diagonal - GGML_ASSERT(a->ne[0] == a->ne[1]); - // B must have same outer dimension as A - GGML_ASSERT(a->ne[1] == b->ne[1]); - - // batch dimensions must be equal - GGML_ASSERT(a->ne[2] == b->ne[2]); - GGML_ASSERT(a->ne[3] == b->ne[3]); - - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_is_contiguous(b)); - - GGML_ASSERT(lower && left && !uni); // TODO: support other variants - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, b->ne[0], b->ne[1], b->ne[2], b->ne[3]); - - result->op = GGML_OP_SOLVE_TRI; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_gated_delta_net - -struct ggml_tensor * ggml_gated_delta_net( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * g, - struct ggml_tensor * beta, - struct ggml_tensor * state) { - GGML_ASSERT(ggml_is_contiguous_rows(q)); - GGML_ASSERT(ggml_is_contiguous_rows(k)); - GGML_ASSERT(ggml_is_contiguous_rows(v)); - GGML_ASSERT(ggml_is_contiguous(g)); - GGML_ASSERT(ggml_is_contiguous(beta)); - GGML_ASSERT(ggml_is_contiguous(state)); - - GGML_ASSERT(q->type == GGML_TYPE_F32); - GGML_ASSERT(k->type == GGML_TYPE_F32); - GGML_ASSERT(v->type == GGML_TYPE_F32); - GGML_ASSERT(g->type == GGML_TYPE_F32); - GGML_ASSERT(beta->type == GGML_TYPE_F32); - GGML_ASSERT(state->type == GGML_TYPE_F32); - - const int64_t S_v = v->ne[0]; - const int64_t H = v->ne[1]; - const int64_t n_tokens = v->ne[2]; - const int64_t n_seqs = v->ne[3]; - - // gate: scalar [1, H, T, B] or vector [S_v, H, T, B] (KDA) - GGML_ASSERT(g->ne[0] == 1 || g->ne[0] == S_v); - GGML_ASSERT(beta->ne[0] == 1); - - GGML_ASSERT(ggml_nelements(state) == S_v * S_v * H * n_seqs); - - // concat output and new_state into a single tensor - // output: S_v * H * n_tokens * n_seqs, state: S_v * S_v * H * n_seqs - const int64_t ne[4] = { S_v * H, n_tokens * n_seqs + S_v * n_seqs, 1, 1 }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_GATED_DELTA_NET; - result->src[0] = q; - result->src[1] = k; - result->src[2] = v; - result->src[3] = g; - result->src[4] = beta; - result->src[5] = state; - - return result; -} - -//////////////////////////////////////////////////////////////////////////////// - -struct ggml_hash_set ggml_hash_set_new(size_t size) { - size = ggml_hash_size(size); - struct ggml_hash_set result; - result.size = size; - result.keys = GGML_MALLOC(sizeof(struct ggml_tensor *) * size); - result.used = GGML_CALLOC(ggml_bitset_size(size), sizeof(ggml_bitset_t)); - return result; -} - -void ggml_hash_set_reset(struct ggml_hash_set * hash_set) { - memset(hash_set->used, 0, sizeof(ggml_bitset_t) * ggml_bitset_size(hash_set->size)); -} - -void ggml_hash_set_free(struct ggml_hash_set * hash_set) { - GGML_FREE(hash_set->used); - GGML_FREE(hash_set->keys); -} - -size_t ggml_hash_size(size_t min_sz) { - // next primes after powers of two - static const size_t primes[] = { - 2, 3, 5, 11, 17, 37, 67, 131, 257, 521, 1031, - 2053, 4099, 8209, 16411, 32771, 65537, 131101, - 262147, 524309, 1048583, 2097169, 4194319, 8388617, - 16777259, 33554467, 67108879, 134217757, 268435459, - 536870923, 1073741827, 2147483659 - }; - static const size_t n_primes = sizeof(primes)/sizeof(primes[0]); - - // find the smallest prime that is larger or equal than min_sz - size_t l = 0; - size_t r = n_primes; - while (l < r) { - size_t m = (l + r)/2; - if (primes[m] < min_sz) { - l = m + 1; - } else { - r = m; - } - } - size_t sz = l < n_primes ? primes[l] : min_sz | 1; - return sz; -} - -struct hash_map { - struct ggml_hash_set set; - struct ggml_tensor ** vals; -}; - -static struct hash_map * ggml_new_hash_map(size_t size) { - struct hash_map * result = GGML_MALLOC(sizeof(struct hash_map)); - result->set = ggml_hash_set_new(size); - result->vals = GGML_CALLOC(result->set.size, sizeof(struct ggml_tensor *)); - return result; -} - -static void ggml_hash_map_free(struct hash_map * map) { - ggml_hash_set_free(&map->set); - GGML_FREE(map->vals); - GGML_FREE(map); -} - -// utility functions to change gradients -// isrc is the index of tensor in cgraph->visited_has_set.keys -// the corresponding gradient (accumulators) are also at position isrc -// if tensor has a gradient accumulator, modify that accumulator in-place -// else if there is no gradient for tensor, set the corresponding value -// else, just add/subtract/etc. the gradients - -static void ggml_add_or_set( - struct ggml_context * ctx, - struct ggml_cgraph * cgraph, - size_t isrc, - struct ggml_tensor * tensor) { - struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; - GGML_ASSERT(src); - if (cgraph->grads[isrc]) { - cgraph->grads[isrc] = ggml_add_impl(ctx, cgraph->grads[isrc], tensor, /*inplace =*/ cgraph->grad_accs[isrc]); - } else { - cgraph->grads[isrc] = tensor; - } - ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name); - ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); -} - -static void ggml_acc_or_set( - struct ggml_context * ctx, - struct ggml_cgraph * cgraph, - size_t isrc, - struct ggml_tensor * tensor, - const size_t nb1, - const size_t nb2, - const size_t nb3, - const size_t offset) { - struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; - GGML_ASSERT(src); - if (cgraph->grads[isrc]) { - cgraph->grads[isrc] = ggml_acc_impl(ctx, cgraph->grads[isrc], tensor, nb1, nb2, nb3, offset, cgraph->grad_accs[isrc]); - } else { - struct ggml_tensor * a_zero = ggml_scale(ctx, src, 0.0f); // FIXME this is going to produce NaN if a contains inf/NaN - cgraph->grads[isrc] = ggml_acc_impl(ctx, a_zero, tensor, nb1, nb2, nb3, offset, false); - } - ggml_format_name(cgraph->grads[isrc], "grad for %s", cgraph->visited_hash_set.keys[isrc]->name); - ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); -} - -static void ggml_add1_or_set( - struct ggml_context * ctx, - struct ggml_cgraph * cgraph, - size_t isrc, - struct ggml_tensor * tensor) { - struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; - GGML_ASSERT(src); - if (cgraph->grads[isrc]) { - cgraph->grads[isrc] = ggml_add1_impl(ctx, cgraph->grads[isrc], tensor, cgraph->grad_accs[isrc]); - } else { - cgraph->grads[isrc] = ggml_repeat(ctx, tensor, src); - } - ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name); - ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); -} - -static void ggml_sub_or_set( - struct ggml_context * ctx, - struct ggml_cgraph * cgraph, - size_t isrc, - struct ggml_tensor * tensor) { - struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; - GGML_ASSERT(src); - if (cgraph->grads[isrc]) { - cgraph->grads[isrc] = ggml_sub_impl(ctx, cgraph->grads[isrc], tensor, cgraph->grad_accs[isrc]); - } else { - cgraph->grads[isrc] = ggml_neg(ctx, tensor); - } - ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name); - ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); -} - -static void ggml_compute_backward( - struct ggml_context * ctx, struct ggml_cgraph * cgraph, int i, const bool * grads_needed) { - struct ggml_tensor * tensor = cgraph->nodes[i]; - struct ggml_tensor * grad = ggml_graph_get_grad(cgraph, tensor); - - if (!grad) { - return; - } - - struct ggml_tensor * src0 = tensor->src[0]; - struct ggml_tensor * src1 = tensor->src[1]; - struct ggml_tensor * src2 = tensor->src[2]; - struct ggml_hash_set * hash_set = &cgraph->visited_hash_set; - const size_t isrc0 = src0 ? ggml_hash_find(hash_set, src0) : (size_t) -1; - const size_t isrc1 = src1 ? ggml_hash_find(hash_set, src1) : (size_t) -1; - const size_t isrc2 = src2 ? ggml_hash_find(hash_set, src2) : (size_t) -1; - const bool src0_needs_grads = src0 && isrc0 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc0) && grads_needed[isrc0]; - const bool src1_needs_grads = src1 && isrc1 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc1) && grads_needed[isrc1]; - const bool src2_needs_grads = src2 && isrc2 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc2) && grads_needed[isrc2]; - - switch (tensor->op) { - case GGML_OP_DUP: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, grad); - } - } break; - case GGML_OP_ADD: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, grad); - } - if (src1_needs_grads) { - struct ggml_tensor * tmp = grad; - if (!ggml_are_same_shape(src0, src1)) { - tmp = ggml_repeat_back(ctx, tmp, src1); - } - ggml_add_or_set(ctx, cgraph, isrc1, tmp); - } - } break; - case GGML_OP_ADD1: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, grad); - } - if (src1_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc1, ggml_mean(ctx, grad)); // TODO: should probably be sum instead of mean - } - } break; - case GGML_OP_ACC: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, grad); - } - if (src1_needs_grads) { - const size_t nb1 = ((int32_t *) tensor->op_params)[0]; - const size_t nb2 = ((int32_t *) tensor->op_params)[1]; - const size_t nb3 = ((int32_t *) tensor->op_params)[2]; - const size_t offset = ((int32_t *) tensor->op_params)[3]; - - struct ggml_tensor * tensor_grad_view = ggml_view_4d(ctx, - grad, src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], - nb1, nb2, nb3, offset); - - ggml_add_or_set(ctx, cgraph, isrc1, ggml_reshape(ctx, ggml_cont(ctx, tensor_grad_view), src1)); - } - } break; - case GGML_OP_SUB: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, grad); - } - if (src1_needs_grads) { - ggml_sub_or_set(ctx, cgraph, isrc1, grad); - } - } break; - case GGML_OP_MUL: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, src1)); - } - if (src1_needs_grads) { - struct ggml_tensor * tmp = ggml_mul(ctx, src0, grad); - if (!ggml_are_same_shape(src0, src1)) { - tmp = ggml_repeat_back(ctx, tmp, src1); - } - ggml_add_or_set(ctx, cgraph, isrc1, tmp); - } - } break; - case GGML_OP_DIV: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_div(ctx, grad, src1)); - } - if (src1_needs_grads) { - ggml_sub_or_set(ctx, cgraph, isrc1, ggml_mul(ctx, grad, ggml_div(ctx, tensor, src1))); - } - } break; - case GGML_OP_SQR: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_scale(ctx, ggml_mul(ctx, src0, grad), 2.0f)); - } - } break; - case GGML_OP_SQRT: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_scale(ctx, ggml_div(ctx, grad, tensor), 0.5f)); - } - } break; - case GGML_OP_LOG: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_div(ctx, grad, src0)); - } - } break; - case GGML_OP_SIN: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_cos(ctx, src0))); - } - } break; - case GGML_OP_COS: { - if (src0_needs_grads) { - ggml_sub_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_sin(ctx, src0))); - } - } break; - case GGML_OP_SUM: { - if (src0_needs_grads) { - ggml_add1_or_set(ctx, cgraph, isrc0, grad); - } - } break; - case GGML_OP_SUM_ROWS: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat(ctx, grad, src0)); - } - } break; - case GGML_OP_MEAN: { - if (src0_needs_grads) { - ggml_add1_or_set(ctx, cgraph, isrc0, ggml_scale_impl(ctx, grad, 1.0f/src0->ne[0], 0.0, false)); - } - } break; - case GGML_OP_REPEAT: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat_back(ctx, grad, src0)); - } - } break; - case GGML_OP_REPEAT_BACK: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat(ctx, grad, src0)); - } - } break; - case GGML_OP_RMS_NORM: { - if (src0_needs_grads) { - float eps; - memcpy(&eps, tensor->op_params, sizeof(float)); - ggml_add_or_set(ctx, cgraph, isrc0, ggml_rms_norm_back(ctx, grad, src0, eps)); - } - } break; - case GGML_OP_MUL_MAT: { - // https://cs231n.github.io/optimization-2/#staged - // # forward pass - // s0 = np.random.randn(5, 10) - // s1 = np.random.randn(10, 3) - // t = s0.dot(s1) - - // # now suppose we had the gradient on t from above in the circuit - // dt = np.random.randn(*t.shape) # same shape as t - // ds0 = dt.dot(s1.T) #.T gives the transpose of the matrix - // ds1 = t.T.dot(dt) - - // tensor.shape [m,p,qq,rr] - // src0.shape [n,m,q1,r1] - // src1.shape [n,p,qq,rr] - - if (src0_needs_grads) { - GGML_ASSERT(grad->ne[2] == src1->ne[2]); - GGML_ASSERT(grad->ne[3] == src1->ne[3]); - struct ggml_tensor * tmp = - ggml_out_prod(ctx, // [n,m,qq,rr] - src1, // [n,p,qq,rr] - grad); // [m,p,qq,rr] - if (!ggml_are_same_shape(tmp, src0)) { - GGML_ASSERT(tmp->ne[0] == src0->ne[0]); - GGML_ASSERT(tmp->ne[1] == src0->ne[1]); - GGML_ASSERT(tmp->ne[3] == 1); - - const int64_t nr2 = tmp->ne[2] / src0->ne[2]; - const size_t nb2 = tmp->nb[2] * nr2; - const size_t nb3 = tmp->nb[2]; - - tmp = ggml_view_4d(ctx, tmp, src0->ne[0], src0->ne[1], src0->ne[2], nr2, tmp->nb[1], nb2, nb3, 0); - tmp = ggml_repeat_back(ctx, tmp, src0); - } - ggml_add_or_set(ctx, cgraph, isrc0, tmp); - } - if (src1_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc1, - // ggml_mul_mat(ctx, // [n,p,qq,rr] - // ggml_cont(ctx, // [m,n,q1,r1] - // ggml_transpose(ctx, src0)), // [m,n,q1,r1] - // grad), // [m,p,qq,rr] - - // when src0 is bigger than tensor->grad (this is mostly the case in llama), - // avoid transpose of src0, rather transpose smaller tensor->grad - // and then use ggml_out_prod - ggml_out_prod(ctx, // [n,p,qq,rr] - src0, // [n,m,q1,r1] - ggml_transpose(ctx, // [p,m,qq,rr] - grad))); // [m,p,qq,rr] - } - } break; - case GGML_OP_SCALE: { - if (src0_needs_grads) { - float s; - memcpy(&s, tensor->op_params, sizeof(float)); - ggml_add_or_set(ctx, cgraph, isrc0, ggml_scale_impl(ctx, grad, s, 0.0, false)); - } - } break; - case GGML_OP_SET: { - const size_t nb1 = ((const int32_t *) tensor->op_params)[0]; - const size_t nb2 = ((const int32_t *) tensor->op_params)[1]; - const size_t nb3 = ((const int32_t *) tensor->op_params)[2]; - const size_t offset = ((const int32_t *) tensor->op_params)[3]; - - struct ggml_tensor * tensor_grad_view = NULL; - - if (src0_needs_grads || src1_needs_grads) { - GGML_ASSERT(src0->type == tensor->type); - GGML_ASSERT(!cgraph->grads[isrc0] || cgraph->grads[isrc0]->type == grad->type); - GGML_ASSERT(!cgraph->grads[isrc1] || !src1_needs_grads || cgraph->grads[isrc1]->type == grad->type); - - tensor_grad_view = ggml_view_4d(ctx, - grad, src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], - nb1, nb2, nb3, offset); - } - - if (src0_needs_grads) { - struct ggml_tensor * tmp = ggml_neg(ctx, tensor_grad_view); - ggml_add_or_set(ctx, cgraph, isrc0, ggml_acc_impl(ctx, grad, tmp, nb1, nb2, nb3, offset, false)); - } - - if (src1_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc1, ggml_reshape(ctx, ggml_cont(ctx, tensor_grad_view), src1)); - } - } break; - case GGML_OP_CPY: { - // cpy overwrites value of src1 by src0 and returns view(src1) - // the overwriting is mathematically equivalent to: - // tensor = src0 * 1 + src1 * 0 - if (src0_needs_grads) { - // dsrc0 = dtensor * 1 - ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape(ctx, grad, src0)); - } - if (src1_needs_grads) { - // dsrc1 = dtensor * 0 -> noop - } - } break; - case GGML_OP_CONT: { - // same as cpy - if (src0_needs_grads) { - GGML_ASSERT(!cgraph->grads[isrc0] || ggml_is_contiguous(cgraph->grads[isrc0])); - GGML_ASSERT(ggml_is_contiguous(grad)); - GGML_ASSERT(ggml_nelements(tensor) == ggml_nelements(src0)); - ggml_add_or_set(ctx, cgraph, isrc0, - ggml_are_same_shape(tensor, src0) ? grad : ggml_reshape(ctx, grad, src0)); - } - } break; - case GGML_OP_RESHAPE: { - if (src0_needs_grads) { - struct ggml_tensor * grad_cont = ggml_is_contiguous(grad) ? grad : ggml_cont(ctx, grad); - ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape(ctx, grad_cont, src0)); - } - } break; - case GGML_OP_VIEW: { - if (src0_needs_grads) { - size_t offset; - - memcpy(&offset, tensor->op_params, sizeof(offset)); - - size_t nb1 = tensor->nb[1]; - size_t nb2 = tensor->nb[2]; - size_t nb3 = tensor->nb[3]; - - if (cgraph->grads[isrc0] && src0->type != cgraph->grads[isrc0]->type) { - // gradient is typically F32, but src0 could be other type - size_t ng = ggml_element_size(cgraph->grads[isrc0]); - size_t n0 = ggml_element_size(src0); - GGML_ASSERT(offset % n0 == 0); - GGML_ASSERT(nb1 % n0 == 0); - GGML_ASSERT(nb2 % n0 == 0); - GGML_ASSERT(nb3 % n0 == 0); - offset = (offset / n0) * ng; - nb1 = (nb1 / n0) * ng; - nb2 = (nb2 / n0) * ng; - nb3 = (nb3 / n0) * ng; - } - - ggml_acc_or_set(ctx, cgraph, isrc0, grad, nb1, nb2, nb3, offset); - } - } break; - case GGML_OP_PERMUTE: { - if (src0_needs_grads) { - const int32_t * axes = (const int32_t *) tensor->op_params; - const int axis0 = axes[0] & 0x3; - const int axis1 = axes[1] & 0x3; - const int axis2 = axes[2] & 0x3; - const int axis3 = axes[3] & 0x3; - int axb[4] = {0,0,0,0}; // axes backward - axb[axis0] = 0; - axb[axis1] = 1; - axb[axis2] = 2; - axb[axis3] = 3; - ggml_add_or_set(ctx, cgraph, isrc0, ggml_permute(ctx, grad, axb[0], axb[1], axb[2], axb[3])); - } - } break; - case GGML_OP_TRANSPOSE: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_transpose(ctx, grad)); - } - } break; - case GGML_OP_GET_ROWS: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_get_rows_back(ctx, grad, src1, src0)); - } - if (src1_needs_grads) { - // noop - } - } break; - case GGML_OP_DIAG_MASK_INF: { - if (src0_needs_grads) { - /* ggml_diag_mask_inf_impl() shouldn't be here */ - /* ref: https://github.com/ggml-org/llama.cpp/pull/4203#discussion_r1412377992 */ - const int n_past = ((const int32_t *) tensor->op_params)[0]; - ggml_add_or_set(ctx, cgraph, isrc0, ggml_diag_mask_zero_impl(ctx, grad, n_past, false)); - } - } break; - case GGML_OP_DIAG_MASK_ZERO: { - if (src0_needs_grads) { - const int n_past = ((const int32_t *) tensor->op_params)[0]; - ggml_add_or_set(ctx, cgraph, isrc0, ggml_diag_mask_zero_impl(ctx, grad, n_past, false)); - } - } break; - case GGML_OP_SOFT_MAX: { - if (src0_needs_grads) { - float scale = 1.0f; - float max_bias = 0.0f; - - memcpy(&scale, (const float *) tensor->op_params + 0, sizeof(float)); - memcpy(&max_bias, (const float *) tensor->op_params + 1, sizeof(float)); - - ggml_add_or_set(ctx, cgraph, isrc0, ggml_soft_max_ext_back(ctx, grad, tensor, scale, max_bias)); - } - GGML_ASSERT((!src1 || !src1_needs_grads) && "backward pass for softmax mask not implemented"); - } break; - case GGML_OP_ROPE: { - if (src0_needs_grads) { - //const int n_past = ((int32_t *) tensor->op_params)[0]; - const int n_dims = ((const int32_t *) tensor->op_params)[1]; - const int mode = ((const int32_t *) tensor->op_params)[2]; - //const int n_ctx = ((int32_t *) tensor->op_params)[3]; - const int n_ctx_orig = ((const int32_t *) tensor->op_params)[4]; - float freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow; - int sections[4] = {0, 0, 0, 0}; - - memcpy(&freq_base, (const float *) tensor->op_params + 5, sizeof(float)); - memcpy(&freq_scale, (const float *) tensor->op_params + 6, sizeof(float)); - memcpy(&ext_factor, (const float *) tensor->op_params + 7, sizeof(float)); - memcpy(&attn_factor, (const float *) tensor->op_params + 8, sizeof(float)); - memcpy(&beta_fast, (const float *) tensor->op_params + 9, sizeof(float)); - memcpy(&beta_slow, (const float *) tensor->op_params + 10, sizeof(float)); - memcpy(§ions, tensor->op_params + 11, sizeof(sections)); - - struct ggml_tensor * rope_back = grad->ne[2] == src1->ne[0] ? - ggml_rope_ext_back(ctx, grad, src1, src2, n_dims, - mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow) : - ggml_rope_multi_back(ctx, grad, src1, src2, n_dims, sections, - mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - ggml_add_or_set(ctx, cgraph, isrc0, rope_back); - } - GGML_ASSERT((!src2 || !src2_needs_grads) && "gradients for freq factors not implemented"); - } break; - case GGML_OP_IM2COL: { - if (src1_needs_grads) { - const int32_t s0 = ggml_get_op_params_i32(tensor, 0); - const int32_t s1 = ggml_get_op_params_i32(tensor, 1); - const int32_t p0 = ggml_get_op_params_i32(tensor, 2); - const int32_t p1 = ggml_get_op_params_i32(tensor, 3); - const int32_t d0 = ggml_get_op_params_i32(tensor, 4); - const int32_t d1 = ggml_get_op_params_i32(tensor, 5); - const bool is_2D = ggml_get_op_params_i32(tensor, 6) == 1; - - ggml_add_or_set(ctx, cgraph, isrc1, ggml_im2col_back(ctx, grad, src0, src1->ne, s0, s1, p0, p1, d0, d1, is_2D)); - } - } break; - case GGML_OP_POOL_2D: { - if (src0_needs_grads) { - const enum ggml_op_pool op = ggml_get_op_params_i32(tensor, 0); - const int32_t k0 = ggml_get_op_params_i32(tensor, 1); - const int32_t k1 = ggml_get_op_params_i32(tensor, 2); - const int32_t s0 = ggml_get_op_params_i32(tensor, 3); - const int32_t s1 = ggml_get_op_params_i32(tensor, 4); - const int32_t p0 = ggml_get_op_params_i32(tensor, 5); - const int32_t p1 = ggml_get_op_params_i32(tensor, 6); - - ggml_add_or_set(ctx, cgraph, isrc0, ggml_pool_2d_back(ctx, grad, src0, op, k0, k1, s0, s1, p0, p1)); - } - } break; - case GGML_OP_WIN_PART: - case GGML_OP_WIN_UNPART: - case GGML_OP_UNARY: { - switch (ggml_get_unary_op(tensor)) { - case GGML_UNARY_OP_ABS: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, ggml_sgn(ctx, src0), grad)); - } - } break; - case GGML_UNARY_OP_SGN: { - // noop - } break; - case GGML_UNARY_OP_NEG: { - if (src0_needs_grads) { - ggml_sub_or_set(ctx, cgraph, isrc0, grad); - } - } break; - case GGML_UNARY_OP_STEP: { - // noop - } break; - case GGML_UNARY_OP_RELU: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, ggml_step(ctx, src0), grad)); - } - } break; - case GGML_UNARY_OP_SILU: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_silu_back(ctx, grad, src0)); - } - } break; - case GGML_UNARY_OP_EXP: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, tensor, grad)); - } - } break; - case GGML_UNARY_OP_EXPM1: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_exp(ctx, src0))); - } - } break; - case GGML_UNARY_OP_SOFTPLUS: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_sigmoid(ctx, src0))); - } - } break; - default: { - fprintf(stderr, "%s: unsupported unary op for backward pass: %s\n", - __func__, ggml_unary_op_name(ggml_get_unary_op(tensor))); - GGML_ABORT("fatal error"); - } //break; - } - } break; - case GGML_OP_CROSS_ENTROPY_LOSS: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_cross_entropy_loss_back(ctx, grad, src0, src1)); - } - GGML_ASSERT(!src1_needs_grads && "backward pass for labels not implemented"); - } break; - case GGML_OP_GLU: { - switch (ggml_get_glu_op(tensor)) { - case GGML_GLU_OP_SWIGLU: { - if (src0_needs_grads) { - GGML_ASSERT(src1 && "backward pass only implemented for split swiglu"); - ggml_add_or_set(ctx, cgraph, isrc0, ggml_silu_back(ctx, ggml_mul(ctx, grad, src1), src0)); - } - if (src1_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc1, ggml_mul(ctx, ggml_silu(ctx, src0), grad)); - } - } break; - default: { - GGML_ABORT("unsupported glu op for backward pass: %s", ggml_glu_op_name(ggml_get_glu_op(tensor))); - } //break; - } - } break; - case GGML_OP_NONE: { - // noop - } break; - case GGML_OP_COUNT: - default: { - GGML_ABORT("%s: unsupported ggml op for backward pass: %s\n", __func__, ggml_op_name(tensor->op)); - } //break; - } - - GGML_ASSERT(!src0_needs_grads || ggml_are_same_shape(src0, cgraph->grads[isrc0])); - GGML_ASSERT(!src1_needs_grads || ggml_are_same_shape(src1, cgraph->grads[isrc1])); - GGML_ASSERT(!src2_needs_grads || ggml_are_same_shape(src2, cgraph->grads[isrc2])); -} - -static size_t ggml_visit_parents_graph(struct ggml_cgraph * cgraph, struct ggml_tensor * node, bool compute) { - if (node->op != GGML_OP_NONE && compute) { - node->flags |= GGML_TENSOR_FLAG_COMPUTE; - } - - const size_t node_hash_pos = ggml_hash_find(&cgraph->visited_hash_set, node); - GGML_ASSERT(node_hash_pos != GGML_HASHSET_FULL); - - if (ggml_bitset_get(cgraph->visited_hash_set.used, node_hash_pos)) { - // already visited - - if (compute) { - // update the compute flag regardless - for (int i = 0; i < GGML_MAX_SRC; ++i) { - struct ggml_tensor * src = node->src[i]; - if (src && ((src->flags & GGML_TENSOR_FLAG_COMPUTE) == 0)) { - ggml_visit_parents_graph(cgraph, src, true); - } - } - } - - return node_hash_pos; - } - - // This is the first time we see this node in the current graph. - cgraph->visited_hash_set.keys[node_hash_pos] = node; - ggml_bitset_set(cgraph->visited_hash_set.used, node_hash_pos); - cgraph->use_counts[node_hash_pos] = 0; - - for (int i = 0; i < GGML_MAX_SRC; ++i) { - const int k = - (cgraph->order == GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT) ? i : - (cgraph->order == GGML_CGRAPH_EVAL_ORDER_RIGHT_TO_LEFT) ? (GGML_MAX_SRC-1-i) : - /* unknown order, just fall back to using i */ i; - - struct ggml_tensor * src = node->src[k]; - if (src) { - const size_t src_hash_pos = ggml_visit_parents_graph(cgraph, src, compute); - - // Update the use count for this operand. - cgraph->use_counts[src_hash_pos]++; - } - } - - if (node->op == GGML_OP_NONE && !(node->flags & GGML_TENSOR_FLAG_PARAM)) { - // reached a leaf node, not part of the gradient graph (e.g. a constant) - GGML_ASSERT(cgraph->n_leafs < cgraph->size); - - if (strlen(node->name) == 0) { - ggml_format_name(node, "leaf_%d", cgraph->n_leafs); - } - - cgraph->leafs[cgraph->n_leafs] = node; - cgraph->n_leafs++; - } else { - GGML_ASSERT(cgraph->n_nodes < cgraph->size); - - if (strlen(node->name) == 0) { - ggml_format_name(node, "node_%d", cgraph->n_nodes); - } - - cgraph->nodes[cgraph->n_nodes] = node; - cgraph->n_nodes++; - } - - return node_hash_pos; -} - -static void ggml_build_forward_impl(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor, bool expand, bool compute) { - if (!expand) { - // TODO: this branch isn't accessible anymore, maybe move this to ggml_build_forward_expand - ggml_graph_clear(cgraph); - } - - const int n_old = cgraph->n_nodes; - - ggml_visit_parents_graph(cgraph, tensor, compute); - - const int n_new = cgraph->n_nodes - n_old; - GGML_PRINT_DEBUG("%s: visited %d new nodes\n", __func__, n_new); - - if (n_new > 0) { - // the last added node should always be starting point - GGML_ASSERT(cgraph->nodes[cgraph->n_nodes - 1] == tensor); - } -} - -struct ggml_tensor * ggml_build_forward_select( - struct ggml_cgraph * cgraph, - struct ggml_tensor ** tensors, - int n_tensors, - int idx) { - GGML_ASSERT(idx >= 0 && idx < n_tensors); - - for (int i = 0; i < n_tensors; i++) { - ggml_build_forward_impl(cgraph, tensors[i], true, i == idx ? true : false); - } - - return tensors[idx]; -} - -void ggml_build_forward_expand(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor) { - ggml_build_forward_impl(cgraph, tensor, true, true); -} - -void ggml_build_backward_expand( - struct ggml_context * ctx, - struct ggml_cgraph * cgraph, - struct ggml_tensor ** grad_accs) { - GGML_ASSERT(cgraph->n_nodes > 0); - GGML_ASSERT(cgraph->grads); - GGML_ASSERT(cgraph->grad_accs); - - const int n_nodes_f = cgraph->n_nodes; - - memset(cgraph->grads, 0, cgraph->visited_hash_set.size*sizeof(struct ggml_tensor *)); - memset(cgraph->grad_accs, 0, cgraph->visited_hash_set.size*sizeof(struct ggml_tensor *)); - bool * grads_needed = calloc(cgraph->visited_hash_set.size, sizeof(bool)); - - { - bool any_params = false; - bool any_loss = false; - for (int i = 0; i < n_nodes_f; ++i) { - struct ggml_tensor * node = cgraph->nodes[i]; - any_params = any_params || (node->flags & GGML_TENSOR_FLAG_PARAM); - any_loss = any_loss || (node->flags & GGML_TENSOR_FLAG_LOSS); - } - GGML_ASSERT(any_params && "no trainable parameters found, did you forget to call ggml_set_param?"); - GGML_ASSERT(any_loss && "no training loss found, did you forget to call ggml_set_loss?"); - } - - for (int i = 0; i < n_nodes_f; ++i) { - struct ggml_tensor * node = cgraph->nodes[i]; - - if (node->type == GGML_TYPE_I32) { - continue; - } - - bool node_needs_grad = (node->flags & GGML_TENSOR_FLAG_PARAM) || (node->flags & GGML_TENSOR_FLAG_LOSS); - bool ignore_src[GGML_MAX_SRC] = {false}; - switch (node->op) { - // gradients in node->src[0] for one reason or another have no effect on output gradients - case GGML_OP_IM2COL: // only used for its shape - case GGML_OP_IM2COL_BACK: // same as IM2COL - ignore_src[0] = true; - break; - case GGML_OP_UNARY: { - const enum ggml_unary_op uop = ggml_get_unary_op(node); - // SGN and STEP unary ops are piecewise constant - if (uop == GGML_UNARY_OP_SGN || uop == GGML_UNARY_OP_STEP) { - ignore_src[0] = true; - } - } break; - - // gradients in node->src[1] for one reason or another have no effect on output gradients - case GGML_OP_CPY: // gradients in CPY target are irrelevant - case GGML_OP_GET_ROWS: // row indices not differentiable - case GGML_OP_GET_ROWS_BACK: // same as for GET_ROWS - case GGML_OP_ROPE: // positions not differentiable - ignore_src[1] = true; - break; - - default: - break; - } - for (int j = 0; j < GGML_MAX_SRC; ++j) { - if (!node->src[j] || ignore_src[j] || !grads_needed[ggml_hash_find(&cgraph->visited_hash_set, node->src[j])]) { - continue; - } - GGML_ASSERT(node->src[j]->type == GGML_TYPE_F32 || node->src[j]->type == GGML_TYPE_F16); - node_needs_grad = true; - break; - } - if (!node_needs_grad) { - continue; - } - - // inplace operations are currently not supported - GGML_ASSERT(!node->view_src || node->op == GGML_OP_CPY || node->op == GGML_OP_VIEW || - node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE); - - const size_t ihash = ggml_hash_find(&cgraph->visited_hash_set, node); - GGML_ASSERT(ihash != GGML_HASHSET_FULL); - GGML_ASSERT(ggml_bitset_get(cgraph->visited_hash_set.used, ihash)); - if (grad_accs && grad_accs[i]) { - cgraph->grad_accs[ihash] = grad_accs[i]; - cgraph->grads[ihash] = cgraph->grad_accs[ihash]; - } else if (node->flags & GGML_TENSOR_FLAG_LOSS) { - // loss tensors always need a gradient accumulator - cgraph->grad_accs[ihash] = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, node->ne); - cgraph->grads[ihash] = cgraph->grad_accs[ihash]; - } - grads_needed[ihash] = true; - } - - for (int i = n_nodes_f - 1; i >= 0; --i) { - // inplace operations to add gradients are not created by ggml_compute_backward except for gradient accumulation - // use allocator to automatically make inplace operations - ggml_compute_backward(ctx, cgraph, i, grads_needed); - } - - free(grads_needed); -} - -static void * incr_ptr_aligned(void ** p, size_t size, size_t align) { - void * ptr = *p; - ptr = (void *) GGML_PAD((uintptr_t) ptr, align); - *p = (void *) ((char *) ptr + size); - return ptr; -} - -static size_t ggml_graph_nbytes(size_t size, bool grads) { - size_t hash_size = ggml_hash_size(size * 2); - void * p = 0; - incr_ptr_aligned(&p, sizeof(struct ggml_cgraph), 1); - incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // nodes - incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // leafs - incr_ptr_aligned(&p, hash_size * sizeof(int32_t), sizeof(int32_t)); // use_counts - incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // hash keys - if (grads) { - incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // grads - incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // grad_accs - } - incr_ptr_aligned(&p, ggml_bitset_size(hash_size) * sizeof(ggml_bitset_t), sizeof(ggml_bitset_t)); - - size_t nbytes = (size_t) p; - return nbytes; -} - -size_t ggml_graph_overhead_custom(size_t size, bool grads) { - return GGML_OBJECT_SIZE + GGML_PAD(ggml_graph_nbytes(size, grads), GGML_MEM_ALIGN); -} - -size_t ggml_graph_overhead(void) { - return ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, false); -} - -struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t size, bool grads) { - const size_t obj_size = ggml_graph_nbytes(size, grads); - struct ggml_object * obj = ggml_new_object(ctx, GGML_OBJECT_TYPE_GRAPH, obj_size); - struct ggml_cgraph * cgraph = (struct ggml_cgraph *) ((char *) ctx->mem_buffer + obj->offs); - - // the size of the hash table is doubled since it needs to hold both nodes and leafs - size_t hash_size = ggml_hash_size(size * 2); - - void * p = cgraph + 1; - - struct ggml_tensor ** nodes_ptr = incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); - struct ggml_tensor ** leafs_ptr = incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); - int32_t * use_counts_ptr = incr_ptr_aligned(&p, hash_size * sizeof(int32_t), sizeof(int32_t)); - struct ggml_tensor ** hash_keys_ptr = incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); - struct ggml_tensor ** grads_ptr = grads ? incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)) : NULL; - struct ggml_tensor ** grad_accs_ptr = grads ? incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)) : NULL; - - ggml_bitset_t * hash_used = incr_ptr_aligned(&p, ggml_bitset_size(hash_size) * sizeof(ggml_bitset_t), sizeof(ggml_bitset_t)); - - // check that we allocated the correct amount of memory - assert(obj_size == (size_t)((char *)p - (char *)cgraph)); - - *cgraph = (struct ggml_cgraph) { - /*.size =*/ size, - /*.n_nodes =*/ 0, - /*.n_leafs =*/ 0, - /*.nodes =*/ nodes_ptr, - /*.grads =*/ grads_ptr, - /*.grad_accs =*/ grad_accs_ptr, - /*.leafs =*/ leafs_ptr, - /*.use_counts =*/ use_counts_ptr, - /*.hash_table =*/ { hash_size, hash_used, hash_keys_ptr }, - /*.order =*/ GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT, - /*.uid =*/ 0, - }; - - ggml_hash_set_reset(&cgraph->visited_hash_set); - if (grads) { - memset(cgraph->grads, 0, hash_size*sizeof(struct ggml_tensor *)); - memset(cgraph->grad_accs, 0, hash_size*sizeof(struct ggml_tensor *)); - } - - return cgraph; -} - -struct ggml_cgraph * ggml_new_graph(struct ggml_context * ctx) { - return ggml_new_graph_custom(ctx, GGML_DEFAULT_GRAPH_SIZE, false); -} - -struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph0, int i0, int i1) { - struct ggml_cgraph cgraph = { - /*.size =*/ 0, - /*.n_nodes =*/ i1 - i0, - /*.n_leafs =*/ 0, - /*.nodes =*/ cgraph0->nodes + i0, - /*.grads =*/ NULL, // gradients would need visited_hash_set - /*.grad_accs =*/ NULL, - /*.leafs =*/ NULL, - /*.use_counts =*/ cgraph0->use_counts, - /*.visited_hash_set =*/ cgraph0->visited_hash_set, - /*.order =*/ cgraph0->order, - /*.uid =*/ 0 - }; - - return cgraph; -} - -void ggml_graph_cpy(struct ggml_cgraph * src, struct ggml_cgraph * dst) { - GGML_ASSERT(dst->size >= src->n_leafs); - GGML_ASSERT(dst->size >= src->n_nodes); - GGML_ASSERT(dst->visited_hash_set.size >= src->visited_hash_set.size); - - dst->n_leafs = src->n_leafs; - dst->n_nodes = src->n_nodes; - dst->order = src->order; - - for (int i = 0; i < src->n_leafs; ++i) { - dst->leafs[i] = src->leafs[i]; - } - - for (int i = 0; i < src->n_nodes; ++i) { - dst->nodes[i] = src->nodes[i]; - } - - for (size_t i = 0; i < src->visited_hash_set.size; ++i) { - // copy all hashset keys (tensors) that are in use - if (ggml_bitset_get(src->visited_hash_set.used, i)) { - size_t new_hash_pos = ggml_hash_insert(&dst->visited_hash_set, src->visited_hash_set.keys[i]); - dst->use_counts[new_hash_pos] = src->use_counts[i]; - } - } - - if (dst->grads) { - memset(dst->grads, 0, dst->visited_hash_set.size*sizeof(struct ggml_tensor *)); - memset(dst->grad_accs, 0, dst->visited_hash_set.size*sizeof(struct ggml_tensor *)); - } - if (src->grads) { - GGML_ASSERT(dst->grads != NULL); - GGML_ASSERT(dst->grad_accs != NULL); - for (int i = 0; i < src->n_nodes; ++i) { - const size_t igrad_src = ggml_hash_find(&src->visited_hash_set, src->nodes[i]); - const size_t igrad_dst = ggml_hash_find(&dst->visited_hash_set, dst->nodes[i]); - - GGML_ASSERT(igrad_src != GGML_HASHSET_FULL); - GGML_ASSERT(ggml_bitset_get(src->visited_hash_set.used, igrad_src)); - GGML_ASSERT(igrad_dst != GGML_HASHSET_FULL); - GGML_ASSERT(ggml_bitset_get(dst->visited_hash_set.used, igrad_dst)); - - dst->grads[igrad_dst] = src->grads[igrad_src]; - dst->grad_accs[igrad_dst] = src->grad_accs[igrad_src]; - } - } -} - -struct ggml_cgraph * ggml_graph_dup(struct ggml_context * ctx, struct ggml_cgraph * cgraph, bool force_grads) { - struct ggml_cgraph * result = ggml_new_graph_custom(ctx, cgraph->size, cgraph->grads || force_grads); - ggml_graph_cpy(cgraph, result); - return result; -} - -struct ggml_tensor * ggml_set_zero(struct ggml_tensor * tensor) { - if (ggml_is_empty(tensor)) { - return tensor; - } - if (tensor->buffer) { - ggml_backend_tensor_memset(tensor, 0, 0, ggml_nbytes(tensor)); - } else { - GGML_ASSERT(tensor->data); - memset(tensor->data, 0, ggml_nbytes(tensor)); - } - return tensor; -} - -void ggml_graph_reset(struct ggml_cgraph * cgraph) { - if (!cgraph) { - return; - } - GGML_ASSERT(cgraph->grads != NULL); - - for (int i = 0; i < cgraph->n_nodes; i++) { - struct ggml_tensor * node = cgraph->nodes[i]; - struct ggml_tensor * grad_acc = ggml_graph_get_grad_acc(cgraph, node); - - if (node->op == GGML_OP_OPT_STEP_ADAMW) { - // clear momenta - ggml_set_zero(node->src[2]); - ggml_set_zero(node->src[3]); - } - - // initial gradients of loss should be 1, 0 otherwise - if (grad_acc) { - if (node->flags & GGML_TENSOR_FLAG_LOSS) { - GGML_ASSERT(grad_acc->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_scalar(grad_acc)); - - const float onef = 1.0f; - if (grad_acc->buffer) { - ggml_backend_tensor_set(grad_acc, &onef, 0, sizeof(float)); - } else { - GGML_ASSERT(grad_acc->data); - *((float *) grad_acc->data) = onef; - } - } else { - ggml_set_zero(grad_acc); - } - } - } -} - -void ggml_graph_clear(struct ggml_cgraph * cgraph) { - cgraph->n_leafs = 0; - cgraph->n_nodes = 0; - ggml_hash_set_reset(&cgraph->visited_hash_set); -} - -int ggml_graph_size(struct ggml_cgraph * cgraph) { - return cgraph->size; -} - -struct ggml_tensor * ggml_graph_node(struct ggml_cgraph * cgraph, int i) { - if (i < 0) { - GGML_ASSERT(cgraph->n_nodes + i >= 0); - return cgraph->nodes[cgraph->n_nodes + i]; - } - - GGML_ASSERT(i < cgraph->n_nodes); - return cgraph->nodes[i]; -} - -struct ggml_tensor ** ggml_graph_nodes(struct ggml_cgraph * cgraph) { - return cgraph->nodes; -} - -int ggml_graph_n_nodes(struct ggml_cgraph * cgraph) { - return cgraph->n_nodes; -} - -void ggml_graph_add_node(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor) { - GGML_ASSERT(cgraph->size > cgraph->n_nodes); - cgraph->nodes[cgraph->n_nodes] = tensor; - cgraph->n_nodes++; -} - -struct ggml_tensor * ggml_graph_get_tensor(const struct ggml_cgraph * cgraph, const char * name) { - for (int i = 0; i < cgraph->n_leafs; i++) { - struct ggml_tensor * leaf = cgraph->leafs[i]; - - if (strcmp(leaf->name, name) == 0) { - return leaf; - } - } - - for (int i = 0; i < cgraph->n_nodes; i++) { - struct ggml_tensor * node = cgraph->nodes[i]; - - if (strcmp(node->name, name) == 0) { - return node; - } - } - - return NULL; -} - -struct ggml_tensor * ggml_graph_get_grad(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { - const size_t igrad = ggml_hash_find(&cgraph->visited_hash_set, node); - return igrad != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, igrad) && cgraph->grads ? cgraph->grads[igrad] : NULL; -} - -struct ggml_tensor * ggml_graph_get_grad_acc(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { - const size_t igrad = ggml_hash_find(&cgraph->visited_hash_set, node); - return igrad != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, igrad) && cgraph->grad_accs ? cgraph->grad_accs[igrad] : NULL; -} - -void ggml_graph_print(const struct ggml_cgraph * cgraph) { - GGML_LOG_INFO("=== GRAPH ===\n"); - - GGML_LOG_INFO("n_nodes = %d\n", cgraph->n_nodes); - for (int i = 0; i < cgraph->n_nodes; i++) { - struct ggml_tensor * node = cgraph->nodes[i]; - - GGML_LOG_INFO(" - %3d: [ %5" PRId64 ", %5" PRId64 ", %5" PRId64 "] %16s %s\n", - i, - node->ne[0], node->ne[1], node->ne[2], - ggml_op_name(node->op), (node->flags & GGML_TENSOR_FLAG_PARAM) ? "x" : - ggml_graph_get_grad(cgraph, node) ? "g" : " "); - } - - GGML_LOG_INFO("n_leafs = %d\n", cgraph->n_leafs); - for (int i = 0; i < cgraph->n_leafs; i++) { - struct ggml_tensor * node = cgraph->leafs[i]; - - GGML_LOG_INFO(" - %3d: [ %5" PRId64 ", %5" PRId64 "] %8s %16s\n", - i, - node->ne[0], node->ne[1], - ggml_op_name(node->op), - ggml_get_name(node)); - } - - GGML_LOG_INFO("========================================\n"); -} - -static int ggml_node_list_find_tensor(const struct ggml_cgraph * cgraph, - const int * idxs, - int count, - const struct ggml_tensor * tensor) { - GGML_ASSERT(cgraph && idxs); - for (int i = 0; i < count; ++i) { - const int node_idx = idxs[i]; - - if (node_idx >= cgraph->n_nodes) { - return -1; - } - if (cgraph->nodes[node_idx] == tensor) { - return i; - } - } - return -1; -} - -bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph, - const int * node_idxs, - int count, - const enum ggml_op * ops, - const int * outputs, - int num_outputs) { - GGML_ASSERT(outputs && num_outputs > 0); - - for (int i = 0; i < count; ++i) { - if (node_idxs[i] >= cgraph->n_nodes) { - return false; - } - - const struct ggml_tensor * node = cgraph->nodes[node_idxs[i]]; - - if (node->op != ops[i]) { - return false; - } - - if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { - return false; - } - - if (ggml_node_list_find_tensor(cgraph, outputs, num_outputs, node) != -1) { - continue; - } - - if (node->flags & GGML_TENSOR_FLAG_OUTPUT) { - return false; - } - - int subgraph_uses = 0; - for (int j = i + 1; j < count; ++j) { - const struct ggml_tensor * other_node = cgraph->nodes[node_idxs[j]]; - for (int src_idx = 0; src_idx < GGML_MAX_SRC; src_idx++) { - if (other_node->src[src_idx] == node) { - subgraph_uses++; - } - } - } - - if (subgraph_uses != ggml_node_get_use_count(cgraph, node_idxs[i])) { - return false; - } - - // if node is a view, check if the view_src and all it's parent view_srcs are within the subgraph - struct ggml_tensor * view_src = node->view_src; - while (view_src) { - if (ggml_node_list_find_tensor(cgraph, node_idxs, count, view_src) == -1) { - return false; - } - view_src = view_src->view_src; - } - } - - return true; -} - -// check if node is part of the graph -static bool ggml_graph_find(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { - if (cgraph == NULL) { - return true; - } - - for (int i = 0; i < cgraph->n_nodes; i++) { - if (cgraph->nodes[i] == node) { - return true; - } - } - - return false; -} - -static struct ggml_tensor * ggml_graph_get_parent(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { - for (int i = 0; i < cgraph->n_nodes; i++) { - struct ggml_tensor * parent = cgraph->nodes[i]; - struct ggml_tensor * grad = ggml_graph_get_grad(cgraph, parent); - - if (grad == node) { - return parent; - } - } - - return NULL; -} - -static void ggml_graph_dump_dot_node_edge(FILE * fp, const struct ggml_cgraph * gb, struct ggml_tensor * node, struct ggml_tensor * parent, const char * label) { - struct ggml_tensor * gparent = ggml_graph_get_parent(gb, node); - struct ggml_tensor * gparent0 = ggml_graph_get_parent(gb, parent); - fprintf(fp, " \"%p\" -> \"%p\" [ arrowhead = %s; style = %s; label = \"%s\"; ]\n", - gparent0 ? (void *) gparent0 : (void *) parent, - gparent ? (void *) gparent : (void *) node, - gparent ? "empty" : "vee", - gparent ? "dashed" : "solid", - label); -} - -static void ggml_graph_dump_dot_leaf_edge(FILE * fp, struct ggml_tensor * node, struct ggml_tensor * parent, const char * label) { - fprintf(fp, " \"%p\" -> \"%p\" [ label = \"%s\"; ]\n", - (void *) parent, - (void *) node, - label); -} - -void ggml_graph_dump_dot(const struct ggml_cgraph * gb, const struct ggml_cgraph * cgraph, const char * filename) { - char color[16]; - - FILE * fp = ggml_fopen(filename, "w"); - GGML_ASSERT(fp); - - fprintf(fp, "digraph G {\n"); - fprintf(fp, " newrank = true;\n"); - fprintf(fp, " rankdir = TB;\n"); - - for (int i = 0; i < gb->n_nodes; i++) { - struct ggml_tensor * node = gb->nodes[i]; - struct ggml_tensor * grad = ggml_graph_get_grad(gb, node); - - if (ggml_graph_get_parent(gb, node) != NULL) { - continue; - } - - if (node->flags & GGML_TENSOR_FLAG_PARAM) { - snprintf(color, sizeof(color), "yellow"); - } else if (grad) { - if (ggml_graph_find(cgraph, node)) { - snprintf(color, sizeof(color), "green"); - } else { - snprintf(color, sizeof(color), "lightblue"); - } - } else { - snprintf(color, sizeof(color), "white"); - } - - fprintf(fp, " \"%p\" [ " - "style = filled; fillcolor = %s; shape = record; " - "label=\"", - (void *) node, color); - - if (strlen(node->name) > 0) { - fprintf(fp, "%s (%s)|", node->name, ggml_type_name(node->type)); - } else { - fprintf(fp, "(%s)|", ggml_type_name(node->type)); - } - - if (ggml_is_matrix(node)) { - fprintf(fp, "%d [%" PRId64 ", %" PRId64 "] | %s", i, node->ne[0], node->ne[1], ggml_op_symbol(node->op)); - } else { - fprintf(fp, "%d [%" PRId64 ", %" PRId64 ", %" PRId64 "] | %s", i, node->ne[0], node->ne[1], node->ne[2], ggml_op_symbol(node->op)); - } - - if (grad) { - fprintf(fp, " | %s\"; ]\n", ggml_op_symbol(grad->op)); - } else { - fprintf(fp, "\"; ]\n"); - } - } - - for (int i = 0; i < gb->n_leafs; i++) { - struct ggml_tensor * node = gb->leafs[i]; - - snprintf(color, sizeof(color), "pink"); - - fprintf(fp, " \"%p\" [ " - "style = filled; fillcolor = %s; shape = record; " - "label=\"", - (void *) node, color); - - if (strlen(node->name) > 0) { - fprintf(fp, "%s (%s)|", node->name, ggml_type_name(node->type)); - } else { - fprintf(fp, "(%s)|", ggml_type_name(node->type)); - } - - fprintf(fp, "CONST %d [%" PRId64 ", %" PRId64 "]", i, node->ne[0], node->ne[1]); - if (ggml_nelements(node) < 5 && node->data != NULL) { - fprintf(fp, " | ("); - for (int j = 0; j < ggml_nelements(node); j++) { - // FIXME: use ggml-backend to obtain the tensor data - //if (node->type == GGML_TYPE_I8 || node->type == GGML_TYPE_I16 || node->type == GGML_TYPE_I32) { - // fprintf(fp, "%d", ggml_get_i32_1d(node, j)); - //} - //else if (node->type == GGML_TYPE_F32 || - // node->type == GGML_TYPE_F16 || - // node->type == GGML_TYPE_BF16) { - // fprintf(fp, "%.1e", (double)ggml_get_f32_1d(node, j)); - //} - //else - { - fprintf(fp, "#"); - } - if (j < ggml_nelements(node) - 1) { - fprintf(fp, ", "); - } - } - fprintf(fp, ")"); - } - fprintf(fp, "\"; ]\n"); - } - - for (int i = 0; i < gb->n_nodes; i++) { - struct ggml_tensor * node = gb->nodes[i]; - - for (int j = 0; j < GGML_MAX_SRC; j++) { - if (node->src[j]) { - char label[16]; - snprintf(label, sizeof(label), "src %d", j); - ggml_graph_dump_dot_node_edge(fp, gb, node, node->src[j], label); - } - } - } - - for (int i = 0; i < gb->n_leafs; i++) { - struct ggml_tensor * node = gb->leafs[i]; - - for (int j = 0; j < GGML_MAX_SRC; j++) { - if (node->src[j]) { - char label[16]; - snprintf(label, sizeof(label), "src %d", j); - ggml_graph_dump_dot_leaf_edge(fp, node, node->src[j], label); - } - } - } - - fprintf(fp, "}\n"); - - fclose(fp); - - GGML_LOG_INFO("%s: dot -Tpng %s -o %s.png && open %s.png\n", __func__, filename, filename, filename); -} - -//////////////////////////////////////////////////////////////////////////////// - -void ggml_set_input(struct ggml_tensor * tensor) { - tensor->flags |= GGML_TENSOR_FLAG_INPUT; -} - -void ggml_set_output(struct ggml_tensor * tensor) { - tensor->flags |= GGML_TENSOR_FLAG_OUTPUT; -} - -void ggml_set_param(struct ggml_tensor * tensor) { - GGML_ASSERT(tensor->op == GGML_OP_NONE); - tensor->flags |= GGML_TENSOR_FLAG_PARAM; -} - -void ggml_set_loss(struct ggml_tensor * tensor) { - GGML_ASSERT(ggml_is_scalar(tensor)); - GGML_ASSERT(tensor->type == GGML_TYPE_F32); - tensor->flags |= GGML_TENSOR_FLAG_LOSS; -} - -//////////////////////////////////////////////////////////////////////////////// - -void ggml_quantize_init(enum ggml_type type) { - ggml_critical_section_start(); - - switch (type) { - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ2_S: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: iq2xs_init_impl(type); break; - case GGML_TYPE_IQ3_XXS: iq3xs_init_impl(256); break; - case GGML_TYPE_IQ3_S: iq3xs_init_impl(512); break; - default: // nothing - break; - } - - ggml_critical_section_end(); -} - -void ggml_quantize_free(void) { - ggml_critical_section_start(); - - iq2xs_free_impl(GGML_TYPE_IQ2_XXS); - iq2xs_free_impl(GGML_TYPE_IQ2_XS); - iq2xs_free_impl(GGML_TYPE_IQ2_S); - iq2xs_free_impl(GGML_TYPE_IQ1_S); - iq2xs_free_impl(GGML_TYPE_IQ1_M); - iq3xs_free_impl(256); - iq3xs_free_impl(512); - - ggml_critical_section_end(); -} - -bool ggml_quantize_requires_imatrix(enum ggml_type type) { - return - type == GGML_TYPE_IQ2_XXS || - type == GGML_TYPE_IQ2_XS || - type == GGML_TYPE_IQ1_S;// || - //type == GGML_TYPE_IQ1_M; -} - -size_t ggml_quantize_chunk( - enum ggml_type type, - const float * src, - void * dst, - int64_t start, - int64_t nrows, - int64_t n_per_row, - const float * imatrix) { - const int64_t n = nrows * n_per_row; - - if (ggml_quantize_requires_imatrix(type)) { - GGML_ASSERT(imatrix != NULL); - } - - GGML_ASSERT(start % type_traits[type].blck_size == 0); - GGML_ASSERT(start % n_per_row == 0); - - ggml_quantize_init(type); // this is noop if already initialized - - const size_t start_row = start / n_per_row; - const size_t row_size = ggml_row_size(type, n_per_row); - - size_t result = 0; - - switch (type) { - case GGML_TYPE_Q1_0: result = quantize_q1_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q4_0: result = quantize_q4_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q4_1: result = quantize_q4_1 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q5_0: result = quantize_q5_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q5_1: result = quantize_q5_1 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q8_0: result = quantize_q8_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_MXFP4: result = quantize_mxfp4 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_NVFP4: result = quantize_nvfp4 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q2_K: result = quantize_q2_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q3_K: result = quantize_q3_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q4_K: result = quantize_q4_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q5_K: result = quantize_q5_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q6_K: result = quantize_q6_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_TQ1_0: result = quantize_tq1_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_TQ2_0: result = quantize_tq2_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ2_XXS: result = quantize_iq2_xxs(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ2_XS: result = quantize_iq2_xs (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ3_XXS: result = quantize_iq3_xxs(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ3_S: result = quantize_iq3_s (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ2_S: result = quantize_iq2_s (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ1_S: result = quantize_iq1_s (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ1_M: result = quantize_iq1_m (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ4_NL: result = quantize_iq4_nl (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ4_XS: result = quantize_iq4_xs (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_F16: - { - size_t elemsize = sizeof(ggml_fp16_t); - ggml_fp32_to_fp16_row(src + start, (ggml_fp16_t *)dst + start, n); - result = n * elemsize; - } break; - case GGML_TYPE_BF16: - { - size_t elemsize = sizeof(ggml_bf16_t); - ggml_fp32_to_bf16_row_ref(src + start, (ggml_bf16_t *)dst + start, n); - result = n * elemsize; - } break; - case GGML_TYPE_F32: - { - size_t elemsize = sizeof(float); - result = n * elemsize; - memcpy((uint8_t *)dst + start * elemsize, src + start, result); - } break; - default: - assert(false); - } - - GGML_ASSERT(result == nrows * row_size); - - return result; -} - -//////////////////////////////////////////////////////////////////////////////// - -void ggml_log_get(ggml_log_callback * log_callback, void ** user_data) { - *log_callback = g_logger_state.log_callback; - *user_data = g_logger_state.log_callback_user_data; -} - -void ggml_log_set(ggml_log_callback log_callback, void * user_data) { - g_logger_state.log_callback = log_callback ? log_callback : ggml_log_callback_default; - g_logger_state.log_callback_user_data = user_data; -} - -void ggml_threadpool_params_init(struct ggml_threadpool_params * p, int n_threads) { - p->n_threads = n_threads; - p->prio = 0; // default priority (usually means normal or inherited) - p->poll = 50; // hybrid-polling enabled - p->strict_cpu = false; // no strict placement (all threads share same cpumask) - p->paused = false; // threads are ready to go - memset(p->cpumask, 0, GGML_MAX_N_THREADS); // all-zero means use the default affinity (usually inherited) -} - -struct ggml_threadpool_params ggml_threadpool_params_default(int n_threads) { - struct ggml_threadpool_params p; - ggml_threadpool_params_init(&p, n_threads); - return p; -} - -bool ggml_threadpool_params_match(const struct ggml_threadpool_params * p0, const struct ggml_threadpool_params * p1) { - if (p0->n_threads != p1->n_threads ) return false; - if (p0->prio != p1->prio ) return false; - if (p0->poll != p1->poll ) return false; - if (p0->strict_cpu != p1->strict_cpu ) return false; - return memcmp(p0->cpumask, p1->cpumask, GGML_MAX_N_THREADS) == 0; -} diff --git a/ggml/src/ggml.c.inc b/ggml/src/ggml.c.inc new file mode 100644 index 000000000000..fead5970eedd --- /dev/null +++ b/ggml/src/ggml.c.inc @@ -0,0 +1,14500 @@ +#define _CRT_SECURE_NO_DEPRECATE // Disables "unsafe" warnings on Windows +#define _USE_MATH_DEFINES // For M_PI on MSVC + +#include "ggml-backend.h.inc" +#include "ggml-impl-defs.inc" +#include "ggml-cpu.h.inc" +#include "ggml.h.inc" + +// FIXME: required here for quantization functions +#include "ggml-quants-defs.inc" + +#ifdef GGML_USE_CPU_HBM +#include +#endif + +#if defined(_MSC_VER) || defined(__MINGW32__) +#include // using malloc.h with MSC/MINGW +#elif !defined(__FreeBSD__) && !defined(__NetBSD__) && !defined(__OpenBSD__) +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#if defined(__gnu_linux__) +#include +#endif + +GGML_API void ggml_critical_section_start(void); +GGML_API void ggml_critical_section_end(void); + +#if defined(__APPLE__) +#include +#include +#include +#endif + +#if defined(_WIN32) +#define WIN32_LEAN_AND_MEAN +#ifndef NOMINMAX + #define NOMINMAX +#endif +#include +#endif + +#define UNUSED GGML_UNUSED + +uint64_t ggml_graph_next_uid(void) { +#ifdef _MSC_VER +#if defined(_WIN32) + static volatile LONG counter = 1; + return (uint64_t) InterlockedIncrement(&counter) - 1; +#else + static volatile long long counter = 1; + return (uint64_t) _InterlockedIncrement64(&counter) - 1; +#endif +#else + static uint64_t counter = 1; + return __atomic_fetch_add(&counter, 1, __ATOMIC_RELAXED); +#endif +} + +// Needed for ggml_fp32_to_bf16_row() +#if defined(__AVX512BF16__) +#if defined(_MSC_VER) +#define m512i(p) p +#else +#include +#define m512i(p) (__m512i)(p) +#endif // defined(_MSC_VER) +#endif // defined(__AVX512BF16__) + +#if defined(__linux__) || \ + defined(__FreeBSD__) || defined(__NetBSD__) || defined(__OpenBSD__) || \ + (defined(__APPLE__) && !TARGET_OS_TV && !TARGET_OS_WATCH) + +#include +#include +#include +#include +#if defined(__linux__) +#include +#endif + +#if defined(__ANDROID__) +#include +#include +#include + +struct backtrace_state { + void ** current; + void ** end; +}; + +static _Unwind_Reason_Code unwind_callback(struct _Unwind_Context* context, void* arg) { + struct backtrace_state * state = (struct backtrace_state *)arg; + uintptr_t pc = _Unwind_GetIP(context); + if (pc) { + if (state->current == state->end) { + return _URC_END_OF_STACK; + } else { + *state->current++ = (void*)pc; + } + } + return _URC_NO_REASON; +} + +static void ggml_print_backtrace_symbols(void) { + const int max = 100; + void* buffer[max]; + + struct backtrace_state state = {buffer, buffer + max}; + _Unwind_Backtrace(unwind_callback, &state); + + int count = state.current - buffer; + + for (int idx = 0; idx < count; ++idx) { + const void * addr = buffer[idx]; + const char * symbol = ""; + + Dl_info info; + if (dladdr(addr, &info) && info.dli_sname) { + symbol = info.dli_sname; + } + + fprintf(stderr, "%d: %p %s\n", idx, addr, symbol); + } +} +#elif defined(__linux__) && defined(__GLIBC__) +#include +static void ggml_print_backtrace_symbols(void) { + void * trace[100]; + int nptrs = backtrace(trace, sizeof(trace)/sizeof(trace[0])); + backtrace_symbols_fd(trace, nptrs, STDERR_FILENO); +} +#elif defined(__APPLE__) +#include +static void ggml_print_backtrace_symbols(void) { + void * trace[100]; + int nptrs = backtrace(trace, sizeof(trace)/sizeof(trace[0])); + backtrace_symbols_fd(trace, nptrs, STDERR_FILENO); +} +#else +static void ggml_print_backtrace_symbols(void) { + // platform not supported +} +#endif + +void ggml_print_backtrace(void) { + const char * GGML_NO_BACKTRACE = getenv("GGML_NO_BACKTRACE"); + if (GGML_NO_BACKTRACE) { + return; + } +#if defined(__APPLE__) + // On macOS, fork+debugger attachment is problematic due to: + // 1. libdispatch "poisons" forked child processes + // 2. lldb has issues attaching to parent from forked child + // Use simple backtrace() instead to avoid Terminal.app crashes + const char * GGML_BACKTRACE_LLDB = getenv("GGML_BACKTRACE_LLDB"); + if (!GGML_BACKTRACE_LLDB) { + fprintf(stderr, "WARNING: Using native backtrace. Set GGML_BACKTRACE_LLDB for more info.\n"); + fprintf(stderr, "WARNING: GGML_BACKTRACE_LLDB may cause native MacOS Terminal.app to crash.\n"); + fprintf(stderr, "See: https://github.com/ggml-org/llama.cpp/pull/17869\n"); + ggml_print_backtrace_symbols(); + return; + } +#endif +#if defined(__linux__) + FILE * f = fopen("/proc/self/status", "r"); + size_t size = 0; + char * line = NULL; + ssize_t length = 0; + while ((length = getline(&line, &size, f)) > 0) { + if (!strncmp(line, "TracerPid:", sizeof("TracerPid:") - 1) && + (length != sizeof("TracerPid:\t0\n") - 1 || line[length - 2] != '0')) { + // Already being debugged, and the breakpoint is the later abort() + free(line); + fclose(f); + return; + } + } + free(line); + fclose(f); + int lock[2] = { -1, -1 }; + (void) !pipe(lock); // Don't start gdb until after PR_SET_PTRACER +#endif + const int parent_pid = getpid(); + const int child_pid = fork(); + if (child_pid < 0) { // error +#if defined(__linux__) + close(lock[1]); + close(lock[0]); +#endif + return; + } else if (child_pid == 0) { // child + char attach[32]; + snprintf(attach, sizeof(attach), "attach %d", parent_pid); +#if defined(__linux__) + close(lock[1]); + (void) !read(lock[0], lock, 1); + close(lock[0]); +#endif + // try gdb + execlp("gdb", "gdb", "--batch", + "-ex", "set style enabled on", + "-ex", attach, + "-ex", "bt -frame-info source-and-location", + "-ex", "detach", + "-ex", "quit", + (char *) NULL); + // try lldb + execlp("lldb", "lldb", "--batch", + "-o", "bt", + "-o", "quit", + "-p", &attach[sizeof("attach ") - 1], + (char *) NULL); + // gdb failed, fallback to backtrace_symbols + ggml_print_backtrace_symbols(); + _Exit(0); + } else { // parent +#if defined(__linux__) + prctl(PR_SET_PTRACER, child_pid); + close(lock[1]); + close(lock[0]); +#endif + waitpid(child_pid, NULL, 0); + } +} +#else +void ggml_print_backtrace(void) { + // platform not supported +} +#endif + +static ggml_abort_callback_t g_abort_callback = NULL; + +// Set the abort callback (passing null will restore original abort functionality: printing a message to stdout) +GGML_API ggml_abort_callback_t ggml_set_abort_callback(ggml_abort_callback_t callback) { + ggml_abort_callback_t ret_val = g_abort_callback; + g_abort_callback = callback; + return ret_val; +} + +void ggml_abort(const char * file, int line, const char * fmt, ...) { + fflush(stdout); + + char message[2048]; + int offset = snprintf(message, sizeof(message), "%s:%d: ", file, line); + + va_list args; + va_start(args, fmt); + vsnprintf(message + offset, sizeof(message) - offset, fmt, args); + va_end(args); + + if (g_abort_callback) { + g_abort_callback(message); + } else { + // default: print error and backtrace to stderr + fprintf(stderr, "%s\n", message); + ggml_print_backtrace(); + } + + abort(); +} + +// +// logging +// + +struct ggml_logger_state { + ggml_log_callback log_callback; + void * log_callback_user_data; +}; +static struct ggml_logger_state g_logger_state = {ggml_log_callback_default, NULL}; + +static void ggml_log_internal_v(enum ggml_log_level level, const char * format, va_list args) { + if (format == NULL) { + return; + } + va_list args_copy; + va_copy(args_copy, args); + char buffer[128]; + int len = vsnprintf(buffer, 128, format, args); + if (len < 128) { + g_logger_state.log_callback(level, buffer, g_logger_state.log_callback_user_data); + } else { + char * buffer2 = (char *) calloc(len + 1, sizeof(char)); + vsnprintf(buffer2, len + 1, format, args_copy); + buffer2[len] = 0; + g_logger_state.log_callback(level, buffer2, g_logger_state.log_callback_user_data); + free(buffer2); + } + va_end(args_copy); +} + +void ggml_log_internal(enum ggml_log_level level, const char * format, ...) { + va_list args; + va_start(args, format); + ggml_log_internal_v(level, format, args); + va_end(args); +} + +void ggml_log_callback_default(enum ggml_log_level level, const char * text, void * user_data) { + (void) level; + (void) user_data; + fputs(text, stderr); + fflush(stderr); +} + +// +// end of logging block +// + +#ifdef GGML_USE_ACCELERATE +// uncomment to use vDSP for soft max computation +// note: not sure if it is actually faster +//#define GGML_SOFT_MAX_ACCELERATE +#endif + + +void * ggml_aligned_malloc(size_t size) { +#if defined(__s390x__) + const int alignment = 256; +#else + const int alignment = 64; +#endif + +#if defined(_MSC_VER) || defined(__MINGW32__) + return _aligned_malloc(size, alignment); +#else + if (size == 0) { + GGML_LOG_WARN("Behavior may be unexpected when allocating 0 bytes for ggml_aligned_malloc!\n"); + return NULL; + } + void * aligned_memory = NULL; + #ifdef GGML_USE_CPU_HBM + int result = hbw_posix_memalign(&aligned_memory, alignment, size); + #elif TARGET_OS_OSX + GGML_UNUSED(alignment); + kern_return_t alloc_status = vm_allocate((vm_map_t) mach_task_self(), (vm_address_t *) &aligned_memory, size, VM_FLAGS_ANYWHERE); + int result = EFAULT; + switch (alloc_status) { + case KERN_SUCCESS: + result = 0; + break; + case KERN_INVALID_ADDRESS: + result = EINVAL; + break; + case KERN_NO_SPACE: + result = ENOMEM; + break; + default: + result = EFAULT; + break; + } + #else + int result = posix_memalign(&aligned_memory, alignment, size); + #endif + if (result != 0) { + // Handle allocation failure + const char *error_desc = "unknown allocation error"; + switch (result) { + case EINVAL: + error_desc = "invalid alignment value"; + break; + case ENOMEM: + error_desc = "insufficient memory"; + break; + } + GGML_LOG_ERROR("%s: %s (attempted to allocate %6.2f MB)\n", __func__, error_desc, size/(1024.0*1024.0)); + return NULL; + } + return aligned_memory; +#endif +} + +void ggml_aligned_free(void * ptr, size_t size) { + GGML_UNUSED(size); +#if defined(_MSC_VER) || defined(__MINGW32__) + _aligned_free(ptr); +#elif GGML_USE_CPU_HBM + if (ptr != NULL) { + hbw_free(ptr); + } +#elif TARGET_OS_OSX + if (ptr != NULL) { + vm_deallocate((vm_map_t)mach_task_self(), (vm_address_t)ptr, size); + } +#else + free(ptr); +#endif +} + + +inline static void * ggml_malloc(size_t size) { + if (size == 0) { + GGML_LOG_WARN("Behavior may be unexpected when allocating 0 bytes for ggml_malloc!\n"); + return NULL; + } + void * result = malloc(size); + if (result == NULL) { + GGML_LOG_ERROR("%s: failed to allocate %6.2f MB\n", __func__, size/(1024.0*1024.0)); + GGML_ABORT("fatal error"); + } + return result; +} + +// calloc +inline static void * ggml_calloc(size_t num, size_t size) { + if (num == 0 || size == 0) { + GGML_LOG_WARN("Behavior may be unexpected when allocating 0 bytes for ggml_calloc!\n"); + return NULL; + } + void * result = calloc(num, size); + if (result == NULL) { + GGML_LOG_ERROR("%s: failed to allocate %6.2f MB\n", __func__, size/(1024.0*1024.0)); + GGML_ABORT("fatal error"); + } + return result; +} + +#define GGML_MALLOC(size) ggml_malloc(size) +#define GGML_CALLOC(num, size) ggml_calloc(num, size) + +#define GGML_FREE(ptr) free(ptr) + +const char * ggml_status_to_string(enum ggml_status status) { + switch (status) { + case GGML_STATUS_ALLOC_FAILED: return "GGML status: error (failed to allocate memory)"; + case GGML_STATUS_FAILED: return "GGML status: error (operation failed)"; + case GGML_STATUS_SUCCESS: return "GGML status: success"; + case GGML_STATUS_ABORTED: return "GGML status: warning (operation aborted)"; + } + + return "GGML status: unknown"; +} + +float ggml_fp16_to_fp32(ggml_fp16_t x) { +#define ggml_fp16_to_fp32 do_not_use__ggml_fp16_to_fp32__in_ggml + return GGML_FP16_TO_FP32(x); +} + +ggml_fp16_t ggml_fp32_to_fp16(float x) { +#define ggml_fp32_to_fp16 do_not_use__ggml_fp32_to_fp16__in_ggml + return GGML_FP32_TO_FP16(x); +} + +float ggml_bf16_to_fp32(ggml_bf16_t x) { +#define ggml_bf16_to_fp32 do_not_use__ggml_bf16_to_fp32__in_ggml + return GGML_BF16_TO_FP32(x); // it just left shifts +} + +ggml_bf16_t ggml_fp32_to_bf16(float x) { +#define ggml_fp32_to_bf16 do_not_use__ggml_fp32_to_bf16__in_ggml + return GGML_FP32_TO_BF16(x); +} + +void ggml_fp16_to_fp32_row(const ggml_fp16_t * x, float * y, int64_t n) { + for (int64_t i = 0; i < n; i++) { + y[i] = GGML_FP16_TO_FP32(x[i]); + } +} + +void ggml_fp32_to_fp16_row(const float * x, ggml_fp16_t * y, int64_t n) { + int i = 0; + for (; i < n; ++i) { + y[i] = GGML_FP32_TO_FP16(x[i]); + } +} + +void ggml_bf16_to_fp32_row(const ggml_bf16_t * x, float * y, int64_t n) { + int i = 0; + for (; i < n; ++i) { + y[i] = GGML_BF16_TO_FP32(x[i]); + } +} + +void ggml_fp32_to_bf16_row_ref(const float * x, ggml_bf16_t * y, int64_t n) { + for (int i = 0; i < n; i++) { + y[i] = ggml_compute_fp32_to_bf16(x[i]); + } +} + +void ggml_fp32_to_bf16_row(const float * x, ggml_bf16_t * y, int64_t n) { + int i = 0; +#if defined(__AVX512BF16__) + // subnormals are flushed to zero on this platform + for (; i + 32 <= n; i += 32) { + _mm512_storeu_si512( + (__m512i *)(y + i), + m512i(_mm512_cvtne2ps_pbh(_mm512_loadu_ps(x + i + 16), + _mm512_loadu_ps(x + i)))); + } +#endif + for (; i < n; i++) { + y[i] = GGML_FP32_TO_BF16(x[i]); + } +} + +bool ggml_guid_matches(ggml_guid_t guid_a, ggml_guid_t guid_b) { + return memcmp(guid_a, guid_b, sizeof(ggml_guid)) == 0; +} + +const char * ggml_version(void) { + return GGML_VERSION; +} + +const char * ggml_commit(void) { + return GGML_COMMIT; +} + +// +// timing +// + +#if defined(_MSC_VER) || defined(__MINGW32__) +static int64_t timer_freq, timer_start; +void ggml_time_init(void) { + LARGE_INTEGER t; + QueryPerformanceFrequency(&t); + timer_freq = t.QuadPart; + + // The multiplication by 1000 or 1000000 below can cause an overflow if timer_freq + // and the uptime is high enough. + // We subtract the program start time to reduce the likelihood of that happening. + QueryPerformanceCounter(&t); + timer_start = t.QuadPart; +} +int64_t ggml_time_ms(void) { + LARGE_INTEGER t; + QueryPerformanceCounter(&t); + return ((t.QuadPart-timer_start) * 1000) / timer_freq; +} +int64_t ggml_time_us(void) { + LARGE_INTEGER t; + QueryPerformanceCounter(&t); + return ((t.QuadPart-timer_start) * 1000000) / timer_freq; +} +#else +void ggml_time_init(void) {} +int64_t ggml_time_ms(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return (int64_t)ts.tv_sec*1000 + (int64_t)ts.tv_nsec/1000000; +} + +int64_t ggml_time_us(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return (int64_t)ts.tv_sec*1000000 + (int64_t)ts.tv_nsec/1000; +} +#endif + +int64_t ggml_cycles(void) { + return clock(); +} + +int64_t ggml_cycles_per_ms(void) { + return CLOCKS_PER_SEC/1000; +} + +// +// cross-platform UTF-8 file paths +// + +#ifdef _WIN32 +static wchar_t * ggml_mbstowcs(const char * mbs) { + int wlen = MultiByteToWideChar(CP_UTF8, 0, mbs, -1, NULL, 0); + if (!wlen) { + errno = EINVAL; + return NULL; + } + + wchar_t * wbuf = GGML_MALLOC(wlen * sizeof(wchar_t)); + wlen = MultiByteToWideChar(CP_UTF8, 0, mbs, -1, wbuf, wlen); + if (!wlen) { + GGML_FREE(wbuf); + errno = EINVAL; + return NULL; + } + + return wbuf; +} +#endif + +FILE * ggml_fopen(const char * fname, const char * mode) { +#ifdef _WIN32 + FILE * file = NULL; + + // convert fname (UTF-8) + wchar_t * wfname = ggml_mbstowcs(fname); + if (wfname) { + // convert mode (ANSI) + wchar_t * wmode = GGML_MALLOC((strlen(mode) + 1) * sizeof(wchar_t)); + wchar_t * wmode_p = wmode; + do { + *wmode_p++ = (wchar_t)*mode; + } while (*mode++); + + // open file + file = _wfopen(wfname, wmode); + + GGML_FREE(wfname); + GGML_FREE(wmode); + } + + return file; +#else + return fopen(fname, mode); +#endif + +} + +static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { + [GGML_TYPE_I8] = { + .type_name = "i8", + .blck_size = 1, + .type_size = sizeof(int8_t), + .is_quantized = false, + }, + [GGML_TYPE_I16] = { + .type_name = "i16", + .blck_size = 1, + .type_size = sizeof(int16_t), + .is_quantized = false, + }, + [GGML_TYPE_I32] = { + .type_name = "i32", + .blck_size = 1, + .type_size = sizeof(int32_t), + .is_quantized = false, + }, + [GGML_TYPE_I64] = { + .type_name = "i64", + .blck_size = 1, + .type_size = sizeof(int64_t), + .is_quantized = false, + }, + [GGML_TYPE_F64] = { + .type_name = "f64", + .blck_size = 1, + .type_size = sizeof(double), + .is_quantized = false, + }, + [GGML_TYPE_F32] = { + .type_name = "f32", + .blck_size = 1, + .type_size = sizeof(float), + .is_quantized = false, + }, + [GGML_TYPE_F16] = { + .type_name = "f16", + .blck_size = 1, + .type_size = sizeof(ggml_fp16_t), + .is_quantized = false, + .to_float = (ggml_to_float_t) ggml_fp16_to_fp32_row, + .from_float_ref = (ggml_from_float_t) ggml_fp32_to_fp16_row, + }, + [GGML_TYPE_Q1_0] = { + .type_name = "q1_0", + .blck_size = QK1_0, + .type_size = sizeof(block_q1_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q1_0, + .from_float_ref = (ggml_from_float_t) quantize_row_q1_0_ref, + }, + [GGML_TYPE_Q4_0] = { + .type_name = "q4_0", + .blck_size = QK4_0, + .type_size = sizeof(block_q4_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q4_0, + .from_float_ref = (ggml_from_float_t) quantize_row_q4_0_ref, + }, + [GGML_TYPE_Q4_1] = { + .type_name = "q4_1", + .blck_size = QK4_1, + .type_size = sizeof(block_q4_1), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q4_1, + .from_float_ref = (ggml_from_float_t) quantize_row_q4_1_ref, + }, + [4] = { // GGML_TYPE_Q4_2 + .type_name = "DEPRECATED", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [5] = { // GGML_TYPE_Q4_3 + .type_name = "DEPRECATED", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [GGML_TYPE_Q5_0] = { + .type_name = "q5_0", + .blck_size = QK5_0, + .type_size = sizeof(block_q5_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q5_0, + .from_float_ref = (ggml_from_float_t) quantize_row_q5_0_ref, + }, + [GGML_TYPE_Q5_1] = { + .type_name = "q5_1", + .blck_size = QK5_1, + .type_size = sizeof(block_q5_1), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q5_1, + .from_float_ref = (ggml_from_float_t) quantize_row_q5_1_ref, + }, + [GGML_TYPE_Q8_0] = { + .type_name = "q8_0", + .blck_size = QK8_0, + .type_size = sizeof(block_q8_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q8_0, + .from_float_ref = (ggml_from_float_t) quantize_row_q8_0_ref, + }, + [GGML_TYPE_Q8_1] = { + .type_name = "q8_1", + .blck_size = QK8_1, + .type_size = sizeof(block_q8_1), + .is_quantized = true, + .from_float_ref = (ggml_from_float_t) quantize_row_q8_1_ref, + }, + [GGML_TYPE_MXFP4] = { + .type_name = "mxfp4", + .blck_size = QK_MXFP4, + .type_size = sizeof(block_mxfp4), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_mxfp4, + .from_float_ref = (ggml_from_float_t)quantize_row_mxfp4_ref, + }, + [GGML_TYPE_NVFP4] = { + .type_name = "nvfp4", + .blck_size = QK_NVFP4, + .type_size = sizeof(block_nvfp4), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_nvfp4, + .from_float_ref = (ggml_from_float_t)quantize_row_nvfp4_ref, + }, + [GGML_TYPE_Q2_K] = { + .type_name = "q2_K", + .blck_size = QK_K, + .type_size = sizeof(block_q2_K), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q2_K, + .from_float_ref = (ggml_from_float_t) quantize_row_q2_K_ref, + }, + [GGML_TYPE_Q3_K] = { + .type_name = "q3_K", + .blck_size = QK_K, + .type_size = sizeof(block_q3_K), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q3_K, + .from_float_ref = (ggml_from_float_t) quantize_row_q3_K_ref, + }, + [GGML_TYPE_Q4_K] = { + .type_name = "q4_K", + .blck_size = QK_K, + .type_size = sizeof(block_q4_K), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q4_K, + .from_float_ref = (ggml_from_float_t) quantize_row_q4_K_ref, + }, + [GGML_TYPE_Q5_K] = { + .type_name = "q5_K", + .blck_size = QK_K, + .type_size = sizeof(block_q5_K), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q5_K, + .from_float_ref = (ggml_from_float_t) quantize_row_q5_K_ref, + }, + [GGML_TYPE_Q6_K] = { + .type_name = "q6_K", + .blck_size = QK_K, + .type_size = sizeof(block_q6_K), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q6_K, + .from_float_ref = (ggml_from_float_t) quantize_row_q6_K_ref, + }, + [GGML_TYPE_IQ2_XXS] = { + .type_name = "iq2_xxs", + .blck_size = QK_K, + .type_size = sizeof(block_iq2_xxs), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq2_xxs, + .from_float_ref = NULL, + }, + [GGML_TYPE_IQ2_XS] = { + .type_name = "iq2_xs", + .blck_size = QK_K, + .type_size = sizeof(block_iq2_xs), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq2_xs, + .from_float_ref = NULL, + }, + [GGML_TYPE_IQ3_XXS] = { + .type_name = "iq3_xxs", + .blck_size = QK_K, + .type_size = sizeof(block_iq3_xxs), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq3_xxs, + .from_float_ref = (ggml_from_float_t)quantize_row_iq3_xxs_ref, + }, + [GGML_TYPE_IQ3_S] = { + .type_name = "iq3_s", + .blck_size = QK_K, + .type_size = sizeof(block_iq3_s), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq3_s, + .from_float_ref = (ggml_from_float_t)quantize_row_iq3_s_ref, + }, + [GGML_TYPE_IQ2_S] = { + .type_name = "iq2_s", + .blck_size = QK_K, + .type_size = sizeof(block_iq2_s), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq2_s, + .from_float_ref = (ggml_from_float_t)quantize_row_iq2_s_ref, + }, + [GGML_TYPE_IQ1_S] = { + .type_name = "iq1_s", + .blck_size = QK_K, + .type_size = sizeof(block_iq1_s), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq1_s, + .from_float_ref = NULL, + }, + [GGML_TYPE_IQ1_M] = { + .type_name = "iq1_m", + .blck_size = QK_K, + .type_size = sizeof(block_iq1_m), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq1_m, + .from_float_ref = NULL, + }, + [GGML_TYPE_IQ4_NL] = { + .type_name = "iq4_nl", + .blck_size = QK4_NL, + .type_size = sizeof(block_iq4_nl), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq4_nl, + .from_float_ref = (ggml_from_float_t)quantize_row_iq4_nl_ref, + }, + [GGML_TYPE_IQ4_XS] = { + .type_name = "iq4_xs", + .blck_size = QK_K, + .type_size = sizeof(block_iq4_xs), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq4_xs, + .from_float_ref = (ggml_from_float_t)quantize_row_iq4_xs_ref, + }, + [GGML_TYPE_Q8_K] = { + .type_name = "q8_K", + .blck_size = QK_K, + .type_size = sizeof(block_q8_K), + .is_quantized = true, + }, + [GGML_TYPE_BF16] = { + .type_name = "bf16", + .blck_size = 1, + .type_size = sizeof(ggml_bf16_t), + .is_quantized = false, + .to_float = (ggml_to_float_t) ggml_bf16_to_fp32_row, + .from_float_ref = (ggml_from_float_t) ggml_fp32_to_bf16_row_ref, + }, + [31] = { // GGML_TYPE_Q4_0_4_4 + .type_name = "TYPE_Q4_0_4_4 REMOVED, use Q4_0 with runtime repacking", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [32] = { // GGML_TYPE_Q4_0_4_8 + .type_name = "TYPE_Q4_0_4_8 REMOVED, use Q4_0 with runtime repacking", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [33] = { // GGML_TYPE_Q4_0_8_8 + .type_name = "TYPE_Q4_0_8_8 REMOVED, use Q4_0 with runtime repacking", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [GGML_TYPE_TQ1_0] = { + .type_name = "tq1_0", + .blck_size = QK_K, + .type_size = sizeof(block_tq1_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_tq1_0, + .from_float_ref = (ggml_from_float_t) quantize_row_tq1_0_ref, + }, + [GGML_TYPE_TQ2_0] = { + .type_name = "tq2_0", + .blck_size = QK_K, + .type_size = sizeof(block_tq2_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_tq2_0, + .from_float_ref = (ggml_from_float_t) quantize_row_tq2_0_ref, + }, + [36] = { // GGML_TYPE_IQ4_NL_4_4 + .type_name = "TYPE_IQ4_NL_4_4 REMOVED, use IQ4_NL with runtime repacking", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [37] = { // GGML_TYPE_IQ4_NL_4_8 + .type_name = "TYPE_IQ4_NL_4_8 REMOVED, use IQ4_NL with runtime repacking", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [38] = { // GGML_TYPE_IQ4_NL_8_8 + .type_name = "TYPE_IQ4_NL_8_8 REMOVED, use IQ4_NL with runtime repacking", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, +}; + +const struct ggml_type_traits * ggml_get_type_traits(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + return &type_traits[type]; +} + +// +// ggml object +// + +struct ggml_object { + size_t offs; + size_t size; + + struct ggml_object * next; + + enum ggml_object_type type; + + char padding[4]; +}; + +static const size_t GGML_OBJECT_SIZE = sizeof(struct ggml_object); + +// +// ggml context +// + +struct ggml_context { + size_t mem_size; + void * mem_buffer; + bool mem_buffer_owned; + bool no_alloc; + + int n_objects; + + struct ggml_object * objects_begin; + struct ggml_object * objects_end; +}; + +// +// data types +// + +static const char * GGML_OP_NAME[GGML_OP_COUNT] = { + "NONE", + + "DUP", + "ADD", + "ADD_ID", + "ADD1", + "ACC", + "SUB", + "MUL", + "DIV", + "SQR", + "SQRT", + "LOG", + "SIN", + "COS", + "SUM", + "SUM_ROWS", + "CUMSUM", + "MEAN", + "ARGMAX", + "COUNT_EQUAL", + "REPEAT", + "REPEAT_BACK", + "CONCAT", + "SILU_BACK", + "NORM", + "RMS_NORM", + "RMS_NORM_BACK", + "GROUP_NORM", + "L2_NORM", + + "MUL_MAT", + "MUL_MAT_ID", + "OUT_PROD", + + "SCALE", + "SET", + "CPY", + "CONT", + "RESHAPE", + "VIEW", + "PERMUTE", + "TRANSPOSE", + "GET_ROWS", + "GET_ROWS_BACK", + "SET_ROWS", + "DIAG", + "DIAG_MASK_INF", + "DIAG_MASK_ZERO", + "SOFT_MAX", + "SOFT_MAX_BACK", + "ROPE", + "ROPE_BACK", + "CLAMP", + "CONV_TRANSPOSE_1D", + "IM2COL", + "IM2COL_BACK", + "IM2COL_3D", + "CONV_2D", + "CONV_3D", + "CONV_2D_DW", + "CONV_TRANSPOSE_2D", + "POOL_1D", + "POOL_2D", + "POOL_2D_BACK", + "UPSCALE", + "PAD", + "PAD_REFLECT_1D", + "ROLL", + "ARANGE", + "TIMESTEP_EMBEDDING", + "ARGSORT", + "TOP_K", + "LEAKY_RELU", + "TRI", + "FILL", + + "FLASH_ATTN_EXT", + "FLASH_ATTN_BACK", + "SSM_CONV", + "SSM_SCAN", + "WIN_PART", + "WIN_UNPART", + "GET_REL_POS", + "ADD_REL_POS", + "RWKV_WKV6", + "GATED_LINEAR_ATTN", + "RWKV_WKV7", + "SOLVE_TRI", + "GATED_DELTA_NET", + + "UNARY", + + "MAP_CUSTOM1", + "MAP_CUSTOM2", + "MAP_CUSTOM3", + + "CUSTOM", + + "CROSS_ENTROPY_LOSS", + "CROSS_ENTROPY_LOSS_BACK", + "OPT_STEP_ADAMW", + "OPT_STEP_SGD", + + "GLU", +}; + +static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT != 96"); + +static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "none", + + "x", + "x+y", + "x[i]+y", + "x+y", + "view(x,nb,offset)+=y->x", + "x-y", + "x*y", + "x/y", + "x^2", + "√x", + "log(x)", + "sin(x)", + "cos(x)", + "Σx", + "Σx_k", + "cumsum(x)", + "Σx/n", + "argmax(x)", + "count_equal(x)", + "repeat(x)", + "repeat_back(x)", + "concat(x, y)", + "silu_back(x)", + "norm(x)", + "rms_norm(x)", + "rms_norm_back(x)", + "group_norm(x)", + "l2_norm(x)", + + "X*Y", + "X[i]*Y", + "X*Y", + + "x*v", + "y-\\>view(x)", + "x-\\>y", + "cont(x)", + "reshape(x)", + "view(x)", + "permute(x)", + "transpose(x)", + "get_rows(x)", + "get_rows_back(x)", + "set_rows(x)", + "diag(x)", + "diag_mask_inf(x)", + "diag_mask_zero(x)", + "soft_max(x)", + "soft_max_back(x)", + "rope(x)", + "rope_back(x)", + "clamp(x)", + "conv_transpose_1d(x)", + "im2col(x)", + "im2col_back(x)", + "im2col_3d(x)", + "conv_2d(x)", + "conv_3d(x)", + "conv_2d_dw(x)", + "conv_transpose_2d(x)", + "pool_1d(x)", + "pool_2d(x)", + "pool_2d_back(x)", + "upscale(x)", + "pad(x)", + "pad_reflect_1d(x)", + "roll(x)", + "arange(start, stop, step)", + "timestep_embedding(timesteps, dim, max_period)", + "argsort(x)", + "top_k(x)", + "leaky_relu(x)", + "tri(x)", + "fill(x, c)", + + "flash_attn_ext(x)", + "flash_attn_back(x)", + "ssm_conv(x)", + "ssm_scan(x)", + "win_part(x)", + "win_unpart(x)", + "get_rel_pos(x)", + "add_rel_pos(x)", + "rwkv_wkv6(k, v, r, tf, td, s)", + "gated_linear_attn(k, v, q, gate, s)", + "rwkv_wkv7(r, w, k, v, a, b, s)", + "A X = B, A triangular, solve X", + "gated_delta_net(q, k, v, g, beta, s)", + + "unary(x)", + + "map_custom(x)", + "map_custom(x,y)", + "map_custom(x,y,z)", + + "custom(x)", + + "cross_entropy_loss(x,y)", + "cross_entropy_loss_back(x,y)", + "adamw(x)", + "sgd(x)", + + "glu(x)", +}; + +static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT != 96"); + +static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); + +static const char * GGML_UNARY_OP_NAME[GGML_UNARY_OP_COUNT] = { + "ABS", + "SGN", + "NEG", + "STEP", + "TANH", + "ELU", + "RELU", + "SIGMOID", + "GELU", + "GELU_QUICK", + "SILU", + "HARDSWISH", + "HARDSIGMOID", + "EXP", + "EXPM1", + "SOFTPLUS", + "GELU_ERF", + "XIELU", + "FLOOR", + "CEIL", + "ROUND", + "TRUNC", +}; + +static_assert(GGML_UNARY_OP_COUNT == 22, "GGML_UNARY_OP_COUNT != 22"); + +static const char * GGML_GLU_OP_NAME[GGML_GLU_OP_COUNT] = { + "REGLU", + "GEGLU", + "SWIGLU", + "SWIGLU_OAI", + "GEGLU_ERF", + "GEGLU_QUICK", +}; + +static_assert(GGML_GLU_OP_COUNT == 6, "GGML_GLU_OP_COUNT != 6"); + + +static_assert(sizeof(struct ggml_object)%GGML_MEM_ALIGN == 0, "ggml_object size must be a multiple of GGML_MEM_ALIGN"); +static_assert(sizeof(struct ggml_tensor)%GGML_MEM_ALIGN == 0, "ggml_tensor size must be a multiple of GGML_MEM_ALIGN"); + + +//////////////////////////////////////////////////////////////////////////////// + +void ggml_print_object(const struct ggml_object * obj) { + GGML_LOG_INFO(" - ggml_object: type = %d, offset = %zu, size = %zu, next = %p\n", + obj->type, obj->offs, obj->size, (const void *) obj->next); +} + +void ggml_print_objects(const struct ggml_context * ctx) { + struct ggml_object * obj = ctx->objects_begin; + + GGML_LOG_INFO("%s: objects in context %p:\n", __func__, (const void *) ctx); + + while (obj != NULL) { + ggml_print_object(obj); + obj = obj->next; + } + + GGML_LOG_INFO("%s: --- end ---\n", __func__); +} + +int64_t ggml_nelements(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return tensor->ne[0]*tensor->ne[1]*tensor->ne[2]*tensor->ne[3]; +} + +int64_t ggml_nrows(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return tensor->ne[1]*tensor->ne[2]*tensor->ne[3]; +} + +size_t ggml_nbytes(const struct ggml_tensor * tensor) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (tensor->ne[i] <= 0) { + return 0; + } + } + + size_t nbytes; + const size_t blck_size = ggml_blck_size(tensor->type); + if (blck_size == 1) { + nbytes = ggml_type_size(tensor->type); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + nbytes += (tensor->ne[i] - 1)*tensor->nb[i]; + } + } + else { + nbytes = tensor->ne[0]*tensor->nb[0]/blck_size; + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + nbytes += (tensor->ne[i] - 1)*tensor->nb[i]; + } + } + + return nbytes; +} + +size_t ggml_nbytes_pad(const struct ggml_tensor * tensor) { + return GGML_PAD(ggml_nbytes(tensor), GGML_MEM_ALIGN); +} + +int64_t ggml_blck_size(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + return type_traits[type].blck_size; +} + +size_t ggml_type_size(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + return type_traits[type].type_size; +} + +size_t ggml_row_size(enum ggml_type type, int64_t ne) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + assert(ne % ggml_blck_size(type) == 0); + return ggml_type_size(type)*ne/ggml_blck_size(type); +} + +double ggml_type_sizef(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + return ((double)(type_traits[type].type_size))/type_traits[type].blck_size; +} + +const char * ggml_type_name(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + return type_traits[type].type_name; +} + +bool ggml_is_quantized(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + return type_traits[type].is_quantized; +} + +const char * ggml_op_name(enum ggml_op op) { + return GGML_OP_NAME[op]; +} + +const char * ggml_op_symbol(enum ggml_op op) { + return GGML_OP_SYMBOL[op]; +} + +const char * ggml_unary_op_name(enum ggml_unary_op op) { + return GGML_UNARY_OP_NAME[op]; +} + +const char * ggml_glu_op_name(enum ggml_glu_op op) { + return GGML_GLU_OP_NAME[op]; +} + +const char * ggml_op_desc(const struct ggml_tensor * t) { + if (t->op == GGML_OP_UNARY) { + enum ggml_unary_op uop = ggml_get_unary_op(t); + return ggml_unary_op_name(uop); + } + if (t->op == GGML_OP_GLU) { + enum ggml_glu_op gop = ggml_get_glu_op(t); + return ggml_glu_op_name(gop); + } + return ggml_op_name(t->op); +} + +size_t ggml_element_size(const struct ggml_tensor * tensor) { + return ggml_type_size(tensor->type); +} + +bool ggml_is_scalar(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return tensor->ne[0] == 1 && tensor->ne[1] == 1 && tensor->ne[2] == 1 && tensor->ne[3] == 1; +} + +bool ggml_is_vector(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return tensor->ne[1] == 1 && tensor->ne[2] == 1 && tensor->ne[3] == 1; +} + +bool ggml_is_matrix(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return tensor->ne[2] == 1 && tensor->ne[3] == 1; +} + +bool ggml_is_3d(const struct ggml_tensor * tensor) { + return tensor->ne[3] == 1; +} + +int ggml_n_dims(const struct ggml_tensor * tensor) { + for (int i = GGML_MAX_DIMS - 1; i >= 1; --i) { + if (tensor->ne[i] > 1) { + return i + 1; + } + } + return 1; +} + +enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype) { + enum ggml_type wtype = GGML_TYPE_COUNT; + + switch (ftype) { + case GGML_FTYPE_ALL_F32: wtype = GGML_TYPE_F32; break; + case GGML_FTYPE_MOSTLY_F16: wtype = GGML_TYPE_F16; break; + case GGML_FTYPE_MOSTLY_BF16: wtype = GGML_TYPE_BF16; break; + case GGML_FTYPE_MOSTLY_Q4_0: wtype = GGML_TYPE_Q4_0; break; + case GGML_FTYPE_MOSTLY_Q4_1: wtype = GGML_TYPE_Q4_1; break; + case GGML_FTYPE_MOSTLY_Q1_0: wtype = GGML_TYPE_Q1_0; break; + case GGML_FTYPE_MOSTLY_Q5_0: wtype = GGML_TYPE_Q5_0; break; + case GGML_FTYPE_MOSTLY_Q5_1: wtype = GGML_TYPE_Q5_1; break; + case GGML_FTYPE_MOSTLY_Q8_0: wtype = GGML_TYPE_Q8_0; break; + case GGML_FTYPE_MOSTLY_MXFP4: wtype = GGML_TYPE_MXFP4; break; + case GGML_FTYPE_MOSTLY_NVFP4: wtype = GGML_TYPE_NVFP4; break; + case GGML_FTYPE_MOSTLY_Q2_K: wtype = GGML_TYPE_Q2_K; break; + case GGML_FTYPE_MOSTLY_Q3_K: wtype = GGML_TYPE_Q3_K; break; + case GGML_FTYPE_MOSTLY_Q4_K: wtype = GGML_TYPE_Q4_K; break; + case GGML_FTYPE_MOSTLY_Q5_K: wtype = GGML_TYPE_Q5_K; break; + case GGML_FTYPE_MOSTLY_Q6_K: wtype = GGML_TYPE_Q6_K; break; + case GGML_FTYPE_MOSTLY_IQ2_XXS: wtype = GGML_TYPE_IQ2_XXS; break; + case GGML_FTYPE_MOSTLY_IQ2_XS: wtype = GGML_TYPE_IQ2_XS; break; + case GGML_FTYPE_MOSTLY_IQ3_XXS: wtype = GGML_TYPE_IQ3_XXS; break; + case GGML_FTYPE_MOSTLY_IQ1_S: wtype = GGML_TYPE_IQ1_S; break; + case GGML_FTYPE_MOSTLY_IQ1_M: wtype = GGML_TYPE_IQ1_M; break; + case GGML_FTYPE_MOSTLY_IQ4_NL: wtype = GGML_TYPE_IQ4_NL; break; + case GGML_FTYPE_MOSTLY_IQ4_XS: wtype = GGML_TYPE_IQ4_XS; break; + case GGML_FTYPE_MOSTLY_IQ3_S: wtype = GGML_TYPE_IQ3_S; break; + case GGML_FTYPE_MOSTLY_IQ2_S: wtype = GGML_TYPE_IQ2_S; break; + case GGML_FTYPE_UNKNOWN: wtype = GGML_TYPE_COUNT; break; + case GGML_FTYPE_MOSTLY_Q4_1_SOME_F16: wtype = GGML_TYPE_COUNT; break; + } + + GGML_ASSERT(wtype != GGML_TYPE_COUNT); + + return wtype; +} + +size_t ggml_tensor_overhead(void) { + return GGML_OBJECT_SIZE + GGML_TENSOR_SIZE; +} + +bool ggml_is_transposed(const struct ggml_tensor * tensor) { + return tensor->nb[0] > tensor->nb[1]; +} + +static bool ggml_is_contiguous_n(const struct ggml_tensor * tensor, int n) { + size_t next_nb = ggml_type_size(tensor->type); + if (tensor->ne[0] != ggml_blck_size(tensor->type) && tensor->nb[0] != next_nb) { + return false; + } + next_nb *= tensor->ne[0]/ggml_blck_size(tensor->type); + for (int i = 1; i < GGML_MAX_DIMS; i++) { + if (i > n) { + if (tensor->ne[i] != 1 && tensor->nb[i] != next_nb) { + return false; + } + next_nb *= tensor->ne[i]; + } else { + // this dimension does not need to be contiguous + next_nb = tensor->ne[i]*tensor->nb[i]; + } + } + return true; +} + +bool ggml_is_contiguous(const struct ggml_tensor * tensor) { + return ggml_is_contiguous_0(tensor); +} + +bool ggml_is_contiguous_0(const struct ggml_tensor * tensor) { + return ggml_is_contiguous_n(tensor, 0); +} + +bool ggml_is_contiguous_1(const struct ggml_tensor * tensor) { + return ggml_is_contiguous_n(tensor, 1); +} + +bool ggml_is_contiguous_2(const struct ggml_tensor * tensor) { + return ggml_is_contiguous_n(tensor, 2); +} + +bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor) { + return ggml_nbytes(tensor) == ggml_nelements(tensor) * ggml_type_size(tensor->type)/ggml_blck_size(tensor->type); +} + +bool ggml_is_permuted(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return tensor->nb[0] > tensor->nb[1] || tensor->nb[1] > tensor->nb[2] || tensor->nb[2] > tensor->nb[3]; +} + +bool ggml_is_contiguous_channels(const struct ggml_tensor * tensor) { + return + tensor->nb[0] > tensor->nb[2] && + tensor->nb[1] > tensor->nb[0] && + tensor->nb[2] == ggml_type_size(tensor->type); +} + +bool ggml_is_contiguous_rows(const struct ggml_tensor * tensor) { + return + tensor->ne[0] == ggml_blck_size(tensor->type) || + tensor->nb[0] == ggml_type_size(tensor->type); +} + +static inline bool ggml_is_padded_1d(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return + tensor->nb[0] == ggml_type_size(tensor->type) && + tensor->nb[2] == tensor->nb[1]*tensor->ne[1] && + tensor->nb[3] == tensor->nb[2]*tensor->ne[2]; +} + +bool ggml_is_empty(const struct ggml_tensor * tensor) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (tensor->ne[i] == 0) { + // empty if any dimension has no elements + return true; + } + } + return false; +} + +bool ggml_are_same_shape(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return + (t0->ne[0] == t1->ne[0]) && + (t0->ne[1] == t1->ne[1]) && + (t0->ne[2] == t1->ne[2]) && + (t0->ne[3] == t1->ne[3]); +} + +bool ggml_are_same_stride(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return + (t0->nb[0] == t1->nb[0]) && + (t0->nb[1] == t1->nb[1]) && + (t0->nb[2] == t1->nb[2]) && + (t0->nb[3] == t1->nb[3]); +} + +bool ggml_is_view(const struct ggml_tensor * t) { + return ggml_impl_is_view(t); +} + +// check if t1 can be represented as a repetition of t0 +bool ggml_can_repeat(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return ggml_is_empty(t0) ? ggml_is_empty(t1) : + (t1->ne[0]%t0->ne[0] == 0) && + (t1->ne[1]%t0->ne[1] == 0) && + (t1->ne[2]%t0->ne[2] == 0) && + (t1->ne[3]%t0->ne[3] == 0); +} + +static inline bool ggml_can_repeat_rows(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return (t0->ne[0] == t1->ne[0]) && ggml_can_repeat(t0, t1); +} + +// assert that pointer is aligned to GGML_MEM_ALIGN +#define GGML_ASSERT_ALIGNED(ptr) \ + GGML_ASSERT(((uintptr_t) (ptr))%GGML_MEM_ALIGN == 0) + +//////////////////////////////////////////////////////////////////////////////// + +struct ggml_context * ggml_init(struct ggml_init_params params) { + static bool is_first_call = true; + + ggml_critical_section_start(); + + if (is_first_call) { + // initialize time system (required on Windows) + ggml_time_init(); + + is_first_call = false; + } + + ggml_critical_section_end(); + + struct ggml_context * ctx = GGML_MALLOC(sizeof(struct ggml_context)); + + // allow to call ggml_init with 0 size + if (params.mem_size == 0) { + params.mem_size = GGML_MEM_ALIGN; + } + + const size_t mem_size = params.mem_buffer ? params.mem_size : GGML_PAD(params.mem_size, GGML_MEM_ALIGN); + + *ctx = (struct ggml_context) { + /*.mem_size =*/ mem_size, + /*.mem_buffer =*/ params.mem_buffer ? params.mem_buffer : ggml_aligned_malloc(mem_size), + /*.mem_buffer_owned =*/ params.mem_buffer ? false : true, + /*.no_alloc =*/ params.no_alloc, + /*.n_objects =*/ 0, + /*.objects_begin =*/ NULL, + /*.objects_end =*/ NULL, + }; + + GGML_ASSERT(ctx->mem_buffer != NULL); + + GGML_ASSERT_ALIGNED(ctx->mem_buffer); + + GGML_PRINT_DEBUG("%s: context initialized\n", __func__); + + return ctx; +} + +void ggml_reset(struct ggml_context * ctx) { + if (ctx == NULL) { + return; + } + + ctx->n_objects = 0; + ctx->objects_begin = NULL; + ctx->objects_end = NULL; +} + +void ggml_free(struct ggml_context * ctx) { + if (ctx == NULL) { + return; + } + + if (ctx->mem_buffer_owned) { + ggml_aligned_free(ctx->mem_buffer, ctx->mem_size); + } + + GGML_FREE(ctx); +} + +size_t ggml_used_mem(const struct ggml_context * ctx) { + return ctx->objects_end == NULL ? 0 : ctx->objects_end->offs + ctx->objects_end->size; +} + +bool ggml_get_no_alloc(struct ggml_context * ctx) { + return ctx->no_alloc; +} + +void ggml_set_no_alloc(struct ggml_context * ctx, bool no_alloc) { + ctx->no_alloc = no_alloc; +} + +void * ggml_get_mem_buffer(const struct ggml_context * ctx) { + return ctx->mem_buffer; +} + +size_t ggml_get_mem_size(const struct ggml_context * ctx) { + return ctx->mem_size; +} + +size_t ggml_get_max_tensor_size(const struct ggml_context * ctx) { + size_t max_size = 0; + + for (struct ggml_tensor * tensor = ggml_get_first_tensor(ctx); tensor != NULL; tensor = ggml_get_next_tensor(ctx, tensor)) { + size_t bytes = ggml_nbytes(tensor); + max_size = MAX(max_size, bytes); + } + + return max_size; +} + +//////////////////////////////////////////////////////////////////////////////// + +static struct ggml_object * ggml_new_object(struct ggml_context * ctx, enum ggml_object_type type, size_t size) { + // always insert objects at the end of the context's memory pool + struct ggml_object * obj_cur = ctx->objects_end; + + const size_t cur_offs = obj_cur == NULL ? 0 : obj_cur->offs; + const size_t cur_size = obj_cur == NULL ? 0 : obj_cur->size; + const size_t cur_end = cur_offs + cur_size; + + // align to GGML_MEM_ALIGN + GGML_ASSERT(size <= SIZE_MAX - (GGML_MEM_ALIGN - 1)); + size_t size_needed = GGML_PAD(size, GGML_MEM_ALIGN); + + char * const mem_buffer = ctx->mem_buffer; + struct ggml_object * const obj_new = (struct ggml_object *)(mem_buffer + cur_end); + + // integer overflow checks + if (cur_end > SIZE_MAX - size_needed) { + GGML_LOG_WARN("%s: overflow detected in cur_end (%zu) + size_needed (%zu)\n", __func__, cur_end, size_needed); + return NULL; + } + if (cur_end + size_needed > SIZE_MAX - GGML_OBJECT_SIZE) { + GGML_LOG_WARN("%s: overflow detected in cur_end (%zu) + size_needed (%zu) + GGML_OBJECT_SIZE (%zu)\n", __func__, + cur_end, size_needed, (size_t) GGML_OBJECT_SIZE); + return NULL; + } + + if (cur_end + size_needed + GGML_OBJECT_SIZE > ctx->mem_size) { + GGML_LOG_WARN("%s: not enough space in the context's memory pool (needed %zu, available %zu)\n", + __func__, cur_end + size_needed + GGML_OBJECT_SIZE, ctx->mem_size); +#ifndef NDEBUG + GGML_ABORT("not enough space in the context's memory pool"); +#endif + return NULL; + } + + *obj_new = (struct ggml_object) { + .offs = cur_end + GGML_OBJECT_SIZE, + .size = size_needed, + .next = NULL, + .type = type, + }; + + GGML_ASSERT_ALIGNED(mem_buffer + obj_new->offs); + + if (obj_cur != NULL) { + obj_cur->next = obj_new; + } else { + // this is the first object in this context + ctx->objects_begin = obj_new; + } + + ctx->objects_end = obj_new; + + //printf("%s: inserted new object at %zu, size = %zu\n", __func__, cur_end, obj_new->size); + + return obj_new; +} + +static struct ggml_tensor * ggml_new_tensor_impl( + struct ggml_context * ctx, + enum ggml_type type, + int n_dims, + const int64_t * ne, + struct ggml_tensor * view_src, + size_t view_offs) { + + GGML_ASSERT(type >= 0 && type < GGML_TYPE_COUNT); + GGML_ASSERT(n_dims >= 1 && n_dims <= GGML_MAX_DIMS); + + // find the base tensor and absolute offset + if (view_src != NULL && view_src->view_src != NULL) { + view_offs += view_src->view_offs; + view_src = view_src->view_src; + } + + size_t data_size = ggml_row_size(type, ne[0]); + for (int i = 1; i < n_dims; i++) { + data_size *= ne[i]; + } + + GGML_ASSERT(view_src == NULL || data_size == 0 || data_size + view_offs <= ggml_nbytes(view_src)); + + void * data = view_src != NULL ? view_src->data : NULL; + if (data != NULL) { + data = (char *) data + view_offs; + } + + size_t obj_alloc_size = 0; + + if (view_src == NULL && !ctx->no_alloc) { + // allocate tensor data in the context's memory pool + obj_alloc_size = data_size; + } + + GGML_ASSERT(GGML_TENSOR_SIZE <= SIZE_MAX - obj_alloc_size); + + struct ggml_object * const obj_new = ggml_new_object(ctx, GGML_OBJECT_TYPE_TENSOR, GGML_TENSOR_SIZE + obj_alloc_size); + GGML_ASSERT(obj_new); + + struct ggml_tensor * const result = (struct ggml_tensor *)((char *)ctx->mem_buffer + obj_new->offs); + + *result = (struct ggml_tensor) { + /*.type =*/ type, + /*.buffer =*/ NULL, + /*.ne =*/ { 1, 1, 1, 1 }, + /*.nb =*/ { 0, 0, 0, 0 }, + /*.op =*/ GGML_OP_NONE, + /*.op_params =*/ { 0 }, + /*.flags =*/ 0, + /*.src =*/ { NULL }, + /*.view_src =*/ view_src, + /*.view_offs =*/ view_offs, + /*.data =*/ obj_alloc_size > 0 ? (void *)(result + 1) : data, + /*.name =*/ { 0 }, + /*.extra =*/ NULL, + /*.padding =*/ { 0 }, + }; + + // TODO: this should not be needed as long as we don't rely on aligned SIMD loads + //GGML_ASSERT_ALIGNED(result->data); + + for (int i = 0; i < n_dims; i++) { + result->ne[i] = ne[i]; + } + + result->nb[0] = ggml_type_size(type); + result->nb[1] = result->nb[0]*(result->ne[0]/ggml_blck_size(type)); + for (int i = 2; i < GGML_MAX_DIMS; i++) { + result->nb[i] = result->nb[i - 1]*result->ne[i - 1]; + } + + ctx->n_objects++; + + return result; +} + +struct ggml_tensor * ggml_new_tensor( + struct ggml_context * ctx, + enum ggml_type type, + int n_dims, + const int64_t * ne) { + return ggml_new_tensor_impl(ctx, type, n_dims, ne, NULL, 0); +} + +struct ggml_tensor * ggml_new_tensor_1d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0) { + return ggml_new_tensor(ctx, type, 1, &ne0); +} + +struct ggml_tensor * ggml_new_tensor_2d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1) { + const int64_t ne[2] = { ne0, ne1 }; + return ggml_new_tensor(ctx, type, 2, ne); +} + +struct ggml_tensor * ggml_new_tensor_3d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1, + int64_t ne2) { + const int64_t ne[3] = { ne0, ne1, ne2 }; + return ggml_new_tensor(ctx, type, 3, ne); +} + +struct ggml_tensor * ggml_new_tensor_4d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3) { + const int64_t ne[4] = { ne0, ne1, ne2, ne3 }; + return ggml_new_tensor(ctx, type, 4, ne); +} + +void * ggml_new_buffer(struct ggml_context * ctx, size_t nbytes) { + struct ggml_object * obj = ggml_new_object(ctx, GGML_OBJECT_TYPE_WORK_BUFFER, nbytes); + + return (uint8_t *)ctx->mem_buffer + obj->offs; +} + +struct ggml_tensor * ggml_dup_tensor(struct ggml_context * ctx, const struct ggml_tensor * src) { + return ggml_new_tensor(ctx, src->type, GGML_MAX_DIMS, src->ne); +} + +void ggml_unravel_index(const struct ggml_tensor * tensor, int64_t i, int64_t * i0, int64_t * i1, int64_t * i2, int64_t * i3) { + const int64_t ne2 = tensor->ne[2]; + const int64_t ne1 = tensor->ne[1]; + const int64_t ne0 = tensor->ne[0]; + + const int64_t i3_ = (i/(ne2*ne1*ne0)); + const int64_t i2_ = (i - i3_*ne2*ne1*ne0)/(ne1*ne0); + const int64_t i1_ = (i - i3_*ne2*ne1*ne0 - i2_*ne1*ne0)/ne0; + const int64_t i0_ = (i - i3_*ne2*ne1*ne0 - i2_*ne1*ne0 - i1_*ne0); + + if (i0) { + * i0 = i0_; + } + if (i1) { + * i1 = i1_; + } + if (i2) { + * i2 = i2_; + } + if (i3) { + * i3 = i3_; + } +} + +void * ggml_get_data(const struct ggml_tensor * tensor) { + return tensor->data; +} + +float * ggml_get_data_f32(const struct ggml_tensor * tensor) { + assert(tensor->type == GGML_TYPE_F32); + return (float *)(tensor->data); +} + +enum ggml_unary_op ggml_get_unary_op(const struct ggml_tensor * tensor) { + GGML_ASSERT(tensor->op == GGML_OP_UNARY); + return (enum ggml_unary_op) ggml_get_op_params_i32(tensor, 0); +} + +enum ggml_glu_op ggml_get_glu_op(const struct ggml_tensor * tensor) { + GGML_ASSERT(tensor->op == GGML_OP_GLU); + return (enum ggml_glu_op) ggml_get_op_params_i32(tensor, 0); +} + +const char * ggml_get_name(const struct ggml_tensor * tensor) { + return tensor->name; +} + +struct ggml_tensor * ggml_set_name(struct ggml_tensor * tensor, const char * name) { + size_t i; + for (i = 0; i < sizeof(tensor->name) - 1 && name[i] != '\0'; i++) { + tensor->name[i] = name[i]; + } + tensor->name[i] = '\0'; + return tensor; +} + +struct ggml_tensor * ggml_format_name(struct ggml_tensor * tensor, const char * fmt, ...) { + va_list args; + va_start(args, fmt); + vsnprintf(tensor->name, sizeof(tensor->name), fmt, args); + va_end(args); + return tensor; +} + +struct ggml_tensor * ggml_view_tensor( + struct ggml_context * ctx, + struct ggml_tensor * src) { + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, src->type, GGML_MAX_DIMS, src->ne, src, 0); + ggml_format_name(result, "%s (view)", src->name); + + for (int i = 0; i < GGML_MAX_DIMS; i++) { + result->nb[i] = src->nb[i]; + } + + return result; +} + +struct ggml_tensor * ggml_get_first_tensor(const struct ggml_context * ctx) { + struct ggml_object * obj = ctx->objects_begin; + + char * const mem_buffer = ctx->mem_buffer; + + while (obj != NULL) { + if (obj->type == GGML_OBJECT_TYPE_TENSOR) { + return (struct ggml_tensor *)(mem_buffer + obj->offs); + } + + obj = obj->next; + } + + return NULL; +} + +struct ggml_tensor * ggml_get_next_tensor(const struct ggml_context * ctx, struct ggml_tensor * tensor) { + struct ggml_object * obj = (struct ggml_object *) ((char *)tensor - GGML_OBJECT_SIZE); + obj = obj->next; + + char * const mem_buffer = ctx->mem_buffer; + + while (obj != NULL) { + if (obj->type == GGML_OBJECT_TYPE_TENSOR) { + return (struct ggml_tensor *)(mem_buffer + obj->offs); + } + + obj = obj->next; + } + + return NULL; +} + +struct ggml_tensor * ggml_get_tensor(struct ggml_context * ctx, const char * name) { + struct ggml_object * obj = ctx->objects_begin; + + char * const mem_buffer = ctx->mem_buffer; + + while (obj != NULL) { + if (obj->type == GGML_OBJECT_TYPE_TENSOR) { + struct ggml_tensor * cur = (struct ggml_tensor *)(mem_buffer + obj->offs); + if (strcmp(cur->name, name) == 0) { + return cur; + } + } + + obj = obj->next; + } + + return NULL; +} + +//////////////////////////////////////////////////////////////////////////////// + +// ggml_dup + +static struct ggml_tensor * ggml_dup_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_DUP; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_dup( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_dup_impl(ctx, a, false); +} + +struct ggml_tensor * ggml_dup_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_dup_impl(ctx, a, true); +} + +// ggml_add + +static struct ggml_tensor * ggml_add_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool inplace) { + GGML_ASSERT(ggml_can_repeat(b, a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_ADD; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_add( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_add_impl(ctx, a, b, false); +} + +struct ggml_tensor * ggml_add_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_add_impl(ctx, a, b, true); +} + +// ggml_add_cast + +static struct ggml_tensor * ggml_add_cast_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + enum ggml_type type) { + // TODO: support less-strict constraint + // GGML_ASSERT(ggml_can_repeat(b, a)); + GGML_ASSERT(ggml_can_repeat_rows(b, a)); + + // currently only supported for quantized input and f16 + GGML_ASSERT(ggml_is_quantized(a->type) || + a->type == GGML_TYPE_F16 || + a->type == GGML_TYPE_BF16); + + struct ggml_tensor * result = ggml_new_tensor(ctx, type, GGML_MAX_DIMS, a->ne); + + result->op = GGML_OP_ADD; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_add_cast( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + enum ggml_type type) { + return ggml_add_cast_impl(ctx, a, b, type); +} + +struct ggml_tensor * ggml_add_id( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * ids) { + + GGML_ASSERT(a->ne[0] == b->ne[0]); + GGML_ASSERT(a->ne[1] == ids->ne[0]); + GGML_ASSERT(a->ne[2] == ids->ne[1]); + GGML_ASSERT(ids->type == GGML_TYPE_I32); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_ADD_ID; + result->src[0] = a; + result->src[1] = b; + result->src[2] = ids; + + return result; +} + +// ggml_add1 + +static struct ggml_tensor * ggml_add1_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool inplace) { + GGML_ASSERT(ggml_is_scalar(b)); + GGML_ASSERT(ggml_is_padded_1d(a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_ADD1; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_add1( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_add1_impl(ctx, a, b, false); +} + +struct ggml_tensor * ggml_add1_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_add1_impl(ctx, a, b, true); +} + +// ggml_acc + +static struct ggml_tensor * ggml_acc_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset, + bool inplace) { + GGML_ASSERT(ggml_nelements(b) <= ggml_nelements(a)); + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(a->type == GGML_TYPE_F32); + GGML_ASSERT(b->type == GGML_TYPE_F32); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + int32_t params[] = { nb1, nb2, nb3, offset, inplace ? 1 : 0 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_ACC; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_acc( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset) { + return ggml_acc_impl(ctx, a, b, nb1, nb2, nb3, offset, false); +} + +struct ggml_tensor * ggml_acc_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset) { + return ggml_acc_impl(ctx, a, b, nb1, nb2, nb3, offset, true); +} + +// ggml_sub + +static struct ggml_tensor * ggml_sub_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool inplace) { + GGML_ASSERT(ggml_can_repeat(b, a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_SUB; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_sub( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_sub_impl(ctx, a, b, false); +} + +struct ggml_tensor * ggml_sub_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_sub_impl(ctx, a, b, true); +} + +// ggml_mul + +static struct ggml_tensor * ggml_mul_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool inplace) { + GGML_ASSERT(ggml_can_repeat(b, a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_MUL; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_mul( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_mul_impl(ctx, a, b, false); +} + +struct ggml_tensor * ggml_mul_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_mul_impl(ctx, a, b, true); +} + +// ggml_div + +static struct ggml_tensor * ggml_div_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool inplace) { + GGML_ASSERT(ggml_can_repeat(b, a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_DIV; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_div( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_div_impl(ctx, a, b, false); +} + +struct ggml_tensor * ggml_div_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_div_impl(ctx, a, b, true); +} + +// ggml_sqr + +static struct ggml_tensor * ggml_sqr_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_SQR; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_sqr( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_sqr_impl(ctx, a, false); +} + +struct ggml_tensor * ggml_sqr_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_sqr_impl(ctx, a, true); +} + +// ggml_sqrt + +static struct ggml_tensor * ggml_sqrt_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_SQRT; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_sqrt( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_sqrt_impl(ctx, a, false); +} + +struct ggml_tensor * ggml_sqrt_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_sqrt_impl(ctx, a, true); +} + +// ggml_log + +static struct ggml_tensor * ggml_log_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_LOG; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_log( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_log_impl(ctx, a, false); +} + +struct ggml_tensor * ggml_log_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_log_impl(ctx, a, true); +} + +struct ggml_tensor * ggml_expm1( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_EXPM1); +} + +struct ggml_tensor * ggml_expm1_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_EXPM1); +} + +struct ggml_tensor * ggml_softplus( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_SOFTPLUS); +} + +struct ggml_tensor * ggml_softplus_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SOFTPLUS); +} + +// ggml_sin + +static struct ggml_tensor * ggml_sin_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_SIN; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_sin( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_sin_impl(ctx, a, false); +} + +struct ggml_tensor * ggml_sin_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_sin_impl(ctx, a, true); +} + +// ggml_cos + +static struct ggml_tensor * ggml_cos_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_COS; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_cos( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_cos_impl(ctx, a, false); +} + +struct ggml_tensor * ggml_cos_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_cos_impl(ctx, a, true); +} + +// ggml_sum + +struct ggml_tensor * ggml_sum( + struct ggml_context * ctx, + struct ggml_tensor * a) { + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, a->type, 1); + + result->op = GGML_OP_SUM; + result->src[0] = a; + + return result; +} + +// ggml_sum_rows + +struct ggml_tensor * ggml_sum_rows( + struct ggml_context * ctx, + struct ggml_tensor * a) { + int64_t ne[GGML_MAX_DIMS] = { 1 }; + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + ne[i] = a->ne[i]; + } + + struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne); + + result->op = GGML_OP_SUM_ROWS; + result->src[0] = a; + + return result; +} + +// ggml_cumsum + +struct ggml_tensor * ggml_cumsum( + struct ggml_context * ctx, + struct ggml_tensor * a) { + GGML_ASSERT(a->type == GGML_TYPE_F32); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_CUMSUM; + result->src[0] = a; + + return result; +} + +// ggml_mean + +struct ggml_tensor * ggml_mean( + struct ggml_context * ctx, + struct ggml_tensor * a) { + int64_t ne[4] = { 1, a->ne[1], a->ne[2], a->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_MEAN; + result->src[0] = a; + + return result; +} + +// ggml_argmax + +struct ggml_tensor * ggml_argmax( + struct ggml_context * ctx, + struct ggml_tensor * a) { + GGML_ASSERT(ggml_is_matrix(a)); + GGML_ASSERT(a->ne[0] <= INT32_MAX); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, a->ne[1]); + + result->op = GGML_OP_ARGMAX; + result->src[0] = a; + + return result; +} + +// ggml_count_equal + +struct ggml_tensor * ggml_count_equal( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_are_same_shape(a, b)); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); + + result->op = GGML_OP_COUNT_EQUAL; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_repeat + +struct ggml_tensor * ggml_repeat( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_can_repeat(a, b)); + + struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, b->ne); + + result->op = GGML_OP_REPEAT; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_repeat_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + const bool can_repeat = ggml_is_empty(a) || ( + (ne0 % a->ne[0] == 0) && + (ne1 % a->ne[1] == 0) && + (ne2 % a->ne[2] == 0) && + (ne3 % a->ne[3] == 0) + ); + GGML_ASSERT(can_repeat); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, ne0, ne1, ne2, ne3); + + result->op = GGML_OP_REPEAT; + result->src[0] = a; + + return result; +} + +// ggml_repeat_back + +struct ggml_tensor * ggml_repeat_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_can_repeat(b, a)); + + struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, b->ne); + + result->op = GGML_OP_REPEAT_BACK; + result->src[0] = a; + + return result; +} + +// ggml_concat + +struct ggml_tensor * ggml_concat( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int dim) { + GGML_ASSERT(dim >= 0 && dim < GGML_MAX_DIMS); + GGML_ASSERT(a->type == b->type); + + int64_t ne[GGML_MAX_DIMS]; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (d == dim) { + ne[d] = a->ne[d] + b->ne[d]; + continue; + } + GGML_ASSERT(a->ne[d] == b->ne[d]); + ne[d] = a->ne[d]; + } + + struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne); + + ggml_set_op_params_i32(result, 0, dim); + + result->op = GGML_OP_CONCAT; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_abs + +struct ggml_tensor * ggml_abs( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_ABS); +} + +struct ggml_tensor * ggml_abs_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_ABS); +} + +// ggml_sgn + +struct ggml_tensor * ggml_sgn( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_SGN); +} + +struct ggml_tensor * ggml_sgn_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SGN); +} + +// ggml_neg + +struct ggml_tensor * ggml_neg( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_NEG); +} + +struct ggml_tensor * ggml_neg_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_NEG); +} + +// ggml_step + +struct ggml_tensor * ggml_step( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_STEP); +} + +struct ggml_tensor * ggml_step_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_STEP); +} + +// ggml_tanh + +struct ggml_tensor * ggml_tanh( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_TANH); +} + +struct ggml_tensor * ggml_tanh_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_TANH); +} + +// ggml_elu + +struct ggml_tensor * ggml_elu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_ELU); +} + +struct ggml_tensor * ggml_elu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_ELU); +} + +// ggml_relu + +struct ggml_tensor * ggml_relu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_RELU); +} + +struct ggml_tensor * ggml_relu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_RELU); +} + +// ggml_leaky_relu + +struct ggml_tensor * ggml_leaky_relu( + struct ggml_context * ctx, + struct ggml_tensor * a, + float negative_slope, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params(result, &negative_slope, sizeof(negative_slope)); + + result->op = GGML_OP_LEAKY_RELU; + result->src[0] = a; + + return result; +} + +// ggml_sigmoid + +struct ggml_tensor * ggml_sigmoid( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_SIGMOID); +} + +struct ggml_tensor * ggml_sigmoid_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SIGMOID); +} + +// ggml_gelu + +struct ggml_tensor * ggml_gelu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_GELU); +} + +struct ggml_tensor * ggml_gelu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_GELU); +} + +// ggml_gelu_erf + +struct ggml_tensor * ggml_gelu_erf( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_GELU_ERF); +} + +struct ggml_tensor * ggml_gelu_erf_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_GELU_ERF); +} + +// ggml_gelu_quick + +struct ggml_tensor * ggml_gelu_quick( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_GELU_QUICK); +} + +struct ggml_tensor * ggml_gelu_quick_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_GELU_QUICK); +} + +// ggml_silu + +struct ggml_tensor * ggml_silu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_SILU); +} + +struct ggml_tensor * ggml_silu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SILU); +} + +// ggml_xielu + +struct ggml_tensor * ggml_xielu( + struct ggml_context * ctx, + struct ggml_tensor * a, + float alpha_n, + float alpha_p, + float beta, + float eps) { + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + ggml_set_op_params_i32(result, 0, (int32_t) GGML_UNARY_OP_XIELU); + ggml_set_op_params_f32(result, 1, beta + ggml_compute_softplus_f32(alpha_n)); + ggml_set_op_params_f32(result, 2, ggml_compute_softplus_f32(alpha_p)); + ggml_set_op_params_f32(result, 3, beta); + ggml_set_op_params_f32(result, 4, eps); + + result->op = GGML_OP_UNARY; + result->src[0] = a; + + return result; +} + +// ggml_silu_back + +struct ggml_tensor * ggml_silu_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_SILU_BACK; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml hardswish + +struct ggml_tensor * ggml_hardswish( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_HARDSWISH); +} + +// ggml hardsigmoid + +struct ggml_tensor * ggml_hardsigmoid( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_HARDSIGMOID); +} + +// ggml exp + +struct ggml_tensor * ggml_exp( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_EXP); +} + +struct ggml_tensor * ggml_exp_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_EXP); +} + +// ggml_glu + +static struct ggml_tensor * ggml_glu_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + enum ggml_glu_op op, + bool swapped) { + GGML_ASSERT(ggml_is_contiguous_1(a)); + + if (b) { + GGML_ASSERT(ggml_is_contiguous_1(b)); + GGML_ASSERT(ggml_are_same_shape(a, b)); + GGML_ASSERT(a->type == b->type); + } + + int64_t ne[GGML_MAX_DIMS] = { a->ne[0] / 2 }; for (int i = 1; i < GGML_MAX_DIMS; i++) ne[i] = a->ne[i]; + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b ? a->ne : ne, NULL, 0); + + ggml_set_op_params_i32(result, 0, (int32_t) op); + ggml_set_op_params_i32(result, 1, (int32_t) swapped); + + result->op = GGML_OP_GLU; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_floor + +struct ggml_tensor * ggml_floor( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_FLOOR); +} + +struct ggml_tensor * ggml_floor_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_FLOOR); +} + +// ggml_ceil + +struct ggml_tensor * ggml_ceil( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_CEIL); +} + +struct ggml_tensor * ggml_ceil_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_CEIL); +} + +//ggml_round + +struct ggml_tensor * ggml_round( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_ROUND); +} + +struct ggml_tensor * ggml_round_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_ROUND); +} + +//ggml_trunc + +struct ggml_tensor * ggml_trunc( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_TRUNC); +} + +struct ggml_tensor * ggml_trunc_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_TRUNC); +} + +struct ggml_tensor * ggml_glu( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_glu_op op, + bool swapped) { + return ggml_glu_impl(ctx, a, NULL, op, swapped); +} + +struct ggml_tensor * ggml_glu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + enum ggml_glu_op op) { + return ggml_glu_impl(ctx, a, b, op, false); +} + +// ggml_reglu + +struct ggml_tensor * ggml_reglu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_REGLU, false); +} + +struct ggml_tensor * ggml_reglu_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_REGLU, true); +} + +struct ggml_tensor * ggml_reglu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_REGLU, false); +} + +// ggml_geglu + +struct ggml_tensor * ggml_geglu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU, false); +} + +struct ggml_tensor * ggml_geglu_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU, true); +} + +struct ggml_tensor * ggml_geglu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_GEGLU, false); +} + +// ggml_swiglu + +struct ggml_tensor * ggml_swiglu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_SWIGLU, false); +} + +struct ggml_tensor * ggml_swiglu_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_SWIGLU, true); +} + +struct ggml_tensor * ggml_swiglu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_SWIGLU, false); +} + +// ggml_geglu_erf + +struct ggml_tensor * ggml_geglu_erf( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_ERF, false); +} + +struct ggml_tensor * ggml_geglu_erf_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_ERF, true); +} + +struct ggml_tensor * ggml_geglu_erf_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_GEGLU_ERF, false); +} + +// ggml_geglu_quick + +struct ggml_tensor * ggml_geglu_quick( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_QUICK, false); +} + +struct ggml_tensor * ggml_geglu_quick_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_QUICK, true); +} + +struct ggml_tensor * ggml_geglu_quick_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_GEGLU_QUICK, false); +} + +struct ggml_tensor * ggml_swiglu_oai( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float alpha, + float limit) { + struct ggml_tensor * result = ggml_glu_impl(ctx, a, b, GGML_GLU_OP_SWIGLU_OAI, false); + ggml_set_op_params_f32(result, 2, alpha); + ggml_set_op_params_f32(result, 3, limit); + + return result; +} + +// ggml_norm + +static struct ggml_tensor * ggml_norm_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params(result, &eps, sizeof(eps)); + + result->op = GGML_OP_NORM; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps) { + return ggml_norm_impl(ctx, a, eps, false); +} + +struct ggml_tensor * ggml_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps) { + return ggml_norm_impl(ctx, a, eps, true); +} + +// ggml_rms_norm + +static struct ggml_tensor * ggml_rms_norm_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params(result, &eps, sizeof(eps)); + + result->op = GGML_OP_RMS_NORM; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_rms_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps) { + return ggml_rms_norm_impl(ctx, a, eps, false); +} + +struct ggml_tensor * ggml_rms_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps) { + return ggml_rms_norm_impl(ctx, a, eps, true); +} + +// ggml_rms_norm_back + +struct ggml_tensor * ggml_rms_norm_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float eps) { + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + ggml_set_op_params(result, &eps, sizeof(eps)); + + result->op = GGML_OP_RMS_NORM_BACK; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_group_norm + +static struct ggml_tensor * ggml_group_norm_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_groups, + float eps, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params_i32(result, 0, n_groups); + ggml_set_op_params_f32(result, 1, eps); + + result->op = GGML_OP_GROUP_NORM; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_group_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_groups, + float eps) { + return ggml_group_norm_impl(ctx, a, n_groups, eps, false); +} + +struct ggml_tensor * ggml_group_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_groups, + float eps) { + return ggml_group_norm_impl(ctx, a, n_groups, eps, true); +} + +// ggml_l2_norm + +static struct ggml_tensor * ggml_l2_norm_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params_f32(result, 0, eps); + + result->op = GGML_OP_L2_NORM; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_l2_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps) { + return ggml_l2_norm_impl(ctx, a, eps, false); +} + +struct ggml_tensor * ggml_l2_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps) { + return ggml_l2_norm_impl(ctx, a, eps, true); +} + +// ggml_mul_mat + +static inline bool ggml_can_mul_mat(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return (t0->ne[0] == t1->ne[0]) && + (t1->ne[2]%t0->ne[2] == 0) && // verify t0 is broadcastable + (t1->ne[3]%t0->ne[3] == 0); +} + +struct ggml_tensor * ggml_mul_mat( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_can_mul_mat(a, b)); + GGML_ASSERT(!ggml_is_transposed(a)); + + const int64_t ne[4] = { a->ne[1], b->ne[1], b->ne[2], b->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_MUL_MAT; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +void ggml_mul_mat_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec) { + GGML_ASSERT(a->op == GGML_OP_MUL_MAT); + + const int32_t prec_i32 = (int32_t) prec; + + ggml_set_op_params_i32(a, 0, prec_i32); +} + +// ggml_mul_mat_id + +/* + c = ggml_mul_mat_id(ctx, as, b, ids); + + as -> [cols, rows, n_expert] + b -> [cols, n_expert_used, n_tokens] + ids -> [n_expert_used, n_tokens] (i32) + c -> [rows, n_expert_used, n_tokens] + + in b, n_expert_used can be broadcasted to match the n_expert_used of ids + + c ~= as[:,:,i] @ b[:,i%r,t], i = ids[e,t] for all e,t in ids +*/ +struct ggml_tensor * ggml_mul_mat_id( + struct ggml_context * ctx, + struct ggml_tensor * as, + struct ggml_tensor * b, + struct ggml_tensor * ids) { + GGML_ASSERT(!ggml_is_transposed(as)); + GGML_ASSERT(ids->type == GGML_TYPE_I32); + + GGML_ASSERT(as->ne[3] == 1); // as is 3d (one matrix per expert) + GGML_ASSERT(b->ne[3] == 1); // b is 3d + GGML_ASSERT(ids->ne[2] == 1 && ids->ne[3] == 1); // ids is 2d + GGML_ASSERT(ids->ne[1] == b->ne[2]); // must have an expert list per b row + GGML_ASSERT(as->ne[0] == b->ne[0]); // can_mul_mat + GGML_ASSERT(ids->ne[0] % b->ne[1] == 0); // can broadcast + + const int64_t ne[4] = { as->ne[1], ids->ne[0], b->ne[2], 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_MUL_MAT_ID; + result->src[0] = as; + result->src[1] = b; + result->src[2] = ids; + + return result; +} + +// ggml_out_prod + +static inline bool ggml_can_out_prod(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return (t0->ne[1] == t1->ne[1]) && + (t1->ne[2]%t0->ne[2] == 0) && // verify t0 is broadcastable + (t1->ne[3]%t0->ne[3] == 0); +} + +struct ggml_tensor * ggml_out_prod( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_can_out_prod(a, b)); + GGML_ASSERT(!ggml_is_transposed(a)); + + // a is broadcastable to b for ne[2] and ne[3] -> use b->ne[2] and b->ne[3] + const int64_t ne[4] = { a->ne[0], b->ne[0], b->ne[2], b->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_OUT_PROD; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_scale + +static struct ggml_tensor * ggml_scale_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s, + float b, + bool inplace) { + GGML_ASSERT(ggml_is_padded_1d(a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + float params[2] = { s, b }; + ggml_set_op_params(result, ¶ms, sizeof(params)); + + result->op = GGML_OP_SCALE; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_scale( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s) { + return ggml_scale_impl(ctx, a, s, 0.0, false); +} + +struct ggml_tensor * ggml_scale_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s) { + return ggml_scale_impl(ctx, a, s, 0.0, true); +} + +struct ggml_tensor * ggml_scale_bias( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s, + float b) { + return ggml_scale_impl(ctx, a, s, b, false); +} + +struct ggml_tensor * ggml_scale_bias_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s, + float b) { + return ggml_scale_impl(ctx, a, s, b, true); +} + +// ggml_set + +static struct ggml_tensor * ggml_set_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset, + bool inplace) { + GGML_ASSERT(ggml_nelements(a) >= ggml_nelements(b)); + + // make a view of the destination + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + GGML_ASSERT(offset < (size_t)(1 << 30)); + int32_t params[] = { nb1, nb2, nb3, offset, inplace ? 1 : 0 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_SET; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_set( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset) { + return ggml_set_impl(ctx, a, b, nb1, nb2, nb3, offset, false); +} + +struct ggml_tensor * ggml_set_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset) { + return ggml_set_impl(ctx, a, b, nb1, nb2, nb3, offset, true); +} + +struct ggml_tensor * ggml_set_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t offset) { + return ggml_set_impl(ctx, a, b, a->nb[1], a->nb[2], a->nb[3], offset, false); +} + +struct ggml_tensor * ggml_set_1d_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t offset) { + return ggml_set_impl(ctx, a, b, a->nb[1], a->nb[2], a->nb[3], offset, true); +} + +struct ggml_tensor * ggml_set_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t offset) { + return ggml_set_impl(ctx, a, b, nb1, a->nb[2], a->nb[3], offset, false); +} + +struct ggml_tensor * ggml_set_2d_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t offset) { + return ggml_set_impl(ctx, a, b, nb1, a->nb[2], a->nb[3], offset, true); +} + +// ggml_cpy + +static struct ggml_tensor * ggml_cpy_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b)); + + // make a view of the destination + struct ggml_tensor * result = ggml_view_tensor(ctx, b); + if (strlen(b->name) > 0) { + ggml_format_name(result, "%s (copy of %s)", b->name, a->name); + } else { + ggml_format_name(result, "%s (copy)", a->name); + } + + result->op = GGML_OP_CPY; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_cpy( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_cpy_impl(ctx, a, b); +} + +struct ggml_tensor * ggml_cast( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_type type) { + struct ggml_tensor * result = ggml_new_tensor(ctx, type, GGML_MAX_DIMS, a->ne); + ggml_format_name(result, "%s (copy)", a->name); + + result->op = GGML_OP_CPY; + result->src[0] = a; + result->src[1] = result; // note: this self-reference might seem redundant, but it's actually needed by some + // backends for consistency with ggml_cpy_impl() above + + return result; +} + +// ggml_cont + +static struct ggml_tensor * ggml_cont_impl( + struct ggml_context * ctx, + struct ggml_tensor * a) { + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + ggml_format_name(result, "%s (cont)", a->name); + + result->op = GGML_OP_CONT; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_cont( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_cont_impl(ctx, a); +} + +// make contiguous, with new shape +GGML_API struct ggml_tensor * ggml_cont_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0) { + return ggml_cont_4d(ctx, a, ne0, 1, 1, 1); +} + +GGML_API struct ggml_tensor * ggml_cont_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1) { + return ggml_cont_4d(ctx, a, ne0, ne1, 1, 1); +} + +GGML_API struct ggml_tensor * ggml_cont_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2) { + return ggml_cont_4d(ctx, a, ne0, ne1, ne2, 1); +} + +struct ggml_tensor * ggml_cont_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3) { + GGML_ASSERT(ggml_nelements(a) == (ne0*ne1*ne2*ne3)); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, ne0, ne1, ne2, ne3); + ggml_format_name(result, "%s (cont)", a->name); + + result->op = GGML_OP_CONT; + result->src[0] = a; + + return result; +} + +// ggml_reshape + +struct ggml_tensor * ggml_reshape( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_is_contiguous(a)); + // as only the shape of b is relevant, and not its memory layout, b is allowed to be non contiguous. + GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b)); + + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b->ne, a, 0); + ggml_format_name(result, "%s (reshaped)", a->name); + + result->op = GGML_OP_RESHAPE; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_reshape_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0) { + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_nelements(a) == ne0); + + const int64_t ne[1] = { ne0 }; + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 1, ne, a, 0); + ggml_format_name(result, "%s (reshaped)", a->name); + + result->op = GGML_OP_RESHAPE; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_reshape_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1) { + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_nelements(a) == ne0*ne1); + + const int64_t ne[2] = { ne0, ne1 }; + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 2, ne, a, 0); + ggml_format_name(result, "%s (reshaped)", a->name); + + result->op = GGML_OP_RESHAPE; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_reshape_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2) { + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_nelements(a) == ne0*ne1*ne2); + + const int64_t ne[3] = { ne0, ne1, ne2 }; + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 3, ne, a, 0); + ggml_format_name(result, "%s (reshaped)", a->name); + + result->op = GGML_OP_RESHAPE; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_reshape_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3) { + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_nelements(a) == ne0*ne1*ne2*ne3); + + const int64_t ne[4] = { ne0, ne1, ne2, ne3 }; + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 4, ne, a, 0); + ggml_format_name(result, "%s (reshaped)", a->name); + + result->op = GGML_OP_RESHAPE; + result->src[0] = a; + + return result; +} + +static struct ggml_tensor * ggml_view_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_dims, + const int64_t * ne, + size_t offset) { + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, n_dims, ne, a, offset); + ggml_format_name(result, "%s (view)", a->name); + + ggml_set_op_params(result, &offset, sizeof(offset)); + + result->op = GGML_OP_VIEW; + result->src[0] = a; + + return result; +} + +// ggml_view_1d + +struct ggml_tensor * ggml_view_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + size_t offset) { + struct ggml_tensor * result = ggml_view_impl(ctx, a, 1, &ne0, offset); + + return result; +} + +// ggml_view_2d + +struct ggml_tensor * ggml_view_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + size_t nb1, + size_t offset) { + const int64_t ne[2] = { ne0, ne1 }; + + struct ggml_tensor * result = ggml_view_impl(ctx, a, 2, ne, offset); + + result->nb[1] = nb1; + result->nb[2] = result->nb[1]*ne1; + result->nb[3] = result->nb[2]; + + return result; +} + +// ggml_view_3d + +struct ggml_tensor * ggml_view_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + size_t nb1, + size_t nb2, + size_t offset) { + const int64_t ne[3] = { ne0, ne1, ne2 }; + + struct ggml_tensor * result = ggml_view_impl(ctx, a, 3, ne, offset); + + result->nb[1] = nb1; + result->nb[2] = nb2; + result->nb[3] = result->nb[2]*ne2; + + return result; +} + +// ggml_view_4d + +struct ggml_tensor * ggml_view_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset) { + const int64_t ne[4] = { ne0, ne1, ne2, ne3 }; + + struct ggml_tensor * result = ggml_view_impl(ctx, a, 4, ne, offset); + + result->nb[1] = nb1; + result->nb[2] = nb2; + result->nb[3] = nb3; + + return result; +} + +// ggml_permute + +struct ggml_tensor * ggml_permute( + struct ggml_context * ctx, + struct ggml_tensor * a, + int axis0, + int axis1, + int axis2, + int axis3) { + GGML_ASSERT(axis0 >= 0 && axis0 < GGML_MAX_DIMS); + GGML_ASSERT(axis1 >= 0 && axis1 < GGML_MAX_DIMS); + GGML_ASSERT(axis2 >= 0 && axis2 < GGML_MAX_DIMS); + GGML_ASSERT(axis3 >= 0 && axis3 < GGML_MAX_DIMS); + + GGML_ASSERT(axis0 != axis1); + GGML_ASSERT(axis0 != axis2); + GGML_ASSERT(axis0 != axis3); + GGML_ASSERT(axis1 != axis2); + GGML_ASSERT(axis1 != axis3); + GGML_ASSERT(axis2 != axis3); + + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + ggml_format_name(result, "%s (permuted)", a->name); + + int ne[GGML_MAX_DIMS]; + int nb[GGML_MAX_DIMS]; + + ne[axis0] = a->ne[0]; + ne[axis1] = a->ne[1]; + ne[axis2] = a->ne[2]; + ne[axis3] = a->ne[3]; + + nb[axis0] = a->nb[0]; + nb[axis1] = a->nb[1]; + nb[axis2] = a->nb[2]; + nb[axis3] = a->nb[3]; + + result->ne[0] = ne[0]; + result->ne[1] = ne[1]; + result->ne[2] = ne[2]; + result->ne[3] = ne[3]; + + result->nb[0] = nb[0]; + result->nb[1] = nb[1]; + result->nb[2] = nb[2]; + result->nb[3] = nb[3]; + + result->op = GGML_OP_PERMUTE; + result->src[0] = a; + + int32_t params[] = { axis0, axis1, axis2, axis3 }; + ggml_set_op_params(result, params, sizeof(params)); + + return result; +} + +// ggml_transpose + +struct ggml_tensor * ggml_transpose( + struct ggml_context * ctx, + struct ggml_tensor * a) { + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + ggml_format_name(result, "%s (transposed)", a->name); + + result->ne[0] = a->ne[1]; + result->ne[1] = a->ne[0]; + + result->nb[0] = a->nb[1]; + result->nb[1] = a->nb[0]; + + result->op = GGML_OP_TRANSPOSE; + result->src[0] = a; + + return result; +} + +// ggml_get_rows + +struct ggml_tensor * ggml_get_rows( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(a->ne[2] == b->ne[1]); + GGML_ASSERT(a->ne[3] == b->ne[2]); + GGML_ASSERT(b->ne[3] == 1); + GGML_ASSERT(b->type == GGML_TYPE_I32); + + // TODO: implement non F32 return + enum ggml_type type = GGML_TYPE_F32; + if (a->type == GGML_TYPE_I32) { + type = a->type; + } + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, type, a->ne[0], b->ne[0], b->ne[1], b->ne[2]); + + result->op = GGML_OP_GET_ROWS; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_get_rows_back + +struct ggml_tensor * ggml_get_rows_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c) { + GGML_ASSERT(ggml_is_matrix(a) && ggml_is_vector(b) && b->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_matrix(c) && (a->ne[0] == c->ne[0])); + + // TODO: implement non F32 return + //struct ggml_tensor * result = ggml_new_tensor_2d(ctx, a->type, a->ne[0], b->ne[0]); + struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, c->ne[0], c->ne[1]); + + result->op = GGML_OP_GET_ROWS_BACK; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_set_rows + +struct ggml_tensor * ggml_set_rows( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c) { + GGML_ASSERT(a->ne[0] == b->ne[0]); + GGML_ASSERT(a->ne[2] == b->ne[2]); + GGML_ASSERT(a->ne[3] == b->ne[3]); + GGML_ASSERT(b->ne[1] == c->ne[0]); + GGML_ASSERT(b->ne[2] % c->ne[1] == 0); + GGML_ASSERT(b->ne[3] % c->ne[2] == 0); + GGML_ASSERT(c->ne[3] == 1); + GGML_ASSERT(b->type == GGML_TYPE_F32); + GGML_ASSERT(c->type == GGML_TYPE_I64 || c->type == GGML_TYPE_I32); + + GGML_ASSERT(ggml_is_contiguous_rows(a)); + GGML_ASSERT(ggml_is_contiguous_rows(b)); + + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + + result->op = GGML_OP_SET_ROWS; + result->src[0] = b; + result->src[1] = c; + result->src[2] = a; // note: order is weird due to legacy reasons (https://github.com/ggml-org/llama.cpp/pull/16063#discussion_r2385795931) + + return result; +} + +// ggml_diag + +struct ggml_tensor * ggml_diag( + struct ggml_context * ctx, + struct ggml_tensor * a) { + GGML_ASSERT(a->ne[1] == 1); + + const int64_t ne[4] = { a->ne[0], a->ne[0], a->ne[2], a->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, 4, ne); + + result->op = GGML_OP_DIAG; + result->src[0] = a; + + return result; +} + +// ggml_diag_mask_inf + +static struct ggml_tensor * ggml_diag_mask_inf_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + int32_t params[] = { n_past }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_DIAG_MASK_INF; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_diag_mask_inf( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past) { + return ggml_diag_mask_inf_impl(ctx, a, n_past, false); +} + +struct ggml_tensor * ggml_diag_mask_inf_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past) { + return ggml_diag_mask_inf_impl(ctx, a, n_past, true); +} + +// ggml_diag_mask_zero + +static struct ggml_tensor * ggml_diag_mask_zero_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + int32_t params[] = { n_past }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_DIAG_MASK_ZERO; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_diag_mask_zero( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past) { + return ggml_diag_mask_zero_impl(ctx, a, n_past, false); +} + +struct ggml_tensor * ggml_diag_mask_zero_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past) { + return ggml_diag_mask_zero_impl(ctx, a, n_past, true); +} + +// ggml_soft_max + +static struct ggml_tensor * ggml_soft_max_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * mask, + float scale, + float max_bias, + bool inplace) { + GGML_ASSERT(ggml_is_contiguous(a)); + + if (mask) { + GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(mask)); + GGML_ASSERT(mask->ne[0] == a->ne[0]); + GGML_ASSERT(mask->ne[1] >= a->ne[1]); + GGML_ASSERT(a->ne[2]%mask->ne[2] == 0); + GGML_ASSERT(a->ne[3]%mask->ne[3] == 0); + } + + if (max_bias > 0.0f) { + GGML_ASSERT(mask); + } + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + float params[] = { scale, max_bias }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_SOFT_MAX; + result->src[0] = a; + result->src[1] = mask; + + return result; +} + +struct ggml_tensor * ggml_soft_max( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_soft_max_impl(ctx, a, NULL, 1.0f, 0.0f, false); +} + +struct ggml_tensor * ggml_soft_max_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_soft_max_impl(ctx, a, NULL, 1.0f, 0.0f, true); +} + +struct ggml_tensor * ggml_soft_max_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * mask, + float scale, + float max_bias) { + return ggml_soft_max_impl(ctx, a, mask, scale, max_bias, false); +} + +struct ggml_tensor * ggml_soft_max_ext_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * mask, + float scale, + float max_bias) { + return ggml_soft_max_impl(ctx, a, mask, scale, max_bias, true); +} + +void ggml_soft_max_add_sinks( + struct ggml_tensor * a, + struct ggml_tensor * sinks) { + if (!sinks) { + a->src[2] = NULL; + return; + } + + GGML_ASSERT(a->op == GGML_OP_SOFT_MAX); + GGML_ASSERT(a->src[2] == NULL); + GGML_ASSERT(a->src[0]->ne[2] == sinks->ne[0]); + GGML_ASSERT(sinks->type == GGML_TYPE_F32); + + a->src[2] = sinks; +} + +// ggml_soft_max_ext_back + +static struct ggml_tensor * ggml_soft_max_ext_back_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float scale, + float max_bias, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_SOFT_MAX_BACK; + result->src[0] = a; + result->src[1] = b; + + memcpy((float *) result->op_params + 0, &scale, sizeof(float)); + memcpy((float *) result->op_params + 1, &max_bias, sizeof(float)); + + return result; +} + +struct ggml_tensor * ggml_soft_max_ext_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float scale, + float max_bias) { + return ggml_soft_max_ext_back_impl(ctx, a, b, scale, max_bias, false); +} + +struct ggml_tensor * ggml_soft_max_ext_back_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float scale, + float max_bias) { + return ggml_soft_max_ext_back_impl(ctx, a, b, scale, max_bias, true); +} + +// ggml_rope + +static struct ggml_tensor * ggml_rope_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[GGML_MROPE_SECTIONS], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool inplace) { + GGML_ASSERT((mode & 1) == 0 && "mode & 1 == 1 is no longer supported"); + + GGML_ASSERT(ggml_is_vector(b)); + GGML_ASSERT(b->type == GGML_TYPE_I32); + + bool mrope_used = mode & GGML_ROPE_TYPE_MROPE; + if (mrope_used) { + GGML_ASSERT(a->ne[2] * 4 == b->ne[0]); // mrope expecting 4 position ids per token + } else { + GGML_ASSERT(a->ne[2] == b->ne[0]); + } + + if (c) { + GGML_ASSERT(c->type == GGML_TYPE_F32); + GGML_ASSERT(c->ne[0] >= n_dims / 2); + } + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + int32_t params[15] = { /*n_past*/ 0, n_dims, mode, /*n_ctx*/ 0, n_ctx_orig }; + memcpy(params + 5, &freq_base, sizeof(float)); + memcpy(params + 6, &freq_scale, sizeof(float)); + memcpy(params + 7, &ext_factor, sizeof(float)); + memcpy(params + 8, &attn_factor, sizeof(float)); + memcpy(params + 9, &beta_fast, sizeof(float)); + memcpy(params + 10, &beta_slow, sizeof(float)); + if (mrope_used && sections) { + memcpy(params + 11, sections, sizeof(int32_t) * GGML_MROPE_SECTIONS); + } else { + memset(params + 11, 0, sizeof(int32_t) * GGML_MROPE_SECTIONS); + } + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_ROPE; + result->src[0] = a; + result->src[1] = b; + result->src[2] = c; + + return result; +} + +struct ggml_tensor * ggml_rope( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode) { + return ggml_rope_impl( + ctx, a, b, NULL, n_dims, NULL, mode, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f, false + ); +} + +struct ggml_tensor * ggml_rope_multi( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[GGML_MROPE_SECTIONS], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ggml_rope_impl( + ctx, a, b, c, n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow, false + ); +} + +struct ggml_tensor * ggml_rope_multi_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[GGML_MROPE_SECTIONS], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ggml_rope_impl( + ctx, a, b, c, n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow, true + ); +} + +struct ggml_tensor * ggml_rope_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode) { + return ggml_rope_impl( + ctx, a, b, NULL, n_dims, NULL, mode, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f, true + ); +} + +struct ggml_tensor * ggml_rope_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ggml_rope_impl( + ctx, a, b, c, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow, false + ); +} + +struct ggml_tensor * ggml_rope_ext_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ggml_rope_impl( + ctx, a, b, c, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow, true + ); +} + +struct ggml_tensor * ggml_rope_custom( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ggml_rope_impl( + ctx, a, b, NULL, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow, false + ); +} + +struct ggml_tensor * ggml_rope_custom_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ggml_rope_impl( + ctx, a, b, NULL, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow, true + ); +} + +// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get +// `corr_dim(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))` +static float ggml_rope_yarn_corr_dim(int n_dims, int n_ctx_orig, float n_rot, float base) { + return n_dims * logf(n_ctx_orig / (n_rot * 2 * (float)M_PI)) / (2 * logf(base)); +} + +void ggml_rope_yarn_corr_dims( + int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2] +) { + // start and end correction dims + float start = floorf(ggml_rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_fast, freq_base)); + float end = ceilf(ggml_rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_slow, freq_base)); + dims[0] = MAX(0, start); + dims[1] = MIN(n_dims - 1, end); +} + +// ggml_rope_back + +struct ggml_tensor * ggml_rope_ext_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + struct ggml_tensor * result = ggml_rope_ext( + ctx, a, b, c, n_dims, mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + result->op = GGML_OP_ROPE_BACK; + return result; +} + +struct ggml_tensor * ggml_rope_multi_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[4], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + struct ggml_tensor * result = ggml_rope_multi( + ctx, a, b, c, n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + result->op = GGML_OP_ROPE_BACK; + return result; +} +// ggml_clamp + +struct ggml_tensor * ggml_clamp( + struct ggml_context * ctx, + struct ggml_tensor * a, + float min, + float max) { + // TODO: when implement backward, fix this: + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + + float params[] = { min, max }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_CLAMP; + result->src[0] = a; + + return result; +} + +static int64_t ggml_calc_conv_output_size(int64_t ins, int64_t ks, int s, int p, int d) { + return (ins + 2 * p - d * (ks - 1) - 1) / s + 1; +} + +// im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW] +// a: [OC，IC, KH, KW] +// b: [N, IC, IH, IW] +// result: [N, OH, OW, IC*KH*KW] +struct ggml_tensor * ggml_im2col( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int s1, + int p0, + int p1, + int d0, + int d1, + bool is_2D, + enum ggml_type dst_type) { + if (is_2D) { + GGML_ASSERT(a->ne[2] == b->ne[2]); + } else { + //GGML_ASSERT(b->ne[1] % a->ne[1] == 0); + GGML_ASSERT(b->ne[1] == a->ne[1]); + GGML_ASSERT(b->ne[3] == 1); + } + + const int64_t OH = is_2D ? ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1) : 0; + const int64_t OW = ggml_calc_conv_output_size(b->ne[0], a->ne[0], s0, p0, d0); + + GGML_ASSERT((!is_2D || OH > 0) && "b too small compared to a"); + GGML_ASSERT((OW > 0) && "b too small compared to a"); + + const int64_t ne[4] = { + is_2D ? (a->ne[2] * a->ne[1] * a->ne[0]) : a->ne[1] * a->ne[0], + OW, + is_2D ? OH : b->ne[2], + is_2D ? b->ne[3] : 1, + }; + + struct ggml_tensor * result = ggml_new_tensor(ctx, dst_type, 4, ne); + int32_t params[] = { s0, s1, p0, p1, d0, d1, (is_2D ? 1 : 0) }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_IM2COL; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_im2col_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int64_t * ne, + int s0, + int s1, + int p0, + int p1, + int d0, + int d1, + bool is_2D) { + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + int32_t params[] = { s0, s1, p0, p1, d0, d1, (is_2D ? 1 : 0) }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_IM2COL_BACK; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_conv_1d + +struct ggml_tensor * ggml_conv_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int p0, + int d0) { + struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16); // [N, OL, IC * K] + + struct ggml_tensor * result = + ggml_mul_mat(ctx, + ggml_reshape_2d(ctx, im2col, im2col->ne[0], (im2col->ne[2] * im2col->ne[1])), // [N, OL, IC * K] => [N*OL, IC * K] + ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1]), a->ne[2])); // [OC，IC, K] => [OC, IC * K] + + result = ggml_reshape_3d(ctx, result, im2col->ne[1], a->ne[2], im2col->ne[2]); // [N, OC, OL] + + return result; +} + +// ggml_conv_1d_ph + +struct ggml_tensor* ggml_conv_1d_ph( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s, + int d) { + return ggml_conv_1d(ctx, a, b, s, a->ne[0] / 2, d); +} + +// ggml_conv_1d_dw + +struct ggml_tensor * ggml_conv_1d_dw( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int p0, + int d0) { + struct ggml_tensor * new_b = ggml_reshape_4d(ctx, b, b->ne[0], 1, b->ne[1], b->ne[2]); + + struct ggml_tensor * im2col = ggml_im2col(ctx, a, new_b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16); + + struct ggml_tensor * result = ggml_mul_mat(ctx, im2col, a); + + result = ggml_reshape_3d(ctx, result, result->ne[0], result->ne[2], 1); + + return result; +} + +// ggml_conv_1d_dw_ph + +struct ggml_tensor * ggml_conv_1d_dw_ph( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int d0) { + return ggml_conv_1d_dw(ctx, a, b, s0, a->ne[0] / 2, d0); +} + +// ggml_conv_transpose_1d + +static int64_t ggml_calc_conv_transpose_1d_output_size(int64_t ins, int64_t ks, int s, int p, int d) { + return (ins - 1) * s - 2 * p + d * (ks - 1) + 1; +} + +GGML_API struct ggml_tensor * ggml_conv_transpose_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int p0, + int d0) { + GGML_ASSERT(ggml_is_matrix(b)); + GGML_ASSERT(a->ne[2] == b->ne[1]); + GGML_ASSERT(a->ne[3] == 1); + + GGML_ASSERT(p0 == 0); + GGML_ASSERT(d0 == 1); + + const int64_t ne[4] = { + ggml_calc_conv_transpose_1d_output_size(b->ne[0], a->ne[0], s0, 0 /*p0*/, 1 /*d0*/), + a->ne[1], b->ne[2], 1, + }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + int32_t params[] = { s0, p0, d0 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_CONV_TRANSPOSE_1D; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_conv_2d + +// a: [OC，IC, KH, KW] +// b: [N, IC, IH, IW] +// result: [N, OC, OH, OW] +struct ggml_tensor * ggml_conv_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int s1, + int p0, + int p1, + int d0, + int d1) { + struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, s1, p0, p1, d0, d1, true, a->type); // [N, OH, OW, IC * KH * KW] + + struct ggml_tensor * result = + ggml_mul_mat(ctx, + ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), // [N, OH, OW, IC * KH * KW] => [N*OH*OW, IC * KH * KW] + ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1] * a->ne[2]), a->ne[3])); // [OC，IC, KH, KW] => [OC, IC * KH * KW] + + result = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], im2col->ne[3], a->ne[3]); // [OC, N, OH, OW] + result = ggml_cont(ctx, ggml_permute(ctx, result, 0, 1, 3, 2)); // [N, OC, OH, OW] + + + return result; +} + +// a: [OC*IC, KD, KH, KW] +// b: [N*IC, ID, IH, IW] +// result: [N*OD, OH, OW, IC * KD * KH * KW] +struct ggml_tensor * ggml_im2col_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int64_t IC, + int s0, // stride width + int s1, // stride height + int s2, // stride depth + int p0, // padding width + int p1, // padding height + int p2, // padding depth + int d0, // dilation width + int d1, // dilation height + int d2, // dilation depth + enum ggml_type dst_type) { + const int64_t N = b->ne[3] / IC; + const int64_t ID = b->ne[2]; + const int64_t IH = b->ne[1]; + const int64_t IW = b->ne[0]; + + const int64_t OC = a->ne[3] / IC; + UNUSED(OC); + const int64_t KD = a->ne[2]; + const int64_t KH = a->ne[1]; + const int64_t KW = a->ne[0]; + const int64_t OD = ggml_calc_conv_output_size(ID, KD, s2, p2, d2); + const int64_t OH = ggml_calc_conv_output_size(IH, KH, s1, p1, d1); + const int64_t OW = ggml_calc_conv_output_size(IW, KW, s0, p0, d0); + + GGML_ASSERT((OD > 0) && "b too small compared to a"); + GGML_ASSERT((OH > 0) && "b too small compared to a"); + GGML_ASSERT((OW > 0) && "b too small compared to a"); + + + const int64_t ne[4] = {KW*KH*KD*IC, OW, OH, OD*N}; + + struct ggml_tensor * result = ggml_new_tensor(ctx, dst_type, 4, ne); + int32_t params[] = { s0, s1, s2, p0, p1, p2, d0, d1, d2, (int32_t)IC}; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_IM2COL_3D; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// a: [OC*IC, KD, KH, KW] +// b: [N*IC, ID, IH, IW] +// result: [N*OC, OD, OH, OW] +struct ggml_tensor * ggml_conv_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int64_t IC, + int s0, // stride width + int s1, // stride height + int s2, // stride depth + int p0, // padding width + int p1, // padding height + int p2, // padding depth + int d0, // dilation width + int d1, // dilation height + int d2 // dilation depth + ) { + struct ggml_tensor * im2col = ggml_im2col_3d(ctx, a, b, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, a->type); // [N*OD, OH, OW, IC * KD * KH * KW] + + int64_t OC = a->ne[3] / IC; + int64_t N = b->ne[3] / IC; + struct ggml_tensor * result = + ggml_mul_mat(ctx, + ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), // [N*OD, OH, OW, IC * KD * KH * KW] => [N*OD*OH*OW, IC * KD * KH * KW] + ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1] * a->ne[2] * IC), OC)); // [OC*IC, KD, KH, KW] => [OC, IC * KD * KH * KW] + + int64_t OD = im2col->ne[3] / N; + result = ggml_reshape_4d(ctx, result, im2col->ne[1]*im2col->ne[2], OD, N, OC); // [OC, N*OD*OH*OW] => [OC, N, OD, OH*OW] + result = ggml_cont(ctx, ggml_permute(ctx, result, 0, 1, 3, 2)); // [N, OC, OD, OH*OW] + result = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], OD, OC * N); // [N*OC, OD, OH, OW] + + return result; +} + +// ggml_conv_2d_sk_p0 + +struct ggml_tensor * ggml_conv_2d_sk_p0( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_conv_2d(ctx, a, b, a->ne[0], a->ne[1], 0, 0, 1, 1); +} + +// ggml_conv_2d_s1_ph + +struct ggml_tensor * ggml_conv_2d_s1_ph( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_conv_2d(ctx, a, b, 1, 1, a->ne[0] / 2, a->ne[1] / 2, 1, 1); +} + +// ggml_conv_2d_dw + +struct ggml_tensor * ggml_conv_2d_dw( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int s1, + int p0, + int p1, + int d0, + int d1) { + struct ggml_tensor * new_a = ggml_reshape_4d(ctx, a, a->ne[0], a->ne[1], 1, a->ne[2] * a->ne[3]); + struct ggml_tensor * im2col = ggml_im2col(ctx, new_a, + ggml_reshape_4d(ctx, b, b->ne[0], b->ne[1], 1, b->ne[2] * b->ne[3]), + s0, s1, p0, p1, d0, d1, true, GGML_TYPE_F16); // [N * IC, OH, OW, KH * KW] + struct ggml_tensor * new_b = ggml_reshape_4d(ctx, im2col, im2col->ne[0], im2col->ne[2] * im2col->ne[1], b->ne[2], b->ne[3]); // [N * IC, OH, OW, KH * KW] => [N, IC, OH * OW, KH * KW] + + new_a = ggml_reshape_4d(ctx, new_a, (new_a->ne[0] * new_a->ne[1]), new_a->ne[2], new_a->ne[3], 1); // [OC，1, KH, KW] => [1, OC, 1, KH * KW] + struct ggml_tensor * result = ggml_mul_mat(ctx, new_a, new_b); + result = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], b->ne[2], b->ne[3]); // [N, OC, OH, OW] + + return result; +} + +// ggml_conv_2d_dw_direct + +struct ggml_tensor * ggml_conv_2d_dw_direct( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int stride0, + int stride1, + int pad0, + int pad1, + int dilation0, + int dilation1) { + GGML_ASSERT(a->ne[2] == 1); + GGML_ASSERT(a->ne[3] == b->ne[2]); + int64_t ne[4]; + ne[0] = ggml_calc_conv_output_size(b->ne[0], a->ne[0], stride0, pad0, dilation0); + ne[1] = ggml_calc_conv_output_size(b->ne[1], a->ne[1], stride1, pad1, dilation1); + ne[2] = b->ne[2]; + ne[3] = b->ne[3]; + + struct ggml_tensor * result = ggml_new_tensor(ctx, b->type, 4, ne); + + if (ggml_is_contiguous_channels(b)) { + // Result will be permuted the same way as input (CWHN order) + const int64_t type_size = ggml_type_size(result->type); + GGML_ASSERT(ggml_blck_size(result->type) == 1); + result->nb[0] = result->ne[2] * type_size; + result->nb[1] = result->ne[0] * result->nb[0]; + result->nb[2] = type_size; + } + + int32_t params[] = { stride0, stride1, pad0, pad1, dilation0, dilation1 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_CONV_2D_DW; + result->src[0] = a; + result->src[1] = b; + return result; +} + +// ggml_conv_2d_direct + +struct ggml_tensor * ggml_conv_2d_direct( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel [KW, KH, IC, OC] + struct ggml_tensor * b, // input data [W, H, C, N] + int s0, // stride dimension 0 + int s1, // stride dimension 1 + int p0, // padding dimension 0 + int p1, // padding dimension 1 + int d0, // dilation dimension 0 + int d1) {// dilation dimension 1 + + GGML_ASSERT(a->ne[2] == b->ne[2]); + //GGML_ASSERT(a->type == b->type); + + int64_t ne[4]; + ne[0] = ggml_calc_conv_output_size(b->ne[0], a->ne[0], s0, p0, d0); + ne[1] = ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1); + ne[2] = a->ne[3]; + ne[3] = b->ne[3]; + + struct ggml_tensor * result = ggml_new_tensor(ctx, b->type, 4, ne); + + ggml_set_op_params_i32(result, 0, s0); + ggml_set_op_params_i32(result, 1, s1); + ggml_set_op_params_i32(result, 2, p0); + ggml_set_op_params_i32(result, 3, p1); + ggml_set_op_params_i32(result, 4, d0); + ggml_set_op_params_i32(result, 5, d1); + + result->op = GGML_OP_CONV_2D; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_conv_3d_direct + +struct ggml_tensor * ggml_conv_3d_direct( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int s1, + int s2, + int p0, + int p1, + int p2, + int d0, + int d1, + int d2, + int c, + int n, + int oc) { + + GGML_ASSERT(a->ne[3] == (int64_t) c * oc); + GGML_ASSERT(b->ne[3] == (int64_t) c * n); + + int64_t ne[4]; + ne[0] = ggml_calc_conv_output_size(b->ne[0], a->ne[0], s0, p0, d0); + ne[1] = ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1); + ne[2] = ggml_calc_conv_output_size(b->ne[2], a->ne[2], s2, p2, d2); + ne[3] = (int64_t) oc * n; + + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + ggml_set_op_params_i32(result, 0, s0); + ggml_set_op_params_i32(result, 1, s1); + ggml_set_op_params_i32(result, 2, s2); + ggml_set_op_params_i32(result, 3, p0); + ggml_set_op_params_i32(result, 4, p1); + ggml_set_op_params_i32(result, 5, p2); + ggml_set_op_params_i32(result, 6, d0); + ggml_set_op_params_i32(result, 7, d1); + ggml_set_op_params_i32(result, 8, d2); + ggml_set_op_params_i32(result, 9, c); + ggml_set_op_params_i32(result, 10, n); + ggml_set_op_params_i32(result, 11, oc); + + result->op = GGML_OP_CONV_3D; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_conv_transpose_2d_p0 + +static int64_t ggml_calc_conv_transpose_output_size(int64_t ins, int64_t ks, int s, int p) { + return (ins - 1) * s - 2 * p + ks; +} + +struct ggml_tensor * ggml_conv_transpose_2d_p0( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int stride) { + GGML_ASSERT(a->ne[3] == b->ne[2]); + + const int64_t ne[4] = { + ggml_calc_conv_transpose_output_size(b->ne[0], a->ne[0], stride, 0 /*p0*/), + ggml_calc_conv_transpose_output_size(b->ne[1], a->ne[1], stride, 0 /*p1*/), + a->ne[2], b->ne[3], + }; + + struct ggml_tensor* result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + ggml_set_op_params_i32(result, 0, stride); + + result->op = GGML_OP_CONV_TRANSPOSE_2D; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_pool_* + +static int64_t ggml_calc_pool_output_size(int64_t ins, int ks, int s, float p) { + return (ins + 2 * p - ks) / s + 1; +} + +// ggml_pool_1d + +struct ggml_tensor * ggml_pool_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_op_pool op, + int k0, + int s0, + int p0) { + const int64_t ne[4] = { + ggml_calc_pool_output_size(a->ne[0], k0, s0, p0), + a->ne[1], + a->ne[2], + a->ne[3], + }; + GGML_ASSERT(ne[0] > 0); + + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + int32_t params[] = { op, k0, s0, p0 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_POOL_1D; + result->src[0] = a; + + return result; +} + +// ggml_pool_2d + +struct ggml_tensor * ggml_pool_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_op_pool op, + int k0, + int k1, + int s0, + int s1, + float p0, + float p1) { + struct ggml_tensor * result; + const int64_t ne[4] = { + ggml_calc_pool_output_size(a->ne[0], k0, s0, p0), + ggml_calc_pool_output_size(a->ne[1], k1, s1, p1), + a->ne[2], + a->ne[3], + }; + GGML_ASSERT(ne[0] > 0); + GGML_ASSERT(ne[1] > 0); + + result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + int32_t params[] = { op, k0, k1, s0, s1, p0, p1 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_POOL_2D; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_pool_2d_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * af, + enum ggml_op_pool op, + int k0, + int k1, + int s0, + int s1, + float p0, + float p1) { + struct ggml_tensor * result; + result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, af->ne); + + int32_t params[] = { op, k0, k1, s0, s1, p0, p1 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_POOL_2D_BACK; + result->src[0] = a; + result->src[1] = af; + + return result; +} + +// ggml_upscale / ggml_interpolate + +static struct ggml_tensor * ggml_interpolate_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + uint32_t mode) { + GGML_ASSERT((mode & 0xFF) < GGML_SCALE_MODE_COUNT); + // TODO: implement antialias for modes other than bilinear + GGML_ASSERT(!(mode & GGML_SCALE_FLAG_ANTIALIAS) || (mode & 0xFF) == GGML_SCALE_MODE_BILINEAR); + GGML_ASSERT(a->type == GGML_TYPE_F32); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, ne0, ne1, ne2, ne3); + + ggml_set_op_params_i32(result, 0, (int32_t)mode); + + result->op = GGML_OP_UPSCALE; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_upscale( + struct ggml_context * ctx, + struct ggml_tensor * a, + int scale_factor, + enum ggml_scale_mode mode) { + GGML_ASSERT(scale_factor > 1); + return ggml_interpolate_impl(ctx, a, a->ne[0] * scale_factor, a->ne[1] * scale_factor, a->ne[2], a->ne[3], mode); +} + +struct ggml_tensor * ggml_upscale_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + int ne0, + int ne1, + int ne2, + int ne3, + enum ggml_scale_mode mode) { + return ggml_interpolate_impl(ctx, a, ne0, ne1, ne2, ne3, mode); +} + +struct ggml_tensor * ggml_interpolate( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + uint32_t mode) { + return ggml_interpolate_impl(ctx, a, ne0, ne1, ne2, ne3, mode); +} + +// ggml_pad + +struct ggml_tensor * ggml_pad( + struct ggml_context * ctx, + struct ggml_tensor * a, + int p0, + int p1, + int p2, + int p3) { + return ggml_pad_ext(ctx, a, 0, p0, 0, p1, 0, p2, 0, p3); +} + +// ggml_pad_circular + +struct ggml_tensor * ggml_pad_circular( + struct ggml_context * ctx, + struct ggml_tensor * a, + int p0, + int p1, + int p2, + int p3) { + return ggml_pad_ext_circular(ctx, a, 0, p0, 0, p1, 0, p2, 0, p3); +} + +struct ggml_tensor * ggml_pad_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + int lp0, + int rp0, + int lp1, + int rp1, + int lp2, + int rp2, + int lp3, + int rp3 + ) { + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, + a->ne[0] + lp0 + rp0, + a->ne[1] + lp1 + rp1, + a->ne[2] + lp2 + rp2, + a->ne[3] + lp3 + rp3); + + ggml_set_op_params_i32(result, 0, lp0); + ggml_set_op_params_i32(result, 1, rp0); + ggml_set_op_params_i32(result, 2, lp1); + ggml_set_op_params_i32(result, 3, rp1); + ggml_set_op_params_i32(result, 4, lp2); + ggml_set_op_params_i32(result, 5, rp2); + ggml_set_op_params_i32(result, 6, lp3); + ggml_set_op_params_i32(result, 7, rp3); + ggml_set_op_params_i32(result, 8, 0); // not circular by default + + + result->op = GGML_OP_PAD; + result->src[0] = a; + + return result; +} + +// ggml_pad_ext_circular + +struct ggml_tensor * ggml_pad_ext_circular( + struct ggml_context * ctx, + struct ggml_tensor * a, + int lp0, + int rp0, + int lp1, + int rp1, + int lp2, + int rp2, + int lp3, + int rp3 + ) { + struct ggml_tensor * result = ggml_pad_ext(ctx, a, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); + ggml_set_op_params_i32(result, 8, 1); // circular + return result; +} + +// ggml_pad_reflect_1d + +struct ggml_tensor * ggml_pad_reflect_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int p0, + int p1) { + GGML_ASSERT(p0 >= 0); + GGML_ASSERT(p1 >= 0); + + GGML_ASSERT(p0 < a->ne[0]); // padding length on each size must be less than the + GGML_ASSERT(p1 < a->ne[0]); // existing length of the dimension being padded + + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(a->type == GGML_TYPE_F32); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, + a->ne[0] + p0 + p1, + a->ne[1], + a->ne[2], + a->ne[3]); + + int32_t params[] = { p0, p1 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_PAD_REFLECT_1D; + result->src[0] = a; + + return result; +} + +// ggml_roll + +struct ggml_tensor * ggml_roll( + struct ggml_context * ctx, + struct ggml_tensor * a, + int shift0, + int shift1, + int shift2, + int shift3) { + GGML_ASSERT(a->nb[0] == ggml_type_size(a->type)); + GGML_ASSERT(abs(shift0) < a->ne[0]); + GGML_ASSERT(abs(shift1) < a->ne[1]); + GGML_ASSERT(abs(shift2) < a->ne[2]); + GGML_ASSERT(abs(shift3) < a->ne[3]); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + ggml_set_op_params_i32(result, 0, shift0); + ggml_set_op_params_i32(result, 1, shift1); + ggml_set_op_params_i32(result, 2, shift2); + ggml_set_op_params_i32(result, 3, shift3); + + result->op = GGML_OP_ROLL; + result->src[0] = a; + + return result; +} + +// ggml_timestep_embedding + +struct ggml_tensor * ggml_timestep_embedding( + struct ggml_context * ctx, + struct ggml_tensor * timesteps, + int dim, + int max_period) { + + struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, timesteps->ne[0]); + + ggml_set_op_params_i32(result, 0, dim); + ggml_set_op_params_i32(result, 1, max_period); + + result->op = GGML_OP_TIMESTEP_EMBEDDING; + result->src[0] = timesteps; + + return result; +} + +// ggml_tri + +struct ggml_tensor * ggml_tri( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_tri_type type) { + GGML_ASSERT(a->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(a->ne[0] == a->ne[1]); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + ggml_set_op_params_i32(result, 0, type); + + result->op = GGML_OP_TRI; + result->src[0] = a; + + return result; +} + +// ggml_fill + +static struct ggml_tensor * ggml_fill_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + float c, + bool inplace) { + GGML_ASSERT(a->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params_f32(result, 0, c); + + result->op = GGML_OP_FILL; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_fill( + struct ggml_context * ctx, + struct ggml_tensor * a, + float c) { + return ggml_fill_impl(ctx, a, c, false); +} + +struct ggml_tensor * ggml_fill_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float c) { + return ggml_fill_impl(ctx, a, c, true); +} + +// ggml_argsort + +struct ggml_tensor * ggml_argsort( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_sort_order order) { + GGML_ASSERT(a->ne[0] <= INT32_MAX); + + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_I32, GGML_MAX_DIMS, a->ne); + + ggml_set_op_params_i32(result, 0, (int32_t) order); + + result->op = GGML_OP_ARGSORT; + result->src[0] = a; + + return result; +} + +// ggml_argsort_top_k + +struct ggml_tensor * ggml_argsort_top_k( + struct ggml_context * ctx, + struct ggml_tensor * a, + int k) { + GGML_ASSERT(a->ne[0] >= k); + + struct ggml_tensor * result = ggml_argsort(ctx, a, GGML_SORT_ORDER_DESC); + + result = ggml_view_4d(ctx, result, + k, result->ne[1], result->ne[2], result->ne[3], + result->nb[1], result->nb[2], result->nb[3], + 0); + + return result; +} + +// ggml_top_k + +struct ggml_tensor * ggml_top_k( + struct ggml_context * ctx, + struct ggml_tensor * a, + int k) { + GGML_ASSERT(a->ne[0] >= k); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, k, a->ne[1], a->ne[2], a->ne[3]); + + result->op = GGML_OP_TOP_K; + result->src[0] = a; + + return result; +} + +// ggml_arange + +struct ggml_tensor * ggml_arange( + struct ggml_context * ctx, + float start, + float stop, + float step) { + GGML_ASSERT(stop > start); + + const int64_t steps = (int64_t) ceilf((stop - start) / step); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, steps); + + ggml_set_op_params_f32(result, 0, start); + ggml_set_op_params_f32(result, 1, stop); + ggml_set_op_params_f32(result, 2, step); + + result->op = GGML_OP_ARANGE; + + return result; +} + +// ggml_flash_attn_ext + +struct ggml_tensor * ggml_flash_attn_ext( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * mask, + float scale, + float max_bias, + float logit_softcap) { + GGML_ASSERT(ggml_can_mul_mat(k, q)); + // TODO: check if vT can be multiplied by (k*qT) + + GGML_ASSERT(q->ne[3] == k->ne[3]); + GGML_ASSERT(q->ne[3] == v->ne[3]); + + if (mask) { + GGML_ASSERT(mask->type == GGML_TYPE_F16); + GGML_ASSERT(ggml_is_contiguous(mask)); + //GGML_ASSERT(ggml_can_repeat_rows(mask, qk)); + + GGML_ASSERT(q->ne[2] % mask->ne[2] == 0); + GGML_ASSERT(q->ne[3] % mask->ne[3] == 0); + } + + if (max_bias > 0.0f) { + GGML_ASSERT(mask); + } + + // permute(0, 2, 1, 3) + int64_t ne[4] = { v->ne[0], q->ne[2], q->ne[1], q->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + float params[] = { scale, max_bias, logit_softcap }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_FLASH_ATTN_EXT; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = mask; + + return result; +} + +void ggml_flash_attn_ext_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + + const int32_t prec_i32 = (int32_t) prec; + + ggml_set_op_params_i32(a, 3, prec_i32); // scale is on first pos, max_bias on second +} + +enum ggml_prec ggml_flash_attn_ext_get_prec( + const struct ggml_tensor * a) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + + const int32_t prec_i32 = ggml_get_op_params_i32(a, 3); + + return (enum ggml_prec) prec_i32; +} + +void ggml_flash_attn_ext_add_sinks( + struct ggml_tensor * a, + struct ggml_tensor * sinks) { + if (!sinks) { + a->src[4] = NULL; + return; + } + + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(a->src[4] == NULL); + GGML_ASSERT(a->src[0]->ne[2] == sinks->ne[0]); + GGML_ASSERT(sinks->type == GGML_TYPE_F32); + + a->src[4] = sinks; +} + +// ggml_flash_attn_back + +struct ggml_tensor * ggml_flash_attn_back( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * d, + bool masked) { + GGML_ABORT("TODO: adapt to ggml_flash_attn_ext() changes"); + + GGML_ASSERT(ggml_can_mul_mat(k, q)); + // TODO: check if vT can be multiplied by (k*qT) + + // d shape [D,N,ne2,ne3] + // q shape [D,N,ne2,ne3] + // k shape [D,M,kvne2,ne3] + // v shape [M,D,kvne2,ne3] + + const int64_t D = q->ne[0]; + const int64_t N = q->ne[1]; + const int64_t M = k->ne[1]; + const int64_t ne2 = q->ne[2]; + const int64_t ne3 = q->ne[3]; + const int64_t kvne2 = k->ne[2]; + + GGML_ASSERT(k->ne[0] == D); + GGML_ASSERT(v->ne[0] == M); + GGML_ASSERT(v->ne[1] == D); + GGML_ASSERT(d->ne[0] == D); + GGML_ASSERT(d->ne[1] == N); + GGML_ASSERT(k->ne[2] == kvne2); + GGML_ASSERT(k->ne[3] == ne3); + GGML_ASSERT(v->ne[2] == kvne2); + GGML_ASSERT(v->ne[3] == ne3); + GGML_ASSERT(d->ne[2] == ne2); + GGML_ASSERT(d->ne[3] == ne3); + + GGML_ASSERT(ne2 % kvne2 == 0); + + // store gradients of q, k and v as continuous tensors concatenated in result. + // note: v and gradv are actually transposed, i.e. v->ne[0] != D. + const int64_t elem_q = ggml_nelements(q); + const int64_t elem_k = ggml_nelements(k); + const int64_t elem_v = ggml_nelements(v); + + enum ggml_type result_type = GGML_TYPE_F32; + GGML_ASSERT(ggml_blck_size(result_type) == 1); + const size_t tsize = ggml_type_size(result_type); + + const size_t offs_q = 0; + const size_t offs_k = offs_q + GGML_PAD(elem_q * tsize, GGML_MEM_ALIGN); + const size_t offs_v = offs_k + GGML_PAD(elem_k * tsize, GGML_MEM_ALIGN); + const size_t end = offs_v + GGML_PAD(elem_v * tsize, GGML_MEM_ALIGN); + + const size_t nelements = (end + tsize - 1)/tsize; + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, nelements); + + int32_t masked_i = masked ? 1 : 0; + ggml_set_op_params(result, &masked_i, sizeof(masked_i)); + + result->op = GGML_OP_FLASH_ATTN_BACK; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = d; + + return result; +} + +// ggml_ssm_conv + +struct ggml_tensor * ggml_ssm_conv( + struct ggml_context * ctx, + struct ggml_tensor * sx, + struct ggml_tensor * c) { + GGML_ASSERT(ggml_is_3d(sx)); + GGML_ASSERT(ggml_is_matrix(c)); + + const int64_t d_conv = c->ne[0]; + const int64_t d_inner = c->ne[1]; + const int64_t n_t = sx->ne[0] - d_conv + 1; // tokens per sequence + const int64_t n_s = sx->ne[2]; + + // TODO: maybe support other strides than 1? + GGML_ASSERT(sx->ne[0] == d_conv - 1 + n_t); + GGML_ASSERT(sx->ne[1] == d_inner); + GGML_ASSERT(n_t >= 0); + + struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, d_inner, n_t, n_s); + + result->op = GGML_OP_SSM_CONV; + result->src[0] = sx; + result->src[1] = c; + + return result; +} + +// ggml_ssm_scan + +struct ggml_tensor * ggml_ssm_scan( + struct ggml_context * ctx, + struct ggml_tensor * s, + struct ggml_tensor * x, + struct ggml_tensor * dt, + struct ggml_tensor * A, + struct ggml_tensor * B, + struct ggml_tensor * C, + struct ggml_tensor * ids) { + GGML_ASSERT(ggml_is_contiguous(s)); + GGML_ASSERT(ggml_is_contiguous(dt)); + GGML_ASSERT(ggml_is_contiguous(A)); + GGML_ASSERT(x->nb[0] == ggml_type_size(x->type)); + GGML_ASSERT(B->nb[0] == ggml_type_size(B->type)); + GGML_ASSERT(C->nb[0] == ggml_type_size(C->type)); + GGML_ASSERT(x->nb[1] == x->ne[0]*x->nb[0]); + GGML_ASSERT(B->nb[1] == B->ne[0]*B->nb[0]); + GGML_ASSERT(C->nb[1] == C->ne[0]*C->nb[0]); + GGML_ASSERT(ggml_are_same_shape(B, C)); + GGML_ASSERT(ids->type == GGML_TYPE_I32); + + { + const int64_t d_state = s->ne[0]; + const int64_t head_dim = x->ne[0]; + const int64_t n_head = x->ne[1]; + const int64_t n_seq_tokens = x->ne[2]; + const int64_t n_seqs = x->ne[3]; + + GGML_ASSERT(dt->ne[0] == n_head); + GGML_ASSERT(dt->ne[1] == n_seq_tokens); + GGML_ASSERT(dt->ne[2] == n_seqs); + GGML_ASSERT(ggml_is_3d(dt)); + GGML_ASSERT(s->ne[1] == head_dim); + GGML_ASSERT(s->ne[2] == n_head); + GGML_ASSERT(B->ne[0] == d_state); + GGML_ASSERT(B->ne[2] == n_seq_tokens); + GGML_ASSERT(B->ne[3] == n_seqs); + GGML_ASSERT(ids->ne[0] == n_seqs); + GGML_ASSERT(ggml_is_vector(ids)); + GGML_ASSERT(A->ne[1] == n_head); + GGML_ASSERT(ggml_is_matrix(A)); + + if (A->ne[0] != 1) { + // Mamba-1 has more granular decay factors + GGML_ASSERT(A->ne[0] == d_state); + } + } + + // concatenated y + ssm_states + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]); + + result->op = GGML_OP_SSM_SCAN; + result->src[0] = s; + result->src[1] = x; + result->src[2] = dt; + result->src[3] = A; + result->src[4] = B; + result->src[5] = C; + result->src[6] = ids; + + return result; +} + +// ggml_win_part + +struct ggml_tensor * ggml_win_part( + struct ggml_context * ctx, + struct ggml_tensor * a, + int w) { + GGML_ASSERT(a->ne[3] == 1); + GGML_ASSERT(a->type == GGML_TYPE_F32); + + // padding + const int px = (w - a->ne[1]%w)%w; + const int py = (w - a->ne[2]%w)%w; + + const int npx = (px + a->ne[1])/w; + const int npy = (py + a->ne[2])/w; + const int np = npx*npy; + + const int64_t ne[4] = { a->ne[0], w, w, np, }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + int32_t params[] = { npx, npy, w }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_WIN_PART; + result->src[0] = a; + + return result; +} + +// ggml_win_unpart + +struct ggml_tensor * ggml_win_unpart( + struct ggml_context * ctx, + struct ggml_tensor * a, + int w0, + int h0, + int w) { + GGML_ASSERT(a->type == GGML_TYPE_F32); + + const int64_t ne[4] = { a->ne[0], w0, h0, 1, }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 3, ne); + + int32_t params[] = { w }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_WIN_UNPART; + result->src[0] = a; + + return result; +} + +// ggml_get_rel_pos + +struct ggml_tensor * ggml_get_rel_pos( + struct ggml_context * ctx, + struct ggml_tensor * a, + int qh, + int kh) { + GGML_ASSERT(qh == kh); + GGML_ASSERT(2*MAX(qh, kh) - 1 == a->ne[1]); + + const int64_t ne[4] = { a->ne[0], kh, qh, 1, }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F16, 3, ne); + + result->op = GGML_OP_GET_REL_POS; + result->src[0] = a; + + return result; +} + +// ggml_add_rel_pos + +static struct ggml_tensor * ggml_add_rel_pos_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * pw, + struct ggml_tensor * ph, + bool inplace) { + GGML_ASSERT(ggml_are_same_shape(pw, ph)); + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_is_contiguous(pw)); + GGML_ASSERT(ggml_is_contiguous(ph)); + GGML_ASSERT(ph->type == GGML_TYPE_F32); + GGML_ASSERT(pw->type == GGML_TYPE_F32); + GGML_ASSERT(pw->ne[3] == a->ne[2]); + GGML_ASSERT(pw->ne[0]*pw->ne[0] == a->ne[0]); + GGML_ASSERT(pw->ne[1]*pw->ne[2] == a->ne[1]); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + ggml_set_op_params_i32(result, 0, inplace ? 1 : 0); + + result->op = GGML_OP_ADD_REL_POS; + result->src[0] = a; + result->src[1] = pw; + result->src[2] = ph; + + return result; +} + +struct ggml_tensor * ggml_add_rel_pos( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * pw, + struct ggml_tensor * ph) { + return ggml_add_rel_pos_impl(ctx, a, pw, ph, false); +} + +struct ggml_tensor * ggml_add_rel_pos_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * pw, + struct ggml_tensor * ph) { + return ggml_add_rel_pos_impl(ctx, a, pw, ph, true); +} + +// ggml_rwkv_wkv6 + +struct ggml_tensor * ggml_rwkv_wkv6( + struct ggml_context * ctx, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * r, + struct ggml_tensor * tf, + struct ggml_tensor * td, + struct ggml_tensor * state) { + GGML_ASSERT(ggml_is_contiguous(k)); + GGML_ASSERT(ggml_is_contiguous(v)); + GGML_ASSERT(ggml_is_contiguous(r)); + GGML_ASSERT(ggml_is_contiguous(tf)); + GGML_ASSERT(ggml_is_contiguous(td)); + GGML_ASSERT(ggml_is_contiguous(state)); + + const int64_t S = k->ne[0]; + const int64_t H = k->ne[1]; + const int64_t n_tokens = k->ne[2]; + const int64_t n_seqs = state->ne[1]; + { + GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); + GGML_ASSERT(r->ne[0] == S && r->ne[1] == H && r->ne[2] == n_tokens); + GGML_ASSERT(td->ne[0] == S && td->ne[1] == H && td->ne[2] == n_tokens); + GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); + } + + // concat output and new_state + const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_RWKV_WKV6; + result->src[0] = k; + result->src[1] = v; + result->src[2] = r; + result->src[3] = tf; + result->src[4] = td; + result->src[5] = state; + + return result; +} + +// ggml_gated_linear_attn + +struct ggml_tensor * ggml_gated_linear_attn( + struct ggml_context * ctx, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * q, + struct ggml_tensor * g, + struct ggml_tensor * state, + float scale) { + GGML_ASSERT(ggml_is_contiguous(k)); + GGML_ASSERT(ggml_is_contiguous(v)); + GGML_ASSERT(ggml_is_contiguous(q)); + GGML_ASSERT(ggml_is_contiguous(g)); + GGML_ASSERT(ggml_is_contiguous(state)); + + const int64_t S = k->ne[0]; + const int64_t H = k->ne[1]; + const int64_t n_tokens = k->ne[2]; + const int64_t n_seqs = state->ne[1]; + { + GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); + GGML_ASSERT(q->ne[0] == S && q->ne[1] == H && q->ne[2] == n_tokens); + GGML_ASSERT(g->ne[0] == S && g->ne[1] == H && g->ne[2] == n_tokens); + GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); + } + + // concat output and new_state + const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + ggml_set_op_params_f32(result, 0, scale); + + result->op = GGML_OP_GATED_LINEAR_ATTN; + result->src[0] = k; + result->src[1] = v; + result->src[2] = q; + result->src[3] = g; + result->src[4] = state; + + return result; +} + +// ggml_rwkv_wkv7 + +struct ggml_tensor * ggml_rwkv_wkv7( + struct ggml_context * ctx, + struct ggml_tensor * r, + struct ggml_tensor * w, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * state) { + GGML_ASSERT(ggml_is_contiguous(r)); + GGML_ASSERT(ggml_is_contiguous(w)); + GGML_ASSERT(ggml_is_contiguous(k)); + GGML_ASSERT(ggml_is_contiguous(v)); + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_is_contiguous(b)); + GGML_ASSERT(ggml_is_contiguous(state)); + + const int64_t S = k->ne[0]; + const int64_t H = k->ne[1]; + const int64_t n_tokens = k->ne[2]; + const int64_t n_seqs = state->ne[1]; + { + GGML_ASSERT(w->ne[0] == S && w->ne[1] == H && w->ne[2] == n_tokens); + GGML_ASSERT(k->ne[0] == S && k->ne[1] == H && k->ne[2] == n_tokens); + GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); + GGML_ASSERT(a->ne[0] == S && a->ne[1] == H && a->ne[2] == n_tokens); + GGML_ASSERT(b->ne[0] == S && b->ne[1] == H && b->ne[2] == n_tokens); + GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); + } + + // concat output and new_state + const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_RWKV_WKV7; + result->src[0] = r; + result->src[1] = w; + result->src[2] = k; + result->src[3] = v; + result->src[4] = a; + result->src[5] = b; + result->src[6] = state; + + return result; +} + +// ggml_unary + +static struct ggml_tensor * ggml_unary_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_unary_op op, + bool inplace) { + GGML_ASSERT(ggml_is_contiguous_rows(a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params_i32(result, 0, (int32_t) op); + + result->op = GGML_OP_UNARY; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_unary( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_unary_op op) { + return ggml_unary_impl(ctx, a, op, false); +} + +struct ggml_tensor * ggml_unary_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_unary_op op) { + return ggml_unary_impl(ctx, a, op, true); +} + +// ggml_map_custom1 + +static struct ggml_tensor * ggml_map_custom1_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + const ggml_custom1_op_t fun, + int n_tasks, + void * userdata, + bool inplace) { + GGML_ASSERT(n_tasks == GGML_N_TASKS_MAX || n_tasks > 0); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + struct ggml_map_custom1_op_params params = { + /*.fun =*/ fun, + /*.n_tasks =*/ n_tasks, + /*.userdata =*/ userdata + }; + ggml_set_op_params(result, ¶ms, sizeof(params)); + + result->op = GGML_OP_MAP_CUSTOM1; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_map_custom1( + struct ggml_context * ctx, + struct ggml_tensor * a, + const ggml_custom1_op_t fun, + int n_tasks, + void * userdata) { + return ggml_map_custom1_impl(ctx, a, fun, n_tasks, userdata, false); +} + +struct ggml_tensor * ggml_map_custom1_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + const ggml_custom1_op_t fun, + int n_tasks, + void * userdata) { + return ggml_map_custom1_impl(ctx, a, fun, n_tasks, userdata, true); +} + +// ggml_map_custom2 + +static struct ggml_tensor * ggml_map_custom2_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + const ggml_custom2_op_t fun, + int n_tasks, + void * userdata, + bool inplace) { + GGML_ASSERT(n_tasks == GGML_N_TASKS_MAX || n_tasks > 0); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + struct ggml_map_custom2_op_params params = { + /*.fun =*/ fun, + /*.n_tasks =*/ n_tasks, + /*.userdata =*/ userdata + }; + ggml_set_op_params(result, ¶ms, sizeof(params)); + + result->op = GGML_OP_MAP_CUSTOM2; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_map_custom2( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + const ggml_custom2_op_t fun, + int n_tasks, + void * userdata) { + return ggml_map_custom2_impl(ctx, a, b, fun, n_tasks, userdata, false); +} + +struct ggml_tensor * ggml_map_custom2_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + const ggml_custom2_op_t fun, + int n_tasks, + void * userdata) { + return ggml_map_custom2_impl(ctx, a, b, fun, n_tasks, userdata, true); +} + +// ggml_map_custom3 + +static struct ggml_tensor * ggml_map_custom3_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + const ggml_custom3_op_t fun, + int n_tasks, + void * userdata, + bool inplace) { + GGML_ASSERT(n_tasks == GGML_N_TASKS_MAX || n_tasks > 0); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + struct ggml_map_custom3_op_params params = { + /*.fun =*/ fun, + /*.n_tasks =*/ n_tasks, + /*.userdata =*/ userdata + }; + ggml_set_op_params(result, ¶ms, sizeof(params)); + + result->op = GGML_OP_MAP_CUSTOM3; + result->src[0] = a; + result->src[1] = b; + result->src[2] = c; + + return result; +} + +struct ggml_tensor * ggml_map_custom3( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + const ggml_custom3_op_t fun, + int n_tasks, + void * userdata) { + return ggml_map_custom3_impl(ctx, a, b, c, fun, n_tasks, userdata, false); +} + +struct ggml_tensor * ggml_map_custom3_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + const ggml_custom3_op_t fun, + int n_tasks, + void * userdata) { + return ggml_map_custom3_impl(ctx, a, b, c, fun, n_tasks, userdata, true); +} + +struct ggml_tensor * ggml_custom_4d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + struct ggml_tensor ** args, + int n_args, + ggml_custom_op_t fun, + int n_tasks, + void * userdata) { + + GGML_ASSERT(n_args < GGML_MAX_SRC); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, type, ne0, ne1, ne2, ne3); + + struct ggml_custom_op_params params = { + /*.fun =*/ fun, + /*.n_tasks =*/ n_tasks, + /*.userdata =*/ userdata + }; + ggml_set_op_params(result, ¶ms, sizeof(params)); + + result->op = GGML_OP_CUSTOM; + for (int i = 0; i < n_args; i++) { + result->src[i] = args[i]; + } + + return result; +} + +struct ggml_tensor * ggml_custom_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor ** args, + int n_args, + ggml_custom_op_t fun, + int n_tasks, + void * userdata) { + + GGML_ASSERT(n_args < GGML_MAX_SRC - 1); + + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + + struct ggml_custom_op_params params = { + /*.fun =*/ fun, + /*.n_tasks =*/ n_tasks, + /*.userdata =*/ userdata + }; + ggml_set_op_params(result, ¶ms, sizeof(params)); + + result->op = GGML_OP_CUSTOM; + result->src[0] = a; + for (int i = 0; i < n_args; i++) { + result->src[i + 1] = args[i]; + } + + return result; +} +// ggml_cross_entropy_loss + +struct ggml_tensor * ggml_cross_entropy_loss( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_are_same_shape(a, b)); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, a->type, 1); + + result->op = GGML_OP_CROSS_ENTROPY_LOSS; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_cross_entropy_loss_back + +struct ggml_tensor * ggml_cross_entropy_loss_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c) { + GGML_ASSERT(ggml_is_scalar(a)); + GGML_ASSERT(ggml_are_same_shape(b, c)); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, b); + + result->op = GGML_OP_CROSS_ENTROPY_LOSS_BACK; + result->src[0] = a; + result->src[1] = b; + result->src[2] = c; + + return result; +} + +// opt_step_adamw + +struct ggml_tensor * ggml_opt_step_adamw( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * grad, + struct ggml_tensor * m, + struct ggml_tensor * v, + struct ggml_tensor * adamw_params) { + GGML_ASSERT(a->flags & GGML_TENSOR_FLAG_PARAM); + GGML_ASSERT(ggml_are_same_shape(a, grad)); + GGML_ASSERT(ggml_are_same_shape(a, m)); + GGML_ASSERT(ggml_are_same_shape(a, v)); + GGML_ASSERT(adamw_params->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_nelements(adamw_params) == 7); + + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + + result->op = GGML_OP_OPT_STEP_ADAMW; + result->src[0] = a; + result->src[1] = grad; + result->src[2] = m; + result->src[3] = v; + result->src[4] = adamw_params; + + return result; +} + +// opt_step_sgd + +struct ggml_tensor * ggml_opt_step_sgd( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * grad, + struct ggml_tensor * params) { + GGML_ASSERT(a->flags & GGML_TENSOR_FLAG_PARAM); + GGML_ASSERT(ggml_are_same_shape(a, grad)); + GGML_ASSERT(params->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_nelements(params) == 2); + + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + + result->op = GGML_OP_OPT_STEP_SGD; + result->src[0] = a; + result->src[1] = grad; + result->src[2] = params; + + return result; +} + +// solve_tri + +struct ggml_tensor * ggml_solve_tri( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool left, + bool lower, + bool uni) { + GGML_ASSERT(a->type == GGML_TYPE_F32); + GGML_ASSERT(b->type == GGML_TYPE_F32); + + // A must be square and lower diagonal + GGML_ASSERT(a->ne[0] == a->ne[1]); + // B must have same outer dimension as A + GGML_ASSERT(a->ne[1] == b->ne[1]); + + // batch dimensions must be equal + GGML_ASSERT(a->ne[2] == b->ne[2]); + GGML_ASSERT(a->ne[3] == b->ne[3]); + + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_is_contiguous(b)); + + GGML_ASSERT(lower && left && !uni); // TODO: support other variants + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, b->ne[0], b->ne[1], b->ne[2], b->ne[3]); + + result->op = GGML_OP_SOLVE_TRI; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_gated_delta_net + +struct ggml_tensor * ggml_gated_delta_net( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state) { + GGML_ASSERT(ggml_is_contiguous_rows(q)); + GGML_ASSERT(ggml_is_contiguous_rows(k)); + GGML_ASSERT(ggml_is_contiguous_rows(v)); + GGML_ASSERT(ggml_is_contiguous(g)); + GGML_ASSERT(ggml_is_contiguous(beta)); + GGML_ASSERT(ggml_is_contiguous(state)); + + GGML_ASSERT(q->type == GGML_TYPE_F32); + GGML_ASSERT(k->type == GGML_TYPE_F32); + GGML_ASSERT(v->type == GGML_TYPE_F32); + GGML_ASSERT(g->type == GGML_TYPE_F32); + GGML_ASSERT(beta->type == GGML_TYPE_F32); + GGML_ASSERT(state->type == GGML_TYPE_F32); + + const int64_t S_v = v->ne[0]; + const int64_t H = v->ne[1]; + const int64_t n_tokens = v->ne[2]; + const int64_t n_seqs = v->ne[3]; + + // gate: scalar [1, H, T, B] or vector [S_v, H, T, B] (KDA) + GGML_ASSERT(g->ne[0] == 1 || g->ne[0] == S_v); + GGML_ASSERT(beta->ne[0] == 1); + + GGML_ASSERT(ggml_nelements(state) == S_v * S_v * H * n_seqs); + + // concat output and new_state into a single tensor + // output: S_v * H * n_tokens * n_seqs, state: S_v * S_v * H * n_seqs + const int64_t ne[4] = { S_v * H, n_tokens * n_seqs + S_v * n_seqs, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_GATED_DELTA_NET; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = g; + result->src[4] = beta; + result->src[5] = state; + + return result; +} + +//////////////////////////////////////////////////////////////////////////////// + +struct ggml_hash_set ggml_hash_set_new(size_t size) { + size = ggml_hash_size(size); + struct ggml_hash_set result; + result.size = size; + result.keys = GGML_MALLOC(sizeof(struct ggml_tensor *) * size); + result.used = GGML_CALLOC(ggml_bitset_size(size), sizeof(ggml_bitset_t)); + return result; +} + +void ggml_hash_set_reset(struct ggml_hash_set * hash_set) { + memset(hash_set->used, 0, sizeof(ggml_bitset_t) * ggml_bitset_size(hash_set->size)); +} + +void ggml_hash_set_free(struct ggml_hash_set * hash_set) { + GGML_FREE(hash_set->used); + GGML_FREE(hash_set->keys); +} + +size_t ggml_hash_size(size_t min_sz) { + // next primes after powers of two + static const size_t primes[] = { + 2, 3, 5, 11, 17, 37, 67, 131, 257, 521, 1031, + 2053, 4099, 8209, 16411, 32771, 65537, 131101, + 262147, 524309, 1048583, 2097169, 4194319, 8388617, + 16777259, 33554467, 67108879, 134217757, 268435459, + 536870923, 1073741827, 2147483659 + }; + static const size_t n_primes = sizeof(primes)/sizeof(primes[0]); + + // find the smallest prime that is larger or equal than min_sz + size_t l = 0; + size_t r = n_primes; + while (l < r) { + size_t m = (l + r)/2; + if (primes[m] < min_sz) { + l = m + 1; + } else { + r = m; + } + } + size_t sz = l < n_primes ? primes[l] : min_sz | 1; + return sz; +} + +struct hash_map { + struct ggml_hash_set set; + struct ggml_tensor ** vals; +}; + +static struct hash_map * ggml_new_hash_map(size_t size) { + struct hash_map * result = GGML_MALLOC(sizeof(struct hash_map)); + result->set = ggml_hash_set_new(size); + result->vals = GGML_CALLOC(result->set.size, sizeof(struct ggml_tensor *)); + return result; +} + +static void ggml_hash_map_free(struct hash_map * map) { + ggml_hash_set_free(&map->set); + GGML_FREE(map->vals); + GGML_FREE(map); +} + +// utility functions to change gradients +// isrc is the index of tensor in cgraph->visited_has_set.keys +// the corresponding gradient (accumulators) are also at position isrc +// if tensor has a gradient accumulator, modify that accumulator in-place +// else if there is no gradient for tensor, set the corresponding value +// else, just add/subtract/etc. the gradients + +static void ggml_add_or_set( + struct ggml_context * ctx, + struct ggml_cgraph * cgraph, + size_t isrc, + struct ggml_tensor * tensor) { + struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; + GGML_ASSERT(src); + if (cgraph->grads[isrc]) { + cgraph->grads[isrc] = ggml_add_impl(ctx, cgraph->grads[isrc], tensor, /*inplace =*/ cgraph->grad_accs[isrc]); + } else { + cgraph->grads[isrc] = tensor; + } + ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name); + ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); +} + +static void ggml_acc_or_set( + struct ggml_context * ctx, + struct ggml_cgraph * cgraph, + size_t isrc, + struct ggml_tensor * tensor, + const size_t nb1, + const size_t nb2, + const size_t nb3, + const size_t offset) { + struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; + GGML_ASSERT(src); + if (cgraph->grads[isrc]) { + cgraph->grads[isrc] = ggml_acc_impl(ctx, cgraph->grads[isrc], tensor, nb1, nb2, nb3, offset, cgraph->grad_accs[isrc]); + } else { + struct ggml_tensor * a_zero = ggml_scale(ctx, src, 0.0f); // FIXME this is going to produce NaN if a contains inf/NaN + cgraph->grads[isrc] = ggml_acc_impl(ctx, a_zero, tensor, nb1, nb2, nb3, offset, false); + } + ggml_format_name(cgraph->grads[isrc], "grad for %s", cgraph->visited_hash_set.keys[isrc]->name); + ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); +} + +static void ggml_add1_or_set( + struct ggml_context * ctx, + struct ggml_cgraph * cgraph, + size_t isrc, + struct ggml_tensor * tensor) { + struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; + GGML_ASSERT(src); + if (cgraph->grads[isrc]) { + cgraph->grads[isrc] = ggml_add1_impl(ctx, cgraph->grads[isrc], tensor, cgraph->grad_accs[isrc]); + } else { + cgraph->grads[isrc] = ggml_repeat(ctx, tensor, src); + } + ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name); + ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); +} + +static void ggml_sub_or_set( + struct ggml_context * ctx, + struct ggml_cgraph * cgraph, + size_t isrc, + struct ggml_tensor * tensor) { + struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; + GGML_ASSERT(src); + if (cgraph->grads[isrc]) { + cgraph->grads[isrc] = ggml_sub_impl(ctx, cgraph->grads[isrc], tensor, cgraph->grad_accs[isrc]); + } else { + cgraph->grads[isrc] = ggml_neg(ctx, tensor); + } + ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name); + ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); +} + +static void ggml_compute_backward( + struct ggml_context * ctx, struct ggml_cgraph * cgraph, int i, const bool * grads_needed) { + struct ggml_tensor * tensor = cgraph->nodes[i]; + struct ggml_tensor * grad = ggml_graph_get_grad(cgraph, tensor); + + if (!grad) { + return; + } + + struct ggml_tensor * src0 = tensor->src[0]; + struct ggml_tensor * src1 = tensor->src[1]; + struct ggml_tensor * src2 = tensor->src[2]; + struct ggml_hash_set * hash_set = &cgraph->visited_hash_set; + const size_t isrc0 = src0 ? ggml_hash_find(hash_set, src0) : (size_t) -1; + const size_t isrc1 = src1 ? ggml_hash_find(hash_set, src1) : (size_t) -1; + const size_t isrc2 = src2 ? ggml_hash_find(hash_set, src2) : (size_t) -1; + const bool src0_needs_grads = src0 && isrc0 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc0) && grads_needed[isrc0]; + const bool src1_needs_grads = src1 && isrc1 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc1) && grads_needed[isrc1]; + const bool src2_needs_grads = src2 && isrc2 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc2) && grads_needed[isrc2]; + + switch (tensor->op) { + case GGML_OP_DUP: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, grad); + } + } break; + case GGML_OP_ADD: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, grad); + } + if (src1_needs_grads) { + struct ggml_tensor * tmp = grad; + if (!ggml_are_same_shape(src0, src1)) { + tmp = ggml_repeat_back(ctx, tmp, src1); + } + ggml_add_or_set(ctx, cgraph, isrc1, tmp); + } + } break; + case GGML_OP_ADD1: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, grad); + } + if (src1_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc1, ggml_mean(ctx, grad)); // TODO: should probably be sum instead of mean + } + } break; + case GGML_OP_ACC: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, grad); + } + if (src1_needs_grads) { + const size_t nb1 = ((int32_t *) tensor->op_params)[0]; + const size_t nb2 = ((int32_t *) tensor->op_params)[1]; + const size_t nb3 = ((int32_t *) tensor->op_params)[2]; + const size_t offset = ((int32_t *) tensor->op_params)[3]; + + struct ggml_tensor * tensor_grad_view = ggml_view_4d(ctx, + grad, src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], + nb1, nb2, nb3, offset); + + ggml_add_or_set(ctx, cgraph, isrc1, ggml_reshape(ctx, ggml_cont(ctx, tensor_grad_view), src1)); + } + } break; + case GGML_OP_SUB: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, grad); + } + if (src1_needs_grads) { + ggml_sub_or_set(ctx, cgraph, isrc1, grad); + } + } break; + case GGML_OP_MUL: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, src1)); + } + if (src1_needs_grads) { + struct ggml_tensor * tmp = ggml_mul(ctx, src0, grad); + if (!ggml_are_same_shape(src0, src1)) { + tmp = ggml_repeat_back(ctx, tmp, src1); + } + ggml_add_or_set(ctx, cgraph, isrc1, tmp); + } + } break; + case GGML_OP_DIV: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_div(ctx, grad, src1)); + } + if (src1_needs_grads) { + ggml_sub_or_set(ctx, cgraph, isrc1, ggml_mul(ctx, grad, ggml_div(ctx, tensor, src1))); + } + } break; + case GGML_OP_SQR: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_scale(ctx, ggml_mul(ctx, src0, grad), 2.0f)); + } + } break; + case GGML_OP_SQRT: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_scale(ctx, ggml_div(ctx, grad, tensor), 0.5f)); + } + } break; + case GGML_OP_LOG: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_div(ctx, grad, src0)); + } + } break; + case GGML_OP_SIN: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_cos(ctx, src0))); + } + } break; + case GGML_OP_COS: { + if (src0_needs_grads) { + ggml_sub_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_sin(ctx, src0))); + } + } break; + case GGML_OP_SUM: { + if (src0_needs_grads) { + ggml_add1_or_set(ctx, cgraph, isrc0, grad); + } + } break; + case GGML_OP_SUM_ROWS: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat(ctx, grad, src0)); + } + } break; + case GGML_OP_MEAN: { + if (src0_needs_grads) { + ggml_add1_or_set(ctx, cgraph, isrc0, ggml_scale_impl(ctx, grad, 1.0f/src0->ne[0], 0.0, false)); + } + } break; + case GGML_OP_REPEAT: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat_back(ctx, grad, src0)); + } + } break; + case GGML_OP_REPEAT_BACK: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat(ctx, grad, src0)); + } + } break; + case GGML_OP_RMS_NORM: { + if (src0_needs_grads) { + float eps; + memcpy(&eps, tensor->op_params, sizeof(float)); + ggml_add_or_set(ctx, cgraph, isrc0, ggml_rms_norm_back(ctx, grad, src0, eps)); + } + } break; + case GGML_OP_MUL_MAT: { + // https://cs231n.github.io/optimization-2/#staged + // # forward pass + // s0 = np.random.randn(5, 10) + // s1 = np.random.randn(10, 3) + // t = s0.dot(s1) + + // # now suppose we had the gradient on t from above in the circuit + // dt = np.random.randn(*t.shape) # same shape as t + // ds0 = dt.dot(s1.T) #.T gives the transpose of the matrix + // ds1 = t.T.dot(dt) + + // tensor.shape [m,p,qq,rr] + // src0.shape [n,m,q1,r1] + // src1.shape [n,p,qq,rr] + + if (src0_needs_grads) { + GGML_ASSERT(grad->ne[2] == src1->ne[2]); + GGML_ASSERT(grad->ne[3] == src1->ne[3]); + struct ggml_tensor * tmp = + ggml_out_prod(ctx, // [n,m,qq,rr] + src1, // [n,p,qq,rr] + grad); // [m,p,qq,rr] + if (!ggml_are_same_shape(tmp, src0)) { + GGML_ASSERT(tmp->ne[0] == src0->ne[0]); + GGML_ASSERT(tmp->ne[1] == src0->ne[1]); + GGML_ASSERT(tmp->ne[3] == 1); + + const int64_t nr2 = tmp->ne[2] / src0->ne[2]; + const size_t nb2 = tmp->nb[2] * nr2; + const size_t nb3 = tmp->nb[2]; + + tmp = ggml_view_4d(ctx, tmp, src0->ne[0], src0->ne[1], src0->ne[2], nr2, tmp->nb[1], nb2, nb3, 0); + tmp = ggml_repeat_back(ctx, tmp, src0); + } + ggml_add_or_set(ctx, cgraph, isrc0, tmp); + } + if (src1_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc1, + // ggml_mul_mat(ctx, // [n,p,qq,rr] + // ggml_cont(ctx, // [m,n,q1,r1] + // ggml_transpose(ctx, src0)), // [m,n,q1,r1] + // grad), // [m,p,qq,rr] + + // when src0 is bigger than tensor->grad (this is mostly the case in llama), + // avoid transpose of src0, rather transpose smaller tensor->grad + // and then use ggml_out_prod + ggml_out_prod(ctx, // [n,p,qq,rr] + src0, // [n,m,q1,r1] + ggml_transpose(ctx, // [p,m,qq,rr] + grad))); // [m,p,qq,rr] + } + } break; + case GGML_OP_SCALE: { + if (src0_needs_grads) { + float s; + memcpy(&s, tensor->op_params, sizeof(float)); + ggml_add_or_set(ctx, cgraph, isrc0, ggml_scale_impl(ctx, grad, s, 0.0, false)); + } + } break; + case GGML_OP_SET: { + const size_t nb1 = ((const int32_t *) tensor->op_params)[0]; + const size_t nb2 = ((const int32_t *) tensor->op_params)[1]; + const size_t nb3 = ((const int32_t *) tensor->op_params)[2]; + const size_t offset = ((const int32_t *) tensor->op_params)[3]; + + struct ggml_tensor * tensor_grad_view = NULL; + + if (src0_needs_grads || src1_needs_grads) { + GGML_ASSERT(src0->type == tensor->type); + GGML_ASSERT(!cgraph->grads[isrc0] || cgraph->grads[isrc0]->type == grad->type); + GGML_ASSERT(!cgraph->grads[isrc1] || !src1_needs_grads || cgraph->grads[isrc1]->type == grad->type); + + tensor_grad_view = ggml_view_4d(ctx, + grad, src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], + nb1, nb2, nb3, offset); + } + + if (src0_needs_grads) { + struct ggml_tensor * tmp = ggml_neg(ctx, tensor_grad_view); + ggml_add_or_set(ctx, cgraph, isrc0, ggml_acc_impl(ctx, grad, tmp, nb1, nb2, nb3, offset, false)); + } + + if (src1_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc1, ggml_reshape(ctx, ggml_cont(ctx, tensor_grad_view), src1)); + } + } break; + case GGML_OP_CPY: { + // cpy overwrites value of src1 by src0 and returns view(src1) + // the overwriting is mathematically equivalent to: + // tensor = src0 * 1 + src1 * 0 + if (src0_needs_grads) { + // dsrc0 = dtensor * 1 + ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape(ctx, grad, src0)); + } + if (src1_needs_grads) { + // dsrc1 = dtensor * 0 -> noop + } + } break; + case GGML_OP_CONT: { + // same as cpy + if (src0_needs_grads) { + GGML_ASSERT(!cgraph->grads[isrc0] || ggml_is_contiguous(cgraph->grads[isrc0])); + GGML_ASSERT(ggml_is_contiguous(grad)); + GGML_ASSERT(ggml_nelements(tensor) == ggml_nelements(src0)); + ggml_add_or_set(ctx, cgraph, isrc0, + ggml_are_same_shape(tensor, src0) ? grad : ggml_reshape(ctx, grad, src0)); + } + } break; + case GGML_OP_RESHAPE: { + if (src0_needs_grads) { + struct ggml_tensor * grad_cont = ggml_is_contiguous(grad) ? grad : ggml_cont(ctx, grad); + ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape(ctx, grad_cont, src0)); + } + } break; + case GGML_OP_VIEW: { + if (src0_needs_grads) { + size_t offset; + + memcpy(&offset, tensor->op_params, sizeof(offset)); + + size_t nb1 = tensor->nb[1]; + size_t nb2 = tensor->nb[2]; + size_t nb3 = tensor->nb[3]; + + if (cgraph->grads[isrc0] && src0->type != cgraph->grads[isrc0]->type) { + // gradient is typically F32, but src0 could be other type + size_t ng = ggml_element_size(cgraph->grads[isrc0]); + size_t n0 = ggml_element_size(src0); + GGML_ASSERT(offset % n0 == 0); + GGML_ASSERT(nb1 % n0 == 0); + GGML_ASSERT(nb2 % n0 == 0); + GGML_ASSERT(nb3 % n0 == 0); + offset = (offset / n0) * ng; + nb1 = (nb1 / n0) * ng; + nb2 = (nb2 / n0) * ng; + nb3 = (nb3 / n0) * ng; + } + + ggml_acc_or_set(ctx, cgraph, isrc0, grad, nb1, nb2, nb3, offset); + } + } break; + case GGML_OP_PERMUTE: { + if (src0_needs_grads) { + const int32_t * axes = (const int32_t *) tensor->op_params; + const int axis0 = axes[0] & 0x3; + const int axis1 = axes[1] & 0x3; + const int axis2 = axes[2] & 0x3; + const int axis3 = axes[3] & 0x3; + int axb[4] = {0,0,0,0}; // axes backward + axb[axis0] = 0; + axb[axis1] = 1; + axb[axis2] = 2; + axb[axis3] = 3; + ggml_add_or_set(ctx, cgraph, isrc0, ggml_permute(ctx, grad, axb[0], axb[1], axb[2], axb[3])); + } + } break; + case GGML_OP_TRANSPOSE: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_transpose(ctx, grad)); + } + } break; + case GGML_OP_GET_ROWS: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_get_rows_back(ctx, grad, src1, src0)); + } + if (src1_needs_grads) { + // noop + } + } break; + case GGML_OP_DIAG_MASK_INF: { + if (src0_needs_grads) { + /* ggml_diag_mask_inf_impl() shouldn't be here */ + /* ref: https://github.com/ggml-org/llama.cpp/pull/4203#discussion_r1412377992 */ + const int n_past = ((const int32_t *) tensor->op_params)[0]; + ggml_add_or_set(ctx, cgraph, isrc0, ggml_diag_mask_zero_impl(ctx, grad, n_past, false)); + } + } break; + case GGML_OP_DIAG_MASK_ZERO: { + if (src0_needs_grads) { + const int n_past = ((const int32_t *) tensor->op_params)[0]; + ggml_add_or_set(ctx, cgraph, isrc0, ggml_diag_mask_zero_impl(ctx, grad, n_past, false)); + } + } break; + case GGML_OP_SOFT_MAX: { + if (src0_needs_grads) { + float scale = 1.0f; + float max_bias = 0.0f; + + memcpy(&scale, (const float *) tensor->op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) tensor->op_params + 1, sizeof(float)); + + ggml_add_or_set(ctx, cgraph, isrc0, ggml_soft_max_ext_back(ctx, grad, tensor, scale, max_bias)); + } + GGML_ASSERT((!src1 || !src1_needs_grads) && "backward pass for softmax mask not implemented"); + } break; + case GGML_OP_ROPE: { + if (src0_needs_grads) { + //const int n_past = ((int32_t *) tensor->op_params)[0]; + const int n_dims = ((const int32_t *) tensor->op_params)[1]; + const int mode = ((const int32_t *) tensor->op_params)[2]; + //const int n_ctx = ((int32_t *) tensor->op_params)[3]; + const int n_ctx_orig = ((const int32_t *) tensor->op_params)[4]; + float freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow; + int sections[4] = {0, 0, 0, 0}; + + memcpy(&freq_base, (const float *) tensor->op_params + 5, sizeof(float)); + memcpy(&freq_scale, (const float *) tensor->op_params + 6, sizeof(float)); + memcpy(&ext_factor, (const float *) tensor->op_params + 7, sizeof(float)); + memcpy(&attn_factor, (const float *) tensor->op_params + 8, sizeof(float)); + memcpy(&beta_fast, (const float *) tensor->op_params + 9, sizeof(float)); + memcpy(&beta_slow, (const float *) tensor->op_params + 10, sizeof(float)); + memcpy(§ions, tensor->op_params + 11, sizeof(sections)); + + struct ggml_tensor * rope_back = grad->ne[2] == src1->ne[0] ? + ggml_rope_ext_back(ctx, grad, src1, src2, n_dims, + mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow) : + ggml_rope_multi_back(ctx, grad, src1, src2, n_dims, sections, + mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + ggml_add_or_set(ctx, cgraph, isrc0, rope_back); + } + GGML_ASSERT((!src2 || !src2_needs_grads) && "gradients for freq factors not implemented"); + } break; + case GGML_OP_IM2COL: { + if (src1_needs_grads) { + const int32_t s0 = ggml_get_op_params_i32(tensor, 0); + const int32_t s1 = ggml_get_op_params_i32(tensor, 1); + const int32_t p0 = ggml_get_op_params_i32(tensor, 2); + const int32_t p1 = ggml_get_op_params_i32(tensor, 3); + const int32_t d0 = ggml_get_op_params_i32(tensor, 4); + const int32_t d1 = ggml_get_op_params_i32(tensor, 5); + const bool is_2D = ggml_get_op_params_i32(tensor, 6) == 1; + + ggml_add_or_set(ctx, cgraph, isrc1, ggml_im2col_back(ctx, grad, src0, src1->ne, s0, s1, p0, p1, d0, d1, is_2D)); + } + } break; + case GGML_OP_POOL_2D: { + if (src0_needs_grads) { + const enum ggml_op_pool op = ggml_get_op_params_i32(tensor, 0); + const int32_t k0 = ggml_get_op_params_i32(tensor, 1); + const int32_t k1 = ggml_get_op_params_i32(tensor, 2); + const int32_t s0 = ggml_get_op_params_i32(tensor, 3); + const int32_t s1 = ggml_get_op_params_i32(tensor, 4); + const int32_t p0 = ggml_get_op_params_i32(tensor, 5); + const int32_t p1 = ggml_get_op_params_i32(tensor, 6); + + ggml_add_or_set(ctx, cgraph, isrc0, ggml_pool_2d_back(ctx, grad, src0, op, k0, k1, s0, s1, p0, p1)); + } + } break; + case GGML_OP_WIN_PART: + case GGML_OP_WIN_UNPART: + case GGML_OP_UNARY: { + switch (ggml_get_unary_op(tensor)) { + case GGML_UNARY_OP_ABS: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, ggml_sgn(ctx, src0), grad)); + } + } break; + case GGML_UNARY_OP_SGN: { + // noop + } break; + case GGML_UNARY_OP_NEG: { + if (src0_needs_grads) { + ggml_sub_or_set(ctx, cgraph, isrc0, grad); + } + } break; + case GGML_UNARY_OP_STEP: { + // noop + } break; + case GGML_UNARY_OP_RELU: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, ggml_step(ctx, src0), grad)); + } + } break; + case GGML_UNARY_OP_SILU: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_silu_back(ctx, grad, src0)); + } + } break; + case GGML_UNARY_OP_EXP: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, tensor, grad)); + } + } break; + case GGML_UNARY_OP_EXPM1: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_exp(ctx, src0))); + } + } break; + case GGML_UNARY_OP_SOFTPLUS: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_sigmoid(ctx, src0))); + } + } break; + default: { + fprintf(stderr, "%s: unsupported unary op for backward pass: %s\n", + __func__, ggml_unary_op_name(ggml_get_unary_op(tensor))); + GGML_ABORT("fatal error"); + } //break; + } + } break; + case GGML_OP_CROSS_ENTROPY_LOSS: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_cross_entropy_loss_back(ctx, grad, src0, src1)); + } + GGML_ASSERT(!src1_needs_grads && "backward pass for labels not implemented"); + } break; + case GGML_OP_GLU: { + switch (ggml_get_glu_op(tensor)) { + case GGML_GLU_OP_SWIGLU: { + if (src0_needs_grads) { + GGML_ASSERT(src1 && "backward pass only implemented for split swiglu"); + ggml_add_or_set(ctx, cgraph, isrc0, ggml_silu_back(ctx, ggml_mul(ctx, grad, src1), src0)); + } + if (src1_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc1, ggml_mul(ctx, ggml_silu(ctx, src0), grad)); + } + } break; + default: { + GGML_ABORT("unsupported glu op for backward pass: %s", ggml_glu_op_name(ggml_get_glu_op(tensor))); + } //break; + } + } break; + case GGML_OP_NONE: { + // noop + } break; + case GGML_OP_COUNT: + default: { + GGML_ABORT("%s: unsupported ggml op for backward pass: %s\n", __func__, ggml_op_name(tensor->op)); + } //break; + } + + GGML_ASSERT(!src0_needs_grads || ggml_are_same_shape(src0, cgraph->grads[isrc0])); + GGML_ASSERT(!src1_needs_grads || ggml_are_same_shape(src1, cgraph->grads[isrc1])); + GGML_ASSERT(!src2_needs_grads || ggml_are_same_shape(src2, cgraph->grads[isrc2])); +} + +static size_t ggml_visit_parents_graph(struct ggml_cgraph * cgraph, struct ggml_tensor * node, bool compute) { + if (node->op != GGML_OP_NONE && compute) { + node->flags |= GGML_TENSOR_FLAG_COMPUTE; + } + + const size_t node_hash_pos = ggml_hash_find(&cgraph->visited_hash_set, node); + GGML_ASSERT(node_hash_pos != GGML_HASHSET_FULL); + + if (ggml_bitset_get(cgraph->visited_hash_set.used, node_hash_pos)) { + // already visited + + if (compute) { + // update the compute flag regardless + for (int i = 0; i < GGML_MAX_SRC; ++i) { + struct ggml_tensor * src = node->src[i]; + if (src && ((src->flags & GGML_TENSOR_FLAG_COMPUTE) == 0)) { + ggml_visit_parents_graph(cgraph, src, true); + } + } + } + + return node_hash_pos; + } + + // This is the first time we see this node in the current graph. + cgraph->visited_hash_set.keys[node_hash_pos] = node; + ggml_bitset_set(cgraph->visited_hash_set.used, node_hash_pos); + cgraph->use_counts[node_hash_pos] = 0; + + for (int i = 0; i < GGML_MAX_SRC; ++i) { + const int k = + (cgraph->order == GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT) ? i : + (cgraph->order == GGML_CGRAPH_EVAL_ORDER_RIGHT_TO_LEFT) ? (GGML_MAX_SRC-1-i) : + /* unknown order, just fall back to using i */ i; + + struct ggml_tensor * src = node->src[k]; + if (src) { + const size_t src_hash_pos = ggml_visit_parents_graph(cgraph, src, compute); + + // Update the use count for this operand. + cgraph->use_counts[src_hash_pos]++; + } + } + + if (node->op == GGML_OP_NONE && !(node->flags & GGML_TENSOR_FLAG_PARAM)) { + // reached a leaf node, not part of the gradient graph (e.g. a constant) + GGML_ASSERT(cgraph->n_leafs < cgraph->size); + + if (strlen(node->name) == 0) { + ggml_format_name(node, "leaf_%d", cgraph->n_leafs); + } + + cgraph->leafs[cgraph->n_leafs] = node; + cgraph->n_leafs++; + } else { + GGML_ASSERT(cgraph->n_nodes < cgraph->size); + + if (strlen(node->name) == 0) { + ggml_format_name(node, "node_%d", cgraph->n_nodes); + } + + cgraph->nodes[cgraph->n_nodes] = node; + cgraph->n_nodes++; + } + + return node_hash_pos; +} + +static void ggml_build_forward_impl(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor, bool expand, bool compute) { + if (!expand) { + // TODO: this branch isn't accessible anymore, maybe move this to ggml_build_forward_expand + ggml_graph_clear(cgraph); + } + + const int n_old = cgraph->n_nodes; + + ggml_visit_parents_graph(cgraph, tensor, compute); + + const int n_new = cgraph->n_nodes - n_old; + GGML_PRINT_DEBUG("%s: visited %d new nodes\n", __func__, n_new); + + if (n_new > 0) { + // the last added node should always be starting point + GGML_ASSERT(cgraph->nodes[cgraph->n_nodes - 1] == tensor); + } +} + +struct ggml_tensor * ggml_build_forward_select( + struct ggml_cgraph * cgraph, + struct ggml_tensor ** tensors, + int n_tensors, + int idx) { + GGML_ASSERT(idx >= 0 && idx < n_tensors); + + for (int i = 0; i < n_tensors; i++) { + ggml_build_forward_impl(cgraph, tensors[i], true, i == idx ? true : false); + } + + return tensors[idx]; +} + +void ggml_build_forward_expand(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor) { + ggml_build_forward_impl(cgraph, tensor, true, true); +} + +void ggml_build_backward_expand( + struct ggml_context * ctx, + struct ggml_cgraph * cgraph, + struct ggml_tensor ** grad_accs) { + GGML_ASSERT(cgraph->n_nodes > 0); + GGML_ASSERT(cgraph->grads); + GGML_ASSERT(cgraph->grad_accs); + + const int n_nodes_f = cgraph->n_nodes; + + memset(cgraph->grads, 0, cgraph->visited_hash_set.size*sizeof(struct ggml_tensor *)); + memset(cgraph->grad_accs, 0, cgraph->visited_hash_set.size*sizeof(struct ggml_tensor *)); + bool * grads_needed = calloc(cgraph->visited_hash_set.size, sizeof(bool)); + + { + bool any_params = false; + bool any_loss = false; + for (int i = 0; i < n_nodes_f; ++i) { + struct ggml_tensor * node = cgraph->nodes[i]; + any_params = any_params || (node->flags & GGML_TENSOR_FLAG_PARAM); + any_loss = any_loss || (node->flags & GGML_TENSOR_FLAG_LOSS); + } + GGML_ASSERT(any_params && "no trainable parameters found, did you forget to call ggml_set_param?"); + GGML_ASSERT(any_loss && "no training loss found, did you forget to call ggml_set_loss?"); + } + + for (int i = 0; i < n_nodes_f; ++i) { + struct ggml_tensor * node = cgraph->nodes[i]; + + if (node->type == GGML_TYPE_I32) { + continue; + } + + bool node_needs_grad = (node->flags & GGML_TENSOR_FLAG_PARAM) || (node->flags & GGML_TENSOR_FLAG_LOSS); + bool ignore_src[GGML_MAX_SRC] = {false}; + switch (node->op) { + // gradients in node->src[0] for one reason or another have no effect on output gradients + case GGML_OP_IM2COL: // only used for its shape + case GGML_OP_IM2COL_BACK: // same as IM2COL + ignore_src[0] = true; + break; + case GGML_OP_UNARY: { + const enum ggml_unary_op uop = ggml_get_unary_op(node); + // SGN and STEP unary ops are piecewise constant + if (uop == GGML_UNARY_OP_SGN || uop == GGML_UNARY_OP_STEP) { + ignore_src[0] = true; + } + } break; + + // gradients in node->src[1] for one reason or another have no effect on output gradients + case GGML_OP_CPY: // gradients in CPY target are irrelevant + case GGML_OP_GET_ROWS: // row indices not differentiable + case GGML_OP_GET_ROWS_BACK: // same as for GET_ROWS + case GGML_OP_ROPE: // positions not differentiable + ignore_src[1] = true; + break; + + default: + break; + } + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (!node->src[j] || ignore_src[j] || !grads_needed[ggml_hash_find(&cgraph->visited_hash_set, node->src[j])]) { + continue; + } + GGML_ASSERT(node->src[j]->type == GGML_TYPE_F32 || node->src[j]->type == GGML_TYPE_F16); + node_needs_grad = true; + break; + } + if (!node_needs_grad) { + continue; + } + + // inplace operations are currently not supported + GGML_ASSERT(!node->view_src || node->op == GGML_OP_CPY || node->op == GGML_OP_VIEW || + node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE); + + const size_t ihash = ggml_hash_find(&cgraph->visited_hash_set, node); + GGML_ASSERT(ihash != GGML_HASHSET_FULL); + GGML_ASSERT(ggml_bitset_get(cgraph->visited_hash_set.used, ihash)); + if (grad_accs && grad_accs[i]) { + cgraph->grad_accs[ihash] = grad_accs[i]; + cgraph->grads[ihash] = cgraph->grad_accs[ihash]; + } else if (node->flags & GGML_TENSOR_FLAG_LOSS) { + // loss tensors always need a gradient accumulator + cgraph->grad_accs[ihash] = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, node->ne); + cgraph->grads[ihash] = cgraph->grad_accs[ihash]; + } + grads_needed[ihash] = true; + } + + for (int i = n_nodes_f - 1; i >= 0; --i) { + // inplace operations to add gradients are not created by ggml_compute_backward except for gradient accumulation + // use allocator to automatically make inplace operations + ggml_compute_backward(ctx, cgraph, i, grads_needed); + } + + free(grads_needed); +} + +static void * incr_ptr_aligned(void ** p, size_t size, size_t align) { + void * ptr = *p; + ptr = (void *) GGML_PAD((uintptr_t) ptr, align); + *p = (void *) ((char *) ptr + size); + return ptr; +} + +static size_t ggml_graph_nbytes(size_t size, bool grads) { + size_t hash_size = ggml_hash_size(size * 2); + void * p = 0; + incr_ptr_aligned(&p, sizeof(struct ggml_cgraph), 1); + incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // nodes + incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // leafs + incr_ptr_aligned(&p, hash_size * sizeof(int32_t), sizeof(int32_t)); // use_counts + incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // hash keys + if (grads) { + incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // grads + incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // grad_accs + } + incr_ptr_aligned(&p, ggml_bitset_size(hash_size) * sizeof(ggml_bitset_t), sizeof(ggml_bitset_t)); + + size_t nbytes = (size_t) p; + return nbytes; +} + +size_t ggml_graph_overhead_custom(size_t size, bool grads) { + return GGML_OBJECT_SIZE + GGML_PAD(ggml_graph_nbytes(size, grads), GGML_MEM_ALIGN); +} + +size_t ggml_graph_overhead(void) { + return ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, false); +} + +struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t size, bool grads) { + const size_t obj_size = ggml_graph_nbytes(size, grads); + struct ggml_object * obj = ggml_new_object(ctx, GGML_OBJECT_TYPE_GRAPH, obj_size); + struct ggml_cgraph * cgraph = (struct ggml_cgraph *) ((char *) ctx->mem_buffer + obj->offs); + + // the size of the hash table is doubled since it needs to hold both nodes and leafs + size_t hash_size = ggml_hash_size(size * 2); + + void * p = cgraph + 1; + + struct ggml_tensor ** nodes_ptr = incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); + struct ggml_tensor ** leafs_ptr = incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); + int32_t * use_counts_ptr = incr_ptr_aligned(&p, hash_size * sizeof(int32_t), sizeof(int32_t)); + struct ggml_tensor ** hash_keys_ptr = incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); + struct ggml_tensor ** grads_ptr = grads ? incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)) : NULL; + struct ggml_tensor ** grad_accs_ptr = grads ? incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)) : NULL; + + ggml_bitset_t * hash_used = incr_ptr_aligned(&p, ggml_bitset_size(hash_size) * sizeof(ggml_bitset_t), sizeof(ggml_bitset_t)); + + // check that we allocated the correct amount of memory + assert(obj_size == (size_t)((char *)p - (char *)cgraph)); + + *cgraph = (struct ggml_cgraph) { + /*.size =*/ size, + /*.n_nodes =*/ 0, + /*.n_leafs =*/ 0, + /*.nodes =*/ nodes_ptr, + /*.grads =*/ grads_ptr, + /*.grad_accs =*/ grad_accs_ptr, + /*.leafs =*/ leafs_ptr, + /*.use_counts =*/ use_counts_ptr, + /*.hash_table =*/ { hash_size, hash_used, hash_keys_ptr }, + /*.order =*/ GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT, + /*.uid =*/ 0, + }; + + ggml_hash_set_reset(&cgraph->visited_hash_set); + if (grads) { + memset(cgraph->grads, 0, hash_size*sizeof(struct ggml_tensor *)); + memset(cgraph->grad_accs, 0, hash_size*sizeof(struct ggml_tensor *)); + } + + return cgraph; +} + +struct ggml_cgraph * ggml_new_graph(struct ggml_context * ctx) { + return ggml_new_graph_custom(ctx, GGML_DEFAULT_GRAPH_SIZE, false); +} + +struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph0, int i0, int i1) { + struct ggml_cgraph cgraph = { + /*.size =*/ 0, + /*.n_nodes =*/ i1 - i0, + /*.n_leafs =*/ 0, + /*.nodes =*/ cgraph0->nodes + i0, + /*.grads =*/ NULL, // gradients would need visited_hash_set + /*.grad_accs =*/ NULL, + /*.leafs =*/ NULL, + /*.use_counts =*/ cgraph0->use_counts, + /*.visited_hash_set =*/ cgraph0->visited_hash_set, + /*.order =*/ cgraph0->order, + /*.uid =*/ 0 + }; + + return cgraph; +} + +void ggml_graph_cpy(struct ggml_cgraph * src, struct ggml_cgraph * dst) { + GGML_ASSERT(dst->size >= src->n_leafs); + GGML_ASSERT(dst->size >= src->n_nodes); + GGML_ASSERT(dst->visited_hash_set.size >= src->visited_hash_set.size); + + dst->n_leafs = src->n_leafs; + dst->n_nodes = src->n_nodes; + dst->order = src->order; + + for (int i = 0; i < src->n_leafs; ++i) { + dst->leafs[i] = src->leafs[i]; + } + + for (int i = 0; i < src->n_nodes; ++i) { + dst->nodes[i] = src->nodes[i]; + } + + for (size_t i = 0; i < src->visited_hash_set.size; ++i) { + // copy all hashset keys (tensors) that are in use + if (ggml_bitset_get(src->visited_hash_set.used, i)) { + size_t new_hash_pos = ggml_hash_insert(&dst->visited_hash_set, src->visited_hash_set.keys[i]); + dst->use_counts[new_hash_pos] = src->use_counts[i]; + } + } + + if (dst->grads) { + memset(dst->grads, 0, dst->visited_hash_set.size*sizeof(struct ggml_tensor *)); + memset(dst->grad_accs, 0, dst->visited_hash_set.size*sizeof(struct ggml_tensor *)); + } + if (src->grads) { + GGML_ASSERT(dst->grads != NULL); + GGML_ASSERT(dst->grad_accs != NULL); + for (int i = 0; i < src->n_nodes; ++i) { + const size_t igrad_src = ggml_hash_find(&src->visited_hash_set, src->nodes[i]); + const size_t igrad_dst = ggml_hash_find(&dst->visited_hash_set, dst->nodes[i]); + + GGML_ASSERT(igrad_src != GGML_HASHSET_FULL); + GGML_ASSERT(ggml_bitset_get(src->visited_hash_set.used, igrad_src)); + GGML_ASSERT(igrad_dst != GGML_HASHSET_FULL); + GGML_ASSERT(ggml_bitset_get(dst->visited_hash_set.used, igrad_dst)); + + dst->grads[igrad_dst] = src->grads[igrad_src]; + dst->grad_accs[igrad_dst] = src->grad_accs[igrad_src]; + } + } +} + +struct ggml_cgraph * ggml_graph_dup(struct ggml_context * ctx, struct ggml_cgraph * cgraph, bool force_grads) { + struct ggml_cgraph * result = ggml_new_graph_custom(ctx, cgraph->size, cgraph->grads || force_grads); + ggml_graph_cpy(cgraph, result); + return result; +} + +struct ggml_tensor * ggml_set_zero(struct ggml_tensor * tensor) { + if (ggml_is_empty(tensor)) { + return tensor; + } + if (tensor->buffer) { + ggml_backend_tensor_memset(tensor, 0, 0, ggml_nbytes(tensor)); + } else { + GGML_ASSERT(tensor->data); + memset(tensor->data, 0, ggml_nbytes(tensor)); + } + return tensor; +} + +void ggml_graph_reset(struct ggml_cgraph * cgraph) { + if (!cgraph) { + return; + } + GGML_ASSERT(cgraph->grads != NULL); + + for (int i = 0; i < cgraph->n_nodes; i++) { + struct ggml_tensor * node = cgraph->nodes[i]; + struct ggml_tensor * grad_acc = ggml_graph_get_grad_acc(cgraph, node); + + if (node->op == GGML_OP_OPT_STEP_ADAMW) { + // clear momenta + ggml_set_zero(node->src[2]); + ggml_set_zero(node->src[3]); + } + + // initial gradients of loss should be 1, 0 otherwise + if (grad_acc) { + if (node->flags & GGML_TENSOR_FLAG_LOSS) { + GGML_ASSERT(grad_acc->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_scalar(grad_acc)); + + const float onef = 1.0f; + if (grad_acc->buffer) { + ggml_backend_tensor_set(grad_acc, &onef, 0, sizeof(float)); + } else { + GGML_ASSERT(grad_acc->data); + *((float *) grad_acc->data) = onef; + } + } else { + ggml_set_zero(grad_acc); + } + } + } +} + +void ggml_graph_clear(struct ggml_cgraph * cgraph) { + cgraph->n_leafs = 0; + cgraph->n_nodes = 0; + ggml_hash_set_reset(&cgraph->visited_hash_set); +} + +int ggml_graph_size(struct ggml_cgraph * cgraph) { + return cgraph->size; +} + +struct ggml_tensor * ggml_graph_node(struct ggml_cgraph * cgraph, int i) { + if (i < 0) { + GGML_ASSERT(cgraph->n_nodes + i >= 0); + return cgraph->nodes[cgraph->n_nodes + i]; + } + + GGML_ASSERT(i < cgraph->n_nodes); + return cgraph->nodes[i]; +} + +struct ggml_tensor ** ggml_graph_nodes(struct ggml_cgraph * cgraph) { + return cgraph->nodes; +} + +int ggml_graph_n_nodes(struct ggml_cgraph * cgraph) { + return cgraph->n_nodes; +} + +void ggml_graph_add_node(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor) { + GGML_ASSERT(cgraph->size > cgraph->n_nodes); + cgraph->nodes[cgraph->n_nodes] = tensor; + cgraph->n_nodes++; +} + +struct ggml_tensor * ggml_graph_get_tensor(const struct ggml_cgraph * cgraph, const char * name) { + for (int i = 0; i < cgraph->n_leafs; i++) { + struct ggml_tensor * leaf = cgraph->leafs[i]; + + if (strcmp(leaf->name, name) == 0) { + return leaf; + } + } + + for (int i = 0; i < cgraph->n_nodes; i++) { + struct ggml_tensor * node = cgraph->nodes[i]; + + if (strcmp(node->name, name) == 0) { + return node; + } + } + + return NULL; +} + +struct ggml_tensor * ggml_graph_get_grad(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { + const size_t igrad = ggml_hash_find(&cgraph->visited_hash_set, node); + return igrad != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, igrad) && cgraph->grads ? cgraph->grads[igrad] : NULL; +} + +struct ggml_tensor * ggml_graph_get_grad_acc(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { + const size_t igrad = ggml_hash_find(&cgraph->visited_hash_set, node); + return igrad != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, igrad) && cgraph->grad_accs ? cgraph->grad_accs[igrad] : NULL; +} + +void ggml_graph_print(const struct ggml_cgraph * cgraph) { + GGML_LOG_INFO("=== GRAPH ===\n"); + + GGML_LOG_INFO("n_nodes = %d\n", cgraph->n_nodes); + for (int i = 0; i < cgraph->n_nodes; i++) { + struct ggml_tensor * node = cgraph->nodes[i]; + + GGML_LOG_INFO(" - %3d: [ %5" PRId64 ", %5" PRId64 ", %5" PRId64 "] %16s %s\n", + i, + node->ne[0], node->ne[1], node->ne[2], + ggml_op_name(node->op), (node->flags & GGML_TENSOR_FLAG_PARAM) ? "x" : + ggml_graph_get_grad(cgraph, node) ? "g" : " "); + } + + GGML_LOG_INFO("n_leafs = %d\n", cgraph->n_leafs); + for (int i = 0; i < cgraph->n_leafs; i++) { + struct ggml_tensor * node = cgraph->leafs[i]; + + GGML_LOG_INFO(" - %3d: [ %5" PRId64 ", %5" PRId64 "] %8s %16s\n", + i, + node->ne[0], node->ne[1], + ggml_op_name(node->op), + ggml_get_name(node)); + } + + GGML_LOG_INFO("========================================\n"); +} + +static int ggml_node_list_find_tensor(const struct ggml_cgraph * cgraph, + const int * idxs, + int count, + const struct ggml_tensor * tensor) { + GGML_ASSERT(cgraph && idxs); + for (int i = 0; i < count; ++i) { + const int node_idx = idxs[i]; + + if (node_idx >= cgraph->n_nodes) { + return -1; + } + if (cgraph->nodes[node_idx] == tensor) { + return i; + } + } + return -1; +} + +bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph, + const int * node_idxs, + int count, + const enum ggml_op * ops, + const int * outputs, + int num_outputs) { + GGML_ASSERT(outputs && num_outputs > 0); + + for (int i = 0; i < count; ++i) { + if (node_idxs[i] >= cgraph->n_nodes) { + return false; + } + + const struct ggml_tensor * node = cgraph->nodes[node_idxs[i]]; + + if (node->op != ops[i]) { + return false; + } + + if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + return false; + } + + if (ggml_node_list_find_tensor(cgraph, outputs, num_outputs, node) != -1) { + continue; + } + + if (node->flags & GGML_TENSOR_FLAG_OUTPUT) { + return false; + } + + int subgraph_uses = 0; + for (int j = i + 1; j < count; ++j) { + const struct ggml_tensor * other_node = cgraph->nodes[node_idxs[j]]; + for (int src_idx = 0; src_idx < GGML_MAX_SRC; src_idx++) { + if (other_node->src[src_idx] == node) { + subgraph_uses++; + } + } + } + + if (subgraph_uses != ggml_node_get_use_count(cgraph, node_idxs[i])) { + return false; + } + + // if node is a view, check if the view_src and all it's parent view_srcs are within the subgraph + struct ggml_tensor * view_src = node->view_src; + while (view_src) { + if (ggml_node_list_find_tensor(cgraph, node_idxs, count, view_src) == -1) { + return false; + } + view_src = view_src->view_src; + } + } + + return true; +} + +// check if node is part of the graph +static bool ggml_graph_find(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { + if (cgraph == NULL) { + return true; + } + + for (int i = 0; i < cgraph->n_nodes; i++) { + if (cgraph->nodes[i] == node) { + return true; + } + } + + return false; +} + +static struct ggml_tensor * ggml_graph_get_parent(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { + for (int i = 0; i < cgraph->n_nodes; i++) { + struct ggml_tensor * parent = cgraph->nodes[i]; + struct ggml_tensor * grad = ggml_graph_get_grad(cgraph, parent); + + if (grad == node) { + return parent; + } + } + + return NULL; +} + +static void ggml_graph_dump_dot_node_edge(FILE * fp, const struct ggml_cgraph * gb, struct ggml_tensor * node, struct ggml_tensor * parent, const char * label) { + struct ggml_tensor * gparent = ggml_graph_get_parent(gb, node); + struct ggml_tensor * gparent0 = ggml_graph_get_parent(gb, parent); + fprintf(fp, " \"%p\" -> \"%p\" [ arrowhead = %s; style = %s; label = \"%s\"; ]\n", + gparent0 ? (void *) gparent0 : (void *) parent, + gparent ? (void *) gparent : (void *) node, + gparent ? "empty" : "vee", + gparent ? "dashed" : "solid", + label); +} + +static void ggml_graph_dump_dot_leaf_edge(FILE * fp, struct ggml_tensor * node, struct ggml_tensor * parent, const char * label) { + fprintf(fp, " \"%p\" -> \"%p\" [ label = \"%s\"; ]\n", + (void *) parent, + (void *) node, + label); +} + +void ggml_graph_dump_dot(const struct ggml_cgraph * gb, const struct ggml_cgraph * cgraph, const char * filename) { + char color[16]; + + FILE * fp = ggml_fopen(filename, "w"); + GGML_ASSERT(fp); + + fprintf(fp, "digraph G {\n"); + fprintf(fp, " newrank = true;\n"); + fprintf(fp, " rankdir = TB;\n"); + + for (int i = 0; i < gb->n_nodes; i++) { + struct ggml_tensor * node = gb->nodes[i]; + struct ggml_tensor * grad = ggml_graph_get_grad(gb, node); + + if (ggml_graph_get_parent(gb, node) != NULL) { + continue; + } + + if (node->flags & GGML_TENSOR_FLAG_PARAM) { + snprintf(color, sizeof(color), "yellow"); + } else if (grad) { + if (ggml_graph_find(cgraph, node)) { + snprintf(color, sizeof(color), "green"); + } else { + snprintf(color, sizeof(color), "lightblue"); + } + } else { + snprintf(color, sizeof(color), "white"); + } + + fprintf(fp, " \"%p\" [ " + "style = filled; fillcolor = %s; shape = record; " + "label=\"", + (void *) node, color); + + if (strlen(node->name) > 0) { + fprintf(fp, "%s (%s)|", node->name, ggml_type_name(node->type)); + } else { + fprintf(fp, "(%s)|", ggml_type_name(node->type)); + } + + if (ggml_is_matrix(node)) { + fprintf(fp, "%d [%" PRId64 ", %" PRId64 "] | %s", i, node->ne[0], node->ne[1], ggml_op_symbol(node->op)); + } else { + fprintf(fp, "%d [%" PRId64 ", %" PRId64 ", %" PRId64 "] | %s", i, node->ne[0], node->ne[1], node->ne[2], ggml_op_symbol(node->op)); + } + + if (grad) { + fprintf(fp, " | %s\"; ]\n", ggml_op_symbol(grad->op)); + } else { + fprintf(fp, "\"; ]\n"); + } + } + + for (int i = 0; i < gb->n_leafs; i++) { + struct ggml_tensor * node = gb->leafs[i]; + + snprintf(color, sizeof(color), "pink"); + + fprintf(fp, " \"%p\" [ " + "style = filled; fillcolor = %s; shape = record; " + "label=\"", + (void *) node, color); + + if (strlen(node->name) > 0) { + fprintf(fp, "%s (%s)|", node->name, ggml_type_name(node->type)); + } else { + fprintf(fp, "(%s)|", ggml_type_name(node->type)); + } + + fprintf(fp, "CONST %d [%" PRId64 ", %" PRId64 "]", i, node->ne[0], node->ne[1]); + if (ggml_nelements(node) < 5 && node->data != NULL) { + fprintf(fp, " | ("); + for (int j = 0; j < ggml_nelements(node); j++) { + // FIXME: use ggml-backend to obtain the tensor data + //if (node->type == GGML_TYPE_I8 || node->type == GGML_TYPE_I16 || node->type == GGML_TYPE_I32) { + // fprintf(fp, "%d", ggml_get_i32_1d(node, j)); + //} + //else if (node->type == GGML_TYPE_F32 || + // node->type == GGML_TYPE_F16 || + // node->type == GGML_TYPE_BF16) { + // fprintf(fp, "%.1e", (double)ggml_get_f32_1d(node, j)); + //} + //else + { + fprintf(fp, "#"); + } + if (j < ggml_nelements(node) - 1) { + fprintf(fp, ", "); + } + } + fprintf(fp, ")"); + } + fprintf(fp, "\"; ]\n"); + } + + for (int i = 0; i < gb->n_nodes; i++) { + struct ggml_tensor * node = gb->nodes[i]; + + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (node->src[j]) { + char label[16]; + snprintf(label, sizeof(label), "src %d", j); + ggml_graph_dump_dot_node_edge(fp, gb, node, node->src[j], label); + } + } + } + + for (int i = 0; i < gb->n_leafs; i++) { + struct ggml_tensor * node = gb->leafs[i]; + + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (node->src[j]) { + char label[16]; + snprintf(label, sizeof(label), "src %d", j); + ggml_graph_dump_dot_leaf_edge(fp, node, node->src[j], label); + } + } + } + + fprintf(fp, "}\n"); + + fclose(fp); + + GGML_LOG_INFO("%s: dot -Tpng %s -o %s.png && open %s.png\n", __func__, filename, filename, filename); +} + +//////////////////////////////////////////////////////////////////////////////// + +void ggml_set_input(struct ggml_tensor * tensor) { + tensor->flags |= GGML_TENSOR_FLAG_INPUT; +} + +void ggml_set_output(struct ggml_tensor * tensor) { + tensor->flags |= GGML_TENSOR_FLAG_OUTPUT; +} + +void ggml_set_param(struct ggml_tensor * tensor) { + GGML_ASSERT(tensor->op == GGML_OP_NONE); + tensor->flags |= GGML_TENSOR_FLAG_PARAM; +} + +void ggml_set_loss(struct ggml_tensor * tensor) { + GGML_ASSERT(ggml_is_scalar(tensor)); + GGML_ASSERT(tensor->type == GGML_TYPE_F32); + tensor->flags |= GGML_TENSOR_FLAG_LOSS; +} + +//////////////////////////////////////////////////////////////////////////////// + +void ggml_quantize_init(enum ggml_type type) { + ggml_critical_section_start(); + + switch (type) { + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: iq2xs_init_impl(type); break; + case GGML_TYPE_IQ3_XXS: iq3xs_init_impl(256); break; + case GGML_TYPE_IQ3_S: iq3xs_init_impl(512); break; + default: // nothing + break; + } + + ggml_critical_section_end(); +} + +void ggml_quantize_free(void) { + ggml_critical_section_start(); + + iq2xs_free_impl(GGML_TYPE_IQ2_XXS); + iq2xs_free_impl(GGML_TYPE_IQ2_XS); + iq2xs_free_impl(GGML_TYPE_IQ2_S); + iq2xs_free_impl(GGML_TYPE_IQ1_S); + iq2xs_free_impl(GGML_TYPE_IQ1_M); + iq3xs_free_impl(256); + iq3xs_free_impl(512); + + ggml_critical_section_end(); +} + +bool ggml_quantize_requires_imatrix(enum ggml_type type) { + return + type == GGML_TYPE_IQ2_XXS || + type == GGML_TYPE_IQ2_XS || + type == GGML_TYPE_IQ1_S;// || + //type == GGML_TYPE_IQ1_M; +} + +size_t ggml_quantize_chunk( + enum ggml_type type, + const float * src, + void * dst, + int64_t start, + int64_t nrows, + int64_t n_per_row, + const float * imatrix) { + const int64_t n = nrows * n_per_row; + + if (ggml_quantize_requires_imatrix(type)) { + GGML_ASSERT(imatrix != NULL); + } + + GGML_ASSERT(start % type_traits[type].blck_size == 0); + GGML_ASSERT(start % n_per_row == 0); + + ggml_quantize_init(type); // this is noop if already initialized + + const size_t start_row = start / n_per_row; + const size_t row_size = ggml_row_size(type, n_per_row); + + size_t result = 0; + + switch (type) { + case GGML_TYPE_Q1_0: result = quantize_q1_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q4_0: result = quantize_q4_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q4_1: result = quantize_q4_1 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q5_0: result = quantize_q5_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q5_1: result = quantize_q5_1 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q8_0: result = quantize_q8_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_MXFP4: result = quantize_mxfp4 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_NVFP4: result = quantize_nvfp4 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q2_K: result = quantize_q2_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q3_K: result = quantize_q3_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q4_K: result = quantize_q4_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q5_K: result = quantize_q5_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q6_K: result = quantize_q6_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_TQ1_0: result = quantize_tq1_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_TQ2_0: result = quantize_tq2_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ2_XXS: result = quantize_iq2_xxs(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ2_XS: result = quantize_iq2_xs (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ3_XXS: result = quantize_iq3_xxs(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ3_S: result = quantize_iq3_s (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ2_S: result = quantize_iq2_s (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ1_S: result = quantize_iq1_s (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ1_M: result = quantize_iq1_m (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ4_NL: result = quantize_iq4_nl (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ4_XS: result = quantize_iq4_xs (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_F16: + { + size_t elemsize = sizeof(ggml_fp16_t); + ggml_fp32_to_fp16_row(src + start, (ggml_fp16_t *)dst + start, n); + result = n * elemsize; + } break; + case GGML_TYPE_BF16: + { + size_t elemsize = sizeof(ggml_bf16_t); + ggml_fp32_to_bf16_row_ref(src + start, (ggml_bf16_t *)dst + start, n); + result = n * elemsize; + } break; + case GGML_TYPE_F32: + { + size_t elemsize = sizeof(float); + result = n * elemsize; + memcpy((uint8_t *)dst + start * elemsize, src + start, result); + } break; + default: + assert(false); + } + + GGML_ASSERT(result == nrows * row_size); + + return result; +} + +//////////////////////////////////////////////////////////////////////////////// + +void ggml_log_get(ggml_log_callback * log_callback, void ** user_data) { + *log_callback = g_logger_state.log_callback; + *user_data = g_logger_state.log_callback_user_data; +} + +void ggml_log_set(ggml_log_callback log_callback, void * user_data) { + g_logger_state.log_callback = log_callback ? log_callback : ggml_log_callback_default; + g_logger_state.log_callback_user_data = user_data; +} + +void ggml_threadpool_params_init(struct ggml_threadpool_params * p, int n_threads) { + p->n_threads = n_threads; + p->prio = 0; // default priority (usually means normal or inherited) + p->poll = 50; // hybrid-polling enabled + p->strict_cpu = false; // no strict placement (all threads share same cpumask) + p->paused = false; // threads are ready to go + memset(p->cpumask, 0, GGML_MAX_N_THREADS); // all-zero means use the default affinity (usually inherited) +} + +struct ggml_threadpool_params ggml_threadpool_params_default(int n_threads) { + struct ggml_threadpool_params p; + ggml_threadpool_params_init(&p, n_threads); + return p; +} + +bool ggml_threadpool_params_match(const struct ggml_threadpool_params * p0, const struct ggml_threadpool_params * p1) { + if (p0->n_threads != p1->n_threads ) return false; + if (p0->prio != p1->prio ) return false; + if (p0->poll != p1->poll ) return false; + if (p0->strict_cpu != p1->strict_cpu ) return false; + return memcmp(p0->cpumask, p1->cpumask, GGML_MAX_N_THREADS) == 0; +} + + +/* ---- allocator implementation ---- */ + +#include "ggml-alloc.h.inc" +#include "ggml-backend-impl-defs.inc" +#include "ggml.h.inc" +#include "ggml-impl-defs.inc" + +#include +#include +#include +#include +#include +#include + +#define MAX(a, b) ((a) > (b) ? (a) : (b)) +#define MAX_FREE_BLOCKS 256 + +//#define GGML_ALLOCATOR_DEBUG + +//#define AT_PRINTF(...) GGML_LOG_DEBUG(__VA_ARGS__) +#define AT_PRINTF(...) + +// ops that return true for this function must not use restrict pointers for their backend implementations +bool ggml_op_can_inplace(enum ggml_op op) { + switch (op) { + case GGML_OP_FILL: + case GGML_OP_SCALE: + case GGML_OP_DIAG_MASK_ZERO: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_ADD: + case GGML_OP_ADD_ID: + case GGML_OP_ADD1: + case GGML_OP_SUB: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_SQR: + case GGML_OP_SQRT: + case GGML_OP_LOG: + case GGML_OP_UNARY: + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: + case GGML_OP_SILU_BACK: + case GGML_OP_RMS_NORM: + case GGML_OP_RMS_NORM_BACK: + case GGML_OP_SOFT_MAX: + case GGML_OP_SOFT_MAX_BACK: + return true; + + default: + return false; + } +} + +#define aligned_offset(buffer, offset, alignment) ggml_aligned_offset_rust(buffer, offset, alignment) + +// tallocr + +struct ggml_tallocr ggml_tallocr_new(ggml_backend_buffer_t buffer) { + void * base = ggml_backend_buffer_get_base(buffer); + size_t align = ggml_backend_buffer_get_alignment(buffer); + + assert(align && !(align & (align - 1))); // power of 2 + + struct ggml_tallocr talloc = (struct ggml_tallocr) { + /*.buffer = */ buffer, + /*.base = */ base, + /*.alignment = */ align, + /*.offset = */ aligned_offset(base, 0, align), + }; + return talloc; +} + +enum ggml_status ggml_tallocr_alloc(struct ggml_tallocr * talloc, struct ggml_tensor * tensor) { + size_t size = ggml_backend_buffer_get_alloc_size(talloc->buffer, tensor); + size = GGML_PAD(size, talloc->alignment); + + if (talloc->offset + size > ggml_backend_buffer_get_size(talloc->buffer)) { + GGML_LOG_ERROR("%s: not enough space in the buffer to allocate %s (needed %zu, available %zu)\n", + __func__, tensor->name, size, ggml_backend_buffer_get_size(talloc->buffer) - talloc->offset); + GGML_ABORT("not enough space in the buffer"); + } + + void * addr = (char *)ggml_backend_buffer_get_base(talloc->buffer) + talloc->offset; + talloc->offset += size; + + assert(((uintptr_t)addr % talloc->alignment) == 0); + + return ggml_backend_tensor_alloc(talloc->buffer, tensor, addr); +} + +// dynamic tensor allocator + +#define GGML_VBUFFER_MAX_CHUNKS 16 + +// relative memory address within an allocation that can be split into multiple buffers (chunks) +struct buffer_address { + int chunk; // index of a backend buffer + size_t offset; // local memory offset within the buffer +}; + +static const struct buffer_address GGML_BUFFER_ADDRESS_INVALID = { -1, SIZE_MAX }; + +struct free_block { + size_t offset; + size_t size; +}; + +struct tallocr_chunk { + struct free_block free_blocks[MAX_FREE_BLOCKS]; + int n_free_blocks; + size_t max_size; +}; + +struct ggml_dyn_tallocr { + size_t alignment; + size_t max_chunk_size; + struct tallocr_chunk * chunks[GGML_VBUFFER_MAX_CHUNKS]; + int n_chunks; + +#ifdef GGML_ALLOCATOR_DEBUG + struct { + const struct ggml_tensor * tensor; + struct buffer_address addr; + } allocated_tensors[1024]; +#endif +}; + +static void ggml_dyn_tallocr_insert_block(struct tallocr_chunk * chunk, size_t offset, size_t size) { + GGML_ASSERT(chunk->n_free_blocks < MAX_FREE_BLOCKS && "out of free blocks"); + // insert the new block in the correct position to keep the array sorted by address (to make merging blocks faster) + int insert_pos = 0; + while (insert_pos < chunk->n_free_blocks && chunk->free_blocks[insert_pos].offset < offset) { + insert_pos++; + } + // shift all blocks from insert_pos onward to make room for the new block + for (int i = chunk->n_free_blocks; i > insert_pos; i--) { + chunk->free_blocks[i] = chunk->free_blocks[i-1]; + } + // insert the new block + chunk->free_blocks[insert_pos].offset = offset; + chunk->free_blocks[insert_pos].size = size; + chunk->n_free_blocks++; +} + +static void ggml_dyn_tallocr_remove_block(struct tallocr_chunk * chunk, int idx) { + // shift all elements after idx by 1 to the left, overwriting the element at idx + for (int i = idx; i < chunk->n_free_blocks; i++) { + chunk->free_blocks[i] = chunk->free_blocks[i+1]; + } + chunk->n_free_blocks--; +} + +static int ggml_dyn_tallocr_new_chunk(struct ggml_dyn_tallocr * alloc, size_t min_size) { + if (alloc->n_chunks >= GGML_VBUFFER_MAX_CHUNKS) { + return -1; + } + struct tallocr_chunk * chunk = calloc(1, sizeof(struct tallocr_chunk)); + chunk->n_free_blocks = 1; + chunk->free_blocks[0].offset = 0; + // available space in a chunk is limited to max_chunk_size, but can be higher if: + // 1. a single tensor exceeds the maximum, and cannot fit any other way + // 2. we are running out of chunks + // backends will either manage to allocate the larger size, or report an error. + chunk->free_blocks[0].size = MAX(min_size, alloc->max_chunk_size); + if (alloc->n_chunks == GGML_VBUFFER_MAX_CHUNKS - 1) { + chunk->free_blocks[0].size = SIZE_MAX/2; + } + alloc->chunks[alloc->n_chunks] = chunk; + alloc->n_chunks++; + return alloc->n_chunks - 1; +} + +#ifdef GGML_ALLOCATOR_DEBUG +static void add_allocated_tensor(struct ggml_dyn_tallocr * alloc, struct buffer_address addr, const struct ggml_tensor * tensor) { + for (int i = 0; i < 1024; i++) { + if (alloc->allocated_tensors[i].tensor == NULL) { + alloc->allocated_tensors[i].tensor = tensor; + alloc->allocated_tensors[i].addr = addr; + return; + } + } + GGML_ABORT("out of allocated_tensors"); +} +static void remove_allocated_tensor(struct ggml_dyn_tallocr * alloc, struct buffer_address addr, const struct ggml_tensor * tensor) { + for (int i = 0; i < 1024; i++) { + if (alloc->allocated_tensors[i].addr.chunk == addr.chunk && alloc->allocated_tensors[i].addr.offset == addr.offset) { + alloc->allocated_tensors[i].tensor = NULL; + return; + } + } + GGML_ABORT("tried to free tensor %s not found\n", tensor->name); +} +#endif + +static struct buffer_address ggml_dyn_tallocr_alloc(struct ggml_dyn_tallocr * alloc, size_t size, const struct ggml_tensor * tensor) { + size = aligned_offset(NULL, size, alloc->alignment); + + AT_PRINTF("%s: allocating %s (%zu bytes) - ", __func__, tensor->name, size); + + int best_fit_chunk = -1; + int best_fit_block = -1; + size_t max_avail = 0; + + // find the best fitting free block besides the last block, within any chunk + for (int c = 0; c < alloc->n_chunks; ++c) { + struct tallocr_chunk * chunk = alloc->chunks[c]; + size_t best_fit_size = SIZE_MAX; + for (int i = 0; i < chunk->n_free_blocks - 1; i++) { + struct free_block * block = &chunk->free_blocks[i]; + max_avail = MAX(max_avail, block->size); + if (block->size >= size && block->size <= best_fit_size) { + best_fit_chunk = c; + best_fit_block = i; + best_fit_size = block->size; + } + } + } + + if (best_fit_block == -1) { + // no suitable block found, try the last block (this may grow a chunks size) + int64_t best_reuse = INT64_MIN; + for (int c = 0; c < alloc->n_chunks; ++c) { + struct tallocr_chunk * chunk = alloc->chunks[c]; + if (chunk->n_free_blocks > 0) { + struct free_block * block = &chunk->free_blocks[chunk->n_free_blocks - 1]; + max_avail = MAX(max_avail, block->size); + int64_t reuse_factor = chunk->max_size - block->offset - size; + // reuse_factor < 0 : amount of extra memory that needs to be allocated + // reuse_factor = 0 : allocated free space exactly matches tensor size + // reuse_factor > 0 : superfluous memory that will remain unused + bool better_reuse = best_reuse < 0 && reuse_factor > best_reuse; + bool better_fit = reuse_factor >= 0 && reuse_factor < best_reuse; + if (block->size >= size && (better_reuse || better_fit)) { + best_fit_chunk = c; + best_fit_block = chunk->n_free_blocks - 1; + best_reuse = reuse_factor; + } + } + } + } + + if (best_fit_block == -1) { + // none of the existing chunks have enough space left + best_fit_chunk = ggml_dyn_tallocr_new_chunk(alloc, size); + best_fit_block = 0; + } + if (best_fit_chunk == -1) { + // since the last chunk always has virtually endless memory, this should never happen + GGML_LOG_ERROR("%s: not enough space in the buffer to allocate %zu bytes, largest block available %zu bytes\n", + __func__, size, max_avail); + GGML_ABORT("graph allocation: failed to reserve memory"); + } + + struct tallocr_chunk * chunk = alloc->chunks[best_fit_chunk]; + struct free_block * block = &chunk->free_blocks[best_fit_block]; + struct buffer_address addr = {.chunk = best_fit_chunk, .offset = block->offset }; + block->offset += size; + block->size -= size; + if (block->size == 0) { + // remove block if empty + ggml_dyn_tallocr_remove_block(chunk, best_fit_block); + } + + AT_PRINTF("block %d, offset %zu, chunk %d\n", best_fit_block, addr.offset, addr.chunk); + +#ifdef GGML_ALLOCATOR_DEBUG + add_allocated_tensor(alloc, addr, tensor); + size_t cur_max = addr.offset + size; + if (cur_max > chunk->max_size) { + // sort allocated_tensors by chunk/offset + for (int i = 0; i < 1024; i++) { + for (int j = i + 1; j < 1024; j++) { + if (ggml_buffer_address_less_rust( + alloc->allocated_tensors[j].addr.chunk, + alloc->allocated_tensors[j].addr.offset, + alloc->allocated_tensors[i].addr.chunk, + alloc->allocated_tensors[i].addr.offset)) { + const struct ggml_tensor * tmp_tensor = alloc->allocated_tensors[i].tensor; + struct buffer_address tmp_addr = alloc->allocated_tensors[i].addr; + alloc->allocated_tensors[i].tensor = alloc->allocated_tensors[j].tensor; + alloc->allocated_tensors[i].addr = alloc->allocated_tensors[j].addr; + alloc->allocated_tensors[j].tensor = tmp_tensor; + alloc->allocated_tensors[j].addr = tmp_addr; + } + } + } + GGML_LOG_DEBUG("max_size[%d] = %.2f MB: tensors: ", addr.chunk, cur_max / 1024.0 / 1024.0); + for (int i = 0; i < 1024; i++) { + if (alloc->allocated_tensors[i].tensor) { + GGML_LOG_DEBUG("%s [%d: %zx-%zx] (%.2f MB) ", alloc->allocated_tensors[i].tensor->name, + alloc->allocated_tensors[i].addr.chunk, + alloc->allocated_tensors[i].addr.offset, + alloc->allocated_tensors[i].addr.offset + ggml_nbytes(alloc->allocated_tensors[i].tensor), + ggml_nbytes(alloc->allocated_tensors[i].tensor) / 1024.0 / 1024.0); + } + } + GGML_LOG_DEBUG("\n"); + } +#endif + + chunk->max_size = MAX(chunk->max_size, addr.offset + size); + + return addr; + + GGML_UNUSED(tensor); +} + +// this is a very naive implementation, but for our case the number of free blocks should be very small +static void ggml_dyn_tallocr_free_bytes(struct ggml_dyn_tallocr * alloc, struct buffer_address addr, size_t size) { + size = aligned_offset(NULL, size, alloc->alignment); + + struct tallocr_chunk * chunk = alloc->chunks[addr.chunk]; + + // see if we can merge with an existing block + for (int i = 0; i < chunk->n_free_blocks; i++) { + struct free_block * block = &chunk->free_blocks[i]; + // check if ptr is at the end of the block + if (block->offset + block->size == addr.offset) { + block->size += size; + // check if we can merge with the next block + if (i < chunk->n_free_blocks - 1) { + struct free_block * next = &chunk->free_blocks[i+1]; + if (block->offset + block->size == next->offset) { + block->size += next->size; + ggml_dyn_tallocr_remove_block(chunk, i+1); + } + } + return; + } + // check if ptr is at the beginning of the block + if (addr.offset + size == block->offset) { + block->offset = addr.offset; + block->size += size; + // check if we can merge with the previous block + if (i > 0) { + struct free_block * prev = &chunk->free_blocks[i-1]; + if (prev->offset + prev->size == block->offset) { + prev->size += block->size; + ggml_dyn_tallocr_remove_block(chunk, i); + } + } + return; + } + } + // otherwise, add a new block + ggml_dyn_tallocr_insert_block(chunk, addr.offset, size); +} + +static void ggml_dyn_tallocr_reset(struct ggml_dyn_tallocr * alloc) { + for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS; i++) { + free(alloc->chunks[i]); + alloc->chunks[i] = NULL; + } + alloc->n_chunks = 0; + +#ifdef GGML_ALLOCATOR_DEBUG + for (int i = 0; i < 1024; i++) { + alloc->allocated_tensors[i].tensor = NULL; + } +#endif +} + +static struct ggml_dyn_tallocr * ggml_dyn_tallocr_new(size_t alignment, size_t max_buffer_size) { + struct ggml_dyn_tallocr * alloc = (struct ggml_dyn_tallocr *)malloc(sizeof(struct ggml_dyn_tallocr)); + + *alloc = (struct ggml_dyn_tallocr) { + /*.alignment = */ alignment, + /*.max_chunk_size = */ MIN(max_buffer_size, SIZE_MAX/2), // clamp to avoid overflows + /*.chunks = */ {NULL}, + /*.n_chunks = */ 0, +#ifdef GGML_ALLOCATOR_DEBUG + /*.allocated_tensors = */ {{0}}, +#endif + }; + + ggml_dyn_tallocr_reset(alloc); + + return alloc; +} + +static void ggml_dyn_tallocr_free(struct ggml_dyn_tallocr * alloc) { + for (int i = 0; i < alloc->n_chunks; ++i) { + free(alloc->chunks[i]); + } + free(alloc); +} + +static size_t ggml_dyn_tallocr_max_size(struct ggml_dyn_tallocr * alloc, int chunk) { + return chunk < alloc->n_chunks ? alloc->chunks[chunk]->max_size : 0; +} + + +// virtual buffer with contiguous memory range, split into multiple backend buffers (chunks) + +struct vbuffer { + ggml_backend_buffer_t chunks[GGML_VBUFFER_MAX_CHUNKS]; +}; + +static void ggml_vbuffer_free(struct vbuffer * buf) { + if (buf == NULL) { + return; + } + for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS; ++i) { + ggml_backend_buffer_free(buf->chunks[i]); + } + free(buf); +} + +static size_t ggml_vbuffer_chunk_size(struct vbuffer * buf, int chunk) { + return buf->chunks[chunk] ? ggml_backend_buffer_get_size(buf->chunks[chunk]) : 0; +} + +static size_t ggml_vbuffer_size(struct vbuffer * buf) { + size_t size = 0; + for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS && buf->chunks[i]; ++i) { + size += ggml_backend_buffer_get_size(buf->chunks[i]); + } + return size; +} + +static struct vbuffer * ggml_vbuffer_alloc(ggml_backend_buffer_type_t buft, const struct ggml_dyn_tallocr * talloc, enum ggml_backend_buffer_usage usage) { + struct vbuffer * buf = (struct vbuffer *)calloc(1, sizeof(struct vbuffer)); + if (buf == NULL) { + return NULL; + } + + for (int n = 0; n < talloc->n_chunks; n++) { + size_t chunk_size = talloc->chunks[n]->max_size; + buf->chunks[n] = ggml_backend_buft_alloc_buffer(buft, chunk_size); + if (buf->chunks[n] == NULL) { + ggml_vbuffer_free(buf); + return NULL; + } + ggml_backend_buffer_set_usage(buf->chunks[n], usage); + } + return buf; +} + +static void ggml_vbuffer_tensor_alloc(struct vbuffer * buf, struct ggml_tensor * tensor, struct buffer_address buf_addr) { + void * base = ggml_backend_buffer_get_base(buf->chunks[buf_addr.chunk]); + void * addr = (char *)base + buf_addr.offset; + ggml_backend_tensor_alloc(buf->chunks[buf_addr.chunk], tensor, addr); +} + +static void ggml_vbuffer_reset(struct vbuffer * buf) { + for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS && buf->chunks[i]; ++i) { + ggml_backend_buffer_reset(buf->chunks[i]); + } +} + + +///////////////////////////////////// + +// graph allocator + +struct hash_node { + int n_children; + int n_views; + int buffer_id; + struct buffer_address addr; + bool allocated; +}; + +struct tensor_alloc { + int buffer_id; + struct buffer_address addr; + size_t size_max; // 0 = pre-allocated, unused, or view +}; + +struct leaf_alloc { + struct tensor_alloc leaf; +}; + +struct node_alloc { + struct tensor_alloc dst; + struct tensor_alloc src[GGML_MAX_SRC]; +}; + +struct ggml_gallocr { + ggml_backend_buffer_type_t * bufts; // [n_buffers] + struct vbuffer ** buffers; // [n_buffers] + struct ggml_dyn_tallocr ** buf_tallocs; // [n_buffers] + int n_buffers; + + struct ggml_hash_set hash_set; + struct hash_node * hash_values; // [hash_set.size] + + struct node_alloc * node_allocs; // [n_nodes] + int n_nodes; + + struct leaf_alloc * leaf_allocs; // [n_leafs] + int n_leafs; +}; + +ggml_gallocr_t ggml_gallocr_new_n(ggml_backend_buffer_type_t * bufts, int n_bufs) { + ggml_gallocr_t galloc = (ggml_gallocr_t)calloc(1, sizeof(struct ggml_gallocr)); + GGML_ASSERT(galloc != NULL); + + galloc->bufts = calloc(n_bufs, sizeof(ggml_backend_buffer_type_t)); + GGML_ASSERT(galloc->bufts != NULL); + + galloc->buffers = calloc(n_bufs, sizeof(struct vbuffer *)); + GGML_ASSERT(galloc->buffers != NULL); + + galloc->buf_tallocs = calloc(n_bufs, sizeof(struct ggml_dyn_tallocr *)); + GGML_ASSERT(galloc->buf_tallocs != NULL); + + for (int i = 0; i < n_bufs; i++) { + galloc->bufts[i] = bufts[i]; + galloc->buffers[i] = NULL; + + // check if the same buffer type is used multiple times and reuse the same allocator + for (int j = 0; j < i; j++) { + if (bufts[i] == bufts[j]) { + galloc->buf_tallocs[i] = galloc->buf_tallocs[j]; + break; + } + } + + if (galloc->buf_tallocs[i] == NULL) { + size_t alignment = ggml_backend_buft_get_alignment(bufts[i]); + size_t max_size = ggml_backend_buft_get_max_size(bufts[i]); + galloc->buf_tallocs[i] = ggml_dyn_tallocr_new(alignment, max_size); + } + } + galloc->n_buffers = n_bufs; + + return galloc; +} + +ggml_gallocr_t ggml_gallocr_new(ggml_backend_buffer_type_t buft) { + return ggml_gallocr_new_n(&buft, 1); +} + +void ggml_gallocr_free(ggml_gallocr_t galloc) { + if (galloc == NULL) { + return; + } + + for (int i = 0; i < galloc->n_buffers; i++) { + if (galloc->buffers != NULL) { + // skip if already freed + bool freed = false; + for (int j = 0; j < i; j++) { + if (galloc->buffers[j] == galloc->buffers[i]) { + freed = true; + break; + } + } + if (!freed) { + ggml_vbuffer_free(galloc->buffers[i]); + } + } + if (galloc->buf_tallocs != NULL) { + // skip if already freed + bool freed = false; + for (int j = 0; j < i; j++) { + if (galloc->buf_tallocs[j] == galloc->buf_tallocs[i]) { + freed = true; + break; + } + } + if (!freed) { + ggml_dyn_tallocr_free(galloc->buf_tallocs[i]); + } + } + } + + ggml_hash_set_free(&galloc->hash_set); + free(galloc->hash_values); + free(galloc->bufts); + free(galloc->buffers); + free(galloc->buf_tallocs); + free(galloc->node_allocs); + free(galloc->leaf_allocs); + free(galloc); +} + +typedef struct ggml_gallocr * ggml_gallocr_t; + +static struct hash_node * ggml_gallocr_hash_get(ggml_gallocr_t galloc, struct ggml_tensor * t) { + size_t i = ggml_hash_find_or_insert(&galloc->hash_set, t); + return &galloc->hash_values[i]; +} + +static bool ggml_gallocr_is_own(ggml_gallocr_t galloc, struct ggml_tensor * t) { + return ggml_gallocr_hash_get(galloc, t)->allocated; +} + +static bool ggml_gallocr_is_allocated(ggml_gallocr_t galloc, struct ggml_tensor * t) { + return t->data != NULL // tensor data already set externally + || t->buffer // tensor on external buffer (but not yet allocated) + || ggml_gallocr_is_own(galloc, t); // tensor will be allocated by galloc +} + +// free the extra space at the end if the new tensor is smaller +static void ggml_gallocr_free_extra_space(ggml_gallocr_t galloc, struct ggml_tensor * node, struct ggml_tensor * parent) { + struct hash_node * hn = ggml_gallocr_hash_get(galloc, node); + struct hash_node * p_hn = ggml_gallocr_hash_get(galloc, parent); + + size_t parent_size = ggml_backend_buft_get_alloc_size(galloc->bufts[p_hn->buffer_id], parent); + size_t node_size = ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], node); + + GGML_ASSERT(parent_size >= node_size); + + // note: we want after the freeing the chunks to continue to be aligned + struct ggml_dyn_tallocr * p_alloc = galloc->buf_tallocs[p_hn->buffer_id]; + parent_size = aligned_offset(NULL, parent_size, p_alloc->alignment); + node_size = aligned_offset(NULL, node_size, p_alloc->alignment); + + if (parent_size > node_size) { + struct buffer_address p_addr = p_hn->addr; + p_addr.offset += node_size; + size_t extra_size = parent_size - node_size; + AT_PRINTF("freeing extra %zu bytes from parent %s for %s\n", extra_size, parent->name, node->name); + ggml_dyn_tallocr_free_bytes(p_alloc, p_addr, extra_size); + } +} + +static void ggml_gallocr_allocate_node(ggml_gallocr_t galloc, struct ggml_tensor * node, int buffer_id) { + GGML_ASSERT(buffer_id >= 0); + struct hash_node * hn = ggml_gallocr_hash_get(galloc, node); + + if (!ggml_gallocr_is_allocated(galloc, node) && !ggml_impl_is_view(node)) { + hn->allocated = true; + assert(hn->addr.offset == 0); + + // try to reuse a parent's buffer (inplace) + if (ggml_op_can_inplace(node->op)) { + for (int i = 0; i < GGML_MAX_SRC; i++) { + struct ggml_tensor * parent = node->src[i]; + if (parent == NULL) { + continue; + } + + // if the node's data is external, then we cannot re-use it + if (!ggml_gallocr_is_own(galloc, parent)) { + AT_PRINTF("not reusing parent %s for %s as %p is external\n", parent->name, node->name, parent->data); + continue; + } + + // outputs cannot be reused + if (parent->flags & GGML_TENSOR_FLAG_OUTPUT || (parent->view_src != NULL && parent->view_src->flags & GGML_TENSOR_FLAG_OUTPUT)) { + AT_PRINTF("not reusing parent %s for %s as it is an output\n", parent->name, node->name); + continue; + } + + if (!ggml_are_same_layout(node, parent)) { + AT_PRINTF("not reusing parent %s for %s as layouts are different\n", parent->name, node->name); + continue; + } + + struct hash_node * p_hn = ggml_gallocr_hash_get(galloc, parent); + if (p_hn->n_children == 1 && p_hn->n_views == 0) { + if (ggml_impl_is_view(parent)) { + struct ggml_tensor * view_src = parent->view_src; + struct hash_node * view_src_hn = ggml_gallocr_hash_get(galloc, view_src); + if (view_src_hn->n_views == 1 && view_src_hn->n_children == 0 && view_src->data == parent->data) { + AT_PRINTF("reusing view parent %s (%s) for %s\n", parent->name, view_src->name, node->name); + assert(view_src_hn->addr.chunk == p_hn->addr.chunk && view_src_hn->addr.offset == p_hn->addr.offset); + hn->buffer_id = p_hn->buffer_id; + hn->addr = p_hn->addr; + p_hn->allocated = false; // avoid freeing the parent + view_src_hn->allocated = false; + ggml_gallocr_free_extra_space(galloc, node, view_src); + return; + } + } else { + AT_PRINTF("reusing parent %s for %s\n", parent->name, node->name); + hn->buffer_id = p_hn->buffer_id; + hn->addr = p_hn->addr; + p_hn->allocated = false; // avoid freeing the parent + ggml_gallocr_free_extra_space(galloc, node, parent); + return; + } + } + } + } + // allocate tensor from the buffer + struct ggml_dyn_tallocr * alloc = galloc->buf_tallocs[buffer_id]; + ggml_backend_buffer_type_t buft = galloc->bufts[buffer_id]; + size_t size = ggml_backend_buft_get_alloc_size(buft, node); + hn->buffer_id = buffer_id; + hn->addr = ggml_dyn_tallocr_alloc(alloc, size, node); + } +} + +static void ggml_gallocr_free_node(ggml_gallocr_t galloc, struct ggml_tensor * node) { + // graph outputs are never freed + if (node->flags & GGML_TENSOR_FLAG_OUTPUT) { + AT_PRINTF("not freeing output %s\n", node->name); + return; + } + + struct hash_node * hn = ggml_gallocr_hash_get(galloc, node); + int buffer_id = hn->buffer_id; + struct ggml_dyn_tallocr * alloc = galloc->buf_tallocs[buffer_id]; + ggml_backend_buffer_type_t buft = galloc->bufts[buffer_id]; + size_t size = ggml_backend_buft_get_alloc_size(buft, node); + + AT_PRINTF("%s: freeing %s at {chunk=%d, offset=%zu} (%zu bytes) - n_free_blocks = %d\n", + __func__, node->name, hn->addr.chunk, hn->addr.offset, size, alloc->chunks[hn->addr.chunk]->n_free_blocks); +#ifdef GGML_ALLOCATOR_DEBUG + remove_allocated_tensor(alloc, hn->addr, node); +#endif + + ggml_dyn_tallocr_free_bytes(alloc, hn->addr, size); + hn->allocated = false; +} + +#define get_node_buffer_id(node_buffer_ids, i) ggml_get_node_buffer_id_rust(node_buffer_ids, i) + +static void ggml_gallocr_alloc_graph_impl(ggml_gallocr_t galloc, struct ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids) { + // clear hash tables + ggml_hash_set_reset(&galloc->hash_set); + memset(galloc->hash_values, 0, sizeof(struct hash_node) * galloc->hash_set.size); + + // allocate leafs + // these may be tensors that the application is not using in the graph, but may still want to allocate for other purposes + for (int i = 0; i < graph->n_leafs; i++) { + struct ggml_tensor * leaf = graph->leafs[i]; + ggml_gallocr_allocate_node(galloc, leaf, get_node_buffer_id(leaf_buffer_ids, i)); + } + + // count number of children and views + // allocate other graph inputs and leafs first to avoid overwriting them + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + + // TODO: better way to add external dependencies + // GGML_OP_NONE does not appear normally in the graph nodes, but is used by ggml-backend to add dependencies to + // control when some tensors are allocated and freed. in this case, the dependencies are in `src`, but the node + // itself is never used and should not be considered a dependency + if (ggml_impl_is_view(node) && node->op != GGML_OP_NONE) { + struct ggml_tensor * view_src = node->view_src; + ggml_gallocr_hash_get(galloc, view_src)->n_views += 1; + } + + if (node->flags & GGML_TENSOR_FLAG_INPUT) { + ggml_gallocr_allocate_node(galloc, graph->nodes[i], get_node_buffer_id(node_buffer_ids, i)); + } + + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (src == NULL) { + continue; + } + + ggml_gallocr_hash_get(galloc, src)->n_children += 1; + + // allocate explicit inputs + if (src->flags & GGML_TENSOR_FLAG_INPUT) { + ggml_gallocr_allocate_node(galloc, src, get_node_buffer_id(node_buffer_ids, i)); + } + } + } + + // allocate tensors + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + int buffer_id = get_node_buffer_id(node_buffer_ids, i); + + // allocate parents (only leafs need to be allocated at this point) + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * parent = node->src[j]; + if (parent == NULL) { + continue; + } + ggml_gallocr_allocate_node(galloc, parent, buffer_id); + } + + // allocate node + ggml_gallocr_allocate_node(galloc, node, buffer_id); + + AT_PRINTF("exec: %s (%s) <= ", ggml_op_desc(node), node->name); + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * parent = node->src[j]; + if (parent == NULL) { + continue; + } + AT_PRINTF("%s", parent->name); + if (j < GGML_MAX_SRC - 1 && node->src[j + 1] != NULL) { + AT_PRINTF(", "); + } + } + AT_PRINTF("\n"); + + // update parents + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * parent = node->src[j]; + if (parent == NULL) { + continue; + } + struct hash_node * p_hn = ggml_gallocr_hash_get(galloc, parent); + p_hn->n_children -= 1; + + AT_PRINTF("parent %s: %d children, %d views, allocated: %d\n", + parent->name, p_hn->n_children, p_hn->n_views, p_hn->allocated); + + if (p_hn->n_children == 0 && p_hn->n_views == 0) { + if (ggml_impl_is_view(parent)) { + struct ggml_tensor * view_src = parent->view_src; + struct hash_node * view_src_hn = ggml_gallocr_hash_get(galloc, view_src); + view_src_hn->n_views -= 1; + AT_PRINTF("view_src %s: %d children, %d views\n", + view_src->name, view_src_hn->n_children, view_src_hn->n_views); + if (view_src_hn->n_views == 0 && view_src_hn->n_children == 0 && view_src_hn->allocated) { + ggml_gallocr_free_node(galloc, view_src); + } + } + else if (p_hn->allocated) { + ggml_gallocr_free_node(galloc, parent); + } + } + AT_PRINTF("\n"); + } + } +} + +static bool ggml_gallocr_reserve_n_impl( + ggml_gallocr_t galloc, struct ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids, bool no_alloc) { + size_t min_hash_size = graph->n_nodes + graph->n_leafs; + // add 25% margin to avoid hash collisions + min_hash_size += min_hash_size / 4; + + // initialize hash table + if (galloc->hash_set.size < min_hash_size) { + ggml_hash_set_free(&galloc->hash_set); + galloc->hash_set = ggml_hash_set_new(min_hash_size); + GGML_ASSERT(galloc->hash_set.keys != NULL); + + free(galloc->hash_values); + galloc->hash_values = malloc(sizeof(struct hash_node) * galloc->hash_set.size); + GGML_ASSERT(galloc->hash_values != NULL); + } + + // reset allocators + for (int i = 0; i < galloc->n_buffers; i++) { + ggml_dyn_tallocr_reset(galloc->buf_tallocs[i]); + } + + // allocate in hash table + ggml_gallocr_alloc_graph_impl(galloc, graph, node_buffer_ids, leaf_buffer_ids); + + // set the node_allocs from the hash table + if (galloc->n_nodes < graph->n_nodes) { + free(galloc->node_allocs); + galloc->node_allocs = calloc(graph->n_nodes, sizeof(struct node_alloc)); + GGML_ASSERT(galloc->node_allocs != NULL); + } + galloc->n_nodes = graph->n_nodes; + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + struct node_alloc * node_alloc = &galloc->node_allocs[i]; + if (node->view_src || node->data) { + node_alloc->dst.buffer_id = -1; + node_alloc->dst.addr = GGML_BUFFER_ADDRESS_INVALID; + node_alloc->dst.size_max = 0; + } else { + struct hash_node * hn = ggml_gallocr_hash_get(galloc, node); + node_alloc->dst.buffer_id = hn->buffer_id; + node_alloc->dst.addr = hn->addr; + node_alloc->dst.size_max = ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], node); + } + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (!src || src->view_src || src->data) { + node_alloc->src[j].buffer_id = -1; + node_alloc->src[j].addr = GGML_BUFFER_ADDRESS_INVALID; + node_alloc->src[j].size_max = 0; + } else { + struct hash_node * hn = ggml_gallocr_hash_get(galloc, src); + node_alloc->src[j].buffer_id = hn->buffer_id; + node_alloc->src[j].addr = hn->addr; + node_alloc->src[j].size_max = ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], src); + } + } + } + if (galloc->n_leafs < graph->n_leafs) { + free(galloc->leaf_allocs); + galloc->leaf_allocs = calloc(graph->n_leafs, sizeof(galloc->leaf_allocs[0])); + GGML_ASSERT(galloc->leaf_allocs != NULL); + } + galloc->n_leafs = graph->n_leafs; + for (int i = 0; i < graph->n_leafs; i++) { + struct ggml_tensor * leaf = graph->leafs[i]; + struct hash_node * hn = ggml_gallocr_hash_get(galloc, leaf); + if (leaf->view_src || leaf->data) { + galloc->leaf_allocs[i].leaf.buffer_id = -1; + galloc->leaf_allocs[i].leaf.addr = GGML_BUFFER_ADDRESS_INVALID; + galloc->leaf_allocs[i].leaf.size_max = 0; + } else { + galloc->leaf_allocs[i].leaf.buffer_id = hn->buffer_id; + galloc->leaf_allocs[i].leaf.addr = hn->addr; + galloc->leaf_allocs[i].leaf.size_max = ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], leaf); + } + } + + // reallocate buffers if needed + for (int i = 0; i < galloc->n_buffers; i++) { + // if the buffer type is used multiple times, we reuse the same buffer + for (int j = 0; j < i; j++) { + if (galloc->buf_tallocs[j] == galloc->buf_tallocs[i]) { + galloc->buffers[i] = galloc->buffers[j]; + break; + } + } + + // even if there are no tensors allocated in this buffer, we still need to allocate it to initialize views + bool realloc = galloc->buffers[i] == NULL; + size_t new_size = 0; + for (int c = 0; c < galloc->buf_tallocs[i]->n_chunks; c++) { + size_t cur_chunk_size = galloc->buffers[i] ? ggml_vbuffer_chunk_size(galloc->buffers[i], c) : 0; + size_t new_chunk_size = ggml_dyn_tallocr_max_size(galloc->buf_tallocs[i], c); + new_size += new_chunk_size; + if (new_chunk_size > cur_chunk_size) { + realloc = true; + } + } + if (realloc) { +#ifndef NDEBUG + { + size_t cur_size = galloc->buffers[i] ? ggml_vbuffer_size(galloc->buffers[i]) : 0; + if (cur_size > 0) { + GGML_LOG_DEBUG("%s: reallocating %s buffer from size %.02f MiB to %.02f MiB\n", + __func__, ggml_backend_buft_name(galloc->bufts[i]), cur_size / 1024.0 / 1024.0, new_size / 1024.0 / 1024.0); + } + } +#endif + ggml_vbuffer_free(galloc->buffers[i]); + if (no_alloc) { + galloc->buffers[i] = NULL; + } else { + galloc->buffers[i] = ggml_vbuffer_alloc(galloc->bufts[i], galloc->buf_tallocs[i], GGML_BACKEND_BUFFER_USAGE_COMPUTE); + if (galloc->buffers[i] == NULL) { + GGML_LOG_ERROR("%s: failed to allocate %s buffer of size %zu\n", __func__, ggml_backend_buft_name(galloc->bufts[i]), new_size); + return false; + } + } + } + } + + return true; +} + +void ggml_gallocr_reserve_n_size( + ggml_gallocr_t galloc, struct ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids, size_t * sizes) { + GGML_ASSERT(ggml_gallocr_reserve_n_impl(galloc, graph, node_buffer_ids, leaf_buffer_ids, /*no_alloc =*/ true)); + for (int i = 0; i < galloc->n_buffers; i++) { + sizes[i] = 0; + for (int c = 0; c < galloc->buf_tallocs[i]->n_chunks; c++) { + sizes[i] += galloc->buf_tallocs[i]->chunks[c]->max_size; + } + } +} + +bool ggml_gallocr_reserve_n(ggml_gallocr_t galloc, struct ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids) { + return ggml_gallocr_reserve_n_impl(galloc, graph, node_buffer_ids, leaf_buffer_ids, /*no_alloc =*/ false); +} + +bool ggml_gallocr_reserve(ggml_gallocr_t galloc, struct ggml_cgraph *graph) { + return ggml_gallocr_reserve_n(galloc, graph, NULL, NULL); +} + +static void ggml_gallocr_init_tensor(ggml_gallocr_t galloc, struct ggml_tensor * tensor, struct tensor_alloc * tensor_alloc) { + int buffer_id = tensor_alloc->buffer_id; + assert(tensor->data || tensor->view_src || ggml_backend_buft_get_alloc_size(galloc->bufts[buffer_id], tensor) <= tensor_alloc->size_max); + + if (tensor->view_src != NULL) { + if (tensor->buffer == NULL) { + assert(tensor_alloc->addr.offset == SIZE_MAX); + if (tensor->view_src->buffer == NULL) { + // this tensor was allocated without ggml-backend + return; + } + ggml_backend_view_init(tensor); + } + } else { + if (tensor->data == NULL) { + assert(tensor_alloc->addr.offset != SIZE_MAX); + assert(ggml_backend_buft_get_alloc_size(galloc->bufts[buffer_id], tensor) <= tensor_alloc->size_max); + ggml_vbuffer_tensor_alloc(galloc->buffers[buffer_id], tensor, tensor_alloc->addr); + } else { + if (tensor->buffer == NULL) { + // this tensor was allocated without ggml-backend + return; + } + } + } +} + +static bool ggml_gallocr_node_needs_realloc(ggml_gallocr_t galloc, struct ggml_tensor * node, struct tensor_alloc * talloc) { + size_t node_size = 0; + if (!node->data && !node->view_src) { + // If we previously had data but don't now then reallocate + if (talloc->buffer_id < 0) { + return false; + } + node_size = ggml_backend_buft_get_alloc_size(galloc->bufts[talloc->buffer_id], node); + } + return talloc->size_max >= node_size; +} + +static bool ggml_gallocr_needs_realloc(ggml_gallocr_t galloc, struct ggml_cgraph * graph) { + if (galloc->n_nodes != graph->n_nodes) { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: graph has different number of nodes\n", __func__); +#endif + return true; + } + + if (galloc->n_leafs != graph->n_leafs) { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: graph has different number of leafs\n", __func__); +#endif + return true; + } + + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + struct node_alloc * node_alloc = &galloc->node_allocs[i]; + + if (!ggml_gallocr_node_needs_realloc(galloc, node, &node_alloc->dst)) { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: node %s is not valid\n", __func__, node->name); +#endif + return true; + } + + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (src == NULL) { + continue; + } + if (!ggml_gallocr_node_needs_realloc(galloc, src, &node_alloc->src[j])) { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: src %d (%s) of node %s is not valid\n", __func__, j, src->name, node->name); +#endif + return true; + } + } + } + + return false; +} + +bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph * graph) { + if (ggml_gallocr_needs_realloc(galloc, graph)) { + if (galloc->n_buffers == 1) { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: reallocating buffers automatically\n", __func__); +#endif + if (!ggml_gallocr_reserve(galloc, graph)) { + return false; + } + } else { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: cannot reallocate multi buffer graph automatically, call reserve\n", __func__); +#endif + return false; + } + } + + // reset buffers + for (int i = 0; i < galloc->n_buffers; i++) { + if (galloc->buffers[i] != NULL) { + ggml_vbuffer_reset(galloc->buffers[i]); + } + } + + // allocate the graph tensors from the previous assignments + // leafs + for (int i = 0; i < graph->n_leafs; i++) { + struct ggml_tensor * leaf = graph->leafs[i]; + struct leaf_alloc * leaf_alloc = &galloc->leaf_allocs[i]; + ggml_gallocr_init_tensor(galloc, leaf, &leaf_alloc->leaf); + } + // nodes + for (int i = 0; i < graph->n_nodes; i++) { + struct ggml_tensor * node = graph->nodes[i]; + struct node_alloc * node_alloc = &galloc->node_allocs[i]; + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (src == NULL) { + continue; + } + ggml_gallocr_init_tensor(galloc, src, &node_alloc->src[j]); + } + ggml_gallocr_init_tensor(galloc, node, &node_alloc->dst); + } + + return true; +} + +size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id) { + GGML_ASSERT(buffer_id >= 0 && buffer_id < galloc->n_buffers); + + if (galloc->buffers[buffer_id] == NULL) { + return 0; + } + + for (int i = 0; i < buffer_id; i++) { + if (galloc->buffers[i] == galloc->buffers[buffer_id]) { + // this buffer is the same as a previous one due to the same buffer type being used multiple times + // only return the buffer size the first time it appears to avoid double counting + return 0; + } + } + + return ggml_vbuffer_size(galloc->buffers[buffer_id]); +} + +// utils + +static void free_buffers(ggml_backend_buffer_t ** buffers, const size_t * n_buffers) { + for (size_t i = 0; i < *n_buffers; i++) { + ggml_backend_buffer_free((*buffers)[i]); + } + free(*buffers); +} + +static bool alloc_tensor_range(struct ggml_context * ctx, + struct ggml_tensor * first, struct ggml_tensor * last, + ggml_backend_buffer_type_t buft, size_t size, + ggml_backend_buffer_t ** buffers, size_t * n_buffers) { + + ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, size); + if (buffer == NULL) { + GGML_LOG_ERROR("%s: failed to allocate %s buffer of size %zu\n", __func__, ggml_backend_buft_name(buft), size); + free_buffers(buffers, n_buffers); + return false; + } + + *buffers = realloc(*buffers, sizeof(ggml_backend_buffer_t) * (*n_buffers + 1)); + (*buffers)[(*n_buffers)++] = buffer; + + struct ggml_tallocr tallocr = ggml_tallocr_new(buffer); + + for (struct ggml_tensor * t = first; t != last; t = ggml_get_next_tensor(ctx, t)) { + enum ggml_status status = GGML_STATUS_SUCCESS; + if (t->data == NULL) { + if (t->view_src == NULL) { + status = ggml_tallocr_alloc(&tallocr, t); + } else if (t->buffer == NULL) { + status = ggml_backend_view_init(t); + } + } else { + if (t->view_src != NULL && t->buffer == NULL) { + // view of a pre-allocated tensor + status = ggml_backend_view_init(t); + } + } + if (status != GGML_STATUS_SUCCESS) { + GGML_LOG_ERROR("%s: failed to initialize tensor %s\n", __func__, t->name); + free_buffers(buffers, n_buffers); + return false; + } + } + + return true; +} + +static ggml_backend_buffer_t ggml_backend_alloc_ctx_tensors_from_buft_impl( + struct ggml_context * ctx, ggml_backend_buffer_type_t buft, size_t * nbytes_total, bool no_alloc) { + GGML_ASSERT(ggml_get_no_alloc(ctx) == true); + + size_t alignment = ggml_backend_buft_get_alignment(buft); + size_t max_size = ggml_backend_buft_get_max_size(buft); + + ggml_backend_buffer_t * buffers = NULL; + size_t n_buffers = 0; + *nbytes_total = 0; + + size_t cur_buf_size = 0; + struct ggml_tensor * first = ggml_get_first_tensor(ctx); + for (struct ggml_tensor * t = first; t != NULL; t = ggml_get_next_tensor(ctx, t)) { + size_t this_size = 0; + if (t->data == NULL && t->view_src == NULL) { + this_size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, t), alignment); + } + + if (cur_buf_size > 0 && (cur_buf_size + this_size) > max_size) { + // allocate tensors in the current buffer + if (!no_alloc && !alloc_tensor_range(ctx, first, t, buft, cur_buf_size, &buffers, &n_buffers)) { + return NULL; + } + first = t; + *nbytes_total += cur_buf_size; + cur_buf_size = this_size; + } else { + cur_buf_size += this_size; + } + } + + // allocate remaining tensors + if (cur_buf_size > 0) { + *nbytes_total += cur_buf_size; + if (!no_alloc && !alloc_tensor_range(ctx, first, NULL, buft, cur_buf_size, &buffers, &n_buffers)) { + return NULL; + } + } + + if (no_alloc) { + return NULL; + } + + if (n_buffers == 0) { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: all tensors in the context are already allocated\n", __func__); +#endif + GGML_ASSERT(!buffers); + return NULL; + } + + ggml_backend_buffer_t buffer; + if (n_buffers == 1) { + buffer = buffers[0]; + } else { + buffer = ggml_backend_multi_buffer_alloc_buffer(buffers, n_buffers); + } + if (buffers) { + free(buffers); // can be NULL if context is empty or no_alloc + } + return buffer; +} + +size_t ggml_backend_alloc_ctx_tensors_from_buft_size(struct ggml_context * ctx, ggml_backend_buffer_type_t buft) { + size_t nbytes_total = 0; + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft_impl(ctx, buft, &nbytes_total, /*no_alloc=*/ true); + GGML_ASSERT(!buf); + return nbytes_total; +} + +ggml_backend_buffer_t ggml_backend_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft) { + size_t nbytes_total = 0; + if (ggml_backend_buft_is_meta(buft)) { + return ggml_backend_meta_alloc_ctx_tensors_from_buft(ctx, buft); + } + return ggml_backend_alloc_ctx_tensors_from_buft_impl(ctx, buft, &nbytes_total, /*no_alloc =*/ false); +} + +ggml_backend_buffer_t ggml_backend_alloc_ctx_tensors(struct ggml_context * ctx, ggml_backend_t backend) { + return ggml_backend_alloc_ctx_tensors_from_buft(ctx, ggml_backend_get_default_buffer_type(backend)); +} + + +/* ---- quantization implementation ---- */ + +#define GGML_COMMON_IMPL_C +#include "ggml-common-defs.inc" + +#include "ggml-quants-defs.inc" +#include "ggml-impl-defs.inc" +#include "ggml-cpu/ggml-cpu-impl-defs.inc" +#include "ggml-cpu.h.inc" + +#include +#include +#include +#include +#include // for qsort +#include // for GGML_ASSERT + +#define GROUP_MAX_EPS 1e-15f +#define GROUP_MAX_EPS_IQ3_XXS 1e-8f +#define GROUP_MAX_EPS_IQ2_S 1e-8f +#define GROUP_MAX_EPS_IQ1_M 1e-7f +#define GROUP_MAX_EPS_IQ1_S 1e-12f + + +static inline int best_index_int8(int n, const int8_t * 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; +} + +// reference implementation for deterministic creation of model files +void quantize_row_q1_0_ref(const float * GGML_RESTRICT x, block_q1_0 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK1_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float sum_abs = 0.0f; + for (int j = 0; j < qk; j++) { + sum_abs += fabsf(x[i*qk + j]); + } + const float d = sum_abs / qk; + + y[i].d = GGML_FP32_TO_FP16(d); + + // Clear all bits first + for (int j = 0; j < qk / 8; ++j) { + y[i].qs[j] = 0; + } + + // Just store sign of each weight directly (no normalization) + for (int j = 0; j < qk; ++j) { + const int bit_index = j; + const int byte_index = bit_index / 8; + const int bit_offset = bit_index % 8; + + if (x[i*qk + j] >= 0.0f) { + y[i].qs[byte_index] |= (1 << bit_offset); + } + } + } +} + +// reference implementation for deterministic creation of model files +void quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK4_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + float max = 0.0f; + + for (int j = 0; j < qk; j++) { + const float v = x[i*qk + j]; + if (amax < fabsf(v)) { + amax = fabsf(v); + max = v; + } + } + + const float d = max / -8; + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + for (int j = 0; j < qk/2; ++j) { + const float x0 = x[i*qk + 0 + j]*id; + const float x1 = x[i*qk + qk/2 + j]*id; + + const uint8_t xi0 = MIN(15, (int8_t)(x0 + 8.5f)); + const uint8_t xi1 = MIN(15, (int8_t)(x1 + 8.5f)); + + y[i].qs[j] = xi0; + y[i].qs[j] |= xi1 << 4; + } + } +} + +void quantize_row_q4_1_ref(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t k) { + const int qk = QK4_1; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float min = FLT_MAX; + float max = -FLT_MAX; + + for (int j = 0; j < qk; j++) { + const float v = x[i*qk + j]; + + if (v < min) min = v; + if (v > max) max = v; + } + + const float d = (max - min) / ((1 << 4) - 1); + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + y[i].m = GGML_FP32_TO_FP16(min); + + for (int j = 0; j < qk/2; ++j) { + const float x0 = (x[i*qk + 0 + j] - min)*id; + const float x1 = (x[i*qk + qk/2 + j] - min)*id; + + const uint8_t xi0 = MIN(15, (int8_t)(x0 + 0.5f)); + const uint8_t xi1 = MIN(15, (int8_t)(x1 + 0.5f)); + + y[i].qs[j] = xi0; + y[i].qs[j] |= xi1 << 4; + } + } +} + +void quantize_row_q5_0_ref(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK5_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + float max = 0.0f; + + for (int j = 0; j < qk; j++) { + const float v = x[i*qk + j]; + if (amax < fabsf(v)) { + amax = fabsf(v); + max = v; + } + } + + const float d = max / -16; + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + uint32_t qh = 0; + + for (int j = 0; j < qk/2; ++j) { + const float x0 = x[i*qk + 0 + j]*id; + const float x1 = x[i*qk + qk/2 + j]*id; + + const uint8_t xi0 = MIN(31, (int8_t)(x0 + 16.5f)); + const uint8_t xi1 = MIN(31, (int8_t)(x1 + 16.5f)); + + y[i].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4); + + // get the 5-th bit and store it in qh at the right position + qh |= ((xi0 & 0x10u) >> 4) << (j + 0); + qh |= ((xi1 & 0x10u) >> 4) << (j + qk/2); + } + + memcpy(&y[i].qh, &qh, sizeof(qh)); + } +} + +void quantize_row_q5_1_ref(const float * GGML_RESTRICT x, block_q5_1 * GGML_RESTRICT y, int64_t k) { + const int qk = QK5_1; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float min = FLT_MAX; + float max = -FLT_MAX; + + for (int j = 0; j < qk; j++) { + const float v = x[i*qk + j]; + + if (v < min) min = v; + if (v > max) max = v; + } + + const float d = (max - min) / ((1 << 5) - 1); + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + y[i].m = GGML_FP32_TO_FP16(min); + + uint32_t qh = 0; + + for (int j = 0; j < qk/2; ++j) { + const float x0 = (x[i*qk + 0 + j] - min)*id; + const float x1 = (x[i*qk + qk/2 + j] - min)*id; + + const uint8_t xi0 = (uint8_t)(x0 + 0.5f); + const uint8_t xi1 = (uint8_t)(x1 + 0.5f); + + y[i].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4); + + // get the 5-th bit and store it in qh at the right position + qh |= ((xi0 & 0x10u) >> 4) << (j + 0); + qh |= ((xi1 & 0x10u) >> 4) << (j + qk/2); + } + + memcpy(&y[i].qh, &qh, sizeof(y[i].qh)); + } +} + +// reference implementation for deterministic creation of model files +void quantize_row_q8_0_ref(const float * GGML_RESTRICT x, block_q8_0 * GGML_RESTRICT y, int64_t k) { + assert(k % QK8_0 == 0); + const int nb = k / QK8_0; + + for (int i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < QK8_0; j++) { + const float v = x[i*QK8_0 + j]; + amax = MAX(amax, fabsf(v)); + } + + const float d = amax / ((1 << 7) - 1); + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + for (int j = 0; j < QK8_0; ++j) { + const float x0 = x[i*QK8_0 + j]*id; + + y[i].qs[j] = roundf(x0); + } + } +} + +// reference implementation for deterministic creation of model files +void quantize_row_q8_1_ref(const float * GGML_RESTRICT x, block_q8_1 * GGML_RESTRICT y, int64_t k) { + assert(QK8_1 == 32); + assert(k % QK8_1 == 0); + const int nb = k / QK8_1; + + for (int i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < QK8_1; j++) { + const float v = x[i*QK8_1 + j]; + amax = MAX(amax, fabsf(v)); + } + + const float d = amax / ((1 << 7) - 1); + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + int sum = 0; + + for (int j = 0; j < QK8_1/2; ++j) { + const float v0 = x[i*QK8_1 + j]*id; + const float v1 = x[i*QK8_1 + QK8_1/2 + j]*id; + + y[i].qs[ j] = roundf(v0); + y[i].qs[QK8_1/2 + j] = roundf(v1); + + sum += y[i].qs[ j]; + sum += y[i].qs[QK8_1/2 + j]; + } + + y[i].s = GGML_FP32_TO_FP16(sum*d); + } +} + +static inline int best_index_mxfp4(float x, float e) { + int best_index = 0; + float best_err = fabsf(kvalues_mxfp4[0]*e - x); + for (int i = 1; i < 16; i++) { + float err = fabsf(kvalues_mxfp4[i]*e - x); + if (err < best_err) { + best_index = i; + best_err = err; + } + } + return best_index; +} + +void quantize_row_mxfp4_ref(const float * GGML_RESTRICT x, block_mxfp4 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK_MXFP4; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < qk; j++) { + const float v = x[i*qk + j]; + + if (amax < fabsf(v)) { + amax = fabsf(v); + } + } + + const uint8_t e = amax > 0.0f ? (uint8_t) (floorf(log2f(amax)) - 2 + 127) : 0; + + const float d = GGML_E8M0_TO_FP32_HALF(e); + + y[i].e = e; + + for (int j = 0; j < qk/2; ++j) { + const uint8_t x0 = best_index_mxfp4(x[i*qk + 0 + j], d); + const uint8_t x1 = best_index_mxfp4(x[i*qk + qk/2 + j], d); + + y[i].qs[j] = x0; + y[i].qs[j] |= x1 << 4; + } + } +} + +void quantize_row_nvfp4_ref(const float * GGML_RESTRICT x, block_nvfp4 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK_NVFP4; + static const int qk_sub = QK_NVFP4_SUB; + static const int n_sub = QK_NVFP4 / QK_NVFP4_SUB; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + for (int s = 0; s < n_sub; s++) { + const float * xb = x + i*qk + s*qk_sub; + + float amax = 0.0f; + for (int j = 0; j < qk_sub; j++) { + if (amax < fabsf(xb[j])) { + amax = fabsf(xb[j]); + } + } + + // UE4M3 scale: amax / 6.0 maps the max E2M1 value (6.0) to amax + const uint8_t ue = ggml_fp32_to_ue4m3(amax / 6.0f); + y[i].d[s] = ue; + const float d = ggml_ue4m3_to_fp32(ue); + + for (int j = 0; j < qk_sub/2; ++j) { + const uint8_t x0 = best_index_mxfp4(xb[0 + j], d); + const uint8_t x1 = best_index_mxfp4(xb[qk_sub/2 + j], d); + + y[i].qs[s*(qk_sub/2) + j] = x0 | (x1 << 4); + } + } + } +} + +void dequantize_row_q1_0(const block_q1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK1_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + const float neg_d = -d; + + for (int j = 0; j < qk; ++j) { + const int byte_index = j / 8; + const int bit_offset = j % 8; + const uint8_t bit = (x[i].qs[byte_index] >> bit_offset) & 1; + y[i*qk + j] = bit ? d : neg_d; + } + } +} + +void dequantize_row_q4_0(const block_q4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK4_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (int j = 0; j < qk/2; ++j) { + const int x0 = (x[i].qs[j] & 0x0F) - 8; + const int x1 = (x[i].qs[j] >> 4) - 8; + + y[i*qk + j + 0 ] = x0*d; + y[i*qk + j + qk/2] = x1*d; + } + } +} + +void dequantize_row_q4_1(const block_q4_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK4_1; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + const float m = GGML_FP16_TO_FP32(x[i].m); + + for (int j = 0; j < qk/2; ++j) { + const int x0 = (x[i].qs[j] & 0x0F); + const int x1 = (x[i].qs[j] >> 4); + + y[i*qk + j + 0 ] = x0*d + m; + y[i*qk + j + qk/2] = x1*d + m; + } + } +} + +void dequantize_row_q5_0(const block_q5_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK5_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + + uint32_t qh; + memcpy(&qh, x[i].qh, sizeof(qh)); + + for (int j = 0; j < qk/2; ++j) { + const uint8_t xh_0 = ((qh >> (j + 0)) << 4) & 0x10; + const uint8_t xh_1 = ((qh >> (j + 12)) ) & 0x10; + + const int32_t x0 = ((x[i].qs[j] & 0x0F) | xh_0) - 16; + const int32_t x1 = ((x[i].qs[j] >> 4) | xh_1) - 16; + + y[i*qk + j + 0 ] = x0*d; + y[i*qk + j + qk/2] = x1*d; + } + } +} + +void dequantize_row_q5_1(const block_q5_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK5_1; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + const float m = GGML_FP16_TO_FP32(x[i].m); + + uint32_t qh; + memcpy(&qh, x[i].qh, sizeof(qh)); + + for (int j = 0; j < qk/2; ++j) { + const uint8_t xh_0 = ((qh >> (j + 0)) << 4) & 0x10; + const uint8_t xh_1 = ((qh >> (j + 12)) ) & 0x10; + + const int x0 = (x[i].qs[j] & 0x0F) | xh_0; + const int x1 = (x[i].qs[j] >> 4) | xh_1; + + y[i*qk + j + 0 ] = x0*d + m; + y[i*qk + j + qk/2] = x1*d + m; + } + } +} + +void dequantize_row_q8_0(const block_q8_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK8_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (int j = 0; j < qk; ++j) { + y[i*qk + j] = x[i].qs[j]*d; + } + } +} + +void dequantize_row_mxfp4(const block_mxfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK_MXFP4; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_E8M0_TO_FP32_HALF(x[i].e); + + for (int j = 0; j < qk/2; ++j) { + const int8_t x0 = kvalues_mxfp4[x[i].qs[j] & 0x0F]; + const int8_t x1 = kvalues_mxfp4[x[i].qs[j] >> 4]; + + y[i*qk + j + 0 ] = x0*d; + y[i*qk + j + qk/2] = x1*d; + } + } +} + +void dequantize_row_nvfp4(const block_nvfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK_NVFP4; + static const int qk_sub = QK_NVFP4_SUB; + static const int n_sub = QK_NVFP4 / QK_NVFP4_SUB; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + for (int s = 0; s < n_sub; s++) { + const float d = ggml_ue4m3_to_fp32(x[i].d[s]); + float * yb = y + i*qk + s*qk_sub; + + for (int j = 0; j < qk_sub/2; ++j) { + const int8_t v0 = kvalues_mxfp4[x[i].qs[s*(qk_sub/2) + j] & 0x0F]; + const int8_t v1 = kvalues_mxfp4[x[i].qs[s*(qk_sub/2) + j] >> 4]; + + yb[j + 0 ] = v0*d; + yb[j + qk_sub/2] = v1*d; + } + } + } +} + +// +// 2-6 bit quantization in super-blocks +// + +// +// ===================== Helper functions +// +static inline int nearest_int(float fval) { + assert(fabsf(fval) <= 4194303.f); + float val = fval + 12582912.f; + int i; memcpy(&i, &val, sizeof(int)); + return (i & 0x007fffff) - 0x00400000; +} + +static float make_qx_quants(int n, int nmax, const float * GGML_RESTRICT x, int8_t * GGML_RESTRICT L, int rmse_type, + const float * GGML_RESTRICT 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 < GROUP_MAX_EPS) { // all zero + for (int i = 0; i < n; ++i) { + L[i] = 0; + } + return 0.f; + } + float iscale = -nmax / max; + if (rmse_type == 0) { + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + L[i] = nmax + MAX(-nmax, MIN(nmax-1, l)); + } + return 1/iscale; + } + bool return_early = false; + if (rmse_type < 0) { + rmse_type = -rmse_type; + return_early = true; + } + float sumlx = 0; + float suml2 = 0; +#ifdef HAVE_BUGGY_APPLE_LINKER + // use 'volatile' to prevent unroll and work around a bug in Apple ld64 1015.7 + for (volatile int i = 0; i < n; ++i) { +#else + for (int i = 0; i < n; ++i) { +#endif + int l = nearest_int(iscale * x[i]); + l = MAX(-nmax, MIN(nmax-1, l)); + L[i] = l + nmax; + float w = qw ? qw[i] : rmse_type == 1 ? x[i] * x[i] : rmse_type == 2 ? 1 : rmse_type == 3 ? fabsf(x[i]) : sqrtf(fabsf(x[i])); + sumlx += w*x[i]*l; + suml2 += w*l*l; + } + float scale = suml2 ? sumlx/suml2 : 0.0f; + if (return_early) return suml2 > 0 ? 0.5f*(scale + 1/iscale) : 1/iscale; + 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 = MAX(-nmax, MIN(nmax-1, l)); + float w = qw ? qw[i] : rmse_type == 1 ? x[i] * x[i] : rmse_type == 2 ? 1 : rmse_type == 3 ? fabsf(x[i]) : sqrtf(fabsf(x[i])); + sumlx += w*x[i]*l; + suml2 += w*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 + MAX(-nmax, MIN(nmax-1, l)); + } + scale = sumlx/suml2; best = scale*sumlx; + } + } + return scale; +} + +static float make_q3_quants(int n, int nmax, const float * GGML_RESTRICT x, int8_t * GGML_RESTRICT L, bool do_rmse) { + 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 < GROUP_MAX_EPS) { // all zero + for (int i = 0; i < n; ++i) { L[i] = 0; } + return 0.f; + } + float iscale = -nmax / max; + if (do_rmse) { + float sumlx = 0; + float suml2 = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + l = MAX(-nmax, MIN(nmax-1, l)); + L[i] = l; + float w = x[i]*x[i]; + sumlx += w*x[i]*l; + suml2 += w*l*l; + } + for (int itry = 0; itry < 5; ++itry) { + int n_changed = 0; + for (int i = 0; i < n; ++i) { + float w = x[i]*x[i]; + float slx = sumlx - w*x[i]*L[i]; + if (slx > 0) { + float sl2 = suml2 - w*L[i]*L[i]; + int new_l = nearest_int(x[i] * sl2 / slx); + new_l = MAX(-nmax, MIN(nmax-1, new_l)); + if (new_l != L[i]) { + slx += w*x[i]*new_l; + sl2 += w*new_l*new_l; + if (sl2 > 0 && slx*slx*suml2 > sumlx*sumlx*sl2) { + L[i] = new_l; sumlx = slx; suml2 = sl2; + ++n_changed; + } + } + } + } + if (!n_changed) { + break; + } + } + for (int i = 0; i < n; ++i) { + L[i] += nmax; + } + return suml2 > 0.0f ? sumlx / suml2 : 0.0f; + } + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + l = MAX(-nmax, MIN(nmax-1, l)); + L[i] = l + nmax; + } + return 1/iscale; +} + +static float make_qkx1_quants(int n, int nmax, const float * GGML_RESTRICT x, uint8_t * GGML_RESTRICT L, float * GGML_RESTRICT the_min, + int ntry, float alpha) { + float min = x[0]; + float max = x[0]; + for (int i = 1; i < n; ++i) { + if (x[i] < min) min = x[i]; + if (x[i] > max) max = x[i]; + } + if (max == min) { + for (int i = 0; i < n; ++i) L[i] = 0; + *the_min = 0; + return 0.f; + } + if (min > 0) min = 0; + float iscale = nmax/(max - min); + float scale = 1/iscale; + for (int itry = 0; itry < ntry; ++itry) { + float sumlx = 0; int suml2 = 0; + bool did_change = false; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale*(x[i] - min)); + l = MAX(0, MIN(nmax, l)); + if (l != L[i]) { + L[i] = l; + did_change = true; + } + sumlx += (x[i] - min)*l; + suml2 += l*l; + } + scale = sumlx/suml2; + float sum = 0; + for (int i = 0; i < n; ++i) { + sum += x[i] - scale*L[i]; + } + min = alpha*min + (1 - alpha)*sum/n; + if (min > 0) min = 0; + iscale = 1/scale; + if (!did_change) break; + } + *the_min = -min; + return scale; +} + +static float make_qkx2_quants(int n, int nmax, const float * GGML_RESTRICT x, const float * GGML_RESTRICT weights, + uint8_t * GGML_RESTRICT L, float * GGML_RESTRICT the_min, uint8_t * GGML_RESTRICT Laux, + float rmin, float rdelta, int nstep, bool use_mad) { + float min = x[0]; + float max = x[0]; + float sum_w = weights[0]; + float sum_x = sum_w * x[0]; +#ifdef HAVE_BUGGY_APPLE_LINKER + // use 'volatile' to prevent unroll and work around a bug in Apple ld64 1015.7 + for (volatile int i = 1; i < n; ++i) { +#else + for (int i = 1; i < n; ++i) { +#endif + if (x[i] < min) min = x[i]; + if (x[i] > max) max = x[i]; + float w = weights[i]; + sum_w += w; + sum_x += w * x[i]; + } + if (min > 0) min = 0; + if (max == min) { + for (int i = 0; i < n; ++i) L[i] = 0; + *the_min = -min; + return 0.f; + } + float iscale = nmax/(max - min); + float scale = 1/iscale; + float best_error = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale*(x[i] - min)); + L[i] = MAX(0, MIN(nmax, l)); + float diff = scale * L[i] + min - x[i]; + diff = use_mad ? fabsf(diff) : diff * diff; + float w = weights[i]; + best_error += w * diff; + } + if (nstep < 1) { + *the_min = -min; + return scale; + } + for (int is = 0; is <= nstep; ++is) { + iscale = (rmin + rdelta*is + nmax)/(max - min); + float sum_l = 0, sum_l2 = 0, sum_xl = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale*(x[i] - min)); + l = MAX(0, MIN(nmax, l)); + Laux[i] = l; + float w = weights[i]; + sum_l += w*l; + sum_l2 += w*l*l; + sum_xl += w*l*x[i]; + } + float D = sum_w * sum_l2 - sum_l * sum_l; + if (D > 0) { + float this_scale = (sum_w * sum_xl - sum_x * sum_l)/D; + float this_min = (sum_l2 * sum_x - sum_l * sum_xl)/D; + if (this_min > 0) { + this_min = 0; + this_scale = sum_xl / sum_l2; + } + float cur_error = 0; + for (int i = 0; i < n; ++i) { + float diff = this_scale * Laux[i] + this_min - x[i]; + diff = use_mad ? fabsf(diff) : diff * diff; + float w = weights[i]; + cur_error += w * diff; + } + if (cur_error < best_error) { + for (int i = 0; i < n; ++i) { + L[i] = Laux[i]; + } + best_error = cur_error; + scale = this_scale; + min = this_min; + } + } + } + *the_min = -min; + return scale; +} + +static inline void get_scale_min_k4(int j, const uint8_t * GGML_RESTRICT q, uint8_t * GGML_RESTRICT d, uint8_t * GGML_RESTRICT 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); + } +} + +//========================- 2-bit (de)-quantization + +void quantize_row_q2_K_ref(const float * GGML_RESTRICT x, block_q2_K * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + uint8_t L[QK_K]; + uint8_t Laux[16]; + float weights[16]; + float mins[QK_K/16]; + float scales[QK_K/16]; + + const float q4scale = 15.f; + + for (int i = 0; i < nb; i++) { + float max_scale = 0; // as we are deducting the min, scales are always positive + float max_min = 0; + for (int j = 0; j < QK_K/16; ++j) { + for (int l = 0; l < 16; ++l) weights[l] = fabsf(x[16*j + l]); + scales[j] = make_qkx2_quants(16, 3, x + 16*j, weights, L + 16*j, &mins[j], Laux, -0.5f, 0.1f, 15, true); + float scale = scales[j]; + if (scale > max_scale) { + max_scale = scale; + } + float min = mins[j]; + if (min > max_min) { + max_min = min; + } + } + + if (max_scale > 0) { + float iscale = q4scale/max_scale; + for (int j = 0; j < QK_K/16; ++j) { + int l = nearest_int(iscale*scales[j]); + y[i].scales[j] = l; + } + y[i].d = GGML_FP32_TO_FP16(max_scale/q4scale); + } else { + for (int j = 0; j < QK_K/16; ++j) y[i].scales[j] = 0; + y[i].d = GGML_FP32_TO_FP16(0.f); + } + if (max_min > 0) { + float iscale = q4scale/max_min; + for (int j = 0; j < QK_K/16; ++j) { + int l = nearest_int(iscale*mins[j]); + y[i].scales[j] |= (l << 4); + } + y[i].dmin = GGML_FP32_TO_FP16(max_min/q4scale); + } else { + y[i].dmin = GGML_FP32_TO_FP16(0.f); + } + for (int j = 0; j < QK_K/16; ++j) { + const float d = GGML_FP16_TO_FP32(y[i].d) * (y[i].scales[j] & 0xF); + if (!d) continue; + const float dm = GGML_FP16_TO_FP32(y[i].dmin) * (y[i].scales[j] >> 4); + for (int ii = 0; ii < 16; ++ii) { + int l = nearest_int((x[16*j + ii] + dm)/d); + l = MAX(0, MIN(3, l)); + L[16*j + ii] = l; + } + } + + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6); + } + } + + x += QK_K; + } +} + +void dequantize_row_q2_K(const block_q2_K * GGML_RESTRICT x, float * GGML_RESTRICT 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 float min = GGML_FP16_TO_FP32(x[i].dmin); + + const uint8_t * q = x[i].qs; + + int is = 0; + float dl, ml; + for (int n = 0; n < QK_K; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + + uint8_t sc = x[i].scales[is++]; + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *y++ = dl * ((int8_t)((q[l] >> shift) & 3)) - ml; + + sc = x[i].scales[is++]; + dl = d * (sc & 0xF); ml = min * (sc >> 4); + for (int l = 0; l < 16; ++l) *y++ = dl * ((int8_t)((q[l+16] >> shift) & 3)) - ml; + + shift += 2; + } + q += 32; + } + } +} + +static float make_qkx3_quants(int n, int nmax, const float * GGML_RESTRICT x, const float * GGML_RESTRICT weights, + uint8_t * GGML_RESTRICT L, float * GGML_RESTRICT the_min, uint8_t * GGML_RESTRICT Laux, + float rmin, float rdelta, int nstep, bool use_mad) { + float min = x[0]; + float max = x[0]; + float sum_w = weights ? weights[0] : x[0]*x[0]; + float sum_x = sum_w * x[0]; +#ifdef HAVE_BUGGY_APPLE_LINKER + // use 'volatile' to prevent unroll and work around a bug in Apple ld64 1015.7 + for (volatile int i = 1; i < n; ++i) { +#else + for (int i = 1; i < n; ++i) { +#endif + if (x[i] < min) min = x[i]; + if (x[i] > max) max = x[i]; + float w = weights ? weights[i] : x[i]*x[i]; + sum_w += w; + sum_x += w * x[i]; + } + if (min > 0) { + min = 0; + } + if (max <= min) { + memset(L, 0, n); + *the_min = -min; + return 0.f; + } + float iscale = nmax/(max - min); + float scale = 1/iscale; + float best_mad = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale*(x[i] - min)); + L[i] = MAX(0, MIN(nmax, l)); + float diff = scale * L[i] + min - x[i]; + diff = use_mad ? fabsf(diff) : diff*diff; + float w = weights ? weights[i] : x[i]*x[i]; + best_mad += w * diff; + } + if (nstep < 1) { + *the_min = -min; + return scale; + } + for (int is = 0; is <= nstep; ++is) { + iscale = (rmin + rdelta*is + nmax)/(max - min); + float sum_l = 0, sum_l2 = 0, sum_xl = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale*(x[i] - min)); + l = MAX(0, MIN(nmax, l)); + Laux[i] = l; + float w = weights ? weights[i] : x[i]*x[i]; + sum_l += w*l; + sum_l2 += w*l*l; + sum_xl += w*l*x[i]; + } + float D = sum_w * sum_l2 - sum_l * sum_l; + if (D > 0) { + float this_scale = (sum_w * sum_xl - sum_x * sum_l)/D; + float this_min = (sum_l2 * sum_x - sum_l * sum_xl)/D; + if (this_min > 0) { + this_min = 0; + this_scale = sum_xl / sum_l2; + } + float mad = 0; + for (int i = 0; i < n; ++i) { + float diff = this_scale * Laux[i] + this_min - x[i]; + diff = use_mad ? fabsf(diff) : diff*diff; + float w = weights ? weights[i] : x[i]*x[i]; + mad += w * diff; + } + if (mad < best_mad) { + for (int i = 0; i < n; ++i) { + L[i] = Laux[i]; + } + best_mad = mad; + scale = this_scale; + min = this_min; + } + } + } + *the_min = -min; + return scale; +} + +static float make_qp_quants(int n, int nmax, const float * GGML_RESTRICT x, uint8_t * GGML_RESTRICT L, const float * quant_weights) { + float max = 0; + for (int i = 0; i < n; ++i) { + max = MAX(max, x[i]); + } + if (max < GROUP_MAX_EPS) { // all zero + for (int i = 0; i < n; ++i) { L[i] = 0; } + return 0.f; + } + float iscale = nmax / max; + for (int i = 0; i < n; ++i) { + L[i] = nearest_int(iscale * x[i]); + } + float scale = 1/iscale; + float best_mse = 0; + for (int i = 0; i < n; ++i) { + float diff = x[i] - scale*L[i]; + float w = quant_weights[i]; + best_mse += w*diff*diff; + } + for (int is = -4; is <= 4; ++is) { + if (is == 0) continue; + float iscale_is = (0.1f*is + nmax)/max; + float scale_is = 1/iscale_is; + float mse = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale_is*x[i]); + l = MIN(nmax, l); + float diff = x[i] - scale_is*l; + float w = quant_weights[i]; + mse += w*diff*diff; + } + if (mse < best_mse) { + best_mse = mse; + iscale = iscale_is; + } + } + float sumlx = 0; + float suml2 = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + l = MIN(nmax, l); + L[i] = l; + float w = quant_weights[i]; + sumlx += w*x[i]*l; + suml2 += w*l*l; + } + for (int itry = 0; itry < 5; ++itry) { + int n_changed = 0; + for (int i = 0; i < n; ++i) { + float w = quant_weights[i]; + float slx = sumlx - w*x[i]*L[i]; + float sl2 = suml2 - w*L[i]*L[i]; + if (slx > 0 && sl2 > 0) { + int new_l = nearest_int(x[i] * sl2 / slx); + new_l = MIN(nmax, new_l); + if (new_l != L[i]) { + slx += w*x[i]*new_l; + sl2 += w*new_l*new_l; + if (slx*slx*suml2 > sumlx*sumlx*sl2) { + L[i] = new_l; sumlx = slx; suml2 = sl2; + ++n_changed; + } + } + } + } + if (!n_changed) { + break; + } + } + return suml2 > 0.0f ? sumlx / suml2 : 0.0f; +} + +static void quantize_row_q2_K_impl(const float * GGML_RESTRICT x, block_q2_K * GGML_RESTRICT y, int k, const float * GGML_RESTRICT quant_weights) { + GGML_ASSERT(quant_weights); + assert(k % QK_K == 0); + const int nb = k / QK_K; + const bool requantize = true; + + uint8_t L[QK_K]; + uint8_t Laux[16]; + float mins[QK_K/16]; + float scales[QK_K/16]; + float sw[QK_K/16]; + float weight[16]; + uint8_t Ls[QK_K/16], Lm[QK_K/16]; + + for (int i = 0; i < nb; i++) { + memset(sw, 0, QK_K/16*sizeof(float)); + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += x[j]*x[j]; + float sigma2 = sumx2/QK_K; + for (int j = 0; j < QK_K/16; ++j) { + const float * GGML_RESTRICT qw = quant_weights + QK_K * i + 16*j; + for (int l = 0; l < 16; ++l) weight[l] = qw[l] * sqrtf(sigma2 + x[16*j + l]*x[16*j + l]); + for (int l = 0; l < QK_K/16; ++l) sw[j] += weight[l]; + scales[j] = make_qkx3_quants(16, 3, x + 16*j, weight, L + 16*j, &mins[j], Laux, -0.9f, 0.05f, 36, false); + } + + float dm, mm; + dm = make_qp_quants(QK_K/16, 15, scales, Ls, sw); + mm = make_qp_quants(QK_K/16, 15, mins, Lm, sw); + + y[i].d = GGML_FP32_TO_FP16(dm); + y[i].dmin = GGML_FP32_TO_FP16(mm); + dm = GGML_FP16_TO_FP32(y[i].d); + mm = GGML_FP16_TO_FP32(y[i].dmin); + + for (int j = 0; j < QK_K/16; ++j) { + y[i].scales[j] = Ls[j] | (Lm[j] << 4); + } + + if (requantize) { + for (int j = 0; j < QK_K/16; ++j) { + const float d = dm * (y[i].scales[j] & 0xF); + if (!d) continue; + const float m = mm * (y[i].scales[j] >> 4); + for (int ii = 0; ii < 16; ++ii) { + int l = nearest_int((x[16*j + ii] + m)/d); + l = MAX(0, MIN(3, l)); + L[16*j + ii] = l; + } + } + } + + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6); + } + } + + x += QK_K; + } +} + +size_t quantize_q2_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + size_t row_size = ggml_row_size(GGML_TYPE_Q2_K, n_per_row); + if (!quant_weights) { + quantize_row_q2_K_ref(src, dst, (int64_t)nrow*n_per_row); + } + else { + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q2_K_impl(src, (block_q2_K*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + } + return nrow * row_size; +} + +//========================= 3-bit (de)-quantization + +void quantize_row_q3_K_ref(const float * GGML_RESTRICT x, block_q3_K * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + int8_t L[QK_K]; + float scales[QK_K / 16]; + + for (int i = 0; i < nb; i++) { + + float max_scale = 0; + float amax = 0; + for (int j = 0; j < QK_K/16; ++j) { + scales[j] = make_q3_quants(16, 4, x + 16*j, L + 16*j, true); + float scale = fabsf(scales[j]); + if (scale > amax) { + amax = scale; max_scale = scales[j]; + } + } + + memset(y[i].scales, 0, 12); + if (max_scale) { + float iscale = -32.f/max_scale; + for (int j = 0; j < QK_K/16; ++j) { + int8_t l = nearest_int(iscale*scales[j]); + l = MAX(-32, MIN(31, l)) + 32; + if (j < 8) { + y[i].scales[j] = l & 0xF; + } else { + y[i].scales[j-8] |= ((l & 0xF) << 4); + } + l >>= 4; + y[i].scales[j%4 + 8] |= (l << (2*(j/4))); + } + y[i].d = GGML_FP32_TO_FP16(1/iscale); + } else { + y[i].d = GGML_FP32_TO_FP16(0.f); + } + + int8_t sc; + for (int j = 0; j < QK_K/16; ++j) { + sc = j < 8 ? y[i].scales[j] & 0xF : y[i].scales[j-8] >> 4; + sc = (sc | (((y[i].scales[8 + j%4] >> (2*(j/4))) & 3) << 4)) - 32; + float d = GGML_FP16_TO_FP32(y[i].d) * sc; + if (!d) { + continue; + } + for (int ii = 0; ii < 16; ++ii) { + int l = nearest_int(x[16*j + ii]/d); + l = MAX(-4, MIN(3, l)); + L[16*j + ii] = l + 4; + } + } + + memset(y[i].hmask, 0, QK_K/8); + // We put the high-bit for the 1st 8 quants into bit 0, the next 8 into bit 1, etc. + int m = 0; + uint8_t hm = 1; + for (int j = 0; j < QK_K; ++j) { + if (L[j] > 3) { + y[i].hmask[m] |= hm; + L[j] -= 4; + } + if (++m == QK_K/8) { + m = 0; hm <<= 1; + } + } + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6); + } + } + + x += QK_K; + } +} + +void dequantize_row_q3_K(const block_q3_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + const uint32_t kmask1 = 0x03030303; + const uint32_t kmask2 = 0x0f0f0f0f; + + uint32_t aux[4]; + const int8_t * scales = (const int8_t*)aux; + + for (int i = 0; i < nb; i++) { + + const float d_all = GGML_FP16_TO_FP32(x[i].d); + + const uint8_t * GGML_RESTRICT q = x[i].qs; + const uint8_t * GGML_RESTRICT hm = x[i].hmask; + uint8_t m = 1; + + memcpy(aux, x[i].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); + + int is = 0; + float dl; + for (int n = 0; n < QK_K; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *y++ = dl * ((int8_t)((q[l+ 0] >> shift) & 3) - ((hm[l+ 0] & m) ? 0 : 4)); + } + + dl = d_all * (scales[is++] - 32); + for (int l = 0; l < 16; ++l) { + *y++ = dl * ((int8_t)((q[l+16] >> shift) & 3) - ((hm[l+16] & m) ? 0 : 4)); + } + + shift += 2; + m <<= 1; + } + q += 32; + } + + } +} + +static void quantize_row_q3_K_impl(const float * GGML_RESTRICT x, block_q3_K * GGML_RESTRICT y, int64_t n_per_row, const float * GGML_RESTRICT quant_weights) { + assert(n_per_row % QK_K == 0); + const int nb = n_per_row / QK_K; + + int8_t L[QK_K]; + float scales[QK_K / 16]; + float weight[16]; + float sw[QK_K / 16]; + int8_t Ls[QK_K / 16]; + + for (int i = 0; i < nb; i++) { + + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += x[j]*x[j]; + float sigma2 = 2*sumx2/QK_K; + + for (int j = 0; j < QK_K/16; ++j) { + if (quant_weights) { + const float * qw = quant_weights + QK_K * i + 16*j; + for (int l = 0; l < 16; ++l) weight[l] = qw[l] * sqrtf(sigma2 + x[16*j+l]*x[16*j+l]); + } else { + for (int l = 0; l < 16; ++l) weight[l] = x[16*j+l]*x[16*j+l]; + } + float sumw = 0; + for (int l = 0; l < 16; ++l) sumw += weight[l]; + sw[j] = sumw; + + scales[j] = make_qx_quants(16, 4, x + 16*j, L + 16*j, 1, weight); + + } + + memset(y[i].scales, 0, 12); + + float d_block = make_qx_quants(QK_K/16, 32, scales, Ls, 1, sw); + for (int j = 0; j < QK_K/16; ++j) { + int l = Ls[j]; + if (j < 8) { + y[i].scales[j] = l & 0xF; + } else { + y[i].scales[j-8] |= ((l & 0xF) << 4); + } + l >>= 4; + y[i].scales[j%4 + 8] |= (l << (2*(j/4))); + } + y[i].d = GGML_FP32_TO_FP16(d_block); + + int8_t sc; + for (int j = 0; j < QK_K/16; ++j) { + sc = j < 8 ? y[i].scales[j] & 0xF : y[i].scales[j-8] >> 4; + sc = (sc | (((y[i].scales[8 + j%4] >> (2*(j/4))) & 3) << 4)) - 32; + float d = GGML_FP16_TO_FP32(y[i].d) * sc; + if (!d) { + continue; + } + for (int ii = 0; ii < 16; ++ii) { + int l = nearest_int(x[16*j + ii]/d); + l = MAX(-4, MIN(3, l)); + L[16*j + ii] = l + 4; + } + } + + memset(y[i].hmask, 0, QK_K/8); + // We put the high-bit for the 1st 8 quants into bit 0, the next 8 into bit 1, etc. + int m = 0; + uint8_t hm = 1; + for (int j = 0; j < QK_K; ++j) { + if (L[j] > 3) { + y[i].hmask[m] |= hm; + L[j] -= 4; + } + if (++m == QK_K/8) { + m = 0; hm <<= 1; + } + } + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + y[i].qs[j/4 + l] = L[j + l] | (L[j + l + 32] << 2) | (L[j + l + 64] << 4) | (L[j + l + 96] << 6); + } + } + + x += QK_K; + } +} + +size_t quantize_q3_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + size_t row_size = ggml_row_size(GGML_TYPE_Q3_K, n_per_row); + if (!quant_weights) { + quantize_row_q3_K_ref(src, dst, (int64_t)nrow*n_per_row); + } + else { + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q3_K_impl(src, (block_q3_K*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + } + return nrow * row_size; +} + +// ====================== 4-bit (de)-quantization + +void quantize_row_q4_K_ref(const float * GGML_RESTRICT x, block_q4_K * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + uint8_t L[QK_K]; + uint8_t Laux[32]; + float weights[32]; + float mins[QK_K/32]; + float scales[QK_K/32]; + + for (int i = 0; i < nb; i++) { + float max_scale = 0; // as we are deducting the min, scales are always positive + float max_min = 0; + for (int j = 0; j < QK_K/32; ++j) { + //scales[j] = make_qkx1_quants(32, 15, x + 32*j, L + 32*j, &mins[j], 9, 0.5f); + float sum_x2 = 0; + for (int l = 0; l < 32; ++l) sum_x2 += x[32*j + l] * x[32*j + l]; + float av_x = sqrtf(sum_x2/32); + for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]); + scales[j] = make_qkx2_quants(32, 15, x + 32*j, weights, L + 32*j, &mins[j], Laux, -1.f, 0.1f, 20, false); + float scale = scales[j]; + if (scale > max_scale) { + max_scale = scale; + } + float min = mins[j]; + if (min > max_min) { + max_min = min; + } + } + + float inv_scale = max_scale > 0 ? 63.f/max_scale : 0.f; + float inv_min = max_min > 0 ? 63.f/max_min : 0.f; + for (int j = 0; j < QK_K/32; ++j) { + uint8_t ls = nearest_int(inv_scale*scales[j]); + uint8_t lm = nearest_int(inv_min*mins[j]); + ls = MIN(63, ls); + lm = MIN(63, lm); + if (j < 4) { + y[i].scales[j] = ls; + y[i].scales[j+4] = lm; + } else { + y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4); + y[i].scales[j-4] |= ((ls >> 4) << 6); + y[i].scales[j-0] |= ((lm >> 4) << 6); + } + } + y[i].d = GGML_FP32_TO_FP16(max_scale/63.f); + y[i].dmin = GGML_FP32_TO_FP16(max_min/63.f); + + uint8_t sc, m; + for (int j = 0; j < QK_K/32; ++j) { + get_scale_min_k4(j, y[i].scales, &sc, &m); + const float d = GGML_FP16_TO_FP32(y[i].d) * sc; + if (!d) continue; + const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m; + for (int ii = 0; ii < 32; ++ii) { + int l = nearest_int((x[32*j + ii] + dm)/d); + l = MAX(0, MIN(15, l)); + L[32*j + ii] = l; + } + } + + uint8_t * q = y[i].qs; + for (int j = 0; j < QK_K; j += 64) { + for (int l = 0; l < 32; ++l) q[l] = L[j + l] | (L[j + l + 32] << 4); + q += 32; + } + + x += QK_K; + } +} + +void dequantize_row_q4_K(const block_q4_K * GGML_RESTRICT x, float * GGML_RESTRICT 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 * q = x[i].qs; + + const float d = GGML_FP16_TO_FP32(x[i].d); + const float min = GGML_FP16_TO_FP32(x[i].dmin); + + int is = 0; + uint8_t sc, m; + for (int j = 0; j < QK_K; j += 64) { + get_scale_min_k4(is + 0, x[i].scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, x[i].scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *y++ = d1 * (q[l] & 0xF) - m1; + for (int l = 0; l < 32; ++l) *y++ = d2 * (q[l] >> 4) - m2; + q += 32; is += 2; + } + } +} + +static void quantize_row_q4_K_impl(const float * GGML_RESTRICT x, block_q4_K * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + assert(n_per_row % QK_K == 0); + const int64_t nb = n_per_row / QK_K; + + uint8_t L[QK_K]; + uint8_t Laux[32]; + uint8_t Ls[QK_K/32]; + uint8_t Lm[QK_K/32]; + float weights[32]; + float sw[QK_K/32]; + float mins[QK_K/32]; + float scales[QK_K/32]; + + for (int i = 0; i < nb; i++) { + + float sum_x2 = 0; + for (int l = 0; l < QK_K; ++l) sum_x2 += x[l] * x[l]; + float sigma2 = 2*sum_x2/QK_K; + float av_x = sqrtf(sigma2); + + for (int j = 0; j < QK_K/32; ++j) { + if (quant_weights) { + const float * qw = quant_weights + QK_K*i + 32*j; + for (int l = 0; l < 32; ++l) weights[l] = qw[l] * sqrtf(sigma2 + x[32*j + l]*x[32*j + l]); + } else { + for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]); + } + float sumw = 0; + for (int l = 0; l < 32; ++l) sumw += weights[l]; + sw[j] = sumw; + scales[j] = make_qkx3_quants(32, 15, x + 32*j, weights, L + 32*j, &mins[j], Laux, -0.9f, 0.05f, 36, false); + } + + float d_block = make_qp_quants(QK_K/32, 63, scales, Ls, sw); + float m_block = make_qp_quants(QK_K/32, 63, mins, Lm, sw); + for (int j = 0; j < QK_K/32; ++j) { + uint8_t ls = Ls[j]; + uint8_t lm = Lm[j]; + if (j < 4) { + y[i].scales[j] = ls; + y[i].scales[j+4] = lm; + } else { + y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4); + y[i].scales[j-4] |= ((ls >> 4) << 6); + y[i].scales[j-0] |= ((lm >> 4) << 6); + } + } + y[i].d = GGML_FP32_TO_FP16(d_block); + y[i].dmin = GGML_FP32_TO_FP16(m_block); + + uint8_t sc, m; + for (int j = 0; j < QK_K/32; ++j) { + get_scale_min_k4(j, y[i].scales, &sc, &m); + const float d = GGML_FP16_TO_FP32(y[i].d) * sc; + if (!d) continue; + const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m; + for (int ii = 0; ii < 32; ++ii) { + int l = nearest_int((x[32*j + ii] + dm)/d); + l = MAX(0, MIN(15, l)); + L[32*j + ii] = l; + } + } + uint8_t * q = y[i].qs; + for (int j = 0; j < QK_K; j += 64) { + for (int l = 0; l < 32; ++l) q[l] = L[j + l] | (L[j + l + 32] << 4); + q += 32; + } + + x += QK_K; + + } +} + +size_t quantize_q4_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + size_t row_size = ggml_row_size(GGML_TYPE_Q4_K, n_per_row); + if (!quant_weights) { + quantize_row_q4_K_ref(src, dst, (int64_t)nrow*n_per_row); + } + else { + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q4_K_impl(src, (block_q4_K*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + } + return nrow * row_size; +} + +// ====================== 5-bit (de)-quantization + +void quantize_row_q5_K_ref(const float * GGML_RESTRICT x, block_q5_K * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + uint8_t L[QK_K]; + float mins[QK_K/32]; + float scales[QK_K/32]; + float weights[32]; + uint8_t Laux[32]; + + for (int i = 0; i < nb; i++) { + float max_scale = 0; // as we are deducting the min, scales are always positive + float max_min = 0; + for (int j = 0; j < QK_K/32; ++j) { + //scales[j] = make_qkx1_quants(32, 31, x + 32*j, L + 32*j, &mins[j], 9, 0.5f); + float sum_x2 = 0; + for (int l = 0; l < 32; ++l) sum_x2 += x[32*j + l] * x[32*j + l]; + float av_x = sqrtf(sum_x2/32); + for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]); + scales[j] = make_qkx2_quants(32, 31, x + 32*j, weights, L + 32*j, &mins[j], Laux, -0.5f, 0.1f, 15, false); + float scale = scales[j]; + if (scale > max_scale) { + max_scale = scale; + } + float min = mins[j]; + if (min > max_min) { + max_min = min; + } + } + + float inv_scale = max_scale > 0 ? 63.f/max_scale : 0.f; + float inv_min = max_min > 0 ? 63.f/max_min : 0.f; + for (int j = 0; j < QK_K/32; ++j) { + uint8_t ls = nearest_int(inv_scale*scales[j]); + uint8_t lm = nearest_int(inv_min*mins[j]); + ls = MIN(63, ls); + lm = MIN(63, lm); + if (j < 4) { + y[i].scales[j] = ls; + y[i].scales[j+4] = lm; + } else { + y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4); + y[i].scales[j-4] |= ((ls >> 4) << 6); + y[i].scales[j-0] |= ((lm >> 4) << 6); + } + } + y[i].d = GGML_FP32_TO_FP16(max_scale/63.f); + y[i].dmin = GGML_FP32_TO_FP16(max_min/63.f); + + uint8_t sc, m; + for (int j = 0; j < QK_K/32; ++j) { + get_scale_min_k4(j, y[i].scales, &sc, &m); + const float d = GGML_FP16_TO_FP32(y[i].d) * sc; + if (!d) continue; + const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m; + for (int ii = 0; ii < 32; ++ii) { + int l = nearest_int((x[32*j + ii] + dm)/d); + l = MAX(0, MIN(31, l)); + L[32*j + ii] = l; + } + } + + uint8_t * GGML_RESTRICT qh = y[i].qh; + uint8_t * GGML_RESTRICT ql = y[i].qs; + memset(qh, 0, QK_K/8); + + uint8_t m1 = 1, m2 = 2; + for (int n = 0; n < QK_K; n += 64) { + for (int j = 0; j < 32; ++j) { + int l1 = L[n + j]; + if (l1 > 15) { + l1 -= 16; qh[j] |= m1; + } + int l2 = L[n + j + 32]; + if (l2 > 15) { + l2 -= 16; qh[j] |= m2; + } + ql[j] = l1 | (l2 << 4); + } + m1 <<= 2; m2 <<= 2; + ql += 32; + } + + x += QK_K; + } +} + +void dequantize_row_q5_K(const block_q5_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + const uint8_t * ql = x[i].qs; + const uint8_t * qh = x[i].qh; + + const float d = GGML_FP16_TO_FP32(x[i].d); + const float min = GGML_FP16_TO_FP32(x[i].dmin); + + int is = 0; + uint8_t sc, m; + uint8_t u1 = 1, u2 = 2; + for (int j = 0; j < QK_K; j += 64) { + get_scale_min_k4(is + 0, x[i].scales, &sc, &m); + const float d1 = d * sc; const float m1 = min * m; + get_scale_min_k4(is + 1, x[i].scales, &sc, &m); + const float d2 = d * sc; const float m2 = min * m; + for (int l = 0; l < 32; ++l) *y++ = d1 * ((ql[l] & 0xF) + (qh[l] & u1 ? 16 : 0)) - m1; + for (int l = 0; l < 32; ++l) *y++ = d2 * ((ql[l] >> 4) + (qh[l] & u2 ? 16 : 0)) - m2; + ql += 32; is += 2; + u1 <<= 2; u2 <<= 2; + } + } +} + +static void quantize_row_q5_K_impl(const float * GGML_RESTRICT x, block_q5_K * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + assert(n_per_row % QK_K == 0); + const int64_t nb = n_per_row / QK_K; + + uint8_t L[QK_K]; + uint8_t Laux[32]; + uint8_t Ls[QK_K/32]; + uint8_t Lm[QK_K/32]; + float mins[QK_K/32]; + float scales[QK_K/32]; + float sw[QK_K/32]; + float weights[32]; + + for (int i = 0; i < nb; i++) { + + float sum_x2 = 0; + for (int l = 0; l < QK_K; ++l) sum_x2 += x[l] * x[l]; + float sigma2 = 2*sum_x2/QK_K; + float av_x = sqrtf(sigma2); + + for (int j = 0; j < QK_K/32; ++j) { + if (quant_weights) { + const float * qw = quant_weights + QK_K*i + 32*j; + for (int l = 0; l < 32; ++l) weights[l] = qw[l] * sqrtf(sigma2 + x[32*j + l]*x[32*j + l]); + } else { + for (int l = 0; l < 32; ++l) weights[l] = av_x + fabsf(x[32*j + l]); + } + float sumw = 0; + for (int l = 0; l < 32; ++l) sumw += weights[l]; + sw[j] = sumw; + + scales[j] = make_qkx3_quants(32, 31, x + 32*j, weights, L + 32*j, &mins[j], Laux, -0.9f, 0.05f, 36, false); + } + + float d_block = make_qp_quants(QK_K/32, 63, scales, Ls, sw); + float m_block = make_qp_quants(QK_K/32, 63, mins, Lm, sw); + + for (int j = 0; j < QK_K/32; ++j) { + uint8_t ls = Ls[j]; + uint8_t lm = Lm[j]; + ls = MIN(63, ls); + lm = MIN(63, lm); + if (j < 4) { + y[i].scales[j] = ls; + y[i].scales[j+4] = lm; + } else { + y[i].scales[j+4] = (ls & 0xF) | ((lm & 0xF) << 4); + y[i].scales[j-4] |= ((ls >> 4) << 6); + y[i].scales[j-0] |= ((lm >> 4) << 6); + } + } + y[i].d = GGML_FP32_TO_FP16(d_block); + y[i].dmin = GGML_FP32_TO_FP16(m_block); + + uint8_t sc, m; + for (int j = 0; j < QK_K/32; ++j) { + get_scale_min_k4(j, y[i].scales, &sc, &m); + const float d = GGML_FP16_TO_FP32(y[i].d) * sc; + if (!d) continue; + const float dm = GGML_FP16_TO_FP32(y[i].dmin) * m; + for (int ii = 0; ii < 32; ++ii) { + int l = nearest_int((x[32*j + ii] + dm)/d); + l = MAX(0, MIN(31, l)); + L[32*j + ii] = l; + } + } + + uint8_t * GGML_RESTRICT qh = y[i].qh; + uint8_t * GGML_RESTRICT ql = y[i].qs; + memset(qh, 0, QK_K/8); + + uint8_t m1 = 1, m2 = 2; + for (int n = 0; n < QK_K; n += 64) { + for (int j = 0; j < 32; ++j) { + int l1 = L[n + j]; + if (l1 > 15) { + l1 -= 16; qh[j] |= m1; + } + int l2 = L[n + j + 32]; + if (l2 > 15) { + l2 -= 16; qh[j] |= m2; + } + ql[j] = l1 | (l2 << 4); + } + m1 <<= 2; m2 <<= 2; + ql += 32; + } + + x += QK_K; + + } +} + +size_t quantize_q5_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + size_t row_size = ggml_row_size(GGML_TYPE_Q5_K, n_per_row); + if (!quant_weights) { + quantize_row_q5_K_ref(src, dst, (int64_t)nrow*n_per_row); + } + else { + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q5_K_impl(src, (block_q5_K*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + } + return nrow * row_size; +} + +// ====================== 6-bit (de)-quantization + +void quantize_row_q6_K_ref(const float * GGML_RESTRICT x, block_q6_K * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + int8_t L[QK_K]; + float scales[QK_K/16]; + + for (int i = 0; i < nb; i++) { + + float max_scale = 0; + float max_abs_scale = 0; + + for (int ib = 0; ib < QK_K/16; ++ib) { + + const float scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, NULL); + scales[ib] = scale; + + const float abs_scale = fabsf(scale); + if (abs_scale > max_abs_scale) { + max_abs_scale = abs_scale; + max_scale = scale; + } + + } + + if (max_abs_scale < GROUP_MAX_EPS) { + memset(&y[i], 0, sizeof(block_q6_K)); + y[i].d = GGML_FP32_TO_FP16(0.f); + x += QK_K; + continue; + } + + float iscale = -128.f/max_scale; + y[i].d = GGML_FP32_TO_FP16(1/iscale); + for (int ib = 0; ib < QK_K/16; ++ib) { + y[i].scales[ib] = MIN(127, nearest_int(iscale*scales[ib])); + } + + for (int j = 0; j < QK_K/16; ++j) { + float d = GGML_FP16_TO_FP32(y[i].d) * y[i].scales[j]; + if (!d) { + continue; + } + for (int ii = 0; ii < 16; ++ii) { + int l = nearest_int(x[16*j + ii]/d); + l = MAX(-32, MIN(31, l)); + L[16*j + ii] = l + 32; + } + } + + uint8_t * GGML_RESTRICT ql = y[i].ql; + uint8_t * GGML_RESTRICT qh = y[i].qh; + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + const uint8_t q1 = L[j + l + 0] & 0xF; + const uint8_t q2 = L[j + l + 32] & 0xF; + const uint8_t q3 = L[j + l + 64] & 0xF; + const uint8_t q4 = L[j + l + 96] & 0xF; + ql[l+ 0] = q1 | (q3 << 4); + ql[l+32] = q2 | (q4 << 4); + qh[l] = (L[j + l] >> 4) | ((L[j + l + 32] >> 4) << 2) | ((L[j + l + 64] >> 4) << 4) | ((L[j + l + 96] >> 4) << 6); + } + ql += 64; + qh += 32; + } + + x += QK_K; + } +} + +void dequantize_row_q6_K(const block_q6_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + + const uint8_t * GGML_RESTRICT ql = x[i].ql; + const uint8_t * GGML_RESTRICT qh = x[i].qh; + const int8_t * GGML_RESTRICT sc = x[i].scales; + + for (int n = 0; n < QK_K; n += 128) { + for (int l = 0; l < 32; ++l) { + int is = l/16; + const int8_t q1 = (int8_t)((ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32; + const int8_t q2 = (int8_t)((ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32; + const int8_t q3 = (int8_t)((ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32; + const int8_t q4 = (int8_t)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32; + y[l + 0] = d * sc[is + 0] * q1; + y[l + 32] = d * sc[is + 2] * q2; + y[l + 64] = d * sc[is + 4] * q3; + y[l + 96] = d * sc[is + 6] * q4; + } + y += 128; + ql += 64; + qh += 32; + sc += 8; + } + } +} + +static void quantize_row_q6_K_impl(const float * GGML_RESTRICT x, block_q6_K * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + assert(n_per_row % QK_K == 0); + const int64_t nb = n_per_row / QK_K; + + int8_t L[QK_K]; + float scales[QK_K/16]; + //float weights[16]; + + for (int i = 0; i < nb; i++) { + + //float sum_x2 = 0; + //for (int j = 0; j < QK_K; ++j) sum_x2 += x[j]*x[j]; + //float sigma2 = sum_x2/QK_K; + + float max_scale = 0; + float max_abs_scale = 0; + + for (int ib = 0; ib < QK_K/16; ++ib) { + + float scale; + if (quant_weights) { + const float * qw = quant_weights + QK_K*i + 16*ib; + //for (int j = 0; j < 16; ++j) weights[j] = qw[j] * sqrtf(sigma2 + x[16*ib + j]*x[16*ib + j]); + //scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, weights); + scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, qw); + } else { + scale = make_qx_quants(16, 32, x + 16*ib, L + 16*ib, 1, NULL); + } + scales[ib] = scale; + + const float abs_scale = fabsf(scale); + if (abs_scale > max_abs_scale) { + max_abs_scale = abs_scale; + max_scale = scale; + } + + } + + if (max_abs_scale < GROUP_MAX_EPS) { + memset(&y[i], 0, sizeof(block_q6_K)); + y[i].d = GGML_FP32_TO_FP16(0.f); + x += QK_K; + continue; + } + + float iscale = -128.f/max_scale; + y[i].d = GGML_FP32_TO_FP16(1/iscale); + for (int ib = 0; ib < QK_K/16; ++ib) { + y[i].scales[ib] = MIN(127, nearest_int(iscale*scales[ib])); + } + + for (int j = 0; j < QK_K/16; ++j) { + float d = GGML_FP16_TO_FP32(y[i].d) * y[i].scales[j]; + if (!d) { + continue; + } + for (int ii = 0; ii < 16; ++ii) { + int l = nearest_int(x[16*j + ii]/d); + l = MAX(-32, MIN(31, l)); + L[16*j + ii] = l + 32; + } + } + + uint8_t * GGML_RESTRICT ql = y[i].ql; + uint8_t * GGML_RESTRICT qh = y[i].qh; + for (int j = 0; j < QK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + const uint8_t q1 = L[j + l + 0] & 0xF; + const uint8_t q2 = L[j + l + 32] & 0xF; + const uint8_t q3 = L[j + l + 64] & 0xF; + const uint8_t q4 = L[j + l + 96] & 0xF; + ql[l+ 0] = q1 | (q3 << 4); + ql[l+32] = q2 | (q4 << 4); + qh[l] = (L[j + l] >> 4) | ((L[j + l + 32] >> 4) << 2) | ((L[j + l + 64] >> 4) << 4) | ((L[j + l + 96] >> 4) << 6); + } + ql += 64; + qh += 32; + } + + x += QK_K; + + } +} + +size_t quantize_q6_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + size_t row_size = ggml_row_size(GGML_TYPE_Q6_K, n_per_row); + if (!quant_weights) { + quantize_row_q6_K_ref(src, dst, (int64_t)nrow*n_per_row); + } + else { + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q6_K_impl(src, (block_q6_K*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + } + return nrow * row_size; +} + +static void quantize_row_q4_0_impl(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + static_assert(QK4_0 == 32, "QK4_0 must be 32"); + + if (!quant_weights) { + quantize_row_q4_0_ref(x, y, n_per_row); + return; + } + + float weight[QK4_0]; + int8_t L[QK4_0]; + + float sum_x2 = 0; + for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j]; + float sigma2 = sum_x2/n_per_row; + + const int64_t nb = n_per_row/QK4_0; + for (int ib = 0; ib < nb; ++ib) { + const float * xb = x + QK4_0 * ib; + const float * qw = quant_weights + QK4_0 * ib; + for (int j = 0; j < QK4_0; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + float d = make_qx_quants(QK4_0, 8, xb, L, 1, weight); + y[ib].d = GGML_FP32_TO_FP16(d); + for (int j = 0; j < 16; ++j) { + y[ib].qs[j] = L[j] | (L[j+16] << 4); + } + } +} + +size_t quantize_q1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + if (!quant_weights) { + quantize_row_q1_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_Q1_0, n_per_row); + } + size_t row_size = ggml_row_size(GGML_TYPE_Q1_0, n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q1_0_ref(src, (block_q1_0*)qrow, n_per_row); + src += n_per_row; + qrow += row_size; + } + return nrow * row_size; +} + + +size_t quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + if (!quant_weights) { + quantize_row_q4_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_Q4_0, n_per_row); + } + size_t row_size = ggml_row_size(GGML_TYPE_Q4_0, n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q4_0_impl(src, (block_q4_0*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + return nrow * row_size; +} + +static void quantize_row_q4_1_impl(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + static_assert(QK4_1 == 32, "QK4_1 must be 32"); + + if (!quant_weights) { + quantize_row_q4_1_ref(x, y, n_per_row); + return; + } + + float weight[QK4_1]; + uint8_t L[QK4_1], Laux[QK4_1]; + + float sum_x2 = 0; + for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j]; + float sigma2 = sum_x2/n_per_row; + + const int64_t nb = n_per_row/QK4_1; + for (int ib = 0; ib < nb; ++ib) { + const float * xb = x + QK4_1 * ib; + const float * qw = quant_weights + QK4_1 * ib; + for (int j = 0; j < QK4_1; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + float min; + float d = make_qkx3_quants(QK4_1, 15, xb, weight, L, &min, Laux, -0.9f, 0.05f, 36, false); + y[ib].d = GGML_FP32_TO_FP16(d); + y[ib].m = GGML_FP32_TO_FP16(-min); + for (int j = 0; j < 16; ++j) { + y[ib].qs[j] = L[j] | (L[j+16] << 4); + } + } +} + +size_t quantize_q4_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + if (!quant_weights) { + quantize_row_q4_1_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_Q4_1, n_per_row); + } + size_t row_size = ggml_row_size(GGML_TYPE_Q4_1, n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q4_1_impl(src, (block_q4_1*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + return nrow * row_size; +} + +static void quantize_row_q5_0_impl(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + static_assert(QK5_0 == 32, "QK5_0 must be 32"); + + if (!quant_weights) { + quantize_row_q5_0_ref(x, y, n_per_row); + return; + } + + float weight[QK5_0]; + int8_t L[QK5_0]; + + float sum_x2 = 0; + for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j]; + float sigma2 = sum_x2/n_per_row; + + const int64_t nb = n_per_row/QK5_0; + for (int ib = 0; ib < nb; ++ib) { + const float * xb = x + QK5_0 * ib; + const float * qw = quant_weights + QK5_0 * ib; + for (int j = 0; j < QK5_0; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + float d = make_qx_quants(QK5_0, 16, xb, L, 1, weight); + y[ib].d = GGML_FP32_TO_FP16(d); + + uint32_t qh = 0; + + for (int j = 0; j < 16; ++j) { + const uint8_t xi0 = L[j]; + const uint8_t xi1 = L[j+16]; + y[ib].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4); + + // get the 5-th bit and store it in qh at the right position + qh |= ((xi0 & 0x10u) >> 4) << (j + 0); + qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_0/2); + } + + memcpy(&y[ib].qh, &qh, sizeof(qh)); + } +} + +size_t quantize_q5_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + if (!quant_weights) { + quantize_row_q5_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_Q5_0, n_per_row); + } + size_t row_size = ggml_row_size(GGML_TYPE_Q5_0, n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q5_0_impl(src, (block_q5_0*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + return nrow * row_size; +} + +static void quantize_row_q5_1_impl(const float * GGML_RESTRICT x, block_q5_1 * GGML_RESTRICT y, int64_t n_per_row, const float * quant_weights) { + static_assert(QK5_1 == 32, "QK5_1 must be 32"); + + if (!quant_weights) { + quantize_row_q5_1_ref(x, y, n_per_row); + return; + } + + float weight[QK5_1]; + uint8_t L[QK5_1], Laux[QK5_1]; + + float sum_x2 = 0; + for (int j = 0; j < n_per_row; ++j) sum_x2 += x[j]*x[j]; + float sigma2 = sum_x2/n_per_row; + + const int64_t nb = n_per_row/QK5_1; + for (int ib = 0; ib < nb; ++ib) { + const float * xb = x + QK5_1 * ib; + const float * qw = quant_weights + QK5_1 * ib; + for (int j = 0; j < QK5_1; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + float min; + float d = make_qkx3_quants(QK5_1, 31, xb, weight, L, &min, Laux, -0.9f, 0.05f, 36, false); + y[ib].d = GGML_FP32_TO_FP16(d); + y[ib].m = GGML_FP32_TO_FP16(-min); + + uint32_t qh = 0; + for (int j = 0; j < 16; ++j) { + const uint8_t xi0 = L[j]; + const uint8_t xi1 = L[j+16]; + y[ib].qs[j] = (xi0 & 0x0F) | ((xi1 & 0x0F) << 4); + // get the 5-th bit and store it in qh at the right position + qh |= ((xi0 & 0x10u) >> 4) << (j + 0); + qh |= ((xi1 & 0x10u) >> 4) << (j + QK5_0/2); + } + memcpy(&y[ib].qh, &qh, sizeof(qh)); + } +} + +size_t quantize_q5_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + if (!quant_weights) { + quantize_row_q5_1_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_Q5_1, n_per_row); + } + size_t row_size = ggml_row_size(GGML_TYPE_Q5_1, n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q5_1_impl(src, (block_q5_1*)qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += row_size; + } + return nrow * row_size; +} + +size_t quantize_q8_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + (void)quant_weights; // not used + const size_t row_size = ggml_row_size(GGML_TYPE_Q8_0, n_per_row); + quantize_row_q8_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * row_size; +} + +size_t quantize_mxfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_UNUSED(quant_weights); + quantize_row_mxfp4_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_MXFP4, n_per_row); +} + +size_t quantize_nvfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_UNUSED(quant_weights); + quantize_row_nvfp4_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_NVFP4, n_per_row); +} + +// ====================== Ternary (de)-quantization (BitNet b1.58 and TriLMs) + +void quantize_row_tq1_0_ref(const float * GGML_RESTRICT x, block_tq1_0 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int64_t i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < QK_K; j++) { + const float v = x[j]; + amax = MAX(amax, fabsf(v)); + } + + const float d = amax; + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + // 5 elements per byte, along 32 bytes + for (size_t j = 0; j < sizeof(y->qs) - sizeof(y->qs) % 32; j += 32) { + for (size_t m = 0; m < 32; ++m) { + uint8_t q = 0; + for (size_t n = 0; n < 5; ++n) { + int xi = lroundf(x[m + n*32] * id) + 1; // -1, 0, 1 -> 0, 1, 2 + q *= 3; + q += xi; + } + // ceiling division (243 == pow(3, 5)) + q = ((uint16_t)q * 256 + (243 - 1)) / 243; + y[i].qs[j + m] = q; + } + x += 5*32; + } + // along 16 bytes + for (size_t j = sizeof(y->qs) - sizeof(y->qs) % 32; j < sizeof(y->qs); j += 16) { + for (size_t m = 0; m < 16; ++m) { + uint8_t q = 0; + for (size_t n = 0; n < 5; ++n) { + int xi = lroundf(x[m + n*16] * id) + 1; // -1, 0, 1 -> 0, 1, 2 + q *= 3; + q += xi; + } + // ceiling division (243 == pow(3, 5)) + q = ((uint16_t)q * 256 + (243 - 1)) / 243; + y[i].qs[j + m] = q; + } + x += 5*16; + } + // 4 elements per byte + for (size_t j = 0; j < sizeof(y->qh); ++j) { + uint8_t q = 0; + for (size_t m = 0; m < 4; ++m) { + // -1, 0, 1 -> 0, 1, 2 + int xi = lroundf(x[j + m*sizeof(y->qh)] * id) + 1; + q *= 3; + q += xi; + } + // shift the first value to the most significant trit + q *= 3; + // ceiling division (243 == pow(3, 5)) + q = ((uint16_t)q * 256 + (243 - 1)) / 243; + y[i].qh[j] = q; + } + x += 4*sizeof(y->qh); + } +} + +void quantize_row_tq2_0_ref(const float * GGML_RESTRICT x, block_tq2_0 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int64_t i = 0; i < nb; i++) { + float amax = 0.0f; // absolute max + + for (int j = 0; j < QK_K; j++) { + const float v = x[j]; + amax = MAX(amax, fabsf(v)); + } + + const float d = amax; + const float id = d ? 1.0f/d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + for (size_t j = 0; j < sizeof(y->qs); j += 32) { + for (size_t m = 0; m < 32; ++m) { + uint8_t q = 0; + for (size_t n = 0; n < 4; ++n) { + // -1, 0, 1 -> 0, 1, 2 + int xi = lroundf(x[m + n*32] * id) + 1; + q += (xi & 3) << (2*n); + } + y[i].qs[j + m] = q; + } + x += 4*32; + } + } +} + +size_t quantize_tq1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + (void)quant_weights; // not used + const size_t row_size = ggml_row_size(GGML_TYPE_TQ1_0, n_per_row); + quantize_row_tq1_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * row_size; +} + +size_t quantize_tq2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + (void)quant_weights; // not used + const size_t row_size = ggml_row_size(GGML_TYPE_TQ2_0, n_per_row); + quantize_row_tq2_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * row_size; +} + +void dequantize_row_tq1_0(const block_tq1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + const uint8_t pow3[6] = {1, 3, 9, 27, 81, 243}; + + for (int64_t i = 0; i < nb; ++i) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (size_t j = 0; j < sizeof(x->qs) - sizeof(x->qs) % 32; j += 32) { + for (size_t n = 0; n < 5; ++n) { + for (size_t m = 0; m < 32; ++m) { + uint8_t q = x[i].qs[j + m] * pow3[n]; + int16_t xi = ((uint16_t) q * 3) >> 8; + *y++ = (float) (xi - 1) * d; + } + } + } + for (size_t j = sizeof(x->qs) - sizeof(x->qs) % 32; j < sizeof(x->qs); j += 16) { + for (size_t n = 0; n < 5; ++n) { + for (size_t m = 0; m < 16; ++m) { + uint8_t q = x[i].qs[j + m] * pow3[n]; + int16_t xi = ((uint16_t) q * 3) >> 8; + *y++ = (float) (xi - 1) * d; + } + } + } + + for (size_t n = 0; n < 4; ++n) { + for (size_t j = 0; j < sizeof(x->qh); ++j) { + uint8_t q = x[i].qh[j] * pow3[n]; + int16_t xi = ((uint16_t) q * 3) >> 8; + *y++ = (float) (xi - 1) * d; + } + } + } +} + +void dequantize_row_tq2_0(const block_tq2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int64_t i = 0; i < nb; ++i) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (size_t j = 0; j < sizeof(x->qs); j += 32) { + for (size_t l = 0; l < 4; ++l) { + for (size_t m = 0; m < 32; ++m) { + int8_t q = (x[i].qs[j + m] >> (l*2)) & 3; + *y++ = (float) (q - 1) * d; + } + } + } + } +} + +// ====================== "True" 2-bit (de)-quantization + +void dequantize_row_iq2_xxs(const block_iq2_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + uint32_t aux32[2]; + const uint8_t * aux8 = (const uint8_t *)aux32; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + memcpy(aux32, x[i].qs + 4*ib32, 2*sizeof(uint32_t)); + const float db = d * (0.5f + (aux32[1] >> 28)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[l]); + const uint8_t signs = ksigns_iq2xs[(aux32[1] >> 7*l) & 127]; + for (int j = 0; j < 8; ++j) { + y[j] = db * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); + } + y += 8; + } + } + } +} + +// ====================== 2.3125 bpw (de)-quantization + +void dequantize_row_iq2_xs(const block_iq2_xs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + float db[2]; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + db[0] = d * (0.5f + (x[i].scales[ib32] & 0xf)) * 0.25f; + db[1] = d * (0.5f + (x[i].scales[ib32] >> 4)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (x[i].qs[4*ib32 + l] & 511)); + const uint8_t signs = ksigns_iq2xs[x[i].qs[4*ib32 + l] >> 9]; + for (int j = 0; j < 8; ++j) { + y[j] = db[l/2] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); + } + y += 8; + } + } + } +} + +// ====================== 2.5625 bpw (de)-quantization + +void dequantize_row_iq2_s(const block_iq2_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + float db[2]; + + 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 * signs = qs + QK_K/8; + + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + db[0] = d * (0.5f + (x[i].scales[ib32] & 0xf)) * 0.25f; + db[1] = d * (0.5f + (x[i].scales[ib32] >> 4)) * 0.25f; + for (int l = 0; l < 4; ++l) { + const float dl = db[l/2]; + const uint8_t * grid = (const uint8_t *)(iq2s_grid + (qs[l] | (qh[ib32] << (8-2*l) & 0x300))); + for (int j = 0; j < 8; ++j) { + y[j] = dl * grid[j] * (signs[l] & kmask_iq2xs[j] ? -1.f : 1.f); + } + y += 8; + } + qs += 4; + signs += 4; + } + } +} + +// ====================== 3.0625 bpw (de)-quantization + +void dequantize_row_iq3_xxs(const block_iq3_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + uint32_t aux32; + + 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 * scales_and_signs = qs + QK_K/4; + + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + memcpy(&aux32, scales_and_signs + 4*ib32, sizeof(uint32_t)); + const float db = d * (0.5f + (aux32 >> 28)) * 0.5f; + for (int l = 0; l < 4; ++l) { + const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*l) & 127]; + const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + qs[2*l+0]); + const uint8_t * grid2 = (const uint8_t *)(iq3xxs_grid + qs[2*l+1]); + for (int j = 0; j < 4; ++j) { + y[j+0] = db * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f); + y[j+4] = db * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f); + } + y += 8; + } + qs += 8; + } + } +} + +// ====================== 3.3125 bpw (de)-quantization + +void dequantize_row_iq3_s(const block_iq3_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t 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 * signs = x[i].signs; + + for (int ib32 = 0; ib32 < QK_K/32; ib32 += 2) { + const float db1 = d * (1 + 2*(x[i].scales[ib32/2] & 0xf)); + const float db2 = d * (1 + 2*(x[i].scales[ib32/2] >> 4)); + for (int l = 0; l < 4; ++l) { + const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[0] << (8-2*l)) & 256))); + const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[0] << (7-2*l)) & 256))); + for (int j = 0; j < 4; ++j) { + y[j+0] = db1 * grid1[j] * (signs[l] & kmask_iq2xs[j+0] ? -1.f : 1.f); + y[j+4] = db1 * grid2[j] * (signs[l] & kmask_iq2xs[j+4] ? -1.f : 1.f); + } + y += 8; + } + qs += 8; + signs += 4; + for (int l = 0; l < 4; ++l) { + const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*l+0] | ((qh[1] << (8-2*l)) & 256))); + const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*l+1] | ((qh[1] << (7-2*l)) & 256))); + for (int j = 0; j < 4; ++j) { + y[j+0] = db2 * grid1[j] * (signs[l] & kmask_iq2xs[j+0] ? -1.f : 1.f); + y[j+4] = db2 * grid2[j] * (signs[l] & kmask_iq2xs[j+4] ? -1.f : 1.f); + } + y += 8; + } + qh += 2; + qs += 8; + signs += 4; + } + } +} + +// ====================== 1.5625 bpw (de)-quantization + +void dequantize_row_iq1_s(const block_iq1_s * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t 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 uint16_t * qh = x[i].qh; + + for (int ib = 0; ib < QK_K/32; ++ib) { + const float dl = d * (2*((qh[ib] >> 12) & 7) + 1); + const float delta = qh[ib] & 0x8000 ? -IQ1S_DELTA : IQ1S_DELTA; + for (int l = 0; l < 4; ++l) { + const int8_t * grid = (const int8_t *)(iq1s_grid + (qs[l] | (((qh[ib] >> 3*l) & 7) << 8))); + for (int j = 0; j < 8; ++j) { + y[j] = dl * (grid[j] + delta); + } + y += 8; + } + qs += 4; + } + } +} + +void dequantize_row_iq1_m(const block_iq1_m * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + float delta[4]; + uint16_t idx[4]; + + iq1m_scale_t scale; + + for (int i = 0; i < nb; i++) { + + const uint16_t * sc = (const uint16_t *)x[i].scales; + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); + const float d = GGML_FP16_TO_FP32(scale.f16); + + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + + for (int ib = 0; ib < QK_K/32; ++ib) { + const float dl1 = d * (2*((sc[ib/2] >> (6*(ib%2)+0)) & 0x7) + 1); + const float dl2 = d * (2*((sc[ib/2] >> (6*(ib%2)+3)) & 0x7) + 1); + + idx[0] = qs[0] | ((qh[0] << 8) & 0x700); + idx[1] = qs[1] | ((qh[0] << 4) & 0x700); + idx[2] = qs[2] | ((qh[1] << 8) & 0x700); + idx[3] = qs[3] | ((qh[1] << 4) & 0x700); + delta[0] = qh[0] & 0x08 ? -IQ1S_DELTA : IQ1S_DELTA; + delta[1] = qh[0] & 0x80 ? -IQ1S_DELTA : IQ1S_DELTA; + delta[2] = qh[1] & 0x08 ? -IQ1S_DELTA : IQ1S_DELTA; + delta[3] = qh[1] & 0x80 ? -IQ1S_DELTA : IQ1S_DELTA; + for (int l = 0; l < 2; ++l) { + const int8_t * grid = (const int8_t *)(iq1s_grid + idx[l]); + for (int j = 0; j < 8; ++j) { + y[j] = dl1 * (grid[j] + delta[l]); + } + y += 8; + } + for (int l = 2; l < 4; ++l) { + const int8_t * grid = (const int8_t *)(iq1s_grid + idx[l]); + for (int j = 0; j < 8; ++j) { + y[j] = dl2 * (grid[j] + delta[l]); + } + y += 8; + } + qs += 4; + qh += 2; + } + } +} + +void dequantize_row_iq4_nl(const block_iq4_nl * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK4_NL == 0); + const int64_t nb = k / QK4_NL; + + for (int i = 0; i < nb; i++) { + + const uint8_t * qs = x[i].qs; + + const float d = GGML_FP16_TO_FP32(x[i].d); + for (int j = 0; j < QK4_NL/2; ++j) { + y[j+ 0] = d * kvalues_iq4nl[qs[j] & 0xf]; + y[j+QK4_NL/2] = d * kvalues_iq4nl[qs[j] >> 4]; + } + y += QK4_NL; + qs += QK4_NL/2; + } +} + +void dequantize_row_iq4_xs(const block_iq4_xs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t 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); + + for (int ib = 0; ib < QK_K/32; ++ib) { + const int ls = ((x[i].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[i].scales_h >> 2*ib) & 3) << 4); + const float dl = d * (ls - 32); + for (int j = 0; j < 16; ++j) { + y[j+ 0] = dl * kvalues_iq4nl[qs[j] & 0xf]; + y[j+16] = dl * kvalues_iq4nl[qs[j] >> 4]; + } + y += 32; + qs += 16; + } + } +} + +//===================================== Q8_K ============================================== + +void quantize_row_q8_K_ref(const float * GGML_RESTRICT x, block_q8_K * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + 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); + } + 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; + } + y[i].d = 1/iscale; + x += QK_K; + } +} + +void dequantize_row_q8_K(const block_q8_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + const int64_t nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + for (int j = 0; j < QK_K; ++j) { + *y++ = x[i].d * x[i].qs[j]; + } + } +} + +// ================================ IQ2 quantization ============================================= + +typedef struct { + uint64_t * grid; + int * map; + uint16_t * neighbours; +} iq2_entry_t; + +static iq2_entry_t iq2_data[4] = { + {NULL, NULL, NULL}, + {NULL, NULL, NULL}, + {NULL, NULL, NULL}, + {NULL, NULL, NULL}, +}; + +static inline int iq2_data_index(enum ggml_type type) { + GGML_ASSERT(type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M || type == GGML_TYPE_IQ2_S); + return type == GGML_TYPE_IQ2_XXS ? 0 : + type == GGML_TYPE_IQ2_XS ? 1 : + type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? 2 : 3; +} + +static inline int iq2_grid_size(enum ggml_type type) { + GGML_ASSERT(type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M || type == GGML_TYPE_IQ2_S); + return type == GGML_TYPE_IQ2_XXS ? 256 : + type == GGML_TYPE_IQ2_XS ? 512 : + type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? NGRID_IQ1S : 1024; +} + +static int iq2_compare_func(const void * left, const void * right) { + const int * l = (const int *)left; + const int * r = (const int *)right; + return l[0] < r[0] ? -1 : l[0] > r[0] ? 1 : l[1] < r[1] ? -1 : l[1] > r[1] ? 1 : 0; +} + +void iq2xs_init_impl(enum ggml_type type) { + const int gindex = iq2_data_index(type); + const int grid_size = iq2_grid_size(type); + if (iq2_data[gindex].grid) { + return; + } + static const uint16_t kgrid_2bit_256[256] = { + 0, 2, 5, 8, 10, 17, 20, 32, 34, 40, 42, 65, 68, 80, 88, 97, + 100, 128, 130, 138, 162, 257, 260, 272, 277, 320, 388, 408, 512, 514, 546, 642, + 1025, 1028, 1040, 1057, 1060, 1088, 1090, 1096, 1120, 1153, 1156, 1168, 1188, 1280, 1282, 1288, + 1312, 1350, 1385, 1408, 1425, 1545, 1552, 1600, 1668, 1700, 2048, 2053, 2056, 2068, 2088, 2113, + 2116, 2128, 2130, 2184, 2308, 2368, 2562, 2580, 4097, 4100, 4112, 4129, 4160, 4192, 4228, 4240, + 4245, 4352, 4360, 4384, 4432, 4442, 4480, 4644, 4677, 5120, 5128, 5152, 5157, 5193, 5248, 5400, + 5474, 5632, 5654, 6145, 6148, 6160, 6208, 6273, 6400, 6405, 6560, 6737, 8192, 8194, 8202, 8260, + 8289, 8320, 8322, 8489, 8520, 8704, 8706, 9217, 9220, 9232, 9280, 9302, 9472, 9537, 9572, 9872, + 10248, 10272, 10388, 10820, 16385, 16388, 16400, 16408, 16417, 16420, 16448, 16456, 16470, 16480, 16513, 16516, + 16528, 16640, 16672, 16737, 16768, 16773, 16897, 16912, 16968, 16982, 17000, 17408, 17416, 17440, 17536, 17561, + 17682, 17700, 17920, 18433, 18436, 18448, 18496, 18501, 18688, 18776, 18785, 18818, 19013, 19088, 20480, 20488, + 20497, 20505, 20512, 20608, 20616, 20740, 20802, 20900, 21137, 21648, 21650, 21770, 22017, 22100, 22528, 22545, + 22553, 22628, 22848, 23048, 24580, 24592, 24640, 24680, 24832, 24917, 25112, 25184, 25600, 25605, 25872, 25874, + 25988, 26690, 32768, 32770, 32778, 32833, 32898, 33028, 33048, 33088, 33297, 33793, 33796, 33808, 33813, 33856, + 33888, 34048, 34118, 34196, 34313, 34368, 34400, 34818, 35076, 35345, 36868, 36880, 36900, 36928, 37025, 37142, + 37248, 37445, 37888, 37922, 37956, 38225, 39041, 39200, 40962, 41040, 41093, 41225, 41472, 42008, 43088, 43268, + }; + static const uint16_t kgrid_2bit_512[512] = { + 0, 2, 5, 8, 10, 17, 20, 22, 25, 32, 34, 37, 40, 65, 68, 70, + 73, 80, 82, 85, 88, 97, 100, 128, 130, 133, 136, 145, 148, 153, 160, 257, + 260, 262, 265, 272, 274, 277, 280, 282, 289, 292, 320, 322, 325, 328, 337, 340, + 352, 360, 385, 388, 400, 512, 514, 517, 520, 529, 532, 544, 577, 580, 592, 597, + 640, 650, 1025, 1028, 1030, 1033, 1040, 1042, 1045, 1048, 1057, 1060, 1088, 1090, 1093, 1096, + 1105, 1108, 1110, 1120, 1153, 1156, 1168, 1280, 1282, 1285, 1288, 1297, 1300, 1312, 1345, 1348, + 1360, 1377, 1408, 1537, 1540, 1552, 1574, 1600, 1602, 1668, 2048, 2050, 2053, 2056, 2058, 2065, + 2068, 2080, 2085, 2113, 2116, 2128, 2136, 2176, 2208, 2218, 2305, 2308, 2320, 2368, 2433, 2441, + 2560, 2592, 2600, 2710, 2720, 4097, 4100, 4102, 4105, 4112, 4114, 4117, 4120, 4129, 4132, 4160, + 4162, 4165, 4168, 4177, 4180, 4192, 4202, 4225, 4228, 4240, 4352, 4354, 4357, 4360, 4369, 4372, + 4384, 4417, 4420, 4432, 4480, 4500, 4502, 4609, 4612, 4614, 4624, 4672, 4704, 5120, 5122, 5125, + 5128, 5137, 5140, 5152, 5185, 5188, 5193, 5200, 5220, 5248, 5377, 5380, 5392, 5440, 5632, 5652, + 5705, 6145, 6148, 6160, 6162, 6208, 6228, 6278, 6400, 6405, 6502, 6737, 6825, 8192, 8194, 8197, + 8200, 8202, 8209, 8212, 8224, 8257, 8260, 8272, 8320, 8352, 8449, 8452, 8464, 8512, 8520, 8549, + 8704, 8738, 8832, 8872, 9217, 9220, 9232, 9257, 9280, 9472, 9537, 9554, 9625, 9729, 9754, 9894, + 10240, 10248, 10250, 10272, 10325, 10376, 10402, 10600, 10640, 10760, 10784, 10882, 10888, 10890, 16385, 16388, + 16390, 16393, 16400, 16402, 16405, 16408, 16417, 16420, 16448, 16450, 16453, 16456, 16458, 16465, 16468, 16480, + 16485, 16513, 16516, 16528, 16640, 16642, 16645, 16648, 16657, 16660, 16672, 16705, 16708, 16720, 16768, 16773, + 16802, 16897, 16900, 16912, 16914, 16937, 16960, 17408, 17410, 17413, 17416, 17425, 17428, 17433, 17440, 17473, + 17476, 17488, 17536, 17556, 17665, 17668, 17680, 17700, 17728, 17818, 17920, 17930, 17988, 18000, 18433, 18436, + 18448, 18496, 18501, 18516, 18530, 18688, 18705, 18756, 18768, 18793, 18948, 20480, 20482, 20485, 20488, 20497, + 20500, 20512, 20520, 20545, 20548, 20560, 20608, 20737, 20740, 20752, 20757, 20800, 20802, 20992, 21060, 21162, + 21505, 21508, 21520, 21537, 21568, 21600, 21633, 21665, 21760, 21768, 21888, 21896, 22049, 22120, 22177, 22528, + 22548, 22593, 22608, 22681, 22810, 22848, 22850, 23173, 24577, 24580, 24592, 24640, 24660, 24674, 24710, 24745, + 24832, 25124, 25162, 25234, 25600, 25622, 25872, 25920, 25925, 26020, 26625, 26730, 26917, 27142, 27220, 27234, + 32768, 32770, 32773, 32776, 32785, 32788, 32800, 32810, 32833, 32836, 32848, 32896, 32898, 32936, 32938, 33025, + 33028, 33030, 33040, 33088, 33105, 33113, 33280, 33312, 33408, 33410, 33440, 33448, 33793, 33796, 33808, 33810, + 33813, 33856, 33888, 33929, 34048, 34116, 34213, 34328, 34410, 34816, 34824, 34853, 34906, 34944, 34946, 34984, + 35078, 35362, 35456, 35464, 35478, 35496, 36865, 36868, 36880, 36928, 36950, 36996, 37120, 37154, 37220, 37462, + 37513, 37888, 37893, 37956, 37968, 37976, 38185, 38288, 38290, 38465, 38993, 39078, 39241, 39445, 39520, 40960, + 40962, 40968, 40970, 40992, 41002, 41120, 41297, 41305, 41382, 41472, 41474, 41480, 41514, 41600, 41632, 42048, + 42133, 42597, 42648, 43018, 43040, 43042, 43048, 43168, 43176, 43268, 43396, 43398, 43560, 43562, 43665, 43690, + }; + static const uint16_t kgrid_1bit_2048[NGRID_IQ1S] = { + 0, 2, 5, 8, 10, 17, 21, 32, 34, 40, 42, 69, 81, 84, 86, 101, + 128, 130, 136, 138, 149, 160, 162, 168, 170, 260, 261, 273, 276, 278, 281, 282, + 293, 321, 326, 329, 338, 341, 346, 353, 356, 358, 360, 389, 401, 404, 406, 421, + 512, 514, 520, 522, 533, 544, 546, 552, 554, 581, 593, 601, 612, 617, 640, 642, + 648, 650, 657, 661, 665, 672, 674, 680, 682, 1041, 1044, 1046, 1061, 1089, 1097, 1109, + 1114, 1124, 1125, 1169, 1177, 1189, 1281, 1284, 1285, 1286, 1301, 1304, 1306, 1321, 1344, 1349, + 1354, 1360, 1361, 1364, 1365, 1366, 1369, 1376, 1378, 1381, 1384, 1386, 1409, 1425, 1429, 1432, + 1434, 1441, 1444, 1445, 1446, 1449, 1556, 1561, 1601, 1604, 1616, 1618, 1621, 1624, 1632, 1633, + 1638, 1641, 1669, 1681, 1684, 1689, 2048, 2050, 2056, 2058, 2069, 2080, 2082, 2088, 2090, 2117, + 2129, 2134, 2149, 2176, 2178, 2184, 2186, 2197, 2208, 2210, 2216, 2218, 2309, 2321, 2324, 2329, + 2340, 2341, 2369, 2384, 2385, 2389, 2401, 2404, 2409, 2449, 2452, 2454, 2457, 2469, 2560, 2562, + 2568, 2570, 2581, 2592, 2594, 2600, 2602, 2629, 2641, 2649, 2657, 2661, 2688, 2690, 2693, 2696, + 2698, 2709, 2720, 2722, 2728, 2730, 4112, 4113, 4116, 4121, 4132, 4133, 4161, 4164, 4176, 4181, + 4184, 4193, 4196, 4197, 4201, 4241, 4244, 4246, 4257, 4261, 4353, 4356, 4358, 4361, 4368, 4370, + 4373, 4376, 4385, 4388, 4393, 4421, 4426, 4432, 4433, 4434, 4436, 4437, 4438, 4441, 4448, 4453, + 4484, 4498, 4501, 4513, 4516, 4625, 4628, 4630, 4645, 4672, 4678, 4681, 4690, 4693, 4696, 4698, + 4708, 4710, 4741, 4753, 4756, 4758, 4773, 5121, 5126, 5129, 5140, 5141, 5144, 5145, 5153, 5158, + 5185, 5189, 5190, 5192, 5194, 5201, 5204, 5205, 5206, 5209, 5218, 5221, 5224, 5252, 5257, 5264, + 5268, 5269, 5272, 5273, 5274, 5281, 5284, 5285, 5289, 5378, 5381, 5386, 5393, 5396, 5397, 5398, + 5401, 5408, 5410, 5413, 5416, 5418, 5441, 5444, 5445, 5446, 5457, 5458, 5460, 5461, 5462, 5465, + 5466, 5473, 5476, 5477, 5478, 5481, 5504, 5506, 5508, 5509, 5512, 5514, 5520, 5521, 5524, 5525, + 5526, 5529, 5530, 5536, 5538, 5541, 5633, 5636, 5637, 5638, 5653, 5654, 5656, 5658, 5665, 5670, + 5696, 5698, 5700, 5701, 5704, 5706, 5713, 5717, 5718, 5720, 5721, 5729, 5732, 5733, 5736, 5737, + 5738, 5766, 5770, 5778, 5781, 5796, 5801, 6161, 6166, 6181, 6209, 6212, 6214, 6217, 6224, 6229, + 6232, 6234, 6240, 6241, 6244, 6246, 6249, 6277, 6289, 6292, 6309, 6416, 6418, 6421, 6426, 6433, + 6437, 6466, 6468, 6469, 6472, 6481, 6484, 6485, 6486, 6489, 6490, 6496, 6501, 6506, 6537, 6545, + 6546, 6549, 6552, 6561, 6566, 6569, 6665, 6678, 6692, 6694, 6724, 6726, 6729, 6736, 6738, 6741, + 6744, 6753, 6758, 6761, 6789, 6801, 6806, 6810, 8192, 8194, 8200, 8202, 8213, 8224, 8226, 8229, + 8232, 8234, 8261, 8273, 8281, 8289, 8293, 8320, 8322, 8328, 8330, 8341, 8352, 8354, 8357, 8360, + 8362, 8453, 8465, 8468, 8473, 8485, 8514, 8516, 8521, 8533, 8536, 8538, 8545, 8548, 8549, 8550, + 8581, 8592, 8598, 8601, 8613, 8705, 8712, 8714, 8721, 8725, 8736, 8738, 8744, 8746, 8773, 8785, + 8790, 8793, 8805, 8833, 8840, 8842, 8849, 8853, 8864, 8866, 8872, 8874, 9221, 9236, 9238, 9241, + 9253, 9284, 9285, 9286, 9289, 9298, 9301, 9304, 9306, 9318, 9349, 9361, 9364, 9369, 9377, 9381, + 9481, 9493, 9505, 9513, 9536, 9541, 9544, 9553, 9556, 9557, 9561, 9570, 9573, 9576, 9609, 9616, + 9620, 9621, 9624, 9626, 9633, 9636, 9638, 9641, 9733, 9744, 9746, 9753, 9765, 9793, 9801, 9813, + 9824, 9825, 9833, 9860, 9862, 9872, 9882, 10240, 10242, 10248, 10250, 10261, 10272, 10274, 10280, 10282, + 10309, 10321, 10324, 10341, 10368, 10370, 10376, 10378, 10400, 10402, 10408, 10410, 10505, 10513, 10516, 10521, + 10533, 10566, 10569, 10578, 10581, 10593, 10596, 10598, 10601, 10629, 10640, 10646, 10649, 10660, 10661, 10752, + 10754, 10760, 10762, 10784, 10786, 10792, 10794, 10821, 10833, 10838, 10841, 10853, 10880, 10882, 10888, 10890, + 10901, 10912, 10914, 10920, 10922, 16389, 16401, 16406, 16421, 16457, 16466, 16469, 16472, 16474, 16481, 16484, + 16486, 16532, 16537, 16545, 16550, 16640, 16641, 16644, 16646, 16649, 16658, 16661, 16662, 16664, 16666, 16673, + 16678, 16681, 16709, 16712, 16714, 16721, 16724, 16725, 16726, 16729, 16730, 16741, 16744, 16746, 16769, 16772, + 16774, 16784, 16786, 16789, 16800, 16801, 16802, 16901, 16913, 16916, 16918, 16933, 16961, 16978, 16981, 16986, + 16996, 17001, 17033, 17044, 17061, 17409, 17429, 17433, 17449, 17477, 17480, 17482, 17489, 17492, 17493, 17494, + 17505, 17506, 17509, 17512, 17514, 17537, 17542, 17545, 17552, 17554, 17557, 17568, 17569, 17577, 17665, 17666, + 17669, 17674, 17681, 17684, 17685, 17686, 17689, 17696, 17701, 17706, 17729, 17732, 17733, 17734, 17737, 17744, + 17745, 17748, 17749, 17750, 17752, 17753, 17761, 17764, 17765, 17766, 17769, 17794, 17796, 17797, 17800, 17809, + 17812, 17813, 17814, 17817, 17818, 17829, 17832, 17834, 17921, 17925, 17929, 17940, 17941, 17944, 17946, 17953, + 17956, 17961, 17984, 17986, 17989, 17992, 18000, 18001, 18002, 18005, 18006, 18009, 18018, 18021, 18024, 18049, + 18053, 18058, 18068, 18069, 18081, 18084, 18086, 18437, 18449, 18453, 18458, 18469, 18498, 18505, 18512, 18517, + 18520, 18529, 18532, 18534, 18537, 18565, 18577, 18580, 18582, 18585, 18597, 18689, 18693, 18694, 18698, 18704, + 18708, 18709, 18712, 18721, 18724, 18726, 18752, 18757, 18762, 18769, 18770, 18772, 18773, 18774, 18777, 18784, + 18786, 18789, 18790, 18794, 18822, 18825, 18834, 18837, 18838, 18840, 18849, 18852, 18854, 18857, 18966, 19012, + 19014, 19017, 19029, 19032, 19034, 19044, 19049, 19092, 19109, 20481, 20484, 20485, 20486, 20489, 20498, 20501, + 20506, 20513, 20516, 20521, 20544, 20549, 20552, 20561, 20564, 20565, 20566, 20569, 20581, 20584, 20614, 20617, + 20629, 20632, 20640, 20641, 20646, 20649, 20741, 20744, 20745, 20746, 20753, 20756, 20757, 20758, 20760, 20761, + 20768, 20773, 20774, 20776, 20778, 20801, 20804, 20805, 20806, 20809, 20816, 20817, 20818, 20820, 20821, 20822, + 20824, 20825, 20826, 20833, 20836, 20837, 20838, 20841, 20866, 20869, 20881, 20884, 20885, 20886, 20889, 20896, + 20901, 20906, 20993, 20998, 21010, 21013, 21018, 21025, 21028, 21058, 21061, 21066, 21073, 21076, 21077, 21078, + 21081, 21090, 21093, 21125, 21136, 21138, 21141, 21145, 21146, 21156, 21508, 21509, 21521, 21524, 21525, 21526, + 21528, 21529, 21537, 21541, 21544, 21546, 21569, 21572, 21573, 21574, 21577, 21578, 21584, 21585, 21588, 21589, + 21590, 21592, 21593, 21594, 21601, 21602, 21604, 21605, 21606, 21609, 21632, 21640, 21642, 21649, 21652, 21653, + 21654, 21657, 21665, 21668, 21669, 21674, 21761, 21762, 21764, 21765, 21766, 21769, 21776, 21777, 21778, 21780, + 21781, 21782, 21785, 21786, 21793, 21796, 21797, 21798, 21801, 21824, 21825, 21826, 21828, 21829, 21830, 21832, + 21833, 21840, 21841, 21842, 21844, 21845, 21846, 21848, 21849, 21850, 21856, 21857, 21860, 21861, 21862, 21864, + 21865, 21866, 21889, 21892, 21893, 21897, 21898, 21904, 21905, 21908, 21909, 21910, 21912, 21913, 21921, 21924, + 21925, 21926, 21929, 22016, 22017, 22018, 22020, 22022, 22024, 22025, 22033, 22036, 22037, 22040, 22041, 22048, + 22049, 22050, 22052, 22053, 22054, 22056, 22057, 22081, 22085, 22086, 22088, 22089, 22090, 22096, 22097, 22098, + 22100, 22101, 22102, 22104, 22105, 22106, 22113, 22116, 22117, 22121, 22146, 22149, 22150, 22152, 22153, 22154, + 22161, 22165, 22170, 22178, 22181, 22182, 22184, 22185, 22532, 22533, 22534, 22537, 22544, 22549, 22552, 22561, + 22570, 22597, 22600, 22602, 22609, 22612, 22613, 22614, 22616, 22617, 22624, 22626, 22628, 22629, 22658, 22665, + 22672, 22674, 22677, 22680, 22689, 22697, 22785, 22786, 22789, 22794, 22801, 22804, 22805, 22806, 22809, 22821, + 22849, 22852, 22853, 22854, 22857, 22864, 22865, 22866, 22868, 22869, 22870, 22872, 22873, 22874, 22881, 22884, + 22885, 22886, 22889, 22913, 22917, 22921, 22929, 22932, 22933, 22934, 22936, 22937, 22949, 23044, 23048, 23061, + 23066, 23072, 23077, 23078, 23081, 23109, 23112, 23113, 23121, 23125, 23126, 23128, 23129, 23138, 23141, 23144, + 23146, 23169, 23178, 23186, 23189, 23190, 23192, 23194, 23201, 24581, 24596, 24598, 24601, 24613, 24644, 24656, + 24661, 24662, 24664, 24666, 24673, 24676, 24678, 24681, 24705, 24726, 24741, 24833, 24836, 24838, 24841, 24850, + 24853, 24865, 24866, 24870, 24873, 24901, 24905, 24913, 24917, 24918, 24921, 24933, 24934, 24938, 24964, 24970, + 24978, 24981, 24993, 24998, 25001, 25105, 25110, 25113, 25152, 25153, 25158, 25173, 25174, 25176, 25184, 25221, + 25233, 25238, 25253, 25617, 25618, 25621, 25622, 25626, 25633, 25638, 25641, 25664, 25666, 25669, 25672, 25674, + 25681, 25684, 25685, 25686, 25689, 25690, 25696, 25698, 25701, 25732, 25733, 25737, 25744, 25746, 25748, 25749, + 25750, 25752, 25754, 25761, 25764, 25769, 25861, 25864, 25866, 25873, 25877, 25878, 25881, 25924, 25925, 25926, + 25929, 25936, 25937, 25940, 25941, 25942, 25945, 25953, 25956, 25957, 25958, 25961, 25990, 25993, 25994, 26001, + 26005, 26006, 26009, 26010, 26018, 26021, 26022, 26024, 26114, 26121, 26133, 26144, 26150, 26152, 26153, 26176, + 26181, 26184, 26186, 26193, 26196, 26197, 26198, 26200, 26202, 26208, 26213, 26216, 26240, 26242, 26245, 26250, + 26260, 26262, 26264, 26265, 26272, 26276, 26278, 26282, 26646, 26649, 26661, 26689, 26706, 26709, 26714, 26721, + 26729, 26757, 26769, 26776, 26790, 26881, 26884, 26896, 26901, 26913, 26916, 26918, 26921, 26944, 26945, 26949, + 26950, 26952, 26961, 26964, 26965, 26966, 26969, 26976, 26981, 26986, 27010, 27012, 27018, 27029, 27041, 27044, + 27045, 27049, 27153, 27158, 27160, 27201, 27204, 27209, 27216, 27221, 27224, 27226, 27236, 27237, 27241, 27270, + 27284, 27288, 27290, 27302, 32768, 32770, 32776, 32778, 32800, 32802, 32808, 32810, 32837, 32848, 32849, 32852, + 32854, 32857, 32869, 32896, 32898, 32904, 32906, 32917, 32928, 32930, 32936, 32938, 33029, 33041, 33044, 33046, + 33049, 33061, 33089, 33092, 33097, 33104, 33106, 33109, 33110, 33112, 33113, 33124, 33126, 33129, 33157, 33161, + 33172, 33174, 33177, 33189, 33280, 33282, 33288, 33290, 33301, 33312, 33314, 33320, 33322, 33361, 33364, 33369, + 33381, 33408, 33410, 33416, 33418, 33429, 33440, 33442, 33448, 33450, 33812, 33817, 33857, 33860, 33873, 33877, + 33882, 33889, 33892, 33897, 33940, 33945, 34049, 34057, 34066, 34069, 34074, 34086, 34089, 34112, 34113, 34117, + 34120, 34129, 34132, 34133, 34134, 34137, 34138, 34149, 34150, 34152, 34154, 34177, 34180, 34182, 34185, 34192, + 34194, 34197, 34200, 34214, 34321, 34326, 34329, 34341, 34369, 34372, 34377, 34378, 34384, 34389, 34393, 34394, + 34401, 34406, 34410, 34437, 34449, 34458, 34468, 34816, 34818, 34824, 34826, 34837, 34848, 34850, 34856, 34858, + 34881, 34885, 34897, 34900, 34905, 34917, 34921, 34944, 34946, 34952, 34954, 34965, 34976, 34978, 34984, 34986, + 35077, 35078, 35089, 35092, 35094, 35109, 35137, 35140, 35142, 35145, 35152, 35154, 35157, 35162, 35169, 35172, + 35205, 35222, 35225, 35237, 35328, 35330, 35336, 35338, 35349, 35360, 35362, 35368, 35370, 35397, 35409, 35412, + 35414, 35456, 35458, 35464, 35466, 35477, 35488, 35490, 35496, 35498, 36869, 36881, 36886, 36888, 36889, 36901, + 36929, 36934, 36937, 36949, 36952, 36954, 36969, 36970, 36997, 37009, 37012, 37014, 37017, 37029, 37121, 37124, + 37126, 37129, 37136, 37141, 37144, 37146, 37153, 37156, 37158, 37161, 37184, 37189, 37200, 37201, 37204, 37205, + 37206, 37209, 37218, 37221, 37252, 37254, 37266, 37269, 37272, 37281, 37284, 37286, 37289, 37381, 37393, 37396, + 37401, 37413, 37444, 37446, 37449, 37456, 37458, 37461, 37464, 37478, 37481, 37509, 37524, 37526, 37545, 37889, + 37892, 37894, 37904, 37909, 37912, 37926, 37952, 37962, 37969, 37972, 37973, 37974, 37976, 37977, 37984, 37985, + 37986, 37989, 38020, 38022, 38034, 38036, 38037, 38040, 38049, 38057, 38144, 38149, 38152, 38154, 38160, 38161, + 38164, 38165, 38166, 38169, 38177, 38181, 38185, 38186, 38209, 38212, 38213, 38214, 38217, 38224, 38225, 38226, + 38228, 38229, 38230, 38232, 38233, 38234, 38241, 38244, 38245, 38246, 38249, 38273, 38277, 38280, 38289, 38290, + 38292, 38293, 38294, 38297, 38298, 38304, 38306, 38309, 38312, 38314, 38401, 38404, 38416, 38421, 38425, 38432, + 38438, 38441, 38469, 38472, 38473, 38481, 38482, 38485, 38486, 38489, 38501, 38504, 38530, 38532, 38537, 38538, + 38546, 38548, 38549, 38564, 38566, 38569, 38917, 38934, 38937, 38949, 38977, 38982, 38992, 38994, 38997, 38998, + 39002, 39012, 39013, 39045, 39057, 39062, 39065, 39077, 39172, 39174, 39177, 39184, 39186, 39189, 39192, 39194, + 39200, 39201, 39204, 39206, 39232, 39234, 39237, 39240, 39242, 39249, 39252, 39253, 39254, 39257, 39266, 39269, + 39270, 39274, 39297, 39300, 39312, 39314, 39317, 39322, 39329, 39334, 39429, 39445, 39461, 39492, 39494, 39497, + 39504, 39509, 39512, 39521, 39557, 39569, 39572, 39573, 39574, 40960, 40962, 40968, 40970, 40981, 40992, 40994, + 41000, 41002, 41029, 41041, 41044, 41046, 41049, 41088, 41090, 41096, 41098, 41109, 41120, 41122, 41128, 41130, + 41221, 41225, 41233, 41236, 41238, 41241, 41242, 41286, 41289, 41297, 41301, 41304, 41306, 41313, 41316, 41349, + 41360, 41362, 41366, 41369, 41474, 41480, 41482, 41488, 41497, 41506, 41512, 41514, 41541, 41553, 41558, 41561, + 41573, 41600, 41602, 41608, 41610, 41621, 41632, 41634, 41640, 41642, 42009, 42021, 42049, 42052, 42064, 42068, + 42069, 42072, 42074, 42081, 42085, 42086, 42088, 42089, 42117, 42246, 42249, 42256, 42258, 42261, 42264, 42278, + 42281, 42306, 42309, 42321, 42324, 42325, 42326, 42329, 42341, 42346, 42369, 42372, 42373, 42374, 42377, 42386, + 42389, 42392, 42501, 42513, 42518, 42522, 42529, 42533, 42564, 42566, 42570, 42578, 42581, 42582, 42584, 42592, + 42594, 42630, 42640, 42645, 42646, 42649, 42657, 42660, 42662, 43008, 43010, 43016, 43018, 43040, 43042, 43048, + 43050, 43089, 43092, 43094, 43097, 43136, 43138, 43144, 43146, 43157, 43168, 43170, 43176, 43178, 43269, 43284, + 43289, 43297, 43301, 43329, 43344, 43349, 43354, 43361, 43366, 43369, 43408, 43414, 43520, 43522, 43528, 43530, + 43552, 43554, 43560, 43562, 43601, 43604, 43606, 43648, 43650, 43656, 43658, 43669, 43680, 43682, 43688, 43690, + }; + static const uint16_t kgrid_2bit_1024[1024] = { + 0, 2, 5, 8, 10, 17, 20, 22, 25, 32, 34, 37, 40, 65, 68, 70, + 73, 80, 82, 85, 88, 97, 100, 102, 105, 128, 130, 133, 136, 145, 148, 160, + 165, 170, 257, 260, 262, 265, 272, 274, 277, 280, 289, 292, 320, 322, 325, 328, + 337, 340, 342, 345, 352, 357, 360, 385, 388, 400, 402, 405, 417, 420, 512, 514, + 517, 520, 529, 532, 544, 554, 577, 580, 582, 585, 592, 597, 640, 645, 650, 660, + 674, 1025, 1028, 1030, 1033, 1040, 1042, 1045, 1048, 1057, 1060, 1062, 1065, 1088, 1090, 1093, + 1096, 1098, 1105, 1108, 1110, 1113, 1120, 1122, 1125, 1153, 1156, 1158, 1161, 1168, 1173, 1176, + 1185, 1188, 1280, 1282, 1285, 1288, 1290, 1297, 1300, 1302, 1305, 1312, 1317, 1320, 1345, 1348, + 1350, 1353, 1360, 1362, 1365, 1368, 1377, 1380, 1408, 1410, 1413, 1416, 1425, 1428, 1440, 1537, + 1540, 1542, 1545, 1552, 1557, 1600, 1605, 1608, 1617, 1620, 1632, 1665, 1668, 1680, 2048, 2050, + 2053, 2056, 2065, 2068, 2070, 2073, 2080, 2085, 2090, 2113, 2116, 2118, 2121, 2128, 2130, 2133, + 2136, 2145, 2148, 2176, 2181, 2196, 2218, 2305, 2308, 2320, 2322, 2325, 2328, 2337, 2368, 2373, + 2376, 2385, 2388, 2400, 2433, 2448, 2560, 2577, 2580, 2594, 2600, 2602, 2640, 2713, 4097, 4100, + 4102, 4105, 4112, 4114, 4117, 4120, 4129, 4132, 4134, 4160, 4162, 4165, 4168, 4177, 4180, 4182, + 4185, 4192, 4194, 4197, 4200, 4225, 4228, 4230, 4240, 4245, 4248, 4257, 4260, 4352, 4354, 4357, + 4360, 4362, 4369, 4372, 4374, 4377, 4384, 4386, 4389, 4392, 4417, 4420, 4422, 4425, 4432, 4434, + 4437, 4440, 4449, 4452, 4480, 4482, 4485, 4488, 4497, 4500, 4609, 4612, 4617, 4624, 4629, 4641, + 4644, 4672, 4677, 4689, 4692, 4737, 4740, 4752, 5120, 5122, 5125, 5128, 5137, 5140, 5142, 5145, + 5152, 5157, 5160, 5185, 5188, 5190, 5193, 5200, 5202, 5205, 5208, 5217, 5220, 5248, 5250, 5253, + 5256, 5265, 5268, 5280, 5377, 5380, 5382, 5385, 5392, 5394, 5397, 5400, 5409, 5412, 5440, 5442, + 5445, 5448, 5457, 5460, 5472, 5505, 5508, 5520, 5632, 5637, 5640, 5649, 5652, 5664, 5697, 5700, + 5712, 5760, 5802, 6145, 6148, 6150, 6153, 6160, 6165, 6168, 6177, 6208, 6210, 6213, 6216, 6225, + 6228, 6240, 6273, 6276, 6400, 6402, 6405, 6408, 6417, 6420, 6432, 6465, 6468, 6480, 6505, 6562, + 6660, 6672, 6720, 6742, 8192, 8194, 8197, 8200, 8209, 8212, 8214, 8217, 8224, 8229, 8234, 8257, + 8260, 8272, 8274, 8277, 8292, 8320, 8330, 8340, 8362, 8449, 8452, 8464, 8466, 8469, 8481, 8512, + 8514, 8517, 8529, 8532, 8544, 8577, 8580, 8592, 8704, 8714, 8738, 8744, 8746, 8772, 8784, 8840, + 8842, 8872, 9217, 9220, 9222, 9225, 9232, 9237, 9240, 9249, 9252, 9280, 9282, 9285, 9288, 9297, + 9300, 9312, 9345, 9348, 9360, 9472, 9477, 9480, 9489, 9492, 9504, 9537, 9540, 9552, 9574, 9600, + 9729, 9732, 9744, 9792, 9817, 10240, 10245, 10257, 10260, 10305, 10308, 10320, 10378, 10410, 10497, 10500, + 10512, 10645, 10762, 10786, 10852, 10888, 10890, 16385, 16388, 16390, 16393, 16400, 16402, 16405, 16408, 16410, + 16417, 16420, 16422, 16448, 16450, 16453, 16456, 16458, 16465, 16468, 16470, 16473, 16480, 16482, 16485, 16513, + 16516, 16528, 16533, 16536, 16545, 16548, 16640, 16642, 16645, 16648, 16657, 16660, 16662, 16665, 16672, 16674, + 16677, 16705, 16708, 16710, 16713, 16720, 16722, 16725, 16728, 16737, 16740, 16768, 16770, 16773, 16776, 16785, + 16788, 16800, 16897, 16900, 16912, 16914, 16917, 16920, 16932, 16960, 16965, 16968, 16977, 16980, 16992, 17025, + 17028, 17408, 17410, 17413, 17416, 17418, 17425, 17428, 17430, 17433, 17440, 17442, 17445, 17448, 17473, 17476, + 17478, 17481, 17488, 17490, 17493, 17496, 17505, 17508, 17536, 17538, 17541, 17544, 17553, 17556, 17568, 17665, + 17668, 17670, 17673, 17680, 17682, 17685, 17688, 17697, 17700, 17728, 17730, 17733, 17736, 17745, 17748, 17760, + 17770, 17793, 17796, 17808, 17920, 17922, 17925, 17928, 17937, 17940, 17952, 17985, 17988, 18000, 18048, 18085, + 18433, 18436, 18441, 18448, 18450, 18453, 18456, 18465, 18468, 18496, 18498, 18501, 18504, 18513, 18516, 18528, + 18564, 18576, 18688, 18690, 18693, 18696, 18705, 18708, 18720, 18753, 18756, 18768, 18816, 18838, 18945, 18948, + 18960, 19008, 20480, 20482, 20485, 20488, 20497, 20500, 20502, 20505, 20512, 20514, 20517, 20520, 20545, 20548, + 20550, 20553, 20560, 20562, 20565, 20568, 20577, 20580, 20608, 20610, 20613, 20616, 20625, 20628, 20737, 20740, + 20742, 20745, 20752, 20754, 20757, 20760, 20769, 20772, 20800, 20802, 20805, 20808, 20817, 20820, 20832, 20865, + 20868, 20880, 20992, 20997, 21000, 21009, 21012, 21024, 21057, 21060, 21072, 21097, 21120, 21505, 21508, 21510, + 21513, 21520, 21522, 21525, 21528, 21537, 21540, 21568, 21570, 21573, 21576, 21585, 21588, 21600, 21633, 21636, + 21648, 21760, 21762, 21765, 21768, 21777, 21780, 21792, 21825, 21828, 21840, 21888, 22017, 22020, 22032, 22054, + 22080, 22528, 22530, 22533, 22536, 22545, 22548, 22560, 22593, 22596, 22608, 22618, 22656, 22785, 22788, 22800, + 22848, 23040, 23065, 23173, 23208, 24577, 24580, 24582, 24592, 24594, 24597, 24600, 24609, 24612, 24640, 24645, + 24648, 24657, 24660, 24672, 24708, 24720, 24832, 24834, 24837, 24840, 24849, 24852, 24864, 24897, 24900, 24912, + 24960, 24985, 25092, 25104, 25152, 25174, 25249, 25600, 25605, 25608, 25617, 25620, 25632, 25665, 25668, 25680, + 25728, 25857, 25860, 25872, 25920, 25930, 25960, 26002, 26112, 26260, 26625, 26628, 26640, 26725, 26776, 26880, + 26922, 27202, 27297, 32768, 32770, 32773, 32776, 32785, 32788, 32793, 32800, 32805, 32833, 32836, 32848, 32850, + 32853, 32856, 32865, 32896, 32901, 32913, 32916, 33025, 33028, 33033, 33040, 33042, 33045, 33048, 33057, 33060, + 33088, 33090, 33093, 33096, 33105, 33108, 33153, 33156, 33168, 33193, 33280, 33285, 33290, 33297, 33300, 33345, + 33348, 33360, 33793, 33796, 33798, 33801, 33808, 33810, 33813, 33816, 33825, 33856, 33858, 33861, 33864, 33873, + 33876, 33888, 33921, 33924, 33936, 34048, 34050, 34053, 34056, 34065, 34068, 34080, 34113, 34116, 34128, 34176, + 34186, 34305, 34308, 34320, 34345, 34368, 34816, 34821, 34833, 34836, 34881, 34884, 34896, 34978, 35073, 35076, + 35136, 35173, 35362, 35416, 35418, 35458, 35490, 36865, 36868, 36873, 36880, 36882, 36885, 36888, 36900, 36928, + 36930, 36933, 36936, 36945, 36948, 36960, 36993, 36996, 37008, 37120, 37125, 37137, 37140, 37185, 37188, 37200, + 37210, 37377, 37380, 37392, 37440, 37542, 37888, 37890, 37893, 37896, 37905, 37908, 37920, 37953, 37956, 37968, + 38016, 38038, 38145, 38148, 38160, 38208, 38296, 38305, 38400, 38470, 38500, 38913, 38916, 38928, 38950, 38976, + 39081, 39168, 39241, 39250, 39568, 40960, 40965, 40970, 40980, 40994, 41002, 41025, 41028, 41040, 41122, 41130, + 41280, 41317, 41474, 41482, 41506, 41512, 41514, 41602, 41608, 41610, 41640, 41985, 41988, 42000, 42048, 42121, + 42148, 42240, 42265, 42577, 43018, 43048, 43170, 43348, 43398, 43528, 43530, 43552, 43554, 43560, 43656, 43690, + }; + + const int kmap_size = 43692; + //const int nwant = type == GGML_TYPE_IQ1_S ? 3 : 2; + const int nwant = type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? 3 : type == GGML_TYPE_IQ2_S ? 1 : 2; + const uint16_t * kgrid = type == GGML_TYPE_IQ2_XXS ? kgrid_2bit_256 : + type == GGML_TYPE_IQ2_XS ? kgrid_2bit_512 : + type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M ? kgrid_1bit_2048 : kgrid_2bit_1024; + uint64_t * kgrid_q2xs; + int * kmap_q2xs; + uint16_t * kneighbors_q2xs; + + //printf("================================================================= %s(grid_size = %d)\n", __func__, grid_size); + uint64_t * the_grid = (uint64_t *)malloc(grid_size*sizeof(uint64_t)); + for (int k = 0; k < grid_size; ++k) { + int8_t * pos = (int8_t *)(the_grid + k); + for (int i = 0; i < 8; ++i) { + int l = (kgrid[k] >> 2*i) & 0x3; + pos[i] = 2*l + 1; + } + } + kgrid_q2xs = the_grid; + iq2_data[gindex].grid = the_grid; + kmap_q2xs = (int *)malloc(kmap_size*sizeof(int)); + iq2_data[gindex].map = kmap_q2xs; + for (int i = 0; i < kmap_size; ++i) kmap_q2xs[i] = -1; + uint64_t aux64; + uint8_t * aux8 = (uint8_t *)&aux64; + for (int i = 0; i < grid_size; ++i) { + aux64 = kgrid_q2xs[i]; + uint16_t index = 0; + for (int k=0; k<8; ++k) { + uint16_t q = (aux8[k] - 1)/2; + index |= (q << 2*k); + } + kmap_q2xs[index] = i; + } + int8_t pos[8]; + int * dist2 = (int *)malloc(2*grid_size*sizeof(int)); + int num_neighbors = 0, num_not_in_map = 0; + for (int i = 0; i < kmap_size; ++i) { + if (kmap_q2xs[i] >= 0) continue; + ++num_not_in_map; + for (int k = 0; k < 8; ++k) { + int l = (i >> 2*k) & 0x3; + pos[k] = 2*l + 1; + } + for (int j = 0; j < grid_size; ++j) { + const int8_t * pg = (const int8_t *)(kgrid_q2xs + j); + int d2 = 0; + for (int k = 0; k < 8; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]); + dist2[2*j+0] = d2; + dist2[2*j+1] = j; + } + qsort(dist2, grid_size, 2*sizeof(int), iq2_compare_func); + int n = 0; int d2 = dist2[0]; + int nhave = 1; + for (int j = 0; j < grid_size; ++j) { + if (dist2[2*j] > d2) { + if (nhave == nwant) break; + d2 = dist2[2*j]; + ++nhave; + } + ++n; + } + num_neighbors += n; + } + //printf("%s: %d neighbours in total\n", __func__, num_neighbors); + kneighbors_q2xs = (uint16_t *)malloc((num_neighbors + num_not_in_map)*sizeof(uint16_t)); + iq2_data[gindex].neighbours = kneighbors_q2xs; + int counter = 0; + for (int i = 0; i < kmap_size; ++i) { + if (kmap_q2xs[i] >= 0) continue; + for (int k = 0; k < 8; ++k) { + int l = (i >> 2*k) & 0x3; + pos[k] = 2*l + 1; + } + for (int j = 0; j < grid_size; ++j) { + const int8_t * pg = (const int8_t *)(kgrid_q2xs + j); + int d2 = 0; + for (int k = 0; k < 8; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]); + dist2[2*j+0] = d2; + dist2[2*j+1] = j; + } + qsort(dist2, grid_size, 2*sizeof(int), iq2_compare_func); + kmap_q2xs[i] = -(counter + 1); + int d2 = dist2[0]; + uint16_t * start = &kneighbors_q2xs[counter++]; + int n = 0, nhave = 1; + for (int j = 0; j < grid_size; ++j) { + if (dist2[2*j] > d2) { + if (nhave == nwant) break; + d2 = dist2[2*j]; + ++nhave; + } + kneighbors_q2xs[counter++] = dist2[2*j+1]; + ++n; + } + *start = n; + } + free(dist2); +} + +void iq2xs_free_impl(enum ggml_type type) { + GGML_ASSERT(type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M || type == GGML_TYPE_IQ2_S); + const int gindex = iq2_data_index(type); + if (iq2_data[gindex].grid) { + free(iq2_data[gindex].grid); iq2_data[gindex].grid = NULL; + free(iq2_data[gindex].map); iq2_data[gindex].map = NULL; + free(iq2_data[gindex].neighbours); iq2_data[gindex].neighbours = NULL; + } +} + +static int iq2_find_best_neighbour(const uint16_t * GGML_RESTRICT neighbours, const uint64_t * GGML_RESTRICT grid, + const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float scale, int8_t * GGML_RESTRICT L) { + int num_neighbors = neighbours[0]; + GGML_ASSERT(num_neighbors > 0); + float best_d2 = FLT_MAX; + int grid_index = -1; + for (int j = 1; j <= num_neighbors; ++j) { + const int8_t * pg = (const int8_t *)(grid + neighbours[j]); + float d2 = 0; + for (int i = 0; i < 8; ++i) { + float q = pg[i]; + float diff = scale*q - xval[i]; + d2 += weight[i]*diff*diff; + } + if (d2 < best_d2) { + best_d2 = d2; grid_index = neighbours[j]; + } + } + GGML_ASSERT(grid_index >= 0); + const int8_t * pg = (const int8_t *)(grid + grid_index); + for (int i = 0; i < 8; ++i) L[i] = (pg[i] - 1)/2; + return grid_index; +} + +static void quantize_row_iq2_xxs_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights) { + + const int gindex = iq2_data_index(GGML_TYPE_IQ2_XXS); + + const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; + const int * kmap_q2xs = iq2_data[gindex].map; + const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; + + GGML_ASSERT(quant_weights && "missing quantization weights"); + GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + const int kMaxQ = 3; + + const int64_t nbl = n/QK_K; + + block_iq2_xxs * y = vy; + + float scales[QK_K/32]; + float weight[32]; + float xval[32]; + int8_t L[32]; + int8_t Laux[32]; + float waux[32]; + uint8_t block_signs[4]; + uint32_t q2[2*(QK_K/32)]; + + for (int ibl = 0; ibl < nbl; ++ibl) { + + y[ibl].d = GGML_FP32_TO_FP16(0.f); + memset(q2, 0, QK_K/4); + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = sumx2/QK_K; + + for (int ib = 0; ib < QK_K/32; ++ib) { + const float * xb = xbl + 32*ib; + const float * qw = quant_weights + QK_K*ibl + 32*ib; + for (int i = 0; i < 32; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + for (int i = 0; i < 32; ++i) waux[i] = sqrtf(weight[i]); + for (int k = 0; k < 4; ++k) { + int nflip = 0; + uint8_t s = 0; + for (int i = 0; i < 8; ++i) { + if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; + else { + xval[8*k + i] = -xb[8*k + i]; ++nflip; s |= (1 << i); + } + } + if (nflip%2) { + int imin = 0; float min = weight[8*k+imin]*xb[8*k+imin]*xb[8*k+imin]; + for (int i = 1; i < 8; ++i) { + float ax = weight[8*k+i]*xb[8*k+i]*xb[8*k+i]; + if (ax < min) { + min = ax; imin = i; + } + } + xval[8*k+imin] = -xval[8*k+imin]; + s ^= (1 << imin); + } + block_signs[k] = s & 127; + } + float max = xval[0]; + for (int i = 1; i < 32; ++i) max = MAX(max, xval[i]); + if (max < GROUP_MAX_EPS) { + scales[ib] = 0; + memset(L, 0, 32); + continue; + } + float scale = make_qp_quants(32, kMaxQ+1, xval, (uint8_t*)L, weight); + float eff_max = scale*kMaxQ; + if (eff_max <= 0) { + scales[ib] = 0; + memset(L, 0, 32); + continue; + } + float best = 0; + for (int is = -6; is <= 6; ++is) { + float id = (2*kMaxQ-1+is*0.1f)/eff_max; + float this_scale = 1/id; + for (int k = 0; k < 4; ++k) { + for (int i = 0; i < 8; ++i) { + int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); + Laux[8*k+i] = MAX(0, MIN(kMaxQ-1, l)); + } + uint16_t u = 0; + for (int i = 0; i < 8; ++i) u |= (Laux[8*k+i] << 2*i); + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, this_scale, Laux + 8*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 32; ++i) { + float w = weight[i]; + float q = 2*Laux[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + scale = sumqx/sumq2; best = scale*sumqx; + memcpy(L, Laux, 32); + } + } + if (scale > 0) { + float id = 1/scale; + for (int k = 0; k < 4; ++k) { + uint16_t u = 0; + for (int i = 0; i < 8; ++i) { + int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); + l = MAX(0, MIN(kMaxQ-1, l)); + u |= (l << 2*i); + } + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, scale, L + 8*k); + } + const int8_t * pg = (const int8_t *)(kgrid_q2xs + grid_index); + for (int i = 0; i < 8; ++i) L[8*k+i] = (pg[i] - 1)/2; + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 32; ++i) { + float w = weight[i]; + float q = 2*L[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0) scale = sumqx/sumq2; + } + if (scale < 0) { + // This should never happen, but just in case, flip scale so that it is positive (we use uint's to encode the scale) + // and correspondingly flip quant signs. + scale = -scale; + for (int k = 0; k < 4; ++k) block_signs[k] = (~block_signs[k]) & 127; + } + for (int k = 0; k < 4; ++k) { + uint16_t u = 0; + for (int i = 0; i < 8; ++i) u |= (L[8*k+i] << 2*i); + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + printf("Oops: found point %u not on grid:", u); + for (int i = 0; i < 8; ++i) printf(" %d", L[8*k+i]); + printf("\n"); + GGML_ABORT("fatal error"); + } + q2[2*ib+0] |= ((uint32_t) grid_index << 8*k); + q2[2*ib+1] |= (block_signs[k] << 7*k); + } + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + memset(y[ibl].qs, 0, QK_K/4); + continue; + } + + float d = max_scale/31; + y[ibl].d = GGML_FP32_TO_FP16(d); + float id = 1/d; + for (int ib = 0; ib < QK_K/32; ++ib) { + int l = nearest_int(0.5f*(id*scales[ib]-1)); + l = MAX(0, MIN(15, l)); + q2[2*ib+1] |= ((uint32_t)l << 28); + } + memcpy(y[ibl].qs, q2, QK_K/4); + } +} + +static void quantize_row_iq2_xs_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights) { + + const int gindex = iq2_data_index(GGML_TYPE_IQ2_XS); + + const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; + const int * kmap_q2xs = iq2_data[gindex].map; + const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; + + GGML_ASSERT(quant_weights && "missing quantization weights"); + GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + const int kMaxQ = 3; + + const int64_t nbl = n/QK_K; + + block_iq2_xs * y = vy; + + float scales[QK_K/16]; + float weight[16]; + float xval[16]; + int8_t L[16]; + int8_t Laux[16]; + float waux[16]; + bool is_on_grid[2]; + bool is_on_grid_aux[2]; + uint8_t block_signs[2]; + uint16_t q2[2*(QK_K/16)]; + + for (int ibl = 0; ibl < nbl; ++ibl) { + + y[ibl].d = GGML_FP32_TO_FP16(0.f); + memset(q2, 0, QK_K/4); + memset(y[ibl].scales, 0, QK_K/32); + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = sumx2/QK_K; + + for (int ib = 0; ib < QK_K/16; ++ib) { + const float * xb = xbl + 16*ib; + const float * qw = quant_weights + QK_K*ibl + 16*ib; + for (int i = 0; i < 16; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + for (int i = 0; i < 16; ++i) waux[i] = sqrtf(weight[i]); + for (int k = 0; k < 2; ++k) { + int nflip = 0; + uint8_t s = 0; + for (int i = 0; i < 8; ++i) { + if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; + else { + xval[8*k + i] = -xb[8*k + i]; ++nflip; s |= (1 << i); + } + } + if (nflip%2) { + int imin = 0; float min = weight[8*k+imin]*xb[8*k+imin]*xb[8*k+imin]; + for (int i = 1; i < 8; ++i) { + float ax = weight[8*k+i]*xb[8*k+i]*xb[8*k+i]; + if (ax < min) { + min = ax; imin = i; + } + } + xval[8*k+imin] = -xval[8*k+imin]; + s ^= (1 << imin); + } + block_signs[k] = s & 127; + } + float max = xval[0]; + for (int i = 1; i < 16; ++i) max = MAX(max, xval[i]); + memset(L, 0, 16); + if (max < GROUP_MAX_EPS) { + scales[ib] = 0; + continue; + } + float best = 0; + float scale = max/(2*kMaxQ-1); + is_on_grid[0] = is_on_grid[1] = true; + for (int is = -9; is <= 9; ++is) { + float id = (2*kMaxQ-1+is*0.1f)/max; + float this_scale = 1/id; + for (int k = 0; k < 2; ++k) { + for (int i = 0; i < 8; ++i) { + int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); + Laux[8*k+i] = MAX(0, MIN(kMaxQ-1, l)); + } + uint16_t u = 0; + for (int i = 0; i < 8; ++i) u |= (Laux[8*k+i] << 2*i); + int grid_index = kmap_q2xs[u]; + is_on_grid_aux[k] = true; + if (grid_index < 0) { + is_on_grid_aux[k] = false; + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, this_scale, Laux + 8*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 16; ++i) { + float w = weight[i]; + float q = 2*Laux[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + scale = sumqx/sumq2; best = scale*sumqx; + for (int i = 0; i < 16; ++i) L[i] = Laux[i]; + for (int k = 0; k < 2; ++k) is_on_grid[k] = is_on_grid_aux[k]; + } + } + int n_not_ongrid = 0; + for (int k = 0; k < 2; ++k) if (!is_on_grid[k]) ++n_not_ongrid; + if (n_not_ongrid > 0 && scale > 0) { + float id = 1/scale; + for (int k = 0; k < 2; ++k) { + if (is_on_grid[k]) continue; + uint16_t u = 0; + for (int i = 0; i < 8; ++i) { + int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); + l = MAX(0, MIN(kMaxQ-1, l)); + u |= (l << 2*i); + L[8*k + i] = l; + } + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, scale, L + 8*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 16; ++i) { + float w = weight[i]; + float q = 2*L[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0) scale = sumqx/sumq2; + } + if (scale < 0) { + scale = -scale; + for (int k = 0; k < 2; ++k) block_signs[k] = (~block_signs[k]) & 127; + } + for (int k = 0; k < 2; ++k) { + uint16_t u = 0; + for (int i = 0; i < 8; ++i) u |= (L[8*k+i] << 2*i); + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + printf("Oops: found point %u not on grid:", u); + for (int i = 0; i < 8; ++i) printf(" %d", L[8*k+i]); + printf("\n"); + GGML_ABORT("fatal error"); + } + q2[2*ib+k] = grid_index | (block_signs[k] << 9); + } + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + memset(y[ibl].qs, 0, QK_K/4); + continue; + } + + float d = max_scale/31; + y[ibl].d = GGML_FP32_TO_FP16(d); + float id = 1/d; + for (int ib = 0; ib < QK_K/16; ++ib) { + int l = nearest_int(0.5f*(id*scales[ib]-1)); + l = MAX(0, MIN(15, l)); + if (ib%2 == 0) y[ibl].scales[ib/2] = l; + else y[ibl].scales[ib/2] |= (l << 4); + } + memcpy(y[ibl].qs, q2, QK_K/4); + + } +} + +size_t quantize_iq2_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq2_xxs_impl(src, qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += nblock*sizeof(block_iq2_xxs); + } + return nrow * nblock * sizeof(block_iq2_xxs); +} + +size_t quantize_iq2_xs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq2_xs_impl(src, qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += nblock*sizeof(block_iq2_xs); + } + return nrow * nblock * sizeof(block_iq2_xs); +} + +// +// ============================================= 3-bit using D4 lattice +// + +typedef struct { + uint32_t * grid; + int * map; + uint16_t * neighbours; +} iq3_entry_t; + +static iq3_entry_t iq3_data[2] = { + {NULL, NULL, NULL}, + {NULL, NULL, NULL}, +}; + +static inline int iq3_data_index(int grid_size) { + (void)grid_size; + GGML_ASSERT(grid_size == 256 || grid_size == 512); + return grid_size == 256 ? 0 : 1; +} + +static int iq3_compare_func(const void * left, const void * right) { + const int * l = (const int *)left; + const int * r = (const int *)right; + return l[0] < r[0] ? -1 : l[0] > r[0] ? 1 : l[1] < r[1] ? -1 : l[1] > r[1] ? 1 : 0; +} + +void iq3xs_init_impl(int grid_size) { + const int gindex = iq3_data_index(grid_size); + if (iq3_data[gindex].grid) { + return; + } + static const uint16_t kgrid_256[256] = { + 0, 2, 4, 9, 11, 15, 16, 18, 25, 34, 59, 61, 65, 67, 72, 74, + 81, 85, 88, 90, 97, 108, 120, 128, 130, 132, 137, 144, 146, 153, 155, 159, + 169, 175, 189, 193, 199, 200, 202, 213, 248, 267, 287, 292, 303, 315, 317, 321, + 327, 346, 362, 413, 436, 456, 460, 462, 483, 497, 513, 515, 520, 522, 529, 531, + 536, 538, 540, 551, 552, 576, 578, 585, 592, 594, 641, 643, 648, 650, 657, 664, + 698, 704, 706, 720, 729, 742, 758, 769, 773, 808, 848, 852, 870, 889, 901, 978, + 992, 1024, 1026, 1033, 1035, 1040, 1042, 1046, 1049, 1058, 1089, 1091, 1093, 1096, 1098, 1105, + 1112, 1139, 1143, 1144, 1152, 1154, 1161, 1167, 1168, 1170, 1183, 1184, 1197, 1217, 1224, 1228, + 1272, 1276, 1309, 1323, 1347, 1367, 1377, 1404, 1473, 1475, 1486, 1509, 1537, 1544, 1546, 1553, + 1555, 1576, 1589, 1594, 1600, 1602, 1616, 1625, 1636, 1638, 1665, 1667, 1672, 1685, 1706, 1722, + 1737, 1755, 1816, 1831, 1850, 1856, 1862, 1874, 1901, 1932, 1950, 1971, 2011, 2032, 2052, 2063, + 2077, 2079, 2091, 2095, 2172, 2192, 2207, 2208, 2224, 2230, 2247, 2277, 2308, 2345, 2356, 2389, + 2403, 2424, 2501, 2504, 2506, 2520, 2570, 2593, 2616, 2624, 2630, 2646, 2669, 2700, 2714, 2746, + 2754, 2795, 2824, 2835, 2839, 2874, 2882, 2905, 2984, 3028, 3042, 3092, 3108, 3110, 3124, 3153, + 3185, 3215, 3252, 3288, 3294, 3364, 3397, 3434, 3483, 3523, 3537, 3587, 3589, 3591, 3592, 3610, + 3626, 3670, 3680, 3722, 3749, 3754, 3776, 3789, 3803, 3824, 3857, 3873, 3904, 3906, 3924, 3992, + }; + static const uint16_t kgrid_512[512] = { + 0, 1, 2, 5, 7, 8, 9, 10, 12, 14, 16, 17, 21, 27, 32, 34, + 37, 39, 41, 43, 48, 50, 57, 60, 63, 64, 65, 66, 68, 72, 73, 77, + 80, 83, 87, 89, 93, 100, 113, 117, 122, 128, 129, 133, 135, 136, 139, 142, + 145, 149, 152, 156, 162, 165, 167, 169, 171, 184, 187, 195, 201, 205, 208, 210, + 217, 219, 222, 228, 232, 234, 247, 249, 253, 256, 267, 271, 273, 276, 282, 288, + 291, 297, 312, 322, 324, 336, 338, 342, 347, 353, 357, 359, 374, 379, 390, 393, + 395, 409, 426, 441, 448, 450, 452, 464, 466, 470, 475, 488, 492, 512, 513, 514, + 516, 520, 521, 523, 525, 527, 528, 530, 537, 540, 542, 556, 558, 561, 570, 576, + 577, 579, 582, 584, 588, 593, 600, 603, 609, 616, 618, 632, 638, 640, 650, 653, + 655, 656, 660, 666, 672, 675, 685, 688, 698, 705, 708, 711, 712, 715, 721, 727, + 728, 732, 737, 754, 760, 771, 773, 778, 780, 793, 795, 802, 806, 808, 812, 833, + 840, 843, 849, 856, 858, 873, 912, 916, 919, 932, 934, 961, 963, 968, 970, 977, + 989, 993, 1010, 1016, 1024, 1025, 1027, 1029, 1031, 1032, 1034, 1036, 1038, 1041, 1043, 1047, + 1048, 1050, 1057, 1059, 1061, 1064, 1066, 1079, 1080, 1083, 1085, 1088, 1090, 1096, 1099, 1103, + 1106, 1109, 1113, 1116, 1122, 1129, 1153, 1156, 1159, 1169, 1171, 1176, 1183, 1185, 1195, 1199, + 1209, 1212, 1216, 1218, 1221, 1225, 1234, 1236, 1241, 1243, 1250, 1256, 1270, 1281, 1287, 1296, + 1299, 1306, 1309, 1313, 1338, 1341, 1348, 1353, 1362, 1375, 1376, 1387, 1400, 1408, 1410, 1415, + 1425, 1453, 1457, 1477, 1481, 1494, 1496, 1507, 1512, 1538, 1545, 1547, 1549, 1551, 1554, 1561, + 1563, 1565, 1570, 1572, 1575, 1577, 1587, 1593, 1601, 1603, 1605, 1612, 1617, 1619, 1632, 1648, + 1658, 1662, 1664, 1674, 1680, 1690, 1692, 1704, 1729, 1736, 1740, 1745, 1747, 1751, 1752, 1761, + 1763, 1767, 1773, 1787, 1795, 1801, 1806, 1810, 1817, 1834, 1840, 1844, 1857, 1864, 1866, 1877, + 1882, 1892, 1902, 1915, 1934, 1953, 1985, 1987, 2000, 2002, 2013, 2048, 2052, 2058, 2064, 2068, + 2071, 2074, 2081, 2088, 2104, 2114, 2119, 2121, 2123, 2130, 2136, 2141, 2147, 2153, 2157, 2177, + 2179, 2184, 2189, 2193, 2203, 2208, 2223, 2226, 2232, 2244, 2249, 2251, 2256, 2258, 2265, 2269, + 2304, 2306, 2324, 2335, 2336, 2361, 2373, 2375, 2385, 2418, 2443, 2460, 2480, 2504, 2509, 2520, + 2531, 2537, 2562, 2568, 2572, 2578, 2592, 2596, 2599, 2602, 2614, 2620, 2625, 2627, 2629, 2634, + 2641, 2650, 2682, 2688, 2697, 2707, 2712, 2718, 2731, 2754, 2759, 2760, 2775, 2788, 2793, 2805, + 2811, 2817, 2820, 2832, 2842, 2854, 2890, 2902, 2921, 2923, 2978, 3010, 3012, 3026, 3081, 3083, + 3085, 3097, 3099, 3120, 3136, 3152, 3159, 3188, 3210, 3228, 3234, 3245, 3250, 3256, 3264, 3276, + 3281, 3296, 3349, 3363, 3378, 3392, 3395, 3420, 3440, 3461, 3488, 3529, 3531, 3584, 3588, 3591, + 3600, 3602, 3614, 3616, 3628, 3634, 3650, 3657, 3668, 3683, 3685, 3713, 3716, 3720, 3726, 3729, + 3736, 3753, 3778, 3802, 3805, 3819, 3841, 3845, 3851, 3856, 3880, 3922, 3938, 3970, 3993, 4032, + }; + + const int kmap_size = 4096; + const int nwant = grid_size == 256 ? 2 : 3; + const uint16_t * kgrid = grid_size == 256 ? kgrid_256 : kgrid_512; + uint32_t * kgrid_q3xs; + int * kmap_q3xs; + uint16_t * kneighbors_q3xs; + + //printf("================================================================= %s(grid_size = %d)\n", __func__, grid_size); + uint32_t * the_grid = (uint32_t *)malloc(grid_size*sizeof(uint32_t)); + for (int k = 0; k < grid_size; ++k) { + int8_t * pos = (int8_t *)(the_grid + k); + for (int i = 0; i < 4; ++i) { + int l = (kgrid[k] >> 3*i) & 0x7; + pos[i] = 2*l + 1; + } + } + kgrid_q3xs = the_grid; + iq3_data[gindex].grid = the_grid; + kmap_q3xs = (int *)malloc(kmap_size*sizeof(int)); + iq3_data[gindex].map = kmap_q3xs; + for (int i = 0; i < kmap_size; ++i) kmap_q3xs[i] = -1; + uint32_t aux32; + uint8_t * aux8 = (uint8_t *)&aux32; + for (int i = 0; i < grid_size; ++i) { + aux32 = kgrid_q3xs[i]; + uint16_t index = 0; + for (int k=0; k<4; ++k) { + uint16_t q = (aux8[k] - 1)/2; + index |= (q << 3*k); + } + kmap_q3xs[index] = i; + } + int8_t pos[4]; + int * dist2 = (int *)malloc(2*grid_size*sizeof(int)); + int num_neighbors = 0, num_not_in_map = 0; + for (int i = 0; i < kmap_size; ++i) { + if (kmap_q3xs[i] >= 0) continue; + ++num_not_in_map; + for (int k = 0; k < 4; ++k) { + int l = (i >> 3*k) & 0x7; + pos[k] = 2*l + 1; + } + for (int j = 0; j < grid_size; ++j) { + const int8_t * pg = (const int8_t *)(kgrid_q3xs + j); + int d2 = 0; + for (int k = 0; k < 4; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]); + dist2[2*j+0] = d2; + dist2[2*j+1] = j; + } + qsort(dist2, grid_size, 2*sizeof(int), iq3_compare_func); + int n = 0; int d2 = dist2[0]; + int nhave = 1; + for (int j = 0; j < grid_size; ++j) { + if (dist2[2*j] > d2) { + if (nhave == nwant) break; + d2 = dist2[2*j]; + ++nhave; + } + ++n; + } + num_neighbors += n; + } + //printf("%s: %d neighbours in total\n", __func__, num_neighbors); + kneighbors_q3xs = (uint16_t *)malloc((num_neighbors + num_not_in_map)*sizeof(uint16_t)); + iq3_data[gindex].neighbours = kneighbors_q3xs; + int counter = 0; + for (int i = 0; i < kmap_size; ++i) { + if (kmap_q3xs[i] >= 0) continue; + for (int k = 0; k < 4; ++k) { + int l = (i >> 3*k) & 0x7; + pos[k] = 2*l + 1; + } + for (int j = 0; j < grid_size; ++j) { + const int8_t * pg = (const int8_t *)(kgrid_q3xs + j); + int d2 = 0; + for (int k = 0; k < 4; ++k) d2 += (pg[k] - pos[k])*(pg[k] - pos[k]); + dist2[2*j+0] = d2; + dist2[2*j+1] = j; + } + qsort(dist2, grid_size, 2*sizeof(int), iq3_compare_func); + kmap_q3xs[i] = -(counter + 1); + int d2 = dist2[0]; + uint16_t * start = &kneighbors_q3xs[counter++]; + int n = 0, nhave = 1; + for (int j = 0; j < grid_size; ++j) { + if (dist2[2*j] > d2) { + if (nhave == nwant) break; + d2 = dist2[2*j]; + ++nhave; + } + kneighbors_q3xs[counter++] = dist2[2*j+1]; + ++n; + } + *start = n; + } + free(dist2); +} + +void iq3xs_free_impl(int grid_size) { + GGML_ASSERT(grid_size == 256 || grid_size == 512); + const int gindex = iq3_data_index(grid_size); + if (iq3_data[gindex].grid) { + free(iq3_data[gindex].grid); iq3_data[gindex].grid = NULL; + free(iq3_data[gindex].map); iq3_data[gindex].map = NULL; + free(iq3_data[gindex].neighbours); iq3_data[gindex].neighbours = NULL; + } +} + +static int iq3_find_best_neighbour(const uint16_t * GGML_RESTRICT neighbours, const uint32_t * GGML_RESTRICT grid, + const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float scale, int8_t * GGML_RESTRICT L) { + int num_neighbors = neighbours[0]; + GGML_ASSERT(num_neighbors > 0); + float best_d2 = FLT_MAX; + int grid_index = -1; + for (int j = 1; j <= num_neighbors; ++j) { + const int8_t * pg = (const int8_t *)(grid + neighbours[j]); + float d2 = 0; + for (int i = 0; i < 4; ++i) { + float q = pg[i]; + float diff = scale*q - xval[i]; + d2 += weight[i]*diff*diff; + } + if (d2 < best_d2) { + best_d2 = d2; grid_index = neighbours[j]; + } + } + GGML_ASSERT(grid_index >= 0); + const int8_t * pg = (const int8_t *)(grid + grid_index); + for (int i = 0; i < 4; ++i) L[i] = (pg[i] - 1)/2; + return grid_index; +} + +static void quantize_row_iq3_xxs_impl(int grid_size, const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, + const float * GGML_RESTRICT quant_weights) { + + const int gindex = iq3_data_index(grid_size); + + const uint32_t * kgrid_q3xs = iq3_data[gindex].grid; + const int * kmap_q3xs = iq3_data[gindex].map; + const uint16_t * kneighbors_q3xs = iq3_data[gindex].neighbours; + + //GGML_ASSERT(quant_weights && "missing quantization weights"); + GGML_ASSERT(kgrid_q3xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kmap_q3xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q3xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + const int kMaxQ = 8; + + const int64_t nbl = n/QK_K; + + ggml_fp16_t * dh; + uint8_t * qs; + int block_size; + if (grid_size == 256) { + block_iq3_xxs * y = vy; + dh = &y->d; + qs = y->qs; + block_size = sizeof(block_iq3_xxs); + } else { + block_iq3_s * y = vy; + dh = &y->d; + qs = y->qs; + block_size = sizeof(block_iq3_s); + } + int quant_size = block_size - sizeof(ggml_fp16_t); + + float scales[QK_K/32]; + float weight[32]; + float xval[32]; + int8_t L[32]; + int8_t Laux[32]; + float waux[32]; + bool is_on_grid[8]; + bool is_on_grid_aux[8]; + uint8_t block_signs[8]; + uint8_t q3[3*(QK_K/8)+QK_K/32]; + uint32_t * scales_and_signs = (uint32_t *)(q3 + QK_K/4); + uint8_t * qh = q3 + 3*(QK_K/8); + + for (int ibl = 0; ibl < nbl; ++ibl) { + + dh[0] = GGML_FP32_TO_FP16(0.f); + memset(q3, 0, 3*QK_K/8+QK_K/32); + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = 2*sumx2/QK_K; + + for (int ib = 0; ib < QK_K/32; ++ib) { + const float * xb = xbl + 32*ib; + if (quant_weights) { + const float * qw = quant_weights + QK_K*ibl + 32*ib; + for (int i = 0; i < 32; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + } else { + for (int i = 0; i < 32; ++i) weight[i] = xb[i]*xb[i]; + } + for (int i = 0; i < 32; ++i) waux[i] = sqrtf(weight[i]); + for (int k = 0; k < 4; ++k) { + int nflip = 0; + uint8_t s = 0; + for (int i = 0; i < 8; ++i) { + if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; + else { + xval[8*k + i] = -xb[8*k + i]; ++nflip; s |= (1 << i); + } + } + if (nflip%2) { + int imin = 0; float min = weight[8*k+imin]*xb[8*k+imin]*xb[8*k+imin]; + for (int i = 1; i < 8; ++i) { + float ax = weight[8*k+i]*xb[8*k+i]*xb[8*k+i]; + if (ax < min) { + min = ax; imin = i; + } + } + xval[8*k+imin] = -xval[8*k+imin]; + s ^= (1 << imin); + } + block_signs[k] = s & 127; + } + float max = xval[0]; + for (int i = 1; i < 32; ++i) max = MAX(max, xval[i]); + memset(L, 0, 32); + if (max < GROUP_MAX_EPS_IQ3_XXS) { + scales[ib] = 0; + continue; + } + float best = 0; + float scale = max/(2*kMaxQ-1); + for (int k = 0; k < 8; ++k) is_on_grid[k] = true; + for (int is = -15; is <= 15; ++is) { + float id = (2*kMaxQ-1+is*0.2f)/max; + float this_scale = 1/id; + for (int k = 0; k < 8; ++k) { + for (int i = 0; i < 4; ++i) { + int l = nearest_int(0.5f*(id*xval[4*k+i]-1)); + Laux[4*k+i] = MAX(0, MIN(kMaxQ-1, l)); + } + uint16_t u = 0; + for (int i = 0; i < 4; ++i) u |= (Laux[4*k+i] << 3*i); + int grid_index = kmap_q3xs[u]; + is_on_grid_aux[k] = true; + if (grid_index < 0) { + is_on_grid_aux[k] = false; + const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1; + grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, this_scale, Laux + 4*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 32; ++i) { + float w = weight[i]; + float q = 2*Laux[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + scale = sumqx/sumq2; best = scale*sumqx; + for (int i = 0; i < 32; ++i) L[i] = Laux[i]; + for (int k = 0; k < 8; ++k) is_on_grid[k] = is_on_grid_aux[k]; + } + } + int n_not_ongrid = 0; + for (int k = 0; k < 8; ++k) if (!is_on_grid[k]) ++n_not_ongrid; + if (n_not_ongrid > 0 && scale > 0) { + float id = 1/scale; + for (int k = 0; k < 8; ++k) { + if (is_on_grid[k]) continue; + uint16_t u = 0; + for (int i = 0; i < 4; ++i) { + int l = nearest_int(0.5f*(id*xval[4*k+i]-1)); + l = MAX(0, MIN(kMaxQ-1, l)); + u |= (l << 3*i); + } + int grid_index = kmap_q3xs[u]; + if (grid_index < 0) { + const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1; + grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, scale, L + 4*k); + } + const int8_t * pg = (const int8_t *)(kgrid_q3xs + grid_index); + for (int i = 0; i < 4; ++i) L[4*k+i] = (pg[i] - 1)/2; + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 32; ++i) { + float w = weight[i]; + float q = 2*L[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0) scale = sumqx/sumq2; + } + if (scale < 0) { + // This should never happen, but just in case, flip scale so that it is positive (we use uint's to encode the scale) + // and correspondingly flip quant signs. + scale = -scale; + for (int k = 0; k < 4; ++k) block_signs[k] = (~block_signs[k]) & 127; + } + for (int k = 0; k < 8; ++k) { + uint16_t u = 0; + for (int i = 0; i < 4; ++i) u |= (L[4*k+i] << 3*i); + int grid_index = kmap_q3xs[u]; + if (grid_index < 0) { + printf("Oops: found point %u not on grid:", u); + for (int i = 0; i < 4; ++i) printf(" %d", L[4*k+i]); + printf("\n"); + GGML_ABORT("fatal error"); + } + if (grid_size == 256) { + q3[8*ib+k] = grid_index; + } else { + q3[8*ib+k] = grid_index & 255; + qh[ib] |= ((grid_index >> 8) << k); + } + + } + scales_and_signs[ib] = block_signs[0] | (block_signs[1] << 7) | (block_signs[2] << 14) | (block_signs[3] << 21); + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + memset(qs, 0, quant_size); + dh += block_size/sizeof(ggml_fp16_t); + qs += block_size; + continue; + } + + float d = max_scale/31; + dh[0] = GGML_FP32_TO_FP16(d * 1.0125f); // small improvement via this fudge factor + float id = 1/d; + for (int ib = 0; ib < QK_K/32; ++ib) { + int l = nearest_int(0.5f*(id*scales[ib]-1)); + l = MAX(0, MIN(15, l)); + scales_and_signs[ib] |= ((uint32_t)l << 28); + } + memcpy(qs, q3, quant_size); + + dh += block_size/sizeof(ggml_fp16_t); + qs += block_size; + + } +} + +size_t quantize_iq3_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq3_xxs_impl(256, src, qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += nblock*sizeof(block_iq3_xxs); + } + return nrow * nblock * sizeof(block_iq3_xxs); +} + +void quantize_row_iq3_xxs_ref(const float * GGML_RESTRICT x, block_iq3_xxs * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_row_iq3_xxs_impl(256, x, y, k, NULL); +} + +static void quantize_row_iq3_s_impl(int block_size, const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int n, + const float * GGML_RESTRICT quant_weights, + float * scales, + float * weight, + float * xval, + int8_t * L, + int8_t * Laux, + float * waux, + bool * is_on_grid, + bool * is_on_grid_aux, + uint8_t * block_signs) { + + const int gindex = iq3_data_index(512); + + const uint32_t * kgrid_q3xs = iq3_data[gindex].grid; + const int * kmap_q3xs = iq3_data[gindex].map; + const uint16_t * kneighbors_q3xs = iq3_data[gindex].neighbours; + + //GGML_ASSERT(quant_weights && "missing quantization weights"); + GGML_ASSERT(kgrid_q3xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kmap_q3xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q3xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + const int kMaxQ = 8; + + const int64_t nbl = n/QK_K; + + block_iq3_s * y = vy; + + const int bs4 = block_size/4; + const int bs8 = block_size/8; + + for (int ibl = 0; ibl < nbl; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq3_s)); + y[ibl].d = GGML_FP32_TO_FP16(0.f); + + uint8_t * qs = y[ibl].qs; + uint8_t * qh = y[ibl].qh; + uint8_t * signs = y[ibl].signs; + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = 2*sumx2/QK_K; + + for (int ib = 0; ib < QK_K/block_size; ++ib) { + const float * xb = xbl + block_size*ib; + if (quant_weights) { + const float * qw = quant_weights + QK_K*ibl + block_size*ib; + for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + } else { + for (int i = 0; i < block_size; ++i) weight[i] = xb[i]*xb[i]; + } + for (int i = 0; i < block_size; ++i) waux[i] = sqrtf(weight[i]); + for (int k = 0; k < bs8; ++k) { + uint8_t s = 0; + for (int i = 0; i < 8; ++i) { + if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; + else { + xval[8*k + i] = -xb[8*k + i]; s |= (1 << i); + } + } + block_signs[k] = s; + } + float max = xval[0]; + for (int i = 1; i < block_size; ++i) max = MAX(max, xval[i]); + memset(L, 0, block_size); + if (!max) { + scales[ib] = 0; + continue; + } + float best = 0; + float scale = max/(2*kMaxQ-1); + for (int k = 0; k < bs4; ++k) is_on_grid[k] = false; + for (int is = -9; is <= 9; ++is) { + float id = (2*kMaxQ-1+is*0.2f)/max; + float this_scale = 1/id; + for (int k = 0; k < bs4; ++k) { + for (int i = 0; i < 4; ++i) { + int l = nearest_int(0.5f*(id*xval[4*k+i]-1)); + Laux[4*k+i] = MAX(0, MIN(kMaxQ-1, l)); + } + uint16_t u = 0; + for (int i = 0; i < 4; ++i) u |= (Laux[4*k+i] << 3*i); + int grid_index = kmap_q3xs[u]; + is_on_grid_aux[k] = true; + if (grid_index < 0) { + is_on_grid_aux[k] = false; + const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1; + grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, this_scale, Laux + 4*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < block_size; ++i) { + float w = weight[i]; + float q = 2*Laux[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + scale = sumqx/sumq2; best = scale*sumqx; + for (int i = 0; i < block_size; ++i) L[i] = Laux[i]; + for (int k = 0; k < bs4; ++k) is_on_grid[k] = is_on_grid_aux[k]; + } + } + int n_not_ongrid = 0; + for (int k = 0; k < bs4; ++k) if (!is_on_grid[k]) ++n_not_ongrid; + if (n_not_ongrid > 0 && scale > 0) { + float id = 1/scale; + for (int k = 0; k < bs4; ++k) { + //if (is_on_grid[k]) continue; + uint16_t u = 0; + for (int i = 0; i < 4; ++i) { + int l = nearest_int(0.5f*(id*xval[4*k+i]-1)); + l = MAX(0, MIN(kMaxQ-1, l)); + u |= (l << 3*i); + } + int grid_index = kmap_q3xs[u]; + if (grid_index < 0) { + const uint16_t * neighbours = kneighbors_q3xs - kmap_q3xs[u] - 1; + grid_index = iq3_find_best_neighbour(neighbours, kgrid_q3xs, xval + 4*k, waux + 4*k, scale, L + 4*k); + } + const int8_t * pg = (const int8_t *)(kgrid_q3xs + grid_index); + for (int i = 0; i < 4; ++i) L[4*k+i] = (pg[i] - 1)/2; + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < block_size; ++i) { + float w = weight[i]; + float q = 2*L[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0) scale = sumqx/sumq2; + } + if (scale < 0) { + // This should never happen, but just in case, flip scale so that it is positive (we use uint's to encode the scale) + // and correspondingly flip quant signs. + scale = -scale; + for (int k = 0; k < bs8; ++k) block_signs[k] = ~block_signs[k]; + } + for (int k = 0; k < bs4; ++k) { + uint16_t u = 0; + for (int i = 0; i < 4; ++i) u |= (L[4*k+i] << 3*i); + int grid_index = kmap_q3xs[u]; + if (grid_index < 0) { + printf("Oops: found point %u not on grid:", u); + for (int i = 0; i < 4; ++i) printf(" %d", L[4*k+i]); + printf("\n"); + GGML_ABORT("fatal error"); + } + qs[k] = grid_index & 255; + qh[(ib*bs4+k)/8] |= ((grid_index >> 8) << ((ib*bs4+k)%8)); + } + qs += bs4; + for (int k = 0; k < bs8; ++k) signs[k] = block_signs[k]; + signs += bs8; + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + continue; + } + + float d = max_scale/31; + y[ibl].d = GGML_FP32_TO_FP16(d * 1.033f); + float id = 1/d; + for (int ib = 0; ib < QK_K/block_size; ib += 2) { + int l1 = nearest_int(0.5f*(id*scales[ib+0]-1)); + l1 = MAX(0, MIN(15, l1)); + int l2 = nearest_int(0.5f*(id*scales[ib+1]-1)); + l2 = MAX(0, MIN(15, l2)); + y[ibl].scales[ib/2] = l1 | (l2 << 4); + } + + } +} + +#define IQ3S_BLOCK_SIZE 32 +size_t quantize_iq3_s(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + int64_t nblock = n_per_row/QK_K; + float scales[QK_K/IQ3S_BLOCK_SIZE]; + float weight[IQ3S_BLOCK_SIZE]; + float xval[IQ3S_BLOCK_SIZE]; + int8_t L[IQ3S_BLOCK_SIZE]; + int8_t Laux[IQ3S_BLOCK_SIZE]; + float waux[IQ3S_BLOCK_SIZE]; + bool is_on_grid[IQ3S_BLOCK_SIZE/4]; + bool is_on_grid_aux[IQ3S_BLOCK_SIZE/4]; + uint8_t block_signs[IQ3S_BLOCK_SIZE/8]; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq3_s_impl(IQ3S_BLOCK_SIZE, src, qrow, n_per_row, quant_weights, + scales, weight, xval, L, Laux, waux, is_on_grid, is_on_grid_aux, block_signs); + src += n_per_row; + qrow += nblock*sizeof(block_iq3_s); + } + return nrow * nblock * sizeof(block_iq3_s); +} + +void quantize_row_iq3_s_ref(const float * GGML_RESTRICT x, block_iq3_s * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq3_s(x, y, 1, k, NULL); +} + + +// =================================== 1.5 bpw =================================================== + +static int iq1_find_best_neighbour(const uint16_t * GGML_RESTRICT neighbours, const uint64_t * GGML_RESTRICT grid, + const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float * scale, int8_t * GGML_RESTRICT L, int ngrid) { + int num_neighbors = neighbours[0]; + GGML_ASSERT(num_neighbors > 0); + float best_score = -FLT_MAX; + int grid_index = -1; + for (int j = 1; j <= num_neighbors; ++j) { + const int8_t * pg = (const int8_t *)(grid + neighbours[j]); + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 8; ++i) { + float q = (pg[i] - 3)/2; + float w = weight[i]; + sumqx += w*q*xval[i]; + sumq2 += w*q*q; + } + if (sumqx > 0 && sumq2 > 0 && sumqx*sumqx > best_score*sumq2) { + *scale = sumqx/sumq2; best_score = *scale * sumqx; + grid_index = neighbours[j]; + } + } + if (grid_index < 0) { + for (int i = 0; i < ngrid; ++i) { + const int8_t * grid_i = (const int8_t *)(grid + i); + float sumqx = 0, sumq2 = 0; + for (int j = 0; j < 8; ++j) { + float w = weight[j]; + float q = (grid_i[j] - 3)/2; + sumqx += w*q*xval[j]; + sumq2 += w*q*q; + } + if (sumqx > 0 && sumq2 > 0 && sumqx*sumqx > best_score*sumq2) { + *scale = sumqx/sumq2; best_score = *scale*sumqx; + grid_index = i; + } + } + } + if (grid_index < 0) { + printf("Oops, did not find grid point\n"); + printf("Have %d neighbours\n", num_neighbors); + for (int j = 1; j <= num_neighbors; ++j) { + const int8_t * pg = (const int8_t *)(grid + neighbours[j]); + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 8; ++i) { + float q = (pg[i] - 3)/2; + float w = weight[i]; + sumqx += w*q*xval[i]; + sumq2 += w*q*q; + } + printf(" neighbour %d: sumqx = %g sumq2 = %g\n", j, (double)sumqx, (double)sumq2); + } + } + GGML_ASSERT(grid_index >= 0); + //!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! + *scale *= 1.05f; // This is a fudge factor. Don't ask me why it improves the result. + //!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! + const int8_t * pg = (const int8_t *)(grid + grid_index); + for (int i = 0; i < 8; ++i) L[i] = (pg[i] - 1)/2; + return grid_index; +} + +static int iq1_find_best_neighbour2(const uint16_t * GGML_RESTRICT neighbours, const uint64_t * GGML_RESTRICT grid, + const float * GGML_RESTRICT xval, const float * GGML_RESTRICT weight, float scale, const float * GGML_RESTRICT xg, int8_t * GGML_RESTRICT L, int ngrid) { + int num_neighbors = neighbours[0]; + GGML_ASSERT(num_neighbors > 0); + float best_score = FLT_MAX; + int grid_index = -1; + for (int j = 1; j <= num_neighbors; ++j) { + const int8_t * pg = (const int8_t *)(grid + neighbours[j]); + float d2 = 0; + for (int i = 0; i < 8; ++i) { + float q = xg[(pg[i] - 1)/2]; + float w = weight[i]; + float diff = scale*q - xval[i]; + d2 += w*diff*diff; + } + if (d2 < best_score) { + best_score = d2; + grid_index = neighbours[j]; + } + } + if (grid_index < 0) { + for (int i = 0; i < ngrid; ++i) { + const int8_t * grid_i = (const int8_t *)(grid + i); + float d2 = 0; + for (int j = 0; j < 8; ++j) { + float w = weight[j]; + float q = xg[(grid_i[j] - 1)/2]; + float diff = scale*q - xval[i]; + d2 += w*diff*diff; + } + if (d2 < best_score) { + best_score = d2; + grid_index = i; + } + } + } + if (grid_index < 0) { + printf("Oops, did not find grid point\n"); + printf("Have %d neighbours\n", num_neighbors); + for (int j = 1; j <= num_neighbors; ++j) { + const int8_t * pg = (const int8_t *)(grid + neighbours[j]); + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 8; ++i) { + float q = xg[(pg[i] - 1)/2]; + float w = weight[i]; + sumqx += w*q*xval[i]; + sumq2 += w*q*q; + } + printf(" neighbour %d: sumqx = %g sumq2 = %g\n", j, (double)sumqx, (double)sumq2); + } + } + GGML_ASSERT(grid_index >= 0); + const int8_t * pg = (const int8_t *)(grid + grid_index); + for (int i = 0; i < 8; ++i) L[i] = (pg[i] - 1)/2; + return grid_index; +} + +static int iq1_sort_helper(const void * left, const void * right) { + const float * l = left; + const float * r = right; + return *l < *r ? -1 : *l > *r ? 1 : 0; +} + +#define IQ1S_BLOCK_SIZE 32 +#define IQ1M_BLOCK_SIZE 16 +static void quantize_row_iq1_s_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights, + float * scales, + float * weight, + float * sumx, + float * sumw, + float * pairs, + int8_t * L, + uint16_t * index, + int8_t * shifts) { + + const int gindex = iq2_data_index(GGML_TYPE_IQ1_S); + + const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; + const int * kmap_q2xs = iq2_data[gindex].map; + const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; + + GGML_ASSERT(quant_weights && "missing quantization weights"); + GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + block_iq1_s * y = vy; + + const int64_t nbl = n/QK_K; + + const int block_size = IQ1S_BLOCK_SIZE; + + const float x_p[3] = {-1 + IQ1S_DELTA, IQ1S_DELTA, 1 + IQ1S_DELTA}; + const float x_m[3] = {-1 - IQ1S_DELTA, -IQ1S_DELTA, 1 - IQ1S_DELTA}; + + + int * idx = (int *)(pairs + 1); + + for (int ibl = 0; ibl < nbl; ++ibl) { + + y[ibl].d = GGML_FP32_TO_FP16(0.f); + memset(y[ibl].qs, 0, QK_K/8); + memset(y[ibl].qh, 0, QK_K/16); + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = 2*sumx2/QK_K; + + for (int ib = 0; ib < QK_K/block_size; ++ib) { + const float * xb = xbl + block_size*ib; + const float * qw = quant_weights + QK_K*ibl + block_size*ib; + for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + float max = fabsf(xb[0]); + for (int i = 1; i < block_size; ++i) max = MAX(max, fabsf(xb[i])); + if (max < GROUP_MAX_EPS_IQ1_S) { + scales[ib] = 0; + shifts[ib] = 1; + memset(L, 1, block_size); + continue; + } + // Here we solve exactly the sum of squared difference (SSD) weighted minimization problem. + // With just 3 allowed quant values (-1, 0, 1), we can search exhaustively for the two + // boundaries that split the weights xb[i] into 3 groups. To do so, we sort the weights + // in ascending order, compute Si = sum[weight[j] xb[j], j = 0...i] and + // Wi = sum[weight[j], j = 0...i], and use these to quckly get get the optimum scale + // for each possible and score for each split. + for (int j = 0; j < block_size; ++j) { + pairs[2*j] = xb[j]; + idx[2*j] = j; + } + qsort(pairs, block_size, 2*sizeof(float), iq1_sort_helper); + { + sumx[0] = sumw[0] = 0; + for (int j = 0; j < block_size; ++j) { + int i = idx[2*j]; + sumx[j+1] = sumx[j] + weight[i]*xb[i]; + sumw[j+1] = sumw[j] + weight[i]; + } + } + float best_score = -FLT_MAX, scale = max; + int besti1 = -1, besti2 = -1, best_shift = 0; + for (int i1 = 0; i1 <= block_size; ++i1) { + for (int i2 = i1; i2 <= block_size; ++i2) { + float sumqx = (sumx[i1] - sumx[0])*x_p[0] + (sumx[i2] - sumx[i1])*x_p[1] + (sumx[block_size] - sumx[i2])*x_p[2]; + float sumq2 = (sumw[i1] - sumw[0])*x_p[0]*x_p[0] + (sumw[i2] - sumw[i1])*x_p[1]*x_p[1] + (sumw[block_size] - sumw[i2])*x_p[2]*x_p[2]; + if (sumq2 > 0 && sumqx*sumqx > best_score*sumq2) { + scale = sumqx/sumq2; best_score = scale*sumqx; + besti1 = i1; besti2 = i2; best_shift = 1; + } + sumqx = (sumx[i1] - sumx[0])*x_m[0] + (sumx[i2] - sumx[i1])*x_m[1] + (sumx[block_size] - sumx[i2])*x_m[2]; + sumq2 = (sumw[i1] - sumw[0])*x_m[0]*x_m[0] + (sumw[i2] - sumw[i1])*x_m[1]*x_m[1] + (sumw[block_size] - sumw[i2])*x_m[2]*x_m[2]; + if (sumq2 > 0 && sumqx*sumqx > best_score*sumq2) { + scale = sumqx/sumq2; best_score = scale*sumqx; + besti1 = i1; besti2 = i2; best_shift = -1; + } + } + } + if (besti1 < 0 || besti2 < 0 || best_shift == 0) { + scales[ib] = 0; + shifts[ib] = 1; + memset(L, 1, block_size); + continue; + } + for (int j = 0; j < besti1; ++j) L[idx[2*j]] = 0; + for (int j = besti1; j < besti2; ++j) L[idx[2*j]] = 1; + for (int j = besti2; j < block_size; ++j) L[idx[2*j]] = 2; + if (scale < 0) { + for (int j = 0; j < block_size; ++j) L[j] = 2 - L[j]; + scale = -scale; best_shift = -best_shift; + } + bool all_on_grid = true; + const float * xx = best_shift == 1 ? x_p : x_m; + for (int k = 0; k < block_size/8; ++k) { + uint16_t u = 0; + for (int j = 0; j < 8; ++j) u |= (L[8*k+j] << 2*j); + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + all_on_grid = false; + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq1_find_best_neighbour2(neighbours, kgrid_q2xs, xb + 8*k, weight + 8*k, scale, xx, L + 8*k, NGRID_IQ1S); + GGML_ASSERT(grid_index >= 0); + } + index[k] = grid_index; + } + if (!all_on_grid) { + float sumqx = 0, sumq2 = 0; + for (int k = 0; k < block_size/8; ++k) { + const int8_t * pg = (const int8_t *)(kgrid_q2xs + index[k]); + for (int j = 0; j < 8; ++j) { + float w = weight[8*k + j]; + float q = xx[(pg[j] - 1)/2]; + sumqx += w*q*xb[8*k+j]; + sumq2 += w*q*q; + } + } + if (sumqx > 0 && sumq2 > 0) scale = sumqx/sumq2; + } + uint16_t h = 0; + for (int k = 0; k < block_size/8; ++k) { + y[ibl].qs[(block_size/8)*ib + k] = index[k] & 255; + h |= (index[k] >> 8) << 3*k; + } + y[ibl].qh[ib] = h; + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + shifts[ib] = best_shift; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + continue; + } + + float d = max_scale/15; + y[ibl].d = GGML_FP32_TO_FP16(d*1.125f); // 1.125f is another fudge factor. Don't ask me why it is needed. + float id = 1/d; + for (int ib = 0; ib < QK_K/block_size; ++ib) { + int l = nearest_int(0.5f*(id*scales[ib]-1)); + l = MAX(0, MIN(7, l)); + if (shifts[ib] == -1) l |= 8; + y[ibl].qh[ib] |= (l << 12); + } + } +} + +size_t quantize_iq1_s(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + float scales[QK_K/IQ1S_BLOCK_SIZE]; + float weight[IQ1S_BLOCK_SIZE]; + int8_t L[IQ1S_BLOCK_SIZE]; + float sumx[IQ1S_BLOCK_SIZE+1]; + float sumw[IQ1S_BLOCK_SIZE+1]; + float pairs[2*IQ1S_BLOCK_SIZE]; + uint16_t index[IQ1S_BLOCK_SIZE/8]; + int8_t shifts[QK_K/IQ1S_BLOCK_SIZE]; + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq1_s_impl(src, qrow, n_per_row, quant_weights, scales, weight, sumx, sumw, pairs, L, index, shifts); + src += n_per_row; + qrow += nblock*sizeof(block_iq1_s); + } + return nrow * nblock * sizeof(block_iq1_s); +} + +static void quantize_row_iq1_m_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights, + float * scales, + float * weight, + float * pairs, + int8_t * L, + uint16_t * index, + int8_t * shifts) { + + const int gindex = iq2_data_index(GGML_TYPE_IQ1_M); + + const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; + const int * kmap_q2xs = iq2_data[gindex].map; + const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; + + //GGML_ASSERT(quant_weights && "missing quantization weights"); + GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + block_iq1_m * y = vy; + + const int64_t nbl = n/QK_K; + + const int block_size = IQ1M_BLOCK_SIZE; + + const float x_p[3] = {-1 + IQ1M_DELTA, IQ1M_DELTA, 1 + IQ1M_DELTA}; + const float x_m[3] = {-1 - IQ1M_DELTA, -IQ1M_DELTA, 1 - IQ1M_DELTA}; + const uint8_t masks[4] = {0x00, 0x80, 0x08, 0x88}; + + int * idx = (int *)(pairs + 1); + + float sumqx[4], sumq2[4]; + + iq1m_scale_t s; + const float * xx; + + for (int ibl = 0; ibl < nbl; ++ibl) { + memset(y[ibl].qs, 0, QK_K/8); + memset(y[ibl].qh, 0, QK_K/16); + memset(y[ibl].scales, 0, QK_K/32); + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = 2*sumx2/QK_K; + + for (int ib = 0; ib < QK_K/block_size; ++ib) { + const float * xb = xbl + block_size*ib; + if (quant_weights) { + const float * qw = quant_weights + QK_K*ibl + block_size*ib; + for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + } else { + for (int i = 0; i < block_size; ++i) weight[i] = xb[i]*xb[i]; + } + float max = fabsf(xb[0]); + for (int i = 1; i < block_size; ++i) max = MAX(max, fabsf(xb[i])); + if (max < GROUP_MAX_EPS_IQ1_M) { + scales[ib] = 0; + shifts[ib] = 0; + memset(L, 1, block_size); + continue; + } + // Here we solve exactly the sum of squared difference (SSD) weighted minimization problem. + // With just 3 allowed quant values (-1, 0, 1), we can search exhaustively for the two + // boundaries that split the weights xb[i] into 3 groups. To do so, we sort the weights + // in ascending order, compute Si = sum[weight[j] xb[j], j = 0...i] and + // Wi = sum[weight[j], j = 0...i], and use these to quckly get get the optimum scale + // for each possible and score for each split. + for (int j = 0; j < block_size; ++j) { + pairs[2*j] = xb[j]; + idx[2*j] = j; + } + qsort(pairs, block_size, 2*sizeof(float), iq1_sort_helper); + float best_score = -FLT_MAX, scale = max; + int besti1 = -1, besti2 = -1, best_k = -1; + // 0: +, + + // 1: +, - + // 2: -, + + // 3: -, - + for (int i1 = 0; i1 <= block_size; ++i1) { + for (int i2 = i1; i2 <= block_size; ++i2) { + memset(sumqx, 0, 4*sizeof(float)); + memset(sumq2, 0, 4*sizeof(float)); + for (int j = 0; j < i1; ++j) { + int i = idx[2*j]; + if (i < block_size/2) { + sumqx[0] += weight[i]*x_p[0]*xb[i]; + sumqx[1] += weight[i]*x_p[0]*xb[i]; + sumqx[2] += weight[i]*x_m[0]*xb[i]; + sumqx[3] += weight[i]*x_m[0]*xb[i]; + sumq2[0] += weight[i]*x_p[0]*x_p[0]; + sumq2[1] += weight[i]*x_p[0]*x_p[0]; + sumq2[2] += weight[i]*x_m[0]*x_m[0]; + sumq2[3] += weight[i]*x_m[0]*x_m[0]; + } else { + sumqx[0] += weight[i]*x_p[0]*xb[i]; + sumqx[2] += weight[i]*x_p[0]*xb[i]; + sumqx[1] += weight[i]*x_m[0]*xb[i]; + sumqx[3] += weight[i]*x_m[0]*xb[i]; + sumq2[0] += weight[i]*x_p[0]*x_p[0]; + sumq2[2] += weight[i]*x_p[0]*x_p[0]; + sumq2[1] += weight[i]*x_m[0]*x_m[0]; + sumq2[3] += weight[i]*x_m[0]*x_m[0]; + } + } + for (int j = i1; j < i2; ++j) { + int i = idx[2*j]; + if (i < block_size/2) { + sumqx[0] += weight[i]*x_p[1]*xb[i]; + sumqx[1] += weight[i]*x_p[1]*xb[i]; + sumqx[2] += weight[i]*x_m[1]*xb[i]; + sumqx[3] += weight[i]*x_m[1]*xb[i]; + sumq2[0] += weight[i]*x_p[1]*x_p[1]; + sumq2[1] += weight[i]*x_p[1]*x_p[1]; + sumq2[2] += weight[i]*x_m[1]*x_m[1]; + sumq2[3] += weight[i]*x_m[1]*x_m[1]; + } else { + sumqx[0] += weight[i]*x_p[1]*xb[i]; + sumqx[2] += weight[i]*x_p[1]*xb[i]; + sumqx[1] += weight[i]*x_m[1]*xb[i]; + sumqx[3] += weight[i]*x_m[1]*xb[i]; + sumq2[0] += weight[i]*x_p[1]*x_p[1]; + sumq2[2] += weight[i]*x_p[1]*x_p[1]; + sumq2[1] += weight[i]*x_m[1]*x_m[1]; + sumq2[3] += weight[i]*x_m[1]*x_m[1]; + } + } + for (int j = i2; j < block_size; ++j) { + int i = idx[2*j]; + if (i < block_size/2) { + sumqx[0] += weight[i]*x_p[2]*xb[i]; + sumqx[1] += weight[i]*x_p[2]*xb[i]; + sumqx[2] += weight[i]*x_m[2]*xb[i]; + sumqx[3] += weight[i]*x_m[2]*xb[i]; + sumq2[0] += weight[i]*x_p[2]*x_p[2]; + sumq2[1] += weight[i]*x_p[2]*x_p[2]; + sumq2[2] += weight[i]*x_m[2]*x_m[2]; + sumq2[3] += weight[i]*x_m[2]*x_m[2]; + } else { + sumqx[0] += weight[i]*x_p[2]*xb[i]; + sumqx[2] += weight[i]*x_p[2]*xb[i]; + sumqx[1] += weight[i]*x_m[2]*xb[i]; + sumqx[3] += weight[i]*x_m[2]*xb[i]; + sumq2[0] += weight[i]*x_p[2]*x_p[2]; + sumq2[2] += weight[i]*x_p[2]*x_p[2]; + sumq2[1] += weight[i]*x_m[2]*x_m[2]; + sumq2[3] += weight[i]*x_m[2]*x_m[2]; + } + } + for (int k = 0; k < 4; ++k) { + if (sumq2[k] > 0 && sumqx[k]*sumqx[k] > best_score*sumq2[k]) { + scale = sumqx[k]/sumq2[k]; best_score = scale*sumqx[k]; + besti1 = i1; besti2 = i2; best_k = k; + } + } + } + } + if (besti1 < 0 || besti2 < 0 || best_k < 0) { + scales[ib] = 0; + shifts[ib] = 0; + memset(L, 1, block_size); + continue; + } + for (int j = 0; j < besti1; ++j) L[idx[2*j]] = 0; + for (int j = besti1; j < besti2; ++j) L[idx[2*j]] = 1; + for (int j = besti2; j < block_size; ++j) L[idx[2*j]] = 2; + if (scale < 0) { + for (int j = 0; j < block_size; ++j) L[j] = 2 - L[j]; + scale = -scale; + best_k = best_k == 0 ? 3 : best_k == 1 ? 2 : best_k == 2 ? 1 : 0; + } + bool all_on_grid = true; + for (int k = 0; k < block_size/8; ++k) { + if (k == 0) xx = best_k < 2 ? x_p : x_m; + else xx = best_k%2 == 0 ? x_p : x_m; + uint16_t u = 0; + for (int j = 0; j < 8; ++j) u |= (L[8*k+j] << 2*j); + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + all_on_grid = false; + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq1_find_best_neighbour2(neighbours, kgrid_q2xs, xb + 8*k, weight + 8*k, scale, xx, L + 8*k, NGRID_IQ1S); + GGML_ASSERT(grid_index >= 0); + } + index[k] = grid_index; + } + if (!all_on_grid) { + float sumqx_f = 0, sumq2_f = 0; + for (int k = 0; k < block_size/8; ++k) { + if (k == 0) xx = best_k < 2 ? x_p : x_m; + else xx = best_k%2 == 0 ? x_p : x_m; + const int8_t * pg = (const int8_t *)(kgrid_q2xs + index[k]); + for (int j = 0; j < 8; ++j) { + float w = weight[8*k + j]; + float q = xx[(pg[j] - 1)/2]; + sumqx_f += w*q*xb[8*k+j]; + sumq2_f += w*q*q; + } + } + if (sumqx_f > 0 && sumq2_f > 0) scale = sumqx_f/sumq2_f; + } + y[ibl].qs[2*ib + 0] = index[0] & 255; + y[ibl].qs[2*ib + 1] = index[1] & 255; + y[ibl].qh[ib] = (index[0] >> 8) | ((index[1] >> 8) << 4); + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + shifts[ib] = best_k; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + continue; + } + + uint16_t * sc = (uint16_t *)y[ibl].scales; + float d = max_scale/15; + float id = 1/d; + float sumqx_f = 0, sumq2_f = 0; + for (int ib = 0; ib < QK_K/block_size; ++ib) { + int l = nearest_int(0.5f*(id*scales[ib+0]-1)); + l = MAX(0, MIN(7, l)); + sc[ib/4] |= (l << 3*(ib%4)); + y[ibl].qh[ib] |= masks[shifts[ib]]; + const float * xb = xbl + block_size*ib; + if (quant_weights) { + const float * qw = quant_weights + QK_K*ibl + block_size*ib; + for (int i = 0; i < block_size; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + } else { + for (int i = 0; i < block_size; ++i) weight[i] = xb[i]*xb[i]; + } + for (int k = 0; k < block_size/8; ++k) { + if (k == 0) xx = shifts[ib] < 2 ? x_p : x_m; + else xx = shifts[ib]%2 == 0 ? x_p : x_m; + const int8_t * pg = (const int8_t *)(kgrid_q2xs + y[ibl].qs[2*ib+k] + ((y[ibl].qh[ib] << (8 - 4*k)) & 0x700)); + for (int j = 0; j < 8; ++j) { + float w = weight[8*k + j]; + float q = xx[(pg[j] - 1)/2]*(2*l+1); + sumqx_f += w*q*xb[8*k+j]; + sumq2_f += w*q*q; + } + } + } + if (sumq2_f > 0) d = sumqx_f/sumq2_f; + s.f16 = GGML_FP32_TO_FP16(d*1.1125f); // 1.1125f is another fudge factor. Don't ask me why it is needed. + sc[0] |= ((s.u16 & 0x000f) << 12); + sc[1] |= ((s.u16 & 0x00f0) << 8); + sc[2] |= ((s.u16 & 0x0f00) << 4); + sc[3] |= ((s.u16 & 0xf000) << 0); + } +} + +size_t quantize_iq1_m(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + float scales[QK_K/IQ1M_BLOCK_SIZE]; + float weight[IQ1M_BLOCK_SIZE]; + int8_t L[IQ1M_BLOCK_SIZE]; + float pairs[2*IQ1M_BLOCK_SIZE]; + uint16_t index[IQ1M_BLOCK_SIZE/8]; + int8_t shifts[QK_K/IQ1M_BLOCK_SIZE]; + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq1_m_impl(src, qrow, n_per_row, quant_weights, scales, weight, pairs, L, index, shifts); + src += n_per_row; + qrow += nblock*sizeof(block_iq1_m); + } + return nrow * nblock * sizeof(block_iq1_m); +} + +// ============================ 4-bit non-linear quants + +static void quantize_row_iq4_nl_impl(const int super_block_size, const int block_size, const float * GGML_RESTRICT x, + ggml_fp16_t * dh, uint8_t * q4, uint16_t * scales_h, uint8_t * scales_l, + float * scales, float * weight, uint8_t * L, + const int8_t * values, + const float * quant_weights, + const int ntry) { + + float sigma2 = 0; + for (int j = 0; j < super_block_size; ++j) sigma2 += x[j]*x[j]; + sigma2 *= 2.f/super_block_size; + + memset(q4, 0, super_block_size/2); + dh[0] = GGML_FP32_TO_FP16(0.f); + + float max_scale = 0, amax_scale = 0; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + const float * xb = x + ib*block_size; + uint8_t * Lb = L + 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 < GROUP_MAX_EPS) { + scales[ib] = 0; + continue; + } + float d = ntry > 0 ? -max/values[0] : max/values[0]; + float id = 1/d; + float sumqx = 0, sumq2 = 0; + for (int j = 0; j < block_size; ++j) { + float al = id*xb[j]; + int l = best_index_int8(16, values, al); + Lb[j] = l; + float q = values[l]; + float w = weight[j]; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + d = sumq2 > 0 ? sumqx/sumq2 : 0.f; + float best = d*sumqx; + for (int itry = -ntry; itry <= ntry; ++itry) { + id = (itry + values[0])/max; + sumqx = sumq2 = 0; + for (int j = 0; j < block_size; ++j) { + float al = id*xb[j]; + int l = best_index_int8(16, values, al); + float q = values[l]; + float w = weight[j]; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d * sumqx; + } + } + scales[ib] = d; + float abs_d = fabsf(d); + if (abs_d > amax_scale) { + amax_scale = abs_d; max_scale = d; + } + } + + if (super_block_size/block_size > 1) { + int nb = super_block_size/block_size; + memset(scales_h, 0, ((nb+7)/8)*sizeof(uint16_t)); + float d = -max_scale/32; + dh[0] = GGML_FP32_TO_FP16(d); + float id = d ? 1/d : 0.f; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + 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; + for (int j = 0; j < block_size; ++j) { + Lb[j] = best_index_int8(16, values, idl*xb[j]); + } + l += 32; + uint8_t l_l = l & 0xf; + uint8_t l_h = l >> 4; + if (ib%2 == 0) scales_l[ib/2] = l_l; + else scales_l[ib/2] |= (l_l << 4); + scales_h[ib/8] |= (l_h << 2*(ib%8)); + } + } else { + dh[0] = GGML_FP32_TO_FP16(scales[0]); + if (ntry > 0) { + float id = scales[0] ? 1/scales[0] : 0; + for (int j = 0; j < super_block_size; ++j) { + L[j] = best_index_int8(16, values, id*x[j]); + } + } + } + + for (int i = 0; i < super_block_size/32; ++i) { + for (int j = 0; j < 16; ++j) { + q4[16*i + j] = L[32*i + j] | (L[32*i + 16 + j] << 4); + } + } +} + +size_t quantize_iq4_nl(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK4_NL == 0); + int64_t nblock = n_per_row/QK4_NL; + char * qrow = (char *)dst; + uint8_t L[QK4_NL]; + float weight[QK4_NL]; + uint16_t unused_h; + uint8_t * unused_l = NULL; + float scale; + for (int64_t row = 0; row < nrow; ++row) { + block_iq4_nl * iq4 = (block_iq4_nl *)qrow; + for (int ibl = 0; ibl < nblock; ++ibl) { + const float * qw = quant_weights ? quant_weights + QK4_NL*ibl : NULL; + quantize_row_iq4_nl_impl(QK4_NL, 32, src + QK4_NL*ibl, &iq4[ibl].d, iq4[ibl].qs, &unused_h, unused_l, + &scale, weight, L, kvalues_iq4nl, qw, 7); + } + src += n_per_row; + qrow += nblock*sizeof(block_iq4_nl); + } + return nrow * nblock * sizeof(block_iq4_nl); +} + +//void quantize_row_iq4_nl_ref(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { +void quantize_row_iq4_nl_ref(const float * GGML_RESTRICT x, block_iq4_nl * GGML_RESTRICT y, int64_t k) { + GGML_ASSERT(k%QK4_NL == 0); + int64_t nblock = k/QK4_NL; + uint8_t L[QK4_NL]; + float weight[QK4_NL]; + uint16_t unused_h; + uint8_t * unused_l = NULL; + float scale; + block_iq4_nl * iq4 = y; + for (int ibl = 0; ibl < nblock; ++ibl) { + quantize_row_iq4_nl_impl(QK4_NL, 32, x + QK4_NL*ibl, &iq4[ibl].d, iq4[ibl].qs, &unused_h, unused_l, + &scale, weight, L, kvalues_iq4nl, NULL, -1); + } +} + +size_t quantize_iq4_xs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + uint8_t L[QK_K]; + float weight[32]; + float scales[QK_K/32]; + for (int64_t row = 0; row < nrow; ++row) { + block_iq4_xs * iq4 = (block_iq4_xs *)qrow; + for (int ibl = 0; ibl < nblock; ++ibl) { + const float * qw = quant_weights ? quant_weights + QK_K*ibl : NULL; + quantize_row_iq4_nl_impl(QK_K, 32, src + QK_K*ibl, &iq4[ibl].d, iq4[ibl].qs, &iq4[ibl].scales_h, iq4[ibl].scales_l, + scales, weight, L, kvalues_iq4nl, qw, 7); + } + src += n_per_row; + qrow += nblock*sizeof(block_iq4_xs); + } + return nrow * nblock * sizeof(block_iq4_xs); +} + +void quantize_row_iq4_xs_ref(const float * GGML_RESTRICT x, block_iq4_xs * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq4_xs(x, y, 1, k, NULL); +} + +// =============================== 2.5625 bpw + +static void quantize_row_iq2_s_impl(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t n, const float * GGML_RESTRICT quant_weights) { + + const int gindex = iq2_data_index(GGML_TYPE_IQ2_S); + + const uint64_t * kgrid_q2xs = iq2_data[gindex].grid; + const int * kmap_q2xs = iq2_data[gindex].map; + const uint16_t * kneighbors_q2xs = iq2_data[gindex].neighbours; + + GGML_ASSERT(kmap_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kgrid_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(kneighbors_q2xs && "forgot to call ggml_quantize_init()?"); + GGML_ASSERT(n%QK_K == 0); + + const int kMaxQ = 3; + + const int64_t nbl = n/QK_K; + + block_iq2_s * y = vy; + + float scales[QK_K/16]; + float weight[16]; + float xval[16]; + int8_t L[16]; + int8_t Laux[16]; + float waux[16]; + bool is_on_grid[2]; + bool is_on_grid_aux[2]; + uint8_t block_signs[2]; + + for (int ibl = 0; ibl < nbl; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq2_s)); + y[ibl].d = GGML_FP32_TO_FP16(0.f); + + float max_scale = 0; + + const float * xbl = x + QK_K*ibl; + float sumx2 = 0; + for (int i = 0; i < QK_K; ++i) sumx2 += xbl[i]*xbl[i]; + float sigma2 = 2*sumx2/QK_K; + + for (int ib = 0; ib < QK_K/16; ++ib) { + const float * xb = xbl + 16*ib; + if (quant_weights) { + const float * qw = quant_weights + QK_K*ibl + 16*ib; + for (int i = 0; i < 16; ++i) weight[i] = qw[i] * sqrtf(sigma2 + xb[i]*xb[i]); + } else { + for (int i = 0; i < 16; ++i) weight[i] = 0.25f*sigma2 + xb[i]*xb[i]; + } + for (int i = 0; i < 16; ++i) waux[i] = sqrtf(weight[i]); + for (int k = 0; k < 2; ++k) { + uint8_t s = 0; + for (int i = 0; i < 8; ++i) { + if (xb[8*k + i] >= 0) xval[8*k + i] = xb[8*k + i]; + else { + xval[8*k + i] = -xb[8*k + i]; s |= (1 << i); + } + } + block_signs[k] = s; + } + float max = xval[0]; + for (int i = 1; i < 16; ++i) max = MAX(max, xval[i]); + memset(L, 0, 16); + if (max < GROUP_MAX_EPS_IQ2_S) { + scales[ib] = 0; + continue; + } + float best = 0; + float scale = max/(2*kMaxQ-1); + is_on_grid[0] = is_on_grid[1] = true; + for (int is = -9; is <= 9; ++is) { + float id = (2*kMaxQ-1+is*0.1f)/max; + float this_scale = 1/id; + for (int k = 0; k < 2; ++k) { + for (int i = 0; i < 8; ++i) { + int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); + Laux[8*k+i] = MAX(0, MIN(kMaxQ-1, l)); + } + uint16_t u = 0; + for (int i = 0; i < 8; ++i) u |= (Laux[8*k+i] << 2*i); + int grid_index = kmap_q2xs[u]; + is_on_grid_aux[k] = true; + if (grid_index < 0) { + is_on_grid_aux[k] = false; + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, this_scale, Laux + 8*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 16; ++i) { + float w = weight[i]; + float q = 2*Laux[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + scale = sumqx/sumq2; best = scale*sumqx; + for (int i = 0; i < 16; ++i) L[i] = Laux[i]; + for (int k = 0; k < 2; ++k) is_on_grid[k] = is_on_grid_aux[k]; + } + } + int n_not_ongrid = 0; + for (int k = 0; k < 2; ++k) if (!is_on_grid[k]) ++n_not_ongrid; + if (n_not_ongrid > 0 && scale > 0) { + float id = 1/scale; + for (int k = 0; k < 2; ++k) { + if (is_on_grid[k]) continue; + uint16_t u = 0; + for (int i = 0; i < 8; ++i) { + int l = nearest_int(0.5f*(id*xval[8*k+i]-1)); + l = MAX(0, MIN(kMaxQ-1, l)); + u |= (l << 2*i); + L[8*k + i] = l; + } + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + const uint16_t * neighbours = kneighbors_q2xs - kmap_q2xs[u] - 1; + grid_index = iq2_find_best_neighbour(neighbours, kgrid_q2xs, xval + 8*k, waux + 8*k, scale, L + 8*k); + } + } + float sumqx = 0, sumq2 = 0; + for (int i = 0; i < 16; ++i) { + float w = weight[i]; + float q = 2*L[i] + 1; + sumqx += w*xval[i]*q; + sumq2 += w*q*q; + } + if (sumq2 > 0) scale = sumqx/sumq2; + } + if (scale < 0) { + scale = -scale; + for (int k = 0; k < 2; ++k) block_signs[k] = ~block_signs[k]; + } + for (int k = 0; k < 2; ++k) { + uint16_t u = 0; + for (int i = 0; i < 8; ++i) u |= (L[8*k+i] << 2*i); + int grid_index = kmap_q2xs[u]; + if (grid_index < 0) { + printf("Oops: found point %u not on grid:", u); + for (int i = 0; i < 8; ++i) printf(" %d", L[8*k+i]); + printf("\n"); + GGML_ABORT("fatal error"); + } + const int i8 = 2*ib + k; + y[ibl].qs[i8] = grid_index & 255; + y[ibl].qh[i8/4] |= ((grid_index >> 8) << 2*(i8%4)); + y[ibl].qs[QK_K/8 + i8] = block_signs[k]; + } + GGML_ASSERT(scale >= 0); + scales[ib] = scale; + max_scale = MAX(max_scale, scale); + } + + if (!max_scale) { + continue; + } + + float d = max_scale/31; + y[ibl].d = GGML_FP32_TO_FP16(d * 0.9875f); + float id = 1/d; + for (int ib = 0; ib < QK_K/16; ++ib) { + int l = nearest_int(0.5f*(id*scales[ib]-1)); + l = MAX(0, MIN(15, l)); + if (ib%2 == 0) y[ibl].scales[ib/2] = l; + else y[ibl].scales[ib/2] |= (l << 4); + } + } +} + +size_t quantize_iq2_s(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_ASSERT(n_per_row%QK_K == 0); + int64_t nblock = n_per_row/QK_K; + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_iq2_s_impl(src, qrow, n_per_row, quant_weights); + src += n_per_row; + qrow += nblock*sizeof(block_iq2_s); + } + return nrow * nblock * sizeof(block_iq2_s); +} + +void quantize_row_iq2_s_ref(const float * GGML_RESTRICT x, block_iq2_s * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq2_s(x, y, 1, k, NULL); +} + +// =============================== data validation + +static bool validate_float(float f, size_t i) { + if (isinf(f)) { + fprintf(stderr, "ggml_validate_row_data: found inf value at block %zu\n", i); + return false; + } + + if (isnan(f)) { + fprintf(stderr, "ggml_validate_row_data: found nan value at block %zu\n", i); + return false; + } + + return true; +} + +#define isinf_fp16(f) ggml_isinf_fp16_rust(f) +#define isnan_fp16(f) ggml_isnan_fp16_rust(f) + +static bool validate_fp16(ggml_fp16_t f, size_t i) { + if (isinf_fp16(f)) { + fprintf(stderr, "ggml_validate_row_data: found inf value at block %zu\n", i); + return false; + } + + if (isnan_fp16(f)) { + fprintf(stderr, "ggml_validate_row_data: found nan value at block %zu\n", i); + return false; + } + + return true; +} + +static bool validate_e_e8m0(uint8_t e, size_t i) { + if (ggml_is_invalid_e8m0_rust(e)) { + fprintf(stderr, "ggml_validate_row_data: found invalid e value %d at block %zu\n", e, i); + return false; + } + + return true; +} + +#define VALIDATE_ROW_DATA_D_F16_IMPL(type, data, nb) \ + const type * q = (const type *) (data); \ + for (size_t i = 0; i < (nb); ++i) { \ + if (!validate_fp16(q[i].d, i)) { \ + return false; \ + } \ + } + +#define VALIDATE_ROW_DATA_DM_F16_IMPL(type, data, nb, d, m) \ + const type * q = (const type *) (data); \ + for (size_t i = 0; i < (nb); ++i) { \ + if (!validate_fp16(q[i].d, i) || !validate_fp16(q[i].m, i)) { \ + return false; \ + } \ + } + +#define VALIDATE_ROW_DATA_E_E8M0_IMPL(type, data, nb) \ + const type * q = (const type *) (data); \ + for (size_t i = 0; i < (nb); ++i) { \ + if (!validate_e_e8m0(q[i].e, i)) { \ + return false; \ + } \ + } + +#define VALIDATE_ROW_DATA_DVEC_F16_IMPL(type, data, nb, nr) \ + const type * q = (const type *) (data); \ + for (size_t i = 0; i < (nb); ++i) { \ + for (size_t j = 0; j < (nr); ++j) { \ + if (!validate_fp16(q[i].d[j], i)) { \ + return false; \ + } \ + } \ + } + +bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbytes) { + if (type < 0 || type >= GGML_TYPE_COUNT) { + fprintf(stderr, "%s: invalid type %d\n", __func__, type); + return false; + } + + if (nbytes % ggml_type_size(type) != 0) { + fprintf(stderr, "%s: invalid size %zu for type %s (type size = %zu)\n", __func__, nbytes, ggml_type_name(type), ggml_type_size(type)); + return false; + } + + const size_t nb = nbytes/ggml_type_size(type); + + switch (type) { + case GGML_TYPE_BF16: + { + int nans = 0; + int infs = 0; + const unsigned short * f = (const unsigned short *) data; + for (size_t i = 0; i < nb; ++i) { + nans += (f[i] & 0x7fff) > 0x7f80; + infs += (f[i] & 0x7fff) == 0x7f80; + } + if (nans) { + fprintf(stderr, "%s: found %d NaNs in row of %zu BF16 values\n", __func__, nans, nb); + return false; + } + if (infs) { + fprintf(stderr, "%s: found %d infinities in row of %zu BF16 values\n", __func__, infs, nb); + return false; + } + } break; + case GGML_TYPE_F16: + { + const ggml_fp16_t * f = (const ggml_fp16_t *) data; + size_t i = 0; +#if defined(__AVX2__) + for (; i + 15 < nb; i += 16) { + __m256i v = _mm256_loadu_si256((const __m256i *)(f + i)); + __m256i vexp = _mm256_and_si256(v, _mm256_set1_epi16(0x7c00)); + __m256i cmp = _mm256_cmpeq_epi16(vexp, _mm256_set1_epi16(0x7c00)); + int mask = _mm256_movemask_epi8(cmp); + if (mask) { + for (size_t j = 0; j < 16; ++j) { + if (!validate_fp16(f[i + j], i + j)) { + return false; + } + } + GGML_UNREACHABLE(); + } + } +#elif defined(__ARM_NEON) + for (; i + 7 < nb; i += 8) { + uint16x8_t v = vld1q_u16(f + i); + uint16x8_t vexp = vandq_u16(v, vdupq_n_u16(0x7c00)); + uint16x8_t cmp = vceqq_u16(vexp, vdupq_n_u16(0x7c00)); + uint64_t mask = vget_lane_u64(vreinterpret_u64_u8(vshrn_n_u16(cmp, 4)), 0); + if (mask) { + for (size_t j = 0; j < 8; ++j) { + if (!validate_fp16(f[i + j], i + j)) { + return false; + } + } + GGML_UNREACHABLE(); + } + } +#endif + for (; i < nb; ++i) { + if (!validate_fp16(f[i], i)) { + return false; + } + } + } break; + case GGML_TYPE_F32: + { + const float * f = (const float *) data; + size_t i = 0; +#if defined(__AVX2__) + for (; i + 7 < nb; i += 8) { + __m256i v = _mm256_loadu_si256((const __m256i *)(f + i)); + __m256i vexp = _mm256_and_si256(v, _mm256_set1_epi32(0x7f800000)); + __m256i cmp = _mm256_cmpeq_epi32(vexp, _mm256_set1_epi32(0x7f800000)); + int mask = _mm256_movemask_epi8(cmp); + if (mask) { + for (size_t j = 0; j < 8; ++j) { + if (!validate_float(f[i + j], i + j)) { + return false; + } + } + GGML_UNREACHABLE(); + } + } +#elif defined(__ARM_NEON) + for (; i + 3 < nb; i += 4) { + uint32x4_t v = vld1q_u32((const uint32_t *)f + i); + uint32x4_t vexp = vandq_u32(v, vdupq_n_u32(0x7f800000)); + uint32x4_t cmp = vceqq_u32(vexp, vdupq_n_u32(0x7f800000)); + uint64_t mask = vget_lane_u64(vreinterpret_u64_u16(vshrn_n_u32(cmp, 8)), 0); + if (mask) { + for (size_t j = 0; j < 4; ++j) { + if (!validate_float(f[i + j], i + j)) { + return false; + } + } + GGML_UNREACHABLE(); + } + } +#endif + for (; i < nb; ++i) { + if (!validate_float(f[i], i)) { + return false; + } + } + } break; + case GGML_TYPE_F64: + { + const double * f = (const double *) data; + for (size_t i = 0; i < nb; ++i) { + if (!validate_float(f[i], i)) { + return false; + } + } + } break; + case GGML_TYPE_Q1_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q1_0, data, nb); + } break; + case GGML_TYPE_Q4_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q4_0, data, nb); + } break; + case GGML_TYPE_Q4_1: + { + VALIDATE_ROW_DATA_DM_F16_IMPL(block_q4_1, data, nb, d, m); + } break; + case GGML_TYPE_Q5_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q5_0, data, nb); + } break; + case GGML_TYPE_Q5_1: + { + VALIDATE_ROW_DATA_DM_F16_IMPL(block_q5_1, data, nb, d, m); + } break; + case GGML_TYPE_Q8_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q8_0, data, nb); + } break; + case GGML_TYPE_MXFP4: + { + VALIDATE_ROW_DATA_E_E8M0_IMPL(block_mxfp4, data, nb); + } break; + case GGML_TYPE_NVFP4: + { + // UE4M3 scales are uint8_t — all byte values are valid + GGML_UNUSED(data); + GGML_UNUSED(nb); + } break; + case GGML_TYPE_Q2_K: + { + VALIDATE_ROW_DATA_DM_F16_IMPL(block_q2_K, data, nb, d, dmin); + } break; + case GGML_TYPE_Q3_K: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q3_K, data, nb); + } break; + case GGML_TYPE_Q4_K: + { + VALIDATE_ROW_DATA_DM_F16_IMPL(block_q4_K, data, nb, d, dmin); + } break; + case GGML_TYPE_Q5_K: + { + VALIDATE_ROW_DATA_DM_F16_IMPL(block_q5_K, data, nb, d, dmin); + } break; + case GGML_TYPE_Q6_K: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q6_K, data, nb); + } break; + case GGML_TYPE_Q8_K: + { + const block_q8_K * q = (const block_q8_K *) data; + for (size_t i = 0; i < nb; ++i) { + if (!validate_float(q[i].d, i)) { + return false; + } + } + } break; + case GGML_TYPE_TQ1_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_tq1_0, data, nb); + } break; + case GGML_TYPE_TQ2_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_tq2_0, data, nb); + } break; + case GGML_TYPE_IQ1_S: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq1_s, data, nb); + } break; + case GGML_TYPE_IQ1_M: + { + const block_iq1_m * q = (const block_iq1_m *) data; + for (size_t i = 0; i < nb; ++i) { + iq1m_scale_t scale; + const uint16_t * sc = (const uint16_t *)q[i].scales; + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); + if (!validate_fp16(scale.f16, i)) { + return false; + } + } + } break; + case GGML_TYPE_IQ2_XXS: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq2_xxs, data, nb); + } break; + case GGML_TYPE_IQ2_XS: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq2_xs, data, nb); + } break; + case GGML_TYPE_IQ2_S: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq2_s, data, nb); + } break; + case GGML_TYPE_IQ3_XXS: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq3_xxs, data, nb); + } break; + + case GGML_TYPE_IQ3_S: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq3_s, data, nb); + } break; + case GGML_TYPE_IQ4_XS: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq4_xs, data, nb); + } break; + case GGML_TYPE_IQ4_NL: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_iq4_nl, data, nb); + } break; + + case GGML_TYPE_I8: + case GGML_TYPE_I16: + case GGML_TYPE_I32: + case GGML_TYPE_I64: + // nothing to validate + break; + default: + { + fprintf(stderr, "%s: invalid type %d\n", __func__, type); + return false; + } + } + + return true; +} diff --git a/ggml/src/ggml.cpp b/ggml/src/ggml.cpp deleted file mode 100644 index 0d388d45536d..000000000000 --- a/ggml/src/ggml.cpp +++ /dev/null @@ -1,26 +0,0 @@ -#include "ggml-impl.h" - -#include -#include - -static std::terminate_handler previous_terminate_handler; - -GGML_NORETURN static void ggml_uncaught_exception() { - ggml_print_backtrace(); - if (previous_terminate_handler) { - previous_terminate_handler(); - } - abort(); // unreachable unless previous_terminate_handler was nullptr -} - -static bool ggml_uncaught_exception_init = []{ - const char * GGML_NO_BACKTRACE = getenv("GGML_NO_BACKTRACE"); - if (GGML_NO_BACKTRACE) { - return false; - } - const auto prev{std::get_terminate()}; - GGML_ASSERT(prev != ggml_uncaught_exception); - previous_terminate_handler = prev; - std::set_terminate(ggml_uncaught_exception); - return true; -}(); diff --git a/ggml/src/gguf.cpp b/ggml/src/gguf.cpp deleted file mode 100644 index ab3cc9748678..000000000000 --- a/ggml/src/gguf.cpp +++ /dev/null @@ -1,1556 +0,0 @@ -#include "ggml.h" -#include "ggml-backend.h" -#include "ggml-impl.h" -#include "gguf.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#define GGUF_MAX_STRING_LENGTH (1024*1024*1024) -#define GGUF_MAX_ARRAY_ELEMENTS (1024*1024*1024) - -#ifdef _WIN32 -# define gguf_ftell _ftelli64 -# define gguf_fseek _fseeki64 -#else -# define gguf_ftell ftello -# define gguf_fseek fseeko -#endif - -template -struct type_to_gguf_type; - -template <> -struct type_to_gguf_type { - static constexpr enum gguf_type value = GGUF_TYPE_UINT8; -}; - -template <> -struct type_to_gguf_type { - static constexpr enum gguf_type value = GGUF_TYPE_INT8; -}; - -template <> -struct type_to_gguf_type { - static constexpr enum gguf_type value = GGUF_TYPE_UINT16; -}; - -template <> -struct type_to_gguf_type { - static constexpr enum gguf_type value = GGUF_TYPE_INT16; -}; - -template <> -struct type_to_gguf_type { - static constexpr enum gguf_type value = GGUF_TYPE_UINT32; -}; - -template <> -struct type_to_gguf_type { - static constexpr enum gguf_type value = GGUF_TYPE_INT32; -}; - -template <> -struct type_to_gguf_type { - static constexpr enum gguf_type value = GGUF_TYPE_FLOAT32; -}; - -template <> -struct type_to_gguf_type { - static constexpr enum gguf_type value = GGUF_TYPE_BOOL; -}; - -template <> -struct type_to_gguf_type { - static constexpr enum gguf_type value = GGUF_TYPE_STRING; -}; - -template <> -struct type_to_gguf_type { - static constexpr enum gguf_type value = GGUF_TYPE_UINT64; -}; - -template <> -struct type_to_gguf_type { - static constexpr enum gguf_type value = GGUF_TYPE_INT64; -}; - -template <> -struct type_to_gguf_type { - static constexpr enum gguf_type value = GGUF_TYPE_FLOAT64; -}; - -static const std::map GGUF_TYPE_SIZE = { - {GGUF_TYPE_UINT8, sizeof(uint8_t)}, - {GGUF_TYPE_INT8, sizeof(int8_t)}, - {GGUF_TYPE_UINT16, sizeof(uint16_t)}, - {GGUF_TYPE_INT16, sizeof(int16_t)}, - {GGUF_TYPE_UINT32, sizeof(uint32_t)}, - {GGUF_TYPE_INT32, sizeof(int32_t)}, - {GGUF_TYPE_FLOAT32, sizeof(float)}, - {GGUF_TYPE_BOOL, sizeof(int8_t)}, - {GGUF_TYPE_STRING, 0}, // undefined - {GGUF_TYPE_ARRAY, 0}, // undefined - {GGUF_TYPE_UINT64, sizeof(uint64_t)}, - {GGUF_TYPE_INT64, sizeof(int64_t)}, - {GGUF_TYPE_FLOAT64, sizeof(double)}, -}; -static_assert(GGUF_TYPE_COUNT == 13, "GGUF_TYPE_COUNT != 13"); - -static const std::map GGUF_TYPE_NAME = { - {GGUF_TYPE_UINT8, "u8"}, - {GGUF_TYPE_INT8, "i8"}, - {GGUF_TYPE_UINT16, "u16"}, - {GGUF_TYPE_INT16, "i16"}, - {GGUF_TYPE_UINT32, "u32"}, - {GGUF_TYPE_INT32, "i32"}, - {GGUF_TYPE_FLOAT32, "f32"}, - {GGUF_TYPE_BOOL, "bool"}, - {GGUF_TYPE_STRING, "str"}, - {GGUF_TYPE_ARRAY, "arr"}, - {GGUF_TYPE_UINT64, "u64"}, - {GGUF_TYPE_INT64, "i64"}, - {GGUF_TYPE_FLOAT64, "f64"}, -}; -static_assert(GGUF_TYPE_COUNT == 13, "GGUF_TYPE_COUNT != 13"); - -size_t gguf_type_size(enum gguf_type type) { - auto it = GGUF_TYPE_SIZE.find(type); - return it == GGUF_TYPE_SIZE.end() ? 0 : it->second; -} - -struct gguf_kv { - std::string key; - - bool is_array; - enum gguf_type type; - - std::vector data; - std::vector data_string; - - template - gguf_kv(const std::string & key, const T value) - : key(key), is_array(false), type(type_to_gguf_type::value) { - GGML_ASSERT(!key.empty()); - data.resize(sizeof(T)); - memcpy(data.data(), &value, sizeof(T)); - } - - template - gguf_kv(const std::string & key, const std::vector & value) - : key(key), is_array(true), type(type_to_gguf_type::value) { - GGML_ASSERT(!key.empty()); - data.resize(value.size()*sizeof(T)); - for (size_t i = 0; i < value.size(); ++i) { - const T tmp = value[i]; - memcpy(data.data() + i*sizeof(T), &tmp, sizeof(T)); - } - } - - gguf_kv(const std::string & key, const std::string & value) - : key(key), is_array(false), type(GGUF_TYPE_STRING) { - GGML_ASSERT(!key.empty()); - data_string.push_back(value); - } - - gguf_kv(const std::string & key, const std::vector & value) - : key(key), is_array(true), type(GGUF_TYPE_STRING) { - GGML_ASSERT(!key.empty()); - data_string = value; - } - - const std::string & get_key() const { - return key; - } - - const enum gguf_type & get_type() const { - return type; - } - - size_t get_ne() const { - if (type == GGUF_TYPE_STRING) { - const size_t ne = data_string.size(); - GGML_ASSERT(is_array || ne == 1); - return ne; - } - const size_t type_size = gguf_type_size(type); - GGML_ASSERT(data.size() % type_size == 0); - const size_t ne = data.size() / type_size; - GGML_ASSERT(is_array || ne == 1); - return ne; - } - - template - const T & get_val(const size_t i = 0) const { - GGML_ASSERT(type_to_gguf_type::value == type); - if constexpr (std::is_same::value) { - GGML_ASSERT(data_string.size() >= i+1); - return data_string[i]; - } - const size_t type_size = gguf_type_size(type); - GGML_ASSERT(data.size() % type_size == 0); - GGML_ASSERT(data.size() >= (i+1)*type_size); - return reinterpret_cast(data.data())[i]; - } - - void cast(const enum gguf_type new_type) { - const size_t new_type_size = gguf_type_size(new_type); - GGML_ASSERT(data.size() % new_type_size == 0); - type = new_type; - } -}; - -struct gguf_tensor_info { - struct ggml_tensor t; // for holding the equivalent info - uint64_t offset; // offset from start of `data`, must be a multiple of `ALIGNMENT` -}; - -struct gguf_context { - uint32_t version = GGUF_VERSION; - - std::vector kv; - std::vector info; - - size_t alignment = GGUF_DEFAULT_ALIGNMENT; - size_t offset = 0; // offset of `data` from beginning of file - size_t size = 0; // size of `data` in bytes - - void * data = nullptr; -}; - -struct gguf_reader { - gguf_reader(FILE * file) : file(file) { - // read the remaining bytes once and update on each read - nbytes_remain = file_remain(file); - } - - // helper for remaining bytes in a file - static uint64_t file_remain(FILE * file) { - const int64_t cur = gguf_ftell(file); - if (cur < 0) { - return 0; - } - if (gguf_fseek(file, 0, SEEK_END) != 0) { - gguf_fseek(file, cur, SEEK_SET); - - return 0; - } - const int64_t end = gguf_ftell(file); - if (end < 0) { - gguf_fseek(file, cur, SEEK_SET); - - return 0; - } - gguf_fseek(file, cur, SEEK_SET); - return static_cast(end - cur); - } - - template - bool read(T & dst) const { - const size_t size = sizeof(dst); - if (nbytes_remain < size) { - return false; - } - const size_t nread = fread(&dst, 1, size, file); - nbytes_remain -= nread; - return nread == size; - } - - template - bool read(std::vector & dst, const size_t n) const { - if (n > GGUF_MAX_ARRAY_ELEMENTS) { - return false; - } - if constexpr (std::is_same::value) { - // strings are prefixed with their length, so we need to account for that - if (n > SIZE_MAX / sizeof(uint64_t)) { - return false; - } - if (nbytes_remain < n * sizeof(uint64_t)) { - return false; - } - } else { - if (n > SIZE_MAX / sizeof(T)) { - return false; - } - if (nbytes_remain < n * sizeof(T)) { - return false; - } - } - dst.resize(n); - for (size_t i = 0; i < dst.size(); ++i) { - if constexpr (std::is_same::value) { - bool tmp; - if (!read(tmp)) { - return false; - } - dst[i] = tmp; - } else { - if (!read(dst[i])) { - return false; - } - } - } - return true; - } - - bool read(bool & dst) const { - int8_t tmp = -1; - if (!read(tmp)) { - return false; - } - dst = tmp != 0; - return true; - } - - bool read(enum ggml_type & dst) const { - int32_t tmp = -1; - if (!read(tmp)) { - return false; - } - dst = ggml_type(tmp); - return true; - } - - bool read(enum gguf_type & dst) const { - int32_t tmp = -1; - if (!read(tmp)) { - return false; - } - dst = gguf_type(tmp); - return true; - } - - bool read(std::string & dst) const { - uint64_t size = 0; - if (!read(size)) { - return false; - } - if (size > GGUF_MAX_STRING_LENGTH) { - GGML_LOG_ERROR("%s: string length %" PRIu64 " exceeds maximum %" PRIu64 "\n", __func__, size, (uint64_t) GGUF_MAX_STRING_LENGTH); - return false; - } - if (size > nbytes_remain) { - GGML_LOG_ERROR("%s: string length %" PRIu64 " exceeds remaining file size %" PRIu64 " bytes\n", __func__, size, nbytes_remain); - return false; - } - dst.resize(static_cast(size)); - const size_t nread = fread(dst.data(), 1, size, file); - nbytes_remain -= nread; - return nread == size; - } - - bool read(void * dst, const size_t size) const { - if (size > nbytes_remain) { - return false; - } - const size_t nread = fread(dst, 1, size, file); - nbytes_remain -= nread; - return nread == size; - } - -private: - FILE * file; - - mutable uint64_t nbytes_remain; -}; - -struct gguf_context * gguf_init_empty(void) { - return new gguf_context; -} - -template -bool gguf_read_emplace_helper(const struct gguf_reader & gr, std::vector & kv, const std::string & key, const bool is_array, const size_t n) { - if (is_array) { - std::vector value; - try { - if (!gr.read(value, n)) { - return false; - } - } catch (std::length_error &) { - GGML_LOG_ERROR("%s: encountered length_error while reading value for key '%s'\n", __func__, key.c_str()); - return false; - } catch (std::bad_alloc &) { - GGML_LOG_ERROR("%s: encountered bad_alloc error while reading value for key '%s'\n", __func__, key.c_str()); - return false; - } - kv.emplace_back(key, value); - } else { - T value; - if (!gr.read(value)) { - return false; - } - kv.emplace_back(key, value); - } - return true; -} - -struct gguf_context * gguf_init_from_file_ptr(FILE * file, struct gguf_init_params params) { - if (!file) { - return nullptr; - } - - const struct gguf_reader gr(file); - struct gguf_context * ctx = new gguf_context; - - bool ok = true; - - // file magic - { - std::vector magic; - ok = ok && gr.read(magic, 4); - - if (!ok) { - GGML_LOG_ERROR("%s: failed to read magic\n", __func__); - gguf_free(ctx); - return nullptr; - } - - for (uint32_t i = 0; i < magic.size(); i++) { - if (magic[i] != GGUF_MAGIC[i]) { - char c0 = isprint(magic[0]) ? magic[0] : '?'; - char c1 = isprint(magic[1]) ? magic[1] : '?'; - char c2 = isprint(magic[2]) ? magic[2] : '?'; - char c3 = isprint(magic[3]) ? magic[3] : '?'; - GGML_LOG_ERROR("%s: invalid magic characters: '%c%c%c%c', expected 'GGUF'\n", __func__, c0, c1, c2, c3); - gguf_free(ctx); - return nullptr; - } - } - } - - // header - int64_t n_kv = 0; - int64_t n_tensors = 0; - - if (ok && gr.read(ctx->version)) { - if (ok && ctx->version == 0) { - GGML_LOG_ERROR("%s: bad GGUF version: %" PRIu32 "\n", __func__, ctx->version); - ok = false; - } - - /* - * bit layout is different when reading non-native endian models. - * assuming that the GGUF version is 3, the non-native endian model - * would read it as 0x30000000. we can use the AND operation against - * the last 4 hexadecimal digits to check if the model is the same - * endianness as the host system. - */ - if (ok && (ctx->version & 0x0000FFFF) == 0x00000000) { - GGML_LOG_ERROR("%s: failed to load model: this GGUF file version %" PRIu32 " is extremely large, is there a mismatch between the host and model endianness?\n", __func__, ctx->version); - ok = false; - } - - if (ok && ctx->version == 1) { - GGML_LOG_ERROR("%s: GGUFv1 is no longer supported, please use a more up-to-date version\n", __func__); - ok = false; - } - if (ok && ctx->version > GGUF_VERSION) { - GGML_LOG_ERROR("%s: this GGUF file is version %" PRIu32 " but this software only supports up to version %d\n", - __func__, ctx->version, GGUF_VERSION); - ok = false; - } - } else { - ok = false; - } - - if (ok && gr.read(n_tensors)) { - static_assert(sizeof(size_t) <= 8 && sizeof(gguf_tensor_info) >= 2, "int64_t insufficient for indexing"); - if (n_tensors < 0 || n_tensors > int64_t(SIZE_MAX/sizeof(gguf_tensor_info))) { - GGML_LOG_ERROR("%s: number of tensors is %" PRIi64 " but must be in [0, %zu]\n", - __func__, n_tensors, SIZE_MAX/sizeof(gguf_tensor_info)); - ok = false; - } - } else { - ok = false; - } - - if (ok && gr.read(n_kv)) { - static_assert(sizeof(size_t) <= 8 && sizeof(gguf_tensor_info) >= 2, "int64_t insufficient for indexing"); - if (n_kv < 0 || n_kv > int64_t(SIZE_MAX/sizeof(gguf_kv))) { - GGML_LOG_ERROR("%s: number of key value pairs is %" PRIi64 " but must be in [0, %zu]\n", - __func__, n_kv, SIZE_MAX/sizeof(gguf_kv)); - ok = false; - } - } else { - ok = false; - } - - if (!ok) { - GGML_LOG_ERROR("%s: failed to read header\n", __func__); - gguf_free(ctx); - return nullptr; - } - - // KV pairs - { - for (int64_t i = 0; ok && i < n_kv; ++i) { - std::string key; - gguf_type type = gguf_type(-1); - bool is_array = false; - uint64_t n = 1; - - try { - ok = ok && gr.read(key); - } catch (std::length_error &) { - GGML_LOG_ERROR("%s: encountered length_error while reading key %" PRIi64 "\n", __func__, i); - ok = false; - } catch (std::bad_alloc &) { - GGML_LOG_ERROR("%s: encountered bad_alloc error while reading key %" PRIi64 "\n", __func__, i); - ok = false; - } - for (size_t j = 0; ok && j < ctx->kv.size(); ++j) { - if (key == ctx->kv[j].key) { - GGML_LOG_ERROR("%s: duplicate key '%s' for tensors %zu and %" PRIi64 " \n", __func__, key.c_str(), j, i); - ok = false; - } - } - if (!ok) { - break; - } - - ok = ok && gr.read(type); - if (type == GGUF_TYPE_ARRAY) { - is_array = true; - ok = ok && gr.read(type); - ok = ok && gr.read(n); - } - if (!ok) { - break; - } - - switch (type) { - case GGUF_TYPE_UINT8: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; - case GGUF_TYPE_INT8: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; - case GGUF_TYPE_UINT16: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; - case GGUF_TYPE_INT16: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; - case GGUF_TYPE_UINT32: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; - case GGUF_TYPE_INT32: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; - case GGUF_TYPE_FLOAT32: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; - case GGUF_TYPE_BOOL: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; - case GGUF_TYPE_STRING: ok = ok && gguf_read_emplace_helper(gr, ctx->kv, key, is_array, n); break; - case GGUF_TYPE_UINT64: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; - case GGUF_TYPE_INT64: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; - case GGUF_TYPE_FLOAT64: ok = ok && gguf_read_emplace_helper (gr, ctx->kv, key, is_array, n); break; - case GGUF_TYPE_ARRAY: - default: - { - GGML_LOG_ERROR("%s: key '%s' has invalid GGUF type %d\n", __func__, key.c_str(), type); - ok = false; - } break; - } - } - - if (!ok) { - GGML_LOG_ERROR("%s: failed to read key-value pairs\n", __func__); - gguf_free(ctx); - return nullptr; - } - GGML_ASSERT(int64_t(ctx->kv.size()) == n_kv); - - const int alignment_idx = gguf_find_key(ctx, GGUF_KEY_GENERAL_ALIGNMENT); - ctx->alignment = alignment_idx == -1 ? GGUF_DEFAULT_ALIGNMENT : gguf_get_val_u32(ctx, alignment_idx); - - if (ctx->alignment == 0 || (ctx->alignment & (ctx->alignment - 1)) != 0) { - GGML_LOG_ERROR("%s: alignment %zu is not a power of 2\n", __func__, ctx->alignment); - gguf_free(ctx); - return nullptr; - } - } - - // read the tensor info - for (int64_t i = 0; ok && i < n_tensors; ++i) { - struct gguf_tensor_info info; - - // tensor name - { - std::string name; - try { - ok = ok && gr.read(name); - } catch (std::length_error &) { - GGML_LOG_ERROR("%s: encountered length_error while reading tensor name %" PRIi64 "\n", __func__, i); - ok = false; - } catch (std::bad_alloc &) { - GGML_LOG_ERROR("%s: encountered bad_alloc error while reading tensor name %" PRIi64 "\n", __func__, i); - ok = false; - } - if (name.length() >= GGML_MAX_NAME) { - GGML_LOG_ERROR("%s: tensor name %" PRIi64 " is too long: %zu >= %d\n", __func__, i, name.length(), GGML_MAX_NAME); - ok = false; - break; - } - ggml_set_name(&info.t, name.c_str()); - - // make sure there are no duplicate tensor names - for (int64_t j = 0; ok && j < i; ++j) { - if (strcmp(info.t.name, ctx->info[j].t.name) == 0) { - GGML_LOG_ERROR("%s: duplicate tensor name '%s' for tensors %" PRIi64 " and %" PRIi64 "\n", __func__, info.t.name, j, i); - ok = false; - break; - } - } - } - if (!ok) { - break; - } - - // tensor shape - { - uint32_t n_dims = 0; - ok = ok && gr.read(n_dims); - if (n_dims > GGML_MAX_DIMS) { - GGML_LOG_ERROR("%s: tensor '%s' has invalid number of dimensions: %" PRIu32 " > %" PRIu32 "\n", - __func__, info.t.name, n_dims, GGML_MAX_DIMS); - ok = false; - break; - } - for (uint32_t j = 0; ok && j < GGML_MAX_DIMS; ++j) { - info.t.ne[j] = 1; - if (j < n_dims) { - ok = ok && gr.read(info.t.ne[j]); - } - - // check that all ne are non-negative - if (info.t.ne[j] < 0) { - GGML_LOG_ERROR("%s: tensor '%s' dimension %" PRIu32 " has invalid number of elements: %" PRIi64 " < 0\n", - __func__, info.t.name, j, info.t.ne[j]); - ok = false; - break; - } - } - - // check that the total number of elements is representable - if (ok && ((INT64_MAX/info.t.ne[1] <= info.t.ne[0]) || - (INT64_MAX/info.t.ne[2] <= info.t.ne[0]*info.t.ne[1]) || - (INT64_MAX/info.t.ne[3] <= info.t.ne[0]*info.t.ne[1]*info.t.ne[2]))) { - - GGML_LOG_ERROR("%s: total number of elements in tensor '%s' with shape " - "(%" PRIi64 ", %" PRIi64 ", %" PRIi64 ", %" PRIi64 ") is >= %" PRIi64 "\n", - __func__, info.t.name, info.t.ne[0], info.t.ne[1], info.t.ne[2], info.t.ne[3], INT64_MAX); - ok = false; - break; - } - } - if (!ok) { - break; - } - - // tensor type - { - ok = ok && gr.read(info.t.type); - - // check that tensor type is within defined range - if (info.t.type < 0 || info.t.type >= GGML_TYPE_COUNT) { - GGML_LOG_ERROR("%s: tensor '%s' has invalid ggml type %d. should be in [0, %d)\n", - __func__, info.t.name, info.t.type, GGML_TYPE_COUNT); - ok = false; - break; - } - const size_t type_size = ggml_type_size(info.t.type); - const int64_t blck_size = ggml_blck_size(info.t.type); - - // check that row size is divisible by block size - if (blck_size == 0 || info.t.ne[0] % blck_size != 0) { - GGML_LOG_ERROR("%s: tensor '%s' of type %d (%s) has %" PRId64 " elements per row, " - "not a multiple of block size (%" PRId64 ")\n", - __func__, info.t.name, (int) info.t.type, ggml_type_name(info.t.type), info.t.ne[0], blck_size); - ok = false; - break; - } - - // check that the size of the tensor in bytes is representable - if (ok && uint64_t(ggml_nelements(&info.t)/ggml_blck_size(info.t.type)) > SIZE_MAX/ggml_type_size(info.t.type)) { - GGML_LOG_ERROR("%s: tensor '%s' with shape (%" PRIi64 ", %" PRIi64 ", %" PRIi64 ", %" PRIi64 ") has a size in bytes > %zu\n", - __func__, info.t.name, info.t.ne[0], info.t.ne[1], info.t.ne[2], info.t.ne[3], SIZE_MAX); - ok = false; - break; - } - - // calculate byte offsets given the tensor shape and type - info.t.nb[0] = type_size; - info.t.nb[1] = info.t.nb[0]*(info.t.ne[0]/blck_size); - for (int j = 2; j < GGML_MAX_DIMS; ++j) { - info.t.nb[j] = info.t.nb[j - 1]*info.t.ne[j - 1]; - } - } - if (!ok) { - break; - } - - // tensor data offset within buffer - ok = ok && gr.read(info.offset); - - ctx->info.push_back(info); - } - - if (!ok) { - GGML_LOG_ERROR("%s: failed to read tensor info\n", __func__); - gguf_free(ctx); - return nullptr; - } - GGML_ASSERT(int64_t(ctx->info.size()) == n_tensors); - - // we require the data section to be aligned, so take into account any padding - if (gguf_fseek(file, GGML_PAD(gguf_ftell(file), ctx->alignment), SEEK_SET) != 0) { - GGML_LOG_ERROR("%s: failed to seek to beginning of data section\n", __func__); - gguf_free(ctx); - return nullptr; - } - - // store the current file offset - this is where the data section starts - ctx->offset = gguf_ftell(file); - - // compute the total size of the data section, taking into account the alignment - { - ctx->size = 0; - for (size_t i = 0; i < ctx->info.size(); ++i) { - const gguf_tensor_info & ti = ctx->info[i]; - if (ti.offset != ctx->size) { - GGML_LOG_ERROR("%s: tensor '%s' has offset %" PRIu64 ", expected %zu\n", - __func__, ti.t.name, ti.offset, ctx->size); - GGML_LOG_ERROR("%s: failed to read tensor data\n", __func__); - gguf_free(ctx); - return nullptr; - } - size_t padded_size = GGML_PAD(ggml_nbytes(&ti.t), ctx->alignment); - if (SIZE_MAX - ctx->size < padded_size) { - GGML_LOG_ERROR("%s: tensor '%s' size overflow, cannot accumulate size %zu + %zu\n", - __func__, ti.t.name, ctx->size, padded_size); - gguf_free(ctx); - return nullptr; - } - ctx->size += padded_size; - } - } - - // load the tensor data only if requested - if (params.ctx != nullptr) { - // if the provided gguf_context is no_alloc, then we create "empty" tensors and do not read the binary blob - // otherwise, we load the binary blob into the created ggml_context as well, and point the "data" members of - // the ggml_tensor structs to the appropriate locations in the binary blob - - // compute the exact size needed for the new ggml_context - size_t mem_size = 0; - if (params.no_alloc) { - if (n_tensors != 0 && SIZE_MAX / n_tensors < ggml_tensor_overhead()) { - GGML_LOG_ERROR("%s: memory size overflow while allocating ggml context\n", __func__); - gguf_free(ctx); - return nullptr; - } - - const size_t overhead = n_tensors * ggml_tensor_overhead(); - - mem_size = overhead; - } else { - if ((n_tensors + 1) != 0 && SIZE_MAX / (n_tensors + 1) < ggml_tensor_overhead()) { - GGML_LOG_ERROR("%s: memory size overflow while allocating ggml context\n", __func__); - gguf_free(ctx); - return nullptr; - } - - const size_t overhead = (n_tensors + 1) * ggml_tensor_overhead(); - - if (SIZE_MAX - overhead < ctx->size) { - GGML_LOG_ERROR("%s: memory size overflow while allocating ggml context\n", __func__); - gguf_free(ctx); - return nullptr; - } - - mem_size = overhead + ctx->size; - } - - struct ggml_init_params pdata = { - /*mem_size =*/ mem_size, - /*mem_buffer =*/ nullptr, - /*no_alloc =*/ params.no_alloc, - }; - - *params.ctx = ggml_init(pdata); - if (*params.ctx == nullptr) { - GGML_LOG_ERROR("%s: failed to initialize ggml context for storing tensors\n", __func__); - gguf_free(ctx); - return nullptr; - } - - struct ggml_context * ctx_data = *params.ctx; - - struct ggml_tensor * data = nullptr; - - if (!params.no_alloc) { - data = ggml_new_tensor_1d(ctx_data, GGML_TYPE_I8, ctx->size); - - ok = ok && data != nullptr; - - if (ok) { - ggml_set_name(data, "GGUF tensor data binary blob"); - } - - // read the binary blob with the tensor data - ok = ok && gr.read(data->data, ctx->size); - - if (!ok) { - GGML_LOG_ERROR("%s: failed to read tensor data binary blob\n", __func__); - ggml_free(ctx_data); - *params.ctx = nullptr; - gguf_free(ctx); - return nullptr; - } - - ctx->data = data->data; - } - - ggml_set_no_alloc(ctx_data, true); - - // create the tensors - for (size_t i = 0; i < ctx->info.size(); ++i) { - const struct gguf_tensor_info & info = ctx->info[i]; - - struct ggml_tensor * cur = ggml_new_tensor(ctx_data, info.t.type, GGML_MAX_DIMS, info.t.ne); - - ok = ok && cur != nullptr; - - if (!ok) { - break; - } - - ggml_set_name(cur, info.t.name); - - // point the data member to the appropriate location in the binary blob using the tensor info - if (!params.no_alloc) { - cur->data = (char *) data->data + info.offset; - } - } - - if (!ok) { - GGML_LOG_ERROR("%s: failed to create tensors\n", __func__); - ggml_free(ctx_data); - *params.ctx = nullptr; - gguf_free(ctx); - return nullptr; - } - - ggml_set_no_alloc(ctx_data, params.no_alloc); - } - - return ctx; -} - -struct gguf_context * gguf_init_from_file(const char * fname, struct gguf_init_params params) { - FILE * file = ggml_fopen(fname, "rb"); - - if (!file) { - GGML_LOG_ERROR("%s: failed to open GGUF file '%s' (%s)\n", __func__, fname, strerror(errno)); - return nullptr; - } - - struct gguf_context * result = gguf_init_from_file_ptr(file, params); - fclose(file); - return result; -} - -void gguf_free(struct gguf_context * ctx) { - if (ctx == nullptr) { - return; - } - delete ctx; -} - -const char * gguf_type_name(enum gguf_type type) { - auto it = GGUF_TYPE_NAME.find(type); - return it == GGUF_TYPE_NAME.end() ? nullptr : it->second; -} - -uint32_t gguf_get_version(const struct gguf_context * ctx) { - return ctx->version; -} - -size_t gguf_get_alignment(const struct gguf_context * ctx) { - return ctx->alignment; -} - -size_t gguf_get_data_offset(const struct gguf_context * ctx) { - return ctx->offset; -} - -int64_t gguf_get_n_kv(const struct gguf_context * ctx) { - return ctx->kv.size(); -} - -int64_t gguf_find_key(const struct gguf_context * ctx, const char * key) { - // return -1 if key not found - int64_t keyfound = -1; - - const int64_t n_kv = gguf_get_n_kv(ctx); - - for (int64_t i = 0; i < n_kv; ++i) { - if (strcmp(key, gguf_get_key(ctx, i)) == 0) { - keyfound = i; - break; - } - } - - return keyfound; -} - -const char * gguf_get_key(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - return ctx->kv[key_id].get_key().c_str(); -} - -enum gguf_type gguf_get_kv_type(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - return ctx->kv[key_id].is_array ? GGUF_TYPE_ARRAY : ctx->kv[key_id].get_type(); -} - -enum gguf_type gguf_get_arr_type(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].is_array); - return ctx->kv[key_id].get_type(); -} - -const void * gguf_get_arr_data(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_type() != GGUF_TYPE_STRING); - return ctx->kv[key_id].data.data(); -} - -const char * gguf_get_arr_str(const struct gguf_context * ctx, int64_t key_id, size_t i) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_type() == GGUF_TYPE_STRING); - return ctx->kv[key_id].data_string[i].c_str(); -} - -size_t gguf_get_arr_n(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - - if (ctx->kv[key_id].type == GGUF_TYPE_STRING) { - return ctx->kv[key_id].data_string.size(); - } - - const size_t type_size = gguf_type_size(ctx->kv[key_id].type); - GGML_ASSERT(ctx->kv[key_id].data.size() % type_size == 0); - return ctx->kv[key_id].data.size() / type_size; -} - -uint8_t gguf_get_val_u8(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - return ctx->kv[key_id].get_val(); -} - -int8_t gguf_get_val_i8(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - return ctx->kv[key_id].get_val(); -} - -uint16_t gguf_get_val_u16(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - return ctx->kv[key_id].get_val(); -} - -int16_t gguf_get_val_i16(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - return ctx->kv[key_id].get_val(); -} - -uint32_t gguf_get_val_u32(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - return ctx->kv[key_id].get_val(); -} - -int32_t gguf_get_val_i32(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - return ctx->kv[key_id].get_val(); -} - -float gguf_get_val_f32(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - return ctx->kv[key_id].get_val(); -} - -uint64_t gguf_get_val_u64(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - return ctx->kv[key_id].get_val(); -} - -int64_t gguf_get_val_i64(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - return ctx->kv[key_id].get_val(); -} - -double gguf_get_val_f64(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - return ctx->kv[key_id].get_val(); -} - -bool gguf_get_val_bool(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - return ctx->kv[key_id].get_val(); -} - -const char * gguf_get_val_str(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - return ctx->kv[key_id].get_val().c_str(); -} - -const void * gguf_get_val_data(const struct gguf_context * ctx, int64_t key_id) { - GGML_ASSERT(key_id >= 0 && key_id < gguf_get_n_kv(ctx)); - GGML_ASSERT(ctx->kv[key_id].get_ne() == 1); - GGML_ASSERT(ctx->kv[key_id].get_type() != GGUF_TYPE_STRING); - return ctx->kv[key_id].data.data(); -} - -int64_t gguf_get_n_tensors(const struct gguf_context * ctx) { - return ctx->info.size(); -} - -int64_t gguf_find_tensor(const struct gguf_context * ctx, const char * name) { - // return -1 if tensor not found - int64_t tensor_id = -1; - - const int64_t n_tensors = gguf_get_n_tensors(ctx); - - for (int64_t i = 0; i < n_tensors; ++i) { - if (strcmp(name, gguf_get_tensor_name(ctx, i)) == 0) { - tensor_id = i; - break; - } - } - - return tensor_id; -} - -size_t gguf_get_tensor_offset(const struct gguf_context * ctx, int64_t tensor_id) { - GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx)); - return ctx->info[tensor_id].offset; -} - -const char * gguf_get_tensor_name(const struct gguf_context * ctx, int64_t tensor_id) { - GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx)); - return ctx->info[tensor_id].t.name; -} - -enum ggml_type gguf_get_tensor_type(const struct gguf_context * ctx, int64_t tensor_id) { - GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx)); - return ctx->info[tensor_id].t.type; -} - -size_t gguf_get_tensor_size(const struct gguf_context * ctx, int64_t tensor_id) { - GGML_ASSERT(tensor_id >= 0 && tensor_id < gguf_get_n_tensors(ctx)); - return ggml_nbytes(&ctx->info[tensor_id].t); -} - -int64_t gguf_remove_key(struct gguf_context * ctx, const char * key) { - const int64_t key_id = gguf_find_key(ctx, key); - if (key_id >= 0) { - ctx->kv.erase(ctx->kv.begin() + key_id); - } - return key_id; -} - -template -static void gguf_check_reserved_keys(const std::string & key, const T val) { - if (key == GGUF_KEY_GENERAL_ALIGNMENT) { - if constexpr (std::is_same::value) { - GGML_ASSERT(val > 0 && (val & (val - 1)) == 0 && GGUF_KEY_GENERAL_ALIGNMENT " must be power of 2"); - } else { - GGML_UNUSED(val); - GGML_ABORT(GGUF_KEY_GENERAL_ALIGNMENT " must be type u32"); - } - } -} - -void gguf_set_val_u8(struct gguf_context * ctx, const char * key, uint8_t val) { - gguf_check_reserved_keys(key, val); - gguf_remove_key(ctx, key); - ctx->kv.emplace_back(key, val); -} - -void gguf_set_val_i8(struct gguf_context * ctx, const char * key, int8_t val) { - gguf_check_reserved_keys(key, val); - gguf_remove_key(ctx, key); - ctx->kv.emplace_back(key, val); -} - -void gguf_set_val_u16(struct gguf_context * ctx, const char * key, uint16_t val) { - gguf_check_reserved_keys(key, val); - gguf_remove_key(ctx, key); - ctx->kv.emplace_back(key, val); -} - -void gguf_set_val_i16(struct gguf_context * ctx, const char * key, int16_t val) { - gguf_check_reserved_keys(key, val); - gguf_remove_key(ctx, key); - ctx->kv.emplace_back(key, val); -} - -void gguf_set_val_u32(struct gguf_context * ctx, const char * key, uint32_t val) { - gguf_check_reserved_keys(key, val); - gguf_remove_key(ctx, key); - ctx->kv.emplace_back(key, val); -} - -void gguf_set_val_i32(struct gguf_context * ctx, const char * key, int32_t val) { - gguf_check_reserved_keys(key, val); - gguf_remove_key(ctx, key); - ctx->kv.emplace_back(key, val); -} - -void gguf_set_val_f32(struct gguf_context * ctx, const char * key, float val) { - gguf_check_reserved_keys(key, val); - gguf_remove_key(ctx, key); - ctx->kv.emplace_back(key, val); -} - -void gguf_set_val_u64(struct gguf_context * ctx, const char * key, uint64_t val) { - gguf_check_reserved_keys(key, val); - gguf_remove_key(ctx, key); - ctx->kv.emplace_back(key, val); -} - -void gguf_set_val_i64(struct gguf_context * ctx, const char * key, int64_t val) { - gguf_check_reserved_keys(key, val); - gguf_remove_key(ctx, key); - ctx->kv.emplace_back(key, val); -} - -void gguf_set_val_f64(struct gguf_context * ctx, const char * key, double val) { - gguf_check_reserved_keys(key, val); - gguf_remove_key(ctx, key); - ctx->kv.emplace_back(key, val); -} - -void gguf_set_val_bool(struct gguf_context * ctx, const char * key, bool val) { - gguf_check_reserved_keys(key, val); - gguf_remove_key(ctx, key); - ctx->kv.emplace_back(key, val); -} - -void gguf_set_val_str(struct gguf_context * ctx, const char * key, const char * val) { - gguf_check_reserved_keys(key, val); - gguf_remove_key(ctx, key); - ctx->kv.emplace_back(key, std::string(val)); -} - -void gguf_set_arr_data(struct gguf_context * ctx, const char * key, enum gguf_type type, const void * data, size_t n) { - gguf_check_reserved_keys(key, data); - gguf_remove_key(ctx, key); - - const size_t nbytes = n*gguf_type_size(type); - std::vector tmp(nbytes); - if (!tmp.empty()) { - memcpy(tmp.data(), data, nbytes); - } - ctx->kv.emplace_back(key, tmp); - ctx->kv.back().cast(type); -} - -void gguf_set_arr_str(struct gguf_context * ctx, const char * key, const char ** data, size_t n) { - gguf_check_reserved_keys(key, data); - gguf_remove_key(ctx, key); - - std::vector tmp(n); - for (size_t i = 0; i < n; ++i) { - tmp[i] = data[i]; - } - ctx->kv.emplace_back(key, tmp); -} - -// set or add KV pairs from another context -void gguf_set_kv(struct gguf_context * ctx, const struct gguf_context * src) { - const int64_t n_kv = gguf_get_n_kv(src); - for (int64_t i = 0; i < n_kv; ++i) { - const struct gguf_kv & kv = src->kv[i]; - - if (!kv.is_array) { - switch (kv.get_type()) { - case GGUF_TYPE_UINT8: gguf_set_val_u8 (ctx, kv.get_key().c_str(), kv.get_val()); break; - case GGUF_TYPE_INT8: gguf_set_val_i8 (ctx, kv.get_key().c_str(), kv.get_val()); break; - case GGUF_TYPE_UINT16: gguf_set_val_u16 (ctx, kv.get_key().c_str(), kv.get_val()); break; - case GGUF_TYPE_INT16: gguf_set_val_i16 (ctx, kv.get_key().c_str(), kv.get_val()); break; - case GGUF_TYPE_UINT32: gguf_set_val_u32 (ctx, kv.get_key().c_str(), kv.get_val()); break; - case GGUF_TYPE_INT32: gguf_set_val_i32 (ctx, kv.get_key().c_str(), kv.get_val()); break; - case GGUF_TYPE_FLOAT32: gguf_set_val_f32 (ctx, kv.get_key().c_str(), kv.get_val()); break; - case GGUF_TYPE_UINT64: gguf_set_val_u64 (ctx, kv.get_key().c_str(), kv.get_val()); break; - case GGUF_TYPE_INT64: gguf_set_val_i64 (ctx, kv.get_key().c_str(), kv.get_val()); break; - case GGUF_TYPE_FLOAT64: gguf_set_val_f64 (ctx, kv.get_key().c_str(), kv.get_val()); break; - case GGUF_TYPE_BOOL: gguf_set_val_bool(ctx, kv.get_key().c_str(), kv.get_val()); break; - case GGUF_TYPE_STRING: gguf_set_val_str (ctx, kv.get_key().c_str(), kv.get_val().c_str()); break; - case GGUF_TYPE_ARRAY: - default: GGML_ABORT("invalid type"); - } - continue; - } - - const size_t ne = kv.get_ne(); - - switch (kv.get_type()) { - case GGUF_TYPE_UINT8: - case GGUF_TYPE_INT8: - case GGUF_TYPE_UINT16: - case GGUF_TYPE_INT16: - case GGUF_TYPE_UINT32: - case GGUF_TYPE_INT32: - case GGUF_TYPE_FLOAT32: - case GGUF_TYPE_UINT64: - case GGUF_TYPE_INT64: - case GGUF_TYPE_FLOAT64: - case GGUF_TYPE_BOOL: { - gguf_set_arr_data(ctx, kv.get_key().c_str(), kv.get_type(), kv.data.data(), ne); - } break; - case GGUF_TYPE_STRING: { - std::vector tmp(ne); - for (size_t j = 0; j < ne; ++j) { - tmp[j] = kv.data_string[j].c_str(); - } - gguf_set_arr_str(ctx, kv.get_key().c_str(), tmp.data(), ne); - } break; - case GGUF_TYPE_ARRAY: - default: GGML_ABORT("invalid type"); - } - } -} - -void gguf_add_tensor( - struct gguf_context * ctx, - const struct ggml_tensor * tensor) { - GGML_ASSERT(tensor); - if (gguf_find_tensor(ctx, tensor->name) != -1) { - GGML_ABORT("duplicate tensor name: %s", tensor->name); - } - - struct gguf_tensor_info ti; - ti.t = *tensor; - ti.offset = ctx->info.empty() ? 0 : - ctx->info.back().offset + GGML_PAD(ggml_nbytes(&ctx->info.back().t), ctx->alignment); - ctx->info.push_back(ti); -} - -void gguf_set_tensor_type(struct gguf_context * ctx, const char * name, enum ggml_type type) { - const int64_t tensor_id = gguf_find_tensor(ctx, name); - if (tensor_id < 0) { - GGML_ABORT("tensor not found: %s", name); - } - struct ggml_tensor * tensor = &ctx->info[tensor_id].t; - const size_t type_size = ggml_type_size(type); - const int64_t blck_size = ggml_blck_size(type); - - tensor->type = type; - GGML_ASSERT(tensor->ne[0] % blck_size == 0 && "tensor row size not divisible by block size of new type"); - - tensor->nb[0] = type_size; - tensor->nb[1] = tensor->nb[0]*(tensor->ne[0]/blck_size); - for (int i = 2; i < GGML_MAX_DIMS; i++) { - tensor->nb[i] = tensor->nb[i - 1]*tensor->ne[i - 1]; - } - - // update offsets - const int64_t n_tensors = gguf_get_n_tensors(ctx); - for (int64_t i = tensor_id + 1; i < n_tensors; ++i) { - ctx->info[i].offset = ctx->info[i - 1].offset + GGML_PAD(ggml_nbytes(&ctx->info[i - 1].t), ctx->alignment); - } -} - -void gguf_set_tensor_data(struct gguf_context * ctx, const char * name, const void * data) { - const int64_t tensor_id = gguf_find_tensor(ctx, name); - if (tensor_id < 0) { - GGML_ABORT("tensor not found: %s", name); - } - - ctx->info[tensor_id].t.data = (void *)(uintptr_t)data; // double cast suppresses warning about casting away const -} - -struct gguf_writer_base { - size_t written_bytes {0u}; - - ~gguf_writer_base(void) = default; - - // we bet on devirtualization - virtual void write(int8_t val) = 0; - virtual void write(const std::vector & val) = 0; - virtual void write_tensor_data(const struct gguf_tensor_info & info, size_t offset_data, size_t alignment) = 0; - - template - void write(const T & val) { - for (size_t i = 0; i < sizeof(val); ++i) { - write(reinterpret_cast(&val)[i]); - } - } - - void write(const bool & val) { - const int8_t val8 = val ? 1 : 0; - write(val8); - } - - void write(const std::string & val) { - { - const uint64_t n = val.length(); - write(n); - } - for (size_t i = 0; i < val.length(); ++i) { - write((val.data())[i]); - } - } - - void write(const char * val) { - write(std::string(val)); - } - - void write(const enum ggml_type & val) { - write(int32_t(val)); - } - - void write(const enum gguf_type & val) { - write(int32_t(val)); - } - - void write(const struct gguf_kv & kv) { - const uint64_t ne = kv.get_ne(); - - write(kv.get_key()); - - if (kv.is_array) { - write(GGUF_TYPE_ARRAY); - write(kv.get_type()); - write(ne); - } else { - write(kv.get_type()); - } - - switch (kv.get_type()) { - case GGUF_TYPE_UINT8: - case GGUF_TYPE_INT8: - case GGUF_TYPE_UINT16: - case GGUF_TYPE_INT16: - case GGUF_TYPE_UINT32: - case GGUF_TYPE_INT32: - case GGUF_TYPE_FLOAT32: - case GGUF_TYPE_UINT64: - case GGUF_TYPE_INT64: - case GGUF_TYPE_FLOAT64: { - write(kv.data); - } break; - case GGUF_TYPE_BOOL: { - for (size_t i = 0; i < ne; ++i) { - write(kv.get_val(i)); - } - } break; - case GGUF_TYPE_STRING: { - for (size_t i = 0; i < ne; ++i) { - write(kv.get_val(i)); - } - } break; - case GGUF_TYPE_ARRAY: - default: GGML_ABORT("invalid type"); - } - } - - void write_tensor_meta(const struct gguf_tensor_info & info) { - write(info.t.name); - - const uint32_t n_dims = ggml_n_dims(&info.t); - write(n_dims); - - for (uint32_t j = 0; j < n_dims; ++j) { - write(info.t.ne[j]); - } - write(info.t.type); - write(info.offset); - } - - void pad(const size_t alignment) { - while (written_bytes % alignment != 0) { - const int8_t zero = 0; - write(zero); - } - } -}; - -// vector buffer based writer -struct gguf_writer_buf final : public gguf_writer_base { - std::vector & buf; - - gguf_writer_buf(std::vector & buf) : buf(buf) {} - - using gguf_writer_base::write; - - void write(const int8_t val) override { - buf.push_back(val); - written_bytes++; - } - - void write(const std::vector & val) override { - buf.insert(buf.end(), val.begin(), val.end()); - written_bytes += val.size(); - } - - void write_tensor_data(const struct gguf_tensor_info & info, const size_t offset_data, const size_t alignment) override { - GGML_ASSERT(buf.size() - offset_data == info.offset); - - GGML_ASSERT(ggml_is_contiguous(&info.t)); - const size_t offset = buf.size(); - const size_t nbytes = ggml_nbytes(&info.t); - - buf.resize(offset + nbytes); - if (info.t.buffer) { - ggml_backend_tensor_get(&info.t, buf.data() + offset, 0, nbytes); - } else { - GGML_ASSERT(info.t.data); - memcpy(buf.data() + offset, info.t.data, nbytes); - } - written_bytes += nbytes; - - pad(alignment); - } -}; - -// file based writer -struct gguf_writer_file final : public gguf_writer_base { - FILE * file; - - gguf_writer_file(FILE* file) : file(file) {} - - using gguf_writer_base::write; - - void write(const int8_t val) override { - const auto real_val = static_cast(val); - const auto ret = fputc(real_val, file); - written_bytes++; - if (ret != real_val) { - throw std::runtime_error("unexpected fputc result '" + std::to_string(ret) + "' instead of '" + std::to_string((int)real_val) + "'"); - } - } - - void write(const std::vector & val) override { - const auto ret = fwrite(val.data(), 1, val.size(), file); - written_bytes += val.size(); - if (ret != val.size()) { - throw std::runtime_error("unexpected fwrite number of bytes written, '" + std::to_string(ret) + "' instead of '" + std::to_string(val.size()) + "'"); - } - } - - void write_tensor_data(const struct gguf_tensor_info & info, const size_t offset_data, const size_t alignment) override { - GGML_ASSERT(written_bytes - offset_data == info.offset); - - GGML_ASSERT(ggml_is_contiguous(&info.t)); - const size_t nbytes = ggml_nbytes(&info.t); - - std::vector buf(nbytes); - if (info.t.buffer) { - ggml_backend_tensor_get(&info.t, buf.data(), 0, nbytes); - } else { - GGML_ASSERT(info.t.data); - memcpy(buf.data(), info.t.data, nbytes); - } - write(buf); - - pad(alignment); - } -}; - -template -static void gguf_write_out(const struct gguf_context * ctx, writer_t & gw, bool only_meta) { - const int64_t n_kv = gguf_get_n_kv(ctx); - const int64_t n_tensors = gguf_get_n_tensors(ctx); - - // write header - gw.write(GGUF_MAGIC[0]); - gw.write(GGUF_MAGIC[1]); - gw.write(GGUF_MAGIC[2]); - gw.write(GGUF_MAGIC[3]); - gw.write(ctx->version); - gw.write(n_tensors); - gw.write(n_kv); - - // write key-value pairs - for (int64_t i = 0; i < n_kv; ++i) { - gw.write(ctx->kv[i]); - } - - // write tensor info - for (int64_t i = 0; i < n_tensors; ++i) { - gw.write_tensor_meta(ctx->info[i]); - } - - // we require the data section to be aligned - gw.pad(ctx->alignment); - - if (only_meta) { - return; - } - - const size_t offset_data = gw.written_bytes; - - // write tensor data - for (int64_t i = 0; i < n_tensors; ++i) { - gw.write_tensor_data(ctx->info[i], offset_data, ctx->alignment); - } -} - -void gguf_write_to_buf(const struct gguf_context * ctx, std::vector & buf, bool only_meta) { - gguf_writer_buf gw(buf); - gguf_write_out(ctx, gw, only_meta); -} - -bool gguf_write_to_file_ptr(const struct gguf_context * ctx, FILE * file, bool only_meta) { - GGML_ASSERT(file); - - try { - gguf_writer_file gw(file); - gguf_write_out(ctx, gw, only_meta); - } catch (const std::runtime_error& ex) { - GGML_LOG_ERROR("%s: failed to write GGUF data: %s\n", __func__, ex.what()); - return false; - } - return true; -} - -bool gguf_write_to_file(const struct gguf_context * ctx, const char * fname, bool only_meta) { - FILE * file = ggml_fopen(fname, "wb"); - - if (!file) { - GGML_LOG_ERROR("%s: failed to open file '%s' for writing GGUF data\n", __func__, fname); - return false; - } - - const bool success = gguf_write_to_file_ptr(ctx, file, only_meta); - if (!success) { - GGML_LOG_ERROR("%s: failed to write GGUF data into '%s'\n", __func__, fname); - } - - fclose(file); - return success; -} - -size_t gguf_get_meta_size(const struct gguf_context * ctx) { - // only return size - std::vector buf; - gguf_write_to_buf(ctx, buf, /*only_meta =*/ true); - return buf.size(); -} - -void gguf_get_meta_data(const struct gguf_context * ctx, void * data) { - std::vector buf; - gguf_write_to_buf(ctx, buf, /*only_meta =*/ true); - memcpy(data, buf.data(), buf.size()); -} diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 83ae51ce9ce3..9aea21063922 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -10,7 +10,7 @@ GGUF_MAGIC = 0x46554747 # "GGUF" GGUF_VERSION = 3 GGUF_DEFAULT_ALIGNMENT = 32 -GGML_QUANT_VERSION = 2 # GGML_QNT_VERSION from ggml.h +GGML_QUANT_VERSION = 2 # GGML_QNT_VERSION from ggml.h.inc # # metadata keys @@ -4032,10 +4032,10 @@ class ExpertGatingFuncType(IntEnum): SIGMOID = 2 -# TODO: add GGMLFileType from ggml_ftype in ggml.h +# TODO: add GGMLFileType from ggml_ftype in ggml.h.inc -# from llama_ftype in llama.h +# from llama_ftype in llama.h.inc # ALL VALUES SHOULD BE THE SAME HERE AS THEY ARE OVER THERE. class LlamaFileType(IntEnum): ALL_F32 = 0 diff --git a/grammars/README.md b/grammars/README.md index dcd28648b157..10f16892101d 100644 --- a/grammars/README.md +++ b/grammars/README.md @@ -136,7 +136,7 @@ You can use GBNF grammars: - In [llama-server](../tools/server)'s completion endpoints, passed as the `grammar` body field - In [llama-cli](../tools/cli) and [llama-completion](../tools/completion), passed as the `--grammar` & `--grammar-file` flags -- With [test-gbnf-validator](../tests/test-gbnf-validator.cpp), to test them against strings. +- With `test-gbnf-validator`, to test them against strings. ## JSON Schemas → GBNF diff --git a/include/llama-cpp.h b/include/llama-cpp.h deleted file mode 100644 index 8f6368177de0..000000000000 --- a/include/llama-cpp.h +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once - -#ifndef __cplusplus -#error "This header is for C++ only" -#endif - -#include - -#include "llama.h" - -struct llama_model_deleter { - void operator()(llama_model * model) { llama_model_free(model); } -}; - -struct llama_context_deleter { - void operator()(llama_context * context) { llama_free(context); } -}; - -struct llama_sampler_deleter { - void operator()(llama_sampler * sampler) { llama_sampler_free(sampler); } -}; - -struct llama_adapter_lora_deleter { - void operator()(llama_adapter_lora * adapter) { llama_adapter_lora_free(adapter); } -}; - -typedef std::unique_ptr llama_model_ptr; -typedef std::unique_ptr llama_context_ptr; -typedef std::unique_ptr llama_sampler_ptr; -typedef std::unique_ptr llama_adapter_lora_ptr; diff --git a/include/llama.h b/include/llama.h.inc similarity index 99% rename from include/llama.h rename to include/llama.h.inc index eb8698140978..2b607c2d08de 100644 --- a/include/llama.h +++ b/include/llama.h.inc @@ -1,11 +1,11 @@ #ifndef LLAMA_H #define LLAMA_H -#include "ggml.h" -#include "ggml-cpu.h" -#include "ggml-backend.h" -#include "ggml-opt.h" -#include "gguf.h" +#include "ggml.h.inc" +#include "ggml-cpu.h.inc" +#include "ggml-backend.h.inc" +#include "ggml-opt.h.inc" +#include "gguf.h.inc" #include #include diff --git a/pocs/vdot/CMakeLists.txt b/pocs/vdot/CMakeLists.txt index f3776268ab69..ebd909cab26a 100644 --- a/pocs/vdot/CMakeLists.txt +++ b/pocs/vdot/CMakeLists.txt @@ -1,9 +1,44 @@ set(TARGET llama-vdot) -add_executable(${TARGET} vdot.cpp) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(VDOT_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/vdot-rust/Cargo.toml) +set(VDOT_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/vdot) +set(VDOT_RUST_BIN ${VDOT_RUST_TARGET_DIR}/debug/llama-vdot${CMAKE_EXECUTABLE_SUFFIX}) +set(VDOT_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${VDOT_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=${CMAKE_RUNTIME_OUTPUT_DIRECTORY} + ${CARGO_EXECUTABLE} build --manifest-path ${VDOT_RUST_MANIFEST} --target-dir ${VDOT_RUST_TARGET_DIR} --bin llama-vdot + COMMAND ${CMAKE_COMMAND} -E copy ${VDOT_RUST_BIN} ${VDOT_RUST_OUTPUT} + DEPENDS + ggml-cpu + ${VDOT_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/vdot-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/vdot-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/vdot-rust/src/main.rs +) + +add_custom_target(${TARGET} DEPENDS ${VDOT_RUST_OUTPUT}) set(TARGET llama-q8dot) -add_executable(${TARGET} q8dot.cpp) -target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) -target_compile_features(${TARGET} PRIVATE cxx_std_17) + +set(Q8DOT_RUST_MANIFEST ${CMAKE_CURRENT_SOURCE_DIR}/q8dot-rust/Cargo.toml) +set(Q8DOT_RUST_TARGET_DIR ${CMAKE_BINARY_DIR}/cargo/q8dot) +set(Q8DOT_RUST_BIN ${Q8DOT_RUST_TARGET_DIR}/debug/llama-q8dot${CMAKE_EXECUTABLE_SUFFIX}) +set(Q8DOT_RUST_OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${TARGET}${CMAKE_EXECUTABLE_SUFFIX}) + +add_custom_command( + OUTPUT ${Q8DOT_RUST_OUTPUT} + COMMAND ${CMAKE_COMMAND} -E env LLAMA_CPP_LIB_DIR=${CMAKE_RUNTIME_OUTPUT_DIRECTORY} + ${CARGO_EXECUTABLE} build --manifest-path ${Q8DOT_RUST_MANIFEST} --target-dir ${Q8DOT_RUST_TARGET_DIR} --bin llama-q8dot + COMMAND ${CMAKE_COMMAND} -E copy ${Q8DOT_RUST_BIN} ${Q8DOT_RUST_OUTPUT} + DEPENDS + ggml-cpu + ${Q8DOT_RUST_MANIFEST} + ${CMAKE_CURRENT_SOURCE_DIR}/q8dot-rust/build.rs + ${CMAKE_CURRENT_SOURCE_DIR}/q8dot-rust/src/lib.rs + ${CMAKE_CURRENT_SOURCE_DIR}/q8dot-rust/src/main.rs +) + +add_custom_target(${TARGET} DEPENDS ${Q8DOT_RUST_OUTPUT}) diff --git a/pocs/vdot/q8dot-rust/Cargo.toml b/pocs/vdot/q8dot-rust/Cargo.toml new file mode 100644 index 000000000000..3a8ae1a2fce1 --- /dev/null +++ b/pocs/vdot/q8dot-rust/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "llama-q8dot-rust" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] + +[[bin]] +name = "llama-q8dot" +path = "src/main.rs" diff --git a/pocs/vdot/q8dot-rust/build.rs b/pocs/vdot/q8dot-rust/build.rs new file mode 100644 index 000000000000..dbff8b4d787b --- /dev/null +++ b/pocs/vdot/q8dot-rust/build.rs @@ -0,0 +1,46 @@ +use std::env; +use std::path::PathBuf; + +fn main() { + if let Ok(lib_dir) = env::var("LLAMA_CPP_LIB_DIR") { + println!("cargo:rustc-link-search=native={lib_dir}"); + println!("cargo:rustc-link-lib=ggml-cpu"); + println!("cargo:rustc-link-lib=ggml-base"); + emit_rpath(&lib_dir); + return; + } + + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let root = manifest_dir.join("../../.."); + let candidates = [ + root.join("build/bin"), + root.join("build/ggml/src"), + root.join("build"), + ]; + + for candidate in candidates { + if candidate.join("libggml-cpu.so").exists() + || candidate.join("libggml-cpu.dylib").exists() + || candidate.join("ggml-cpu.dll").exists() + || candidate.join("ggml-cpu.lib").exists() + || candidate.join("libggml-cpu.a").exists() + { + println!("cargo:rustc-link-search=native={}", candidate.display()); + println!("cargo:rustc-link-lib=ggml-cpu"); + println!("cargo:rustc-link-lib=ggml-base"); + emit_rpath(&candidate.display().to_string()); + return; + } + } + + println!( + "cargo:warning=libggml-cpu was not found; build llama.cpp before building llama-q8dot" + ); +} + +fn emit_rpath(lib_dir: &str) { + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + if target_os != "windows" { + println!("cargo:rustc-link-arg=-Wl,-rpath,{lib_dir}"); + } +} diff --git a/pocs/vdot/q8dot-rust/src/lib.rs b/pocs/vdot/q8dot-rust/src/lib.rs new file mode 100644 index 000000000000..20034661b5f6 --- /dev/null +++ b/pocs/vdot/q8dot-rust/src/lib.rs @@ -0,0 +1,174 @@ +pub const QK4_0: usize = 32; +pub const QK4_1: usize = 32; +pub const QK8_0: usize = 32; + +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct BlockQ4_0 { + pub d: f32, + pub qs: [u8; QK4_0 / 2], +} + +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct BlockQ4_1 { + pub d: f32, + pub m: f32, + pub qs: [u8; QK4_1 / 2], +} + +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct BlockQ8_0 { + pub d: f32, + pub s: f32, + pub qs: [i8; QK8_0], +} + +pub struct Rng(u32); + +impl Rng { + pub fn new(seed: u32) -> Self { + Self(seed) + } + + pub fn next_u32(&mut self) -> u32 { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + self.0 = x; + x + } +} + +pub fn fill_q4_0(blocks: &mut [BlockQ4_0], rng: &mut Rng) { + for block in blocks { + block.d = 1.0; + for q in &mut block.qs { + let v1 = (rng.next_u32() >> 28) as u8; + let v2 = (rng.next_u32() >> 28) as u8; + *q = v1 | (v2 << 4); + } + } +} + +pub fn fill_q4_1(blocks: &mut [BlockQ4_1], rng: &mut Rng) { + for block in blocks { + block.d = 1.0; + block.m = 1.0; + for q in &mut block.qs { + let v1 = (rng.next_u32() >> 28) as u8; + let v2 = (rng.next_u32() >> 28) as u8; + *q = v1 | (v2 << 4); + } + } +} + +pub fn fill_q8_0(blocks: &mut [BlockQ8_0], rng: &mut Rng) { + for block in blocks { + block.d = 1.0; + let mut sum = 0_i32; + for q in &mut block.qs { + *q = ((rng.next_u32() >> 24) as i32 - 128) as i8; + sum += i32::from(*q); + } + block.s = block.d * sum as f32; + } +} + +pub fn simple_dot_q4_0(x: &BlockQ4_0, y: &BlockQ8_0) -> f32 { + let mut s1 = 0_i32; + for i in (0..QK4_0 / 2).step_by(2) { + let v1 = i32::from(x.qs[i] & 0x0f); + let v2 = i32::from(x.qs[i] >> 4); + let v3 = i32::from(x.qs[i + 1] & 0x0f); + let v4 = i32::from(x.qs[i + 1] >> 4); + let j = 2 * i; + s1 += v1 * i32::from(y.qs[j]) + + v2 * i32::from(y.qs[j + 1]) + + v3 * i32::from(y.qs[j + 2]) + + v4 * i32::from(y.qs[j + 3]); + } + y.d * x.d * s1 as f32 - 8.0 * x.d * y.s +} + +pub fn simple_dot_q4_1(x: &BlockQ4_1, y: &BlockQ8_0) -> f32 { + let mut s1 = 0_i32; + for i in (0..QK4_1 / 2).step_by(2) { + let v1 = i32::from(x.qs[i] & 0x0f); + let v2 = i32::from(x.qs[i] >> 4); + let v3 = i32::from(x.qs[i + 1] & 0x0f); + let v4 = i32::from(x.qs[i + 1] >> 4); + let j = 2 * i; + s1 += v1 * i32::from(y.qs[j]) + + v2 * i32::from(y.qs[j + 1]) + + v3 * i32::from(y.qs[j + 2]) + + v4 * i32::from(y.qs[j + 3]); + } + y.d * x.d * s1 as f32 + y.s * x.m +} + +#[derive(Default)] +pub struct Stat { + pub sum: f64, + pub sumt: f64, + pub sumt2: f64, + pub maxt: f64, + pub nloop: i32, +} + +impl Stat { + pub fn add_result(&mut self, s: f64, t: f64) { + self.sum += s; + self.sumt += t; + self.sumt2 += t * t; + self.maxt = self.maxt.max(t); + self.nloop += 1; + } + + pub fn report(&self, title: &str) { + if self.nloop < 1 { + println!("report_result({title}): no result"); + return; + } + + let t = self.sumt / f64::from(self.nloop); + let mut dt = self.sumt2 / f64::from(self.nloop) - t * t; + if dt > 0.0 { + dt = dt.sqrt(); + } + + println!("============ {title}"); + println!(" = {}", self.sum / f64::from(self.nloop)); + println!("